fix(account-deletion): resume the workspace storage sweep from its saved listing position

apps/sim/lib/uploads/core/prefix-sweep.ts

@@ -94,10 +94,13 @@ async function localKeys(
 }
 
 /** One provider page of object keys; `done` once the listing has nothing after this page. */
-interface StorageKeyPage {
+export interface StorageKeyPage {
   keys: string[]
   done: boolean
-  /** Provider cursor for the following page, valid only within the current run. */
+  /**
+   * Opaque provider cursor for the following page. A caller may persist it across runs, but a
+   * stored one can stop listing (a provider expiring or rejecting it), which that caller handles.
+   */
   nextPageToken?: string
 }
 
@@ -170,6 +173,8 @@ export async function listStorageKeyPage(
         : undefined
     )
   }
+  if (pageToken && !pageToken.startsWith(prefix))
+    throw new Error('Local storage page token is not a key under the prefix')
   const keys = await localKeys(prefix, signal, pageToken)
   return keyPage(keys, keys.length < STORAGE_KEY_PAGE_SIZE ? undefined : keys[keys.length - 1])
 }

apps/sim/lib/users/storage-sweep.test.ts

@@ -16,9 +16,14 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 vi.mock('@/lib/uploads/core/setup.server', () => uploadsSetupMock)
 vi.mock('@/lib/uploads/config', () => uploadsConfigMock)
 vi.mock('@/lib/uploads/server/metadata', () => uploadsMetadataMock)
+vi.mock('@/lib/uploads/core/prefix-sweep', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@/lib/uploads/core/prefix-sweep')>()
+  return { ...actual, listStorageKeyPage: vi.fn(actual.listStorageKeyPage) }
+})
 
 import type { OutboxEventContext } from '@/lib/core/outbox/service'
 import type { DbTransaction } from '@/lib/db/types'
+import { listStorageKeyPage } from '@/lib/uploads/core/prefix-sweep'
 import { createStorageReferenceGuard } from '@/lib/users/storage-reference-guard'
 import {
   accountStorageSweepOutboxHandlers,
@@ -47,11 +52,59 @@ function outboxContext(overrides: Partial<OutboxEventContext> = {}): OutboxEvent
   }
 }
 
+/** A context whose deadline is near once it has listed `listings` storage pages. */
+function deadlineAfter(listings: number, overrides: Partial<OutboxEventContext> = {}) {
+  const context = outboxContext(overrides)
+  const listed = vi.mocked(listStorageKeyPage).mock.calls.length
+  Object.defineProperty(context, 'deadlineAt', {
+    get: () =>
+      vi.mocked(listStorageKeyPage).mock.calls.length - listed >= listings
+        ? 1
+        : Number.POSITIVE_INFINITY,
+  })
+  return context
+}
+
+/** Re-runs the sweep the way the outbox does, merging each checkpoint into the next payload. */
+async function sweepUntilSettled(payload: Record<string, unknown>, listingsPerRun: number) {
+  let current = payload
+  const outcomes: Awaited<ReturnType<typeof sweep>>[] = []
+  for (let run = 0; run < 12; run++) {
+    const outcome = await sweep(
+      current,
+      deadlineAfter(listingsPerRun, {
+        checkpointPayload: vi.fn(async (patch: Record<string, unknown>) => {
+          current = { ...current, ...patch }
+        }),
+      })
+    )
+    outcomes.push(outcome)
+    if (outcome === undefined) break
+  }
+  return outcomes
+}
+
 async function put(key: string, body = 'x'): Promise<void> {
   await mkdir(dirname(join(root, key)), { recursive: true })
   await writeFile(join(root, key), body)
 }
 
+/** Writes `count` objects a surviving document names, answering `lookups` reference checks. */
+async function putKept(count: number, lookups = 1): Promise<string[]> {
+  const width = String(count - 1).length
+  const keys = Array.from(
+    { length: count },
+    (_, index) => `workspace/wsdoomed1/a-kept/${String(index).padStart(width, '0')}.bin`
+  )
+  for (const key of keys) await put(key)
+  for (let lookup = 0; lookup < lookups; lookup++)
+    queueTableRows(
+      schemaMock.document,
+      keys.map((key) => ({ key }))
+    )
+  return keys
+}
+
 async function exists(key: string): Promise<boolean> {
   try {
     await stat(join(root, key))
@@ -249,12 +302,132 @@ describe('deleted workspace storage sweep', () => {
     for (const key of keptKeys) expect(await exists(key)).toBe(true)
   })
 
+  it('reaches objects behind more kept pages than one run can walk, spending no attempt', async () => {
+    const keptKeys = await putKept(1000, 30)
+    await put('workspace/wsdoomed1/z-unkept/0.bin')
+
+    const outcomes = await sweepUntilSettled({ version: 1, workspaceId: 'wsdoomed1' }, 2)
+
+    expect(outcomes.at(-1)).toBeUndefined()
+    for (const outcome of outcomes.slice(0, -1))
+      expect(outcome).toMatchObject({ consumeAttempt: false })
+    expect(await exists('workspace/wsdoomed1/z-unkept/0.bin')).toBe(false)
+    for (const key of keptKeys) expect(await exists(key)).toBe(true)
+  })
+
+  it('resumes a family from its saved listing position', async () => {
+    const checkpointPayload = vi.fn(async () => {})
+
+    await expect(
+      sweep(
+        { version: 1, workspaceId: 'wsdoomed1', family: 0, pageToken: 'workspace/wsdoomed1/a.txt' },
+        outboxContext({ checkpointPayload })
+      )
+    ).resolves.toBe(undefined)
+
+    expect(await exists('workspace/wsdoomed1/a.txt')).toBe(true)
+    expect(await exists('workspace/wsdoomed1/nested/b.txt')).toBe(false)
+    expect(await exists('table-import/wsdoomed1/import-1/rows.csv')).toBe(false)
+    expect(checkpointPayload).toHaveBeenCalledWith({ family: 1, pageToken: null })
+  })
+
+  it('retries a saved position that failed to list once instead of walking from the start', async () => {
+    let payload: Record<string, unknown> = {
+      version: 1,
+      workspaceId: 'wsdoomed1',
+      pageToken: 'workspace/wsdoomed1/a.txt',
+    }
+    const checkpointPayload = vi.fn(async (patch: Record<string, unknown>) => {
+      payload = { ...payload, ...patch }
+    })
+    vi.mocked(listStorageKeyPage).mockRejectedValueOnce(new Error('Service unavailable'))
+
+    await expect(sweep(payload, outboxContext({ checkpointPayload }))).rejects.toThrow(
+      'Service unavailable'
+    )
+    for (const key of DOOMED_KEYS) expect(await exists(key)).toBe(true)
+    expect(payload).toMatchObject({ pageToken: 'workspace/wsdoomed1/a.txt', tokenFailures: 1 })
+
+    await expect(sweep(payload, outboxContext({ checkpointPayload }))).resolves.toBe(undefined)
+    expect(await exists('workspace/wsdoomed1/a.txt')).toBe(true)
+    expect(await exists('workspace/wsdoomed1/nested/b.txt')).toBe(false)
+    expect(payload).toMatchObject({ tokenFailures: 0 })
+  })
+
+  it.each([
+    { name: 'no longer lists', pageToken: 'expired-token', rejects: true },
+    { name: 'is not a key under the prefix', pageToken: 'zz-other-provider', rejects: false },
+  ])(
+    'walks the family from the start once its saved position $name on a second attempt',
+    async ({ pageToken, rejects }) => {
+      let payload: Record<string, unknown> = { version: 1, workspaceId: 'wsdoomed1', pageToken }
+      const checkpointPayload = vi.fn(async (patch: Record<string, unknown>) => {
+        payload = { ...payload, ...patch }
+      })
+      const invalid = new Error('The continuation token is invalid')
+      if (rejects)
+        vi.mocked(listStorageKeyPage).mockRejectedValueOnce(invalid).mockRejectedValueOnce(invalid)
+
+      await expect(sweep(payload, outboxContext({ checkpointPayload }))).rejects.toThrow()
+      for (const key of DOOMED_KEYS) expect(await exists(key)).toBe(true)
+
+      await expect(sweep(payload, outboxContext({ checkpointPayload }))).resolves.toBe(undefined)
+      for (const key of DOOMED_KEYS) expect(await exists(key)).toBe(false)
+    }
+  )
+
+  it('spends an attempt when a restart from a rejected position only walks kept objects', async () => {
+    const keptKeys = await putKept(500)
+    vi.mocked(listStorageKeyPage).mockRejectedValueOnce(
+      new Error('The continuation token is invalid')
+    )
+    const checkpointPayload = vi.fn(async () => {})
+
+    const outcome = await sweep(
+      { version: 1, workspaceId: 'wsdoomed1', pageToken: 'expired-token', tokenFailures: 1 },
+      deadlineAfter(2, { checkpointPayload })
+    )
+
+    expect(outcome).toMatchObject({ outcome: 'deferred' })
+    expect(outcome).not.toHaveProperty('consumeAttempt')
+    expect(checkpointPayload).toHaveBeenCalledWith({
+      family: 0,
+      pageToken: keptKeys.at(-1),
+      resumes: 0,
+    })
+  })
+
+  it('spends an attempt on a moved position once its resume budget is gone', async () => {
+    const keptKeys = await putKept(1000)
+    const checkpointPayload = vi.fn(async () => {})
+
+    const outcome = await sweep(
+      { version: 1, workspaceId: 'wsdoomed1', resumes: 1000 },
+      deadlineAfter(1, { checkpointPayload })
+    )
+
+    expect(outcome).toMatchObject({ outcome: 'deferred' })
+    expect(outcome).not.toHaveProperty('consumeAttempt')
+    expect(checkpointPayload).toHaveBeenCalledWith({
+      family: 0,
+      pageToken: keptKeys[499],
+      resumes: 1001,
+    })
+  })
+
   it.each([
     { version: 1, workspaceId: '' },
     { version: 1, workspaceId: '../wsdoomed1' },
     { version: 1, workspaceId: 'wsdoomed1/' },
     { version: 1, workspaceId: '*' },
     { workspaceId: 'wsdoomed1' },
+    { version: 1, workspaceId: 'wsdoomed1', pageToken: 42 },
+    { version: 1, workspaceId: 'wsdoomed1', pageToken: '' },
+    { version: 1, workspaceId: 'wsdoomed1', pageToken: 'x'.repeat(4097) },
+    { version: 1, workspaceId: 'wsdoomed1', resumes: -1 },
+    { version: 1, workspaceId: 'wsdoomed1', resumes: '1' },
+    { version: 1, workspaceId: 'wsdoomed1', tokenFailures: 2 },
+    { version: 1, workspaceId: 'wsdoomed1', tokenFailures: '1' },
   ])('rejects an unbounded or malformed payload %o', async (payload) => {
     await expect(sweep(payload, outboxContext())).rejects.toThrow()
     for (const key of DOOMED_KEYS) expect(await exists(key)).toBe(true)