diff --git a/.changeset/atomic-counter-primitive.md b/.changeset/atomic-counter-primitive.md new file mode 100644 index 000000000..c49ca0095 --- /dev/null +++ b/.changeset/atomic-counter-primitive.md @@ -0,0 +1,5 @@ +--- +"@aws-blocks/bb-distributed-data": minor +--- + +Add an atomic `Counter` primitive to `DistributedDatabase`. `db.counter(name).next()` returns a race-free monotonic sequence value (with `.current()` and `.reset()`), backed by a single `INSERT ... ON CONFLICT DO UPDATE ... RETURNING` upsert against a framework-managed `_blocks_counters` table (created automatically on deploy). This replaces the racy `SELECT MAX(seq) + 1` read-modify-write pattern that DSQL otherwise forces, since it has no sequences (`SERIAL` / `BIGSERIAL`). diff --git a/packages/bb-distributed-data/API.md b/packages/bb-distributed-data/API.md index bfea8d3d7..c9268e961 100644 --- a/packages/bb-distributed-data/API.md +++ b/packages/bb-distributed-data/API.md @@ -6,6 +6,7 @@ import type { ChildLogger } from '@aws-blocks/bb-logger'; import { createKyselyAdapter } from '@aws-blocks/data-common'; +import type { DatabaseBase } from '@aws-blocks/data-common'; import { DatabaseEngine } from '@aws-blocks/data-common'; import { Scope } from '@aws-blocks/core'; import type { ScopeParent } from '@aws-blocks/core'; @@ -13,11 +14,21 @@ import { sql } from '@aws-blocks/data-common'; import { SqlQuery } from '@aws-blocks/data-common'; import { Transaction } from '@aws-blocks/data-common'; +// @public +export class Counter { + // @internal + constructor(resolveBase: () => Promise, name: string); + current(): Promise; + next(delta?: number): Promise; + reset(value?: number): Promise; +} + export { createKyselyAdapter } // @public (undocumented) export class DistributedDatabase extends Scope { constructor(scope: ScopeParent, id: string, _options?: DistributedDatabaseOptions); + counter(name: string): Counter; // (undocumented) execute(query: SqlQuery): Promise<{ rowCount: number; diff --git a/packages/bb-distributed-data/DESIGN.md b/packages/bb-distributed-data/DESIGN.md index 330cb8bb6..0425f77a4 100644 --- a/packages/bb-distributed-data/DESIGN.md +++ b/packages/bb-distributed-data/DESIGN.md @@ -120,6 +120,28 @@ async transaction(fn, options?) { Default: no retry (honest, predictable). `retryOnConflict: true` is explicit opt-in with JSDoc warning about side effects. +## Atomic Counters + +DSQL has no sequences (`SERIAL` / `BIGSERIAL` are rejected by the validator), so code that needs a monotonic sequence number is otherwise pushed into a racy `SELECT MAX(seq) + 1` read-modify-write: two concurrent callers read the same maximum and both write the same next value, producing duplicate or skipped numbers. + +`db.counter(name)` replaces that pattern with a single atomic upsert against a framework-managed table: + +```sql +CREATE TABLE IF NOT EXISTS _blocks_counters (name TEXT PRIMARY KEY, value BIGINT NOT NULL DEFAULT 0); + +-- Counter.next(delta) — atomic read-and-increment in one statement: +INSERT INTO _blocks_counters (name, value) VALUES ($name, $delta) + ON CONFLICT (name) DO UPDATE SET value = _blocks_counters.value + $delta + RETURNING value; +``` + +- **Atomicity** — the increment is one `INSERT ... ON CONFLICT DO UPDATE ... RETURNING` statement, so there is no read-then-write window for a competing writer. It runs inside `transaction({ retryOnConflict: true })` so a DSQL OCC conflict on the row retries transparently. +- **Managed table** — `_blocks_counters` is created by the migration Lambda (admin) on every deploy, like `_migrations`. The app runtime is DML-only and never issues the DDL. Because the table is admin-owned and not named `_migrations`, `provisionAppRole` automatically grants the app role `SELECT/INSERT/UPDATE/DELETE` on it. +- **Mock parity** — `DsqlMockEngine` creates the same table inside its `withDdl` init window, so `db.counter()` behaves identically against PGlite locally. +- **Range** — values are stored as `BIGINT` and returned as `number` (exact up to 2^53 − 1). + +API: `Counter.next(delta = 1)`, `Counter.current()`, `Counter.reset(value = 0)`. + ## Error Translation Happens in the engine layer (same pattern as bb-data engines): @@ -155,6 +177,7 @@ The `DistributedDatabase` class does not wrap errors — engines handle translat - Retries with exponential backoff on transient connection errors (cluster may take a moment after creation) - `.sql` files bundled into the Lambda package via CDK `commandHooks` - `migrationsHash` property triggers re-invocation when files change +- **Ensures the `_blocks_counters` table** exists (see [Atomic Counters](#atomic-counters)) — created here as admin so the DML grants below cover it - **Provisions custom DB role** on every deploy (idempotent): 1. `CREATE ROLE "app_role" WITH LOGIN` (if not exists) 2. `AWS IAM GRANT "app_role" TO 'arn:aws:iam::...:role/...'` (maps IAM to DB role) diff --git a/packages/bb-distributed-data/README.md b/packages/bb-distributed-data/README.md index 046312e3b..15d8ac43c 100644 --- a/packages/bb-distributed-data/README.md +++ b/packages/bb-distributed-data/README.md @@ -61,6 +61,25 @@ interface TransactionOptions { } ``` +## Atomic Counters + +DSQL has no sequences (`SERIAL` / `BIGSERIAL`), so a monotonic sequence number would otherwise force a racy `SELECT MAX(seq) + 1` read-modify-write — two concurrent callers read the same maximum and both write the same next value, producing duplicate or skipped numbers. `db.counter(name)` gives you a race-free named counter instead: each operation is a single atomic upsert wrapped in OCC retry. + +```typescript +// Per-user monotonic sequence, safe under concurrency: +const seq = await db.counter(`notes:${userId}`).next(); +await db.execute( + sql`INSERT INTO notes (id, user_id, seq, body) VALUES (${id}, ${userId}, ${seq}, ${body})` +); + +await db.counter('signups').next(); // increment by 1, returns the new value +await db.counter('score').next(10); // increment by a custom (integer) delta +await db.counter('signups').current(); // read without incrementing (0 if never seen) +await db.counter('signups').reset(); // set back to 0 +``` + +Counters are stored in a framework-managed `_blocks_counters` table that is created automatically on deploy — no migration required. Values are stored as `BIGINT` and returned as `number` (exact up to `Number.MAX_SAFE_INTEGER`). + ## Migrations One DDL statement per file. DML in separate files. This matches DSQL's transaction constraints. @@ -102,7 +121,7 @@ DSQL is a subset of PostgreSQL. The local mock enforces these restrictions so co | Triggers | Event-driven logic (EventBridge, Lambda) | | Views | CTEs or application-layer query composition | | PL/pgSQL functions | `LANGUAGE SQL` functions or app logic | -| SERIAL / BIGSERIAL | UUIDs (`gen_random_uuid()`) | +| SERIAL / BIGSERIAL | UUIDs (`gen_random_uuid()`) for IDs; `db.counter()` for monotonic sequences | | TRUNCATE | `DELETE FROM` | | Temporary tables | CTEs or subqueries | | LISTEN / NOTIFY | AppSync Events, EventBridge, or polling | diff --git a/packages/bb-distributed-data/src/counter.test.ts b/packages/bb-distributed-data/src/counter.test.ts new file mode 100644 index 000000000..b8b4c72b9 --- /dev/null +++ b/packages/bb-distributed-data/src/counter.test.ts @@ -0,0 +1,158 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** + * Tests for the atomic Counter primitive (DSQL upsert semantics). + * + * Runs against the mock engine (PGlite). The atomic upsert replaces the racy + * `SELECT MAX(seq) + 1` read-modify-write that DSQL otherwise forces (no + * sequences), so the headline test asserts no duplicate values under + * concurrent increments. + */ + +import { describe, it, before, after, beforeEach } from 'node:test'; +import assert from 'node:assert/strict'; +import { rmSync } from 'node:fs'; +import { DatabaseBase, sql } from '@aws-blocks/data-common'; +import { DsqlMockEngine } from './engines/dsql-mock-engine.js'; +import { Counter, ensureCounterTable, COUNTER_TABLE } from './counter.js'; +import { DistributedDatabase } from './index.mock.js'; + +const DIR = '.bb-data/__test_counter__'; + +describe('Counter (mock)', () => { + let engine: DsqlMockEngine; + let base: DatabaseBase; + const counter = (name: string) => new Counter(async () => base, name); + + before(async () => { + rmSync(DIR, { recursive: true, force: true }); + engine = new DsqlMockEngine(DIR); + base = new DatabaseBase(engine); + // The counter table is DDL — created by admin/migrations in prod. Locally + // it goes through the same withDdl escape hatch the migration runner uses. + await engine.withDdl(() => ensureCounterTable(engine)); + }); + + after(async () => { + await engine.destroy(); + rmSync(DIR, { recursive: true, force: true }); + }); + + beforeEach(async () => { + await base.execute(sql`DELETE FROM _blocks_counters`); + }); + + it('next() returns a monotonic sequence starting at 1', async () => { + const c = counter('seq'); + assert.equal(await c.next(), 1); + assert.equal(await c.next(), 2); + assert.equal(await c.next(), 3); + }); + + it('next(delta) adds an arbitrary integer, including negatives', async () => { + const c = counter('by-tens'); + assert.equal(await c.next(10), 10); + assert.equal(await c.next(5), 15); + assert.equal(await c.next(-3), 12); + }); + + it('current() reads without incrementing; unknown counter is 0', async () => { + const c = counter('reads'); + assert.equal(await c.current(), 0); + await c.next(); + await c.next(); + assert.equal(await c.current(), 2); + assert.equal(await c.current(), 2); + }); + + it('reset() sets an explicit value', async () => { + const c = counter('resettable'); + await c.next(); + await c.next(); + await c.reset(100); + assert.equal(await c.current(), 100); + assert.equal(await c.next(), 101); + await c.reset(); + assert.equal(await c.current(), 0); + }); + + it('named counters are independent', async () => { + const a = counter('a'); + const b = counter('b'); + await a.next(); + await a.next(); + await a.next(); + await b.next(); + assert.equal(await a.current(), 3); + assert.equal(await b.current(), 1); + }); + + it('recovers from an OCC conflict via retry without double counting', async () => { + const c = counter('occ'); + assert.equal(await c.next(), 1); + engine.simulateConflict(); + // First commit throws 40001 and rolls back; retry succeeds. The rolled-back + // attempt must NOT leak its increment. + assert.equal(await c.next(), 2); + assert.equal(await c.current(), 2); + }); + + it('concurrent next() calls never produce a duplicate value', async () => { + const c = counter('race'); + const N = 25; + const results = await Promise.all(Array.from({ length: N }, () => c.next())); + const unique = new Set(results); + assert.equal( + unique.size, + N, + `expected ${N} unique values, got ${unique.size}: ${[...results].sort((x, y) => x - y).join(',')}`, + ); + assert.equal(Math.max(...results), N); + assert.equal(Math.min(...results), 1); + assert.equal(await c.current(), N); + }); + + it('rejects an empty counter name', () => { + assert.throws(() => counter(''), /non-empty string/); + }); + + it('rejects a non-integer delta or reset value', async () => { + const c = counter('validate'); + await assert.rejects(() => c.next(1.5), /must be an integer/); + await assert.rejects(() => c.reset(0.5), /must be an integer/); + }); +}); + +describe('DistributedDatabase.counter (mock integration)', () => { + const scope = { id: 'test' }; + let db: InstanceType; + + before(async () => { + db = new DistributedDatabase(scope as any, 'counterint'); + // Clear any state persisted from a previous run of this data dir. + await db.execute(sql`DELETE FROM _blocks_counters`); + }); + + after(async () => { + await (db as any).mockEngine.destroy(); + }); + + it('auto-creates the counter table with no migrations configured', async () => { + // No migrationsPath was provided, yet counter() works — proving the table + // is created during init (mirrors the migration Lambda creating it in prod). + const seq = db.counter('notes:user-1'); + assert.equal(await seq.next(), 1); + assert.equal(await seq.next(), 2); + assert.equal(await db.counter('notes:user-2').next(), 1); + }); + + it('persists to the framework-managed _blocks_counters table', async () => { + assert.equal(COUNTER_TABLE, '_blocks_counters'); + await db.counter('direct-check').next(7); + const row = await db.queryOne<{ value: string | number }>( + sql`SELECT value FROM _blocks_counters WHERE name = ${'direct-check'}`, + ); + assert.equal(Number(row?.value), 7); + }); +}); diff --git a/packages/bb-distributed-data/src/counter.ts b/packages/bb-distributed-data/src/counter.ts new file mode 100644 index 000000000..c76453c39 --- /dev/null +++ b/packages/bb-distributed-data/src/counter.ts @@ -0,0 +1,141 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** + * Atomic named counters for {@link DistributedDatabase} (Aurora DSQL). + * + * Aurora DSQL does not support sequences (`SERIAL`, `BIGSERIAL`, `CREATE SEQUENCE`), + * so code that needs a monotonic sequence number is otherwise forced into a racy + * `SELECT MAX(seq) + 1` read-modify-write: two concurrent callers read the same + * max and both write the same next value, producing duplicate or skipped numbers. + * + * This primitive replaces that pattern with a single atomic upsert against a + * framework-managed `_blocks_counters` table, wrapped in an OCC-retrying + * transaction so concurrent callers never observe a duplicate value. + */ + +import type { DatabaseBase, DatabaseEngine } from '@aws-blocks/data-common'; +import { sql } from '@aws-blocks/data-common'; +import { transactionWithRetry } from './transaction.js'; + +/** Name of the framework-managed table backing all atomic counters. */ +export const COUNTER_TABLE = '_blocks_counters'; + +/** + * Idempotent DDL for the framework-managed counter table. + * + * Created by the migration Lambda (as admin) in production and by the mock + * engine during local init — never by the app runtime, which is DML-only. + */ +export const COUNTER_TABLE_DDL = + `CREATE TABLE IF NOT EXISTS ${COUNTER_TABLE} (name TEXT PRIMARY KEY, value BIGINT NOT NULL DEFAULT 0)`; + +/** + * Create the framework-managed counter table if it does not already exist. + * + * Must run with DDL privileges — the admin role via the migration Lambda in + * production, or inside `DsqlMockEngine.withDdl(...)` locally. The app runtime + * only has DML access and must not call this. + */ +export async function ensureCounterTable(engine: DatabaseEngine): Promise { + await engine.execute(COUNTER_TABLE_DDL); +} + +function assertValidName(name: string): void { + if (typeof name !== 'string' || name.length === 0) { + throw new Error('Counter name must be a non-empty string'); + } +} + +function assertInteger(label: string, n: number): void { + if (!Number.isInteger(n)) { + throw new Error(`Counter ${label} must be an integer, received: ${n}`); + } +} + +/** + * A named, atomic counter backed by {@link DistributedDatabase} (Aurora DSQL). + * + * Obtain one via {@link DistributedDatabase.counter}. Every mutation is a single + * atomic upsert wrapped in an OCC-retrying transaction, so concurrent callers + * never observe a duplicate or skipped value — the correct replacement for a + * racy `SELECT MAX(seq) + 1` read-modify-write. + * + * Values are stored as `BIGINT` and returned as `number`; they are exact up to + * `Number.MAX_SAFE_INTEGER` (2^53 − 1). + * + * @example + * ```ts + * // Per-user monotonic sequence, race-free under concurrency: + * const seq = await db.counter(`notes:${userId}`).next(); + * await db.execute(sql` + * INSERT INTO notes (id, user_id, seq, body) VALUES (${id}, ${userId}, ${seq}, ${body}) + * `); + * ``` + */ +export class Counter { + /** @internal Use {@link DistributedDatabase.counter} to construct. */ + constructor( + private readonly resolveBase: () => Promise, + private readonly name: string, + ) { + assertValidName(name); + } + + /** + * Atomically add `delta` to the counter and return its new value. + * The first call on a never-seen counter returns `delta` (default 1). + * + * @param delta - Integer amount to add. Defaults to 1. May be negative. + * @returns The counter value after applying `delta`. + */ + async next(delta = 1): Promise { + assertInteger('delta', delta); + const base = await this.resolveBase(); + return transactionWithRetry( + base, + async (tx) => { + const rows = await tx.query<{ value: string | number | bigint }>(sql` + INSERT INTO _blocks_counters (name, value) VALUES (${this.name}, ${delta}) + ON CONFLICT (name) DO UPDATE SET value = _blocks_counters.value + ${delta} + RETURNING value + `); + return Number(rows[0]!.value); + }, + { retryOnConflict: true }, + ); + } + + /** + * Read the current counter value without modifying it. + * Returns 0 if the counter has never been incremented. + */ + async current(): Promise { + const base = await this.resolveBase(); + const row = await base.queryOne<{ value: string | number | bigint }>( + sql`SELECT value FROM _blocks_counters WHERE name = ${this.name}`, + ); + return row ? Number(row.value) : 0; + } + + /** + * Atomically set the counter to an explicit value (default 0). + * The next {@link Counter.next} call returns `value + delta`. + * + * @param value - Integer value to store. Defaults to 0. + */ + async reset(value = 0): Promise { + assertInteger('value', value); + const base = await this.resolveBase(); + await transactionWithRetry( + base, + async (tx) => { + await tx.execute(sql` + INSERT INTO _blocks_counters (name, value) VALUES (${this.name}, ${value}) + ON CONFLICT (name) DO UPDATE SET value = ${value} + `); + }, + { retryOnConflict: true }, + ); + } +} diff --git a/packages/bb-distributed-data/src/index.aws.ts b/packages/bb-distributed-data/src/index.aws.ts index 9843ddc43..56a6a5c6a 100644 --- a/packages/bb-distributed-data/src/index.aws.ts +++ b/packages/bb-distributed-data/src/index.aws.ts @@ -12,6 +12,7 @@ import { DatabaseBase, type SqlQuery, type Transaction } from '@aws-blocks/data- import { DsqlSigner } from '@aws-sdk/dsql-signer'; import { DsqlEngine } from './engines/dsql-engine.js'; import { transactionWithRetry } from './transaction.js'; +import { Counter } from './counter.js'; import type { DistributedDatabaseOptions, TransactionOptions } from './types.js'; import { ENV_SANITIZE, sanitizeDbRoleName } from './constants.js'; import { Logger } from '@aws-blocks/bb-logger'; @@ -64,6 +65,18 @@ export class DistributedDatabase extends Scope { return transactionWithRetry(this.base, fn, options); } + /** + * Get a named atomic counter backed by this database. + * + * Use this instead of a racy `SELECT MAX(seq) + 1` read-modify-write when you + * need a monotonic sequence number — DSQL has no sequences (SERIAL / BIGSERIAL). + * + * @param name - Stable identifier for the counter (e.g. `notes:${userId}`). + */ + counter(name: string): Counter { + return new Counter(async () => this.base, name); + } + /** @internal */ getEngine() { return this.base.getEngine(); } } @@ -71,4 +84,5 @@ export class DistributedDatabase extends Scope { export { sql, createKyselyAdapter } from '@aws-blocks/data-common'; export type { SqlQuery, Transaction } from '@aws-blocks/data-common'; export { DistributedDatabaseErrors } from './errors.js'; +export { Counter } from './counter.js'; export type { DistributedDatabaseOptions, TransactionOptions } from './types.js'; diff --git a/packages/bb-distributed-data/src/index.mock.ts b/packages/bb-distributed-data/src/index.mock.ts index 2d1aac07e..2430c57dd 100644 --- a/packages/bb-distributed-data/src/index.mock.ts +++ b/packages/bb-distributed-data/src/index.mock.ts @@ -12,6 +12,7 @@ import { DatabaseBase, type SqlQuery, type Transaction } from '@aws-blocks/data- import { DsqlMockEngine } from './engines/dsql-mock-engine.js'; import { runMigrations, loadMigrationsFromDir } from './migrations.js'; import { transactionWithRetry } from './transaction.js'; +import { Counter, ensureCounterTable } from './counter.js'; import type { DistributedDatabaseOptions, TransactionOptions } from './types.js'; import { Logger } from '@aws-blocks/bb-logger'; import type { ChildLogger } from '@aws-blocks/bb-logger'; @@ -19,7 +20,7 @@ import type { ChildLogger } from '@aws-blocks/bb-logger'; export class DistributedDatabase extends Scope { private base: DatabaseBase; private mockEngine: DsqlMockEngine; - private migrationsRun: Promise | null = null; + private initPromise: Promise; /** @internal Logger for internal operations. Defaults to error-level when not provided. */ protected log: ChildLogger; @@ -31,15 +32,19 @@ export class DistributedDatabase extends Scope { this.base = new DatabaseBase(this.mockEngine); registerSdkIdentifiers(this.fullId, { clusterEndpoint: `mock-endpoint-${this.fullId}` }); - if (options?.migrationsPath) { - const path = options.migrationsPath; - this.migrationsRun = loadMigrationsFromDir(path) - .then(m => this.mockEngine.withDdl(() => runMigrations(this.mockEngine, m))) - .then(() => {}); - } + // Create the framework-managed counter table (and run migrations, if any) + // in a single DDL window — the mock engine rejects DDL outside withDdl, + // mirroring the app runtime's DML-only access. + const migrationsPath = options?.migrationsPath; + this.initPromise = this.mockEngine.withDdl(async () => { + await ensureCounterTable(this.mockEngine); + if (migrationsPath) { + await runMigrations(this.mockEngine, await loadMigrationsFromDir(migrationsPath)); + } + }); } - private async ready(): Promise { if (this.migrationsRun) await this.migrationsRun; } + private async ready(): Promise { await this.initPromise; } query(query: SqlQuery): Promise { return this.ready().then(() => this.base.query(query)); } queryOne(query: SqlQuery): Promise { return this.ready().then(() => this.base.queryOne(query)); } @@ -56,6 +61,18 @@ export class DistributedDatabase extends Scope { return transactionWithRetry(this.base, fn, options); } + /** + * Get a named atomic counter backed by this database. + * + * Use this instead of a racy `SELECT MAX(seq) + 1` read-modify-write when you + * need a monotonic sequence number — DSQL has no sequences (SERIAL / BIGSERIAL). + * + * @param name - Stable identifier for the counter (e.g. `notes:${userId}`). + */ + counter(name: string): Counter { + return new Counter(async () => { await this.ready(); return this.base; }, name); + } + /** Test helper: simulate OCC conflict on next commit. */ simulateConflict(): void { this.mockEngine.simulateConflict(); } @@ -66,4 +83,5 @@ export class DistributedDatabase extends Scope { export { sql, createKyselyAdapter } from '@aws-blocks/data-common'; export type { SqlQuery, Transaction } from '@aws-blocks/data-common'; export { DistributedDatabaseErrors } from './errors.js'; +export { Counter } from './counter.js'; export type { DistributedDatabaseOptions, TransactionOptions } from './types.js'; diff --git a/packages/bb-distributed-data/src/migration-lambda.ts b/packages/bb-distributed-data/src/migration-lambda.ts index 4114dcc13..6f27bea7f 100644 --- a/packages/bb-distributed-data/src/migration-lambda.ts +++ b/packages/bb-distributed-data/src/migration-lambda.ts @@ -19,6 +19,7 @@ import { existsSync } from 'node:fs'; import type { CloudFormationCustomResourceEvent } from 'aws-lambda'; import { DsqlEngine } from './engines/dsql-engine.js'; import { runMigrations, loadMigrationsFromDir } from './migrations.js'; +import { ensureCounterTable } from './counter.js'; import { LAMBDA_MIGRATIONS_DIR } from './constants.js'; // Set by the CDK construct's environment config. Default is the Lambda deployment root @@ -147,6 +148,10 @@ export const handler = async (event: CloudFormationCustomResourceEvent): Promise console.log('[bb-distributed-data] No migrations directory bundled, skipping migrations'); } + // 1.5 Ensure the framework-managed atomic counter table exists (admin-owned). + // provisionAppRole below grants the app role DML on it (it is != _migrations). + await withRetry(() => ensureCounterTable(engine)); + // 2. Provision the app Lambda's custom DB role (least-privilege DML access) await withRetry(() => provisionAppRole(engine, dbRoleName, appRoleArn));