use crate::cmd::Cmd;
use crate::config::{GlideClientConfiguration, GlideClusterClientConfiguration};
use crate::error::GlideError;
use crate::executor::CommandExecutor;
use crate::pipeline::dispatch_pipeline;
use crate::pipeline_options::PipelineOptions;
use crate::routes::Route;
use crate::value::FromValkeyValue;
use crate::{ValkeyFuture, ValkeyResult, ValkeyValue};
use async_trait::async_trait;
use bytes::Bytes;
use glide_core::client::{Client as CoreClient, FINISHED_SCAN_CURSOR};
use glide_core::cluster_scan_container::{get_cluster_scan_cursor, remove_scan_state_cursor};
use redis::cluster_routing::RoutingInfo;
use redis::{ClusterScanArgs, PushInfo, PushKind, ScanStateRC, Value};
use std::sync::Arc;
use tokio::sync::Mutex as AsyncMutex;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
mod pipeline;
pub use pipeline::PipelineExt;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PubSubMessageKind {
Message,
PMessage,
SMessage,
}
#[derive(Debug, Clone)]
pub struct PubSubMessage {
pub kind: PubSubMessageKind,
pub channel: Bytes,
pub payload: Bytes,
pub pattern: Option<Bytes>,
}
type PushRx = Arc<AsyncMutex<UnboundedReceiver<PushInfo>>>;
fn make_push_channel(
has_subscriptions: bool,
) -> (Option<UnboundedSender<PushInfo>>, Option<PushRx>) {
if has_subscriptions {
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
(Some(tx), Some(Arc::new(AsyncMutex::new(rx))))
} else {
(None, None)
}
}
fn pubsub_rx(rx: &Option<PushRx>) -> ValkeyResult<&PushRx> {
rx.as_ref()
.ok_or_else(|| GlideError::Request("client has no configured pub/sub subscriptions".into()))
}
async fn recv_pubsub_message(rx: &Option<PushRx>) -> ValkeyResult<PubSubMessage> {
let mut guard = pubsub_rx(rx)?.lock().await;
loop {
match guard.recv().await {
Some(push) => {
if let Some(msg) = push_to_message(push) {
return Ok(msg);
}
}
None => return Err(GlideError::Request("pub/sub channel closed".into())),
}
}
}
async fn try_recv_pubsub_message(rx: &Option<PushRx>) -> ValkeyResult<Option<PubSubMessage>> {
let mut guard = pubsub_rx(rx)?.lock().await;
loop {
match guard.try_recv() {
Ok(push) => {
if let Some(msg) = push_to_message(push) {
return Ok(Some(msg));
}
}
Err(tokio::sync::mpsc::error::TryRecvError::Empty) => return Ok(None),
Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {
return Err(GlideError::Request("pub/sub channel closed".into()));
}
}
}
}
fn push_to_message(push: PushInfo) -> Option<PubSubMessage> {
let (kind, has_pattern) = match push.kind {
PushKind::Message => (PubSubMessageKind::Message, false),
PushKind::SMessage => (PubSubMessageKind::SMessage, false),
PushKind::PMessage => (PubSubMessageKind::PMessage, true),
_ => return None,
};
let mut data = push.data.into_iter();
let message = (|| {
let mut next = || match data.next()? {
Value::BulkString(b) => Some(b),
_ => None,
};
let pattern = if has_pattern { Some(next()?) } else { None };
let channel = next()?;
let payload = next()?;
Some(PubSubMessage {
kind,
channel,
payload,
pattern,
})
})();
if message.is_none() {
glide_logger::log_warn(
"pubsub",
format!("skipping malformed {kind:?} push: missing or non-string channel/payload"),
);
}
message
}
async fn run_pipeline(
core: &CoreClient,
pipeline: &crate::pipeline::Pipeline,
routing: Option<RoutingInfo>,
raise_on_error: bool,
options: &PipelineOptions,
) -> ValkeyResult<Vec<ValkeyValue>> {
let reply = dispatch_pipeline(core, pipeline, routing, raise_on_error, options).await?;
match reply {
ValkeyValue::Array(items) => Ok(items),
ValkeyValue::Nil => Ok(Vec::new()),
other => unreachable!("unexpected pipeline reply from Valkey: {other:?}"),
}
}
#[derive(Debug, PartialEq, Eq)]
pub struct ClusterScanCursor(String);
impl ClusterScanCursor {
const INITIAL_CURSOR: &'static str = "";
pub fn new() -> Self {
ClusterScanCursor(Self::INITIAL_CURSOR.to_owned())
}
pub fn id(&self) -> &str {
&self.0
}
fn is_initial(&self) -> bool {
self.0 == Self::INITIAL_CURSOR
}
pub fn is_finished(&self) -> bool {
self.0 == FINISHED_SCAN_CURSOR
}
}
impl Default for ClusterScanCursor {
fn default() -> Self {
Self::new()
}
}
impl Drop for ClusterScanCursor {
fn drop(&mut self) {
if !self.is_initial() && !self.is_finished() {
remove_scan_state_cursor(std::mem::take(&mut self.0));
}
}
}
#[derive(Clone)]
pub struct GlideClient {
inner: CoreClient,
pubsub_rx: Option<PushRx>,
}
impl GlideClient {
pub async fn connect(config: GlideClientConfiguration) -> ValkeyResult<Self> {
let request = config.to_request();
let has_subs = config
.pubsub_subscriptions
.as_ref()
.is_some_and(|s| !s.is_empty());
let (sender, pubsub_rx) = make_push_channel(has_subs || config.force_pubsub_channel);
let inner = CoreClient::new(request, sender)
.await
.map_err(GlideError::from_connection_error)?;
Ok(GlideClient { inner, pubsub_rx })
}
pub async fn get_pubsub_message(&self) -> ValkeyResult<PubSubMessage> {
recv_pubsub_message(&self.pubsub_rx).await
}
pub async fn try_get_pubsub_message(&self) -> ValkeyResult<Option<PubSubMessage>> {
try_recv_pubsub_message(&self.pubsub_rx).await
}
pub async fn exec(
&self,
pipeline: &crate::pipeline::Pipeline,
raise_on_error: bool,
options: &PipelineOptions,
) -> ValkeyResult<Vec<ValkeyValue>> {
run_pipeline(&self.inner, pipeline, None, raise_on_error, options).await
}
pub async fn update_connection_password(
&self,
password: Option<String>,
immediate_auth: bool,
) -> ValkeyResult<()> {
let mut client = self.inner.clone();
client
.update_connection_password(password, immediate_auth)
.await
.map_err(GlideError::from_redis_error)?;
Ok(())
}
}
#[async_trait]
impl CommandExecutor for GlideClient {
async fn execute_command(
&self,
mut cmd: Cmd,
route: Option<Route>,
) -> ValkeyResult<ValkeyValue> {
let routing = route.map(|r| r.to_routing_info(Some(cmd.as_redis())));
let mut client = self.inner.clone();
let value = client
.send_command(cmd.as_redis_mut(), routing)
.await
.map_err(GlideError::from_redis_error)?;
ValkeyValue::from_redis(value)
}
}
#[derive(Clone)]
pub struct GlideClusterClient {
inner: CoreClient,
pubsub_rx: Option<PushRx>,
}
impl GlideClusterClient {
pub async fn connect(config: GlideClusterClientConfiguration) -> ValkeyResult<Self> {
let request = config.to_request();
let has_subs = config
.pubsub_subscriptions
.as_ref()
.is_some_and(|s| !s.is_empty());
let (sender, pubsub_rx) = make_push_channel(has_subs || config.force_pubsub_channel);
let inner = CoreClient::new(request, sender)
.await
.map_err(GlideError::from_connection_error)?;
Ok(GlideClusterClient { inner, pubsub_rx })
}
pub async fn get_pubsub_message(&self) -> ValkeyResult<PubSubMessage> {
recv_pubsub_message(&self.pubsub_rx).await
}
pub async fn try_get_pubsub_message(&self) -> ValkeyResult<Option<PubSubMessage>> {
try_recv_pubsub_message(&self.pubsub_rx).await
}
pub async fn route_command(&self, mut cmd: Cmd, route: Route) -> ValkeyResult<ValkeyValue> {
let routing = route.to_routing_info(Some(cmd.as_redis()));
let mut client = self.inner.clone();
let value = client
.send_command(cmd.as_redis_mut(), Some(routing))
.await
.map_err(GlideError::from_redis_error)?;
ValkeyValue::from_redis(value)
}
pub async fn exec(
&self,
pipeline: &crate::pipeline::Pipeline,
raise_on_error: bool,
route: Option<Route>,
options: &PipelineOptions,
) -> ValkeyResult<Vec<ValkeyValue>> {
let routing = route.map(|r| r.to_routing_info(None));
run_pipeline(&self.inner, pipeline, routing, raise_on_error, options).await
}
pub async fn update_connection_password(
&self,
password: Option<String>,
immediate_auth: bool,
) -> ValkeyResult<()> {
let mut client = self.inner.clone();
client
.update_connection_password(password, immediate_auth)
.await
.map_err(GlideError::from_redis_error)?;
Ok(())
}
pub async fn cluster_scan(
&self,
cursor: &ClusterScanCursor,
match_pattern: Option<&[u8]>,
count: Option<u32>,
object_type: Option<crate::commands::options::ObjectType>,
) -> ValkeyResult<(ClusterScanCursor, Vec<Bytes>)> {
let scan_state = if cursor.is_initial() {
ScanStateRC::new()
} else {
get_cluster_scan_cursor(cursor.0.clone()).map_err(GlideError::from_redis_error)?
};
let mut builder = ClusterScanArgs::builder();
if let Some(p) = match_pattern {
builder = builder.with_match_pattern(p.to_vec());
}
if let Some(c) = count {
builder = builder.with_count(c);
}
if let Some(t) = object_type {
builder = builder.with_object_type(t.to_redis());
}
let args = builder.build();
let mut client = self.inner.clone();
let reply = client
.cluster_scan(&scan_state, args)
.await
.map_err(GlideError::from_redis_error)?;
let items = match ValkeyValue::from_redis(reply)? {
ValkeyValue::Array(items) => items,
other => {
return Err(GlideError::Request(format!(
"unexpected cluster scan reply: {other:?}"
)));
}
};
let [cursor_val, keys_val] = <[ValkeyValue; 2]>::try_from(items).map_err(|items| {
GlideError::Request(format!("unexpected cluster scan reply arity: {items:?}"))
})?;
let next = ClusterScanCursor(String::from_owned_valkey_value(cursor_val)?);
let keys = match keys_val {
ValkeyValue::Array(elems) => elems
.into_iter()
.map(Bytes::from_owned_valkey_value)
.collect::<ValkeyResult<Vec<_>>>()?,
ValkeyValue::Nil => Vec::new(),
other => vec![Bytes::from_owned_valkey_value(other)?],
};
Ok((next, keys))
}
}
#[async_trait]
impl CommandExecutor for GlideClusterClient {
async fn execute_command(
&self,
mut cmd: Cmd,
route: Option<Route>,
) -> ValkeyResult<ValkeyValue> {
let routing = route.map(|r| r.to_routing_info(Some(cmd.as_redis())));
let mut client = self.inner.clone();
let value = client
.send_command(cmd.as_redis_mut(), routing)
.await
.map_err(GlideError::from_redis_error)?;
ValkeyValue::from_redis(value)
}
}
macro_rules! impl_async_command {
($ty:ty) => {
#[::sealed::sealed]
impl crate::commands::core::CommandDispatch for $ty {
fn glide_dispatch_command<'a>(&'a self, mut cmd: Cmd) -> ValkeyFuture<'a, ValkeyValue> {
let mut client = self.inner.clone();
Box::pin(async move {
let value = client
.send_command(cmd.as_redis_mut(), None)
.await
.map_err(GlideError::from_redis_error)?;
ValkeyValue::from_redis(value)
})
}
}
};
}
impl_async_command!(GlideClient);
impl_async_command!(GlideClusterClient);
macro_rules! impl_script_invoke {
($ty:ty) => {
#[::sealed::sealed]
impl crate::script::ScriptInvoke for $ty {
fn glide_invoke_script<'a>(
&'a self,
hash: &'a str,
keys: &'a [Vec<u8>],
args: &'a [Vec<u8>],
) -> ValkeyFuture<'a, ValkeyValue> {
let mut client = self.inner.clone();
Box::pin(async move {
let key_refs: Vec<&[u8]> = keys.iter().map(Vec::as_slice).collect();
let arg_refs: Vec<&[u8]> = args.iter().map(Vec::as_slice).collect();
let value = client
.invoke_script(hash, &key_refs, &arg_refs, None)
.await
.map_err(GlideError::from_redis_error)?;
ValkeyValue::from_redis(value)
})
}
}
};
}
impl_script_invoke!(GlideClient);
impl_script_invoke!(GlideClusterClient);
macro_rules! impl_pipeline_dispatch {
($ty:ty) => {
#[::sealed::sealed]
impl pipeline::PipelineDispatch for $ty {
fn glide_dispatch_pipeline<'a>(
&self,
pipeline: &'a crate::pipeline::Pipeline,
) -> ValkeyFuture<'a, ValkeyValue> {
let core = self.inner.clone();
Box::pin(async move {
let opts = &PipelineOptions::default();
dispatch_pipeline(&core, pipeline, None, true, opts).await
})
}
}
};
}
impl_pipeline_dispatch!(GlideClient);
impl_pipeline_dispatch!(GlideClusterClient);
#[cfg(test)]
mod push_tests {
use super::*;
fn bulk(s: &str) -> Value {
Value::BulkString(Bytes::copy_from_slice(s.as_bytes()))
}
#[test]
fn message_push_parses_channel_and_payload() {
let push = PushInfo {
kind: PushKind::Message,
data: vec![bulk("chan"), bulk("hello")],
};
let msg = push_to_message(push).expect("message should parse");
assert_eq!(msg.kind, PubSubMessageKind::Message);
assert_eq!(msg.channel, Bytes::from_static(b"chan"));
assert_eq!(msg.payload, Bytes::from_static(b"hello"));
assert_eq!(msg.pattern, None);
}
#[test]
fn pmessage_push_parses_pattern() {
let push = PushInfo {
kind: PushKind::PMessage,
data: vec![bulk("ch.*"), bulk("chan"), bulk("hello")],
};
let msg = push_to_message(push).expect("pmessage should parse");
assert_eq!(msg.kind, PubSubMessageKind::PMessage);
assert_eq!(msg.pattern, Some(Bytes::from_static(b"ch.*")));
assert_eq!(msg.channel, Bytes::from_static(b"chan"));
assert_eq!(msg.payload, Bytes::from_static(b"hello"));
}
#[test]
fn non_message_kind_is_skipped() {
let push = PushInfo {
kind: PushKind::Subscribe,
data: vec![bulk("chan"), Value::Int(1)],
};
assert!(push_to_message(push).is_none());
}
#[test]
fn non_string_payload_is_skipped_not_fabricated() {
let push = PushInfo {
kind: PushKind::Message,
data: vec![bulk("chan"), Value::Int(5)],
};
assert!(push_to_message(push).is_none());
}
#[test]
fn short_message_push_is_skipped() {
let push = PushInfo {
kind: PushKind::Message,
data: vec![bulk("chan")],
};
assert!(push_to_message(push).is_none());
}
}
#[cfg(test)]
mod cluster_scan_cursor_tests {
use super::*;
use glide_core::cluster_scan_container::{get_cluster_scan_cursor, insert_cluster_scan_cursor};
#[test]
fn is_initial() {
let cases = [
(ClusterScanCursor::new(), true),
(ClusterScanCursor::default(), true),
(ClusterScanCursor("id".to_owned()), false),
(ClusterScanCursor(FINISHED_SCAN_CURSOR.to_owned()), false),
];
for (cursor, expected) in &cases {
assert_eq!(cursor.is_initial(), *expected);
}
}
#[test]
fn is_finished() {
let cases = [
(ClusterScanCursor::new(), false),
(ClusterScanCursor::default(), false),
(ClusterScanCursor("id".to_owned()), false),
(ClusterScanCursor(FINISHED_SCAN_CURSOR.to_owned()), true),
];
for (cursor, expected) in &cases {
assert_eq!(cursor.is_finished(), *expected);
}
}
#[test]
fn drop_releases_container_entry() {
let cursor = ClusterScanCursor::new();
let id = cursor.id().to_owned();
assert!(get_cluster_scan_cursor(id.clone()).is_err());
drop(cursor);
assert!(get_cluster_scan_cursor(id).is_err());
let id = insert_cluster_scan_cursor(ScanStateRC::new());
assert!(get_cluster_scan_cursor(id.clone()).is_ok());
drop(ClusterScanCursor(id.clone()));
assert!(get_cluster_scan_cursor(id).is_err());
let cursor = ClusterScanCursor(FINISHED_SCAN_CURSOR.to_owned());
let id = cursor.id().to_owned();
assert!(get_cluster_scan_cursor(id.clone()).is_err());
drop(cursor);
assert!(get_cluster_scan_cursor(id).is_err());
}
}