orez 0.1.21 → 0.1.23

package/src/integration/replication-latency.test.ts

@@ -0,0 +1,428 @@
+/**
+ * replication latency stress test.
+ *
+ * measures the end-to-end time from a proxy write to the zero-cache
+ * websocket poke arriving at the client. this is the critical path
+ * that determines whether UI re-renders overlap with user interactions.
+ *
+ * run: vitest run src/integration/replication-latency.test.ts
+ */
+
+import postgres from 'postgres'
+import { describe, expect, test, beforeAll, afterAll } from 'vitest'
+import WebSocket from 'ws'
+
+import { startZeroLite } from '../index.js'
+import { installChangeTracking } from '../replication/change-tracker.js'
+import {
+  ensureTablesInPublications,
+  installAllowAllPermissions,
+} from './test-permissions.js'
+
+import type { PGlite } from '@electric-sql/pglite'
+
+const SYNC_PROTOCOL_VERSION = 45
+
+function encodeSecProtocols(
+  initConnectionMessage: unknown,
+  authToken: string | undefined
+): string {
+  const payload = JSON.stringify({ initConnectionMessage, authToken })
+  return encodeURIComponent(Buffer.from(payload, 'utf-8').toString('base64'))
+}
+
+class Queue<T> {
+  private items: T[] = []
+  private waiters: Array<{
+    resolve: (v: T) => void
+    timer?: ReturnType<typeof setTimeout>
+  }> = []
+
+  enqueue(item: T) {
+    const waiter = this.waiters.shift()
+    if (waiter) {
+      if (waiter.timer) clearTimeout(waiter.timer)
+      waiter.resolve(item)
+    } else {
+      this.items.push(item)
+    }
+  }
+
+  dequeue(fallback?: T, timeoutMs = 10000): Promise<T> {
+    if (this.items.length > 0) {
+      return Promise.resolve(this.items.shift()!)
+    }
+    return new Promise<T>((resolve) => {
+      const waiter: { resolve: (v: T) => void; timer?: ReturnType<typeof setTimeout> } = {
+        resolve,
+      }
+      if (fallback !== undefined) {
+        waiter.timer = setTimeout(() => {
+          const idx = this.waiters.indexOf(waiter)
+          if (idx >= 0) this.waiters.splice(idx, 1)
+          resolve(fallback)
+        }, timeoutMs)
+      }
+      this.waiters.push(waiter)
+    })
+  }
+}
+
+describe('replication latency', { timeout: 120000 }, () => {
+  let db: PGlite
+  let zeroPort: number
+  let pgPort: number
+  let shutdown: () => Promise<void>
+  let resetZeroFull: (() => Promise<void>) | undefined
+  let dataDir: string
+  let sql: ReturnType<typeof postgres>
+
+  beforeAll(async () => {
+    const testPgPort = 24000 + Math.floor(Math.random() * 1000)
+    const testZeroPort = testPgPort + 100
+
+    dataDir = `.orez-latency-test-${Date.now()}`
+    const result = await startZeroLite({
+      pgPort: testPgPort,
+      zeroPort: testZeroPort,
+      dataDir,
+      logLevel: 'info',
+      skipZeroCache: false,
+    })
+
+    db = result.db
+    zeroPort = result.zeroPort
+    pgPort = result.pgPort
+    shutdown = result.stop
+    resetZeroFull = result.resetZeroFull
+
+    // create test table
+    await db.exec(`
+      CREATE TABLE IF NOT EXISTS latency_test (
+        id TEXT PRIMARY KEY,
+        value TEXT,
+        ts BIGINT
+      );
+    `)
+    await ensureTablesInPublications(db, ['latency_test'])
+    const pubName = process.env.ZERO_APP_PUBLICATIONS?.trim()
+    if (pubName) {
+      const quotedPub = '"' + pubName.replace(/"/g, '""') + '"'
+      await db
+        .exec(`ALTER PUBLICATION ${quotedPub} ADD TABLE "public"."latency_test"`)
+        .catch(() => {})
+      await installChangeTracking(db)
+    }
+    await installAllowAllPermissions(db, ['latency_test'])
+    if (resetZeroFull) await resetZeroFull()
+
+    // wait for zero-cache ready
+    await waitForZero(zeroPort, 90000)
+
+    // connect via wire protocol (like a real app would)
+    sql = postgres(`postgresql://user:password@127.0.0.1:${pgPort}/postgres`, {
+      max: 1,
+      idle_timeout: 0,
+    })
+  }, 120000)
+
+  afterAll(async () => {
+    if (sql) await sql.end()
+    if (shutdown) await shutdown()
+    if (dataDir) {
+      const { rmSync } = await import('node:fs')
+      try {
+        rmSync(dataDir, { recursive: true, force: true })
+      } catch {}
+    }
+  })
+
+  test('measure write-to-poke latency (single inserts)', async () => {
+    const downstream = new Queue<unknown>()
+    const ws = connectAndSubscribe(zeroPort, downstream)
+    await drainInitialPokes(downstream)
+
+    const NUM_WRITES = 20
+    const latencies: number[] = []
+
+    for (let i = 0; i < NUM_WRITES; i++) {
+      const id = `latency-${i}-${Date.now()}`
+      const writeStart = performance.now()
+
+      // write through the wire protocol proxy (like a real app)
+      await sql`INSERT INTO latency_test (id, value, ts) VALUES (${id}, ${'test'}, ${Date.now()})`
+
+      // wait for the poke containing our row
+      const poke = await waitForPokeWithRow(downstream, 'latency_test', id, 10000)
+      const latencyMs = performance.now() - writeStart
+
+      expect(poke).toBeTruthy()
+      latencies.push(latencyMs)
+    }
+
+    ws.close()
+
+    // report
+    latencies.sort((a, b) => a - b)
+    const avg = latencies.reduce((s, v) => s + v, 0) / latencies.length
+    const p50 = latencies[Math.floor(latencies.length * 0.5)]
+    const p95 = latencies[Math.floor(latencies.length * 0.95)]
+    const p99 = latencies[Math.floor(latencies.length * 0.99)]
+    const max = latencies[latencies.length - 1]
+
+    console.log(`\n[replication latency] ${NUM_WRITES} single inserts via wire protocol:`)
+    console.log(
+      `  avg=${avg.toFixed(1)}ms  p50=${p50.toFixed(1)}ms  p95=${p95.toFixed(1)}ms  p99=${p99.toFixed(1)}ms  max=${max.toFixed(1)}ms`
+    )
+    console.log(`  all: ${latencies.map((l) => l.toFixed(0)).join(', ')}ms`)
+
+    // assert reasonable latency — under 200ms avg means the UI re-render
+    // arrives before a user can interact with the element
+    expect(avg).toBeLessThan(200)
+    // no single write should take more than 500ms
+    expect(max).toBeLessThan(500)
+  })
+
+  test('count poke batches per single write', async () => {
+    // theory: orez causes 2+ poke batches per write because zero-cache
+    // writes shard updates back through the proxy, creating a separate
+    // replication batch. real postgres doesn't have this round-trip.
+    const downstream = new Queue<unknown>()
+    const ws = connectAndSubscribe(zeroPort, downstream)
+    await drainInitialPokes(downstream)
+
+    const id = `poke-count-${Date.now()}`
+    await sql`INSERT INTO latency_test (id, value, ts) VALUES (${id}, ${'count-test'}, ${Date.now()})`
+
+    // collect ALL messages for 2 seconds after the write
+    const messages: any[] = []
+    const deadline = Date.now() + 2000
+    while (Date.now() < deadline) {
+      const remaining = Math.max(100, deadline - Date.now())
+      const msg = (await downstream.dequeue('timeout' as any, remaining)) as any
+      if (msg !== 'timeout') messages.push(msg)
+    }
+
+    const pokeStarts = messages.filter((m) => Array.isArray(m) && m[0] === 'pokeStart')
+    const pokeEnds = messages.filter((m) => Array.isArray(m) && m[0] === 'pokeEnd')
+    const pokeParts = messages.filter((m) => Array.isArray(m) && m[0] === 'pokePart')
+
+    console.log(`\n[poke batches] after 1 INSERT:`)
+    console.log(
+      `  pokeStart=${pokeStarts.length}  pokePart=${pokeParts.length}  pokeEnd=${pokeEnds.length}`
+    )
+    console.log(`  total messages: ${messages.length}`)
+    for (const msg of messages) {
+      if (Array.isArray(msg)) {
+        const type = msg[0]
+        if (type === 'pokePart' && msg[1]?.rowsPatch) {
+          const tables = msg[1].rowsPatch
+            .map((r: any) => `${r.op}:${r.tableName}`)
+            .join(', ')
+          console.log(`    pokePart: ${tables}`)
+        } else {
+          console.log(`    ${type}`)
+        }
+      }
+    }
+
+    // ideally just 1 poke cycle per write, but we want to measure reality
+    expect(pokeStarts.length).toBeGreaterThanOrEqual(1)
+
+    ws.close()
+  })
+
+  test('count poke batches when shard tables update', async () => {
+    // simulate what happens in the real app: zero-cache writes to shard
+    // tables (clients.lastMutationID) after processing a mutation.
+    // these shard writes go through the proxy and trigger replication.
+    const downstream = new Queue<unknown>()
+    const ws = connectAndSubscribe(zeroPort, downstream)
+    await drainInitialPokes(downstream)
+
+    const id = `shard-test-${Date.now()}`
+    // insert via proxy (triggers replication)
+    await sql`INSERT INTO latency_test (id, value, ts) VALUES (${id}, ${'shard'}, ${Date.now()})`
+
+    // now simulate a shard write (like zero-cache updating clients table)
+    // check if any shard schemas exist
+    const shardSchemas = await sql`
+      SELECT nspname FROM pg_namespace
+      WHERE nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast', 'public', '_orez')
+        AND nspname NOT LIKE 'pg_%'
+        AND nspname NOT LIKE 'zero_%'
+        AND nspname NOT LIKE '_zero_%'
+        AND nspname NOT LIKE '%/%'
+    `
+
+    // collect messages for 3 seconds
+    const messages: any[] = []
+    const deadline = Date.now() + 3000
+    while (Date.now() < deadline) {
+      const remaining = Math.max(100, deadline - Date.now())
+      const msg = (await downstream.dequeue('timeout' as any, remaining)) as any
+      if (msg !== 'timeout') messages.push(msg)
+    }
+
+    const pokeStarts = messages.filter((m) => Array.isArray(m) && m[0] === 'pokeStart')
+    const pokeParts = messages.filter((m) => Array.isArray(m) && m[0] === 'pokePart')
+
+    console.log(
+      `\n[shard poke batches] after INSERT + shard schemas=${shardSchemas.length}:`
+    )
+    console.log(`  pokeStart=${pokeStarts.length}  pokePart=${pokeParts.length}`)
+    for (const msg of messages) {
+      if (Array.isArray(msg) && msg[0] === 'pokePart' && msg[1]?.rowsPatch) {
+        const tables = msg[1].rowsPatch
+          .map((r: any) => `${r.op}:${r.tableName}`)
+          .join(', ')
+        console.log(`    pokePart: ${tables}`)
+      }
+    }
+
+    expect(pokeStarts.length).toBeGreaterThanOrEqual(1)
+    ws.close()
+  })
+
+  test('measure rapid sequential write latency', async () => {
+    const downstream = new Queue<unknown>()
+    const ws = connectAndSubscribe(zeroPort, downstream)
+    await drainInitialPokes(downstream)
+
+    // simulate rapid sequential writes (like a chat app sending messages)
+    const NUM_WRITES = 10
+    const ids: string[] = []
+    const writeStart = performance.now()
+
+    for (let i = 0; i < NUM_WRITES; i++) {
+      const id = `rapid-${i}-${Date.now()}`
+      ids.push(id)
+      await sql`INSERT INTO latency_test (id, value, ts) VALUES (${id}, ${'rapid'}, ${Date.now()})`
+    }
+
+    const writeEnd = performance.now()
+
+    // wait for ALL rows to arrive
+    const receivedIds = new Set<string>()
+    const deadline = Date.now() + 30000
+    while (receivedIds.size < NUM_WRITES && Date.now() < deadline) {
+      const msg = (await downstream.dequeue('timeout' as any, 5000)) as any
+      if (msg === 'timeout') continue
+      if (Array.isArray(msg) && msg[0] === 'pokePart' && msg[1]?.rowsPatch) {
+        for (const row of msg[1].rowsPatch) {
+          if (row.op === 'put' && row.tableName === 'latency_test' && row.value?.id) {
+            receivedIds.add(row.value.id)
+          }
+        }
+      }
+    }
+
+    const totalMs = performance.now() - writeStart
+    const writeMs = writeEnd - writeStart
+    const replicationMs = totalMs - writeMs
+
+    console.log(`\n[replication latency] ${NUM_WRITES} rapid sequential inserts:`)
+    console.log(
+      `  write=${writeMs.toFixed(1)}ms  replication=${replicationMs.toFixed(1)}ms  total=${totalMs.toFixed(1)}ms`
+    )
+    console.log(`  received ${receivedIds.size}/${NUM_WRITES} rows`)
+
+    expect(receivedIds.size).toBe(NUM_WRITES)
+    for (const id of ids) {
+      expect(receivedIds.has(id)).toBe(true)
+    }
+    // all 10 writes + replication should complete in under 3s
+    expect(totalMs).toBeLessThan(3000)
+  })
+
+  // --- helpers ---
+
+  function connectAndSubscribe(port: number, downstream: Queue<unknown>): WebSocket {
+    const cg = `latency-cg-${Date.now()}`
+    const cid = `latency-client-${Date.now()}`
+    const secProtocol = encodeSecProtocols(
+      [
+        'initConnection',
+        {
+          desiredQueriesPatch: [
+            {
+              op: 'put',
+              hash: 'q1',
+              ast: {
+                table: 'latency_test',
+                orderBy: [['id', 'asc']],
+              },
+            },
+          ],
+          clientSchema: {
+            tables: {
+              latency_test: {
+                columns: {
+                  id: { type: 'string' },
+                  value: { type: 'string' },
+                  ts: { type: 'number' },
+                },
+                primaryKey: ['id'],
+              },
+            },
+          },
+        },
+      ],
+      undefined
+    )
+    const ws = new WebSocket(
+      `ws://localhost:${port}/sync/v${SYNC_PROTOCOL_VERSION}/connect` +
+        `?clientGroupID=${cg}&clientID=${cid}&wsid=ws1&schemaVersion=1&baseCookie=&ts=${Date.now()}&lmid=0`,
+      secProtocol
+    )
+    ws.on('message', (data) => downstream.enqueue(JSON.parse(data.toString())))
+    return ws
+  }
+
+  async function drainInitialPokes(downstream: Queue<unknown>) {
+    let settled = false
+    const timeout = Date.now() + 30000
+    while (!settled && Date.now() < timeout) {
+      const msg = (await downstream.dequeue('timeout' as any, 3000)) as any
+      if (msg === 'timeout') {
+        settled = true
+      } else if (Array.isArray(msg) && msg[0] === 'pokeEnd') {
+        const next = (await downstream.dequeue('timeout' as any, 2000)) as any
+        if (next === 'timeout') settled = true
+      }
+    }
+  }
+
+  async function waitForPokeWithRow(
+    downstream: Queue<unknown>,
+    tableName: string,
+    rowId: string,
+    timeoutMs = 10000
+  ): Promise<Record<string, any> | null> {
+    const deadline = Date.now() + timeoutMs
+    while (Date.now() < deadline) {
+      const remaining = Math.max(500, deadline - Date.now())
+      const msg = (await downstream.dequeue('timeout' as any, remaining)) as any
+      if (msg === 'timeout') return null
+      if (Array.isArray(msg) && msg[0] === 'pokePart' && msg[1]?.rowsPatch) {
+        const match = msg[1].rowsPatch.find(
+          (r: any) => r.op === 'put' && r.tableName === tableName && r.value?.id === rowId
+        )
+        if (match) return match
+      }
+    }
+    return null
+  }
+})
+
+async function waitForZero(port: number, timeoutMs = 60000): Promise<void> {
+  const deadline = Date.now() + timeoutMs
+  while (Date.now() < deadline) {
+    try {
+      const res = await fetch(`http://localhost:${port}/`)
+      if (res.ok) return
+    } catch {}
+    await new Promise((r) => setTimeout(r, 500))
+  }
+  throw new Error(`zero-cache did not become ready within ${timeoutMs}ms`)
+}

package/src/pg-proxy.ts

@@ -470,16 +470,14 @@ export async function startPgProxy(
     return { db: instances.postgres, mutex: mutexes.postgres, txState: txStates.postgres }
   }
 
-  // shared debounce timer for extended protocol write signaling.
-  // 2ms trailing-edge: each Sync resets the timer, so the signal fires
-  // 2ms after the LAST write in a pipeline (coalesces rapid sequential writes).
-  // 2ms is enough because postgres.js pipelines are sub-ms per statement.
+  // signal replication handler after writes complete.
+  // 2ms trailing-edge debounce coalesces rapid sequential writes and ensures
+  // the socket response is fully flushed before the handler acquires the mutex.
   let signalTimer: ReturnType<typeof setTimeout> | null = null
-  function debouncedSignal() {
+  function signalWrite() {
     if (signalTimer) clearTimeout(signalTimer)
     signalTimer = setTimeout(() => {
       signalTimer = null
-      log.debug.proxy('ext-write: debounced signal firing')
       signalReplicationChange()
     }, 2)
   }
@@ -497,6 +495,10 @@ export async function startPgProxy(
     // track extended protocol writes (Parse with INSERT/UPDATE/DELETE/COPY/TRUNCATE)
     // so we can signal replication on Sync (0x53) after the pipeline completes
     let extWritePending = false
+    // hold mutex across entire extended protocol pipeline (Parse→Sync).
+    // prevents other connections from interleaving and corrupting PGlite's
+    // unnamed portal/statement state during the pipeline.
+    let pipelineMutexHeld = false
     // clean up pglite transaction state when a client disconnects
     socket.on('close', async () => {
       // replication sockets don't own a transaction — skip ROLLBACK
@@ -583,6 +585,94 @@ export async function startPgProxy(
           const msgType = data[0]
           const { db, mutex, txState } = getDbContext(dbName)
 
+          // extended protocol pipeline: hold mutex across Parse→Sync to prevent
+          // other connections from interleaving and corrupting unnamed portal state.
+          // 0x50=Parse, 0x42=Bind, 0x44=Describe, 0x45=Execute, 0x43=Close, 0x48=Flush
+          const isExtendedMsg =
+            msgType === 0x50 ||
+            msgType === 0x42 ||
+            msgType === 0x44 ||
+            msgType === 0x45 ||
+            msgType === 0x43 ||
+            msgType === 0x48
+          const isSyncInPipeline = msgType === 0x53 && pipelineMutexHeld
+
+          if (isExtendedMsg || isSyncInPipeline) {
+            // acquire mutex on first message of pipeline
+            if (!pipelineMutexHeld) {
+              const t0 = performance.now()
+              await mutex.acquire()
+              proxyStats.totalWaitMs += performance.now() - t0
+              pipelineMutexHeld = true
+              // auto-rollback stale transactions from other connections
+              if (txState.status === 0x45 && txState.owner !== socket) {
+                try {
+                  await db.exec('ROLLBACK')
+                } catch {}
+                txState.status = 0x49
+                txState.owner = null
+              }
+            }
+
+            // detect extended protocol writes for replication signaling
+            if (dbName === 'postgres' && msgType === 0x50) {
+              const q = extractParseQuery(data)?.trimStart().toLowerCase()
+              if (q && /^(insert|update|delete|copy|truncate)/.test(q)) {
+                extWritePending = true
+                log.debug.proxy(`ext-write: detected ${q.slice(0, 40)}`)
+              }
+            }
+
+            // apply query rewrites
+            data = interceptQuery(data)
+
+            const t1 = performance.now()
+            let result: Uint8Array
+            try {
+              result = await db.execProtocolRaw(data, { syncToFs: false })
+            } catch (err) {
+              mutex.release()
+              pipelineMutexHeld = false
+              throw err
+            }
+            const t2 = performance.now()
+            proxyStats.totalExecMs += t2 - t1
+            proxyStats.count++
+
+            // update transaction state
+            const rfqStatus = getReadyForQueryStatus(result)
+            if (rfqStatus !== null) {
+              txState.status = rfqStatus
+              txState.owner = rfqStatus === 0x49 ? null : socket
+            }
+
+            // release mutex on Sync (end of pipeline)
+            if (msgType === 0x53) {
+              mutex.release()
+              pipelineMutexHeld = false
+              proxyStats.batches++
+
+              // signal replication handler on postgres writes
+              if (dbName === 'postgres' && extWritePending) {
+                extWritePending = false
+                signalWrite()
+              }
+            } else {
+              // strip ReadyForQuery from non-Sync pipeline messages
+              result = stripResponseMessages(result, true)
+            }
+
+            if (proxyStats.count % 200 === 0) {
+              log.debug.proxy(
+                `perf: ${proxyStats.count} ops (${proxyStats.batches} batches) | mutex ${proxyStats.totalWaitMs.toFixed(0)}ms | pglite ${proxyStats.totalExecMs.toFixed(0)}ms`
+              )
+            }
+
+            return result
+          }
+
+          // Simple Query (0x51) or standalone Sync — per-message mutex
+
           // check for no-op queries (only SimpleQuery has queries worth intercepting)
           if (isNoopQuery(data)) {
             if (msgType === 0x51) {
@@ -590,15 +680,6 @@ export async function startPgProxy(
             }
           }
 
-          // detect extended protocol writes on postgres db for replication signaling
-          if (dbName === 'postgres' && msgType === 0x50) {
-            const q = extractParseQuery(data)?.trimStart().toLowerCase()
-            if (q && /^(insert|update|delete|copy|truncate)/.test(q)) {
-              extWritePending = true
-              log.debug.proxy(`ext-write: detected ${q.slice(0, 40)}`)
-            }
-          }
-
           // intercept and rewrite queries
           data = interceptQuery(data)
 
@@ -620,10 +701,6 @@ export async function startPgProxy(
           const execute = async (): Promise<Uint8Array> => {
             const t0 = performance.now()
             await mutex.acquire()
-            // pglite is single-connection: if a previous connection left an aborted
-            // transaction, this connection will inherit it. auto-rollback stale
-            // transactions from OTHER connections, but let the SAME connection
-            // handle its own errors (e.g. ROLLBACK TO SAVEPOINT in migrations).
             if (txState.status === 0x45 && txState.owner !== socket) {
               try {
                 await db.exec('ROLLBACK')
@@ -639,7 +716,6 @@ export async function startPgProxy(
               mutex.release()
               throw err
             }
-            // update transaction state tracking
             const rfqStatus = getReadyForQueryStatus(result)
             if (rfqStatus !== null) {
               txState.status = rfqStatus
@@ -689,25 +765,16 @@ export async function startPgProxy(
           result = stripResponseMessages(result, stripRfq)
 
           // signal replication handler on postgres writes for instant sync
-          if (dbName === 'postgres') {
-            if (isSimpleQuery && queryText) {
-              // immediate signal for SimpleQuery writes
-              const q = queryText.trimStart().toLowerCase()
-              if (
-                q.startsWith('insert') ||
-                q.startsWith('update') ||
-                q.startsWith('delete') ||
-                q.startsWith('copy') ||
-                q.startsWith('truncate')
-              ) {
-                signalReplicationChange()
-              }
-            } else if (msgType === 0x53 && extWritePending) {
-              // debounced signal for extended protocol writes.
-              // fires 5ms after the last Sync in a pipeline, coalescing
-              // rapid sequential writes (e.g. server→channel→member)
-              extWritePending = false
-              debouncedSignal()
+          if (dbName === 'postgres' && isSimpleQuery && queryText) {
+            const q = queryText.trimStart().toLowerCase()
+            if (
+              q.startsWith('insert') ||
+              q.startsWith('update') ||
+              q.startsWith('delete') ||
+              q.startsWith('copy') ||
+              q.startsWith('truncate')
+            ) {
+              signalReplicationChange()
             }
           }