orez 0.5.9 → 0.5.12

package/dist/pg-proxy-browser.js

@@ -14,29 +14,10 @@
 import { PostgresConnection } from 'pg-gateway';
 import { log } from './log.js';
 import { Mutex } from './mutex.js';
-import { createReplicationFeedbackParser, handleReplicationQuery, handleStartReplication, signalReplicationChange, } from './replication/handler.js';
+import { createReplicationFeedbackParser, handleReplicationQuery, handleStartReplication, noteConfirmedFlushLsn, signalReplicationChange, } from './replication/handler.js';
 // shared encoder/decoder instances
 const textEncoder = new TextEncoder();
 const textDecoder = new TextDecoder();
-// verbose wire-protocol logging. off by default — every postgres message hits
-// the hot path, so unconditional console.debug calls are surprisingly expensive
-// even when devtools is set to hide them (the interpolated strings still get
-// built, and devtools retains them).
-//
-// re-evaluated on every call so you can flip it at runtime from any iframe,
-// worker, or console without restarting. toggle via either:
-//   - env:    OREZ_DEBUG_WIRE=1 (picked up at process start)
-//   - global: globalThis.__OREZ_DEBUG_WIRE__ = true (works anywhere, anytime)
-function isDebugWire() {
-    try {
-        if (globalThis.__OREZ_DEBUG_WIRE__ === true)
-            return true;
-        if (typeof process !== 'undefined' && process.env?.OREZ_DEBUG_WIRE)
-            return true;
-    }
-    catch { }
-    return false;
-}
 // postgres frontend/backend message type bytes — only used for debug-wire logging
 const PG_MSG_TYPE_NAMES = {
     0x50: 'Parse',
@@ -109,8 +90,6 @@ function extractQueryText(data) {
     }
     return null;
 }
-// abort previous replication handler when a new one starts
-let abortPreviousReplication = null;
 // clean version string: strip emscripten compiler info that breaks pg_restore/pg_dump
 const PG_VERSION_STRING = "'PostgreSQL 17.4 on x86_64-pc-linux-gnu, compiled by gcc (GCC) 12.2.0, 64-bit'";
 // query rewrites: make pglite look like real postgres with logical replication
@@ -552,7 +531,6 @@ function stripResponseMessages(data, stripRfq) {
  * readable receives Uint8Array messages from the port.
  * writable sends Uint8Array messages via the port.
  */
-let _globalWriteCount = 0;
 function messagePortToDuplexWithInject(port) {
     let readController;
     let msgCount = 0;
@@ -593,52 +571,6 @@ function messagePortToDuplexWithInject(port) {
     };
     return { duplex: { readable, writable }, rawWrite, injectMessage };
 }
-function messagePortToDuplex(port) {
-    let msgCount = 0;
-    const readable = new ReadableStream({
-        start(controller) {
-            port.onmessage = (ev) => {
-                msgCount++;
-                if (isDebugWire() && msgCount <= 3) {
-                    console.debug(`[pg-proxy-duplex] msg#${msgCount} type=${typeof ev.data} isAB=${ev.data instanceof ArrayBuffer} isU8=${ev.data instanceof Uint8Array} len=${ev.data?.byteLength ?? ev.data?.length ?? '?'}`);
-                }
-                if (ev.data instanceof ArrayBuffer) {
-                    controller.enqueue(new Uint8Array(ev.data));
-                }
-                else if (ev.data instanceof Uint8Array) {
-                    controller.enqueue(ev.data);
-                }
-                else {
-                    console.warn(`[pg-proxy-duplex] unexpected data type:`, typeof ev.data, ev.data);
-                }
-            };
-        },
-        cancel() {
-            port.close();
-        },
-    });
-    const writable = new WritableStream({
-        write(chunk) {
-            _globalWriteCount++;
-            if (isDebugWire() && _globalWriteCount <= 200) {
-                console.debug(`[pg-proxy-ws-write] #${_globalWriteCount} len=${chunk.byteLength}`);
-            }
-            // CF workerd MessagePort does not support transfer lists.
-            const buf = chunk.buffer.slice(chunk.byteOffset, chunk.byteOffset + chunk.byteLength);
-            port.postMessage(buf);
-        },
-        close() {
-            port.close();
-        },
-    });
-    // raw write function for injecting data outside of pg-gateway's stream
-    // (e.g. parameter status messages during onAuthenticated)
-    const rawWrite = (data) => {
-        const buf = data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength);
-        port.postMessage(buf);
-    };
-    return { duplex: { readable, writable }, rawWrite };
-}
 function hasProtocolSessionFactory(db) {
     return (typeof db.createProtocolSession === 'function');
 }
@@ -660,6 +592,8 @@ function createConnectionContext(ctx) {
     };
 }
 export async function createBrowserProxy(dbInput, config) {
+    const debugWire = config.debugWire === true;
+    const isDebugWire = () => debugWire;
     // normalize input: single PGlite instance = use it for all databases (backwards compat for tests)
     const instances = 'postgres' in dbInput
         ? dbInput
@@ -767,16 +701,39 @@ export async function createBrowserProxy(dbInput, config) {
         signalPending = true;
         queueMicrotask(() => {
             signalPending = false;
-            signalReplicationChange();
+            signalReplicationChange(config.instanceId);
         });
     }
     let closed = false;
+    let closePromise = null;
     const activeConnectionClosers = new Set();
+    const unreportedConnectionCleanupErrors = new Set();
+    let abortPreviousReplication = null;
+    const activeReplicationTasks = new Set();
+    const replicationRuntime = {
+        replace(abort) {
+            abortPreviousReplication?.();
+            abortPreviousReplication = abort;
+        },
+        track(promise, abort, onError) {
+            let tracked;
+            tracked = promise.catch(onError).finally(() => {
+                activeReplicationTasks.delete(tracked);
+                if (abortPreviousReplication === abort)
+                    abortPreviousReplication = null;
+            });
+            activeReplicationTasks.add(tracked);
+        },
+    };
     function handleConnection(port) {
         if (closed) {
             port.close();
             return;
         }
+        const closePendingConnection = async () => {
+            port.close();
+        };
+        activeConnectionClosers.add(closePendingConnection);
         port.start();
         // peek at the first message to detect replication connections.
         // replication connections bypass pg-gateway entirely and are handled
@@ -786,6 +743,10 @@ export async function createBrowserProxy(dbInput, config) {
         const buffered = [];
         let handlerInstalled = false;
         port.onmessage = (ev) => {
+            if (closed) {
+                port.close();
+                return;
+            }
             if (handlerInstalled)
                 return; // shouldn't happen — handler replaced port.onmessage
             buffered.push(ev);
@@ -795,6 +756,7 @@ export async function createBrowserProxy(dbInput, config) {
             if (!(data instanceof Uint8Array) || data.length < 8) {
                 handlerInstalled = true;
                 handleRegularConnection(port, ev);
+                activeConnectionClosers.delete(closePendingConnection);
                 // flush buffered messages to the new handler
                 for (let i = 1; i < buffered.length; i++) {
                     port.onmessage?.(buffered[i]);
@@ -813,6 +775,7 @@ export async function createBrowserProxy(dbInput, config) {
                 // in browser Web Workers (same root cause as patches #9, #18, #20).
                 handleRawConnection(port, data, params, getDbContext(dbName), isRepl);
                 handlerInstalled = true;
+                activeConnectionClosers.delete(closePendingConnection);
                 // flush any messages that arrived while we were processing the startup
                 for (let i = 1; i < buffered.length; i++) {
                     if (port.onmessage) {
@@ -855,43 +818,53 @@ export async function createBrowserProxy(dbInput, config) {
         const connId = {};
         const dbName = params.database || 'postgres';
         let connClosed = false;
-        let sessionClosed = false;
-        const closeSession = async () => {
-            if (sessionClosed)
-                return;
-            sessionClosed = true;
-            activeConnectionClosers.delete(closeSession);
-            // a connection torn down mid-pipeline still OWNS the db mutex. a dead
-            // pipeline can never send its Sync, so without this release the mutex
-            // is leaked forever and every later connection on this db (including
-            // the next embed generation after an idle-hibernation teardown) blocks
-            // in mutex.acquire() before its first statement. when a statement is
-            // still executing, the pipeline continuation performs this cleanup
-            // after it settles instead (pipelineExecInFlight).
-            if (pipelineMutexHeld && !pipelineExecInFlight) {
-                pipelineMutexHeld = false;
-                try {
-                    await rollbackOwnTransaction();
+        let closeSessionPromise = null;
+        const closeSession = () => {
+            if (closeSessionPromise)
+                return closeSessionPromise;
+            connClosed = true;
+            closeSessionPromise = (async () => {
+                port.close();
+                // a connection torn down mid-pipeline still OWNS the db mutex. a dead
+                // pipeline can never send its Sync, so without this release the mutex
+                // is leaked forever and every later connection on this db (including
+                // the next embed generation after an idle-hibernation teardown) blocks
+                // in mutex.acquire() before its first statement. when a statement is
+                // still executing, the pipeline continuation performs this cleanup
+                // after it settles instead (pipelineExecInFlight).
+                if (pipelineExecInFlight)
+                    await pipelineIdlePromise;
+                if (pipelineMutexHeld) {
+                    pipelineMutexHeld = false;
+                    try {
+                        await rollbackOwnTransaction();
+                    }
+                    finally {
+                        mutex.release();
+                    }
                 }
-                finally {
-                    mutex.release();
+                else {
+                    await mutex.acquire();
+                    try {
+                        await rollbackOwnTransaction();
+                    }
+                    finally {
+                        mutex.release();
+                    }
                 }
-            }
-            try {
                 await scopedContext.close();
-            }
-            catch { }
+            })();
+            void closeSessionPromise.then(() => activeConnectionClosers.delete(closeSession), (error) => {
+                activeConnectionClosers.delete(closeSession);
+                unreportedConnectionCleanupErrors.add(error);
+            });
+            return closeSessionPromise;
         };
         activeConnectionClosers.add(closeSession);
         // roll back this connection's open transaction. caller must hold the mutex.
         const rollbackOwnTransaction = async () => {
             if (txState.owner === connId && txState.status !== 0x49) {
-                try {
-                    await db.exec('ROLLBACK');
-                }
-                catch {
-                    // db is closed or rollback failed — ignore
-                }
+                await db.exec('ROLLBACK');
                 txState.status = 0x49;
                 txState.owner = null;
             }
@@ -959,6 +932,18 @@ export async function createBrowserProxy(dbInput, config) {
         // statement currently executing inside the held pipeline — closeSession
         // must not release the mutex under it (see closeSession above)
         let pipelineExecInFlight = false;
+        let resolvePipelineIdle = () => { };
+        let pipelineIdlePromise = Promise.resolve();
+        const beginPipelineExec = () => {
+            pipelineExecInFlight = true;
+            pipelineIdlePromise = new Promise((resolve) => {
+                resolvePipelineIdle = resolve;
+            });
+        };
+        const endPipelineExec = () => {
+            pipelineExecInFlight = false;
+            resolvePipelineIdle();
+        };
         let extWritePending = false;
         let pipelineBuffer = [];
         // remember the most recent Parse / SimpleQuery for diagnostic logging when
@@ -1029,9 +1014,15 @@ export async function createBrowserProxy(dbInput, config) {
                         .replace(/\0$/, '');
                     const upper = query.trim().toUpperCase();
                     if (upper.startsWith('START_REPLICATION')) {
-                        if (abortPreviousReplication)
-                            abortPreviousReplication();
                         let aborted = false;
+                        const abort = () => {
+                            if (aborted)
+                                return;
+                            aborted = true;
+                            connClosed = true;
+                            port.close();
+                            void closeSession();
+                        };
                         const writer = {
                             write(chunk) {
                                 if (!connClosed && !aborted) {
@@ -1047,20 +1038,15 @@ export async function createBrowserProxy(dbInput, config) {
                                 return connClosed || aborted;
                             },
                             close() {
-                                abortPreviousReplication?.();
+                                abort();
                             },
                         };
-                        abortPreviousReplication = () => {
-                            aborted = true;
-                            connClosed = true;
-                            port.close();
-                            void closeSession();
-                        };
+                        replicationRuntime.replace(abort);
                         // client→server traffic on a replication connection carries
                         // standby status updates: the consumer's confirmed-flush LSN
                         // drives purging of _zero_changes (rows are retained until the
                         // consumer durably commits them, like real pg WAL retention).
-                        const feedbackParser = createReplicationFeedbackParser();
+                        const feedbackParser = createReplicationFeedbackParser((lsn) => noteConfirmedFlushLsn(lsn, config.instanceId));
                         port.onmessage = (ev) => {
                             const incoming = ev.data instanceof ArrayBuffer
                                 ? new Uint8Array(ev.data)
@@ -1068,7 +1054,7 @@ export async function createBrowserProxy(dbInput, config) {
                             if (incoming instanceof Uint8Array)
                                 feedbackParser(incoming);
                         };
-                        handleStartReplication(query, writer, db, transactionAwareMutex(db, mutex, txState, null)).catch(() => { });
+                        replicationRuntime.track(handleStartReplication(query, writer, db, transactionAwareMutex(db, mutex, txState, null), config.instanceId), abort, () => { });
                         return;
                     }
                     // replication queries (IDENTIFY_SYSTEM, SET ..., CREATE/DROP SLOT).
@@ -1086,7 +1072,7 @@ export async function createBrowserProxy(dbInput, config) {
                     // open tx and commits with it.
                     await mutex.acquire();
                     try {
-                        const response = await handleReplicationQuery(query, db);
+                        const response = await handleReplicationQuery(query, db, config.instanceId);
                         if (response) {
                             write(response);
                             return;
@@ -1163,7 +1149,7 @@ export async function createBrowserProxy(dbInput, config) {
                             const combined = concatUint8Arrays(pipelineBuffer);
                             pipelineBuffer = [];
                             let flushResult;
-                            pipelineExecInFlight = true;
+                            beginPipelineExec();
                             try {
                                 flushResult = await db.execProtocolRaw(combined, { syncToFs: false });
                             }
@@ -1174,19 +1160,7 @@ export async function createBrowserProxy(dbInput, config) {
                                 return;
                             }
                             finally {
-                                pipelineExecInFlight = false;
-                            }
-                            if (sessionClosed && pipelineMutexHeld) {
-                                // session torn down while this statement was executing —
-                                // closeSession skipped its cleanup, finish it here
-                                pipelineMutexHeld = false;
-                                try {
-                                    await rollbackOwnTransaction();
-                                }
-                                finally {
-                                    mutex.release();
-                                }
-                                return;
+                                endPipelineExec();
                             }
                             flushResult = stripResponseMessages(flushResult, true);
                             write(flushResult);
@@ -1199,7 +1173,7 @@ export async function createBrowserProxy(dbInput, config) {
                     pipelineBuffer = [];
                     let result;
                     const t0 = performance.now();
-                    pipelineExecInFlight = true;
+                    beginPipelineExec();
                     try {
                         result = await db.execProtocolRaw(combined, { syncToFs: false });
                     }
@@ -1210,7 +1184,7 @@ export async function createBrowserProxy(dbInput, config) {
                         return;
                     }
                     finally {
-                        pipelineExecInFlight = false;
+                        endPipelineExec();
                     }
                     const dt = performance.now() - t0;
                     if (isDebugWire() && dt > 100)
@@ -1290,8 +1264,9 @@ export async function createBrowserProxy(dbInput, config) {
                     // signal writes
                     if (dbName === 'postgres' && msgType === 0x51) {
                         const qn = extractQueryText(data)?.trimStart().toLowerCase();
-                        if (qn && isWriteNormalized(qn))
-                            signalReplicationChange();
+                        if (qn && isWriteNormalized(qn)) {
+                            signalReplicationChange(config.instanceId);
+                        }
                     }
                     write(result);
                 }
@@ -1327,17 +1302,25 @@ export async function createBrowserProxy(dbInput, config) {
         // not release the mutex under it; the pipeline continuation finishes the
         // cleanup after the statement settles.
         let pipelineExecInFlight = false;
+        let resolvePipelineIdle = () => { };
+        let pipelineIdlePromise = Promise.resolve();
+        const beginPipelineExec = () => {
+            pipelineExecInFlight = true;
+            pipelineIdlePromise = new Promise((resolve) => {
+                resolvePipelineIdle = resolve;
+            });
+        };
+        const endPipelineExec = () => {
+            pipelineExecInFlight = false;
+            resolvePipelineIdle();
+        };
         // connection closed flag
         let connClosed = false;
+        let cleanupPromise = null;
         // roll back this connection's open transaction. caller must hold the mutex.
         const rollbackOwnTransaction = async (db, txState) => {
             if (txState.owner === connId && txState.status !== 0x49) {
-                try {
-                    await db.exec('ROLLBACK');
-                }
-                catch {
-                    // db is closed or rollback failed — ignore
-                }
+                await db.exec('ROLLBACK');
                 txState.status = 0x49;
                 txState.owner = null;
             }
@@ -1350,32 +1333,30 @@ export async function createBrowserProxy(dbInput, config) {
         // pollution that 2eb1873 fixed on the node side. browser cleanup was
         // missed and showed up as "current transaction is aborted, commands
         // ignored" when a foreign connection closed mid-someone-else's-txn.
-        const cleanup = async () => {
-            if (connClosed)
-                return;
+        const cleanup = () => {
+            if (cleanupPromise)
+                return cleanupPromise;
             connClosed = true;
-            activeConnectionClosers.delete(cleanup);
-            // replication connections don't own a transaction — skip ROLLBACK
-            if (isReplicationConnection)
-                return;
-            try {
+            cleanupPromise = (async () => {
+                port.close();
+                // replication connections don't own a transaction — skip ROLLBACK
+                if (isReplicationConnection)
+                    return;
                 const { db, mutex, txState } = getDbContext(dbName);
+                if (pipelineExecInFlight)
+                    await pipelineIdlePromise;
                 if (pipelineMutexHeld) {
                     // this connection died mid-pipeline and still OWNS the mutex. a
                     // dead pipeline can never send its Sync, so without this release
                     // the mutex is leaked forever and every later connection on this
                     // db (including the next embed generation after idle-hibernation
                     // teardown) blocks in mutex.acquire() before its first statement.
-                    // when a statement is in flight, the pipeline continuation does
-                    // this exact cleanup after it settles.
-                    if (!pipelineExecInFlight) {
-                        pipelineMutexHeld = false;
-                        try {
-                            await rollbackOwnTransaction(db, txState);
-                        }
-                        finally {
-                            mutex.release();
-                        }
+                    pipelineMutexHeld = false;
+                    try {
+                        await rollbackOwnTransaction(db, txState);
+                    }
+                    finally {
+                        mutex.release();
                     }
                     return;
                 }
@@ -1386,10 +1367,12 @@ export async function createBrowserProxy(dbInput, config) {
                 finally {
                     mutex.release();
                 }
-            }
-            catch {
-                // instance may have been replaced during reset, ignore
-            }
+            })();
+            void cleanupPromise.then(() => activeConnectionClosers.delete(cleanup), (error) => {
+                activeConnectionClosers.delete(cleanup);
+                unreportedConnectionCleanupErrors.add(error);
+            });
+            return cleanupPromise;
         };
         activeConnectionClosers.add(cleanup);
         try {
@@ -1452,7 +1435,7 @@ export async function createBrowserProxy(dbInput, config) {
                         return handleReplicationMessageBrowser(data, rawWrite, () => connClosed, () => {
                             connClosed = true;
                             port.close();
-                        }, instances.postgres, mutexes.postgres, connection);
+                        }, instances.postgres, mutexes.postgres, connection, config.instanceId, replicationRuntime, debugWire);
                     }
                     const msgType = data[0];
                     const { db, mutex, txState } = getDbContext(dbName);
@@ -1486,7 +1469,7 @@ export async function createBrowserProxy(dbInput, config) {
                         data = interceptQuery(data);
                         const t1 = performance.now();
                         let result;
-                        pipelineExecInFlight = true;
+                        beginPipelineExec();
                         try {
                             result = await db.execProtocolRaw(data, { syncToFs: false });
                         }
@@ -1496,7 +1479,7 @@ export async function createBrowserProxy(dbInput, config) {
                             throw err;
                         }
                         finally {
-                            pipelineExecInFlight = false;
+                            endPipelineExec();
                         }
                         const t2 = performance.now();
                         proxyStats.totalExecMs += t2 - t1;
@@ -1518,18 +1501,6 @@ export async function createBrowserProxy(dbInput, config) {
                                 signalWrite();
                             }
                         }
-                        else if (connClosed) {
-                            // the connection died while this statement was executing, so
-                            // cleanup skipped its release (pipelineExecInFlight). no Sync
-                            // can ever arrive — roll back and release here instead.
-                            pipelineMutexHeld = false;
-                            try {
-                                await rollbackOwnTransaction(db, txState);
-                            }
-                            finally {
-                                mutex.release();
-                            }
-                        }
                         else {
                             // strip ReadyForQuery from non-Sync pipeline messages
                             // result = stripResponseMessages(result, true) // disabled for debugging
@@ -1627,7 +1598,7 @@ export async function createBrowserProxy(dbInput, config) {
                     result = stripResponseMessages(result, stripRfq);
                     // signal replication handler on postgres writes for instant sync
                     if (dbName === 'postgres' && queryNorm && isWriteNormalized(queryNorm)) {
-                        signalReplicationChange();
+                        signalReplicationChange(config.instanceId);
                     }
                     return result;
                 },
@@ -1707,11 +1678,59 @@ export async function createBrowserProxy(dbInput, config) {
         txState.owner = owner;
     }
     const defaultExternalOwner = {};
+    const closedExternalOwners = new WeakSet();
+    function createExternalOwnerCloser(owner, dbNames) {
+        let closing = null;
+        const closeOwner = () => {
+            if (closing)
+                return closing;
+            closedExternalOwners.add(owner);
+            closing = (async () => {
+                const contexts = new Map();
+                for (const dbName of dbNames) {
+                    const context = getDbContext(dbName);
+                    contexts.set(context.txState, context);
+                }
+                const results = await Promise.allSettled([...contexts.values()].map(async ({ db, mutex, txState }) => {
+                    await mutex.acquire();
+                    try {
+                        if (txState.owner === owner && txState.status !== 0x49) {
+                            await db.exec('ROLLBACK');
+                            txState.status = 0x49;
+                            txState.owner = null;
+                        }
+                    }
+                    finally {
+                        mutex.release();
+                    }
+                }));
+                const errors = results.flatMap((result) => result.status === 'rejected' ? [result.reason] : []);
+                if (errors.length === 1)
+                    throw errors[0];
+                if (errors.length > 1) {
+                    throw new AggregateError(errors, 'browser proxy external cleanup failed');
+                }
+            })();
+            void closing.then(() => activeConnectionClosers.delete(closeOwner), (error) => {
+                activeConnectionClosers.delete(closeOwner);
+                unreportedConnectionCleanupErrors.add(error);
+            });
+            return closing;
+        };
+        activeConnectionClosers.add(closeOwner);
+        return closeOwner;
+    }
+    createExternalOwnerCloser(defaultExternalOwner, ['postgres', 'zero_cvr', 'zero_cdb']);
     async function externalQuery(dbName, sql, params, owner = defaultExternalOwner) {
-        if (closed)
+        if (closed || closedExternalOwners.has(owner)) {
             throw new Error('pg-proxy is closed');
+        }
         const { db, mutex, txState } = getDbContext(dbName);
         await acquireForOwner(db, mutex, txState, owner);
+        if (closed || closedExternalOwners.has(owner)) {
+            mutex.release();
+            throw new Error('pg-proxy is closed');
+        }
         try {
             const result = await db.query(sql, params);
             updateExternalTxStateAfterSuccess(txState, owner, sql);
@@ -1730,10 +1749,15 @@ export async function createBrowserProxy(dbInput, config) {
         }
     }
     async function externalExec(dbName, sql, owner = defaultExternalOwner) {
-        if (closed)
+        if (closed || closedExternalOwners.has(owner)) {
             throw new Error('pg-proxy is closed');
+        }
         const { db, mutex, txState } = getDbContext(dbName);
         await acquireForOwner(db, mutex, txState, owner);
+        if (closed || closedExternalOwners.has(owner)) {
+            mutex.release();
+            throw new Error('pg-proxy is closed');
+        }
         try {
             const result = await db.exec(sql);
             updateExternalTxStateAfterSuccess(txState, owner, sql);
@@ -1752,59 +1776,68 @@ export async function createBrowserProxy(dbInput, config) {
     }
     function createExternalSession(dbName = 'postgres') {
         const owner = {};
+        let sessionClosed = false;
+        const closeOwner = createExternalOwnerCloser(owner, [dbName]);
         return {
             query(sql, params) {
+                if (sessionClosed)
+                    return Promise.reject(new Error('pg-proxy session is closed'));
                 return externalQuery(dbName, sql, params, owner);
             },
             exec(sql) {
+                if (sessionClosed)
+                    return Promise.reject(new Error('pg-proxy session is closed'));
                 return externalExec(dbName, sql, owner);
             },
-            async close() {
-                const { db, mutex, txState } = getDbContext(dbName);
-                await mutex.acquire();
-                try {
-                    if (txState.owner === owner && txState.status !== 0x49) {
-                        try {
-                            await db.exec('ROLLBACK');
-                        }
-                        catch { }
-                        txState.status = 0x49;
-                        txState.owner = null;
-                    }
-                }
-                finally {
-                    mutex.release();
-                }
+            close() {
+                sessionClosed = true;
+                return closeOwner();
             },
         };
     }
     return {
         handleConnection,
         close() {
+            if (closePromise)
+                return closePromise;
             closed = true;
             signalPending = false;
-            for (const closeConnection of activeConnectionClosers) {
-                void closeConnection();
-            }
+            closePromise = (async () => {
+                abortPreviousReplication?.();
+                abortPreviousReplication = null;
+                signalReplicationChange(config.instanceId);
+                const results = await Promise.allSettled([...activeConnectionClosers].map((closeConnection) => Promise.resolve().then(() => closeConnection())));
+                const errors = new Set(unreportedConnectionCleanupErrors);
+                for (const result of results) {
+                    if (result.status === 'rejected')
+                        errors.add(result.reason);
+                }
+                await Promise.all([...activeReplicationTasks]);
+                unreportedConnectionCleanupErrors.clear();
+                if (errors.size > 0) {
+                    throw new AggregateError([...errors], 'browser proxy connection cleanup failed');
+                }
+            })();
+            return closePromise;
         },
         createExternalSession,
         query: externalQuery,
         exec: externalExec,
     };
 }
-async function handleReplicationMessageBrowser(data, rawWrite, isClosed, closeConn, db, mutex, connection) {
-    if (isDebugWire()) {
+async function handleReplicationMessageBrowser(data, rawWrite, isClosed, closeConn, db, mutex, connection, instanceId, replicationRuntime, debugWire) {
+    if (debugWire) {
         console.debug(`[pg-proxy-repl] ENTRY type=0x${data[0].toString(16)} len=${data.length}`);
     }
     // for non-SimpleQuery messages (extended protocol), execute against PGlite directly.
     if (data[0] !== 0x51) {
-        if (isDebugWire()) {
+        if (debugWire) {
             console.debug(`[pg-proxy-repl] ext protocol msg type=0x${data[0].toString(16)} len=${data.length}`);
         }
         await mutex.acquire();
         try {
             const result = await db.execProtocolRaw(data, { syncToFs: false });
-            if (isDebugWire()) {
+            if (debugWire) {
                 console.debug(`[pg-proxy-repl] ext protocol result len=${result.length}`);
             }
             return result;
@@ -1820,11 +1853,6 @@ async function handleReplicationMessageBrowser(data, rawWrite, isClosed, closeCo
     // check if this is a START_REPLICATION command
     if (upper.startsWith('START_REPLICATION')) {
         await connection.detach();
-        // abort any previous replication handler to prevent zombies
-        if (abortPreviousReplication) {
-            log.proxy('aborting previous replication handler');
-            abortPreviousReplication();
-        }
         let aborted = false;
         const writer = {
             write(chunk) {
@@ -1846,30 +1874,32 @@ async function handleReplicationMessageBrowser(data, rawWrite, isClosed, closeCo
             },
         };
         const abort = () => {
+            if (aborted)
+                return;
             aborted = true;
             closeConn();
         };
-        abortPreviousReplication = abort;
-        handleStartReplication(query, writer, db, mutex).catch((err) => {
+        replicationRuntime.replace(abort);
+        replicationRuntime.track(handleStartReplication(query, writer, db, mutex, instanceId), abort, (err) => {
             log.proxy(`replication stream ended: ${err}`);
         });
         return undefined;
     }
     // handle replication queries + fallthrough to pglite, all under mutex
-    if (isDebugWire()) {
+    if (debugWire) {
         console.debug(`[pg-proxy-repl] query: ${query.slice(0, 100)}`);
         console.debug(`[pg-proxy-repl] acquiring mutex...`);
     }
     await mutex.acquire();
-    if (isDebugWire())
+    if (debugWire)
         console.debug(`[pg-proxy-repl] mutex acquired, testing db access...`);
     try {
         const testResult = await db.query('SELECT 1 as test');
-        if (isDebugWire()) {
+        if (debugWire) {
             console.debug(`[pg-proxy-repl] db.query works: ${JSON.stringify(testResult.rows)}`);
         }
-        const response = await handleReplicationQuery(query, db);
-        if (isDebugWire()) {
+        const response = await handleReplicationQuery(query, db, instanceId);
+        if (debugWire) {
             console.debug(`[pg-proxy-repl] handleReplicationQuery result: ${response ? 'bytes(' + response.length + ')' : 'null'}`);
         }
         if (response)