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>
This commit is contained in:
2026-04-28 09:14:24 +08:00
parent 93cf734673
commit 7605890a14
10 changed files with 600 additions and 48 deletions
+1 -1
View File
@@ -76,7 +76,7 @@ func main() {
app.Config.MonitoringDispatch, app.Config.MonitoringDispatch,
app.Config.BrandLibrary, app.Config.BrandLibrary,
app.Logger, app.Logger,
) ).WithRedis(app.Redis)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
+1 -1
View File
@@ -32,7 +32,7 @@ func main() {
app.DesktopTaskStreams.Run(workerCtx) app.DesktopTaskStreams.Run(workerCtx)
app.DesktopDispatch.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) monitoringCallbackService := tenantapp.NewMonitoringCallbackService(app.DB, app.MonitoringDB, app.RabbitMQ, app.LLM, app.Logger)
tenantapp.NewMonitoringCollectOutboxWorker(monitoringService, app.Logger).Start(workerCtx) tenantapp.NewMonitoringCollectOutboxWorker(monitoringService, app.Logger).Start(workerCtx)
tenantapp.NewMonitoringResultIngestWorker(monitoringCallbackService, app.RabbitMQ, app.Logger, app.Config.MonitoringWorkers).Start(workerCtx) tenantapp.NewMonitoringResultIngestWorker(monitoringCallbackService, app.RabbitMQ, app.Logger, app.Config.MonitoringWorkers).Start(workerCtx)
@@ -11,6 +11,7 @@ import (
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/pgx/v5/pgxpool"
goredis "github.com/redis/go-redis/v9"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/geo-platform/tenant-api/internal/shared/auth" "github.com/geo-platform/tenant-api/internal/shared/auth"
@@ -26,6 +27,7 @@ type DesktopTaskService struct {
monitoringPool *pgxpool.Pool monitoringPool *pgxpool.Pool
repo repository.DesktopTaskRepository repo repository.DesktopTaskRepository
cache sharedcache.Cache cache sharedcache.Cache
redis *goredis.Client
messaging *rabbitmq.Client messaging *rabbitmq.Client
logger *zap.Logger logger *zap.Logger
} }
@@ -45,6 +47,13 @@ func (s *DesktopTaskService) WithCache(c sharedcache.Cache) *DesktopTaskService
return s return s
} }
func (s *DesktopTaskService) WithRedis(redis *goredis.Client) *DesktopTaskService {
if s != nil {
s.redis = redis
}
return s
}
type DesktopTaskView struct { type DesktopTaskView struct {
ID string `json:"id"` ID string `json:"id"`
JobID string `json:"job_id"` JobID string `json:"job_id"`
@@ -163,7 +172,9 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt
var task *repository.DesktopTask var task *repository.DesktopTask
switch { switch {
case taskID != nil: 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: default:
task, err = s.repo.LeaseNextQueued(ctx, client.ID, req.Kind, leaseParams) 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 }, 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) { func (s *DesktopTaskService) Extend(ctx context.Context, client *repository.DesktopClient, desktopID uuid.UUID, req ExtendDesktopTaskRequest) (*DesktopTaskView, error) {
if client == nil { if client == nil {
return nil, response.ErrUnauthorized(40108, "desktop_client_missing", "desktop client context is required") 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") 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") 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") 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 { func buildDesktopTaskView(task *repository.DesktopTask) DesktopTaskView {
var activeAttemptID *string var activeAttemptID *string
if task.ActiveAttemptID != nil { if task.ActiveAttemptID != nil {
@@ -374,7 +374,7 @@ func (s *MonitoringService) generateDailyMonitoringTasksForBrand(
if materializeLimit <= 0 && dispatchLimit <= 0 { if materializeLimit <= 0 && dispatchLimit <= 0 {
return 0, 0, 0, 0, nil return 0, 0, 0, 0, nil
} }
if executionOwner != "desktop_tasks" || plan.PrimaryClientID == nil { if executionOwner != "desktop_tasks" {
return 0, 0, 0, 0, nil return 0, 0, 0, 0, nil
} }
@@ -435,7 +435,6 @@ func (s *MonitoringService) generateDailyMonitoringTasksForBrand(
businessDay, businessDay,
now.Add(defaultMonitoringDailyLookahead), now.Add(defaultMonitoringDailyLookahead),
now.Add(-defaultMonitoringDailyDesktopRetryCooldown), now.Add(-defaultMonitoringDailyDesktopRetryCooldown),
*plan.PrimaryClientID,
dispatchLimit, dispatchLimit,
) )
if err != nil { if err != nil {
@@ -577,7 +576,9 @@ func (s *MonitoringService) loadDailyMonitoringPrimaryClientPlans(
ORDER BY c.tenant_id ASC, c.workspace_id ASC ORDER BY c.tenant_id ASC, c.workspace_id ASC
`, freeBrandLimit, paidBrandLimit, questionLimit, supportedPlatforms) `, freeBrandLimit, paidBrandLimit, questionLimit, supportedPlatforms)
if err != nil { 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() 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 ORDER BY p.tenant_id ASC, p.workspace_id ASC, p.client_id ASC
`, tenantIDs, workspaceIDs, clientIDs, supportedPlatforms) `, tenantIDs, workspaceIDs, clientIDs, supportedPlatforms)
if err != nil { 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() defer rows.Close()
@@ -865,15 +868,15 @@ func (s *MonitoringService) ensureDailyMonitoringCollectTasks(
'pending', 'pending',
100, 100,
'normal', 'normal',
$5, NULL,
input.dispatch_after, input.dispatch_after,
$6, $5,
$7 $6
FROM unnest( FROM unnest(
$8::bigint[], $7::bigint[],
$9::bytea[], $8::bytea[],
$10::text[], $9::text[],
$11::timestamptz[] $10::timestamptz[]
) AS input(question_id, question_hash, ai_platform_id, dispatch_after) ) AS input(question_id, question_hash, ai_platform_id, dispatch_after)
ON CONFLICT ( ON CONFLICT (
tenant_id, brand_id, question_id, ai_platform_id, collector_type, run_mode, business_date 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, brand.BrandID,
monitoringCollectorType, monitoringCollectorType,
businessDate, businessDate,
nullableUUID(plan.PrimaryClientID),
shardKey, shardKey,
executionOwner, executionOwner,
questionIDs, questionIDs,
@@ -906,7 +908,6 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs(
businessDate time.Time, businessDate time.Time,
cutoff time.Time, cutoff time.Time,
retryBefore time.Time, retryBefore time.Time,
targetClientID uuid.UUID,
limit int, limit int,
) ([]monitorDesktopTaskSpec, error) { ) ([]monitorDesktopTaskSpec, error) {
if q == nil || limit <= 0 { if q == nil || limit <= 0 {
@@ -925,11 +926,10 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs(
AND t.trigger_source = 'automatic' AND t.trigger_source = 'automatic'
AND t.run_mode = 'desktop_standard' AND t.run_mode = 'desktop_standard'
AND COALESCE(t.execution_owner, 'legacy') = 'desktop_tasks' AND COALESCE(t.execution_owner, 'legacy') = 'desktop_tasks'
AND t.target_client_id = $5 AND (t.dispatch_after IS NULL OR t.dispatch_after <= $5)
AND (t.dispatch_after IS NULL OR t.dispatch_after <= $6) AND (t.last_dispatched_at IS NULL OR t.last_dispatched_at <= $6)
AND (t.last_dispatched_at IS NULL OR t.last_dispatched_at <= $7)
ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC
LIMIT $8 LIMIT $7
FOR UPDATE SKIP LOCKED FOR UPDATE SKIP LOCKED
), ),
claimed_tasks AS ( claimed_tasks AS (
@@ -964,9 +964,9 @@ func (s *MonitoringService) loadDueDailyMonitorDesktopTaskSpecs(
ORDER BY s.projected_at DESC, s.id DESC ORDER BY s.projected_at DESC, s.id DESC
LIMIT 1 LIMIT 1
), '') AS question_text_snapshot ), '') AS question_text_snapshot
FROM claimed_tasks t FROM claimed_tasks t
ORDER BY t.dispatch_after ASC NULLS FIRST, t.dispatch_priority DESC, t.id ASC 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) `, tenantID, brandID, monitoringCollectorType, monitoringBusinessDateText(businessDate), cutoff.UTC(), retryBefore.UTC(), limit)
if err != nil { if err != nil {
return nil, response.ErrInternal(50041, "query_failed", "failed to load due daily monitor desktop task specs") 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") return nil, response.ErrInternal(50041, "scan_failed", "failed to parse due daily monitor desktop task specs")
} }
spec.WorkspaceID = workspaceID spec.WorkspaceID = workspaceID
spec.TargetClientID = targetClientID
spec.BusinessDate = businessDay.Format("2006-01-02") spec.BusinessDate = businessDay.Format("2006-01-02")
spec.QuestionHash = encodeQuestionHash(questionHash) spec.QuestionHash = encodeQuestionHash(questionHash)
spec.SchedulerGroupKey = monitoringSchedulerGroupKey(spec.QuestionHash, spec.QuestionID, spec.QuestionText) spec.SchedulerGroupKey = monitoringSchedulerGroupKey(spec.QuestionHash, spec.QuestionID, spec.QuestionText)
@@ -26,6 +26,7 @@ type monitorDesktopTaskSpec struct {
MonitorTaskID int64 MonitorTaskID int64
TenantID int64 TenantID int64
WorkspaceID int64 WorkspaceID int64
TargetAccountID uuid.UUID
TargetClientID uuid.UUID TargetClientID uuid.UUID
PlatformID string PlatformID string
BusinessDate string BusinessDate string
@@ -40,6 +41,19 @@ type monitorDesktopTaskSpec struct {
InterruptGeneration int 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 { type monitorDesktopInterruptTarget struct {
TaskID string TaskID string
MonitorTaskID int64 MonitorTaskID int64
@@ -181,7 +195,7 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks(
return nil, nil, nil 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 { if err != nil {
return nil, nil, err return nil, nil, err
} }
@@ -196,7 +210,15 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks(
publishableTasks := make([]*repository.DesktopTask, 0, len(specs)) publishableTasks := make([]*repository.DesktopTask, 0, len(specs))
fallbackLegacyTaskIDs := make([]int64, 0) fallbackLegacyTaskIDs := make([]int64, 0)
for _, spec := range specs { 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 { if upsertErr != nil {
return nil, nil, upsertErr return nil, nil, upsertErr
} }
@@ -215,20 +237,32 @@ func (s *MonitoringService) dispatchMonitorDesktopTasks(
return publishableTasks, fallbackLegacyTaskIDs, nil return publishableTasks, fallbackLegacyTaskIDs, nil
} }
func (s *MonitoringService) loadMonitorDesktopAccountMap( func (s *MonitoringService) loadMonitorDesktopTaskTargets(
ctx context.Context, ctx context.Context,
tenantID, workspaceID int64, tenantID, workspaceID int64,
targetClientID uuid.UUID, specs []monitorDesktopTaskSpec,
) (map[string]uuid.UUID, error) { ) (map[string]monitorDesktopTaskTarget, error) {
platformIDs := uniqueMonitorDesktopTaskPlatformIDs(specs)
if len(platformIDs) == 0 {
return nil, nil
}
rows, err := s.businessPool.Query(ctx, ` rows, err := s.businessPool.Query(ctx, `
SELECT DISTINCT ON (platform_id) SELECT
platform_id, 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 FROM platform_accounts
WHERE tenant_id = $1 WHERE tenant_id = $1
AND workspace_id = $2 AND workspace_id = $2
AND client_id = $3
AND deleted_at IS NULL AND deleted_at IS NULL
AND platform_id = ANY($3::text[])
ORDER BY platform_id ASC, ORDER BY platform_id ASC,
CASE health CASE health
WHEN 'live' THEN 0 WHEN 'live' THEN 0
@@ -239,35 +273,158 @@ func (s *MonitoringService) loadMonitorDesktopAccountMap(
verified_at DESC NULLS LAST, verified_at DESC NULLS LAST,
updated_at DESC, updated_at DESC,
created_at DESC created_at DESC
`, tenantID, workspaceID, targetClientID) `, tenantID, workspaceID, platformIDs)
if err != nil { if err != nil {
return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to resolve monitor desktop accounts") return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to resolve monitor desktop accounts")
} }
defer rows.Close() defer rows.Close()
result := make(map[string]uuid.UUID) candidates := make([]monitorDesktopAccountCandidate, 0)
for rows.Next() { for rows.Next() {
var ( var (
platformID string item monitorDesktopAccountCandidate
desktopID uuid.UUID 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") 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 { if err := rows.Err(); err != nil {
return nil, response.ErrInternal(50117, "desktop_account_lookup_failed", "failed to iterate monitor desktop accounts") 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 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( func (s *MonitoringService) upsertMonitorDesktopTask(
ctx context.Context, ctx context.Context,
tx pgx.Tx, tx pgx.Tx,
repo repository.DesktopTaskRepository, repo repository.DesktopTaskRepository,
spec monitorDesktopTaskSpec, spec monitorDesktopTaskSpec,
accountMap map[string]uuid.UUID,
) (*repository.DesktopTask, bool, bool, error) { ) (*repository.DesktopTask, bool, bool, error) {
var ( var (
existingDesktopID uuid.UUID existingDesktopID uuid.UUID
@@ -297,9 +454,9 @@ func (s *MonitoringService) upsertMonitorDesktopTask(
if existingStatus != "queued" && !expiredInProgress { if existingStatus != "queued" && !expiredInProgress {
return task, false, false, nil return task, false, false, nil
} }
targetAccountID := task.TargetAccountID targetAccountID := spec.TargetAccountID
if override, ok := accountMap[spec.PlatformID]; ok { if targetAccountID == uuid.Nil {
targetAccountID = override targetAccountID = task.TargetAccountID
} }
payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec) payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec)
if payloadErr != nil { 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") return nil, false, false, response.ErrInternal(50118, "desktop_task_lookup_failed", "failed to inspect active phase2 monitor desktop task")
} }
targetAccountID, ok := accountMap[spec.PlatformID] targetAccountID := spec.TargetAccountID
if !ok { if targetAccountID == uuid.Nil || spec.TargetClientID == uuid.Nil {
return nil, false, true, nil return nil, false, true, nil
} }
payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec) payloadJSON, payloadErr := marshalMonitorDesktopTaskPayload(spec)
@@ -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)
}
@@ -14,6 +14,7 @@ import (
"github.com/google/uuid" "github.com/google/uuid"
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/pgx/v5/pgxpool"
goredis "github.com/redis/go-redis/v9"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/geo-platform/tenant-api/internal/shared/auth" "github.com/geo-platform/tenant-api/internal/shared/auth"
@@ -40,6 +41,7 @@ type MonitoringService struct {
businessPool *pgxpool.Pool businessPool *pgxpool.Pool
monitoringPool *pgxpool.Pool monitoringPool *pgxpool.Pool
rabbitMQ *rabbitmq.Client rabbitMQ *rabbitmq.Client
redis *goredis.Client
logger *zap.Logger logger *zap.Logger
desktopTasksRolloutPercentage int desktopTasksRolloutPercentage int
brandLibraryConfig config.BrandLibraryConfig 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 { type MonitoringDashboardCompositeResponse struct {
Overview MonitoringOverview `json:"overview"` Overview MonitoringOverview `json:"overview"`
Runtime MonitoringDashboardRuntime `json:"runtime"` Runtime MonitoringDashboardRuntime `json:"runtime"`
@@ -20,7 +20,7 @@ func NewDesktopClientHandler(a *bootstrap.App) *DesktopClientHandler {
a.MonitoringDB, a.MonitoringDB,
a.RabbitMQ, a.RabbitMQ,
a.Logger, a.Logger,
).WithCache(a.Cache) ).WithCache(a.Cache).WithRedis(a.Redis)
return &DesktopClientHandler{ return &DesktopClientHandler{
svc: app.NewDesktopClientService( svc: app.NewDesktopClientService(
repository.NewDesktopClientRepository(a.DB), repository.NewDesktopClientRepository(a.DB),
@@ -27,7 +27,7 @@ func NewDesktopTaskHandler(a *bootstrap.App) *DesktopTaskHandler {
a.MonitoringDB, a.MonitoringDB,
a.RabbitMQ, a.RabbitMQ,
a.Logger, a.Logger,
).WithCache(a.Cache), ).WithCache(a.Cache).WithRedis(a.Redis),
publishSvc: app.NewPublishJobService( publishSvc: app.NewPublishJobService(
a.DB, a.DB,
a.RabbitMQ, a.RabbitMQ,
@@ -26,7 +26,7 @@ type monitoringCollectNowRequest struct {
func NewMonitoringHandler(a *bootstrap.App) *MonitoringHandler { func NewMonitoringHandler(a *bootstrap.App) *MonitoringHandler {
return &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),
} }
} }