diff --git a/tpu_sync/telemetry/BUILD b/tpu_sync/telemetry/BUILD index 7bdae27a..c75f0129 100644 --- a/tpu_sync/telemetry/BUILD +++ b/tpu_sync/telemetry/BUILD @@ -161,3 +161,32 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "base_shm_exporter", + srcs = ["base_shm_exporter.cc"], + hdrs = ["base_shm_exporter.h"], + deps = [ + ":metrics_backend", + "//tpu_sync/telemetry/shm:shm_collector", + "//tpu_sync/telemetry/shm:shm_writer", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "base_shm_exporter_test", + srcs = ["base_shm_exporter_test.cc"], + deps = [ + ":base_shm_exporter", + ":metrics_backend", + "//tpu_sync/telemetry/shm:shm_collector", + "//tpu_sync/telemetry/shm:shm_layout", + "@com_google_absl//absl/cleanup", + "@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..8f83b60e --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter.cc @@ -0,0 +1,118 @@ +// 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/algorithm/container.h" +#include "absl/container/inlined_vector.h" +#include "absl/log/check.h" +#include "absl/strings/numbers.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_collector.h" +#include "tpu_sync/telemetry/shm/shm_writer.h" + +namespace tpu_raiden::telemetry { +namespace { + +// Validates that `options.local_rank` is a valid non-negative integer +// (canonicalizing any leading zeros) and that `options.shm_dir` is an absolute, +// non-empty directory path with trailing slashes stripped, rejecting the root +// path "/". +ExporterOptions NormalizeAndValidateOptions(ExporterOptions options) { + CHECK(options.local_rank.has_value() && !options.local_rank->empty()) + << "options.local_rank must be specified and non-empty for " + "BaseShmExporter"; + + int rank_val = -1; + CHECK(absl::SimpleAtoi(*options.local_rank, &rank_val) && rank_val >= 0) + << "options.local_rank must be a valid non-negative integer for " + "BaseShmExporter, got: '" + << *options.local_rank << "'"; + options.local_rank = std::to_string(rank_val); + + while (options.shm_dir.has_value() && options.shm_dir->ends_with("/")) { + options.shm_dir->pop_back(); + } + CHECK(options.shm_dir.has_value() && !options.shm_dir->empty()) + << "options.shm_dir must be specified and non-empty and a valid path"; + + return options; +} + +} // namespace + +BaseShmExporter::BaseShmExporter(ExporterOptions options) + : options_(NormalizeAndValidateOptions(std::move(options))), + shm_writer_(ShmWriterOptions{ + .shm_dir = *options_.shm_dir, + .local_rank = *options_.local_rank, + }), + shm_collector_(ShmCollectorOptions{ + .shm_dir = *options_.shm_dir, + }) {} + +BaseShmExporter::~BaseShmExporter() = default; + +void BaseShmExporter::IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const { + shm_writer_.IncrementCounter(name, labels, val); +} + +void BaseShmExporter::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + if (!std::isfinite(val)) return; + + const MetricLabel rank_label = local_rank_label(); + if (labels.empty()) { + shm_writer_.SetGauge(name, LabelSpan(&rank_label, 1), val); + return; + } + if (absl::c_any_of(labels, [](const MetricLabel& label) { + return label.key == metric_labels::kLocalRank; + })) { + shm_writer_.SetGauge(name, labels, val); + return; + } + absl::InlinedVector + augmented_labels(labels.begin(), labels.end()); + if (absl::c_is_sorted(labels)) { + auto it = std::lower_bound(augmented_labels.begin(), augmented_labels.end(), + rank_label); + augmented_labels.insert(it, rank_label); + } else { + augmented_labels.push_back(rank_label); + absl::c_sort(augmented_labels); + } + shm_writer_.SetGauge(name, augmented_labels, val); +} + +void BaseShmExporter::ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const { + shm_writer_.ObserveHistogram(name, labels, val); +} + +std::string BaseShmExporter::GetTextSnapshot() const { return std::string(); } + +AggregatedMetrics BaseShmExporter::CollectMetrics() const { + return shm_collector_.CollectMetrics(); +} + +} // 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..c547a474 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter.h @@ -0,0 +1,89 @@ +// 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 "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_collector.h" +#include "tpu_sync/telemetry/shm/shm_writer.h" + +namespace tpu_raiden::telemetry { + +// Base class for multi-process shared-memory telemetry exporters. +// Manages an underlying ShmWriter for low-overhead metric publishing and an +// ShmCollector for multi-worker aggregation. +// +// Preconditions: +// Requires valid local_rank (via options.local_rank) and shm_dir (via +// options.shm_dir). Fails fast with CHECK if either is missing or invalid. +// Trailing slashes in shm_dir are stripped; the root path "/" is prohibited. +// Note: custom_buckets is ignored by shared-memory backends in favor of fixed +// compile-time bucket layouts (kDefaultHistogramBuckets). +// +// Thread-safety & Destruction contract: +// Callers must ensure all concurrent metric recording and CollectMetrics +// calls have finished prior to exporter destruction. +class BaseShmExporter : public MetricsBackend { + public: + explicit BaseShmExporter(ExporterOptions options); + ~BaseShmExporter() override; + + BaseShmExporter(const BaseShmExporter&) = delete; + BaseShmExporter& operator=(const BaseShmExporter&) = delete; + BaseShmExporter(BaseShmExporter&&) = delete; + BaseShmExporter& operator=(BaseShmExporter&&) = delete; + + // Metric recording methods. Safe for concurrent execution across worker + // threads. Writes directly to the memory-mapped shared segment. + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const override; + // Records a gauge value. Automatically attaches the "local_rank" label if not + // already present in caller labels. + void SetGauge(absl::string_view name, LabelSpan labels, + double val) const override; + // Records a histogram observation. Safe for concurrent execution across + // worker threads. Writes directly to the memory-mapped shared segment. + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const override; + + // Returns an empty string. BaseShmExporter is un-opinionated and does not + // prescribe an exposition format (Prometheus vs OTel); this method is + // intended to be overridden by format-specific subclasses. + std::string GetTextSnapshot() const override; + + // Scans the shared-memory directory and returns aggregated metrics across + // all local worker processes. + AggregatedMetrics CollectMetrics() const; + + // Returns the validated exporter configuration options. + const ExporterOptions& options() const { return options_; } + + private: + MetricLabel local_rank_label() const { + return MetricLabel{metric_labels::kLocalRank, *options_.local_rank}; + } + + ExporterOptions options_; + ShmWriter shm_writer_; + ShmCollector shm_collector_; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_BASE_SHM_EXPORTER_H_ diff --git a/tpu_sync/telemetry/base_shm_exporter_test.cc b/tpu_sync/telemetry/base_shm_exporter_test.cc new file mode 100644 index 00000000..48a279f0 --- /dev/null +++ b/tpu_sync/telemetry/base_shm_exporter_test.cc @@ -0,0 +1,382 @@ +// 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 +#include // NOLINT(build/c++17) +#include +#include +#include +#include // NOLINT(build/c++11) +#include // NOLINT(build/c++11) +#include + +#include +#include "absl/cleanup/cleanup.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_collector.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { +namespace { + +bool HasWorkerSegment(absl::string_view dir, absl::string_view prefix) { + std::error_code ec; + for (const std::filesystem::directory_entry& entry : + std::filesystem::directory_iterator(std::filesystem::path(dir), ec)) { + if (ec) break; + if (absl::StartsWith(entry.path().filename().string(), prefix)) { + return true; + } + } + return false; +} + +class BaseShmExporterTest : public testing::Test { + protected: + void SetUp() override { + test_dir_ = absl::StrCat( + testing::TempDir(), "/base_shm_exporter_test_", getpid(), "_", + testing::UnitTest::GetInstance()->current_test_info()->name()); + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + std::filesystem::create_directories(test_dir_, ec); + ASSERT_FALSE(ec); + } + + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + } + + std::string test_dir_; +}; + +TEST_F(BaseShmExporterTest, CheckFailsWhenLocalRankMissingOrInvalid) { + ExporterOptions options; + options.shm_dir = test_dir_; + options.local_rank = std::nullopt; + + EXPECT_DEATH((BaseShmExporter(options)), "local_rank"); + + options.local_rank = ""; + EXPECT_DEATH((BaseShmExporter(options)), "local_rank"); + + options.local_rank = "-1"; + EXPECT_DEATH((BaseShmExporter(options)), "local_rank"); + + options.local_rank = "../0"; + EXPECT_DEATH((BaseShmExporter(options)), "local_rank"); + + options.local_rank = "abc"; + EXPECT_DEATH((BaseShmExporter(options)), "local_rank"); +} + +TEST_F(BaseShmExporterTest, NormalizesAndCanonicalizesOptions) { + ExporterOptions options; + options.shm_dir = absl::StrCat(test_dir_, "///"); + options.local_rank = "007"; + + BaseShmExporter exporter(options); + EXPECT_EQ(exporter.options().shm_dir, test_dir_); + EXPECT_EQ(exporter.options().local_rank, "7"); + EXPECT_TRUE(HasWorkerSegment(test_dir_, absl::StrCat(kShmFilePrefix, "7_"))); +} + +TEST_F(BaseShmExporterTest, CheckFailsWhenShmDirMissingOrInvalid) { + ExporterOptions options; + options.local_rank = "0"; + + options.shm_dir = std::nullopt; + EXPECT_DEATH((BaseShmExporter(options)), "shm_dir"); + + options.shm_dir = ""; + EXPECT_DEATH((BaseShmExporter(options)), "shm_dir"); + + options.shm_dir = "/"; + EXPECT_DEATH((BaseShmExporter(options)), "shm_dir"); + + options.shm_dir = "///"; + EXPECT_DEATH((BaseShmExporter(options)), "shm_dir"); +} + +TEST_F(BaseShmExporterTest, RecordsAndAggregatesMetricsAcrossWorkers) { + BaseShmExporter exporter0(ExporterOptions{ + .local_rank = "0", + .shm_dir = test_dir_, + }); + + constexpr std::array push_labels = { + MetricLabel{metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter0.IncrementCounter(metric_names::kSentBytesTotal, push_labels, 1024); + exporter0.IncrementCounter(metric_names::kSentBytesTotal, push_labels, 512); + + constexpr std::array pull_labels = { + MetricLabel{metric_labels::kDirection, metric_labels::kDirectionPull}}; + exporter0.SetGauge(metric_names::kBufferAllocatedBytes, pull_labels, 1024.0); + exporter0.SetGauge("worker_state", {}, 1.0); + + exporter0.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 12.5); + EXPECT_TRUE(exporter0.GetTextSnapshot().empty()); + + { + BaseShmExporter exporter1(ExporterOptions{ + .local_rank = "1", + .shm_dir = test_dir_, + }); + exporter1.IncrementCounter(metric_names::kSentBytesTotal, push_labels, + 2000); + exporter1.SetGauge(metric_names::kBufferAllocatedBytes, pull_labels, + 2048.0); + exporter1.SetGauge("worker_state", {}, 2.0); + exporter1.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 25.0); + + AggregatedMetrics metrics = exporter0.CollectMetrics(); + + EXPECT_EQ(metrics.counters["sent_bytes_total"]["direction=push"], 3536); + EXPECT_DOUBLE_EQ( + metrics.gauges["buffer_allocated_bytes"]["direction=pull;local_rank=0"], + 1024.0); + EXPECT_DOUBLE_EQ( + metrics.gauges["buffer_allocated_bytes"]["direction=pull;local_rank=1"], + 2048.0); + EXPECT_DOUBLE_EQ(metrics.gauges["worker_state"]["local_rank=0"], 1.0); + EXPECT_DOUBLE_EQ(metrics.gauges["worker_state"]["local_rank=1"], 2.0); + + const HistogramData& hist = + metrics.histograms["transfer_duration_ms"]["direction=push"]; + EXPECT_EQ(hist.sample_count, 2); + EXPECT_DOUBLE_EQ(hist.sample_sum, 37.5); + // Both 12.5 and 25.0 fall into bucket 7 (le=25.0) + EXPECT_EQ(hist.bucket_counts[7], 2); + } + + AggregatedMetrics metrics = exporter0.CollectMetrics(); + + EXPECT_EQ(metrics.counters["sent_bytes_total"]["direction=push"], 1536); + EXPECT_DOUBLE_EQ( + metrics.gauges["buffer_allocated_bytes"]["direction=pull;local_rank=0"], + 1024.0); + EXPECT_FALSE(metrics.gauges["buffer_allocated_bytes"].contains( + "direction=pull;local_rank=1")); + EXPECT_DOUBLE_EQ(metrics.gauges["worker_state"]["local_rank=0"], 1.0); + EXPECT_FALSE(metrics.gauges["worker_state"].contains("local_rank=1")); + + const HistogramData& hist = + metrics.histograms["transfer_duration_ms"]["direction=push"]; + EXPECT_EQ(hist.sample_count, 1); + EXPECT_DOUBLE_EQ(hist.sample_sum, 12.5); + // 12.5 falls into bucket 7 (le=25.0) + EXPECT_EQ(hist.bucket_counts[7], 1); +} + +TEST_F(BaseShmExporterTest, ReapsDeadWorkerSegmentOnCollection) { + int pipe_fds[2]; + ASSERT_EQ(pipe(pipe_fds), 0); + absl::Cleanup pipe_cleanup_both = [&pipe_fds] { + close(pipe_fds[0]); + close(pipe_fds[1]); + }; + pid_t pid = fork(); + ASSERT_GE(pid, 0); + std::move(pipe_cleanup_both).Cancel(); + if (pid == 0) { + close(pipe_fds[0]); + ExporterOptions dead_options; + dead_options.shm_dir = test_dir_; + dead_options.local_rank = "99"; + BaseShmExporter dead_exporter(dead_options); + constexpr std::array labels = { + MetricLabel{metric_labels::kDirection, metric_labels::kDirectionPush}}; + dead_exporter.IncrementCounter(metric_names::kSentBytesTotal, labels, 500); + char ready = 'R'; + if (write(pipe_fds[1], &ready, 1) != 1) { + _exit(1); + } + close(pipe_fds[1]); + _exit(0); // Terminate without running ~BaseShmExporter() so segment file + // remains. + } + close(pipe_fds[1]); + absl::Cleanup pipe_cleanup = [read_fd = pipe_fds[0]] { close(read_fd); }; + bool child_reaped = false; + int status = 0; + absl::Cleanup child_cleanup = [pid, &status, &child_reaped] { + if (!child_reaped && pid > 0) { + waitpid(pid, &status, 0); + } + }; + char ready = 0; + ssize_t bytes_read = read(pipe_fds[0], &ready, 1); + std::move(pipe_cleanup).Cancel(); + close(pipe_fds[0]); + pid_t waited_pid = waitpid(pid, &status, 0); + child_reaped = true; + ASSERT_EQ(bytes_read, 1); + ASSERT_EQ(waited_pid, pid); + ASSERT_TRUE(WIFEXITED(status)); + ASSERT_EQ(WEXITSTATUS(status), 0); + + ASSERT_TRUE(HasWorkerSegment(test_dir_, absl::StrCat(kShmFilePrefix, "99_"))); + + ExporterOptions live_options; + live_options.shm_dir = test_dir_; + live_options.local_rank = "0"; + BaseShmExporter live_exporter(live_options); + + constexpr std::array labels = { + MetricLabel{metric_labels::kDirection, metric_labels::kDirectionPush}}; + live_exporter.IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + + EXPECT_EQ(live_exporter.CollectMetrics() + .counters["sent_bytes_total"]["direction=push"], + 100); + EXPECT_FALSE( + HasWorkerSegment(test_dir_, absl::StrCat(kShmFilePrefix, "99_"))); +} + +TEST_F(BaseShmExporterTest, SetGaugeDropsNonFiniteValues) { + BaseShmExporter exporter(ExporterOptions{ + .local_rank = "0", + .shm_dir = test_dir_, + }); + + exporter.SetGauge("test_nan", {}, std::numeric_limits::quiet_NaN()); + exporter.SetGauge("test_inf", {}, std::numeric_limits::infinity()); + exporter.SetGauge("test_neg_inf", {}, + -std::numeric_limits::infinity()); + + constexpr std::array labels = { + MetricLabel{metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.SetGauge("test_nan_labels", labels, + std::numeric_limits::quiet_NaN()); + exporter.SetGauge("test_inf_labels", labels, + std::numeric_limits::infinity()); + exporter.SetGauge("test_neg_inf_labels", labels, + -std::numeric_limits::infinity()); + + AggregatedMetrics metrics = exporter.CollectMetrics(); + EXPECT_FALSE(metrics.gauges.contains("test_nan")); + EXPECT_FALSE(metrics.gauges.contains("test_inf")); + EXPECT_FALSE(metrics.gauges.contains("test_neg_inf")); + EXPECT_FALSE(metrics.gauges.contains("test_nan_labels")); + EXPECT_FALSE(metrics.gauges.contains("test_inf_labels")); + EXPECT_FALSE(metrics.gauges.contains("test_neg_inf_labels")); +} + +TEST_F(BaseShmExporterTest, ConcurrentMetricRecording) { + BaseShmExporter exporter(ExporterOptions{ + .local_rank = "0", + .shm_dir = test_dir_, + }); + + constexpr int kNumThreads = 4; + constexpr int kOpsPerThread = 1000; + std::vector threads; + threads.reserve(kNumThreads); + + for (int t = 0; t < kNumThreads; ++t) { + threads.emplace_back([&exporter, t]() { + std::string t_str = absl::StrCat(t); + MetricLabel thread_label{"thread", t_str}; + for (int i = 0; i < kOpsPerThread; ++i) { + exporter.IncrementCounter(metric_names::kSentBytesTotal, + LabelSpan(&thread_label, 1), 1); + exporter.SetGauge(metric_names::kBufferAllocatedBytes, + LabelSpan(&thread_label, 1), static_cast(i)); + exporter.ObserveHistogram(metric_names::kTransferDurationMs, + LabelSpan(&thread_label, 1), 5.0); + } + }); + } + + for (auto& thread : threads) { + thread.join(); + } + + AggregatedMetrics metrics = exporter.CollectMetrics(); + for (int t = 0; t < kNumThreads; ++t) { + std::string counter_key = absl::StrCat("thread=", t); + EXPECT_EQ(metrics.counters[metric_names::kSentBytesTotal][counter_key], + kOpsPerThread); + std::string gauge_key = absl::StrCat("local_rank=0;thread=", t); + EXPECT_TRUE(metrics.gauges[metric_names::kBufferAllocatedBytes].contains( + gauge_key)); + EXPECT_DOUBLE_EQ( + metrics.gauges[metric_names::kBufferAllocatedBytes][gauge_key], + static_cast(kOpsPerThread - 1)); + EXPECT_EQ(metrics.histograms[metric_names::kTransferDurationMs][counter_key] + .sample_count, + kOpsPerThread); + EXPECT_DOUBLE_EQ( + metrics.histograms[metric_names::kTransferDurationMs][counter_key] + .sample_sum, + static_cast(kOpsPerThread) * 5.0); + } +} + +TEST_F(BaseShmExporterTest, SetGaugeLabelHandlingAndCapacity) { + BaseShmExporter exporter(ExporterOptions{ + .local_rank = "2", + .shm_dir = test_dir_, + }); + + // 1. Empty labels -> attaches local_rank. + exporter.SetGauge("empty_label_gauge", {}, 10.0); + + // 2. Caller already provides local_rank -> preserves caller's local_rank. + constexpr std::array custom_rank_labels = { + MetricLabel{metric_labels::kLocalRank, "99"}}; + exporter.SetGauge("custom_rank_gauge", custom_rank_labels, 30.0); + + // 3. Pre-sorted 8-label set (capacity boundary: 8 -> 9 slots, sorted + // insertion). + const std::string expected = "a=1;b=2;c=3;d=4;local_rank=2;m=5;n=6;y=7;z=8"; + const std::array sorted_labels = { + MetricLabel{"a", "1"}, MetricLabel{"b", "2"}, MetricLabel{"c", "3"}, + MetricLabel{"d", "4"}, MetricLabel{"m", "5"}, MetricLabel{"n", "6"}, + MetricLabel{"y", "7"}, MetricLabel{"z", "8"}, + }; + exporter.SetGauge("sorted_gauge", sorted_labels, 42.0); + + // 4. Unsorted 8-label set (triggers in-place sort path). + const std::array unsorted_labels = { + MetricLabel{"z", "8"}, MetricLabel{"y", "7"}, MetricLabel{"n", "6"}, + MetricLabel{"m", "5"}, MetricLabel{"d", "4"}, MetricLabel{"c", "3"}, + MetricLabel{"b", "2"}, MetricLabel{"a", "1"}, + }; + exporter.SetGauge("unsorted_gauge", unsorted_labels, 42.0); + + AggregatedMetrics metrics = exporter.CollectMetrics(); + EXPECT_DOUBLE_EQ(metrics.gauges["empty_label_gauge"]["local_rank=2"], 10.0); + EXPECT_DOUBLE_EQ(metrics.gauges["custom_rank_gauge"]["local_rank=99"], 30.0); + EXPECT_FALSE(metrics.gauges["custom_rank_gauge"].contains("local_rank=2")); + EXPECT_DOUBLE_EQ(metrics.gauges["sorted_gauge"][expected], 42.0); + EXPECT_DOUBLE_EQ(metrics.gauges["unsorted_gauge"][expected], 42.0); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/buffered_metrics_exporter.cc b/tpu_sync/telemetry/buffered_metrics_exporter.cc index 42f9c089..1d6fd867 100644 --- a/tpu_sync/telemetry/buffered_metrics_exporter.cc +++ b/tpu_sync/telemetry/buffered_metrics_exporter.cc @@ -15,7 +15,6 @@ #include "tpu_sync/telemetry/buffered_metrics_exporter.h" #include -#include #include #include #include @@ -34,7 +33,6 @@ namespace tpu_raiden::telemetry { namespace { constexpr absl::string_view kMetricPrefix = "tpu_raiden_"; -constexpr size_t kDefaultInlinedLabelCapacity = 4; std::string EscapeLabelValue(absl::string_view value) { std::string escaped; diff --git a/tpu_sync/telemetry/label_util.cc b/tpu_sync/telemetry/label_util.cc index 01b798db..e4a53fce 100644 --- a/tpu_sync/telemetry/label_util.cc +++ b/tpu_sync/telemetry/label_util.cc @@ -30,7 +30,6 @@ namespace tpu_raiden::telemetry { namespace { -constexpr size_t kDefaultInlinedLabelCapacity = 8; constexpr size_t kDefaultPrometheusStackBufferSize = 256; // Lightweight buffer writer that bounds-checks appends into a char span. diff --git a/tpu_sync/telemetry/metrics_backend.h b/tpu_sync/telemetry/metrics_backend.h index 4d630e6a..1731dffb 100644 --- a/tpu_sync/telemetry/metrics_backend.h +++ b/tpu_sync/telemetry/metrics_backend.h @@ -15,6 +15,7 @@ #ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_METRICS_BACKEND_H_ #define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_METRICS_BACKEND_H_ +#include #include #include #include @@ -50,11 +51,19 @@ struct ExporterOptions { int port = kDefaultExporterPort; // Non-owning view of histogram bucket boundaries. Defaults to // kDefaultHistogramBuckets and is copied by the exporter during construction. + // Note: Shared-memory backends (e.g. BaseShmExporter and its derived + // exporters like PrometheusShmExporter) ignore custom bucket spans and use + // fixed compile-time bucket layouts (kDefaultHistogramBuckets). absl::Span custom_buckets = kDefaultHistogramBuckets; - // Worker local rank identifier in distributed multi-rank environments. If - // unset (std::nullopt) or empty, telemetry initialization falls back to the - // LOCAL_RANK environment variable. + // Worker local rank identifier in distributed multi-rank environments. Note: + // Fallback to the LOCAL_RANK environment variable is handled at the + // metrics_api initialization layer; backend constructors (such as + // BaseShmExporter) expect an already-resolved, valid rank string. std::optional local_rank; + // Base directory for POSIX shared-memory segments for inter-worker telemetry + // aggregation via BaseShmExporter and its derived exporters. Trailing slashes + // are stripped; the root path "/" is prohibited. + std::optional shm_dir; }; // Structure defining centralized metadata for a Raiden metric. @@ -159,6 +168,7 @@ inline constexpr absl::string_view kDirectionPull = "pull"; inline constexpr absl::string_view kDirectionPullResponse = "pull_response"; inline constexpr absl::string_view kErrorCode = "error_code"; +inline constexpr absl::string_view kLocalRank = "local_rank"; } // namespace metric_labels // Structure defining a metric key-value label pair. @@ -172,6 +182,10 @@ struct MetricLabel { // Allocation-free label view span type definition using LabelSpan = absl::Span; +// Standard inline capacity for label containers across telemetry pipelines. +// Guarantees zero heap allocation for label sets up to this size. +inline constexpr size_t kDefaultInlinedLabelCapacity = 8; + // Abstract Dual-Backend Interface class MetricsBackend { public: diff --git a/tpu_sync/telemetry/shm/BUILD b/tpu_sync/telemetry/shm/BUILD index 8dc99457..02efafbe 100644 --- a/tpu_sync/telemetry/shm/BUILD +++ b/tpu_sync/telemetry/shm/BUILD @@ -68,3 +68,35 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "shm_collector", + srcs = ["shm_collector.cc"], + hdrs = ["shm_collector.h"], + deps = [ + ":shm_layout", + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "shm_collector_test", + srcs = ["shm_collector_test.cc"], + deps = [ + ":shm_collector", + ":shm_layout", + ":shm_writer", + "//tpu_sync/telemetry:metrics_backend", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + "@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..d272fdea --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.cc @@ -0,0 +1,311 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/shm/shm_collector.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/cleanup/cleanup.h" +#include "absl/container/flat_hash_map.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/strings/match.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { + +namespace { + +// Validates that the memory region beginning at `slot_bytes` has at least +// `sizeof(T)` byte capacity and meets the `alignof(T)` natural alignment +// required for atomic operations on slot type `T`. +template +bool IsSlotValid(const uint8_t* slot_bytes, uint32_t entry_size) { + return entry_size >= sizeof(T) && + (reinterpret_cast(slot_bytes) % alignof(T) == 0); +} + +// Retrieves a reference to an existing metric value in `map` without heap +// allocations via string_view lookup, or emplacing a new entry if absent. +template +ValueType& GetOrEmplaceMetric( + absl::flat_hash_map>& map, + absl::string_view metric_name, absl::string_view encoded_labels) { + auto metric_it = map.find(metric_name); + if (metric_it == map.end()) { + metric_it = map.emplace(std::string(metric_name), + absl::flat_hash_map{}) + .first; + } + auto label_it = metric_it->second.find(encoded_labels); + if (label_it == metric_it->second.end()) { + label_it = + metric_it->second.emplace(std::string(encoded_labels), ValueType{}) + .first; + } + return label_it->second; +} + +// Aggregates metrics from a mapped shared-memory segment into `metrics`. +// +// Preconditions: +// - `raw_segment` must be non-null and aligned to +// `alignof(ShmSegmentLayout)`. +// - `raw_segment` must point to at least `kSegmentTotalFileSize` bytes of +// valid readable memory. +// +// Reads TOC entries using acquire barriers to guard against uninitialized or +// partially written entries. Defensive checks reject corrupt offsets, +// non-null-terminated strings, and invalid scalar values. +void AggregateSegment(const void* raw_segment, AggregatedMetrics& metrics) { + if (raw_segment == nullptr || + (reinterpret_cast(raw_segment) % alignof(ShmSegmentLayout) != + 0)) { + return; + } + const ShmSegmentLayout* segment = + static_cast(raw_segment); + + if (segment->header.magic.load(std::memory_order_acquire) != + kRaidenShmMagic) { + return; + } + if (segment->header.max_toc_entries != kMaxTocEntries) { + return; + } + + const uint32_t data_pool_offset = segment->header.data_pool_offset; + if (data_pool_offset < sizeof(ShmSegmentLayout) || + data_pool_offset >= kSegmentTotalFileSize || + data_pool_offset % kMetricSlotAlignment != 0) { + return; + } + const size_t num_entries = + segment->header.toc_entry_count.load(std::memory_order_acquire); + if (num_entries > kMaxTocEntries) { + return; + } + + for (size_t i = 0; i < num_entries; ++i) { + const ShmTocEntry& toc_entry = segment->toc[i]; + if (toc_entry.entry_state.load(std::memory_order_acquire) != + TocEntryState::kCommitted) { + continue; + } + + const MetricType metric_type = toc_entry.type; + const uint32_t entry_offset = toc_entry.offset; + const uint32_t entry_size = toc_entry.size; + + if (metric_type != MetricType::kCounter && + metric_type != MetricType::kGauge && + metric_type != MetricType::kHistogram) { + continue; + } + if (entry_offset < data_pool_offset || + static_cast(entry_offset) + entry_size > + kSegmentTotalFileSize) { + continue; + } + + const void* const name_nul = + std::memchr(toc_entry.metric_name, '\0', sizeof(toc_entry.metric_name)); + if (name_nul == nullptr || name_nul == toc_entry.metric_name) { + continue; + } + absl::string_view metric_name( + toc_entry.metric_name, + static_cast(name_nul) - toc_entry.metric_name); + + const void* const labels_nul = std::memchr( + toc_entry.encoded_labels, '\0', sizeof(toc_entry.encoded_labels)); + if (labels_nul == nullptr) { + continue; + } + absl::string_view encoded_labels( + toc_entry.encoded_labels, + static_cast(labels_nul) - toc_entry.encoded_labels); + + const uint8_t* const slot_bytes = + reinterpret_cast(segment) + entry_offset; + + switch (metric_type) { + case MetricType::kCounter: { + if (IsSlotValid>(slot_bytes, entry_size)) { + GetOrEmplaceMetric(metrics.counters, metric_name, encoded_labels) += + reinterpret_cast*>(slot_bytes) + ->load(std::memory_order_relaxed); + } + break; + } + case MetricType::kGauge: { + if (IsSlotValid>(slot_bytes, entry_size)) { + const double gauge_value = + reinterpret_cast*>(slot_bytes) + ->load(std::memory_order_relaxed); + if (std::isfinite(gauge_value)) { + GetOrEmplaceMetric(metrics.gauges, metric_name, encoded_labels) += + gauge_value; + } + } + break; + } + case MetricType::kHistogram: { + if (IsSlotValid(slot_bytes, entry_size)) { + const ShmHistogramSlot* const histogram_slot = + reinterpret_cast(slot_bytes); + const uint64_t count = + histogram_slot->sample_count.load(std::memory_order_relaxed); + const double sum = + histogram_slot->sample_sum.load(std::memory_order_relaxed); + if (!std::isfinite(sum)) { + break; + } + + HistogramData& hist_data = GetOrEmplaceMetric( + metrics.histograms, metric_name, encoded_labels); + hist_data.sample_count += count; + hist_data.sample_sum += sum; + + for (size_t bucket_index = 0; + bucket_index < hist_data.bucket_counts.size(); ++bucket_index) { + hist_data.bucket_counts[bucket_index] += + histogram_slot->bucket_counts[bucket_index].load( + std::memory_order_relaxed); + } + } + break; + } + } + } +} + +} // namespace + +ShmCollector::ShmCollector(ShmCollectorOptions options) + : options_(std::move(options)) { + CHECK(!options_.shm_dir.empty()) + << "ShmCollector requires a non-empty shm_dir"; +} + +AggregatedMetrics ShmCollector::CollectMetrics() const { + AggregatedMetrics metrics; + + DIR* const dir_stream = opendir(options_.shm_dir.c_str()); + if (dir_stream == nullptr) { + PLOG_EVERY_N_SEC(WARNING, 10) + << "Failed to open shared-memory directory " << options_.shm_dir; + return metrics; + } + absl::Cleanup close_dir = [dir_stream] { closedir(dir_stream); }; + + while (const dirent* dir_entry = readdir(dir_stream)) { + absl::string_view filename(dir_entry->d_name); + const bool is_tmp = absl::EndsWith(filename, kShmTmpFileExtension); + if (!absl::StartsWith(filename, kShmFilePrefix) || + (!absl::EndsWith(filename, kShmFileExtension) && !is_tmp)) { + continue; + } + + const int fd = openat(dirfd(dir_stream), dir_entry->d_name, + O_RDONLY | O_CLOEXEC | O_NOFOLLOW | O_NONBLOCK); + if (fd < 0) { + if (errno != ENOENT) { + PLOG_EVERY_N_SEC(WARNING, 10) + << "Failed to open shared-memory segment at " << options_.shm_dir + << "/" << filename; + } + continue; + } + absl::Cleanup close_fd = [fd] { close(fd); }; + + struct stat file_stat{}; + if (fstat(fd, &file_stat) != 0 || !S_ISREG(file_stat.st_mode)) { + continue; + } + + if (flock(fd, LOCK_EX | LOCK_NB) == 0) { + // Process is dead. Dead worker metrics are intentionally discarded upon + // reaping; only metrics from live processes are aggregated. + // Re-verify fstat under lock: if already unlinked by a concurrent reaper, + // abort early. + if (fstat(fd, &file_stat) != 0 || file_stat.st_nlink == 0) { + continue; + } + // Security note: The shared memory directory is expected to be protected + // with sticky-bit or restricted permissions to prevent unprivileged + // symlink replacement attacks. We verify inode equality under exclusive + // lock and use unlinkat() relative to dir_stream to avoid path traversal + // TOCTOU. + struct stat current_stat{}; + if (fstatat(dirfd(dir_stream), dir_entry->d_name, ¤t_stat, + AT_SYMLINK_NOFOLLOW) == 0 && + current_stat.st_ino == file_stat.st_ino && + current_stat.st_dev == file_stat.st_dev) { + if (unlinkat(dirfd(dir_stream), dir_entry->d_name, 0) != 0 && + errno != ENOENT) { + PLOG_EVERY_N_SEC(WARNING, 10) + << "Failed to unlink dead worker shared-memory segment at " + << options_.shm_dir << "/" << filename; + } + } + // Do NOT call flock(LOCK_UN); close_fd will release the lock. + } else if (!is_tmp && flock(fd, LOCK_SH | LOCK_NB) == 0) { + // Memory Safety / Fault Tolerance: + // In this architecture, TPU workers and the collector operate in a + // trusted local domain. Shared advisory locks protect against collecting + // partially initialized segments. We re-verify file size and link count + // under the lock before mapping to guard against undersized or + // concurrently unlinked dead worker files. + if (fstat(fd, &file_stat) != 0 || file_stat.st_nlink == 0 || + file_stat.st_size < static_cast(kSegmentTotalFileSize)) { + continue; + } + void* const mapped_address = + mmap(nullptr, kSegmentTotalFileSize, PROT_READ, MAP_SHARED, fd, 0); + if (mapped_address == MAP_FAILED) { + PLOG_EVERY_N_SEC(WARNING, 10) + << "Failed to mmap shared-memory segment at " << options_.shm_dir + << "/" << filename; + continue; + } + absl::Cleanup unmap_segment = [mapped_address] { + munmap(mapped_address, kSegmentTotalFileSize); + }; + AggregateSegment(mapped_address, metrics); + } + } + + return metrics; +} + +} // 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..46aebd43 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.h @@ -0,0 +1,96 @@ +// 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 +#include + +#include "absl/container/flat_hash_map.h" +#include "tpu_sync/telemetry/shm/shm_layout.h" + +namespace tpu_raiden::telemetry { + +struct HistogramData { + uint64_t sample_count = 0; + double sample_sum = 0.0; + // Stores NON-CUMULATIVE counts directly as read from SHM to avoid premature + // cumulative conversion. + // Indices 0 to kNumHistogramBuckets - 1 correspond to upper bounds in + // kDefaultHistogramBuckets. Index kNumHistogramBuckets corresponds to +Inf + // (the overflow bucket). + std::array bucket_counts{}; +}; + +struct AggregatedMetrics { + // Metric Name -> Encoded Labels -> Value + absl::flat_hash_map> + counters; + absl::flat_hash_map> + gauges; + absl::flat_hash_map> + histograms; + + [[nodiscard]] bool empty() const { + return counters.empty() && gauges.empty() && histograms.empty(); + } +}; + +// 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. +// +// Thread-safety note: `ShmCollector` instances are thread-safe and const +// member methods may be called concurrently from multiple threads. +class ShmCollector { + public: + explicit ShmCollector(ShmCollectorOptions options); + + // Scans the configured shared-memory directory, aggregates metric totals + // across all live worker processes, and reaps dead worker files. + // + // Dead worker processes (whose exclusive lock can be acquired) are reaped + // by unlinking their segment files. Dead worker metrics are intentionally + // discarded upon reaping; only metrics from live processes are aggregated. + // Downstream consumers should note that cumulative counters from reaped + // workers are not retained. + // + // Gauges are aggregated across workers by addition. For non-additive gauges + // (e.g. utilization or fraction), callers should supply rank-distinguishing + // labels to avoid cross-worker aggregation. + // + // Histogram slots are read using relaxed atomic operations. Active writer + // updates during scraping are eventually consistent. + // + // Thread safety: Concurrent invocations from multiple threads are safe. + [[nodiscard]] AggregatedMetrics CollectMetrics() 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..27204ee4 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector_test.cc @@ -0,0 +1,795 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/shm/shm_collector.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include // NOLINT(build/c++17) +#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 "absl/synchronization/notification.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 { +namespace { + +using ::testing::SizeIs; + +constexpr MetricLabel kPush{metric_labels::kDirection, + metric_labels::kDirectionPush}; +constexpr MetricLabel kPull{metric_labels::kDirection, + metric_labels::kDirectionPull}; +constexpr std::array kPushLabels = {kPush}; +constexpr std::array kPullLabels = {kPull}; + +class ShmCollectorTest : public testing::Test { + protected: + void SetUp() override { + test_dir_ = (std::filesystem::path(testing::TempDir()) / + absl::StrCat("shm_col_", getpid(), "_", + reinterpret_cast(this))) + .string(); + std::error_code ec; + std::filesystem::create_directories(test_dir_, ec); + ASSERT_FALSE(ec); + } + + void TearDown() override { + for (int fd : held_fds_) { + if (fd >= 0) close(fd); + } + held_fds_.clear(); + 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 FilePath(absl::string_view filename) const { + return absl::StrCat(test_dir_, "/", filename); + } + + std::string ShmName(absl::string_view name, + absl::string_view ext = kShmFileExtension) const { + return absl::StrCat(kShmFilePrefix, name, ext); + } + + AggregatedMetrics Collect(absl::string_view sub_dir = "") const { + const std::string dir = sub_dir.empty() ? test_dir_ : FilePath(sub_dir); + return ShmCollector(ShmCollectorOptions{.shm_dir = dir}).CollectMetrics(); + } + + void CollectInto(AggregatedMetrics& metrics) const { + metrics = ShmCollector(CollectorOptions()).CollectMetrics(); + } + + void WriteFile(absl::string_view filename, absl::string_view data = "", + size_t file_size = 0) { + const std::string path = FilePath(filename); + const int fd = + open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644); + ASSERT_GE(fd, 0); + absl::Cleanup close_fd = [fd] { close(fd); }; + if (!data.empty()) { + ASSERT_EQ(write(fd, data.data(), data.size()), + static_cast(data.size())); + } + if (file_size > data.size()) { + ASSERT_EQ(ftruncate(fd, file_size), 0); + } + } + + int HoldSharedLock(absl::string_view filename) { + const std::string path = FilePath(filename); + const int fd = open(path.c_str(), O_RDONLY | O_CLOEXEC); + if (fd < 0) return -1; + if (flock(fd, LOCK_SH | LOCK_NB) != 0) { + close(fd); + return -1; + } + held_fds_.push_back(fd); + return fd; + } + + int WriteAndLockFile(absl::string_view filename, absl::string_view data = "", + size_t file_size = 0) { + WriteFile(filename, data, file_size); + return HoldSharedLock(filename); + } + + void ReleaseHeldLock(int fd) { + auto it = std::find(held_fds_.begin(), held_fds_.end(), fd); + if (it != held_fds_.end()) { + close(*it); + held_fds_.erase(it); + } + } + + bool FileExists(absl::string_view filename) const { + std::error_code ec; + const std::filesystem::file_status status = + std::filesystem::symlink_status(FilePath(filename), ec); + return !ec && std::filesystem::status_known(status) && + status.type() != std::filesystem::file_type::not_found; + } + + int CountMatchingFiles(absl::string_view substring) const { + int count = 0; + std::error_code ec; + for (const std::filesystem::directory_entry& entry : + std::filesystem::directory_iterator(test_dir_, ec)) { + if (absl::StrContains(entry.path().filename().string(), substring)) { + ++count; + } + } + return count; + } + + template + void CreateDeadWorkerSegment(absl::string_view rank, F&& populate) { + std::array pipe_fds; + ASSERT_EQ(pipe2(pipe_fds.data(), O_CLOEXEC), 0); + const pid_t pid = fork(); + ASSERT_GE(pid, 0); + if (pid == 0) { + prctl(PR_SET_PDEATHSIG, SIGKILL); + close(pipe_fds[0]); + // Allocate writer on heap without deleting so ~ShmWriter() does not + // unlink the segment file, simulating an abrupt process termination. + auto* writer = new ShmWriter(WriterOptions(rank)); + populate(*writer); + char ready = 'R'; + if (write(pipe_fds[1], &ready, 1) != 1) { + _exit(1); + } + close(pipe_fds[1]); + _exit(0); + } + close(pipe_fds[1]); + absl::Cleanup cleanup_parent = [pipe_fd = pipe_fds[0], pid] { + close(pipe_fd); + kill(pid, SIGKILL); + int status = 0; + waitpid(pid, &status, 0); + }; + char sync = 0; + ASSERT_EQ(read(pipe_fds[0], &sync, 1), 1); + std::move(cleanup_parent).Cancel(); + close(pipe_fds[0]); + int status = 0; + ASSERT_EQ(waitpid(pid, &status, 0), pid); + ASSERT_TRUE(WIFEXITED(status)); + ASSERT_EQ(WEXITSTATUS(status), 0); + } + + std::string test_dir_; + std::vector held_fds_; +}; + +using ShmCollectorDeathTest = ShmCollectorTest; + +// Lightweight helper to build memory-mapped segments for unit testing. +struct MockSegment { + struct alignas(64) AlignedSegment { + alignas(64) char data[kSegmentTotalFileSize]; + }; + + std::unique_ptr storage; + ShmSegmentLayout* layout; + uint32_t next_offset = sizeof(ShmSegmentLayout); + + MockSegment() + : storage(std::make_unique()), + layout(reinterpret_cast(storage->data)) { + std::memset(storage->data, 0, sizeof(storage->data)); + layout->header.magic.store(kRaidenShmMagic); + layout->header.max_toc_entries = kMaxTocEntries; + layout->header.data_pool_offset = sizeof(ShmSegmentLayout); + } + + uint32_t Alloc(uint32_t size, uint32_t align = 64) { + next_offset = (next_offset + align - 1) & ~(align - 1); + const uint32_t offset = next_offset; + next_offset += size; + return offset; + } + + ShmTocEntry& AddEntry(absl::string_view name, MetricType type, + std::optional offset = std::nullopt, + uint32_t size = 8, absl::string_view labels = "", + TocEntryState state = TocEntryState::kCommitted) { + const size_t idx = + layout->header.toc_entry_count.load(std::memory_order_relaxed); + ShmTocEntry& entry = layout->toc[idx]; + entry.type = type; + entry.offset = offset.has_value() ? *offset : Alloc(size); + entry.size = size; + SetBounded(entry.metric_name, name); + SetBounded(entry.encoded_labels, labels); + entry.entry_state.store(state, std::memory_order_release); + layout->header.toc_entry_count.store(idx + 1, std::memory_order_release); + return entry; + } + + template + void StoreValue(uint32_t offset, T value) { + if (offset % alignof(T) == 0 && + offset + sizeof(T) <= sizeof(storage->data)) { + ::new (static_cast(storage->data + offset)) std::atomic(value); + } + } + + void AddCounter(absl::string_view name, uint64_t value, + absl::string_view labels = "", + TocEntryState state = TocEntryState::kCommitted, + std::optional offset = std::nullopt, + uint32_t size = sizeof(std::atomic)) { + const ShmTocEntry& entry = + AddEntry(name, MetricType::kCounter, offset, size, labels, state); + StoreValue(entry.offset, value); + } + + void AddGauge(absl::string_view name, double value, + absl::string_view labels = "", + TocEntryState state = TocEntryState::kCommitted, + std::optional offset = std::nullopt, + uint32_t size = sizeof(std::atomic)) { + const ShmTocEntry& entry = + AddEntry(name, MetricType::kGauge, offset, size, labels, state); + StoreValue(entry.offset, value); + } + + void AddHistogram(absl::string_view name, double sum, uint64_t count, + absl::string_view labels = "", + std::optional offset = std::nullopt) { + const ShmTocEntry& entry = AddEntry(name, MetricType::kHistogram, offset, + sizeof(ShmHistogramSlot), labels); + if (entry.offset % alignof(ShmHistogramSlot) == 0 && + entry.offset + sizeof(ShmHistogramSlot) <= sizeof(storage->data)) { + auto* slot = ::new (static_cast(storage->data + entry.offset)) + ShmHistogramSlot(); + slot->sample_sum.store(sum, std::memory_order_relaxed); + slot->sample_count.store(count, std::memory_order_relaxed); + } + } + + absl::string_view data() const { + return absl::string_view(storage->data, sizeof(storage->data)); + } + + private: + template + static void SetBounded(char (&dest)[N], absl::string_view src) { + std::memset(dest, 0, N); + if (!src.empty()) { + std::memcpy(dest, src.data(), std::min(src.size(), N - 1)); + } + } +}; + +TEST_F(ShmCollectorDeathTest, RejectsEmptyOrUnsetShmDir) { + EXPECT_DEATH(ShmCollector(ShmCollectorOptions{}), + "ShmCollector requires a non-empty shm_dir"); + EXPECT_DEATH(ShmCollector(ShmCollectorOptions{.shm_dir = ""}), + "ShmCollector requires a non-empty shm_dir"); +} + +TEST_F(ShmCollectorTest, HandlesNonExistentShmDir) { + EXPECT_TRUE(Collect("nonexistent").empty()); +} + +TEST_F(ShmCollectorTest, AggregatesLiveWorkersAndReapsDead) { + ASSERT_NO_FATAL_FAILURE( + CreateDeadWorkerSegment("dead", [](const ShmWriter& dead_writer) { + dead_writer.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, + 500); + dead_writer.SetGauge(metric_names::kBufferAllocatedBytes, {}, 512.0); + })); + + ShmWriter writer_0(WriterOptions("0")); + ShmWriter writer_1(WriterOptions("1")); + const std::array failure_labels = { + kPull, MetricLabel{metric_labels::kErrorCode, "DEADLINE_EXCEEDED"}}; + const std::array hist_labels = { + kPush, MetricLabel{metric_labels::kErrorCode, "RESOURCE_EXHAUSTED"}}; + + writer_0.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, 100); + writer_1.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, 200); + writer_0.IncrementCounter(metric_names::kSentBytesTotal, kPullLabels, 50); + writer_1.IncrementCounter(metric_names::kSentBytesTotal, kPullLabels, 75); + writer_1.IncrementCounter(metric_names::kTransferFailuresTotal, + failure_labels, 5); + writer_0.SetGauge(metric_names::kBufferAllocatedBytes, {}, 1000.0); + writer_1.SetGauge(metric_names::kBufferAllocatedBytes, {}, 2000.0); + writer_0.ObserveHistogram(metric_names::kTransferDurationMs, {}, 0.05); + writer_1.ObserveHistogram(metric_names::kTransferDurationMs, {}, 0.15); + writer_1.ObserveHistogram(metric_names::kTransferDurationMs, hist_labels, + 0.25); + + AggregatedMetrics metrics = Collect(); + + // Verify Counters + EXPECT_EQ(metrics.counters[metric_names::kSentBytesTotal]["direction=push"], + 300); + EXPECT_EQ(metrics.counters[metric_names::kSentBytesTotal]["direction=pull"], + 125); + EXPECT_EQ(metrics.counters[metric_names::kTransferFailuresTotal] + ["direction=pull;error_code=DEADLINE_EXCEEDED"], + 5); + + // Verify Gauges + EXPECT_DOUBLE_EQ(metrics.gauges[metric_names::kBufferAllocatedBytes][""], + 3000.0); + + // Verify Histograms (Non-cumulative) + const HistogramData& hist_unlabeled = + metrics.histograms[metric_names::kTransferDurationMs][""]; + EXPECT_EQ(hist_unlabeled.sample_count, 2); + EXPECT_DOUBLE_EQ(hist_unlabeled.sample_sum, 0.20); + // writer_0 observed 0.05 -> bucket 0 (le=0.1) + // writer_1 observed 0.15 -> bucket 1 (le=0.25) + EXPECT_EQ(hist_unlabeled.bucket_counts[0], 1); + EXPECT_EQ(hist_unlabeled.bucket_counts[1], 1); + EXPECT_EQ(hist_unlabeled.bucket_counts[2], 0); + + const HistogramData& hist_labeled = + metrics.histograms[metric_names::kTransferDurationMs] + ["direction=push;error_code=RESOURCE_EXHAUSTED"]; + EXPECT_EQ(hist_labeled.sample_count, 1); + EXPECT_DOUBLE_EQ(hist_labeled.sample_sum, 0.25); + // writer_1 observed 0.25 -> bucket 1 (le=0.25) + EXPECT_EQ(hist_labeled.bucket_counts[0], 0); + EXPECT_EQ(hist_labeled.bucket_counts[1], 1); + + EXPECT_EQ(CountMatchingFiles("worker_rank_dead_"), 0); + EXPECT_EQ(CountMatchingFiles("worker_rank_0_"), 1); + EXPECT_EQ(CountMatchingFiles("worker_rank_1_"), 1); +} + +TEST_F(ShmCollectorTest, ConcurrencyAndProcessCrashHandling) { + ASSERT_NO_FATAL_FAILURE( + CreateDeadWorkerSegment("crashed", [](const ShmWriter& writer) { + writer.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, + 777); + })); + for (int i = 0; i < 3; ++i) { + ASSERT_NO_FATAL_FAILURE(CreateDeadWorkerSegment( + absl::StrCat("dead_", i), [](const ShmWriter& writer) { + writer.IncrementCounter(metric_names::kSentBytesTotal, {}, 100); + })); + } + ShmWriter live(WriterOptions("live")); + live.IncrementCounter(metric_names::kSentBytesTotal, {}, 42); + + absl::Notification start_notification; + std::vector threads; + threads.reserve(8); + for (int i = 0; i < 8; ++i) { + threads.emplace_back([this, &start_notification]() { + start_notification.WaitForNotification(); + AggregatedMetrics thread_metrics = + ShmCollector(CollectorOptions()).CollectMetrics(); + if (thread_metrics.counters.contains("sent_bytes_total")) { + EXPECT_FALSE(thread_metrics.counters.at("sent_bytes_total") + .contains("direction=push")); + } + EXPECT_EQ(thread_metrics.counters["sent_bytes_total"][""], 42); + }); + } + start_notification.Notify(); + for (std::thread& thread : threads) { + thread.join(); + } + + AggregatedMetrics final_metrics = Collect(); + EXPECT_EQ(final_metrics.counters["sent_bytes_total"][""], 42); + if (final_metrics.counters.contains("sent_bytes_total")) { + EXPECT_FALSE(final_metrics.counters.at("sent_bytes_total") + .contains("direction=push")); + } + EXPECT_EQ(CountMatchingFiles("crashed"), 0); + EXPECT_EQ(CountMatchingFiles("dead_"), 0); + EXPECT_EQ(CountMatchingFiles("live"), 1); +} + +TEST_F(ShmCollectorTest, AggregatesMultiChunkWorker) { + ShmWriter writer(WriterOptions("multi")); + for (size_t i = 0; i < kMaxTocEntries + 20; ++i) { + writer.IncrementCounter(absl::StrCat("metric_", i), {}, 1); + } + + AggregatedMetrics metrics = Collect(); + EXPECT_THAT(metrics.counters, SizeIs(kMaxTocEntries + 20)); + EXPECT_EQ(metrics.counters["metric_0"][""], 1); + EXPECT_EQ(metrics.counters[absl::StrCat("metric_", kMaxTocEntries - 1)][""], + 1); + EXPECT_EQ(metrics.counters[absl::StrCat("metric_", kMaxTocEntries)][""], 1); + EXPECT_EQ(metrics.counters[absl::StrCat("metric_", kMaxTocEntries + 19)][""], + 1); +} + +TEST_F(ShmCollectorTest, AggregatesMultiChunkHistogramsOnPoolExhaustion) { + constexpr size_t kNumHistograms = 350; + ShmWriter writer(WriterOptions("multi_hist")); + for (size_t i = 0; i < kNumHistograms; ++i) { + writer.ObserveHistogram(absl::StrCat("hist_", i), {}, 0.05); + } + + EXPECT_GE(CountMatchingFiles("multi_hist"), 2); + AggregatedMetrics metrics = Collect(); + EXPECT_THAT(metrics.histograms, SizeIs(kNumHistograms)); + EXPECT_EQ(metrics.histograms["hist_0"][""].sample_count, 1); + EXPECT_EQ(metrics.histograms["hist_0"][""].bucket_counts[0], 1); + EXPECT_EQ(metrics.histograms[absl::StrCat("hist_", kNumHistograms - 1)][""] + .sample_count, + 1); +} + +TEST_F(ShmCollectorTest, ReapsAllChunksOfMultiChunkDeadWorker) { + ASSERT_NO_FATAL_FAILURE( + CreateDeadWorkerSegment("dead_multi", [](const ShmWriter& writer) { + for (size_t i = 0; i < kMaxTocEntries + 20; ++i) { + writer.IncrementCounter(absl::StrCat("dead_metric_", i), {}, 1); + } + })); + + ASSERT_GE(CountMatchingFiles("dead_multi"), 2); + EXPECT_TRUE(Collect().empty()); + EXPECT_EQ(CountMatchingFiles("dead_multi"), 0); +} + +TEST_F(ShmCollectorTest, ClearsPrePopulatedMetrics) { + ShmWriter live(WriterOptions("live")); + live.IncrementCounter(metric_names::kSentBytesTotal, {}, 10); + + AggregatedMetrics metrics; + metrics.counters["stale_a"][""] = 999; + metrics.gauges["stale_b"][""] = 123.0; + + CollectInto(metrics); + + EXPECT_THAT(metrics.counters, SizeIs(1)); + EXPECT_EQ(metrics.counters["sent_bytes_total"][""], 10); + EXPECT_TRUE(metrics.gauges.empty()); +} + +TEST_F(ShmCollectorTest, SkipsDefensiveSlotValidationFailures) { + constexpr uint32_t kBase = sizeof(ShmSegmentLayout); + MockSegment mock; + + mock.AddCounter("valid_counter", 100); + // 1. Unaligned slot pointer (offset + 1 is not 8-byte aligned) + mock.AddCounter("unaligned_counter", 0, "", TocEntryState::kCommitted, + kBase + 65); + // 2. TOC offset below data_pool_offset + mock.AddCounter("below_pool_offset", 0, "", TocEntryState::kCommitted, + kBase - 64); + // 3. TOC offset + size exceeds kSegmentTotalFileSize + mock.AddCounter("exceeds_filesize", 0, "", TocEntryState::kCommitted, + kSegmentTotalFileSize - 4); + // 3b. 32-bit unsigned offset wraparound overflow + mock.AddCounter("overflow_offset", 0, "", TocEntryState::kCommitted, + 0xFFFFFFC0, 64); + // 4. Non-finite gauge value (NaN) + mock.AddGauge("nan_gauge", std::numeric_limits::quiet_NaN()); + // 5. Non-finite gauge value (Inf) + mock.AddGauge("inf_gauge", std::numeric_limits::infinity()); + // 6. Non-finite histogram sample sum (Inf) + mock.AddHistogram("inf_hist", std::numeric_limits::infinity(), 5); + // 7. Undersized slot descriptor for Counter (< 8 bytes) + mock.AddCounter("undersized_counter", 0, "", TocEntryState::kCommitted, + std::nullopt, 4); + // 8. Histogram with 0 observations + mock.AddHistogram("zero_hist", 0.0, 0); + // 9. Non-null-terminated metric name filling all 64 bytes + ShmTocEntry& bad_name = mock.AddEntry("", MetricType::kCounter, std::nullopt, + 8, "", TocEntryState::kWriting); + std::memset(bad_name.metric_name, 'x', sizeof(bad_name.metric_name)); + bad_name.entry_state.store(TocEntryState::kCommitted, + std::memory_order_release); + // 10. Non-null-terminated encoded labels filling all 128 bytes + ShmTocEntry& bad_labels = + mock.AddEntry("valid_name_bad_labels", MetricType::kCounter, std::nullopt, + 8, "", TocEntryState::kWriting); + std::memset(bad_labels.encoded_labels, 'y', + sizeof(bad_labels.encoded_labels)); + bad_labels.entry_state.store(TocEntryState::kCommitted, + std::memory_order_release); + mock.AddCounter("uninit_metric", 0, "", TocEntryState::kUninitialized); + mock.AddCounter("writing_metric", 0, "", TocEntryState::kWriting); + + ASSERT_GE(WriteAndLockFile(ShmName("defensive_slots"), mock.data()), 0); + + AggregatedMetrics metrics = Collect(); + EXPECT_THAT(metrics.counters, SizeIs(1)); + EXPECT_THAT(metrics.histograms, SizeIs(1)); + + EXPECT_EQ(metrics.counters["valid_counter"][""], 100); + EXPECT_EQ(metrics.histograms["zero_hist"][""].sample_count, 0); + + for (absl::string_view name : + {"unaligned_counter", "below_pool_offset", "exceeds_filesize", + "overflow_offset", "undersized_counter", "valid_name_bad_labels", + "uninit_metric", "writing_metric"}) { + EXPECT_FALSE(metrics.counters.contains(name)) + << "Found unexpected counter: " << name; + } + + for (absl::string_view name : {"nan_gauge", "inf_gauge"}) { + EXPECT_FALSE(metrics.gauges.contains(name)) + << "Found unexpected gauge: " << name; + } + + EXPECT_FALSE(metrics.histograms.contains("inf_hist")); +} + +TEST_F(ShmCollectorTest, ResilientToCorruptFiles) { + ShmWriter live(WriterOptions("good")); + live.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, 500); + + WriteFile(ShmName("zero")); + WriteFile(ShmName("truncated"), "", sizeof(ShmSegmentLayout) + 128); + EXPECT_GE(WriteAndLockFile(ShmName("live_truncated"), "short"), 0); + + auto write_bad_header = [&](absl::string_view name, uint32_t magic, + uint32_t max_toc, uint32_t pool_offset, + uint32_t toc_count = 0) { + ShmTocHeader header{}; + header.magic.store(magic); + header.max_toc_entries = max_toc; + header.data_pool_offset = pool_offset; + header.toc_entry_count.store(toc_count); + const absl::string_view data(reinterpret_cast(&header), + sizeof(header)); + WriteFile(ShmName(name), data, kSegmentTotalFileSize); + EXPECT_GE(WriteAndLockFile(ShmName(absl::StrCat("live_", name)), data, + kSegmentTotalFileSize), + 0); + }; + + constexpr uint32_t kInvalidMagic = 0xDEADBEEF; + write_bad_header("bad_magic", kInvalidMagic, kMaxTocEntries, + sizeof(ShmSegmentLayout)); + write_bad_header("bad_max_toc", kRaidenShmMagic, kMaxTocEntries + 1, + sizeof(ShmSegmentLayout)); + write_bad_header("bad_offset", kRaidenShmMagic, kMaxTocEntries, + kSegmentTotalFileSize); + write_bad_header("bad_offset_low", kRaidenShmMagic, kMaxTocEntries, + sizeof(ShmSegmentLayout) - 1); + write_bad_header("bad_offset_unaligned", kRaidenShmMagic, kMaxTocEntries, + sizeof(ShmSegmentLayout) + 1); + write_bad_header("bad_count", kRaidenShmMagic, kMaxTocEntries, + sizeof(ShmSegmentLayout), kMaxTocEntries + 1); + + WriteFile("other_file.txt", "non-shm content"); + std::error_code ec; + std::filesystem::create_directories(FilePath(ShmName("dir")), ec); + ASSERT_FALSE(ec); + ASSERT_EQ(symlink("/dev/null", FilePath(ShmName("symlink")).c_str()), 0); + + const std::string unreadable = ShmName("unreadable"); + const std::string unreadable_path = FilePath(unreadable); + const bool is_root = (geteuid() == 0); + if (!is_root) { + WriteFile(unreadable, "", kSegmentTotalFileSize); + ASSERT_EQ(chmod(unreadable_path.c_str(), 0000), 0); + } + absl::Cleanup restore_permissions = [unreadable_path, is_root] { + if (!is_root) { + chmod(unreadable_path.c_str(), 0644); + } + }; + + AggregatedMetrics metrics = Collect(); + EXPECT_THAT(metrics.counters, SizeIs(1)); + EXPECT_EQ(metrics.counters["sent_bytes_total"]["direction=push"], 500); + + EXPECT_TRUE(FileExists("other_file.txt")); + EXPECT_TRUE(FileExists(ShmName("dir"))); + EXPECT_TRUE(FileExists(ShmName("symlink"))); + if (!is_root) { + EXPECT_TRUE(FileExists(unreadable)); + } + + for (absl::string_view name : + {"zero", "truncated", "bad_magic", "bad_max_toc", "bad_offset", + "bad_offset_low", "bad_offset_unaligned", "bad_count"}) { + SCOPED_TRACE(absl::StrCat("Checking deletion of corrupt file: ", name)); + EXPECT_FALSE(FileExists(ShmName(name))); + } + for (absl::string_view name : + {"bad_magic", "bad_max_toc", "bad_offset", "bad_offset_low", + "bad_offset_unaligned", "bad_count"}) { + SCOPED_TRACE(absl::StrCat("Checking preservation of live file: ", name)); + EXPECT_TRUE(FileExists(ShmName(absl::StrCat("live_", name)))); + } + EXPECT_TRUE(FileExists(ShmName("live_truncated"))); +} + +TEST_F(ShmCollectorTest, SkipsFifoWithoutBlocking) { + const std::string fifo = ShmName("fifo"); + ASSERT_EQ(mkfifo(FilePath(fifo).c_str(), 0644), 0); + + ShmWriter live(WriterOptions("live")); + live.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, 100); + + AggregatedMetrics metrics = Collect(); + EXPECT_THAT(metrics.counters, SizeIs(1)); + EXPECT_EQ(metrics.counters["sent_bytes_total"]["direction=push"], 100); + EXPECT_TRUE(FileExists(fifo)); +} + +TEST_F(ShmCollectorTest, ReapsOrphanedTmpFilesAndPreservesActiveTmpFiles) { + ShmWriter live(WriterOptions("live")); + live.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, 100); + + const std::string dead_chunk = ShmName("dead_chunk_0", kShmTmpFileExtension); + const std::string dead_plain = ShmName("dead_plain", ".tmp"); + const std::string active_chunk = + ShmName("active_chunk_0", kShmTmpFileExtension); + + WriteFile(dead_chunk, "orphaned chunk tmp"); + WriteFile(dead_plain, "orphaned plain tmp"); + + MockSegment mock; + mock.AddCounter("tmp_uncommitted_counter", 999); + const int active_fd = WriteAndLockFile(active_chunk, mock.data()); + ASSERT_GE(active_fd, 0); + + AggregatedMetrics metrics = Collect(); + EXPECT_THAT(metrics.counters, SizeIs(1)); + EXPECT_EQ(metrics.counters["sent_bytes_total"]["direction=push"], 100); + EXPECT_FALSE(FileExists(dead_chunk)); + EXPECT_TRUE(FileExists(dead_plain)); + EXPECT_TRUE(FileExists(active_chunk)); + + ReleaseHeldLock(active_fd); + + metrics = Collect(); + EXPECT_THAT(metrics.counters, SizeIs(1)); + EXPECT_EQ(metrics.counters["sent_bytes_total"]["direction=push"], 100); + EXPECT_FALSE(FileExists(active_chunk)); + EXPECT_TRUE(FileExists(dead_plain)); +} + +#if defined(__linux__) +struct UnlinkContext { + static inline std::atomic path{nullptr}; + static inline std::atomic invoked{0}; + + static void HandleSigio(int /*signum*/) { + const char* target = path.exchange(nullptr, std::memory_order_acq_rel); + if (target != nullptr) { + unlink(target); + invoked.store(1, std::memory_order_release); + } + } +}; + +TEST_F(ShmCollectorTest, SkipsUnlinkedSegmentUnderSharedLock) { + // Execute in an isolated child subprocess so asynchronous signals and + // handler registrations never leak into the test harness or interfere with + // thread sanitizers and test watchdog threads. + const pid_t pid = fork(); + ASSERT_GE(pid, 0); + if (pid == 0) { + const std::string unlinked_file = ShmName("unlinked"); + const std::string path = FilePath(unlinked_file); + + MockSegment mock; + mock.AddCounter("unlinked_counter", 999); + const int holder_fd = WriteAndLockFile(unlinked_file, mock.data()); + if (holder_fd < 0) { + _exit(1); + } + + ShmWriter live(WriterOptions("live")); + live.IncrementCounter(metric_names::kSentBytesTotal, kPushLabels, 100); + + UnlinkContext::path.store(path.c_str(), std::memory_order_release); + UnlinkContext::invoked.store(0, std::memory_order_release); + + struct sigaction sa{}; + sa.sa_handler = UnlinkContext::HandleSigio; + sa.sa_flags = SA_RESTART; + sigemptyset(&sa.sa_mask); + if (sigaction(SIGIO, &sa, nullptr) != 0) { + _exit(2); + } + + const int inotify_fd = inotify_init1(IN_CLOEXEC); + if (inotify_fd < 0) { + _exit(3); + } + if (inotify_add_watch(inotify_fd, path.c_str(), IN_OPEN) < 0) { + _exit(4); + } + if (fcntl(inotify_fd, F_SETOWN, getpid()) != 0) { + _exit(5); + } + const int flags = fcntl(inotify_fd, F_GETFL); + if (flags < 0 || fcntl(inotify_fd, F_SETFL, flags | O_ASYNC) != 0) { + _exit(6); + } + + const AggregatedMetrics metrics = Collect(); + const bool correct_metrics = + metrics.counters.size() == 1 && + metrics.counters.contains("sent_bytes_total") && + metrics.counters.at("sent_bytes_total").contains("direction=push") && + metrics.counters.at("sent_bytes_total").at("direction=push") == 100; + const bool unlinked = + UnlinkContext::invoked.load(std::memory_order_acquire) == 1; + + struct stat holder_stat{}; + const bool link_zero = + fstat(holder_fd, &holder_stat) == 0 && holder_stat.st_nlink == 0; + + close(inotify_fd); + close(holder_fd); + + _exit((correct_metrics && unlinked && link_zero) ? 0 : 7); + } + + int status = 0; + ASSERT_EQ(waitpid(pid, &status, 0), pid); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 0); +} +#endif + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/shm/shm_layout.h b/tpu_sync/telemetry/shm/shm_layout.h index c4b8eed2..d8e6f340 100644 --- a/tpu_sync/telemetry/shm/shm_layout.h +++ b/tpu_sync/telemetry/shm/shm_layout.h @@ -34,6 +34,7 @@ namespace tpu_raiden::telemetry { inline constexpr uint32_t kRaidenShmMagic = 0xABCD1234; inline constexpr absl::string_view kShmFilePrefix = "worker_rank_"; inline constexpr absl::string_view kShmFileExtension = ".mmap"; +inline constexpr absl::string_view kShmTmpFileExtension = ".mmap.tmp"; inline constexpr size_t kMaxTocEntries = 1024; // Data pool capacity for metric slots in a single chunk (64 KB).