feat: add worker tracking and graceful shutdown for edge functions (#513)

# Implement Worker Tracking for Improved Reliability

This PR adds functionality to track edge function workers more effectively, improving the reliability of worker management:

1. Added two new methods to the `Queries` class:
   - `trackWorkerFunction()` - Registers an edge function for monitoring by the `ensure_workers()` cron job
   - `markWorkerStopped()` - Marks a worker as stopped during graceful shutdown

2. Updated worker lifecycle implementations to call `trackWorkerFunction()` early in the startup process, which:
   - Sets `last_invoked_at` to prevent cron from pinging during startup (debounce)
   - Ensures proper monitoring of the worker by the cron job

3. Enhanced the shutdown process in `SupabasePlatformAdapter`:
   - Replaced the previous approach of spawning a new edge function
   - Now signals worker death to the cron job by setting `stopped_at`
   - Allows immediate detection of worker termination

4. Added comprehensive unit tests for the new functionality

These changes improve worker reliability by providing better signaling between workers and the monitoring system, ensuring faster recovery from failures and more efficient worker management.

Author: jumski

pkgs/edge-worker/deno.lock

@@ -73,4 +73,25 @@ export class Queries {
     `;
     return result.result;
   }
+
+  /**
+   * Registers an edge function for monitoring by ensure_workers() cron.
+   * Called by workers on startup. Sets last_invoked_at to prevent cron from
+   * pinging during startup (debounce).
+   */
+  async trackWorkerFunction(functionName: string): Promise<void> {
+    await this.sql`
+      SELECT pgflow.track_worker_function(${functionName})
+    `;
+  }
+
+  /**
+   * Marks a worker as stopped for graceful shutdown signaling.
+   * Called by workers on beforeunload to allow ensure_workers() to detect death immediately.
+   */
+  async markWorkerStopped(workerId: string): Promise<void> {
+    await this.sql`
+      SELECT pgflow.mark_worker_stopped(${workerId}::uuid)
+    `;
+  }
 }

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

@@ -29,6 +29,10 @@ export class WorkerLifecycle<IMessage extends Json> implements ILifecycle {
   async acknowledgeStart(workerBootstrap: WorkerBootstrap): Promise<void> {
     this.workerState.transitionTo(States.Starting);
 
+    // Register this edge function for monitoring by ensure_workers() cron.
+    // Must be called early to set last_invoked_at (debounce) before heartbeat timeout.
+    await this.queries.trackWorkerFunction(workerBootstrap.edgeFunctionName);
+
     this.logger.debug(`Ensuring queue '${this.queue.queueName}' exists...`);
     await this.queue.safeCreate();
 

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

@@ -41,6 +41,10 @@ export class FlowWorkerLifecycle<TFlow extends AnyFlow> implements ILifecycle {
     // Store workerId for supplier pattern
     this._workerId = workerBootstrap.workerId;
 
+    // Register this edge function for monitoring by ensure_workers() cron.
+    // Must be called early to set last_invoked_at (debounce) before heartbeat timeout.
+    await this.queries.trackWorkerFunction(workerBootstrap.edgeFunctionName);
+
     // Compile/verify flow as part of Starting (before registering worker)
     if (this.ensureCompiledOnStartup) {
       await this.ensureFlowCompiled();

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

@@ -10,6 +10,7 @@ import {
   resolveConnectionString,
   resolveSqlConnection,
 } from './resolveConnection.js';
+import { Queries } from '../core/Queries.js';
 
 /**
  * Supabase Edge Runtime type (without global augmentation to comply with JSR)
@@ -47,6 +48,7 @@ export class SupabasePlatformAdapter implements PlatformAdapter<SupabaseResource
   private _platformResources: SupabaseResources;
   private validatedEnv: SupabaseEnv;
   private _connectionString: string | undefined;
+  private queries: Queries;
 
   // Logging factory with dynamic workerId support
   private loggingFactory = createLoggingFactory();
@@ -79,6 +81,9 @@ export class SupabasePlatformAdapter implements PlatformAdapter<SupabaseResource
       sql: resolveSqlConnection(env, options),
       supabase: createServiceSupabaseClient(this.validatedEnv)
     };
+
+    // Create Queries instance for shutdown handler
+    this.queries = new Queries(this._platformResources.sql);
   }
 
   /**
@@ -158,35 +163,6 @@ export class SupabasePlatformAdapter implements PlatformAdapter<SupabaseResource
     return this._platformResources;
   }
 
-  private async spawnNewEdgeFunction(): Promise<void> {
-    if (!this.edgeFunctionName) {
-      throw new Error('functionName cannot be null or empty');
-    }
-
-    const supabaseUrl = this.validatedEnv.SUPABASE_URL;
-
-    this.logger.debug('Spawning a new Edge Function...');
-
-    const response = await fetch(
-      `${supabaseUrl}/functions/v1/${this.edgeFunctionName}`,
-      {
-        method: 'POST',
-        headers: {
-          Authorization: `Bearer ${this.validatedEnv.SUPABASE_ANON_KEY}`,
-          'Content-Type': 'application/json',
-        },
-      }
-    );
-
-    this.logger.debug('Edge Function spawned successfully!');
-
-    if (!response.ok) {
-      throw new Error(
-        `Edge function returned non-OK status: ${response.status} ${response.statusText}`
-      );
-    }
-  }
-
   private extractFunctionName(req: Request): string {
     return new URL(req.url).pathname.replace(/^\/+|\/+$/g, '');
   }
@@ -196,7 +172,10 @@ export class SupabasePlatformAdapter implements PlatformAdapter<SupabaseResource
       this.logger.debug('Shutting down...');
 
       if (this.worker) {
-        await this.spawnNewEdgeFunction();
+        // Signal death to ensure_workers() cron by setting stopped_at.
+        // This allows the cron to immediately ping for a replacement worker.
+        const workerId = this.validatedEnv.SB_EXECUTION_ID;
+        await this.queries.markWorkerStopped(workerId);
       }
 
       await this.stopWorker();

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

@@ -0,0 +1,28 @@
+import { EdgeWorker } from '@pgflow/edge-worker';
+import { crypto } from 'jsr:@std/crypto';
+import { sql } from '../utils.ts';
+
+async function cpuIntensiveTask(payload: { debug?: boolean }) {
+  let data = new TextEncoder().encode('burn');
+  const timeId = `stopped_at_test_${Math.random()}`;
+
+  if (payload.debug) {
+    console.time(timeId);
+  }
+
+  for (let i = 0; i < 10000; i++) {
+    data = new Uint8Array(await crypto.subtle.digest('SHA-256', data));
+  }
+
+  if (payload.debug) {
+    console.timeEnd(timeId);
+    console.log(
+      '[stopped_at_test] last_val = ',
+      await sql`SELECT nextval('stopped_at_test_seq')`
+    );
+  } else {
+    await sql`SELECT nextval('stopped_at_test_seq')`;
+  }
+}
+
+EdgeWorker.start(cpuIntensiveTask, { queueName: 'stopped_at_test' });