Merge pull request #1997 from RolandSherwin/local_metrics

feat(manager): optionally set the different ports when running local network

sn_node_manager/src/add_services/config.rs

@@ -23,20 +23,46 @@ pub enum PortRange {
     Range(u16, u16),
 }
 
-pub fn parse_port_range(s: &str) -> Result<PortRange> {
-    if let Ok(port) = u16::from_str(s) {
-        Ok(PortRange::Single(port))
-    } else {
-        let parts: Vec<&str> = s.split('-').collect();
-        if parts.len() != 2 {
-            return Err(eyre!("Port range must be in the format 'start-end'"));
+impl PortRange {
+    pub fn parse(s: &str) -> Result<Self> {
+        if let Ok(port) = u16::from_str(s) {
+            Ok(Self::Single(port))
+        } else {
+            let parts: Vec<&str> = s.split('-').collect();
+            if parts.len() != 2 {
+                return Err(eyre!("Port range must be in the format 'start-end'"));
+            }
+            let start = parts[0].parse::<u16>()?;
+            let end = parts[1].parse::<u16>()?;
+            if start >= end {
+                return Err(eyre!("End port must be greater than start port"));
+            }
+            Ok(Self::Range(start, end))
         }
-        let start = parts[0].parse::<u16>()?;
-        let end = parts[1].parse::<u16>()?;
-        if start >= end {
-            return Err(eyre!("End port must be greater than start port"));
+    }
+
+    /// Validate the port range against a count to make sure the correct number of ports are provided.
+    pub fn validate(&self, count: u16) -> Result<()> {
+        match self {
+            Self::Single(_) => {
+                if count != 1 {
+                    error!("The count ({count}) does not match the number of ports (1)");
+                    return Err(eyre!(
+                        "The count ({count}) does not match the number of ports (1)"
+                    ));
+                }
+            }
+            Self::Range(start, end) => {
+                let port_count = end - start + 1;
+                if count != port_count {
+                    error!("The count ({count}) does not match the number of ports ({port_count})");
+                    return Err(eyre!(
+                        "The count ({count}) does not match the number of ports ({port_count})"
+                    ));
+                }
+            }
         }
-        Ok(PortRange::Range(start, end))
+        Ok(())
     }
 }
 

sn_node_manager/src/add_services/mod.rs

@@ -12,10 +12,11 @@ mod tests;
 use self::config::{
     AddAuditorServiceOptions, AddDaemonServiceOptions, AddFaucetServiceOptions,
     AddNodeServiceOptions, InstallAuditorServiceCtxBuilder, InstallFaucetServiceCtxBuilder,
-    InstallNodeServiceCtxBuilder, PortRange,
+    InstallNodeServiceCtxBuilder,
 };
 use crate::{
     config::{create_owned_dir, get_user_safenode_data_dir},
+    helpers::{check_port_availability, get_start_port_if_applicable, increment_port_option},
     VerbosityLevel, DAEMON_SERVICE_NAME,
 };
 use color_eyre::{
@@ -62,39 +63,18 @@ pub async fn add_node(
         }
     }
 
-    if let Some(ref port_range) = options.node_port {
-        match port_range {
-            PortRange::Single(_) => {
-                let count = options.count.unwrap_or(1);
-                if count != 1 {
-                    error!("The number of services to add ({count}) does not match the number of ports (1)");
-                    return Err(eyre!(
-                        "The number of services to add ({count}) does not match the number of ports (1)"
-                    ));
-                }
-            }
-            PortRange::Range(start, end) => {
-                let port_count = end - start + 1;
-                let service_count = options.count.unwrap_or(1);
-                if port_count != service_count {
-                    error!("The number of services to add ({service_count}) does not match the number of ports ({port_count})");
-                    return Err(eyre!(
-                        "The number of services to add ({service_count}) does not match the number of ports ({port_count})"
-                    ));
-                }
-            }
-        }
-    }
-
     if let Some(port_option) = &options.node_port {
+        port_option.validate(options.count.unwrap_or(1))?;
         check_port_availability(port_option, &node_registry.nodes)?;
     }
 
     if let Some(port_option) = &options.metrics_port {
+        port_option.validate(options.count.unwrap_or(1))?;
         check_port_availability(port_option, &node_registry.nodes)?;
     }
 
     if let Some(port_option) = &options.rpc_port {
+        port_option.validate(options.count.unwrap_or(1))?;
         check_port_availability(port_option, &node_registry.nodes)?;
     }
 
@@ -575,52 +555,3 @@ pub fn add_faucet(
         }
     }
 }
-
-fn get_start_port_if_applicable(range: Option<PortRange>) -> Option<u16> {
-    if let Some(port) = range {
-        match port {
-            PortRange::Single(val) => return Some(val),
-            PortRange::Range(start, _) => return Some(start),
-        }
-    }
-    None
-}
-
-fn increment_port_option(port: Option<u16>) -> Option<u16> {
-    if let Some(port) = port {
-        let incremented_port = port + 1;
-        return Some(incremented_port);
-    }
-    None
-}
-
-fn check_port_availability(port_option: &PortRange, nodes: &[NodeServiceData]) -> Result<()> {
-    let mut all_ports = Vec::new();
-    for node in nodes {
-        if let Some(port) = node.metrics_port {
-            all_ports.push(port);
-        }
-        if let Some(port) = node.node_port {
-            all_ports.push(port);
-        }
-        all_ports.push(node.rpc_socket_addr.port());
-    }
-
-    match port_option {
-        PortRange::Single(port) => {
-            if all_ports.iter().any(|p| *p == *port) {
-                error!("Port {port} is being used by another service");
-                return Err(eyre!("Port {port} is being used by another service"));
-            }
-        }
-        PortRange::Range(start, end) => {
-            for i in *start..=*end {
-                if all_ports.iter().any(|p| *p == i) {
-                    error!("Port {i} is being used by another service");
-                    return Err(eyre!("Port {i} is being used by another service"));
-                }
-            }
-        }
-    }
-    Ok(())
-}

sn_node_manager/src/add_services/tests.rs

@@ -1568,7 +1568,7 @@ async fn add_node_should_return_an_error_if_port_and_node_count_do_not_match() -
         Ok(_) => panic!("This test should result in an error"),
         Err(e) => {
             assert_eq!(
-                format!("The number of services to add (2) does not match the number of ports (3)"),
+                format!("The count (2) does not match the number of ports (3)"),
                 e.to_string()
             )
         }
@@ -1639,7 +1639,7 @@ async fn add_node_should_return_an_error_if_multiple_services_are_specified_with
         Ok(_) => panic!("This test should result in an error"),
         Err(e) => {
             assert_eq!(
-                format!("The number of services to add (2) does not match the number of ports (1)"),
+                format!("The count (2) does not match the number of ports (1)"),
                 e.to_string()
             )
         }
@@ -2514,7 +2514,7 @@ async fn add_node_should_return_an_error_if_duplicate_custom_rpc_port_in_range_i
             auto_restart: false,
             auto_set_nat_flags: false,
             bootstrap_peers: vec![],
-            count: None,
+            count: Some(2),
             delete_safenode_src: true,
             enable_metrics_server: false,
             env_variables: None,