Use timeout from query when set

2 files changed, 12 insertions, 1 deletions

diff --git a/tenant-cd/src/main/java/ai/vespa/hosted/cd/Query.java b/tenant-cd/src/main/java/ai/vespa/hosted/cd/Query.java
index d421dc14322..4bdc63fca1a 100644
--- a/tenant-cd/src/main/java/ai/vespa/hosted/cd/Query.java
+++ b/tenant-cd/src/main/java/ai/vespa/hosted/cd/Query.java
@@ -1,6 +1,9 @@
 package ai.vespa.hosted.cd;
 
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Stream;
 
 import static java.util.Map.copyOf;
@@ -57,4 +60,11 @@ public class Query {
     /** Returns the parameters of this query. */
     public Map<String, String> parameters() { return parameters; }
 
+    /** Returns the timeout parameter of the request, if one is set. */
+    public Optional<Duration> timeout() {
+        return Optional.ofNullable(parameters.get("timeout"))
+                       .map(timeout -> Duration.of(Long.parseLong(timeout.replaceAll("\\s*m?s", "")),
+                                                   timeout.contains("ms") ? ChronoUnit.MILLIS : ChronoUnit.SECONDS));
+    }
+
 }
diff --git a/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java b/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java
index af66e20bbc7..4fafa65773d 100644
--- a/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java
+++ b/tenant-cd/src/main/java/ai/vespa/hosted/cd/http/HttpEndpoint.java
@@ -56,7 +56,8 @@ public class HttpEndpoint implements TestEndpoint {
         try {
             URI target = endpoint.resolve(searchApiPath).resolve("?" + query.rawQuery());
             HttpRequest request = HttpRequest.newBuilder()
-                                             .timeout(Duration.ofSeconds(5))
+                                             .timeout(query.timeout().orElse(Duration.ofMillis(500))
+                                                           .plus(Duration.ofSeconds(1)))
                                              .uri(target)
                                              .build();
             HttpResponse<byte[]> response = client.send(request, HttpResponse.BodyHandlers.ofByteArray());