From f4af9cdefa7d1fb5e50f1b67af4a284ee29238d0 Mon Sep 17 00:00:00 2001
From: Tyler Levine <tyler.levine@falconx.io>
Date: Tue, 30 Jul 2024 12:47:35 -0700
Subject: [PATCH] chore: lazily add a comma if needed before writing family
 labels

---
 examples/custom-metric.rs |  4 ++--
 src/encoding.rs           |  3 ++-
 src/encoding/text.rs      | 35 ++++++++++++++++++++++++++++-------
 3 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/examples/custom-metric.rs b/examples/custom-metric.rs
index 9ad17a5a..a61ff8a7 100644
--- a/examples/custom-metric.rs
+++ b/examples/custom-metric.rs
@@ -1,4 +1,4 @@
-use prometheus_client::encoding::{text::encode, EncodeMetric, MetricEncoder};
+use prometheus_client::encoding::{text::encode, EncodeMetric, MetricEncoder, NoLabelSet};
 use prometheus_client::metrics::MetricType;
 use prometheus_client::registry::Registry;
 
@@ -20,7 +20,7 @@ impl EncodeMetric for MyCustomMetric {
         // E.g. every CPU cycle spend in this method delays the response send to
         // the Prometheus server.
 
-        encoder.encode_counter::<(), _, u64>(&rand::random::<u64>(), None)
+        encoder.encode_counter::<NoLabelSet, _, u64>(&rand::random::<u64>(), None)
     }
 
     fn metric_type(&self) -> prometheus_client::metrics::MetricType {
diff --git a/src/encoding.rs b/src/encoding.rs
index 8f8c7dea..c644f82b 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -248,7 +248,8 @@ pub trait EncodeLabel {
 pub struct LabelEncoder<'a>(LabelEncoderInner<'a>);
 
 /// Uninhabited type to represent the lack of a label set for a metric
-pub(crate) enum NoLabelSet {}
+#[derive(Debug)]
+pub enum NoLabelSet {}
 
 #[derive(Debug)]
 enum LabelEncoderInner<'a> {
diff --git a/src/encoding/text.rs b/src/encoding/text.rs
index 5d8b991e..8af67a5f 100644
--- a/src/encoding/text.rs
+++ b/src/encoding/text.rs
@@ -512,18 +512,39 @@ impl<'a> MetricEncoder<'a> {
             additional_labels.encode(LabelSetEncoder::new(self.writer).into())?;
         }
 
-        if let Some(labels) = &self.family_labels {
-            let mut string_writer = String::new();
-            labels.encode(LabelSetEncoder::new(&mut string_writer).into())?;
+        /// Writer impl which prepends a comma on the first call to write output to the wrapped writer
+        struct CommaPrependingWriter<'a> {
+            writer: &'a mut dyn Write,
+            should_prepend: bool,
+        }
 
-            if !string_writer.is_empty() {
-                if !self.const_labels.is_empty() || additional_labels.is_some() {
-                    self.writer.write_str(",")?;
+        impl Write for CommaPrependingWriter<'_> {
+            fn write_str(&mut self, s: &str) -> std::fmt::Result {
+                if self.should_prepend {
+                    self.writer.write_char(',')?;
+                    self.should_prepend = false;
                 }
-                self.writer.write_str(string_writer.as_str())?;
+                self.writer.write_str(s)
             }
         }
 
+        if let Some(labels) = self.family_labels {
+            // if const labels or additional labels have been written, a comma must be prepended before writing the family labels.
+            // However, it could be the case that the family labels are `Some` and yet empty, so the comma should _only_
+            // be prepended if one of the `Write` methods are actually called when attempting to write the family labels.
+            // Therefore, wrap the writer on `Self` with a CommaPrependingWriter if other labels have been written and
+            // there may be a need to prepend an extra comma before writing additional labels.
+            if !self.const_labels.is_empty() || additional_labels.is_some() {
+                let mut writer = CommaPrependingWriter {
+                    writer: self.writer,
+                    should_prepend: true,
+                };
+                labels.encode(LabelSetEncoder::new(&mut writer).into())?;
+            } else {
+                labels.encode(LabelSetEncoder::new(self.writer).into())?;
+            };
+        }
+
         self.writer.write_str("}")?;
 
         Ok(())