diff --git a/tpu_sync/telemetry/BUILD b/tpu_sync/telemetry/BUILD index d3a57d89..ad64a480 100644 --- a/tpu_sync/telemetry/BUILD +++ b/tpu_sync/telemetry/BUILD @@ -132,3 +132,29 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "base_shm_exporter", + srcs = ["base_shm_exporter.cc"], + hdrs = ["base_shm_exporter.h"], + deps = [ + ":metrics_backend", + "//tpu_sync/telemetry/shm:shm_collector", + "//tpu_sync/telemetry/shm:shm_writer", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "base_shm_exporter_test", + srcs = ["base_shm_exporter_test.cc"], + deps = [ + ":base_shm_exporter", + ":metrics_backend", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_sync/telemetry/base_shm_exporter.cc b/tpu_sync/telemetry/base_shm_exporter.cc new file mode 100644 index 00000000..656a3b74 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter.cc @@ -0,0 +1,98 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/base_shm_exporter.h" + +#include +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/log/check.h" +#include "absl/strings/numbers.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_collector.h" +#include "tpu_sync/telemetry/shm/shm_writer.h" + +namespace tpu_raiden::telemetry { + +BaseShmExporter::BaseShmExporter(ExporterOptions options) + : options_(std::move(options)) { + CHECK(options_.local_rank.has_value() && !options_.local_rank->empty()) + << "options.local_rank must be specified and non-empty for " + "BaseShmExporter"; + + int rank_val = -1; + CHECK(absl::SimpleAtoi(*options_.local_rank, &rank_val) && rank_val >= 0) + << "options.local_rank must be a valid non-negative integer for " + "BaseShmExporter, got: '" + << *options_.local_rank << "'"; + + CHECK(options_.shm_dir.has_value() && !options_.shm_dir->empty()) + << "options.shm_dir must be specified and non-empty for BaseShmExporter"; + + while (options_.shm_dir->size() > 1 && options_.shm_dir->back() == '/') { + options_.shm_dir->pop_back(); + } + + shm_writer_ = std::make_unique(ShmWriterOptions{ + .shm_dir = *options_.shm_dir, + .local_rank = *options_.local_rank, + }); + + collector_ = std::make_unique(ShmCollectorOptions{ + .shm_dir = *options_.shm_dir, + }); +} + +BaseShmExporter::~BaseShmExporter() = default; + +void BaseShmExporter::IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const { + if (shm_writer_ != nullptr) { + shm_writer_->IncrementCounter(name, labels, val); + } +} + +void BaseShmExporter::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + if (shm_writer_ != nullptr) { + shm_writer_->SetGauge(name, labels, val); + } +} + +void BaseShmExporter::ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const { + if (shm_writer_ != nullptr) { + shm_writer_->ObserveHistogram(name, labels, val); + } +} + +std::string BaseShmExporter::GetTextSnapshot() const { + // Shared-memory exporters intentionally do not produce in-process text + // snapshots because metrics are gathered out-of-process via memory-mapped + // segments. + return ""; +} + +void BaseShmExporter::CollectMetrics( + absl::flat_hash_map& totals) const { + if (collector_ != nullptr) { + collector_->CollectMetrics(totals); + } +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/base_shm_exporter.h b/tpu_sync/telemetry/base_shm_exporter.h new file mode 100644 index 00000000..015629f7 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter.h @@ -0,0 +1,80 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_ + +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_collector.h" +#include "tpu_sync/telemetry/shm/shm_writer.h" + +namespace tpu_raiden::telemetry { + +// Base class for multi-process shared-memory telemetry exporters. +// Manages an underlying ShmWriter for low-overhead metric publishing and an +// ShmCollector for multi-worker aggregation. +// +// Preconditions: +// Requires valid local_rank (via options.local_rank) and shm_dir (via +// options.shm_dir). Fails fast with CHECK if either is missing or invalid. +// Trailing slashes in shm_dir are normalized while preserving "/". +// +// Thread-safety & Destruction contract: +// Callers must ensure all concurrent metric recording has finished prior to +// exporter destruction. +class BaseShmExporter : public MetricsBackend { + public: + explicit BaseShmExporter(ExporterOptions options); + ~BaseShmExporter() override; + + BaseShmExporter(const BaseShmExporter&) = delete; + BaseShmExporter& operator=(const BaseShmExporter&) = delete; + BaseShmExporter(BaseShmExporter&&) = delete; + BaseShmExporter& operator=(BaseShmExporter&&) = delete; + + // Metric recording methods. Safe for concurrent execution across worker + // threads. Writes directly to the memory-mapped shared segment. + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const override; + void SetGauge(absl::string_view name, LabelSpan labels, + double val) const override; + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const override; + + // Returns an empty string. Shared-memory exporters intentionally do not + // produce in-process text snapshots because metrics are gathered + // out-of-process via memory-mapped segments. + std::string GetTextSnapshot() const override; + + // Scans the shared-memory directory and aggregates metric totals across all + // local worker processes into `totals`. + void CollectMetrics(absl::flat_hash_map& totals) const; + + const ExporterOptions& GetOptions() const { return options_; } + + private: + ExporterOptions options_; + std::unique_ptr shm_writer_; + std::unique_ptr collector_; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_ diff --git a/tpu_sync/telemetry/base_shm_exporter_test.cc b/tpu_sync/telemetry/base_shm_exporter_test.cc new file mode 100644 index 00000000..ed544896 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter_test.cc @@ -0,0 +1,220 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/base_shm_exporter.h" + +#include + +#include // NOLINT(build/c++17) +#include +#include +#include + +#include +#include +#include "absl/container/flat_hash_map.h" +#include "absl/strings/str_cat.h" +#include "tpu_sync/telemetry/metrics_backend.h" + +namespace tpu_raiden::telemetry { +namespace { + +using testing::DoubleEq; +using testing::Pair; + +class BaseShmExporterTest : public testing::Test { + protected: + void SetUp() override { + test_dir_ = + absl::StrCat(testing::TempDir(), "/base_shm_exporter_test_", getpid()); + std::filesystem::remove_all(test_dir_); + std::filesystem::create_directories(test_dir_); + } + + void TearDown() override { std::filesystem::remove_all(test_dir_); } + + std::string test_dir_; +}; + +TEST_F(BaseShmExporterTest, CheckFailsWhenLocalRankMissingOrInvalid) { + ExporterOptions options; + options.shm_dir = test_dir_; + options.local_rank = std::nullopt; + + EXPECT_DEATH(BaseShmExporter exporter(options), "local_rank"); + + options.local_rank = ""; + EXPECT_DEATH(BaseShmExporter exporter(options), "local_rank"); + + options.local_rank = "-1"; + EXPECT_DEATH(BaseShmExporter exporter(options), "local_rank"); + + options.local_rank = "../0"; + EXPECT_DEATH(BaseShmExporter exporter(options), "local_rank"); + + options.local_rank = "abc"; + EXPECT_DEATH(BaseShmExporter exporter(options), "local_rank"); +} + +TEST_F(BaseShmExporterTest, ValidatesAndUsesLocalRank) { + ExporterOptions options; + options.shm_dir = test_dir_; + options.local_rank = "0"; + + BaseShmExporter exporter(options); + EXPECT_EQ(exporter.GetOptions().local_rank, "0"); + EXPECT_EQ(exporter.GetOptions().shm_dir, test_dir_); + + bool found_rank_0 = false; + for (const auto& entry : std::filesystem::directory_iterator(test_dir_)) { + const std::string filename = entry.path().filename().string(); + if (filename.rfind("worker_rank_0_", 0) == 0) { + found_rank_0 = true; + } + } + EXPECT_TRUE(found_rank_0); +} + +TEST_F(BaseShmExporterTest, CheckFailsWhenShmDirMissingOrInvalid) { + ExporterOptions options; + options.local_rank = "0"; + options.shm_dir = std::nullopt; + + EXPECT_DEATH(BaseShmExporter exporter(options), "shm_dir"); + + options.shm_dir = ""; + EXPECT_DEATH(BaseShmExporter exporter(options), "shm_dir"); + + // When neither local_rank nor shm_dir is specified in options, death occurs + // on local_rank first. + ExporterOptions default_options; + EXPECT_DEATH(BaseShmExporter exporter(default_options), "local_rank"); +} + +TEST_F(BaseShmExporterTest, NormalizesShmDir) { + // Normalization of trailing slashes on option and segment creation. + ExporterOptions slash_options; + slash_options.shm_dir = absl::StrCat(test_dir_, "///"); + slash_options.local_rank = "0"; + BaseShmExporter slash_exporter(slash_options); + EXPECT_EQ(slash_exporter.GetOptions().shm_dir, test_dir_); + + bool found_segment = false; + for (const auto& entry : std::filesystem::directory_iterator(test_dir_)) { + const std::string filename = entry.path().filename().string(); + if (filename.rfind("worker_rank_0_", 0) == 0) { + found_segment = true; + } + } + EXPECT_TRUE(found_segment); + + // Preserves root directory "/" when trailing slashes stripped from option. + ExporterOptions root_slash_options; + root_slash_options.shm_dir = "///"; + root_slash_options.local_rank = "0"; + BaseShmExporter root_slash_exporter(root_slash_options); + EXPECT_EQ(root_slash_exporter.GetOptions().shm_dir, "/"); + + // Preserves root directory "/" directly on option. + ExporterOptions root_options; + root_options.shm_dir = "/"; + root_options.local_rank = "0"; + BaseShmExporter root_exporter(root_options); + EXPECT_EQ(root_exporter.GetOptions().shm_dir, "/"); +} + +TEST_F(BaseShmExporterTest, SnapshotExtractionMethodsReturnEmpty) { + ExporterOptions options; + options.shm_dir = test_dir_; + options.local_rank = "0"; + + BaseShmExporter exporter(options); + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + exporter.SetGauge(metric_names::kBufferAllocatedBytes, labels, 2048.0); + exporter.ObserveHistogram(metric_names::kTransferDurationMs, labels, 10.0); + + EXPECT_EQ(exporter.GetTextSnapshot(), ""); + EXPECT_TRUE(exporter.GetAndResetMetricSamples().empty()); +} + +TEST_F(BaseShmExporterTest, RecordsAndAggregatesMetricsAcrossWorkers) { + ExporterOptions options0; + options0.shm_dir = test_dir_; + options0.local_rank = "0"; + BaseShmExporter exporter0(options0); + + auto exporter1 = std::make_unique(ExporterOptions{ + .local_rank = "1", + .shm_dir = test_dir_, + }); + + const MetricLabel push_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter0.IncrementCounter(metric_names::kSentBytesTotal, push_labels, 1024); + exporter0.IncrementCounter(metric_names::kSentBytesTotal, push_labels, 512); + exporter1->IncrementCounter(metric_names::kSentBytesTotal, push_labels, 2000); + + const MetricLabel pull_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}}; + exporter0.SetGauge(metric_names::kBufferAllocatedBytes, pull_labels, 1024.0); + exporter1->SetGauge(metric_names::kBufferAllocatedBytes, pull_labels, 2048.0); + + exporter0.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 12.5); + exporter1->ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 25.0); + + absl::flat_hash_map totals; + exporter0.CollectMetrics(totals); + + // Both ranks are live; metrics (counters, gauges, histograms) are aggregated + // across segments. + EXPECT_THAT( + totals, + testing::AllOf( + testing::Contains( + Pair("sent_bytes_total/direction=push", DoubleEq(3536.0))), + testing::Contains( + Pair("buffer_allocated_bytes/direction=pull", DoubleEq(3072.0))), + testing::Contains( + Pair("transfer_duration_ms/direction=push/count", DoubleEq(2.0))), + testing::Contains( + Pair("transfer_duration_ms/direction=push/sum", DoubleEq(37.5))), + testing::Contains(Pair("transfer_duration_ms/direction=push/bucket_7", + DoubleEq(2.0))))); + + // Terminate rank 1 process (releasing its flock on segment). + exporter1.reset(); + + // Collector reaps the dead worker segment and aggregates only live workers. + exporter0.CollectMetrics(totals); + EXPECT_THAT( + totals, + testing::AllOf( + testing::Contains( + Pair("sent_bytes_total/direction=push", DoubleEq(1536.0))), + testing::Contains( + Pair("buffer_allocated_bytes/direction=pull", DoubleEq(1024.0))), + testing::Contains( + Pair("transfer_duration_ms/direction=push/count", DoubleEq(1.0))), + testing::Contains( + Pair("transfer_duration_ms/direction=push/sum", DoubleEq(12.5))), + testing::Contains(Pair("transfer_duration_ms/direction=push/bucket_7", + DoubleEq(1.0))))); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/metrics_backend.h b/tpu_sync/telemetry/metrics_backend.h index 30f92bef..343105de 100644 --- a/tpu_sync/telemetry/metrics_backend.h +++ b/tpu_sync/telemetry/metrics_backend.h @@ -55,6 +55,9 @@ struct ExporterOptions { // unset (std::nullopt) or empty, telemetry initialization falls back to the // LOCAL_RANK environment variable. std::optional local_rank; + // Base directory for POSIX shared-memory segments for inter-worker telemetry + // aggregation. + std::optional shm_dir; }; // Structure defining centralized metadata for a Raiden metric. diff --git a/tpu_sync/telemetry/shm/BUILD b/tpu_sync/telemetry/shm/BUILD index cab9a9c5..31d61980 100644 --- a/tpu_sync/telemetry/shm/BUILD +++ b/tpu_sync/telemetry/shm/BUILD @@ -34,3 +34,67 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "shm_writer", + srcs = ["shm_writer.cc"], + hdrs = ["shm_writer.h"], + deps = [ + ":shm_layout", + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/hash", + "@com_google_absl//absl/log", + "@com_google_absl//absl/random", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "shm_writer_test", + srcs = ["shm_writer_test.cc"], + deps = [ + ":shm_layout", + ":shm_writer", + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "shm_collector", + srcs = ["shm_collector.cc"], + hdrs = ["shm_collector.h"], + deps = [ + ":shm_layout", + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "shm_collector_test", + srcs = ["shm_collector_test.cc"], + deps = [ + ":shm_collector", + ":shm_layout", + ":shm_writer", + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_sync/telemetry/shm/shm_collector.cc b/tpu_sync/telemetry/shm/shm_collector.cc new file mode 100644 index 00000000..36143d52 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.cc @@ -0,0 +1,238 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/shm/shm_collector.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/cleanup/cleanup.h" +#include "absl/container/flat_hash_map.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { + +ShmCollector::ShmCollector(ShmCollectorOptions options) + : options_(std::move(options)) { + CHECK(!options_.shm_dir.empty()) + << "ShmCollector requires a non-empty shm_dir"; +} + +namespace internal { + +void AggregateSegment(const void* raw_seg, + absl::flat_hash_map& totals) { + if (raw_seg == nullptr || + (reinterpret_cast(raw_seg) % alignof(ShmSegmentLayout) != 0)) { + return; + } + const auto* seg = static_cast(raw_seg); + + if (seg->header.magic.load(std::memory_order_acquire) != kRaidenShmMagic) { + return; + } + if (seg->header.max_toc_entries != kMaxTocEntries) { + return; + } + + const uint32_t data_pool_offset = seg->header.data_pool_offset; + if (data_pool_offset < sizeof(ShmSegmentLayout) || + data_pool_offset >= kSegmentTotalFileSize) { + return; + } + size_t count = seg->header.toc_entry_count.load(std::memory_order_acquire); + size_t num_entries = std::min(count, kMaxTocEntries); + + for (size_t i = 0; i < num_entries; ++i) { + const ShmTocEntry& toc_entry = seg->toc[i]; + if (toc_entry.entry_state.load(std::memory_order_acquire) != + TocEntryState::kCommitted) { + continue; + } + if (toc_entry.type != MetricType::kCounter && + toc_entry.type != MetricType::kGauge && + toc_entry.type != MetricType::kHistogram) { + continue; + } + if (toc_entry.offset < data_pool_offset || + static_cast(toc_entry.offset) + toc_entry.size > + kSegmentTotalFileSize) { + continue; + } + + absl::string_view metric_name( + toc_entry.metric_name, + strnlen(toc_entry.metric_name, sizeof(toc_entry.metric_name))); + if (metric_name.empty()) { + continue; + } + + absl::string_view encoded_labels( + toc_entry.encoded_labels, + strnlen(toc_entry.encoded_labels, sizeof(toc_entry.encoded_labels))); + + const uint8_t* raw = + reinterpret_cast(seg) + toc_entry.offset; + auto get_key = [&] { + return encoded_labels.empty() + ? std::string(metric_name) + : absl::StrCat(metric_name, "/", encoded_labels); + }; + + switch (toc_entry.type) { + case MetricType::kCounter: { + if (toc_entry.size >= sizeof(std::atomic) && + (reinterpret_cast(raw) % + alignof(std::atomic) == + 0)) { + totals[get_key()] += static_cast( + reinterpret_cast*>(raw)->load( + std::memory_order_relaxed)); + } + break; + } + case MetricType::kGauge: { + if (toc_entry.size >= sizeof(std::atomic) && + (reinterpret_cast(raw) % alignof(std::atomic) == + 0)) { + double gauge_value = + reinterpret_cast*>(raw)->load( + std::memory_order_relaxed); + if (std::isfinite(gauge_value)) { + totals[get_key()] += gauge_value; + } + } + break; + } + case MetricType::kHistogram: { + if (toc_entry.size >= sizeof(ShmHistogramSlot) && + (reinterpret_cast(raw) % alignof(ShmHistogramSlot) == + 0)) { + auto* histogram_slot = reinterpret_cast(raw); + double sum = + histogram_slot->sample_sum.load(std::memory_order_relaxed); + if (!std::isfinite(sum)) { + break; + } + std::string key = get_key(); + totals[absl::StrCat(key, "/count")] += static_cast( + histogram_slot->sample_count.load(std::memory_order_relaxed)); + totals[absl::StrCat(key, "/sum")] += sum; + + // TODO: Optimize dynamic string allocation churn in + // hot scrape loop across histogram buckets. + std::string bucket_prefix = absl::StrCat(key, "/bucket_"); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + totals[absl::StrCat(bucket_prefix, b)] += + static_cast(histogram_slot->bucket_counts[b].load( + std::memory_order_relaxed)); + } + } + break; + } + } + } +} + +} // namespace internal + +void ShmCollector::CollectMetrics( + absl::flat_hash_map& totals) const { + totals.clear(); + + DIR* dir = opendir(options_.shm_dir.c_str()); + if (dir == nullptr) { + return; + } + absl::Cleanup close_dir = [dir] { closedir(dir); }; + + struct dirent* entry = nullptr; + while ((entry = readdir(dir)) != nullptr) { + absl::string_view filename(entry->d_name); + if (!absl::StartsWith(filename, kShmFilePrefix) || + !absl::EndsWith(filename, kShmFileExtension)) { + continue; + } + + const std::string path = absl::StrCat(options_.shm_dir, "/", filename); + int fd = open(path.c_str(), O_RDONLY | O_CLOEXEC | O_NOFOLLOW); + if (fd < 0) { + continue; + } + absl::Cleanup close_fd = [fd] { close(fd); }; + + struct stat file_stat = {}; + if (fstat(fd, &file_stat) != 0 || !S_ISREG(file_stat.st_mode)) { + continue; + } + + if (flock(fd, LOCK_EX | LOCK_NB) == 0) { + // Process is dead. Dead worker metrics are intentionally discarded upon + // reaping; only metrics from live processes are aggregated. + // Re-verify fstat under lock: if already unlinked by a concurrent reaper, + // abort early. + if (fstat(fd, &file_stat) != 0 || file_stat.st_nlink == 0) { + continue; + } + // Verify inode before unlinking to prevent TOCTOU. + struct stat current_stat = {}; + if (lstat(path.c_str(), ¤t_stat) == 0 && + current_stat.st_ino == file_stat.st_ino && + current_stat.st_dev == file_stat.st_dev) { + if (unlink(path.c_str()) != 0) { + PLOG(WARNING) + << "Failed to unlink dead worker shared-memory segment at " + << path; + } + } + // Do NOT call flock(LOCK_UN); close_fd will release the lock. + } else if (flock(fd, LOCK_SH | LOCK_NB) == 0) { + // Process is live. Re-verify fstat after acquiring shared lock to prevent + // SIGBUS from newly created, undersized files during worker startup. + if (fstat(fd, &file_stat) != 0 || + file_stat.st_size < static_cast(kSegmentTotalFileSize)) { + continue; + } + void* addr = + mmap(nullptr, kSegmentTotalFileSize, PROT_READ, MAP_SHARED, fd, 0); + if (addr != MAP_FAILED) { + absl::Cleanup unmap_addr = [addr] { + munmap(addr, kSegmentTotalFileSize); + }; + internal::AggregateSegment(addr, totals); + } + } + } +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/shm/shm_collector.h b/tpu_sync/telemetry/shm/shm_collector.h new file mode 100644 index 00000000..c4a34843 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.h @@ -0,0 +1,77 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_COLLECTOR_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_COLLECTOR_H_ + +#include + +#include "absl/container/flat_hash_map.h" + +namespace tpu_raiden::telemetry { + +namespace internal { + +// Aggregates metric entries from a single memory-mapped segment layout into +// `totals`. Exposed for testing only. +// Parameterized as `const void* raw_seg` to accept arbitrary addresses safely +// without invoking undefined behavior at the call boundary. +void AggregateSegment(const void* raw_seg, + absl::flat_hash_map& totals); + +} // namespace internal + +// Configuration options for the shared-memory telemetry collector. +struct ShmCollectorOptions { + // Directory where shared-memory segment files (.mmap) are stored (e.g. + // "/dev/shm" or "/tmp"). + std::string shm_dir; +}; + +// Thread-safe shared-memory telemetry collector. +// Scans for memory-mapped segment files in /dev/shm created by active TPU +// workers, acquires shared reader locks, aggregates counters, gauges, and +// histograms across worker processes, and cleans up dead worker files. +class ShmCollector { + public: + explicit ShmCollector(ShmCollectorOptions options); + + // Scans the configured shared-memory directory, aggregates metric totals + // across all live worker processes into `totals`, and reaps dead worker + // files. Clears and populates `totals`, reusing existing bucket allocation + // across calls. + // + // Dead worker processes (whose exclusive lock can be acquired) are reaped + // by unlinking their segment files. Dead worker metrics are intentionally + // discarded upon reaping; only metrics from live processes are aggregated. + // Downstream consumers should note that cumulative counters from reaped + // workers are not retained. + // + // Gauges are aggregated across workers by addition. For non-additive gauges + // (e.g. utilization or fraction), callers should supply rank-distinguishing + // labels to avoid cross-worker aggregation. + // + // Metric keys in `totals` are formatted as: + // - "metric_name" (unlabeled) or "metric_name/label1=val1;label2=val2" + // - For histograms, suffixes "/count", "/sum", and "/bucket_" are + // appended (e.g. "transfer_duration_ms/bucket_0"). + void CollectMetrics(absl::flat_hash_map& totals) const; + + private: + ShmCollectorOptions options_; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_COLLECTOR_H_ diff --git a/tpu_sync/telemetry/shm/shm_collector_test.cc b/tpu_sync/telemetry/shm/shm_collector_test.cc new file mode 100644 index 00000000..3f44372a --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector_test.cc @@ -0,0 +1,453 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/shm/shm_collector.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include +#include +#include +#include // NOLINT(build/c++11) +#include // NOLINT(build/c++11) +#include +#include + +#include +#include +#include "absl/cleanup/cleanup.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" +#include "tpu_sync/telemetry/shm/shm_writer.h" + +namespace tpu_raiden::telemetry { + +class ShmCollectorTest : public testing::Test { + protected: + void SetUp() override { + test_dir_ = absl::StrCat(testing::TempDir(), "/shm_col_", getpid(), "_", + reinterpret_cast(this)); + std::filesystem::create_directories(test_dir_); + } + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + } + ShmWriterOptions WriterOptions(absl::string_view rank) const { + return {.shm_dir = test_dir_, .local_rank = std::string(rank)}; + } + ShmCollectorOptions CollectorOptions() const { + return {.shm_dir = test_dir_}; + } + + static double MetricValue( + const absl::flat_hash_map& metrics, + absl::string_view key) { + auto it = metrics.find(key); + if (it == metrics.end()) { + return std::numeric_limits::quiet_NaN(); + } + return it->second; + } + + std::string test_dir_; +}; + +namespace { + +using ::testing::DoubleEq; + +TEST_F(ShmCollectorTest, RejectsEmptyShmDir) { + EXPECT_DEATH(ShmCollector(ShmCollectorOptions{.shm_dir = ""}), + "ShmCollector requires a non-empty shm_dir"); +} + +TEST_F(ShmCollectorTest, AggregatesLiveWorkersAndReapsDead) { + // Dead worker segment created and destroyed before collection. + { + ShmWriter dead_writer(WriterOptions("dead")); + MetricLabel l{metric_labels::kDirection, metric_labels::kDirectionPush}; + dead_writer.IncrementCounter(metric_names::kSentBytesTotal, {&l, 1}, 500); + dead_writer.SetGauge(metric_names::kBufferAllocatedBytes, {}, 512.0); + } + + // Active live workers. + ShmWriter w0(WriterOptions("0")), w1(WriterOptions("1")); + MetricLabel push{metric_labels::kDirection, metric_labels::kDirectionPush}; + MetricLabel pull{metric_labels::kDirection, metric_labels::kDirectionPull}; + const std::array fail = { + MetricLabel{metric_labels::kDirection, metric_labels::kDirectionPull}, + MetricLabel{metric_labels::kErrorCode, "DEADLINE_EXCEEDED"}}; + + w0.IncrementCounter(metric_names::kSentBytesTotal, {&push, 1}, 100); + w1.IncrementCounter(metric_names::kSentBytesTotal, {&push, 1}, 200); + w0.IncrementCounter(metric_names::kSentBytesTotal, {&pull, 1}, 50); + w1.IncrementCounter(metric_names::kSentBytesTotal, {&pull, 1}, 75); + w1.IncrementCounter(metric_names::kTransferFailuresTotal, + {fail.data(), fail.size()}, 5); + w0.SetGauge(metric_names::kBufferAllocatedBytes, {}, 1000.0); + w1.SetGauge(metric_names::kBufferAllocatedBytes, {}, 2000.0); + w0.ObserveHistogram(metric_names::kTransferDurationMs, {}, 0.05); + w1.ObserveHistogram(metric_names::kTransferDurationMs, {}, 0.15); + MetricLabel err{metric_labels::kErrorCode, "RESOURCE_EXHAUSTED"}; + const std::array labels = {push, err}; + w1.ObserveHistogram(metric_names::kTransferDurationMs, + {labels.data(), labels.size()}, 0.25); + + ShmCollector initial(CollectorOptions()); + ShmCollector collector(std::move(initial)); + absl::flat_hash_map totals; + collector.CollectMetrics(totals); + + EXPECT_THAT(totals["sent_bytes_total/direction=push"], DoubleEq(300.0)); + EXPECT_THAT(totals["sent_bytes_total/direction=pull"], DoubleEq(125.0)); + EXPECT_THAT(totals["transfer_failures_total/" + "direction=pull;error_code=DEADLINE_EXCEEDED"], + DoubleEq(5.0)); + EXPECT_THAT(totals["buffer_allocated_bytes"], DoubleEq(3000.0)); + EXPECT_THAT(totals["transfer_duration_ms/count"], DoubleEq(2.0)); + EXPECT_THAT(totals["transfer_duration_ms/sum"], DoubleEq(0.20)); + EXPECT_THAT(totals["transfer_duration_ms/bucket_0"], DoubleEq(1.0)); + EXPECT_THAT(totals["transfer_duration_ms/bucket_1"], DoubleEq(1.0)); + std::string labeled_prefix = + "transfer_duration_ms/direction=push;error_code=RESOURCE_EXHAUSTED"; + EXPECT_THAT(totals[absl::StrCat(labeled_prefix, "/count")], DoubleEq(1.0)); + EXPECT_THAT(totals[absl::StrCat(labeled_prefix, "/sum")], DoubleEq(0.25)); + EXPECT_THAT(totals[absl::StrCat(labeled_prefix, "/bucket_1")], DoubleEq(1.0)); + + // Dead worker file is reaped (unlinked) while live worker files remain. + int dead_files = 0; + int live_files = 0; + for (const auto& entry : std::filesystem::directory_iterator(test_dir_)) { + const std::string fn = entry.path().filename().string(); + if (absl::StrContains(fn, "worker_rank_dead_")) { + ++dead_files; + } + if (absl::StrContains(fn, "worker_rank_0_") || + absl::StrContains(fn, "worker_rank_1_")) { + ++live_files; + } + } + EXPECT_EQ(dead_files, 0); + EXPECT_EQ(live_files, 2); +} + +TEST_F(ShmCollectorTest, ConcurrencyAndProcessCrashHandling) { + int pfd[2]; + ASSERT_EQ(pipe(pfd), 0); + pid_t pid = fork(); + ASSERT_GE(pid, 0); + + if (pid == 0) { + prctl(PR_SET_PDEATHSIG, SIGKILL); + close(pfd[0]); + ShmWriter child(WriterOptions("crashed")); + MetricLabel l{metric_labels::kDirection, metric_labels::kDirectionPush}; + child.IncrementCounter(metric_names::kSentBytesTotal, {&l, 1}, 777); + char ready = 'R'; + if (write(pfd[1], &ready, 1) != 1) { + _exit(1); + } + close(pfd[1]); + while (true) pause(); + } + + close(pfd[1]); + char sync = 0; + ASSERT_EQ(read(pfd[0], &sync, 1), 1); + close(pfd[0]); + kill(pid, SIGKILL); + int status = 0; + ASSERT_EQ(waitpid(pid, &status, 0), pid); + + for (int i = 0; i < 3; ++i) { + ShmWriter(WriterOptions(absl::StrCat("dead_", i))) + .IncrementCounter(metric_names::kSentBytesTotal, {}, 100); + } + ShmWriter live(WriterOptions("live")); + live.IncrementCounter(metric_names::kSentBytesTotal, {}, 42); + + // Concurrent collector threads race to aggregate and reap dead files. + std::vector threads; + threads.reserve(8); + for (int i = 0; i < 8; ++i) { + threads.emplace_back([this]() { + absl::flat_hash_map t; + ShmCollector(CollectorOptions()).CollectMetrics(t); + }); + } + for (std::thread& t : threads) { + t.join(); + } + + absl::flat_hash_map final_totals; + ShmCollector(CollectorOptions()).CollectMetrics(final_totals); + EXPECT_THAT(final_totals["sent_bytes_total"], DoubleEq(42.0)); + EXPECT_FALSE(final_totals.contains("sent_bytes_total/direction=push")); + + int dead_files = 0; + for (const auto& e : std::filesystem::directory_iterator(test_dir_)) { + const std::string fn = e.path().filename().string(); + dead_files += + absl::StrContains(fn, "crashed") || absl::StrContains(fn, "dead_"); + } + EXPECT_EQ(dead_files, 0); +} + +TEST_F(ShmCollectorTest, AggregatesMultiChunkWorker) { + ShmWriter writer(WriterOptions("multi")); + for (size_t i = 0; i < kMaxTocEntries + 20; ++i) { + writer.IncrementCounter(absl::StrCat("metric_", i), {}, 1); + } + + absl::flat_hash_map totals; + ShmCollector(CollectorOptions()).CollectMetrics(totals); + + EXPECT_THAT(totals["metric_0"], DoubleEq(1.0)); + EXPECT_THAT(totals[absl::StrCat("metric_", kMaxTocEntries + 19)], + DoubleEq(1.0)); +} + +TEST_F(ShmCollectorTest, HandlesDeadAndLiveTruncatedFiles) { + // Dead truncated file (no lock held) must be reaped. + const std::string dead_trunc_path = absl::StrCat( + test_dir_, "/", kShmFilePrefix, "dead_trunc", kShmFileExtension); + { + std::ofstream ofs(dead_trunc_path, std::ios::binary); + ofs.write("short", 5); + } + + // Live truncated file (shared lock held) must be skipped and preserved. + const std::string live_trunc_path = absl::StrCat( + test_dir_, "/", kShmFilePrefix, "live_trunc", kShmFileExtension); + int fd = open(live_trunc_path.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0644); + ASSERT_GE(fd, 0); + absl::Cleanup close_fd = [fd] { close(fd); }; + ASSERT_EQ(write(fd, "short", 5), 5); + ASSERT_EQ(flock(fd, LOCK_SH | LOCK_NB), 0); + + absl::flat_hash_map totals; + ShmCollector(CollectorOptions()).CollectMetrics(totals); + + EXPECT_TRUE(totals.empty()); + EXPECT_FALSE(std::filesystem::exists(dead_trunc_path)); + EXPECT_TRUE(std::filesystem::exists(live_trunc_path)); +} + +TEST_F(ShmCollectorTest, ResilientToCorruptedFilesAndNonExistentDir) { + // Non-existent directory handling. + absl::flat_hash_map missing_totals; + ShmCollector({.shm_dir = absl::StrCat(test_dir_, "/missing")}) + .CollectMetrics(missing_totals); + EXPECT_TRUE(missing_totals.empty()); + + ShmWriter live_writer(WriterOptions("good")); + MetricLabel push{metric_labels::kDirection, metric_labels::kDirectionPush}; + live_writer.IncrementCounter(metric_names::kSentBytesTotal, {&push, 1}, 500); + + auto write_file = [&](absl::string_view name, absl::string_view data, + size_t pad = 0) { + std::ofstream ofs(absl::StrCat(test_dir_, "/", name), std::ios::binary); + ofs.write(data.data(), data.size()); + if (pad > 0) ofs.write(std::string(pad, '\0').data(), pad); + }; + auto write_hdr = [&](absl::string_view name, uint32_t magic) { + ShmTocHeader h{}; + h.magic.store(magic); + write_file(absl::StrCat(kShmFilePrefix, name), + {reinterpret_cast(&h), sizeof(h)}, + kSegmentTotalFileSize - sizeof(h)); + }; + + write_file(absl::StrCat(kShmFilePrefix, "zero.mmap"), ""); + write_file(absl::StrCat(kShmFilePrefix, "short.mmap"), "short_hdr"); + write_file(absl::StrCat(kShmFilePrefix, "truncated.mmap"), "", + sizeof(ShmSegmentLayout) + 128); + write_hdr("bad_magic.mmap", 0xDEADBEEF); + write_file("other_file.txt", "non-shm content"); + std::filesystem::create_directories( + absl::StrCat(test_dir_, "/", kShmFilePrefix, "dir.mmap")); + ASSERT_EQ( + symlink( + "/dev/null", + absl::StrCat(test_dir_, "/", kShmFilePrefix, "symlink.mmap").c_str()), + 0); + + const std::string unreadable = + absl::StrCat(test_dir_, "/", kShmFilePrefix, "unreadable.mmap"); + write_file(absl::StrCat(kShmFilePrefix, "unreadable.mmap"), "", + kSegmentTotalFileSize); + chmod(unreadable.c_str(), 0000); + absl::Cleanup restore_perms = [&] { chmod(unreadable.c_str(), 0644); }; + + absl::flat_hash_map totals; + ShmCollector(CollectorOptions()).CollectMetrics(totals); + EXPECT_THAT(totals["sent_bytes_total/direction=push"], DoubleEq(500.0)); + EXPECT_TRUE( + std::filesystem::exists(absl::StrCat(test_dir_, "/other_file.txt"))); +} + +TEST_F(ShmCollectorTest, DirectAggregateSegmentEdgeCases) { + absl::flat_hash_map totals; + + internal::AggregateSegment(nullptr, totals); + EXPECT_TRUE(totals.empty()); + + struct alignas(64) Buffer { + uint8_t data[kSegmentTotalFileSize]{}; + }; + auto mem = std::make_unique(); + auto* seg = reinterpret_cast(mem->data); + seg->header.magic.store(kRaidenShmMagic, std::memory_order_release); + seg->header.max_toc_entries = kMaxTocEntries; + seg->header.data_pool_offset = sizeof(ShmSegmentLayout); + + // Unaligned segment pointer (safe to pass as const void*). + const void* unaligned_seg = + reinterpret_cast(reinterpret_cast(seg) + 1); + internal::AggregateSegment(unaligned_seg, totals); + EXPECT_TRUE(totals.empty()); + + // Corrupt magic. + seg->header.magic.store(0xDEADBEEF, std::memory_order_release); + internal::AggregateSegment(seg, totals); + EXPECT_TRUE(totals.empty()); + seg->header.magic.store(kRaidenShmMagic, std::memory_order_release); + + // Corrupt max_toc_entries. + seg->header.max_toc_entries = kMaxTocEntries + 1; + internal::AggregateSegment(seg, totals); + EXPECT_TRUE(totals.empty()); + seg->header.max_toc_entries = kMaxTocEntries; + + // Corrupt data_pool_offset below sizeof(ShmSegmentLayout). + seg->header.data_pool_offset = sizeof(ShmSegmentLayout) - 1; + internal::AggregateSegment(seg, totals); + EXPECT_TRUE(totals.empty()); + + // Corrupt data_pool_offset >= kSegmentTotalFileSize. + seg->header.data_pool_offset = kSegmentTotalFileSize; + internal::AggregateSegment(seg, totals); + EXPECT_TRUE(totals.empty()); + + // Restore valid data_pool_offset. + seg->header.data_pool_offset = sizeof(ShmSegmentLayout); + + const uint32_t pool = sizeof(ShmSegmentLayout); + new (mem->data + pool) std::atomic(42.5); + new (mem->data + pool + kMetricSlotAlignment) std::atomic(100); + auto* histogram_slot = + new (mem->data + pool + 2 * kMetricSlotAlignment) ShmHistogramSlot(); + histogram_slot->sample_count.store(1); + histogram_slot->sample_sum.store(0.5); + histogram_slot->bucket_counts[2].store(1); + + const uint32_t nan_hist_pool = pool + 3 * kMetricSlotAlignment; + auto* nan_hist = new (mem->data + nan_hist_pool) ShmHistogramSlot(); + nan_hist->sample_count.store(10); + nan_hist->sample_sum.store(std::numeric_limits::quiet_NaN()); + nan_hist->bucket_counts[0].store(10); + + const uint32_t inf_hist_pool = pool + 4 * kMetricSlotAlignment; + auto* inf_hist = new (mem->data + inf_hist_pool) ShmHistogramSlot(); + inf_hist->sample_count.store(20); + inf_hist->sample_sum.store(std::numeric_limits::infinity()); + inf_hist->bucket_counts[0].store(20); + + struct EntrySpec { + absl::string_view name; + MetricType type; + uint32_t off; + uint32_t size; + TocEntryState state; + }; + const std::array specs = { + EntrySpec{"uncomm", MetricType::kCounter, pool, 8, + TocEntryState::kWriting}, + EntrySpec{"gauge_m", MetricType::kGauge, pool, 8, + TocEntryState::kCommitted}, + EntrySpec{"counter_m", MetricType::kCounter, pool + kMetricSlotAlignment, + 8, TocEntryState::kCommitted}, + EntrySpec{"hist_m", MetricType::kHistogram, + pool + 2 * kMetricSlotAlignment, sizeof(ShmHistogramSlot), + TocEntryState::kCommitted}, + EntrySpec{"nan_hist", MetricType::kHistogram, nan_hist_pool, + sizeof(ShmHistogramSlot), TocEntryState::kCommitted}, + EntrySpec{"inf_hist", MetricType::kHistogram, inf_hist_pool, + sizeof(ShmHistogramSlot), TocEntryState::kCommitted}, + EntrySpec{"oob", MetricType::kCounter, kSegmentTotalFileSize + 100, 8, + TocEntryState::kCommitted}, + EntrySpec{"overflow", MetricType::kCounter, pool, UINT32_MAX, + TocEntryState::kCommitted}, + EntrySpec{"unaligned", MetricType::kCounter, pool + 3, 8, + TocEntryState::kCommitted}, + EntrySpec{"", MetricType::kCounter, pool + kMetricSlotAlignment, 8, + TocEntryState::kCommitted}, + EntrySpec{"underflow", MetricType::kCounter, + sizeof(ShmSegmentLayout) - 64, 8, TocEntryState::kCommitted}, + EntrySpec{"corrupt_type", static_cast(99), pool, 8, + TocEntryState::kCommitted}, + }; + for (size_t i = 0; i < specs.size(); ++i) { + ShmTocEntry& toc_entry = seg->toc[i]; + snprintf(toc_entry.metric_name, sizeof(toc_entry.metric_name), "%.*s", + static_cast(specs[i].name.size()), specs[i].name.data()); + toc_entry.type = specs[i].type; + toc_entry.offset = specs[i].off; + toc_entry.size = specs[i].size; + toc_entry.entry_state.store(specs[i].state); + } + seg->header.toc_entry_count.store(specs.size()); + + internal::AggregateSegment(seg, totals); + EXPECT_FALSE(totals.contains("uncomm")); + EXPECT_FALSE(totals.contains("oob")); + EXPECT_FALSE(totals.contains("overflow")); + EXPECT_FALSE(totals.contains("unaligned")); + EXPECT_FALSE(totals.contains("")); + EXPECT_FALSE(totals.contains("underflow")); + EXPECT_FALSE(totals.contains("corrupt_type")); + EXPECT_FALSE(totals.contains("nan_hist/count")); + EXPECT_FALSE(totals.contains("nan_hist/sum")); + EXPECT_FALSE(totals.contains("inf_hist/count")); + EXPECT_FALSE(totals.contains("inf_hist/sum")); + EXPECT_THAT(totals["gauge_m"], DoubleEq(42.5)); + EXPECT_THAT(totals["counter_m"], DoubleEq(100.0)); + EXPECT_THAT(totals["hist_m/count"], DoubleEq(1.0)); + EXPECT_THAT(totals["hist_m/sum"], DoubleEq(0.5)); + EXPECT_THAT(totals["hist_m/bucket_2"], DoubleEq(1.0)); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/shm/shm_writer.cc b/tpu_sync/telemetry/shm/shm_writer.cc new file mode 100644 index 00000000..c271f6ec --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_writer.cc @@ -0,0 +1,347 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/shm/shm_writer.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include // NOLINT(build/c++17) +#include +#include // NOLINT(build/c++11) +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/inlined_vector.h" +#include "absl/log/log.h" +#include "absl/random/random.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { + +namespace { + +// Maximum number of label pairs sorted inline on the stack before falling back +// to dynamic heap allocation. 8 labels * 32 bytes = 256 bytes stack overhead. +constexpr size_t kInlineLabelCapacity = 8; + +void AppendEscaped(absl::string_view s, std::string& out) { + for (char c : s) { + if (c == '\\' || c == '=' || c == ';') { + out.push_back('\\'); + } + out.push_back(c); + } +} + +} // namespace + +std::string EncodeLabels(LabelSpan labels) { + if (labels.empty()) return ""; + if (labels.size() == 1) { + std::string encoded; + AppendEscaped(labels[0].key, encoded); + encoded.push_back('='); + AppendEscaped(labels[0].value, encoded); + return encoded; + } + + absl::InlinedVector, + kInlineLabelCapacity> + sorted; + sorted.reserve(labels.size()); + for (const auto& label : labels) { + sorted.emplace_back(label.key, label.value); + } + absl::c_sort(sorted); + + std::string encoded; + bool first = true; + for (const auto& [key, value] : sorted) { + if (!first) { + encoded.push_back(';'); + } + first = false; + AppendEscaped(key, encoded); + encoded.push_back('='); + AppendEscaped(value, encoded); + } + return encoded; +} + +ShmWriter::ShmWriter(const ShmWriterOptions& options) : options_(options) { + if (options_.shm_dir.empty() || options_.local_rank.empty()) { + LOG(WARNING) << "ShmWriter disabled: shm_dir or local_rank is empty"; + return; + } + + std::error_code ec; + std::filesystem::create_directories(options_.shm_dir, ec); + if (ec) { + LOG(ERROR) << "ShmWriter failed to create directory " << options_.shm_dir + << ": " << ec.message(); + return; + } + + absl::BitGen bitgen; + uuid_ = absl::StrFormat("%08x", absl::Uniform(bitgen)); + + absl::MutexLock lock(mutex_); + chunks_.reserve(kMaxChunks); + if (!AllocateNewChunk()) { + LOG(ERROR) << "ShmWriter failed to allocate initial shared-memory chunk in " + << options_.shm_dir; + } +} + +ShmWriter::~ShmWriter() { + absl::MutexLock lock(mutex_); + for (auto& chunk : chunks_) { + if (chunk.segment) { + munmap(chunk.segment, kSegmentTotalFileSize); + chunk.segment = nullptr; + } + if (chunk.fd >= 0) { + close(chunk.fd); + chunk.fd = -1; + } + } + chunks_.clear(); + counter_cache_.clear(); + gauge_cache_.clear(); + histogram_cache_.clear(); +} + +bool ShmWriter::AllocateNewChunk() const { + if (chunks_.size() >= kMaxChunks) { + LOG(ERROR) << "ShmWriter reached maximum chunk limit of " << kMaxChunks + << " chunks. Cannot allocate additional shared-memory chunks."; + return false; + } + + uint32_t chunk_idx = static_cast(chunks_.size()); + std::string path = + absl::StrCat(options_.shm_dir, "/", kShmFilePrefix, options_.local_rank, + "_", uuid_, "_chunk_", chunk_idx, kShmFileExtension); + std::string tmp_path = absl::StrCat(path, ".tmp"); + + int fd = open(tmp_path.c_str(), + O_RDWR | O_CREAT | O_TRUNC | O_NOFOLLOW | O_CLOEXEC, + options_.file_mode); + if (fd < 0) { + PLOG(ERROR) << "ShmWriter failed to open temporary file " << tmp_path; + return false; + } + + if (ftruncate(fd, kSegmentTotalFileSize) != 0) { + PLOG(ERROR) << "ShmWriter failed to ftruncate file " << tmp_path; + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + void* addr = mmap(nullptr, kSegmentTotalFileSize, PROT_READ | PROT_WRITE, + MAP_SHARED, fd, 0); + if (addr == MAP_FAILED) { + PLOG(ERROR) << "ShmWriter failed to mmap file " << tmp_path; + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + auto* segment = reinterpret_cast(addr); + std::memset(segment, 0, kSegmentTotalFileSize); + + ShmTocHeader& header = segment->header; + header.pid = static_cast(getpid()); + header.chunk_index = chunk_idx; + header.toc_entry_count.store(0, std::memory_order_relaxed); + header.max_toc_entries = static_cast(kMaxTocEntries); + header.data_pool_offset = static_cast(sizeof(ShmSegmentLayout)); + header.data_pool_bytes.store(0, std::memory_order_relaxed); + + header.magic.store(kRaidenShmMagic, std::memory_order_release); + + if (flock(fd, LOCK_SH | LOCK_NB) != 0) { + PLOG(ERROR) << "ShmWriter failed to flock file " << tmp_path; + munmap(addr, kSegmentTotalFileSize); + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + if (rename(tmp_path.c_str(), path.c_str()) != 0) { + PLOG(ERROR) << "ShmWriter failed to rename file from " << tmp_path << " to " + << path; + munmap(addr, kSegmentTotalFileSize); + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + chunks_.push_back({fd, segment, path, chunk_idx}); + return true; +} + +bool ShmWriter::ValidateMetric(absl::string_view name, + absl::string_view encoded_labels) const { + if (chunks_.empty()) return false; + + if (name.size() >= sizeof(ShmTocEntry::metric_name)) { + LOG(ERROR) << "Metric name rejected: " << name << " exceeds buffer size " + << sizeof(ShmTocEntry::metric_name); + return false; + } + if (encoded_labels.size() >= sizeof(ShmTocEntry::encoded_labels)) { + LOG(ERROR) << "Metric " << name << " labels rejected: encoded length " + << encoded_labels.size() << " exceeds buffer size " + << sizeof(ShmTocEntry::encoded_labels); + return false; + } + return true; +} + +template +SlotT* ShmWriter::AllocateSlotAndPublishEntry(absl::string_view name, + absl::string_view encoded_labels, + MetricKey key) const { + constexpr uint32_t kAlign = alignof(SlotT) < kMetricSlotAlignment + ? static_cast(kMetricSlotAlignment) + : static_cast(alignof(SlotT)); + constexpr uint32_t kSlotSize = + sizeof(SlotT) < kMetricSlotAlignment + ? static_cast(kMetricSlotAlignment) + : static_cast(sizeof(SlotT)); + Chunk* active_chunk = &chunks_.back(); + ShmTocHeader* header = &active_chunk->segment->header; + uint32_t current_count = + header->toc_entry_count.load(std::memory_order_relaxed); + uint32_t current_bytes = + header->data_pool_bytes.load(std::memory_order_relaxed); + uint32_t aligned_bytes = (current_bytes + kAlign - 1) & ~(kAlign - 1); + + if (current_count >= header->max_toc_entries || + aligned_bytes + kSlotSize > kMaxDataPoolBytes) { + if (!AllocateNewChunk()) return nullptr; + active_chunk = &chunks_.back(); + header = &active_chunk->segment->header; + current_count = header->toc_entry_count.load(std::memory_order_relaxed); + current_bytes = header->data_pool_bytes.load(std::memory_order_relaxed); + aligned_bytes = (current_bytes + kAlign - 1) & ~(kAlign - 1); + } + + uint32_t offset = header->data_pool_offset + aligned_bytes; + uint8_t* raw = reinterpret_cast(active_chunk->segment) + offset; + auto* slot = new (raw) SlotT(); + + ShmTocEntry& entry = active_chunk->segment->toc[current_count]; + entry.entry_state.store(TocEntryState::kWriting, std::memory_order_relaxed); + snprintf(entry.metric_name, sizeof(entry.metric_name), "%.*s", + static_cast(name.size()), name.data()); + snprintf(entry.encoded_labels, sizeof(entry.encoded_labels), "%.*s", + static_cast(encoded_labels.size()), encoded_labels.data()); + entry.type = Type; + entry.offset = offset; + entry.size = kSlotSize; + + entry.entry_state.store(TocEntryState::kCommitted, std::memory_order_release); + header->data_pool_bytes.store(aligned_bytes + kSlotSize, + std::memory_order_relaxed); + header->toc_entry_count.fetch_add(1, std::memory_order_release); + + auto& cache = GetCache(); + cache.emplace(std::move(key), slot); + return slot; +} + +template +SlotT* ShmWriter::GetOrCreateSlot(absl::string_view name, + LabelSpan labels) const { + // TODO: Performance: GetOrCreateSlot calls EncodeLabels(labels) + // unconditionally on every invocation, allocating a dynamic std::string on + // the heap even when the metric stream is already present in the cache. + // Consider formatting the encoded labels into a stack-allocated scratch + // buffer (e.g. char scratch[128] or absl::InlinedVector) so that + // MetricKeyView can perform transparent hash lookups on cache hits with zero + // heap allocations. + std::string encoded_labels = EncodeLabels(labels); + MetricKey key{std::string(name), encoded_labels}; + + // 1. Fast-path: check cache under reader lock. + { + absl::ReaderMutexLock read_lock(mutex_); + const auto& cache = GetCache(); + auto it = cache.find(key); + if (it != cache.end()) return it->second; + } + + // 2. Slow-path: acquire exclusive lock to validate, allocate, and publish. + absl::MutexLock write_lock(mutex_); + auto& cache = GetCache(); + auto it = cache.find(key); + if (it != cache.end()) return it->second; + + if (!ValidateMetric(name, encoded_labels)) return nullptr; + + return AllocateSlotAndPublishEntry(name, encoded_labels, + std::move(key)); +} + +void ShmWriter::IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const { + std::atomic* slot = + GetOrCreateSlot, MetricType::kCounter>(name, + labels); + if (slot) { + slot->fetch_add(val, std::memory_order_relaxed); + } +} + +void ShmWriter::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + std::atomic* slot = + GetOrCreateSlot, MetricType::kGauge>(name, labels); + if (slot && std::isfinite(val)) { + slot->store(val, std::memory_order_relaxed); + } +} + +void ShmWriter::ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const { + ShmHistogramSlot* slot = + GetOrCreateSlot(name, labels); + if (slot) { + slot->Observe(val); + } +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/shm/shm_writer.h b/tpu_sync/telemetry/shm/shm_writer.h new file mode 100644 index 00000000..14dc74ad --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_writer.h @@ -0,0 +1,188 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_WRITER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_WRITER_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_map.h" +#include "absl/hash/hash.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { + +// Configuration options for the shared-memory telemetry writer. +struct ShmWriterOptions { + // Directory where shared-memory segment files (.mmap) are stored (e.g. + // "/dev/shm" or "/tmp"). + std::string shm_dir; + // Worker local rank identifier in distributed multi-rank environments. + std::string local_rank; + // POSIX file creation mode/permissions for segment files. Defaults to 0644 so + // out-of-process metric collectors can read shared-memory segments. + mode_t file_mode = 0644; +}; + +// Thread-safe shared-memory telemetry writer using memory-mapped segment files. +// Uses dynamic on-demand Table of Contents, multi-chunk expansion, and +// reader-lock protected pointer caching. Intended for high-throughput, +// low-overhead metric recording for out-of-process collection and aggregation. +// +// NOTE: ShmWriter is intentionally standalone (not derived from MetricsBackend) +// to avoid virtual dispatch overhead and because shared-memory telemetry does +// not maintain in-process textual snapshots or stateful sample reset buffers. +class ShmWriter { + public: + static constexpr size_t kMaxChunks = 16; + + explicit ShmWriter(const ShmWriterOptions& options = {}); + ~ShmWriter(); + + ShmWriter(const ShmWriter&) = delete; + ShmWriter& operator=(const ShmWriter&) = delete; + ShmWriter(ShmWriter&&) = delete; + ShmWriter& operator=(ShmWriter&&) = delete; + + // Metric recording methods. + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val = 1) const; + void SetGauge(absl::string_view name, LabelSpan labels, double val) const; + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const; + + private: + // Owning key stored in the metric pointer caches (counter_cache_, + // gauge_cache_, histogram_cache_). Allocates owning strings only on cache + // misses when a new metric stream is first registered. + struct MetricKey { + std::string name; + std::string encoded_labels; + + template + friend H AbslHashValue(H h, const MetricKey& k) { + return H::combine(std::move(h), k.name, k.encoded_labels); + } + + bool operator==(const MetricKey& other) const = default; + }; + + // Non-owning view borrowing string pointers and lengths. Used on the hot path + // during steady-state metric recordings to look up cached slot pointers with + // zero dynamic heap allocations. + struct MetricKeyView { + absl::string_view name; + absl::string_view encoded_labels; + + template + friend H AbslHashValue(H h, const MetricKeyView& k) { + return H::combine(std::move(h), k.name, k.encoded_labels); + } + + bool operator==(const MetricKeyView& other) const = default; + }; + + // Transparent hasher enabling heterogeneous lookup in absl::flat_hash_map. + // Allows querying the map with MetricKeyView without constructing an owning + // MetricKey. + struct MetricKeyHash { + using is_transparent = void; + template + size_t operator()(const T& k) const { + return absl::HashOf(k.name, k.encoded_labels); + } + }; + + // Transparent equality comparator enabling heterogeneous lookup between + // MetricKey and MetricKeyView in absl::flat_hash_map. + struct MetricKeyEq { + using is_transparent = void; + template + bool operator()(const T& a, const U& b) const { + return a.name == b.name && a.encoded_labels == b.encoded_labels; + } + }; + + struct Chunk { + int fd = -1; + ShmSegmentLayout* segment = nullptr; + std::string file_path; + uint32_t chunk_index = 0; + }; + + bool AllocateNewChunk() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + template + auto& GetCache() const ABSL_SHARED_LOCKS_REQUIRED(mutex_) { + if constexpr (std::is_same_v>) { + return counter_cache_; + } else if constexpr (std::is_same_v>) { + return gauge_cache_; + } else if constexpr (std::is_same_v) { + return histogram_cache_; + } + } + + bool ValidateMetric(absl::string_view name, + absl::string_view encoded_labels) const + ABSL_SHARED_LOCKS_REQUIRED(mutex_); + + template + SlotT* AllocateSlotAndPublishEntry(absl::string_view name, + absl::string_view encoded_labels, + MetricKey key) const + ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + template + SlotT* GetOrCreateSlot(absl::string_view name, LabelSpan labels) const; + + ShmWriterOptions options_; + std::string uuid_; + + mutable absl::Mutex mutex_; + mutable std::vector chunks_ ABSL_GUARDED_BY(mutex_); + mutable absl::flat_hash_map*, MetricKeyHash, + MetricKeyEq> + counter_cache_ ABSL_GUARDED_BY(mutex_); + mutable absl::flat_hash_map*, MetricKeyHash, + MetricKeyEq> + gauge_cache_ ABSL_GUARDED_BY(mutex_); + mutable absl::flat_hash_map + histogram_cache_ ABSL_GUARDED_BY(mutex_); +}; + +// Encodes a span of metric labels into a canonical semicolon-delimited string +// representation (e.g., "key1=val1;key2=val2"). +// +// Labels are sorted lexicographically by key to guarantee deterministic output +// regardless of caller-specified label order. Delimiter characters ('\', '=', +// ';') in keys and values are escaped with backslashes (e.g. '\=', '\;', '\\') +// to prevent collision and enable unambiguous round-trip parsing by collectors. +std::string EncodeLabels(LabelSpan labels); + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_WRITER_H_ diff --git a/tpu_sync/telemetry/shm/shm_writer_test.cc b/tpu_sync/telemetry/shm/shm_writer_test.cc new file mode 100644 index 00000000..5d44c073 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_writer_test.cc @@ -0,0 +1,794 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/shm/shm_writer.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include // NOLINT(build/c++17) +#include +#include +#include +#include // NOLINT(build/c++11) +#include // NOLINT(build/c++11) +#include +#include + +#include +#include +#include "absl/strings/match.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_replace.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::DoubleEq; +using ::testing::EndsWith; +using ::testing::HasSubstr; +using ::testing::IsNull; +using ::testing::NotNull; + +inline constexpr mode_t kDefaultFileMode = + S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH; // 0644 +inline constexpr mode_t kRestrictedFileMode = S_IRUSR | S_IWUSR; // 0600 + +static_assert(!std::is_copy_constructible_v && + !std::is_copy_assignable_v && + !std::is_move_constructible_v && + !std::is_move_assignable_v); + +// Streamlined RAII helper to discover, mmap, and inspect ShmWriter segments. +class MappedSegment { + public: + MappedSegment() = default; + explicit MappedSegment(absl::string_view path, absl::string_view rank = "0", + uint32_t chunk_idx = 0) { + std::error_code ec; + if (!std::filesystem::is_directory(path, ec) || ec) { + return; + } + std::string file_prefix = absl::StrCat(kShmFilePrefix, rank, "_"); + std::string file_suffix = + absl::StrCat("_chunk_", chunk_idx, kShmFileExtension); + auto dir_it = std::filesystem::directory_iterator(path, ec); + if (ec) { + return; + } + for (const std::filesystem::directory_entry& entry : dir_it) { + if (ec) { + break; + } + std::string filename = entry.path().filename().string(); + if (absl::StartsWith(filename, file_prefix) && + absl::EndsWith(filename, file_suffix)) { + file_path_ = entry.path().string(); + break; + } + } + if (file_path_.empty()) { + return; + } + fd_ = open(file_path_.c_str(), O_RDONLY | O_CLOEXEC); + struct stat file_stat; + if (fd_ >= 0 && fstat(fd_, &file_stat) == 0 && + file_stat.st_size >= kSegmentTotalFileSize) { + void* addr = + mmap(nullptr, kSegmentTotalFileSize, PROT_READ, MAP_SHARED, fd_, 0); + if (addr != MAP_FAILED) { + segment_ = static_cast(addr); + } + } + } + ~MappedSegment() { + if (segment_ != nullptr) { + munmap(const_cast(segment_), kSegmentTotalFileSize); + } + if (fd_ >= 0) { + close(fd_); + } + } + MappedSegment(const MappedSegment&) = delete; + MappedSegment& operator=(const MappedSegment&) = delete; + + bool is_valid() const { + return segment_ && segment_->header.magic.load(std::memory_order_acquire) == + kRaidenShmMagic; + } + const ShmSegmentLayout* segment() const { return segment_; } + const std::string& file_path() const { return file_path_; } + int fd() const { return fd_; } + + static absl::string_view BoundedString(const char* buf, size_t max_len) { + return absl::string_view(buf, strnlen(buf, max_len)); + } + + const ShmTocEntry* FindToc( + absl::string_view name, LabelSpan labels = {}, + std::optional type = std::nullopt) const { + if (!is_valid()) { + return nullptr; + } + std::string encoded = EncodeLabels(labels); + uint32_t entry_count = std::min( + segment_->header.toc_entry_count.load(std::memory_order_acquire), + static_cast(kMaxTocEntries)); + for (uint32_t i = 0; i < entry_count; ++i) { + const ShmTocEntry& entry = segment_->toc[i]; + if (entry.entry_state.load(std::memory_order_acquire) == + TocEntryState::kCommitted && + (!type || entry.type == *type) && + BoundedString(entry.metric_name, sizeof(entry.metric_name)) == name && + BoundedString(entry.encoded_labels, sizeof(entry.encoded_labels)) == + encoded) { + return &entry; + } + } + return nullptr; + } + + template + const T* ReadSlot(absl::string_view name, LabelSpan labels = {}) const { + const ShmTocEntry* entry = FindToc(name, labels, ExpectedType); + if (entry == nullptr) { + return nullptr; + } + if (entry->offset % alignof(T) != 0 || + entry->offset < sizeof(ShmSegmentLayout) || + entry->offset + sizeof(T) > kSegmentTotalFileSize || + entry->size < sizeof(T)) { + return nullptr; + } + return reinterpret_cast( + reinterpret_cast(segment_) + entry->offset); + } + + uint64_t ReadCounter(absl::string_view name, LabelSpan labels = {}) const { + const std::atomic* slot = + ReadSlot, MetricType::kCounter>(name, labels); + return slot != nullptr ? slot->load(std::memory_order_relaxed) : 0; + } + double ReadGauge(absl::string_view name, LabelSpan labels = {}) const { + const std::atomic* slot = + ReadSlot, MetricType::kGauge>(name, labels); + return slot != nullptr ? slot->load(std::memory_order_relaxed) : 0.0; + } + const ShmHistogramSlot* ReadHistogram(absl::string_view name, + LabelSpan labels = {}) const { + return ReadSlot(name, labels); + } + + private: + int fd_ = -1; + const ShmSegmentLayout* segment_ = nullptr; + std::string file_path_; +}; + +class ShmWriterTest : public testing::Test { + protected: + void SetUp() override { + old_umask_ = umask(0022); + test_dir_ = absl::StrCat(testing::TempDir(), "/shm_writer_test_", getpid()); + std::error_code ec; + std::filesystem::create_directories(test_dir_, ec); + } + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + umask(old_umask_); + } + ShmWriterOptions DefaultOptions(absl::string_view rank = "0") const { + return ShmWriterOptions{.shm_dir = test_dir_, + .local_rank = std::string(rank)}; + } + int FilePermissions(const std::string& path) const { + struct stat file_stat; + return stat(path.c_str(), &file_stat) == 0 ? (file_stat.st_mode & 0777) + : -1; + } + std::string test_dir_; + mode_t old_umask_ = 0; +}; + +// 1. Validates header magic, pid, initial TOC limits, contiguity, and +// permissions (0600 vs 0644). +TEST_F(ShmWriterTest, LifecycleAndSegmentLayout) { + { + ShmWriter writer(ShmWriterOptions{.shm_dir = test_dir_, + .local_rank = "c", + .file_mode = kRestrictedFileMode}); + MappedSegment mapped(test_dir_, "c"); + ASSERT_TRUE(mapped.is_valid()); + EXPECT_EQ(FilePermissions(mapped.file_path()), kRestrictedFileMode); + } + + ShmWriter writer(DefaultOptions()); + MappedSegment mapped(test_dir_); + ASSERT_TRUE(mapped.is_valid()); + EXPECT_EQ(FilePermissions(mapped.file_path()), kDefaultFileMode); + EXPECT_THAT(mapped.file_path(), HasSubstr("worker_rank_0_")); + EXPECT_THAT(mapped.file_path(), EndsWith("_chunk_0.mmap")); + + const ShmSegmentLayout* segment = mapped.segment(); + ASSERT_THAT(segment, NotNull()); + EXPECT_EQ(segment->header.magic.load(), kRaidenShmMagic); + EXPECT_EQ(segment->header.pid, getpid()); + EXPECT_EQ(segment->header.max_toc_entries, kMaxTocEntries); + EXPECT_EQ(segment->header.data_pool_offset, sizeof(ShmSegmentLayout)); + EXPECT_EQ(segment->header.toc_entry_count.load(), 0); + EXPECT_EQ(segment->header.data_pool_bytes.load(), 0); + + MetricLabel push_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + MetricLabel pull_label{metric_labels::kDirection, + metric_labels::kDirectionPull}; + writer.IncrementCounter(metric_names::kSentBytesTotal, {&push_label, 1}, 100); + writer.SetGauge(metric_names::kBufferAllocatedBytes, {&pull_label, 1}, 5.0); + writer.ObserveHistogram(metric_names::kTransferDurationMs, {}, 0.05); + + EXPECT_EQ(segment->header.toc_entry_count.load(), 3); + EXPECT_GT(segment->header.data_pool_bytes.load(), 0); + + uint32_t expected_offset = sizeof(ShmSegmentLayout); + for (uint32_t i = 0; i < 3; ++i) { + expected_offset = (expected_offset + kMetricSlotAlignment - 1) & + ~(kMetricSlotAlignment - 1); + EXPECT_EQ(segment->toc[i].offset, expected_offset); + expected_offset += segment->toc[i].size; + } + + for (int i = 0; i < 5; ++i) { + std::string rank = absl::StrCat("r", i); + ShmWriter rank_writer(DefaultOptions(rank)); + MappedSegment rank_mapped(test_dir_, rank); + ASSERT_TRUE(rank_mapped.is_valid()); + rank_writer.IncrementCounter(metric_names::kSentBytesTotal, + {&push_label, 1}, i); + EXPECT_EQ(rank_mapped.ReadCounter(metric_names::kSentBytesTotal, + {&push_label, 1}), + i); + } +} + +// 2. Consolidates counter, gauge, histogram mutations, label sorting, delimiter +// escaping, and multi-type same-name resolution. +TEST_F(ShmWriterTest, MetricRecordingAndLabelEncoding) { + ShmWriter writer(DefaultOptions()); + MappedSegment mapped(test_dir_); + ASSERT_TRUE(mapped.is_valid()); + + MetricLabel custom_label{"k", "v"}; + writer.IncrementCounter("c", {&custom_label, 1}, 100); + writer.SetGauge("g", {&custom_label, 1}, 3.14); + writer.ObserveHistogram("h", {&custom_label, 1}, 1.5); + EXPECT_EQ(mapped.ReadCounter("c", {&custom_label, 1}), 100); + EXPECT_EQ(mapped.ReadCounter("c", {}), 0); + EXPECT_THAT(mapped.ReadGauge("g", {&custom_label, 1}), DoubleEq(3.14)); + EXPECT_THAT(mapped.ReadGauge("g", {}), DoubleEq(0.0)); + const ShmHistogramSlot* hist = mapped.ReadHistogram("h", {&custom_label, 1}); + ASSERT_THAT(hist, NotNull()); + EXPECT_EQ(hist->sample_count.load(), 1); + EXPECT_THAT(hist->sample_sum.load(), DoubleEq(1.5)); + EXPECT_THAT(mapped.ReadHistogram("h", {}), IsNull()); + + MetricLabel wrong_label{"k", "w"}; + EXPECT_EQ(mapped.ReadCounter("c", {&wrong_label, 1}), 0); + + uint32_t toc_count = mapped.segment()->header.toc_entry_count.load(); + writer.IncrementCounter("c", {&custom_label, 1}, 50); + EXPECT_EQ(mapped.ReadCounter("c", {&custom_label, 1}), 150); + EXPECT_EQ(mapped.segment()->header.toc_entry_count.load(), toc_count); + + std::array canonical_labels = { + {{metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "INTERNAL"}}}; + std::array reversed_labels = { + {{metric_labels::kErrorCode, "INTERNAL"}, + {metric_labels::kDirection, metric_labels::kDirectionPull}}}; + writer.IncrementCounter("sorted", canonical_labels, 10); + writer.IncrementCounter("sorted", reversed_labels, 15); + EXPECT_EQ(mapped.ReadCounter("sorted", canonical_labels), 25); + EXPECT_EQ(mapped.ReadCounter("sorted", reversed_labels), 25); + + std::array delimiter_label = {{{"k", "v1;k2=v2"}}}; + std::array multi_labels = {{{"k", "v1"}, {"k2", "v2"}}}; + writer.IncrementCounter("delim", delimiter_label, 10); + writer.IncrementCounter("delim", multi_labels, 20); + EXPECT_EQ(mapped.ReadCounter("delim", delimiter_label), 10); + EXPECT_EQ(mapped.ReadCounter("delim", multi_labels), 20); + + writer.IncrementCounter("same", {}, 10); + writer.SetGauge("same", {}, 2.718); + writer.ObserveHistogram("same", {}, 42.0); + EXPECT_EQ(mapped.ReadCounter("same"), 10); + EXPECT_THAT(mapped.ReadGauge("same"), DoubleEq(2.718)); + const ShmHistogramSlot* same_histogram = mapped.ReadHistogram("same"); + ASSERT_THAT(same_histogram, NotNull()); + EXPECT_THAT(same_histogram->sample_sum.load(), DoubleEq(42.0)); +} + +// 3. Verifies bucket threshold binning, cumulative counts, and boundary +// placement. +TEST_F(ShmWriterTest, HistogramDistributionAndBoundaries) { + ShmWriter writer(DefaultOptions()); + MappedSegment mapped(test_dir_); + ASSERT_TRUE(mapped.is_valid()); + + writer.ObserveHistogram("hb", {}, 0.05); // bucket 0 (< 0.1) + writer.ObserveHistogram("hb", {}, 0.1); // bucket 0 (== 0.1) + writer.ObserveHistogram("hb", {}, 25.0); // bucket 7 (== 25.0) + writer.ObserveHistogram("hb", {}, 100000.0); // overflow bucket + + const ShmHistogramSlot* slot = mapped.ReadHistogram("hb"); + ASSERT_THAT(slot, NotNull()); + EXPECT_EQ(slot->sample_count.load(), 4); + EXPECT_THAT(slot->sample_sum.load(), DoubleEq(100025.15)); + EXPECT_EQ(slot->bucket_counts[0].load(), 2); + EXPECT_EQ(slot->bucket_counts[7].load(), 1); + EXPECT_EQ(slot->bucket_counts[kNumHistogramBuckets].load(), 1); + + for (size_t bucket = 0; bucket <= kNumHistogramBuckets; ++bucket) { + if (bucket != 0 && bucket != 7 && bucket != kNumHistogramBuckets) { + EXPECT_EQ(slot->bucket_counts[bucket].load(), 0); + } + } +} + +// 4. Multi-threaded stress: 32 writer threads + 4 reader threads verifying +// deterministic sums (80k / 400k / 800k / 8M) and barrier integrity. +TEST_F(ShmWriterTest, ConcurrentWriterAndReaderObservationStress) { + ShmWriter writer(DefaultOptions()); + MappedSegment mapped(test_dir_); + ASSERT_TRUE(mapped.is_valid()); + + constexpr int kNumWriterThreads = 32; + constexpr int kNumReaderThreads = 4; + constexpr int kIterationsPerWriter = 10000; + + MetricLabel push_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + MetricLabel pull_label{metric_labels::kDirection, + metric_labels::kDirectionPull}; + std::array fail_labels = { + {{metric_labels::kDirection, metric_labels::kDirectionPullResponse}, + {metric_labels::kErrorCode, "UNAVAILABLE"}}}; + + std::atomic stop{false}; + std::atomic invariant_failures{0}; + std::atomic read_cycles{0}; + std::vector readers; + for (int reader_idx = 0; reader_idx < kNumReaderThreads; ++reader_idx) { + readers.emplace_back([&]() { + const ShmSegmentLayout* segment = mapped.segment(); + while (!stop.load(std::memory_order_relaxed)) { + read_cycles.fetch_add(1, std::memory_order_relaxed); + if (segment->header.magic.load(std::memory_order_acquire) != + kRaidenShmMagic) { + continue; + } + uint32_t entry_count = std::min( + segment->header.toc_entry_count.load(std::memory_order_acquire), + static_cast(kMaxTocEntries)); + for (uint32_t i = 0; i < entry_count; ++i) { + const ShmTocEntry& entry = segment->toc[i]; + if (entry.entry_state.load(std::memory_order_acquire) == + TocEntryState::kCommitted) { + if (entry.type != MetricType::kCounter && + entry.type != MetricType::kGauge && + entry.type != MetricType::kHistogram) { + invariant_failures.fetch_add(1, std::memory_order_relaxed); + } + if (entry.offset < sizeof(ShmSegmentLayout) || + entry.offset + entry.size > kSegmentTotalFileSize || + entry.offset % kMetricSlotAlignment != 0 || + MappedSegment::BoundedString(entry.metric_name, + sizeof(entry.metric_name)) + .empty()) { + invariant_failures.fetch_add(1, std::memory_order_relaxed); + } + if (entry.type == MetricType::kCounter && + entry.offset % alignof(std::atomic) == 0 && + entry.offset + sizeof(std::atomic) <= + kSegmentTotalFileSize) { + const std::atomic* counter_slot = + reinterpret_cast*>( + reinterpret_cast(segment) + entry.offset); + uint64_t value = counter_slot->load(std::memory_order_relaxed); + absl::string_view name = MappedSegment::BoundedString( + entry.metric_name, sizeof(entry.metric_name)); + if (name == metric_names::kSentBytesTotal && value > 80000) { + invariant_failures.fetch_add(1, std::memory_order_relaxed); + } else if (name == metric_names::kTransferFailuresTotal && + (value > 400000 || value % 5 != 0)) { + invariant_failures.fetch_add(1, std::memory_order_relaxed); + } else if (name == metric_names::kReceivedBytesTotal && + (value > 800000 || value % 10 != 0)) { + invariant_failures.fetch_add(1, std::memory_order_relaxed); + } else if (name == "dyn_c" && + (value > 8000000 || value % 100 != 0)) { + invariant_failures.fetch_add(1, std::memory_order_relaxed); + } + } + } + } + } + }); + } + + std::vector writers; + for (int thread_idx = 0; thread_idx < kNumWriterThreads; ++thread_idx) { + writers.emplace_back([&, thread_idx]() { + for (int i = 0; i < kIterationsPerWriter; ++i) { + switch ((thread_idx + i) % 4) { + case 0: + writer.IncrementCounter(metric_names::kSentBytesTotal, + {&push_label, 1}, 1); + break; + case 1: + writer.IncrementCounter(metric_names::kTransferFailuresTotal, + fail_labels, 5); + break; + case 2: + writer.IncrementCounter(metric_names::kReceivedBytesTotal, + {&push_label, 1}, 10); + break; + default: + writer.IncrementCounter("dyn_c", {&pull_label, 1}, 100); + writer.SetGauge("dyn_g", {&pull_label, 1}, 3.14); + writer.ObserveHistogram("dyn_h", {&pull_label, 1}, 1.0); + break; + } + } + }); + } + + for (std::thread& writer_thread : writers) { + writer_thread.join(); + } + stop.store(true, std::memory_order_relaxed); + for (std::thread& reader_thread : readers) { + reader_thread.join(); + } + + EXPECT_EQ(invariant_failures.load(), 0); + EXPECT_GT(read_cycles.load(), 1000); + EXPECT_EQ(mapped.ReadCounter(metric_names::kSentBytesTotal, {&push_label, 1}), + 80000); + EXPECT_EQ( + mapped.ReadCounter(metric_names::kTransferFailuresTotal, fail_labels), + 400000); + EXPECT_EQ( + mapped.ReadCounter(metric_names::kReceivedBytesTotal, {&push_label, 1}), + 800000); + EXPECT_EQ(mapped.ReadCounter("dyn_c", {&pull_label, 1}), 8000000); + EXPECT_THAT(mapped.ReadGauge("dyn_g", {&pull_label, 1}), DoubleEq(3.14)); + const ShmHistogramSlot* dyn_hist = + mapped.ReadHistogram("dyn_h", {&pull_label, 1}); + ASSERT_THAT(dyn_hist, NotNull()); + EXPECT_EQ(dyn_hist->sample_count.load(), 80000); + EXPECT_THAT(dyn_hist->sample_sum.load(), DoubleEq(80000.0)); +} + +// 5. Consolidates flock contention, mutation persistence, and release on +// process destruction. +TEST_F(ShmWriterTest, AdvisoryLockingAndCrossProcessLifecycle) { + auto run_child = [](auto&& fn) -> bool { + pid_t pid = fork(); + if (pid < 0) { + return false; + } + if (pid == 0) { + _exit(fn() ? 0 : 1); + } + int child_status = 0; + if (waitpid(pid, &child_status, 0) != pid) { + return false; + } + return WIFEXITED(child_status) && WEXITSTATUS(child_status) == 0; + }; + auto can_lock_exclusive = [](const std::string& file_path) -> bool { + int fd = open(file_path.c_str(), O_RDWR | O_CLOEXEC); + if (fd < 0) { + return false; + } + bool lock_succeeded = flock(fd, LOCK_EX | LOCK_NB) == 0; + if (lock_succeeded) { + flock(fd, LOCK_UN); + } + close(fd); + return lock_succeeded; + }; + + std::string path; + { + ShmWriter writer(DefaultOptions()); + MappedSegment mapped(test_dir_); + ASSERT_TRUE(mapped.is_valid()); + path = mapped.file_path(); + std::error_code ec; + EXPECT_TRUE(std::filesystem::exists(path, ec)); + EXPECT_FALSE(ec); + + auto has_cloexec_descriptor = [&](const std::string& target_path) -> bool { + std::error_code dir_ec; + auto dir_it = + std::filesystem::directory_iterator("/proc/self/fd", dir_ec); + if (dir_ec) { + return false; + } + for (const std::filesystem::directory_entry& entry : dir_it) { + if (dir_ec) { + return false; + } + int fd_number = 0; + std::error_code equiv_ec; + if (std::filesystem::equivalent(entry.path(), target_path, equiv_ec) && + !equiv_ec && + absl::SimpleAtoi(entry.path().filename().string(), &fd_number) && + fd_number != mapped.fd()) { + if ((fcntl(fd_number, F_GETFD) & FD_CLOEXEC) != 0) { + return true; + } + } + } + return false; + }; + EXPECT_TRUE(has_cloexec_descriptor(path)); + + int probe_fd = open(path.c_str(), O_RDWR | O_CLOEXEC); + ASSERT_GE(probe_fd, 0); + auto assert_contention = [&](int fd) { + EXPECT_EQ(flock(fd, LOCK_EX | LOCK_NB), -1); + EXPECT_EQ(errno, EWOULDBLOCK); + EXPECT_EQ(flock(fd, LOCK_SH | LOCK_NB), 0); + flock(fd, LOCK_UN); + }; + assert_contention(probe_fd); + + MetricLabel label{metric_labels::kDirection, metric_labels::kDirectionPush}; + for (int i = 0; i < 100; ++i) { + writer.IncrementCounter(metric_names::kSentBytesTotal, {&label, 1}, 10); + if (i % 25 == 0) { + assert_contention(probe_fd); + } + } + close(probe_fd); + + EXPECT_TRUE(run_child([&]() { + int fd = open(path.c_str(), O_RDWR | O_CLOEXEC); + bool ok = fd >= 0 && flock(fd, LOCK_EX | LOCK_NB) == -1 && + errno == EWOULDBLOCK && flock(fd, LOCK_SH | LOCK_NB) == 0; + if (fd >= 0) { + flock(fd, LOCK_UN); + close(fd); + } + return ok; + })); + } + + std::error_code exists_ec; + EXPECT_TRUE(std::filesystem::exists(path, exists_ec)); + EXPECT_FALSE(exists_ec); + EXPECT_TRUE(can_lock_exclusive(path)); + + EXPECT_TRUE(run_child([&]() { + ShmWriter child_writer(DefaultOptions("child")); + MetricLabel child_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + child_writer.IncrementCounter(metric_names::kSentBytesTotal, + {&child_label, 1}, 777); + return true; + })); + + MappedSegment child_mapped(test_dir_, "child"); + ASSERT_TRUE(child_mapped.is_valid()); + MetricLabel child_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + EXPECT_EQ(child_mapped.ReadCounter(metric_names::kSentBytesTotal, + {&child_label, 1}), + 777); + EXPECT_TRUE(can_lock_exclusive(child_mapped.file_path())); +} + +// 6. Tests TOC slot overflow (1024), 64 KB data pool overflow (342 histograms), +// and 16 chunks limit. +TEST_F(ShmWriterTest, MultiChunkExpansionAndLimitHandling) { + { + ShmWriter writer(DefaultOptions()); + MappedSegment chunk0_mapped(test_dir_); + ASSERT_TRUE(chunk0_mapped.is_valid()); + + constexpr size_t total_metrics = ShmWriter::kMaxChunks * kMaxTocEntries; + for (size_t i = 0; i < total_metrics; ++i) { + writer.IncrementCounter(absl::StrCat("m_", i), {}, i + 1); + } + EXPECT_EQ(chunk0_mapped.ReadCounter("m_0"), 1); + writer.IncrementCounter("m_0", {}, 10); + EXPECT_EQ(chunk0_mapped.ReadCounter("m_0"), 11); + + MappedSegment chunk1_mapped(test_dir_, "0", 1); + ASSERT_TRUE(chunk1_mapped.is_valid()); + EXPECT_EQ(chunk1_mapped.segment()->header.chunk_index, 1); + EXPECT_EQ(chunk1_mapped.ReadCounter("m_1024"), 1025); + + MappedSegment chunk15_mapped(test_dir_, "0", 15); + ASSERT_TRUE(chunk15_mapped.is_valid()); + EXPECT_EQ(chunk15_mapped.segment()->header.chunk_index, 15); + EXPECT_EQ(chunk15_mapped.ReadCounter(absl::StrCat("m_", total_metrics - 1)), + total_metrics); + + writer.IncrementCounter("overflow_m", {}, 10); + writer.SetGauge("overflow_g", {}, 3.14); + writer.ObserveHistogram("overflow_h", {}, 1.0); + EXPECT_EQ(chunk15_mapped.ReadCounter("overflow_m"), 0); + EXPECT_THAT(chunk15_mapped.ReadGauge("overflow_g"), DoubleEq(0.0)); + EXPECT_THAT(chunk15_mapped.ReadHistogram("overflow_h"), IsNull()); + EXPECT_EQ(chunk15_mapped.segment()->header.toc_entry_count.load(), + kMaxTocEntries); + + std::error_code exists_ec; + std::string chunk16_path = + absl::StrReplaceAll(chunk0_mapped.file_path(), + {{absl::StrCat("_chunk_0", kShmFileExtension), + absl::StrCat("_chunk_16", kShmFileExtension)}}); + EXPECT_FALSE(std::filesystem::exists(chunk16_path, exists_ec)); + EXPECT_FALSE(exists_ec); + } + + { + ShmWriter hist_writer(DefaultOptions("hist")); + MappedSegment hist_chunk0(test_dir_, "hist", 0); + ASSERT_TRUE(hist_chunk0.is_valid()); + for (size_t i = 0; i < 342; ++i) { + hist_writer.ObserveHistogram(absl::StrCat("h_", i), {}, + static_cast(i + 1)); + } + EXPECT_EQ(hist_chunk0.segment()->header.toc_entry_count.load(), 341); + EXPECT_LT(hist_chunk0.segment()->header.toc_entry_count.load(), + kMaxTocEntries); + + MappedSegment hist_chunk1(test_dir_, "hist", 1); + ASSERT_TRUE(hist_chunk1.is_valid()); + EXPECT_EQ(hist_chunk1.segment()->header.chunk_index, 1); + EXPECT_EQ(hist_chunk1.segment()->header.toc_entry_count.load(), 1); + + const ShmHistogramSlot* h0_slot = hist_chunk0.ReadHistogram("h_0"); + ASSERT_THAT(h0_slot, NotNull()); + EXPECT_THAT(h0_slot->sample_sum.load(), DoubleEq(1.0)); + + const ShmHistogramSlot* h341_slot = hist_chunk1.ReadHistogram("h_341"); + ASSERT_THAT(h341_slot, NotNull()); + EXPECT_THAT(h341_slot->sample_sum.load(), DoubleEq(342.0)); + } +} + +// 7. Consolidates 63/64-byte name boundaries, 127/128-byte label boundaries, +// NaN/Inf rejection, and invalid dir. +TEST_F(ShmWriterTest, InputValidationAndFailureResilience) { + ShmWriter writer(DefaultOptions()); + MappedSegment mapped(test_dir_); + ASSERT_TRUE(mapped.is_valid()); + + std::array label_1 = {{{"k1", "v1"}}}; + std::array labels_8 = {{{"k1", "v1"}, + {"k2", "v2"}, + {"k3", "v3"}, + {"k4", "v4"}, + {"k5", "v5"}, + {"k6", "v6"}, + {"k7", "v7"}, + {"k8", "v8"}}}; + std::array label_keys; + std::array labels_20; + for (int i = 0; i < 20; ++i) { + label_keys[i] = std::string(1, 'a' + i); + labels_20[i] = {label_keys[i], "1"}; + } + writer.IncrementCounter("c0", {}, 1); + writer.IncrementCounter("c1", label_1, 2); + writer.IncrementCounter("c8", labels_8, 3); + writer.IncrementCounter("c20", labels_20, 4); + EXPECT_EQ(mapped.ReadCounter("c0"), 1); + EXPECT_EQ(mapped.ReadCounter("c1", label_1), 2); + EXPECT_EQ(mapped.ReadCounter("c8", labels_8), 3); + EXPECT_EQ(mapped.ReadCounter("c20", labels_20), 4); + for (uint32_t i = 0; i < 4; ++i) { + const ShmTocEntry& entry = mapped.segment()->toc[i]; + EXPECT_GE(entry.offset, sizeof(ShmSegmentLayout)); + EXPECT_LT(entry.offset + entry.size, kSegmentTotalFileSize); + } + + std::string name_63(63, 'a'); + std::string name_64(64, 'b'); + std::string name_oversized(74, 'c'); + writer.IncrementCounter(name_63, {}, 10); + writer.IncrementCounter(name_64, {}, 20); + writer.IncrementCounter(name_oversized, {}, 30); + EXPECT_EQ(mapped.ReadCounter(name_63), 10); + EXPECT_EQ(mapped.ReadCounter(name_64), 0); + EXPECT_EQ(mapped.ReadCounter(name_oversized), 0); + + std::string val_125(125, 'v'); + std::string val_126(126, 'w'); + std::string val_oversized(148, 'x'); + std::array label_127 = {{{"k", val_125}}}; + std::array label_128 = {{{"k", val_126}}}; + std::array label_oversized = {{{"k", val_oversized}}}; + writer.IncrementCounter("b_lbl", label_127, 30); + writer.IncrementCounter("b_lbl", label_128, 40); + writer.IncrementCounter("b_lbl", label_oversized, 50); + EXPECT_EQ(mapped.ReadCounter("b_lbl", label_127), 30); + EXPECT_EQ(mapped.ReadCounter("b_lbl", label_128), 0); + EXPECT_EQ(mapped.ReadCounter("b_lbl", label_oversized), 0); + + writer.SetGauge("g_nan", {}, 42.0); + for (double val : {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + writer.ObserveHistogram("h_nan", {}, val); + writer.SetGauge("g_nan", {}, val); + } + const ShmHistogramSlot* nan_histogram = mapped.ReadHistogram("h_nan"); + ASSERT_THAT(nan_histogram, NotNull()); + EXPECT_EQ(nan_histogram->sample_count.load(), 0); + EXPECT_THAT(mapped.ReadGauge("g_nan"), DoubleEq(42.0)); + + writer.ObserveHistogram("h_nan", {}, 10.5); + writer.ObserveHistogram("h_nan", {}, + std::numeric_limits::denorm_min()); + EXPECT_EQ(nan_histogram->sample_count.load(), 2); + EXPECT_THAT(nan_histogram->sample_sum.load(), DoubleEq(10.5)); + + std::string regular_file = absl::StrCat(test_dir_, "/regular_file"); + int fd = open(regular_file.c_str(), O_CREAT | O_WRONLY, kDefaultFileMode); + ASSERT_GE(fd, 0); + close(fd); + + MetricLabel dummy_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + for (const ShmWriterOptions& opt : + {ShmWriterOptions{.shm_dir = "", .local_rank = "0"}, + ShmWriterOptions{.shm_dir = test_dir_, .local_rank = ""}, + ShmWriterOptions{.shm_dir = absl::StrCat(regular_file, "/x"), + .local_rank = "0"}}) { + ShmWriter bad_writer(opt); + bad_writer.IncrementCounter(metric_names::kSentBytesTotal, + {&dummy_label, 1}, 10); + bad_writer.SetGauge("any_g", {&dummy_label, 1}, 1.0); + bad_writer.ObserveHistogram("any_h", {&dummy_label, 1}, 2.0); + } + + EXPECT_FALSE(MappedSegment("").is_valid()); + EXPECT_FALSE(MappedSegment(test_dir_, "").is_valid()); + EXPECT_FALSE(MappedSegment(regular_file).is_valid()); +} + +} // namespace +} // namespace tpu_raiden::telemetry