Introduce PrometheusLabelView in label_util to support zero-allocation hot-path Prometheus label formatting and lookup, previous FormatPrometheusLabels was using heap allocation even in cold path

tpu_sync/telemetry/BUILD

@@ -148,6 +148,7 @@ cc_library(
     deps = [
         ":metrics_backend",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",

tpu_sync/telemetry/label_util.cc

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/base/attributes.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
@@ -30,8 +31,6 @@
 namespace tpu_raiden::telemetry {
 namespace {
 
-constexpr size_t kDefaultPrometheusStackBufferSize = 256;
-
 // Lightweight buffer writer that bounds-checks appends into a char span.
 class BufferWriter {
  public:
@@ -92,12 +91,40 @@ class BufferWriter {
   size_t offset_ = 0;
 };
 
-// Common stack-allocated sorting helper for multi-label sets.
-void SortLabels(
-    LabelSpan labels,
-    absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity>& out) {
-  out.assign(labels.begin(), labels.end());
-  absl::c_sort(out);
+// Returns labels ordered by key. When `labels` is not already sorted, the
+// sorted copy is stored in `scratch` and the returned span aliases it;
+// otherwise the returned span aliases `labels`. Both arguments must outlive
+// the returned span.
+LabelSpan EnsureSortedLabels(
+    LabelSpan labels ABSL_ATTRIBUTE_LIFETIME_BOUND,
+    absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity>& scratch
+        ABSL_ATTRIBUTE_LIFETIME_BOUND) {
+  if (absl::c_is_sorted(labels)) return labels;
+  scratch.assign(labels.begin(), labels.end());
+  absl::c_sort(scratch);
+  return scratch;
+}
+
+// Formats pre-sorted labels into Prometheus canonical format directly into
+// output_buffer without checking sort order.
+std::optional<absl::string_view> FormatSortedPrometheusLabelsToBuffer(
+    LabelSpan sorted_labels, absl::Span<char> output_buffer) {
+  if (sorted_labels.empty()) {
+    return absl::string_view(output_buffer.data(), 0);
+  }
+
+  BufferWriter writer(output_buffer);
+  if (!writer.Append('{')) return std::nullopt;
+  for (size_t i = 0; i < sorted_labels.size(); ++i) {
+    if (i > 0 && !writer.Append(',')) return std::nullopt;
+    const MetricLabel& label = sorted_labels[i];
+    if (!writer.Append(label.key) || !writer.Append("=\"") ||
+        !writer.AppendPrometheusEscaped(label.value) || !writer.Append('"')) {
+      return std::nullopt;
+    }
+  }
+  if (!writer.Append('}')) return std::nullopt;
+  return writer.view();
 }
 
 // Returns the byte length of Prometheus label value after escaping ('\', '"',
@@ -123,30 +150,6 @@ size_t ComputePrometheusLabelsSize(LabelSpan labels) {
   return total;
 }
 
-// Generic 2-stage stack-to-heap allocation fallback orchestrator.
-template <size_t StackBufferSize, typename SizeFn, typename BufferFormatFn>
-std::string FormatWithStackBufferFallback(LabelSpan labels, SizeFn size_fn,
-                                          BufferFormatFn format_fn) {
-  if (labels.empty()) return "";
-
-  char stack_buffer[StackBufferSize];
-  if (std::optional<absl::string_view> formatted =
-          format_fn(labels, absl::MakeSpan(stack_buffer));
-      formatted.has_value()) {
-    return std::string(*formatted);
-  }
-
-  // Exact-size dynamic fallback for label sets exceeding stack buffer.
-  std::string result(size_fn(labels), '\0');
-  std::optional<absl::string_view> formatted =
-      format_fn(labels, absl::MakeSpan(result));
-  if (!formatted.has_value()) {
-    return "";
-  }
-  result.resize(formatted->size());
-  return result;
-}
-
 }  // namespace
 
 std::optional<absl::string_view> FormatShmLabelsToBuffer(
@@ -157,11 +160,8 @@ std::optional<absl::string_view> FormatShmLabelsToBuffer(
 
   BufferWriter writer(output_buffer);
 
-  absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity> sorted;
-  if (!absl::c_is_sorted(labels)) {
-    SortLabels(labels, sorted);
-    labels = sorted;
-  }
+  absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity> scratch;
+  labels = EnsureSortedLabels(labels, scratch);
 
   for (size_t i = 0; i < labels.size(); ++i) {
     if (i > 0 && !writer.Append(';')) return std::nullopt;
@@ -233,30 +233,42 @@ std::optional<absl::string_view> FormatPrometheusLabelsToBuffer(
     return absl::string_view(output_buffer.data(), 0);
   }
 
-  BufferWriter writer(output_buffer);
+  absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity> scratch;
+  labels = EnsureSortedLabels(labels, scratch);
+  return FormatSortedPrometheusLabelsToBuffer(labels, output_buffer);
+}
 
-  absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity> sorted;
-  if (!absl::c_is_sorted(labels)) {
-    SortLabels(labels, sorted);
-    labels = sorted;
-  }
+PrometheusLabelView::PrometheusLabelView(LabelSpan labels) {
+  if (labels.empty()) return;
 
-  if (!writer.Append('{')) return std::nullopt;
-  for (size_t i = 0; i < labels.size(); ++i) {
-    if (i > 0 && !writer.Append(',')) return std::nullopt;
-    const MetricLabel& label = labels[i];
-    if (!writer.Append(label.key) || !writer.Append("=\"") ||
-        !writer.AppendPrometheusEscaped(label.value) || !writer.Append('"')) {
-      return std::nullopt;
+  absl::InlinedVector<MetricLabel, kDefaultInlinedLabelCapacity> scratch;
+  labels = EnsureSortedLabels(labels, scratch);
+
+  std::optional<absl::string_view> formatted =
+      FormatSortedPrometheusLabelsToBuffer(labels, absl::MakeSpan(stack_buf_));
+  if (formatted.has_value()) {
+    view_ = *formatted;
+  } else {
+    heap_fallback_.resize(ComputePrometheusLabelsSize(labels));
+    std::optional<absl::string_view> heap_formatted =
+        FormatSortedPrometheusLabelsToBuffer(labels,
+                                             absl::MakeSpan(heap_fallback_));
+    if (heap_formatted.has_value()) {
+      heap_fallback_.resize(heap_formatted->size());
+      view_ = heap_fallback_;
+    } else {
+      heap_fallback_.clear();
     }
   }
-  if (!writer.Append('}')) return std::nullopt;
-  return writer.view();
 }
 
-std::string FormatPrometheusLabels(LabelSpan labels) {
-  return FormatWithStackBufferFallback<kDefaultPrometheusStackBufferSize>(
-      labels, ComputePrometheusLabelsSize, FormatPrometheusLabelsToBuffer);
+std::string PrometheusLabelView::ToOwned() {
+  const absl::string_view current_view = view_;
+  view_ = "";
+  if (!heap_fallback_.empty()) {
+    return std::exchange(heap_fallback_, std::string());
+  }
+  return std::string(current_view);
 }
 
 }  // namespace tpu_raiden::telemetry

tpu_sync/telemetry/label_util.h

@@ -22,6 +22,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/base/attributes.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "tpu_sync/telemetry/metrics_backend.h"
@@ -66,10 +67,44 @@ std::vector<std::pair<std::string, std::string>> ParseShmLabels(
 std::optional<absl::string_view> FormatPrometheusLabelsToBuffer(
     LabelSpan labels, absl::Span<char> output_buffer);
 
-// Owning std::string wrapper for non-critical paths, testing, and series
-// identification. Performs raw serialization without runtime key syntax
-// validation.
-std::string FormatPrometheusLabels(LabelSpan labels);
+// Default stack buffer capacity (in bytes) for formatted Prometheus labels,
+// sized to accommodate standard multi-label metric descriptors without heap
+// allocation.
+inline constexpr std::size_t kDefaultPrometheusStackBufferSize = 256;
+
+// Zero-allocation RAII stack buffer view for hot-path metric lookup.
+// Formats labels directly into an internal 256-byte stack buffer and provides
+// an absl::string_view for heterogeneous hash map lookup. Falls back to dynamic
+// heap allocation only if the formatted string exceeds 256 bytes or if sorting
+// requires more than 8 unsorted labels.
+//
+// view() is not guaranteed to be null-terminated (the stack path writes
+// exactly the formatted bytes with no trailing '\0', and
+// kDefaultPrometheusStackBufferSize reserves no byte for one). Never pass
+// view().data() to C-style APIs.
+class PrometheusLabelView {
+ public:
+  explicit PrometheusLabelView(LabelSpan labels);
+
+  // Non-copyable and non-movable: holds internal pointers to stack_buf_ /
+  // heap_fallback_.
+  PrometheusLabelView(const PrometheusLabelView&) = delete;
+  PrometheusLabelView& operator=(const PrometheusLabelView&) = delete;
+  PrometheusLabelView(PrometheusLabelView&&) = delete;
+  PrometheusLabelView& operator=(PrometheusLabelView&&) = delete;
+
+  absl::string_view view() const ABSL_ATTRIBUTE_LIFETIME_BOUND { return view_; }
+
+  // Returns an owned string, moving the heap fallback if allocated or
+  // constructing from the stack buffer. Invalidates view() and resets internal
+  // state.
+  std::string ToOwned();
+
+ private:
+  char stack_buf_[kDefaultPrometheusStackBufferSize];
+  std::string heap_fallback_;
+  absl::string_view view_ = {};
+};
 
 // ============================================================================
 // Zero-Allocation Label Resolution for Fixed Schemas (Fixed-Arity)