fix: prevent concurrent step data updates

Author: LouisHaftmann

package.json

@@ -44,6 +44,7 @@
     "exit-hook": "^5.0.1",
     "groupmq": "^1.1.0",
     "ioredis": "^5.9.2",
+    "p-mutex": "^1.0.0",
     "p-retry": "^7.1.1",
     "superjson": "^2.2.6",
     "type-fest": "^5.4.2"

pnpm-lock.yaml

@@ -2,6 +2,7 @@ import type { Span } from '@opentelemetry/api'
 import type { Options } from 'p-retry'
 import type { WorkflowQueueInternal } from './types'
 import { setTimeout } from 'node:timers/promises'
+import Mutex from 'p-mutex'
 import pRetry from 'p-retry'
 import { deserialize, serialize } from './serializer'
 import { Settings } from './settings'
@@ -24,6 +25,7 @@ export class WorkflowStep {
   private queue
   private workflowJobId
   private stepNamePrefix
+  private updateStepDataMutex = new Mutex()
 
   constructor(opts: {
     queue: WorkflowQueueInternal<unknown>
@@ -171,13 +173,17 @@ export class WorkflowStep {
     return stepData as Extract<WorkflowStepData, { type: T }> | undefined
   }
 
+  // TODO maybe update this to use lua script for atomic operation
   private async updateStepData(stepName: string, data: WorkflowStepData) {
-    const job = await this.getWorkflowJob()
+    // prevent concurrent updates to step data which could lead to race conditions and lost updates
+    await this.updateStepDataMutex.withLock(async () => {
+      const job = await this.getWorkflowJob()
 
-    const jobData = deserialize(job.data)
-    jobData.stepData[stepName] = data
+      const jobData = deserialize(job.data)
+      jobData.stepData[stepName] = data
 
-    await job.updateData(serialize(jobData))
+      await job.updateData(serialize(jobData))
+    })
   }
 
   private async getWorkflowJob() {

src/step.ts

@@ -4,6 +4,7 @@ import { beforeAll, describe, expect, test, vi } from 'vitest'
 import { createRedis, Settings, Workflow } from '../src'
 
 beforeAll(() => {
+  Settings.logger = console
   Settings.defaultConnection = async () =>
     createRedis({
       host: 'localhost',
@@ -57,12 +58,21 @@ describe('input', () => {
 
 describe('step', () => {
   test('only runs once', async () => {
-    const stepHandler = vi.fn()
+    const stepHandler1 = vi.fn()
+    const stepHandler2 = vi.fn()
+    const stepHandler3 = vi.fn()
     const handler = vi.fn()
     const workflow = new Workflow({
       id: randomUUID(),
+      queueOptions: {
+        maxAttempts: 10,
+      },
       run: async ({ step }) => {
-        await step.do('test-step', stepHandler)
+        await Promise.all([
+          step.do('test-step1', stepHandler1),
+          step.do('test-step2', stepHandler2),
+          step.do('test-step3', stepHandler3),
+        ])
         await handler()
         throw new Error('error')
       },
@@ -74,8 +84,10 @@ describe('step', () => {
     await workflow.run(undefined)
 
     await vi.waitFor(() => {
-      expect(stepHandler).toHaveBeenCalledOnce()
-      expect(handler).toHaveBeenCalledTimes(3)
+      expect(stepHandler1).toHaveBeenCalledOnce()
+      expect(stepHandler2).toHaveBeenCalledOnce()
+      expect(stepHandler3).toHaveBeenCalledOnce()
+      expect(handler).toHaveBeenCalledTimes(10)
     })
   })
   test('caches output', async () => {