refactor: improve weighted load balancing implementation

- Refactor the adapter selection logic for better clarity and maintainability.
- Add a test case to verify the weighted selection logic.

Author: DaMandal0rian

chain/ethereum/src/network.rs

@@ -197,46 +197,83 @@ impl EthereumNetworkAdapters {
         Self::available_with_capabilities(all, required_capabilities)
     }
 
-    // handle adapter selection from a list, implements the availability checking with an abstracted
+    """    // handle adapter selection from a list, implements the availability checking with an abstracted
     // source of the adapter list.
     fn cheapest_from(
         &self,
         input: Vec<&EthereumNetworkAdapter>,
         required_capabilities: &NodeCapabilities,
     ) -> Result<Arc<EthereumAdapter>, Error> {
-        let retest_rng: f64 = (&mut rand::rng()).random();
+        // First, try to re-test an errored adapter
+        if let Some(adapter) = self.retest_errored_adapter(&input) {
+            return Ok(adapter);
+        }
+
+        // If no adapter was re-tested, select the best available adapter
+        self.select_best_adapter(input, required_capabilities)
+    }
 
+    fn retest_errored_adapter(&self, input: &[&EthereumNetworkAdapter]) -> Option<Arc<EthereumAdapter>> {
+        let retest_rng: f64 = (&mut rand::rng()).random();
         if retest_rng < self.retest_percent {
-            if let Some(adapter) = input.iter().max_by_key(|a| a.current_error_count()) {
-                return Ok(adapter.adapter.clone());
-            }
+            input
+                .iter()
+                .max_by_key(|a| a.current_error_count())
+                .map(|adapter| adapter.adapter.clone())
+        } else {
+            None
         }
+    }
 
+    fn select_best_adapter(
+        &self,
+        input: Vec<&EthereumNetworkAdapter>,
+        required_capabilities: &NodeCapabilities,
+    ) -> Result<Arc<EthereumAdapter>, Error> {
         if self.weighted {
-            if input.is_empty() {
-                return Err(anyhow!(
-                    "A matching Ethereum network with {:?} was not found.",
-                    required_capabilities
-                ));
-            }
-            let weights: Vec<_> = input.iter().map(|a| a.weight).collect();
-            if let Ok(dist) = WeightedIndex::new(&weights) {
-                let idx = dist.sample(&mut rand::rng());
-                return Ok(input[idx].adapter.clone());
-            }
+            self.select_weighted_adapter(input, required_capabilities)
         } else {
-            let choices = input.into_iter().choose_multiple(&mut rand::rng(), 3);
-            if let Some(adapter) = choices.iter().min_by_key(|a| a.current_error_count()) {
-                return Ok(adapter.adapter.clone());
-            }
+            self.select_random_adapter(input, required_capabilities)
         }
+    }
 
-        Err(anyhow!(
-            "A matching Ethereum network with {:?} was not found.",
-            required_capabilities
-        ))
+    fn select_weighted_adapter(
+        &self,
+        input: Vec<&EthereumNetworkAdapter>,
+        required_capabilities: &NodeCapabilities,
+    ) -> Result<Arc<EthereumAdapter>, Error> {
+        if input.is_empty() {
+            return Err(anyhow!(
+                "A matching Ethereum network with {:?} was not found.",
+                required_capabilities
+            ));
+        }
+        let weights: Vec<_> = input.iter().map(|a| a.weight).collect();
+        if let Ok(dist) = WeightedIndex::new(&weights) {
+            let idx = dist.sample(&mut rand::rng());
+            Ok(input[idx].adapter.clone())
+        } else {
+            // Fallback to random selection if weights are invalid
+            self.select_random_adapter(input, required_capabilities)
+        }
     }
 
+    fn select_random_adapter(
+        &self,
+        input: Vec<&EthereumNetworkAdapter>,
+        required_capabilities: &NodeCapabilities,
+    ) -> Result<Arc<EthereumAdapter>, Error> {
+        let choices = input.into_iter().choose_multiple(&mut rand::rng(), 3);
+        if let Some(adapter) = choices.iter().min_by_key(|a| a.current_error_count()) {
+            Ok(adapter.adapter.clone())
+        } else {
+            Err(anyhow!(
+                "A matching Ethereum network with {:?} was not found.",
+                required_capabilities
+            ))
+        }
+    }""
+
     pub(crate) fn unverified_cheapest_with(
         &self,
         required_capabilities: &NodeCapabilities,
@@ -976,4 +1013,96 @@ mod tests {
             .await,
         )
     }
+
+    #[tokio::test]
+    async fn test_weighted_adapter_selection() {
+        let metrics = Arc::new(EndpointMetrics::mock());
+        let logger = graph::log::logger(true);
+        let mock_registry = Arc::new(MetricsRegistry::mock());
+        let transport = Transport::new_rpc(
+            Url::parse("http://127.0.0.1").unwrap(),
+            HeaderMap::new(),
+            metrics.clone(),
+            "",
+        );
+        let provider_metrics = Arc::new(ProviderEthRpcMetrics::new(mock_registry.clone()));
+
+        let adapter1 = Arc::new(
+            EthereumAdapter::new(
+                logger.clone(),
+                "adapter1".to_string(),
+                transport.clone(),
+                provider_metrics.clone(),
+                true,
+                false,
+            )
+            .await,
+        );
+
+        let adapter2 = Arc::new(
+            EthereumAdapter::new(
+                logger.clone(),
+                "adapter2".to_string(),
+                transport.clone(),
+                provider_metrics.clone(),
+                true,
+                false,
+            )
+            .await,
+        );
+
+        let adapters = EthereumNetworkAdapters::for_testing(
+            vec![
+                EthereumNetworkAdapter::new(
+                    metrics.cheap_clone(),
+                    NodeCapabilities {
+                        archive: true,
+                        traces: false,
+                    },
+                    adapter1.clone(),
+                    SubgraphLimit::Unlimited,
+                    0.2,
+                ),
+                EthereumNetworkAdapter::new(
+                    metrics.cheap_clone(),
+                    NodeCapabilities {
+                        archive: true,
+                        traces: false,
+                    },
+                    adapter2.clone(),
+                    SubgraphLimit::Unlimited,
+                    0.8,
+                ),
+            ],
+            vec![],
+        )
+        .await;
+
+        let mut adapters = adapters;
+        adapters.weighted = true;
+
+        let mut adapter1_count = 0;
+        let mut adapter2_count = 0;
+
+        for _ in 0..1000 {
+            let selected_adapter = adapters
+                .cheapest_with(&NodeCapabilities {
+                    archive: true,
+                    traces: false,
+                })
+                .await
+                .unwrap();
+
+            if selected_adapter.provider() == "adapter1" {
+                adapter1_count += 1;
+            } else {
+                adapter2_count += 1;
+            }
+        }
+
+        // Check that the selection is roughly proportional to the weights.
+        // Allow for a 10% tolerance.
+        assert!(adapter1_count > 100 && adapter1_count < 300);
+        assert!(adapter2_count > 700 && adapter2_count < 900);
+    }
 }