Flush metrics on span exit

This is wasteful, but is required for tests only
akoshelev · Sep 17, 2024 · a54e0dd · a54e0dd
1 parent 3ff5ed5
commit a54e0dd
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/ipa-metrics-tracing/src/layer.rs b/ipa-metrics-tracing/src/layer.rs
@@ -1,5 +1,6 @@
 use std::fmt::Debug;
 
+use ipa_metrics::MetricsCurrentThreadContext;
 use tracing::{
     field::{Field, Visit},
     instrument::WithSubscriber,
@@ -59,6 +60,13 @@ impl<S: Subscriber + for<'s> LookupSpan<'s>> Layer<S> for MetricsPartitioningLay
             ipa_metrics::set_or_unset_partition(prev.take())
         };
     }
+
+    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
+        let span = ctx.span(&id).expect("Span should exists before closing it");
+        if let Some(MetricPartition(_, _)) = span.extensions().get() {
+            MetricsCurrentThreadContext::flush();
+        };
+    }
 }
 
 #[cfg(test)]