@pgflow/edge-worker 0.0.5-prealpha.0

package/tests/e2e/performance.test.ts

@@ -0,0 +1,60 @@
+import { sql } from '../sql.ts';
+import {
+  waitFor,
+  sendBatch,
+  waitForSeqToIncrementBy,
+  startWorker,
+  log,
+} from './_helpers.ts';
+
+const MESSAGES_TO_SEND = 20000;
+const WORKER_NAME = 'max_concurrency';
+
+Deno.test(
+  'worker can handle tens of thousands of jobs queued at once',
+  async () => {
+    await sql`CREATE SEQUENCE IF NOT EXISTS test_seq`;
+    await sql`ALTER SEQUENCE test_seq RESTART WITH 1`;
+    await sql`SELECT pgmq.create(${WORKER_NAME})`;
+    await sql`SELECT pgmq.drop_queue(${WORKER_NAME})`;
+    await sql`SELECT pgmq.create(${WORKER_NAME})`;
+    await startWorker(WORKER_NAME);
+    await waitFor(
+      async () => {
+        const [{ worker_count }] = await sql`
+        SELECT COUNT(*)::integer AS worker_count
+        FROM edge_worker.active_workers
+        WHERE function_name = ${WORKER_NAME}
+      `;
+
+        log('worker_count', worker_count);
+        return worker_count === 1;
+      },
+      { description: 'Waiting for exacly one worker' }
+    );
+
+    try {
+      // worker sleeps for 1s for each message
+      // se we will expect roughly 1 message per second
+      const startTime = Date.now();
+
+      await sendBatch(MESSAGES_TO_SEND, WORKER_NAME);
+      await waitForSeqToIncrementBy(MESSAGES_TO_SEND, {
+        timeoutMs: MESSAGES_TO_SEND * 1000 + 1000,
+        pollIntervalMs: 1000,
+      });
+
+      const endTime = Date.now();
+      const totalMs = Math.round(endTime - startTime);
+      const totalS = totalMs / 1000;
+      const msgsPerSecond = MESSAGES_TO_SEND / totalS;
+
+      log('');
+      log('');
+      log(`Total time:`, totalMs);
+      log(`msgs/second:`, msgsPerSecond);
+    } finally {
+      await sql.end();
+    }
+  }
+);

package/tests/e2e/restarts.test.ts

@@ -0,0 +1,56 @@
+import { sql } from '../sql.ts';
+import { assertGreater, assertGreaterOrEqual } from 'jsr:@std/assert';
+import {
+  fetchWorkers,
+  sendBatch,
+  seqLastValue,
+  startWorker,
+  waitForSeqToIncrementBy,
+} from './_helpers.ts';
+
+const WORKER_NAME = 'cpu_intensive';
+
+// TODO: document relation between CPU clock limit, amount of time to process
+//       single message and amount of messages to send
+const MESSAGES_TO_SEND = 30;
+
+Deno.test('should spawn next worker when CPU clock limit hits', async () => {
+  await sql`CREATE SEQUENCE IF NOT EXISTS test_seq`;
+  await sql`ALTER SEQUENCE test_seq RESTART WITH 1`;
+  try {
+    await sql`SELECT pgmq.drop_queue(${WORKER_NAME})`;
+  } catch {
+    // ignore
+  }
+  await sql`SELECT pgmq.create(${WORKER_NAME})`;
+  await sql`
+    DELETE FROM edge_worker.workers
+    WHERE worker_id IN (
+      SELECT worker_id
+      FROM edge_worker.inactive_workers
+    )`;
+  await startWorker(WORKER_NAME);
+
+  try {
+    await sendBatch(MESSAGES_TO_SEND, WORKER_NAME);
+    await waitForSeqToIncrementBy(MESSAGES_TO_SEND, {
+      timeoutMs: 35000,
+      pollIntervalMs: 300,
+    });
+
+    assertGreaterOrEqual(
+      await seqLastValue(),
+      MESSAGES_TO_SEND,
+      'Sequence value should be greater than or equal to the number of messages sent'
+    );
+
+    const workers = await fetchWorkers(WORKER_NAME);
+    assertGreater(
+      workers.length,
+      1,
+      'expected worker to spawn another but there is only 1 worker'
+    );
+  } finally {
+    await sql.end();
+  }
+});

package/tests/helpers.ts

@@ -0,0 +1,22 @@
+import type { PgmqMessageRecord } from "../src/queue/types.ts";
+import type { postgres } from "./sql.ts";
+
+export async function sendBatch(count: number, queueName: string, sql: postgres.Sql) {
+  return await sql`
+    SELECT pgmq.send_batch(
+      ${queueName},
+      ARRAY(
+        SELECT '{}'::jsonb
+        FROM generate_series(1, ${count}::integer)
+      )
+    )`;
+}
+
+/**
+ * Fetches archived messages from the queue
+ */
+export async function getArchivedMessages(sql: postgres.Sql, queueName: string) {
+  return await sql<PgmqMessageRecord[]>
+    `SELECT * FROM ${sql('pgmq.a_' + queueName)}`;
+}
+

package/tests/integration/_helpers.ts

@@ -0,0 +1,43 @@
+import type { AnyFlow, ExtractFlowInput, Flow, Json } from '@pgflow/dsl';
+import {
+  createFlowWorker,
+  type FlowWorkerConfig,
+} from '../../src/flow/createFlowWorker.ts';
+import type { postgres } from '../sql.ts';
+import { PgflowSqlClient } from '../../../core/src/PgflowSqlClient.ts';
+
+export async function startFlow<TFlow extends AnyFlow>(
+  sql: postgres.Sql,
+  flow: TFlow,
+  input: ExtractFlowInput<TFlow>
+) {
+  const pgflow = new PgflowSqlClient<TFlow>(sql);
+
+  return await pgflow.startFlow(flow, input);
+}
+
+export function startWorker<
+  T extends Json,
+  S extends Record<string, Json> = Record<never, never>,
+  D extends Record<string, string[]> = Record<string, string[]>
+>(sql: postgres.Sql, flow: Flow<T, S, D>, options: FlowWorkerConfig) {
+  const defaultOptions = {
+    sql,
+    maxConcurrent: 1,
+    batchSize: 10,
+  };
+
+  const mergedOptions = {
+    ...defaultOptions,
+    ...options,
+  };
+
+  const worker = createFlowWorker(flow, mergedOptions);
+
+  worker.startOnlyOnce({
+    edgeFunctionName: 'test_flow',
+    workerId: crypto.randomUUID(),
+  });
+
+  return worker;
+}

package/tests/integration/creating_queue.test.ts

@@ -0,0 +1,32 @@
+import { assertEquals } from "@std/assert";
+import { createQueueWorker } from '../../src/queue/createQueueWorker.ts';
+import { withTransaction } from "../db.ts";
+import { delay } from "@std/async";
+
+Deno.test('creates queue when starting worker', withTransaction(async (sql) => {
+  const worker = createQueueWorker(console.log, {
+    sql,
+    maxPollSeconds: 1,
+    queueName: 'custom_queue'
+  });
+
+  worker.startOnlyOnce({
+    edgeFunctionName: 'test',
+    // random uuid
+    workerId: crypto.randomUUID(),
+  });
+
+  await delay(100);
+
+  try {
+    const result: {queue_name: string}[] = await sql`select queue_name from pgmq.list_queues();`;
+
+    assertEquals(
+      [...result],
+      [{ queue_name: 'custom_queue' }],
+      'queue "custom_queue" was created'
+    );
+  } finally {
+    await worker.stop();
+  }
+}));

package/tests/integration/flow/minimalFlow.test.ts

@@ -0,0 +1,121 @@
+import { assert, assertEquals } from '@std/assert';
+import { withPgNoTransaction } from '../../db.ts';
+import { Flow } from '../../../../dsl/src/dsl.ts';
+import { waitFor } from '../../e2e/_helpers.ts';
+import { delay } from '@std/async';
+import type { Json } from '@pgflow/core';
+import { startFlow, startWorker } from '../_helpers.ts';
+
+// Define a minimal flow with two steps:
+// 1. Convert a number to a string
+// 2. Wrap the string in an array
+const MinimalFlow = new Flow<number>({ slug: 'test_minimal_flow' })
+  .step({ slug: 'toStringStep' }, async (input) => {
+    await delay(1);
+    return input.run.toString();
+  })
+  .step(
+    { slug: 'wrapInArrayStep', dependsOn: ['toStringStep'] },
+    async (input) => {
+      await delay(1);
+      return [input.toStringStep];
+    }
+  );
+
+Deno.test(
+  'minimal flow executes successfully',
+  withPgNoTransaction(async (sql) => {
+    await sql`select pgflow_tests.reset_db();`;
+
+    const worker = startWorker(sql, MinimalFlow, {
+      maxConcurrent: 1,
+      batchSize: 10,
+    });
+
+    try {
+      await sql`select pgflow.create_flow('test_minimal_flow');`;
+      await sql`select pgflow.add_step('test_minimal_flow', 'toStringStep');`;
+      await sql`select pgflow.add_step('test_minimal_flow', 'wrapInArrayStep', deps_slugs => ARRAY['toStringStep']::text[]);`;
+
+      // Start a flow run with input value 42
+      const flowRun = await startFlow(sql, MinimalFlow, 42);
+
+      let i = 0;
+      // Wait for the run to complete with a timeout
+      const polledRun = await waitFor(
+        async () => {
+          // Check run status
+          const [run] = await sql`
+            SELECT * FROM pgflow.runs WHERE run_id = ${flowRun.run_id};
+          `;
+
+          i += 1;
+          console.log(`Run ${i}`, run);
+
+          if (run.status != 'completed' && run.status != 'failed') {
+            return false;
+          }
+
+          return run;
+        },
+        {
+          pollIntervalMs: 500,
+          timeoutMs: 5000,
+          description: `flow run ${flowRun.run_id} to be 'completed'`,
+        }
+      );
+
+      console.log('Polled run', polledRun);
+
+      assert(polledRun.status === 'completed', 'Run should be completed');
+
+      // Verify step_states are all completed
+      const stepStates = await sql<{ step_slug: string; status: string }[]>`
+        SELECT step_slug, status FROM pgflow.step_states
+        WHERE run_id = ${flowRun.run_id}
+        ORDER BY step_slug;
+      `;
+
+      console.log('Step states:', stepStates);
+      assertEquals(
+        stepStates.map((s) => s.status),
+        ['completed', 'completed'],
+        'All step states should be completed'
+      );
+
+      // Verify step_tasks are all succeeded
+      const stepTasks = await sql<
+        { step_slug: string; status: string; output: Json }[]
+      >`
+        SELECT step_slug, status, output FROM pgflow.step_tasks
+        WHERE run_id = ${flowRun.run_id}
+        ORDER BY step_slug;
+      `;
+
+      console.log('Step tasks:', stepTasks);
+      assertEquals(
+        stepTasks.map((s) => s.status),
+        ['completed', 'completed'],
+        'All step tasks should be succeeded'
+      );
+
+      // Verify run is succeeded
+      const [finalRun] = await sql<{ status: string; output: unknown }[]>`
+        SELECT status, output FROM pgflow.runs WHERE run_id = ${flowRun.run_id};
+      `;
+
+      console.log('Final run:', finalRun);
+      assertEquals(finalRun.status, 'completed', 'Run should be succeeded');
+
+      // Verify run output matches expected ["42"]
+      assertEquals(
+        finalRun.output,
+        { wrapInArrayStep: ['42'] },
+        'Run output should match expected value'
+      );
+    } finally {
+      // Stop the worker
+      await worker.stop();
+    }
+  })
+);

package/tests/integration/maxConcurrent.test.ts

@@ -0,0 +1,76 @@
+import { assertEquals, assertGreaterOrEqual } from '@std/assert';
+import { createQueueWorker } from '../../src/queue/createQueueWorker.ts';
+import { withTransaction } from '../db.ts';
+import { waitFor } from '../e2e/_helpers.ts';
+import type { PgmqMessageRecord } from '../../src/queue/types.ts';
+import { delay } from '@std/async';
+import { sendBatch } from '../helpers.ts';
+
+const QUEUE_NAME = 'max_concurrent';
+const MESSAGES_TO_SEND = 3;
+
+async function sleepFor1s() {
+  await delay(1000);
+}
+
+Deno.test(
+  'maxConcurrent option is respected',
+  withTransaction(async (sql) => {
+    const worker = createQueueWorker(sleepFor1s, {
+      sql,
+      maxConcurrent: 1,
+      maxPollSeconds: 1,
+      visibilityTimeout: 5,
+      queueName: QUEUE_NAME,
+    });
+
+    try {
+      worker.startOnlyOnce({
+        edgeFunctionName: 'test',
+        // random uuid
+        workerId: crypto.randomUUID(),
+      });
+
+      await sendBatch(MESSAGES_TO_SEND, QUEUE_NAME, sql);
+
+      // worker sleeps for 1s for each message
+      // se we will expect roughly 1 message per second
+      const startTime = Date.now();
+
+      const messages = await waitFor(
+        async () => {
+          const archivedMessages = await sql<
+            PgmqMessageRecord[]
+          >`SELECT * FROM ${sql('pgmq.a_' + QUEUE_NAME)}`;
+
+          return (
+            archivedMessages.length >= MESSAGES_TO_SEND && archivedMessages
+          );
+        },
+        {
+          timeoutMs: 5000,
+        }
+      );
+
+      assertEquals(messages.length, 3, 'there should be 3 archived messages');
+      assertEquals(
+        messages.map((m) => m.read_ct),
+        [1, 1, 1],
+        'each message should be read exacly once'
+      );
+
+      const endTime = Date.now();
+      const totalMs = Math.round(endTime - startTime);
+
+      assertGreaterOrEqual(
+        totalMs,
+        MESSAGES_TO_SEND * 1000, // 3 messages, each takes 1s
+        `Should take at least ${MESSAGES_TO_SEND}s to process all messages, took ${totalMs}ms instead`
+      );
+    } catch (error) {
+      throw error;
+    } finally {
+      await worker.stop();
+    }
+  })
+);

package/tests/integration/retries.test.ts

@@ -0,0 +1,78 @@
+import { assertEquals, assertGreaterOrEqual } from '@std/assert';
+import { createQueueWorker } from '../../src/queue/createQueueWorker.ts';
+import { withTransaction } from "../db.ts";
+import { log, waitFor } from "../e2e/_helpers.ts";
+import { getArchivedMessages, sendBatch } from "../helpers.ts";
+
+const workerConfig = {
+  maxPollSeconds: 1,
+  retryDelay: 2,      // seconds between retries
+  retryLimit: 2,      // number of retries
+  queueName: 'failing_always',
+} as const;
+
+/**
+ * Creates a handler that always fails and logs elapsed time
+ * @param startTime - reference time for elapsed calculations
+ */
+function createFailingHandler(startTime: number) {
+  return function failingHandler() {
+    const elapsedSec = ((Date.now() - startTime) / 1000).toFixed(2);
+    log(`[elapsed: ${elapsedSec}s] Failed as expected (╯°□°)╯︵ ┻━┻`);
+    throw new Error('Intentional failure');
+  };
+}
+
+/**
+ * Test verifies that:
+ * 1. Message processing takes at least RETRY_LIMIT * RETRY_DELAY seconds
+ * 2. Message is read exactly RETRY_LIMIT + 1 times (initial + retries)
+ */
+Deno.test('message retry mechanism works correctly', withTransaction(async (sql) => {
+  const startTime = Date.now();
+  const worker = createQueueWorker(createFailingHandler(startTime), {
+    sql,
+    ...workerConfig
+  });
+
+  try {
+    // Start worker and send test message
+    worker.startOnlyOnce({
+      edgeFunctionName: 'test',
+      workerId: crypto.randomUUID(),
+    });
+    await sendBatch(1, workerConfig.queueName, sql);
+
+    // Calculate expected processing time
+    const expectedMinimumMs = workerConfig.retryLimit * workerConfig.retryDelay * 1000;
+
+    // Wait for message to be archived
+    const [message] = await waitFor(
+      async () => {
+        const messages = await getArchivedMessages(sql, workerConfig.queueName);
+        return messages.length >= 1 && messages;
+      },
+      {
+        timeoutMs: expectedMinimumMs + 500,
+      }
+    );
+
+    // Verify timing
+    const totalMs = Date.now() - startTime;
+    assertGreaterOrEqual(
+      totalMs,
+      expectedMinimumMs,
+      `Processing time ${totalMs}ms was shorter than minimum ${expectedMinimumMs}ms`
+    );
+
+    // Verify retry count
+    const expectedReads = workerConfig.retryLimit + 1;
+    assertEquals(
+      message.read_ct,
+      expectedReads,
+      `Message should be read ${expectedReads} times (1 initial + ${workerConfig.retryLimit} retries)`
+    );
+  } finally {
+    await worker.stop();
+  }
+}));

package/tests/integration/starting_worker.test.ts

@@ -0,0 +1,35 @@
+import { createQueueWorker } from '../../src/queue/createQueueWorker.ts';
+import { withTransaction } from "../db.ts";
+import { delay } from "@std/async";
+
+Deno.test('Starting worker', withTransaction(async (sql) => {
+  const worker = createQueueWorker(console.log, {
+    sql,
+    maxPollSeconds: 1
+  });
+
+  worker.startOnlyOnce({
+    edgeFunctionName: 'test',
+    // random uuid
+    workerId: crypto.randomUUID(),
+  });
+
+  await delay(100);
+
+  try {
+    const workers = await sql`select * from edge_worker.workers`;
+
+    console.log(workers);
+  } finally {
+    await worker.stop();
+  }
+}));
+
+Deno.test('check pgmq version', withTransaction(async (sql) => {
+  const result = await sql`
+    SELECT extversion
+    FROM pg_extension
+    WHERE extname = 'pgmq'
+  `;
+  console.log('pgmq version:', result);
+}));

package/tests/sql.ts

@@ -0,0 +1,46 @@
+import postgres from 'postgres';
+
+const DB_URL = 'postgresql://postgres:postgres@127.0.0.1:50322/postgres';
+
+export function createSql() {
+  return postgres(DB_URL, {
+    prepare: false,
+    onnotice(_: unknown) {
+      // no-op to silence notices
+    },
+  });
+}
+
+export async function withRollback<T>(
+  callback: (sql: postgres.Sql) => Promise<T>
+): Promise<T> {
+  const sql = createSql();
+  try {
+    const result = (await sql.begin(
+      'read write',
+      async (sqlTx: postgres.Sql) => {
+        const callbackResult = await callback(sqlTx);
+        await sqlTx`ROLLBACK`;
+        return callbackResult;
+      }
+    )) as T;
+    return result;
+  } finally {
+    await sql.end();
+  }
+}
+
+export async function withSql<T>(
+  callback: (sql: postgres.Sql) => Promise<T>
+): Promise<T> {
+  const sql = createSql();
+  try {
+    return await callback(sql);
+  } finally {
+    await sql.end();
+  }
+}
+
+const sql = createSql();
+
+export { type postgres, sql };

package/tests/unit/WorkerState.test.ts

@@ -0,0 +1,74 @@
+import { assertEquals, assertThrows } from '@std/assert';
+import { WorkerState, States, TransitionError } from '../../src/core/WorkerState.ts';
+
+Deno.test('WorkerState - initial state should be Created', () => {
+  const state = new WorkerState();
+  assertEquals(state.current, States.Created);
+  assertEquals(state.isCreated, true);
+});
+
+Deno.test('WorkerState - valid state transitions', () => {
+  const state = new WorkerState();
+
+  // Created -> Starting
+  state.transitionTo(States.Starting);
+  assertEquals(state.current, States.Starting);
+  assertEquals(state.isStarting, true);
+
+  // Starting -> Running
+  state.transitionTo(States.Running);
+  assertEquals(state.current, States.Running);
+  assertEquals(state.isRunning, true);
+
+  // Running -> Stopping
+  state.transitionTo(States.Stopping);
+  assertEquals(state.current, States.Stopping);
+  assertEquals(state.isStopping, true);
+
+  // Stopping -> Stopped
+  state.transitionTo(States.Stopped);
+  assertEquals(state.current, States.Stopped);
+});
+
+Deno.test('WorkerState - invalid state transitions should throw', () => {
+  const state = new WorkerState();
+
+  // Cannot transition from Created to Running
+  assertThrows(
+    () => {
+      state.transitionTo(States.Running);
+    },
+    TransitionError,
+    'Cannot transition from created to running'
+  );
+
+  // Cannot transition from Created to Stopped
+  assertThrows(
+    () => {
+      state.transitionTo(States.Stopped);
+    },
+    TransitionError,
+    'Cannot transition from created to stopped'
+  );
+});
+
+Deno.test('WorkerState - transitioning to same state should be no-op', () => {
+  const state = new WorkerState();
+
+  // Transition to Starting first
+  state.transitionTo(States.Starting);
+  assertEquals(state.current, States.Starting);
+
+  // Transition to Starting again
+  state.transitionTo(States.Starting);
+  assertEquals(state.current, States.Starting);
+});
+
+Deno.test('WorkerState - state getters', () => {
+  const state = new WorkerState();
+
+  assertEquals(state.isCreated, true);
+  assertEquals(state.isStarting, false);
+  assertEquals(state.isRunning, false);
+  assertEquals(state.isStopping, false);
+});

package/tsconfig.lib.json

@@ -0,0 +1,23 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "outDir": "dist",
+    "tsBuildInfoFile": "dist/tsconfig.lib.tsbuildinfo",
+    "emitDeclarationOnly": false,
+    "types": ["node"],
+    "target": "es2020",
+    "module": "commonjs",
+    "moduleResolution": "node"
+  },
+  "include": ["src/**/*.ts"],
+  "exclude": [
+    "src/**/*.test.ts",
+    "src/**/*.spec.ts",
+    "src/**/*.test.tsx",
+    "src/**/*.spec.tsx",
+    "src/**/*.test.js",
+    "src/**/*.spec.js",
+    "src/**/*.test.jsx",
+    "src/**/*.spec.jsx"
+  ]
+}