feat(desktop): push publish tasks via AMQP topic dispatch WebSocket
Introduce a desktop.task.dispatch topic exchange with per-client routing keys. Tenant-api publishes a task_available frame keyed by target client ID when a publish job is created; every instance binds its own transient queue and forwards to the matching WebSocket (/api/desktop/dispatch) it owns. Desktop-client now prefers this push channel and only falls back to HTTP /lease polling when the socket is down. Also drop admin-web monitoring-plugin remnants and scope the publish modal account list to publish platforms. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,274 @@
|
||||
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"`
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
amqp "github.com/rabbitmq/amqp091-go"
|
||||
)
|
||||
|
||||
func TestDispatchHub_PublishRoutesToLocalSubscriber(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"})
|
||||
sub, cancel := hub.Subscribe("client-a")
|
||||
defer cancel()
|
||||
|
||||
event := DesktopDispatchEvent{
|
||||
Type: "task_available",
|
||||
TaskID: "task-1",
|
||||
TargetClientID: "client-a",
|
||||
Kind: "publish",
|
||||
Status: "queued",
|
||||
}
|
||||
hub.Publish(event)
|
||||
|
||||
select {
|
||||
case got := <-sub.Outbound():
|
||||
if got.TaskID != "task-1" {
|
||||
t.Fatalf("unexpected task id: %s", got.TaskID)
|
||||
}
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
t.Fatal("expected to receive event within 200ms")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchHub_DropsEventsForUnknownClient(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"})
|
||||
|
||||
// Subscribe for a different client, publish for another one.
|
||||
_, cancel := hub.Subscribe("client-b")
|
||||
defer cancel()
|
||||
|
||||
hub.Publish(DesktopDispatchEvent{
|
||||
Type: "task_available",
|
||||
TargetClientID: "client-a",
|
||||
})
|
||||
|
||||
ids := hub.ActiveClientIDs()
|
||||
if len(ids) != 1 || ids[0] != "client-b" {
|
||||
t.Fatalf("expected only client-b subscribed, got %v", ids)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchHub_CancelRemovesSubscriber(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"})
|
||||
sub, cancel := hub.Subscribe("client-a")
|
||||
|
||||
cancel()
|
||||
|
||||
// Subsequent publish must not panic even though the subscriber is gone;
|
||||
// channel is closed, reads must yield zero value with !ok.
|
||||
hub.Publish(DesktopDispatchEvent{TargetClientID: "client-a"})
|
||||
|
||||
select {
|
||||
case _, ok := <-sub.Outbound():
|
||||
if ok {
|
||||
t.Fatal("expected outbound channel to be closed after cancel")
|
||||
}
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
t.Fatal("expected closed channel to yield immediately")
|
||||
}
|
||||
|
||||
if ids := hub.ActiveClientIDs(); len(ids) != 0 {
|
||||
t.Fatalf("expected no subscribers after cancel, got %v", ids)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchHub_MultipleSubscribersSameClient(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"})
|
||||
subA, cancelA := hub.Subscribe("client-a")
|
||||
defer cancelA()
|
||||
subB, cancelB := hub.Subscribe("client-a")
|
||||
defer cancelB()
|
||||
|
||||
hub.Publish(DesktopDispatchEvent{
|
||||
Type: "task_available",
|
||||
TaskID: "task-1",
|
||||
TargetClientID: "client-a",
|
||||
})
|
||||
|
||||
for i, ch := range []<-chan DesktopDispatchEvent{subA.Outbound(), subB.Outbound()} {
|
||||
select {
|
||||
case got := <-ch:
|
||||
if got.TaskID != "task-1" {
|
||||
t.Fatalf("subscriber %d got unexpected task id: %s", i, got.TaskID)
|
||||
}
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
t.Fatalf("subscriber %d did not receive event", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchHub_HandleDeliveryFallsBackToRoutingKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
hub := NewDesktopDispatchHub(nil, nil, []string{"publish.*"})
|
||||
sub, cancel := hub.Subscribe("client-c")
|
||||
defer cancel()
|
||||
|
||||
// Body intentionally omits target_client_id; the hub should recover it
|
||||
// from the routing key "<prefix>.<client_id>".
|
||||
delivery := amqp.Delivery{
|
||||
RoutingKey: "publish.client-c",
|
||||
Body: []byte(`{"type":"task_available","task_id":"task-9","kind":"publish","status":"queued"}`),
|
||||
}
|
||||
hub.handleDelivery(delivery)
|
||||
|
||||
select {
|
||||
case got := <-sub.Outbound():
|
||||
if got.TargetClientID != "client-c" {
|
||||
t.Fatalf("expected fallback to set target client id, got %q", got.TargetClientID)
|
||||
}
|
||||
if got.TaskID != "task-9" {
|
||||
t.Fatalf("unexpected task id: %s", got.TaskID)
|
||||
}
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
t.Fatal("expected fallback-routed event")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user