at-container-registry/pkg/labeler/hub.go

package labeler

import (
	"sync"
)

// hubSubscriber is one connected subscribeLabels client. The hub fans out new labels
// to each subscriber's bounded channel; if a slow client fills the buffer, the hub
// drops them rather than blocking the writer.
type hubSubscriber struct {
	ch     chan *Label
	closed bool
}

// Hub broadcasts newly-inserted labels to all live subscribeLabels clients.
type Hub struct {
	mu   sync.Mutex
	subs map[*hubSubscriber]struct{}
}

// NewHub returns an empty hub ready to accept subscribers.
func NewHub() *Hub {
	return &Hub{subs: make(map[*hubSubscriber]struct{})}
}

// subscribe registers a new subscriber and returns its event channel + a cancel func.
// The buffer size bounds backpressure tolerance per client.
func (h *Hub) subscribe(buffer int) (*hubSubscriber, func()) {
	s := &hubSubscriber{ch: make(chan *Label, buffer)}
	h.mu.Lock()
	h.subs[s] = struct{}{}
	h.mu.Unlock()
	return s, func() { h.unsubscribe(s) }
}

func (h *Hub) unsubscribe(s *hubSubscriber) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if _, ok := h.subs[s]; !ok {
		return
	}
	delete(h.subs, s)
	if !s.closed {
		s.closed = true
		close(s.ch)
	}
}

// Broadcast sends a copy of the label to every live subscriber. Subscribers whose
// buffer is full are evicted on the spot rather than slowing down the writer.
func (h *Hub) Broadcast(l *Label) {
	if l == nil {
		return
	}
	h.mu.Lock()
	dead := make([]*hubSubscriber, 0)
	for s := range h.subs {
		select {
		case s.ch <- l:
		default:
			dead = append(dead, s)
		}
	}
	h.mu.Unlock()
	for _, s := range dead {
		h.unsubscribe(s)
	}
}

// Len returns the number of live subscribers (mostly for tests / metrics).
func (h *Hub) Len() int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return len(h.subs)
}