redis: register the subscribe() listener only after the server confirms the subscription

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<Box<PendingSubscription>>,
 }
 
 // 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<Entry, crate::Error> {
+    pub fn create(
+        command: &Command<'_>,
+        promise: Promise,
+        pending_subscription: Option<Box<PendingSubscription>>,
+    ) -> Result<Entry, crate::Error> {
         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<Box<PendingSubscription>>,
 }
 
 // See `entry` note above.

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<Box<command::PendingSubscription>>,
     ) -> 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

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))

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;
 }

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<Box<command::PendingSubscription>>,
     ) -> 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<Box<command::PendingSubscription>>,
     ) -> 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,