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

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
}