Idle tolerance + Adaptive floor by default + RPC Proxy Req

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/transport/middle_proxy/health.rs

@@ -18,6 +18,10 @@ const JITTER_FRAC_NUM: u64 = 2; // jitter up to 50% of backoff
 #[allow(dead_code)]
 const MAX_CONCURRENT_PER_DC_DEFAULT: usize = 1;
 const SHADOW_ROTATE_RETRY_SECS: u64 = 30;
+const IDLE_REFRESH_TRIGGER_BASE_SECS: u64 = 45;
+const IDLE_REFRESH_TRIGGER_JITTER_SECS: u64 = 5;
+const IDLE_REFRESH_RETRY_SECS: u64 = 8;
+const IDLE_REFRESH_SUCCESS_GUARD_SECS: u64 = 5;
 
 pub async fn me_health_monitor(pool: Arc<MePool>, rng: Arc<SecureRandom>, _min_connections: usize) {
     let mut backoff: HashMap<(i32, IpFamily), u64> = HashMap::new();
@@ -27,6 +31,7 @@ pub async fn me_health_monitor(pool: Arc<MePool>, rng: Arc<SecureRandom>, _min_c
     let mut outage_next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new();
     let mut single_endpoint_outage: HashSet<(i32, IpFamily)> = HashSet::new();
     let mut shadow_rotate_deadline: HashMap<(i32, IpFamily), Instant> = HashMap::new();
+    let mut idle_refresh_next_attempt: HashMap<(i32, IpFamily), Instant> = HashMap::new();
     let mut adaptive_idle_since: HashMap<(i32, IpFamily), Instant> = HashMap::new();
     let mut adaptive_recover_until: HashMap<(i32, IpFamily), Instant> = HashMap::new();
     loop {
@@ -43,6 +48,7 @@ pub async fn me_health_monitor(pool: Arc<MePool>, rng: Arc<SecureRandom>, _min_c
             &mut outage_next_attempt,
             &mut single_endpoint_outage,
             &mut shadow_rotate_deadline,
+            &mut idle_refresh_next_attempt,
             &mut adaptive_idle_since,
             &mut adaptive_recover_until,
         )
@@ -58,6 +64,7 @@ pub async fn me_health_monitor(pool: Arc<MePool>, rng: Arc<SecureRandom>, _min_c
             &mut outage_next_attempt,
             &mut single_endpoint_outage,
             &mut shadow_rotate_deadline,
+            &mut idle_refresh_next_attempt,
             &mut adaptive_idle_since,
             &mut adaptive_recover_until,
         )
@@ -76,6 +83,7 @@ async fn check_family(
     outage_next_attempt: &mut HashMap<(i32, IpFamily), Instant>,
     single_endpoint_outage: &mut HashSet<(i32, IpFamily)>,
     shadow_rotate_deadline: &mut HashMap<(i32, IpFamily), Instant>,
+    idle_refresh_next_attempt: &mut HashMap<(i32, IpFamily), Instant>,
     adaptive_idle_since: &mut HashMap<(i32, IpFamily), Instant>,
     adaptive_recover_until: &mut HashMap<(i32, IpFamily), Instant>,
 ) {
@@ -120,6 +128,7 @@ async fn check_family(
             .or_default()
             .push(writer.id);
     }
+    let writer_idle_since = pool.registry.writer_idle_since_snapshot().await;
 
     for (dc, endpoints) in dc_endpoints {
         if endpoints.is_empty() {
@@ -171,6 +180,7 @@ async fn check_family(
             outage_backoff.remove(&key);
             outage_next_attempt.remove(&key);
             shadow_rotate_deadline.remove(&key);
+            idle_refresh_next_attempt.remove(&key);
             adaptive_idle_since.remove(&key);
             adaptive_recover_until.remove(&key);
             info!(
@@ -184,6 +194,20 @@ async fn check_family(
         }
 
         if alive >= required {
+            maybe_refresh_idle_writer_for_dc(
+                pool,
+                rng,
+                key,
+                dc,
+                family,
+                &endpoints,
+                alive,
+                required,
+                &live_writer_ids_by_addr,
+                &writer_idle_since,
+                idle_refresh_next_attempt,
+            )
+            .await;
             maybe_rotate_single_endpoint_shadow(
                 pool,
                 rng,
@@ -287,6 +311,113 @@ async fn check_family(
     }
 }
 
+async fn maybe_refresh_idle_writer_for_dc(
+    pool: &Arc<MePool>,
+    rng: &Arc<SecureRandom>,
+    key: (i32, IpFamily),
+    dc: i32,
+    family: IpFamily,
+    endpoints: &[SocketAddr],
+    alive: usize,
+    required: usize,
+    live_writer_ids_by_addr: &HashMap<SocketAddr, Vec<u64>>,
+    writer_idle_since: &HashMap<u64, u64>,
+    idle_refresh_next_attempt: &mut HashMap<(i32, IpFamily), Instant>,
+) {
+    if alive < required {
+        return;
+    }
+
+    let now = Instant::now();
+    if let Some(next) = idle_refresh_next_attempt.get(&key)
+        && now < *next
+    {
+        return;
+    }
+
+    let now_epoch_secs = MePool::now_epoch_secs();
+    let mut candidate: Option<(u64, SocketAddr, u64, u64)> = None;
+    for endpoint in endpoints {
+        let Some(writer_ids) = live_writer_ids_by_addr.get(endpoint) else {
+            continue;
+        };
+        for writer_id in writer_ids {
+            let Some(idle_since_epoch_secs) = writer_idle_since.get(writer_id).copied() else {
+                continue;
+            };
+            let idle_age_secs = now_epoch_secs.saturating_sub(idle_since_epoch_secs);
+            let threshold_secs = IDLE_REFRESH_TRIGGER_BASE_SECS
+                + (*writer_id % (IDLE_REFRESH_TRIGGER_JITTER_SECS + 1));
+            if idle_age_secs < threshold_secs {
+                continue;
+            }
+            if candidate
+                .as_ref()
+                .map(|(_, _, age, _)| idle_age_secs > *age)
+                .unwrap_or(true)
+            {
+                candidate = Some((*writer_id, *endpoint, idle_age_secs, threshold_secs));
+            }
+        }
+    }
+
+    let Some((old_writer_id, endpoint, idle_age_secs, threshold_secs)) = candidate else {
+        return;
+    };
+
+    let rotate_ok = match tokio::time::timeout(pool.me_one_timeout, pool.connect_one(endpoint, rng.as_ref())).await {
+        Ok(Ok(())) => true,
+        Ok(Err(error)) => {
+            debug!(
+                dc = %dc,
+                ?family,
+                %endpoint,
+                old_writer_id,
+                idle_age_secs,
+                threshold_secs,
+                %error,
+                "Idle writer pre-refresh connect failed"
+            );
+            false
+        }
+        Err(_) => {
+            debug!(
+                dc = %dc,
+                ?family,
+                %endpoint,
+                old_writer_id,
+                idle_age_secs,
+                threshold_secs,
+                "Idle writer pre-refresh connect timed out"
+            );
+            false
+        }
+    };
+
+    if !rotate_ok {
+        idle_refresh_next_attempt.insert(key, now + Duration::from_secs(IDLE_REFRESH_RETRY_SECS));
+        return;
+    }
+
+    pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false)
+        .await;
+    idle_refresh_next_attempt.insert(
+        key,
+        now + Duration::from_secs(IDLE_REFRESH_SUCCESS_GUARD_SECS),
+    );
+    info!(
+        dc = %dc,
+        ?family,
+        %endpoint,
+        old_writer_id,
+        idle_age_secs,
+        threshold_secs,
+        alive,
+        required,
+        "Idle writer refreshed before upstream idle timeout"
+    );
+}
+
 async fn should_reduce_floor_for_idle(
     pool: &Arc<MePool>,
     key: (i32, IpFamily),