diff --git a/src/runtime/valkey_jsc/ValkeyCommand.rs b/src/runtime/valkey_jsc/ValkeyCommand.rs index 54d7e94e39c3..86eaf6ff6c94 100644 --- a/src/runtime/valkey_jsc/ValkeyCommand.rs +++ b/src/runtime/valkey_jsc/ValkeyCommand.rs @@ -104,11 +104,21 @@ impl<'a> Command<'a> { } } +/// The channels and listener a `SUBSCRIBE` asked for, held until the server confirms the +/// subscription. The listener is only wired into the receive-handler map on confirmation, +/// so a SUBSCRIBE that fails leaves the map untouched. +pub struct PendingSubscription { + /// A channel name, or an array of them. + pub channels: jsc::Strong, + pub listener: jsc::Strong, +} + /// Command stored in offline queue when disconnected pub struct Entry { pub serialized_data: Box<[u8]>, // Pre-serialized RESP protocol bytes pub meta: Meta, pub promise: Promise, + pub pending_subscription: Option>, } // Inherent associated @@ -119,13 +129,18 @@ pub mod entry { impl Entry { // Create an Offline by serializing the Valkey command directly - pub fn create(command: &Command<'_>, promise: Promise) -> Result { + pub fn create( + command: &Command<'_>, + promise: Promise, + pending_subscription: Option>, + ) -> Result { Ok(Entry { serialized_data: command.serialize()?, // We should be calling .check against command here but due // to a hack introduced to let SUBSCRIBE work, we are not doing that for now. meta: command.meta, promise, + pending_subscription, }) } } @@ -228,6 +243,7 @@ impl Promise { pub struct PromisePair { pub meta: Meta, pub promise: Promise, + pub pending_subscription: Option>, } // See `entry` note above. diff --git a/src/runtime/valkey_jsc/js_valkey.rs b/src/runtime/valkey_jsc/js_valkey.rs index 93b3e215d00f..35aa2405e776 100644 --- a/src/runtime/valkey_jsc/js_valkey.rs +++ b/src/runtime/valkey_jsc/js_valkey.rs @@ -213,6 +213,28 @@ impl SubscriptionCtx { Ok(Some(new_length as usize)) } + /// Wire up the listener a `subscribe()` asked for, once the server has confirmed the + /// subscription. Registering up-front instead would leave the listener (and the event + /// loop ref it implies) behind when the SUBSCRIBE fails. + pub fn register_subscription( + &self, + global_object: &JSGlobalObject, + pending: &command::PendingSubscription, + ) -> JsResult<()> { + let channels = pending.channels.get(); + let listener = pending.listener.get(); + + if !channels.is_array() { + return self.upsert_receive_handler(global_object, channels, listener); + } + + let mut channels_iter = channels.array_iterator(global_object)?; + while let Some(channel) = channels_iter.next()? { + self.upsert_receive_handler(global_object, channel, listener)?; + } + Ok(()) + } + /// Add a handler for receiving messages on a specific channel pub fn upsert_receive_handler( &self, @@ -1691,6 +1713,7 @@ impl JSValkeyClient { global_this: &JSGlobalObject, _this_value: JSValue, command: &Command, + pending_subscription: Option>, ) -> Result<*mut JSPromise, crate::Error> { if self.client.get().flags.needs_to_open_socket { bun_core::hint::cold(); @@ -1713,7 +1736,8 @@ impl JSValkeyClient { let self_br = BackRef::new(self); let _update = scopeguard::guard(self_br, |p| p.update_poll_ref()); - self.client_mut().send(global_this, command) + self.client_mut() + .send(global_this, command, pending_subscription) } // Getter for memory cost - useful for diagnostics diff --git a/src/runtime/valkey_jsc/js_valkey_functions.rs b/src/runtime/valkey_jsc/js_valkey_functions.rs index d77eabb00ecb..77710b5b72b1 100644 --- a/src/runtime/valkey_jsc/js_valkey_functions.rs +++ b/src/runtime/valkey_jsc/js_valkey_functions.rs @@ -2,14 +2,16 @@ use crate::node::BlobOrStringOrBuffer as JSArgument; use bun_collections::VecExt as _; use bun_core::OwnedString; use bun_jsc::{ - self as jsc, CallFrame, ErrorCode, JSGlobalObject, JSPromise, JSPropertyIterator, JSValue, - JsRef, JsResult, + self as jsc, CallFrame, ErrorCode, JSArray, JSGlobalObject, JSPromise, JSPropertyIterator, + JSValue, JsRef, JsResult, }; use super::js_valkey::{JSValkeyClient, SubscriptionCtx}; use super::protocol_jsc as protocol; use super::valkey; -use super::valkey_command_body::{Args as CommandArgs, Command, Meta as CommandMeta}; +use super::valkey_command_body::{ + Args as CommandArgs, Command, Meta as CommandMeta, PendingSubscription, +}; type Slice = bun_jsc::ZigStringSlice; @@ -125,6 +127,7 @@ fn send_cmd( args, meta, }, + None, ) { Ok(p) => Ok(promise_to_js(p)), Err(err) => send_err_to_js(global, err_msg, &err), @@ -470,7 +473,7 @@ impl JSValkeyClient { let checked_meta = cmd.meta.check(&cmd); cmd.meta = checked_meta; // Send command with slices directly - let promise = match this.send(global, frame.this(), &cmd) { + let promise = match this.send(global, frame.this(), &cmd, None) { Ok(p) => p, Err(err) => { return send_err_to_js(global, "Failed to send command", &err); @@ -1651,6 +1654,11 @@ impl JSValkeyClient { return Err(global.throw_invalid_argument_type("subscribe", "listener", "function")); } + // The listener is registered only once the server confirms the subscription (see + // `SubscriptionCtx::register_subscription`), so a failed SUBSCRIBE leaves nothing + // behind, and the channel list has to be snapshotted out of the caller's reach. + let subscribed_channels: jsc::Strong; + // The first argument given is the channel or may be an array of channels. if channel_or_many.is_array() { if channel_or_many.get_length(global)? == 0 { @@ -1659,6 +1667,7 @@ impl JSValkeyClient { ))); } redis_channels.ensure_total_capacity(channel_or_many.get_length(global)? as usize); + let snapshot = jsc::Strong::create(JSArray::create_empty(global, 0)?, global); let mut array_iter = channel_or_many.array_iterator(global)?; while let Some(channel_arg) = array_iter.next()? { @@ -1670,31 +1679,16 @@ impl JSValkeyClient { )); }; redis_channels.push(channel); - - // What we do here is add our receive handler. Notice that this doesn't really do anything until the - // "SUBSCRIBE" command is sent to redis and we get a response. - // - // This is less-than-ideal, still, because this assumes a happy path. What happens if - // the SUBSCRIBE command fails? We have no way to roll back the addition of the - // handler. - this._subscription_ctx.get().upsert_receive_handler( - global, - channel_arg, - handler_callback, - )?; + snapshot.get().push(global, channel_arg)?; } + subscribed_channels = snapshot; } else if channel_or_many.is_string() { // It is a single string channel let Some(channel) = from_js(global, channel_or_many)? else { return Err(global.throw_invalid_argument_type("subscribe", "channel", "string")); }; redis_channels.push(channel); - - this._subscription_ctx.get().upsert_receive_handler( - global, - channel_or_many, - handler_callback, - )?; + subscribed_channels = jsc::Strong::create(channel_or_many, global); } else { return Err(global.throw_invalid_argument_type( "subscribe", @@ -1703,20 +1697,19 @@ impl JSValkeyClient { )); } + let pending_subscription = Box::new(PendingSubscription { + channels: subscribed_channels, + listener: jsc::Strong::create(handler_callback, global), + }); + let command = Command { command: b"SUBSCRIBE", args: CommandArgs::Args(&redis_channels), meta: CommandMeta::default() | CommandMeta::SUBSCRIPTION_REQUEST, }; - let promise = match this.send(global, frame.this(), &command) { + let promise = match this.send(global, frame.this(), &command, Some(pending_subscription)) { Ok(p) => p, - Err(err) => { - // If we catch an error, we need to clean up any handlers we may have added and fall out of subscription mode - this._subscription_ctx - .get() - .clear_all_receive_handlers(global)?; - return send_err_to_js(global, "Failed to send SUBSCRIBE command", &err); - } + Err(err) => return send_err_to_js(global, "Failed to send SUBSCRIBE command", &err), }; Ok(promise_to_js(promise)) diff --git a/src/runtime/valkey_jsc/mod.rs b/src/runtime/valkey_jsc/mod.rs index 20b86bef4547..d7732bbf51b7 100644 --- a/src/runtime/valkey_jsc/mod.rs +++ b/src/runtime/valkey_jsc/mod.rs @@ -57,7 +57,9 @@ pub use valkey_context::ValkeyContext; // the body's items so `command::PromisePair` / `command::Entry` resolve, and // alias it as `ValkeyCommand` for callers that use that spelling. pub mod valkey_command { - pub use super::valkey_command_body::{Entry, Meta, Promise, PromisePair, entry, promise_pair}; + pub use super::valkey_command_body::{ + Entry, Meta, PendingSubscription, Promise, PromisePair, entry, promise_pair, + }; // `index.rs` re-exports `super::valkey_command::ValkeyCommand`. pub use super::ValkeyCommand; } diff --git a/src/runtime/valkey_jsc/valkey.rs b/src/runtime/valkey_jsc/valkey.rs index 8af08821c23f..8da823d7df0c 100644 --- a/src/runtime/valkey_jsc/valkey.rs +++ b/src/runtime/valkey_jsc/valkey.rs @@ -448,6 +448,7 @@ impl ValkeyClient { .write_item(command::PromisePair { meta: cmd.meta, promise: cmd.promise, + pending_subscription: cmd.pending_subscription, }) .unwrap_or_oom(); self.write_buffer @@ -923,10 +924,6 @@ impl ValkeyClient { } RESPValue::Push(push) => { let p = self.parent(); - let sub_count = p - ._subscription_ctx - .get() - .channels_subscribed_to_count(&global_this)?; if let Some(msg_type) = protocol::SubscriptionPushMessage::from_bytes(&push.kind) { match msg_type { @@ -935,6 +932,21 @@ impl ValkeyClient { Ok(SubscribeHandled::Handled) } protocol::SubscriptionPushMessage::Subscribe => { + // Wire up the listener now that the server has confirmed the + // subscription, and count the channels only afterwards. + let pending = pair + .as_deref() + .and_then(|req_pair| req_pair.pending_subscription.as_ref()); + if let Some(pending) = pending { + p._subscription_ctx + .get() + .register_subscription(&global_this, pending)?; + } + let sub_count = p + ._subscription_ctx + .get() + .channels_subscribed_to_count(&global_this)?; + p.add_subscription(); self.on_valkey_subscribe(value); @@ -1327,6 +1339,7 @@ impl ValkeyClient { .write_item(command::PromisePair { meta: offline_cmd.meta, promise: offline_cmd.promise, + pending_subscription: offline_cmd.pending_subscription, }) .unwrap_or_oom(); let data = offline_cmd.serialized_data; @@ -1363,6 +1376,7 @@ impl ValkeyClient { &mut self, command: &Command, mut promise: command::Promise, + pending_subscription: Option>, ) -> Result<(), crate::Error> { let can_pipeline = command .meta @@ -1389,7 +1403,7 @@ impl ValkeyClient { || can_pipeline { // We serialize the bytes in here, so we don't need to worry about the lifetime of the Command itself. - let entry = command::Entry::create(command, promise)?; + let entry = command::Entry::create(command, promise, pending_subscription)?; self.queue.write_item(entry)?; // If we're connected and using auto pipelining, schedule a flush @@ -1414,6 +1428,7 @@ impl ValkeyClient { let cmd_pair = command::PromisePair { meta: command.meta, promise, + pending_subscription, }; // Add to queue with command type @@ -1427,6 +1442,7 @@ impl ValkeyClient { &mut self, global_this: &JSGlobalObject, command: &Command, + pending_subscription: Option>, ) -> Result<*mut JSPromise, crate::Error> { // FIX: Check meta before using it for routing decisions let mut checked_command = *command; @@ -1449,7 +1465,7 @@ impl ValkeyClient { // Handle disconnected state with offline queue match self.status { Status::Connected => { - self.enqueue(&checked_command, promise)?; + self.enqueue(&checked_command, promise, pending_subscription)?; // Schedule auto-flushing to process this command if pipelining is enabled if self.flags.enable_auto_pipelining @@ -1465,7 +1481,7 @@ impl ValkeyClient { Status::Connecting | Status::Disconnected => { // Only queue if offline queue is enabled if self.flags.enable_offline_queue { - self.enqueue(&checked_command, promise)?; + self.enqueue(&checked_command, promise, pending_subscription)?; } else { let _ = promise.reject( global_this, diff --git a/test/js/valkey/valkey-subscribe-listener.test.ts b/test/js/valkey/valkey-subscribe-listener.test.ts new file mode 100644 index 000000000000..540752c76a6a --- /dev/null +++ b/test/js/valkey/valkey-subscribe-listener.test.ts @@ -0,0 +1,207 @@ +import { RedisClient, type SocketHandler, type TCPSocketListener } from "bun"; +import { expect, test } from "bun:test"; +import { bunEnv, bunExe } from "harness"; + +// A scripted RESP3 server: it frames the RESP arrays the client sends and hands each +// complete command to `reply`, which writes back whatever the test wants. No real redis. + +const bulk = (s: string) => `$${Buffer.byteLength(s)}\r\n${s}\r\n`; +const push = (...parts: string[]) => `>${parts.length}\r\n${parts.join("")}`; +const pushSubscribe = (channel: string, count: number) => push(bulk("subscribe"), bulk(channel), `:${count}\r\n`); +const pushMessage = (channel: string, message: string) => push(bulk("message"), bulk(channel), bulk(message)); + +const HELLO_REPLY = `%3\r\n${bulk("server")}${bulk("redis")}${bulk("version")}${bulk("7.2.0")}${bulk("proto")}:3\r\n`; + +/** Consume one complete RESP array of bulk strings from the front of `buf`. */ +function parseCommand(buf: Buffer): { args: string[]; consumed: number } | null { + if (buf.length < 4 || buf[0] !== 0x2a /* '*' */) return null; + let eol = buf.indexOf("\r\n"); + if (eol < 0) return null; + const count = parseInt(buf.subarray(1, eol).toString("latin1"), 10); + let off = eol + 2; + const args: string[] = []; + for (let i = 0; i < count; i++) { + if (off >= buf.length || buf[off] !== 0x24 /* '$' */) return null; + eol = buf.indexOf("\r\n", off); + if (eol < 0) return null; + const len = parseInt(buf.subarray(off + 1, eol).toString("latin1"), 10); + const start = eol + 2; + if (start + len + 2 > buf.length) return null; + args.push(buf.subarray(start, start + len).toString("latin1")); + off = start + len + 2; + } + return { args, consumed: off }; +} + +type Reply = (command: string, args: string[], write: (resp: string) => void) => void; + +function scriptedRedis(reply: Reply): TCPSocketListener { + const buffers = new WeakMap[0], Buffer>(); + return Bun.listen({ + hostname: "127.0.0.1", + port: 0, + socket: { + open(socket) { + buffers.set(socket, Buffer.alloc(0)); + }, + data(socket, chunk) { + let buf = Buffer.concat([buffers.get(socket)!, chunk]); + let frame: ReturnType; + while ((frame = parseCommand(buf))) { + buf = buf.subarray(frame.consumed); + reply(frame.args[0].toUpperCase(), frame.args.slice(1), resp => socket.write(resp)); + } + buffers.set(socket, buf); + }, + error() {}, + close() {}, + }, + }); +} + +/** Replies to HELLO, confirms every SUBSCRIBE, and publishes one message per channel. */ +function confirmingRedis(): TCPSocketListener { + return scriptedRedis((command, args, write) => { + if (command === "HELLO") return write(HELLO_REPLY); + if (command !== "SUBSCRIBE") return; + for (const [i, channel] of args.entries()) write(pushSubscribe(channel, i + 1)); + for (const channel of args) write(pushMessage(channel, "hi")); + }); +} + +test.concurrent("a subscribe() rejected before it is sent does not register its listener", async () => { + using server = confirmingRedis(); + + // With the offline queue disabled the first subscribe is rejected while the connection + // is still being established, so nothing ever reaches the server. + const redis = new RedisClient(`redis://127.0.0.1:${server.port}`, { + enableOfflineQueue: false, + autoReconnect: false, + }); + const received: string[] = []; + const { promise: delivered, resolve: onDelivered } = Promise.withResolvers(); + const listener = (message: string) => { + received.push(message); + onDelivered(); + }; + + try { + let rejected: Error | undefined; + try { + await redis.subscribe("ch", listener); + } catch (error) { + rejected = error as Error; + } + expect(rejected?.message).toBe("Connection is closed and offline queue is disabled"); + + // The rejected subscribe must not have left its listener behind, or this one registers + // it a second time and the single published message arrives twice. + await redis.connect(); + expect(await redis.subscribe("ch", listener)).toBe(1); + await delivered; + expect(received).toEqual(["hi"]); + } finally { + redis.close(); + } +}); + +test.concurrent("a subscribe() abandoned by a closing connection does not register its listener", async () => { + const { promise: sawSubscribe, resolve: onSawSubscribe } = Promise.withResolvers(); + let subscribesSeen = 0; + using server = scriptedRedis((command, args, write) => { + if (command === "HELLO") return write(HELLO_REPLY); + if (command !== "SUBSCRIBE") return; + // Leave the first SUBSCRIBE on the wire, unanswered. + if (subscribesSeen++ === 0) return void onSawSubscribe(); + write(pushSubscribe(args[0], 1)); + write(pushMessage(args[0], "hi")); + }); + + const redis = new RedisClient(`redis://127.0.0.1:${server.port}`, { autoReconnect: false }); + const received: string[] = []; + const { promise: delivered, resolve: onDelivered } = Promise.withResolvers(); + const listener = (message: string) => { + received.push(message); + onDelivered(); + }; + + try { + const pending = redis.subscribe("ch", listener); + await sawSubscribe; + redis.close(); + + let rejected: Error | undefined; + try { + await pending; + } catch (error) { + rejected = error as Error; + } + expect(rejected?.message).toBe("Connection closed"); + + // Same as above: the SUBSCRIBE never got its confirmation, so it must not have left a + // listener behind for the next one to duplicate. + await redis.connect(); + expect(await redis.subscribe("ch", listener)).toBe(1); + await delivered; + expect(received).toEqual(["hi"]); + } finally { + redis.close(); + } +}); + +test.concurrent("a subscribe() rejected before it is sent does not keep the event loop alive", async () => { + using server = confirmingRedis(); + + // The client is deliberately never closed: it holds no subscription and no pending + // command once the subscribe is rejected, so the process must exit on its own. The + // unref'd timer only fires if something is still holding the loop open. + await using proc = Bun.spawn({ + cmd: [ + bunExe(), + "-e", + ` + const redis = new Bun.RedisClient("redis://127.0.0.1:${server.port}", { + enableOfflineQueue: false, + autoReconnect: false, + }); + setTimeout(() => { console.log("event loop pinned"); process.exit(7); }, 2500).unref(); + try { + await redis.subscribe("ch", () => {}); + console.log("unexpectedly resolved"); + } catch (e) { + console.log("rejected:", e.message); + } + `, + ], + env: bunEnv, + stdout: "pipe", + stderr: "pipe", + }); + + const [stdout, stderr, exitCode] = await Promise.all([proc.stdout.text(), proc.stderr.text(), proc.exited]); + expect({ stdout: stdout.trim(), stderr, exitCode, signalCode: proc.signalCode }).toEqual({ + stdout: "rejected: Connection is closed and offline queue is disabled", + stderr: "", + exitCode: 0, + signalCode: null, + }); +}); + +test.concurrent("mutating the channel array after subscribe() does not change what gets registered", async () => { + using server = confirmingRedis(); + + const redis = new RedisClient(`redis://127.0.0.1:${server.port}`, { autoReconnect: false }); + const channels = ["a"]; + const { promise: delivered, resolve: onDelivered } = Promise.withResolvers(); + + try { + const subscribed = redis.subscribe(channels, onDelivered); + channels[0] = "b"; + expect(await subscribed).toBe(1); + + // The listener belongs to "a", the channel that was actually sent to the server. + expect(await delivered).toBe("hi"); + } finally { + redis.close(); + } +});