1 files changed, 300 insertions, 0 deletions
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpResponseStatisticsCollector.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpResponseStatisticsCollector.java
new file mode 100644
index 00000000000..82c445c7ca9
--- /dev/null
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/HttpResponseStatisticsCollector.java
@@ -0,0 +1,300 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.jdisc.http.server.jetty;
+
+import com.yahoo.jdisc.http.HttpRequest;
+import org.eclipse.jetty.http.HttpStatus;
+import org.eclipse.jetty.server.AsyncContextEvent;
+import org.eclipse.jetty.server.Handler;
+import org.eclipse.jetty.server.HttpChannelState;
+import org.eclipse.jetty.server.Request;
+import org.eclipse.jetty.server.handler.HandlerWrapper;
+import org.eclipse.jetty.util.FutureCallback;
+import org.eclipse.jetty.util.component.Graceful;
+
+import javax.servlet.AsyncEvent;
+import javax.servlet.AsyncListener;
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.LongAdder;
+
+/**
+ * HttpResponseStatisticsCollector collects statistics about HTTP response types aggregated by category
+ * (1xx, 2xx, etc). It is similar to {@link org.eclipse.jetty.server.handler.StatisticsHandler}
+ * with the distinction that this class collects response type statistics grouped
+ * by HTTP method and only collects the numbers that are reported as metrics from Vespa.
+ *
+ * @author ollivir
+ */
+public class HttpResponseStatisticsCollector extends HandlerWrapper implements Graceful {
+
+    static final String requestTypeAttribute = "requestType";
+
+    private final AtomicReference<FutureCallback> shutdown = new AtomicReference<>();
+    private final List<String> monitoringHandlerPaths;
+    private final List<String> searchHandlerPaths;
+
+    public enum HttpMethod {
+        GET, PATCH, POST, PUT, DELETE, OPTIONS, HEAD, OTHER
+    }
+
+    public enum HttpScheme {
+        HTTP, HTTPS, OTHER
+    }
+
+    private static final String[] HTTP_RESPONSE_GROUPS = {
+            MetricDefinitions.RESPONSES_1XX,
+            MetricDefinitions.RESPONSES_2XX,
+            MetricDefinitions.RESPONSES_3XX,
+            MetricDefinitions.RESPONSES_4XX,
+            MetricDefinitions.RESPONSES_5XX,
+            MetricDefinitions.RESPONSES_401,
+            MetricDefinitions.RESPONSES_403
+    };
+
+    private final AtomicLong inFlight = new AtomicLong();
+    private final LongAdder[][][][] statistics;
+
+    public HttpResponseStatisticsCollector(List<String> monitoringHandlerPaths, List<String> searchHandlerPaths) {
+        this.monitoringHandlerPaths = monitoringHandlerPaths;
+        this.searchHandlerPaths = searchHandlerPaths;
+        statistics = new LongAdder[HttpScheme.values().length][HttpMethod.values().length][][];
+        for (int scheme = 0; scheme < HttpScheme.values().length; ++scheme) {
+            for (int method = 0; method < HttpMethod.values().length; method++) {
+                statistics[scheme][method] = new LongAdder[HTTP_RESPONSE_GROUPS.length][];
+                for (int group = 0; group < HTTP_RESPONSE_GROUPS.length; group++) {
+                    statistics[scheme][method][group] = new LongAdder[HttpRequest.RequestType.values().length];
+                    for (int requestType = 0; requestType < HttpRequest.RequestType.values().length; requestType++) {
+                        statistics[scheme][method][group][requestType] = new LongAdder();
+                    }
+                }
+            }
+        }
+    }
+
+    private final AsyncListener completionWatcher = new AsyncListener() {
+
+        @Override
+        public void onTimeout(AsyncEvent event) { }
+
+        @Override
+        public void onStartAsync(AsyncEvent event) {
+            event.getAsyncContext().addListener(this);
+        }
+
+        @Override
+        public void onError(AsyncEvent event) { }
+
+        @Override
+        public void onComplete(AsyncEvent event) throws IOException {
+            HttpChannelState state = ((AsyncContextEvent) event).getHttpChannelState();
+            Request request = state.getBaseRequest();
+
+            observeEndOfRequest(request, null);
+        }
+    };
+
+    @Override
+    public void handle(String path, Request baseRequest, HttpServletRequest request, HttpServletResponse response)
+            throws IOException, ServletException {
+        inFlight.incrementAndGet();
+
+        try {
+            Handler handler = getHandler();
+            if (handler != null && shutdown.get() == null && isStarted()) {
+                handler.handle(path, baseRequest, request, response);
+            } else if ( ! baseRequest.isHandled()) {
+                baseRequest.setHandled(true);
+                response.sendError(HttpStatus.SERVICE_UNAVAILABLE_503);
+            }
+        } finally {
+            HttpChannelState state = baseRequest.getHttpChannelState();
+            if (state.isSuspended()) {
+                if (state.isInitial()) {
+                    state.addListener(completionWatcher);
+                }
+            } else if (state.isInitial()) {
+                observeEndOfRequest(baseRequest, response);
+            }
+        }
+    }
+
+    private void observeEndOfRequest(Request request, HttpServletResponse flushableResponse) throws IOException {
+        int group = groupIndex(request);
+        if (group >= 0) {
+            HttpScheme scheme = getScheme(request);
+            HttpMethod method = getMethod(request);
+            HttpRequest.RequestType requestType = getRequestType(request);
+
+            statistics[scheme.ordinal()][method.ordinal()][group][requestType.ordinal()].increment();
+            if (group == 5 || group == 6) { // if 401/403, also increment 4xx
+                statistics[scheme.ordinal()][method.ordinal()][3][requestType.ordinal()].increment();
+            }
+        }
+
+        long live = inFlight.decrementAndGet();
+        FutureCallback shutdownCb = shutdown.get();
+        if (shutdownCb != null) {
+            if (flushableResponse != null) {
+                flushableResponse.flushBuffer();
+            }
+            if (live == 0) {
+                shutdownCb.succeeded();
+            }
+        }
+    }
+
+    private int groupIndex(Request request) {
+        int index = request.getResponse().getStatus();
+        if (index == 401) {
+            return 5;
+        }
+        if (index == 403) {
+            return 6;
+        }
+
+        index = index / 100 - 1; // 1xx = 0, 2xx = 1 etc.
+        if (index < 0 || index >= statistics[0].length) {
+            return -1;
+        } else {
+            return index;
+        }
+    }
+
+    private HttpScheme getScheme(Request request) {
+        switch (request.getScheme()) {
+            case "http":
+                return HttpScheme.HTTP;
+            case "https":
+                return HttpScheme.HTTPS;
+            default:
+                return HttpScheme.OTHER;
+        }
+    }
+
+    private HttpMethod getMethod(Request request) {
+        switch (request.getMethod()) {
+        case "GET":
+            return HttpMethod.GET;
+        case "PATCH":
+            return HttpMethod.PATCH;
+        case "POST":
+            return HttpMethod.POST;
+        case "PUT":
+            return HttpMethod.PUT;
+        case "DELETE":
+            return HttpMethod.DELETE;
+        case "OPTIONS":
+            return HttpMethod.OPTIONS;
+        case "HEAD":
+            return HttpMethod.HEAD;
+        default:
+            return HttpMethod.OTHER;
+        }
+    }
+
+    private HttpRequest.RequestType getRequestType(Request request) {
+        HttpRequest.RequestType requestType = (HttpRequest.RequestType)request.getAttribute(requestTypeAttribute);
+        if (requestType != null) return requestType;
+
+        // Deduce from path and method:
+        String path = request.getRequestURI();
+        for (String monitoringHandlerPath : monitoringHandlerPaths) {
+            if (path.startsWith(monitoringHandlerPath)) return HttpRequest.RequestType.MONITORING;
+        }
+        for (String searchHandlerPath : searchHandlerPaths) {
+            if (path.startsWith(searchHandlerPath)) return HttpRequest.RequestType.READ;
+        }
+        if ("GET".equals(request.getMethod())) {
+            return HttpRequest.RequestType.READ;
+        } else {
+            return HttpRequest.RequestType.WRITE;
+        }
+    }
+
+    public List<StatisticsEntry> takeStatistics() {
+        var ret = new ArrayList<StatisticsEntry>();
+        for (HttpScheme scheme : HttpScheme.values()) {
+            int schemeIndex = scheme.ordinal();
+            for (HttpMethod method : HttpMethod.values()) {
+                int methodIndex = method.ordinal();
+                for (int group = 0; group < HTTP_RESPONSE_GROUPS.length; group++) {
+                    for (HttpRequest.RequestType type : HttpRequest.RequestType.values()) {
+                        long value = statistics[schemeIndex][methodIndex][group][type.ordinal()].sumThenReset();
+                        if (value > 0) {
+                            ret.add(new StatisticsEntry(scheme.name().toLowerCase(), method.name(), HTTP_RESPONSE_GROUPS[group], type.name().toLowerCase(), value));
+                        }
+                    }
+                }
+            }
+        }
+        return ret;
+    }
+
+    @Override
+    protected void doStart() throws Exception {
+        shutdown.set(null);
+        super.doStart();
+    }
+
+    @Override
+    protected void doStop() throws Exception {
+        super.doStop();
+        FutureCallback shutdownCb = shutdown.get();
+        if ( ! shutdownCb.isDone()) {
+            shutdownCb.failed(new TimeoutException());
+        }
+    }
+
+    @Override
+    public Future<Void> shutdown() {
+        FutureCallback shutdownCb = new FutureCallback(false);
+        shutdown.compareAndSet(null, shutdownCb);
+        shutdownCb = shutdown.get();
+        if (inFlight.get() == 0) {
+            shutdownCb.succeeded();
+        }
+        return shutdownCb;
+    }
+
+    @Override
+    public boolean isShutdown() {
+        FutureCallback futureCallback = shutdown.get();
+        return futureCallback != null && futureCallback.isDone();
+    }
+
+    public static class StatisticsEntry {
+
+        public final String scheme;
+        public final String method;
+        public final String name;
+        public final String requestType;
+        public final long value;
+
+        public StatisticsEntry(String scheme, String method, String name, String requestType, long value) {
+            this.scheme = scheme;
+            this.method = method;
+            this.name = name;
+            this.requestType = requestType;
+            this.value = value;
+        }
+
+        @Override
+        public String toString() {
+            return "scheme: " + scheme +
+                   ", method: " + method +
+                   ", name: " + name +
+                   ", requestType: " + requestType +
+                   ", value: " + value;
+        }
+
+    }
+
+}