feat: big ahh commit
- 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
This commit is contained in:
@@ -0,0 +1,111 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user