refactor(server/monitoring): route monitor desktop tasks by account presence

Daily monitoring tasks were pinned to a single primary client. If that client went offline the task wedged even when another desktop client of the same workspace was online and bound to the same account. Drop the primary-client constraint on the materialized collect rows, and instead pick a target per-platform from live account/client presence in Redis, falling back to the DB client_id when the desktop client is still flagged online. Desktop task lease for kind=monitor now matches by account ownership in addition to target_client_id, so any client owning the account can drain the queue. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-28 09:14:24 +08:00
parent 93cf734673
commit 7605890a14
10 changed files with 600 additions and 48 deletions
@@ -11,6 +11,7 @@ import (
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/jackc/pgx/v5/pgxpool"
+	goredis "github.com/redis/go-redis/v9"
 	"go.uber.org/zap"

 	"github.com/geo-platform/tenant-api/internal/shared/auth"
@@ -26,6 +27,7 @@ type DesktopTaskService struct {
 	monitoringPool *pgxpool.Pool
 	repo           repository.DesktopTaskRepository
 	cache          sharedcache.Cache
+	redis          *goredis.Client
 	messaging      *rabbitmq.Client
 	logger         *zap.Logger
 }
@@ -45,6 +47,13 @@ func (s *DesktopTaskService) WithCache(c sharedcache.Cache) *DesktopTaskService
 	return s
 }

+func (s *DesktopTaskService) WithRedis(redis *goredis.Client) *DesktopTaskService {
+	if s != nil {
+		s.redis = redis
+	}
+	return s
+}
+
 type DesktopTaskView struct {
 	ID              string          `json:"id"`
 	JobID           string          `json:"job_id"`
@@ -163,7 +172,9 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt
 	var task *repository.DesktopTask
 	switch {
 	case taskID != nil:
-		task, err = s.repo.LeaseQueuedByDesktopID(ctx, *taskID, client.ID, leaseParams)
+		task, err = s.leaseQueuedDesktopTaskByID(ctx, client, *taskID, leaseParams)
+	case req.Kind != nil && strings.TrimSpace(*req.Kind) == "monitor":
+		task, err = s.leaseNextQueuedMonitorTask(ctx, client, leaseParams)
 	default:
 		task, err = s.repo.LeaseNextQueued(ctx, client.ID, req.Kind, leaseParams)
 	}
@@ -219,6 +230,288 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt
 	}, nil
 }

+const desktopTaskRepositoryReturningColumns = `
+	t.desktop_id,
+	t.job_id,
+	t.tenant_id,
+	t.workspace_id,
+	t.target_account_id,
+	t.target_client_id,
+	t.platform_id,
+	t.kind,
+	t.payload,
+	t.status,
+	t.priority,
+	t.lane,
+	t.lane_weight,
+	t.source,
+	t.scheduler_group_key,
+	t.monitor_task_id,
+	t.supersedes_task_id,
+	t.control_flags,
+	t.interrupt_generation,
+	t.dedup_key,
+	t.active_attempt_id,
+	t.lease_expires_at,
+	t.attempts,
+	t.result,
+	t.error,
+	t.started_at,
+	t.interrupted_at,
+	t.interrupt_reason,
+	t.enqueued_at,
+	t.created_at,
+	t.updated_at`
+
+type desktopTaskRepositoryScanner interface {
+	Scan(dest ...any) error
+}
+
+func (s *DesktopTaskService) leaseNextQueuedMonitorTask(
+	ctx context.Context,
+	client *repository.DesktopClient,
+	params repository.DesktopTaskLeaseParams,
+) (*repository.DesktopTask, error) {
+	accountIDs, err := s.loadMonitorLeaseAccountIDs(ctx, client)
+	if err != nil {
+		return nil, err
+	}
+	hasAccountIDs := len(accountIDs) > 0
+
+	row := s.pool.QueryRow(ctx, `
+		WITH candidate AS (
+			SELECT dt.desktop_id
+			FROM desktop_tasks AS dt
+			WHERE dt.tenant_id = $1
+			  AND dt.workspace_id = $2
+			  AND dt.kind = 'monitor'
+			  AND dt.status = 'queued'
+			  AND (
+				dt.target_client_id = $3
+				OR ($4::boolean AND dt.target_account_id = ANY($5::uuid[]))
+			  )
+			ORDER BY dt.lane_weight DESC,
+				dt.priority DESC,
+				COALESCE(dt.enqueued_at, dt.created_at) ASC,
+				dt.created_at ASC,
+				dt.desktop_id ASC
+			LIMIT 1
+			FOR UPDATE SKIP LOCKED
+		)
+		UPDATE desktop_tasks AS t
+		SET target_client_id = $3,
+		    active_attempt_id = $6,
+		    lease_token_hash = $7,
+		    lease_expires_at = now() + interval '10 minutes',
+		    status = 'in_progress',
+		    attempts = t.attempts + 1,
+		    updated_at = now()
+		FROM candidate
+		WHERE t.desktop_id = candidate.desktop_id
+		RETURNING `+desktopTaskRepositoryReturningColumns,
+		client.TenantID,
+		client.WorkspaceID,
+		client.ID,
+		hasAccountIDs,
+		accountIDs,
+		params.AttemptID,
+		params.LeaseTokenHash,
+	)
+	return scanRepositoryDesktopTask(row)
+}
+
+func (s *DesktopTaskService) leaseQueuedDesktopTaskByID(
+	ctx context.Context,
+	client *repository.DesktopClient,
+	desktopID uuid.UUID,
+	params repository.DesktopTaskLeaseParams,
+) (*repository.DesktopTask, error) {
+	accountIDs, err := s.loadMonitorLeaseAccountIDs(ctx, client)
+	if err != nil {
+		return nil, err
+	}
+	hasAccountIDs := len(accountIDs) > 0
+
+	row := s.pool.QueryRow(ctx, `
+		WITH candidate AS (
+			SELECT dt.desktop_id
+			FROM desktop_tasks AS dt
+			WHERE dt.desktop_id = $1
+			  AND dt.workspace_id = $2
+			  AND dt.status = 'queued'
+			  AND (
+				(
+					dt.kind = 'monitor'
+					AND dt.tenant_id = $3
+					AND (
+						dt.target_client_id = $4
+						OR ($5::boolean AND dt.target_account_id = ANY($6::uuid[]))
+					)
+				)
+				OR (
+					dt.kind <> 'monitor'
+					AND dt.target_client_id = $4
+				)
+			  )
+			LIMIT 1
+			FOR UPDATE SKIP LOCKED
+		)
+		UPDATE desktop_tasks AS t
+		SET target_client_id = CASE WHEN t.kind = 'monitor' THEN $4 ELSE t.target_client_id END,
+		    active_attempt_id = $7,
+		    lease_token_hash = $8,
+		    lease_expires_at = now() + interval '10 minutes',
+		    status = 'in_progress',
+		    attempts = t.attempts + 1,
+		    updated_at = now()
+		FROM candidate
+		WHERE t.desktop_id = candidate.desktop_id
+		RETURNING `+desktopTaskRepositoryReturningColumns,
+		desktopID,
+		client.WorkspaceID,
+		client.TenantID,
+		client.ID,
+		hasAccountIDs,
+		accountIDs,
+		params.AttemptID,
+		params.LeaseTokenHash,
+	)
+	return scanRepositoryDesktopTask(row)
+}
+
+func (s *DesktopTaskService) loadMonitorLeaseAccountIDs(ctx context.Context, client *repository.DesktopClient) ([]uuid.UUID, error) {
+	if s == nil || client == nil || s.pool == nil {
+		return nil, nil
+	}
+
+	accountIDs := loadTrackedDesktopClientAccountIDs(ctx, s.redis, client.ID)
+	rows, err := s.pool.Query(ctx, `
+		SELECT desktop_id
+		FROM platform_accounts
+		WHERE tenant_id = $1
+		  AND workspace_id = $2
+		  AND client_id = $3
+		  AND deleted_at IS NULL
+		  AND platform_id = ANY($4::text[])
+	`, client.TenantID, client.WorkspaceID, client.ID, monitoringPlatformIDs(defaultMonitoringPlatforms))
+	if err != nil {
+		return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to resolve monitor lease accounts")
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var accountID uuid.UUID
+		if scanErr := rows.Scan(&accountID); scanErr != nil {
+			return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to parse monitor lease accounts")
+		}
+		accountIDs = append(accountIDs, accountID)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to iterate monitor lease accounts")
+	}
+	return uniqueUUIDs(accountIDs), nil
+}
+
+func scanRepositoryDesktopTask(row desktopTaskRepositoryScanner) (*repository.DesktopTask, error) {
+	var (
+		task              repository.DesktopTask
+		schedulerGroupKey pgtype.Text
+		monitorTaskID     pgtype.Int8
+		supersedesTaskID  pgtype.UUID
+		dedupKey          pgtype.Text
+		activeAttemptID   pgtype.UUID
+		leaseExpiresAt    pgtype.Timestamptz
+		startedAt         pgtype.Timestamptz
+		interruptedAt     pgtype.Timestamptz
+		interruptReason   pgtype.Text
+		enqueuedAt        pgtype.Timestamptz
+	)
+	if 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,
+		&schedulerGroupKey,
+		&monitorTaskID,
+		&supersedesTaskID,
+		&task.ControlFlags,
+		&task.InterruptGeneration,
+		&dedupKey,
+		&activeAttemptID,
+		&leaseExpiresAt,
+		&task.Attempts,
+		&task.Result,
+		&task.Error,
+		&startedAt,
+		&interruptedAt,
+		&interruptReason,
+		&enqueuedAt,
+		&task.CreatedAt,
+		&task.UpdatedAt,
+	); err != nil {
+		return nil, err
+	}
+
+	task.SchedulerGroupKey = desktopTaskNullableText(schedulerGroupKey)
+	task.MonitorTaskID = desktopTaskNullableInt64(monitorTaskID)
+	task.SupersedesTaskID = desktopTaskNullableUUID(supersedesTaskID)
+	task.DedupKey = desktopTaskNullableText(dedupKey)
+	task.ActiveAttemptID = desktopTaskNullableUUID(activeAttemptID)
+	task.LeaseExpiresAt = desktopTaskNullableTime(leaseExpiresAt)
+	task.StartedAt = desktopTaskNullableTime(startedAt)
+	task.InterruptedAt = desktopTaskNullableTime(interruptedAt)
+	task.InterruptReason = desktopTaskNullableText(interruptReason)
+	if enqueuedAt.Valid {
+		task.EnqueuedAt = enqueuedAt.Time
+	}
+	return &task, nil
+}
+
+func desktopTaskNullableText(value pgtype.Text) *string {
+	if !value.Valid {
+		return nil
+	}
+	text := value.String
+	return &text
+}
+
+func desktopTaskNullableInt64(value pgtype.Int8) *int64 {
+	if !value.Valid {
+		return nil
+	}
+	number := value.Int64
+	return &number
+}
+
+func desktopTaskNullableUUID(value pgtype.UUID) *uuid.UUID {
+	if !value.Valid {
+		return nil
+	}
+	resolved, err := uuid.FromBytes(value.Bytes[:])
+	if err != nil {
+		return nil
+	}
+	return &resolved
+}
+
+func desktopTaskNullableTime(value pgtype.Timestamptz) *time.Time {
+	if !value.Valid {
+		return nil
+	}
+	timestamp := value.Time
+	return &timestamp
+}
+
 func (s *DesktopTaskService) Extend(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req ExtendDesktopTaskRequest) (*DesktopTaskView, error) {
 	if client == nil {
 		return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required")
@@ -731,7 +1024,7 @@ func (s *DesktopTaskService) classifyLeaseError(ctx context.Context, client *rep
 		return response.ErrInternal(50095, "desktop_task_lookup_failed", "failed to inspect desktop task state")
 	}

-	if task.TargetClientID != client.ID {
+	if task.TargetClientID != client.ID && !s.clientCanLeaseMonitorTask(ctx, client, task) {
 		return response.ErrForbidden(40381, "desktop_task_not_target_client", "desktop task is assigned to another desktop client")
 	}

@@ -742,6 +1035,22 @@ func (s *DesktopTaskService) classifyLeaseError(ctx context.Context, client *rep
 	return response.ErrConflict(40991, "desktop_task_not_queued", "desktop task is not queued")
 }

+func (s *DesktopTaskService) clientCanLeaseMonitorTask(ctx context.Context, client *repository.DesktopClient, task *repository.DesktopTask) bool {
+	if s == nil || client == nil || task == nil || task.Kind != "monitor" || task.TargetAccountID == uuid.Nil {
+		return false
+	}
+	accountIDs, err := s.loadMonitorLeaseAccountIDs(ctx, client)
+	if err != nil {
+		return false
+	}
+	for _, accountID := range accountIDs {
+		if accountID == task.TargetAccountID {
+			return true
+		}
+	}
+	return false
+}
+
 func buildDesktopTaskView(task *repository.DesktopTask) DesktopTaskView {
 	var activeAttemptID *string
 	if task.ActiveAttemptID != nil {