experimental-a2 0.6.0 → 0.8.0

package/src/store-postgres.ts

@@ -31,6 +31,7 @@ import {
   type EventCause,
   type IdSource,
   type StoredEvent,
+  type StoreStateReadRequest,
   type StoreStateRead,
 } from './store.ts'
 
@@ -120,6 +121,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,
@@ -287,6 +308,103 @@ export function postgres(options: PostgresStoreOptions = {}): PostgresStore {
   })
   const client = connection.get
 
+  const readStateBatch = async (
+    requests: readonly StoreStateReadRequest[],
+    reducerName: string,
+  ): Promise<StoreStateRead[]> => {
+    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<StoreStateRead>(() => ({
+      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
+  }
+
   // Transactions run on one connection. Pools check one out; a
   // single-connection client (PGlite) serializes on this mutex instead.
   let mutex: Promise<unknown> = Promise.resolve()
@@ -797,6 +915,22 @@ export function postgres(options: PostgresStoreOptions = {}): PostgresStore {
              where session_id = $3 and idx = $4`,
             [now, attempt, sessionId, 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'] as string | null) ?? null
           if (parentLane !== null) {
             await query(
@@ -816,10 +950,10 @@ export function postgres(options: PostgresStoreOptions = {}): PostgresStore {
     },
 
     async failAttempt({ sessionId, index, attempt, error, maxFailures }) {
-      return wrap(async () => {
-        const c = await client()
-        const { rows } = await c.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
@@ -867,148 +1001,231 @@ export function postgres(options: PostgresStoreOptions = {}): PostgresStore {
              (select failure_count from updated),
              (select failure_count from current)
            ) as failure_count`,
-          [
-            sessionId,
-            index,
-            attempt,
-            error,
-            maxFailures,
-            clock.now().getTime(),
-          ],
-        )
-        const row = rows[0]
-        if (!row || row['outcome'] === 'missing') {
-          throw new TypeError(
-            `no event at index ${index} in session '${sessionId}'`,
+            [
+              sessionId,
+              index,
+              attempt,
+              error,
+              maxFailures,
+              clock.now().getTime(),
+            ],
           )
-        }
-        return {
-          outcome: row['outcome'] as 'failed' | 'dead_lettered' | 'superseded',
-          failureCount: asNumber(row['failure_count']),
-        }
-      })
+          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'] as
+              'failed' | 'dead_lettered' | 'superseded',
+            failureCount: asNumber(row['failure_count']),
+          }
+        }),
+      )
     },
 
-    async readState(sessionId, reducerName) {
-      return wrap(async () => {
-        const c = await client()
-        const { rows } = await c.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 === undefined
+                    ? {}
+                    : { throughIndex: stateOptions.throughIndex }),
+                  ...(stateOptions?.snapshotThroughIndex === undefined
+                    ? {}
+                    : {
+                        snapshotThroughIndex: stateOptions.snapshotThroughIndex,
+                      }),
+                },
+              ],
+              reducerName,
+            )
+          )[0]!,
+      )
     },
 
-    async readStates(sessionIds, reducerName) {
-      if (sessionIds.length === 0) return []
-      return wrap(async () => {
-        const c = await client()
-        // One statement, one MVCC snapshot — each session's tail is
-        // bounded by the same snapshot row the read returns.
-        const { rows } = await c.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 = new Map<string, StoreStateRead>()
-        const readOf = (sessionId: string): StoreStateRead => {
-          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'] as string)
-          if (row['row_kind'] === 'snapshot') {
-            read.snapshot = {
-              index: asNumber(row['snapshot_index']),
-              state: row['snapshot_state'],
+    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 eventIndexes = [
+              ...new Set(requestedPins.map((pin) => pin.eventIndex)),
+            ]
+            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, eventIndexes],
+            )
+            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)],
+              )
             }
-          } else {
-            read.events.push(toEvent(row))
           }
-        }
-        return sessionIds.map(
-          (sessionId) =>
-            bySession.get(sessionId) ?? { snapshot: null, events: [] },
-        )
-      })
-    },
 
-    async putSnapshot(sessionId, reducerName, index, state) {
-      await wrap(async () => {
-        const c = await client()
-        // Guarded upsert: a slower concurrent writer must never clobber
-        // a further-along snapshot (spec §3).
-        await c.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, index, jsonb(state), clock.now().getTime()],
-        )
-      })
+          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 relevantIndexes = [
+            ...new Set([
+              ...normalized.map((write) => write.index),
+              ...(currentHead ? [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, relevantIndexes],
+          )
+          const pinned = new Set(
+            pinRows.map((row) => asNumber(row['up_to_index'])),
+          )
+          const historical = new Map<
+            number,
+            { index: number; state: unknown; updatedAt: number }
+          >()
+          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,
+              })
+            }
+          }
+          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, 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: {
@@ -1086,106 +1303,6 @@ export function postgres(options: PostgresStoreOptions = {}): PostgresStore {
       },
     },
 
-    inspect: {
-      async listSessions(inspectionOptions) {
-        return wrap(async () => {
-          const c = await client()
-          const conditions = ['left(session_id, length($1)) = $1']
-          const params: unknown[] = [inspectionOptions.prefix]
-          if (inspectionOptions.cursor !== undefined) {
-            params.push(inspectionOptions.cursor)
-            conditions.push(`session_id > $${params.length}`)
-          }
-          params.push(inspectionOptions.limit + 1)
-          const { rows } = await c.query(
-            `select
-               session_id,
-               count(*) as event_count,
-               count(*) filter (where processed_at is null and failed_at is null) as pending_count,
-               count(*) filter (where failed_at is not null) as failed_count,
-               sum(attempt_count) as attempt_count,
-               sum(failure_count) as failure_count,
-               min(created_at) as first_event_at,
-               max(greatest(
-                 created_at,
-                 coalesce(first_claimed_at, 0),
-                 coalesce(last_claimed_at, 0),
-                 coalesce(last_failed_at, 0),
-                 coalesce(processed_at, 0),
-                 coalesce(failed_at, 0)
-               )) as updated_at
-             from a2_events
-             where ${conditions.join(' and ')}
-             group by session_id
-             order by session_id
-             limit $${params.length}`,
-            params,
-          )
-          const hasMore = rows.length > inspectionOptions.limit
-          const visible = rows.slice(0, inspectionOptions.limit)
-          return {
-            sessions: visible.map((row) => ({
-              sessionId: row['session_id'] as string,
-              eventCount: asNumber(row['event_count']),
-              pendingCount: asNumber(row['pending_count']),
-              failedCount: asNumber(row['failed_count']),
-              attemptCount: asNumber(row['attempt_count']),
-              failureCount: asNumber(row['failure_count']),
-              firstEventAt: new Date(asNumber(row['first_event_at'])),
-              updatedAt: new Date(asNumber(row['updated_at'])),
-            })),
-            cursor: hasMore
-              ? ((visible.at(-1)?.['session_id'] as string | undefined) ?? null)
-              : null,
-          }
-        })
-      },
-      async readEvents(sessionId, pageOptions) {
-        return wrap(async () => {
-          const c = await client()
-          const currentFrontier = asNumber(
-            (
-              await c.query(
-                'select coalesce(max(idx), 0) as frontier from a2_events where session_id = $1',
-                [sessionId],
-              )
-            ).rows[0]?.['frontier'] ?? 0,
-          )
-          const throughIndex = Math.min(
-            pageOptions.throughIndex ?? currentFrontier,
-            currentFrontier,
-          )
-          const { rows } = await c.query(
-            `select * from a2_events
-              where session_id = $1 and idx > $2 and idx <= $3
-              order by idx limit $4`,
-            [
-              sessionId,
-              pageOptions.afterIndex,
-              throughIndex,
-              pageOptions.limit,
-            ],
-          )
-          return { events: rows.map(toStored), throughIndex }
-        })
-      },
-      async listSnapshots(sessionId) {
-        return wrap(async () => {
-          const c = await client()
-          const { rows } = await c.query(
-            `select reducer_name, up_to_index, updated_at
-             from a2_snapshots where session_id = $1 order by reducer_name`,
-            [sessionId],
-          )
-          return rows.map((row) => ({
-            reducerName: row['reducer_name'] as string,
-            index: asNumber(row['up_to_index']),
-            updatedAt: new Date(asNumber(row['updated_at'])),
-          }))
-        })
-      },
-    },
-
     stream(sessionId, opts) {
       const readAfter = (afterIndex: number): Promise<StoredEvent[]> =>
         wrap(async () => {