1 files changed, 32 insertions, 8 deletions

diff --git a/src/rpc/ring.rs b/src/rpc/ring.rs
index 85caafeb..2e997523 100644
--- a/src/rpc/ring.rs
+++ b/src/rpc/ring.rs
@@ -5,6 +5,11 @@ use serde::{Deserialize, Serialize};
 
 use garage_util::data::*;
 
+// A partition number is encoded on 16 bits,
+// i.e. we have up to 2**16 partitions.
+// (in practice we have exactly 2**PARTITION_BITS partitions)
+pub type Partition = u16;
+
 // TODO: make this constant parametrizable in the config file
 // For deployments with many nodes it might make sense to bump
 // it up to 10.
@@ -161,29 +166,48 @@ impl Ring {
 			})
 			.collect::<Vec<_>>();
 
-		eprintln!("RING: --");
-		for e in ring.iter() {
-			eprintln!("{:?}", e);
-		}
-		eprintln!("END --");
+		// eprintln!("RING: --");
+		// for e in ring.iter() {
+		// 	eprintln!("{:?}", e);
+		// }
+		// eprintln!("END --");
 
 		Self { config, ring }
 	}
 
+	pub fn partition_of(&self, from: &Hash) -> Partition {
+		let top = u16::from_be_bytes(from.as_slice()[0..2].try_into().unwrap());
+		top >> (16 - PARTITION_BITS)
+	}
+
+	pub fn partitions(&self) -> Vec<(Partition, Hash)> {
+		let mut ret = vec![];
+
+		for (i, entry) in self.ring.iter().enumerate() {
+			ret.push((i as u16, entry.location));
+		}
+		if ret.len() > 0 {
+			assert_eq!(ret[0].1, [0u8; 32].into());
+		}
+
+		ret
+	}
+
 	pub fn walk_ring(&self, from: &Hash, n: usize) -> Vec<UUID> {
 		if self.ring.len() != 1 << PARTITION_BITS {
-			warn!("Ring not yet ready, read/writes will be lost");
+			warn!("Ring not yet ready, read/writes will be lost!");
 			return vec![];
 		}
 
 		let top = u16::from_be_bytes(from.as_slice()[0..2].try_into().unwrap());
-
 		let partition_idx = (top >> (16 - PARTITION_BITS)) as usize;
+		assert_eq!(partition_idx, self.partition_of(from) as usize);
+
 		let partition = &self.ring[partition_idx];
 
 		let partition_top =
 			u16::from_be_bytes(partition.location.as_slice()[0..2].try_into().unwrap());
-		assert!(partition_top & PARTITION_MASK_U16 == top & PARTITION_MASK_U16);
+		assert_eq!(partition_top & PARTITION_MASK_U16, top & PARTITION_MASK_U16);
 
 		assert!(n <= partition.nodes.len());
 		partition.nodes[..n].iter().cloned().collect::<Vec<_>>()