1 files changed, 255 insertions, 20 deletions

diff --git a/src/model/k2v/rpc.rs b/src/model/k2v/rpc.rs
index f64a7984..117103b6 100644
--- a/src/model/k2v/rpc.rs
+++ b/src/model/k2v/rpc.rs
@@ -5,9 +5,10 @@
 //! node does not process the entry directly, as this would
 //! mean the vector clock gets much larger than needed).
 
-use std::collections::HashMap;
-use std::sync::Arc;
-use std::time::Duration;
+use std::collections::{BTreeMap, HashMap};
+use std::convert::TryInto;
+use std::sync::{Arc, Mutex, MutexGuard};
+use std::time::{Duration, Instant};
 
 use async_trait::async_trait;
 use futures::stream::FuturesUnordered;
@@ -15,9 +16,12 @@ use futures::StreamExt;
 use serde::{Deserialize, Serialize};
 use tokio::select;
 
+use garage_db as db;
+
 use garage_util::crdt::*;
 use garage_util::data::*;
 use garage_util::error::*;
+use garage_util::time::now_msec;
 
 use garage_rpc::system::System;
 use garage_rpc::*;
@@ -25,9 +29,15 @@ use garage_rpc::*;
 use garage_table::replication::{TableReplication, TableShardedReplication};
 use garage_table::{PartitionKey, Table};
 
+use crate::helper::error::Error as HelperError;
 use crate::k2v::causality::*;
 use crate::k2v::item_table::*;
-use crate::k2v::poll::*;
+use crate::k2v::seen::*;
+use crate::k2v::sub::*;
+
+const POLL_RANGE_EXTRA_DELAY: Duration = Duration::from_millis(200);
+
+const TIMESTAMP_KEY: &'static [u8] = b"timestamp";
 
 /// RPC messages for K2V
 #[derive(Debug, Serialize, Deserialize)]
@@ -40,7 +50,13 @@ enum K2VRpc {
 		causal_context: CausalContext,
 		timeout_msec: u64,
 	},
+	PollRange {
+		range: PollRange,
+		seen_str: Option<String>,
+		timeout_msec: u64,
+	},
 	PollItemResponse(Option<K2VItem>),
+	PollRangeResponse(Uuid, Vec<K2VItem>),
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -59,6 +75,12 @@ impl Rpc for K2VRpc {
 pub struct K2VRpcHandler {
 	system: Arc<System>,
 	item_table: Arc<Table<K2VItemTable, TableShardedReplication>>,
+
+	// Using a mutex on the local_timestamp_tree is not strictly necessary,
+	// but it helps to not try to do several inserts at the same time,
+	// which would create transaction conflicts and force many useless retries.
+	local_timestamp_tree: Mutex<db::Tree>,
+
 	endpoint: Arc<Endpoint<K2VRpc, Self>>,
 	subscriptions: Arc<SubscriptionManager>,
 }
@@ -66,14 +88,19 @@ pub struct K2VRpcHandler {
 impl K2VRpcHandler {
 	pub fn new(
 		system: Arc<System>,
+		db: &db::Db,
 		item_table: Arc<Table<K2VItemTable, TableShardedReplication>>,
 		subscriptions: Arc<SubscriptionManager>,
 	) -> Arc<Self> {
+		let local_timestamp_tree = db
+			.open_tree("k2v_local_timestamp")
+			.expect("Unable to open DB tree for k2v local timestamp");
 		let endpoint = system.netapp.endpoint("garage_model/k2v/Rpc".to_string());
 
 		let rpc_handler = Arc::new(Self {
 			system,
 			item_table,
+			local_timestamp_tree: Mutex::new(local_timestamp_tree),
 			endpoint,
 			subscriptions,
 		});
@@ -181,7 +208,7 @@ impl K2VRpcHandler {
 		Ok(())
 	}
 
-	pub async fn poll(
+	pub async fn poll_item(
 		&self,
 		bucket_id: Uuid,
 		partition_key: String,
@@ -230,9 +257,7 @@ impl K2VRpcHandler {
 						resp = Some(x);
 					}
 				}
-				K2VRpc::PollItemResponse(None) => {
-					return Ok(None);
-				}
+				K2VRpc::PollItemResponse(None) => (),
 				v => return Err(Error::unexpected_rpc_message(v)),
 			}
 		}
@@ -240,10 +265,117 @@ impl K2VRpcHandler {
 		Ok(resp)
 	}
 
+	pub async fn poll_range(
+		&self,
+		range: PollRange,
+		seen_str: Option<String>,
+		timeout_msec: u64,
+	) -> Result<Option<(BTreeMap<String, K2VItem>, String)>, HelperError> {
+		let has_seen_marker = seen_str.is_some();
+
+		// Parse seen marker, we will use it below. This is also the first check
+		// that it is valid, which returns a bad request error if not.
+		let mut seen = seen_str
+			.as_deref()
+			.map(RangeSeenMarker::decode_helper)
+			.transpose()?
+			.unwrap_or_default();
+		seen.restrict(&range);
+
+		// Prepare PollRange RPC to send to the storage nodes responsible for the parititon
+		let nodes = self
+			.item_table
+			.data
+			.replication
+			.write_nodes(&range.partition.hash());
+		let quorum = self.item_table.data.replication.read_quorum();
+		let msg = K2VRpc::PollRange {
+			range,
+			seen_str,
+			timeout_msec,
+		};
+
+		// Send the request to all nodes, use FuturesUnordered to get the responses in any order
+		let msg = msg.into_req().map_err(netapp::error::Error::from)?;
+		let rs = RequestStrategy::with_priority(PRIO_NORMAL).without_timeout();
+		let mut requests = nodes
+			.iter()
+			.map(|node| self.system.rpc.call(&self.endpoint, *node, msg.clone(), rs))
+			.collect::<FuturesUnordered<_>>();
+
+		// Fetch responses. This procedure stops fetching responses when any of the following
+		// conditions arise:
+		// - we have a response to all requests
+		// - we have a response to a read quorum of requests (e.g. 2/3), and an extra delay
+		//   has passed since the quorum was achieved
+		// - a global RPC timeout expired
+		// The extra delay after a quorum was received is usefull if the third response was to
+		// arrive during this short interval: this would allow us to consider all the data seen
+		// by that last node in the response we produce, and would likely help reduce the
+		// size of the seen marker that we will return (because we would have an info of the
+		// kind: all items produced by that node until time ts have been returned, so we can
+		// bump the entry in the global vector clock and possibly remove some item-specific
+		// vector clocks)
+		let mut deadline =
+			Instant::now() + Duration::from_millis(timeout_msec) + self.system.rpc.rpc_timeout();
+		let mut resps = vec![];
+		let mut errors = vec![];
+		loop {
+			select! {
+				_ =	tokio::time::sleep_until(deadline.into()) => {
+					break;
+				}
+				res = requests.next() => match res {
+					None => break,
+					Some(Err(e)) => errors.push(e),
+					Some(Ok(r)) => {
+						resps.push(r);
+						if resps.len() >= quorum {
+							deadline = std::cmp::min(deadline, Instant::now() + POLL_RANGE_EXTRA_DELAY);
+						}
+					}
+				}
+			}
+		}
+		if errors.len() > nodes.len() - quorum {
+			let errors = errors.iter().map(|e| format!("{}", e)).collect::<Vec<_>>();
+			return Err(Error::Quorum(quorum, resps.len(), nodes.len(), errors).into());
+		}
+
+		// Take all returned items into account to produce the response.
+		let mut new_items = BTreeMap::<String, K2VItem>::new();
+		for v in resps {
+			if let K2VRpc::PollRangeResponse(node, items) = v {
+				seen.mark_seen_node_items(node, items.iter());
+				for item in items.into_iter() {
+					match new_items.get_mut(&item.sort_key) {
+						Some(ent) => {
+							ent.merge(&item);
+						}
+						None => {
+							new_items.insert(item.sort_key.clone(), item);
+						}
+					}
+				}
+			} else {
+				return Err(Error::unexpected_rpc_message(v).into());
+			}
+		}
+
+		if new_items.is_empty() && has_seen_marker {
+			Ok(None)
+		} else {
+			Ok(Some((new_items, seen.encode()?)))
+		}
+	}
+
 	// ---- internal handlers ----
 
 	async fn handle_insert(&self, item: &InsertedItem) -> Result<K2VRpc, Error> {
-		let new = self.local_insert(item)?;
+		let new = {
+			let local_timestamp_tree = self.local_timestamp_tree.lock().unwrap();
+			self.local_insert(&local_timestamp_tree, item)?
+		};
 
 		// Propagate to rest of network
 		if let Some(updated) = new {
@@ -256,11 +388,14 @@ impl K2VRpcHandler {
 	async fn handle_insert_many(&self, items: &[InsertedItem]) -> Result<K2VRpc, Error> {
 		let mut updated_vec = vec![];
 
-		for item in items {
-			let new = self.local_insert(item)?;
+		{
+			let local_timestamp_tree = self.local_timestamp_tree.lock().unwrap();
+			for item in items {
+				let new = self.local_insert(&local_timestamp_tree, item)?;
 
-			if let Some(updated) = new {
-				updated_vec.push(updated);
+				if let Some(updated) = new {
+					updated_vec.push(updated);
+				}
 			}
 		}
 
@@ -272,10 +407,22 @@ impl K2VRpcHandler {
 		Ok(K2VRpc::Ok)
 	}
 
-	fn local_insert(&self, item: &InsertedItem) -> Result<Option<K2VItem>, Error> {
+	fn local_insert(
+		&self,
+		local_timestamp_tree: &MutexGuard<'_, db::Tree>,
+		item: &InsertedItem,
+	) -> Result<Option<K2VItem>, Error> {
+		let now = now_msec();
+
 		self.item_table
 			.data
-			.update_entry_with(&item.partition, &item.sort_key, |ent| {
+			.update_entry_with(&item.partition, &item.sort_key, |tx, ent| {
+				let old_local_timestamp = tx
+					.get(&local_timestamp_tree, TIMESTAMP_KEY)?
+					.and_then(|x| x.try_into().ok())
+					.map(u64::from_be_bytes)
+					.unwrap_or_default();
+
 				let mut ent = ent.unwrap_or_else(|| {
 					K2VItem::new(
 						item.partition.bucket_id,
@@ -283,13 +430,25 @@ impl K2VRpcHandler {
 						item.sort_key.clone(),
 					)
 				});
-				ent.update(self.system.id, &item.causal_context, item.value.clone());
-				ent
+				let new_local_timestamp = ent.update(
+					self.system.id,
+					&item.causal_context,
+					item.value.clone(),
+					std::cmp::max(old_local_timestamp, now),
+				);
+
+				tx.insert(
+					&local_timestamp_tree,
+					TIMESTAMP_KEY,
+					u64::to_be_bytes(new_local_timestamp),
+				)?;
+
+				Ok(ent)
 			})
 	}
 
-	async fn handle_poll(&self, key: &PollKey, ct: &CausalContext) -> Result<K2VItem, Error> {
-		let mut chan = self.subscriptions.subscribe(key);
+	async fn handle_poll_item(&self, key: &PollKey, ct: &CausalContext) -> Result<K2VItem, Error> {
+		let mut chan = self.subscriptions.subscribe_item(key);
 
 		let mut value = self
 			.item_table
@@ -311,6 +470,71 @@ impl K2VRpcHandler {
 
 		Ok(value)
 	}
+
+	async fn handle_poll_range(
+		&self,
+		range: &PollRange,
+		seen_str: &Option<String>,
+	) -> Result<Vec<K2VItem>, Error> {
+		if let Some(seen_str) = seen_str {
+			let seen = RangeSeenMarker::decode(seen_str).ok_or_message("Invalid seenMarker")?;
+
+			// Subscribe now to all changes on that partition,
+			// so that new items that are inserted while we are reading the range
+			// will be seen in the loop below
+			let mut chan = self.subscriptions.subscribe_partition(&range.partition);
+
+			// Check for the presence of any new items already stored in the item table
+			let mut new_items = self.poll_range_read_range(range, &seen)?;
+
+			// If we found no new items, wait for a matching item to arrive
+			// on the channel
+			while new_items.is_empty() {
+				let item = chan.recv().await?;
+				if range.matches(&item) && seen.is_new_item(&item) {
+					new_items.push(item);
+				}
+			}
+
+			Ok(new_items)
+		} else {
+			// If no seen marker was specified, we do not poll for anything.
+			// We return immediately with the set of known items (even if
+			// it is empty), which will give the client an inital view of
+			// the dataset and an initial seen marker for further
+			// PollRange calls.
+			self.poll_range_read_range(range, &RangeSeenMarker::default())
+		}
+	}
+
+	fn poll_range_read_range(
+		&self,
+		range: &PollRange,
+		seen: &RangeSeenMarker,
+	) -> Result<Vec<K2VItem>, Error> {
+		let mut new_items = vec![];
+
+		let partition_hash = range.partition.hash();
+		let first_key = match &range.start {
+			None => partition_hash.to_vec(),
+			Some(sk) => self.item_table.data.tree_key(&range.partition, sk),
+		};
+		for item in self.item_table.data.store.range(first_key..)? {
+			let (key, value) = item?;
+			if &key[..32] != partition_hash.as_slice() {
+				break;
+			}
+			let item = self.item_table.data.decode_entry(&value)?;
+			if !range.matches(&item) {
+				break;
+			}
+			if seen.is_new_item(&item) {
+				new_items.push(item);
+			}
+		}
+
+		Ok(new_items)
+	}
 }
 
 #[async_trait]
@@ -326,10 +550,21 @@ impl EndpointHandler<K2VRpc> for K2VRpcHandler {
 			} => {
 				let delay = tokio::time::sleep(Duration::from_millis(*timeout_msec));
 				select! {
-					ret = self.handle_poll(key, causal_context) => ret.map(Some).map(K2VRpc::PollItemResponse),
+					ret = self.handle_poll_item(key, causal_context) => ret.map(Some).map(K2VRpc::PollItemResponse),
 					_ = delay => Ok(K2VRpc::PollItemResponse(None)),
 				}
 			}
+			K2VRpc::PollRange {
+				range,
+				seen_str,
+				timeout_msec,
+			} => {
+				let delay = tokio::time::sleep(Duration::from_millis(*timeout_msec));
+				select! {
+					ret = self.handle_poll_range(range, seen_str) => ret.map(|items| K2VRpc::PollRangeResponse(self.system.id, items)),
+					_ = delay => Ok(K2VRpc::PollRangeResponse(self.system.id, vec![])),
+				}
+			}
 			m => Err(Error::unexpected_rpc_message(m)),
 		}
 	}