Add RESP3 support to cluster connections.

amazon-contributing · Dec 11, 2023 · a70078d · a70078d
1 parent 63011d6
commit a70078d
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 1 deletion.
diff --git a/redis/src/cluster.rs b/redis/src/cluster.rs
@@ -910,6 +910,7 @@ pub(crate) fn get_connection_info(
         redis: RedisConnectionInfo {
             password: cluster_params.password,
             username: cluster_params.username,
+            use_resp3: cluster_params.use_resp3,
             ..Default::default()
         },
     })

diff --git a/redis/src/cluster_client.rs b/redis/src/cluster_client.rs
@@ -33,6 +33,7 @@ struct BuilderParams {
     retries_configuration: RetryParams,
     connection_timeout: Option<Duration>,
     topology_checks_interval: Option<Duration>,
+    use_resp3: bool,
 }
 
 #[derive(Clone)]
@@ -86,6 +87,7 @@ pub(crate) struct ClusterParams {
     pub(crate) connection_timeout: Duration,
     pub(crate) topology_checks_interval: Option<Duration>,
     pub(crate) tls_params: Option<TlsConnParams>,
+    pub(crate) use_resp3: bool,
 }
 
 impl ClusterParams {
@@ -109,6 +111,7 @@ impl ClusterParams {
             connection_timeout: value.connection_timeout.unwrap_or(Duration::MAX),
             topology_checks_interval: value.topology_checks_interval,
             tls_params,
+            use_resp3: value.use_resp3,
         })
     }
 }
@@ -315,6 +318,15 @@ impl ClusterClientBuilder {
         self
     }
 
+    /// Sets whether the new ClusterClient should connect to the servers using RESP3.
+    ///
+    /// If not set, the default is to use RESP3.
+    #[cfg(any(feature = "tls-native-tls", feature = "tls-rustls"))]
+    pub fn use_resp3(mut self, use_resp3: bool) -> ClusterClientBuilder {
+        self.builder_params.use_resp3 = use_resp3;
+        self
+    }
+
     /// Use `build()`.
     #[deprecated(since = "0.22.0", note = "Use build()")]
     pub fn open(self) -> RedisResult<ClusterClient> {

diff --git a/redis/tests/support/cluster.rs b/redis/tests/support/cluster.rs
@@ -16,6 +16,7 @@ use tempfile::TempDir;
 
 use crate::support::{build_keys_and_certs_for_tls, Module};
 
+use super::use_resp3;
 #[cfg(feature = "tls-rustls")]
 use super::{build_single_client, load_certs_from_file};
 
@@ -343,6 +344,7 @@ impl TestClusterContext {
             .map(RedisServer::connection_info)
             .collect();
         let mut builder = redis::cluster::ClusterClientBuilder::new(initial_nodes.clone());
+        builder = builder.use_resp3(use_resp3());
 
         #[cfg(feature = "tls-rustls")]
         if mtls_enabled {

diff --git a/redis/tests/support/mod.rs b/redis/tests/support/mod.rs
@@ -20,7 +20,7 @@ use redis::{ClientTlsConfig, TlsCertificates};
 use socket2::{Domain, Socket, Type};
 use tempfile::TempDir;
 
-fn use_resp3() -> bool {
+pub fn use_resp3() -> bool {
     env::var("RESP3").unwrap_or_default() == "true"
 }
 

diff --git a/redis/tests/test_cluster.rs b/redis/tests/test_cluster.rs
@@ -136,6 +136,45 @@ fn test_cluster_multi_shard_commands() {
     assert_eq!(res, vec!["bazz", "bar", "foo"]);
 }
 
+#[test]
+fn test_cluster_resp3() {
+    if !use_resp3() {
+        return;
+    }
+    let cluster = TestClusterContext::new(3, 0);
+
+    let mut connection = cluster.connection();
+
+    let _: () = redis::cmd("HSET")
+        .arg("hash")
+        .arg("foo")
+        .arg("baz")
+        .execute(&mut connection);
+    let _: () = redis::cmd("HSET")
+        .arg("hash")
+        .arg("bar")
+        .arg("foobar")
+        .execute(&mut connection);
+    let result: Value = redis::cmd("HGETALL")
+        .arg("hash")
+        .query(&mut connection)
+        .unwrap();
+
+    assert_eq!(
+        result,
+        Value::Map(vec![
+            (
+                Value::BulkString("foo".as_bytes().to_vec()),
+                Value::BulkString("baz".as_bytes().to_vec())
+            ),
+            (
+                Value::BulkString("bar".as_bytes().to_vec()),
+                Value::BulkString("foobar".as_bytes().to_vec())
+            )
+        ])
+    );
+}
+
 #[test]
 #[cfg(feature = "script")]
 fn test_cluster_script() {

diff --git a/redis/tests/test_cluster_async.rs b/redis/tests/test_cluster_async.rs
@@ -203,6 +203,52 @@ fn test_async_cluster_route_info_to_nodes() {
     .unwrap();
 }
 
+#[test]
+fn test_cluster_resp3() {
+    if !use_resp3() {
+        return;
+    }
+    block_on_all(async move {
+        let cluster = TestClusterContext::new(3, 0);
+
+        let mut connection = cluster.async_connection().await;
+
+        let _: () = redis::cmd("HSET")
+            .arg("hash")
+            .arg("foo")
+            .arg("baz")
+            .query_async(&mut connection)
+            .await
+            .unwrap();
+        let _: () = redis::cmd("HSET")
+            .arg("hash")
+            .arg("bar")
+            .arg("foobar")
+            .query_async(&mut connection)
+            .await
+            .unwrap();
+        let result: Value = redis::cmd("HGETALL")
+            .arg("hash")
+            .query_async(&mut connection)
+            .await
+            .unwrap();
+
+        assert_eq!(
+            result,
+            Value::Map(vec![
+                (
+                    Value::BulkString("foo".as_bytes().to_vec()),
+                    Value::BulkString("baz".as_bytes().to_vec())
+                ),
+                (
+                    Value::BulkString("bar".as_bytes().to_vec()),
+                    Value::BulkString("foobar".as_bytes().to_vec())
+                )
+            ])
+        );
+    });
+}
+
 #[ignore] // TODO Handle pipe where the keys do not all go to the same node
 #[test]
 fn test_async_cluster_basic_pipe() {