fix: handle stop during startup with safe state transitions and hard-exit

pkgs/edge-worker/deno.test.json

@@ -11,6 +11,7 @@
     "@std/crypto/timing-safe-equal": "jsr:@std/crypto@^0.224.0/timing-safe-equal",
     "@std/log": "jsr:@std/log@^0.224.13",
     "@std/testing/mock": "jsr:@std/testing@^0.224.0/mock",
+    "@std/testing/time": "jsr:@std/testing@^0.224.0/time",
     "postgres": "jsr:@oscar6echo/postgres@3.4.5-d",
     "@pgflow/core": "../core/src/index.ts",
     "@pgflow/dsl": "../dsl/src/index.ts",

pkgs/edge-worker/src/core/Worker.ts

@@ -1,6 +1,11 @@
 import type { IBatchProcessor, ILifecycle, WorkerBootstrap } from './types.js';
 import type { Logger } from '../platform/types.js';
 
+/** Initial delay before retrying a failed main-loop iteration. */
+const RETRY_DELAY_MS = 100;
+/** Maximum delay for consecutive failed main-loop iterations. */
+const MAX_RETRY_DELAY_MS = 5_000;
+
 export interface WorkerOptions {
   requestShutdown?: () => void;
   cleanup?: () => Promise<void>;
@@ -57,12 +62,17 @@ export class Worker {
   }
 
   private async runMainLoop() {
+    let consecutiveFailures = 0;
+
     try {
       while (this.isMainLoopActive) {
+        let iterationFailed = false;
+
         try {
           await this.lifecycle.sendHeartbeat();
         } catch (error: unknown) {
           this.logger.error(`Error sending heartbeat: ${error}`);
+          iterationFailed = true;
         }
 
         if (!this.isMainLoopActive) {
@@ -77,6 +87,19 @@ export class Worker {
           await this.batchProcessor.processBatch();
         } catch (error: unknown) {
           this.logger.error(`Error processing batch: ${error}`);
+          iterationFailed = true;
+        }
+
+        if (iterationFailed) {
+          consecutiveFailures++;
+          if (this.isMainLoopActive) {
+            await this.waitForRetry(
+              Math.min(RETRY_DELAY_MS * 2 ** (consecutiveFailures - 1), MAX_RETRY_DELAY_MS)
+            );
+          }
+        } else {
+          // Only a fully successful iteration resets the backoff.
+          consecutiveFailures = 0;
         }
       }
     } catch (error) {
@@ -85,6 +108,29 @@ export class Worker {
     }
   }
 
+  /**
+   * Abort-aware backoff wait: resolves after `ms`, or immediately when the
+   * worker's abort signal fires so stop() never waits out a retry delay.
+   */
+  private waitForRetry(ms: number): Promise<void> {
+    return new Promise<void>((resolve) => {
+      if (this.isAborted) {
+        resolve();
+        return;
+      }
+
+      const onAbort = () => {
+        clearTimeout(timer);
+        resolve();
+      };
+      const timer = setTimeout(() => {
+        this.abortController.signal.removeEventListener('abort', onAbort);
+        resolve();
+      }, ms);
+      this.abortController.signal.addEventListener('abort', onAbort, { once: true });
+    });
+  }
+
   onDeprecated(handler: () => void): void {
     this.deprecationHandler = handler;
   }
@@ -99,15 +145,22 @@ export class Worker {
       return;
     }
 
-    this.lifecycle.transitionToStopping();
-
     try {
       this.logDeprecation();
       this.requestShutdown?.();
       this.abortController.abort();
 
+      // Wait for startup to settle before transitioning: a stop during
+      // Starting must not attempt an invalid transition, and the abort
+      // above already keeps the main loop from processing any batch.
+      if (this.startupPromise) {
+        await this.startupPromise;
+      }
+
+      this.lifecycle.transitionToStopping();
+
+      this.logger.debug('-> Waiting for main loop to complete');
       try {
-        this.logger.debug('-> Waiting for main loop to complete');
         await this.mainLoopPromise;
       } catch (error) {
         this.logger.error(

pkgs/edge-worker/src/core/WorkerLifecycle.ts

@@ -47,10 +47,6 @@ export class WorkerLifecycle<IMessage extends Json> implements InternalLifecycle
   acknowledgeStop() {
     this.workerState.transitionTo(States.Stopping);
 
-    if (!this.workerRow) {
-      throw new Error('Cannot stop worker: workerRow not set');
-    }
-
     try {
       this.logger.debug('Acknowledging worker stop...');
 

pkgs/edge-worker/src/core/WorkerState.ts

@@ -22,7 +22,7 @@ export enum States {
 }
 
 export const Transitions: Record<States, States[]> = {
-  [States.Created]: [States.Starting],
+  [States.Created]: [States.Starting, States.Stopping],
   [States.Starting]: [States.Running],
   [States.Running]: [States.Deprecated, States.Stopping],
   [States.Deprecated]: [States.Stopping],

pkgs/edge-worker/src/flow/FlowWorkerLifecycle.ts

@@ -114,10 +114,6 @@ export class FlowWorkerLifecycle<TFlow extends AnyFlow> implements InternalLifec
   acknowledgeStop() {
     this.workerState.transitionTo(States.Stopping);
 
-    if (!this.workerRow) {
-      throw new Error('Cannot stop worker: workerRow not set');
-    }
-
     try {
       this.logger.debug('Acknowledging worker stop...');
       this.workerState.transitionTo(States.Stopped);

pkgs/edge-worker/src/flow/StepTaskPoller.ts

@@ -115,7 +115,9 @@ export class StepTaskPoller<TFlow extends AnyFlow>
       return taskWithMessages;
     } catch (err: unknown) {
       this.logger.error(`Error in two-phase polling for flow tasks: ${err}`);
-      return [];
+      // Rethrow so Worker can distinguish a failed poll (which drives its
+      // retry backoff) from an empty successful poll.
+      throw err;
     }
   }
 

pkgs/edge-worker/src/platform/ProcessPlatformAdapter.ts

@@ -43,6 +43,7 @@ export class ProcessPlatformAdapter implements PlatformAdapter<SupabaseResources
   private cleanupPromise: Promise<void> | null = null;
   private signalHandlersRegistered = false;
   private signalCount = 0;
+  private startupCompleted = false;
   private readonly signalHandler = () => this.handleSignal();
 
   constructor(
@@ -75,7 +76,12 @@ export class ProcessPlatformAdapter implements PlatformAdapter<SupabaseResources
     this.registerSignalHandlers();
     this.startupPromise ??= this.performStartWorker(createWorkerFn).catch(
       async (error) => {
-        await this.cleanup();
+        try {
+          await this.cleanup();
+        } catch (cleanupError) {
+          // A cleanup failure must not replace the startup failure.
+          this.logger.error('Cleanup after startup failure failed', cleanupError);
+        }
         throw error;
       }
     );
@@ -141,6 +147,7 @@ export class ProcessPlatformAdapter implements PlatformAdapter<SupabaseResources
       workerId,
       startMode: 'process',
     });
+    this.startupCompleted = true;
   }
 
   private registerSignalHandlers(): void {
@@ -164,6 +171,7 @@ export class ProcessPlatformAdapter implements PlatformAdapter<SupabaseResources
   private async performStopWorker(): Promise<void> {
     this.requestShutdown();
 
+    let operationError: { error: unknown } | null = null;
     try {
       await this.startupPromise;
       if (this.worker) {
@@ -172,8 +180,22 @@ export class ProcessPlatformAdapter implements PlatformAdapter<SupabaseResources
       if (this.workerId) {
         await this.queries.markWorkerStopped(this.workerId);
       }
-    } finally {
+    } catch (error) {
+      operationError = { error };
+    }
+
+    try {
       await this.cleanup();
+    } catch (cleanupError) {
+      if (!operationError) {
+        throw cleanupError;
+      }
+      // A cleanup failure must not replace the drain or marking failure.
+      this.logger.error('Cleanup during shutdown failed', cleanupError);
+    }
+
+    if (operationError) {
+      throw operationError.error;
     }
   }
 
@@ -212,6 +234,15 @@ export class ProcessPlatformAdapter implements PlatformAdapter<SupabaseResources
       this.deps.exit(1);
     }
 
+    if (!this.startupCompleted) {
+      // No batch loop is ready, so no accepted task needs draining, and the
+      // hung bootstrap may be exactly what is blocking termination. Hard-exit
+      // now; the OS closes process resources after exit.
+      this.requestShutdown();
+      this.deps.setExitCode(0);
+      this.deps.exit(0);
+    }
+
     await this.gracefulExit();
   }