JAVA-1994: Use already-retained AsyncConnection to asynchronously kill cursor, instead of checking out a new one while the existing one has not yet been released.

jyemin · jyemin · commit 8a3c72d8fc17 · 2015-10-02T08:10:23.000-04:00
This avoids unnecessary pressure on the connection pool.
diff --git a/driver-core/src/main/com/mongodb/operation/AggregateOperation.java b/driver-core/src/main/com/mongodb/operation/AggregateOperation.java
@@ -298,7 +298,7 @@ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final Async
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result) {
                 QueryResult<T> queryResult = createQueryResult(result, connection.getDescription());
-                return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize != null ? batchSize : 0, decoder, source);
+                return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize != null ? batchSize : 0, decoder, source, connection);
             }
         };
     }
diff --git a/driver-core/src/main/com/mongodb/operation/AsyncQueryBatchCursor.java b/driver-core/src/main/com/mongodb/operation/AsyncQueryBatchCursor.java
@@ -47,11 +47,11 @@ 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);
+        this(firstBatch, limit, batchSize, decoder, null, null);
     }
 
     AsyncQueryBatchCursor(final QueryResult<T> firstBatch, final int limit, final int batchSize,
-                          final Decoder<T> decoder, final AsyncConnectionSource connectionSource) {
+                          final Decoder<T> decoder, final AsyncConnectionSource connectionSource, final AsyncConnection connection) {
         this.namespace = firstBatch.getNamespace();
         this.firstBatch = firstBatch;
         this.limit = limit;
@@ -60,6 +60,7 @@ class AsyncQueryBatchCursor<T> implements AsyncBatchCursor<T> {
         this.cursor = firstBatch.getCursor();
         if (this.cursor != null) {
             notNull("connectionSource", connectionSource);
+            notNull("connection", connection);
         }
         if (connectionSource != null) {
             this.connectionSource = connectionSource.retain();
@@ -68,15 +69,15 @@ class AsyncQueryBatchCursor<T> implements AsyncBatchCursor<T> {
         }
         this.count += firstBatch.getResults().size();
         if (limitReached()) {
-            killCursor();
+             killCursor(connection);
         }
     }
 
     @Override
     public void close() {
         if (!closed) {
             closed = true;
-            killCursor();
+            killCursor(null);
         }
     }
 
@@ -133,32 +134,42 @@ public void onResult(final AsyncConnection connection, final Throwable t) {
         });
     }
 
-    private void killCursor() {
+    private void killCursor(final AsyncConnection connection) {
         if (cursor != null) {
             final ServerCursor localCursor = cursor;
             final AsyncConnectionSource localConnectionSource = connectionSource;
             cursor = null;
             connectionSource = null;
-            localConnectionSource.getConnection(new SingleResultCallback<AsyncConnection>() {
-                @Override
-                public void onResult(final AsyncConnection connection, final Throwable connectionException) {
-                    if (connection != null) {
-                        connection.killCursorAsync(namespace, singletonList(localCursor.getId()), new SingleResultCallback<Void>() {
-                            @Override
-                            public void onResult(final Void result, final Throwable t) {
-                                connection.release();
-                                localConnectionSource.release();
-                            }
-                        });
+            if (connection != null) {
+                connection.retain();
+                killCursorAsynchronouslyAndReleaseConnectionAndSource(connection, localCursor, localConnectionSource);
+            } else {
+                localConnectionSource.getConnection(new SingleResultCallback<AsyncConnection>() {
+                    @Override
+                    public void onResult(final AsyncConnection connection, final Throwable connectionException) {
+                        if (connectionException == null) {
+                            killCursorAsynchronouslyAndReleaseConnectionAndSource(connection, localCursor, localConnectionSource);
+                        }
                     }
-                }
-            });
+                });
+            }
         } else if (connectionSource != null) {
             connectionSource.release();
             connectionSource = null;
         }
     }
 
+    private void killCursorAsynchronouslyAndReleaseConnectionAndSource(final AsyncConnection connection, final ServerCursor localCursor,
+                                                                       final AsyncConnectionSource localConnectionSource) {
+        connection.killCursorAsync(namespace, singletonList(localCursor.getId()), new SingleResultCallback<Void>() {
+            @Override
+            public void onResult(final Void result, final Throwable t) {
+                connection.release();
+                localConnectionSource.release();
+            }
+        });
+    }
+
     private class QueryResultSingleResultCallback implements SingleResultCallback<QueryResult<T>> {
         private final AsyncConnection connection;
         private final SingleResultCallback<List<T>> callback;
@@ -181,7 +192,7 @@ public void onResult(final QueryResult<T> result, final Throwable t) {
                 cursor = result.getCursor();
                 count += result.getResults().size();
                 if (limitReached()) {
-                    killCursor();
+                    killCursor(connection);
                 }
                 connection.release();
                 if (result.getResults().isEmpty()) {
diff --git a/driver-core/src/main/com/mongodb/operation/DistinctOperation.java b/driver-core/src/main/com/mongodb/operation/DistinctOperation.java
@@ -141,7 +141,7 @@ public void call(final AsyncConnectionSource source, final AsyncConnection conne
                     errorHandlingCallback(callback).onResult(null, t);
                 } else {
                     executeWrappedCommandProtocolAsync(binding, namespace.getDatabaseName(), getCommand(), createCommandDecoder(),
-                            connection, asyncTransformer(source, connection),
+                            connection, asyncTransformer(connection.getDescription()),
                             releasingCallback(errorHandlingCallback(callback), source, connection));
                 }
             }
@@ -167,13 +167,12 @@ public BatchCursor<T> apply(final BsonDocument result) {
         };
     }
 
-    private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source,
-                                                                         final AsyncConnection connection) {
+    private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final ConnectionDescription connectionDescription) {
         return new Function<BsonDocument, AsyncBatchCursor<T>>() {
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result) {
-                QueryResult<T> queryResult = createQueryResult(result, connection.getDescription());
-                return new AsyncQueryBatchCursor<T>(queryResult, 0, 0, null, source);
+                QueryResult<T> queryResult = createQueryResult(result, connectionDescription);
+                return new AsyncQueryBatchCursor<T>(queryResult, 0, 0, null);
             }
         };
     }
diff --git a/driver-core/src/main/com/mongodb/operation/FindOperation.java b/driver-core/src/main/com/mongodb/operation/FindOperation.java
@@ -431,7 +431,9 @@ public void onResult(final QueryResult<T> result, final Throwable t) {
                             if (t != null) {
                                 wrappedCallback.onResult(null, t);
                             } else {
-                                wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, limit, batchSize, decoder, source), null);
+                                wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, limit, batchSize, decoder, source,
+                                                                                      connection),
+                                                         null);
                             }
                         }
                     });
diff --git a/driver-core/src/main/com/mongodb/operation/FindOperationHelper.java b/driver-core/src/main/com/mongodb/operation/FindOperationHelper.java
diff --git a/driver-core/src/main/com/mongodb/operation/ListCollectionsOperation.java b/driver-core/src/main/com/mongodb/operation/ListCollectionsOperation.java
@@ -202,7 +202,7 @@ public void call(final AsyncConnectionSource source, final AsyncConnection conne
                                                                                                         source, connection);
                     if (serverIsAtLeastVersionThreeDotZero(connection.getDescription())) {
                         executeWrappedCommandProtocolAsync(binding, databaseName, getCommand(), createCommandDecoder(),
-                                connection, asyncTransformer(source),
+                                connection, asyncTransformer(source, connection),
                                 new SingleResultCallback<AsyncBatchCursor<T>>() {
                                     @Override
                                     public void onResult(final AsyncBatchCursor<T> result, final Throwable t) {
@@ -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)
+                                                            batchSize, new BsonDocumentCodec(), source, connection)
                                             ), null);
                                         }
                                     }
@@ -243,11 +243,12 @@ private MongoNamespace createNamespace() {
         return new MongoNamespace(databaseName, "$cmd.listCollections");
     }
 
-    private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source) {
+    private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source,
+                                                                         final AsyncConnection connection) {
         return new Function<BsonDocument, AsyncBatchCursor<T>>() {
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result) {
-                return cursorDocumentToAsyncBatchCursor(result.getDocument("cursor"), decoder, source, batchSize);
+                return cursorDocumentToAsyncBatchCursor(result.getDocument("cursor"), decoder, source, connection, batchSize);
             }
         };
     }
diff --git a/driver-core/src/main/com/mongodb/operation/ListDatabasesOperation.java b/driver-core/src/main/com/mongodb/operation/ListDatabasesOperation.java
@@ -136,7 +136,8 @@ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final Async
         return new Function<BsonDocument, AsyncBatchCursor<T>>() {
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result) {
-                return new AsyncQueryBatchCursor<T>(createQueryResult(result, connection.getDescription()), 0, 0, decoder, source);
+                return new AsyncQueryBatchCursor<T>(createQueryResult(result, connection.getDescription()), 0, 0, decoder, source,
+                                                    connection);
             }
         };
     }
diff --git a/driver-core/src/main/com/mongodb/operation/ListIndexesOperation.java b/driver-core/src/main/com/mongodb/operation/ListIndexesOperation.java
@@ -166,7 +166,7 @@ public void call(final AsyncConnectionSource source, final AsyncConnection conne
                                                                                                         source, connection);
                     if (serverIsAtLeastVersionThreeDotZero(connection.getDescription())) {
                         executeWrappedCommandProtocolAsync(binding, namespace.getDatabaseName(), getCommand(), createCommandDecoder(),
-                                connection, asyncTransformer(source),
+                                connection, asyncTransformer(source, connection),
                                 new SingleResultCallback<AsyncBatchCursor<T>>() {
                                     @Override
                                     public void onResult(final AsyncBatchCursor<T> result,
@@ -190,7 +190,8 @@ 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, decoder, source,
+                                                                                                  connection),
                                                     null);
                                         }
                                     }
@@ -238,11 +239,12 @@ public BatchCursor<T> apply(final BsonDocument result) {
         };
     }
 
-    private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source) {
+    private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source,
+                                                                         final AsyncConnection connection) {
         return new Function<BsonDocument, AsyncBatchCursor<T>>() {
             @Override
             public AsyncBatchCursor<T> apply(final BsonDocument result) {
-                return cursorDocumentToAsyncBatchCursor(result.getDocument("cursor"), decoder, source, batchSize);
+                return cursorDocumentToAsyncBatchCursor(result.getDocument("cursor"), decoder, source, connection, batchSize);
             }
         };
     }
diff --git a/driver-core/src/main/com/mongodb/operation/OperationHelper.java b/driver-core/src/main/com/mongodb/operation/OperationHelper.java
@@ -94,10 +94,11 @@ static <T> BatchCursor<T> cursorDocumentToBatchCursor(final BsonDocument cursorD
     }
 
     static <T> AsyncBatchCursor<T> cursorDocumentToAsyncBatchCursor(final BsonDocument cursorDocument, final Decoder<T> decoder,
-                                                                    final AsyncConnectionSource source, final int batchSize) {
+                                                                    final AsyncConnectionSource source, final AsyncConnection connection,
+                                                                    final int batchSize) {
         return new AsyncQueryBatchCursor<T>(OperationHelper.<T>cursorDocumentToQueryResult(cursorDocument,
                                                                                            source.getServerDescription().getAddress()),
-                                            0, batchSize, decoder, source);
+                                            0, batchSize, decoder, source, connection);
     }
 
 
diff --git a/driver-core/src/main/com/mongodb/operation/ParallelCollectionScanOperation.java b/driver-core/src/main/com/mongodb/operation/ParallelCollectionScanOperation.java
diff --git a/driver-core/src/test/functional/com/mongodb/connection/ServerHelper.java b/driver-core/src/test/functional/com/mongodb/connection/ServerHelper.java
diff --git a/driver-core/src/test/functional/com/mongodb/operation/AsyncQueryBatchCursorSpecification.groovy b/driver-core/src/test/functional/com/mongodb/operation/AsyncQueryBatchCursorSpecification.groovy

Original file line number	Diff line number	Diff line change
`@@ -298,7 +298,7 @@ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final Async`
`298`	`298`	`@Override`
`299`	`299`	`public AsyncBatchCursor<T> apply(final BsonDocument result) {`
`300`	`300`	`QueryResult<T> queryResult = createQueryResult(result, connection.getDescription());`
`301`		`- return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize != null ? batchSize : 0, decoder, source);`
	`301`	`+ return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize != null ? batchSize : 0, decoder, source, connection);`
`302`	`302`	`}`
`303`	`303`	`};`
`304`	`304`	`}`
Original file line number	Diff line number	Diff line change
`@@ -141,7 +141,7 @@ public void call(final AsyncConnectionSource source, final AsyncConnection conne`
`141`	`141`	`errorHandlingCallback(callback).onResult(null, t);`
`142`	`142`	`} else {`
`143`	`143`	`executeWrappedCommandProtocolAsync(binding, namespace.getDatabaseName(), getCommand(), createCommandDecoder(),`
`144`		`- connection, asyncTransformer(source, connection),`
	`144`	`+ connection, asyncTransformer(connection.getDescription()),`
`145`	`145`	`releasingCallback(errorHandlingCallback(callback), source, connection));`
`146`	`146`	`}`
`147`	`147`	`}`
`@@ -167,13 +167,12 @@ public BatchCursor<T> apply(final BsonDocument result) {`
`167`	`167`	`};`
`168`	`168`	`}`
`169`	`169`
`170`		`- private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final AsyncConnectionSource source,`
`171`		`- final AsyncConnection connection) {`
	`170`	`+ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final ConnectionDescription connectionDescription) {`
`172`	`171`	`return new Function<BsonDocument, AsyncBatchCursor<T>>() {`
`173`	`172`	`@Override`
`174`	`173`	`public AsyncBatchCursor<T> apply(final BsonDocument result) {`
`175`		`- QueryResult<T> queryResult = createQueryResult(result, connection.getDescription());`
`176`		`- return new AsyncQueryBatchCursor<T>(queryResult, 0, 0, null, source);`
	`174`	`+ QueryResult<T> queryResult = createQueryResult(result, connectionDescription);`
	`175`	`+ return new AsyncQueryBatchCursor<T>(queryResult, 0, 0, null);`
`177`	`176`	`}`
`178`	`177`	`};`
`179`	`178`	`}`
Original file line number	Diff line number	Diff line change
`@@ -431,7 +431,9 @@ public void onResult(final QueryResult<T> result, final Throwable t) {`
`431`	`431`	`if (t != null) {`
`432`	`432`	`wrappedCallback.onResult(null, t);`
`433`	`433`	`} else {`
`434`		`- wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, limit, batchSize, decoder, source), null);`
	`434`	`+ wrappedCallback.onResult(new AsyncQueryBatchCursor<T>(result, limit, batchSize, decoder, source,`
	`435`	`+ connection),`
	`436`	`+ null);`
`435`	`437`	`}`
`436`	`438`	`}`
`437`	`439`	`});`
Original file line number	Diff line number	Diff line change
`@@ -136,7 +136,8 @@ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final Async`
`136`	`136`	`return new Function<BsonDocument, AsyncBatchCursor<T>>() {`
`137`	`137`	`@Override`
`138`	`138`	`public AsyncBatchCursor<T> apply(final BsonDocument result) {`
`139`		`- return new AsyncQueryBatchCursor<T>(createQueryResult(result, connection.getDescription()), 0, 0, decoder, source);`
	`139`	`+ return new AsyncQueryBatchCursor<T>(createQueryResult(result, connection.getDescription()), 0, 0, decoder, source,`
	`140`	`+ connection);`
`140`	`141`	`}`
`141`	`142`	`};`
`142`	`143`	`}`
Original file line number	Diff line number	Diff line change
`@@ -94,10 +94,11 @@ static <T> BatchCursor<T> cursorDocumentToBatchCursor(final BsonDocument cursorD`
`94`	`94`	`}`
`95`	`95`
`96`	`96`	`static <T> AsyncBatchCursor<T> cursorDocumentToAsyncBatchCursor(final BsonDocument cursorDocument, final Decoder<T> decoder,`
`97`		`- final AsyncConnectionSource source, final int batchSize) {`
	`97`	`+ final AsyncConnectionSource source, final AsyncConnection connection,`
	`98`	`+ final int batchSize) {`
`98`	`99`	`return new AsyncQueryBatchCursor<T>(OperationHelper.<T>cursorDocumentToQueryResult(cursorDocument,`
`99`	`100`	`source.getServerDescription().getAddress()),`
`100`		`- 0, batchSize, decoder, source);`
	`101`	`+ 0, batchSize, decoder, source, connection);`
`101`	`102`	`}`
`102`	`103`
`103`	`104`