Persistence: a hydration that starts inside a commit-reserving source transaction rejects that transaction's commit

// @vitest-environment node
// Repro against @tanstack/db 0.15.2 + @tanstack/db-sqlite-persistence-core 0.4.12. Run: npx vitest run

import { createCollection, type SyncConfig } from '@tanstack/db'
import {
  eagerStorageTarget,
  type ProtocolEnvelope,
  persistedCollectionOptions,
  SingleProcessCoordinator,
  SQLITE_DRIVER_SHARED_LOGICAL_SCHEDULING_KEY,
  SQLiteCorePersistenceAdapter,
  type SQLiteDriver,
  type TxCommitted,
} from '@tanstack/db-sqlite-persistence-core'
import { afterEach, describe, expect, it } from 'vitest'

import { DatabaseSync, type SQLInputValue } from 'node:sqlite'

interface ShapeRow extends Record<string, unknown> {
  readonly id: string
  readonly value: string
}

/**
 * A real in-memory SQLite driver whose next write transaction can be held to keep the persistence queue busy.
 * `sharedScheduling` marks it as one physical database, as the browser driver does, so every adapter on it
 * shares one write/hydrate scheduler.
 */
const openHoldableDatabase = ({ sharedScheduling = false }: { sharedScheduling?: boolean } = {}) => {
  const database = new DatabaseSync(':memory:')
  let hold: { reached: () => void; released: Promise<void> } | undefined
  let transactionQueue: Promise<unknown> = Promise.resolve()
  const driver: SQLiteDriver = {
    ...(sharedScheduling ? { [SQLITE_DRIVER_SHARED_LOGICAL_SCHEDULING_KEY]: {} } : {}),
    exec: async sql => {
      database.exec(sql)
    },
    query: async <T>(sql: string, params: readonly unknown[] = []) =>
      // SQLiteDriver's generic contract erases SQL parameter and result types.
      database.prepare(sql).all(...(params as SQLInputValue[])) as T[],
    run: async (sql, params = []) => {
      database.prepare(sql).run(...(params as SQLInputValue[]))
    },
    transaction: work => {
      const transaction = transactionQueue.then(async () => {
        const current = hold
        hold = undefined
        if (current) {
          current.reached()
          await current.released
        }
        database.exec('BEGIN IMMEDIATE')
        try {
          const result = await work(driver)
          database.exec('COMMIT')
          return result
        } catch (error) {
          database.exec('ROLLBACK')
          throw error
        }
      })
      transactionQueue = transaction.catch(() => undefined)
      return transaction
    },
  }
  const holdNextWrite = () => {
    let release!: () => void
    let reached!: () => void
    const wasReached = new Promise<void>(resolve => {
      reached = resolve
    })
    hold = { reached, released: new Promise<void>(resolve => (release = resolve)) }
    return { release, wasReached }
  }
  return { database, driver, holdNextWrite }
}

const waitFor = async (assertion: () => void, timeoutMs: number = 2_000): Promise<void> => {
  const deadline = Date.now() + timeoutMs
  while (true) {
    try {
      assertion()
      return
    } catch (error) {
      if (Date.now() >= deadline) throw error
      await new Promise(resolve => setTimeout(resolve, 5))
    }
  }
}

/** A coordinator for one tab that lets the test deliver another tab's broadcast, as BroadcastChannel does. */
class PeerTabCoordinator extends SingleProcessCoordinator {
  private readonly listeners = new Map<string, (message: ProtocolEnvelope<unknown>) => void>()

  override subscribe(collectionId?: string, onMessage?: (message: ProtocolEnvelope<unknown>) => void): () => void {
    if (collectionId === undefined || onMessage === undefined) return () => undefined
    this.listeners.set(collectionId, onMessage)
    return () => this.listeners.delete(collectionId)
  }

  deliverFromPeer(collectionId: string, payload: TxCommitted): void {
    const listener = this.listeners.get(collectionId)
    if (!listener) throw new Error(`No runtime subscribed to ${collectionId}.`)
    listener({ collectionId, dbName: 'repro', payload, senderId: 'peer-tab', ts: Date.now(), v: 1 })
  }
}

const COLLECTION_ID = 'items'

type SourceSyncParams = Parameters<SyncConfig<ShapeRow, string>['sync']>[0]

/**
 * A persisted collection whose source the test drives by hand, with another tab committing to the same SQLite
 * database. `reserve()` opens a source transaction that reserves its commit turn, as Electric's refetch does.
 */
const openPeerTabHarness = async ({ sharedScheduling = false }: { sharedScheduling?: boolean } = {}) => {
  const { database, driver, holdNextWrite } = openHoldableDatabase({ sharedScheduling })
  // The concrete class exposes the stream-position and hydration-scope methods the factory's interface marks optional.
  const adapter = new SQLiteCorePersistenceAdapter({ driver })
  const coordinator = new PeerTabCoordinator()
  let source: SourceSyncParams | undefined
  const collection = createCollection(
    persistedCollectionOptions<ShapeRow, string>({
      id: COLLECTION_ID,
      persistence: { adapter, coordinator },
      schemaVersion: 1,
      getKey: row => row.id,
      sync: {
        sync: params => {
          source = params
          params.begin()
          params.write({ type: 'insert', value: { id: 'row-1', value: 'initial' } })
          params.commit()
          params.markReady()
          return {}
        },
      },
    })
  )
  const requireSource = (): SourceSyncParams => {
    if (!source) throw new Error('The collection has not started its source.')
    return source
  }
  const start = async () => {
    await collection.preload()
    await waitFor(() => expect(collection.get('row-1')).toMatchObject({ value: 'initial' }))
  }
  /** Commits a row change from the peer tab into the shared database and returns its broadcast notice. */
  const commitFromPeer = async (value: string): Promise<TxCommitted> => {
    const position = await adapter.getStreamPosition(eagerStorageTarget(COLLECTION_ID))
    const tx = {
      mutations: [{ key: 'row-1', type: 'update' as const, value: { id: 'row-1', value } }],
      rowVersion: position.latestRowVersion + 1,
      seq: position.latestSeq + 1,
      term: position.latestTerm,
      txId: value,
    }
    await adapter.applyCommittedTx(eagerStorageTarget(COLLECTION_ID), tx)
    return {
      latestRowVersion: tx.rowVersion,
      requiresFullReload: true,
      seq: tx.seq,
      term: tx.term,
      txId: tx.txId,
      type: 'tx:committed',
    }
  }
  const reserve = (value?: string) => {
    const params = requireSource()
    params.begin()
    if (value) params.write({ type: 'update', value: { id: 'row-1', value } })
    params.metadata?.persistence?.reserveCommitTurn()
  }
  const commit = () => requireSource().commit()
  /** The peer commits and announces while this tab's reserved transaction is open, so the notice is deferred. */
  const deferPeerCommit = async (value: string) => {
    const notice = await commitFromPeer(value)
    reserve()
    coordinator.deliverFromPeer(COLLECTION_ID, notice)
    await new Promise(resolve => setTimeout(resolve, 30))
    expect(collection.get('row-1')).not.toMatchObject({ value })
  }
  await start()
  return { adapter, coordinator, collection, commit, commitFromPeer, database, deferPeerCommit, holdNextWrite, reserve, start }
}


describe('a hydration that starts inside a commit-reserving source transaction', () => {
  const cleanups: Array<() => Promise<void> | void> = []
  afterEach(async () => {
    for (const cleanup of cleanups.splice(0).reverse()) await cleanup()
  })
  const open = async (options?: { sharedScheduling?: boolean }) => {
    const harness = await openPeerTabHarness(options)
    cleanups.push(() => harness.database.close())
    cleanups.push(() => harness.collection.cleanup())
    return harness
  }

  // Two tabs on one database. The other tab commits and broadcasts a full-reload notice while this tab's source
  // holds an open transaction that reserved its commit turn (as Electric does across a 409 must-refetch).
  it('commits the reserved transaction when a peer full-reload notice arrives while it is open', async () => {
    const harness = await open()
    const notice = await harness.commitFromPeer('peer edit')
    harness.reserve('refetched')
    harness.coordinator.deliverFromPeer(COLLECTION_ID, notice)
    await new Promise(resolve => setTimeout(resolve, 30))
    await harness.commit()
    await waitFor(() => expect(harness.collection.get('row-1')).toMatchObject({ value: 'refetched' }))
    expect(harness.collection.status).not.toBe('error')
  })

  // The shared scheduler can hold a hydration behind another collection's write after its mutex turn began. The
  // reservation must be checked when the hydration starts, not when its turn was admitted.
  it('commits the reserved transaction when a reload waited behind another collection in the shared scheduler', async () => {
    const harness = await open({ sharedScheduling: true })
    expect(harness.adapter.isHydrationScopeScheduled()).toBe(true)
    await harness.adapter.loadSubset(eagerStorageTarget('other_collection'), {})
    const notice = await harness.commitFromPeer('peer edit')
    const otherWrite = harness.holdNextWrite()
    const blocker = harness.adapter.applyCommittedTx(eagerStorageTarget('other_collection'), { mutations: [], rowVersion: 1, seq: 1, term: 1, txId: 'blocker' })
    await otherWrite.wasReached
    let reloadAdmitted!: () => void
    const admission = new Promise<void>(resolve => {
      reloadAdmitted = resolve
    })
    const runInHydrationScope = harness.adapter.runInHydrationScope.bind(harness.adapter)
    harness.adapter.runInHydrationScope = task => {
      reloadAdmitted()
      return runInHydrationScope(task)
    }
    harness.coordinator.deliverFromPeer(COLLECTION_ID, notice)
    await admission
    harness.reserve('refetched')
    otherWrite.release()
    await blocker
    await new Promise(resolve => setTimeout(resolve, 30))
    await harness.commit()
    await waitFor(() => expect(harness.collection.get('row-1')).toMatchObject({ value: 'refetched' }))
    expect(harness.collection.status).not.toBe('error')
  })
})