orez 0.1.21 → 0.1.22

package/src/replication/handler.test.ts

@@ -7,6 +7,7 @@ import {
   handleReplicationQuery,
   handleStartReplication,
   resetReplicationState,
+  signalReplicationChange,
   type ReplicationWriter,
 } from './handler'
 
@@ -137,6 +138,8 @@ describe('handleStartReplication', () => {
   afterEach(async () => {
     // closing db causes poll loop to exit with 'closed' error
     await db.close()
+    // wake handler from idle sleep so it hits the closed db and exits
+    signalReplicationChange()
     // wait for the replication promise to settle
     await replicationPromise?.catch(() => {})
   })
@@ -213,6 +216,7 @@ describe('handleStartReplication', () => {
 
     await new Promise((r) => setTimeout(r, 100))
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('streamed', 123)`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     const types = written.flatMap(extractPayloadTypes)
@@ -245,12 +249,15 @@ describe('handleStartReplication', () => {
     await new Promise((r) => setTimeout(r, 100))
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('mut', 1)`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     await db.exec(`UPDATE public.items SET value = 2 WHERE name = 'mut'`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     await db.exec(`DELETE FROM public.items WHERE name = 'mut'`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     const types = written.flatMap(extractPayloadTypes)
@@ -272,9 +279,11 @@ describe('handleStartReplication', () => {
     await new Promise((r) => setTimeout(r, 100))
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('a', 1)`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('b', 2)`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     const types = written.flatMap(extractPayloadTypes)
@@ -299,6 +308,7 @@ describe('handleStartReplication', () => {
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('a', 1)`)
     await db.exec(`INSERT INTO public.other (label) VALUES ('b')`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     const types = written.flatMap(extractPayloadTypes)
@@ -321,8 +331,9 @@ describe('handleStartReplication', () => {
     for (let i = 0; i < 20; i++) {
       await db.exec(`INSERT INTO public.items (name, value) VALUES ('r${i}', ${i})`)
     }
+    signalReplicationChange()
 
-    // wait multiple poll cycles
+    // wait for handler to process
     await new Promise((r) => setTimeout(r, 1500))
 
     const inserts = written.flatMap(extractPayloadTypes).filter((t) => t === 0x49)
@@ -342,9 +353,11 @@ describe('handleStartReplication', () => {
     await new Promise((r) => setTimeout(r, 100))
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('tx1', 1)`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('tx2', 2)`)
+    signalReplicationChange()
     await new Promise((r) => setTimeout(r, 700))
 
     const types = written.flatMap(extractPayloadTypes)

package/src/replication/handler.ts

@@ -11,6 +11,7 @@ import { log } from '../log.js'
 const textEncoder = new TextEncoder()
 import {
   getChangesSince,
+  getCurrentWatermark,
   purgeConsumedChanges,
   installTriggersOnShardTables,
   type ChangeRecord,
@@ -237,8 +238,16 @@ export async function handleReplicationQuery(
     const lsn = lsnToString(nextLsn())
     const snapshotName = `00000003-00000001-1`
 
-    // fresh slot = fresh zero-cache instance, reset watermark
-    lastStreamedWatermark = 0
+    // set watermark to current DB state so replication only delivers changes
+    // that happen AFTER this point. this mirrors real postgres behavior where
+    // CREATE_REPLICATION_SLOT creates a consistent snapshot — the initial copy
+    // captures everything up to this point, and replication picks up from here.
+    // on reconnect this is effectively a no-op since the watermark is already
+    // at or past the current DB state.
+    const currentWm = await getCurrentWatermark(db)
+    if (currentWm > lastStreamedWatermark) {
+      lastStreamedWatermark = currentWm
+    }
 
     // persist slot so pg_replication_slots queries find it
     await db.query(
@@ -527,20 +536,24 @@ export async function handleStartReplication(
 
   // event-driven replication: proxy signals changes directly via signalReplicationChange(),
   // pg_notify as secondary signal, polling as final fallback.
-  const pollIntervalIdle = 100
+  const pollIntervalIdle = 5000
   const batchSize = 50000
-  const purgeEveryN = 5
-  const shardRescanEveryN = 40
+  const purgeEveryN = 1
+  const shardRescanIntervalMs = 10_000
   let running = true
   let pollsSincePurge = 0
   let tryAcquireFailures = 0
-  let pollsSinceShardRescan = 0
+  let lastShardRescan = -shardRescanIntervalMs
   let hasStreamedOnce = false
 
-  // promise-based wakeup mechanism
+  // promise-based wakeup mechanism.
+  // signalPending captures signals that arrive while the handler is
+  // processing (not in waitForWakeup), preventing signal loss.
   let wakeupResolve: (() => void) | null = null
+  let signalPending = false
   let lastWakeupTime = 0
   const wakeup = () => {
+    signalPending = true
     if (wakeupResolve) {
       lastWakeupTime = performance.now()
       log.debug.repl('signal received, waking up')
@@ -548,15 +561,15 @@ export async function handleStartReplication(
       wakeupResolve = null
     }
   }
-  const waitForWakeup = (timeoutMs: number): Promise<void> => {
+  const waitForWakeup = (timeoutMs: number): Promise<boolean> => {
     return new Promise((resolve) => {
       const timer = setTimeout(() => {
         wakeupResolve = null
-        resolve()
+        resolve(false)
       }, timeoutMs)
       wakeupResolve = () => {
         clearTimeout(timer)
-        resolve()
+        resolve(true)
       }
     })
   }
@@ -574,6 +587,9 @@ export async function handleStartReplication(
   }
 
   const poll = async () => {
+    let queryPending = true // query immediately on first iteration
+    let idleTimeoutCount = 0
+
     while (running) {
       // check if the connection was closed
       if (writer.closed) {
@@ -583,16 +599,59 @@ export async function handleStartReplication(
       }
 
       try {
+        // when no query is pending, wait for a signal or timeout.
+        // signals fire instantly when the proxy processes a write,
+        // so we only hit the timeout when truly idle.
+        if (!queryPending) {
+          // check if a signal arrived while we were processing
+          if (!signalPending) {
+            const wasSignaled = await waitForWakeup(pollIntervalIdle)
+            if (writer.closed) {
+              running = false
+              break
+            }
+            if (!wasSignaled) {
+              idleTimeoutCount++
+              // send keepalive on every timeout
+              writer.write(encodeKeepalive(currentLsn, nowMicros(), false))
+              log.debug.repl(`idle keepalive (lastWatermark=${lastWatermark})`)
+              // re-scan for new shard schemas during idle
+              if (performance.now() - lastShardRescan > shardRescanIntervalMs) {
+                if (mutex.tryAcquire()) {
+                  lastShardRescan = performance.now()
+                  try {
+                    await installTriggersOnShardTables(db)
+                  } finally {
+                    mutex.release()
+                  }
+                }
+              }
+              // safety poll every ~30s to catch edge cases (6 * 5000ms)
+              if (idleTimeoutCount < 6) continue
+              idleTimeoutCount = 0
+              log.debug.repl('safety poll')
+              // fall through to query
+            } else {
+              idleTimeoutCount = 0
+            }
+          } else {
+            idleTimeoutCount = 0
+          }
+          signalPending = false
+        }
+        queryPending = false
+
         // periodically re-scan for new shard schemas (e.g. chat_0 created by zero-cache)
-        pollsSinceShardRescan++
-        if (pollsSinceShardRescan >= shardRescanEveryN) {
+        if (performance.now() - lastShardRescan > shardRescanIntervalMs) {
           if (mutex.tryAcquire()) {
-            pollsSinceShardRescan = 0
+            lastShardRescan = performance.now()
             try {
               await installTriggersOnShardTables(db)
             } finally {
               mutex.release()
             }
+          } else {
+            log.debug.repl('shard rescan skipped: mutex busy')
           }
         }
 
@@ -610,6 +669,7 @@ export async function handleStartReplication(
             if (tryAcquireFailures < 10) {
               // pre-sync: yield so zero-cache initial copy can finish
               await waitForWakeup(Math.min(10 * tryAcquireFailures, 100))
+              queryPending = true
               continue
             }
             await mutex.acquire()
@@ -673,9 +733,9 @@ export async function handleStartReplication(
             lastWatermark = batchEnd
             lastStreamedWatermark = batchEnd
             // all changes were filtered out (e.g. shard internal tables).
-            // sleep briefly to avoid a tight loop when zero-cache is
-            // continuously writing internal state.
-            await waitForWakeup(pollIntervalIdle)
+            // brief wait to avoid tight loop, then recheck.
+            await waitForWakeup(200)
+            queryPending = true
             continue
           }
 
@@ -707,17 +767,15 @@ export async function handleStartReplication(
           }
 
           // got changes - continue immediately to check for more
+          queryPending = true
           continue
         }
 
-        // send keepalive
+        // no changes: send keepalive
         const ts = nowMicros()
         writer.write(encodeKeepalive(currentLsn, ts, false))
-
         log.debug.repl(`idle (lastWatermark=${lastWatermark})`)
-
-        // no changes: wait for notify signal or poll interval
-        await waitForWakeup(pollIntervalIdle)
+        // next iteration will wait for signal at the top
       } catch (err: unknown) {
         const msg = err instanceof Error ? err.message : String(err)
         log.repl(`replication poll error: ${msg}`)
@@ -734,7 +792,10 @@ export async function handleStartReplication(
   try {
     await poll()
   } finally {
-    _replicationWakeup = null
+    // only clear if still pointing to our wakeup (a new handler may have replaced it)
+    if (_replicationWakeup === wakeup) {
+      _replicationWakeup = null
+    }
     if (unsubscribe) {
       await unsubscribe().catch(() => {})
     }

package/src/replication/tcp-replication.test.ts

@@ -17,6 +17,7 @@ import { describe, it, expect, beforeEach, afterEach } from 'vitest'
 import { getConfig } from '../config'
 import { startPgProxy } from '../pg-proxy'
 import { installChangeTracking } from './change-tracker'
+import { signalReplicationChange } from './handler'
 
 import type { Server, AddressInfo } from 'node:net'
 
@@ -559,8 +560,9 @@ describe('tcp replication', () => {
       "START_REPLICATION SLOT \"stream_test\" LOGICAL 0/0 (proto_version '1', publication_names 'zero_takeout')"
     )
 
-    // insert data right away - the poll loop will pick it up once it starts
+    // insert data right away - signal so the handler picks it up immediately
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('tcp_streamed', 42)`)
+    signalReplicationChange()
 
     // collect everything for long enough to catch the change
     const allRaw = await replClient.collectStream(3000)
@@ -626,9 +628,11 @@ describe('tcp replication', () => {
     await replClient.collectStream(200) // skip CopyBothResponse
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('upd_target', 10)`)
+    signalReplicationChange()
     await replClient.collectStream(1500)
 
     await db.exec(`UPDATE public.items SET value = 20 WHERE name = 'upd_target'`)
+    signalReplicationChange()
     const stream = await replClient.collectStream(1500)
 
     const decoded: PgOutputMessage[] = []
@@ -670,9 +674,11 @@ describe('tcp replication', () => {
     await replClient.collectStream(200)
 
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('del_target', 99)`)
+    signalReplicationChange()
     await replClient.collectStream(1500)
 
     await db.exec(`DELETE FROM public.items WHERE name = 'del_target'`)
+    signalReplicationChange()
     const stream = await replClient.collectStream(1500)
 
     const decoded: PgOutputMessage[] = []
@@ -715,6 +721,7 @@ describe('tcp replication', () => {
     // insert into both tables
     await db.exec(`INSERT INTO public.items (name, value) VALUES ('multi1', 1)`)
     await db.exec(`INSERT INTO public.products (label) VALUES ('multi2')`)
+    signalReplicationChange()
 
     // collect until we see both relations (with timeout)
     const allDecoded: PgOutputMessage[] = []
@@ -763,6 +770,7 @@ describe('tcp replication', () => {
     for (let i = 0; i < count; i++) {
       await db.exec(`INSERT INTO public.items (name, value) VALUES ('rapid${i}', ${i})`)
     }
+    signalReplicationChange()
 
     // give enough time for all changes to stream
     const stream = await replClient.collectStream(2000)

package/src/replication/zero-compat.test.ts

@@ -17,7 +17,8 @@ import { describe, it, expect, beforeEach, afterEach } from 'vitest'
 
 import { getConfig } from '../config'
 import { startPgProxy } from '../pg-proxy'
-import { installChangeTracking } from './change-tracker'
+import { installChangeTracking, resetShardSchemaCache } from './change-tracker'
+import { signalReplicationChange } from './handler'
 
 import type { Server, AddressInfo } from 'node:net'
 
@@ -447,6 +448,7 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
   let port: number
 
   beforeEach(async () => {
+    resetShardSchemaCache()
     db = new PGlite()
     await db.waitReady
     await db.exec(`
@@ -467,12 +469,23 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
     await db.exec(`CREATE PUBLICATION zero_pub FOR ALL TABLES`)
     await installChangeTracking(db)
 
+    // auto-signal the replication handler after every db.exec() call.
+    // in production, writes go through the TCP proxy which signals automatically.
+    // in tests, db.exec() bypasses the proxy, so we signal explicitly.
+    const origExec = db.exec.bind(db)
+    ;(db as any).exec = async (sql: string) => {
+      const result = await origExec(sql)
+      signalReplicationChange()
+      return result
+    }
+
     const config = { ...getConfig(), pgPort: 0 }
     server = await startPgProxy(db, config)
     port = (server.address() as AddressInfo).port
   })
 
   afterEach(async () => {
+    signalReplicationChange()
     server?.close()
     await db?.close()
   })
@@ -968,6 +981,9 @@ describe('zero-cache pgoutput compatibility', { timeout: 30000 }, () => {
     const s = await stream()
     const q = s.messages
 
+    // give handler time to finish setup (trigger installation)
+    await new Promise((r) => setTimeout(r, 300))
+
     // insert into all three shard tables + a public table
     await db.exec(
       `INSERT INTO chat_0.clients ("clientGroupID", "clientID", "lastMutationID") VALUES ('cg1', 'c1', 1)`