diff --git a/docs/runtime/sql.mdx b/docs/runtime/sql.mdx index ddef2b99e6f5..bd52220a3ab9 100644 --- a/docs/runtime/sql.mdx +++ b/docs/runtime/sql.mdx @@ -43,6 +43,7 @@ const sqliteResults = await sqlite` - Binary protocol support makes it faster - TLS support (and auth mode) - Automatic configuration with environment variables +- PostgreSQL `LISTEN` / `NOTIFY` --- @@ -984,6 +985,63 @@ try { --- +## LISTEN / NOTIFY (PostgreSQL) + +`sql.listen()` subscribes to a PostgreSQL notification channel and `sql.notify()` publishes to one. Processes sharing a database can use this as a lightweight message bus: cache invalidation, waking workers when rows are inserted (typically from a trigger calling `pg_notify`), fan-out of small events. + +```ts +import { sql } from "bun"; + +const subscription = await sql.listen("orders", payload => { + console.log("new order", JSON.parse(payload)); +}); + +await sql.notify("orders", JSON.stringify({ id: 42 })); + +await subscription.unlisten(); +``` + +`listen()` resolves once PostgreSQL has acknowledged the `LISTEN`, so a `notify()` issued after it resolves is delivered. The subscription it resolves to is also an async disposable: + +```ts +{ + await using subscription = await sql.listen("orders", handleOrder); + await doWork(); +} // unlistened here +``` + +### How it works + +- All subscriptions on a client share one dedicated connection. The first `listen()` opens it and removing the last subscription closes it, so a client that never listens never pays for it, and unlistening everything lets the process exit without `sql.close()`. +- While anything is subscribed, that connection keeps the process alive, like a listening server. +- If the connection drops, it is re-established with exponential backoff (250ms doubling to 32s, with jitter) and every channel is re-subscribed. PostgreSQL only delivers to connected listeners, so notifications sent in between are lost; the optional third argument to `listen()` runs on the initial subscribe and after every reconnect, which is the place to catch up: + +```ts +await sql.listen("orders", handleOrder, async () => { + for (const order of await sql`SELECT * FROM orders WHERE processed = false`) handleOrder(order); +}); +``` + +- Every `listen()` call is its own subscription. Several on one channel share a single server-side `LISTEN`, each callback receives every notification, and each handle's `unlisten()` removes only what its own call registered. A callback (either argument) that throws is reported as an uncaught exception and stays subscribed. +- Channel names are quoted as identifiers for you; like any PostgreSQL identifier they are limited to 63 bytes, and longer names are rejected rather than silently truncated. PostgreSQL limits payloads to 8000 bytes by default. + +### `notify()` + +`notify()` is an ordinary query (`SELECT pg_notify($1, $2)`) on whichever handle you call it through. On `sql` it uses the pool; inside `sql.begin()` it runs in the transaction, so PostgreSQL delivers it on `COMMIT` and drops it on `ROLLBACK`, which is how to announce a change only once it is visible: + +```ts +await sql.begin(async tx => { + const [order] = await tx`INSERT INTO orders ${sql(data)} RETURNING id`; + await tx.notify("orders", String(order.id)); // delivered when the transaction commits +}); +``` + +The payload is optional: `sql.notify("cache-invalidated")` is PostgreSQL's bare `NOTIFY`. Reserved and transaction handles also have `listen()`; it always uses the client's shared listen connection. + +`LISTEN`/`NOTIFY` is PostgreSQL-only; on MySQL and SQLite these methods reject. + +--- + ## Prepared Statements By default, Bun's SQL client creates named prepared statements for queries it can infer are static, which is faster. To disable this, set `prepare: false` in the connection options: @@ -1303,8 +1361,6 @@ We haven't implemented `LOAD DATA INFILE` support yet. We haven't implemented these yet: - `COPY` support -- `LISTEN` support -- `NOTIFY` support We also haven't implemented some of the more uncommon features like: diff --git a/packages/bun-types/sql.d.ts b/packages/bun-types/sql.d.ts index f0761a239f7e..f8e6f8310cbe 100644 --- a/packages/bun-types/sql.d.ts +++ b/packages/bun-types/sql.d.ts @@ -433,6 +433,17 @@ declare module "bun" { */ type Options = SQLiteOptions | PostgresOrMySQLOptions; + /** One registration made by {@link SQL.listen}. */ + interface ListenSubscription extends AsyncDisposable { + readonly channel: string; + /** + * Remove this registration. Resolves once the channel is no longer + * subscribed, or immediately when other registrations on it remain. + * Idempotent; `await using` calls it at the end of the scope. + */ + unlisten(): Promise; + } + /** * A pending SQL query. Extends `Promise`, so it can be awaited, and adds * methods to control how it runs. @@ -959,6 +970,57 @@ declare module "bun" { * ``` */ file(filename: string, values?: any[] | Record): SQL.Query; + + /** + * Subscribe to a PostgreSQL `LISTEN` channel. Resolves once the server has + * acknowledged the subscription, with a handle that removes it again. + * + * Every call is its own registration: several on one channel share a + * single server-side subscription and each receives every notification. + * All of them share one dedicated connection, opened by the first + * `listen()` and closed when the last registration is removed. If it + * drops, it is reconnected with exponential backoff and every channel is + * re-subscribed; `onlisten` runs again each time. + * + * A throwing `onnotify` or `onlisten` is reported as an uncaught exception. + * + * @param channel - Channel name, quoted for you; at most 63 bytes, the + * PostgreSQL identifier limit + * @param onnotify - Receives each notification's payload + * @param onlisten - Runs once the `LISTEN` is acknowledged, initially and + * after every reconnect + * + * @example + * ```ts + * const subscription = await sql.listen("events", payload => console.log(payload)); + * await sql.notify("events", "hello"); + * await subscription.unlisten(); + * ``` + * + * @example + * ```ts + * await using subscription = await sql.listen("events", handle); + * ``` + */ + listen( + channel: string, + onnotify: (payload: string) => void, + onlisten?: () => void, + ): Promise; + + /** + * Send a PostgreSQL `NOTIFY` via `pg_notify`. Runs as a normal query on + * this handle, so on a `sql.begin()` transaction it is delivered on commit + * and discarded on rollback. Omitting `payload` sends an empty one, like a + * bare `NOTIFY channel`. + * + * @example + * ```ts + * await sql.notify("events", JSON.stringify({ id: 1 })); + * await sql.notify("cache-invalidated"); + * ``` + */ + notify(channel: string, payload?: string): Promise; } /** diff --git a/src/js/bun/sql.ts b/src/js/bun/sql.ts index c3a74582a259..5f663e876a36 100644 --- a/src/js/bun/sql.ts +++ b/src/js/bun/sql.ts @@ -237,6 +237,43 @@ const SQL: typeof Bun.SQL = function SQL( } } + const listenable = "listen" in pool ? pool : null; + function validateChannel(channel: unknown): asserts channel is string { + if (typeof channel !== "string" || channel.length === 0) { + throw $ERR_INVALID_ARG_VALUE("channel", channel, "must be a non-empty string"); + } + if (channel.includes("\0")) { + throw $ERR_INVALID_ARG_VALUE("channel", channel, "must not contain null bytes"); + } + // PostgreSQL truncates longer identifiers, so notifications would arrive under a different name. + if (Buffer.byteLength(channel) > 63) { + throw $ERR_INVALID_ARG_VALUE("channel", channel, "must be at most 63 bytes"); + } + } + function validateCallback(name: string, fn: unknown, optional: boolean) { + if (!$isCallable(fn) && !(optional && fn === undefined)) { + throw $ERR_INVALID_ARG_TYPE(name, "function", fn); + } + } + const listenUnsupported = () => + Promise.$reject(new Error("LISTEN/NOTIFY is not supported by this adapter (PostgreSQL only)")); + const listen: Bun.SQL["listen"] = async (channel, onnotify, onlisten) => { + validateChannel(channel); + validateCallback("onnotify", onnotify, false); + validateCallback("onlisten", onlisten, true); + return listenable ? listenable.listen(channel, onnotify, onlisten) : listenUnsupported(); + }; + // .execute(): queries are lazy, and notify() must send even when not awaited. + function makeNotify(target: { unsafe: Bun.SQL["unsafe"] }): Bun.SQL["notify"] { + return (channel, payload) => { + validateChannel(channel); + if (payload === undefined) payload = ""; + else if (typeof payload !== "string") throw $ERR_INVALID_ARG_TYPE("payload", "string", payload); + if (!listenable) return listenUnsupported(); + return target.unsafe("SELECT pg_notify($1, $2)", [channel, payload]).execute() as unknown as Promise; + }; + } + function onReserveConnected(this: Query, err: Error | null, pooledConnection) { const { resolve, reject } = this; @@ -317,6 +354,8 @@ const SQL: typeof Bun.SQL = function SQL( // this matchs the behavior of the postgres package reserved_sql.reserve = () => sql.reserve(); reserved_sql.array = sql.array; + reserved_sql.listen = listen; + reserved_sql.notify = makeNotify(reserved_sql); function onTransactionFinished(transaction_promise: Promise) { reservedTransaction.delete(transaction_promise); } @@ -592,6 +631,8 @@ const SQL: typeof Bun.SQL = function SQL( // this matchs the behavior of the postgres package transaction_sql.reserve = () => sql.reserve(); transaction_sql.array = sql.array; + transaction_sql.listen = listen; + transaction_sql.notify = makeNotify(transaction_sql); transaction_sql.connect = () => { if (state.connectionState & ReservedConnectionState.closed) { @@ -938,6 +979,8 @@ const SQL: typeof Bun.SQL = function SQL( sql.transaction = sql.begin; sql.distributed = sql.beginDistributed; sql.end = sql.close; + sql.listen = listen; + sql.notify = makeNotify(sql); return sql; }; @@ -1017,6 +1060,14 @@ defaultSQLObject.flush = (...args: Parameters) => { ensureDefaultSQL(); return lazyDefaultSQL.flush(...args); }; +defaultSQLObject.listen = (...args: Parameters) => { + ensureDefaultSQL(); + return lazyDefaultSQL.listen(...args); +}; +defaultSQLObject.notify = (...args: Parameters) => { + ensureDefaultSQL(); + return lazyDefaultSQL.notify(...args); +}; //define lazy properties defineProperties(defaultSQLObject, { options: { diff --git a/src/js/internal/sql/postgres.ts b/src/js/internal/sql/postgres.ts index 897f6526c82b..c8907501063a 100644 --- a/src/js/internal/sql/postgres.ts +++ b/src/js/internal/sql/postgres.ts @@ -523,6 +523,342 @@ class PostgresAdapter } return pushBindParam(this, value, binding_values, index); } + + #listener: ListenConnection | null = null; + + listen(channel: string, onnotify: Listener, onlisten: OnListen | undefined): Promise { + if (this.closed) return Promise.$reject(this.connectionClosedError()); + return (this.#listener ??= new ListenConnection(this)).listen(channel, onnotify, onlisten); + } + + protected closeDedicatedConnections() { + this.#listener?.close(); + this.#listener = null; + } +} + +type Listener = (payload: string) => void; +type OnListen = () => void; +type ListenHandle = $ZigGeneratedClasses.PostgresSQLConnection; + +/** Resolved from `sql.listen()`; removes the registration that call made. */ +class ListenSubscription { + readonly channel: string; + #connection: ListenConnection | null; + readonly #onnotify: Listener; + readonly #onlisten: OnListen | undefined; + + constructor(connection: ListenConnection, channel: string, onnotify: Listener, onlisten: OnListen | undefined) { + this.channel = channel; + this.#connection = connection; + this.#onnotify = onnotify; + this.#onlisten = onlisten; + } + + unlisten(): Promise { + const connection = this.#connection; + if (connection === null) return Promise.$resolve(undefined); + this.#connection = null; + return connection.unsubscribe(this.channel, this.#onnotify, this.#onlisten); + } + + [Symbol.asyncDispose](): Promise { + return this.unlisten(); + } +} + +// One registration per listen() call (a callback may appear twice); the entry is deleted with its last one, so +// code resuming after an await compares its entry against the map to learn whether it is still the live one. +class Channel { + // A lone listener is stored bare; arrays are replaced, never mutated, so a dispatch in progress is unaffected. + listeners: Listener | readonly Listener[]; + onlisten: readonly OnListen[] | null = null; + /** LISTEN round trip on the current connection; reset to null on disconnect. */ + ready: Promise | null = null; + + constructor(listener: Listener) { + this.listeners = listener; + } + + add(listener: Listener) { + const current = this.listeners; + this.listeners = typeof current === "function" ? [current, listener] : [...current, listener]; + } + + /** @returns true when the channel has no registrations left */ + remove(listener: Listener): boolean { + const current = this.listeners; + if (typeof current === "function") return current === listener; + const index = current.indexOf(listener); + if (index !== -1) this.listeners = current.length === 2 ? current[1 - index] : current.toSpliced(index, 1); + return false; + } + + addOnlisten(callback: OnListen) { + this.onlisten = this.onlisten === null ? [callback] : [...this.onlisten, callback]; + } + + removeOnlisten(callback: OnListen) { + const current = this.onlisten; + if (current === null) return; + const index = current.indexOf(callback); + if (index !== -1) this.onlisten = current.length === 1 ? null : current.toSpliced(index, 1); + } + + fireOnlisten() { + const callbacks = this.onlisten; + if (callbacks === null) return; + for (let i = 0; i < callbacks.length; i++) invoke(callbacks[i]); + } +} + +// A throwing callback is reported as uncaught; it must not skip the callbacks +// after it, reject listen(), or look like a failed LISTEN to #sweep. +function invoke(callback: (arg?: T) => void, arg?: T) { + try { + callback(arg); + } catch (err) { + reportError(err); + } +} + +// The members onResolvePostgresQuery/onRejectPostgresQuery read off a query. +class ListenQuery { + resolve!: () => void; + reject!: (err: unknown) => void; + [_results] = null; + [_handle]: $ZigGeneratedClasses.PostgresSQLQuery; + + constructor(handle: $ZigGeneratedClasses.PostgresSQLQuery) { + this[_handle] = handle; + } + + static run(conn: ListenHandle, sql: string): Promise { + const handle = createPostgresQuery(sql, [], new SQLResultArray(), undefined, false, true); + const query = new ListenQuery(handle); + const { promise, resolve, reject } = Promise.withResolvers(); + query.resolve = resolve; + query.reject = err => reject(wrapPostgresError(err as Error)); + handle.run(conn, query as any); + return promise; + } +} + +function quoteChannel(channel: string) { + return '"' + channel.replaceAll('"', '""') + '"'; +} + +const RECONNECT_MIN_MS = 250; +const RECONNECT_MAX_MS = 32_000; + +// Every failure (connect error, drop, rejected LISTEN) is repaired by +// #scheduleSweep(): after a backoff, #sweep() re-issues LISTEN for each +// channel whose `ready` is null. +class ListenConnection { + readonly #adapter: PostgresAdapter; + readonly #channels = new Map(); + + #conn: ListenHandle | null = null; + #connecting: Promise | null = null; + /** Mid-handshake handle, so close() can abort it. */ + #handshake: ListenHandle | null = null; + #sweepTimer: ReturnType | null = null; + #backoffMs = RECONNECT_MIN_MS; + + constructor(adapter: PostgresAdapter) { + this.#adapter = adapter; + } + + readonly #onNotification = (channel: string, payload: string) => { + const entry = this.#channels.get(channel); + if (entry === undefined) return; + const listeners = entry.listeners; + if (typeof listeners === "function") { + listeners(payload); // nothing to shield from a throw; native reports it the same way + return; + } + for (let i = 0; i < listeners.length; i++) invoke(listeners[i], payload); + }; + + async listen(channel: string, onnotify: Listener, onlisten: OnListen | undefined): Promise { + let entry = this.#channels.get(channel); + if (entry === undefined) { + entry = new Channel(onnotify); + this.#channels.set(channel, entry); + } else { + entry.add(onnotify); + } + + try { + await (entry.ready ??= this.#subscribe(channel, entry)); + } catch (err) { + if (this.#channels.get(channel) === entry) { + if (entry.remove(onnotify)) { + this.#channels.delete(channel); + this.#closeIfIdle(); + } else { + this.#scheduleSweep(); // the channel's other listeners lost this round trip too + } + } + throw err; + } + + // Our registration pins the entry, so only close() can have removed it meanwhile. + if (this.#channels.get(channel) !== entry) throw this.#adapter.connectionClosedError(); + if (onlisten !== undefined) { + entry.addOnlisten(onlisten); + invoke(onlisten); + } + return new ListenSubscription(this, channel, onnotify, onlisten); + } + + async unsubscribe(channel: string, onnotify: Listener, onlisten: OnListen | undefined): Promise { + const entry = this.#channels.get(channel); + if (entry === undefined) return; + if (onlisten !== undefined) entry.removeOnlisten(onlisten); + if (!entry.remove(onnotify)) return; + + this.#channels.delete(channel); + if (this.#channels.size === 0) { + this.#closeIfIdle(); + return; + } + const conn = this.#conn; + if (conn === null) return; + try { + await ListenQuery.run(conn, "UNLISTEN " + quoteChannel(channel)); + } catch { + // The connection dropped, which unsubscribed everything anyway. + } + } + + close() { + this.#clearSweep(); + this.#channels.clear(); + const conn = this.#conn; + const handshake = this.#handshake; + this.#conn = this.#handshake = this.#connecting = null; + conn?.close(); + handshake?.close(); + } + + async #subscribe(channel: string, entry: Channel): Promise { + try { + const conn = await this.#connection(); + if (this.#channels.get(channel) !== entry) { + this.#closeIfIdle(); + return; + } + await ListenQuery.run(conn, "LISTEN " + quoteChannel(channel)); + if (this.#channels.get(channel) !== entry) return; + entry.fireOnlisten(); + } catch (err) { + entry.ready = null; + throw err; + } + } + + #connection(): Promise { + if (this.#conn !== null) return Promise.$resolve(this.#conn); + return (this.#connecting ??= this.#connect().finally(() => { + this.#connecting = null; + })); + } + + #connect(): Promise { + const adapter = this.#adapter; + const { promise, resolve, reject } = Promise.withResolvers(); + let live: ListenHandle | null = null; + + createPooledConnectionHandle( + createPostgresConnection, + { ...adapter.connectionInfo, idleTimeout: 0, maxLifetime: 0 }, + (err, conn) => { + this.#handshake = null; + if (err) return reject(wrapPostgresError(err)); + if (adapter.closed) { + conn.close(); + return reject(adapter.connectionClosedError()); + } + live = this.#conn = conn; + this.#backoffMs = RECONNECT_MIN_MS; + conn.onnotification = this.#onNotification; + conn.ref(); + resolve(conn); + this.#clearSweep(); + this.#sweep(); + }, + err => { + if (live === null) { + this.#handshake = null; + return reject(wrapPostgresError(err ?? adapter.connectionClosedError())); + } + if (this.#conn !== live) return; + this.#conn = null; + for (const entry of this.#channels.values()) entry.ready = null; + this.#scheduleSweep(); + }, + ).then(handle => { + if (handle === null || live !== null) return; + if (adapter.closed) handle.close(); + else this.#handshake = handle; + }); + + return promise; + } + + #sweep() { + if (this.#adapter.closed) return; + let failed = false; + let pending = 1; // released after the loop, so a sweep with nothing to do also settles + const settle = () => { + if (--pending !== 0) return; + if (failed) this.#scheduleSweep(); + else this.#backoffMs = RECONNECT_MIN_MS; + }; + for (const [channel, entry] of this.#channels) { + if (entry.ready !== null) continue; + pending++; + (entry.ready = this.#subscribe(channel, entry)).then(settle, err => { + failed = true; + if (this.#channels.get(channel) === entry) { + console.warn(`bun:sql LISTEN ${quoteChannel(channel)} failed, retrying: ${(err as Error)?.message ?? err}`); + } + settle(); + }); + } + settle(); + } + + #scheduleSweep() { + if (this.#sweepTimer !== null || this.#adapter.closed || this.#channels.size === 0) return; + // Ref'd on purpose: while disconnected, this timer keeps the process alive. + const delay = this.#backoffMs * (0.75 + Math.random() * 0.5); + this.#backoffMs = Math.min(this.#backoffMs * 2, RECONNECT_MAX_MS); + this.#sweepTimer = setTimeout(() => { + this.#sweepTimer = null; + this.#sweep(); + }, delay); + } + + #clearSweep() { + if (this.#sweepTimer !== null) { + clearTimeout(this.#sweepTimer); + this.#sweepTimer = null; + } + } + + // An in-flight handshake is left alone: the #subscribe awaiting it lands here. + #closeIfIdle() { + if (this.#channels.size !== 0) return; + this.#clearSweep(); + this.#backoffMs = RECONNECT_MIN_MS; + const conn = this.#conn; + if (conn !== null) { + this.#conn = null; + conn.close(); + } + } } export default { diff --git a/src/js/internal/sql/shared.ts b/src/js/internal/sql/shared.ts index 7dc59a8b0290..46877975e10e 100644 --- a/src/js/internal/sql/shared.ts +++ b/src/js/internal/sql/shared.ts @@ -1234,19 +1234,27 @@ abstract class BaseSQLAdapter { if (this.closed) { return; } let timeout = options?.timeout; - if (timeout) { + const hasTimeout = !!timeout; + if (hasTimeout) { timeout = Number(timeout); if (timeout > 2 ** 31 || timeout < 0 || timeout !== timeout) { throw $ERR_INVALID_ARG_VALUE("options.timeout", timeout, "must be a non-negative integer less than 2^31"); } + } + + this.closed = true; + this.closeDedicatedConnections(); - this.closed = true; + if (hasTimeout) { if (timeout === 0 || !this.hasPendingQueries()) { // close immediately await this.#close(); @@ -1268,7 +1276,6 @@ abstract class BaseSQLAdapter( @@ -12,15 +15,12 @@ impl NotificationResponse { if remaining < 4 { return Err(AnyPostgresError::InvalidMessage); } - // pid reader.int4()?; remaining -= 4; - // channel - let (_, consumed) = reader.string_within(remaining)?; + let (channel, consumed) = reader.string_within(remaining)?; remaining -= consumed; - // payload - reader.string_within(remaining)?; + let (payload, _) = reader.string_within(remaining)?; - Ok(Self {}) + Ok(Self { channel, payload }) } } diff --git a/src/sql/shared/ConnectionFlags.rs b/src/sql/shared/ConnectionFlags.rs index b06ab88756ae..ec85975fffe3 100644 --- a/src/sql/shared/ConnectionFlags.rs +++ b/src/sql/shared/ConnectionFlags.rs @@ -8,6 +8,8 @@ bitflags! { const USE_UNNAMED_PREPARED_STATEMENTS = 1 << 2; const WAITING_TO_PREPARE = 1 << 3; const HAS_BACKPRESSURE = 1 << 4; + /// `ref()` was called; `on_data` must not unref the idle connection. + const KEEP_ALIVE_REQUESTED = 1 << 5; } } diff --git a/src/sql_jsc/jsc.rs b/src/sql_jsc/jsc.rs index 1991c743fa4c..3341a30f2775 100644 --- a/src/sql_jsc/jsc.rs +++ b/src/sql_jsc/jsc.rs @@ -29,7 +29,7 @@ use core::ptr::NonNull; pub use bun_jsc::{ ArrayBuffer, CallFrame, CoerceTo, ErrorBuilder, ErrorCode, ExternColumnIdentifier, ExternColumnIdentifierValue, GlobalRef, JSArrayIterator, JSCell, JSGlobalObject, JSObject, - JSType, JSValue, JsCell, JsError, JsRef, JsResult, MarkedArgumentBuffer, StringJsc, + JSType, JSValue, JsCell, JsError, JsRef, JsResult, MarkedArgumentBuffer, StringJsc, Strong, StrongOptional, ThrowFmtArgs, ZigStringJsc, bun_string_jsc, host_fn, }; @@ -619,7 +619,7 @@ pub use bun_jsc::JsClass; pub mod codegen { ::bun_jsc::js_class_module!(JSPostgresSQLConnection = "PostgresSQLConnection" - as crate::postgres::PostgresSQLConnection { queries, onconnect, onclose }); + as crate::postgres::PostgresSQLConnection { queries, onconnect, onclose, onnotification }); ::bun_jsc::js_class_module!( JSPostgresSQLQuery = "PostgresSQLQuery" as crate::postgres::PostgresSQLQuery, impl_js_class { diff --git a/src/sql_jsc/postgres/PostgresSQLConnection.rs b/src/sql_jsc/postgres/PostgresSQLConnection.rs index 8aa1246304ed..a558acb40d89 100644 --- a/src/sql_jsc/postgres/PostgresSQLConnection.rs +++ b/src/sql_jsc/postgres/PostgresSQLConnection.rs @@ -7,8 +7,8 @@ use core::sync::atomic::{AtomicU32, Ordering}; use crate::jsc::EventLoopTimer; use crate::jsc::webcore::AutoFlusher; use crate::jsc::{ - self as jsc, CallFrame, HasAutoFlush, JSGlobalObject, JSValue, JsResult, VirtualMachine, - VirtualMachineSqlExt as _, + self as jsc, CallFrame, HasAutoFlush, JSGlobalObject, JSValue, JsResult, Strong, + VirtualMachine, VirtualMachineSqlExt as _, bun_string_jsc, }; use bun_boringssl as BoringSSL; use bun_collections::{OffsetByteList, StringHashMap, StringMap}; @@ -174,6 +174,14 @@ pub struct PostgresSQLConnection { // Private — see `timer`; recovered via [`Self::from_max_lifetime_timer_ptr`]. max_lifetime_timer: JsCell, pub(crate) auto_flusher: JsCell, + + /// Interned notification channel names; see `channel_name_js`. + channel_names: JsCell>, +} + +struct InternedChannel { + bytes: Box<[u8]>, + js: Strong, } bun_event_loop::impl_timer_owner!(PostgresSQLConnection; @@ -408,6 +416,7 @@ impl PostgresSQLConnection { lazy_array(get_queries => queries_get_cached, queries_set_cached), (get_on_connect, set_on_connect => onconnect_get_cached, onconnect_set_cached), (get_on_close, set_on_close => onclose_get_cached, onclose_set_cached), + (get_on_notification, set_on_notification => onnotification_get_cached, onnotification_set_cached), } pub(crate) fn setup_tls(&self) { @@ -1026,6 +1035,10 @@ impl PostgresSQLConnection { event_loop.exit(); // === defer block === if self.status.get() == Status::Connected + && !self + .flags + .get() + .contains(ConnectionFlags::KEEP_ALIVE_REQUESTED) && !self.has_query_running() && self.write_buffer.get().remaining().is_empty() { @@ -1210,6 +1223,7 @@ pub(crate) fn call(global_object: &JSGlobalObject, callframe: &CallFrame) -> JsR EventLoopTimerTag::PostgresSQLConnectionMaxLifetime, )), auto_flusher: JsCell::new(AutoFlusher::default()), + channel_names: JsCell::new(Vec::new()), })); // `heap::into_raw` is `Box::into_raw` — never null. Sole owner until @@ -1343,11 +1357,19 @@ impl SocketHandler { } impl PostgresSQLConnection { - bun_jsc::poll_ref_hostfns!( - field = poll_ref, - ctx = vm_ctx, - after = |this: &Self| this.update_has_pending_activity(), - ); + pub fn do_ref(this: &Self, _: &JSGlobalObject, _: &CallFrame) -> JsResult { + this.update_flags(|f| f.insert(ConnectionFlags::KEEP_ALIVE_REQUESTED)); + this.poll_ref.with_mut(|p| p.ref_(this.vm_ctx())); + this.update_has_pending_activity(); + Ok(JSValue::UNDEFINED) + } + + pub fn do_unref(this: &Self, _: &JSGlobalObject, _: &CallFrame) -> JsResult { + this.update_flags(|f| f.remove(ConnectionFlags::KEEP_ALIVE_REQUESTED)); + this.poll_ref.with_mut(|p| p.unref(this.vm_ctx())); + this.update_has_pending_activity(); + Ok(JSValue::UNDEFINED) + } pub fn do_flush(this: &Self, _: &JSGlobalObject, _: &CallFrame) -> JsResult { this.register_auto_flusher(); @@ -3026,8 +3048,8 @@ impl PostgresSQLConnection { // _resp dropped at scope end } MessageType::NotificationResponse => { - debug!("UNSUPPORTED NotificationResponse"); - let _resp = protocol::NotificationResponse::decode_internal(reader.reborrow())?; + let resp = protocol::NotificationResponse::decode_internal(reader.reborrow())?; + self.on_notification(resp.channel.slice(), resp.payload.slice()); } MessageType::EmptyQueryResponse => { reader.eat_message(&protocol::EMPTY_QUERY_RESPONSE)?; @@ -3076,6 +3098,58 @@ impl PostgresSQLConnection { JSValue::from(this.status.get() == Status::Connected) } + const MAX_INTERNED_CHANNELS: usize = 256; + + fn channel_name_js(&self, global: &JSGlobalObject, channel: &[u8]) -> Option { + if let Some(entry) = self + .channel_names + .get() + .iter() + .find(|entry| entry.bytes.as_ref() == channel) + { + return Some(entry.js.get()); + } + let js = match bun_string_jsc::create_utf8_for_js(global, channel) { + Ok(js) => js, + Err(e) => { + global.report_active_exception_as_unhandled(e); + return None; + } + }; + if self.channel_names.get().len() < Self::MAX_INTERNED_CHANNELS { + self.channel_names.with_mut(|names| { + names.push(InternedChannel { + bytes: channel.into(), + js: Strong::create(js, global), + }) + }); + } + Some(js) + } + + fn on_notification(&self, channel: &[u8], payload: &[u8]) { + let Some(this_value) = self.js_value.get().try_get() else { + return; + }; + let Some(callback) = js::onnotification_get_cached(this_value) else { + return; + }; + let global = self.global(); + let Some(channel_js) = self.channel_name_js(global, channel) else { + return; + }; + let payload_js = match bun_string_jsc::create_utf8_for_js(global, payload) { + Ok(js) => js, + Err(e) => return global.report_active_exception_as_unhandled(e), + }; + self.event_loop().run_callback( + callback, + global, + JSValue::UNDEFINED, + &[channel_js, payload_js], + ); + } + pub(crate) fn consume_on_connect_callback( &self, global_object: &JSGlobalObject, diff --git a/test/js/sql/postgres-listen-notify.test.ts b/test/js/sql/postgres-listen-notify.test.ts new file mode 100644 index 000000000000..731ace6d4957 --- /dev/null +++ b/test/js/sql/postgres-listen-notify.test.ts @@ -0,0 +1,1238 @@ +// sql.listen() / subscription.unlisten() / sql.notify() (PostgreSQL LISTEN/NOTIFY). +// +// Wire-level behavior (which LISTEN/UNLISTEN statements reach the server, how +// NotificationResponse frames are routed, reconnect) is tested against a +// scripted backend: it runs without docker and can produce interleavings a +// real server cannot be made to produce on demand. Behavior only observable +// at process level (exit, uncaught exceptions, RSS) runs in a subprocess. +// End-to-end behavior runs against the docker-compose postgres service. + +import { SQL } from "bun"; +import { describe, expect, test } from "bun:test"; +import { bunEnv, bunExe, describeWithContainer, isDockerEnabled, tempDir } from "harness"; +import net from "node:net"; +import path from "node:path"; +import { + listeningServer, + neverAnsweringServer, + pgAuthenticationOk, + pgCommandComplete, + pgErrorResponse, + pgInt32, + pgNotificationResponse, + pgRaw, + pgReadFrontendMessages, + pgReadyForQuery, +} from "./wire-frames"; + +const PID = 4242; +const SECRET = 99; +const pgBackendKeyData = (pid: number, secret: number) => pgRaw("K", Buffer.concat([pgInt32(pid), pgInt32(secret)])); +const pgError = (message: string) => + Buffer.concat([pgErrorResponse({ S: "ERROR", C: "42601", M: message }), pgReadyForQuery()]); + +// Simple queries ('Q', which LISTEN/UNLISTEN use) are recorded per connection +// and acked, unless their channel is held (ack parked until release()) or +// armed to fail once. Extended-protocol queries (which notify() uses) are +// recorded from their Parse and rejected so they settle. +async function mockServer() { + const connections: string[][] = []; + const sockets = new Set(); + const queryWaiters: Array<() => void> = []; + const closeWaiters: Array<() => void> = []; + const heldAcks: Array<() => void> = []; + let held: string[] = []; + const failOnce = new Set(); + let closed = 0; + + const { port, server } = await listeningServer(socket => { + const queries: string[] = []; + connections.push(queries); + sockets.add(socket); + let buffered = Buffer.alloc(0); + const record = (query: string) => { + queries.push(query); + for (const wake of queryWaiters.splice(0)) wake(); + }; + socket.once("data", () => { + socket.write(Buffer.concat([pgAuthenticationOk(), pgBackendKeyData(PID, SECRET), pgReadyForQuery()])); + socket.on("data", data => { + buffered = pgReadFrontendMessages(Buffer.concat([buffered, data]), (type, body) => { + if (type === 0x50 /* Parse: name\0 query\0 ... */) { + const nameEnd = body.indexOf(0); + record(body.toString("utf8", nameEnd + 1, body.indexOf(0, nameEnd + 1))); + socket.write(pgError("mock rejects extended-protocol queries")); + return; + } + if (type !== 0x51 /* Query: query\0 */) return; + const query = body.toString("utf8", 0, body.length - 1); + record(query); + const space = query.indexOf(" "); + const verb = query.slice(0, space); + const channel = query.slice(space + 2, -1).replaceAll('""', '"'); + if (verb === "LISTEN" && failOnce.delete(channel)) { + socket.write(pgError(`cannot LISTEN ${channel}`)); + return; + } + const ack = Buffer.concat([pgCommandComplete(verb), pgReadyForQuery()]); + if (held.includes(channel)) heldAcks.push(() => socket.write(ack)); + else socket.write(ack); + }); + }); + }); + socket.on("close", () => { + sockets.delete(socket); + closed++; + for (const wake of closeWaiters.splice(0)) wake(); + }); + socket.on("error", () => {}); + }); + + const all = () => connections.flat(); + const waitUntil = (waiters: Array<() => void>, done: () => boolean) => + done() + ? Promise.resolve() + : new Promise(resolve => { + const check = () => (done() ? resolve() : waiters.push(check)); + waiters.push(check); + }); + + return { + url: `postgres://u@127.0.0.1:${port}/db`, + /** Every query received so far, across connections, in arrival order. */ + get queries() { + return all(); + }, + /** Queries received on each connection, in accept order. */ + connections, + get liveConnections() { + return sockets.size; + }, + untilQuery: (pred: (queries: string[]) => boolean) => waitUntil(queryWaiters, () => pred(all())), + untilClosed: (count: number) => waitUntil(closeWaiters, () => closed >= count), + notify(channel: string, payload: string) { + this.notifyMany([[channel, payload]]); + }, + /** Several notifications in one segment, so they arrive in one read. */ + notifyMany(frames: Array<[channel: string, payload: string]>) { + const blob = Buffer.concat(frames.map(([channel, payload]) => pgNotificationResponse(PID, channel, payload))); + for (const socket of sockets) socket.write(blob); + }, + hold: (...channels: string[]) => void (held = channels), + release() { + held = []; + for (const ack of heldAcks.splice(0)) ack(); + }, + failNextListen: (channel: string) => void failOnce.add(channel), + dropConnections() { + for (const socket of sockets) socket.destroy(); + }, + [Symbol.asyncDispose]: () => new Promise(resolve => server.close(() => resolve())), + }; +} + +const client = (url: string) => new SQL(url, { max: 1, connectionTimeout: 5, idleTimeout: 5 }); + +/** A promise plus ways to settle it: `open()` now, or `after(n)` on the n-th call of the returned function. */ +function gate() { + const { promise, resolve } = Promise.withResolvers(); + let calls = 0; + return Object.assign(promise, { + open: () => resolve(), + after: (count: number) => () => void (++calls === count && resolve()), + }); +} + +describe("listen", () => { + test("routes notifications to the channel's listener, in order", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const got: string[] = []; + const third = gate(); + const count = third.after(3); + + const subscription = await sql.listen("orders", payload => { + got.push(payload); + count(); + }); + expect(server.queries).toEqual(['LISTEN "orders"']); + expect(subscription.channel).toBe("orders"); + + server.notifyMany([ + ["orders", "1"], + ["other", "not subscribed"], + ["orders", ""], + ["orders", "3"], + ]); + await third; + expect(got).toEqual(["1", "", "3"]); + }); + + test("thousands of notifications across channels arriving in one read", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const total = 3000; + const counts = { a: 0, b: 0, c: 0 }; + const done = gate(); + const count = done.after(total); + const listener = (channel: keyof typeof counts) => () => { + counts[channel]++; + count(); + }; + await Promise.all([sql.listen("a", listener("a")), sql.listen("b", listener("b")), sql.listen("c", listener("c"))]); + + const channels = ["a", "b", "c"] as const; + server.notifyMany(Array.from({ length: total }, (_, i) => [channels[i % 3], String(i)])); + await done; + expect(counts).toEqual({ a: 1000, b: 1000, c: 1000 }); + }); + + test("channel names and payloads are UTF-8", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const got = Promise.withResolvers(); + await sql.listen("канал", got.resolve); + expect(server.queries).toEqual(['LISTEN "канал"']); + server.notify("канал", "héllo → 世界 🚀"); + expect(await got.promise).toBe("héllo → 世界 🚀"); + }); + + test("channel names are quoted as identifiers", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const channel = 'weird "name"; DROP TABLE x'; + const got = Promise.withResolvers(); + await sql.listen(channel, got.resolve); + expect(server.queries).toEqual(['LISTEN "weird ""name""; DROP TABLE x"']); + server.notify(channel, "ok"); + expect(await got.promise).toBe("ok"); + }); + + test("channel names are limited to PostgreSQL's 63 identifier bytes, counted in UTF-8", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const ascii63 = Buffer.alloc(63, "c").toString(); + const cjk63 = Buffer.alloc(63, "字").toString(); // 21 characters + const cjk66 = Buffer.alloc(66, "字").toString(); // 22 characters, 66 bytes + await expect(sql.listen(ascii63 + "c", () => {})).rejects.toThrow(/63 bytes/); + await expect(sql.listen(cjk66, () => {})).rejects.toThrow(/63 bytes/); + expect(() => sql.notify(cjk66)).toThrow(/63 bytes/); + + await sql.listen(ascii63, () => {}); + await sql.listen(cjk63, () => {}); + expect(server.queries).toEqual([`LISTEN "${ascii63}"`, `LISTEN "${cjk63}"`]); + }); + + test("several listeners on one channel share one LISTEN and each receives", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const got: Record = { a: [], b: [], c: [] }; + const all = gate(); + const count = all.after(3); + for (const name of Object.keys(got)) { + await sql.listen("ch", payload => { + got[name].push(payload); + count(); + }); + } + expect(server.queries).toEqual(['LISTEN "ch"']); + server.notify("ch", "x"); + await all; + expect(got).toEqual({ a: ["x"], b: ["x"], c: ["x"] }); + }); + + test("concurrent listen() calls on a new channel share its round trip", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const all = gate(); + const count = all.after(3); + const subscriptions = await Promise.all([1, 2, 3].map(() => sql.listen("ch", () => count()))); + expect(server.queries).toEqual(['LISTEN "ch"']); + expect(subscriptions.map(subscription => subscription.channel)).toEqual(["ch", "ch", "ch"]); + expect(new Set(subscriptions).size).toBe(3); + server.notify("ch", "x"); + await all; + }); + + test("a listener that unlistens itself mid-dispatch does not starve the others", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const calls: string[] = []; + let count = () => {}; + const oneShot = await sql.listen("ch", payload => { + calls.push("oneShot:" + payload); + void oneShot.unlisten(); + count(); + }); + for (const name of ["b", "c"]) { + await sql.listen("ch", payload => { + calls.push(`${name}:${payload}`); + count(); + }); + } + + const first = gate(); + count = first.after(3); + server.notify("ch", "1"); + await first; + expect(calls).toEqual(["oneShot:1", "b:1", "c:1"]); + + const second = gate(); + count = second.after(2); + calls.length = 0; + server.notify("ch", "2"); + await second; + expect(calls).toEqual(["b:2", "c:2"]); + }); + + test("a listener added from inside a callback receives the next notification, not the current one", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const calls: string[] = []; + const done = gate(); + const count = done.after(5); // a:1 b:1 a:2 b:2 late:2 + await sql.listen("ch", payload => { + calls.push("a:" + payload); + if (payload === "1") { + // Registers synchronously (the channel is already subscribed); the + // dispatch in progress must not pick it up. + void sql.listen("ch", p => { + calls.push("late:" + p); + count(); + }); + } + count(); + }); + await sql.listen("ch", payload => { + calls.push("b:" + payload); + count(); + }); + + server.notifyMany([ + ["ch", "1"], + ["ch", "2"], + ]); + await done; + expect(calls).toEqual(["a:1", "b:1", "a:2", "b:2", "late:2"]); + }); + + test("registering the same callback twice is two subscriptions, each removed by its own handle", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + let calls = 0; + let count = () => {}; + const callback = () => { + calls++; + count(); + }; + const first = await sql.listen("ch", callback); + const second = await sql.listen("ch", callback); + expect(first).not.toBe(second); + + const both = gate(); + count = both.after(2); + server.notify("ch", "x"); + await both; + expect(calls).toBe(2); + + await first.unlisten(); + const one = gate(); + count = one.open; + server.notify("ch", "y"); + await one; + // This round trip is ordered after any second delivery of "y". + await sql.listen("barrier", () => {}); + expect(calls).toBe(3); + expect(server.queries).toEqual(['LISTEN "ch"', 'LISTEN "barrier"']); + + await second.unlisten(); + expect(server.queries).toEqual(['LISTEN "ch"', 'LISTEN "barrier"', 'UNLISTEN "ch"']); + }); + + test("onlisten runs after the LISTEN ack and before listen() resolves", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const events: string[] = []; + const subscription = await sql.listen( + "ch", + () => {}, + () => events.push(`onlisten after ${server.queries.length} queries`), + ); + events.push("resolved"); + expect(events).toEqual(["onlisten after 1 queries", "resolved"]); + expect(subscription.channel).toBe("ch"); + }); + + test("a later subscriber's onlisten runs once although no LISTEN is sent", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + await sql.listen("ch", () => {}); + let onlisten = 0; + await sql.listen( + "ch", + () => {}, + () => onlisten++, + ); + expect(onlisten).toBe(1); + expect(server.queries).toEqual(['LISTEN "ch"']); + }); +}); + +describe("unlisten", () => { + test("removing one of several listeners is local; removing the last sends UNLISTEN", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + await sql.listen("keep", () => {}); + const a = await sql.listen("ch", () => {}); + const b = await sql.listen("ch", () => {}); + + await a.unlisten(); + expect(server.queries).toEqual(['LISTEN "keep"', 'LISTEN "ch"']); + await b.unlisten(); + expect(server.queries).toEqual(['LISTEN "keep"', 'LISTEN "ch"', 'UNLISTEN "ch"']); + }); + + test("unlisten() and `await using` remove exactly that subscription; unlisten() is idempotent", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const kept: string[] = []; + const arrived = gate(); + await sql.listen("ch", payload => { + kept.push(payload); + arrived.open(); + }); + + const removed: string[] = []; + { + await using _subscription = await sql.listen("ch", payload => removed.push(payload)); + } + const subscription = await sql.listen("ch", payload => removed.push(payload)); + await subscription.unlisten(); + await subscription.unlisten(); + + server.notify("ch", "x"); + await arrived; + expect(kept).toEqual(["x"]); + expect(removed).toEqual([]); + expect(server.queries).toEqual(['LISTEN "ch"']); + }); + + test("delivery stops at unlisten(), before its round trip completes", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + await sql.listen("keep", () => {}); + const got: string[] = []; + const subscription = await sql.listen("ch", payload => got.push(payload)); + + server.hold("ch"); + const unlistening = subscription.unlisten(); + server.notify("ch", "late"); + server.release(); + // The UNLISTEN ack is written after the notification, so by the time + // this resolves the notification has been processed. + await unlistening; + expect(got).toEqual([]); + }); + + test("removing the last subscription closes the connection instead of sending UNLISTEN", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const subscription = await sql.listen("ch", () => {}); + expect(server.liveConnections).toBe(1); + await subscription.unlisten(); + await server.untilClosed(1); + expect(server.queries).toEqual(['LISTEN "ch"']); + + await sql.listen("ch", () => {}); + expect(server.connections).toEqual([['LISTEN "ch"'], ['LISTEN "ch"']]); + }); + + test("unlisten() resolves when the connection drops during its round trip", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + await sql.listen("keep", () => {}); + const subscription = await sql.listen("ch", () => {}); + server.hold("ch"); + const unlistening = subscription.unlisten(); + await server.untilQuery(queries => queries.includes('UNLISTEN "ch"')); + server.dropConnections(); + await unlistening; + }); + + test("a handle from before close() is a no-op afterwards, also for a re-opened client", async () => { + await using server = await mockServer(); + const sql = client(server.url); + const stale = await sql.listen("ch", () => {}); + await sql.close(); + await stale.unlisten(); + + await using reopened = client(server.url); + const got = Promise.withResolvers(); + await reopened.listen("ch", got.resolve); + await stale.unlisten(); + server.notify("ch", "still subscribed"); + expect(await got.promise).toBe("still subscribed"); + expect(server.queries).toEqual(['LISTEN "ch"', 'LISTEN "ch"']); + }); +}); + +describe("listen/unlisten interleavings", () => { + test("a re-listen() while the UNLISTEN is in flight gets its own LISTEN, ordered after it", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + await sql.listen("keep", () => {}); + const first = await sql.listen("ch", () => {}); + server.hold("ch"); + const unlistening = first.unlisten(); + const got = Promise.withResolvers(); + const relistening = sql.listen("ch", got.resolve); + server.release(); + await Promise.all([unlistening, relistening]); + expect(server.queries).toEqual(['LISTEN "keep"', 'LISTEN "ch"', 'UNLISTEN "ch"', 'LISTEN "ch"']); + server.notify("ch", "after"); + expect(await got.promise).toBe("after"); + }); + + test("a rejected LISTEN rejects that listen() and leaves nothing registered", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + await sql.listen("keep", () => {}); + server.failNextListen("bad"); + await expect(sql.listen("bad", () => {})).rejects.toThrow("cannot LISTEN bad"); + await sql.listen("bad", () => {}); + expect(server.queries).toEqual(['LISTEN "keep"', 'LISTEN "bad"', 'LISTEN "bad"']); + }); + + test("a rejected LISTEN shared by concurrent callers rejects all of them and releases the connection", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + server.hold("ch"); + server.failNextListen("ch"); + const a = sql.listen("ch", () => {}); + const b = sql.listen("ch", () => {}); + server.release(); + const results = await Promise.allSettled([a, b]); + expect(results.map(result => result.status)).toEqual(["rejected", "rejected"]); + await server.untilClosed(1); + }); + + test("a connection failure rejects listen()", async () => { + const { port, server } = await listeningServer(socket => socket.destroy()); + try { + await using sql = client(`postgres://u@127.0.0.1:${port}/db`); + await expect(sql.listen("ch", () => {})).rejects.toThrow(); + } finally { + server.close(); + } + }); +}); + +describe("reconnect", () => { + test("re-subscribes every channel, runs onlisten again, resumes delivery", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const bothAgain = gate(); + const onlisten = bothAgain.after(4); // two channels, twice each + const got = Promise.withResolvers(); + await sql.listen("a", got.resolve, onlisten); + await sql.listen("b", () => {}, onlisten); + + server.dropConnections(); + await bothAgain; + expect(server.connections).toHaveLength(2); + expect(server.connections[1].toSorted()).toEqual(['LISTEN "a"', 'LISTEN "b"']); + + server.notify("a", "after"); + expect(await got.promise).toBe("after"); + }); + + test("an unlistened subscription's onlisten does not run again on reconnect", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const keptAgain = gate(); + const removed = { onlisten: 0 }; + await sql.listen("ch", () => {}, keptAgain.after(2)); + const subscription = await sql.listen( + "ch", + () => {}, + () => removed.onlisten++, + ); + expect(removed.onlisten).toBe(1); + await subscription.unlisten(); + + server.dropConnections(); + await keptAgain; + expect(removed.onlisten).toBe(1); + }); + + test("a listen() during the backoff brings the connection up and re-subscribes the rest on it", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const oldAgain = gate(); + await sql.listen("old", () => {}, oldAgain.after(2)); + server.dropConnections(); + await server.untilClosed(1); + + // The adapter may observe the drop a moment after the server did; a + // listen() landing on the dying connection rejects, so retry. + let added = false; + while (!added) { + await sql + .listen("new", () => {}) + .then( + () => (added = true), + () => {}, + ); + } + await oldAgain; + const replacement = server.connections.at(-1)!; + expect(replacement.toSorted()).toEqual(['LISTEN "new"', 'LISTEN "old"']); + }); + + test("a channel whose re-LISTEN is rejected is retried with backoff and warns", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const warnings: string[] = []; + const originalWarn = console.warn; + console.warn = (...args: unknown[]) => void warnings.push(args.join(" ")); + try { + const resubscribed = gate(); + await sql.listen("flaky", () => {}, resubscribed.after(2)); + server.failNextListen("flaky"); + server.dropConnections(); + await resubscribed; + expect(server.connections.at(-1)!.filter(query => query === 'LISTEN "flaky"')).toHaveLength(2); + expect(warnings).toEqual([expect.stringContaining('LISTEN "flaky" failed, retrying: cannot LISTEN flaky')]); + } finally { + console.warn = originalWarn; + } + }); + + test("unlistening everything during the backoff cancels the reconnect", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const subscription = await sql.listen("ch", () => {}); + server.dropConnections(); + await server.untilClosed(1); + await subscription.unlisten(); + + // Span at least one backoff period deterministically: a second client + // goes through a full drop-and-reconnect against the same server. Had the + // first client's timer still been armed, it would have fired within it. + await using other = client(server.url); + const reconnected = gate(); + await other.listen("other", () => {}, reconnected.after(2)); + server.dropConnections(); + await reconnected; + + expect(server.queries.filter(query => query === 'LISTEN "ch"')).toHaveLength(1); + expect(server.liveConnections).toBe(1); + }); + + test("a listen() that reconnects but has its LISTEN rejected still gets the channel's other listener repaired", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const repaired = gate(); + const got = Promise.withResolvers(); + await sql.listen("ch", got.resolve, repaired.after(2)); + + // Plain awaits rather than expect().rejects: that one runs the event loop + // re-entrantly, and this code is still inside the dropped connection's + // data callback; on Windows the inner run frees the socket under it. + const outcome = (promise: Promise) => + promise.then( + () => "resolved", + (err: Error) => err.message, + ); + + // Issued in the same tick as the drop, this lands on the dying connection + // and rejects once the client has processed the drop, which is the moment + // the client considers "ch" unsubscribed and has armed its backoff. + server.dropConnections(); + expect(await outcome(sql.listen("probe", () => {}))).toMatch(/closed/i); + + // This listen() brings the connection back (cancelling the backoff) and is + // the one sending LISTEN "ch"; its rejection must not strand the first listener. + server.failNextListen("ch"); + expect(await outcome(sql.listen("ch", () => {}))).toBe("cannot LISTEN ch"); + + await repaired; + expect(server.connections).toEqual([['LISTEN "ch"'], ['LISTEN "ch"', 'LISTEN "ch"']]); + server.notify("ch", "after"); + expect(await got.promise).toBe("after"); + }); +}); + +describe("close()", () => { + test("tears down the listen connection and rejects an in-flight listen()", async () => { + await using server = await mockServer(); + const sql = client(server.url); + const subscription = await sql.listen("ch", () => {}); + server.hold("pending"); + const pending = sql.listen("pending", () => {}); + await server.untilQuery(queries => queries.includes('LISTEN "pending"')); + + await sql.close(); + await expect(pending).rejects.toThrow(); + await server.untilClosed(1); + await expect(sql.listen("ch", () => {})).rejects.toThrow("Connection closed"); + await subscription.unlisten(); + }); + + test("in the same tick as a listen() on the live connection rejects it before any LISTEN is sent", async () => { + await using server = await mockServer(); + const sql = client(server.url); + await sql.listen("keep", () => {}); + // Plain awaits: this continuation is inside the live connection's data callback (see the reconnect suite). + const outcome = sql + .listen("ch", () => {}) + .then( + () => "resolved", + (err: Error) => err.message, + ); + const closing = sql.close(); + expect(await outcome).toBe("Connection closed"); + await closing; + expect(server.queries).toEqual(['LISTEN "keep"']); + }); + + test("during the handshake aborts it and rejects the listen()", async () => { + const { port, server } = await neverAnsweringServer(); + try { + const sql = new SQL(`postgres://u@127.0.0.1:${port}/db`, { max: 1, connectionTimeout: 60 }); + const listening = sql.listen("ch", () => {}); + await sql.close(); + await expect(listening).rejects.toThrow(); + } finally { + server.close(); + } + }); + + test("an invalid timeout rejects before touching the subscriptions", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + const got = Promise.withResolvers(); + await sql.listen("ch", got.resolve); + await expect(sql.close({ timeout: -1 })).rejects.toThrow(); + server.notify("ch", "still subscribed"); + expect(await got.promise).toBe("still subscribed"); + }); +}); + +describe("arguments", () => { + test.each(["postgres", "sqlite"] as const)("%s: invalid arguments fail before any I/O", async adapter => { + await using sql = adapter === "postgres" ? client("postgres://u@127.0.0.1:1/db") : new SQL("sqlite://:memory:"); + const callback = () => {}; + await expect(sql.listen("", callback)).rejects.toThrow(/non-empty/); + await expect(sql.listen("a\0b", callback)).rejects.toThrow(/null bytes/); + await expect(sql.listen("ch", 1 as any)).rejects.toThrow(/onnotify/); + await expect(sql.listen("ch", callback, 1 as any)).rejects.toThrow(/onlisten/); + expect(() => sql.notify("", "p")).toThrow(/non-empty/); + expect(() => sql.notify("ch", null as any)).toThrow(/payload/); + expect(() => sql.notify("ch", 1 as any)).toThrow(/payload/); + }); + + test("non-Postgres adapters reject with a clear error", async () => { + await using sql = new SQL("sqlite://:memory:"); + await expect(sql.listen("ch", () => {})).rejects.toThrow("PostgreSQL only"); + await expect(sql.notify("ch", "p")).rejects.toThrow("PostgreSQL only"); + await expect(sql.notify("ch")).rejects.toThrow("PostgreSQL only"); + }); + + test("reserved connections expose listen() on the client's shared listen connection; unlisten lives on the subscription only", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + using reserved = await sql.reserve(); + expect("unlisten" in sql).toBe(false); + expect("unlisten" in reserved).toBe(false); + const viaReserved = await reserved.listen("ch", () => {}); + expect(typeof viaReserved.unlisten).toBe("function"); + const viaClient = await sql.listen("ch", () => {}); + expect(server.queries).toEqual(['LISTEN "ch"']); + await viaReserved.unlisten(); + await viaClient.unlisten(); + await server.untilClosed(1); + expect(server.queries).toEqual(['LISTEN "ch"']); + expect(typeof reserved.notify).toBe("function"); + }); +}); + +describe("notify()", () => { + test("is sent through the pool without being awaited, and opens no listen connection", async () => { + await using server = await mockServer(); + await using sql = client(server.url); + // Both calls share one prepared statement, so a single Parse arrives and + // the mock's rejection of it settles both. Payload values are checked end + // to end in the docker suite. + const settled = Promise.allSettled([sql.notify("ch", "payload"), sql.notify("signal")]); + await server.untilQuery(queries => queries.some(query => query.includes("pg_notify($1, $2)"))); + expect((await settled).map(result => result.status)).toEqual(["rejected", "rejected"]); + expect(server.liveConnections).toBe(1); + }); +}); + +describe("in a subprocess", () => { + const wireFrames = path.join(import.meta.dir, "wire-frames.ts"); + // The mock runs inside the child and is unref'd, so only the subscription + // under test can hold the child open. `notifyMany`, `dropConnections`, + // `sockets`, `pgNotificationResponse` and `sql` are in scope for `body`. + async function run(body: string, env: Record = {}) { + using dir = tempDir("pg-listen-subprocess", { + "fixture.ts": ` + import { SQL } from "bun"; + import { listeningServer, pgAuthenticationOk, pgReadyForQuery, pgCommandComplete, pgNotificationResponse, pgReadFrontendMessages, pgRaw, pgInt32 } from ${JSON.stringify(wireFrames)}; + const sockets = new Set(); + const { port, server } = await listeningServer(socket => { + sockets.add(socket); + socket.unref(); + let buffered = Buffer.alloc(0); + socket.once("data", () => { + socket.write(Buffer.concat([pgAuthenticationOk(), pgRaw("K", Buffer.concat([pgInt32(1), pgInt32(2)])), pgReadyForQuery()])); + socket.on("data", data => { + buffered = pgReadFrontendMessages(Buffer.concat([buffered, data]), (type, body) => { + if (type !== 0x51) return; + socket.write(Buffer.concat([pgCommandComplete(body.toString("utf8", 0, body.indexOf(" "))), pgReadyForQuery()])); + }); + }); + }); + socket.on("close", () => sockets.delete(socket)); + }); + server.unref(); + const notifyMany = frames => { + const blob = Buffer.concat(frames.map(([channel, payload]) => pgNotificationResponse(1, channel, payload))); + for (const socket of sockets) socket.write(blob); + }; + const dropConnections = () => { for (const socket of sockets) socket.destroy(); }; + const sql = new SQL("postgres://u@127.0.0.1:" + port + "/db", { max: 1 }); + ${body} + `, + }); + await using proc = Bun.spawn({ + cmd: [bunExe(), "fixture.ts"], + cwd: String(dir), + env: { ...bunEnv, ...env }, + stderr: "pipe", + timeout: 30_000, + killSignal: "SIGKILL", + }); + const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]); + return { + stdout, + stderr: stderr + .split("\n") + .filter(line => line && !line.startsWith("WARNING: ASAN")) + .join("\n"), + exitCode, + signal: proc.signalCode, + }; + } + const clean = (stdout: string) => ({ stdout, stderr: "", exitCode: 0, signal: null }); + + test("a subscription keeps the process alive until it is removed", async () => { + const result = await run(` + const subscription = await sql.listen("ch", async payload => { + console.log("got " + payload); + await subscription.unlisten(); + console.log("unlistened"); + }); + console.log("subscribed"); + // Sent from an unref'd timer after top-level code finishes, so the + // subscription is the only thing keeping the process here. + setTimeout(() => notifyMany([["ch", "wake"]]), 20).unref(); + `); + expect(result).toEqual(clean("subscribed\ngot wake\nunlistened\n")); + }); + + test("a dropped connection keeps the process alive through the backoff and reconnects", async () => { + const result = await run(` + let subscribes = 0; + const subscription = await sql.listen("ch", () => {}, () => { + if (++subscribes === 2) { console.log("reconnected"); subscription.unlisten(); } + }); + console.log("subscribed"); + dropConnections(); + `); + expect(result).toEqual(clean("subscribed\nreconnected\n")); + }); + + test("sql.close() releases a subscription so the process exits", async () => { + const result = await run(` + await sql.listen("ch", () => {}); + await sql.close(); + console.log("closed"); + `); + expect(result).toEqual(clean("closed\n")); + }); + + test("a throwing listener surfaces as uncaughtException; the channel's other listeners and later notifications are unaffected", async () => { + const result = await run(` + process.on("uncaughtException", err => console.log("uncaught: " + err.message)); + const first = await sql.listen("ch", payload => { + console.log("first got " + payload); + if (payload === "bad") throw new Error("listener failed"); + }); + const second = await sql.listen("ch", payload => { + console.log("second got " + payload); + if (payload === "good") { first.unlisten(); second.unlisten(); } + }); + notifyMany([["ch", "bad"], ["ch", "good"]]); + `); + expect(result).toEqual( + clean( + ["first got bad", "uncaught: listener failed", "second got bad", "first got good", "second got good", ""].join( + "\n", + ), + ), + ); + }); + + test("a lone throwing listener surfaces as uncaughtException and later notifications still arrive", async () => { + const result = await run(` + process.on("uncaughtException", err => console.log("uncaught: " + err.message)); + const subscription = await sql.listen("ch", payload => { + console.log("got " + payload); + if (payload === "bad") throw new Error("listener failed"); + subscription.unlisten(); + }); + notifyMany([["ch", "bad"], ["ch", "good"]]); + `); + expect(result).toEqual(clean("got bad\nuncaught: listener failed\ngot good\n")); + }); + + test("a throwing onlisten surfaces as uncaughtException; listen() resolves and the reconnect is not retried", async () => { + // An empty stderr is the assertion that the reconnect sweep did not take + // the second throw for a failed LISTEN (which it warns about and retries). + const result = await run(` + process.on("uncaughtException", err => console.log("uncaught: " + err.message)); + let calls = 0; + const subscription = await sql.listen("ch", () => {}, () => { + console.log("onlisten " + ++calls); + if (calls === 2) subscription.unlisten().then(() => console.log("unlistened")); + throw new Error("onlisten failed " + calls); + }); + console.log("subscribed to " + subscription.channel); + dropConnections(); + `); + expect(result).toEqual( + clean( + [ + "onlisten 1", + "uncaught: onlisten failed 1", + "subscribed to ch", + "onlisten 2", + "uncaught: onlisten failed 2", + "unlistened", + "", + ].join("\n"), + ), + ); + }); + + test("delivering many notifications retains nothing", async () => { + // Measured by RSS so a leaked native string backing (invisible to the JS + // heap) is caught. Each phase delivers its volume twice and reports the + // growth across the second pass: the first pass takes RSS to its steady + // state (the allocator keeps the GC's between-collection high-water mark + // of payload strings as free memory), after which a correct + // implementation shows roughly no growth and leaking one string per + // notification shows the pass's full payload volume. + const result = await run( + ` + const channels = ["a", "b", "c", "d"]; + let listener; + for (const channel of channels) await sql.listen(channel, payload => listener(payload)); + const segment = (count, payload) => [count, Buffer.concat(Array.from({ length: count }, (_, i) => pgNotificationResponse(1, channels[i % 4], payload)))]; + const deliver = ([count, blob]) => new Promise(done => { + let remaining = count; + listener = () => { if (--remaining === 0) done(); }; + for (const socket of sockets) socket.write(blob); + }); + const rss = () => { Bun.gc(true); return process.memoryUsage.rss(); }; + const measure = async (seg, rounds) => { + const pass = async () => { for (let i = 0; i < rounds; i++) await deliver(seg); }; + await pass(); + const base = rss(); + await pass(); + return Math.round((rss() - base) / 1024 / 1024); + }; + // 256 KiB per segment, 96 segments = 24 MiB per pass. + const small = await measure(segment(256, Buffer.alloc(1024, 0x61).toString()), 96); + const large = await measure(segment(4, Buffer.alloc(64 * 1024, 0x62).toString()), 96); + console.log(JSON.stringify({ small, large })); + await sql.close(); + `, + { ASAN_OPTIONS: "quarantine_size_mb=4:detect_leaks=0" }, + ); + expect(result.stderr).toBe(""); + expect(result.exitCode).toBe(0); + const growth = JSON.parse(result.stdout); + // Steady state drifts by a few MiB between passes; a leak adds ~24 MiB. + expect(growth).toEqual({ small: expect.any(Number), large: expect.any(Number) }); + expect(growth.small, JSON.stringify(growth)).toBeLessThan(6); + expect(growth.large, JSON.stringify(growth)).toBeLessThan(6); + }, 30_000); +}); + +if (isDockerEnabled()) { + describeWithContainer("postgres", { image: "postgres_plain" }, container => { + const connect = () => + new SQL(`postgres://bun_sql_test@${container.host}:${container.port}/bun_sql_test`, { max: 2 }); + // The listening backend is the session whose last statement is our LISTEN. + const terminateListeningBackend = async (sql: SQL, channel: string) => { + const terminated = await sql` + SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE query = ${`LISTEN "${channel}"`} + `; + expect(terminated).toHaveLength(1); + }; + + test("notify() round trips payloads, including empty and omitted", async () => { + await container.ready; + await using sql = connect(); + const got: string[] = []; + const three = gate(); + const count = three.after(3); + await using subscription = await sql.listen("e2e", payload => { + got.push(payload); + count(); + }); + expect(subscription.channel).toBe("e2e"); + await sql.notify("e2e", JSON.stringify({ n: 1 })); + await sql.notify("e2e", ""); + await sql.notify("e2e"); + await three; + expect(got).toEqual(['{"n":1}', "", ""]); + }); + + test("notify() in a transaction is delivered on commit and dropped on rollback", async () => { + await container.ready; + await using sql = connect(); + const got: string[] = []; + const barrier = gate(); + await using _subscription = await sql.listen("e2e_tx", payload => { + got.push(payload); + if (payload === "barrier") barrier.open(); + }); + + await expect( + sql.begin(async tx => { + await tx.notify("e2e_tx", "rolled back"); + throw new Error("abort"); + }), + ).rejects.toThrow("abort"); + await sql.begin(tx => tx.notify("e2e_tx", "committed")); + await sql.notify("e2e_tx", "barrier"); + await barrier; + expect(got).toEqual(["committed", "barrier"]); + }); + + // The LISTEN/NOTIFY tests of postgres.js (tests/index.js), adapted to the + // subscription object, with its delay() calls replaced by waiting for the + // deliveries themselves. Notifications to one connection arrive in order, + // so a later delivery proves an earlier notification was or was not + // delivered. + describe("ported from postgres.js", () => { + test("listen and notify", async () => { + await container.ready; + await using sql = connect(); + const result = Promise.withResolvers(); + await sql.listen("pgjs_hello", result.resolve); + await sql.notify("pgjs_hello", "works"); + expect(await result.promise).toBe("works"); + }); + + test("double listen", async () => { + await container.ready; + await using sql = connect(); + let count = 0; + for (let i = 0; i < 2; i++) { + const received = Promise.withResolvers(); + await sql.listen("pgjs_hello", received.resolve); + await sql.notify("pgjs_hello", "world"); + await received.promise; + count++; + } + await sql.listen("pgjs_weee", () => {}); + expect(count).toBe(2); + }); + + test("multiple listeners work after a reconnect", async () => { + await container.ready; + await using sql = connect(); + const xs: string[] = []; + let count = () => {}; + const resubscribed = gate(); + await sql.listen( + "pgjs_reconnect_multi", + x => { + xs.push("1" + x); + count(); + }, + resubscribed.after(2), + ); + await sql.listen("pgjs_reconnect_multi", x => { + xs.push("2" + x); + count(); + }); + + const a = gate(); + count = a.after(2); + await sql.notify("pgjs_reconnect_multi", "a"); + await a; + await terminateListeningBackend(sql, "pgjs_reconnect_multi"); + await resubscribed; + const b = gate(); + count = b.after(2); + await sql.notify("pgjs_reconnect_multi", "b"); + await b; + expect(xs.join("")).toBe("1a2a1b2b"); + }); + + test("listen and notify with weird name", async () => { + await container.ready; + await using sql = connect(); + const channel = "wat-;.ø.§"; + const got: string[] = []; + const first = gate(); + const subscription = await sql.listen(channel, payload => { + got.push(payload); + first.open(); + }); + await sql.notify(channel, "works"); + await first; + await subscription.unlisten(); + + const barrier = gate(); + await sql.listen("pgjs_barrier", barrier.open); + await sql.notify(channel, "after unlisten"); + await sql.notify("pgjs_barrier", ""); + await barrier; + expect(got).toEqual(["works"]); + }); + + test("listen and notify with upper case", async () => { + await container.ready; + await using sql = connect(); + const result = Promise.withResolvers(); + await sql.listen("withUpperChar", result.resolve); + await sql.notify("withUpperChar", "works"); + expect(await result.promise).toBe("works"); + }); + + test("listen reconnects", async () => { + await container.ready; + await using sql = connect(); + const a = gate(); + const b = gate(); + const resolvers: Record void> = { a: a.open, b: b.open }; + let connects = 0; + const reconnected = gate(); + await sql.listen( + "pgjs_reconnect", + x => resolvers[x]?.(), + () => { + if (++connects === 2) reconnected.open(); + }, + ); + await sql.notify("pgjs_reconnect", "a"); + await a; + await terminateListeningBackend(sql, "pgjs_reconnect"); + await reconnected; + await sql.notify("pgjs_reconnect", "b"); + await b; + expect(connects).toBe(2); + }); + + test("listen result reports correct connection state after reconnection", async () => { + await container.ready; + await using sql = connect(); + const listeningPids = async () => + (await sql`SELECT pid FROM pg_stat_activity WHERE query = ${'LISTEN "pgjs_state"'}`).map( + (row: { pid: number }) => row.pid, + ); + const resubscribed = gate(); + await sql.listen("pgjs_state", () => {}, resubscribed.after(2)); + const [initialPid] = await listeningPids(); + expect(initialPid).toBeNumber(); + + await terminateListeningBackend(sql, "pgjs_state"); + await resubscribed; + // The terminated backend may still be winding down, so only require a new one. + expect((await listeningPids()).some(pid => pid !== initialPid)).toBe(true); + }); + + test("unlisten removes subscription", async () => { + await container.ready; + await using sql = connect(); + const xs: string[] = []; + const a = gate(); + const subscription = await sql.listen("pgjs_test", x => { + xs.push(x); + a.open(); + }); + await sql.notify("pgjs_test", "a"); + await a; + await subscription.unlisten(); + + const barrier = gate(); + await sql.listen("pgjs_barrier", barrier.open); + await sql.notify("pgjs_test", "b"); + await sql.notify("pgjs_barrier", ""); + await barrier; + expect(xs.join("")).toBe("a"); + }); + + test("listen after unlisten", async () => { + await container.ready; + await using sql = connect(); + const xs: string[] = []; + let received = () => {}; + const listener = (x: string) => { + xs.push(x); + received(); + }; + + const a = gate(); + received = a.open; + const subscription = await sql.listen("pgjs_test", listener); + await sql.notify("pgjs_test", "a"); + await a; + await subscription.unlisten(); + await sql.notify("pgjs_test", "b"); + + const c = gate(); + received = c.open; + await sql.listen("pgjs_test", listener); + await sql.notify("pgjs_test", "c"); + await c; + expect(xs.join("")).toBe("ac"); + }); + + test("multiple listeners and unlisten one", async () => { + await container.ready; + await using sql = connect(); + const xs: string[] = []; + let count = () => {}; + await sql.listen("pgjs_test", x => { + xs.push("1" + x); + count(); + }); + const s2 = await sql.listen("pgjs_test", x => { + xs.push("2" + x); + count(); + }); + + const a = gate(); + count = a.after(2); + await sql.notify("pgjs_test", "a"); + await a; + await s2.unlisten(); + + const b = gate(); + count = b.open; + await sql.notify("pgjs_test", "b"); + await b; + expect(xs.join("")).toBe("1a2a1b"); + }); + }); + }); +}