Improve timeout handling in /doc/v1

Set timeout based on qhen HTTP request was connected to the container, instead of when
it is dispatched to the document API; this should reduce the chance of the client
going away before we close the request from the handler

2 files changed, 18 insertions, 14 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
index 5126f7e0f43..9f53007ce74 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
@@ -200,7 +200,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                                 ClusterListConfig clusterListConfig,
                                 AllClustersBucketSpacesConfig bucketSpacesConfig,
                                 DocumentOperationExecutorConfig executorConfig) {
-        this(Clock.systemUTC(), Duration.ofSeconds(5), metric, metricReceiver, documentAccess,
+        this(Clock.systemUTC(), Duration.ofMillis(500), metric, metricReceiver, documentAccess,
              documentManagerConfig, executorConfig, clusterListConfig, bucketSpacesConfig);
     }
 
@@ -237,9 +237,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         HttpRequest request = (HttpRequest) rawRequest;
         try {
             // Set a higher HTTP layer timeout than the document API timeout, to prefer triggering the latter.
-            request.setTimeout(  getProperty(request, TIMEOUT, timeoutMillisParser).orElse(defaultTimeout.toMillis())
-                               + handlerTimeout.toMillis(),
-                               MILLISECONDS);
+            request.setTimeout(doomMillis(request) - clock.millis() + handlerTimeout.toMillis(), MILLISECONDS);
 
             Path requestPath = Path.withoutValidation(request.getUri()); // No segment validation here, as document IDs can be anything.
             for (String path : handlers.keySet()) {
@@ -267,7 +265,8 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
 
     @Override
     public void handleTimeout(Request request, ResponseHandler responseHandler) {
-        timeout((HttpRequest) request, "Timeout after " + (request.getTimeout(MILLISECONDS) - handlerTimeout.toMillis()) + "ms", responseHandler);
+        HttpRequest httpRequest = (HttpRequest) request;
+        timeout(httpRequest, "Timeout after " + (getProperty(httpRequest, TIMEOUT, timeoutMillisParser).orElse(defaultTimeout.toMillis())) + "ms", responseHandler);
     }
 
     @Override
@@ -523,9 +522,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
     private DocumentOperationParameters parametersFromRequest(HttpRequest request, String... names) {
         DocumentOperationParameters parameters = getProperty(request, TRACELEVEL, integerParser).map(parameters()::withTraceLevel)
                                                                                                 .orElse(parameters());
-        parameters = getProperty(request, TIMEOUT, timeoutMillisParser).map(clock.instant()::plusMillis)
-                                                                       .map(parameters::withDeadline)
-                                                                       .orElse(parameters);
+        parameters = parameters.withDeadline(Instant.ofEpochMilli(doomMillis(request)));
         for (String name : names)
             parameters = switch (name) {
                 case CLUSTER ->
@@ -1168,7 +1165,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         parameters.setFieldSet(getProperty(request, FIELD_SET).orElse(path.documentType().map(type -> type + ":[document]").orElse(DocumentOnly.NAME)));
         parameters.setMaxTotalHits(wantedDocumentCount);
         parameters.visitInconsistentBuckets(true);
-        long timeoutMs = Math.max(1, request.getTimeout(MILLISECONDS) - handlerTimeout.toMillis());
+        long timeoutMs = Math.max(1, doomMillis(request) - clock.millis());
         if (streamed) {
             StaticThrottlePolicy throttlePolicy = new DynamicThrottlePolicy().setMinWindowSize(1).setWindowSizeIncrement(1);
             concurrency.ifPresent(throttlePolicy::setMaxPendingCount);
@@ -1188,7 +1185,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         VisitorParameters parameters = parseCommonParameters(request, path, Optional.of(requireProperty(request, CLUSTER)));
         parameters.setThrottlePolicy(new DynamicThrottlePolicy().setMinWindowSize(1).setWindowSizeIncrement(1));
         long timeChunk = getProperty(request, TIME_CHUNK, timeoutMillisParser).orElse(60_000L);
-        parameters.setSessionTimeoutMs(Math.max(1, Math.min(timeChunk, request.getTimeout(MILLISECONDS) - handlerTimeout.toMillis())));
+        parameters.setSessionTimeoutMs(Math.max(1, Math.min(timeChunk, doomMillis(request) - clock.millis())));
         return parameters;
     }
 
@@ -1345,7 +1342,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             AtomicReference<String> error = new AtomicReference<>(); // Set if error occurs during processing of visited documents.
             callback.onStart(response, fullyApplied);
             VisitorControlHandler controller = new VisitorControlHandler() {
-                final ScheduledFuture<?> abort = streaming ? visitDispatcher.schedule(this::abort, request.getTimeout(MILLISECONDS), MILLISECONDS) : null;
+                final ScheduledFuture<?> abort = streaming ? visitDispatcher.schedule(this::abort, doomMillis(request) - clock.millis(), MILLISECONDS) : null;
                 final AtomicReference<VisitorSession> session = new AtomicReference<>();
                 @Override public void setSession(VisitorControlSession session) { // Workaround for broken session API ಠ_ಠ
                     super.setSession(session);
@@ -1426,6 +1423,12 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
 
     // ------------------------------------------------ Helpers ------------------------------------------------
 
+    private static long doomMillis(HttpRequest request) {
+        long connectedAtMillis = request.getConnectedAt(MILLISECONDS);
+        long requestTimeoutMillis = getProperty(request, TIMEOUT, timeoutMillisParser).orElse(defaultTimeout.toMillis());
+        return connectedAtMillis + requestTimeoutMillis;
+    }
+
     private static String requireProperty(HttpRequest request, String name) {
         return getProperty(request, name)
                 .orElseThrow(() -> new IllegalArgumentException("Must specify '" + name + "' at '" + request.getUri().getRawPath() + "'"));
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
index 7696fd2196c..1d81c45daf1 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
@@ -89,6 +89,7 @@ import static com.yahoo.jdisc.http.HttpRequest.Method.POST;
 import static com.yahoo.jdisc.http.HttpRequest.Method.PUT;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -221,7 +222,7 @@ public class DocumentV1ApiTest {
             assertEquals(100, ((StaticThrottlePolicy) parameters.getThrottlePolicy()).getMaxPendingCount());
             assertEquals("[id]", parameters.getFieldSet());
             assertEquals("(all the things)", parameters.getDocumentSelection());
-            assertEquals(6000, parameters.getSessionTimeoutMs());
+            assertTrue(6000 <= parameters.getSessionTimeoutMs()); // Static clock in handler < connected time for request, test artefact.
             assertEquals(9, parameters.getTraceLevel());
             assertEquals(1_000_000, parameters.getFromTimestamp());
             assertEquals(2_000_000, parameters.getToTimestamp());
@@ -283,7 +284,7 @@ public class DocumentV1ApiTest {
             assertEquals(1, ((StaticThrottlePolicy) parameters.getThrottlePolicy()).getMaxPendingCount());
             assertEquals("[id]", parameters.getFieldSet());
             assertEquals("(all the things)", parameters.getDocumentSelection());
-            assertEquals(6000, parameters.getTimeoutMs());
+            assertTrue(6000 <= parameters.getTimeoutMs()); // Static clock in handler < connected time for request, test artefact.
             assertEquals(4, parameters.getSlices());
             assertEquals(1, parameters.getSliceId());
             assertEquals(0, parameters.getFromTimestamp()); // not set; 0 is default
@@ -812,7 +813,7 @@ public class DocumentV1ApiTest {
 
         // TIMEOUT is a 504
         access.session.expect((id, parameters) -> {
-            assertEquals(clock.instant().plusSeconds(1000), parameters.deadline().get());
+            assertFalse(clock.instant().plusSeconds(1000).isAfter(parameters.deadline().get())); // Static clock in handler vs real clock in Request.
             parameters.responseHandler().get().handleResponse(new Response(0, "timeout", Response.Outcome.TIMEOUT));
             return new Result();
         });