KtorMetrics.kt

import io.ktor.application.*
import io.ktor.pipeline.*
import io.ktor.routing.*
import io.ktor.util.*
import io.micrometer.core.instrument.*
import io.micrometer.core.instrument.binder.MeterBinder
import io.micrometer.core.instrument.binder.jvm.*
import io.micrometer.core.instrument.binder.system.*
import java.util.concurrent.atomic.AtomicInteger

class Metrics(val registry: MeterRegistry) {
	val baseName: String = "ktor.calls"
	private val active = registry.gauge("$baseName.active", AtomicInteger(0))
	private val exceptions = registry.counter("$baseName.exceptions")

	class Configuration {
		lateinit var registry: MeterRegistry

		var meterBinders: List<MeterBinder> = listOf(
				ClassLoaderMetrics(),
				JvmMemoryMetrics(),
				JvmGcMetrics(),
				ProcessorMetrics(),
				JvmThreadMetrics(),
				FileDescriptorMetrics()
		)
	}

	companion object Feature : ApplicationFeature<Application, Configuration, Metrics> {
		override val key = AttributeKey<Metrics>("metrics")

		private class RoutingMetrics(val name: String, val port: String, val context: Timer.Sample)

		private val routingMetricsKey = AttributeKey<RoutingMetrics>("metrics")

		override fun install(pipeline: Application, configure: Configuration.() -> Unit): Metrics {
			val configuration = Configuration().apply(configure)
			val feature = Metrics(configuration.registry)

			configuration.meterBinders.forEach { it.bindTo(configuration.registry) }

			val phase = PipelinePhase("Metrics")
			pipeline.insertPhaseBefore(ApplicationCallPipeline.Infrastructure, phase)
			pipeline.intercept(phase) {
				feature.before(call)
				try {
					proceed()
				} catch (e: Exception) {
					feature.exception(call, e)
					throw e
				} finally {
					feature.after(call)
				}
			}

			pipeline.environment.monitor.subscribe(Routing.RoutingCallStarted) { call ->
				val port: String = call.request.local.port.toString()

				val name = call.route.toString()
				val meter = feature.registry.counter("${feature.baseName}.meter")
				val timer = Timer.start(feature.registry)
				meter.increment()
				call.attributes.put(routingMetricsKey, RoutingMetrics(name, port, timer))
			}

			pipeline.environment.monitor.subscribe(Routing.RoutingCallFinished) { call ->
				val routingMetrics = call.attributes.take(routingMetricsKey)
				val status = call.response.status()?.value ?: 0

				routingMetrics.context.stop(feature.registry.timer("${feature.baseName}.timer",
				"path", routingMetrics.name, "port", routingMetrics.port, "status", status.toString()
				))
			}

			return feature
		}
	}


	private data class CallMeasure(val timer: Timer.Sample)

	private val measureKey = AttributeKey<CallMeasure>("metrics")

	private fun before(call: ApplicationCall) {
		active.incrementAndGet()

		call.attributes.put(measureKey, CallMeasure(Timer.start(registry)))
	}

	private fun after(call: ApplicationCall) {
		active.decrementAndGet()

		registry.counter("$baseName.status",
				"status", (call.response.status()?.value ?: 0).toString()
		).increment()

		call.attributes.getOrNull(measureKey)?.apply {
			timer.stop(registry.timer("$baseName.duration"))
		}
	}

	@Suppress("UNUSED_PARAMETER")
	private fun exception(call: ApplicationCall, e: Throwable) {
		exceptions.increment()
	}
}