HOL Minimized + Random conn_id + Target DC Magics

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

src/transport/middle_proxy/send.rs

@@ -6,7 +6,7 @@ use tokio::sync::Mutex;
 use tracing::{debug, warn};
 
 use crate::error::{ProxyError, Result};
-use crate::protocol::constants::RPC_CLOSE_EXT_U32;
+use crate::protocol::constants::{RPC_CLOSE_EXT_U32, TG_MIDDLE_PROXIES_V4};
 
 use super::MePool;
 use super::codec::RpcWriter;
@@ -16,6 +16,7 @@ impl MePool {
     pub async fn send_proxy_req(
         &self,
         conn_id: u64,
+        target_dc: i16,
         client_addr: SocketAddr,
         our_addr: SocketAddr,
         data: &[u8],
@@ -35,14 +36,20 @@ impl MePool {
             if ws.is_empty() {
                 return Err(ProxyError::Proxy("All ME connections dead".into()));
             }
-            let writers: Vec<Arc<Mutex<RpcWriter>>> = ws.iter().cloned().collect();
-            let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % writers.len();
+            let writers: Vec<(SocketAddr, Arc<Mutex<RpcWriter>>)> = ws.iter().cloned().collect();
             drop(ws);
 
+            let candidate_indices = candidate_indices_for_dc(&writers, target_dc);
+            if candidate_indices.is_empty() {
+                return Err(ProxyError::Proxy("No ME writers available for target DC".into()));
+            }
+            let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len();
+
             // Prefer immediately available writer to avoid waiting on stalled connection.
-            for offset in 0..writers.len() {
-                let idx = (start + offset) % writers.len();
-                let w = writers[idx].clone();
+            for offset in 0..candidate_indices.len() {
+                let cidx = (start + offset) % candidate_indices.len();
+                let idx = candidate_indices[cidx];
+                let w = writers[idx].1.clone();
                 if let Ok(mut guard) = w.try_lock() {
                     let send_res = guard.send(&payload).await;
                     drop(guard);
@@ -51,7 +58,7 @@ impl MePool {
                         Err(e) => {
                             warn!(error = %e, "ME write failed, removing dead conn");
                             let mut ws = self.writers.write().await;
-                            ws.retain(|o| !Arc::ptr_eq(o, &w));
+                            ws.retain(|(_, o)| !Arc::ptr_eq(o, &w));
                             if ws.is_empty() {
                                 return Err(ProxyError::Proxy("All ME connections dead".into()));
                             }
@@ -62,13 +69,13 @@ impl MePool {
             }
 
             // All writers are currently busy, wait for the selected one.
-            let w = writers[start].clone();
+            let w = writers[candidate_indices[start]].1.clone();
             match w.lock().await.send(&payload).await {
                 Ok(()) => return Ok(()),
                 Err(e) => {
                     warn!(error = %e, "ME write failed, removing dead conn");
                     let mut ws = self.writers.write().await;
-                    ws.retain(|o| !Arc::ptr_eq(o, &w));
+                    ws.retain(|(_, o)| !Arc::ptr_eq(o, &w));
                     if ws.is_empty() {
                         return Err(ProxyError::Proxy("All ME connections dead".into()));
                     }
@@ -80,7 +87,7 @@ impl MePool {
     pub async fn send_close(&self, conn_id: u64) -> Result<()> {
         let ws = self.writers.read().await;
         if !ws.is_empty() {
-            let w = ws[0].clone();
+            let w = ws[0].1.clone();
             drop(ws);
             let mut p = Vec::with_capacity(12);
             p.extend_from_slice(&RPC_CLOSE_EXT_U32.to_le_bytes());
@@ -88,7 +95,7 @@ impl MePool {
             if let Err(e) = w.lock().await.send(&p).await {
                 debug!(error = %e, "ME close write failed");
                 let mut ws = self.writers.write().await;
-                ws.retain(|o| !Arc::ptr_eq(o, &w));
+                ws.retain(|(_, o)| !Arc::ptr_eq(o, &w));
             }
         }
 
@@ -100,3 +107,40 @@ impl MePool {
         self.writers.try_read().map(|w| w.len()).unwrap_or(0)
     }
 }
+
+fn candidate_indices_for_dc(
+    writers: &[(SocketAddr, Arc<Mutex<RpcWriter>>)],
+    target_dc: i16,
+) -> Vec<usize> {
+    let mut preferred = Vec::<SocketAddr>::new();
+    let key = target_dc as i32;
+    if let Some(v) = TG_MIDDLE_PROXIES_V4.get(&key) {
+        preferred.extend(v.iter().map(|(ip, port)| SocketAddr::new(*ip, *port)));
+    }
+    if preferred.is_empty() {
+        let abs = key.abs();
+        if let Some(v) = TG_MIDDLE_PROXIES_V4.get(&abs) {
+            preferred.extend(v.iter().map(|(ip, port)| SocketAddr::new(*ip, *port)));
+        }
+    }
+    if preferred.is_empty() {
+        let abs = key.abs();
+        if let Some(v) = TG_MIDDLE_PROXIES_V4.get(&-abs) {
+            preferred.extend(v.iter().map(|(ip, port)| SocketAddr::new(*ip, *port)));
+        }
+    }
+    if preferred.is_empty() {
+        return (0..writers.len()).collect();
+    }
+
+    let mut out = Vec::new();
+    for (idx, (addr, _)) in writers.iter().enumerate() {
+        if preferred.iter().any(|p| p == addr) {
+            out.push(idx);
+        }
+    }
+    if out.is_empty() {
+        return (0..writers.len()).collect();
+    }
+    out
+}