64 lines
1.8 KiB
Go
64 lines
1.8 KiB
Go
package grpc
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
grpcpkg "google.golang.org/grpc"
|
|
proto "stream.api/internal/api/proto/agent/v1"
|
|
"stream.api/internal/dto"
|
|
"stream.api/internal/service"
|
|
)
|
|
|
|
type Server struct {
|
|
proto.UnimplementedWoodpeckerServer
|
|
proto.UnimplementedWoodpeckerAuthServer
|
|
jobService *service.JobService
|
|
agentManager *AgentManager
|
|
agentSecret string
|
|
sessions sync.Map
|
|
agentJobs sync.Map
|
|
onAgentEvent func(string, *dto.AgentWithStats)
|
|
}
|
|
|
|
func NewServer(jobService *service.JobService, agentSecret string) *Server {
|
|
return &Server{jobService: jobService, agentManager: NewAgentManager(), agentSecret: agentSecret}
|
|
}
|
|
|
|
func (s *Server) SetAgentEventHandler(handler func(string, *dto.AgentWithStats)) {
|
|
s.onAgentEvent = handler
|
|
}
|
|
|
|
func (s *Server) Register(grpcServer grpcpkg.ServiceRegistrar) {
|
|
proto.RegisterWoodpeckerServer(grpcServer, s)
|
|
proto.RegisterWoodpeckerAuthServer(grpcServer, s)
|
|
}
|
|
|
|
func (s *Server) SendCommand(agentID string, cmd string) bool {
|
|
return s.agentManager.SendCommand(agentID, cmd)
|
|
}
|
|
|
|
func (s *Server) ListAgents() []*dto.Agent { return s.agentManager.ListAll() }
|
|
|
|
func (s *Server) ListAgentsWithStats() []*dto.AgentWithStats {
|
|
agents := s.agentManager.ListAll()
|
|
result := make([]*dto.AgentWithStats, 0, len(agents))
|
|
for _, agent := range agents {
|
|
result = append(result, &dto.AgentWithStats{Agent: agent, ActiveJobCount: int64(len(s.getAgentJobs(agent.ID)))})
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (s *Server) getAgentWithStats(agentID string) *dto.AgentWithStats {
|
|
for _, agent := range s.ListAgentsWithStats() {
|
|
if agent != nil && agent.Agent != nil && agent.Agent.ID == agentID {
|
|
return agent
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *Server) Version(context.Context, *proto.Empty) (*proto.VersionResponse, error) {
|
|
return &proto.VersionResponse{GrpcVersion: 15, ServerVersion: "stream.api"}, nil
|
|
}
|