chore: replace wsconncache with a single tailnet (#8176)

tailnet/multiagent.go

@@ -0,0 +1,167 @@
+package tailnet
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
+)
+
+type MultiAgentConn interface {
+	UpdateSelf(node *Node) error
+	SubscribeAgent(agentID uuid.UUID) error
+	UnsubscribeAgent(agentID uuid.UUID) error
+	NextUpdate(ctx context.Context) ([]*Node, bool)
+	AgentIsLegacy(agentID uuid.UUID) bool
+	Close() error
+	IsClosed() bool
+}
+
+type MultiAgent struct {
+	mu sync.RWMutex
+
+	closed bool
+
+	ID     uuid.UUID
+	Logger slog.Logger
+
+	AgentIsLegacyFunc func(agentID uuid.UUID) bool
+	OnSubscribe       func(enq Queue, agent uuid.UUID) (*Node, error)
+	OnUnsubscribe     func(enq Queue, agent uuid.UUID) error
+	OnNodeUpdate      func(id uuid.UUID, node *Node) error
+	OnRemove          func(id uuid.UUID)
+
+	updates   chan []*Node
+	closeOnce sync.Once
+	start     int64
+	lastWrite int64
+	// Client nodes normally generate a unique id for each connection so
+	// overwrites are really not an issue, but is provided for compatibility.
+	overwrites int64
+}
+
+func (m *MultiAgent) Init() *MultiAgent {
+	m.updates = make(chan []*Node, 128)
+	m.start = time.Now().Unix()
+	return m
+}
+
+func (m *MultiAgent) UniqueID() uuid.UUID {
+	return m.ID
+}
+
+func (m *MultiAgent) AgentIsLegacy(agentID uuid.UUID) bool {
+	return m.AgentIsLegacyFunc(agentID)
+}
+
+var ErrMultiAgentClosed = xerrors.New("multiagent is closed")
+
+func (m *MultiAgent) UpdateSelf(node *Node) error {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	if m.closed {
+		return ErrMultiAgentClosed
+	}
+
+	return m.OnNodeUpdate(m.ID, node)
+}
+
+func (m *MultiAgent) SubscribeAgent(agentID uuid.UUID) error {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	if m.closed {
+		return ErrMultiAgentClosed
+	}
+
+	node, err := m.OnSubscribe(m, agentID)
+	if err != nil {
+		return err
+	}
+
+	if node != nil {
+		return m.enqueueLocked([]*Node{node})
+	}
+
+	return nil
+}
+
+func (m *MultiAgent) UnsubscribeAgent(agentID uuid.UUID) error {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	if m.closed {
+		return ErrMultiAgentClosed
+	}
+
+	return m.OnUnsubscribe(m, agentID)
+}
+
+func (m *MultiAgent) NextUpdate(ctx context.Context) ([]*Node, bool) {
+	select {
+	case <-ctx.Done():
+		return nil, false
+
+	case nodes, ok := <-m.updates:
+		return nodes, ok
+	}
+}
+
+func (m *MultiAgent) Enqueue(nodes []*Node) error {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	if m.closed {
+		return nil
+	}
+
+	return m.enqueueLocked(nodes)
+}
+
+func (m *MultiAgent) enqueueLocked(nodes []*Node) error {
+	atomic.StoreInt64(&m.lastWrite, time.Now().Unix())
+
+	select {
+	case m.updates <- nodes:
+		return nil
+	default:
+		return ErrWouldBlock
+	}
+}
+
+func (m *MultiAgent) Name() string {
+	return m.ID.String()
+}
+
+func (m *MultiAgent) Stats() (start int64, lastWrite int64) {
+	return m.start, atomic.LoadInt64(&m.lastWrite)
+}
+
+func (m *MultiAgent) Overwrites() int64 {
+	return m.overwrites
+}
+
+func (m *MultiAgent) IsClosed() bool {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.closed
+}
+
+func (m *MultiAgent) CoordinatorClose() error {
+	m.mu.Lock()
+	if !m.closed {
+		m.closed = true
+		close(m.updates)
+	}
+	m.mu.Unlock()
+	return nil
+}
+
+func (m *MultiAgent) Close() error {
+	_ = m.CoordinatorClose()
+	m.closeOnce.Do(func() { m.OnRemove(m.ID) })
+	return nil
+}