so-many-crows/hub.go

package main

import (
	"encoding/json"
	"sync"
	"time"
)

// sseMsg is one event, already encoded, numbered within its list.
type sseMsg struct {
	ID   int
	Name string
	Data []byte
}

// listStream holds everything one list has said so far, plus who is listening.
type listStream struct {
	requester string
	created   time.Time
	finished  bool
	history   []sseMsg
	subs      map[chan sseMsg]struct{}
}

// Hub fans events out to SSE subscribers per list_id and keeps each list's
// history in memory so late subscribers can catch up.
type Hub struct {
	mu    sync.Mutex
	lists map[string]*listStream
}

func NewHub() *Hub {
	return &Hub{lists: map[string]*listStream{}}
}

// subBuffer is how far a subscriber may fall behind before it is cut off.
// It reconnects with Last-Event-ID and catches up from history.
const subBuffer = 256

// keepLists caps how many lists the hub remembers; the oldest finished go first.
const keepLists = 200

func (h *Hub) Open(listID, requester string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.prune()
	h.lists[listID] = &listStream{requester: requester, created: time.Now(), subs: map[chan sseMsg]struct{}{}}
}

func (h *Hub) prune() {
	for len(h.lists) >= keepLists {
		var oldest string
		for id, ls := range h.lists {
			if ls.finished && (oldest == "" || ls.created.Before(h.lists[oldest].created)) {
				oldest = id
			}
		}
		if oldest == "" {
			return // every list is still flying; keep them all
		}
		h.drop(oldest)
	}
}

func (h *Hub) drop(id string) {
	for ch := range h.lists[id].subs {
		close(ch)
	}
	delete(h.lists, id)
}

// ActiveLists returns the requester's lists that are still flying.
func (h *Hub) ActiveLists(requester string) []string {
	h.mu.Lock()
	defer h.mu.Unlock()
	var ids []string
	for id, ls := range h.lists {
		if ls.requester == requester && !ls.finished {
			ids = append(ids, id)
		}
	}
	return ids
}

// Publish records an event for a list and sends it to current subscribers.
func (h *Hub) Publish(listID, name string, fields map[string]any) {
	payload := map[string]any{}
	for k, v := range fields {
		payload[k] = v
	}
	payload["list_id"] = listID
	payload["ts"] = time.Now().UTC().Format(time.RFC3339)
	data, err := json.Marshal(payload)
	if err != nil {
		return
	}

	h.mu.Lock()
	defer h.mu.Unlock()
	ls := h.lists[listID]
	if ls == nil {
		return
	}
	m := sseMsg{ID: len(ls.history) + 1, Name: name, Data: data}
	ls.history = append(ls.history, m)
	if name == "list.finished" {
		ls.finished = true
	}
	for ch := range ls.subs {
		select {
		case ch <- m:
		default: // too slow: cut it off rather than stall the flock
			delete(ls.subs, ch)
			close(ch)
		}
	}
}

// Subscribe returns the history after lastID and a channel for what comes next.
// The channel closes if the subscriber falls too far behind.
func (h *Hub) Subscribe(listID string, lastID int) ([]sseMsg, chan sseMsg, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	ls := h.lists[listID]
	if ls == nil {
		return nil, nil, false
	}
	var past []sseMsg
	if lastID < len(ls.history) {
		past = append(past, ls.history[max(lastID, 0):]...)
	}
	ch := make(chan sseMsg, subBuffer)
	ls.subs[ch] = struct{}{}
	return past, ch, true
}

func (h *Hub) Unsubscribe(listID string, ch chan sseMsg) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if ls := h.lists[listID]; ls != nil {
		if _, ok := ls.subs[ch]; ok {
			delete(ls.subs, ch)
			close(ch)
		}
	}
}

// Emitter binds the hub to one list.
func (h *Hub) Emitter(listID string) Emitter { return hubEmitter{h, listID} }

type hubEmitter struct {
	h  *Hub
	id string
}

func (e hubEmitter) Emit(name string, fields map[string]any) { e.h.Publish(e.id, name, fields) }