SpringBootMetricsCollector.java

package me.kondratovich.actuator.prometheus;

import static io.micrometer.core.instrument.Statistic.Count;
import static io.micrometer.core.instrument.Statistic.Total;
import static io.micrometer.core.instrument.Statistic.TotalTime;
import static io.micrometer.core.instrument.Statistic.Value;

import com.google.common.base.Objects;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimaps;
import io.micrometer.core.instrument.Measurement;
import io.micrometer.core.instrument.Meter;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Statistic;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.composite.CompositeMeterRegistry;
import io.prometheus.client.Collector;
import io.prometheus.client.Collector.MetricFamilySamples.Sample;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

public class SpringBootMetricsCollector extends Collector {
    private final MeterRegistry registry;

    public SpringBootMetricsCollector(MeterRegistry registry) {
        this.registry = registry;
    }

    @Override
    public List<MetricFamilySamples> collect() {
        List<Meter> meters = flatten(registry)
            .flatMap(mr -> mr.getMeters().stream())
            .collect(Collectors.toList());

        ListMultimap<MetricId, Meter> namedMeters = Multimaps.index(meters, MetricId::new);
        return namedMeters.asMap().entrySet().stream()
            .map(entry -> {
                MetricId metricId = entry.getKey();

                String name = metricId.getOuterName();
                String help = "Spring metric " + metricId.getInnerName();

                return metricId.getType().map(type -> {
                    List<Sample> samples = entry.getValue().stream()
                        .flatMap(m -> samples(name, m)
                            .map(Collection::stream)
                            .orElse(Stream.empty()))
                        .collect(Collectors.toList());

                    return new MetricFamilySamples(name, type, help, samples);
                });
            })
            .flatMap(optional -> optional
                .map(Stream::of)
                .orElse(Stream.empty()))
            .collect(Collectors.toList());
    }

    private static Optional<List<Sample>> samples(String name, Meter meter) {
        Iterable<Measurement> measurements = meter.measure();

        List<String> labels = labels(meter.getId().getTags());
        List<String> values = values(meter.getId().getTags());

        switch (meter.getType()) {
            case Counter: {
                return measurement(measurements, Count).map(m -> {
                    Sample sample = new Sample(name, labels, values, m.getValue());
                    return Collections.singletonList(sample);
                });
            }
            case Gauge:
                return measurement(measurements, Value).map(m -> {
                    Sample sample = new Sample(name, labels, values, m.getValue());
                    return Collections.singletonList(sample);
                });
            case Timer:
                return measurement(measurements, Count).flatMap(m0 ->
                    measurement(measurements, TotalTime).map(m1 ->
                        Lists.newArrayList(
                            new Sample(name + "_count", labels, values, m0.getValue()),
                            new Sample(name + "_sum", labels, values, m1.getValue())
                        )
                    )
                );
            case DistributionSummary:
                return measurement(measurements, Count).flatMap(m0 ->
                    measurement(measurements, Total).map(m1 ->
                        Lists.newArrayList(
                            new Sample(name + "_count", labels, values, m0.getValue()),
                            new Sample(name + "_sum", labels, values, m1.getValue())
                        )
                    )
                );
            default:
                return Optional.empty();
        }
    }

    private static Optional<Type> type(Meter meter) {
        switch (meter.getType()) {
            case Counter:
                return Optional.of(Type.COUNTER);
            case Gauge:
                return Optional.of(Type.GAUGE);
            case Timer:
            case DistributionSummary:
                return Optional.of(Type.SUMMARY);
            default:
                return Optional.empty();
        }
    }

    private static Optional<Measurement> measurement(Iterable<Measurement> measurements, Statistic statistic) {
        return StreamSupport.stream(measurements.spliterator(), false)
            .filter(m -> m.getStatistic() == statistic)
            .findFirst();
    }

    private static List<String> labels(Iterable<Tag> tags) {
        return StreamSupport.stream(tags.spliterator(), false)
            .map(Tag::getKey)
            .collect(Collectors.toList());
    }

    private static List<String> values(Iterable<Tag> tags) {
        return StreamSupport.stream(tags.spliterator(), false)
            .map(Tag::getValue)
            .collect(Collectors.toList());
    }

    private static Stream<MeterRegistry> flatten(MeterRegistry meterRegistry) {
        if (meterRegistry instanceof CompositeMeterRegistry) {
            return ((CompositeMeterRegistry) meterRegistry).getRegistries().stream()
                .flatMap(SpringBootMetricsCollector::flatten);
        }

        return Stream.of(meterRegistry);
    }

    private static class MetricId {
        private final Meter meter;

        private MetricId(Meter meter) {
            this.meter = meter;
        }

        private String getInnerName() {
            return meter.getId().getName();
        }

        private String getOuterName() {
            String underscored = getInnerName().replaceAll("\\.", "_");
            return (underscored + "_" + meter.getType()).toLowerCase();
        }

        private Optional<Type> getType() {
            return type(meter);
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (!(o instanceof MetricId)) return false;
            MetricId metricId = (MetricId) o;
            return Objects.equal(getOuterName(), metricId.getOuterName());
        }

        @Override
        public int hashCode() {
            return Objects.hashCode(getOuterName());
        }
    }
}