JAVA-2020: Add support for maxAwaitTime, which for a tailable cursor specifies the maximum time for the getMore command to wait for more documents before replying

Author: jyemin

driver-async/src/main/com/mongodb/async/client/FindIterable.java

@@ -65,6 +65,26 @@ public interface FindIterable<T> extends MongoIterable<T> {
      */
     FindIterable<T> maxTime(long maxTime, TimeUnit timeUnit);
 
+    /**
+     * The maximum amount of time for the server to wait on new documents to satisfy a tailable cursor
+     * query. This only applies to a TAILABLE_AWAIT cursor. When the cursor is not a TAILABLE_AWAIT cursor,
+     * this option is ignored.
+     *
+     * On servers >= 3.2, this option will be specified on the getMore command as "maxTimeMS". The default
+     * is no value: no "maxTimeMS" is sent to the server with the getMore command.
+     *
+     * On servers < 3.2, this option is ignored, and indicates that the driver should respect the server's default value
+     *
+     * A zero value will be ignored.
+     *
+     * @param maxAwaitTime  the max await time
+     * @param timeUnit the time unit to return the result in
+     * @return the maximum await execution time in the given time unit
+     * @mongodb.driver.manual reference/method/cursor.maxTimeMS/#cursor.maxTimeMS Max Time
+     * @since 3.2
+     */
+    FindIterable<T> maxAwaitTime(long maxAwaitTime, TimeUnit timeUnit);
+
     /**
      * Sets the query modifiers to apply to this operation.
      *

driver-async/src/main/com/mongodb/async/client/FindIterableImpl.java

@@ -87,6 +87,13 @@ public FindIterable<TResult> maxTime(final long maxTime, final TimeUnit timeUnit
         return this;
     }
 
+    @Override
+    public FindIterable<TResult> maxAwaitTime(final long maxAwaitTime, final TimeUnit timeUnit) {
+        notNull("timeUnit", timeUnit);
+        findOptions.maxAwaitTime(maxAwaitTime, timeUnit);
+        return this;
+    }
+
     @Override
     public FindIterable<TResult> batchSize(final int batchSize) {
         findOptions.batchSize(batchSize);
@@ -175,6 +182,7 @@ private FindOperation<TResult> createQueryOperation() {
                .skip(findOptions.getSkip())
                .limit(findOptions.getLimit())
                .maxTime(findOptions.getMaxTime(MILLISECONDS), MILLISECONDS)
+               .maxAwaitTime(findOptions.getMaxAwaitTime(MILLISECONDS), MILLISECONDS)
                .modifiers(toBsonDocument(findOptions.getModifiers()))
                .projection(toBsonDocument(findOptions.getProjection()))
                .sort(toBsonDocument(findOptions.getSort()))

driver-async/src/test/unit/com/mongodb/async/client/FindIterableSpecification.groovy

@@ -37,6 +37,7 @@ import spock.lang.Specification
 import static com.mongodb.CustomMatchers.isTheSameAs
 import static com.mongodb.ReadPreference.secondary
 import static java.util.concurrent.TimeUnit.MILLISECONDS
+import static java.util.concurrent.TimeUnit.SECONDS
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders
 import static spock.util.matcher.HamcrestSupport.expect
 
@@ -58,7 +59,8 @@ class FindIterableSpecification extends Specification {
         def findOptions = new FindOptions().sort(new Document('sort', 1))
                 .modifiers(new Document('modifier', 1))
                 .projection(new Document('projection', 1))
-                .maxTime(1000, MILLISECONDS)
+                .maxTime(10, SECONDS)
+                .maxAwaitTime(20, SECONDS)
                 .batchSize(100)
                 .limit(100)
                 .skip(10)
@@ -81,7 +83,8 @@ class FindIterableSpecification extends Specification {
                 .sort(new BsonDocument('sort', new BsonInt32(1)))
                 .modifiers(new BsonDocument('modifier', new BsonInt32(1)))
                 .projection(new BsonDocument('projection', new BsonInt32(1)))
-                .maxTime(1000, MILLISECONDS)
+                .maxTime(10000, MILLISECONDS)
+                .maxAwaitTime(20000, MILLISECONDS)
                 .batchSize(100)
                 .limit(100)
                 .skip(10)
@@ -95,7 +98,8 @@ class FindIterableSpecification extends Specification {
                 .sort(new Document('sort', 2))
                 .modifiers(new Document('modifier', 2))
                 .projection(new Document('projection', 2))
-                .maxTime(999, MILLISECONDS)
+                .maxTime(9, SECONDS)
+                .maxAwaitTime(18, SECONDS)
                 .batchSize(99)
                 .limit(99)
                 .skip(9)
@@ -113,7 +117,8 @@ class FindIterableSpecification extends Specification {
                 .sort(new BsonDocument('sort', new BsonInt32(2)))
                 .modifiers(new BsonDocument('modifier', new BsonInt32(2)))
                 .projection(new BsonDocument('projection', new BsonInt32(2)))
-                .maxTime(999, MILLISECONDS)
+                .maxTime(9000, MILLISECONDS)
+                .maxAwaitTime(18000, MILLISECONDS)
                 .batchSize(99)
                 .limit(99)
                 .skip(9)

driver-core/src/main/com/mongodb/client/model/FindOptions.java

@@ -21,6 +21,7 @@
 
 import java.util.concurrent.TimeUnit;
 
+import static com.mongodb.assertions.Assertions.isTrueArgument;
 import static com.mongodb.assertions.Assertions.notNull;
 
 /**
@@ -36,6 +37,7 @@ public final class FindOptions {
     private Bson modifiers;
     private Bson projection;
     private long maxTimeMS;
+    private long maxAwaitTimeMS;
     private int skip;
     private Bson sort;
     private CursorType cursorType = CursorType.NonTailable;
@@ -59,6 +61,7 @@ public FindOptions(final FindOptions from) {
         modifiers = from.modifiers;
         projection = from.projection;
         maxTimeMS = from.maxTimeMS;
+        maxAwaitTimeMS = from.maxAwaitTimeMS;
         skip = from.skip;
         sort = from.sort;
         cursorType = from.cursorType;
@@ -133,10 +136,50 @@ public long getMaxTime(final TimeUnit timeUnit) {
      */
     public FindOptions maxTime(final long maxTime, final TimeUnit timeUnit) {
         notNull("timeUnit", timeUnit);
+        isTrueArgument("maxTime > = 0", maxTime >= 0);
         this.maxTimeMS = TimeUnit.MILLISECONDS.convert(maxTime, timeUnit);
         return this;
     }
 
+    /**
+     * The maximum amount of time for the server to wait on new documents to satisfy a tailable cursor
+     * query. This only applies to a TAILABLE_AWAIT cursor. When the cursor is not a TAILABLE_AWAIT cursor,
+     * this option is ignored.
+     *
+     * On servers >= 3.2, this option will be specified on the getMore command as "maxTimeMS". The default
+     * is no value: no "maxTimeMS" is sent to the server with the getMore command.
+     *
+     * On servers < 3.2, this option is ignored, and indicates that the driver should respect the server's default value
+     *
+     * A zero value will be ignored.
+     *
+     * @param timeUnit the time unit to return the result in
+     * @return the maximum await execution time in the given time unit
+     * @since 3.2
+     * @mongodb.driver.manual reference/method/cursor.maxTimeMS/#cursor.maxTimeMS Max Time
+     */
+    public long getMaxAwaitTime(final TimeUnit timeUnit) {
+        notNull("timeUnit", timeUnit);
+        return timeUnit.convert(maxAwaitTimeMS, TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Sets the maximum await execution time on the server for this operation.
+     *
+     * @param maxAwaitTime  the max await time.  A zero value will be ignored, and indicates that the driver should respect the server's
+     *                      default value
+     * @param timeUnit the time unit, which may not be null
+     * @return this
+     * @since 3.2
+     * @mongodb.driver.manual reference/method/cursor.maxTimeMS/#cursor.maxTimeMS Max Time
+     */
+    public FindOptions maxAwaitTime(final long maxAwaitTime, final TimeUnit timeUnit) {
+        notNull("timeUnit", timeUnit);
+        isTrueArgument("maxAwaitTime > = 0", maxAwaitTime >= 0);
+        this.maxAwaitTimeMS = TimeUnit.MILLISECONDS.convert(maxAwaitTime, timeUnit);
+        return this;
+    }
+
     /**
      * Gets the number of documents to return per batch.  Default to 0, which indicates that the server chooses an appropriate batch
      * size.

driver-core/src/main/com/mongodb/operation/AggregateOperation.java

@@ -338,7 +338,7 @@ private CommandTransformer<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(f
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result, final ServerAddress serverAddress) {
                 QueryResult<T> queryResult = createQueryResult(result, connection.getDescription());
-                return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize != null ? batchSize : 0, decoder, source, connection);
+                return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize != null ? batchSize : 0, 0, decoder, source, connection);
             }
         };
     }

driver-core/src/main/com/mongodb/operation/AsyncQueryBatchCursor.java

@@ -34,6 +34,7 @@
 import java.util.List;
 
 import static com.mongodb.assertions.Assertions.isTrue;
+import static com.mongodb.assertions.Assertions.isTrueArgument;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
 import static com.mongodb.operation.CursorHelper.getNumberToReturn;
@@ -47,6 +48,7 @@ class AsyncQueryBatchCursor<T> implements AsyncBatchCursor<T> {
     private final MongoNamespace namespace;
     private final int limit;
     private final Decoder<T> decoder;
+    private final long maxTimeMS;
     private volatile AsyncConnectionSource connectionSource;
     private volatile QueryResult<T> firstBatch;
     private volatile int batchSize;
@@ -56,11 +58,13 @@ class AsyncQueryBatchCursor<T> implements AsyncBatchCursor<T> {
 
     AsyncQueryBatchCursor(final QueryResult<T> firstBatch, final int limit, final int batchSize,
                           final Decoder<T> decoder) {
-        this(firstBatch, limit, batchSize, decoder, null, null);
+        this(firstBatch, limit, batchSize, 0, decoder, null, null);
     }
 
-    AsyncQueryBatchCursor(final QueryResult<T> firstBatch, final int limit, final int batchSize,
+    AsyncQueryBatchCursor(final QueryResult<T> firstBatch, final int limit, final int batchSize, final long maxTimeMS,
                           final Decoder<T> decoder, final AsyncConnectionSource connectionSource, final AsyncConnection connection) {
+        isTrueArgument("maxTimeMS >= 0", maxTimeMS >= 0);
+        this.maxTimeMS = maxTimeMS;
         this.namespace = firstBatch.getNamespace();
         this.firstBatch = firstBatch;
         this.limit = limit;
@@ -166,7 +170,9 @@ private BsonDocument asGetMoreCommandDocument() {
         if (batchSizeForGetMoreCommand != 0) {
             document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand));
         }
-
+        if (maxTimeMS != 0) {
+            document.append("maxTimeMS", new BsonInt64(maxTimeMS));
+        }
         return document;
     }
 

driver-core/src/main/com/mongodb/operation/FindOperation.java

@@ -53,6 +53,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.ReadPreference.primary;
+import static com.mongodb.assertions.Assertions.isTrueArgument;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.connection.ServerType.SHARD_ROUTER;
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
@@ -84,6 +85,7 @@ public class FindOperation<T> implements AsyncReadOperation<AsyncBatchCursor<T>>
     private BsonDocument modifiers;
     private BsonDocument projection;
     private long maxTimeMS;
+    private long maxAwaitTimeMS;
     private int skip;
     private BsonDocument sort;
     private CursorType cursorType = CursorType.NonTailable;
@@ -252,10 +254,49 @@ public long getMaxTime(final TimeUnit timeUnit) {
      */
     public FindOperation<T> maxTime(final long maxTime, final TimeUnit timeUnit) {
         notNull("timeUnit", timeUnit);
+        isTrueArgument("maxTime >= 0", maxTime >= 0);
         this.maxTimeMS = TimeUnit.MILLISECONDS.convert(maxTime, timeUnit);
         return this;
     }
 
+    /**
+     * The maximum amount of time for the server to wait on new documents to satisfy a tailable cursor
+     * query. This only applies to a TAILABLE_AWAIT cursor. When the cursor is not a TAILABLE_AWAIT cursor,
+     * this option is ignored.
+     *
+     * On servers >= 3.2, this option will be specified on the getMore command as "maxTimeMS". The default
+     * is no value: no "maxTimeMS" is sent to the server with the getMore command.
+     *
+     * On servers < 3.2, this option is ignored, and indicates that the driver should respect the server's default value
+     *
+     * A zero value will be ignored.
+     *
+     * @param timeUnit the time unit to return the result in
+     * @return the maximum await execution time in the given time unit
+     * @since 3.2
+     * @mongodb.driver.manual reference/method/cursor.maxTimeMS/#cursor.maxTimeMS Max Time
+     */
+    public long getMaxAwaitTime(final TimeUnit timeUnit) {
+        notNull("timeUnit", timeUnit);
+        return timeUnit.convert(maxAwaitTimeMS, TimeUnit.MILLISECONDS);
+    }
+
+    /**
+     * Sets the maximum await execution time on the server for this operation.
+     *
+     * @param maxAwaitTime  the max await time.  A zero value will be ignored, and indicates that the driver should respect the server's
+     *                      default value
+     * @param timeUnit the time unit, which may not be null
+     * @return this
+     * @since 3.2
+     * @mongodb.driver.manual reference/method/cursor.maxTimeMS/#cursor.maxTimeMS Max Time
+     */
+    public FindOperation<T> maxAwaitTime(final long maxAwaitTime, final TimeUnit timeUnit) {
+        notNull("timeUnit", timeUnit);
+        isTrueArgument("maxAwaitTime >= 0", maxAwaitTime >= 0);
+        this.maxAwaitTimeMS = TimeUnit.MILLISECONDS.convert(maxAwaitTime, timeUnit);
+        return this;
+    }
     /**
      * Gets the number of documents to skip.  The default is 0.
      *
@@ -460,7 +501,7 @@ public BatchCursor<T> call(final ConnectionSource source, final Connection conne
                                                                   isPartial(),
                                                                   isOplogReplay(),
                                                                   decoder);
-                    return new QueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source, connection);
+                    return new QueryBatchCursor<T>(queryResult, limit, batchSize, getMaxTimeForCursor(), decoder, source, connection);
                 }
             }
         });
@@ -499,6 +540,7 @@ public void onResult(final QueryResult<T> result, final Throwable t) {
                                                         wrappedCallback.onResult(null, t);
                                                     } else {
                                                         wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, limit, batchSize,
+                                                                        getMaxTimeForCursor(),
                                                                         decoder, source, connection), null);
                                                     }
                                                 }
@@ -771,19 +813,23 @@ private CommandTransformer<BsonDocument, BatchCursor<T>> transformer(final Conne
             public BatchCursor<T> apply(final BsonDocument result, final ServerAddress serverAddress) {
                 QueryResult<T> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),
                                                                          connection.getDescription().getServerAddress());
-                return new QueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source);
+                return new QueryBatchCursor<T>(queryResult, limit, batchSize, getMaxTimeForCursor(), decoder, source, connection);
             }
         };
     }
 
+    private long getMaxTimeForCursor() {
+        return cursorType == CursorType.TailableAwait ? maxAwaitTimeMS : 0;
+    }
+
     private CommandTransformer<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source,
                                                                          final AsyncConnection connection) {
         return new CommandTransformer<BsonDocument, AsyncBatchCursor<T>>() {
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result, final ServerAddress serverAddress) {
                 QueryResult<T> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),
                                                                          connection.getDescription().getServerAddress());
-                return new AsyncQueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source, connection);
+                return new AsyncQueryBatchCursor<T>(queryResult, limit, batchSize, getMaxTimeForCursor(), decoder, source, connection);
             }
         };
     }

driver-core/src/main/com/mongodb/operation/ListCollectionsOperation.java

@@ -224,7 +224,7 @@ public void onResult(final QueryResult<BsonDocument> result, final Throwable t)
                                         } else {
                                             wrappedCallback.onResult(new ProjectingAsyncBatchCursor(
                                                     new AsyncQueryBatchCursor<BsonDocument>(result, 0,
-                                                            batchSize, new BsonDocumentCodec(), source, connection)
+                                                            batchSize, 0, new BsonDocumentCodec(), source, connection)
                                             ), null);
                                         }
                                     }

driver-core/src/main/com/mongodb/operation/ListDatabasesOperation.java

@@ -137,7 +137,7 @@ private CommandTransformer<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(f
         return new CommandTransformer<BsonDocument, AsyncBatchCursor<T>>() {
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result, final ServerAddress serverAddress) {
-                return new AsyncQueryBatchCursor<T>(createQueryResult(result, connection.getDescription()), 0, 0, decoder, source,
+                return new AsyncQueryBatchCursor<T>(createQueryResult(result, connection.getDescription()), 0, 0, 0, decoder, source,
                                                     connection);
             }
         };

driver-core/src/main/com/mongodb/operation/ListIndexesOperation.java

@@ -191,7 +191,7 @@ public void onResult(final QueryResult<T> result, final Throwable t) {
                                         if (t != null) {
                                             wrappedCallback.onResult(null, t);
                                         } else {
-                                            wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, 0, batchSize, decoder, source,
+                                            wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, 0, batchSize, 0, decoder, source,
                                                                                                   connection),
                                                     null);
                                         }