experimental-a2 0.7.0 → 0.8.0

package/dist/store-postgres.js

@@ -71,6 +71,26 @@ create table if not exists a2_snapshots (
   primary key (session_id, reducer_name)
 );
 
+create table if not exists a2_snapshot_history (
+  session_id   text   not null,
+  reducer_name text   not null,
+  up_to_index  bigint not null,
+  state        jsonb  not null,
+  updated_at   bigint not null,
+  primary key (session_id, reducer_name, up_to_index)
+);
+
+create table if not exists a2_snapshot_pins (
+  session_id   text   not null,
+  event_index  bigint not null,
+  reducer_name text   not null,
+  up_to_index  bigint not null,
+  primary key (session_id, event_index, reducer_name, up_to_index)
+);
+
+create index if not exists a2_snapshot_pins_checkpoint
+  on a2_snapshot_pins (session_id, reducer_name, up_to_index);
+
 create table if not exists a2_presence (
   ns          text   not null,
   participant text   not null,
@@ -175,6 +195,90 @@ function postgres(options = {}) {
 		}
 	});
 	const client = connection.get;
+	const readStateBatch = async (requests, reducerName) => {
+		if (requests.length === 0) return [];
+		const c = await client();
+		const input = requests.map((request, position) => ({
+			position,
+			sessionId: request.sessionId,
+			throughIndex: request.throughIndex ?? null,
+			snapshotThroughIndex: request.snapshotThroughIndex ?? request.throughIndex ?? null
+		}));
+		const { rows } = await c.query(`with request as materialized (
+         select * from jsonb_to_recordset($1::jsonb) as value(
+           position integer,
+           "sessionId" text,
+           "throughIndex" bigint,
+           "snapshotThroughIndex" bigint
+         )
+       ), selected as materialized (
+         select request.position,
+                request."sessionId" as session_id,
+                request."throughIndex" as through_index,
+                head.up_to_index as head_index,
+                candidate.up_to_index as snapshot_index,
+                candidate.state as snapshot_state
+           from request
+           left join a2_snapshots head
+             on head.session_id = request."sessionId"
+            and head.reducer_name = $2
+           left join lateral (
+             select up_to_index, state from (
+               select head.up_to_index, head.state
+                where request."snapshotThroughIndex" is null
+                   or head.up_to_index <= request."snapshotThroughIndex"
+               union all
+               select history.up_to_index, history.state
+                 from a2_snapshot_history history
+                where history.session_id = request."sessionId"
+                  and history.reducer_name = $2
+                  and (
+                    request."snapshotThroughIndex" is null
+                    or history.up_to_index <= request."snapshotThroughIndex"
+                  )
+             ) candidates
+             order by up_to_index desc
+             limit 1
+           ) candidate on true
+       )
+       select selected.position, 0 as row_order, selected.session_id,
+              selected.head_index, selected.snapshot_index,
+              selected.snapshot_state,
+              null::bigint as idx, null::text as event_type,
+              null::jsonb as payload, null::text as event_id,
+              null::bigint as created_at
+         from selected
+       union all
+       select selected.position, 1 as row_order, event.session_id,
+              selected.head_index, null::bigint as snapshot_index,
+              null::jsonb as snapshot_state,
+              event.idx, event.event_type, event.payload,
+              event.event_id, event.created_at
+         from selected
+         join a2_events event on event.session_id = selected.session_id
+        where event.idx > coalesce(selected.snapshot_index, 0)
+          and (
+            selected.through_index is null
+            or event.idx <= selected.through_index
+          )
+       order by position, row_order, idx`, [jsonb(input), reducerName]);
+		const reads = requests.map(() => ({
+			headIndex: null,
+			snapshot: null,
+			events: []
+		}));
+		for (const row of rows) {
+			const read = reads[asNumber(row["position"])];
+			if (row["row_order"] === 0 || row["row_order"] === "0") {
+				read.headIndex = row["head_index"] === null ? null : asNumber(row["head_index"]);
+				if (row["snapshot_index"] !== null) read.snapshot = {
+					index: asNumber(row["snapshot_index"]),
+					state: row["snapshot_state"]
+				};
+			} else read.events.push(toEvent(row));
+		}
+		return reads;
+	};
 	let mutex = Promise.resolve();
 	const withTx = async (fn) => {
 		const c = await client();
@@ -520,6 +624,16 @@ function postgres(options = {}) {
 					index,
 					jsonb(ids)
 				]);
+				await query(`delete from a2_snapshot_pins
+              where session_id = $1 and event_index = $2`, [sessionId, index]);
+				await query(`delete from a2_snapshot_history history
+              where history.session_id = $1
+                and not exists (
+                  select 1 from a2_snapshot_pins pin
+                   where pin.session_id = history.session_id
+                     and pin.reducer_name = history.reducer_name
+                     and pin.up_to_index = history.up_to_index
+                )`, [sessionId]);
 				const parentLane = parent["lane"] ?? null;
 				if (parentLane !== null) await query(`update a2_events set lane_ready = true
                 where session_id = $1
@@ -535,8 +649,8 @@ function postgres(options = {}) {
 			}));
 		},
 		async failAttempt({ sessionId, index, attempt, error, maxFailures }) {
-			return wrap(async () => {
-				const { rows } = await (await client()).query(`with current as materialized (
+			return wrap(() => withTx(async (query) => {
+				const { rows } = await query(`with current as materialized (
              select e.* from a2_events e
              where e.session_id = $1 and e.idx = $2
              for update of e
@@ -593,116 +707,131 @@ function postgres(options = {}) {
 				]);
 				const row = rows[0];
 				if (!row || row["outcome"] === "missing") throw new TypeError(`no event at index ${index} in session '${sessionId}'`);
+				if (row["outcome"] === "dead_lettered") {
+					await query(`delete from a2_snapshot_pins
+                where session_id = $1 and event_index = $2`, [sessionId, index]);
+					await query(`delete from a2_snapshot_history history
+                where history.session_id = $1
+                  and not exists (
+                    select 1 from a2_snapshot_pins pin
+                     where pin.session_id = history.session_id
+                       and pin.reducer_name = history.reducer_name
+                       and pin.up_to_index = history.up_to_index
+                  )`, [sessionId]);
+				}
 				return {
 					outcome: row["outcome"],
 					failureCount: asNumber(row["failure_count"])
 				};
-			});
+			}));
 		},
-		async readState(sessionId, reducerName) {
-			return wrap(async () => {
-				const { rows } = await (await client()).query(`with snapshot as materialized (
-             select up_to_index, state
-             from a2_snapshots
-             where session_id = $1 and reducer_name = $2
-           )
-           select 0 as row_order, 'snapshot' as row_kind,
-                  snapshot.up_to_index as snapshot_index,
-                  snapshot.state as snapshot_state,
-                  null::text as session_id, null::bigint as idx,
-                  null::text as event_type, null::jsonb as payload,
-                  null::text as event_id, null::bigint as created_at
-           from snapshot
-           union all
-           select 1 as row_order, 'event' as row_kind,
-                  null::bigint as snapshot_index,
-                  null::jsonb as snapshot_state,
-                  event.session_id, event.idx, event.event_type,
-                  event.payload, event.event_id, event.created_at
-           from a2_events event
-           where event.session_id = $1
-             and event.idx > coalesce((select up_to_index from snapshot), 0)
-           order by row_order, idx`, [sessionId, reducerName]);
-				const snapshot = rows.find((row) => row["row_kind"] === "snapshot");
-				return {
-					snapshot: snapshot ? {
-						index: asNumber(snapshot["snapshot_index"]),
-						state: snapshot["snapshot_state"]
-					} : null,
-					events: rows.filter((row) => row["row_kind"] === "event").map(toEvent)
-				};
-			});
+		async readState(sessionId, reducerName, stateOptions) {
+			return wrap(async () => (await readStateBatch([{
+				sessionId,
+				...stateOptions?.throughIndex === void 0 ? {} : { throughIndex: stateOptions.throughIndex },
+				...stateOptions?.snapshotThroughIndex === void 0 ? {} : { snapshotThroughIndex: stateOptions.snapshotThroughIndex }
+			}], reducerName))[0]);
 		},
-		async readStates(sessionIds, reducerName) {
-			if (sessionIds.length === 0) return [];
-			return wrap(async () => {
-				const { rows } = await (await client()).query(`with snapshot as materialized (
-             select session_id, up_to_index, state
-             from a2_snapshots
-             where session_id = any($1::text[]) and reducer_name = $2
-           )
-           select 0 as row_order, 'snapshot' as row_kind,
-                  snapshot.session_id,
-                  snapshot.up_to_index as snapshot_index,
-                  snapshot.state as snapshot_state,
-                  null::bigint as idx, null::text as event_type,
-                  null::jsonb as payload, null::text as event_id,
-                  null::bigint as created_at
-           from snapshot
-           union all
-           select 1 as row_order, 'event' as row_kind,
-                  event.session_id,
-                  null::bigint as snapshot_index,
-                  null::jsonb as snapshot_state,
-                  event.idx, event.event_type, event.payload,
-                  event.event_id, event.created_at
-           from a2_events event
-           left join snapshot on snapshot.session_id = event.session_id
-           where event.session_id = any($1::text[])
-             and event.idx > coalesce(snapshot.up_to_index, 0)
-           order by row_order, session_id, idx`, [sessionIds, reducerName]);
-				const bySession = /* @__PURE__ */ new Map();
-				const readOf = (sessionId) => {
-					let read = bySession.get(sessionId);
-					if (!read) {
-						read = {
-							snapshot: null,
-							events: []
-						};
-						bySession.set(sessionId, read);
-					}
-					return read;
-				};
-				for (const row of rows) {
-					const read = readOf(row["session_id"]);
-					if (row["row_kind"] === "snapshot") read.snapshot = {
-						index: asNumber(row["snapshot_index"]),
-						state: row["snapshot_state"]
-					};
-					else read.events.push(toEvent(row));
+		async readStates(requests, reducerName) {
+			return wrap(() => readStateBatch(requests, reducerName));
+		},
+		async putSnapshots(sessionId, reducerName, writes) {
+			await wrap(() => withTx(async (query) => {
+				const normalized = [...new Map(writes.map((write) => [write.index, write])).values()].toSorted((a, b) => a.index - b.index);
+				if (normalized.length === 0) return;
+				await query("select pg_advisory_xact_lock(hashtext($1), hashtext($2))", [sessionId, reducerName]);
+				const requestedPins = normalized.flatMap((write) => (write.pinEventIndexes ?? []).map((eventIndex) => ({
+					eventIndex,
+					snapshotIndex: write.index
+				})));
+				if (requestedPins.length > 0) {
+					const { rows: events } = await query(`select idx, processed_at, failed_at from a2_events
+                where session_id = $1 and idx = any($2::bigint[])
+                for update`, [sessionId, [...new Set(requestedPins.map((pin) => pin.eventIndex))]]);
+					const unfinished = new Set(events.filter((row) => row["processed_at"] === null && row["failed_at"] === null).map((row) => asNumber(row["idx"])));
+					const accepted = requestedPins.filter((pin) => unfinished.has(pin.eventIndex));
+					if (accepted.length > 0) await query(`insert into a2_snapshot_pins
+                   (session_id, event_index, reducer_name, up_to_index)
+                 select $1, pin."eventIndex", $2, pin."snapshotIndex"
+                   from jsonb_to_recordset($3::jsonb) as pin(
+                     "eventIndex" bigint,
+                     "snapshotIndex" bigint
+                   )
+                 on conflict do nothing`, [
+						sessionId,
+						reducerName,
+						jsonb(accepted)
+					]);
 				}
-				return sessionIds.map((sessionId) => bySession.get(sessionId) ?? {
-					snapshot: null,
-					events: []
+				const { rows: headRows } = await query(`select up_to_index, state, updated_at from a2_snapshots
+              where session_id = $1 and reducer_name = $2
+              for update`, [sessionId, reducerName]);
+				const currentHead = headRows[0];
+				const latestWrite = normalized.at(-1);
+				const writesHead = !currentHead || latestWrite.index >= asNumber(currentHead["up_to_index"]);
+				const advancesHead = !currentHead || latestWrite.index > asNumber(currentHead["up_to_index"]);
+				const finalHeadIndex = advancesHead ? latestWrite.index : asNumber(currentHead["up_to_index"]);
+				const { rows: pinRows } = await query(`select distinct up_to_index from a2_snapshot_pins
+              where session_id = $1 and reducer_name = $2
+                and up_to_index = any($3::bigint[])`, [
+					sessionId,
+					reducerName,
+					[.../* @__PURE__ */ new Set([...normalized.map((write) => write.index), ...currentHead ? [asNumber(currentHead["up_to_index"])] : []])]
+				]);
+				const pinned = new Set(pinRows.map((row) => asNumber(row["up_to_index"])));
+				const historical = /* @__PURE__ */ new Map();
+				if (advancesHead && currentHead && pinned.has(asNumber(currentHead["up_to_index"]))) {
+					const index = asNumber(currentHead["up_to_index"]);
+					historical.set(index, {
+						index,
+						state: currentHead["state"],
+						updatedAt: asNumber(currentHead["updated_at"])
+					});
+				}
+				const updatedAt = clock.now().getTime();
+				for (const write of normalized) if (write.index < finalHeadIndex && pinned.has(write.index)) historical.set(write.index, {
+					index: write.index,
+					state: write.state,
+					updatedAt
 				});
-			});
-		},
-		async putSnapshot(sessionId, reducerName, index, state) {
-			await wrap(async () => {
-				await (await client()).query(`insert into a2_snapshots (session_id, reducer_name, up_to_index, state, updated_at)
-           values ($1, $2, $3, $4::jsonb, $5)
-           on conflict (session_id, reducer_name) do update set
-             up_to_index = excluded.up_to_index,
-             state = excluded.state,
-             updated_at = excluded.updated_at
-           where excluded.up_to_index > a2_snapshots.up_to_index`, [
+				if (historical.size > 0) await query(`insert into a2_snapshot_history
+                 (session_id, reducer_name, up_to_index, state, updated_at)
+               select $1, $2, snapshot.index, snapshot.state, snapshot."updatedAt"
+                 from jsonb_to_recordset($3::jsonb) as snapshot(
+                   index bigint,
+                   state jsonb,
+                   "updatedAt" bigint
+                 )
+               on conflict (session_id, reducer_name, up_to_index) do update set
+                 state = excluded.state,
+                 updated_at = excluded.updated_at`, [
 					sessionId,
 					reducerName,
-					index,
-					jsonb(state),
-					clock.now().getTime()
+					jsonb([...historical.values()])
 				]);
-			});
+				if (writesHead) {
+					await query(`insert into a2_snapshots
+                 (session_id, reducer_name, up_to_index, state, updated_at)
+               values ($1, $2, $3, $4::jsonb, $5)
+               on conflict (session_id, reducer_name) do update set
+                 up_to_index = excluded.up_to_index,
+                 state = excluded.state,
+                 updated_at = excluded.updated_at
+               where excluded.up_to_index >= a2_snapshots.up_to_index`, [
+						sessionId,
+						reducerName,
+						latestWrite.index,
+						jsonb(latestWrite.state),
+						clock.now().getTime()
+					]);
+					await query(`delete from a2_snapshot_history
+                where session_id = $1 and reducer_name = $2 and up_to_index = $3`, [
+						sessionId,
+						reducerName,
+						latestWrite.index
+					]);
+				}
+			}));
 		},
 		presence: {
 			async set(ns, participant, values, meta) {