Refactor Network Tracing into separate data source

This splits the polling and DataSource aspects into separate classes
in preparation for making polling a singleton. This change is not
expected to change behavior.

Note:
* OnSetup saves the poll duration and passes it to Start to simplify
  the Poller's life cycle.
* OnStart's handling of mTaskRunner/Loop is moved to the end of Start.
* OnStop's handling of mTaskRunner is moved to the end of Stop and Stop
  is changed to continue on failure (so that the task runner is reset in
  all cases as it was before this change).

Bug: 246985031
Test: atest libnetworkstats_test
Change-Id: I521eb50cd00916aa4d98174f0db461508b532732
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
index aeadb4a..894ff7d 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
@@ -56,10 +56,10 @@
 }
 
 NetworkTraceHandler::NetworkTraceHandler()
-    : NetworkTraceHandler([this](const PacketTrace& pkt) {
+    : mPoller([](const PacketTrace& pkt) {
         NetworkTraceHandler::Trace(
-            [this, pkt](NetworkTraceHandler::TraceContext ctx) {
-              Fill(pkt, *ctx.NewTracePacket());
+            [pkt](NetworkTraceHandler::TraceContext ctx) {
+              NetworkTraceHandler::Fill(pkt, *ctx.NewTracePacket());
             });
       }) {}
 
@@ -74,22 +74,16 @@
   }
 }
 
-void NetworkTraceHandler::OnStart(const StartArgs&) {
-  if (!Start()) return;
-  mTaskRunner = perfetto::Platform::GetDefaultPlatform()->CreateTaskRunner({});
-  Loop();
-}
+void NetworkTraceHandler::OnStart(const StartArgs&) { mPoller.Start(mPollMs); }
+void NetworkTraceHandler::OnStop(const StopArgs&) { mPoller.Stop(); }
 
-void NetworkTraceHandler::OnStop(const StopArgs&) {
-  Stop();
-  mTaskRunner.reset();
-}
-
-void NetworkTraceHandler::Loop() {
-  mTaskRunner->PostDelayedTask([this]() { Loop(); }, mPollMs);
+void NetworkTracePoller::SchedulePolling() {
+  // Schedules another run of ourselves to recursively poll periodically.
+  mTaskRunner->PostDelayedTask([this]() { SchedulePolling(); }, mPollMs);
   ConsumeAll();
 }
 
+// static class method
 void NetworkTraceHandler::Fill(const PacketTrace& src, TracePacket& dst) {
   dst.set_timestamp(src.timestampNs);
   auto* event = dst.set_network_packet();
@@ -113,7 +107,7 @@
   }
 }
 
-bool NetworkTraceHandler::Start() {
+bool NetworkTracePoller::Start(uint32_t pollMs) {
   ALOGD("Starting datasource");
 
   auto status = mConfigurationMap.init(PACKET_TRACE_ENABLED_MAP_PATH);
@@ -136,24 +130,29 @@
     return false;
   }
 
+  // Start a task runner to run ConsumeAll every mPollMs milliseconds.
+  mTaskRunner = perfetto::Platform::GetDefaultPlatform()->CreateTaskRunner({});
+  mPollMs = pollMs;
+  SchedulePolling();
+
   return true;
 }
 
-bool NetworkTraceHandler::Stop() {
+bool NetworkTracePoller::Stop() {
   ALOGD("Stopping datasource");
 
   auto res = mConfigurationMap.writeValue(0, false, BPF_ANY);
   if (!res.ok()) {
     ALOGW("Failed to disable tracing: %s", res.error().message().c_str());
-    return false;
   }
 
+  mTaskRunner.reset();
   mRingBuffer.reset();
 
-  return true;
+  return res.ok();
 }
 
-bool NetworkTraceHandler::ConsumeAll() {
+bool NetworkTracePoller::ConsumeAll() {
   if (mRingBuffer == nullptr) {
     ALOGW("Tracing is not active");
     return false;
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
index 560194f..8d03e2a 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
@@ -39,6 +39,9 @@
 
 namespace android {
 namespace bpf {
+// Use uint32 max to cause the handler to never Loop. Instead, the tests will
+// manually drive things by calling ConsumeAll explicitly.
+constexpr uint32_t kNeverPoll = std::numeric_limits<uint32_t>::max();
 
 __be16 bindAndListen(int s) {
   sockaddr_in sin = {.sin_family = AF_INET};
@@ -83,7 +86,7 @@
   }
 };
 
-class NetworkTraceHandlerTest : public testing::Test {
+class NetworkTracePollerTest : public testing::Test {
  protected:
   void SetUp() {
     if (access(PACKET_TRACE_RINGBUF_PATH, R_OK)) {
@@ -95,31 +98,31 @@
   }
 };
 
-TEST_F(NetworkTraceHandlerTest, PollWhileInactive) {
-  NetworkTraceHandler handler([&](const PacketTrace& pkt) {});
+TEST_F(NetworkTracePollerTest, PollWhileInactive) {
+  NetworkTracePoller handler([&](const PacketTrace& pkt) {});
 
   // One succeed after start and before stop.
   EXPECT_FALSE(handler.ConsumeAll());
-  ASSERT_TRUE(handler.Start());
+  ASSERT_TRUE(handler.Start(kNeverPoll));
   EXPECT_TRUE(handler.ConsumeAll());
   ASSERT_TRUE(handler.Stop());
   EXPECT_FALSE(handler.ConsumeAll());
 }
 
-TEST_F(NetworkTraceHandlerTest, TraceTcpSession) {
+TEST_F(NetworkTracePollerTest, TraceTcpSession) {
   __be16 server_port = 0;
   std::vector<PacketTrace> packets;
 
   // Record all packets with the bound address and current uid. This callback is
   // involked only within ConsumeAll, at which point the port should have
   // already been filled in and all packets have been processed.
-  NetworkTraceHandler handler([&](const PacketTrace& pkt) {
+  NetworkTracePoller handler([&](const PacketTrace& pkt) {
     if (pkt.sport != server_port && pkt.dport != server_port) return;
     if (pkt.uid != getuid()) return;
     packets.push_back(pkt);
   });
 
-  ASSERT_TRUE(handler.Start());
+  ASSERT_TRUE(handler.Start(kNeverPoll));
   const uint32_t kClientTag = 2468;
   const uint32_t kServerTag = 1357;
 
diff --git a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
index c257aa0..1ab35ef 100644
--- a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
+++ b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
@@ -31,37 +31,22 @@
 namespace android {
 namespace bpf {
 
-class NetworkTraceHandler : public perfetto::DataSource<NetworkTraceHandler> {
+// NetworkTracePoller is responsible for interactions with the BPF ring buffer
+// including polling. This class is an internal helper for NetworkTraceHandler,
+// it is not meant to be used elsewhere.
+class NetworkTracePoller {
  public:
-  // Registers this DataSource.
-  static void RegisterDataSource();
-
-  // Connects to the system Perfetto daemon and registers the trace handler.
-  static void InitPerfettoTracing();
-
-  // Initialize with the default Perfetto callback.
-  NetworkTraceHandler();
-
   // Testonly: initialize with a callback capable of intercepting data.
-  NetworkTraceHandler(std::function<void(const PacketTrace&)> callback)
+  NetworkTracePoller(std::function<void(const PacketTrace&)> callback)
       : mCallback(std::move(callback)) {}
 
   // Testonly: standalone functions without perfetto dependency.
-  bool Start();
+  bool Start(uint32_t pollMs);
   bool Stop();
   bool ConsumeAll();
 
-  // perfetto::DataSource overrides:
-  void OnSetup(const SetupArgs&) override;
-  void OnStart(const StartArgs&) override;
-  void OnStop(const StopArgs&) override;
-
-  // Convert a PacketTrace into a Perfetto trace packet.
-  void Fill(const PacketTrace& src,
-            ::perfetto::protos::pbzero::TracePacket& dst);
-
  private:
-  void Loop();
+  void SchedulePolling();
 
   // How often to poll the ring buffer, defined by the trace config.
   uint32_t mPollMs;
@@ -80,5 +65,33 @@
   std::unique_ptr<perfetto::base::TaskRunner> mTaskRunner;
 };
 
+// NetworkTraceHandler implements the android.network_packets data source. This
+// class is registered with Perfetto and is instantiated when tracing starts and
+// destroyed when tracing ends. There is one instance per trace session.
+class NetworkTraceHandler : public perfetto::DataSource<NetworkTraceHandler> {
+ public:
+  // Registers this DataSource.
+  static void RegisterDataSource();
+
+  // Connects to the system Perfetto daemon and registers the trace handler.
+  static void InitPerfettoTracing();
+
+  // Initialize with the default Perfetto callback.
+  NetworkTraceHandler();
+
+  // perfetto::DataSource overrides:
+  void OnSetup(const SetupArgs& args) override;
+  void OnStart(const StartArgs&) override;
+  void OnStop(const StopArgs&) override;
+
+ private:
+  // Convert a PacketTrace into a Perfetto trace packet.
+  static void Fill(const PacketTrace& src,
+                   ::perfetto::protos::pbzero::TracePacket& dst);
+
+  NetworkTracePoller mPoller;
+  uint32_t mPollMs;
+};
+
 }  // namespace bpf
 }  // namespace android