diff --git a/tpu_sync/telemetry/shm/BUILD b/tpu_sync/telemetry/shm/BUILD index 8dc99457..25ffa520 100644 --- a/tpu_sync/telemetry/shm/BUILD +++ b/tpu_sync/telemetry/shm/BUILD @@ -47,6 +47,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/hash", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/random", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -68,3 +69,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..d57ac909 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.cc @@ -0,0 +1,301 @@ +// 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) { + absl::flat_hash_map& label_map = + map.try_emplace(metric_name).first->second; + return label_map.try_emplace(encoded_labels).first->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); + // Skip non-shm or .tmp shm files. + if (!absl::StartsWith(filename, kShmFilePrefix) || + !absl::EndsWith(filename, kShmFileExtension)) { + 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 && errno != ELOOP) { + 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 (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..1e533837 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector.h @@ -0,0 +1,127 @@ +// 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 { + +// Aggregated sample statistics and non-cumulative bucket counts for a histogram +// metric stream, aggregated across all live worker shared-memory segments. +// +// Invariants: +// - `bucket_counts` stores raw, non-cumulative counts per bucket as scraped +// from `ShmHistogramSlot` to avoid loss of precision during aggregation. +// - Indices `0` through `kNumHistogramBuckets - 1` correspond to the upper +// bounds in `kDefaultHistogramBuckets`. +// - Index `kNumHistogramBuckets` represents the `+Inf` overflow bucket. +// - `sample_count` is guaranteed to equal the sum of all elements in +// `bucket_counts`. +// +// Consumption: +// - Exporters (e.g., Prometheus) must compute running cumulative sums when +// emitting standard cumulative histogram bucket formats. +struct HistogramData { + uint64_t sample_count = 0; + double sample_sum = 0.0; + std::array bucket_counts{}; +}; + +// Point-in-time snapshot of telemetry metrics collected and aggregated across +// all active worker shared-memory segments. +// +// Data Organization: +// - Two-level map structure: Metric Name -> Encoded Labels -> Value. +// - The inner label key is the canonical representation generated by +// `EncodeLabels` (comma-separated `key=value` pairs sorted lexicographically +// by key). Unlabeled metrics use the empty string (`""`) as their key. +// +// Invariants: +// - Aggregation is additive: for matching metric names and label sets across +// workers, counters, gauges, and histogram buckets/counts/sums are summed. +// - Contains only metrics from valid, non-corrupted segments where shared-lock +// acquisition and TOC validation succeeded. +// +// Thread Safety & Consumption: +// - Move-only/copyable value object returned by +// `ShmCollector::CollectMetrics()`. +// - Not internally synchronized; callers own the returned instance. +// - Typically consumed by telemetry exporters (e.g., `PrometheusShmExporter`) +// to generate scrape endpoints or metrics exposition text. +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..9d5099e6 --- /dev/null +++ b/tpu_sync/telemetry/shm/shm_collector_test.cc @@ -0,0 +1,828 @@ +// 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"}}; + const std::array host_pool_labels = { + MetricLabel{"pool", "host"}}; + const std::array device_pool_labels = { + MetricLabel{"pool", "device"}}; + + 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.SetGauge(metric_names::kBufferAllocatedBytes, host_pool_labels, + 500.0); + writer_1.SetGauge(metric_names::kBufferAllocatedBytes, host_pool_labels, + 700.0); + writer_1.SetGauge(metric_names::kBufferAllocatedBytes, device_pool_labels, + 250.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); + EXPECT_DOUBLE_EQ( + metrics.gauges[metric_names::kBufferAllocatedBytes]["pool=host"], 1200.0); + EXPECT_DOUBLE_EQ( + metrics.gauges[metric_names::kBufferAllocatedBytes]["pool=device"], + 250.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(); + const auto it = + thread_metrics.counters.find(metric_names::kSentBytesTotal); + ASSERT_NE(it, thread_metrics.counters.end()); + EXPECT_FALSE(it->second.contains("direction=push")); + const auto label_it = it->second.find(""); + ASSERT_NE(label_it, it->second.end()); + EXPECT_EQ(label_it->second, 42); + }); + } + start_notification.Notify(); + for (std::thread& thread : threads) { + thread.join(); + } + + AggregatedMetrics final_metrics = Collect(); + const auto it = final_metrics.counters.find(metric_names::kSentBytesTotal); + ASSERT_NE(it, final_metrics.counters.end()); + EXPECT_FALSE(it->second.contains("direction=push")); + const auto label_it = it->second.find(""); + ASSERT_NE(label_it, it->second.end()); + EXPECT_EQ(label_it->second, 42); + 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[metric_names::kSentBytesTotal][""], 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[metric_names::kSentBytesTotal]["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[metric_names::kSentBytesTotal]["direction=push"], + 100); + EXPECT_TRUE(FileExists(fifo)); +} + +TEST_F(ShmCollectorTest, IgnoresTmpFilesAndPreservesActiveWorkers) { + 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[metric_names::kSentBytesTotal]["direction=push"], + 100); + EXPECT_TRUE(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[metric_names::kSentBytesTotal]["direction=push"], + 100); + EXPECT_TRUE(FileExists(dead_chunk)); + EXPECT_TRUE(FileExists(dead_plain)); + EXPECT_TRUE(FileExists(active_chunk)); +} + +#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(); + + // Precondition: inotify / SIGIO handler must have executed. + if (UnlinkContext::invoked.load(std::memory_order_acquire) != 1) { + _exit(7); + } + + // Precondition: file must be unlinked while holder_fd is open. + struct stat holder_stat{}; + if (fstat(holder_fd, &holder_stat) != 0 || holder_stat.st_nlink != 0) { + _exit(8); + } + + // Assertion: metrics aggregation must have skipped unlinked segment and + // only aggregated the live worker. + if (!(metrics.counters.size() == 1 && + metrics.counters.contains(metric_names::kSentBytesTotal) && + metrics.counters.at(metric_names::kSentBytesTotal) + .contains("direction=push") && + metrics.counters.at(metric_names::kSentBytesTotal) + .at("direction=push") == 100)) { + _exit(9); + } + + close(inotify_fd); + close(holder_fd); + _exit(0); + } + + 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..f78d659f 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). @@ -101,6 +102,7 @@ struct alignas(64) ShmTocEntry { // updating distinct metric streams on different CPU cores. struct alignas(64) ShmTocHeader { std::atomic magic{0}; + uint32_t reserved{0}; int64_t pid{0}; std::atomic toc_entry_count{0}; uint32_t max_toc_entries{kMaxTocEntries}; diff --git a/tpu_sync/telemetry/shm/shm_writer.cc b/tpu_sync/telemetry/shm/shm_writer.cc index 88b08fc9..ee3e0ba2 100644 --- a/tpu_sync/telemetry/shm/shm_writer.cc +++ b/tpu_sync/telemetry/shm/shm_writer.cc @@ -18,7 +18,6 @@ #include #include #include -#include #include #include @@ -34,6 +33,7 @@ #include #include "absl/container/flat_hash_map.h" +#include "absl/log/check.h" #include "absl/log/log.h" #include "absl/random/random.h" #include "absl/strings/str_cat.h" @@ -48,10 +48,8 @@ namespace tpu_raiden::telemetry { 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; - } + CHECK(!options_.shm_dir.empty() && !options_.local_rank.empty()) + << "ShmWriter requires non-empty shm_dir and local_rank"; std::error_code ec; std::filesystem::create_directories(options_.shm_dir, ec); @@ -102,10 +100,11 @@ bool ShmWriter::AllocateNewChunk() const { } uint32_t chunk_idx = static_cast(chunks_.size()); - std::string path = + std::string base_path = absl::StrCat(options_.shm_dir, "/", kShmFilePrefix, options_.local_rank, - "_", uuid_, "_chunk_", chunk_idx, kShmFileExtension); - std::string tmp_path = absl::StrCat(path, ".tmp"); + "_", uuid_, "_chunk_", chunk_idx); + const std::string path = absl::StrCat(base_path, kShmFileExtension); + std::string tmp_path = absl::StrCat(base_path, kShmTmpFileExtension); int fd = open(tmp_path.c_str(), O_RDWR | O_CREAT | O_TRUNC | O_NOFOLLOW | O_CLOEXEC, diff --git a/tpu_sync/telemetry/shm/shm_writer.h b/tpu_sync/telemetry/shm/shm_writer.h index e66e12b5..89a1caae 100644 --- a/tpu_sync/telemetry/shm/shm_writer.h +++ b/tpu_sync/telemetry/shm/shm_writer.h @@ -58,7 +58,7 @@ class ShmWriter { public: static constexpr size_t kMaxChunks = 16; - explicit ShmWriter(const ShmWriterOptions& options = {}); + explicit ShmWriter(const ShmWriterOptions& options); ~ShmWriter(); ShmWriter(const ShmWriter&) = delete; diff --git a/tpu_sync/telemetry/shm/shm_writer_test.cc b/tpu_sync/telemetry/shm/shm_writer_test.cc index 6fe8b493..c236b4d4 100644 --- a/tpu_sync/telemetry/shm/shm_writer_test.cc +++ b/tpu_sync/telemetry/shm/shm_writer_test.cc @@ -765,19 +765,20 @@ TEST_F(ShmWriterTest, InputValidationAndFailureResilience) { ASSERT_GE(fd, 0); close(fd); + EXPECT_DEATH(ShmWriter(ShmWriterOptions{.shm_dir = "", .local_rank = "0"}), + "ShmWriter requires non-empty shm_dir and local_rank"); + EXPECT_DEATH( + ShmWriter(ShmWriterOptions{.shm_dir = test_dir_, .local_rank = ""}), + "ShmWriter requires non-empty shm_dir and local_rank"); + MetricLabel dummy_label{metric_labels::kDirection, metric_labels::kDirectionPush}; - for (const ShmWriterOptions& opt : - {ShmWriterOptions{.shm_dir = "", .local_rank = "0"}, - ShmWriterOptions{.shm_dir = test_dir_, .local_rank = ""}, - ShmWriterOptions{.shm_dir = absl::StrCat(regular_file, "/x"), - .local_rank = "0"}}) { - ShmWriter bad_writer(opt); - bad_writer.IncrementCounter(metric_names::kSentBytesTotal, - {&dummy_label, 1}, 10); - bad_writer.SetGauge("any_g", {&dummy_label, 1}, 1.0); - bad_writer.ObserveHistogram("any_h", {&dummy_label, 1}, 2.0); - } + ShmWriter bad_writer(ShmWriterOptions{ + .shm_dir = absl::StrCat(regular_file, "/x"), .local_rank = "0"}); + bad_writer.IncrementCounter(metric_names::kSentBytesTotal, {&dummy_label, 1}, + 10); + bad_writer.SetGauge("any_g", {&dummy_label, 1}, 1.0); + bad_writer.ObserveHistogram("any_h", {&dummy_label, 1}, 2.0); EXPECT_FALSE(MappedSegment("").is_valid()); EXPECT_FALSE(MappedSegment(test_dir_, "").is_valid());