diff --git a/tpu_sync/telemetry/BUILD b/tpu_sync/telemetry/BUILD index c75f0129..65686791 100644 --- a/tpu_sync/telemetry/BUILD +++ b/tpu_sync/telemetry/BUILD @@ -87,6 +87,7 @@ cc_library( # Disables Clang header modules due to incompatibility with prometheus-cpp / CivetWeb headers. features = ["-use_header_modules"], deps = [ + ":exporter_util", ":metrics_backend", "@com_github_jupp0r_prometheus_cpp//core", "@com_github_jupp0r_prometheus_cpp//pull", @@ -114,6 +115,7 @@ cc_library( srcs = ["buffered_metrics_exporter.cc"], hdrs = ["buffered_metrics_exporter.h"], deps = [ + ":exporter_util", ":metrics_backend", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", @@ -190,3 +192,61 @@ cc_test( "@com_google_googletest//:gtest_main", ], ) + +cc_library( + name = "prometheus_shm_exporter", + srcs = ["prometheus_shm_exporter.cc"], + hdrs = ["prometheus_shm_exporter.h"], + # -fexceptions is required because prometheus-cpp Exposer throws C++ + # exceptions (std::runtime_error) on socket binding and initialization failure. + copts = ["-fexceptions"], + # Disables Clang header modules due to incompatibility with prometheus-cpp / CivetWeb headers. + features = ["-use_header_modules"], + deps = [ + ":base_shm_exporter", + ":exporter_util", + ":label_util", + ":metrics_backend", + "//tpu_sync/telemetry/shm:shm_collector", + "//tpu_sync/telemetry/shm:shm_layout", + "@com_github_jupp0r_prometheus_cpp//core", + "@com_github_jupp0r_prometheus_cpp//pull", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "prometheus_shm_exporter_test", + srcs = ["prometheus_shm_exporter_test.cc"], + deps = [ + ":metrics_backend", + ":prometheus_shm_exporter", + ":test_util", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/synchronization", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "exporter_util", + srcs = ["exporter_util.cc"], + hdrs = ["exporter_util.h"], + deps = [ + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "exporter_util_test", + srcs = ["exporter_util_test.cc"], + deps = [ + ":exporter_util", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_sync/telemetry/buffered_metrics_exporter.cc b/tpu_sync/telemetry/buffered_metrics_exporter.cc index 1d6fd867..e7f1c72f 100644 --- a/tpu_sync/telemetry/buffered_metrics_exporter.cc +++ b/tpu_sync/telemetry/buffered_metrics_exporter.cc @@ -26,14 +26,13 @@ #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "tpu_sync/telemetry/exporter_util.h" #include "tpu_sync/telemetry/metrics_backend.h" namespace tpu_raiden::telemetry { namespace { -constexpr absl::string_view kMetricPrefix = "tpu_raiden_"; - std::string EscapeLabelValue(absl::string_view value) { std::string escaped; escaped.reserve(value.size()); @@ -154,8 +153,8 @@ BufferedMetricsExporter::GetAndResetMetricSamples() { LockFreeCounterAccumulator* counter) { uint64_t delta = counter->ExchangeAndReset(); if (delta > 0) { - std::string full_name = - absl::StrCat(kMetricPrefix, name, canonical_labels); + const std::string full_name = + absl::StrCat(kPrometheusMetricPrefix, name, canonical_labels); result[full_name].push_back(static_cast(delta)); } }); @@ -166,8 +165,8 @@ BufferedMetricsExporter::GetAndResetMetricSamples() { [&](absl::string_view canonical_labels, QueueBuffer<>* gauge) { std::vector samples = gauge->ExtractAndReset(); if (!samples.empty()) { - std::string full_name = - absl::StrCat(kMetricPrefix, name, canonical_labels); + const std::string full_name = + absl::StrCat(kPrometheusMetricPrefix, name, canonical_labels); result[full_name] = std::move(samples); } }); @@ -179,7 +178,7 @@ BufferedMetricsExporter::GetAndResetMetricSamples() { std::vector samples = histogram->ExtractAndReset(); if (!samples.empty()) { std::string full_name = - absl::StrCat(kMetricPrefix, name, canonical_labels); + absl::StrCat(kPrometheusMetricPrefix, name, canonical_labels); result[full_name] = std::move(samples); } }); diff --git a/tpu_sync/telemetry/exporter_util.cc b/tpu_sync/telemetry/exporter_util.cc new file mode 100644 index 00000000..b67033f4 --- /dev/null +++ b/tpu_sync/telemetry/exporter_util.cc @@ -0,0 +1,32 @@ +// 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/exporter_util.h" + +#include + +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" + +namespace tpu_raiden::telemetry { + +std::string JoinHostPort(absl::string_view host, int port) { + if (absl::StrContains(host, ':') && !absl::StartsWith(host, "[")) { + return absl::StrCat("[", host, "]:", port); + } + return absl::StrCat(host, ":", port); +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/exporter_util.h b/tpu_sync/telemetry/exporter_util.h new file mode 100644 index 00000000..afa64a19 --- /dev/null +++ b/tpu_sync/telemetry/exporter_util.h @@ -0,0 +1,40 @@ +// 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_EXPORTER_UTIL_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_EXPORTER_UTIL_H_ + +#include + +#include "absl/strings/string_view.h" + +namespace tpu_raiden::telemetry { + +// Common metric namespace prefix prepended to all Prometheus metric family +// names. +inline constexpr absl::string_view kPrometheusMetricPrefix = "tpu_raiden_"; + +// Reserved Prometheus histogram bucket upper-bound label ("less than or +// equal"). +inline constexpr absl::string_view kPrometheusLeLabel = "le"; + +// Formats a host string and port into a valid endpoint address (e.g. +// "127.0.0.1:8080"). If the host contains ':' and is not already bracketed +// (IPv6 address), it wraps the host in brackets (e.g. "[::1]:8080") per RFC +// 3986. +std::string JoinHostPort(absl::string_view host, int port); + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_EXPORTER_UTIL_H_ diff --git a/tpu_sync/telemetry/exporter_util_test.cc b/tpu_sync/telemetry/exporter_util_test.cc new file mode 100644 index 00000000..942cf269 --- /dev/null +++ b/tpu_sync/telemetry/exporter_util_test.cc @@ -0,0 +1,54 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/exporter_util.h" + +#include +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" + +namespace tpu_raiden::telemetry { +namespace { + +TEST(ExporterUtilTest, JoinHostPort) { + struct TestCase { + absl::string_view host; + int port; + absl::string_view expected; + }; + + constexpr TestCase kTestCases[] = { + // Standard IPv4 + {"127.0.0.1", 8080, "127.0.0.1:8080"}, + {"0.0.0.0", 9090, "0.0.0.0:9090"}, + // IPv6 without brackets (must be bracketed per RFC 3986) + {"::1", 8080, "[::1]:8080"}, + {"2001:db8::1", 9090, "[2001:db8::1]:9090"}, + // Pre-bracketed IPv6 (must not be double-bracketed) + {"[::1]", 8080, "[::1]:8080"}, + {"[2001:db8::1]", 9090, "[2001:db8::1]:9090"}, + // Edge cases + {"", 8080, ":8080"}, + {"localhost", 0, "localhost:0"}, + {"localhost", 65535, "localhost:65535"}, + }; + + for (const auto& [host, port, expected] : kTestCases) { + SCOPED_TRACE(absl::StrCat("host: '", host, "', port: ", port)); + EXPECT_EQ(JoinHostPort(host, port), expected); + } +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/prometheus_exporter.cc b/tpu_sync/telemetry/prometheus_exporter.cc index 66882975..25c7c3bb 100644 --- a/tpu_sync/telemetry/prometheus_exporter.cc +++ b/tpu_sync/telemetry/prometheus_exporter.cc @@ -30,17 +30,15 @@ #include "prometheus/text_serializer.h" #include "absl/container/flat_hash_map.h" #include "absl/log/log.h" -#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/exporter_util.h" #include "tpu_sync/telemetry/metrics_backend.h" namespace tpu_raiden::telemetry { namespace { -constexpr absl::string_view kMetricPrefix = "tpu_raiden_"; - std::map ConvertLabels(LabelSpan labels) { if (labels.empty()) { return {}; @@ -52,13 +50,6 @@ std::map ConvertLabels(LabelSpan labels) { return result; } -std::string JoinHostPort(absl::string_view host, int port) { - if (absl::StrContains(host, ':') && !absl::StartsWith(host, "[")) { - return absl::StrCat("[", host, "]:", port); - } - return absl::StrCat(host, ":", port); -} - } // namespace void PrometheusExporter::RegisterKnownFamilies() { @@ -68,7 +59,8 @@ void PrometheusExporter::RegisterKnownFamilies() { histogram_families_.contains(meta.name)) { continue; } - std::string prometheus_name = absl::StrCat(kMetricPrefix, meta.name); + const std::string prometheus_name = + absl::StrCat(kPrometheusMetricPrefix, meta.name); switch (meta.type) { case MetricType::kCounter: { auto* family = &prometheus::BuildCounter() @@ -118,7 +110,7 @@ PrometheusExporter::PrometheusExporter(const ExporterOptions& options) << ": " << e.what(); exposer_.reset(); } - } else if (options_.port != 0) { + } else if (options_.port > 0) { LOG(WARNING) << "Invalid port configured for Prometheus HTTP exporter: " << options_.port << ". Expected port in range [" << kMinPort << ", " << kMaxPort << "]."; diff --git a/tpu_sync/telemetry/prometheus_shm_exporter.cc b/tpu_sync/telemetry/prometheus_shm_exporter.cc new file mode 100644 index 00000000..ee0abe6f --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter.cc @@ -0,0 +1,248 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/prometheus_shm_exporter.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "prometheus/client_metric.h" +#include "prometheus/collectable.h" +#include "prometheus/exposer.h" +#include "prometheus/metric_family.h" +#include "prometheus/metric_type.h" +#include "prometheus/text_serializer.h" +#include "absl/algorithm/container.h" +#include "absl/log/log.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "tpu_sync/telemetry/base_shm_exporter.h" +#include "tpu_sync/telemetry/exporter_util.h" +#include "tpu_sync/telemetry/label_util.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 { + +// Parses compact semicolon-encoded SHM labels into Prometheus ClientMetric +// labels, reusing ParseShmLabels() and sorting label pairs lexicographically by +// name per the Prometheus / OpenMetrics exposition standard. +std::vector ParseLabels( + absl::string_view encoded_labels) { + if (encoded_labels.empty()) { + return {}; + } + std::vector> raw_label_pairs = + ParseShmLabels(encoded_labels); + std::vector labels; + labels.reserve(raw_label_pairs.size()); + for (auto& [key, value] : raw_label_pairs) { + labels.push_back({.name = std::move(key), .value = std::move(value)}); + } + absl::c_sort(labels); + return labels; +} + +prometheus::MetricFamily CreateFamily(const MetricMetadata& metadata) { + prometheus::MetricFamily family; + family.name = absl::StrCat(kPrometheusMetricPrefix, metadata.name); + family.help = metadata.description; + switch (metadata.type) { + case MetricType::kCounter: + family.type = prometheus::MetricType::Counter; + break; + case MetricType::kGauge: + family.type = prometheus::MetricType::Gauge; + break; + case MetricType::kHistogram: + family.type = prometheus::MetricType::Histogram; + break; + } + return family; +} + +// Converts aggregated shared memory metrics into Prometheus MetricFamily +// structures for exposition. +std::vector CollectMetricFamilies( + const AggregatedMetrics& metrics) { + if (metrics.empty()) { + return {}; + } + + std::vector families; + families.reserve(std::size(metric_metadata::kAllMetrics)); + + for (const MetricMetadata& metadata : metric_metadata::kAllMetrics) { + switch (metadata.type) { + case MetricType::kCounter: { + auto counter_it = metrics.counters.find(metadata.name); + if (counter_it == metrics.counters.end() || + counter_it->second.empty()) { + break; + } + const auto& label_counts = counter_it->second; + prometheus::MetricFamily family = CreateFamily(metadata); + family.metric.reserve(label_counts.size()); + for (const auto& [encoded_labels, count] : label_counts) { + prometheus::ClientMetric& metric = family.metric.emplace_back(); + metric.label = ParseLabels(encoded_labels); + metric.counter.value = static_cast(count); + } + families.push_back(std::move(family)); + break; + } + case MetricType::kGauge: { + auto gauge_it = metrics.gauges.find(metadata.name); + if (gauge_it == metrics.gauges.end() || gauge_it->second.empty()) { + break; + } + const auto& label_values = gauge_it->second; + prometheus::MetricFamily family = CreateFamily(metadata); + family.metric.reserve(label_values.size()); + for (const auto& [encoded_labels, gauge_value] : label_values) { + prometheus::ClientMetric& metric = family.metric.emplace_back(); + metric.label = ParseLabels(encoded_labels); + metric.gauge.value = gauge_value; + } + families.push_back(std::move(family)); + break; + } + case MetricType::kHistogram: { + auto hist_it = metrics.histograms.find(metadata.name); + if (hist_it == metrics.histograms.end() || hist_it->second.empty()) { + break; + } + const auto& label_histograms = hist_it->second; + prometheus::MetricFamily family = CreateFamily(metadata); + family.metric.reserve(label_histograms.size()); + for (const auto& [encoded_labels, histogram_data] : label_histograms) { + prometheus::ClientMetric& metric = family.metric.emplace_back(); + metric.label = ParseLabels(encoded_labels); + std::erase_if(metric.label, + [](const prometheus::ClientMetric::Label& label) { + return label.name == kPrometheusLeLabel; + }); + + metric.histogram.sample_sum = histogram_data.sample_sum; + metric.histogram.bucket.reserve(kNumHistogramBuckets + 1); + + uint64_t cumulative_count = 0; + for (size_t bucket_idx = 0; bucket_idx < kNumHistogramBuckets; + ++bucket_idx) { + cumulative_count += histogram_data.bucket_counts[bucket_idx]; + metric.histogram.bucket.push_back({ + .cumulative_count = cumulative_count, + .upper_bound = kDefaultHistogramBuckets[bucket_idx], + }); + } + // Add +Inf bucket + cumulative_count += + histogram_data.bucket_counts[kNumHistogramBuckets]; + metric.histogram.bucket.push_back({ + .cumulative_count = cumulative_count, + .upper_bound = std::numeric_limits::infinity(), + }); + metric.histogram.sample_count = cumulative_count; + } + families.push_back(std::move(family)); + break; + } + } + } + + // Sort metrics within families + for (prometheus::MetricFamily& family : families) { + absl::c_sort(family.metric, [](const prometheus::ClientMetric& metric_a, + const prometheus::ClientMetric& metric_b) { + return metric_a.label < metric_b.label; + }); + } + + return families; +} + +// Adapter class implementing prometheus::Collectable by querying the exporter's +// aggregated metrics and converting them via CollectMetricFamilies(). +class PrometheusShmCollectable final : public prometheus::Collectable { + public: + explicit PrometheusShmCollectable(const PrometheusShmExporter& exporter) + : exporter_(exporter) {} + std::vector Collect() const override { + return CollectMetricFamilies(exporter_.CollectMetrics()); + } + + private: + const PrometheusShmExporter& exporter_; +}; + +} // namespace + +// Holds HTTP exposition state. +// Note: `collectable` is declared before `exposer` so that `exposer` destructs +// first (in reverse declaration order), synchronously stopping the CivetWeb +// server and joining worker threads before `collectable` is released. +struct PrometheusShmExporter::ExposerState { + std::shared_ptr collectable; + std::unique_ptr exposer; +}; + +PrometheusShmExporter::PrometheusShmExporter(ExporterOptions exporter_options) + : BaseShmExporter(std::move(exporter_options)) { + const int port = options().port; + if (port >= kMinPort && port <= kMaxPort) { + const std::string endpoint = JoinHostPort(options().bind_address, port); + try { + std::unique_ptr exposer = + std::make_unique(endpoint); + std::shared_ptr collectable = + std::make_shared(*this); + exposer->RegisterCollectable(collectable); + exposer_state_ = std::make_unique(); + exposer_state_->collectable = std::move(collectable); + exposer_state_->exposer = std::move(exposer); + LOG(INFO) << "Prometheus SHM exporter listening on http://" << endpoint + << "/metrics"; + } catch (const std::exception& e) { + LOG(INFO) << "Failed to bind Prometheus HTTP exposer on " << endpoint + << ": " << e.what(); + exposer_state_.reset(); + } + } else if (port > 0) { + LOG(WARNING) << "Invalid port configured for Prometheus SHM HTTP exposer: " + << port << ". Expected port in range [" << kMinPort << ", " + << kMaxPort << "]."; + } +} + +PrometheusShmExporter::~PrometheusShmExporter() = default; + +bool PrometheusShmExporter::IsServerRunningForTesting() const { + return exposer_state_ != nullptr && exposer_state_->exposer != nullptr; +} + +std::string PrometheusShmExporter::GetTextSnapshot() const { + return prometheus::TextSerializer().Serialize( + CollectMetricFamilies(CollectMetrics())); +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_sync/telemetry/prometheus_shm_exporter.h b/tpu_sync/telemetry/prometheus_shm_exporter.h new file mode 100644 index 00000000..f9b41455 --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter.h @@ -0,0 +1,66 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_PROMETHEUS_SHM_EXPORTER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_PROMETHEUS_SHM_EXPORTER_H_ + +#include +#include + +#include "tpu_sync/telemetry/base_shm_exporter.h" +#include "tpu_sync/telemetry/metrics_backend.h" + +namespace tpu_raiden::telemetry { + +// PrometheusShmExporter exposes shared memory metrics collected across +// multi-worker processes via an HTTP /metrics endpoint for Prometheus scrapers. +// +// Inherits from BaseShmExporter to aggregate counters, gauges, and histograms +// from shared memory segments. The first process to bind the designated port +// runs the HTTP exposer, while secondary workers gracefully disable their HTTP +// server and continue writing/aggregating metrics via shared memory. All +// instances can generate text snapshots independently. +// +// Operational note: In this decentralized port-contention model, the HTTP +// exposer is bound during construction. If the primary worker serving the HTTP +// endpoint terminates, surviving workers continue aggregating and serving +// metrics via shared memory and GetTextSnapshot(). A newly initialized exporter +// or worker can bind the port once the primary terminates. +// +// Thread safety: Thread-safe for concurrent calls to public methods. +class PrometheusShmExporter final : public BaseShmExporter { + public: + explicit PrometheusShmExporter(ExporterOptions exporter_options); + ~PrometheusShmExporter() override; + PrometheusShmExporter(const PrometheusShmExporter&) = delete; + PrometheusShmExporter& operator=(const PrometheusShmExporter&) = delete; + PrometheusShmExporter(PrometheusShmExporter&&) = delete; + PrometheusShmExporter& operator=(PrometheusShmExporter&&) = delete; + + // Collects and serializes all shared memory metrics into Prometheus text + // exposition format (version 0.0.4). + std::string GetTextSnapshot() const override; + + // Returns true if the HTTP server is actively running and listening for + // scrape requests. For testing purposes only. + bool IsServerRunningForTesting() const; + + private: + struct ExposerState; + std::unique_ptr exposer_state_; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_SYNC_TELEMETRY_PROMETHEUS_SHM_EXPORTER_H_ diff --git a/tpu_sync/telemetry/prometheus_shm_exporter_test.cc b/tpu_sync/telemetry/prometheus_shm_exporter_test.cc new file mode 100644 index 00000000..7fb4362e --- /dev/null +++ b/tpu_sync/telemetry/prometheus_shm_exporter_test.cc @@ -0,0 +1,343 @@ +// Copyright 2026 Google LLC. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tpu_sync/telemetry/prometheus_shm_exporter.h" + +#include + +#include // NOLINT(build/c++17) +#include +#include +#include +#include // NOLINT(build/c++11) +#include // NOLINT(build/c++11) +#include + +#include +#include +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/notification.h" +#include "tpu_sync/telemetry/metrics_backend.h" +#include "tpu_sync/telemetry/test_util.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::HasSubstr; +using ::testing::Not; + +// Test fixture for PrometheusShmExporter validating shared memory metric +// aggregation, HTTP exposer lifecycle, port contention, and format conversions +// (counters, gauges, histograms). +class PrometheusShmExporterTest : public testing::Test { + protected: + void SetUp() override { + std::filesystem::path temp_dir = testing::TempDir(); + test_dir_ = (temp_dir / + absl::StrCat("prom_shm_test_", getpid(), "_", absl::Hex(this))) + .string(); + std::error_code ec; + std::filesystem::create_directories(test_dir_, ec); + ASSERT_FALSE(ec) << ec.message(); + } + + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(test_dir_, ec); + } + + ExporterOptions DefaultOptions(absl::string_view rank = "0", + int port = 0) const { + ExporterOptions options; + options.shm_dir = test_dir_; + options.local_rank = std::string(rank); + options.bind_address = "127.0.0.1"; + options.port = port; + return options; + } + + std::string test_dir_; +}; + +// Options Validation & Port Edge Cases + +TEST_F(PrometheusShmExporterTest, InvalidOrZeroPortGracefullyDisablesExposer) { + for (int port : {0, -100, -1, 65536, 100000}) { + PrometheusShmExporter exporter(DefaultOptions("0", port)); + EXPECT_FALSE(exporter.IsServerRunningForTesting()); + } + + // Metric emission and local snapshot function even when exposer is disabled. + PrometheusShmExporter exporter(DefaultOptions("0", 0)); + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, labels, 42); + EXPECT_THAT(exporter.GetTextSnapshot(), + HasSubstr(R"(tpu_raiden_sent_bytes_total{direction="push"} 42)")); +} + +// HTTP Port Lifecycle, Collision, Recovery & Multi-threaded Contention + +TEST_F(PrometheusShmExporterTest, + HttpExposerLifecyclePortCollisionAndFallbackRecovery) { + int port = 0; + std::unique_ptr exp1; + for (int attempt = 0; attempt < 3; ++attempt) { + port = PickUnusedPort(); + if (port <= 0) continue; + exp1 = std::make_unique(DefaultOptions("0", port)); + if (exp1->IsServerRunningForTesting()) break; + exp1.reset(); + } + if (!exp1 || !exp1->IsServerRunningForTesting()) { + GTEST_SKIP() << "No free port available for HTTP exposer test"; + } + + ExporterOptions opt2 = DefaultOptions("1", port); + + // Secondary exporter detects collision, degrades gracefully without crash. + PrometheusShmExporter exp2(opt2); + EXPECT_FALSE(exp2.IsServerRunningForTesting()); + + // Both workers can record metrics into shared memory. + const MetricLabel labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exp1->IncrementCounter(metric_names::kSentBytesTotal, labels, 100); + exp2.IncrementCounter(metric_names::kSentBytesTotal, labels, 200); + + // Primary serves aggregated metrics from both workers. + EXPECT_TRUE(exp1->IsServerRunningForTesting()); + EXPECT_THAT(exp1->GetTextSnapshot(), + HasSubstr(R"(tpu_raiden_sent_bytes_total{)" + R"(direction="push"} 300)")); + + // In-process snapshot from secondary also reflects the aggregated total. + EXPECT_THAT(exp2.GetTextSnapshot(), + HasSubstr(R"(tpu_raiden_sent_bytes_total{)" + R"(direction="push"} 300)")); + + // Destroy primary exporter to release the port. + exp1.reset(); + + // A newly constructed exporter on that port can now bind and recover. + PrometheusShmExporter exp3(DefaultOptions("0", port)); + EXPECT_TRUE(exp3.IsServerRunningForTesting()); + EXPECT_THAT(exp3.GetTextSnapshot(), + HasSubstr(R"(tpu_raiden_sent_bytes_total{)" + R"(direction="push"} 200)")); +} + +TEST_F(PrometheusShmExporterTest, MultiThreadedPortContentionExactlyOneBinds) { + constexpr int kNumExporters = 8; + const int port = PickUnusedPort(); + if (port <= 0) { + GTEST_SKIP() << "No free port available for contention test"; + } + + std::vector> exporters(kNumExporters); + absl::Notification start_gate; + std::vector threads; + threads.reserve(kNumExporters); + + for (int i = 0; i < kNumExporters; ++i) { + threads.emplace_back([this, port, i, &exporters, &start_gate]() { + start_gate.WaitForNotification(); + exporters[i] = std::make_unique( + DefaultOptions(absl::StrCat(i), port)); + }); + } + + start_gate.Notify(); + for (std::thread& t : threads) { + t.join(); + } + + std::vector running_ranks; + for (int i = 0; i < kNumExporters; ++i) { + if (exporters[i] != nullptr && exporters[i]->IsServerRunningForTesting()) { + running_ranks.push_back(i); + } + } + + // Exactly 1 exporter must succeed in binding; the other 7 fail gracefully. + EXPECT_EQ(running_ranks.size(), 1) + << "Port contention test failed on port " << port << "; running ranks: [" + << absl::StrJoin(running_ranks, ", ") << "]"; +} + +// Metric Translation, Aggregation & Exposition Snapshots + +TEST_F(PrometheusShmExporterTest, MetricFamilyTranslationWithHistograms) { + PrometheusShmExporter exporter(DefaultOptions()); + + const MetricLabel push_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPush}}; + exporter.IncrementCounter(metric_names::kSentBytesTotal, push_labels, 1024); + + const MetricLabel fail_labels[] = { + {metric_labels::kDirection, metric_labels::kDirectionPull}, + {metric_labels::kErrorCode, "DEADLINE_EXCEEDED"}}; + exporter.IncrementCounter(metric_names::kTransferFailuresTotal, fail_labels, + 3); + exporter.IncrementCounter(metric_names::kSentBytesTotal, {}, 512); + exporter.SetGauge(metric_names::kBufferAllocatedBytes, {}, 65536.0); + + // Histograms: boundary values and rejection of non-finite values (NaN, Inf). + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 0.05); // bucket <= 0.1 + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 0.1); // bucket <= 0.1 exact boundary + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 5.5); // bucket <= 10.0 + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 50000.0); // bucket <= 50000.0 exact boundary + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + 60000.0); // bucket <= +Inf + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + std::numeric_limits::quiet_NaN()); + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + std::numeric_limits::infinity()); + exporter.ObserveHistogram(metric_names::kTransferDurationMs, push_labels, + -std::numeric_limits::infinity()); + + // Unlabelled histogram observation. + exporter.ObserveHistogram(metric_names::kH2dTransferTimeMs, {}, 15.5); + + const std::string snapshot = exporter.GetTextSnapshot(); + + // Counters: unlabelled and labelled translations. + EXPECT_THAT(snapshot, HasSubstr("# HELP tpu_raiden_sent_bytes_total")); + EXPECT_THAT(snapshot, + HasSubstr("# TYPE tpu_raiden_sent_bytes_total counter")); + EXPECT_THAT(snapshot, HasSubstr("tpu_raiden_sent_bytes_total 512\n")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_sent_bytes_total{direction="push"} 1024)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_failures_total{direction="pull",)" + R"(error_code="DEADLINE_EXCEEDED"} 3)")); + + // Gauge: automatic local_rank label attachment. + EXPECT_THAT(snapshot, + HasSubstr("# TYPE tpu_raiden_buffer_allocated_bytes gauge")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_buffer_allocated_bytes{local_rank="0"} 65536)")); + + // Histogram: sample count, sum, and cumulative bucket progression. + EXPECT_THAT(snapshot, + HasSubstr("# TYPE tpu_raiden_transfer_duration_ms histogram")); + EXPECT_THAT( + snapshot, + HasSubstr( + R"(tpu_raiden_transfer_duration_ms_count{direction="push"} 5)")); + EXPECT_THAT(snapshot, HasSubstr(R"(tpu_raiden_transfer_duration_ms_sum{)" + R"(direction="push"} 110005.65)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(le="0.1"} 2)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(le="5"} 2)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(le="10"} 3)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(le="25000"} 3)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(le="50000"} 4)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(le="+Inf"} 5)")); + + // Non-finite values (NaN, -Inf) must not pollute the snapshot. Note: +Inf + // rejection cannot be verified by substring check because the Prometheus + // histogram bucket label le="+Inf" is legitimately present; its rejection is + // verified above by sample_count == 5 and the exact finite sample_sum. + EXPECT_THAT(snapshot, Not(HasSubstr("nan"))); + EXPECT_THAT(snapshot, Not(HasSubstr("NaN"))); + EXPECT_THAT(snapshot, Not(HasSubstr("-Inf"))); + + // Unlabelled histogram exposition. + EXPECT_THAT(snapshot, + HasSubstr("# TYPE tpu_raiden_h2d_transfer_time_ms histogram")); + EXPECT_THAT(snapshot, HasSubstr("tpu_raiden_h2d_transfer_time_ms_count 1")); + EXPECT_THAT(snapshot, HasSubstr("tpu_raiden_h2d_transfer_time_ms_sum 15.5")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_h2d_transfer_time_ms_bucket{le="25"} 1)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_h2d_transfer_time_ms_bucket{le="+Inf"} 1)")); +} + +TEST_F(PrometheusShmExporterTest, LabelParsingSortingAndSanitization) { + PrometheusShmExporter exporter(DefaultOptions()); + + const MetricLabel unsorted_labels[] = { + {"zebra", "last"}, {"combo", R"(val;1=2\3)"}, {"empty_val", ""}, + {"apple", "first"}, {"path", R"(C:\dir\file)"}, + }; + exporter.IncrementCounter(metric_names::kSentBytesTotal, unsorted_labels, 42); + + const MetricLabel hist_labels[] = { + {"direction", "push"}, + {"le", "custom_val"}, + {"test_tag", "esc;le=val"}, + }; + exporter.ObserveHistogram(metric_names::kTransferDurationMs, hist_labels, + 10.0); + + const std::string snapshot = exporter.GetTextSnapshot(); + + // Escaped delimiters, empty values, and lexicographical label sorting: + // apple < combo < empty_val < path < zebra. + EXPECT_THAT( + snapshot, + HasSubstr( + R"(tpu_raiden_sent_bytes_total{apple="first",combo="val;1=2\\3",)" + R"(empty_val="",path="C:\\dir\\file",zebra="last"} 42)")); + + // Reserved "le" label must be sanitized out while preserving escaped + // delimiters in other labels. + EXPECT_THAT(snapshot, Not(HasSubstr("custom_val"))); + EXPECT_THAT(snapshot, HasSubstr(R"(test_tag="esc;le=val")")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_count{direction="push",)" + R"(test_tag="esc;le=val"} 1)")); + EXPECT_THAT( + snapshot, + HasSubstr(R"(tpu_raiden_transfer_duration_ms_bucket{direction="push",)" + R"(test_tag="esc;le=val",le="10"} 1)")); +} + +TEST_F(PrometheusShmExporterTest, EmptyExporterReturnsNoFamiliesOrSnapshot) { + PrometheusShmExporter exporter(DefaultOptions()); + EXPECT_EQ(exporter.GetTextSnapshot(), ""); +} + +} // namespace +} // namespace tpu_raiden::telemetry