feat: allow low peers outside mainnet

ultrasoundmoney · Dec 1, 2023 · 334c3ef · 334c3ef
1 parent 7b45514
commit 334c3ef
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 54 deletions.
diff --git a/.env.example b/.env.example
@@ -2,4 +2,5 @@
 BEACON_URL=http://54.38.38.164:5052
 BIND_PUBLIC_INTERFACE=false
 GETH_URL=http://54.38.38.164:8545
+NETWORK=mainnet
 RUST_LOG=node_health=debug
diff --git a/src/env.rs b/src/env.rs
@@ -1,10 +1,10 @@
 //! Fns to read variables from the environment more conveniently and help other functions figure
 //! out what environment they're running in.
 
-use std::env;
+use std::{env, fmt};
 
 use lazy_static::lazy_static;
-use tracing::debug;
+use tracing::{debug, warn};
 
 const SECRET_LOG_BLACKLIST: [&str; 0] = [];
 
@@ -57,35 +57,35 @@ pub fn get_env_bool(key: &str) -> Option<bool> {
     })
 }
 
-// #[derive(Debug, Clone, PartialEq, Eq)]
-// pub enum Network {
-//     Mainnet,
-//     Goerli,
-// }
-
-// impl Display for Network {
-//     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-//         match self {
-//             Network::Mainnet => write!(f, "mainnet"),
-//             Network::Goerli => write!(f, "goerli"),
-//         }
-//     }
-// }
-
-// pub fn get_network() -> Network {
-//     let network_str = get_env_var("NETWORK");
-//     match network_str {
-//         None => {
-//             warn!("no NETWORK in env, assuming Mainnet");
-//             Network::Mainnet
-//         }
-//         Some(str) => match str.to_lowercase().as_ref() {
-//             "mainnet" => Network::Mainnet,
-//             "goerli" => Network::Goerli,
-//             _ => panic!("NETWORK present: {str}, but not one of [mainnet, goerli], panicking!"),
-//         },
-//     }
-// }
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum Network {
+    Mainnet,
+    Goerli,
+}
+
+impl fmt::Display for Network {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Network::Mainnet => write!(f, "mainnet"),
+            Network::Goerli => write!(f, "goerli"),
+        }
+    }
+}
+
+pub fn get_network() -> Network {
+    let network_str = get_env_var("NETWORK");
+    match network_str {
+        None => {
+            warn!("no NETWORK in env, assuming Mainnet");
+            Network::Mainnet
+        }
+        Some(str) => match str.to_lowercase().as_ref() {
+            "mainnet" => Network::Mainnet,
+            "goerli" => Network::Goerli,
+            _ => panic!("NETWORK present: {str}, but not one of [mainnet, goerli], panicking!"),
+        },
+    }
+}
 
 #[derive(Debug, Clone)]
 pub struct EnvConfig {
@@ -94,7 +94,7 @@ pub struct EnvConfig {
     pub geth_url: String,
     pub log_json: bool,
     pub log_perf: bool,
-    // pub network: Network,
+    pub network: Network,
 }
 
 fn get_env_config() -> EnvConfig {
@@ -104,7 +104,7 @@ fn get_env_config() -> EnvConfig {
         geth_url: get_env_var("GETH_URL").expect("GETH_URL not set"),
         log_json: get_env_bool("LOG_JSON").unwrap_or(false),
         log_perf: get_env_bool("LOG_PERF").unwrap_or(false),
-        // network: get_network(),
+        network: get_network(),
     }
 }
 
@@ -186,28 +186,28 @@ mod tests {
         );
     }
 
-    // #[test]
-    // fn test_get_network() {
-    //     std::env::set_var("NETWORK", "mainnet");
-    //     assert_eq!(get_network(), Network::Mainnet);
+    #[test]
+    fn test_get_network() {
+        std::env::set_var("NETWORK", "mainnet");
+        assert_eq!(get_network(), Network::Mainnet);
 
-    //     std::env::set_var("NETWORK", "goerli");
-    //     assert_eq!(get_network(), Network::Goerli);
+        std::env::set_var("NETWORK", "goerli");
+        assert_eq!(get_network(), Network::Goerli);
 
-    //     std::env::set_var("NETWORK", "Mainnet");
-    //     assert_eq!(get_network(), Network::Mainnet);
+        std::env::set_var("NETWORK", "Mainnet");
+        assert_eq!(get_network(), Network::Mainnet);
 
-    //     std::env::set_var("NETWORK", "Goerli");
-    //     assert_eq!(get_network(), Network::Goerli);
+        std::env::set_var("NETWORK", "Goerli");
+        assert_eq!(get_network(), Network::Goerli);
 
-    //     std::env::remove_var("NETWORK");
-    //     assert_eq!(get_network(), Network::Mainnet);
-    // }
+        std::env::remove_var("NETWORK");
+        assert_eq!(get_network(), Network::Mainnet);
+    }
 
-    // #[test]
-    // #[ignore = "this test breaks NETWORK for parallel tests"]
-    // fn test_get_network_panics() {
-    //     std::env::set_var("NETWORK", "invalid_network");
-    //     get_network();
-    // }
+    #[test]
+    #[ignore = "this test breaks NETWORK for parallel tests"]
+    fn test_get_network_panics() {
+        std::env::set_var("NETWORK", "invalid_network");
+        get_network();
+    }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -5,7 +5,10 @@ use std::{
     time::{Duration, SystemTime},
 };
 
-use node_health::{geth, log};
+use node_health::{
+    env::{Network, ENV_CONFIG},
+    geth, log,
+};
 use tokio::{spawn, sync::Notify, time::sleep};
 use tracing::{debug, info};
 
@@ -65,7 +68,12 @@ async fn main() -> anyhow::Result<()> {
         }
 
         let geth_peer_count = geth::peer_count(&geth_client).await?;
-        if geth_peer_count < 10 {
+        let min_peer_count = if ENV_CONFIG.network == Network::Mainnet {
+            10
+        } else {
+            2
+        };
+        if geth_peer_count < min_peer_count {
             info!(geth_peer_count, "geth has less than 10 peers, not ready");
             is_ready.store(false, std::sync::atomic::Ordering::Relaxed);
             sleep(Duration::from_secs(4)).await;