diff --git a/tpu_sync/telemetry/BUILD b/tpu_sync/telemetry/BUILD index d3a57d89..886700b3 100644 --- a/tpu_sync/telemetry/BUILD +++ b/tpu_sync/telemetry/BUILD @@ -132,3 +132,93 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "base_shm_exporter", + srcs = ["base_shm_exporter.cc"], + hdrs = ["base_shm_exporter.h"], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + 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", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "base_shm_exporter_test", + srcs = ["base_shm_exporter_test.cc"], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + deps = [ + ":base_shm_exporter", + ":metrics_backend", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "prometheus_shm_exporter", + srcs = ["prometheus_shm_exporter.cc"], + hdrs = ["prometheus_shm_exporter.h"], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + deps = [ + ":base_shm_exporter", + ":metrics_backend", + "//tpu_sync/telemetry/shm:shm_collector", + "//tpu_sync/telemetry/shm:shm_layout", + "@com_github_jupp0r_prometheus_cpp//core", + "@com_github_jupp0r_prometheus_cpp//pull", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + ], +) + +cc_test( + name = "prometheus_shm_exporter_test", + srcs = ["prometheus_shm_exporter_test.cc"], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + deps = [ + ":base_shm_exporter", + ":metrics_api", + ":metrics_backend", + ":prometheus_shm_exporter", + ":test_util", + "@com_github_jupp0r_prometheus_cpp//core", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + +cc_test( + name = "prometheus_shm_exporter_challenger_test", + srcs = ["prometheus_shm_exporter_challenger_test.cc"], + copts = ["-fexceptions"], + features = ["-use_header_modules"], + deps = [ + ":base_shm_exporter", + ":metrics_api", + ":metrics_backend", + ":prometheus_shm_exporter", + ":test_util", + "//tpu_sync/telemetry/shm:shm_layout", + "@com_github_jupp0r_prometheus_cpp//core", + "@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..cb750557 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter.cc @@ -0,0 +1,91 @@ +// 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 + +#include "absl/container/flat_hash_map.h" +#include "absl/log/log.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.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(const ExporterOptions& options) + : options_(options) { + std::string rank; + if (options_.local_rank.has_value() && !options_.local_rank->empty()) { + rank = *options_.local_rank; + } else { + const char* env_rank = std::getenv("LOCAL_RANK"); + if (env_rank != nullptr && *env_rank != '\0') { + rank = env_rank; + } + } + + if (rank.empty()) { + LOG(ERROR) << "LOCAL_RANK must be specified for BaseShmExporter (via " + "options.local_rank or LOCAL_RANK environment variable)."; + throw std::invalid_argument( + "LOCAL_RANK must be specified for BaseShmExporter (via " + "options.local_rank or LOCAL_RANK environment variable)."); + } + options_.local_rank = rank; + const std::string dir = options_.GetShmDir(); + + ShmWriterOptions w_opts; + w_opts.shm_dir = dir; + w_opts.local_rank = rank; + shm_writer_ = std::make_unique(w_opts); + + ShmCollectorOptions c_opts; + c_opts.shm_dir = dir; + collector_ = std::make_unique(c_opts); +} + +BaseShmExporter::~BaseShmExporter() { Stop(); } + +void BaseShmExporter::IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const { + if (shm_writer_) shm_writer_->IncrementCounter(name, labels, val); +} + +void BaseShmExporter::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + if (shm_writer_) shm_writer_->SetGauge(name, labels, val); +} + +void BaseShmExporter::ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const { + if (shm_writer_) shm_writer_->ObserveHistogram(name, labels, val); +} + +std::string BaseShmExporter::GetTextSnapshot() const { + return ""; +} + +void BaseShmExporter::CollectMetrics( + absl::flat_hash_map& totals) const { + if (collector_) 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..1aaead46 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter.h @@ -0,0 +1,74 @@ +// 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 + +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.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 { + +// Abstract base class for multi-process shared-memory telemetry exporters. +// Manages underlying ShmWriter for low-overhead metric publishing and +// ShmCollector for multi-worker aggregation. +class BaseShmExporter : public MetricsBackend { + public: + explicit BaseShmExporter(const ExporterOptions& options = {}); + ~BaseShmExporter() override; + + BaseShmExporter(const BaseShmExporter&) = delete; + BaseShmExporter& operator=(const BaseShmExporter&) = delete; + BaseShmExporter(BaseShmExporter&&) = delete; + BaseShmExporter& operator=(BaseShmExporter&&) = delete; + + 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; + + std::string GetTextSnapshot() const override; + void CollectMetrics(absl::flat_hash_map& totals) const; + + virtual void Start() { is_running_.store(true, std::memory_order_release); } + virtual void Stop() { is_running_.store(false, std::memory_order_release); } + virtual bool IsRunning() const { + return is_running_.load(std::memory_order_acquire); + } + + const ExporterOptions& GetOptions() const { return options_; } + const ShmWriter* GetWriter() const { return shm_writer_.get(); } + const ShmCollector* GetCollector() const { return collector_.get(); } + + protected: + ExporterOptions options_; + std::unique_ptr shm_writer_; + std::unique_ptr collector_; + std::atomic is_running_{false}; +}; + +} // 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..8da6711c --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter_test.cc @@ -0,0 +1,191 @@ +// 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 "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::NotNull; + +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 { + unsetenv("LOCAL_RANK"); + unsetenv("SHM_DIR"); + std::filesystem::remove_all(test_dir_); + } + + std::string test_dir_; +}; + +TEST_F(BaseShmExporterTest, ThrowsWhenLocalRankMissing) { + unsetenv("LOCAL_RANK"); + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = std::nullopt; + + EXPECT_THROW(BaseShmExporter exporter(options), std::invalid_argument); + + options.local_rank = ""; + EXPECT_THROW(BaseShmExporter exporter(options), std::invalid_argument); + + setenv("LOCAL_RANK", "", 1); + options.local_rank = std::nullopt; + EXPECT_THROW(BaseShmExporter exporter(options), std::invalid_argument); +} + +TEST_F(BaseShmExporterTest, ResolvesLocalRankFromOptions) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + + BaseShmExporter exporter(options); + EXPECT_THAT(exporter.GetWriter(), NotNull()); + EXPECT_THAT(exporter.GetCollector(), NotNull()); + EXPECT_EQ(exporter.GetOptions().local_rank, "0"); +} + +TEST_F(BaseShmExporterTest, ResolvesLocalRankFromEnvironment) { + setenv("LOCAL_RANK", "3", 1); + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = std::nullopt; + + BaseShmExporter exporter(options); + EXPECT_THAT(exporter.GetWriter(), NotNull()); + EXPECT_THAT(exporter.GetCollector(), NotNull()); + EXPECT_EQ(exporter.GetOptions().local_rank, "3"); +} + +TEST_F(BaseShmExporterTest, ResolvesShmDirFromOption) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + + EXPECT_EQ(options.GetShmDir(), test_dir_); + BaseShmExporter exporter(options); + EXPECT_EQ(exporter.GetOptions().GetShmDir(), test_dir_); +} + +TEST_F(BaseShmExporterTest, ResolvesShmDirFromEnvironment) { + setenv("SHM_DIR", test_dir_.c_str(), 1); + ExporterOptions options; + options.local_rank = "0"; + + EXPECT_EQ(options.GetShmDir(), test_dir_); + BaseShmExporter exporter(options); + EXPECT_EQ(exporter.GetOptions().GetShmDir(), test_dir_); +} + +TEST_F(BaseShmExporterTest, ResolvesShmDirDefaultFallback) { + unsetenv("SHM_DIR"); + ExporterOptions options; + EXPECT_EQ(options.GetShmDir(), "/dev/shm"); + + setenv("SHM_DIR", "", 1); + EXPECT_EQ(options.GetShmDir(), "/dev/shm"); +} + +TEST_F(BaseShmExporterTest, LifecycleStartStopIsRunning) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + + BaseShmExporter exporter(options); + EXPECT_FALSE(exporter.IsRunning()); + + exporter.Start(); + EXPECT_TRUE(exporter.IsRunning()); + + exporter.Stop(); + EXPECT_FALSE(exporter.IsRunning()); +} + +TEST_F(BaseShmExporterTest, GetTextSnapshotReturnsEmpty) { + ExporterOptions options; + options.base_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); + + EXPECT_EQ(exporter.GetTextSnapshot(), ""); +} + +TEST_F(BaseShmExporterTest, GetAndResetMetricSamplesReturnsEmpty) { + ExporterOptions options; + options.base_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); + + EXPECT_TRUE(exporter.GetAndResetMetricSamples().empty()); +} + +TEST_F(BaseShmExporterTest, MetricRecordingAndCollection) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + + BaseShmExporter exporter(options); + + const MetricLabel sent_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, sent_labels, 1024); + exporter.IncrementCounter(metric_names::kSentBytesTotal, sent_labels, 512); + + const MetricLabel gauge_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}}; + exporter.SetGauge(metric_names::kBufferAllocatedBytes, gauge_labels, 4096.0); + + exporter.ObserveHistogram(metric_names::kTransferDurationMs, sent_labels, + 12.5); + + absl::flat_hash_map totals; + exporter.CollectMetrics(totals); + + EXPECT_EQ(totals["sent_bytes_total/direction=push"], 1536.0); + EXPECT_EQ(totals["buffer_allocated_bytes/direction=pull"], 4096.0); + EXPECT_EQ(totals["transfer_duration_ms/direction=push/count"], 1.0); + EXPECT_EQ(totals["transfer_duration_ms/direction=push/sum"], 12.5); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/metrics_backend.h b/tpu_sync/telemetry/metrics_backend.h index 550c52f6..4f448227 100644 --- a/tpu_sync/telemetry/metrics_backend.h +++ b/tpu_sync/telemetry/metrics_backend.h @@ -16,6 +16,7 @@ #define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_METRICS_BACKEND_H_ #include +#include #include #include #include @@ -26,7 +27,7 @@ namespace tpu_raiden::telemetry { -enum class MetricType { +enum class MetricType : uint8_t { kCounter, kGauge, kHistogram, @@ -55,6 +56,22 @@ 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. If empty, falls back to + // the SHM_DIR environment variable, or "/dev/shm". + std::string base_shm_dir; + + // Returns the resolved shared-memory directory, checking `base_shm_dir`, + // the SHM_DIR environment variable, and falling back to "/dev/shm". + std::string GetShmDir() const { + if (!base_shm_dir.empty()) { + return base_shm_dir; + } + const char* env_shm = std::getenv("SHM_DIR"); + if (env_shm != nullptr && *env_shm != '\0') { + return env_shm; + } + return "/dev/shm"; + } }; // Structure defining centralized metadata for a Raiden metric. diff --git a/tpu_sync/telemetry/prometheus_shm_exporter.cc b/tpu_sync/telemetry/prometheus_shm_exporter.cc new file mode 100644 index 00000000..ab5ff049 --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter.cc @@ -0,0 +1,201 @@ +// 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/prometheus_shm_exporter.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "prometheus/client_metric.h" +#include "prometheus/collectable.h" +#include "prometheus/exposer.h" +#include "prometheus/metric_family.h" +#include "prometheus/metric_type.h" +#include "prometheus/text_serializer.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "tpu_sync/telemetry/base_shm_exporter.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_collector.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { +namespace { + +std::string JoinHostPort(absl::string_view host, int port) { + if (absl::StrContains(host, ':') && !absl::StartsWith(host, "[")) { + return absl::StrCat("[", host, "]:", port); + } + return absl::StrCat(host, ":", port); +} + +std::vector ParseLabels( + absl::string_view enc) { + std::vector labels; + if (enc.empty()) return labels; + for (absl::string_view p : absl::StrSplit(enc, ';')) { + std::pair kv = + absl::StrSplit(p, absl::MaxSplits('=', 1)); + if (!kv.first.empty()) { + labels.push_back({std::string(kv.first), std::string(kv.second)}); + } + } + std::sort(labels.begin(), labels.end()); + return labels; +} + +class PrometheusShmCollectable : public prometheus::Collectable { + public: + explicit PrometheusShmCollectable(const PrometheusShmExporter* exp) + : exp_(exp) {} + std::vector Collect() const override { + return exp_ ? exp_->CollectMetricFamilies() + : std::vector{}; + } + + private: + const PrometheusShmExporter* exp_; +}; + +} // namespace + +PrometheusShmExporter::PrometheusShmExporter(const ExporterOptions& options) + : BaseShmExporter(options) { + if (options_.port > 0) Start(); +} + +PrometheusShmExporter::~PrometheusShmExporter() { Stop(); } + +void PrometheusShmExporter::Start() { + BaseShmExporter::Start(); + absl::MutexLock lock(&mutex_); + if (exposer_ != nullptr || options_.port <= 0) return; + + if (options_.port >= kMinPort && options_.port <= kMaxPort) { + std::string endpoint = JoinHostPort(options_.bind_address, options_.port); + try { + exposer_ = std::make_unique(endpoint); + collectable_ = std::make_shared(this); + exposer_->RegisterCollectable(collectable_); + LOG(INFO) << "Prometheus SHM exporter listening on http://" << endpoint + << "/metrics"; + } catch (const std::exception& e) { + LOG(ERROR) << "Failed to bind Prometheus HTTP exposer on " << endpoint + << ": " << e.what(); + exposer_.reset(); + collectable_.reset(); + } + } else { + LOG(WARNING) << "Invalid port configured for Prometheus SHM HTTP exposer: " + << options_.port; + } +} + +void PrometheusShmExporter::Stop() { + absl::MutexLock lock(&mutex_); + exposer_.reset(); + collectable_.reset(); + BaseShmExporter::Stop(); +} + +bool PrometheusShmExporter::IsExposerRunning() const { + absl::MutexLock lock(&mutex_); + return exposer_ != nullptr; +} + +std::vector +PrometheusShmExporter::CollectMetricFamilies() const { + std::vector families; + if (!collector_) return families; + + absl::flat_hash_map totals; + collector_->CollectMetrics(totals); + + for (const MetricMetadata& meta : metric_metadata::kAllMetrics) { + prometheus::MetricFamily family; + family.name = absl::StrCat("tpu_raiden_", meta.name); + family.help = std::string(meta.description); + family.type = + (meta.type == MetricType::kCounter) ? prometheus::MetricType::Counter + : (meta.type == MetricType::kGauge) ? prometheus::MetricType::Gauge + : prometheus::MetricType::Histogram; + + std::string prefix = absl::StrCat(meta.name, "/"); + + if (meta.type == MetricType::kCounter || meta.type == MetricType::kGauge) { + for (const auto& [raw_key, val] : totals) { + if (!absl::StartsWith(raw_key, prefix)) continue; + prometheus::ClientMetric cm; + cm.label = ParseLabels(raw_key.substr(prefix.length())); + if (meta.type == MetricType::kCounter) { + cm.counter.value = val; + } else { + cm.gauge.value = val; + } + family.metric.push_back(std::move(cm)); + } + } else if (meta.type == MetricType::kHistogram) { + absl::flat_hash_set seen_enc; + for (const auto& [raw_key, _] : totals) { + if (!absl::StartsWith(raw_key, prefix)) continue; + std::string sub = raw_key.substr(prefix.length()); + size_t slash_pos = sub.find('/'); + std::string enc = + (slash_pos != std::string::npos) ? sub.substr(0, slash_pos) : ""; + if (!seen_enc.insert(enc).second) continue; + + prometheus::ClientMetric cm; + cm.label = ParseLabels(enc); + std::string base_key = absl::StrCat(prefix, enc); + auto count_it = totals.find(absl::StrCat(base_key, "/count")); + if (count_it != totals.end()) { + cm.histogram.sample_count = static_cast(count_it->second); + } + auto sum_it = totals.find(absl::StrCat(base_key, "/sum")); + if (sum_it != totals.end()) cm.histogram.sample_sum = sum_it->second; + for (size_t b = 0; b < kNumHistogramBuckets; ++b) { + auto b_it = totals.find(absl::StrCat(base_key, "/bucket_", b)); + uint64_t b_count = + (b_it != totals.end()) ? static_cast(b_it->second) : 0; + cm.histogram.bucket.push_back({b_count, kDefaultHistogramBuckets[b]}); + } + family.metric.push_back(std::move(cm)); + } + } + + std::sort( + family.metric.begin(), family.metric.end(), + [](const prometheus::ClientMetric& a, + const prometheus::ClientMetric& b) { return a.label < b.label; }); + families.push_back(std::move(family)); + } + return families; +} + +std::string PrometheusShmExporter::GetTextSnapshot() const { + return prometheus::TextSerializer().Serialize(CollectMetricFamilies()); +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/prometheus_shm_exporter.h b/tpu_sync/telemetry/prometheus_shm_exporter.h new file mode 100644 index 00000000..8c12df9c --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter.h @@ -0,0 +1,54 @@ +// 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_PROMETHEUS_SHM_EXPORTER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_PROMETHEUS_SHM_EXPORTER_H_ + +#include +#include +#include + +#include "prometheus/collectable.h" +#include "prometheus/exposer.h" +#include "prometheus/metric_family.h" +#include "absl/base/thread_annotations.h" +#include "absl/synchronization/mutex.h" +#include "tpu_sync/telemetry/base_shm_exporter.h" +#include "tpu_sync/telemetry/metrics_backend.h" + +namespace tpu_raiden::telemetry { + +class PrometheusShmExporter : public BaseShmExporter { + public: + explicit PrometheusShmExporter(const ExporterOptions& options = {}); + ~PrometheusShmExporter() override; + PrometheusShmExporter(const PrometheusShmExporter&) = delete; + PrometheusShmExporter& operator=(const PrometheusShmExporter&) = delete; + + void Start() override; + void Stop() override; + + std::string GetTextSnapshot() const override; + bool IsExposerRunning() const; + std::vector CollectMetricFamilies() const; + + private: + mutable absl::Mutex mutex_; + std::unique_ptr exposer_ ABSL_GUARDED_BY(mutex_); + std::shared_ptr collectable_ ABSL_GUARDED_BY(mutex_); +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_PROMETHEUS_SHM_EXPORTER_H_ diff --git a/tpu_sync/telemetry/prometheus_shm_exporter_challenger_test.cc b/tpu_sync/telemetry/prometheus_shm_exporter_challenger_test.cc new file mode 100644 index 00000000..c364bd23 --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter_challenger_test.cc @@ -0,0 +1,627 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#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 "third_party/prometheus_cpp_client/core/include/prometheus/metric_family.h" +#include "tpu_sync/telemetry/base_shm_exporter.h" +#include "tpu_sync/telemetry/metrics_api.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/prometheus_shm_exporter.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" +#include "tpu_sync/telemetry/test_util.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::DoubleEq; +using ::testing::DoubleNear; +using ::testing::Ge; +using ::testing::HasSubstr; +using ::testing::Le; +using ::testing::Not; + +std::string HttpGet(int port, const std::string& path) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) return ""; + + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr); + + if (connect(fd, reinterpret_cast(&addr), sizeof(addr)) < 0) { + close(fd); + return ""; + } + + std::string req = absl::StrCat("GET ", path, + " HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: " + "close\r\n\r\n"); + send(fd, req.data(), req.size(), 0); + + std::string response; + char buf[4096]; + ssize_t n = 0; + while ((n = recv(fd, buf, sizeof(buf), 0)) > 0) { + response.append(buf, n); + } + close(fd); + return response; +} + +class PrometheusShmExporterChallengerTest : public ::testing::Test { + protected: + void SetUp() override { + const char* tmp = std::getenv("TEST_TMPDIR"); + test_dir_ = tmp ? absl::StrCat(tmp, "/prom_challenger_", getpid(), "_", + reinterpret_cast(this)) + : absl::StrCat("/tmp/prom_challenger_", getpid(), "_", + reinterpret_cast(this)); + std::filesystem::create_directories(test_dir_); + } + + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + } + + std::string test_dir_; +}; + +// ============================================================================ +// CHALLENGE AREA 1: Port Collisions & Binding Resilience +// ============================================================================ + +TEST_F(PrometheusShmExporterChallengerTest, + PortCollisionTwoExportersGracefulDegradation) { + int port = PickUnusedPort(); + ASSERT_GT(port, 0); + + ExporterOptions opt1; + opt1.base_shm_dir = test_dir_; + opt1.local_rank = "0"; + opt1.bind_address = "127.0.0.1"; + opt1.port = port; + + ExporterOptions opt2; + opt2.base_shm_dir = test_dir_; + opt2.local_rank = "1"; + opt2.bind_address = "127.0.0.1"; + opt2.port = port; + + PrometheusShmExporter exp1(opt1); + EXPECT_TRUE(exp1.IsExposerRunning()); + + // Second exporter attempts to bind to the identical port. + // Must catch exception internally, log ERROR, and NOT crash. + PrometheusShmExporter exp2(opt2); + EXPECT_FALSE(exp2.IsExposerRunning()); + + // Both writers must still function normally for shared-memory writes. + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exp1.IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + exp2.IncrementCounter(metric_names::kSentBytesTotal, labels, 200); + + // Exporter 1 serves HTTP request containing aggregated metrics. + std::string http_res = HttpGet(port, "/metrics"); + EXPECT_THAT(http_res, HasSubstr("HTTP/1.1 200 OK")); + EXPECT_THAT(http_res, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 300")); + + // In-process snapshot from exp2 still works despite exposer failure. + std::string snapshot2 = exp2.GetTextSnapshot(); + EXPECT_THAT(snapshot2, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 300")); +} + +TEST_F(PrometheusShmExporterChallengerTest, + PortCollisionRecoveryAfterPrimaryStops) { + int port = PickUnusedPort(); + ASSERT_GT(port, 0); + + ExporterOptions opt1; + opt1.base_shm_dir = test_dir_; + opt1.local_rank = "0"; + opt1.bind_address = "127.0.0.1"; + opt1.port = port; + + ExporterOptions opt2; + opt2.base_shm_dir = test_dir_; + opt2.local_rank = "1"; + opt2.bind_address = "127.0.0.1"; + opt2.port = port; + + auto exp1 = std::make_unique(opt1); + EXPECT_TRUE(exp1->IsExposerRunning()); + + auto exp2 = std::make_unique(opt2); + EXPECT_FALSE(exp2->IsExposerRunning()); + + // Stop primary exporter to release port. + exp1->Stop(); + EXPECT_FALSE(exp1->IsExposerRunning()); + + // Now retry starting secondary exporter. + exp2->Start(); + EXPECT_TRUE(exp2->IsExposerRunning()); + + std::string http_res = HttpGet(port, "/metrics"); + EXPECT_THAT(http_res, HasSubstr("HTTP/1.1 200 OK")); + + exp2->Stop(); +} + +TEST_F(PrometheusShmExporterChallengerTest, + MassivePortCollisionContentionMultiThreaded) { + int port = PickUnusedPort(); + ASSERT_GT(port, 0); + + constexpr int kNumExporters = 8; + std::vector> exporters(kNumExporters); + std::vector threads; + threads.reserve(kNumExporters); + + for (int i = 0; i < kNumExporters; ++i) { + threads.emplace_back([this, port, i, &exporters]() { + ExporterOptions opt; + opt.base_shm_dir = test_dir_; + opt.local_rank = absl::StrCat(i); + opt.bind_address = "127.0.0.1"; + opt.port = port; + exporters[i] = std::make_unique(opt); + }); + } + + for (auto& t : threads) { + t.join(); + } + + int running_count = 0; + for (int i = 0; i < kNumExporters; ++i) { + if (exporters[i]->IsExposerRunning()) { + running_count++; + } + } + + // Exactly 1 exporter must succeed in binding; the other 7 fail gracefully. + EXPECT_EQ(running_count, 1); + + // All exporters clean up cleanly without hanging or crashing. + for (int i = 0; i < kNumExporters; ++i) { + exporters[i]->Stop(); + } +} + +TEST_F(PrometheusShmExporterChallengerTest, + InvalidPortConfigurationsHandledGracefully) { + const int invalid_ports[] = {-100, -1, 65536, 100000}; + for (int p : invalid_ports) { + ExporterOptions opt; + opt.base_shm_dir = test_dir_; + opt.local_rank = "0"; + opt.port = p; + + PrometheusShmExporter exp(opt); + EXPECT_FALSE(exp.IsExposerRunning()); + } +} + +// ============================================================================ +// CHALLENGE AREA 2: Disabled Exposer (port = 0) +// ============================================================================ + +TEST_F(PrometheusShmExporterChallengerTest, + DisabledExposerZeroPortFullLocalFunctionality) { + ExporterOptions opt; + opt.base_shm_dir = test_dir_; + opt.local_rank = "0"; + opt.port = 0; + + PrometheusShmExporter exp(opt); + EXPECT_FALSE(exp.IsExposerRunning()); + + // Start() explicitly called should not start exposer if port == 0. + exp.Start(); + EXPECT_FALSE(exp.IsExposerRunning()); + + // Metric emission works. + const MetricLabel sent_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exp.IncrementCounter(metric_names::kSentBytesTotal, sent_labels, 555); + + const MetricLabel recv_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}}; + exp.IncrementCounter(metric_names::kReceivedBytesTotal, recv_labels, 888); + + std::string snapshot = exp.GetTextSnapshot(); + EXPECT_THAT(snapshot, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 555")); + EXPECT_THAT( + snapshot, + HasSubstr("tpu_raiden_received_bytes_total{direction=\"pull\"} 888")); + + auto families = exp.CollectMetricFamilies(); + EXPECT_FALSE(families.empty()); + + exp.Stop(); + EXPECT_FALSE(exp.IsExposerRunning()); +} + +TEST_F(PrometheusShmExporterChallengerTest, + MultipleZeroPortExportersConcurrentSharedMemoryAggregation) { + constexpr int kWorkers = 4; + std::vector> exps; + exps.reserve(kWorkers); + + for (int i = 0; i < kWorkers; ++i) { + ExporterOptions opt; + opt.base_shm_dir = test_dir_; + opt.local_rank = absl::StrCat(i); + opt.port = 0; + exps.push_back(std::make_unique(opt)); + } + + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + + for (int i = 0; i < kWorkers; ++i) { + exps[i]->IncrementCounter(metric_names::kSentBytesTotal, labels, 1000); + } + + // Any worker querying snapshot must see the full aggregate (4 * 1000 = 4000). + for (int i = 0; i < kWorkers; ++i) { + std::string snap = exps[i]->GetTextSnapshot(); + EXPECT_THAT( + snap, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 4000")); + } +} + +// ============================================================================ +// CHALLENGE AREA 3: Dead Worker Reaping Integration via Snapshot Calls +// ============================================================================ + +TEST_F(PrometheusShmExporterChallengerTest, + DeadWorkerReapingIntegratedInTextSnapshotAndMetricFamilies) { + ExporterOptions live_opt; + live_opt.base_shm_dir = test_dir_; + live_opt.local_rank = "0"; + live_opt.port = 0; + + PrometheusShmExporter live_exp(live_opt); + + // Simulate dead worker by creating it in a scope and destroying it. + { + ExporterOptions dead_opt; + dead_opt.base_shm_dir = test_dir_; + dead_opt.local_rank = "1"; + dead_opt.port = 0; + + PrometheusShmExporter dead_exp(dead_opt); + const MetricLabel sent_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + dead_exp.IncrementCounter(metric_names::kSentBytesTotal, sent_labels, 450); + + const MetricLabel fail_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "UNAVAILABLE"}}; + dead_exp.IncrementCounter(metric_names::kTransferFailuresTotal, fail_labels, + 7); + } + + // Live worker emits metrics before calling snapshot. + const MetricLabel live_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + live_exp.IncrementCounter(metric_names::kSentBytesTotal, live_labels, 150); + + // Calling CollectMetricFamilies() should trigger dead worker reaping. + auto families = live_exp.CollectMetricFamilies(); + bool found_sent = false; + + for (const auto& fam : families) { + if (fam.name == "tpu_raiden_sent_bytes_total") { + found_sent = true; + for (const auto& m : fam.metric) { + if (!m.label.empty() && m.label[0].value == "push") { + // Expected: Live (150); dead worker unlinked. + EXPECT_DOUBLE_EQ(m.counter.value, 150.0); + } + } + } + } + + EXPECT_TRUE(found_sent); + + // Text snapshot format reflects live worker total. + std::string snapshot = live_exp.GetTextSnapshot(); + EXPECT_THAT(snapshot, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 150")); +} + +TEST_F(PrometheusShmExporterChallengerTest, + SequentialDeadWorkerCascadesUnlinkedOnSweep) { + ExporterOptions survivor_opt; + survivor_opt.base_shm_dir = test_dir_; + survivor_opt.local_rank = "0"; + survivor_opt.port = 0; + + PrometheusShmExporter survivor(survivor_opt); + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + + // 5 workers start, record 100, and die sequentially. + for (int i = 1; i <= 5; ++i) { + std::string dead_file; + { + ExporterOptions dead_opt; + dead_opt.base_shm_dir = test_dir_; + dead_opt.local_rank = absl::StrCat(i); + dead_opt.port = 0; + PrometheusShmExporter dead_exp(dead_opt); + dead_exp.IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + } + // Snapshot unlinks dead worker. + std::string snap = survivor.GetTextSnapshot(); + survivor.IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + snap = survivor.GetTextSnapshot(); + EXPECT_THAT(snap, HasSubstr(absl::StrCat( + "tpu_raiden_sent_bytes_total{direction=\"push\"} ", + i * 100))); + } +} + +TEST_F(PrometheusShmExporterChallengerTest, DeadWorkerReapingOptionPreserved) { + ExporterOptions live_opt; + live_opt.base_shm_dir = test_dir_; + live_opt.local_rank = "0"; + live_opt.port = 0; + + PrometheusShmExporter live_exp(live_opt); + + { + ExporterOptions dead_opt; + dead_opt.base_shm_dir = test_dir_; + dead_opt.local_rank = "1"; + dead_opt.port = 0; + PrometheusShmExporter dead_exp(dead_opt); + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + dead_exp.IncrementCounter(metric_names::kSentBytesTotal, labels, 500); + } + + // Dead worker is unlinked and excluded from live snapshot. + std::string snapshot = live_exp.GetTextSnapshot(); + EXPECT_THAT( + snapshot, + Not(HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 500"))); +} + +// ============================================================================ +// CHALLENGE AREA 4: Histogram Bucket Ordering, Boundary Values & Formatting +// ============================================================================ + +TEST_F(PrometheusShmExporterChallengerTest, + HistogramSlotCumulativeBucketOrderingAndBoundaries) { + ShmHistogramSlot slot; + EXPECT_EQ(slot.sample_count.load(), 0); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + + // Default bucket boundaries: + // [0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, + // 750.0, 1000.0, 2500.0, 5000.0, 7500.0, 10000.0, 25000.0, 50000.0] + + // Value 0.05 falls into bucket 0 (<= 0.1) and all cumulative buckets. + slot.Observe(0.05); + EXPECT_EQ(slot.sample_count.load(), 1); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.05)); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 1); + } + + // Value 5.5 falls into bucket 6 (<= 10.0), skipping buckets 0..5. + slot.Observe(5.5); + EXPECT_EQ(slot.sample_count.load(), 2); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(5.55)); + for (size_t b = 0; b <= 5; ++b) EXPECT_EQ(slot.bucket_counts[b].load(), 1); + for (size_t b = 6; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 2); + } + + // Value 60000.0 exceeds highest boundary (50000.0), updating only +Inf bucket + // (index 20). + slot.Observe(60000.0); + EXPECT_EQ(slot.sample_count.load(), 3); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(60005.55)); + for (size_t b = 0; b <= 5; ++b) EXPECT_EQ(slot.bucket_counts[b].load(), 1); + for (size_t b = 6; b < kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 2); + } + EXPECT_EQ(slot.bucket_counts[kNumHistogramBuckets].load(), 3); + + // Monotonic cumulative invariant: + for (size_t i = 0; i < kNumHistogramBuckets; ++i) { + EXPECT_LE(slot.bucket_counts[i].load(), slot.bucket_counts[i + 1].load()); + } +} + +TEST_F(PrometheusShmExporterChallengerTest, + HistogramSlotExactBoundaryMatching) { + ShmHistogramSlot slot; + slot.Observe(0.1); // Exactly on bucket 0 boundary + slot.Observe(0.1000001); // Just above bucket 0 -> bucket 1 + slot.Observe(50000.0); // Exactly on bucket 19 boundary + slot.Observe(50000.0001); // Just above bucket 19 -> bucket 20 (+Inf) + + EXPECT_EQ(slot.sample_count.load(), 4); + EXPECT_EQ(slot.bucket_counts[0].load(), 1); + EXPECT_EQ(slot.bucket_counts[1].load(), 2); + EXPECT_EQ(slot.bucket_counts[19].load(), 3); + EXPECT_EQ(slot.bucket_counts[20].load(), 4); + + for (size_t i = 0; i < kNumHistogramBuckets; ++i) { + EXPECT_LE(slot.bucket_counts[i].load(), slot.bucket_counts[i + 1].load()); + } +} + +TEST_F(PrometheusShmExporterChallengerTest, + HistogramSlotNonFiniteAndAdversarialValues) { + ShmHistogramSlot slot; + + // 1. Non-finite values discarded + slot.Observe(std::numeric_limits::quiet_NaN()); + slot.Observe(std::numeric_limits::signaling_NaN()); + slot.Observe(-std::numeric_limits::quiet_NaN()); + slot.Observe(std::numeric_limits::infinity()); + slot.Observe(-std::numeric_limits::infinity()); + + EXPECT_EQ(slot.sample_count.load(), 0); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 0); + } + + // 2. Signed zeros + slot.Observe(+0.0); + slot.Observe(-0.0); + EXPECT_EQ(slot.sample_count.load(), 2); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 2); + } + + // 3. Subnormals + slot.Observe(std::numeric_limits::denorm_min()); + EXPECT_EQ(slot.sample_count.load(), 3); + EXPECT_GT(slot.sample_sum.load(), 0.0); + + // 4. Negative numbers + slot.Observe(-50.0); + EXPECT_EQ(slot.sample_count.load(), 4); + EXPECT_LT(slot.sample_sum.load(), 0.0); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 4); + } +} + +TEST_F(PrometheusShmExporterChallengerTest, + HistogramSlotConcurrentContentionCasLoop) { + ShmHistogramSlot slot; + constexpr int kNumThreads = 16; + constexpr int kNumIters = 25000; // Total 400,000 observations + std::atomic start_signal{false}; + + std::vector threads; + threads.reserve(kNumThreads); + + for (int t = 0; t < kNumThreads; ++t) { + threads.emplace_back([&slot, &start_signal, t]() { + while (!start_signal.load(std::memory_order_acquire)) { + } + for (int i = 0; i < kNumIters; ++i) { + double val = ((t * 17 + i) % 20) * 2.5 + 0.05; + slot.Observe(val); + } + }); + } + + start_signal.store(true, std::memory_order_release); + for (auto& th : threads) th.join(); + + EXPECT_EQ(slot.sample_count.load(), + static_cast(kNumThreads) * kNumIters); + EXPECT_EQ(slot.bucket_counts[kNumHistogramBuckets].load(), + static_cast(kNumThreads) * kNumIters); + + // Monotonic cumulative invariant: + for (size_t i = 0; i < kNumHistogramBuckets; ++i) { + EXPECT_LE(slot.bucket_counts[i].load(), slot.bucket_counts[i + 1].load()); + } +} + +TEST_F(PrometheusShmExporterChallengerTest, + HistogramDeadWorkerUnlinkedOnSweep) { + ShmCollectorOptions collector_opts; + collector_opts.shm_dir = test_dir_; + ShmCollector collector(collector_opts); + + // Step 2: Create a dead worker file with a histogram slot. + std::string worker_path = absl::StrCat(test_dir_, "/worker_9999_hist.mmap"); + int fd = + open(worker_path.c_str(), O_RDWR | O_CREAT | O_TRUNC | O_CLOEXEC, 0600); + ASSERT_GE(fd, 0); + ASSERT_EQ(ftruncate(fd, kSegmentTotalFileSize), 0); + + void* addr = mmap(nullptr, kSegmentTotalFileSize, PROT_READ | PROT_WRITE, + MAP_SHARED, fd, 0); + ASSERT_NE(addr, MAP_FAILED); + + auto* seg = reinterpret_cast(addr); + std::memset(seg, 0, kSegmentTotalFileSize); + ShmTocHeader& h = seg->header; + h.version = kSupportedVersion; + h.pid = 9999; + h.max_toc_entries = kMaxTocEntries; + h.data_pool_offset = sizeof(ShmSegmentLayout); + + ShmTocEntry& e = seg->toc[h.toc_entry_count++]; + snprintf(e.metric_name, sizeof(e.metric_name), "sent_bytes_total"); + snprintf(e.encoded_labels, sizeof(e.encoded_labels), "direction=push"); + e.type = MetricType::kCounter; + e.offset = h.data_pool_offset + h.data_pool_bytes; + e.size = sizeof(std::atomic); + e.entry_state.store(TocEntryState::kCommitted); + uint8_t* raw = reinterpret_cast(seg) + e.offset; + new (raw) std::atomic(12345); + h.data_pool_bytes += e.size; + + h.magic.store(kRaidenShmMagic, std::memory_order_release); + munmap(addr, kSegmentTotalFileSize); + close(fd); + + // Sweep unlinks dead worker + absl::flat_hash_map totals; + collector.CollectMetrics(totals); + + EXPECT_EQ(totals["sent_bytes_total/direction=push"], 0.0); + EXPECT_FALSE(std::filesystem::exists(worker_path)); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/prometheus_shm_exporter_test.cc b/tpu_sync/telemetry/prometheus_shm_exporter_test.cc new file mode 100644 index 00000000..551e1838 --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter_test.cc @@ -0,0 +1,339 @@ +// 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/prometheus_shm_exporter.h" + +#include +#include +#include +#include + +#include +#include +#include // NOLINT(build/c++17) +#include +#include +#include + +#include +#include +#include "absl/container/flat_hash_map.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "third_party/prometheus_cpp_client/core/include/prometheus/metric_family.h" +#include "tpu_sync/telemetry/base_shm_exporter.h" +#include "tpu_sync/telemetry/metrics_api.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/test_util.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::HasSubstr; + +std::string HttpGet(int port, const std::string& path) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) return ""; + + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr); + + if (connect(fd, reinterpret_cast(&addr), sizeof(addr)) < 0) { + close(fd); + return ""; + } + + std::string req = absl::StrCat("GET ", path, + " HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: " + "close\r\n\r\n"); + send(fd, req.data(), req.size(), 0); + + std::string response; + char buf[4096]; + ssize_t n = 0; + while ((n = recv(fd, buf, sizeof(buf), 0)) > 0) { + response.append(buf, n); + } + close(fd); + return response; +} + +class PrometheusShmExporterTest : public testing::Test { + protected: + void SetUp() override { + const char* tmp = std::getenv("TEST_TMPDIR"); + test_dir_ = tmp ? absl::StrCat(tmp, "/prom_shm_test_", getpid()) + : absl::StrCat("/tmp/prom_shm_test_", getpid()); + std::filesystem::create_directories(test_dir_); + } + + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + } + + std::string test_dir_; +}; + +TEST_F(PrometheusShmExporterTest, ThrowsWhenLocalRankNotSpecified) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = std::nullopt; + EXPECT_THROW(PrometheusShmExporter exporter(options), std::invalid_argument); +} + +TEST_F(PrometheusShmExporterTest, BaseExporterLifecycle) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + + BaseShmExporter exporter(options); + EXPECT_EQ(exporter.GetOptions().port, 0); + EXPECT_NE(exporter.GetWriter(), nullptr); + EXPECT_NE(exporter.GetCollector(), nullptr); + EXPECT_FALSE(exporter.IsRunning()); + + exporter.Start(); + EXPECT_TRUE(exporter.IsRunning()); + + exporter.Stop(); + EXPECT_FALSE(exporter.IsRunning()); +} + +TEST_F(PrometheusShmExporterTest, BaseExporterMetricForwarding) { + ExporterOptions options; + options.base_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, 2048); + + absl::flat_hash_map totals; + exporter.CollectMetrics(totals); + EXPECT_EQ(totals["sent_bytes_total/direction=push"], 2048); +} + +TEST_F(PrometheusShmExporterTest, TextSnapshotContainsMetricFormatting) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + options.port = 0; + + PrometheusShmExporter exporter(options); + const MetricLabel sent_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, sent_labels, 1024); + + const MetricLabel fail_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "DEADLINE_EXCEEDED"}}; + exporter.IncrementCounter(metric_names::kTransferFailuresTotal, fail_labels, + 3); + + std::string snapshot = exporter.GetTextSnapshot(); + EXPECT_THAT(snapshot, HasSubstr("# HELP tpu_raiden_sent_bytes_total")); + EXPECT_THAT(snapshot, + HasSubstr("# TYPE tpu_raiden_sent_bytes_total counter")); + EXPECT_THAT( + snapshot, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 1024")); + EXPECT_THAT( + snapshot, + HasSubstr("tpu_raiden_transfer_failures_total{direction=\"pull\",error_" + "code=\"DEADLINE_EXCEEDED\"} 3")); +} + +TEST_F(PrometheusShmExporterTest, ExposerDisabledWhenPortZero) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + options.port = 0; + + PrometheusShmExporter exporter(options); + EXPECT_FALSE(exporter.IsExposerRunning()); + exporter.Start(); + EXPECT_FALSE(exporter.IsExposerRunning()); + exporter.Stop(); +} + +TEST_F(PrometheusShmExporterTest, ExposerStartsAndServesHttpScrape) { + int port = PickUnusedPort(); + if (port <= 0) { + GTEST_SKIP() << "No free port available for HTTP exposer test"; + } + + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + options.bind_address = "127.0.0.1"; + options.port = port; + + PrometheusShmExporter exporter(options); + EXPECT_TRUE(exporter.IsExposerRunning()); + + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, labels, 777); + + std::string response = HttpGet(port, "/metrics"); + EXPECT_THAT(response, HasSubstr("HTTP/1.1 200 OK")); + EXPECT_THAT(response, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 777")); + + exporter.Stop(); + EXPECT_FALSE(exporter.IsExposerRunning()); +} + +TEST_F(PrometheusShmExporterTest, ExposerGracefulDegradationOnPortCollision) { + int port = PickUnusedPort(); + if (port <= 0) { + GTEST_SKIP() << "No free port available for port collision test"; + } + + int sock = socket(AF_INET, SOCK_STREAM, 0); + ASSERT_GE(sock, 0); + int on = 1; + setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr); + ASSERT_EQ(bind(sock, reinterpret_cast(&addr), sizeof(addr)), 0); + ASSERT_EQ(listen(sock, 1), 0); + + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + options.bind_address = "127.0.0.1"; + options.port = port; + + PrometheusShmExporter exporter(options); + EXPECT_FALSE(exporter.IsExposerRunning()); + + close(sock); +} + +TEST_F(PrometheusShmExporterTest, CollectMetricFamiliesDataIntegrity) { + ExporterOptions options; + options.base_shm_dir = test_dir_; + options.local_rank = "0"; + options.port = 0; + + PrometheusShmExporter exporter(options); + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, labels, 500); + + std::vector families = + exporter.CollectMetricFamilies(); + EXPECT_FALSE(families.empty()); + + bool found_sent = false; + for (const auto& fam : families) { + if (fam.name == "tpu_raiden_sent_bytes_total") { + found_sent = true; + EXPECT_EQ(fam.type, prometheus::MetricType::Counter); + ASSERT_FALSE(fam.metric.empty()); + bool found_push = false; + for (const auto& m : fam.metric) { + for (const auto& lbl : m.label) { + if (lbl.name == "direction" && lbl.value == "push") { + found_push = true; + EXPECT_EQ(m.counter.value, 500.0); + } + } + } + EXPECT_TRUE(found_push); + } + } + EXPECT_TRUE(found_sent); +} + +TEST_F(PrometheusShmExporterTest, DecentralizedApproachANoLeaderLock) { + ExporterOptions options1; + options1.base_shm_dir = test_dir_; + options1.local_rank = "0"; + options1.port = 0; + + ExporterOptions options2; + options2.base_shm_dir = test_dir_; + options2.local_rank = "1"; + options2.port = 0; + + PrometheusShmExporter exporter1(options1); + PrometheusShmExporter exporter2(options2); + + EXPECT_FALSE(std::filesystem::exists(test_dir_ + "/leader.lock")); +} + +TEST_F(PrometheusShmExporterTest, MultiWorkerAggregationInTextSnapshot) { + ExporterOptions options1; + options1.base_shm_dir = test_dir_; + options1.local_rank = "0"; + options1.port = 0; + + ExporterOptions options2; + options2.base_shm_dir = test_dir_; + options2.local_rank = "1"; + options2.port = 0; + + PrometheusShmExporter exporter1(options1); + PrometheusShmExporter exporter2(options2); + + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter1.IncrementCounter(metric_names::kSentBytesTotal, labels, 150); + exporter2.IncrementCounter(metric_names::kSentBytesTotal, labels, 350); + + std::string snapshot = exporter1.GetTextSnapshot(); + EXPECT_THAT(snapshot, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 500")); +} + +TEST_F(PrometheusShmExporterTest, DeadWorkerUnlinkedOnSweepInTextSnapshot) { + ExporterOptions options1; + options1.base_shm_dir = test_dir_; + options1.local_rank = "0"; + options1.port = 0; + + PrometheusShmExporter exporter1(options1); + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter1.IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + + std::string file2; + { + ExporterOptions options2; + options2.base_shm_dir = test_dir_; + options2.local_rank = "1"; + options2.port = 0; + + PrometheusShmExporter exporter2(options2); + file2 = exporter2.GetWriter()->file_path(); + exporter2.IncrementCounter(metric_names::kSentBytesTotal, labels, 400); + } + + std::string snapshot = exporter1.GetTextSnapshot(); + // Exporter 1 remains (100); dead exporter 2 unlinked + EXPECT_THAT(snapshot, + HasSubstr("tpu_raiden_sent_bytes_total{direction=\"push\"} 100")); + EXPECT_FALSE(std::filesystem::exists(file2)); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/shm/BUILD b/tpu_sync/telemetry/shm/BUILD new file mode 100644 index 00000000..71ccad3d --- /dev/null +++ b/tpu_sync/telemetry/shm/BUILD @@ -0,0 +1,100 @@ +# 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. + +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "shm_layout", + hdrs = ["shm_layout.h"], + deps = [ + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "shm_layout_test", + srcs = ["shm_layout_test.cc"], + deps = [ + ":shm_layout", + "//tpu_sync/telemetry:metrics_backend", + "@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_absl//absl/synchronization", + "@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: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..3b20d460 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.cc @@ -0,0 +1,221 @@ +// 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 // NOLINT(build/c++17) +#include +#include +#include // NOLINT(build/c++11) +#include + +#include "absl/cleanup/cleanup.h" +#include "absl/container/flat_hash_map.h" +#include "absl/log/check.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 ShmSegmentLayout* seg, + absl::flat_hash_map& totals) { + if (seg == nullptr || + (reinterpret_cast(seg) % alignof(ShmSegmentLayout) != 0)) { + return; + } + if (seg->header.data_pool_offset < sizeof(ShmSegmentLayout) || + seg->header.data_pool_offset >= kSegmentTotalFileSize) { + return; + } + uint32_t count = seg->header.toc_entry_count.load(std::memory_order_acquire); + uint32_t num_entries = std::min(count, static_cast(kMaxTocEntries)); + + for (uint32_t i = 0; i < num_entries; ++i) { + const ShmTocEntry& e = seg->toc[i]; + if (e.entry_state.load(std::memory_order_acquire) != + TocEntryState::kCommitted) { + continue; + } + if (e.offset < seg->header.data_pool_offset || + static_cast(e.offset) + e.size > kSegmentTotalFileSize) { + continue; + } + + absl::string_view metric_name( + e.metric_name, strnlen(e.metric_name, sizeof(e.metric_name))); + if (metric_name.empty()) { + continue; + } + + absl::string_view encoded_labels( + e.encoded_labels, strnlen(e.encoded_labels, sizeof(e.encoded_labels))); + + const uint8_t* raw = reinterpret_cast(seg) + e.offset; + std::string key(metric_name); + if (!encoded_labels.empty()) { + absl::StrAppend(&key, "/", encoded_labels); + } + + switch (e.type) { + case MetricType::kCounter: { + if (e.size >= sizeof(std::atomic) && + (reinterpret_cast(raw) % + alignof(std::atomic) == + 0)) { + totals[key] += static_cast( + reinterpret_cast*>(raw)->load( + std::memory_order_relaxed)); + } + break; + } + case MetricType::kGauge: { + if (e.size >= sizeof(std::atomic) && + (reinterpret_cast(raw) % alignof(std::atomic) == + 0)) { + double val = reinterpret_cast*>(raw)->load( + std::memory_order_relaxed); + if (std::isfinite(val)) { + totals[key] += val; + } + } + break; + } + case MetricType::kHistogram: { + if (e.size >= sizeof(ShmHistogramSlot) && + (reinterpret_cast(raw) % alignof(ShmHistogramSlot) == + 0)) { + auto* h = reinterpret_cast(raw); + totals[absl::StrCat(key, "/count")] += static_cast( + h->sample_count.load(std::memory_order_relaxed)); + double sum = h->sample_sum.load(std::memory_order_relaxed); + if (std::isfinite(sum)) { + 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( + h->bucket_counts[b].load(std::memory_order_relaxed)); + } + } + break; + } + } + } +} + +} // namespace internal + +void ShmCollector::CollectMetrics( + absl::flat_hash_map& totals) const { + totals.clear(); + std::error_code ec; + if (!std::filesystem::exists(options_.shm_dir, ec) || ec) { + return; + } + + auto it = std::filesystem::directory_iterator(options_.shm_dir, ec); + if (ec) { + return; + } + + for (; it != std::filesystem::directory_iterator();) { + absl::Cleanup advance = [&] { + it.increment(ec); + if (ec) { + ec.clear(); + } + }; + + std::error_code entry_ec; + if (!it->is_regular_file(entry_ec) || entry_ec) { + continue; + } + std::string fn = it->path().filename().string(); + if (!absl::StartsWith(fn, kShmFilePrefix) || + !absl::EndsWith(fn, kShmFileExtension)) { + continue; + } + + std::string path = it->path().string(); + 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 st = {}; + if (fstat(fd, &st) != 0) { + 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. + // Verify inode before unlinking to prevent TOCTOU. + struct stat st_current = {}; + if (st.st_nlink > 0 && lstat(path.c_str(), &st_current) == 0 && + st_current.st_ino == st.st_ino && st_current.st_dev == st.st_dev) { + unlink(path.c_str()); + } + // Do NOT call flock(LOCK_UN); close_fd will release the lock. + } else if (flock(fd, LOCK_SH | LOCK_NB) == 0) { + // Process is live. Prevent SIGBUS: must be at least + // kSegmentTotalFileSize. + if (st.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); + }; + auto* live = reinterpret_cast(addr); + if (live->header.magic.load(std::memory_order_acquire) == + kRaidenShmMagic && + live->header.version == kSupportedVersion) { + internal::AggregateSegment(live, 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..8d067d15 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.h @@ -0,0 +1,72 @@ +// 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" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { + +namespace internal { + +// Aggregates metric entries from a single memory-mapped segment layout into +// `totals`. Exposed for testing only. +void AggregateSegment(const ShmSegmentLayout* 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); + ~ShmCollector() = default; + + ShmCollector(const ShmCollector&) = default; + ShmCollector& operator=(const ShmCollector&) = default; + ShmCollector(ShmCollector&&) noexcept = default; + ShmCollector& operator=(ShmCollector&&) noexcept = default; + + // 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. + // + // 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..7aa90571 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector_test.cc @@ -0,0 +1,395 @@ +// 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 // 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_}; + } + + 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); + + 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)); + + // 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(); + dead_files += absl::StrContains(fn, "worker_rank_dead_"); + live_files += absl::StrContains(fn, "worker_rank_0_") || + absl::StrContains(fn, "worker_rank_1_"); + } + 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) { + 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'; + (void)write(pfd[1], &ready, 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, AggregatesLabeledHistograms) { + ShmWriter writer(WriterOptions("hist_worker")); + MetricLabel push{metric_labels::kDirection, metric_labels::kDirectionPush}; + MetricLabel err{metric_labels::kErrorCode, "RESOURCE_EXHAUSTED"}; + const std::array labels = {push, err}; + + writer.ObserveHistogram(metric_names::kTransferDurationMs, + {labels.data(), labels.size()}, 0.05); + writer.ObserveHistogram(metric_names::kTransferDurationMs, + {labels.data(), labels.size()}, 0.25); + + absl::flat_hash_map totals; + ShmCollector(CollectorOptions()).CollectMetrics(totals); + + std::string prefix = + "transfer_duration_ms/direction=push;error_code=RESOURCE_EXHAUSTED"; + EXPECT_THAT(totals[absl::StrCat(prefix, "/count")], DoubleEq(2.0)); + EXPECT_THAT(totals[absl::StrCat(prefix, "/sum")], DoubleEq(0.30)); + EXPECT_THAT(totals[absl::StrCat(prefix, "/bucket_0")], DoubleEq(1.0)); + EXPECT_THAT(totals[absl::StrCat(prefix, "/bucket_1")], DoubleEq(1.0)); +} + +TEST_F(ShmCollectorTest, ReapsDeadTruncatedFiles) { + const std::string trunc_path = absl::StrCat(test_dir_, "/", kShmFilePrefix, + "dead_trunc", kShmFileExtension); + { + std::ofstream ofs(trunc_path, std::ios::binary); + ofs.write("short", 5); + } + absl::flat_hash_map totals; + ShmCollector(CollectorOptions()).CollectMetrics(totals); + + EXPECT_FALSE(std::filesystem::exists(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, uint32_t ver) { + ShmTocHeader h{.version = ver}; + 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, kSupportedVersion); + write_hdr("bad_ver.mmap", kRaidenShmMagic, 999); + write_file("other_file.txt", "non-shm content"); + std::filesystem::create_directories( + absl::StrCat(test_dir_, "/", kShmFilePrefix, "dir.mmap")); + symlink("/dev/null", + absl::StrCat(test_dir_, "/", kShmFilePrefix, "symlink.mmap").c_str()); + + 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_EQ(totals["sent_bytes_total/direction=push"], 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.version = kSupportedVersion; + seg->header.max_toc_entries = kMaxTocEntries; + seg->header.data_pool_offset = sizeof(ShmSegmentLayout); + + // Unaligned segment pointer. + auto unaligned_seg = reinterpret_cast( + reinterpret_cast(seg) + 1); + internal::AggregateSegment(unaligned_seg, totals); + EXPECT_TRUE(totals.empty()); + + // 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 + 64) std::atomic(100); + auto* h = reinterpret_cast(mem->data + pool + 128); + h->sample_count.store(1); + h->sample_sum.store(0.5); + h->bucket_counts[2].store(1); + + 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 + 64, 8, + TocEntryState::kCommitted}, + EntrySpec{"hist_m", MetricType::kHistogram, pool + 128, + 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 + 64, 8, + TocEntryState::kCommitted}, + EntrySpec{"underflow", MetricType::kCounter, + sizeof(ShmSegmentLayout) - 64, 8, TocEntryState::kCommitted}, + }; + for (size_t i = 0; i < specs.size(); ++i) { + ShmTocEntry& e = seg->toc[i]; + snprintf(e.metric_name, sizeof(e.metric_name), "%.*s", + static_cast(specs[i].name.size()), specs[i].name.data()); + e.type = specs[i].type; + e.offset = specs[i].off; + e.size = specs[i].size; + e.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_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_layout.h b/tpu_sync/telemetry/shm/shm_layout.h new file mode 100644 index 00000000..b89c51f5 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_layout.h @@ -0,0 +1,185 @@ +// 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_LAYOUT_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_LAYOUT_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" + +namespace tpu_raiden::telemetry { + +inline constexpr uint32_t kRaidenShmMagic = 0xABCD1234; +inline constexpr uint32_t kSupportedVersion = 2; +inline constexpr absl::string_view kShmFilePrefix = "worker_rank_"; +inline constexpr absl::string_view kShmFileExtension = ".mmap"; + +inline constexpr size_t kMaxTocEntries = 1024; +// Data pool capacity for metric slots in a single chunk (64 KB). +// Sized to accommodate 1024 64-byte aligned counter/gauge slots, or up to 341 +// 192-byte histogram slots. Multi-chunk dynamic expansion allocates additional +// chunk files on demand if data pool capacity is exceeded. +inline constexpr size_t kMaxDataPoolBytes = 64 * 1024; // 64 KB +inline constexpr size_t kMetricSlotAlignment = 64; +inline constexpr size_t kNumHistogramBuckets = + std::size(kDefaultHistogramBuckets); + +// State machine transitions for a Table of Contents (TOC) entry descriptor: +// - kUninitialized: Entry slot is empty or unallocated. +// - kWriting: Writer is populating entry fields (name, labels, type, offset). +// - kCommitted: Entry is fully published and valid for reader consumption. +// +// Writers transition from kWriting to kCommitted using release memory ordering +// after all entry fields and data pool slot initialization are complete. +// Readers load entry_state using acquire memory ordering to safely synchronize +// and read published metric descriptors and slot contents without locks. +enum class TocEntryState : uint32_t { + kUninitialized = 0, + kWriting = 1, + kCommitted = 2, +}; + +// Table of Contents entry descriptor for a single metric stream. +// +// Publication protocol: The writer fills metric_name, encoded_labels, type, +// offset, and size, then publishes the entry by setting entry_state to +// TocEntryState::kCommitted with release memory ordering. Readers must load +// entry_state with acquire memory ordering before accessing the metric slot at +// offset. +struct alignas(64) ShmTocEntry { + char metric_name[64]; + char encoded_labels[128]; + MetricType type; + // Explicitly reserve 3 bytes to ensure 64-byte alignment of the TOC entry. + uint8_t reserved[3]{}; + uint32_t offset; + uint32_t size; + std::atomic entry_state{TocEntryState::kUninitialized}; + uint8_t padding[48]{}; +}; + +// 64-byte aligned header for the shared-memory telemetry segment. +// +// Initialization and publication protocol: +// The creator process zeroes or mmaps the shared memory segment, sets all +// non-atomic fields (version, pid, max_toc_entries, data_pool_offset, +// chunk_index), and finally publishes the segment by writing kRaidenShmMagic to +// magic using release memory ordering. Readers must verify that magic loaded +// with acquire memory ordering equals kRaidenShmMagic before reading any header +// or TOC fields. +// +// Data pool metric slots are aligned to kMetricSlotAlignment (64 bytes) to +// prevent intra-process false sharing between multiple threads concurrently +// updating distinct metric streams on different CPU cores. +struct alignas(64) ShmTocHeader { + std::atomic magic{0}; + uint32_t version{kSupportedVersion}; + int64_t pid{0}; + std::atomic toc_entry_count{0}; + uint32_t max_toc_entries{kMaxTocEntries}; + uint32_t data_pool_offset{0}; + std::atomic data_pool_bytes{0}; + uint32_t chunk_index{0}; + uint8_t padding[28]{}; +}; + +// Lock-free shared-memory slot for a single histogram metric stream. +// +// Uses the standard uniform histogram bucket distribution +// (`kDefaultHistogramBuckets`) to ensure a fixed 192-byte standard layout with +// 64-byte alignment and zero serialization overhead across heterogeneous +// exporter pipelines. +// +// Stores non-cumulative (differential) bucket counts to eliminate write +// amplification (3 atomic operations per Observe instead of 23) and prevent +// reader-writer monotonicity race conditions. Exporters compute cumulative +// bucket counts on-the-fly, which is inherently monotonic since all bucket +// counts are non-negative. +struct alignas(64) ShmHistogramSlot { + std::atomic sample_count{0}; + std::atomic sample_sum{0.0}; + std::atomic bucket_counts[kNumHistogramBuckets + 1]{}; + uint8_t padding[8]{}; + + // Records an observation value into the histogram. Non-finite values (NaN, + // +/-Inf) are ignored. Updates sample count, sample sum, and the specific + // non-cumulative bucket count atomically with relaxed memory ordering. + void Observe(double value) { + if (!std::isfinite(value)) { + return; + } + sample_count.fetch_add(1, std::memory_order_relaxed); + sample_sum.fetch_add(value, std::memory_order_relaxed); + auto it = std::lower_bound(std::begin(kDefaultHistogramBuckets), + std::end(kDefaultHistogramBuckets), value); + size_t bucket_idx = static_cast( + std::distance(std::begin(kDefaultHistogramBuckets), it)); + bucket_counts[bucket_idx].fetch_add(1, std::memory_order_relaxed); + } +}; + +// Complete 64-byte aligned segment layout containing header and TOC entries. +struct alignas(64) ShmSegmentLayout { + ShmTocHeader header; + ShmTocEntry toc[kMaxTocEntries]; +}; + +// Calculates the total shared-memory segment file size in bytes required for +// `max_toc_entries` metric stream descriptors and `data_pool_bytes` of metric +// slot storage. Returns the total segment size including header, TOC, and data +// pool. +constexpr size_t CalculateChunkFileSize(size_t max_toc_entries, + size_t data_pool_bytes) { + return sizeof(ShmTocHeader) + (sizeof(ShmTocEntry) * max_toc_entries) + + data_pool_bytes; +} + +inline constexpr size_t kSegmentTotalFileSize = + CalculateChunkFileSize(kMaxTocEntries, kMaxDataPoolBytes); + +static_assert(std::is_standard_layout_v && + std::is_standard_layout_v && + std::is_standard_layout_v && + std::is_standard_layout_v); + +static_assert(std::is_trivially_copyable_v && + std::is_trivially_copyable_v && + std::is_trivially_copyable_v && + std::is_trivially_copyable_v); + +static_assert(alignof(ShmTocHeader) == 64 && sizeof(ShmTocHeader) == 64); +static_assert(alignof(ShmTocEntry) == 64 && sizeof(ShmTocEntry) == 256); +static_assert(alignof(ShmHistogramSlot) == 64 && + sizeof(ShmHistogramSlot) == 192); +static_assert(alignof(ShmSegmentLayout) == 64 && + sizeof(ShmSegmentLayout) == 262208); + +static_assert(std::atomic::is_always_lock_free && + std::atomic::is_always_lock_free && + std::atomic::is_always_lock_free && + std::atomic::is_always_lock_free && + std::atomic::is_always_lock_free); + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_SHM_SHM_LAYOUT_H_ diff --git a/tpu_sync/telemetry/shm/shm_layout_test.cc b/tpu_sync/telemetry/shm/shm_layout_test.cc new file mode 100644 index 00000000..f61b4fd1 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_layout_test.cc @@ -0,0 +1,514 @@ +// 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_layout.h" + +#include +#include // NOLINT +#include +#include +#include +#include +#include +#include +#include +#include // NOLINT +#include +#include + +#include +#include +#include "tpu_sync/telemetry/metrics_backend.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::DoubleEq; + +TEST(ShmLayoutTest, StaticLayoutInvariants) { + EXPECT_TRUE(std::is_standard_layout_v); + EXPECT_TRUE(std::is_standard_layout_v); + EXPECT_TRUE(std::is_standard_layout_v); + EXPECT_TRUE(std::is_standard_layout_v); + + EXPECT_TRUE(std::is_trivially_copyable_v); + EXPECT_TRUE(std::is_trivially_copyable_v); + EXPECT_TRUE(std::is_trivially_copyable_v); + EXPECT_TRUE(std::is_trivially_copyable_v); + + EXPECT_EQ(alignof(ShmTocHeader), 64); + EXPECT_EQ(sizeof(ShmTocHeader), 64); + EXPECT_EQ(alignof(ShmSegmentLayout), 64); + EXPECT_EQ(sizeof(ShmSegmentLayout), 262208); + EXPECT_EQ(alignof(ShmTocEntry), 64); + EXPECT_EQ(sizeof(ShmTocEntry), 256); + EXPECT_EQ(alignof(ShmHistogramSlot), 64); + EXPECT_EQ(sizeof(ShmHistogramSlot), 192); + + EXPECT_EQ(sizeof(ShmTocEntry::metric_name), 64); + EXPECT_EQ(sizeof(ShmTocEntry::encoded_labels), 128); +} + +TEST(ShmLayoutTest, StructMemberOffsetsAndNoInternalPadding) { + // ShmTocEntry offset checks + EXPECT_EQ(offsetof(ShmTocEntry, metric_name), 0); + EXPECT_EQ(offsetof(ShmTocEntry, encoded_labels), 64); + EXPECT_EQ(offsetof(ShmTocEntry, type), 192); + EXPECT_EQ(offsetof(ShmTocEntry, offset), 196); + EXPECT_EQ(offsetof(ShmTocEntry, size), 200); + EXPECT_EQ(offsetof(ShmTocEntry, entry_state), 204); + EXPECT_EQ(offsetof(ShmTocEntry, padding), 208); + EXPECT_EQ(sizeof(ShmTocEntry), 256); + + // ShmTocHeader offset checks + EXPECT_EQ(offsetof(ShmTocHeader, magic), 0); + EXPECT_EQ(offsetof(ShmTocHeader, version), 4); + EXPECT_EQ(offsetof(ShmTocHeader, pid), 8); + EXPECT_EQ(offsetof(ShmTocHeader, toc_entry_count), 16); + EXPECT_EQ(offsetof(ShmTocHeader, max_toc_entries), 20); + EXPECT_EQ(offsetof(ShmTocHeader, data_pool_offset), 24); + EXPECT_EQ(offsetof(ShmTocHeader, data_pool_bytes), 28); + EXPECT_EQ(offsetof(ShmTocHeader, chunk_index), 32); + EXPECT_EQ(offsetof(ShmTocHeader, padding), 36); + EXPECT_EQ(sizeof(ShmTocHeader), 64); + + // ShmHistogramSlot offset checks + EXPECT_EQ(offsetof(ShmHistogramSlot, sample_count), 0); + EXPECT_EQ(offsetof(ShmHistogramSlot, sample_sum), 8); + EXPECT_EQ(offsetof(ShmHistogramSlot, bucket_counts), 16); + EXPECT_EQ(offsetof(ShmHistogramSlot, padding), 184); + EXPECT_EQ(sizeof(ShmHistogramSlot), 192); + + // ShmSegmentLayout offset checks + EXPECT_EQ(offsetof(ShmSegmentLayout, header), 0); + EXPECT_EQ(offsetof(ShmSegmentLayout, toc), 64); + EXPECT_EQ(sizeof(ShmSegmentLayout), 262208); +} + +TEST(ShmLayoutTest, ConstantsAndOffsets) { + EXPECT_EQ(kRaidenShmMagic, 0xABCD1234); + EXPECT_EQ(kSupportedVersion, 2); + EXPECT_EQ(kShmFilePrefix, "worker_rank_"); + EXPECT_EQ(kShmFileExtension, ".mmap"); + EXPECT_EQ(kMaxTocEntries, 1024); + EXPECT_EQ(kMaxDataPoolBytes, 65536); + EXPECT_EQ(kMetricSlotAlignment, 64); + EXPECT_EQ(kNumHistogramBuckets, 20); + + size_t expected_total_size = sizeof(ShmTocHeader) + + (sizeof(ShmTocEntry) * kMaxTocEntries) + + kMaxDataPoolBytes; + EXPECT_EQ(kSegmentTotalFileSize, expected_total_size); + EXPECT_EQ(kSegmentTotalFileSize, 327744); + EXPECT_EQ(CalculateChunkFileSize(kMaxTocEntries, kMaxDataPoolBytes), + expected_total_size); +} + +TEST(ShmLayoutTest, TocEntryStateEnumValues) { + EXPECT_EQ(static_cast(TocEntryState::kUninitialized), 0); + EXPECT_EQ(static_cast(TocEntryState::kWriting), 1); + EXPECT_EQ(static_cast(TocEntryState::kCommitted), 2); +} + +TEST(ShmLayoutTest, HistogramSingleThreadObservation) { + ShmHistogramSlot slot; + EXPECT_EQ(slot.sample_count.load(), 0); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + + // Value 0.05 falls into bucket 0 (<= 0.1). + slot.Observe(0.05); + EXPECT_EQ(slot.sample_count.load(), 1); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.05)); + EXPECT_EQ(slot.bucket_counts[0].load(), 1); + for (size_t b = 1; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 0); + } + + // Value 5.5 falls into bucket 6 (<= 10.0, > 5.0). + slot.Observe(5.5); + EXPECT_EQ(slot.sample_count.load(), 2); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(5.55)); + EXPECT_EQ(slot.bucket_counts[0].load(), 1); + EXPECT_EQ(slot.bucket_counts[6].load(), 1); + for (size_t b = 1; b <= kNumHistogramBuckets; ++b) { + if (b != 6) { + EXPECT_EQ(slot.bucket_counts[b].load(), 0); + } + } + + // Value 60000.0 exceeds highest boundary (50000.0), falling into bucket 20 + // (+Inf). + slot.Observe(60000.0); + EXPECT_EQ(slot.sample_count.load(), 3); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(60005.55)); + EXPECT_EQ(slot.bucket_counts[0].load(), 1); + EXPECT_EQ(slot.bucket_counts[6].load(), 1); + EXPECT_EQ(slot.bucket_counts[kNumHistogramBuckets].load(), 1); + + // Cumulative view verification (as computed by exporter) + uint64_t cumulative_sum = 0; + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + cumulative_sum += slot.bucket_counts[b].load(); + if (b < 6) { + EXPECT_EQ(cumulative_sum, 1); + } else if (b < kNumHistogramBuckets) { + EXPECT_EQ(cumulative_sum, 2); + } else { + EXPECT_EQ(cumulative_sum, 3); + } + } +} + +TEST(ShmLayoutTest, HistogramNonFiniteObservationsIgnored) { + ShmHistogramSlot slot; + slot.Observe(std::numeric_limits::quiet_NaN()); + slot.Observe(std::numeric_limits::infinity()); + slot.Observe(-std::numeric_limits::infinity()); + + EXPECT_EQ(slot.sample_count.load(), 0); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 0); + } +} + +TEST(ShmLayoutTest, HistogramConcurrentObservations) { + ShmHistogramSlot slot; + constexpr int kNumThreads = 8; + constexpr int kNumIters = 1000; + constexpr double kObsValue = 0.5; + + std::vector threads; + threads.reserve(kNumThreads); + for (int t = 0; t < kNumThreads; ++t) { + threads.emplace_back([&slot]() { + for (int i = 0; i < kNumIters; ++i) { + slot.Observe(kObsValue); + } + }); + } + for (auto& th : threads) { + th.join(); + } + + EXPECT_EQ(slot.sample_count.load(), kNumThreads * kNumIters); + EXPECT_THAT(slot.sample_sum.load(), + DoubleEq(kNumThreads * kNumIters * kObsValue)); + // 0.5 falls into bucket 2 (0.25 < 0.5 <= 0.5). + EXPECT_EQ(slot.bucket_counts[0].load(), 0); + EXPECT_EQ(slot.bucket_counts[1].load(), 0); + EXPECT_EQ(slot.bucket_counts[2].load(), kNumThreads * kNumIters); + for (size_t b = 3; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 0); + } +} + +TEST(ShmLayoutTest, HistogramBucketSemantics) { + ShmHistogramSlot slot; + slot.Observe(0.1); // Exactly on bucket 0 boundary + slot.Observe(0.1000001); // Just above bucket 0 -> bucket 1 + slot.Observe(50000.0); // Exactly on bucket 19 boundary + slot.Observe(50000.0001); // Just above bucket 19 -> bucket 20 (+Inf) + + EXPECT_EQ(slot.sample_count.load(), 4); + EXPECT_EQ(slot.bucket_counts[0].load(), 1); + EXPECT_EQ(slot.bucket_counts[1].load(), 1); + EXPECT_EQ(slot.bucket_counts[19].load(), 1); + EXPECT_EQ(slot.bucket_counts[20].load(), 1); + + // Cumulative view is monotonic + uint64_t prev_cumulative = 0; + uint64_t cumulative = 0; + for (size_t i = 0; i <= kNumHistogramBuckets; ++i) { + cumulative += slot.bucket_counts[i].load(); + EXPECT_GE(cumulative, prev_cumulative); + prev_cumulative = cumulative; + } + EXPECT_EQ(cumulative, 4); +} + +TEST(ShmLayoutTest, AdversarialNonFiniteAndEdgeValues) { + ShmHistogramSlot slot; + + // 1. NaN variants + slot.Observe(std::numeric_limits::quiet_NaN()); + slot.Observe(std::numeric_limits::signaling_NaN()); + slot.Observe(-std::numeric_limits::quiet_NaN()); + + // 2. Infinities + slot.Observe(std::numeric_limits::infinity()); + slot.Observe(-std::numeric_limits::infinity()); + + EXPECT_EQ(slot.sample_count.load(), 0); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + EXPECT_EQ(slot.bucket_counts[b].load(), 0); + } + + // 3. Signed Zeros + slot.Observe(+0.0); + slot.Observe(-0.0); + EXPECT_EQ(slot.sample_count.load(), 2); + EXPECT_THAT(slot.sample_sum.load(), DoubleEq(0.0)); + // 0.0 <= 0.1 so placed in bucket 0 + EXPECT_EQ(slot.bucket_counts[0].load(), 2); + + // 4. Subnormals + slot.Observe(std::numeric_limits::denorm_min()); + EXPECT_EQ(slot.sample_count.load(), 3); + EXPECT_GT(slot.sample_sum.load(), 0.0); + EXPECT_EQ(slot.bucket_counts[0].load(), 3); + + // 5. Negative finite numbers + slot.Observe(-50.0); + EXPECT_EQ(slot.sample_count.load(), 4); + EXPECT_LT(slot.sample_sum.load(), 0.0); + // -50.0 <= 0.1, placed in bucket 0 + EXPECT_EQ(slot.bucket_counts[0].load(), 4); + + // 6. Extreme max finite number + ShmHistogramSlot max_slot; + max_slot.Observe(std::numeric_limits::max()); + EXPECT_EQ(max_slot.sample_count.load(), 1); + EXPECT_EQ(max_slot.sample_sum.load(), std::numeric_limits::max()); + // max exceeds 50000.0, placed in +Inf bucket (index 20) + for (size_t b = 0; b < kNumHistogramBuckets; ++b) { + EXPECT_EQ(max_slot.bucket_counts[b].load(), 0); + } + EXPECT_EQ(max_slot.bucket_counts[kNumHistogramBuckets].load(), 1); + + // 7. Overflow to +Inf in sum accumulation + max_slot.Observe(std::numeric_limits::max()); + EXPECT_EQ(max_slot.sample_count.load(), 2); + EXPECT_TRUE(std::isinf(max_slot.sample_sum.load())); +} + +TEST(ShmLayoutTest, HighContentionThreadScalingAndCASLivelock) { + ShmHistogramSlot slot; + constexpr int kNumThreads = 32; + constexpr int kNumIters = 25000; // Total 800,000 observations + std::atomic start_signal{false}; + + std::vector threads; + threads.reserve(kNumThreads); + + for (int t = 0; t < kNumThreads; ++t) { + threads.emplace_back([&slot, &start_signal, t]() { + while (!start_signal.load(std::memory_order_acquire)) { + } + for (int i = 0; i < kNumIters; ++i) { + double val = ((t * 17 + i) % 20) * 2.5 + 0.05; + slot.Observe(val); + } + }); + } + + start_signal.store(true, std::memory_order_release); + for (auto& th : threads) { + th.join(); + } + + EXPECT_EQ(slot.sample_count.load(), + static_cast(kNumThreads) * kNumIters); + + uint64_t total_bucket_samples = 0; + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + total_bucket_samples += slot.bucket_counts[b].load(); + } + EXPECT_EQ(total_bucket_samples, + static_cast(kNumThreads) * kNumIters); +} + +TEST(ShmLayoutTest, FloatingPointPrecisionAndAbsorptionStress) { + ShmHistogramSlot slot; + + // Case 1: Absorption under extreme dynamic range (1e16 + 1.0) + slot.Observe(1e16); + constexpr int kSmallAdds = 10000; + for (int i = 0; i < kSmallAdds; ++i) { + slot.Observe(1.0); + } + EXPECT_EQ(slot.sample_count.load(), 1 + kSmallAdds); + + // Case 2: Multi-precision comparison for normal dynamic range + ShmHistogramSlot normal_slot; + constexpr int kNormalIters = 100000; + for (int i = 0; i < kNormalIters; ++i) { + normal_slot.Observe(0.125); // Exact power of 2 fraction (1/8) + } + EXPECT_EQ(normal_slot.sample_count.load(), kNormalIters); + EXPECT_THAT(normal_slot.sample_sum.load(), DoubleEq(kNormalIters * 0.125)); +} + +TEST(ShmLayoutTest, ConcurrentReaderCumulativeMonotonicityRace) { + ShmHistogramSlot slot; + constexpr int kNumWriters = 8; + constexpr int kNumReaders = 4; + constexpr int kWritesPerThread = 50000; + std::atomic stop_readers{false}; + std::atomic start_signal{false}; + std::atomic non_monotonic_snapshots{0}; + std::atomic total_snapshots{0}; + + std::vector writers; + writers.reserve(kNumWriters); + for (int w = 0; w < kNumWriters; ++w) { + writers.emplace_back([&slot, &start_signal, w]() { + while (!start_signal.load(std::memory_order_acquire)) { + } + for (int i = 0; i < kWritesPerThread; ++i) { + double val = ((w * 31 + i) % 22) * 2500.0 + 0.05; + slot.Observe(val); + } + }); + } + + std::vector readers; + readers.reserve(kNumReaders); + for (int r = 0; r < kNumReaders; ++r) { + readers.emplace_back([&slot, &start_signal, &stop_readers, + &non_monotonic_snapshots, &total_snapshots]() { + while (!start_signal.load(std::memory_order_acquire)) { + } + uint64_t local_counts[kNumHistogramBuckets + 1]; + while (!stop_readers.load(std::memory_order_relaxed)) { + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + local_counts[b] = + slot.bucket_counts[b].load(std::memory_order_relaxed); + } + total_snapshots.fetch_add(1, std::memory_order_relaxed); + // Compute cumulative counts on-the-fly and check monotonicity + uint64_t prev_cum = 0; + uint64_t cum = 0; + bool monotonic = true; + for (size_t b = 0; b <= kNumHistogramBuckets; ++b) { + cum += local_counts[b]; + if (cum < prev_cum) { + monotonic = false; + break; + } + prev_cum = cum; + } + if (!monotonic) { + non_monotonic_snapshots.fetch_add(1, std::memory_order_relaxed); + } + } + }); + } + + start_signal.store(true, std::memory_order_release); + for (auto& w : writers) { + w.join(); + } + stop_readers.store(true, std::memory_order_release); + for (auto& r : readers) { + r.join(); + } + + EXPECT_EQ(non_monotonic_snapshots.load(), 0); + EXPECT_GT(total_snapshots.load(), 0); +} + +TEST(ShmLayoutTest, BinarySerializationAndNoGaps) { + alignas(64) uint8_t raw_buffer[sizeof(ShmTocEntry)]; + std::memset(raw_buffer, 0xAA, sizeof(raw_buffer)); + + auto* entry = reinterpret_cast(raw_buffer); + std::memset(entry->metric_name, 'A', sizeof(entry->metric_name)); + entry->metric_name[sizeof(entry->metric_name) - 1] = '\0'; + std::memset(entry->encoded_labels, 'B', sizeof(entry->encoded_labels)); + entry->encoded_labels[sizeof(entry->encoded_labels) - 1] = '\0'; + entry->type = MetricType::kHistogram; + entry->offset = 0x12345678; + entry->size = 0x87654321; + entry->entry_state.store(TocEntryState::kCommitted); + + EXPECT_EQ(raw_buffer[0], 'A'); + EXPECT_EQ(raw_buffer[63], '\0'); + EXPECT_EQ(raw_buffer[64], 'B'); + EXPECT_EQ(raw_buffer[191], '\0'); + EXPECT_EQ(*reinterpret_cast(&raw_buffer[192]), + MetricType::kHistogram); + EXPECT_EQ(*reinterpret_cast(&raw_buffer[196]), 0x12345678); + EXPECT_EQ(*reinterpret_cast(&raw_buffer[200]), 0x87654321); + EXPECT_EQ(*reinterpret_cast(&raw_buffer[204]), + TocEntryState::kCommitted); + + EXPECT_EQ(sizeof(ShmHistogramSlot), 8 + 8 + (21 * 8) + 8); +} + +TEST(ShmLayoutTest, CacheLineBoundaryAndFalseSharingAnalysis) { + EXPECT_EQ(alignof(ShmTocHeader), 64); + EXPECT_EQ(sizeof(ShmTocHeader), 64); + EXPECT_EQ(sizeof(ShmTocHeader) % 64, 0); + + EXPECT_EQ(alignof(ShmSegmentLayout), 64); + EXPECT_EQ(sizeof(ShmSegmentLayout), 262208); + EXPECT_EQ(sizeof(ShmSegmentLayout) % 64, 0); + + EXPECT_EQ(kSegmentTotalFileSize, 327744); + EXPECT_EQ(kSegmentTotalFileSize % 64, 0); + EXPECT_EQ(kMetricSlotAlignment, 64); +} + +TEST(ShmLayoutTest, FullSegmentArrayNoOverlap) { + auto segment = std::make_unique(); + std::memset(segment.get(), 0, sizeof(ShmSegmentLayout)); + + segment->header.magic.store(kRaidenShmMagic); + segment->header.version = kSupportedVersion; + segment->header.pid = 123456; + segment->header.toc_entry_count = kMaxTocEntries; + segment->header.max_toc_entries = kMaxTocEntries; + segment->header.data_pool_offset = sizeof(ShmSegmentLayout); + segment->header.data_pool_bytes = kMaxDataPoolBytes; + + for (size_t i = 0; i < kMaxTocEntries; ++i) { + snprintf(segment->toc[i].metric_name, sizeof(segment->toc[i].metric_name), + "metric_name_stream_%04zu", i); + snprintf(segment->toc[i].encoded_labels, + sizeof(segment->toc[i].encoded_labels), + "cluster=tpuv4,task=%04zu,direction=push", i); + segment->toc[i].type = + (i % 2 == 0) ? MetricType::kCounter : MetricType::kHistogram; + segment->toc[i].offset = + static_cast(segment->header.data_pool_offset + (i * 192)); + segment->toc[i].size = (i % 2 == 0) ? 8 : 192; + segment->toc[i].entry_state.store(TocEntryState::kCommitted); + } + + EXPECT_EQ(segment->header.magic.load(), kRaidenShmMagic); + EXPECT_EQ(segment->header.version, kSupportedVersion); + EXPECT_EQ(segment->header.pid, 123456); + EXPECT_EQ(segment->header.toc_entry_count, kMaxTocEntries); + + for (size_t i = 0; i < kMaxTocEntries; ++i) { + char expected_name[64]; + char expected_labels[128]; + snprintf(expected_name, sizeof(expected_name), "metric_name_stream_%04zu", + i); + snprintf(expected_labels, sizeof(expected_labels), + "cluster=tpuv4,task=%04zu,direction=push", i); + + EXPECT_STREQ(segment->toc[i].metric_name, expected_name); + EXPECT_STREQ(segment->toc[i].encoded_labels, expected_labels); + EXPECT_EQ(segment->toc[i].offset, + segment->header.data_pool_offset + (i * 192)); + EXPECT_EQ(segment->toc[i].size, (i % 2 == 0) ? 8 : 192); + EXPECT_EQ(segment->toc[i].entry_state.load(), TocEntryState::kCommitted); + } +} + +} // 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..baefddbb --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_writer.cc @@ -0,0 +1,336 @@ +// 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 // 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) { + LOG(ERROR) << "ShmWriter failed to open temporary file " << tmp_path << ": " + << std::strerror(errno); + return false; + } + + if (ftruncate(fd, kSegmentTotalFileSize) != 0) { + LOG(ERROR) << "ShmWriter failed to ftruncate file " << tmp_path << ": " + << std::strerror(errno); + 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) { + LOG(ERROR) << "ShmWriter failed to mmap file " << tmp_path << ": " + << std::strerror(errno); + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + auto* segment = reinterpret_cast(addr); + std::memset(segment, 0, kSegmentTotalFileSize); + + ShmTocHeader& header = segment->header; + header.version = kSupportedVersion; + 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) { + LOG(ERROR) << "ShmWriter failed to flock file " << tmp_path << ": " + << std::strerror(errno); + munmap(addr, kSegmentTotalFileSize); + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + if (rename(tmp_path.c_str(), path.c_str()) != 0) { + LOG(ERROR) << "ShmWriter failed to rename file from " << tmp_path << " to " + << path << ": " << std::strerror(errno); + munmap(addr, kSegmentTotalFileSize); + close(fd); + unlink(tmp_path.c_str()); + return false; + } + + chunks_.push_back({fd, segment, path, chunk_idx}); + return true; +} + +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}; + + { + absl::ReaderMutexLock read_lock(mutex_); + const auto& cache = GetCache(); + auto it = cache.find(key); + if (it != cache.end()) return it->second; + } + + absl::MutexLock write_lock(mutex_); + auto& cache = GetCache(); + auto it = cache.find(key); + if (it != cache.end()) return it->second; + + if (chunks_.empty()) return nullptr; + + if (name.size() >= sizeof(ShmTocEntry::metric_name)) { + LOG(ERROR) << "Metric name rejected: " << name << " exceeds buffer size " + << sizeof(ShmTocEntry::metric_name); + return nullptr; + } + 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 nullptr; + } + + 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); + + cache.emplace(std::move(key), slot); + return slot; +} + +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..45383617 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_writer.h @@ -0,0 +1,169 @@ +// 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: + friend class ShmWriterTest; + + 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; + }; + + 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; + }; + + struct MetricKeyHash { + using is_transparent = void; + template + size_t operator()(const T& k) const { + return absl::HashOf(k.name, k.encoded_labels); + } + }; + + 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_; + } + } + + 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..147e1674 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_writer_test.cc @@ -0,0 +1,695 @@ +// 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 // NOLINT(build/c++17) +#include +#include +#include // NOLINT(build/c++11) +#include +#include + +#include +#include +#include "absl/strings/match.h" +#include "absl/strings/str_cat.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 { + +class ShmWriterTest : public testing::Test { + protected: + void SetUp() override { + test_dir_ = absl::StrCat(testing::TempDir(), "/shm_writer_test_", getpid()); + std::filesystem::create_directories(test_dir_); + } + + void TearDown() override { std::filesystem::remove_all(test_dir_); } + + ShmWriterOptions DefaultOptions(absl::string_view rank = "0") const { + return ShmWriterOptions{.shm_dir = test_dir_, + .local_rank = std::string(rank)}; + } + + static bool IsValid(const ShmWriter& writer) { + absl::MutexLock lock(writer.mutex_); + return !writer.chunks_.empty() && writer.chunks_[0].segment != nullptr; + } + + static const ShmSegmentLayout* GetSegment(const ShmWriter& writer) { + absl::MutexLock lock(writer.mutex_); + return writer.chunks_.empty() ? nullptr : writer.chunks_[0].segment; + } + + static std::string GetFilePath(const ShmWriter& writer) { + absl::MutexLock lock(writer.mutex_); + return writer.chunks_.empty() ? "" : writer.chunks_[0].file_path; + } + + static std::string GetUuid(const ShmWriter& writer) { return writer.uuid_; } + + static int GetFd(const ShmWriter& writer) { + absl::MutexLock lock(writer.mutex_); + return writer.chunks_.empty() ? -1 : writer.chunks_[0].fd; + } + + static const ShmTocEntry* FindTocEntry(const ShmSegmentLayout* segment, + absl::string_view name, + LabelSpan labels = {}) { + if (!segment) return nullptr; + std::string encoded = EncodeLabels(labels); + uint32_t count = + segment->header.toc_entry_count.load(std::memory_order_acquire); + for (uint32_t i = 0; i < count; ++i) { + if (segment->toc[i].metric_name == name) { + if (labels.empty() || segment->toc[i].encoded_labels == encoded) { + return &segment->toc[i]; + } + } + } + return nullptr; + } + + template + static const T* ReadSlot(const ShmWriter& writer, absl::string_view name, + LabelSpan labels = {}) { + const ShmSegmentLayout* segment = GetSegment(writer); + const ShmTocEntry* entry = FindTocEntry(segment, name, labels); + if (!entry) return nullptr; + return reinterpret_cast( + reinterpret_cast(segment) + entry->offset); + } + + static uint64_t ReadCounter(const ShmWriter& writer, absl::string_view name, + LabelSpan labels = {}) { + const std::atomic* slot = + ReadSlot>(writer, name, labels); + return slot ? slot->load(std::memory_order_relaxed) : 0; + } + + static double ReadGauge(const ShmWriter& writer, absl::string_view name, + LabelSpan labels = {}) { + const std::atomic* slot = + ReadSlot>(writer, name, labels); + return slot ? slot->load(std::memory_order_relaxed) : 0.0; + } + + static const ShmHistogramSlot* ReadHistogram(const ShmWriter& writer, + absl::string_view name, + LabelSpan labels = {}) { + return ReadSlot(writer, name, labels); + } + + std::string test_dir_; +}; + +namespace { + +using ::testing::DoubleEq; +using ::testing::IsNull; +using ::testing::NotNull; + +TEST_F(ShmWriterTest, InitializesEmptyTOCAndExpandsDynamically) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + EXPECT_EQ(GetFilePath(writer), + absl::StrCat(test_dir_, "/", kShmFilePrefix, "0_", GetUuid(writer), + "_chunk_0", kShmFileExtension)); + + const ShmSegmentLayout* segment = GetSegment(writer); + ASSERT_THAT(segment, NotNull()); + + EXPECT_EQ(segment->header.magic.load(), kRaidenShmMagic); + EXPECT_EQ(segment->header.version, kSupportedVersion); + 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); + + // Dynamic allocation on demand. + MetricLabel sent_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + MetricLabel fail_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "INTERNAL"}}; + + writer.IncrementCounter(metric_names::kSentBytesTotal, {&sent_label, 1}, 100); + writer.IncrementCounter(metric_names::kTransferFailuresTotal, fail_labels, 1); + writer.SetGauge(metric_names::kBufferAllocatedBytes, {}, 5.0); + + EXPECT_EQ(segment->header.toc_entry_count.load(), 3); + EXPECT_GT(segment->header.data_pool_bytes.load(), 0); + + EXPECT_EQ(segment->toc[0].type, MetricType::kCounter); + EXPECT_EQ(absl::string_view(segment->toc[0].metric_name), + metric_names::kSentBytesTotal); + EXPECT_TRUE( + absl::StrContains(segment->toc[0].encoded_labels, "direction=push")); + + EXPECT_EQ(segment->toc[1].type, MetricType::kCounter); + EXPECT_EQ(absl::string_view(segment->toc[1].metric_name), + metric_names::kTransferFailuresTotal); + EXPECT_TRUE( + absl::StrContains(segment->toc[1].encoded_labels, "direction=pull")); + EXPECT_TRUE( + absl::StrContains(segment->toc[1].encoded_labels, "error_code=INTERNAL")); + + EXPECT_EQ(segment->toc[2].type, MetricType::kGauge); + EXPECT_EQ(absl::string_view(segment->toc[2].metric_name), + metric_names::kBufferAllocatedBytes); +} + +TEST_F(ShmWriterTest, DataPoolOffsetMonotonicityAndContiguity) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + MetricLabel push{metric_labels::kDirection, metric_labels::kDirectionPush}; + MetricLabel pull{metric_labels::kDirection, metric_labels::kDirectionPull}; + writer.IncrementCounter(metric_names::kSentBytesTotal, {&push, 1}, 1); + writer.IncrementCounter(metric_names::kSentBytesTotal, {&pull, 1}, 2); + writer.ObserveHistogram(metric_names::kTransferDurationMs, {}, 0.05); + + const ShmSegmentLayout* segment = GetSegment(writer); + ASSERT_THAT(segment, NotNull()); + EXPECT_EQ(segment->header.toc_entry_count.load(), 3); + + uint32_t expected_offset = sizeof(ShmSegmentLayout); + for (uint32_t i = 0; i < segment->header.toc_entry_count.load(); ++i) { + const ShmTocEntry& entry = segment->toc[i]; + uint32_t align = kMetricSlotAlignment; + expected_offset = (expected_offset + align - 1) & ~(align - 1); + EXPECT_EQ(entry.offset, expected_offset); + EXPECT_GT(entry.size, 0); + expected_offset += entry.size; + } +} + +TEST_F(ShmWriterTest, PointerCacheResolutionAndMultiLabelUpdates) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + MetricLabel push{metric_labels::kDirection, metric_labels::kDirectionPush}; + MetricLabel pull{metric_labels::kDirection, metric_labels::kDirectionPull}; + + writer.IncrementCounter(metric_names::kSentBytesTotal, {&push, 1}, 42); + writer.IncrementCounter(metric_names::kSentBytesTotal, {&pull, 1}, 100); + + EXPECT_EQ(GetSegment(writer)->header.toc_entry_count.load(), 2); + EXPECT_EQ(ReadCounter(writer, metric_names::kSentBytesTotal, {&push, 1}), 42); + EXPECT_EQ(ReadCounter(writer, metric_names::kSentBytesTotal, {&pull, 1}), + 100); + + // Subsequent updates to cached entries do not allocate new TOC slots. + writer.IncrementCounter(metric_names::kSentBytesTotal, {&push, 1}, 58); + EXPECT_EQ(GetSegment(writer)->header.toc_entry_count.load(), 2); + EXPECT_EQ(ReadCounter(writer, metric_names::kSentBytesTotal, {&push, 1}), + 100); + EXPECT_EQ(ReadCounter(writer, metric_names::kSentBytesTotal, {&pull, 1}), + 100); +} + +TEST_F(ShmWriterTest, HoldsAdvisorySharedFlockAndReleasesOnDestruction) { + std::string path; + { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + path = GetFilePath(writer); + EXPECT_TRUE(std::filesystem::exists(path)); + + // Verify O_CLOEXEC flag on writer descriptor. + int flags = fcntl(GetFd(writer), F_GETFD); + EXPECT_TRUE(flags & FD_CLOEXEC); + + // Open duplicate descriptor to probe advisory locking. + int probe_fd = open(path.c_str(), O_RDWR | O_CLOEXEC); + ASSERT_GE(probe_fd, 0); + + // Exclusive lock must fail with EWOULDBLOCK while writer is active. + EXPECT_EQ(flock(probe_fd, LOCK_EX | LOCK_NB), -1); + EXPECT_EQ(errno, EWOULDBLOCK); + + // Concurrent shared lock must succeed. + EXPECT_EQ(flock(probe_fd, LOCK_SH | LOCK_NB), 0); + flock(probe_fd, LOCK_UN); + close(probe_fd); + } + + // File remains on disk after destruction for ShmCollector dead reaping. + EXPECT_TRUE(std::filesystem::exists(path)); + + // Lock must now be completely free, allowing exclusive acquisition. + int dead_fd = open(path.c_str(), O_RDWR | O_CLOEXEC); + ASSERT_GE(dead_fd, 0); + EXPECT_EQ(flock(dead_fd, LOCK_EX | LOCK_NB), 0); + flock(dead_fd, LOCK_UN); + close(dead_fd); +} + +TEST_F(ShmWriterTest, RecordsMetricsWithMixedLabelOrdering) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + MetricLabel labels_canonical[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "INTERNAL"}}; + MetricLabel labels_reversed[] = { + {metric_labels::kErrorCode, "INTERNAL"}, + {metric_labels::kDirection, metric_labels::kDirectionPull}}; + + writer.IncrementCounter(metric_names::kTransferFailuresTotal, + labels_canonical, 10); + writer.IncrementCounter(metric_names::kTransferFailuresTotal, labels_reversed, + 15); + + EXPECT_EQ(GetSegment(writer)->header.toc_entry_count.load(), 1); + EXPECT_EQ(ReadCounter(writer, metric_names::kTransferFailuresTotal), 25); +} + +TEST_F(ShmWriterTest, DynamicMetricsAndLabelsAllocatedSafely) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + MetricLabel label{"custom_key", "custom_val"}; + writer.IncrementCounter("custom_metric", {&label, 1}, 100); + writer.SetGauge("custom_gauge", {&label, 1}, 3.14); + writer.ObserveHistogram("custom_hist", {&label, 1}, 1.5); + + EXPECT_EQ(ReadCounter(writer, "custom_metric"), 100); + EXPECT_THAT(ReadGauge(writer, "custom_gauge"), DoubleEq(3.14)); + const ShmHistogramSlot* histogram_slot = ReadHistogram(writer, "custom_hist"); + ASSERT_THAT(histogram_slot, NotNull()); + EXPECT_EQ(histogram_slot->sample_count.load(), 1); +} + +TEST(ShmWriterStaticTest, NonCopyableAndNonMovable) { + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(!std::is_move_constructible_v); + static_assert(!std::is_move_assignable_v); +} + +TEST_F(ShmWriterTest, ConcurrentMultiThreadedIncrementCounter) { + constexpr int kNumThreads = 16; + constexpr int kItersPerThread = 25000; + + MetricLabel sent_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + MetricLabel fail_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "RESOURCE_EXHAUSTED"}}; + + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + std::vector workers; + workers.reserve(kNumThreads); + for (int t = 0; t < kNumThreads; ++t) { + workers.emplace_back([&writer, &sent_label, &fail_labels]() { + for (int j = 0; j < kItersPerThread; ++j) { + writer.IncrementCounter(metric_names::kSentBytesTotal, {&sent_label, 1}, + 1); + writer.IncrementCounter(metric_names::kTransferFailuresTotal, + fail_labels, 2); + } + }); + } + + for (auto& w : workers) { + w.join(); + } + + EXPECT_EQ(ReadCounter(writer, metric_names::kSentBytesTotal), + static_cast(kNumThreads) * kItersPerThread); + EXPECT_EQ(ReadCounter(writer, metric_names::kTransferFailuresTotal), + static_cast(kNumThreads) * kItersPerThread * 2); +} + +TEST_F(ShmWriterTest, ConcurrentMultiThreadedMixedWorkloadStress) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + constexpr int kNumThreads = 32; + constexpr int kIterations = 10000; + + MetricLabel push_label{metric_labels::kDirection, + metric_labels::kDirectionPush}; + MetricLabel pull_label{metric_labels::kDirection, + metric_labels::kDirectionPull}; + MetricLabel fail_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPullResponse}, + {metric_labels::kErrorCode, "UNAVAILABLE"}}; + + std::vector workers; + workers.reserve(kNumThreads); + for (int t = 0; t < kNumThreads; ++t) { + workers.emplace_back([&writer, &push_label, &pull_label, &fail_labels, + t]() { + for (int i = 0; i < kIterations; ++i) { + int op = (t + i) % 4; + switch (op) { + 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; + case 3: + writer.IncrementCounter("dynamic_metric", {&pull_label, 1}, 100); + writer.SetGauge("dynamic_metric", {&pull_label, 1}, 3.14); + writer.ObserveHistogram("dynamic_metric", {&pull_label, 1}, 1.0); + break; + } + } + }); + } + + for (auto& w : workers) { + w.join(); + } + + EXPECT_GT(ReadCounter(writer, metric_names::kSentBytesTotal), 0); + EXPECT_GT(ReadCounter(writer, metric_names::kTransferFailuresTotal), 0); + EXPECT_GT(ReadCounter(writer, metric_names::kReceivedBytesTotal), 0); + EXPECT_GT(ReadCounter(writer, "dynamic_metric"), 0); +} + +TEST_F(ShmWriterTest, BoundaryLabelsAndAdversarialCases) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + // Empty label span allocates valid slot. + writer.IncrementCounter(metric_names::kSentBytesTotal, {}, 1); + + // Exactly 1 label (dedicated fast-path branch in EncodeLabels). + MetricLabel label_1[1] = {{"k1", "v1"}}; + writer.IncrementCounter(metric_names::kSentBytesTotal, {label_1, 1}, 2); + + // Exactly 8 labels (stack array upper bound in EncodeLabels). + MetricLabel labels_8[8] = {{"k1", "v1"}, {"k2", "v2"}, {"k3", "v3"}, + {"k4", "v4"}, {"k5", "v5"}, {"k6", "v6"}, + {"k7", "v7"}, {"k8", "v8"}}; + writer.IncrementCounter(metric_names::kSentBytesTotal, labels_8, 3); + + // > 8 labels (dynamically allocated vector fallback in EncodeLabels). + MetricLabel labels_20[20] = { + {"l01", "1"}, {"l02", "2"}, {"l03", "3"}, {"l04", "4"}, {"l05", "5"}, + {"l06", "6"}, {"l07", "7"}, {"l08", "8"}, {"l09", "9"}, {"l10", "0"}, + {"l11", "1"}, {"l12", "2"}, {"l13", "3"}, {"l14", "4"}, {"l15", "5"}, + {"l16", "6"}, {"l17", "7"}, {"l18", "8"}, {"l19", "9"}, {"l20", "0"}}; + writer.IncrementCounter(metric_names::kSentBytesTotal, labels_20, 4); + + // Verify memory boundaries for all registered slots + const ShmSegmentLayout* segment = GetSegment(writer); + ASSERT_THAT(segment, NotNull()); + const uint8_t* seg_start = reinterpret_cast(segment); + const uint8_t* seg_end = seg_start + kSegmentTotalFileSize; + + EXPECT_EQ(segment->header.toc_entry_count.load(), 4); + for (uint32_t i = 0; i < segment->header.toc_entry_count.load(); ++i) { + const ShmTocEntry& entry = segment->toc[i]; + EXPECT_GE(entry.offset, sizeof(ShmSegmentLayout)); + EXPECT_LT(entry.offset + entry.size, kSegmentTotalFileSize); + const uint8_t* slot_addr = seg_start + entry.offset; + EXPECT_GE(slot_addr, seg_start); + EXPECT_LE(slot_addr + entry.size, seg_end); + } +} + +TEST_F(ShmWriterTest, EscapedDelimiterLabelsAndCollisionFreeTOC) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + // Case 1: Label value containing ';' and '=' + MetricLabel label1[1] = {{"k", "v1;k2=v2"}}; + // Case 2: Two discrete labels that would collide with unescaped + // representation + MetricLabel label2[2] = {{"k", "v1"}, {"k2", "v2"}}; + + writer.IncrementCounter("delim_test", {label1, 1}, 10); + writer.IncrementCounter("delim_test", label2, 20); + + const ShmSegmentLayout* segment = GetSegment(writer); + ASSERT_THAT(segment, NotNull()); + EXPECT_EQ(segment->header.toc_entry_count.load(), 2); + + EXPECT_EQ(ReadCounter(writer, "delim_test", {label1, 1}), 10); + EXPECT_EQ(ReadCounter(writer, "delim_test", label2), 20); +} + +TEST_F(ShmWriterTest, OversizedMetricNameAndLabelsRejection) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + std::string huge_name(70, 'x'); // exceeds 64-byte buffer + writer.IncrementCounter(huge_name, {}, 1); + + const ShmSegmentLayout* segment = GetSegment(writer); + ASSERT_THAT(segment, NotNull()); + EXPECT_EQ(segment->header.toc_entry_count.load(), 0); + + std::string huge_label_val(150, 'y'); // exceeds 128-byte buffer + MetricLabel huge_label[1] = {{"key", huge_label_val}}; + writer.IncrementCounter("valid_name", {huge_label, 1}, 1); + EXPECT_EQ(segment->header.toc_entry_count.load(), 0); +} + +TEST_F(ShmWriterTest, FlockLivenessAcrossMetricMutations) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + const std::string path = GetFilePath(writer); + + int probe_fd = open(path.c_str(), O_RDWR | O_CLOEXEC); + ASSERT_GE(probe_fd, 0); + + // Probe: Exclusive lock must fail immediately with EWOULDBLOCK. + EXPECT_EQ(flock(probe_fd, LOCK_EX | LOCK_NB), -1); + EXPECT_EQ(errno, EWOULDBLOCK); + + // Perform continuous mutations while verifying lock remains held. + MetricLabel label{metric_labels::kDirection, metric_labels::kDirectionPush}; + for (int i = 0; i < 500; ++i) { + writer.IncrementCounter(metric_names::kSentBytesTotal, {&label, 1}, 10); + if (i % 50 == 0) { + EXPECT_EQ(flock(probe_fd, LOCK_EX | LOCK_NB), -1); + EXPECT_EQ(errno, EWOULDBLOCK); + EXPECT_EQ(flock(probe_fd, LOCK_SH | LOCK_NB), 0); + } + } + + close(probe_fd); +} + +TEST_F(ShmWriterTest, CrossProcessExclusiveLockContentionFork) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + const std::string path = GetFilePath(writer); + + pid_t pid = fork(); + ASSERT_GE(pid, 0); + + if (pid == 0) { + int child_fd = open(path.c_str(), O_RDWR | O_CLOEXEC); + if (child_fd < 0) _exit(1); + if (flock(child_fd, LOCK_EX | LOCK_NB) != -1 || errno != EWOULDBLOCK) { + close(child_fd); + _exit(2); + } + if (flock(child_fd, LOCK_SH | LOCK_NB) != 0) { + close(child_fd); + _exit(3); + } + flock(child_fd, LOCK_UN); + close(child_fd); + _exit(0); + } + + int status = 0; + ASSERT_EQ(waitpid(pid, &status, 0), pid); + EXPECT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 0); +} + +TEST_F(ShmWriterTest, CrossProcessLockReleaseOnProcessExit) { + pid_t pid = fork(); + ASSERT_GE(pid, 0); + + if (pid == 0) { + ShmWriter child_writer(DefaultOptions()); + if (!IsValid(child_writer)) _exit(1); + MetricLabel label{metric_labels::kDirection, metric_labels::kDirectionPush}; + child_writer.IncrementCounter(metric_names::kSentBytesTotal, {&label, 1}, + 777); + _exit(0); + } + + int status = 0; + ASSERT_EQ(waitpid(pid, &status, 0), pid); + EXPECT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 0); + + std::string expected_file; + for (const auto& entry : std::filesystem::directory_iterator(test_dir_)) { + if (absl::StrContains(entry.path().string(), "worker_rank_0_") && + absl::EndsWith(entry.path().string(), "_chunk_0.mmap")) { + expected_file = entry.path().string(); + break; + } + } + ASSERT_FALSE(expected_file.empty()); + EXPECT_TRUE(std::filesystem::exists(expected_file)); + + // The OS kernel automatically releases advisory locks on child exit. + int dead_fd = open(expected_file.c_str(), O_RDWR | O_CLOEXEC); + ASSERT_GE(dead_fd, 0); + EXPECT_EQ(flock(dead_fd, LOCK_EX | LOCK_NB), 0); + flock(dead_fd, LOCK_UN); + close(dead_fd); +} + +TEST_F(ShmWriterTest, InvalidDoubleValuesNaNAndInfResilience) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + writer.ObserveHistogram("hist", {}, std::numeric_limits::quiet_NaN()); + writer.ObserveHistogram("hist", {}, std::numeric_limits::infinity()); + writer.ObserveHistogram("hist", {}, -std::numeric_limits::infinity()); + + const ShmHistogramSlot* slot = ReadHistogram(writer, "hist"); + ASSERT_THAT(slot, NotNull()); + EXPECT_EQ(slot->sample_count.load(), 0); + EXPECT_THAT(slot->sample_sum.load(), DoubleEq(0.0)); + + writer.ObserveHistogram("hist", {}, 10.5); + EXPECT_EQ(slot->sample_count.load(), 1); + EXPECT_THAT(slot->sample_sum.load(), DoubleEq(10.5)); + + // Subnormal / extreme finite value resilience. + writer.ObserveHistogram("hist", {}, + std::numeric_limits::denorm_min()); + EXPECT_EQ(slot->sample_count.load(), 2); + + // Gauge ignores non-finite values. + writer.SetGauge("gauge_test", {}, 42.0); + EXPECT_THAT(ReadGauge(writer, "gauge_test"), DoubleEq(42.0)); + writer.SetGauge("gauge_test", {}, std::numeric_limits::quiet_NaN()); + EXPECT_THAT(ReadGauge(writer, "gauge_test"), DoubleEq(42.0)); + writer.SetGauge("gauge_test", {}, std::numeric_limits::infinity()); + EXPECT_THAT(ReadGauge(writer, "gauge_test"), DoubleEq(42.0)); + writer.SetGauge("gauge_test", {}, -std::numeric_limits::infinity()); + EXPECT_THAT(ReadGauge(writer, "gauge_test"), DoubleEq(42.0)); +} + +TEST_F(ShmWriterTest, InitializationFailureGracefulDegradation) { + // 1. Empty shm_dir + ShmWriter empty_dir(ShmWriterOptions{.shm_dir = "", .local_rank = "0"}); + EXPECT_FALSE(IsValid(empty_dir)); + EXPECT_THAT(GetSegment(empty_dir), IsNull()); + + // 2. Empty local_rank + ShmWriter empty_rank( + ShmWriterOptions{.shm_dir = test_dir_, .local_rank = ""}); + EXPECT_FALSE(IsValid(empty_rank)); + EXPECT_THAT(GetSegment(empty_rank), IsNull()); + + // 3. Impossible directory path + std::string regular_file = absl::StrCat(test_dir_, "/regular_file"); + int fd = open(regular_file.c_str(), O_CREAT | O_WRONLY, 0644); + ASSERT_GE(fd, 0); + close(fd); + + ShmWriter invalid_dir_shm_writer( + ShmWriterOptions{.shm_dir = absl::StrCat(regular_file, "/impossible_dir"), + .local_rank = "0"}); + EXPECT_FALSE(IsValid(invalid_dir_shm_writer)); + EXPECT_EQ(GetFd(invalid_dir_shm_writer), -1); + EXPECT_THAT(GetSegment(invalid_dir_shm_writer), IsNull()); + + // Operations on invalid writer must safely no-op without crash. + MetricLabel label{metric_labels::kDirection, metric_labels::kDirectionPush}; + invalid_dir_shm_writer.IncrementCounter(metric_names::kSentBytesTotal, + {&label, 1}, 10); + invalid_dir_shm_writer.SetGauge("any_gauge", {&label, 1}, 1.0); + invalid_dir_shm_writer.ObserveHistogram("any_hist", {&label, 1}, 2.0); +} + +TEST_F(ShmWriterTest, RapidConstructionDestructionCycle) { + MetricLabel label{metric_labels::kDirection, metric_labels::kDirectionPush}; + for (int i = 0; i < 50; ++i) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + writer.IncrementCounter(metric_names::kSentBytesTotal, {&label, 1}, i); + } +} + +TEST_F(ShmWriterTest, DynamicMultiChunkExpansionOnOverflow) { + ShmWriter writer(DefaultOptions()); + ASSERT_TRUE(IsValid(writer)); + + std::string chunk0_path = GetFilePath(writer); + std::string chunk1_path = + absl::StrCat(test_dir_, "/", kShmFilePrefix, "0", "_", GetUuid(writer), + "_chunk_1", kShmFileExtension); + + EXPECT_TRUE(std::filesystem::exists(chunk0_path)); + EXPECT_FALSE(std::filesystem::exists(chunk1_path)); + + // Allocate kMaxTocEntries + 5 unique metrics to force chunk 0 exhaustion and + // allocate chunk 1. + for (size_t i = 0; i < kMaxTocEntries + 5; ++i) { + std::string name = absl::StrCat("metric_", i); + writer.IncrementCounter(name, {}, static_cast(i + 1)); + } + + EXPECT_TRUE(std::filesystem::exists(chunk1_path)); + EXPECT_EQ(ReadCounter(writer, "metric_0"), 1); + + writer.IncrementCounter("metric_0", {}, 10); + EXPECT_EQ(ReadCounter(writer, "metric_0"), 11); +} + +TEST_F(ShmWriterTest, CustomFileModePermissions) { + ShmWriter writer(ShmWriterOptions{ + .shm_dir = test_dir_, .local_rank = "0", .file_mode = 0600}); + ASSERT_TRUE(IsValid(writer)); + std::string path = GetFilePath(writer); + + struct stat st; + ASSERT_EQ(stat(path.c_str(), &st), 0); + EXPECT_EQ(st.st_mode & 0777, 0600); +} + +} // namespace +} // namespace tpu_raiden::telemetry