From 7407d208fd6bd0fb1b44398cd7c86edc660b0080 Mon Sep 17 00:00:00 2001 From: Ayush Bansal Date: Tue, 27 May 2025 01:07:01 +0530 Subject: [PATCH 01/21] Add support for multiple upstream connections and configuration loading - Introduced a new `Config` struct for loading configuration from a file. - Updated `main.rs` to parse command-line arguments and load configuration. - Enhanced `ProxyState` to manage multiple upstream connections with connection status tracking. - Modified `Router` to support round-robin selection of upstreams and connection management. - Added new dependencies: `serde` and `toml` for configuration handling. --- Cargo.lock | 85 ++++++++++++++- Cargo.toml | 2 + config.toml | 9 ++ src/main.rs | 253 +++++++++++++++++++++++++++++++++++++++++---- src/proxy_state.rs | 225 ++++++++++++++++++++++++++++++++++++++++ src/router/mod.rs | 177 +++++++++++++++++++++++++++++-- 6 files changed, 720 insertions(+), 31 deletions(-) create mode 100644 config.toml diff --git a/Cargo.lock b/Cargo.lock index 8ca8babe..dc203dca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -432,7 +432,7 @@ checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" dependencies = [ "cfg-if", "crossbeam-utils", - "hashbrown", + "hashbrown 0.14.5", "lock_api", "once_cell", "parking_lot_core", @@ -460,11 +460,13 @@ dependencies = [ "pid", "rand", "roles_logic_sv2", + "serde", "serde_json", "sha2 0.10.8", "sv1_api", "tokio", "tokio-util", + "toml", "tracing", "tracing-subscriber", ] @@ -519,6 +521,12 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "framing_sv2" version = "3.0.0" @@ -661,6 +669,12 @@ version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +[[package]] +name = "hashbrown" +version = "0.15.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84b26c544d002229e640969970a2e74021aadf6e2f96372b9c58eff97de08eb3" + [[package]] name = "heck" version = "0.5.0" @@ -679,6 +693,16 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "212ab92002354b4819390025006c897e8140934349e8635c9b077f47b4dcbd20" +[[package]] +name = "indexmap" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cea70ddb795996207ad57735b50c5982d8844f38ba9ee5f1aedcfb708a2aa11e" +dependencies = [ + "equivalent", + "hashbrown 0.15.3", +] + [[package]] name = "inout" version = "0.1.3" @@ -1164,6 +1188,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1" +dependencies = [ + "serde", +] + [[package]] name = "sha2" version = "0.9.9" @@ -1351,6 +1384,47 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05ae329d1f08c4d17a59bed7ff5b5a769d062e64a62d34a3261b219e62cd5aae" +dependencies = [ + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", +] + +[[package]] +name = "toml_datetime" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3da5db5a963e24bc68be8b17b6fa82814bb22ee8660f192bb182771d498f09a3" +dependencies = [ + "serde", +] + +[[package]] +name = "toml_edit" +version = "0.22.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "310068873db2c5b3e7659d2cc35d21855dbafa50d1ce336397c666e3cb08137e" +dependencies = [ + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", + "toml_write", + "winnow", +] + +[[package]] +name = "toml_write" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfb942dfe1d8e29a7ee7fcbde5bd2b9a25fb89aa70caea2eba3bee836ff41076" + [[package]] name = "tracing" version = "0.1.41" @@ -1562,6 +1636,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06928c8748d81b05c9be96aad92e1b6ff01833332f281e8cfca3be4b35fc9ec" +dependencies = [ + "memchr", +] + [[package]] name = "zerocopy" version = "0.7.35" diff --git a/Cargo.toml b/Cargo.toml index 1b5b64cf..f0917930 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,8 @@ tokio = {version="^1.36.0",features = ["full","tracing"]} key-utils = "1.0.0" pid ={ version = "4.0.0"} clap={version = "4.5.31", features = ["derive"]} +serde = { version = "1.0", features = ["derive"] } +toml = "0.8" #roles_logic_sv2 = "1.2.1" #sv1_api = "1.0.1" #demand-sv2-connection = "0.0.3" diff --git a/config.toml b/config.toml new file mode 100644 index 00000000..ae2d7491 --- /dev/null +++ b/config.toml @@ -0,0 +1,9 @@ +token="UR01TMkZv6Vs5zbwr0t6" +pool_addresses=["18.193.252.132:2000", "3.74.36.119:2000","18.193.252.132:2000", "3.74.36.119:2000"] +tp_address = "127.0.0.1:8442" +interval = 120_000 +delay = 540 +downstream_hashrate = "200T" +loglevel = "info" +nc_loglevel = "off" +test = true \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 2d95fe58..8875186e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use clap::Parser; +use clap::{Parser, ArgAction}; #[cfg(not(target_os = "windows"))] use jemallocator::Jemalloc; use router::Router; @@ -41,17 +41,27 @@ const TEST_AUTH_PUB_KEY: &str = "9auqWEzQDVyd2oe1JVGFLMLHZtCo2FFqZwtKA5gd9xbuEu7 const DEFAULT_LISTEN_ADDRESS: &str = "0.0.0.0:32767"; lazy_static! { + static ref ARGS: Args = Args::parse(); + // Remove or comment out the old TOKEN line: + // static ref TOKEN: String = std::env::var("TOKEN").expect("Missing TOKEN environment variable"); + + // Other existing lazy_static variables... static ref SV1_DOWN_LISTEN_ADDR: String = std::env::var("SV1_DOWN_LISTEN_ADDR").unwrap_or(DEFAULT_LISTEN_ADDRESS.to_string()); static ref TP_ADDRESS: roles_logic_sv2::utils::Mutex> = - roles_logic_sv2::utils::Mutex::new(std::env::var("TP_ADDRESS").ok()); - static ref EXPECTED_SV1_HASHPOWER: f32 = Args::parse() - .downstream_hashrate - .unwrap_or(DEFAULT_SV1_HASHPOWER); + roles_logic_sv2::utils::Mutex::new(None); // We'll set this from config + static ref EXPECTED_SV1_HASHPOWER: f32 = { + if let Some(value) = ARGS.downstream_hashrate { + value + } else { + let env_var = std::env::var("EXPECTED_SV1_HASHPOWER").ok(); + env_var.and_then(|s| parse_hashrate(&s).ok()) + .unwrap_or(DEFAULT_SV1_HASHPOWER) + } + }; } lazy_static! { - static ref ARGS: Args = Args::parse(); pub static ref POOL_ADDRESS: &'static str = if ARGS.test { TEST_POOL_ADDRESS } else { @@ -63,8 +73,9 @@ lazy_static! { MAIN_AUTH_PUB_KEY }; } -#[derive(Parser)] -struct Args { +#[derive(Parser, Debug)] +#[clap(author, version, about, long_about = None)] +pub struct Args { // Use test enpoint if test flag is provided #[clap(long)] test: bool, @@ -74,6 +85,17 @@ struct Args { loglevel: String, #[clap(long = "nc", short = 'n', default_value = "off")] noise_connection_log: String, + // New argument for multiple upstream servers + #[clap(long = "upstream", short = 'u', action = ArgAction::Append)] + upstream_servers: Option>, + // Option to enable round-robin distribution + #[clap(long = "round-robin", short = 'r')] + round_robin: bool, + // In Args struct +#[clap(long = "monitor-hashrate", short = 'm')] +monitor_hashrate: bool, + #[clap(long = "config", short = 'c')] + pub config_file: Option, } #[tokio::main] @@ -112,9 +134,44 @@ async fn main() { log_level, noise_connection_log_level ))) .init(); - std::env::var("TOKEN").expect("Missing TOKEN environment variable"); + // std::env::var("TOKEN").expect("Missing TOKEN environment variable"); + + // Load configuration first + let config = if let Some(config_path) = &ARGS.config_file { + match Config::from_file(config_path) { + Ok(config) => Some(config), + Err(e) => { + error!("Failed to load config file: {}", e); + std::process::exit(1); + } + } + } else { + None + }; - let hashpower = *EXPECTED_SV1_HASHPOWER; + // Get TOKEN from config or environment + let token = if let Some(ref config) = config { + config.token.clone() + } else { + std::env::var("TOKEN").unwrap_or_else(|_| { + error!("TOKEN must be provided either in config file or as environment variable"); + std::process::exit(1); + }) + }; + + + // Use configured hashrate if available + let hashpower = if let Some(ref config) = config { + if let Some(ref hashrate_str) = config.downstream_hashrate { + parse_hashrate(hashrate_str).unwrap_or(*EXPECTED_SV1_HASHPOWER) + } else { + *EXPECTED_SV1_HASHPOWER + } + } else { + *EXPECTED_SV1_HASHPOWER + }; + + ProxyState::set_total_hashrate(hashpower); if args.downstream_hashrate.is_some() { info!( @@ -127,21 +184,102 @@ async fn main() { HashUnit::format_value(hashpower) ); } - if args.test { - info!("Connecting to test endpoint..."); + + // Handle multiple upstream servers from config or hard-coded fallback + let mut pool_addresses = Vec::new(); + + if let Some(ref config) = config { + // Use configuration file pool addresses + info!("Loading pool addresses from configuration file"); + + for (idx, pool_addr_str) in config.pool_addresses.iter().enumerate() { + // Parse address + let addr = if let Ok(mut addrs) = pool_addr_str.to_socket_addrs() { + addrs.next().unwrap_or_else(|| { + error!("Failed to resolve address: {}", pool_addr_str); + std::process::exit(1); + }) + } else { + error!("Invalid pool address: {}", pool_addr_str); + std::process::exit(1); + }; + + // Use test public key for now (you might want to add this to config) + let pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY + .parse() + .expect("Invalid test public key"); + + info!("Adding upstream server {}: {}", idx + 1, addr); + pool_addresses.push((addr, pubkey.clone())); + + ProxyState::add_upstream_connection( + format!("upstream-{}", idx), + format!("config-pool-{}", idx + 1), + addr, + pubkey, + crate::proxy_state::UpstreamType::JDCMiningUpstream, + ); + } + } else { + // Fallback to hard-coded addresses (your existing code) + info!("Using hard-coded pool addresses"); + + let test_pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY + .parse() + .expect("Invalid test public key"); + + let test_addr = if let Ok(addr) = "3.74.36.119:2000".to_socket_addrs() { + addr.collect::>()[0] + } else { + "3.74.36.119:2000".parse().expect("Invalid IP address") + }; + + // Add hard-coded pools + for i in 0..2 { + info!("Adding upstream server {}: {}", i + 1, test_addr); + pool_addresses.push((test_addr, test_pubkey.clone())); + ProxyState::add_upstream_connection( + format!("upstream-{}", i), + format!("test-pool-{}", i + 1), + test_addr, + test_pubkey.clone(), + crate::proxy_state::UpstreamType::JDCMiningUpstream, + ); + } } - let auth_pub_k: Secp256k1PublicKey = AUTH_PUB_KEY.parse().expect("Invalid public key"); - let address = POOL_ADDRESS - .to_socket_addrs() - .expect("Invalid pool address") - .next() - .expect("Invalid pool address"); + // Direct verification + info!("DIRECT VERIFICATION: Checking hashrate distribution"); + let upstream_connections = ProxyState::get_upstream_connections(); + let active_count = upstream_connections.len(); + info!("Active upstreams: {}", active_count); + + if active_count > 0 && ARGS.round_robin { + let hashrate_per_upstream = hashpower / active_count as f32; + info!("Round-robin hashrate per upstream: {}", HashUnit::format_value(hashrate_per_upstream)); + + for (upstream_id, addr, _) in upstream_connections { + info!("Upstream {}: {} - allocated {}", + upstream_id, + addr, + HashUnit::format_value(hashrate_per_upstream)); + } + } - // We will add upstream addresses here - let pool_addresses = vec![address]; + // Set all upstreams as initially connected + for (idx, (addr, _)) in pool_addresses.iter().enumerate() { + let id = if idx == 0 && pool_addresses.len() == 1 { + "upstream-default".to_string() + } else { + format!("upstream-{}", idx) + }; + ProxyState::set_upstream_connection_status(&id, true); + } - let mut router = router::Router::new(pool_addresses, auth_pub_k, None, None); + let pool_socket_addresses: Vec = pool_addresses.iter().map(|(addr, _)| *addr).collect(); + let auth_pub_keys: Vec = pool_addresses.iter().map(|(_, pubkey)| pubkey.clone()).collect(); + let auth_pub_k = auth_pub_keys.first().expect("No public keys available").clone(); + let mut router = router::Router::new(pool_socket_addresses, auth_pub_k, None, None); let epsilon = Duration::from_millis(10); let best_upstream = router.select_pool_connect().await; initialize_proxy(&mut router, best_upstream, epsilon).await; @@ -280,10 +418,54 @@ async fn monitor( epsilon: Duration, ) -> Reconnect { let mut should_check_upstreams_latency = 0; + let mut distribution_check_counter = 0; + loop { + if distribution_check_counter >= 100 { + distribution_check_counter = 0; + + // Add debug output to confirm this code runs + info!("Generating hashrate distribution report..."); + + info!("Hashrate Distribution Report:"); + info!("Total hashrate: {}", HashUnit::format_value(*EXPECTED_SV1_HASHPOWER)); + + // Get upstream info directly from ProxyState + if ARGS.round_robin { + // In round-robin mode, check how many active upstreams we have + let upstream_connections = ProxyState::get_upstream_connections(); + let active_count = upstream_connections.len(); + + if active_count > 0 { + let hashrate_per_upstream = *EXPECTED_SV1_HASHPOWER / active_count as f32; + info!("Round-robin mode: {} active upstreams", active_count); + info!("Each upstream allocated: {}", HashUnit::format_value(hashrate_per_upstream)); + + // List each upstream and its allocated hashrate + for (upstream_id, addr, _) in upstream_connections { + info!("Upstream {}: {} - allocated {}", + upstream_id, + addr, + HashUnit::format_value(hashrate_per_upstream)); + } + } else { + info!("No active upstreams found"); + } + } else { + // In latency-based mode + info!("Latency-based mode: Using best upstream"); + if let Some(current_addr) = router.get_current_upstream() { + info!("Current upstream: {} - allocated {}", + current_addr, + HashUnit::format_value(*EXPECTED_SV1_HASHPOWER)); + } else { + info!("No upstream currently selected"); + } + } + } + // Check if a better upstream exist every 100 seconds if should_check_upstreams_latency == 10 * 100 { - should_check_upstreams_latency = 0; if let Some(new_upstream) = router.monitor_upstream(epsilon).await { info!("Faster upstream detected. Reinitializing proxy..."); drop(abort_handles); @@ -328,7 +510,11 @@ async fn monitor( return Reconnect::NoUpstream; } + // Increment the counters should_check_upstreams_latency += 1; + distribution_check_counter += 1; + + // Make sure this line exists to give time for the loop tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; } } @@ -410,3 +596,26 @@ impl HashUnit { } } } +use serde::{Deserialize, Serialize}; +use std::path::Path; + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct Config { + pub token: String, + pub pool_addresses: Vec, + pub tp_address: Option, + pub interval: Option, + pub delay: Option, + pub downstream_hashrate: Option, + pub loglevel: Option, + pub nc_loglevel: Option, + pub test: Option, +} + +impl Config { + pub fn from_file>(path: P) -> Result> { + let contents = std::fs::read_to_string(path)?; + let config: Config = toml::from_str(&contents)?; + Ok(config) + } +} diff --git a/src/proxy_state.rs b/src/proxy_state.rs index 6e6a90b1..8b04c32c 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -1,4 +1,6 @@ +use std::collections::HashMap; use std::sync::Arc; +use std::time::Instant; use lazy_static::lazy_static; use roles_logic_sv2::utils::Mutex; @@ -84,6 +86,19 @@ pub enum UpstreamType { TranslatorUpstream, } +/// Create an UpstreamConnection struct to store connection info +#[derive(Debug, Clone)] +pub struct UpstreamConnection { + pub url: String, + pub address: std::net::SocketAddr, + pub auth_key: key_utils::Secp256k1PublicKey, + pub connection_type: UpstreamType, + pub is_connected: bool, + pub shares_submitted: u64, + pub shares_accepted: u64, + pub last_used: Instant, +} + /// Represents global proxy state #[derive(Debug)] pub struct ProxyState { @@ -95,6 +110,11 @@ pub struct ProxyState { pub inconsistency: Option, pub downstream: DownstreamState, pub upstream: UpstreamState, + + // New fields for multiple upstream support + pub upstream_connections: HashMap, + pub total_hashrate: f32, + pub current_upstream_index: usize, } impl ProxyState { @@ -108,6 +128,11 @@ impl ProxyState { inconsistency: None, downstream: DownstreamState::Up, upstream: UpstreamState::Up, + + // Initialize new fields + upstream_connections: HashMap::new(), + total_hashrate: 0.0, + current_upstream_index: 0, } } @@ -286,4 +311,204 @@ impl ProxyState { Ok(errors) } } + + /// Add a new upstream connection + pub fn add_upstream_connection( + id: String, + url: String, + address: std::net::SocketAddr, + auth_key: key_utils::Secp256k1PublicKey, + connection_type: UpstreamType, + ) { + info!("Adding upstream connection: {} at {}", id, url); + + if PROXY_STATE + .safe_lock(|state| { + let connection = UpstreamConnection { + url, + address, + auth_key, + connection_type, + is_connected: false, + shares_submitted: 0, + shares_accepted: 0, + last_used: Instant::now(), + }; + + state.upstream_connections.insert(id, connection); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Update connection status for an upstream + pub fn set_upstream_connection_status(id: &str, connected: bool) { + if PROXY_STATE + .safe_lock(|state| { + if let Some(conn) = state.upstream_connections.get_mut(id) { + conn.is_connected = connected; + if connected { + info!("Upstream {} is now connected", id); + } else { + info!("Upstream {} is now disconnected", id); + } + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Set the total hashrate to be distributed among upstreams + pub fn set_total_hashrate(hashrate: f32) { + info!("Setting total hashrate to: {} h/s", hashrate); + if PROXY_STATE + .safe_lock(|state| { + state.total_hashrate = hashrate; + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Get the next upstream in round-robin fashion + pub fn get_next_upstream() -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + let mut result = None; + + if PROXY_STATE + .safe_lock(|state| { + // Get IDs of all connected upstreams + let active_upstreams: Vec<&String> = state.upstream_connections + .iter() + .filter(|(_, conn)| conn.is_connected) + .map(|(id, _)| id) + .collect(); + + if active_upstreams.is_empty() { + return; + } + + // Use round-robin to select the next upstream + if state.current_upstream_index >= active_upstreams.len() { + state.current_upstream_index = 0; + } + + let id = active_upstreams[state.current_upstream_index].clone(); + if let Some(conn) = state.upstream_connections.get(&id) { + result = Some(( + id.clone(), + conn.address, + conn.auth_key.clone() + )); + } + + // Update index for next call + state.current_upstream_index = (state.current_upstream_index + 1) % active_upstreams.len(); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + result + } + + /// Get the hashrate for a specific upstream (equal distribution) + pub fn get_hashrate_for_upstream(id: Option<&str>) -> f32 { + let mut hashrate = 0.0; + + if PROXY_STATE + .safe_lock(|state| { + // Count active connections + let active_count = state.upstream_connections + .values() + .filter(|conn| conn.is_connected) + .count(); + + if active_count > 0 { + // Equal distribution - each upstream gets the same portion + hashrate = state.total_hashrate / active_count as f32; + + // If a specific ID was provided, check if it's active + if let Some(id) = id { + if let Some(conn) = state.upstream_connections.get(id) { + if !conn.is_connected { + // If this specific upstream isn't connected, return 0 + hashrate = 0.0; + } + } else { + // If this ID doesn't exist, return 0 + hashrate = 0.0; + } + } + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + hashrate + } + + /// Record a share submission to an upstream + pub fn record_share_submission(upstream_id: &str) { + if PROXY_STATE + .safe_lock(|state| { + if let Some(conn) = state.upstream_connections.get_mut(upstream_id) { + conn.shares_submitted += 1; + conn.last_used = Instant::now(); + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Record a share acceptance from an upstream + pub fn record_share_acceptance(upstream_id: &str) { + if PROXY_STATE + .safe_lock(|state| { + if let Some(conn) = state.upstream_connections.get_mut(upstream_id) { + conn.shares_accepted += 1; + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Get all upstream connections + pub fn get_upstream_connections() -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + let mut connections = Vec::new(); + + if PROXY_STATE + .safe_lock(|state| { + connections = state.upstream_connections + .iter() + .filter(|(_, conn)| conn.is_connected) + .map(|(id, conn)| (id.clone(), conn.address, conn.auth_key.clone())) + .collect(); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + connections + } } diff --git a/src/router/mod.rs b/src/router/mod.rs index 2e913515..0a0cc037 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -18,16 +18,19 @@ use tracing::{error, info}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, + proxy_state::ProxyState, // Add ProxyState import shared::utils::AbortOnDrop, }; /// Router handles connection to Multiple upstreams. pub struct Router { pool_addresses: Vec, + auth_keys: Vec, // Store auth keys separately to map to addresses current_pool: Option, auth_pub_k: Secp256k1PublicKey, setup_connection_msg: Option>, timer: Option, + use_round_robin: bool, // Add flag to control selection strategy } impl Router { @@ -42,17 +45,82 @@ impl Router { // If None, default time of 5s is used. timer: Option, ) -> Self { + // Create auth_keys vector with the same key for all addresses + let auth_keys = vec![auth_pub_k.clone(); pool_addresses.len()]; + Self { pool_addresses, + auth_keys, current_pool: None, auth_pub_k, setup_connection_msg, timer, + use_round_robin: false, // Default to latency-based selection } } + + /// Creates a new Router with multiple upstream addresses and auth keys + pub fn new_multi( + pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, + setup_connection_msg: Option>, + timer: Option, + use_round_robin: bool, + ) -> Result { + if pool_address_keys.is_empty() { + return Err("Cannot create router with empty pool_address_keys"); + } + + let mut addresses = Vec::new(); + let mut keys = Vec::new(); + + for (addr, key) in &pool_address_keys { + addresses.push(*addr); + keys.push(key.clone()); + } + + // Use the first key as the default auth_pub_k + let default_key = pool_address_keys[0].1.clone(); + + // Register all upstreams in the ProxyState + for (idx, (addr, key)) in pool_address_keys.iter().enumerate() { + let id = format!("upstream-{}", idx); + ProxyState::add_upstream_connection( + id.clone(), + format!("{:?}", addr), + *addr, + key.clone(), + crate::proxy_state::UpstreamType::JDCMiningUpstream, + ); + } + + Ok(Self { + pool_addresses: addresses, + auth_keys: keys, + current_pool: None, + auth_pub_k: default_key, + setup_connection_msg, + timer, + use_round_robin, + }) + } + + /// Enable or disable round-robin upstream selection + pub fn set_round_robin(&mut self, enabled: bool) { + self.use_round_robin = enabled; + } /// Internal function to select pool with the least latency. async fn select_pool(&self) -> Option<(SocketAddr, Duration)> { + // If round-robin is enabled, use ProxyState to get the next upstream + if self.use_round_robin { + if let Some((id, addr, _)) = ProxyState::get_next_upstream() { + info!("Round-robin selected upstream {}: {:?}", id, addr); + // Return with a dummy zero latency + return Some((addr, Duration::from_millis(0))); + } + } + + // Fall back to latency-based selection let mut best_pool = None; let mut least_latency = Duration::MAX; @@ -69,19 +137,61 @@ impl Router { } /// Select the best pool for connection - pub async fn select_pool_connect(&self) -> Option { - info!("Selecting the best upstream "); + pub async fn select_pool_connect(&mut self) -> Option { + info!("Selecting the best upstream"); + + if self.use_round_robin { + // Check if we have any registered upstreams in ProxyState + if let Some((id, addr, _)) = ProxyState::get_next_upstream() { + info!("Round-robin selected upstream {}: {:?}", id, addr); + return Some(addr); + } + + // If no upstreams registered in ProxyState yet, register them now + if !self.pool_addresses.is_empty() { + for (idx, (addr, key)) in self.pool_addresses.iter().zip(self.auth_keys.iter()).enumerate() { + let id = format!("upstream-{}", idx); + ProxyState::add_upstream_connection( + id, + format!("{:?}", addr), + *addr, + key.clone(), + crate::proxy_state::UpstreamType::JDCMiningUpstream, + ); + // Mark all as initially connected + ProxyState::set_upstream_connection_status(&format!("upstream-{}", idx), true); + } + + // Now try again to get a round-robin selection + if let Some((_, addr, _)) = ProxyState::get_next_upstream() { + return Some(addr); + } + } + } + + // Fall back to latency-based selection if let Some((pool, latency)) = self.select_pool().await { info!("Latency for upstream {:?} is {:?}", pool, latency); Some(pool) } else { - //info!("No available pool"); None } } /// Select the best pool for monitoring async fn select_pool_monitor(&self, epsilon: Duration) -> Option { + // If using round-robin, just get the next upstream + if self.use_round_robin { + if let Some((_, addr, _)) = ProxyState::get_next_upstream() { + // Only switch if it's different from the current one + if Some(addr) != self.current_pool { + return Some(addr); + } + return None; + } + } + + // Otherwise use latency-based selection if let Some((best_pool, best_pool_latency)) = self.select_pool().await { if let Some(current_pool) = self.current_pool { if best_pool == current_pool { @@ -137,30 +247,56 @@ impl Router { self.current_pool = Some(pool); info!("Upstream {:?} selected", pool); + + // Find the matching auth key for this address + let auth_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { + self.auth_keys[index].clone() + } else { + self.auth_pub_k.clone() + }; match minin_pool_connection::connect_pool( pool, - self.auth_pub_k, + auth_key, self.setup_connection_msg.clone(), self.timer, ) .await { Ok((send_to_pool, recv_from_pool, pool_connection_abortable)) => { + // Update ProxyState with successful connection + let upstream_id = format!("upstream-{}", + self.pool_addresses.iter().position(|&a| a == pool).unwrap_or(0)); + ProxyState::set_upstream_connection_status(&upstream_id, true); + Ok((send_to_pool, recv_from_pool, pool_connection_abortable)) } - Err(e) => Err(e), + Err(e) => { + // Update ProxyState with failed connection + let upstream_id = format!("upstream-{}", + self.pool_addresses.iter().position(|&a| a == pool).unwrap_or(0)); + ProxyState::set_upstream_connection_status(&upstream_id, false); + + Err(e) + } } } /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { + // Find the auth key for this address + let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool_address) { + self.auth_keys[index].clone() + } else { + self.auth_pub_k.clone() + }; + let mut pool = PoolLatency::new(pool_address); let setup_connection_msg = self.setup_connection_msg.as_ref(); let timer = self.timer.as_ref(); - let auth_pub_key = self.auth_pub_k; + // Rest of the function remains the same tokio::time::timeout( Duration::from_secs(15), PoolLatency::get_mining_setup_latencies( @@ -178,6 +314,7 @@ impl Router { ); })??; + // Rest of the function remains unchanged if (PoolLatency::get_mining_setup_latencies( &mut pool, setup_connection_msg.cloned(), @@ -213,9 +350,21 @@ impl Router { /// Checks for faster upstream switch to it if found pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { + // If using round-robin, just get the next upstream + if self.use_round_robin { + if let Some((_, addr, _)) = ProxyState::get_next_upstream() { + if Some(addr) != self.current_pool { + info!("Switching to next round-robin upstream {:?}", addr); + return Some(addr); + } + return None; + } + } + + // Otherwise use latency-based selection if let Some(best_pool) = self.select_pool_monitor(epsilon).await { if Some(best_pool) != self.current_pool { - info!("Switching to faster upstreamn {:?}", best_pool); + info!("Switching to faster upstream {:?}", best_pool); return Some(best_pool); } else { return None; @@ -223,6 +372,18 @@ impl Router { } None } + + pub fn get_current_upstream(&self) -> Option { + self.current_pool + } + + // For compatibility with existing code + pub fn is_current_upstream(&self, addr: &std::net::SocketAddr) -> bool { + if let Some(current) = self.current_pool { + return current == *addr; + } + false + } } /// Track latencies for various stages of pool connection setup. @@ -466,4 +627,4 @@ async fn initialize_mining_connections( let setup_connection_msg = setup_connection_msg.unwrap_or(get_mining_setup_connection_msg(true)); Ok((receiver, sender, setup_connection_msg)) -} +} \ No newline at end of file From e480555b8f997f36551957eb123015a4dec8b79e Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Tue, 27 May 2025 14:31:49 +0530 Subject: [PATCH 02/21] Refactor upstream connection handling and router initialization for multi-upstream support --- Cargo.lock | 84 ++++++++++++++++++++-------------------------- src/main.rs | 44 +++++++++++++++++++++--- src/proxy_state.rs | 2 +- src/router/mod.rs | 8 ++--- 4 files changed, 81 insertions(+), 57 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9913d924..886924ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -907,16 +907,6 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "212ab92002354b4819390025006c897e8140934349e8635c9b077f47b4dcbd20" -[[package]] -name = "indexmap" -version = "2.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cea70ddb795996207ad57735b50c5982d8844f38ba9ee5f1aedcfb708a2aa11e" -dependencies = [ - "equivalent", - "hashbrown 0.15.3", -] - [[package]] name = "hex-conservative" version = "0.2.1" @@ -1665,23 +1655,23 @@ dependencies = [ ] [[package]] -name = "serde_urlencoded" -version = "0.7.1" +name = "serde_spanned" +version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1" dependencies = [ - "form_urlencoded", - "itoa", - "ryu", "serde", ] [[package]] -name = "serde_spanned" -version = "0.6.8" +name = "serde_urlencoded" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87607cb1398ed59d48732e575a4c28a7a8ebf2454b964fe3f224f2afc07909e1" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" dependencies = [ + "form_urlencoded", + "itoa", + "ryu", "serde", ] @@ -1904,34 +1894,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "tower" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" -dependencies = [ - "futures-core", - "futures-util", - "pin-project-lite", - "sync_wrapper", - "tokio", - "tower-layer", - "tower-service", - "tracing", -] - -[[package]] -name = "tower-layer" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" - -[[package]] -name = "tower-service" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" - [[package]] name = "toml" version = "0.8.22" @@ -1973,6 +1935,34 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bfb942dfe1d8e29a7ee7fcbde5bd2b9a25fb89aa70caea2eba3bee836ff41076" +[[package]] +name = "tower" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + [[package]] name = "tracing" version = "0.1.41" diff --git a/src/main.rs b/src/main.rs index c4633bd6..eca76ece 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,6 +14,7 @@ use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; use std::{net::ToSocketAddrs, time::Duration}; use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; +use std::net::SocketAddr; mod api; mod ingress; pub mod jd_client; @@ -262,10 +263,13 @@ async fn main() { let hashrate_per_upstream = hashpower / active_count as f32; info!("Round-robin hashrate per upstream: {}", HashUnit::format_value(hashrate_per_upstream)); - for (upstream_id, addr, _) in upstream_connections { + let upstream_connections: Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> = ProxyState::get_upstream_connections(); + + // And update the usage in the loop below: + for (upstream_id, addr, _auth_key) in upstream_connections { info!("Upstream {}: {} - allocated {}", upstream_id, - addr, + addr, HashUnit::format_value(hashrate_per_upstream)); } } @@ -283,7 +287,34 @@ async fn main() { let pool_socket_addresses: Vec = pool_addresses.iter().map(|(addr, _)| *addr).collect(); let auth_pub_keys: Vec = pool_addresses.iter().map(|(_, pubkey)| pubkey.clone()).collect(); let auth_pub_k = auth_pub_keys.first().expect("No public keys available").clone(); - let mut router = router::Router::new(pool_socket_addresses, auth_pub_k, None, None); + // Create the router based on whether we're using round-robin or not +let mut router = if !pool_addresses.is_empty() && pool_addresses.len() > 1 { + // Multiple upstreams - use new_multi + match router::Router::new_multi( + pool_addresses.clone(), + None, // setup_connection_msg + None, // timer + ARGS.round_robin + ) { + Ok(router) => router, + Err(e) => { + error!("Failed to create multi-upstream router: {}", e); + std::process::exit(1); + } + } +} else { + // Single upstream - use regular constructor + let pool_socket_addresses: Vec = pool_addresses.iter().map(|(addr, _)| *addr).collect(); + let auth_pub_k = pool_addresses.first().expect("No pool addresses available").1.clone(); + + router::Router::new( + pool_socket_addresses, + auth_pub_k, + None, // setup_connection_msg + None // timer + ) +}; + let epsilon = Duration::from_millis(10); let best_upstream = router.select_pool_connect().await; initialize_proxy(&mut router, best_upstream, epsilon).await; @@ -451,10 +482,10 @@ async fn monitor( info!("Each upstream allocated: {}", HashUnit::format_value(hashrate_per_upstream)); // List each upstream and its allocated hashrate - for (upstream_id, addr, _) in upstream_connections { + for (upstream_id, addr, _auth_key) in upstream_connections { info!("Upstream {}: {} - allocated {}", upstream_id, - addr, + addr, HashUnit::format_value(hashrate_per_upstream)); } } else { @@ -532,6 +563,7 @@ async fn monitor( } } + /// Parses a hashrate string (e.g., "10T", "2.5P", "500E") into an f32 value in h/s. fn parse_hashrate(hashrate_str: &str) -> Result { let hashrate_str = hashrate_str.trim(); @@ -565,6 +597,8 @@ fn parse_hashrate(hashrate_str: &str) -> Result { Ok(hashrate) } + + pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream diff --git a/src/proxy_state.rs b/src/proxy_state.rs index 8efb43dc..402a624f 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -101,7 +101,7 @@ pub struct UpstreamConnection { } /// Represents global proxy state -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug)] pub struct ProxyState { pub pool: PoolState, pub tp: TpState, diff --git a/src/router/mod.rs b/src/router/mod.rs index 5b4d2e3e..8705636b 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -17,7 +17,7 @@ use tokio::{ watch, }, }; -use tracing::{error, info}; +use tracing::{error, info,warn}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, @@ -36,7 +36,6 @@ pub struct Router { timer: Option, latency_tx: watch::Sender>, pub latency_rx: watch::Receiver>, - timer: Option, use_round_robin: bool, // Add flag to control selection strategy } @@ -80,6 +79,7 @@ impl Router { return Err("Cannot create router with empty pool_address_keys"); } + let (latency_tx, latency_rx) = watch::channel(None); let mut addresses = Vec::new(); let mut keys = Vec::new(); @@ -108,11 +108,11 @@ impl Router { auth_keys: keys, current_pool: None, auth_pub_k: default_key, - setup_connection_msg, + setup_connection_msg, timer, latency_tx, latency_rx, - use_round_robin: false, // Default to latency-based selection + use_round_robin, }) } From 0067de8a10c6a71db814748e89b3ac5dadf02329 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Tue, 27 May 2025 15:02:31 +0530 Subject: [PATCH 03/21] Refactor code for improved readability and maintainability; remove unused variables and streamline function signatures --- src/main.rs | 175 ++++++++++--------- src/proxy_state.rs | 30 ++-- src/router/mod.rs | 110 +++++++----- src/translator/downstream/diff_management.rs | 24 +-- 4 files changed, 188 insertions(+), 151 deletions(-) diff --git a/src/main.rs b/src/main.rs index eca76ece..b1228857 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use clap::{Parser, ArgAction}; +use clap::{ArgAction, Parser}; #[cfg(not(target_os = "windows"))] use jemallocator::Jemalloc; use router::Router; @@ -14,7 +14,6 @@ use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; use std::{net::ToSocketAddrs, time::Duration}; use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; -use std::net::SocketAddr; mod api; mod ingress; pub mod jd_client; @@ -45,7 +44,7 @@ lazy_static! { static ref ARGS: Args = Args::parse(); // Remove or comment out the old TOKEN line: // static ref TOKEN: String = std::env::var("TOKEN").expect("Missing TOKEN environment variable"); - + // Other existing lazy_static variables... static ref SV1_DOWN_LISTEN_ADDR: String = std::env::var("SV1_DOWN_LISTEN_ADDR").unwrap_or(DEFAULT_LISTEN_ADDRESS.to_string()); @@ -97,8 +96,8 @@ pub struct Args { #[clap(long = "round-robin", short = 'r')] round_robin: bool, // In Args struct -#[clap(long = "monitor-hashrate", short = 'm')] -monitor_hashrate: bool, + #[clap(long = "monitor-hashrate", short = 'm')] + monitor_hashrate: bool, #[clap(long = "config", short = 'c')] pub config_file: Option, } @@ -155,7 +154,7 @@ async fn main() { }; // Get TOKEN from config or environment - let token = if let Some(ref config) = config { + let _token = if let Some(ref config) = config { config.token.clone() } else { std::env::var("TOKEN").unwrap_or_else(|_| { @@ -164,7 +163,6 @@ async fn main() { }) }; - // Use configured hashrate if available let hashpower = if let Some(ref config) = config { if let Some(ref hashrate_str) = config.downstream_hashrate { @@ -175,7 +173,7 @@ async fn main() { } else { *EXPECTED_SV1_HASHPOWER }; - + ProxyState::set_total_hashrate(hashpower); if args.downstream_hashrate.is_some() { @@ -192,11 +190,11 @@ async fn main() { // Handle multiple upstream servers from config or hard-coded fallback let mut pool_addresses = Vec::new(); - + if let Some(ref config) = config { // Use configuration file pool addresses info!("Loading pool addresses from configuration file"); - + for (idx, pool_addr_str) in config.pool_addresses.iter().enumerate() { // Parse address let addr = if let Ok(mut addrs) = pool_addr_str.to_socket_addrs() { @@ -208,15 +206,14 @@ async fn main() { error!("Invalid pool address: {}", pool_addr_str); std::process::exit(1); }; - + // Use test public key for now (you might want to add this to config) - let pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY - .parse() - .expect("Invalid test public key"); - + let pubkey: Secp256k1PublicKey = + TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key"); + info!("Adding upstream server {}: {}", idx + 1, addr); - pool_addresses.push((addr, pubkey.clone())); - + pool_addresses.push((addr, pubkey)); + ProxyState::add_upstream_connection( format!("upstream-{}", idx), format!("config-pool-{}", idx + 1), @@ -228,26 +225,25 @@ async fn main() { } else { // Fallback to hard-coded addresses (your existing code) info!("Using hard-coded pool addresses"); - - let test_pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY - .parse() - .expect("Invalid test public key"); - + + let test_pubkey: Secp256k1PublicKey = + TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key"); + let test_addr = if let Ok(addr) = "3.74.36.119:2000".to_socket_addrs() { addr.collect::>()[0] } else { "3.74.36.119:2000".parse().expect("Invalid IP address") }; - + // Add hard-coded pools for i in 0..2 { info!("Adding upstream server {}: {}", i + 1, test_addr); - pool_addresses.push((test_addr, test_pubkey.clone())); + pool_addresses.push((test_addr, test_pubkey)); ProxyState::add_upstream_connection( format!("upstream-{}", i), format!("test-pool-{}", i + 1), test_addr, - test_pubkey.clone(), + test_pubkey, crate::proxy_state::UpstreamType::JDCMiningUpstream, ); } @@ -261,59 +257,69 @@ async fn main() { if active_count > 0 && ARGS.round_robin { let hashrate_per_upstream = hashpower / active_count as f32; - info!("Round-robin hashrate per upstream: {}", HashUnit::format_value(hashrate_per_upstream)); - - let upstream_connections: Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> = ProxyState::get_upstream_connections(); + info!( + "Round-robin hashrate per upstream: {}", + HashUnit::format_value(hashrate_per_upstream) + ); + + let upstream_connections: Vec<( + String, + std::net::SocketAddr, + key_utils::Secp256k1PublicKey, + )> = ProxyState::get_upstream_connections(); // And update the usage in the loop below: for (upstream_id, addr, _auth_key) in upstream_connections { - info!("Upstream {}: {} - allocated {}", - upstream_id, - addr, - HashUnit::format_value(hashrate_per_upstream)); + info!( + "Upstream {}: {} - allocated {}", + upstream_id, + addr, + HashUnit::format_value(hashrate_per_upstream) + ); } } // Set all upstreams as initially connected - for (idx, (addr, _)) in pool_addresses.iter().enumerate() { + for (idx, (_addr, _)) in pool_addresses.iter().enumerate() { let id = if idx == 0 && pool_addresses.len() == 1 { "upstream-default".to_string() } else { - format!("upstream-{}", idx) + format!("upstream-{}", idx) }; ProxyState::set_upstream_connection_status(&id, true); } - let pool_socket_addresses: Vec = pool_addresses.iter().map(|(addr, _)| *addr).collect(); - let auth_pub_keys: Vec = pool_addresses.iter().map(|(_, pubkey)| pubkey.clone()).collect(); - let auth_pub_k = auth_pub_keys.first().expect("No public keys available").clone(); // Create the router based on whether we're using round-robin or not -let mut router = if !pool_addresses.is_empty() && pool_addresses.len() > 1 { - // Multiple upstreams - use new_multi - match router::Router::new_multi( - pool_addresses.clone(), - None, // setup_connection_msg - None, // timer - ARGS.round_robin - ) { - Ok(router) => router, - Err(e) => { - error!("Failed to create multi-upstream router: {}", e); - std::process::exit(1); + let mut router = if !pool_addresses.is_empty() && pool_addresses.len() > 1 { + // Multiple upstreams - use new_multi + match router::Router::new_multi( + pool_addresses.clone(), + None, // setup_connection_msg + None, // timer + ARGS.round_robin, + ) { + Ok(router) => router, + Err(e) => { + error!("Failed to create multi-upstream router: {}", e); + std::process::exit(1); + } } - } -} else { - // Single upstream - use regular constructor - let pool_socket_addresses: Vec = pool_addresses.iter().map(|(addr, _)| *addr).collect(); - let auth_pub_k = pool_addresses.first().expect("No pool addresses available").1.clone(); - - router::Router::new( - pool_socket_addresses, - auth_pub_k, - None, // setup_connection_msg - None // timer - ) -}; + } else { + // Single upstream - use regular constructor + let pool_socket_addresses: Vec = + pool_addresses.iter().map(|(addr, _)| *addr).collect(); + let auth_pub_k = pool_addresses + .first() + .expect("No pool addresses available") + .1; + + router::Router::new( + pool_socket_addresses, + auth_pub_k, + None, // setup_connection_msg + None, // timer + ) + }; let epsilon = Duration::from_millis(10); let best_upstream = router.select_pool_connect().await; @@ -459,34 +465,42 @@ async fn monitor( ) -> Reconnect { let mut should_check_upstreams_latency = 0; let mut distribution_check_counter = 0; - + loop { if distribution_check_counter >= 100 { distribution_check_counter = 0; - + // Add debug output to confirm this code runs info!("Generating hashrate distribution report..."); - + info!("Hashrate Distribution Report:"); - info!("Total hashrate: {}", HashUnit::format_value(*EXPECTED_SV1_HASHPOWER)); - + info!( + "Total hashrate: {}", + HashUnit::format_value(*EXPECTED_SV1_HASHPOWER) + ); + // Get upstream info directly from ProxyState if ARGS.round_robin { // In round-robin mode, check how many active upstreams we have let upstream_connections = ProxyState::get_upstream_connections(); let active_count = upstream_connections.len(); - + if active_count > 0 { let hashrate_per_upstream = *EXPECTED_SV1_HASHPOWER / active_count as f32; info!("Round-robin mode: {} active upstreams", active_count); - info!("Each upstream allocated: {}", HashUnit::format_value(hashrate_per_upstream)); - + info!( + "Each upstream allocated: {}", + HashUnit::format_value(hashrate_per_upstream) + ); + // List each upstream and its allocated hashrate for (upstream_id, addr, _auth_key) in upstream_connections { - info!("Upstream {}: {} - allocated {}", - upstream_id, - addr, - HashUnit::format_value(hashrate_per_upstream)); + info!( + "Upstream {}: {} - allocated {}", + upstream_id, + addr, + HashUnit::format_value(hashrate_per_upstream) + ); } } else { info!("No active upstreams found"); @@ -495,9 +509,11 @@ async fn monitor( // In latency-based mode info!("Latency-based mode: Using best upstream"); if let Some(current_addr) = router.get_current_upstream() { - info!("Current upstream: {} - allocated {}", - current_addr, - HashUnit::format_value(*EXPECTED_SV1_HASHPOWER)); + info!( + "Current upstream: {} - allocated {}", + current_addr, + HashUnit::format_value(*EXPECTED_SV1_HASHPOWER) + ); } else { info!("No upstream currently selected"); } @@ -557,13 +573,12 @@ async fn monitor( // Increment the counters should_check_upstreams_latency += 1; distribution_check_counter += 1; - + // Make sure this line exists to give time for the loop tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; } } - /// Parses a hashrate string (e.g., "10T", "2.5P", "500E") into an f32 value in h/s. fn parse_hashrate(hashrate_str: &str) -> Result { let hashrate_str = hashrate_str.trim(); @@ -597,8 +612,6 @@ fn parse_hashrate(hashrate_str: &str) -> Result { Ok(hashrate) } - - pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream diff --git a/src/proxy_state.rs b/src/proxy_state.rs index 402a624f..af14c2ec 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -88,6 +88,7 @@ pub enum UpstreamType { } /// Create an UpstreamConnection struct to store connection info +#[allow(dead_code)] #[derive(Debug, Clone)] pub struct UpstreamConnection { pub url: String, @@ -380,13 +381,15 @@ impl ProxyState { } /// Get the next upstream in round-robin fashion - pub fn get_next_upstream() -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + pub fn get_next_upstream( + ) -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { let mut result = None; if PROXY_STATE .safe_lock(|state| { // Get IDs of all connected upstreams - let active_upstreams: Vec<&String> = state.upstream_connections + let active_upstreams: Vec<&String> = state + .upstream_connections .iter() .filter(|(_, conn)| conn.is_connected) .map(|(id, _)| id) @@ -403,15 +406,12 @@ impl ProxyState { let id = active_upstreams[state.current_upstream_index].clone(); if let Some(conn) = state.upstream_connections.get(&id) { - result = Some(( - id.clone(), - conn.address, - conn.auth_key.clone() - )); + result = Some((id.clone(), conn.address, conn.auth_key)); } // Update index for next call - state.current_upstream_index = (state.current_upstream_index + 1) % active_upstreams.len(); + state.current_upstream_index = + (state.current_upstream_index + 1) % active_upstreams.len(); }) .is_err() { @@ -423,13 +423,15 @@ impl ProxyState { } /// Get the hashrate for a specific upstream (equal distribution) + #[allow(dead_code)] pub fn get_hashrate_for_upstream(id: Option<&str>) -> f32 { let mut hashrate = 0.0; if PROXY_STATE .safe_lock(|state| { // Count active connections - let active_count = state.upstream_connections + let active_count = state + .upstream_connections .values() .filter(|conn| conn.is_connected) .count(); @@ -462,6 +464,7 @@ impl ProxyState { } /// Record a share submission to an upstream + #[allow(dead_code)] pub fn record_share_submission(upstream_id: &str) { if PROXY_STATE .safe_lock(|state| { @@ -478,6 +481,7 @@ impl ProxyState { } /// Record a share acceptance from an upstream + #[allow(dead_code)] pub fn record_share_acceptance(upstream_id: &str) { if PROXY_STATE .safe_lock(|state| { @@ -493,15 +497,17 @@ impl ProxyState { } /// Get all upstream connections - pub fn get_upstream_connections() -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + pub fn get_upstream_connections( + ) -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { let mut connections = Vec::new(); if PROXY_STATE .safe_lock(|state| { - connections = state.upstream_connections + connections = state + .upstream_connections .iter() .filter(|(_, conn)| conn.is_connected) - .map(|(id, conn)| (id.clone(), conn.address, conn.auth_key.clone())) + .map(|(id, conn)| (id.clone(), conn.address, conn.auth_key)) .collect(); }) .is_err() diff --git a/src/router/mod.rs b/src/router/mod.rs index 8705636b..062bcbbb 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -17,7 +17,7 @@ use tokio::{ watch, }, }; -use tracing::{error, info,warn}; +use tracing::{error, info}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, @@ -53,8 +53,8 @@ impl Router { ) -> Self { let (latency_tx, latency_rx) = watch::channel(None); // Create auth_keys vector with the same key for all addresses - let auth_keys = vec![auth_pub_k.clone(); pool_addresses.len()]; - + let auth_keys = vec![auth_pub_k; pool_addresses.len()]; + Self { pool_addresses, auth_keys, @@ -67,7 +67,7 @@ impl Router { use_round_robin: false, // Default to latency-based selection } } - + /// Creates a new Router with multiple upstream addresses and auth keys pub fn new_multi( pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, @@ -78,19 +78,19 @@ impl Router { if pool_address_keys.is_empty() { return Err("Cannot create router with empty pool_address_keys"); } - + let (latency_tx, latency_rx) = watch::channel(None); let mut addresses = Vec::new(); let mut keys = Vec::new(); - + for (addr, key) in &pool_address_keys { addresses.push(*addr); - keys.push(key.clone()); + keys.push(*key); } - + // Use the first key as the default auth_pub_k - let default_key = pool_address_keys[0].1.clone(); - + let default_key = pool_address_keys[0].1; + // Register all upstreams in the ProxyState for (idx, (addr, key)) in pool_address_keys.iter().enumerate() { let id = format!("upstream-{}", idx); @@ -98,11 +98,11 @@ impl Router { id.clone(), format!("{:?}", addr), *addr, - key.clone(), + *key, crate::proxy_state::UpstreamType::JDCMiningUpstream, ); } - + Ok(Self { pool_addresses: addresses, auth_keys: keys, @@ -115,12 +115,22 @@ impl Router { use_round_robin, }) } - + /// Enable or disable round-robin upstream selection + #[allow(dead_code)] pub fn set_round_robin(&mut self, enabled: bool) { self.use_round_robin = enabled; } + // For compatibility with existing code + #[allow(dead_code)] + pub fn is_current_upstream(&self, addr: &std::net::SocketAddr) -> bool { + if let Some(current) = self.current_pool { + return current == *addr; + } + false + } + /// Internal function to select pool with the least latency. async fn select_pool(&self) -> Option<(SocketAddr, Duration)> { // If round-robin is enabled, use ProxyState to get the next upstream @@ -131,7 +141,7 @@ impl Router { return Some((addr, Duration::from_millis(0))); } } - + // Fall back to latency-based selection let mut best_pool = None; let mut least_latency = Duration::MAX; @@ -151,36 +161,41 @@ impl Router { /// Select the best pool for connection pub async fn select_pool_connect(&mut self) -> Option { info!("Selecting the best upstream"); - + if self.use_round_robin { // Check if we have any registered upstreams in ProxyState if let Some((id, addr, _)) = ProxyState::get_next_upstream() { info!("Round-robin selected upstream {}: {:?}", id, addr); return Some(addr); } - + // If no upstreams registered in ProxyState yet, register them now if !self.pool_addresses.is_empty() { - for (idx, (addr, key)) in self.pool_addresses.iter().zip(self.auth_keys.iter()).enumerate() { + for (idx, (addr, key)) in self + .pool_addresses + .iter() + .zip(self.auth_keys.iter()) + .enumerate() + { let id = format!("upstream-{}", idx); ProxyState::add_upstream_connection( id, format!("{:?}", addr), *addr, - key.clone(), + *key, crate::proxy_state::UpstreamType::JDCMiningUpstream, ); // Mark all as initially connected ProxyState::set_upstream_connection_status(&format!("upstream-{}", idx), true); } - + // Now try again to get a round-robin selection if let Some((_, addr, _)) = ProxyState::get_next_upstream() { return Some(addr); } } } - + // Fall back to latency-based selection if let Some((pool, latency)) = self.select_pool().await { info!("Latency for upstream {:?} is {:?}", pool, latency); @@ -203,7 +218,7 @@ impl Router { return None; } } - + // Otherwise use latency-based selection if let Some((best_pool, best_pool_latency)) = self.select_pool().await { if let Some(current_pool) = self.current_pool { @@ -260,12 +275,12 @@ impl Router { self.current_pool = Some(pool); info!("Upstream {:?} selected", pool); - + // Find the matching auth key for this address let auth_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { - self.auth_keys[index].clone() + self.auth_keys[index] } else { - self.auth_pub_k.clone() + self.auth_pub_k }; match minin_pool_connection::connect_pool( @@ -278,19 +293,29 @@ impl Router { { Ok((send_to_pool, recv_from_pool, pool_connection_abortable)) => { // Update ProxyState with successful connection - let upstream_id = format!("upstream-{}", - self.pool_addresses.iter().position(|&a| a == pool).unwrap_or(0)); + let upstream_id = format!( + "upstream-{}", + self.pool_addresses + .iter() + .position(|&a| a == pool) + .unwrap_or(0) + ); ProxyState::set_upstream_connection_status(&upstream_id, true); - + Ok((send_to_pool, recv_from_pool, pool_connection_abortable)) } Err(e) => { // Update ProxyState with failed connection - let upstream_id = format!("upstream-{}", - self.pool_addresses.iter().position(|&a| a == pool).unwrap_or(0)); + let upstream_id = format!( + "upstream-{}", + self.pool_addresses + .iter() + .position(|&a| a == pool) + .unwrap_or(0) + ); ProxyState::set_upstream_connection_status(&upstream_id, false); - + Err(e) } } @@ -299,12 +324,13 @@ impl Router { /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { // Find the auth key for this address - let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool_address) { - self.auth_keys[index].clone() - } else { - self.auth_pub_k.clone() - }; - + let auth_pub_key = + if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool_address) { + self.auth_keys[index] + } else { + self.auth_pub_k + }; + let mut pool = PoolLatency::new(pool_address); let setup_connection_msg = self.setup_connection_msg.as_ref(); let timer = self.timer.as_ref(); @@ -373,7 +399,7 @@ impl Router { return None; } } - + // Otherwise use latency-based selection if let Some(best_pool) = self.select_pool_monitor(epsilon).await { if Some(best_pool) != self.current_pool { @@ -389,14 +415,6 @@ impl Router { pub fn get_current_upstream(&self) -> Option { self.current_pool } - - // For compatibility with existing code - pub fn is_current_upstream(&self, addr: &std::net::SocketAddr) -> bool { - if let Some(current) = self.current_pool { - return current == *addr; - } - false - } } /// Track latencies for various stages of pool connection setup. @@ -640,4 +658,4 @@ async fn initialize_mining_connections( let setup_connection_msg = setup_connection_msg.unwrap_or(get_mining_setup_connection_msg(true)); Ok((receiver, sender, setup_connection_msg)) -} \ No newline at end of file +} diff --git a/src/translator/downstream/diff_management.rs b/src/translator/downstream/diff_management.rs index 8931fd37..9eaae7cd 100644 --- a/src/translator/downstream/diff_management.rs +++ b/src/translator/downstream/diff_management.rs @@ -291,6 +291,14 @@ fn diff_to_sv1_message(diff: f64) -> ProxyResult<'static, (json_rpc::Message, [u Ok((message, target)) } +pub fn nearest_power_of_10(x: f32) -> f32 { + if x <= 0.0 { + return 0.001; + } + let exponent = x.log10().round() as i32; + 10f32.powi(exponent) +} + #[cfg(test)] mod test { use super::super::super::upstream::diff_management::UpstreamDifficultyConfig; @@ -314,7 +322,7 @@ mod test { let initial_nominal_hashrate = dbg!(measure_hashrate(10)); let target = match roles_logic_sv2::utils::hash_rate_to_target( initial_nominal_hashrate, - expected_shares_per_minute.into(), + expected_shares_per_minute, ) { Ok(target) => target, Err(_) => panic!(), @@ -334,7 +342,7 @@ mod test { let calculated_share_per_min = count as f32 / (elapsed.as_secs_f32() / 60.0); // This is the error margin for a confidence of 99% given the expect number of shares per // minute TODO the review the math under it - let error_margin = get_error(expected_shares_per_minute.into()); + let error_margin = get_error(expected_shares_per_minute); let error = (dbg!(calculated_share_per_min) - dbg!(expected_shares_per_minute as f32)).abs(); assert!( @@ -371,8 +379,8 @@ mod test { let elapsed_secs = start_time.elapsed().as_secs_f64(); let hashrate = hashes as f64 / elapsed_secs; - let nominal_hash_rate = hashrate; - nominal_hash_rate + + hashrate } fn hash(share: &mut [u8; 80]) -> Target { @@ -492,11 +500,3 @@ mod test { // a share but we try to updated the estimated hash power every 2 seconds and updated the // target consequentially this shuold start to provide shares within a normal amount of time } - -pub fn nearest_power_of_10(x: f32) -> f32 { - if x <= 0.0 { - return 0.001; - } - let exponent = x.log10().round() as i32; - 10f32.powi(exponent) -} From 2bd9ea1847c60a44ccd0873bf28a6fab74b9f8f1 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Wed, 4 Jun 2025 01:33:52 +0530 Subject: [PATCH 04/21] Refactor ProxyState and Router for Multi-Upstream Management - Removed deprecated methods and fields related to round-robin upstream selection in ProxyState and Router. - Introduced MultiUpstreamManager to handle multiple upstream connections concurrently. - Updated Router to utilize MultiUpstreamManager for managing upstream connections and sending messages. - Enhanced error handling and logging throughout the upstream connection process. - Added methods for retrieving connection statistics and managing upstream states. - Improved code organization and readability by separating concerns into the new MultiUpstreamManager module. --- config.toml | 2 +- src/main.rs | 569 +++++++++++---------------- src/proxy_state.rs | 293 ++++++++++---- src/router/mod.rs | 418 +++++++++++--------- src/router/multi_upstream_manager.rs | 369 +++++++++++++++++ src/shared/utils.rs | 7 + src/translator/mod.rs | 4 +- 7 files changed, 1069 insertions(+), 593 deletions(-) create mode 100644 src/router/multi_upstream_manager.rs diff --git a/config.toml b/config.toml index ae2d7491..01320f14 100644 --- a/config.toml +++ b/config.toml @@ -1,5 +1,5 @@ token="UR01TMkZv6Vs5zbwr0t6" -pool_addresses=["18.193.252.132:2000", "3.74.36.119:2000","18.193.252.132:2000", "3.74.36.119:2000"] +pool_addresses=["18.193.252.132:2000", "3.74.36.119:2000"] tp_address = "127.0.0.1:8442" interval = 120_000 delay = 540 diff --git a/src/main.rs b/src/main.rs index b1228857..2bf613dc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,14 +3,17 @@ use clap::{ArgAction, Parser}; use jemallocator::Jemalloc; use router::Router; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -#[cfg(not(target_os = "windows"))] +// Add these missing imports +use serde::{Deserialize, Serialize}; +use std::path::Path; + #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; use crate::shared::utils::AbortOnDrop; use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; -use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; +use proxy_state::{PoolState, ProxyState}; use std::{net::ToSocketAddrs, time::Duration}; use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; @@ -91,15 +94,38 @@ pub struct Args { adjustment_interval: u64, // New argument for multiple upstream servers #[clap(long = "upstream", short = 'u', action = ArgAction::Append)] - upstream_servers: Option>, - // Option to enable round-robin distribution - #[clap(long = "round-robin", short = 'r')] - round_robin: bool, + // In Args struct - #[clap(long = "monitor-hashrate", short = 'm')] - monitor_hashrate: bool, + // #[clap(long = "monitor-hashrate", short = 'm')] + // monitor_hashrate: bool, #[clap(long = "config", short = 'c')] pub config_file: Option, + + + /// Enable parallel upstream usage (sends to all upstreams simultaneously) + #[clap(long = "parallel", short = 'p')] + pub parallel: bool, +} + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct Config { + pub token: String, + pub pool_addresses: Vec, + pub tp_address: Option, + pub interval: Option, + pub delay: Option, + pub downstream_hashrate: Option, + pub loglevel: Option, + pub nc_loglevel: Option, + pub test: Option, +} + +impl Config { + pub fn from_file>(path: P) -> Result> { + let contents = std::fs::read_to_string(path)?; + let config: Config = toml::from_str(&contents)?; + Ok(config) + } } #[tokio::main] @@ -150,9 +176,9 @@ async fn main() { } } } else { + info!("No config file provided, using default settings"); None }; - // Get TOKEN from config or environment let _token = if let Some(ref config) = config { config.token.clone() @@ -207,13 +233,13 @@ async fn main() { std::process::exit(1); }; - // Use test public key for now (you might want to add this to config) let pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key"); info!("Adding upstream server {}: {}", idx + 1, addr); pool_addresses.push((addr, pubkey)); + // ONLY add to ProxyState, don't create connections yet ProxyState::add_upstream_connection( format!("upstream-{}", idx), format!("config-pool-{}", idx + 1), @@ -222,363 +248,251 @@ async fn main() { crate::proxy_state::UpstreamType::JDCMiningUpstream, ); } - } else { - // Fallback to hard-coded addresses (your existing code) - info!("Using hard-coded pool addresses"); + } + else { + // Fallback to hard-coded address + info!("Using hard-coded fallback pool address"); + + let test_addr = match "3.74.36.119:2000".to_socket_addrs() { + Ok(mut addrs) => addrs.next().unwrap_or_else(|| { + error!("Failed to resolve fallback pool address"); + std::process::exit(1); + }), + Err(_) => { + error!("Invalid fallback pool address format"); + std::process::exit(1); + } + }; - let test_pubkey: Secp256k1PublicKey = - TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key"); + let test_pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY + .parse() + .expect("Invalid fallback public key"); - let test_addr = if let Ok(addr) = "3.74.36.119:2000".to_socket_addrs() { - addr.collect::>()[0] - } else { - "3.74.36.119:2000".parse().expect("Invalid IP address") - }; + info!("Adding fallback upstream server: {}", test_addr); + pool_addresses.push((test_addr, test_pubkey)); - // Add hard-coded pools - for i in 0..2 { - info!("Adding upstream server {}: {}", i + 1, test_addr); - pool_addresses.push((test_addr, test_pubkey)); - ProxyState::add_upstream_connection( - format!("upstream-{}", i), - format!("test-pool-{}", i + 1), - test_addr, - test_pubkey, - crate::proxy_state::UpstreamType::JDCMiningUpstream, - ); - } + ProxyState::add_upstream_connection( + "upstream-0".to_string(), + "fallback-pool-1".to_string(), + test_addr, + test_pubkey, + crate::proxy_state::UpstreamType::JDCMiningUpstream + ); } + // Direct verification info!("DIRECT VERIFICATION: Checking hashrate distribution"); let upstream_connections = ProxyState::get_upstream_connections(); let active_count = upstream_connections.len(); info!("Active upstreams: {}", active_count); - if active_count > 0 && ARGS.round_robin { - let hashrate_per_upstream = hashpower / active_count as f32; - info!( - "Round-robin hashrate per upstream: {}", - HashUnit::format_value(hashrate_per_upstream) - ); - - let upstream_connections: Vec<( - String, - std::net::SocketAddr, - key_utils::Secp256k1PublicKey, - )> = ProxyState::get_upstream_connections(); - - // And update the usage in the loop below: - for (upstream_id, addr, _auth_key) in upstream_connections { - info!( - "Upstream {}: {} - allocated {}", - upstream_id, - addr, - HashUnit::format_value(hashrate_per_upstream) - ); - } - } - - // Set all upstreams as initially connected - for (idx, (_addr, _)) in pool_addresses.iter().enumerate() { - let id = if idx == 0 && pool_addresses.len() == 1 { - "upstream-default".to_string() - } else { - format!("upstream-{}", idx) - }; - ProxyState::set_upstream_connection_status(&id, true); - } - - // Create the router based on whether we're using round-robin or not - let mut router = if !pool_addresses.is_empty() && pool_addresses.len() > 1 { - // Multiple upstreams - use new_multi + // Remove round-robin hashrate calculation + // if active_count > 0 && ARGS.round_robin { + // let hashrate_per_upstream = hashpower / active_count as f32; + // info!( + // "Round-robin hashrate per upstream: {}", + // HashUnit::format_value(hashrate_per_upstream) + // ); + // + // let upstream_connections: Vec<( + // String, + // std::net::SocketAddr, + // key_utils::Secp256k1PublicKey, + // )> = ProxyState::get_upstream_connections(); + // + // // And update the usage in the loop below: + // for (upstream_id, addr, _auth_key) in upstream_connections { + // info!( + // "Upstream {}: {} - allocated {}", + // upstream_id, + // addr, + // HashUnit::format_value(hashrate_per_upstream) + // ); + // } + // } + + // Create the router - always use multi upstream with parallel mode + let mut router = if !pool_addresses.is_empty() { + // Always use multi-upstream mode (even for single upstream) match router::Router::new_multi( pool_addresses.clone(), None, // setup_connection_msg None, // timer - ARGS.round_robin, - ) { - Ok(router) => router, + true, // Always enable parallel mode + ).await { + Ok(mut router) => { + if let Err(e) = router.initialize_upstream_connections().await { + error!("Failed to initialize upstream connections: {}", e); + std::process::exit(1); + } + router + }, Err(e) => { error!("Failed to create multi-upstream router: {}", e); std::process::exit(1); } } } else { - // Single upstream - use regular constructor - let pool_socket_addresses: Vec = - pool_addresses.iter().map(|(addr, _)| *addr).collect(); - let auth_pub_k = pool_addresses - .first() - .expect("No pool addresses available") - .1; - - router::Router::new( - pool_socket_addresses, - auth_pub_k, - None, // setup_connection_msg - None, // timer - ) + error!("No pool addresses configured. Please provide pool addresses via config file."); + std::process::exit(1); }; - let epsilon = Duration::from_millis(10); - let best_upstream = router.select_pool_connect().await; - initialize_proxy(&mut router, best_upstream, epsilon).await; - info!("exiting"); - tokio::time::sleep(tokio::time::Duration::from_secs(60)).await; + let epsilon = Duration::from_millis(100); + + // Always use multi-upstream mode + initialize_proxy(&mut router, None, epsilon).await; + + // Start monitoring task for multi-upstream + let router_clone = router.clone(); + tokio::spawn(async move { + monitor_multi_upstream(router_clone, epsilon).await; + }); + + // Keep main thread alive + loop { + tokio::time::sleep(Duration::from_secs(60)).await; + info!("Multi-upstream proxy running..."); + } } async fn initialize_proxy( router: &mut Router, - mut pool_addr: Option, - epsilon: Duration, + mut _pool_addr: Option, // Add underscore to indicate unused + _epsilon: Duration, // Add underscore to indicate unused ) { - loop { - // Initial setup for the proxy - let stats_sender = api::stats::StatsSender::new(); - - let (send_to_pool, recv_from_pool, pool_connection_abortable) = - match router.connect_pool(pool_addr).await { - Ok(connection) => connection, - Err(_) => { - error!("No upstream available. Retrying..."); - warn!("Are you using the correct TOKEN??"); - let mut secs = 10; - while secs > 0 { - tracing::warn!("Retrying in {} seconds...", secs); - tokio::time::sleep(Duration::from_secs(1)).await; - secs -= 1; - } - // Restart loop, esentially restarting proxy - continue; + // Add a static flag to prevent multiple downstream listeners + static DOWNSTREAM_STARTED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + + // For multi-upstream mode, don't loop - just set up and wait + if router.is_multi_upstream_enabled() { + info!("Multi-upstream mode: using aggregated message handling"); + + // Only start downstream listener once + if !DOWNSTREAM_STARTED.swap(true, std::sync::atomic::Ordering::SeqCst) { + info!("Starting downstream listener (multi-upstream mode)"); + + // Create channel for downstream connections + let (downstreams_tx, mut downstreams_rx) = tokio::sync::mpsc::channel(10); + + // Start the downstream listener + let _abort_handle = ingress::sv1_ingress::start_listen_for_downstream(downstreams_tx); + + // Handle incoming downstream connections + tokio::spawn(async move { + while let Some((send_to_downstream, recv_from_downstream, client_addr)) = downstreams_rx.recv().await { + info!("New downstream client connected: {}", client_addr); + + // Here you would typically start the translator for this downstream connection + // For now, we'll just log the connection + tokio::spawn(async move { + // Handle this specific downstream connection + // You'll need to integrate this with your translator logic + let mut recv = recv_from_downstream; + while let Some(message) = recv.recv().await { + info!("Received from downstream {}: {}", client_addr, message); + // Process the message and potentially send to upstreams + } + info!("Downstream client {} disconnected", client_addr); + }); } - }; - - let (downs_sv1_tx, downs_sv1_rx) = channel(10); - let sv1_ingress_abortable = ingress::sv1_ingress::start_listen_for_downstream(downs_sv1_tx); - - let (translator_up_tx, mut translator_up_rx) = channel(10); - let translator_abortable = - match translator::start(downs_sv1_rx, translator_up_tx, stats_sender.clone()).await { - Ok(abortable) => abortable, - Err(e) => { - error!("Impossible to initialize translator: {e}"); - // Impossible to start the proxy so we restart proxy - ProxyState::update_translator_state(TranslatorState::Down); - ProxyState::update_tp_state(TpState::Down); - return; + }); + } + + // Get aggregated receiver if available + if let Some(aggregated_receiver) = router.get_aggregated_receiver() { + tokio::spawn(async move { + let mut receiver = aggregated_receiver; + while let Some(message) = receiver.recv().await { + info!("Received aggregated message from upstream: {:?}", message); } - }; - - let (from_jdc_to_share_accounter_send, from_jdc_to_share_accounter_recv) = channel(10); - let (from_share_accounter_to_jdc_send, from_share_accounter_to_jdc_recv) = channel(10); - let (jdc_to_translator_sender, jdc_from_translator_receiver, _) = translator_up_rx - .recv() - .await - .expect("Translator failed before initialization"); - - let jdc_abortable: Option; - let share_accounter_abortable; - let tp = match TP_ADDRESS.safe_lock(|tp| tp.clone()) { - Ok(tp) => tp, - Err(e) => { - error!("TP_ADDRESS Mutex Corrupted: {e}"); - return; - } - }; + }); + } + + // Instead of looping, just wait indefinitely + info!("Multi-upstream proxy initialized successfully"); + return; // Exit the function, don't loop + } + + // For single upstream mode - just log and exit since we're not supporting it + error!("Single upstream mode is not supported in this version. Please use multi-upstream mode with --config option."); + std::process::exit(1); +} - if let Some(_tp_addr) = tp { - jdc_abortable = jd_client::start( - jdc_from_translator_receiver, - jdc_to_translator_sender, - from_share_accounter_to_jdc_recv, - from_jdc_to_share_accounter_send, - ) - .await; - if jdc_abortable.is_none() { - ProxyState::update_tp_state(TpState::Down); - }; - share_accounter_abortable = match share_accounter::start( - from_jdc_to_share_accounter_recv, - from_share_accounter_to_jdc_send, - recv_from_pool, - send_to_pool, - ) - .await - { - Ok(abortable) => abortable, - Err(_) => { - error!("Failed to start share_accounter"); - return; +async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { // Remove mut + let mut distribution_check_counter = 0; + + loop { + tokio::time::sleep(Duration::from_secs(10)).await; + distribution_check_counter += 1; + + if distribution_check_counter >= 10 { // Every 100 seconds + distribution_check_counter = 0; + + let total_hashrate = ProxyState::get_total_hashrate(); + + info!("=== Hashrate Distribution Report ==="); + info!("Total hashrate: {}", HashUnit::format_value(total_hashrate)); + + let (total_connections, active_connections) = router.get_connection_stats().await; + info!("Multi-upstream mode: {} total, {} active connections", total_connections, active_connections); + + if active_connections > 0 { + info!("Parallel mode: Using ALL upstreams simultaneously"); + info!("Total hashrate distributed across {} upstreams: {}", + active_connections, HashUnit::format_value(total_hashrate)); + + let active_upstreams = router.get_active_upstreams().await; + for upstream_id in active_upstreams { + info!("Upstream {}: receiving full hashrate {}", + upstream_id, HashUnit::format_value(total_hashrate)); } + } else { + warn!("No active upstream connections!"); } - } else { - jdc_abortable = None; - - share_accounter_abortable = match share_accounter::start( - jdc_from_translator_receiver, - jdc_to_translator_sender, - recv_from_pool, - send_to_pool, - ) - .await - { - Ok(abortable) => abortable, - Err(_) => { - error!("Failed to start share_accounter"); - return; - } - }; - }; - - // Collecting all abort handles - let mut abort_handles = vec![ - (pool_connection_abortable, "pool_connection".to_string()), - (sv1_ingress_abortable, "sv1_ingress".to_string()), - (translator_abortable, "translator".to_string()), - (share_accounter_abortable, "share_accounter".to_string()), - ]; - if let Some(jdc_handle) = jdc_abortable { - abort_handles.push((jdc_handle, "jdc".to_string())); + info!("====================================="); } - let server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); - match monitor(router, abort_handles, epsilon, server_handle).await { - Reconnect::NewUpstream(new_pool_addr) => { - ProxyState::update_proxy_state_up(); - pool_addr = Some(new_pool_addr); - continue; - } - Reconnect::NoUpstream => { - ProxyState::update_proxy_state_up(); - pool_addr = None; - continue; - } - }; } } - async fn monitor( router: &mut Router, abort_handles: Vec<(AbortOnDrop, std::string::String)>, epsilon: Duration, server_handle: tokio::task::JoinHandle<()>, ) -> Reconnect { - let mut should_check_upstreams_latency = 0; - let mut distribution_check_counter = 0; - + // Since we only support multi-upstream mode now, this monitor function + // is simplified to just handle multi-upstream monitoring loop { - if distribution_check_counter >= 100 { - distribution_check_counter = 0; - - // Add debug output to confirm this code runs - info!("Generating hashrate distribution report..."); - - info!("Hashrate Distribution Report:"); + tokio::time::sleep(Duration::from_secs(30)).await; + + // Generate periodic reports for multi-upstream + let (total_connections, active_connections) = router.get_connection_stats().await; + let total_hashrate = ProxyState::get_total_hashrate(); + + info!("Multi-upstream mode: {} total, {} active connections", total_connections, active_connections); + + if active_connections > 0 { + info!("Parallel mode: Using ALL upstreams simultaneously"); info!( - "Total hashrate: {}", - HashUnit::format_value(*EXPECTED_SV1_HASHPOWER) + "Total hashrate distributed across {} upstreams: {}", + active_connections, + HashUnit::format_value(total_hashrate) ); - - // Get upstream info directly from ProxyState - if ARGS.round_robin { - // In round-robin mode, check how many active upstreams we have - let upstream_connections = ProxyState::get_upstream_connections(); - let active_count = upstream_connections.len(); - - if active_count > 0 { - let hashrate_per_upstream = *EXPECTED_SV1_HASHPOWER / active_count as f32; - info!("Round-robin mode: {} active upstreams", active_count); - info!( - "Each upstream allocated: {}", - HashUnit::format_value(hashrate_per_upstream) - ); - - // List each upstream and its allocated hashrate - for (upstream_id, addr, _auth_key) in upstream_connections { - info!( - "Upstream {}: {} - allocated {}", - upstream_id, - addr, - HashUnit::format_value(hashrate_per_upstream) - ); - } - } else { - info!("No active upstreams found"); - } - } else { - // In latency-based mode - info!("Latency-based mode: Using best upstream"); - if let Some(current_addr) = router.get_current_upstream() { - info!( - "Current upstream: {} - allocated {}", - current_addr, - HashUnit::format_value(*EXPECTED_SV1_HASHPOWER) - ); - } else { - info!("No upstream currently selected"); - } - } - } - - // Check if a better upstream exist every 100 seconds - if should_check_upstreams_latency == 10 * 100 { - if let Some(new_upstream) = router.monitor_upstream(epsilon).await { - info!("Faster upstream detected. Reinitializing proxy..."); - drop(abort_handles); - server_handle.abort(); // abort server - - // Needs a little to time to drop - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; - return Reconnect::NewUpstream(new_upstream); - } - } - - // Monitor finished tasks - if let Some((_handle, name)) = abort_handles - .iter() - .find(|(handle, _name)| handle.is_finished()) - { - error!("Task {:?} finished, Closing connection", name); - for (handle, _name) in abort_handles { - drop(handle); - } - server_handle.abort(); // abort server - - // Check if the proxy state is down, and if so, reinitialize the proxy. - let is_proxy_down = ProxyState::is_proxy_down(); - if is_proxy_down.0 { - error!( - "Status: {:?}. Reinitializing proxy...", - is_proxy_down.1.unwrap_or("Proxy".to_string()) + + let active_upstreams = router.get_active_upstreams().await; + for upstream_id in active_upstreams { + info!( + "Upstream {}: receiving full hashrate {}", + upstream_id, + HashUnit::format_value(total_hashrate) ); - return Reconnect::NoUpstream; - } else { - return Reconnect::NoUpstream; } + } else { + warn!("No active upstream connections!"); + // In a real implementation, you might want to try reconnecting here } - - // Check if the proxy state is down, and if so, reinitialize the proxy. - let is_proxy_down = ProxyState::is_proxy_down(); - if is_proxy_down.0 { - error!( - "{:?} is DOWN. Reinitializing proxy...", - is_proxy_down.1.unwrap_or("Proxy".to_string()) - ); - drop(abort_handles); // Drop all abort handles - server_handle.abort(); // abort server - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; // Needs a little to time to drop - return Reconnect::NoUpstream; - } - - // Increment the counters - should_check_upstreams_latency += 1; - distribution_check_counter += 1; - - // Make sure this line exists to give time for the loop - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; } } - /// Parses a hashrate string (e.g., "10T", "2.5P", "500E") into an f32 value in h/s. fn parse_hashrate(hashrate_str: &str) -> Result { let hashrate_str = hashrate_str.trim(); @@ -656,26 +570,5 @@ impl HashUnit { } } } -use serde::{Deserialize, Serialize}; -use std::path::Path; -#[derive(Debug, Deserialize, Serialize, Clone)] -pub struct Config { - pub token: String, - pub pool_addresses: Vec, - pub tp_address: Option, - pub interval: Option, - pub delay: Option, - pub downstream_hashrate: Option, - pub loglevel: Option, - pub nc_loglevel: Option, - pub test: Option, -} -impl Config { - pub fn from_file>(path: P) -> Result> { - let contents = std::fs::read_to_string(path)?; - let config: Config = toml::from_str(&contents)?; - Ok(config) - } -} diff --git a/src/proxy_state.rs b/src/proxy_state.rs index af14c2ec..a469ecbc 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -265,55 +265,6 @@ impl ProxyState { } } - pub fn is_proxy_down() -> (bool, Option) { - let errors = Self::get_errors(); - if errors.is_ok() && errors.as_ref().unwrap().is_empty() { - (false, None) - } else { - let error_descriptions: Vec = - errors.iter().map(|e| format!("{:?}", e)).collect(); - (true, Some(error_descriptions.join(", "))) - } - } - - pub fn get_errors() -> Result, ()> { - let mut errors = Vec::new(); - if PROXY_STATE - .safe_lock(|state| { - if state.pool == PoolState::Down { - errors.push(ProxyStates::Pool(state.pool)); - } - if state.tp == TpState::Down { - errors.push(ProxyStates::Tp(state.tp)); - } - if state.jd == JdState::Down { - errors.push(ProxyStates::Jd(state.jd)); - } - if state.share_accounter == ShareAccounterState::Down { - errors.push(ProxyStates::ShareAccounter(state.share_accounter)); - } - if state.translator == TranslatorState::Down { - errors.push(ProxyStates::Translator(state.translator)); - } - if let Some(inconsistency) = state.inconsistency { - errors.push(ProxyStates::InternalInconsistency(inconsistency)); - } - if matches!(state.downstream, DownstreamState::Down(_)) { - errors.push(ProxyStates::Downstream(state.downstream.clone())); - } - if matches!(state.upstream, UpstreamState::Down(_)) { - errors.push(ProxyStates::Upstream(state.upstream.clone())); - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } else { - Ok(errors) - } - } - /// Add a new upstream connection pub fn add_upstream_connection( id: String, @@ -346,18 +297,14 @@ impl ProxyState { } } - /// Update connection status for an upstream - pub fn set_upstream_connection_status(id: &str, connected: bool) { + + + /// Set the total hashrate to be distributed among upstreams + pub fn set_total_hashrate(hashrate: f32) { + info!("Setting total hashrate to: {} h/s", hashrate); if PROXY_STATE .safe_lock(|state| { - if let Some(conn) = state.upstream_connections.get_mut(id) { - conn.is_connected = connected; - if connected { - info!("Upstream {} is now connected", id); - } else { - info!("Upstream {} is now disconnected", id); - } - } + state.total_hashrate = hashrate; }) .is_err() { @@ -366,18 +313,19 @@ impl ProxyState { } } - /// Set the total hashrate to be distributed among upstreams - pub fn set_total_hashrate(hashrate: f32) { - info!("Setting total hashrate to: {} h/s", hashrate); + /// Get the total hashrate + pub fn get_total_hashrate() -> f32 { + let mut hashrate = 0.0; if PROXY_STATE .safe_lock(|state| { - state.total_hashrate = hashrate; + hashrate = state.total_hashrate; }) .is_err() { error!("Global Proxy Mutex Corrupted"); std::process::exit(1); } + hashrate } /// Get the next upstream in round-robin fashion @@ -496,9 +444,193 @@ impl ProxyState { } } - /// Get all upstream connections - pub fn get_upstream_connections( - ) -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + + /// Update connection status for an upstream with timestamp + pub fn set_upstream_connection_status(id: &str, connected: bool) { + if PROXY_STATE + .safe_lock(|state| { + if let Some(conn) = state.upstream_connections.get_mut(id) { + conn.is_connected = connected; + if connected { + conn.last_used = Instant::now(); + info!("Upstream {} is now connected", id); + } else { + info!("Upstream {} is now disconnected", id); + } + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Get connection count + pub fn get_upstream_connection_count() -> usize { + let mut count = 0; + + if PROXY_STATE + .safe_lock(|state| { + count = state.upstream_connections.len(); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + count + } + + /// Get active connection count + pub fn get_active_upstream_count() -> usize { + let mut count = 0; + + if PROXY_STATE + .safe_lock(|state| { + count = state.upstream_connections + .values() + .filter(|conn| conn.is_connected) + .count(); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + count + } + + /// Update upstream shares + pub fn update_upstream_shares(upstream_id: &str, submitted: u64, accepted: u64) { + if PROXY_STATE + .safe_lock(|state| { + if let Some(conn) = state.upstream_connections.get_mut(upstream_id) { + conn.shares_submitted += submitted; + conn.shares_accepted += accepted; + conn.last_used = Instant::now(); + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Remove an upstream connection + pub fn remove_upstream_connection(upstream_id: &str) { + info!("Removing upstream connection: {}", upstream_id); + + if PROXY_STATE + .safe_lock(|state| { + state.upstream_connections.remove(upstream_id); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + } + + /// Check if proxy is down + pub fn is_proxy_down() -> (bool, Option) { + let errors = Self::get_errors(); + if errors.is_ok() && errors.as_ref().unwrap().is_empty() { + (false, None) + } else { + let error_descriptions: Vec = + errors.iter().map(|e| format!("{:?}", e)).collect(); + (true, Some(error_descriptions.join(", "))) + } + } + + pub fn get_errors() -> Result, ()> { + let mut errors = Vec::new(); + if PROXY_STATE + .safe_lock(|state| { + if state.pool == PoolState::Down { + errors.push(ProxyStates::Pool(state.pool)); + } + if state.tp == TpState::Down { + errors.push(ProxyStates::Tp(state.tp)); + } + if state.jd == JdState::Down { + errors.push(ProxyStates::Jd(state.jd)); + } + if state.share_accounter == ShareAccounterState::Down { + errors.push(ProxyStates::ShareAccounter(state.share_accounter)); + } + if state.translator == TranslatorState::Down { + errors.push(ProxyStates::Translator(state.translator)); + } + if let Some(inconsistency) = state.inconsistency { + errors.push(ProxyStates::InternalInconsistency(inconsistency)); + } + if matches!(state.downstream, DownstreamState::Down(_)) { + errors.push(ProxyStates::Downstream(state.downstream.clone())); + } + if matches!(state.upstream, UpstreamState::Down(_)) { + errors.push(ProxyStates::Upstream(state.upstream.clone())); + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } else { + Ok(errors) + } + } + + + /// Get upstream statistics + pub fn get_upstream_stats() -> Vec<(String, bool, u64, u64)> { + let mut stats = Vec::new(); + + if PROXY_STATE + .safe_lock(|state| { + stats = state.upstream_connections + .iter() + .map(|(id, conn)| { + (id.clone(), conn.is_connected, conn.shares_submitted, conn.shares_accepted) + }) + .collect(); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + stats + } + + /// Get all upstream connections (including inactive ones) + pub fn get_all_upstream_connections() -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey, bool)> { + let mut connections = Vec::new(); + + if PROXY_STATE + .safe_lock(|state| { + connections = state + .upstream_connections + .iter() + .map(|(id, conn)| (id.clone(), conn.address, conn.auth_key, conn.is_connected)) + .collect(); + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + connections + } + + /// Get only active upstream connections (existing method is fine) + pub fn get_upstream_connections() -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { let mut connections = Vec::new(); if PROXY_STATE @@ -518,4 +650,33 @@ impl ProxyState { connections } + + /// Mark an upstream as inactive (disconnect it) + pub fn mark_upstream_inactive(upstream_id: &str) { + Self::set_upstream_connection_status(upstream_id, false); + } + + /// Get best upstream based on latency or other criteria + /// For now, just returns the first active upstream + pub fn get_best_upstream() -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + let mut result = None; + + if PROXY_STATE + .safe_lock(|state| { + // Find the first active upstream + for (id, conn) in &state.upstream_connections { + if conn.is_connected { + result = Some((id.clone(), conn.address, conn.auth_key)); + break; + } + } + }) + .is_err() + { + error!("Global Proxy Mutex Corrupted"); + std::process::exit(1); + } + + result + } } diff --git a/src/router/mod.rs b/src/router/mod.rs index 062bcbbb..7a985210 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -21,22 +21,49 @@ use tracing::{error, info}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, - proxy_state::ProxyState, // Add ProxyState import + proxy_state::ProxyState, shared::utils::AbortOnDrop, }; +// Add the new module +pub mod multi_upstream_manager; +use multi_upstream_manager::MultiUpstreamManager; + /// Router handles connection to Multiple upstreams. -#[derive(Clone)] pub struct Router { - pool_addresses: Vec, - auth_keys: Vec, // Store auth keys separately to map to addresses + pub pool_socket_addresses: Vec, + pub keys: Vec, pub current_pool: Option, - auth_pub_k: Secp256k1PublicKey, - setup_connection_msg: Option>, - timer: Option, - latency_tx: watch::Sender>, - pub latency_rx: watch::Receiver>, - use_round_robin: bool, // Add flag to control selection strategy + pub upstream_manager: Option, + pub aggregated_receiver: Option>>, + + // Keep these fields for backward compatibility with single upstream mode + pub auth_pub_k: Secp256k1PublicKey, + pub setup_connection_msg: Option>, + pub timer: Option, + pub latency_tx: tokio::sync::watch::Sender>, + pub latency_rx: tokio::sync::watch::Receiver>, + + // Remove round-robin fields entirely + // use_round_robin: bool, + // use_parallel: bool, // This will be determined by presence of upstream_manager +} +// Remove Clone derive since Receiver can't be cloned +impl Clone for Router { + fn clone(&self) -> Self { + Self { + pool_socket_addresses: self.pool_socket_addresses.clone(), + keys: self.keys.clone(), + current_pool: self.current_pool, + upstream_manager: self.upstream_manager.clone(), + aggregated_receiver: None, // Can't clone receiver + auth_pub_k: self.auth_pub_k, + setup_connection_msg: self.setup_connection_msg.clone(), + timer: self.timer, + latency_tx: self.latency_tx.clone(), + latency_rx: self.latency_rx.clone(), + } + } } impl Router { @@ -44,210 +71,129 @@ impl Router { pub fn new( pool_addresses: Vec, auth_pub_k: Secp256k1PublicKey, - // Configuration msg used to setup connection between client and pool - // If not, present `get_mining_setup_connection_msg()` is called to generated default values setup_connection_msg: Option>, - // Max duration for pool setup after which it times out. - // If None, default time of 5s is used. timer: Option, - ) -> Self { + ) -> Self{ let (latency_tx, latency_rx) = watch::channel(None); - // Create auth_keys vector with the same key for all addresses let auth_keys = vec![auth_pub_k; pool_addresses.len()]; - + Self { - pool_addresses, - auth_keys, + pool_socket_addresses: pool_addresses, + keys: auth_keys, current_pool: None, + upstream_manager: None, + aggregated_receiver: None, auth_pub_k, setup_connection_msg, timer, latency_tx, latency_rx, - use_round_robin: false, // Default to latency-based selection } } + /// Creates a new Router with multiple upstream addresses and auth keys - pub fn new_multi( + pub async fn new_multi( pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, setup_connection_msg: Option>, timer: Option, - use_round_robin: bool, + _use_parallel: bool, // Parameter kept for compatibility but always use parallel ) -> Result { - if pool_address_keys.is_empty() { - return Err("Cannot create router with empty pool_address_keys"); - } - - let (latency_tx, latency_rx) = watch::channel(None); - let mut addresses = Vec::new(); - let mut keys = Vec::new(); + let pool_socket_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); + let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); - for (addr, key) in &pool_address_keys { - addresses.push(*addr); - keys.push(*key); - } + // Create with parallel mode enabled + let (upstream_manager, aggregated_receiver) = MultiUpstreamManager::new(true); + + // Use first key as default auth key for backward compatibility + let auth_pub_k = keys.first().copied().ok_or("No authentication keys provided")?; - // Use the first key as the default auth_pub_k - let default_key = pool_address_keys[0].1; - - // Register all upstreams in the ProxyState - for (idx, (addr, key)) in pool_address_keys.iter().enumerate() { - let id = format!("upstream-{}", idx); - ProxyState::add_upstream_connection( - id.clone(), - format!("{:?}", addr), - *addr, - *key, - crate::proxy_state::UpstreamType::JDCMiningUpstream, - ); - } + let (latency_tx, latency_rx) = watch::channel(None); Ok(Self { - pool_addresses: addresses, - auth_keys: keys, + pool_socket_addresses, + keys, current_pool: None, - auth_pub_k: default_key, + upstream_manager: Some(upstream_manager), + aggregated_receiver: Some(aggregated_receiver), + auth_pub_k, setup_connection_msg, timer, latency_tx, latency_rx, - use_round_robin, }) } + + // Remove round-robin methods + // pub fn set_round_robin(&mut self, enabled: bool) { ... } - /// Enable or disable round-robin upstream selection - #[allow(dead_code)] - pub fn set_round_robin(&mut self, enabled: bool) { - self.use_round_robin = enabled; + /// Check if the router is using the MultiUpstreamManager + pub fn is_multi_upstream_enabled(&self) -> bool { + self.upstream_manager.is_some() } - - // For compatibility with existing code - #[allow(dead_code)] - pub fn is_current_upstream(&self, addr: &std::net::SocketAddr) -> bool { - if let Some(current) = self.current_pool { - return current == *addr; - } - false + + /// Check if the router is using parallel mode - true when MultiUpstreamManager is present + pub fn is_parallel_mode(&self) -> bool { + self.upstream_manager.is_some() } - /// Internal function to select pool with the least latency. - async fn select_pool(&self) -> Option<(SocketAddr, Duration)> { - // If round-robin is enabled, use ProxyState to get the next upstream - if self.use_round_robin { - if let Some((id, addr, _)) = ProxyState::get_next_upstream() { - info!("Round-robin selected upstream {}: {:?}", id, addr); - // Return with a dummy zero latency - return Some((addr, Duration::from_millis(0))); + /// Checks for faster upstream and switches to it if found + pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { + // For multi-upstream mode, we don't switch since we use all simultaneously + if self.is_multi_upstream_enabled() { + // In parallel mode, we don't need to switch upstreams + // All upstreams are used simultaneously + return None; + } + + // For single upstream mode, check for better latency + if let Some(best_pool) = self.select_pool_monitor(epsilon).await { + if Some(best_pool) != self.current_pool { + info!("Switching to faster upstream {:?}", best_pool); + return Some(best_pool); } } + None + } - // Fall back to latency-based selection - let mut best_pool = None; - let mut least_latency = Duration::MAX; - - for &pool_addr in &self.pool_addresses { - if let Ok(latency) = self.get_latency(pool_addr).await { - if latency < least_latency { - least_latency = latency; - best_pool = Some(pool_addr) + /// Select the best pool for monitoring (checks latency differences) + async fn select_pool_monitor(&self, epsilon: Duration) -> Option { + if let Some(current_pool) = self.current_pool { + if let Ok(current_latency) = self.get_latency(current_pool).await { + // Check if there's a significantly better pool + for &pool_addr in &self.pool_socket_addresses { + if pool_addr != current_pool { + if let Ok(pool_latency) = self.get_latency(pool_addr).await { + // Switch if the new pool is significantly faster + if current_latency > pool_latency + epsilon { + info!( + "Found faster upstream: {:?} (latency: {:?}) vs current {:?} (latency: {:?})", + pool_addr, pool_latency, current_pool, current_latency + ); + return Some(pool_addr); + } + } + } } } } - - best_pool.map(|pool| (pool, least_latency)) + None } /// Select the best pool for connection pub async fn select_pool_connect(&mut self) -> Option { info!("Selecting the best upstream"); - if self.use_round_robin { - // Check if we have any registered upstreams in ProxyState - if let Some((id, addr, _)) = ProxyState::get_next_upstream() { - info!("Round-robin selected upstream {}: {:?}", id, addr); - return Some(addr); - } - - // If no upstreams registered in ProxyState yet, register them now - if !self.pool_addresses.is_empty() { - for (idx, (addr, key)) in self - .pool_addresses - .iter() - .zip(self.auth_keys.iter()) - .enumerate() - { - let id = format!("upstream-{}", idx); - ProxyState::add_upstream_connection( - id, - format!("{:?}", addr), - *addr, - *key, - crate::proxy_state::UpstreamType::JDCMiningUpstream, - ); - // Mark all as initially connected - ProxyState::set_upstream_connection_status(&format!("upstream-{}", idx), true); - } - - // Now try again to get a round-robin selection - if let Some((_, addr, _)) = ProxyState::get_next_upstream() { - return Some(addr); - } - } - } - - // Fall back to latency-based selection + // Remove round-robin logic - only use latency-based selection if let Some((pool, latency)) = self.select_pool().await { info!("Latency for upstream {:?} is {:?}", pool, latency); - self.latency_tx.send_replace(Some(latency)); // update latency + self.latency_tx.send_replace(Some(latency)); Some(pool) } else { None } } - /// Select the best pool for monitoring - async fn select_pool_monitor(&self, epsilon: Duration) -> Option { - // If using round-robin, just get the next upstream - if self.use_round_robin { - if let Some((_, addr, _)) = ProxyState::get_next_upstream() { - // Only switch if it's different from the current one - if Some(addr) != self.current_pool { - return Some(addr); - } - return None; - } - } - - // Otherwise use latency-based selection - if let Some((best_pool, best_pool_latency)) = self.select_pool().await { - if let Some(current_pool) = self.current_pool { - if best_pool == current_pool { - return None; - } - let current_latency = match self.get_latency(current_pool).await { - Ok(latency) => latency, - Err(e) => { - error!("Failed to get latency: {:?}", e); - return None; - } - }; - if best_pool_latency < (current_latency - epsilon) { - info!( - "Found faster pool: {:?} with latency {:?}", - best_pool, best_pool_latency - ); - return Some(best_pool); - } else { - return None; - } - } else { - return Some(best_pool); - } - } - None - } - /// Selects the best upstream and connects to. /// Uses minin_pool_connection::connect_pool pub async fn connect_pool( @@ -276,9 +222,9 @@ impl Router { info!("Upstream {:?} selected", pool); - // Find the matching auth key for this address - let auth_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { - self.auth_keys[index] + // Find the matching auth key for this address - fix field name + let auth_key = if let Some(index) = self.pool_socket_addresses.iter().position(|&a| a == pool) { + self.keys[index] } else { self.auth_pub_k }; @@ -295,7 +241,7 @@ impl Router { // Update ProxyState with successful connection let upstream_id = format!( "upstream-{}", - self.pool_addresses + self.pool_socket_addresses .iter() .position(|&a| a == pool) .unwrap_or(0) @@ -309,7 +255,7 @@ impl Router { // Update ProxyState with failed connection let upstream_id = format!( "upstream-{}", - self.pool_addresses + self.pool_socket_addresses .iter() .position(|&a| a == pool) .unwrap_or(0) @@ -323,10 +269,10 @@ impl Router { /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { - // Find the auth key for this address + // Find the auth key for this address - fix field names let auth_pub_key = - if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool_address) { - self.auth_keys[index] + if let Some(index) = self.pool_socket_addresses.iter().position(|&a| a == pool_address) { + self.keys[index] } else { self.auth_pub_k }; @@ -387,33 +333,133 @@ impl Router { Ok(sum_of_latencies) } - /// Checks for faster upstream switch to it if found - pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { - // If using round-robin, just get the next upstream - if self.use_round_robin { - if let Some((_, addr, _)) = ProxyState::get_next_upstream() { - if Some(addr) != self.current_pool { - info!("Switching to next round-robin upstream {:?}", addr); - return Some(addr); + pub fn get_current_upstream(&self) -> Option { + self.current_pool + } + + /// Initialize upstream connections for the manager + pub async fn initialize_upstream_connections(&mut self) -> Result<(), String> { + if let Some(ref manager) = self.upstream_manager { + info!("Initializing {} upstream connections", self.pool_socket_addresses.len()); + + // Add each unique upstream only once - fix field names + for (idx, (addr, key)) in self.pool_socket_addresses.iter().zip(self.keys.iter()).enumerate() { + let id = format!("upstream-{}", idx); + + info!("Adding upstream {}: {} ({})", id, addr, key); + + if let Err(e) = manager.add_upstream( + id.clone(), + *addr, + *key, + self.setup_connection_msg.clone(), + self.timer, + ).await { + error!("Failed to initialize upstream {}: {}", addr, e); + // Continue with other upstreams even if one fails + } else { + info!("Successfully initialized upstream {}: {}", id, addr); } - return None; } + + // Wait a moment for connections to stabilize + tokio::time::sleep(Duration::from_millis(100)).await; + + let (total, active) = self.get_connection_stats().await; + info!("Initialized {} upstream connections ({} active)", total, active); } + Ok(()) + } - // Otherwise use latency-based selection - if let Some(best_pool) = self.select_pool_monitor(epsilon).await { - if Some(best_pool) != self.current_pool { - info!("Switching to faster upstream {:?}", best_pool); - return Some(best_pool); + /// Send a message to a specific upstream using the manager + pub async fn send_to_upstream(&self, upstream_id: &str, message: PoolExtMessages<'static>) -> Result<(), String> { + if let Some(ref manager) = self.upstream_manager { + manager.send_to_upstream(upstream_id, message).await + } else { + Err("MultiUpstreamManager not initialized".to_string()) + } + } + + /// Send a message to the next upstream (round-robin or best available) + pub async fn send_to_next_upstream(&self, message: PoolExtMessages<'static>) -> Result<(), String> { + if let Some(ref manager) = self.upstream_manager { + manager.send_to_next_upstream(message).await + } else { + Err("MultiUpstreamManager not initialized".to_string()) + } + } + + /// Broadcast a message to all active upstreams + pub async fn broadcast_to_all_upstreams(&self, message: PoolExtMessages<'static>) -> Vec { + if let Some(ref manager) = self.upstream_manager { + manager.broadcast(message).await + } else { + vec!["MultiUpstreamManager not initialized".to_string()] + } + } + + /// Get aggregated receiver for multi-upstream mode + pub fn get_aggregated_receiver(&mut self) -> Option>> { + self.aggregated_receiver.take() + } + + /// Get connection statistics for multi-upstream mode + pub async fn get_connection_stats(&self) -> (usize, usize) { + if let Some(ref manager) = self.upstream_manager { + // Get stats from the manager + let upstreams = manager.get_upstreams().await; + let total = upstreams.len(); + let active = upstreams.values().filter(|u| u.is_active).count(); + (total, active) + } else { + (self.pool_socket_addresses.len(), 1) // Single upstream mode + } + } + + /// Get list of active upstream IDs + pub async fn get_active_upstreams(&self) -> Vec { + if let Some(ref manager) = self.upstream_manager { + manager.get_active_upstream_ids().await + } else { + vec!["upstream-0".to_string()] // Single upstream mode + } + } + + /// Switch to a specific upstream by ID + pub async fn switch_to_upstream(&mut self, upstream_id: &str) -> Result<(), String> { + // Find the upstream address by ID + if let Ok(idx) = upstream_id.strip_prefix("upstream-") + .and_then(|s| s.parse::().ok()) + .ok_or("Invalid upstream ID format".to_string()) + { + if idx < self.pool_socket_addresses.len() { // Fix field name + let addr = self.pool_socket_addresses[idx]; // Fix field name + self.current_pool = Some(addr); + info!("Switched to upstream {}: {:?}", upstream_id, addr); + Ok(()) } else { - return None; + Err(format!("Upstream index {} out of range", idx)) } + } else { + Err("Invalid upstream ID format".to_string()) } - None } - pub fn get_current_upstream(&self) -> Option { - self.current_pool + /// Select the best pool based on latency + async fn select_pool(&mut self) -> Option<(SocketAddr, Duration)> { + let mut best_pool = None; + let mut best_latency = Duration::from_secs(u64::MAX); + + for &pool_addr in &self.pool_socket_addresses { + if let Ok(latency) = self.get_latency(pool_addr).await { + if latency < best_latency { + best_latency = latency; + best_pool = Some(pool_addr); + } + } + } + + best_pool.map(|pool| (pool, best_latency)) } } diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs new file mode 100644 index 00000000..d193e450 --- /dev/null +++ b/src/router/multi_upstream_manager.rs @@ -0,0 +1,369 @@ +use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; +use tokio::sync::{mpsc, RwLock, Mutex}; +use tracing::{error, info, warn}; +use key_utils::Secp256k1PublicKey; +use demand_share_accounting_ext::parser::PoolExtMessages; +use roles_logic_sv2::common_messages_sv2::SetupConnection; +use crate::{ + minin_pool_connection::{connect_pool}, + proxy_state::ProxyState, + shared::utils::AbortOnDrop, +}; +use futures::future::join_all; + +#[derive(Clone)] +pub struct UpstreamConnection { + pub id: String, + pub address: SocketAddr, + pub auth_key: Secp256k1PublicKey, + pub sender: mpsc::Sender>, + pub is_active: bool, + pub connection_handle: Option, +} + +#[derive(Clone)] +pub struct MultiUpstreamManager { + upstreams: Arc>>, + aggregated_sender: tokio::sync::mpsc::Sender>, + // Remove round-robin fields + // current_index: Arc>, + // use_round_robin: bool, +} + +impl MultiUpstreamManager { + pub fn new( + _use_parallel: bool, // Parameter kept for compatibility but always use parallel + ) -> (Self, tokio::sync::mpsc::Receiver>) { + let (sender, receiver) = tokio::sync::mpsc::channel(1000); + + ( + Self { + upstreams: Arc::new(Mutex::new(HashMap::new())), + aggregated_sender: sender, + // Remove round-robin fields + // current_index: Arc::new(Mutex::new(0)), + // use_round_robin: false, // Always use parallel + }, + receiver, + ) + } + + /// Add a new upstream connection + pub async fn add_upstream( + &self, + id: String, + address: SocketAddr, + auth_key: Secp256k1PublicKey, + setup_connection_msg: Option>, + timer: Option, + ) -> Result<(), String> { + info!("🔍 ADD_UPSTREAM CALLED: {} -> {}", id, address); + + let mut upstreams = self.upstreams.lock().await; + + // Check if upstream already exists + if upstreams.contains_key(&id) { + warn!("Upstream {} already exists, skipping", id); + return Ok(()); + } + + // Create upstream connection + let upstream_connection = UpstreamConnection { + id: id.clone(), + address, + auth_key, + is_active: false, // Start as inactive, will be set to true when connected + sender: self.aggregated_sender.clone(), // Use a proper sender here + connection_handle: None, // Set to None initially + }; + + upstreams.insert(id.clone(), upstream_connection); + info!("✅ Added upstream {} to manager", id); + + // Start connection task + let upstreams_clone = self.upstreams.clone(); + let sender_clone = self.aggregated_sender.clone(); + + tokio::spawn(async move { + Self::connect_upstream( + upstreams_clone, + id, + address, + auth_key, + setup_connection_msg, + timer, + sender_clone, + ).await + }); + + Ok(()) + } + + /// Connect to a specific upstream + async fn connect_upstream( + upstreams: Arc>>, + id: String, + address: SocketAddr, + auth_key: Secp256k1PublicKey, + setup_connection_msg: Option>, + timer: Option, + sender: tokio::sync::mpsc::Sender>, + ) { + info!("Connecting to upstream {}: {}", id, address); + + match crate::minin_pool_connection::connect_pool( + address, + auth_key, + setup_connection_msg, + timer, + ).await { + Ok((send_to_pool, recv_from_pool, _abort_handle)) => { + info!("Successfully connected to upstream {}: {}", id, address); + + // Update upstream status in both MultiUpstreamManager and ProxyState + { + let mut upstreams_lock = upstreams.lock().await; + if let Some(upstream) = upstreams_lock.get_mut(&id) { + upstream.is_active = true; + } + } + + // Also update ProxyState + crate::proxy_state::ProxyState::set_upstream_connection_status(&id, true); + + // Handle messages from this upstream + Self::handle_upstream_messages(id.clone(), recv_from_pool, sender).await; + + // If we reach here, connection was lost + warn!("Connection to upstream {} lost", id); + + // Mark as inactive + { + let mut upstreams_lock = upstreams.lock().await; + if let Some(upstream) = upstreams_lock.get_mut(&id) { + upstream.is_active = false; + } + } + crate::proxy_state::ProxyState::set_upstream_connection_status(&id, false); + } + Err(e) => { + error!("Failed to connect to upstream {}: {}", id, e); + + // Mark as inactive + let mut upstreams_lock = upstreams.lock().await; + if let Some(upstream) = upstreams_lock.get_mut(&id) { + upstream.is_active = false; + } + crate::proxy_state::ProxyState::set_upstream_connection_status(&id, false); + } + } + } + + /// Handle messages from an upstream + async fn handle_upstream_messages( + upstream_id: String, + mut receiver: tokio::sync::mpsc::Receiver>, + sender: tokio::sync::mpsc::Sender>, + ) { + info!("Starting message handler for upstream {}", upstream_id); + + while let Some(message) = receiver.recv().await { + if let Err(e) = sender.send(message).await { + error!("Failed to forward message from upstream {}: {}", upstream_id, e); + break; + } + } + + warn!("Message handler for upstream {} stopped", upstream_id); + } + + /// Send a message to a specific upstream + pub async fn send_to_upstream(&self, upstream_id: &str, message: PoolExtMessages<'static>) -> Result<(), String> { + let connections = self.upstreams.lock().await; + if let Some(connection) = connections.get(upstream_id) { + if connection.is_active { + connection.sender.send(message).await + .map_err(|e| format!("Failed to send to upstream {}: {}", upstream_id, e))?; + Ok(()) + } else { + Err(format!("Upstream {} is not active", upstream_id)) + } + } else { + Err(format!("Upstream {} not found", upstream_id)) + } + } + + /// Send a message to the next upstream using the configured strategy + pub async fn send_to_next_upstream(&self, message: PoolExtMessages<'static>) -> Result<(), String> { + info!("Broadcasting to all upstreams (parallel mode)"); + let results = self.broadcast(message).await; + + // Check if any sends were successful + let success_count = results.iter() + .filter(|result| result.contains("Success")) + .count(); + + if success_count > 0 { + info!("Successfully sent to {} upstreams", success_count); + Ok(()) + } else { + Err(format!("Failed to send to any upstreams: {:?}", results)) + } + } + + // Remove round-robin send method + // async fn send_round_robin(&self, message: PoolExtMessages<'static>) -> Result<(), String> { ... } + + // /// Send to the best upstream (first active one for now) + // async fn send_to_best_upstream(&self, message: PoolExtMessages<'static>) -> Result<(), String> { + // let connections = self.upstreams.lock().await; + // for connection in connections.values() { + // if connection.is_active { + // return connection.sender.send(message).await + // .map_err(|e| format!("Failed to send to upstream {}: {}", connection.id, e)); + // } + // } + // Err("No active upstreams available".to_string()) + // } + /// Connect to upstream with retry logic + async fn connect_upstream_with_retry( + upstreams: Arc>>, + id: String, + address: SocketAddr, + auth_key: Secp256k1PublicKey, + setup_connection_msg: Option>, + timer: Option, + sender: tokio::sync::mpsc::Sender>, + ) { + let mut retry_count = 0; + const MAX_RETRIES: u32 = 3; + const RETRY_DELAY: Duration = Duration::from_secs(5); + + while retry_count < MAX_RETRIES { + info!("Connecting to upstream {} (attempt {}/{}): {}", id, retry_count + 1, MAX_RETRIES, address); + + match crate::minin_pool_connection::connect_pool( + address, + auth_key, + setup_connection_msg.clone(), + timer, + ).await { + Ok((send_to_pool, recv_from_pool, _abort_handle)) => { + info!("Successfully connected to upstream {}: {}", id, address); + + // Update upstream status + { + let mut upstreams_lock = upstreams.lock().await; + if let Some(upstream) = upstreams_lock.get_mut(&id) { + upstream.is_active = true; + } + } + + // Handle messages from this upstream + Self::handle_upstream_messages(id.clone(), recv_from_pool, sender).await; + return; // Success, exit retry loop + } + Err(e) => { + error!("Failed to connect to upstream {} (attempt {}): {}", id, retry_count + 1, e); + retry_count += 1; + + if retry_count < MAX_RETRIES { + info!("Retrying connection to {} in {:?}", id, RETRY_DELAY); + tokio::time::sleep(RETRY_DELAY).await; + } + } + } + } + + error!("Failed to connect to upstream {} after {} attempts", id, MAX_RETRIES); + + // Mark as inactive after all retries failed + let mut upstreams_lock = upstreams.lock().await; + if let Some(upstream) = upstreams_lock.get_mut(&id) { + upstream.is_active = false; + } + } + /// Broadcast to all active upstreams (parallel execution) + pub async fn broadcast(&self, message: PoolExtMessages<'static>) -> Vec { + let upstreams = self.upstreams.lock().await; + let mut results = Vec::new(); + + if upstreams.is_empty() { + results.push("No upstreams configured".to_string()); + return results; + } + + let active_upstreams: Vec<_> = upstreams.values() + .filter(|upstream| upstream.is_active) + .collect(); + + if active_upstreams.is_empty() { + results.push("No active upstreams available".to_string()); + return results; + } + + for upstream in active_upstreams { + match upstream.sender.send(message.clone()).await { + Ok(_) => { + results.push(format!("Sent to {}: Success", upstream.id)); + } + Err(e) => { + results.push(format!("Sent to {}: Failed - {}", upstream.id, e)); + // Mark upstream as inactive if send fails + warn!("Upstream {} send failed, may need reconnection: {}", upstream.id, e); + } + } + } + + results + } + + /// Get all upstreams (for stats) + pub async fn get_upstreams(&self) -> std::collections::HashMap { + let upstreams = self.upstreams.lock().await; + upstreams.clone() + } + + /// Get list of active upstream IDs + pub async fn get_active_upstream_ids(&self) -> Vec { + let upstreams = self.upstreams.lock().await; + upstreams.values() + .filter(|upstream| upstream.is_active) + .map(|upstream| upstream.id.clone()) + .collect() + } + + /// Mark an upstream as inactive + pub async fn mark_upstream_inactive(&self, upstream_id: &str) { + let mut connections = self.upstreams.lock().await; + if let Some(connection) = connections.get_mut(upstream_id) { + connection.is_active = false; + ProxyState::set_upstream_connection_status(upstream_id, false); + info!("Marked upstream {} as inactive", upstream_id); + } + } + + /// Remove an upstream connection + pub async fn remove_upstream(&self, upstream_id: &str) { + let mut connections = self.upstreams.lock().await; + if let Some(connection) = connections.remove(upstream_id) { + if let Some(handle) = connection.connection_handle { + drop(handle); // This will abort the connection + } + ProxyState::set_upstream_connection_status(upstream_id, false); + info!("Removed upstream {}", upstream_id); + } + } + + /// Get connection count + pub async fn connection_count(&self) -> usize { + let connections = self.upstreams.lock().await; + connections.len() + } + + /// Get active connection count + pub async fn active_connection_count(&self) -> usize { + let connections = self.upstreams.lock().await; + connections.values().filter(|conn| conn.is_active).count() + } +} \ No newline at end of file diff --git a/src/shared/utils.rs b/src/shared/utils.rs index 476cbbf0..c3bc7397 100644 --- a/src/shared/utils.rs +++ b/src/shared/utils.rs @@ -8,6 +8,13 @@ use tokio::task::JoinHandle; pub struct AbortOnDrop { abort_handle: AbortHandle, } +impl Clone for AbortOnDrop { + fn clone(&self) -> Self { + Self { + abort_handle: self.abort_handle.clone(), + } + } +} impl AbortOnDrop { pub fn new(handle: JoinHandle) -> Self { diff --git a/src/translator/mod.rs b/src/translator/mod.rs index b861b46f..a5c6ed35 100644 --- a/src/translator/mod.rs +++ b/src/translator/mod.rs @@ -1,8 +1,8 @@ mod downstream; mod error; -mod proxy; -mod upstream; +pub mod proxy; // Make public +pub mod upstream; // Make public mod utils; use bitcoin::Address; From a02ba6274d14371899aa9a389a259cff044cf595 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Wed, 4 Jun 2025 01:42:26 +0530 Subject: [PATCH 05/21] Enhance multi-upstream functionality: update pool addresses, add detailed connection stats, and improve hashrate reporting --- config.toml | 2 +- src/main.rs | 85 +++++++++++++++++++++++++++++++++++++---------- src/router/mod.rs | 24 +++++++++++++ 3 files changed, 92 insertions(+), 19 deletions(-) diff --git a/config.toml b/config.toml index 01320f14..7da55468 100644 --- a/config.toml +++ b/config.toml @@ -1,5 +1,5 @@ token="UR01TMkZv6Vs5zbwr0t6" -pool_addresses=["18.193.252.132:2000", "3.74.36.119:2000"] +pool_addresses=["3.74.36.119:2000"] tp_address = "127.0.0.1:8442" interval = 120_000 delay = 540 diff --git a/src/main.rs b/src/main.rs index 2bf613dc..b7b8e945 100644 --- a/src/main.rs +++ b/src/main.rs @@ -343,6 +343,49 @@ async fn main() { // Always use multi-upstream mode initialize_proxy(&mut router, None, epsilon).await; + // Wait a moment for connections to establish + tokio::time::sleep(Duration::from_secs(2)).await; + + // Show immediate hashrate distribution with network verification + tokio::time::sleep(Duration::from_secs(2)).await; + + info!("🔍 VERIFYING NETWORK CONNECTIONS:"); + + // Check network connections + let output = std::process::Command::new("ss") + .args(&["-tn"]) + .output(); + + if let Ok(output) = output { + let connections = String::from_utf8_lossy(&output.stdout); + let pool_connections: Vec<&str> = connections + .lines() + .filter(|line| line.contains("18.193.252.132:2000") || line.contains("3.74.36.119:2000")) + .collect(); + + info!("📡 Active pool connections: {}", pool_connections.len()); + for conn in pool_connections { + info!(" 🔗 {}", conn.trim()); + } + } + + // Show hashrate distribution + let total_hashrate = ProxyState::get_total_hashrate(); + let detailed_stats = router.get_detailed_connection_stats().await; + + info!("🚀 === INITIAL HASHRATE DISTRIBUTION ==="); + info!("🔋 Total configured hashrate: {}", HashUnit::format_value(total_hashrate)); + info!("🌐 Mode: Parallel (each upstream gets full hashrate)"); + + for (upstream_id, is_active, hashrate) in detailed_stats { + if is_active { + info!(" ✅ {}: receiving {} (100% of total)", upstream_id, HashUnit::format_value(hashrate)); + } else { + info!(" ❌ {}: {} (INACTIVE)", upstream_id, HashUnit::format_value(hashrate)); + } + } + info!("========================================"); + // Start monitoring task for multi-upstream let router_clone = router.clone(); tokio::spawn(async move { @@ -419,38 +462,44 @@ async fn initialize_proxy( std::process::exit(1); } -async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { // Remove mut +async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { let mut distribution_check_counter = 0; loop { tokio::time::sleep(Duration::from_secs(10)).await; distribution_check_counter += 1; - if distribution_check_counter >= 10 { // Every 100 seconds + // Show report every 30 seconds + if distribution_check_counter >= 3 { distribution_check_counter = 0; let total_hashrate = ProxyState::get_total_hashrate(); + let detailed_stats = router.get_detailed_connection_stats().await; - info!("=== Hashrate Distribution Report ==="); - info!("Total hashrate: {}", HashUnit::format_value(total_hashrate)); - - let (total_connections, active_connections) = router.get_connection_stats().await; - info!("Multi-upstream mode: {} total, {} active connections", total_connections, active_connections); + info!("📊 === HASHRATE DISTRIBUTION REPORT ==="); + info!("🔋 Total configured hashrate: {}", HashUnit::format_value(total_hashrate)); + info!("🌐 Parallel mode: Each upstream receives FULL hashrate"); + info!("📡 Upstream details:"); - if active_connections > 0 { - info!("Parallel mode: Using ALL upstreams simultaneously"); - info!("Total hashrate distributed across {} upstreams: {}", - active_connections, HashUnit::format_value(total_hashrate)); - - let active_upstreams = router.get_active_upstreams().await; - for upstream_id in active_upstreams { - info!("Upstream {}: receiving full hashrate {}", - upstream_id, HashUnit::format_value(total_hashrate)); + let mut active_count = 0; + for (upstream_id, is_active, hashrate) in detailed_stats { + if is_active { + active_count += 1; + info!(" ✅ {}: {} (ACTIVE)", upstream_id, HashUnit::format_value(hashrate)); + } else { + info!(" ❌ {}: {} (INACTIVE)", upstream_id, HashUnit::format_value(hashrate)); } + } + + if active_count > 0 { + let total_distributed = total_hashrate * active_count as f32; + info!("🚀 Total hashrate being sent: {} across {} upstreams", + HashUnit::format_value(total_distributed), active_count); + info!("💡 Note: In parallel mode, each upstream receives the full hashrate simultaneously"); } else { - warn!("No active upstream connections!"); + warn!("⚠️ WARNING: No active upstream connections!"); } - info!("====================================="); + info!("=========================================="); } } } diff --git a/src/router/mod.rs b/src/router/mod.rs index 7a985210..38ba61ca 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -461,6 +461,30 @@ impl Router { best_pool.map(|pool| (pool, best_latency)) } + + /// Get detailed connection statistics with individual upstream info + pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { + let mut results = Vec::new(); + + if let Some(ref manager) = self.upstream_manager { + let upstreams = manager.get_upstreams().await; + let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); + + for (id, upstream) in upstreams { + // In parallel mode, each upstream gets the full hashrate + let hashrate = if upstream.is_active { total_hashrate } else { 0.0 }; + results.push((id, upstream.is_active, hashrate)); + } + } else { + // Single upstream mode + if let Some(current) = self.current_pool { + let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); + results.push(("upstream-0".to_string(), true, total_hashrate)); + } + } + + results + } } /// Track latencies for various stages of pool connection setup. From c2e6088b4561f83aad23df6b5b4c4c583d483e9d Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Wed, 4 Jun 2025 01:55:43 +0530 Subject: [PATCH 06/21] Update pool addresses in config and refactor multi-upstream handling for equal distribution --- config.toml | 2 +- src/main.rs | 179 ++++++++++++++++----------- src/proxy_state.rs | 34 +++-- src/router/mod.rs | 148 ++++++++++++++-------- src/router/multi_upstream_manager.rs | 153 ++++++++++++++--------- src/translator/mod.rs | 2 +- 6 files changed, 320 insertions(+), 198 deletions(-) diff --git a/config.toml b/config.toml index 7da55468..01320f14 100644 --- a/config.toml +++ b/config.toml @@ -1,5 +1,5 @@ token="UR01TMkZv6Vs5zbwr0t6" -pool_addresses=["3.74.36.119:2000"] +pool_addresses=["18.193.252.132:2000", "3.74.36.119:2000"] tp_address = "127.0.0.1:8442" interval = 120_000 delay = 540 diff --git a/src/main.rs b/src/main.rs index b7b8e945..45991cb2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,7 +15,6 @@ use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; use proxy_state::{PoolState, ProxyState}; use std::{net::ToSocketAddrs, time::Duration}; -use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; mod api; mod ingress; @@ -94,13 +93,11 @@ pub struct Args { adjustment_interval: u64, // New argument for multiple upstream servers #[clap(long = "upstream", short = 'u', action = ArgAction::Append)] - // In Args struct // #[clap(long = "monitor-hashrate", short = 'm')] // monitor_hashrate: bool, #[clap(long = "config", short = 'c')] pub config_file: Option, - /// Enable parallel upstream usage (sends to all upstreams simultaneously) #[clap(long = "parallel", short = 'p')] @@ -248,8 +245,7 @@ async fn main() { crate::proxy_state::UpstreamType::JDCMiningUpstream, ); } - } - else { + } else { // Fallback to hard-coded address info!("Using hard-coded fallback pool address"); @@ -276,11 +272,10 @@ async fn main() { "fallback-pool-1".to_string(), test_addr, test_pubkey, - crate::proxy_state::UpstreamType::JDCMiningUpstream + crate::proxy_state::UpstreamType::JDCMiningUpstream, ); } - // Direct verification info!("DIRECT VERIFICATION: Checking hashrate distribution"); let upstream_connections = ProxyState::get_upstream_connections(); @@ -311,23 +306,25 @@ async fn main() { // ); // } // } - - // Create the router - always use multi upstream with parallel mode + + // Create the router - always use multi upstream with equal distribution mode let mut router = if !pool_addresses.is_empty() { // Always use multi-upstream mode (even for single upstream) match router::Router::new_multi( pool_addresses.clone(), None, // setup_connection_msg None, // timer - true, // Always enable parallel mode - ).await { + false, // Changed from true to false - disable parallel mode for equal distribution + ) + .await + { Ok(mut router) => { if let Err(e) = router.initialize_upstream_connections().await { error!("Failed to initialize upstream connections: {}", e); std::process::exit(1); } router - }, + } Err(e) => { error!("Failed to create multi-upstream router: {}", e); std::process::exit(1); @@ -342,56 +339,74 @@ async fn main() { // Always use multi-upstream mode initialize_proxy(&mut router, None, epsilon).await; - + // Wait a moment for connections to establish tokio::time::sleep(Duration::from_secs(2)).await; - + // Show immediate hashrate distribution with network verification tokio::time::sleep(Duration::from_secs(2)).await; - + info!("🔍 VERIFYING NETWORK CONNECTIONS:"); - + // Check network connections - let output = std::process::Command::new("ss") - .args(&["-tn"]) - .output(); - + let output = std::process::Command::new("ss").args(&["-tn"]).output(); + if let Ok(output) = output { let connections = String::from_utf8_lossy(&output.stdout); let pool_connections: Vec<&str> = connections .lines() - .filter(|line| line.contains("18.193.252.132:2000") || line.contains("3.74.36.119:2000")) + .filter(|line| { + line.contains("18.193.252.132:2000") || line.contains("3.74.36.119:2000") + }) .collect(); - + info!("📡 Active pool connections: {}", pool_connections.len()); for conn in pool_connections { info!(" 🔗 {}", conn.trim()); } } - - // Show hashrate distribution + // Show hashrate distribution let total_hashrate = ProxyState::get_total_hashrate(); let detailed_stats = router.get_detailed_connection_stats().await; - + info!("🚀 === INITIAL HASHRATE DISTRIBUTION ==="); - info!("🔋 Total configured hashrate: {}", HashUnit::format_value(total_hashrate)); - info!("🌐 Mode: Parallel (each upstream gets full hashrate)"); + info!( + "🔋 Total configured hashrate: {}", + HashUnit::format_value(total_hashrate) + ); - for (upstream_id, is_active, hashrate) in detailed_stats { - if is_active { - info!(" ✅ {}: receiving {} (100% of total)", upstream_id, HashUnit::format_value(hashrate)); - } else { - info!(" ❌ {}: {} (INACTIVE)", upstream_id, HashUnit::format_value(hashrate)); + let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); + + if active_count > 0 { + let percentage_per_upstream = 100.0 / active_count as f32; + info!("🌐 Mode: Equal Distribution ({:.1}% per upstream)", percentage_per_upstream); + + for (upstream_id, is_active, hashrate) in detailed_stats { + if is_active { + info!( + " ✅ {}: receiving {} ({:.1}% of total)", + upstream_id, + HashUnit::format_value(hashrate), + percentage_per_upstream + ); + } else { + info!( + " ❌ {}: {} (INACTIVE)", + upstream_id, + HashUnit::format_value(hashrate) + ); + } } + } else { + warn!("❌ No active upstream connections!"); } info!("========================================"); - // Start monitoring task for multi-upstream let router_clone = router.clone(); tokio::spawn(async move { monitor_multi_upstream(router_clone, epsilon).await; }); - + // Keep main thread alive loop { tokio::time::sleep(Duration::from_secs(60)).await; @@ -402,30 +417,33 @@ async fn main() { async fn initialize_proxy( router: &mut Router, mut _pool_addr: Option, // Add underscore to indicate unused - _epsilon: Duration, // Add underscore to indicate unused + _epsilon: Duration, // Add underscore to indicate unused ) { // Add a static flag to prevent multiple downstream listeners - static DOWNSTREAM_STARTED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); - + static DOWNSTREAM_STARTED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + // For multi-upstream mode, don't loop - just set up and wait if router.is_multi_upstream_enabled() { info!("Multi-upstream mode: using aggregated message handling"); - + // Only start downstream listener once if !DOWNSTREAM_STARTED.swap(true, std::sync::atomic::Ordering::SeqCst) { info!("Starting downstream listener (multi-upstream mode)"); - + // Create channel for downstream connections let (downstreams_tx, mut downstreams_rx) = tokio::sync::mpsc::channel(10); - + // Start the downstream listener let _abort_handle = ingress::sv1_ingress::start_listen_for_downstream(downstreams_tx); - + // Handle incoming downstream connections tokio::spawn(async move { - while let Some((send_to_downstream, recv_from_downstream, client_addr)) = downstreams_rx.recv().await { + while let Some((send_to_downstream, recv_from_downstream, client_addr)) = + downstreams_rx.recv().await + { info!("New downstream client connected: {}", client_addr); - + // Here you would typically start the translator for this downstream connection // For now, we'll just log the connection tokio::spawn(async move { @@ -441,7 +459,7 @@ async fn initialize_proxy( } }); } - + // Get aggregated receiver if available if let Some(aggregated_receiver) = router.get_aggregated_receiver() { tokio::spawn(async move { @@ -451,12 +469,12 @@ async fn initialize_proxy( } }); } - + // Instead of looping, just wait indefinitely info!("Multi-upstream proxy initialized successfully"); return; // Exit the function, don't loop } - + // For single upstream mode - just log and exit since we're not supporting it error!("Single upstream mode is not supported in this version. Please use multi-upstream mode with --config option."); std::process::exit(1); @@ -464,38 +482,54 @@ async fn initialize_proxy( async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { let mut distribution_check_counter = 0; - + loop { tokio::time::sleep(Duration::from_secs(10)).await; distribution_check_counter += 1; - + // Show report every 30 seconds if distribution_check_counter >= 3 { distribution_check_counter = 0; - + let total_hashrate = ProxyState::get_total_hashrate(); let detailed_stats = router.get_detailed_connection_stats().await; - + info!("📊 === HASHRATE DISTRIBUTION REPORT ==="); - info!("🔋 Total configured hashrate: {}", HashUnit::format_value(total_hashrate)); - info!("🌐 Parallel mode: Each upstream receives FULL hashrate"); - info!("📡 Upstream details:"); + info!( + "🔋 Total configured hashrate: {}", + HashUnit::format_value(total_hashrate) + ); - let mut active_count = 0; - for (upstream_id, is_active, hashrate) in detailed_stats { - if is_active { - active_count += 1; - info!(" ✅ {}: {} (ACTIVE)", upstream_id, HashUnit::format_value(hashrate)); - } else { - info!(" ❌ {}: {} (INACTIVE)", upstream_id, HashUnit::format_value(hashrate)); - } - } + let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); if active_count > 0 { - let total_distributed = total_hashrate * active_count as f32; - info!("🚀 Total hashrate being sent: {} across {} upstreams", - HashUnit::format_value(total_distributed), active_count); - info!("💡 Note: In parallel mode, each upstream receives the full hashrate simultaneously"); + let percentage_per_upstream = 100.0 / active_count as f32; + info!("🌐 Equal Distribution mode: Each upstream gets {:.1}% of total hashrate", percentage_per_upstream); + info!("📡 Upstream details:"); + + for (upstream_id, is_active, hashrate) in detailed_stats { + if is_active { + info!( + " ✅ {}: {} ({:.1}% of total)", + upstream_id, + HashUnit::format_value(hashrate), + percentage_per_upstream + ); + } else { + info!( + " ❌ {}: {} (INACTIVE)", + upstream_id, + HashUnit::format_value(hashrate) + ); + } + } + + info!( + "🚀 Total hashrate distributed: {} across {} upstreams", + HashUnit::format_value(total_hashrate), + active_count + ); + info!("💡 Note: Hashrate is split equally among all active upstreams"); } else { warn!("⚠️ WARNING: No active upstream connections!"); } @@ -517,9 +551,12 @@ async fn monitor( // Generate periodic reports for multi-upstream let (total_connections, active_connections) = router.get_connection_stats().await; let total_hashrate = ProxyState::get_total_hashrate(); - - info!("Multi-upstream mode: {} total, {} active connections", total_connections, active_connections); - + + info!( + "Multi-upstream mode: {} total, {} active connections", + total_connections, active_connections + ); + if active_connections > 0 { info!("Parallel mode: Using ALL upstreams simultaneously"); info!( @@ -527,7 +564,7 @@ async fn monitor( active_connections, HashUnit::format_value(total_hashrate) ); - + let active_upstreams = router.get_active_upstreams().await; for upstream_id in active_upstreams { info!( @@ -619,5 +656,3 @@ impl HashUnit { } } } - - diff --git a/src/proxy_state.rs b/src/proxy_state.rs index a469ecbc..d9bcca46 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -297,8 +297,6 @@ impl ProxyState { } } - - /// Set the total hashrate to be distributed among upstreams pub fn set_total_hashrate(hashrate: f32) { info!("Setting total hashrate to: {} h/s", hashrate); @@ -444,7 +442,6 @@ impl ProxyState { } } - /// Update connection status for an upstream with timestamp pub fn set_upstream_connection_status(id: &str, connected: bool) { if PROXY_STATE @@ -489,7 +486,8 @@ impl ProxyState { if PROXY_STATE .safe_lock(|state| { - count = state.upstream_connections + count = state + .upstream_connections .values() .filter(|conn| conn.is_connected) .count(); @@ -535,7 +533,7 @@ impl ProxyState { } } - /// Check if proxy is down + /// Check if proxy is down pub fn is_proxy_down() -> (bool, Option) { let errors = Self::get_errors(); if errors.is_ok() && errors.as_ref().unwrap().is_empty() { @@ -547,7 +545,7 @@ impl ProxyState { } } - pub fn get_errors() -> Result, ()> { + pub fn get_errors() -> Result, ()> { let mut errors = Vec::new(); if PROXY_STATE .safe_lock(|state| { @@ -585,17 +583,22 @@ impl ProxyState { } } - /// Get upstream statistics pub fn get_upstream_stats() -> Vec<(String, bool, u64, u64)> { let mut stats = Vec::new(); if PROXY_STATE .safe_lock(|state| { - stats = state.upstream_connections + stats = state + .upstream_connections .iter() .map(|(id, conn)| { - (id.clone(), conn.is_connected, conn.shares_submitted, conn.shares_accepted) + ( + id.clone(), + conn.is_connected, + conn.shares_submitted, + conn.shares_accepted, + ) }) .collect(); }) @@ -609,7 +612,12 @@ impl ProxyState { } /// Get all upstream connections (including inactive ones) - pub fn get_all_upstream_connections() -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey, bool)> { + pub fn get_all_upstream_connections() -> Vec<( + String, + std::net::SocketAddr, + key_utils::Secp256k1PublicKey, + bool, + )> { let mut connections = Vec::new(); if PROXY_STATE @@ -630,7 +638,8 @@ impl ProxyState { } /// Get only active upstream connections (existing method is fine) - pub fn get_upstream_connections() -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + pub fn get_upstream_connections( + ) -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { let mut connections = Vec::new(); if PROXY_STATE @@ -658,7 +667,8 @@ impl ProxyState { /// Get best upstream based on latency or other criteria /// For now, just returns the first active upstream - pub fn get_best_upstream() -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { + pub fn get_best_upstream( + ) -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { let mut result = None; if PROXY_STATE diff --git a/src/router/mod.rs b/src/router/mod.rs index 38ba61ca..2fa33329 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -36,14 +36,13 @@ pub struct Router { pub current_pool: Option, pub upstream_manager: Option, pub aggregated_receiver: Option>>, - + // Keep these fields for backward compatibility with single upstream mode pub auth_pub_k: Secp256k1PublicKey, pub setup_connection_msg: Option>, pub timer: Option, pub latency_tx: tokio::sync::watch::Sender>, pub latency_rx: tokio::sync::watch::Receiver>, - // Remove round-robin fields entirely // use_round_robin: bool, // use_parallel: bool, // This will be determined by presence of upstream_manager @@ -73,10 +72,10 @@ impl Router { auth_pub_k: Secp256k1PublicKey, setup_connection_msg: Option>, timer: Option, - ) -> Self{ + ) -> Self { let (latency_tx, latency_rx) = watch::channel(None); let auth_keys = vec![auth_pub_k; pool_addresses.len()]; - + Self { pool_socket_addresses: pool_addresses, keys: auth_keys, @@ -90,23 +89,26 @@ impl Router { latency_rx, } } - /// Creates a new Router with multiple upstream addresses and auth keys - pub async fn new_multi( + pub async fn new_multi( pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, setup_connection_msg: Option>, timer: Option, _use_parallel: bool, // Parameter kept for compatibility but always use parallel ) -> Result { - let pool_socket_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); + let pool_socket_addresses: Vec = + pool_address_keys.iter().map(|(addr, _)| *addr).collect(); let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); // Create with parallel mode enabled let (upstream_manager, aggregated_receiver) = MultiUpstreamManager::new(true); - + // Use first key as default auth key for backward compatibility - let auth_pub_k = keys.first().copied().ok_or("No authentication keys provided")?; + let auth_pub_k = keys + .first() + .copied() + .ok_or("No authentication keys provided")?; let (latency_tx, latency_rx) = watch::channel(None); @@ -123,7 +125,7 @@ impl Router { latency_rx, }) } - + // Remove round-robin methods // pub fn set_round_robin(&mut self, enabled: bool) { ... } @@ -131,7 +133,7 @@ impl Router { pub fn is_multi_upstream_enabled(&self) -> bool { self.upstream_manager.is_some() } - + /// Check if the router is using parallel mode - true when MultiUpstreamManager is present pub fn is_parallel_mode(&self) -> bool { self.upstream_manager.is_some() @@ -145,7 +147,7 @@ impl Router { // All upstreams are used simultaneously return None; } - + // For single upstream mode, check for better latency if let Some(best_pool) = self.select_pool_monitor(epsilon).await { if Some(best_pool) != self.current_pool { @@ -223,11 +225,12 @@ impl Router { info!("Upstream {:?} selected", pool); // Find the matching auth key for this address - fix field name - let auth_key = if let Some(index) = self.pool_socket_addresses.iter().position(|&a| a == pool) { - self.keys[index] - } else { - self.auth_pub_k - }; + let auth_key = + if let Some(index) = self.pool_socket_addresses.iter().position(|&a| a == pool) { + self.keys[index] + } else { + self.auth_pub_k + }; match minin_pool_connection::connect_pool( pool, @@ -270,12 +273,15 @@ impl Router { /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { // Find the auth key for this address - fix field names - let auth_pub_key = - if let Some(index) = self.pool_socket_addresses.iter().position(|&a| a == pool_address) { - self.keys[index] - } else { - self.auth_pub_k - }; + let auth_pub_key = if let Some(index) = self + .pool_socket_addresses + .iter() + .position(|&a| a == pool_address) + { + self.keys[index] + } else { + self.auth_pub_k + }; let mut pool = PoolLatency::new(pool_address); let setup_connection_msg = self.setup_connection_msg.as_ref(); @@ -337,42 +343,60 @@ impl Router { self.current_pool } - /// Initialize upstream connections for the manager + /// Initialize upstream connections for the manager pub async fn initialize_upstream_connections(&mut self) -> Result<(), String> { if let Some(ref manager) = self.upstream_manager { - info!("Initializing {} upstream connections", self.pool_socket_addresses.len()); - + info!( + "Initializing {} upstream connections", + self.pool_socket_addresses.len() + ); + // Add each unique upstream only once - fix field names - for (idx, (addr, key)) in self.pool_socket_addresses.iter().zip(self.keys.iter()).enumerate() { + for (idx, (addr, key)) in self + .pool_socket_addresses + .iter() + .zip(self.keys.iter()) + .enumerate() + { let id = format!("upstream-{}", idx); - + info!("Adding upstream {}: {} ({})", id, addr, key); - - if let Err(e) = manager.add_upstream( - id.clone(), - *addr, - *key, - self.setup_connection_msg.clone(), - self.timer, - ).await { + + if let Err(e) = manager + .add_upstream( + id.clone(), + *addr, + *key, + self.setup_connection_msg.clone(), + self.timer, + ) + .await + { error!("Failed to initialize upstream {}: {}", addr, e); // Continue with other upstreams even if one fails } else { info!("Successfully initialized upstream {}: {}", id, addr); } } - + // Wait a moment for connections to stabilize tokio::time::sleep(Duration::from_millis(100)).await; - + let (total, active) = self.get_connection_stats().await; - info!("Initialized {} upstream connections ({} active)", total, active); + info!( + "Initialized {} upstream connections ({} active)", + total, active + ); } Ok(()) } /// Send a message to a specific upstream using the manager - pub async fn send_to_upstream(&self, upstream_id: &str, message: PoolExtMessages<'static>) -> Result<(), String> { + pub async fn send_to_upstream( + &self, + upstream_id: &str, + message: PoolExtMessages<'static>, + ) -> Result<(), String> { if let Some(ref manager) = self.upstream_manager { manager.send_to_upstream(upstream_id, message).await } else { @@ -381,7 +405,10 @@ impl Router { } /// Send a message to the next upstream (round-robin or best available) - pub async fn send_to_next_upstream(&self, message: PoolExtMessages<'static>) -> Result<(), String> { + pub async fn send_to_next_upstream( + &self, + message: PoolExtMessages<'static>, + ) -> Result<(), String> { if let Some(ref manager) = self.upstream_manager { manager.send_to_next_upstream(message).await } else { @@ -390,7 +417,10 @@ impl Router { } /// Broadcast a message to all active upstreams - pub async fn broadcast_to_all_upstreams(&self, message: PoolExtMessages<'static>) -> Vec { + pub async fn broadcast_to_all_upstreams( + &self, + message: PoolExtMessages<'static>, + ) -> Vec { if let Some(ref manager) = self.upstream_manager { manager.broadcast(message).await } else { @@ -399,10 +429,12 @@ impl Router { } /// Get aggregated receiver for multi-upstream mode - pub fn get_aggregated_receiver(&mut self) -> Option>> { + pub fn get_aggregated_receiver( + &mut self, + ) -> Option>> { self.aggregated_receiver.take() } - + /// Get connection statistics for multi-upstream mode pub async fn get_connection_stats(&self) -> (usize, usize) { if let Some(ref manager) = self.upstream_manager { @@ -415,7 +447,7 @@ impl Router { (self.pool_socket_addresses.len(), 1) // Single upstream mode } } - + /// Get list of active upstream IDs pub async fn get_active_upstreams(&self) -> Vec { if let Some(ref manager) = self.upstream_manager { @@ -428,11 +460,13 @@ impl Router { /// Switch to a specific upstream by ID pub async fn switch_to_upstream(&mut self, upstream_id: &str) -> Result<(), String> { // Find the upstream address by ID - if let Ok(idx) = upstream_id.strip_prefix("upstream-") + if let Ok(idx) = upstream_id + .strip_prefix("upstream-") .and_then(|s| s.parse::().ok()) - .ok_or("Invalid upstream ID format".to_string()) + .ok_or("Invalid upstream ID format".to_string()) { - if idx < self.pool_socket_addresses.len() { // Fix field name + if idx < self.pool_socket_addresses.len() { + // Fix field name let addr = self.pool_socket_addresses[idx]; // Fix field name self.current_pool = Some(addr); info!("Switched to upstream {}: {:?}", upstream_id, addr); @@ -462,7 +496,7 @@ impl Router { best_pool.map(|pool| (pool, best_latency)) } - /// Get detailed connection statistics with individual upstream info + /// Get detailed connection statistics with equal distribution pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { let mut results = Vec::new(); @@ -470,13 +504,23 @@ impl Router { let upstreams = manager.get_upstreams().await; let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); + // Count active upstreams first + let active_count = upstreams.values().filter(|upstream| upstream.is_active).count(); + + // Calculate hashrate per upstream (equal distribution) + let hashrate_per_upstream = if active_count > 0 { + total_hashrate / active_count as f32 + } else { + 0.0 + }; + for (id, upstream) in upstreams { - // In parallel mode, each upstream gets the full hashrate - let hashrate = if upstream.is_active { total_hashrate } else { 0.0 }; + // Each active upstream gets an equal share + let hashrate = if upstream.is_active { hashrate_per_upstream } else { 0.0 }; results.push((id, upstream.is_active, hashrate)); } } else { - // Single upstream mode + // Single upstream mode (fallback) if let Some(current) = self.current_pool { let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); results.push(("upstream-0".to_string(), true, total_hashrate)); diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index d193e450..77390788 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -1,15 +1,10 @@ -use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; -use tokio::sync::{mpsc, RwLock, Mutex}; -use tracing::{error, info, warn}; -use key_utils::Secp256k1PublicKey; +use crate::{proxy_state::ProxyState, shared::utils::AbortOnDrop}; use demand_share_accounting_ext::parser::PoolExtMessages; +use key_utils::Secp256k1PublicKey; use roles_logic_sv2::common_messages_sv2::SetupConnection; -use crate::{ - minin_pool_connection::{connect_pool}, - proxy_state::ProxyState, - shared::utils::AbortOnDrop, -}; -use futures::future::join_all; +use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; +use tokio::sync::{mpsc, Mutex}; +use tracing::{error, info, warn}; #[derive(Clone)] pub struct UpstreamConnection { @@ -35,7 +30,7 @@ impl MultiUpstreamManager { _use_parallel: bool, // Parameter kept for compatibility but always use parallel ) -> (Self, tokio::sync::mpsc::Receiver>) { let (sender, receiver) = tokio::sync::mpsc::channel(1000); - + ( Self { upstreams: Arc::new(Mutex::new(HashMap::new())), @@ -58,15 +53,15 @@ impl MultiUpstreamManager { timer: Option, ) -> Result<(), String> { info!("🔍 ADD_UPSTREAM CALLED: {} -> {}", id, address); - + let mut upstreams = self.upstreams.lock().await; - + // Check if upstream already exists if upstreams.contains_key(&id) { warn!("Upstream {} already exists, skipping", id); return Ok(()); } - + // Create upstream connection let upstream_connection = UpstreamConnection { id: id.clone(), @@ -76,14 +71,14 @@ impl MultiUpstreamManager { sender: self.aggregated_sender.clone(), // Use a proper sender here connection_handle: None, // Set to None initially }; - + upstreams.insert(id.clone(), upstream_connection); info!("✅ Added upstream {} to manager", id); - + // Start connection task let upstreams_clone = self.upstreams.clone(); let sender_clone = self.aggregated_sender.clone(); - + tokio::spawn(async move { Self::connect_upstream( upstreams_clone, @@ -93,12 +88,13 @@ impl MultiUpstreamManager { setup_connection_msg, timer, sender_clone, - ).await + ) + .await }); - + Ok(()) } - + /// Connect to a specific upstream async fn connect_upstream( upstreams: Arc>>, @@ -110,16 +106,18 @@ impl MultiUpstreamManager { sender: tokio::sync::mpsc::Sender>, ) { info!("Connecting to upstream {}: {}", id, address); - + match crate::minin_pool_connection::connect_pool( address, auth_key, setup_connection_msg, timer, - ).await { + ) + .await + { Ok((send_to_pool, recv_from_pool, _abort_handle)) => { info!("Successfully connected to upstream {}: {}", id, address); - + // Update upstream status in both MultiUpstreamManager and ProxyState { let mut upstreams_lock = upstreams.lock().await; @@ -127,16 +125,16 @@ impl MultiUpstreamManager { upstream.is_active = true; } } - + // Also update ProxyState crate::proxy_state::ProxyState::set_upstream_connection_status(&id, true); - + // Handle messages from this upstream Self::handle_upstream_messages(id.clone(), recv_from_pool, sender).await; - + // If we reach here, connection was lost warn!("Connection to upstream {} lost", id); - + // Mark as inactive { let mut upstreams_lock = upstreams.lock().await; @@ -148,7 +146,7 @@ impl MultiUpstreamManager { } Err(e) => { error!("Failed to connect to upstream {}: {}", id, e); - + // Mark as inactive let mut upstreams_lock = upstreams.lock().await; if let Some(upstream) = upstreams_lock.get_mut(&id) { @@ -158,7 +156,7 @@ impl MultiUpstreamManager { } } } - + /// Handle messages from an upstream async fn handle_upstream_messages( upstream_id: String, @@ -166,23 +164,33 @@ impl MultiUpstreamManager { sender: tokio::sync::mpsc::Sender>, ) { info!("Starting message handler for upstream {}", upstream_id); - + while let Some(message) = receiver.recv().await { if let Err(e) = sender.send(message).await { - error!("Failed to forward message from upstream {}: {}", upstream_id, e); + error!( + "Failed to forward message from upstream {}: {}", + upstream_id, e + ); break; } } - + warn!("Message handler for upstream {} stopped", upstream_id); } /// Send a message to a specific upstream - pub async fn send_to_upstream(&self, upstream_id: &str, message: PoolExtMessages<'static>) -> Result<(), String> { + pub async fn send_to_upstream( + &self, + upstream_id: &str, + message: PoolExtMessages<'static>, + ) -> Result<(), String> { let connections = self.upstreams.lock().await; if let Some(connection) = connections.get(upstream_id) { if connection.is_active { - connection.sender.send(message).await + connection + .sender + .send(message) + .await .map_err(|e| format!("Failed to send to upstream {}: {}", upstream_id, e))?; Ok(()) } else { @@ -194,15 +202,19 @@ impl MultiUpstreamManager { } /// Send a message to the next upstream using the configured strategy - pub async fn send_to_next_upstream(&self, message: PoolExtMessages<'static>) -> Result<(), String> { + pub async fn send_to_next_upstream( + &self, + message: PoolExtMessages<'static>, + ) -> Result<(), String> { info!("Broadcasting to all upstreams (parallel mode)"); let results = self.broadcast(message).await; - + // Check if any sends were successful - let success_count = results.iter() + let success_count = results + .iter() .filter(|result| result.contains("Success")) .count(); - + if success_count > 0 { info!("Successfully sent to {} upstreams", success_count); Ok(()) @@ -238,19 +250,27 @@ impl MultiUpstreamManager { let mut retry_count = 0; const MAX_RETRIES: u32 = 3; const RETRY_DELAY: Duration = Duration::from_secs(5); - + while retry_count < MAX_RETRIES { - info!("Connecting to upstream {} (attempt {}/{}): {}", id, retry_count + 1, MAX_RETRIES, address); - + info!( + "Connecting to upstream {} (attempt {}/{}): {}", + id, + retry_count + 1, + MAX_RETRIES, + address + ); + match crate::minin_pool_connection::connect_pool( address, auth_key, setup_connection_msg.clone(), timer, - ).await { + ) + .await + { Ok((send_to_pool, recv_from_pool, _abort_handle)) => { info!("Successfully connected to upstream {}: {}", id, address); - + // Update upstream status { let mut upstreams_lock = upstreams.lock().await; @@ -258,15 +278,20 @@ impl MultiUpstreamManager { upstream.is_active = true; } } - + // Handle messages from this upstream Self::handle_upstream_messages(id.clone(), recv_from_pool, sender).await; return; // Success, exit retry loop } Err(e) => { - error!("Failed to connect to upstream {} (attempt {}): {}", id, retry_count + 1, e); + error!( + "Failed to connect to upstream {} (attempt {}): {}", + id, + retry_count + 1, + e + ); retry_count += 1; - + if retry_count < MAX_RETRIES { info!("Retrying connection to {} in {:?}", id, RETRY_DELAY); tokio::time::sleep(RETRY_DELAY).await; @@ -274,9 +299,12 @@ impl MultiUpstreamManager { } } } - - error!("Failed to connect to upstream {} after {} attempts", id, MAX_RETRIES); - + + error!( + "Failed to connect to upstream {} after {} attempts", + id, MAX_RETRIES + ); + // Mark as inactive after all retries failed let mut upstreams_lock = upstreams.lock().await; if let Some(upstream) = upstreams_lock.get_mut(&id) { @@ -284,24 +312,25 @@ impl MultiUpstreamManager { } } /// Broadcast to all active upstreams (parallel execution) - pub async fn broadcast(&self, message: PoolExtMessages<'static>) -> Vec { + pub async fn broadcast(&self, message: PoolExtMessages<'static>) -> Vec { let upstreams = self.upstreams.lock().await; let mut results = Vec::new(); - + if upstreams.is_empty() { results.push("No upstreams configured".to_string()); return results; } - - let active_upstreams: Vec<_> = upstreams.values() + + let active_upstreams: Vec<_> = upstreams + .values() .filter(|upstream| upstream.is_active) .collect(); - + if active_upstreams.is_empty() { results.push("No active upstreams available".to_string()); return results; } - + for upstream in active_upstreams { match upstream.sender.send(message.clone()).await { Ok(_) => { @@ -310,11 +339,14 @@ impl MultiUpstreamManager { Err(e) => { results.push(format!("Sent to {}: Failed - {}", upstream.id, e)); // Mark upstream as inactive if send fails - warn!("Upstream {} send failed, may need reconnection: {}", upstream.id, e); + warn!( + "Upstream {} send failed, may need reconnection: {}", + upstream.id, e + ); } } } - + results } @@ -323,11 +355,12 @@ impl MultiUpstreamManager { let upstreams = self.upstreams.lock().await; upstreams.clone() } - + /// Get list of active upstream IDs pub async fn get_active_upstream_ids(&self) -> Vec { let upstreams = self.upstreams.lock().await; - upstreams.values() + upstreams + .values() .filter(|upstream| upstream.is_active) .map(|upstream| upstream.id.clone()) .collect() @@ -366,4 +399,4 @@ impl MultiUpstreamManager { let connections = self.upstreams.lock().await; connections.values().filter(|conn| conn.is_active).count() } -} \ No newline at end of file +} diff --git a/src/translator/mod.rs b/src/translator/mod.rs index a5c6ed35..a6090d9a 100644 --- a/src/translator/mod.rs +++ b/src/translator/mod.rs @@ -1,7 +1,7 @@ mod downstream; mod error; -pub mod proxy; // Make public +pub mod proxy; // Make public pub mod upstream; // Make public mod utils; From 534e4353cea425d767d2f3cec395e307d82e0396 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Wed, 4 Jun 2025 15:13:55 +0530 Subject: [PATCH 07/21] Implement custom hashrate distribution for multi-upstream management; update configuration handling and validation --- config.toml | 10 +- src/main.rs | 111 ++++++++++++++--- src/router/mod.rs | 180 ++++++++++++--------------- src/router/multi_upstream_manager.rs | 91 +++++++++----- 4 files changed, 246 insertions(+), 146 deletions(-) diff --git a/config.toml b/config.toml index 01320f14..44cfd441 100644 --- a/config.toml +++ b/config.toml @@ -1,9 +1,13 @@ token="UR01TMkZv6Vs5zbwr0t6" -pool_addresses=["18.193.252.132:2000", "3.74.36.119:2000"] +pool_addresses=["18.193.252.132:2000","3.74.36.119:2000"] tp_address = "127.0.0.1:8442" interval = 120_000 delay = 540 -downstream_hashrate = "200T" +downstream_hashrate = "500T" loglevel = "info" nc_loglevel = "off" -test = true \ No newline at end of file +test = true + +# Custom hashrate distribution (percentages for each upstream) +# This array should match the number of pool_addresses +hashrate_distribution = [30.0, 70.0] \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 45991cb2..3f56e342 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,7 +10,6 @@ use std::path::Path; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; -use crate::shared::utils::AbortOnDrop; use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; use proxy_state::{PoolState, ProxyState}; @@ -115,6 +114,8 @@ pub struct Config { pub loglevel: Option, pub nc_loglevel: Option, pub test: Option, + pub hashrate_distribution: Option>, // Array of percentages for each upstream + } impl Config { @@ -213,10 +214,24 @@ async fn main() { // Handle multiple upstream servers from config or hard-coded fallback let mut pool_addresses = Vec::new(); + let mut custom_distribution: Option> = None; - if let Some(ref config) = config { + if let Some(ref config) = config { // Use configuration file pool addresses info!("Loading pool addresses from configuration file"); + + // Load custom distribution if provided + if let Some(ref distribution) = config.hashrate_distribution { + info!("Found custom hashrate distribution in config: {:?}", distribution); + + // Validate the distribution + if let Err(e) = router::MultiUpstreamManager::validate_distribution(distribution) { + error!("Invalid hashrate distribution in config: {}", e); + std::process::exit(1); + } + + custom_distribution = Some(distribution.clone()); + } for (idx, pool_addr_str) in config.pool_addresses.iter().enumerate() { // Parse address @@ -245,6 +260,19 @@ async fn main() { crate::proxy_state::UpstreamType::JDCMiningUpstream, ); } + + // Validate that distribution count matches pool count + if let Some(ref distribution) = custom_distribution { + if distribution.len() != pool_addresses.len() { + error!( + "Hashrate distribution count ({}) doesn't match pool addresses count ({})", + distribution.len(), + pool_addresses.len() + ); + std::process::exit(1); + } + } + } else { // Fallback to hard-coded address info!("Using hard-coded fallback pool address"); @@ -307,18 +335,27 @@ async fn main() { // } // } - // Create the router - always use multi upstream with equal distribution mode + // Create the router - always use multi upstream with custom distribution mode let mut router = if !pool_addresses.is_empty() { // Always use multi-upstream mode (even for single upstream) match router::Router::new_multi( pool_addresses.clone(), None, // setup_connection_msg None, // timer - false, // Changed from true to false - disable parallel mode for equal distribution + false, // Changed from true to false - disable parallel mode for custom distribution ) .await { Ok(mut router) => { + // Set custom distribution if provided + if let Some(distribution) = custom_distribution { + info!("Setting custom hashrate distribution: {:?}", distribution); + if let Err(e) = router.set_hashrate_distribution(distribution).await { + error!("Failed to set custom hashrate distribution: {}", e); + std::process::exit(1); + } + } + if let Err(e) = router.initialize_upstream_connections().await { error!("Failed to initialize upstream connections: {}", e); std::process::exit(1); @@ -378,16 +415,33 @@ async fn main() { let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); if active_count > 0 { - let percentage_per_upstream = 100.0 / active_count as f32; - info!("🌐 Mode: Equal Distribution ({:.1}% per upstream)", percentage_per_upstream); + // Check if using custom distribution + let mut is_custom_distribution = false; + for (_, is_active, hashrate) in &detailed_stats { + if *is_active { + let percentage = (*hashrate / total_hashrate) * 100.0; + if (percentage - (100.0 / active_count as f32)).abs() > 1.0 { + is_custom_distribution = true; + break; + } + } + } + + if is_custom_distribution { + info!("🎯 Mode: Custom Distribution (from config)"); + } else { + let percentage_per_upstream = 100.0 / active_count as f32; + info!("🌐 Mode: Equal Distribution ({:.1}% per upstream)", percentage_per_upstream); + } for (upstream_id, is_active, hashrate) in detailed_stats { if is_active { + let percentage = (hashrate / total_hashrate) * 100.0; info!( " ✅ {}: receiving {} ({:.1}% of total)", upstream_id, HashUnit::format_value(hashrate), - percentage_per_upstream + percentage ); } else { info!( @@ -461,7 +515,7 @@ async fn initialize_proxy( } // Get aggregated receiver if available - if let Some(aggregated_receiver) = router.get_aggregated_receiver() { + if let Some(aggregated_receiver) = router.get_aggregated_receiver().await { tokio::spawn(async move { let mut receiver = aggregated_receiver; while let Some(message) = receiver.recv().await { @@ -503,17 +557,35 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); if active_count > 0 { - let percentage_per_upstream = 100.0 / active_count as f32; - info!("🌐 Equal Distribution mode: Each upstream gets {:.1}% of total hashrate", percentage_per_upstream); + // Check if using custom distribution + let mut is_custom_distribution = false; + + for (_, is_active, hashrate) in &detailed_stats { + if *is_active { + let percentage = (*hashrate / total_hashrate) * 100.0; + if (percentage - (100.0 / active_count as f32)).abs() > 1.0 { + is_custom_distribution = true; + } + } + } + + if is_custom_distribution { + info!("🎯 Custom Distribution mode: Using configured percentages"); + } else { + let percentage_per_upstream = 100.0 / active_count as f32; + info!("🌐 Equal Distribution mode: Each upstream gets {:.1}% of total hashrate", percentage_per_upstream); + } + info!("📡 Upstream details:"); for (upstream_id, is_active, hashrate) in detailed_stats { if is_active { + let percentage = (hashrate / total_hashrate) * 100.0; info!( " ✅ {}: {} ({:.1}% of total)", upstream_id, HashUnit::format_value(hashrate), - percentage_per_upstream + percentage ); } else { info!( @@ -529,7 +601,12 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { HashUnit::format_value(total_hashrate), active_count ); - info!("💡 Note: Hashrate is split equally among all active upstreams"); + + if is_custom_distribution { + info!("💡 Note: Using custom distribution from config file"); + } else { + info!("💡 Note: Hashrate is split equally among all active upstreams"); + } } else { warn!("⚠️ WARNING: No active upstream connections!"); } @@ -539,7 +616,7 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { } async fn monitor( router: &mut Router, - abort_handles: Vec<(AbortOnDrop, std::string::String)>, + abort_handles: Vec<(std::string::String)>, epsilon: Duration, server_handle: tokio::task::JoinHandle<()>, ) -> Reconnect { @@ -549,7 +626,10 @@ async fn monitor( tokio::time::sleep(Duration::from_secs(30)).await; // Generate periodic reports for multi-upstream - let (total_connections, active_connections) = router.get_connection_stats().await; + // Replace with a method that exists, e.g., get_detailed_connection_stats + let detailed_stats = router.get_detailed_connection_stats().await; + let total_connections = detailed_stats.len(); + let active_connections = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); let total_hashrate = ProxyState::get_total_hashrate(); info!( @@ -615,8 +695,7 @@ fn parse_hashrate(hashrate_str: &str) -> Result { pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream -} - +} enum HashUnit { Tera, Peta, diff --git a/src/router/mod.rs b/src/router/mod.rs index 2fa33329..e2ea8f51 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -27,7 +27,7 @@ use crate::{ // Add the new module pub mod multi_upstream_manager; -use multi_upstream_manager::MultiUpstreamManager; +pub use multi_upstream_manager::MultiUpstreamManager; /// Router handles connection to Multiple upstreams. pub struct Router { @@ -95,14 +95,14 @@ impl Router { pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, setup_connection_msg: Option>, timer: Option, - _use_parallel: bool, // Parameter kept for compatibility but always use parallel + use_parallel: bool, ) -> Result { let pool_socket_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); // Create with parallel mode enabled - let (upstream_manager, aggregated_receiver) = MultiUpstreamManager::new(true); + let (upstream_manager, aggregated_receiver) = MultiUpstreamManager::new(use_parallel); // Use first key as default auth key for backward compatibility let auth_pub_k = keys @@ -126,9 +126,74 @@ impl Router { }) } - // Remove round-robin methods - // pub fn set_round_robin(&mut self, enabled: bool) { ... } + /// Get aggregated receiver for multi-upstream mode + pub async fn get_aggregated_receiver(&self) -> Option>> { + if let Some(ref manager) = self.upstream_manager { + manager.get_aggregated_receiver().await + } else { + None + } + } + + /// Set hashrate distribution for multi-upstream mode + pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result<(), String> { + if let Some(ref manager) = self.upstream_manager { + manager.set_hashrate_distribution(distribution).await; + Ok(()) + } else { + Err("Multi-upstream manager not available".to_string()) + } + } + + /// Get detailed connection statistics with custom distribution + pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { + let mut results = Vec::new(); + + if let Some(ref manager) = self.upstream_manager { + let upstreams = manager.get_upstreams().await; + let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); + let distribution = manager.get_hashrate_distribution().await; + + // Check if we have custom distribution + if distribution.len() == upstreams.len() && !distribution.is_empty() { + // Use custom distribution + let upstream_ids: Vec = upstreams.keys().cloned().collect(); + for (idx, (id, upstream)) in upstreams.iter().enumerate() { + let percentage = distribution.get(idx).copied().unwrap_or(0.0); + let hashrate = if upstream.is_active { + total_hashrate * (percentage / 100.0) + } else { + 0.0 + }; + results.push((id.clone(), upstream.is_active, hashrate)); + } + } else { + // Fall back to equal distribution + let active_count = upstreams.values().filter(|upstream| upstream.is_active).count(); + let hashrate_per_upstream = if active_count > 0 { + total_hashrate / active_count as f32 + } else { + 0.0 + }; + + for (id, upstream) in upstreams { + let hashrate = if upstream.is_active { hashrate_per_upstream } else { 0.0 }; + results.push((id, upstream.is_active, hashrate)); + } + } + } else { + // Single upstream mode (fallback) + if let Some(current) = self.current_pool { + let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); + results.push(("upstream-0".to_string(), true, total_hashrate)); + } + } + + results + } + + /// Check if the router is using the MultiUpstreamManager pub fn is_multi_upstream_enabled(&self) -> bool { self.upstream_manager.is_some() @@ -213,7 +278,7 @@ impl Router { Some(addr) => addr, None => match self.select_pool_connect().await { Some(addr) => addr, - // Called when we initialize the proxy, without a valid pool we can not start mine and we + // Called when we initialize // return Err None => { return Err(minin_pool_connection::errors::Error::Unrecoverable); @@ -372,73 +437,24 @@ impl Router { ) .await { - error!("Failed to initialize upstream {}: {}", addr, e); - // Continue with other upstreams even if one fails - } else { - info!("Successfully initialized upstream {}: {}", id, addr); + error!("Failed to add upstream {}: {:?}", id, e); } } - // Wait a moment for connections to stabilize - tokio::time::sleep(Duration::from_millis(100)).await; - - let (total, active) = self.get_connection_stats().await; - info!( - "Initialized {} upstream connections ({} active)", - total, active - ); - } - Ok(()) - } + // Start the manager after adding all upstreams + // manager.start().await?; + // If MultiUpstreamManager requires initialization, call it here. Otherwise, remove this. - /// Send a message to a specific upstream using the manager - pub async fn send_to_upstream( - &self, - upstream_id: &str, - message: PoolExtMessages<'static>, - ) -> Result<(), String> { - if let Some(ref manager) = self.upstream_manager { - manager.send_to_upstream(upstream_id, message).await - } else { - Err("MultiUpstreamManager not initialized".to_string()) - } - } - - /// Send a message to the next upstream (round-robin or best available) - pub async fn send_to_next_upstream( - &self, - message: PoolExtMessages<'static>, - ) -> Result<(), String> { - if let Some(ref manager) = self.upstream_manager { - manager.send_to_next_upstream(message).await + info!("Upstream connections initialized"); + Ok(()) } else { - Err("MultiUpstreamManager not initialized".to_string()) + Err("No upstream manager available".to_string()) } } - /// Broadcast a message to all active upstreams - pub async fn broadcast_to_all_upstreams( - &self, - message: PoolExtMessages<'static>, - ) -> Vec { + /// Get count of total and active upstreams + pub async fn get_upstream_counts(&self) -> (usize, usize) { if let Some(ref manager) = self.upstream_manager { - manager.broadcast(message).await - } else { - vec!["MultiUpstreamManager not initialized".to_string()] - } - } - - /// Get aggregated receiver for multi-upstream mode - pub fn get_aggregated_receiver( - &mut self, - ) -> Option>> { - self.aggregated_receiver.take() - } - - /// Get connection statistics for multi-upstream mode - pub async fn get_connection_stats(&self) -> (usize, usize) { - if let Some(ref manager) = self.upstream_manager { - // Get stats from the manager let upstreams = manager.get_upstreams().await; let total = upstreams.len(); let active = upstreams.values().filter(|u| u.is_active).count(); @@ -496,39 +512,9 @@ impl Router { best_pool.map(|pool| (pool, best_latency)) } - /// Get detailed connection statistics with equal distribution - pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { - let mut results = Vec::new(); - - if let Some(ref manager) = self.upstream_manager { - let upstreams = manager.get_upstreams().await; - let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); - - // Count active upstreams first - let active_count = upstreams.values().filter(|upstream| upstream.is_active).count(); - - // Calculate hashrate per upstream (equal distribution) - let hashrate_per_upstream = if active_count > 0 { - total_hashrate / active_count as f32 - } else { - 0.0 - }; - - for (id, upstream) in upstreams { - // Each active upstream gets an equal share - let hashrate = if upstream.is_active { hashrate_per_upstream } else { 0.0 }; - results.push((id, upstream.is_active, hashrate)); - } - } else { - // Single upstream mode (fallback) - if let Some(current) = self.current_pool { - let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); - results.push(("upstream-0".to_string(), true, total_hashrate)); - } - } - - results - } + + + } /// Track latencies for various stages of pool connection setup. diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 77390788..0d9fb1cc 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -20,27 +20,34 @@ pub struct UpstreamConnection { pub struct MultiUpstreamManager { upstreams: Arc>>, aggregated_sender: tokio::sync::mpsc::Sender>, - // Remove round-robin fields - // current_index: Arc>, - // use_round_robin: bool, + aggregated_receiver: Arc>>>>, + // New field for custom hashrate distribution + hashrate_distribution: Arc>>, } impl MultiUpstreamManager { - pub fn new( - _use_parallel: bool, // Parameter kept for compatibility but always use parallel - ) -> (Self, tokio::sync::mpsc::Receiver>) { + pub fn new(use_parallel: bool) -> (Self, tokio::sync::mpsc::Receiver>) { let (sender, receiver) = tokio::sync::mpsc::channel(1000); + let manager = Self { + upstreams: Arc::new(Mutex::new(HashMap::new())), + aggregated_sender: sender, + aggregated_receiver: Arc::new(Mutex::new(None)), // Start with None + hashrate_distribution: Arc::new(Mutex::new(Vec::new())), // Initialize empty + }; + + (manager, receiver) + } - ( - Self { - upstreams: Arc::new(Mutex::new(HashMap::new())), - aggregated_sender: sender, - // Remove round-robin fields - // current_index: Arc::new(Mutex::new(0)), - // use_round_robin: false, // Always use parallel - }, - receiver, - ) + /// Get the aggregated receiver (for use in the main loop) + pub async fn get_aggregated_receiver(&self) -> Option>> { + let mut rx = self.aggregated_receiver.lock().await; + rx.take() + } + + /// Set the aggregated receiver (if needed to put it back) + pub async fn set_aggregated_receiver(&self, receiver: tokio::sync::mpsc::Receiver>) { + let mut rx = self.aggregated_receiver.lock().await; + *rx = Some(receiver); } /// Add a new upstream connection @@ -223,20 +230,6 @@ impl MultiUpstreamManager { } } - // Remove round-robin send method - // async fn send_round_robin(&self, message: PoolExtMessages<'static>) -> Result<(), String> { ... } - - // /// Send to the best upstream (first active one for now) - // async fn send_to_best_upstream(&self, message: PoolExtMessages<'static>) -> Result<(), String> { - // let connections = self.upstreams.lock().await; - // for connection in connections.values() { - // if connection.is_active { - // return connection.sender.send(message).await - // .map_err(|e| format!("Failed to send to upstream {}: {}", connection.id, e)); - // } - // } - // Err("No active upstreams available".to_string()) - // } /// Connect to upstream with retry logic async fn connect_upstream_with_retry( upstreams: Arc>>, @@ -399,4 +392,42 @@ impl MultiUpstreamManager { let connections = self.upstreams.lock().await; connections.values().filter(|conn| conn.is_active).count() } + + /// Set custom hashrate distribution percentages + pub async fn set_hashrate_distribution(&self, distribution: Vec) { + let mut dist = self.hashrate_distribution.lock().await; + *dist = distribution; + info!("Set custom hashrate distribution: {:?}", *dist); + } + + /// Get current hashrate distribution + pub async fn get_hashrate_distribution(&self) -> Vec { + self.hashrate_distribution.lock().await.clone() + } + + /// Validate that distribution percentages add up to 100% + pub fn validate_distribution(distribution: &[f32]) -> Result<(), String> { + if distribution.is_empty() { + return Err("Distribution cannot be empty".to_string()); + } + + let total: f32 = distribution.iter().sum(); + if (total - 100.0).abs() > 0.1 { // Allow small floating point errors + return Err(format!( + "Distribution percentages must add up to 100%, got {:.1}%", + total + )); + } + + for (idx, &percentage) in distribution.iter().enumerate() { + if percentage < 0.0 || percentage > 100.0 { + return Err(format!( + "Invalid percentage at index {}: {:.1}%. Must be between 0% and 100%", + idx, percentage + )); + } + } + + Ok(()) + } } From 5b10e1fd696481e94556148d640094d4a32c8990 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Wed, 4 Jun 2025 15:17:16 +0530 Subject: [PATCH 08/21] Fix type annotation for abort_handles in monitor function --- src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main.rs b/src/main.rs index 3f56e342..7ff65e2b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -616,7 +616,7 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { } async fn monitor( router: &mut Router, - abort_handles: Vec<(std::string::String)>, + abort_handles: Vec, epsilon: Duration, server_handle: tokio::task::JoinHandle<()>, ) -> Reconnect { From 8162d3e049ff3113c99391286e7a96f4ea4c319a Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Thu, 5 Jun 2025 16:59:29 +0530 Subject: [PATCH 09/21] feat(router): Enhance Router functionality for multi-pool support - Added support for multiple upstream pools with different authentication keys. - Introduced latency-based selection for upstream pools. - Updated `new_with_keys` and `new_multi` methods to handle custom distribution. - Implemented `select_best_pools` and `select_best_pools_for_distribution` for optimal pool selection based on latency. - Enhanced logging with additional tracing for better debugging. refactor(multi_upstream_manager): Streamline upstream connection management - Simplified upstream connection handling and removed unnecessary fields. - Updated connection logic to maintain connections based on hashrate distribution. - Improved error handling and logging for upstream connection attempts. fix(config): Introduce comprehensive configuration management - Added a new configuration structure to handle command-line arguments, environment variables, and configuration files. - Implemented parsing for hashrate and socket addresses with error handling. - Ensured backward compatibility with existing configuration setups. chore(translator): Adjust module visibility and clean up imports - Changed visibility of proxy and upstream modules to private. - Cleaned up unnecessary imports and ensured consistent formatting. style(share_accounter): Ensure consistent formatting - Added missing newline at the end of the file for consistency. --- config.toml | 5 +- src/config.rs | 345 ++++++++++++ src/ingress/sv1_ingress.rs | 2 +- src/main.rs | 777 +++++++++++++++------------ src/router/mod.rs | 218 +++++--- src/router/multi_upstream_manager.rs | 545 ++++++------------- src/share_accounter/mod.rs | 2 +- src/translator/mod.rs | 6 +- 8 files changed, 1096 insertions(+), 804 deletions(-) create mode 100644 src/config.rs diff --git a/config.toml b/config.toml index 44cfd441..dae76260 100644 --- a/config.toml +++ b/config.toml @@ -8,6 +8,5 @@ loglevel = "info" nc_loglevel = "off" test = true -# Custom hashrate distribution (percentages for each upstream) -# This array should match the number of pool_addresses -hashrate_distribution = [30.0, 70.0] \ No newline at end of file +# Uncomment this line to enable hashrate distribution +#hashrate_distribution = [30.0, 70.0] \ No newline at end of file diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 00000000..57e957cc --- /dev/null +++ b/src/config.rs @@ -0,0 +1,345 @@ +use clap::Parser; +use lazy_static::lazy_static; +use serde::{Deserialize, Serialize}; +use std::{ + net::{SocketAddr, ToSocketAddrs}, + path::PathBuf, +}; +use tracing::{error, info, warn}; + +use crate::{HashUnit, DEFAULT_SV1_HASHPOWER}; +lazy_static! { + pub static ref CONFIG: Configuration = Configuration::load_config(); +} +#[derive(Parser)] +struct Args { + #[clap(long)] + test: bool, + #[clap(long = "d", short = 'd', value_parser = parse_hashrate)] + downstream_hashrate: Option, + #[clap(long = "loglevel", short = 'l')] + loglevel: Option, + #[clap(long = "nc", short = 'n')] + noise_connection_log: Option, + #[clap(long = "delay")] + delay: Option, + #[clap(long = "interval", short = 'i')] + adjustment_interval: Option, + #[clap(long = "pool", short = 'p', value_delimiter = ',')] + pool_addresses: Option>, + #[clap(long = "test-pool", value_delimiter = ',')] + test_pool_addresses: Option>, + #[clap(long)] + token: Option, + #[clap(long)] + tp_address: Option, + #[clap(long)] + listening_addr: Option, + #[clap(long = "config", short = 'c')] + config_file: Option, + #[clap(long = "api-server-port", short = 's')] + api_server_port: Option, +} + +#[derive(Serialize, Deserialize)] +struct ConfigFile { + token: Option, + tp_address: Option, + pool_addresses: Option>, + test_pool_addresses: Option>, + interval: Option, + delay: Option, + downstream_hashrate: Option, + loglevel: Option, + nc_loglevel: Option, + test: Option, + listening_addr: Option, + api_server_port: Option, +} + +pub struct Configuration { + token: Option, + tp_address: Option, + pool_addresses: Option>, + test_pool_addresses: Option>, + interval: u64, + delay: u64, + downstream_hashrate: f32, + loglevel: String, + nc_loglevel: String, + test: bool, + listening_addr: Option, + api_server_port: String, +} +impl Configuration { + pub fn token() -> Option { + CONFIG.token.clone() + } + + pub fn tp_address() -> Option { + CONFIG.tp_address.clone() + } + + pub fn pool_address() -> Option> { + if CONFIG.test { + CONFIG.test_pool_addresses.clone() // Return test pool addresses in test mode + } else { + CONFIG.pool_addresses.clone() + } + } + + pub fn adjustment_interval() -> u64 { + CONFIG.interval + } + + pub fn delay() -> u64 { + CONFIG.delay + } + + pub fn downstream_hashrate() -> f32 { + CONFIG.downstream_hashrate + } + + pub fn downstream_listening_addr() -> Option { + CONFIG.listening_addr.clone() + } + pub fn api_server_port() -> String { + CONFIG.api_server_port.clone() + } + + pub fn loglevel() -> &'static str { + match CONFIG.loglevel.to_lowercase().as_str() { + "trace" | "debug" | "info" | "warn" | "error" | "off" => &CONFIG.loglevel, + _ => { + eprintln!( + "Invalid log level '{}'. Defaulting to 'info'.", + CONFIG.loglevel + ); + "info" + } + } + } + + pub fn nc_loglevel() -> &'static str { + match CONFIG.nc_loglevel.as_str() { + "trace" | "debug" | "info" | "warn" | "error" | "off" => &CONFIG.nc_loglevel, + _ => { + eprintln!( + "Invalid log level for noise_connection '{}' Defaulting to 'off'.", + &CONFIG.nc_loglevel + ); + "off" + } + } + } + + pub fn test() -> bool { + CONFIG.test + } + + // Loads config from CLI, file, or env vars with precedence: CLI > file > env. + fn load_config() -> Self { + let args = Args::parse(); + let config_path: PathBuf = args.config_file.unwrap_or("config.toml".into()); + let config: ConfigFile = std::fs::read_to_string(&config_path) + .ok() + .and_then(|content| toml::from_str(&content).ok()) + .unwrap_or(ConfigFile { + token: None, + tp_address: None, + pool_addresses: None, + test_pool_addresses: None, + interval: None, + delay: None, + downstream_hashrate: None, + loglevel: None, + nc_loglevel: None, + test: None, + listening_addr: None, + api_server_port: None, + }); + + let token = args + .token + .or(config.token) + .or_else(|| std::env::var("TOKEN").ok()); + + let tp_address = args + .tp_address + .or(config.tp_address) + .or_else(|| std::env::var("TP_ADDRESS").ok()); + + let pool_addresses: Option> = args + .pool_addresses + .map(|addresses| { + addresses + .into_iter() + .map(parse_address) + .collect::>() + }) + .or_else(|| { + config.pool_addresses.map(|addresses| { + addresses + .into_iter() + .map(parse_address) + .collect::>() + }) + }) + .or_else(|| { + std::env::var("POOL_ADDRESSES").ok().map(|s| { + s.split(',') + .map(|s| parse_address(s.trim().to_string())) + .collect::>() + }) + }); + + let test_pool_addresses: Option> = args + .test_pool_addresses + .map(|addresses| { + addresses + .into_iter() + .map(parse_address) + .collect::>() + }) + .or_else(|| { + config.test_pool_addresses.map(|addresses| { + addresses + .into_iter() + .map(parse_address) + .collect::>() + }) + }) + .or_else(|| { + std::env::var("TEST_POOL_ADDRESSES").ok().map(|s| { + s.split(',') + .map(|s| parse_address(s.trim().to_string())) + .collect::>() + }) + }); + + let interval = args + .adjustment_interval + .or(config.interval) + .or_else(|| std::env::var("INTERVAL").ok().and_then(|s| s.parse().ok())) + .unwrap_or(120_000); + + let delay = args + .delay + .or(config.delay) + .or_else(|| std::env::var("DELAY").ok().and_then(|s| s.parse().ok())) + .unwrap_or(0); + + let expected_hashrate = args + .downstream_hashrate + .or_else(|| { + config + .downstream_hashrate + .as_deref() + .and_then(|d| parse_hashrate(d).ok()) + }) + .or_else(|| { + std::env::var("DOWNSTREAM_HASHRATE") + .ok() + .and_then(|s| s.parse().ok()) + }); + let downstream_hashrate; + if let Some(hashpower) = expected_hashrate { + downstream_hashrate = hashpower; + info!( + "Using downstream hashrate: {}h/s", + HashUnit::format_value(hashpower) + ); + } else { + downstream_hashrate = DEFAULT_SV1_HASHPOWER; + warn!( + "No downstream hashrate provided, using default value: {}h/s", + HashUnit::format_value(DEFAULT_SV1_HASHPOWER) + ); + } + + let listening_addr = args.listening_addr.or(config.listening_addr).or_else(|| { + std::env::var("DOWNSTREAM_HASHRATE") + .ok() + .and_then(|s| s.parse().ok()) + }); + let api_server_port = args + .api_server_port + .or(config.api_server_port) + .or_else(|| { + std::env::var("API_SERVER_PORT") + .ok() + .and_then(|s| s.parse().ok()) + }) + .unwrap_or("3001".to_string()); + + let loglevel = args + .loglevel + .or(config.loglevel) + .or_else(|| std::env::var("LOGLEVEL").ok()) + .unwrap_or("info".to_string()); + + let nc_loglevel = args + .noise_connection_log + .or(config.nc_loglevel) + .or_else(|| std::env::var("NC_LOGLEVEL").ok()) + .unwrap_or("off".to_string()); + + let test = args.test || config.test.unwrap_or(false) || std::env::var("TEST").is_ok(); + + Configuration { + token, + tp_address, + pool_addresses, + test_pool_addresses, + interval, + delay, + downstream_hashrate, + loglevel, + nc_loglevel, + test, + listening_addr, + api_server_port, + } + } +} + +/// Parses a hashrate string (e.g., "10T", "2.5P", "500E") into an f32 value in h/s. +fn parse_hashrate(hashrate_str: &str) -> Result { + let hashrate_str = hashrate_str.trim(); + if hashrate_str.is_empty() { + return Err("Hashrate cannot be empty. Expected format: '' (e.g., '10T', '2.5P', '5E'".to_string()); + } + + let unit = hashrate_str.chars().last().unwrap_or(' ').to_string(); + let num = &hashrate_str[..hashrate_str.len().saturating_sub(1)]; + + let num: f32 = num.parse().map_err(|_| { + format!( + "Invalid number '{}'. Expected format: '' (e.g., '10T', '2.5P', '5E')", + num + ) + })?; + + let multiplier = HashUnit::from_str(&unit) + .map(|unit| unit.multiplier()) + .ok_or_else(|| format!( + "Invalid unit '{}'. Expected 'T' (Terahash), 'P' (Petahash), or 'E' (Exahash). Example: '10T', '2.5P', '5E'", + unit + ))?; + + let hashrate = num * multiplier; + + if hashrate.is_infinite() || hashrate.is_nan() { + return Err("Hashrate too large or invalid".to_string()); + } + + Ok(hashrate) +} + +fn parse_address(addr: String) -> SocketAddr { + addr.to_socket_addrs() + .map_err(|e| error!("Invalid socket address: {}", e)) + .expect("Failed to parse socket address") + .next() + .expect("No socket address resolved") +} \ No newline at end of file diff --git a/src/ingress/sv1_ingress.rs b/src/ingress/sv1_ingress.rs index 4bd2c6f8..1404bb1c 100644 --- a/src/ingress/sv1_ingress.rs +++ b/src/ingress/sv1_ingress.rs @@ -117,4 +117,4 @@ impl Downstream { Err(_) => Sv1IngressError::TaskFailed, } } -} +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 7ff65e2b..a13210bd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,21 +1,24 @@ -use clap::{ArgAction, Parser}; #[cfg(not(target_os = "windows"))] use jemallocator::Jemalloc; use router::Router; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -// Add these missing imports +use clap::{ArgAction, Parser}; use serde::{Deserialize, Serialize}; use std::path::Path; - +use crate::shared::utils::AbortOnDrop; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; + use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; -use proxy_state::{PoolState, ProxyState}; -use std::{net::ToSocketAddrs, time::Duration}; +use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; +use std::{net::SocketAddr, time::Duration}; +use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; mod api; + +mod config; mod ingress; pub mod jd_client; mod minin_pool_connection; @@ -25,6 +28,7 @@ mod share_accounter; mod shared; mod translator; + const TRANSLATOR_BUFFER_SIZE: usize = 32; const MIN_EXTRANONCE_SIZE: u16 = 6; const MIN_EXTRANONCE2_SIZE: u16 = 5; @@ -36,7 +40,7 @@ const MAX_LEN_DOWN_MSG: u32 = 10000; const MAIN_POOL_ADDRESS: &str = "mining.dmnd.work:2000"; //const TEST_POOL_ADDRESS: &str = "127.0.0.1:20000"; const TEST_POOL_ADDRESS: &str = "18.193.252.132:2000"; -const MAIN_AUTH_PUB_KEY: &str = "9bQHWXsQ2J9TRFTaxRh3KjoxdyLRfWVEy25YHtKF8y8gotLoCZZ"; +const MAIN_AUTH_PUB_KEY: &str = "9bQHWsQ2J9TRFTaxRh3KjoxdyLRfWVEy25YHtKF8y8gotLoCZZ"; const TEST_AUTH_PUB_KEY: &str = "9auqWEzQDVyd2oe1JVGFLMLHZtCo2FFqZwtKA5gd9xbuEu7PH72"; //const TP_ADDRESS: &str = "127.0.0.1:8442"; const DEFAULT_LISTEN_ADDRESS: &str = "0.0.0.0:32767"; @@ -74,10 +78,11 @@ lazy_static! { MAIN_AUTH_PUB_KEY }; } + #[derive(Parser, Debug)] #[clap(author, version, about, long_about = None)] pub struct Args { - // Use test enpoint if test flag is provided + // Use test #[clap(long)] test: bool, #[clap(long ="d", short ='d', value_parser = parse_hashrate)] @@ -90,17 +95,11 @@ pub struct Args { delay: u64, #[clap(long = "interval", short = 'i', default_value = "120000")] adjustment_interval: u64, - // New argument for multiple upstream servers + // Add the missing upstream field #[clap(long = "upstream", short = 'u', action = ArgAction::Append)] - // In Args struct - // #[clap(long = "monitor-hashrate", short = 'm')] - // monitor_hashrate: bool, + upstream: Option>, #[clap(long = "config", short = 'c')] pub config_file: Option, - - /// Enable parallel upstream usage (sends to all upstreams simultaneously) - #[clap(long = "parallel", short = 'p')] - pub parallel: bool, } #[derive(Debug, Deserialize, Serialize, Clone)] @@ -114,8 +113,7 @@ pub struct Config { pub loglevel: Option, pub nc_loglevel: Option, pub test: Option, - pub hashrate_distribution: Option>, // Array of percentages for each upstream - + pub hashrate_distribution: Option>, // New field for custom distribution } impl Config { @@ -164,7 +162,7 @@ async fn main() { .init(); // std::env::var("TOKEN").expect("Missing TOKEN environment variable"); - // Load configuration first + // Load configuration and get pool addresses let config = if let Some(config_path) = &ARGS.config_file { match Config::from_file(config_path) { Ok(config) => Some(config), @@ -174,208 +172,198 @@ async fn main() { } } } else { - info!("No config file provided, using default settings"); None }; - // Get TOKEN from config or environment - let _token = if let Some(ref config) = config { - config.token.clone() - } else { - std::env::var("TOKEN").unwrap_or_else(|_| { - error!("TOKEN must be provided either in config file or as environment variable"); - std::process::exit(1); - }) - }; - - // Use configured hashrate if available - let hashpower = if let Some(ref config) = config { + println!("Config: {:?}", config); + + // Set the total hashrate from config BEFORE creating router + if let Some(ref config) = config { if let Some(ref hashrate_str) = config.downstream_hashrate { - parse_hashrate(hashrate_str).unwrap_or(*EXPECTED_SV1_HASHPOWER) - } else { - *EXPECTED_SV1_HASHPOWER - } - } else { - *EXPECTED_SV1_HASHPOWER - }; - - ProxyState::set_total_hashrate(hashpower); - - if args.downstream_hashrate.is_some() { - info!( - "Using downstream hashrate: {}h/s", - HashUnit::format_value(hashpower) - ); - } else { - warn!( - "No downstream hashrate provided, using default value: {}h/s", - HashUnit::format_value(hashpower) - ); - } - - // Handle multiple upstream servers from config or hard-coded fallback - let mut pool_addresses = Vec::new(); - let mut custom_distribution: Option> = None; - - if let Some(ref config) = config { - // Use configuration file pool addresses - info!("Loading pool addresses from configuration file"); - - // Load custom distribution if provided - if let Some(ref distribution) = config.hashrate_distribution { - info!("Found custom hashrate distribution in config: {:?}", distribution); - - // Validate the distribution - if let Err(e) = router::MultiUpstreamManager::validate_distribution(distribution) { - error!("Invalid hashrate distribution in config: {}", e); - std::process::exit(1); + match parse_hashrate(hashrate_str) { + Ok(hashrate) => { + ProxyState::set_total_hashrate(hashrate); + info!("Set total hashrate to: {}", HashUnit::format_value(hashrate)); + } + Err(e) => { + error!("Failed to parse downstream_hashrate: {}", e); + } } - - custom_distribution = Some(distribution.clone()); } + } else if let Some(hashrate) = ARGS.downstream_hashrate { + ProxyState::set_total_hashrate(hashrate); + info!("Set total hashrate from args to: {}", HashUnit::format_value(hashrate)); + } - for (idx, pool_addr_str) in config.pool_addresses.iter().enumerate() { - // Parse address - let addr = if let Ok(mut addrs) = pool_addr_str.to_socket_addrs() { - addrs.next().unwrap_or_else(|| { - error!("Failed to resolve address: {}", pool_addr_str); + // Get pool addresses with auth keys + let pool_address_keys = if let Some(ref config) = config { + config.pool_addresses.iter().map(|addr_str| { + let addr = addr_str.parse::() + .unwrap_or_else(|_| { + error!("Invalid pool address: {}", addr_str); std::process::exit(1); - }) + }); + + let auth_key: Secp256k1PublicKey = if config.test.unwrap_or(false) { + TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key") } else { - error!("Invalid pool address: {}", pool_addr_str); - std::process::exit(1); + MAIN_AUTH_PUB_KEY.parse().expect("Invalid main public key") }; - - let pubkey: Secp256k1PublicKey = - TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key"); - - info!("Adding upstream server {}: {}", idx + 1, addr); - pool_addresses.push((addr, pubkey)); - - // ONLY add to ProxyState, don't create connections yet - ProxyState::add_upstream_connection( - format!("upstream-{}", idx), - format!("config-pool-{}", idx + 1), - addr, - pubkey, - crate::proxy_state::UpstreamType::JDCMiningUpstream, - ); - } + + (addr, auth_key) + }).collect() + } else { + // Fallback to single upstream + let addr = if ARGS.test { + TEST_POOL_ADDRESS.parse::().expect("Invalid test address") + } else { + MAIN_POOL_ADDRESS.parse::().expect("Invalid main address") + }; - // Validate that distribution count matches pool count - if let Some(ref distribution) = custom_distribution { - if distribution.len() != pool_addresses.len() { - error!( - "Hashrate distribution count ({}) doesn't match pool addresses count ({})", - distribution.len(), - pool_addresses.len() - ); + let auth_key = if ARGS.test { + TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key") + } else { + MAIN_AUTH_PUB_KEY.parse().expect("Invalid main public key") + }; + + vec![(addr, auth_key)] + }; + + // Determine if we want hashrate distribution + let wants_distribution = config.as_ref() + .and_then(|c| c.hashrate_distribution.as_ref()) + .is_some() || ARGS.parallel; + + // Create router based on configuration + let mut router = if wants_distribution { + // User wants hashrate distribution - use multi-upstream mode + let pool_address_keys_clone = pool_address_keys.clone(); + match Router::new_multi( + pool_address_keys_clone, + None, // setup_connection_msg + None, // timer + true, // use_distribution = true + ).await { + Ok(router) => { + info!("Created multi-upstream router with {} upstreams for hashrate distribution", + pool_address_keys.len()); + router + } + Err(e) => { + error!("Failed to create multi-upstream router: {}", e); std::process::exit(1); } } - + } else if pool_address_keys.len() > 1 { + // Multiple pools but no distribution specified - use latency-based selection (single upstream mode) + info!("Created latency-based router with {} pools (will select best latency)", pool_address_keys.len()); + Router::new_with_keys( + pool_address_keys.clone(), + None, // setup_connection_msg + None, // timer + ) } else { - // Fallback to hard-coded address - info!("Using hard-coded fallback pool address"); + // Single pool - use single upstream mode + let (addr, auth_key) = pool_address_keys[0]; + info!("Created single upstream router"); + Router::new( + vec![addr], + auth_key, + None, // setup_connection_msg + None, // timer + ) + }; - let test_addr = match "3.74.36.119:2000".to_socket_addrs() { - Ok(mut addrs) => addrs.next().unwrap_or_else(|| { - error!("Failed to resolve fallback pool address"); - std::process::exit(1); - }), - Err(_) => { - error!("Invalid fallback pool address format"); - std::process::exit(1); - } - }; + // Add this right after router creation, before the if statements: + let epsilon = Duration::from_millis(10); - let test_pubkey: Secp256k1PublicKey = TEST_AUTH_PUB_KEY - .parse() - .expect("Invalid fallback public key"); + // Handle the three different scenarios + if wants_distribution { + info!("=== HASHRATE DISTRIBUTION MODE ==="); + + // Get the desired distribution + let distribution = config.as_ref() + .and_then(|c| c.hashrate_distribution.clone()) + .unwrap_or_else(|| { + // Default equal distribution if parallel flag is set but no custom distribution + let count = pool_address_keys.len(); + let equal_percentage = 100.0 / count as f32; + vec![equal_percentage; count] + }); + println!("Using hashrate distribution: {:?}", distribution); - info!("Adding fallback upstream server: {}", test_addr); - pool_addresses.push((test_addr, test_pubkey)); + let desired_pool_count = distribution.len(); + info!("Desired pool count for distribution: {}", desired_pool_count); - ProxyState::add_upstream_connection( - "upstream-0".to_string(), - "fallback-pool-1".to_string(), - test_addr, - test_pubkey, - crate::proxy_state::UpstreamType::JDCMiningUpstream, - ); - } + // Select best pools based on latency (implements condition 2 and 3) + let selected_pools = if pool_address_keys.len() <= desired_pool_count { + // Condition 3: Use all available pools (skip latency testing) + info!("Using all {} available pools (condition 3)", pool_address_keys.len()); + pool_address_keys.clone() + } else { + // Condition 2: Select best pools based on latency + info!("Selecting {} best pools from {} available (condition 2)", desired_pool_count, pool_address_keys.len()); + + // Create temporary router for latency testing + let mut temp_router = Router::new_with_keys( + pool_address_keys.clone(), + None, + None, + ); + + temp_router.select_best_pools_for_distribution(desired_pool_count).await + }; - // Direct verification - info!("DIRECT VERIFICATION: Checking hashrate distribution"); - let upstream_connections = ProxyState::get_upstream_connections(); - let active_count = upstream_connections.len(); - info!("Active upstreams: {}", active_count); - - // Remove round-robin hashrate calculation - // if active_count > 0 && ARGS.round_robin { - // let hashrate_per_upstream = hashpower / active_count as f32; - // info!( - // "Round-robin hashrate per upstream: {}", - // HashUnit::format_value(hashrate_per_upstream) - // ); - // - // let upstream_connections: Vec<( - // String, - // std::net::SocketAddr, - // key_utils::Secp256k1PublicKey, - // )> = ProxyState::get_upstream_connections(); - // - // // And update the usage in the loop below: - // for (upstream_id, addr, _auth_key) in upstream_connections { - // info!( - // "Upstream {}: {} - allocated {}", - // upstream_id, - // addr, - // HashUnit::format_value(hashrate_per_upstream) - // ); - // } - // } - - // Create the router - always use multi upstream with custom distribution mode - let mut router = if !pool_addresses.is_empty() { - // Always use multi-upstream mode (even for single upstream) - match router::Router::new_multi( - pool_addresses.clone(), - None, // setup_connection_msg - None, // timer - false, // Changed from true to false - disable parallel mode for custom distribution - ) - .await - { - Ok(mut router) => { - // Set custom distribution if provided - if let Some(distribution) = custom_distribution { - info!("Setting custom hashrate distribution: {:?}", distribution); - if let Err(e) = router.set_hashrate_distribution(distribution).await { - error!("Failed to set custom hashrate distribution: {}", e); - std::process::exit(1); - } - } - - if let Err(e) = router.initialize_upstream_connections().await { - error!("Failed to initialize upstream connections: {}", e); - std::process::exit(1); - } + // Now create the final multi-upstream router with selected pools + router = match Router::new_multi( + selected_pools, + None, + None, + true, + ).await { + Ok(router) => { + info!("Created distribution router with {} selected pools", desired_pool_count); router } Err(e) => { - error!("Failed to create multi-upstream router: {}", e); + error!("Failed to create distribution router: {}", e); std::process::exit(1); } + }; + + // Initialize upstream connections + if let Err(e) = router.initialize_upstream_connections().await { + error!("Failed to initialize upstream connections: {}", e); + std::process::exit(1); } - } else { - error!("No pool addresses configured. Please provide pool addresses via config file."); - std::process::exit(1); - }; + // Set the distribution + if let Err(e) = router.set_hashrate_distribution(distribution.clone()).await { + error!("Failed to set hashrate distribution: {}", e); + std::process::exit(1); + } + info!("Set hashrate distribution: {:?}", distribution); - let epsilon = Duration::from_millis(100); + + info!("🔧 About to call initialize_proxy"); + initialize_proxy(&mut router, None, epsilon).await; + } else if pool_address_keys.len() > 1 { + info!("=== LATENCY-BASED SELECTION MODE (Condition 1) ==="); + info!("Multiple pools available, will select best latency"); + + // Actually select the best latency pool + let best_upstream = router.select_pool_connect().await; + if best_upstream.is_some() { + info!("Selected best upstream: {:?}", best_upstream); + } + initialize_proxy(&mut router, best_upstream, epsilon).await; + } else { + info!("=== SINGLE UPSTREAM MODE ==="); + + // Single pool mode + let best_upstream = router.select_pool_connect().await; + initialize_proxy(&mut router, best_upstream, epsilon).await; + } - // Always use multi-upstream mode - initialize_proxy(&mut router, None, epsilon).await; + info!("Proxy initialization complete"); // Wait a moment for connections to establish tokio::time::sleep(Duration::from_secs(2)).await; @@ -402,7 +390,8 @@ async fn main() { info!(" 🔗 {}", conn.trim()); } } - // Show hashrate distribution + + // Show hashrate distribution let total_hashrate = ProxyState::get_total_hashrate(); let detailed_stats = router.get_detailed_connection_stats().await; @@ -436,6 +425,7 @@ async fn main() { for (upstream_id, is_active, hashrate) in detailed_stats { if is_active { + info!("{} = hashrate", hashrate); let percentage = (hashrate / total_hashrate) * 100.0; info!( " ✅ {}: receiving {} ({:.1}% of total)", @@ -451,138 +441,271 @@ async fn main() { ); } } - } else { - warn!("❌ No active upstream connections!"); - } - info!("========================================"); - // Start monitoring task for multi-upstream - let router_clone = router.clone(); - tokio::spawn(async move { - monitor_multi_upstream(router_clone, epsilon).await; - }); - - // Keep main thread alive - loop { - tokio::time::sleep(Duration::from_secs(60)).await; - info!("Multi-upstream proxy running..."); + info!("========================================"); + + // Start monitoring task for multi-upstream + let router_clone = router.clone(); + tokio::spawn(async move { + monitor_multi_upstream(router_clone, epsilon).await; + }); + + // Keep main thread alive + loop { + tokio::time::sleep(Duration::from_secs(60)).await; + info!("Multi-upstream proxy running..."); + } } } async fn initialize_proxy( router: &mut Router, - mut _pool_addr: Option, // Add underscore to indicate unused - _epsilon: Duration, // Add underscore to indicate unused + mut pool_addr: Option, + epsilon: Duration, ) { - // Add a static flag to prevent multiple downstream listeners - static DOWNSTREAM_STARTED: std::sync::atomic::AtomicBool = - std::sync::atomic::AtomicBool::new(false); - - // For multi-upstream mode, don't loop - just set up and wait + // Check if we're in multi-upstream mode if router.is_multi_upstream_enabled() { - info!("Multi-upstream mode: using aggregated message handling"); - - // Only start downstream listener once - if !DOWNSTREAM_STARTED.swap(true, std::sync::atomic::Ordering::SeqCst) { - info!("Starting downstream listener (multi-upstream mode)"); - - // Create channel for downstream connections - let (downstreams_tx, mut downstreams_rx) = tokio::sync::mpsc::channel(10); - - // Start the downstream listener - let _abort_handle = ingress::sv1_ingress::start_listen_for_downstream(downstreams_tx); - - // Handle incoming downstream connections - tokio::spawn(async move { - while let Some((send_to_downstream, recv_from_downstream, client_addr)) = - downstreams_rx.recv().await - { - info!("New downstream client connected: {}", client_addr); - - // Here you would typically start the translator for this downstream connection - // For now, we'll just log the connection - tokio::spawn(async move { - // Handle this specific downstream connection - // You'll need to integrate this with your translator logic - let mut recv = recv_from_downstream; - while let Some(message) = recv.recv().await { - info!("Received from downstream {}: {}", client_addr, message); - // Process the message and potentially send to upstreams - } - info!("Downstream client {} disconnected", client_addr); - }); + info!("Initializing proxy in HASHRATE DISTRIBUTION mode"); + info!("🎯 Focus: Distribute hashrate allocation across upstreams"); + + // Start API server for monitoring + let stats_sender = api::stats::StatsSender::new(); + let _server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); + + // Start the monitor task in the background + let router_clone = router.clone(); + tokio::spawn(async move { + monitor_multi_upstream(router_clone, epsilon).await; + }); + + info!("✅ Hashrate distribution monitor started"); + return; + } + + // Single upstream mode only + loop { + let stats_sender = api::stats::StatsSender::new(); + + let (send_to_pool, recv_from_pool, pool_connection_abortable) = + match router.connect_pool(pool_addr).await { + Ok(connection) => connection, + Err(_) => { + error!("No upstream available. Retrying..."); + warn!("Are you using the correct TOKEN??"); + let mut secs = 10; + while secs > 0 { + tracing::warn!("Retrying in {} seconds...", secs); + tokio::time::sleep(Duration::from_secs(1)).await; + secs -= 1; + } + continue; } - }); - } + }; - // Get aggregated receiver if available - if let Some(aggregated_receiver) = router.get_aggregated_receiver().await { - tokio::spawn(async move { - let mut receiver = aggregated_receiver; - while let Some(message) = receiver.recv().await { - info!("Received aggregated message from upstream: {:?}", message); + let (downs_sv1_tx, downs_sv1_rx) = channel(10); + let sv1_ingress_abortable = ingress::sv1_ingress::start_listen_for_downstream(downs_sv1_tx); + + let (translator_up_tx, mut translator_up_rx) = channel(10); + let translator_abortable = + match translator::start(downs_sv1_rx, translator_up_tx, stats_sender.clone()).await { + Ok(abortable) => abortable, + Err(e) => { + error!("Impossible to initialize translator: {e}"); + // Impossible to start the proxy so we restart proxy + ProxyState::update_translator_state(TranslatorState::Down); + ProxyState::update_tp_state(TpState::Down); + return; } - }); - } + }; - // Instead of looping, just wait indefinitely - info!("Multi-upstream proxy initialized successfully"); - return; // Exit the function, don't loop - } + let (from_jdc_to_share_accounter_send, from_jdc_to_share_accounter_recv) = channel(10); + let (from_share_accounter_to_jdc_send, from_share_accounter_to_jdc_recv) = channel(10); + let (jdc_to_translator_sender, jdc_from_translator_receiver, _) = translator_up_rx + .recv() + .await + .expect("Translator failed before initialization"); + + let jdc_abortable: Option; + let share_accounter_abortable; + let tp = match TP_ADDRESS.safe_lock(|tp| tp.clone()) { + Ok(tp) => tp, + Err(e) => { + error!("TP_ADDRESS Mutex Corrupted: {e}"); + return; + } + }; - // For single upstream mode - just log and exit since we're not supporting it - error!("Single upstream mode is not supported in this version. Please use multi-upstream mode with --config option."); - std::process::exit(1); -} + if let Some(_tp_addr) = tp { + jdc_abortable = jd_client::start( + jdc_from_translator_receiver, + jdc_to_translator_sender, + from_share_accounter_to_jdc_recv, + from_jdc_to_share_accounter_send, + ) + .await; + if jdc_abortable.is_none() { + ProxyState::update_tp_state(TpState::Down); + }; + share_accounter_abortable = match share_accounter::start( + from_jdc_to_share_accounter_recv, + from_share_accounter_to_jdc_send, + recv_from_pool, + send_to_pool, + ) + .await + { + Ok(abortable) => abortable, + Err(_) => { + error!("Failed to start share_accounter"); + return; + } + } + } else { + jdc_abortable = None; + + share_accounter_abortable = match share_accounter::start( + jdc_from_translator_receiver, + jdc_to_translator_sender, + recv_from_pool, + send_to_pool, + ) + .await + { + Ok(abortable) => abortable, + Err(_) => { + error!("Failed to start share_accounter"); + return; + } + }; + }; -async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { - let mut distribution_check_counter = 0; + // Collecting all abort handles + let mut abort_handles = vec![ + (pool_connection_abortable, "pool_connection".to_string()), + (sv1_ingress_abortable, "sv1_ingress".to_string()), + (translator_abortable, "translator".to_string()), + (share_accounter_abortable, "share_accounter".to_string()), + ]; + if let Some(jdc_handle) = jdc_abortable { + abort_handles.push((jdc_handle, "jdc".to_string())); + } + let server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); + match monitor(router, abort_handles, epsilon, server_handle).await { + Reconnect::NewUpstream(new_pool_addr) => { + ProxyState::update_proxy_state_up(); + pool_addr = Some(new_pool_addr); + continue; + } + Reconnect::NoUpstream => { + ProxyState::update_proxy_state_up(); + pool_addr = None; + continue; + } + } + } +} +async fn monitor( + router: &mut Router, + abort_handles: Vec<(AbortOnDrop, std::string::String)>, + epsilon: Duration, + server_handle: tokio::task::JoinHandle<()>, +) -> Reconnect { + let mut should_check_upstreams_latency = 0; loop { - tokio::time::sleep(Duration::from_secs(10)).await; - distribution_check_counter += 1; + // Check if a better upstream exist every 100 seconds + if should_check_upstreams_latency == 10 * 100 { + should_check_upstreams_latency = 0; + if let Some(new_upstream) = router.monitor_upstream(epsilon).await { + info!("Faster upstream detected. Reinitializing proxy..."); + drop(abort_handles); + server_handle.abort(); // abort server + + // Needs a little to time to drop + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + return Reconnect::NewUpstream(new_upstream); + } + } - // Show report every 30 seconds - if distribution_check_counter >= 3 { - distribution_check_counter = 0; + // Monitor finished tasks + if let Some((_handle, name)) = abort_handles + .iter() + .find(|(handle, _name)| handle.is_finished()) + { + error!("Task {:?} finished, Closing connection", name); + for (handle, _name) in abort_handles { + drop(handle); + } + server_handle.abort(); // abort server - let total_hashrate = ProxyState::get_total_hashrate(); - let detailed_stats = router.get_detailed_connection_stats().await; + // Check if the proxy state is down, and if so, reinitialize the proxy. + let is_proxy_down = ProxyState::is_proxy_down(); + if is_proxy_down.0 { + error!( + "Status: {:?}. Reinitializing proxy...", + is_proxy_down.1.unwrap_or("Proxy".to_string()) + ); + return Reconnect::NoUpstream; + } else { + return Reconnect::NoUpstream; + } + } - info!("📊 === HASHRATE DISTRIBUTION REPORT ==="); - info!( - "🔋 Total configured hashrate: {}", - HashUnit::format_value(total_hashrate) + // Check if the proxy state is down, and if so, reinitialize the proxy. + let is_proxy_down = ProxyState::is_proxy_down(); + if is_proxy_down.0 { + error!( + "{:?} is DOWN. Reinitializing proxy...", + is_proxy_down.1.unwrap_or("Proxy".to_string()) ); + drop(abort_handles); // Drop all abort handles + server_handle.abort(); // abort server + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; // Needs a little to time to drop + return Reconnect::NoUpstream; + } + + should_check_upstreams_latency += 1; + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + } +} + +// Consolidated monitoring function for multi-upstream mode +async fn monitor_multi_upstream(router: Router, epsilon: Duration) { + let mut stats_report_counter = 0; + + loop { + tokio::time::sleep(Duration::from_secs(10)).await; + stats_report_counter += 1; + + // Generate periodic reports every 30 seconds (3 * 10 seconds) + if stats_report_counter >= 3 { + stats_report_counter = 0; + + let (total_count, active_count) = router.get_upstream_counts().await; + let total_hashrate = ProxyState::get_total_hashrate(); - let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); + info!("🚀 === MULTI-UPSTREAM HASHRATE REPORT ==="); + info!("📊 Total upstreams: {}, Active: {}", total_count, active_count); + info!("🔋 Total hashrate: {}", HashUnit::format_value(total_hashrate)); if active_count > 0 { - // Check if using custom distribution - let mut is_custom_distribution = false; + let detailed_stats = router.get_detailed_connection_stats().await; + let distribution = router.get_hashrate_distribution().await; - for (_, is_active, hashrate) in &detailed_stats { - if *is_active { - let percentage = (*hashrate / total_hashrate) * 100.0; - if (percentage - (100.0 / active_count as f32)).abs() > 1.0 { - is_custom_distribution = true; - } - } - } + // Check if using custom distribution + let is_custom = !distribution.is_empty() && distribution.len() == detailed_stats.len(); - if is_custom_distribution { - info!("🎯 Custom Distribution mode: Using configured percentages"); + if is_custom { + info!("🎯 Distribution mode: Custom"); + info!("📈 Distribution percentages: {:?}", distribution); } else { - let percentage_per_upstream = 100.0 / active_count as f32; - info!("🌐 Equal Distribution mode: Each upstream gets {:.1}% of total hashrate", percentage_per_upstream); + info!("⚖️ Distribution mode: Equal (automatic)"); } - - info!("📡 Upstream details:"); for (upstream_id, is_active, hashrate) in detailed_stats { if is_active { let percentage = (hashrate / total_hashrate) * 100.0; info!( - " ✅ {}: {} ({:.1}% of total)", + " ✅ {}: receiving {} ({:.1}% of total)", upstream_id, HashUnit::format_value(hashrate), percentage @@ -595,67 +718,17 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { ); } } - - info!( - "🚀 Total hashrate distributed: {} across {} upstreams", - HashUnit::format_value(total_hashrate), - active_count - ); - - if is_custom_distribution { - info!("💡 Note: Using custom distribution from config file"); - } else { - info!("💡 Note: Hashrate is split equally among all active upstreams"); - } } else { - warn!("⚠️ WARNING: No active upstream connections!"); + warn!("❌ No active upstream connections!"); } - info!("=========================================="); + info!("========================================="); } - } -} -async fn monitor( - router: &mut Router, - abort_handles: Vec, - epsilon: Duration, - server_handle: tokio::task::JoinHandle<()>, -) -> Reconnect { - // Since we only support multi-upstream mode now, this monitor function - // is simplified to just handle multi-upstream monitoring - loop { - tokio::time::sleep(Duration::from_secs(30)).await; - - // Generate periodic reports for multi-upstream - // Replace with a method that exists, e.g., get_detailed_connection_stats - let detailed_stats = router.get_detailed_connection_stats().await; - let total_connections = detailed_stats.len(); - let active_connections = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); - let total_hashrate = ProxyState::get_total_hashrate(); - - info!( - "Multi-upstream mode: {} total, {} active connections", - total_connections, active_connections - ); - - if active_connections > 0 { - info!("Parallel mode: Using ALL upstreams simultaneously"); - info!( - "Total hashrate distributed across {} upstreams: {}", - active_connections, - HashUnit::format_value(total_hashrate) - ); - let active_upstreams = router.get_active_upstreams().await; - for upstream_id in active_upstreams { - info!( - "Upstream {}: receiving full hashrate {}", - upstream_id, - HashUnit::format_value(total_hashrate) - ); - } - } else { - warn!("No active upstream connections!"); - // In a real implementation, you might want to try reconnecting here + // Check proxy state + let (is_down, error_msg) = ProxyState::is_proxy_down(); + if is_down { + error!("⚠️ Proxy state is DOWN: {:?}", error_msg); + break; } } } diff --git a/src/router/mod.rs b/src/router/mod.rs index e2ea8f51..00846c9f 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -17,7 +17,7 @@ use tokio::{ watch, }, }; -use tracing::{error, info}; +use tracing::{error, info, warn}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, @@ -67,6 +67,7 @@ impl Clone for Router { impl Router { /// Creates a new `Router` instance with the specified upstream addresses. + /// Now supports multiple pools with different auth keys for latency-based selection pub fn new( pool_addresses: Vec, auth_pub_k: Secp256k1PublicKey, @@ -91,20 +92,60 @@ impl Router { } /// Creates a new Router with multiple upstream addresses and auth keys + /// Support both latency-based selection and custom distribution + pub fn new_with_keys( + pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, + setup_connection_msg: Option>, + timer: Option, + ) -> Self { + let pool_socket_addresses: Vec = + pool_address_keys.iter().map(|(addr, _)| *addr).collect(); + let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); + let auth_pub_k = keys[0]; // Use first key as primary + + let (latency_tx, latency_rx) = watch::channel(None); + + Self { + pool_socket_addresses, + keys, + current_pool: None, + upstream_manager: None, + aggregated_receiver: None, + auth_pub_k, + setup_connection_msg, + timer, + latency_tx, + latency_rx, + } + } + + /// Creates a new Router with multiple upstream addresses and auth keys + /// This now supports both latency-based selection and custom distribution pub async fn new_multi( pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, setup_connection_msg: Option>, timer: Option, - use_parallel: bool, + use_distribution: bool, // Changed from use_parallel to use_distribution ) -> Result { let pool_socket_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); - // Create with parallel mode enabled - let (upstream_manager, aggregated_receiver) = MultiUpstreamManager::new(use_parallel); + // Create upstream manager only if we want custom distribution + let (upstream_manager, aggregated_receiver) = if use_distribution { + let manager = MultiUpstreamManager::new( + pool_socket_addresses.clone(), + keys[0], // or use a key per upstream if needed + setup_connection_msg.clone(), + timer.clone(), + ); + let receiver = None; // or whatever is appropriate for your design + (Some(manager), receiver) + } else { + (None, None) + }; - // Use first key as default auth key for backward compatibility + // Use let auth_pub_k = keys .first() .copied() @@ -116,8 +157,8 @@ impl Router { pool_socket_addresses, keys, current_pool: None, - upstream_manager: Some(upstream_manager), - aggregated_receiver: Some(aggregated_receiver), + upstream_manager, + aggregated_receiver, auth_pub_k, setup_connection_msg, timer, @@ -134,67 +175,39 @@ impl Router { None } } - - /// Set hashrate distribution for multi-upstream mode - pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result<(), String> { + + /// Get aggregated sender for multi-upstream mode + pub async fn get_aggregated_sender(&self) -> tokio::sync::mpsc::Sender> { if let Some(ref manager) = self.upstream_manager { - manager.set_hashrate_distribution(distribution).await; - Ok(()) + manager.get_aggregated_sender().await } else { - Err("Multi-upstream manager not available".to_string()) + // Return a dummy sender if no manager exists + let (sender, _) = tokio::sync::mpsc::channel(1); + sender } } - - /// Get detailed connection statistics with custom distribution + + + + /// Get detailed connection statistics pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { - let mut results = Vec::new(); - if let Some(ref manager) = self.upstream_manager { - let upstreams = manager.get_upstreams().await; - let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); - let distribution = manager.get_hashrate_distribution().await; - - // Check if we have custom distribution - if distribution.len() == upstreams.len() && !distribution.is_empty() { - // Use custom distribution - let upstream_ids: Vec = upstreams.keys().cloned().collect(); - for (idx, (id, upstream)) in upstreams.iter().enumerate() { - let percentage = distribution.get(idx).copied().unwrap_or(0.0); - let hashrate = if upstream.is_active { - total_hashrate * (percentage / 100.0) - } else { - 0.0 - }; - results.push((id.clone(), upstream.is_active, hashrate)); - } - } else { - // Fall back to equal distribution - let active_count = upstreams.values().filter(|upstream| upstream.is_active).count(); - let hashrate_per_upstream = if active_count > 0 { - total_hashrate / active_count as f32 - } else { - 0.0 - }; - - for (id, upstream) in upstreams { - let hashrate = if upstream.is_active { hashrate_per_upstream } else { 0.0 }; - results.push((id, upstream.is_active, hashrate)); - } - } + manager.get_detailed_connection_stats().await } else { - // Single upstream mode (fallback) - if let Some(current) = self.current_pool { - let total_hashrate = crate::proxy_state::ProxyState::get_total_hashrate(); - results.push(("upstream-0".to_string(), true, total_hashrate)); - } + vec![] } - - results } - - - /// Check if the router is using the MultiUpstreamManager + /// Get current hashrate distribution + pub async fn get_hashrate_distribution(&self) -> Vec { + if let Some(ref manager) = self.upstream_manager { + manager.get_hashrate_distribution().await + } else { + vec![] + } + } + + /// Check if multi-upstream is enabled pub fn is_multi_upstream_enabled(&self) -> bool { self.upstream_manager.is_some() } @@ -204,6 +217,29 @@ impl Router { self.upstream_manager.is_some() } + /// Check if the router is using custom distribution mode + pub fn is_distribution_mode(&self) -> bool { + self.upstream_manager.is_some() + } + + /// Get the best pools based on latency and desired count + pub async fn select_best_pools(&mut self, desired_count: usize) -> Vec<(SocketAddr, Duration)> { + let mut pool_latencies = Vec::new(); + + // Get latencies for all pools + for &pool_addr in &self.pool_socket_addresses { + if let Ok(latency) = self.get_latency(pool_addr).await { + pool_latencies.push((pool_addr, latency)); + } + } + + // Sort by latency (lowest first) + pool_latencies.sort_by_key(|(_, latency)| *latency); + + // Return the best pools up to desired count + pool_latencies.into_iter().take(desired_count).collect() + } + /// Checks for faster upstream and switches to it if found pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { // For multi-upstream mode, we don't switch since we use all simultaneously @@ -441,11 +477,11 @@ impl Router { } } - // Start the manager after adding all upstreams - // manager.start().await?; - // If MultiUpstreamManager requires initialization, call it here. Otherwise, remove this. + // IMPORTANT: Initialize connections BEFORE any hashrate distribution is set + // This ensures add_upstream() calls are complete and won't overwrite distribution + manager.initialize_connections().await; - info!("Upstream connections initialized"); + info!("Upstream connections initialized - ready for hashrate distribution"); Ok(()) } else { Err("No upstream manager available".to_string()) @@ -463,6 +499,20 @@ impl Router { (self.pool_socket_addresses.len(), 1) // Single upstream mode } } + /// Sets the hashrate distribution for the upstream manager. +pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result<(), &'static str> { + info!("🔧 Router::set_hashrate_distribution called with: {:?}", distribution); + + if let Some(ref manager) = self.upstream_manager { + info!("🔧 Calling manager.set_hashrate_distribution"); + manager.set_hashrate_distribution(distribution).await; + info!("🔧 Manager.set_hashrate_distribution returned"); + Ok(()) + } else { + error!("❌ No upstream manager available"); + Err("No upstream manager available") + } + } /// Get list of active upstream IDs pub async fn get_active_upstreams(&self) -> Vec { @@ -512,9 +562,49 @@ impl Router { best_pool.map(|pool| (pool, best_latency)) } - + /// Select best pools for hashrate distribution based on latency + /// This implements condition 2: choose N best pools from available pools + pub async fn select_best_pools_for_distribution(&mut self, desired_count: usize) -> Vec<(SocketAddr, Secp256k1PublicKey)> { + // If we have exactly the desired count or fewer, use all + if self.pool_socket_addresses.len() <= desired_count { + info!("Using all {} available pools (desired: {})", self.pool_socket_addresses.len(), desired_count); + return self.pool_socket_addresses + .iter() + .zip(self.keys.iter()) + .map(|(&addr, &key)| (addr, key)) + .collect(); + } + + // Get latencies for all pools + let mut pool_latencies = Vec::new(); + info!("Testing latencies for {} pools to select {} best", self.pool_socket_addresses.len(), desired_count); + + for (i, &pool_addr) in self.pool_socket_addresses.iter().enumerate() { + match self.get_latency(pool_addr).await { + Ok(latency) => { + info!("Pool {}: latency {:?}", pool_addr, latency); + pool_latencies.push((pool_addr, self.keys[i], latency)); + } + Err(_) => { + warn!("Failed to get latency for pool {}", pool_addr); + } + } + } - + // Sort by latency (lowest first) and take the best ones + pool_latencies.sort_by_key(|(_, _, latency)| *latency); + let selected: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_latencies + .into_iter() + .take(desired_count) + .map(|(addr, key, latency)| { + info!("Selected pool {} with latency {:?}", addr, latency); + (addr, key) + }) + .collect(); + + info!("Selected {} pools based on latency", selected.len()); + selected + } } /// Track latencies for various stages of pool connection setup. diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 0d9fb1cc..e963d48c 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -1,433 +1,218 @@ -use crate::{proxy_state::ProxyState, shared::utils::AbortOnDrop}; +use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; +use tokio::sync::Mutex; +use tracing::{error, info}; +use crate::{ + minin_pool_connection, + shared::utils::AbortOnDrop, + proxy_state::ProxyState, + HashUnit, +}; use demand_share_accounting_ext::parser::PoolExtMessages; use key_utils::Secp256k1PublicKey; -use roles_logic_sv2::common_messages_sv2::SetupConnection; -use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; -use tokio::sync::{mpsc, Mutex}; -use tracing::{error, info, warn}; - -#[derive(Clone)] -pub struct UpstreamConnection { - pub id: String, - pub address: SocketAddr, - pub auth_key: Secp256k1PublicKey, - pub sender: mpsc::Sender>, - pub is_active: bool, - pub connection_handle: Option, -} #[derive(Clone)] pub struct MultiUpstreamManager { upstreams: Arc>>, - aggregated_sender: tokio::sync::mpsc::Sender>, - aggregated_receiver: Arc>>>>, - // New field for custom hashrate distribution - hashrate_distribution: Arc>>, + auth_pub_k: Secp256k1PublicKey, + setup_connection_msg: Option>, + timer: Option, } -impl MultiUpstreamManager { - pub fn new(use_parallel: bool) -> (Self, tokio::sync::mpsc::Receiver>) { - let (sender, receiver) = tokio::sync::mpsc::channel(1000); - let manager = Self { - upstreams: Arc::new(Mutex::new(HashMap::new())), - aggregated_sender: sender, - aggregated_receiver: Arc::new(Mutex::new(None)), // Start with None - hashrate_distribution: Arc::new(Mutex::new(Vec::new())), // Initialize empty - }; - - (manager, receiver) - } - - /// Get the aggregated receiver (for use in the main loop) - pub async fn get_aggregated_receiver(&self) -> Option>> { - let mut rx = self.aggregated_receiver.lock().await; - rx.take() - } - - /// Set the aggregated receiver (if needed to put it back) - pub async fn set_aggregated_receiver(&self, receiver: tokio::sync::mpsc::Receiver>) { - let mut rx = self.aggregated_receiver.lock().await; - *rx = Some(receiver); - } - - /// Add a new upstream connection - pub async fn add_upstream( - &self, - id: String, - address: SocketAddr, - auth_key: Secp256k1PublicKey, - setup_connection_msg: Option>, - timer: Option, - ) -> Result<(), String> { - info!("🔍 ADD_UPSTREAM CALLED: {} -> {}", id, address); - - let mut upstreams = self.upstreams.lock().await; - - // Check if upstream already exists - if upstreams.contains_key(&id) { - warn!("Upstream {} already exists, skipping", id); - return Ok(()); - } +#[derive(Clone)] +pub struct UpstreamConnection { + pub address: SocketAddr, + pub allocated_hashrate: f64, + pub allocated_percentage: f32, + pub is_active: bool, +} - // Create upstream connection - let upstream_connection = UpstreamConnection { - id: id.clone(), +impl UpstreamConnection { + fn new(address: SocketAddr) -> Self { + Self { address, - auth_key, - is_active: false, // Start as inactive, will be set to true when connected - sender: self.aggregated_sender.clone(), // Use a proper sender here - connection_handle: None, // Set to None initially - }; - - upstreams.insert(id.clone(), upstream_connection); - info!("✅ Added upstream {} to manager", id); - - // Start connection task - let upstreams_clone = self.upstreams.clone(); - let sender_clone = self.aggregated_sender.clone(); - - tokio::spawn(async move { - Self::connect_upstream( - upstreams_clone, - id, - address, - auth_key, - setup_connection_msg, - timer, - sender_clone, - ) - .await - }); + allocated_hashrate: 0.0, + allocated_percentage: 0.0, + is_active: false, + } + } - Ok(()) + fn update_allocation(&mut self, percentage: f32, total_hashrate: f64) { + self.allocated_percentage = percentage; + self.allocated_hashrate = total_hashrate * percentage as f64 / 100.0; } +} - /// Connect to a specific upstream - async fn connect_upstream( - upstreams: Arc>>, - id: String, - address: SocketAddr, - auth_key: Secp256k1PublicKey, - setup_connection_msg: Option>, +impl MultiUpstreamManager { + pub fn new( + upstreams: Vec, + auth_pub_k: Secp256k1PublicKey, + setup_connection_msg: Option>, timer: Option, - sender: tokio::sync::mpsc::Sender>, - ) { - info!("Connecting to upstream {}: {}", id, address); - - match crate::minin_pool_connection::connect_pool( - address, - auth_key, + ) -> Self { + let upstream_map = upstreams + .into_iter() + .enumerate() + .map(|(i, addr)| (format!("upstream-{}", i), UpstreamConnection::new(addr))) + .collect(); + + Self { + upstreams: Arc::new(Mutex::new(upstream_map)), + auth_pub_k, setup_connection_msg, timer, - ) - .await - { - Ok((send_to_pool, recv_from_pool, _abort_handle)) => { - info!("Successfully connected to upstream {}: {}", id, address); - - // Update upstream status in both MultiUpstreamManager and ProxyState - { - let mut upstreams_lock = upstreams.lock().await; - if let Some(upstream) = upstreams_lock.get_mut(&id) { - upstream.is_active = true; - } - } - - // Also update ProxyState - crate::proxy_state::ProxyState::set_upstream_connection_status(&id, true); - - // Handle messages from this upstream - Self::handle_upstream_messages(id.clone(), recv_from_pool, sender).await; - - // If we reach here, connection was lost - warn!("Connection to upstream {} lost", id); - - // Mark as inactive - { - let mut upstreams_lock = upstreams.lock().await; - if let Some(upstream) = upstreams_lock.get_mut(&id) { - upstream.is_active = false; - } - } - crate::proxy_state::ProxyState::set_upstream_connection_status(&id, false); - } - Err(e) => { - error!("Failed to connect to upstream {}: {}", id, e); - - // Mark as inactive - let mut upstreams_lock = upstreams.lock().await; - if let Some(upstream) = upstreams_lock.get_mut(&id) { - upstream.is_active = false; - } - crate::proxy_state::ProxyState::set_upstream_connection_status(&id, false); - } - } - } - - /// Handle messages from an upstream - async fn handle_upstream_messages( - upstream_id: String, - mut receiver: tokio::sync::mpsc::Receiver>, - sender: tokio::sync::mpsc::Sender>, - ) { - info!("Starting message handler for upstream {}", upstream_id); - - while let Some(message) = receiver.recv().await { - if let Err(e) = sender.send(message).await { - error!( - "Failed to forward message from upstream {}: {}", - upstream_id, e - ); - break; - } } - - warn!("Message handler for upstream {} stopped", upstream_id); } - /// Send a message to a specific upstream - pub async fn send_to_upstream( - &self, - upstream_id: &str, - message: PoolExtMessages<'static>, - ) -> Result<(), String> { - let connections = self.upstreams.lock().await; - if let Some(connection) = connections.get(upstream_id) { - if connection.is_active { - connection - .sender - .send(message) + /// Start and maintain all upstream connections according to hashrate distribution + pub async fn maintain_connections(&self) { + let upstreams = self.upstreams.clone(); + let auth_pub_k = self.auth_pub_k; + let setup_connection_msg = self.setup_connection_msg.clone(); + let timer = self.timer; + + let upstreams_guard = upstreams.lock().await; + for (id, conn) in upstreams_guard.iter() { + let id = id.clone(); + let address = conn.address; + let upstreams = upstreams.clone(); + let setup_connection_msg = setup_connection_msg.clone(); + + tokio::spawn(async move { + loop { + info!("Connecting to upstream {}: {}", id, address); + match minin_pool_connection::connect_pool( + address, + auth_pub_k, + setup_connection_msg.clone(), + timer, + ) .await - .map_err(|e| format!("Failed to send to upstream {}: {}", upstream_id, e))?; - Ok(()) - } else { - Err(format!("Upstream {} is not active", upstream_id)) - } - } else { - Err(format!("Upstream {} not found", upstream_id)) - } - } - - /// Send a message to the next upstream using the configured strategy - pub async fn send_to_next_upstream( - &self, - message: PoolExtMessages<'static>, - ) -> Result<(), String> { - info!("Broadcasting to all upstreams (parallel mode)"); - let results = self.broadcast(message).await; - - // Check if any sends were successful - let success_count = results - .iter() - .filter(|result| result.contains("Success")) - .count(); - - if success_count > 0 { - info!("Successfully sent to {} upstreams", success_count); - Ok(()) - } else { - Err(format!("Failed to send to any upstreams: {:?}", results)) - } - } - - /// Connect to upstream with retry logic - async fn connect_upstream_with_retry( - upstreams: Arc>>, - id: String, - address: SocketAddr, - auth_key: Secp256k1PublicKey, - setup_connection_msg: Option>, - timer: Option, - sender: tokio::sync::mpsc::Sender>, - ) { - let mut retry_count = 0; - const MAX_RETRIES: u32 = 3; - const RETRY_DELAY: Duration = Duration::from_secs(5); - - while retry_count < MAX_RETRIES { - info!( - "Connecting to upstream {} (attempt {}/{}): {}", - id, - retry_count + 1, - MAX_RETRIES, - address - ); - - match crate::minin_pool_connection::connect_pool( - address, - auth_key, - setup_connection_msg.clone(), - timer, - ) - .await - { - Ok((send_to_pool, recv_from_pool, _abort_handle)) => { - info!("Successfully connected to upstream {}: {}", id, address); - - // Update upstream status { - let mut upstreams_lock = upstreams.lock().await; - if let Some(upstream) = upstreams_lock.get_mut(&id) { - upstream.is_active = true; + Ok((_send, _recv, _abortable)) => { + info!("✅ Connected to upstream {} - maintaining connection", id); + Self::update_upstream_status(&upstreams, &id, true).await; + // Keep the connection alive (simulate mining) + tokio::time::sleep(Duration::from_secs(30)).await; + } + Err(e) => { + error!("Failed to connect to upstream {}: {:?}", id, e); + Self::update_upstream_status(&upstreams, &id, false).await; + tokio::time::sleep(Duration::from_secs(5)).await; } - } - - // Handle messages from this upstream - Self::handle_upstream_messages(id.clone(), recv_from_pool, sender).await; - return; // Success, exit retry loop - } - Err(e) => { - error!( - "Failed to connect to upstream {} (attempt {}): {}", - id, - retry_count + 1, - e - ); - retry_count += 1; - - if retry_count < MAX_RETRIES { - info!("Retrying connection to {} in {:?}", id, RETRY_DELAY); - tokio::time::sleep(RETRY_DELAY).await; } } - } - } - - error!( - "Failed to connect to upstream {} after {} attempts", - id, MAX_RETRIES - ); - - // Mark as inactive after all retries failed - let mut upstreams_lock = upstreams.lock().await; - if let Some(upstream) = upstreams_lock.get_mut(&id) { - upstream.is_active = false; + }); } } - /// Broadcast to all active upstreams (parallel execution) - pub async fn broadcast(&self, message: PoolExtMessages<'static>) -> Vec { - let upstreams = self.upstreams.lock().await; - let mut results = Vec::new(); - - if upstreams.is_empty() { - results.push("No upstreams configured".to_string()); - return results; - } - let active_upstreams: Vec<_> = upstreams - .values() - .filter(|upstream| upstream.is_active) - .collect(); - - if active_upstreams.is_empty() { - results.push("No active upstreams available".to_string()); - return results; + /// Helper method to update upstream connection status + async fn update_upstream_status( + upstreams: &Arc>>, + id: &str, + is_active: bool, + ) { + let mut upstreams = upstreams.lock().await; + if let Some(upstream) = upstreams.get_mut(id) { + upstream.is_active = is_active; } + } - for upstream in active_upstreams { - match upstream.sender.send(message.clone()).await { - Ok(_) => { - results.push(format!("Sent to {}: Success", upstream.id)); - } - Err(e) => { - results.push(format!("Sent to {}: Failed - {}", upstream.id, e)); - // Mark upstream as inactive if send fails - warn!( - "Upstream {} send failed, may need reconnection: {}", - upstream.id, e - ); - } + /// Set hashrate distribution for each upstream (percentages) + pub async fn set_hashrate_distribution(&self, distribution: Vec) { + println!("calling set_hashrate_distribution in multi manager with: {:?}", distribution); + + let mut upstreams = self.upstreams.lock().await; + let total_hashrate = ProxyState::get_total_hashrate() as f64; + + println!("Total hashrate: {}", HashUnit::format_value(total_hashrate as f32)); + + for (i, (_id, conn)) in upstreams.iter_mut().enumerate() { + if let Some(&percentage) = distribution.get(i) { + println!("setting upstream {} to {}%", i, percentage); + conn.update_allocation(percentage, total_hashrate); + info!( + "Upstream {}: allocated {} ({:.1}%)", + conn.address, + HashUnit::format_value(conn.allocated_hashrate as f32), + percentage + ); } } + } - results + pub async fn get_upstreams(&self) -> HashMap { + self.upstreams.lock().await.clone() } - /// Get all upstreams (for stats) - pub async fn get_upstreams(&self) -> std::collections::HashMap { + pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { let upstreams = self.upstreams.lock().await; - upstreams.clone() + + upstreams + .iter() + .map(|(id, conn)| { + println!( + "Inside get_detailed_connection_stats Upstream {}: is_active={}, allocated_percentage={:.2}%", + id, + conn.is_active, + conn.allocated_percentage + ); + (id.clone(), conn.is_active, conn.allocated_percentage) + }) + .collect() } - /// Get list of active upstream IDs - pub async fn get_active_upstream_ids(&self) -> Vec { + pub async fn get_hashrate_distribution(&self) -> Vec { let upstreams = self.upstreams.lock().await; upstreams .values() - .filter(|upstream| upstream.is_active) - .map(|upstream| upstream.id.clone()) + .map(|conn| conn.allocated_percentage) .collect() } - /// Mark an upstream as inactive - pub async fn mark_upstream_inactive(&self, upstream_id: &str) { - let mut connections = self.upstreams.lock().await; - if let Some(connection) = connections.get_mut(upstream_id) { - connection.is_active = false; - ProxyState::set_upstream_connection_status(upstream_id, false); - info!("Marked upstream {} as inactive", upstream_id); - } - } - - /// Remove an upstream connection - pub async fn remove_upstream(&self, upstream_id: &str) { - let mut connections = self.upstreams.lock().await; - if let Some(connection) = connections.remove(upstream_id) { - if let Some(handle) = connection.connection_handle { - drop(handle); // This will abort the connection - } - ProxyState::set_upstream_connection_status(upstream_id, false); - info!("Removed upstream {}", upstream_id); - } + pub async fn get_active_upstream_ids(&self) -> Vec { + let upstreams = self.upstreams.lock().await; + upstreams + .iter() + .filter(|(_, conn)| conn.is_active) + .map(|(id, _)| id.clone()) + .collect() } - /// Get connection count - pub async fn connection_count(&self) -> usize { - let connections = self.upstreams.lock().await; - connections.len() + // Dummy implementations for sender/receiver if needed + pub async fn get_aggregated_receiver(&self) -> Option>> { + None } - /// Get active connection count - pub async fn active_connection_count(&self) -> usize { - let connections = self.upstreams.lock().await; - connections.values().filter(|conn| conn.is_active).count() + pub async fn get_aggregated_sender(&self) -> tokio::sync::mpsc::Sender> { + let (sender, _) = tokio::sync::mpsc::channel(1); + sender } - /// Set custom hashrate distribution percentages - pub async fn set_hashrate_distribution(&self, distribution: Vec) { - let mut dist = self.hashrate_distribution.lock().await; - *dist = distribution; - info!("Set custom hashrate distribution: {:?}", *dist); - } + // Add upstream (called from Router) + pub async fn add_upstream( + &self, + id: String, + address: SocketAddr, + _key: Secp256k1PublicKey, + _setup_connection_msg: Option>, + _timer: Option, +) -> Result<(), String> { + let mut upstreams = self.upstreams.lock().await; - /// Get current hashrate distribution - pub async fn get_hashrate_distribution(&self) -> Vec { - self.hashrate_distribution.lock().await.clone() + // Check if upstream already exists - if so, don't overwrite it + if upstreams.contains_key(&id) { + println!("⚠️ Upstream {} already exists, not overwriting", id); + return Ok(()); } - /// Validate that distribution percentages add up to 100% - pub fn validate_distribution(distribution: &[f32]) -> Result<(), String> { - if distribution.is_empty() { - return Err("Distribution cannot be empty".to_string()); - } - - let total: f32 = distribution.iter().sum(); - if (total - 100.0).abs() > 0.1 { // Allow small floating point errors - return Err(format!( - "Distribution percentages must add up to 100%, got {:.1}%", - total - )); - } - - for (idx, &percentage) in distribution.iter().enumerate() { - if percentage < 0.0 || percentage > 100.0 { - return Err(format!( - "Invalid percentage at index {}: {:.1}%. Must be between 0% and 100%", - idx, percentage - )); - } - } - - Ok(()) - } + // Only insert if it doesn't exist + upstreams.insert(id.clone(), UpstreamConnection::new(address)); + println!("✅ Added new upstream {}", id); + Ok(()) } + + // Maintain connections (called from Router) + pub async fn initialize_connections(&self) { + self.maintain_connections().await; + } +} \ No newline at end of file diff --git a/src/share_accounter/mod.rs b/src/share_accounter/mod.rs index 08004da6..ecda99ae 100644 --- a/src/share_accounter/mod.rs +++ b/src/share_accounter/mod.rs @@ -124,4 +124,4 @@ fn relay_down( } }); task.into() -} +} \ No newline at end of file diff --git a/src/translator/mod.rs b/src/translator/mod.rs index a6090d9a..04603ef7 100644 --- a/src/translator/mod.rs +++ b/src/translator/mod.rs @@ -1,8 +1,8 @@ mod downstream; mod error; -pub mod proxy; // Make public -pub mod upstream; // Make public +mod proxy; +mod upstream; mod utils; use bitcoin::Address; @@ -207,4 +207,4 @@ pub async fn start( .map_err(|_| Error::TranslatorTaskManagerFailed)?; Ok(abortable) -} +} \ No newline at end of file From 893bc22a14896a73bdec709fe01161c087c82ad6 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Fri, 6 Jun 2025 02:05:10 +0530 Subject: [PATCH 10/21] feat(router): Add hashrate distribution check and improve detailed connection stats logging --- src/main.rs | 42 +++++++++++++++------------- src/router/multi_upstream_manager.rs | 33 +++++++++++----------- 2 files changed, 38 insertions(+), 37 deletions(-) diff --git a/src/main.rs b/src/main.rs index a13210bd..cecbd8b6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -122,6 +122,9 @@ impl Config { let config: Config = toml::from_str(&contents)?; Ok(config) } + pub fn wants_hashrate_distribution(&self) -> bool { + self.hashrate_distribution.is_some() + } } #[tokio::main] @@ -230,8 +233,8 @@ async fn main() { // Determine if we want hashrate distribution let wants_distribution = config.as_ref() - .and_then(|c| c.hashrate_distribution.as_ref()) - .is_some() || ARGS.parallel; + .map(|c| c.wants_hashrate_distribution()) + .unwrap_or(false); // Create router based on configuration let mut router = if wants_distribution { @@ -423,24 +426,23 @@ async fn main() { info!("🌐 Mode: Equal Distribution ({:.1}% per upstream)", percentage_per_upstream); } - for (upstream_id, is_active, hashrate) in detailed_stats { - if is_active { - info!("{} = hashrate", hashrate); - let percentage = (hashrate / total_hashrate) * 100.0; - info!( - " ✅ {}: receiving {} ({:.1}% of total)", - upstream_id, - HashUnit::format_value(hashrate), - percentage - ); - } else { - info!( - " ❌ {}: {} (INACTIVE)", - upstream_id, - HashUnit::format_value(hashrate) - ); - } - } + for (upstream_id, is_active, percentage) in detailed_stats { + if is_active { + // Calculate actual hashrate from percentage + let actual_hashrate = (percentage / 100.0) * total_hashrate; + info!( + " ✅ {}: allocated {} ({:.1}% of total)", + upstream_id, + HashUnit::format_value(actual_hashrate), + percentage + ); + } else { + info!( + " ❌ {}: 0.00T (INACTIVE - 0.0% of total)", + upstream_id + ); + } +} info!("========================================"); // Start monitoring task for multi-upstream diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index e963d48c..51eaaf00 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -143,23 +143,22 @@ impl MultiUpstreamManager { pub async fn get_upstreams(&self) -> HashMap { self.upstreams.lock().await.clone() } - - pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { - let upstreams = self.upstreams.lock().await; - - upstreams - .iter() - .map(|(id, conn)| { - println!( - "Inside get_detailed_connection_stats Upstream {}: is_active={}, allocated_percentage={:.2}%", - id, - conn.is_active, - conn.allocated_percentage - ); - (id.clone(), conn.is_active, conn.allocated_percentage) - }) - .collect() - } +pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { + let upstreams = self.upstreams.lock().await; + + upstreams + .iter() + .map(|(id, conn)| { + println!( + "Upstream {}: {}% allocation", + id, + conn.allocated_percentage + ); + // Return the percentage, not the hashrate + (id.clone(), conn.is_active, conn.allocated_percentage) + }) + .collect() +} pub async fn get_hashrate_distribution(&self) -> Vec { let upstreams = self.upstreams.lock().await; From 9d4f3cb3af98c114c30940094148c1514585acd5 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Fri, 6 Jun 2025 02:07:14 +0530 Subject: [PATCH 11/21] feat(router): Refactor hashrate distribution logging for clarity and accuracy --- src/main.rs | 49 ++++++++++++++++++++++++------------------------- 1 file changed, 24 insertions(+), 25 deletions(-) diff --git a/src/main.rs b/src/main.rs index cecbd8b6..e7f9dc06 100644 --- a/src/main.rs +++ b/src/main.rs @@ -409,15 +409,14 @@ async fn main() { if active_count > 0 { // Check if using custom distribution let mut is_custom_distribution = false; - for (_, is_active, hashrate) in &detailed_stats { - if *is_active { - let percentage = (*hashrate / total_hashrate) * 100.0; - if (percentage - (100.0 / active_count as f32)).abs() > 1.0 { - is_custom_distribution = true; - break; - } - } + for (_, is_active, percentage) in &detailed_stats { + if *is_active { + if (percentage - (100.0 / active_count as f32)).abs() > 1.0 { + is_custom_distribution = true; + break; } + } +} if is_custom_distribution { info!("🎯 Mode: Custom Distribution (from config)"); @@ -703,23 +702,23 @@ async fn monitor_multi_upstream(router: Router, epsilon: Duration) { info!("⚖️ Distribution mode: Equal (automatic)"); } - for (upstream_id, is_active, hashrate) in detailed_stats { - if is_active { - let percentage = (hashrate / total_hashrate) * 100.0; - info!( - " ✅ {}: receiving {} ({:.1}% of total)", - upstream_id, - HashUnit::format_value(hashrate), - percentage - ); - } else { - info!( - " ❌ {}: {} (INACTIVE)", - upstream_id, - HashUnit::format_value(hashrate) - ); - } - } + for (upstream_id, is_active, percentage) in detailed_stats { + if is_active { + // Calculate actual hashrate from percentage + let allocated_hashrate = (percentage / 100.0) * total_hashrate; + info!( + " ✅ {}: allocated {} ({:.1}% of total)", + upstream_id, + HashUnit::format_value(allocated_hashrate), + percentage + ); + } else { + info!( + " ❌ {}: 0.00T (INACTIVE - 0.0% of total)", + upstream_id + ); + } +} } else { warn!("❌ No active upstream connections!"); } From b0eae5be64224eee2b5489be9fa405ee90dbb8d2 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Fri, 6 Jun 2025 02:10:24 +0530 Subject: [PATCH 12/21] fix(main): Remove debug print statement for config in main function --- src/main.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/main.rs b/src/main.rs index e7f9dc06..a38ab81e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -177,7 +177,6 @@ async fn main() { } else { None }; - println!("Config: {:?}", config); // Set the total hashrate from config BEFORE creating router if let Some(ref config) = config { From b24c082e463b6664b6a21010fa172a70da6a528f Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Fri, 6 Jun 2025 18:14:24 +0530 Subject: [PATCH 13/21] feat: Implement hashrate distribution with multi-upstream support - Added optional hashrate_distribution configuration in config.toml - Implemented three operational modes: 1. Latency-based selection (default when no distribution specified) 2. Hashrate distribution (when percentages provided in config) 3. Single upstream (when only one pool configured) - Enhanced Configuration struct with wants_hashrate_distribution() method - Added multi-upstream router with percentage-based hashrate allocation - Improved router initialization and upstream connection management - Added detailed monitoring and reporting for multi-upstream mode - Fixed pool address parsing and configuration loading from TOML - Enhanced logging with connection stats and allocation percentages - Added comprehensive error handling and fallback mechanisms Breaking changes: None (backward compatible) New features: hashrate_distribution config option, multi-upstream routing Fixes: Pool address configuration parsing, test mode pool selection --- config.toml | 6 +- src/config.rs | 103 +++-- src/jd_client/mod.rs | 21 +- src/main.rs | 432 ++++++------------- src/proxy_state.rs | 8 +- src/router/multi_upstream_manager.rs | 2 +- src/translator/downstream/diff_management.rs | 8 +- src/translator/downstream/notify.rs | 10 +- 8 files changed, 222 insertions(+), 368 deletions(-) diff --git a/config.toml b/config.toml index dae76260..43a25367 100644 --- a/config.toml +++ b/config.toml @@ -1,6 +1,6 @@ token="UR01TMkZv6Vs5zbwr0t6" -pool_addresses=["18.193.252.132:2000","3.74.36.119:2000"] -tp_address = "127.0.0.1:8442" +test_pool_addresses=["3.74.36.119:2000","18.193.252.132:2000"] +#tp_address = "127.0.0.1:8442" interval = 120_000 delay = 540 downstream_hashrate = "500T" @@ -9,4 +9,4 @@ nc_loglevel = "off" test = true # Uncomment this line to enable hashrate distribution -#hashrate_distribution = [30.0, 70.0] \ No newline at end of file +hashrate_distribution = [70.0, 30.0] \ No newline at end of file diff --git a/src/config.rs b/src/config.rs index 57e957cc..693b3946 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,8 @@ use clap::Parser; use lazy_static::lazy_static; use serde::{Deserialize, Serialize}; +use tracing_subscriber::field::delimited::VisitDelimited; +use core::hash; use std::{ net::{SocketAddr, ToSocketAddrs}, path::PathBuf, @@ -8,55 +10,61 @@ use std::{ use tracing::{error, info, warn}; use crate::{HashUnit, DEFAULT_SV1_HASHPOWER}; + lazy_static! { pub static ref CONFIG: Configuration = Configuration::load_config(); } + #[derive(Parser)] -struct Args { +pub struct Args { #[clap(long)] - test: bool, + pub test: bool, #[clap(long = "d", short = 'd', value_parser = parse_hashrate)] - downstream_hashrate: Option, + pub downstream_hashrate: Option, #[clap(long = "loglevel", short = 'l')] - loglevel: Option, + pub loglevel: Option, #[clap(long = "nc", short = 'n')] - noise_connection_log: Option, + pub noise_connection_log: Option, #[clap(long = "delay")] - delay: Option, + pub delay: Option, #[clap(long = "interval", short = 'i')] - adjustment_interval: Option, + pub adjustment_interval: Option, #[clap(long = "pool", short = 'p', value_delimiter = ',')] - pool_addresses: Option>, + pub pool_addresses: Option>, #[clap(long = "test-pool", value_delimiter = ',')] - test_pool_addresses: Option>, + pub test_pool_addresses: Option>, #[clap(long)] - token: Option, + pub token: Option, #[clap(long)] - tp_address: Option, + pub tp_address: Option, #[clap(long)] - listening_addr: Option, + pub listening_addr: Option, #[clap(long = "config", short = 'c')] - config_file: Option, + pub config_file: Option, #[clap(long = "api-server-port", short = 's')] - api_server_port: Option, + pub api_server_port: Option, + #[clap(long = "hashrate-distribution", value_delimiter = ',')] + pub hashrate_distribution: Option>, } #[derive(Serialize, Deserialize)] -struct ConfigFile { - token: Option, - tp_address: Option, - pool_addresses: Option>, - test_pool_addresses: Option>, - interval: Option, - delay: Option, - downstream_hashrate: Option, - loglevel: Option, - nc_loglevel: Option, - test: Option, - listening_addr: Option, - api_server_port: Option, +pub struct ConfigFile { + pub token: Option, + pub tp_address: Option, + pub pool_addresses: Option>, + pub test_pool_addresses: Option>, + pub interval: Option, + pub delay: Option, + pub downstream_hashrate: Option, + pub loglevel: Option, + pub nc_loglevel: Option, + pub test: Option, + pub listening_addr: Option, + pub api_server_port: Option, + pub hashrate_distribution: Option>, } + pub struct Configuration { token: Option, tp_address: Option, @@ -70,8 +78,11 @@ pub struct Configuration { test: bool, listening_addr: Option, api_server_port: String, + hashrate_distribution: Option>, } + impl Configuration { + pub fn token() -> Option { CONFIG.token.clone() } @@ -82,7 +93,7 @@ impl Configuration { pub fn pool_address() -> Option> { if CONFIG.test { - CONFIG.test_pool_addresses.clone() // Return test pool addresses in test mode + CONFIG.test_pool_addresses.clone() } else { CONFIG.pool_addresses.clone() } @@ -103,6 +114,7 @@ impl Configuration { pub fn downstream_listening_addr() -> Option { CONFIG.listening_addr.clone() } + pub fn api_server_port() -> String { CONFIG.api_server_port.clone() } @@ -137,10 +149,19 @@ impl Configuration { CONFIG.test } + pub fn hashrate_distribution() -> Option> { + CONFIG.hashrate_distribution.clone() + } + + + pub fn wants_hashrate_distribution() -> bool { + CONFIG.hashrate_distribution.is_some() + } + // Loads config from CLI, file, or env vars with precedence: CLI > file > env. fn load_config() -> Self { let args = Args::parse(); - let config_path: PathBuf = args.config_file.unwrap_or("config.toml".into()); + let config_path: PathBuf = args.config_file.clone().unwrap_or("config.toml".into()); let config: ConfigFile = std::fs::read_to_string(&config_path) .ok() .and_then(|content| toml::from_str(&content).ok()) @@ -157,6 +178,7 @@ impl Configuration { test: None, listening_addr: None, api_server_port: None, + hashrate_distribution: None, }); let token = args @@ -285,11 +307,23 @@ impl Configuration { .unwrap_or("off".to_string()); let test = args.test || config.test.unwrap_or(false) || std::env::var("TEST").is_ok(); - +let hashrate_distribution = args + .hashrate_distribution + .or(config.hashrate_distribution) + .or_else(|| { + std::env::var("HASHRATE_DISTRIBUTION") + .ok() + .and_then(|s| { + s.split(',') + .map(|s| s.trim().parse::().ok()) + .collect::>>() + }) + }); + Configuration { token, tp_address, - pool_addresses, + pool_addresses, test_pool_addresses, interval, delay, @@ -299,12 +333,13 @@ impl Configuration { test, listening_addr, api_server_port, - } + hashrate_distribution, + } } } /// Parses a hashrate string (e.g., "10T", "2.5P", "500E") into an f32 value in h/s. -fn parse_hashrate(hashrate_str: &str) -> Result { +pub fn parse_hashrate(hashrate_str: &str) -> Result { let hashrate_str = hashrate_str.trim(); if hashrate_str.is_empty() { return Err("Hashrate cannot be empty. Expected format: '' (e.g., '10T', '2.5P', '5E'".to_string()); @@ -342,4 +377,4 @@ fn parse_address(addr: String) -> SocketAddr { .expect("Failed to parse socket address") .next() .expect("No socket address resolved") -} \ No newline at end of file +} diff --git a/src/jd_client/mod.rs b/src/jd_client/mod.rs index 42a4c842..97c9767c 100644 --- a/src/jd_client/mod.rs +++ b/src/jd_client/mod.rs @@ -48,8 +48,6 @@ use std::{ sync::Arc, }; -use std::net::ToSocketAddrs; - use crate::shared::utils::AbortOnDrop; pub async fn start( @@ -111,10 +109,19 @@ async fn initialize_jd( let port_tp = parts.next().expect("The passed value for TP address is not valid. Terminating.... TP_ADDRESS should be in this format `127.0.0.1:8442`").parse::().expect("This operation should not fail because a valid port_tp should always be converted to U16"); let auth_pub_k: Secp256k1PublicKey = crate::AUTH_PUB_KEY.parse().expect("Invalid public key"); - let address = crate::POOL_ADDRESS - .to_socket_addrs() - .expect("The passed Pool Address is not valid") - .next()?; + let address = match crate::POOL_ADDRESS.safe_lock(|address| *address) { + Ok(Some(address)) => address, + Ok(None) => { + error!("Pool address is missing"); + ProxyState::update_inconsistency(Some(1)); + return None; + } + Err(e) => { + error!("Pool address mutex is poisoned: {e:?}"); + ProxyState::update_inconsistency(Some(1)); + return None; + } + }; let (jd, jd_abortable) = match JobDeclarator::new(address, auth_pub_k.into_bytes(), upstream.clone(), true).await { @@ -264,4 +271,4 @@ async fn retry_connection(address: String) { break; } } -} +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index a38ab81e..0f54237f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,22 +2,21 @@ use jemallocator::Jemalloc; use router::Router; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -use clap::{ArgAction, Parser}; -use serde::{Deserialize, Serialize}; -use std::path::Path; use crate::shared::utils::AbortOnDrop; +use crate::config::{Configuration, parse_hashrate}; // Import from config module + +#[cfg(not(target_os = "windows"))] #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; - use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; use std::{net::SocketAddr, time::Duration}; use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; -mod api; +mod api; mod config; mod ingress; pub mod jd_client; @@ -46,112 +45,25 @@ const TEST_AUTH_PUB_KEY: &str = "9auqWEzQDVyd2oe1JVGFLMLHZtCo2FFqZwtKA5gd9xbuEu7 const DEFAULT_LISTEN_ADDRESS: &str = "0.0.0.0:32767"; lazy_static! { - static ref ARGS: Args = Args::parse(); - // Remove or comment out the old TOKEN line: - // static ref TOKEN: String = std::env::var("TOKEN").expect("Missing TOKEN environment variable"); - - // Other existing lazy_static variables... static ref SV1_DOWN_LISTEN_ADDR: String = - std::env::var("SV1_DOWN_LISTEN_ADDR").unwrap_or(DEFAULT_LISTEN_ADDRESS.to_string()); + Configuration::downstream_listening_addr().unwrap_or(DEFAULT_LISTEN_ADDRESS.to_string()); static ref TP_ADDRESS: roles_logic_sv2::utils::Mutex> = - roles_logic_sv2::utils::Mutex::new(None); // We'll set this from config - static ref EXPECTED_SV1_HASHPOWER: f32 = { - if let Some(value) = ARGS.downstream_hashrate { - value - } else { - let env_var = std::env::var("EXPECTED_SV1_HASHPOWER").ok(); - env_var.and_then(|s| parse_hashrate(&s).ok()) - .unwrap_or(DEFAULT_SV1_HASHPOWER) - } - }; -} - -lazy_static! { - pub static ref POOL_ADDRESS: &'static str = if ARGS.test { - TEST_POOL_ADDRESS - } else { - MAIN_POOL_ADDRESS - }; - pub static ref AUTH_PUB_KEY: &'static str = if ARGS.test { + roles_logic_sv2::utils::Mutex::new(Configuration::tp_address()); + static ref POOL_ADDRESS: roles_logic_sv2::utils::Mutex> = + roles_logic_sv2::utils::Mutex::new(None); // Connected pool address + static ref EXPECTED_SV1_HASHPOWER: f32 = Configuration::downstream_hashrate(); + static ref API_SERVER_PORT: String = Configuration::api_server_port(); + static ref AUTH_PUB_KEY: &'static str = if Configuration::test() { TEST_AUTH_PUB_KEY } else { MAIN_AUTH_PUB_KEY }; } -#[derive(Parser, Debug)] -#[clap(author, version, about, long_about = None)] -pub struct Args { - // Use test - #[clap(long)] - test: bool, - #[clap(long ="d", short ='d', value_parser = parse_hashrate)] - downstream_hashrate: Option, - #[clap(long = "loglevel", short = 'l', default_value = "info")] - loglevel: String, - #[clap(long = "nc", short = 'n', default_value = "off")] - noise_connection_log: String, - #[clap(long = "delay", default_value = "0")] - delay: u64, - #[clap(long = "interval", short = 'i', default_value = "120000")] - adjustment_interval: u64, - // Add the missing upstream field - #[clap(long = "upstream", short = 'u', action = ArgAction::Append)] - upstream: Option>, - #[clap(long = "config", short = 'c')] - pub config_file: Option, -} - -#[derive(Debug, Deserialize, Serialize, Clone)] -pub struct Config { - pub token: String, - pub pool_addresses: Vec, - pub tp_address: Option, - pub interval: Option, - pub delay: Option, - pub downstream_hashrate: Option, - pub loglevel: Option, - pub nc_loglevel: Option, - pub test: Option, - pub hashrate_distribution: Option>, // New field for custom distribution -} - -impl Config { - pub fn from_file>(path: P) -> Result> { - let contents = std::fs::read_to_string(path)?; - let config: Config = toml::from_str(&contents)?; - Ok(config) - } - pub fn wants_hashrate_distribution(&self) -> bool { - self.hashrate_distribution.is_some() - } -} - #[tokio::main] async fn main() { - let args = Args::parse(); - - let log_level = match args.loglevel.to_lowercase().as_str() { - "trace" | "debug" | "info" | "warn" | "error" => args.loglevel, - _ => { - error!( - "Invalid log level '{}'. Defaulting to 'info'.", - args.loglevel - ); - "info".to_string() - } - }; - - let noise_connection_log_level = match args.noise_connection_log.as_str() { - "trace" | "debug" | "info" | "warn" | "error" => args.noise_connection_log, - _ => { - error!( - "Invalid log level for noise_connection '{}' Defaulting to 'off'.", - args.noise_connection_log - ); - "off".to_string() - } - }; + let log_level = Configuration::loglevel(); + let noise_connection_log_level = Configuration::nc_loglevel(); //Disable noise_connection error (for now) because: // 1. It produce logs that are not very user friendly and also bloat the logs @@ -163,87 +75,53 @@ async fn main() { log_level, noise_connection_log_level ))) .init(); - // std::env::var("TOKEN").expect("Missing TOKEN environment variable"); - // Load configuration and get pool addresses - let config = if let Some(config_path) = &ARGS.config_file { - match Config::from_file(config_path) { - Ok(config) => Some(config), - Err(e) => { - error!("Failed to load config file: {}", e); - std::process::exit(1); - } - } - } else { - None - }; - - // Set the total hashrate from config BEFORE creating router - if let Some(ref config) = config { - if let Some(ref hashrate_str) = config.downstream_hashrate { - match parse_hashrate(hashrate_str) { - Ok(hashrate) => { - ProxyState::set_total_hashrate(hashrate); - info!("Set total hashrate to: {}", HashUnit::format_value(hashrate)); - } - Err(e) => { - error!("Failed to parse downstream_hashrate: {}", e); - } - } - } - } else if let Some(hashrate) = ARGS.downstream_hashrate { - ProxyState::set_total_hashrate(hashrate); - info!("Set total hashrate from args to: {}", HashUnit::format_value(hashrate)); + Configuration::token().expect("TOKEN is not set"); + + if Configuration::test() { + info!("Connecting to test endpoint..."); } - // Get pool addresses with auth keys - let pool_address_keys = if let Some(ref config) = config { - config.pool_addresses.iter().map(|addr_str| { - let addr = addr_str.parse::() - .unwrap_or_else(|_| { - error!("Invalid pool address: {}", addr_str); - std::process::exit(1); - }); - - let auth_key: Secp256k1PublicKey = if config.test.unwrap_or(false) { - TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key") + let auth_pub_k: Secp256k1PublicKey = AUTH_PUB_KEY.parse().expect("Invalid public key"); + + // Use Configuration methods consistently + let pool_addresses = Configuration::pool_address() + .filter(|p| !p.is_empty()) + .unwrap_or_else(|| { + if Configuration::test() { + panic!("Test pool address is missing"); } else { - MAIN_AUTH_PUB_KEY.parse().expect("Invalid main public key") - }; - - (addr, auth_key) - }).collect() - } else { - // Fallback to single upstream - let addr = if ARGS.test { - TEST_POOL_ADDRESS.parse::().expect("Invalid test address") - } else { - MAIN_POOL_ADDRESS.parse::().expect("Invalid main address") - }; - - let auth_key = if ARGS.test { - TEST_AUTH_PUB_KEY.parse().expect("Invalid test public key") - } else { - MAIN_AUTH_PUB_KEY.parse().expect("Invalid main public key") - }; - - vec![(addr, auth_key)] - }; + panic!("Pool address is missing"); + } + }); + + // Set downstream hashrate using Configuration pattern + let downstream_hashrate = Configuration::downstream_hashrate(); + ProxyState::set_downstream_hashrate(downstream_hashrate); + info!("Set downstream hashrate to: {}", HashUnit::format_value(downstream_hashrate)); + + // Get pool addresses with auth keys using Configuration pattern + let pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_addresses.iter().map(|&addr| { + (addr, auth_pub_k) + }).collect(); + + // Determine hashrate distribution using Configuration pattern + let wants_distribution = Configuration::wants_hashrate_distribution(); - // Determine if we want hashrate distribution - let wants_distribution = config.as_ref() - .map(|c| c.wants_hashrate_distribution()) - .unwrap_or(false); +if wants_distribution { + info!("Hashrate distribution is enabled"); +} else { + info!("Hashrate distribution is disabled"); +} // Create router based on configuration let mut router = if wants_distribution { - // User wants hashrate distribution - use multi-upstream mode - let pool_address_keys_clone = pool_address_keys.clone(); + // Multi-upstream with distribution match Router::new_multi( - pool_address_keys_clone, - None, // setup_connection_msg - None, // timer - true, // use_distribution = true + pool_address_keys.clone(), + None, + None, + true, ).await { Ok(router) => { info!("Created multi-upstream router with {} upstreams for hashrate distribution", @@ -256,115 +134,83 @@ async fn main() { } } } else if pool_address_keys.len() > 1 { - // Multiple pools but no distribution specified - use latency-based selection (single upstream mode) - info!("Created latency-based router with {} pools (will select best latency)", pool_address_keys.len()); - Router::new_with_keys( - pool_address_keys.clone(), - None, // setup_connection_msg - None, // timer - ) + info!("Created latency-based router with {} pools", pool_address_keys.len()); + Router::new_with_keys(pool_address_keys.clone(), None, None) } else { - // Single pool - use single upstream mode let (addr, auth_key) = pool_address_keys[0]; info!("Created single upstream router"); - Router::new( - vec![addr], - auth_key, - None, // setup_connection_msg - None, // timer - ) + Router::new(vec![addr], auth_key, None, None) }; - // Add this right after router creation, before the if statements: let epsilon = Duration::from_millis(10); // Handle the three different scenarios + if wants_distribution { info!("=== HASHRATE DISTRIBUTION MODE ==="); - // Get the desired distribution - let distribution = config.as_ref() - .and_then(|c| c.hashrate_distribution.clone()) + // Get the distribution from config + let distribution = Configuration::hashrate_distribution() .unwrap_or_else(|| { - // Default equal distribution if parallel flag is set but no custom distribution let count = pool_address_keys.len(); let equal_percentage = 100.0 / count as f32; vec![equal_percentage; count] }); - println!("Using hashrate distribution: {:?}", distribution); - let desired_pool_count = distribution.len(); - info!("Desired pool count for distribution: {}", desired_pool_count); - - // Select best pools based on latency (implements condition 2 and 3) - let selected_pools = if pool_address_keys.len() <= desired_pool_count { - // Condition 3: Use all available pools (skip latency testing) - info!("Using all {} available pools (condition 3)", pool_address_keys.len()); - pool_address_keys.clone() - } else { - // Condition 2: Select best pools based on latency - info!("Selecting {} best pools from {} available (condition 2)", desired_pool_count, pool_address_keys.len()); - - // Create temporary router for latency testing - let mut temp_router = Router::new_with_keys( - pool_address_keys.clone(), - None, - None, - ); - - temp_router.select_best_pools_for_distribution(desired_pool_count).await - }; - - // Now create the final multi-upstream router with selected pools - router = match Router::new_multi( - selected_pools, - None, - None, - true, - ).await { - Ok(router) => { - info!("Created distribution router with {} selected pools", desired_pool_count); - router - } - Err(e) => { - error!("Failed to create distribution router: {}", e); - std::process::exit(1); - } - }; + info!("Using hashrate distribution: {:?}", distribution); // Initialize upstream connections if let Err(e) = router.initialize_upstream_connections().await { error!("Failed to initialize upstream connections: {}", e); std::process::exit(1); } + // Set the distribution if let Err(e) = router.set_hashrate_distribution(distribution.clone()).await { error!("Failed to set hashrate distribution: {}", e); std::process::exit(1); } - info!("Set hashrate distribution: {:?}", distribution); - - info!("🔧 About to call initialize_proxy"); + info!("Starting proxy with hashrate distribution..."); initialize_proxy(&mut router, None, epsilon).await; + } else if pool_address_keys.len() > 1 { - info!("=== LATENCY-BASED SELECTION MODE (Condition 1) ==="); - info!("Multiple pools available, will select best latency"); + info!("=== LATENCY-BASED SELECTION MODE ==="); + info!("Testing pool latencies and selecting best pool..."); - // Actually select the best latency pool + // Test latency and select best pool let best_upstream = router.select_pool_connect().await; - if best_upstream.is_some() { - info!("Selected best upstream: {:?}", best_upstream); + + if let Some(ref upstream) = best_upstream { + info!("Selected best upstream: {:?}", upstream); + } else { + error!("Failed to connect to any upstream pool"); + std::process::exit(1); } + + info!("Starting proxy with latency-based selection..."); initialize_proxy(&mut router, best_upstream, epsilon).await; + } else { info!("=== SINGLE UPSTREAM MODE ==="); - // Single pool mode + // Connect to single pool let best_upstream = router.select_pool_connect().await; + + if let Some(ref upstream) = best_upstream { + info!("Connected to single upstream: {:?}", upstream); + } else { + error!("Failed to connect to upstream pool"); + std::process::exit(1); + } + + info!("Starting proxy with single upstream..."); initialize_proxy(&mut router, best_upstream, epsilon).await; } + // This line should never be reached since initialize_proxy runs forever + info!("Proxy stopped unexpectedly"); + // This code should never be reached because initialize_proxy runs forever info!("Proxy initialization complete"); // Wait a moment for connections to establish @@ -376,14 +222,24 @@ async fn main() { info!("🔍 VERIFYING NETWORK CONNECTIONS:"); // Check network connections - let output = std::process::Command::new("ss").args(&["-tn"]).output(); + let output = std::process::Command::new("sh") + .arg("-c") + .arg(&format!("ss -tn | grep -E \"({})\"", + pool_address_keys.iter() + .map(|(addr, _)| addr.to_string()) + .collect::>() + .join("|") + )) + .output(); if let Ok(output) = output { let connections = String::from_utf8_lossy(&output.stdout); let pool_connections: Vec<&str> = connections .lines() .filter(|line| { - line.contains("18.193.252.132:2000") || line.contains("3.74.36.119:2000") + pool_address_keys.iter().any(|(addr, _)| { + line.contains(&addr.to_string()) + }) }) .collect(); @@ -394,13 +250,13 @@ async fn main() { } // Show hashrate distribution - let total_hashrate = ProxyState::get_total_hashrate(); + let downstream_hashrate = ProxyState::get_downstream_hashrate(); let detailed_stats = router.get_detailed_connection_stats().await; info!("🚀 === INITIAL HASHRATE DISTRIBUTION ==="); info!( - "🔋 Total configured hashrate: {}", - HashUnit::format_value(total_hashrate) + "🔋 Total configured downstream hashrate: {}", + HashUnit::format_value(downstream_hashrate) ); let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); @@ -427,16 +283,16 @@ async fn main() { for (upstream_id, is_active, percentage) in detailed_stats { if is_active { // Calculate actual hashrate from percentage - let actual_hashrate = (percentage / 100.0) * total_hashrate; + let actual_hashrate = (percentage / 100.0) * downstream_hashrate; info!( - " ✅ {}: allocated {} ({:.1}% of total)", + " ✅ {}: allocated {} ({:.1}% of downstream)", upstream_id, HashUnit::format_value(actual_hashrate), percentage ); } else { info!( - " ❌ {}: 0.00T (INACTIVE - 0.0% of total)", + " ❌ {}: 0.00T (INACTIVE - 0.0% of downstream)", upstream_id ); } @@ -457,6 +313,7 @@ async fn main() { } } + async fn initialize_proxy( router: &mut Router, mut pool_addr: Option, @@ -681,43 +538,32 @@ async fn monitor_multi_upstream(router: Router, epsilon: Duration) { stats_report_counter = 0; let (total_count, active_count) = router.get_upstream_counts().await; - let total_hashrate = ProxyState::get_total_hashrate(); + let downstream_hashrate = ProxyState::get_downstream_hashrate(); info!("🚀 === MULTI-UPSTREAM HASHRATE REPORT ==="); info!("📊 Total upstreams: {}, Active: {}", total_count, active_count); - info!("🔋 Total hashrate: {}", HashUnit::format_value(total_hashrate)); + info!("🔋 Downstream hashrate: {}", HashUnit::format_value(downstream_hashrate)); if active_count > 0 { let detailed_stats = router.get_detailed_connection_stats().await; - let distribution = router.get_hashrate_distribution().await; - - // Check if using custom distribution - let is_custom = !distribution.is_empty() && distribution.len() == detailed_stats.len(); - if is_custom { - info!("🎯 Distribution mode: Custom"); - info!("📈 Distribution percentages: {:?}", distribution); - } else { - info!("⚖️ Distribution mode: Equal (automatic)"); + for (upstream_id, is_active, percentage) in detailed_stats { + if is_active { + // Calculate actual hashrate from percentage + let allocated_hashrate = (percentage / 100.0) * downstream_hashrate; + info!( + " ✅ {}: allocated {} ({:.1}% of downstream)", + upstream_id, + HashUnit::format_value(allocated_hashrate), + percentage + ); + } else { + info!( + " ❌ {}: 0.00T (INACTIVE - 0.0% of downstream)", + upstream_id + ); + } } - - for (upstream_id, is_active, percentage) in detailed_stats { - if is_active { - // Calculate actual hashrate from percentage - let allocated_hashrate = (percentage / 100.0) * total_hashrate; - info!( - " ✅ {}: allocated {} ({:.1}% of total)", - upstream_id, - HashUnit::format_value(allocated_hashrate), - percentage - ); - } else { - info!( - " ❌ {}: 0.00T (INACTIVE - 0.0% of total)", - upstream_id - ); - } -} } else { warn!("❌ No active upstream connections!"); } @@ -732,38 +578,6 @@ async fn monitor_multi_upstream(router: Router, epsilon: Duration) { } } } -/// Parses a hashrate string (e.g., "10T", "2.5P", "500E") into an f32 value in h/s. -fn parse_hashrate(hashrate_str: &str) -> Result { - let hashrate_str = hashrate_str.trim(); - if hashrate_str.is_empty() { - return Err("Hashrate cannot be empty. Expected format: '' (e.g., '10T', '2.5P', '5E'".to_string()); - } - - let unit = hashrate_str.chars().last().unwrap_or(' ').to_string(); - let num = &hashrate_str[..hashrate_str.len().saturating_sub(1)]; - - let num: f32 = num.parse().map_err(|_| { - format!( - "Invalid number '{}'. Expected format: '' (e.g., '10T', '2.5P', '5E')", - num - ) - })?; - - let multiplier = HashUnit::from_str(&unit) - .map(|unit| unit.multiplier()) - .ok_or_else(|| format!( - "Invalid unit '{}'. Expected 'T' (Terahash), 'P' (Petahash), or 'E' (Exahash). Example: '10T', '2.5P', '5E'", - unit - ))?; - - let hashrate = num * multiplier; - - if hashrate.is_infinite() || hashrate.is_nan() { - return Err("Hashrate too large or invalid".to_string()); - } - - Ok(hashrate) -} pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream diff --git a/src/proxy_state.rs b/src/proxy_state.rs index d9bcca46..ac73bd86 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -297,8 +297,8 @@ impl ProxyState { } } - /// Set the total hashrate to be distributed among upstreams - pub fn set_total_hashrate(hashrate: f32) { + /// Set the downstream hashrate to be distributed among upstreams + pub fn set_downstream_hashrate(hashrate: f32) { info!("Setting total hashrate to: {} h/s", hashrate); if PROXY_STATE .safe_lock(|state| { @@ -311,8 +311,8 @@ impl ProxyState { } } - /// Get the total hashrate - pub fn get_total_hashrate() -> f32 { + /// Get the downstream hashrate + pub fn get_downstream_hashrate() -> f32 { let mut hashrate = 0.0; if PROXY_STATE .safe_lock(|state| { diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 51eaaf00..88871817 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -122,7 +122,7 @@ impl MultiUpstreamManager { println!("calling set_hashrate_distribution in multi manager with: {:?}", distribution); let mut upstreams = self.upstreams.lock().await; - let total_hashrate = ProxyState::get_total_hashrate() as f64; + let total_hashrate = ProxyState::get_downstream_hashrate() as f64; println!("Total hashrate: {}", HashUnit::format_value(total_hashrate as f32)); diff --git a/src/translator/downstream/diff_management.rs b/src/translator/downstream/diff_management.rs index 56e4d6ef..ea6c4e86 100644 --- a/src/translator/downstream/diff_management.rs +++ b/src/translator/downstream/diff_management.rs @@ -29,7 +29,7 @@ impl Downstream { let (message, _) = diff_to_sv1_message(diff as f64)?; Downstream::send_message_downstream(self_.clone(), message.clone()).await; - let total_delay = Duration::from_secs(crate::ARGS.delay); + let total_delay = Duration::from_secs(crate::Configuration::delay()); let repeat_interval = Duration::from_secs(30); let self_clone = self_.clone(); @@ -379,9 +379,7 @@ mod test { } let elapsed_secs = start_time.elapsed().as_secs_f64(); - let hashrate = hashes as f64 / elapsed_secs; - let nominal_hash_rate = hashrate; - nominal_hash_rate + hashes as f64 / elapsed_secs } fn hash(share: &mut [u8; 80]) -> Target { @@ -486,4 +484,4 @@ mod test { // TODO make a test where unknown donwstream is simulated and we do not wait for it to produce // a share but we try to updated the estimated hash power every 2 seconds and updated the // target consequentially this shuold start to provide shares within a normal amount of time -} +} \ No newline at end of file diff --git a/src/translator/downstream/notify.rs b/src/translator/downstream/notify.rs index 39957906..5a81cecd 100644 --- a/src/translator/downstream/notify.rs +++ b/src/translator/downstream/notify.rs @@ -148,17 +148,17 @@ async fn start_update( connection_id: u32, ) -> Result<(), Error<'static>> { let handle = task::spawn(async move { - // Prevent difficulty adjustments until after crate::ARGS.delay elapses - tokio::time::sleep(std::time::Duration::from_secs(crate::ARGS.delay)).await; + // Prevent difficulty adjustments until after delay elapses + tokio::time::sleep(std::time::Duration::from_secs(crate::Configuration::delay())).await; loop { let share_count = crate::translator::utils::get_share_count(connection_id); let sleep_duration = if share_count >= crate::SHARE_PER_MIN * 3.0 || share_count <= crate::SHARE_PER_MIN / 3.0 { // TODO: this should only apply when after the first share has been received - std::time::Duration::from_millis(crate::ARGS.adjustment_interval) + std::time::Duration::from_millis(crate::Configuration::adjustment_interval()) } else { - std::time::Duration::from_millis(crate::ARGS.adjustment_interval) + std::time::Duration::from_millis(crate::Configuration::adjustment_interval()) }; tokio::time::sleep(sleep_duration).await; @@ -182,4 +182,4 @@ async fn start_update( TaskManager::add_update(task_manager, handle.into()) .await .map_err(|_| Error::TranslatorTaskManagerFailed) -} +} \ No newline at end of file From 9c464ec2d2b49889f4ac043c2277844992dd3d51 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Fri, 6 Jun 2025 22:57:38 +0530 Subject: [PATCH 14/21] feat: Remove unused aggregated_receiver and clean up code - Remove unused aggregated_receiver field from Router struct - Clean up unused imports and methods - Fix clippy warnings for better code quality - Remove dead code after initialize_proxy calls --- src/main.rs | 16 +- src/proxy_state.rs | 332 +-------------------------- src/router/mod.rs | 170 +------------- src/router/multi_upstream_manager.rs | 29 --- 4 files changed, 23 insertions(+), 524 deletions(-) diff --git a/src/main.rs b/src/main.rs index 1af4cb9b..2a38bee3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -221,7 +221,7 @@ if wants_distribution { // Check network connections let output = std::process::Command::new("sh") .arg("-c") - .arg(&format!("ss -tn | grep -E \"({})\"", + .arg(format!("ss -tn | grep -E \"({})\"", pool_address_keys.iter() .map(|(addr, _)| addr.to_string()) .collect::>() @@ -262,11 +262,9 @@ if wants_distribution { // Check if using custom distribution let mut is_custom_distribution = false; for (_, is_active, percentage) in &detailed_stats { - if *is_active { - if (percentage - (100.0 / active_count as f32)).abs() > 1.0 { - is_custom_distribution = true; - break; - } + if *is_active && (percentage - (100.0 / active_count as f32)).abs() > 1.0 { + is_custom_distribution = true; + break; } } @@ -299,7 +297,7 @@ if wants_distribution { // Start monitoring task for multi-upstream let router_clone = router.clone(); tokio::spawn(async move { - monitor_multi_upstream(router_clone, epsilon).await; + monitor_multi_upstream(router_clone).await; }); // Keep main thread alive @@ -328,7 +326,7 @@ async fn initialize_proxy( // Start the monitor task in the background let router_clone = router.clone(); tokio::spawn(async move { - monitor_multi_upstream(router_clone, epsilon).await; + monitor_multi_upstream(router_clone).await; }); info!("✅ Hashrate distribution monitor started"); @@ -523,7 +521,7 @@ async fn monitor( } // Consolidated monitoring function for multi-upstream mode -async fn monitor_multi_upstream(router: Router, epsilon: Duration) { +async fn monitor_multi_upstream(router: Router) { let mut stats_report_counter = 0; loop { diff --git a/src/proxy_state.rs b/src/proxy_state.rs index ac73bd86..36fc7552 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -116,7 +116,6 @@ pub struct ProxyState { // New fields for multiple upstream support pub upstream_connections: HashMap, pub total_hashrate: f32, - pub current_upstream_index: usize, } impl ProxyState { @@ -134,7 +133,6 @@ impl ProxyState { // Initialize new fields upstream_connections: HashMap::new(), total_hashrate: 0.0, - current_upstream_index: 0, } } @@ -265,38 +263,7 @@ impl ProxyState { } } - /// Add a new upstream connection - pub fn add_upstream_connection( - id: String, - url: String, - address: std::net::SocketAddr, - auth_key: key_utils::Secp256k1PublicKey, - connection_type: UpstreamType, - ) { - info!("Adding upstream connection: {} at {}", id, url); - - if PROXY_STATE - .safe_lock(|state| { - let connection = UpstreamConnection { - url, - address, - auth_key, - connection_type, - is_connected: false, - shares_submitted: 0, - shares_accepted: 0, - last_used: Instant::now(), - }; - - state.upstream_connections.insert(id, connection); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - + /// Set the downstream hashrate to be distributed among upstreams pub fn set_downstream_hashrate(hashrate: f32) { info!("Setting total hashrate to: {} h/s", hashrate); @@ -326,122 +293,7 @@ impl ProxyState { hashrate } - /// Get the next upstream in round-robin fashion - pub fn get_next_upstream( - ) -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { - let mut result = None; - - if PROXY_STATE - .safe_lock(|state| { - // Get IDs of all connected upstreams - let active_upstreams: Vec<&String> = state - .upstream_connections - .iter() - .filter(|(_, conn)| conn.is_connected) - .map(|(id, _)| id) - .collect(); - - if active_upstreams.is_empty() { - return; - } - - // Use round-robin to select the next upstream - if state.current_upstream_index >= active_upstreams.len() { - state.current_upstream_index = 0; - } - - let id = active_upstreams[state.current_upstream_index].clone(); - if let Some(conn) = state.upstream_connections.get(&id) { - result = Some((id.clone(), conn.address, conn.auth_key)); - } - - // Update index for next call - state.current_upstream_index = - (state.current_upstream_index + 1) % active_upstreams.len(); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - result - } - - /// Get the hashrate for a specific upstream (equal distribution) - #[allow(dead_code)] - pub fn get_hashrate_for_upstream(id: Option<&str>) -> f32 { - let mut hashrate = 0.0; - - if PROXY_STATE - .safe_lock(|state| { - // Count active connections - let active_count = state - .upstream_connections - .values() - .filter(|conn| conn.is_connected) - .count(); - - if active_count > 0 { - // Equal distribution - each upstream gets the same portion - hashrate = state.total_hashrate / active_count as f32; - - // If a specific ID was provided, check if it's active - if let Some(id) = id { - if let Some(conn) = state.upstream_connections.get(id) { - if !conn.is_connected { - // If this specific upstream isn't connected, return 0 - hashrate = 0.0; - } - } else { - // If this ID doesn't exist, return 0 - hashrate = 0.0; - } - } - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - hashrate - } - - /// Record a share submission to an upstream - #[allow(dead_code)] - pub fn record_share_submission(upstream_id: &str) { - if PROXY_STATE - .safe_lock(|state| { - if let Some(conn) = state.upstream_connections.get_mut(upstream_id) { - conn.shares_submitted += 1; - conn.last_used = Instant::now(); - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - - /// Record a share acceptance from an upstream - #[allow(dead_code)] - pub fn record_share_acceptance(upstream_id: &str) { - if PROXY_STATE - .safe_lock(|state| { - if let Some(conn) = state.upstream_connections.get_mut(upstream_id) { - conn.shares_accepted += 1; - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - + /// Update connection status for an upstream with timestamp pub fn set_upstream_connection_status(id: &str, connected: bool) { if PROXY_STATE @@ -463,77 +315,7 @@ impl ProxyState { } } - /// Get connection count - pub fn get_upstream_connection_count() -> usize { - let mut count = 0; - - if PROXY_STATE - .safe_lock(|state| { - count = state.upstream_connections.len(); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - count - } - - /// Get active connection count - pub fn get_active_upstream_count() -> usize { - let mut count = 0; - - if PROXY_STATE - .safe_lock(|state| { - count = state - .upstream_connections - .values() - .filter(|conn| conn.is_connected) - .count(); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - count - } - - /// Update upstream shares - pub fn update_upstream_shares(upstream_id: &str, submitted: u64, accepted: u64) { - if PROXY_STATE - .safe_lock(|state| { - if let Some(conn) = state.upstream_connections.get_mut(upstream_id) { - conn.shares_submitted += submitted; - conn.shares_accepted += accepted; - conn.last_used = Instant::now(); - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - - /// Remove an upstream connection - pub fn remove_upstream_connection(upstream_id: &str) { - info!("Removing upstream connection: {}", upstream_id); - - if PROXY_STATE - .safe_lock(|state| { - state.upstream_connections.remove(upstream_id); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - - /// Check if proxy is down + /// Check if proxy is down pub fn is_proxy_down() -> (bool, Option) { let errors = Self::get_errors(); if errors.is_ok() && errors.as_ref().unwrap().is_empty() { @@ -583,110 +365,4 @@ impl ProxyState { } } - /// Get upstream statistics - pub fn get_upstream_stats() -> Vec<(String, bool, u64, u64)> { - let mut stats = Vec::new(); - - if PROXY_STATE - .safe_lock(|state| { - stats = state - .upstream_connections - .iter() - .map(|(id, conn)| { - ( - id.clone(), - conn.is_connected, - conn.shares_submitted, - conn.shares_accepted, - ) - }) - .collect(); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - stats - } - - /// Get all upstream connections (including inactive ones) - pub fn get_all_upstream_connections() -> Vec<( - String, - std::net::SocketAddr, - key_utils::Secp256k1PublicKey, - bool, - )> { - let mut connections = Vec::new(); - - if PROXY_STATE - .safe_lock(|state| { - connections = state - .upstream_connections - .iter() - .map(|(id, conn)| (id.clone(), conn.address, conn.auth_key, conn.is_connected)) - .collect(); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - connections - } - - /// Get only active upstream connections (existing method is fine) - pub fn get_upstream_connections( - ) -> Vec<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { - let mut connections = Vec::new(); - - if PROXY_STATE - .safe_lock(|state| { - connections = state - .upstream_connections - .iter() - .filter(|(_, conn)| conn.is_connected) - .map(|(id, conn)| (id.clone(), conn.address, conn.auth_key)) - .collect(); - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - connections - } - - /// Mark an upstream as inactive (disconnect it) - pub fn mark_upstream_inactive(upstream_id: &str) { - Self::set_upstream_connection_status(upstream_id, false); - } - - /// Get best upstream based on latency or other criteria - /// For now, just returns the first active upstream - pub fn get_best_upstream( - ) -> Option<(String, std::net::SocketAddr, key_utils::Secp256k1PublicKey)> { - let mut result = None; - - if PROXY_STATE - .safe_lock(|state| { - // Find the first active upstream - for (id, conn) in &state.upstream_connections { - if conn.is_connected { - result = Some((id.clone(), conn.address, conn.auth_key)); - break; - } - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - - result - } -} + } diff --git a/src/router/mod.rs b/src/router/mod.rs index f8391882..97873c93 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -17,7 +17,7 @@ use tokio::{ watch, }, }; -use tracing::{error, info, warn}; +use tracing::{error, info}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, @@ -35,7 +35,6 @@ pub struct Router { pub keys: Vec, pub current_pool: Option, pub upstream_manager: Option, - pub aggregated_receiver: Option>>, // Keep these fields for backward compatibility with single upstream mode pub auth_pub_k: Secp256k1PublicKey, @@ -43,11 +42,7 @@ pub struct Router { pub timer: Option, pub latency_tx: tokio::sync::watch::Sender>, pub latency_rx: tokio::sync::watch::Receiver>, - // Remove round-robin fields entirely - // use_round_robin: bool, - // use_parallel: bool, // This will be determined by presence of upstream_manager -} -// Remove Clone derive since Receiver can't be cloned + } impl Clone for Router { fn clone(&self) -> Self { Self { @@ -55,7 +50,6 @@ impl Clone for Router { keys: self.keys.clone(), current_pool: self.current_pool, upstream_manager: self.upstream_manager.clone(), - aggregated_receiver: None, // Can't clone receiver auth_pub_k: self.auth_pub_k, setup_connection_msg: self.setup_connection_msg.clone(), timer: self.timer, @@ -82,7 +76,6 @@ impl Router { keys: auth_keys, current_pool: None, upstream_manager: None, - aggregated_receiver: None, auth_pub_k, setup_connection_msg, timer, @@ -110,7 +103,6 @@ impl Router { keys, current_pool: None, upstream_manager: None, - aggregated_receiver: None, auth_pub_k, setup_connection_msg, timer, @@ -132,17 +124,15 @@ impl Router { let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); // Create upstream manager only if we want custom distribution - let (upstream_manager, aggregated_receiver) = if use_distribution { - let manager = MultiUpstreamManager::new( + let upstream_manager = if use_distribution { + Some(MultiUpstreamManager::new( pool_socket_addresses.clone(), - keys[0], // or use a key per upstream if needed + keys[0], setup_connection_msg.clone(), - timer.clone(), - ); - let receiver = None; // or whatever is appropriate for your design - (Some(manager), receiver) + timer, + )) } else { - (None, None) + None }; // Use @@ -158,36 +148,13 @@ impl Router { keys, current_pool: None, upstream_manager, - aggregated_receiver, auth_pub_k, setup_connection_msg, timer, latency_tx, latency_rx, }) - } - - /// Get aggregated receiver for multi-upstream mode - pub async fn get_aggregated_receiver(&self) -> Option>> { - if let Some(ref manager) = self.upstream_manager { - manager.get_aggregated_receiver().await - } else { - None - } - } - - /// Get aggregated sender for multi-upstream mode - pub async fn get_aggregated_sender(&self) -> tokio::sync::mpsc::Sender> { - if let Some(ref manager) = self.upstream_manager { - manager.get_aggregated_sender().await - } else { - // Return a dummy sender if no manager exists - let (sender, _) = tokio::sync::mpsc::channel(1); - sender - } - } - - + } /// Get detailed connection statistics pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { @@ -198,48 +165,12 @@ impl Router { } } - /// Get current hashrate distribution - pub async fn get_hashrate_distribution(&self) -> Vec { - if let Some(ref manager) = self.upstream_manager { - manager.get_hashrate_distribution().await - } else { - vec![] - } - } - /// Check if multi-upstream is enabled pub fn is_multi_upstream_enabled(&self) -> bool { self.upstream_manager.is_some() } - /// Check if the router is using parallel mode - true when MultiUpstreamManager is present - pub fn is_parallel_mode(&self) -> bool { - self.upstream_manager.is_some() - } - - /// Check if the router is using custom distribution mode - pub fn is_distribution_mode(&self) -> bool { - self.upstream_manager.is_some() - } - - /// Get the best pools based on latency and desired count - pub async fn select_best_pools(&mut self, desired_count: usize) -> Vec<(SocketAddr, Duration)> { - let mut pool_latencies = Vec::new(); - - // Get latencies for all pools - for &pool_addr in &self.pool_socket_addresses { - if let Ok(latency) = self.get_latency(pool_addr).await { - pool_latencies.push((pool_addr, latency)); - } - } - - // Sort by latency (lowest first) - pool_latencies.sort_by_key(|(_, latency)| *latency); - - // Return the best pools up to desired count - pool_latencies.into_iter().take(desired_count).collect() - } - + /// Checks for faster upstream and switches to it if found pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { // For multi-upstream mode, we don't switch since we use all simultaneously @@ -451,9 +382,7 @@ impl Router { Ok(sum_of_latencies) } - pub fn get_current_upstream(&self) -> Option { - self.current_pool - } + /// Initialize upstream connections for the manager pub async fn initialize_upstream_connections(&mut self) -> Result<(), String> { @@ -525,38 +454,7 @@ pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result< } } - /// Get list of active upstream IDs - pub async fn get_active_upstreams(&self) -> Vec { - if let Some(ref manager) = self.upstream_manager { - manager.get_active_upstream_ids().await - } else { - vec!["upstream-0".to_string()] // Single upstream mode - } - } - - /// Switch to a specific upstream by ID - pub async fn switch_to_upstream(&mut self, upstream_id: &str) -> Result<(), String> { - // Find the upstream address by ID - if let Ok(idx) = upstream_id - .strip_prefix("upstream-") - .and_then(|s| s.parse::().ok()) - .ok_or("Invalid upstream ID format".to_string()) - { - if idx < self.pool_socket_addresses.len() { - // Fix field name - let addr = self.pool_socket_addresses[idx]; // Fix field name - self.current_pool = Some(addr); - info!("Switched to upstream {}: {:?}", upstream_id, addr); - Ok(()) - } else { - Err(format!("Upstream index {} out of range", idx)) - } - } else { - Err("Invalid upstream ID format".to_string()) - } - } - - /// Select the best pool based on latency +/// Select the best pool based on latency async fn select_pool(&mut self) -> Option<(SocketAddr, Duration)> { let mut best_pool = None; let mut best_latency = Duration::from_secs(u64::MAX); @@ -572,50 +470,6 @@ pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result< best_pool.map(|pool| (pool, best_latency)) } - - /// Select best pools for hashrate distribution based on latency - /// This implements condition 2: choose N best pools from available pools - pub async fn select_best_pools_for_distribution(&mut self, desired_count: usize) -> Vec<(SocketAddr, Secp256k1PublicKey)> { - // If we have exactly the desired count or fewer, use all - if self.pool_socket_addresses.len() <= desired_count { - info!("Using all {} available pools (desired: {})", self.pool_socket_addresses.len(), desired_count); - return self.pool_socket_addresses - .iter() - .zip(self.keys.iter()) - .map(|(&addr, &key)| (addr, key)) - .collect(); - } - - // Get latencies for all pools - let mut pool_latencies = Vec::new(); - info!("Testing latencies for {} pools to select {} best", self.pool_socket_addresses.len(), desired_count); - - for (i, &pool_addr) in self.pool_socket_addresses.iter().enumerate() { - match self.get_latency(pool_addr).await { - Ok(latency) => { - info!("Pool {}: latency {:?}", pool_addr, latency); - pool_latencies.push((pool_addr, self.keys[i], latency)); - } - Err(_) => { - warn!("Failed to get latency for pool {}", pool_addr); - } - } - } - - // Sort by latency (lowest first) and take the best ones - pool_latencies.sort_by_key(|(_, _, latency)| *latency); - let selected: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_latencies - .into_iter() - .take(desired_count) - .map(|(addr, key, latency)| { - info!("Selected pool {} with latency {:?}", addr, latency); - (addr, key) - }) - .collect(); - - info!("Selected {} pools based on latency", selected.len()); - selected - } } /// Track latencies for various stages of pool connection setup. diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 88871817..6a8bc0ee 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -3,11 +3,9 @@ use tokio::sync::Mutex; use tracing::{error, info}; use crate::{ minin_pool_connection, - shared::utils::AbortOnDrop, proxy_state::ProxyState, HashUnit, }; -use demand_share_accounting_ext::parser::PoolExtMessages; use key_utils::Secp256k1PublicKey; #[derive(Clone)] @@ -160,33 +158,6 @@ pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { .collect() } - pub async fn get_hashrate_distribution(&self) -> Vec { - let upstreams = self.upstreams.lock().await; - upstreams - .values() - .map(|conn| conn.allocated_percentage) - .collect() - } - - pub async fn get_active_upstream_ids(&self) -> Vec { - let upstreams = self.upstreams.lock().await; - upstreams - .iter() - .filter(|(_, conn)| conn.is_active) - .map(|(id, _)| id.clone()) - .collect() - } - - // Dummy implementations for sender/receiver if needed - pub async fn get_aggregated_receiver(&self) -> Option>> { - None - } - - pub async fn get_aggregated_sender(&self) -> tokio::sync::mpsc::Sender> { - let (sender, _) = tokio::sync::mpsc::channel(1); - sender - } - // Add upstream (called from Router) pub async fn add_upstream( &self, From bca55c258fc463d0fd7c877872addd165b11d85a Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Sat, 7 Jun 2025 15:14:42 +0530 Subject: [PATCH 15/21] refactor: Clean up code formatting and remove unnecessary whitespace across multiple files --- src/config.rs | 29 ++-- src/ingress/sv1_ingress.rs | 2 +- src/jd_client/mod.rs | 2 +- src/main.rs | 166 ++++++++++--------- src/proxy_state.rs | 7 +- src/router/mod.rs | 89 +++++----- src/router/multi_upstream_manager.rs | 102 ++++++------ src/share_accounter/mod.rs | 2 +- src/translator/downstream/diff_management.rs | 2 +- src/translator/downstream/notify.rs | 2 +- src/translator/mod.rs | 2 +- 11 files changed, 200 insertions(+), 205 deletions(-) diff --git a/src/config.rs b/src/config.rs index e7277b9d..0c112f55 100644 --- a/src/config.rs +++ b/src/config.rs @@ -59,10 +59,9 @@ pub struct ConfigFile { pub test: Option, pub listening_addr: Option, pub api_server_port: Option, - pub hashrate_distribution: Option>, + pub hashrate_distribution: Option>, } - pub struct Configuration { token: Option, tp_address: Option, @@ -80,7 +79,6 @@ pub struct Configuration { } impl Configuration { - pub fn token() -> Option { CONFIG.token.clone() } @@ -150,11 +148,10 @@ impl Configuration { CONFIG.hashrate_distribution.clone() } - pub fn wants_hashrate_distribution() -> bool { CONFIG.hashrate_distribution.is_some() } - + // Loads config from CLI, file, or env vars with precedence: CLI > file > env. fn load_config() -> Self { let args = Args::parse(); @@ -175,7 +172,7 @@ impl Configuration { test: None, listening_addr: None, api_server_port: None, - hashrate_distribution: None, + hashrate_distribution: None, }); let token = args @@ -304,23 +301,21 @@ impl Configuration { .unwrap_or("off".to_string()); let test = args.test || config.test.unwrap_or(false) || std::env::var("TEST").is_ok(); -let hashrate_distribution = args + let hashrate_distribution = args .hashrate_distribution .or(config.hashrate_distribution) .or_else(|| { - std::env::var("HASHRATE_DISTRIBUTION") - .ok() - .and_then(|s| { - s.split(',') - .map(|s| s.trim().parse::().ok()) - .collect::>>() - }) + std::env::var("HASHRATE_DISTRIBUTION").ok().and_then(|s| { + s.split(',') + .map(|s| s.trim().parse::().ok()) + .collect::>>() + }) }); - + Configuration { token, tp_address, - pool_addresses, + pool_addresses, test_pool_addresses, interval, delay, @@ -331,7 +326,7 @@ let hashrate_distribution = args listening_addr, api_server_port, hashrate_distribution, - } + } } } diff --git a/src/ingress/sv1_ingress.rs b/src/ingress/sv1_ingress.rs index 1404bb1c..4bd2c6f8 100644 --- a/src/ingress/sv1_ingress.rs +++ b/src/ingress/sv1_ingress.rs @@ -117,4 +117,4 @@ impl Downstream { Err(_) => Sv1IngressError::TaskFailed, } } -} \ No newline at end of file +} diff --git a/src/jd_client/mod.rs b/src/jd_client/mod.rs index 97c9767c..79fd87fa 100644 --- a/src/jd_client/mod.rs +++ b/src/jd_client/mod.rs @@ -271,4 +271,4 @@ async fn retry_connection(address: String) { break; } } -} \ No newline at end of file +} diff --git a/src/main.rs b/src/main.rs index d8e31405..c24c11f1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -27,7 +27,6 @@ mod share_accounter; mod shared; mod translator; - const TRANSLATOR_BUFFER_SIZE: usize = 32; const MIN_EXTRANONCE_SIZE: u16 = 6; const MIN_EXTRANONCE2_SIZE: u16 = 5; @@ -95,38 +94,37 @@ async fn main() { // Set downstream hashrate using Configuration pattern let downstream_hashrate = Configuration::downstream_hashrate(); ProxyState::set_downstream_hashrate(downstream_hashrate); - info!("Set downstream hashrate to: {}", HashUnit::format_value(downstream_hashrate)); + info!( + "Set downstream hashrate to: {}", + HashUnit::format_value(downstream_hashrate) + ); // Get pool addresses with auth keys using Configuration pattern - let pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_addresses.iter().map(|&addr| { - (addr, auth_pub_k) - }).collect(); + let pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_addresses + .iter() + .map(|&addr| (addr, auth_pub_k)) + .collect(); // Determine hashrate distribution using Configuration pattern - let wants_distribution = Configuration::wants_hashrate_distribution(); + let wants_distribution = Configuration::wants_hashrate_distribution(); -if wants_distribution { - info!("Hashrate distribution is enabled"); -} else { - info!("Hashrate distribution is disabled"); -} + if wants_distribution { + info!("Hashrate distribution is enabled"); + } else { + info!("Hashrate distribution is disabled"); + } // Create router based on configuration let mut router = if wants_distribution { // Multi-upstream with distribution - match Router::new_multi( - pool_address_keys.clone(), - None, - None, - true, - ).await { + match Router::new_multi(pool_address_keys.clone(), None, None, true).await { Ok(router) => router, Err(e) => { error!("Failed to create multi-upstream router: {}", e); std::process::exit(1); } } - } else if pool_address_keys.len() > 1 { + } else if pool_address_keys.len() > 1 { // Multiple pools, latency-based selection Router::new_with_keys(pool_address_keys.clone(), None, None) } else { @@ -137,17 +135,16 @@ if wants_distribution { let epsilon = Duration::from_millis(30_000); // Handle the three different scenarios - + if wants_distribution { info!("=== HASHRATE DISTRIBUTION MODE ==="); - + // Get the distribution from config - let distribution = Configuration::hashrate_distribution() - .unwrap_or_else(|| { - let count = pool_address_keys.len(); - let equal_percentage = 100.0 / count as f32; - vec![equal_percentage; count] - }); + let distribution = Configuration::hashrate_distribution().unwrap_or_else(|| { + let count = pool_address_keys.len(); + let equal_percentage = 100.0 / count as f32; + vec![equal_percentage; count] + }); info!("Using hashrate distribution: {:?}", distribution); @@ -156,7 +153,7 @@ if wants_distribution { error!("Failed to initialize upstream connections: {}", e); std::process::exit(1); } - + // Set the distribution if let Err(e) = router.set_hashrate_distribution(distribution.clone()).await { error!("Failed to set hashrate distribution: {}", e); @@ -165,14 +162,13 @@ if wants_distribution { info!("Starting proxy with hashrate distribution..."); initialize_proxy(&mut router, None, epsilon).await; - } else if pool_address_keys.len() > 1 { info!("=== LATENCY-BASED SELECTION MODE ==="); info!("Testing pool latencies and selecting best pool..."); - + // Test latency and select best pool let best_upstream = router.select_pool_connect().await; - + if let Some(ref upstream) = best_upstream { info!("Selected best upstream: {:?}", upstream); } else { @@ -182,13 +178,12 @@ if wants_distribution { info!("Starting proxy with latency-based selection..."); initialize_proxy(&mut router, best_upstream, epsilon).await; - } else { info!("=== SINGLE UPSTREAM MODE ==="); - + // Connect to single pool let best_upstream = router.select_pool_connect().await; - + if let Some(ref upstream) = best_upstream { info!("Connected to single upstream: {:?}", upstream); } else { @@ -216,8 +211,10 @@ if wants_distribution { // Check network connections let output = std::process::Command::new("sh") .arg("-c") - .arg(format!("ss -tn | grep -E \"({})\"", - pool_address_keys.iter() + .arg(format!( + "ss -tn | grep -E \"({})\"", + pool_address_keys + .iter() .map(|(addr, _)| addr.to_string()) .collect::>() .join("|") @@ -229,9 +226,9 @@ if wants_distribution { let pool_connections: Vec<&str> = connections .lines() .filter(|line| { - pool_address_keys.iter().any(|(addr, _)| { - line.contains(&addr.to_string()) - }) + pool_address_keys + .iter() + .any(|(addr, _)| line.contains(&addr.to_string())) }) .collect(); @@ -240,7 +237,7 @@ if wants_distribution { info!(" 🔗 {}", conn.trim()); } } - + // Show hashrate distribution let downstream_hashrate = ProxyState::get_downstream_hashrate(); let detailed_stats = router.get_detailed_connection_stats().await; @@ -250,45 +247,51 @@ if wants_distribution { "🔋 Total configured downstream hashrate: {}", HashUnit::format_value(downstream_hashrate) ); - - let active_count = detailed_stats.iter().filter(|(_, is_active, _)| *is_active).count(); - + + let active_count = detailed_stats + .iter() + .filter(|(_, is_active, _)| *is_active) + .count(); + if active_count > 0 { // Check if using custom distribution let mut is_custom_distribution = false; for (_, is_active, percentage) in &detailed_stats { - if *is_active && (percentage - (100.0 / active_count as f32)).abs() > 1.0 { - is_custom_distribution = true; - break; - } -} - + if *is_active && (percentage - (100.0 / active_count as f32)).abs() > 1.0 { + is_custom_distribution = true; + break; + } + } + if is_custom_distribution { info!("🎯 Mode: Custom Distribution (from config)"); } else { let percentage_per_upstream = 100.0 / active_count as f32; - info!("🌐 Mode: Equal Distribution ({:.1}% per upstream)", percentage_per_upstream); + info!( + "🌐 Mode: Equal Distribution ({:.1}% per upstream)", + percentage_per_upstream + ); } - for (upstream_id, is_active, percentage) in detailed_stats { - if is_active { - // Calculate actual hashrate from percentage - let actual_hashrate = (percentage / 100.0) * downstream_hashrate; - info!( - " ✅ {}: allocated {} ({:.1}% of downstream)", - upstream_id, - HashUnit::format_value(actual_hashrate), - percentage - ); - } else { - info!( - " ❌ {}: 0.00T (INACTIVE - 0.0% of downstream)", - upstream_id - ); - } -} + for (upstream_id, is_active, percentage) in detailed_stats { + if is_active { + // Calculate actual hashrate from percentage + let actual_hashrate = (percentage / 100.0) * downstream_hashrate; + info!( + " ✅ {}: allocated {} ({:.1}% of downstream)", + upstream_id, + HashUnit::format_value(actual_hashrate), + percentage + ); + } else { + info!( + " ❌ {}: 0.00T (INACTIVE - 0.0% of downstream)", + upstream_id + ); + } + } info!("========================================"); - + // Start monitoring task for multi-upstream let router_clone = router.clone(); tokio::spawn(async move { @@ -303,7 +306,6 @@ if wants_distribution { } } - async fn initialize_proxy( router: &mut Router, mut pool_addr: Option, @@ -313,7 +315,7 @@ async fn initialize_proxy( if router.is_multi_upstream_enabled() { info!("Initializing proxy in HASHRATE DISTRIBUTION mode"); info!("🎯 Focus: Distribute hashrate allocation across upstreams"); - + // Start API server for monitoring let stats_sender = api::stats::StatsSender::new(); let _server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); @@ -327,11 +329,11 @@ async fn initialize_proxy( info!("✅ Hashrate distribution monitor started"); return; } - + // Single upstream mode only loop { let stats_sender = api::stats::StatsSender::new(); - + let (send_to_pool, recv_from_pool, pool_connection_abortable) = match router.connect_pool(pool_addr).await { Ok(connection) => connection, @@ -518,7 +520,7 @@ async fn monitor( // Consolidated monitoring function for multi-upstream mode async fn monitor_multi_upstream(router: Router) { let mut stats_report_counter = 0; - + loop { tokio::time::sleep(Duration::from_secs(10)).await; stats_report_counter += 1; @@ -526,17 +528,23 @@ async fn monitor_multi_upstream(router: Router) { // Generate periodic reports every 30 seconds (3 * 10 seconds) if stats_report_counter >= 3 { stats_report_counter = 0; - + let (total_count, active_count) = router.get_upstream_counts().await; let downstream_hashrate = ProxyState::get_downstream_hashrate(); - + info!("🚀 === MULTI-UPSTREAM HASHRATE REPORT ==="); - info!("📊 Total upstreams: {}, Active: {}", total_count, active_count); - info!("🔋 Downstream hashrate: {}", HashUnit::format_value(downstream_hashrate)); - + info!( + "📊 Total upstreams: {}, Active: {}", + total_count, active_count + ); + info!( + "🔋 Downstream hashrate: {}", + HashUnit::format_value(downstream_hashrate) + ); + if active_count > 0 { let detailed_stats = router.get_detailed_connection_stats().await; - + for (upstream_id, is_active, percentage) in detailed_stats { if is_active { // Calculate actual hashrate from percentage @@ -572,7 +580,7 @@ async fn monitor_multi_upstream(router: Router) { pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream -} +} enum HashUnit { Tera, Peta, diff --git a/src/proxy_state.rs b/src/proxy_state.rs index 36fc7552..9677acac 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -263,7 +263,6 @@ impl ProxyState { } } - /// Set the downstream hashrate to be distributed among upstreams pub fn set_downstream_hashrate(hashrate: f32) { info!("Setting total hashrate to: {} h/s", hashrate); @@ -293,7 +292,6 @@ impl ProxyState { hashrate } - /// Update connection status for an upstream with timestamp pub fn set_upstream_connection_status(id: &str, connected: bool) { if PROXY_STATE @@ -315,7 +313,7 @@ impl ProxyState { } } - /// Check if proxy is down + /// Check if proxy is down pub fn is_proxy_down() -> (bool, Option) { let errors = Self::get_errors(); if errors.is_ok() && errors.as_ref().unwrap().is_empty() { @@ -364,5 +362,4 @@ impl ProxyState { Ok(errors) } } - - } +} diff --git a/src/router/mod.rs b/src/router/mod.rs index 010878de..ffb45767 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -31,7 +31,7 @@ pub use multi_upstream_manager::MultiUpstreamManager; /// Router handles connection to Multiple upstreams. pub struct Router { - pub pool_addresses: Vec, + pub pool_addresses: Vec, pub keys: Vec, pub current_pool: Option, pub upstream_manager: Option, @@ -42,11 +42,11 @@ pub struct Router { pub timer: Option, pub latency_tx: tokio::sync::watch::Sender>, pub latency_rx: tokio::sync::watch::Receiver>, - } +} impl Clone for Router { fn clone(&self) -> Self { Self { - pool_addresses: self. pool_addresses.clone(), + pool_addresses: self.pool_addresses.clone(), keys: self.keys.clone(), current_pool: self.current_pool, upstream_manager: self.upstream_manager.clone(), @@ -72,7 +72,7 @@ impl Router { let auth_keys = vec![auth_pub_k; pool_addresses.len()]; Self { - pool_addresses: pool_addresses, + pool_addresses: pool_addresses, keys: auth_keys, current_pool: None, upstream_manager: None, @@ -91,7 +91,7 @@ impl Router { setup_connection_msg: Option>, timer: Option, ) -> Self { - let pool_addresses: Vec = + let pool_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); let auth_pub_k = keys[0]; // Use first key as primary @@ -99,7 +99,7 @@ impl Router { let (latency_tx, latency_rx) = watch::channel(None); Self { - pool_addresses, + pool_addresses, keys, current_pool: None, upstream_manager: None, @@ -119,14 +119,14 @@ impl Router { timer: Option, use_distribution: bool, // Changed from use_parallel to use_distribution ) -> Result { - let pool_addresses: Vec = + let pool_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); // Create upstream manager only if we want custom distribution - let upstream_manager = if use_distribution { + let upstream_manager = if use_distribution { Some(MultiUpstreamManager::new( - pool_addresses.clone(), + pool_addresses.clone(), keys[0], setup_connection_msg.clone(), timer, @@ -144,7 +144,7 @@ impl Router { let (latency_tx, latency_rx) = watch::channel(None); Ok(Self { - pool_addresses, + pool_addresses, keys, current_pool: None, upstream_manager, @@ -154,7 +154,7 @@ impl Router { latency_tx, latency_rx, }) - } + } /// Get detailed connection statistics pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { @@ -170,9 +170,8 @@ impl Router { self.upstream_manager.is_some() } - /// Checks for faster upstream and switches to it if found - pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { + pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { // For multi-upstream mode, we don't switch since we use all simultaneously if self.is_multi_upstream_enabled() { // In distribution mode, we don't need to switch upstreams @@ -217,7 +216,8 @@ impl Router { Ok(latency) => latency, Err(e) => { error!("Failed to get latency: {:?}", e); - return Some(best_pool); } + return Some(best_pool); + } }; // saturating_sub is used to avoid panic on negative duration result if best_pool_latency < current_latency.saturating_sub(epsilon) { @@ -236,7 +236,6 @@ impl Router { None } - /// Select the best pool for connection pub async fn select_pool_connect(&mut self) -> Option { info!("Selecting the best upstream"); @@ -280,12 +279,11 @@ impl Router { info!("Upstream {:?} selected", pool); // Find the matching auth key for this address - fix field name - let auth_key = - if let Some(index) = self. pool_addresses.iter().position(|&a| a == pool) { - self.keys[index] - } else { - self.auth_pub_k - }; + let auth_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { + self.keys[index] + } else { + self.auth_pub_k + }; match minin_pool_connection::connect_pool( pool, @@ -299,14 +297,14 @@ impl Router { // Update ProxyState with successful connection let upstream_id = format!( "upstream-{}", - self. pool_addresses + self.pool_addresses .iter() .position(|&a| a == pool) .unwrap_or(0) ); ProxyState::set_upstream_connection_status(&upstream_id, true); - // Update current pool address + // Update current pool address crate::POOL_ADDRESS .safe_lock(|pool_address| { *pool_address = Some(pool); @@ -323,7 +321,7 @@ impl Router { // Update ProxyState with failed connection let upstream_id = format!( "upstream-{}", - self. pool_addresses + self.pool_addresses .iter() .position(|&a| a == pool) .unwrap_or(0) @@ -338,15 +336,12 @@ impl Router { /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { // Find the auth key for this address - fix field names - let auth_pub_key = if let Some(index) = self - . pool_addresses - .iter() - .position(|&a| a == pool_address) - { - self.keys[index] - } else { - self.auth_pub_k - }; + let auth_pub_key = + if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool_address) { + self.keys[index] + } else { + self.auth_pub_k + }; let mut pool = PoolLatency::new(pool_address); let setup_connection_msg = self.setup_connection_msg.as_ref(); @@ -404,23 +399,16 @@ impl Router { Ok(sum_of_latencies) } - - /// Initialize upstream connections for the manager pub async fn initialize_upstream_connections(&mut self) -> Result<(), String> { if let Some(ref manager) = self.upstream_manager { info!( "Initializing {} upstream connections", - self. pool_addresses.len() + self.pool_addresses.len() ); // Add each unique upstream only once - fix field names - for (idx, (addr, key)) in self - . pool_addresses - .iter() - .zip(self.keys.iter()) - .enumerate() - { + for (idx, (addr, key)) in self.pool_addresses.iter().zip(self.keys.iter()).enumerate() { let id = format!("upstream-{}", idx); info!("Adding upstream {}: {} ({})", id, addr, key); @@ -439,7 +427,7 @@ impl Router { } } - // IMPORTANT: Initialize connections BEFORE any hashrate distribution is set + // IMPORTANT: Initialize connections BEFORE any hashrate distribution is set // This ensures add_upstream() calls are complete and won't overwrite distribution manager.initialize_connections().await; @@ -458,13 +446,19 @@ impl Router { let active = upstreams.values().filter(|u| u.is_active).count(); (total, active) } else { - (self. pool_addresses.len(), 1) // Single upstream mode + (self.pool_addresses.len(), 1) // Single upstream mode } } /// Sets the hashrate distribution for the upstream manager. -pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result<(), &'static str> { - info!("🔧 Router::set_hashrate_distribution called with: {:?}", distribution); - + pub async fn set_hashrate_distribution( + &self, + distribution: Vec, + ) -> Result<(), &'static str> { + info!( + "🔧 Router::set_hashrate_distribution called with: {:?}", + distribution + ); + if let Some(ref manager) = self.upstream_manager { info!("🔧 Calling manager.set_hashrate_distribution"); manager.set_hashrate_distribution(distribution).await; @@ -475,7 +469,6 @@ pub async fn set_hashrate_distribution(&self, distribution: Vec) -> Result< Err("No upstream manager available") } } - } /// Track latencies for various stages of pool connection setup. diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 6a8bc0ee..402a2ae0 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -1,12 +1,8 @@ +use crate::{minin_pool_connection, proxy_state::ProxyState, HashUnit}; +use key_utils::Secp256k1PublicKey; use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; use tokio::sync::Mutex; use tracing::{error, info}; -use crate::{ - minin_pool_connection, - proxy_state::ProxyState, - HashUnit, -}; -use key_utils::Secp256k1PublicKey; #[derive(Clone)] pub struct MultiUpstreamManager { @@ -44,7 +40,9 @@ impl MultiUpstreamManager { pub fn new( upstreams: Vec, auth_pub_k: Secp256k1PublicKey, - setup_connection_msg: Option>, + setup_connection_msg: Option< + roles_logic_sv2::common_messages_sv2::SetupConnection<'static>, + >, timer: Option, ) -> Self { let upstream_map = upstreams @@ -52,7 +50,7 @@ impl MultiUpstreamManager { .enumerate() .map(|(i, addr)| (format!("upstream-{}", i), UpstreamConnection::new(addr))) .collect(); - + Self { upstreams: Arc::new(Mutex::new(upstream_map)), auth_pub_k, @@ -117,13 +115,19 @@ impl MultiUpstreamManager { /// Set hashrate distribution for each upstream (percentages) pub async fn set_hashrate_distribution(&self, distribution: Vec) { - println!("calling set_hashrate_distribution in multi manager with: {:?}", distribution); - + println!( + "calling set_hashrate_distribution in multi manager with: {:?}", + distribution + ); + let mut upstreams = self.upstreams.lock().await; let total_hashrate = ProxyState::get_downstream_hashrate() as f64; - - println!("Total hashrate: {}", HashUnit::format_value(total_hashrate as f32)); - + + println!( + "Total hashrate: {}", + HashUnit::format_value(total_hashrate as f32) + ); + for (i, (_id, conn)) in upstreams.iter_mut().enumerate() { if let Some(&percentage) = distribution.get(i) { println!("setting upstream {} to {}%", i, percentage); @@ -141,48 +145,46 @@ impl MultiUpstreamManager { pub async fn get_upstreams(&self) -> HashMap { self.upstreams.lock().await.clone() } -pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { - let upstreams = self.upstreams.lock().await; - - upstreams - .iter() - .map(|(id, conn)| { - println!( - "Upstream {}: {}% allocation", - id, - conn.allocated_percentage - ); - // Return the percentage, not the hashrate - (id.clone(), conn.is_active, conn.allocated_percentage) - }) - .collect() -} + pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { + let upstreams = self.upstreams.lock().await; + + upstreams + .iter() + .map(|(id, conn)| { + println!("Upstream {}: {}% allocation", id, conn.allocated_percentage); + // Return the percentage, not the hashrate + (id.clone(), conn.is_active, conn.allocated_percentage) + }) + .collect() + } // Add upstream (called from Router) - pub async fn add_upstream( - &self, - id: String, - address: SocketAddr, - _key: Secp256k1PublicKey, - _setup_connection_msg: Option>, - _timer: Option, -) -> Result<(), String> { - let mut upstreams = self.upstreams.lock().await; - - // Check if upstream already exists - if so, don't overwrite it - if upstreams.contains_key(&id) { - println!("⚠️ Upstream {} already exists, not overwriting", id); - return Ok(()); + pub async fn add_upstream( + &self, + id: String, + address: SocketAddr, + _key: Secp256k1PublicKey, + _setup_connection_msg: Option< + roles_logic_sv2::common_messages_sv2::SetupConnection<'static>, + >, + _timer: Option, + ) -> Result<(), String> { + let mut upstreams = self.upstreams.lock().await; + + // Check if upstream already exists - if so, don't overwrite it + if upstreams.contains_key(&id) { + println!("⚠️ Upstream {} already exists, not overwriting", id); + return Ok(()); + } + + // Only insert if it doesn't exist + upstreams.insert(id.clone(), UpstreamConnection::new(address)); + println!("✅ Added new upstream {}", id); + Ok(()) } - - // Only insert if it doesn't exist - upstreams.insert(id.clone(), UpstreamConnection::new(address)); - println!("✅ Added new upstream {}", id); - Ok(()) -} // Maintain connections (called from Router) pub async fn initialize_connections(&self) { self.maintain_connections().await; } -} \ No newline at end of file +} diff --git a/src/share_accounter/mod.rs b/src/share_accounter/mod.rs index ecda99ae..08004da6 100644 --- a/src/share_accounter/mod.rs +++ b/src/share_accounter/mod.rs @@ -124,4 +124,4 @@ fn relay_down( } }); task.into() -} \ No newline at end of file +} diff --git a/src/translator/downstream/diff_management.rs b/src/translator/downstream/diff_management.rs index ea6c4e86..858f3afc 100644 --- a/src/translator/downstream/diff_management.rs +++ b/src/translator/downstream/diff_management.rs @@ -484,4 +484,4 @@ mod test { // TODO make a test where unknown donwstream is simulated and we do not wait for it to produce // a share but we try to updated the estimated hash power every 2 seconds and updated the // target consequentially this shuold start to provide shares within a normal amount of time -} \ No newline at end of file +} diff --git a/src/translator/downstream/notify.rs b/src/translator/downstream/notify.rs index 5a81cecd..673122cf 100644 --- a/src/translator/downstream/notify.rs +++ b/src/translator/downstream/notify.rs @@ -182,4 +182,4 @@ async fn start_update( TaskManager::add_update(task_manager, handle.into()) .await .map_err(|_| Error::TranslatorTaskManagerFailed) -} \ No newline at end of file +} diff --git a/src/translator/mod.rs b/src/translator/mod.rs index 04603ef7..b861b46f 100644 --- a/src/translator/mod.rs +++ b/src/translator/mod.rs @@ -207,4 +207,4 @@ pub async fn start( .map_err(|_| Error::TranslatorTaskManagerFailed)?; Ok(abortable) -} \ No newline at end of file +} From c476a18278b9ad8e520f2e782d0085bd50dae6f3 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Sat, 7 Jun 2025 19:13:54 +0530 Subject: [PATCH 16/21] refactor: Remove redundant logging and improve hashrate distribution logic in proxy and router --- src/main.rs | 190 +++++---------------------- src/proxy_state.rs | 2 +- src/router/mod.rs | 34 +---- src/router/multi_upstream_manager.rs | 175 +++++++++++++++++++----- 4 files changed, 188 insertions(+), 213 deletions(-) diff --git a/src/main.rs b/src/main.rs index c24c11f1..a7a658a4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -94,10 +94,6 @@ async fn main() { // Set downstream hashrate using Configuration pattern let downstream_hashrate = Configuration::downstream_hashrate(); ProxyState::set_downstream_hashrate(downstream_hashrate); - info!( - "Set downstream hashrate to: {}", - HashUnit::format_value(downstream_hashrate) - ); // Get pool addresses with auth keys using Configuration pattern let pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_addresses @@ -108,12 +104,6 @@ async fn main() { // Determine hashrate distribution using Configuration pattern let wants_distribution = Configuration::wants_hashrate_distribution(); - if wants_distribution { - info!("Hashrate distribution is enabled"); - } else { - info!("Hashrate distribution is disabled"); - } - // Create router based on configuration let mut router = if wants_distribution { // Multi-upstream with distribution @@ -164,19 +154,16 @@ async fn main() { initialize_proxy(&mut router, None, epsilon).await; } else if pool_address_keys.len() > 1 { info!("=== LATENCY-BASED SELECTION MODE ==="); - info!("Testing pool latencies and selecting best pool..."); // Test latency and select best pool let best_upstream = router.select_pool_connect().await; - if let Some(ref upstream) = best_upstream { - info!("Selected best upstream: {:?}", upstream); + if let Some(ref _upstream) = best_upstream { } else { error!("Failed to connect to any upstream pool"); std::process::exit(1); } - info!("Starting proxy with latency-based selection..."); initialize_proxy(&mut router, best_upstream, epsilon).await; } else { info!("=== SINGLE UPSTREAM MODE ==="); @@ -191,117 +178,31 @@ async fn main() { std::process::exit(1); } - info!("Starting proxy with single upstream..."); initialize_proxy(&mut router, best_upstream, epsilon).await; } - // This line should never be reached since initialize_proxy runs forever - info!("Proxy stopped unexpectedly"); - // This code should never be reached because initialize_proxy runs forever - info!("Proxy initialization complete"); - // Wait a moment for connections to establish tokio::time::sleep(Duration::from_secs(2)).await; - // Show immediate hashrate distribution with network verification - tokio::time::sleep(Duration::from_secs(2)).await; - - info!("🔍 VERIFYING NETWORK CONNECTIONS:"); - - // Check network connections - let output = std::process::Command::new("sh") - .arg("-c") - .arg(format!( - "ss -tn | grep -E \"({})\"", - pool_address_keys - .iter() - .map(|(addr, _)| addr.to_string()) - .collect::>() - .join("|") - )) - .output(); - - if let Ok(output) = output { - let connections = String::from_utf8_lossy(&output.stdout); - let pool_connections: Vec<&str> = connections - .lines() - .filter(|line| { - pool_address_keys - .iter() - .any(|(addr, _)| line.contains(&addr.to_string())) - }) - .collect(); - - info!("📡 Active pool connections: {}", pool_connections.len()); - for conn in pool_connections { - info!(" 🔗 {}", conn.trim()); - } - } - // Show hashrate distribution - let downstream_hashrate = ProxyState::get_downstream_hashrate(); let detailed_stats = router.get_detailed_connection_stats().await; - - info!("🚀 === INITIAL HASHRATE DISTRIBUTION ==="); - info!( - "🔋 Total configured downstream hashrate: {}", - HashUnit::format_value(downstream_hashrate) - ); - let active_count = detailed_stats .iter() .filter(|(_, is_active, _)| *is_active) .count(); if active_count > 0 { - // Check if using custom distribution - let mut is_custom_distribution = false; - for (_, is_active, percentage) in &detailed_stats { - if *is_active && (percentage - (100.0 / active_count as f32)).abs() > 1.0 { - is_custom_distribution = true; - break; - } - } - - if is_custom_distribution { - info!("🎯 Mode: Custom Distribution (from config)"); - } else { - let percentage_per_upstream = 100.0 / active_count as f32; - info!( - "🌐 Mode: Equal Distribution ({:.1}% per upstream)", - percentage_per_upstream - ); - } - - for (upstream_id, is_active, percentage) in detailed_stats { - if is_active { - // Calculate actual hashrate from percentage - let actual_hashrate = (percentage / 100.0) * downstream_hashrate; - info!( - " ✅ {}: allocated {} ({:.1}% of downstream)", - upstream_id, - HashUnit::format_value(actual_hashrate), - percentage - ); - } else { - info!( - " ❌ {}: 0.00T (INACTIVE - 0.0% of downstream)", - upstream_id - ); - } - } - info!("========================================"); + info!("Proxy initialized with {} active upstream(s)", active_count); // Start monitoring task for multi-upstream let router_clone = router.clone(); tokio::spawn(async move { - monitor_multi_upstream(router_clone).await; + monitor_multi_upstream(router_clone, epsilon).await; }); // Keep main thread alive loop { - tokio::time::sleep(Duration::from_secs(60)).await; - info!("Multi-upstream proxy running..."); + tokio::time::sleep(Duration::from_secs(300)).await; } } } @@ -314,7 +215,6 @@ async fn initialize_proxy( // Check if we're in multi-upstream mode if router.is_multi_upstream_enabled() { info!("Initializing proxy in HASHRATE DISTRIBUTION mode"); - info!("🎯 Focus: Distribute hashrate allocation across upstreams"); // Start API server for monitoring let stats_sender = api::stats::StatsSender::new(); @@ -323,10 +223,9 @@ async fn initialize_proxy( // Start the monitor task in the background let router_clone = router.clone(); tokio::spawn(async move { - monitor_multi_upstream(router_clone).await; + monitor_multi_upstream(router_clone, epsilon).await; }); - info!("✅ Hashrate distribution monitor started"); return; } @@ -518,65 +417,48 @@ async fn monitor( } // Consolidated monitoring function for multi-upstream mode -async fn monitor_multi_upstream(router: Router) { - let mut stats_report_counter = 0; +async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { + let mut report_counter = 0; + let mut last_stats: Option> = None; loop { - tokio::time::sleep(Duration::from_secs(10)).await; - stats_report_counter += 1; - - // Generate periodic reports every 30 seconds (3 * 10 seconds) - if stats_report_counter >= 3 { - stats_report_counter = 0; + tokio::time::sleep(Duration::from_secs(30)).await; - let (total_count, active_count) = router.get_upstream_counts().await; - let downstream_hashrate = ProxyState::get_downstream_hashrate(); - - info!("🚀 === MULTI-UPSTREAM HASHRATE REPORT ==="); - info!( - "📊 Total upstreams: {}, Active: {}", - total_count, active_count - ); - info!( - "🔋 Downstream hashrate: {}", - HashUnit::format_value(downstream_hashrate) - ); + let detailed_stats = router.get_detailed_connection_stats().await; + let active_count = detailed_stats + .iter() + .filter(|(_, active, _)| *active) + .count(); + + // Only log every 10 cycles (5 minutes) or if stats changed significantly + let stats_changed = match &last_stats { + Some(prev) => { + prev.len() != detailed_stats.len() + || prev.iter().zip(&detailed_stats).any( + |((_, prev_active, prev_pct), (_, active, pct))| { + prev_active != active || (prev_pct - pct).abs() > 1.0 + }, + ) + } + None => true, + }; + if stats_changed || report_counter >= 10 { if active_count > 0 { - let detailed_stats = router.get_detailed_connection_stats().await; - - for (upstream_id, is_active, percentage) in detailed_stats { - if is_active { - // Calculate actual hashrate from percentage - let allocated_hashrate = (percentage / 100.0) * downstream_hashrate; - info!( - " ✅ {}: allocated {} ({:.1}% of downstream)", - upstream_id, - HashUnit::format_value(allocated_hashrate), - percentage - ); - } else { - info!( - " ❌ {}: 0.00T (INACTIVE - 0.0% of downstream)", - upstream_id - ); + info!("Active pools: {}", active_count); + for (upstream_id, is_active, percentage) in &detailed_stats { + if *is_active && *percentage > 0.0 { + info!(" {} {:.1}%", upstream_id, percentage); } } - } else { - warn!("❌ No active upstream connections!"); } - info!("========================================="); - } - - // Check proxy state - let (is_down, error_msg) = ProxyState::is_proxy_down(); - if is_down { - error!("⚠️ Proxy state is DOWN: {:?}", error_msg); - break; + report_counter = 0; + last_stats = Some(detailed_stats); + } else { + report_counter += 1; } } } - pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream diff --git a/src/proxy_state.rs b/src/proxy_state.rs index 9677acac..ade632d7 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -265,7 +265,7 @@ impl ProxyState { /// Set the downstream hashrate to be distributed among upstreams pub fn set_downstream_hashrate(hashrate: f32) { - info!("Setting total hashrate to: {} h/s", hashrate); + info!("Setting downstream hashrate to: {} h/s", hashrate); if PROXY_STATE .safe_lock(|state| { state.total_hashrate = hashrate; diff --git a/src/router/mod.rs b/src/router/mod.rs index ffb45767..da1254c4 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -17,7 +17,7 @@ use tokio::{ watch, }, }; -use tracing::{error, info}; +use tracing::{debug, error, info}; use crate::{ minin_pool_connection::{self, get_mining_setup_connection_msg, mining_setup_connection}, @@ -72,7 +72,7 @@ impl Router { let auth_keys = vec![auth_pub_k; pool_addresses.len()]; Self { - pool_addresses: pool_addresses, + pool_addresses, keys: auth_keys, current_pool: None, upstream_manager: None, @@ -402,7 +402,7 @@ impl Router { /// Initialize upstream connections for the manager pub async fn initialize_upstream_connections(&mut self) -> Result<(), String> { if let Some(ref manager) = self.upstream_manager { - info!( + debug!( "Initializing {} upstream connections", self.pool_addresses.len() ); @@ -411,7 +411,7 @@ impl Router { for (idx, (addr, key)) in self.pool_addresses.iter().zip(self.keys.iter()).enumerate() { let id = format!("upstream-{}", idx); - info!("Adding upstream {}: {} ({})", id, addr, key); + debug!("Adding upstream {}: {} ({})", id, addr, key); if let Err(e) = manager .add_upstream( @@ -431,41 +431,21 @@ impl Router { // This ensures add_upstream() calls are complete and won't overwrite distribution manager.initialize_connections().await; - info!("Upstream connections initialized - ready for hashrate distribution"); Ok(()) } else { Err("No upstream manager available".to_string()) } } - /// Get count of total and active upstreams - pub async fn get_upstream_counts(&self) -> (usize, usize) { - if let Some(ref manager) = self.upstream_manager { - let upstreams = manager.get_upstreams().await; - let total = upstreams.len(); - let active = upstreams.values().filter(|u| u.is_active).count(); - (total, active) - } else { - (self.pool_addresses.len(), 1) // Single upstream mode - } - } /// Sets the hashrate distribution for the upstream manager. pub async fn set_hashrate_distribution( - &self, + &mut self, distribution: Vec, ) -> Result<(), &'static str> { - info!( - "🔧 Router::set_hashrate_distribution called with: {:?}", - distribution - ); - - if let Some(ref manager) = self.upstream_manager { - info!("🔧 Calling manager.set_hashrate_distribution"); - manager.set_hashrate_distribution(distribution).await; - info!("🔧 Manager.set_hashrate_distribution returned"); + if let Some(ref mut manager) = self.upstream_manager { + let _ = manager.set_hashrate_distribution(distribution).await; Ok(()) } else { - error!("❌ No upstream manager available"); Err("No upstream manager available") } } diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 402a2ae0..e60890cb 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -2,7 +2,7 @@ use crate::{minin_pool_connection, proxy_state::ProxyState, HashUnit}; use key_utils::Secp256k1PublicKey; use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; use tokio::sync::Mutex; -use tracing::{error, info}; +use tracing::{debug, error, info, warn}; #[derive(Clone)] pub struct MultiUpstreamManager { @@ -15,6 +15,7 @@ pub struct MultiUpstreamManager { #[derive(Clone)] pub struct UpstreamConnection { pub address: SocketAddr, + #[allow(dead_code)] pub allocated_hashrate: f64, pub allocated_percentage: f32, pub is_active: bool, @@ -29,7 +30,7 @@ impl UpstreamConnection { is_active: false, } } - + #[allow(dead_code)] // This method is kept for future use fn update_allocation(&mut self, percentage: f32, total_hashrate: f64) { self.allocated_percentage = percentage; self.allocated_hashrate = total_hashrate * percentage as f64 / 100.0; @@ -75,7 +76,6 @@ impl MultiUpstreamManager { tokio::spawn(async move { loop { - info!("Connecting to upstream {}: {}", id, address); match minin_pool_connection::connect_pool( address, auth_pub_k, @@ -85,7 +85,6 @@ impl MultiUpstreamManager { .await { Ok((_send, _recv, _abortable)) => { - info!("✅ Connected to upstream {} - maintaining connection", id); Self::update_upstream_status(&upstreams, &id, true).await; // Keep the connection alive (simulate mining) tokio::time::sleep(Duration::from_secs(30)).await; @@ -113,35 +112,122 @@ impl MultiUpstreamManager { } } - /// Set hashrate distribution for each upstream (percentages) - pub async fn set_hashrate_distribution(&self, distribution: Vec) { - println!( - "calling set_hashrate_distribution in multi manager with: {:?}", - distribution - ); + pub async fn set_hashrate_distribution( + &mut self, + mut distribution: Vec, + ) -> Result<(), String> { + let upstream_count = { + let upstreams = self.upstreams.lock().await; + upstreams.len() + }; + + // Pad with zeros if needed + if distribution.len() < upstream_count { + let zeros_needed = upstream_count - distribution.len(); + distribution.extend(vec![0.0; zeros_needed]); + debug!("Padded distribution with {} zeros", zeros_needed); + } else if distribution.len() > upstream_count { + return Err(format!( + "Distribution array has {} values but only {} pools available", + distribution.len(), + upstream_count + )); + } + + // Sort distribution highest to lowest + distribution.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal)); + // Get pools ranked by latency (best first) + let ranked_pools = self.rank_by_latency().await; + + // Normalize percentages + let total: f32 = distribution.iter().sum(); + if (total - 100.0).abs() > 0.1 { + warn!("Distribution sum: {:.1}%, normalizing to 100%", total); + if total > 0.0 { + for percentage in &mut distribution { + *percentage = (*percentage / total) * 100.0; + } + } else { + return Err("All distribution percentages are zero or negative".to_string()); + } + } + + // Apply distribution to ranked pools let mut upstreams = self.upstreams.lock().await; - let total_hashrate = ProxyState::get_downstream_hashrate() as f64; - println!( - "Total hashrate: {}", - HashUnit::format_value(total_hashrate as f32) - ); + // Reset all to 0% + for upstream in upstreams.values_mut() { + upstream.allocated_percentage = 0.0; + } + + // Assign percentages to best latency pools first + let downstream_hashrate = ProxyState::get_downstream_hashrate(); - for (i, (_id, conn)) in upstreams.iter_mut().enumerate() { + for (i, (pool_id, _address, latency)) in ranked_pools.iter().enumerate() { if let Some(&percentage) = distribution.get(i) { - println!("setting upstream {} to {}%", i, percentage); - conn.update_allocation(percentage, total_hashrate); - info!( - "Upstream {}: allocated {} ({:.1}%)", - conn.address, - HashUnit::format_value(conn.allocated_hashrate as f32), - percentage - ); + if let Some(upstream) = upstreams.get_mut(pool_id) { + upstream.allocated_percentage = percentage; + + if percentage > 0.0 { + let allocated_hashrate = (percentage / 100.0) * downstream_hashrate; + info!( + "{}: {} ({:.1}%) - {}ms", + pool_id, + HashUnit::format_value(allocated_hashrate), + percentage, + latency.as_millis() + ); + } + } } } + + Ok(()) } + // Add this method to test latency + async fn test_pool_latency(address: &SocketAddr) -> Result { + let start = std::time::Instant::now(); + + match tokio::time::timeout( + Duration::from_secs(2), + tokio::net::TcpStream::connect(address), + ) + .await + { + Ok(Ok(_)) => Ok(start.elapsed()), + Ok(Err(e)) => Err(format!("Connection failed: {}", e)), + Err(_) => Err("Timeout".to_string()), + } + } + + // Add this method to rank pools by latency + pub async fn rank_by_latency(&self) -> Vec<(String, SocketAddr, Duration)> { + let upstreams = self.upstreams.lock().await; + let mut pool_latencies = Vec::new(); + + info!("Testing latency for {} pools", upstreams.len()); + + for (id, conn) in upstreams.iter() { + match Self::test_pool_latency(&conn.address).await { + Ok(latency) => { + debug!("{}: {}ms", id, latency.as_millis()); + pool_latencies.push((id.clone(), conn.address, latency)); + } + Err(e) => { + warn!("{}: latency test failed - {}", id, e); + pool_latencies.push((id.clone(), conn.address, Duration::from_secs(999))); + } + } + } + + // Sort by latency (lowest first) + pool_latencies.sort_by_key(|(_, _, latency)| *latency); + + pool_latencies + } + #[allow(dead_code)] // This method is kept for future use pub async fn get_upstreams(&self) -> HashMap { self.upstreams.lock().await.clone() } @@ -150,11 +236,7 @@ impl MultiUpstreamManager { upstreams .iter() - .map(|(id, conn)| { - println!("Upstream {}: {}% allocation", id, conn.allocated_percentage); - // Return the percentage, not the hashrate - (id.clone(), conn.is_active, conn.allocated_percentage) - }) + .map(|(id, conn)| (id.clone(), conn.is_active, conn.allocated_percentage)) .collect() } @@ -173,13 +255,11 @@ impl MultiUpstreamManager { // Check if upstream already exists - if so, don't overwrite it if upstreams.contains_key(&id) { - println!("⚠️ Upstream {} already exists, not overwriting", id); return Ok(()); } // Only insert if it doesn't exist upstreams.insert(id.clone(), UpstreamConnection::new(address)); - println!("✅ Added new upstream {}", id); Ok(()) } @@ -187,4 +267,37 @@ impl MultiUpstreamManager { pub async fn initialize_connections(&self) { self.maintain_connections().await; } + #[allow(dead_code)] // This method is kept for future use + pub async fn handle_upstream_failure(&mut self, upstream_id: &str) { + { + let mut upstreams = self.upstreams.lock().await; + + if let Some(upstream) = upstreams.get_mut(upstream_id) { + upstream.is_active = false; + warn!( + "Upstream {} disconnected, attempting reconnection", + upstream_id + ); + } + } // Drop the guard here + + // Redistribute hashrate among remaining active upstreams + self.rebalance_on_failure().await; + } + #[allow(dead_code)] // This method is kept for future use + async fn rebalance_on_failure(&mut self) { + let upstreams = self.upstreams.lock().await; + let active_count = upstreams.values().filter(|u| u.is_active).count(); + + if active_count == 0 { + error!("All upstreams disconnected - proxy will retry connections"); + return; + } + + info!( + "Rebalancing hashrate across {} active upstreams", + active_count + ); + // Implement rebalancing logic here + } } From dd564270c93b96eceeeab8ad8a12b987639c8630 Mon Sep 17 00:00:00 2001 From: Ayush Bansal <116891199+Bansalayush247@users.noreply.github.com> Date: Sat, 7 Jun 2025 19:22:31 +0530 Subject: [PATCH 17/21] Delete config.toml --- config.toml | 12 ------------ 1 file changed, 12 deletions(-) delete mode 100644 config.toml diff --git a/config.toml b/config.toml deleted file mode 100644 index 43a25367..00000000 --- a/config.toml +++ /dev/null @@ -1,12 +0,0 @@ -token="UR01TMkZv6Vs5zbwr0t6" -test_pool_addresses=["3.74.36.119:2000","18.193.252.132:2000"] -#tp_address = "127.0.0.1:8442" -interval = 120_000 -delay = 540 -downstream_hashrate = "500T" -loglevel = "info" -nc_loglevel = "off" -test = true - -# Uncomment this line to enable hashrate distribution -hashrate_distribution = [70.0, 30.0] \ No newline at end of file From 4007acad83be7ff9a379b5b3d0fec914b957b218 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Sat, 7 Jun 2025 20:24:38 +0530 Subject: [PATCH 18/21] refactor: clean up configuration structures in code --- config.toml | 12 --------- cpuminer-opt | 1 + src/config.rs | 62 ++++++++++++++++++++++----------------------- src/main.rs | 8 ++++-- src/router/mod.rs | 28 +++++++++++--------- src/shared/utils.rs | 7 ----- 6 files changed, 54 insertions(+), 64 deletions(-) delete mode 100644 config.toml create mode 160000 cpuminer-opt diff --git a/config.toml b/config.toml deleted file mode 100644 index 43a25367..00000000 --- a/config.toml +++ /dev/null @@ -1,12 +0,0 @@ -token="UR01TMkZv6Vs5zbwr0t6" -test_pool_addresses=["3.74.36.119:2000","18.193.252.132:2000"] -#tp_address = "127.0.0.1:8442" -interval = 120_000 -delay = 540 -downstream_hashrate = "500T" -loglevel = "info" -nc_loglevel = "off" -test = true - -# Uncomment this line to enable hashrate distribution -hashrate_distribution = [70.0, 30.0] \ No newline at end of file diff --git a/cpuminer-opt b/cpuminer-opt new file mode 160000 index 00000000..dd99580a --- /dev/null +++ b/cpuminer-opt @@ -0,0 +1 @@ +Subproject commit dd99580a4c05414562ccc1bffca7cdce2073fa29 diff --git a/src/config.rs b/src/config.rs index 0c112f55..ba520cfa 100644 --- a/src/config.rs +++ b/src/config.rs @@ -13,53 +13,53 @@ lazy_static! { pub static ref CONFIG: Configuration = Configuration::load_config(); } -#[derive(Parser)] -pub struct Args { +#[derive(Parser)] +struct Args { #[clap(long)] - pub test: bool, + test: bool, #[clap(long = "d", short = 'd', value_parser = parse_hashrate)] - pub downstream_hashrate: Option, + downstream_hashrate: Option, #[clap(long = "loglevel", short = 'l')] - pub loglevel: Option, + loglevel: Option, #[clap(long = "nc", short = 'n')] - pub noise_connection_log: Option, + noise_connection_log: Option, #[clap(long = "delay")] - pub delay: Option, + delay: Option, #[clap(long = "interval", short = 'i')] - pub adjustment_interval: Option, + adjustment_interval: Option, #[clap(long = "pool", short = 'p', value_delimiter = ',')] - pub pool_addresses: Option>, + pool_addresses: Option>, #[clap(long = "test-pool", value_delimiter = ',')] - pub test_pool_addresses: Option>, + test_pool_addresses: Option>, #[clap(long)] - pub token: Option, + token: Option, #[clap(long)] - pub tp_address: Option, + tp_address: Option, #[clap(long)] - pub listening_addr: Option, + listening_addr: Option, #[clap(long = "config", short = 'c')] - pub config_file: Option, + config_file: Option, #[clap(long = "api-server-port", short = 's')] - pub api_server_port: Option, + api_server_port: Option, #[clap(long = "hashrate-distribution", value_delimiter = ',')] - pub hashrate_distribution: Option>, + hashrate_distribution: Option>, } #[derive(Serialize, Deserialize)] -pub struct ConfigFile { - pub token: Option, - pub tp_address: Option, - pub pool_addresses: Option>, - pub test_pool_addresses: Option>, - pub interval: Option, - pub delay: Option, - pub downstream_hashrate: Option, - pub loglevel: Option, - pub nc_loglevel: Option, - pub test: Option, - pub listening_addr: Option, - pub api_server_port: Option, - pub hashrate_distribution: Option>, +struct ConfigFile { + tp_address: Option, + token: Option, + pool_addresses: Option>, + test_pool_addresses: Option>, + interval: Option, + delay: Option, + downstream_hashrate: Option, + loglevel: Option, + nc_loglevel: Option, + test: Option, + listening_addr: Option, + api_server_port: Option, + hashrate_distribution: Option>, } pub struct Configuration { @@ -155,7 +155,7 @@ impl Configuration { // Loads config from CLI, file, or env vars with precedence: CLI > file > env. fn load_config() -> Self { let args = Args::parse(); - let config_path: PathBuf = args.config_file.clone().unwrap_or("config.toml".into()); + let config_path: PathBuf = args.config_file.unwrap_or("config.toml".into()); let config: ConfigFile = std::fs::read_to_string(&config_path) .ok() .and_then(|content| toml::from_str(&content).ok()) diff --git a/src/main.rs b/src/main.rs index a7a658a4..80c902c3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,7 +2,6 @@ use jemallocator::Jemalloc; use router::Router; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; - #[cfg(not(target_os = "windows"))] #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; @@ -15,8 +14,8 @@ use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; use std::{net::SocketAddr, time::Duration}; use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; - mod api; + mod config; mod ingress; pub mod jd_client; @@ -61,6 +60,7 @@ lazy_static! { #[tokio::main] async fn main() { let log_level = Configuration::loglevel(); + let noise_connection_log_level = Configuration::nc_loglevel(); //Disable noise_connection error (for now) because: @@ -231,6 +231,7 @@ async fn initialize_proxy( // Single upstream mode only loop { + // Initial setup for the proxy let stats_sender = api::stats::StatsSender::new(); let (send_to_pool, recv_from_pool, pool_connection_abortable) = @@ -245,6 +246,7 @@ async fn initialize_proxy( tokio::time::sleep(Duration::from_secs(1)).await; secs -= 1; } + // Restart loop, esentially restarting proxy continue; } }; @@ -459,10 +461,12 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { } } } + pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream } + enum HashUnit { Tera, Peta, diff --git a/src/router/mod.rs b/src/router/mod.rs index da1254c4..3c0eb0b2 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -31,16 +31,16 @@ pub use multi_upstream_manager::MultiUpstreamManager; /// Router handles connection to Multiple upstreams. pub struct Router { - pub pool_addresses: Vec, - pub keys: Vec, + pool_addresses: Vec, + keys: Vec, pub current_pool: Option, - pub upstream_manager: Option, + upstream_manager: Option, // Keep these fields for backward compatibility with single upstream mode - pub auth_pub_k: Secp256k1PublicKey, - pub setup_connection_msg: Option>, - pub timer: Option, - pub latency_tx: tokio::sync::watch::Sender>, + auth_pub_k: Secp256k1PublicKey, + setup_connection_msg: Option>, + timer: Option, + latency_tx: tokio::sync::watch::Sender>, pub latency_rx: tokio::sync::watch::Receiver>, } impl Clone for Router { @@ -65,15 +65,19 @@ impl Router { pub fn new( pool_addresses: Vec, auth_pub_k: Secp256k1PublicKey, + // Configuration msg used to setup connection between client and pool + // If not, present `get_mining_setup_connection_msg()` is called to generated default values setup_connection_msg: Option>, + // Max duration for pool setup after which it times out. + // If None, default time of 5s is used. timer: Option, ) -> Self { let (latency_tx, latency_rx) = watch::channel(None); - let auth_keys = vec![auth_pub_k; pool_addresses.len()]; + let auth_pub_keys = vec![auth_pub_k; pool_addresses.len()]; Self { pool_addresses, - keys: auth_keys, + keys: auth_pub_keys, current_pool: None, upstream_manager: None, auth_pub_k, @@ -267,7 +271,7 @@ impl Router { Some(addr) => addr, None => match self.select_pool_connect().await { Some(addr) => addr, - // Called when we initialize + // Called when we initialize the proxy, without a valid pool we can not start mine and we // return Err None => { return Err(minin_pool_connection::errors::Error::Unrecoverable); @@ -279,7 +283,7 @@ impl Router { info!("Upstream {:?} selected", pool); // Find the matching auth key for this address - fix field name - let auth_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { + let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { self.keys[index] } else { self.auth_pub_k @@ -287,7 +291,7 @@ impl Router { match minin_pool_connection::connect_pool( pool, - auth_key, + auth_pub_key, self.setup_connection_msg.clone(), self.timer, ) diff --git a/src/shared/utils.rs b/src/shared/utils.rs index c3bc7397..476cbbf0 100644 --- a/src/shared/utils.rs +++ b/src/shared/utils.rs @@ -8,13 +8,6 @@ use tokio::task::JoinHandle; pub struct AbortOnDrop { abort_handle: AbortHandle, } -impl Clone for AbortOnDrop { - fn clone(&self) -> Self { - Self { - abort_handle: self.abort_handle.clone(), - } - } -} impl AbortOnDrop { pub fn new(handle: JoinHandle) -> Self { From 79edbb260fdcd8dd002679a0e3f81aab9e1f0bfd Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Sun, 8 Jun 2025 21:37:00 +0530 Subject: [PATCH 19/21] feat: Add multi-upstream hashrate distribution support - Implement MultiUpstreamManager for load balancing across multiple pools - Add hashrate distribution configuration with percentage-based routing - Support both latency-based pool ranking and custom distribution - Add JDC bridge for Template Provider integration - Clean up verbose logging and remove unused code - Fix all clippy warnings and optimize performance - Add comprehensive API endpoints for monitoring pool stats - Support graceful failover and connection monitoring Features: - Multi-pool hashrate distribution (e.g., 30%/70% split) - Automatic latency-based pool ranking - Share routing based on configured percentages - Pool health monitoring and reconnection - API server for real-time statistics - Compatible with existing single-upstream mode --- src/config.rs | 2 +- src/main.rs | 227 +++++++++++++-------- src/router/mod.rs | 39 +++- src/router/multi_upstream_manager.rs | 295 +++++++++++++++++++++++++-- 4 files changed, 460 insertions(+), 103 deletions(-) diff --git a/src/config.rs b/src/config.rs index ba520cfa..882f6823 100644 --- a/src/config.rs +++ b/src/config.rs @@ -13,7 +13,7 @@ lazy_static! { pub static ref CONFIG: Configuration = Configuration::load_config(); } -#[derive(Parser)] +#[derive(Parser)] struct Args { #[clap(long)] test: bool, diff --git a/src/main.rs b/src/main.rs index 80c902c3..58ce2586 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,7 +10,7 @@ use crate::shared::utils::AbortOnDrop; use config::Configuration; use key_utils::Secp256k1PublicKey; use lazy_static::lazy_static; -use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; +use proxy_state::{PoolState, ProxyState, TpState, TranslatorState}; // Add ProxyStates use std::{net::SocketAddr, time::Duration}; use tokio::sync::mpsc::channel; use tracing::{error, info, warn}; @@ -83,17 +83,10 @@ async fn main() { let pool_addresses = Configuration::pool_address() .filter(|p| !p.is_empty()) - .unwrap_or_else(|| { - if Configuration::test() { - panic!("Test pool address is missing"); - } else { - panic!("Pool address is missing"); - } - }); + .unwrap_or_else(|| panic!("Pool address is missing")); // Set downstream hashrate using Configuration pattern - let downstream_hashrate = Configuration::downstream_hashrate(); - ProxyState::set_downstream_hashrate(downstream_hashrate); + ProxyState::set_downstream_hashrate(Configuration::downstream_hashrate()); // Get pool addresses with auth keys using Configuration pattern let pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_addresses @@ -127,8 +120,6 @@ async fn main() { // Handle the three different scenarios if wants_distribution { - info!("=== HASHRATE DISTRIBUTION MODE ==="); - // Get the distribution from config let distribution = Configuration::hashrate_distribution().unwrap_or_else(|| { let count = pool_address_keys.len(); @@ -150,11 +141,10 @@ async fn main() { std::process::exit(1); } - info!("Starting proxy with hashrate distribution..."); + // Wait for connections to establish + tokio::time::sleep(Duration::from_secs(3)).await; initialize_proxy(&mut router, None, epsilon).await; } else if pool_address_keys.len() > 1 { - info!("=== LATENCY-BASED SELECTION MODE ==="); - // Test latency and select best pool let best_upstream = router.select_pool_connect().await; @@ -166,45 +156,13 @@ async fn main() { initialize_proxy(&mut router, best_upstream, epsilon).await; } else { - info!("=== SINGLE UPSTREAM MODE ==="); - - // Connect to single pool let best_upstream = router.select_pool_connect().await; - - if let Some(ref upstream) = best_upstream { - info!("Connected to single upstream: {:?}", upstream); - } else { + if best_upstream.is_none() { error!("Failed to connect to upstream pool"); std::process::exit(1); } - initialize_proxy(&mut router, best_upstream, epsilon).await; } - - // Wait a moment for connections to establish - tokio::time::sleep(Duration::from_secs(2)).await; - - // Show hashrate distribution - let detailed_stats = router.get_detailed_connection_stats().await; - let active_count = detailed_stats - .iter() - .filter(|(_, is_active, _)| *is_active) - .count(); - - if active_count > 0 { - info!("Proxy initialized with {} active upstream(s)", active_count); - - // Start monitoring task for multi-upstream - let router_clone = router.clone(); - tokio::spawn(async move { - monitor_multi_upstream(router_clone, epsilon).await; - }); - - // Keep main thread alive - loop { - tokio::time::sleep(Duration::from_secs(300)).await; - } - } } async fn initialize_proxy( @@ -214,18 +172,105 @@ async fn initialize_proxy( ) { // Check if we're in multi-upstream mode if router.is_multi_upstream_enabled() { - info!("Initializing proxy in HASHRATE DISTRIBUTION mode"); - - // Start API server for monitoring + // Initialize all the same components as single upstream mode let stats_sender = api::stats::StatsSender::new(); - let _server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); - // Start the monitor task in the background - let router_clone = router.clone(); - tokio::spawn(async move { - monitor_multi_upstream(router_clone, epsilon).await; - }); + // Start SV1 ingress to handle downstream (mining device) connections + let (downs_sv1_tx, downs_sv1_rx) = channel(10); + let sv1_ingress_abortable = ingress::sv1_ingress::start_listen_for_downstream(downs_sv1_tx); + + // Start translator to handle SV1 <-> SV2 translation + let (translator_up_tx, mut translator_up_rx) = channel(10); + let translator_abortable = + match translator::start(downs_sv1_rx, translator_up_tx, stats_sender.clone()).await { + Ok(abortable) => abortable, + Err(e) => { + error!("Impossible to initialize translator: {e}"); + ProxyState::update_translator_state(TranslatorState::Down); + ProxyState::update_tp_state(TpState::Down); + return; + } + }; + + // Get translator channels + let (jdc_to_translator_sender, jdc_from_translator_receiver, _) = translator_up_rx + .recv() + .await + .expect("Translator failed before initialization"); + + // Setup JDC channels with correct types + let (from_jdc_to_share_accounter_send, from_jdc_to_share_accounter_recv) = + channel::>(10); + let (from_share_accounter_to_jdc_send, from_share_accounter_to_jdc_recv) = + channel::>(10); + + // Check if TP is available first + let tp_available = TP_ADDRESS.safe_lock(|tp| tp.clone()).ok().flatten(); + + let (share_accounter_abortable, jdc_abortable) = if let Some(_tp_addr) = tp_available { + // WITH TP: Start JDC first, then share accounting with JDC channels + let jdc_handle = jd_client::start( + jdc_from_translator_receiver, + jdc_to_translator_sender, + from_share_accounter_to_jdc_recv, + from_jdc_to_share_accounter_send, + ) + .await; + if jdc_handle.is_some() { + // JDC started successfully, use JDC channels + match router + .start_multi_upstream_share_accounting_with_jdc( + from_jdc_to_share_accounter_recv, + from_share_accounter_to_jdc_send, + ) + .await + { + Ok(handle) => (handle, jdc_handle), + Err(e) => { + error!( + "Failed to start multi-upstream share accounting with JDC: {}", + e + ); + return; + } + } + } else { + error!("Failed to start JDC with TP"); + return; + } + } else { + // WITHOUT TP: Use translator channels directly + match router + .start_multi_upstream_share_accounting( + jdc_from_translator_receiver, + jdc_to_translator_sender, + ) + .await + { + Ok(handle) => (handle, None), + Err(e) => { + error!("Failed to start multi-upstream share accounting: {}", e); + return; + } + } + }; + + // Start API server + let server_handle = tokio::spawn(api::start(router.clone(), stats_sender)); + + // Collect abort handles for monitoring + let mut abort_handles = vec![ + (sv1_ingress_abortable, "sv1_ingress".to_string()), + (translator_abortable, "translator".to_string()), + (share_accounter_abortable, "share_accounter".to_string()), + ]; + if let Some(jdc_handle) = jdc_abortable { + abort_handles.push((jdc_handle, "jdc".to_string())); + } + + // Use combined monitoring function + monitor_multi_upstream(router.clone(), abort_handles, server_handle, epsilon).await; return; } @@ -267,13 +312,14 @@ async fn initialize_proxy( } }; - let (from_jdc_to_share_accounter_send, from_jdc_to_share_accounter_recv) = channel(10); - let (from_share_accounter_to_jdc_send, from_share_accounter_to_jdc_recv) = channel(10); + // This creates the channels but translator has no upstream to connect to let (jdc_to_translator_sender, jdc_from_translator_receiver, _) = translator_up_rx .recv() .await .expect("Translator failed before initialization"); + let (from_jdc_to_share_accounter_send, from_jdc_to_share_accounter_recv) = channel(10); + let (from_share_accounter_to_jdc_send, from_share_accounter_to_jdc_recv) = channel(10); let jdc_abortable: Option; let share_accounter_abortable; let tp = match TP_ADDRESS.safe_lock(|tp| tp.clone()) { @@ -418,34 +464,55 @@ async fn monitor( } } -// Consolidated monitoring function for multi-upstream mode -async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { +// Combined monitoring function for multi-upstream mode +async fn monitor_multi_upstream( + router: Router, + abort_handles: Vec<(AbortOnDrop, String)>, + server_handle: tokio::task::JoinHandle<()>, + _epsilon: Duration, +) { let mut report_counter = 0; - let mut last_stats: Option> = None; loop { tokio::time::sleep(Duration::from_secs(30)).await; - let detailed_stats = router.get_detailed_connection_stats().await; - let active_count = detailed_stats + // Monitor finished tasks (critical components) + if let Some((_handle, name)) = abort_handles .iter() - .filter(|(_, active, _)| *active) - .count(); - - // Only log every 10 cycles (5 minutes) or if stats changed significantly - let stats_changed = match &last_stats { - Some(prev) => { - prev.len() != detailed_stats.len() - || prev.iter().zip(&detailed_stats).any( - |((_, prev_active, prev_pct), (_, active, pct))| { - prev_active != active || (prev_pct - pct).abs() > 1.0 - }, - ) + .find(|(handle, _name)| handle.is_finished()) + { + if name != "jdc" { + error!("Critical task {:?} failed, restarting", name); + for (handle, _name) in abort_handles { + drop(handle); + } + server_handle.abort(); + return; } - None => true, - }; + } + + // Check proxy state - but be more lenient for multi-upstream mode + let is_proxy_down = ProxyState::is_proxy_down(); + if is_proxy_down.0 { + if let Some(ref status) = is_proxy_down.1 { + let status_str = format!("{:?}", status); + if !status_str.contains("Tp(Down)") && !status_str.contains("InternalInconsistency") + { + error!("{:?} is DOWN, restarting", status); + drop(abort_handles); + server_handle.abort(); + return; + } + } + } + + if report_counter >= 10 { + let detailed_stats = router.get_detailed_connection_stats().await; + let active_count = detailed_stats + .iter() + .filter(|(_, active, _)| *active) + .count(); - if stats_changed || report_counter >= 10 { if active_count > 0 { info!("Active pools: {}", active_count); for (upstream_id, is_active, percentage) in &detailed_stats { @@ -455,13 +522,11 @@ async fn monitor_multi_upstream(router: Router, _epsilon: Duration) { } } report_counter = 0; - last_stats = Some(detailed_stats); } else { report_counter += 1; } } } - pub enum Reconnect { NewUpstream(std::net::SocketAddr), // Reconnecting with a new upstream NoUpstream, // Reconnecting without upstream diff --git a/src/router/mod.rs b/src/router/mod.rs index 3c0eb0b2..848b4c5c 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -65,7 +65,7 @@ impl Router { pub fn new( pool_addresses: Vec, auth_pub_k: Secp256k1PublicKey, - // Configuration msg used to setup connection between client and pool + // Configuration msg used to setup connection between client and pool // If not, present `get_mining_setup_connection_msg()` is called to generated default values setup_connection_msg: Option>, // Max duration for pool setup after which it times out. @@ -283,7 +283,8 @@ impl Router { info!("Upstream {:?} selected", pool); // Find the matching auth key for this address - fix field name - let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { + let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) + { self.keys[index] } else { self.auth_pub_k @@ -337,6 +338,23 @@ impl Router { } } + /// Start multi-upstream share accounting + pub async fn start_multi_upstream_share_accounting( + &self, + from_translator_recv: tokio::sync::mpsc::Receiver< + roles_logic_sv2::parsers::Mining<'static>, + >, + to_translator_send: tokio::sync::mpsc::Sender>, + ) -> Result> { + if let Some(ref manager) = self.upstream_manager { + manager + .start_multi_upstream_share_accounting(from_translator_recv, to_translator_send) + .await + } else { + Err("Multi-upstream manager not initialized".into()) + } + } + /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { // Find the auth key for this address - fix field names @@ -453,6 +471,23 @@ impl Router { Err("No upstream manager available") } } + pub async fn start_multi_upstream_share_accounting_with_jdc( + &self, + from_jdc_recv: Receiver>, + to_jdc_send: Sender>, + ) -> Result> { + if let Some(ref manager) = self.upstream_manager { + // Direct pass-through since types already match + let manager_handle = manager + .start_multi_upstream_share_accounting(from_jdc_recv, to_jdc_send) + .await?; + + // No type conversion needed - just return the handle + Ok(manager_handle) + } else { + Err("Multi-upstream manager not initialized".into()) + } + } } /// Track latencies for various stages of pool connection setup. diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index e60890cb..230e250a 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -1,9 +1,9 @@ use crate::{minin_pool_connection, proxy_state::ProxyState, HashUnit}; use key_utils::Secp256k1PublicKey; +use rand; use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; use tokio::sync::Mutex; use tracing::{debug, error, info, warn}; - #[derive(Clone)] pub struct MultiUpstreamManager { upstreams: Arc>>, @@ -12,31 +12,31 @@ pub struct MultiUpstreamManager { timer: Option, } -#[derive(Clone)] +#[derive(Debug)] pub struct UpstreamConnection { pub address: SocketAddr, - #[allow(dead_code)] - pub allocated_hashrate: f64, - pub allocated_percentage: f32, pub is_active: bool, + pub allocated_percentage: f32, + // Add connection channels if needed + pub sender: Option< + tokio::sync::mpsc::Sender>, + >, + pub receiver: Option< + tokio::sync::mpsc::Receiver>, + >, } impl UpstreamConnection { - fn new(address: SocketAddr) -> Self { + pub fn new(address: SocketAddr) -> Self { Self { address, - allocated_hashrate: 0.0, - allocated_percentage: 0.0, is_active: false, + allocated_percentage: 0.0, + sender: None, + receiver: None, } } - #[allow(dead_code)] // This method is kept for future use - fn update_allocation(&mut self, percentage: f32, total_hashrate: f64) { - self.allocated_percentage = percentage; - self.allocated_hashrate = total_hashrate * percentage as f64 / 100.0; - } } - impl MultiUpstreamManager { pub fn new( upstreams: Vec, @@ -76,6 +76,7 @@ impl MultiUpstreamManager { tokio::spawn(async move { loop { + info!("Attempting to connect to upstream: {} ({})", id, address); match minin_pool_connection::connect_pool( address, auth_pub_k, @@ -84,10 +85,24 @@ impl MultiUpstreamManager { ) .await { - Ok((_send, _recv, _abortable)) => { - Self::update_upstream_status(&upstreams, &id, true).await; - // Keep the connection alive (simulate mining) - tokio::time::sleep(Duration::from_secs(30)).await; + Ok((send, recv, _abortable)) => { + info!("Successfully connected to upstream: {} ({})", id, address); + + // Store both sender and receiver + { + let mut upstreams_guard = upstreams.lock().await; + if let Some(upstream) = upstreams_guard.get_mut(&id) { + upstream.sender = Some(send); + upstream.receiver = Some(recv); // Store receiver too! + upstream.is_active = true; + info!("Stored sender and receiver for upstream: {}", id); + } + } + + // Keep connection alive + loop { + tokio::time::sleep(Duration::from_secs(30)).await; + } } Err(e) => { error!("Failed to connect to upstream {}: {:?}", id, e); @@ -98,6 +113,7 @@ impl MultiUpstreamManager { } }); } + drop(upstreams_guard); } /// Helper method to update upstream connection status @@ -229,7 +245,22 @@ impl MultiUpstreamManager { } #[allow(dead_code)] // This method is kept for future use pub async fn get_upstreams(&self) -> HashMap { - self.upstreams.lock().await.clone() + let upstreams = self.upstreams.lock().await; + upstreams + .iter() + .map(|(k, v)| { + ( + k.clone(), + UpstreamConnection { + address: v.address, + is_active: v.is_active, + allocated_percentage: v.allocated_percentage, + sender: v.sender.clone(), + receiver: None, // receiver can't be cloned, set to None + }, + ) + }) + .collect() } pub async fn get_detailed_connection_stats(&self) -> Vec<(String, bool, f32)> { let upstreams = self.upstreams.lock().await; @@ -300,4 +331,230 @@ impl MultiUpstreamManager { ); // Implement rebalancing logic here } + + pub async fn start_multi_upstream_share_accounting( + &self, + mut from_translator_recv: tokio::sync::mpsc::Receiver< + roles_logic_sv2::parsers::Mining<'static>, + >, + to_translator_send: tokio::sync::mpsc::Sender>, // Remove underscore + ) -> Result> { + info!("Starting multi-upstream share accounting"); + + // Clone the manager for use in the async task + let manager = self.clone(); + + let share_accounting_task = tokio::spawn({ + async move { + info!("Multi-upstream share accounting task started"); + + let mut message_count = 0; + let response_count = 0; + + // Start response handlers for each upstream + manager + .start_upstream_response_handlers(to_translator_send.clone()) + .await; + + loop { + tokio::select! { + // Handle outbound messages FROM translator TO upstreams + Some(message) = from_translator_recv.recv() => { + message_count += 1; + info!("Received message #{} from translator", message_count); + + // Route message to appropriate upstream based on hashrate distribution + manager.route_message_to_upstream(message).await; + } + + // Keep task alive + _ = tokio::time::sleep(tokio::time::Duration::from_secs(10)) => { + debug!("Multi-upstream share accounting keepalive (out: {}, in: {} messages)", message_count, response_count); + } + + else => { + warn!("All channels closed, stopping share accounting"); + break; + } + } + } + + info!( + "Multi-upstream share accounting task finished (out: {}, in: {} messages)", + message_count, response_count + ); + } + }); + + // Create AbortOnDrop to manage the task + Ok(crate::shared::utils::AbortOnDrop::new( + share_accounting_task, + )) + } + + pub async fn route_message_to_upstream( + &self, + message: roles_logic_sv2::parsers::Mining<'static>, + ) { + let upstreams_guard = self.upstreams.lock().await; + let active_upstreams: Vec<(String, f32)> = upstreams_guard + .iter() + .filter_map(|(id, conn)| { + if conn.is_active { + Some((id.clone(), conn.allocated_percentage)) + } else { + None + } + }) + .collect(); + + if active_upstreams.is_empty() { + warn!("No active upstreams available for message routing"); + return; + } + + let total_percentage: f32 = active_upstreams.iter().map(|(_, pct)| pct).sum(); + if total_percentage <= 0.0 { + warn!("Total upstream percentage is 0, cannot distribute"); + return; + } + + // Weighted random + let random_value = rand::random::() * total_percentage; + let mut cumulative = 0.0; + let mut selected_upstream: Option = None; + + for (upstream_id, percentage) in &active_upstreams { + cumulative += percentage; + if random_value <= cumulative { + selected_upstream = Some(upstream_id.clone()); + break; + } + } + + if let Some(upstream_id) = selected_upstream { + let percentage = active_upstreams + .iter() + .find(|(id, _)| id == &upstream_id) + .map(|(_, pct)| *pct) + .unwrap_or(0.0); + + let message_type = match &message { + roles_logic_sv2::parsers::Mining::OpenStandardMiningChannel(_) => { + "OpenStandardMiningChannel" + } + roles_logic_sv2::parsers::Mining::OpenExtendedMiningChannel(_) => { + "OpenExtendedMiningChannel" + } + roles_logic_sv2::parsers::Mining::SubmitSharesStandard(_) => "SubmitSharesStandard", + roles_logic_sv2::parsers::Mining::SubmitSharesExtended(_) => "SubmitSharesExtended", + _ => "Other", + }; + + debug!( + "Routing {} message to upstream: {} ({:.1}%)", + message_type, upstream_id, percentage + ); + + if let Some(upstream) = upstreams_guard.get(&upstream_id) { + if let Some(ref sender) = upstream.sender { + let pool_message = + demand_share_accounting_ext::parser::PoolExtMessages::Mining(message); + if let Err(e) = sender.send(pool_message).await { + warn!("Failed to send message to upstream {}: {}", upstream_id, e); + } else { + info!( + "Successfully routed {} to {} ({:.1}%)", + message_type, upstream_id, percentage + ); + } + } else { + warn!("No sender available for upstream {}", upstream_id); + } + } + } + + drop(upstreams_guard); + } + /// Start response handlers for all upstream connections + async fn start_upstream_response_handlers( + &self, + to_translator_send: tokio::sync::mpsc::Sender>, + ) { + let upstreams = self.upstreams.lock().await; + + for (upstream_id, _conn) in upstreams.iter() { + let upstream_id = upstream_id.clone(); + let upstreams = self.upstreams.clone(); + let to_translator_send = to_translator_send.clone(); + + tokio::spawn(async move { + loop { + // Try to get the receiver for this upstream + let mut receiver_opt = None; + { + let mut upstreams_guard = upstreams.lock().await; + if let Some(upstream) = upstreams_guard.get_mut(&upstream_id) { + if upstream.is_active && upstream.receiver.is_some() { + receiver_opt = upstream.receiver.take(); // Take ownership + } + } + } + + if let Some(mut receiver) = receiver_opt { + info!("Starting response handler for upstream: {}", upstream_id); + + // Listen for responses from this upstream + loop { + tokio::select! { + Some(response) = receiver.recv() => { + debug!("Received response from upstream: {}", upstream_id); + + // Convert PoolExtMessages back to Mining message + if let demand_share_accounting_ext::parser::PoolExtMessages::Mining(mining_response) = response { + // Forward response back to translator (and then to cpuminer) + if let Err(e) = to_translator_send.send(mining_response).await { + warn!("Failed to send response to translator from {}: {}", upstream_id, e); + break; + } else { + debug!("Successfully forwarded response from {} to translator", upstream_id); + } + } + } + + _ = tokio::time::sleep(tokio::time::Duration::from_secs(30)) => { + // Check if upstream is still active + let upstreams_guard = upstreams.lock().await; + if let Some(upstream) = upstreams_guard.get(&upstream_id) { + if !upstream.is_active { + warn!("Upstream {} became inactive, stopping response handler", upstream_id); + break; + } + } + } + + else => { + warn!("Response channel closed for upstream: {}", upstream_id); + break; + } + } + } + + // Put receiver back if upstream is still active + { + let mut upstreams_guard = upstreams.lock().await; + if let Some(upstream) = upstreams_guard.get_mut(&upstream_id) { + if upstream.is_active { + upstream.receiver = Some(receiver); + } + } + } + } else { + // No receiver available, wait and retry + tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; + } + } + }); + } + } } From 1b3e12b0af34b3451176d135ed6f04f5a618bd06 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Sun, 8 Jun 2025 21:44:40 +0530 Subject: [PATCH 20/21] Remove cpuminer-opt directory --- cpuminer-opt | 1 - 1 file changed, 1 deletion(-) delete mode 160000 cpuminer-opt diff --git a/cpuminer-opt b/cpuminer-opt deleted file mode 160000 index dd99580a..00000000 --- a/cpuminer-opt +++ /dev/null @@ -1 +0,0 @@ -Subproject commit dd99580a4c05414562ccc1bffca7cdce2073fa29 From 1302954b443d76e2c5852a00ead3841a341df615 Mon Sep 17 00:00:00 2001 From: bansalayush247 Date: Wed, 11 Jun 2025 15:26:30 +0530 Subject: [PATCH 21/21] feat: Implement ConnectionManager for managing upstream connections and downstream hashrate --- .gitignore | 1 + src/main.rs | 3 +- src/proxy_state.rs | 77 +------------------ src/router/connection_manager.rs | 98 ++++++++++++++++++++++++ src/router/mod.rs | 108 ++++++++++++--------------- src/router/multi_upstream_manager.rs | 4 +- 6 files changed, 150 insertions(+), 141 deletions(-) create mode 100644 src/router/connection_manager.rs diff --git a/.gitignore b/.gitignore index ea8c4bf7..1e6eccae 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ /target +/cpuminer-opt \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 58ce2586..fb34d9fc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,7 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; +use crate::router::connection_manager::ConnectionManager; // Add this import use crate::shared::utils::AbortOnDrop; use config::Configuration; use key_utils::Secp256k1PublicKey; @@ -86,7 +87,7 @@ async fn main() { .unwrap_or_else(|| panic!("Pool address is missing")); // Set downstream hashrate using Configuration pattern - ProxyState::set_downstream_hashrate(Configuration::downstream_hashrate()); + ConnectionManager::set_downstream_hashrate(Configuration::downstream_hashrate()); // Get pool addresses with auth keys using Configuration pattern let pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)> = pool_addresses diff --git a/src/proxy_state.rs b/src/proxy_state.rs index ade632d7..aac0740c 100644 --- a/src/proxy_state.rs +++ b/src/proxy_state.rs @@ -1,6 +1,4 @@ -use std::collections::HashMap; use std::sync::Arc; -use std::time::Instant; use lazy_static::lazy_static; use roles_logic_sv2::utils::Mutex; @@ -87,22 +85,8 @@ pub enum UpstreamType { TranslatorUpstream, } -/// Create an UpstreamConnection struct to store connection info -#[allow(dead_code)] -#[derive(Debug, Clone)] -pub struct UpstreamConnection { - pub url: String, - pub address: std::net::SocketAddr, - pub auth_key: key_utils::Secp256k1PublicKey, - pub connection_type: UpstreamType, - pub is_connected: bool, - pub shares_submitted: u64, - pub shares_accepted: u64, - pub last_used: Instant, -} - /// Represents global proxy state -#[derive(Debug)] +#[derive(Debug, Serialize, Deserialize)] pub struct ProxyState { pub pool: PoolState, pub tp: TpState, @@ -112,10 +96,6 @@ pub struct ProxyState { pub inconsistency: Option, pub downstream: DownstreamState, pub upstream: UpstreamState, - - // New fields for multiple upstream support - pub upstream_connections: HashMap, - pub total_hashrate: f32, } impl ProxyState { @@ -129,10 +109,6 @@ impl ProxyState { inconsistency: None, downstream: DownstreamState::Up, upstream: UpstreamState::Up, - - // Initialize new fields - upstream_connections: HashMap::new(), - total_hashrate: 0.0, } } @@ -263,57 +239,6 @@ impl ProxyState { } } - /// Set the downstream hashrate to be distributed among upstreams - pub fn set_downstream_hashrate(hashrate: f32) { - info!("Setting downstream hashrate to: {} h/s", hashrate); - if PROXY_STATE - .safe_lock(|state| { - state.total_hashrate = hashrate; - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - - /// Get the downstream hashrate - pub fn get_downstream_hashrate() -> f32 { - let mut hashrate = 0.0; - if PROXY_STATE - .safe_lock(|state| { - hashrate = state.total_hashrate; - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - hashrate - } - - /// Update connection status for an upstream with timestamp - pub fn set_upstream_connection_status(id: &str, connected: bool) { - if PROXY_STATE - .safe_lock(|state| { - if let Some(conn) = state.upstream_connections.get_mut(id) { - conn.is_connected = connected; - if connected { - conn.last_used = Instant::now(); - info!("Upstream {} is now connected", id); - } else { - info!("Upstream {} is now disconnected", id); - } - } - }) - .is_err() - { - error!("Global Proxy Mutex Corrupted"); - std::process::exit(1); - } - } - - /// Check if proxy is down pub fn is_proxy_down() -> (bool, Option) { let errors = Self::get_errors(); if errors.is_ok() && errors.as_ref().unwrap().is_empty() { diff --git a/src/router/connection_manager.rs b/src/router/connection_manager.rs new file mode 100644 index 00000000..aafe7457 --- /dev/null +++ b/src/router/connection_manager.rs @@ -0,0 +1,98 @@ +use lazy_static::lazy_static; +use roles_logic_sv2::utils::Mutex; +use std::collections::HashMap; +use tracing::{error, info}; + +lazy_static! { + static ref DOWNSTREAM_HASHRATE: Mutex = Mutex::new(0.0); + static ref UPSTREAM_CONNECTIONS: Mutex> = + Mutex::new(HashMap::new()); +} + +#[derive(Debug, Clone)] +pub struct ConnectionInfo { + pub connected: bool, + pub address: std::net::SocketAddr, + #[allow(dead_code)] + pub allocated_percentage: f32, + #[allow(dead_code)] + pub last_seen: std::time::Instant, +} + +pub struct ConnectionManager; + +impl ConnectionManager { + pub fn set_downstream_hashrate(hashrate: f32) { + info!("Setting downstream hashrate to: {} h/s", hashrate); + if DOWNSTREAM_HASHRATE.safe_lock(|h| *h = hashrate).is_err() { + error!("Failed to set downstream hashrate"); + } + } + + pub fn get_downstream_hashrate() -> f32 { + let mut hashrate = 0.0; + if DOWNSTREAM_HASHRATE.safe_lock(|h| hashrate = *h).is_err() { + error!("Failed to get downstream hashrate"); + } + hashrate + } + + pub fn set_upstream_connection(id: &str, info: ConnectionInfo) { + if UPSTREAM_CONNECTIONS + .safe_lock(|conns| { + conns.insert(id.to_string(), info.clone()); + }) + .is_err() + { + error!("Failed to update upstream connection"); + } + + if info.connected { + info!("Upstream {} connected to {}", id, info.address); + } else { + info!("Upstream {} disconnected", id); + } + } + + #[allow(dead_code)] + pub fn get_upstream_connection(id: &str) -> Option { + let mut result = None; + if UPSTREAM_CONNECTIONS + .safe_lock(|conns| { + result = conns.get(id).cloned(); + }) + .is_err() + { + error!("Failed to get upstream connection"); + } + result + } + + #[allow(dead_code)] + pub fn get_all_upstream_connections() -> HashMap { + let mut result = HashMap::new(); + if UPSTREAM_CONNECTIONS + .safe_lock(|conns| { + result = conns.clone(); + }) + .is_err() + { + error!("Failed to get all upstream connections"); + } + result + } + + #[allow(dead_code)] + pub fn update_upstream_percentage(id: &str, percentage: f32) { + if UPSTREAM_CONNECTIONS + .safe_lock(|conns| { + if let Some(info) = conns.get_mut(id) { + info.allocated_percentage = percentage; + } + }) + .is_err() + { + error!("Failed to update upstream percentage"); + } + } +} diff --git a/src/router/mod.rs b/src/router/mod.rs index 848b4c5c..ea513066 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -25,8 +25,11 @@ use crate::{ shared::utils::AbortOnDrop, }; -// Add the new module +// Add the new modules +pub mod connection_manager; pub mod multi_upstream_manager; + +pub use connection_manager::ConnectionManager; pub use multi_upstream_manager::MultiUpstreamManager; /// Router handles connection to Multiple upstreams. @@ -43,6 +46,7 @@ pub struct Router { latency_tx: tokio::sync::watch::Sender>, pub latency_rx: tokio::sync::watch::Receiver>, } + impl Clone for Router { fn clone(&self) -> Self { Self { @@ -65,11 +69,7 @@ impl Router { pub fn new( pool_addresses: Vec, auth_pub_k: Secp256k1PublicKey, - // Configuration msg used to setup connection between client and pool - // If not, present `get_mining_setup_connection_msg()` is called to generated default values setup_connection_msg: Option>, - // Max duration for pool setup after which it times out. - // If None, default time of 5s is used. timer: Option, ) -> Self { let (latency_tx, latency_rx) = watch::channel(None); @@ -98,7 +98,7 @@ impl Router { let pool_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); let keys: Vec = pool_address_keys.iter().map(|(_, key)| *key).collect(); - let auth_pub_k = keys[0]; // Use first key as primary + let auth_pub_k = keys[0]; let (latency_tx, latency_rx) = watch::channel(None); @@ -121,7 +121,7 @@ impl Router { pool_address_keys: Vec<(SocketAddr, Secp256k1PublicKey)>, setup_connection_msg: Option>, timer: Option, - use_distribution: bool, // Changed from use_parallel to use_distribution + use_distribution: bool, ) -> Result { let pool_addresses: Vec = pool_address_keys.iter().map(|(addr, _)| *addr).collect(); @@ -139,7 +139,6 @@ impl Router { None }; - // Use let auth_pub_k = keys .first() .copied() @@ -178,8 +177,6 @@ impl Router { pub async fn monitor_upstream(&mut self, epsilon: Duration) -> Option { // For multi-upstream mode, we don't switch since we use all simultaneously if self.is_multi_upstream_enabled() { - // In distribution mode, we don't need to switch upstreams - // All upstreams are used simultaneously return None; } @@ -209,6 +206,7 @@ impl Router { best_pool.map(|pool| (pool, least_latency)) } + /// Select the best pool for monitoring async fn select_pool_monitor(&self, epsilon: Duration) -> Option { if let Some((best_pool, best_pool_latency)) = self.select_pool().await { @@ -223,7 +221,6 @@ impl Router { return Some(best_pool); } }; - // saturating_sub is used to avoid panic on negative duration result if best_pool_latency < current_latency.saturating_sub(epsilon) { info!( "Found faster pool: {:?} with latency {:?}", @@ -244,7 +241,6 @@ impl Router { pub async fn select_pool_connect(&mut self) -> Option { info!("Selecting the best upstream"); - // Remove round-robin logic - only use latency-based selection if let Some((pool, latency)) = self.select_pool().await { info!("Latency for upstream {:?} is {:?}", pool, latency); self.latency_tx.send_replace(Some(latency)); @@ -254,7 +250,7 @@ impl Router { } } - /// Selects the best upstream and connects to. + /// Selects the best upstream and connects to it. /// Uses minin_pool_connection::connect_pool pub async fn connect_pool( &mut self, @@ -271,8 +267,6 @@ impl Router { Some(addr) => addr, None => match self.select_pool_connect().await { Some(addr) => addr, - // Called when we initialize the proxy, without a valid pool we can not start mine and we - // return Err None => { return Err(minin_pool_connection::errors::Error::Unrecoverable); } @@ -282,7 +276,7 @@ impl Router { info!("Upstream {:?} selected", pool); - // Find the matching auth key for this address - fix field name + // Find the matching auth key for this address let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool) { self.keys[index] @@ -299,7 +293,7 @@ impl Router { .await { Ok((send_to_pool, recv_from_pool, pool_connection_abortable)) => { - // Update ProxyState with successful connection + // Update ConnectionManager with successful connection let upstream_id = format!( "upstream-{}", self.pool_addresses @@ -307,7 +301,14 @@ impl Router { .position(|&a| a == pool) .unwrap_or(0) ); - ProxyState::set_upstream_connection_status(&upstream_id, true); + + let connection_info = connection_manager::ConnectionInfo { + connected: true, + address: pool, + allocated_percentage: 100.0, // Single upstream gets 100% + last_seen: std::time::Instant::now(), + }; + ConnectionManager::set_upstream_connection(&upstream_id, connection_info); // Update current pool address crate::POOL_ADDRESS @@ -316,14 +317,14 @@ impl Router { }) .unwrap_or_else(|_| { error!("Pool address Mutex corrupt"); - crate::proxy_state::ProxyState::update_inconsistency(Some(1)); + ProxyState::update_inconsistency(Some(1)); }); Ok((send_to_pool, recv_from_pool, pool_connection_abortable)) } Err(e) => { - // Update ProxyState with failed connection + // Update ConnectionManager with failed connection let upstream_id = format!( "upstream-{}", self.pool_addresses @@ -331,7 +332,14 @@ impl Router { .position(|&a| a == pool) .unwrap_or(0) ); - ProxyState::set_upstream_connection_status(&upstream_id, false); + + let connection_info = connection_manager::ConnectionInfo { + connected: false, + address: pool, + allocated_percentage: 0.0, + last_seen: std::time::Instant::now(), + }; + ConnectionManager::set_upstream_connection(&upstream_id, connection_info); Err(e) } @@ -355,9 +363,25 @@ impl Router { } } + /// Start multi-upstream share accounting with JDC + pub async fn start_multi_upstream_share_accounting_with_jdc( + &self, + from_jdc_recv: Receiver>, + to_jdc_send: Sender>, + ) -> Result> { + if let Some(ref manager) = self.upstream_manager { + let manager_handle = manager + .start_multi_upstream_share_accounting(from_jdc_recv, to_jdc_send) + .await?; + Ok(manager_handle) + } else { + Err("Multi-upstream manager not initialized".into()) + } + } + /// Returns the sum all the latencies for a given upstream async fn get_latency(&self, pool_address: SocketAddr) -> Result { - // Find the auth key for this address - fix field names + // Find the auth key for this address let auth_pub_key = if let Some(index) = self.pool_addresses.iter().position(|&a| a == pool_address) { self.keys[index] @@ -369,7 +393,6 @@ impl Router { let setup_connection_msg = self.setup_connection_msg.as_ref(); let timer = self.timer.as_ref(); - // Rest of the function remains the same tokio::time::timeout( Duration::from_secs(15), PoolLatency::get_mining_setup_latencies( @@ -387,7 +410,6 @@ impl Router { ); })??; - // Rest of the function remains unchanged if (PoolLatency::get_mining_setup_latencies( &mut pool, setup_connection_msg.cloned(), @@ -416,7 +438,6 @@ impl Router { pool.open_sv2_jd_connection, pool.get_a_mining_token, ]; - // Get sum of all latencies for pool let sum_of_latencies: Duration = latencies.iter().flatten().sum(); Ok(sum_of_latencies) } @@ -429,10 +450,8 @@ impl Router { self.pool_addresses.len() ); - // Add each unique upstream only once - fix field names for (idx, (addr, key)) in self.pool_addresses.iter().zip(self.keys.iter()).enumerate() { let id = format!("upstream-{}", idx); - debug!("Adding upstream {}: {} ({})", id, addr, key); if let Err(e) = manager @@ -449,10 +468,7 @@ impl Router { } } - // IMPORTANT: Initialize connections BEFORE any hashrate distribution is set - // This ensures add_upstream() calls are complete and won't overwrite distribution manager.initialize_connections().await; - Ok(()) } else { Err("No upstream manager available".to_string()) @@ -471,23 +487,6 @@ impl Router { Err("No upstream manager available") } } - pub async fn start_multi_upstream_share_accounting_with_jdc( - &self, - from_jdc_recv: Receiver>, - to_jdc_send: Sender>, - ) -> Result> { - if let Some(ref manager) = self.upstream_manager { - // Direct pass-through since types already match - let manager_handle = manager - .start_multi_upstream_share_accounting(from_jdc_recv, to_jdc_send) - .await?; - - // No type conversion needed - just return the handle - Ok(manager_handle) - } else { - Err("Multi-upstream manager not initialized".into()) - } - } } /// Track latencies for various stages of pool connection setup. @@ -503,7 +502,6 @@ struct PoolLatency { } impl PoolLatency { - // Create new `PoolLatency` given an upstream address fn new(pool: SocketAddr) -> PoolLatency { Self { pool, @@ -516,15 +514,12 @@ impl PoolLatency { } } - /// Sets the `PoolLatency`'s `open_sv2_mining_connection`, `setup_channel_timer`, `receive_first_job`, - /// and `receive_first_set_new_prev_hash` async fn get_mining_setup_latencies( &mut self, setup_connection_msg: Option>, timer: Option, authority_public_key: Secp256k1PublicKey, ) -> Result<(), ()> { - // Set open_sv2_mining_connection latency let open_sv2_mining_connection_timer = Instant::now(); match TcpStream::connect(self.pool).await { Ok(stream) => { @@ -538,7 +533,6 @@ impl PoolLatency { ) .await?; - // Set setup_channel latency let setup_channel_timer = Instant::now(); let result = mining_setup_connection( &mut receiver, @@ -575,18 +569,15 @@ impl PoolLatency { _new_ext_job, )) = message.clone() { - // Set receive_first_job latency self.receive_first_job = Some(timer.elapsed()); received_new_job = true; } if let PoolExtMessages::Mining(Mining::SetNewPrevHash(_new_prev_hash)) = message.clone() { - // Set receive_first_set_new_prev_hash latency self.receive_first_set_new_prev_hash = Some(timer.elapsed()); received_prev_hash = true; } - // Both latencies have been set so we break the loop if received_new_job && received_prev_hash { break; } @@ -612,14 +603,11 @@ impl PoolLatency { } } - /// Sets the `PoolLatency`'s `open_sv2_jd_connection` and `get_a_mining_token` async fn get_jd_latencies( &mut self, authority_public_key: Secp256k1PublicKey, ) -> Result<(), ()> { let address = self.pool; - - // Set open_sv2_jd_connection latency let open_sv2_jd_connection_timer = Instant::now(); match tokio::time::timeout(Duration::from_secs(2), TcpStream::connect(address)).await { @@ -629,7 +617,6 @@ impl PoolLatency { .map_err(|_| error!(" TP_ADDRESS Mutex Corrupted"))?; if let Some(_tp_addr) = tp { let initiator = Initiator::from_raw_k(authority_public_key.into_bytes()) - // Safe expect Key is a constant and must be right .expect("Unable to create initialtor"); let (mut receiver, mut sender, _, _) = match Connection::new(stream, HandshakeRole::Initiator(initiator)).await { @@ -673,7 +660,6 @@ impl PoolLatency { } }; - // Set get_a_mining_token latency let get_a_mining_token_timer = Instant::now(); let _token = JobDeclarator::get_last_token(&job_declarator).await; self.get_a_mining_token = Some(get_a_mining_token_timer.elapsed()); @@ -698,7 +684,6 @@ fn open_channel() -> Mining<'static> { user_identity: "ABC" .to_string() .try_into() - // This can never fail .expect("Failed to convert user identity to string"), nominal_hash_rate: 0.0, }, @@ -718,7 +703,6 @@ async fn initialize_mining_connections( (), > { let initiator = - // Safe expect Key is a constant and must be right Initiator::from_raw_k(authority_public_key.into_bytes()).expect("Invalid authority key"); let (receiver, sender, _, _) = match Connection::new(stream, HandshakeRole::Initiator(initiator)).await { diff --git a/src/router/multi_upstream_manager.rs b/src/router/multi_upstream_manager.rs index 230e250a..4f927947 100644 --- a/src/router/multi_upstream_manager.rs +++ b/src/router/multi_upstream_manager.rs @@ -1,4 +1,4 @@ -use crate::{minin_pool_connection, proxy_state::ProxyState, HashUnit}; +use crate::{minin_pool_connection, router::connection_manager::ConnectionManager, HashUnit}; use key_utils::Secp256k1PublicKey; use rand; use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration}; @@ -178,7 +178,7 @@ impl MultiUpstreamManager { } // Assign percentages to best latency pools first - let downstream_hashrate = ProxyState::get_downstream_hashrate(); + let downstream_hashrate = ConnectionManager::get_downstream_hashrate(); for (i, (pool_id, _address, latency)) in ranked_pools.iter().enumerate() { if let Some(&percentage) = distribution.get(i) {