Allow user specified timeout for SyncSession

3 files changed, 104 insertions, 7 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/SyncParameters.java b/documentapi/src/main/java/com/yahoo/documentapi/SyncParameters.java
index 5a4b4abbf1e..a7bbf1243ae 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/SyncParameters.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/SyncParameters.java
@@ -1,11 +1,42 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.documentapi;
 
+import java.time.temporal.TemporalAmount;
+import java.util.Optional;
+
 /**
  * Parameters for creating a synchronous session
  *
+ * @author bjorncs
  * @author <a href="mailto:simon@yahoo-inc.com">Simon Thoresen</a>
  */
 public class SyncParameters extends Parameters {
-    // empty
+    private final TemporalAmount defaultTimeout;
+
+    public SyncParameters() {
+        this(null);
+    }
+
+    private SyncParameters(TemporalAmount defaultTimeout) {
+        this.defaultTimeout = defaultTimeout;
+    }
+
+    public Optional<TemporalAmount> defaultTimeout() {
+        return Optional.ofNullable(defaultTimeout);
+    }
+
+    public static class Builder {
+        private TemporalAmount defaultTimeout;
+
+        /**
+         * Set default timeout for all messagebus operations.
+         */
+        public void setDefaultTimeout(TemporalAmount defaultTimeout) {
+            this.defaultTimeout = defaultTimeout;
+        }
+
+        public SyncParameters build() {
+            return new SyncParameters(defaultTimeout);
+        }
+    }
 }
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/SyncSession.java b/documentapi/src/main/java/com/yahoo/documentapi/SyncSession.java
index 2a985212449..e04fb7a4e07 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/SyncSession.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/SyncSession.java
@@ -5,16 +5,17 @@ import com.yahoo.document.Document;
 import com.yahoo.document.DocumentId;
 import com.yahoo.document.DocumentPut;
 import com.yahoo.document.DocumentRemove;
-import com.yahoo.document.DocumentType;
 import com.yahoo.document.DocumentUpdate;
-import com.yahoo.document.TestAndSetCondition;
 import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
 
+import java.time.temporal.TemporalAmount;
+
 /**
  * <p>A session for synchronous access to a document repository. This class
  * provides simple document access where throughput is not a concern.</p>
  *
  * @author <a href="mailto:simon@yahoo-inc.com">Simon Thoresen</a>
+ * @author bjorncs
  */
 public interface SyncSession extends Session {
 
@@ -60,6 +61,41 @@ public interface SyncSession extends Session {
     Document get(DocumentId id, String fieldSet, DocumentProtocol.Priority priority);
 
     /**
+     * <p>Gets a document with timeout.</p>
+     *
+     * @param id The id of the document to get.
+     * @param timeout Timeout. If timeout is null, an unspecified default will be used.
+     * @return The known document having this id, or null if there is no
+     *         document having this id.
+     * @throws UnsupportedOperationException Thrown if this access does not
+     *                                       support retrieving.
+     * @throws RuntimeException If operation times out.
+     */
+    // TODO Vespa 7: Remove default implementation. Consider removing get() overloads without timeout.
+    default Document get(DocumentId id, TemporalAmount timeout) {
+        return get(id);
+    }
+
+    /**
+     * <p>Gets a document with timeout. </p>
+     *
+     * @param id       The id of the document to get.
+     * @param fieldSet A comma-separated list of fields to retrieve
+     * @param priority The priority with which to perform this operation.
+     * @param timeout Timeout. If timeout is null, an unspecified default will be used.
+     * @return The known document having this id, or null if there is no
+     *         document having this id.
+     * @throws UnsupportedOperationException Thrown if this access does not
+     *                                       support retrieving.
+     * * @throws RuntimeException If operation times out.
+     */
+    // TODO Vespa 7: Remove default implementation. Consider removing get() overloads without timeout.
+    default Document get(DocumentId id, String fieldSet, DocumentProtocol.Priority priority,
+                         TemporalAmount timeout) {
+        return get(id, fieldSet, priority);
+    }
+
+    /**
      * <p>Removes a document if it is present and condition is fulfilled.</p>
      * @param documentRemove document to delete
      * @return true If the document with this id was removed, false otherwise.
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusSyncSession.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusSyncSession.java
index 2d0f1d9db82..f2b1816a410 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusSyncSession.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/MessageBusSyncSession.java
@@ -12,21 +12,32 @@ import com.yahoo.documentapi.Response;
 import com.yahoo.documentapi.Result;
 import com.yahoo.documentapi.SyncParameters;
 import com.yahoo.documentapi.SyncSession;
-import com.yahoo.documentapi.messagebus.protocol.*;
 import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
+import com.yahoo.documentapi.messagebus.protocol.GetDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.GetDocumentReply;
+import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.RemoveDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.RemoveDocumentReply;
+import com.yahoo.documentapi.messagebus.protocol.UpdateDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.UpdateDocumentReply;
 import com.yahoo.messagebus.Message;
 import com.yahoo.messagebus.MessageBus;
 import com.yahoo.messagebus.Reply;
 import com.yahoo.messagebus.ReplyHandler;
 
+import java.time.temporal.ChronoUnit;
+import java.time.temporal.TemporalAmount;
+
 /**
  * An implementation of the SyncSession interface running over message bus.
  *
  * @author <a href="mailto:simon@yahoo-inc.com">Simon Thoresen</a>
+ * @author bjorncs
  */
 public class MessageBusSyncSession implements MessageBusSession, SyncSession, ReplyHandler {
 
-    private MessageBusAsyncSession session;
+    private final MessageBusAsyncSession session;
+    private final TemporalAmount defaultTimeout;
 
     /**
      * Creates a new sync session running on message bus logic.
@@ -37,6 +48,7 @@ public class MessageBusSyncSession implements MessageBusSession, SyncSession, Re
      */
     MessageBusSyncSession(SyncParameters syncParams, MessageBus bus, MessageBusParams mbusParams) {
         session = new MessageBusAsyncSession(new AsyncParameters(), bus, mbusParams, this);
+        defaultTimeout = syncParams.defaultTimeout().orElse(null);
     }
 
     @Override
@@ -72,6 +84,13 @@ public class MessageBusSyncSession implements MessageBusSession, SyncSession, Re
      * @return The reply received.
      */
     public Reply syncSend(Message msg) {
+        return syncSend(msg, defaultTimeout);
+    }
+
+    private Reply syncSend(Message msg, TemporalAmount timeout) {
+        if (timeout != null) {
+            msg.setTimeRemaining(timeout.get(ChronoUnit.MILLIS));
+        }
         try {
             RequestMonitor monitor = new RequestMonitor();
             msg.setContext(monitor);
@@ -107,15 +126,26 @@ public class MessageBusSyncSession implements MessageBusSession, SyncSession, Re
 
     @Override
     public Document get(DocumentId id) {
-        return get(id, "[all]", DocumentProtocol.Priority.NORMAL_1);
+        return get(id, null);
     }
 
     @Override
     public Document get(DocumentId id, String fieldSet, DocumentProtocol.Priority pri) {
+        return get(id, fieldSet, pri, null);
+    }
+
+    @Override
+    public Document get(DocumentId id, TemporalAmount timeout) {
+        return get(id, "[all]", DocumentProtocol.Priority.NORMAL_1, timeout);
+    }
+
+    @Override
+    public Document get(DocumentId id, String fieldSet, DocumentProtocol.Priority pri,
+                        TemporalAmount timeout) {
         GetDocumentMessage msg = new GetDocumentMessage(id, fieldSet);
         msg.setPriority(pri);
 
-        Reply reply = syncSend(msg);
+        Reply reply = syncSend(msg, timeout != null ? timeout : defaultTimeout);
         if (reply.hasErrors()) {
             throw new DocumentAccessException(MessageBusAsyncSession.getErrorMessage(reply));
         }