implemented separate clients for monitor and predict | some refactoring

Author: pushkar1262

src/client.rs

@@ -1,6 +1,6 @@
 use crate::{
 	error::{Error, TimeoutError},
-	prelude::{ChainClient, LOG_TARGET},
+	prelude::{ChainClient, Config, LOG_TARGET},
 	prometheus,
 };
 use std::{
@@ -11,6 +11,7 @@ use std::{
 	time::Duration,
 };
 use subxt::backend::{
+	chain_head::{ChainHeadBackend, ChainHeadBackendBuilder},
 	legacy::LegacyBackend,
 	rpc::reconnecting_rpc_client::{ExponentialBackoff, RpcClient as ReconnectingRpcClient},
 };
@@ -52,7 +53,7 @@ pub struct Client {
 }
 
 impl Client {
-	/// Create a new client from a comma-separated list of RPC endpoints.
+	/// Create a new client from a comma-separated list of RPC endpoints using ChainHeadBackend.
 	///
 	/// The client will try each endpoint in sequence until one connects successfully.
 	/// Multiple endpoints can be specified for failover:
@@ -72,7 +73,38 @@ impl Client {
 			log::info!(target: LOG_TARGET, "RPC endpoint pool: {} endpoint(s)", endpoints.len());
 		}
 
-		let (chain_api, connected_index) = Self::connect_with_failover(&endpoints, 0).await?;
+		let (chain_api, connected_index) =
+			Self::connect_with_failover(&endpoints, 0, false).await?;
+
+		Ok(Self {
+			chain_api: Arc::new(RwLock::new(chain_api)),
+			endpoints: Arc::new(endpoints),
+			current_endpoint_index: Arc::new(AtomicUsize::new(connected_index)),
+			reconnect_generation: Arc::new(AtomicUsize::new(0)),
+		})
+	}
+
+	/// Create a new client from a comma-separated list of RPC endpoints using LegacyBackend.
+	///
+	/// The client will try each endpoint in sequence until one connects successfully.
+	/// Multiple endpoints can be specified for failover:
+	/// "wss://rpc1.example.com,wss://rpc2.example.com"
+	pub async fn new_with_legacy_backend(uris: &str) -> Result<Self, Error> {
+		let endpoints: Vec<String> = uris
+			.split(',')
+			.map(|s| s.trim().to_string())
+			.filter(|s| !s.is_empty())
+			.collect();
+
+		if endpoints.is_empty() {
+			return Err(Error::Other("No RPC endpoints provided".into()));
+		}
+
+		if endpoints.len() > 1 {
+			log::info!(target: LOG_TARGET, "RPC endpoint pool: {} endpoint(s)", endpoints.len());
+		}
+
+		let (chain_api, connected_index) = Self::connect_with_failover(&endpoints, 0, true).await?;
 
 		Ok(Self {
 			chain_api: Arc::new(RwLock::new(chain_api)),
@@ -95,6 +127,7 @@ impl Client {
 	async fn connect_with_failover(
 		endpoints: &[String],
 		start_index: usize,
+		use_legacy: bool,
 	) -> Result<(ChainClient, usize), Error> {
 		let mut last_error = None;
 		let total = endpoints.len();
@@ -112,7 +145,7 @@ impl Client {
 					"attempting to connect to {uri:?} (endpoint {endpoint_num}/{total}, attempt {attempt}/{max_attempts})"
 				);
 
-				match Self::try_connect(uri).await {
+				match Self::try_connect(uri, use_legacy).await {
 					Ok(client) => {
 						if total > 1 {
 							log::info!(
@@ -150,8 +183,8 @@ impl Client {
 	}
 
 	/// Try to connect to a single endpoint with timeout.
-	async fn try_connect(uri: &str) -> Result<ChainClient, Error> {
-		let connect_future = async {
+	async fn try_connect(uri: &str, use_legacy: bool) -> Result<ChainClient, Error> {
+		let connect_future = async move {
 			let reconnecting_rpc = ReconnectingRpcClient::builder()
 				.retry_policy(
 					ExponentialBackoff::from_millis(500).max_delay(Duration::from_secs(10)).take(3),
@@ -160,11 +193,18 @@ impl Client {
 				.await
 				.map_err(|e| Error::Other(format!("Failed to connect: {e:?}")))?;
 
-			let backend = LegacyBackend::builder().build(reconnecting_rpc.clone());
-
-			let chain_api = ChainClient::from_backend(Arc::new(backend)).await?;
-
-			log::info!(target: LOG_TARGET, "Connected to {uri} with Legacy backend");
+			let chain_api = if use_legacy {
+				let backend = LegacyBackend::builder().build(reconnecting_rpc.clone());
+				let client = ChainClient::from_backend(Arc::new(backend)).await?;
+				log::info!(target: LOG_TARGET, "Connected to {uri} with Legacy backend");
+				client
+			} else {
+				let backend: ChainHeadBackend<Config> = ChainHeadBackendBuilder::default()
+					.build_with_background_driver(reconnecting_rpc);
+				let client = ChainClient::from_backend(Arc::new(backend)).await?;
+				log::info!(target: LOG_TARGET, "Connected to {uri} with ChainHead backend");
+				client
+			};
 
 			Ok::<ChainClient, Error>(chain_api)
 		};
@@ -216,7 +256,7 @@ impl Client {
 
 		// Establish new connection before acquiring write lock
 		let (new_client, connected_idx) =
-			Self::connect_with_failover(&self.endpoints, start_idx).await?;
+			Self::connect_with_failover(&self.endpoints, start_idx, false).await?;
 
 		// Acquire write lock and check if another task already reconnected
 		let mut guard = self.chain_api.write().await;

src/commands/predict.rs

@@ -79,7 +79,9 @@ where
 				.map_err(|e| {
 					Error::Other(format!("Failed to fetch Desired Targets from chain: {e}"))
 				})?
-				.expect("Error in fetching desired validators from chain")
+				.ok_or_else(|| {
+					Error::Other("Desired validators not found in chain storage".to_string())
+				})?
 		},
 	};
 
@@ -100,7 +102,10 @@ where
 	let (target_snapshot, mut voter_snapshot) =
 		convert_election_data_to_snapshots::<T>(candidates, nominators)?;
 
-	// Fix the order for staking data source
+	// When fetching from staking data, voters come from BagsList in descending order (highest
+	// stake first). The SDK expects page 0 (lsp) to contain lowest stake voters and page n-1
+	// (msp) to contain highest stake voters. Reversing ensures correct page assignment during
+	// pagination.
 	if matches!(data_source, ElectionDataSource::Staking) {
 		voter_snapshot.reverse();
 	}
@@ -114,7 +119,7 @@ where
 	);
 
 	// Use actual voter page count, not the chain's max pages
-	// Staking data may have added some pages
+	// Staking/Overridden data may have added some pages
 	let n_pages = n_pages.max(voter_snapshot.len() as u32);
 
 	// Mine the solution with timeout to prevent indefinite hanging
@@ -181,7 +186,7 @@ where
 	let validators_output = output_dir.join("validators_prediction.json");
 	let nominators_output = output_dir.join("nominators_prediction.json");
 	// Save validators prediction
-	write_data_to_json_file(&validators_prediction, validators_output.to_str().unwrap()).await?;
+	write_data_to_json_file(&validators_prediction, &validators_output).await?;
 
 	log::info!(
 		target: LOG_TARGET,
@@ -190,7 +195,7 @@ where
 	);
 
 	// Save nominators prediction
-	write_data_to_json_file(&nominators_prediction, nominators_output.to_str().unwrap()).await?;
+	write_data_to_json_file(&nominators_prediction, &nominators_output).await?;
 
 	log::info!(
 		target: LOG_TARGET,

src/dynamic/election_data.rs

@@ -267,11 +267,10 @@ where
 	let validators_with_only_self_vote: HashSet<AccountId> = all_voters
 		.iter()
 		.filter(|(nominator, _, targets)| {
-			// validator has only self-vote if:
+			// validator has only self-vote if either:
 			// 1. They are a validator (in active_set)
 			// 2. Their only target is themselves
-			// NOTE: Reverted to your original logic as requested, assuming you want strictly this
-			// behavior.
+
 			active_set.contains(nominator) || (targets.len() == 1 && targets[0] == *nominator)
 		})
 		.map(|(nominator, _, _)| nominator.clone())

src/dynamic/staking.rs

@@ -185,7 +185,7 @@ pub(crate) async fn fetch_voter_list(
 	// Increase voter limit to have a buffer for filtering ineligible voters later
 	let extended_voter_limit = voter_limit.saturating_add(100);
 
-	log::info!(target: LOG_TARGET, "Fetching From Voter List");
+	log::info!(target: LOG_TARGET, "Fetching From VoterList");
 
 	// Fetch all bags (ListBags) - store as HashMap with bag_upper as key
 	log::trace!(target: LOG_TARGET, "Fetching ListBags...");
@@ -371,7 +371,7 @@ pub(crate) async fn fetch_voter_list(
 
 	log::info!(
 		target: LOG_TARGET,
-		"Voter List Fetch Completed"
+		"VoterList Fetch Completed"
 	);
 	log::info!(
 		target: LOG_TARGET,

src/main.rs

@@ -103,7 +103,11 @@ async fn main() -> Result<(), Error> {
 	// Initialize the timestamp so that if connection hangs, the stall detection alert can fire.
 	prometheus::set_last_block_processing_time();
 
-	let client = Client::new(&uri).await?;
+	// Create client with appropriate backend based on command type
+	let client = match command {
+		Command::Predict(_) => Client::new_with_legacy_backend(&uri).await?,
+		_ => Client::new(&uri).await?,
+	};
 
 	let version_bytes = client
 		.chain_api()

src/utils.rs

@@ -123,11 +123,12 @@ pub async fn wait_tx_in_finalized_block(
 }
 
 /// Write data to a JSON file
-pub async fn write_data_to_json_file<T>(data: &T, file_path: &str) -> Result<(), Error>
+pub async fn write_data_to_json_file<T, P>(data: &T, file_path: &P) -> Result<(), Error>
 where
 	T: Serialize,
+	P: AsRef<Path>,
 {
-	let path = Path::new(file_path);
+	let path = file_path.as_ref();
 	if let Some(parent) = path.parent() &&
 		!parent.as_os_str().is_empty()
 	{
@@ -145,11 +146,12 @@ where
 }
 
 /// Read data from a JSON file
-pub async fn read_data_from_json_file<T>(file_path: &str) -> Result<T, Error>
+pub async fn read_data_from_json_file<T, P>(file_path: P) -> Result<T, Error>
 where
 	T: DeserializeOwned,
+	P: AsRef<Path>,
 {
-	let path = Path::new(file_path);
+	let path = file_path.as_ref();
 
 	let mut file = File::open(path)?;
 	let mut content = String::new();
@@ -345,12 +347,11 @@ mod tests {
 	async fn test_read_write_json_file() {
 		let dir = tempfile::tempdir().unwrap();
 		let file_path = dir.path().join("test.json");
-		let file_path_str = file_path.to_str().unwrap();
 
 		let data = vec![1, 2, 3];
-		write_data_to_json_file(&data, file_path_str).await.unwrap();
+		write_data_to_json_file(&data, &file_path).await.unwrap();
 
-		let read_data: Vec<i32> = read_data_from_json_file(file_path_str).await.unwrap();
+		let read_data: Vec<i32> = read_data_from_json_file(&file_path).await.unwrap();
 		assert_eq!(data, read_data);
 	}