1 files changed, 136 insertions, 62 deletions

diff --git a/core/realm.rs b/core/realm.rs
index 375f74088..fc0ff4b4b 100644
--- a/core/realm.rs
+++ b/core/realm.rs
@@ -4,8 +4,11 @@ use crate::bindings;
 use crate::modules::ModuleCode;
 use crate::ops::OpCtx;
 use crate::runtime::exception_to_err_result;
+use crate::runtime::JsRuntimeState;
 use crate::JsRuntime;
+use crate::OpCall;
 use anyhow::Error;
+use futures::stream::FuturesUnordered;
 use std::cell::RefCell;
 use std::collections::HashSet;
 use std::collections::VecDeque;
@@ -45,9 +48,11 @@ pub(crate) struct ContextState {
   pub(crate) pending_promise_rejections:
     VecDeque<(v8::Global<v8::Promise>, v8::Global<v8::Value>)>,
   pub(crate) unrefed_ops: HashSet<i32, BuildHasherDefault<IdentityHasher>>,
+  pub(crate) pending_ops: FuturesUnordered<OpCall>,
   // We don't explicitly re-read this prop but need the slice to live alongside
   // the context
   pub(crate) op_ctxs: Box<[OpCtx]>,
+  pub(crate) isolate: Option<*mut v8::OwnedIsolate>,
 }
 
 /// A representation of a JavaScript realm tied to a [`JsRuntime`], that allows
@@ -95,28 +100,110 @@ pub(crate) struct ContextState {
 /// keep the underlying V8 context alive even if it would have otherwise been
 /// garbage collected.
 #[derive(Clone)]
-pub struct JsRealm(Rc<v8::Global<v8::Context>>);
-impl JsRealm {
-  pub fn new(context: v8::Global<v8::Context>) -> Self {
-    JsRealm(Rc::new(context))
+#[repr(transparent)]
+pub struct JsRealm(pub(crate) JsRealmInner);
+
+#[derive(Clone)]
+pub(crate) struct JsRealmInner {
+  context_state: Rc<RefCell<ContextState>>,
+  context: Rc<v8::Global<v8::Context>>,
+  runtime_state: Rc<RefCell<JsRuntimeState>>,
+  is_global: bool,
+}
+
+impl JsRealmInner {
+  pub(crate) fn new(
+    context_state: Rc<RefCell<ContextState>>,
+    context: v8::Global<v8::Context>,
+    runtime_state: Rc<RefCell<JsRuntimeState>>,
+    is_global: bool,
+  ) -> Self {
+    Self {
+      context_state,
+      context: context.into(),
+      runtime_state,
+      is_global,
+    }
+  }
+
+  pub fn num_pending_ops(&self) -> usize {
+    self.context_state.borrow().pending_ops.len()
+  }
+
+  pub fn num_unrefed_ops(&self) -> usize {
+    self.context_state.borrow().unrefed_ops.len()
   }
 
   #[inline(always)]
   pub fn context(&self) -> &v8::Global<v8::Context> {
-    &self.0
+    &self.context
+  }
+
+  #[inline(always)]
+  pub(crate) fn state(&self) -> Rc<RefCell<ContextState>> {
+    self.context_state.clone()
   }
 
+  /// For info on the [`v8::Isolate`] parameter, check [`JsRealm#panics`].
   #[inline(always)]
-  pub(crate) fn state(
+  pub fn handle_scope<'s>(
     &self,
-    isolate: &mut v8::Isolate,
-  ) -> Rc<RefCell<ContextState>> {
-    self
-      .context()
-      .open(isolate)
-      .get_slot::<Rc<RefCell<ContextState>>>(isolate)
-      .unwrap()
-      .clone()
+    isolate: &'s mut v8::Isolate,
+  ) -> v8::HandleScope<'s> {
+    v8::HandleScope::with_context(isolate, &*self.context)
+  }
+
+  pub(crate) fn check_promise_rejections(
+    &self,
+    scope: &mut v8::HandleScope,
+  ) -> Result<(), Error> {
+    let Some((_, handle)) = self.context_state.borrow_mut().pending_promise_rejections.pop_front() else {
+      return Ok(());
+    };
+
+    let exception = v8::Local::new(scope, handle);
+    let state_rc = JsRuntime::state(scope);
+    let state = state_rc.borrow();
+    if let Some(inspector) = &state.inspector {
+      let inspector = inspector.borrow();
+      inspector.exception_thrown(scope, exception, true);
+      if inspector.has_blocking_sessions() {
+        return Ok(());
+      }
+    }
+    exception_to_err_result(scope, exception, true)
+  }
+
+  pub(crate) fn is_same(&self, other: &Rc<v8::Global<v8::Context>>) -> bool {
+    Rc::ptr_eq(&self.context, other)
+  }
+
+  pub fn destroy(self) {
+    let state = self.state();
+    let raw_ptr = self.state().borrow().isolate.unwrap();
+    // SAFETY: We know the isolate outlives the realm
+    let isolate = unsafe { raw_ptr.as_mut().unwrap() };
+    let mut realm_state = state.borrow_mut();
+    // These globals will prevent snapshots from completing, take them
+    std::mem::take(&mut realm_state.js_event_loop_tick_cb);
+    std::mem::take(&mut realm_state.js_build_custom_error_cb);
+    std::mem::take(&mut realm_state.js_promise_reject_cb);
+    std::mem::take(&mut realm_state.js_format_exception_cb);
+    std::mem::take(&mut realm_state.js_wasm_streaming_cb);
+    // The OpCtx slice may contain a circular reference
+    std::mem::take(&mut realm_state.op_ctxs);
+
+    self.context().open(isolate).clear_all_slots(isolate);
+
+    // Expect that this context is dead (we only check this in debug mode)
+    // TODO(mmastrac): This check fails for some tests, will need to fix this
+    // debug_assert_eq!(Rc::strong_count(&self.context), 1, "Realm was still alive when we wanted to destory it. Not dropped?");
+  }
+}
+
+impl JsRealm {
+  pub(crate) fn new(inner: JsRealmInner) -> Self {
+    Self(inner)
   }
 
   #[inline(always)]
@@ -130,13 +217,28 @@ impl JsRealm {
       .clone()
   }
 
+  #[inline(always)]
+  pub fn num_pending_ops(&self) -> usize {
+    self.0.num_pending_ops()
+  }
+
+  #[inline(always)]
+  pub fn num_unrefed_ops(&self) -> usize {
+    self.0.num_unrefed_ops()
+  }
+
   /// For info on the [`v8::Isolate`] parameter, check [`JsRealm#panics`].
   #[inline(always)]
   pub fn handle_scope<'s>(
     &self,
     isolate: &'s mut v8::Isolate,
   ) -> v8::HandleScope<'s> {
-    v8::HandleScope::with_context(isolate, &*self.0)
+    self.0.handle_scope(isolate)
+  }
+
+  #[inline(always)]
+  pub fn context(&self) -> &v8::Global<v8::Context> {
+    self.0.context()
   }
 
   /// For info on the [`v8::Isolate`] parameter, check [`JsRealm#panics`].
@@ -144,8 +246,8 @@ impl JsRealm {
     &self,
     isolate: &'s mut v8::Isolate,
   ) -> v8::Local<'s, v8::Object> {
-    let scope = &mut self.handle_scope(isolate);
-    self.0.open(scope).global(scope)
+    let scope = &mut self.0.handle_scope(isolate);
+    self.0.context.open(scope).global(scope)
   }
 
   fn string_from_code<'a>(
@@ -206,7 +308,7 @@ impl JsRealm {
     name: &'static str,
     source_code: ModuleCode,
   ) -> Result<v8::Global<v8::Value>, Error> {
-    let scope = &mut self.handle_scope(isolate);
+    let scope = &mut self.0.handle_scope(isolate);
 
     let source = Self::string_from_code(scope, &source_code).unwrap();
     debug_assert!(name.is_ascii());
@@ -240,51 +342,23 @@ impl JsRealm {
   // TODO(andreubotella): `mod_evaluate`, `load_main_module`, `load_side_module`
 }
 
-pub struct JsRealmLocal<'s>(v8::Local<'s, v8::Context>);
-impl<'s> JsRealmLocal<'s> {
-  pub fn new(context: v8::Local<'s, v8::Context>) -> Self {
-    JsRealmLocal(context)
-  }
-
-  #[inline(always)]
-  pub fn context(&self) -> v8::Local<v8::Context> {
-    self.0
-  }
-
-  #[inline(always)]
-  pub(crate) fn state(
-    &self,
-    isolate: &mut v8::Isolate,
-  ) -> Rc<RefCell<ContextState>> {
-    self
-      .context()
-      .get_slot::<Rc<RefCell<ContextState>>>(isolate)
-      .unwrap()
-      .clone()
-  }
-
-  pub(crate) fn check_promise_rejections(
-    &self,
-    scope: &mut v8::HandleScope,
-  ) -> Result<(), Error> {
-    let context_state_rc = self.state(scope);
-    let mut context_state = context_state_rc.borrow_mut();
-
-    let Some((_, handle)) = context_state.pending_promise_rejections.pop_front() else {
-      return Ok(());
-    };
-    drop(context_state);
+impl Drop for JsRealm {
+  fn drop(&mut self) {
+    // Don't do anything special with the global realm
+    if self.0.is_global {
+      return;
+    }
 
-    let exception = v8::Local::new(scope, handle);
-    let state_rc = JsRuntime::state(scope);
-    let state = state_rc.borrow();
-    if let Some(inspector) = &state.inspector {
-      let inspector = inspector.borrow();
-      inspector.exception_thrown(scope, exception, true);
-      if inspector.has_blocking_sessions() {
-        return Ok(());
-      }
+    // There's us and there's the runtime
+    if Rc::strong_count(&self.0.context) == 2 {
+      self
+        .0
+        .runtime_state
+        .borrow_mut()
+        .remove_realm(&self.0.context);
+      assert_eq!(Rc::strong_count(&self.0.context), 1);
+      self.0.clone().destroy();
+      assert_eq!(Rc::strong_count(&self.0.context_state), 1);
     }
-    exception_to_err_result(scope, exception, true)
   }
 }