Collecting AWS client metrics into codahale metrics, in java

AwsClientMetricsCollector.java

package ...;

import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.amazonaws.Request;
import com.amazonaws.Response;
import com.amazonaws.metrics.AwsSdkMetrics;
import com.amazonaws.metrics.MetricCollector;
import com.amazonaws.metrics.RequestMetricCollector;
import com.amazonaws.metrics.ServiceMetricCollector;
import com.amazonaws.util.AWSRequestMetrics;
import com.amazonaws.util.AWSRequestMetrics.Field;
import com.amazonaws.util.TimingInfo;
import com.codahale.metrics.Counter;
import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricRegistry;
import com.totango.util.Assert;

/**
 * Collects some of the AWS client metrics into codahale metrics.
 *
 * @author ran
 */
public enum AwsClientMetricsCollector {
    ;

    public static final Set<Field> SUPPORTED_METRIC_TYPES = new HashSet<>();
    static {
        SUPPORTED_METRIC_TYPES.add(Field.RequestCount);
        SUPPORTED_METRIC_TYPES.add(Field.RetryCount);
    }

    private static final Logger log = LoggerFactory.getLogger(AwsClientMetricsCollector.class);

    private static MetricRegistry codahaleMetrics;

    private final static Set<Field> collectedMetrics = new HashSet<>();

    public static void collectMetric(final Field awsMetricType) {
        if (!SUPPORTED_METRIC_TYPES.contains(awsMetricType)) {
            throw new IllegalArgumentException(String.format(
                    "The type %s is not supported. The supported types are: %s", awsMetricType,
                    SUPPORTED_METRIC_TYPES));
        }
        collectedMetrics.add(awsMetricType);
    }

    public static void uncollectMetric(final Field awsMetricType) {
        collectedMetrics.remove(awsMetricType);
    }

    public Set<Field> getCollectedMetrics() {
        return Collections.unmodifiableSet(collectedMetrics);
    }

    public static void setup(final MetricRegistry codahaleMetricsRegistry) {
        Assert.notNull(codahaleMetricsRegistry);
        codahaleMetrics = codahaleMetricsRegistry;
        final MetricCollector metricCollector = AwsSdkMetrics.getMetricCollector();
        final RequestMetricCollector requestMetricCollector = metricCollector.getRequestMetricCollector();
        final ServiceMetricCollector serviceMetricCollector = metricCollector.getServiceMetricCollector();


        AwsSdkMetrics.setMetricCollector(new MetricCollector() {

            @Override
            public boolean stop() {
                return metricCollector.stop();
            }

            @Override
            public boolean start() {
                return metricCollector.start();
            }

            @Override
            public boolean isEnabled() {
                return metricCollector.isEnabled();
            }

            @Override
            public ServiceMetricCollector getServiceMetricCollector() {
                return serviceMetricCollector;
            }

            @Override
            public RequestMetricCollector getRequestMetricCollector() {
                return new RequestMetricCollector() {

                    @Override
                    public void collectMetrics(final Request<?> request, final Response<?> response) {
                        requestMetricCollector.collectMetrics(request, response);
                        // Right now we only support RetryCount and RequestCount
                        for (final Field metricType : collectedMetrics) {
                            switch (metricType) {
                            case RequestCount:
                            case RetryCount:
                                final int cnt = metricOfRequestOrRetryCount(metricType, request, response);
                                count(request.getServiceName(), metricType.name(), cnt);
                                break;
                            default:
                                log.error("Cannot count type {}", metricType);
                            }
                        }
                    }

                    /**
                     * Returns the value of RequestCount or RetryCount.
                     * If metricType is neither RequestCount nor RetryCount then returns 0
                     *
                     * Copied and modified from PredefinedMetricTransformer.metricOfRequestOrRetryCount
                     * If you want to support more types look there.
                     * @param metricType
                     *            must be either {@link Field#RequestCount} or
                     *            {@link Field#RetryCount}; or else return just 0
                     *
                     *
                     */
                    protected int metricOfRequestOrRetryCount(
                            final Field metricType, final Request<?> req, final Object resp) {
                        final AWSRequestMetrics m = req.getAWSRequestMetrics();
                        final TimingInfo ti = m.getTimingInfo();
                        // Always retrieve the request count even for retry which is equivalent
                        // to the number of requests minus one.
                        final Number counter = ti.getCounter(Field.RequestCount.name());
                        if (counter == null) {
                            // this is possible if one of the request handlers screwed up
                            return 0;
                        }
                        final int requestCount = counter.intValue();
                        if (requestCount < 1) {
                            log.warn("request count must be at least one");
                            return 0;
                        }
                        final int count = metricType == Field.RequestCount
                                           ? requestCount
                                           : requestCount-1 // retryCount = requestCount - 1
                                           ;
                        return count;
                    }
                };
            }
        });
    }

    private static void count(final String serviceName, final String counterName, final int cnt) {
        final String key = String.format("aws.%s.%s", serviceName, counterName);
        final Metric metric = codahaleMetrics.counter(key);
        ((Counter) metric).inc(cnt);
    }
}