DDI/server/internal/agenthub/hub.go

// SPDX-License-Identifier: LicenseRef-DDI-Proprietary

// Package agenthub tracks the open Connect streams of the agents, so that
// the service layer can push commands to them and close the stream of a
// revoked agent (ADR 0021).
package agenthub

import (
	"crypto/rand"
	"fmt"
	"sync"

	agentv1 "git.beugsiud.de/beug.k/DDI/proto/gen/ddi/agent/v1"
)

// queueSize bounds the commands waiting for one stream. A full queue drops
// commands; the agent catches up through GetDesiredState.
const queueSize = 16

// Mode values (agents.mode).
const (
	ModeAutoPull = "auto_pull"
	ModeManual   = "manual"
)

// Reasons for closing a stream from the server side.
const (
	ReasonRevoked  = "agent revoked"
	ReasonDeleted  = "agent deleted"
	ReasonReplaced = "replaced by a newer connection"
	ReasonShutdown = "server shutting down"
	// ReasonModeChanged closes the stream after a mode change; the agent
	// reconnects and receives the new mode in Welcome (ADR 0022).
	ReasonModeChanged = "operating mode changed"
	reasonEnded       = "stream ended"
)

// Hub is the registry of open streams. It is safe for concurrent use.
type Hub struct {
	mu    sync.Mutex
	conns map[string]*Conn
}

// New creates an empty hub.
func New() *Hub { return &Hub{conns: map[string]*Conn{}} }

// Conn is one open stream.
type Conn struct {
	AgentID string
	Role    string
	Mode    string

	cmds   chan *agentv1.Command
	done   chan struct{}
	once   sync.Once
	reason string
}

// Commands delivers the commands to send on the stream.
func (c *Conn) Commands() <-chan *agentv1.Command { return c.cmds }

// Done is closed when the server ends the stream (revocation, deletion or a
// newer stream of the same agent).
func (c *Conn) Done() <-chan struct{} { return c.done }

// Reason says why Done was closed.
func (c *Conn) Reason() string {
	<-c.done
	return c.reason
}

func (c *Conn) close(reason string) {
	c.once.Do(func() {
		c.reason = reason
		close(c.done)
	})
}

// Attach registers a stream. An older stream of the same agent is closed.
func (h *Hub) Attach(agentID, role, mode string) *Conn {
	c := &Conn{AgentID: agentID, Role: role, Mode: mode, cmds: make(chan *agentv1.Command, queueSize), done: make(chan struct{})}
	h.mu.Lock()
	old := h.conns[agentID]
	h.conns[agentID] = c
	h.mu.Unlock()
	if old != nil {
		old.close(ReasonReplaced)
	}
	return c
}

// CloseAll closes every stream (shutdown).
func (h *Hub) CloseAll(reason string) {
	h.mu.Lock()
	conns := h.conns
	h.conns = map[string]*Conn{}
	h.mu.Unlock()
	for _, c := range conns {
		c.close(reason)
	}
}

// Detach removes c if it is still the registered stream of its agent.
func (h *Hub) Detach(c *Conn) {
	h.mu.Lock()
	if h.conns[c.AgentID] == c {
		delete(h.conns, c.AgentID)
	}
	h.mu.Unlock()
	c.close(reasonEnded)
}

// Connected reports whether the agent has an open stream.
func (h *Hub) Connected(agentID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	_, ok := h.conns[agentID]
	return ok
}

// Disconnect closes the agent's stream, if any.
func (h *Hub) Disconnect(agentID, reason string) {
	h.mu.Lock()
	c := h.conns[agentID]
	delete(h.conns, agentID)
	h.mu.Unlock()
	if c != nil {
		c.close(reason)
	}
}

// SetMode updates the mode of the agent's open stream (the mode the agent
// was told in Welcome), which decides whether new revisions are pushed.
func (h *Hub) SetMode(agentID, mode string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if c := h.conns[agentID]; c != nil {
		c.Mode = mode
	}
}

// Count returns the number of open streams per role.
func (h *Hub) Count() map[string]int {
	h.mu.Lock()
	defer h.mu.Unlock()
	out := map[string]int{}
	for _, c := range h.conns {
		out[c.Role]++
	}
	return out
}

// NotifyRevision sends an ApplyRevisionCommand to every connected agent in
// auto_pull mode and returns how many were notified.
func (h *Hub) NotifyRevision(revision uint64) int {
	h.mu.Lock()
	targets := make([]*Conn, 0, len(h.conns))
	for _, c := range h.conns {
		if c.Mode == ModeAutoPull {
			targets = append(targets, c)
		}
	}
	h.mu.Unlock()
	n := 0
	for _, c := range targets {
		if c.enqueue(ApplyCommand(revision, NewID())) {
			n++
		}
	}
	return n
}

// ApplyCommand builds an ApplyRevisionCommand with a new command ID.
func ApplyCommand(revision uint64, deploymentID string) *agentv1.Command {
	return &agentv1.Command{
		CommandId: NewID(),
		Action: &agentv1.Command_ApplyRevision{ApplyRevision: &agentv1.ApplyRevisionCommand{
			Revision: revision, DeploymentId: deploymentID,
		}},
	}
}

// Send queues cmd for the agent's stream and reports whether it has one
// (and its queue had room). Without, the agent picks the work up when it
// reconnects or polls.
func (h *Hub) Send(agentID string, cmd *agentv1.Command) bool {
	h.mu.Lock()
	c := h.conns[agentID]
	h.mu.Unlock()
	return c != nil && c.enqueue(cmd)
}

func (c *Conn) enqueue(cmd *agentv1.Command) bool {
	select {
	case c.cmds <- cmd:
		return true
	default: // queue full: the agent will pull
		return false
	}
}

// NewID returns a random UUID (version 4) in canonical lowercase form.
func NewID() string {
	var b [16]byte
	_, _ = rand.Read(b[:]) // never fails (crypto/rand)
	b[6] = b[6]&0x0f | 0x40
	b[8] = b[8]&0x3f | 0x80
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
}