Let handlers dedfine a default request type

4 files changed, 38 insertions, 6 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/jdisc/ThreadedRequestHandler.java b/container-core/src/main/java/com/yahoo/container/jdisc/ThreadedRequestHandler.java
index 323935e2a26..691675dfec9 100644
--- a/container-core/src/main/java/com/yahoo/container/jdisc/ThreadedRequestHandler.java
+++ b/container-core/src/main/java/com/yahoo/container/jdisc/ThreadedRequestHandler.java
@@ -94,7 +94,7 @@ public abstract class ThreadedRequestHandler extends AbstractRequestHandler {
             }
         }
         BufferedContentChannel content = new BufferedContentChannel();
-        final RequestTask command = new RequestTask(request, content, responseHandler);
+        RequestTask command = new RequestTask(request, content, responseHandler);
         try {
             executor.execute(command);
         } catch (RejectedExecutionException e) {
@@ -106,6 +106,18 @@ public abstract class ThreadedRequestHandler extends AbstractRequestHandler {
         return content;
     }
 
+    /**
+     * <p>Returns the request type classification to use for requests to this handler.
+     * This overrides the default classification based on request method, and can in turn
+     * be overridden by setting a request type on individual responses in handleRequest
+     * whenever it is invoked (i.e not for requests that are rejected early e.g due to overload).</p>
+     *
+     * <p>This default implementation returns null.</p>
+     *
+     * @return the request type to set, or null to not override the default classification based on request method
+     */
+    protected Request.RequestType getRequestType() { return null; }
+
     public Duration getTimeout() {
         return TIMEOUT;
     }
@@ -145,7 +157,9 @@ public abstract class ThreadedRequestHandler extends AbstractRequestHandler {
      * A subclass may override this method to define a custom response.
      */
     protected void writeErrorResponseOnOverload(Request request, ResponseHandler responseHandler) {
-        ResponseDispatch.newInstance(Response.Status.SERVICE_UNAVAILABLE).dispatch(responseHandler);
+        Response response = new Response(Response.Status.SERVICE_UNAVAILABLE);
+        response.setRequestType(getRequestType());
+        ResponseDispatch.newInstance(response).dispatch(responseHandler);
     }
 
     private class RequestTask implements ResponseHandler, Runnable {
@@ -188,6 +202,8 @@ public abstract class ThreadedRequestHandler extends AbstractRequestHandler {
         @Override
         public ContentChannel handleResponse(Response response) {
             if ( tryHasResponded()) throw new IllegalStateException("Response already handled");
+            if (response.getRequestType() == null)
+                response.setRequestType(getRequestType());
             ContentChannel cc = responseHandler.handleResponse(response);
             HandlerMetricContextUtil.onHandled(request, metric, getClass());
             return cc;
diff --git a/container-search/src/main/java/com/yahoo/search/handler/SearchHandler.java b/container-search/src/main/java/com/yahoo/search/handler/SearchHandler.java
index bee25fbb47f..80dc7a7b489 100644
--- a/container-search/src/main/java/com/yahoo/search/handler/SearchHandler.java
+++ b/container-search/src/main/java/com/yahoo/search/handler/SearchHandler.java
@@ -251,6 +251,9 @@ public class SearchHandler extends LoggingRequestHandler {
         }
     }
 
+    @Override
+    public Request.RequestType getRequestType() { return Request.RequestType.READ; }
+
     private int getHttpResponseStatus(com.yahoo.container.jdisc.HttpRequest httpRequest, Result result) {
         boolean benchmarkOutput = VespaHeaders.benchmarkOutput(httpRequest);
         if (benchmarkOutput) {
@@ -326,9 +329,7 @@ public class SearchHandler extends LoggingRequestHandler {
         HttpSearchResponse response = new HttpSearchResponse(getHttpResponseStatus(request, result),
                                                              result, query, renderer,
                                                              extractTraceNode(query));
-        response.setRequestType(Request.RequestType.READ);
-        if (hostResponseHeaderKey.isPresent())
-            response.headers().add(hostResponseHeaderKey.get(), selfHostname);
+        hostResponseHeaderKey.ifPresent(key -> response.headers().add(key, selfHostname));
 
         if (benchmarking)
             VespaHeaders.benchmarkOutput(response.headers(), benchmarkCoverage, response.getTiming(),
diff --git a/container-search/src/test/java/com/yahoo/search/handler/test/SearchHandlerTestCase.java b/container-search/src/test/java/com/yahoo/search/handler/test/SearchHandlerTestCase.java
index 7d03e065fa9..2eb5901b786 100644
--- a/container-search/src/test/java/com/yahoo/search/handler/test/SearchHandlerTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/handler/test/SearchHandlerTestCase.java
@@ -8,6 +8,7 @@ import com.yahoo.container.jdisc.HttpResponse;
 import com.yahoo.container.jdisc.RequestHandlerTestDriver;
 import com.yahoo.container.jdisc.ThreadedHttpRequestHandler;
 import com.yahoo.io.IOUtils;
+import com.yahoo.jdisc.Request;
 import com.yahoo.jdisc.handler.RequestHandler;
 import com.yahoo.net.HostName;
 import com.yahoo.search.Query;
@@ -179,10 +180,24 @@ public class SearchHandlerTestCase {
                     "http://localhost/search/?yql=select%20*%20from%20foo%20where%20bar%20%3E%201453501295%27%3B");
             responseHandler.readAll();
             assertThat(responseHandler.getStatus(), is(400));
+            assertEquals(Request.RequestType.READ, responseHandler.getResponse().getRequestType());
         }
     }
 
+    @Test
+    public void testRequestType() throws Exception {
+        IOUtils.copyDirectory(new File(testDir, "config_yql"), new File(tempDir), 1);
+        generateComponentsConfigForActive();
+        configurer.reloadConfig();
 
+        SearchHandler newSearchHandler = fetchSearchHandler(configurer);
+        try (RequestHandlerTestDriver newDriver = new RequestHandlerTestDriver(newSearchHandler)) {
+            RequestHandlerTestDriver.MockResponseHandler responseHandler = newDriver.sendRequest(
+                    "http://localhost/search/?query=foo");
+            responseHandler.readAll();
+            assertEquals(Request.RequestType.READ, responseHandler.getResponse().getRequestType());
+        }
+    }
 
     // Query handling takes a different code path when a query profile is active, so we test both paths.
     @Test
diff --git a/jdisc_core/src/main/java/com/yahoo/jdisc/handler/ResponseDispatch.java b/jdisc_core/src/main/java/com/yahoo/jdisc/handler/ResponseDispatch.java
index 22a77968dfd..bd519f05a57 100644
--- a/jdisc_core/src/main/java/com/yahoo/jdisc/handler/ResponseDispatch.java
+++ b/jdisc_core/src/main/java/com/yahoo/jdisc/handler/ResponseDispatch.java
@@ -161,7 +161,7 @@ public abstract class ResponseDispatch extends ForwardingListenableFuture<Boolea
      * @param content  The provider of the Response's ByteBuffer content.
      * @return The created ResponseDispatch.
      */
-    public static ResponseDispatch newInstance(final Response response, final Iterable<ByteBuffer> content) {
+    public static ResponseDispatch newInstance(Response response, Iterable<ByteBuffer> content) {
         return new ResponseDispatch() {
 
             @Override