Add meters to custom slot supplier contexts (#913)

Author: Sushisource

core-api/src/telemetry/metrics.rs

@@ -2,6 +2,7 @@ use std::{
     any::Any,
     borrow::Cow,
     fmt::Debug,
+    ops::Deref,
     sync::{Arc, OnceLock},
     time::Duration,
 };
@@ -63,6 +64,13 @@ pub struct TemporalMeter {
     pub default_attribs: NewAttributes,
 }
 
+impl Deref for TemporalMeter {
+    type Target = dyn CoreMeter;
+    fn deref(&self) -> &Self::Target {
+        self.inner.as_ref()
+    }
+}
+
 impl CoreMeter for Arc<dyn CoreMeter> {
     fn new_attributes(&self, attribs: NewAttributes) -> MetricAttributes {
         self.as_ref().new_attributes(attribs)

core-api/src/worker.rs

@@ -282,10 +282,6 @@ pub trait WorkerTuner {
     fn nexus_task_slot_supplier(
         &self,
     ) -> Arc<dyn SlotSupplier<SlotKind = NexusSlotKind> + Send + Sync>;
-
-    /// Core will call this at worker initialization time, allowing the implementation to hook up to
-    /// metrics if any are configured. If not, it will not be called.
-    fn attach_metrics(&self, metrics: TemporalMeter);
 }
 
 /// Implementing this trait allows users to customize how many tasks of certain kinds the worker
@@ -339,6 +335,11 @@ pub trait SlotReservationContext: Send + Sync {
 
     /// Returns true iff this is a sticky poll for a workflow task
     fn is_sticky(&self) -> bool;
+
+    /// Returns the metrics meter if metrics are enabled
+    fn get_metrics_meter(&self) -> Option<TemporalMeter> {
+        None
+    }
 }
 
 pub trait SlotMarkUsedContext: Send + Sync {
@@ -347,6 +348,11 @@ pub trait SlotMarkUsedContext: Send + Sync {
     fn permit(&self) -> &SlotSupplierPermit;
     /// Returns the info of slot that was marked as used
     fn info(&self) -> &<Self::SlotKind as SlotKind>::Info;
+
+    /// Returns the metrics meter if metrics are enabled
+    fn get_metrics_meter(&self) -> Option<TemporalMeter> {
+        None
+    }
 }
 
 pub trait SlotReleaseContext: Send + Sync {
@@ -355,6 +361,11 @@ pub trait SlotReleaseContext: Send + Sync {
     fn permit(&self) -> &SlotSupplierPermit;
     /// Returns the info of slot that was released, if it was used
     fn info(&self) -> Option<&<Self::SlotKind as SlotKind>::Info>;
+
+    /// Returns the metrics meter if metrics are enabled
+    fn get_metrics_meter(&self) -> Option<TemporalMeter> {
+        None
+    }
 }
 
 #[derive(Default, Debug)]

core/src/abstractions.rs

@@ -10,9 +10,12 @@ use std::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
     },
 };
-use temporal_sdk_core_api::worker::{
-    SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
-    SlotSupplierPermit, WorkerDeploymentVersion, WorkflowSlotKind,
+use temporal_sdk_core_api::{
+    telemetry::metrics::TemporalMeter,
+    worker::{
+        SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
+        SlotSupplierPermit, WorkerDeploymentVersion, WorkflowSlotKind,
+    },
 };
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
@@ -36,6 +39,7 @@ pub(crate) struct MeteredPermitDealer<SK: SlotKind> {
     /// there will need to be some associated refactoring.
     max_permits: Option<usize>,
     metrics_ctx: MetricsContext,
+    meter: Option<TemporalMeter>,
     /// Only applies to permit dealers for workflow tasks. True if this permit dealer is associated
     /// with a sticky queue poller.
     is_sticky_poller: bool,
@@ -59,12 +63,14 @@ where
         metrics_ctx: MetricsContext,
         max_permits: Option<usize>,
         context_data: Arc<PermitDealerContextData>,
+        meter: Option<TemporalMeter>,
     ) -> Self {
         Self {
             supplier,
             unused_claimants: Arc::new(AtomicUsize::new(0)),
             extant_permits: watch::channel(0),
             metrics_ctx,
+            meter,
             max_permits,
             is_sticky_poller: false,
             context_data,
@@ -141,6 +147,7 @@ where
             release_ctx: ReleaseCtx {
                 permit: res,
                 stored_info: None,
+                meter: self.meter.clone(),
             },
             use_fn: Box::new(move |info| {
                 supp_c.mark_slot_used(info);
@@ -182,11 +189,16 @@ impl<SK: SlotKind> SlotReservationContext for MeteredPermitDealer<SK> {
     fn is_sticky(&self) -> bool {
         self.is_sticky_poller
     }
+
+    fn get_metrics_meter(&self) -> Option<TemporalMeter> {
+        self.meter.clone()
+    }
 }
 
 struct UseCtx<'a, SK: SlotKind> {
     stored_info: &'a SK::Info,
     permit: &'a SlotSupplierPermit,
+    meter: Option<TemporalMeter>,
 }
 
 impl<SK: SlotKind> SlotMarkUsedContext for UseCtx<'_, SK> {
@@ -199,11 +211,16 @@ impl<SK: SlotKind> SlotMarkUsedContext for UseCtx<'_, SK> {
     fn info(&self) -> &<Self::SlotKind as SlotKind>::Info {
         self.stored_info
     }
+
+    fn get_metrics_meter(&self) -> Option<TemporalMeter> {
+        self.meter.clone()
+    }
 }
 
 struct ReleaseCtx<SK: SlotKind> {
     permit: SlotSupplierPermit,
     stored_info: Option<SK::Info>,
+    meter: Option<TemporalMeter>,
 }
 
 impl<SK: SlotKind> SlotReleaseContext for ReleaseCtx<SK> {
@@ -216,6 +233,10 @@ impl<SK: SlotKind> SlotReleaseContext for ReleaseCtx<SK> {
     fn info(&self) -> Option<&<Self::SlotKind as SlotKind>::Info> {
         self.stored_info.as_ref()
     }
+
+    fn get_metrics_meter(&self) -> Option<TemporalMeter> {
+        self.meter.clone()
+    }
 }
 
 /// A version of [MeteredPermitDealer] that can be closed and supports waiting for close to complete.
@@ -353,6 +374,7 @@ impl<SK: SlotKind> OwnedMeteredSemPermit<SK> {
         let ctx = UseCtx {
             stored_info: &info,
             permit: &self.release_ctx.permit,
+            meter: self.release_ctx.meter.clone(),
         };
         (self.use_fn)(&ctx);
         self.release_ctx.stored_info = Some(info);
@@ -386,6 +408,7 @@ pub(crate) mod tests {
             MetricsContext::no_op(),
             None,
             Arc::new(Default::default()),
+            None,
         )
     }
 

core/src/worker/activities/local_activities.rs

@@ -255,6 +255,7 @@ impl LocalActivityManager {
                 MetricsContext::no_op(),
                 None,
                 Arc::new(Default::default()),
+                None,
             ),
             hb_tx,
             MetricsContext::no_op(),

core/src/worker/mod.rs

@@ -318,9 +318,6 @@ impl Worker {
             .unwrap_or_else(|| Arc::new(TunerBuilder::from_config(&config).build()));
 
         metrics.worker_registered();
-        if let Some(meter) = meter {
-            tuner.attach_metrics(meter.clone());
-        }
         let shutdown_token = CancellationToken::new();
         let slot_context_data = Arc::new(PermitDealerContextData {
             task_queue: config.task_queue.clone(),
@@ -338,13 +335,15 @@ impl Worker {
                 None
             },
             slot_context_data.clone(),
+            meter.clone(),
         );
         let wft_permits = wft_slots.get_extant_count_rcv();
         let act_slots = MeteredPermitDealer::new(
             tuner.activity_task_slot_supplier(),
             metrics.with_new_attrs([activity_worker_type()]),
             None,
             slot_context_data.clone(),
+            meter.clone(),
         );
         let act_permits = act_slots.get_extant_count_rcv();
         let (external_wft_tx, external_wft_rx) = unbounded_channel();
@@ -353,6 +352,7 @@ impl Worker {
             metrics.with_new_attrs([nexus_worker_type()]),
             None,
             slot_context_data.clone(),
+            meter.clone(),
         );
         let (wft_stream, act_poller, nexus_poller) = match task_pollers {
             TaskPollers::Real => {
@@ -438,6 +438,7 @@ impl Worker {
             metrics.with_new_attrs([local_activity_worker_type()]),
             None,
             slot_context_data,
+            meter.clone(),
         );
         let la_permits = la_pemit_dealer.get_extant_count_rcv();
         let local_act_mgr = Arc::new(LocalActivityManager::new(

core/src/worker/tuner.rs

@@ -7,13 +7,10 @@ pub use resource_based::{
     ResourceSlotOptions,
 };
 
-use std::sync::{Arc, OnceLock};
-use temporal_sdk_core_api::{
-    telemetry::metrics::TemporalMeter,
-    worker::{
-        ActivitySlotKind, LocalActivitySlotKind, NexusSlotKind, SlotKind, SlotSupplier,
-        WorkerConfig, WorkerTuner, WorkflowSlotKind,
-    },
+use std::sync::Arc;
+use temporal_sdk_core_api::worker::{
+    ActivitySlotKind, LocalActivitySlotKind, NexusSlotKind, SlotKind, SlotSupplier, WorkerConfig,
+    WorkerTuner, WorkflowSlotKind,
 };
 
 /// Allows for the composition of different slot suppliers into a [WorkerTuner]
@@ -22,7 +19,6 @@ pub struct TunerHolder {
     act_supplier: Arc<dyn SlotSupplier<SlotKind = ActivitySlotKind> + Send + Sync>,
     la_supplier: Arc<dyn SlotSupplier<SlotKind = LocalActivitySlotKind> + Send + Sync>,
     nexus_supplier: Arc<dyn SlotSupplier<SlotKind = NexusSlotKind> + Send + Sync>,
-    metrics: OnceLock<TemporalMeter>,
 }
 
 /// Can be used to construct a [TunerHolder] without needing to manually construct each
@@ -245,7 +241,6 @@ impl TunerBuilder {
                 .nexus_slot_supplier
                 .clone()
                 .unwrap_or_else(|| Arc::new(FixedSizeSlotSupplier::new(100))),
-            metrics: OnceLock::new(),
         }
     }
 }
@@ -274,8 +269,4 @@ impl WorkerTuner for TunerHolder {
     ) -> Arc<dyn SlotSupplier<SlotKind = NexusSlotKind> + Send + Sync> {
         self.nexus_supplier.clone()
     }
-
-    fn attach_metrics(&self, m: TemporalMeter) {
-        let _ = self.metrics.set(m);
-    }
 }

core/src/worker/tuner/resource_based.rs

@@ -262,6 +262,9 @@ where
     type SlotKind = SK;
 
     async fn reserve_slot(&self, ctx: &dyn SlotReservationContext) -> SlotSupplierPermit {
+        if let Some(m) = ctx.get_metrics_meter() {
+            self.inner.attach_metrics(m);
+        }
         loop {
             if let Some(value) = self.issue_if_below_minimums(ctx) {
                 return value;
@@ -282,6 +285,9 @@ where
     }
 
     fn try_reserve_slot(&self, ctx: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
+        if let Some(m) = ctx.get_metrics_meter() {
+            self.inner.attach_metrics(m);
+        }
         if let v @ Some(_) = self.issue_if_below_minimums(ctx) {
             return v;
         }
@@ -396,10 +402,6 @@ impl<MI: SystemResourceInfo + Sync + Send + 'static> WorkerTuner for ResourceBas
         let o = self.nexus_opts.unwrap_or(DEFAULT_NEXUS_SLOT_OPTS);
         self.slots.as_kind(o)
     }
-
-    fn attach_metrics(&self, metrics: TemporalMeter) {
-        self.slots.attach_metrics(metrics);
-    }
 }
 
 impl<MI: SystemResourceInfo + Sync + Send> ResourceController<MI> {
@@ -596,6 +598,7 @@ mod tests {
             MetricsContext::no_op(),
             None,
             Arc::new(Default::default()),
+            None,
         );
         let pd_s = pd.clone().into_sticky();
         // Start with too high usage
@@ -625,6 +628,7 @@ mod tests {
             MetricsContext::no_op(),
             None,
             Arc::new(Default::default()),
+            None,
         );
         // Start with too high usage
         used.store(90_000, Ordering::Release);
@@ -649,6 +653,7 @@ mod tests {
             MetricsContext::no_op(),
             None,
             Arc::new(Default::default()),
+            None,
         );
         let pd_s = pd.clone().into_sticky();
         let order = crossbeam_queue::ArrayQueue::new(2);
@@ -684,6 +689,7 @@ mod tests {
             MetricsContext::no_op(),
             None,
             Arc::new(Default::default()),
+            None,
         );
         let order = crossbeam_queue::ArrayQueue::new(2);
         let waits_free = async {
@@ -713,6 +719,7 @@ mod tests {
             MetricsContext::no_op(),
             None,
             Arc::new(Default::default()),
+            None,
         );
         used.store(90_000, Ordering::Release);
         let _p1 = pd.try_acquire_owned().unwrap();