Merge pull request #22 from nihohit/fix-slot-map

Fix SlotMap inconsistency.

redis/src/cluster.rs

@@ -45,7 +45,7 @@ use rand::{seq::IteratorRandom, thread_rng, Rng};
 
 use crate::cluster_pipeline::UNROUTABLE_ERROR;
 use crate::cluster_routing::{MultipleNodeRoutingInfo, SingleNodeRoutingInfo, SlotAddr};
-use crate::cluster_topology::{build_slot_map, parse_slots, SlotMap, SLOT_SIZE};
+use crate::cluster_topology::{parse_slots, SlotMap, SLOT_SIZE};
 use crate::cmd::{cmd, Cmd};
 use crate::connection::{
     connect, Connection, ConnectionAddr, ConnectionInfo, ConnectionLike, RedisConnectionInfo,
@@ -143,7 +143,7 @@ where
     ) -> RedisResult<Self> {
         let connection = Self {
             connections: RefCell::new(HashMap::new()),
-            slots: RefCell::new(SlotMap::new()),
+            slots: RefCell::new(SlotMap::new(vec![])),
             auto_reconnect: RefCell::new(true),
             read_from_replicas: cluster_params.read_from_replicas,
             cluster_params,
@@ -290,18 +290,15 @@ where
         let mut rng = thread_rng();
         let len = connections.len();
         let mut samples = connections.values_mut().choose_multiple(&mut rng, len);
-        let mut new_slots = SlotMap::new();
         let mut result = Err(RedisError::from((
             ErrorKind::ResponseError,
             "Slot refresh error.",
             "didn't get any slots from server".to_string(),
         )));
         for conn in samples.iter_mut() {
             let value = conn.req_command(&slot_cmd())?;
-            match parse_slots(&value, self.cluster_params.tls)
-                .map(|v| build_slot_map(&mut new_slots, v))
-            {
-                Ok(_) => {
+            match parse_slots(&value, self.cluster_params.tls).map(SlotMap::new) {
+                Ok(new_slots) => {
                     result = Ok(new_slots);
                     break;
                 }

redis/src/cluster_topology.rs

@@ -76,14 +76,15 @@ impl SlotMapValue {
 pub(crate) struct SlotMap(BTreeMap<u16, SlotMapValue>);
 
 impl SlotMap {
-    pub fn new() -> Self {
-        Self(BTreeMap::new())
-    }
-
-    pub fn fill_slots(&mut self, slots: Vec<Slot>) {
-        for slot in slots {
-            self.0.insert(slot.end(), SlotMapValue::from_slot(slot));
-        }
+    pub fn new(slots: Vec<Slot>) -> Self {
+        let mut this = Self(BTreeMap::new());
+        this.0.extend(
+            slots
+                .into_iter()
+                .map(|slot| (slot.end(), SlotMapValue::from_slot(slot))),
+        );
+        trace!("{:?}", this);
+        this
     }
 
     pub fn slot_addr_for_route(&self, route: &Route) -> Option<&str> {
@@ -97,10 +98,6 @@ impl SlotMap {
         })
     }
 
-    pub fn clear(&mut self) {
-        self.0.clear();
-    }
-
     pub fn values(&self) -> impl Iterator<Item = &SlotAddrs> {
         self.0.values().map(|slot_value| &slot_value.addrs)
     }
@@ -234,12 +231,6 @@ pub(crate) fn parse_slots(raw_slot_resp: &Value, tls: Option<TlsMode>) -> RedisR
     Ok(result)
 }
 
-pub(crate) fn build_slot_map(slot_map: &mut SlotMap, slots_data: Vec<Slot>) {
-    slot_map.clear();
-    slot_map.fill_slots(slots_data);
-    trace!("{:?}", slot_map);
-}
-
 fn calculate_hash<T: Hash>(t: &T) -> u64 {
     let mut s = DefaultHasher::new();
     t.hash(&mut s);
@@ -315,7 +306,6 @@ pub(crate) fn calculate_topology(
     }
 
     let parse_and_built_result = |mut most_frequent_topology: TopologyView| {
-        let mut new_slots = SlotMap::new();
         most_frequent_topology.parse_and_count_slots(tls_mode);
         let slots_data = most_frequent_topology
             .slots_and_count
@@ -324,8 +314,8 @@ pub(crate) fn calculate_topology(
                 ErrorKind::ResponseError,
                 "Failed to parse the slots on the majority view",
             )))?;
-        build_slot_map(&mut new_slots, slots_data);
-        Ok(new_slots)
+
+        Ok(SlotMap::new(slots_data))
     };
 
     if non_unique_max_node_count {
@@ -529,8 +519,7 @@ mod tests {
 
     #[test]
     fn test_slot_map() {
-        let mut slot_map = SlotMap::new();
-        slot_map.fill_slots(vec![
+        let slot_map = SlotMap::new(vec![
             Slot::new(
                 1,
                 1000,
@@ -544,6 +533,7 @@ mod tests {
                 vec!["replica2:6379".to_owned()],
             ),
         ]);
+
         assert!(slot_map
             .slot_addr_for_route(&Route::new(0, SlotAddr::Master))
             .is_none());