diff --git a/crates/metrics/Cargo.toml b/crates/metrics/Cargo.toml index 6710565..6c588d7 100644 --- a/crates/metrics/Cargo.toml +++ b/crates/metrics/Cargo.toml @@ -9,3 +9,4 @@ prometheus = "0.13" opentelemetry = { version = "0.19", features = ["metrics", "rt-tokio"] } opentelemetry-prometheus = "0.12" once_cell = "1.17" +smallvec = "1.11" diff --git a/crates/metrics/src/task.rs b/crates/metrics/src/task.rs index a364d14..4f99ec8 100644 --- a/crates/metrics/src/task.rs +++ b/crates/metrics/src/task.rs @@ -7,6 +7,7 @@ use { }, once_cell::sync::OnceCell, opentelemetry::metrics::{Counter, Histogram}, + smallvec::SmallVec, std::{ops::Deref, sync::Arc, time::Duration}, }; @@ -14,7 +15,7 @@ use { /// initialized. pub struct TaskMetrics { prefix: &'static str, - inner: OnceCell>, + inner: OnceCell, } impl TaskMetrics { @@ -50,60 +51,57 @@ impl Deref for TaskMetrics { /// The above metrics are tracked using [`opentelemetry`] metrics API and are /// prefixed according to the constructor arguments. #[derive(Clone)] -pub struct OtelTaskMetricsRecorder { +pub struct OtelTaskMetricsRecorder { inner: Arc, - name: Option, + name: &'static str, + attributes: SmallVec<[otel::KeyValue; 2]>, } -impl OtelTaskMetricsRecorder<()> { +impl OtelTaskMetricsRecorder { pub fn new(prefix: &str) -> Self { Self { inner: Arc::new(OtelRecorderInner::new(prefix)), - name: None, + name: "unknown", + attributes: SmallVec::new(), } } -} - -impl OtelTaskMetricsRecorder -where - N: AsTaskName, -{ - #[inline] - fn task_name_kv(&self) -> otel::KeyValue { - let name: &'static str = self - .name - .as_ref() - .map(AsTaskName::as_task_name) - .unwrap_or_default(); - - otel::KeyValue::new("task_name", name) - } -} -impl OtelTaskMetricsRecorder -where - N1: AsTaskName, -{ /// Clones the current recording context with a new task name. - pub fn with_name(&self, name: N2) -> OtelTaskMetricsRecorder + pub fn with_name(&self, name: N) -> Self where - N2: AsTaskName, + N: AsTaskName, { - OtelTaskMetricsRecorder { + Self { inner: self.inner.clone(), - name: Some(name), + name: name.as_task_name(), + attributes: self.attributes.clone(), } } + + /// Clones the current recording context with a new set of attributes. + pub fn with_attributes( + &self, + attributes: impl IntoIterator, + ) -> OtelTaskMetricsRecorder { + Self { + inner: self.inner.clone(), + name: self.name, + attributes: attributes.into_iter().collect(), + } + } + + fn combine_attributes(&self) -> SmallVec<[otel::KeyValue; 4]> { + let name = [otel::KeyValue::new("task_name", self.name)]; + let extra = self.attributes.iter().cloned(); + name.into_iter().chain(extra).collect() + } } -impl TaskMetricsRecorder for OtelTaskMetricsRecorder -where - N: AsTaskName, -{ +impl TaskMetricsRecorder for OtelTaskMetricsRecorder { fn record_task_started(&self) { self.inner .tasks_started - .add(&otel::Context::new(), 1, &[self.task_name_kv()]); + .add(&otel::Context::new(), 1, &self.combine_attributes()); } fn record_task_finished( @@ -115,10 +113,10 @@ where ) { let total_duration_ms = duration_as_millis_f64(total_duration); let poll_duration_ms = duration_as_millis_f64(poll_duration); - let attrs = [ - self.task_name_kv(), - otel::KeyValue::new("completed", completed), - ]; + + let mut attrs = self.combine_attributes(); + attrs.push(otel::KeyValue::new("completed", completed)); + let ctx = otel::Context::new(); self.inner