Track instance_queue_length by reducer (#530)

Co-authored-by: Noa <[email protected]>

Author: 2 people

crates/core/src/host/module_host.rs

@@ -12,8 +12,9 @@ use crate::identity::Identity;
 use crate::json::client_api::{SubscriptionUpdateJson, TableRowOperationJson, TableUpdateJson};
 use crate::protobuf::client_api::{table_row_operation, SubscriptionUpdate, TableRowOperation, TableUpdate};
 use crate::subscription::module_subscription_actor::ModuleSubscriptionManager;
-use crate::util::lending_pool::{Closed, LendingPool, LentResource, PoolClosed};
+use crate::util::lending_pool::{Closed, LendingPool, LentResource, PoolClosed, WaiterGauge};
 use crate::util::notify_once::NotifyOnce;
+use crate::util::prometheus_handle::{GaugeInc, IntGaugeExt};
 use crate::worker_metrics::WORKER_METRICS;
 use base64::{engine::general_purpose::STANDARD as BASE_64_STD, Engine as _};
 use futures::{Future, FutureExt};
@@ -333,7 +334,10 @@ impl fmt::Debug for ModuleHost {
 
 #[async_trait::async_trait]
 trait DynModuleHost: Send + Sync + 'static {
-    async fn get_instance(&self) -> Result<(&HostThreadpool, Box<dyn ModuleInstance>), NoSuchModule>;
+    async fn get_instance(
+        &self,
+        waiter_gauge_context: <InstancePoolGauge as WaiterGauge>::Context<'_>,
+    ) -> Result<(&HostThreadpool, Box<dyn ModuleInstance>), NoSuchModule>;
     fn inject_logs(&self, log_level: LogLevel, message: &str);
     fn one_off_query(
         &self,
@@ -349,7 +353,7 @@ trait DynModuleHost: Send + Sync + 'static {
 struct HostControllerActor<T: Module> {
     module: Arc<T>,
     threadpool: Arc<HostThreadpool>,
-    instance_pool: LendingPool<T::Instance>,
+    instance_pool: LendingPool<T::Instance, InstancePoolGauge>,
     start: NotifyOnce,
 }
 
@@ -377,18 +381,38 @@ async fn select_first<A: Future, B: Future<Output = ()>>(fut_a: A, fut_b: B) ->
     }
 }
 
+#[derive(Clone)]
+struct InstancePoolGauge;
+
+impl WaiterGauge for InstancePoolGauge {
+    type Context<'a> = &'a (&'a Identity, &'a Hash, &'a Address, &'a str);
+    type IncGuard = GaugeInc;
+    fn inc(&self, &(identity, module_hash, database_address, reducer_symbol): Self::Context<'_>) -> Self::IncGuard {
+        WORKER_METRICS
+            .instance_queue_length
+            .with_label_values(identity, module_hash, database_address, reducer_symbol)
+            .inc_scope()
+    }
+}
+
 #[async_trait::async_trait]
 impl<T: Module> DynModuleHost for HostControllerActor<T> {
-    async fn get_instance(&self) -> Result<(&HostThreadpool, Box<dyn ModuleInstance>), NoSuchModule> {
+    async fn get_instance(
+        &self,
+        waiter_gauge_context: <InstancePoolGauge as WaiterGauge>::Context<'_>,
+    ) -> Result<(&HostThreadpool, Box<dyn ModuleInstance>), NoSuchModule> {
         self.start.notified().await;
         // in the future we should do something like in the else branch here -- add more instances based on load.
         // we need to do write-skew retries first - right now there's only ever once instance per module.
         let inst = if true {
-            self.instance_pool.request().await.map_err(|_| NoSuchModule)?
+            self.instance_pool
+                .request_with_context(waiter_gauge_context)
+                .await
+                .map_err(|_| NoSuchModule)?
         } else {
             const GET_INSTANCE_TIMEOUT: Duration = Duration::from_millis(500);
             select_first(
-                self.instance_pool.request(),
+                self.instance_pool.request_with_context(waiter_gauge_context),
                 tokio::time::sleep(GET_INSTANCE_TIMEOUT).map(|()| self.spinup_new_instance()),
             )
             .await
@@ -476,11 +500,7 @@ pub enum InitDatabaseError {
 impl ModuleHost {
     pub fn new(threadpool: Arc<HostThreadpool>, mut module: impl Module) -> Self {
         let info = module.info();
-        let waiter_gauge =
-            WORKER_METRICS
-                .instance_queue_length
-                .with_label_values(&info.identity, &info.module_hash, &info.address);
-        let instance_pool = LendingPool::new(waiter_gauge);
+        let instance_pool = LendingPool::with_gauge(InstancePoolGauge);
         instance_pool.add_multiple(module.initial_instances()).unwrap();
         let inner = Arc::new(HostControllerActor {
             module: Arc::new(module),
@@ -505,12 +525,18 @@ impl ModuleHost {
         &self.info.subscription
     }
 
-    async fn call<F, R>(&self, f: F) -> Result<R, NoSuchModule>
+    async fn call<F, R>(&self, reducer_name: &str, f: F) -> Result<R, NoSuchModule>
     where
         F: FnOnce(&mut dyn ModuleInstance) -> R + Send + 'static,
         R: Send + 'static,
     {
-        let (threadpool, mut inst) = self.inner.get_instance().await?;
+        let waiter_gauge_context = (
+            &self.info.identity,
+            &self.info.module_hash,
+            &self.info.address,
+            reducer_name,
+        );
+        let (threadpool, mut inst) = self.inner.get_instance(&waiter_gauge_context).await?;
 
         let (tx, rx) = oneshot::channel();
         threadpool.spawn(move || {
@@ -561,7 +587,7 @@ impl ModuleHost {
         let args = args.into_tuple(self.info.typespace.with_type(schema))?;
         let caller_address = caller_address.unwrap_or(Address::__DUMMY);
 
-        self.call(move |inst| {
+        self.call(reducer_name, move |inst| {
             inst.call_reducer(CallReducerParams {
                 timestamp: Timestamp::now(),
                 caller_identity,
@@ -619,13 +645,13 @@ impl ModuleHost {
             Some(schema) => args.into_tuple(schema)?,
             _ => ArgsTuple::default(),
         };
-        self.call(move |inst| inst.init_database(fence, args))
+        self.call("<init_database>", move |inst| inst.init_database(fence, args))
             .await?
             .map_err(InitDatabaseError::Other)
     }
 
     pub async fn update_database(&self, fence: u128) -> Result<UpdateDatabaseResult, anyhow::Error> {
-        self.call(move |inst| inst.update_database(fence))
+        self.call("<update_database>", move |inst| inst.update_database(fence))
             .await?
             .map_err(Into::into)
     }

crates/core/src/util/lending_pool.rs

@@ -15,24 +15,32 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use crate::util::prometheus_handle::IntGaugeExt;
 
 use super::notify_once::{NotifiedOnce, NotifyOnce};
+use super::prometheus_handle::GaugeInc;
 
-pub struct LendingPool<T> {
+pub struct LendingPool<T, G = ()> {
     sem: Arc<Semaphore>,
+    waiter_gauge: G,
     inner: Arc<LendingPoolInner<T>>,
 }
 
-impl<T> Clone for LendingPool<T> {
+impl<T, G: WaiterGauge + Default> Default for LendingPool<T, G> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T, G: Clone> Clone for LendingPool<T, G> {
     fn clone(&self) -> Self {
         Self {
             sem: self.sem.clone(),
+            waiter_gauge: self.waiter_gauge.clone(),
             inner: self.inner.clone(),
         }
     }
 }
 
 struct LendingPoolInner<T> {
     closed_notify: NotifyOnce,
-    waiter_gauge: IntGauge,
     vec: Mutex<PoolVec<T>>,
 }
 
@@ -44,18 +52,25 @@ struct PoolVec<T> {
 #[derive(Debug)]
 pub struct PoolClosed;
 
-impl<T> LendingPool<T> {
-    pub fn new(waiter_gauge: IntGauge) -> Self {
-        Self::from_iter(std::iter::empty(), waiter_gauge)
+impl<T, G: WaiterGauge> LendingPool<T, G> {
+    pub fn new() -> Self
+    where
+        G: Default,
+    {
+        Self::with_gauge(G::default())
+    }
+
+    pub fn with_gauge(waiter_gauge: G) -> Self {
+        Self::from_iter_with_gauge(std::iter::empty(), waiter_gauge)
     }
 
-    pub fn from_iter<I: IntoIterator<Item = T>>(iter: I, waiter_gauge: IntGauge) -> Self {
+    pub fn from_iter_with_gauge<I: IntoIterator<Item = T>>(iter: I, waiter_gauge: G) -> Self {
         let deque = VecDeque::from_iter(iter);
         Self {
             sem: Arc::new(Semaphore::new(deque.len())),
+            waiter_gauge,
             inner: Arc::new(LendingPoolInner {
                 closed_notify: NotifyOnce::new(),
-                waiter_gauge,
                 vec: Mutex::new(PoolVec {
                     total_count: deque.len(),
                     deque: Some(deque),
@@ -64,10 +79,20 @@ impl<T> LendingPool<T> {
         }
     }
 
-    pub fn request(&self) -> impl Future<Output = Result<LentResource<T>, PoolClosed>> {
+    pub fn request(&self) -> impl Future<Output = Result<LentResource<T>, PoolClosed>>
+    where
+        G: for<'a> WaiterGauge<Context<'a> = ()>,
+    {
+        self.request_with_context(())
+    }
+
+    pub fn request_with_context(
+        &self,
+        context: G::Context<'_>,
+    ) -> impl Future<Output = Result<LentResource<T>, PoolClosed>> {
         let acq = self.sem.clone().acquire_owned();
         let pool_inner = self.inner.clone();
-        let waiter_guard = pool_inner.waiter_gauge.inc_scope();
+        let waiter_guard = self.waiter_gauge.inc(context);
         async move {
             let permit = acq.await.map_err(|_| PoolClosed)?;
             drop(waiter_guard);
@@ -129,6 +154,12 @@ impl<T> LendingPool<T> {
     }
 }
 
+impl<T, G: WaiterGauge + Default> FromIterator<T> for LendingPool<T, G> {
+    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
+        Self::from_iter_with_gauge(iter, G::default())
+    }
+}
+
 pin_project_lite::pin_project! {
     pub struct Closed<'a> {
         #[pin]
@@ -194,3 +225,31 @@ impl<T> Drop for LentResource<T> {
         }
     }
 }
+
+pub trait WaiterGauge {
+    type Context<'a>;
+    type IncGuard;
+    fn inc(&self, context: Self::Context<'_>) -> Self::IncGuard;
+}
+
+impl WaiterGauge for () {
+    type Context<'a> = ();
+    type IncGuard = ();
+    fn inc(&self, (): ()) -> Self::IncGuard {}
+}
+
+impl<G: WaiterGauge> WaiterGauge for Arc<G> {
+    type Context<'a> = G::Context<'a>;
+    type IncGuard = G::IncGuard;
+    fn inc(&self, context: Self::Context<'_>) -> Self::IncGuard {
+        (**self).inc(context)
+    }
+}
+
+impl WaiterGauge for IntGauge {
+    type Context<'a> = ();
+    type IncGuard = GaugeInc;
+    fn inc(&self, (): ()) -> Self::IncGuard {
+        self.inc_scope()
+    }
+}

crates/core/src/worker_metrics/mod.rs

@@ -66,7 +66,7 @@ metrics_group!(
 
         #[name = spacetime_worker_instance_operation_queue_length]
         #[help = "Length of the wait queue for access to a module instance."]
-        #[labels(identity: Identity, module_hash: Hash, database_address: Address)]
+        #[labels(identity: Identity, module_hash: Hash, database_address: Address, reducer_symbol: str)]
         pub instance_queue_length: IntGaugeVec,
 
         #[name = spacetime_system_disk_space_total_bytes]