InstrumentedFactory

AbstractInstrumentedFactory.java

package io.dropwizard.servlets;

import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.codahale.metrics.annotation.ExceptionMetered;
import com.codahale.metrics.annotation.Metered;
import com.codahale.metrics.annotation.Timed;

import java.lang.reflect.Method;

import static com.codahale.metrics.MetricRegistry.name;

/**
 *
 */
public abstract class AbstractInstrumentedFactory<T> implements InstrumentedFactory<T> {

    private final MetricRegistry metricRegistry;

    public AbstractInstrumentedFactory(MetricRegistry metricRegistry) {
        this.metricRegistry = metricRegistry;
    }

    protected abstract T instrumentTimer(T underlying, Timer timer);

    protected abstract T instrumentMeter(T underlying, Meter meter);

    protected abstract T instrumentExceptionMeter(T underlying, Meter exceptionMeter, Class<? extends Throwable> cause);

    public T createInstrumented(T underlying, Method method, Class<?> target) {
        if(method.isAnnotationPresent(Timed.class)) {
            Timed annotation = method.getAnnotation(Timed.class);
            String name = chooseName(annotation.name(),
                    annotation.absolute(),
                    target);
            Timer timer = metricRegistry.timer(name);
            return instrumentTimer(underlying, timer);
        }

        if(method.isAnnotationPresent(Metered.class)) {
            Metered annotation = method.getAnnotation(Metered.class);
            String name = chooseName(annotation.name(),
                    annotation.absolute(),
                    target);
            Meter meter = metricRegistry.meter(name);
            return instrumentMeter(underlying, meter);
        }

        if(method.isAnnotationPresent(ExceptionMetered.class)) {
            ExceptionMetered annotation = method.getAnnotation(ExceptionMetered.class);
            String name = chooseName(annotation.name(),
                    annotation.absolute(),
                    target,
                    ExceptionMetered.DEFAULT_NAME_SUFFIX);
            Meter exceptionMeter = metricRegistry.meter(name);
            return instrumentExceptionMeter(underlying, exceptionMeter, annotation.cause());
        }

        return underlying;
    }

    protected String chooseName(String explicitName, boolean absolute, Class<?> targetClass, String... suffixes) {
        if (explicitName != null && !explicitName.isEmpty()) {
            if (absolute) {
                return explicitName;
            }
            return name(targetClass, explicitName);
        }

        return name(targetClass, suffixes);
    }

}

InstrumentedFactory.java

package io.dropwizard.servlets;

import java.lang.reflect.Method;

/**
 *
 */
public interface InstrumentedFactory<T> {

    T createInstrumented(T underlying, Method method, Class<?> targetClass);

}

TaskServlet.java

package io.dropwizard.servlets.tasks;

import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Maps;
import com.google.common.net.MediaType;
import io.dropwizard.servlets.AbstractInstrumentedFactory;
import io.dropwizard.servlets.InstrumentedFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.Enumeration;
import java.util.concurrent.ConcurrentMap;

/**
 * A servlet which provides access to administrative {@link Task}s. It only responds to {@code POST}
 * requests, since most {@link Task}s aren't side-effect free, and passes along the query string
 * parameters of the request to the task as a multimap.
 *
 * @see Task
 */
public class TaskServlet extends HttpServlet {
    private static final long serialVersionUID = 7404713218661358124L;
    private static final Logger LOGGER = LoggerFactory.getLogger(TaskServlet.class);
    private final ConcurrentMap<String, Task> tasks;
    private final ConcurrentMap<Task, TaskExecutor> taskExecutors;

    private final MetricRegistry metricRegistry;

    /**
     * Creates a new TaskServlet.
     */
    public TaskServlet(MetricRegistry metricRegistry) {
        this.metricRegistry = metricRegistry;
        this.tasks = Maps.newConcurrentMap();
        this.taskExecutors = Maps.newConcurrentMap();
    }

    public void add(Task task) {
        tasks.put('/' + task.getName(), task);

        TaskExecutor taskExecutor = new TaskExecutor(task);
        try {
            Method executeMethod = task.getClass().getMethod("execute",
                    ImmutableMultimap.class, PrintWriter.class);

            InstrumentedFactory<TaskExecutor> measuredTaskExecutor = 
                new InstrumentedTaskExecutorFactory(metricRegistry);
            taskExecutor = measuredTaskExecutor.createInstrumented(
                taskExecutor, executeMethod, task.getClass());
        } catch (NoSuchMethodException e) {
        }

        taskExecutors.put(task, taskExecutor);
    }

    @Override
    protected void doPost(HttpServletRequest req,
                          HttpServletResponse resp) throws ServletException, IOException {
        final Task task = tasks.get(req.getPathInfo());
        if (task != null) {
            resp.setContentType(MediaType.PLAIN_TEXT_UTF_8.toString());
            final PrintWriter output = resp.getWriter();
            try {
                TaskExecutor taskExecutor = taskExecutors.get(task);
                taskExecutor.executeTask(getParams(req), output);
            } catch (Exception e) {
                LOGGER.error("Error running {}", task.getName(), e);
                resp.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
                output.println();
                output.println(e.getMessage());
                e.printStackTrace(output);
            } finally {
                output.close();
            }
        } else {
            resp.sendError(HttpServletResponse.SC_NOT_FOUND);
        }
    }

    private static ImmutableMultimap<String, String> getParams(HttpServletRequest req) {
        final ImmutableMultimap.Builder<String, String> results = ImmutableMultimap.builder();
        final Enumeration<String> names = req.getParameterNames();
        while (names.hasMoreElements()) {
            final String name = names.nextElement();
            final String[] values = req.getParameterValues(name);
            results.putAll(name, values);
        }
        return results.build();
    }

    public Collection<Task> getTasks() {
        return tasks.values();
    }

    private static class TaskExecutor {
        private final Task task;

        private TaskExecutor(Task task) {
            this.task = task;
        }

        public void executeTask(ImmutableMultimap<String, String> params, PrintWriter output) throws Exception {
            try {
                task.execute(params, output);
            } catch (Exception e) {
                throw e;
            }
        }
    }

    private static class TimedTask extends TaskExecutor {
        private TaskExecutor underlying;
        private final Timer timer;

        private TimedTask(TaskExecutor underlying, Timer timer) {
            super(underlying.task);
            this.underlying = underlying;
            this.timer = timer;
        }

        @Override
        public void executeTask(ImmutableMultimap<String, String> params, PrintWriter output) throws Exception {
            final Timer.Context context = timer.time();
            try {
                underlying.executeTask(params, output);
            } finally {
                context.stop();
            }
        }
    }

    private static class MeteredTask extends TaskExecutor {
        private TaskExecutor underlying;
        private final Meter meter;

        private MeteredTask(TaskExecutor underlying, Meter meter) {
            super(underlying.task);
            this.meter = meter;
            this.underlying = underlying;
        }

        @Override
        public void executeTask(ImmutableMultimap<String, String> params, PrintWriter output) throws Exception {
            meter.mark();
            underlying.executeTask(params, output);
        }
    }

    private static class ExceptionMeteredTask extends TaskExecutor {
        private TaskExecutor underlying;
        private final Meter exceptionMeter;
        private final Class<?> exceptionClass;

        private ExceptionMeteredTask(TaskExecutor underlying,
                                     Meter exceptionMeter, Class<? extends Throwable> exceptionClass) {
            super(underlying.task);
            this.underlying = underlying;
            this.exceptionMeter = exceptionMeter;
            this.exceptionClass = exceptionClass;
        }

        @Override
        public void executeTask(ImmutableMultimap<String, String> params, PrintWriter output) throws Exception {
            try {
                underlying.executeTask(params, output);
            } catch(Exception e) {
                if (exceptionMeter != null && exceptionClass.isAssignableFrom(e.getClass()) ||
                        (e.getCause() != null && exceptionClass.isAssignableFrom(e.getCause().getClass()))) {
                    exceptionMeter.mark();
                }

                throw e;
            }
        }
    }

    private static class InstrumentedTaskExecutorFactory extends AbstractInstrumentedFactory<TaskExecutor> {

        private InstrumentedTaskExecutorFactory(MetricRegistry metricRegistry) {
            super(metricRegistry);
        }

        @Override
        protected TaskExecutor instrumentTimer(TaskExecutor underlying, Timer timer) {
            return new TimedTask(underlying, timer);
        }

        @Override
        protected TaskExecutor instrumentMeter(TaskExecutor underlying, Meter meter) {
            return new MeteredTask(underlying, meter);
        }

        @Override
        protected TaskExecutor instrumentExceptionMeter(TaskExecutor underlying, Meter exceptionMeter, Class<? extends Throwable> cause) {
            return new ExceptionMeteredTask(underlying, exceptionMeter, cause);
        }
    }

}