Reimplement HealthCheckProxyHandler as an asynchronous handler

1 files changed, 95 insertions, 24 deletions

diff --git a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HealthCheckProxyHandler.java b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HealthCheckProxyHandler.java
index 2722c21bce3..5a273548794 100644
--- a/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HealthCheckProxyHandler.java
+++ b/jdisc_http_service/src/main/java/com/yahoo/jdisc/http/server/jetty/HealthCheckProxyHandler.java
@@ -1,7 +1,7 @@
 // Copyright 2019 Oath Inc. 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.Response;
+import com.yahoo.concurrent.DaemonThreadFactory;
 import com.yahoo.jdisc.http.ConnectorConfig;
 import com.yahoo.security.SslContextBuilder;
 import com.yahoo.security.tls.TransportSecurityOptions;
@@ -15,6 +15,7 @@ import org.apache.http.client.methods.HttpGet;
 import org.apache.http.conn.ssl.NoopHostnameVerifier;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.util.EntityUtils;
 import org.eclipse.jetty.server.DetectorConnectionFactory;
 import org.eclipse.jetty.server.Request;
 import org.eclipse.jetty.server.SslConnectionFactory;
@@ -22,8 +23,10 @@ import org.eclipse.jetty.server.handler.HandlerWrapper;
 import org.eclipse.jetty.util.ssl.SslContextFactory;
 
 import javax.net.ssl.SSLContext;
+import javax.servlet.AsyncContext;
 import javax.servlet.ServletException;
 import javax.servlet.ServletOutputStream;
+import javax.servlet.WriteListener;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import java.io.IOException;
@@ -32,10 +35,11 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-import static com.yahoo.jdisc.Response.Status.NOT_FOUND;
 import static com.yahoo.jdisc.http.core.HttpServletRequestUtils.getConnectorLocalPort;
 
 /**
@@ -49,6 +53,7 @@ class HealthCheckProxyHandler extends HandlerWrapper {
 
     private static final String HEALTH_CHECK_PATH = "/status.html";
 
+    private final Executor executor = Executors.newSingleThreadExecutor(new DaemonThreadFactory("health-check-proxy-client"));
     private final Map<Integer, ProxyTarget> portToProxyTargetMapping;
 
     HealthCheckProxyHandler(List<JDiscServerConnector> connectors) {
@@ -89,27 +94,13 @@ class HealthCheckProxyHandler extends HandlerWrapper {
         int localPort = getConnectorLocalPort(servletRequest);
         ProxyTarget proxyTarget = portToProxyTargetMapping.get(localPort);
         if (proxyTarget != null) {
+            AsyncContext asyncContext = request.startAsync();
+            ServletOutputStream out = servletResponse.getOutputStream();
             if (servletRequest.getRequestURI().equals(HEALTH_CHECK_PATH)) {
-                try (CloseableHttpResponse proxyResponse = proxyTarget.requestStatusHtml()) {
-                    servletResponse.setStatus(proxyResponse.getStatusLine().getStatusCode());
-                    servletResponse.setHeader("Vespa-Health-Check-Proxy-Target", Integer.toString(proxyTarget.port));
-                    HttpEntity entity = proxyResponse.getEntity();
-                    if (entity != null) {
-                        Header contentType = entity.getContentType();
-                        if (contentType != null) {
-                            servletResponse.addHeader("Content-Type", contentType.getValue());
-                        }
-                        try (ServletOutputStream output = servletResponse.getOutputStream()) {
-                            entity.getContent().transferTo(output);
-                        }
-                    }
-                } catch (Exception e) { // Typically timeouts which are reported as SSLHandshakeException
-                    String message = String.format("Health check request from port %d to %d failed: %s", localPort, proxyTarget.port, e.getMessage());
-                    log.log(Level.FINE, message, e);
-                    servletResponse.sendError(Response.Status.INTERNAL_SERVER_ERROR, message);
-                }
+                executor.execute(new ProxyRequestTask(asyncContext, proxyTarget, servletResponse, out));
             } else {
-                servletResponse.sendError(NOT_FOUND);
+                servletResponse.setStatus(HttpServletResponse.SC_NOT_FOUND);
+                asyncContext.complete();
             }
         } else {
             _handler.handle(target, request, servletRequest, servletResponse);
@@ -124,11 +115,53 @@ class HealthCheckProxyHandler extends HandlerWrapper {
         super.doStop();
     }
 
+    private static class ProxyRequestTask implements Runnable {
+
+        final AsyncContext asyncContext;
+        final ProxyTarget target;
+        final HttpServletResponse servletResponse;
+        final ServletOutputStream output;
+
+        ProxyRequestTask(AsyncContext asyncContext, ProxyTarget target, HttpServletResponse servletResponse, ServletOutputStream output) {
+            this.asyncContext = asyncContext;
+            this.target = target;
+            this.servletResponse = servletResponse;
+            this.output = output;
+        }
+
+        @Override
+        public void run() {
+            StatusResponse statusResponse = target.requestStatusHtml();
+            servletResponse.setStatus(statusResponse.statusCode);
+            if (statusResponse.contentType != null) {
+                servletResponse.setHeader("Content-Type", statusResponse.contentType);
+            }
+            output.setWriteListener(new WriteListener() {
+                @Override
+                public void onWritePossible() throws IOException {
+                    if (output.isReady()) {
+                        if (statusResponse.content != null) {
+                            output.write(statusResponse.content);
+                        }
+                        asyncContext.complete();
+                    }
+                }
+
+                @Override
+                public void onError(Throwable t) {
+                    log.log(Level.FINE, t, () -> "Failed to write status response: " + t.getMessage());
+                    asyncContext.complete();
+                }
+            });
+        }
+    }
+
     private static class ProxyTarget implements AutoCloseable {
         final int port;
         final Duration timeout;
         final SslContextFactory.Server sslContextFactory;
         volatile CloseableHttpClient client;
+        volatile StatusResponse lastResponse;
 
         ProxyTarget(int port, Duration timeout, SslContextFactory.Server sslContextFactory) {
             this.port = port;
@@ -136,9 +169,32 @@ class HealthCheckProxyHandler extends HandlerWrapper {
             this.sslContextFactory = sslContextFactory;
         }
 
-        CloseableHttpResponse requestStatusHtml() throws IOException {
-            return client()
-                    .execute(new HttpGet("https://localhost:" + port + HEALTH_CHECK_PATH));
+        StatusResponse requestStatusHtml() {
+            StatusResponse response = lastResponse;
+            if (response != null && !response.isExpired()) {
+                return response;
+            }
+            StatusResponse statusResponse = getStatusResponse();
+            lastResponse = statusResponse;
+            return statusResponse;
+        }
+
+        private StatusResponse getStatusResponse() {
+            try (CloseableHttpResponse clientResponse = client().execute(new HttpGet("https://localhost:" + port + HEALTH_CHECK_PATH))) {
+                int statusCode = clientResponse.getStatusLine().getStatusCode();
+                HttpEntity entity = clientResponse.getEntity();
+                if (entity != null) {
+                    Header contentTypeHeader = entity.getContentType();
+                    String contentType = contentTypeHeader != null ? contentTypeHeader.getValue() : null;
+                    byte[] content = EntityUtils.toByteArray(entity);
+                    return new StatusResponse(statusCode, contentType, content);
+                } else {
+                    return new StatusResponse(statusCode, null, null);
+                }
+            } catch (Exception e) {
+                log.log(Level.FINE, e, () -> "Proxy request failed" + e.getMessage());
+                return new StatusResponse(500, "text/plain", e.getMessage().getBytes());
+            }
         }
 
         // Client construction must be delayed to ensure that the SslContextFactory is started before calling getSslContext().
@@ -200,4 +256,19 @@ class HealthCheckProxyHandler extends HandlerWrapper {
             }
         }
     }
+
+    private static class StatusResponse {
+        final long createdAt = System.nanoTime();
+        final int statusCode;
+        final String contentType;
+        final byte[] content;
+
+        StatusResponse(int statusCode, String contentType, byte[] content) {
+            this.statusCode = statusCode;
+            this.contentType = contentType;
+            this.content = content;
+        }
+
+        boolean isExpired() { return System.nanoTime() - createdAt > Duration.ofSeconds(1).toNanos(); }
+    }
 }