Rack Awareness, Sequence Policies, More complex policies and other Fixes

Cargo.toml

@@ -28,7 +28,7 @@ aerospike-sync = {path = "./aerospike-sync", optional = true}
 aerospike-macro = {path = "./aerospike-macro", optional = true}
 
 [features]
-default = ["async", "serialization"]
+default = ["async", "serialization", "rt-tokio"]
 serialization = ["aerospike-core/serialization"]
 async = ["aerospike-core"]
 sync = ["aerospike-sync"]
@@ -50,6 +50,7 @@ bencher = "0.1"
 serde_json = "1.0"
 rand = "0.7"
 lazy_static = "1.4"
+ripemd = "0.1"
 aerospike-macro = {path = "./aerospike-macro"}
 aerospike-rt = {path = "./aerospike-rt"}
 futures = {version = "0.3.16" }

aerospike-core/Cargo.toml

@@ -8,13 +8,13 @@ edition = "2018"
 [dependencies]
 log = "0.4"
 byteorder = "1.3"
-ripemd160 = "0.8"
-base64 = "0.11"
-crossbeam-queue = "0.2"
-rand = "0.7"
+ripemd = "0.1"
+base64 = "0.13"
+crossbeam-queue = "0.3"
+rand = "0.8"
 lazy_static = "1.4"
 error-chain = "0.12"
-pwhash = "0.3"
+pwhash = "1.0"
 serde = { version = "1.0", features = ["derive"], optional = true }
 aerospike-rt = {path = "../aerospike-rt"}
 futures = {version = "0.3.16" }
@@ -26,7 +26,7 @@ rt-tokio = ["aerospike-rt/rt-tokio"]
 rt-async-std = ["aerospike-rt/rt-async-std"]
 
 [dev-dependencies]
-env_logger = "0.9.3"
+env_logger = "0.9"
 hex = "0.4"
 bencher = "0.1"
 serde_json = "1.0"

aerospike-core/src/batch/batch_executor.rs

@@ -13,18 +13,16 @@
 // License for the specific language governing permissions and limitations under
 // the License.
 
-use std::cmp;
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 use crate::batch::BatchRead;
 use crate::cluster::partition::Partition;
 use crate::cluster::{Cluster, Node};
 use crate::commands::BatchReadCommand;
-use crate::errors::{Error, Result};
+use crate::errors::Result;
 use crate::policy::{BatchPolicy, Concurrency};
 use crate::Key;
-use futures::lock::Mutex;
 
 pub struct BatchExecutor {
     cluster: Arc<Cluster>,
@@ -40,80 +38,47 @@ impl BatchExecutor {
         policy: &BatchPolicy,
         batch_reads: Vec<BatchRead>,
     ) -> Result<Vec<BatchRead>> {
-        let mut batch_nodes = self.get_batch_nodes(&batch_reads).await?;
+        let batch_nodes = self.get_batch_nodes(&batch_reads, policy.replica).await?;
         let jobs = batch_nodes
-            .drain()
+            .into_iter()
             .map(|(node, reads)| BatchReadCommand::new(policy, node, reads))
             .collect();
-        let reads = self.execute_batch_jobs(jobs, &policy.concurrency).await?;
-        let mut res: Vec<BatchRead> = vec![];
-        for mut read in reads {
-            res.append(&mut read.batch_reads);
-        }
-        Ok(res)
+        let reads = self.execute_batch_jobs(jobs, policy.concurrency).await?;
+        let mut all_results: Vec<_> = reads.into_iter().flat_map(|cmd|cmd.batch_reads).collect();
+        all_results.sort_by_key(|(_, i)|*i);
+        Ok(all_results.into_iter().map(|(b, _)|b).collect())
     }
 
     async fn execute_batch_jobs(
         &self,
         jobs: Vec<BatchReadCommand>,
-        concurrency: &Concurrency,
+        concurrency: Concurrency,
     ) -> Result<Vec<BatchReadCommand>> {
-        let threads = match *concurrency {
-            Concurrency::Sequential => 1,
-            Concurrency::Parallel => jobs.len(),
-            Concurrency::MaxThreads(max) => cmp::min(max, jobs.len()),
-        };
-        let size = jobs.len() / threads;
-        let mut overhead = jobs.len() % threads;
-        let last_err: Arc<Mutex<Option<Error>>> = Arc::default();
-        let mut slice_index = 0;
-        let mut handles = vec![];
-        let res = Arc::new(Mutex::new(vec![]));
-        for _ in 0..threads {
-            let mut thread_size = size;
-            if overhead >= 1 {
-                thread_size += 1;
-                overhead -= 1;
-            }
-            let slice = Vec::from(&jobs[slice_index..slice_index + thread_size]);
-            slice_index = thread_size + 1;
-            let last_err = last_err.clone();
-            let res = res.clone();
-            let handle = aerospike_rt::spawn(async move {
-                //let next_job = async { jobs.lock().await.next().await};
-                for mut cmd in slice {
-                    if let Err(err) = cmd.execute().await {
-                        *last_err.lock().await = Some(err);
-                    };
-                    res.lock().await.push(cmd);
-                }
-            });
-            handles.push(handle);
-        }
-        futures::future::join_all(handles).await;
-        match Arc::try_unwrap(last_err).unwrap().into_inner() {
-            None => Ok(res.lock().await.to_vec()),
-            Some(err) => Err(err),
+        let handles = jobs.into_iter().map(|job|job.execute(self.cluster.clone()));
+        match concurrency {
+            Concurrency::Sequential => futures::future::join_all(handles).await.into_iter().collect(),
+            Concurrency::Parallel => futures::future::join_all(handles.map(aerospike_rt::spawn)).await.into_iter().map(|value|value.map_err(|e|e.to_string())?).collect(),
         }
     }
 
     async fn get_batch_nodes(
         &self,
         batch_reads: &[BatchRead],
-    ) -> Result<HashMap<Arc<Node>, Vec<BatchRead>>> {
+        replica: crate::policy::Replica,
+    ) -> Result<HashMap<Arc<Node>, Vec<(BatchRead, usize)>>> {
         let mut map = HashMap::new();
-        for (_, batch_read) in batch_reads.iter().enumerate() {
-            let node = self.node_for_key(&batch_read.key).await?;
+        for (index, batch_read) in batch_reads.iter().enumerate() {
+            let node = self.node_for_key(&batch_read.key, replica).await?;
             map.entry(node)
                 .or_insert_with(Vec::new)
-                .push(batch_read.clone());
+                .push((batch_read.clone(), index));
         }
         Ok(map)
     }
 
-    async fn node_for_key(&self, key: &Key) -> Result<Arc<Node>> {
+    async fn node_for_key(&self, key: &Key, replica: crate::policy::Replica) -> Result<Arc<Node>> {
         let partition = Partition::new_by_key(key);
-        let node = self.cluster.get_node(&partition).await?;
+        let node = self.cluster.get_node(&partition, replica, Weak::new()).await?;
         Ok(node)
     }
 }

aerospike-core/src/client.rs

@@ -184,7 +184,7 @@ impl Client {
         T: Into<Bins> + Send + Sync + 'static,
     {
         let bins = bins.into();
-        let mut command = ReadCommand::new(policy, self.cluster.clone(), key, bins);
+        let mut command = ReadCommand::new(&policy.base_policy, self.cluster.clone(), key, bins, policy.replica);
         command.execute().await?;
         Ok(command.record.unwrap())
     }

aerospike-core/src/cluster/mod.rs

@@ -21,7 +21,7 @@ pub mod partition_tokenizer;
 use aerospike_rt::time::{Duration, Instant};
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::vec::Vec;
 
 pub use self::node::Node;
@@ -30,6 +30,7 @@ use self::node_validator::NodeValidator;
 use self::partition::Partition;
 use self::partition_tokenizer::PartitionTokenizer;
 
+use crate::commands::Message;
 use crate::errors::{ErrorKind, Result};
 use crate::net::Host;
 use crate::policy::ClientPolicy;
@@ -38,6 +39,65 @@ use futures::channel::mpsc;
 use futures::channel::mpsc::{Receiver, Sender};
 use futures::lock::Mutex;
 
+#[derive(Debug)]
+pub struct PartitionForNamespace {
+    nodes: Vec<(u32, Option<Arc<Node>>)>,
+    replicas: usize,
+}
+type PartitionTable = HashMap<String, PartitionForNamespace>;
+
+impl Default for PartitionForNamespace {
+    fn default() -> Self {
+        Self { nodes: Vec::default(), replicas: 0 }
+    }
+}
+
+impl PartitionForNamespace {
+    fn all_replicas(&self, index: usize) -> impl Iterator<Item = Option<Arc<Node>>> + '_ {
+        (0..self.replicas).map(move |i|self.nodes.get(i * node::PARTITIONS + index).and_then(|(_, item)|item.clone()))
+    }
+
+    async fn get_node(&self, cluster: &Cluster, partition: &Partition<'_>, replica: crate::policy::Replica, last_tried: Weak<Node>) -> Result<Arc<Node>> {
+        fn get_next_in_sequence<I: Iterator<Item = Arc<Node>>, F: Fn()->I>(get_sequence: F, last_tried: Weak<Node>) -> Option<Arc<Node>> {
+            if let Some(last_tried) = last_tried.upgrade() {
+                // If this isn't the first attempt, try the replica immediately after in sequence (that is actually valid)
+                let mut replicas = get_sequence();
+                while let Some(replica) = replicas.next() {
+                    if Arc::ptr_eq(&replica, &last_tried) {
+                        if let Some(in_sequence_after) = replicas.next() {
+                            return Some(in_sequence_after)
+                        }
+
+                        // No more after this? Drop through to try from the beginning.
+                        break;
+                    }
+                }
+            }
+            // If we get here, we're on the first attempt, the last node is already gone, or there are no more nodes in sequence. Just find the next populated option.
+            get_sequence().next()
+        }
+
+
+        let node = match replica {
+            crate::policy::Replica::Master => self.all_replicas(partition.partition_id).next().flatten(),
+            crate::policy::Replica::Sequence => {
+                get_next_in_sequence(||self.all_replicas(partition.partition_id).flatten(), last_tried)
+            },
+            crate::policy::Replica::PreferRack => {
+                let rack_ids = cluster.client_policy.rack_ids.as_ref().ok_or_else(||"Attempted to use Replica::PreferRack without configuring racks in client policy".to_string())?;
+                get_next_in_sequence(||
+                    self
+                    .all_replicas(partition.partition_id)
+                    .flatten()
+                    .filter(|node|node.is_in_rack(partition.namespace, rack_ids)), last_tried.clone())
+                .or_else(||get_next_in_sequence(||self.all_replicas(partition.partition_id).flatten(), last_tried))
+            },
+        };
+
+        node.ok_or_else(||format!("Cannot get appropriate node for namespace: {} partition: {}", partition.namespace, partition.partition_id).into())
+    }
+}
+
 // Cluster encapsulates the aerospike cluster nodes and manages
 // them.
 #[derive(Debug)]
@@ -51,8 +111,8 @@ pub struct Cluster {
     // Active nodes in cluster.
     nodes: Arc<RwLock<Vec<Arc<Node>>>>,
 
-    // Hints for best node for a partition
-    partition_write_map: Arc<RwLock<HashMap<String, Vec<Arc<Node>>>>>,
+    // Which partition contains the key.
+    partition_write_map: RwLock<PartitionTable>,
 
     // Random node index.
     node_index: AtomicIsize,
@@ -73,7 +133,7 @@ impl Cluster {
             aliases: Arc::new(RwLock::new(HashMap::new())),
             nodes: Arc::new(RwLock::new(vec![])),
 
-            partition_write_map: Arc::new(RwLock::new(HashMap::new())),
+            partition_write_map: RwLock::new(HashMap::default()),
             node_index: AtomicIsize::new(0),
 
             tend_channel: Mutex::new(tx),
@@ -137,6 +197,7 @@ impl Cluster {
         // Refresh all known nodes.
         for node in nodes {
             let old_gen = node.partition_generation();
+            let old_rebalance_gen = node.rebalance_generation();
             if node.is_active() {
                 match node.refresh(self.aliases().await).await {
                     Ok(friends) => {
@@ -147,7 +208,11 @@ impl Cluster {
                         }
 
                         if old_gen != node.partition_generation() {
-                            self.update_partitions(node.clone()).await?;
+                            self.update_partitions(&node).await?;
+                        }
+
+                        if old_rebalance_gen != node.rebalance_generation() {
+                            self.update_rack_ids(&node).await?;
                         }
                     }
                     Err(err) => {
@@ -231,23 +296,14 @@ impl Cluster {
         Ok(aliases.contains_key(host))
     }
 
-    async fn set_partitions(&self, partitions: HashMap<String, Vec<Arc<Node>>>) {
-        let mut partition_map = self.partition_write_map.write().await;
-        *partition_map = partitions;
-    }
-
-    fn partitions(&self) -> Arc<RwLock<HashMap<String, Vec<Arc<Node>>>>> {
-        self.partition_write_map.clone()
-    }
 
     pub async fn node_partitions(&self, node: &Node, namespace: &str) -> Vec<u16> {
         let mut res: Vec<u16> = vec![];
-        let partitions = self.partitions();
-        let partitions = partitions.read().await;
+        let partitions = self.partition_write_map.read().await;
 
         if let Some(node_array) = partitions.get(namespace) {
-            for (i, tnode) in node_array.iter().enumerate() {
-                if node == tnode.as_ref() {
+            for (i, (_, tnode)) in node_array.nodes.iter().enumerate().take(node::PARTITIONS) {
+                if tnode.as_ref().map_or(false, |tnode|tnode.as_ref() == node) {
                     res.push(i as u16);
                 }
             }
@@ -256,15 +312,29 @@ impl Cluster {
         res
     }
 
-    pub async fn update_partitions(&self, node: Arc<Node>) -> Result<()> {
+    pub async fn update_partitions(&self, node: &Arc<Node>) -> Result<()> {
         let mut conn = node.get_connection().await?;
-        let tokens = PartitionTokenizer::new(&mut conn).await.map_err(|e| {
+        let tokens = PartitionTokenizer::new(&mut conn, node).await.map_err(|e| {
             conn.invalidate();
             e
         })?;
 
-        let nmap = tokens.update_partition(self.partitions(), node).await?;
-        self.set_partitions(nmap).await;
+        let mut partitions = self.partition_write_map.write().await;
+        tokens.update_partition(&mut partitions, node)?;
+
+        Ok(())
+    }
+
+    pub async fn update_rack_ids(&self, node: &Arc<Node>) -> Result<()> {
+        const RACK_IDS: &str = "rack-ids";
+        let mut conn = node.get_connection().await?;
+        let info_map = Message::info(&mut conn, &[RACK_IDS, node::REBALANCE_GENERATION]).await?;
+        if let Some(buf) = info_map.get(RACK_IDS) {
+            node.parse_rack(buf.as_str())?;
+        }
+
+        // We re-update the rebalance generation right now (in case its changed since it was last polled)
+        node.update_rebalance_generation(&info_map)?;
 
         Ok(())
     }
@@ -440,10 +510,11 @@ impl Cluster {
     }
 
     async fn find_node_in_partition_map(&self, filter: Arc<Node>) -> bool {
+        let filter = Some(filter);
         let partitions = self.partition_write_map.read().await;
         (*partitions)
             .values()
-            .any(|map| map.iter().any(|node| *node == filter))
+            .any(|map| map.nodes.iter().any(|(_, node)| *node == filter))
     }
 
     async fn add_nodes(&self, friend_list: &[Arc<Node>]) {
@@ -492,17 +563,14 @@ impl Cluster {
         *nodes = new_nodes;
     }
 
-    pub async fn get_node(&self, partition: &Partition<'_>) -> Result<Arc<Node>> {
-        let partitions = self.partitions();
-        let partitions = partitions.read().await;
-
-        if let Some(node_array) = partitions.get(partition.namespace) {
-            if let Some(node) = node_array.get(partition.partition_id) {
-                return Ok(node.clone());
-            }
-        }
+    pub async fn get_node(&self, partition: &Partition<'_>, replica: crate::policy::Replica, last_tried: Weak<Node>) -> Result<Arc<Node>> {
+        let partitions = self.partition_write_map.read().await;
 
-        self.get_random_node().await
+        let namespace = partitions
+            .get(partition.namespace)
+            .ok_or_else(||format!("Cannot get appropriate node for namespace: {}", partition.namespace))?;
+
+        namespace.get_node(self, partition, replica, last_tried).await
     }
 
     pub async fn get_random_node(&self) -> Result<Arc<Node>> {