Set message bus message timeout from http request, and 504 on mbus timeout

2 files changed, 21 insertions, 0 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 1081f654d04..8cb927b8989 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
@@ -479,6 +479,9 @@ 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);
         for (String name : names) switch (name) {
             case CLUSTER:
                 parameters = getProperty(request, CLUSTER).map(cluster -> resolveCluster(Optional.of(cluster), clusters).name())
@@ -904,6 +907,9 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                     case INSUFFICIENT_STORAGE:
                         jsonResponse.commit(Response.Status.INSUFFICIENT_STORAGE);
                         break;
+                    case TIMEOUT:
+                        jsonResponse.commit(Response.Status.GATEWAY_TIMEOUT);
+                        break;
                     default:
                         log.log(WARNING, "Unexpected document API operation outcome '" + response.outcome() + "'");
                     case ERROR:
@@ -1022,6 +1028,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                                     break; // This is all OK — the latter two are due to mitigating races.
                                 case ERROR:
                                 case INSUFFICIENT_STORAGE:
+                                case TIMEOUT:
                                     onError.accept(operationResponse.getTextMessage());
                                     break;
                                 default:
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 c4777b26072..6f1b0466350 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
@@ -577,6 +577,20 @@ public class DocumentV1ApiTest {
                        "}", response.readAll());
         assertEquals(400, response.getStatus());
 
+        // TIMEOUT is a 504
+        access.session.expect((id, parameters) -> {
+            assertEquals(clock.instant().plusSeconds(1000), parameters.deadline().get());
+            parameters.responseHandler().get().handleResponse(new Response(0, "timeout", Response.Outcome.TIMEOUT));
+            return new Result(Result.ResultType.SUCCESS, null);
+        });
+        response = driver.sendRequest("http://localhost/document/v1/space/music/number/1/two?timeout=1ks");
+        assertSameJson("{" +
+                       "  \"pathId\": \"/document/v1/space/music/number/1/two\"," +
+                       "  \"id\": \"id:space:music:n=1:two\"," +
+                       "  \"message\": \"timeout\"" +
+                       "}", response.readAll());
+        assertEquals(504, response.getStatus());
+
         // INSUFFICIENT_STORAGE is a 507
         access.session.expect((id, parameters) -> {
             parameters.responseHandler().get().handleResponse(new Response(0, "disk full", Response.Outcome.INSUFFICIENT_STORAGE));