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

nyannyacha · nyannyacha · commit 3476038bd742 · 2024-07-04T00:25:38.000Z
diff --git a/crates/base/src/rt_worker/worker_ctx.rs b/crates/base/src/rt_worker/worker_ctx.rs
@@ -575,8 +575,7 @@ pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
     maybe_request_idle_timeout: Option<u64>,
 ) -> 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
@@ -260,11 +260,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()
             {
@@ -437,7 +437,7 @@ impl WorkerPool {
                         permit: permit.map(Arc::new),
                         status: status.clone(),
                         exit: ctx.exit,
-                        cancel,
+                        cancel: cancel.clone(),
                     };
 
                     if worker_pool_msgs_tx
@@ -446,12 +446,20 @@ 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")
                     };
 
                     status.demand.fetch_add(1, Ordering::Release);
                 }
+
                 Err(e) => {
                     if tx.send(Err(e)).is_err() {
                         error!("main worker receiver dropped")
@@ -622,7 +630,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;
         }
@@ -639,18 +651,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
@@ -221,6 +221,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
@@ -43,6 +43,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",]
@@ -82,11 +84,11 @@ 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>>();
@@ -182,7 +184,13 @@ pub async fn op_user_worker_create(
     let result = result.unwrap();
     match result {
         Err(e) => Err(custom_error("InvalidWorkerCreation", e.to_string())),
-        Ok(res) => Ok(res.key.to_string()),
+        Ok(CreateUserWorkerResult { key, token }) => Ok((
+            key.to_string(),
+            state
+                .borrow_mut()
+                .resource_table
+                .add(UserWorkerCancellationToken(token.clone())),
+        )),
     }
 }
 
@@ -559,3 +567,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
`@@ -221,6 +221,7 @@ pub type SendRequestResult = (Response<Body>, mpsc::UnboundedSender<()>);`
`221`	`221`	`#[derive(Debug)]`
`222`	`222`	`pub struct CreateUserWorkerResult {`
`223`	`223`	`pub key: Uuid,`
	`224`	`+ pub token: CancellationToken,`
`224`	`225`	`}`
`225`	`226`
`226`	`227`	`#[derive(Debug)]`