121 lines
2.4 KiB
Go
121 lines
2.4 KiB
Go
package grpc
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
agentv1 "github.com/containarr/server/internal/proto/agentv1"
|
|
)
|
|
|
|
type AgentState struct {
|
|
ID string
|
|
Hostname string
|
|
Alias string
|
|
IPAddress string
|
|
Arch string
|
|
OS string
|
|
LastSeenAt time.Time
|
|
Containers []*agentv1.ContainerInfo
|
|
Images []*agentv1.ImageInfo
|
|
Volumes []*agentv1.VolumeInfo
|
|
Networks []*agentv1.NetworkInfo
|
|
|
|
cmdCh chan *agentv1.ServerMessage
|
|
}
|
|
|
|
type Registry struct {
|
|
mu sync.RWMutex
|
|
agents map[string]*AgentState
|
|
}
|
|
|
|
func NewRegistry() *Registry {
|
|
return &Registry{agents: make(map[string]*AgentState)}
|
|
}
|
|
|
|
func (r *Registry) Register(id, hostname, alias, ipAddress, arch, os string) *AgentState {
|
|
state := &AgentState{
|
|
ID: id,
|
|
Hostname: hostname,
|
|
Alias: alias,
|
|
IPAddress: ipAddress,
|
|
Arch: arch,
|
|
OS: os,
|
|
cmdCh: make(chan *agentv1.ServerMessage, 16),
|
|
}
|
|
r.mu.Lock()
|
|
r.agents[id] = state
|
|
r.mu.Unlock()
|
|
return state
|
|
}
|
|
|
|
func (r *Registry) Deregister(id string) {
|
|
r.mu.Lock()
|
|
if s, ok := r.agents[id]; ok {
|
|
close(s.cmdCh)
|
|
delete(r.agents, id)
|
|
}
|
|
r.mu.Unlock()
|
|
}
|
|
|
|
func (r *Registry) Get(id string) (*AgentState, bool) {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
s, ok := r.agents[id]
|
|
return s, ok
|
|
}
|
|
|
|
func (r *Registry) List() []*AgentState {
|
|
r.mu.RLock()
|
|
defer r.mu.RUnlock()
|
|
out := make([]*AgentState, 0, len(r.agents))
|
|
for _, s := range r.agents {
|
|
out = append(out, s)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (r *Registry) UpdateContainers(id string, containers []*agentv1.ContainerInfo) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if s, ok := r.agents[id]; ok {
|
|
s.Containers = containers
|
|
s.LastSeenAt = time.Now()
|
|
}
|
|
}
|
|
|
|
func (r *Registry) UpdateResources(id string, containers []*agentv1.ContainerInfo, images []*agentv1.ImageInfo, volumes []*agentv1.VolumeInfo, networks []*agentv1.NetworkInfo) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if s, ok := r.agents[id]; ok {
|
|
s.Containers = containers
|
|
s.Images = images
|
|
s.Volumes = volumes
|
|
s.Networks = networks
|
|
s.LastSeenAt = time.Now()
|
|
}
|
|
}
|
|
|
|
// UpdateAlias refreshes the alias for a live agent (called after an admin update).
|
|
func (r *Registry) UpdateAlias(id, alias string) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if s, ok := r.agents[id]; ok {
|
|
s.Alias = alias
|
|
}
|
|
}
|
|
|
|
func (r *Registry) Send(agentID string, msg *agentv1.ServerMessage) bool {
|
|
r.mu.RLock()
|
|
s, ok := r.agents[agentID]
|
|
r.mu.RUnlock()
|
|
if !ok {
|
|
return false
|
|
}
|
|
select {
|
|
case s.cmdCh <- msg:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|