1 files changed, 29 insertions, 50 deletions

diff --git a/src/garage/cli/layout.rs b/src/garage/cli/layout.rs
index e76f7737..3884bb92 100644
--- a/src/garage/cli/layout.rs
+++ b/src/garage/cli/layout.rs
@@ -1,6 +1,6 @@
 use garage_util::crdt::Crdt;
-use garage_util::data::*;
 use garage_util::error::*;
+use garage_util::formater::format_table;
 
 use garage_rpc::layout::*;
 use garage_rpc::system::*;
@@ -36,21 +36,29 @@ pub async fn cmd_assign_role(
 	args: AssignRoleOpt,
 ) -> Result<(), Error> {
 	let status = match rpc_cli
-		.call(&rpc_host, &SystemRpc::GetKnownNodes, PRIO_NORMAL)
+		.call(&rpc_host, SystemRpc::GetKnownNodes, PRIO_NORMAL)
 		.await??
 	{
 		SystemRpc::ReturnKnownNodes(nodes) => nodes,
 		resp => return Err(Error::Message(format!("Invalid RPC response: {:?}", resp))),
 	};
 
+	let mut layout = fetch_layout(rpc_cli, rpc_host).await?;
+
 	let added_nodes = args
 		.node_ids
 		.iter()
-		.map(|node_id| find_matching_node(status.iter().map(|adv| adv.id), node_id))
+		.map(|node_id| {
+			find_matching_node(
+				status
+					.iter()
+					.map(|adv| adv.id)
+					.chain(layout.node_ids().iter().cloned()),
+				node_id,
+			)
+		})
 		.collect::<Result<Vec<_>, _>>()?;
 
-	let mut layout = fetch_layout(rpc_cli, rpc_host).await?;
-
 	let mut roles = layout.roles.clone();
 	roles.merge(&layout.staging);
 
@@ -203,31 +211,9 @@ pub async fn cmd_apply_layout(
 	rpc_host: NodeID,
 	apply_opt: ApplyLayoutOpt,
 ) -> Result<(), Error> {
-	let mut layout = fetch_layout(rpc_cli, rpc_host).await?;
-
-	match apply_opt.version {
-		None => {
-			println!("Please pass the --version flag to ensure that you are writing the correct version of the cluster layout.");
-			println!("To know the correct value of the --version flag, invoke `garage layout show` and review the proposed changes.");
-			return Err(Error::Message("--version flag is missing".into()));
-		}
-		Some(v) => {
-			if v != layout.version + 1 {
-				return Err(Error::Message("Invalid value of --version flag".into()));
-			}
-		}
-	}
-
-	layout.roles.merge(&layout.staging);
-
-	if !layout.calculate_partition_assignation() {
-		return Err(Error::Message("Could not calculate new assignation of partitions to nodes. This can happen if there are less nodes than the desired number of copies of your data (see the replication_mode configuration parameter).".into()));
-	}
+	let layout = fetch_layout(rpc_cli, rpc_host).await?;
 
-	layout.staging.clear();
-	layout.staging_hash = blake2sum(&rmp_to_vec_all_named(&layout.staging).unwrap()[..]);
-
-	layout.version += 1;
+	let layout = layout.apply_staged_changes(apply_opt.version)?;
 
 	send_layout(rpc_cli, rpc_host, layout).await?;
 
@@ -242,25 +228,9 @@ pub async fn cmd_revert_layout(
 	rpc_host: NodeID,
 	revert_opt: RevertLayoutOpt,
 ) -> Result<(), Error> {
-	let mut layout = fetch_layout(rpc_cli, rpc_host).await?;
-
-	match revert_opt.version {
-		None => {
-			println!("Please pass the --version flag to ensure that you are writing the correct version of the cluster layout.");
-			println!("To know the correct value of the --version flag, invoke `garage layout show` and review the proposed changes.");
-			return Err(Error::Message("--version flag is missing".into()));
-		}
-		Some(v) => {
-			if v != layout.version + 1 {
-				return Err(Error::Message("Invalid value of --version flag".into()));
-			}
-		}
-	}
-
-	layout.staging.clear();
-	layout.staging_hash = blake2sum(&rmp_to_vec_all_named(&layout.staging).unwrap()[..]);
+	let layout = fetch_layout(rpc_cli, rpc_host).await?;
 
-	layout.version += 1;
+	let layout = layout.revert_staged_changes(revert_opt.version)?;
 
 	send_layout(rpc_cli, rpc_host, layout).await?;
 
@@ -275,7 +245,7 @@ pub async fn fetch_layout(
 	rpc_host: NodeID,
 ) -> Result<ClusterLayout, Error> {
 	match rpc_cli
-		.call(&rpc_host, &SystemRpc::PullClusterLayout, PRIO_NORMAL)
+		.call(&rpc_host, SystemRpc::PullClusterLayout, PRIO_NORMAL)
 		.await??
 	{
 		SystemRpc::AdvertiseClusterLayout(t) => Ok(t),
@@ -291,7 +261,7 @@ pub async fn send_layout(
 	rpc_cli
 		.call(
 			&rpc_host,
-			&SystemRpc::AdvertiseClusterLayout(layout),
+			SystemRpc::AdvertiseClusterLayout(layout),
 			PRIO_NORMAL,
 		)
 		.await??;
@@ -323,11 +293,20 @@ pub fn print_cluster_layout(layout: &ClusterLayout) -> bool {
 }
 
 pub fn print_staging_role_changes(layout: &ClusterLayout) -> bool {
-	if !layout.staging.items().is_empty() {
+	let has_changes = layout
+		.staging
+		.items()
+		.iter()
+		.any(|(k, _, v)| layout.roles.get(k) != Some(v));
+
+	if has_changes {
 		println!();
 		println!("==== STAGED ROLE CHANGES ====");
 		let mut table = vec!["ID\tTags\tZone\tCapacity".to_string()];
 		for (id, _, role) in layout.staging.items().iter() {
+			if layout.roles.get(id) == Some(role) {
+				continue;
+			}
 			if let Some(role) = &role.0 {
 				let tags = role.tags.join(",");
 				table.push(format!(