Add maxPollSeconds and pollIntervalMs options to flow worker (#90)

* feat: add configurable polling intervals and batch size options to EdgeWorker and related
components

- Introduced maxPollSeconds and pollIntervalMs options in EdgeWorker configuration
- Updated StepTaskPoller to accept and utilize maxPollSeconds and pollIntervalMs
- Extended createFlowWorker to support polling interval and max poll seconds
- Modified tests to include new polling configuration options for better control and testing
- Ensured default values are set when options are not provided
- Updated helper functions to pass new configuration parameters for testing purposes

* fix(core): move visibilityTimeout as last opt in pollForTasks

* refactor(worker): add promise tracking for main loop to improve shutdown handling

- Introduced mainLoopPromise to track the worker's main loop execution
- Updated startOnlyOnce to prevent multiple concurrent starts
- Enhanced stop method to await main loop completion and handle errors
- Improved worker shutdown to ensure all tasks are finalized before stopping

* fix: improve worker shutdown and task polling behavior

- Make worker.stop() wait for the main loop promise to ensure proper shutdown
- Adjust pollForTasks to make visibilityTimeout the last option, allowing it to be skipped if
not needed

* fix: remove unused import in integration test helpers

- Deleted the import of createFakeLogger from _helpers.ts to clean up unused code
- Minor cleanup to improve code clarity and maintainability

Author: jumski

.changeset/bumpy-words-scream.md

@@ -0,0 +1,5 @@
+---
+'@pgflow/edge-worker': patch
+---
+
+Make worker.stop() wait for the main loop promise

.changeset/dirty-moose-reply.md

@@ -0,0 +1,5 @@
+---
+'@pgflow/core': patch
+---
+
+Make visibilityTimeout the last option to pollForTasks so it can be skipped

pkgs/core/src/PgflowSqlClient.ts

@@ -19,9 +19,9 @@ export class PgflowSqlClient<TFlow extends AnyFlow>
   async pollForTasks(
     queueName: string,
     batchSize = 20,
-    visibilityTimeout = 2,
     maxPollSeconds = 5,
-    pollIntervalMs = 200
+    pollIntervalMs = 200,
+    visibilityTimeout = 2
   ): Promise<StepTaskRecord<TFlow>[]> {
     return await this.sql<StepTaskRecord<TFlow>[]>`
       SELECT *

pkgs/edge-worker/src/EdgeWorker.ts

@@ -186,6 +186,12 @@ export class EdgeWorker {
    *
    *   // batch size for polling messages
    *   batchSize: 10,
+   *
+   *   // in-worker polling interval
+   *   maxPollSeconds: 2,
+   *
+   *   // in-database polling interval
+   *   pollIntervalMs: 100,
    * });
    * ```
    */
@@ -204,6 +210,8 @@ export class EdgeWorker {
       maxConcurrent: config.maxConcurrent ?? 10,
       maxPgConnections: config.maxPgConnections ?? 4,
       batchSize: config.batchSize ?? 10,
+      maxPollSeconds: config.maxPollSeconds ?? 2,
+      pollIntervalMs: config.pollIntervalMs ?? 100,
       connectionString:
         config.connectionString || this.platform.getConnectionString(),
     };

pkgs/edge-worker/src/core/Worker.ts

@@ -9,6 +9,7 @@ export class Worker {
 
   private batchProcessor: IBatchProcessor;
   private sql: postgres.Sql;
+  private mainLoopPromise: Promise<void> | undefined;
 
   constructor(
     batchProcessor: IBatchProcessor,
@@ -22,13 +23,15 @@ export class Worker {
     this.logger = logger;
   }
 
-  async startOnlyOnce(workerBootstrap: WorkerBootstrap) {
+  startOnlyOnce(workerBootstrap: WorkerBootstrap) {
     if (this.lifecycle.isRunning) {
       this.logger.debug('Worker already running, ignoring start request');
       return;
     }
 
-    await this.start(workerBootstrap);
+    if (!this.mainLoopPromise) {
+      this.mainLoopPromise = this.start(workerBootstrap);
+    }
   }
 
   private async start(workerBootstrap: WorkerBootstrap) {
@@ -68,6 +71,16 @@ export class Worker {
       this.logger.info('-> Stopped accepting new messages');
       this.abortController.abort();
 
+      try {
+        this.logger.debug('-> Waiting for main loop to complete');
+        await this.mainLoopPromise;
+      } catch (error) {
+        this.logger.error(
+          `Error in main loop: ${error}. Continuing to stop worker`
+        );
+        throw error;
+      }
+
       this.logger.info('-> Waiting for pending tasks to complete...');
       await this.batchProcessor.awaitCompletion();
       this.logger.info('-> Pending tasks completed!');

pkgs/edge-worker/src/flow/StepTaskPoller.ts

@@ -6,6 +6,8 @@ import type { AnyFlow } from '@pgflow/dsl';
 export interface StepTaskPollerConfig {
   batchSize: number;
   queueName: string;
+  maxPollSeconds?: number;
+  pollIntervalMs?: number;
 }
 
 /**
@@ -32,13 +34,16 @@ export class StepTaskPoller<TFlow extends AnyFlow>
     }
 
     this.logger.debug(
-      `Polling for flow tasks with batch size ${this.config.batchSize}`
+      `Polling for flow tasks with batch size ${this.config.batchSize}, maxPollSeconds: ${this.config.maxPollSeconds}, pollIntervalMs: ${this.config.pollIntervalMs}`
     );
 
     try {
+      // Pass polling configuration to the adapter if they're provided
       const tasks = await this.adapter.pollForTasks(
         this.config.queueName,
-        this.config.batchSize
+        this.config.batchSize,
+        this.config.maxPollSeconds,
+        this.config.pollIntervalMs
       );
       this.logger.debug(`Retrieved ${tasks.length} flow tasks`);
       return tasks;

pkgs/edge-worker/src/flow/createFlowWorker.ts

@@ -44,6 +44,18 @@ export type FlowWorkerConfig = {
    * @default 10
    */
   batchSize?: number;
+
+  /**
+   * In-worker polling interval in seconds
+   * @default 2
+   */
+  maxPollSeconds?: number;
+
+  /**
+   * In-database polling interval in milliseconds
+   * @default 100
+   */
+  pollIntervalMs?: number;
 };
 
 /**
@@ -96,6 +108,8 @@ export function createFlowWorker<TFlow extends AnyFlow>(
   const pollerConfig: StepTaskPollerConfig = {
     batchSize: config.batchSize || 10,
     queueName: flow.slug,
+    maxPollSeconds: config.maxPollSeconds || 2,
+    pollIntervalMs: config.pollIntervalMs || 100,
   };
   const poller = new StepTaskPoller<TFlow>(
     pgflowAdapter,

pkgs/edge-worker/tests/integration/_helpers.ts

@@ -4,7 +4,6 @@ import {
   type FlowWorkerConfig,
 } from '../../src/flow/createFlowWorker.ts';
 import type { postgres } from '../sql.ts';
-import { createFakeLogger } from '../fakes.ts';
 import { PgflowSqlClient } from '@pgflow/core';
 
 export async function startFlow<TFlow extends AnyFlow>(
@@ -33,7 +32,14 @@ export function startWorker<
     ...options,
   };
 
-  const worker = createFlowWorker(flow, mergedOptions, createFakeLogger);
+  const consoleLogger = {
+    debug: console.log,
+    info: console.log,
+    warn: console.warn,
+    error: console.error,
+  };
+
+  const worker = createFlowWorker(flow, mergedOptions, () => consoleLogger);
 
   worker.startOnlyOnce({
     edgeFunctionName: 'test_flow',

pkgs/edge-worker/tests/integration/flow/minimalFlow.test.ts

@@ -30,6 +30,8 @@ Deno.test(
     const worker = startWorker(sql, MinimalFlow, {
       maxConcurrent: 1,
       batchSize: 10,
+      maxPollSeconds: 1,
+      pollIntervalMs: 200,
     });
 
     try {
@@ -114,8 +116,10 @@ Deno.test(
         'Run output should match expected value'
       );
     } finally {
+      console.log('Stopping worker');
       // Stop the worker
       await worker.stop();
+      console.log('Worker stopped');
     }
   })
 );