fix(workflow-engine-*): Cleanup expired executions and reduce redis storage usage (#12795)

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

@@ -21,9 +21,9 @@ import {
   MedusaError,
   TransactionState,
   TransactionStepState,
-  isDefined,
   isPresent,
 } from "@medusajs/framework/utils"
+import { raw } from "@mikro-orm/core"
 import { WorkflowOrchestratorService } from "@services"
 import { type CronExpression, parseExpression } from "cron-parser"
 import { WorkflowExecution } from "../models/workflow-execution"
@@ -61,7 +61,8 @@ export class InMemoryDistributedTransactionStorage
   private logger_: Logger
   private workflowOrchestratorService_: WorkflowOrchestratorService
 
-  private storage: Map<string, TransactionCheckpoint> = new Map()
+  private storage: Map<string, Omit<TransactionCheckpoint, "context">> =
+    new Map()
   private scheduled: Map<
     string,
     {
@@ -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
   }
@@ -122,17 +137,6 @@ export class InMemoryDistributedTransactionStorage
       isCancelling?: boolean
     }
   ): Promise<TransactionCheckpoint | undefined> {
-    const data = this.storage.get(key)
-
-    if (data) {
-      return data
-    }
-
-    const { idempotent, store, retentionTime } = options ?? {}
-    if (!idempotent && !(store && isDefined(retentionTime))) {
-      return
-    }
-
     const [_, workflowId, transactionId] = key.split(":")
     const trx: InferEntityType<typeof WorkflowExecution> | undefined =
       await this.workflowExecutionService_
@@ -153,6 +157,7 @@ export class InMemoryDistributedTransactionStorage
         .catch(() => undefined)
 
     if (trx) {
+      const { idempotent } = options ?? {}
       const execution = trx.execution as TransactionFlow
 
       if (!idempotent) {
@@ -187,10 +192,6 @@ export class InMemoryDistributedTransactionStorage
     return
   }
 
-  async list(): Promise<TransactionCheckpoint[]> {
-    return Array.from(this.storage.values())
-  }
-
   async save(
     key: string,
     data: TransactionCheckpoint,
@@ -237,7 +238,11 @@ export class InMemoryDistributedTransactionStorage
       }
     }
 
-    this.storage.set(key, data)
+    const { flow, errors } = data
+    this.storage.set(key, {
+      flow,
+      errors,
+    })
 
     // Optimize DB operations - only perform when necessary
     if (hasFinished) {
@@ -272,14 +277,22 @@ export class InMemoryDistributedTransactionStorage
      */
     const currentFlow = data.flow
 
-    const getOptions = {
-      ...options,
-      isCancelling: !!data.flow.cancelledAt,
-    } as Parameters<typeof this.get>[1]
+    const rawData = this.storage.get(key)
+    let data_ = {} as TransactionCheckpoint
+    if (rawData) {
+      data_ = rawData as TransactionCheckpoint
+    } else {
+      const getOptions = {
+        ...options,
+        isCancelling: !!data.flow.cancelledAt,
+      } as Parameters<typeof this.get>[1]
 
-    const { flow: latestUpdatedFlow } =
-      (await this.get(key, getOptions)) ??
-      ({ flow: {} } as { flow: TransactionFlow })
+      data_ =
+        (await this.get(key, getOptions)) ??
+        ({ flow: {} } as TransactionCheckpoint)
+    }
+
+    const { flow: latestUpdatedFlow } = data_
 
     if (!isInitialCheckpoint && !isPresent(latestUpdatedFlow)) {
       /**
@@ -613,4 +626,25 @@ export class InMemoryDistributedTransactionStorage
       throw e
     }
   }
+
+  async clearExpiredExecutions(): Promise<void> {
+    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,
+        ],
+      },
+    })
+  }
 }