@pikku/kysely 0.13.18 → 0.13.19

package/CHANGELOG.md

@@ -1,3 +1,26 @@
+## 0.13.19
+
+### Patch Changes
+
+- 2d21628: fix(kysely): claim a workflow step atomically in every SQL dialect
+
+  The workflow engine's "atomic claim" was a read-then-write guarded by
+  `withStepLock`, and `@pikku/kysely` inherited a silent pass-through for that
+  lock — so on every dialect but Postgres and MySQL a redelivered queue job could
+  claim a step another dispatch was already running, executing a side-effecting
+  step twice.
+
+  `@pikku/kysely` now claims the step with a status-guarded `UPDATE` and reads the
+  affected-row count, which is atomic in every SQL dialect without an
+  advisory-lock primitive. Relay redispatch is enabled for all Kysely dialects as
+  a result, not just Postgres and MySQL.
+
+- Updated dependencies [9687ad1]
+- Updated dependencies [2d21628]
+- Updated dependencies [985b87b]
+- Updated dependencies [3a83f85]
+  - @pikku/core@0.12.87
+
 ## 0.13.18
 
 ### Patch Changes

package/dist/src/kysely-workflow-service.d.ts

@@ -47,7 +47,36 @@ export declare class KyselyWorkflowService extends PikkuWorkflowService {
     protected setStepErrorImpl(stepId: string, error: Error): Promise<void>;
     protected createRetryAttemptImpl(stepId: string, status: 'pending' | 'running'): Promise<StepState>;
     withRunLock<T>(_id: string, fn: () => Promise<T>): Promise<T>;
+    /**
+     * A pass-through, and deliberately so: the one decision that must exclude —
+     * claiming a step to execute it — is made by `claimStepForExecution` below as
+     * a single conditional write, which needs no lock to be exclusive.
+     *
+     * What is left are the suspend and approval sections in the engine, where the
+     * lock is genuine mutual exclusion rather than a claim. A dialect with a real
+     * primitive overrides this to cover them — `kysely-postgres` with
+     * `pg_advisory_xact_lock`, `kysely-mysql` with `GET_LOCK`.
+     */
     withStepLock<T>(_runId: string, _stepName: string, fn: () => Promise<T>): Promise<T>;
+    /**
+     * Claim the step with a status-guarded `UPDATE` and read the affected-row
+     * count: the database decides the winner in one statement, so two dispatches
+     * racing for the same step cannot both proceed.
+     *
+     * This replaces the read-then-write the base engine does under `withStepLock`,
+     * which is only as exclusive as that lock — and every dialect but Postgres and
+     * MySQL inherits a pass-through. A conditional update needs no advisory-lock
+     * primitive, so every SQL dialect gets the same guarantee.
+     *
+     * The winner then goes through the ordinary transition methods, so history
+     * rows and the mirror see exactly what they saw before.
+     */
+    protected claimStepForExecution(runId: string, stepName: string, rpcName: string): Promise<StepState | null>;
+    /**
+     * Move a step to `running` only if it is still in one of `from`, reporting
+     * whether this caller is the one that moved it.
+     */
+    private claimStepStatus;
     getCompletedGraphState(runId: string): Promise<{
         completedNodeIds: string[];
         failedNodeIds: string[];
@@ -81,16 +110,18 @@ export declare class KyselyWorkflowService extends PikkuWorkflowService {
     private jsonPathFor;
     protected updateRunStateImpl(runId: string, name: string, value: unknown): Promise<void>;
     /**
-     * The undispatched-step query, shared by the dialects that can safely relay.
+     * The undispatched-step query, shared by every dialect.
      *
-     * Deliberately NOT an override of `findUndispatchedSteps`: this class's
-     * `withStepLock` is a pass-through, so a subclass inheriting it (kysely-sqlite)
-     * has no atomic claim, and the relay's redundant dispatches would become
-     * double executions. A dialect opts in by overriding `findUndispatchedSteps`
-     * to call this — which `kysely-postgres` and `kysely-mysql` do, having real
-     * locks.
+     * It used to be deliberately NOT an override, because this class's
+     * `withStepLock` is a pass-through: a subclass inheriting it (kysely-sqlite)
+     * had no atomic claim, so the relay's redundant dispatches would have become
+     * double executions, and only `kysely-postgres` and `kysely-mysql` opted in
+     * on the strength of their real locks. `claimStepForExecution` no longer
+     * needs that lock — the claim is a status-guarded `UPDATE`, atomic in every
+     * dialect — so the relay is safe here for all of them and the override lives
+     * in the base.
      */
-    protected queryUndispatchedSteps(before: Date, limit: number): Promise<Array<{
+    protected findUndispatchedSteps(before: Date, limit: number): Promise<Array<{
         runId: string;
         stepId: string;
     }>>;

package/dist/src/kysely-workflow-service.js

@@ -1,4 +1,4 @@
-import { PikkuWorkflowService } from '@pikku/core/workflow';
+import { PikkuWorkflowService, WorkflowStepFunctionMismatchError, } from '@pikku/core/workflow';
 import { sql } from 'kysely';
 import { KyselyWorkflowRunService } from './kysely-workflow-run-service.js';
 import { parseJson } from './kysely-json.js';
@@ -363,9 +363,73 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
     async withRunLock(_id, fn) {
         return fn();
     }
+    /**
+     * A pass-through, and deliberately so: the one decision that must exclude —
+     * claiming a step to execute it — is made by `claimStepForExecution` below as
+     * a single conditional write, which needs no lock to be exclusive.
+     *
+     * What is left are the suspend and approval sections in the engine, where the
+     * lock is genuine mutual exclusion rather than a claim. A dialect with a real
+     * primitive overrides this to cover them — `kysely-postgres` with
+     * `pg_advisory_xact_lock`, `kysely-mysql` with `GET_LOCK`.
+     */
     async withStepLock(_runId, _stepName, fn) {
         return fn();
     }
+    /**
+     * Claim the step with a status-guarded `UPDATE` and read the affected-row
+     * count: the database decides the winner in one statement, so two dispatches
+     * racing for the same step cannot both proceed.
+     *
+     * This replaces the read-then-write the base engine does under `withStepLock`,
+     * which is only as exclusive as that lock — and every dialect but Postgres and
+     * MySQL inherits a pass-through. A conditional update needs no advisory-lock
+     * primitive, so every SQL dialect gets the same guarantee.
+     *
+     * The winner then goes through the ordinary transition methods, so history
+     * rows and the mirror see exactly what they saw before.
+     */
+    async claimStepForExecution(runId, stepName, rpcName) {
+        const stepState = await this.getStepState(runId, stepName);
+        // knowledge: decisions/security/a-step-runs-the-function-the-workflow-dispatched-it-with.md
+        if (stepState.rpcName !== undefined &&
+            stepState.rpcName !== (rpcName ?? null)) {
+            throw new WorkflowStepFunctionMismatchError(runId, stepName);
+        }
+        if (stepState.status === 'succeeded' || stepState.status === 'running') {
+            return null;
+        }
+        if (stepState.status === 'failed') {
+            if (!(await this.claimStepStatus(stepState.stepId, ['failed']))) {
+                return null;
+            }
+            return this.createRetryAttempt(stepState.stepId, 'running');
+        }
+        if (stepState.status === 'pending' || stepState.status === 'scheduled') {
+            if (!(await this.claimStepStatus(stepState.stepId, [
+                'pending',
+                'scheduled',
+            ]))) {
+                return null;
+            }
+            await this.setStepRunning(stepState.stepId);
+            return { ...stepState, status: 'running' };
+        }
+        return stepState;
+    }
+    /**
+     * Move a step to `running` only if it is still in one of `from`, reporting
+     * whether this caller is the one that moved it.
+     */
+    async claimStepStatus(stepId, from) {
+        const claimed = await this.db
+            .updateTable('workflowStep')
+            .set({ status: 'running', updatedAt: new Date() })
+            .where('workflowStepId', '=', stepId)
+            .where('status', 'in', from)
+            .executeTakeFirst();
+        return Number(claimed?.numUpdatedRows ?? 0n) > 0;
+    }
     async getCompletedGraphState(runId) {
         const results = await this.db
             .selectFrom('workflowStep as ws')
@@ -470,16 +534,18 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
             .execute();
     }
     /**
-     * The undispatched-step query, shared by the dialects that can safely relay.
+     * The undispatched-step query, shared by every dialect.
      *
-     * Deliberately NOT an override of `findUndispatchedSteps`: this class's
-     * `withStepLock` is a pass-through, so a subclass inheriting it (kysely-sqlite)
-     * has no atomic claim, and the relay's redundant dispatches would become
-     * double executions. A dialect opts in by overriding `findUndispatchedSteps`
-     * to call this — which `kysely-postgres` and `kysely-mysql` do, having real
-     * locks.
+     * It used to be deliberately NOT an override, because this class's
+     * `withStepLock` is a pass-through: a subclass inheriting it (kysely-sqlite)
+     * had no atomic claim, so the relay's redundant dispatches would have become
+     * double executions, and only `kysely-postgres` and `kysely-mysql` opted in
+     * on the strength of their real locks. `claimStepForExecution` no longer
+     * needs that lock — the claim is a status-guarded `UPDATE`, atomic in every
+     * dialect — so the relay is safe here for all of them and the override lives
+     * in the base.
      */
-    async queryUndispatchedSteps(before, limit) {
+    async findUndispatchedSteps(before, limit) {
         const rows = await this.db
             .selectFrom('workflowStep as s')
             .innerJoin('workflowRuns as r', 'r.workflowRunId', 's.workflowRunId')

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pikku/kysely",
-  "version": "0.13.18",
+  "version": "0.13.19",
   "description": "Kysely-backed Pikku services — workflow state, secrets, sessions, scopes and audit",
   "author": "yasser.fadl@gmail.com",
   "license": "MIT",
@@ -21,7 +21,7 @@
     "prepublishOnly": "yarn build"
   },
   "peerDependencies": {
-    "@pikku/core": "^0.12.86"
+    "@pikku/core": "^0.12.87"
   },
   "dependencies": {
     "kysely": "^0.29.5"

package/src/kysely-workflow-service.test.ts

@@ -3,6 +3,7 @@ import assert from 'node:assert/strict'
 import { CamelCasePlugin, Kysely, SqliteDialect, sql } from 'kysely'
 import Database from 'better-sqlite3'
 
+import type { StepState } from '@pikku/core/ecosystem/workflow'
 import type { KyselyPikkuDB } from './kysely-tables.js'
 import { SerializePlugin } from './serialize-plugin.js'
 import { KyselyWorkflowService } from './kysely-workflow-service.js'
@@ -428,3 +429,114 @@ describe('KyselyWorkflowService — a transition that addresses no history row',
     assert.equal(history[0]!.status, 'succeeded')
   })
 })
+
+/**
+ * Dispatch is at-least-once — the relay re-dispatches a step it believes was
+ * dropped, and a queue can redeliver a job it already handed out — so the claim
+ * is the only thing standing between a duplicate dispatch and a second
+ * execution of a side-effecting step.
+ */
+describe('KyselyWorkflowService — claiming a step for execution', () => {
+  const claim = (runId: string, stepName: string, rpcName = 'rpc.fn') =>
+    (service as any).claimStepForExecution(
+      runId,
+      stepName,
+      rpcName
+    ) as Promise<StepState | null>
+
+  test('two dispatches racing for the same pending step: exactly one wins', async () => {
+    const runId = await seedRun()
+    await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+
+    const claims = await Promise.all([claim(runId, 's1'), claim(runId, 's1')])
+    const winners = claims.filter((c) => c !== null)
+
+    assert.equal(
+      winners.length,
+      1,
+      `both dispatches claimed the step, so a side-effecting step would run twice: ${JSON.stringify(claims)}`
+    )
+    assert.equal((await service.getStepState(runId, 's1')).status, 'running')
+  })
+
+  test('two dispatches racing to retry the same failed step: exactly one wins', async () => {
+    const runId = await seedRun()
+    const step = await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+    await service.setStepRunning(step.stepId)
+    await service.setStepError(step.stepId, new Error('boom'))
+
+    const claims = await Promise.all([claim(runId, 's1'), claim(runId, 's1')])
+    const winners = claims.filter((c) => c !== null)
+
+    assert.equal(
+      winners.length,
+      1,
+      `both dispatches started a retry, so the failed step would be retried twice concurrently: ${JSON.stringify(claims)}`
+    )
+    assert.equal(
+      (await service.getStepState(runId, 's1')).attemptCount,
+      2,
+      'the retry raced itself and burned more than one attempt'
+    )
+  })
+
+  test('a step already claimed is not claimable again', async () => {
+    const runId = await seedRun()
+    await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+
+    assert.notEqual(await claim(runId, 's1'), null)
+    assert.equal(
+      await claim(runId, 's1'),
+      null,
+      'a redelivered job re-claimed a step that is already running'
+    )
+  })
+
+  test('a succeeded step is not claimable', async () => {
+    const runId = await seedRun()
+    const step = await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+    await service.setStepRunning(step.stepId)
+    await service.setStepResult(step.stepId, { ok: true })
+
+    assert.equal(await claim(runId, 's1'), null)
+  })
+
+  test('a claim that fails releases the step, so it is retryable', async () => {
+    const runId = await seedRun()
+    const step = await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+
+    const first = await claim(runId, 's1')
+    assert.notEqual(first, null)
+    await service.setStepError(step.stepId, new Error('boom'))
+
+    const retry = await claim(runId, 's1')
+    assert.notEqual(
+      retry,
+      null,
+      'the failed step stayed claimed, so it can never be retried'
+    )
+    assert.equal(retry?.status, 'running')
+    assert.equal(retry?.attemptCount, 2)
+  })
+
+  test('a claim for a different function than the step was dispatched with is refused', async () => {
+    const runId = await seedRun()
+    await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+
+    await assert.rejects(() => claim(runId, 's1', 'rpc.other'))
+    assert.equal(
+      (await service.getStepState(runId, 's1')).status,
+      'pending',
+      'the mismatched claim moved the step anyway'
+    )
+  })
+
+  test('a scheduled step is claimable, and only once', async () => {
+    const runId = await seedRun()
+    const step = await service.insertStepState(runId, 's1', 'rpc.fn', { x: 1 })
+    await service.setStepScheduled(step.stepId)
+
+    const claims = await Promise.all([claim(runId, 's1'), claim(runId, 's1')])
+    assert.equal(claims.filter((c) => c !== null).length, 1)
+  })
+})

package/src/kysely-workflow-service.ts

@@ -1,5 +1,8 @@
 import type { SerializedError } from '@pikku/core/errors'
-import { PikkuWorkflowService } from '@pikku/core/workflow'
+import {
+  PikkuWorkflowService,
+  WorkflowStepFunctionMismatchError,
+} from '@pikku/core/workflow'
 import type {
   WorkflowPlannedStep,
   WorkflowQueueOptions,
@@ -460,6 +463,16 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
     return fn()
   }
 
+  /**
+   * A pass-through, and deliberately so: the one decision that must exclude —
+   * claiming a step to execute it — is made by `claimStepForExecution` below as
+   * a single conditional write, which needs no lock to be exclusive.
+   *
+   * What is left are the suspend and approval sections in the engine, where the
+   * lock is genuine mutual exclusion rather than a claim. A dialect with a real
+   * primitive overrides this to cover them — `kysely-postgres` with
+   * `pg_advisory_xact_lock`, `kysely-mysql` with `GET_LOCK`.
+   */
   async withStepLock<T>(
     _runId: string,
     _stepName: string,
@@ -468,6 +481,79 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
     return fn()
   }
 
+  /**
+   * Claim the step with a status-guarded `UPDATE` and read the affected-row
+   * count: the database decides the winner in one statement, so two dispatches
+   * racing for the same step cannot both proceed.
+   *
+   * This replaces the read-then-write the base engine does under `withStepLock`,
+   * which is only as exclusive as that lock — and every dialect but Postgres and
+   * MySQL inherits a pass-through. A conditional update needs no advisory-lock
+   * primitive, so every SQL dialect gets the same guarantee.
+   *
+   * The winner then goes through the ordinary transition methods, so history
+   * rows and the mirror see exactly what they saw before.
+   */
+  protected override async claimStepForExecution(
+    runId: string,
+    stepName: string,
+    rpcName: string
+  ): Promise<StepState | null> {
+    const stepState = await this.getStepState(runId, stepName)
+
+    // knowledge: decisions/security/a-step-runs-the-function-the-workflow-dispatched-it-with.md
+    if (
+      stepState.rpcName !== undefined &&
+      stepState.rpcName !== (rpcName ?? null)
+    ) {
+      throw new WorkflowStepFunctionMismatchError(runId, stepName)
+    }
+
+    if (stepState.status === 'succeeded' || stepState.status === 'running') {
+      return null
+    }
+
+    if (stepState.status === 'failed') {
+      if (!(await this.claimStepStatus(stepState.stepId, ['failed']))) {
+        return null
+      }
+      return this.createRetryAttempt(stepState.stepId, 'running')
+    }
+
+    if (stepState.status === 'pending' || stepState.status === 'scheduled') {
+      if (
+        !(await this.claimStepStatus(stepState.stepId, [
+          'pending',
+          'scheduled',
+        ]))
+      ) {
+        return null
+      }
+      await this.setStepRunning(stepState.stepId)
+      return { ...stepState, status: 'running' }
+    }
+
+    return stepState
+  }
+
+  /**
+   * Move a step to `running` only if it is still in one of `from`, reporting
+   * whether this caller is the one that moved it.
+   */
+  private async claimStepStatus(
+    stepId: string,
+    from: StepStatus[]
+  ): Promise<boolean> {
+    const claimed = await this.db
+      .updateTable('workflowStep')
+      .set({ status: 'running', updatedAt: new Date() })
+      .where('workflowStepId', '=', stepId)
+      .where('status', 'in', from)
+      .executeTakeFirst()
+
+    return Number(claimed?.numUpdatedRows ?? 0n) > 0
+  }
+
   async getCompletedGraphState(runId: string): Promise<{
     completedNodeIds: string[]
     failedNodeIds: string[]
@@ -606,16 +692,18 @@ export class KyselyWorkflowService extends PikkuWorkflowService {
   }
 
   /**
-   * The undispatched-step query, shared by the dialects that can safely relay.
+   * The undispatched-step query, shared by every dialect.
    *
-   * Deliberately NOT an override of `findUndispatchedSteps`: this class's
-   * `withStepLock` is a pass-through, so a subclass inheriting it (kysely-sqlite)
-   * has no atomic claim, and the relay's redundant dispatches would become
-   * double executions. A dialect opts in by overriding `findUndispatchedSteps`
-   * to call this — which `kysely-postgres` and `kysely-mysql` do, having real
-   * locks.
+   * It used to be deliberately NOT an override, because this class's
+   * `withStepLock` is a pass-through: a subclass inheriting it (kysely-sqlite)
+   * had no atomic claim, so the relay's redundant dispatches would have become
+   * double executions, and only `kysely-postgres` and `kysely-mysql` opted in
+   * on the strength of their real locks. `claimStepForExecution` no longer
+   * needs that lock — the claim is a status-guarded `UPDATE`, atomic in every
+   * dialect — so the relay is safe here for all of them and the override lives
+   * in the base.
    */
-  protected async queryUndispatchedSteps(
+  protected override async findUndispatchedSteps(
     before: Date,
     limit: number
   ): Promise<Array<{ runId: string; stepId: string }>> {