diff --git a/apps/sim/lib/uploads/core/prefix-sweep.ts b/apps/sim/lib/uploads/core/prefix-sweep.ts index 40a183ee3dd..7e7719f422c 100644 --- a/apps/sim/lib/uploads/core/prefix-sweep.ts +++ b/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]) } diff --git a/apps/sim/lib/users/storage-sweep.test.ts b/apps/sim/lib/users/storage-sweep.test.ts index c9cb20e1d88..6d702a43c01 100644 --- a/apps/sim/lib/users/storage-sweep.test.ts +++ b/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() + 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 = {}): OutboxEvent } } +/** A context whose deadline is near once it has listed `listings` storage pages. */ +function deadlineAfter(listings: number, overrides: Partial = {}) { + 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, listingsPerRun: number) { + let current = payload + const outcomes: Awaited>[] = [] + for (let run = 0; run < 12; run++) { + const outcome = await sweep( + current, + deadlineAfter(listingsPerRun, { + checkpointPayload: vi.fn(async (patch: Record) => { + current = { ...current, ...patch } + }), + }) + ) + outcomes.push(outcome) + if (outcome === undefined) break + } + return outcomes +} + async function put(key: string, body = 'x'): Promise { 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 { + 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 { 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 = { + version: 1, + workspaceId: 'wsdoomed1', + pageToken: 'workspace/wsdoomed1/a.txt', + } + const checkpointPayload = vi.fn(async (patch: Record) => { + 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 = { version: 1, workspaceId: 'wsdoomed1', pageToken } + const checkpointPayload = vi.fn(async (patch: Record) => { + 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) diff --git a/apps/sim/lib/users/storage-sweep.ts b/apps/sim/lib/users/storage-sweep.ts index ef649fda616..04b57f8aa80 100644 --- a/apps/sim/lib/users/storage-sweep.ts +++ b/apps/sim/lib/users/storage-sweep.ts @@ -42,6 +42,7 @@ import { deleteStorageKeysUnderPrefix, listStorageKeyPage, outboxDeadlineNear, + type StorageKeyPage, storageCleanupAvailableAt, } from '@/lib/uploads/core/prefix-sweep' import { LOCAL_UPLOAD_METADATA_SUFFIX } from '@/lib/uploads/core/storage-key' @@ -58,6 +59,14 @@ const logger = createLogger('AccountStorageSweep') const ROW_PAGE_SIZE = 1000 const OBJECTS_PER_EVENT = 100 +const MAX_PAGE_TOKEN_LENGTH = 4096 + +/** + * Runs one workspace sweep may continue on a moved listing position alone. Provider cursors are + * opaque, so nothing here proves one never revisits a position; past this, a run that only moved + * its position spends an attempt, keeping the sweep bounded even by a cursor that cycles. + */ +const MAX_POSITION_RESUMES = 1000 /** * Ceiling on exact keys one deletion records for durable cleanup. Past it the rest are logged @@ -97,6 +106,15 @@ interface StorageObject { interface WorkspaceSweepPayload { workspaceId: string family: number + /** Listing position inside `family` a stopped run saved; absent from events queued before it. */ + pageToken?: string + /** Runs continued on a moved listing position alone; see {@link MAX_POSITION_RESUMES}. */ + resumes: number + /** + * Set once listing from `pageToken` failed. A second failure from the same token restarts the + * family; a single one is retried in place, since providers do not reliably mark a dead cursor. + */ + tokenFailed: boolean } interface ObjectCleanupPayload { @@ -123,7 +141,27 @@ function parseWorkspaceSweep(value: unknown): WorkspaceSweepPayload { family > WORKSPACE_STORAGE_FAMILIES.length ) throw new Error('Invalid workspace storage sweep cursor') - return { workspaceId: value.workspaceId, family } + const pageToken = value.pageToken ?? undefined + if ( + pageToken !== undefined && + (typeof pageToken !== 'string' || + pageToken.length === 0 || + pageToken.length > MAX_PAGE_TOKEN_LENGTH) + ) + throw new Error('Invalid workspace storage sweep page token') + const resumes = value.resumes ?? 0 + if (typeof resumes !== 'number' || !Number.isSafeInteger(resumes) || resumes < 0) + throw new Error('Invalid workspace storage sweep resume count') + const tokenFailures = value.tokenFailures ?? 0 + if (tokenFailures !== 0 && tokenFailures !== 1) + throw new Error('Invalid workspace storage sweep page token failure count') + return { + workspaceId: value.workspaceId, + family, + resumes, + tokenFailed: tokenFailures === 1, + ...(pageToken ? { pageToken } : {}), + } } function isSafeKey(key: string): boolean { @@ -205,24 +243,60 @@ function sidecarOwner(key: string): string | null { : null } +interface PrefixSweepResult { + deleted: number + status: 'done' | 'partial' | 'stalled' | 'refused' + /** On `partial`, the listing position the next run resumes from; absent means the start. */ + pageToken?: string + /** The saved position could not be listed from, so this run walked from the start instead. */ + restarted: boolean +} + /** - * Walks one prefix to its end, deleting every listed object the reference guard does not keep, - * following the listing's page token within a run so kept objects never hold back later ones. A + * Walks one prefix to its end from `startToken`, deleting every listed object the reference guard + * does not keep, following the listing's page token so kept objects never hold back later ones. A * page that neither advances the token nor removes anything is `stalled` rather than re-listed - * forever. + * forever. A saved token is provider-specific and may no longer list (an expired or rejected + * cursor). The first failed listing from one is recorded and rethrown, so the outbox retries from + * the same position; when that token has failed before, the walk restarts from the prefix start. */ async function sweepWorkspacePrefix( prefix: string, storageContext: StorageContext, context: OutboxEventContext, isReferenced: (keys: readonly string[]) => Promise>, - workspaceExists: () => Promise -): Promise<{ deleted: number; status: 'done' | 'partial' | 'stalled' | 'refused' }> { + workspaceExists: () => Promise, + startToken: string | undefined, + startTokenFailed: boolean +): Promise { let deleted = 0 - let pageToken: string | undefined + let pageToken = startToken + let restarted = false + let failureRecorded = startTokenFailed for (;;) { - if (outboxDeadlineNear(context)) return { deleted, status: 'partial' } - const page = await listStorageKeyPage(prefix, storageContext, context.signal, pageToken) + if (outboxDeadlineNear(context)) return { deleted, status: 'partial', pageToken, restarted } + let page: StorageKeyPage + try { + page = await listStorageKeyPage(prefix, storageContext, context.signal, pageToken) + } catch (error) { + if (pageToken === undefined || pageToken !== startToken || restarted) throw error + context.signal.throwIfAborted() + if (!failureRecorded) { + await context.checkpointPayload({ tokenFailures: 1 }) + throw error + } + logger.warn('Workspace storage sweep could not resume its saved listing position', { + prefix, + error: describeError(error), + }) + restarted = true + pageToken = undefined + continue + } + if (failureRecorded) { + await context.checkpointPayload({ tokenFailures: 0 }) + failureRecorded = false + } if (page.keys.some((key) => !key.startsWith(prefix))) throw new Error('Storage provider returned a key outside the workspace prefix') const owners = page.keys.flatMap((key) => sidecarOwner(key) ?? []) @@ -231,7 +305,7 @@ async function sweepWorkspacePrefix( const keptCount = page.keys.filter(kept).length if (keptCount > 0) logger.warn('Keeping objects a surviving reference still names', { prefix, kept: keptCount }) - if (await workspaceExists()) return { deleted, status: 'refused' } + if (await workspaceExists()) return { deleted, status: 'refused', restarted } const result = await deleteStorageKeysUnderPrefix( prefix, storageContext, @@ -239,14 +313,14 @@ async function sweepWorkspacePrefix( context ) deleted += result.deleted - if (!result.complete) return { deleted, status: 'partial' } - if (page.done) return { deleted, status: 'done' } + if (!result.complete) return { deleted, status: 'partial', pageToken, restarted } + if (page.done) return { deleted, status: 'done', restarted } const advanced = page.nextPageToken !== undefined && page.nextPageToken !== pageToken if (advanced) { pageToken = page.nextPageToken } else if (result.deleted === 0) { logger.error('Workspace storage sweep made no progress on a full listing page', { prefix }) - return { deleted, status: 'stalled' } + return { deleted, status: 'stalled', restarted } } } } @@ -255,7 +329,10 @@ async function sweepWorkspacePrefix( * Deletes a deleted workspace's objects by prefix in each bucket. A workspace row with the same * id means the database no longer agrees the workspace is gone (a restore, a point-in-time * recovery), so the sweep stops without touching storage; it looks again before every page it - * deletes from. + * deletes from. A run stopped by its deadline saves its listing position, so a family whose early + * pages are all kept still reaches the objects behind them; moving that position counts as + * progress, except after a restart from the start or past {@link MAX_POSITION_RESUMES}, so a + * cursor that never lists or never ends stays bounded by the attempt budget. */ const sweepWorkspaceStorage: OutboxHandler = async (raw, context) => { const payload = parseWorkspaceSweep(raw) @@ -277,22 +354,32 @@ const sweepWorkspaceStorage: OutboxHandler = async (raw, context) => { for (let family = payload.family; family < WORKSPACE_STORAGE_FAMILIES.length; family++) { const { root, context: storageContext } = WORKSPACE_STORAGE_FAMILIES[family] const prefix = workspaceStoragePrefix(root, payload.workspaceId) + const startToken = family === payload.family ? payload.pageToken : undefined const result = await sweepWorkspacePrefix( prefix, storageContext, context, isReferenced, - workspaceExists + workspaceExists, + startToken, + family === payload.family && payload.tokenFailed ) if (result.status === 'refused') return if (result.deleted > 0) progressed = true if (result.status === 'stalled') return deferOutboxHandler('Workspace storage sweep made no progress on a full listing page') - if (result.status === 'partial') + if (result.status === 'partial') { + if (result.pageToken !== startToken) { + const movedOnly = !progressed && !result.restarted + const resumes = payload.resumes + (movedOnly ? 1 : 0) + await context.checkpointPayload({ family, pageToken: result.pageToken ?? null, resumes }) + if (movedOnly && resumes <= MAX_POSITION_RESUMES) progressed = true + } return progressed ? continueOutboxHandler('Continuing workspace storage sweep') : deferOutboxHandler('Workspace storage sweep made no progress before its deadline') - await context.checkpointPayload({ family: family + 1 }) + } + await context.checkpointPayload({ family: family + 1, pageToken: null }) progressed = true } }