prometheus type-safe

package.scala

package ru.scala

import io.prometheus.client.{Collector, Counter, Gauge, SimpleCollector}
import shapeless.{HNil, LabelledGeneric}
import shapeless._
import shapeless.labelled.FieldType

import scala.annotation.implicitNotFound

package object prometheus {

  /** These prometheus wrappers enable type-safe metric labeling.
    * For a given `YourMetric` which extends `LabelledMetric[YourLabels]`
    * you can only access a collector by providing an instance of `YourLabels`.
    * Given that `YourLabels` is a case class (recommended), prometheus label names and values will be provided automatically
    * Label name will be same as case class field name, label value will be fieldValue.toString, but can be
    * customized by providing `LabelValue` type-class instance for `fieldValue` type
    * When `YourLabels` is not a case class, `AsLabels[YourLabels]` implicit instance is required.
    * Example usage:
    * {{{
    *
    *   sealed trait HttpMethod
    *   case object Get    extends HttpMethod
    *   case object Post   extends HttpMethod
    *   case object Delete extends HttpMethod
    *
    *   // we can customize desired repr
    *   implicit val methodLabelValue: LabelValue[HttpMethod] = _.toString.toUpperCase
    *   case class RequestLabels(method: HttpMethod, success: Boolean) // strict types: HttpMethod, Boolean
    *
    *   object RequestMetric extends LabelledMetric[RequestLabels](
    *     "requests_total",
    *     "Requests"
    *   ) with PrometheusCounter
    *
    *   // client code:
    *   val typeSafeLabels: RequestLabels(Post, true)
    *   RequestMetric.getCollector(typeSafeLabels).inc()
    * }}}
    * What happened internally:
    *   A private counter is created for `RequestMetric` singleton:
    *   {{{ Counter.build("requests_total", "Requests").labelNames("method", "success").register() }}}
    *   When `.getCollector(typeSafeLabels).inc()` is called, it translates to:
    *   {{{ counter.labels("POST", "true").inc() }}}
    *
    * Errors this type-safety prevents:
    *   - `counter.labels` fails at runtime when length of label values != length of label names
    *   - an arbitrary string can be provided as label value (e.g. `requests_total{method="INSERT", success="YES"})`
    *
    * Accordingly, instances of `NoLabelsMetric` will have no-argument `getCollector`, meaning we can't
    * provide any non-declared labels
    *
    * Also, a similar `PrometheusGauge` trait can be used to get a properly-labeled Gauge, via `getGauge` method
    *
    * Inspired by http://limansky.me/posts/2017-02-02-generating-sql-queries-with-shapeless.html
    */
  trait LabelValue[-V] {
    def get(v: V): String
  }
  object LabelValue {
    implicit val default: LabelValue[Any] = _.toString
  }

  trait LowPriorityAsLabels {
    implicit def headLabel[K <: Symbol, H, T <: HList](
      implicit witness: Witness.Aux[K],
      stringValue: LabelValue[H],
      tailAsLabels: AsLabels[T]
    ): AsLabels[FieldType[K, H] :: T] = new AsLabels[FieldType[K, H] :: T] {
      override val names: List[String] = witness.value.name :: tailAsLabels.names
      override def values(a: FieldType[K, H] :: T): List[String] =
        stringValue.get(a.head) :: tailAsLabels.values(a.tail)
    }
  }

  @implicitNotFound(
    "Cannot use ${A} as labels. It is recommended to be a case class, otherwise implicit value of AsLabels[${A}] type should be provided"
  )
  trait AsLabels[A] {
    val names: List[String]
    def values(a: A): List[String]
    def render(a: A): String = names.zip(values(a)).map(kv => s"""${kv._1}="${kv._2}"""").mkString("{", ", ", "}")
  }
  object AsLabels extends LowPriorityAsLabels {
    implicit def product[A, R](
      implicit gen: LabelledGeneric.Aux[A, R],
      asLabel: Lazy[AsLabels[R]]
    ): AsLabels[A] = new AsLabels[A] {
      override val names: List[String]        = asLabel.value.names
      override def values(a: A): List[String] = asLabel.value.values(gen.to(a))
    }
    def apply[T](implicit ev: AsLabels[T]): AsLabels[T] = ev

    implicit val hnilLabels: AsLabels[HNil] = new AsLabels[HNil] {
      override val names: List[String]           = Nil
      override def values(a: HNil): List[String] = Nil
    }
    implicit def hconsLabels[K, H, T <: HList](
      implicit
      hLister: Lazy[AsLabels[H]],
      tLister: AsLabels[T]
    ): AsLabels[FieldType[K, H] :: T] = new AsLabels[FieldType[K, H] :: T] {
      override val names: List[String] = hLister.value.names ++ tLister.names
      override def values(
        a: FieldType[K, H] :: T
      ): List[String] = hLister.value.values(a.head) ++ tLister.values(a.tail)
    }
  }

  sealed trait Metric {
    val labelNames: List[String]
    val name: String
    val help: String
  }

  /* Subclasses has to be singleton objects */
  abstract class LabelledMetric[L](val name: String, val help: String)(implicit val al: AsLabels[L]) extends Metric {
    val labelNames: List[String]             = al.names
    def labelValues(labels: L): List[String] = al.values(labels)
  }

  case class NoLabelsMetric(name: String, help: String) extends Metric {
    val labelNames: List[String] = List.empty
  }

  trait HasCollector[C <: Collector] {
    protected val collector: C
  }

  sealed trait PrometheusSimpleCollector[Child, C <: SimpleCollector[Child], B <: SimpleCollector.Builder[B, C]]
      extends HasCollector[C] {
    self: Metric =>
    val builder: B
    override protected lazy val collector: C = PrometheusSimpleCollector.create[C, B](
      builder,
      self
    )
    def getCollector[L](l: L)(implicit ev: self.type <:< LabelledMetric[L]): Child =
      collector.labels(this.labelValues(l): _*)
    def getCollector(implicit ev: self.type <:< NoLabelsMetric): Child = collector.labels()
  }
  object PrometheusSimpleCollector {
    def create[C <: SimpleCollector[_], B <: SimpleCollector.Builder[B, C]](builder: B, m: Metric): C =
      builder
        .name(
          m.name
        )
        .help(
          m.help
        )
        .labelNames(
          m.labelNames: _*
        )
        .create()
  }

  trait PrometheusCounter extends PrometheusSimpleCollector[Counter.Child, Counter, Counter.Builder] {
    self: Metric =>
    override val builder: Counter.Builder = Counter.build()
  }

  trait PrometheusGauge extends PrometheusSimpleCollector[Gauge.Child, Gauge, Gauge.Builder] {
    self: Metric =>
    override val builder: Gauge.Builder = Gauge.build()
  }

  trait AutoRegistration {
    self: Metric with HasCollector[_ <: Collector] with Singleton =>
    collector.register()
  }
}