geo/server/internal/tenant/transport/desktop_dispatch_handler.go

package transport

import (
	"context"
	"encoding/json"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/bootstrap"
	"github.com/geo-platform/tenant-api/internal/shared/stream"
	tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app"
	"github.com/geo-platform/tenant-api/internal/tenant/repository"
)

const (
	dispatchWriteWait       = 10 * time.Second
	dispatchPongWait        = 60 * time.Second
	dispatchPingPeriod      = 25 * time.Second
	dispatchReadLimitBytes  = 2048
	dispatchCloseGracePause = 500 * time.Millisecond
	dispatchReplayLimit     = 20
)

type DesktopDispatchHandler struct {
	app      *bootstrap.App
	upgrader websocket.Upgrader
}

func NewDesktopDispatchHandler(a *bootstrap.App) *DesktopDispatchHandler {
	return &DesktopDispatchHandler{
		app: a,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin: func(_ *http.Request) bool {
				// Origin enforcement happens at the gateway / LB layer; this
				// endpoint is already gated by the bearer token middleware.
				return true
			},
		},
	}
}

// Dispatch upgrades the HTTP request into a WebSocket and streams dispatch
// events for the current desktop client. The contract intentionally mirrors
// the existing SSE events endpoint: connect with a bearer token, receive JSON
// frames whose shape matches DesktopDispatchEvent, and reconnect on drop.
func (h *DesktopDispatchHandler) Dispatch(c *gin.Context) {
	client := MustDesktopClient(c.Request.Context())
	hub := h.app.DesktopDispatch
	if hub == nil {
		c.AbortWithStatus(http.StatusServiceUnavailable)
		return
	}

	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		if h.app.Logger != nil {
			h.app.Logger.Warn("desktop dispatch upgrade failed",
				zap.String("client_id", client.ID.String()),
				zap.Error(err),
			)
		}
		return
	}

	subscriber, cancel := hub.Subscribe(client.ID.String())
	defer cancel()
	tenantapp.MarkDesktopTaskConsumerPresent(c.Request.Context(), h.app.Redis, client.ID)

	conn.SetReadLimit(dispatchReadLimitBytes)
	_ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait))
	conn.SetPongHandler(func(string) error {
		_ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait))
		return nil
	})

	// Reader goroutine: drains control frames and signals disconnect. We do
	// not expect application data from the client over this socket.
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	serverTime := time.Now().UTC()
	connected := stream.DesktopDispatchEvent{
		Type:           "connected",
		TargetClientID: client.ID.String(),
		WorkspaceID:    client.WorkspaceID,
		UpdatedAt:      serverTime,
		ServerTime:     &serverTime,
	}
	if err := writeDispatchFrame(conn, connected); err != nil {
		_ = conn.Close()
		return
	}
	h.replayQueuedPublishTasks(c.Request.Context(), client, conn)

	pingTicker := time.NewTicker(dispatchPingPeriod)
	defer pingTicker.Stop()

	defer func() {
		_ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait))
		_ = conn.WriteMessage(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		time.Sleep(dispatchCloseGracePause)
		_ = conn.Close()
	}()

	outbound := subscriber.Outbound()
	for {
		select {
		case <-c.Request.Context().Done():
			return
		case <-readerDone:
			return
		case event, ok := <-outbound:
			if !ok {
				return
			}
			if err := writeDispatchFrame(conn, event); err != nil {
				return
			}
		case <-pingTicker.C:
			_ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait))
			if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func writeDispatchFrame(conn *websocket.Conn, event stream.DesktopDispatchEvent) error {
	body, err := json.Marshal(event)
	if err != nil {
		return err
	}
	_ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait))
	return conn.WriteMessage(websocket.TextMessage, body)
}

func (h *DesktopDispatchHandler) replayQueuedPublishTasks(ctx context.Context, client *repository.DesktopClient, conn *websocket.Conn) {
	if h == nil || h.app == nil || h.app.DB == nil || client == nil || conn == nil {
		return
	}

	rows, err := h.app.DB.Query(ctx, `
		SELECT
			dt.desktop_id,
			dt.job_id,
			dt.tenant_id,
			dt.workspace_id,
			dt.target_account_id,
			dt.target_client_id,
			dt.platform_id,
			dt.kind,
			dt.payload,
			dt.status,
			dt.priority,
			dt.lane,
			dt.lane_weight,
			dt.source,
			dt.scheduler_group_key,
			dt.monitor_task_id,
			dt.supersedes_task_id,
			dt.control_flags,
			dt.interrupt_generation,
			dt.dedup_key,
			dt.active_attempt_id,
			dt.lease_expires_at,
			dt.attempts,
			dt.result,
			dt.error,
			dt.started_at,
			dt.interrupted_at,
			dt.interrupt_reason,
			dt.enqueued_at,
			dt.created_at,
			dt.updated_at
		FROM desktop_tasks AS dt
		WHERE dt.tenant_id = $1
		  AND dt.workspace_id = $2
		  AND dt.target_client_id = $3
		  AND dt.kind = 'publish'
		  AND dt.status = 'queued'
		  AND dt.attempts < 3
		  AND NOT EXISTS (
			SELECT 1
			FROM desktop_publish_jobs AS j
			WHERE j.desktop_id = dt.job_id
			  AND j.status <> 'queued'
		  )
		ORDER BY dt.priority DESC,
			COALESCE(dt.enqueued_at, dt.created_at) ASC,
			dt.created_at ASC,
			dt.desktop_id ASC
		LIMIT $4
	`, client.TenantID, client.WorkspaceID, client.ID, dispatchReplayLimit)
	if err != nil {
		if h.app.Logger != nil {
			h.app.Logger.Warn("desktop dispatch replay query failed",
				zap.String("client_id", client.ID.String()),
				zap.Error(err),
			)
		}
		return
	}
	defer rows.Close()

	replayed := 0
	for rows.Next() {
		task, scanErr := scanDispatchReplayDesktopTask(rows)
		if scanErr != nil {
			if h.app.Logger != nil {
				h.app.Logger.Warn("desktop dispatch replay scan failed",
					zap.String("client_id", client.ID.String()),
					zap.Error(scanErr),
				)
			}
			return
		}
		if err := writeDispatchFrame(conn, tenantapp.DesktopDispatchEventFromTask(task, "task_available")); err != nil {
			if h.app.Logger != nil {
				h.app.Logger.Warn("desktop dispatch replay write failed",
					zap.String("client_id", client.ID.String()),
					zap.String("task_id", task.DesktopID.String()),
					zap.Error(err),
				)
			}
			return
		}
		replayed++
	}
	if err := rows.Err(); err != nil && h.app.Logger != nil {
		h.app.Logger.Warn("desktop dispatch replay rows failed",
			zap.String("client_id", client.ID.String()),
			zap.Error(err),
		)
	}
	if replayed > 0 && h.app.Logger != nil {
		h.app.Logger.Info("desktop dispatch replayed queued publish tasks",
			zap.String("client_id", client.ID.String()),
			zap.Int("count", replayed),
		)
	}
}

type dispatchReplayDesktopTaskScanner interface {
	Scan(dest ...any) error
}

func scanDispatchReplayDesktopTask(row dispatchReplayDesktopTaskScanner) (*repository.DesktopTask, error) {
	var task repository.DesktopTask
	err := row.Scan(
		&task.DesktopID,
		&task.JobID,
		&task.TenantID,
		&task.WorkspaceID,
		&task.TargetAccountID,
		&task.TargetClientID,
		&task.Platform,
		&task.Kind,
		&task.Payload,
		&task.Status,
		&task.Priority,
		&task.Lane,
		&task.LaneWeight,
		&task.Source,
		&task.SchedulerGroupKey,
		&task.MonitorTaskID,
		&task.SupersedesTaskID,
		&task.ControlFlags,
		&task.InterruptGeneration,
		&task.DedupKey,
		&task.ActiveAttemptID,
		&task.LeaseExpiresAt,
		&task.Attempts,
		&task.Result,
		&task.Error,
		&task.StartedAt,
		&task.InterruptedAt,
		&task.InterruptReason,
		&task.EnqueuedAt,
		&task.CreatedAt,
		&task.UpdatedAt,
	)
	if err != nil {
		return nil, err
	}
	return &task, nil
}