Process-wide concurrency + Cancellation ownership fixes

src/transport/middle_proxy/send/selection.rs

@@ -1,5 +1,6 @@
 use std::cmp::Reverse;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
+use std::net::SocketAddr;
 use std::sync::atomic::Ordering;
 
 use super::super::MePool;
@@ -10,6 +11,18 @@ use super::{
 };
 use crate::config::MeWriterPickMode;
 
+struct WriterSelectionKey {
+    index: usize,
+    contour_rank: usize,
+    stale: usize,
+    degraded: usize,
+    idle_rank: usize,
+    queue_remaining: usize,
+    addr: SocketAddr,
+    id: u64,
+    pick_score: u64,
+}
+
 impl MePool {
     pub(super) async fn candidate_indices_for_dc(
         &self,
@@ -72,24 +85,20 @@ impl MePool {
         }
     }
 
-    pub(super) fn writer_contour_rank_for_selection(
-        &self,
-        writer: &super::super::pool::MeWriter,
-    ) -> usize {
-        match WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) {
+    fn writer_contour_rank_for_selection(contour: WriterContour) -> usize {
+        match contour {
             WriterContour::Active => 0,
             WriterContour::Warm => 1,
             WriterContour::Draining => 2,
         }
     }
 
-    pub(super) fn writer_idle_rank_for_selection(
-        &self,
-        writer: &super::super::pool::MeWriter,
+    fn writer_idle_rank_for_selection(
+        writer_id: u64,
         idle_since_by_writer: &HashMap<u64, u64>,
         now_epoch_secs: u64,
     ) -> usize {
-        let Some(idle_since) = idle_since_by_writer.get(&writer.id).copied() else {
+        let Some(idle_since) = idle_since_by_writer.get(&writer_id).copied() else {
             return 0;
         };
         let idle_age_secs = now_epoch_secs.saturating_sub(idle_since);
@@ -102,51 +111,67 @@ impl MePool {
         }
     }
 
-    pub(super) fn writer_pick_score(
+    fn capture_writer_selection_key(
         &self,
+        index: usize,
         writer: &super::super::pool::MeWriter,
         idle_since_by_writer: &HashMap<u64, u64>,
         now_epoch_secs: u64,
-    ) -> u64 {
-        let contour_penalty = match WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) {
+        current_generation: u64,
+    ) -> WriterSelectionKey {
+        let contour = WriterContour::from_u8(writer.contour.load(Ordering::Relaxed));
+        let contour_rank = Self::writer_contour_rank_for_selection(contour);
+        let contour_penalty = match contour {
             WriterContour::Active => 0,
             WriterContour::Warm => PICK_PENALTY_WARM,
             WriterContour::Draining => PICK_PENALTY_DRAINING,
         };
-        let stale_penalty = if writer.generation < self.current_generation() {
+        let stale = (writer.generation < current_generation) as usize;
+        let stale_penalty = if stale != 0 {
             PICK_PENALTY_STALE
         } else {
             0
         };
-        let degraded_penalty = if writer.degraded.load(Ordering::Relaxed) {
+        let degraded = writer.degraded.load(Ordering::Relaxed) as usize;
+        let degraded_penalty = if degraded != 0 {
             PICK_PENALTY_DEGRADED
         } else {
             0
         };
-        let idle_penalty =
-            (self.writer_idle_rank_for_selection(writer, idle_since_by_writer, now_epoch_secs)
-                as u64)
-                * 100;
+        let idle_rank =
+            Self::writer_idle_rank_for_selection(writer.id, idle_since_by_writer, now_epoch_secs);
+        let idle_penalty = (idle_rank as u64) * 100;
         let queue_cap = self.writer_lifecycle.writer_cmd_channel_capacity.max(1) as u64;
-        let queue_remaining = writer.tx.capacity() as u64;
-        let queue_used = queue_cap.saturating_sub(queue_remaining.min(queue_cap));
+        let queue_remaining = writer.tx.capacity();
+        let queue_used = queue_cap.saturating_sub((queue_remaining as u64).min(queue_cap));
         let queue_util_pct = queue_used.saturating_mul(100) / queue_cap;
         let queue_penalty = queue_util_pct.saturating_mul(4);
         let rtt_penalty =
             ((writer.rtt_ema_ms_x10.load(Ordering::Relaxed) as u64).saturating_add(5) / 10)
                 .min(400);
 
-        contour_penalty
+        let pick_score = contour_penalty
             .saturating_add(stale_penalty)
             .saturating_add(degraded_penalty)
             .saturating_add(idle_penalty)
             .saturating_add(queue_penalty)
-            .saturating_add(rtt_penalty)
+            .saturating_add(rtt_penalty);
+        WriterSelectionKey {
+            index,
+            contour_rank,
+            stale,
+            degraded,
+            idle_rank,
+            queue_remaining,
+            addr: writer.addr,
+            id: writer.id,
+            pick_score,
+        }
     }
 
     pub(super) fn p2c_ordered_candidate_indices(
         &self,
-        candidate_indices: &[usize],
+        mut candidate_indices: Vec<usize>,
         writers_snapshot: &[super::super::pool::MeWriter],
         idle_since_by_writer: &HashMap<u64, u64>,
         now_epoch_secs: u64,
@@ -158,33 +183,26 @@ impl MePool {
             return Vec::new();
         }
 
-        let mut sampled = Vec::<usize>::with_capacity(sample_size.min(total));
-        let mut seen = HashSet::<usize>::with_capacity(total);
-        for offset in 0..sample_size.min(total) {
-            let idx = candidate_indices[(start + offset) % total];
-            if seen.insert(idx) {
-                sampled.push(idx);
-            }
+        candidate_indices.rotate_left(start % total);
+        let current_generation = self.current_generation();
+        let sample_size = sample_size.min(total);
+        let mut sampled = candidate_indices[..sample_size]
+            .iter()
+            .map(|idx| {
+                self.capture_writer_selection_key(
+                    *idx,
+                    &writers_snapshot[*idx],
+                    idle_since_by_writer,
+                    now_epoch_secs,
+                    current_generation,
+                )
+            })
+            .collect::<Vec<_>>();
+        sampled.sort_by_key(|candidate| (candidate.pick_score, candidate.addr, candidate.id));
+        for (target, candidate) in candidate_indices[..sample_size].iter_mut().zip(sampled) {
+            *target = candidate.index;
         }
-
-        sampled.sort_by_key(|idx| {
-            let writer = &writers_snapshot[*idx];
-            (
-                self.writer_pick_score(writer, idle_since_by_writer, now_epoch_secs),
-                writer.addr,
-                writer.id,
-            )
-        });
-
-        let mut ordered = Vec::<usize>::with_capacity(total);
-        ordered.extend(sampled.iter().copied());
-        for offset in 0..total {
-            let idx = candidate_indices[(start + offset) % total];
-            if seen.insert(idx) {
-                ordered.push(idx);
-            }
-        }
-        ordered
+        candidate_indices
     }
 
     pub(super) async fn ordered_candidate_indices(
@@ -206,7 +224,7 @@ impl MePool {
         let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len();
         if pick_mode == MeWriterPickMode::P2c {
             return self.p2c_ordered_candidate_indices(
-                &candidate_indices,
+                candidate_indices,
                 writers_snapshot,
                 &writer_idle_since,
                 now_epoch_secs,
@@ -220,60 +238,65 @@ impl MePool {
             .me_deterministic_writer_sort
             .load(Ordering::Relaxed)
         {
-            candidate_indices.sort_by(|lhs, rhs| {
-                let left = &writers_snapshot[*lhs];
-                let right = &writers_snapshot[*rhs];
-                let left_key = (
-                    self.writer_contour_rank_for_selection(left),
-                    (left.generation < self.current_generation()) as usize,
-                    left.degraded.load(Ordering::Relaxed) as usize,
-                    self.writer_idle_rank_for_selection(
-                        left,
+            let current_generation = self.current_generation();
+            let mut captured = candidate_indices
+                .iter()
+                .map(|idx| {
+                    self.capture_writer_selection_key(
+                        *idx,
+                        &writers_snapshot[*idx],
                         &writer_idle_since,
                         now_epoch_secs,
-                    ),
-                    Reverse(left.tx.capacity()),
-                    left.addr,
-                    left.id,
-                );
-                let right_key = (
-                    self.writer_contour_rank_for_selection(right),
-                    (right.generation < self.current_generation()) as usize,
-                    right.degraded.load(Ordering::Relaxed) as usize,
-                    self.writer_idle_rank_for_selection(
-                        right,
-                        &writer_idle_since,
-                        now_epoch_secs,
-                    ),
-                    Reverse(right.tx.capacity()),
-                    right.addr,
-                    right.id,
-                );
-                left_key.cmp(&right_key)
-            });
-        } else {
-            candidate_indices.sort_by_key(|idx| {
-                let writer = &writers_snapshot[*idx];
-                let degraded = writer.degraded.load(Ordering::Relaxed);
-                let stale = (writer.generation < self.current_generation()) as usize;
+                        current_generation,
+                    )
+                })
+                .collect::<Vec<_>>();
+            captured.sort_by_key(|candidate| {
                 (
-                    self.writer_contour_rank_for_selection(writer),
-                    stale,
-                    degraded as usize,
-                    self.writer_idle_rank_for_selection(
-                        writer,
-                        &writer_idle_since,
-                        now_epoch_secs,
-                    ),
-                    Reverse(writer.tx.capacity()),
+                    candidate.contour_rank,
+                    candidate.stale,
+                    candidate.degraded,
+                    candidate.idle_rank,
+                    Reverse(candidate.queue_remaining),
+                    candidate.addr,
+                    candidate.id,
                 )
             });
+            for (target, candidate) in candidate_indices.iter_mut().zip(captured) {
+                *target = candidate.index;
+            }
+        } else {
+            let current_generation = self.current_generation();
+            let mut captured = candidate_indices
+                .iter()
+                .map(|idx| {
+                    self.capture_writer_selection_key(
+                        *idx,
+                        &writers_snapshot[*idx],
+                        &writer_idle_since,
+                        now_epoch_secs,
+                        current_generation,
+                    )
+                })
+                .collect::<Vec<_>>();
+            captured.sort_by_key(|candidate| {
+                (
+                    candidate.contour_rank,
+                    candidate.stale,
+                    candidate.degraded,
+                    candidate.idle_rank,
+                    Reverse(candidate.queue_remaining),
+                )
+            });
+            for (target, candidate) in candidate_indices.iter_mut().zip(captured) {
+                *target = candidate.index;
+            }
         }
 
-        let mut ordered = Vec::<usize>::with_capacity(candidate_indices.len());
-        for offset in 0..candidate_indices.len() {
-            ordered.push(candidate_indices[(start + offset) % candidate_indices.len()]);
+        if !candidate_indices.is_empty() {
+            let len = candidate_indices.len();
+            candidate_indices.rotate_left(start % len);
         }
-        ordered
+        candidate_indices
     }
 }