@pgflow/edge-worker 0.0.5-prealpha.0

package/src/queue/createQueueWorker.ts

@@ -0,0 +1,145 @@
+import { ExecutionController } from '../core/ExecutionController.ts';
+import { MessageExecutor } from './MessageExecutor.ts';
+import { Queries } from '../core/Queries.ts';
+import { Queue } from './Queue.ts';
+import { ReadWithPollPoller } from './ReadWithPollPoller.ts';
+import type { Json } from '../core/types.ts';
+import type { PgmqMessageRecord } from './types.ts';
+import { Worker } from '../core/Worker.ts';
+import postgres from 'postgres';
+import { WorkerLifecycle } from '../core/WorkerLifecycle.ts';
+import { BatchProcessor } from '../core/BatchProcessor.ts';
+
+/**
+ * Configuration for the queue worker
+ */
+export type QueueWorkerConfig = {
+  /**
+   * PostgreSQL connection string.
+   * If not provided, it will be read from the EDGE_WORKER_DB_URL environment variable.
+   */
+  connectionString?: string;
+
+  /**
+   * Name of the queue to poll for messages
+   * @default 'tasks'
+   */
+  queueName?: string;
+
+  /**
+   * How many tasks are processed at the same time
+   * @default 10
+   */
+  maxConcurrent?: number;
+
+  /**
+   * How many connections to the database are opened
+   * @default 4
+   */
+  maxPgConnections?: number;
+
+  /**
+   * In-worker polling interval in seconds
+   * @default 5
+   */
+  maxPollSeconds?: number;
+
+  /**
+   * In-database polling interval in milliseconds
+   * @default 200
+   */
+  pollIntervalMs?: number;
+
+  /**
+   * How long to wait before retrying a failed job in seconds
+   * @default 5
+   */
+  retryDelay?: number;
+
+  /**
+   * How many times to retry a failed job
+   * @default 5
+   */
+  retryLimit?: number;
+
+  /**
+   * How long a job is invisible after reading in seconds.
+   * If not successful, will reappear after this time.
+   * @default 3
+   */
+  visibilityTimeout?: number;
+
+  /**
+   * Batch size for polling messages
+   * @default 10
+   */
+  batchSize?: number;
+
+  /**
+   * Optional SQL client instance
+   */
+  sql?: postgres.Sql;
+};
+
+/**
+ * Creates a new Worker instance for processing queue messages.
+ *
+ * @param handler - The message handler function that processes each message from the queue
+ * @param config - Configuration options for the worker
+ * @returns A configured Worker instance ready to be started
+ */
+export function createQueueWorker<TPayload extends Json>(
+  handler: (message: TPayload) => Promise<void> | void,
+  config: QueueWorkerConfig
+): Worker {
+  type QueueMessage = PgmqMessageRecord<TPayload>;
+
+  const abortController = new AbortController();
+  const abortSignal = abortController.signal;
+
+  // Use provided SQL connection if available, otherwise create one from connection string
+  const sql =
+    config.sql ||
+    postgres(config.connectionString || '', {
+      max: config.maxPgConnections,
+      prepare: false,
+    });
+
+  const queue = new Queue<TPayload>(sql, config.queueName || 'tasks');
+  const queries = new Queries(sql);
+
+  const lifecycle = new WorkerLifecycle<TPayload>(queries, queue);
+
+  const executorFactory = (record: QueueMessage, signal: AbortSignal) => {
+    return new MessageExecutor(
+      queue,
+      record,
+      handler,
+      signal,
+      config.retryLimit || 5,
+      config.retryDelay || 3
+    );
+  };
+
+  const poller = new ReadWithPollPoller(queue, abortSignal, {
+    batchSize: config.batchSize || config.maxConcurrent || 10,
+    maxPollSeconds: config.maxPollSeconds || 5,
+    pollIntervalMs: config.pollIntervalMs || 200,
+    visibilityTimeout: config.visibilityTimeout || 3,
+  });
+
+  const executionController = new ExecutionController<QueueMessage>(
+    executorFactory,
+    abortSignal,
+    {
+      maxConcurrent: config.maxConcurrent || 10,
+    }
+  );
+  const batchProcessor = new BatchProcessor<QueueMessage>(
+    executionController,
+    poller,
+    abortSignal
+  );
+
+  return new Worker(batchProcessor, lifecycle, sql);
+}

package/src/queue/types.ts

@@ -0,0 +1,14 @@
+import type { Json, IMessage } from "../core/types.ts";
+
+/**
+ * Fields are nullable because types in postgres does not allow NOT NULL,
+ * but all those values except `message` come from queue table columns,
+ * which are explicitely marked as NOT NULL.
+ */
+export interface PgmqMessageRecord<TPayload extends Json | null = Json> extends IMessage {
+  read_ct: number;
+  enqueued_at: string;
+  vt: string;
+  message: TPayload;
+}
+

package/src/spawnNewEdgeFunction.ts

@@ -0,0 +1,33 @@
+import { getLogger } from './core/Logger.ts';
+
+// @ts-ignore - TODO: fix the types
+const SUPABASE_URL = Deno.env.get('SUPABASE_URL') as string;
+const SUPABASE_ANON_KEY = Deno.env.get('SUPABASE_ANON_KEY') as string;
+
+const logger = getLogger('spawnNewEdgeFunction');
+
+export default async function spawnNewEdgeFunction(
+  functionName: string = 'edge-worker'
+): Promise<void> {
+  if (!functionName) {
+    throw new Error('functionName cannot be null or empty');
+  }
+
+  logger.debug('Spawning a new Edge Function...');
+
+  const response = await fetch(`${SUPABASE_URL}/functions/v1/${functionName}`, {
+    method: 'POST',
+    headers: {
+      Authorization: `Bearer ${SUPABASE_ANON_KEY}`,
+      'Content-Type': 'application/json',
+    }
+  });
+
+  logger.debug('Edge Function spawned successfully!');
+
+  if (!response.ok) {
+    throw new Error(
+      `Edge function returned non-OK status: ${response.status} ${response.statusText}`
+    );
+  }
+}

package/supabase/call

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# get env vars from supabase status and eval them for usage
+supabase_env=$(supabase status --output env 2>/dev/null)
+
+if [ $? -eq 0 ]; then
+  eval "$supabase_env"
+fi
+
+if [ -f ".env" ]; then
+  source .env
+fi
+
+function_name="$1"
+data="$2"
+
+API_URL=http://localhost:50321
+curl \
+  --request POST \
+  "${API_URL}/functions/v1/${function_name}" \
+  --header 'Authorization: Bearer '${ANON_KEY}'' \
+  --header 'Content-Type: application/json' \
+  --data-raw "$data"

package/supabase/cli

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+pnpm supabase --workdir pkgs/edge-worker $@

package/supabase/config.toml

@@ -0,0 +1,42 @@
+project_id = "edge-worker"
+
+[api]
+enabled = true
+port = 50321
+
+[db]
+port = 50322
+shadow_port = 50320
+major_version = 15
+
+[db.pooler]
+enabled = true
+port = 50329
+pool_mode = "transaction"
+default_pool_size = 200
+max_client_conn = 250
+
+[db.seed]
+enabled = true
+sql_paths = ['./seed.sql']
+
+[edge_runtime]
+enabled = true
+policy = "per_worker"
+inspector_port = 8083
+
+# disable unused features
+[realtime]
+enabled = false
+[studio]
+enabled = false
+[inbucket]
+enabled = false
+[analytics]
+enabled = false
+[storage]
+enabled = false
+[auth]
+enabled = false
+
+

package/supabase/functions/cpu_intensive/index.ts

@@ -0,0 +1,20 @@
+import { EdgeWorker } from '../_src/EdgeWorker.ts';
+import { crypto } from 'jsr:@std/crypto';
+import { sql } from '../utils.ts';
+
+async function cpuIntensiveTask() {
+  let data = new TextEncoder().encode('burn');
+  const timeId = `cpu_intensive_${Math.random()}`;
+  console.time(timeId);
+  for (let i = 0; i < 10000; i++) {
+    data = new Uint8Array(await crypto.subtle.digest('SHA-256', data));
+  }
+  console.timeEnd(timeId);
+
+  console.log(
+    '[cpu_intensive] last_val = ',
+    await sql`SELECT nextval('test_seq')`
+  );
+}
+
+EdgeWorker.start(cpuIntensiveTask, { queueName: 'cpu_intensive' });

package/supabase/functions/creating_queue/index.ts

@@ -0,0 +1,5 @@
+import { EdgeWorker } from '../_src/EdgeWorker.ts';
+
+EdgeWorker.start(console.log, {
+  queueName: 'creating_queue',
+});

package/supabase/functions/failing_always/index.ts

@@ -0,0 +1,13 @@
+import { EdgeWorker } from '../_src/EdgeWorker.ts';
+
+function failingAlways() {
+  console.log('(╯°□°)╯︵ ┻━┻');
+  throw new Error('(╯°□°)╯︵ ┻━┻');
+}
+
+EdgeWorker.start(failingAlways, {
+  queueName: 'failing_always',
+  retryLimit: 2,
+  retryDelay: 2,
+  maxPollSeconds: 1,
+});

package/supabase/functions/increment_sequence/index.ts

@@ -0,0 +1,14 @@
+import { EdgeWorker } from '../_src/EdgeWorker.ts';
+import { sql } from '../utils.ts';
+
+// await sql`CREATE SEQUENCE IF NOT EXISTS test_seq`;
+// await sql`SELECT pgmq.create('increment_sequence')`;
+
+async function incrementCounter() {
+  console.log(
+    '[increment_sequence] next_seq =',
+    await sql`SELECT nextval('test_seq')`
+  );
+}
+
+EdgeWorker.start(incrementCounter, { queueName: 'increment_sequence' });

package/supabase/functions/max_concurrency/index.ts

@@ -0,0 +1,17 @@
+import { EdgeWorker } from '../_src/EdgeWorker.ts';
+import { sleep, sql } from '../utils.ts';
+
+async function incrementSeq() {
+  await sleep(50);
+
+  console.log(
+    '[max_concurrency] last_val =',
+    await sql`SELECT nextval('test_seq')`
+  );
+}
+
+EdgeWorker.start(incrementSeq, {
+  queueName: 'max_concurrency',
+  maxConcurrent: 10,
+  maxPgConnections: 4,
+});

package/supabase/functions/serial_sleep/index.ts

@@ -0,0 +1,16 @@
+import { EdgeWorker } from '../_src/EdgeWorker.ts';
+import { sql, sleep } from '../utils.ts';
+
+const sleep1s = async () => {
+  console.time('Task time');
+  const lastVal = await sql`SELECT nextval('test_seq')`;
+  console.log('[serial_sleep] lastVal =', lastVal);
+  await sleep(1000);
+  console.timeEnd('Task time');
+};
+
+EdgeWorker.start(sleep1s, {
+  queueName: 'serial_sleep',
+  maxConcurrent: 1,
+  visibilityTimeout: 5, // higher than the delay()
+});

package/supabase/functions/utils.ts

@@ -0,0 +1,13 @@
+import postgres from 'postgres';
+import { delay } from '@std/async';
+
+const EDGE_WORKER_DB_URL = Deno.env.get('EDGE_WORKER_DB_URL')!;
+console.log('EDGE_WORKER_DB_URL', EDGE_WORKER_DB_URL);
+
+export const sql = postgres(EDGE_WORKER_DB_URL, { prepare: false });
+
+export const sleep = delay;
+
+export function randomInt(min: number, max: number) {
+  return Math.floor(Math.random() * (max - min + 1)) + min;
+}

package/supabase/seed.sql

@@ -0,0 +1,2 @@
+select vault.create_secret('http://host.docker.internal:50321', 'app_url');
+select vault.create_secret('eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZS1kZW1vIiwicm9sZSI6ImFub24iLCJleHAiOjE5ODM4MTI5OTZ9.CRXP1A7WOeoJeXxjNni43kdQwgnWNReilDMblYTn_I0', 'supabase_anon_key');

package/tests/db/compose.yaml

@@ -0,0 +1,20 @@
+services:
+  db:
+    # image: postgres:15.8-alpine
+    # image: supabase/postgres:15.8.1.039
+    # image: supabase/postgres:15.1.0.37
+
+    image: supabase/postgres:15.8.1.020
+    # image: supabase/postgres:15.8.1.018
+    ports:
+      - '5432:5432'
+    volumes:
+      - ./migrations/edge_worker.sql:/docker-entrypoint-initdb.d/migrations/950_edge_worker.sql
+    environment:
+      POSTGRES_DB: postgres
+      # commented out because of the problems with supabase_admin permissions
+      # see: https://github.com/supabase/postgres/issues/1219#issuecomment-2362955730
+      # POSTGRES_USER: postgres
+      POSTGRES_PASSWORD: postgres
+      POSTGRES_HOST: /var/run/postgresql
+      POSTGRES_PORT: 5432

package/tests/db.ts

@@ -0,0 +1,71 @@
+import postgres from 'postgres';
+
+function createSql(dbUrl: string) {
+  return postgres(dbUrl, {
+    prepare: false,
+    onnotice(_: unknown) {
+      // no-op to silence notices
+    },
+  });
+}
+
+export function withTransaction(
+  callback: (sql: postgres.Sql) => Promise<unknown>
+) {
+  const dbUrl = `postgresql://supabase_admin:postgres@localhost:5432/postgres`;
+  const localSql = createSql(dbUrl);
+
+  return async () => {
+    try {
+      console.log('calling callback');
+
+      let callbackError: unknown = null;
+
+      await localSql.begin(async (sql: postgres.Sql) => {
+        // Add no-op end() method to transaction-local sql
+        const wrappedSql = Object.assign(sql, {
+          end: async () => {
+            /* no-op */
+          },
+        });
+
+        try {
+          await callback(wrappedSql);
+        } catch (error) {
+          callbackError = error;
+        } finally {
+          console.log('Rolling back transaction');
+          // Using ROLLBACK AND CHAIN to avoid "no transaction in progress" warning
+          await sql`ROLLBACK AND CHAIN`;
+        }
+      });
+
+      console.log('callback called');
+
+      if (callbackError) {
+        throw callbackError;
+      }
+    } catch (err) {
+      console.error('Error in withTransaction:', err);
+      throw err;
+    } finally {
+      console.log('Closing connection');
+      await localSql.end();
+    }
+  };
+}
+
+export function withPgNoTransaction(
+  callback: (sql: postgres.Sql) => Promise<unknown>
+) {
+  const dbUrl = 'postgresql://supabase_admin:postgres@localhost:5432/postgres';
+  const sql = createSql(dbUrl);
+
+  return async () => {
+    try {
+      await callback(sql);
+    } finally {
+      await sql.end();
+    }
+  };
+}

package/tests/e2e/README.md

@@ -0,0 +1,54 @@
+# E2E Testing Strategy
+
+We use a real Supabase instance running locally (in `supabase/`) to perform end-to-end testing of the entire EdgeWorker stack. The tests interact with purpose-built test workers that simulate different behaviors:
+
+- Different queues with specific behaviors
+- Workers that always fail
+- Workers with varying retry patterns
+- etc.
+
+Lot of workers increment `test_seq` as a mean to identify the completion of a job.
+This also allows to identify how many times handlers were called if any retries were attempted.
+
+The test flow is straightforward:
+
+1. Put messages on specific queues
+2. Worker calls handlers, handlers increment `test_seq`
+3. `await` until `test_seq` was incremented to expected value (number of messages)
+4. Assert that workers behaved as expected (retries, failures, etc.)
+
+This approach lets us verify the entire stack from message enqueueing through worker processing, retries, and completion.
+
+## Core Test Scenarios Needed
+
+#### Glossary
+
+- `worker` - instance of given worker edge function that is subject to CPU and memory limits and can be killed
+- `worker function` - edge function within supabase app that uses EdgeWorker instead of serving requests
+- `queue` - pgmq queue that workers can pull from
+- `message` - PGMQ `message_record` that contains metadata (`msg_id`, `read_ct`, `vt`) and payload (`message JSONB`)
+
+### [ ] Happy Path
+
+- [x] Worker picks messages from queue
+- [ ] Worker calls handler function with each message
+- [x] Worker can process big amounts of messages (restarts itself when CPU clock limit hits)
+- [x] Different worker functions can pull from different queues
+- [ ] Different worker functions can pull from the same queue
+
+### [x] Worker Lifecycle
+
+- [x] Worker registers on start
+- [x] Worker sends heartbeats every 5s
+- [x] Worker updates function_name with heartbeat
+
+### [ ] Retries & Failures
+
+- [ ] Worker retries failed jobs n-times and succeeds
+- [ ] Worker uses exponential backoff for each subsequent retry
+- [x] Worker uses proper number of retries for each job
+- [x] Worker archives jobs that will not be retried
+
+### [x] Concurrency
+
+- [x] Worker respects maxConcurrent and processes messages in serial when set to 1

package/tests/e2e/_helpers.ts

@@ -0,0 +1,135 @@
+import { sql } from '../sql.ts';
+import { delay } from '@std/async';
+import ProgressBar from 'jsr:@deno-library/progress';
+import { dim } from 'https://deno.land/std@0.224.0/fmt/colors.ts';
+
+interface WaitForOptions {
+  pollIntervalMs?: number;
+  timeoutMs?: number;
+  description?: string;
+}
+
+export function log(message: string, ...args: unknown[]) {
+  console.log(dim(` -> ${message}`), ...args);
+}
+
+export async function waitFor<T>(
+  predicate: () => Promise<T | false>,
+  options: WaitForOptions = {}
+): Promise<T> {
+  const {
+    pollIntervalMs = 250,
+    timeoutMs = 30000,
+    description = 'condition',
+  } = options;
+
+  const startTime = Date.now();
+
+  while (true) {
+    const result = await predicate();
+
+    if (result) return result;
+
+    if (Date.now() - startTime > timeoutMs) {
+      throw new Error(
+        `Timeout after ${timeoutMs}ms waiting for ${description}`
+      );
+    }
+
+    await delay(pollIntervalMs);
+  }
+}
+
+export async function sendBatch(count: number, queueName: string) {
+  return await sql`
+    SELECT pgmq.send_batch(
+      ${queueName},
+      ARRAY(
+        SELECT '{}'::jsonb
+        FROM generate_series(1, ${count}::integer)
+      )
+    )`;
+}
+
+export async function seqLastValue(
+  seqName: string = 'test_seq'
+): Promise<number> {
+  // Postgres sequences are initialized with a value of 1,
+  // but incrementing them for the first time does not increment the last_value,
+  // only sets is_called to true
+  const seqResult = await sql`
+    SELECT
+      CASE
+        WHEN is_called THEN last_value::integer
+        ELSE 0
+      END as last_value
+    FROM ${sql(seqName)}`;
+  return seqResult[0].last_value;
+}
+
+interface WaitForSeqValueOptions {
+  pollIntervalMs?: number;
+  seqName?: string;
+  timeoutMs?: number;
+}
+
+export async function waitForSeqToIncrementBy(
+  value: number,
+  options: WaitForSeqValueOptions = {}
+): Promise<number> {
+  const { seqName = 'test_seq' } = options;
+
+  const perSecond = 0;
+
+  const progress = new ProgressBar({
+    title: `${seqName} (${perSecond}/s)`,
+    total: value,
+    width: 20,
+    display: dim(
+      ` -> incrementing "${seqName}": :completed/:total (:eta left) [:bar] :percent`
+    ),
+    prettyTime: true,
+  });
+
+  const startVal = await seqLastValue(seqName);
+  let lastVal = startVal;
+
+  return await waitFor(
+    async () => {
+      lastVal = await seqLastValue(seqName);
+      progress.render(lastVal);
+      const incrementedBy = lastVal - startVal;
+
+      return incrementedBy >= value ? lastVal : false;
+    },
+    {
+      ...options,
+      description: `sequence ${seqName} to reach value ${value}`,
+    }
+  );
+}
+
+export async function waitForActiveWorker() {
+  return await waitFor(
+    async () => {
+      const [{ has_active: hasActiveWorker }] =
+        await sql`SELECT count(*) > 0 AS has_active FROM edge_worker.active_workers`;
+      log('waiting for active worker ', hasActiveWorker);
+      return hasActiveWorker;
+    },
+    {
+      pollIntervalMs: 300,
+      description: 'active worker',
+    }
+  );
+}
+
+export async function fetchWorkers(functionName: string) {
+  return await sql`SELECT * FROM edge_worker.workers WHERE function_name = ${functionName}`;
+}
+
+export async function startWorker(workerName: string) {
+  await sql`SELECT edge_worker.spawn(${workerName}::text)`;
+  await waitForActiveWorker();
+  log('worker spawned!');
+}