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

2025-03-12 09:53:34 -03:00
parent c97eaa0e0d
commit 72d2cf9207
24 changed files with 1130 additions and 235 deletions
@@ -8,7 +8,7 @@ import { setTimeout } from "timers/promises"
 const step_1 = createStep(
  "step_1",
  jest.fn(async (input) => {
-    await setTimeout(200)
+    await setTimeout(1000)

    return new StepResponse(input, { compensate: 123 })
  })
@@ -42,6 +42,8 @@ createWorkflow(
 createWorkflow(
  {
    name: "workflow_step_timeout_async",
+    idempotent: true,
+    retentionTime: 5,
  },
  function (input) {
    const resp = step_1_async(input)
@@ -33,6 +33,8 @@ createWorkflow(
  {
    name: "workflow_transaction_timeout_async",
    timeout: 0.1, // 0.1 second
+    idempotent: true,
+    retentionTime: 5,
  },
  function (input) {
    const resp = step_1(input).config({
@@ -335,7 +335,7 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
            throwOnError: false,
          })

-          await setTimeout(200)
+          await setTimeout(2000)

          const { transaction, result, errors } = (await workflowOrcModule.run(
            "workflow_step_timeout_async",
@@ -569,7 +569,6 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
          )
        })

-        // TODO: investigate why it fails intermittently
        it.skip("the scheduled workflow should have access to the shared container", async () => {
          const wait = times(1)
          sharedContainer_.register("test-value", asValue("test"))
@@ -0,0 +1,203 @@
+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(999900000)
+
+const failTrap = (done) => {
+  setTimeoutSync(() => {
+    // REF:https://stackoverflow.com/questions/78028715/jest-async-test-with-event-emitter-isnt-ending
+    console.warn(
+      "Jest is breaking the event emit with its debouncer. This allows to continue the test by managing the timeout of the test manually."
+    )
+    done()
+  }, 5000)
+}
+
+// REF:https://stackoverflow.com/questions/78028715/jest-async-test-with-event-emitter-isnt-ending
+
+moduleIntegrationTestRunner<IWorkflowEngineService>({
+  moduleName: Modules.WORKFLOW_ENGINE,
+  resolve: __dirname + "/../..",
+  moduleOptions: {
+    redis: {
+      url: "localhost:6379",
+    },
+  },
+  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 transactionId = "transaction_id"
+
+        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",
+          transactionId,
+          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", { transactionId })
+          .then(({ result }) => {
+            expect(result).toBe("result from step 0")
+          })
+
+        failTrap(done)
+      })
+
+      it("should prevent race continuation of the workflow compensation during retryIntervalAwaiting in background execution", (done) => {
+        const transactionId = "transaction_id"
+        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(500)
+            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,
+          transactionId,
+          subscriber: (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)
+              done()
+            }
+          },
+        })
+
+        workflowOrcModule
+          .run(workflowId, { transactionId })
+          .then(({ result }) => {
+            expect(result).toBe("result from step 0")
+          })
+
+        failTrap(done)
+      })
+    })
+  },
+})
@@ -12,13 +12,9 @@ import {
  Logger,
  MedusaContainer,
 } from "@medusajs/framework/types"
+import { isString, TransactionState } from "@medusajs/framework/utils"
 import {
-  InjectSharedContext,
-  isString,
-  MedusaContext,
-  TransactionState,
-} from "@medusajs/framework/utils"
-import {
+  FlowCancelOptions,
  FlowRunOptions,
  MedusaWorkflow,
  resolveValue,
@@ -37,6 +33,11 @@ export type WorkflowOrchestratorRunOptions<T> = Omit<
  container?: ContainerLike
 }

+export type WorkflowOrchestratorCancelOptions = Omit<
+  FlowCancelOptions,
+  "transaction"
+>
+
 type RegisterStepSuccessOptions<T> = Omit<
  WorkflowOrchestratorRunOptions<T>,
  "transactionId" | "input"
@@ -183,11 +184,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,
@@ -239,6 +238,7 @@ export class WorkflowOrchestratorService {
    const hasFinished = ret.transaction.hasFinished()
    const metadata = ret.transaction.getFlow().metadata
    const { parentStepIdempotencyKey } = metadata ?? {}
+
    const hasFailed = [
      TransactionState.REVERTED,
      TransactionState.FAILED,
@@ -271,12 +271,115 @@ 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?: { context?: Context }
  ): Promise<DistributedTransactionType> {
    let { context } = options ?? {}

@@ -289,7 +392,6 @@ export class WorkflowOrchestratorService {
    }

    context ??= {}
-    context.transactionId ??= transactionId

    const exportedWorkflow: any = MedusaWorkflow.getWorkflow(workflowId)
    if (!exportedWorkflow) {
@@ -304,19 +406,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,
@@ -375,19 +473,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,
@@ -446,11 +540,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()

@@ -487,11 +582,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[]) => {
@@ -75,7 +75,9 @@ export class WorkflowsModuleService<

      await this.clearExpiredExecutions()
      this.clearTimeout_ = setInterval(async () => {
-        await this.clearExpiredExecutions()
+        try {
+          await this.clearExpiredExecutions()
+        } catch {}
      }, 1000 * 60 * 60)
    },
  }
@@ -90,11 +92,13 @@ 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
  }
@@ -108,7 +112,7 @@ export class WorkflowsModuleService<
    return await this.workflowOrchestratorService_.getRunningTransaction(
      workflowId,
      transactionId,
-      context
+      { context }
    )
  }

@@ -125,14 +129,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()
@@ -148,14 +152,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()
@@ -168,7 +172,7 @@ export class WorkflowsModuleService<
    },
    @MedusaContext() context: Context = {}
  ) {
-    return this.workflowOrchestratorService_.subscribe(args as any, context)
+    return this.workflowOrchestratorService_.subscribe(args as any)
  }

  @InjectSharedContext()
@@ -180,7 +184,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() {
@@ -4,15 +4,19 @@ import {
  IDistributedSchedulerStorage,
  IDistributedTransactionStorage,
  SchedulerOptions,
+  SkipExecutionError,
  TransactionCheckpoint,
+  TransactionFlow,
  TransactionOptions,
  TransactionStep,
 } from "@medusajs/framework/orchestration"
 import { Logger, ModulesSdkTypes } from "@medusajs/framework/types"
 import {
+  isPresent,
  MedusaError,
  promiseAll,
  TransactionState,
+  TransactionStepState,
 } from "@medusajs/framework/utils"
 import { WorkflowOrchestratorService } from "@services"
 import { Queue, Worker } from "bullmq"
@@ -28,7 +32,6 @@ enum JobType {
 export class RedisDistributedTransactionStorage
  implements IDistributedTransactionStorage, IDistributedSchedulerStorage
 {
-  private static TTL_AFTER_COMPLETED = 60 * 2 // 2 minutes
  private workflowExecutionService_: ModulesSdkTypes.IMedusaInternalService<any>
  private logger_: Logger
  private workflowOrchestratorService_: WorkflowOrchestratorService
@@ -263,6 +266,12 @@ export class RedisDistributedTransactionStorage

    const { retentionTime, idempotent } = options ?? {}

+    await this.#preventRaceConditionExecutionIfNecessary({
+      data,
+      key,
+      options,
+    })
+
    if (hasFinished) {
      Object.assign(data, {
        retention_time: retentionTime,
@@ -286,12 +295,7 @@ export class RedisDistributedTransactionStorage
    }

    if (hasFinished) {
-      await this.redisClient.set(
-        key,
-        stringifiedData,
-        "EX",
-        RedisDistributedTransactionStorage.TTL_AFTER_COMPLETED
-      )
+      await this.redisClient.unlink(key)
    }
  }

@@ -457,4 +461,116 @@ export class RedisDistributedTransactionStorage
      repeatableJobs.map((job) => this.jobQueue?.removeRepeatableByKey(job.key))
    )
  }
+
+  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
+      : 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")
+    }
+  }
 }