diff options
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 |