[tracing] Fix flow id plumbing through histogram

Following up on
https://chromium-review.googlesource.com/c/chromium/src/+/6233616

I realized that the callback is invoked as a task, so the thread local
event id  isn't plumbed correctly.
The solution I'm proposing here is to grab the event id directly in histogram impl (FindAndRunCallbacks), and forward it to callbacks, without changing the histogram API.
Drive-by: IWYU fixes.

Bug: 40257548
Change-Id: Ibf5ac1c4e96f0eff63b108d1acf3e550579f6ee5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6302393
Reviewed-by: Luc Nguyen <lucnguyen@google.com>
Reviewed-by: Etienne Bergeron <etienneb@chromium.org>
Commit-Queue: Etienne Pierre-Doray <etiennep@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1426581}

base/metrics/histogram_base.cc

@@ -26,6 +26,7 @@
 #include "base/rand_util.h"
 #include "base/strings/stringprintf.h"
 #include "base/synchronization/lock.h"
+#include "base/trace_event/histogram_scope.h"  // no-presubmit-check
 #include "base/values.h"
 
 namespace base {
@@ -187,10 +188,11 @@ void HistogramBase::WriteJSON(std::string* output,
 }
 
 void HistogramBase::FindAndRunCallbacks(HistogramBase::Sample32 sample) const {
+  auto event_id = trace_event::HistogramScope::GetFlowId();
   StatisticsRecorder::GlobalSampleCallback global_sample_callback =
       StatisticsRecorder::global_sample_callback();
   if (global_sample_callback) {
-    global_sample_callback(histogram_name(), name_hash(), sample);
+    global_sample_callback(histogram_name(), name_hash(), sample, event_id);
   }
 
   // We check the flag first since it is very cheap and we can avoid the
@@ -200,7 +202,8 @@ void HistogramBase::FindAndRunCallbacks(HistogramBase::Sample32 sample) const {
   }
 
   StatisticsRecorder::FindAndRunHistogramCallbacks(
-      base::PassKey<HistogramBase>(), histogram_name(), name_hash(), sample);
+      base::PassKey<HistogramBase>(), histogram_name(), name_hash(), sample,
+      event_id);
 }
 
 HistogramBase::CountAndBucketData HistogramBase::GetCountAndBucketData() const {

base/metrics/statistics_recorder.cc

@@ -58,7 +58,16 @@ std::atomic<StatisticsRecorder::GlobalSampleCallback>
 StatisticsRecorder::ScopedHistogramSampleObserver::
     ScopedHistogramSampleObserver(std::string_view name,
                                   OnSampleCallback callback)
-    : histogram_name_(name), callback_(callback) {
+    : histogram_name_(name),
+      callback_(
+          base::IgnoreArgs<std::optional<uint64_t>>(std::move(callback))) {
+  StatisticsRecorder::AddHistogramSampleObserver(histogram_name_, this);
+}
+
+StatisticsRecorder::ScopedHistogramSampleObserver::
+    ScopedHistogramSampleObserver(std::string_view name,
+                                  OnSampleWithEventCallback callback)
+    : histogram_name_(name), callback_(std::move(callback)) {
   StatisticsRecorder::AddHistogramSampleObserver(histogram_name_, this);
 }
 
@@ -70,8 +79,9 @@ StatisticsRecorder::ScopedHistogramSampleObserver::
 void StatisticsRecorder::ScopedHistogramSampleObserver::RunCallback(
     std::string_view histogram_name,
     uint64_t name_hash,
-    HistogramBase::Sample32 sample) {
-  callback_.Run(histogram_name, name_hash, sample);
+    HistogramBase::Sample32 sample,
+    std::optional<uint64_t> event_id) {
+  callback_.Run(event_id, histogram_name, name_hash, sample);
 }
 
 StatisticsRecorder::~StatisticsRecorder() {
@@ -386,7 +396,8 @@ void StatisticsRecorder::FindAndRunHistogramCallbacks(
     base::PassKey<HistogramBase>,
     std::string_view histogram_name,
     uint64_t name_hash,
-    HistogramBase::Sample32 sample) {
+    HistogramBase::Sample32 sample,
+    std::optional<uint64_t> event_id) {
   DCHECK_EQ(name_hash, HashMetricName(histogram_name));
 
   const AutoLock auto_lock(GetLock());
@@ -406,7 +417,7 @@ void StatisticsRecorder::FindAndRunHistogramCallbacks(
   }
 
   it->second->Notify(FROM_HERE, &ScopedHistogramSampleObserver::RunCallback,
-                     histogram_name, name_hash, sample);
+                     histogram_name, name_hash, sample, event_id);
 }
 
 // static

base/metrics/statistics_recorder.h

@@ -77,6 +77,11 @@ class BASE_EXPORT StatisticsRecorder {
       base::RepeatingCallback<void(std::string_view /*=histogram_name*/,
                                    uint64_t /*=name_hash*/,
                                    HistogramBase::Sample32)>;
+  using OnSampleWithEventCallback =
+      base::RepeatingCallback<void(std::optional<uint64_t> event_id,
+                                   std::string_view /*=histogram_name*/,
+                                   uint64_t /*=name_hash*/,
+                                   HistogramBase::Sample32)>;
 
   // An observer that gets notified whenever a new sample is recorded for a
   // particular histogram. Clients only need to construct it with the histogram
@@ -91,6 +96,8 @@ class BASE_EXPORT StatisticsRecorder {
     // be invoked when a sample is recorded.
     explicit ScopedHistogramSampleObserver(std::string_view histogram_name,
                                            OnSampleCallback callback);
+    explicit ScopedHistogramSampleObserver(std::string_view histogram_name,
+                                           OnSampleWithEventCallback callback);
     ~ScopedHistogramSampleObserver();
 
    private:
@@ -99,14 +106,15 @@ class BASE_EXPORT StatisticsRecorder {
     // Runs the callback.
     void RunCallback(std::string_view histogram_name,
                      uint64_t name_hash,
-                     HistogramBase::Sample32 sample);
+                     HistogramBase::Sample32 sample,
+                     std::optional<uint64_t> event_id);
 
     // The name of the histogram to observe.
     const std::string histogram_name_;
 
     // The client supplied callback that is invoked when the histogram sample is
     // collected.
-    const OnSampleCallback callback_;
+    const OnSampleWithEventCallback callback_;
   };
 
   typedef std::vector<HistogramBase*> Histograms;
@@ -218,7 +226,8 @@ class BASE_EXPORT StatisticsRecorder {
   static void FindAndRunHistogramCallbacks(base::PassKey<HistogramBase>,
                                            std::string_view histogram_name,
                                            uint64_t name_hash,
-                                           HistogramBase::Sample32 sample);
+                                           HistogramBase::Sample32 sample,
+                                           std::optional<uint64_t> event_id);
 
   // Returns the number of known histograms.
   //
@@ -280,7 +289,8 @@ class BASE_EXPORT StatisticsRecorder {
 
   using GlobalSampleCallback = void (*)(std::string_view /*=histogram_name*/,
                                         uint64_t /*=name_hash*/,
-                                        HistogramBase::Sample32);
+                                        HistogramBase::Sample32,
+                                        std::optional<uint64_t> /*=event_id*/);
   // Installs a global callback which will be called for every added
   // histogram sample. The given callback is a raw function pointer in order
   // to be accessed lock-free and can be called on any thread.

base/metrics/statistics_recorder_unittest.cc

@@ -22,6 +22,7 @@
 #include "base/metrics/record_histogram_checker.h"
 #include "base/metrics/sparse_histogram.h"
 #include "base/test/task_environment.h"
+#include "base/trace_event/histogram_scope.h"  // no-presubmit-check
 #include "base/values.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -477,7 +478,16 @@ struct CallbackCheckWrapper {
     last_histogram_value = histogram_value;
   }
 
+  void OnHistogramChangedWithEventId(std::optional<uint64_t> event_id,
+                                     std::string_view histogram_name,
+                                     uint64_t name_hash,
+                                     HistogramBase::Sample32 histogram_value) {
+    last_event_id = event_id;
+    OnHistogramChanged(histogram_name, name_hash, histogram_value);
+  }
+
   bool called = false;
+  std::optional<uint64_t> last_event_id;
   std::string last_histogram_name = "";
   uint64_t last_name_hash;
   base::HistogramBase::Sample32 last_histogram_value = 0;
@@ -730,6 +740,44 @@ TEST_P(StatisticsRecorderTest, CallbackUsedBeforeHistogramCreatedTest) {
   EXPECT_EQ(callback_wrapper.last_histogram_value, 1);
 }
 
+// Check that setting a callback before the histogram exists works.
+TEST_P(StatisticsRecorderTest, CallbackWithEventIdTest) {
+  test::TaskEnvironment task_environment;
+  CallbackCheckWrapper callback_wrapper;
+
+  auto callback =
+      std::make_unique<base::StatisticsRecorder::ScopedHistogramSampleObserver>(
+          "TestHistogram",
+          base::BindRepeating(
+              &CallbackCheckWrapper::OnHistogramChangedWithEventId,
+              base::Unretained(&callback_wrapper)));
+
+  HistogramBase* histogram = Histogram::FactoryGet("TestHistogram", 1, 1000, 10,
+                                                   HistogramBase::kNoFlags);
+  EXPECT_TRUE(histogram);
+  {
+    base::trace_event::HistogramScope histogram_trace_scope(42);
+    histogram->Add(1);
+    base::RunLoop().RunUntilIdle();
+
+    EXPECT_TRUE(callback_wrapper.called);
+    EXPECT_EQ(callback_wrapper.last_histogram_name, "TestHistogram");
+    EXPECT_EQ(callback_wrapper.last_event_id, 42);
+    EXPECT_EQ(callback_wrapper.last_name_hash, HashMetricName("TestHistogram"));
+    EXPECT_EQ(callback_wrapper.last_histogram_value, 1);
+  }
+
+  callback_wrapper = CallbackCheckWrapper();  // clear previous callback data
+  histogram->Add(2);
+  base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(callback_wrapper.called);
+  EXPECT_EQ(callback_wrapper.last_histogram_name, "TestHistogram");
+  EXPECT_EQ(callback_wrapper.last_event_id, std::nullopt);
+  EXPECT_EQ(callback_wrapper.last_name_hash, HashMetricName("TestHistogram"));
+  EXPECT_EQ(callback_wrapper.last_histogram_value, 2);
+}
+
 TEST_P(StatisticsRecorderTest, GlobalCallbackCalled) {
   HistogramBase* histogram = Histogram::FactoryGet("TestHistogram", 1, 1000, 10,
                                                    HistogramBase::kNoFlags);
@@ -741,7 +789,8 @@ TEST_P(StatisticsRecorderTest, GlobalCallbackCalled) {
   static size_t callback_callcount;
   callback_callcount = 0;
   auto callback = [](std::string_view histogram_name, uint64_t name_hash,
-                     HistogramBase::Sample32 sample) {
+                     HistogramBase::Sample32 sample,
+                     std::optional<uint64_t> event_id) {
     EXPECT_EQ(histogram_name, "TestHistogram");
     EXPECT_EQ(sample, 1);
     ++callback_callcount;

base/trace_event/histogram_scope.cc

@@ -4,6 +4,8 @@
 
 #include "base/trace_event/histogram_scope.h"
 
+#include "base/check_op.h"
+
 namespace base::trace_event {
 namespace {
 

base/trace_event/histogram_scope.h

@@ -11,7 +11,6 @@
 #include <string>
 
 #include "base/base_export.h"
-#include "base/trace_event/base_tracing.h"
 
 namespace base::trace_event {
 

content/browser/scheduler/responsiveness/calculator.cc

@@ -13,6 +13,7 @@
 #include "base/metrics/histogram_macros.h"
 #include "base/trace_event/histogram_scope.h"
 #include "base/trace_event/trace_event.h"
+#include "base/trace_event/trace_id_helper.h"
 #include "build/build_config.h"
 #include "content/public/browser/browser_thread.h"
 

content/browser/tracing/background_tracing_rule.cc

@@ -268,6 +268,7 @@ class HistogramRule : public BackgroundTracingRule,
   void OnHistogramChangedCallback(
       base::Histogram::Sample32 reference_lower_value,
       base::Histogram::Sample32 reference_upper_value,
+      std::optional<uint64_t> event_id,
       std::string_view histogram_name,
       uint64_t name_hash,
       base::Histogram::Sample32 actual_value) {
@@ -277,8 +278,8 @@ class HistogramRule : public BackgroundTracingRule,
       return;
     }
 
-    uint64_t flow_id = base::trace_event::HistogramScope::GetFlowId().value_or(
-        base::trace_event::GetNextGlobalTraceId());
+    uint64_t flow_id =
+        event_id.value_or(base::trace_event::GetNextGlobalTraceId());
 
     // Add the histogram name and its corresponding value to the trace.
     const auto trace_details = [&](perfetto::EventContext& ctx) {

services/tracing/public/cpp/background_tracing/background_tracing_agent_impl.cc

@@ -80,6 +80,7 @@ void BackgroundTracingAgentImpl::OnHistogramChanged(
     const std::string& rule_id,
     base::Histogram::Sample32 histogram_lower_value,
     base::Histogram::Sample32 histogram_upper_value,
+    std::optional<uint64_t> event_id,
     std::string_view histogram_name,
     uint64_t name_hash,
     base::Histogram::Sample32 actual_value) {
@@ -89,8 +90,8 @@ void BackgroundTracingAgentImpl::OnHistogramChanged(
   }
 
   auto track = perfetto::NamedTrack("HistogramSamples");
-  uint64_t flow_id = base::trace_event::HistogramScope::GetFlowId().value_or(
-      base::trace_event::GetNextGlobalTraceId());
+  uint64_t flow_id =
+      event_id.value_or(base::trace_event::GetNextGlobalTraceId());
   TRACE_EVENT("toplevel,latency", "HistogramSampleTrigger", track,
               [&](perfetto::EventContext ctx) {
                 perfetto::protos::pbzero::ChromeHistogramSample* new_sample =

services/tracing/public/cpp/background_tracing/background_tracing_agent_impl.h

@@ -49,6 +49,7 @@ class COMPONENT_EXPORT(BACKGROUND_TRACING_CPP) BackgroundTracingAgentImpl
   void OnHistogramChanged(const std::string& rule_id,
                           base::Histogram::Sample32 reference_lower_value,
                           base::Histogram::Sample32 reference_upper_value,
+                          std::optional<uint64_t> event_id,
                           std::string_view histogram_name,
                           uint64_t name_hash,
                           base::Histogram::Sample32 actual_value);

services/tracing/public/cpp/perfetto/histogram_samples_data_source.cc

@@ -136,6 +136,7 @@ void HistogramSamplesDataSource::OnStop(const StopArgs&) {
 void HistogramSamplesDataSource::OnMetricSample(
     std::optional<base::HistogramBase::Sample32> reference_lower_value,
     std::optional<base::HistogramBase::Sample32> reference_upper_value,
+    std::optional<uint64_t> event_id,
     std::string_view histogram_name,
     uint64_t name_hash,
     base::HistogramBase::Sample32 sample) {
@@ -143,18 +144,20 @@ void HistogramSamplesDataSource::OnMetricSample(
       (reference_upper_value && sample > reference_upper_value)) {
     return;
   }
-  OnMetricSampleImpl(histogram_name, name_hash, sample,
+  OnMetricSampleImpl(event_id, histogram_name, name_hash, sample,
                      reinterpret_cast<uintptr_t>(this));
 }
 
 void HistogramSamplesDataSource::OnAnyMetricSample(
     std::string_view histogram_name,
     uint64_t name_hash,
-    base::HistogramBase::Sample32 sample) {
-  OnMetricSampleImpl(histogram_name, name_hash, sample, std::nullopt);
+    base::HistogramBase::Sample32 sample,
+    std::optional<uint64_t> event_id) {
+  OnMetricSampleImpl(event_id, histogram_name, name_hash, sample, std::nullopt);
 }
 
 void HistogramSamplesDataSource::OnMetricSampleImpl(
+    std::optional<uint64_t> event_id,
     std::string_view histogram_name,
     uint64_t name_hash,
     base::HistogramBase::Sample32 sample,
@@ -186,6 +189,9 @@ void HistogramSamplesDataSource::OnMetricSampleImpl(
     event->set_name_iid(1);
     event->add_category_iids(1);
     event->set_type(::perfetto::protos::pbzero::TrackEvent::TYPE_INSTANT);
+    if (event_id) {
+      event->add_flow_ids(*event_id);
+    }
 
     perfetto::protos::pbzero::ChromeHistogramSample* new_sample =
         event->set_chrome_histogram_sample();

services/tracing/public/cpp/perfetto/histogram_samples_data_source.h

@@ -46,15 +46,18 @@ class COMPONENT_EXPORT(TRACING_CPP) HistogramSamplesDataSource
   void OnMetricSample(
       std::optional<base::HistogramBase::Sample32> reference_lower_value,
       std::optional<base::HistogramBase::Sample32> reference_upper_value,
+      std::optional<uint64_t> event_id,
       std::string_view histogram_name,
       uint64_t name_hash,
       base::HistogramBase::Sample32 actual_value);
   static void OnAnyMetricSample(std::string_view histogram_name,
                                 uint64_t name_hash,
-                                base::HistogramBase::Sample32 sample);
+                                base::HistogramBase::Sample32 sample,
+                                std::optional<uint64_t> event_id);
   // `instance` identifies the instance that registered a callback, or nullopt
   // if this is a global callback.
-  static void OnMetricSampleImpl(std::string_view histogram_name,
+  static void OnMetricSampleImpl(std::optional<uint64_t> event_id,
+                                 std::string_view histogram_name,
                                  uint64_t name_hash,
                                  base::HistogramBase::Sample32 sample,
                                  std::optional<uintptr_t> instance);