feat: Add support for scheduled workflows (#7651)

We still need to:
But wanted to open the PR for early feedback on the approach

packages/modules/workflow-engine-inmemory/integration-tests/__fixtures__/workflow_scheduled.ts

@@ -0,0 +1,23 @@
+import { SchedulerOptions } from "@medusajs/orchestration"
+import {
+  createStep,
+  createWorkflow,
+  StepResponse,
+} from "@medusajs/workflows-sdk"
+
+export const createScheduled = (name: string, schedule?: SchedulerOptions) => {
+  const workflowScheduledStepInvoke = jest.fn((input, context) => {
+    return new StepResponse({})
+  })
+
+  const step = createStep("step_1", workflowScheduledStepInvoke)
+
+  createWorkflow(
+    { name, schedule: schedule ?? "* * * * * *" },
+    function (input) {
+      return step(input)
+    }
+  )
+
+  return workflowScheduledStepInvoke
+}

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

@@ -3,9 +3,10 @@ import { RemoteJoinerQuery } from "@medusajs/types"
 import { TransactionHandlerType } from "@medusajs/utils"
 import { IWorkflowEngineService } from "@medusajs/workflows-sdk"
 import { knex } from "knex"
-import { setTimeout } from "timers/promises"
+import { setTimeout as setTimeoutPromise } from "timers/promises"
 import "../__fixtures__"
 import { workflow2Step2Invoke, workflow2Step3Invoke } from "../__fixtures__"
+import { createScheduled } from "../__fixtures__/workflow_scheduled"
 import { DB_URL, TestDatabase } from "../utils"
 
 const sharedPgConnection = knex<any, any>({
@@ -24,42 +25,41 @@ const afterEach_ = async () => {
 jest.setTimeout(50000)
 
 describe("Workflow Orchestrator module", function () {
-  describe("Testing basic workflow", function () {
-    let workflowOrcModule: IWorkflowEngineService
-    let query: (
-      query: string | RemoteJoinerQuery | object,
-      variables?: Record<string, unknown>
-    ) => Promise<any>
+  let workflowOrcModule: IWorkflowEngineService
+  let query: (
+    query: string | RemoteJoinerQuery | object,
+    variables?: Record<string, unknown>
+  ) => Promise<any>
 
-    afterEach(afterEach_)
+  afterEach(afterEach_)
 
-    beforeAll(async () => {
-      const {
-        runMigrations,
-        query: remoteQuery,
-        modules,
-      } = await MedusaApp({
-        sharedResourcesConfig: {
-          database: {
-            connection: sharedPgConnection,
-          },
+  beforeAll(async () => {
+    const {
+      runMigrations,
+      query: remoteQuery,
+      modules,
+    } = await MedusaApp({
+      sharedResourcesConfig: {
+        database: {
+          connection: sharedPgConnection,
         },
-        modulesConfig: {
-          workflows: {
-            resolve: __dirname + "/../..",
-          },
+      },
+      modulesConfig: {
+        workflows: {
+          resolve: __dirname + "/../..",
         },
-      })
-
-      query = remoteQuery
-
-      await runMigrations()
-
-      workflowOrcModule = modules.workflows as unknown as IWorkflowEngineService
+      },
     })
 
-    afterEach(afterEach_)
+    query = remoteQuery
 
+    await runMigrations()
+
+    workflowOrcModule = modules.workflows as unknown as IWorkflowEngineService
+  })
+
+  afterEach(afterEach_)
+  describe("Testing basic workflow", function () {
     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: {
@@ -168,7 +168,7 @@ describe("Workflow Orchestrator module", function () {
         }
       )
 
-      await setTimeout(200)
+      await setTimeoutPromise(200)
 
       expect(transaction.flow.state).toEqual("reverted")
     })
@@ -201,4 +201,43 @@ describe("Workflow Orchestrator module", function () {
       expect(onFinish).toHaveBeenCalledTimes(0)
     })
   })
+
+  describe("Scheduled workflows", () => {
+    beforeAll(() => {
+      jest.useFakeTimers()
+      jest.spyOn(global, "setTimeout")
+    })
+
+    afterAll(() => {
+      jest.useRealTimers()
+    })
+
+    it("should execute a scheduled workflow", async () => {
+      const spy = createScheduled("standard")
+
+      await jest.runOnlyPendingTimersAsync()
+      expect(setTimeout).toHaveBeenCalledTimes(2)
+      expect(spy).toHaveBeenCalledTimes(1)
+
+      await jest.runOnlyPendingTimersAsync()
+      expect(setTimeout).toHaveBeenCalledTimes(3)
+      expect(spy).toHaveBeenCalledTimes(2)
+    })
+
+    it("should stop executions after the set number of executions", async () => {
+      const spy = await createScheduled("num-executions", {
+        cron: "* * * * * *",
+        numberOfExecutions: 2,
+      })
+
+      await jest.runOnlyPendingTimersAsync()
+      expect(spy).toHaveBeenCalledTimes(1)
+
+      await jest.runOnlyPendingTimersAsync()
+      expect(spy).toHaveBeenCalledTimes(2)
+
+      await jest.runOnlyPendingTimersAsync()
+      expect(spy).toHaveBeenCalledTimes(2)
+    })
+  })
 })

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

@@ -53,6 +53,7 @@
     "@mikro-orm/migrations": "5.9.7",
     "@mikro-orm/postgresql": "5.9.7",
     "awilix": "^8.0.0",
+    "cron-parser": "^4.9.0",
     "dotenv": "^16.4.5",
     "knex": "2.4.2"
   }

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

@@ -3,6 +3,7 @@ import {
   DistributedTransactionEvents,
   TransactionHandlerType,
   TransactionStep,
+  WorkflowScheduler,
 } from "@medusajs/orchestration"
 import { ContainerLike, Context, MedusaContainer } from "@medusajs/types"
 import { InjectSharedContext, MedusaContext, isString } from "@medusajs/utils"
@@ -83,6 +84,7 @@ export class WorkflowOrchestratorService {
   }) {
     inMemoryDistributedTransactionStorage.setWorkflowOrchestratorService(this)
     DistributedTransaction.setStorage(inMemoryDistributedTransactionStorage)
+    WorkflowScheduler.setStorage(inMemoryDistributedTransactionStorage)
   }
 
   @InjectSharedContext()

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

@@ -1,19 +1,34 @@
 import {
   DistributedTransaction,
   DistributedTransactionStorage,
+  IDistributedSchedulerStorage,
+  IDistributedTransactionStorage,
+  SchedulerOptions,
   TransactionCheckpoint,
   TransactionStep,
 } from "@medusajs/orchestration"
 import { ModulesSdkTypes } from "@medusajs/types"
 import { TransactionState } from "@medusajs/utils"
 import { WorkflowOrchestratorService } from "@services"
+import { CronExpression, parseExpression } from "cron-parser"
 
 // eslint-disable-next-line max-len
-export class InMemoryDistributedTransactionStorage extends DistributedTransactionStorage {
+export class InMemoryDistributedTransactionStorage
+  implements IDistributedTransactionStorage, IDistributedSchedulerStorage
+{
   private workflowExecutionService_: ModulesSdkTypes.InternalModuleService<any>
   private workflowOrchestratorService_: WorkflowOrchestratorService
 
   private storage: Map<string, TransactionCheckpoint> = new Map()
+  private scheduled: Map<
+    string,
+    {
+      timer: NodeJS.Timeout
+      expression: CronExpression
+      numberOfExecutions: number
+      config: SchedulerOptions
+    }
+  > = new Map()
   private retries: Map<string, unknown> = new Map()
   private timeouts: Map<string, unknown> = new Map()
 
@@ -22,8 +37,6 @@ export class InMemoryDistributedTransactionStorage extends DistributedTransactio
   }: {
     workflowExecutionService: ModulesSdkTypes.InternalModuleService<any>
   }) {
-    super()
-
     this.workflowExecutionService_ = workflowExecutionService
   }
 
@@ -215,4 +228,78 @@ export class InMemoryDistributedTransactionStorage extends DistributedTransactio
       this.timeouts.delete(key)
     }
   }
+
+  /* Scheduler storage methods */
+  async schedule(
+    jobDefinition: string | { jobId: string },
+    schedulerOptions: SchedulerOptions
+  ): Promise<void> {
+    const jobId =
+      typeof jobDefinition === "string" ? jobDefinition : jobDefinition.jobId
+
+    // In order to ensure that the schedule configuration is always up to date, we first cancel an existing job, if there was one
+    // any only then we add the new one.
+    await this.remove(jobId)
+    const expression = parseExpression(schedulerOptions.cron)
+    const nextExecution = expression.next().getTime() - Date.now()
+
+    const timer = setTimeout(async () => {
+      this.jobHandler(jobId)
+    }, nextExecution)
+
+    this.scheduled.set(jobId, {
+      timer,
+      expression,
+      numberOfExecutions: 0,
+      config: schedulerOptions,
+    })
+  }
+
+  async remove(jobId: string): Promise<void> {
+    const job = this.scheduled.get(jobId)
+    if (!job) {
+      return
+    }
+
+    clearTimeout(job.timer)
+    this.scheduled.delete(jobId)
+  }
+
+  async removeAll(): Promise<void> {
+    this.scheduled.forEach((_, key) => {
+      this.remove(key)
+    })
+  }
+
+  async jobHandler(jobId: string) {
+    const job = this.scheduled.get(jobId)
+    if (!job) {
+      return
+    }
+
+    if (
+      job.config?.numberOfExecutions !== undefined &&
+      job.config.numberOfExecutions <= job.numberOfExecutions
+    ) {
+      this.scheduled.delete(jobId)
+      return
+    }
+
+    const nextExecution = job.expression.next().getTime() - Date.now()
+    const timer = setTimeout(async () => {
+      this.jobHandler(jobId)
+    }, nextExecution)
+
+    this.scheduled.set(jobId, {
+      timer,
+      expression: job.expression,
+      numberOfExecutions: (job.numberOfExecutions ?? 0) + 1,
+      config: job.config,
+    })
+
+    // With running the job after setting a new timer we basically allow for concurrent runs, unless we add idempotency keys once they are supported.
+    await this.workflowOrchestratorService_.run(jobId, {
+      throwOnError: false,
+    })
+  }
 }