Merge pull request #26826 from vespa-engine/jonmv/improve-doc-v1-timeouts

Improve timeout handling in /doc/v1

2 files changed, 28 insertions, 18 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..45c93ec0755 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
@@ -132,6 +132,7 @@ import static java.util.stream.Collectors.toUnmodifiableMap;
 public class DocumentV1ApiHandler extends AbstractRequestHandler {
 
     private static final Duration defaultTimeout = Duration.ofSeconds(180); // Match document API default timeout.
+    private static final Duration handlerTimeout = Duration.ofMillis(100); // Extra time to allow for handler, JDisc and jetty to complete.
 
     private static final Logger log = Logger.getLogger(DocumentV1ApiHandler.class.getName());
     private static final Parser<Integer> integerParser = Integer::parseInt;
@@ -175,7 +176,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
     private static final String TO_TIMESTAMP = "toTimestamp";
 
     private final Clock clock;
-    private final Duration handlerTimeout;
+    private final Duration visitTimeout;
     private final Metric metric;
     private final DocumentApiMetrics metrics;
     private final DocumentOperationParser parser;
@@ -204,11 +205,11 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
              documentManagerConfig, executorConfig, clusterListConfig, bucketSpacesConfig);
     }
 
-    DocumentV1ApiHandler(Clock clock, Duration handlerTimeout, Metric metric, MetricReceiver metricReceiver, DocumentAccess access,
+    DocumentV1ApiHandler(Clock clock, Duration visitTimeout, Metric metric, MetricReceiver metricReceiver, DocumentAccess access,
                          DocumentmanagerConfig documentmanagerConfig, DocumentOperationExecutorConfig executorConfig,
                          ClusterListConfig clusterListConfig, AllClustersBucketSpacesConfig bucketSpacesConfig) {
         this.clock = clock;
-        this.handlerTimeout = handlerTimeout;
+        this.visitTimeout = visitTimeout;
         this.parser = new DocumentOperationParser(documentmanagerConfig);
         this.metric = metric;
         this.metrics = new DocumentApiMetrics(metricReceiver, "documentV1");
@@ -237,9 +238,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(), MILLISECONDS);
 
             Path requestPath = Path.withoutValidation(request.getUri()); // No segment validation here, as document IDs can be anything.
             for (String path : handlers.keySet()) {
@@ -267,7 +266,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 +523,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)).minus(handlerTimeout));
         for (String name : names)
             parameters = switch (name) {
                 case CLUSTER ->
@@ -1168,16 +1166,15 @@ 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());
         if (streamed) {
             StaticThrottlePolicy throttlePolicy = new DynamicThrottlePolicy().setMinWindowSize(1).setWindowSizeIncrement(1);
             concurrency.ifPresent(throttlePolicy::setMaxPendingCount);
             parameters.setThrottlePolicy(throttlePolicy);
-            parameters.setTimeoutMs(timeoutMs); // Ensure visitor eventually completes.
+            parameters.setTimeoutMs(visitTimeout(request)); // Ensure visitor eventually completes.
         }
         else {
             parameters.setThrottlePolicy(new StaticThrottlePolicy().setMaxPendingCount(Math.min(100, concurrency.orElse(1))));
-            parameters.setSessionTimeoutMs(timeoutMs);
+            parameters.setSessionTimeoutMs(visitTimeout(request));
         }
         return parameters;
     }
@@ -1188,10 +1185,16 @@ 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.min(timeChunk, visitTimeout(request)));
         return parameters;
     }
 
+    private long visitTimeout(HttpRequest request) {
+        return Math.max(1,
+                        Math.max(doomMillis(request) - clock.millis() - visitTimeout.toMillis(),
+                                 9 * (doomMillis(request) - clock.millis()) / 10 - handlerTimeout.toMillis()));
+    }
+
     private VisitorParameters parseCommonParameters(HttpRequest request, DocumentPath path, Optional<String> cluster) {
         VisitorParameters parameters = new VisitorParameters(Stream.of(getProperty(request, SELECTION),
                                                                        path.documentType(),
@@ -1345,7 +1348,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, visitTimeout(request), MILLISECONDS) : null;
                 final AtomicReference<VisitorSession> session = new AtomicReference<>();
                 @Override public void setSession(VisitorControlSession session) { // Workaround for broken session API ಠ_ಠ
                     super.setSession(session);
@@ -1426,6 +1429,12 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
 
     // ------------------------------------------------ Helpers ------------------------------------------------
 
+    private static long doomMillis(HttpRequest request) {
+        long createdAtMillis = request.creationTime(MILLISECONDS);
+        long requestTimeoutMillis = getProperty(request, TIMEOUT, timeoutMillisParser).orElse(defaultTimeout.toMillis());
+        return createdAtMillis + 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();
         });