@pgflow/edge-worker 0.0.5-prealpha.0

package/src/EdgeWorker.ts

@@ -0,0 +1,172 @@
+import type { Worker } from './core/Worker.ts';
+import spawnNewEdgeFunction from './spawnNewEdgeFunction.ts';
+import type { Json } from './core/types.ts';
+import { getLogger, setupLogger } from './core/Logger.ts';
+import {
+  createQueueWorker,
+  type QueueWorkerConfig,
+} from './queue/createQueueWorker.ts';
+
+/**
+ * Configuration options for the EdgeWorker.
+ */
+export type EdgeWorkerConfig = QueueWorkerConfig;
+
+/**
+ * EdgeWorker is the main entry point for creating and starting edge workers.
+ *
+ * It provides a simple interface for starting a worker that processes messages from a queue.
+ *
+ * @example
+ * ```typescript
+ * import { EdgeWorker } from '@pgflow/edge-worker';
+ *
+ * EdgeWorker.start(async (message) => {
+ *   // Process the message
+ *   console.log('Processing message:', message);
+ * }, {
+ *   queueName: 'my-queue',
+ *   maxConcurrent: 5,
+ *   retryLimit: 3
+ * });
+ * ```
+ */
+export class EdgeWorker {
+  private static logger = getLogger('EdgeWorker');
+  private static wasCalled = false;
+
+  /**
+   * Start the EdgeWorker with the given message handler and configuration.
+   *
+   * @param handler - Function that processes each message from the queue
+   * @param config - Configuration options for the worker
+   *
+   * @example
+   * ```typescript
+   * EdgeWorker.start(handler, {
+   *   // name of the queue to poll for messages
+   *   queueName: 'tasks',
+   *
+   *   // how many tasks are processed at the same time
+   *   maxConcurrent: 10,
+   *
+   *   // how many connections to the database are opened
+   *   maxPgConnections: 4,
+   *
+   *   // in-worker polling interval
+   *   maxPollSeconds: 5,
+   *
+   *   // in-database polling interval
+   *   pollIntervalMs: 200,
+   *
+   *   // how long to wait before retrying a failed job
+   *   retryDelay: 5,
+   *
+   *   // how many times to retry a failed job
+   *   retryLimit: 5,
+   *
+   *   // how long a job is invisible after reading
+   *   // if not successful, will reappear after this time
+   *   visibilityTimeout: 3,
+   * });
+   * ```
+   */
+  static start<TPayload extends Json = Json>(
+    handler: (message: TPayload) => Promise<void> | void,
+    config: EdgeWorkerConfig = {}
+  ) {
+    this.ensureFirstCall();
+
+    // Get connection string from config or environment
+    const connectionString =
+      config.connectionString || this.getConnectionString();
+
+    // Create a complete configuration object with defaults
+    const completeConfig: EdgeWorkerConfig = {
+      // Pass through any config options first
+      ...config,
+
+      // Then override with defaults for missing values
+      queueName: config.queueName || 'tasks',
+      maxConcurrent: config.maxConcurrent ?? 10,
+      maxPgConnections: config.maxPgConnections ?? 4,
+      maxPollSeconds: config.maxPollSeconds ?? 5,
+      pollIntervalMs: config.pollIntervalMs ?? 200,
+      retryDelay: config.retryDelay ?? 5,
+      retryLimit: config.retryLimit ?? 5,
+      visibilityTimeout: config.visibilityTimeout ?? 3,
+
+      // Ensure connectionString is always set
+      connectionString,
+    };
+
+    this.setupRequestHandler(handler, completeConfig);
+  }
+
+  private static ensureFirstCall() {
+    if (this.wasCalled) {
+      throw new Error('EdgeWorker.start() can only be called once');
+    }
+    this.wasCalled = true;
+  }
+
+  private static getConnectionString(): string {
+    // @ts-ignore - TODO: fix the types
+    const connectionString = Deno.env.get('EDGE_WORKER_DB_URL');
+    if (!connectionString) {
+      const message =
+        'EDGE_WORKER_DB_URL is not set!\n' +
+        'See https://pgflow.pages.dev/edge-worker/prepare-environment/#prepare-connection-string';
+      throw new Error(message);
+    }
+    return connectionString;
+  }
+
+  private static setupShutdownHandler(worker: Worker) {
+    globalThis.onbeforeunload = async () => {
+      if (worker.edgeFunctionName) {
+        await spawnNewEdgeFunction(worker.edgeFunctionName);
+      }
+
+      worker.stop();
+    };
+
+    // use waitUntil to prevent the function from exiting
+    // @ts-ignore: TODO: fix the types
+    EdgeRuntime.waitUntil(new Promise(() => {}));
+  }
+
+  private static setupRequestHandler<TPayload extends Json>(
+    handler: (message: TPayload) => Promise<void> | void,
+    workerConfig: EdgeWorkerConfig
+  ) {
+    let worker: Worker | null = null;
+
+    Deno.serve({}, (req) => {
+      if (!worker) {
+        const edgeFunctionName = this.extractFunctionName(req);
+        const sbExecutionId = Deno.env.get('SB_EXECUTION_ID')!;
+        setupLogger(sbExecutionId);
+
+        this.logger.info(`HTTP Request: ${edgeFunctionName}`);
+        // Create the worker with all configuration options
+
+        worker = createQueueWorker(handler, workerConfig);
+        worker.startOnlyOnce({
+          edgeFunctionName,
+          workerId: sbExecutionId,
+        });
+
+        this.setupShutdownHandler(worker);
+      }
+
+      return new Response('ok', {
+        headers: { 'Content-Type': 'application/json' },
+      });
+    });
+  }
+
+  private static extractFunctionName(req: Request): string {
+    return new URL(req.url).pathname.replace(/^\/+|\/+$/g, '');
+  }
+}

package/src/core/BatchProcessor.ts

@@ -0,0 +1,38 @@
+import type { ExecutionController } from './ExecutionController.ts';
+import type { IMessage, IPoller } from './types.ts';
+import { getLogger } from './Logger.ts';
+
+export class BatchProcessor<TMessage extends IMessage> {
+  private logger = getLogger('BatchProcessor');
+
+  constructor(
+    private executionController: ExecutionController<TMessage>,
+    private poller: IPoller<TMessage>,
+    private signal: AbortSignal
+  ) {
+    this.executionController = executionController;
+    this.signal = signal;
+    this.poller = poller;
+  }
+
+  async processBatch() {
+    this.logger.debug('Polling for new batch of messages...');
+    const messageRecords = await this.poller.poll();
+
+    if (this.signal.aborted) {
+      this.logger.info('Discarding messageRecords because worker is stopping');
+      return;
+    }
+
+    this.logger.debug(`Starting ${messageRecords.length} messages`);
+
+    const startPromises = messageRecords.map(
+      (message) => this.executionController.start(message)
+    );
+    await Promise.all(startPromises);
+  }
+
+  async awaitCompletion() {
+    return await this.executionController.awaitCompletion();
+  }
+}

package/src/core/ExecutionController.ts

@@ -0,0 +1,51 @@
+import { newQueue, type Queue as PromiseQueue } from '@jsr/henrygd__queue';
+import type { IExecutor, IMessage } from './types.ts';
+import { getLogger } from './Logger.ts';
+
+export interface ExecutionConfig {
+  maxConcurrent: number;
+}
+
+export class ExecutionController<TMessage extends IMessage> {
+  private logger = getLogger('ExecutionController');
+  private promiseQueue: PromiseQueue;
+  private signal: AbortSignal;
+  private createExecutor: (record: TMessage, signal: AbortSignal) => IExecutor;
+
+  constructor(
+    executorFactory: (record: TMessage, signal: AbortSignal) => IExecutor,
+    abortSignal: AbortSignal,
+    config: ExecutionConfig
+  ) {
+    this.signal = abortSignal;
+    this.createExecutor = executorFactory;
+    this.promiseQueue = newQueue(config.maxConcurrent);
+  }
+
+  async start(record: TMessage) {
+    const executor = this.createExecutor(record, this.signal);
+
+    this.logger.info(`Scheduling execution of task ${executor.msgId}`);
+
+    return await this.promiseQueue.add(async () => {
+      try {
+        this.logger.debug(`Executing task ${executor.msgId}...`);
+        await executor.execute();
+        this.logger.debug(`Execution successful for ${executor.msgId}`);
+      } catch (error) {
+        this.logger.error(`Execution failed for ${executor.msgId}:`, error);
+        throw error;
+      }
+    });
+  }
+
+  async awaitCompletion() {
+    const active = this.promiseQueue.active();
+    const all = this.promiseQueue.size();
+
+    this.logger.debug(
+      `Awaiting completion of all tasks... (active/all: ${active}}/${all})`
+    );
+    await this.promiseQueue.done();
+  }
+}

package/src/core/Heartbeat.ts

@@ -0,0 +1,23 @@
+import type { Queries } from './Queries.ts';
+import type { WorkerRow } from './types.ts';
+import { getLogger } from './Logger.ts';
+
+export class Heartbeat {
+  private logger = getLogger('Heartbeat');
+  private lastHeartbeat = 0;
+
+  constructor(
+    private interval: number,
+    private queries: Queries,
+    private workerRow: WorkerRow
+  ) {}
+
+  async send(): Promise<void> {
+    const now = Date.now();
+    if (now - this.lastHeartbeat >= this.interval) {
+      await this.queries.sendHeartbeat(this.workerRow);
+      this.logger.debug('OK');
+      this.lastHeartbeat = now;
+    }
+  }
+}

package/src/core/Logger.ts

@@ -0,0 +1,69 @@
+import * as log from '@jsr/std__log';
+
+function getLogLevelFromEnv(): log.LevelName {
+  const validLevels = [
+    // 'NOTSET',
+    'DEBUG',
+    'INFO',
+    // 'WARNING',
+    'ERROR',
+    // 'CRITICAL',
+  ];
+  const logLevel = Deno.env.get('EDGE_WORKER_LOG_LEVEL')?.toUpperCase();
+
+  if (logLevel && !validLevels.includes(logLevel)) {
+    console.warn(`Invalid log level "${logLevel}". Using "INFO" instead.`);
+    return 'INFO';
+  }
+
+  return (logLevel as log.LevelName) || 'INFO';
+}
+
+const defaultLoggerConfig: log.LoggerConfig = {
+  level: 'DEBUG',
+  handlers: ['console'],
+};
+
+export function setupLogger(workerId: string) {
+  log.setup({
+    handlers: {
+      console: new log.ConsoleHandler(getLogLevelFromEnv(), {
+        formatter: (record: {
+          loggerName: string;
+          msg: string;
+          args: unknown[];
+        }) => {
+          const prefix = `worker_id=${workerId}`;
+          const module = record.loggerName;
+          const msg = record.msg;
+
+          // If there are additional args, pretty print them using console.log
+          if (record.args.length > 0) {
+            return `${prefix} [${module}] ${msg}`;
+          }
+
+          return `${prefix} [${module}] ${msg}`;
+        },
+        useColors: true,
+      }),
+    },
+
+    loggers: {
+      BatchProcessor: defaultLoggerConfig,
+      EdgeWorker: defaultLoggerConfig,
+      ExecutionController: defaultLoggerConfig,
+      Heartbeat: defaultLoggerConfig,
+      Logger: defaultLoggerConfig,
+      MessageExecutor: defaultLoggerConfig,
+      Worker: defaultLoggerConfig,
+      WorkerLifecycle: defaultLoggerConfig,
+      WorkerState: defaultLoggerConfig,
+      spawnNewEdgeFunction: defaultLoggerConfig,
+    },
+  });
+}
+
+// Helper function to get logger for specific module
+export function getLogger(module: string) {
+  return log.getLogger(module);
+}

package/src/core/Queries.ts

@@ -0,0 +1,44 @@
+import type postgres from 'postgres';
+import type { WorkerRow } from './types.ts';
+
+export class Queries {
+  constructor(private readonly sql: postgres.Sql) {}
+
+  async onWorkerStarted({
+    queueName,
+    workerId,
+    edgeFunctionName,
+  }: {
+    queueName: string;
+    workerId: string;
+    edgeFunctionName: string;
+  }): Promise<WorkerRow> {
+    const [worker] = await this.sql<WorkerRow[]>`
+      INSERT INTO edge_worker.workers (queue_name, worker_id, function_name)
+      VALUES (${queueName}, ${workerId}, ${edgeFunctionName})
+      RETURNING *;
+    `;
+
+    return worker;
+  }
+
+  async onWorkerStopped(workerRow: WorkerRow): Promise<WorkerRow> {
+    const [worker] = await this.sql<WorkerRow[]>`
+      UPDATE edge_worker.workers AS w
+      SET stopped_at = clock_timestamp(), last_heartbeat_at = clock_timestamp()
+      WHERE w.worker_id = ${workerRow.worker_id}
+      RETURNING *;
+    `;
+
+    return worker;
+  }
+
+  async sendHeartbeat(workerRow: WorkerRow): Promise<void> {
+    await this.sql<WorkerRow[]>`
+      UPDATE edge_worker.workers AS w
+      SET last_heartbeat_at = clock_timestamp()
+      WHERE w.worker_id = ${workerRow.worker_id}
+      RETURNING *;
+    `;
+  }
+}

package/src/core/Worker.ts

@@ -0,0 +1,102 @@
+import type postgres from 'postgres';
+import type { IBatchProcessor, ILifecycle, WorkerBootstrap } from './types.ts';
+import { getLogger, setupLogger } from './Logger.ts';
+
+export class Worker {
+  private lifecycle: ILifecycle;
+  private logger = getLogger('Worker');
+  private abortController = new AbortController();
+
+  private batchProcessor: IBatchProcessor;
+  private sql: postgres.Sql;
+
+  constructor(
+    batchProcessor: IBatchProcessor,
+    lifecycle: ILifecycle,
+    sql: postgres.Sql,
+  ) {
+    this.sql = sql;
+
+    this.lifecycle = lifecycle;
+
+    this.batchProcessor = batchProcessor;
+  }
+
+  async startOnlyOnce(workerBootstrap: WorkerBootstrap) {
+    if (this.lifecycle.isRunning) {
+      this.logger.debug('Worker already running, ignoring start request');
+      return;
+    }
+
+    await this.start(workerBootstrap);
+  }
+
+  private async start(workerBootstrap: WorkerBootstrap) {
+    setupLogger(workerBootstrap.workerId);
+
+    try {
+      await this.lifecycle.acknowledgeStart(workerBootstrap);
+
+      while (this.isMainLoopActive) {
+        try {
+          await this.lifecycle.sendHeartbeat();
+        } catch (error: unknown) {
+          this.logger.error(`Error sending heartbeat: ${error}`);
+          // Continue execution - a failed heartbeat shouldn't stop processing
+        }
+
+        try {
+          await this.batchProcessor.processBatch();
+        } catch (error: unknown) {
+          this.logger.error(`Error processing batch: ${error}`);
+          // Continue to next iteration - failed batch shouldn't stop the worker
+        }
+      }
+    } catch (error) {
+      this.logger.error(`Error in worker main loop: ${error}`);
+      throw error;
+    }
+  }
+
+  async stop() {
+    // If the worker is already stopping or stopped, do nothing
+    if (this.lifecycle.isStopping || this.lifecycle.isStopped) {
+      return;
+    }
+
+    this.lifecycle.transitionToStopping();
+
+    try {
+      this.logger.info('-> Stopped accepting new messages');
+      this.abortController.abort();
+
+      this.logger.info('-> Waiting for pending tasks to complete...');
+      await this.batchProcessor.awaitCompletion();
+      this.logger.info('-> Pending tasks completed!');
+
+      this.lifecycle.acknowledgeStop();
+
+      this.logger.info('-> Closing SQL connection...');
+      await this.sql.end();
+      this.logger.info('-> SQL connection closed!');
+    } catch (error) {
+      this.logger.info(`Error during worker stop: ${error}`);
+      throw error;
+    }
+  }
+
+  get edgeFunctionName() {
+    return this.lifecycle.edgeFunctionName;
+  }
+
+  /**
+   * Returns true if worker state is Running and worker was not stopped
+   */
+  private get isMainLoopActive() {
+    return this.lifecycle.isRunning && !this.isAborted;
+  }
+
+  private get isAborted() {
+    return this.abortController.signal.aborted;
+  }
+}

package/src/core/WorkerLifecycle.ts

@@ -0,0 +1,93 @@
+import { Heartbeat } from './Heartbeat.ts';
+import { getLogger } from './Logger.ts';
+import type { Queries } from './Queries.ts';
+import type { Queue } from '../queue/Queue.ts';
+import type { ILifecycle, Json, WorkerBootstrap, WorkerRow } from './types.ts';
+import { States, WorkerState } from './WorkerState.ts';
+
+export interface LifecycleConfig {
+  queueName: string;
+}
+
+export class WorkerLifecycle<IMessage extends Json> implements ILifecycle {
+  private workerState: WorkerState = new WorkerState();
+  private heartbeat?: Heartbeat;
+  private logger = getLogger('WorkerLifecycle');
+  private queries: Queries;
+  private queue: Queue<IMessage>;
+  private workerRow?: WorkerRow;
+
+  constructor(queries: Queries, queue: Queue<IMessage>) {
+    this.queries = queries;
+    this.queue = queue;
+  }
+
+  async acknowledgeStart(workerBootstrap: WorkerBootstrap): Promise<void> {
+    this.workerState.transitionTo(States.Starting);
+
+    this.logger.info(`Ensuring queue '${this.queue.queueName}' exists...`);
+    await this.queue.safeCreate();
+
+    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...');
+
+      // TODO: commented out because we can live without this
+      //       but it is causing problems with DbHandler - workes does not have
+      //       enough time to fire this query before hard-terimnated
+      //       We can always check the heartbeat to see if it is still running
+      //
+      // await this.queries.onWorkerStopped(this.workerRow);
+
+      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.queue.queueName;
+  }
+
+  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/core/WorkerState.ts

@@ -0,0 +1,85 @@
+import { getLogger } from './Logger.ts';
+
+export enum States {
+  /** The worker has been created but has not yet started. */
+  Created = 'created',
+
+  /** The worker is starting but has not yet started processing messages. */
+  Starting = 'starting',
+
+  /** The worker is processing messages. */
+  Running = 'running',
+
+  /** The worker stopped processing messages but is still releasing resources. */
+  Stopping = 'stopping',
+
+  /** The worker has stopped processing messages and released resources
+   * and can be discarded. */
+  Stopped = 'stopped',
+}
+
+export const Transitions: Record<States, States[]> = {
+  [States.Created]: [States.Starting],
+  [States.Starting]: [States.Running],
+  [States.Running]: [States.Stopping],
+  [States.Stopping]: [States.Stopped],
+  [States.Stopped]: [], // Terminal state - no valid transitions from here
+};
+
+export class TransitionError extends Error {
+  constructor(options: { from: States; to: States }) {
+    super(`Cannot transition from ${options.from} to ${options.to}`);
+  }
+}
+
+/**
+ * Represents the state of a worker and exposes method for doing allowed transitions
+ */
+export class WorkerState {
+  private logger = getLogger('WorkerState');
+  private state: States = States.Created;
+
+  get current() {
+    return this.state;
+  }
+
+  get isCreated() {
+    return this.state === States.Created;
+  }
+
+  get isStarting() {
+    return this.state === States.Starting;
+  }
+
+  get isRunning() {
+    return this.state === States.Running;
+  }
+
+  get isStopping() {
+    return this.state === States.Stopping;
+  }
+
+  get isStopped() {
+    return this.state === States.Stopped;
+  }
+
+  transitionTo(state: States) {
+    this.logger.debug(
+      `[WorkerState] Starting transition to '${state}' (current state: ${this.state})`
+    );
+
+    if (this.state === state) {
+      return;
+    }
+
+    if (Transitions[this.state].includes(state)) {
+      this.state = state;
+      this.logger.debug(`[WorkerState] Transitioned to '${state}'`);
+    } else {
+      throw new TransitionError({
+        from: this.state,
+        to: state,
+      });
+    }
+  }
+}