Improve memory consumption of prometheus handler by about 60% (#33)

zio · Feb 26, 2023 · 951a3d9 · 951a3d9
1 parent ed1ae96
commit 951a3d9
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 58 deletions.
diff --git a/core/src/main/scala/zio/metrics/connectors/prometheus/PrometheusEncoder.scala b/core/src/main/scala/zio/metrics/connectors/prometheus/PrometheusEncoder.scala
@@ -17,51 +17,75 @@ case object PrometheusEncoder {
     timestamp: Instant,
     descriptionKey: Option[String],
   ): Chunk[String] = {
+    val name = key.name.replaceAll("-", "_").trim
+
+    // The header required for all Prometheus metrics
+    val prometheusType = state match {
+      case _: MetricState.Counter   => "counter"
+      case _: MetricState.Gauge     => "gauge"
+      case _: MetricState.Histogram => "histogram"
+      case _: MetricState.Summary   => "summary"
+      case _: MetricState.Frequency => "counter"
+    }
+
+    val encodeHead = {
+      val description = descriptionKey.flatMap(d => key.tags.find(_.key == d)).fold("")(l => s" ${l.value}")
+      Chunk(
+        s"# TYPE $name $prometheusType",
+        s"# HELP $name$description",
+      )
+    }
+
+    val encodeTimestamp = s"${timestamp.toEpochMilli}"
+
+    def encodeLabels(allLabels: Set[MetricLabel]) =
+      if (allLabels.isEmpty) new StringBuilder("")
+      else
+        allLabels
+          .foldLeft(new StringBuilder(256).append("{")) { case (sb, l) =>
+            sb.append(l.key).append("=\"").append(l.value).append("\",")
+          }
+          .append("}")
+
+    val tagsWithoutDescription = descriptionKey.fold(key.tags)(d => key.tags.filter(_.key != d))
+
+    val baseLabels = encodeLabels(tagsWithoutDescription)
+
+    def encodeExtraLabels(extraLabels: Set[MetricLabel]) =
+      if (extraLabels.isEmpty) baseLabels else encodeLabels(tagsWithoutDescription ++ extraLabels)
 
     def encodeCounter(c: MetricState.Counter, extraLabels: MetricLabel*): String =
-      s"${encodeName(key.name)}${encodeLabels(extraLabels.toSet)} ${c.count} $encodeTimestamp"
+      s"$name${encodeExtraLabels(extraLabels.toSet)} ${c.count} $encodeTimestamp"
 
     def encodeGauge(g: MetricState.Gauge): String =
-      s"${encodeName(key.name)}${encodeLabels()} ${g.value} $encodeTimestamp"
+      s"$name$baseLabels ${g.value} $encodeTimestamp"
 
     def encodeHistogram(h: MetricState.Histogram): Chunk[String] =
       encodeSamples(sampleHistogram(h), suffix = "_bucket")
 
     def encodeSummary(s: MetricState.Summary): Chunk[String] =
       encodeSamples(sampleSummary(s), suffix = "")
 
-    // The header required for all Prometheus metrics
-    def encodeHead: Chunk[String] = {
-      val description = descriptionKey.flatMap(d => key.tags.find(_.key == d)).fold("")(l => s" ${l.value}")
-      Chunk(
-        s"# TYPE ${encodeName(key.name)} $prometheusType",
-        s"# HELP ${encodeName(key.name)}$description",
-      )
-    }
-
-    def encodeName(s: String): String =
-      s.replaceAll("-", "_")
-
-    def encodeLabels(extraLabels: Set[MetricLabel] = Set.empty): String = {
-      val allLabels = descriptionKey.fold(key.tags)(d => key.tags.filter(_.key != d)) ++ extraLabels
-
-      if (allLabels.isEmpty) ""
-      else allLabels.map(l => l.key + "=\"" + l.value + "\"").mkString("{", ",", "}")
-    }
-
     def encodeSamples(samples: SampleResult, suffix: String): Chunk[String] =
-      samples.buckets.map { b =>
-        s"${encodeName(key.name)}$suffix${encodeLabels(b._1)} ${b._2.map(_.toString).getOrElse("NaN")} $encodeTimestamp"
-          .trim()
-      } ++ Chunk(
-        s"${encodeName(key.name)}_sum${encodeLabels()} ${samples.sum} $encodeTimestamp".trim(),
-        s"${encodeName(key.name)}_count${encodeLabels()} ${samples.count} $encodeTimestamp".trim(),
-        s"${encodeName(key.name)}_min${encodeLabels()} ${samples.min} $encodeTimestamp".trim(),
-        s"${encodeName(key.name)}_max${encodeLabels()} ${samples.max} $encodeTimestamp".trim(),
+      Chunk(
+        samples.buckets
+          .foldLeft(new StringBuilder(samples.buckets.size * 100)) { case (sb, (l, v)) =>
+            sb.append(name)
+              .append(suffix)
+              .append(encodeExtraLabels(l))
+              .append(" ")
+              .append(v.map(_.toString).getOrElse("NaN"))
+              .append(" ")
+              .append(encodeTimestamp)
+              .append("\n")
+          }
+          .toString,
+        s"${name}_sum$baseLabels ${samples.sum} $encodeTimestamp",
+        s"${name}_count$baseLabels ${samples.count} $encodeTimestamp",
+        s"${name}_min$baseLabels ${samples.min} $encodeTimestamp",
+        s"${name}_max$baseLabels ${samples.max} $encodeTimestamp",
       )
 
-    def encodeTimestamp = s"${timestamp.toEpochMilli()}"
-
     def sampleHistogram(h: MetricState.Histogram): SampleResult =
       SampleResult(
         count = h.count.doubleValue(),
@@ -90,14 +114,6 @@ case object PrometheusEncoder {
         ),
       )
 
-    def prometheusType: String = state match {
-      case _: MetricState.Counter   => "counter"
-      case _: MetricState.Gauge     => "gauge"
-      case _: MetricState.Histogram => "histogram"
-      case _: MetricState.Summary   => "summary"
-      case _: MetricState.Frequency => "counter"
-    }
-
     def encodeDetails: Chunk[String] = state match {
       case c: MetricState.Counter   => Chunk(encodeCounter(c))
       case g: MetricState.Gauge     => Chunk(encodeGauge(g))

diff --git a/core/src/main/scala/zio/metrics/connectors/prometheus/package.scala b/core/src/main/scala/zio/metrics/connectors/prometheus/package.scala
@@ -19,16 +19,13 @@ package object prometheus {
   private def prometheusHandler(clt: PrometheusPublisher, config: MetricsConfig): Iterable[MetricEvent] => UIO[Unit] =
     events =>
       for {
-        reportComplete <- ZIO.foreach(events)(evt =>
-                            for {
-                              reportEvent <-
-                                PrometheusEncoder
-                                  .encode(evt, descriptionKey = Some(config.descriptionKey))
-                                  .map(_.mkString("\n"))
-                                  .catchAll(_ => ZIO.succeed(""))
-                            } yield reportEvent,
-                          )
-        _              <- clt.set(reportComplete.mkString("\n"))
+        old            <- clt.get
+        reportComplete <- ZIO.foreach(Chunk.fromIterable(events)) { e =>
+                            PrometheusEncoder.encode(e, descriptionKey = Some(config.descriptionKey)).catchAll { t =>
+                              ZIO.succeed(Chunk.empty)
+                            }
+                          }
+        _              <- clt.set(reportComplete.flatten.addString(new StringBuilder(old.length), "\n").toString())
       } yield ()
 
 }
diff --git a/core/src/test/scala/zio/metrics/connectors/prometheus/PrometheusEncoderSpec.scala b/core/src/test/scala/zio/metrics/connectors/prometheus/PrometheusEncoderSpec.scala
@@ -23,7 +23,7 @@ object PrometheusEncoderSpec extends ZIOSpecDefault with Generators {
   private def labelString(key: MetricKey.Untyped, extra: (String, String)*) = {
     val tags = key.tags.filter(_.key != descriptionKey) ++ extra.map(x => MetricLabel(x._1, x._2)).toSet
     if (tags.isEmpty) ""
-    else tags.map(l => s"""${l.key}="${l.value}"""").mkString("{", ",", "}")
+    else tags.toList.map(l => s"""${l.key}="${l.value}",""").mkString("{", "", "}")
   }
 
   private val encodeCounter = test("Encode a Counter")(check(genCounter) { case (pair, state) =>
@@ -87,10 +87,10 @@ object PrometheusEncoderSpec extends ZIOSpecDefault with Generators {
       text == Chunk(
         s"# TYPE $name summary",
         s"# HELP $name$help",
-      ) ++ state.quantiles.map { case (k, v) =>
+      ) ++ Chunk(state.quantiles.map { case (k, v) =>
         val labelsWithExtra = labelString(pair.metricKey, "quantile" -> k.toString, "error" -> state.error.toString)
-        s"""$name$labelsWithExtra ${v.getOrElse(Double.NaN)} $epochMilli"""
-      } ++ Chunk(
+        s"""$name$labelsWithExtra ${v.getOrElse(Double.NaN)} $epochMilli\n"""
+      }.mkString) ++ Chunk(
         s"${name}_sum$labels ${state.sum} $epochMilli",
         s"${name}_count$labels ${state.count.toDouble} $epochMilli",
         s"${name}_min$labels ${state.min} $epochMilli",
@@ -112,11 +112,15 @@ object PrometheusEncoderSpec extends ZIOSpecDefault with Generators {
       text == Chunk(
         s"# TYPE $name histogram",
         s"# HELP $name$help",
-      ) ++ state.buckets.filter(_._1 < Double.MaxValue).map { case (k, v) =>
-        val labelsWithExtra = labelString(pair.metricKey, "le" -> k.toString)
-        s"""${name}_bucket$labelsWithExtra ${v.toDouble} $epochMilli"""
-      } ++ Chunk(
-        s"""${name}_bucket$labelsWithInf ${state.count.toDouble} $epochMilli""",
+      ) ++ Chunk(
+        (state.buckets
+          .filter(_._1 < Double.MaxValue)
+          .map { case (k, v) =>
+            val labelsWithExtra = labelString(pair.metricKey, "le" -> k.toString)
+            s"""${name}_bucket$labelsWithExtra ${v.toDouble} $epochMilli\n"""
+          }
+          ++ s"""${name}_bucket$labelsWithInf ${state.count.toDouble} $epochMilli\n""").mkString,
+      ) ++ Chunk(
         s"${name}_sum$labels ${state.sum} $epochMilli",
         s"${name}_count$labels ${state.count.toDouble} $epochMilli",
         s"${name}_min$labels ${state.min} $epochMilli",