From f27a6235698cba843a9563851b81ed28b33844e1 Mon Sep 17 00:00:00 2001
From: Nyannyacha <meow@nnc.gg>
Date: Mon, 25 Nov 2024 01:02:05 +0000
Subject: [PATCH] perf(base): sent low memory notification when worker switches
 to retired state

---
 crates/base/src/rt_worker/supervisor/mod.rs              | 7 +++++++
 .../src/rt_worker/supervisor/strategy_per_request.rs     | 9 +++++++--
 .../base/src/rt_worker/supervisor/strategy_per_worker.rs | 8 ++++++--
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/crates/base/src/rt_worker/supervisor/mod.rs b/crates/base/src/rt_worker/supervisor/mod.rs
index c2b5308b2..ee3531e0c 100644
--- a/crates/base/src/rt_worker/supervisor/mod.rs
+++ b/crates/base/src/rt_worker/supervisor/mod.rs
@@ -14,6 +14,7 @@ use tokio::sync::{
     oneshot,
 };
 use tokio_util::sync::CancellationToken;
+use tracing::{debug, instrument};
 use uuid::Uuid;
 
 use crate::{
@@ -182,3 +183,9 @@ extern "C" fn v8_handle_wall_clock_beforeunload(
         );
     }
 }
+
+#[instrument(level = "debug", skip_all)]
+extern "C" fn v8_handle_early_retire(isolate: &mut v8::Isolate, _data: *mut std::ffi::c_void) {
+    isolate.low_memory_notification();
+    debug!("sent low mem notification");
+}
diff --git a/crates/base/src/rt_worker/supervisor/strategy_per_request.rs b/crates/base/src/rt_worker/supervisor/strategy_per_request.rs
index 56afb6ea8..6383edc2b 100644
--- a/crates/base/src/rt_worker/supervisor/strategy_per_request.rs
+++ b/crates/base/src/rt_worker/supervisor/strategy_per_request.rs
@@ -9,8 +9,9 @@ use sb_workers::context::{Timing, TimingStatus, UserWorkerMsgs};
 use tokio::time::Instant;
 
 use crate::rt_worker::supervisor::{
-    create_wall_clock_beforeunload_alert, v8_handle_termination, v8_handle_wall_clock_beforeunload,
-    wait_cpu_alarm, CPUUsage, CPUUsageMetrics, Tokens, V8HandleTerminationData,
+    create_wall_clock_beforeunload_alert, v8_handle_early_retire, v8_handle_termination,
+    v8_handle_wall_clock_beforeunload, wait_cpu_alarm, CPUUsage, CPUUsageMetrics, Tokens,
+    V8HandleTerminationData,
 };
 
 use super::Arguments;
@@ -46,6 +47,10 @@ pub async fn supervise(args: Arguments, oneshot: bool) -> (ShutdownReason, i64)
     let (_, hard_limit_ms) = cpu_timer_param.limits();
 
     let _guard = scopeguard::guard(is_retired, |v| {
+        if thread_safe_handle.request_interrupt(v8_handle_early_retire, std::ptr::null_mut()) {
+            waker.wake();
+        }
+
         v.raise();
     });
 
diff --git a/crates/base/src/rt_worker/supervisor/strategy_per_worker.rs b/crates/base/src/rt_worker/supervisor/strategy_per_worker.rs
index f52e5e051..9fbc65c0d 100644
--- a/crates/base/src/rt_worker/supervisor/strategy_per_worker.rs
+++ b/crates/base/src/rt_worker/supervisor/strategy_per_worker.rs
@@ -8,8 +8,8 @@ use log::error;
 use sb_workers::context::{Timing, TimingStatus, UserWorkerMsgs};
 
 use crate::rt_worker::supervisor::{
-    create_wall_clock_beforeunload_alert, v8_handle_wall_clock_beforeunload, wait_cpu_alarm,
-    CPUUsage, Tokens,
+    create_wall_clock_beforeunload_alert, v8_handle_early_retire,
+    v8_handle_wall_clock_beforeunload, wait_cpu_alarm, CPUUsage, Tokens,
 };
 
 use super::{v8_handle_termination, Arguments, CPUUsageMetrics, V8HandleTerminationData};
@@ -85,6 +85,10 @@ pub async fn supervise(args: Arguments) -> (ShutdownReason, i64) {
     );
 
     let early_retire_fn = || {
+        if thread_safe_handle.request_interrupt(v8_handle_early_retire, std::ptr::null_mut()) {
+            waker.wake();
+        }
+
         // we should raise a retire signal because subsequent incoming requests are unlikely to get
         // enough wall clock time or cpu time
         guard.raise();