001package io.prometheus.client.dropwizard.samplebuilder;
002
003import io.prometheus.client.Collector;
004
005import java.util.ArrayList;
006import java.util.Collections;
007import java.util.List;
008import java.util.Map;
009
010/**
011 * Custom {@link SampleBuilder} implementation to allow Dropwizard metrics to be translated to Prometheus metrics including custom labels and names.
012 * Prometheus metric name and labels are extracted from the Dropwizard name based on the provided list of {@link MapperConfig}s.
013 * The FIRST matching config will be used.
014 * If no config is matched, the {@link DefaultSampleBuilder} is used.
015 */
016public class CustomMappingSampleBuilder implements SampleBuilder {
017    private final List<CompiledMapperConfig> compiledMapperConfigs;
018    private final DefaultSampleBuilder defaultMetricSampleBuilder = new DefaultSampleBuilder();
019
020    public CustomMappingSampleBuilder(final List<MapperConfig> mapperConfigs) {
021        if (mapperConfigs == null || mapperConfigs.isEmpty()) {
022            throw new IllegalArgumentException("CustomMappingSampleBuilder needs some mapper configs!");
023        }
024
025        this.compiledMapperConfigs = new ArrayList<CompiledMapperConfig>(mapperConfigs.size());
026        for (MapperConfig config : mapperConfigs) {
027            this.compiledMapperConfigs.add(new CompiledMapperConfig(config));
028        }
029    }
030
031    @Override
032    public Collector.MetricFamilySamples.Sample createSample(final String dropwizardName, final String nameSuffix, final List<String> additionalLabelNames, final List<String> additionalLabelValues, final double value) {
033        if (dropwizardName == null) {
034            throw new IllegalArgumentException("Dropwizard metric name cannot be null");
035        }
036
037        CompiledMapperConfig matchingConfig = null;
038        for (CompiledMapperConfig config : this.compiledMapperConfigs) {
039            if (config.pattern.matches(dropwizardName)) {
040                matchingConfig = config;
041                break;
042            }
043        }
044
045        if (matchingConfig != null) {
046            final Map<String, String> params = matchingConfig.pattern.extractParameters(dropwizardName);
047            final NameAndLabels nameAndLabels = getNameAndLabels(matchingConfig.mapperConfig, params);
048            nameAndLabels.labelNames.addAll(additionalLabelNames);
049            nameAndLabels.labelValues.addAll(additionalLabelValues);
050            return defaultMetricSampleBuilder.createSample(
051                    nameAndLabels.name, nameSuffix,
052                    nameAndLabels.labelNames,
053                    nameAndLabels.labelValues,
054                    value
055            );
056        }
057
058
059        return defaultMetricSampleBuilder.createSample(
060                dropwizardName, nameSuffix,
061                additionalLabelNames,
062                additionalLabelValues,
063                value
064        );
065    }
066
067    protected NameAndLabels getNameAndLabels(final MapperConfig config, final Map<String, String> parameters) {
068        final String metricName = formatTemplate(config.getName(), parameters);
069        final List<String> labels = new ArrayList<String>(config.getLabels().size());
070        final List<String> labelValues = new ArrayList<String>(config.getLabels().size());
071        for (Map.Entry<String, String> entry : config.getLabels().entrySet()) {
072            labels.add(entry.getKey());
073            labelValues.add(formatTemplate(entry.getValue(), parameters));
074        }
075
076        return new NameAndLabels(metricName, labels, labelValues);
077    }
078
079    private String formatTemplate(final String template, final Map<String, String> params) {
080        String result = template;
081        for (Map.Entry<String, String> entry : params.entrySet()) {
082            result = result.replace(entry.getKey(), entry.getValue());
083        }
084
085        return result;
086    }
087
088    static class CompiledMapperConfig {
089        final MapperConfig mapperConfig;
090        final GraphiteNamePattern pattern;
091
092        CompiledMapperConfig(final MapperConfig mapperConfig) {
093            this.mapperConfig = mapperConfig;
094            this.pattern = new GraphiteNamePattern(mapperConfig.getMatch());
095        }
096    }
097
098    static class NameAndLabels {
099        final String name;
100        final List<String> labelNames;
101        final List<String> labelValues;
102
103        NameAndLabels(final String name, final List<String> labelNames, final List<String> labelValues) {
104            this.name = name;
105            this.labelNames = labelNames;
106            this.labelValues = labelValues;
107        }
108    }
109}