From 8f698060ed36a89c6fd928b4a41adfa39369fca2 Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:42:25 +0000 Subject: [PATCH 1/5] feat(durable): write a Postgres commit as one set-based statement Each record used to cost three round trips (id check, id claim, insert) and each document action read its record twice. A commit now checks and claims ids as sets, inserts each table's rows in one multi-row statement, reads the documents it touches once, and clears and adds revisions together, all as data-modifying CTEs of the statement that advances the session. Conversation records are cached per storage, which is the Session's only writer. Admitting one input goes from 22 queries to 6, creating a conversation from 37 to 13, and a plain model turn from 61 to 28. PostgresStorage.open takes migrate: false for deployments that migrate in an explicit step. --- packages/durable/CHANGELOG.md | 7 + packages/durable/README.md | 2 +- .../durable/src/storage/postgres/storage.ts | 388 +++++++++++------- 3 files changed, 240 insertions(+), 157 deletions(-) diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 0fa6e445b8a6..898df8eb9b66 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -2,6 +2,13 @@ ## [Unreleased] +### Changed + +- Postgres storage commits a batch in one transaction of a few set-based statements instead of one statement per record: ids are checked and claimed as sets, each table's rows are inserted in one multi-row statement, documents are read once, and revisions are cleared and added together, all in a single data-modifying CTE statement that also advances the session. Admitting one input drops from 22 queries to 6; creating a conversation from 37 to 13. +- Postgres storage caches the conversation records it has read; the storage is a Session's only writer and conversation records are only inserted, so a record stays valid until this storage writes it. +- `PostgresStorage.open` takes `migrate: false` to skip the schema check for deployments that migrate in an explicit step. + + ### Breaking Changes - The portable SQLite facade in `@earendil-works/pi-durable/storage/sqlite` is asynchronous: `SqliteDatabase` extends the new `SqliteExecutor` (`exec`, `run`, `get`, `all` by SQL text), `prepare` and `SqliteStatement` are removed, `transaction` takes an async callback that receives a transaction handle, and `close()` returns a promise. Custom adapters must be rewritten ([#10232](https://github.com/earendil-works/pi/pull/10232) by [@christianklotz](https://github.com/christianklotz)). diff --git a/packages/durable/README.md b/packages/durable/README.md index db1b1f5c8fe3..3fe97b6d649f 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -523,7 +523,7 @@ const usage = await harness.usage(context); // { models: { "openai/gpt-6-sol": U | Memory | `MemoryStorage` from the package root | Nothing is persisted. | | SQLite | `openNodeSqliteStorage(file)` from `@earendil-works/pi-durable/storage/sqlite/node` | One database file. WAL mode with `synchronous = NORMAL`: commits survive process crashes; the newest may be lost on power or host failure. | | JSONL | `openNodeJsonlStorage(directory, context)` from `@earendil-works/pi-durable/storage/jsonl/node` | Append-only files in one directory. Pass `{ fsync: true }` to flush before each commit marker. | -| Postgres | `PostgresStorage.open(pool, { session })` from `@earendil-works/pi-durable/storage/postgres` | One Session per `session` name; many Sessions share one schema. Locked to one process. | +| Postgres | `PostgresStorage.open(pool, { session })` from `@earendil-works/pi-durable/storage/postgres` | One Session per `session` name; many Sessions share one schema. Locked to one process. Each commit is one transaction of a few set-based statements. Pass `migrate: false` when the schema is migrated in an explicit step, so an open costs two statements. | One process owns a storage at a time; SQLite and JSONL do no cross-process locking. The portable SQLite and JSONL cores (`/storage/sqlite`, `/storage/jsonl`) run without Node APIs, for example on Bun or in Cloudflare Durable Objects, given an asynchronous `SqliteDatabase` facade or a `FileSystem` from `@earendil-works/pi-durable/env`. diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts index 30ba686af228..2a4f9ed548bb 100644 --- a/packages/durable/src/storage/postgres/storage.ts +++ b/packages/durable/src/storage/postgres/storage.ts @@ -53,6 +53,13 @@ type RevisionRow = { }; type IdRow = { readonly id: string }; type MetadataRow = { readonly next_id: string; readonly next_seq: string }; +type DocumentWrites = { + readonly claims: DocumentId[]; + readonly inserts: PostgresValue[][]; + readonly retirements: PostgresValue[][]; + readonly clearRevisions: DocumentId[]; + readonly revisions: PostgresValue[][]; +}; type DocumentAction = { create?: DocumentCreate; copy?: Extract["source"]; @@ -79,6 +86,12 @@ export interface PostgresStorageOptions { readonly session: string; /** Schema holding the tables, created on first open. Defaults to `pi_durable`. */ readonly schema?: string; + /** + * Apply schema migrations on open. Defaults to true. A deployment that migrates in an explicit step before + * starting, or opens many Sessions per process, passes false so an open costs two statements and takes no + * schema-wide lock. + */ + readonly migrate?: boolean; } const parseJson = (value: string): T => JSON.parse(value) as T; @@ -134,6 +147,9 @@ const isAliveAt = (record: DocumentRecord, at: DocumentPoint): boolean => { const isCurrentOnly = (record: DocumentRecord): boolean => record.scope.kind !== "conversation" || record.history === "latest"; +/** A Postgres array literal of record ids, for `= ANY($n::bigint[])` with the adapter's scalar parameter type. */ +const arrayLiteral = (ids: readonly number[]): string => `{${[...new Set(ids)].join(",")}}`; + const writeId = (write: StorageWrite): Id | undefined => { switch (write.type) { case "conversation": @@ -171,6 +187,11 @@ export class PostgresStorage implements Storage { private closed = false; private lost: Error | undefined; private closing: Promise | undefined; + /** + * Conversation records read so far. This storage is the session's only writer and conversation records are + * only ever inserted, so a record stays valid until a commit of this storage writes that conversation. + */ + private readonly conversations = new Map(); private admitted = 0; private drained: (() => void) | undefined; @@ -206,7 +227,7 @@ export class PostgresStorage implements Storage { static async open(pool: PostgresPool, options: PostgresStorageOptions): Promise { if (options.session.length === 0) throw new TypeError("PostgresStorage requires a non-empty session"); const schema = options.schema ?? DEFAULT_POSTGRES_SCHEMA; - await applyPostgresMigrations(pool, schema); + if (options.migrate !== false) await applyPostgresMigrations(pool, schema); const lockClient = await pool.connect(); let storage: PostgresStorage | undefined; @@ -282,17 +303,23 @@ export class PostgresStorage implements Storage { } const committedSeq = seqFromNumber(Number(metadata.next_seq)); await this.checkGlobalIds(client, writes); - await this.checkDocumentActions(client, documentActions); - for (const write of writes) await this.applyTableWrite(client, write, committedSeq); - await this.applyDocumentActions(client, documentActions, committedSeq); - await client.query(`UPDATE ${this.tables.sessions} SET next_id = $2, next_seq = $3 WHERE session_id = $1`, [ - this.session, - Math.max(Number(metadata.next_id), candidateNextId), - committedSeq + 1, - ]); + const existingDocuments = await this.checkDocumentActions(client, documentActions); + const documentWrites = await this.prepareDocumentWrites( + client, + documentActions, + existingDocuments, + committedSeq, + ); + await this.writeBatch(client, writes, documentWrites, committedSeq, { + nextId: Math.max(Number(metadata.next_id), candidateNextId), + nextSeq: committedSeq + 1, + }); return committedSeq; }); this.nextId = Math.max(this.nextId, candidateNextId); + for (const write of writes) { + if (write.type === "conversation") this.conversations.set(write.value.id, write.value); + } return seq; }); } @@ -643,8 +670,12 @@ export class PostgresStorage implements Storage { return record === undefined ? undefined : parseJson(record); } - private readConversation(id: ConversationId): Promise { - return this.recordById(this.tables.conversations, id); + private async readConversation(id: ConversationId): Promise { + const cached = this.conversations.get(id); + if (cached !== undefined) return cached; + const record = await this.recordById(this.tables.conversations, id); + if (record !== undefined) this.conversations.set(id, record); + return record; } private async materializeDocument( @@ -695,18 +726,22 @@ export class PostgresStorage implements Storage { } private async checkGlobalIds(executor: PostgresQueryable, writes: readonly StorageWrite[]): Promise { - const claimed = new Map, TableName>(); + const planned: Array<{ readonly id: Id; readonly table: TableName }> = []; for (const write of writes) { if (write.type === "document.change" || write.type === "document.retire") continue; const document = write.type === "document.create" || write.type === "document.copy"; - const table: TableName = document ? "document" : write.type; - const id = document ? write.record.id : write.value.id; - const existing = ( - await executor.query( - `SELECT record_type FROM ${this.tables.recordIds} WHERE session_id = $1 AND id = $2`, - [this.session, id], - ) - ).rows[0]?.record_type; + planned.push({ id: document ? write.record.id : write.value.id, table: document ? "document" : write.type }); + } + if (planned.length === 0) return; + const existingByID = new Map, TableName>(); + const rows = await executor.query( + `SELECT id, record_type FROM ${this.tables.recordIds} WHERE session_id = $1 AND id = ANY($2::bigint[])`, + [this.session, arrayLiteral(planned.map((entry) => entry.id))], + ); + for (const row of rows.rows) existingByID.set(idFromNumber>(Number(row.id)), row.record_type); + const claimed = new Map, TableName>(); + for (const { id, table } of planned) { + const existing = existingByID.get(id); const earlier = claimed.get(id); if (table === "conversation" || table === "entry" || table === "document") { if (existing !== undefined) throw new Error(`ID ${id} already belongs to ${existing}`); @@ -767,16 +802,26 @@ export class PostgresStorage implements Storage { return actions; } + /** Validates the batch's document actions and returns the records of the documents it touches that already exist. */ private async checkDocumentActions( executor: PostgresQueryable, actions: ReadonlyMap, - ): Promise { + ): Promise> { + const existingByID = new Map(); + if (actions.size === 0) return existingByID; + const rows = await executor.query<{ readonly id: string; readonly record: string }>( + `SELECT id, record::text AS record FROM ${this.tables.documents} WHERE session_id = $1 AND id = ANY($2::bigint[])`, + [this.session, arrayLiteral([...actions.keys()])], + ); + for (const row of rows.rows) { + existingByID.set(idFromNumber(Number(row.id)), parseJson(row.record)); + } const liveCounts = new Map(); for (const [id, action] of actions) { if (action.copy !== undefined && actions.has(action.copy.id)) { throw new StorageRejected(`Document copy ${id} source is changed in the copy batch`); } - const existing = await this.recordById(this.tables.documents, id, executor); + const existing = existingByID.get(id); if (action.create === undefined && existing === undefined) throw new Error(`Unknown document: ${id}`); if (action.create !== undefined && existing !== undefined) throw new Error(`Document ${id} already exists`); if (existing?.retiredAt !== undefined) throw new Error(`Document ${id} is retired`); @@ -793,6 +838,9 @@ export class PostgresStorage implements Storage { throw new Error(`Document ${id} version transition requires a base`); } } + // Only a create or retire can change how many incarnations an address has; a content change on an + // existing document needs no lookup. + if (action.create === undefined && !action.retire) continue; const record = action.create ?? existing!; const key = addressKey(record); let live = liveCounts.get(key); @@ -804,6 +852,7 @@ export class PostgresStorage implements Storage { for (const live of liveCounts.values()) { if (live > 1) throw new Error("Document address already has a current incarnation"); } + return existingByID; } private async currentDocumentId( @@ -823,96 +872,18 @@ export class PostgresStorage implements Storage { return id === undefined ? undefined : idFromNumber(Number(id)); } - private async applyTableWrite(executor: PostgresQueryable, write: StorageWrite, seq: Seq): Promise { - switch (write.type) { - case "conversation": - await this.claimId(executor, write.value.id, "conversation"); - await executor.query( - `INSERT INTO ${this.tables.conversations} - (session_id, id, owner_conversation_id, owner_task_id, record) VALUES ($1, $2, $3, $4, $5)`, - [ - this.session, - write.value.id, - write.value.owner?.conversationId ?? null, - write.value.owner?.taskId ?? null, - encodeJson(write.value), - ], - ); - break; - case "entry": - await this.claimId(executor, write.value.id, "entry"); - await executor.query( - `INSERT INTO ${this.tables.entries} (session_id, id, conversation_id, head, commit_seq, record) - VALUES ($1, $2, $3, $4, $5, $6)`, - [ - this.session, - write.value.id, - write.value.conversationId, - write.value.head ?? null, - seq, - encodeJson(write.value), - ], - ); - break; - case "task": - await this.claimId(executor, write.value.id, "task"); - await executor.query( - `INSERT INTO ${this.tables.tasks} - (session_id, id, conversation_id, kind, status, abort_requested, background, record) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8) - ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, - kind = excluded.kind, status = excluded.status, abort_requested = excluded.abort_requested, - background = excluded.background, record = excluded.record`, - [ - this.session, - write.value.id, - write.value.conversationId, - encodeIndexedString(write.value.kind), - write.value.state.status, - write.value.abortRequested, - write.value.background, - encodeJson(write.value), - ], - ); - break; - case "submission": - await this.claimId(executor, write.value.id, "submission"); - await executor.query( - `INSERT INTO ${this.tables.submissions} (session_id, id, conversation_id, request_id, status, record) - VALUES ($1, $2, $3, $4, $5, $6) - ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, - request_id = excluded.request_id, status = excluded.status, record = excluded.record`, - [ - this.session, - write.value.id, - write.value.conversationId, - write.value.requestId === undefined ? null : encodeIndexedString(write.value.requestId), - write.value.status, - encodeJson(write.value), - ], - ); - break; - case "document.create": - case "document.copy": - case "document.change": - case "document.retire": - break; - } - } - - private async claimId(executor: PostgresQueryable, id: Id, table: TableName): Promise { - await executor.query( - `INSERT INTO ${this.tables.recordIds} (session_id, id, record_type) VALUES ($1, $2, $3) - ON CONFLICT (session_id, id) DO NOTHING`, - [this.session, id, table], - ); - } - - private async applyDocumentActions( + /** + * Resolves the batch's document actions into rows: records to insert, retirements to apply to existing + * documents, revision histories to clear, and revisions to add. Copy sources are read here; the rows are + * written with everything else in `writeBatch`. + */ + private async prepareDocumentWrites( executor: PostgresQueryable, actions: ReadonlyMap, + existing: ReadonlyMap, seq: Seq, - ): Promise { + ): Promise { + const writes: DocumentWrites = { claims: [], inserts: [], retirements: [], clearRevisions: [], revisions: [] }; for (const [id, action] of actions) { let content = action.content; if (action.copy !== undefined) { @@ -938,62 +909,167 @@ export class PostgresStorage implements Storage { } let record: DocumentRecord; if (action.create !== undefined) { - record = { - ...action.create, - createdAt: seq, - ...(action.retire ? { retiredAt: seq } : {}), - }; + record = { ...action.create, createdAt: seq, ...(action.retire ? { retiredAt: seq } : {}) }; const parts = addressParts(record); - await this.claimId(executor, id, "document"); - await executor.query( - `INSERT INTO ${this.tables.documents} - (session_id, id, kind, family, key_value, scope_kind, owner_id, created_at, retired_at, record) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, - [ - this.session, - id, - parts.kind, - parts.family, - parts.keyValue, - parts.scopeKind, - parts.ownerId, - seq, - action.retire ? seq : null, - encodeJson(record), - ], - ); + writes.claims.push(id); + writes.inserts.push([ + this.session, + id, + parts.kind, + parts.family, + parts.keyValue, + parts.scopeKind, + parts.ownerId, + seq, + action.retire ? seq : null, + encodeJson(record), + ]); } else { - record = (await this.recordById(this.tables.documents, id, executor))!; + record = existing.get(id)!; + if (action.retire) { + record = { ...record, retiredAt: seq }; + writes.retirements.push([this.session, id, seq, encodeJson(record)]); + } } - - if (content !== undefined) { - if (content.kind === "base" && isCurrentOnly(record)) await this.deleteRevisions(executor, id); + // A retired current-only document keeps no revisions, so one written in the same batch would only be + // deleted again and is skipped. A base otherwise replaces a current-only document's history. + const retiredCurrentOnly = action.retire && isCurrentOnly(record); + if (retiredCurrentOnly || (content?.kind === "base" && isCurrentOnly(record))) writes.clearRevisions.push(id); + if (content !== undefined && !retiredCurrentOnly) { const encodedContent = content.kind === "base" ? encodeJson(content.value) : encodeJson(content.ops); - await executor.query( - `INSERT INTO ${this.tables.revisions} (session_id, document_id, seq, kind, version, content) - VALUES ($1, $2, $3, $4, $5, $6)`, - [this.session, id, seq, content.kind, content.version, encodedContent], - ); + writes.revisions.push([this.session, id, seq, content.kind, content.version, encodedContent]); } + } + return writes; + } - if (action.retire) { - if (action.create === undefined) { - record = { ...record, retiredAt: seq }; - await executor.query( - `UPDATE ${this.tables.documents} SET retired_at = $3, record = $4 WHERE session_id = $1 AND id = $2`, - [this.session, id, seq, encodeJson(record)], - ); - } - if (isCurrentOnly(record)) await this.deleteRevisions(executor, id); + /** + * Writes the whole batch in one statement: data-modifying CTEs claim every new id, insert each table's rows, + * clear and add revisions, and the statement's final UPDATE advances the session. A task or submission written + * twice in one batch keeps its last write, as serial upserts did. Retiring an existing document is the one + * action written separately, since it updates rows that may also be read by the batch's other actions. + */ + private async writeBatch( + executor: PostgresQueryable, + writes: readonly StorageWrite[], + documents: DocumentWrites, + seq: Seq, + session: { readonly nextId: number; readonly nextSeq: number }, + ): Promise { + const claims = new Map, TableName>(); + const conversations: PostgresValue[][] = []; + const entries: PostgresValue[][] = []; + const tasks = new Map, PostgresValue[]>(); + const submissions = new Map, PostgresValue[]>(); + for (const write of writes) { + switch (write.type) { + case "conversation": + claims.set(write.value.id, "conversation"); + conversations.push([ + this.session, + write.value.id, + write.value.owner?.conversationId ?? null, + write.value.owner?.taskId ?? null, + encodeJson(write.value), + ]); + break; + case "entry": + claims.set(write.value.id, "entry"); + entries.push([ + this.session, + write.value.id, + write.value.conversationId, + write.value.head ?? null, + seq, + encodeJson(write.value), + ]); + break; + case "task": + claims.set(write.value.id, "task"); + tasks.set(write.value.id, [ + this.session, + write.value.id, + write.value.conversationId, + encodeIndexedString(write.value.kind), + write.value.state.status, + write.value.abortRequested, + write.value.background, + encodeJson(write.value), + ]); + break; + case "submission": + claims.set(write.value.id, "submission"); + submissions.set(write.value.id, [ + this.session, + write.value.id, + write.value.conversationId, + write.value.requestId === undefined ? null : encodeIndexedString(write.value.requestId), + write.value.status, + encodeJson(write.value), + ]); + break; + case "document.create": + case "document.copy": + case "document.change": + case "document.retire": + break; } } - } + for (const id of documents.claims) claims.set(id, "document"); - private async deleteRevisions(executor: PostgresQueryable, id: DocumentId): Promise { - await executor.query(`DELETE FROM ${this.tables.revisions} WHERE session_id = $1 AND document_id = $2`, [ - this.session, - id, - ]); + for (const [sessionID, id, retiredAt, record] of documents.retirements) { + await executor.query( + `UPDATE ${this.tables.documents} SET retired_at = $3, record = $4 WHERE session_id = $1 AND id = $2`, + [sessionID, id, retiredAt, record], + ); + } + + const params: PostgresValue[] = [this.session]; + const values = (rows: readonly (readonly PostgresValue[])[]): string => + rows.map((row) => `(${row.map((value) => `$${params.push(value)}`).join(", ")})`).join(", "); + const parts: string[] = []; + if (claims.size > 0) { + parts.push(`claims AS (INSERT INTO ${this.tables.recordIds} (session_id, id, record_type) + VALUES ${values([...claims].map(([id, table]) => [this.session, id, table]))} + ON CONFLICT (session_id, id) DO NOTHING)`); + } + if (conversations.length > 0) { + parts.push(`conversations AS (INSERT INTO ${this.tables.conversations} + (session_id, id, owner_conversation_id, owner_task_id, record) VALUES ${values(conversations)})`); + } + if (entries.length > 0) { + parts.push(`entries AS (INSERT INTO ${this.tables.entries} + (session_id, id, conversation_id, head, commit_seq, record) VALUES ${values(entries)})`); + } + if (tasks.size > 0) { + parts.push(`tasks AS (INSERT INTO ${this.tables.tasks} + (session_id, id, conversation_id, kind, status, abort_requested, background, record) + VALUES ${values([...tasks.values()])} + ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, + kind = excluded.kind, status = excluded.status, abort_requested = excluded.abort_requested, + background = excluded.background, record = excluded.record)`); + } + if (submissions.size > 0) { + parts.push(`submissions AS (INSERT INTO ${this.tables.submissions} + (session_id, id, conversation_id, request_id, status, record) VALUES ${values([...submissions.values()])} + ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, + request_id = excluded.request_id, status = excluded.status, record = excluded.record)`); + } + if (documents.inserts.length > 0) { + parts.push(`documents AS (INSERT INTO ${this.tables.documents} + (session_id, id, kind, family, key_value, scope_kind, owner_id, created_at, retired_at, record) + VALUES ${values(documents.inserts)})`); + } + if (documents.clearRevisions.length > 0) { + parts.push(`cleared AS (DELETE FROM ${this.tables.revisions} + WHERE session_id = $1 AND document_id = ANY($${params.push(arrayLiteral(documents.clearRevisions))}::bigint[]))`); + } + if (documents.revisions.length > 0) { + parts.push(`revisions AS (INSERT INTO ${this.tables.revisions} + (session_id, document_id, seq, kind, version, content) VALUES ${values(documents.revisions)})`); + } + const update = `UPDATE ${this.tables.sessions} SET next_id = $${params.push(session.nextId)}, next_seq = $${params.push(session.nextSeq)} WHERE session_id = $1`; + await executor.query(parts.length === 0 ? update : `WITH ${parts.join(",\n")} ${update}`, params); } private assertOpen(): void { From d1972f233832db094c8c33498726b234a99c7d8e Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:55:15 +0000 Subject: [PATCH 2/5] feat(durable): commit a Postgres batch in two round trips One statement reads the session row, the batch's claimed ids, and its documents as JSON; one data-modifying CTE statement writes the batch. Every CTE is gated on the session still carrying this owner at the sequence the read observed, and the final UPDATE advances it, so a takeover by another process writes nothing and is reported as before. A single statement is atomic, so no transaction brackets the commit. --- .../durable/src/storage/postgres/storage.ts | 204 +++++++++++------- 1 file changed, 128 insertions(+), 76 deletions(-) diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts index 2a4f9ed548bb..37f09fd3c321 100644 --- a/packages/durable/src/storage/postgres/storage.ts +++ b/packages/durable/src/storage/postgres/storage.ts @@ -42,7 +42,6 @@ import { applyPostgresMigrations, DEFAULT_POSTGRES_SCHEMA, quoteIdentifier } fro type StoredTask = TaskRecord; type TableName = "conversation" | "entry" | "task" | "submission" | "document"; -type RecordIdRow = { readonly record_type: TableName }; type JsonRow = { readonly record: string }; type EntryJsonRow = { readonly record: string; readonly commit_seq: string }; type RevisionRow = { @@ -53,6 +52,13 @@ type RevisionRow = { }; type IdRow = { readonly id: string }; type MetadataRow = { readonly next_id: string; readonly next_seq: string }; +type CommitState = { + readonly session: + | { readonly next_id: number | string; readonly next_seq: number | string; readonly owner: string } + | undefined; + readonly ids: ReadonlyMap, TableName>; + readonly documents: ReadonlyMap; +}; type DocumentWrites = { readonly claims: DocumentId[]; readonly inserts: PostgresValue[][]; @@ -147,6 +153,15 @@ const isAliveAt = (record: DocumentRecord, at: DocumentPoint): boolean => { const isCurrentOnly = (record: DocumentRecord): boolean => record.scope.kind !== "conversation" || record.history === "latest"; +// Column types of each table's insert list, for casting VALUES parameters. +const RECORD_ID_TYPES = ["text", "bigint", "text"]; +const CONVERSATION_TYPES = ["text", "bigint", "bigint", "bigint", "json"]; +const ENTRY_TYPES = ["text", "bigint", "bigint", "bigint", "bigint", "json"]; +const TASK_TYPES = ["text", "bigint", "bigint", "text", "text", "boolean", "boolean", "json"]; +const SUBMISSION_TYPES = ["text", "bigint", "bigint", "text", "text", "json"]; +const DOCUMENT_TYPES = ["text", "bigint", "text", "smallint", "text", "text", "bigint", "bigint", "bigint", "json"]; +const REVISION_TYPES = ["text", "bigint", "bigint", "text", "integer", "json"]; + /** A Postgres array literal of record ids, for `= ANY($n::bigint[])` with the adapter's scalar parameter type. */ const arrayLiteral = (ids: readonly number[]): string => `{${[...new Set(ids)].join(",")}}`; @@ -288,42 +303,82 @@ export class PostgresStorage implements Storage { return this.admit(async () => { const documentActions = this.prepareDocumentActions(writes); const candidateNextId = this.candidateNextId(writes); - const seq = await inTransaction(this.pool, "BEGIN", async (client) => { - const metadata = ( - await client.query( - `SELECT next_id, next_seq, owner FROM ${this.tables.sessions} WHERE session_id = $1 FOR UPDATE`, - [this.session], - ) - ).rows[0]; - if (metadata === undefined) throw new Error("Durable Postgres session metadata is missing"); - if (metadata.owner !== this.owner) { - const error = new Error(`Durable Postgres session ${this.session} was opened by another process`); - this.loseLock(error); - throw error; - } - const committedSeq = seqFromNumber(Number(metadata.next_seq)); - await this.checkGlobalIds(client, writes); - const existingDocuments = await this.checkDocumentActions(client, documentActions); - const documentWrites = await this.prepareDocumentWrites( - client, - documentActions, - existingDocuments, - committedSeq, - ); - await this.writeBatch(client, writes, documentWrites, committedSeq, { - nextId: Math.max(Number(metadata.next_id), candidateNextId), - nextSeq: committedSeq + 1, - }); - return committedSeq; + // Two round trips: one read of everything the checks need, then one statement that writes the whole + // batch. A single statement is atomic by itself, so no transaction brackets it; the write is gated on + // this storage still owning the session at the sequence the read observed, which is what the row lock + // used to guarantee. + const state = await this.readCommitState(writes, documentActions); + if (state.session === undefined) throw new Error("Durable Postgres session metadata is missing"); + if (state.session.owner !== this.owner) { + const error = new Error(`Durable Postgres session ${this.session} was opened by another process`); + this.loseLock(error); + throw error; + } + const committedSeq = seqFromNumber(Number(state.session.next_seq)); + this.checkGlobalIds(writes, state.ids); + await this.checkDocumentActions(this.pool, documentActions, state.documents); + const documentWrites = await this.prepareDocumentWrites( + this.pool, + documentActions, + state.documents, + committedSeq, + ); + const written = await this.writeBatch(this.pool, writes, documentWrites, committedSeq, { + nextId: Math.max(Number(state.session.next_id), candidateNextId), + nextSeq: committedSeq + 1, + expectedSeq: committedSeq, }); + if (!written) { + const error = new Error(`Durable Postgres session ${this.session} was opened by another process`); + this.loseLock(error); + throw error; + } this.nextId = Math.max(this.nextId, candidateNextId); for (const write of writes) { if (write.type === "conversation") this.conversations.set(write.value.id, write.value); } - return seq; + return committedSeq; }); } + /** Everything a commit checks before writing, read in one statement: the session row, the batch's claimed ids, and its documents. */ + private async readCommitState( + writes: readonly StorageWrite[], + actions: ReadonlyMap, + ): Promise { + const ids = new Set(); + for (const write of writes) { + const id = writeId(write); + if (id !== undefined) ids.add(id); + } + const row = ( + await this.pool.query<{ readonly state: string }>( + `SELECT json_build_object( + 'session', (SELECT row_to_json(s) FROM ${this.tables.sessions} s WHERE s.session_id = $1), + 'ids', (SELECT coalesce(json_agg(json_build_object('id', r.id, 'type', r.record_type)), '[]'::json) + FROM ${this.tables.recordIds} r WHERE r.session_id = $1 AND r.id = ANY($2::bigint[])), + 'documents', (SELECT coalesce(json_agg(json_build_object('id', d.id, 'record', d.record)), '[]'::json) + FROM ${this.tables.documents} d WHERE d.session_id = $1 AND d.id = ANY($3::bigint[])) + )::text AS state`, + [this.session, arrayLiteral([...ids]), arrayLiteral([...actions.keys()])], + ) + ).rows[0]; + const raw = parseJson<{ + readonly session: { + readonly next_id: number | string; + readonly next_seq: number | string; + readonly owner: string; + } | null; + readonly ids: readonly { readonly id: number | string; readonly type: TableName }[]; + readonly documents: readonly { readonly id: number | string; readonly record: DocumentRecord }[]; + }>(row.state); + const existingIds = new Map, TableName>(); + for (const entry of raw.ids) existingIds.set(idFromNumber>(Number(entry.id)), entry.type); + const documents = new Map(); + for (const entry of raw.documents) documents.set(idFromNumber(Number(entry.id)), entry.record); + return { session: raw.session ?? undefined, ids: existingIds, documents }; + } + async mintId>(): Promise { this.assertOpen(); if (!Number.isSafeInteger(this.nextId)) throw new Error("ID space is exhausted"); @@ -725,22 +780,13 @@ export class PostgresStorage implements Storage { return nextId; } - private async checkGlobalIds(executor: PostgresQueryable, writes: readonly StorageWrite[]): Promise { - const planned: Array<{ readonly id: Id; readonly table: TableName }> = []; + private checkGlobalIds(writes: readonly StorageWrite[], existingByID: ReadonlyMap, TableName>): void { + const claimed = new Map, TableName>(); for (const write of writes) { if (write.type === "document.change" || write.type === "document.retire") continue; const document = write.type === "document.create" || write.type === "document.copy"; - planned.push({ id: document ? write.record.id : write.value.id, table: document ? "document" : write.type }); - } - if (planned.length === 0) return; - const existingByID = new Map, TableName>(); - const rows = await executor.query( - `SELECT id, record_type FROM ${this.tables.recordIds} WHERE session_id = $1 AND id = ANY($2::bigint[])`, - [this.session, arrayLiteral(planned.map((entry) => entry.id))], - ); - for (const row of rows.rows) existingByID.set(idFromNumber>(Number(row.id)), row.record_type); - const claimed = new Map, TableName>(); - for (const { id, table } of planned) { + const table: TableName = document ? "document" : write.type; + const id = document ? write.record.id : write.value.id; const existing = existingByID.get(id); const earlier = claimed.get(id); if (table === "conversation" || table === "entry" || table === "document") { @@ -802,20 +848,12 @@ export class PostgresStorage implements Storage { return actions; } - /** Validates the batch's document actions and returns the records of the documents it touches that already exist. */ + /** Validates the batch's document actions against the records read for them. */ private async checkDocumentActions( executor: PostgresQueryable, actions: ReadonlyMap, - ): Promise> { - const existingByID = new Map(); - if (actions.size === 0) return existingByID; - const rows = await executor.query<{ readonly id: string; readonly record: string }>( - `SELECT id, record::text AS record FROM ${this.tables.documents} WHERE session_id = $1 AND id = ANY($2::bigint[])`, - [this.session, arrayLiteral([...actions.keys()])], - ); - for (const row of rows.rows) { - existingByID.set(idFromNumber(Number(row.id)), parseJson(row.record)); - } + existingByID: ReadonlyMap, + ): Promise { const liveCounts = new Map(); for (const [id, action] of actions) { if (action.copy !== undefined && actions.has(action.copy.id)) { @@ -852,7 +890,6 @@ export class PostgresStorage implements Storage { for (const live of liveCounts.values()) { if (live > 1) throw new Error("Document address already has a current incarnation"); } - return existingByID; } private async currentDocumentId( @@ -954,8 +991,8 @@ export class PostgresStorage implements Storage { writes: readonly StorageWrite[], documents: DocumentWrites, seq: Seq, - session: { readonly nextId: number; readonly nextSeq: number }, - ): Promise { + session: { readonly nextId: number; readonly nextSeq: number; readonly expectedSeq: number }, + ): Promise { const claims = new Map, TableName>(); const conversations: PostgresValue[][] = []; const entries: PostgresValue[][] = []; @@ -1017,59 +1054,74 @@ export class PostgresStorage implements Storage { } for (const id of documents.claims) claims.set(id, "document"); - for (const [sessionID, id, retiredAt, record] of documents.retirements) { - await executor.query( - `UPDATE ${this.tables.documents} SET retired_at = $3, record = $4 WHERE session_id = $1 AND id = $2`, - [sessionID, id, retiredAt, record], - ); + // Every write is gated on the guard: the session row still carries this owner at the sequence the + // pre-read observed. When another process has taken the session over, every CTE writes nothing and the + // final UPDATE affects no row. + const params: PostgresValue[] = [this.session, this.owner, session.expectedSeq]; + const guard = `WHERE EXISTS (SELECT 1 FROM guard)`; + // Parameters in a VALUES subselect carry no type inference, so each column is cast to its table type. + const values = (rows: readonly (readonly PostgresValue[])[], casts: readonly string[]): string => { + const names = casts.map((_, index) => `c${index}`).join(", "); + const list = rows + .map((row) => `(${row.map((value, index) => `$${params.push(value)}::${casts[index]}`).join(", ")})`) + .join(", "); + return `SELECT ${names} FROM (VALUES ${list}) AS v(${names}) ${guard}`; + }; + const parts: string[] = [ + `guard AS (SELECT 1 FROM ${this.tables.sessions} WHERE session_id = $1 AND owner = $2 AND next_seq = $3)`, + ]; + for (const [, id, retiredAt, record] of documents.retirements) { + parts.push(`retire_${id} AS (UPDATE ${this.tables.documents} SET retired_at = $${params.push(retiredAt)}::bigint, + record = $${params.push(record)}::json WHERE session_id = $1 AND id = $${params.push(id)}::bigint AND EXISTS (SELECT 1 FROM guard))`); } - - const params: PostgresValue[] = [this.session]; - const values = (rows: readonly (readonly PostgresValue[])[]): string => - rows.map((row) => `(${row.map((value) => `$${params.push(value)}`).join(", ")})`).join(", "); - const parts: string[] = []; if (claims.size > 0) { parts.push(`claims AS (INSERT INTO ${this.tables.recordIds} (session_id, id, record_type) - VALUES ${values([...claims].map(([id, table]) => [this.session, id, table]))} + ${values( + [...claims].map(([id, table]) => [this.session, id, table]), + RECORD_ID_TYPES, + )} ON CONFLICT (session_id, id) DO NOTHING)`); } if (conversations.length > 0) { parts.push(`conversations AS (INSERT INTO ${this.tables.conversations} - (session_id, id, owner_conversation_id, owner_task_id, record) VALUES ${values(conversations)})`); + (session_id, id, owner_conversation_id, owner_task_id, record) ${values(conversations, CONVERSATION_TYPES)})`); } if (entries.length > 0) { parts.push(`entries AS (INSERT INTO ${this.tables.entries} - (session_id, id, conversation_id, head, commit_seq, record) VALUES ${values(entries)})`); + (session_id, id, conversation_id, head, commit_seq, record) ${values(entries, ENTRY_TYPES)})`); } if (tasks.size > 0) { parts.push(`tasks AS (INSERT INTO ${this.tables.tasks} (session_id, id, conversation_id, kind, status, abort_requested, background, record) - VALUES ${values([...tasks.values()])} + ${values([...tasks.values()], TASK_TYPES)} ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, kind = excluded.kind, status = excluded.status, abort_requested = excluded.abort_requested, background = excluded.background, record = excluded.record)`); } if (submissions.size > 0) { parts.push(`submissions AS (INSERT INTO ${this.tables.submissions} - (session_id, id, conversation_id, request_id, status, record) VALUES ${values([...submissions.values()])} + (session_id, id, conversation_id, request_id, status, record) ${values([...submissions.values()], SUBMISSION_TYPES)} ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, request_id = excluded.request_id, status = excluded.status, record = excluded.record)`); } if (documents.inserts.length > 0) { parts.push(`documents AS (INSERT INTO ${this.tables.documents} (session_id, id, kind, family, key_value, scope_kind, owner_id, created_at, retired_at, record) - VALUES ${values(documents.inserts)})`); + ${values(documents.inserts, DOCUMENT_TYPES)})`); } if (documents.clearRevisions.length > 0) { parts.push(`cleared AS (DELETE FROM ${this.tables.revisions} - WHERE session_id = $1 AND document_id = ANY($${params.push(arrayLiteral(documents.clearRevisions))}::bigint[]))`); + WHERE session_id = $1 AND document_id = ANY($${params.push(arrayLiteral(documents.clearRevisions))}::bigint[]) + AND EXISTS (SELECT 1 FROM guard))`); } if (documents.revisions.length > 0) { parts.push(`revisions AS (INSERT INTO ${this.tables.revisions} - (session_id, document_id, seq, kind, version, content) VALUES ${values(documents.revisions)})`); + (session_id, document_id, seq, kind, version, content) ${values(documents.revisions, REVISION_TYPES)})`); } - const update = `UPDATE ${this.tables.sessions} SET next_id = $${params.push(session.nextId)}, next_seq = $${params.push(session.nextSeq)} WHERE session_id = $1`; - await executor.query(parts.length === 0 ? update : `WITH ${parts.join(",\n")} ${update}`, params); + const update = `UPDATE ${this.tables.sessions} SET next_id = $${params.push(session.nextId)}, next_seq = $${params.push(session.nextSeq)} + WHERE session_id = $1 AND owner = $2 AND next_seq = $3`; + const result = await executor.query(`WITH ${parts.join(",\n")} ${update}`, params); + return (result.rowCount ?? 0) > 0; } private assertOpen(): void { From 718dff04d29c4720f4babe00dcacb2ff3ec239d7 Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:02:15 +0000 Subject: [PATCH 3/5] feat(durable): serve repeated Postgres reads from memory The storage is the Session's only writer, so a record it has read stays valid until it writes that record. Tasks and submissions by id and request, each conversation's newest entries and head markers, current document scans and addresses, and the materialized current content of documents are cached and kept current by every commit. Cached values are JSON-encoded so callers get detached objects, as from the database. The delta version check joins the commit's pre-read. A plain turn is now 9 round trips, from 61. --- packages/durable/CHANGELOG.md | 6 +- packages/durable/README.md | 2 +- .../durable/src/storage/postgres/storage.ts | 380 +++++++++++++++--- 3 files changed, 328 insertions(+), 60 deletions(-) diff --git a/packages/durable/CHANGELOG.md b/packages/durable/CHANGELOG.md index 898df8eb9b66..5e275ec11723 100644 --- a/packages/durable/CHANGELOG.md +++ b/packages/durable/CHANGELOG.md @@ -4,9 +4,9 @@ ### Changed -- Postgres storage commits a batch in one transaction of a few set-based statements instead of one statement per record: ids are checked and claimed as sets, each table's rows are inserted in one multi-row statement, documents are read once, and revisions are cleared and added together, all in a single data-modifying CTE statement that also advances the session. Admitting one input drops from 22 queries to 6; creating a conversation from 37 to 13. -- Postgres storage caches the conversation records it has read; the storage is a Session's only writer and conversation records are only inserted, so a record stays valid until this storage writes it. -- `PostgresStorage.open` takes `migrate: false` to skip the schema check for deployments that migrate in an explicit step. +- Postgres storage is built for a database that is a network hop away. A commit is two round trips: one statement reads the session row, the batch's claimed ids, its documents, and their latest revision versions as JSON; one data-modifying CTE statement writes every table's rows, clears and adds revisions, and advances the session, each part gated on the session still carrying this owner at the sequence the read observed. A single statement is atomic, so no transaction brackets a commit, and a takeover by another process writes nothing and is reported as before. +- Postgres storage serves repeated reads from memory. It is the Session's only writer, so a record it has read stays valid until it writes that record: conversation records, tasks and submissions by id and request, each conversation's newest entries and head markers, current document scans and addresses, and the materialized current content of documents are cached and kept current by every commit. Measured on one faux-model turn, admitting an input drops from 22 queries to 3, creating a conversation from 37 to 11, a plain turn from 61 to 9, and a turn with one tool call from 189 to 31. +- `PostgresStorage.open` takes `migrate: false` to skip the schema check for deployments that migrate in an explicit step; an open is then two statements. ### Breaking Changes diff --git a/packages/durable/README.md b/packages/durable/README.md index 3fe97b6d649f..2d1174f827f1 100644 --- a/packages/durable/README.md +++ b/packages/durable/README.md @@ -523,7 +523,7 @@ const usage = await harness.usage(context); // { models: { "openai/gpt-6-sol": U | Memory | `MemoryStorage` from the package root | Nothing is persisted. | | SQLite | `openNodeSqliteStorage(file)` from `@earendil-works/pi-durable/storage/sqlite/node` | One database file. WAL mode with `synchronous = NORMAL`: commits survive process crashes; the newest may be lost on power or host failure. | | JSONL | `openNodeJsonlStorage(directory, context)` from `@earendil-works/pi-durable/storage/jsonl/node` | Append-only files in one directory. Pass `{ fsync: true }` to flush before each commit marker. | -| Postgres | `PostgresStorage.open(pool, { session })` from `@earendil-works/pi-durable/storage/postgres` | One Session per `session` name; many Sessions share one schema. Locked to one process. Each commit is one transaction of a few set-based statements. Pass `migrate: false` when the schema is migrated in an explicit step, so an open costs two statements. | +| Postgres | `PostgresStorage.open(pool, { session })` from `@earendil-works/pi-durable/storage/postgres` | One Session per `session` name; many Sessions share one schema. Locked to one process. Built for a remote database: a commit is two round trips and repeated reads are served from memory. Pass `migrate: false` when the schema is migrated in an explicit step, so an open costs two statements. | One process owns a storage at a time; SQLite and JSONL do no cross-process locking. The portable SQLite and JSONL cores (`/storage/sqlite`, `/storage/jsonl`) run without Node APIs, for example on Bun or in Cloudflare Durable Objects, given an asynchronous `SqliteDatabase` facade or a `FileSystem` from `@earendil-works/pi-durable/env`. diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts index 37f09fd3c321..3efc943ca024 100644 --- a/packages/durable/src/storage/postgres/storage.ts +++ b/packages/durable/src/storage/postgres/storage.ts @@ -52,12 +52,16 @@ type RevisionRow = { }; type IdRow = { readonly id: string }; type MetadataRow = { readonly next_id: string; readonly next_seq: string }; +/** A conversation's own entries over one contiguous id range ending at its newest entry, newest-last. */ +type EntryWindow = { ids: EntryId[]; complete: boolean }; +type DocumentState = { readonly record: string; readonly version: number; value: JsonObject; deltasSinceBase: number }; type CommitState = { readonly session: | { readonly next_id: number | string; readonly next_seq: number | string; readonly owner: string } | undefined; readonly ids: ReadonlyMap, TableName>; readonly documents: ReadonlyMap; + readonly versions: ReadonlyMap; }; type DocumentWrites = { readonly claims: DocumentId[]; @@ -65,6 +69,11 @@ type DocumentWrites = { readonly retirements: PostgresValue[][]; readonly clearRevisions: DocumentId[]; readonly revisions: PostgresValue[][]; + /** The record as written and the content applied, per document, for the caches. */ + readonly applied: Map< + DocumentId, + { readonly record: DocumentRecord; readonly content: DocumentContent | undefined } + >; }; type DocumentAction = { create?: DocumentCreate; @@ -162,6 +171,24 @@ const SUBMISSION_TYPES = ["text", "bigint", "bigint", "text", "text", "json"]; const DOCUMENT_TYPES = ["text", "bigint", "text", "smallint", "text", "text", "bigint", "bigint", "bigint", "json"]; const REVISION_TYPES = ["text", "bigint", "bigint", "text", "integer", "json"]; +/** Entries loaded per conversation on first use, and the most kept as commits extend the window. */ +const ENTRY_WINDOW_LOAD = 256; +const ENTRY_WINDOW_MAX = 1024; + +const requestKey = (conversationId: ConversationId, requestId: string): string => `${conversationId}:${requestId}`; + +/** Index of the first element greater than `value` in an ascending array. */ +const upperBound = (values: readonly number[], value: number): number => { + let low = 0; + let high = values.length; + while (low < high) { + const middle = (low + high) >>> 1; + if (values[middle] <= value) low = middle + 1; + else high = middle; + } + return low; +}; + /** A Postgres array literal of record ids, for `= ANY($n::bigint[])` with the adapter's scalar parameter type. */ const arrayLiteral = (ids: readonly number[]): string => `{${[...new Set(ids)].join(",")}}`; @@ -207,6 +234,18 @@ export class PostgresStorage implements Storage { * only ever inserted, so a record stays valid until a commit of this storage writes that conversation. */ private readonly conversations = new Map(); + // Read-through caches with the same justification: this storage is the Session's only writer, so a record it + // has read stays valid until it writes that record, and every write passes through `commit` below. Values are + // kept JSON-encoded so a caller gets a detached object each time, as from the database. + private readonly tasks = new Map(); + private readonly submissions = new Map(); + private readonly submissionsByRequest = new Map(); + private readonly entries = new Map(); + private readonly entryWindows = new Map(); + private readonly headMarkers = new Map(); + private readonly documentStates = new Map(); + private documentScans = new Map(); + private documentAddresses = new Map(); private admitted = 0; private drained: (() => void) | undefined; @@ -316,7 +355,7 @@ export class PostgresStorage implements Storage { } const committedSeq = seqFromNumber(Number(state.session.next_seq)); this.checkGlobalIds(writes, state.ids); - await this.checkDocumentActions(this.pool, documentActions, state.documents); + await this.checkDocumentActions(this.pool, documentActions, state.documents, state.versions); const documentWrites = await this.prepareDocumentWrites( this.pool, documentActions, @@ -334,9 +373,7 @@ export class PostgresStorage implements Storage { throw error; } this.nextId = Math.max(this.nextId, candidateNextId); - for (const write of writes) { - if (write.type === "conversation") this.conversations.set(write.value.id, write.value); - } + this.absorb(writes, documentActions, documentWrites, committedSeq); return committedSeq; }); } @@ -351,6 +388,7 @@ export class PostgresStorage implements Storage { const id = writeId(write); if (id !== undefined) ids.add(id); } + const deltas = [...actions].filter(([, action]) => action.content?.kind === "delta").map(([id]) => id); const row = ( await this.pool.query<{ readonly state: string }>( `SELECT json_build_object( @@ -358,9 +396,12 @@ export class PostgresStorage implements Storage { 'ids', (SELECT coalesce(json_agg(json_build_object('id', r.id, 'type', r.record_type)), '[]'::json) FROM ${this.tables.recordIds} r WHERE r.session_id = $1 AND r.id = ANY($2::bigint[])), 'documents', (SELECT coalesce(json_agg(json_build_object('id', d.id, 'record', d.record)), '[]'::json) - FROM ${this.tables.documents} d WHERE d.session_id = $1 AND d.id = ANY($3::bigint[])) + FROM ${this.tables.documents} d WHERE d.session_id = $1 AND d.id = ANY($3::bigint[])), + 'versions', (SELECT coalesce(json_agg(json_build_object('id', v.document_id, 'version', v.version)), '[]'::json) + FROM (SELECT DISTINCT ON (document_id) document_id, version FROM ${this.tables.revisions} + WHERE session_id = $1 AND document_id = ANY($4::bigint[]) ORDER BY document_id, seq DESC) v) )::text AS state`, - [this.session, arrayLiteral([...ids]), arrayLiteral([...actions.keys()])], + [this.session, arrayLiteral([...ids]), arrayLiteral([...actions.keys()]), arrayLiteral(deltas)], ) ).rows[0]; const raw = parseJson<{ @@ -371,12 +412,15 @@ export class PostgresStorage implements Storage { } | null; readonly ids: readonly { readonly id: number | string; readonly type: TableName }[]; readonly documents: readonly { readonly id: number | string; readonly record: DocumentRecord }[]; + readonly versions: readonly { readonly id: number | string; readonly version: number }[]; }>(row.state); const existingIds = new Map, TableName>(); for (const entry of raw.ids) existingIds.set(idFromNumber>(Number(entry.id)), entry.type); const documents = new Map(); for (const entry of raw.documents) documents.set(idFromNumber(Number(entry.id)), entry.record); - return { session: raw.session ?? undefined, ids: existingIds, documents }; + const versions = new Map(); + for (const entry of raw.versions) versions.set(idFromNumber(Number(entry.id)), entry.version); + return { session: raw.session ?? undefined, ids: existingIds, documents, versions }; } async mintId>(): Promise { @@ -460,14 +504,19 @@ export class PostgresStorage implements Storage { conversation = await this.readConversation(conversationId); if (conversation === undefined) throw new Error(`Unknown conversation: ${conversationId}`); } - const row = ( - await this.pool.query( - `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} WHERE session_id = $1 AND id = $2`, - [this.session, id], - ) - ).rows[0]; - if (row === undefined) return undefined; - const entry = parseJson(row.record); + let stored = this.entries.get(id); + if (stored === undefined) { + const row = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} WHERE session_id = $1 AND id = $2`, + [this.session, id], + ) + ).rows[0]; + if (row === undefined) return undefined; + stored = { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }; + this.entries.set(id, stored); + } + const entry = parseJson(stored.record); if (conversation !== undefined) { let upperEntryId = Number.POSITIVE_INFINITY; while (conversation.id !== entry.conversationId) { @@ -477,7 +526,7 @@ export class PostgresStorage implements Storage { } if (entry.id > upperEntryId) return undefined; } - return { entry, commitSeq: seqFromNumber(Number(row.commit_seq)) }; + return { entry, commitSeq: stored.commitSeq }; } private async readLatestHeadMarker( @@ -486,23 +535,39 @@ export class PostgresStorage implements Storage { ): Promise<(EntryRecord & { readonly head: EntryId }) | undefined> { let conversation = await this.readConversation(conversationId); if (conversation === undefined) throw new Error(`Unknown conversation: ${conversationId}`); - let upper: number | undefined = atOrBeforeEntryId; + let upper = atOrBeforeEntryId ?? Number.POSITIVE_INFINITY; while (true) { - const params: PostgresValue[] = [this.session, conversation.id]; - const bound = upper === undefined ? "" : ` AND id <= $${params.push(upper)}`; - const [record] = await this.records( - `SELECT record::text AS record FROM ${this.tables.entries} - WHERE session_id = $1 AND conversation_id = $2 AND head IS NOT NULL${bound} - ORDER BY id DESC LIMIT 1`, - params, - ); - if (record !== undefined) return parseJson(record); + const ids = await this.headMarkerIds(conversation.id); + const index = upperBound(ids, upper) - 1; + if (index >= 0) + return parseJson(this.entries.get(ids[index])!.record); if (conversation.parent === undefined) return undefined; - upper = upper === undefined ? conversation.parent.at : Math.min(upper, conversation.parent.at); + upper = Math.min(upper, conversation.parent.at); conversation = (await this.readConversation(conversation.parent.conversationId))!; } } + /** Every head marker of one conversation, ascending; loaded once, extended by commits. */ + private async headMarkerIds(conversationId: ConversationId): Promise { + let ids = this.headMarkers.get(conversationId); + if (ids !== undefined) return ids; + const rows = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} + WHERE session_id = $1 AND conversation_id = $2 AND head IS NOT NULL ORDER BY id`, + [this.session, conversationId], + ) + ).rows; + ids = []; + for (const row of rows) { + const entry = parseJson(row.record); + this.entries.set(entry.id, { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }); + ids.push(entry.id); + } + this.headMarkers.set(conversationId, ids); + return ids; + } + private async readEntries( query: EntryQuery, limit: number, @@ -515,16 +580,8 @@ export class PostgresStorage implements Storage { if (after !== undefined) upper = Math.min(upper ?? Number.MAX_SAFE_INTEGER, after - 1); const values: EntryRecord[] = []; while (true) { - const params: PostgresValue[] = [this.session, conversation.id]; - const clauses = ["session_id = $1", "conversation_id = $2"]; - if (query.minEntryId !== undefined) clauses.push(`id >= $${params.push(query.minEntryId)}`); - if (upper !== undefined) clauses.push(`id <= $${params.push(upper)}`); - const rows = await this.records( - `SELECT record::text AS record FROM ${this.tables.entries} WHERE ${clauses.join(" AND ")} - ORDER BY id DESC LIMIT $${params.push(limit + 1 - values.length)}`, - params, - ); - values.push(...rows.map(parseJson)); + const rows = await this.entriesBelow(conversation.id, query.minEntryId, upper, limit + 1 - values.length); + values.push(...rows); if (values.length > limit || conversation.parent === undefined) break; upper = upper === undefined ? conversation.parent.at : Math.min(upper, conversation.parent.at); if (query.minEntryId !== undefined && upper < query.minEntryId) break; @@ -533,8 +590,164 @@ export class PostgresStorage implements Storage { return page(values, limit); } + /** + * Up to `limit` of one conversation's own entries with ids in `[min, upper]`, newest first. Served from the + * conversation's window when it covers the range; otherwise the part below the window is read, and extends + * the window when it is contiguous with it. + */ + private async entriesBelow( + conversationId: ConversationId, + min: number | undefined, + upper: number | undefined, + limit: number, + ): Promise { + const window = await this.entryWindow(conversationId); + const out: EntryRecord[] = []; + let index = upper === undefined ? window.ids.length - 1 : upperBound(window.ids, upper) - 1; + for (; index >= 0 && out.length < limit; index--) { + const id = window.ids[index]; + if (min !== undefined && id < min) return out; + out.push(parseJson(this.entries.get(id)!.record)); + } + if (out.length >= limit || window.complete) return out; + const oldest = window.ids[0]; + const belowWindow = upper === undefined || oldest === undefined || upper >= oldest; + const dbUpper = belowWindow && oldest !== undefined ? oldest - 1 : upper; + const params: PostgresValue[] = [this.session, conversationId]; + const clauses = ["session_id = $1", "conversation_id = $2"]; + if (min !== undefined) clauses.push(`id >= $${params.push(min)}`); + if (dbUpper !== undefined) clauses.push(`id <= $${params.push(dbUpper)}`); + const wanted = limit - out.length; + const rows = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} WHERE ${clauses.join(" AND ")} + ORDER BY id DESC LIMIT $${params.push(wanted)}`, + params, + ) + ).rows; + const fetched = rows.map((row) => ({ entry: parseJson(row.record), row })); + if (belowWindow) { + // Contiguous with the window: keep these so a later scan of the same range needs no read. Fewer rows + // than asked for, with no lower bound cutting them off, means the conversation starts here. + for (const { entry, row } of fetched) { + this.entries.set(entry.id, { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }); + } + window.ids.unshift(...fetched.map(({ entry }) => entry.id).reverse()); + if (fetched.length < wanted && min === undefined) window.complete = true; + } + out.push(...fetched.map(({ entry }) => entry)); + return out; + } + + /** The conversation's entry window, loaded from its newest entries on first use. */ + private async entryWindow(conversationId: ConversationId): Promise { + let window = this.entryWindows.get(conversationId); + if (window !== undefined) return window; + const rows = ( + await this.pool.query( + `SELECT record::text AS record, commit_seq FROM ${this.tables.entries} + WHERE session_id = $1 AND conversation_id = $2 ORDER BY id DESC LIMIT $3`, + [this.session, conversationId, ENTRY_WINDOW_LOAD + 1], + ) + ).rows; + const ids: EntryId[] = []; + for (const row of rows.slice(0, ENTRY_WINDOW_LOAD).reverse()) { + const entry = parseJson(row.record); + this.entries.set(entry.id, { record: row.record, commitSeq: seqFromNumber(Number(row.commit_seq)) }); + ids.push(entry.id); + } + window = { ids, complete: rows.length <= ENTRY_WINDOW_LOAD }; + this.entryWindows.set(conversationId, window); + return window; + } + + /** Reflects a committed batch in the caches. Called only after the write statement reported success. */ + private absorb( + writes: readonly StorageWrite[], + actions: ReadonlyMap, + documents: DocumentWrites, + seq: Seq, + ): void { + for (const write of writes) { + switch (write.type) { + case "conversation": + this.conversations.set(write.value.id, parseJson(encodeJson(write.value))); + break; + case "task": + this.tasks.set(write.value.id, encodeJson(write.value)); + break; + case "submission": + this.rememberSubmission(encodeJson(write.value)); + break; + case "entry": { + const record = encodeJson(write.value); + this.entries.set(write.value.id, { record, commitSeq: seq }); + const window = this.entryWindows.get(write.value.conversationId); + if (window !== undefined) { + window.ids.push(write.value.id); + if (window.ids.length > ENTRY_WINDOW_MAX) { + const heads = new Set(this.headMarkers.get(write.value.conversationId) ?? []); + for (const id of window.ids.splice(0, window.ids.length - ENTRY_WINDOW_MAX)) { + if (!heads.has(id)) this.entries.delete(id); + } + window.complete = false; + } + } + if (write.value.head !== undefined) + this.headMarkers.get(write.value.conversationId)?.push(write.value.id); + break; + } + default: + break; + } + } + if (actions.size === 0) return; + for (const [id, { record, content }] of documents.applied) { + if (record.retiredAt !== undefined) { + this.documentStates.delete(id); + continue; + } + const state = this.documentStates.get(id); + if (content?.kind === "base") { + this.documentStates.set(id, { + record: encodeJson(record), + version: content.version, + value: parseJson(encodeJson(content.value)), + deltasSinceBase: 0, + }); + } else if (content?.kind === "delta" && state !== undefined && state.version === content.version) { + state.value = apply(state.value, parseJson(encodeJson(content.ops))) as JsonObject; + state.deltasSinceBase++; + } else if (content !== undefined) { + this.documentStates.delete(id); + } + } + // Creating or retiring a document changes what an address or scope scan resolves to. + if (documents.inserts.length > 0 || documents.retirements.length > 0) { + this.documentScans = new Map(); + this.documentAddresses = new Map(); + } + } + + private rememberSubmission(record: string): void { + const submission = parseJson(record); + this.submissions.set(submission.id, record); + if (submission.requestId !== undefined) { + this.submissionsByRequest.set(requestKey(submission.conversationId, submission.requestId), submission.id); + } + } + task(id: TaskId, _context: Context): Promise { - return this.admit(() => this.recordById(this.tables.tasks, id)); + return this.admit(async () => { + const cached = this.tasks.get(id); + if (cached !== undefined) return parseJson(cached); + const [record] = await this.records( + `SELECT record::text AS record FROM ${this.tables.tasks} WHERE session_id = $1 AND id = $2`, + [this.session, id], + ); + if (record !== undefined) this.tasks.set(id, record); + return record === undefined ? undefined : parseJson(record); + }); } scanTasks( @@ -564,7 +777,16 @@ export class PostgresStorage implements Storage { } submission(id: SubmissionId, _context: Context): Promise { - return this.admit(() => this.recordById(this.tables.submissions, id)); + return this.admit(async () => { + const cached = this.submissions.get(id); + if (cached !== undefined) return parseJson(cached); + const [record] = await this.records( + `SELECT record::text AS record FROM ${this.tables.submissions} WHERE session_id = $1 AND id = $2`, + [this.session, id], + ); + if (record !== undefined) this.rememberSubmission(record); + return record === undefined ? undefined : parseJson(record); + }); } scanSubmissions( @@ -594,17 +816,26 @@ export class PostgresStorage implements Storage { _context: Context, ): Promise { return this.admit(async () => { + const cachedId = this.submissionsByRequest.get(requestKey(conversationId, requestId)); + const cached = cachedId === undefined ? undefined : this.submissions.get(cachedId); + if (cached !== undefined) return parseJson(cached); const [record] = await this.records( `SELECT record::text AS record FROM ${this.tables.submissions} WHERE session_id = $1 AND conversation_id = $2 AND request_id = $3`, [this.session, conversationId, encodeIndexedString(requestId)], ); + if (record !== undefined) this.rememberSubmission(record); return record === undefined ? undefined : parseJson(record); }); } findDocument(address: DocumentAddress, at: DocumentPoint, _context: Context): Promise { return this.admit(async () => { + const key = at === "current" ? addressKey(address) : undefined; + if (key !== undefined) { + const cached = this.documentAddresses.get(key); + if (cached !== undefined) return parseJson(cached); + } const parts = addressParts(address); const params: PostgresValue[] = [ this.session, @@ -625,17 +856,38 @@ export class PostgresStorage implements Storage { ORDER BY created_at DESC LIMIT 1`, params, ); + if (key !== undefined && record !== undefined) this.documentAddresses.set(key, record); return record === undefined ? undefined : parseJson(record); }); } document(id: DocumentId, at: DocumentPoint, _context: Context): Promise { - // The record and revision queries must observe one committed state; a commit between them can replace the base. - return this.admit(() => - inTransaction(this.pool, "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY", (client) => + return this.admit(async () => { + if (at === "current") { + const state = this.documentStates.get(id); + if (state !== undefined) { + return { + record: parseJson(state.record), + version: state.version, + value: parseJson(encodeJson(state.value)), + deltasSinceBase: state.deltasSinceBase, + }; + } + } + // The record and revision queries must observe one committed state; a commit between them can replace the base. + const stored = await inTransaction(this.pool, "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY", (client) => this.materializeDocument(client, id, at), - ), - ); + ); + if (at === "current" && stored !== undefined) { + this.documentStates.set(id, { + record: encodeJson(stored.record), + version: stored.version, + value: parseJson(encodeJson(stored.value)), + deltasSinceBase: stored.deltasSinceBase, + }); + } + return stored; + }); } scanDocuments( @@ -646,7 +898,16 @@ export class PostgresStorage implements Storage { ): Promise> { return this.admit(async () => { const scope = scopeColumns(query.scope); - const params: PostgresValue[] = [this.session, scope.scopeKind, scope.ownerId, cursorId(cursor) ?? -1]; + const after = cursorId(cursor) ?? -1; + const key = query.at === "current" ? `${scope.scopeKind}:${scope.ownerId}:${query.kind ?? ""}` : undefined; + if (key !== undefined) { + const cached = this.documentScans.get(key); + if (cached !== undefined) { + const records = cached.map(parseJson).filter((record) => record.id > after); + return page(records.slice(0, limit + 1), limit); + } + } + const params: PostgresValue[] = [this.session, scope.scopeKind, scope.ownerId, after]; const clauses = ["session_id = $1", "scope_kind = $2", "owner_id = $3", "id > $4"]; if (query.kind !== undefined) clauses.push(`kind = $${params.push(encodeIndexedString(query.kind))}`); if (query.at === "current") { @@ -657,12 +918,16 @@ export class PostgresStorage implements Storage { `(retired_at IS NULL OR retired_at > $${params.length})`, ); } + // A current scan is cached whole: it is read from its start without a page limit, and a scope's live + // documents are few. + const whole = key !== undefined && after === -1; const rows = await this.records( `SELECT record::text AS record FROM ${this.tables.documents} WHERE ${clauses.join(" AND ")} - ORDER BY id LIMIT $${params.push(limit + 1)}`, + ORDER BY id${whole ? "" : ` LIMIT $${params.push(limit + 1)}`}`, params, ); - return page(rows.map(parseJson), limit); + if (whole) this.documentScans.set(key!, rows); + return page(rows.slice(0, limit + 1).map(parseJson), limit); }); } @@ -853,6 +1118,7 @@ export class PostgresStorage implements Storage { executor: PostgresQueryable, actions: ReadonlyMap, existingByID: ReadonlyMap, + versions: ReadonlyMap, ): Promise { const liveCounts = new Map(); for (const [id, action] of actions) { @@ -864,15 +1130,9 @@ export class PostgresStorage implements Storage { if (action.create !== undefined && existing !== undefined) throw new Error(`Document ${id} already exists`); if (existing?.retiredAt !== undefined) throw new Error(`Document ${id} is retired`); if (action.content?.kind === "delta") { - const previous = ( - await executor.query<{ readonly version: number }>( - `SELECT version FROM ${this.tables.revisions} - WHERE session_id = $1 AND document_id = $2 ORDER BY seq DESC LIMIT 1`, - [this.session, id], - ) - ).rows[0]; + const previous = versions.get(id); if (previous === undefined) throw new Error(`Document ${id} delta has no base`); - if (previous.version !== action.content.version) { + if (previous !== action.content.version) { throw new Error(`Document ${id} version transition requires a base`); } } @@ -920,7 +1180,14 @@ export class PostgresStorage implements Storage { existing: ReadonlyMap, seq: Seq, ): Promise { - const writes: DocumentWrites = { claims: [], inserts: [], retirements: [], clearRevisions: [], revisions: [] }; + const writes: DocumentWrites = { + claims: [], + inserts: [], + retirements: [], + clearRevisions: [], + revisions: [], + applied: new Map(), + }; for (const [id, action] of actions) { let content = action.content; if (action.copy !== undefined) { @@ -976,6 +1243,7 @@ export class PostgresStorage implements Storage { const encodedContent = content.kind === "base" ? encodeJson(content.value) : encodeJson(content.ops); writes.revisions.push([this.session, id, seq, content.kind, content.version, encodedContent]); } + writes.applied.set(id, { record, content: retiredCurrentOnly ? undefined : content }); } return writes; } From 3293dc12157694bea8a9e7598b2b1ec40854033c Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:11:07 +0000 Subject: [PATCH 4/5] feat(durable): commit a Postgres batch in one round trip through a schema function A plpgsql function created by schema version 2 locks the session row, checks ownership at the expected sequence, validates the batch against the stored state, writes every table, and advances the session. Each rule raises the adapter's message with its own SQLSTATE, which the adapter maps back to its errors. Records and content travel inside the batch as JSON text and are cast on insert, because Postgres rejects lone-surrogate escapes in any function that decodes string values. Historical document reads are one statement, and a request key with no submission is cached. --- .../src/storage/postgres/migrations.ts | 164 ++++- .../durable/src/storage/postgres/storage.ts | 604 ++++++------------ 2 files changed, 366 insertions(+), 402 deletions(-) diff --git a/packages/durable/src/storage/postgres/migrations.ts b/packages/durable/src/storage/postgres/migrations.ts index 26870217c800..546e51df4d62 100644 --- a/packages/durable/src/storage/postgres/migrations.ts +++ b/packages/durable/src/storage/postgres/migrations.ts @@ -101,8 +101,170 @@ const INITIAL_SCHEMA: readonly string[] = [ "CREATE INDEX document_revisions_by_kind ON {schema}.document_revisions (session_id, document_id, kind, seq DESC)", ]; +// One commit is one call. The function locks the session row, checks that the caller still owns the session at +// the sequence it expects, validates the batch against the stored state, writes every table, and advances the +// session. Each rule raises the adapter's message with its own SQLSTATE so the adapter can map it back; see +// `commitErrorCodes` in storage.ts. +// +// Records and content travel inside the batch as JSON text and are cast to `json` on insert. Postgres accepts a +// lone UTF-16 surrogate escape when it stores `json` but rejects it in any function that decodes string values, +// such as `json_to_recordset` and `->>`; so nothing here decodes a record, and the current-only test matches the +// stored record's text. +const COMMIT_FUNCTION = `CREATE OR REPLACE FUNCTION {schema}.commit_batch( + p_session TEXT, p_owner TEXT, p_seq BIGINT, p_next_id BIGINT, p_batch JSON +) RETURNS BIGINT LANGUAGE plpgsql AS $fn$ +DECLARE + v_owner TEXT; + v_seq BIGINT; + v_id BIGINT; + v_type TEXT; + v_exists BOOLEAN; + v_retired BIGINT; + v_version INTEGER; + v_key TEXT; + v_live JSONB := '{}'::jsonb; + v_count INTEGER; + doc RECORD; + stored RECORD; + live RECORD; +BEGIN + SELECT owner, next_seq INTO v_owner, v_seq FROM {schema}.sessions WHERE session_id = p_session FOR UPDATE; + IF NOT FOUND THEN + RAISE EXCEPTION 'Durable Postgres session metadata is missing' USING ERRCODE = 'DU001'; + END IF; + IF v_owner <> p_owner OR v_seq <> p_seq THEN + RAISE EXCEPTION 'Durable Postgres session % was opened by another process', p_session USING ERRCODE = 'DU002'; + END IF; + + SELECT c.id, r.record_type INTO v_id, v_type + FROM json_to_recordset(p_batch->'claims') AS c(id BIGINT, type TEXT) + JOIN {schema}.record_ids r ON r.session_id = p_session AND r.id = c.id + WHERE c.type IN ('conversation', 'entry', 'document') OR r.record_type <> c.type + LIMIT 1; + IF FOUND THEN + RAISE EXCEPTION 'ID % already belongs to %', v_id, v_type USING ERRCODE = 'DU003'; + END IF; + + FOR doc IN SELECT * FROM json_to_recordset(p_batch->'documents') AS x( + id BIGINT, is_create BOOLEAN, retire BOOLEAN, kind TEXT, family SMALLINT, key_value TEXT, + scope_kind TEXT, owner_id BIGINT, content_kind TEXT, version INTEGER) + LOOP + SELECT x.retired_at, x.kind, x.family, x.key_value, x.scope_kind, x.owner_id + INTO stored FROM {schema}.documents x WHERE x.session_id = p_session AND x.id = doc.id; + v_exists := FOUND; + IF NOT doc.is_create AND NOT v_exists THEN + RAISE EXCEPTION 'Unknown document: %', doc.id USING ERRCODE = 'DU004'; + END IF; + IF doc.is_create AND v_exists THEN + RAISE EXCEPTION 'Document % already exists', doc.id USING ERRCODE = 'DU005'; + END IF; + IF v_exists AND stored.retired_at IS NOT NULL THEN + RAISE EXCEPTION 'Document % is retired', doc.id USING ERRCODE = 'DU006'; + END IF; + IF doc.content_kind = 'delta' THEN + SELECT r.version INTO v_version FROM {schema}.document_revisions r + WHERE r.session_id = p_session AND r.document_id = doc.id ORDER BY r.seq DESC LIMIT 1; + IF NOT FOUND THEN + RAISE EXCEPTION 'Document % delta has no base', doc.id USING ERRCODE = 'DU007'; + END IF; + IF v_version <> doc.version THEN + RAISE EXCEPTION 'Document % version transition requires a base', doc.id USING ERRCODE = 'DU008'; + END IF; + END IF; + IF doc.is_create OR doc.retire THEN + IF doc.is_create THEN + v_key := concat_ws('|', doc.kind, doc.scope_kind, doc.owner_id, doc.family, doc.key_value); + ELSE + v_key := concat_ws('|', stored.kind, stored.scope_kind, stored.owner_id, stored.family, stored.key_value); + END IF; + IF NOT v_live ? v_key THEN + IF doc.is_create THEN + SELECT count(*) INTO v_count FROM {schema}.documents x WHERE x.session_id = p_session + AND x.kind = doc.kind AND x.scope_kind = doc.scope_kind AND x.owner_id = doc.owner_id + AND x.family = doc.family AND x.key_value = doc.key_value AND x.retired_at IS NULL; + ELSE + SELECT count(*) INTO v_count FROM {schema}.documents x WHERE x.session_id = p_session + AND x.kind = stored.kind AND x.scope_kind = stored.scope_kind AND x.owner_id = stored.owner_id + AND x.family = stored.family AND x.key_value = stored.key_value AND x.retired_at IS NULL; + END IF; + v_count := LEAST(v_count, 1); + ELSE + v_count := (v_live->>v_key)::integer; + END IF; + IF doc.retire AND v_exists THEN v_count := v_count - 1; END IF; + IF doc.is_create AND NOT doc.retire THEN v_count := v_count + 1; END IF; + v_live := v_live || jsonb_build_object(v_key, v_count); + END IF; + END LOOP; + FOR live IN SELECT value FROM jsonb_each_text(v_live) LOOP + IF live.value::integer > 1 THEN + RAISE EXCEPTION 'Document address already has a current incarnation' USING ERRCODE = 'DU009'; + END IF; + END LOOP; + + INSERT INTO {schema}.record_ids (session_id, id, record_type) + SELECT p_session, c.id, c.type FROM json_to_recordset(p_batch->'claims') AS c(id BIGINT, type TEXT) + ON CONFLICT (session_id, id) DO NOTHING; + INSERT INTO {schema}.conversations (session_id, id, owner_conversation_id, owner_task_id, record) + SELECT p_session, x.id, x.owner_conversation_id, x.owner_task_id, x.record::json + FROM json_to_recordset(p_batch->'conversations') + AS x(id BIGINT, owner_conversation_id BIGINT, owner_task_id BIGINT, record TEXT); + INSERT INTO {schema}.entries (session_id, id, conversation_id, head, commit_seq, record) + SELECT p_session, x.id, x.conversation_id, x.head, p_seq, x.record::json + FROM json_to_recordset(p_batch->'entries') AS x(id BIGINT, conversation_id BIGINT, head BIGINT, record TEXT); + INSERT INTO {schema}.tasks (session_id, id, conversation_id, kind, status, abort_requested, background, record) + SELECT p_session, x.id, x.conversation_id, x.kind, x.status, x.abort_requested, x.background, x.record::json + FROM json_to_recordset(p_batch->'tasks') AS x(id BIGINT, conversation_id BIGINT, kind TEXT, status TEXT, + abort_requested BOOLEAN, background BOOLEAN, record TEXT) + ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, kind = excluded.kind, + status = excluded.status, abort_requested = excluded.abort_requested, background = excluded.background, + record = excluded.record; + INSERT INTO {schema}.submissions (session_id, id, conversation_id, request_id, status, record) + SELECT p_session, x.id, x.conversation_id, x.request_id, x.status, x.record::json + FROM json_to_recordset(p_batch->'submissions') + AS x(id BIGINT, conversation_id BIGINT, request_id TEXT, status TEXT, record TEXT) + ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, + request_id = excluded.request_id, status = excluded.status, record = excluded.record; + + -- A retired record gains retiredAt by appending to its text, which keeps the rest of the record verbatim. + UPDATE {schema}.documents x SET retired_at = p_seq, + record = (left(x.record::text, length(x.record::text) - 1) || ',"retiredAt":' || p_seq || '}')::json + FROM json_to_recordset(p_batch->'documents') AS d(id BIGINT, is_create BOOLEAN, retire BOOLEAN) + WHERE x.session_id = p_session AND x.id = d.id AND d.retire AND NOT d.is_create; + INSERT INTO {schema}.documents + (session_id, id, kind, family, key_value, scope_kind, owner_id, created_at, retired_at, record) + SELECT p_session, d.id, d.kind, d.family, d.key_value, d.scope_kind, d.owner_id, p_seq, + CASE WHEN d.retire THEN p_seq END, d.record::json + FROM json_to_recordset(p_batch->'documents') AS d(id BIGINT, is_create BOOLEAN, retire BOOLEAN, kind TEXT, + family SMALLINT, key_value TEXT, scope_kind TEXT, owner_id BIGINT, record TEXT) + WHERE d.is_create; + -- A current-only document keeps no history: a base replaces its revisions, and a retirement clears them. The + -- record is matched as text; its only object values are the adapter's own, so the pair cannot appear elsewhere. + DELETE FROM {schema}.document_revisions r + USING json_to_recordset(p_batch->'documents') AS d(id BIGINT, retire BOOLEAN, content_kind TEXT) + JOIN {schema}.documents x ON x.session_id = p_session AND x.id = d.id + WHERE r.session_id = p_session AND r.document_id = d.id + AND (x.scope_kind <> 'conversation' OR x.record::text LIKE '%"history":"latest"%') + AND (d.retire OR d.content_kind = 'base'); + INSERT INTO {schema}.document_revisions (session_id, document_id, seq, kind, version, content) + SELECT p_session, d.id, p_seq, d.content_kind, d.version, d.content::json + FROM json_to_recordset(p_batch->'documents') + AS d(id BIGINT, retire BOOLEAN, content_kind TEXT, version INTEGER, content TEXT) + JOIN {schema}.documents x ON x.session_id = p_session AND x.id = d.id + WHERE d.content_kind IS NOT NULL + AND NOT (d.retire AND (x.scope_kind <> 'conversation' OR x.record::text LIKE '%"history":"latest"%')); + + UPDATE {schema}.sessions SET next_id = GREATEST(next_id, p_next_id), next_seq = p_seq + 1 + WHERE session_id = p_session; + RETURN p_seq; +END +$fn$`; + /** Immutable, ordered schema history. Append new migrations after the initial schema ships. */ -export const POSTGRES_MIGRATIONS: readonly PostgresMigration[] = [{ version: 1, statements: INITIAL_SCHEMA }]; +export const POSTGRES_MIGRATIONS: readonly PostgresMigration[] = [ + { version: 1, statements: INITIAL_SCHEMA }, + { version: 2, statements: [COMMIT_FUNCTION] }, +]; export const CURRENT_POSTGRES_SCHEMA_VERSION = POSTGRES_MIGRATIONS.at(-1)?.version ?? 0; diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts index 3efc943ca024..7ef87f3d4c40 100644 --- a/packages/durable/src/storage/postgres/storage.ts +++ b/packages/durable/src/storage/postgres/storage.ts @@ -31,50 +31,43 @@ import type { TaskQuery, TaskRecord, } from "../../types.ts"; -import { - inTransaction, - type PostgresPool, - type PostgresPoolClient, - type PostgresQueryable, - type PostgresValue, -} from "./database.ts"; +import type { PostgresPool, PostgresPoolClient, PostgresQueryable, PostgresValue } from "./database.ts"; import { applyPostgresMigrations, DEFAULT_POSTGRES_SCHEMA, quoteIdentifier } from "./migrations.ts"; type StoredTask = TaskRecord; type TableName = "conversation" | "entry" | "task" | "submission" | "document"; type JsonRow = { readonly record: string }; type EntryJsonRow = { readonly record: string; readonly commit_seq: string }; -type RevisionRow = { - readonly seq: string; - readonly kind: DocumentContent["kind"]; - readonly version: number; - readonly content: string; -}; -type IdRow = { readonly id: string }; type MetadataRow = { readonly next_id: string; readonly next_seq: string }; /** A conversation's own entries over one contiguous id range ending at its newest entry, newest-last. */ type EntryWindow = { ids: EntryId[]; complete: boolean }; type DocumentState = { readonly record: string; readonly version: number; value: JsonObject; deltasSinceBase: number }; -type CommitState = { - readonly session: - | { readonly next_id: number | string; readonly next_seq: number | string; readonly owner: string } - | undefined; - readonly ids: ReadonlyMap, TableName>; - readonly documents: ReadonlyMap; - readonly versions: ReadonlyMap; -}; type DocumentWrites = { readonly claims: DocumentId[]; - readonly inserts: PostgresValue[][]; - readonly retirements: PostgresValue[][]; - readonly clearRevisions: DocumentId[]; - readonly revisions: PostgresValue[][]; - /** The record as written and the content applied, per document, for the caches. */ + readonly rows: unknown[]; + /** What each document action means for the caches once the commit succeeds. */ readonly applied: Map< DocumentId, - { readonly record: DocumentRecord; readonly content: DocumentContent | undefined } + { + readonly created: DocumentRecord | undefined; + readonly retired: boolean; + readonly content: DocumentContent | undefined; + } >; }; + +/** SQLSTATEs the commit function raises, each for one adapter error. */ +const commitErrorCodes: Readonly> = { + DU001: true, + DU002: true, + DU003: true, + DU004: true, + DU005: true, + DU006: true, + DU007: true, + DU008: true, + DU009: true, +}; type DocumentAction = { create?: DocumentCreate; copy?: Extract["source"]; @@ -162,15 +155,6 @@ const isAliveAt = (record: DocumentRecord, at: DocumentPoint): boolean => { const isCurrentOnly = (record: DocumentRecord): boolean => record.scope.kind !== "conversation" || record.history === "latest"; -// Column types of each table's insert list, for casting VALUES parameters. -const RECORD_ID_TYPES = ["text", "bigint", "text"]; -const CONVERSATION_TYPES = ["text", "bigint", "bigint", "bigint", "json"]; -const ENTRY_TYPES = ["text", "bigint", "bigint", "bigint", "bigint", "json"]; -const TASK_TYPES = ["text", "bigint", "bigint", "text", "text", "boolean", "boolean", "json"]; -const SUBMISSION_TYPES = ["text", "bigint", "bigint", "text", "text", "json"]; -const DOCUMENT_TYPES = ["text", "bigint", "text", "smallint", "text", "text", "bigint", "bigint", "bigint", "json"]; -const REVISION_TYPES = ["text", "bigint", "bigint", "text", "integer", "json"]; - /** Entries loaded per conversation on first use, and the most kept as commits extend the window. */ const ENTRY_WINDOW_LOAD = 256; const ENTRY_WINDOW_MAX = 1024; @@ -189,9 +173,6 @@ const upperBound = (values: readonly number[], value: number): number => { return low; }; -/** A Postgres array literal of record ids, for `= ANY($n::bigint[])` with the adapter's scalar parameter type. */ -const arrayLiteral = (ids: readonly number[]): string => `{${[...new Set(ids)].join(",")}}`; - const writeId = (write: StorageWrite): Id | undefined => { switch (write.type) { case "conversation": @@ -226,6 +207,9 @@ export class PostgresStorage implements Storage { private readonly owner: string; private readonly tables: Tables; private nextId: number; + /** The sequence the next commit takes. This storage is the session's only writer, so it advances by one per commit. */ + private nextSeq: number; + private readonly commitFunction: string; private closed = false; private lost: Error | undefined; private closing: Promise | undefined; @@ -239,7 +223,9 @@ export class PostgresStorage implements Storage { // kept JSON-encoded so a caller gets a detached object each time, as from the database. private readonly tasks = new Map(); private readonly submissions = new Map(); - private readonly submissionsByRequest = new Map(); + // A request key with no submission is remembered too: nothing else can create one, so a miss stays a miss + // until this storage commits that submission. + private readonly submissionsByRequest = new Map(); private readonly entries = new Map(); private readonly entryWindows = new Map(); private readonly headMarkers = new Map(); @@ -253,7 +239,7 @@ export class PostgresStorage implements Storage { pool: PostgresPool, lockClient: PostgresPoolClient, onLockError: (error: Error) => void, - options: { session: string; schema: string; lockKey: string; owner: string; nextId: number }, + options: { session: string; schema: string; lockKey: string; owner: string; nextId: number; nextSeq: number }, ) { this.pool = pool; this.lockClient = lockClient; @@ -262,7 +248,9 @@ export class PostgresStorage implements Storage { this.lockKey = options.lockKey; this.owner = options.owner; this.nextId = options.nextId; + this.nextSeq = options.nextSeq; const schema = quoteIdentifier(options.schema); + this.commitFunction = `${schema}.commit_batch`; this.tables = { sessions: `${schema}.sessions`, recordIds: `${schema}.record_ids`, @@ -319,6 +307,7 @@ export class PostgresStorage implements Storage { lockKey: key, owner, nextId: Number(metadata.next_id), + nextSeq: Number(metadata.next_seq), }); if (lockError !== undefined) storage.loseLock(lockError); return storage; @@ -342,85 +331,107 @@ export class PostgresStorage implements Storage { return this.admit(async () => { const documentActions = this.prepareDocumentActions(writes); const candidateNextId = this.candidateNextId(writes); - // Two round trips: one read of everything the checks need, then one statement that writes the whole - // batch. A single statement is atomic by itself, so no transaction brackets it; the write is gated on - // this storage still owning the session at the sequence the read observed, which is what the row lock - // used to guarantee. - const state = await this.readCommitState(writes, documentActions); - if (state.session === undefined) throw new Error("Durable Postgres session metadata is missing"); - if (state.session.owner !== this.owner) { - const error = new Error(`Durable Postgres session ${this.session} was opened by another process`); - this.loseLock(error); - throw error; - } - const committedSeq = seqFromNumber(Number(state.session.next_seq)); - this.checkGlobalIds(writes, state.ids); - await this.checkDocumentActions(this.pool, documentActions, state.documents, state.versions); - const documentWrites = await this.prepareDocumentWrites( - this.pool, - documentActions, - state.documents, - committedSeq, - ); - const written = await this.writeBatch(this.pool, writes, documentWrites, committedSeq, { - nextId: Math.max(Number(state.session.next_id), candidateNextId), - nextSeq: committedSeq + 1, - expectedSeq: committedSeq, - }); - if (!written) { - const error = new Error(`Durable Postgres session ${this.session} was opened by another process`); - this.loseLock(error); - throw error; + this.checkBatchIds(writes); + const seq = seqFromNumber(this.nextSeq); + const documentWrites = await this.prepareDocumentWrites(documentActions, seq); + // One round trip: the commit function locks the session row, checks that this storage still owns it + // at the sequence it expects, validates the batch against the stored state, and writes every table. + try { + await this.pool.query(`SELECT ${this.commitFunction}($1, $2, $3, $4, $5::json)`, [ + this.session, + this.owner, + seq, + candidateNextId, + this.encodeBatch(writes, documentWrites), + ]); + } catch (error) { + throw this.commitError(error); } + this.nextSeq = seq + 1; this.nextId = Math.max(this.nextId, candidateNextId); - this.absorb(writes, documentActions, documentWrites, committedSeq); - return committedSeq; + this.absorb(writes, documentActions, documentWrites, seq); + return seq; }); } - /** Everything a commit checks before writing, read in one statement: the session row, the batch's claimed ids, and its documents. */ - private async readCommitState( - writes: readonly StorageWrite[], - actions: ReadonlyMap, - ): Promise { - const ids = new Set(); + /** Maps an error raised by the commit function to the adapter's own errors; other errors pass through. */ + private commitError(error: unknown): unknown { + const code = (error as { readonly code?: unknown }).code; + const message = (error as { readonly message?: unknown }).message; + if (typeof code !== "string" || typeof message !== "string" || !(code in commitErrorCodes)) return error; + if (code === "DU002") { + const lost = new Error(message); + this.loseLock(lost); + return lost; + } + return new Error(message); + } + + /** The batch as the commit function reads it. Records and content go as JSON text, cast to `json` on insert. */ + private encodeBatch(writes: readonly StorageWrite[], documents: DocumentWrites): string { + const claims = new Map, TableName>(); + const conversations: unknown[] = []; + const entries: unknown[] = []; + const tasks = new Map, unknown>(); + const submissions = new Map, unknown>(); for (const write of writes) { - const id = writeId(write); - if (id !== undefined) ids.add(id); + switch (write.type) { + case "conversation": + claims.set(write.value.id, "conversation"); + conversations.push({ + id: write.value.id, + owner_conversation_id: write.value.owner?.conversationId ?? null, + owner_task_id: write.value.owner?.taskId ?? null, + record: encodeJson(write.value), + }); + break; + case "entry": + claims.set(write.value.id, "entry"); + entries.push({ + id: write.value.id, + conversation_id: write.value.conversationId, + head: write.value.head ?? null, + record: encodeJson(write.value), + }); + break; + case "task": + claims.set(write.value.id, "task"); + tasks.set(write.value.id, { + id: write.value.id, + conversation_id: write.value.conversationId, + kind: encodeIndexedString(write.value.kind), + status: write.value.state.status, + abort_requested: write.value.abortRequested, + background: write.value.background, + record: encodeJson(write.value), + }); + break; + case "submission": + claims.set(write.value.id, "submission"); + submissions.set(write.value.id, { + id: write.value.id, + conversation_id: write.value.conversationId, + request_id: write.value.requestId === undefined ? null : encodeIndexedString(write.value.requestId), + status: write.value.status, + record: encodeJson(write.value), + }); + break; + case "document.create": + case "document.copy": + case "document.change": + case "document.retire": + break; + } } - const deltas = [...actions].filter(([, action]) => action.content?.kind === "delta").map(([id]) => id); - const row = ( - await this.pool.query<{ readonly state: string }>( - `SELECT json_build_object( - 'session', (SELECT row_to_json(s) FROM ${this.tables.sessions} s WHERE s.session_id = $1), - 'ids', (SELECT coalesce(json_agg(json_build_object('id', r.id, 'type', r.record_type)), '[]'::json) - FROM ${this.tables.recordIds} r WHERE r.session_id = $1 AND r.id = ANY($2::bigint[])), - 'documents', (SELECT coalesce(json_agg(json_build_object('id', d.id, 'record', d.record)), '[]'::json) - FROM ${this.tables.documents} d WHERE d.session_id = $1 AND d.id = ANY($3::bigint[])), - 'versions', (SELECT coalesce(json_agg(json_build_object('id', v.document_id, 'version', v.version)), '[]'::json) - FROM (SELECT DISTINCT ON (document_id) document_id, version FROM ${this.tables.revisions} - WHERE session_id = $1 AND document_id = ANY($4::bigint[]) ORDER BY document_id, seq DESC) v) - )::text AS state`, - [this.session, arrayLiteral([...ids]), arrayLiteral([...actions.keys()]), arrayLiteral(deltas)], - ) - ).rows[0]; - const raw = parseJson<{ - readonly session: { - readonly next_id: number | string; - readonly next_seq: number | string; - readonly owner: string; - } | null; - readonly ids: readonly { readonly id: number | string; readonly type: TableName }[]; - readonly documents: readonly { readonly id: number | string; readonly record: DocumentRecord }[]; - readonly versions: readonly { readonly id: number | string; readonly version: number }[]; - }>(row.state); - const existingIds = new Map, TableName>(); - for (const entry of raw.ids) existingIds.set(idFromNumber>(Number(entry.id)), entry.type); - const documents = new Map(); - for (const entry of raw.documents) documents.set(idFromNumber(Number(entry.id)), entry.record); - const versions = new Map(); - for (const entry of raw.versions) versions.set(idFromNumber(Number(entry.id)), entry.version); - return { session: raw.session ?? undefined, ids: existingIds, documents, versions }; + for (const id of documents.claims) claims.set(id, "document"); + return encodeJson({ + claims: [...claims].map(([id, type]) => ({ id, type })), + conversations, + entries, + tasks: [...tasks.values()], + submissions: [...submissions.values()], + documents: documents.rows, + }); } async mintId>(): Promise { @@ -702,15 +713,16 @@ export class PostgresStorage implements Storage { } } if (actions.size === 0) return; - for (const [id, { record, content }] of documents.applied) { - if (record.retiredAt !== undefined) { + for (const [id, { created, retired, content }] of documents.applied) { + if (retired) { this.documentStates.delete(id); continue; } const state = this.documentStates.get(id); - if (content?.kind === "base") { + const record = created !== undefined ? encodeJson(created) : state?.record; + if (content?.kind === "base" && record !== undefined) { this.documentStates.set(id, { - record: encodeJson(record), + record, version: content.version, value: parseJson(encodeJson(content.value)), deltasSinceBase: 0, @@ -723,7 +735,11 @@ export class PostgresStorage implements Storage { } } // Creating or retiring a document changes what an address or scope scan resolves to. - if (documents.inserts.length > 0 || documents.retirements.length > 0) { + if ( + documents.rows.some( + (row) => (row as { is_create: boolean; retire: boolean }).is_create || (row as { retire: boolean }).retire, + ) + ) { this.documentScans = new Map(); this.documentAddresses = new Map(); } @@ -816,7 +832,9 @@ export class PostgresStorage implements Storage { _context: Context, ): Promise { return this.admit(async () => { - const cachedId = this.submissionsByRequest.get(requestKey(conversationId, requestId)); + const key = requestKey(conversationId, requestId); + const cachedId = this.submissionsByRequest.get(key); + if (cachedId === null) return undefined; const cached = cachedId === undefined ? undefined : this.submissions.get(cachedId); if (cached !== undefined) return parseJson(cached); const [record] = await this.records( @@ -825,6 +843,7 @@ export class PostgresStorage implements Storage { [this.session, conversationId, encodeIndexedString(requestId)], ); if (record !== undefined) this.rememberSubmission(record); + else this.submissionsByRequest.set(key, null); return record === undefined ? undefined : parseJson(record); }); } @@ -874,10 +893,7 @@ export class PostgresStorage implements Storage { }; } } - // The record and revision queries must observe one committed state; a commit between them can replace the base. - const stored = await inTransaction(this.pool, "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY", (client) => - this.materializeDocument(client, id, at), - ); + const stored = await this.materializeDocument(id, at); if (at === "current" && stored !== undefined) { this.documentStates.set(id, { record: encodeJson(stored.record), @@ -998,40 +1014,46 @@ export class PostgresStorage implements Storage { return record; } - private async materializeDocument( - executor: PostgresQueryable, - id: DocumentId, - at: DocumentPoint, - ): Promise { - const record = await this.recordById(this.tables.documents, id, executor); - if (record === undefined) return undefined; - if (at !== "current" && isCurrentOnly(record)) { - throw new Error(`Document ${id} does not retain historical content`); - } - if (!isAliveAt(record, at)) return undefined; + /** One statement reads the record and the revisions from the newest base at or before `at`. */ + private async materializeDocument(id: DocumentId, at: DocumentPoint): Promise { const upper = at === "current" ? Number.MAX_SAFE_INTEGER : at; - const base = ( - await executor.query( - `SELECT seq, kind, version, content::text AS content FROM ${this.tables.revisions} - WHERE session_id = $1 AND document_id = $2 AND kind = 'base' AND seq <= $3 - ORDER BY seq DESC LIMIT 1`, + const row = ( + await this.pool.query<{ readonly record: string | null; readonly revisions: string | null }>( + `SELECT (SELECT x.record::text FROM ${this.tables.documents} x WHERE x.session_id = $1 AND x.id = $2) AS record, + (SELECT json_agg(json_build_object('seq', r.seq, 'kind', r.kind, 'version', r.version, 'content', r.content) + ORDER BY r.seq)::text + FROM ${this.tables.revisions} r WHERE r.session_id = $1 AND r.document_id = $2 AND r.seq <= $3 + AND r.seq >= (SELECT b.seq FROM ${this.tables.revisions} b WHERE b.session_id = $1 AND b.document_id = $2 + AND b.kind = 'base' AND b.seq <= $3 ORDER BY b.seq DESC LIMIT 1)) AS revisions`, [this.session, id, upper], ) ).rows[0]; - if (base === undefined) throw new Error(`Document ${id} is missing a required base`); - let value = parseJson(base.content); - const tail = ( - await executor.query( - `SELECT seq, kind, version, content::text AS content FROM ${this.tables.revisions} - WHERE session_id = $1 AND document_id = $2 AND seq > $3 AND seq <= $4 ORDER BY seq`, - [this.session, id, Number(base.seq), upper], - ) - ).rows; + if (row?.record === null || row?.record === undefined) return undefined; + const record = parseJson(row.record); + if (at !== "current" && isCurrentOnly(record)) { + throw new Error(`Document ${id} does not retain historical content`); + } + if (!isAliveAt(record, at)) return undefined; + const revisions = + row.revisions === null + ? [] + : parseJson< + readonly { + readonly seq: number; + readonly kind: DocumentContent["kind"]; + readonly version: number; + readonly content: JsonValue; + }[] + >(row.revisions); + const base = revisions[0]; + if (base === undefined || base.kind !== "base") throw new Error(`Document ${id} is missing a required base`); + let value = base.content as JsonObject; + const tail = revisions.slice(1); for (const revision of tail) { if (revision.kind !== "delta" || revision.version !== base.version) { throw new Error(`Document ${id} crosses a stored version boundary without a base`); } - value = apply(value, parseJson(revision.content)) as JsonObject; + value = apply(value, revision.content as unknown as readonly Op[]) as JsonObject; } return { record, version: base.version, value, deltasSinceBase: tail.length }; } @@ -1045,22 +1067,19 @@ export class PostgresStorage implements Storage { return nextId; } - private checkGlobalIds(writes: readonly StorageWrite[], existingByID: ReadonlyMap, TableName>): void { + /** The batch-internal id rules; whether an id already belongs to a stored record is the commit function's check. */ + private checkBatchIds(writes: readonly StorageWrite[]): void { const claimed = new Map, TableName>(); for (const write of writes) { if (write.type === "document.change" || write.type === "document.retire") continue; const document = write.type === "document.create" || write.type === "document.copy"; const table: TableName = document ? "document" : write.type; const id = document ? write.record.id : write.value.id; - const existing = existingByID.get(id); const earlier = claimed.get(id); if (table === "conversation" || table === "entry" || table === "document") { - if (existing !== undefined) throw new Error(`ID ${id} already belongs to ${existing}`); if (earlier !== undefined) throw new Error(`ID ${id} is written more than once`); - } else { - if (existing !== undefined && existing !== table) - throw new Error(`ID ${id} already belongs to ${existing}`); - if (earlier !== undefined && earlier !== table) throw new Error(`ID ${id} is written as two record types`); + } else if (earlier !== undefined && earlier !== table) { + throw new Error(`ID ${id} is written as two record types`); } claimed.set(id, table); } @@ -1113,86 +1132,24 @@ export class PostgresStorage implements Storage { return actions; } - /** Validates the batch's document actions against the records read for them. */ - private async checkDocumentActions( - executor: PostgresQueryable, - actions: ReadonlyMap, - existingByID: ReadonlyMap, - versions: ReadonlyMap, - ): Promise { - const liveCounts = new Map(); - for (const [id, action] of actions) { - if (action.copy !== undefined && actions.has(action.copy.id)) { - throw new StorageRejected(`Document copy ${id} source is changed in the copy batch`); - } - const existing = existingByID.get(id); - if (action.create === undefined && existing === undefined) throw new Error(`Unknown document: ${id}`); - if (action.create !== undefined && existing !== undefined) throw new Error(`Document ${id} already exists`); - if (existing?.retiredAt !== undefined) throw new Error(`Document ${id} is retired`); - if (action.content?.kind === "delta") { - const previous = versions.get(id); - if (previous === undefined) throw new Error(`Document ${id} delta has no base`); - if (previous !== action.content.version) { - throw new Error(`Document ${id} version transition requires a base`); - } - } - // Only a create or retire can change how many incarnations an address has; a content change on an - // existing document needs no lookup. - if (action.create === undefined && !action.retire) continue; - const record = action.create ?? existing!; - const key = addressKey(record); - let live = liveCounts.get(key); - if (live === undefined) live = (await this.currentDocumentId(executor, record)) === undefined ? 0 : 1; - if (action.retire && existing !== undefined) live--; - if (action.create !== undefined && !action.retire) live++; - liveCounts.set(key, live); - } - for (const live of liveCounts.values()) { - if (live > 1) throw new Error("Document address already has a current incarnation"); - } - } - - private async currentDocumentId( - executor: PostgresQueryable, - address: DocumentAddress | DocumentCreate | DocumentRecord, - ): Promise { - const parts = addressParts(address); - const id = ( - await executor.query( - `SELECT id FROM ${this.tables.documents} - WHERE session_id = $1 AND kind = $2 AND scope_kind = $3 AND owner_id = $4 AND family = $5 - AND key_value = $6 AND retired_at IS NULL - LIMIT 1`, - [this.session, parts.kind, parts.scopeKind, parts.ownerId, parts.family, parts.keyValue], - ) - ).rows[0]?.id; - return id === undefined ? undefined : idFromNumber(Number(id)); - } - /** - * Resolves the batch's document actions into rows: records to insert, retirements to apply to existing - * documents, revision histories to clear, and revisions to add. Copy sources are read here; the rows are - * written with everything else in `writeBatch`. + * Resolves the batch's document actions into the rows the commit function takes, reading copy sources here. + * A created record carries its sequence; an existing document's retirement and any history clearing are the + * function's to apply, since it holds the stored record. */ private async prepareDocumentWrites( - executor: PostgresQueryable, actions: ReadonlyMap, - existing: ReadonlyMap, seq: Seq, ): Promise { - const writes: DocumentWrites = { - claims: [], - inserts: [], - retirements: [], - clearRevisions: [], - revisions: [], - applied: new Map(), - }; + const writes: DocumentWrites = { claims: [], rows: [], applied: new Map() }; for (const [id, action] of actions) { + if (action.copy !== undefined && actions.has(action.copy.id)) { + throw new StorageRejected(`Document copy ${id} source is changed in the copy batch`); + } let content = action.content; if (action.copy !== undefined) { try { - const stored = await this.materializeDocument(executor, action.copy.id, action.copy.at); + const stored = await this.materializeDocument(action.copy.id, action.copy.at); if (stored === undefined) throw new Error(`Fork source document ${action.copy.id} cannot be read`); const create = action.create!; if ( @@ -1211,187 +1168,32 @@ export class PostgresStorage implements Storage { throw new StorageRejected(`Document copy ${id} was rejected`, { cause: error }); } } - let record: DocumentRecord; + let created: DocumentRecord | undefined; + let parts: ReturnType | undefined; if (action.create !== undefined) { - record = { ...action.create, createdAt: seq, ...(action.retire ? { retiredAt: seq } : {}) }; - const parts = addressParts(record); + created = { ...action.create, createdAt: seq, ...(action.retire ? { retiredAt: seq } : {}) }; + parts = addressParts(created); writes.claims.push(id); - writes.inserts.push([ - this.session, - id, - parts.kind, - parts.family, - parts.keyValue, - parts.scopeKind, - parts.ownerId, - seq, - action.retire ? seq : null, - encodeJson(record), - ]); - } else { - record = existing.get(id)!; - if (action.retire) { - record = { ...record, retiredAt: seq }; - writes.retirements.push([this.session, id, seq, encodeJson(record)]); - } } - // A retired current-only document keeps no revisions, so one written in the same batch would only be - // deleted again and is skipped. A base otherwise replaces a current-only document's history. - const retiredCurrentOnly = action.retire && isCurrentOnly(record); - if (retiredCurrentOnly || (content?.kind === "base" && isCurrentOnly(record))) writes.clearRevisions.push(id); - if (content !== undefined && !retiredCurrentOnly) { - const encodedContent = content.kind === "base" ? encodeJson(content.value) : encodeJson(content.ops); - writes.revisions.push([this.session, id, seq, content.kind, content.version, encodedContent]); - } - writes.applied.set(id, { record, content: retiredCurrentOnly ? undefined : content }); + writes.rows.push({ + id, + is_create: created !== undefined, + retire: action.retire, + kind: parts?.kind ?? null, + family: parts?.family ?? null, + key_value: parts?.keyValue ?? null, + scope_kind: parts?.scopeKind ?? null, + owner_id: parts?.ownerId ?? null, + record: created === undefined ? null : encodeJson(created), + content_kind: content?.kind ?? null, + version: content?.version ?? null, + content: content === undefined ? null : encodeJson(content.kind === "base" ? content.value : content.ops), + }); + writes.applied.set(id, { created, retired: action.retire, content }); } return writes; } - /** - * Writes the whole batch in one statement: data-modifying CTEs claim every new id, insert each table's rows, - * clear and add revisions, and the statement's final UPDATE advances the session. A task or submission written - * twice in one batch keeps its last write, as serial upserts did. Retiring an existing document is the one - * action written separately, since it updates rows that may also be read by the batch's other actions. - */ - private async writeBatch( - executor: PostgresQueryable, - writes: readonly StorageWrite[], - documents: DocumentWrites, - seq: Seq, - session: { readonly nextId: number; readonly nextSeq: number; readonly expectedSeq: number }, - ): Promise { - const claims = new Map, TableName>(); - const conversations: PostgresValue[][] = []; - const entries: PostgresValue[][] = []; - const tasks = new Map, PostgresValue[]>(); - const submissions = new Map, PostgresValue[]>(); - for (const write of writes) { - switch (write.type) { - case "conversation": - claims.set(write.value.id, "conversation"); - conversations.push([ - this.session, - write.value.id, - write.value.owner?.conversationId ?? null, - write.value.owner?.taskId ?? null, - encodeJson(write.value), - ]); - break; - case "entry": - claims.set(write.value.id, "entry"); - entries.push([ - this.session, - write.value.id, - write.value.conversationId, - write.value.head ?? null, - seq, - encodeJson(write.value), - ]); - break; - case "task": - claims.set(write.value.id, "task"); - tasks.set(write.value.id, [ - this.session, - write.value.id, - write.value.conversationId, - encodeIndexedString(write.value.kind), - write.value.state.status, - write.value.abortRequested, - write.value.background, - encodeJson(write.value), - ]); - break; - case "submission": - claims.set(write.value.id, "submission"); - submissions.set(write.value.id, [ - this.session, - write.value.id, - write.value.conversationId, - write.value.requestId === undefined ? null : encodeIndexedString(write.value.requestId), - write.value.status, - encodeJson(write.value), - ]); - break; - case "document.create": - case "document.copy": - case "document.change": - case "document.retire": - break; - } - } - for (const id of documents.claims) claims.set(id, "document"); - - // Every write is gated on the guard: the session row still carries this owner at the sequence the - // pre-read observed. When another process has taken the session over, every CTE writes nothing and the - // final UPDATE affects no row. - const params: PostgresValue[] = [this.session, this.owner, session.expectedSeq]; - const guard = `WHERE EXISTS (SELECT 1 FROM guard)`; - // Parameters in a VALUES subselect carry no type inference, so each column is cast to its table type. - const values = (rows: readonly (readonly PostgresValue[])[], casts: readonly string[]): string => { - const names = casts.map((_, index) => `c${index}`).join(", "); - const list = rows - .map((row) => `(${row.map((value, index) => `$${params.push(value)}::${casts[index]}`).join(", ")})`) - .join(", "); - return `SELECT ${names} FROM (VALUES ${list}) AS v(${names}) ${guard}`; - }; - const parts: string[] = [ - `guard AS (SELECT 1 FROM ${this.tables.sessions} WHERE session_id = $1 AND owner = $2 AND next_seq = $3)`, - ]; - for (const [, id, retiredAt, record] of documents.retirements) { - parts.push(`retire_${id} AS (UPDATE ${this.tables.documents} SET retired_at = $${params.push(retiredAt)}::bigint, - record = $${params.push(record)}::json WHERE session_id = $1 AND id = $${params.push(id)}::bigint AND EXISTS (SELECT 1 FROM guard))`); - } - if (claims.size > 0) { - parts.push(`claims AS (INSERT INTO ${this.tables.recordIds} (session_id, id, record_type) - ${values( - [...claims].map(([id, table]) => [this.session, id, table]), - RECORD_ID_TYPES, - )} - ON CONFLICT (session_id, id) DO NOTHING)`); - } - if (conversations.length > 0) { - parts.push(`conversations AS (INSERT INTO ${this.tables.conversations} - (session_id, id, owner_conversation_id, owner_task_id, record) ${values(conversations, CONVERSATION_TYPES)})`); - } - if (entries.length > 0) { - parts.push(`entries AS (INSERT INTO ${this.tables.entries} - (session_id, id, conversation_id, head, commit_seq, record) ${values(entries, ENTRY_TYPES)})`); - } - if (tasks.size > 0) { - parts.push(`tasks AS (INSERT INTO ${this.tables.tasks} - (session_id, id, conversation_id, kind, status, abort_requested, background, record) - ${values([...tasks.values()], TASK_TYPES)} - ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, - kind = excluded.kind, status = excluded.status, abort_requested = excluded.abort_requested, - background = excluded.background, record = excluded.record)`); - } - if (submissions.size > 0) { - parts.push(`submissions AS (INSERT INTO ${this.tables.submissions} - (session_id, id, conversation_id, request_id, status, record) ${values([...submissions.values()], SUBMISSION_TYPES)} - ON CONFLICT (session_id, id) DO UPDATE SET conversation_id = excluded.conversation_id, - request_id = excluded.request_id, status = excluded.status, record = excluded.record)`); - } - if (documents.inserts.length > 0) { - parts.push(`documents AS (INSERT INTO ${this.tables.documents} - (session_id, id, kind, family, key_value, scope_kind, owner_id, created_at, retired_at, record) - ${values(documents.inserts, DOCUMENT_TYPES)})`); - } - if (documents.clearRevisions.length > 0) { - parts.push(`cleared AS (DELETE FROM ${this.tables.revisions} - WHERE session_id = $1 AND document_id = ANY($${params.push(arrayLiteral(documents.clearRevisions))}::bigint[]) - AND EXISTS (SELECT 1 FROM guard))`); - } - if (documents.revisions.length > 0) { - parts.push(`revisions AS (INSERT INTO ${this.tables.revisions} - (session_id, document_id, seq, kind, version, content) ${values(documents.revisions, REVISION_TYPES)})`); - } - const update = `UPDATE ${this.tables.sessions} SET next_id = $${params.push(session.nextId)}, next_seq = $${params.push(session.nextSeq)} - WHERE session_id = $1 AND owner = $2 AND next_seq = $3`; - const result = await executor.query(`WITH ${parts.join(",\n")} ${update}`, params); - return (result.rowCount ?? 0) > 0; - } - private assertOpen(): void { if (this.closed) throw new Error("PostgresStorage is closed"); if (this.lost !== undefined) { From 84d5e610701292ac62e611aa42f5078ad26eec0b Mon Sep 17 00:00:00 2001 From: rgarcia <72655+rgarcia@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:12:57 +0000 Subject: [PATCH 5/5] feat(durable): invalidate Postgres document caches per scope --- .../durable/src/storage/postgres/storage.ts | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/packages/durable/src/storage/postgres/storage.ts b/packages/durable/src/storage/postgres/storage.ts index 7ef87f3d4c40..014e3cb008c3 100644 --- a/packages/durable/src/storage/postgres/storage.ts +++ b/packages/durable/src/storage/postgres/storage.ts @@ -734,15 +734,34 @@ export class PostgresStorage implements Storage { this.documentStates.delete(id); } } - // Creating or retiring a document changes what an address or scope scan resolves to. - if ( - documents.rows.some( - (row) => (row as { is_create: boolean; retire: boolean }).is_create || (row as { retire: boolean }).retire, - ) - ) { - this.documentScans = new Map(); - this.documentAddresses = new Map(); + // Creating or retiring a document changes what an address or scope scan in its scope resolves to. A + // retirement of a document whose scope is not known here invalidates every scope. + for (const [id, { created, retired }] of documents.applied) { + if (created === undefined && !retired) continue; + const scope = created !== undefined ? scopeColumns(created.scope) : this.cachedDocumentScope(id); + if (scope === undefined) { + this.documentScans = new Map(); + this.documentAddresses = new Map(); + break; + } + const scanPrefix = `${scope.scopeKind}:${scope.ownerId}:`; + for (const key of [...this.documentScans.keys()]) + if (key.startsWith(scanPrefix)) this.documentScans.delete(key); + const addressMark = `,"${scope.scopeKind}",${scope.ownerId},`; + for (const key of [...this.documentAddresses.keys()]) + if (key.includes(addressMark)) this.documentAddresses.delete(key); + } + } + + /** The scope of a document this storage has a record for, from its state or address caches. */ + private cachedDocumentScope(id: DocumentId): ScopeColumns | undefined { + const state = this.documentStates.get(id); + if (state !== undefined) return scopeColumns(parseJson(state.record).scope); + for (const record of this.documentAddresses.values()) { + const parsed = parseJson(record); + if (parsed.id === id) return scopeColumns(parsed.scope); } + return undefined; } private rememberSubmission(record: string): void {