stamp: expose cancel token for capturing the timing when the user worker is terminated

nyannyacha · nyannyacha · commit b8de2b01578e · 2025-01-02T05:32:39.000Z
diff --git a/crates/base/src/worker/pool.rs b/crates/base/src/worker/pool.rs
@@ -263,11 +263,13 @@ impl WorkerPool {
             .as_user_worker()
             .map_or(false, |it| !is_oneshot_policy && it.force_create);
 
-        if let Some(ref active_worker_uuid) = self.maybe_active_worker(&service_path, force_create)
+        if let Some((ref active_worker_uuid, profile)) =
+            self.maybe_active_worker(&service_path, force_create)
         {
             if tx
                 .send(Ok(CreateUserWorkerResult {
                     key: *active_worker_uuid,
+                    token: profile.cancel.clone(),
                 }))
                 .is_err()
             {
@@ -448,7 +450,7 @@ impl WorkerPool {
                         permit: permit.map(Arc::new),
                         status: status.clone(),
                         exit: surface.exit,
-                        cancel,
+                        cancel: cancel.clone(),
                     };
 
                     if worker_pool_msgs_tx
@@ -457,7 +459,13 @@ impl WorkerPool {
                     {
                         error!("user worker msgs receiver dropped")
                     }
-                    if tx.send(Ok(CreateUserWorkerResult { key: uuid })).is_err() {
+                    if tx
+                        .send(Ok(CreateUserWorkerResult {
+                            key: uuid,
+                            token: cancel,
+                        }))
+                        .is_err()
+                    {
                         error!("main worker receiver dropped")
                     };
                 }
@@ -632,7 +640,11 @@ impl WorkerPool {
         }
     }
 
-    fn maybe_active_worker(&mut self, service_path: &String, force_create: bool) -> Option<Uuid> {
+    fn maybe_active_worker(
+        &mut self,
+        service_path: &String,
+        force_create: bool,
+    ) -> Option<(Uuid, &UserWorkerProfile)> {
         if force_create {
             return None;
         }
@@ -648,11 +660,13 @@ impl WorkerPool {
             .get(&worker_uuid)
             .map(|it| it.status.is_retired.clone())
         {
-            Some(is_retired) if !is_retired.is_raised() => Some(worker_uuid),
+            Some(is_retired) if !is_retired.is_raised() => {
+                Some((worker_uuid, self.user_workers.get(&worker_uuid).unwrap()))
+            }
 
             _ => {
                 self.retire(&worker_uuid);
-                self.maybe_active_worker(service_path, force_create)
+                self.maybe_active_worker(service_path, false)
             }
         }
     }
diff --git a/crates/base/src/worker/worker_surface_creation.rs b/crates/base/src/worker/worker_surface_creation.rs
@@ -14,7 +14,6 @@ use sb_workers::context::{
     WorkerExit, WorkerRequestMsg, WorkerRuntimeOpts,
 };
 use tokio::sync::{mpsc, oneshot};
-use tokio_util::sync::CancellationToken;
 
 use crate::{inspector_server::Inspector, server::ServerFlags};
 
@@ -329,8 +328,7 @@ impl WorkerSurfaceBuilder {
             worker_builder_hook,
         } = self;
 
-        let (worker_boot_result_tx, worker_boot_result_rx) =
-            oneshot::channel::<Result<(MetricSource, CancellationToken), anyhow::Error>>();
+        let (worker_boot_result_tx, worker_boot_result_rx) = oneshot::channel();
 
         let flags = flags.unwrap_or_default();
         let init_opts = init_opts.context("init_opts must be specified")?;
diff --git a/ext/workers/context.rs b/ext/workers/context.rs
@@ -245,6 +245,7 @@ pub type SendRequestResult = (Response<Body>, mpsc::UnboundedSender<()>);
 #[derive(Debug)]
 pub struct CreateUserWorkerResult {
     pub key: Uuid,
+    pub token: CancellationToken,
 }
 
 #[derive(Debug)]
diff --git a/ext/workers/lib.rs b/ext/workers/lib.rs
@@ -94,16 +94,13 @@ pub struct UserWorkerCreateOptions {
 }
 
 #[op2(async)]
-#[string]
+#[serde]
 pub async fn op_user_worker_create(
     state: Rc<RefCell<OpState>>,
     #[serde] opts: UserWorkerCreateOptions,
-) -> Result<String, AnyError> {
+) -> Result<(String, ResourceId), AnyError> {
     let result_rx = {
-        let op_state = state.borrow();
-        let tx = op_state.borrow::<mpsc::UnboundedSender<UserWorkerMsgs>>();
-        let (result_tx, result_rx) = oneshot::channel::<Result<CreateUserWorkerResult, Error>>();
-
+        let (tx, rx) = oneshot::channel::<Result<CreateUserWorkerResult, Error>>();
         let UserWorkerCreateOptions {
             service_path,
             env_vars,
@@ -183,7 +180,7 @@ pub async fn op_user_worker_create(
                             deno_core::resolve_url_or_path(
                                 // FIXME: The type alias does not have a unique
                                 // type id and should not be used here.
-                                op_state.borrow::<ModuleSpecifier>().as_str(),
+                                state.borrow().borrow::<ModuleSpecifier>().as_str(),
                                 std::env::current_dir()?.as_path(),
                             )?
                         },
@@ -195,8 +192,12 @@ pub async fn op_user_worker_create(
             maybe_tmp_fs_config,
         };
 
-        tx.send(UserWorkerMsgs::Create(user_worker_options, result_tx))?;
-        result_rx
+        state
+            .borrow()
+            .borrow::<mpsc::UnboundedSender<UserWorkerMsgs>>()
+            .send(UserWorkerMsgs::Create(user_worker_options, tx))?;
+
+        rx
     };
 
     match result_rx.await {
@@ -209,7 +210,13 @@ pub async fn op_user_worker_create(
         )),
 
         Ok(Err(err)) => Err(custom_error("InvalidWorkerCreation", format!("{err:#}"))),
-        Ok(Ok(v)) => Ok(v.key.to_string()),
+        Ok(Ok(CreateUserWorkerResult { key, token })) => Ok((
+            key.to_string(),
+            state
+                .borrow_mut()
+                .resource_table
+                .add(UserWorkerCancellationToken(token.clone())),
+        )),
     }
 }
 
@@ -586,3 +593,37 @@ impl Stream for BodyStream {
         self.0.poll_recv(cx)
     }
 }
+
+struct UserWorkerCancellationToken(CancellationToken);
+
+impl Resource for UserWorkerCancellationToken {
+    fn name(&self) -> std::borrow::Cow<str> {
+        std::any::type_name::<Self>().into()
+    }
+}
+
+#[op2(async)]
+#[serde]
+pub async fn op_user_user_worker_wait_token_cancelled(
+    state: Rc<RefCell<OpState>>,
+    #[smi] rid: ResourceId,
+) -> Result<(), AnyError> {
+    let token = state
+        .borrow()
+        .resource_table
+        .get::<UserWorkerCancellationToken>(rid)?
+        .0
+        .clone();
+
+    token.cancelled().await;
+    Ok(())
+}
+
+#[op2(fast)]
+pub fn op_user_worker_is_active(state: &mut OpState, #[smi] rid: ResourceId) -> bool {
+    state
+        .resource_table
+        .get::<UserWorkerCancellationToken>(rid)
+        .map(|it| !it.0.is_cancelled())
+        .unwrap_or_default()
+}

Original file line number	Diff line number	Diff line change
`@@ -263,11 +263,13 @@ impl WorkerPool {`
`263`	`263`	`.as_user_worker()`
`264`	`264`	`.map_or(false, \|it\| !is_oneshot_policy && it.force_create);`
`265`	`265`
`266`		`- if let Some(ref active_worker_uuid) = self.maybe_active_worker(&service_path, force_create)`
	`266`	`+ if let Some((ref active_worker_uuid, profile)) =`
	`267`	`+ self.maybe_active_worker(&service_path, force_create)`
`267`	`268`	`{`
`268`	`269`	`if tx`
`269`	`270`	`.send(Ok(CreateUserWorkerResult {`
`270`	`271`	`key: *active_worker_uuid,`
	`272`	`+ token: profile.cancel.clone(),`
`271`	`273`	`}))`
`272`	`274`	`.is_err()`
`273`	`275`	`{`
`@@ -448,7 +450,7 @@ impl WorkerPool {`
`448`	`450`	`permit: permit.map(Arc::new),`
`449`	`451`	`status: status.clone(),`
`450`	`452`	`exit: surface.exit,`
`451`		`- cancel,`
	`453`	`+ cancel: cancel.clone(),`
`452`	`454`	`};`
`453`	`455`
`454`	`456`	`if worker_pool_msgs_tx`
`@@ -457,7 +459,13 @@ impl WorkerPool {`
`457`	`459`	`{`
`458`	`460`	`error!("user worker msgs receiver dropped")`
`459`	`461`	`}`
`460`		`- if tx.send(Ok(CreateUserWorkerResult { key: uuid })).is_err() {`
	`462`	`+ if tx`
	`463`	`+ .send(Ok(CreateUserWorkerResult {`
	`464`	`+ key: uuid,`
	`465`	`+ token: cancel,`
	`466`	`+ }))`
	`467`	`+ .is_err()`
	`468`	`+ {`
`461`	`469`	`error!("main worker receiver dropped")`
`462`	`470`	`};`
`463`	`471`	`}`
`@@ -632,7 +640,11 @@ impl WorkerPool {`
`632`	`640`	`}`
`633`	`641`	`}`
`634`	`642`
`635`		`- fn maybe_active_worker(&mut self, service_path: &String, force_create: bool) -> Option<Uuid> {`
	`643`	`+ fn maybe_active_worker(`
	`644`	`+ &mut self,`
	`645`	`+ service_path: &String,`
	`646`	`+ force_create: bool,`
	`647`	`+ ) -> Option<(Uuid, &UserWorkerProfile)> {`
`636`	`648`	`if force_create {`
`637`	`649`	`return None;`
`638`	`650`	`}`
`@@ -648,11 +660,13 @@ impl WorkerPool {`
`648`	`660`	`.get(&worker_uuid)`
`649`	`661`	`.map(\|it\| it.status.is_retired.clone())`
`650`	`662`	`{`
`651`		`- Some(is_retired) if !is_retired.is_raised() => Some(worker_uuid),`
	`663`	`+ Some(is_retired) if !is_retired.is_raised() => {`
	`664`	`+ Some((worker_uuid, self.user_workers.get(&worker_uuid).unwrap()))`
	`665`	`+ }`
`652`	`666`
`653`	`667`	`_ => {`
`654`	`668`	`self.retire(&worker_uuid);`
`655`		`- self.maybe_active_worker(service_path, force_create)`
	`669`	`+ self.maybe_active_worker(service_path, false)`
`656`	`670`	`}`
`657`	`671`	`}`
`658`	`672`	`}`
Original file line number	Diff line number	Diff line change
`@@ -245,6 +245,7 @@ pub type SendRequestResult = (Response<Body>, mpsc::UnboundedSender<()>);`
`245`	`245`	`#[derive(Debug)]`
`246`	`246`	`pub struct CreateUserWorkerResult {`
`247`	`247`	`pub key: Uuid,`
	`248`	`+ pub token: CancellationToken,`
`248`	`249`	`}`
`249`	`250`
`250`	`251`	`#[derive(Debug)]`