shinyzero0
7be99f8e91
- agent+proto+backend: transfer service status - agent: fix returning empty message on nonzero exit status - backend: refactor collector+commander and handlers dependent on them: implement agent accounting via grpc stats handler
112 lines
2.4 KiB
Go
112 lines
2.4 KiB
Go
package collector
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
|
|
"google.golang.org/grpc/metadata"
|
|
"google.golang.org/grpc/stats"
|
|
)
|
|
|
|
// ConnTracker tracks connected Collector agents and handles cleanup on disconnect.
|
|
// It implements grpc.StatsHandler for disconnect detection.
|
|
type ConnTracker struct {
|
|
mu sync.RWMutex
|
|
agents map[string]*Agent
|
|
}
|
|
|
|
func NewConnTracker() *ConnTracker {
|
|
return &ConnTracker{
|
|
agents: make(map[string]*Agent),
|
|
}
|
|
}
|
|
|
|
// Register adds an agent to the tracker. Called by Collector.Stream().
|
|
func (t *ConnTracker) Register(agent *Agent) {
|
|
t.mu.Lock()
|
|
t.agents[agent.ID] = agent
|
|
t.mu.Unlock()
|
|
log.Printf("[collector] agent registered: %s", agent.ID)
|
|
}
|
|
|
|
// Unregister removes an agent from the tracker.
|
|
func (t *ConnTracker) Unregister(id string) {
|
|
t.mu.Lock()
|
|
delete(t.agents, id)
|
|
t.mu.Unlock()
|
|
log.Printf("[collector] agent unregistered: %s", id)
|
|
}
|
|
|
|
// GetAgent returns the agent for the given ID.
|
|
func (t *ConnTracker) GetAgent(id string) (*Agent, bool) {
|
|
t.mu.RLock()
|
|
defer t.mu.RUnlock()
|
|
a, ok := t.agents[id]
|
|
return a, ok
|
|
}
|
|
|
|
// Agents returns all connected agents.
|
|
func (t *ConnTracker) Agents() []*Agent {
|
|
t.mu.RLock()
|
|
defer t.mu.RUnlock()
|
|
result := make([]*Agent, 0, len(t.agents))
|
|
for _, a := range t.agents {
|
|
result = append(result, a)
|
|
}
|
|
return result
|
|
}
|
|
|
|
// grpc.StatsHandler implementation.
|
|
|
|
func (t *ConnTracker) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
|
|
return ctx
|
|
}
|
|
|
|
func (t *ConnTracker) HandleRPC(ctx context.Context, _ stats.RPCStats) {}
|
|
|
|
func (t *ConnTracker) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
|
|
return ctx
|
|
}
|
|
|
|
func (t *ConnTracker) HandleConn(ctx context.Context, s stats.ConnStats) {
|
|
switch s.(type) {
|
|
case *stats.ConnEnd:
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
|
if !ok {
|
|
return
|
|
}
|
|
whoamiVals := md["whoami"]
|
|
if len(whoamiVals) == 0 {
|
|
return
|
|
}
|
|
t.Unregister(whoamiVals[0])
|
|
}
|
|
}
|
|
|
|
// UpdateServices updates the services list for the given agent.
|
|
func (t *ConnTracker) UpdateServices(id string, services []Service) bool {
|
|
t.mu.Lock()
|
|
defer t.mu.Unlock()
|
|
agent, ok := t.agents[id]
|
|
if !ok {
|
|
return false
|
|
}
|
|
agent.Services = services
|
|
return true
|
|
}
|
|
|
|
// Service represents a named service with its current status.
|
|
type Service struct {
|
|
Name, Status string
|
|
}
|
|
|
|
// Agent represents a connected agent streaming logs to the collector.
|
|
type Agent struct {
|
|
ID string
|
|
Label string
|
|
Services []Service
|
|
ConnectedAt time.Time
|
|
}
|