From b4549edd6eb436ce8d59acdb798c9fa79b6c8392 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Thu, 6 Jan 2022 12:21:21 -0800
Subject: [PATCH] put back poll duration histograms

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 console-subscriber/src/aggregator/mod.rs | 15 +--------------
 console-subscriber/src/stats.rs          | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/console-subscriber/src/aggregator/mod.rs b/console-subscriber/src/aggregator/mod.rs
index be88bf4d6..6f21e4d54 100644
--- a/console-subscriber/src/aggregator/mod.rs
+++ b/console-subscriber/src/aggregator/mod.rs
@@ -21,11 +21,6 @@ use std::{
 };
 use tracing_core::{span::Id, Metadata};
 
-use hdrhistogram::{
-    serialization::{Serializer, V2SerializeError, V2Serializer},
-    Histogram,
-};
-
 mod id_data;
 mod shrink;
 use self::id_data::{IdData, Include};
@@ -329,8 +324,7 @@ impl Aggregator {
                 && subscription.update(&proto::tasks::TaskDetails {
                     task_id: Some(id.clone().into()),
                     now: Some(now.into()),
-                    poll_times_histogram: None, // TODO(eliza): put back
-                                                // poll_times_histogram: serialize_histogram(&stats.poll_times_histogram).ok(),
+                    poll_times_histogram: stats.serialize_histogram(),
                 })
             {
                 self.details_watchers
@@ -720,13 +714,6 @@ impl From<AttributeUpdate> for Attribute {
     }
 }
 
-fn serialize_histogram(histogram: &Histogram<u64>) -> Result<Vec<u8>, V2SerializeError> {
-    let mut serializer = V2Serializer::new();
-    let mut buf = Vec::new();
-    serializer.serialize(histogram, &mut buf)?;
-    Ok(buf)
-}
-
 fn update_attribute(attribute: &mut Attribute, update: &AttributeUpdate) {
     use proto::field::Value::*;
     let attribute_val = attribute.field.as_mut().and_then(|a| a.value.as_mut());
diff --git a/console-subscriber/src/stats.rs b/console-subscriber/src/stats.rs
index 160f9dc43..714ce19e1 100644
--- a/console-subscriber/src/stats.rs
+++ b/console-subscriber/src/stats.rs
@@ -1,5 +1,12 @@
-use crate::{attribute, sync::Mutex, ToProto};
-use hdrhistogram::Histogram;
+use crate::{
+    attribute,
+    sync::{Mutex, MutexGuard},
+    ToProto,
+};
+use hdrhistogram::{
+    serialization::{Serializer, V2Serializer},
+    Histogram,
+};
 use std::cmp;
 use std::collections::HashMap;
 use std::sync::{
@@ -218,6 +225,15 @@ impl TaskStats {
     fn make_dirty(&self) {
         self.is_dirty.swap(true, AcqRel);
     }
+
+    pub(crate) fn serialize_histogram(&self) -> Option<Vec<u8>> {
+        let poll_timestamps = self.poll_stats.timestamps.lock();
+        let histogram = poll_timestamps.histogram.as_ref()?;
+        let mut serializer = V2Serializer::new();
+        let mut buf = Vec::new();
+        serializer.serialize(histogram, &mut buf).ok()?;
+        Some(buf)
+    }
 }
 
 impl ToProto for TaskStats {