From 4109b086f62328ae502c566601af07b9ecd69c48 Mon Sep 17 00:00:00 2001 From: Googler Date: Fri, 31 Jul 2026 00:01:30 -0700 Subject: [PATCH] [tpu_raiden] Add C++ 3P Prometheus exporter backend PiperOrigin-RevId: 956957621 --- tpu_raiden/telemetry/BUILD | 71 ++++ tpu_raiden/telemetry/metrics_api.cc | 133 ++++++++ tpu_raiden/telemetry/metrics_api.h | 135 ++++++++ tpu_raiden/telemetry/metrics_api_test.cc | 219 ++++++++++++ tpu_raiden/telemetry/prometheus_exporter.cc | 313 ++++++++++++++++++ tpu_raiden/telemetry/prometheus_exporter.h | 143 ++++++++ .../telemetry/prometheus_exporter_test.cc | 268 +++++++++++++++ 7 files changed, 1282 insertions(+) create mode 100644 tpu_raiden/telemetry/BUILD create mode 100644 tpu_raiden/telemetry/metrics_api.cc create mode 100644 tpu_raiden/telemetry/metrics_api.h create mode 100644 tpu_raiden/telemetry/metrics_api_test.cc create mode 100644 tpu_raiden/telemetry/prometheus_exporter.cc create mode 100644 tpu_raiden/telemetry/prometheus_exporter.h create mode 100644 tpu_raiden/telemetry/prometheus_exporter_test.cc diff --git a/tpu_raiden/telemetry/BUILD b/tpu_raiden/telemetry/BUILD new file mode 100644 index 00000000..11bdb0c3 --- /dev/null +++ b/tpu_raiden/telemetry/BUILD @@ -0,0 +1,71 @@ +# 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. + +# Copyright 2026 Google LLC + +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") + +package(default_visibility = ["//visibility:public"]) + +licenses(["notice"]) + +cc_library( + name = "metrics_api", + srcs = ["metrics_api.cc"], + hdrs = ["metrics_api.h"], + deps = [ + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "metrics_api_test", + srcs = ["metrics_api_test.cc"], + deps = [ + ":metrics_api", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "metrics_3p_prometheus_exporter", + srcs = ["prometheus_exporter.cc"], + hdrs = ["prometheus_exporter.h"], + deps = [ + ":metrics_api", + "//third_party/prometheus_cpp_client:prometheus_client_core", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/synchronization", + ], +) + +cc_test( + name = "prometheus_exporter_test", + srcs = ["prometheus_exporter_test.cc"], + deps = [ + ":metrics_3p_prometheus_exporter", + ":metrics_api", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_raiden/telemetry/metrics_api.cc b/tpu_raiden/telemetry/metrics_api.cc new file mode 100644 index 00000000..f628f6d5 --- /dev/null +++ b/tpu_raiden/telemetry/metrics_api.cc @@ -0,0 +1,133 @@ +// 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. + +// 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_raiden/telemetry/metrics_api.h" + +#include +#include +#include +#include +#include +#include + +#include "absl/base/no_destructor.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" + +namespace tpu_raiden::telemetry { + +absl::string_view GetMetricDescription(absl::string_view name) { + static const absl::NoDestructor< + absl::flat_hash_map> + kDescriptionMap({ + {metric_names::kSentBytesTotal, + "Total count of bytes sent over TPU Raiden interfaces."}, + {metric_names::kReceivedBytesTotal, + "Total count of bytes received over TPU Raiden interfaces."}, + {metric_names::kTransferDurationSeconds, + "Histogram of TPU Raiden transfer durations in seconds."}, + {metric_names::kStageLatencySeconds, + "Histogram of TPU Raiden pipeline stage latencies in seconds."}, + {metric_names::kActiveTransfers, + "Number of currently active TPU Raiden data transfers."}, + {metric_names::kBufferOccupancyBytes, + "Gauge of TPU Raiden buffer occupancy in bytes."}, + {metric_names::kTransferFailuresTotal, + "Total count of failed TPU Raiden data transfers."}, + }); + + auto it = kDescriptionMap->find(name); + if (it != kDescriptionMap->end()) { + return it->second; + } + return name; +} + +RaidenMetricStore& RaidenMetricStore::GetGlobalMetricStore() { + static absl::NoDestructor global_store; + return *global_store; +} + +void RaidenMetricStore::AddBackend(std::unique_ptr backend) { + if (!backend) return; + absl::MutexLock lock(mutex_); + backends_.push_back(std::move(backend)); + has_backends_.store(true, std::memory_order_release); +} + +void RaidenMetricStore::ClearBackends() { + absl::MutexLock lock(mutex_); + has_backends_.store(false, std::memory_order_release); + backends_.clear(); +} + +bool RaidenMetricStore::HasBackends() const { + return has_backends_.load(std::memory_order_acquire); +} + +void RaidenMetricStore::IncrementCounter(absl::string_view name, + LabelSpan labels, uint64_t val) const { + if (!has_backends_.load(std::memory_order_acquire)) return; + // TODO: Explore RCU optimization for lock-free reads. + absl::ReaderMutexLock lock(mutex_); + for (const auto& backend : backends_) { + backend->IncrementCounter(name, labels, val); + } +} + +void RaidenMetricStore::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + if (!has_backends_.load(std::memory_order_acquire)) return; + absl::ReaderMutexLock lock(mutex_); + for (const auto& backend : backends_) { + backend->SetGauge(name, labels, val); + } +} + +void RaidenMetricStore::ObserveHistogram(absl::string_view name, + LabelSpan labels, double val) const { + if (!has_backends_.load(std::memory_order_acquire)) return; + absl::ReaderMutexLock lock(mutex_); + for (const auto& backend : backends_) { + backend->ObserveHistogram(name, labels, val); + } +} + +std::string RaidenMetricStore::GetTextSnapshot() const { + if (!has_backends_.load(std::memory_order_acquire)) return ""; + absl::ReaderMutexLock lock(mutex_); + std::string result; + for (const auto& backend : backends_) { + absl::StrAppend(&result, backend->GetTextSnapshot()); + } + return result; +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/metrics_api.h b/tpu_raiden/telemetry/metrics_api.h new file mode 100644 index 00000000..fdc31845 --- /dev/null +++ b/tpu_raiden/telemetry/metrics_api.h @@ -0,0 +1,135 @@ +// 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. + +// 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_RAIDEN_TELEMETRY_METRICS_API_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_METRICS_API_H_ + +#include +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" + +namespace tpu_raiden::telemetry { + +namespace metric_names { + +inline constexpr absl::string_view kSentBytesTotal = "sent_bytes_total"; +inline constexpr absl::string_view kReceivedBytesTotal = "received_bytes_total"; +inline constexpr absl::string_view kTransferDurationSeconds = + "transfer_duration_seconds"; +inline constexpr absl::string_view kStageLatencySeconds = + "stage_latency_seconds"; +inline constexpr absl::string_view kActiveTransfers = "active_transfers"; +inline constexpr absl::string_view kBufferOccupancyBytes = + "buffer_occupancy_bytes"; +inline constexpr absl::string_view kTransferFailuresTotal = + "transfer_failures_total"; + +} // namespace metric_names + +// Returns the common human-readable description for a given metric name, +// falling back to the metric name itself if unmapped. This description is +// common across all metric exporter backends (e.g., Prometheus and Streamz). +absl::string_view GetMetricDescription(absl::string_view name); + +// Structure defining a metric key-value label pair. +struct MetricLabel { + absl::string_view key; + absl::string_view value; +}; + +// Allocation-free label view span type definition +using LabelSpan = absl::Span; + +// Abstract Dual-Backend Interface +class MetricsBackend { + public: + MetricsBackend() = default; + MetricsBackend(const MetricsBackend&) = delete; + MetricsBackend& operator=(const MetricsBackend&) = delete; + MetricsBackend(MetricsBackend&&) = delete; + MetricsBackend& operator=(MetricsBackend&&) = delete; + + virtual ~MetricsBackend() = default; + + virtual void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const = 0; + + virtual void SetGauge(absl::string_view name, LabelSpan labels, + double val) const = 0; + + virtual void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const = 0; + + virtual std::string GetTextSnapshot() const = 0; +}; + +// Central Telemetry Facade for managing metrics across registered backends. +// This class is thread-safe for all concurrent operations. +class RaidenMetricStore { + public: + static RaidenMetricStore& GetGlobalMetricStore(); + + RaidenMetricStore() = default; + ~RaidenMetricStore() = default; + + RaidenMetricStore(const RaidenMetricStore&) = delete; + RaidenMetricStore& operator=(const RaidenMetricStore&) = delete; + RaidenMetricStore(RaidenMetricStore&&) = delete; + RaidenMetricStore& operator=(RaidenMetricStore&&) = delete; + + void AddBackend(std::unique_ptr backend); + void ClearBackends(); + bool HasBackends() const; + + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val = 1) const; + + void SetGauge(absl::string_view name, LabelSpan labels, double val) const; + + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const; + + std::string GetTextSnapshot() const; + + private: + mutable absl::Mutex mutex_; + std::vector> backends_ + ABSL_GUARDED_BY(mutex_); + std::atomic has_backends_{false}; +}; + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_METRICS_API_H_ diff --git a/tpu_raiden/telemetry/metrics_api_test.cc b/tpu_raiden/telemetry/metrics_api_test.cc new file mode 100644 index 00000000..e77de3f4 --- /dev/null +++ b/tpu_raiden/telemetry/metrics_api_test.cc @@ -0,0 +1,219 @@ +// 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. + +// 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_raiden/telemetry/metrics_api.h" + +#include +#include +#include +#include // NOLINT(build/c++11) +#include +#include +#include + +#include +#include +#include "absl/strings/string_view.h" + +namespace tpu_raiden::telemetry { +namespace { + +static_assert(!std::is_copy_constructible_v, + "MetricsBackend must not be copy constructible"); +static_assert(!std::is_copy_assignable_v, + "MetricsBackend must not be copy assignable"); +static_assert(!std::is_move_constructible_v, + "MetricsBackend must not be move constructible"); +static_assert(!std::is_move_assignable_v, + "MetricsBackend must not be move assignable"); + +using testing::_; +using testing::Eq; +using testing::Return; + +class MockMetricsBackend : public MetricsBackend { + public: + MOCK_METHOD(void, IncrementCounter, + (absl::string_view name, LabelSpan labels, uint64_t val), + (override, const)); + MOCK_METHOD(void, SetGauge, + (absl::string_view name, LabelSpan labels, double val), + (override, const)); + MOCK_METHOD(void, ObserveHistogram, + (absl::string_view name, LabelSpan labels, double val), + (override, const)); + MOCK_METHOD(std::string, GetTextSnapshot, (), (override, const)); +}; + +class MetricsApiTest : public testing::Test { + protected: + RaidenMetricStore store_; +}; + +TEST_F(MetricsApiTest, GlobalMetricStoreSingleton) { + RaidenMetricStore& global1 = RaidenMetricStore::GetGlobalMetricStore(); + RaidenMetricStore& global2 = RaidenMetricStore::GetGlobalMetricStore(); + EXPECT_EQ(&global1, &global2); +} + +TEST_F(MetricsApiTest, GetMetricDescription) { + EXPECT_EQ(GetMetricDescription(metric_names::kSentBytesTotal), + "Total count of bytes sent over TPU Raiden interfaces."); + EXPECT_EQ(GetMetricDescription(metric_names::kReceivedBytesTotal), + "Total count of bytes received over TPU Raiden interfaces."); + EXPECT_EQ(GetMetricDescription(metric_names::kTransferDurationSeconds), + "Histogram of TPU Raiden transfer durations in seconds."); + EXPECT_EQ(GetMetricDescription(metric_names::kStageLatencySeconds), + "Histogram of TPU Raiden pipeline stage latencies in seconds."); + EXPECT_EQ(GetMetricDescription(metric_names::kActiveTransfers), + "Number of currently active TPU Raiden data transfers."); + EXPECT_EQ(GetMetricDescription(metric_names::kBufferOccupancyBytes), + "Gauge of TPU Raiden buffer occupancy in bytes."); + EXPECT_EQ(GetMetricDescription(metric_names::kTransferFailuresTotal), + "Total count of failed TPU Raiden data transfers."); + EXPECT_EQ(GetMetricDescription("unknown_metric"), "unknown_metric"); +} + +TEST_F(MetricsApiTest, FastPathExitWhenNoBackends) { + EXPECT_FALSE(store_.HasBackends()); + + store_.IncrementCounter(metric_names::kSentBytesTotal, {}, 1024); + store_.SetGauge(metric_names::kActiveTransfers, {}, 5); + store_.ObserveHistogram(metric_names::kTransferDurationSeconds, {}, 0.0125); + store_.ObserveHistogram(metric_names::kStageLatencySeconds, {}, 0.025); + store_.SetGauge(metric_names::kBufferOccupancyBytes, {}, 4096); + store_.IncrementCounter(metric_names::kTransferFailuresTotal, {}, 1); +} + +TEST_F(MetricsApiTest, DispatchesToRegisteredBackend) { + auto mock_backend = std::make_unique(); + MockMetricsBackend* raw_mock = mock_backend.get(); + + EXPECT_CALL(*raw_mock, + IncrementCounter(Eq(metric_names::kSentBytesTotal), _, 2048)) + .Times(1); + EXPECT_CALL(*raw_mock, SetGauge(Eq(metric_names::kActiveTransfers), _, 3)) + .Times(1); + EXPECT_CALL( + *raw_mock, + ObserveHistogram(Eq(metric_names::kTransferDurationSeconds), _, 0.005)) + .Times(1); + EXPECT_CALL(*raw_mock, ObserveHistogram( + Eq(metric_names::kStageLatencySeconds), _, 0.015)) + .Times(1); + EXPECT_CALL(*raw_mock, + SetGauge(Eq(metric_names::kBufferOccupancyBytes), _, 8192)) + .Times(1); + EXPECT_CALL(*raw_mock, + IncrementCounter(Eq(metric_names::kTransferFailuresTotal), _, 2)) + .Times(1); + EXPECT_CALL(*raw_mock, GetTextSnapshot()).WillOnce(Return("# HELP mock\n")); + + store_.AddBackend(std::move(mock_backend)); + EXPECT_TRUE(store_.HasBackends()); + + store_.IncrementCounter(metric_names::kSentBytesTotal, {}, 2048); + store_.SetGauge(metric_names::kActiveTransfers, {}, 3); + store_.ObserveHistogram(metric_names::kTransferDurationSeconds, {}, 0.005); + store_.ObserveHistogram(metric_names::kStageLatencySeconds, {}, 0.015); + store_.SetGauge(metric_names::kBufferOccupancyBytes, {}, 8192); + store_.IncrementCounter(metric_names::kTransferFailuresTotal, {}, 2); + EXPECT_EQ(store_.GetTextSnapshot(), "# HELP mock\n"); +} + +TEST_F(MetricsApiTest, ClearBackendsResetsFastPath) { + auto mock_backend = std::make_unique(); + MockMetricsBackend* raw_mock = mock_backend.get(); + + EXPECT_CALL(*raw_mock, IncrementCounter(_, _, _)).Times(0); + + store_.AddBackend(std::move(mock_backend)); + EXPECT_TRUE(store_.HasBackends()); + + store_.ClearBackends(); + EXPECT_FALSE(store_.HasBackends()); + + store_.IncrementCounter(metric_names::kReceivedBytesTotal, {}, 1); +} + +TEST_F(MetricsApiTest, ConcurrentTelemetryEmissions) { + auto backend = std::make_unique(); + MockMetricsBackend* raw_backend = backend.get(); + + constexpr int kNumThreads = 8; + constexpr int kIterations = 100; + constexpr int kTotalCalls = kNumThreads * kIterations; + + EXPECT_CALL(*raw_backend, IncrementCounter(Eq("counter"), _, 1)) + .Times(kTotalCalls); + EXPECT_CALL(*raw_backend, SetGauge(Eq("gauge"), _, 42)) + .Times(kTotalCalls); + EXPECT_CALL(*raw_backend, ObserveHistogram(Eq("histogram"), _, 3.14)) + .Times(kTotalCalls); + EXPECT_CALL(*raw_backend, GetTextSnapshot()) + .Times(kTotalCalls) + .WillRepeatedly(Return("snapshot\n")); + + store_.AddBackend(std::move(backend)); + + std::vector threads; + threads.reserve(kNumThreads); + + for (int i = 0; i < kNumThreads; ++i) { + threads.emplace_back([this] { + for (int j = 0; j < kIterations; ++j) { + store_.IncrementCounter("counter", {}, 1); + store_.SetGauge("gauge", {}, 42); + store_.ObserveHistogram("histogram", {}, 3.14); + (void)store_.GetTextSnapshot(); + } + }); + } + + for (auto& t : threads) { + t.join(); + } +} + +TEST_F(MetricsApiTest, ConstMetricsBackendReference) { + MockMetricsBackend backend; + const MetricsBackend& const_ref = backend; + + EXPECT_CALL(backend, IncrementCounter(Eq("counter"), _, 5)).Times(1); + EXPECT_CALL(backend, SetGauge(Eq("gauge"), _, 10)).Times(1); + EXPECT_CALL(backend, ObserveHistogram(Eq("histogram"), _, 1.23)).Times(1); + EXPECT_CALL(backend, GetTextSnapshot()).WillOnce(Return("snapshot\n")); + + const_ref.IncrementCounter("counter", {}, 5); + const_ref.SetGauge("gauge", {}, 10); + const_ref.ObserveHistogram("histogram", {}, 1.23); + EXPECT_EQ(const_ref.GetTextSnapshot(), "snapshot\n"); +} + +} // namespace +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/prometheus_exporter.cc b/tpu_raiden/telemetry/prometheus_exporter.cc new file mode 100644 index 00000000..7758adae --- /dev/null +++ b/tpu_raiden/telemetry/prometheus_exporter.cc @@ -0,0 +1,313 @@ +// 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. + +// 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_raiden/telemetry/prometheus_exporter.h" + +#include +#include +#include +#include +#include + +#include "prometheus/counter.h" +#include "prometheus/family.h" +#include "prometheus/gauge.h" +#include "prometheus/histogram.h" +#include "prometheus/registry.h" +#include "prometheus/text_serializer.h" +#include "absl/base/no_destructor.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "tpu_raiden/telemetry/metrics_api.h" + +namespace tpu_raiden::telemetry { + +namespace { + +std::map ConvertLabels(LabelSpan labels) { + if (labels.empty()) { + return {}; + } + std::map result; + for (const auto& [key, value] : labels) { + result.emplace(key, value); + } + return result; +} + +} // namespace + +namespace prometheus_names { + +inline constexpr absl::string_view kSentBytesTotal = + "tpu_raiden_sent_bytes_total"; +inline constexpr absl::string_view kReceivedBytesTotal = + "tpu_raiden_received_bytes_total"; +inline constexpr absl::string_view kTransferDurationSeconds = + "tpu_raiden_transfer_duration_seconds"; +inline constexpr absl::string_view kStageLatencySeconds = + "tpu_raiden_stage_latency_seconds"; +inline constexpr absl::string_view kActiveTransfers = + "tpu_raiden_active_transfers"; +inline constexpr absl::string_view kBufferOccupancyBytes = + "tpu_raiden_buffer_occupancy_bytes"; +inline constexpr absl::string_view kTransferFailuresTotal = + "tpu_raiden_transfer_failures_total"; + +} // namespace prometheus_names + +PrometheusMetricInfo GetPrometheusMetricInfo(absl::string_view name) { + absl::string_view prom_name = name; + absl::string_view desc = GetMetricDescription(name); + if (name == metric_names::kSentBytesTotal || + name == prometheus_names::kSentBytesTotal) { + prom_name = prometheus_names::kSentBytesTotal; + desc = GetMetricDescription(metric_names::kSentBytesTotal); + } else if (name == metric_names::kReceivedBytesTotal || + name == prometheus_names::kReceivedBytesTotal) { + prom_name = prometheus_names::kReceivedBytesTotal; + desc = GetMetricDescription(metric_names::kReceivedBytesTotal); + } else if (name == metric_names::kTransferDurationSeconds || + name == prometheus_names::kTransferDurationSeconds) { + prom_name = prometheus_names::kTransferDurationSeconds; + desc = GetMetricDescription(metric_names::kTransferDurationSeconds); + } else if (name == metric_names::kStageLatencySeconds || + name == prometheus_names::kStageLatencySeconds) { + prom_name = prometheus_names::kStageLatencySeconds; + desc = GetMetricDescription(metric_names::kStageLatencySeconds); + } else if (name == metric_names::kActiveTransfers || + name == prometheus_names::kActiveTransfers) { + prom_name = prometheus_names::kActiveTransfers; + desc = GetMetricDescription(metric_names::kActiveTransfers); + } else if (name == metric_names::kBufferOccupancyBytes || + name == prometheus_names::kBufferOccupancyBytes) { + prom_name = prometheus_names::kBufferOccupancyBytes; + desc = GetMetricDescription(metric_names::kBufferOccupancyBytes); + } else if (name == metric_names::kTransferFailuresTotal || + name == prometheus_names::kTransferFailuresTotal) { + prom_name = prometheus_names::kTransferFailuresTotal; + desc = GetMetricDescription(metric_names::kTransferFailuresTotal); + } + return PrometheusMetricInfo{.prometheus_name = prom_name, .help_text = desc}; +} + +void PrometheusExporter::RegisterKnownFamilies() { + sent_bytes_family_ = + &prometheus::BuildCounter() + .Name(std::string(prometheus_names::kSentBytesTotal)) + .Help(std::string( + GetMetricDescription(metric_names::kSentBytesTotal))) + .Register(*registry_); + received_bytes_family_ = + &prometheus::BuildCounter() + .Name(std::string(prometheus_names::kReceivedBytesTotal)) + .Help(std::string( + GetMetricDescription(metric_names::kReceivedBytesTotal))) + .Register(*registry_); + transfer_failures_family_ = + &prometheus::BuildCounter() + .Name(std::string(prometheus_names::kTransferFailuresTotal)) + .Help(std::string( + GetMetricDescription(metric_names::kTransferFailuresTotal))) + .Register(*registry_); + + active_transfers_family_ = + &prometheus::BuildGauge() + .Name(std::string(prometheus_names::kActiveTransfers)) + .Help(std::string( + GetMetricDescription(metric_names::kActiveTransfers))) + .Register(*registry_); + buffer_occupancy_family_ = + &prometheus::BuildGauge() + .Name(std::string(prometheus_names::kBufferOccupancyBytes)) + .Help(std::string( + GetMetricDescription(metric_names::kBufferOccupancyBytes))) + .Register(*registry_); + + transfer_duration_family_ = + &prometheus::BuildHistogram() + .Name(std::string(prometheus_names::kTransferDurationSeconds)) + .Help(std::string(GetMetricDescription( + metric_names::kTransferDurationSeconds))) + .Register(*registry_); + stage_latency_family_ = + &prometheus::BuildHistogram() + .Name(std::string(prometheus_names::kStageLatencySeconds)) + .Help(std::string( + GetMetricDescription(metric_names::kStageLatencySeconds))) + .Register(*registry_); +} + +PrometheusExporter::PrometheusExporter( + prometheus::Histogram::BucketBoundaries custom_buckets) + : registry_(std::make_shared()), + default_buckets_(std::move(custom_buckets)) { + absl::MutexLock lock(&mutex_); + RegisterKnownFamilies(); +} + +prometheus::Family* PrometheusExporter::GetCounterFamily( + absl::string_view name) const { + if (name == metric_names::kSentBytesTotal || + name == prometheus_names::kSentBytesTotal) { + return sent_bytes_family_; + } + if (name == metric_names::kReceivedBytesTotal || + name == prometheus_names::kReceivedBytesTotal) { + return received_bytes_family_; + } + if (name == metric_names::kTransferFailuresTotal || + name == prometheus_names::kTransferFailuresTotal) { + return transfer_failures_family_; + } + PrometheusMetricInfo info = GetPrometheusMetricInfo(name); + auto it = counter_families_.find(info.prometheus_name); + if (it != counter_families_.end()) { + return it->second; + } + auto* family = &prometheus::BuildCounter() + .Name(std::string(info.prometheus_name)) + .Help(std::string(info.help_text)) + .Register(*registry_); + counter_families_.emplace(std::string(info.prometheus_name), family); + return family; +} + +prometheus::Family* PrometheusExporter::GetGaugeFamily( + absl::string_view name) const { + if (name == metric_names::kActiveTransfers || + name == prometheus_names::kActiveTransfers) { + return active_transfers_family_; + } + if (name == metric_names::kBufferOccupancyBytes || + name == prometheus_names::kBufferOccupancyBytes) { + return buffer_occupancy_family_; + } + PrometheusMetricInfo info = GetPrometheusMetricInfo(name); + auto it = gauge_families_.find(info.prometheus_name); + if (it != gauge_families_.end()) { + return it->second; + } + auto* family = &prometheus::BuildGauge() + .Name(std::string(info.prometheus_name)) + .Help(std::string(info.help_text)) + .Register(*registry_); + gauge_families_.emplace(std::string(info.prometheus_name), family); + return family; +} + +prometheus::Family* +PrometheusExporter::GetHistogramFamily(absl::string_view name) const { + if (name == metric_names::kTransferDurationSeconds || + name == prometheus_names::kTransferDurationSeconds) { + return transfer_duration_family_; + } + if (name == metric_names::kStageLatencySeconds || + name == prometheus_names::kStageLatencySeconds) { + return stage_latency_family_; + } + PrometheusMetricInfo info = GetPrometheusMetricInfo(name); + auto it = histogram_families_.find(info.prometheus_name); + if (it != histogram_families_.end()) { + return it->second; + } + auto* family = &prometheus::BuildHistogram() + .Name(std::string(info.prometheus_name)) + .Help(std::string(info.help_text)) + .Register(*registry_); + histogram_families_.emplace(std::string(info.prometheus_name), family); + return family; +} + +void PrometheusExporter::IncrementCounter(absl::string_view name, + LabelSpan labels, + uint64_t val) const { + std::shared_ptr current_registry; + prometheus::Family* family = nullptr; + { + absl::MutexLock lock(&mutex_); + current_registry = registry_; + family = GetCounterFamily(name); + } + prometheus::Counter& counter = family->Add(ConvertLabels(labels)); + counter.Increment(static_cast(val)); +} + +void PrometheusExporter::SetGauge(absl::string_view name, LabelSpan labels, + double val) const { + std::shared_ptr current_registry; + prometheus::Family* family = nullptr; + { + absl::MutexLock lock(&mutex_); + current_registry = registry_; + family = GetGaugeFamily(name); + } + prometheus::Gauge& gauge = family->Add(ConvertLabels(labels)); + gauge.Set(val); +} + +void PrometheusExporter::ObserveHistogram(absl::string_view name, + LabelSpan labels, double val) const { + std::shared_ptr current_registry; + prometheus::Family* family = nullptr; + { + absl::MutexLock lock(&mutex_); + current_registry = registry_; + family = GetHistogramFamily(name); + } + prometheus::Histogram& histogram = + family->Add(ConvertLabels(labels), default_buckets_); + histogram.Observe(val); +} + +void PrometheusExporter::Reset() { + absl::MutexLock lock(&mutex_); + registry_ = std::make_shared(); + counter_families_.clear(); + gauge_families_.clear(); + histogram_families_.clear(); + RegisterKnownFamilies(); +} + +std::string PrometheusExporter::GetTextSnapshot() const { + std::shared_ptr reg; + { + absl::MutexLock lock(&mutex_); + reg = registry_; + } + prometheus::TextSerializer serializer; + return serializer.Serialize(reg->Collect()); +} + +PrometheusExporter& GetGlobalPrometheusExporter() { + static absl::NoDestructor global_exporter; + return *global_exporter; +} + +} // namespace tpu_raiden::telemetry diff --git a/tpu_raiden/telemetry/prometheus_exporter.h b/tpu_raiden/telemetry/prometheus_exporter.h new file mode 100644 index 00000000..8969af5b --- /dev/null +++ b/tpu_raiden/telemetry/prometheus_exporter.h @@ -0,0 +1,143 @@ +// 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. + +// 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_RAIDEN_TELEMETRY_PROMETHEUS_EXPORTER_H_ +#define THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_PROMETHEUS_EXPORTER_H_ + +#include +#include +#include + +#include "prometheus/counter.h" +#include "prometheus/family.h" +#include "prometheus/gauge.h" +#include "prometheus/histogram.h" +#include "prometheus/registry.h" +#include "absl/base/no_destructor.h" +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "tpu_raiden/telemetry/metrics_api.h" + +namespace tpu_raiden::telemetry { + +// Mapping structure for 3P Prometheus metric name and help text. +// The help_text uses the common metric description from GetMetricDescription. +struct PrometheusMetricInfo { + absl::string_view prometheus_name; + absl::string_view help_text; +}; + +// Returns PrometheusMetricInfo for a given general metric name/key, +// using GetMetricDescription for common human-readable help text. +PrometheusMetricInfo GetPrometheusMetricInfo(absl::string_view name); + +inline const prometheus::Histogram::BucketBoundaries& +DefaultHistogramBuckets() { + static const absl::NoDestructor + kBuckets(prometheus::Histogram::BucketBoundaries{ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0}); + return *kBuckets; +} + +class PrometheusExporter : public MetricsBackend { + public: + explicit PrometheusExporter(prometheus::Histogram::BucketBoundaries + custom_buckets = DefaultHistogramBuckets()); + + ~PrometheusExporter() override = default; + + PrometheusExporter(const PrometheusExporter&) = delete; + PrometheusExporter& operator=(const PrometheusExporter&) = delete; + PrometheusExporter(PrometheusExporter&&) = delete; + PrometheusExporter& operator=(PrometheusExporter&&) = delete; + + void IncrementCounter(absl::string_view name, LabelSpan labels, + uint64_t val) const override; + + void SetGauge(absl::string_view name, LabelSpan labels, + double val) const override; + + void ObserveHistogram(absl::string_view name, LabelSpan labels, + double val) const override; + + std::string GetTextSnapshot() const override; + + void Reset(); + + std::shared_ptr GetRegistry() const { + absl::ReaderMutexLock lock(mutex_); + return registry_; + } + + private: + void RegisterKnownFamilies() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + prometheus::Family* GetCounterFamily( + absl::string_view name) const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + prometheus::Family* GetGaugeFamily(absl::string_view name) + const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + prometheus::Family* GetHistogramFamily( + absl::string_view name) const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + mutable absl::Mutex mutex_; + std::shared_ptr registry_ ABSL_GUARDED_BY(mutex_); + prometheus::Histogram::BucketBoundaries default_buckets_; + + prometheus::Family* sent_bytes_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + prometheus::Family* received_bytes_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + prometheus::Family* transfer_failures_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + prometheus::Family* active_transfers_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + prometheus::Family* buffer_occupancy_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + prometheus::Family* transfer_duration_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + prometheus::Family* stage_latency_family_ + ABSL_GUARDED_BY(mutex_) = nullptr; + + mutable absl::flat_hash_map*> + counter_families_ ABSL_GUARDED_BY(mutex_); + mutable absl::flat_hash_map*> + gauge_families_ ABSL_GUARDED_BY(mutex_); + mutable absl::flat_hash_map*> + histogram_families_ ABSL_GUARDED_BY(mutex_); +}; + +PrometheusExporter& GetGlobalPrometheusExporter(); + +} // namespace tpu_raiden::telemetry + +#endif // THIRD_PARTY_TPU_RAIDEN_TPU_RAIDEN_TELEMETRY_PROMETHEUS_EXPORTER_H_ diff --git a/tpu_raiden/telemetry/prometheus_exporter_test.cc b/tpu_raiden/telemetry/prometheus_exporter_test.cc new file mode 100644 index 00000000..11c3d986 --- /dev/null +++ b/tpu_raiden/telemetry/prometheus_exporter_test.cc @@ -0,0 +1,268 @@ +// 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. + +// 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_raiden/telemetry/prometheus_exporter.h" + +#include +#include +#include // NOLINT(build/c++11) +#include +#include + +#include +#include +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "tpu_raiden/telemetry/metrics_api.h" + +namespace tpu_raiden::telemetry { +namespace { + +using ::testing::HasSubstr; + +TEST(PrometheusExporterTest, RecordAndExportFormat) { + auto exporter = std::make_unique(); + RaidenMetricStore store; + store.AddBackend(std::move(exporter)); + + MetricLabel label1{"interface", "ICI"}; + MetricLabel label2{"direction", "OUTBOUND"}; + MetricLabel label3{"status", "OK"}; + + store.IncrementCounter(metric_names::kSentBytesTotal, {label1}, 1024); + store.SetGauge(metric_names::kActiveTransfers, {label2}, 5); + store.ObserveHistogram(metric_names::kTransferDurationSeconds, {label3}, + 0.0055); + + std::string output = store.GetTextSnapshot(); + + EXPECT_TRUE( + absl::StrContains(output, "# TYPE tpu_raiden_sent_bytes_total counter")); + EXPECT_TRUE(absl::StrContains( + output, "tpu_raiden_sent_bytes_total{interface=\"ICI\"} 1024")); + EXPECT_TRUE( + absl::StrContains(output, "# TYPE tpu_raiden_active_transfers gauge")); + EXPECT_TRUE(absl::StrContains( + output, "tpu_raiden_active_transfers{direction=\"OUTBOUND\"} 5")); + EXPECT_TRUE(absl::StrContains( + output, "# TYPE tpu_raiden_transfer_duration_seconds histogram")); + EXPECT_TRUE(absl::StrContains( + output, + "tpu_raiden_transfer_duration_seconds_sum{status=\"OK\"} 0.0055")); + EXPECT_TRUE(absl::StrContains( + output, "tpu_raiden_transfer_duration_seconds_count{status=\"OK\"} 1")); +} + +TEST(PrometheusExporterTest, GetPrometheusMetricInfoMapped) { + PrometheusMetricInfo info1 = + GetPrometheusMetricInfo(metric_names::kSentBytesTotal); + EXPECT_EQ(info1.prometheus_name, "tpu_raiden_sent_bytes_total"); + EXPECT_THAT( + info1.help_text, + HasSubstr("Total count of bytes sent over TPU Raiden interfaces.")); + + PrometheusMetricInfo info2 = + GetPrometheusMetricInfo(metric_names::kReceivedBytesTotal); + EXPECT_EQ(info2.prometheus_name, "tpu_raiden_received_bytes_total"); + + PrometheusMetricInfo info3 = + GetPrometheusMetricInfo(metric_names::kTransferDurationSeconds); + EXPECT_EQ(info3.prometheus_name, "tpu_raiden_transfer_duration_seconds"); + + PrometheusMetricInfo info4 = + GetPrometheusMetricInfo(metric_names::kStageLatencySeconds); + EXPECT_EQ(info4.prometheus_name, "tpu_raiden_stage_latency_seconds"); + + PrometheusMetricInfo info5 = + GetPrometheusMetricInfo(metric_names::kActiveTransfers); + EXPECT_EQ(info5.prometheus_name, "tpu_raiden_active_transfers"); + + PrometheusMetricInfo info6 = + GetPrometheusMetricInfo(metric_names::kBufferOccupancyBytes); + EXPECT_EQ(info6.prometheus_name, "tpu_raiden_buffer_occupancy_bytes"); + + PrometheusMetricInfo info7 = + GetPrometheusMetricInfo(metric_names::kTransferFailuresTotal); + EXPECT_EQ(info7.prometheus_name, "tpu_raiden_transfer_failures_total"); + + // Also verify alias lookup via 3P string + PrometheusMetricInfo info_alias = + GetPrometheusMetricInfo("tpu_raiden_sent_bytes_total"); + EXPECT_EQ(info_alias.prometheus_name, "tpu_raiden_sent_bytes_total"); + EXPECT_EQ(info_alias.help_text, info1.help_text); +} + +TEST(PrometheusExporterTest, AllPhase0MetricsMapped) { + PrometheusExporter exporter; + + exporter.IncrementCounter(metric_names::kSentBytesTotal, {}, 100); + exporter.IncrementCounter(metric_names::kReceivedBytesTotal, {}, 200); + exporter.ObserveHistogram(metric_names::kTransferDurationSeconds, {}, 0.5); + exporter.ObserveHistogram(metric_names::kStageLatencySeconds, {}, 0.1); + exporter.SetGauge(metric_names::kActiveTransfers, {}, 10); + exporter.SetGauge(metric_names::kBufferOccupancyBytes, {}, 2048); + exporter.IncrementCounter(metric_names::kTransferFailuresTotal, {}, 1); + + std::string output = exporter.GetTextSnapshot(); + + EXPECT_TRUE( + absl::StrContains(output, + "# HELP tpu_raiden_sent_bytes_total Total count of " + "bytes sent over TPU Raiden interfaces.")); + EXPECT_TRUE( + absl::StrContains(output, + "# HELP tpu_raiden_received_bytes_total Total count of " + "bytes received over TPU Raiden interfaces.")); + EXPECT_TRUE( + absl::StrContains(output, + "# HELP tpu_raiden_transfer_duration_seconds Histogram " + "of TPU Raiden transfer durations in seconds.")); + EXPECT_TRUE( + absl::StrContains(output, + "# HELP tpu_raiden_stage_latency_seconds Histogram of " + "TPU Raiden pipeline stage latencies in seconds.")); + EXPECT_TRUE(absl::StrContains(output, + "# HELP tpu_raiden_active_transfers Number of " + "currently active TPU Raiden data transfers.")); + EXPECT_TRUE( + absl::StrContains(output, + "# HELP tpu_raiden_buffer_occupancy_bytes Gauge of TPU " + "Raiden buffer occupancy in bytes.")); + EXPECT_TRUE( + absl::StrContains(output, + "# HELP tpu_raiden_transfer_failures_total Total count " + "of failed TPU Raiden data transfers.")); + + EXPECT_TRUE(absl::StrContains(output, "tpu_raiden_sent_bytes_total 100")); + EXPECT_TRUE(absl::StrContains(output, "tpu_raiden_received_bytes_total 200")); + EXPECT_TRUE(absl::StrContains( + output, "tpu_raiden_transfer_duration_seconds_sum 0.5")); + EXPECT_TRUE( + absl::StrContains(output, "tpu_raiden_stage_latency_seconds_sum 0.1")); + EXPECT_TRUE(absl::StrContains(output, "tpu_raiden_active_transfers 10")); + EXPECT_TRUE( + absl::StrContains(output, "tpu_raiden_buffer_occupancy_bytes 2048")); + EXPECT_TRUE( + absl::StrContains(output, "tpu_raiden_transfer_failures_total 1")); +} + +TEST(PrometheusExporterTest, UnmappedCustomMetricFallback) { + PrometheusMetricInfo custom_info = + GetPrometheusMetricInfo("custom_unmapped_metric"); + EXPECT_EQ(custom_info.prometheus_name, "custom_unmapped_metric"); + EXPECT_EQ(custom_info.help_text, "custom_unmapped_metric"); + + PrometheusExporter exporter; + exporter.IncrementCounter("custom_unmapped_counter", {}, 42); + exporter.SetGauge("custom_unmapped_gauge", {}, 99); + exporter.ObserveHistogram("custom_unmapped_histogram", {}, 1.23); + + std::string output = exporter.GetTextSnapshot(); + + EXPECT_TRUE(absl::StrContains( + output, "# HELP custom_unmapped_counter custom_unmapped_counter")); + EXPECT_TRUE( + absl::StrContains(output, "# TYPE custom_unmapped_counter counter")); + EXPECT_TRUE(absl::StrContains(output, "custom_unmapped_counter 42")); + + EXPECT_TRUE(absl::StrContains( + output, "# HELP custom_unmapped_gauge custom_unmapped_gauge")); + EXPECT_TRUE(absl::StrContains(output, "# TYPE custom_unmapped_gauge gauge")); + EXPECT_TRUE(absl::StrContains(output, "custom_unmapped_gauge 99")); + + EXPECT_TRUE(absl::StrContains( + output, "# HELP custom_unmapped_histogram custom_unmapped_histogram")); + EXPECT_TRUE( + absl::StrContains(output, "# TYPE custom_unmapped_histogram histogram")); + EXPECT_TRUE(absl::StrContains(output, "custom_unmapped_histogram_sum 1.23")); +} + +TEST(PrometheusExporterTest, ResetAndGlobalExporter) { + auto& global_exporter = GetGlobalPrometheusExporter(); + global_exporter.Reset(); + global_exporter.IncrementCounter(metric_names::kReceivedBytesTotal, {}, 3); + + std::string text = global_exporter.GetTextSnapshot(); + EXPECT_THAT(text, HasSubstr("tpu_raiden_received_bytes_total 3")); + + global_exporter.Reset(); + EXPECT_THAT(global_exporter.GetTextSnapshot(), + Not(HasSubstr("tpu_raiden_received_bytes_total 3"))); + EXPECT_THAT( + global_exporter.GetTextSnapshot(), + HasSubstr("# HELP tpu_raiden_received_bytes_total Total count of bytes " + "received over TPU Raiden interfaces.")); +} + +TEST(PrometheusExporterTest, ConstReferenceAccess) { + PrometheusExporter exporter; + const PrometheusExporter& const_exporter = exporter; + const MetricsBackend& const_backend = exporter; + + const_exporter.IncrementCounter(metric_names::kSentBytesTotal, {}, 500); + const_backend.SetGauge(metric_names::kActiveTransfers, {}, 2); + const_backend.ObserveHistogram(metric_names::kTransferDurationSeconds, {}, + 0.01); + + std::string snapshot = const_backend.GetTextSnapshot(); + EXPECT_TRUE(absl::StrContains(snapshot, "tpu_raiden_sent_bytes_total 500")); + EXPECT_TRUE(absl::StrContains(snapshot, "tpu_raiden_active_transfers 2")); + EXPECT_TRUE(absl::StrContains( + snapshot, "tpu_raiden_transfer_duration_seconds_sum 0.01")); +} + +TEST(PrometheusExporterTest, ConcurrentMetricUpdates) { + PrometheusExporter exporter; + constexpr int kNumThreads = 8; + constexpr int kIterations = 1000; + std::vector threads; + threads.reserve(kNumThreads); + + for (int i = 0; i < kNumThreads; ++i) { + threads.emplace_back([&exporter]() { + for (int j = 0; j < kIterations; ++j) { + exporter.IncrementCounter(metric_names::kSentBytesTotal, {}, 1); + exporter.SetGauge(metric_names::kActiveTransfers, {}, j); + exporter.ObserveHistogram(metric_names::kTransferDurationSeconds, {}, + 0.001); + } + }); + } + + for (auto& t : threads) { + t.join(); + } + + std::string snapshot = exporter.GetTextSnapshot(); + EXPECT_TRUE(absl::StrContains( + snapshot, + absl::StrCat("tpu_raiden_sent_bytes_total ", kNumThreads * kIterations))); +} + +} // namespace +} // namespace tpu_raiden::telemetry