Ensure connectionSource is released in QueryBatchCursor

Also ensure killCursors does not throw an exception.

JAVA-3774

Author: rozza

driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java

@@ -112,16 +112,14 @@ class QueryBatchCursor<T> implements AggregateResponseBatchCursor<T> {
 
         initFromQueryResult(firstQueryResult);
         firstBatchEmpty = firstQueryResult.getResults().isEmpty();
-        if (limitReached()) {
-            killCursor(connection);
-        }
-        if (serverCursor == null && this.connectionSource != null) {
-            this.connectionSource.release();
-            this.connectionSource = null;
-        }
+
         if (connection != null) {
             this.maxWireVersion = connection.getDescription().getMaxWireVersion();
+            if (limitReached()) {
+                killCursor(connection);
+            }
         }
+        releaseConnectionSourceIfNoServerCursor();
     }
 
     @Override
@@ -283,12 +281,9 @@ private void getMore() {
             if (limitReached()) {
                 killCursor(connection);
             }
-            if (serverCursor == null) {
-                this.connectionSource.release();
-                this.connectionSource = null;
-            }
         } finally {
             connection.release();
+            releaseConnectionSourceIfNoServerCursor();
         }
     }
 
@@ -343,13 +338,23 @@ private void killCursor() {
     private void killCursor(final Connection connection) {
         if (serverCursor != null) {
             notNull("connection", connection);
-            if (serverIsAtLeastVersionThreeDotTwo(connection.getDescription())) {
-                connection.command(namespace.getDatabaseName(), asKillCursorsCommandDocument(), NO_OP_FIELD_NAME_VALIDATOR,
-                        ReadPreference.primary(), new BsonDocumentCodec(), connectionSource.getSessionContext());
-            } else {
-                connection.killCursor(namespace, singletonList(serverCursor.getId()));
+            try {
+                if (serverIsAtLeastVersionThreeDotTwo(connection.getDescription())) {
+                    connection.command(namespace.getDatabaseName(), asKillCursorsCommandDocument(), NO_OP_FIELD_NAME_VALIDATOR,
+                            ReadPreference.primary(), new BsonDocumentCodec(), connectionSource.getSessionContext());
+                } else {
+                    connection.killCursor(namespace, singletonList(serverCursor.getId()));
+                }
+            } finally {
+                serverCursor = null;
             }
-            serverCursor = null;
+        }
+    }
+
+    private void releaseConnectionSourceIfNoServerCursor() {
+        if (serverCursor == null && connectionSource != null) {
+            connectionSource.release();
+            connectionSource = null;
         }
     }