Code review improvements & gauges for in-progress requests

3 files changed, 160 insertions, 28 deletions

diff --git a/docs/backends/prometheus.rst b/docs/backends/prometheus.rst
index a86b6e6..8b89a23 100644
--- a/docs/backends/prometheus.rst
+++ b/docs/backends/prometheus.rst
@@ -13,6 +13,18 @@ The Prometheus backend wraps any other backend, for example::
 
   implicit val sttpBackend = PrometheusBackend(AkkaHttpBackend())
 
-It uses by default ``sttp_request_latency`` histogram name, defined in ``PrometheusBackend.DefaultHistogramName``. It is possible to define custom histograms name by passing function mapping request to histogram name::
+It gathers request execution times in ``Histogram``. It uses by default ``sttp_request_latency`` name, defined in ``PrometheusBackend.DefaultHistogramName``. It is possible to define custom histograms name by passing function mapping request to histogram name::
 
-  implicit val sttpBackend = PrometheusBackend(AkkaHttpBackend(), Some(request => request.uri.toString))
\ No newline at end of file
+  implicit val sttpBackend = PrometheusBackend(AkkaHttpBackend(), request => Some(request.uri.host))
+
+You can disable request histograms by passing ``None`` returning function::
+
+  implicit val sttpBackend = PrometheusBackend(AkkaHttpBackend(), _ => None)
+
+This backend also offers ``Gauge`` with currently in-progress requests number. It uses by default ``sttp_requests_in_progress`` name, defined in ``PrometheusBackend.DefaultRequestsInProgressGaugeName``. It is possible to define custom gauge name by passing function mapping request to gauge name::
+
+   implicit val sttpBackend = PrometheusBackend(AkkaHttpBackend(), requestToHistogramNameMapper = request => Some(request.uri.host))
+
+You can disable request in-progress gauges by passing ``None`` returning function::
+
+  implicit val sttpBackend = PrometheusBackend(AkkaHttpBackend(), requestToInProgressGaugeNameMapper = _ => None)
\ No newline at end of file
diff --git a/metrics/prometheus-backend/src/main/scala/com/softwaremill/sttp/prometheus/PrometheusBackend.scala b/metrics/prometheus-backend/src/main/scala/com/softwaremill/sttp/prometheus/PrometheusBackend.scala
index d1712db..173d9c6 100644
--- a/metrics/prometheus-backend/src/main/scala/com/softwaremill/sttp/prometheus/PrometheusBackend.scala
+++ b/metrics/prometheus-backend/src/main/scala/com/softwaremill/sttp/prometheus/PrometheusBackend.scala
@@ -3,33 +3,42 @@ package com.softwaremill.sttp.prometheus
 import java.util.concurrent.ConcurrentHashMap
 
 import com.softwaremill.sttp.{FollowRedirectsBackend, MonadError, Request, Response, SttpBackend}
-import io.prometheus.client.Histogram
+import io.prometheus.client.{Gauge, Histogram}
 
 import scala.collection.mutable
 import scala.language.higherKinds
 import scala.collection.JavaConverters._
 
 class PrometheusBackend[R[_], S] private (delegate: SttpBackend[R, S],
-                                          requestToHistogramNameMapper: Option[Request[_, S] => String])
+                                          requestToHistogramNameMapper: Request[_, S] => Option[String],
+                                          requestToInProgressGaugeNameMapper: Request[_, S] => Option[String])
     extends SttpBackend[R, S] {
 
-  import PrometheusBackend._
-
   private[this] val histograms: mutable.Map[String, Histogram] = new ConcurrentHashMap[String, Histogram]().asScala
+  private[this] val gauges: mutable.Map[String, Gauge] = new ConcurrentHashMap[String, Gauge]().asScala
 
   override def send[T](request: Request[T, S]): R[Response[T]] = {
-    val histogramName = getHistogramName(request)
-    val histogram = histograms.getOrElseUpdate(histogramName, createNewHistogram(histogramName))
-    val requestTimer = histogram.startTimer()
+    val requestTimer: Option[Histogram.Timer] = for {
+      histogramName: String <- requestToHistogramNameMapper(request)
+      histogram: Histogram = histograms.getOrElseUpdate(histogramName, createNewHistogram(histogramName))
+    } yield histogram.startTimer()
+
+    val gauge: Option[Gauge] = for {
+      gaugeName: String <- requestToInProgressGaugeNameMapper(request)
+    } yield gauges.getOrElseUpdate(gaugeName, createNewGauge(gaugeName))
+
+    gauge.foreach(_.inc())
 
     responseMonad.handleError(
       responseMonad.map(delegate.send(request)) { response =>
-        requestTimer.observeDuration()
+        requestTimer.foreach(_.observeDuration())
+        gauge.foreach(_.dec())
         response
       }
     ) {
       case e: Exception =>
-        requestTimer.observeDuration()
+        requestTimer.foreach(_.observeDuration())
+        gauge.foreach(_.dec())
         responseMonad.error(e)
     }
   }
@@ -38,19 +47,23 @@ class PrometheusBackend[R[_], S] private (delegate: SttpBackend[R, S],
 
   override def responseMonad: MonadError[R] = delegate.responseMonad
 
-  private[this] def getHistogramName(request: Request[_, S]): String =
-    requestToHistogramNameMapper.map(_.apply(request)).getOrElse(DefaultHistogramName)
-
   private[this] def createNewHistogram(name: String): Histogram = Histogram.build().name(name).help(name).register()
+
+  private[this] def createNewGauge(name: String): Gauge = Gauge.build().name(name).help(name).register()
 }
 
 object PrometheusBackend {
 
   val DefaultHistogramName = "sttp_request_latency"
+  val DefaultRequestsInProgressGaugeName = "sttp_requests_in_progress"
 
   def apply[R[_], S](delegate: SttpBackend[R, S],
-                     requestToHistogramNameMapper: Option[Request[_, S] => String] = None): SttpBackend[R, S] = {
-    // redirects should be handled before brave tracing, hence adding the follow-redirects backend on top
-    new FollowRedirectsBackend(new PrometheusBackend(delegate, requestToHistogramNameMapper))
+                     requestToHistogramNameMapper: Request[_, S] => Option[String] = (_: Request[_, S]) =>
+                       Some(DefaultHistogramName),
+                     requestToInProgressGaugeNameMapper: Request[_, S] => Option[String] = (_: Request[_, S]) =>
+                       Some(DefaultRequestsInProgressGaugeName)): SttpBackend[R, S] = {
+    // redirects should be handled before prometheus
+    new FollowRedirectsBackend(
+      new PrometheusBackend(delegate, requestToHistogramNameMapper, requestToInProgressGaugeNameMapper))
   }
 }
diff --git a/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala b/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala
index e6476de..33477a0 100644
--- a/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala
+++ b/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala
@@ -1,12 +1,22 @@
 package com.softwaremill.sttp.prometheus
 
-import com.softwaremill.sttp.{HttpURLConnectionBackend, Id, sttp}
-import io.prometheus.client.CollectorRegistry
-import org.scalatest.{BeforeAndAfter, FlatSpec, Matchers}
-import com.softwaremill.sttp._
+import java.lang
+import java.util.concurrent.CountDownLatch
+
 import com.softwaremill.sttp.testing.SttpBackendStub
+import com.softwaremill.sttp.{HttpURLConnectionBackend, Id, sttp, _}
+import io.prometheus.client.CollectorRegistry
+import org.scalatest.concurrent.{Eventually, IntegrationPatience}
+import org.scalatest.{BeforeAndAfter, FlatSpec, Matchers, OptionValues}
 
-class PrometheusBackendTest extends FlatSpec with Matchers with BeforeAndAfter {
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.Future
+
+class PrometheusBackendTest extends FlatSpec with Matchers with BeforeAndAfter with Eventually with OptionValues {
+
+  before {
+    CollectorRegistry.defaultRegistry.clear()
+  }
 
   it should "use default histogram name" in {
     // given
@@ -18,22 +28,119 @@ class PrometheusBackendTest extends FlatSpec with Matchers with BeforeAndAfter {
     (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
 
     // then
-    val result = CollectorRegistry.defaultRegistry.getSampleValue(s"${PrometheusBackend.DefaultHistogramName}_count")
-    result shouldBe requestsNumber
+    getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count").value shouldBe requestsNumber
   }
 
   it should "use mapped request to histogram name" in {
     // given
-    val customHistogramName = "my-custom-histogram"
+    val customHistogramName = "my_custom_histogram"
     val backend =
-      PrometheusBackend[Id, Nothing](SttpBackendStub(HttpURLConnectionBackend()), Some(_ => customHistogramName))
+      PrometheusBackend[Id, Nothing](SttpBackendStub(HttpURLConnectionBackend()), _ => Some(customHistogramName))
     val requestsNumber = 5
 
     // when
     (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
 
     // then
-    CollectorRegistry.defaultRegistry.getSampleValue(s"${PrometheusBackend.DefaultHistogramName}_count") shouldBe null
-    CollectorRegistry.defaultRegistry.getSampleValue(s"${customHistogramName}_count") shouldBe requestsNumber
+    getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count") shouldBe empty
+    getMetricVale(s"${customHistogramName}_count").value shouldBe requestsNumber
+  }
+
+  it should "disable histograms" in {
+    // given
+    val backend =
+      PrometheusBackend[Id, Nothing](SttpBackendStub(HttpURLConnectionBackend()), _ => None)
+    val requestsNumber = 6
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count") shouldBe empty
+  }
+
+  it should "use default gauge name" in {
+    // given
+    val requestsNumber = 10
+    val countDownLatch = new CountDownLatch(1)
+    val backendStub = SttpBackendStub.asynchronousFuture.whenAnyRequest.thenRespondWrapped {
+      Future {
+        countDownLatch.await()
+        Response(Right(""), 200, "", Nil, Nil)
+      }
+    }
+    val backend = PrometheusBackend[Future, Nothing](backendStub)
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName).value shouldBe requestsNumber
+    }
+
+    countDownLatch.countDown()
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName).value shouldBe 0
+    }
+  }
+
+  it should "use mapped request to gauge name" in {
+    // given
+    val customGaugeName = "my_custom_gauge"
+    val requestsNumber = 10
+    val countDownLatch = new CountDownLatch(1)
+    val backendStub = SttpBackendStub.asynchronousFuture.whenAnyRequest.thenRespondWrapped {
+      Future {
+        countDownLatch.await()
+        Response(Right(""), 200, "", Nil, Nil)
+      }
+    }
+    val backend =
+      PrometheusBackend[Future, Nothing](backendStub, requestToInProgressGaugeNameMapper = _ => Some(customGaugeName))
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+      getMetricVale(customGaugeName).value shouldBe requestsNumber
+    }
+
+    countDownLatch.countDown()
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+      getMetricVale(customGaugeName).value shouldBe 0
+    }
   }
+
+  it should "disable gauge" in {
+    // given
+    val requestsNumber = 10
+    val countDownLatch = new CountDownLatch(1)
+    val backendStub = SttpBackendStub.asynchronousFuture.whenAnyRequest.thenRespondWrapped {
+      Future {
+        countDownLatch.await()
+        Response(Right(""), 200, "", Nil, Nil)
+      }
+    }
+    val backend = PrometheusBackend[Future, Nothing](backendStub, requestToInProgressGaugeNameMapper = _ => None)
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+
+    countDownLatch.countDown()
+    eventually {
+      getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count").value shouldBe requestsNumber
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+    }
+  }
+
+  private[this] def getMetricVale(name: String): Option[lang.Double] =
+    Option(CollectorRegistry.defaultRegistry.getSampleValue(name))
+
 }