[record & replay] Add rr for eBPF maps & arrays. (#1714)

Summary: We add the record & replay capability for eBPF maps & arrays. PerCPU array and stack table will be added in a future PR. Type of change: /kind feature Test Plan: We extend the test case: `rr_bpf_test`. --------- Signed-off-by: Pete Stevenson <[email protected]>
pixie-io · Oct 3, 2023 · 4218e02 · 4218e02
1 parent ba1010e
commit 4218e02
Show file tree

Hide file tree

Showing 5 changed files with 453 additions and 5 deletions.
diff --git a/src/stirling/bpf_tools/bcc_wrapper.cc b/src/stirling/bpf_tools/bcc_wrapper.cc
@@ -485,7 +485,9 @@ std::unique_ptr<WrappedBCCStackTable> WrappedBCCStackTable::Create(bpf_tools::BC
                                                                    const std::string& name) {
   using BaseT = WrappedBCCStackTable;
   using ImplT = WrappedBCCStackTableImpl;
-  return CreateBCCWrappedMapOrArray<BaseT, ImplT>(bcc, name);
+
+  // TODO(jps): Impl. rr for stack table.
+  return CreateBCCWrappedMapOrArray<BaseT, ImplT, ImplT, ImplT>(bcc, name);
 }
 
 }  // namespace bpf_tools

diff --git a/src/stirling/bpf_tools/bcc_wrapper.h b/src/stirling/bpf_tools/bcc_wrapper.h
@@ -499,6 +499,44 @@ class WrappedBCCArrayTableImpl : public WrappedBCCArrayTable<T> {
   std::unique_ptr<U> underlying_;
 };
 
+template <typename T>
+class RecordingWrappedBCCArrayTableImpl : public WrappedBCCArrayTableImpl<T> {
+ public:
+  using Super = WrappedBCCArrayTableImpl<T>;
+
+  StatusOr<T> GetValue(const uint32_t idx) override {
+    PX_ASSIGN_OR_RETURN(const T value, Super::GetValue(idx));
+    recorder_.RecordBPFArrayTableGetValueEvent(this->name_, idx, sizeof(value), &value);
+    return value;
+  }
+
+  RecordingWrappedBCCArrayTableImpl(bpf_tools::BCCWrapper* bcc, const std::string& name)
+      : WrappedBCCArrayTableImpl<T>(bcc, name),
+        recorder_(*bcc->GetBPFRecorder().ConsumeValueOrDie()) {}
+
+ private:
+  BPFRecorder& recorder_;
+};
+
+template <typename T>
+class ReplayingWrappedBCCArrayTableImpl : public WrappedBCCArrayTable<T> {
+ public:
+  StatusOr<T> GetValue(const uint32_t idx) override {
+    T value;
+    PX_RETURN_IF_ERROR(replayer_.ReplayArrayGetValue(this->name_, idx, sizeof(T), &value));
+    return value;
+  }
+
+  Status SetValue(const uint32_t, const T&) override { return Status::OK(); }
+
+  ReplayingWrappedBCCArrayTableImpl(bpf_tools::BCCWrapper* bcc, const std::string& name)
+      : replayer_(*bcc->GetBPFReplayer().ConsumeValueOrDie()), name_(name) {}
+
+ private:
+  BPFReplayer& replayer_;
+  const std::string name_;
+};
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // Map / BPF Hash Table
@@ -598,6 +636,88 @@ class WrappedBCCMapImpl : public WrappedBCCMap<K, V, kUserSpaceManaged> {
   absl::flat_hash_set<K> shadow_keys_;
 };
 
+template <typename K, typename V, bool kUserSpaceManaged = false>
+class RecordingWrappedBCCMapImpl : public WrappedBCCMapImpl<K, V, kUserSpaceManaged> {
+ public:
+  using Super = WrappedBCCMapImpl<K, V, kUserSpaceManaged>;
+
+  StatusOr<V> GetValue(const K& key) const override {
+    PX_ASSIGN_OR_RETURN(const V value, Super::GetValue(key));
+    recorder_.RecordBPFMapGetValueEvent(this->name_, sizeof(key), &key, sizeof(value), &value);
+    return value;
+  }
+
+  size_t capacity() const override {
+    const size_t n = Super::capacity();
+    recorder_.RecordBPFMapCapacityEvent(this->name_, n);
+    return n;
+  }
+
+  std::vector<std::pair<K, V>> GetTableOffline(const bool clear_table = false) override {
+    const auto r = Super::GetTableOffline(clear_table);
+
+    // Synthesize a "get table offline" recording by:
+    // 1. Recording how many key/value pairs were returned.
+    // 2. Recording each key/value pair individually.
+    recorder_.RecordBPFMapGetTableOfflineEvent(this->name_, r.size());
+    for (const auto& [key, value] : r) {
+      recorder_.RecordBPFMapGetValueEvent(this->name_, sizeof(key), &key, sizeof(value), &value);
+    }
+    return r;
+  }
+
+  RecordingWrappedBCCMapImpl(bpf_tools::BCCWrapper* bcc, const std::string& name)
+      : WrappedBCCMapImpl<K, V, kUserSpaceManaged>(bcc, name),
+        recorder_(*bcc->GetBPFRecorder().ConsumeValueOrDie()) {}
+
+ private:
+  BPFRecorder& recorder_;
+};
+
+template <typename K, typename V, bool kUserSpaceManaged = false>
+class ReplayingWrappedBCCMapImpl : public WrappedBCCMap<K, V, kUserSpaceManaged> {
+ public:
+  StatusOr<V> GetValue(const K& k) const override {
+    V v;
+    PX_RETURN_IF_ERROR(replayer_.ReplayMapGetValue(this->name_, sizeof(K), &k, sizeof(V), &v));
+    return v;
+  }
+
+  Status SetValue(const K&, const V&) override { return Status::OK(); }
+  Status RemoveValue(const K&) override { return Status::OK(); }
+
+  std::vector<std::pair<K, V>> GetTableOffline(const bool) override {
+    std::vector<std::pair<K, V>> r;
+    auto status_or_size = replayer_.ReplayBPFMapGetTableOfflineEvent(name_);
+    if (!status_or_size.ok()) {
+      return r;
+    }
+    const int n = status_or_size.ConsumeValueOrDie();
+    PX_UNUSED(n);
+    for (int i = 0; i < n; ++i) {
+      K k;
+      V v;
+      auto s = replayer_.ReplayMapGetKeyAndValue(this->name_, sizeof(K), &k, sizeof(V), &v);
+      if (!s.ok()) {
+        return r;
+      }
+      r.push_back({k, v});
+    }
+    return r;
+  }
+
+  size_t capacity() const override {
+    return replayer_.ReplayBPFMapCapacityEvent(name_).ConsumeValueOr(0);
+  }
+
+  ReplayingWrappedBCCMapImpl(BCCWrapper* bcc, const std::string& name)
+      : name_(name), replayer_(*bcc->GetBPFReplayer().ConsumeValueOrDie()) {}
+
+ private:
+  const std::string name_;
+  BPFReplayer& replayer_;
+};
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // Per CPU Array Table
@@ -687,8 +807,14 @@ class WrappedBCCStackTableImpl : public WrappedBCCStackTable {
 //   return std::make_unique<ImplT>(bcc, name);
 // }
 
-template <typename BaseT, typename ImplT>
+template <typename BaseT, typename ImplT, typename RecordingT, typename ReplayingT>
 std::unique_ptr<BaseT> CreateBCCWrappedMapOrArray(BCCWrapper* bcc, const std::string& name) {
+  if (bcc->IsRecording()) {
+    return std::make_unique<RecordingT>(bcc, name);
+  }
+  if (bcc->IsReplaying()) {
+    return std::make_unique<ReplayingT>(bcc, name);
+  }
   return std::make_unique<ImplT>(bcc, name);
 }
 
@@ -697,23 +823,29 @@ std::unique_ptr<WrappedBCCArrayTable<T>> WrappedBCCArrayTable<T>::Create(BCCWrap
                                                                          const std::string& name) {
   using BaseT = WrappedBCCArrayTable<T>;
   using ImplT = WrappedBCCArrayTableImpl<T>;
-  return CreateBCCWrappedMapOrArray<BaseT, ImplT>(bcc, name);
+  using RecordingT = RecordingWrappedBCCArrayTableImpl<T>;
+  using ReplayingT = ReplayingWrappedBCCArrayTableImpl<T>;
+  return CreateBCCWrappedMapOrArray<BaseT, ImplT, RecordingT, ReplayingT>(bcc, name);
 }
 
 template <typename K, typename V, bool U>
 std::unique_ptr<WrappedBCCMap<K, V, U>> WrappedBCCMap<K, V, U>::Create(BCCWrapper* bcc,
                                                                        const std::string& name) {
   using BaseT = WrappedBCCMap<K, V, U>;
   using ImplT = WrappedBCCMapImpl<K, V, U>;
-  return CreateBCCWrappedMapOrArray<BaseT, ImplT>(bcc, name);
+  using RecordingT = RecordingWrappedBCCMapImpl<K, V, U>;
+  using ReplayingT = ReplayingWrappedBCCMapImpl<K, V, U>;
+  return CreateBCCWrappedMapOrArray<BaseT, ImplT, RecordingT, ReplayingT>(bcc, name);
 }
 
 template <typename T>
 std::unique_ptr<WrappedBCCPerCPUArrayTable<T>> WrappedBCCPerCPUArrayTable<T>::Create(
     BCCWrapper* bcc, const std::string& name) {
   using BaseT = WrappedBCCPerCPUArrayTable<T>;
   using ImplT = WrappedBCCPerCPUArrayTableImpl<T>;
-  return CreateBCCWrappedMapOrArray<BaseT, ImplT>(bcc, name);
+
+  // TODO(jps): Impl. rr for per cpu array.
+  return CreateBCCWrappedMapOrArray<BaseT, ImplT, ImplT, ImplT>(bcc, name);
 }
 
 }  // namespace bpf_tools

diff --git a/src/stirling/bpf_tools/rr/rr.cc b/src/stirling/bpf_tools/rr/rr.cc
@@ -24,6 +24,52 @@ namespace px {
 namespace stirling {
 namespace bpf_tools {
 
+void BPFRecorder::RecordBPFArrayTableGetValueEvent(const std::string& name, const int32_t idx,
+                                                   const uint32_t data_size,
+                                                   void const* const data) {
+  auto event = events_proto_.add_event()->mutable_array_table_get_value_event();
+  auto event_name = event->mutable_name();
+  auto event_data = event->mutable_data();
+
+  const std::string data_as_string(reinterpret_cast<const char*>(data), data_size);
+
+  event->set_idx(idx);
+  *event_name = name;
+  *event_data = data_as_string;
+}
+
+void BPFRecorder::RecordBPFMapGetValueEvent(const std::string& name, const uint32_t key_size,
+                                            void const* const key, const uint32_t val_size,
+                                            void const* const val) {
+  auto event = events_proto_.add_event()->mutable_map_get_value_event();
+  auto event_name = event->mutable_name();
+  auto event_key = event->mutable_key();
+  auto event_val = event->mutable_value();
+
+  const std::string key_as_string(reinterpret_cast<const char*>(key), key_size);
+  const std::string val_as_string(reinterpret_cast<const char*>(val), val_size);
+
+  *event_name = name;
+  *event_key = key_as_string;
+  *event_val = val_as_string;
+}
+
+void BPFRecorder::RecordBPFMapGetTableOfflineEvent(const std::string& name, const uint32_t size) {
+  auto event = events_proto_.add_event()->mutable_map_get_table_offline_event();
+  auto event_name = event->mutable_name();
+
+  event->set_size(size);
+  *event_name = name;
+}
+
+void BPFRecorder::RecordBPFMapCapacityEvent(const std::string& name, const int32_t n) {
+  auto event = events_proto_.add_event()->mutable_map_capacity_event();
+  auto event_name = event->mutable_name();
+
+  event->set_capacity(n);
+  *event_name = name;
+}
+
 void BPFRecorder::WriteProto(const std::string& proto_buf_file_path) {
   if (!recording_written_) {
     LOG(INFO) << "Writing BPF events pb to file: " << proto_buf_file_path;
@@ -99,6 +145,147 @@ void BPFReplayer::ReplayPerfBufferEvents(const PerfBufferSpec& perf_buffer_spec)
   }
 }
 
+Status BPFReplayer::ReplayArrayGetValue(const std::string& name, const int32_t idx,
+                                        const uint32_t data_size, void* value) {
+  if (PlabackComplete()) {
+    return error::Internal("Playback complete.");
+  }
+
+  const auto event_wrapper = events_proto_.event(playback_event_idx_);
+  if (!event_wrapper.has_array_table_get_value_event()) {
+    return error::Internal("Array table event not available.");
+  }
+
+  const auto event = event_wrapper.array_table_get_value_event();
+
+  if (name != event.name()) {
+    const char* const msg = "Mismatched eBPF array name. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.name(), name));
+  }
+  if (idx != event.idx()) {
+    const char* const msg = "Mismatched array index. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.idx(), idx));
+  }
+  if (data_size != event.data().size()) {
+    const char* const msg = "Mismatched data size. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.data().size(), data_size));
+  }
+
+  memcpy(value, event.data().data(), data_size);
+  ++playback_event_idx_;
+
+  return Status::OK();
+}
+
+Status BPFReplayer::ReplayMapGetValue(const std::string& name, const uint32_t key_size,
+                                      void const* const key, const uint32_t val_size, void* value) {
+  if (PlabackComplete()) {
+    return error::Internal("Playback complete.");
+  }
+
+  const auto event_wrapper = events_proto_.event(playback_event_idx_);
+  if (!event_wrapper.has_map_get_value_event()) {
+    return error::Internal("Map event not available.");
+  }
+
+  const auto event = event_wrapper.map_get_value_event();
+
+  if (name != event.name()) {
+    const char* const msg = "Mismatched eBPF map name. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.name(), name));
+  }
+  if (key_size != event.key().size()) {
+    const char* const msg = "Mismatched key size. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.key().size(), key_size));
+  }
+  if (val_size != event.value().size()) {
+    const char* const msg = "Mismatched value size. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.value().size(), val_size));
+  }
+  if (0 != memcmp(key, event.key().data(), key_size)) {
+    return error::Internal("Mismatched key.");
+  }
+
+  memcpy(value, event.value().data(), val_size);
+  ++playback_event_idx_;
+
+  return Status::OK();
+}
+
+Status BPFReplayer::ReplayMapGetKeyAndValue(const std::string& name, const uint32_t key_size,
+                                            void* key, const uint32_t val_size, void* val) {
+  if (PlabackComplete()) {
+    return error::Internal("Playback complete.");
+  }
+
+  const auto event_wrapper = events_proto_.event(playback_event_idx_);
+  if (!event_wrapper.has_map_get_value_event()) {
+    return error::Internal("Map event not available.");
+  }
+
+  const auto event = event_wrapper.map_get_value_event();
+
+  if (name != event.name()) {
+    const char* const msg = "Mismatched eBPF map name. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.name(), name));
+  }
+  if (key_size != event.key().size()) {
+    const char* const msg = "Mismatched key size. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.key().size(), key_size));
+  }
+  if (val_size != event.value().size()) {
+    const char* const msg = "Mismatched value size. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.value().size(), val_size));
+  }
+  memcpy(key, event.key().data(), key_size);
+  memcpy(val, event.value().data(), val_size);
+  ++playback_event_idx_;
+
+  return Status::OK();
+}
+
+StatusOr<int32_t> BPFReplayer::ReplayBPFMapCapacityEvent(const std::string& name) {
+  if (PlabackComplete()) {
+    return error::Internal("Playback complete.");
+  }
+
+  const auto event_wrapper = events_proto_.event(playback_event_idx_);
+  if (!event_wrapper.has_map_capacity_event()) {
+    return error::Internal("Map event not available.");
+  }
+
+  const auto event = event_wrapper.map_capacity_event();
+
+  if (name != event.name()) {
+    const char* const msg = "Mismatched eBPF map name. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.name(), name));
+  }
+  ++playback_event_idx_;
+
+  return event.capacity();
+}
+
+StatusOr<int32_t> BPFReplayer::ReplayBPFMapGetTableOfflineEvent(const std::string& name) {
+  if (PlabackComplete()) {
+    return error::Internal("Playback complete.");
+  }
+
+  const auto event_wrapper = events_proto_.event(playback_event_idx_);
+  if (!event_wrapper.has_map_get_table_offline_event()) {
+    return error::Internal("Map event not available.");
+  }
+
+  const auto event = event_wrapper.map_get_table_offline_event();
+
+  if (name != event.name()) {
+    const char* const msg = "Mismatched eBPF map name. Expected: $0, requested: $1.";
+    return error::Internal(absl::Substitute(msg, event.name(), name));
+  }
+  ++playback_event_idx_;
+
+  return event.size();
+}
+
 Status BPFReplayer::OpenReplayProtobuf(const std::string& replay_events_pb_file_path) {
   LOG(INFO) << absl::Substitute("replay_events_pb_file_path: $0.", replay_events_pb_file_path);
   std::fstream input(replay_events_pb_file_path, std::ios::in | std::ios::binary);