geo/server/internal/shared/stream/desktop_dispatch.go

package stream

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq"
)

// DesktopDispatchEvent is the payload pushed to a desktop client over the
// dispatch WebSocket. The shape intentionally mirrors DesktopTaskEvent so that
// the existing SSE handler on the desktop client can treat both channels the
// same way and dedupe by task_id.
type DesktopDispatchEvent struct {
	Type                string    `json:"type"`
	TaskID              string    `json:"task_id"`
	JobID               string    `json:"job_id,omitempty"`
	WorkspaceID         int64     `json:"workspace_id"`
	TargetAccountID     string    `json:"target_account_id,omitempty"`
	TargetClientID      string    `json:"target_client_id"`
	Platform            string    `json:"platform,omitempty"`
	Title               *string   `json:"title,omitempty"`
	BusinessDate        *string   `json:"business_date,omitempty"`
	SchedulerGroupKey   *string   `json:"scheduler_group_key,omitempty"`
	QuestionText        *string   `json:"question_text,omitempty"`
	Status              string    `json:"status"`
	Kind                string    `json:"kind"`
	Priority            int       `json:"priority,omitempty"`
	Lane                string    `json:"lane,omitempty"`
	InterruptGeneration int       `json:"interrupt_generation,omitempty"`
	SignalOnly          bool      `json:"signal_only,omitempty"`
	Control             string    `json:"control,omitempty"`
	Reason              string    `json:"reason,omitempty"`
	ReplacementTaskID   string    `json:"replacement_task_id,omitempty"`
	UpdatedAt           time.Time `json:"updated_at"`
}

// DispatchSubscriber represents a live WebSocket connection bound to a single
// desktop client id. A client may have multiple subscribers (e.g. during a
// reconnect race), so the hub stores them as a set per client id.
type DispatchSubscriber struct {
	id     uint64
	out    chan DesktopDispatchEvent
	closed chan struct{}
}

// Outbound returns the channel the handler goroutine should read from to write
// frames to the WebSocket. The channel is closed when the subscriber is
// unregistered; readers should range over it.
func (s *DispatchSubscriber) Outbound() <-chan DesktopDispatchEvent {
	return s.out
}

type dispatchClientStream struct {
	subscribers map[uint64]*DispatchSubscriber
}

// DesktopDispatchHub delivers per-client dispatch messages pulled from the
// desktop.task.dispatch topic exchange. Every tenant-api instance owns its own
// autodelete queue bound to the exchange, so every instance sees every message
// and routes it to the subscribers it currently owns. Clients not connected to
// this instance are simply filtered out; the client reconnects via the load
// balancer and lands on whichever instance currently holds its WebSocket.
type DesktopDispatchHub struct {
	mu           sync.RWMutex
	clients      map[string]*dispatchClientStream
	rabbitMQ     *rabbitmq.Client
	logger       *zap.Logger
	instanceID   string
	consumerName string
	bindingKeys  []string
	nextSubID    uint64
}

// NewDesktopDispatchHub wires the hub to the shared rabbitmq client. The
// supplied binding keys decide which dispatch routing keys this instance
// receives; initially that is just `publish.*`, and the monitor slice will be
// added when the monitor chain is migrated.
func NewDesktopDispatchHub(rabbitMQClient *rabbitmq.Client, logger *zap.Logger, bindingKeys []string) *DesktopDispatchHub {
	instanceID := fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano())
	keys := make([]string, 0, len(bindingKeys))
	for _, key := range bindingKeys {
		trimmed := strings.TrimSpace(key)
		if trimmed == "" {
			continue
		}
		keys = append(keys, trimmed)
	}
	if len(keys) == 0 {
		keys = []string{"publish.*"}
	}
	return &DesktopDispatchHub{
		clients:      make(map[string]*dispatchClientStream),
		rabbitMQ:     rabbitMQClient,
		logger:       logger,
		instanceID:   instanceID,
		consumerName: "desktop-dispatch-" + instanceID,
		bindingKeys:  keys,
	}
}

// Run starts the background AMQP consumer. It is safe to call Run on a hub
// without a rabbitmq client — the hub simply behaves as an in-process broker
// used in tests, useful for unit-level wiring.
func (h *DesktopDispatchHub) Run(ctx context.Context) {
	if h == nil || h.rabbitMQ == nil {
		return
	}
	go h.run(ctx)
}

// Subscribe registers a new WebSocket consumer for the given client id. The
// returned cancel closes the outbound channel and removes the subscriber; the
// WebSocket handler must call it on disconnect.
func (h *DesktopDispatchHub) Subscribe(clientID string) (*DispatchSubscriber, func()) {
	if h == nil || clientID == "" {
		return nil, func() {}
	}

	h.mu.Lock()
	stream, ok := h.clients[clientID]
	if !ok {
		stream = &dispatchClientStream{subscribers: make(map[uint64]*DispatchSubscriber)}
		h.clients[clientID] = stream
	}
	h.nextSubID++
	sub := &DispatchSubscriber{
		id:     h.nextSubID,
		out:    make(chan DesktopDispatchEvent, 32),
		closed: make(chan struct{}),
	}
	stream.subscribers[sub.id] = sub
	h.mu.Unlock()

	cancel := func() {
		h.mu.Lock()
		stream, ok := h.clients[clientID]
		if !ok {
			h.mu.Unlock()
			return
		}
		existing, present := stream.subscribers[sub.id]
		if !present {
			h.mu.Unlock()
			return
		}
		delete(stream.subscribers, sub.id)
		if len(stream.subscribers) == 0 {
			delete(h.clients, clientID)
		}
		h.mu.Unlock()

		select {
		case <-existing.closed:
			// already closed
		default:
			close(existing.closed)
			close(existing.out)
		}
	}

	return sub, cancel
}

// Publish delivers the event to every subscriber currently registered for the
// event's target client id on this instance. Events for clients that are not
// on this instance are dropped silently, by design.
func (h *DesktopDispatchHub) Publish(event DesktopDispatchEvent) {
	if h == nil || event.TargetClientID == "" {
		return
	}
	h.mu.RLock()
	stream, ok := h.clients[event.TargetClientID]
	if !ok {
		h.mu.RUnlock()
		return
	}
	subs := make([]*DispatchSubscriber, 0, len(stream.subscribers))
	for _, sub := range stream.subscribers {
		subs = append(subs, sub)
	}
	h.mu.RUnlock()

	for _, sub := range subs {
		select {
		case <-sub.closed:
			continue
		case sub.out <- event:
		default:
			// Subscriber's outbound channel is full: the reader has fallen
			// behind. We drop the message rather than block the dispatch
			// consumer — the client will resync via the lease API on
			// reconnect.
			if h.logger != nil {
				h.logger.Warn("desktop dispatch subscriber slow, dropping event",
					zap.String("client_id", event.TargetClientID),
					zap.String("task_id", event.TaskID),
				)
			}
		}
	}
}

// ActiveClientIDs returns the set of client ids currently connected on this
// instance. Intended for diagnostics and the /api/health/ready endpoint.
func (h *DesktopDispatchHub) ActiveClientIDs() []string {
	if h == nil {
		return nil
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	out := make([]string, 0, len(h.clients))
	for id := range h.clients {
		out = append(out, id)
	}
	return out
}

func (h *DesktopDispatchHub) run(ctx context.Context) {
	for {
		deliveries, ch, err := h.rabbitMQ.ConsumeDesktopDispatch(h.consumerName, h.bindingKeys)
		if err != nil {
			if h.logger != nil {
				h.logger.Warn("desktop dispatch consume failed, retrying",
					zap.Error(err),
				)
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(2 * time.Second):
				continue
			}
		}

		closed := false
		for !closed {
			select {
			case <-ctx.Done():
				_ = ch.Close()
				return
			case delivery, ok := <-deliveries:
				if !ok {
					closed = true
					continue
				}
				h.handleDelivery(delivery)
			}
		}
		_ = ch.Close()
	}
}

func (h *DesktopDispatchHub) handleDelivery(delivery amqp.Delivery) {
	var event DesktopDispatchEvent
	if err := json.Unmarshal(delivery.Body, &event); err != nil {
		if h.logger != nil {
			h.logger.Warn("desktop dispatch event decode failed",
				zap.Error(err),
				zap.String("routing_key", delivery.RoutingKey),
			)
		}
		return
	}
	if event.TargetClientID == "" {
		// Fall back to the routing key: `{prefix}.{client_id}` -> client_id.
		if idx := strings.LastIndex(delivery.RoutingKey, "."); idx >= 0 && idx < len(delivery.RoutingKey)-1 {
			event.TargetClientID = delivery.RoutingKey[idx+1:]
		}
	}
	h.Publish(event)
}