tracing: Add a way to observe track event sessions

This patch adds a new TrackEventSessionObserver interface in order to
allow applications to listen to track event tracing session lifecycle
events.

Bug: chromium:1006541

Change-Id: I0849924a782883a505dd42efe0eaea2fd74d24d0
diff --git a/CHANGELOG b/CHANGELOG
index 8841f7c..65d16b4 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -3,6 +3,8 @@
     * Added trace packet interceptor API for rerouting trace data into
       non-Perfetto systems.
     * Added support for printing track events to the console.
+    * Added a way to observe track event tracing sessions starting and
+      stopping.
   Trace Processor:
     *
   UI:
diff --git a/docs/instrumentation/track-events.md b/docs/instrumentation/track-events.md
index 1e673a0..6cf61e7 100644
--- a/docs/instrumentation/track-events.md
+++ b/docs/instrumentation/track-events.md
@@ -15,8 +15,8 @@
 section of the Tracing SDK page for instructions on how to check out and
 build the SDK.
 
-TIP: The code from this example is also available as a
-     [GitHub repository](https://github.com/skyostil/perfetto-sdk-example).
+TIP: The code from these examples is also available [in the
+repository](/examples/sdk/README.md).
 
 There are a few main types of track events:
 
@@ -437,4 +437,62 @@
 Note that interned data is strongly typed, i.e., each class of interned data
 uses a separate namespace for identifiers.
 
-[RAII]: https://en.cppreference.com/w/cpp/language/raii
+### Tracing session observers
+
+The session observer interface allows applications to be notified when track
+event tracing starts and stops:
+
+```C++
+class Observer : public perfetto::TrackEventSessionObserver {
+  public:
+  ~Observer() override = default;
+
+  void OnSetup(const perfetto::DataSourceBase::SetupArgs&) override {
+    // Called when tracing session is configured. Note tracing isn't active yet,
+    // so track events emitted here won't be recorded.
+  }
+
+  void OnStart(const DataSourceBase::SetupArgs&) override {
+    // Called when a tracing session is started. It is possible to emit track
+    // events from this callback.
+  }
+
+  void OnStop(const DataSourceBase::StartArgs&) override {
+    // Called when a tracing session is stopped. It is still possible to emit
+    // track events from this callback.
+  }
+};
+```
+
+Note that all methods of the interface are called on an internal Perfetto
+thread.
+
+For example, here's how to wait for any tracing session to start:
+
+```C++
+class Observer : public perfetto::TrackEventSessionObserver {
+ public:
+  Observer() { perfetto::TrackEvent::AddSessionObserver(this); }
+  ~Observer() { perfetto::TrackEvent::RemoveSessionObserver(this); }
+
+  void OnStart(const perfetto::DataSourceBase::StartArgs&) override {
+    std::unique_lock<std::mutex> lock(mutex);
+    cv.notify_one();
+  }
+
+  void WaitForTracingStart() {
+    printf("Waiting for tracing to start...\n");
+    std::unique_lock<std::mutex> lock(mutex);
+    cv.wait(lock, [] { return perfetto::TrackEvent::IsEnabled(); });
+    printf("Tracing started\n");
+  }
+
+  std::mutex mutex;
+  std::condition_variable cv;
+};
+
+Observer observer;
+observer.WaitForTracingToStart();
+```
+
+[RAII]: https://en.cppreference.com/w/cpp/language/raii
\ No newline at end of file
diff --git a/examples/sdk/example_system_wide.cc b/examples/sdk/example_system_wide.cc
index f755d37..b7101de 100644
--- a/examples/sdk/example_system_wide.cc
+++ b/examples/sdk/example_system_wide.cc
@@ -19,9 +19,31 @@
 #include "trace_categories.h"
 
 #include <chrono>
+#include <condition_variable>
 #include <fstream>
 #include <thread>
 
+class Observer : public perfetto::TrackEventSessionObserver {
+ public:
+  Observer() { perfetto::TrackEvent::AddSessionObserver(this); }
+  ~Observer() { perfetto::TrackEvent::RemoveSessionObserver(this); }
+
+  void OnStart(const perfetto::DataSourceBase::StartArgs&) override {
+    std::unique_lock<std::mutex> lock(mutex);
+    cv.notify_one();
+  }
+
+  void WaitForTracingStart() {
+    PERFETTO_LOG("Waiting for tracing to start...");
+    std::unique_lock<std::mutex> lock(mutex);
+    cv.wait(lock, [] { return perfetto::TrackEvent::IsEnabled(); });
+    PERFETTO_LOG("Tracing started");
+  }
+
+  std::mutex mutex;
+  std::condition_variable cv;
+};
+
 void InitializePerfetto() {
   perfetto::TracingInitArgs args;
   // The backends determine where trace events are recorded. For this example we
@@ -33,14 +55,6 @@
   perfetto::TrackEvent::Register();
 }
 
-void WaitForTracingStart() {
-  PERFETTO_LOG("Waiting for tracing to start...");
-  while (!TRACE_EVENT_CATEGORY_ENABLED("rendering")) {
-    std::this_thread::sleep_for(std::chrono::milliseconds(100));
-  }
-  PERFETTO_LOG("Tracing started");
-}
-
 void DrawPlayer(int player_number) {
   TRACE_EVENT("rendering", "DrawPlayer", "player_number", player_number);
   // Sleep to simulate a long computation.
@@ -55,7 +69,9 @@
 
 int main(int, const char**) {
   InitializePerfetto();
-  WaitForTracingStart();
+
+  Observer observer;
+  observer.WaitForTracingStart();
 
   // Simulate some work that emits trace events.
   // Note that we don't start and stop tracing here; for system-wide tracing
diff --git a/include/perfetto/tracing/internal/track_event_data_source.h b/include/perfetto/tracing/internal/track_event_data_source.h
index bbae795..aa48467 100644
--- a/include/perfetto/tracing/internal/track_event_data_source.h
+++ b/include/perfetto/tracing/internal/track_event_data_source.h
@@ -91,25 +91,46 @@
   using Base = DataSource<DataSourceType, TrackEventDataSourceTraits>;
 
  public:
+  // Add or remove a session observer for this track event data source. The
+  // observer will be notified about started and stopped tracing sessions.
+  // Returns |true| if the observer was succesfully added (i.e., the maximum
+  // number of observers wasn't exceeded).
+  static bool AddSessionObserver(TrackEventSessionObserver* observer) {
+    return TrackEventInternal::AddSessionObserver(observer);
+  }
+
+  static void RemoveSessionObserver(TrackEventSessionObserver* observer) {
+    TrackEventInternal::RemoveSessionObserver(observer);
+  }
+
   // DataSource implementation.
   void OnSetup(const DataSourceBase::SetupArgs& args) override {
     auto config_raw = args.config->track_event_config_raw();
     bool ok = config_.ParseFromArray(config_raw.data(), config_raw.size());
     PERFETTO_DCHECK(ok);
-    TrackEventInternal::EnableTracing(*Registry, config_,
-                                      args.internal_instance_index);
+    TrackEventInternal::EnableTracing(*Registry, config_, args);
   }
 
-  void OnStart(const DataSourceBase::StartArgs&) override {}
+  void OnStart(const DataSourceBase::StartArgs& args) override {
+    TrackEventInternal::OnStart(args);
+  }
 
   void OnStop(const DataSourceBase::StopArgs& args) override {
-    TrackEventInternal::DisableTracing(*Registry, args.internal_instance_index);
+    TrackEventInternal::DisableTracing(*Registry, args);
   }
 
   static void Flush() {
     Base::template Trace([](typename Base::TraceContext ctx) { ctx.Flush(); });
   }
 
+  // Determine if *any* tracing category is enabled.
+  static bool IsEnabled() {
+    bool enabled = false;
+    Base::template CallIfEnabled(
+        [&](uint32_t /*instances*/) { enabled = true; });
+    return enabled;
+  }
+
   // Determine if tracing for the given static category is enabled.
   template <size_t CategoryIndex>
   static bool IsCategoryEnabled() {
diff --git a/include/perfetto/tracing/internal/track_event_internal.h b/include/perfetto/tracing/internal/track_event_internal.h
index c5f3020..cfd2be2 100644
--- a/include/perfetto/tracing/internal/track_event_internal.h
+++ b/include/perfetto/tracing/internal/track_event_internal.h
@@ -20,6 +20,7 @@
 #include "perfetto/base/flat_set.h"
 #include "perfetto/protozero/scattered_heap_buffer.h"
 #include "perfetto/tracing/core/forward_decls.h"
+#include "perfetto/tracing/data_source.h"
 #include "perfetto/tracing/debug_annotation.h"
 #include "perfetto/tracing/trace_writer_base.h"
 #include "perfetto/tracing/track.h"
@@ -31,6 +32,7 @@
 
 namespace perfetto {
 class EventContext;
+class TrackEventSessionObserver;
 struct Category;
 namespace protos {
 namespace gen {
@@ -41,6 +43,24 @@
 }  // namespace pbzero
 }  // namespace protos
 
+// A callback interface for observing track event tracing sessions starting and
+// stopping. See TrackEvent::{Add,Remove}SessionObserver. Note that all methods
+// will be called on an internal Perfetto thread.
+class TrackEventSessionObserver {
+ public:
+  virtual ~TrackEventSessionObserver();
+  // Called when a track event tracing session is configured. Note tracing isn't
+  // active yet, so track events emitted here won't be recorded. See
+  // DataSourceBase::OnSetup.
+  virtual void OnSetup(const DataSourceBase::SetupArgs&);
+  // Called when a track event tracing session is started. It is possible to
+  // emit track events from this callback.
+  virtual void OnStart(const DataSourceBase::StartArgs&);
+  // Called when a track event tracing session is stopped. It is still possible
+  // to emit track events from this callback.
+  virtual void OnStop(const DataSourceBase::StopArgs&);
+};
+
 namespace internal {
 class TrackEventCategoryRegistry;
 
@@ -99,11 +119,15 @@
       const TrackEventCategoryRegistry&,
       bool (*register_data_source)(const DataSourceDescriptor&));
 
+  static bool AddSessionObserver(TrackEventSessionObserver*);
+  static void RemoveSessionObserver(TrackEventSessionObserver*);
+
   static void EnableTracing(const TrackEventCategoryRegistry& registry,
                             const protos::gen::TrackEventConfig& config,
-                            uint32_t instance_index);
+                            const DataSourceBase::SetupArgs&);
+  static void OnStart(const DataSourceBase::StartArgs&);
   static void DisableTracing(const TrackEventCategoryRegistry& registry,
-                             uint32_t instance_index);
+                             const DataSourceBase::StopArgs&);
   static bool IsCategoryEnabled(const TrackEventCategoryRegistry& registry,
                                 const protos::gen::TrackEventConfig& config,
                                 const Category& category);
diff --git a/src/tracing/internal/track_event_internal.cc b/src/tracing/internal/track_event_internal.cc
index 94c2a77..77ecf18 100644
--- a/src/tracing/internal/track_event_internal.cc
+++ b/src/tracing/internal/track_event_internal.cc
@@ -31,6 +31,12 @@
 #include "protos/perfetto/trace/track_event/track_descriptor.pbzero.h"
 
 namespace perfetto {
+
+TrackEventSessionObserver::~TrackEventSessionObserver() = default;
+void TrackEventSessionObserver::OnSetup(const DataSourceBase::SetupArgs&) {}
+void TrackEventSessionObserver::OnStart(const DataSourceBase::StartArgs&) {}
+void TrackEventSessionObserver::OnStop(const DataSourceBase::StopArgs&) {}
+
 namespace internal {
 
 BaseTrackEventInternedDataIndex::~BaseTrackEventInternedDataIndex() = default;
@@ -42,6 +48,19 @@
 static constexpr const char kSlowTag[] = "slow";
 static constexpr const char kDebugTag[] = "debug";
 
+void ForEachObserver(
+    std::function<bool(TrackEventSessionObserver*&)> callback) {
+  // Session observers, shared by all track event data source instances.
+  static constexpr int kMaxObservers = 8;
+  static std::recursive_mutex* mutex = new std::recursive_mutex{};  // Leaked.
+  static std::array<TrackEventSessionObserver*, kMaxObservers> observers{};
+  std::unique_lock<std::recursive_mutex> lock(*mutex);
+  for (auto& o : observers) {
+    if (!callback(o))
+      break;
+  }
+}
+
 struct InternedEventCategory
     : public TrackEventInternedDataIndex<
           InternedEventCategory,
@@ -152,22 +171,68 @@
 }
 
 // static
+bool TrackEventInternal::AddSessionObserver(
+    TrackEventSessionObserver* observer) {
+  bool result = false;
+  ForEachObserver([&](TrackEventSessionObserver*& o) {
+    if (!o) {
+      o = observer;
+      result = true;
+      return false;
+    }
+    return true;
+  });
+  return result;
+}
+
+// static
+void TrackEventInternal::RemoveSessionObserver(
+    TrackEventSessionObserver* observer) {
+  ForEachObserver([&](TrackEventSessionObserver*& o) {
+    if (o == observer) {
+      o = nullptr;
+      return false;
+    }
+    return true;
+  });
+}
+
+// static
 void TrackEventInternal::EnableTracing(
     const TrackEventCategoryRegistry& registry,
     const protos::gen::TrackEventConfig& config,
-    uint32_t instance_index) {
+    const DataSourceBase::SetupArgs& args) {
   for (size_t i = 0; i < registry.category_count(); i++) {
     if (IsCategoryEnabled(registry, config, *registry.GetCategory(i)))
-      registry.EnableCategoryForInstance(i, instance_index);
+      registry.EnableCategoryForInstance(i, args.internal_instance_index);
   }
+  ForEachObserver([&](TrackEventSessionObserver*& o) {
+    if (o)
+      o->OnSetup(args);
+    return true;
+  });
+}
+
+// static
+void TrackEventInternal::OnStart(const DataSourceBase::StartArgs& args) {
+  ForEachObserver([&](TrackEventSessionObserver*& o) {
+    if (o)
+      o->OnStart(args);
+    return true;
+  });
 }
 
 // static
 void TrackEventInternal::DisableTracing(
     const TrackEventCategoryRegistry& registry,
-    uint32_t instance_index) {
+    const DataSourceBase::StopArgs& args) {
+  ForEachObserver([&](TrackEventSessionObserver*& o) {
+    if (o)
+      o->OnStop(args);
+    return true;
+  });
   for (size_t i = 0; i < registry.category_count(); i++)
-    registry.DisableCategoryForInstance(i, instance_index);
+    registry.DisableCategoryForInstance(i, args.internal_instance_index);
 }
 
 // static
diff --git a/src/tracing/test/api_integrationtest.cc b/src/tracing/test/api_integrationtest.cc
index 5525876..aa9824e 100644
--- a/src/tracing/test/api_integrationtest.cc
+++ b/src/tracing/test/api_integrationtest.cc
@@ -3242,6 +3242,130 @@
   EXPECT_THAT(lines, ContainerEq(golden_lines));
 }
 
+TEST_P(PerfettoApiTest, TrackEventObserver) {
+  class Observer : public perfetto::TrackEventSessionObserver {
+   public:
+    ~Observer() override = default;
+
+    void OnSetup(const perfetto::DataSourceBase::SetupArgs&) {
+      // Since other tests here register multiple track event data sources,
+      // ignore all but the first notifications.
+      if (setup_called)
+        return;
+      setup_called = true;
+      if (unsubscribe_at_setup)
+        perfetto::TrackEvent::RemoveSessionObserver(this);
+      // This event isn't recorded in the trace because tracing isn't active yet
+      // when OnSetup is called.
+      TRACE_EVENT_INSTANT("foo", "OnSetup");
+      // However the active tracing categories have already been updated at this
+      // point.
+      EXPECT_TRUE(perfetto::TrackEvent::IsEnabled());
+      EXPECT_TRUE(TRACE_EVENT_CATEGORY_ENABLED("foo"));
+    }
+
+    void OnStart(const perfetto::DataSourceBase::StartArgs&) {
+      if (start_called)
+        return;
+      start_called = true;
+      EXPECT_TRUE(perfetto::TrackEvent::IsEnabled());
+      EXPECT_TRUE(TRACE_EVENT_CATEGORY_ENABLED("foo"));
+      TRACE_EVENT_INSTANT("foo", "OnStart");
+    }
+
+    void OnStop(const perfetto::DataSourceBase::StopArgs&) {
+      if (stop_called)
+        return;
+      stop_called = true;
+      EXPECT_TRUE(perfetto::TrackEvent::IsEnabled());
+      EXPECT_TRUE(TRACE_EVENT_CATEGORY_ENABLED("foo"));
+      TRACE_EVENT_INSTANT("foo", "OnStop");
+      perfetto::TrackEvent::Flush();
+    }
+
+    bool setup_called{};
+    bool start_called{};
+    bool stop_called{};
+    bool unsubscribe_at_setup{};
+  };
+
+  EXPECT_FALSE(perfetto::TrackEvent::IsEnabled());
+  {
+    Observer observer;
+    perfetto::TrackEvent::AddSessionObserver(&observer);
+
+    auto* tracing_session = NewTraceWithCategories({"foo"});
+    tracing_session->get()->StartBlocking();
+    EXPECT_TRUE(observer.setup_called);
+    EXPECT_TRUE(observer.start_called);
+    tracing_session->get()->StopBlocking();
+    EXPECT_TRUE(observer.stop_called);
+    perfetto::TrackEvent::RemoveSessionObserver(&observer);
+    auto slices = ReadSlicesFromTrace(tracing_session->get());
+    EXPECT_THAT(slices, ElementsAre("I:foo.OnStart", "I:foo.OnStop"));
+  }
+
+  // No notifications after removing observer.
+  {
+    Observer observer;
+    perfetto::TrackEvent::AddSessionObserver(&observer);
+    perfetto::TrackEvent::RemoveSessionObserver(&observer);
+    auto* tracing_session = NewTraceWithCategories({"foo"});
+    tracing_session->get()->StartBlocking();
+    EXPECT_FALSE(observer.setup_called);
+    EXPECT_FALSE(observer.start_called);
+    tracing_session->get()->StopBlocking();
+    EXPECT_FALSE(observer.stop_called);
+  }
+
+  // Removing observer in a callback.
+  {
+    Observer observer;
+    observer.unsubscribe_at_setup = true;
+    perfetto::TrackEvent::AddSessionObserver(&observer);
+    auto* tracing_session = NewTraceWithCategories({"foo"});
+    tracing_session->get()->StartBlocking();
+    EXPECT_TRUE(observer.setup_called);
+    EXPECT_FALSE(observer.start_called);
+    tracing_session->get()->StopBlocking();
+    EXPECT_FALSE(observer.stop_called);
+    perfetto::TrackEvent::RemoveSessionObserver(&observer);
+  }
+
+  // Multiple observers.
+  {
+    Observer observer1;
+    Observer observer2;
+    perfetto::TrackEvent::AddSessionObserver(&observer1);
+    perfetto::TrackEvent::AddSessionObserver(&observer2);
+    auto* tracing_session = NewTraceWithCategories({"foo"});
+    tracing_session->get()->StartBlocking();
+    tracing_session->get()->StopBlocking();
+    perfetto::TrackEvent::RemoveSessionObserver(&observer1);
+    perfetto::TrackEvent::RemoveSessionObserver(&observer2);
+    auto slices = ReadSlicesFromTrace(tracing_session->get());
+    EXPECT_THAT(slices, ElementsAre("I:foo.OnStart", "I:foo.OnStart",
+                                    "I:foo.OnStop", "I:foo.OnStop"));
+  }
+
+  // Multiple observers with one being removed midway.
+  {
+    Observer observer1;
+    Observer observer2;
+    perfetto::TrackEvent::AddSessionObserver(&observer1);
+    perfetto::TrackEvent::AddSessionObserver(&observer2);
+    auto* tracing_session = NewTraceWithCategories({"foo"});
+    tracing_session->get()->StartBlocking();
+    perfetto::TrackEvent::RemoveSessionObserver(&observer1);
+    tracing_session->get()->StopBlocking();
+    perfetto::TrackEvent::RemoveSessionObserver(&observer2);
+    auto slices = ReadSlicesFromTrace(tracing_session->get());
+    EXPECT_THAT(slices,
+                ElementsAre("I:foo.OnStart", "I:foo.OnStart", "I:foo.OnStop"));
+  }
+  EXPECT_FALSE(perfetto::TrackEvent::IsEnabled());
+}
+
 struct BackendTypeAsString {
   std::string operator()(
       const ::testing::TestParamInfo<perfetto::BackendType>& info) const {