orez 0.1.13 → 0.1.15

package/src/replication/handler.ts

@@ -7,9 +7,10 @@
  */
 
 import { log } from '../log.js'
+
+const textEncoder = new TextEncoder()
 import {
   getChangesSince,
-  getCurrentWatermark,
   purgeConsumedChanges,
   installTriggersOnShardTables,
   type ChangeRecord,
@@ -22,11 +23,9 @@ import {
   encodeUpdate,
   encodeDelete,
   encodeKeepalive,
-  wrapXLogData,
-  wrapCopyData,
+  encodeWrappedChange,
   getTableOid,
   inferColumns,
-  type ColumnInfo,
 } from './pgoutput-encoder.js'
 
 import type { Mutex } from '../mutex.js'
@@ -34,6 +33,7 @@ import type { PGlite } from '@electric-sql/pglite'
 
 export interface ReplicationWriter {
   write(data: Uint8Array): void
+  readonly closed?: boolean
 }
 
 // current lsn counter
@@ -56,7 +56,7 @@ function nowMicros(): bigint {
 // build a wire protocol row description + data row response
 function buildSimpleResponse(columns: string[], values: string[]): Uint8Array {
   const parts: Uint8Array[] = []
-  const encoder = new TextEncoder()
+  const encoder = textEncoder
 
   // RowDescription (0x54)
   let rdSize = 6 // int32 len + int16 numFields
@@ -140,7 +140,7 @@ function buildSimpleResponse(columns: string[], values: string[]): Uint8Array {
 }
 
 function buildCommandComplete(tag: string): Uint8Array {
-  const encoder = new TextEncoder()
+  const encoder = textEncoder
   const tagBytes = encoder.encode(tag + '\0')
   const cc = new Uint8Array(1 + 4 + tagBytes.length)
   cc[0] = 0x43
@@ -159,7 +159,7 @@ function buildCommandComplete(tag: string): Uint8Array {
 }
 
 function buildErrorResponse(message: string): Uint8Array {
-  const encoder = new TextEncoder()
+  const encoder = textEncoder
   const msgBytes = encoder.encode(message)
   // S(severity) + M(message) + null terminator
   const fields = new Uint8Array(2 + 6 + 2 + msgBytes.length + 1 + 1) // S + ERROR\0 + M + msg\0 + terminator
@@ -334,6 +334,9 @@ export async function handleStartReplication(
         FOR EACH STATEMENT EXECUTE FUNCTION public._zero_notify_change();
     `)
     }
+    if (tables.length > 0) {
+      log.proxy(`installed notify triggers on ${tables.length} public table(s)`)
+    }
 
     // discover shard schemas (e.g. chat_0) and install NOTIFY triggers
     const shardSchemas = await db.query<{ nspname: string }>(
@@ -463,34 +466,76 @@ export async function handleStartReplication(
   const sentRelations = new Set<string>()
   let txCounter = 1
 
-  // polling + notification loop
-  // adaptive: poll fast when catching up, slow when idle
-  const pollIntervalIdle = 500
-  const pollIntervalCatchUp = 20
+  // event-driven replication with promise-based wakeup
+  // uses pg_notify to wake up immediately, polling only as fallback
+  const pollIntervalIdle = 200
   const batchSize = 2000
-  const purgeEveryN = 10
+  const purgeEveryN = 5
   const shardRescanEveryN = 20
   let running = true
   let pollsSincePurge = 0
   let pollsSinceShardRescan = 0
 
+  // promise-based wakeup mechanism
+  let wakeupResolve: (() => void) | null = null
+  const wakeup = () => {
+    if (wakeupResolve) {
+      wakeupResolve()
+      wakeupResolve = null
+    }
+  }
+  const waitForWakeup = (timeoutMs: number): Promise<void> => {
+    return new Promise((resolve) => {
+      const timer = setTimeout(() => {
+        wakeupResolve = null
+        resolve()
+      }, timeoutMs)
+      wakeupResolve = () => {
+        clearTimeout(timer)
+        resolve()
+      }
+    })
+  }
+
+  // set up LISTEN to wake up immediately on changes
+  let unsubscribe: (() => Promise<void>) | null = null
+  try {
+    unsubscribe = await db.listen('_zero_changes', wakeup)
+    log.debug.proxy('replication: listening for _zero_changes notifications')
+  } catch {
+    log.debug.proxy('replication: LISTEN not available, using polling only')
+  }
+
   const poll = async () => {
     while (running) {
+      // check if the connection was closed
+      if (writer.closed) {
+        log.debug.proxy('replication: writer closed, exiting poll loop')
+        running = false
+        break
+      }
+
       try {
         // periodically re-scan for new shard schemas (e.g. chat_0 created by zero-cache)
         pollsSinceShardRescan++
         if (pollsSinceShardRescan >= shardRescanEveryN) {
-          pollsSinceShardRescan = 0
-          await mutex.acquire()
-          try {
-            await installTriggersOnShardTables(db)
-          } finally {
-            mutex.release()
+          if (mutex.tryAcquire()) {
+            pollsSinceShardRescan = 0
+            try {
+              await installTriggersOnShardTables(db)
+            } finally {
+              mutex.release()
+            }
           }
         }
 
-        // acquire mutex to avoid conflicting with proxy connections
-        await mutex.acquire()
+        // try to acquire mutex without blocking proxy connections.
+        // if the mutex is busy (proxy handling zero-cache queries), skip
+        // this iteration to avoid starving initial sync.
+        if (!mutex.tryAcquire()) {
+          await waitForWakeup(pollIntervalIdle)
+          continue
+        }
         let changes: Awaited<ReturnType<typeof getChangesSince>>
         try {
           changes = await getChangesSince(db, lastWatermark, batchSize)
@@ -531,9 +576,8 @@ export async function handleStartReplication(
 
           // purge consumed changes periodically to free wasm memory
           pollsSincePurge++
-          if (pollsSincePurge >= purgeEveryN) {
+          if (pollsSincePurge >= purgeEveryN && mutex.tryAcquire()) {
             pollsSincePurge = 0
-            await mutex.acquire()
             try {
               const purged = await purgeConsumedChanges(db, lastWatermark)
               if (purged > 0) {
@@ -543,15 +587,17 @@ export async function handleStartReplication(
               mutex.release()
             }
           }
+
+          // got changes - continue immediately to check for more
+          continue
         }
 
         // send keepalive
         const ts = nowMicros()
         writer.write(encodeKeepalive(currentLsn, ts, false))
 
-        // if we got a full batch, there's likely more - poll fast
-        const delay = changes.length >= batchSize ? pollIntervalCatchUp : pollIntervalIdle
-        await new Promise((resolve) => setTimeout(resolve, delay))
+        // no changes: wait for notify signal or poll interval
+        await waitForWakeup(pollIntervalIdle)
       } catch (err: unknown) {
         const msg = err instanceof Error ? err.message : String(err)
         log.debug.proxy(`replication poll error: ${msg}`)
@@ -565,10 +611,19 @@ export async function handleStartReplication(
   }
 
   log.debug.proxy('replication: starting poll loop')
-  await poll()
+  try {
+    await poll()
+  } finally {
+    if (unsubscribe) {
+      await unsubscribe().catch(() => {})
+    }
+  }
   log.debug.proxy('replication: poll loop exited')
 }
 
+// cache column info per table to avoid per-change allocation
+const cachedColumns = new Map<string, ReturnType<typeof inferColumns>>()
+
 async function streamChanges(
   changes: ChangeRecord[],
   writer: ReplicationWriter,
@@ -582,8 +637,7 @@ async function streamChanges(
   const lsn = nextLsn()
 
   // BEGIN
-  const beginMsg = wrapXLogData(lsn, lsn, ts, encodeBegin(lsn, ts, txId))
-  writer.write(wrapCopyData(beginMsg))
+  writer.write(encodeWrappedChange(lsn, lsn, ts, encodeBegin(lsn, ts, txId)))
 
   for (const change of changes) {
     // parse schema-qualified name (schema.table or bare table)
@@ -626,18 +680,23 @@ async function streamChanges(
     const row = rowData || oldData
     if (!row) continue
 
-    const keySet = tableKeyColumns.get(qualifiedKey)
-    const typeOids = columnTypeOids.get(qualifiedKey)
-    const columns = inferColumns(row).map((col) => ({
-      ...col,
-      typeOid: typeOids?.get(col.name) ?? col.typeOid,
-      isKey: keySet?.has(col.name) ?? false,
-    }))
+    // use cached columns or build and cache them
+    let columns = cachedColumns.get(qualifiedKey)
+    if (!columns) {
+      const keySet = tableKeyColumns.get(qualifiedKey)
+      const typeOids = columnTypeOids.get(qualifiedKey)
+      columns = inferColumns(row).map((col) => ({
+        ...col,
+        typeOid: typeOids?.get(col.name) ?? col.typeOid,
+        isKey: keySet?.has(col.name) ?? false,
+      }))
+      cachedColumns.set(qualifiedKey, columns)
+    }
 
     // send RELATION if not yet sent
     if (!sentRelations.has(qualifiedKey)) {
       const relMsg = encodeRelation(tableOid, schema, tableName, 0x64, columns)
-      writer.write(wrapCopyData(wrapXLogData(lsn, lsn, ts, relMsg)))
+      writer.write(encodeWrappedChange(lsn, lsn, ts, relMsg))
       sentRelations.add(qualifiedKey)
     }
 
@@ -660,13 +719,12 @@ async function streamChanges(
         continue
     }
 
-    writer.write(wrapCopyData(wrapXLogData(lsn, lsn, ts, changeMsg)))
+    writer.write(encodeWrappedChange(lsn, lsn, ts, changeMsg))
   }
 
   // COMMIT
   const endLsn = nextLsn()
-  const commitMsg = wrapXLogData(endLsn, endLsn, ts, encodeCommit(0, lsn, endLsn, ts))
-  writer.write(wrapCopyData(commitMsg))
+  writer.write(encodeWrappedChange(endLsn, endLsn, ts, encodeCommit(0, lsn, endLsn, ts)))
 }
 
 function normalizeShardClientsRow(

package/src/replication/pgoutput-encoder.ts

@@ -10,6 +10,9 @@
 // postgres epoch: 2000-01-01 in microseconds from unix epoch
 const PG_EPOCH_MICROS = 946684800000000n
 
+// shared encoder instance - avoids per-call allocation
+const encoder = new TextEncoder()
+
 // table oid tracking
 const tableOids = new Map<string, number>()
 let nextOid = 16384
@@ -39,10 +42,27 @@ export function inferColumns(row: Record<string, unknown>): ColumnInfo[] {
   }))
 }
 
-function encodeString(str: string): Uint8Array {
-  return new TextEncoder().encode(str)
+// reusable scratch buffer for building messages (64KB, grows if needed)
+let scratch = new Uint8Array(65536)
+let scratchView = new DataView(scratch.buffer)
+
+function ensureScratch(size: number): void {
+  if (scratch.length < size) {
+    const newSize = Math.max(size, scratch.length * 2)
+    scratch = new Uint8Array(newSize)
+    scratchView = new DataView(scratch.buffer)
+  }
+}
+
+function writeInt16At(offset: number, val: number): void {
+  scratchView.setInt16(offset, val)
+}
+
+function writeInt32At(offset: number, val: number): void {
+  scratchView.setInt32(offset, val)
 }
 
+// legacy helpers for standalone buffers
 function writeInt16(buf: Uint8Array, offset: number, val: number): void {
   new DataView(buf.buffer, buf.byteOffset).setInt16(offset, val)
 }
@@ -89,14 +109,14 @@ export function encodeRelation(
   replicaIdentity: number,
   columns: ColumnInfo[]
 ): Uint8Array {
-  const schemaBytes = encodeString(schema)
-  const nameBytes = encodeString(tableName)
+  const schemaBytes = encoder.encode(schema)
+  const nameBytes = encoder.encode(tableName)
 
   // calculate column sizes
   let columnsSize = 0
   const colNameBytes: Uint8Array[] = []
   for (const col of columns) {
-    const nb = encodeString(col.name)
+    const nb = encoder.encode(col.name)
     colNameBytes.push(nb)
     columnsSize += 1 + nb.length + 1 + 4 + 4 // flags + name + null + typeOid + typeMod
   }
@@ -133,19 +153,25 @@ export function encodeRelation(
   return buf
 }
 
-function encodeTupleData(
+// encode tuple data directly into scratch buffer starting at given offset.
+// returns the number of bytes written.
+function encodeTupleDataInto(
   row: Record<string, unknown>,
-  columns: ColumnInfo[]
-): Uint8Array {
-  const parts: Uint8Array[] = []
-  let totalSize = 2 // ncolumns (int16)
+  columns: ColumnInfo[],
+  startOffset: number
+): number {
+  let pos = startOffset
+
+  // reserve space for ncolumns
+  ensureScratch(pos + 2 + columns.length * 32)
+  writeInt16At(pos, columns.length)
+  pos += 2
 
-  const values: (Uint8Array | null)[] = []
   for (const col of columns) {
     const val = row[col.name]
     if (val === null || val === undefined) {
-      values.push(null)
-      totalSize += 1 // 'n' byte
+      ensureScratch(pos + 1)
+      scratch[pos++] = 0x6e // 'n' for null
     } else {
       // convert to postgresql text format
       let strVal: string
@@ -156,28 +182,48 @@ function encodeTupleData(
       } else {
         strVal = String(val)
       }
-      const bytes = encodeString(strVal)
-      values.push(bytes)
-      totalSize += 1 + 4 + bytes.length // 't' + len + data
+      const bytes = encoder.encode(strVal)
+      ensureScratch(pos + 1 + 4 + bytes.length)
+      scratch[pos++] = 0x74 // 't' for text
+      writeInt32At(pos, bytes.length)
+      pos += 4
+      scratch.set(bytes, pos)
+      pos += bytes.length
     }
   }
 
+  return pos - startOffset
+}
+
+/**
+ * encode a complete change message wrapped in CopyData(XLogData(...)).
+ * avoids intermediate buffer allocations by writing directly into one buffer.
+ */
+export function encodeWrappedChange(
+  walStart: bigint,
+  walEnd: bigint,
+  timestamp: bigint,
+  changeData: Uint8Array
+): Uint8Array {
+  // CopyData header: 'd' + int32 len
+  // XLogData header: 'w' + int64 walStart + int64 walEnd + int64 timestamp
+  // then changeData
+  const xlogSize = 1 + 8 + 8 + 8 + changeData.length
+  const totalSize = 1 + 4 + xlogSize
   const buf = new Uint8Array(totalSize)
-  let pos = 0
-  writeInt16(buf, pos, columns.length)
-  pos += 2
 
-  for (const val of values) {
-    if (val === null) {
-      buf[pos++] = 0x6e // 'n' for null
-    } else {
-      buf[pos++] = 0x74 // 't' for text
-      writeInt32(buf, pos, val.length)
-      pos += 4
-      buf.set(val, pos)
-      pos += val.length
-    }
-  }
+  // CopyData
+  buf[0] = 0x64 // 'd'
+  writeInt32(buf, 1, 4 + xlogSize)
+
+  // XLogData
+  buf[5] = 0x77 // 'w'
+  writeInt64(buf, 6, walStart)
+  writeInt64(buf, 14, walEnd)
+  writeInt64(buf, 22, timestamp - PG_EPOCH_MICROS)
+
+  // change payload
+  buf.set(changeData, 30)
 
   return buf
 }
@@ -188,13 +234,14 @@ export function encodeInsert(
   row: Record<string, unknown>,
   columns: ColumnInfo[]
 ): Uint8Array {
-  const tuple = encodeTupleData(row, columns)
-  const buf = new Uint8Array(1 + 4 + 1 + tuple.length)
-  buf[0] = 0x49 // 'I'
-  writeInt32(buf, 1, tableOid)
-  buf[5] = 0x4e // 'N' for new tuple
-  buf.set(tuple, 6)
-  return buf
+  // write header + tuple directly into scratch
+  const headerSize = 1 + 4 + 1 // 'I' + oid + 'N'
+  ensureScratch(headerSize + 2 + columns.length * 64)
+  scratch[0] = 0x49 // 'I'
+  writeInt32At(1, tableOid)
+  scratch[5] = 0x4e // 'N' for new tuple
+  const tupleLen = encodeTupleDataInto(row, columns, 6)
+  return scratch.slice(0, 6 + tupleLen)
 }
 
 // encode an UPDATE message
@@ -204,26 +251,21 @@ export function encodeUpdate(
   oldRow: Record<string, unknown> | null,
   columns: ColumnInfo[]
 ): Uint8Array {
-  const newTuple = encodeTupleData(row, columns)
+  ensureScratch(1 + 4 + 1 + columns.length * 128)
+  scratch[0] = 0x55 // 'U'
+  writeInt32At(1, tableOid)
 
   if (oldRow) {
-    const oldTuple = encodeTupleData(oldRow, columns)
-    const buf = new Uint8Array(1 + 4 + 1 + oldTuple.length + 1 + newTuple.length)
-    buf[0] = 0x55 // 'U'
-    writeInt32(buf, 1, tableOid)
-    buf[5] = 0x4f // 'O' for old tuple
-    buf.set(oldTuple, 6)
-    buf[6 + oldTuple.length] = 0x4e // 'N' for new tuple
-    buf.set(newTuple, 7 + oldTuple.length)
-    return buf
+    scratch[5] = 0x4f // 'O' for old tuple
+    const oldLen = encodeTupleDataInto(oldRow, columns, 6)
+    scratch[6 + oldLen] = 0x4e // 'N' for new tuple
+    const newLen = encodeTupleDataInto(row, columns, 7 + oldLen)
+    return scratch.slice(0, 7 + oldLen + newLen)
   }
 
-  const buf = new Uint8Array(1 + 4 + 1 + newTuple.length)
-  buf[0] = 0x55 // 'U'
-  writeInt32(buf, 1, tableOid)
-  buf[5] = 0x4e // 'N'
-  buf.set(newTuple, 6)
-  return buf
+  scratch[5] = 0x4e // 'N'
+  const newLen = encodeTupleDataInto(row, columns, 6)
+  return scratch.slice(0, 6 + newLen)
 }
 
 // encode a DELETE message
@@ -232,13 +274,12 @@ export function encodeDelete(
   oldRow: Record<string, unknown>,
   columns: ColumnInfo[]
 ): Uint8Array {
-  const tuple = encodeTupleData(oldRow, columns)
-  const buf = new Uint8Array(1 + 4 + 1 + tuple.length)
-  buf[0] = 0x44 // 'D'
-  writeInt32(buf, 1, tableOid)
-  buf[5] = 0x4b // 'K' for key tuple
-  buf.set(tuple, 6)
-  return buf
+  ensureScratch(1 + 4 + 1 + columns.length * 64)
+  scratch[0] = 0x44 // 'D'
+  writeInt32At(1, tableOid)
+  scratch[5] = 0x4b // 'K' for key tuple
+  const tupleLen = encodeTupleDataInto(oldRow, columns, 6)
+  return scratch.slice(0, 6 + tupleLen)
 }
 
 // wrap a pgoutput message in XLogData format