diff --git a/src/stirling/bpf_tools/bcc_wrapper.cc b/src/stirling/bpf_tools/bcc_wrapper.cc index c1b60140c8c..e063cbaf9b5 100644 --- a/src/stirling/bpf_tools/bcc_wrapper.cc +++ b/src/stirling/bpf_tools/bcc_wrapper.cc @@ -485,7 +485,9 @@ std::unique_ptr WrappedBCCStackTable::Create(bpf_tools::BC const std::string& name) { using BaseT = WrappedBCCStackTable; using ImplT = WrappedBCCStackTableImpl; - return CreateBCCWrappedMapOrArray(bcc, name); + + // TODO(jps): Impl. rr for stack table. + return CreateBCCWrappedMapOrArray(bcc, name); } } // namespace bpf_tools diff --git a/src/stirling/bpf_tools/bcc_wrapper.h b/src/stirling/bpf_tools/bcc_wrapper.h index 5a1963a4ee3..7f60a5cd057 100644 --- a/src/stirling/bpf_tools/bcc_wrapper.h +++ b/src/stirling/bpf_tools/bcc_wrapper.h @@ -499,6 +499,44 @@ class WrappedBCCArrayTableImpl : public WrappedBCCArrayTable { std::unique_ptr underlying_; }; +template +class RecordingWrappedBCCArrayTableImpl : public WrappedBCCArrayTableImpl { + public: + using Super = WrappedBCCArrayTableImpl; + + StatusOr 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(bcc, name), + recorder_(*bcc->GetBPFRecorder().ConsumeValueOrDie()) {} + + private: + BPFRecorder& recorder_; +}; + +template +class ReplayingWrappedBCCArrayTableImpl : public WrappedBCCArrayTable { + public: + StatusOr 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 { absl::flat_hash_set shadow_keys_; }; +template +class RecordingWrappedBCCMapImpl : public WrappedBCCMapImpl { + public: + using Super = WrappedBCCMapImpl; + + StatusOr 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> 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(bcc, name), + recorder_(*bcc->GetBPFRecorder().ConsumeValueOrDie()) {} + + private: + BPFRecorder& recorder_; +}; + +template +class ReplayingWrappedBCCMapImpl : public WrappedBCCMap { + public: + StatusOr 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> GetTableOffline(const bool) override { + std::vector> 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(bcc, name); // } -template +template std::unique_ptr CreateBCCWrappedMapOrArray(BCCWrapper* bcc, const std::string& name) { + if (bcc->IsRecording()) { + return std::make_unique(bcc, name); + } + if (bcc->IsReplaying()) { + return std::make_unique(bcc, name); + } return std::make_unique(bcc, name); } @@ -697,7 +823,9 @@ std::unique_ptr> WrappedBCCArrayTable::Create(BCCWrap const std::string& name) { using BaseT = WrappedBCCArrayTable; using ImplT = WrappedBCCArrayTableImpl; - return CreateBCCWrappedMapOrArray(bcc, name); + using RecordingT = RecordingWrappedBCCArrayTableImpl; + using ReplayingT = ReplayingWrappedBCCArrayTableImpl; + return CreateBCCWrappedMapOrArray(bcc, name); } template @@ -705,7 +833,9 @@ std::unique_ptr> WrappedBCCMap::Create(BCCWrappe const std::string& name) { using BaseT = WrappedBCCMap; using ImplT = WrappedBCCMapImpl; - return CreateBCCWrappedMapOrArray(bcc, name); + using RecordingT = RecordingWrappedBCCMapImpl; + using ReplayingT = ReplayingWrappedBCCMapImpl; + return CreateBCCWrappedMapOrArray(bcc, name); } template @@ -713,7 +843,9 @@ std::unique_ptr> WrappedBCCPerCPUArrayTable::Cr BCCWrapper* bcc, const std::string& name) { using BaseT = WrappedBCCPerCPUArrayTable; using ImplT = WrappedBCCPerCPUArrayTableImpl; - return CreateBCCWrappedMapOrArray(bcc, name); + + // TODO(jps): Impl. rr for per cpu array. + return CreateBCCWrappedMapOrArray(bcc, name); } } // namespace bpf_tools diff --git a/src/stirling/bpf_tools/rr/rr.cc b/src/stirling/bpf_tools/rr/rr.cc index 2ea3661675d..e26268717d7 100644 --- 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(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(key), key_size); + const std::string val_as_string(reinterpret_cast(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 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 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); diff --git a/src/stirling/bpf_tools/rr/rr.h b/src/stirling/bpf_tools/rr/rr.h index bf1d0b15ca1..6f774c9c637 100644 --- a/src/stirling/bpf_tools/rr/rr.h +++ b/src/stirling/bpf_tools/rr/rr.h @@ -31,6 +31,13 @@ namespace bpf_tools { class BPFRecorder : public NotCopyMoveable { public: + void RecordBPFArrayTableGetValueEvent(const std::string& name, const int32_t idx, + const uint32_t data_size, void const* const data); + void RecordBPFMapGetValueEvent(const std::string& name, const uint32_t key_size, + void const* const key, const uint32_t data_size, + void const* const value); + void RecordBPFMapGetTableOfflineEvent(const std::string& name, const uint32_t size); + void RecordBPFMapCapacityEvent(const std::string& name, const int32_t n); void RecordPerfBufferEvent(PerfBufferSpec* pb_spec, void const* const data, const int data_size); void WriteProto(const std::string& proto_buf_file_path); @@ -42,6 +49,14 @@ class BPFRecorder : public NotCopyMoveable { class BPFReplayer : public NotCopyMoveable { public: void ReplayPerfBufferEvents(const PerfBufferSpec& perf_buffer_spec); + Status ReplayArrayGetValue(const std::string& name, const int32_t idx, const uint32_t data_size, + void* data); + Status ReplayMapGetValue(const std::string& name, const uint32_t key_size, void const* const key, + const uint32_t val_size, void* value); + Status ReplayMapGetKeyAndValue(const std::string& name, const uint32_t key_size, void* key, + const uint32_t val_size, void* value); + StatusOr ReplayBPFMapCapacityEvent(const std::string& name); + StatusOr ReplayBPFMapGetTableOfflineEvent(const std::string& name); Status OpenReplayProtobuf(const std::string& replay_events_pb_file_path); bool PlabackComplete() const { return playback_event_idx_ >= events_proto_.event_size(); } diff --git a/src/stirling/bpf_tools/rr/rr_bpf_test.cc b/src/stirling/bpf_tools/rr/rr_bpf_test.cc index a51a0120719..bc6cef2c9b8 100644 --- a/src/stirling/bpf_tools/rr/rr_bpf_test.cc +++ b/src/stirling/bpf_tools/rr/rr_bpf_test.cc @@ -173,6 +173,118 @@ TEST_F(BasicRecorderTest, PerfBufferRRTest) { ASSERT_OK(replaying_bcc_->OpenReplayProtobuf(pb_file_name)); replaying_bcc_->PollPerfBuffers(); EXPECT_EQ(test::test_idx, test::gold_data.size()); + + // TODO(jps): add the expectations. +} + +TEST_F(BasicRecorderTest, BPFArrayRRTest) { + auto recording_bpf_array = WrappedBCCArrayTable::Create(recording_bcc_.get(), "results"); + + constexpr uint32_t kLoopIters = 16; + + for (uint32_t i = 0; i < kLoopIters; ++i) { + // Invoking Foo() or Bar() triggers our eBPF uprobe, which will capture the function argument. + PX_UNUSED(::test::Foo(2 * i + 0)); + PX_UNUSED(::test::Bar(2 * i + 1)); + } + for (uint32_t i = 0; i < 2 * kLoopIters; ++i) { + ASSERT_OK_AND_ASSIGN(const int r, recording_bpf_array->GetValue(i)); + EXPECT_EQ(r, i); + } + + constexpr int k100 = 0; + constexpr int k200 = 100; + ASSERT_OK(recording_bpf_array->SetValue(k100, 100)); + ASSERT_OK(recording_bpf_array->SetValue(k200, 200)); + + { + ASSERT_OK_AND_ASSIGN(const int r100, recording_bpf_array->GetValue(k100)); + ASSERT_OK_AND_ASSIGN(const int r200, recording_bpf_array->GetValue(k200)); + EXPECT_EQ(r100, 100); + EXPECT_EQ(r200, 200); + } + + const std::string pb_file_name = "bpf_array_replay_test.pb"; + + // Write out the protobuf file and close the recording BCC wrapper. + recording_bcc_->WriteProto(pb_file_name); + recording_bcc_->Close(); + + // Open the protobuf file in the replaying BCC wrapper. + ASSERT_OK(replaying_bcc_->OpenReplayProtobuf(pb_file_name)); + + // Get a pointer to the replay mode BPF map. + auto replaying_bpf_array = WrappedBCCArrayTable::Create(replaying_bcc_.get(), "results"); + + // Test the first N map get value results. + // These key/val pairs were set by triggering the BPF program. + for (uint32_t i = 0; i < 2 * kLoopIters; ++i) { + ASSERT_OK_AND_ASSIGN(const int r, replaying_bpf_array->GetValue(i)); + EXPECT_EQ(r, i); + } + + // Test the the last two map get value results where we set the value from user space. + { + ASSERT_OK_AND_ASSIGN(const int r100, replaying_bpf_array->GetValue(k100)); + ASSERT_OK_AND_ASSIGN(const int r200, replaying_bpf_array->GetValue(k200)); + EXPECT_EQ(r100, 100); + EXPECT_EQ(r200, 200); + } +} + +TEST_F(BasicRecorderTest, BPFMapRRTest) { + auto recording_bpf_map = WrappedBCCMap::Create(recording_bcc_.get(), "map"); + + constexpr uint32_t kLoopIters = 16; + + for (uint32_t i = 0; i < kLoopIters; ++i) { + // Invoking Foo() or Bar() triggers our eBPF uprobe, which will capture the function argument. + PX_UNUSED(::test::Foo(2 * i + 0)); + PX_UNUSED(::test::Bar(2 * i + 1)); + } + for (uint32_t i = 0; i < 2 * kLoopIters; ++i) { + ASSERT_OK_AND_ASSIGN(const int r, recording_bpf_map->GetValue(i)); + EXPECT_EQ(r, i); + } + + constexpr int k100 = 0; + constexpr int k200 = 100; + ASSERT_OK(recording_bpf_map->SetValue(k100, 100)); + ASSERT_OK(recording_bpf_map->SetValue(k200, 200)); + + { + ASSERT_OK_AND_ASSIGN(const int r100, recording_bpf_map->GetValue(k100)); + ASSERT_OK_AND_ASSIGN(const int r200, recording_bpf_map->GetValue(k200)); + EXPECT_EQ(r100, 100); + EXPECT_EQ(r200, 200); + } + + const std::string pb_file_name = "bpf_map_replay_test.pb"; + + // Write out the protobuf file and close the recording BCC wrapper. + recording_bcc_->WriteProto(pb_file_name); + recording_bcc_->Close(); + + // Open the protobuf file in the replaying BCC wrapper. + ASSERT_OK(replaying_bcc_->OpenReplayProtobuf(pb_file_name)); + + // Get a pointer to the replay mode BPF map. + auto replaying_bpf_map = WrappedBCCMap::Create(replaying_bcc_.get(), "map"); + + // Test the first N map get value results. + // These key/val pairs were set by triggering the BPF program. + for (uint32_t i = 0; i < 2 * kLoopIters; ++i) { + ASSERT_OK_AND_ASSIGN(const int r, replaying_bpf_map->GetValue(i)); + EXPECT_EQ(r, i); + } + + // Test the the last two map get value results where we set the value from user space. + { + ASSERT_OK_AND_ASSIGN(const int r100, recording_bpf_map->GetValue(k100)); + ASSERT_OK_AND_ASSIGN(const int r200, recording_bpf_map->GetValue(k200)); + EXPECT_EQ(r100, 100); + EXPECT_EQ(r200, 200); + } } } // namespace stirling