fix(cleanup): batch orphaned snapshot deletes to avoid slow-query spike (#4348)

* fix(cleanup): batch orphaned snapshot deletes to avoid slow-query spike

* fix(cleanup): recheck orphan condition in delete to close TOCTOU gap

Author: waleedlatif1

apps/sim/lib/logs/execution/snapshot/service.test.ts

@@ -533,4 +533,73 @@ describe('SnapshotService', () => {
       expect(databaseMock.db.select).not.toHaveBeenCalled()
     })
   })
+
+  describe('cleanupOrphanedSnapshots', () => {
+    function setupCleanupMocks(selectBatches: Array<Array<{ id: string }>>) {
+      const limitFn = vi.fn()
+      for (const batch of selectBatches) limitFn.mockResolvedValueOnce(batch)
+      limitFn.mockResolvedValue([])
+      const whereSelect = vi.fn().mockReturnValue({ limit: limitFn })
+      const fromFn = vi.fn().mockReturnValue({ where: whereSelect })
+      databaseMock.db.select = vi.fn().mockReturnValue({ from: fromFn })
+
+      const returningFn = vi.fn().mockImplementation(() => Promise.resolve([]))
+      const whereDelete = vi.fn().mockReturnValue({ returning: returningFn })
+      let batchIdx = 0
+      const deleteFn = vi.fn().mockImplementation(() => {
+        const batch = selectBatches[batchIdx] ?? []
+        batchIdx++
+        returningFn.mockImplementationOnce(() => Promise.resolve(batch.map((r) => ({ id: r.id }))))
+        return { where: whereDelete }
+      })
+      databaseMock.db.delete = deleteFn
+
+      return { deleteFn }
+    }
+
+    it('returns 0 and skips delete when nothing is orphaned', async () => {
+      const service = new SnapshotService()
+      const { deleteFn } = setupCleanupMocks([])
+
+      const count = await service.cleanupOrphanedSnapshots(7)
+
+      expect(count).toBe(0)
+      expect(deleteFn).not.toHaveBeenCalled()
+    })
+
+    it('stops after the first short batch', async () => {
+      const service = new SnapshotService()
+      const partial = Array.from({ length: 3 }, (_, i) => ({ id: `s${i}` }))
+      const { deleteFn } = setupCleanupMocks([partial])
+
+      const count = await service.cleanupOrphanedSnapshots(7)
+
+      expect(count).toBe(3)
+      expect(deleteFn).toHaveBeenCalledTimes(1)
+    })
+
+    it('loops through multiple full batches until exhausted', async () => {
+      const service = new SnapshotService()
+      const fullBatch = Array.from({ length: 1000 }, (_, i) => ({ id: `s${i}` }))
+      const tail = [{ id: 'tail-1' }]
+      const { deleteFn } = setupCleanupMocks([fullBatch, fullBatch, tail])
+
+      const count = await service.cleanupOrphanedSnapshots(7)
+
+      expect(count).toBe(2001)
+      expect(deleteFn).toHaveBeenCalledTimes(3)
+    })
+
+    it('caps at MAX_BATCHES (20 × 1000) even when more rows remain', async () => {
+      const service = new SnapshotService()
+      const fullBatch = Array.from({ length: 1000 }, (_, i) => ({ id: `s${i}` }))
+      const batches = Array.from({ length: 25 }, () => fullBatch)
+      const { deleteFn } = setupCleanupMocks(batches)
+
+      const count = await service.cleanupOrphanedSnapshots(7)
+
+      expect(count).toBe(20_000)
+      expect(deleteFn).toHaveBeenCalledTimes(20)
+    })
+  })
 })

apps/sim/lib/logs/execution/snapshot/service.ts

@@ -3,7 +3,7 @@ import { workflowExecutionLogs, workflowExecutionSnapshots } from '@sim/db/schem
 import { createLogger } from '@sim/logger'
 import { sha256Hex } from '@sim/security/hash'
 import { generateId } from '@sim/utils/id'
-import { and, eq, lt, notExists, sql } from 'drizzle-orm'
+import { and, eq, inArray, lt, notExists, sql } from 'drizzle-orm'
 import type {
   SnapshotService as ISnapshotService,
   SnapshotCreationResult,
@@ -92,24 +92,57 @@ export class SnapshotService implements ISnapshotService {
     const cutoffDate = new Date()
     cutoffDate.setDate(cutoffDate.getDate() - olderThanDays)
 
-    const deletedSnapshots = await db
-      .delete(workflowExecutionSnapshots)
-      .where(
-        and(
-          lt(workflowExecutionSnapshots.createdAt, cutoffDate),
-          notExists(
-            db
-              .select({ id: workflowExecutionLogs.id })
-              .from(workflowExecutionLogs)
-              .where(eq(workflowExecutionLogs.stateSnapshotId, workflowExecutionSnapshots.id))
+    const BATCH_SIZE = 1000
+    const MAX_BATCHES = 20
+
+    let totalDeleted = 0
+    let stoppedEarly = false
+
+    for (let batch = 0; batch < MAX_BATCHES; batch++) {
+      const candidates = await db
+        .select({ id: workflowExecutionSnapshots.id })
+        .from(workflowExecutionSnapshots)
+        .where(
+          and(
+            lt(workflowExecutionSnapshots.createdAt, cutoffDate),
+            notExists(
+              db
+                .select({ one: sql`1` })
+                .from(workflowExecutionLogs)
+                .where(eq(workflowExecutionLogs.stateSnapshotId, workflowExecutionSnapshots.id))
+            )
           )
         )
-      )
-      .returning({ id: workflowExecutionSnapshots.id })
+        .limit(BATCH_SIZE)
+
+      if (candidates.length === 0) break
+
+      const ids = candidates.map((c) => c.id)
+      const deleted = await db
+        .delete(workflowExecutionSnapshots)
+        .where(
+          and(
+            inArray(workflowExecutionSnapshots.id, ids),
+            notExists(
+              db
+                .select({ one: sql`1` })
+                .from(workflowExecutionLogs)
+                .where(eq(workflowExecutionLogs.stateSnapshotId, workflowExecutionSnapshots.id))
+            )
+          )
+        )
+        .returning({ id: workflowExecutionSnapshots.id })
 
-    const deletedCount = deletedSnapshots.length
-    logger.info(`Cleaned up ${deletedCount} orphaned snapshots older than ${olderThanDays} days`)
-    return deletedCount
+      totalDeleted += deleted.length
+
+      if (candidates.length < BATCH_SIZE) break
+      if (batch === MAX_BATCHES - 1) stoppedEarly = true
+    }
+
+    logger.info(
+      `Cleaned up ${totalDeleted} orphaned snapshots older than ${olderThanDays} days${stoppedEarly ? ' (batch cap reached, remainder deferred to next run)' : ''}`
+    )
+    return totalDeleted
   }
 }
 

packages/testing/src/mocks/database.mock.ts

@@ -49,6 +49,8 @@ export function createMockSqlOperators() {
     isNotNull: vi.fn((column) => ({ type: 'isNotNull', column })),
     inArray: vi.fn((column, values) => ({ type: 'inArray', column, values })),
     notInArray: vi.fn((column, values) => ({ type: 'notInArray', column, values })),
+    exists: vi.fn((subquery) => ({ type: 'exists', subquery })),
+    notExists: vi.fn((subquery) => ({ type: 'notExists', subquery })),
     like: vi.fn((column, pattern) => ({ type: 'like', column, pattern })),
     ilike: vi.fn((column, pattern) => ({ type: 'ilike', column, pattern })),
     desc: vi.fn((column) => ({ type: 'desc', column })),