fix(workflow-engines): race condition when retry interval is used (#11771)

packages/modules/workflow-engine-inmemory/integration-tests/__tests__/index.spec.ts

@@ -30,7 +30,7 @@ import {
 } from "../__fixtures__/workflow_event_group_id"
 import { createScheduled } from "../__fixtures__/workflow_scheduled"
 
-jest.setTimeout(100000)
+jest.setTimeout(3000000)
 
 moduleIntegrationTestRunner<IWorkflowEngineService>({
   moduleName: Modules.WORKFLOW_ENGINE,

packages/modules/workflow-engine-inmemory/integration-tests/__tests__/race.spec.ts

@@ -0,0 +1,182 @@
+import { IWorkflowEngineService } from "@medusajs/framework/types"
+import { Modules } from "@medusajs/framework/utils"
+import {
+  createStep,
+  createWorkflow,
+  StepResponse,
+  transform,
+  WorkflowResponse,
+} from "@medusajs/framework/workflows-sdk"
+import { moduleIntegrationTestRunner } from "@medusajs/test-utils"
+import { setTimeout as setTimeoutSync } from "timers"
+import { setTimeout } from "timers/promises"
+import "../__fixtures__"
+
+jest.setTimeout(3000000)
+
+moduleIntegrationTestRunner<IWorkflowEngineService>({
+  moduleName: Modules.WORKFLOW_ENGINE,
+  resolve: __dirname + "/../..",
+  testSuite: ({ service: workflowOrcModule, medusaApp }) => {
+    describe("Testing race condition of the workflow during retry", () => {
+      it("should prevent race continuation of the workflow during retryIntervalAwaiting in background execution", (done) => {
+        const step0InvokeMock = jest.fn()
+        const step1InvokeMock = jest.fn()
+        const step2InvokeMock = jest.fn()
+        const transformMock = jest.fn()
+
+        const step0 = createStep("step0", async (_) => {
+          step0InvokeMock()
+          return new StepResponse("result from step 0")
+        })
+
+        const step1 = createStep("step1", async (_) => {
+          step1InvokeMock()
+          await setTimeout(2000)
+          return new StepResponse({ isSuccess: true })
+        })
+
+        const step2 = createStep("step2", async (input: any) => {
+          step2InvokeMock()
+          return new StepResponse({ result: input })
+        })
+
+        const subWorkflow = createWorkflow("sub-workflow-1", function () {
+          const status = step1()
+          return new WorkflowResponse(status)
+        })
+
+        createWorkflow("workflow-1", function () {
+          const build = step0()
+
+          const status = subWorkflow.runAsStep({} as any).config({
+            async: true,
+            compensateAsync: true,
+            backgroundExecution: true,
+            retryIntervalAwaiting: 1,
+          })
+
+          const transformedResult = transform({ status }, (data) => {
+            transformMock()
+            return {
+              status: data.status,
+            }
+          })
+
+          step2(transformedResult)
+          return new WorkflowResponse(build)
+        })
+
+        void workflowOrcModule.subscribe({
+          workflowId: "workflow-1",
+
+          subscriber: (event) => {
+            if (event.eventType === "onFinish") {
+              expect(step0InvokeMock).toHaveBeenCalledTimes(1)
+              expect(step1InvokeMock.mock.calls.length).toBeGreaterThan(1)
+              expect(step2InvokeMock).toHaveBeenCalledTimes(1)
+              expect(transformMock).toHaveBeenCalledTimes(1)
+              setTimeoutSync(done, 500)
+            }
+          },
+        })
+
+        workflowOrcModule
+          .run("workflow-1", { throwOnError: false })
+          .then(({ result }) => {
+            expect(result).toBe("result from step 0")
+          })
+          .catch((e) => e)
+      })
+
+      it("should prevent race continuation of the workflow compensation during retryIntervalAwaiting in background execution", (done) => {
+        const workflowId = "RACE_workflow-1"
+
+        const step0InvokeMock = jest.fn()
+        const step0CompensateMock = jest.fn()
+        const step1InvokeMock = jest.fn()
+        const step1CompensateMock = jest.fn()
+        const step2InvokeMock = jest.fn()
+        const transformMock = jest.fn()
+
+        const step0 = createStep(
+          "RACE_step0",
+          async (_) => {
+            step0InvokeMock()
+            return new StepResponse("result from step 0")
+          },
+          () => {
+            step0CompensateMock()
+          }
+        )
+
+        const step1 = createStep(
+          "RACE_step1",
+          async (_) => {
+            step1InvokeMock()
+            await setTimeout(300)
+            throw new Error("error from step 1")
+          },
+          () => {
+            step1CompensateMock()
+          }
+        )
+
+        const step2 = createStep("RACE_step2", async (input: any) => {
+          step2InvokeMock()
+          return new StepResponse({ result: input })
+        })
+
+        const subWorkflow = createWorkflow("RACE_sub-workflow-1", function () {
+          const status = step1()
+          return new WorkflowResponse(status)
+        })
+
+        createWorkflow(workflowId, function () {
+          const build = step0()
+
+          const status = subWorkflow.runAsStep({} as any).config({
+            async: true,
+            compensateAsync: true,
+            backgroundExecution: true,
+            retryIntervalAwaiting: 0.1,
+          })
+
+          const transformedResult = transform({ status }, (data) => {
+            transformMock()
+            return {
+              status: data.status,
+            }
+          })
+
+          step2(transformedResult)
+          return new WorkflowResponse(build)
+        })
+
+        void workflowOrcModule.subscribe({
+          workflowId: workflowId,
+          subscriber: async (event) => {
+            if (event.eventType === "onFinish") {
+              expect(step0InvokeMock).toHaveBeenCalledTimes(1)
+              expect(step0CompensateMock).toHaveBeenCalledTimes(1)
+              expect(step1InvokeMock.mock.calls.length).toBeGreaterThan(2)
+              expect(step1CompensateMock).toHaveBeenCalledTimes(1)
+              expect(step2InvokeMock).toHaveBeenCalledTimes(0)
+              expect(transformMock).toHaveBeenCalledTimes(0)
+              setTimeoutSync(done, 500)
+            }
+          },
+        })
+
+        workflowOrcModule
+          .run(workflowId, {
+            throwOnError: false,
+          })
+          .then(({ result }) => {
+            expect(result).toBe("result from step 0")
+          })
+          .catch((e) => e)
+      })
+    })
+  },
+})

packages/modules/workflow-engine-inmemory/package.json

@@ -29,7 +29,7 @@
     "resolve:aliases": "tsc --showConfig -p tsconfig.json > tsconfig.resolved.json && tsc-alias -p tsconfig.resolved.json && rimraf tsconfig.resolved.json",
     "build": "rimraf dist && tsc --build && npm run resolve:aliases",
     "test": "jest --passWithNoTests --runInBand --bail --forceExit -- src",
-    "test:integration": "jest --silent --forceExit -- integration-tests/**/__tests__/**/*.ts",
+    "test:integration": "jest --forceExit -- integration-tests/**/__tests__/**/*.ts",
     "migration:initial": " MIKRO_ORM_CLI_CONFIG=./mikro-orm.config.dev.ts medusa-mikro-orm migration:create --initial",
     "migration:create": " MIKRO_ORM_CLI_CONFIG=./mikro-orm.config.dev.ts medusa-mikro-orm migration:create",
     "migration:up": " MIKRO_ORM_CLI_CONFIG=./mikro-orm.config.dev.ts medusa-mikro-orm migration:up",

packages/modules/workflow-engine-inmemory/src/services/workflow-orchestrator.ts

@@ -6,15 +6,9 @@ import {
   TransactionStep,
   WorkflowScheduler,
 } from "@medusajs/framework/orchestration"
+import { ContainerLike, MedusaContainer } from "@medusajs/framework/types"
 import {
-  ContainerLike,
-  Context,
-  MedusaContainer,
-} from "@medusajs/framework/types"
-import {
-  InjectSharedContext,
   isString,
-  MedusaContext,
   MedusaError,
   TransactionState,
 } from "@medusajs/framework/utils"
@@ -26,6 +20,7 @@ import {
 } from "@medusajs/framework/workflows-sdk"
 import { ulid } from "ulid"
 import { InMemoryDistributedTransactionStorage } from "../utils"
+import { WorkflowOrchestratorCancelOptions } from "@types"
 
 export type WorkflowOrchestratorRunOptions<T> = Omit<
   FlowRunOptions<T>,
@@ -125,11 +120,9 @@ export class WorkflowOrchestratorService {
     }
   }
 
-  @InjectSharedContext()
   async run<T = unknown>(
     workflowIdOrWorkflow: string | ReturnWorkflow<any, any, any>,
-    options?: WorkflowOrchestratorRunOptions<T>,
-    @MedusaContext() sharedContext: Context = {}
+    options?: WorkflowOrchestratorRunOptions<T>
   ) {
     const {
       input,
@@ -217,12 +210,116 @@ export class WorkflowOrchestratorService {
     return { acknowledgement, ...ret }
   }
 
-  @InjectSharedContext()
+  async cancel(
+    workflowIdOrWorkflow: string | ReturnWorkflow<any, any, any>,
+    options?: WorkflowOrchestratorCancelOptions
+  ) {
+    const {
+      transactionId,
+      logOnError,
+      events: eventHandlers,
+      container,
+    } = options ?? {}
+
+    let { throwOnError, context } = options ?? {}
+
+    throwOnError ??= true
+    context ??= {}
+
+    const workflowId = isString(workflowIdOrWorkflow)
+      ? workflowIdOrWorkflow
+      : workflowIdOrWorkflow.getName()
+
+    if (!workflowId) {
+      throw new Error("Workflow ID is required")
+    }
+
+    if (!transactionId) {
+      throw new Error("Transaction ID is required")
+    }
+
+    const events: FlowRunOptions["events"] = this.buildWorkflowEvents({
+      customEventHandlers: eventHandlers,
+      workflowId,
+      transactionId: transactionId,
+    })
+
+    const exportedWorkflow = MedusaWorkflow.getWorkflow(workflowId)
+    if (!exportedWorkflow) {
+      throw new Error(`Workflow with id "${workflowId}" not found.`)
+    }
+
+    const originalOnFinishHandler = events.onFinish!
+    delete events.onFinish
+
+    const transaction = await this.getRunningTransaction(
+      workflowId,
+      transactionId,
+      options
+    )
+
+    if (!transaction) {
+      if (!throwOnError) {
+        return {
+          acknowledgement: {
+            transactionId,
+            workflowId,
+            exists: false,
+          },
+        }
+      }
+      throw new Error("Transaction not found")
+    }
+
+    const ret = await exportedWorkflow.cancel({
+      transaction,
+      throwOnError: false,
+      logOnError,
+      context,
+      events,
+      container: container ?? this.container_,
+    })
+
+    const hasFinished = ret.transaction.hasFinished()
+    const metadata = ret.transaction.getFlow().metadata
+    const { parentStepIdempotencyKey } = metadata ?? {}
+
+    const hasFailed = [TransactionState.FAILED].includes(
+      ret.transaction.getFlow().state
+    )
+
+    const acknowledgement = {
+      transactionId: context.transactionId,
+      workflowId: workflowId,
+      parentStepIdempotencyKey,
+      hasFinished,
+      hasFailed,
+      exists: true,
+    }
+
+    if (hasFinished) {
+      const { result, errors } = ret
+
+      await originalOnFinishHandler({
+        transaction: ret.transaction,
+        result,
+        errors,
+      })
+
+      await this.triggerParentStep(ret.transaction, result)
+    }
+
+    if (throwOnError && ret.thrownError) {
+      throw ret.thrownError
+    }
+
+    return { acknowledgement, ...ret }
+  }
+
   async getRunningTransaction(
     workflowId: string,
     transactionId: string,
-    options?: WorkflowOrchestratorRunOptions<undefined>,
-    @MedusaContext() sharedContext: Context = {}
+    options?: WorkflowOrchestratorRunOptions<unknown>
   ): Promise<DistributedTransactionType> {
     let { context, container } = options ?? {}
 
@@ -251,19 +348,15 @@ export class WorkflowOrchestratorService {
     return transaction
   }
 
-  @InjectSharedContext()
-  async setStepSuccess<T = unknown>(
-    {
-      idempotencyKey,
-      stepResponse,
-      options,
-    }: {
-      idempotencyKey: string | IdempotencyKeyParts
-      stepResponse: unknown
-      options?: RegisterStepSuccessOptions<T>
-    },
-    @MedusaContext() sharedContext: Context = {}
-  ) {
+  async setStepSuccess<T = unknown>({
+    idempotencyKey,
+    stepResponse,
+    options,
+  }: {
+    idempotencyKey: string | IdempotencyKeyParts
+    stepResponse: unknown
+    options?: RegisterStepSuccessOptions<T>
+  }) {
     const {
       context,
       logOnError,
@@ -321,19 +414,15 @@ export class WorkflowOrchestratorService {
     return ret
   }
 
-  @InjectSharedContext()
-  async setStepFailure<T = unknown>(
-    {
-      idempotencyKey,
-      stepResponse,
-      options,
-    }: {
-      idempotencyKey: string | IdempotencyKeyParts
-      stepResponse: unknown
-      options?: RegisterStepSuccessOptions<T>
-    },
-    @MedusaContext() sharedContext: Context = {}
-  ) {
+  async setStepFailure<T = unknown>({
+    idempotencyKey,
+    stepResponse,
+    options,
+  }: {
+    idempotencyKey: string | IdempotencyKeyParts
+    stepResponse: unknown
+    options?: RegisterStepSuccessOptions<T>
+  }) {
     const {
       context,
       logOnError,
@@ -391,11 +480,12 @@ export class WorkflowOrchestratorService {
     return ret
   }
 
-  @InjectSharedContext()
-  subscribe(
-    { workflowId, transactionId, subscriber, subscriberId }: SubscribeOptions,
-    @MedusaContext() sharedContext: Context = {}
-  ) {
+  subscribe({
+    workflowId,
+    transactionId,
+    subscriber,
+    subscriberId,
+  }: SubscribeOptions) {
     subscriber._id = subscriberId
     const subscribers = this.subscribers.get(workflowId) ?? new Map()
 
@@ -427,11 +517,11 @@ export class WorkflowOrchestratorService {
     this.subscribers.set(workflowId, subscribers)
   }
 
-  @InjectSharedContext()
-  unsubscribe(
-    { workflowId, transactionId, subscriberOrId }: UnsubscribeOptions,
-    @MedusaContext() sharedContext: Context = {}
-  ) {
+  unsubscribe({
+    workflowId,
+    transactionId,
+    subscriberOrId,
+  }: UnsubscribeOptions) {
     const subscribers = this.subscribers.get(workflowId) ?? new Map()
 
     const filterSubscribers = (handlers: SubscriberHandler[]) => {

packages/modules/workflow-engine-inmemory/src/services/workflows-module.ts

@@ -62,7 +62,9 @@ export class WorkflowsModuleService<
       await this.clearExpiredExecutions()
 
       this.clearTimeout_ = setInterval(async () => {
-        await this.clearExpiredExecutions()
+        try {
+          await this.clearExpiredExecutions()
+        } catch {}
       }, 1000 * 60 * 60)
     },
     onApplicationShutdown: async () => {
@@ -80,11 +82,14 @@ export class WorkflowsModuleService<
     > = {},
     @MedusaContext() context: Context = {}
   ) {
+    options ??= {}
+    options.context ??= context
+
     const ret = await this.workflowOrchestratorService_.run<
       TWorkflow extends ReturnWorkflow<any, any, any>
         ? UnwrapWorkflowInputDataType<TWorkflow>
         : unknown
-    >(workflowIdOrWorkflow, options, context)
+    >(workflowIdOrWorkflow, options)
 
     return ret as any
   }
@@ -115,14 +120,14 @@ export class WorkflowsModuleService<
     },
     @MedusaContext() context: Context = {}
   ) {
-    return await this.workflowOrchestratorService_.setStepSuccess(
-      {
-        idempotencyKey,
-        stepResponse,
-        options,
-      } as any,
-      context
-    )
+    options ??= {}
+    options.context ??= context
+
+    return await this.workflowOrchestratorService_.setStepSuccess({
+      idempotencyKey,
+      stepResponse,
+      options,
+    } as any)
   }
 
   @InjectSharedContext()
@@ -138,14 +143,14 @@ export class WorkflowsModuleService<
     },
     @MedusaContext() context: Context = {}
   ) {
-    return await this.workflowOrchestratorService_.setStepFailure(
-      {
-        idempotencyKey,
-        stepResponse,
-        options,
-      } as any,
-      context
-    )
+    options ??= {}
+    options.context ??= context
+
+    return await this.workflowOrchestratorService_.setStepFailure({
+      idempotencyKey,
+      stepResponse,
+      options,
+    } as any)
   }
 
   @InjectSharedContext()
@@ -158,7 +163,7 @@ export class WorkflowsModuleService<
     },
     @MedusaContext() context: Context = {}
   ) {
-    return this.workflowOrchestratorService_.subscribe(args as any, context)
+    return this.workflowOrchestratorService_.subscribe(args as any)
   }
 
   @InjectSharedContext()
@@ -170,7 +175,7 @@ export class WorkflowsModuleService<
     },
     @MedusaContext() context: Context = {}
   ) {
-    return this.workflowOrchestratorService_.unsubscribe(args as any, context)
+    return this.workflowOrchestratorService_.unsubscribe(args as any)
   }
 
   private async clearExpiredExecutions() {

packages/modules/workflow-engine-inmemory/src/types/index.ts

@@ -1,5 +1,14 @@
+import { ContainerLike } from "@medusajs/framework"
 import { Logger } from "@medusajs/framework/types"
+import { FlowCancelOptions } from "@medusajs/framework/workflows-sdk"
 
 export type InitializeModuleInjectableDependencies = {
   logger?: Logger
 }
+
+export type WorkflowOrchestratorCancelOptions = Omit<
+  FlowCancelOptions,
+  "transaction" | "container"
+> & {
+  container?: ContainerLike
+}

packages/modules/workflow-engine-inmemory/src/utils/workflow-orchestrator-storage.ts

@@ -3,12 +3,19 @@ import {
   IDistributedSchedulerStorage,
   IDistributedTransactionStorage,
   SchedulerOptions,
+  SkipExecutionError,
   TransactionCheckpoint,
+  TransactionFlow,
   TransactionOptions,
   TransactionStep,
 } from "@medusajs/framework/orchestration"
 import { Logger, ModulesSdkTypes } from "@medusajs/framework/types"
-import { MedusaError, TransactionState } from "@medusajs/framework/utils"
+import {
+  MedusaError,
+  TransactionState,
+  TransactionStepState,
+  isPresent,
+} from "@medusajs/framework/utils"
 import { WorkflowOrchestratorService } from "@services"
 import { CronExpression, parseExpression } from "cron-parser"
 
@@ -121,8 +128,6 @@ export class InMemoryDistributedTransactionStorage
     ttl?: number,
     options?: TransactionOptions
   ): Promise<void> {
-    this.storage.set(key, data)
-
     /**
      * Store the retention time only if the transaction is done, failed or reverted.
      * From that moment, this tuple can be later on archived or deleted after the retention time.
@@ -135,11 +140,16 @@ export class InMemoryDistributedTransactionStorage
 
     const { retentionTime, idempotent } = options ?? {}
 
-    if (hasFinished) {
-      Object.assign(data, {
-        retention_time: retentionTime,
-      })
-    }
+    await this.#preventRaceConditionExecutionIfNecessary({
+      data,
+      key,
+      options,
+    })
+
+    Object.assign(data, {
+      retention_time: retentionTime,
+    })
+    this.storage.set(key, data)
 
     if (hasFinished && !retentionTime && !idempotent) {
       await this.deleteFromDb(data)
@@ -152,6 +162,118 @@ export class InMemoryDistributedTransactionStorage
     }
   }
 
+  async #preventRaceConditionExecutionIfNecessary({
+    data,
+    key,
+    options,
+  }: {
+    data: TransactionCheckpoint
+    key: string
+    options?: TransactionOptions
+  }) {
+    let isInitialCheckpoint = false
+
+    if (data.flow.state === TransactionState.NOT_STARTED) {
+      isInitialCheckpoint = true
+    }
+
+    /**
+     * In case many execution can succeed simultaneously, we need to ensure that the latest
+     * execution does continue if a previous execution is considered finished
+     */
+    const currentFlow = data.flow
+    const { flow: latestUpdatedFlow } =
+      (await this.get(key, options)) ??
+      ({ flow: {} } as { flow: TransactionFlow })
+
+    if (!isInitialCheckpoint && !isPresent(latestUpdatedFlow)) {
+      /**
+       * the initial checkpoint expect no other checkpoint to have been stored.
+       * In case it is not the initial one and another checkpoint is trying to
+       * find if a concurrent execution has finished, we skip the execution.
+       * The already finished execution would have deleted the checkpoint already.
+       */
+      throw new SkipExecutionError("Already finished by another execution")
+    }
+
+    const currentFlowLastInvokingStepIndex = Object.values(
+      currentFlow.steps
+    ).findIndex((step) => {
+      return [
+        TransactionStepState.INVOKING,
+        TransactionStepState.NOT_STARTED,
+      ].includes(step.invoke?.state)
+    })
+
+    const latestUpdatedFlowLastInvokingStepIndex = !latestUpdatedFlow.steps
+      ? 1 // There is no other execution, so the current execution is the latest
+      : Object.values(
+          (latestUpdatedFlow.steps as Record<string, TransactionStep>) ?? {}
+        ).findIndex((step) => {
+          return [
+            TransactionStepState.INVOKING,
+            TransactionStepState.NOT_STARTED,
+          ].includes(step.invoke?.state)
+        })
+
+    const currentFlowLastCompensatingStepIndex = Object.values(
+      currentFlow.steps
+    )
+      .reverse()
+      .findIndex((step) => {
+        return [
+          TransactionStepState.COMPENSATING,
+          TransactionStepState.NOT_STARTED,
+        ].includes(step.compensate?.state)
+      })
+
+    const latestUpdatedFlowLastCompensatingStepIndex = !latestUpdatedFlow.steps
+      ? -1 // There is no other execution, so the current execution is the latest
+      : Object.values(
+          (latestUpdatedFlow.steps as Record<string, TransactionStep>) ?? {}
+        )
+          .reverse()
+          .findIndex((step) => {
+            return [
+              TransactionStepState.COMPENSATING,
+              TransactionStepState.NOT_STARTED,
+            ].includes(step.compensate?.state)
+          })
+
+    const isLatestExecutionFinishedIndex = -1
+    const invokeShouldBeSkipped =
+      (latestUpdatedFlowLastInvokingStepIndex ===
+        isLatestExecutionFinishedIndex ||
+        currentFlowLastInvokingStepIndex <
+          latestUpdatedFlowLastInvokingStepIndex) &&
+      currentFlowLastInvokingStepIndex !== isLatestExecutionFinishedIndex
+
+    const compensateShouldBeSkipped =
+      currentFlowLastCompensatingStepIndex <
+        latestUpdatedFlowLastCompensatingStepIndex &&
+      currentFlowLastCompensatingStepIndex !== isLatestExecutionFinishedIndex &&
+      latestUpdatedFlowLastCompensatingStepIndex !==
+        isLatestExecutionFinishedIndex
+
+    if (
+      (data.flow.state !== TransactionState.COMPENSATING &&
+        invokeShouldBeSkipped) ||
+      (data.flow.state === TransactionState.COMPENSATING &&
+        compensateShouldBeSkipped) ||
+      (latestUpdatedFlow.state === TransactionState.COMPENSATING &&
+        ![TransactionState.REVERTED, TransactionState.FAILED].includes(
+          currentFlow.state
+        ) &&
+        currentFlow.state !== latestUpdatedFlow.state) ||
+      (latestUpdatedFlow.state === TransactionState.REVERTED &&
+        currentFlow.state !== TransactionState.REVERTED) ||
+      (latestUpdatedFlow.state === TransactionState.FAILED &&
+        currentFlow.state !== TransactionState.FAILED)
+    ) {
+      throw new SkipExecutionError("Already finished by another execution")
+    }
+  }
+
   async scheduleRetry(
     transaction: DistributedTransactionType,
     step: TransactionStep,