fix(workflow-engine-*): Cleanup expired executions

Author: adrien2p

packages/core/orchestration/src/transaction/datastore/abstract-storage.ts

@@ -150,4 +150,8 @@ export abstract class DistributedTransactionStorage
   ): Promise<void> {
     throw new Error("Method 'clearStepTimeout' not implemented.")
   }
+
+  async clearExpiredExecutions(): Promise<void> {
+    throw new Error("Method 'clearExpiredExecutions' not implemented.")
+  }
 }

packages/core/orchestration/src/transaction/datastore/base-in-memory-storage.ts

@@ -41,4 +41,6 @@ export class BaseInMemoryDistributedTransactionStorage extends DistributedTransa
       this.storage.set(key, data)
     }
   }
+
+  async clearExpiredExecutions(): Promise<void> {}
 }

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

@@ -931,6 +931,120 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
           expect(executionsListAfter).toHaveLength(1)
         })
       })
+
+      describe.only("Cleaner job", function () {
+        it("should remove expired executions of finished workflows and keep the others", async () => {
+          const doneWorkflowId = "done-workflow-" + ulid()
+          createWorkflow({ name: doneWorkflowId, retentionTime: 1 }, () => {
+            return new WorkflowResponse("done")
+          })
+
+          const failingWorkflowId = "failing-workflow-" + ulid()
+          const failingStep = createStep("failing-step", () => {
+            throw new Error("I am failing")
+          })
+          createWorkflow({ name: failingWorkflowId, retentionTime: 1 }, () => {
+            failingStep()
+          })
+
+          const revertingStep = createStep(
+            "reverting-step",
+            () => {
+              throw new Error("I am reverting")
+            },
+            () => {
+              return new StepResponse("reverted")
+            }
+          )
+
+          const revertingWorkflowId = "reverting-workflow-" + ulid()
+          createWorkflow(
+            { name: revertingWorkflowId, retentionTime: 1 },
+            () => {
+              revertingStep()
+              return new WorkflowResponse("reverted")
+            }
+          )
+
+          const runningWorkflowId = "running-workflow-" + ulid()
+          const longRunningStep = createStep("long-running-step", async () => {
+            await setTimeoutPromise(10000)
+            return new StepResponse("long running finished")
+          })
+          createWorkflow({ name: runningWorkflowId, retentionTime: 1 }, () => {
+            longRunningStep().config({ async: true, backgroundExecution: true })
+            return new WorkflowResponse("running workflow started")
+          })
+
+          const notExpiredWorkflowId = "not-expired-workflow-" + ulid()
+          createWorkflow(
+            { name: notExpiredWorkflowId, retentionTime: 1000 },
+            () => {
+              return new WorkflowResponse("not expired")
+            }
+          )
+
+          const trx_done = "trx-done-" + ulid()
+          const trx_failed = "trx-failed-" + ulid()
+          const trx_reverting = "trx-reverting-" + ulid()
+          const trx_running = "trx-running-" + ulid()
+          const trx_not_expired = "trx-not-expired-" + ulid()
+
+          // run workflows
+          await workflowOrcModule.run(doneWorkflowId, {
+            transactionId: trx_done,
+          })
+
+          await workflowOrcModule.run(failingWorkflowId, {
+            transactionId: trx_failed,
+            throwOnError: false,
+          })
+
+          await workflowOrcModule.run(revertingWorkflowId, {
+            transactionId: trx_reverting,
+            throwOnError: false,
+          })
+
+          await workflowOrcModule.run(runningWorkflowId, {
+            transactionId: trx_running,
+          })
+
+          await workflowOrcModule.run(notExpiredWorkflowId, {
+            transactionId: trx_not_expired,
+          })
+
+          let executions = await workflowOrcModule.listWorkflowExecutions()
+          expect(executions).toHaveLength(5)
+
+          await setTimeoutPromise(2000)
+
+          // Manually trigger cleaner
+          await (workflowOrcModule as any).workflowOrchestratorService_[
+            "inMemoryDistributedTransactionStorage_"
+          ]["clearExpiredExecutions"]()
+
+          let remainingExecutions =
+            await workflowOrcModule.listWorkflowExecutions()
+
+          expect(remainingExecutions).toHaveLength(2)
+
+          const remainingTrxIds = remainingExecutions
+            .map((e) => e.transaction_id)
+            .sort()
+
+          expect(remainingTrxIds).toEqual([trx_not_expired, trx_running].sort())
+
+          const notExpiredExec = remainingExecutions.find(
+            (e) => e.transaction_id === trx_not_expired
+          )
+          expect(notExpiredExec?.state).toBe(TransactionState.DONE)
+
+          const runningExec = remainingExecutions.find(
+            (e) => e.transaction_id === trx_running
+          )
+          expect(runningExec?.state).toBe(TransactionState.INVOKING)
+        })
+      })
     })
   },
 })

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

@@ -87,6 +87,7 @@ const AnySubscriber = "any"
 export class WorkflowOrchestratorService {
   private subscribers: Subscribers = new Map()
   private container_: MedusaContainer
+  private inMemoryDistributedTransactionStorage_: InMemoryDistributedTransactionStorage
 
   constructor({
     inMemoryDistributedTransactionStorage,
@@ -97,11 +98,21 @@ export class WorkflowOrchestratorService {
     sharedContainer: MedusaContainer
   }) {
     this.container_ = sharedContainer
+    this.inMemoryDistributedTransactionStorage_ =
+      inMemoryDistributedTransactionStorage
     inMemoryDistributedTransactionStorage.setWorkflowOrchestratorService(this)
     DistributedTransaction.setStorage(inMemoryDistributedTransactionStorage)
     WorkflowScheduler.setStorage(inMemoryDistributedTransactionStorage)
   }
 
+  async onApplicationStart() {
+    await this.inMemoryDistributedTransactionStorage_.onApplicationStart()
+  }
+
+  async onApplicationShutdown() {
+    await this.inMemoryDistributedTransactionStorage_.onApplicationShutdown()
+  }
+
   private async triggerParentStep(transaction, result) {
     const metadata = transaction.flow.metadata
     const { parentStepIdempotencyKey } = metadata ?? {}

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

@@ -43,7 +43,6 @@ export class WorkflowsModuleService<
   protected workflowExecutionService_: ModulesSdkTypes.IMedusaInternalService<TWorkflowExecution>
   protected workflowOrchestratorService_: WorkflowOrchestratorService
   protected manager_: SqlEntityManager
-  private clearTimeout_: NodeJS.Timeout
 
   constructor(
     {
@@ -65,16 +64,10 @@ export class WorkflowsModuleService<
 
   __hooks = {
     onApplicationStart: async () => {
-      await this.clearExpiredExecutions()
-
-      this.clearTimeout_ = setInterval(async () => {
-        try {
-          await this.clearExpiredExecutions()
-        } catch {}
-      }, 1000 * 60 * 60)
+      await this.workflowOrchestratorService_.onApplicationStart()
     },
     onApplicationShutdown: async () => {
-      clearInterval(this.clearTimeout_)
+      await this.workflowOrchestratorService_.onApplicationShutdown()
     },
   }
 
@@ -289,14 +282,6 @@ export class WorkflowsModuleService<
     return this.workflowOrchestratorService_.unsubscribe(args as any)
   }
 
-  private async clearExpiredExecutions() {
-    return this.manager_.execute(`
-      DELETE FROM workflow_execution
-      WHERE retention_time IS NOT NULL AND
-      updated_at <= (CURRENT_TIMESTAMP - INTERVAL '1 second' * retention_time);
-    `)
-  }
-
   @InjectSharedContext()
   async cancel<TWorkflow extends string | ReturnWorkflow<any, any, any>>(
     workflowIdOrWorkflow: TWorkflow,

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

@@ -27,6 +27,7 @@ import {
 import { WorkflowOrchestratorService } from "@services"
 import { type CronExpression, parseExpression } from "cron-parser"
 import { WorkflowExecution } from "../models/workflow-execution"
+import { raw } from "@mikro-orm/core"
 
 function parseNextExecution(
   optionsOrExpression: SchedulerOptions | CronExpression | string | number
@@ -74,6 +75,8 @@ export class InMemoryDistributedTransactionStorage
   private retries: Map<string, unknown> = new Map()
   private timeouts: Map<string, unknown> = new Map()
 
+  private clearTimeout_: NodeJS.Timeout
+
   constructor({
     workflowExecutionService,
     logger,
@@ -85,6 +88,18 @@ export class InMemoryDistributedTransactionStorage
     this.logger_ = logger
   }
 
+  async onApplicationStart() {
+    this.clearTimeout_ = setInterval(async () => {
+      try {
+        await this.clearExpiredExecutions()
+      } catch {}
+    }, 1000 * 60 * 60)
+  }
+
+  async onApplicationShutdown() {
+    clearInterval(this.clearTimeout_)
+  }
+
   setWorkflowOrchestratorService(workflowOrchestratorService) {
     this.workflowOrchestratorService_ = workflowOrchestratorService
   }
@@ -613,4 +628,26 @@ export class InMemoryDistributedTransactionStorage
       throw e
     }
   }
+
+  // TODO: Move
+  async clearExpiredExecutions(): Promise<void> {
+    await this.workflowExecutionService_.delete({
+      retention_time: {
+        $ne: null,
+      },
+      created_at: {
+        $lte: raw(
+          (alias) =>
+            `CURRENT_TIMESTAMP - (INTERVAL '1 second' * retention_time)`
+        ),
+      },
+      state: {
+        $in: [
+          TransactionState.DONE,
+          TransactionState.FAILED,
+          TransactionState.REVERTED,
+        ],
+      },
+    })
+  }
 }

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

@@ -773,15 +773,15 @@ export class RedisDistributedTransactionStorage
     }
   }
 
-  private async clearExpiredExecutions() {
+  async clearExpiredExecutions() {
     await this.workflowExecutionService_.delete({
       retention_time: {
         $ne: null,
       },
-      updated_at: {
+      created_at: {
         $lte: raw(
           (alias) =>
-            `CURRENT_TIMESTAMP - (INTERVAL '1 second' * retention_time)`
+            `CURRENT_TIMESTAMP - (INTERVAL '1 second' * "retention_time")`
         ),
       },
       state: {