From 7605890a1490a6be4f803d0bf0ddc6603dcd6a79 Mon Sep 17 00:00:00 2001 From: liangxu Date: Tue, 28 Apr 2026 09:14:24 +0800 Subject: [PATCH] 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) --- server/cmd/scheduler/main.go | 2 +- server/cmd/tenant-api/main.go | 2 +- .../tenant/app/desktop_task_service.go | 313 +++++++++++++++++- .../app/monitoring_daily_task_worker.go | 41 ++- .../app/monitoring_phase2_desktop_tasks.go | 197 +++++++++-- .../monitoring_phase2_desktop_tasks_test.go | 78 +++++ .../internal/tenant/app/monitoring_service.go | 9 + .../transport/desktop_client_handler.go | 2 +- .../tenant/transport/desktop_task_handler.go | 2 +- .../tenant/transport/monitoring_handler.go | 2 +- 10 files changed, 600 insertions(+), 48 deletions(-) create mode 100644 server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go diff --git a/server/cmd/scheduler/main.go b/server/cmd/scheduler/main.go index 06b80b0..e4aa002 100644 --- a/server/cmd/scheduler/main.go +++ b/server/cmd/scheduler/main.go @@ -76,7 +76,7 @@ func main() { app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger, - ) + ).WithRedis(app.Redis) ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/server/cmd/tenant-api/main.go b/server/cmd/tenant-api/main.go index 454ee0c..76902d9 100644 --- a/server/cmd/tenant-api/main.go +++ b/server/cmd/tenant-api/main.go @@ -32,7 +32,7 @@ func main() { app.DesktopTaskStreams.Run(workerCtx) app.DesktopDispatch.Run(workerCtx) - monitoringService := tenantapp.NewMonitoringService(app.DB, app.MonitoringDB, app.RabbitMQ, app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger) + monitoringService := tenantapp.NewMonitoringService(app.DB, app.MonitoringDB, app.RabbitMQ, app.Config.MonitoringDispatch, app.Config.BrandLibrary, app.Logger).WithRedis(app.Redis) monitoringCallbackService := tenantapp.NewMonitoringCallbackService(app.DB, app.MonitoringDB, app.RabbitMQ, app.LLM, app.Logger) tenantapp.NewMonitoringCollectOutboxWorker(monitoringService, app.Logger).Start(workerCtx) tenantapp.NewMonitoringResultIngestWorker(monitoringCallbackService, app.RabbitMQ, app.Logger, app.Config.MonitoringWorkers).Start(workerCtx) diff --git a/server/internal/tenant/app/desktop_task_service.go b/server/internal/tenant/app/desktop_task_service.go index c5e50ea..18681f4 100644 --- a/server/internal/tenant/app/desktop_task_service.go +++ b/server/internal/tenant/app/desktop_task_service.go @@ -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 ×tamp +} + 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 { diff --git a/server/internal/tenant/app/monitoring_daily_task_worker.go b/server/internal/tenant/app/monitoring_daily_task_worker.go index 9515b9a..bcf638d 100644 --- a/server/internal/tenant/app/monitoring_daily_task_worker.go +++ b/server/internal/tenant/app/monitoring_daily_task_worker.go @@ -374,7 +374,7 @@ func (s *MonitoringService) generateDailyMonitoringTasksForBrand( if materializeLimit <= 0 && dispatchLimit <= 0 { return 0, 0, 0, 0, nil } - if executionOwner != "desktop_tasks" || plan.PrimaryClientID == nil { + if executionOwner != "desktop_tasks" { return 0, 0, 0, 0, nil } @@ -435,7 +435,6 @@ func (s *MonitoringService) generateDailyMonitoringTasksForBrand( businessDay, now.Add(defaultMonitoringDailyLookahead), now.Add(-defaultMonitoringDailyDesktopRetryCooldown), - *plan.PrimaryClientID, dispatchLimit, ) if err != nil { @@ -577,7 +576,9 @@ func (s *MonitoringService) loadDailyMonitoringPrimaryClientPlans( ORDER BY c.tenant_id ASC, c.workspace_id ASC `, freeBrandLimit, paidBrandLimit, questionLimit, supportedPlatforms) if err != nil { - return nil, response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plans") + appErr := response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plans") + appErr.Cause = err + return nil, appErr } defer rows.Close() @@ -661,7 +662,9 @@ func (s *MonitoringService) loadDailyMonitoringPlanPlatformJSON( ORDER BY p.tenant_id ASC, p.workspace_id ASC, p.client_id ASC `, tenantIDs, workspaceIDs, clientIDs, supportedPlatforms) if err != nil { - return nil, response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plan platforms") + appErr := response.ErrInternal(50041, "query_failed", "failed to load daily monitoring plan platforms") + appErr.Cause = err + return nil, appErr } defer rows.Close() @@ -865,15 +868,15 @@ func (s *MonitoringService) ensureDailyMonitoringCollectTasks( 'pending', 100, 'normal', - $5, + NULL, input.dispatch_after, - $6, - $7 + $5, + $6 FROM unnest( - $8::bigint[], - $9::bytea[], - $10::text[], - $11::timestamptz[] + $7::bigint[], + $8::bytea[], + $9::text[], + $10::timestamptz[] ) AS input(question_id, question_hash, ai_platform_id, dispatch_after) ON CONFLICT ( tenant_id, brand_id, question_id, ai_platform_id, collector_type, run_mode, business_date @@ -884,7 +887,6 @@ func (s *MonitoringService) ensureDailyMonitoringCollectTasks( brand.BrandID, monitoringCollectorType, businessDate, - nullableUUID(plan.PrimaryClientID), shardKey, executionOwner, questionIDs, @@ -906,7 +908,6 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( businessDate time.Time, cutoff time.Time, retryBefore time.Time, - targetClientID uuid.UUID, limit int, ) ([]monitorDesktopTaskSpec, error) { if q == nil || limit <= 0 { @@ -925,11 +926,10 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( AND t.trigger_source = 'automatic' AND t.run_mode = 'desktop_standard' AND COALESCE(t.execution_owner, 'legacy') = 'desktop_tasks' - AND t.target_client_id = $5 - AND (t.dispatch_after IS NULL OR t.dispatch_after <= $6) - AND (t.last_dispatched_at IS NULL OR t.last_dispatched_at <= $7) + AND (t.dispatch_after IS NULL OR t.dispatch_after <= $5) + AND (t.last_dispatched_at IS NULL OR t.last_dispatched_at <= $6) ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC - LIMIT $8 + LIMIT $7 FOR UPDATE SKIP LOCKED ), claimed_tasks AS ( @@ -964,9 +964,9 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( ORDER BY s.projected_at DESC, s.id DESC LIMIT 1 ), '') AS question_text_snapshot - FROM claimed_tasks t - ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC - `, tenantID, brandID, monitoringCollectorType, monitoringBusinessDateText(businessDate), targetClientID, cutoff.UTC(), retryBefore.UTC(), limit) + FROM claimed_tasks t + ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC + `, tenantID, brandID, monitoringCollectorType, monitoringBusinessDateText(businessDate), cutoff.UTC(), retryBefore.UTC(), limit) if err != nil { return nil, response.ErrInternal(50041, "query_failed", "failed to load due daily monitor desktop task specs") } @@ -996,7 +996,6 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs( return nil, response.ErrInternal(50041, "scan_failed", "failed to parse due daily monitor desktop task specs") } spec.WorkspaceID = workspaceID - spec.TargetClientID = targetClientID spec.BusinessDate = businessDay.Format("2006-01-02") spec.QuestionHash = encodeQuestionHash(questionHash) spec.SchedulerGroupKey = monitoringSchedulerGroupKey(spec.QuestionHash, spec.QuestionID, spec.QuestionText) diff --git a/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go b/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go index 9070470..4ea9759 100644 --- a/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go +++ b/server/internal/tenant/app/monitoring_phase2_desktop_tasks.go @@ -26,6 +26,7 @@ type monitorDesktopTaskSpec struct { MonitorTaskID int64 TenantID int64 WorkspaceID int64 + TargetAccountID uuid.UUID TargetClientID uuid.UUID PlatformID string BusinessDate string @@ -40,6 +41,19 @@ type monitorDesktopTaskSpec struct { InterruptGeneration int } +type monitorDesktopAccountCandidate struct { + PlatformID string + AccountID uuid.UUID + DBClientID *uuid.UUID + HealthRank int + RecordOrder int +} + +type monitorDesktopTaskTarget struct { + AccountID uuid.UUID + ClientID uuid.UUID +} + type monitorDesktopInterruptTarget struct { TaskID string MonitorTaskID int64 @@ -181,7 +195,7 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks( return nil, nil, nil } - accountMap, err := s.loadMonitorDesktopAccountMap(ctx, specs[0].TenantID, specs[0].WorkspaceID, specs[0].TargetClientID) + targets, err := s.loadMonitorDesktopTaskTargets(ctx, specs[0].TenantID, specs[0].WorkspaceID, specs) if err != nil { return nil, nil, err } @@ -196,7 +210,15 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks( publishableTasks := make([]*repository.DesktopTask, 0, len(specs)) fallbackLegacyTaskIDs := make([]int64, 0) for _, spec := range specs { - task, shouldPublish, fallbackLegacy, upsertErr := s.upsertMonitorDesktopTask(ctx, tx, repo, spec, accountMap) + target, ok := targets[normalizeMonitoringPlatformID(spec.PlatformID)] + if !ok || target.AccountID == uuid.Nil || target.ClientID == uuid.Nil { + fallbackLegacyTaskIDs = append(fallbackLegacyTaskIDs, spec.MonitorTaskID) + continue + } + spec.TargetAccountID = target.AccountID + spec.TargetClientID = target.ClientID + + task, shouldPublish, fallbackLegacy, upsertErr := s.upsertMonitorDesktopTask(ctx, tx, repo, spec) if upsertErr != nil { return nil, nil, upsertErr } @@ -215,20 +237,32 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks( return publishableTasks, fallbackLegacyTaskIDs, nil } -func (s *MonitoringService) loadMonitorDesktopAccountMap( +func (s *MonitoringService) loadMonitorDesktopTaskTargets( ctx context.Context, tenantID, workspaceID int64, - targetClientID uuid.UUID, -) (map[string]uuid.UUID, error) { + specs []monitorDesktopTaskSpec, +) (map[string]monitorDesktopTaskTarget, error) { + platformIDs := uniqueMonitorDesktopTaskPlatformIDs(specs) + if len(platformIDs) == 0 { + return nil, nil + } + rows, err := s.businessPool.Query(ctx, ` - SELECT DISTINCT ON (platform_id) + SELECT platform_id, - desktop_id + desktop_id, + COALESCE(client_id::text, '') AS client_id, + CASE health + WHEN 'live' THEN 0 + WHEN 'risk' THEN 1 + WHEN 'captcha' THEN 2 + ELSE 3 + END AS health_rank FROM platform_accounts WHERE tenant_id = $1 AND workspace_id = $2 - AND client_id = $3 AND deleted_at IS NULL + AND platform_id = ANY($3::text[]) ORDER BY platform_id ASC, CASE health WHEN 'live' THEN 0 @@ -239,35 +273,158 @@ func (s *MonitoringService) loadMonitorDesktopAccountMap( verified_at DESC NULLS LAST, updated_at DESC, created_at DESC - `, tenantID, workspaceID, targetClientID) + `, tenantID, workspaceID, platformIDs) if err != nil { return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to resolve monitor desktop accounts") } defer rows.Close() - result := make(map[string]uuid.UUID) + candidates := make([]monitorDesktopAccountCandidate, 0) for rows.Next() { var ( - platformID string - desktopID uuid.UUID + item monitorDesktopAccountCandidate + clientIDText string ) - if scanErr := rows.Scan(&platformID, &desktopID); scanErr != nil { + if scanErr := rows.Scan(&item.PlatformID, &item.AccountID, &clientIDText, &item.HealthRank); scanErr != nil { return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to parse monitor desktop accounts") } - result[platformID] = desktopID + if parsedClientID, parseErr := uuid.Parse(strings.TrimSpace(clientIDText)); parseErr == nil { + item.DBClientID = &parsedClientID + } + item.PlatformID = normalizeMonitoringPlatformID(item.PlatformID) + item.RecordOrder = len(candidates) + candidates = append(candidates, item) } if err := rows.Err(); err != nil { return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to iterate monitor desktop accounts") } + + accountIDs := monitorDesktopAccountCandidateAccountIDs(candidates) + accountPresence := loadDesktopAccountPresence(ctx, s.redis, accountIDs) + clientIDs := monitorDesktopAccountCandidateClientIDs(candidates, accountPresence) + onlineClients, err := s.loadOnlineMonitorDesktopClients(ctx, workspaceID, clientIDs) + if err != nil { + return nil, err + } + return selectMonitorDesktopTaskTargets(candidates, accountPresence, onlineClients), nil +} + +func uniqueMonitorDesktopTaskPlatformIDs(specs []monitorDesktopTaskSpec) []string { + if len(specs) == 0 { + return nil + } + seen := make(map[string]struct{}, len(specs)) + result := make([]string, 0, len(specs)) + for _, spec := range specs { + platformID := normalizeMonitoringPlatformID(spec.PlatformID) + if platformID == "" { + continue + } + if _, ok := seen[platformID]; ok { + continue + } + seen[platformID] = struct{}{} + result = append(result, platformID) + } + return result +} + +func monitorDesktopAccountCandidateAccountIDs(candidates []monitorDesktopAccountCandidate) []uuid.UUID { + result := make([]uuid.UUID, 0, len(candidates)) + for _, candidate := range candidates { + if candidate.AccountID != uuid.Nil { + result = append(result, candidate.AccountID) + } + } + return uniqueUUIDs(result) +} + +func monitorDesktopAccountCandidateClientIDs(candidates []monitorDesktopAccountCandidate, accountPresence map[uuid.UUID]uuid.UUID) []uuid.UUID { + result := make([]uuid.UUID, 0, len(candidates)*2) + for _, candidate := range candidates { + if candidate.DBClientID != nil && *candidate.DBClientID != uuid.Nil { + result = append(result, *candidate.DBClientID) + } + if accountPresence != nil { + if clientID, ok := accountPresence[candidate.AccountID]; ok && clientID != uuid.Nil { + result = append(result, clientID) + } + } + } + return uniqueUUIDs(result) +} + +func (s *MonitoringService) loadOnlineMonitorDesktopClients(ctx context.Context, workspaceID int64, clientIDs []uuid.UUID) (map[uuid.UUID]bool, error) { + result := make(map[uuid.UUID]bool) + clientIDs = uniqueUUIDs(clientIDs) + if len(clientIDs) == 0 { + return result, nil + } + + for clientID, online := range loadDesktopClientPresence(ctx, s.redis, clientIDs) { + if online { + result[clientID] = true + } + } + + rows, err := s.businessPool.Query(ctx, ` + SELECT id + FROM desktop_clients + WHERE workspace_id = $1 + AND id = ANY($2::uuid[]) + AND revoked_at IS NULL + AND last_seen_at IS NOT NULL + AND last_seen_at >= NOW() - $3::interval + `, workspaceID, clientIDs, formatPgInterval(desktopClientPresenceTTL)) + if err != nil { + return nil, response.ErrInternal(50117, "desktop_client_lookup_failed", "failed to resolve online monitor desktop clients") + } + defer rows.Close() + + for rows.Next() { + var clientID uuid.UUID + if scanErr := rows.Scan(&clientID); scanErr != nil { + return nil, response.ErrInternal(50117, "desktop_client_lookup_failed", "failed to parse online monitor desktop clients") + } + result[clientID] = true + } + if err := rows.Err(); err != nil { + return nil, response.ErrInternal(50117, "desktop_client_lookup_failed", "failed to iterate online monitor desktop clients") + } return result, nil } +func selectMonitorDesktopTaskTargets( + candidates []monitorDesktopAccountCandidate, + accountPresence map[uuid.UUID]uuid.UUID, + onlineClients map[uuid.UUID]bool, +) map[string]monitorDesktopTaskTarget { + result := make(map[string]monitorDesktopTaskTarget) + for _, candidate := range candidates { + platformID := normalizeMonitoringPlatformID(candidate.PlatformID) + if platformID == "" || candidate.AccountID == uuid.Nil { + continue + } + if _, exists := result[platformID]; exists { + continue + } + + if clientID, ok := accountPresence[candidate.AccountID]; ok && clientID != uuid.Nil && onlineClients[clientID] { + result[platformID] = monitorDesktopTaskTarget{AccountID: candidate.AccountID, ClientID: clientID} + continue + } + if candidate.DBClientID != nil && *candidate.DBClientID != uuid.Nil && onlineClients[*candidate.DBClientID] { + result[platformID] = monitorDesktopTaskTarget{AccountID: candidate.AccountID, ClientID: *candidate.DBClientID} + } + } + return result +} + func (s *MonitoringService) upsertMonitorDesktopTask( ctx context.Context, tx pgx.Tx, repo repository.DesktopTaskRepository, spec monitorDesktopTaskSpec, - accountMap map[string]uuid.UUID, ) (*repository.DesktopTask, bool, bool, error) { var ( existingDesktopID uuid.UUID @@ -297,9 +454,9 @@ func (s *MonitoringService) upsertMonitorDesktopTask( if existingStatus != "queued" && !expiredInProgress { return task, false, false, nil } - targetAccountID := task.TargetAccountID - if override, ok := accountMap[spec.PlatformID]; ok { - targetAccountID = override + targetAccountID := spec.TargetAccountID + if targetAccountID == uuid.Nil { + targetAccountID = task.TargetAccountID } payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec) if payloadErr != nil { @@ -372,8 +529,8 @@ func (s *MonitoringService) upsertMonitorDesktopTask( return nil, false, false, response.ErrInternal(50118, "desktop_task_lookup_failed", "failed to inspect active phase2 monitor desktop task") } - targetAccountID, ok := accountMap[spec.PlatformID] - if !ok { + targetAccountID := spec.TargetAccountID + if targetAccountID == uuid.Nil || spec.TargetClientID == uuid.Nil { return nil, false, true, nil } payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec) diff --git a/server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go b/server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go new file mode 100644 index 0000000..8f518c7 --- /dev/null +++ b/server/internal/tenant/app/monitoring_phase2_desktop_tasks_test.go @@ -0,0 +1,78 @@ +package app + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" +) + +func TestSelectMonitorDesktopTaskTargetsPrefersPresentAccountClient(t *testing.T) { + accountID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000101") + staleClientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000201") + liveClientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000202") + + targets := selectMonitorDesktopTaskTargets( + []monitorDesktopAccountCandidate{ + { + PlatformID: "qwen", + AccountID: accountID, + DBClientID: &staleClientID, + }, + }, + map[uuid.UUID]uuid.UUID{ + accountID: liveClientID, + }, + map[uuid.UUID]bool{ + liveClientID: true, + }, + ) + + assert.Equal(t, monitorDesktopTaskTarget{ + AccountID: accountID, + ClientID: liveClientID, + }, targets["qwen"]) +} + +func TestSelectMonitorDesktopTaskTargetsFallsBackToOnlineDBClient(t *testing.T) { + accountID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000102") + clientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000203") + + targets := selectMonitorDesktopTaskTargets( + []monitorDesktopAccountCandidate{ + { + PlatformID: "deepseek", + AccountID: accountID, + DBClientID: &clientID, + }, + }, + nil, + map[uuid.UUID]bool{ + clientID: true, + }, + ) + + assert.Equal(t, monitorDesktopTaskTarget{ + AccountID: accountID, + ClientID: clientID, + }, targets["deepseek"]) +} + +func TestSelectMonitorDesktopTaskTargetsSkipsOfflineClients(t *testing.T) { + accountID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000103") + clientID := mustParseDailyMonitoringUUID(t, "00000000-0000-0000-0000-000000000204") + + targets := selectMonitorDesktopTaskTargets( + []monitorDesktopAccountCandidate{ + { + PlatformID: "kimi", + AccountID: accountID, + DBClientID: &clientID, + }, + }, + nil, + map[uuid.UUID]bool{}, + ) + + assert.Empty(t, targets) +} diff --git a/server/internal/tenant/app/monitoring_service.go b/server/internal/tenant/app/monitoring_service.go index f8e1621..fca725c 100644 --- a/server/internal/tenant/app/monitoring_service.go +++ b/server/internal/tenant/app/monitoring_service.go @@ -14,6 +14,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "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" @@ -40,6 +41,7 @@ type MonitoringService struct { businessPool *pgxpool.Pool monitoringPool *pgxpool.Pool rabbitMQ *rabbitmq.Client + redis *goredis.Client logger *zap.Logger desktopTasksRolloutPercentage int brandLibraryConfig config.BrandLibraryConfig @@ -62,6 +64,13 @@ func NewMonitoringService( } } +func (s *MonitoringService) WithRedis(redis *goredis.Client) *MonitoringService { + if s != nil { + s.redis = redis + } + return s +} + type MonitoringDashboardCompositeResponse struct { Overview MonitoringOverview `json:"overview"` Runtime MonitoringDashboardRuntime `json:"runtime"` diff --git a/server/internal/tenant/transport/desktop_client_handler.go b/server/internal/tenant/transport/desktop_client_handler.go index 458382e..476705c 100644 --- a/server/internal/tenant/transport/desktop_client_handler.go +++ b/server/internal/tenant/transport/desktop_client_handler.go @@ -20,7 +20,7 @@ func NewDesktopClientHandler(a *bootstrap.App) *DesktopClientHandler { a.MonitoringDB, a.RabbitMQ, a.Logger, - ).WithCache(a.Cache) + ).WithCache(a.Cache).WithRedis(a.Redis) return &DesktopClientHandler{ svc: app.NewDesktopClientService( repository.NewDesktopClientRepository(a.DB), diff --git a/server/internal/tenant/transport/desktop_task_handler.go b/server/internal/tenant/transport/desktop_task_handler.go index 6bc63a3..74e6596 100644 --- a/server/internal/tenant/transport/desktop_task_handler.go +++ b/server/internal/tenant/transport/desktop_task_handler.go @@ -27,7 +27,7 @@ func NewDesktopTaskHandler(a *bootstrap.App) *DesktopTaskHandler { a.MonitoringDB, a.RabbitMQ, a.Logger, - ).WithCache(a.Cache), + ).WithCache(a.Cache).WithRedis(a.Redis), publishSvc: app.NewPublishJobService( a.DB, a.RabbitMQ, diff --git a/server/internal/tenant/transport/monitoring_handler.go b/server/internal/tenant/transport/monitoring_handler.go index 3319569..1d43381 100644 --- a/server/internal/tenant/transport/monitoring_handler.go +++ b/server/internal/tenant/transport/monitoring_handler.go @@ -26,7 +26,7 @@ type monitoringCollectNowRequest struct { func NewMonitoringHandler(a *bootstrap.App) *MonitoringHandler { return &MonitoringHandler{ - svc: app.NewMonitoringService(a.DB, a.MonitoringDB, a.RabbitMQ, a.Config.MonitoringDispatch, a.Config.BrandLibrary, a.Logger), + svc: app.NewMonitoringService(a.DB, a.MonitoringDB, a.RabbitMQ, a.Config.MonitoringDispatch, a.Config.BrandLibrary, a.Logger).WithRedis(a.Redis), } }