Queue Drivers + Direct wf/steps proxy

Signed-off-by: paulhenri-l <25308170+paulhenri-l@users.noreply.github.com>

Author: paulhenri-l

packages/world-postgres/src/config.ts

@@ -1,5 +1,11 @@
-export interface PostgresWorldConfig {
+import type { QueueDriver } from './queue-drivers/types.js';
+
+export interface QueueConfig {
   connectionString: string;
   jobPrefix?: string;
   queueConcurrency?: number;
 }
+
+export type PostgresWorldConfig = QueueConfig & {
+  queueFactory: (config: QueueConfig) => QueueDriver;
+};

packages/world-postgres/src/index.ts

@@ -1,9 +1,9 @@
 import type { Storage, World } from '@workflow/world';
-import PgBoss from 'pg-boss';
 import createPostgres from 'postgres';
 import type { PostgresWorldConfig } from './config.js';
 import { createClient, type Drizzle } from './drizzle/index.js';
 import { createQueue } from './queue.js';
+import { createPgBossQueue } from './queue-drivers/pg-boss.js';
 import {
   createEventsStorage,
   createHooksStorage,
@@ -30,14 +30,14 @@ export function createWorld(
     queueConcurrency:
       parseInt(process.env.WORKFLOW_POSTGRES_WORKER_CONCURRENCY || '10', 10) ||
       10,
+    queueFactory: (config) => createPgBossQueue(config),
   }
 ): World & { start(): Promise<void> } {
-  const boss = new PgBoss({
-    connectionString: config.connectionString,
-  });
+  const queueDriver = config.queueFactory(config);
   const postgres = createPostgres(config.connectionString);
   const drizzle = createClient(postgres);
-  const queue = createQueue(boss, config);
+
+  const queue = createQueue(queueDriver);
   const storage = createStorage(drizzle);
   const streamer = createStreamer(postgres, drizzle);
 
@@ -46,7 +46,7 @@ export function createWorld(
     ...streamer,
     ...queue,
     async start() {
-      await queue.start();
+      await queueDriver.start();
     },
   };
 }

packages/world-postgres/src/queue-drivers/pg-boss.ts

@@ -0,0 +1,86 @@
+import PgBoss from 'pg-boss';
+import type { QueueConfig } from '../config.js';
+import { MessageData } from './types.js';
+import { proxyStep, proxyWorkflow } from './wkf-proxy.js';
+
+export function createPgBossQueue(config: QueueConfig) {
+  let startPromise: Promise<unknown> | null = null;
+  const boss = new PgBoss(config.connectionString);
+
+  const stepQueueName = 'workflow_steps';
+  const workflowQueueName = 'workflow_flows';
+
+  const ensureStarted = async () => {
+    if (!startPromise) {
+      startPromise = boss.start().then(() => {
+        return Promise.all([
+          boss.createQueue(workflowQueueName),
+          boss.createQueue(stepQueueName),
+        ]);
+      });
+    }
+
+    await startPromise;
+  };
+
+  return {
+    pushStep: async (message: MessageData) => {
+      await ensureStarted();
+
+      await boss.send(stepQueueName, MessageData.encode(message), {
+        singletonKey: message?.idempotencyKey ?? message.messageId,
+        retryLimit: 3,
+      });
+    },
+
+    pushFlow: async (message: MessageData) => {
+      await ensureStarted();
+
+      await boss.send(workflowQueueName, MessageData.encode(message), {
+        singletonKey: message?.idempotencyKey ?? message.messageId,
+        retryLimit: 3,
+      });
+    },
+
+    start: async () => {
+      await ensureStarted();
+
+      const workflowHandler = async ([job]: PgBoss.Job[]) => {
+        const message = MessageData.parse(job.data);
+
+        console.log(`[${job.id}] running: ${message.queueName}`);
+
+        try {
+          await proxyWorkflow(message);
+        } catch (error) {
+          console.error(
+            `[${job.id}] Error handling workflow: ${message.queueName}`,
+            error
+          );
+          throw error;
+        }
+      };
+
+      const stepHandler = async ([job]: PgBoss.Job[]) => {
+        const message = MessageData.parse(job.data);
+
+        console.log(`[${job.id}] running: ${message.queueName}`);
+
+        try {
+          await proxyStep(message);
+        } catch (error) {
+          console.error(
+            `[${job.id}] Error handling step: ${message.queueName}`,
+            error
+          );
+          throw error;
+        }
+      };
+
+      for (let i = 0; i < (config.queueConcurrency || 10); i++) {
+        await boss.work(workflowQueueName, workflowHandler);
+        await boss.work(stepQueueName, stepHandler);
+      }
+    },
+  };
+}

packages/world-postgres/src/boss.ts renamed to packages/world-postgres/src/queue-drivers/types.ts

@@ -1,21 +1,30 @@
 import { MessageId } from '@workflow/world';
 import * as z from 'zod';
-import { Base64Buffer } from './zod.js';
+import { Base64Buffer } from '../zod.js';
 
 /**
-/* pgboss is using JSON under the hood, so we need to base64 encode
-/* the body to ensure binary safety
-/* maybe later we can have a `blobs` table for larger payloads
- **/
+/* Most queues are using JSON under the hood, so we need to base64
+/* encode the body to ensure binary safety maybe later we can
+/* have a `blobs` table for larger payloads
+**/
 export const MessageData = z.object({
-  attempt: z.number().describe('The attempt number of the message'),
-  messageId: MessageId.describe('The unique ID of the message'),
-  idempotencyKey: z.string().optional(),
   id: z
     .string()
     .describe(
       "The ID of the sub-queue. For workflows, it's the workflow name. For steps, it's the step name."
     ),
+
+  idempotencyKey: z.string().optional(),
+  queueName: z.string().describe('The name of the queue'),
   data: Base64Buffer.describe('The message that was sent'),
+  messageId: MessageId.describe('The unique ID of the message'),
+  attempt: z.number().describe('The attempt number of the message'),
 });
+
 export type MessageData = z.infer<typeof MessageData>;
+
+export interface QueueDriver {
+  pushStep: (message: MessageData) => Promise<void>;
+  pushFlow: (message: MessageData) => Promise<void>;
+  start: () => Promise<void>;
+}

packages/world-postgres/src/queue-drivers/wkf-proxy.ts

@@ -0,0 +1,29 @@
+import { createRequire } from 'node:module';
+import Path from 'node:path';
+import { MessageData } from './types.js';
+
+const require = createRequire(Path.join(process.cwd(), 'index.js'));
+
+export async function proxyWorkflow(message: MessageData) {
+  const workflows = require(process.env.WORKFLOW_POSTGRES_FLOWS!);
+
+  const request = new Request('https://world-postgres.local/wkf-direct-call', {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify(MessageData.encode(message)),
+  });
+
+  return (await workflows.__wkf_entrypoint(request)) as Promise<Response>;
+}
+
+export async function proxyStep(message: MessageData) {
+  const steps = require(process.env.WORKFLOW_POSTGRES_STEPS!);
+
+  const request = new Request('https://world-postgres.local/wkf-direct-call', {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify(MessageData.encode(message)),
+  });
+
+  return (await steps.__wkf_entrypoint(request)) as Promise<Response>;
+}