refactor

Author: lvboudre

src/channel.rs

@@ -2,6 +2,7 @@ use {
     etcd_client::{BalancedChannelBuilder, Channel, Client, ConnectOptions, Error},
     std::{
         collections::{HashMap, HashSet},
+        future,
         sync::{Arc, Mutex},
         time::Duration,
     },
@@ -31,7 +32,6 @@ pub struct ReliableChannelStats {
 pub struct ReliableChannelStatsRegistry {
     inner: Arc<Mutex<ReliableChannelStats>>,
     endpoint_status: Arc<Mutex<HashMap<String, EndpointStatus>>>,
-    probe_notify: Arc<Notify>,
     call_counter: Arc<dashmap::DashMap<Uri, usize>>,
 }
 
@@ -43,10 +43,6 @@ impl ReliableChannelStatsRegistry {
             .expect("reliable channel stats lock poisoned")
     }
 
-    pub fn trigger_probe(&self) {
-        self.probe_notify.notify_one();
-    }
-
     pub fn call_counts_snapshot(&self) -> HashMap<String, usize> {
         self.call_counter
             .iter()
@@ -77,6 +73,25 @@ impl ReliableChannelStatsRegistry {
     }
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct ProbeTrigger {
+    notify: Arc<Notify>,
+}
+
+impl ProbeTrigger {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn trigger_probe(&self) {
+        self.notify.notify_one();
+    }
+
+    async fn notified(&self) {
+        self.notify.notified().await;
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum EndpointStatus {
     Active,
@@ -94,6 +109,7 @@ pub struct ReliableBalancedChannelBuilder {
     pub probe_timeout: Duration,
     pub quarantine_retry_interval: Duration,
     stats_registry: ReliableChannelStatsRegistry,
+    probe_trigger: Option<ProbeTrigger>,
 }
 
 impl Default for ReliableBalancedChannelBuilder {
@@ -102,14 +118,35 @@ impl Default for ReliableBalancedChannelBuilder {
             probe_timeout: Duration::from_secs(5),
             quarantine_retry_interval: Duration::from_secs(15),
             stats_registry: ReliableChannelStatsRegistry::default(),
+            probe_trigger: None,
         }
     }
 }
 
 impl ReliableBalancedChannelBuilder {
+    pub fn new_with_probe_trigger(probe_trigger: ProbeTrigger) -> Self {
+        Self {
+            probe_trigger: Some(probe_trigger),
+            ..Self::default()
+        }
+    }
+
     pub fn stats_registry(&self) -> ReliableChannelStatsRegistry {
         self.stats_registry.clone()
     }
+
+    pub fn probe_trigger(&self) -> Option<ProbeTrigger> {
+        self.probe_trigger.clone()
+    }
+
+    pub fn set_probe_trigger(&mut self, probe_trigger: Option<ProbeTrigger>) {
+        self.probe_trigger = probe_trigger;
+    }
+
+    pub fn with_probe_trigger(mut self, probe_trigger: ProbeTrigger) -> Self {
+        self.probe_trigger = Some(probe_trigger);
+        self
+    }
 }
 
 impl BalancedChannelBuilder for ReliableBalancedChannelBuilder {
@@ -138,6 +175,7 @@ impl BalancedChannelBuilder for ReliableBalancedChannelBuilder {
             discover_updater,
             spy_quarantine_rx,
             self.stats_registry.clone(),
+            self.probe_trigger.clone(),
             self.probe_timeout,
             self.quarantine_retry_interval,
         ));
@@ -151,6 +189,7 @@ async fn run_endpoint_manager(
     discover_updater: EndpointUpdater,
     mut spy_quarantine_rx: mpsc::UnboundedReceiver<Uri>,
     stats_registry: ReliableChannelStatsRegistry,
+    probe_trigger: Option<ProbeTrigger>,
     probe_timeout: Duration,
     quarantine_retry_interval: Duration,
 ) {
@@ -186,7 +225,13 @@ async fn run_endpoint_manager(
                     &stats_registry,
                 );
             }
-            _ = stats_registry.probe_notify.notified() => {
+            _ = async {
+                if let Some(trigger) = &probe_trigger {
+                    trigger.notified().await;
+                } else {
+                    future::pending::<()>().await;
+                }
+            } => {
                 sweep_active(
                     &desired,
                     &mut active,

tests/test_channel.rs

@@ -5,7 +5,8 @@ use tonic::transport::{Endpoint, channel::Change};
 use tower::{Service, ServiceExt};
 
 use rust_etcd_utils::channel::{
-    EndpointStatus, ReliableBalancedChannelBuilder, connect_with_reliable_balanced_channel,
+    EndpointStatus, ProbeTrigger, ReliableBalancedChannelBuilder,
+    connect_with_reliable_balanced_channel,
 };
 
 mod common;
@@ -233,6 +234,8 @@ async fn request_fails_when_mock_endpoint_goes_down_then_recovers_after_restart(
     let mut builder = ReliableBalancedChannelBuilder::default();
     builder.probe_timeout = Duration::from_millis(100);
     builder.quarantine_retry_interval = Duration::from_millis(150);
+    let probe_trigger = ProbeTrigger::new();
+    builder.set_probe_trigger(Some(probe_trigger.clone()));
     let stats_registry = builder.stats_registry();
 
     let (mut channel, updater) = builder
@@ -307,7 +310,7 @@ async fn request_fails_when_mock_endpoint_goes_down_then_recovers_after_restart(
     );
 
     // Phase 3: trigger a one-shot health probe and verify quarantine.
-    stats_registry.trigger_probe();
+    probe_trigger.trigger_probe();
     tokio::time::timeout(Duration::from_secs(3), async {
         loop {
             if stats_registry
@@ -487,6 +490,8 @@ async fn fair_share_then_failover_drains_to_1000_successful_requests() {
     let mut builder = ReliableBalancedChannelBuilder::default();
     builder.probe_timeout = Duration::from_millis(100);
     builder.quarantine_retry_interval = Duration::from_millis(150);
+    let probe_trigger = ProbeTrigger::new();
+    builder.set_probe_trigger(Some(probe_trigger.clone()));
     let stats_registry = builder.stats_registry();
 
     let (mut channel, updater) = builder
@@ -590,7 +595,7 @@ async fn fair_share_then_failover_drains_to_1000_successful_requests() {
             .await;
     }
 
-    stats_registry.trigger_probe();
+    probe_trigger.trigger_probe();
 
     let mut successful_requests = first_phase_requests;
     let deadline = Instant::now() + Duration::from_secs(30);