diff --git a/server/internal/tenant/app/desktop_task_service.go b/server/internal/tenant/app/desktop_task_service.go index a676515..a88bbb0 100644 --- a/server/internal/tenant/app/desktop_task_service.go +++ b/server/internal/tenant/app/desktop_task_service.go @@ -28,6 +28,7 @@ import ( const ( desktopPublishMaxAttempts = 3 maxConcurrentMonitorPlatformsPerDesktopClient = 2 + monitorLeaseCandidateScanLimitSQL = "16" monitorSucceededTaskCooldownSQL = "2 seconds" monitorUnhealthyTaskCooldownSQL = "30 seconds" ) @@ -519,12 +520,6 @@ func (s *DesktopTaskService) leaseNextQueuedMonitorTask( 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 - tx, err := beginDesktopMonitorLeaseTx(ctx, s.pool, client.ID) if err != nil { return nil, err @@ -535,8 +530,6 @@ func (s *DesktopTaskService) leaseNextQueuedMonitorTask( client.TenantID, client.WorkspaceID, client.ID, - hasAccountIDs, - accountIDs, params.AttemptID, params.LeaseTokenHash, maxConcurrentMonitorPlatformsPerDesktopClient, @@ -553,37 +546,71 @@ func (s *DesktopTaskService) leaseNextQueuedMonitorTask( func leaseNextQueuedMonitorTaskSQL() string { return ` - 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 NOT EXISTS ( + WITH current_accounts AS ( + SELECT desktop_id, user_id, platform_id, platform_uid, account_fingerprint + FROM platform_accounts + WHERE tenant_id = $1 + AND workspace_id = $2 + AND client_id = $3 + AND deleted_at IS NULL + AND delete_requested_at IS NULL + ), + candidate_pool AS ( + SELECT dt.desktop_id, + ca.desktop_id AS lease_account_id, + dt.lane_weight, + dt.priority, + dt.enqueued_at, + dt.created_at, + ` + monitorAccountLeaseLockSQL("target_account") + ` AS account_slot_lock_key + FROM current_accounts AS ca + JOIN platform_accounts AS target_account + ON target_account.tenant_id = $1 + AND target_account.workspace_id = $2 + AND target_account.user_id = ca.user_id + AND target_account.platform_id = ca.platform_id + AND ` + monitorAccountIdentityMatchSQL("ca", "target_account") + ` + JOIN desktop_tasks AS dt + ON dt.target_account_id = target_account.desktop_id + AND dt.tenant_id = $1 + AND dt.workspace_id = $2 + AND dt.platform_id = ca.platform_id + AND dt.kind = 'monitor' + AND dt.status = 'queued' + WHERE NOT EXISTS ( SELECT 1 FROM desktop_tasks AS active + JOIN platform_accounts AS active_account + ON active_account.desktop_id = active.target_account_id + AND active_account.tenant_id = active.tenant_id + AND active_account.workspace_id = active.workspace_id WHERE active.tenant_id = dt.tenant_id AND active.workspace_id = dt.workspace_id AND active.kind = 'monitor' - AND active.target_client_id = $3 AND active.platform_id = dt.platform_id AND active.status = 'in_progress' AND active.desktop_id <> dt.desktop_id + AND active_account.user_id = target_account.user_id + AND ` + monitorAccountIdentityMatchSQL("active_account", "target_account") + ` ) AND NOT EXISTS ( SELECT 1 FROM desktop_tasks AS recent + JOIN platform_accounts AS recent_account + ON recent_account.desktop_id = recent.target_account_id + AND recent_account.tenant_id = recent.tenant_id + AND recent_account.workspace_id = recent.workspace_id WHERE recent.tenant_id = dt.tenant_id AND recent.workspace_id = dt.workspace_id AND recent.kind = 'monitor' - AND recent.target_client_id = $3 AND recent.platform_id = dt.platform_id AND ( (recent.status = 'succeeded' AND recent.updated_at >= now() - interval '` + monitorSucceededTaskCooldownSQL + `') OR (recent.status IN ('failed', 'unknown', 'aborted') AND recent.updated_at >= now() - interval '` + monitorUnhealthyTaskCooldownSQL + `') ) AND recent.desktop_id <> dt.desktop_id + AND recent_account.user_id = target_account.user_id + AND ` + monitorAccountIdentityMatchSQL("recent_account", "target_account") + ` ) AND ( SELECT COUNT(DISTINCT active_platform.platform_id) @@ -594,23 +621,44 @@ func leaseNextQueuedMonitorTaskSQL() string { AND active_platform.target_client_id = $3 AND active_platform.status = 'in_progress' AND active_platform.desktop_id <> dt.desktop_id - ) < $8::integer - AND ( - dt.target_client_id = $3 - OR ($4::boolean AND dt.target_account_id = ANY($5::uuid[])) + ) < $6::integer + AND NOT EXISTS ( + SELECT 1 + FROM desktop_tasks AS active_client_platform + WHERE active_client_platform.tenant_id = dt.tenant_id + AND active_client_platform.workspace_id = dt.workspace_id + AND active_client_platform.kind = 'monitor' + AND active_client_platform.target_client_id = $3 + AND active_client_platform.platform_id = dt.platform_id + AND active_client_platform.status = 'in_progress' + AND active_client_platform.desktop_id <> dt.desktop_id ) 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 + dt.desktop_id ASC, + ca.desktop_id ASC + LIMIT ` + monitorLeaseCandidateScanLimitSQL + ` + FOR UPDATE OF dt SKIP LOCKED + ), + candidate AS ( + SELECT desktop_id, lease_account_id + FROM candidate_pool + WHERE pg_try_advisory_xact_lock(account_slot_lock_key) + ORDER BY lane_weight DESC, + priority DESC, + COALESCE(enqueued_at, created_at) ASC, + created_at ASC, + desktop_id ASC, + lease_account_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, + target_account_id = candidate.lease_account_id, + active_attempt_id = $4, + lease_token_hash = $5, lease_expires_at = now() + interval '10 minutes', status = 'in_progress', attempts = t.attempts + 1, @@ -620,6 +668,51 @@ func leaseNextQueuedMonitorTaskSQL() string { RETURNING ` + desktopTaskRepositoryReturningColumns } +func monitorAccountLeaseLockSQL(alias string) string { + return fmt.Sprintf(`hashtextextended( + concat_ws( + ':', + 'monitor_desktop_account_slot', + %s.tenant_id::text, + %s.workspace_id::text, + %s.user_id::text, + %s.platform_id, + CASE + WHEN btrim(COALESCE(%s.account_fingerprint, '')) <> '' THEN 'fp:' || %s.account_fingerprint + WHEN btrim(COALESCE(%s.platform_uid, '')) <> '' THEN 'uid:' || %s.platform_uid + ELSE 'id:' || %s.desktop_id::text + END + ), + 0 + )`, + alias, + alias, + alias, + alias, + alias, alias, + alias, alias, + alias, + ) +} + +func monitorAccountIdentityMatchSQL(leftAlias, rightAlias string) string { + return fmt.Sprintf(`( + %s.desktop_id = %s.desktop_id + OR ( + btrim(COALESCE(%s.account_fingerprint, '')) <> '' + AND %s.account_fingerprint = %s.account_fingerprint + ) + OR ( + btrim(COALESCE(%s.platform_uid, '')) <> '' + AND %s.platform_uid = %s.platform_uid + ) + )`, + leftAlias, rightAlias, + leftAlias, leftAlias, rightAlias, + leftAlias, leftAlias, rightAlias, + ) +} + func (s *DesktopTaskService) leaseNextQueuedPublishTask( ctx context.Context, client *repository.DesktopClient, @@ -683,12 +776,6 @@ func (s *DesktopTaskService) leaseQueuedDesktopTaskByID( 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 - tx, err := beginDesktopMonitorLeaseTx(ctx, s.pool, client.ID) if err != nil { return nil, err @@ -700,8 +787,6 @@ func (s *DesktopTaskService) leaseQueuedDesktopTaskByID( client.WorkspaceID, client.TenantID, client.ID, - hasAccountIDs, - accountIDs, params.AttemptID, params.LeaseTokenHash, desktopPublishMaxAttempts, @@ -719,13 +804,35 @@ func (s *DesktopTaskService) leaseQueuedDesktopTaskByID( func leaseQueuedDesktopTaskByIDSQL() string { return ` - WITH candidate AS ( - SELECT dt.desktop_id + WITH current_accounts AS ( + SELECT desktop_id, user_id, platform_id, platform_uid, account_fingerprint + FROM platform_accounts + WHERE tenant_id = $3 + AND workspace_id = $2 + AND client_id = $4 + AND deleted_at IS NULL + AND delete_requested_at IS NULL + ), + candidate AS ( + SELECT dt.desktop_id, + dt.kind, + ca.desktop_id AS lease_account_id, + ` + monitorAccountLeaseLockSQL("target_account") + ` AS account_slot_lock_key FROM desktop_tasks AS dt + LEFT JOIN platform_accounts AS target_account + ON dt.kind = 'monitor' + AND target_account.desktop_id = dt.target_account_id + AND target_account.tenant_id = dt.tenant_id + AND target_account.workspace_id = dt.workspace_id + LEFT JOIN current_accounts AS ca + ON dt.kind = 'monitor' + AND ca.user_id = target_account.user_id + AND ca.platform_id = dt.platform_id + AND ` + monitorAccountIdentityMatchSQL("ca", "target_account") + ` WHERE dt.desktop_id = $1 AND dt.workspace_id = $2 AND dt.status = 'queued' - AND (dt.kind <> 'publish' OR dt.attempts < $9) + AND (dt.kind <> 'publish' OR dt.attempts < $7) AND ( dt.kind <> 'publish' OR EXISTS ( @@ -738,18 +845,27 @@ func leaseQueuedDesktopTaskByIDSQL() string { AND pa.delete_requested_at IS NULL ) ) + AND ( + dt.kind <> 'monitor' + OR ca.desktop_id IS NOT NULL + ) AND ( dt.kind <> 'monitor' OR NOT EXISTS ( SELECT 1 FROM desktop_tasks AS active + JOIN platform_accounts AS active_account + ON active_account.desktop_id = active.target_account_id + AND active_account.tenant_id = active.tenant_id + AND active_account.workspace_id = active.workspace_id WHERE active.tenant_id = dt.tenant_id AND active.workspace_id = dt.workspace_id AND active.kind = 'monitor' - AND active.target_client_id = $4 AND active.platform_id = dt.platform_id AND active.status = 'in_progress' AND active.desktop_id <> dt.desktop_id + AND active_account.user_id = target_account.user_id + AND ` + monitorAccountIdentityMatchSQL("active_account", "target_account") + ` ) ) AND ( @@ -757,16 +873,21 @@ func leaseQueuedDesktopTaskByIDSQL() string { OR NOT EXISTS ( SELECT 1 FROM desktop_tasks AS recent + JOIN platform_accounts AS recent_account + ON recent_account.desktop_id = recent.target_account_id + AND recent_account.tenant_id = recent.tenant_id + AND recent_account.workspace_id = recent.workspace_id WHERE recent.tenant_id = dt.tenant_id AND recent.workspace_id = dt.workspace_id AND recent.kind = 'monitor' - AND recent.target_client_id = $4 AND recent.platform_id = dt.platform_id AND ( (recent.status = 'succeeded' AND recent.updated_at >= now() - interval '` + monitorSucceededTaskCooldownSQL + `') OR (recent.status IN ('failed', 'unknown', 'aborted') AND recent.updated_at >= now() - interval '` + monitorUnhealthyTaskCooldownSQL + `') ) AND recent.desktop_id <> dt.desktop_id + AND recent_account.user_id = target_account.user_id + AND ` + monitorAccountIdentityMatchSQL("recent_account", "target_account") + ` ) ) AND ( @@ -780,36 +901,55 @@ func leaseQueuedDesktopTaskByIDSQL() string { AND active_platform.target_client_id = $4 AND active_platform.status = 'in_progress' AND active_platform.desktop_id <> dt.desktop_id - ) < $10::integer + ) < $8::integer + ) + AND ( + dt.kind <> 'monitor' + OR NOT EXISTS ( + SELECT 1 + FROM desktop_tasks AS active_client_platform + WHERE active_client_platform.tenant_id = dt.tenant_id + AND active_client_platform.workspace_id = dt.workspace_id + AND active_client_platform.kind = 'monitor' + AND active_client_platform.target_client_id = $4 + AND active_client_platform.platform_id = dt.platform_id + AND active_client_platform.status = 'in_progress' + AND active_client_platform.desktop_id <> dt.desktop_id + ) ) 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 ) ) + ORDER BY ca.desktop_id ASC NULLS LAST + LIMIT 1 + FOR UPDATE OF dt SKIP LOCKED + ), + locked_candidate AS ( + SELECT desktop_id, lease_account_id + FROM candidate + WHERE kind <> 'monitor' + OR pg_try_advisory_xact_lock(account_slot_lock_key) 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, + target_account_id = CASE WHEN t.kind = 'monitor' THEN locked_candidate.lease_account_id ELSE t.target_account_id END, + active_attempt_id = $5, + lease_token_hash = $6, lease_expires_at = now() + CASE WHEN t.kind = 'publish' THEN interval '3 minutes' ELSE interval '10 minutes' END, status = 'in_progress', attempts = t.attempts + 1, started_at = COALESCE(t.started_at, now()), updated_at = now() - FROM candidate - WHERE t.desktop_id = candidate.desktop_id + FROM locked_candidate + WHERE t.desktop_id = locked_candidate.desktop_id RETURNING ` + desktopTaskRepositoryReturningColumns } @@ -828,39 +968,6 @@ func beginDesktopMonitorLeaseTx(ctx context.Context, pool *pgxpool.Pool, clientI return tx, nil } -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 @@ -1662,16 +1769,28 @@ func (s *DesktopTaskService) clientCanLeaseMonitorTask(ctx context.Context, clie 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 { + var canLease bool + if err := s.pool.QueryRow(ctx, ` + SELECT EXISTS ( + SELECT 1 + FROM platform_accounts AS target_account + JOIN platform_accounts AS client_account + ON client_account.tenant_id = target_account.tenant_id + AND client_account.workspace_id = target_account.workspace_id + AND client_account.user_id = target_account.user_id + AND client_account.platform_id = target_account.platform_id + AND client_account.client_id = $4 + AND client_account.deleted_at IS NULL + AND client_account.delete_requested_at IS NULL + AND `+monitorAccountIdentityMatchSQL("client_account", "target_account")+` + WHERE target_account.desktop_id = $1 + AND target_account.tenant_id = $2 + AND target_account.workspace_id = $3 + ) + `, task.TargetAccountID, task.TenantID, task.WorkspaceID, client.ID).Scan(&canLease); err != nil { return false } - for _, accountID := range accountIDs { - if accountID == task.TargetAccountID { - return true - } - } - return false + return canLease } func buildDesktopTaskView(task *repository.DesktopTask) DesktopTaskView { diff --git a/server/internal/tenant/app/desktop_task_service_test.go b/server/internal/tenant/app/desktop_task_service_test.go index 2c5bb45..b3ff740 100644 --- a/server/internal/tenant/app/desktop_task_service_test.go +++ b/server/internal/tenant/app/desktop_task_service_test.go @@ -239,19 +239,28 @@ func TestBuildCountPublishTasksByStatusesQueryHidesDeletedPublishRecords(t *test } } -func TestLeaseNextQueuedMonitorTaskSQLSerializesPerClientPlatform(t *testing.T) { +func TestLeaseNextQueuedMonitorTaskSQLUsesAccountIdentityAndClientSlots(t *testing.T) { t.Parallel() normalized := normalizeSQLForDesktopTaskTest(leaseNextQueuedMonitorTaskSQL()) for _, fragment := range []string{ + "WITH current_accounts AS", + "client_id = $3", + "dt.target_account_id = target_account.desktop_id", + "target_account.account_fingerprint", + "target_account.platform_uid", + "hashtextextended", + "pg_try_advisory_xact_lock(account_slot_lock_key)", + "SET target_client_id = $3, target_account_id = candidate.lease_account_id", "NOT EXISTS ( SELECT 1 FROM desktop_tasks AS active", - "active.target_client_id = $3", "active.platform_id = dt.platform_id", "active.status = 'in_progress'", "active.desktop_id <> dt.desktop_id", "COUNT(DISTINCT active_platform.platform_id)", "active_platform.target_client_id = $3", - "< $8::integer", + "< $6::integer", + "active_client_platform.target_client_id = $3", + "active_client_platform.platform_id = dt.platform_id", "recent.status = 'succeeded'", "recent.updated_at >= now() - interval '2 seconds'", "recent.status IN ('failed', 'unknown', 'aborted')", @@ -264,19 +273,30 @@ func TestLeaseNextQueuedMonitorTaskSQLSerializesPerClientPlatform(t *testing.T) } } -func TestLeaseQueuedDesktopTaskByIDSQLSerializesMonitorPerClientPlatform(t *testing.T) { +func TestLeaseQueuedDesktopTaskByIDSQLUsesAccountIdentityAndClientSlots(t *testing.T) { t.Parallel() normalized := normalizeSQLForDesktopTaskTest(leaseQueuedDesktopTaskByIDSQL()) for _, fragment := range []string{ + "WITH current_accounts AS", + "client_id = $4", + "dt.kind <> 'monitor' OR ca.desktop_id IS NOT NULL", + "target_account.account_fingerprint", + "target_account.platform_uid", + "SELECT dt.desktop_id, dt.kind, ca.desktop_id AS lease_account_id", + "locked_candidate AS", + "pg_try_advisory_xact_lock(account_slot_lock_key)", + "FROM locked_candidate WHERE t.desktop_id = locked_candidate.desktop_id", + "target_account_id = CASE WHEN t.kind = 'monitor' THEN locked_candidate.lease_account_id ELSE t.target_account_id END", "dt.kind <> 'monitor' OR NOT EXISTS", - "active.target_client_id = $4", "active.platform_id = dt.platform_id", "active.status = 'in_progress'", "active.desktop_id <> dt.desktop_id", "COUNT(DISTINCT active_platform.platform_id)", "active_platform.target_client_id = $4", - "< $10::integer", + "< $8::integer", + "active_client_platform.target_client_id = $4", + "active_client_platform.platform_id = dt.platform_id", "recent.status = 'succeeded'", "recent.updated_at >= now() - interval '2 seconds'", "recent.status IN ('failed', 'unknown', 'aborted')", diff --git a/server/migrations/20260622190000_add_monitor_desktop_lease_indexes.down.sql b/server/migrations/20260622190000_add_monitor_desktop_lease_indexes.down.sql new file mode 100644 index 0000000..3695204 --- /dev/null +++ b/server/migrations/20260622190000_add_monitor_desktop_lease_indexes.down.sql @@ -0,0 +1,5 @@ +DROP INDEX CONCURRENTLY IF EXISTS idx_platform_accounts_monitor_identity_uid; +DROP INDEX CONCURRENTLY IF EXISTS idx_platform_accounts_monitor_identity_fingerprint; +DROP INDEX CONCURRENTLY IF EXISTS idx_desktop_tasks_monitor_active_client_platform; +DROP INDEX CONCURRENTLY IF EXISTS idx_desktop_tasks_monitor_active_platform_account; +DROP INDEX CONCURRENTLY IF EXISTS idx_desktop_tasks_monitor_queue_account_order; diff --git a/server/migrations/20260622190000_add_monitor_desktop_lease_indexes.up.sql b/server/migrations/20260622190000_add_monitor_desktop_lease_indexes.up.sql new file mode 100644 index 0000000..26abf13 --- /dev/null +++ b/server/migrations/20260622190000_add_monitor_desktop_lease_indexes.up.sql @@ -0,0 +1,62 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_desktop_tasks_monitor_queue_account_order + ON desktop_tasks ( + tenant_id, + workspace_id, + target_account_id, + platform_id, + lane_weight DESC, + priority DESC, + enqueued_at ASC, + created_at ASC, + desktop_id ASC + ) + WHERE kind = 'monitor' + AND status = 'queued'; + +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_desktop_tasks_monitor_active_platform_account + ON desktop_tasks ( + tenant_id, + workspace_id, + platform_id, + status, + target_account_id, + updated_at DESC + ) + WHERE kind = 'monitor' + AND status IN ('in_progress', 'succeeded', 'failed', 'unknown', 'aborted'); + +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_desktop_tasks_monitor_active_client_platform + ON desktop_tasks ( + tenant_id, + workspace_id, + target_client_id, + platform_id, + status, + desktop_id + ) + WHERE kind = 'monitor' + AND status = 'in_progress'; + +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_platform_accounts_monitor_identity_fingerprint + ON platform_accounts ( + tenant_id, + workspace_id, + user_id, + platform_id, + account_fingerprint + ) + WHERE deleted_at IS NULL + AND delete_requested_at IS NULL + AND btrim(COALESCE(account_fingerprint, '')) <> ''; + +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_platform_accounts_monitor_identity_uid + ON platform_accounts ( + tenant_id, + workspace_id, + user_id, + platform_id, + platform_uid + ) + WHERE deleted_at IS NULL + AND delete_requested_at IS NULL + AND btrim(COALESCE(platform_uid, '')) <> '';