Skip to main content

glide/commands/
pubsub.rs

1// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
2//! Pub/Sub subscription management and introspection commands.
3//!
4//! This covers runtime (un)subscription (`SUBSCRIBE`/`PSUBSCRIBE`/`SSUBSCRIBE`
5//! and counterparts), shard publishing (`SPUBLISH`), and the `PUBSUB`
6//! introspection subcommands. Plain `PUBLISH` lives in the unified command API
7//! ([`crate::AsyncCommands::publish`]). Receiving messages via a subscription
8//! is delivered through the client (`get_pubsub_message`), not through these
9//! commands.
10
11use crate::ValkeyResult;
12use crate::cmd::Cmd;
13use crate::executor::CommandExecutor;
14use crate::value::FromValkeyValue;
15use crate::value::ValkeyValue;
16use crate::value::to_glide_error;
17use crate::write::ToValkeyArgs;
18use async_trait::async_trait;
19use bytes::Bytes;
20
21/// Pub/Sub commands (`SUBSCRIBE`/`UNSUBSCRIBE` at runtime, `SPUBLISH`,
22/// `PUBSUB ...`).
23#[async_trait]
24pub trait PubSubCommands: CommandExecutor {
25    /// Subscribe to one or more exact channels at runtime (`SUBSCRIBE`).
26    ///
27    /// Received messages are delivered via `get_pubsub_message` /
28    /// `try_get_pubsub_message`. The client must have the Pub/Sub push channel
29    /// enabled — either by configuring connect-time `subscriptions` or by calling
30    /// `enable_pubsub()` on the configuration — otherwise messages are not
31    /// captured.
32    ///
33    /// Note: runtime subscriptions are session-scoped and are not automatically
34    /// restored after a reconnect (connect-time subscriptions are).
35    async fn subscribe<C: ToValkeyArgs + Send + Sync>(&self, channels: &[C]) -> ValkeyResult<()> {
36        self.pubsub_subscribe_impl("SUBSCRIBE", channels).await
37    }
38
39    /// Unsubscribe from exact channels (`UNSUBSCRIBE`). An empty slice
40    /// unsubscribes from all exact channels.
41    async fn unsubscribe<C: ToValkeyArgs + Send + Sync>(&self, channels: &[C]) -> ValkeyResult<()> {
42        self.pubsub_subscribe_impl("UNSUBSCRIBE", channels).await
43    }
44
45    /// Subscribe to one or more glob patterns at runtime (`PSUBSCRIBE`). See
46    /// [`Self::subscribe`] for delivery requirements.
47    async fn psubscribe<C: ToValkeyArgs + Send + Sync>(&self, patterns: &[C]) -> ValkeyResult<()> {
48        self.pubsub_subscribe_impl("PSUBSCRIBE", patterns).await
49    }
50
51    /// Unsubscribe from patterns (`PUNSUBSCRIBE`). An empty slice unsubscribes
52    /// from all patterns.
53    async fn punsubscribe<C: ToValkeyArgs + Send + Sync>(
54        &self,
55        patterns: &[C],
56    ) -> ValkeyResult<()> {
57        self.pubsub_subscribe_impl("PUNSUBSCRIBE", patterns).await
58    }
59
60    /// Subscribe to one or more shard channels at runtime (`SSUBSCRIBE`, cluster
61    /// only). See [`Self::subscribe`] for delivery requirements.
62    async fn ssubscribe<C: ToValkeyArgs + Send + Sync>(&self, channels: &[C]) -> ValkeyResult<()> {
63        self.pubsub_subscribe_impl("SSUBSCRIBE", channels).await
64    }
65
66    /// Unsubscribe from shard channels (`SUNSUBSCRIBE`). An empty slice
67    /// unsubscribes from all shard channels.
68    async fn sunsubscribe<C: ToValkeyArgs + Send + Sync>(
69        &self,
70        channels: &[C],
71    ) -> ValkeyResult<()> {
72        self.pubsub_subscribe_impl("SUNSUBSCRIBE", channels).await
73    }
74
75    #[doc(hidden)]
76    async fn pubsub_subscribe_impl<C: ToValkeyArgs + Send + Sync>(
77        &self,
78        keyword: &'static str,
79        channels: &[C],
80    ) -> ValkeyResult<()> {
81        let mut cmd = Cmd::new();
82        cmd.arg(keyword);
83        for c in channels {
84            cmd.arg(c);
85        }
86        self.execute_command(cmd, None).await?;
87        Ok(())
88    }
89
90    // NOTE: `PUBLISH` and `SPUBLISH` are implemented in the
91    // unified command table (`crate::AsyncCommands::publish`).
92
93    /// List active channels, optionally matching `pattern` (`PUBSUB CHANNELS`).
94    async fn pubsub_channels(&self, pattern: Option<&[u8]>) -> ValkeyResult<Vec<Bytes>> {
95        self.pubsub_channels_impl("CHANNELS", pattern).await
96    }
97
98    /// List active shard channels, optionally matching `pattern`
99    /// (`PUBSUB SHARDCHANNELS`).
100    async fn pubsub_shardchannels(&self, pattern: Option<&[u8]>) -> ValkeyResult<Vec<Bytes>> {
101        self.pubsub_channels_impl("SHARDCHANNELS", pattern).await
102    }
103
104    /// Get the number of subscriptions to patterns (`PUBSUB NUMPAT`).
105    async fn pubsub_numpat(&self) -> ValkeyResult<i64> {
106        let mut cmd = Cmd::new();
107        cmd.arg("PUBSUB").arg("NUMPAT");
108        i64::from_owned_valkey_value(self.execute_command(cmd, None).await?)
109    }
110
111    /// Get the number of subscribers per channel (`PUBSUB NUMSUB`).
112    async fn pubsub_numsub<C: ToValkeyArgs + Send + Sync>(
113        &self,
114        channels: &[C],
115    ) -> ValkeyResult<Vec<(Bytes, i64)>> {
116        self.pubsub_numsub_impl("NUMSUB", channels).await
117    }
118
119    /// Get the number of subscribers per shard channel (`PUBSUB SHARDNUMSUB`).
120    async fn pubsub_shardnumsub<C: ToValkeyArgs + Send + Sync>(
121        &self,
122        channels: &[C],
123    ) -> ValkeyResult<Vec<(Bytes, i64)>> {
124        self.pubsub_numsub_impl("SHARDNUMSUB", channels).await
125    }
126
127    #[doc(hidden)]
128    async fn pubsub_channels_impl(
129        &self,
130        sub: &'static str,
131        pattern: Option<&[u8]>,
132    ) -> ValkeyResult<Vec<Bytes>> {
133        let mut cmd = Cmd::new();
134        cmd.arg("PUBSUB").arg(sub);
135        if let Some(p) = pattern {
136            cmd.arg(p);
137        }
138        match self.execute_command(cmd, None).await? {
139            ValkeyValue::Array(items) => items
140                .into_iter()
141                .map(Bytes::from_owned_valkey_value)
142                .collect(),
143            ValkeyValue::Nil => Ok(Vec::new()),
144            other => Ok(vec![Bytes::from_owned_valkey_value(other)?]),
145        }
146    }
147
148    #[doc(hidden)]
149    async fn pubsub_numsub_impl<C: ToValkeyArgs + Send + Sync>(
150        &self,
151        sub: &'static str,
152        channels: &[C],
153    ) -> ValkeyResult<Vec<(Bytes, i64)>> {
154        let mut cmd = Cmd::new();
155        cmd.arg("PUBSUB").arg(sub);
156        for c in channels {
157            cmd.arg(c);
158        }
159        parse_numsub(self.execute_command(cmd, None).await?)
160    }
161}
162
163/// Parse a `PUBSUB NUMSUB` reply (flat `[channel, count, ...]` or RESP3 map).
164fn parse_numsub(v: ValkeyValue) -> ValkeyResult<Vec<(Bytes, i64)>> {
165    match v {
166        ValkeyValue::Nil => Ok(Vec::new()),
167        ValkeyValue::Map(pairs) => pairs
168            .into_iter()
169            .map(|(c, n)| {
170                Ok((
171                    Bytes::from_owned_valkey_value(c)?,
172                    i64::from_owned_valkey_value(n)?,
173                ))
174            })
175            .collect(),
176        ValkeyValue::Array(items) => {
177            let mut out = Vec::with_capacity(items.len() / 2);
178            let mut iter = items.into_iter();
179            while let (Some(c), Some(n)) = (iter.next(), iter.next()) {
180                out.push((
181                    Bytes::from_owned_valkey_value(c)?,
182                    i64::from_owned_valkey_value(n)?,
183                ));
184            }
185            Ok(out)
186        }
187        other => Err(to_glide_error(other, "Unexpected PUBSUB NUMSUB reply.")),
188    }
189}
190
191impl<T: CommandExecutor + ?Sized> PubSubCommands for T {}