From cb8fe1614cc40a0d44bfc04a1a0a3321e7f3f075 Mon Sep 17 00:00:00 2001
From: Arpad Borsos <arpad.borsos@sentry.io>
Date: Thu, 11 Apr 2024 12:16:55 +0200
Subject: [PATCH] Avoid intermediate `HashMap` in `gets`

This uses a `Vec` instead, plus iterator chaining, and avoids a `Clone` of the `Connection`.
---
 src/client.rs | 40 ++++++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 446b86d..f8e63ed 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -136,6 +136,26 @@ impl Client {
         &self.connections[(self.hash_function)(key) as usize % connections_count]
     }
 
+    /// Distributes the input `keys` to the available `connections`.
+    ///
+    /// This uses the `hash_function` internally, and the returned [`Vec`] matches
+    /// the available `connections`.
+    fn distribute_keys<'a>(&self, keys: &[&'a str]) -> Result<Vec<Vec<&'a str>>, MemcacheError> {
+        for key in keys {
+            check_key_len(key)?;
+        }
+
+        let connections_count = self.connections.len();
+        let mut con_keys = Vec::new();
+        con_keys.resize_with(connections_count, Vec::new);
+        for key in keys {
+            let connection_index = (self.hash_function)(key) as usize % connections_count;
+            con_keys[connection_index].push(*key);
+        }
+
+        Ok(con_keys)
+    }
+
     /// Set the socket read timeout for TCP connections.
     ///
     /// Example:
@@ -247,21 +267,13 @@ impl Client {
     /// assert_eq!(result["foo"], "42");
     /// ```
     pub fn gets<V: FromMemcacheValueExt>(&self, keys: &[&str]) -> Result<HashMap<String, V>, MemcacheError> {
-        for key in keys {
-            check_key_len(key)?;
-        }
-        let mut con_keys: HashMap<usize, Vec<&str>> = HashMap::new();
-        let mut result: HashMap<String, V> = HashMap::new();
-        let connections_count = self.connections.len();
+        let distributed_keys = self.distribute_keys(keys)?;
 
-        for key in keys {
-            let connection_index = (self.hash_function)(key) as usize % connections_count;
-            let array = con_keys.entry(connection_index).or_default();
-            array.push(key);
-        }
-        for (&connection_index, keys) in con_keys.iter() {
-            let connection = self.connections[connection_index].clone();
-            result.extend(connection.get()?.gets(keys)?);
+        let mut result: HashMap<String, V> = HashMap::new();
+        for (connection, keys) in self.connections.iter().zip(distributed_keys) {
+            if !keys.is_empty() {
+                result.extend(connection.get()?.gets(&keys)?);
+            }
         }
         Ok(result)
     }