1use crate::cmd::Cmd;
9use crate::config::{GlideClientConfiguration, GlideClusterClientConfiguration};
10use crate::error::GlideError;
11use crate::executor::CommandExecutor;
12use crate::pipeline::dispatch_pipeline;
13use crate::pipeline_options::PipelineOptions;
14use crate::routes::Route;
15use crate::value::FromValkeyValue;
16use crate::{ValkeyFuture, ValkeyResult, ValkeyValue};
17use async_trait::async_trait;
18use bytes::Bytes;
19use glide_core::client::{Client as CoreClient, FINISHED_SCAN_CURSOR};
20use glide_core::cluster_scan_container::{get_cluster_scan_cursor, remove_scan_state_cursor};
21use redis::cluster_routing::RoutingInfo;
22use redis::{ClusterScanArgs, PushInfo, PushKind, ScanStateRC, Value};
23use std::sync::Arc;
24use tokio::sync::Mutex as AsyncMutex;
25use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
26
27mod pipeline;
28pub use pipeline::PipelineExt;
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum PubSubMessageKind {
33 Message,
35 PMessage,
37 SMessage,
39}
40
41#[derive(Debug, Clone)]
45pub struct PubSubMessage {
46 pub kind: PubSubMessageKind,
48 pub channel: Bytes,
50 pub payload: Bytes,
52 pub pattern: Option<Bytes>,
54}
55
56type PushRx = Arc<AsyncMutex<UnboundedReceiver<PushInfo>>>;
66
67fn make_push_channel(
69 has_subscriptions: bool,
70) -> (Option<UnboundedSender<PushInfo>>, Option<PushRx>) {
71 if has_subscriptions {
72 let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
73 (Some(tx), Some(Arc::new(AsyncMutex::new(rx))))
74 } else {
75 (None, None)
76 }
77}
78
79fn pubsub_rx(rx: &Option<PushRx>) -> ValkeyResult<&PushRx> {
82 rx.as_ref()
83 .ok_or_else(|| GlideError::Request("client has no configured pub/sub subscriptions".into()))
84}
85
86async fn recv_pubsub_message(rx: &Option<PushRx>) -> ValkeyResult<PubSubMessage> {
90 let mut guard = pubsub_rx(rx)?.lock().await;
91 loop {
92 match guard.recv().await {
93 Some(push) => {
94 if let Some(msg) = push_to_message(push) {
95 return Ok(msg);
96 }
97 }
98 None => return Err(GlideError::Request("pub/sub channel closed".into())),
99 }
100 }
101}
102
103async fn try_recv_pubsub_message(rx: &Option<PushRx>) -> ValkeyResult<Option<PubSubMessage>> {
106 let mut guard = pubsub_rx(rx)?.lock().await;
107 loop {
108 match guard.try_recv() {
109 Ok(push) => {
110 if let Some(msg) = push_to_message(push) {
111 return Ok(Some(msg));
112 }
113 }
114 Err(tokio::sync::mpsc::error::TryRecvError::Empty) => return Ok(None),
115 Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {
116 return Err(GlideError::Request("pub/sub channel closed".into()));
117 }
118 }
119 }
120}
121
122fn push_to_message(push: PushInfo) -> Option<PubSubMessage> {
126 let (kind, has_pattern) = match push.kind {
127 PushKind::Message => (PubSubMessageKind::Message, false),
128 PushKind::SMessage => (PubSubMessageKind::SMessage, false),
129 PushKind::PMessage => (PubSubMessageKind::PMessage, true),
130 _ => return None,
131 };
132
133 let mut data = push.data.into_iter();
134 let message = (|| {
135 let mut next = || match data.next()? {
136 Value::BulkString(b) => Some(b),
137 _ => None,
138 };
139 let pattern = if has_pattern { Some(next()?) } else { None };
140 let channel = next()?;
141 let payload = next()?;
142 Some(PubSubMessage {
143 kind,
144 channel,
145 payload,
146 pattern,
147 })
148 })();
149
150 if message.is_none() {
151 glide_logger::log_warn(
152 "pubsub",
153 format!("skipping malformed {kind:?} push: missing or non-string channel/payload"),
154 );
155 }
156 message
157}
158
159async fn run_pipeline(
161 core: &CoreClient,
162 pipeline: &crate::pipeline::Pipeline,
163 routing: Option<RoutingInfo>,
164 raise_on_error: bool,
165 options: &PipelineOptions,
166) -> ValkeyResult<Vec<ValkeyValue>> {
167 let reply = dispatch_pipeline(core, pipeline, routing, raise_on_error, options).await?;
168 match reply {
169 ValkeyValue::Array(items) => Ok(items),
170 ValkeyValue::Nil => Ok(Vec::new()),
171 other => unreachable!("unexpected pipeline reply from Valkey: {other:?}"),
172 }
173}
174
175#[derive(Debug, PartialEq, Eq)]
181pub struct ClusterScanCursor(String);
182
183impl ClusterScanCursor {
184 const INITIAL_CURSOR: &'static str = "";
186
187 pub fn new() -> Self {
189 ClusterScanCursor(Self::INITIAL_CURSOR.to_owned())
190 }
191
192 pub fn id(&self) -> &str {
194 &self.0
195 }
196
197 fn is_initial(&self) -> bool {
199 self.0 == Self::INITIAL_CURSOR
200 }
201
202 pub fn is_finished(&self) -> bool {
204 self.0 == FINISHED_SCAN_CURSOR
205 }
206}
207
208impl Default for ClusterScanCursor {
209 fn default() -> Self {
210 Self::new()
211 }
212}
213
214impl Drop for ClusterScanCursor {
215 fn drop(&mut self) {
216 if !self.is_initial() && !self.is_finished() {
218 remove_scan_state_cursor(std::mem::take(&mut self.0));
219 }
220 }
221}
222
223#[derive(Clone)]
228pub struct GlideClient {
229 inner: CoreClient,
230 pubsub_rx: Option<PushRx>,
231}
232
233impl GlideClient {
234 pub async fn connect(config: GlideClientConfiguration) -> ValkeyResult<Self> {
236 let request = config.to_request();
237 let has_subs = config
238 .pubsub_subscriptions
239 .as_ref()
240 .is_some_and(|s| !s.is_empty());
241 let (sender, pubsub_rx) = make_push_channel(has_subs || config.force_pubsub_channel);
242 let inner = CoreClient::new(request, sender)
243 .await
244 .map_err(GlideError::from_connection_error)?;
245 Ok(GlideClient { inner, pubsub_rx })
246 }
247
248 pub async fn get_pubsub_message(&self) -> ValkeyResult<PubSubMessage> {
254 recv_pubsub_message(&self.pubsub_rx).await
255 }
256
257 pub async fn try_get_pubsub_message(&self) -> ValkeyResult<Option<PubSubMessage>> {
261 try_recv_pubsub_message(&self.pubsub_rx).await
262 }
263
264 pub async fn exec(
272 &self,
273 pipeline: &crate::pipeline::Pipeline,
274 raise_on_error: bool,
275 options: &PipelineOptions,
276 ) -> ValkeyResult<Vec<ValkeyValue>> {
277 run_pipeline(&self.inner, pipeline, None, raise_on_error, options).await
278 }
279
280 pub async fn update_connection_password(
293 &self,
294 password: Option<String>,
295 immediate_auth: bool,
296 ) -> ValkeyResult<()> {
297 let mut client = self.inner.clone();
300 client
301 .update_connection_password(password, immediate_auth)
302 .await
303 .map_err(GlideError::from_redis_error)?;
304 Ok(())
305 }
306}
307
308#[async_trait]
309impl CommandExecutor for GlideClient {
310 async fn execute_command(
311 &self,
312 mut cmd: Cmd,
313 route: Option<Route>,
314 ) -> ValkeyResult<ValkeyValue> {
315 let routing = route.map(|r| r.to_routing_info(Some(cmd.as_redis())));
316 let mut client = self.inner.clone();
319 let value = client
320 .send_command(cmd.as_redis_mut(), routing)
321 .await
322 .map_err(GlideError::from_redis_error)?;
323 ValkeyValue::from_redis(value)
324 }
325}
326
327#[derive(Clone)]
333pub struct GlideClusterClient {
334 inner: CoreClient,
335 pubsub_rx: Option<PushRx>,
336}
337
338impl GlideClusterClient {
339 pub async fn connect(config: GlideClusterClientConfiguration) -> ValkeyResult<Self> {
341 let request = config.to_request();
342 let has_subs = config
343 .pubsub_subscriptions
344 .as_ref()
345 .is_some_and(|s| !s.is_empty());
346 let (sender, pubsub_rx) = make_push_channel(has_subs || config.force_pubsub_channel);
347 let inner = CoreClient::new(request, sender)
348 .await
349 .map_err(GlideError::from_connection_error)?;
350 Ok(GlideClusterClient { inner, pubsub_rx })
351 }
352
353 pub async fn get_pubsub_message(&self) -> ValkeyResult<PubSubMessage> {
356 recv_pubsub_message(&self.pubsub_rx).await
357 }
358
359 pub async fn try_get_pubsub_message(&self) -> ValkeyResult<Option<PubSubMessage>> {
362 try_recv_pubsub_message(&self.pubsub_rx).await
363 }
364
365 pub async fn route_command(&self, mut cmd: Cmd, route: Route) -> ValkeyResult<ValkeyValue> {
367 let routing = route.to_routing_info(Some(cmd.as_redis()));
368 let mut client = self.inner.clone();
369 let value = client
370 .send_command(cmd.as_redis_mut(), Some(routing))
371 .await
372 .map_err(GlideError::from_redis_error)?;
373 ValkeyValue::from_redis(value)
374 }
375
376 pub async fn exec(
379 &self,
380 pipeline: &crate::pipeline::Pipeline,
381 raise_on_error: bool,
382 route: Option<Route>,
383 options: &PipelineOptions,
384 ) -> ValkeyResult<Vec<ValkeyValue>> {
385 let routing = route.map(|r| r.to_routing_info(None));
386 run_pipeline(&self.inner, pipeline, routing, raise_on_error, options).await
387 }
388
389 pub async fn update_connection_password(
399 &self,
400 password: Option<String>,
401 immediate_auth: bool,
402 ) -> ValkeyResult<()> {
403 let mut client = self.inner.clone();
404 client
405 .update_connection_password(password, immediate_auth)
406 .await
407 .map_err(GlideError::from_redis_error)?;
408 Ok(())
409 }
410
411 pub async fn cluster_scan(
419 &self,
420 cursor: &ClusterScanCursor,
421 match_pattern: Option<&[u8]>,
422 count: Option<u32>,
423 object_type: Option<crate::commands::options::ObjectType>,
424 ) -> ValkeyResult<(ClusterScanCursor, Vec<Bytes>)> {
425 let scan_state = if cursor.is_initial() {
427 ScanStateRC::new()
428 } else {
429 get_cluster_scan_cursor(cursor.0.clone()).map_err(GlideError::from_redis_error)?
430 };
431
432 let mut builder = ClusterScanArgs::builder();
433 if let Some(p) = match_pattern {
434 builder = builder.with_match_pattern(p.to_vec());
435 }
436 if let Some(c) = count {
437 builder = builder.with_count(c);
438 }
439 if let Some(t) = object_type {
440 builder = builder.with_object_type(t.to_redis());
441 }
442 let args = builder.build();
443
444 let mut client = self.inner.clone();
445 let reply = client
446 .cluster_scan(&scan_state, args)
447 .await
448 .map_err(GlideError::from_redis_error)?;
449
450 let items = match ValkeyValue::from_redis(reply)? {
452 ValkeyValue::Array(items) => items,
453 other => {
454 return Err(GlideError::Request(format!(
455 "unexpected cluster scan reply: {other:?}"
456 )));
457 }
458 };
459 let [cursor_val, keys_val] = <[ValkeyValue; 2]>::try_from(items).map_err(|items| {
460 GlideError::Request(format!("unexpected cluster scan reply arity: {items:?}"))
461 })?;
462 let next = ClusterScanCursor(String::from_owned_valkey_value(cursor_val)?);
463 let keys = match keys_val {
464 ValkeyValue::Array(elems) => elems
465 .into_iter()
466 .map(Bytes::from_owned_valkey_value)
467 .collect::<ValkeyResult<Vec<_>>>()?,
468 ValkeyValue::Nil => Vec::new(),
469 other => vec![Bytes::from_owned_valkey_value(other)?],
470 };
471 Ok((next, keys))
472 }
473}
474
475#[async_trait]
476impl CommandExecutor for GlideClusterClient {
477 async fn execute_command(
478 &self,
479 mut cmd: Cmd,
480 route: Option<Route>,
481 ) -> ValkeyResult<ValkeyValue> {
482 let routing = route.map(|r| r.to_routing_info(Some(cmd.as_redis())));
483 let mut client = self.inner.clone();
484 let value = client
485 .send_command(cmd.as_redis_mut(), routing)
486 .await
487 .map_err(GlideError::from_redis_error)?;
488 ValkeyValue::from_redis(value)
489 }
490}
491
492macro_rules! impl_async_command {
495 ($ty:ty) => {
496 #[::sealed::sealed]
497 impl crate::commands::core::CommandDispatch for $ty {
498 fn glide_dispatch_command<'a>(&'a self, mut cmd: Cmd) -> ValkeyFuture<'a, ValkeyValue> {
499 let mut client = self.inner.clone();
500 Box::pin(async move {
501 let value = client
502 .send_command(cmd.as_redis_mut(), None)
503 .await
504 .map_err(GlideError::from_redis_error)?;
505 ValkeyValue::from_redis(value)
506 })
507 }
508 }
509 };
510}
511
512impl_async_command!(GlideClient);
513impl_async_command!(GlideClusterClient);
514
515macro_rules! impl_script_invoke {
518 ($ty:ty) => {
519 #[::sealed::sealed]
520 impl crate::script::ScriptInvoke for $ty {
521 fn glide_invoke_script<'a>(
522 &'a self,
523 hash: &'a str,
524 keys: &'a [Vec<u8>],
525 args: &'a [Vec<u8>],
526 ) -> ValkeyFuture<'a, ValkeyValue> {
527 let mut client = self.inner.clone();
528 Box::pin(async move {
529 let key_refs: Vec<&[u8]> = keys.iter().map(Vec::as_slice).collect();
530 let arg_refs: Vec<&[u8]> = args.iter().map(Vec::as_slice).collect();
531 let value = client
532 .invoke_script(hash, &key_refs, &arg_refs, None)
533 .await
534 .map_err(GlideError::from_redis_error)?;
535 ValkeyValue::from_redis(value)
536 })
537 }
538 }
539 };
540}
541
542impl_script_invoke!(GlideClient);
543impl_script_invoke!(GlideClusterClient);
544
545macro_rules! impl_pipeline_dispatch {
548 ($ty:ty) => {
549 #[::sealed::sealed]
550 impl pipeline::PipelineDispatch for $ty {
551 fn glide_dispatch_pipeline<'a>(
552 &self,
553 pipeline: &'a crate::pipeline::Pipeline,
554 ) -> ValkeyFuture<'a, ValkeyValue> {
555 let core = self.inner.clone();
558 Box::pin(async move {
559 let opts = &PipelineOptions::default();
560 dispatch_pipeline(&core, pipeline, None, true, opts).await
561 })
562 }
563 }
564 };
565}
566
567impl_pipeline_dispatch!(GlideClient);
568impl_pipeline_dispatch!(GlideClusterClient);
569
570#[cfg(test)]
573mod push_tests {
574 use super::*;
575
576 fn bulk(s: &str) -> Value {
577 Value::BulkString(Bytes::copy_from_slice(s.as_bytes()))
578 }
579
580 #[test]
581 fn message_push_parses_channel_and_payload() {
582 let push = PushInfo {
583 kind: PushKind::Message,
584 data: vec![bulk("chan"), bulk("hello")],
585 };
586 let msg = push_to_message(push).expect("message should parse");
587 assert_eq!(msg.kind, PubSubMessageKind::Message);
588 assert_eq!(msg.channel, Bytes::from_static(b"chan"));
589 assert_eq!(msg.payload, Bytes::from_static(b"hello"));
590 assert_eq!(msg.pattern, None);
591 }
592
593 #[test]
594 fn pmessage_push_parses_pattern() {
595 let push = PushInfo {
596 kind: PushKind::PMessage,
597 data: vec![bulk("ch.*"), bulk("chan"), bulk("hello")],
598 };
599 let msg = push_to_message(push).expect("pmessage should parse");
600 assert_eq!(msg.kind, PubSubMessageKind::PMessage);
601 assert_eq!(msg.pattern, Some(Bytes::from_static(b"ch.*")));
602 assert_eq!(msg.channel, Bytes::from_static(b"chan"));
603 assert_eq!(msg.payload, Bytes::from_static(b"hello"));
604 }
605
606 #[test]
607 fn non_message_kind_is_skipped() {
608 let push = PushInfo {
609 kind: PushKind::Subscribe,
610 data: vec![bulk("chan"), Value::Int(1)],
611 };
612 assert!(push_to_message(push).is_none());
613 }
614
615 #[test]
616 fn non_string_payload_is_skipped_not_fabricated() {
617 let push = PushInfo {
620 kind: PushKind::Message,
621 data: vec![bulk("chan"), Value::Int(5)],
622 };
623 assert!(push_to_message(push).is_none());
624 }
625
626 #[test]
627 fn short_message_push_is_skipped() {
628 let push = PushInfo {
629 kind: PushKind::Message,
630 data: vec![bulk("chan")],
631 };
632 assert!(push_to_message(push).is_none());
633 }
634}
635
636#[cfg(test)]
637mod cluster_scan_cursor_tests {
638 use super::*;
639 use glide_core::cluster_scan_container::{get_cluster_scan_cursor, insert_cluster_scan_cursor};
640
641 #[test]
642 fn is_initial() {
643 let cases = [
645 (ClusterScanCursor::new(), true),
646 (ClusterScanCursor::default(), true),
647 (ClusterScanCursor("id".to_owned()), false),
648 (ClusterScanCursor(FINISHED_SCAN_CURSOR.to_owned()), false),
649 ];
650
651 for (cursor, expected) in &cases {
652 assert_eq!(cursor.is_initial(), *expected);
653 }
654 }
655
656 #[test]
657 fn is_finished() {
658 let cases = [
660 (ClusterScanCursor::new(), false),
661 (ClusterScanCursor::default(), false),
662 (ClusterScanCursor("id".to_owned()), false),
663 (ClusterScanCursor(FINISHED_SCAN_CURSOR.to_owned()), true),
664 ];
665
666 for (cursor, expected) in &cases {
667 assert_eq!(cursor.is_finished(), *expected);
668 }
669 }
670
671 #[test]
672 fn drop_releases_container_entry() {
673 let cursor = ClusterScanCursor::new();
675 let id = cursor.id().to_owned();
676 assert!(get_cluster_scan_cursor(id.clone()).is_err());
677
678 drop(cursor);
679 assert!(get_cluster_scan_cursor(id).is_err());
680
681 let id = insert_cluster_scan_cursor(ScanStateRC::new());
683 assert!(get_cluster_scan_cursor(id.clone()).is_ok());
684
685 drop(ClusterScanCursor(id.clone()));
686 assert!(get_cluster_scan_cursor(id).is_err());
687
688 let cursor = ClusterScanCursor(FINISHED_SCAN_CURSOR.to_owned());
690 let id = cursor.id().to_owned();
691 assert!(get_cluster_scan_cursor(id.clone()).is_err());
692
693 drop(cursor);
694 assert!(get_cluster_scan_cursor(id).is_err());
695 }
696}