Restore Debug hooks

Summary: This needs to be on the new TransferData() as well as the old one.

Test Plan: Manual.

Reviewers: #stirling, yzhao

Reviewed By: #stirling, yzhao

Differential Revision: https://phab.corp.pixielabs.ai/D8428

GitOrigin-RevId: 12bf5dc8b2d858f17b969a7d3355538b1fbec3cd

src/stirling/e2e_tests/bpf_map_leak_test.cc

@@ -83,7 +83,7 @@ TEST_F(BPFMapLeakTest, unclosed_connection) {
   // And because it was killed, it should have leaked a BPF map entry.
 
   // For testing, make sure Stirling cleans up BPF entries right away.
-  // Without this flag, Stirling delays clean-up to acccumulate a clean-up batch.
+  // Without this flag, Stirling delays clean-up to accumulate a clean-up batch.
   FLAGS_stirling_conn_map_cleanup_threshold = 1;
 
   DataTable data_table(kHTTPTable);

src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc

@@ -330,6 +330,15 @@ void SocketTraceConnector::CachedUpdateCommonState(ConnectorContext* ctx, uint32
   table_access_history_.reset(table_num);
 }
 
+void SocketTraceConnector::UpdateTrackerTraceLevel(ConnTracker* tracker) {
+  if (pids_to_trace_.contains(tracker->conn_id().upid.pid)) {
+    tracker->SetDebugTrace(2);
+  }
+  if (pids_to_trace_disable_.contains(tracker->conn_id().upid.pid)) {
+    tracker->SetDebugTrace(0);
+  }
+}
+
 void SocketTraceConnector::TransferDataImpl(ConnectorContext* ctx, uint32_t table_num,
                                             DataTable* data_table) {
   DCHECK_LT(table_num, kTables.size())
@@ -393,6 +402,8 @@ void SocketTraceConnector::TransferDataImpl(ConnectorContext* ctx,
         continue;
       }
 
+      UpdateTrackerTraceLevel(conn_tracker.get());
+
       conn_tracker->IterationPreTick(iteration_start_time_, cluster_cidrs, proc_parser_.get(),
                                      socket_info_mgr_.get());
       if (transfer_spec.enabled && transfer_spec.transfer_fn) {
@@ -401,6 +412,9 @@ void SocketTraceConnector::TransferDataImpl(ConnectorContext* ctx,
       conn_tracker->IterationPostTick();
     }
   }
+
+  // Once we've cleared all the debug trace levels for this pid, we can remove it from the list.
+  pids_to_trace_disable_.clear();
 }
 
 bool SocketTraceConnector::output_multi_tables() const {
@@ -895,15 +909,6 @@ void SocketTraceConnector::WriteDataEvent(const SocketDataEvent& event) {
 // TransferData Helpers
 //-----------------------------------------------------------------------------
 
-void SocketTraceConnector::UpdateTrackerTraceLevel(ConnTracker* tracker) {
-  if (pids_to_trace_.contains(tracker->conn_id().upid.pid)) {
-    tracker->SetDebugTrace(2);
-  }
-  if (pids_to_trace_disable_.contains(tracker->conn_id().upid.pid)) {
-    tracker->SetDebugTrace(0);
-  }
-}
-
 void SocketTraceConnector::TransferStreams(ConnectorContext* ctx, uint32_t table_num,
                                            DataTable* data_table) {
   std::vector<CIDRBlock> cluster_cidrs = ctx->GetClusterCIDRs();