1 files changed, 24 insertions, 58 deletions

diff --git a/src/rpc/layout/history.rs b/src/rpc/layout/history.rs
index e59c9e9c..9ae28887 100644
--- a/src/rpc/layout/history.rs
+++ b/src/rpc/layout/history.rs
@@ -1,5 +1,3 @@
-use std::cmp::Ordering;
-
 use garage_util::crdt::{Crdt, Lww, LwwMap};
 use garage_util::data::*;
 use garage_util::encode::nonversioned_encode;
@@ -12,14 +10,15 @@ impl LayoutHistory {
 	pub fn new(replication_factor: usize) -> Self {
 		let version = LayoutVersion::new(replication_factor);
 
-		let staging_parameters = Lww::<LayoutParameters>::new(version.parameters);
-		let empty_lwwmap = LwwMap::new();
+		let staging = LayoutStaging {
+			parameters: Lww::<LayoutParameters>::new(version.parameters),
+			roles: LwwMap::new(),
+		};
 
 		let mut ret = LayoutHistory {
 			versions: vec![version].into_boxed_slice().into(),
 			update_trackers: Default::default(),
-			staging_parameters,
-			staging_roles: empty_lwwmap,
+			staging: Lww::raw(0, staging),
 			staging_hash: [0u8; 32].into(),
 		};
 		ret.staging_hash = ret.calculate_staging_hash();
@@ -31,8 +30,7 @@ impl LayoutHistory {
 	}
 
 	pub(crate) fn calculate_staging_hash(&self) -> Hash {
-		let hashed_tuple = (&self.staging_roles, &self.staging_parameters);
-		blake2sum(&nonversioned_encode(&hashed_tuple).unwrap()[..])
+		blake2sum(&nonversioned_encode(&self.staging).unwrap()[..])
 	}
 
 	// ================== updates to layout, public interface ===================
@@ -41,26 +39,10 @@ impl LayoutHistory {
 		let mut changed = false;
 
 		// Merge staged layout changes
-		match other.current().version.cmp(&self.current().version) {
-			Ordering::Greater => {
-				self.staging_parameters = other.staging_parameters.clone();
-				self.staging_roles = other.staging_roles.clone();
-				self.staging_hash = other.staging_hash;
-				changed = true;
-			}
-			Ordering::Equal => {
-				self.staging_parameters.merge(&other.staging_parameters);
-				self.staging_roles.merge(&other.staging_roles);
-
-				let new_staging_hash = self.calculate_staging_hash();
-				if new_staging_hash != self.staging_hash {
-					changed = true;
-				}
-
-				self.staging_hash = new_staging_hash;
-			}
-			Ordering::Less => (),
+		if self.staging != other.staging {
+			changed = true;
 		}
+		self.staging.merge(&other.staging);
 
 		// Add any new versions to history
 		for v2 in other.versions.iter() {
@@ -102,50 +84,34 @@ To know the correct value of the new layout version, invoke `garage layout show`
 			}
 		}
 
+		// Compute new version and add it to history
 		let mut new_version = self.current().clone();
 		new_version.version += 1;
 
-		new_version.roles.merge(&self.staging_roles);
+		new_version.roles.merge(&self.staging.get().roles);
 		new_version.roles.retain(|(_, _, v)| v.0.is_some());
-		new_version.parameters = *self.staging_parameters.get();
-
-		self.staging_roles.clear();
-		self.staging_hash = self.calculate_staging_hash();
+		new_version.parameters = *self.staging.get().parameters.get();
 
 		let msg = new_version.calculate_partition_assignment()?;
-
 		self.versions.push(new_version);
 
+		// Reset the staged layout changes
+		self.staging.update(LayoutStaging {
+			parameters: self.staging.get().parameters.clone(),
+			roles: LwwMap::new(),
+		});
+		self.staging_hash = self.calculate_staging_hash();
+
 		Ok((self, msg))
 	}
 
-	pub fn revert_staged_changes(mut self, version: Option<u64>) -> Result<Self, Error> {
-		match version {
-			None => {
-				let error = r#"
-Please pass the new layout version number to ensure that you are writing the correct version of the cluster layout.
-To know the correct value of the new layout version, invoke `garage layout show` and review the proposed changes.
-				"#;
-				return Err(Error::Message(error.into()));
-			}
-			Some(v) => {
-				if v != self.current().version + 1 {
-					return Err(Error::Message("Invalid new layout version".into()));
-				}
-			}
-		}
-
-		self.staging_roles.clear();
-		self.staging_parameters.update(self.current().parameters);
+	pub fn revert_staged_changes(mut self) -> Result<Self, Error> {
+		self.staging.update(LayoutStaging {
+			parameters: Lww::new(self.current().parameters.clone()),
+			roles: LwwMap::new(),
+		});
 		self.staging_hash = self.calculate_staging_hash();
 
-		// TODO this is stupid, we should have a separate version counter/LWW
-		// for the staging params
-		let mut new_version = self.current().clone();
-		new_version.version += 1;
-
-		self.versions.push(new_version);
-
 		Ok(self)
 	}