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

Author: adrien2p

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

@@ -56,6 +56,7 @@ export interface IDistributedTransactionStorage {
     transaction: DistributedTransactionType,
     step: TransactionStep
   ): Promise<void>
+  clearExpiredExecutions(): Promise<void>
 }
 
 export abstract class DistributedSchedulerStorage

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

@@ -1107,6 +1107,120 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
           })
         })
       })
+
+      describe("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 setTimeout(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 setTimeout(2000)
+
+          // Manually trigger cleaner
+          await (workflowOrcModule as any).workflowOrchestratorService_[
+            "redisDistributedTransactionStorage_"
+          ]["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-redis/src/services/workflows-module.ts

@@ -47,7 +47,6 @@ export class WorkflowsModuleService<
   protected workflowOrchestratorService_: WorkflowOrchestratorService
   protected redisDisconnectHandler_: () => Promise<void>
   protected manager_: SqlEntityManager
-  private clearTimeout_: NodeJS.Timeout
 
   constructor(
     {
@@ -72,21 +71,13 @@ export class WorkflowsModuleService<
   __hooks = {
     onApplicationStart: async () => {
       await this.workflowOrchestratorService_.onApplicationStart()
-
-      await this.clearExpiredExecutions()
-      this.clearTimeout_ = setInterval(async () => {
-        try {
-          await this.clearExpiredExecutions()
-        } catch {}
-      }, 1000 * 60 * 60)
     },
     onApplicationPrepareShutdown: async () => {
       await this.workflowOrchestratorService_.onApplicationPrepareShutdown()
     },
     onApplicationShutdown: async () => {
       await this.workflowOrchestratorService_.onApplicationShutdown()
       await this.redisDisconnectHandler_()
-      clearInterval(this.clearTimeout_)
     },
   }
 
@@ -301,14 +292,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(
     workflowId: string,

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

@@ -22,6 +22,7 @@ import {
   TransactionState,
   TransactionStepState,
 } from "@medusajs/framework/utils"
+import { raw } from "@mikro-orm/core"
 import { WorkflowOrchestratorService } from "@services"
 import { Queue, RepeatOptions, Worker } from "bullmq"
 import Redis from "ioredis"
@@ -33,6 +34,9 @@ enum JobType {
   TRANSACTION_TIMEOUT = "transaction_timeout",
 }
 
+const ONE_HOUR_IN_MS = 1000 * 60 * 60
+const REPEATABLE_CLEARER_JOB_ID = "clear-expired-executions"
+
 export class RedisDistributedTransactionStorage
   implements IDistributedTransactionStorage, IDistributedSchedulerStorage
 {
@@ -48,6 +52,9 @@ export class RedisDistributedTransactionStorage
   private jobQueue?: Queue
   private worker: Worker
   private jobWorker?: Worker
+  private cleanerQueueName: string
+  private cleanerWorker_: Worker
+  private cleanerQueue_?: Queue
 
   #isWorkerMode: boolean = false
 
@@ -72,6 +79,7 @@ export class RedisDistributedTransactionStorage
     this.logger_ = logger
     this.redisClient = redisConnection
     this.redisWorkerConnection = redisWorkerConnection
+    this.cleanerQueueName = "workflows-cleaner"
     this.queueName = redisQueueName
     this.jobQueueName = redisJobQueueName
     this.queue = new Queue(redisQueueName, { connection: this.redisClient })
@@ -80,18 +88,33 @@ export class RedisDistributedTransactionStorage
           connection: this.redisClient,
         })
       : undefined
+    this.cleanerQueue_ = isWorkerMode
+      ? new Queue(this.cleanerQueueName, {
+          connection: this.redisClient,
+        })
+      : undefined
     this.#isWorkerMode = isWorkerMode
   }
 
   async onApplicationPrepareShutdown() {
     // Close worker gracefully, i.e. wait for the current jobs to finish
     await this.worker?.close()
     await this.jobWorker?.close()
+
+    const repeatableJobs = (await this.cleanerQueue_?.getRepeatableJobs()) ?? []
+    for (const job of repeatableJobs) {
+      if (job.id === REPEATABLE_CLEARER_JOB_ID) {
+        await this.cleanerQueue_?.removeRepeatableByKey(job.key)
+      }
+    }
+
+    await this.cleanerWorker_?.close()
   }
 
   async onApplicationShutdown() {
     await this.queue?.close()
     await this.jobQueue?.close()
+    await this.cleanerQueue_?.close()
   }
 
   async onApplicationStart() {
@@ -151,6 +174,27 @@ export class RedisDistributedTransactionStorage
         },
         workerOptions
       )
+
+      this.cleanerWorker_ = new Worker(
+        this.cleanerQueueName,
+        async () => {
+          await this.clearExpiredExecutions()
+        },
+        { connection: this.redisClient }
+      )
+
+      await this.cleanerQueue_?.add(
+        "cleaner",
+        {},
+        {
+          repeat: {
+            every: ONE_HOUR_IN_MS,
+          },
+          jobId: REPEATABLE_CLEARER_JOB_ID,
+          removeOnComplete: true,
+          removeOnFail: true,
+        }
+      )
     }
   }
 
@@ -728,4 +772,25 @@ export class RedisDistributedTransactionStorage
       throw new SkipExecutionError("Already finished by another execution")
     }
   }
+
+  private async clearExpiredExecutions() {
+    await this.workflowExecutionService_.delete({
+      retention_time: {
+        $ne: null,
+      },
+      updated_at: {
+        $lte: raw(
+          (alias) =>
+            `CURRENT_TIMESTAMP - (INTERVAL '1 second' * retention_time)`
+        ),
+      },
+      state: {
+        $in: [
+          TransactionState.DONE,
+          TransactionState.FAILED,
+          TransactionState.REVERTED,
+        ],
+      },
+    })
+  }
 }