perf(base): sent low memory notification when worker switches to reti…

…red state (#450)
supabase · Nov 25, 2024 · 388d2ea · 388d2ea
1 parent 13ee2a2
commit 388d2ea
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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();