feat: add EdgeRuntime.userWorkers.memStats (#560)

* feat: add `EdgeRuntime.userWorkers.memStats`

* chore: update `rust-toolchain.toml`

Author: nyannyacha

Cargo.lock

@@ -3,6 +3,7 @@ use crate::server::ServerFlags;
 use crate::worker::WorkerSurfaceBuilder;
 
 use anyhow::{anyhow, bail, Context, Error};
+use base_mem_check::WorkerHeapStatistics;
 use deno_config::JsxImportSourceConfig;
 use either::Either::Left;
 use enum_as_inner::EnumAsInner;
@@ -449,6 +450,7 @@ impl WorkerPool {
                         permit: permit.map(Arc::new),
                         status: status.clone(),
                         exit: surface.exit,
+                        mem_check: surface.mem_check.clone(),
                         cancel,
                     };
 
@@ -657,6 +659,23 @@ impl WorkerPool {
             }
         }
     }
+
+    fn memory_usage(&self, tx: Sender<HashMap<Uuid, Option<WorkerHeapStatistics>>>) {
+        let mem_checks = self
+            .user_workers
+            .iter()
+            .map(|it| (it.0.clone(), it.1.mem_check.clone()))
+            .collect::<Vec<_>>();
+
+        drop(tokio::task::spawn_blocking(move || {
+            let mut results = HashMap::new();
+            for (uuid, mem_check) in mem_checks {
+                results.insert(uuid, mem_check.read().ok().map(|it| it.current.clone()));
+            }
+
+            let _ = tx.send(results);
+        }));
+    }
 }
 
 pub async fn create_user_worker_pool(
@@ -750,6 +769,10 @@ pub async fn create_user_worker_pool(
                                     break;
                                 }
                             }
+
+                            Some(UserWorkerMsgs::InqueryMemoryUsage(tx)) => {
+                                worker_pool.memory_usage(tx);
+                            }
                         }
                     }
                 }

crates/base/src/worker/pool.rs

@@ -4,6 +4,7 @@ use crate::server::ServerFlags;
 use crate::worker::utils::{get_event_metadata, send_event_if_event_worker_available};
 
 use anyhow::Error;
+use base_mem_check::MemCheckState;
 use base_rt::error::CloneableError;
 use deno_core::unsync::MaskFutureAsSend;
 use futures_util::FutureExt;
@@ -17,7 +18,7 @@ use sb_workers::context::{
     WorkerRequestMsg,
 };
 use std::future::ready;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use tokio::io;
 use tokio::sync::{mpsc, oneshot};
 use tokio::time::Instant;
@@ -183,7 +184,9 @@ impl std::ops::Deref for Worker {
 impl Worker {
     pub fn start(
         self,
-        booter_signal: oneshot::Sender<Result<(MetricSource, CancellationToken), Error>>,
+        booter_signal: oneshot::Sender<
+            Result<(MetricSource, Arc<RwLock<MemCheckState>>, CancellationToken), Error>,
+        >,
         exit: WorkerExit,
     ) {
         let worker_name = self.worker_name.clone();
@@ -242,7 +245,11 @@ impl Worker {
                 }
             };
 
-            let _ = booter_signal.send(Ok((metric_src, runtime.drop_token.clone())));
+            let _ = booter_signal.send(Ok((
+                metric_src,
+                runtime.mem_check_state(),
+                runtime.drop_token.clone(),
+            )));
             let supervise_fut = match imp.clone().supervise(&mut runtime) {
                 Some(v) => v.boxed(),
                 None if worker_kind.is_user_worker() => return None,
@@ -325,4 +332,5 @@ pub struct WorkerSurface {
     pub msg_tx: mpsc::UnboundedSender<WorkerRequestMsg>,
     pub exit: WorkerExit,
     pub cancel: CancellationToken,
+    pub mem_check: Arc<RwLock<MemCheckState>>,
 }

crates/base/src/worker/worker_inner.rs

@@ -1,9 +1,10 @@
 use std::{
     path::{Path, PathBuf},
-    sync::Arc,
+    sync::{Arc, RwLock},
 };
 
 use anyhow::Context;
+use base_mem_check::MemCheckState;
 use deno_config::JsxImportSourceConfig;
 use either::Either;
 use graph::{DecoratorType, EszipPayloadKind};
@@ -329,8 +330,9 @@ 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::<
+            Result<(MetricSource, Arc<RwLock<MemCheckState>>, CancellationToken), anyhow::Error>,
+        >();
 
         let flags = flags.unwrap_or_default();
         let init_opts = init_opts.context("init_opts must be specified")?;
@@ -377,7 +379,7 @@ impl WorkerSurfaceBuilder {
 
         // wait for worker to be successfully booted
         match worker_boot_result_rx.await? {
-            Ok((metric, cancel)) => {
+            Ok((metric, mem_check_state, cancel)) => {
                 let elapsed = cx.worker_boot_start_time.elapsed().as_millis();
 
                 send_event_if_event_worker_available(
@@ -393,6 +395,7 @@ impl WorkerSurfaceBuilder {
                     msg_tx: worker_req_tx,
                     exit,
                     cancel,
+                    mem_check: mem_check_state,
                 })
             }
 

crates/base/src/worker/worker_surface_creation.rs

@@ -19,6 +19,7 @@ fs.workspace = true
 
 sb_core.workspace = true
 sb_event_worker.workspace = true
+base_mem_check.workspace = true
 
 anyhow.workspace = true
 uuid.workspace = true

ext/workers/Cargo.toml

@@ -1,4 +1,5 @@
 use anyhow::{anyhow, Error};
+use base_mem_check::{MemCheckState, WorkerHeapStatistics};
 use deno_config::JsxImportSourceConfig;
 use deno_core::FastString;
 use enum_as_inner::EnumAsInner;
@@ -10,6 +11,7 @@ use sb_core::{MetricSource, SharedMetricSource};
 use sb_event_worker::events::{UncaughtExceptionEvent, WorkerEventWithMetadata};
 use std::path::PathBuf;
 use std::sync::atomic::AtomicUsize;
+use std::sync::RwLock;
 use std::{collections::HashMap, sync::Arc};
 use tokio::sync::mpsc::unbounded_channel;
 use tokio::sync::{mpsc, oneshot, Mutex, Notify, OwnedSemaphorePermit};
@@ -121,6 +123,7 @@ pub struct UserWorkerProfile {
     pub cancel: CancellationToken,
     pub status: TimingStatus,
     pub exit: WorkerExit,
+    pub mem_check: Arc<RwLock<MemCheckState>>,
 }
 
 #[derive(Debug, Clone)]
@@ -238,6 +241,7 @@ pub enum UserWorkerMsgs {
     ),
     Idle(Uuid),
     Shutdown(Uuid),
+    InqueryMemoryUsage(oneshot::Sender<HashMap<Uuid, Option<WorkerHeapStatistics>>>),
 }
 
 pub type SendRequestResult = (Response<Body>, mpsc::UnboundedSender<()>);

ext/workers/context.rs

@@ -6,6 +6,7 @@ use crate::context::{
     WorkerRuntimeOpts,
 };
 use anyhow::Error;
+use base_mem_check::WorkerHeapStatistics;
 use context::SendRequestResult;
 use deno_config::JsxImportSourceConfig;
 use deno_core::error::{custom_error, type_error, AnyError};
@@ -31,6 +32,7 @@ use once_cell::sync::Lazy;
 use sb_core::conn_sync::ConnWatcher;
 use serde::{Deserialize, Serialize};
 use std::cell::RefCell;
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::rc::Rc;
@@ -45,6 +47,7 @@ deno_core::extension!(
         op_user_worker_create,
         op_user_worker_fetch_build,
         op_user_worker_fetch_send,
+        op_user_worker_mem_stats,
     ],
     esm_entry_point = "ext:sb_user_workers/user_workers.js",
     esm = ["user_workers.js",]
@@ -579,6 +582,24 @@ pub async fn op_user_worker_fetch_send(
     Ok(response)
 }
 
+#[op2(async)]
+#[serde]
+pub async fn op_user_worker_mem_stats(
+    state: Rc<RefCell<OpState>>,
+) -> Result<HashMap<Uuid, Option<WorkerHeapStatistics>>, Error> {
+    let mem_rx = {
+        let op_state = state.borrow();
+        let tx = op_state.borrow::<mpsc::UnboundedSender<UserWorkerMsgs>>();
+
+        let (mem_tx, mem_rx) = oneshot::channel();
+        let _ = tx.send(UserWorkerMsgs::InqueryMemoryUsage(mem_tx));
+
+        mem_rx
+    };
+
+    Ok(mem_rx.await?)
+}
+
 /// Wraps a [`mpsc::Receiver`] in a [`Stream`] that can be used as a Hyper [`Body`].
 pub struct BodyStream(pub mpsc::Receiver<Result<bytes::Bytes, Error>>);
 

ext/workers/lib.rs

@@ -9,6 +9,7 @@ const { TypeError } = primordials;
 const {
 	op_user_worker_fetch_send,
 	op_user_worker_create,
+	op_user_worker_mem_stats,
 } = ops;
 
 const NO_SUPABASE_TAG_WARN_MSG = `Unable to find the supabase tag from the request instance.\n\
@@ -140,6 +141,10 @@ class UserWorker {
 
 		return new UserWorker(key);
 	}
+
+	static async memStats() {
+		return await op_user_worker_mem_stats();
+	}
 }
 
 const SUPABASE_USER_WORKERS = UserWorker;

ext/workers/user_workers.js

@@ -1,2 +1,3 @@
 [toolchain]
-channel = "1.79.0"
+channel = "1.79.0"
+components = ["rustfmt", "clippy"]

rust-toolchain.toml

@@ -96,6 +96,7 @@ declare namespace EdgeRuntime {
 
         fetch(request: Request, options?: UserWorkerFetchOptions): Promise<Response>;
         static create(opts: UserWorkerCreateOptions): Promise<UserWorker>;
+        static memStats(): Promise<HeapStatistics>;
     }
 
     export function scheduleTermination(): void;