Event-driven Wakeup for ME Admission-gate

src/maestro/admission.rs

@@ -17,6 +17,7 @@ pub(crate) async fn configure_admission_gate(
     route_runtime: Arc<RouteRuntimeController>,
     admission_tx: &watch::Sender<bool>,
     config_rx: watch::Receiver<Arc<ProxyConfig>>,
+    me_ready_rx: watch::Receiver<u64>,
 ) {
     if config.general.use_middle_proxy {
         if let Some(pool) = me_pool.as_ref() {
@@ -52,6 +53,7 @@ pub(crate) async fn configure_admission_gate(
             let admission_tx_gate = admission_tx.clone();
             let route_runtime_gate = route_runtime.clone();
             let mut config_rx_gate = config_rx.clone();
+            let mut me_ready_rx_gate = me_ready_rx;
             let mut admission_poll_ms = config.general.me_admission_poll_ms.max(1);
             tokio::spawn(async move {
                 let mut gate_open = initial_gate_open;
@@ -74,6 +76,11 @@ pub(crate) async fn configure_admission_gate(
                             fast_fallback_enabled = cfg.general.me2dc_fallback && cfg.general.me2dc_fast;
                             continue;
                         }
+                        changed = me_ready_rx_gate.changed() => {
+                            if changed.is_err() {
+                                break;
+                            }
+                        }
                         _ = tokio::time::sleep(Duration::from_millis(admission_poll_ms)) => {}
                     }
                     let ready = pool_for_gate.admission_ready_conditional_cast().await;

src/maestro/me_startup.rs

@@ -3,7 +3,7 @@
 use std::sync::Arc;
 use std::time::Duration;
 
-use tokio::sync::RwLock;
+use tokio::sync::{RwLock, watch};
 use tracing::{error, info, warn};
 
 use crate::config::ProxyConfig;
@@ -29,6 +29,7 @@ pub(crate) async fn initialize_me_pool(
     rng: Arc<SecureRandom>,
     stats: Arc<Stats>,
     api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
+    me_ready_tx: watch::Sender<u64>,
 ) -> Option<Arc<MePool>> {
     if !use_middle_proxy {
         return None;
@@ -314,6 +315,7 @@ pub(crate) async fn initialize_me_pool(
                     let pool_bg = pool.clone();
                     let rng_bg = rng.clone();
                     let startup_tracker_bg = startup_tracker.clone();
+                    let me_ready_tx_bg = me_ready_tx.clone();
                     let retry_limit = if me_init_retry_attempts == 0 {
                         String::from("unlimited")
                     } else {
@@ -347,6 +349,9 @@ pub(crate) async fn initialize_me_pool(
                                         startup_tracker_bg
                                             .set_me_status(StartupMeStatus::Ready, "ready")
                                             .await;
+                                        me_ready_tx_bg.send_modify(|version| {
+                                            *version = version.saturating_add(1);
+                                        });
                                         info!(
                                             attempt = init_attempt,
                                             "Middle-End pool initialized successfully"
@@ -474,6 +479,9 @@ pub(crate) async fn initialize_me_pool(
                                 startup_tracker
                                     .set_me_status(StartupMeStatus::Ready, "ready")
                                     .await;
+                                me_ready_tx.send_modify(|version| {
+                                    *version = version.saturating_add(1);
+                                });
                                 info!(
                                     attempt = init_attempt,
                                     "Middle-End pool initialized successfully"

src/maestro/mod.rs

@@ -660,6 +660,8 @@ async fn run_telemt_core(
             .await;
     }
 
+    let (me_ready_tx, me_ready_rx) = watch::channel(0_u64);
+
     let me_pool: Option<Arc<MePool>> = me_startup::initialize_me_pool(
         use_middle_proxy,
         &config,
@@ -670,6 +672,7 @@ async fn run_telemt_core(
         rng.clone(),
         stats.clone(),
         api_me_pool.clone(),
+        me_ready_tx.clone(),
     )
     .await;
 
@@ -743,6 +746,7 @@ async fn run_telemt_core(
         api_config_tx.clone(),
         me_pool.clone(),
         shared_state.clone(),
+        me_ready_tx.clone(),
     )
     .await;
     let config_rx = runtime_watches.config_rx;
@@ -756,6 +760,7 @@ async fn run_telemt_core(
         route_runtime.clone(),
         &admission_tx,
         config_rx.clone(),
+        me_ready_rx,
     )
     .await;
     let _admission_tx_hold = admission_tx;

src/maestro/runtime_tasks.rs

@@ -53,6 +53,7 @@ pub(crate) async fn spawn_runtime_tasks(
     api_config_tx: watch::Sender<Arc<ProxyConfig>>,
     me_pool_for_policy: Option<Arc<MePool>>,
     shared_state: Arc<ProxySharedState>,
+    me_ready_tx: watch::Sender<u64>,
 ) -> RuntimeWatches {
     let um_clone = upstream_manager.clone();
     let dc_overrides_for_health = config.dc_overrides.clone();
@@ -250,12 +251,14 @@ pub(crate) async fn spawn_runtime_tasks(
         let pool_clone_sched = pool.clone();
         let rng_clone_sched = rng.clone();
         let config_rx_clone_sched = config_rx.clone();
+        let me_ready_tx_sched = me_ready_tx.clone();
         tokio::spawn(async move {
             crate::transport::middle_proxy::me_reinit_scheduler(
                 pool_clone_sched,
                 rng_clone_sched,
                 config_rx_clone_sched,
                 reinit_rx,
+                me_ready_tx_sched,
             )
             .await;
         });