stamp: initial impl for worker heap metrics

Author: nyannyacha

crates/base/src/deno_runtime.rs

@@ -573,7 +573,12 @@ mod test {
             maybe_module_code: Some(FastString::from(String::from(
                 "Deno.serve((req) => new Response('Hello World'));",
             ))),
-            conf: { WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts { worker_pool_tx }) },
+            conf: {
+                WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
+                    worker_pool_tx,
+                    event_worker_metric_src: None,
+                })
+            },
         })
         .await
         .expect("It should not panic");
@@ -612,7 +617,12 @@ mod test {
             maybe_eszip: Some(EszipPayloadKind::VecKind(eszip_code)),
             maybe_entrypoint: None,
             maybe_module_code: None,
-            conf: { WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts { worker_pool_tx }) },
+            conf: {
+                WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
+                    worker_pool_tx,
+                    event_worker_metric_src: None,
+                })
+            },
         })
         .await;
 
@@ -673,7 +683,12 @@ mod test {
             maybe_eszip: Some(EszipPayloadKind::VecKind(eszip_code)),
             maybe_entrypoint: None,
             maybe_module_code: None,
-            conf: { WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts { worker_pool_tx }) },
+            conf: {
+                WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
+                    worker_pool_tx,
+                    event_worker_metric_src: None,
+                })
+            },
         })
         .await;
 
@@ -731,7 +746,10 @@ mod test {
                 if let Some(uc) = user_conf {
                     uc
                 } else {
-                    WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts { worker_pool_tx })
+                    WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
+                        worker_pool_tx,
+                        event_worker_metric_src: None,
+                    })
                 }
             },
         })

crates/base/src/rt_worker/worker.rs

@@ -11,6 +11,7 @@ use event_worker::events::{
 use futures_util::FutureExt;
 use log::{debug, error};
 use sb_core::conn_sync::ConnSync;
+use sb_core::{RuntimeMetricSource, WorkerMetricSource};
 use sb_workers::context::{UserWorkerMsgs, WorkerContextInitOpts};
 use std::any::Any;
 use std::future::{pending, Future};
@@ -87,7 +88,7 @@ impl Worker {
             UnboundedSender<UnixStreamEntry>,
             UnboundedReceiver<UnixStreamEntry>,
         ),
-        booter_signal: Sender<Result<(), Error>>,
+        booter_signal: Sender<Result<WorkerMetricSource, Error>>,
         termination_token: Option<TerminationToken>,
     ) {
         let worker_name = self.worker_name.clone();
@@ -101,24 +102,47 @@ impl Worker {
 
         let method_cloner = self.clone();
         let timing = opts.timing.take();
-        let is_user_worker = opts.conf.is_user_worker();
+        let worker_kind = opts.conf.to_worker_kind();
+        let maybe_event_worker_metric_src = opts
+            .conf
+            .as_main_worker()
+            .as_ref()
+            .and_then(|it| it.event_worker_metric_src.clone());
 
         let cancel = self.cancel.clone();
-        let rt = if is_user_worker {
+        let rt = if worker_kind.is_user_worker() {
             &rt::USER_WORKER_RT
         } else {
             &rt::PRIMARY_WORKER_RT
         };
 
         let _worker_handle = rt.spawn_pinned(move || {
             tokio::task::spawn_local(async move {
-                let (maybe_cpu_usage_metrics_tx, maybe_cpu_usage_metrics_rx) = is_user_worker
+                let (maybe_cpu_usage_metrics_tx, maybe_cpu_usage_metrics_rx) = worker_kind
+                    .is_user_worker()
                     .then(unbounded_channel::<CPUUsageMetrics>)
                     .unzip();
 
                 let result = match DenoRuntime::new(opts).await {
                     Ok(mut new_runtime) => {
-                        let _ = booter_signal.send(Ok(()));
+                        let metric = {
+                            let js_runtime = &mut new_runtime.js_runtime;
+                            let metric = WorkerMetricSource::from_js_runtime(js_runtime);
+
+                            if worker_kind.is_main_worker() {
+                                let state = js_runtime.op_state();
+                                let mut state_mut = state.borrow_mut();
+
+                                state_mut.put(RuntimeMetricSource::new(
+                                    metric.clone(),
+                                    maybe_event_worker_metric_src,
+                                ));
+                            }
+
+                            metric
+                        };
+
+                        let _ = booter_signal.send(Ok(metric));
 
                         // CPU TIMER
                         let (termination_event_tx, termination_event_rx) =
@@ -127,7 +151,7 @@ impl Worker {
                         let _cpu_timer;
 
                         // TODO: Allow customization of supervisor
-                        let termination_fut = if is_user_worker {
+                        let termination_fut = if worker_kind.is_user_worker() {
                             // cputimer is returned from supervisor and assigned here to keep it in scope.
                             let Ok(maybe_timer) = create_supervisor(
                                 worker_key.unwrap_or(Uuid::nil()),
@@ -209,7 +233,7 @@ impl Worker {
                         let result = data.await;
 
                         if let Some(token) = termination_token.as_ref() {
-                            if !is_user_worker {
+                            if !worker_kind.is_user_worker() {
                                 let _ = termination_fut.await;
                             }
 

crates/base/src/rt_worker/worker_ctx.rs

@@ -12,6 +12,7 @@ use event_worker::events::{
 use hyper::{Body, Request, Response};
 use log::{debug, error};
 use sb_core::conn_sync::ConnSync;
+use sb_core::WorkerMetricSource;
 use sb_graph::EszipPayloadKind;
 use sb_workers::context::{
     EventWorkerRuntimeOpts, MainWorkerRuntimeOpts, Timing, UserWorkerMsgs, WorkerContextInitOpts,
@@ -309,9 +310,10 @@ impl CreateWorkerArgs {
 
 pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
     init_opts: Opt,
-) -> Result<mpsc::UnboundedSender<WorkerRequestMsg>, Error> {
-    let (worker_boot_result_tx, worker_boot_result_rx) = oneshot::channel::<Result<(), Error>>();
+) -> Result<(WorkerMetricSource, mpsc::UnboundedSender<WorkerRequestMsg>), Error> {
     let (unix_stream_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStreamEntry>();
+    let (worker_boot_result_tx, worker_boot_result_rx) =
+        oneshot::channel::<Result<WorkerMetricSource, Error>>();
 
     let CreateWorkerArgs(init_opts, maybe_supervisor_policy, maybe_termination_token) =
         init_opts.into();
@@ -370,7 +372,7 @@ pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
 
                 bail!(err)
             }
-            Ok(_) => {
+            Ok(metric) => {
                 let elapsed = worker_struct_ref
                     .worker_boot_start_time
                     .elapsed()
@@ -384,7 +386,7 @@ pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
                     worker_struct_ref.event_metadata.clone(),
                 );
 
-                Ok(worker_req_tx)
+                Ok((metric, worker_req_tx))
             }
         }
     } else {
@@ -422,7 +424,7 @@ pub async fn create_main_worker(
     main_worker_path: PathBuf,
     import_map_path: Option<String>,
     no_module_cache: bool,
-    user_worker_msgs_tx: mpsc::UnboundedSender<UserWorkerMsgs>,
+    runtime_opts: MainWorkerRuntimeOpts,
     maybe_entrypoint: Option<String>,
     termination_token: Option<TerminationToken>,
 ) -> Result<mpsc::UnboundedSender<WorkerRequestMsg>, Error> {
@@ -435,7 +437,7 @@ pub async fn create_main_worker(
         }
     }
 
-    let main_worker_req_tx = create_worker((
+    let (_, sender) = create_worker((
         WorkerContextInitOpts {
             service_path,
             import_map_path,
@@ -445,17 +447,15 @@ pub async fn create_main_worker(
             maybe_eszip,
             maybe_entrypoint,
             maybe_module_code: None,
-            conf: WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
-                worker_pool_tx: user_worker_msgs_tx,
-            }),
+            conf: WorkerRuntimeOpts::MainWorker(runtime_opts),
             env_vars: std::env::vars().collect(),
         },
         termination_token,
     ))
     .await
     .map_err(|err| anyhow!("main worker boot error: {}", err))?;
 
-    Ok(main_worker_req_tx)
+    Ok(sender)
 }
 
 pub async fn create_events_worker(
@@ -464,7 +464,13 @@ pub async fn create_events_worker(
     no_module_cache: bool,
     maybe_entrypoint: Option<String>,
     termination_token: Option<TerminationToken>,
-) -> Result<mpsc::UnboundedSender<WorkerEventWithMetadata>, Error> {
+) -> Result<
+    (
+        WorkerMetricSource,
+        mpsc::UnboundedSender<WorkerEventWithMetadata>,
+    ),
+    Error,
+> {
     let (events_tx, events_rx) = mpsc::unbounded_channel::<WorkerEventWithMetadata>();
 
     let mut service_path = events_worker_path.clone();
@@ -478,7 +484,7 @@ pub async fn create_events_worker(
         }
     }
 
-    let _ = create_worker((
+    let (metric, _) = create_worker((
         WorkerContextInitOpts {
             service_path,
             no_module_cache,
@@ -496,7 +502,7 @@ pub async fn create_events_worker(
     .await
     .map_err(|err| anyhow!("events worker boot error: {}", err))?;
 
-    Ok(events_tx)
+    Ok((metric, events_tx))
 }
 
 pub async fn create_user_worker_pool(

crates/base/src/rt_worker/worker_pool.rs

@@ -404,7 +404,7 @@ impl WorkerPool {
             match create_worker((worker_options, supervisor_policy, termination_token.clone()))
                 .await
             {
-                Ok(worker_request_msg_tx) => {
+                Ok((_, worker_request_msg_tx)) => {
                     let profile = UserWorkerProfile {
                         worker_request_msg_tx,
                         timing_tx_pair: (req_start_timing_tx, req_end_timing_tx),

crates/base/src/server.rs

@@ -8,7 +8,7 @@ use futures_util::Stream;
 use hyper::{server::conn::Http, service::Service, Body, Request, Response};
 use log::{debug, error, info};
 use sb_core::conn_sync::ConnSync;
-use sb_workers::context::WorkerRequestMsg;
+use sb_workers::context::{MainWorkerRuntimeOpts, WorkerRequestMsg};
 use std::future::Future;
 use std::net::IpAddr;
 use std::net::Ipv4Addr;
@@ -180,17 +180,17 @@ impl Server {
         entrypoints: WorkerEntrypoints,
         termination_token: Option<TerminationToken>,
     ) -> Result<Self, Error> {
-        let mut worker_events_sender: Option<mpsc::UnboundedSender<WorkerEventWithMetadata>> = None;
+        let mut worker_events_tx: Option<mpsc::UnboundedSender<WorkerEventWithMetadata>> = None;
         let maybe_events_entrypoint = entrypoints.events;
         let maybe_main_entrypoint = entrypoints.main;
         let termination_token = termination_token.unwrap_or_default();
 
         // Create Event Worker
-        if let Some(events_service_path) = maybe_events_service_path {
+        let event_worker_metric_src = if let Some(events_service_path) = maybe_events_service_path {
             let events_path = Path::new(&events_service_path);
             let events_path_buf = events_path.to_path_buf();
 
-            let events_worker = create_events_worker(
+            let (metric, sender) = create_events_worker(
                 events_path_buf,
                 import_map_path.clone(),
                 no_module_cache,
@@ -199,13 +199,16 @@ impl Server {
             )
             .await?;
 
-            worker_events_sender = Some(events_worker);
-        }
+            worker_events_tx = Some(sender);
+            Some(metric)
+        } else {
+            None
+        };
 
         // Create a user worker pool
-        let user_worker_msgs_tx = create_user_worker_pool(
+        let worker_pool_tx = create_user_worker_pool(
             maybe_user_worker_policy.unwrap_or_default(),
-            worker_events_sender,
+            worker_events_tx,
             None,
         )
         .await?;
@@ -216,7 +219,10 @@ impl Server {
             main_worker_path,
             import_map_path.clone(),
             no_module_cache,
-            user_worker_msgs_tx,
+            MainWorkerRuntimeOpts {
+                worker_pool_tx,
+                event_worker_metric_src,
+            },
             maybe_main_entrypoint,
             Some(termination_token.child_token()),
         )

crates/base/src/utils/integration_test_helper.rs

@@ -192,7 +192,7 @@ impl TestBedBuilder {
     }
 
     pub async fn build(self) -> TestBed {
-        let (worker_pool_msg_tx, pool_termination_token) = {
+        let (worker_pool_tx, pool_termination_token) = {
             let token = TerminationToken::new();
             (
                 create_user_worker_pool(
@@ -218,12 +218,13 @@ impl TestBedBuilder {
             maybe_entrypoint: None,
             maybe_module_code: None,
             conf: WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
-                worker_pool_tx: worker_pool_msg_tx,
+                worker_pool_tx,
+                event_worker_metric_src: None,
             }),
         };
 
         let main_termination_token = TerminationToken::new();
-        let main_worker_msg_tx =
+        let (_, main_worker_msg_tx) =
             create_worker((main_worker_init_opts, main_termination_token.clone()))
                 .await
                 .unwrap();
@@ -293,20 +294,24 @@ pub async fn create_test_user_worker<Opt: Into<CreateTestUserWorkerArgs>>(
         ..Default::default()
     });
 
-    Ok((
-        create_worker(
+    Ok({
+        let (_, sender) = create_worker(
             opts.with_policy(policy)
                 .with_termination_token(termination_token.clone()),
         )
-        .await?,
-        RequestScope {
-            policy,
-            req_start_tx,
-            req_end_tx,
-            termination_token,
-            conn: (Some(conn_tx), conn_rx),
-        },
-    ))
+        .await?;
+
+        (
+            sender,
+            RequestScope {
+                policy,
+                req_start_tx,
+                req_end_tx,
+                termination_token,
+                conn: (Some(conn_tx), conn_rx),
+            },
+        )
+    })
 }
 
 pub fn test_user_worker_pool_policy() -> WorkerPoolPolicy {