perf: partially recovers the regression of request throughput on the per_worker policy

(cherry picked from commit 0aac44b)

Author: nyannyacha

crates/base/src/rt_worker/supervisor/strategy_per_worker.rs

@@ -21,7 +21,7 @@ pub async fn supervise(args: Arguments) -> ShutdownReason {
 
     let Timing {
         status: TimingStatus { demand, is_retired },
-        req: (mut req_start_rx, mut req_end_rx),
+        req: (_, mut req_end_rx),
     } = timing.unwrap_or_default();
 
     let is_retired = is_retired.unwrap();
@@ -51,12 +51,6 @@ pub async fn supervise(args: Arguments) -> ShutdownReason {
                 if !cpu_time_soft_limit_reached {
                     // retire worker
                     is_retired.raise();
-                    if let Some(tx) = pool_msg_tx.clone() {
-                        if tx.send(UserWorkerMsgs::Retire(key)).is_err() {
-                            error!("failed to send retire msg to pool: {:?}", key);
-                        }
-                    }
-
                     error!("CPU time soft limit reached. isolate: {:?}", key);
                     cpu_time_soft_limit_reached = true;
 
@@ -82,10 +76,6 @@ pub async fn supervise(args: Arguments) -> ShutdownReason {
                 }
             }
 
-            Some(notify) = req_start_rx.recv() => {
-                notify.notify_one();
-            }
-
             Some(_) = req_end_rx.recv() => {
                 req_ack_count += 1;
 
@@ -119,11 +109,6 @@ pub async fn supervise(args: Arguments) -> ShutdownReason {
                 } else if wall_clock_alerts == 1 {
                     // retire worker
                     is_retired.raise();
-                    if let Some(tx) = pool_msg_tx.clone() {
-                        if tx.send(UserWorkerMsgs::Retire(key)).is_err() {
-                            error!("failed to send retire msg to pool: {:?}", key);
-                        }
-                    }
                     error!("wall clock duration warning. isolate: {:?}", key);
                     wall_clock_alerts += 1;
                 } else {

crates/base/src/rt_worker/worker_ctx.rs

@@ -431,9 +431,6 @@ pub async fn create_user_worker_pool(
                 Some(UserWorkerMsgs::SendRequest(key, req, res_tx, conn_watch)) => {
                     worker_pool.send_request(&key, req, res_tx, conn_watch);
                 }
-                Some(UserWorkerMsgs::Retire(key)) => {
-                    worker_pool.retire(&key);
-                }
                 Some(UserWorkerMsgs::Idle(key)) => {
                     worker_pool.idle(&key);
                 }

crates/base/src/rt_worker/worker_pool.rs

@@ -134,7 +134,7 @@ impl ActiveWorkerRegistry {
         }
     }
 
-    fn mark_used_and_try_advance(&mut self) -> Option<&Uuid> {
+    fn mark_used_and_try_advance(&mut self, policy: SupervisorPolicy) -> Option<&Uuid> {
         if self.workers.is_empty() {
             let _ = self.next.take();
             return None;
@@ -147,25 +147,40 @@ impl ActiveWorkerRegistry {
             .unwrap_or(0);
 
         match self.workers.iter().nth(idx).cloned() {
-            Some(WorkerId(key, true)) => {
-                let key = self
-                    .workers
-                    .replace(WorkerId(key, false))
-                    .and_then(|WorkerId(ref key, _)| self.workers.get(key).map(|it| &it.0));
-
-                self.next = self.workers.iter().position(|it| it.1);
-                key
-            }
+            Some(WorkerId(key, true)) => match policy {
+                SupervisorPolicy::PerWorker => {
+                    self.next = Some(idx + 1);
+                    self.workers.get(&key).map(|it| &it.0)
+                }
+
+                SupervisorPolicy::PerRequest { .. } => {
+                    let key = self
+                        .workers
+                        .replace(WorkerId(key, false))
+                        .and_then(|WorkerId(ref key, _)| self.workers.get(key).map(|it| &it.0));
+
+                    self.next = self.workers.iter().position(|it| it.1);
+                    key
+                }
+            },
+
             _ => {
                 let _ = self.next.take();
                 None
             }
         }
     }
 
-    fn mark_idle(&mut self, key: &Uuid) {
-        if let Some(WorkerId(key, false)) = self.workers.get(key).cloned() {
-            let _ = self.workers.replace(WorkerId(key, true));
+    fn mark_idle(&mut self, key: &Uuid, policy: SupervisorPolicy) {
+        if let Some(WorkerId(key, mark)) = self.workers.get(key).cloned() {
+            if policy.is_per_request() {
+                if mark != false {
+                    return;
+                }
+
+                let _ = self.workers.replace(WorkerId(key, true));
+            }
+
             let (notify_tx, _) = self.notify_pair.clone();
             let _ = notify_tx.send(Some(key));
         }
@@ -418,7 +433,10 @@ impl WorkerPool {
             .entry(profile.service_path.clone())
             .or_insert_with(|| ActiveWorkerRegistry::new(self.policy.max_parallelism));
 
-        registry.workers.insert(WorkerId(key, false));
+        registry
+            .workers
+            .insert(WorkerId(key, self.policy.supervisor_policy.is_per_worker()));
+
         self.user_workers.insert(key, profile);
     }
 
@@ -431,21 +449,24 @@ impl WorkerPool {
     ) {
         let _: Result<(), Error> = match self.user_workers.get(key) {
             Some(worker) => {
+                let policy = self.policy.supervisor_policy;
                 let profile = worker.clone();
                 let cancel = worker.cancel.clone();
                 let (req_start_tx, req_end_tx) = profile.timing_tx_pair.clone();
 
                 // Create a closure to handle the request and send the response
                 let request_handler = async move {
-                    let fence = Arc::new(Notify::const_new());
+                    if !policy.is_per_worker() {
+                        let fence = Arc::new(Notify::const_new());
 
-                    if let Err(ex) = req_start_tx.send(fence.clone()) {
-                        error!("failed to notify the fence to the supervisor");
-                        return Err(ex)
-                            .with_context(|| "failed to notify the fence to the supervisor");
-                    }
+                        if let Err(ex) = req_start_tx.send(fence.clone()) {
+                            error!("failed to notify the fence to the supervisor");
+                            return Err(ex)
+                                .with_context(|| "failed to notify the fence to the supervisor");
+                        }
 
-                    fence.notified().await;
+                        fence.notified().await;
+                    }
 
                     let result = send_user_worker_request(
                         profile.worker_request_msg_tx,
@@ -487,26 +508,13 @@ impl WorkerPool {
         };
     }
 
-    pub fn retire(&mut self, key: &Uuid) {
-        if let Some(profile) = self.user_workers.get(key) {
-            let registry = self
-                .active_workers
-                .get_mut(&profile.service_path)
-                .expect("registry must be initialized at this point");
-
-            if registry.workers.contains(key) {
-                registry.workers.remove(key);
-            }
-        }
-    }
-
     pub fn idle(&mut self, key: &Uuid) {
         if let Some(registry) = self
             .user_workers
             .get_mut(key)
             .and_then(|it| self.active_workers.get_mut(&it.service_path))
         {
-            registry.mark_idle(key);
+            registry.mark_idle(key, self.policy.supervisor_policy);
         }
     }
 
@@ -525,6 +533,19 @@ impl WorkerPool {
         let _ = notify_tx.send(None);
     }
 
+    fn retire(&mut self, key: &Uuid) {
+        if let Some(profile) = self.user_workers.get(key) {
+            let registry = self
+                .active_workers
+                .get_mut(&profile.service_path)
+                .expect("registry must be initialized at this point");
+
+            if registry.workers.contains(key) {
+                registry.workers.remove(key);
+            }
+        }
+    }
+
     fn maybe_active_worker(&mut self, service_path: &String, force_create: bool) -> Option<Uuid> {
         if force_create {
             return None;
@@ -534,33 +555,35 @@ impl WorkerPool {
             return None;
         };
 
-        let mut advance_fn = move || registry.mark_used_and_try_advance().copied();
+        let policy = self.policy.supervisor_policy;
+        let mut advance_fn = move || registry.mark_used_and_try_advance(policy).copied();
 
-        if self.policy.supervisor_policy.is_per_request() {
+        if policy.is_per_request() {
             return advance_fn();
         }
 
-        loop {
-            let Some(worker_uuid) = advance_fn() else {
-                return None;
-            };
+        let Some(worker_uuid) = advance_fn() else {
+            return None;
+        };
 
-            match self
-                .user_workers
-                .get(&worker_uuid)
-                .and_then(|it| it.status.is_retired.as_ref())
-            {
-                Some(is_retired) if !is_retired.is_raised() => {
-                    self.user_workers
-                        .get(&worker_uuid)
-                        .map(|it| it.status.demand.as_ref())
-                        .unwrap()
-                        .fetch_add(1, Ordering::Release);
-
-                    return Some(worker_uuid);
-                }
+        match self
+            .user_workers
+            .get(&worker_uuid)
+            .and_then(|it| it.status.is_retired.as_ref())
+        {
+            Some(is_retired) if !is_retired.is_raised() => {
+                self.user_workers
+                    .get(&worker_uuid)
+                    .map(|it| it.status.demand.as_ref())
+                    .unwrap()
+                    .fetch_add(1, Ordering::Release);
+
+                return Some(worker_uuid);
+            }
 
-                _ => {}
+            _ => {
+                self.retire(&worker_uuid);
+                self.maybe_active_worker(service_path, force_create)
             }
         }
     }

crates/sb_workers/context.rs

@@ -141,7 +141,6 @@ pub enum UserWorkerMsgs {
         oneshot::Sender<Result<SendRequestResult, Error>>,
         Option<watch::Receiver<ConnSync>>,
     ),
-    Retire(Uuid),
     Idle(Uuid),
     Shutdown(Uuid),
 }