diff --git a/docs/rfcs/0002-active-passive-fleet-ha.md b/docs/rfcs/0002-active-passive-fleet-ha.md index b6748521f..1d7f700ef 100644 --- a/docs/rfcs/0002-active-passive-fleet-ha.md +++ b/docs/rfcs/0002-active-passive-fleet-ha.md @@ -3,7 +3,7 @@ - **Status**: approved - **Author(s)**: Ankit Goswami (@ankitgoswami) - **Created**: 2026-07-13 -- **Last updated**: 2026-07-14 +- **Last updated**: 2026-08-04 ## Summary @@ -29,7 +29,7 @@ This RFC deliberately scopes the HA promise to the real-time control plane. Curt The HA design needs to satisfy four requirements: - **No single point of failure** in the supported HA topology. -- **Automatic recovery** of curtailment dispatch, with a target recovery time under 60 seconds. +- **Automatic recovery** of curtailment dispatch, with a target recovery time under 180 seconds. - **Correctness over convenience**: never run two active control dispatchers at once. - **Deployment simplicity** across on-prem and cloud installs. @@ -205,11 +205,11 @@ Rules: - Acquire succeeds when the row is absent, expired according to DB time, or already held by the same process incarnation. - Renew succeeds only when `name`, `holder_id`, and `lease_epoch` still match and the lease has not expired according to DB time. - A process that resumes after its lease expires must reacquire the lease and receive a new epoch before active runtime can resume. -- Active renews every few seconds with a short TTL. +- Active renews every few seconds with a 10-second TTL. - Lease renewal and `/health/active` depend on critical active-runtime health. If the active instance cannot run required control loops, it must stop passing active health and stop renewing or relinquish the lease. -- Every active-only runtime loop carries the activation epoch and confirms the local coordinator still owns that epoch before dispatching or claiming work. -- External side effects must be fenced at the nearest durable boundary. Effects that cannot be receiver-fenced must be idempotent, cancelable, or reconciled before the implementation can claim zero dual dispatch. -- A stalled old active that resumes after takeover cannot dispatch because its epoch is stale. +- Command write transactions are bounded to 5 seconds, and state transitions require the expected prior state. +- When command execution starts, it preserves `PENDING` work and fails `PROCESSING` work because the device outcome is unknown. The existing reaper then handles stale work and terminal batches on its normal schedule. +- After an active process loses ownership or critical runtime health, it exits instead of trying to demote and restart active work in place. The service supervisor restarts it in passive mode. This lease model intentionally allows future per-subsystem leases, but v1 does not use them. Active/active Fleet is out of scope. @@ -226,7 +226,7 @@ Passive is deliberately dumb in v1. It does not need a carefully audited read-on ### Runtime supervisor -Add a runtime supervisor that starts and stops active-only services on lease transitions. Startup moves from "all runtime work starts when the process starts" to "active runtime work starts only after activation." +Add a runtime supervisor that starts active-only services after lease acquisition. In HA mode, ownership loss terminates the process so the external service supervisor provides a clean restart. Active-only work in v1 includes: @@ -323,7 +323,7 @@ The exact installer flags, templates, compose files, and runbook commands belong | Failure | Expected behavior | | ------- | ----------------- | | Active Fleet app process dies | Lease expires; peer activates and starts active runtime. | -| Active Fleet app hangs | Lease renewal stops or epoch checks fail; peer takes over; old active cannot dispatch with stale epoch. | +| Active Fleet app hangs | Lease renewal stops; the process exits or is replaced after lease expiry, and the peer takes over. | | DB primary dies | Patroni promotes standby; new Fleet DB connections select the read-write host; lease renews or is reacquired. | | Sync standby dies | Fleet continues in async-degraded mode after the configured behavior; HA status reports `FAILOVER READY: NO`. | | Witness host dies | Service continues; quorum tolerance is degraded; HA status reports degraded readiness. | @@ -341,12 +341,12 @@ The HA mode is not supported until these gates pass. Activation and fencing: - Two Fleet app processes against the same DB produce exactly one active holder. -- Killing active Fleet activates the peer within 15 seconds. +- Killing active Fleet keeps end-to-end curtailment recovery under 180 seconds. - Partial active-runtime failure causes `/health/active` to fail and allows the peer to take over. -- Network partition tests produce zero dual dispatch. +- Network partition tests prove guarded transitions cannot overwrite command state after ownership loss. - Passive mode rejects all product traffic, including non-RPC HTTP routes, while preserving the explicit health and operator-status bypasses. - Lease loss terminates already accepted product streams and active-scoped request work. -- A stalled active process cannot dispatch after lease loss, emit stale external side effects, or renew an expired lease with its old epoch. +- A stalled active process cannot overwrite terminal command state or renew an expired lease with its old epoch. - Fleet-scale reconnect tests avoid synchronized ControlStream reconnect storms during failover. Database and durability: @@ -359,8 +359,8 @@ Database and durability: Real-time control: -- Full curtailment dispatch recovery after a single failure completes within 60 seconds across repeated trials. -- MQTT curtailment intake failover resumes on the new active Fleet app within the RTO target without dual processing. +- Full curtailment dispatch recovery after a single failure completes within 180 seconds across repeated trials. +- MQTT curtailment intake failover resumes on the new active Fleet app within the RTO target; pending command work resumes and interrupted attempts are failed. Deployment and diagnostics: @@ -397,7 +397,7 @@ Deployment and diagnostics: ## Unresolved questions -- **Exact Patroni timings**. Initial targets are chosen to fit the 60s RTO, but final `ttl`, `loop_wait`, and `retry_timeout` values must be set from lab measurements. +- **Exact Patroni timings**. Initial targets are chosen to fit the 180-second curtailment RTO, but final `ttl`, `loop_wait`, and `retry_timeout` values must be confirmed by lab measurements. - **Critical write classification**. The implementation must audit write paths and decide which writes require critical durability and which are best-effort. - **Artifact HA boundary**. Firmware and command artifact behavior after failover needs a product decision: document re-upload in v1, add rsync, or require shared storage for covered command types. - **Grafana HA datasource**. Grafana is out of the RTO path, but the install should still decide whether dashboards use best-effort single-host datasource config, multi-host config if supported, or an explicit "not HA" warning. diff --git a/server/cmd/fleetd/runtime_jobs_test.go b/server/cmd/fleetd/runtime_jobs_test.go index 74137d3f2..7ff742bfa 100644 --- a/server/cmd/fleetd/runtime_jobs_test.go +++ b/server/cmd/fleetd/runtime_jobs_test.go @@ -21,7 +21,6 @@ func (noopLifecycle) Stop(context.Context) error { return nil } type funcLifecycle struct { start func(context.Context) error stop func(context.Context) error - abort func() } func (l funcLifecycle) Start(ctx context.Context) error { @@ -38,12 +37,6 @@ func (l funcLifecycle) Stop(ctx context.Context) error { return l.stop(ctx) } -func (l funcLifecycle) Abort() { - if l.abort != nil { - l.abort() - } -} - type scriptedRuntimeJobGroupStopper struct { stop func(context.Context) error contexts []context.Context @@ -176,30 +169,6 @@ func TestRuntimeJobGroupKeepsCommandExecutionAliveWhileProducersDrain(t *testing } } -func TestRuntimeJobGroupAbortsCommandExecution(t *testing.T) { - commandAborted := false - jobs, err := newRuntimeJobs(runtimeJobLifecycles{ - identityStateCleanup: noopLifecycle{}, - commandArtifactCleanup: noopLifecycle{}, - diagnosticsErrorCloser: noopLifecycle{}, - telemetry: noopLifecycle{}, - ipScanner: noopLifecycle{}, - commandExecution: funcLifecycle{abort: func() { commandAborted = true }}, - scheduleProcessor: noopLifecycle{}, - curtailmentReconciler: noopLifecycle{}, - curtailmentMQTTSubscriber: noopLifecycle{}, - curtailmentRigConfig: noopLifecycle{}, - chunkedUploadCleanup: noopLifecycle{}, - }) - require.NoError(t, err) - group, err := runtimejobs.NewGroup(jobs) - require.NoError(t, err) - - group.Abort() - - require.True(t, commandAborted) -} - func TestBackgroundLoopCanRestartAfterDraining(t *testing.T) { started := make(chan struct{}, 2) loop := newBackgroundLoop(func(ctx context.Context) { diff --git a/server/generated/sqlc/command.sql.go b/server/generated/sqlc/command.sql.go index 9c03a6b58..b248d93db 100644 --- a/server/generated/sqlc/command.sql.go +++ b/server/generated/sqlc/command.sql.go @@ -246,41 +246,73 @@ func (q *Queries) ListBatchDeviceResults(ctx context.Context, arg ListBatchDevic return items, nil } -const markCommandBatchFinished = `-- name: MarkCommandBatchFinished :exec +const lockCommandBatch = `-- name: LockCommandBatch :one +SELECT status +FROM command_batch_log +WHERE uuid = $1 +FOR UPDATE +` + +func (q *Queries) LockCommandBatch(ctx context.Context, uuid string) (BatchStatusEnum, error) { + row := q.queryRow(ctx, q.lockCommandBatchStmt, lockCommandBatch, uuid) + var status BatchStatusEnum + err := row.Scan(&status) + return status, err +} + +const markCommandBatchFinished = `-- name: MarkCommandBatchFinished :execrows UPDATE command_batch_log SET status = 'FINISHED', finished_at = NOW() WHERE uuid = $1 + AND status IN ('PENDING', 'PROCESSING') ` -func (q *Queries) MarkCommandBatchFinished(ctx context.Context, uuid string) error { - _, err := q.exec(ctx, q.markCommandBatchFinishedStmt, markCommandBatchFinished, uuid) - return err +func (q *Queries) MarkCommandBatchFinished(ctx context.Context, uuid string) (int64, error) { + result, err := q.exec(ctx, q.markCommandBatchFinishedStmt, markCommandBatchFinished, uuid) + if err != nil { + return 0, err + } + return result.RowsAffected() } -const markCommandBatchFinishedWithStartedAt = `-- name: MarkCommandBatchFinishedWithStartedAt :exec +const markCommandBatchFinishedWithStartedAt = `-- name: MarkCommandBatchFinishedWithStartedAt :execrows UPDATE command_batch_log SET status = 'FINISHED', started_at = NOW(), finished_at = NOW() WHERE uuid = $1 + AND status = 'PENDING' ` -func (q *Queries) MarkCommandBatchFinishedWithStartedAt(ctx context.Context, uuid string) error { - _, err := q.exec(ctx, q.markCommandBatchFinishedWithStartedAtStmt, markCommandBatchFinishedWithStartedAt, uuid) - return err +func (q *Queries) MarkCommandBatchFinishedWithStartedAt(ctx context.Context, uuid string) (int64, error) { + result, err := q.exec(ctx, q.markCommandBatchFinishedWithStartedAtStmt, markCommandBatchFinishedWithStartedAt, uuid) + if err != nil { + return 0, err + } + return result.RowsAffected() } -const markCommandBatchProcessing = `-- name: MarkCommandBatchProcessing :exec -UPDATE command_batch_log +const markCommandBatchProcessing = `-- name: MarkCommandBatchProcessing :execrows +UPDATE command_batch_log AS batch SET status = 'PROCESSING', started_at = NOW() -WHERE uuid = $1 +WHERE batch.uuid = $1 + AND batch.status = 'PENDING' + AND EXISTS ( + SELECT 1 + FROM queue_message AS message + WHERE message.command_batch_log_uuid = batch.uuid + AND message.status = 'PROCESSING' + ) ` -func (q *Queries) MarkCommandBatchProcessing(ctx context.Context, uuid string) error { - _, err := q.exec(ctx, q.markCommandBatchProcessingStmt, markCommandBatchProcessing, uuid) - return err +func (q *Queries) MarkCommandBatchProcessing(ctx context.Context, uuid string) (int64, error) { + result, err := q.exec(ctx, q.markCommandBatchProcessingStmt, markCommandBatchProcessing, uuid) + if err != nil { + return 0, err + } + return result.RowsAffected() } const upsertCommandOnDeviceLog = `-- name: UpsertCommandOnDeviceLog :exec diff --git a/server/generated/sqlc/db.go b/server/generated/sqlc/db.go index 236c99b15..83f06f8ef 100644 --- a/server/generated/sqlc/db.go +++ b/server/generated/sqlc/db.go @@ -219,6 +219,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.countOrgScopeSuperAdminsExcludingUserStmt, err = db.PrepareContext(ctx, countOrgScopeSuperAdminsExcludingUser); err != nil { return nil, fmt.Errorf("error preparing query CountOrgScopeSuperAdminsExcludingUser: %w", err) } + if q.countQueueMessagesByBatchStmt, err = db.PrepareContext(ctx, countQueueMessagesByBatch); err != nil { + return nil, fmt.Errorf("error preparing query CountQueueMessagesByBatch: %w", err) + } if q.countRacksBySiteStmt, err = db.PrepareContext(ctx, countRacksBySite); err != nil { return nil, fmt.Errorf("error preparing query CountRacksBySite: %w", err) } @@ -267,6 +270,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.createQueueMessageStmt, err = db.PrepareContext(ctx, createQueueMessage); err != nil { return nil, fmt.Errorf("error preparing query CreateQueueMessage: %w", err) } + if q.createQueueMessagesStmt, err = db.PrepareContext(ctx, createQueueMessages); err != nil { + return nil, fmt.Errorf("error preparing query CreateQueueMessages: %w", err) + } if q.createRackExtensionStmt, err = db.PrepareContext(ctx, createRackExtension); err != nil { return nil, fmt.Errorf("error preparing query CreateRackExtension: %w", err) } @@ -372,6 +378,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.findDevicesWithSiteOrBuildingStmt, err = db.PrepareContext(ctx, findDevicesWithSiteOrBuilding); err != nil { return nil, fmt.Errorf("error preparing query FindDevicesWithSiteOrBuilding: %w", err) } + if q.finishTerminalCommandBatchesStmt, err = db.PrepareContext(ctx, finishTerminalCommandBatches); err != nil { + return nil, fmt.Errorf("error preparing query FinishTerminalCommandBatches: %w", err) + } if q.forceReleaseCurtailmentEventStmt, err = db.PrepareContext(ctx, forceReleaseCurtailmentEvent); err != nil { return nil, fmt.Errorf("error preparing query ForceReleaseCurtailmentEvent: %w", err) } @@ -750,6 +759,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.getPoolStmt, err = db.PrepareContext(ctx, getPool); err != nil { return nil, fmt.Errorf("error preparing query GetPool: %w", err) } + if q.getQueueMessagesByBatchStmt, err = db.PrepareContext(ctx, getQueueMessagesByBatch); err != nil { + return nil, fmt.Errorf("error preparing query GetQueueMessagesByBatch: %w", err) + } if q.getRackDetailsForDevicesStmt, err = db.PrepareContext(ctx, getRackDetailsForDevices); err != nil { return nil, fmt.Errorf("error preparing query GetRackDetailsForDevices: %w", err) } @@ -876,9 +888,6 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.isBatchFinishedStmt, err = db.PrepareContext(ctx, isBatchFinished); err != nil { return nil, fmt.Errorf("error preparing query IsBatchFinished: %w", err) } - if q.isBatchProcessingStmt, err = db.PrepareContext(ctx, isBatchProcessing); err != nil { - return nil, fmt.Errorf("error preparing query IsBatchProcessing: %w", err) - } if q.isDeviceOwnedByFleetNodeStmt, err = db.PrepareContext(ctx, isDeviceOwnedByFleetNode); err != nil { return nil, fmt.Errorf("error preparing query IsDeviceOwnedByFleetNode: %w", err) } @@ -1092,6 +1101,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.lockBuildingsBySiteForWriteStmt, err = db.PrepareContext(ctx, lockBuildingsBySiteForWrite); err != nil { return nil, fmt.Errorf("error preparing query LockBuildingsBySiteForWrite: %w", err) } + if q.lockCommandBatchStmt, err = db.PrepareContext(ctx, lockCommandBatch); err != nil { + return nil, fmt.Errorf("error preparing query LockCommandBatch: %w", err) + } if q.lockCurtailmentEventByUUIDForWriteStmt, err = db.PrepareContext(ctx, lockCurtailmentEventByUUIDForWrite); err != nil { return nil, fmt.Errorf("error preparing query LockCurtailmentEventByUUIDForWrite: %w", err) } @@ -1173,11 +1185,8 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.queryErrorsStmt, err = db.PrepareContext(ctx, queryErrors); err != nil { return nil, fmt.Errorf("error preparing query QueryErrors: %w", err) } - if q.reapStuckFirmwareUpdateMessagesStmt, err = db.PrepareContext(ctx, reapStuckFirmwareUpdateMessages); err != nil { - return nil, fmt.Errorf("error preparing query ReapStuckFirmwareUpdateMessages: %w", err) - } - if q.reapStuckProcessingMessagesStmt, err = db.PrepareContext(ctx, reapStuckProcessingMessages); err != nil { - return nil, fmt.Errorf("error preparing query ReapStuckProcessingMessages: %w", err) + if q.reapMessagesStmt, err = db.PrepareContext(ctx, reapMessages); err != nil { + return nil, fmt.Errorf("error preparing query ReapMessages: %w", err) } if q.reassignDevicesUnderBuildingStmt, err = db.PrepareContext(ctx, reassignDevicesUnderBuilding); err != nil { return nil, fmt.Errorf("error preparing query ReassignDevicesUnderBuilding: %w", err) @@ -1230,6 +1239,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.resetCurtailmentTargetsForRestoreStmt, err = db.PrepareContext(ctx, resetCurtailmentTargetsForRestore); err != nil { return nil, fmt.Errorf("error preparing query ResetCurtailmentTargetsForRestore: %w", err) } + if q.resetReapedFirmwareStatusesStmt, err = db.PrepareContext(ctx, resetReapedFirmwareStatuses); err != nil { + return nil, fmt.Errorf("error preparing query ResetReapedFirmwareStatuses: %w", err) + } if q.resumeCurtailmentFromRestoringStmt, err = db.PrepareContext(ctx, resumeCurtailmentFromRestoring); err != nil { return nil, fmt.Errorf("error preparing query ResumeCurtailmentFromRestoring: %w", err) } @@ -1278,6 +1290,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.setInfrastructureControlSubnetsStmt, err = db.PrepareContext(ctx, setInfrastructureControlSubnets); err != nil { return nil, fmt.Errorf("error preparing query SetInfrastructureControlSubnets: %w", err) } + if q.setLocalTransactionTimeoutStmt, err = db.PrepareContext(ctx, setLocalTransactionTimeout); err != nil { + return nil, fmt.Errorf("error preparing query SetLocalTransactionTimeout: %w", err) + } if q.setMQTTSourceConfigEnabledStmt, err = db.PrepareContext(ctx, setMQTTSourceConfigEnabled); err != nil { return nil, fmt.Errorf("error preparing query SetMQTTSourceConfigEnabled: %w", err) } @@ -1923,6 +1938,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing countOrgScopeSuperAdminsExcludingUserStmt: %w", cerr) } } + if q.countQueueMessagesByBatchStmt != nil { + if cerr := q.countQueueMessagesByBatchStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing countQueueMessagesByBatchStmt: %w", cerr) + } + } if q.countRacksBySiteStmt != nil { if cerr := q.countRacksBySiteStmt.Close(); cerr != nil { err = fmt.Errorf("error closing countRacksBySiteStmt: %w", cerr) @@ -2003,6 +2023,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing createQueueMessageStmt: %w", cerr) } } + if q.createQueueMessagesStmt != nil { + if cerr := q.createQueueMessagesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing createQueueMessagesStmt: %w", cerr) + } + } if q.createRackExtensionStmt != nil { if cerr := q.createRackExtensionStmt.Close(); cerr != nil { err = fmt.Errorf("error closing createRackExtensionStmt: %w", cerr) @@ -2178,6 +2203,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing findDevicesWithSiteOrBuildingStmt: %w", cerr) } } + if q.finishTerminalCommandBatchesStmt != nil { + if cerr := q.finishTerminalCommandBatchesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing finishTerminalCommandBatchesStmt: %w", cerr) + } + } if q.forceReleaseCurtailmentEventStmt != nil { if cerr := q.forceReleaseCurtailmentEventStmt.Close(); cerr != nil { err = fmt.Errorf("error closing forceReleaseCurtailmentEventStmt: %w", cerr) @@ -2808,6 +2838,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing getPoolStmt: %w", cerr) } } + if q.getQueueMessagesByBatchStmt != nil { + if cerr := q.getQueueMessagesByBatchStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getQueueMessagesByBatchStmt: %w", cerr) + } + } if q.getRackDetailsForDevicesStmt != nil { if cerr := q.getRackDetailsForDevicesStmt.Close(); cerr != nil { err = fmt.Errorf("error closing getRackDetailsForDevicesStmt: %w", cerr) @@ -3018,11 +3053,6 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing isBatchFinishedStmt: %w", cerr) } } - if q.isBatchProcessingStmt != nil { - if cerr := q.isBatchProcessingStmt.Close(); cerr != nil { - err = fmt.Errorf("error closing isBatchProcessingStmt: %w", cerr) - } - } if q.isDeviceOwnedByFleetNodeStmt != nil { if cerr := q.isDeviceOwnedByFleetNodeStmt.Close(); cerr != nil { err = fmt.Errorf("error closing isDeviceOwnedByFleetNodeStmt: %w", cerr) @@ -3378,6 +3408,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing lockBuildingsBySiteForWriteStmt: %w", cerr) } } + if q.lockCommandBatchStmt != nil { + if cerr := q.lockCommandBatchStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing lockCommandBatchStmt: %w", cerr) + } + } if q.lockCurtailmentEventByUUIDForWriteStmt != nil { if cerr := q.lockCurtailmentEventByUUIDForWriteStmt.Close(); cerr != nil { err = fmt.Errorf("error closing lockCurtailmentEventByUUIDForWriteStmt: %w", cerr) @@ -3513,14 +3548,9 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing queryErrorsStmt: %w", cerr) } } - if q.reapStuckFirmwareUpdateMessagesStmt != nil { - if cerr := q.reapStuckFirmwareUpdateMessagesStmt.Close(); cerr != nil { - err = fmt.Errorf("error closing reapStuckFirmwareUpdateMessagesStmt: %w", cerr) - } - } - if q.reapStuckProcessingMessagesStmt != nil { - if cerr := q.reapStuckProcessingMessagesStmt.Close(); cerr != nil { - err = fmt.Errorf("error closing reapStuckProcessingMessagesStmt: %w", cerr) + if q.reapMessagesStmt != nil { + if cerr := q.reapMessagesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing reapMessagesStmt: %w", cerr) } } if q.reassignDevicesUnderBuildingStmt != nil { @@ -3608,6 +3638,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing resetCurtailmentTargetsForRestoreStmt: %w", cerr) } } + if q.resetReapedFirmwareStatusesStmt != nil { + if cerr := q.resetReapedFirmwareStatusesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing resetReapedFirmwareStatusesStmt: %w", cerr) + } + } if q.resumeCurtailmentFromRestoringStmt != nil { if cerr := q.resumeCurtailmentFromRestoringStmt.Close(); cerr != nil { err = fmt.Errorf("error closing resumeCurtailmentFromRestoringStmt: %w", cerr) @@ -3688,6 +3723,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing setInfrastructureControlSubnetsStmt: %w", cerr) } } + if q.setLocalTransactionTimeoutStmt != nil { + if cerr := q.setLocalTransactionTimeoutStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing setLocalTransactionTimeoutStmt: %w", cerr) + } + } if q.setMQTTSourceConfigEnabledStmt != nil { if cerr := q.setMQTTSourceConfigEnabledStmt.Close(); cerr != nil { err = fmt.Errorf("error closing setMQTTSourceConfigEnabledStmt: %w", cerr) @@ -4317,6 +4357,7 @@ type Queries struct { countMinersByStateStmt *sql.Stmt countNonTerminalCurtailmentEventsByInfrastructureDevicesStmt *sql.Stmt countOrgScopeSuperAdminsExcludingUserStmt *sql.Stmt + countQueueMessagesByBatchStmt *sql.Stmt countRacksBySiteStmt *sql.Stmt countRacksInBuildingStmt *sql.Stmt countResponseProfilesByInfrastructureDeviceStmt *sql.Stmt @@ -4333,6 +4374,7 @@ type Queries struct { createPendingEnrollmentStmt *sql.Stmt createPoolStmt *sql.Stmt createQueueMessageStmt *sql.Stmt + createQueueMessagesStmt *sql.Stmt createRackExtensionStmt *sql.Stmt createScheduleStmt *sql.Stmt createScheduleTargetStmt *sql.Stmt @@ -4368,6 +4410,7 @@ type Queries struct { findDevicesInBuildingLessPlacedRacksStmt *sql.Stmt findDevicesInSiteLessRacksStmt *sql.Stmt findDevicesWithSiteOrBuildingStmt *sql.Stmt + finishTerminalCommandBatchesStmt *sql.Stmt forceReleaseCurtailmentEventStmt *sql.Stmt getActiveSchedulesStmt *sql.Stmt getActiveUnpairedDiscoveredDevicesStmt *sql.Stmt @@ -4494,6 +4537,7 @@ type Queries struct { getPermissionByKeyStmt *sql.Stmt getPermissionsByKeysStmt *sql.Stmt getPoolStmt *sql.Stmt + getQueueMessagesByBatchStmt *sql.Stmt getRackDetailsForDevicesStmt *sql.Stmt getRackInfoStmt *sql.Stmt getRackInfoBatchStmt *sql.Stmt @@ -4536,7 +4580,6 @@ type Queries struct { insertNotificationHistoryStmt *sql.Stmt insertNotificationMetricSamplesStmt *sql.Stmt isBatchFinishedStmt *sql.Stmt - isBatchProcessingStmt *sql.Stmt isDeviceOwnedByFleetNodeStmt *sql.Stmt listActiveCurtailedDevicesByOrgStmt *sql.Stmt listActiveCurtailmentEventsStmt *sql.Stmt @@ -4608,6 +4651,7 @@ type Queries struct { lockAndCountOrgScopeSuperAdminsStmt *sql.Stmt lockBuildingForWriteStmt *sql.Stmt lockBuildingsBySiteForWriteStmt *sql.Stmt + lockCommandBatchStmt *sql.Stmt lockCurtailmentEventByUUIDForWriteStmt *sql.Stmt lockCurtailmentEventForFanCommandStmt *sql.Stmt lockCurtailmentFanDeviceForWriteStmt *sql.Stmt @@ -4635,8 +4679,7 @@ type Queries struct { queryComponentKeysWithErrorsStmt *sql.Stmt queryDeviceIDsWithErrorsStmt *sql.Stmt queryErrorsStmt *sql.Stmt - reapStuckFirmwareUpdateMessagesStmt *sql.Stmt - reapStuckProcessingMessagesStmt *sql.Stmt + reapMessagesStmt *sql.Stmt reassignDevicesUnderBuildingStmt *sql.Stmt reassignDevicesUnderBuildingsBulkStmt *sql.Stmt reassignRacksUnderBuildingStmt *sql.Stmt @@ -4654,6 +4697,7 @@ type Queries struct { requeueRigConfigReconciliationAfterTerminalFailureStmt *sql.Stmt resetCurtailmentTargetsForRecurtailStmt *sql.Stmt resetCurtailmentTargetsForRestoreStmt *sql.Stmt + resetReapedFirmwareStatusesStmt *sql.Stmt resumeCurtailmentFromRestoringStmt *sql.Stmt resumePausedScheduleStmt *sql.Stmt retryRigConfigReconciliationStmt *sql.Stmt @@ -4670,6 +4714,7 @@ type Queries struct { setDevicePairingAuthNeededIfNotPairedStmt *sql.Stmt setFleetNodeEnrollmentStatusStmt *sql.Stmt setInfrastructureControlSubnetsStmt *sql.Stmt + setLocalTransactionTimeoutStmt *sql.Stmt setMQTTSourceConfigEnabledStmt *sql.Stmt setRackBuildingPositionStmt *sql.Stmt setRackBuildingPositionBulkClearStmt *sql.Stmt @@ -4846,6 +4891,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { countMinersByStateStmt: q.countMinersByStateStmt, countNonTerminalCurtailmentEventsByInfrastructureDevicesStmt: q.countNonTerminalCurtailmentEventsByInfrastructureDevicesStmt, countOrgScopeSuperAdminsExcludingUserStmt: q.countOrgScopeSuperAdminsExcludingUserStmt, + countQueueMessagesByBatchStmt: q.countQueueMessagesByBatchStmt, countRacksBySiteStmt: q.countRacksBySiteStmt, countRacksInBuildingStmt: q.countRacksInBuildingStmt, countResponseProfilesByInfrastructureDeviceStmt: q.countResponseProfilesByInfrastructureDeviceStmt, @@ -4862,6 +4908,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { createPendingEnrollmentStmt: q.createPendingEnrollmentStmt, createPoolStmt: q.createPoolStmt, createQueueMessageStmt: q.createQueueMessageStmt, + createQueueMessagesStmt: q.createQueueMessagesStmt, createRackExtensionStmt: q.createRackExtensionStmt, createScheduleStmt: q.createScheduleStmt, createScheduleTargetStmt: q.createScheduleTargetStmt, @@ -4897,6 +4944,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { findDevicesInBuildingLessPlacedRacksStmt: q.findDevicesInBuildingLessPlacedRacksStmt, findDevicesInSiteLessRacksStmt: q.findDevicesInSiteLessRacksStmt, findDevicesWithSiteOrBuildingStmt: q.findDevicesWithSiteOrBuildingStmt, + finishTerminalCommandBatchesStmt: q.finishTerminalCommandBatchesStmt, forceReleaseCurtailmentEventStmt: q.forceReleaseCurtailmentEventStmt, getActiveSchedulesStmt: q.getActiveSchedulesStmt, getActiveUnpairedDiscoveredDevicesStmt: q.getActiveUnpairedDiscoveredDevicesStmt, @@ -5023,6 +5071,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { getPermissionByKeyStmt: q.getPermissionByKeyStmt, getPermissionsByKeysStmt: q.getPermissionsByKeysStmt, getPoolStmt: q.getPoolStmt, + getQueueMessagesByBatchStmt: q.getQueueMessagesByBatchStmt, getRackDetailsForDevicesStmt: q.getRackDetailsForDevicesStmt, getRackInfoStmt: q.getRackInfoStmt, getRackInfoBatchStmt: q.getRackInfoBatchStmt, @@ -5065,7 +5114,6 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { insertNotificationHistoryStmt: q.insertNotificationHistoryStmt, insertNotificationMetricSamplesStmt: q.insertNotificationMetricSamplesStmt, isBatchFinishedStmt: q.isBatchFinishedStmt, - isBatchProcessingStmt: q.isBatchProcessingStmt, isDeviceOwnedByFleetNodeStmt: q.isDeviceOwnedByFleetNodeStmt, listActiveCurtailedDevicesByOrgStmt: q.listActiveCurtailedDevicesByOrgStmt, listActiveCurtailmentEventsStmt: q.listActiveCurtailmentEventsStmt, @@ -5137,6 +5185,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { lockAndCountOrgScopeSuperAdminsStmt: q.lockAndCountOrgScopeSuperAdminsStmt, lockBuildingForWriteStmt: q.lockBuildingForWriteStmt, lockBuildingsBySiteForWriteStmt: q.lockBuildingsBySiteForWriteStmt, + lockCommandBatchStmt: q.lockCommandBatchStmt, lockCurtailmentEventByUUIDForWriteStmt: q.lockCurtailmentEventByUUIDForWriteStmt, lockCurtailmentEventForFanCommandStmt: q.lockCurtailmentEventForFanCommandStmt, lockCurtailmentFanDeviceForWriteStmt: q.lockCurtailmentFanDeviceForWriteStmt, @@ -5164,8 +5213,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { queryComponentKeysWithErrorsStmt: q.queryComponentKeysWithErrorsStmt, queryDeviceIDsWithErrorsStmt: q.queryDeviceIDsWithErrorsStmt, queryErrorsStmt: q.queryErrorsStmt, - reapStuckFirmwareUpdateMessagesStmt: q.reapStuckFirmwareUpdateMessagesStmt, - reapStuckProcessingMessagesStmt: q.reapStuckProcessingMessagesStmt, + reapMessagesStmt: q.reapMessagesStmt, reassignDevicesUnderBuildingStmt: q.reassignDevicesUnderBuildingStmt, reassignDevicesUnderBuildingsBulkStmt: q.reassignDevicesUnderBuildingsBulkStmt, reassignRacksUnderBuildingStmt: q.reassignRacksUnderBuildingStmt, @@ -5183,6 +5231,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { requeueRigConfigReconciliationAfterTerminalFailureStmt: q.requeueRigConfigReconciliationAfterTerminalFailureStmt, resetCurtailmentTargetsForRecurtailStmt: q.resetCurtailmentTargetsForRecurtailStmt, resetCurtailmentTargetsForRestoreStmt: q.resetCurtailmentTargetsForRestoreStmt, + resetReapedFirmwareStatusesStmt: q.resetReapedFirmwareStatusesStmt, resumeCurtailmentFromRestoringStmt: q.resumeCurtailmentFromRestoringStmt, resumePausedScheduleStmt: q.resumePausedScheduleStmt, retryRigConfigReconciliationStmt: q.retryRigConfigReconciliationStmt, @@ -5199,6 +5248,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { setDevicePairingAuthNeededIfNotPairedStmt: q.setDevicePairingAuthNeededIfNotPairedStmt, setFleetNodeEnrollmentStatusStmt: q.setFleetNodeEnrollmentStatusStmt, setInfrastructureControlSubnetsStmt: q.setInfrastructureControlSubnetsStmt, + setLocalTransactionTimeoutStmt: q.setLocalTransactionTimeoutStmt, setMQTTSourceConfigEnabledStmt: q.setMQTTSourceConfigEnabledStmt, setRackBuildingPositionStmt: q.setRackBuildingPositionStmt, setRackBuildingPositionBulkClearStmt: q.setRackBuildingPositionBulkClearStmt, diff --git a/server/generated/sqlc/querier.go b/server/generated/sqlc/querier.go index ae0d4414b..f24ee89a4 100644 --- a/server/generated/sqlc/querier.go +++ b/server/generated/sqlc/querier.go @@ -285,6 +285,7 @@ type Querier interface { // deactivated. Same liveness filters as above so a deactivated user // never inflates the count. CountOrgScopeSuperAdminsExcludingUser(ctx context.Context, arg CountOrgScopeSuperAdminsExcludingUserParams) (int64, error) + CountQueueMessagesByBatch(ctx context.Context, commandBatchLogUuid string) (int64, error) CountRacksBySite(ctx context.Context, arg CountRacksBySiteParams) (int64, error) // Total live racks currently assigned to a building (placed or // unplaced — membership, not grid occupancy). Used by @@ -316,6 +317,7 @@ type Querier interface { CreatePendingEnrollment(ctx context.Context, arg CreatePendingEnrollmentParams) (PendingEnrollment, error) CreatePool(ctx context.Context, arg CreatePoolParams) (int64, error) CreateQueueMessage(ctx context.Context, arg CreateQueueMessageParams) error + CreateQueueMessages(ctx context.Context, arg CreateQueueMessagesParams) error // org_id is denormalized onto device_set_rack so the building FK can be // composite-keyed; inherit it from device_set so the caller's org_id // must match. site_id / building_id are NULL for unassigned racks. @@ -421,6 +423,7 @@ type Querier interface { // miner with only a direct building (site NULL, building set, e.g. one // assigned to a site-less building) must trip the confirm too. FindDevicesWithSiteOrBuilding(ctx context.Context, arg FindDevicesWithSiteOrBuildingParams) ([]string, error) + FinishTerminalCommandBatches(ctx context.Context, finishLimit int32) (int64, error) // Last-resort recovery: persistently releases curtailment ownership for any // non-terminal event row. Unlike AdminTerminateCurtailmentEvent, this // intentionally supports ACTIVE events and has no in-flight command gate because @@ -700,6 +703,7 @@ type Querier interface { GetPermissionByKey(ctx context.Context, key string) (Permission, error) GetPermissionsByKeys(ctx context.Context, keys []string) ([]Permission, error) GetPool(ctx context.Context, arg GetPoolParams) (Pool, error) + GetQueueMessagesByBatch(ctx context.Context, commandBatchLogUuid string) ([]GetQueueMessagesByBatchRow, error) // Batch query to get rack label and formatted slot position for multiple devices at once. // Returns at most one rack per device due to partial unique index. GetRackDetailsForDevices(ctx context.Context, arg GetRackDetailsForDevicesParams) ([]GetRackDetailsForDevicesRow, error) @@ -803,7 +807,6 @@ type Querier interface { // the in-process metrics provider on every flush. InsertNotificationMetricSamples(ctx context.Context, arg InsertNotificationMetricSamplesParams) error IsBatchFinished(ctx context.Context, commandBatchLogUuid string) (bool, error) - IsBatchProcessing(ctx context.Context, commandBatchLogUuid string) (bool, error) IsDeviceOwnedByFleetNode(ctx context.Context, arg IsDeviceOwnedByFleetNodeParams) (bool, error) // Devices locked in a non-terminal event; excluded from candidates to // enforce the per-device single-writer rule. @@ -1125,6 +1128,7 @@ type Querier interface { // the locked ids (result is informational; the FOR UPDATE side-effect // is what matters). LockBuildingsBySiteForWrite(ctx context.Context, arg LockBuildingsBySiteForWriteParams) ([]int64, error) + LockCommandBatch(ctx context.Context, uuid string) (BatchStatusEnum, error) LockCurtailmentEventByUUIDForWrite(ctx context.Context, arg LockCurtailmentEventByUUIDForWriteParams) (CurtailmentEvent, error) // Physical fan commands run only while this exact lifecycle phase remains // current. Holding the row lock through the command serializes Force Release's @@ -1201,9 +1205,9 @@ type Querier interface { // between the existence check and the cascade write. Returns the // site id when alive; sql.ErrNoRows when soft-deleted or missing. LockSiteForWrite(ctx context.Context, arg LockSiteForWriteParams) (int64, error) - MarkCommandBatchFinished(ctx context.Context, uuid string) error - MarkCommandBatchFinishedWithStartedAt(ctx context.Context, uuid string) error - MarkCommandBatchProcessing(ctx context.Context, uuid string) error + MarkCommandBatchFinished(ctx context.Context, uuid string) (int64, error) + MarkCommandBatchFinishedWithStartedAt(ctx context.Context, uuid string) (int64, error) + MarkCommandBatchProcessing(ctx context.Context, uuid string) (int64, error) NegateSchedulePriorities(ctx context.Context, arg NegateSchedulePrioritiesParams) error PairDeviceToFleetNode(ctx context.Context, arg PairDeviceToFleetNodeParams) (int64, error) PasswordUpdatedAt(ctx context.Context, id int64) (sql.NullTime, error) @@ -1233,8 +1237,9 @@ type Querier interface { // Time range and include_closed are always applied as base filters. // Uses cursor-based pagination with (severity, last_seen_at, error_id) ordering. QueryErrors(ctx context.Context, arg QueryErrorsParams) ([]QueryErrorsRow, error) - ReapStuckFirmwareUpdateMessages(ctx context.Context, arg ReapStuckFirmwareUpdateMessagesParams) ([]ReapStuckFirmwareUpdateMessagesRow, error) - ReapStuckProcessingMessages(ctx context.Context, arg ReapStuckProcessingMessagesParams) ([]ReapStuckProcessingMessagesRow, error) + // Startup reaping fails every PROCESSING row left by the previous process. + // Periodic reaping only fails rows that exceeded their command-specific cutoff. + ReapMessages(ctx context.Context, arg ReapMessagesParams) ([]ReapMessagesRow, error) // Sets device.site_id = $target for every device in any live rack of // the given building. Caller wraps this in the same tx as the building // UPDATE. The JOIN on device_set with deleted_at IS NULL skips @@ -1318,6 +1323,7 @@ type Querier interface { // desired_state='active' and clears phase-local cursors so the restorer // has an unambiguous queue. Terminal states are untouched. ResetCurtailmentTargetsForRestore(ctx context.Context, curtailmentEventID int64) error + ResetReapedFirmwareStatuses(ctx context.Context, deviceIds []int64) error // Restore reversal: go back through pending so the curtail dispatcher picks // up reset targets. Preserve fan_off_sent_at and fan_last_error until the // active reconciler has positively reopened airflow; clearing them here can @@ -1345,6 +1351,7 @@ type Querier interface { // Explicitly replaces the commissioned OT allowlist. Empty text // decommissions the site. Canonicalization happens in the sites domain. SetInfrastructureControlSubnets(ctx context.Context, arg SetInfrastructureControlSubnetsParams) (string, error) + SetLocalTransactionTimeout(ctx context.Context, timeoutMilliseconds int64) error SetMQTTSourceConfigEnabled(ctx context.Context, arg SetMQTTSourceConfigEnabledParams) (SetMQTTSourceConfigEnabledRow, error) // Writes the rack's grid placement (aisle_index, position_in_aisle). // Caller must have already set building_id via UpdateRackPlacement — diff --git a/server/generated/sqlc/queue.sql.go b/server/generated/sqlc/queue.sql.go index 56d0c2756..3cfc7accf 100644 --- a/server/generated/sqlc/queue.sql.go +++ b/server/generated/sqlc/queue.sql.go @@ -10,6 +10,7 @@ import ( "database/sql" "time" + "github.com/lib/pq" "github.com/sqlc-dev/pqtype" ) @@ -25,6 +26,19 @@ func (q *Queries) ClaimMessageForProcessing(ctx context.Context, id int64) (sql. return q.exec(ctx, q.claimMessageForProcessingStmt, claimMessageForProcessing, id) } +const countQueueMessagesByBatch = `-- name: CountQueueMessagesByBatch :one +SELECT COUNT(*) +FROM queue_message +WHERE command_batch_log_uuid = $1 +` + +func (q *Queries) CountQueueMessagesByBatch(ctx context.Context, commandBatchLogUuid string) (int64, error) { + row := q.queryRow(ctx, q.countQueueMessagesByBatchStmt, countQueueMessagesByBatch, commandBatchLogUuid) + var count int64 + err := row.Scan(&count) + return count, err +} + const createQueueMessage = `-- name: CreateQueueMessage :exec INSERT INTO queue_message ( command_batch_log_uuid, @@ -64,6 +78,80 @@ func (q *Queries) CreateQueueMessage(ctx context.Context, arg CreateQueueMessage return err } +const createQueueMessages = `-- name: CreateQueueMessages :exec +INSERT INTO queue_message ( + command_batch_log_uuid, + command_type, + device_id, + status, + retry_count, + payload +) +SELECT + $1, + $2, + devices.device_id, + 'PENDING'::queue_status_enum, + 0, + payloads.payload::JSONB +FROM unnest($3::BIGINT[]) WITH ORDINALITY AS devices(device_id, ord) +JOIN unnest($4::TEXT[]) WITH ORDINALITY AS payloads(payload, ord) USING (ord) +` + +type CreateQueueMessagesParams struct { + CommandBatchLogUuid string + CommandType string + DeviceIds []int64 + Payloads []string +} + +func (q *Queries) CreateQueueMessages(ctx context.Context, arg CreateQueueMessagesParams) error { + _, err := q.exec(ctx, q.createQueueMessagesStmt, createQueueMessages, + arg.CommandBatchLogUuid, + arg.CommandType, + pq.Array(arg.DeviceIds), + pq.Array(arg.Payloads), + ) + return err +} + +const finishTerminalCommandBatches = `-- name: FinishTerminalCommandBatches :execrows +WITH candidates AS MATERIALIZED ( + SELECT batch.id + FROM command_batch_log AS batch + WHERE batch.status IN ('PENDING', 'PROCESSING') + AND EXISTS ( + SELECT 1 + FROM queue_message AS message + WHERE message.command_batch_log_uuid = batch.uuid + ) + AND NOT EXISTS ( + SELECT 1 + FROM queue_message AS message + WHERE message.command_batch_log_uuid = batch.uuid + AND message.status IN ('PENDING', 'PROCESSING') + ) + ORDER BY batch.id + LIMIT $1 + FOR UPDATE +) +UPDATE command_batch_log AS batch +SET + status = 'FINISHED'::batch_status_enum, + finished_at = CURRENT_TIMESTAMP +FROM candidates +WHERE batch.id = candidates.id + AND batch.status IN ('PENDING', 'PROCESSING') +` + +func (q *Queries) FinishTerminalCommandBatches(ctx context.Context, finishLimit int32) (int64, error) { + result, err := q.exec(ctx, q.finishTerminalCommandBatchesStmt, finishTerminalCommandBatches, finishLimit) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + const getMessagesToProcess = `-- name: GetMessagesToProcess :many SELECT m.id, m.command_batch_log_uuid, m.device_id, m.command_type, m.status, m.retry_count, m.error_info, m.payload, m.created_at, m.updated_at, @@ -137,92 +225,35 @@ func (q *Queries) GetMessagesToProcess(ctx context.Context, arg GetMessagesToPro return items, nil } -const isBatchFinished = `-- name: IsBatchFinished :one -SELECT - CASE - WHEN COUNT(*) = 0 THEN false - WHEN COUNT(*) = SUM(CASE WHEN status IN ('SUCCESS', 'FAILED') THEN 1 ELSE 0 END) THEN true - ELSE false - END AS is_finished -FROM queue_message -WHERE command_batch_log_uuid = $1 -` - -func (q *Queries) IsBatchFinished(ctx context.Context, commandBatchLogUuid string) (bool, error) { - row := q.queryRow(ctx, q.isBatchFinishedStmt, isBatchFinished, commandBatchLogUuid) - var is_finished bool - err := row.Scan(&is_finished) - return is_finished, err -} - -const isBatchProcessing = `-- name: IsBatchProcessing :one -SELECT - CASE - WHEN COUNT(*) > 0 THEN true - ELSE false - END AS is_processing +const getQueueMessagesByBatch = `-- name: GetQueueMessagesByBatch :many +SELECT id, device_id, status, error_info, payload FROM queue_message WHERE command_batch_log_uuid = $1 - AND status = 'PROCESSING' ` -func (q *Queries) IsBatchProcessing(ctx context.Context, commandBatchLogUuid string) (bool, error) { - row := q.queryRow(ctx, q.isBatchProcessingStmt, isBatchProcessing, commandBatchLogUuid) - var is_processing bool - err := row.Scan(&is_processing) - return is_processing, err -} - -const reapStuckFirmwareUpdateMessages = `-- name: ReapStuckFirmwareUpdateMessages :many -WITH stuck AS ( - SELECT m.id FROM queue_message m - WHERE m.status = 'PROCESSING' - AND m.updated_at < $1 - AND m.command_type = 'FirmwareUpdate' - LIMIT $2 -) -UPDATE queue_message -SET status = 'FAILED'::queue_status_enum, - error_info = 'reaped: firmware update stuck in PROCESSING beyond timeout', - updated_at = CURRENT_TIMESTAMP -FROM stuck, device -WHERE queue_message.id = stuck.id - AND queue_message.status = 'PROCESSING' - AND queue_message.device_id = device.id -RETURNING queue_message.id, queue_message.device_id, queue_message.command_batch_log_uuid, - queue_message.error_info, queue_message.command_type, device.org_id -` - -type ReapStuckFirmwareUpdateMessagesParams struct { - Cutoff time.Time - ReapLimit int32 -} - -type ReapStuckFirmwareUpdateMessagesRow struct { - ID int64 - DeviceID int64 - CommandBatchLogUuid string - ErrorInfo sql.NullString - CommandType string - OrgID int64 +type GetQueueMessagesByBatchRow struct { + ID int64 + DeviceID int64 + Status QueueStatusEnum + ErrorInfo sql.NullString + Payload pqtype.NullRawMessage } -func (q *Queries) ReapStuckFirmwareUpdateMessages(ctx context.Context, arg ReapStuckFirmwareUpdateMessagesParams) ([]ReapStuckFirmwareUpdateMessagesRow, error) { - rows, err := q.query(ctx, q.reapStuckFirmwareUpdateMessagesStmt, reapStuckFirmwareUpdateMessages, arg.Cutoff, arg.ReapLimit) +func (q *Queries) GetQueueMessagesByBatch(ctx context.Context, commandBatchLogUuid string) ([]GetQueueMessagesByBatchRow, error) { + rows, err := q.query(ctx, q.getQueueMessagesByBatchStmt, getQueueMessagesByBatch, commandBatchLogUuid) if err != nil { return nil, err } defer rows.Close() - var items []ReapStuckFirmwareUpdateMessagesRow + var items []GetQueueMessagesByBatchRow for rows.Next() { - var i ReapStuckFirmwareUpdateMessagesRow + var i GetQueueMessagesByBatchRow if err := rows.Scan( &i.ID, &i.DeviceID, - &i.CommandBatchLogUuid, + &i.Status, &i.ErrorInfo, - &i.CommandType, - &i.OrgID, + &i.Payload, ); err != nil { return nil, err } @@ -237,49 +268,101 @@ func (q *Queries) ReapStuckFirmwareUpdateMessages(ctx context.Context, arg ReapS return items, nil } -const reapStuckProcessingMessages = `-- name: ReapStuckProcessingMessages :many -WITH stuck AS ( - SELECT m.id FROM queue_message m - WHERE m.status = 'PROCESSING' - AND m.updated_at < $1 - AND m.command_type != 'FirmwareUpdate' - LIMIT $2 +const isBatchFinished = `-- name: IsBatchFinished :one +SELECT + CASE + WHEN COUNT(*) = 0 THEN false + WHEN COUNT(*) = SUM(CASE WHEN status IN ('SUCCESS', 'FAILED') THEN 1 ELSE 0 END) THEN true + ELSE false + END AS is_finished +FROM queue_message +WHERE command_batch_log_uuid = $1 +` + +func (q *Queries) IsBatchFinished(ctx context.Context, commandBatchLogUuid string) (bool, error) { + row := q.queryRow(ctx, q.isBatchFinishedStmt, isBatchFinished, commandBatchLogUuid) + var is_finished bool + err := row.Scan(&is_finished) + return is_finished, err +} + +const reapMessages = `-- name: ReapMessages :many +WITH candidates AS ( + SELECT message.id + FROM queue_message AS message + WHERE message.status = 'PROCESSING' + AND ( + $3::BOOLEAN + OR message.updated_at < CASE + WHEN message.command_type = 'FirmwareUpdate' + THEN $4::TIMESTAMPTZ + ELSE $5::TIMESTAMPTZ + END + ) + ORDER BY message.updated_at, message.id + LIMIT $6 + FOR UPDATE ) -UPDATE queue_message -SET status = 'FAILED'::queue_status_enum, - error_info = 'reaped: stuck in PROCESSING beyond timeout', +UPDATE queue_message AS message +SET + status = 'FAILED'::queue_status_enum, + error_info = CASE + WHEN message.command_type = 'FirmwareUpdate' + THEN $1::TEXT + ELSE $2::TEXT + END, updated_at = CURRENT_TIMESTAMP -FROM stuck, device -WHERE queue_message.id = stuck.id - AND queue_message.status = 'PROCESSING' - AND queue_message.device_id = device.id -RETURNING queue_message.id, queue_message.device_id, queue_message.command_batch_log_uuid, - queue_message.error_info, queue_message.command_type, device.org_id +FROM candidates, device +WHERE message.id = candidates.id + AND message.status = 'PROCESSING' + AND message.device_id = device.id +RETURNING + message.id, + message.device_id, + message.command_batch_log_uuid, + message.error_info, + message.command_type, + device.org_id, + device.site_id ` -type ReapStuckProcessingMessagesParams struct { - Cutoff time.Time - ReapLimit int32 +type ReapMessagesParams struct { + FirmwareErrorInfo string + ErrorInfo string + IncludeFresh bool + FirmwareCutoff time.Time + Cutoff time.Time + ReapLimit int32 } -type ReapStuckProcessingMessagesRow struct { +type ReapMessagesRow struct { ID int64 DeviceID int64 CommandBatchLogUuid string ErrorInfo sql.NullString CommandType string OrgID int64 + SiteID sql.NullInt64 } -func (q *Queries) ReapStuckProcessingMessages(ctx context.Context, arg ReapStuckProcessingMessagesParams) ([]ReapStuckProcessingMessagesRow, error) { - rows, err := q.query(ctx, q.reapStuckProcessingMessagesStmt, reapStuckProcessingMessages, arg.Cutoff, arg.ReapLimit) +// Startup reaping fails every PROCESSING row left by the previous process. +// Periodic reaping only fails rows that exceeded their command-specific cutoff. +func (q *Queries) ReapMessages(ctx context.Context, arg ReapMessagesParams) ([]ReapMessagesRow, error) { + rows, err := q.query(ctx, q.reapMessagesStmt, reapMessages, + arg.FirmwareErrorInfo, + arg.ErrorInfo, + arg.IncludeFresh, + arg.FirmwareCutoff, + arg.Cutoff, + arg.ReapLimit, + ) if err != nil { return nil, err } defer rows.Close() - var items []ReapStuckProcessingMessagesRow + var items []ReapMessagesRow for rows.Next() { - var i ReapStuckProcessingMessagesRow + var i ReapMessagesRow if err := rows.Scan( &i.ID, &i.DeviceID, @@ -287,6 +370,7 @@ func (q *Queries) ReapStuckProcessingMessages(ctx context.Context, arg ReapStuck &i.ErrorInfo, &i.CommandType, &i.OrgID, + &i.SiteID, ); err != nil { return nil, err } @@ -301,6 +385,34 @@ func (q *Queries) ReapStuckProcessingMessages(ctx context.Context, arg ReapStuck return items, nil } +const resetReapedFirmwareStatuses = `-- name: ResetReapedFirmwareStatuses :exec +UPDATE device_status +SET + status = 'ACTIVE'::device_status_enum, + status_timestamp = CURRENT_TIMESTAMP, + status_details = NULL +WHERE device_id = ANY($1::BIGINT[]) + AND status IN ('UPDATING', 'REBOOT_REQUIRED') +` + +func (q *Queries) ResetReapedFirmwareStatuses(ctx context.Context, deviceIds []int64) error { + _, err := q.exec(ctx, q.resetReapedFirmwareStatusesStmt, resetReapedFirmwareStatuses, pq.Array(deviceIds)) + return err +} + +const setLocalTransactionTimeout = `-- name: SetLocalTransactionTimeout :exec +SELECT set_config( + 'transaction_timeout', + $1::BIGINT::TEXT || 'ms', + TRUE +) +` + +func (q *Queries) SetLocalTransactionTimeout(ctx context.Context, timeoutMilliseconds int64) error { + _, err := q.exec(ctx, q.setLocalTransactionTimeoutStmt, setLocalTransactionTimeout, timeoutMilliseconds) + return err +} + const updateMessageAfterFailure = `-- name: UpdateMessageAfterFailure :execresult UPDATE queue_message SET status = CASE diff --git a/server/generated/sqlc/retrying_querier.gen.go b/server/generated/sqlc/retrying_querier.gen.go index 2f35d6781..b1fa836bd 100644 --- a/server/generated/sqlc/retrying_querier.gen.go +++ b/server/generated/sqlc/retrying_querier.gen.go @@ -750,6 +750,18 @@ func (q *retryingQuerier) CountOrgScopeSuperAdminsExcludingUser(ctx context.Cont return result, err } +func (q *retryingQuerier) CountQueueMessagesByBatch(ctx context.Context, commandBatchLogUuid string) (int64, error) { + var result int64 + err := q.retrier.RetryQuery(ctx, "CountQueueMessagesByBatch", func() error { + callResult, callErr := q.next.CountQueueMessagesByBatch(ctx, commandBatchLogUuid) + if callErr == nil { + result = callResult + } + return callErr + }) + return result, err +} + func (q *retryingQuerier) CountRacksBySite(ctx context.Context, arg CountRacksBySiteParams) (int64, error) { var result int64 err := q.retrier.RetryQuery(ctx, "CountRacksBySite", func() error { @@ -924,6 +936,12 @@ func (q *retryingQuerier) CreateQueueMessage(ctx context.Context, arg CreateQueu }) } +func (q *retryingQuerier) CreateQueueMessages(ctx context.Context, arg CreateQueueMessagesParams) error { + return q.retrier.RetryQuery(ctx, "CreateQueueMessages", func() error { + return q.next.CreateQueueMessages(ctx, arg) + }) +} + func (q *retryingQuerier) CreateRackExtension(ctx context.Context, arg CreateRackExtensionParams) error { return q.retrier.RetryQuery(ctx, "CreateRackExtension", func() error { return q.next.CreateRackExtension(ctx, arg) @@ -1284,6 +1302,18 @@ func (q *retryingQuerier) FindDevicesWithSiteOrBuilding(ctx context.Context, arg return result, err } +func (q *retryingQuerier) FinishTerminalCommandBatches(ctx context.Context, finishLimit int32) (int64, error) { + var result int64 + err := q.retrier.RetryQuery(ctx, "FinishTerminalCommandBatches", func() error { + callResult, callErr := q.next.FinishTerminalCommandBatches(ctx, finishLimit) + if callErr == nil { + result = callResult + } + return callErr + }) + return result, err +} + func (q *retryingQuerier) ForceReleaseCurtailmentEvent(ctx context.Context, arg ForceReleaseCurtailmentEventParams) (CurtailmentEvent, error) { var result CurtailmentEvent err := q.retrier.RetryQuery(ctx, "ForceReleaseCurtailmentEvent", func() error { @@ -2796,6 +2826,18 @@ func (q *retryingQuerier) GetPool(ctx context.Context, arg GetPoolParams) (Pool, return result, err } +func (q *retryingQuerier) GetQueueMessagesByBatch(ctx context.Context, commandBatchLogUuid string) ([]GetQueueMessagesByBatchRow, error) { + var result []GetQueueMessagesByBatchRow + err := q.retrier.RetryQuery(ctx, "GetQueueMessagesByBatch", func() error { + callResult, callErr := q.next.GetQueueMessagesByBatch(ctx, commandBatchLogUuid) + if callErr == nil { + result = callResult + } + return callErr + }) + return result, err +} + func (q *retryingQuerier) GetRackDetailsForDevices(ctx context.Context, arg GetRackDetailsForDevicesParams) ([]GetRackDetailsForDevicesRow, error) { var result []GetRackDetailsForDevicesRow err := q.retrier.RetryQuery(ctx, "GetRackDetailsForDevices", func() error { @@ -3264,18 +3306,6 @@ func (q *retryingQuerier) IsBatchFinished(ctx context.Context, commandBatchLogUu return result, err } -func (q *retryingQuerier) IsBatchProcessing(ctx context.Context, commandBatchLogUuid string) (bool, error) { - var result bool - err := q.retrier.RetryQuery(ctx, "IsBatchProcessing", func() error { - callResult, callErr := q.next.IsBatchProcessing(ctx, commandBatchLogUuid) - if callErr == nil { - result = callResult - } - return callErr - }) - return result, err -} - func (q *retryingQuerier) IsDeviceOwnedByFleetNode(ctx context.Context, arg IsDeviceOwnedByFleetNodeParams) (bool, error) { var result bool err := q.retrier.RetryQuery(ctx, "IsDeviceOwnedByFleetNode", func() error { @@ -4128,6 +4158,18 @@ func (q *retryingQuerier) LockBuildingsBySiteForWrite(ctx context.Context, arg L return result, err } +func (q *retryingQuerier) LockCommandBatch(ctx context.Context, uuid string) (BatchStatusEnum, error) { + var result BatchStatusEnum + err := q.retrier.RetryQuery(ctx, "LockCommandBatch", func() error { + callResult, callErr := q.next.LockCommandBatch(ctx, uuid) + if callErr == nil { + result = callResult + } + return callErr + }) + return result, err +} + func (q *retryingQuerier) LockCurtailmentEventByUUIDForWrite(ctx context.Context, arg LockCurtailmentEventByUUIDForWriteParams) (CurtailmentEvent, error) { var result CurtailmentEvent err := q.retrier.RetryQuery(ctx, "LockCurtailmentEventByUUIDForWrite", func() error { @@ -4296,22 +4338,40 @@ func (q *retryingQuerier) LockSiteForWrite(ctx context.Context, arg LockSiteForW return result, err } -func (q *retryingQuerier) MarkCommandBatchFinished(ctx context.Context, uuid string) error { - return q.retrier.RetryQuery(ctx, "MarkCommandBatchFinished", func() error { - return q.next.MarkCommandBatchFinished(ctx, uuid) +func (q *retryingQuerier) MarkCommandBatchFinished(ctx context.Context, uuid string) (int64, error) { + var result int64 + err := q.retrier.RetryQuery(ctx, "MarkCommandBatchFinished", func() error { + callResult, callErr := q.next.MarkCommandBatchFinished(ctx, uuid) + if callErr == nil { + result = callResult + } + return callErr }) + return result, err } -func (q *retryingQuerier) MarkCommandBatchFinishedWithStartedAt(ctx context.Context, uuid string) error { - return q.retrier.RetryQuery(ctx, "MarkCommandBatchFinishedWithStartedAt", func() error { - return q.next.MarkCommandBatchFinishedWithStartedAt(ctx, uuid) +func (q *retryingQuerier) MarkCommandBatchFinishedWithStartedAt(ctx context.Context, uuid string) (int64, error) { + var result int64 + err := q.retrier.RetryQuery(ctx, "MarkCommandBatchFinishedWithStartedAt", func() error { + callResult, callErr := q.next.MarkCommandBatchFinishedWithStartedAt(ctx, uuid) + if callErr == nil { + result = callResult + } + return callErr }) + return result, err } -func (q *retryingQuerier) MarkCommandBatchProcessing(ctx context.Context, uuid string) error { - return q.retrier.RetryQuery(ctx, "MarkCommandBatchProcessing", func() error { - return q.next.MarkCommandBatchProcessing(ctx, uuid) +func (q *retryingQuerier) MarkCommandBatchProcessing(ctx context.Context, uuid string) (int64, error) { + var result int64 + err := q.retrier.RetryQuery(ctx, "MarkCommandBatchProcessing", func() error { + callResult, callErr := q.next.MarkCommandBatchProcessing(ctx, uuid) + if callErr == nil { + result = callResult + } + return callErr }) + return result, err } func (q *retryingQuerier) NegateSchedulePriorities(ctx context.Context, arg NegateSchedulePrioritiesParams) error { @@ -4398,22 +4458,10 @@ func (q *retryingQuerier) QueryErrors(ctx context.Context, arg QueryErrorsParams return result, err } -func (q *retryingQuerier) ReapStuckFirmwareUpdateMessages(ctx context.Context, arg ReapStuckFirmwareUpdateMessagesParams) ([]ReapStuckFirmwareUpdateMessagesRow, error) { - var result []ReapStuckFirmwareUpdateMessagesRow - err := q.retrier.RetryQuery(ctx, "ReapStuckFirmwareUpdateMessages", func() error { - callResult, callErr := q.next.ReapStuckFirmwareUpdateMessages(ctx, arg) - if callErr == nil { - result = callResult - } - return callErr - }) - return result, err -} - -func (q *retryingQuerier) ReapStuckProcessingMessages(ctx context.Context, arg ReapStuckProcessingMessagesParams) ([]ReapStuckProcessingMessagesRow, error) { - var result []ReapStuckProcessingMessagesRow - err := q.retrier.RetryQuery(ctx, "ReapStuckProcessingMessages", func() error { - callResult, callErr := q.next.ReapStuckProcessingMessages(ctx, arg) +func (q *retryingQuerier) ReapMessages(ctx context.Context, arg ReapMessagesParams) ([]ReapMessagesRow, error) { + var result []ReapMessagesRow + err := q.retrier.RetryQuery(ctx, "ReapMessages", func() error { + callResult, callErr := q.next.ReapMessages(ctx, arg) if callErr == nil { result = callResult } @@ -4608,6 +4656,12 @@ func (q *retryingQuerier) ResetCurtailmentTargetsForRestore(ctx context.Context, }) } +func (q *retryingQuerier) ResetReapedFirmwareStatuses(ctx context.Context, deviceIds []int64) error { + return q.retrier.RetryQuery(ctx, "ResetReapedFirmwareStatuses", func() error { + return q.next.ResetReapedFirmwareStatuses(ctx, deviceIds) + }) +} + func (q *retryingQuerier) ResumeCurtailmentFromRestoring(ctx context.Context, id int64) (CurtailmentEvent, error) { var result CurtailmentEvent err := q.retrier.RetryQuery(ctx, "ResumeCurtailmentFromRestoring", func() error { @@ -4758,6 +4812,12 @@ func (q *retryingQuerier) SetInfrastructureControlSubnets(ctx context.Context, a return result, err } +func (q *retryingQuerier) SetLocalTransactionTimeout(ctx context.Context, timeoutMilliseconds int64) error { + return q.retrier.RetryQuery(ctx, "SetLocalTransactionTimeout", func() error { + return q.next.SetLocalTransactionTimeout(ctx, timeoutMilliseconds) + }) +} + func (q *retryingQuerier) SetMQTTSourceConfigEnabled(ctx context.Context, arg SetMQTTSourceConfigEnabledParams) (SetMQTTSourceConfigEnabledRow, error) { var result SetMQTTSourceConfigEnabledRow err := q.retrier.RetryQuery(ctx, "SetMQTTSourceConfigEnabled", func() error { diff --git a/server/internal/domain/command/enqueue_failure_integration_test.go b/server/internal/domain/command/enqueue_failure_integration_test.go new file mode 100644 index 000000000..8e48f743a --- /dev/null +++ b/server/internal/domain/command/enqueue_failure_integration_test.go @@ -0,0 +1,138 @@ +package command_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + commonpb "github.com/block/proto-fleet/server/generated/grpc/common/v1" + pb "github.com/block/proto-fleet/server/generated/grpc/minercommand/v1" + "github.com/block/proto-fleet/server/generated/sqlc" + "github.com/block/proto-fleet/server/internal/domain/commandtype" + "github.com/block/proto-fleet/server/internal/infrastructure/queue" + "github.com/block/proto-fleet/server/internal/testutil" +) + +type failingEnqueueQueue struct { + err error + delegate queue.MessageQueue + batchUUID *string + statusChecks chan string +} + +func (q failingEnqueueQueue) Enqueue(ctx context.Context, batchUUID string, commandType commandtype.Type, deviceIDs []int64, payload interface{}) error { + if q.batchUUID != nil { + *q.batchUUID = batchUUID + } + if q.delegate != nil { + if err := q.delegate.Enqueue(ctx, batchUUID, commandType, deviceIDs, payload); err != nil { + return err + } + } + return q.err +} + +func (q failingEnqueueQueue) EnqueueMany(ctx context.Context, batchUUID string, commandType commandtype.Type, messages []queue.EnqueueMessage) error { + if q.batchUUID != nil { + *q.batchUUID = batchUUID + } + if q.delegate != nil { + if err := q.delegate.EnqueueMany(ctx, batchUUID, commandType, messages); err != nil { + return err + } + } + return q.err +} + +func (failingEnqueueQueue) Dequeue(context.Context, int32) ([]queue.Message, error) { + return nil, nil +} + +func (q failingEnqueueQueue) IsBatchFinished(_ context.Context, batchUUID string) (bool, error) { + if q.statusChecks != nil { + select { + case q.statusChecks <- batchUUID: + default: + } + } + return false, nil +} + +func (failingEnqueueQueue) MaxFailureRetries() int32 { + return 0 +} + +func TestCommandEnqueueFailureFinishesCreatedBatch(t *testing.T) { + if testing.Short() { + t.Skip("Skipping database integration test in short mode") + } + + // Arrange + conn, dbService, user := setupRetentionTest(t) + device := dbService.CreateDevice(user.OrganizationID, "proto") + enqueueErr := errors.New("queue unavailable") + var batchUUID string + svc := newDispatchIntegrationTestService(t, conn, failingEnqueueQueue{err: enqueueErr, batchUUID: &batchUUID}) + ctx := testutil.MockAuthContextForTesting(t.Context(), user.DatabaseID, user.OrganizationID) + + // Act + result, err := svc.BlinkLED(ctx, &pb.DeviceSelector{ + SelectionType: &pb.DeviceSelector_IncludeDevices{ + IncludeDevices: &commonpb.DeviceIdentifierList{DeviceIdentifiers: []string{device.ID}}, + }, + }) + + // Assert + require.Error(t, err) + assert.Nil(t, result) + assert.ErrorContains(t, err, enqueueErr.Error()) + queries := sqlc.New(conn) + batch, err := queries.GetBatchLog(t.Context(), batchUUID) + require.NoError(t, err) + assert.Equal(t, sqlc.BatchStatusEnumFINISHED, batch.Status) + messages, err := queries.GetQueueMessagesByBatch(t.Context(), batchUUID) + require.NoError(t, err) + assert.Empty(t, messages) +} + +func TestCommandEnqueueCommittedBeforeErrorReturnsSuccessAndTracksBatch(t *testing.T) { + if testing.Short() { + t.Skip("Skipping database integration test in short mode") + } + + // Arrange + conn, dbService, user := setupRetentionTest(t) + device := dbService.CreateDevice(user.OrganizationID, "proto") + statusChecks := make(chan string, 1) + messageQueue := failingEnqueueQueue{ + err: errors.New("connection lost after commit"), + delegate: queue.NewDatabaseMessageQueue(&queue.Config{}, conn), + statusChecks: statusChecks, + } + svc := newDispatchIntegrationTestService(t, conn, messageQueue) + ctx := testutil.MockAuthContextForTesting(t.Context(), user.DatabaseID, user.OrganizationID) + + // Act + result, err := svc.BlinkLED(ctx, &pb.DeviceSelector{ + SelectionType: &pb.DeviceSelector_IncludeDevices{ + IncludeDevices: &commonpb.DeviceIdentifierList{DeviceIdentifiers: []string{device.ID}}, + }, + }) + + // Assert + require.NoError(t, err) + require.NotNil(t, result) + messages, err := sqlc.New(conn).GetQueueMessagesByBatch(t.Context(), result.BatchIdentifier) + require.NoError(t, err) + require.Len(t, messages, 1) + select { + case trackedBatch := <-statusChecks: + assert.Equal(t, result.BatchIdentifier, trackedBatch) + case <-time.After(time.Second): + t.Fatal("batch status tracking did not start") + } +} diff --git a/server/internal/domain/command/execution_service.go b/server/internal/domain/command/execution_service.go index 9e40f1759..7f85e4c44 100644 --- a/server/internal/domain/command/execution_service.go +++ b/server/internal/domain/command/execution_service.go @@ -32,6 +32,7 @@ import ( "github.com/block/proto-fleet/server/internal/infrastructure/files" "github.com/block/proto-fleet/server/internal/infrastructure/queue" "github.com/block/proto-fleet/server/internal/runtimejobs" + "github.com/block/proto-fleet/server/internal/runtimepolicy" ) const ( @@ -153,6 +154,10 @@ func (es *ExecutionService) Start(ctx context.Context) error { es.lifecycleMu.Unlock() return nil } + if err := es.reapAfterRestart(ctx); err != nil { + es.lifecycleMu.Unlock() + return err + } run := newExecutionRun(ctx) es.run = run @@ -228,8 +233,7 @@ func (es *ExecutionService) IsRunning() bool { return es.run != nil && es.run.accepting && es.run.admissionCtx.Err() == nil } -// Abort cancels both new and already-admitted command work. HA demotion uses -// this before graceful group cleanup so commands cannot outlive ownership. +// Abort immediately cancels admitted work before a fatal HA process exit. func (es *ExecutionService) Abort() { es.lifecycleMu.Lock() run := es.run @@ -292,7 +296,7 @@ func (es *ExecutionService) startStuckMessageReaper(ctx context.Context) { continue } reapCtx, reapCancel := context.WithTimeout(ctx, dbWriteTimeout) - reaped, fwDeviceIDs, err := es.reapStuckMessages(reapCtx) + reaped, err := es.reapMessages(reapCtx, reapModeStuck) reapCancel() if err != nil { slog.Error("stuck message reaper error", "error", err) @@ -303,8 +307,8 @@ func (es *ExecutionService) startStuckMessageReaper(ctx context.Context) { slog.Warn("stuck message reaper moved messages to FAILED", "count", len(reaped)) } es.emitReapedCommandMetrics(ctx, reaped) - for _, deviceID := range fwDeviceIDs { - es.clearFirmwareUpdateStatus(ctx, deviceID) + if _, err := es.finishTerminalCommandBatches(ctx); err != nil { + slog.Error("finish terminal command batches", "error", err) } reportProgress() } @@ -313,37 +317,86 @@ func (es *ExecutionService) startStuckMessageReaper(ctx context.Context) { type reapedCommand struct { orgID int64 + siteID int64 commandType string } -// reapStuckMessages atomically marks stuck PROCESSING messages as FAILED and -// writes the corresponding audit log entries in a single transaction. -// Firmware update messages use a longer cutoff since they include install polling. -// Returns the reaped commands' metric metadata and the device IDs from reaped -// firmware update messages (so callers can clean up stuck device statuses). -func (es *ExecutionService) reapStuckMessages(ctx context.Context) ([]reapedCommand, []int64, error) { - cutoff := time.Now().Add(-es.config.StuckMessageTimeout) - var reapedCmds []reapedCommand - var fwDeviceIDs []int64 - err := db.WithTransactionNoResult(ctx, es.conn, func(q sqlc.Querier) error { - reaped, err := q.ReapStuckProcessingMessages(ctx, sqlc.ReapStuckProcessingMessagesParams{ - Cutoff: cutoff, - ReapLimit: 100, - }) +type reapMode uint8 + +const ( + reapModeStuck reapMode = iota + reapModeRestart +) + +const ( + reaperBatchSize = 100 + stuckMessageReason = "reaped: stuck in PROCESSING beyond timeout" + stuckFirmwareReason = "reaped: firmware update stuck in PROCESSING beyond timeout" + commandRestartReason = "Interrupted by Fleet restart; device outcome may be unknown" +) + +// reapAfterRestart fails PROCESSING work whose device outcome is unknown. +// PENDING work remains queued for this process to execute. +func (es *ExecutionService) reapAfterRestart(ctx context.Context) error { + if es.conn == nil { + return nil + } + + for { + reaped, err := es.reapMessages(ctx, reapModeRestart) + if err != nil { + return fmt.Errorf("reap commands after restart: %w", err) + } + es.emitReapedCommandMetrics(ctx, reaped) + if len(reaped) < reaperBatchSize { + break + } + } + for { + finished, err := es.finishTerminalCommandBatches(ctx) if err != nil { return err } + if finished < reaperBatchSize { + return nil + } + } +} - fwCutoff := time.Now().Add(-es.config.FirmwareUpdateStuckTimeout) - fwReaped, err := q.ReapStuckFirmwareUpdateMessages(ctx, sqlc.ReapStuckFirmwareUpdateMessagesParams{ - Cutoff: fwCutoff, - ReapLimit: 100, +func (es *ExecutionService) finishTerminalCommandBatches(ctx context.Context) (int64, error) { + return db.WithTransactionTimeout(ctx, es.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) (int64, error) { + return q.FinishTerminalCommandBatches(ctx, reaperBatchSize) + }) +} + +// reapMessages marks one bounded batch FAILED, writes its audit rows, and +// clears firmware-owned status. +func (es *ExecutionService) reapMessages(ctx context.Context, mode reapMode) ([]reapedCommand, error) { + now := time.Now() + errorInfo := stuckMessageReason + firmwareErrorInfo := stuckFirmwareReason + if mode == reapModeRestart { + errorInfo = commandRestartReason + firmwareErrorInfo = commandRestartReason + } + + var reapedCmds []reapedCommand + err := db.WithTransactionTimeoutNoResult(ctx, es.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) error { + reaped, err := q.ReapMessages(ctx, sqlc.ReapMessagesParams{ + IncludeFresh: mode == reapModeRestart, + Cutoff: now.Add(-es.config.StuckMessageTimeout), + FirmwareCutoff: now.Add(-es.config.FirmwareUpdateStuckTimeout), + ReapLimit: reaperBatchSize, + ErrorInfo: errorInfo, + FirmwareErrorInfo: firmwareErrorInfo, }) if err != nil { return err } - reapedCmds = make([]reapedCommand, 0, len(reaped)+len(fwReaped)) + reapedCmds = make([]reapedCommand, 0, len(reaped)) + firmwareDeviceIDs := make([]int64, 0) + requeueOrganizations := make(map[int64]struct{}) for _, msg := range reaped { if err := q.UpsertCommandOnDeviceLog(ctx, sqlc.UpsertCommandOnDeviceLogParams{ Uuid: msg.CommandBatchLogUuid, @@ -356,31 +409,37 @@ func (es *ExecutionService) reapStuckMessages(ctx context.Context) ([]reapedComm } kind, kindErr := commandtype.FromString(msg.CommandType) if kindErr == nil && kind == commandtype.ApplyCurtailmentConfig { - if err := q.RequeueRigConfigReconciliationAfterTerminalFailure(ctx, msg.OrgID); err != nil { - return fmt.Errorf("requeue reaped rig config reconciliation: %w", err) - } + requeueOrganizations[msg.OrgID] = struct{}{} + } + if kindErr == nil && kind == commandtype.FirmwareUpdate { + firmwareDeviceIDs = append(firmwareDeviceIDs, msg.DeviceID) + } + siteID := int64(0) + if msg.SiteID.Valid { + siteID = msg.SiteID.Int64 } - reapedCmds = append(reapedCmds, reapedCommand{orgID: msg.OrgID, commandType: msg.CommandType}) + reapedCmds = append(reapedCmds, reapedCommand{ + orgID: msg.OrgID, + siteID: siteID, + commandType: msg.CommandType, + }) } - for _, msg := range fwReaped { - if err := q.UpsertCommandOnDeviceLog(ctx, sqlc.UpsertCommandOnDeviceLogParams{ - Uuid: msg.CommandBatchLogUuid, - DeviceID: msg.DeviceID, - Status: sqlc.DeviceCommandStatusEnumFAILED, - UpdatedAt: time.Now(), - ErrorInfo: msg.ErrorInfo, - }); err != nil { - return err + for orgID := range requeueOrganizations { + if err := q.RequeueRigConfigReconciliationAfterTerminalFailure(ctx, orgID); err != nil { + return fmt.Errorf("requeue reaped rig config reconciliation: %w", err) + } + } + if len(firmwareDeviceIDs) > 0 { + if err := q.ResetReapedFirmwareStatuses(ctx, firmwareDeviceIDs); err != nil { + return fmt.Errorf("reset reaped firmware statuses: %w", err) } - reapedCmds = append(reapedCmds, reapedCommand{orgID: msg.OrgID, commandType: msg.CommandType}) - fwDeviceIDs = append(fwDeviceIDs, msg.DeviceID) } return nil }) - return reapedCmds, fwDeviceIDs, err + return reapedCmds, err } -var errReapedStuck = errors.New("reaped: stuck in PROCESSING beyond timeout") +var errReapedCommand = errors.New("reaped command") // records a result="failure" sample for each reaped command. func (es *ExecutionService) emitReapedCommandMetrics(ctx context.Context, reaped []reapedCommand) { @@ -394,7 +453,7 @@ func (es *ExecutionService) emitReapedCommandMetrics(ctx context.Context, reaped "command_type", r.commandType, "error", err) continue } - emitTerminalCommand(ctx, es.metricsEmitter, r.orgID, 0, kind, errReapedStuck) + emitTerminalCommand(ctx, es.metricsEmitter, r.orgID, r.siteID, kind, errReapedCommand) } } @@ -539,7 +598,7 @@ func (es *ExecutionService) workerProcessCommand(ctx context.Context, message qu queueUpdated bool queueTerminal bool ) - txErr := db.WithTransactionNoResult(dbCtx, es.conn, func(q sqlc.Querier) error { + txErr := db.WithTransactionTimeoutNoResult(dbCtx, es.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) error { // First: transition queue_message status (detects staleness via rowsAffected). updated, terminal, err := es.markQueueMessageStatus(dbCtx, q, message, workerError) if err != nil { @@ -552,6 +611,9 @@ func (es *ExecutionService) workerProcessCommand(ctx context.Context, message qu "message_id", message.ID, "device_id", message.DeviceID) return nil } + if !terminal { + return nil + } // Second: write device log only if the queue transition succeeded. // Persist a sanitized reason so the activity-log detail RPC can surface @@ -634,7 +696,10 @@ func (es *ExecutionService) markQueueMessageStatus(ctx context.Context, q sqlc.Q if err != nil { return false, false, fleeterror.NewInternalErrorf("failed to update queue message status: %v", err) } - rowsAffected, _ := result.RowsAffected() + rowsAffected, err := result.RowsAffected() + if err != nil { + return false, false, fleeterror.NewInternalErrorf("failed to read queue message transition result: %v", err) + } return rowsAffected > 0, terminal, nil } @@ -753,6 +818,10 @@ func (es *ExecutionService) executeCommandOnDevice(ctx context.Context, commandT break } if shouldReboot { + if ctx.Err() != nil { + err = ctx.Err() + break + } err = es.rebootAfterFirmwareInstall(ctx, minerInfo, message.DeviceID) } case commandtype.Unpair: @@ -1286,7 +1355,11 @@ func (es *ExecutionService) clearFirmwareUpdateStatusForDevice(ctx context.Conte // already verified reboot support for firmware updates. For polling-capable // devices, status transitions to UPDATING while installation runs, then // REBOOT_REQUIRED on success. -func (es *ExecutionService) pollFirmwareInstallStatus(ctx context.Context, minerInfo interfaces.Miner, deviceID int64) (bool, error) { +func (es *ExecutionService) pollFirmwareInstallStatus( + ctx context.Context, + minerInfo interfaces.Miner, + deviceID int64, +) (bool, error) { provider, canPoll := minerInfo.(interfaces.FirmwareUpdateStatusProvider) if !canPoll { slog.Info("firmware update status provider unavailable, rebooting after upload", "device_id", deviceID) diff --git a/server/internal/domain/command/reaper_integration_test.go b/server/internal/domain/command/reaper_integration_test.go index 8d6d4063f..ccc7fbd78 100644 --- a/server/internal/domain/command/reaper_integration_test.go +++ b/server/internal/domain/command/reaper_integration_test.go @@ -27,19 +27,28 @@ func setupReaperTest(t *testing.T) (*sql.DB, *testutil.DatabaseService, *testuti return dbService.DB, dbService, user } -func createBatchLog(t *testing.T, conn *sql.DB, batchUUID string, userID int64, deviceCount int32) { +func createBatchLog(t *testing.T, conn *sql.DB, batchUUID string, userID int64, deviceCount int32, status sqlc.BatchStatusEnum) { + t.Helper() + createBatchLogs(t, conn, []string{batchUUID}, userID, deviceCount, status) +} + +func createBatchLogs(t *testing.T, conn *sql.DB, batchUUIDs []string, userID int64, deviceCount int32, status sqlc.BatchStatusEnum) { t.Helper() err := db2.WithTransactionNoResult(context.Background(), conn, func(q sqlc.Querier) error { - _, err := q.CreateCommandBatchLog(context.Background(), sqlc.CreateCommandBatchLogParams{ - Uuid: batchUUID, - Type: "REBOOT", - CreatedBy: userID, - CreatedAt: time.Now(), - Status: sqlc.BatchStatusEnumPROCESSING, - DevicesCount: deviceCount, - Payload: pqtype.NullRawMessage{Valid: false}, - }) - return err + for _, batchUUID := range batchUUIDs { + if _, err := q.CreateCommandBatchLog(context.Background(), sqlc.CreateCommandBatchLogParams{ + Uuid: batchUUID, + Type: "REBOOT", + CreatedBy: userID, + CreatedAt: time.Now(), + Status: status, + DevicesCount: deviceCount, + Payload: pqtype.NullRawMessage{Valid: false}, + }); err != nil { + return err + } + } + return nil }) require.NoError(t, err) } @@ -126,17 +135,9 @@ func (n *noopMessageQueue) Dequeue(ctx context.Context, _ int32) ([]queue.Messag <-ctx.Done() return nil, fmt.Errorf("dequeue cancelled: %w", ctx.Err()) } -func (n *noopMessageQueue) MarkSuccess(_ context.Context, _ int64) error { return nil } -func (n *noopMessageQueue) MarkFailed(_ context.Context, _ int64, _ string) error { return nil } -func (n *noopMessageQueue) MarkPermanentlyFailed(_ context.Context, _ int64, _ string) error { - return nil -} func (n *noopMessageQueue) IsBatchFinished(_ context.Context, _ string) (bool, error) { return false, nil } -func (n *noopMessageQueue) IsBatchProcessing(_ context.Context, _ string) (bool, error) { - return false, nil -} func (n *noopMessageQueue) MaxFailureRetries() int32 { return 5 } func TestReaperIntegration(t *testing.T) { @@ -148,28 +149,24 @@ func TestReaperIntegration(t *testing.T) { // Arrange conn, dbService, user := setupReaperTest(t) device := dbService.CreateDevice(user.OrganizationID, "proto") - - batchUUID := "reap-test-batch-1" - createBatchLog(t, conn, batchUUID, user.DatabaseID, 1) - createStuckMessage(t, conn, batchUUID, device.DatabaseID, 10*time.Minute) - svc := command.NewExecutionService(&command.Config{ MaxWorkers: 5, MasterPollingInterval: 100 * time.Millisecond, StuckMessageTimeout: 5 * time.Minute, ReaperInterval: 50 * time.Millisecond, }, conn, &noopMessageQueue{}, nil, nil, nil, nil, nil, nil) - - // Act — start the service, the reaper should fire within 50ms err := svc.Start(t.Context()) require.NoError(t, err) + batchUUID := "reap-test-batch-1" + createBatchLog(t, conn, batchUUID, user.DatabaseID, 1, sqlc.BatchStatusEnumPROCESSING) + createStuckMessage(t, conn, batchUUID, device.DatabaseID, 10*time.Minute) - // Assert — reaper should mark the stuck message as FAILED + // Act assert.Eventually(t, func() bool { return getQueueMessageStatus(t, conn, batchUUID, device.DatabaseID) == sqlc.QueueStatusEnumFAILED }, 500*time.Millisecond, 25*time.Millisecond, "reaper should mark stuck message as FAILED") - // Audit log should also be written + // Assert auditStatus, found := getAuditLogStatus(t, conn, batchUUID, device.DatabaseID) assert.True(t, found, "audit log should exist for reaped message") assert.Equal(t, sqlc.DeviceCommandStatusEnumFAILED, auditStatus) @@ -185,23 +182,19 @@ func TestReaperIntegration(t *testing.T) { // Arrange conn, dbService, user := setupReaperTest(t) device := dbService.CreateDevice(user.OrganizationID, "proto") - - batchUUID := "reap-test-batch-2" - createBatchLog(t, conn, batchUUID, user.DatabaseID, 1) - createStuckMessage(t, conn, batchUUID, device.DatabaseID, 1*time.Minute) - svc := command.NewExecutionService(&command.Config{ MaxWorkers: 5, MasterPollingInterval: 100 * time.Millisecond, StuckMessageTimeout: 5 * time.Minute, ReaperInterval: 50 * time.Millisecond, }, conn, &noopMessageQueue{}, nil, nil, nil, nil, nil, nil) - - // Act err := svc.Start(t.Context()) require.NoError(t, err) + batchUUID := "reap-test-batch-2" + createBatchLog(t, conn, batchUUID, user.DatabaseID, 1, sqlc.BatchStatusEnumPROCESSING) + createStuckMessage(t, conn, batchUUID, device.DatabaseID, 1*time.Minute) - // Wait for a few reaper ticks + // Act time.Sleep(200 * time.Millisecond) // Assert — message should still be PROCESSING @@ -215,7 +208,7 @@ func TestReaperIntegration(t *testing.T) { device := dbService.CreateDevice(user.OrganizationID, "proto") batchUUID := "reap-test-batch-3" - createBatchLog(t, conn, batchUUID, user.DatabaseID, 1) + createBatchLog(t, conn, batchUUID, user.DatabaseID, 1, sqlc.BatchStatusEnumPROCESSING) // Create a message in SUCCESS state with an old timestamp ctx := context.Background() @@ -261,31 +254,27 @@ func TestReaperIntegration(t *testing.T) { conn, dbService, user := setupReaperTest(t) device1 := dbService.CreateDevice(user.OrganizationID, "proto") device2 := dbService.CreateDevice(user.OrganizationID, "proto") - - batchUUID := "reap-test-batch-4" - createBatchLog(t, conn, batchUUID, user.DatabaseID, 2) - createStuckMessage(t, conn, batchUUID, device1.DatabaseID, 10*time.Minute) - createStuckMessage(t, conn, batchUUID, device2.DatabaseID, 10*time.Minute) - svc := command.NewExecutionService(&command.Config{ MaxWorkers: 5, MasterPollingInterval: 100 * time.Millisecond, StuckMessageTimeout: 5 * time.Minute, ReaperInterval: 50 * time.Millisecond, }, conn, &noopMessageQueue{}, nil, nil, nil, nil, nil, nil) - - // Act err := svc.Start(t.Context()) require.NoError(t, err) + batchUUID := "reap-test-batch-4" + createBatchLog(t, conn, batchUUID, user.DatabaseID, 2, sqlc.BatchStatusEnumPROCESSING) + createStuckMessage(t, conn, batchUUID, device1.DatabaseID, 10*time.Minute) + createStuckMessage(t, conn, batchUUID, device2.DatabaseID, 10*time.Minute) - // Assert — both should be reaped + // Act assert.Eventually(t, func() bool { s1 := getQueueMessageStatus(t, conn, batchUUID, device1.DatabaseID) s2 := getQueueMessageStatus(t, conn, batchUUID, device2.DatabaseID) return s1 == sqlc.QueueStatusEnumFAILED && s2 == sqlc.QueueStatusEnumFAILED }, 500*time.Millisecond, 25*time.Millisecond) - // Both should have audit log entries + // Assert status1, found1 := getAuditLogStatus(t, conn, batchUUID, device1.DatabaseID) status2, found2 := getAuditLogStatus(t, conn, batchUUID, device2.DatabaseID) assert.True(t, found1) @@ -294,3 +283,72 @@ func TestReaperIntegration(t *testing.T) { assert.Equal(t, sqlc.DeviceCommandStatusEnumFAILED, status2) }) } + +func TestExecutionServiceStartupPreservesPendingAndFailsProcessing(t *testing.T) { + if testing.Short() { + t.Skip("Skipping database integration test in short mode") + } + + // Arrange + conn, dbService, user := setupReaperTest(t) + queries := sqlc.New(conn) + commandDevice := dbService.CreateDevice(user.OrganizationID, "proto") + firmwareDevice := dbService.CreateDevice(user.OrganizationID, "proto") + commandBatch := "restart-pending-command-batch" + firmwareBatch := "restart-processing-firmware-batch" + rebootCommand := commandtype.Reboot + firmwareCommand := commandtype.FirmwareUpdate + createBatchLog(t, conn, commandBatch, user.DatabaseID, 1, sqlc.BatchStatusEnumPENDING) + createBatchLog(t, conn, firmwareBatch, user.DatabaseID, 1, sqlc.BatchStatusEnumPROCESSING) + require.NoError(t, queries.CreateQueueMessage(t.Context(), sqlc.CreateQueueMessageParams{ + CommandBatchLogUuid: commandBatch, + CommandType: rebootCommand.String(), + DeviceID: commandDevice.DatabaseID, + Status: sqlc.QueueStatusEnumPENDING, + Payload: pqtype.NullRawMessage{}, + })) + require.NoError(t, queries.CreateQueueMessage(t.Context(), sqlc.CreateQueueMessageParams{ + CommandBatchLogUuid: firmwareBatch, + CommandType: firmwareCommand.String(), + DeviceID: firmwareDevice.DatabaseID, + Status: sqlc.QueueStatusEnumPROCESSING, + Payload: pqtype.NullRawMessage{}, + })) + require.NoError(t, queries.UpsertDeviceStatus(t.Context(), sqlc.UpsertDeviceStatusParams{ + DeviceID: firmwareDevice.DatabaseID, + Status: sqlc.DeviceStatusEnumUPDATING, + })) + + svc := command.NewExecutionService( + &command.Config{MaxWorkers: 1}, + conn, + &noopMessageQueue{}, + nil, + nil, + nil, + nil, + nil, + nil, + ) + + // Act + err := svc.Start(t.Context()) + + // Assert + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, svc.Stop(context.Background())) }) + require.Equal(t, sqlc.QueueStatusEnumPENDING, getQueueMessageStatus(t, conn, commandBatch, commandDevice.DatabaseID)) + require.Equal(t, sqlc.QueueStatusEnumFAILED, getQueueMessageStatus(t, conn, firmwareBatch, firmwareDevice.DatabaseID)) + reason, found := getAuditLogErrorInfo(t, conn, firmwareBatch, firmwareDevice.DatabaseID) + require.True(t, found) + require.Equal(t, "Interrupted by Fleet restart; device outcome may be unknown", reason) + firmwareStatus, err := queries.GetDeviceStatus(t.Context(), firmwareDevice.DatabaseID) + require.NoError(t, err) + require.Equal(t, sqlc.DeviceStatusEnumACTIVE, firmwareStatus) + pendingBatch, err := queries.GetBatchLog(t.Context(), commandBatch) + require.NoError(t, err) + require.Equal(t, sqlc.BatchStatusEnumPENDING, pendingBatch.Status) + batch, err := queries.GetBatchLog(t.Context(), firmwareBatch) + require.NoError(t, err) + require.Equal(t, sqlc.BatchStatusEnumFINISHED, batch.Status) +} diff --git a/server/internal/domain/command/results_integration_test.go b/server/internal/domain/command/results_integration_test.go index 05d717dbe..ab21fa95f 100644 --- a/server/internal/domain/command/results_integration_test.go +++ b/server/internal/domain/command/results_integration_test.go @@ -51,12 +51,16 @@ func newResultsTestService(conn *sql.DB) *command.Service { // by the results-RPC tests instead of seedFinishedBatch because several tests // need PENDING / PROCESSING. func seedBatchInState(t *testing.T, conn *sql.DB, batchUUID string, userID, orgID int64, deviceCount int32, status sqlc.BatchStatusEnum) { + seedBatchOfTypeInState(t, conn, batchUUID, "REBOOT", userID, orgID, deviceCount, status) +} + +func seedBatchOfTypeInState(t *testing.T, conn *sql.DB, batchUUID, commandType string, userID, orgID int64, deviceCount int32, status sqlc.BatchStatusEnum) { t.Helper() ctx := context.Background() err := db2.WithTransactionNoResult(ctx, conn, func(q sqlc.Querier) error { _, err := q.CreateCommandBatchLog(ctx, sqlc.CreateCommandBatchLogParams{ Uuid: batchUUID, - Type: "REBOOT", + Type: commandType, CreatedBy: userID, CreatedAt: time.Now(), Status: status, @@ -74,6 +78,51 @@ func seedBatchInState(t *testing.T, conn *sql.DB, batchUUID string, userID, orgI } } +func TestGetCommandBatchLogBundle_NotFoundForCrossOrg(t *testing.T) { + if testing.Short() { + t.Skip("Skipping database integration test in short mode") + } + + // Arrange + conn, dbService, orgAUser := setupRetentionTest(t) + orgBUser := dbService.CreateSuperAdminUser2() + batchUUID := "log-bundle-cross-org-1" + seedBatchOfTypeInState(t, conn, batchUUID, "DownloadLogs", orgAUser.DatabaseID, orgAUser.OrganizationID, 1, sqlc.BatchStatusEnumFINISHED) + svc := newResultsTestService(conn) + ctx := testutil.MockAuthContextForTesting(context.Background(), orgBUser.DatabaseID, orgBUser.OrganizationID) + + // Act + _, err := svc.GetCommandBatchLogBundle(ctx, batchUUID) + + // Assert + require.Error(t, err) + var fleetErr fleeterror.FleetError + require.True(t, errors.As(err, &fleetErr), "expected FleetError, got %T", err) + assert.Equal(t, connect.CodeNotFound, fleetErr.GRPCCode) +} + +func TestGetCommandBatchLogBundle_NotFoundForWrongCommandType(t *testing.T) { + if testing.Short() { + t.Skip("Skipping database integration test in short mode") + } + + // Arrange + conn, _, user := setupRetentionTest(t) + batchUUID := "log-bundle-wrong-type-1" + seedBatchInState(t, conn, batchUUID, user.DatabaseID, user.OrganizationID, 1, sqlc.BatchStatusEnumFINISHED) + svc := newResultsTestService(conn) + ctx := testutil.MockAuthContextForTesting(context.Background(), user.DatabaseID, user.OrganizationID) + + // Act + _, err := svc.GetCommandBatchLogBundle(ctx, batchUUID) + + // Assert + require.Error(t, err) + var fleetErr fleeterror.FleetError + require.True(t, errors.As(err, &fleetErr), "expected FleetError, got %T", err) + assert.Equal(t, connect.CodeNotFound, fleetErr.GRPCCode) +} + func TestGetCommandBatchDeviceResults_HappyPath(t *testing.T) { if testing.Short() { t.Skip("Skipping database integration test in short mode") @@ -268,8 +317,8 @@ func TestGetCommandBatchDeviceResults_TruncatesLargeBatchesWithConsistentCounts( // TestGetCommandBatchDeviceResults_DeviceSnapshot exercises the audit-capture // feature end-to-end: the first Upsert records the raw device-identity fields // (custom_name, manufacturer, model, IP, MAC) onto the codl row, and later -// Upserts (retries, reap-after-success) update status/error_info but must -// never overwrite those captured values — even if the underlying device is +// Upserts for the same terminal result update error_info but must never +// overwrite those captured values, even if the underlying device is // renamed or moves to a new IP between the two writes. func TestGetCommandBatchDeviceResults_DeviceSnapshot(t *testing.T) { if testing.Short() { @@ -299,7 +348,7 @@ func TestGetCommandBatchDeviceResults_DeviceSnapshot(t *testing.T) { } // 1. First write captures the identity. - upsert(sqlc.DeviceCommandStatusEnumSUCCESS, sql.NullString{}) + upsert(sqlc.DeviceCommandStatusEnumFAILED, sql.NullString{String: "initial failure", Valid: true}) // 2. Rename the device and move it to a new IP. Represents a legitimate // operator action that happens between the first Upsert and the reaper's @@ -311,7 +360,7 @@ func TestGetCommandBatchDeviceResults_DeviceSnapshot(t *testing.T) { dev.DatabaseID) require.NoError(t, err) - // 3. Reaper-style second Upsert: status/error_info flip, snapshot must not. + // 3. A retry of the same terminal result updates the error, not the snapshot. upsert(sqlc.DeviceCommandStatusEnumFAILED, sql.NullString{String: "reaper timeout", Valid: true}) svc := newResultsTestService(conn) diff --git a/server/internal/domain/command/service.go b/server/internal/domain/command/service.go index 593b376d3..5d6f8a1a1 100644 --- a/server/internal/domain/command/service.go +++ b/server/internal/domain/command/service.go @@ -36,6 +36,7 @@ import ( "github.com/block/proto-fleet/server/internal/infrastructure/db" id "github.com/block/proto-fleet/server/internal/infrastructure/id" "github.com/block/proto-fleet/server/internal/infrastructure/queue" + "github.com/block/proto-fleet/server/internal/runtimepolicy" sdk "github.com/block/proto-fleet/server/sdk/v1" commonpb "github.com/block/proto-fleet/server/generated/grpc/common/v1" @@ -480,7 +481,7 @@ func (s *Service) saveCommandBatchLogToDB(ctx context.Context, userID, organizat return "", fleeterror.NewInternalErrorf("cannot create command batch: session missing organization_id") } - return db.WithTransaction(ctx, s.conn, func(q sqlc.Querier) (string, error) { + return db.WithTransactionTimeout(ctx, s.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) (string, error) { timeNow := time.Now() newUUID := id.GenerateID() @@ -503,40 +504,74 @@ func (s *Service) saveCommandBatchLogToDB(ctx context.Context, userID, organizat } func (s *Service) statusUpdateIsProcessingBranch(ctx context.Context, commandBatchLogUUID string) (bool, error) { - isProcessing, err := s.messageQueue.IsBatchProcessing(ctx, commandBatchLogUUID) + updated, err := db.WithTransactionTimeout(ctx, s.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) (bool, error) { + rowsAffected, updateErr := q.MarkCommandBatchProcessing(ctx, commandBatchLogUUID) + return rowsAffected > 0, updateErr + }) if err != nil { - return false, fleeterror.NewInternalErrorf("error asking isProcessing: %v", err) - } - if isProcessing { - err = db.WithTransactionNoResult(ctx, s.conn, func(q sqlc.Querier) error { - return q.MarkCommandBatchProcessing(ctx, commandBatchLogUUID) - }) - if err != nil { - return false, fleeterror.NewInternalErrorf("error marking batch: %v", err) - } - return true, nil + return false, fleeterror.NewInternalErrorf("error marking batch: %v", err) } - return false, nil + return updated, nil } func (s *Service) getMarkFinishedBatchFunction(processingMarkedInDB bool) func(ctx context.Context, commandBatchLogUUID string) error { return func(ctx context.Context, commandBatchLogUUID string) error { - return db.WithTransactionNoResult(ctx, s.conn, func(q sqlc.Querier) error { + updated, err := db.WithTransactionTimeout(ctx, s.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) (bool, error) { + var rowsAffected int64 + var updateErr error if processingMarkedInDB { - return q.MarkCommandBatchFinished(ctx, commandBatchLogUUID) + rowsAffected, updateErr = q.MarkCommandBatchFinished(ctx, commandBatchLogUUID) + } else { + rowsAffected, updateErr = q.MarkCommandBatchFinishedWithStartedAt(ctx, commandBatchLogUUID) } - return q.MarkCommandBatchFinishedWithStartedAt(ctx, commandBatchLogUUID) + return rowsAffected > 0, updateErr }) + if err != nil { + return err + } + if !updated { + slog.Debug("command batch already left expected state", "batch_uuid", commandBatchLogUUID) + } + return nil } } -func (s *Service) finishUnenqueuedCommandBatch(ctx context.Context, commandBatchLogUUID string) error { +func (s *Service) reconcileFailedEnqueue(ctx context.Context, commandBatchLogUUID string, expectedMessages int, enqueueErr error) error { cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), dbWriteTimeout) defer cancel() - if err := s.getMarkFinishedBatchFunction(false)(cleanupCtx, commandBatchLogUUID); err != nil { - return fleeterror.NewInternalErrorf("command execution stopped before enqueue; failed to finish command batch: %v", err) + + enqueueCommitted, err := db.WithTransactionTimeout(cleanupCtx, s.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) (bool, error) { + status, err := q.LockCommandBatch(cleanupCtx, commandBatchLogUUID) + if err != nil { + return false, err + } + messageCount, err := q.CountQueueMessagesByBatch(cleanupCtx, commandBatchLogUUID) + if err != nil { + return false, err + } + if messageCount == int64(expectedMessages) { + return true, nil // the enqueue committed before returning an ambiguous error + } + if messageCount != 0 { + return false, fmt.Errorf("enqueue created %d of %d expected queue messages", messageCount, expectedMessages) + } + + if status != sqlc.BatchStatusEnumPENDING { + return false, nil + } + _, err = q.MarkCommandBatchFinishedWithStartedAt(cleanupCtx, commandBatchLogUUID) + return false, err + }) + if err != nil { + return fleeterror.NewInternalErrorf( + "command enqueue failed and reconciliation also failed: %w", + errors.Join(enqueueErr, err), + ) + } + if enqueueCommitted { + return nil } - return fleeterror.NewInternalError("command execution service stopped before enqueue") + return enqueueErr } func (s *Service) statusUpdateIsFinishedBranch(ctx context.Context, commandBatchLogUUID string) (bool, error) { @@ -1051,14 +1086,19 @@ func (s *Service) processCommand(ctx context.Context, command *Command) (*Comman } return s.messageQueue.EnqueueMany(workCtx, batchLogIdentifier, command.commandType, queuePayloads) }) - if errors.Is(err, errExecutionStoppedBeforeEnqueue) { - return nil, s.finishUnenqueuedCommandBatch(ctx, batchLogIdentifier) - } if err != nil { - if len(queuePayloads) == 0 { - return nil, fleeterror.NewInternalErrorf("error enqueuing a batch of commands: %v", err) + var enqueueErr error + switch { + case errors.Is(err, errExecutionStoppedBeforeEnqueue): + enqueueErr = fleeterror.NewInternalError("command execution service stopped before enqueue") + case len(queuePayloads) == 0: + enqueueErr = fleeterror.NewInternalErrorf("error enqueuing a batch of commands: %v", err) + default: + enqueueErr = fleeterror.NewInternalErrorf("error enqueuing per-device command payloads: %v", err) + } + if err := s.reconcileFailedEnqueue(ctx, batchLogIdentifier, len(deviceIDs), enqueueErr); err != nil { + return nil, err } - return nil, fleeterror.NewInternalErrorf("error enqueuing per-device command payloads: %v", err) } return &CommandResult{ @@ -1536,11 +1576,13 @@ func (s *Service) ReapplyCurrentPoolsWithWorkerNames( err = s.executionService.withAdmission(ctx, func(workCtx context.Context) error { return s.enqueueWorkerNameReapplyMessages(workCtx, commandBatchLogUUID, deviceIdentifiers, deviceIDsByIdentifier, desiredWorkerNamesByDeviceIdentifier) }) - if errors.Is(err, errExecutionStoppedBeforeEnqueue) { - return "", s.finishUnenqueuedCommandBatch(ctx, commandBatchLogUUID) - } if err != nil { - return "", err + if errors.Is(err, errExecutionStoppedBeforeEnqueue) { + err = fleeterror.NewInternalError("command execution service stopped before enqueue") + } + if err := s.reconcileFailedEnqueue(ctx, commandBatchLogUUID, len(deviceIdentifiers), err); err != nil { + return "", err + } } s.initializeStatusUpdateRoutine(commandBatchLogUUID, nil) @@ -1572,30 +1614,22 @@ func (s *Service) enqueueWorkerNameReapplyMessages( deviceIDsByIdentifier map[string]int64, desiredWorkerNamesByDeviceIdentifier map[string]string, ) error { - return db.WithTransactionNoResult(ctx, s.conn, func(q sqlc.Querier) error { - commandType := commandtype.UpdateMiningPools - for _, deviceIdentifier := range deviceIdentifiers { - payloadBytes, err := json.Marshal(dto.UpdateMiningPoolsPayload{ + messages := make([]queue.EnqueueMessage, 0, len(deviceIdentifiers)) + for _, deviceIdentifier := range deviceIdentifiers { + messages = append(messages, queue.EnqueueMessage{ + DeviceID: deviceIDsByIdentifier[deviceIdentifier], + Payload: dto.UpdateMiningPoolsPayload{ ReapplyCurrentPoolsWithStoredWorkerName: true, DesiredWorkerName: desiredWorkerNamesByDeviceIdentifier[deviceIdentifier], - }) - if err != nil { - return fleeterror.NewInternalErrorf("failed to marshal worker-name reapply payload: %v", err) - } - - if err := q.CreateQueueMessage(ctx, sqlc.CreateQueueMessageParams{ - CommandBatchLogUuid: commandBatchLogUUID, - CommandType: commandType.String(), - DeviceID: deviceIDsByIdentifier[deviceIdentifier], - Status: sqlc.QueueStatusEnumPENDING, - RetryCount: 0, - Payload: pqtype.NullRawMessage{RawMessage: payloadBytes, Valid: true}, - }); err != nil { - return fleeterror.NewInternalErrorf("failed to enqueue worker-name reapply message: %v", err) - } - } - return nil - }) + }, + }) + } + return s.messageQueue.EnqueueMany( + ctx, + commandBatchLogUUID, + commandtype.UpdateMiningPools, + messages, + ) } func (s *Service) DownloadLogs(ctx context.Context, deviceSelector *pb.DeviceSelector) (*CommandResult, error) { @@ -1882,7 +1916,39 @@ func (s *Service) StreamCommandBatchUpdates(ctx context.Context, msg *pb.StreamC return responseChan, nil } -func (s *Service) GetCommandBatchLogBundle(batchUUID string) (*pb.GetCommandBatchLogBundleResponse, error) { +func (s *Service) GetCommandBatchLogBundle(ctx context.Context, batchUUID string) (*pb.GetCommandBatchLogBundleResponse, error) { + info, err := session.GetInfo(ctx) + if err != nil { + return nil, fleeterror.NewInternalErrorf("error getting session info: %v", err) + } + + batch, err := db.WithTransaction(ctx, s.conn, func(q sqlc.Querier) (sqlc.GetBatchHeaderForOrgRow, error) { + header, queryErr := q.GetBatchHeaderForOrg(ctx, sqlc.GetBatchHeaderForOrgParams{ + Uuid: batchUUID, + OrganizationID: sql.NullInt64{Int64: info.OrganizationID, Valid: true}, + }) + if errors.Is(queryErr, sql.ErrNoRows) { + return header, fleeterror.NewNotFoundErrorf("command batch %s not found", batchUUID) + } + return header, queryErr + }) + if err != nil { + if fleeterror.IsNotFoundError(err) { + return nil, err + } + return nil, fleeterror.NewInternalErrorf("error reading command batch: %v", err) + } + downloadLogs := commandtype.DownloadLogs + if batch.Type != downloadLogs.String() { + return nil, fleeterror.NewNotFoundErrorf("command batch %s not found", batchUUID) + } + if batch.Status != sqlc.BatchStatusEnumFINISHED { + return nil, fleeterror.NewInternalError("log bundle is not available yet, please try again later") + } + if err := s.filesService.EnsureBatchLogBundle(batchUUID); err != nil { + return nil, fleeterror.NewInternalErrorf("error bundling logs: %v", err) + } + file, err := s.filesService.GetBatchLogBundleFile(batchUUID) if err != nil { return nil, err diff --git a/server/internal/domain/command/service_curtail_test.go b/server/internal/domain/command/service_curtail_test.go index 29b94b2d0..2cefd533f 100644 --- a/server/internal/domain/command/service_curtail_test.go +++ b/server/internal/domain/command/service_curtail_test.go @@ -54,21 +54,9 @@ func (f *fakeMessageQueue) EnqueueMany(_ context.Context, batchUUID string, ct c func (f *fakeMessageQueue) Dequeue(context.Context, int32) ([]queue.Message, error) { panic("Dequeue not used") } -func (f *fakeMessageQueue) MarkSuccess(context.Context, int64) error { - panic("MarkSuccess not used") -} -func (f *fakeMessageQueue) MarkFailed(context.Context, int64, string) error { - panic("MarkFailed not used") -} -func (f *fakeMessageQueue) MarkPermanentlyFailed(context.Context, int64, string) error { - panic("MarkPermanentlyFailed not used") -} func (f *fakeMessageQueue) IsBatchFinished(context.Context, string) (bool, error) { return true, nil } -func (f *fakeMessageQueue) IsBatchProcessing(context.Context, string) (bool, error) { - return false, nil -} func (f *fakeMessageQueue) MaxFailureRetries() int32 { return 0 } // newCurtailDispatchService builds a Service wired against in-memory test diff --git a/server/internal/domain/command/zero_target_integration_test.go b/server/internal/domain/command/zero_target_integration_test.go index 9ad6bedf1..b46ab347b 100644 --- a/server/internal/domain/command/zero_target_integration_test.go +++ b/server/internal/domain/command/zero_target_integration_test.go @@ -23,6 +23,14 @@ import ( func newZeroTargetDispatchTestService(t *testing.T, conn *sql.DB) *command.Service { t.Helper() + return newDispatchIntegrationTestService(t, conn, queue.NewDatabaseMessageQueue(&queue.Config{ + DequeLimit: 10, + MaxFailureRetries: 1, + }, conn)) +} + +func newDispatchIntegrationTestService(t *testing.T, conn *sql.DB, messageQueue queue.MessageQueue) *command.Service { + t.Helper() commandConfig := &command.Config{ MaxWorkers: 1, @@ -33,11 +41,6 @@ func newZeroTargetDispatchTestService(t *testing.T, conn *sql.DB) *command.Servi StuckMessageTimeout: time.Hour, ReaperInterval: time.Hour, } - queueConfig := &queue.Config{ - DequeLimit: 10, - MaxFailureRetries: 1, - } - messageQueue := queue.NewDatabaseMessageQueue(queueConfig, conn) executionCtx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) diff --git a/server/internal/ha/coordinator.go b/server/internal/ha/coordinator.go index af04373cf..61ee2fb33 100644 --- a/server/internal/ha/coordinator.go +++ b/server/internal/ha/coordinator.go @@ -40,16 +40,16 @@ type Coordinator struct { config CoordinatorConfig holderID uuid.UUID - mu sync.RWMutex - ownership Ownership - activeCtx context.Context //nolint:containedctx // The coordinator owns this explicit active-lifetime context. - cancelActive context.CancelFunc - leaseTimer *time.Timer - leaseVersion uint64 - stateChanged chan struct{} - acquirePaused bool - lastError string - updatedAt time.Time + mu sync.RWMutex + ownership Ownership + activeCtx context.Context //nolint:containedctx // The coordinator owns this explicit active-lifetime context. + cancelActive context.CancelCauseFunc + leaseTimer *time.Timer + leaseVersion uint64 + stateChanged chan struct{} + acquireAfter time.Time + lastError string + updatedAt time.Time } func NewCoordinator( @@ -159,106 +159,128 @@ func (c *Coordinator) WaitForActive(ctx context.Context) (context.Context, Token } } -// RequestDemotion stops renewal, cancels the current active lifetime, and -// pauses acquisition until runtime cleanup succeeds. -func (c *Coordinator) RequestDemotion(cause error) { - c.deactivate(cause) -} - -// ResumeAcquisition allows a new active lifetime after runtime cleanup. -func (c *Coordinator) ResumeAcquisition() { - c.mu.Lock() - defer c.mu.Unlock() - if !c.acquirePaused { - return - } - c.acquirePaused = false - c.signalStateChangedLocked() -} - -// Run continuously observes and renews until its parent context ends. Failures -// demote immediately and are retried while remaining passive. +// Run retries while passive. After activation, ownership loss is terminal so +// the process supervisor can restart Fleet in a clean passive state. func (c *Coordinator) Run(ctx context.Context) error { for { - err := c.step(ctx) - delay := c.config.RetryInterval - if err == nil && c.Snapshot().State == StateActive { - delay = c.config.RenewInterval + activated, _ := c.tryAcquire(ctx) + if ctx.Err() != nil { + c.deactivate(ctx.Err()) + return fmt.Errorf("HA coordinator stopped: %w", ctx.Err()) } - timer := time.NewTimer(delay) - c.mu.RLock() - var activeDone <-chan struct{} - if c.activeCtx != nil { - activeDone = c.activeCtx.Done() + if activated { + activeCtx, _, active := c.ActiveLifetime() + if !active { + return fmt.Errorf("active Fleet ownership ended: %w", ErrOwnershipLost) + } + return c.renewUntilStopped(ctx, activeCtx) } - stateChanged := c.stateChanged - c.mu.RUnlock() + timer := time.NewTimer(c.config.RetryInterval) select { case <-ctx.Done(): timer.Stop() c.deactivate(ctx.Err()) return fmt.Errorf("HA coordinator stopped: %w", ctx.Err()) - case <-activeDone: - timer.Stop() - case <-stateChanged: - timer.Stop() case <-timer.C: } } } -func (c *Coordinator) step(ctx context.Context) error { +func (c *Coordinator) renewUntilStopped(ctx, activeCtx context.Context) error { + ticker := time.NewTicker(c.config.RenewInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + c.deactivate(ctx.Err()) + return fmt.Errorf("HA coordinator stopped: %w", ctx.Err()) + case <-activeCtx.Done(): + cause := context.Cause(activeCtx) + if cause == nil { + cause = ErrOwnershipLost + } + return fmt.Errorf("active Fleet ownership ended: %w", cause) + case <-ticker.C: + if err := c.renewActive(ctx, activeCtx); err != nil { + return fmt.Errorf("active Fleet ownership ended: %w", err) + } + } + } +} + +func (c *Coordinator) tryAcquire(ctx context.Context) (bool, error) { c.mu.RLock() - activeCtx := c.activeCtx - current := c.ownership holderID := c.holderID - acquirePaused := c.acquirePaused + acquireAfter := c.acquireAfter c.mu.RUnlock() + if time.Now().Before(acquireAfter) { + return false, nil + } - if activeCtx == nil && acquirePaused { - return nil - } - - stepCtx := activeCtx - if stepCtx == nil { - var cancel context.CancelFunc - stepCtx, cancel = context.WithTimeout(ctx, c.config.LeaseDuration) - defer cancel() - } - - if activeCtx == nil { - var ( - ownership Ownership - requestStarted time.Time - ) - observed, err := c.observer.ObserveAndRun( - stepCtx, - func(actionCtx context.Context, observed WriterObservation) error { - requestStarted = time.Now() - var acquireErr error - ownership, acquireErr = c.store.Acquire( - actionCtx, - observed, - holderID, - c.config.LeaseDuration, - ) + takeoverCtx, cancelTakeover := context.WithTimeout(ctx, 2*c.config.LeaseDuration) + defer cancelTakeover() + var ownership Ownership + var candidateExpiresAt time.Time + acquired := false + observed, err := c.observer.ObserveAndRun( + takeoverCtx, + func(actionCtx context.Context, observed WriterObservation) error { + acquireCtx, cancelAcquire := context.WithTimeout(actionCtx, c.config.LeaseDuration) + var acquireErr error + ownership, acquireErr = c.store.Acquire( + acquireCtx, + observed, + holderID, + c.config.LeaseDuration, + ) + cancelAcquire() + if acquireErr != nil { return acquireErr - }, - ) - if err != nil { - c.deactivate(err) - return err - } - if err := c.activate( - ctx, - ownership, - requestStarted, - observed.DCSProofDeadline, - ); err != nil { + } + acquired = true + candidateExpiresAt = time.Now().Add( + ownership.ExpiresAt.Sub(ownership.DatabaseTime), + ) + return nil + }, + ) + if err != nil { + if acquired { + c.abandonCandidate(candidateExpiresAt, err) + } else { c.deactivate(err) - return err } - return nil + return false, err + } + activationCtx, cancelActivation := context.WithDeadline(ctx, observed.DCSProofDeadline) + defer cancelActivation() + requestStarted := time.Now() + renewed, err := c.store.Renew( + activationCtx, + observed, + ownership, + c.config.LeaseDuration, + ) + if err != nil { + c.abandonCandidate(candidateExpiresAt, err) + return false, err + } + ownership = renewed + candidateExpiresAt = time.Now().Add(renewed.ExpiresAt.Sub(renewed.DatabaseTime)) + if err := c.activate(ctx, ownership, requestStarted, observed.DCSProofDeadline); err != nil { + c.abandonCandidate(candidateExpiresAt, err) + return false, err + } + return true, nil +} + +func (c *Coordinator) renewActive(ctx context.Context, expectedCtx context.Context) error { + c.mu.RLock() + activeCtx := c.activeCtx + current := c.ownership + c.mu.RUnlock() + if activeCtx != expectedCtx { + return ErrOwnershipLost } var ( @@ -266,7 +288,7 @@ func (c *Coordinator) step(ctx context.Context) error { requestStarted time.Time ) observed, err := c.observer.ObserveAndRun( - stepCtx, + activeCtx, func(actionCtx context.Context, observed WriterObservation) error { if observed.DCSClusterID != current.DCSClusterID || observed.WriterGeneration != current.Token.WriterGeneration { @@ -327,11 +349,10 @@ func (c *Coordinator) activate( return fmt.Errorf("activate HA coordinator: %w", err) } if c.cancelActive != nil { - c.cancelActive() + c.cancelActive(ErrOwnershipLost) } - c.activeCtx, c.cancelActive = context.WithCancel(parent) + c.activeCtx, c.cancelActive = context.WithCancelCause(parent) c.ownership = ownership - c.acquirePaused = false c.lastError = "" c.updatedAt = time.Now() c.resetLeaseTimerLocked(deadline) @@ -340,13 +361,13 @@ func (c *Coordinator) activate( } func (c *Coordinator) updateActive( - ownership Ownership, + renewed Ownership, expected Ownership, requestStarted time.Time, dcsProofDeadline time.Time, ) error { deadline, err := localOwnershipDeadline( - ownership, + renewed, requestStarted, dcsProofDeadline, ) @@ -362,7 +383,7 @@ func (c *Coordinator) updateActive( c.ownership.HolderID != expected.HolderID { return ErrOwnershipLost } - c.ownership = ownership + c.ownership = renewed c.updatedAt = time.Now() c.lastError = "" c.resetLeaseTimerLocked(deadline) @@ -421,17 +442,24 @@ func (c *Coordinator) deactivate(cause error) { c.deactivateLocked(cause) } +func (c *Coordinator) abandonCandidate(leaseExpiresAt time.Time, cause error) { + c.mu.Lock() + defer c.mu.Unlock() + c.holderID = uuid.New() + // Give another process one normal retry interval after this lease expires. + c.acquireAfter = leaseExpiresAt.Add(c.config.RetryInterval) + c.deactivateLocked(cause) +} + func (c *Coordinator) deactivateLocked(cause error) { wasActive := c.activeCtx != nil c.stopLeaseTimerLocked() if c.cancelActive != nil { - c.cancelActive() + c.cancelActive(cause) c.cancelActive = nil c.activeCtx = nil } if wasActive { - c.holderID = uuid.New() - c.acquirePaused = true c.signalStateChangedLocked() } c.ownership = Ownership{} diff --git a/server/internal/ha/coordinator_test.go b/server/internal/ha/coordinator_test.go index f99261220..bf6c7f822 100644 --- a/server/internal/ha/coordinator_test.go +++ b/server/internal/ha/coordinator_test.go @@ -22,21 +22,6 @@ func TestNewCoordinatorUsesRandomProcessIncarnation(t *testing.T) { require.NotEqual(t, first.HolderID(), second.HolderID()) } -func TestTokenCompareOrdersWriterGenerationBeforeLeaseEpoch(t *testing.T) { - require.Equal(t, -1, (Token{WriterGeneration: 41, LeaseEpoch: 99}).Compare( - Token{WriterGeneration: 42, LeaseEpoch: 1}, - )) - require.Equal(t, -1, (Token{WriterGeneration: 42, LeaseEpoch: 1}).Compare( - Token{WriterGeneration: 42, LeaseEpoch: 2}, - )) - require.Equal(t, 0, (Token{WriterGeneration: 42, LeaseEpoch: 2}).Compare( - Token{WriterGeneration: 42, LeaseEpoch: 2}, - )) - require.Equal(t, 1, (Token{WriterGeneration: 43, LeaseEpoch: 1}).Compare( - Token{WriterGeneration: 42, LeaseEpoch: 100}, - )) -} - func TestCoordinatorActivatesAndExposesLifetime(t *testing.T) { holder := uuid.New() store := &fakeLeaseStore{} @@ -55,7 +40,35 @@ func TestCoordinatorActivatesAndExposesLifetime(t *testing.T) { require.Equal(t, StateActive, coordinator.Snapshot().State) } -func TestCoordinatorWaitForActiveUnblocksOnActivationAndRequestedDemotion(t *testing.T) { +func TestCoordinatorRenewsAfterClosingDCSProof(t *testing.T) { + // Arrange + closingProofDone := false + store := &fakeLeaseStore{} + coordinator := newCoordinatorWithHolder( + &closingProofObserver{ + observation: coordinatorObservation("cluster-a", 41, time.Second), + completed: &closingProofDone, + }, + store, + coordinatorTestConfig(), + uuid.New(), + ) + + // Act + err := coordinator.step(t.Context()) + + // Assert + require.NoError(t, err) + require.True(t, closingProofDone) + require.Equal( + t, + []string{"acquire", "renew"}, + store.callSequence(), + ) + require.Equal(t, StateActive, coordinator.Snapshot().State) +} + +func TestCoordinatorWaitForActiveUnblocksOnActivationAndOwnershipLoss(t *testing.T) { coordinator := newCoordinatorWithHolder( staticObserver{observation: coordinatorObservation("cluster-a", 41, time.Second)}, &fakeLeaseStore{}, @@ -75,34 +88,12 @@ func TestCoordinatorWaitForActiveUnblocksOnActivationAndRequestedDemotion(t *tes require.NoError(t, coordinator.step(t.Context())) result := <-waitResult require.NoError(t, result.err) - coordinator.RequestDemotion(errors.New("runtime unhealthy")) + coordinator.deactivate(ErrOwnershipLost) require.Eventually(t, func() bool { return result.ctx.Err() != nil }, time.Second, time.Millisecond) require.Equal(t, StatePassive, coordinator.Snapshot().State) } -func TestCoordinatorWaitsForRuntimeCleanupBeforeReacquiring(t *testing.T) { - store := &fakeLeaseStore{} - coordinator := newCoordinatorWithHolder( - staticObserver{observation: coordinatorObservation("cluster-a", 41, time.Second)}, - store, - coordinatorTestConfig(), - uuid.New(), - ) - - require.NoError(t, coordinator.step(t.Context())) - require.Equal(t, 1, store.acquireCount()) - coordinator.RequestDemotion(errors.New("runtime unhealthy")) - - require.NoError(t, coordinator.step(t.Context())) - require.Equal(t, 1, store.acquireCount()) - - coordinator.ResumeAcquisition() - require.NoError(t, coordinator.step(t.Context())) - require.Equal(t, 2, store.acquireCount()) - require.Equal(t, StateActive, coordinator.Snapshot().State) -} - func TestCoordinatorCancelsLifetimeOnObservationLoss(t *testing.T) { holder := uuid.New() observer := &sequenceObserver{ @@ -121,29 +112,48 @@ func TestCoordinatorCancelsLifetimeOnObservationLoss(t *testing.T) { require.Error(t, coordinator.step(t.Context())) require.Error(t, activeCtx.Err()) require.Equal(t, StatePassive, coordinator.Snapshot().State) - require.NotEqual(t, holder, coordinator.HolderID()) + require.Equal(t, holder, coordinator.HolderID()) } -func TestCoordinatorKeepsHolderAfterPassiveAcquisitionProofFailure(t *testing.T) { +func TestCoordinatorGivesPeersAnAcquisitionWindowAfterPassiveProofFailure(t *testing.T) { + // Arrange holder := uuid.New() proofErr := errors.New("closing DCS proof failed") + store := &fakeLeaseStore{} + config := coordinatorTestConfig() coordinator := newCoordinatorWithHolder( actionThenErrorObserver{ observation: coordinatorObservation("cluster-a", 41, time.Second), err: proofErr, }, - &fakeLeaseStore{}, - coordinatorTestConfig(), + store, + config, holder, ) + // Act require.ErrorIs(t, coordinator.step(t.Context()), proofErr) + require.NoError(t, coordinator.step(t.Context())) + + // Assert require.Equal(t, StatePassive, coordinator.Snapshot().State) - require.Equal(t, holder, coordinator.HolderID()) + require.NotEqual(t, holder, coordinator.HolderID()) + require.Equal(t, 1, store.acquireCount()) + coordinator.mu.RLock() + acquireAfter := coordinator.acquireAfter + coordinator.mu.RUnlock() + require.True(t, acquireAfter.After(time.Now().Add(config.RetryInterval))) + + coordinator.mu.Lock() + coordinator.acquireAfter = time.Now().Add(-time.Millisecond) + coordinator.mu.Unlock() + + require.ErrorIs(t, coordinator.step(t.Context()), proofErr) + require.Equal(t, 2, store.acquireCount()) } func TestCoordinatorCancelsLifetimeOnRenewalLoss(t *testing.T) { - store := &fakeLeaseStore{renewErr: ErrOwnershipLost} + store := &fakeLeaseStore{} coordinator := newCoordinatorWithHolder( staticObserver{observation: coordinatorObservation("cluster-a", 41, time.Second)}, store, @@ -153,12 +163,60 @@ func TestCoordinatorCancelsLifetimeOnRenewalLoss(t *testing.T) { require.NoError(t, coordinator.step(t.Context())) activeCtx, _, active := coordinator.ActiveLifetime() require.True(t, active) + store.setRenewError(ErrOwnershipLost) require.ErrorIs(t, coordinator.step(t.Context()), ErrOwnershipLost) require.Error(t, activeCtx.Err()) require.Equal(t, StatePassive, coordinator.Snapshot().State) } +func TestCoordinatorRunStopsAfterActiveOwnershipLoss(t *testing.T) { + // Arrange + store := &fakeLeaseStore{} + config := coordinatorTestConfig() + config.RenewInterval = time.Millisecond + coordinator := newCoordinatorWithHolder( + staticObserver{observation: coordinatorObservation("cluster-a", 41, time.Second)}, + store, + config, + uuid.New(), + ) + runResult := make(chan error, 1) + go func() { runResult <- coordinator.Run(t.Context()) }() + activeCtx, _, err := coordinator.WaitForActive(t.Context()) + require.NoError(t, err) + + // Act + store.setRenewError(ErrOwnershipLost) + + // Assert + require.ErrorIs(t, <-runResult, ErrOwnershipLost) + require.ErrorIs(t, context.Cause(activeCtx), ErrOwnershipLost) + require.Equal(t, StatePassive, coordinator.Snapshot().State) +} + +func TestCoordinatorRenewalCannotReacquireAfterWatchdogDemotion(t *testing.T) { + // Arrange + store := &fakeLeaseStore{} + coordinator := newCoordinatorWithHolder( + staticObserver{observation: coordinatorObservation("cluster-a", 41, time.Second)}, + store, + coordinatorTestConfig(), + uuid.New(), + ) + require.NoError(t, coordinator.step(t.Context())) + activeCtx, _, active := coordinator.ActiveLifetime() + require.True(t, active) + coordinator.deactivate(ErrOwnershipExpired) + + // Act + err := coordinator.renewActive(t.Context(), activeCtx) + + // Assert + require.ErrorIs(t, err, ErrOwnershipLost) + require.Equal(t, 1, store.acquireCount()) +} + func TestCoordinatorCancelsLifetimeOnWriterGenerationChange(t *testing.T) { observer := &sequenceObserver{ results: []observerResult{ @@ -244,7 +302,7 @@ func TestCoordinatorActiveRenewalStopsAtWatchdogDeadline(t *testing.T) { require.Equal(t, StatePassive, coordinator.Snapshot().State) } -func TestCoordinatorRejectsLeaseThatExpiredBeforeAcquireReturned(t *testing.T) { +func TestCoordinatorCannotRenewLeaseThatExpiredBeforeAcquireReturned(t *testing.T) { config := coordinatorTestConfig() config.LeaseDuration = 10 * time.Millisecond config.RenewInterval = 5 * time.Millisecond @@ -256,7 +314,7 @@ func TestCoordinatorRejectsLeaseThatExpiredBeforeAcquireReturned(t *testing.T) { uuid.New(), ) - require.ErrorIs(t, coordinator.step(t.Context()), ErrOwnershipExpired) + require.ErrorIs(t, coordinator.step(t.Context()), ErrOwnershipLost) _, _, active := coordinator.ActiveLifetime() require.False(t, active) require.Equal(t, StatePassive, coordinator.Snapshot().State) @@ -338,11 +396,36 @@ func coordinatorTestConfig() CoordinatorConfig { } } +func (c *Coordinator) step(ctx context.Context) error { + activeCtx, _, active := c.ActiveLifetime() + if active { + return c.renewActive(ctx, activeCtx) + } + _, err := c.tryAcquire(ctx) + return err +} + type staticObserver struct { observation WriterObservation err error } +type closingProofObserver struct { + observation WriterObservation + completed *bool +} + +func (o *closingProofObserver) ObserveAndRun( + ctx context.Context, + action func(context.Context, WriterObservation) error, +) (WriterObservation, error) { + if err := action(ctx, o.observation); err != nil { + return WriterObservation{}, err + } + *o.completed = true + return o.observation, nil +} + func (s staticObserver) ObserveAndRun( ctx context.Context, action func(context.Context, WriterObservation) error, @@ -437,7 +520,7 @@ func (b *activateThenBlockObserver) ObserveAndRun( type fakeLeaseStore struct { mu sync.Mutex - ownership Ownership + calls []string acquireErr error renewErr error acquireDelay time.Duration @@ -453,6 +536,7 @@ func (f *fakeLeaseStore) Acquire( f.mu.Lock() defer f.mu.Unlock() f.acquires++ + f.calls = append(f.calls, "acquire") if f.acquireErr != nil { return Ownership{}, f.acquireErr } @@ -460,7 +544,7 @@ func (f *fakeLeaseStore) Acquire( if f.acquireDelay > 0 { time.Sleep(f.acquireDelay) } - f.ownership = Ownership{ + active := Ownership{ DCSClusterID: observed.DCSClusterID, Token: Token{ WriterGeneration: observed.WriterGeneration, @@ -470,7 +554,7 @@ func (f *fakeLeaseStore) Acquire( DatabaseTime: now, ExpiresAt: now.Add(duration), } - return f.ownership, nil + return active, nil } func (f *fakeLeaseStore) acquireCount() int { @@ -482,17 +566,32 @@ func (f *fakeLeaseStore) acquireCount() int { func (f *fakeLeaseStore) Renew( _ context.Context, _ WriterObservation, - ownership Ownership, + active Ownership, duration time.Duration, ) (Ownership, error) { f.mu.Lock() defer f.mu.Unlock() + f.calls = append(f.calls, "renew") if f.renewErr != nil { return Ownership{}, f.renewErr } now := time.Now() - ownership.DatabaseTime = now - ownership.ExpiresAt = now.Add(duration) - f.ownership = ownership - return ownership, nil + if !active.ExpiresAt.After(now) { + return Ownership{}, ErrOwnershipLost + } + active.DatabaseTime = now + active.ExpiresAt = now.Add(duration) + return active, nil +} + +func (f *fakeLeaseStore) setRenewError(err error) { + f.mu.Lock() + defer f.mu.Unlock() + f.renewErr = err +} + +func (f *fakeLeaseStore) callSequence() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.calls...) } diff --git a/server/internal/ha/runtime.go b/server/internal/ha/runtime.go index c8427375f..757f39d7b 100644 --- a/server/internal/ha/runtime.go +++ b/server/internal/ha/runtime.go @@ -24,15 +24,12 @@ type RuntimeConfig struct { type runtimeOwner interface { Run(ctx context.Context) error WaitForActive(ctx context.Context) (context.Context, Token, error) - RequestDemotion(cause error) - ResumeAcquisition() } type runtimeGroup interface { Start(ctx context.Context) error Abort() Stop(ctx context.Context) error - Err() error } // Runtime serializes lease ownership, the existing runtime-job group, and @@ -145,9 +142,21 @@ func (r *Runtime) runHA(ctx context.Context) error { coordinatorResult <- r.owner.Run(coordinatorCtx) }() - for { + type activationResult struct { + ctx context.Context //nolint:containedctx // Carries the owned lifetime returned by WaitForActive. + err error + } + activation := make(chan activationResult, 1) + go func() { activeCtx, _, err := r.owner.WaitForActive(coordinatorCtx) - if err != nil { + activation <- activationResult{ctx: activeCtx, err: err} + }() + + var activeCtx context.Context + select { + case result := <-activation: + activeCtx = result.ctx + if result.err != nil { if ctx.Err() != nil { return nil } @@ -155,47 +164,36 @@ func (r *Runtime) runHA(ctx context.Context) error { case coordinatorErr := <-coordinatorResult: return fmt.Errorf("run Fleet ownership coordinator: %w", coordinatorErr) default: - return fmt.Errorf("wait for Fleet ownership: %w", err) + return fmt.Errorf("wait for Fleet ownership: %w", result.err) } } + case coordinatorErr := <-coordinatorResult: + return fmt.Errorf("run Fleet ownership coordinator: %w", coordinatorErr) + case <-ctx.Done(): + return nil + } - if err := r.group.Start(activeCtx); err != nil { - r.owner.RequestDemotion(fmt.Errorf("start Fleet runtime: %w", err)) - if terminalErr := r.group.Err(); terminalErr != nil { - return fmt.Errorf("start Fleet runtime left terminal cleanup failure: %w", terminalErr) - } - r.owner.ResumeAcquisition() - continue - } + if err := r.group.Start(activeCtx); err != nil { + return fmt.Errorf("start active Fleet runtime: %w", err) + } + if !r.healthCheck() { + r.group.Abort() + return errCriticalRuntimeUnhealthy + } - if !r.healthCheck() { - r.group.Abort() - if activeCtx.Err() == nil { - r.owner.RequestDemotion(errCriticalRuntimeUnhealthy) - } - if stopErr := r.stopGroup(); stopErr != nil { - return stopErr - } - if ctx.Err() == nil { - r.owner.ResumeAcquisition() - } - continue - } + r.gate.activate(activeCtx) + activeErr := r.waitWhileHealthy(ctx, activeCtx) + admissionDrained := r.gate.deactivate() + if ctx.Err() != nil { + return r.stopGroupAndDrainAdmissions(admissionDrained) + } + r.group.Abort() - r.gate.activate(activeCtx) - activeErr := r.waitWhileHealthy(ctx, activeCtx) - admissionDrained := r.gate.deactivate() - r.group.Abort() - if errors.Is(activeErr, errCriticalRuntimeUnhealthy) { - r.owner.RequestDemotion(activeErr) - } - if err := r.stopGroupAndDrainAdmissions(admissionDrained); err != nil { - return err - } - if ctx.Err() != nil { - return nil - } - r.owner.ResumeAcquisition() + select { + case coordinatorErr := <-coordinatorResult: + return errors.Join(activeErr, coordinatorErr) + default: + return activeErr } } @@ -207,7 +205,7 @@ func (r *Runtime) waitWhileHealthy(parent, activeCtx context.Context) error { case <-parent.Done(): return fmt.Errorf("Fleet runtime stopped: %w", parent.Err()) case <-activeCtx.Done(): - return fmt.Errorf("active lifetime ended: %w", activeCtx.Err()) + return fmt.Errorf("active lifetime ended: %w", context.Cause(activeCtx)) case <-ticker.C: if !r.healthCheck() { return errCriticalRuntimeUnhealthy diff --git a/server/internal/ha/runtime_test.go b/server/internal/ha/runtime_test.go index 71db6d11a..a6935cc49 100644 --- a/server/internal/ha/runtime_test.go +++ b/server/internal/ha/runtime_test.go @@ -2,7 +2,6 @@ package ha import ( "context" - "errors" "fmt" "sync" "sync/atomic" @@ -25,16 +24,10 @@ type runtimeTestActivation struct { type runtimeTestOwner struct { activations chan runtimeTestActivation - demotions chan error - resumed chan struct{} } func newRuntimeTestOwner() *runtimeTestOwner { - return &runtimeTestOwner{ - activations: make(chan runtimeTestActivation, 4), - demotions: make(chan error, 4), - resumed: make(chan struct{}, 4), - } + return &runtimeTestOwner{activations: make(chan runtimeTestActivation, 1)} } func (o *runtimeTestOwner) Run(ctx context.Context) error { @@ -51,39 +44,48 @@ func (o *runtimeTestOwner) WaitForActive(ctx context.Context) (context.Context, } } -func (o *runtimeTestOwner) RequestDemotion(err error) { - o.demotions <- err +type missedActivationOwner struct{} + +func (missedActivationOwner) Run(context.Context) error { + return ErrOwnershipLost } -func (o *runtimeTestOwner) ResumeAcquisition() { - o.resumed <- struct{}{} +func (missedActivationOwner) WaitForActive(ctx context.Context) (context.Context, Token, error) { + <-ctx.Done() + return nil, Token{}, fmt.Errorf("wait for missed activation: %w", ctx.Err()) } type runtimeTestGroup struct { mu sync.Mutex - started int - stopped int - startErr error - stopErr error - terminalErr error - aborted int - abortFirst bool - startedCh chan context.Context - stoppedCh chan struct{} + started int + stopped int + aborted int + startErr error + stopErr error + startedCh chan context.Context + stoppedCh chan struct{} + abortedCh chan struct{} } func newRuntimeTestGroup() *runtimeTestGroup { return &runtimeTestGroup{ - startedCh: make(chan context.Context, 4), - stoppedCh: make(chan struct{}, 4), + startedCh: make(chan context.Context, 1), + stoppedCh: make(chan struct{}, 1), + abortedCh: make(chan struct{}, 1), } } +func (g *runtimeTestGroup) Abort() { + g.mu.Lock() + g.aborted++ + g.mu.Unlock() + g.abortedCh <- struct{}{} +} + func (g *runtimeTestGroup) Start(ctx context.Context) error { g.mu.Lock() g.started++ - g.aborted = 0 err := g.startErr g.mu.Unlock() g.startedCh <- ctx @@ -93,31 +95,12 @@ func (g *runtimeTestGroup) Start(ctx context.Context) error { func (g *runtimeTestGroup) Stop(context.Context) error { g.mu.Lock() g.stopped++ - g.abortFirst = g.aborted > 0 err := g.stopErr g.mu.Unlock() g.stoppedCh <- struct{}{} return err } -func (g *runtimeTestGroup) Abort() { - g.mu.Lock() - defer g.mu.Unlock() - g.aborted++ -} - -func (g *runtimeTestGroup) wasAbortedBeforeStop() bool { - g.mu.Lock() - defer g.mu.Unlock() - return g.abortFirst -} - -func (g *runtimeTestGroup) Err() error { - g.mu.Lock() - defer g.mu.Unlock() - return g.terminalErr -} - func TestNewRuntimeRequiresExplicitCriticalHealth(t *testing.T) { group, err := runtimejobs.NewGroup(nil) require.NoError(t, err) @@ -129,241 +112,138 @@ func TestNewRuntimeRequiresExplicitCriticalHealth(t *testing.T) { require.ErrorContains(t, err, "requires a critical health check") } -func TestRuntimeStartsOnlyForOwnedLifetimeAndDrainsOnDemotion(t *testing.T) { +func TestHARuntimeStartsOnlyAfterOwnership(t *testing.T) { + // Arrange owner := newRuntimeTestOwner() group := newRuntimeTestGroup() runtime := newRuntime(owner, group, alwaysHealthy, runtimeTestConfig()) runCtx, cancelRun := context.WithCancel(t.Context()) runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() + go func() { runResult <- runtime.Run(runCtx) }() + // Act require.Never(t, runtime.Active, 20*time.Millisecond, time.Millisecond) - select { - case <-group.startedCh: - t.Fatal("passive runtime started jobs") - default: - } - activeCtx, cancelActive := context.WithCancel(t.Context()) - owner.activations <- runtimeTestActivation{ - ctx: activeCtx, - token: Token{WriterGeneration: 7, LeaseEpoch: 11}, - } + owner.activations <- runtimeTestActivation{ctx: activeCtx} + + // Assert requireReceiveContext(t, group.startedCh) require.Eventually(t, runtime.Active, eventuallyTimeout, eventuallyInterval) - requestCtx, release, err := runtime.Admit(t.Context()) - require.NoError(t, err) - + cancelRun() cancelActive() - require.Eventually(t, func() bool { return requestCtx.Err() != nil }, eventuallyTimeout, eventuallyInterval) requireReceive(t, group.stoppedCh) - require.Never(t, channelClosed(owner.resumed), 20*time.Millisecond, time.Millisecond) - release() - requireReceive(t, owner.resumed) - require.False(t, runtime.Active()) - - cancelRun() require.NoError(t, <-runResult) } -func TestRuntimeAdmissionDrainTimeoutIsTerminal(t *testing.T) { +func TestHARuntimeReturnsWhenOwnershipEndsBeforeActivationIsObserved(t *testing.T) { + // Arrange + runtime := newRuntime(missedActivationOwner{}, newRuntimeTestGroup(), alwaysHealthy, runtimeTestConfig()) + + // Act + err := runtime.Run(t.Context()) + + // Assert + require.ErrorIs(t, err, ErrOwnershipLost) + require.False(t, runtime.Active()) +} + +func TestHARuntimeExitsOnOwnershipLossWithoutInProcessCleanup(t *testing.T) { + // Arrange owner := newRuntimeTestOwner() group := newRuntimeTestGroup() - config := runtimeTestConfig() - config.CleanupTimeout = 20 * time.Millisecond - runtime := newRuntime(owner, group, alwaysHealthy, config) - runCtx, cancelRun := context.WithCancel(t.Context()) - defer cancelRun() + runtime := newRuntime(owner, group, alwaysHealthy, runtimeTestConfig()) runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - - activeCtx, cancelActive := context.WithCancel(t.Context()) + go func() { runResult <- runtime.Run(t.Context()) }() + activeCtx, cancelActive := context.WithCancelCause(t.Context()) owner.activations <- runtimeTestActivation{ctx: activeCtx} - requireReceiveContext(t, group.startedCh) + jobCtx := requireReceiveContext(t, group.startedCh) require.Eventually(t, runtime.Active, eventuallyTimeout, eventuallyInterval) - _, release, err := runtime.Admit(t.Context()) + requestCtx, release, err := runtime.Admit(t.Context()) require.NoError(t, err) defer release() - cancelActive() - requireReceive(t, group.stoppedCh) - require.ErrorContains(t, <-runResult, "drain active Fleet requests") + // Act + cancelActive(ErrOwnershipLost) + + // Assert + require.ErrorContains(t, <-runResult, "active lifetime ended") + requireReceive(t, group.abortedCh) + require.Eventually(t, func() bool { return requestCtx.Err() != nil }, eventuallyTimeout, eventuallyInterval) + require.False(t, runtime.Active()) + require.Error(t, jobCtx.Err()) select { - case <-owner.resumed: - t.Fatal("terminal admission drain failure resumed acquisition") + case <-group.stoppedCh: + t.Fatal("ownership loss performed graceful in-process cleanup") default: } } -func TestRuntimeDemotesWhenCriticalHealthFailsAfterAdmission(t *testing.T) { +func TestHARuntimeExitsWhenCriticalHealthFails(t *testing.T) { + // Arrange owner := newRuntimeTestOwner() group := newRuntimeTestGroup() var healthy atomic.Bool healthy.Store(true) runtime := newRuntime(owner, group, healthy.Load, runtimeTestConfig()) - runCtx, cancelRun := context.WithCancel(t.Context()) runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - + go func() { runResult <- runtime.Run(t.Context()) }() activeCtx, cancelActive := context.WithCancel(t.Context()) defer cancelActive() owner.activations <- runtimeTestActivation{ctx: activeCtx} requireReceiveContext(t, group.startedCh) require.Eventually(t, runtime.Active, eventuallyTimeout, eventuallyInterval) - requestCtx, release, err := runtime.Admit(t.Context()) - require.NoError(t, err) + // Act healthy.Store(false) - _ = requireReceiveError(t, owner.demotions) - requireReceive(t, group.stoppedCh) - require.Eventually(t, func() bool { return requestCtx.Err() != nil }, eventuallyTimeout, eventuallyInterval) - release() - requireReceive(t, owner.resumed) - require.False(t, runtime.Active()) - require.True(t, group.wasAbortedBeforeStop()) - - cancelRun() - require.NoError(t, <-runResult) -} - -func TestRuntimeDemotesWhenCriticalHealthFailsDuringStartup(t *testing.T) { - owner := newRuntimeTestOwner() - group := newRuntimeTestGroup() - runtime := newRuntime(owner, group, func() bool { return false }, runtimeTestConfig()) - runCtx, cancelRun := context.WithCancel(t.Context()) - runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - - activeCtx, cancelActive := context.WithCancel(t.Context()) - defer cancelActive() - owner.activations <- runtimeTestActivation{ctx: activeCtx} - requireReceiveContext(t, group.startedCh) - _ = requireReceiveError(t, owner.demotions) - requireReceive(t, group.stoppedCh) - requireReceive(t, owner.resumed) - require.False(t, runtime.Active()) - cancelRun() - require.NoError(t, <-runResult) -} - -func TestRuntimeStopsOwnershipAfterTerminalCleanupFailure(t *testing.T) { - owner := newRuntimeTestOwner() - group := newRuntimeTestGroup() - group.stopErr = errors.New("cleanup failed") - group.terminalErr = group.stopErr - runtime := newRuntime(owner, group, alwaysHealthy, runtimeTestConfig()) - runCtx, cancelRun := context.WithCancel(t.Context()) - defer cancelRun() - runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - - activeCtx, cancelActive := context.WithCancel(t.Context()) - owner.activations <- runtimeTestActivation{ctx: activeCtx} - requireReceiveContext(t, group.startedCh) - require.Eventually(t, runtime.Active, eventuallyTimeout, eventuallyInterval) - cancelActive() - - err := <-runResult - require.ErrorContains(t, err, "cleanup failed") - require.False(t, runtime.Active()) -} - -func TestRuntimeStartFailureDemotesWithoutAdmission(t *testing.T) { - owner := newRuntimeTestOwner() - group := newRuntimeTestGroup() - group.startErr = errors.New("start failed") - runtime := newRuntime(owner, group, alwaysHealthy, runtimeTestConfig()) - runCtx, cancelRun := context.WithCancel(t.Context()) - runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - - activeCtx, cancelActive := context.WithCancel(t.Context()) - defer cancelActive() - owner.activations <- runtimeTestActivation{ctx: activeCtx} - requireReceiveContext(t, group.startedCh) - require.ErrorContains(t, requireReceiveError(t, owner.demotions), "start failed") - requireReceive(t, owner.resumed) + // Assert + require.ErrorIs(t, <-runResult, errCriticalRuntimeUnhealthy) + requireReceive(t, group.abortedCh) require.False(t, runtime.Active()) - - cancelRun() - require.NoError(t, <-runResult) -} - -func TestRuntimeTerminalStartFailureStopsCoordinator(t *testing.T) { - owner := newRuntimeTestOwner() - group := newRuntimeTestGroup() - group.startErr = errors.New("start failed") - group.terminalErr = errors.New("rollback failed") - runtime := newRuntime(owner, group, alwaysHealthy, runtimeTestConfig()) - runCtx, cancelRun := context.WithCancel(t.Context()) - defer cancelRun() - runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - - activeCtx, cancelActive := context.WithCancel(t.Context()) - defer cancelActive() - owner.activations <- runtimeTestActivation{ctx: activeCtx} - requireReceiveContext(t, group.startedCh) - err := <-runResult - require.ErrorContains(t, err, "rollback failed") select { - case <-owner.resumed: - t.Fatal("terminal cleanup failure resumed acquisition") + case <-group.stoppedCh: + t.Fatal("critical failure performed graceful in-process cleanup") default: } } -func TestStandaloneRuntimePreservesSingleHostLifecycle(t *testing.T) { +func TestStandaloneRuntimePreservesGracefulLifecycle(t *testing.T) { + // Arrange group := newRuntimeTestGroup() runtime := newRuntime(nil, group, alwaysHealthy, runtimeTestConfig()) runCtx, cancelRun := context.WithCancel(t.Context()) runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() + go func() { runResult <- runtime.Run(runCtx) }() + // Act requireReceiveContext(t, group.startedCh) require.Eventually(t, runtime.Active, eventuallyTimeout, eventuallyInterval) cancelRun() + + // Assert requireReceive(t, group.stoppedCh) require.NoError(t, <-runResult) } func TestStandaloneRuntimeStopsWhenCriticalHealthFails(t *testing.T) { + // Arrange group := newRuntimeTestGroup() var healthy atomic.Bool healthy.Store(true) runtime := newRuntime(nil, group, healthy.Load, runtimeTestConfig()) - runCtx, cancelRun := context.WithCancel(t.Context()) - defer cancelRun() runResult := make(chan error, 1) - go func() { - runResult <- runtime.Run(runCtx) - }() - + go func() { runResult <- runtime.Run(t.Context()) }() requireReceiveContext(t, group.startedCh) require.Eventually(t, runtime.Active, eventuallyTimeout, eventuallyInterval) requestCtx, release, err := runtime.Admit(t.Context()) require.NoError(t, err) defer release() + // Act healthy.Store(false) + + // Assert requireReceive(t, group.stoppedCh) require.Eventually(t, func() bool { return requestCtx.Err() != nil }, eventuallyTimeout, eventuallyInterval) require.False(t, runtime.Active()) @@ -400,14 +280,3 @@ func requireReceiveContext(t *testing.T, ch <-chan context.Context) context.Cont return nil } } - -func requireReceiveError(t *testing.T, ch <-chan error) error { - t.Helper() - select { - case err := <-ch: - return err - case <-time.After(eventuallyTimeout): - t.Fatal("timed out waiting for error") - return nil - } -} diff --git a/server/internal/ha/store.go b/server/internal/ha/store.go index 8f46cf389..8596b3a58 100644 --- a/server/internal/ha/store.go +++ b/server/internal/ha/store.go @@ -76,18 +76,18 @@ func (s *LeaseStore) Acquire( func (s *LeaseStore) Renew( ctx context.Context, observed WriterObservation, - ownership Ownership, + active Ownership, duration time.Duration, ) (Ownership, error) { if err := validateWriterObservation(observed); err != nil { return Ownership{}, err } - if observed.DCSClusterID != ownership.DCSClusterID || - observed.WriterGeneration != ownership.Token.WriterGeneration || - ownership.DCSClusterID == "" || - ownership.Token.WriterGeneration <= 0 || - ownership.Token.LeaseEpoch <= 0 || - ownership.HolderID == uuid.Nil || + if observed.DCSClusterID != active.DCSClusterID || + observed.WriterGeneration != active.Token.WriterGeneration || + active.DCSClusterID == "" || + active.Token.WriterGeneration <= 0 || + active.Token.LeaseEpoch <= 0 || + active.HolderID == uuid.Nil || duration.Milliseconds() <= 0 { return Ownership{}, errors.New("invalid Fleet active lease renewal") } @@ -98,10 +98,10 @@ func (s *LeaseStore) Renew( ServerPort: observed.ServerPort, Timeline: observed.Timeline, LeaseDurationMilliseconds: duration.Milliseconds(), - DcsClusterID: ownership.DCSClusterID, - WriterGeneration: ownership.Token.WriterGeneration, - LeaseEpoch: ownership.Token.LeaseEpoch, - HolderID: ownership.HolderID, + DcsClusterID: active.DCSClusterID, + WriterGeneration: active.Token.WriterGeneration, + LeaseEpoch: active.Token.LeaseEpoch, + HolderID: active.HolderID, }, ) if errors.Is(err, sql.ErrNoRows) { diff --git a/server/internal/ha/store_integration_test.go b/server/internal/ha/store_integration_test.go index 527a452fc..944155fd4 100644 --- a/server/internal/ha/store_integration_test.go +++ b/server/internal/ha/store_integration_test.go @@ -17,22 +17,22 @@ func TestLeaseStoreAcquireAndRenewUseDatabaseTime(t *testing.T) { store, reader := leaseTestSurfaces(t) holder := uuid.New() - ownership, err := store.Acquire( + active, err := store.Acquire( t.Context(), databaseObservation(t, reader, "cluster-a", 41), holder, 2*time.Second, ) require.NoError(t, err) - require.Equal(t, holder, ownership.HolderID) - require.Equal(t, Token{WriterGeneration: 41, LeaseEpoch: 1}, ownership.Token) - require.WithinDuration(t, ownership.DatabaseTime.Add(2*time.Second), ownership.ExpiresAt, 50*time.Millisecond) + require.Equal(t, holder, active.HolderID) + require.Equal(t, Token{WriterGeneration: 41, LeaseEpoch: 1}, active.Token) + require.WithinDuration(t, active.DatabaseTime.Add(2*time.Second), active.ExpiresAt, 50*time.Millisecond) renewed, err := store.Renew( t.Context(), databaseObservation(t, reader, "cluster-a", 41), - ownership, + active, time.Second, ) require.NoError(t, err) - require.Equal(t, ownership.Token, renewed.Token) + require.Equal(t, active.Token, renewed.Token) require.WithinDuration(t, renewed.DatabaseTime.Add(time.Second), renewed.ExpiresAt, 50*time.Millisecond) } @@ -66,42 +66,6 @@ func TestRacingCoordinatorsProduceOneOwner(t *testing.T) { require.Equal(t, 1, countMatchingErrors(errs, ErrLeaseUnavailable)) } -func TestLeaseStoreCoordinatorDemotionRequiresExpiryBeforeSameWriterReacquisition(t *testing.T) { - store, reader := leaseTestSurfaces(t) - observed := databaseObservation(t, reader, "cluster-a", 41) - observer := &sequenceObserver{ - results: []observerResult{ - {observation: observed}, - {err: errors.New("DCS unavailable")}, - {observation: observed}, - {observation: observed}, - }, - } - config := coordinatorTestConfig() - config.LeaseDuration = shortIntegrationLeaseDuration - coordinator := newCoordinatorWithHolder(observer, store, config, uuid.New()) - - require.NoError(t, coordinator.step(t.Context())) - first := coordinator.Snapshot() - firstHolder := first.HolderID - - require.Error(t, coordinator.step(t.Context())) - require.Equal(t, StatePassive, coordinator.Snapshot().State) - require.NotEqual(t, firstHolder, coordinator.HolderID()) - - coordinator.ResumeAcquisition() - require.ErrorIs(t, coordinator.step(t.Context()), ErrLeaseUnavailable) - - waitForLeaseExpiry(t, Ownership{ - DatabaseTime: first.ExpiresAt.Add(-config.LeaseDuration), - ExpiresAt: first.ExpiresAt, - }) - require.NoError(t, coordinator.step(t.Context())) - second := coordinator.Snapshot() - require.Equal(t, StateActive, second.State) - require.Greater(t, second.Token.LeaseEpoch, first.Token.LeaseEpoch) -} - func TestPromotionAfterLostAsyncLeaseStateSupersedesUnexpiredLease(t *testing.T) { store, reader := leaseTestSurfaces(t) // This unexpired generation-41 row represents lease state acknowledged on @@ -202,7 +166,7 @@ func TestLeaseStoreSameHolderAfterExpiryAdvancesEpoch(t *testing.T) { func TestLeaseStoreRenewRequiresExactOwnership(t *testing.T) { store, reader := leaseTestSurfaces(t) - ownership, err := store.Acquire( + active, err := store.Acquire( t.Context(), databaseObservation(t, reader, "cluster-a", 41), uuid.New(), time.Minute, ) require.NoError(t, err) @@ -215,7 +179,7 @@ func TestLeaseStoreRenewRequiresExactOwnership(t *testing.T) { } for name, mutate := range tests { t.Run(name, func(t *testing.T) { - wrong := ownership + wrong := active mutate(&wrong) _, renewErr := store.Renew( t.Context(), @@ -241,24 +205,24 @@ func TestLeaseStoreRejectsDifferentWritableServerIdentity(t *testing.T) { _, err := store.Acquire(t.Context(), observed, uuid.New(), time.Minute) require.ErrorIs(t, err, ErrLeaseUnavailable) - ownership, err := store.Acquire( + active, err := store.Acquire( t.Context(), databaseObservation(t, reader, "cluster-a", 41), uuid.New(), time.Minute, ) require.NoError(t, err) - require.Equal(t, Token{WriterGeneration: 41, LeaseEpoch: 1}, ownership.Token) + require.Equal(t, Token{WriterGeneration: 41, LeaseEpoch: 1}, active.Token) } func TestLeaseStoreRejectsRenewalOnDifferentWritableServerIdentity(t *testing.T) { store, reader := leaseTestSurfaces(t) observed := databaseObservation(t, reader, "cluster-a", 41) - ownership, err := store.Acquire(t.Context(), observed, uuid.New(), time.Minute) + active, err := store.Acquire(t.Context(), observed, uuid.New(), time.Minute) require.NoError(t, err) observed.Timeline++ - _, err = store.Renew(t.Context(), observed, ownership, time.Minute) + _, err = store.Renew(t.Context(), observed, active, time.Minute) require.ErrorIs(t, err, ErrOwnershipLost) } @@ -284,7 +248,8 @@ func databaseObservation( func leaseTestSurfaces(t *testing.T) (*LeaseStore, postgresIdentityReader) { t.Helper() - queries := sqlc.New(testutil.GetTestDB(t)) + conn := testutil.GetTestDB(t) + queries := sqlc.New(conn) return NewLeaseStore(queries), queries } diff --git a/server/internal/ha/types.go b/server/internal/ha/types.go index 480d38d45..f2e91c8db 100644 --- a/server/internal/ha/types.go +++ b/server/internal/ha/types.go @@ -1,28 +1,18 @@ package ha import ( - "cmp" "context" "time" "github.com/google/uuid" ) -// Token totally orders Fleet ownership within one DCS cluster identity. +// Token identifies one Fleet ownership term within a DCS cluster. type Token struct { WriterGeneration int64 LeaseEpoch int64 } -// Compare orders ownership tokens lexicographically. A writer promotion always -// outranks every lease epoch from an older writer generation. -func (t Token) Compare(other Token) int { - return cmp.Or( - cmp.Compare(t.WriterGeneration, other.WriterGeneration), - cmp.Compare(t.LeaseEpoch, other.LeaseEpoch), - ) -} - // WriterObservation is a fail-closed binding between one DCS leader term and // the writable PostgreSQL server reached through Fleet's multi-host DSN. type WriterObservation struct { @@ -61,7 +51,7 @@ type ownershipStore interface { Renew( ctx context.Context, observed WriterObservation, - ownership Ownership, + active Ownership, duration time.Duration, ) (Ownership, error) } diff --git a/server/internal/handlers/command/handler.go b/server/internal/handlers/command/handler.go index 3719fc9bc..92962193c 100644 --- a/server/internal/handlers/command/handler.go +++ b/server/internal/handlers/command/handler.go @@ -282,7 +282,7 @@ func (h *Handler) GetCommandBatchLogBundle( if _, err := middleware.RequirePermission(ctx, authz.PermMinerDownloadLogs, authz.ResourceContext{}); err != nil { return nil, err } - resp, err := h.commandSvc.GetCommandBatchLogBundle(req.Msg.BatchIdentifier) + resp, err := h.commandSvc.GetCommandBatchLogBundle(ctx, req.Msg.BatchIdentifier) if err != nil { return nil, err } diff --git a/server/internal/infrastructure/db/with_transaction.go b/server/internal/infrastructure/db/with_transaction.go index b09f86d6c..bec7958fa 100644 --- a/server/internal/infrastructure/db/with_transaction.go +++ b/server/internal/infrastructure/db/with_transaction.go @@ -100,6 +100,31 @@ func WithTransactionNoResult(ctx context.Context, db *sql.DB, action func(q sqlc return withTransactionNoResultWithRetry(ctx, db, action, DefaultRetryConfig, firstTxOpts(opts)) } +// WithTransactionTimeout bounds the complete operation, including retries, and +// applies the same limit inside each transaction. +func WithTransactionTimeout[T any](ctx context.Context, db *sql.DB, timeout time.Duration, action func(q sqlc.Querier) (T, error)) (T, error) { + timeoutCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + return WithTransaction(timeoutCtx, db, func(q sqlc.Querier) (T, error) { + var zero T + if err := q.SetLocalTransactionTimeout( + timeoutCtx, + timeout.Milliseconds(), + ); err != nil { + return zero, err + } + return action(q) + }) +} + +func WithTransactionTimeoutNoResult(ctx context.Context, db *sql.DB, timeout time.Duration, action func(q sqlc.Querier) error) error { + _, err := WithTransactionTimeout(ctx, db, timeout, func(q sqlc.Querier) (struct{}, error) { + return struct{}{}, action(q) + }) + return err +} + func withTransactionNoResultWithRetry(ctx context.Context, db *sql.DB, action func(q sqlc.Querier) error, config RetryConfig, txOpts *sql.TxOptions) error { _, err := withTransactionWithRetry(ctx, db, func(sq sqlc.Querier) (any, error) { var emptyResult any diff --git a/server/internal/infrastructure/files/service.go b/server/internal/infrastructure/files/service.go index 12fe12f74..d0eaecaf7 100644 --- a/server/internal/infrastructure/files/service.go +++ b/server/internal/infrastructure/files/service.go @@ -95,6 +95,7 @@ type Service struct { commandArtifactCleanupInterval time.Duration mu sync.Mutex + bundleCreationMu sync.Mutex firmwareMetadataReuseMu sync.RWMutex checksumIndex map[string][]string // SHA-256 hex -> reuse-eligible file IDs firmwareChecksumByID map[string]string // fileID -> SHA-256 hex @@ -478,10 +479,25 @@ func (s *Service) GetBatchLogBundleFile(batchLogUUID string) (*FSFile, error) { return &FSFile{Filename: filename, Data: data}, nil } +func (s *Service) EnsureBatchLogBundle(batchLogUUID string) error { + if findBatchBundlePath(batchLogUUID) != "" { + return nil + } + + s.bundleCreationMu.Lock() + defer s.bundleCreationMu.Unlock() + + if findBatchBundlePath(batchLogUUID) != "" { + return nil + } + + _, err := s.bundleLogs(batchLogUUID) + return err +} + func (s *Service) DownloadLogsOnFinishedCallback(batchLogUUID string) func() error { return func() error { - _, err := s.bundleLogs(batchLogUUID) - if err != nil { + if err := s.EnsureBatchLogBundle(batchLogUUID); err != nil { return fleeterror.NewInternalErrorf("error bundling logs: %v", err) } diff --git a/server/internal/infrastructure/files/service_test.go b/server/internal/infrastructure/files/service_test.go index 47797f0f7..e48c4ce3b 100644 --- a/server/internal/infrastructure/files/service_test.go +++ b/server/internal/infrastructure/files/service_test.go @@ -267,6 +267,22 @@ func TestGetBatchLogBundleFile_NotReady(t *testing.T) { assert.Contains(t, err.Error(), "not available yet") } +func TestEnsureBatchLogBundleBuildsMissingBundle(t *testing.T) { + // Arrange + svc := setupService(t) + _, err := svc.SaveLogs("batch-lazy", "cc:dd:ee:ff:00:11", []string{"Time,Message", `2026-01-01T00:00:00Z,"data"`}) + require.NoError(t, err) + + // Act + err = svc.EnsureBatchLogBundle("batch-lazy") + + // Assert + require.NoError(t, err) + file, err := svc.GetBatchLogBundleFile("batch-lazy") + require.NoError(t, err) + assert.NotEmpty(t, file.Data) +} + // TestFindBatchBundlePath_PrefersZIPOverCSV verifies that when both a ZIP and a CSV // happen to exist for the same batch UUID the ZIP path is returned. func TestFindBatchBundlePath_PrefersZIPOverCSV(t *testing.T) { diff --git a/server/internal/infrastructure/queue/interface.go b/server/internal/infrastructure/queue/interface.go index c77d27342..4d524fb59 100644 --- a/server/internal/infrastructure/queue/interface.go +++ b/server/internal/infrastructure/queue/interface.go @@ -2,15 +2,10 @@ package queue import ( "context" - "errors" "github.com/block/proto-fleet/server/internal/domain/commandtype" ) -// ErrStale is returned when a MarkSuccess/MarkFailed/MarkPermanentlyFailed update -// finds 0 rows affected because the message is no longer in PROCESSING state (e.g., already reaped). -var ErrStale = errors.New("stale: message no longer PROCESSING") - type Message struct { ID int64 BatchLogUUID string @@ -37,22 +32,8 @@ type MessageQueue interface { // Dequeue retrieves and locks at most limit commands for processing. Dequeue(ctx context.Context, limit int32) ([]Message, error) - // MarkSuccess updates a command as successfully processed. - // Returns ErrStale if the message is no longer PROCESSING. - MarkSuccess(ctx context.Context, messageID int64) error - - // MarkFailed updates a command as failed with error info (may retry if under max retries). - // Returns ErrStale if the message is no longer PROCESSING. - MarkFailed(ctx context.Context, messageID int64, errorInfo string) error - - // MarkPermanentlyFailed marks a command as failed with no retries (for permanent errors like unsupported capabilities). - // Returns ErrStale if the message is no longer PROCESSING. - MarkPermanentlyFailed(ctx context.Context, messageID int64, errorInfo string) error - IsBatchFinished(ctx context.Context, commandBatchLogUUID string) (bool, error) - IsBatchProcessing(ctx context.Context, commandBatchLogUUID string) (bool, error) - // MaxFailureRetries returns the configured maximum number of retry attempts // before a message is permanently marked FAILED. MaxFailureRetries() int32 diff --git a/server/internal/infrastructure/queue/mocks/mock_message_queue.go b/server/internal/infrastructure/queue/mocks/mock_message_queue.go index d75383f53..5aa9240ca 100644 --- a/server/internal/infrastructure/queue/mocks/mock_message_queue.go +++ b/server/internal/infrastructure/queue/mocks/mock_message_queue.go @@ -100,63 +100,6 @@ func (mr *MockMessageQueueMockRecorder) IsBatchFinished(ctx, commandBatchLogUUID return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsBatchFinished", reflect.TypeOf((*MockMessageQueue)(nil).IsBatchFinished), ctx, commandBatchLogUUID) } -// IsBatchProcessing mocks base method. -func (m *MockMessageQueue) IsBatchProcessing(ctx context.Context, commandBatchLogUUID string) (bool, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "IsBatchProcessing", ctx, commandBatchLogUUID) - ret0, _ := ret[0].(bool) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// IsBatchProcessing indicates an expected call of IsBatchProcessing. -func (mr *MockMessageQueueMockRecorder) IsBatchProcessing(ctx, commandBatchLogUUID any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsBatchProcessing", reflect.TypeOf((*MockMessageQueue)(nil).IsBatchProcessing), ctx, commandBatchLogUUID) -} - -// MarkFailed mocks base method. -func (m *MockMessageQueue) MarkFailed(ctx context.Context, messageID int64, errorInfo string) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkFailed", ctx, messageID, errorInfo) - ret0, _ := ret[0].(error) - return ret0 -} - -// MarkFailed indicates an expected call of MarkFailed. -func (mr *MockMessageQueueMockRecorder) MarkFailed(ctx, messageID, errorInfo any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkFailed", reflect.TypeOf((*MockMessageQueue)(nil).MarkFailed), ctx, messageID, errorInfo) -} - -// MarkPermanentlyFailed mocks base method. -func (m *MockMessageQueue) MarkPermanentlyFailed(ctx context.Context, messageID int64, errorInfo string) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkPermanentlyFailed", ctx, messageID, errorInfo) - ret0, _ := ret[0].(error) - return ret0 -} - -// MarkPermanentlyFailed indicates an expected call of MarkPermanentlyFailed. -func (mr *MockMessageQueueMockRecorder) MarkPermanentlyFailed(ctx, messageID, errorInfo any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkPermanentlyFailed", reflect.TypeOf((*MockMessageQueue)(nil).MarkPermanentlyFailed), ctx, messageID, errorInfo) -} - -// MarkSuccess mocks base method. -func (m *MockMessageQueue) MarkSuccess(ctx context.Context, messageID int64) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "MarkSuccess", ctx, messageID) - ret0, _ := ret[0].(error) - return ret0 -} - -// MarkSuccess indicates an expected call of MarkSuccess. -func (mr *MockMessageQueueMockRecorder) MarkSuccess(ctx, messageID any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkSuccess", reflect.TypeOf((*MockMessageQueue)(nil).MarkSuccess), ctx, messageID) -} - // MaxFailureRetries mocks base method. func (m *MockMessageQueue) MaxFailureRetries() int32 { m.ctrl.T.Helper() diff --git a/server/internal/infrastructure/queue/service.go b/server/internal/infrastructure/queue/service.go index cf16a9f86..679fc4062 100644 --- a/server/internal/infrastructure/queue/service.go +++ b/server/internal/infrastructure/queue/service.go @@ -4,14 +4,12 @@ import ( "context" "database/sql" "encoding/json" - "fmt" - - "github.com/sqlc-dev/pqtype" "github.com/block/proto-fleet/server/generated/sqlc" "github.com/block/proto-fleet/server/internal/domain/commandtype" "github.com/block/proto-fleet/server/internal/domain/fleeterror" "github.com/block/proto-fleet/server/internal/infrastructure/db" + "github.com/block/proto-fleet/server/internal/runtimepolicy" ) type DatabaseMessageQueue struct { @@ -58,19 +56,29 @@ func (d DatabaseMessageQueue) EnqueueMany(ctx context.Context, commandBatchLogUU } func (d DatabaseMessageQueue) enqueueEncoded(ctx context.Context, commandBatchLogUUID string, commandType commandtype.Type, messages []encodedMessage) error { - return db.WithTransactionNoResult(ctx, d.conn, func(q sqlc.Querier) error { - for _, message := range messages { - err := q.CreateQueueMessage(ctx, sqlc.CreateQueueMessageParams{ - CommandBatchLogUuid: commandBatchLogUUID, - CommandType: commandType.String(), - DeviceID: message.deviceID, - Status: sqlc.QueueStatusEnumPENDING, - RetryCount: 0, - Payload: pqtype.NullRawMessage{RawMessage: message.payload, Valid: true}, - }) - if err != nil { - return fleeterror.NewInternalErrorf("failed to enqueue message: %v", err) - } + deviceIDs := make([]int64, len(messages)) + payloads := make([]string, len(messages)) + for i, message := range messages { + deviceIDs[i] = message.deviceID + payloads[i] = string(message.payload) + } + return db.WithTransactionTimeoutNoResult(ctx, d.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) error { + batchStatus, err := q.LockCommandBatch(ctx, commandBatchLogUUID) + if err != nil { + return fleeterror.NewInternalErrorf("failed to lock command batch: %v", err) + } + if batchStatus != sqlc.BatchStatusEnumPENDING { + return fleeterror.NewInternalErrorf("cannot enqueue messages for command batch in %s status", batchStatus) + } + + err = q.CreateQueueMessages(ctx, sqlc.CreateQueueMessagesParams{ + CommandBatchLogUuid: commandBatchLogUUID, + CommandType: commandType.String(), + DeviceIds: deviceIDs, + Payloads: payloads, + }) + if err != nil { + return fleeterror.NewInternalErrorf("failed to enqueue messages: %v", err) } return nil }) @@ -83,7 +91,7 @@ func (d DatabaseMessageQueue) Dequeue(ctx context.Context, limit int32) ([]Messa if d.config.DequeLimit > 0 { limit = min(limit, d.config.DequeLimit) } - messages, err := db.WithTransaction(ctx, d.conn, func(q sqlc.Querier) ([]Message, error) { + messages, err := db.WithTransactionTimeout(ctx, d.conn, runtimepolicy.CommandTransactionBound, func(q sqlc.Querier) ([]Message, error) { dbMessages, err := q.GetMessagesToProcess(ctx, sqlc.GetMessagesToProcessParams{ RetryCount: d.config.MaxFailureRetries, Limit: limit, @@ -129,84 +137,12 @@ func (d DatabaseMessageQueue) Dequeue(ctx context.Context, limit int32) ([]Messa return messages, nil } -func (d DatabaseMessageQueue) MarkSuccess(ctx context.Context, messageID int64) error { - updated, err := db.WithTransaction(ctx, d.conn, func(q sqlc.Querier) (bool, error) { - result, err := q.UpdateMessageStatus(ctx, sqlc.UpdateMessageStatusParams{ - ID: messageID, - Status: sqlc.QueueStatusEnumSUCCESS, - }) - if err != nil { - return false, fleeterror.NewInternalErrorf("failed to mark message as a success: %v", err) - } - rowsAffected, _ := result.RowsAffected() - return rowsAffected > 0, nil - }) - if err != nil { - return err - } - if !updated { - return fmt.Errorf("message %d: %w", messageID, ErrStale) - } - return nil -} - -func (d DatabaseMessageQueue) MarkFailed(ctx context.Context, messageID int64, errorInfo string) error { - updated, err := db.WithTransaction(ctx, d.conn, func(q sqlc.Querier) (bool, error) { - result, err := q.UpdateMessageAfterFailure(ctx, sqlc.UpdateMessageAfterFailureParams{ - ID: messageID, - RetryCount: d.config.MaxFailureRetries, - ErrorInfo: sql.NullString{String: errorInfo, Valid: true}, - }) - if err != nil { - return false, fleeterror.NewInternalErrorf("failed to mark message as failed: %v", err) - } - rowsAffected, _ := result.RowsAffected() - return rowsAffected > 0, nil - }) - if err != nil { - return err - } - if !updated { - return fmt.Errorf("message %d: %w", messageID, ErrStale) - } - return nil -} - -func (d DatabaseMessageQueue) MarkPermanentlyFailed(ctx context.Context, messageID int64, errorInfo string) error { - updated, err := db.WithTransaction(ctx, d.conn, func(q sqlc.Querier) (bool, error) { - result, err := q.UpdateMessagePermanentlyFailed(ctx, sqlc.UpdateMessagePermanentlyFailedParams{ - ID: messageID, - ErrorInfo: sql.NullString{String: errorInfo, Valid: true}, - }) - if err != nil { - return false, fleeterror.NewInternalErrorf("failed to mark message as permanently failed: %v", err) - } - rowsAffected, _ := result.RowsAffected() - return rowsAffected > 0, nil - }) - if err != nil { - return err - } - if !updated { - return fmt.Errorf("message %d: %w", messageID, ErrStale) - } - return nil -} - -type BatchStatusCheckFunc func(ctx context.Context, commandBatchLogID int64) (bool, error) - func (d DatabaseMessageQueue) IsBatchFinished(ctx context.Context, commandBatchLogUUID string) (bool, error) { return db.WithTransaction(ctx, d.conn, func(q sqlc.Querier) (bool, error) { return q.IsBatchFinished(ctx, commandBatchLogUUID) }) } -func (d DatabaseMessageQueue) IsBatchProcessing(ctx context.Context, commandBatchLogUUID string) (bool, error) { - return db.WithTransaction(ctx, d.conn, func(q sqlc.Querier) (bool, error) { - return q.IsBatchProcessing(ctx, commandBatchLogUUID) - }) -} - func (d DatabaseMessageQueue) MaxFailureRetries() int32 { return d.config.MaxFailureRetries } diff --git a/server/internal/infrastructure/queue/service_integration_test.go b/server/internal/infrastructure/queue/service_integration_test.go new file mode 100644 index 000000000..8c0041805 --- /dev/null +++ b/server/internal/infrastructure/queue/service_integration_test.go @@ -0,0 +1,69 @@ +package queue_test + +import ( + "database/sql" + "encoding/json" + "testing" + "time" + + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/block/proto-fleet/server/generated/sqlc" + "github.com/block/proto-fleet/server/internal/domain/commandtype" + "github.com/block/proto-fleet/server/internal/infrastructure/id" + "github.com/block/proto-fleet/server/internal/infrastructure/queue" + "github.com/block/proto-fleet/server/internal/testutil" +) + +func TestDatabaseMessageQueueEnqueueManyInsertsPerDevicePayloads(t *testing.T) { + if testing.Short() { + t.Skip("Skipping database integration test in short mode") + } + + // Arrange + cfg, err := testutil.GetTestConfig() + require.NoError(t, err) + dbService := testutil.NewDatabaseService(t, cfg) + user := dbService.CreateSuperAdminUser() + firstDevice := dbService.CreateDevice(user.OrganizationID, "proto") + secondDevice := dbService.CreateDevice(user.OrganizationID, "proto") + batchUUID := id.GenerateID() + commandType := commandtype.UpdateMiningPools + _, err = sqlc.New(dbService.DB).CreateCommandBatchLog(t.Context(), sqlc.CreateCommandBatchLogParams{ + Uuid: batchUUID, + Type: commandType.String(), + CreatedBy: user.DatabaseID, + CreatedAt: time.Now(), + Status: sqlc.BatchStatusEnumPENDING, + DevicesCount: 2, + Payload: pqtype.NullRawMessage{}, + OrganizationID: sql.NullInt64{Int64: user.OrganizationID, Valid: true}, + }) + require.NoError(t, err) + messageQueue := queue.NewDatabaseMessageQueue(&queue.Config{}, dbService.DB) + messages := []queue.EnqueueMessage{ + {DeviceID: firstDevice.DatabaseID, Payload: map[string]string{"worker_name": "first"}}, + {DeviceID: secondDevice.DatabaseID, Payload: map[string]string{"worker_name": "second"}}, + } + + // Act + err = messageQueue.EnqueueMany(t.Context(), batchUUID, commandType, messages) + + // Assert + require.NoError(t, err) + rows, err := sqlc.New(dbService.DB).GetQueueMessagesByBatch(t.Context(), batchUUID) + require.NoError(t, err) + gotPayloads := make(map[int64]map[string]string) + for _, row := range rows { + var decoded map[string]string + require.NoError(t, json.Unmarshal(row.Payload.RawMessage, &decoded)) + gotPayloads[row.DeviceID] = decoded + assert.Equal(t, sqlc.QueueStatusEnumPENDING, row.Status) + } + assert.Equal(t, map[int64]map[string]string{ + firstDevice.DatabaseID: {"worker_name": "first"}, + secondDevice.DatabaseID: {"worker_name": "second"}, + }, gotPayloads) +} diff --git a/server/internal/runtimejobs/group.go b/server/internal/runtimejobs/group.go index 97a123caa..38a38e0e6 100644 --- a/server/internal/runtimejobs/group.go +++ b/server/internal/runtimejobs/group.go @@ -203,8 +203,7 @@ func (g *Group) Stop(ctx context.Context) error { return nil } -// Abort immediately cancels work that must not survive ownership loss. -// Stop must follow to finish cleanup and make the group restartable. +// Abort immediately cancels work that cannot survive a fatal runtime error. func (g *Group) Abort() { abortJobs(g.jobs) } diff --git a/server/internal/runtimejobs/lifecycle.go b/server/internal/runtimejobs/lifecycle.go index ef1ecf9b8..f110f53e0 100644 --- a/server/internal/runtimejobs/lifecycle.go +++ b/server/internal/runtimejobs/lifecycle.go @@ -17,8 +17,7 @@ type Lifecycle interface { Stop(ctx context.Context) error } -// Aborter cancels work that must not outlive active ownership. Stop must still -// be called afterward to wait for cleanup and make the lifecycle restartable. +// Aborter immediately cancels work that cannot wait for graceful shutdown. type Aborter interface { Abort() } diff --git a/server/internal/runtimepolicy/timing.go b/server/internal/runtimepolicy/timing.go new file mode 100644 index 000000000..a00d76b3c --- /dev/null +++ b/server/internal/runtimepolicy/timing.go @@ -0,0 +1,7 @@ +package runtimepolicy + +import "time" + +const ( + CommandTransactionBound = 5 * time.Second +) diff --git a/server/sqlc/queries/command.sql b/server/sqlc/queries/command.sql index 8e000c574..9d1ccae89 100644 --- a/server/sqlc/queries/command.sql +++ b/server/sqlc/queries/command.sql @@ -22,24 +22,33 @@ INSERT INTO command_batch_log ( $8 ); --- name: MarkCommandBatchProcessing :exec -UPDATE command_batch_log +-- name: MarkCommandBatchProcessing :execrows +UPDATE command_batch_log AS batch SET status = 'PROCESSING', started_at = NOW() -WHERE uuid = $1; +WHERE batch.uuid = $1 + AND batch.status = 'PENDING' + AND EXISTS ( + SELECT 1 + FROM queue_message AS message + WHERE message.command_batch_log_uuid = batch.uuid + AND message.status = 'PROCESSING' + ); --- name: MarkCommandBatchFinished :exec +-- name: MarkCommandBatchFinished :execrows UPDATE command_batch_log SET status = 'FINISHED', finished_at = NOW() -WHERE uuid = $1; +WHERE uuid = $1 + AND status IN ('PENDING', 'PROCESSING'); --- name: MarkCommandBatchFinishedWithStartedAt :exec +-- name: MarkCommandBatchFinishedWithStartedAt :execrows UPDATE command_batch_log SET status = 'FINISHED', started_at = NOW(), finished_at = NOW() -WHERE uuid = $1; +WHERE uuid = $1 + AND status = 'PENDING'; -- name: UpsertCommandOnDeviceLog :exec -- PostgreSQL version using CTE for the subquery. @@ -136,6 +145,12 @@ SELECT FROM command_batch_log cbl WHERE cbl.uuid = $1; +-- name: LockCommandBatch :one +SELECT status +FROM command_batch_log +WHERE uuid = $1 +FOR UPDATE; + -- name: GetBatchHeaderForOrg :one -- Returns the batch header only if its recorded organization_id matches the -- caller's session org. Rows with organization_id IS NULL (pre-migration diff --git a/server/sqlc/queries/queue.sql b/server/sqlc/queries/queue.sql index ec13ed8f3..7e3b7d7bf 100644 --- a/server/sqlc/queries/queue.sql +++ b/server/sqlc/queries/queue.sql @@ -15,6 +15,42 @@ INSERT INTO queue_message ( $6 ); +-- name: CreateQueueMessages :exec +INSERT INTO queue_message ( + command_batch_log_uuid, + command_type, + device_id, + status, + retry_count, + payload +) +SELECT + sqlc.arg('command_batch_log_uuid'), + sqlc.arg('command_type'), + devices.device_id, + 'PENDING'::queue_status_enum, + 0, + payloads.payload::JSONB +FROM unnest(sqlc.arg('device_ids')::BIGINT[]) WITH ORDINALITY AS devices(device_id, ord) +JOIN unnest(sqlc.arg('payloads')::TEXT[]) WITH ORDINALITY AS payloads(payload, ord) USING (ord); + +-- name: GetQueueMessagesByBatch :many +SELECT id, device_id, status, error_info, payload +FROM queue_message +WHERE command_batch_log_uuid = $1; + +-- name: CountQueueMessagesByBatch :one +SELECT COUNT(*) +FROM queue_message +WHERE command_batch_log_uuid = $1; + +-- name: SetLocalTransactionTimeout :exec +SELECT set_config( + 'transaction_timeout', + sqlc.arg('timeout_milliseconds')::BIGINT::TEXT || 'ms', + TRUE +); + -- name: UpdateMessageStatus :execresult UPDATE queue_message SET status = $1, @@ -67,43 +103,83 @@ WHERE m.status = 'PENDING' ORDER BY m.created_at LIMIT $2; --- name: ReapStuckProcessingMessages :many -WITH stuck AS ( - SELECT m.id FROM queue_message m - WHERE m.status = 'PROCESSING' - AND m.updated_at < @cutoff - AND m.command_type != 'FirmwareUpdate' - LIMIT @reap_limit +-- name: ReapMessages :many +-- Startup reaping fails every PROCESSING row left by the previous process. +-- Periodic reaping only fails rows that exceeded their command-specific cutoff. +WITH candidates AS ( + SELECT message.id + FROM queue_message AS message + WHERE message.status = 'PROCESSING' + AND ( + sqlc.arg('include_fresh')::BOOLEAN + OR message.updated_at < CASE + WHEN message.command_type = 'FirmwareUpdate' + THEN sqlc.arg('firmware_cutoff')::TIMESTAMPTZ + ELSE sqlc.arg('cutoff')::TIMESTAMPTZ + END + ) + ORDER BY message.updated_at, message.id + LIMIT sqlc.arg('reap_limit') + FOR UPDATE ) -UPDATE queue_message -SET status = 'FAILED'::queue_status_enum, - error_info = 'reaped: stuck in PROCESSING beyond timeout', +UPDATE queue_message AS message +SET + status = 'FAILED'::queue_status_enum, + error_info = CASE + WHEN message.command_type = 'FirmwareUpdate' + THEN sqlc.arg('firmware_error_info')::TEXT + ELSE sqlc.arg('error_info')::TEXT + END, updated_at = CURRENT_TIMESTAMP -FROM stuck, device -WHERE queue_message.id = stuck.id - AND queue_message.status = 'PROCESSING' - AND queue_message.device_id = device.id -RETURNING queue_message.id, queue_message.device_id, queue_message.command_batch_log_uuid, - queue_message.error_info, queue_message.command_type, device.org_id; +FROM candidates, device +WHERE message.id = candidates.id + AND message.status = 'PROCESSING' + AND message.device_id = device.id +RETURNING + message.id, + message.device_id, + message.command_batch_log_uuid, + message.error_info, + message.command_type, + device.org_id, + device.site_id; + +-- name: ResetReapedFirmwareStatuses :exec +UPDATE device_status +SET + status = 'ACTIVE'::device_status_enum, + status_timestamp = CURRENT_TIMESTAMP, + status_details = NULL +WHERE device_id = ANY(sqlc.arg('device_ids')::BIGINT[]) + AND status IN ('UPDATING', 'REBOOT_REQUIRED'); --- name: ReapStuckFirmwareUpdateMessages :many -WITH stuck AS ( - SELECT m.id FROM queue_message m - WHERE m.status = 'PROCESSING' - AND m.updated_at < @cutoff - AND m.command_type = 'FirmwareUpdate' - LIMIT @reap_limit +-- name: FinishTerminalCommandBatches :execrows +WITH candidates AS MATERIALIZED ( + SELECT batch.id + FROM command_batch_log AS batch + WHERE batch.status IN ('PENDING', 'PROCESSING') + AND EXISTS ( + SELECT 1 + FROM queue_message AS message + WHERE message.command_batch_log_uuid = batch.uuid + ) + AND NOT EXISTS ( + SELECT 1 + FROM queue_message AS message + WHERE message.command_batch_log_uuid = batch.uuid + AND message.status IN ('PENDING', 'PROCESSING') + ) + ORDER BY batch.id + LIMIT sqlc.arg('finish_limit') + FOR UPDATE ) -UPDATE queue_message -SET status = 'FAILED'::queue_status_enum, - error_info = 'reaped: firmware update stuck in PROCESSING beyond timeout', - updated_at = CURRENT_TIMESTAMP -FROM stuck, device -WHERE queue_message.id = stuck.id - AND queue_message.status = 'PROCESSING' - AND queue_message.device_id = device.id -RETURNING queue_message.id, queue_message.device_id, queue_message.command_batch_log_uuid, - queue_message.error_info, queue_message.command_type, device.org_id; +UPDATE command_batch_log AS batch +SET + status = 'FINISHED'::batch_status_enum, + finished_at = CURRENT_TIMESTAMP +FROM candidates +WHERE batch.id = candidates.id + AND batch.status IN ('PENDING', 'PROCESSING'); -- name: IsBatchFinished :one SELECT @@ -114,13 +190,3 @@ SELECT END AS is_finished FROM queue_message WHERE command_batch_log_uuid = $1; - --- name: IsBatchProcessing :one -SELECT - CASE - WHEN COUNT(*) > 0 THEN true - ELSE false - END AS is_processing -FROM queue_message -WHERE command_batch_log_uuid = $1 - AND status = 'PROCESSING';