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

nyannyacha · nyannyacha · commit 9678c214ad35 · 2024-11-25T08:03:29.000+09:00
diff --git a/crates/base/src/rt_worker/worker_ctx.rs b/crates/base/src/rt_worker/worker_ctx.rs
@@ -599,8 +599,7 @@ pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
     inspector: Option<Inspector>,
 ) -> Result<WorkerCtx, Error> {
     let (duplex_stream_tx, duplex_stream_rx) = mpsc::unbounded_channel::<DuplexStreamEntry>();
-    let (worker_boot_result_tx, worker_boot_result_rx) =
-        oneshot::channel::<Result<MetricSource, Error>>();
+    let (worker_boot_result_tx, worker_boot_result_rx) = oneshot::channel();
 
     let CreateWorkerArgs(worker_init_opts, maybe_supervisor_policy, maybe_termination_token) =
         init_opts.into();
diff --git a/crates/base/src/rt_worker/worker_pool.rs b/crates/base/src/rt_worker/worker_pool.rs
@@ -258,11 +258,11 @@ 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((key, profile)) = self.maybe_active_worker(&service_path, force_create) {
             if tx
                 .send(Ok(CreateUserWorkerResult {
-                    key: *active_worker_uuid,
+                    key,
+                    token: profile.cancel.clone(),
                 }))
                 .is_err()
             {
@@ -440,7 +440,7 @@ impl WorkerPool {
                         permit: permit.map(Arc::new),
                         status: status.clone(),
                         exit: ctx.exit,
-                        cancel,
+                        cancel: cancel.clone(),
                     };
 
                     if worker_pool_msgs_tx
@@ -449,7 +449,14 @@ 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")
                     };
 
@@ -624,7 +631,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;
         }
@@ -641,18 +652,15 @@ impl WorkerPool {
             .map(|it| it.status.is_retired.clone())
         {
             Some(is_retired) if !is_retired.is_raised() => {
-                self.user_workers
-                    .get(&worker_uuid)
-                    .map(|it| it.status.demand.as_ref())
-                    .unwrap()
-                    .fetch_add(1, Ordering::Release);
+                let profile = self.user_workers.get(&worker_uuid).unwrap();
 
-                Some(worker_uuid)
+                profile.status.demand.fetch_add(1, Ordering::Release);
+                Some((worker_uuid, profile))
             }
 
             _ => {
                 self.retire(&worker_uuid);
-                self.maybe_active_worker(service_path, force_create)
+                self.maybe_active_worker(service_path, false)
             }
         }
     }
diff --git a/crates/sb_workers/context.rs b/crates/sb_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/crates/sb_workers/lib.rs b/crates/sb_workers/lib.rs
@@ -45,6 +45,8 @@ deno_core::extension!(
         op_user_worker_create,
         op_user_worker_fetch_build,
         op_user_worker_fetch_send,
+        op_user_user_worker_wait_token_cancelled,
+        op_user_worker_is_active,
     ],
     esm_entry_point = "ext:sb_user_workers/user_workers.js",
     esm = ["user_workers.js",]
@@ -94,16 +96,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 +182,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 +194,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 +212,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 +595,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
`@@ -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)]`