Removing ref from Registry.Service Impl (#111)

zio · Jan 4, 2022 · 685d6ff · 685d6ff
1 parent c4b5bb6
commit 685d6ff
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 28 deletions.
diff --git a/prometheus/src/main/scala/zio/metrics/prometheus/package.scala b/prometheus/src/main/scala/zio/metrics/prometheus/package.scala
@@ -1,7 +1,7 @@
 package zio.metrics
 
 import zio.{ Has, ZLayer }
-import zio.{ Ref, Task, UIO }
+import zio.{ Task, UIO }
 
 package object prometheus {
 
@@ -26,43 +26,37 @@ package object prometheus {
     type Percentile = Double
     type Tolerance  = Double
 
-    val explicit: ZLayer[Has[Option[CollectorRegistry]], Nothing, Registry] =
-      ZLayer.fromFunction[Has[Option[CollectorRegistry]], Registry.Service](
-        optionalRegistry =>
+    val explicit: ZLayer[Has[CollectorRegistry], Nothing, Registry] =
+      ZLayer.fromService[CollectorRegistry, Registry.Service](
+        registry =>
           new Service {
-            private val registryRef: UIO[Ref[CollectorRegistry]] = {
-              val registry = optionalRegistry.get
-              Ref.make(registry.getOrElse(CollectorRegistry.defaultRegistry))
-            }
 
-            def getCurrent(): UIO[CollectorRegistry] = registryRef >>= (_.get)
+            def getCurrent(): UIO[CollectorRegistry] = UIO(registry)
 
             def registerCounter[A: Show](label: Label[A]): Task[PCounter] =
-              registryRef >>= (_.modify(r => {
+              Task({
                 val name = Show[A].show(label.name)
-                val c = PCounter
+                PCounter
                   .build()
                   .name(name)
                   .labelNames(label.labels: _*)
                   .help(label.help)
-                  .register(r)
-                (c, r)
-              }))
+                  .register(registry)
+              })
 
             def registerGauge[L: Show](label: Label[L]): Task[PGauge] =
-              registryRef >>= (_.modify(r => {
+              Task({
                 val name = Show[L].show(label.name)
-                val g = PGauge
+                PGauge
                   .build()
                   .name(name)
                   .labelNames(label.labels: _*)
                   .help(label.help)
-                  .register(r)
-                (g, r)
-              }))
+                  .register(registry)
+              })
 
             def registerHistogram[L: Show](label: Label[L], buckets: Buckets): Task[PHistogram] =
-              registryRef >>= (_.modify(r => {
+              Task({
                 val name = Show[L].show(label.name)
                 val hb = PHistogram
                   .build()
@@ -75,25 +69,25 @@ package object prometheus {
                   case LinearBuckets(s, w, c)      => hb.linearBuckets(s, w, c)
                   case ExponentialBuckets(s, f, c) => hb.exponentialBuckets(s, f, c)
                 }
-                (h.register(r), r)
-              }))
+                h.register(registry)
+              })
 
             def registerSummary[L: Show](label: Label[L], quantiles: List[(Percentile, Tolerance)]): Task[PSummary] =
-              registryRef >>= (_.modify(r => {
+              Task({
                 val name = Show[L].show(label.name)
                 val sb = PSummary
                   .build()
                   .name(name)
                   .labelNames(label.labels: _*)
                   .help(label.help)
 
-                val s = quantiles.foldLeft(sb)((acc, c) => acc.quantile(c._1, c._2)).register(r)
-                (s, r)
-              }))
+                quantiles.foldLeft(sb)((acc, c) => acc.quantile(c._1, c._2)).register(registry)
+              })
           }
       )
 
-    val live: ZLayer[Any, Nothing, Has[Registry.Service]] = ZLayer.succeed[Option[CollectorRegistry]](None) >>> explicit
+    val live
+      : ZLayer[Any, Nothing, Has[Registry.Service]] = ZLayer.succeed(CollectorRegistry.defaultRegistry) >>> explicit
 
   }
 }
diff --git a/prometheus/src/test/scala/zio/metrics/ExplicitRegistryLayer.scala b/prometheus/src/test/scala/zio/metrics/ExplicitRegistryLayer.scala
@@ -22,7 +22,7 @@ object ExplicitRegistryLayer {
     .register(myRegistry)
   preCounter.inc(9)
 
-  val myCustomLayer = ZLayer.succeed[Option[CollectorRegistry]](Some(myRegistry)) >>> Registry.explicit
+  val myCustomLayer = ZLayer.succeed(myRegistry) >>> Registry.explicit
 
   val rt = Runtime.unsafeFromLayer(MetricMap.live ++ Exporters.live ++ Console.live)