fm_be/internal/realtime/helicopterfile/hub.go

package helicopterfilerealtime

import (
	"errors"
	"log/slog"
	"strings"
	"sync"
)

const (
	DefaultEventType      = "helicopter_file_changed"
	DefaultSubscriberSize = 32
)

type Event struct {
	EventType         string `json:"event_type"`
	HelicopterID      string `json:"helicopter_id"`
	HelicopterFileID  string `json:"helicopter_file_id"`
	FileID            string `json:"file_id,omitempty"`
	Section           string `json:"section,omitempty"`
	Name              string `json:"name,omitempty"`
	MimeType          string `json:"mime_type,omitempty"`
	AttachmentPending bool   `json:"attachment_pending"`
	Position          int    `json:"position,omitempty"`
	IsMandatory       bool   `json:"is_mandatory,omitempty"`
	UpdatedAt         string `json:"updated_at,omitempty"`
}

type Subscription struct {
	events  <-chan Event
	done    <-chan struct{}
	closeFn func()
}

func (s *Subscription) Events() <-chan Event {
	if s == nil {
		return nil
	}
	return s.events
}

func (s *Subscription) Done() <-chan struct{} {
	if s == nil {
		return nil
	}
	return s.done
}

func (s *Subscription) Close() {
	if s == nil || s.closeFn == nil {
		return
	}
	s.closeFn()
}

type subscriber struct {
	events    chan Event
	done      chan struct{}
	match     func(Event) bool
	closeOnce sync.Once
}

func (s *subscriber) close() {
	if s == nil {
		return
	}
	s.closeOnce.Do(func() {
		close(s.done)
	})
}

type Hub struct {
	mu           sync.RWMutex
	subscribers  map[string]map[uint64]*subscriber
	recentByUser map[string][]Event
	nextID       uint64
	buffer       int
	replayLimit  int
	logger       *slog.Logger
}

func NewHub(buffer int, logger *slog.Logger) *Hub {
	if buffer <= 0 {
		buffer = DefaultSubscriberSize
	}
	if logger == nil {
		logger = slog.Default()
	}
	return &Hub{
		subscribers:  make(map[string]map[uint64]*subscriber),
		recentByUser: make(map[string][]Event),
		buffer:       buffer,
		replayLimit:  maxInt(buffer*4, DefaultSubscriberSize),
		logger:       logger.With("component", "helicopter-file-hub"),
	}
}

func (h *Hub) Subscribe(userUUID string) (*Subscription, error) {
	return h.SubscribeWithFilter(userUUID, nil)
}

func (h *Hub) SubscribeWithFilter(userUUID string, match func(Event) bool) (*Subscription, error) {
	if h == nil {
		return nil, errors.New("helicopter file hub is not configured")
	}
	userUUID = canonicalUserID(userUUID)
	if userUUID == "" {
		return nil, errors.New("user uuid is required")
	}

	sub := &subscriber{
		events: make(chan Event, h.buffer),
		done:   make(chan struct{}),
		match:  match,
	}

	h.mu.Lock()
	h.nextID++
	subID := h.nextID
	if h.subscribers[userUUID] == nil {
		h.subscribers[userUUID] = make(map[uint64]*subscriber)
	}
	h.subscribers[userUUID][subID] = sub
	backlog := append([]Event(nil), h.recentByUser[userUUID]...)
	h.mu.Unlock()

	for i := range backlog {
		if sub.match != nil && !sub.match(backlog[i]) {
			continue
		}
		enqueueSubscriberEvent(sub, backlog[i])
	}

	var closeOnce sync.Once
	closeFn := func() {
		closeOnce.Do(func() {
			h.mu.Lock()
			defer h.mu.Unlock()

			if byUser, ok := h.subscribers[userUUID]; ok {
				delete(byUser, subID)
				if len(byUser) == 0 {
					delete(h.subscribers, userUUID)
				}
			}
			sub.close()
		})
	}

	return &Subscription{
		events:  sub.events,
		done:    sub.done,
		closeFn: closeFn,
	}, nil
}

func (h *Hub) PublishToUsers(userUUIDs []string, event Event) int {
	if h == nil {
		return 0
	}
	if strings.TrimSpace(event.EventType) == "" {
		event.EventType = DefaultEventType
	}
	if strings.TrimSpace(event.HelicopterID) == "" || strings.TrimSpace(event.HelicopterFileID) == "" {
		return 0
	}

	uniqueUsers := make(map[string]struct{}, len(userUUIDs))
	for i := range userUUIDs {
		userID := canonicalUserID(userUUIDs[i])
		if userID == "" {
			continue
		}
		uniqueUsers[userID] = struct{}{}
	}
	if len(uniqueUsers) == 0 {
		return 0
	}

	delivered := 0
	for userID := range uniqueUsers {
		h.mu.Lock()
		h.recentByUser[userID] = appendRecentEvents(h.recentByUser[userID], event, h.replayLimit)
		h.mu.Unlock()

		h.mu.RLock()
		byUser := h.subscribers[userID]
		if len(byUser) == 0 {
			h.mu.RUnlock()
			continue
		}

		for _, sub := range byUser {
			select {
			case <-sub.done:
				continue
			default:
			}
			if sub.match != nil && !sub.match(event) {
				continue
			}

			if enqueueSubscriberEvent(sub, event) {
				delivered++
			}
		}
		h.mu.RUnlock()
	}
	return delivered
}

func canonicalUserID(raw string) string {
	return strings.ToLower(strings.TrimSpace(raw))
}

func enqueueSubscriberEvent(sub *subscriber, event Event) bool {
	if sub == nil {
		return false
	}
	select {
	case <-sub.done:
		return false
	default:
	}

	select {
	case sub.events <- event:
		return true
	default:
		select {
		case <-sub.events:
		default:
		}
		select {
		case sub.events <- event:
			return true
		default:
			return false
		}
	}
}

func appendRecentEvents(current []Event, event Event, limit int) []Event {
	if limit <= 0 {
		return current[:0]
	}
	current = append(current, event)
	if len(current) <= limit {
		return current
	}
	start := len(current) - limit
	copy(current, current[start:])
	return current[:limit]
}

func maxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}