[record & replay] Move callback cookie into perf buffer spec. (#1656)

Summary: To support record & replay, we move the callback cookie into
the perf buffer spec. Previously, we plumbed it through as a separate
argument. Moving the callback cookie into the perf buffer spec. will
allow us to insert a different callback and save the original callback,
i.e. to interpose a recorder for recording of perf buffer events.

Relevant Issues: #1163

Type of change: /kind feature

Test Plan: Fully covered by existing tests.

---------

Signed-off-by: Pete Stevenson <jps@pixielabs.ai>

Author: Pete Stevenson

src/stirling/bpf_tools/bcc_wrapper.cc

@@ -346,7 +346,7 @@ void BCCWrapper::DetachTracepoints() {
   tracepoints_.clear();
 }
 
-Status BCCWrapper::OpenPerfBuffer(const PerfBufferSpec& perf_buffer, void* cb_cookie) {
+Status BCCWrapper::OpenPerfBuffer(const PerfBufferSpec& perf_buffer) {
   const int kPageSizeBytes = system::Config::GetInstance().PageSizeBytes();
   int num_pages = IntRoundUpDivide(perf_buffer.size_bytes, kPageSizeBytes);
 
@@ -358,15 +358,15 @@ Status BCCWrapper::OpenPerfBuffer(const PerfBufferSpec& perf_buffer, void* cb_co
       perf_buffer.ToString(), num_pages, num_pages * kPageSizeBytes);
   PX_RETURN_IF_ERROR(bpf_.open_perf_buffer(std::string(perf_buffer.name),
                                            perf_buffer.probe_output_fn, perf_buffer.probe_loss_fn,
-                                           cb_cookie, num_pages));
+                                           perf_buffer.cb_cookie, num_pages));
   perf_buffers_.push_back(perf_buffer);
   ++num_open_perf_buffers_;
   return Status::OK();
 }
 
-Status BCCWrapper::OpenPerfBuffers(const ArrayView<PerfBufferSpec>& perf_buffers, void* cb_cookie) {
+Status BCCWrapper::OpenPerfBuffers(const ArrayView<PerfBufferSpec>& perf_buffers) {
   for (const PerfBufferSpec& p : perf_buffers) {
-    PX_RETURN_IF_ERROR(OpenPerfBuffer(p, cb_cookie));
+    PX_RETURN_IF_ERROR(OpenPerfBuffer(p));
   }
   return Status::OK();
 }

src/stirling/bpf_tools/bcc_wrapper.h

@@ -142,11 +142,10 @@ class BCCWrapper {
   /**
    * Open a perf buffer for reading events.
    * @param perf_buff Specifications of the perf buffer (name, callback function, etc.).
-   * @param cb_cookie A pointer that is sent to the callback function when triggered by
    * PollPerfBuffer().
    * @return Error if perf buffer cannot be opened (e.g. perf buffer does not exist).
    */
-  Status OpenPerfBuffer(const PerfBufferSpec& perf_buffer, void* cb_cookie = nullptr);
+  Status OpenPerfBuffer(const PerfBufferSpec& perf_buffer);
 
   /**
    * Attach a perf event, which runs a probe every time a perf counter reaches a threshold
@@ -192,10 +191,9 @@ class BCCWrapper {
   /**
    * Convenience function that opens multiple perf buffers.
    * @param probes Vector of perf buffer descriptors.
-   * @param cb_cookie Raw pointer returned on callback, typically used for tracking context.
    * @return Error of first failure (remaining perf buffer opens are not attempted).
    */
-  Status OpenPerfBuffers(const ArrayView<PerfBufferSpec>& perf_buffers, void* cb_cookie);
+  Status OpenPerfBuffers(const ArrayView<PerfBufferSpec>& perf_buffers);
 
   /**
    * Convenience function that opens multiple perf events.

src/stirling/bpf_tools/probe_specs/probe_specs.h

@@ -170,6 +170,9 @@ struct PerfBufferSpec {
   // Function that will be called if there are lost/clobbered perf events.
   perf_reader_lost_cb probe_loss_fn;
 
+  // Used to invoke callback.
+  void* cb_cookie;
+
   // Size of perf buffer. Will be rounded up to and allocated in a power of 2 number of pages.
   int size_bytes = 1024 * 1024;
 

src/stirling/source_connectors/dynamic_tracer/dynamic_trace_connector.cc

@@ -184,9 +184,10 @@ Status DynamicTraceConnector::InitImpl() {
       .name = bcc_program_.perf_buffer_specs.front().name,
       .probe_output_fn = &GenericHandleEvent,
       .probe_loss_fn = &GenericHandleEventLoss,
+      .cb_cookie = this,
   };
 
-  PX_RETURN_IF_ERROR(OpenPerfBuffer(spec, this));
+  PX_RETURN_IF_ERROR(OpenPerfBuffer(spec));
 
   return Status::OK();
 }

src/stirling/source_connectors/perf_profiler/perf_profile_connector.cc

@@ -138,12 +138,12 @@ Status PerfProfileConnector::InitImpl() {
       {"sample_call_stack", static_cast<uint64_t>(stack_trace_sampling_period_.count())});
 
   const auto perf_buffer_specs = MakeArray<bpf_tools::PerfBufferSpec>(
-      {{std::string(kHistogramAName), HandleHistoEvent, HandleHistoLoss, perf_buffer_size},
-       {std::string(kHistogramBName), HandleHistoEvent, HandleHistoLoss, perf_buffer_size}});
+      {{std::string(kHistogramAName), HandleHistoEvent, HandleHistoLoss, this, perf_buffer_size},
+       {std::string(kHistogramBName), HandleHistoEvent, HandleHistoLoss, this, perf_buffer_size}});
 
   PX_RETURN_IF_ERROR(InitBPFProgram(profiler_bcc_script, defines));
   PX_RETURN_IF_ERROR(AttachSamplingProbes(probe_specs));
-  PX_RETURN_IF_ERROR(OpenPerfBuffers(perf_buffer_specs, this));
+  PX_RETURN_IF_ERROR(OpenPerfBuffers(perf_buffer_specs));
 
   stack_traces_a_ = WrappedBCCStackTable::Create(this, "stack_traces_a");
   stack_traces_b_ = WrappedBCCStackTable::Create(this, "stack_traces_b");

src/stirling/source_connectors/proc_exit/proc_exit_connector.cc

@@ -94,12 +94,12 @@ Status ProcExitConnector::InitImpl() {
   }
 
   const auto perf_buffer_specs = MakeArray<bpf_tools::PerfBufferSpec>({
-      {"proc_exit_events", HandleProcExitEvent, HandleProcExitEventLoss, kPerfBufferPerCPUSizeBytes,
-       bpf_tools::PerfBufferSizeCategory::kControl},
+      {"proc_exit_events", HandleProcExitEvent, HandleProcExitEventLoss, this,
+       kPerfBufferPerCPUSizeBytes, bpf_tools::PerfBufferSizeCategory::kControl},
   });
 
   PX_RETURN_IF_ERROR(AttachTracepoints(kTracepointSpecs));
-  PX_RETURN_IF_ERROR(OpenPerfBuffers(perf_buffer_specs, this));
+  PX_RETURN_IF_ERROR(OpenPerfBuffers(perf_buffer_specs));
 
   return Status::OK();
 }

src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc

@@ -394,22 +394,22 @@ auto SocketTraceConnector::InitPerfBufferSpecs() {
 
   auto specs = MakeArray<bpf_tools::PerfBufferSpec>({
       // For data events. The order must be consistent with output tables.
-      {"socket_data_events", HandleDataEvent, HandleDataEventLoss, kTargetDataBufferSize,
+      {"socket_data_events", HandleDataEvent, HandleDataEventLoss, this, kTargetDataBufferSize,
        PerfBufferSizeCategory::kData},
       // For non-data events. Must not mix with the above perf buffers for data events.
-      {"socket_control_events", HandleControlEvent, HandleControlEventLoss,
+      {"socket_control_events", HandleControlEvent, HandleControlEventLoss, this,
        kTargetControlBufferSize, PerfBufferSizeCategory::kControl},
-      {"conn_stats_events", HandleConnStatsEvent, HandleConnStatsEventLoss,
+      {"conn_stats_events", HandleConnStatsEvent, HandleConnStatsEventLoss, this,
        kTargetControlBufferSize, PerfBufferSizeCategory::kControl},
-      {"mmap_events", HandleMMapEvent, HandleMMapEventLoss, kTargetControlBufferSize / 10,
+      {"mmap_events", HandleMMapEvent, HandleMMapEventLoss, this, kTargetControlBufferSize / 10,
        PerfBufferSizeCategory::kControl},
-      {"go_grpc_events", HandleHTTP2Event, HandleHTTP2EventLoss, kTargetDataBufferSize,
+      {"go_grpc_events", HandleHTTP2Event, HandleHTTP2EventLoss, this, kTargetDataBufferSize,
        PerfBufferSizeCategory::kData},
-      {"grpc_c_events", HandleGrpcCEvent, HandleGrpcCDataLoss, kTargetDataBufferSize,
+      {"grpc_c_events", HandleGrpcCEvent, HandleGrpcCDataLoss, this, kTargetDataBufferSize,
        PerfBufferSizeCategory::kData},
-      {"grpc_c_header_events", HandleGrpcCHeaderEvent, HandleGrpcCHeaderDataLoss,
+      {"grpc_c_header_events", HandleGrpcCHeaderEvent, HandleGrpcCHeaderDataLoss, this,
        kTargetDataBufferSize, PerfBufferSizeCategory::kData},
-      {"grpc_c_close_events", HandleGrpcCCloseEvent, HandleGrpcCCloseDataLoss,
+      {"grpc_c_close_events", HandleGrpcCCloseEvent, HandleGrpcCCloseDataLoss, this,
        kTargetDataBufferSize, PerfBufferSizeCategory::kData},
   });
   ResizePerfBufferSpecs(&specs, category_maximums);
@@ -438,9 +438,9 @@ Status SocketTraceConnector::InitBPF() {
   LOG(INFO) << absl::Substitute("Number of kprobes deployed = $0", kProbeSpecs.size());
   LOG(INFO) << "Probes successfully deployed.";
 
-  const auto kPerfBufferSpecs = InitPerfBufferSpecs();
-  PX_RETURN_IF_ERROR(OpenPerfBuffers(kPerfBufferSpecs, this));
-  LOG(INFO) << absl::Substitute("Number of perf buffers opened = $0", kPerfBufferSpecs.size());
+  const auto perf_buffer_specs = InitPerfBufferSpecs();
+  PX_RETURN_IF_ERROR(OpenPerfBuffers(perf_buffer_specs));
+  LOG(INFO) << absl::Substitute("Number of perf buffers opened = $0", perf_buffer_specs.size());
 
   // Set trace role to BPF probes.
   for (const auto& p : magic_enum::enum_values<traffic_protocol_t>()) {

src/stirling/source_connectors/tcp_stats/tcp_stats_connector.cc

@@ -60,17 +60,17 @@ void HandleTcpEventLoss(void* /*cb_cookie*/, uint64_t /*lost*/) {
   // TODO(RagalahariP): Add stats counter.
 }
 
-const auto kPerfBufferSpecs = MakeArray<bpf_tools::PerfBufferSpec>({
-    {"tcp_events", HandleTcpEvent, HandleTcpEventLoss, kPerfBufferPerCPUSizeBytes,
-     bpf_tools::PerfBufferSizeCategory::kData},
-});
-
 Status TCPStatsConnector::InitImpl() {
+  const auto perf_buffer_specs = MakeArray<bpf_tools::PerfBufferSpec>({
+      {"tcp_events", HandleTcpEvent, HandleTcpEventLoss, this, kPerfBufferPerCPUSizeBytes,
+       bpf_tools::PerfBufferSizeCategory::kData},
+  });
+
   sampling_freq_mgr_.set_period(kSamplingPeriod);
   push_freq_mgr_.set_period(kPushPeriod);
   PX_RETURN_IF_ERROR(InitBPFProgram(tcpstats_bcc_script));
   PX_RETURN_IF_ERROR(AttachKProbes(kProbeSpecs));
-  PX_RETURN_IF_ERROR(OpenPerfBuffers(kPerfBufferSpecs, this));
+  PX_RETURN_IF_ERROR(OpenPerfBuffers(perf_buffer_specs));
   LOG(INFO) << absl::Substitute("Successfully deployed $0 kprobes.", kProbeSpecs.size());
   return Status::OK();
 }