Let handlers dedfine a default request type

1 files changed, 18 insertions, 2 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;