SamzaBridgedMetricRegistry.java

import com.codahale.metrics.*;
import com.google.common.base.Preconditions;
import org.apache.samza.metrics.Gauge;
import org.apache.samza.metrics.MetricsRegistry;

import java.util.HashMap;
import java.util.Map;

/**
 * Bridges <a href="http://dropwizard.github.io/metrics/">Dropwizard Metrics</a> (aka Coda Hale metrics)
 * {@link MetricRegistry} to Samza's metrics.
 *
 * Samza defines its own set of interfaces and classes for metrics that are almost identical to Coda Hale metrics.
 * 
 * This allows business logic for Samza job that is not strictly bound to Samza runtime to use
 *  Coda Hale metrics and remain independent from Samza dependencies.
 *
 * Inspired by <a href="https://github.com/quantiply/samza-coda-metrics">quantiply/samza-coda-metrics</a>,
 * but takes a slightly different approach where the metrics are registered
 * on the Coda Hale {@link MetricRegistry} rather than a special {@code MetricAdaptor} class.
 */
class SamzaBridgedMetricRegistry extends MetricRegistry {
    private final MetricsRegistry samzaRegistry;
    private final String groupName;

    SamzaBridgedMetricRegistry(MetricsRegistry samzaRegistry, String groupName) {
        this.samzaRegistry = Preconditions.checkNotNull(samzaRegistry);
        this.groupName = Preconditions.checkNotNull(groupName);
    }

    /**
     * Given a {@link Metric}, registers it under the given name.
     *
     * Additionally registers a metric in Samza's registry that mirrors the metric.
     *
     * @throws IllegalArgumentException if metric type is not handled
     */
    @Override
    public <T extends Metric> T register(String name, T metric) throws IllegalArgumentException {
        metric = super.register(name, metric);

        // Ignore metrics that have already been registered
        if (getMetrics().containsKey(name)) {
            return metric;
        }

        // Ignore MetricSet because it recursively invokes this method and each metric has already been handled.
        if (metric instanceof MetricSet) {
            return metric;
        }

        // Check if metric type is supported
        if (metric instanceof Meter ||
            metric instanceof Histogram ||
            metric instanceof Counter ||
            metric instanceof Timer) {
            throw new IllegalArgumentException(metric.getClass() + "is not currently supported with " + SamzaBridgedMetricRegistry.class);
        }

        samzaRegistry.newGauge(groupName, new MapGauge(name, metric));

        return metric;
    }

    /**
     * Samza's {@link Gauge} metric is a simple wrapper metric that is well-suited for this use case.
     *
     * @see <a href="https://github.com/quantiply/samza-coda-metrics/blob/5b64a60/src/main/java/com/quantiply/samza/MapGauge.java">Original Source</a>
     */
    public static class MapGauge extends Gauge<Map<String, Object>> {
        private Metric _metric;

        public MapGauge(String name, Metric metric) {
            super(name, new HashMap<>());
            this._metric = metric;
        }

        public Map<String, Object> getValue() {
            return this.toMap(this._metric);
        }

        private Map<String, Object> toMap(Metric metric) {
            if (metric instanceof Meter) {
                return this.meter((Meter)metric);
            } else if (metric instanceof Histogram) {
                return this.histogram((Histogram)metric);
            } else if (metric instanceof Counter) {
                return this.counter((Counter)metric);
            } else if (metric instanceof Timer) {
                return this.timer((Timer)metric);
            } else {
                return metric instanceof com.codahale.metrics.Gauge ? this.gauge((com.codahale.metrics.Gauge)metric) : null;
            }
        }

        public Map<String, Object> meter(Meter meter) {
            Map<String, Object> data = new HashMap<>();
            data.put("type", "meter");
            data.put("count", meter.getCount());
            data.put("oneMinuteRate", meter.getOneMinuteRate());
            data.put("fiveMinuteRate", meter.getFiveMinuteRate());
            data.put("fifteenMinuteRate", meter.getFifteenMinuteRate());
            data.put("meanRate", meter.getMeanRate());
            data.put("rateUnit", "SECONDS");
            return data;
        }

        private Map<String, Object> counter(Counter counter) {
            Map<String, Object> data = new HashMap<>();
            data.put("type", "counter");
            data.put("count", counter.getCount());
            return data;
        }

        private Map<String, Object> histogram(Histogram histogram) {
            Map<String, Object> data = new HashMap<>();
            Snapshot snapshot = histogram.getSnapshot();
            data.put("type", "histogram");
            data.put("samples", snapshot.size());
            data.put("min", snapshot.getMin());
            data.put("max", snapshot.getMax());
            data.put("mean", snapshot.getMean());
            data.put("stdDev", snapshot.getStdDev());
            data.put("median", snapshot.getMedian());
            data.put("75thPercentile", snapshot.get75thPercentile());
            data.put("95thPercentile", snapshot.get95thPercentile());
            data.put("98thPercentile", snapshot.get98thPercentile());
            data.put("99thPercentile", snapshot.get99thPercentile());
            data.put("999thPercentile", snapshot.get999thPercentile());
            return data;
        }

        private Map<String, Object> timer(Timer timer) {
            Map<String, Object> data = new HashMap<>();
            Snapshot snapshot = timer.getSnapshot();
            data.put("type", "timer");
            data.put("samples", snapshot.size());
            data.put("min", snapshot.getMin());
            data.put("max", snapshot.getMax());
            data.put("mean", snapshot.getMean());
            data.put("stdDev", snapshot.getStdDev());
            data.put("median", snapshot.getMedian());
            data.put("75thPercentile", snapshot.get75thPercentile());
            data.put("95thPercentile", snapshot.get95thPercentile());
            data.put("98thPercentile", snapshot.get98thPercentile());
            data.put("99thPercentile", snapshot.get99thPercentile());
            data.put("999thPercentile", snapshot.get999thPercentile());
            data.put("durationUnit", "NANOSECONDS");
            return data;
        }

        private Map<String, Object> gauge(com.codahale.metrics.Gauge gauge) {
            Map<String, Object> data = new HashMap<>();
            data.put("type", "gauge");
            data.put("value", gauge.getValue());
            return data;
        }
    }
}