diff --git a/server/internal/tenant/app/desktop_task_service.go b/server/internal/tenant/app/desktop_task_service.go index 4fd2eda..2b1f5f5 100644 --- a/server/internal/tenant/app/desktop_task_service.go +++ b/server/internal/tenant/app/desktop_task_service.go @@ -1363,27 +1363,8 @@ func (s *DesktopTaskService) recoverClientTasks( defer tx.Rollback(ctx) repo := repository.NewDesktopTaskRepository(tx) - query := ` - SELECT - desktop_id, - workspace_id, - kind, - active_attempt_id - FROM desktop_tasks - WHERE target_client_id = $1 - AND workspace_id = $2 - AND status = 'in_progress' - ` + query := desktopTaskRecoverySelectQuery(mode) queryArgs := []any{client.ID, client.WorkspaceID} - if mode == desktopTaskRecoveryModeLeaseExpiry { - query += ` - AND lease_expires_at IS NOT NULL - AND lease_expires_at < NOW() - ` - } - query += ` - FOR UPDATE - ` rows, err := tx.Query(ctx, query, queryArgs...) if err != nil { @@ -1402,11 +1383,13 @@ func (s *DesktopTaskService) recoverClientTasks( &item.Status, &item.ActiveAttempt, ); scanErr != nil { + s.logWarn("desktop task recovery scan failed", scanErr, zap.String("client_id", client.ID.String()), zap.String("mode", string(mode))) return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to parse recovered desktop tasks") } recovered = append(recovered, item) } if err := rows.Err(); err != nil { + s.logWarn("desktop task recovery rows failed", err, zap.String("client_id", client.ID.String()), zap.String("mode", string(mode))) return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to iterate recovered desktop tasks") } @@ -1509,6 +1492,31 @@ func (s *DesktopTaskService) recoverClientTasks( return nil } +func desktopTaskRecoverySelectQuery(mode desktopTaskRecoveryMode) string { + query := ` + SELECT + desktop_id, + workspace_id, + kind, + status, + active_attempt_id + FROM desktop_tasks + WHERE target_client_id = $1 + AND workspace_id = $2 + AND status = 'in_progress' + ` + if mode == desktopTaskRecoveryModeLeaseExpiry { + query += ` + AND lease_expires_at IS NOT NULL + AND lease_expires_at < NOW() + ` + } + query += ` + FOR UPDATE + ` + return query +} + func desktopTaskRecoveryPayload(mode desktopTaskRecoveryMode, kind string) ([]byte, string, string, error) { reason := strings.TrimSpace(string(mode)) message := "desktop task was recovered after the client lost the active lease" diff --git a/server/internal/tenant/app/desktop_task_service_test.go b/server/internal/tenant/app/desktop_task_service_test.go new file mode 100644 index 0000000..e8fc142 --- /dev/null +++ b/server/internal/tenant/app/desktop_task_service_test.go @@ -0,0 +1,63 @@ +package app + +import ( + "regexp" + "strings" + "testing" +) + +func TestDesktopTaskRecoverySelectQueryMatchesRecoveredLeaseScanOrder(t *testing.T) { + t.Parallel() + + query := desktopTaskRecoverySelectQuery(desktopTaskRecoveryModeStartup) + gotColumns := recoverDesktopTaskSelectColumns(query) + wantColumns := []string{ + "desktop_id", + "workspace_id", + "kind", + "status", + "active_attempt_id", + } + + if len(gotColumns) != len(wantColumns) { + t.Fatalf("recover select columns = %v, want %v", gotColumns, wantColumns) + } + for index := range wantColumns { + if gotColumns[index] != wantColumns[index] { + t.Fatalf("recover select columns = %v, want %v", gotColumns, wantColumns) + } + } +} + +func TestDesktopTaskRecoverySelectQueryLeaseExpiryFiltersExpiredLeases(t *testing.T) { + t.Parallel() + + query := desktopTaskRecoverySelectQuery(desktopTaskRecoveryModeLeaseExpiry) + + for _, fragment := range []string{ + "AND lease_expires_at IS NOT NULL", + "AND lease_expires_at < NOW()", + "FOR UPDATE", + } { + if !strings.Contains(query, fragment) { + t.Fatalf("recover lease-expiry query missing %q: %s", fragment, query) + } + } +} + +func recoverDesktopTaskSelectColumns(query string) []string { + re := regexp.MustCompile(`(?is)select\s+(.*?)\s+from\s+desktop_tasks`) + match := re.FindStringSubmatch(query) + if len(match) != 2 { + return nil + } + + rawColumns := strings.Split(match[1], ",") + columns := make([]string, 0, len(rawColumns)) + for _, rawColumn := range rawColumns { + column := strings.TrimSpace(rawColumn) + column = strings.Fields(column)[0] + columns = append(columns, column) + } + return columns +}