1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.metrics.prometheus
10 
11 import com.codahale.metrics._
12 import com.typesafe.config.{Config, ConfigFactory}
13 import io.prometheus.client.Collector.MetricFamilySamples
14 import io.prometheus.client.CollectorRegistry
15 import io.prometheus.client.dropwizard.DropwizardExports
16 import io.prometheus.client.dropwizard.samplebuilder.{DefaultSampleBuilder, SampleBuilder}
17 import io.prometheus.client.exporter.{HTTPServer, PushGateway}
18 import org.locationtech.geomesa.metrics.core.ReporterFactory
19 import org.locationtech.geomesa.utils.io.CloseWithLogging
20 import pureconfig.{ConfigReader, ConfigSource}
21 
22 import java.net.URL
23 import java.util.Locale
24 import java.util.concurrent.TimeUnit
25 
26 class PrometheusReporterFactory extends ReporterFactory {
27 
28   import PrometheusReporterFactory._
29 
30   override def apply(
31       conf: Config,
32       registry: MetricRegistry,
33       rates: TimeUnit,
34       durations: TimeUnit): Option[ScheduledReporter] = {
35     val typ = if (conf.hasPath("type")) { Option(conf.getString("type")).map(_.toLowerCase(Locale.US)) } else { None }
36     typ match {
37       case Some("prometheus") =>
38         val config = ConfigSource.fromConfig(conf.withFallback(Defaults)).loadOrThrow[PrometheusConfig]
39         Some(new PrometheusReporter(registry, config.suffix, config.port))
40 
41       case Some("prometheus-pushgateway") =>
42         val config = ConfigSource.fromConfig(conf.withFallback(Defaults)).loadOrThrow[PrometheusPushgatewayConfig]
43         Some(new PushgatewayReporter(registry, config.suffix, config.gateway, config.jobName))
44 
45       case _ =>
46         None
47     }
48   }
49 }
50 
51 object PrometheusReporterFactory {
52 
53   import pureconfig.generic.semiauto._
54 
55   private val Defaults = ConfigFactory.parseResourcesAnySyntax("prometheus-defaults").resolve()
56 
57   case class PrometheusConfig(port: Int, suffix: Option[String])
58   case class PrometheusPushgatewayConfig(gateway: URL, jobName: String, suffix: Option[String])
59 
60   implicit val PrometheusConfigReader: ConfigReader[PrometheusConfig] = deriveReader[PrometheusConfig]
61   implicit val PrometheusPushgatewayConfigReader: ConfigReader[PrometheusPushgatewayConfig] =
62     deriveReader[PrometheusPushgatewayConfig]
63 
64   /**
65    * Prometheus reporter
66    *
67    * @param registry registry
68    * @param suffix metrics suffix
69    * @param port http server port, or zero to use an ephemeral open port
70    */
71   class PrometheusReporter(registry: MetricRegistry, suffix: Option[String], port: Int)
72       extends BasePrometheusReporter(registry, CollectorRegistry.defaultRegistry, suffix) {
73 
74     private val server = new HTTPServer.Builder().withPort(port).build()
75 
76     def getPort: Int = server.getPort
77 
78     override def close(): Unit = {
79       CloseWithLogging(server)
80       super.close()
81     }
82   }
83 
84   /**
85    * Pushgateway reporter
86    *
87    * Note: we use a new collector registry, as generally with pushgateway you don't want to expose things
88    * from the default registry like memory, etc
89    *
90    * @param registry registry
91    * @param suffix metrics suffix
92    * @param gateway pushgateway url
93    * @param jobName pushgateway job name
94    */
95   class PushgatewayReporter(registry: MetricRegistry, suffix: Option[String], gateway: URL, jobName: String)
96       extends BasePrometheusReporter(registry, new CollectorRegistry(), suffix) {
97     override def close(): Unit = {
98       try { new PushGateway(gateway).pushAdd(collectorRegistry, jobName) } finally {
99         super.close()
100       }
101     }
102   }
103 
104   /**
105    * Placeholder reporter that doesn't report any metrics, but handles the lifecycle of the prometheus
106    * pusher or server
107    *
108    * @param registry registry
109    * @param collectorRegistry prometheus registry
110    * @param suffix metrics suffix
111    */
112   class BasePrometheusReporter(
113       registry: MetricRegistry,
114       protected val collectorRegistry: CollectorRegistry,
115       suffix: Option[String]
116     ) extends ScheduledReporter(registry, "prometheus", MetricFilter.ALL, TimeUnit.MILLISECONDS, TimeUnit.MILLISECONDS) {
117 
118     private val sampler = suffix match {
119       case None => new DefaultSampleBuilder()
120       case Some(s) => new SuffixSampleBuilder(s)
121     }
122     new DropwizardExports(registry, sampler).register(collectorRegistry)
123 
124     // since prometheus scrapes metrics, we don't have to report them here
125 
126     override def start(initialDelay: Long, period: Long, unit: TimeUnit): Unit = {}
127 
128     override def report(
129         gauges: java.util.SortedMap[String, Gauge[_]],
130         counters: java.util.SortedMap[String, Counter],
131         histograms: java.util.SortedMap[String, Histogram],
132         meters: java.util.SortedMap[String, Meter],
133         timers: java.util.SortedMap[String, Timer]): Unit = {}
134   }
135 
136   /**
137    * Adds a suffix to each metric
138    *
139    * @param suffix suffix
140    */
141   class SuffixSampleBuilder(suffix: String) extends SampleBuilder {
142 
143     private val builder = new DefaultSampleBuilder()
144 
145     override def createSample(
146         dropwizardName: String,
147         nameSuffix: String,
148         additionalLabelNames: java.util.List[String],
149         additionalLabelValues: java.util.List[String],
150         value: Double): MetricFamilySamples.Sample = {
151       builder.createSample(dropwizardName, suffix, additionalLabelNames, additionalLabelValues, value)
152     }
153   }
154 }