diff --git a/apps/api/cmd/gateway/main.go b/apps/api/cmd/gateway/main.go index f117ab1..595b762 100644 --- a/apps/api/cmd/gateway/main.go +++ b/apps/api/cmd/gateway/main.go @@ -36,14 +36,17 @@ func main() { } ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) - defer stop() db, err := store.ConnectWithMaxConns(ctx, cfg.DatabaseURL, cfg.DatabaseMaxConns) if err != nil { + stop() logger.Error("connect postgres failed", "error", err) os.Exit(1) } defer db.Close() + // Cancel background LISTEN and worker goroutines before closing the pool. + // This also prevents a startup panic from waiting forever in pgxpool.Close. + defer stop() if recovery, err := db.RecoverInterruptedRuntimeState(ctx); err != nil { logger.Error("recover interrupted runtime state failed", "error", err) os.Exit(1) diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index f21ea55..9f6aadf 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -23,6 +23,8 @@ const ( orphanedRiverJobScanInterval = 15 * time.Second orphanedRiverJobWorkerStaleAfter = 45 * time.Second orphanedRiverJobRecoveryBatchSize = 100 + asyncQueueStartupMaxAttempts = 3 + asyncQueueStartupRetryDelay = 2 * time.Second ) type asyncTaskArgs struct { @@ -83,19 +85,79 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs } func (s *Service) StartAsyncQueueWorker(ctx context.Context) { - if err := s.startRiverQueue(ctx, true); err != nil { + if err := retryAsyncQueueStartup( + ctx, + asyncQueueStartupMaxAttempts, + asyncQueueStartupRetryDelay, + func() error { return s.startRiverQueue(ctx, true) }, + func(attempt int, err error) { + s.logger.Warn("start river async queue failed; retrying", + "attempt", attempt, + "maxAttempts", asyncQueueStartupMaxAttempts, + "error", err, + ) + }, + ); err != nil { s.logger.Error("start river async queue failed", "error", err) panic(err) } } func (s *Service) StartAsyncQueueClient(ctx context.Context) { - if err := s.startRiverQueue(ctx, false); err != nil { + if err := retryAsyncQueueStartup( + ctx, + asyncQueueStartupMaxAttempts, + asyncQueueStartupRetryDelay, + func() error { return s.startRiverQueue(ctx, false) }, + func(attempt int, err error) { + s.logger.Warn("start river async queue client failed; retrying", + "attempt", attempt, + "maxAttempts", asyncQueueStartupMaxAttempts, + "error", err, + ) + }, + ); err != nil { s.logger.Error("start river async queue client failed", "error", err) panic(err) } } +func retryAsyncQueueStartup( + ctx context.Context, + maxAttempts int, + retryDelay time.Duration, + start func() error, + onRetry func(attempt int, err error), +) error { + if maxAttempts < 1 { + maxAttempts = 1 + } + var lastErr error + for attempt := 1; attempt <= maxAttempts; attempt++ { + if err := start(); err != nil { + lastErr = err + if attempt == maxAttempts { + break + } + if onRetry != nil { + onRetry(attempt, err) + } + timer := time.NewTimer(retryDelay) + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return ctx.Err() + case <-timer.C: + } + continue + } + return nil + } + return lastErr +} + func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error { driver := riverpgxv5.New(s.store.Pool()) migrator, err := rivermigrate.New(driver, nil) @@ -465,6 +527,28 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error { func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) { recoverJobs := func() { + runtimeRecovery, err := s.store.RecoverInterruptedRuntimeState(ctx) + if err != nil { + if ctx.Err() == nil { + s.logger.Warn("recover interrupted runtime state from worker failed", "error", err) + } + return + } + if runtimeRecovery.ReleasedConcurrencyLeases > 0 || + runtimeRecovery.ReleasedRateReservations > 0 || + runtimeRecovery.FailedAttempts > 0 || + runtimeRecovery.FailedTasks > 0 || + runtimeRecovery.RequeuedAsyncTasks > 0 || + runtimeRecovery.CleanedTaskAdmissions > 0 { + s.logger.Warn("interrupted runtime state recovered by worker", + "releasedConcurrencyLeases", runtimeRecovery.ReleasedConcurrencyLeases, + "releasedRateReservations", runtimeRecovery.ReleasedRateReservations, + "failedAttempts", runtimeRecovery.FailedAttempts, + "failedTasks", runtimeRecovery.FailedTasks, + "requeuedAsyncTasks", runtimeRecovery.RequeuedAsyncTasks, + "cleanedTaskAdmissions", runtimeRecovery.CleanedTaskAdmissions, + ) + } recovered, err := s.store.RecoverOrphanedAsyncRiverJobs( ctx, orphanedRiverJobWorkerStaleAfter, diff --git a/apps/api/internal/runner/queue_worker_manager_test.go b/apps/api/internal/runner/queue_worker_manager_test.go index b01e993..26c1865 100644 --- a/apps/api/internal/runner/queue_worker_manager_test.go +++ b/apps/api/internal/runner/queue_worker_manager_test.go @@ -21,6 +21,57 @@ type fakeAsyncExecutionClient struct { stopCancelled atomic.Int64 } +func TestRetryAsyncQueueStartupRecoversTransientFailures(t *testing.T) { + transient := errors.New("transient River startup failure") + var attempts atomic.Int32 + var retries atomic.Int32 + err := retryAsyncQueueStartup( + context.Background(), + 3, + time.Millisecond, + func() error { + if attempts.Add(1) < 3 { + return transient + } + return nil + }, + func(attempt int, err error) { + if attempt < 1 || !errors.Is(err, transient) { + t.Fatalf("unexpected retry callback attempt=%d err=%v", attempt, err) + } + retries.Add(1) + }, + ) + if err != nil { + t.Fatalf("retry startup: %v", err) + } + if attempts.Load() != 3 || retries.Load() != 2 { + t.Fatalf("startup attempts=%d retries=%d, want 3/2", attempts.Load(), retries.Load()) + } +} + +func TestRetryAsyncQueueStartupStopsWhenContextIsCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + var attempts atomic.Int32 + err := retryAsyncQueueStartup( + ctx, + 3, + time.Hour, + func() error { + attempts.Add(1) + cancel() + return errors.New("startup failed") + }, + nil, + ) + if !errors.Is(err, context.Canceled) { + t.Fatalf("retry startup error=%v, want context cancelled", err) + } + if attempts.Load() != 1 { + t.Fatalf("startup attempts=%d, want 1", attempts.Load()) + } +} + func (c *fakeAsyncExecutionClient) Start(context.Context) error { c.started.Add(1) return c.startErr