package stream import ( "context" "encoding/json" "fmt" "os" "sync" "time" "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" ) type DesktopTaskEvent struct { Type string `json:"type"` TaskID string `json:"task_id"` JobID string `json:"job_id"` WorkspaceID int64 `json:"workspace_id"` TargetClientID string `json:"target_client_id"` Status string `json:"status"` Kind string `json:"kind"` UpdatedAt time.Time `json:"updated_at"` } type desktopClientStream struct { subscribers map[int]chan DesktopTaskEvent nextSubscriber int } type DesktopTaskHub struct { mu sync.Mutex streams map[string]*desktopClientStream rabbitMQ *rabbitmq.Client instanceID string consumerName string } func NewDesktopTaskHub(rabbitMQClient *rabbitmq.Client) *DesktopTaskHub { instanceID := fmt.Sprintf("%d-%d", os.Getpid(), time.Now().UnixNano()) return &DesktopTaskHub{ streams: make(map[string]*desktopClientStream), rabbitMQ: rabbitMQClient, instanceID: instanceID, consumerName: "desktop-task-stream-" + instanceID, } } func (h *DesktopTaskHub) Run(ctx context.Context) { if h == nil || h.rabbitMQ == nil { return } go h.run(ctx) } func (h *DesktopTaskHub) Publish(event DesktopTaskEvent) { if h == nil || event.TargetClientID == "" { return } h.mu.Lock() stream := h.ensureLocked(event.TargetClientID) subscribers := make([]chan DesktopTaskEvent, 0, len(stream.subscribers)) for _, subscriber := range stream.subscribers { subscribers = append(subscribers, subscriber) } h.mu.Unlock() for _, subscriber := range subscribers { select { case subscriber <- event: default: } } } func (h *DesktopTaskHub) Subscribe(clientID string) (<-chan DesktopTaskEvent, func()) { if h == nil || clientID == "" { return nil, func() {} } h.mu.Lock() stream := h.ensureLocked(clientID) subscriberID := stream.nextSubscriber stream.nextSubscriber++ ch := make(chan DesktopTaskEvent, 64) stream.subscribers[subscriberID] = ch h.mu.Unlock() cancel := func() { h.mu.Lock() defer h.mu.Unlock() current := h.streams[clientID] if current == nil { close(ch) return } delete(current.subscribers, subscriberID) if len(current.subscribers) == 0 { delete(h.streams, clientID) } close(ch) } return ch, cancel } func (h *DesktopTaskHub) run(ctx context.Context) { for { deliveries, ch, err := h.rabbitMQ.ConsumeDesktopTaskEvents(h.consumerName) if err != nil { 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.Body) _ = delivery.Ack(false) } } _ = ch.Close() } } func (h *DesktopTaskHub) handleDelivery(body []byte) { var event DesktopTaskEvent if err := json.Unmarshal(body, &event); err != nil { return } h.Publish(event) } func (h *DesktopTaskHub) ensureLocked(clientID string) *desktopClientStream { stream := h.streams[clientID] if stream == nil { stream = &desktopClientStream{ subscribers: make(map[int]chan DesktopTaskEvent), } h.streams[clientID] = stream } return stream }