Merge pull request #86 from monadicus/feat-agent-flags

feat(agent): flags

Cargo.toml

@@ -25,6 +25,7 @@ http = "1.1"
 indexmap = { version = "2.2", features = ["serde"] }
 indicatif = { version = "0.17", features = ["rayon"] }
 lazy_static = "1.4"
+lasso = { version = "0.7.2", features = ["multi-threaded"] }
 rand = "0.8"
 rand_chacha = "0.3"
 rayon = "1"

crates/snot-agent/src/cli.rs

@@ -1,9 +1,12 @@
 use std::{
+    env,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::PathBuf,
 };
 
 use clap::Parser;
+use http::Uri;
+use snot_common::state::{AgentId, AgentMode, PortConfig};
 
 pub const ENV_ENDPOINT: &str = "SNOT_ENDPOINT";
 pub const ENV_ENDPOINT_DEFAULT: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 1234);
@@ -16,29 +19,63 @@ pub struct Cli {
     /// Control plane endpoint address
     pub endpoint: Option<SocketAddr>,
 
-    #[arg(long, default_value = "./snot-data")]
+    #[arg(long)]
+    pub id: Option<AgentId>,
+
+    #[arg(long, value_delimiter = ',', num_args = 1..)]
+    pub labels: Option<Vec<String>>,
+
     /// Path to the directory containing the stored data and configuration
+    #[arg(long, default_value = "./snot-data")]
     pub path: PathBuf,
 
-    #[arg(long, default_value_t = false)]
     /// Enable the agent to fetch its external address. Necessary to determine
     /// which agents are on shared networks, and for
     /// external-to-external connections
+    #[arg(long)]
     pub external: bool,
 
     #[clap(long = "bind", default_value_t = IpAddr::V4(Ipv4Addr::UNSPECIFIED))]
     pub bind_addr: IpAddr,
-    /// Specify the IP address and port for the node server
-    #[clap(long = "node", default_value_t = 4130)]
-    pub node: u16,
-    /// Specify the IP address and port for the BFT
-    #[clap(long = "bft", default_value = "5000")]
-    pub bft: u16,
-    /// Specify the IP address and port for the REST server
-    #[clap(long = "rest", default_value = "3030")]
-    pub rest: u16,
-    /// Specify the port for the metrics
-    #[clap(long = "metrics", default_value_t = 9000)]
-    pub metrics: u16,
-    // TODO: specify allowed modes (--validator --client --tx-gen)
+
+    #[clap(flatten)]
+    pub ports: PortConfig,
+
+    #[clap(flatten)]
+    pub modes: AgentMode,
+}
+
+impl Cli {
+    pub fn endpoint_and_uri(&self) -> (SocketAddr, Uri) {
+        // get the endpoint
+        let endpoint = self
+            .endpoint
+            .or_else(|| {
+                env::var(ENV_ENDPOINT)
+                    .ok()
+                    .and_then(|s| s.as_str().parse().ok())
+            })
+            .unwrap_or(ENV_ENDPOINT_DEFAULT);
+
+        let mut query = format!("/agent?mode={}", u8::from(self.modes));
+
+        // add ?id=
+        if let Some(id) = self.id {
+            query.push_str(&format!("&id={}", id));
+        }
+
+        // add ?labels= or &labels= if id is present
+        if let Some(labels) = &self.labels {
+            query.push_str(&format!("&labels={}", labels.join(",")));
+        }
+
+        let ws_uri = Uri::builder()
+            .scheme("ws")
+            .authority(endpoint.to_string())
+            .path_and_query(query)
+            .build()
+            .unwrap();
+
+        (endpoint, ws_uri)
+    }
 }

crates/snot-agent/src/main.rs

@@ -6,15 +6,14 @@ mod rpc;
 mod state;
 
 use std::{
-    env,
     os::unix::fs::PermissionsExt,
     path::Path,
     sync::{Arc, Mutex},
     time::Duration,
 };
 
 use clap::Parser;
-use cli::{Cli, ENV_ENDPOINT, ENV_ENDPOINT_DEFAULT};
+use cli::Cli;
 use futures::{executor::block_on, SinkExt};
 use futures_util::stream::{FuturesUnordered, StreamExt};
 use http::HeaderValue;
@@ -27,7 +26,7 @@ use tokio::{
 };
 use tokio_tungstenite::{
     connect_async,
-    tungstenite::{self, client::IntoClientRequest, http::Uri},
+    tungstenite::{self, client::IntoClientRequest},
 };
 use tracing::{error, info, level_filters::LevelFilter, warn};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
@@ -89,21 +88,7 @@ async fn main() {
     }
 
     // get the endpoint
-    let endpoint = args
-        .endpoint
-        .or_else(|| {
-            env::var(ENV_ENDPOINT)
-                .ok()
-                .and_then(|s| s.as_str().parse().ok())
-        })
-        .unwrap_or(ENV_ENDPOINT_DEFAULT);
-
-    let ws_uri = Uri::builder()
-        .scheme("ws")
-        .authority(endpoint.to_string())
-        .path_and_query("/agent")
-        .build()
-        .unwrap();
+    let (endpoint, ws_uri) = args.endpoint_and_uri();
 
     // create the data directory
     tokio::fs::create_dir_all(&args.path)

crates/snot-agent/src/metrics/mod.rs

@@ -41,7 +41,7 @@ pub fn init(state: Arc<GlobalState>) {
                 let response = match client
                     .get(format!(
                         "http://{}/",
-                        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), state.cli.metrics,)
+                        SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), state.cli.ports.metrics)
                     ))
                     .send()
                     .await

crates/snot-agent/src/rpc.rs

@@ -239,13 +239,13 @@ impl AgentService for AgentRpcServer {
                         .arg("--bind")
                         .arg(state.cli.bind_addr.to_string())
                         .arg("--bft")
-                        .arg(state.cli.bft.to_string())
+                        .arg(state.cli.ports.bft.to_string())
                         .arg("--rest")
-                        .arg(state.cli.rest.to_string())
+                        .arg(state.cli.ports.rest.to_string())
                         .arg("--metrics")
-                        .arg(state.cli.metrics.to_string())
+                        .arg(state.cli.ports.metrics.to_string())
                         .arg("--node")
-                        .arg(state.cli.node.to_string());
+                        .arg(state.cli.ports.node.to_string());
 
                     if let Some(pk) = node.private_key {
                         command.arg("--private-key").arg(pk);
@@ -374,11 +374,7 @@ impl AgentService for AgentRpcServer {
 
     async fn get_addrs(self, _: context::Context) -> (PortConfig, Option<IpAddr>, Vec<IpAddr>) {
         (
-            PortConfig {
-                bft: self.state.cli.bft,
-                node: self.state.cli.node,
-                rest: self.state.cli.rest,
-            },
+            self.state.cli.ports.clone(),
             self.state.external_addr,
             self.state.internal_addrs.clone(),
         )
@@ -394,7 +390,7 @@ impl AgentService for AgentRpcServer {
 
         let url = format!(
             "http://127.0.0.1:{}/mainnet/latest/stateRoot",
-            self.state.cli.rest
+            self.state.cli.ports.rest
         );
         let response = reqwest::get(&url)
             .await
@@ -415,7 +411,7 @@ impl AgentService for AgentRpcServer {
 
         let url = format!(
             "http://127.0.0.1:{}/mainnet/transaction/broadcast",
-            self.state.cli.rest
+            self.state.cli.ports.rest
         );
         let response = reqwest::Client::new()
             .post(url)

crates/snot-common/Cargo.toml

@@ -4,7 +4,9 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+clap.workspace = true
 futures.workspace = true
+lasso.workspace = true
 lazy_static.workspace = true
 regex.workspace = true
 serde.workspace = true

crates/snot-common/src/lib.rs

@@ -1,7 +1,12 @@
 pub mod rpc;
 pub mod state;
+pub use lasso;
 
 pub mod prelude {
     pub use crate::rpc::*;
     pub use crate::state::*;
 }
+
+lazy_static::lazy_static! {
+    pub static ref INTERN: lasso::ThreadedRodeo = lasso::ThreadedRodeo::default();
+}