diff --git a/src/backend/clients/dynamodb/DDBClient.test.ts b/src/backend/clients/dynamodb/DDBClient.test.ts index c9a672f3f..c5c560934 100644 --- a/src/backend/clients/dynamodb/DDBClient.test.ts +++ b/src/backend/clients/dynamodb/DDBClient.test.ts @@ -774,6 +774,142 @@ describe('DDBClient — unprocessed batch keys', () => { }); }); +// Answers by operation (X-Amz-Target) and records how many requests overlap, so +// chunk concurrency and the TTL sweep's retries can be asserted. +describe('DDBClient — batch write chunks and TTL sweep', () => { + let server: Server; + let endpoint: string; + let handlers: Record) => unknown>; + let calls: string[]; + let inFlight = 0; + let peakInFlight = 0; + + beforeAll(async () => { + server = createServer((req, res) => { + let body = ''; + req.on('data', (chunk) => { + body += chunk; + }); + req.on('end', () => { + const op = String(req.headers['x-amz-target']).split('.')[1]; + calls.push(op); + inFlight += 1; + peakInFlight = Math.max(peakInFlight, inFlight); + setTimeout(() => { + inFlight -= 1; + const reply = handlers[op]?.(JSON.parse(body)) ?? {}; + const { status = 200, ...payload } = reply as { + status?: number; + }; + res.writeHead(status, { + 'content-type': 'application/x-amz-json-1.0', + }); + res.end(JSON.stringify(payload)); + }, 20); + }); + }); + await new Promise((resolve) => + server.listen(0, '127.0.0.1', resolve), + ); + endpoint = `http://127.0.0.1:${(server.address() as AddressInfo).port}`; + }); + + afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); + }); + + beforeEach(() => { + calls = []; + inFlight = 0; + peakInFlight = 0; + }); + + const stubClient = (extra: Record = {}) => + new DDBClient({ + port: 0, + extensions: [], + dynamo: { + aws: { access_key: 'a', secret_key: 'b', region: 'us-west-2' }, + endpoint, + ...extra, + }, + } as unknown as IConfig); + + it('writes chunks concurrently, a bounded number at a time', async () => { + handlers = { BatchWriteItem: () => ({ UnprocessedItems: {} }) }; + + await stubClient().batchPut( + Array.from({ length: 200 }, (_, i) => ({ + table: TABLE, + item: { pk: `item-${i}` }, + })), + ); + + expect(calls).toHaveLength(8); + expect(peakInFlight).toBeGreaterThan(1); + expect(peakInFlight).toBeLessThanOrEqual(4); + }); + + it('stops starting chunks once one fails', async () => { + handlers = { + BatchWriteItem: () => ({ + status: 400, + __type: 'com.amazon.coral.validate#ValidationException', + message: 'bad item', + }), + }; + + await expect( + stubClient().batchPut( + Array.from({ length: 250 }, (_, i) => ({ + table: TABLE, + item: { pk: `item-${i}` }, + })), + ), + ).rejects.toThrow('bad item'); + // The first wave fails; none of the six chunks queued behind it run. + expect(calls.length).toBeLessThanOrEqual(4); + }); + + it('retries the unprocessed deletes of the TTL sweep', async () => { + const expired = [{ pk: { S: 'old' } }]; + let batchWrites = 0; + handlers = { + CreateTable: () => ({ + status: 400, + __type: 'com.amazonaws.dynamodb.v20120810#ResourceInUseException', + message: 'Table already exists', + }), + Scan: () => ({ Items: expired }), + BatchWriteItem: (body) => { + batchWrites += 1; + return batchWrites === 1 + ? { UnprocessedItems: body.RequestItems } + : { UnprocessedItems: {} }; + }, + }; + + await stubClient({ bootstrapTables: true }).createTableIfNotExists( + { + TableName: TABLE, + KeySchema: [{ AttributeName: 'pk', KeyType: 'HASH' }], + AttributeDefinitions: [ + { AttributeName: 'pk', AttributeType: 'S' }, + ], + BillingMode: 'PAY_PER_REQUEST', + }, + 'ttl', + ); + + expect(calls).toEqual([ + 'CreateTable', + 'Scan', + 'BatchWriteItem', + 'BatchWriteItem', + ]); + }); +}); + // Throttles are a service-side signal the emulator never produces, so these // drive the client against a stub answering in DynamoDB's error shape. describe('DDBClient — throttle logging', () => { diff --git a/src/backend/clients/dynamodb/DDBClient.ts b/src/backend/clients/dynamodb/DDBClient.ts index d9463fe13..560ad761b 100644 --- a/src/backend/clients/dynamodb/DDBClient.ts +++ b/src/backend/clients/dynamodb/DDBClient.ts @@ -44,6 +44,7 @@ import { Agent as httpsAgent } from 'node:https'; import { HttpError } from '../../core/http'; import type { IConfig, IDynamoConfig } from '../../types'; import { BoundedTtlMap } from '../../util/boundedTtlMap.js'; +import { runWithConcurrencyLimit } from '../../util/concurrency.js'; import { Span } from '../../util/span.js'; import { PuterClient } from '../types'; import { @@ -56,10 +57,10 @@ import { attachThrottleLogging } from './throttleLog.js'; const LOCAL_DYNAMO_MEMORY_PREFIX = ':memory:'; const localDynaliteEndpointPromises = new Map>(); const MAX_BATCH_WRITE_ITEMS = 25; -const MAX_BATCH_WRITE_RETRIES = 8; -const BATCH_WRITE_RETRY_BASE_MS = 25; -const MAX_BATCH_GET_RETRIES = 8; -const BATCH_GET_RETRY_BASE_MS = 25; +// Chunks of one batch write in flight at once. +const BATCH_WRITE_CONCURRENCY = 4; +const MAX_BATCH_RETRIES = 8; +const BATCH_RETRY_BASE_MS = 25; // In-memory mode gives each DDBClient its own unique key, which keys // into a fresh dynalite server. This is what test parallelism needs: @@ -119,6 +120,32 @@ const chunkValues = (values: T[], size: number): T[][] => { return chunks; }; +type BatchWriteItems = NonNullable; +type WriteRequest = BatchWriteItems[string][number]; + +/** Consumed capacity summed per table across the requests of one batch. */ +class ConsumedCapacityTally { + #byTable = new Map(); + + add(entries: { TableName?: string; CapacityUnits?: number }[] = []) { + for (const { TableName, CapacityUnits } of entries) { + if (!TableName) continue; + this.#byTable.set( + TableName, + (this.#byTable.get(TableName) ?? 0) + + Number(CapacityUnits ?? 0), + ); + } + } + + entries() { + return Array.from(this.#byTable, ([TableName, CapacityUnits]) => ({ + TableName, + CapacityUnits, + })); + } +} + const sleep = async (ms: number) => { await new Promise((resolve) => setTimeout(resolve, ms)); }; @@ -256,9 +283,7 @@ export class DDBClient extends PuterClient { {} as Record[]>, ); - let requestItems: BatchGetCommandInput['RequestItems'] = Object.entries( - allRequestItemsPerTable, - ).reduce( + const requestItems = Object.entries(allRequestItemsPerTable).reduce( (acc, [table, keyList]) => { acc[table] = { Keys: keyList, @@ -271,53 +296,30 @@ export class DDBClient extends PuterClient { const client = await this.#getDocumentClient(); const responsesByTable = new Map[]>(); - const consumedCapacityByTable = new Map(); + const capacity = new ConsumedCapacityTally(); - for (let attempt = 0; attempt <= MAX_BATCH_GET_RETRIES; attempt++) { - if (Object.keys(requestItems).length === 0) break; - - const response = await client.send( - new BatchGetCommand({ - RequestItems: requestItems, - ReturnConsumedCapacity: 'TOTAL', - }), - ); - - for (const [table, items] of Object.entries( - response.Responses ?? {}, - )) { - const existing = responsesByTable.get(table) ?? []; - existing.push(...items); - responsesByTable.set(table, existing); - } - for (const entry of response.ConsumedCapacity ?? []) { - if (!entry.TableName) continue; - consumedCapacityByTable.set( - entry.TableName, - (consumedCapacityByTable.get(entry.TableName) ?? 0) + - Number(entry.CapacityUnits ?? 0), + const drained = await this.#sendWithUnprocessedRetry( + requestItems, + async (pending) => { + const response = await client.send( + new BatchGetCommand({ + RequestItems: pending, + ReturnConsumedCapacity: 'TOTAL', + }), ); - } + for (const [table, items] of Object.entries( + response.Responses ?? {}, + )) { + const existing = responsesByTable.get(table) ?? []; + existing.push(...items); + responsesByTable.set(table, existing); + } + capacity.add(response.ConsumedCapacity); + return response.UnprocessedKeys as typeof pending | undefined; + }, + ); - const unprocessedKeys = response.UnprocessedKeys ?? {}; - if (Object.keys(unprocessedKeys).length === 0) { - requestItems = {}; - break; - } - - requestItems = unprocessedKeys as NonNullable< - BatchGetCommandInput['RequestItems'] - >; - if (attempt < MAX_BATCH_GET_RETRIES) { - const delayMs = Math.min( - 1000, - BATCH_GET_RETRY_BASE_MS * 2 ** attempt, - ); - await sleep(delayMs); - } - } - - if (Object.keys(requestItems).length > 0) { + if (!drained) { throw new HttpError( 400, 'Failed to batch get all items from DynamoDB', @@ -327,12 +329,7 @@ export class DDBClient extends PuterClient { return { Responses: Object.fromEntries(responsesByTable), - ConsumedCapacity: Array.from(consumedCapacityByTable.entries()).map( - ([TableName, CapacityUnits]) => ({ - TableName, - CapacityUnits, - }), - ), + ConsumedCapacity: capacity.entries(), }; } @@ -360,122 +357,89 @@ export class DDBClient extends PuterClient { ); } - // Shared BatchWriteItem plumbing for batchPut/batchDel: 25-item chunks, - // UnprocessedItems retried with capped exponential backoff, consumed - // capacity accumulated per table across every request. - async #batchWrite( - params: { - table: string; - request: NonNullable< - BatchWriteCommandInput['RequestItems'] - >[string][number]; - }[], - ) { - const consumedCapacityByTable = new Map(); + // Shared BatchWriteItem plumbing for batchPut/batchDel and the TTL sweep: + // 25-item chunks a few at a time, UnprocessedItems retried, consumed + // capacity summed per table across every request. + async #batchWrite(params: { table: string; request: WriteRequest }[]) { if (params.length === 0) { return { ConsumedCapacity: [] }; } - const accumulateConsumedCapacity = ( - consumedCapacityEntries: - | Array<{ TableName?: string; CapacityUnits?: number }> - | undefined, - ) => { - if (!consumedCapacityEntries) { - return; - } - - for (const consumedCapacityEntry of consumedCapacityEntries) { - const table = consumedCapacityEntry.TableName; - if (!table) { - continue; - } - - const existingUsage = consumedCapacityByTable.get(table) ?? 0; - consumedCapacityByTable.set( - table, - existingUsage + - Number(consumedCapacityEntry.CapacityUnits ?? 0), - ); - } - }; - const client = await this.#getDocumentClient(); - const chunks = chunkValues(params, MAX_BATCH_WRITE_ITEMS); + const capacity = new ConsumedCapacityTally(); + // Once a chunk fails, the rest aren't started. + let failed = false; - for (const chunk of chunks) { - let requestItems = chunk.reduce( - (acc, curr) => { - const tableRequests = acc[curr.table] ?? []; - tableRequests.push(curr.request); - acc[curr.table] = tableRequests; - return acc; - }, - {} as NonNullable, - ); - - for ( - let attempt = 0; - attempt <= MAX_BATCH_WRITE_RETRIES; - attempt++ - ) { - if (Object.keys(requestItems).length === 0) { - break; + await runWithConcurrencyLimit( + chunkValues(params, MAX_BATCH_WRITE_ITEMS), + BATCH_WRITE_CONCURRENCY, + async (chunk) => { + if (failed) return; + const requestItems: BatchWriteItems = {}; + for (const { table, request } of chunk) { + (requestItems[table] ??= []).push(request); } - - const response = await sendRepairingValues( - requestItems, - (itemsToWrite) => - client.send( - new BatchWriteCommand({ - RequestItems: itemsToWrite, - ReturnConsumedCapacity: 'TOTAL', - }), - ), - () => - `a batch write to ${Object.keys(requestItems).join(', ')}`, - ); - accumulateConsumedCapacity( - response.ConsumedCapacity as - | Array<{ TableName?: string; CapacityUnits?: number }> - | undefined, - ); - - const unprocessedItems = response.UnprocessedItems ?? {}; - if (Object.keys(unprocessedItems).length === 0) { - requestItems = {}; - break; - } - - requestItems = unprocessedItems as NonNullable< - BatchWriteCommandInput['RequestItems'] - >; - if (attempt < MAX_BATCH_WRITE_RETRIES) { - const delayMs = Math.min( - 1000, - BATCH_WRITE_RETRY_BASE_MS * 2 ** attempt, + try { + const drained = await this.#sendWithUnprocessedRetry( + requestItems, + async (pending) => { + const response = await sendRepairingValues( + pending, + (itemsToWrite) => + client.send( + new BatchWriteCommand({ + RequestItems: itemsToWrite, + ReturnConsumedCapacity: 'TOTAL', + }), + ), + () => + `a batch write to ${Object.keys(pending).join(', ')}`, + ); + capacity.add(response.ConsumedCapacity); + return response.UnprocessedItems as + | BatchWriteItems + | undefined; + }, ); - await sleep(delayMs); + if (!drained) { + throw new HttpError( + 400, + 'Failed to batch write all items to DynamoDB', + { legacyCode: 'bad_request' }, + ); + } + } catch (error) { + failed = true; + throw error; } - } + }, + ); - if (Object.keys(requestItems).length > 0) { - throw new HttpError( - 400, - 'Failed to batch write all items to DynamoDB', - { legacyCode: 'bad_request' }, - ); + return { ConsumedCapacity: capacity.entries() }; + } + + /** + * Send a batch request and resend whatever comes back unprocessed, with + * capped exponential backoff. `send` returns the unprocessed remainder. + * Resolves false when items are still left after the last retry. + */ + async #sendWithUnprocessedRetry( + requestItems: TItems, + send: (pending: TItems) => Promise, + ): Promise { + let pending: TItems | undefined = requestItems; + for (let attempt = 0; attempt <= MAX_BATCH_RETRIES; attempt++) { + if (!pending || Object.keys(pending).length === 0) return true; + pending = await send(pending); + if ( + attempt < MAX_BATCH_RETRIES && + pending && + Object.keys(pending).length > 0 + ) { + await sleep(Math.min(1000, BATCH_RETRY_BASE_MS * 2 ** attempt)); } } - - return { - ConsumedCapacity: Array.from(consumedCapacityByTable.entries()).map( - ([TableName, CapacityUnits]) => ({ - TableName, - CapacityUnits, - }), - ), - }; + return !pending || Object.keys(pending).length === 0; } @Span('ddb.del', (table: string) => ({ 'db.table': table })) @@ -727,34 +691,10 @@ export class DDBClient extends PuterClient { ); } - const region = this.#ddbConfig.aws?.region || 'us-west-2'; - const ddbClient = new DynamoDBClient({ - credentials: { - accessKeyId, - secretAccessKey, - }, - maxAttempts: 3, - requestHandler: new NodeHttpHandler({ - connectionTimeout: 5000, - requestTimeout: 5000, - httpsAgent: new httpsAgent({ keepAlive: true }), - }), - ...(this.#ddbConfig.endpoint - ? { endpoint: this.#ddbConfig.endpoint } - : {}), - region, - }); - attachThrottleLogging(ddbClient, region, (table) => - this.#keyNamesByTable.get(table), - ); - - this.#documentClient = DynamoDBDocumentClient.from(ddbClient, { - marshallOptions: { - removeUndefinedValues: true, - }, - unmarshallOptions: { - wrapNumbers: clampStoredNumber, - }, + this.#bindClient({ + credentials: { accessKeyId, secretAccessKey }, + region: this.#ddbConfig.aws?.region || 'us-west-2', + endpoint: this.#ddbConfig.endpoint, }); } @@ -762,21 +702,36 @@ export class DDBClient extends PuterClient { const endpoint = await getOrCreateLocalDynaliteEndpoint( this.#localPathKey, ); + this.#bindClient({ + credentials: { accessKeyId: 'fake', secretAccessKey: 'fake' }, + region: 'us-west-2', + endpoint, + }); + } + #bindClient({ + credentials, + region, + endpoint, + }: { + credentials: { accessKeyId: string; secretAccessKey: string }; + region: string; + endpoint?: string; + }) { const ddbClient = new DynamoDBClient({ - credentials: { - accessKeyId: 'fake', - secretAccessKey: 'fake', - }, + credentials, maxAttempts: 3, requestHandler: new NodeHttpHandler({ connectionTimeout: 5000, requestTimeout: 5000, httpsAgent: new httpsAgent({ keepAlive: true }), }), - endpoint, - region: 'us-west-2', + ...(endpoint ? { endpoint } : {}), + region, }); + attachThrottleLogging(ddbClient, region, (table) => + this.#keyNamesByTable.get(table), + ); this.#documentClient = DynamoDBDocumentClient.from(ddbClient, { marshallOptions: { @@ -822,26 +777,23 @@ export class DDBClient extends PuterClient { ); lastEvaluatedKey = scan.LastEvaluatedKey as - Record | undefined; + | Record + | undefined; const items = scan.Items; if (!items || items.length === 0) continue; - const chunks = chunkValues(items, MAX_BATCH_WRITE_ITEMS); - for (const chunk of chunks) { - await client.send( - new BatchWriteCommand({ - RequestItems: { - [table]: chunk.map((item) => ({ - DeleteRequest: { - Key: Object.fromEntries( - keyNames.map((key) => [key, item[key]]), - ), - }, - })), + await this.#batchWrite( + items.map((item) => ({ + table, + request: { + DeleteRequest: { + Key: Object.fromEntries( + keyNames.map((key) => [key, item[key]]), + ), }, - }), - ); - } + }, + })), + ); } while (lastEvaluatedKey); } }