[TS] Implement near-heap-limit termination (#4777)

# Description of Changes

Uses `Isolate::add_near_heap_limit_callback` to prevent unbounded heap
growth. Upon nearing the heap limit, we will now:
1. Call `Isolate::terminate_execution`.
2. Request that v8 double the heap limit.

Then, upon finishing a function call, we lower the heap limit back down.

This should hopefully fix the issue where v8 hits the heap limit and
crashes the whole process.

Also improves the way termination requests are checked for and
processed.

# Expected complexity level and risk

2

# Testing

- [ ] Manual testing with memory-leaky modules

Author: coolreader18

crates/core/src/config.rs

@@ -190,8 +190,12 @@ pub struct V8HeapPolicyConfig {
     pub heap_gc_trigger_fraction: f64,
     #[serde(default = "def_retire", deserialize_with = "de_fraction")]
     pub heap_retire_fraction: f64,
-    #[serde(default, rename = "heap-limit-mb", deserialize_with = "de_limit_mb")]
-    pub heap_limit_bytes: Option<usize>,
+    #[serde(
+        default = "def_heap_limit",
+        rename = "heap-limit-mb",
+        deserialize_with = "de_limit_mb"
+    )]
+    pub heap_limit_bytes: usize,
 }
 
 impl Default for V8HeapPolicyConfig {
@@ -201,7 +205,7 @@ impl Default for V8HeapPolicyConfig {
             heap_check_time_interval: def_time_interval(),
             heap_gc_trigger_fraction: def_gc_trigger(),
             heap_retire_fraction: def_retire(),
-            heap_limit_bytes: None,
+            heap_limit_bytes: def_heap_limit(),
         }
     }
 }
@@ -242,6 +246,12 @@ fn def_retire() -> f64 {
     0.75
 }
 
+/// Default heap limit, in bytes
+fn def_heap_limit() -> usize {
+    // 1 GiB
+    1024 * 1024 * 1024
+}
+
 fn de_nz_u64<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
 where
     D: serde::Deserializer<'de>,
@@ -294,22 +304,20 @@ where
     }
 }
 
-fn de_limit_mb<'de, D>(deserializer: D) -> Result<Option<usize>, D::Error>
+fn de_limit_mb<'de, D>(deserializer: D) -> Result<usize, D::Error>
 where
     D: serde::Deserializer<'de>,
 {
     let value = u64::deserialize(deserializer)?;
     if value == 0 {
-        return Ok(None);
+        return Ok(def_heap_limit());
     }
 
     let bytes = value
         .checked_mul(1024 * 1024)
         .ok_or_else(|| serde::de::Error::custom("heap-limit-mb is too large"))?;
 
-    usize::try_from(bytes)
-        .map(Some)
-        .map_err(|_| serde::de::Error::custom("heap-limit-mb does not fit in usize"))
+    usize::try_from(bytes).map_err(|_| serde::de::Error::custom("heap-limit-mb does not fit in usize"))
 }
 
 #[cfg(test)]
@@ -441,7 +449,7 @@ mod tests {
         );
         assert_eq!(config.v8_heap_policy.heap_gc_trigger_fraction, 0.67);
         assert_eq!(config.v8_heap_policy.heap_retire_fraction, 0.75);
-        assert_eq!(config.v8_heap_policy.heap_limit_bytes, None);
+        assert_eq!(config.v8_heap_policy.heap_limit_bytes, 1024 * 1024 * 1024);
     }
 
     #[test]
@@ -464,6 +472,6 @@ mod tests {
         );
         assert_eq!(config.v8_heap_policy.heap_gc_trigger_fraction, 0.6);
         assert_eq!(config.v8_heap_policy.heap_retire_fraction, 0.8);
-        assert_eq!(config.v8_heap_policy.heap_limit_bytes, Some(256 * 1024 * 1024));
+        assert_eq!(config.v8_heap_policy.heap_limit_bytes, 256 * 1024 * 1024);
     }
 }

crates/core/src/host/v8/budget.rs

@@ -69,6 +69,7 @@ pub(super) extern "C" fn cb_noop(_: v8::UnsafeRawIsolatePtr, _: *mut c_void) {}
 ///
 /// Every `callback_every` ticks, `callback` is called.
 fn run_timeout_and_cb_every(
+    // TODO: use RemoteTerminator here once we actually call this function, and make RemoteTerminator thread-safe.
     handle: IsolateHandle,
     callback_every: u64,
     callback: InterruptCallback,

crates/core/src/host/v8/error.rs

@@ -13,7 +13,9 @@ use core::fmt;
 use spacetimedb_data_structures::map::IntMap;
 use spacetimedb_primitives::errno;
 use spacetimedb_sats::Serialize;
+use std::cell::Cell;
 use std::num::NonZero;
+use std::rc::Rc;
 use v8::{tc_scope, Exception, HandleScope, Local, PinScope, PinnedRef, StackFrame, StackTrace, TryCatch, Value};
 
 /// The result of trying to convert a [`Value`] in scope `'scope` to some type `T`.
@@ -107,24 +109,6 @@ impl ArrayTooLongError {
     }
 }
 
-/// A catchable termination error thrown in callbacks to indicate a host error.
-#[derive(Serialize)]
-pub(super) struct TerminationError {
-    __terminated__: String,
-}
-
-impl TerminationError {
-    /// Converts [`anyhow::Error`] to a termination error.
-    pub(super) fn from_error<'scope>(
-        scope: &PinScope<'scope, '_>,
-        error: &anyhow::Error,
-    ) -> ExcResult<ExceptionValue<'scope>> {
-        let __terminated__ = format!("{error}");
-        let error = Self { __terminated__ };
-        serialize_to_js(scope, &error).map(ExceptionValue)
-    }
-}
-
 /// Collapses `res` where the `Ok(x)` where `x` is throwable.
 pub(super) fn collapse_exc_thrown<'scope>(
     scope: &PinScope<'scope, '_>,
@@ -159,38 +143,96 @@ pub type SysCallResult<T> = Result<T, SysCallError>;
 /// A flag set in [`throw_nodes_error`].
 /// The flag should be checked in every module -> host ABI.
 /// If the flag is set, the call is prevented.
-struct TerminationFlag;
+#[derive(Default, Clone)]
+pub(super) struct TerminationFlag(Rc<TerminationFlagInner>);
+
+#[derive(Default)]
+struct TerminationFlagInner {
+    flag: Cell<bool>,
+    reason: Cell<Option<anyhow::Error>>,
+}
+
+impl TerminationFlag {
+    /// Set the terminate-execution flag due to the given host error.
+    pub(super) fn set(&self, err: anyhow::Error) {
+        let already_set = self.0.flag.replace(true);
+        if !already_set {
+            self.0.reason.set(Some(err));
+        }
+    }
+
+    pub(super) fn is_set(&self) -> bool {
+        self.0.flag.get()
+    }
+
+    pub(super) fn clear(&self) -> Option<anyhow::Error> {
+        let flag = self.0.flag.take();
+        let reason = self.0.reason.take();
+        flag.then_some(()).and(reason)
+    }
+}
 
 /// Terminate execution immediately due to the given host error.
 pub(super) fn terminate_execution<'scope>(
     scope: &mut PinScope<'scope, '_>,
-    err: &anyhow::Error,
+    err: anyhow::Error,
 ) -> ExcResult<ExceptionValue<'scope>> {
-    // Terminate execution ASAP and throw a catchable exception (`TerminationError`).
-    // Unfortunately, JS execution won't be terminated once the callback returns,
-    // so we set a slot that all callbacks immediately check
-    // to ensure that the module won't be able to do anything to the host
-    // while it's being terminated (eventually).
-    scope.terminate_execution();
-    scope.set_slot(TerminationFlag);
-    TerminationError::from_error(scope, err)
+    get_or_insert_slot(scope, TerminationFlag::default).set(err);
+    Err(terminate_execution_now(scope))
 }
 
-/// Checks the termination flag and throws a `TerminationError` if set.
-///
-/// Returns whether the flag was set.
-pub(super) fn throw_if_terminated(scope: &PinScope<'_, '_>) -> bool {
+/// A handle to terminate execution of a v8 isolate.
+pub(super) struct RemoteTerminator {
+    pub(super) flag: TerminationFlag,
+    pub(super) handle: v8::IsolateHandle,
+}
+
+impl RemoteTerminator {
+    pub(super) fn new(isolate: &mut v8::Isolate) -> Self {
+        let flag = get_or_insert_slot(isolate, TerminationFlag::default).clone();
+        let handle = isolate.thread_safe_handle();
+        Self { flag, handle }
+    }
+
+    /// Request that v8 terminate execution due to the given host error.
+    pub(super) fn terminate_execution(&self, err: anyhow::Error) {
+        self.flag.set(err);
+        self.handle.terminate_execution();
+    }
+}
+
+/// Checks the termination flag and terminates execution immediately if it was set.
+pub(super) fn check_termination(scope: &PinScope<'_, '_>) -> ExcResult<()> {
     // If the flag was set in `throw_nodes_error`,
     // we need to block all module -> host ABI calls.
-    let set = scope.get_slot::<TerminationFlag>().is_some();
-    if set {
-        let err = anyhow::anyhow!("execution is being terminated");
-        if let Ok(exception) = TerminationError::from_error(scope, &err) {
-            exception.throw(scope);
-        }
+    if let Some(flag) = scope.get_slot::<TerminationFlag>()
+        && flag.is_set()
+    {
+        return Err(terminate_execution_now(scope));
     }
 
-    set
+    Ok(())
+}
+
+/// Terminate execution of the v8 isolate, not as a request, but right now.
+fn terminate_execution_now(scope: &PinScope<'_, '_>) -> ExceptionThrown {
+    if !scope.is_execution_terminating() {
+        scope.terminate_execution();
+        handle_interrupts(scope);
+        assert!(scope.is_execution_terminating());
+    }
+    exception_already_thrown()
+}
+
+/// Check for interrupts, such as an execution termination request, and handle then.
+fn handle_interrupts(scope: &PinScope<'_, '_>) {
+    // This is stupid. `Isolate::TerminateExecution` requests a termination interrupt, meaning that
+    // `isolate.terminate_execution(); isolate.is_execution_terminating()` evaluates to false.
+    // v8 exposes neither the internal, immediate version `i::Isolate::TerminateExecution`, nor
+    // `Isolate::HandleInterrupts`, so we have to call `HandleInterrupts` in a roundabout way:
+    // via `JSON::Stringify`, which checks at the start of the call for any interrupts. Yippee.
+    let obj = v8::Integer::new(scope, 0);
+    let _ = v8::json::stringify(scope, obj.into());
 }
 
 /// A catchable error code thrown in callbacks
@@ -218,7 +260,7 @@ pub(crate) struct ExceptionThrown {
 
 impl ExceptionThrown {
     /// Turns a caught JS exception in `scope` into a [`JSError`].
-    pub(crate) fn into_error(self, scope: &mut PinTryCatch) -> JsError {
+    pub(crate) fn into_error(self, scope: &mut PinTryCatch) -> Result<JsError, UnknownJsError> {
         JsError::from_caught(scope)
     }
 }
@@ -255,12 +297,15 @@ pub(super) enum ErrorOrException<Exc> {
     Exception(Exc),
 }
 
-impl<Exc> ErrorOrException<Exc> {
-    pub(super) fn map_exception<Exc2>(self, f: impl FnOnce(Exc) -> Exc2) -> ErrorOrException<Exc2> {
-        match self {
+impl ErrorOrException<ExceptionThrown> {
+    pub(super) fn exc_into_error(
+        self,
+        scope: &mut PinTryCatch<'_, '_, '_, '_>,
+    ) -> Result<ErrorOrException<JsError>, UnknownJsError> {
+        Ok(match self {
             ErrorOrException::Err(e) => ErrorOrException::Err(e),
-            ErrorOrException::Exception(exc) => ErrorOrException::Exception(f(exc)),
-        }
+            ErrorOrException::Exception(exc) => ErrorOrException::Exception(exc.into_error(scope)?),
+        })
     }
 }
 
@@ -276,6 +321,12 @@ impl From<ExceptionThrown> for ErrorOrException<ExceptionThrown> {
     }
 }
 
+impl From<JsError> for ErrorOrException<JsError> {
+    fn from(e: JsError) -> Self {
+        Self::Exception(e)
+    }
+}
+
 impl From<ErrorOrException<JsError>> for anyhow::Error {
     fn from(err: ErrorOrException<JsError>) -> Self {
         match err {
@@ -508,23 +559,41 @@ fn get_or_insert_slot<T: 'static>(isolate: &mut v8::Isolate, default: impl FnOnc
 
 impl JsError {
     /// Turns a caught JS exception in `scope` into a [`JSError`].
-    fn from_caught(scope: &mut PinTryCatch<'_, '_, '_, '_>) -> Self {
-        match scope.message() {
-            Some(message) => Self {
-                trace: message
-                    .get_stack_trace(scope)
-                    .map(|trace| JsStackTrace::from_trace(scope, trace))
-                    .unwrap_or_default(),
-                msg: message.get(scope).to_rust_string_lossy(scope),
-            },
-            None => Self {
-                trace: JsStackTrace::default(),
-                msg: "unknown error".to_owned(),
-            },
+    fn from_caught(scope: &mut PinTryCatch<'_, '_, '_, '_>) -> Result<Self, UnknownJsError> {
+        let message = scope.message().ok_or(UnknownJsError)?;
+        Ok(Self {
+            trace: message
+                .get_stack_trace(scope)
+                .map(|trace| JsStackTrace::from_trace(scope, trace))
+                .unwrap_or_default(),
+            msg: message.get(scope).to_rust_string_lossy(scope),
+        })
+    }
+}
+
+pub(super) struct UnknownJsError;
+
+impl From<UnknownJsError> for JsError {
+    fn from(_: UnknownJsError) -> Self {
+        Self {
+            trace: JsStackTrace::default(),
+            msg: "unknown error".to_owned(),
         }
     }
 }
 
+impl From<UnknownJsError> for ErrorOrException<JsError> {
+    fn from(e: UnknownJsError) -> Self {
+        Self::Exception(e.into())
+    }
+}
+
+impl From<UnknownJsError> for anyhow::Error {
+    fn from(e: UnknownJsError) -> Self {
+        JsError::from(e).into()
+    }
+}
+
 pub(super) fn log_traceback(replica_ctx: &ReplicaContext, func_type: &str, func: &str, e: &anyhow::Error) {
     // no need to log `JsError` separately; it'll be displayed if it exists in the error.
     log::info!("{func_type} \"{func}\" raised a runtime error: {e:#}");
@@ -553,7 +622,7 @@ pub(super) fn catch_exception<'scope, T>(
     body: impl FnOnce(&mut PinTryCatch<'scope, '_, '_, '_>) -> Result<T, ErrorOrException<ExceptionThrown>>,
 ) -> Result<T, ErrorOrException<JsError>> {
     tc_scope!(scope, scope);
-    body(scope).map_err(|e| e.map_exception(|exc| exc.into_error(scope)))
+    body(scope).map_err(|e| e.exc_into_error(scope).unwrap_or_else(Into::into))
 }
 
 pub(super) type PinTryCatch<'scope, 'iso, 'x, 's> = PinnedRef<'x, TryCatch<'s, 'scope, HandleScope<'iso>>>;