diff --git a/docs/runtime/redis.mdx b/docs/runtime/redis.mdx index 6e81051bf893..df87ec3f17c1 100644 --- a/docs/runtime/redis.mdx +++ b/docs/runtime/redis.mdx @@ -226,8 +226,9 @@ bun run publisher.ts Subscribing takes over the `RedisClient` connection: a client with -subscriptions can only call `RedisClient.prototype.subscribe()`. To send other -commands to Redis, create a separate connection with `.duplicate()`: +subscriptions can only call `.ping()`, `.subscribe()`, `.unsubscribe()`, +`.psubscribe()` and `.punsubscribe()`. To send other commands to Redis, create +a separate connection with `.duplicate()`: ```ts redis.ts icon="/icons/typescript.svg" import { RedisClient } from "bun"; @@ -266,6 +267,25 @@ await client.unsubscribe(channel); // Unsubscribe a particular channel. await client.unsubscribe(channel, listener); // Unsubscribe a particular listener. ``` +### Pattern Subscriptions + +Subscribe to channels matching a glob-style pattern with `.psubscribe()`. The +listener receives the message and the concrete channel it was published to: + +```typescript redis.ts icon="/icons/typescript.svg" +await client.psubscribe("news.*", (message, channel) => { + console.log(`[${channel}] ${message}`); +}); +``` + +Unsubscribe from patterns with `.punsubscribe()`: + +```typescript redis.ts icon="/icons/typescript.svg" +await client.punsubscribe(); // Unsubscribe from all patterns. +await client.punsubscribe(pattern); // Unsubscribe a particular pattern. +await client.punsubscribe(pattern, listener); // Unsubscribe a particular listener. +``` + ## Advanced Usage ### Command Execution and Pipelining @@ -368,7 +388,7 @@ The following commands disable automatic pipelining: - `SUBSCRIBE` - `PSUBSCRIBE` - `UNSUBSCRIBE` -- `UNPSUBSCRIBE` +- `PUNSUBSCRIBE` --- diff --git a/packages/bun-types/redis.d.ts b/packages/bun-types/redis.d.ts index 9f496436e0ed..da247da97e8a 100644 --- a/packages/bun-types/redis.d.ts +++ b/packages/bun-types/redis.d.ts @@ -2695,8 +2695,10 @@ declare module "bun" { * * Subscribing moves the channel to a dedicated subscription state which * prevents most other commands from being executed until unsubscribed. Only - * {@link ping `.ping()`}, {@link subscribe `.subscribe()`}, and - * {@link unsubscribe `.unsubscribe()`} can be called while subscribed. + * {@link ping `.ping()`}, {@link subscribe `.subscribe()`}, + * {@link unsubscribe `.unsubscribe()`}, {@link psubscribe `.psubscribe()`} + * and {@link punsubscribe `.punsubscribe()`} can be called while + * subscribed. * * @param channel The channel to subscribe to. * @param listener The listener to call when a message is received on the @@ -2785,6 +2787,81 @@ declare module "bun" { */ unsubscribe(channels: string[]): Promise; + /** + * Subscribe to a Redis channel pattern. + * + * Patterns use glob-style matching: `?` matches one character, `*` matches + * zero or more characters, and `[...]` matches one character in the set. + * + * Subscribing disables automatic pipelining, so all commands are received + * immediately. + * + * Subscribing moves the client to a dedicated subscription state which + * prevents most other commands from being executed until unsubscribed. Only + * {@link ping `.ping()`}, {@link subscribe `.subscribe()`}, + * {@link unsubscribe `.unsubscribe()`}, {@link psubscribe `.psubscribe()`} + * and {@link punsubscribe `.punsubscribe()`} can be called while + * subscribed. + * + * @param pattern The glob pattern to subscribe to. + * @param listener The listener to call when a message is received on a + * channel matching the pattern. The listener receives the message as the + * first argument and the originating channel name as the second. + * + * @example + * ```ts + * await client.psubscribe("news.*", (message, channel) => { + * console.log(`Received message on ${channel}: ${message}`); + * }); + * ``` + */ + psubscribe(pattern: string, listener: RedisClient.StringPubSubListener): Promise; + + /** + * Subscribe to multiple Redis channel patterns. + * + * @param patterns An array of glob patterns to subscribe to. + * @param listener The listener to call when a message is received on a + * channel matching any of the patterns. + */ + psubscribe(patterns: string[], listener: RedisClient.StringPubSubListener): Promise; + + /** + * Unsubscribe from a single Redis channel pattern. + * + * If there are no more channels or patterns subscribed to, the client + * automatically re-enables pipelining if it was previously enabled and + * leaves the subscription state. + * + * @param pattern The pattern to unsubscribe from. + */ + punsubscribe(pattern: string): Promise; + + /** + * Remove a listener from a given Redis channel pattern. + * + * @param pattern The pattern to unsubscribe from. + * @param listener The listener to remove. This is tested against + * referential equality so you must pass the exact same listener instance as + * when subscribing. + */ + punsubscribe(pattern: string, listener: RedisClient.StringPubSubListener): Promise; + + /** + * Unsubscribe from all registered Redis channel patterns. + * + * Channel subscriptions registered via {@link subscribe `.subscribe()`} are + * unaffected. + */ + punsubscribe(): Promise; + + /** + * Unsubscribe from multiple Redis channel patterns. + * + * @param patterns An array of patterns to unsubscribe from. + */ + punsubscribe(patterns: string[]): Promise; + /** * Create a new RedisClient instance with the same configuration as the * current instance. diff --git a/src/runtime/valkey_jsc/ValkeyCommand.rs b/src/runtime/valkey_jsc/ValkeyCommand.rs index 54d7e94e39c3..fd37bcc6989e 100644 --- a/src/runtime/valkey_jsc/ValkeyCommand.rs +++ b/src/runtime/valkey_jsc/ValkeyCommand.rs @@ -167,7 +167,7 @@ bun_core::comptime_string_set! { b"SUBSCRIBE", b"PSUBSCRIBE", b"UNSUBSCRIBE", - b"UNPSUBSCRIBE", + b"PUNSUBSCRIBE", }; } diff --git a/src/runtime/valkey_jsc/js_valkey.rs b/src/runtime/valkey_jsc/js_valkey.rs index 28fbc111edc7..57663a26e108 100644 --- a/src/runtime/valkey_jsc/js_valkey.rs +++ b/src/runtime/valkey_jsc/js_valkey.rs @@ -71,6 +71,12 @@ pub struct SubscriptionCtx { pub original_enable_auto_pipelining: bool, } +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum SubscriptionKind { + Channel, + Pattern, +} + /// The generate-classes.ts output emits a /// `js_RedisClient` module with snake-case `*_set_cached`/`*_get_cached` /// free-fns plus `to_js`/`from_js`. Re-exported here as `Js`. @@ -82,7 +88,6 @@ bun_core::impl_field_parent! { SubscriptionCtx => JSValkeyClient._subscription_c impl SubscriptionCtx { pub fn init(valkey_parent: &JSValkeyClient) -> JsResult { - let callback_map = JSMap::create(&valkey_parent.global_object); let parent_this = valkey_parent .this_value .get() @@ -92,7 +97,12 @@ impl SubscriptionCtx { Js::subscription_callback_map_set_cached( parent_this, &valkey_parent.global_object, - callback_map, + JSMap::create(&valkey_parent.global_object), + ); + Js::pattern_callback_map_set_cached( + parent_this, + &valkey_parent.global_object, + JSMap::create(&valkey_parent.global_object), ); Ok(SubscriptionCtx { @@ -106,24 +116,32 @@ impl SubscriptionCtx { }) } - fn subscription_callback_map(&self) -> &mut JSMap { + fn callback_map(&self, kind: SubscriptionKind) -> &mut JSMap { let parent_this = self .parent() .this_value .get() .try_get() .expect("unreachable"); - let value_js = Js::subscription_callback_map_get_cached(parent_this).unwrap(); + let value_js = match kind { + SubscriptionKind::Channel => Js::subscription_callback_map_get_cached(parent_this), + SubscriptionKind::Pattern => Js::pattern_callback_map_get_cached(parent_this), + } + .unwrap(); // `JSMap` is an `opaque_ffi!` ZST — `opaque_mut` is the safe deref. // `from_js` returns a non-null heap cell when the slot was set by // `init()`; single JS thread. JSMap::opaque_mut(JSMap::from_js(value_js).unwrap().as_ptr()) } - /// Get the total number of channels that this subscription context is subscribed to. pub fn channels_subscribed_to_count(&self, global_object: &JSGlobalObject) -> JsResult { - let count = self.subscription_callback_map().size(global_object)?; - Ok(count) + let channels = self + .callback_map(SubscriptionKind::Channel) + .size(global_object)?; + let patterns = self + .callback_map(SubscriptionKind::Pattern) + .size(global_object)?; + Ok(channels + patterns) } /// Test whether this context has any subscriptions. It is mandatory to @@ -134,16 +152,21 @@ impl SubscriptionCtx { pub fn clear_receive_handlers( &self, + kind: SubscriptionKind, global_object: &JSGlobalObject, channel_name: JSValue, ) -> JsResult<()> { - let map = self.subscription_callback_map(); + let map = self.callback_map(kind); let _ = map.remove(global_object, channel_name)?; Ok(()) } - pub fn clear_all_receive_handlers(&self, global_object: &JSGlobalObject) -> JsResult<()> { - self.subscription_callback_map().clear(global_object) + pub fn clear_all_receive_handlers( + &self, + kind: SubscriptionKind, + global_object: &JSGlobalObject, + ) -> JsResult<()> { + self.callback_map(kind).clear(global_object) } /// Remove a specific receive handler. @@ -154,11 +177,12 @@ impl SubscriptionCtx { /// Note: This function will empty out the map entry if there are no more handlers registered. pub fn remove_receive_handler( &self, + kind: SubscriptionKind, global_object: &JSGlobalObject, channel_name: JSValue, callback: JSValue, ) -> JsResult> { - let map = self.subscription_callback_map(); + let map = self.callback_map(kind); let existing = map.get(global_object, channel_name)?; if existing.is_undefined_or_null() { @@ -198,6 +222,7 @@ impl SubscriptionCtx { /// Add a handler for receiving messages on a specific channel pub fn upsert_receive_handler( &self, + kind: SubscriptionKind, global_object: &JSGlobalObject, channel_name: JSValue, callback: JSValue, @@ -208,7 +233,7 @@ impl SubscriptionCtx { let _guard = scopeguard::guard(parent_br, |p| { p.on_new_subscription_callback_insert(); }); - let map = self.subscription_callback_map(); + let map = self.callback_map(kind); let handlers_array: JSValue; let mut is_new_channel = false; @@ -244,12 +269,11 @@ impl SubscriptionCtx { pub fn get_callbacks( &self, + kind: SubscriptionKind, global_object: &JSGlobalObject, channel_name: JSValue, ) -> JsResult> { - let result = self - .subscription_callback_map() - .get(global_object, channel_name)?; + let result = self.callback_map(kind).get(global_object, channel_name)?; if result == JSValue::UNDEFINED { return Ok(None); } @@ -260,11 +284,12 @@ impl SubscriptionCtx { /// Handles both single callbacks and arrays of callbacks pub fn invoke_callbacks( &self, + kind: SubscriptionKind, global_object: &JSGlobalObject, channel_name: JSValue, args: &[JSValue], ) -> JsResult<()> { - let Some(callbacks) = self.get_callbacks(global_object, channel_name)? else { + let Some(callbacks) = self.get_callbacks(kind, global_object, channel_name)? else { debug!( "No callbacks found for channel {}", // `JSString` is an `opaque_ffi!` ZST — `opaque_ref` is the safe @@ -321,6 +346,7 @@ impl SubscriptionCtx { global_object, JSValue::UNDEFINED, ); + Js::pattern_callback_map_set_cached(parent_this, global_object, JSValue::UNDEFINED); } } } @@ -1303,7 +1329,7 @@ impl JSValkeyClient { Ok(()) } - pub fn on_valkey_message(&self, value: &mut [protocol::RESPValue]) { + pub fn on_valkey_message(&self, kind: SubscriptionKind, value: &mut [protocol::RESPValue]) { if !self.is_subscriber() { debug!("onMessage called but client is not in subscriber mode"); return; @@ -1312,31 +1338,47 @@ impl JSValkeyClient { let global_object = self.global_object; let _exit = self.vm().enter_event_loop_scope(); - // The message push should be an array with [channel, message] - if value.len() < 2 { + let (channel_idx, message_idx) = match kind { + SubscriptionKind::Channel => (0, 1), + SubscriptionKind::Pattern => (1, 2), + }; + if value.len() <= message_idx { debug!("Message array has insufficient elements: {}", value.len()); return; } - // Extract channel and message - let Ok(channel_value) = protocol_jsc::resp_value_to_js(&mut value[0], &global_object) + let Ok(channel_value) = + protocol_jsc::resp_value_to_js(&mut value[channel_idx], &global_object) else { debug!("Failed to convert channel to JS"); return; }; - let Ok(message_value) = protocol_jsc::resp_value_to_js(&mut value[1], &global_object) + let Ok(message_value) = + protocol_jsc::resp_value_to_js(&mut value[message_idx], &global_object) else { debug!("Failed to convert message to JS"); return; }; + let key_value = if kind == SubscriptionKind::Channel { + channel_value + } else { + match protocol_jsc::resp_value_to_js(&mut value[0], &global_object) { + Ok(v) => v, + Err(_) => { + debug!("Failed to convert pattern to JS"); + return; + } + } + }; // Invoke callbacks for this channel with message and channel as arguments if self ._subscription_ctx .get() .invoke_callbacks( + kind, &global_object, - channel_value, + key_value, &[message_value, channel_value], ) .is_err() diff --git a/src/runtime/valkey_jsc/js_valkey_functions.rs b/src/runtime/valkey_jsc/js_valkey_functions.rs index b63d9eb066b4..f9642dc6b7d7 100644 --- a/src/runtime/valkey_jsc/js_valkey_functions.rs +++ b/src/runtime/valkey_jsc/js_valkey_functions.rs @@ -6,7 +6,7 @@ use bun_jsc::{ JsRef, JsResult, }; -use super::js_valkey::{JSValkeyClient, SubscriptionCtx}; +use super::js_valkey::{JSValkeyClient, SubscriptionCtx, SubscriptionKind}; use super::protocol_jsc as protocol; use super::valkey; use super::valkey_command_body::{Args as CommandArgs, Command, Meta as CommandMeta}; @@ -1594,8 +1594,6 @@ impl JSValkeyClient { NotSubscriber ); cmd_key_varargs!(zrevrank, b"zrevrank", "ZREVRANK", "key", NotSubscriber); - cmd_strings_varargs!(psubscribe, b"psubscribe", "PSUBSCRIBE", DontCare); - cmd_strings_varargs!(punsubscribe, b"punsubscribe", "PUNSUBSCRIBE", DontCare); cmd_strings_varargs!(pubsub, b"pubsub", "PUBSUB", DontCare); cmd_strings_varargs!(copy, b"copy", "COPY", NotSubscriber); cmd_key_varargs!(unlink, b"unlink", "UNLINK", "key", NotSubscriber); @@ -1642,8 +1640,16 @@ impl JSValkeyClient { ) } - #[bun_jsc::host_fn(method)] - pub fn subscribe(this: &Self, global: &JSGlobalObject, frame: &CallFrame) -> JsResult { + fn do_subscribe( + this: &Self, + global: &JSGlobalObject, + frame: &CallFrame, + kind: SubscriptionKind, + fn_name: &'static str, + arg_name: &'static str, + redis_command: &'static [u8], + err_msg: &'static str, + ) -> JsResult { // `upsert_receive_handler`'s exit guard re-enters `on_writable` / // `update_poll_ref` before `send()` is reached; hold a ref so `*this` // stays live across those calls. @@ -1651,16 +1657,17 @@ impl JSValkeyClient { let [channel_or_many, handler_callback] = frame.arguments_as_array::<2>(); let mut redis_channels: Vec = Vec::with_capacity(1); + let mut registered: Vec = Vec::with_capacity(1); if !handler_callback.is_callable() { - return Err(global.throw_invalid_argument_type("subscribe", "listener", "function")); + return Err(global.throw_invalid_argument_type(fn_name, "listener", "function")); } // 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 { return Err(global.throw_invalid_arguments(format_args!( - "subscribe requires at least one channel" + "{fn_name} requires at least one {arg_name}" ))); } redis_channels.ensure_total_capacity(channel_or_many.get_length(global)? as usize); @@ -1668,65 +1675,93 @@ impl JSValkeyClient { let mut array_iter = channel_or_many.array_iterator(global)?; while let Some(channel_arg) = array_iter.next()? { let Some(channel) = from_js(global, channel_arg)? else { - return Err(global.throw_invalid_argument_type( - "subscribe", - "channel", - "string", - )); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string")); }; 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( + kind, global, channel_arg, handler_callback, )?; + registered.push(channel_arg); } } 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")); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string")); }; redis_channels.push(channel); this._subscription_ctx.get().upsert_receive_handler( + kind, global, channel_or_many, handler_callback, )?; + registered.push(channel_or_many); } else { - return Err(global.throw_invalid_argument_type( - "subscribe", - "channel", - "string or array", - )); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string or array")); } let command = Command { - command: b"SUBSCRIBE", + command: redis_command, args: CommandArgs::Args(&redis_channels), meta: CommandMeta::default() | CommandMeta::SUBSCRIPTION_REQUEST, }; let promise = match this.send(global, frame.this(), &command) { 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); + for ch in ®istered { + let _ = this._subscription_ctx.get().remove_receive_handler( + kind, + global, + *ch, + handler_callback, + ); + } + return send_err_to_js(global, err_msg, &err); } }; Ok(promise_to_js(promise)) } + #[bun_jsc::host_fn(method)] + pub fn subscribe(this: &Self, global: &JSGlobalObject, frame: &CallFrame) -> JsResult { + Self::do_subscribe( + this, + global, + frame, + SubscriptionKind::Channel, + "subscribe", + "channel", + b"SUBSCRIBE", + "Failed to send SUBSCRIBE command", + ) + } + + #[bun_jsc::host_fn(method)] + pub fn psubscribe( + this: &Self, + global: &JSGlobalObject, + frame: &CallFrame, + ) -> JsResult { + Self::do_subscribe( + this, + global, + frame, + SubscriptionKind::Pattern, + "psubscribe", + "pattern", + b"PSUBSCRIBE", + "Failed to send PSUBSCRIBE command", + ) + } + /// Send redis the UNSUBSCRIBE RESP command and clean up anything necessary after the unsubscribe commoand. /// /// The subscription context must exist when calling this function. @@ -1734,31 +1769,37 @@ impl JSValkeyClient { this: &Self, this_js: JSValue, global: &JSGlobalObject, + redis_command: &'static [u8], + err_msg: &'static str, redis_channels: &[JSArgument], ) -> JsResult { send_cmd( this, global, this_js, - b"UNSUBSCRIBE", + redis_command, CommandArgs::Args(redis_channels), CommandMeta::default(), - "Failed to send UNSUBSCRIBE command", + err_msg, ) } - #[bun_jsc::host_fn(method)] - pub fn unsubscribe( + fn do_unsubscribe( this: &Self, global: &JSGlobalObject, frame: &CallFrame, + kind: SubscriptionKind, + fn_name: &'static str, + arg_name: &'static str, + redis_command: &'static [u8], + err_msg: &'static str, ) -> JsResult { // Hold a ref so `*this` stays live across the handler-map updates and // the `send()` below. let _guard = this.ref_scope(); // Check if we're in subscription mode - require_subscriber(this, b"unsubscribe")?; + require_subscriber(this, fn_name.as_bytes())?; let args_view = frame.arguments(); @@ -1768,11 +1809,13 @@ impl JSValkeyClient { if args_view.is_empty() { this._subscription_ctx .get() - .clear_all_receive_handlers(global)?; + .clear_all_receive_handlers(kind, global)?; return Self::send_unsubscribe_request_and_cleanup( this, frame.this(), global, + redis_command, + err_msg, &redis_channels, ); } @@ -1793,18 +1836,14 @@ impl JSValkeyClient { // In this case, the first argument is a channel string and the second // argument is the handler to remove. if !channel_or_many.is_string() { - return Err(global.throw_invalid_argument_type("unsubscribe", "channel", "string")); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string")); } let channel = channel_or_many; let listener_cb = frame.argument(1); if !listener_cb.is_callable() { - return Err(global.throw_invalid_argument_type( - "unsubscribe", - "listener", - "function", - )); + return Err(global.throw_invalid_argument_type(fn_name, "listener", "function")); } // Populate the redis_channels list with the single channel to @@ -1812,11 +1851,12 @@ impl JSValkeyClient { // the UNSUBSCRIBE command to redis. Without this, we would end up // unsubscribing from all channels. let Some(ch) = from_js(global, channel)? else { - return Err(global.throw_invalid_argument_type("unsubscribe", "channel", "string")); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string")); }; redis_channels.push(ch); let remaining_listeners = match this._subscription_ctx.get().remove_receive_handler( + kind, global, channel, listener_cb, @@ -1840,6 +1880,8 @@ impl JSValkeyClient { this, frame.this(), global, + redis_command, + err_msg, &redis_channels, ); } @@ -1854,7 +1896,7 @@ impl JSValkeyClient { if channel_or_many.is_array() { if channel_or_many.get_length(global)? == 0 { return Err(global.throw_invalid_arguments(format_args!( - "unsubscribe requires at least one channel" + "{fn_name} requires at least one {arg_name}" ))); } @@ -1865,38 +1907,73 @@ impl JSValkeyClient { let mut array_iter = channel_or_many.array_iterator(global)?; while let Some(channel_arg) = array_iter.next()? { let Some(channel) = from_js(global, channel_arg)? else { - return Err(global.throw_invalid_argument_type( - "unsubscribe", - "channel", - "string", - )); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string")); }; redis_channels.push(channel); // Clear the handlers for this channel this._subscription_ctx .get() - .clear_receive_handlers(global, channel_arg)?; + .clear_receive_handlers(kind, global, channel_arg)?; } } 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("unsubscribe", "channel", "string")); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string")); }; redis_channels.push(channel); // Clear the handlers for this channel this._subscription_ctx .get() - .clear_receive_handlers(global, channel_or_many)?; + .clear_receive_handlers(kind, global, channel_or_many)?; } else { - return Err(global.throw_invalid_argument_type( - "unsubscribe", - "channel", - "string or array", - )); + return Err(global.throw_invalid_argument_type(fn_name, arg_name, "string or array")); } // Now send the unsubscribe command and clean up if necessary - Self::send_unsubscribe_request_and_cleanup(this, frame.this(), global, &redis_channels) + Self::send_unsubscribe_request_and_cleanup( + this, + frame.this(), + global, + redis_command, + err_msg, + &redis_channels, + ) + } + + #[bun_jsc::host_fn(method)] + pub fn unsubscribe( + this: &Self, + global: &JSGlobalObject, + frame: &CallFrame, + ) -> JsResult { + Self::do_unsubscribe( + this, + global, + frame, + SubscriptionKind::Channel, + "unsubscribe", + "channel", + b"UNSUBSCRIBE", + "Failed to send UNSUBSCRIBE command", + ) + } + + #[bun_jsc::host_fn(method)] + pub fn punsubscribe( + this: &Self, + global: &JSGlobalObject, + frame: &CallFrame, + ) -> JsResult { + Self::do_unsubscribe( + this, + global, + frame, + SubscriptionKind::Pattern, + "punsubscribe", + "pattern", + b"PUNSUBSCRIBE", + "Failed to send PUNSUBSCRIBE command", + ) } #[bun_jsc::host_fn(method)] diff --git a/src/runtime/valkey_jsc/valkey.classes.ts b/src/runtime/valkey_jsc/valkey.classes.ts index 3e2cf711c81b..1f153185a780 100644 --- a/src/runtime/valkey_jsc/valkey.classes.ts +++ b/src/runtime/valkey_jsc/valkey.classes.ts @@ -572,9 +572,9 @@ export default [ hstrlen: { fn: "hstrlen" }, zrank: { fn: "zrank" }, zrevrank: { fn: "zrevrank" }, - subscribe: { fn: "subscribe" }, + subscribe: { fn: "subscribe", length: 2 }, duplicate: { fn: "duplicate" }, - psubscribe: { fn: "psubscribe" }, + psubscribe: { fn: "psubscribe", length: 2 }, unsubscribe: { fn: "unsubscribe" }, punsubscribe: { fn: "punsubscribe" }, pubsub: { fn: "pubsub" }, @@ -584,6 +584,6 @@ export default [ rename: { fn: "rename", length: 2 }, renamenx: { fn: "renamenx", length: 2 }, }, - values: ["onconnect", "onclose", "connectionPromise", "hello", "subscriptionCallbackMap"], + values: ["onconnect", "onclose", "connectionPromise", "hello", "subscriptionCallbackMap", "patternCallbackMap"], }), ]; diff --git a/src/runtime/valkey_jsc/valkey.rs b/src/runtime/valkey_jsc/valkey.rs index eed40edb6a8f..8cb90ed88d1c 100644 --- a/src/runtime/valkey_jsc/valkey.rs +++ b/src/runtime/valkey_jsc/valkey.rs @@ -10,7 +10,7 @@ use bun_uws::{self as uws, AnySocket, SocketGroup, SocketKind, SslCtx}; use bun_valkey::valkey_protocol as protocol; use bun_valkey::valkey_protocol::{RESPValue, RedisError}; -use super::js_valkey_body::JSValkeyClient; +use super::js_valkey_body::{JSValkeyClient, SubscriptionKind}; use super::protocol_jsc::{resp_value_to_js, valkey_error_to_js}; use super::valkey_command_body as command; use super::valkey_command_body::{Args, Command}; @@ -920,10 +920,15 @@ impl ValkeyClient { if let Some(msg_type) = protocol::SubscriptionPushMessage::from_bytes(&push.kind) { match msg_type { protocol::SubscriptionPushMessage::Message => { - self.on_valkey_message(&mut push.data); + self.on_valkey_message(SubscriptionKind::Channel, &mut push.data); Ok(SubscribeHandled::Handled) } - protocol::SubscriptionPushMessage::Subscribe => { + protocol::SubscriptionPushMessage::Pmessage => { + self.on_valkey_message(SubscriptionKind::Pattern, &mut push.data); + Ok(SubscribeHandled::Handled) + } + protocol::SubscriptionPushMessage::Subscribe + | protocol::SubscriptionPushMessage::Psubscribe => { p.add_subscription(); self.on_valkey_subscribe(value); @@ -937,7 +942,8 @@ impl ValkeyClient { } Ok(SubscribeHandled::Handled) } - protocol::SubscriptionPushMessage::Unsubscribe => { + protocol::SubscriptionPushMessage::Unsubscribe + | protocol::SubscriptionPushMessage::Punsubscribe => { self.on_valkey_unsubscribe()?; self.parent().remove_subscription(); @@ -1091,14 +1097,11 @@ impl ValkeyClient { // For subscription clients, check if this is a push message that doesn't need a promise pair if let RESPValue::Push(push) = value { match protocol::SubscriptionPushMessage::from_bytes(&push.kind) { - Some(protocol::SubscriptionPushMessage::Message) => { + Some(kind) if kind.is_message() => { // Message pushes never need promise pairs should_consume_promise_pair = false; } - Some( - protocol::SubscriptionPushMessage::Subscribe - | protocol::SubscriptionPushMessage::Unsubscribe, - ) => { + Some(_) => { // Subscribe/unsubscribe pushes only need promise pairs if we have pending commands if self.in_flight.readable_length() == 0 { should_consume_promise_pair = false; @@ -1527,8 +1530,8 @@ impl ValkeyClient { self.parent().on_valkey_unsubscribe() } - pub fn on_valkey_message(&mut self, value: &mut [RESPValue]) { - self.parent().on_valkey_message(value); + pub fn on_valkey_message(&mut self, kind: SubscriptionKind, value: &mut [RESPValue]) { + self.parent().on_valkey_message(kind, value); } pub fn on_valkey_reconnect(&mut self) { diff --git a/src/valkey/valkey_protocol.rs b/src/valkey/valkey_protocol.rs index 479322bca6c7..ee4f978b67da 100644 --- a/src/valkey/valkey_protocol.rs +++ b/src/valkey/valkey_protocol.rs @@ -787,6 +787,9 @@ pub enum SubscriptionPushMessage { Message, Subscribe, Unsubscribe, + Pmessage, + Psubscribe, + Punsubscribe, } bun_core::comptime_string_map! { @@ -794,6 +797,9 @@ bun_core::comptime_string_map! { b"message" => SubscriptionPushMessage::Message, b"subscribe" => SubscriptionPushMessage::Subscribe, b"unsubscribe" => SubscriptionPushMessage::Unsubscribe, + b"pmessage" => SubscriptionPushMessage::Pmessage, + b"psubscribe" => SubscriptionPushMessage::Psubscribe, + b"punsubscribe" => SubscriptionPushMessage::Punsubscribe, }; } @@ -803,17 +809,19 @@ impl SubscriptionPushMessage { SUBSCRIPTION_PUSH_MESSAGES.get(bytes).copied() } - /// Pattern (`p`-prefixed) and sharded (`s`-prefixed) variants of the - /// `Subscribe`/`Unsubscribe` push kinds; the unprefixed kinds are matched by - /// `from_bytes` before this is consulted. #[inline] pub fn is_reply_kind(kind: &[u8]) -> bool { match kind.split_first() { - Some((b'p' | b's', base)) => matches!( + Some((b's', base)) => matches!( Self::from_bytes(base), Some(Self::Subscribe | Self::Unsubscribe) ), _ => false, } } + + #[inline] + pub fn is_message(self) -> bool { + matches!(self, Self::Message | Self::Pmessage) + } } diff --git a/test/js/valkey/reliability/resp-nesting-depth.test.ts b/test/js/valkey/reliability/resp-nesting-depth.test.ts index 600f5005ffae..62c669ac1bde 100644 --- a/test/js/valkey/reliability/resp-nesting-depth.test.ts +++ b/test/js/valkey/reliability/resp-nesting-depth.test.ts @@ -345,10 +345,10 @@ describe("Valkey: RESP push frame routing", () => { }); try { - const psubscribed = client.psubscribe("news.*"); + const psubscribed = client.psubscribe("news.*", () => {}); const pinged = client.send("PING", []); - expect(await psubscribed).toEqual({ type: "psubscribe", data: ["news.*", 1] }); + expect(await psubscribed).toBe(1); expect(await pinged).toBe("PONG"); } finally { client.close(); diff --git a/test/js/valkey/test-utils.ts b/test/js/valkey/test-utils.ts index fba5ad8c1519..46c8a86697af 100644 --- a/test/js/valkey/test-utils.ts +++ b/test/js/valkey/test-utils.ts @@ -365,6 +365,9 @@ export const context: TestContext = { try { await client.unsubscribe(); } catch {} + try { + await client.punsubscribe(); + } catch {} if (client.connected) { client.close(); diff --git a/test/js/valkey/valkey.test.ts b/test/js/valkey/valkey.test.ts index 0e521c0fc6d3..5db8e119ea6b 100644 --- a/test/js/valkey/valkey.test.ts +++ b/test/js/valkey/valkey.test.ts @@ -6574,6 +6574,151 @@ for (const connectionType of [ConnectionType.TLS, ConnectionType.TCP]) { ); }); + test("psubscribe to a pattern receives messages", async () => { + const TEST_MESSAGE_COUNT = 32; + const subscriber = await ctx.newSubscriberClient(connectionType); + const prefix = `${randomUUIDv7()}:`; + const pattern = `${prefix}*`; + + const received: { message: string; channel: string }[] = []; + const counter = awaitableCounter(); + const count = await subscriber.psubscribe(pattern, (message, channel) => { + received.push({ message, channel }); + counter.increment(); + }); + expect(count).toBe(1); + + for (let i = 0; i < TEST_MESSAGE_COUNT; i++) { + expect(await ctx.redis.publish(`${prefix}${i}`, `msg${i}`)).toBe(1); + } + + await counter.untilValue(TEST_MESSAGE_COUNT); + expect(received).toEqual( + Array.from({ length: TEST_MESSAGE_COUNT }, (_, i) => ({ + message: `msg${i}`, + channel: `${prefix}${i}`, + })), + ); + + await subscriber.punsubscribe(pattern); + }); + + test("psubscribe to multiple patterns receives messages", async () => { + const subscriber = await ctx.newSubscriberClient(connectionType); + const prefixA = `${randomUUIDv7()}:a:`; + const prefixB = `${randomUUIDv7()}:b:`; + + const received: { message: string; channel: string }[] = []; + const counter = awaitableCounter(); + await subscriber.psubscribe([`${prefixA}*`, `${prefixB}*`], (message, channel) => { + received.push({ message, channel }); + counter.increment(); + }); + + expect(await ctx.redis.publish(`${prefixA}1`, "a1")).toBe(1); + expect(await ctx.redis.publish(`${prefixB}1`, "b1")).toBe(1); + expect(await ctx.redis.publish(`${prefixA}2`, "a2")).toBe(1); + + await counter.untilValue(3); + expect(received).toEqual([ + { message: "a1", channel: `${prefixA}1` }, + { message: "b1", channel: `${prefixB}1` }, + { message: "a2", channel: `${prefixA}2` }, + ]); + + await subscriber.punsubscribe([`${prefixA}*`, `${prefixB}*`]); + }); + + test("punsubscribe stops delivery for that pattern", async () => { + const subscriber = await ctx.newSubscriberClient(connectionType); + const prefix = `${randomUUIDv7()}:`; + const pattern = `${prefix}*`; + + const counter = awaitableCounter(); + let calls = 0; + await subscriber.psubscribe(pattern, () => { + calls++; + counter.increment(); + }); + + expect(await ctx.redis.publish(`${prefix}x`, "before")).toBe(1); + await counter.untilValue(1); + expect(calls).toBe(1); + + await subscriber.punsubscribe(pattern); + + expect(await ctx.redis.publish(`${prefix}x`, "after")).toBe(0); + expect(calls).toBe(1); + }); + + test("mixing subscribe and psubscribe on one client", async () => { + const subscriber = await ctx.newSubscriberClient(connectionType); + const exactChannel = `exact-${randomUUIDv7()}`; + const prefix = `${randomUUIDv7()}:`; + + const fromChannel: string[] = []; + const fromPattern: string[] = []; + const counter = awaitableCounter(); + + await subscriber.subscribe(exactChannel, msg => { + fromChannel.push(msg); + counter.increment(); + }); + await subscriber.psubscribe(`${prefix}*`, msg => { + fromPattern.push(msg); + counter.increment(); + }); + + expect(await ctx.redis.publish(exactChannel, "to-channel")).toBe(1); + expect(await ctx.redis.publish(`${prefix}1`, "to-pattern")).toBe(1); + await counter.untilValue(2); + + expect(fromChannel).toEqual(["to-channel"]); + expect(fromPattern).toEqual(["to-pattern"]); + + // punsubscribe() clears pattern listeners but leaves the channel + // subscription intact, so the client stays in subscriber mode. + await subscriber.punsubscribe(); + expect(() => subscriber.set("k", "v")).toThrow("subscriber mode"); + + await subscriber.unsubscribe(); + }); + + test("punsubscribe(pattern, listener) removes only that listener", async () => { + const subscriber = await ctx.newSubscriberClient(connectionType); + const prefix = `${randomUUIDv7()}:`; + const pattern = `${prefix}*`; + + const counter = awaitableCounter(); + let a = 0; + let b = 0; + const listenerA = () => { + a++; + counter.increment(); + }; + const listenerB = () => { + b++; + counter.increment(); + }; + + await subscriber.psubscribe(pattern, listenerA); + await subscriber.psubscribe(pattern, listenerB); + + expect(await ctx.redis.publish(`${prefix}x`, "m1")).toBe(1); + await counter.untilValue(2); + expect(a).toBe(1); + expect(b).toBe(1); + + await subscriber.punsubscribe(pattern, listenerA); + + expect(await ctx.redis.publish(`${prefix}x`, "m2")).toBe(1); + await counter.untilValue(3); + expect(a).toBe(1); + expect(b).toBe(2); + + await subscriber.punsubscribe(pattern); + }); + test("high volume pub/sub", async () => { const channel = testChannel(); @@ -6933,3 +7078,137 @@ describe("RedisClient URL parsing", () => { } }); }); + +// Self-contained mock-server coverage for pattern pub/sub so the psubscribe +// listener path is exercised even when a real Redis/Valkey isn't available. +describe("RedisClient PSUBSCRIBE (mock server)", () => { + const CRLF = "\r\n"; + const bulk = (s: string) => `$${Buffer.byteLength(s)}${CRLF}${s}${CRLF}`; + const HELLO = `%3${CRLF}${bulk("server")}${bulk("redis")}${bulk("proto")}:3${CRLF}${bulk("version")}${bulk("7.4.0")}`; + + type Mock = { + port: number; + pushPmessage: (pattern: string, channel: string, payload: string) => void; + stop: () => void; + }; + + function makeServer(): Mock { + let sock: any; + const server = Bun.listen({ + hostname: "127.0.0.1", + port: 0, + socket: { + open(s) { + sock = s; + }, + data(s, raw) { + const t = raw.toString(); + if (t.includes("HELLO")) { + s.write(HELLO); + } else if (t.includes("PUNSUBSCRIBE")) { + const m = /\r\n(.+?)\r\n$/.exec(t); + s.write(`>3${CRLF}${bulk("punsubscribe")}${bulk(m ? m[1] : "")}:0${CRLF}`); + } else if (t.includes("PSUBSCRIBE")) { + const m = /\r\n(.+?)\r\n$/.exec(t); + s.write(`>3${CRLF}${bulk("psubscribe")}${bulk(m ? m[1] : "")}:1${CRLF}`); + } else if (t.includes("PING")) { + s.write(`+PONG${CRLF}`); + } else { + s.write(`+OK${CRLF}`); + } + }, + }, + }); + return { + port: server.port, + pushPmessage(pattern, channel, payload) { + sock.write(`>4${CRLF}${bulk("pmessage")}${bulk(pattern)}${bulk(channel)}${bulk(payload)}`); + }, + stop: () => server.stop(true), + }; + } + + test("psubscribe requires a listener and has arity 2", () => { + expect(typeof RedisClient.prototype.psubscribe).toBe("function"); + expect(RedisClient.prototype.psubscribe.length).toBe(2); + expect(typeof RedisClient.prototype.punsubscribe).toBe("function"); + }); + + test("psubscribe without a listener throws", async () => { + const srv = makeServer(); + const c = new RedisClient(`redis://127.0.0.1:${srv.port}`, { autoReconnect: false }); + try { + await c.connect(); + // @ts-expect-error psubscribe requires a listener + expect(() => c.psubscribe("evt:*")).toThrow("listener"); + } finally { + c.close(); + srv.stop(); + } + }); + + test("psubscribe delivers pmessage pushes to its listener", async () => { + const srv = makeServer(); + const c = new RedisClient(`redis://127.0.0.1:${srv.port}`, { autoReconnect: false }); + try { + await c.connect(); + + const received: [string, string][] = []; + const counter = awaitableCounter(); + const result = await c.psubscribe("evt:*", (message, channel) => { + received.push([message, channel]); + counter.increment(); + }); + + // psubscribe should resolve with the subscription count, not the raw + // RESP3 push frame it used to leak. + expect(result).toBe(1); + + srv.pushPmessage("evt:*", "evt:1", "payload1"); + srv.pushPmessage("evt:*", "evt:2", "payload2"); + srv.pushPmessage("evt:*", "evt:3", "payload3"); + + await counter.untilValue(3); + expect(received).toEqual([ + ["payload1", "evt:1"], + ["payload2", "evt:2"], + ["payload3", "evt:3"], + ]); + + // Connection stays healthy for regular commands. + expect(await c.ping()).toBe("PONG"); + expect(c.connected).toBe(true); + } finally { + c.close(); + srv.stop(); + } + }); + + test("punsubscribe stops delivery and leaves subscriber mode", async () => { + const srv = makeServer(); + const c = new RedisClient(`redis://127.0.0.1:${srv.port}`, { autoReconnect: false }); + try { + await c.connect(); + + let calls = 0; + const counter = awaitableCounter(); + await c.psubscribe("evt:*", () => { + calls++; + counter.increment(); + }); + + srv.pushPmessage("evt:*", "evt:1", "a"); + await counter.untilValue(1); + expect(calls).toBe(1); + + await c.punsubscribe("evt:*"); + + // After punsubscribe the client leaves subscriber mode, so regular + // commands are allowed and answered again. + expect(await c.set("k", "v")).toBe("OK"); + } finally { + c.close(); + srv.stop(); + } + }); +});