feat: run nested async workflows (#9119)

2024-09-16 10:06:45 -03:00
parent 0bcdcccbe2
commit ef8dc4087e
23 changed files with 295 additions and 100 deletions
--- a/packages/modules/workflow-engine-redis/integration-tests/fixtures/workflow_async.ts
+++ b/packages/modules/workflow-engine-redis/integration-tests/fixtures/workflow_async.ts
@@ -1,7 +1,9 @@
 import {
  StepResponse,
+  WorkflowResponse,
  createStep,
  createWorkflow,
+  parallelize,
 } from "@medusajs/workflows-sdk"
 import { setTimeout } from "timers/promises"

@@ -11,15 +13,15 @@ const step_1_background = createStep(
    async: true,
  },
  jest.fn(async (input) => {
-    await setTimeout(200)
+    await setTimeout(Math.random() * 300)

    return new StepResponse(input)
  })
 )

-createWorkflow(
+const nestedWorkflow = createWorkflow(
  {
-    name: "workflow_async_background",
+    name: "nested_sub_flow_async",
  },
  function (input) {
    const resp = step_1_background(input)
@@ -27,3 +29,35 @@ createWorkflow(
    return resp
  }
 )
+
+createWorkflow(
+  {
+    name: "workflow_async_background",
+  },
+  function (input) {
+    const [ret] = parallelize(
+      nestedWorkflow
+        .runAsStep({
+          input,
+        })
+        .config({ name: "step_sub_flow_1" }),
+      nestedWorkflow
+        .runAsStep({
+          input,
+        })
+        .config({ name: "step_sub_flow_2" }),
+      nestedWorkflow
+        .runAsStep({
+          input,
+        })
+        .config({ name: "step_sub_flow_3" }),
+      nestedWorkflow
+        .runAsStep({
+          input,
+        })
+        .config({ name: "step_sub_flow_4" })
+    )
+
+    return new WorkflowResponse(ret)
+  }
+)
--- a/packages/modules/workflow-engine-redis/integration-tests/tests/index.spec.ts
+++ b/packages/modules/workflow-engine-redis/integration-tests/tests/index.spec.ts
@@ -23,7 +23,7 @@ import "../__fixtures__"
 import { createScheduled } from "../__fixtures__/workflow_scheduled"
 import { TestDatabase } from "../utils"

-jest.setTimeout(100000)
+jest.setTimeout(999900000)

 moduleIntegrationTestRunner<IWorkflowEngineService>({
  moduleName: Modules.WORKFLOW_ENGINE,
@@ -35,10 +35,8 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
  },
  testSuite: ({ service: workflowOrcModule, medusaApp }) => {
    describe("Workflow Orchestrator module", function () {
-      const afterEach_ = async () => {
+      beforeEach(async () => {
        await TestDatabase.clearTables()
-      }
-      beforeEach(() => {
        jest.clearAllMocks()
      })

@@ -93,8 +91,6 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
      })

      describe("Testing basic workflow", function () {
-        afterEach(afterEach_)
-
        it("should return a list of workflow executions and remove after completed when there is no retentionTime set", async () => {
          await workflowOrcModule.run("workflow_1", {
            input: {
@@ -271,34 +267,31 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
          ).toBe(true)
        })

-        it("should complete an async workflow that returns a StepResponse", async () => {
-          const { transaction, result } = await workflowOrcModule.run(
-            "workflow_async_background",
-            {
+        it("should complete an async workflow that returns a StepResponse", (done) => {
+          const transactionId = "transaction_1"
+          void workflowOrcModule
+            .run("workflow_async_background", {
              input: {
                myInput: "123",
              },
-              transactionId: "transaction_1",
-              throwOnError: false,
-            }
-          )
+              transactionId,
+              throwOnError: true,
+            })
+            .then(({ transaction, result }) => {
+              expect(transaction.flow.state).toEqual(
+                TransactionStepState.INVOKING
+              )
+              expect(result).toEqual(undefined)
+            })

-          expect(transaction.flow.state).toEqual(TransactionStepState.INVOKING)
-          expect(result).toEqual(undefined)
-
-          await setTimeout(205)
-
-          const trx = await workflowOrcModule.run("workflow_async_background", {
-            input: {
-              myInput: "123",
+          void workflowOrcModule.subscribe({
+            workflowId: "workflow_async_background",
+            transactionId,
+            subscriber: (event) => {
+              if (event.eventType === "onFinish") {
+                done()
+              }
            },
-            transactionId: "transaction_1",
-            throwOnError: false,
-          })
-
-          expect(trx.transaction.flow.state).toEqual(TransactionStepState.DONE)
-          expect(trx.result).toEqual({
-            myInput: "123",
          })
        })

--- a/packages/modules/workflow-engine-redis/package.json
+++ b/packages/modules/workflow-engine-redis/package.json
@@ -25,7 +25,7 @@
    "watch:test": "tsc --build tsconfig.spec.json --watch",
    "build": "rimraf dist && tsc --build && tsc-alias -p tsconfig.json",
    "test": "jest --passWithNoTests --runInBand --bail --forceExit -- src",
-    "test:integration": "jest --forceExit -- integration-tests/**/__tests__/**/*.ts",
+    "test:integration": "jest  --forceExit -- integration-tests/**/__tests__/**/*.ts",
    "migration:generate": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts mikro-orm migration:generate",
    "migration:initial": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts mikro-orm migration:create --initial",
    "migration:create": " MIKRO_ORM_CLI=./mikro-orm.config.dev.ts mikro-orm migration:create",
@@ -47,13 +47,14 @@
    "@medusajs/orchestration": "^0.5.7",
    "@medusajs/utils": "^1.11.9",
    "@medusajs/workflows-sdk": "^0.1.6",
-    "bullmq": "5.12.0",
+    "bullmq": "5.13.0",
    "ioredis": "^5.4.1"
  },
  "peerDependencies": {
    "@mikro-orm/core": "5.9.7",
    "@mikro-orm/migrations": "5.9.7",
    "@mikro-orm/postgresql": "5.9.7",
-    "awilix": "^8.0.1"
+    "awilix": "^8.0.1",
+    "ulid": "^2.3.0"
  }
 }
--- a/packages/modules/workflow-engine-redis/src/services/workflow-orchestrator.ts
+++ b/packages/modules/workflow-engine-redis/src/services/workflow-orchestrator.ts
@@ -12,7 +12,12 @@ import {
  Logger,
  MedusaContainer,
 } from "@medusajs/types"
-import { InjectSharedContext, MedusaContext, isString } from "@medusajs/utils"
+import {
+  InjectSharedContext,
+  MedusaContext,
+  TransactionState,
+  isString,
+} from "@medusajs/utils"
 import {
  FlowRunOptions,
  MedusaWorkflow,
@@ -146,6 +151,29 @@ export class WorkflowOrchestratorService {
    await this.redisDistributedTransactionStorage_.onApplicationStart()
  }

+  private async triggerParentStep(transaction, result) {
+    const metadata = transaction.flow.metadata
+    const { parentStepIdempotencyKey } = metadata ?? {}
+    if (parentStepIdempotencyKey) {
+      const hasFailed = [
+        TransactionState.REVERTED,
+        TransactionState.FAILED,
+      ].includes(transaction.flow.state)
+
+      if (hasFailed) {
+        await this.setStepFailure({
+          idempotencyKey: parentStepIdempotencyKey,
+          stepResponse: result,
+        })
+      } else {
+        await this.setStepSuccess({
+          idempotencyKey: parentStepIdempotencyKey,
+          stepResponse: result,
+        })
+      }
+    }
+  }
+
  @InjectSharedContext()
  async run<T = unknown>(
    workflowIdOrWorkflow: string | ReturnWorkflow<any, any, any>,
@@ -180,16 +208,12 @@ export class WorkflowOrchestratorService {
      transactionId: context.transactionId,
    })

-    const exportedWorkflow: any = MedusaWorkflow.getWorkflow(workflowId)
+    const exportedWorkflow = MedusaWorkflow.getWorkflow(workflowId)
    if (!exportedWorkflow) {
      throw new Error(`Workflow with id "${workflowId}" not found.`)
    }

-    const flow = exportedWorkflow(
-      (container as MedusaContainer) ?? this.container_
-    )
-
-    const ret = await flow.run({
+    const ret = await exportedWorkflow.run({
      input,
      throwOnError,
      logOnError,
@@ -198,14 +222,25 @@ export class WorkflowOrchestratorService {
      events,
    })

-    // TODO: temporary
+    const hasFinished = ret.transaction.hasFinished()
+    const metadata = ret.transaction.getFlow().metadata
+    const { parentStepIdempotencyKey } = metadata ?? {}
+    const hasFailed = [
+      TransactionState.REVERTED,
+      TransactionState.FAILED,
+    ].includes(ret.transaction.getFlow().state)
+
    const acknowledgement = {
      transactionId: context.transactionId,
      workflowId: workflowId,
+      parentStepIdempotencyKey,
+      hasFinished,
+      hasFailed,
    }

    if (ret.transaction.hasFinished()) {
      const { result, errors } = ret
+
      await this.notify({
        eventType: "onFinish",
        workflowId,
@@ -213,6 +248,8 @@ export class WorkflowOrchestratorService {
        result,
        errors,
      })
+
+      await this.triggerParentStep(ret.transaction, result)
    }

    return { acknowledgement, ...ret }
@@ -243,12 +280,11 @@ export class WorkflowOrchestratorService {
      throw new Error(`Workflow with id "${workflowId}" not found.`)
    }

-    const flow = exportedWorkflow(
-      (container as MedusaContainer) ?? this.container_
+    const transaction = await exportedWorkflow.getRunningTransaction(
+      transactionId,
+      context
    )

-    const transaction = await flow.getRunningTransaction(transactionId, context)
-
    return transaction
  }

@@ -282,17 +318,13 @@ export class WorkflowOrchestratorService {
      throw new Error(`Workflow with id "${workflowId}" not found.`)
    }

-    const flow = exportedWorkflow(
-      (container as MedusaContainer) ?? this.container_
-    )
-
    const events = this.buildWorkflowEvents({
      customEventHandlers: eventHandlers,
      transactionId,
      workflowId,
    })

-    const ret = await flow.registerStepSuccess({
+    const ret = await exportedWorkflow.registerStepSuccess({
      idempotencyKey: idempotencyKey_,
      context,
      resultFrom,
@@ -304,6 +336,7 @@ export class WorkflowOrchestratorService {

    if (ret.transaction.hasFinished()) {
      const { result, errors } = ret
+
      await this.notify({
        eventType: "onFinish",
        workflowId,
@@ -311,6 +344,8 @@ export class WorkflowOrchestratorService {
        result,
        errors,
      })
+
+      await this.triggerParentStep(ret.transaction, result)
    }

    return ret
@@ -368,6 +403,7 @@ export class WorkflowOrchestratorService {

    if (ret.transaction.hasFinished()) {
      const { result, errors } = ret
+
      await this.notify({
        eventType: "onFinish",
        workflowId,
@@ -375,6 +411,8 @@ export class WorkflowOrchestratorService {
        result,
        errors,
      })
+
+      await this.triggerParentStep(ret.transaction, result)
    }

    return ret
--- a/packages/modules/workflow-engine-redis/src/utils/workflow-orchestrator-storage.ts
+++ b/packages/modules/workflow-engine-redis/src/utils/workflow-orchestrator-storage.ts
@@ -66,16 +66,16 @@ export class RedisDistributedTransactionStorage
  }

  async onApplicationStart() {
+    const allowedJobs = [
+      JobType.RETRY,
+      JobType.STEP_TIMEOUT,
+      JobType.TRANSACTION_TIMEOUT,
+    ]
+
    this.worker = new Worker(
      this.queueName,
      async (job) => {
-        const allJobs = [
-          JobType.RETRY,
-          JobType.STEP_TIMEOUT,
-          JobType.TRANSACTION_TIMEOUT,
-        ]
-
-        if (allJobs.includes(job.name as JobType)) {
+        if (allowedJobs.includes(job.name as JobType)) {
          await this.executeTransaction(
            job.data.workflowId,
            job.data.transactionId