feat: ExternalOpsTracker (#770)

Allows an embedder to track "external" operations which should still
keep the event loop alive. Will be used to ref napi threadsafe
functions.

core/lib.rs

@@ -127,6 +127,7 @@ pub use crate::modules::SourceCodeCacheInfo;
 pub use crate::modules::StaticModuleLoader;
 pub use crate::modules::ValidateImportAttributesCb;
 pub use crate::normalize_path::normalize_path;
+pub use crate::ops::ExternalOpsTracker;
 pub use crate::ops::OpId;
 pub use crate::ops::OpMetadata;
 pub use crate::ops::OpState;

core/ops.rs

@@ -16,6 +16,8 @@ use std::ops::Deref;
 use std::ops::DerefMut;
 use std::ptr::NonNull;
 use std::rc::Rc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use v8::fast_api::CFunctionInfo;
 use v8::fast_api::CTypeInfo;
@@ -250,12 +252,43 @@ impl OpCtx {
   }
 }
 
+/// Allows an embedder to track operations which should
+/// keep the event loop alive.
+#[derive(Debug, Clone)]
+pub struct ExternalOpsTracker {
+  counter: Arc<AtomicUsize>,
+}
+
+impl ExternalOpsTracker {
+  pub fn ref_op(&self) {
+    self.counter.fetch_add(1, Ordering::Relaxed);
+  }
+
+  pub fn unref_op(&self) {
+    let _ =
+      self
+        .counter
+        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
+          if x == 0 {
+            None
+          } else {
+            Some(x - 1)
+          }
+        });
+  }
+
+  pub(crate) fn has_pending_ops(&self) -> bool {
+    self.counter.load(Ordering::Relaxed) > 0
+  }
+}
+
 /// Maintains the resources and ops inside a JS runtime.
 pub struct OpState {
   pub resource_table: ResourceTable,
   pub(crate) gotham_state: GothamState,
   pub waker: Arc<AtomicWaker>,
   pub feature_checker: Arc<FeatureChecker>,
+  pub external_ops_tracker: ExternalOpsTracker,
 }
 
 impl OpState {
@@ -265,6 +298,9 @@ impl OpState {
       gotham_state: Default::default(),
       waker: Arc::new(AtomicWaker::new()),
       feature_checker: maybe_feature_checker.unwrap_or_default(),
+      external_ops_tracker: ExternalOpsTracker {
+        counter: Arc::new(AtomicUsize::new(0)),
+      },
     }
   }
 

core/runtime/jsrealm.rs

@@ -11,6 +11,7 @@ use crate::modules::ModuleCodeString;
 use crate::modules::ModuleId;
 use crate::modules::ModuleMap;
 use crate::modules::ModuleName;
+use crate::ops::ExternalOpsTracker;
 use crate::ops::OpCtx;
 use crate::stats::RuntimeActivityTraces;
 use crate::tasks::V8TaskSpawnerFactory;
@@ -68,6 +69,7 @@ pub struct ContextState {
   pub(crate) exception_state: Rc<ExceptionState>,
   pub(crate) has_next_tick_scheduled: Cell<bool>,
   pub(crate) get_error_class_fn: GetErrorClassFn,
+  pub(crate) external_ops_tracker: ExternalOpsTracker,
 }
 
 impl ContextState {
@@ -76,6 +78,7 @@ impl ContextState {
     isolate_ptr: *mut v8::OwnedIsolate,
     get_error_class_fn: GetErrorClassFn,
     op_ctxs: Box<[OpCtx]>,
+    external_ops_tracker: ExternalOpsTracker,
   ) -> Self {
     Self {
       isolate: Some(isolate_ptr),
@@ -91,6 +94,7 @@ impl ContextState {
       task_spawner_factory: Default::default(),
       timers: Default::default(),
       unrefed_ops: Default::default(),
+      external_ops_tracker,
     }
   }
 }

core/runtime/jsruntime.rs

@@ -811,6 +811,7 @@ impl JsRuntime {
       isolate_ptr,
       options.get_error_class_fn.unwrap_or(&|_| "Error"),
       op_ctxs,
+      op_state.borrow().external_ops_tracker.clone(),
     ));
 
     // TODO(bartlomieju): factor out
@@ -1809,6 +1810,7 @@ impl JsRuntime {
         || pending_state.has_pending_dyn_imports
         || pending_state.has_pending_dyn_module_evaluation
         || pending_state.has_pending_background_tasks
+        || pending_state.has_pending_external_ops
         || pending_state.has_tick_scheduled
       {
         // pass, will be polled again
@@ -1823,6 +1825,7 @@ impl JsRuntime {
       if pending_state.has_pending_ops
         || pending_state.has_pending_dyn_imports
         || pending_state.has_pending_background_tasks
+        || pending_state.has_pending_external_ops
         || pending_state.has_tick_scheduled
       {
         // pass, will be polled again
@@ -2011,6 +2014,7 @@ pub(crate) struct EventLoopPendingState {
   has_pending_background_tasks: bool,
   has_tick_scheduled: bool,
   has_pending_promise_events: bool,
+  has_pending_external_ops: bool,
 }
 
 impl EventLoopPendingState {
@@ -2053,6 +2057,7 @@ impl EventLoopPendingState {
       has_pending_background_tasks: scope.has_pending_background_tasks(),
       has_tick_scheduled: state.has_next_tick_scheduled.get(),
       has_pending_promise_events,
+      has_pending_external_ops: state.external_ops_tracker.has_pending_ops(),
     }
   }
 
@@ -2071,6 +2076,7 @@ impl EventLoopPendingState {
       || self.has_pending_background_tasks
       || self.has_tick_scheduled
       || self.has_pending_promise_events
+      || self.has_pending_external_ops
   }
 }