feat: Context slots (#937)

src/binding.cc

@@ -1578,6 +1578,20 @@ const v8::Object* v8__Context__Global(const v8::Context& self) {
   return local_to_ptr(ptr_to_local(&self)->Global());
 }
 
+uint32_t v8__Context__GetNumberOfEmbedderDataFields(const v8::Context& self) {
+  return ptr_to_local(&self)->GetNumberOfEmbedderDataFields();
+}
+
+void* v8__Context__GetAlignedPointerFromEmbedderData(const v8::Context& self,
+                                                     int index) {
+  return ptr_to_local(&self)->GetAlignedPointerFromEmbedderData(index);
+}
+
+void v8__Context__SetAlignedPointerInEmbedderData(v8::Context& self, int index,
+                                                  void* value) {
+  ptr_to_local(&self)->SetAlignedPointerInEmbedderData(index, value);
+}
+
 const v8::Data* v8__Context__GetDataFromSnapshotOnce(v8::Context& self,
                                                      size_t index) {
   return maybe_local_to_ptr(

src/context.rs

@@ -1,5 +1,8 @@
 // Copyright 2019-2021 the Deno authors. All rights reserved. MIT license.
+use crate::handle::UnsafeRefHandle;
+use crate::isolate::BuildTypeIdHasher;
 use crate::isolate::Isolate;
+use crate::isolate::RawSlot;
 use crate::Context;
 use crate::Function;
 use crate::HandleScope;
@@ -7,7 +10,12 @@ use crate::Local;
 use crate::Object;
 use crate::ObjectTemplate;
 use crate::Value;
-use std::ptr::null;
+use crate::Weak;
+use libc::c_int;
+use std::any::TypeId;
+use std::collections::HashMap;
+use std::ffi::c_void;
+use std::ptr::{null, null_mut};
 
 extern "C" {
   fn v8__Context__New(
@@ -15,6 +23,7 @@ extern "C" {
     templ: *const ObjectTemplate,
     global_object: *const Value,
   ) -> *const Context;
+  fn v8__Context__GetIsolate(this: *const Context) -> *mut Isolate;
   fn v8__Context__Global(this: *const Context) -> *const Object;
   fn v8__Context__SetPromiseHooks(
     this: *const Context,
@@ -23,9 +32,21 @@ extern "C" {
     after_hook: *const Function,
     resolve_hook: *const Function,
   );
+  fn v8__Context__GetNumberOfEmbedderDataFields(this: *const Context) -> u32;
+  fn v8__Context__GetAlignedPointerFromEmbedderData(
+    this: *const Context,
+    index: c_int,
+  ) -> *mut c_void;
+  fn v8__Context__SetAlignedPointerInEmbedderData(
+    this: *const Context,
+    index: c_int,
+    value: *mut c_void,
+  );
 }
 
 impl Context {
+  const ANNEX_SLOT: c_int = 1;
+
   /// Creates a new context.
   pub fn new<'s>(scope: &mut HandleScope<'s, ()>) -> Local<'s, Context> {
     // TODO: optional arguments;
@@ -84,4 +105,187 @@ impl Context {
       )
     }
   }
+
+  fn get_annex_mut<'a>(
+    &'a self,
+    isolate: &'a mut Isolate,
+    create_if_not_present: bool,
+  ) -> Option<&'a mut ContextAnnex> {
+    assert!(
+      std::ptr::eq(isolate, unsafe { v8__Context__GetIsolate(self) }),
+      "attempted to use Context slots with the wrong Isolate"
+    );
+
+    let num_data_fields =
+      unsafe { v8__Context__GetNumberOfEmbedderDataFields(self) } as c_int;
+    if num_data_fields > Self::ANNEX_SLOT {
+      let annex_ptr = unsafe {
+        v8__Context__GetAlignedPointerFromEmbedderData(self, Self::ANNEX_SLOT)
+      } as *mut ContextAnnex;
+      if !annex_ptr.is_null() {
+        // SAFETY: This reference doesn't outlive the Context, so it can't outlive
+        // the annex itself. Also, any mutations or accesses to the annex after
+        // its creation require a mutable reference to the context's isolate, but
+        // such a mutable reference is consumed by this reference during its
+        // lifetime.
+        return Some(unsafe { &mut *annex_ptr });
+      }
+    }
+
+    if !create_if_not_present {
+      return None;
+    }
+
+    let annex = Box::new(ContextAnnex {
+      slots: Default::default(),
+      // Gets replaced later in the method.
+      self_weak: Weak::empty(isolate),
+    });
+    let annex_ptr = Box::into_raw(annex);
+    unsafe {
+      v8__Context__SetAlignedPointerInEmbedderData(
+        self,
+        Self::ANNEX_SLOT,
+        annex_ptr as *mut _,
+      )
+    };
+    assert!(
+      unsafe { v8__Context__GetNumberOfEmbedderDataFields(self) } as c_int
+        > Self::ANNEX_SLOT
+    );
+
+    // Make sure to drop the annex after the context is dropped, by creating a
+    // weak handle with a finalizer that drops the annex, and storing the weak
+    // in the annex itself.
+    let weak = {
+      // SAFETY: `self` can only have been derived from a `Local` or `Global`,
+      // and assuming the caller is only using safe code, the `Local` or
+      // `Global` must still be alive, so `self_ref_handle` won't outlive it.
+      // We also check above that `isolate` is the context's isolate.
+      let self_ref_handle = unsafe { UnsafeRefHandle::new(self, isolate) };
+
+      Weak::with_finalizer(
+        isolate,
+        self_ref_handle,
+        Box::new(move |_| {
+          // SAFETY: The lifetimes of references to the annex returned by this
+          // method are always tied to the context, and because this is the
+          // context's finalizer, we know there are no living references to
+          // the annex. And since the finalizer is only called once, the annex
+          // can't have been dropped before.
+          let _ = unsafe { Box::from_raw(annex_ptr) };
+        }),
+      )
+    };
+
+    // SAFETY: This reference doesn't outlive the Context, so it can't outlive
+    // the annex itself. Also, any mutations or accesses to the annex after
+    // its creation require a mutable reference to the context's isolate, but
+    // such a mutable reference is consumed by this reference during its
+    // lifetime.
+    let annex_mut = unsafe { &mut *annex_ptr };
+    annex_mut.self_weak = weak;
+    Some(annex_mut)
+  }
+
+  /// Get a reference to embedder data added with [`Self::set_slot()`].
+  pub fn get_slot<'a, T: 'static>(
+    &'a self,
+    isolate: &'a mut Isolate,
+  ) -> Option<&'a T> {
+    if let Some(annex) = self.get_annex_mut(isolate, false) {
+      annex.slots.get(&TypeId::of::<T>()).map(|slot| {
+        // SAFETY: `Self::set_slot` guarantees that only values of type T will be
+        // stored with T's TypeId as their key.
+        unsafe { slot.borrow::<T>() }
+      })
+    } else {
+      None
+    }
+  }
+
+  /// Get a mutable reference to embedder data added with [`Self::set_slot()`].
+  pub fn get_slot_mut<'a, T: 'static>(
+    &'a self,
+    isolate: &'a mut Isolate,
+  ) -> Option<&'a mut T> {
+    if let Some(annex) = self.get_annex_mut(isolate, false) {
+      annex.slots.get_mut(&TypeId::of::<T>()).map(|slot| {
+        // SAFETY: `Self::set_slot` guarantees that only values of type T will be
+        // stored with T's TypeId as their key.
+        unsafe { slot.borrow_mut::<T>() }
+      })
+    } else {
+      None
+    }
+  }
+
+  /// Use with [`Context::get_slot`] and [`Context::get_slot_mut`] to associate
+  /// state with a Context.
+  ///
+  /// This method gives ownership of value to the Context. Exactly one object of
+  /// each type can be associated with a Context. If called more than once with
+  /// an object of the same type, the earlier version will be dropped and
+  /// replaced.
+  ///
+  /// Returns true if value was set without replacing an existing value.
+  ///
+  /// The value will be dropped when the context is garbage collected.
+  pub fn set_slot<'a, T: 'static>(
+    &'a self,
+    isolate: &'a mut Isolate,
+    value: T,
+  ) -> bool {
+    self
+      .get_annex_mut(isolate, true)
+      .unwrap()
+      .slots
+      .insert(TypeId::of::<T>(), RawSlot::new(value))
+      .is_none()
+  }
+
+  /// Removes the embedder data added with [`Self::set_slot()`] and returns it
+  /// if it exists.
+  pub fn remove_slot<'a, T: 'static>(
+    &'a self,
+    isolate: &'a mut Isolate,
+  ) -> Option<T> {
+    if let Some(annex) = self.get_annex_mut(isolate, false) {
+      annex.slots.remove(&TypeId::of::<T>()).map(|slot| {
+        // SAFETY: `Self::set_slot` guarantees that only values of type T will be
+        // stored with T's TypeId as their key.
+        unsafe { slot.into_inner::<T>() }
+      })
+    } else {
+      None
+    }
+  }
+
+  /// Removes all embedder data added with [`Self::set_slot()`], and
+  /// deletes any internal state needed to keep track of such slots.
+  ///
+  /// This is needed to make a snapshot with
+  /// [`SnapshotCreator`](crate::SnapshotCreator), since the internal embedder
+  /// state uses [`Weak`] handles, which cannot be alive at the time of
+  /// snapshotting.
+  pub fn clear_all_slots<'a>(&'a self, isolate: &'a mut Isolate) {
+    if let Some(annex_mut) = self.get_annex_mut(isolate, false) {
+      let annex_ptr = annex_mut as *mut ContextAnnex;
+      let _ = unsafe { Box::from_raw(annex_ptr) };
+      unsafe {
+        v8__Context__SetAlignedPointerInEmbedderData(
+          self,
+          Self::ANNEX_SLOT,
+          null_mut(),
+        )
+      };
+    }
+  }
+}
+
+struct ContextAnnex {
+  slots: HashMap<TypeId, RawSlot, BuildTypeIdHasher>,
+  // In order to run the finalizer that drops the ContextAnnex when the Context
+  // is GC'd, the corresponding Weak must be kept alive until that time.
+  self_weak: Weak<Context>,
 }

src/handle.rs

@@ -215,6 +215,30 @@ impl<T> Drop for Global<T> {
   }
 }
 
+/// An implementation of [`Handle`] that can be constructed unsafely from a
+/// reference.
+pub(crate) struct UnsafeRefHandle<'a, T> {
+  reference: &'a T,
+  isolate_handle: IsolateHandle,
+}
+impl<'a, T> UnsafeRefHandle<'a, T> {
+  /// Constructs an `UnsafeRefHandle`.
+  ///
+  /// # Safety
+  ///
+  /// `reference` must be derived from a [`Local`] or [`Global`] handle, and its
+  /// lifetime must not outlive that handle. Furthermore, `isolate` must be the
+  /// isolate associated with the handle (for [`Local`], the current isolate;
+  /// for [`Global`], the isolate you would pass to the [`Global::open()`]
+  /// method).
+  pub unsafe fn new(reference: &'a T, isolate: &mut Isolate) -> Self {
+    UnsafeRefHandle {
+      reference,
+      isolate_handle: isolate.thread_safe_handle(),
+    }
+  }
+}
+
 pub trait Handle: Sized {
   type Data;
 
@@ -285,6 +309,26 @@ impl<'a, T> Handle for &'a Global<T> {
   }
 }
 
+impl<'a, T> Handle for UnsafeRefHandle<'a, T> {
+  type Data = T;
+  fn get_handle_info(&self) -> HandleInfo<T> {
+    HandleInfo::new(
+      NonNull::from(self.reference),
+      (&self.isolate_handle).into(),
+    )
+  }
+}
+
+impl<'a, T> Handle for &'a UnsafeRefHandle<'_, T> {
+  type Data = T;
+  fn get_handle_info(&self) -> HandleInfo<T> {
+    HandleInfo::new(
+      NonNull::from(self.reference),
+      (&self.isolate_handle).into(),
+    )
+  }
+}
+
 impl<'s, T> Borrow<T> for Local<'s, T> {
   fn borrow(&self) -> &T {
     &**self

src/isolate.rs

@@ -1108,7 +1108,7 @@ const _: () = {
   assert!(align_of::<TypeId>() == size_of::<u64>());
 };
 
-struct RawSlot {
+pub(crate) struct RawSlot {
   data: RawSlotData,
   dtor: Option<RawSlotDtor>,
 }

tests/slots.rs

@@ -10,6 +10,17 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::sync::Once;
 
+fn setup() {
+  static START: Once = Once::new();
+  START.call_once(|| {
+    v8::V8::set_flags_from_string("--expose_gc");
+    v8::V8::initialize_platform(
+      v8::new_default_platform(0, false).make_shared(),
+    );
+    v8::V8::initialize();
+  });
+}
+
 struct CoreIsolate(v8::OwnedIsolate);
 
 struct CoreIsolateState {
@@ -25,13 +36,7 @@ impl Drop for CoreIsolateState {
 
 impl CoreIsolate {
   fn new(drop_count: Rc<AtomicUsize>) -> CoreIsolate {
-    static START: Once = Once::new();
-    START.call_once(|| {
-      v8::V8::initialize_platform(
-        v8::new_default_platform(0, false).make_shared(),
-      );
-      v8::V8::initialize();
-    });
+    setup();
     let mut isolate = v8::Isolate::new(Default::default());
     let state = CoreIsolateState { drop_count, i: 0 };
     isolate.set_slot(state);
@@ -244,3 +249,73 @@ fn slots_auto_boxing() {
   assert_eq!(Some(value2), core_isolate.remove_slot::<Box<TestData>>());
   assert_eq!(None, core_isolate.get_slot::<Box<TestData>>());
 }
+
+#[test]
+fn context_slots() {
+  setup();
+  let isolate = &mut v8::Isolate::new(Default::default());
+  let scope = &mut v8::HandleScope::new(isolate);
+  let context = v8::Context::new(scope);
+
+  assert!(context.set_slot(scope, TestState(0)));
+  assert!(!context.set_slot(scope, TestState(1)));
+
+  context.get_slot_mut::<TestState>(scope).unwrap().0 += 5;
+  assert_eq!(context.get_slot::<TestState>(scope).unwrap().0, 6);
+
+  let value = context.remove_slot::<TestState>(scope).unwrap();
+  assert_eq!(value.0, 6);
+  assert!(context.remove_slot::<TestState>(scope).is_none());
+}
+
+#[test]
+fn dropped_context_slots() {
+  // Test that context slots are dropped when the context is GC'd.
+  use std::cell::Cell;
+
+  struct DropMarker(Rc<Cell<bool>>);
+  impl Drop for DropMarker {
+    fn drop(&mut self) {
+      println!("Dropping the drop marker");
+      self.0.set(true);
+    }
+  }
+
+  let mut isolate = CoreIsolate::new(Default::default());
+  let dropped = Rc::new(Cell::new(false));
+  {
+    let scope = &mut v8::HandleScope::new(isolate.deref_mut());
+    let context = v8::Context::new(scope);
+
+    context.set_slot(scope, DropMarker(dropped.clone()));
+  }
+
+  assert!(isolate.execute("gc()"));
+  assert!(dropped.get());
+}
+
+#[test]
+fn clear_all_context_slots() {
+  setup();
+
+  let mut snapshot_creator = v8::SnapshotCreator::new(None);
+  let mut isolate = unsafe { snapshot_creator.get_owned_isolate() };
+
+  {
+    let scope = &mut v8::HandleScope::new(&mut isolate);
+    let context = v8::Context::new(scope);
+    let scope = &mut v8::ContextScope::new(scope, context);
+
+    context.set_slot(scope, TestState(0));
+    context.clear_all_slots(scope);
+    assert!(context.get_slot::<TestState>(scope).is_none());
+
+    snapshot_creator.set_default_context(context);
+  }
+
+  std::mem::forget(isolate);
+
+  snapshot_creator
+    .create_blob(v8::FunctionCodeHandling::Keep)
+    .unwrap();
+}