Join method

crates/ion/src/env.rs

@@ -1,6 +1,4 @@
-use std::cell::RefCell;
 use std::future::Future;
-use std::rc::Rc;
 use std::sync::Arc;
 
 use flume::Sender;
@@ -27,7 +25,6 @@ pub struct Env {
     pub(crate) context: sys::GlobalContext,
     pub(crate) background_task_manager: Arc<BackgroundTaskManager>,
     pub(crate) global_refs: RefCounter,
-    pub(crate) shutdown_requested: Rc<RefCell<bool>>,
     pub(crate) tx: Sender<JsWorkerEvent>,
     pub(crate) finalizer_registry: FinalizerRegistery,
     pub(crate) global_this: sys::GlobalThis,
@@ -40,7 +37,6 @@ impl Env {
         context: sys::GlobalContext,
         background_task_manager: Arc<BackgroundTaskManager>,
         global_refs: RefCounter,
-        shutdown_requested: Rc<RefCell<bool>>,
         tx: Sender<JsWorkerEvent>,
         finalizer_registry: FinalizerRegistery,
         global_this: sys::GlobalThis,
@@ -53,7 +49,6 @@ impl Env {
             background_task_manager,
             inner: std::ptr::null_mut(),
             global_refs,
-            shutdown_requested,
             finalizer_registry,
             tx,
         });
@@ -85,19 +80,6 @@ impl Env {
 
     pub fn dec_ref(&self) {
         self.global_refs.dec();
-        let shutdown_requested = {
-            let shutdown_requested = self.shutdown_requested.borrow();
-            *shutdown_requested
-        };
-
-        if self.global_refs.count() == 0 && shutdown_requested {
-            self.tx
-                .try_send(JsWorkerEvent::RequestContextShutdown {
-                    id: self.realm_id,
-                    resolve: None,
-                })
-                .unwrap();
-        }
     }
 
     pub fn ref_count(&self) -> usize {

crates/ion/src/error.rs

@@ -13,6 +13,8 @@ pub enum Error {
     PlatformCommunicationError,
     PlatformInitializeError,
     PlatformDisposeError,
+    WorkerAlreadyShutdown,
+    ContextAlreadyShutdown,
     IsolateNotInitializedError,
     EventLoopNotInitializedError,
     WorkerInitializeError,

crates/ion/src/js_context.rs

@@ -1,17 +1,22 @@
+use std::sync::Arc;
+
 use flume::Sender;
 use flume::bounded;
 
 use crate::Env;
 use crate::Error;
 use crate::JsUnknown;
 use crate::platform::worker::JsWorkerEvent;
-use crate::utils::channel::oneshot;
+use crate::platform::worker_handle_state::WorkerHandleState;
+use crate::utils::complete_signal::CompleteSignal;
 
 /// This is a handle to a v8::Context
 #[derive(Debug, Clone)]
 pub struct JsContext {
+    pub(crate) worker_handle_state: Arc<WorkerHandleState>,
     pub(crate) id: usize,
     pub(crate) tx: Sender<JsWorkerEvent>,
+    pub(crate) context_shutdown_sig: CompleteSignal,
 }
 
 impl JsContext {
@@ -83,25 +88,64 @@ impl JsContext {
         let specifier = specifier.as_ref().to_string();
         self.exec_blocking(move |env| env.import(specifier))
     }
-}
 
-impl Drop for JsContext {
-    fn drop(&mut self) {
-        let (tx, rx) = oneshot();
+    /// Wait for the context to complete all activity
+    pub fn join(self) -> crate::Result<()> {
+        if !self.worker_handle_state.worker_handle_active() {
+            return Err(crate::Error::WorkerAlreadyShutdown);
+        }
+
+        self.worker_handle_state
+            .context_handle_set_status(&self.id, false);
 
         if self
             .tx
-            .send(JsWorkerEvent::RequestContextShutdown {
-                id: self.id,
-                resolve: Some(tx),
+            .send(JsWorkerEvent::ContextHandleDeactivated {
+                id: self.id.clone(),
             })
             .is_err()
         {
-            panic!("Cannot drop JsContext 1")
-        };
+            return Err(crate::Error::ContextAlreadyShutdown);
+        }
+
+        self.context_shutdown_sig.wait();
+
+        Ok(())
+    }
+
+    /// Wait for the context to complete all activity
+    pub async fn join_async(self) -> crate::Result<()> {
+        if !self.worker_handle_state.worker_handle_active() {
+            return Err(crate::Error::WorkerAlreadyShutdown);
+        }
 
-        if rx.recv().is_err() {
-            panic!("Cannot drop JsContext 2")
+        self.worker_handle_state
+            .context_handle_set_status(&self.id, false);
+
+        if self
+            .tx
+            .send(JsWorkerEvent::ContextHandleDeactivated {
+                id: self.id.clone(),
+            })
+            .is_err()
+        {
+            return Err(crate::Error::ContextAlreadyShutdown);
         }
+
+        self.context_shutdown_sig.wait_async().await;
+
+        Ok(())
+    }
+}
+
+impl Drop for JsContext {
+    fn drop(&mut self) {
+        self.worker_handle_state
+            .context_handle_set_status(&self.id, false);
+
+        drop(
+            self.tx
+                .try_send(JsWorkerEvent::ContextHandleDropped { id: self.id }),
+        );
     }
 }

crates/ion/src/js_runtime.rs

@@ -14,6 +14,8 @@ use crate::JsWorker;
 use crate::JsWorkerOptions;
 use crate::platform::platform::HAS_INIT;
 use crate::platform::platform::PLATFORM;
+use crate::platform::worker_handle_state::WorkerHandleState;
+use crate::utils::complete_signal::CompleteSignal;
 
 static JS_RUNTIME: OnceLock<crate::Result<Arc<JsRuntime>>> = OnceLock::new();
 
@@ -173,12 +175,17 @@ impl JsRuntime {
     pub fn spawn_worker(
         &self,
         options: JsWorkerOptions,
-    ) -> crate::Result<Arc<JsWorker>> {
+    ) -> crate::Result<JsWorker> {
+        let worker_handle_state = Arc::new(WorkerHandleState::default());
+        let worker_shutdown_sig = CompleteSignal::default();
+
         let (tx, rx) = bounded(1);
 
         if self
             .tx
             .send(PlatformEvent::SpawnWorker {
+                worker_shutdown_sig: worker_shutdown_sig.clone(),
+                worker_handle_state: Arc::clone(&worker_handle_state),
                 extensions: options.extensions,
                 transformers: options.transformers,
                 resolvers: options.resolvers,
@@ -193,7 +200,12 @@ impl JsRuntime {
             return Err(Error::WorkerInitializeError);
         };
 
-        Ok(Arc::new(JsWorker::new(tx, handle)))
+        Ok(JsWorker::new(
+            worker_handle_state,
+            tx,
+            Arc::new(handle),
+            worker_shutdown_sig,
+        ))
     }
 }
 

crates/ion/src/js_worker.rs

@@ -11,7 +11,8 @@ use crate::JsExtension;
 use crate::JsResolver;
 use crate::JsTransformer;
 use crate::platform::worker::JsWorkerEvent;
-use crate::utils::channel::oneshot;
+use crate::platform::worker_handle_state::WorkerHandleState;
+use crate::utils::complete_signal::CompleteSignal;
 
 #[derive(Default)]
 pub struct JsWorkerOptions {
@@ -30,25 +31,39 @@ pub struct JsWorkerOptions {
 /// to be used to execute JavaScript
 #[derive(Debug)]
 pub struct JsWorker {
+    worker_handle_state: Arc<WorkerHandleState>,
     tx: Sender<JsWorkerEvent>,
-    handle: Mutex<Option<JoinHandle<crate::Result<()>>>>,
+    handle: Arc<Mutex<Option<JoinHandle<crate::Result<()>>>>>,
+    worker_shutdown_sig: CompleteSignal,
 }
 
 impl JsWorker {
     pub(crate) fn new(
+        worker_handle_state: Arc<WorkerHandleState>,
         tx: Sender<JsWorkerEvent>,
-        handle: Mutex<Option<JoinHandle<crate::Result<()>>>>,
+        handle: Arc<Mutex<Option<JoinHandle<crate::Result<()>>>>>,
+        worker_shutdown_sig: CompleteSignal,
     ) -> Self {
-        JsWorker { tx, handle }
+        JsWorker {
+            tx,
+            handle,
+            worker_handle_state,
+            worker_shutdown_sig,
+        }
     }
 
     /// Create a handle to a v8::Context associated with this v8::Isolate
-    pub fn create_context(&self) -> crate::Result<Arc<JsContext>> {
+    pub fn create_context(&self) -> crate::Result<JsContext> {
+        let context_shutdown_sig = CompleteSignal::default();
+
         let (tx, rx) = bounded(1);
 
         if self
             .tx
-            .send(JsWorkerEvent::CreateContext { resolve: tx })
+            .send(JsWorkerEvent::CreateContext {
+                resolve: tx,
+                context_shutdown_sig: context_shutdown_sig.clone(),
+            })
             .is_err()
         {
             return Err(Error::WorkerInitializeError);
@@ -58,7 +73,15 @@ impl JsWorker {
             return Err(Error::WorkerInitializeError);
         };
 
-        Ok(Arc::new(JsContext { id, tx }))
+        self.worker_handle_state
+            .context_handle_set_status(&id, true);
+
+        Ok(JsContext {
+            id,
+            tx,
+            worker_handle_state: Arc::clone(&self.worker_handle_state),
+            context_shutdown_sig,
+        })
     }
 
     pub fn run_garbage_collection_for_testing(&self) -> crate::Result<()> {
@@ -74,30 +97,51 @@ impl JsWorker {
 
         Ok(rx.recv()?)
     }
-}
 
-impl Drop for JsWorker {
-    fn drop(&mut self) {
-        let (tx, rx) = oneshot();
+    /// Wait for all of the contexts within the worker to complete all activity
+    pub fn join(self) -> crate::Result<()> {
+        self.worker_handle_state.worker_handle_deactivate();
+        self.tx
+            .send(JsWorkerEvent::WorkerHandleDeactivated)
+            .unwrap();
 
-        if self
-            .tx
-            .send(JsWorkerEvent::RequestShutdown { resolve: tx })
-            .is_err()
-        {
-            panic!("Cannot drop JsWorker 1");
+        self.worker_shutdown_sig.wait();
+
+        let Ok(mut handle) = self.handle.lock() else {
+            panic!("Cannot drop JsWorker 3");
         };
 
-        if rx.recv().is_err() {
-            panic!("Cannot drop JsWorker 2");
+        if let Some(handle) = handle.take() {
+            drop(handle.join().unwrap());
         }
 
+        Ok(())
+    }
+
+    /// Wait for all of the contexts within the worker to complete all activity
+    pub async fn join_async(self) -> crate::Result<()> {
+        self.worker_handle_state.worker_handle_deactivate();
+        self.tx
+            .send(JsWorkerEvent::WorkerHandleDeactivated)
+            .unwrap();
+
+        self.worker_shutdown_sig.wait_async().await;
+
         let Ok(mut handle) = self.handle.lock() else {
             panic!("Cannot drop JsWorker 3");
         };
 
         if let Some(handle) = handle.take() {
             drop(handle.join().unwrap());
         }
+
+        Ok(())
+    }
+}
+
+impl Drop for JsWorker {
+    fn drop(&mut self) {
+        self.worker_handle_state.worker_handle_deactivate();
+        drop(self.tx.try_send(JsWorkerEvent::WorkerHandleDropped));
     }
 }

crates/ion/src/platform/mod.rs

@@ -9,5 +9,6 @@ mod realm;
 pub mod resolve;
 pub(crate) mod sys;
 pub(crate) mod worker;
+pub(crate) mod worker_handle_state;
 
 pub(crate) use realm::*;

crates/ion/src/platform/platform.rs

@@ -16,6 +16,8 @@ use crate::JsTransformer;
 use crate::platform::background_worker::BackgroundTaskManager;
 use crate::platform::worker::JsWorkerEvent;
 use crate::platform::worker::start_js_worker_thread;
+use crate::platform::worker_handle_state::WorkerHandleState;
+use crate::utils::complete_signal::CompleteSignal;
 
 pub(crate) enum PlatformEvent {
     Init {
@@ -26,6 +28,8 @@ pub(crate) enum PlatformEvent {
         transformers: Vec<JsTransformer>,
     },
     SpawnWorker {
+        worker_shutdown_sig: CompleteSignal,
+        worker_handle_state: Arc<WorkerHandleState>,
         extensions: Vec<JsExtension>,
         resolvers: Vec<JsResolver>,
         transformers: Vec<JsTransformer>,
@@ -96,6 +100,8 @@ pub(crate) static PLATFORM: LazyLock<Sender<PlatformEvent>> = LazyLock::new(|| {
                     }
                 }
                 PlatformEvent::SpawnWorker {
+                    worker_shutdown_sig,
+                    worker_handle_state,
                     resolve,
                     extensions: init_extensions,
                     resolvers: init_resolvers,
@@ -117,6 +123,8 @@ pub(crate) static PLATFORM: LazyLock<Sender<PlatformEvent>> = LazyLock::new(|| {
                     }
 
                     let (tx, handle) = start_js_worker_thread(
+                        worker_shutdown_sig,
+                        worker_handle_state,
                         background_task_manager.clone(),
                         worker_extensions,
                         worker_resolvers,