@nicnocquee/dataqueue 1.24.0 → 1.25.0

package/src/processor.ts

@@ -1,5 +1,5 @@
-import { Pool } from 'pg';
 import { Worker } from 'worker_threads';
+import { Pool } from 'pg';
 import {
   JobRecord,
   ProcessorOptions,
@@ -8,15 +8,76 @@ import {
   JobType,
   FailureReason,
   JobHandlers,
+  JobContext,
+  OnTimeoutCallback,
+  WaitSignal,
+  WaitDuration,
+  WaitTokenResult,
 } from './types.js';
+import { QueueBackend } from './backend.js';
+import { PostgresBackend } from './backends/postgres.js';
 import {
-  getNextBatch,
-  completeJob,
-  failJob,
-  setPendingReasonForUnpickedJobs,
+  waitJob,
+  updateStepData,
+  createWaitpoint,
+  getWaitpoint,
 } from './queue.js';
 import { log, setLogContext } from './log-context.js';
 
+/**
+ * Try to extract the underlying pg Pool from a QueueBackend.
+ * Returns null for non-PostgreSQL backends.
+ */
+function tryExtractPool(backend: QueueBackend): Pool | null {
+  if (backend instanceof PostgresBackend) {
+    return backend.getPool();
+  }
+  return null;
+}
+
+/**
+ * Build a JobContext without wait support (for non-PostgreSQL backends).
+ * prolong/onTimeout work normally; wait-related methods throw helpful errors.
+ */
+function buildBasicContext(
+  backend: QueueBackend,
+  jobId: number,
+  baseCtx: {
+    prolong: JobContext['prolong'];
+    onTimeout: JobContext['onTimeout'];
+  },
+): JobContext {
+  const waitError = () =>
+    new Error(
+      'Wait features (waitFor, waitUntil, createToken, waitForToken, ctx.run) are currently only supported with the PostgreSQL backend.',
+    );
+  return {
+    prolong: baseCtx.prolong,
+    onTimeout: baseCtx.onTimeout,
+    run: async <T>(_stepName: string, fn: () => Promise<T>): Promise<T> => {
+      // Without PostgreSQL, just execute the function directly (no persistence)
+      return fn();
+    },
+    waitFor: async () => {
+      throw waitError();
+    },
+    waitUntil: async () => {
+      throw waitError();
+    },
+    createToken: async () => {
+      throw waitError();
+    },
+    waitForToken: async () => {
+      throw waitError();
+    },
+    setProgress: async (percent: number) => {
+      if (percent < 0 || percent > 100)
+        throw new Error('Progress must be between 0 and 100');
+      await backend.updateProgress(jobId, Math.round(percent));
+    },
+  };
+}
+
 /**
  * Validates that a handler can be serialized for worker thread execution.
  * Throws an error with helpful message if serialization fails.
@@ -253,6 +314,254 @@ async function runHandlerInWorker<
   });
 }
 
+/**
+ * Convert a WaitDuration to a target Date.
+ */
+function calculateWaitUntil(duration: WaitDuration): Date {
+  const now = Date.now();
+  let ms = 0;
+  if (duration.seconds) ms += duration.seconds * 1000;
+  if (duration.minutes) ms += duration.minutes * 60 * 1000;
+  if (duration.hours) ms += duration.hours * 60 * 60 * 1000;
+  if (duration.days) ms += duration.days * 24 * 60 * 60 * 1000;
+  if (duration.weeks) ms += duration.weeks * 7 * 24 * 60 * 60 * 1000;
+  if (duration.months) ms += duration.months * 30 * 24 * 60 * 60 * 1000;
+  if (duration.years) ms += duration.years * 365 * 24 * 60 * 60 * 1000;
+  if (ms <= 0) {
+    throw new Error(
+      'waitFor duration must be positive. Provide at least one positive duration field.',
+    );
+  }
+  return new Date(now + ms);
+}
+
+/**
+ * Create a no-op JobContext for cases where prolong/onTimeout are not supported
+ * (e.g. forceKillOnTimeout mode or no timeout set).
+ */
+function createNoOpContext(
+  backend: QueueBackend,
+  jobId: number,
+  reason: string,
+): JobContext {
+  return {
+    prolong: () => {
+      log(`prolong() called but ignored: ${reason}`);
+    },
+    onTimeout: () => {
+      log(`onTimeout() called but ignored: ${reason}`);
+    },
+    run: async <T>(_stepName: string, fn: () => Promise<T>): Promise<T> => {
+      // In no-op context (forceKillOnTimeout), just execute the function directly
+      return fn();
+    },
+    waitFor: async () => {
+      throw new Error(
+        `waitFor() is not supported when forceKillOnTimeout is enabled. ${reason}`,
+      );
+    },
+    waitUntil: async () => {
+      throw new Error(
+        `waitUntil() is not supported when forceKillOnTimeout is enabled. ${reason}`,
+      );
+    },
+    createToken: async () => {
+      throw new Error(
+        `createToken() is not supported when forceKillOnTimeout is enabled. ${reason}`,
+      );
+    },
+    waitForToken: async () => {
+      throw new Error(
+        `waitForToken() is not supported when forceKillOnTimeout is enabled. ${reason}`,
+      );
+    },
+    setProgress: async (percent: number) => {
+      if (percent < 0 || percent > 100)
+        throw new Error('Progress must be between 0 and 100');
+      await backend.updateProgress(jobId, Math.round(percent));
+    },
+  };
+}
+
+/**
+ * Pre-process stepData before handler re-invocation.
+ * Marks pending waits as completed and fetches token outputs.
+ */
+async function resolveCompletedWaits(
+  pool: Pool,
+  stepData: Record<string, any>,
+): Promise<void> {
+  for (const key of Object.keys(stepData)) {
+    if (!key.startsWith('__wait_')) continue;
+    const entry = stepData[key];
+    if (!entry || typeof entry !== 'object' || entry.completed) continue;
+
+    if (entry.type === 'duration' || entry.type === 'date') {
+      // Time-based wait has elapsed (we got picked up, so it must have)
+      stepData[key] = { ...entry, completed: true };
+    } else if (entry.type === 'token' && entry.tokenId) {
+      // Token-based wait -- fetch the waitpoint result
+      const wp = await getWaitpoint(pool, entry.tokenId);
+      if (wp && wp.status === 'completed') {
+        stepData[key] = {
+          ...entry,
+          completed: true,
+          result: { ok: true, output: wp.output },
+        };
+      } else if (wp && wp.status === 'timed_out') {
+        stepData[key] = {
+          ...entry,
+          completed: true,
+          result: { ok: false, error: 'Token timed out' },
+        };
+      }
+      // If still waiting (shouldn't happen), leave as pending
+    }
+  }
+}
+
+/**
+ * Build the extended JobContext with step tracking and wait support.
+ */
+function buildWaitContext(
+  backend: QueueBackend,
+  pool: Pool,
+  jobId: number,
+  stepData: Record<string, any>,
+  baseCtx: {
+    prolong: JobContext['prolong'];
+    onTimeout: JobContext['onTimeout'];
+  },
+): JobContext {
+  // Wait counter always starts at 0 per invocation.
+  // The handler replays from the top each time, so the counter position
+  // must match the order of waitFor/waitUntil/waitForToken calls in code.
+  let waitCounter = 0;
+
+  const ctx: JobContext = {
+    prolong: baseCtx.prolong,
+    onTimeout: baseCtx.onTimeout,
+
+    run: async <T>(stepName: string, fn: () => Promise<T>): Promise<T> => {
+      // Check if step was already completed in a previous invocation
+      const cached = stepData[stepName];
+      if (cached && typeof cached === 'object' && cached.__completed) {
+        log(`Step "${stepName}" replayed from cache for job ${jobId}`);
+        return cached.result as T;
+      }
+
+      // Execute the step
+      const result = await fn();
+
+      // Persist step result
+      stepData[stepName] = { __completed: true, result };
+      await updateStepData(pool, jobId, stepData);
+
+      return result;
+    },
+
+    waitFor: async (duration: WaitDuration): Promise<void> => {
+      const waitKey = `__wait_${waitCounter++}`;
+
+      // Check if this wait was already completed (from a previous invocation)
+      const cached = stepData[waitKey];
+      if (cached && typeof cached === 'object' && cached.completed) {
+        log(`Wait "${waitKey}" already completed for job ${jobId}, skipping`);
+        return;
+      }
+
+      // Calculate when to resume
+      const waitUntilDate = calculateWaitUntil(duration);
+
+      // Record this wait as pending in step data
+      stepData[waitKey] = { type: 'duration', completed: false };
+
+      // Throw WaitSignal to pause the handler
+      throw new WaitSignal('duration', waitUntilDate, undefined, stepData);
+    },
+
+    waitUntil: async (date: Date): Promise<void> => {
+      const waitKey = `__wait_${waitCounter++}`;
+
+      // Check if this wait was already completed
+      const cached = stepData[waitKey];
+      if (cached && typeof cached === 'object' && cached.completed) {
+        log(`Wait "${waitKey}" already completed for job ${jobId}, skipping`);
+        return;
+      }
+
+      // Record this wait as pending
+      stepData[waitKey] = { type: 'date', completed: false };
+
+      // Throw WaitSignal to pause the handler
+      throw new WaitSignal('date', date, undefined, stepData);
+    },
+
+    createToken: async (options?) => {
+      const token = await createWaitpoint(pool, jobId, options);
+      return token;
+    },
+
+    waitForToken: async <T = any>(
+      tokenId: string,
+    ): Promise<WaitTokenResult<T>> => {
+      const waitKey = `__wait_${waitCounter++}`;
+
+      // Check if this wait was already completed
+      const cached = stepData[waitKey];
+      if (cached && typeof cached === 'object' && cached.completed) {
+        log(
+          `Token wait "${waitKey}" already completed for job ${jobId}, returning cached result`,
+        );
+        return cached.result as WaitTokenResult<T>;
+      }
+
+      // Check if the token is already completed (e.g., completed while job was still processing)
+      const wp = await getWaitpoint(pool, tokenId);
+      if (wp && wp.status === 'completed') {
+        const result: WaitTokenResult<T> = {
+          ok: true,
+          output: wp.output as T,
+        };
+        stepData[waitKey] = {
+          type: 'token',
+          tokenId,
+          completed: true,
+          result,
+        };
+        await updateStepData(pool, jobId, stepData);
+        return result;
+      }
+      if (wp && wp.status === 'timed_out') {
+        const result: WaitTokenResult<T> = {
+          ok: false,
+          error: 'Token timed out',
+        };
+        stepData[waitKey] = {
+          type: 'token',
+          tokenId,
+          completed: true,
+          result,
+        };
+        await updateStepData(pool, jobId, stepData);
+        return result;
+      }
+
+      // Token not yet completed -- save pending state and throw WaitSignal
+      stepData[waitKey] = { type: 'token', tokenId, completed: false };
+      throw new WaitSignal('token', undefined, tokenId, stepData);
+    },
+
+    setProgress: async (percent: number) => {
+      if (percent < 0 || percent > 100)
+        throw new Error('Progress must be between 0 and 100');
+      await backend.updateProgress(jobId, Math.round(percent));
+    },
+  };
+
+  return ctx;
+}
+
 /**
  * Process a single job using the provided handler map
  */
@@ -260,20 +569,18 @@ export async function processJobWithHandlers<
   PayloadMap,
   T extends keyof PayloadMap & string,
 >(
-  pool: Pool,
+  backend: QueueBackend,
   job: JobRecord<PayloadMap, T>,
   jobHandlers: JobHandlers<PayloadMap>,
 ): Promise<void> {
   const handler = jobHandlers[job.jobType];
 
   if (!handler) {
-    await setPendingReasonForUnpickedJobs(
-      pool,
+    await backend.setPendingReasonForUnpickedJobs(
       `No handler registered for job type: ${job.jobType}`,
       job.jobType,
     );
-    await failJob(
-      pool,
+    await backend.failJob(
       job.id,
       new Error(`No handler registered for job type: ${job.jobType}`),
       FailureReason.NoHandler,
@@ -281,6 +588,22 @@ export async function processJobWithHandlers<
     return;
   }
 
+  // Load step data (may contain completed steps from previous invocations)
+  const stepData: Record<string, any> = { ...(job.stepData || {}) };
+
+  // Try to get pool for wait features (PostgreSQL-only)
+  const pool = tryExtractPool(backend);
+
+  // If resuming from a wait, resolve any pending wait entries
+  const hasStepHistory = Object.keys(stepData).some((k) =>
+    k.startsWith('__wait_'),
+  );
+  if (hasStepHistory && pool) {
+    await resolveCompletedWaits(pool, stepData);
+    // Persist the resolved step data
+    await updateStepData(pool, job.id, stepData);
+  }
+
   // Per-job timeout logic
   const timeoutMs = job.timeoutMs ?? undefined;
   const forceKillOnTimeout = job.forceKillOnTimeout ?? false;
@@ -288,24 +611,100 @@ export async function processJobWithHandlers<
   const controller = new AbortController();
   try {
     // If forceKillOnTimeout is true, run handler in a worker thread
+    // Note: wait features are not available in forceKillOnTimeout mode
     if (forceKillOnTimeout && timeoutMs && timeoutMs > 0) {
       await runHandlerInWorker(handler, job.payload, timeoutMs, job.jobType);
     } else {
-      // Default: graceful shutdown with AbortController
-      const jobPromise = handler(job.payload, controller.signal);
-      if (timeoutMs && timeoutMs > 0) {
+      // Build the JobContext for prolong/onTimeout support
+      let onTimeoutCallback: OnTimeoutCallback | undefined;
+
+      // Reference to the reject function of the timeout promise so we can re-arm it
+      let timeoutReject: ((error: Error) => void) | undefined;
+
+      /**
+       * Arms (or re-arms) the timeout. When it fires:
+       * 1. If an onTimeout callback is registered, call it first.
+       *    - If it returns a positive number, re-arm with that duration and update DB.
+       *    - Otherwise, proceed with abort.
+       * 2. If no onTimeout callback, proceed with abort.
+       */
+      const armTimeout = (ms: number) => {
+        if (timeoutId) clearTimeout(timeoutId);
+        timeoutId = setTimeout(() => {
+          // Check if an onTimeout callback wants to extend
+          if (onTimeoutCallback) {
+            try {
+              const extension = onTimeoutCallback();
+              if (typeof extension === 'number' && extension > 0) {
+                // Extend: re-arm timeout and update DB
+                backend.prolongJob(job.id).catch(() => {});
+                armTimeout(extension);
+                return;
+              }
+            } catch (callbackError) {
+              log(
+                `onTimeout callback threw for job ${job.id}: ${callbackError}`,
+              );
+              // Treat as "no extension" and proceed with abort
+            }
+          }
+          // No extension -- proceed with abort
+          controller.abort();
+          const timeoutError = new Error(`Job timed out after ${ms} ms`);
+          // @ts-ignore
+          timeoutError.failureReason = FailureReason.Timeout;
+          if (timeoutReject) {
+            timeoutReject(timeoutError);
+          }
+        }, ms);
+      };
+
+      const hasTimeout = timeoutMs != null && timeoutMs > 0;
+
+      // Build base prolong/onTimeout context
+      const baseCtx = hasTimeout
+        ? {
+            prolong: (ms?: number) => {
+              const duration = ms ?? timeoutMs;
+              if (duration != null && duration > 0) {
+                armTimeout(duration);
+                // Update DB locked_at to prevent reclaimStuckJobs
+                backend.prolongJob(job.id).catch(() => {});
+              }
+            },
+            onTimeout: (callback: OnTimeoutCallback) => {
+              onTimeoutCallback = callback;
+            },
+          }
+        : {
+            prolong: () => {
+              log('prolong() called but ignored: job has no timeout set');
+            },
+            onTimeout: () => {
+              log('onTimeout() called but ignored: job has no timeout set');
+            },
+          };
+
+      // Build context: full wait support for PostgreSQL, basic for others
+      const ctx = pool
+        ? buildWaitContext(backend, pool, job.id, stepData, baseCtx)
+        : buildBasicContext(backend, job.id, baseCtx);
+
+      // If forceKillOnTimeout was set but timeoutMs was missing, warn
+      if (forceKillOnTimeout && !hasTimeout) {
+        log(
+          `forceKillOnTimeout is set but no timeoutMs for job ${job.id}, running without force kill`,
+        );
+      }
+
+      const jobPromise = handler(job.payload, controller.signal, ctx);
+
+      if (hasTimeout) {
         await Promise.race([
           jobPromise,
-          new Promise((_, reject) => {
-            timeoutId = setTimeout(() => {
-              controller.abort();
-              const timeoutError = new Error(
-                `Job timed out after ${timeoutMs} ms`,
-              );
-              // @ts-ignore
-              timeoutError.failureReason = FailureReason.Timeout;
-              reject(timeoutError);
-            }, timeoutMs);
+          new Promise<never>((_, reject) => {
+            timeoutReject = reject;
+            armTimeout(timeoutMs!);
           }),
         ]);
       } else {
@@ -313,21 +712,50 @@ export async function processJobWithHandlers<
       }
     }
     if (timeoutId) clearTimeout(timeoutId);
-    await completeJob(pool, job.id);
+
+    // Job completed successfully -- complete via backend
+    await backend.completeJob(job.id);
   } catch (error) {
     if (timeoutId) clearTimeout(timeoutId);
+
+    // Check if this is a WaitSignal (not a real error)
+    if (error instanceof WaitSignal) {
+      if (!pool) {
+        // Wait signals should never happen with non-PostgreSQL backends
+        // since the context methods throw, but guard just in case
+        await backend.failJob(
+          job.id,
+          new Error(
+            'WaitSignal received but wait features require the PostgreSQL backend.',
+          ),
+          FailureReason.HandlerError,
+        );
+        return;
+      }
+      log(
+        `Job ${job.id} entering wait: type=${error.type}, waitUntil=${error.waitUntil?.toISOString() ?? 'none'}, tokenId=${error.tokenId ?? 'none'}`,
+      );
+      await waitJob(pool, job.id, {
+        waitUntil: error.waitUntil,
+        waitTokenId: error.tokenId,
+        stepData: error.stepData,
+      });
+      return;
+    }
+
+    // Real error -- handle as failure
     console.error(`Error processing job ${job.id}:`, error);
     let failureReason = FailureReason.HandlerError;
     if (
       error &&
       typeof error === 'object' &&
       'failureReason' in error &&
-      (error as any).failureReason === FailureReason.Timeout
+      (error as { failureReason?: FailureReason }).failureReason ===
+        FailureReason.Timeout
     ) {
       failureReason = FailureReason.Timeout;
     }
-    await failJob(
-      pool,
+    await backend.failJob(
       job.id,
       error instanceof Error ? error : new Error(String(error)),
       failureReason,
@@ -339,15 +767,15 @@ export async function processJobWithHandlers<
  * Process a batch of jobs using the provided handler map and concurrency limit
  */
 export async function processBatchWithHandlers<PayloadMap>(
-  pool: Pool,
+  backend: QueueBackend,
   workerId: string,
   batchSize: number,
   jobType: string | string[] | undefined,
   jobHandlers: JobHandlers<PayloadMap>,
   concurrency?: number,
+  onError?: (error: Error) => void,
 ): Promise<number> {
-  const jobs = await getNextBatch<PayloadMap, JobType<PayloadMap>>(
-    pool,
+  const jobs = await backend.getNextBatch<PayloadMap, JobType<PayloadMap>>(
     workerId,
     batchSize,
     jobType,
@@ -355,7 +783,7 @@ export async function processBatchWithHandlers<PayloadMap>(
   if (!concurrency || concurrency >= jobs.length) {
     // Default: all in parallel
     await Promise.all(
-      jobs.map((job) => processJobWithHandlers(pool, job, jobHandlers)),
+      jobs.map((job) => processJobWithHandlers(backend, job, jobHandlers)),
     );
     return jobs.length;
   }
@@ -369,7 +797,7 @@ export async function processBatchWithHandlers<PayloadMap>(
       while (running < concurrency && idx < jobs.length) {
         const job = jobs[idx++];
         running++;
-        processJobWithHandlers(pool, job, jobHandlers)
+        processJobWithHandlers(backend, job, jobHandlers)
           .then(() => {
             running--;
             finished++;
@@ -378,6 +806,9 @@ export async function processBatchWithHandlers<PayloadMap>(
           .catch((err) => {
             running--;
             finished++;
+            if (onError) {
+              onError(err instanceof Error ? err : new Error(String(err)));
+            }
             next();
           });
       }
@@ -388,13 +819,13 @@ export async function processBatchWithHandlers<PayloadMap>(
 
 /**
  * Start a job processor that continuously processes jobs
- * @param pool - The database pool
+ * @param backend - The queue backend
  * @param handlers - The job handlers for this processor instance
  * @param options - The processor options. Leave pollInterval empty to run only once. Use jobType to filter jobs by type.
  * @returns {Processor} The processor instance
  */
 export const createProcessor = <PayloadMap = any>(
-  pool: Pool,
+  backend: QueueBackend,
   handlers: JobHandlers<PayloadMap>,
   options: ProcessorOptions = {},
 ): Processor => {
@@ -409,6 +840,7 @@ export const createProcessor = <PayloadMap = any>(
 
   let running = false;
   let intervalId: NodeJS.Timeout | null = null;
+  let currentBatchPromise: Promise<number> | null = null;
 
   setLogContext(options.verbose ?? false);
 
@@ -421,12 +853,13 @@ export const createProcessor = <PayloadMap = any>(
 
     try {
       const processed = await processBatchWithHandlers(
-        pool,
+        backend,
         workerId,
         batchSize,
         jobType,
         handlers,
         concurrency,
+        onError,
       );
       // Only process one batch in start; do not schedule next batch here
       return processed;
@@ -447,27 +880,62 @@ export const createProcessor = <PayloadMap = any>(
 
       log(`Starting job processor with workerId: ${workerId}`);
       running = true;
-      // Background: process batches repeatedly if needed
-      const processBatches = async () => {
+
+      // Single serialized loop: process a batch, then either immediately
+      // continue (if full batch was returned) or wait pollInterval.
+      const scheduleNext = (immediate: boolean) => {
         if (!running) return;
-        const processed = await processJobs();
-        if (processed === batchSize && running) {
-          setImmediate(processBatches);
+        if (immediate) {
+          intervalId = setTimeout(loop, 0);
+        } else {
+          intervalId = setTimeout(loop, pollInterval);
         }
       };
-      processBatches(); // Process immediately on start
-      intervalId = setInterval(processJobs, pollInterval);
+
+      const loop = async () => {
+        if (!running) return;
+        currentBatchPromise = processJobs();
+        const processed = await currentBatchPromise;
+        currentBatchPromise = null;
+        // If we got a full batch, there may be more work — process immediately
+        scheduleNext(processed === batchSize);
+      };
+
+      // Start the first iteration immediately
+      loop();
     },
     /**
-     * Stop the job processor that runs in the background
+     * Stop the job processor that runs in the background.
+     * Does not wait for in-flight jobs.
      */
     stop: () => {
       log(`Stopping job processor with workerId: ${workerId}`);
       running = false;
       if (intervalId) {
-        clearInterval(intervalId);
+        clearTimeout(intervalId);
+        intervalId = null;
+      }
+    },
+    /**
+     * Stop the job processor and wait for all in-flight jobs to complete.
+     * Useful for graceful shutdown (e.g., SIGTERM handling).
+     */
+    stopAndDrain: async (drainTimeoutMs = 30000) => {
+      log(`Stopping and draining job processor with workerId: ${workerId}`);
+      running = false;
+      if (intervalId) {
+        clearTimeout(intervalId);
         intervalId = null;
       }
+      // Wait for current batch to finish, with a timeout
+      if (currentBatchPromise) {
+        await Promise.race([
+          currentBatchPromise.catch(() => {}),
+          new Promise<void>((resolve) => setTimeout(resolve, drainTimeoutMs)),
+        ]);
+        currentBatchPromise = null;
+      }
+      log(`Job processor ${workerId} drained`);
     },
     /**
      * Start the job processor synchronously.