@pgflow/edge-worker 0.0.5-prealpha.0

package/src/core/types.ts

@@ -0,0 +1,47 @@
+export type { Json } from '../../../core/src/types.ts';
+
+export interface IPoller<IMessage> {
+  poll(): Promise<IMessage[]>;
+}
+
+export interface IExecutor {
+  get msgId(): number;
+  execute(): Promise<unknown>;
+}
+
+export interface IMessage {
+  msg_id: number;
+}
+
+export interface ILifecycle {
+  acknowledgeStart(workerBootstrap: WorkerBootstrap): Promise<void>;
+  acknowledgeStop(): void;
+  sendHeartbeat(): Promise<void>;
+
+  get edgeFunctionName(): string | undefined;
+  get queueName(): string;
+  get isRunning(): boolean;
+  get isStopping(): boolean;
+  get isStopped(): boolean;
+
+  transitionToStopping(): void;
+}
+
+export interface IBatchProcessor {
+  processBatch(): Promise<void>;
+  awaitCompletion(): Promise<void>;
+}
+
+export type WorkerRow = {
+  last_heartbeat_at: string;
+  queue_name: string;
+  started_at: string;
+  stopped_at: string | null;
+  worker_id: string;
+  function_name: string;
+};
+
+export interface WorkerBootstrap {
+  edgeFunctionName: string;
+  workerId: string;
+}

package/src/flow/FlowWorkerLifecycle.ts

@@ -0,0 +1,81 @@
+import { Heartbeat } from '../core/Heartbeat.ts';
+import { getLogger } from '../core/Logger.ts';
+import type { Queries } from '../core/Queries.ts';
+import type { ILifecycle, WorkerBootstrap, WorkerRow } from '../core/types.ts';
+import { States, WorkerState } from '../core/WorkerState.ts';
+import type { AnyFlow } from '@pgflow/dsl';
+
+/**
+ * A specialized WorkerLifecycle for Flow-based workers that is aware of the Flow's step types
+ */
+export class FlowWorkerLifecycle<TFlow extends AnyFlow> implements ILifecycle {
+  private workerState: WorkerState = new WorkerState();
+  private heartbeat?: Heartbeat;
+  private logger = getLogger('FlowWorkerLifecycle');
+  private queries: Queries;
+  private workerRow?: WorkerRow;
+  private flow: TFlow;
+
+  constructor(queries: Queries, flow: TFlow) {
+    this.queries = queries;
+    this.flow = flow;
+  }
+
+  async acknowledgeStart(workerBootstrap: WorkerBootstrap): Promise<void> {
+    this.workerState.transitionTo(States.Starting);
+
+    this.workerRow = await this.queries.onWorkerStarted({
+      queueName: this.queueName,
+      ...workerBootstrap,
+    });
+
+    this.heartbeat = new Heartbeat(5000, this.queries, this.workerRow);
+
+    this.workerState.transitionTo(States.Running);
+  }
+
+  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);
+      this.logger.debug('Worker stop acknowledged');
+    } catch (error) {
+      this.logger.debug(`Error acknowledging worker stop: ${error}`);
+      throw error;
+    }
+  }
+
+  get edgeFunctionName() {
+    return this.workerRow?.function_name;
+  }
+
+  get queueName() {
+    return this.flow.slug;
+  }
+
+  async sendHeartbeat() {
+    await this.heartbeat?.send();
+  }
+
+  get isRunning() {
+    return this.workerState.isRunning;
+  }
+
+  get isStopping() {
+    return this.workerState.isStopping;
+  }
+
+  get isStopped() {
+    return this.workerState.isStopped;
+  }
+
+  transitionToStopping() {
+    this.workerState.transitionTo(States.Stopping);
+  }
+}

package/src/flow/StepTaskExecutor.ts

@@ -0,0 +1,87 @@
+import type { AnyFlow } from '@pgflow/dsl';
+import type { StepTaskRecord, IPgflowClient } from './types.ts';
+import type { IExecutor } from '../core/types.ts';
+import { getLogger } from '../core/Logger.ts';
+
+class AbortError extends Error {
+  constructor() {
+    super('Operation aborted');
+    this.name = 'AbortError';
+  }
+}
+
+/**
+ * An executor that processes step tasks using an IPgflowClient
+ * with strong typing for the flow's step handlers
+ */
+export class StepTaskExecutor<TFlow extends AnyFlow> implements IExecutor {
+  private logger = getLogger('StepTaskExecutor');
+
+  constructor(
+    private readonly flow: TFlow,
+    private readonly task: StepTaskRecord<TFlow>,
+    private readonly adapter: IPgflowClient<TFlow>,
+    private readonly signal: AbortSignal
+  ) {}
+
+  get msgId() {
+    return this.task.msg_id;
+  }
+
+  async execute(): Promise<void> {
+    try {
+      if (this.signal.aborted) {
+        throw new AbortError();
+      }
+
+      // Check if already aborted before starting
+      this.signal.throwIfAborted();
+
+      const stepSlug = this.task.step_slug;
+      this.logger.debug(
+        `Executing step task ${this.task.msg_id} for step ${stepSlug}`
+      );
+
+      // Get the step handler from the flow with proper typing
+      const stepDef = this.flow.getStepDefinition(stepSlug);
+
+      if (!stepDef) {
+        throw new Error(`No step definition found for slug=${stepSlug}`);
+      }
+
+      // !!! HANDLER EXECUTION !!!
+      const result = await stepDef.handler(this.task.input);
+      // !!! HANDLER EXECUTION !!!
+
+      this.logger.debug(
+        `step task ${this.task.msg_id} completed successfully, marking as complete`
+      );
+      await this.adapter.completeTask(this.task, result);
+
+      this.logger.debug(`step task ${this.task.msg_id} marked as complete`);
+    } catch (error) {
+      await this.handleExecutionError(error);
+    }
+  }
+
+  /**
+   * Handles the error that occurred during execution.
+   *
+   * If the error is an AbortError, it means that the worker was aborted and stopping,
+   * the task will be picked up by another worker later.
+   *
+   * Otherwise, it marks the task as failed.
+   */
+  private async handleExecutionError(error: unknown) {
+    if (error instanceof Error && error.name === 'AbortError') {
+      this.logger.debug(`Aborted execution for step task ${this.task.msg_id}`);
+      // Do not mark as failed - the worker was aborted and stopping,
+      // the task will be picked up by another worker later
+    } else {
+      this.logger.error(
+        `step task ${this.task.msg_id} failed with error: ${error}`
+      );
+      await this.adapter.failTask(this.task, error);
+    }
+  }
+}

package/src/flow/StepTaskPoller.ts

@@ -0,0 +1,51 @@
+import type { StepTaskRecord, IPgflowClient } from './types.ts';
+import type { IPoller } from '../core/types.ts';
+import { getLogger } from '../core/Logger.ts';
+import type { AnyFlow } from '@pgflow/dsl';
+
+export interface StepTaskPollerConfig {
+  batchSize: number;
+  queueName: string;
+}
+
+/**
+ * A poller that retrieves flow tasks using an IPgflowClient
+ */
+export class StepTaskPoller<TFlow extends AnyFlow>
+  implements IPoller<StepTaskRecord<TFlow>>
+{
+  private logger = getLogger('StepTaskPoller');
+
+  constructor(
+    private readonly adapter: IPgflowClient<TFlow>,
+    private readonly signal: AbortSignal,
+    private readonly config: StepTaskPollerConfig
+  ) {}
+
+  async poll(): Promise<StepTaskRecord<TFlow>[]> {
+    if (this.isAborted()) {
+      this.logger.debug('Polling aborted, returning empty array');
+      return [];
+    }
+
+    this.logger.debug(
+      `Polling for flow tasks with batch size ${this.config.batchSize}`
+    );
+
+    try {
+      const tasks = await this.adapter.pollForTasks(
+        this.config.queueName,
+        this.config.batchSize
+      );
+      this.logger.debug(`Retrieved ${tasks.length} flow tasks`);
+      return tasks;
+    } catch (err: unknown) {
+      this.logger.error(`Error polling for flow tasks: ${err}`);
+      return [];
+    }
+  }
+
+  private isAborted(): boolean {
+    return this.signal.aborted;
+  }
+}

package/src/flow/createFlowWorker.ts

@@ -0,0 +1,105 @@
+import type { AnyFlow } from '@pgflow/dsl';
+import type { EdgeWorkerConfig } from '../EdgeWorker.ts';
+import { ExecutionController } from '../core/ExecutionController.ts';
+import { StepTaskPoller, type StepTaskPollerConfig } from './StepTaskPoller.ts';
+import { StepTaskExecutor } from './StepTaskExecutor.ts';
+import { PgflowSqlClient } from '@pgflow/core';
+import { Queries } from '../core/Queries.ts';
+import type { StepTaskRecord } from './types.ts';
+import type { IExecutor } from '../core/types.ts';
+import { Worker } from '../core/Worker.ts';
+import postgres from 'postgres';
+import { FlowWorkerLifecycle } from './FlowWorkerLifecycle.ts';
+import { BatchProcessor } from '../core/BatchProcessor.ts';
+import { getLogger } from '../core/Logger.ts';
+
+/**
+ * Configuration for the flow worker
+ */
+export type FlowWorkerConfig = EdgeWorkerConfig & {
+  maxConcurrent?: number;
+  connectionString?: string;
+  sql?: postgres.Sql;
+  maxPgConnections?: number;
+  batchSize?: number;
+};
+
+/**
+ * Creates a new Worker instance for processing flow tasks.
+ *
+ * @param flow - The Flow DSL definition
+ * @param config - Configuration options for the worker
+ * @returns A configured Worker instance ready to be started
+ */
+export function createFlowWorker<TFlow extends AnyFlow>(
+  flow: TFlow,
+  config: FlowWorkerConfig
+): Worker {
+  const logger = getLogger('createFlowWorker');
+
+  // Create abort controller for graceful shutdown
+  const abortController = new AbortController();
+  const abortSignal = abortController.signal;
+
+  if (!config.sql && !config.connectionString) {
+    throw new Error(
+      "Either 'sql' or 'connectionString' must be provided in FlowWorkerConfig."
+    );
+  }
+
+  const sql =
+    config.sql ||
+    postgres(config.connectionString!, {
+      max: config.maxPgConnections,
+      prepare: false,
+    });
+
+  // Create the pgflow adapter
+  const pgflowAdapter = new PgflowSqlClient<TFlow>(sql);
+
+  // Use flow slug as queue name, or fallback to 'tasks'
+  const queueName = flow.slug || 'tasks';
+  logger.debug(`Using queue name: ${queueName}`);
+
+  // Create specialized FlowWorkerLifecycle with the proxied queue and flow
+  const queries = new Queries(sql);
+  const lifecycle = new FlowWorkerLifecycle<TFlow>(queries, flow);
+
+  // Create StepTaskPoller
+  const pollerConfig: StepTaskPollerConfig = {
+    batchSize: config.batchSize || 10,
+    queueName: flow.slug,
+  };
+  const poller = new StepTaskPoller<TFlow>(
+    pgflowAdapter,
+    abortSignal,
+    pollerConfig
+  );
+
+  // Create executor factory with proper typing
+  const executorFactory = (
+    record: StepTaskRecord<TFlow>,
+    signal: AbortSignal
+  ): IExecutor => {
+    return new StepTaskExecutor<TFlow>(flow, record, pgflowAdapter, signal);
+  };
+
+  // Create ExecutionController
+  const executionController = new ExecutionController<StepTaskRecord<TFlow>>(
+    executorFactory,
+    abortSignal,
+    {
+      maxConcurrent: config.maxConcurrent || 10,
+    }
+  );
+
+  // Create BatchProcessor
+  const batchProcessor = new BatchProcessor<StepTaskRecord<TFlow>>(
+    executionController,
+    poller,
+    abortSignal
+  );
+
+  // Return Worker
+  return new Worker(batchProcessor, lifecycle, sql);
+}

package/src/flow/types.ts

@@ -0,0 +1 @@
+export * from '../../../core/src/types.ts';

package/src/index.ts

@@ -0,0 +1,15 @@
+// Export existing queue-based worker
+export { createQueueWorker } from './queue/createQueueWorker.ts';
+export { EdgeWorker } from './EdgeWorker.ts';
+
+// Export new flow-based worker
+export { createFlowWorker } from './flow/createFlowWorker.ts';
+export { FlowWorkerLifecycle } from './flow/FlowWorkerLifecycle.ts';
+
+// Export types
+export type { StepTaskRecord } from './flow/types.ts';
+export type { FlowWorkerConfig } from './flow/createFlowWorker.ts';
+export type { StepTaskPollerConfig } from './flow/StepTaskPoller.ts';
+
+// Re-export types from the base system
+export type { Json, IExecutor, IPoller, IMessage, ILifecycle, IBatchProcessor } from './core/types.ts';

package/src/queue/MessageExecutor.ts

@@ -0,0 +1,105 @@
+import type { Json } from '../core/types.ts';
+import type { PgmqMessageRecord } from './types.ts';
+import type { Queue } from './Queue.ts';
+import { getLogger } from '../core/Logger.ts';
+
+class AbortError extends Error {
+  constructor() {
+    super('Operation aborted');
+    this.name = 'AbortError';
+  }
+}
+
+/**
+ * A class that executes a message handler.
+ *
+ * It handles the execution of the message handler and retries or archives the message
+ * based on the retry limit and delay.
+ *
+ * It also handles the abort signal and logs the error.
+ */
+export class MessageExecutor<TPayload extends Json> {
+  private logger = getLogger('MessageExecutor');
+
+  constructor(
+    private readonly queue: Queue<TPayload>,
+    private readonly record: PgmqMessageRecord<TPayload>,
+    private readonly messageHandler: (
+      message: TPayload
+    ) => Promise<void> | void,
+    private readonly signal: AbortSignal,
+    private readonly retryLimit: number,
+    private readonly retryDelay: number
+  ) {}
+
+  get msgId() {
+    return this.record.msg_id;
+  }
+
+  async execute(): Promise<void> {
+    try {
+      if (this.signal.aborted) {
+        throw new AbortError();
+      }
+
+      // Check if already aborted before starting
+      this.signal.throwIfAborted();
+
+      this.logger.debug(`Executing task ${this.msgId}...`);
+      await this.messageHandler(this.record.message!);
+
+      this.logger.debug(
+        `Task ${this.msgId} completed successfully, archiving...`
+      );
+      await this.queue.archive(this.msgId);
+      this.logger.debug(`Archived task ${this.msgId} successfully`);
+    } catch (error) {
+      await this.handleExecutionError(error);
+    }
+  }
+
+  /**
+   * Handles the error that occurred during execution.
+   *
+   * If the error is an AbortError, it means that the worker was aborted and stopping,
+   * the message will reappear after the visibility timeout and be picked up by another worker.
+   *
+   * Otherwise, it proceeds with retry or archiving forever.
+   */
+  private async handleExecutionError(error: unknown) {
+    if (error instanceof Error && error.name === 'AbortError') {
+      this.logger.debug(`Aborted execution for ${this.msgId}`);
+      // Do not throw - the worker was aborted and stopping,
+      // the message will reappear after the visibility timeout
+      // and be picked up by another worker
+    } else {
+      this.logger.debug(`Task ${this.msgId} failed with error: ${error}`);
+      await this.retryOrArchive();
+    }
+  }
+
+  /**
+   * Retries the message if it is available.
+   * Otherwise, archives the message forever and stops processing it.
+   */
+  private async retryOrArchive() {
+    if (this.retryAvailable) {
+      // adjust visibility timeout for message to appear after retryDelay
+      this.logger.debug(`Retrying ${this.msgId} in ${this.retryDelay} seconds`);
+      await this.queue.setVt(this.msgId, this.retryDelay);
+    } else {
+      // archive message forever and stop processing it
+      this.logger.debug(`Archiving ${this.msgId} forever`);
+      await this.queue.archive(this.msgId);
+    }
+  }
+
+  /**
+   * Returns true if the message can be retried.
+   */
+  private get retryAvailable() {
+    const readCountLimit = this.retryLimit + 1; // initial read also counts
+
+    return this.record.read_ct < readCountLimit;
+  }
+}

package/src/queue/Queue.ts

@@ -0,0 +1,92 @@
+import type postgres from 'postgres';
+import type { PgmqMessageRecord } from './types.ts';
+import type { Json } from '../core/types.ts';
+
+export class Queue<TPayload extends Json> {
+  constructor(private readonly sql: postgres.Sql, readonly queueName: string) {}
+
+  /**
+   * Creates a queue if it doesn't exist.
+   * If the queue already exists, this method does nothing.
+   */
+  async safeCreate() {
+    return await this.sql`
+        select * from pgmq.create(${this.queueName})
+        where not exists (
+          select 1 from pgmq.list_queues() where queue_name = ${this.queueName}
+        );
+    `;
+  }
+
+  /**
+   * Drops a queue if it exists.
+   * If the queue doesn't exist, this method does nothing.
+   */
+  async safeDrop() {
+    return await this.sql`
+        select * from pgmq.drop_queue(${this.queueName})
+        where exists (
+          select 1 from pgmq.list_queues() where queue_name = ${this.queueName}
+        );
+    `;
+  }
+
+  async archive(msgId: number): Promise<void> {
+    await this.sql`
+      SELECT pgmq.archive(queue_name => ${this.queueName}, msg_id => ${msgId}::bigint);
+    `;
+  }
+
+  async archiveBatch(msgIds: number[]): Promise<void> {
+    await this.sql`
+      SELECT pgmq.archive(queue_name => ${this.queueName}, msg_ids => ${msgIds}::bigint[]);
+    `;
+  }
+
+  async send(message: TPayload): Promise<void> {
+    const msgJson = JSON.stringify(message);
+    await this.sql`
+      SELECT pgmq.send(queue_name => ${this.queueName}, msg => ${msgJson}::jsonb)
+    `;
+  }
+
+  async readWithPoll(
+    batchSize = 20,
+    visibilityTimeout = 2,
+    maxPollSeconds = 5,
+    pollIntervalMs = 200
+  ) {
+    return await this.sql<PgmqMessageRecord<TPayload>[]>`
+      SELECT *
+      FROM edge_worker.read_with_poll(
+        queue_name => ${this.queueName},
+        vt => ${visibilityTimeout},
+        qty => ${batchSize},
+        max_poll_seconds => ${maxPollSeconds},
+        poll_interval_ms => ${pollIntervalMs}
+      );
+    `;
+  }
+
+  /**
+   * Sets the visibility timeout of a message to the current time plus the given offset.
+   *
+   * This is an inlined version of the pgmq.set_vt in order to fix the bug.
+   * The original uses now() instead of clock_timestamp() which is problematic in transactions.
+   * See more details here: https://github.com/tembo-io/pgmq/issues/367
+   *
+   * The only change made is now() replaced with clock_timestamp().
+   */
+  async setVt(
+    msgId: number,
+    vtOffsetSeconds: number
+  ): Promise<PgmqMessageRecord<TPayload>> {
+    const records = await this.sql<PgmqMessageRecord<TPayload>[]>`
+      UPDATE ${this.sql('pgmq.q_' + this.queueName)}
+      SET vt = (clock_timestamp() + make_interval(secs => ${vtOffsetSeconds}))
+      WHERE msg_id = ${msgId}::bigint
+      RETURNING *;
+    `;
+    return records[0];
+  }
+}

package/src/queue/ReadWithPollPoller.ts

@@ -0,0 +1,35 @@
+import type { Queue } from './Queue.ts';
+import type { PgmqMessageRecord } from './types.ts';
+import type { Json } from '../core/types.ts';
+
+export interface PollerConfig {
+  batchSize: number;
+  maxPollSeconds: number;
+  pollIntervalMs: number;
+  visibilityTimeout: number;
+}
+
+export class ReadWithPollPoller<TPayload extends Json> {
+  constructor(
+    protected readonly queue: Queue<TPayload>,
+    protected readonly signal: AbortSignal,
+    protected readonly config: PollerConfig
+  ) {}
+
+  async poll(): Promise<PgmqMessageRecord<TPayload>[]> {
+    if (this.isAborted()) {
+      return [];
+    }
+
+    return await this.queue.readWithPoll(
+      this.config.batchSize,
+      this.config.visibilityTimeout,
+      this.config.maxPollSeconds,
+      this.config.pollIntervalMs
+    );
+  }
+
+  private isAborted(): boolean {
+    return this.signal.aborted;
+  }
+}