CSHARP-1951: Use killCursors command.

Author: DmitryLukyanov

src/MongoDB.Driver.Core/Core/Misc/Feature.cs

@@ -53,6 +53,7 @@ public class Feature
         private static readonly Feature __findCommand = new Feature("FindCommand", new SemanticVersion(3, 2, 0));
         private static readonly Feature __geoNearCommand = new Feature("GeoNearCommand", new SemanticVersion(1, 0, 0), new SemanticVersion(4, 1, 0, ""));
         private static readonly Feature __groupCommand = new Feature("GroupCommand", new SemanticVersion(1, 0, 0), new SemanticVersion(4, 0, 0, "rc1"));
+        private static readonly Feature __killCursorsCommand = new Feature("KillCursorsCommand", new SemanticVersion(3, 2, 0));
         private static readonly Feature __listCollectionsCommand = new Feature("ListCollectionsCommand", new SemanticVersion(3, 0, 0));
         private static readonly Feature __listDatabasesFilter = new Feature("ListDatabasesFilter", new SemanticVersion(3, 4, 2));
         private static readonly Feature __listDatabasesNameOnlyOption = new Feature("ListDatabasesNameOnlyOption", new SemanticVersion(3, 4, 3));
@@ -221,6 +222,11 @@ public class Feature
         /// </summary>
         public static Feature GroupCommand => __groupCommand;
 
+        /// <summary>
+        /// Get the killCursors command feature.
+        /// </summary>
+        public static Feature KillCursorsCommand => __killCursorsCommand;
+
         /// <summary>
         /// Gets the index options defaults feature.
         /// </summary>

src/MongoDB.Driver.Core/Core/Operations/AsyncCursor.cs

@@ -23,6 +23,7 @@
 using MongoDB.Bson.Serialization;
 using MongoDB.Bson.Serialization.Serializers;
 using MongoDB.Driver.Core.Bindings;
+using MongoDB.Driver.Core.Connections;
 using MongoDB.Driver.Core.Events;
 using MongoDB.Driver.Core.Misc;
 using MongoDB.Driver.Core.WireProtocol;
@@ -48,6 +49,7 @@ public class AsyncCursor<TDocument> : IAsyncCursor<TDocument>
         private readonly int? _batchSize;
         private readonly CollectionNamespace _collectionNamespace;
         private IChannelSource _channelSource;
+        private bool _closed;
         private int _count;
         private IReadOnlyList<TDocument> _currentBatch;
         private long _cursorId;
@@ -133,6 +135,39 @@ private int CalculateGetMoreProtocolNumberToReturn()
             return numberToReturn;
         }
 
+        /// <summary>
+        /// Closes the cursor.
+        /// </summary>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        public void Close(CancellationToken cancellationToken = default(CancellationToken))
+        {
+            try
+            {
+                CloseIfNotAlreadyClosed(cancellationToken);
+            }
+            finally
+            {
+                Dispose();
+            }
+        }
+
+        /// <summary>
+        /// Closes the cursor.
+        /// </summary>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>A task.</returns>
+        public async Task CloseAsync(CancellationToken cancellationToken = default(CancellationToken))
+        {
+            try
+            {
+                await CloseIfNotAlreadyClosedAsync(cancellationToken).ConfigureAwait(false);
+            }
+            finally
+            {
+                Dispose();
+            }
+        }
+
         private CursorBatch<TDocument> CreateCursorBatch(BsonDocument result)
         {
             var cursorDocument = result["cursor"].AsBsonDocument;
@@ -159,6 +194,17 @@ private BsonDocument CreateGetMoreCommand()
             return command;
         }
 
+        private BsonDocument CreateKillCursorsCommand()
+        {
+            var command = new BsonDocument
+            {
+                { "killCursors", _collectionNamespace.CollectionName },
+                { "cursors", new BsonArray { _cursorId } }
+            };
+
+            return command;
+        }
+
         private CursorBatch<TDocument> ExecuteGetMoreCommand(IChannelHandle channel, CancellationToken cancellationToken)
         {
             var command = CreateGetMoreCommand();
@@ -227,6 +273,47 @@ private Task<CursorBatch<TDocument>> ExecuteGetMoreProtocolAsync(IChannelHandle
                 cancellationToken);
         }
 
+        private void ExecuteKillCursorsCommand(IChannelHandle channel, CancellationToken cancellationToken)
+        {
+            var command = CreateKillCursorsCommand();
+            var result = channel.Command(
+                _channelSource.Session,
+                null, // readPreference
+                _collectionNamespace.DatabaseNamespace,
+                command,
+                null, // commandPayloads
+                NoOpElementNameValidator.Instance,
+                null, // additionalOptions
+                null, // postWriteAction
+                CommandResponseHandling.Return,
+                BsonDocumentSerializer.Instance,
+                _messageEncoderSettings,
+                cancellationToken);
+
+            ThrowIfKillCursorsCommandFailed(result, channel.ConnectionDescription.ConnectionId);
+        }
+
+        private async Task ExecuteKillCursorsCommandAsync(IChannelHandle channel, CancellationToken cancellationToken)
+        {
+            var command = CreateKillCursorsCommand();
+            var result = await channel.CommandAsync(
+                _channelSource.Session,
+                null, // readPreference
+                _collectionNamespace.DatabaseNamespace,
+                command,
+                null, // commandPayloads
+                NoOpElementNameValidator.Instance,
+                null, // additionalOptions
+                null, // postWriteAction
+                CommandResponseHandling.Return,
+                BsonDocumentSerializer.Instance,
+                _messageEncoderSettings,
+                cancellationToken)
+                .ConfigureAwait(false);
+
+            ThrowIfKillCursorsCommandFailed(result, channel.ConnectionDescription.ConnectionId);
+        }
+
         private void ExecuteKillCursorsProtocol(IChannelHandle channel, CancellationToken cancellationToken)
         {
             channel.KillCursors(
@@ -235,6 +322,14 @@ private void ExecuteKillCursorsProtocol(IChannelHandle channel, CancellationToke
                 cancellationToken);
         }
 
+        private Task ExecuteKillCursorsProtocolAsync(IChannelHandle channel, CancellationToken cancellationToken)
+        {
+            return channel.KillCursorsAsync(
+                new[] { _cursorId },
+                _messageEncoderSettings,
+                cancellationToken);
+        }
+
         /// <inheritdoc/>
         public void Dispose()
         {
@@ -252,27 +347,66 @@ protected virtual void Dispose(bool disposing)
             {
                 if (!_disposed)
                 {
-                    try
+                    CloseIfNotAlreadyClosedFromDispose();
+                    
+                    if (_channelSource != null)
                     {
-                        if (_cursorId != 0)
-                        {
-                            using (var source = new CancellationTokenSource(TimeSpan.FromSeconds(10)))
-                            {
-                                KillCursor(_cursorId, source.Token);
-                            }
-                        }
+                        _channelSource.Dispose();
                     }
-                    catch
+                    _disposed = true;
+                }
+            }
+        }
+
+        private void CloseIfNotAlreadyClosed(CancellationToken cancellationToken)
+        {
+            if (!_closed)
+            {
+                try
+                {
+                    if (_cursorId != 0)
                     {
-                        // ignore exceptions
+                        KillCursors(cancellationToken);
                     }
-                    if (_channelSource != null)
+                }
+                finally
+                {
+                    _closed = true;
+                }
+            }
+        }
+
+        private async Task CloseIfNotAlreadyClosedAsync(CancellationToken cancellationToken)
+        {
+            if (!_closed)
+            {
+                try
+                {
+                    if (_cursorId != 0)
                     {
-                        _channelSource.Dispose();
+                        await KillCursorsAsync(cancellationToken).ConfigureAwait(false);
                     }
                 }
+                finally
+                {
+                    _closed = true;
+                }
+            }
+        }
+
+        private void CloseIfNotAlreadyClosedFromDispose()
+        {
+            try
+            {
+                using (var source = new CancellationTokenSource(TimeSpan.FromSeconds(10)))
+                {
+                    CloseIfNotAlreadyClosed(source.Token);
+                }
+            }
+            catch
+            {
+                // ignore any exceptions from CloseIfNotAlreadyClosed when called from Dispose
             }
-            _disposed = true;
         }
 
         private void DisposeChannelSourceIfNoLongerNeeded()
@@ -316,21 +450,39 @@ private async Task<CursorBatch<TDocument>> GetNextBatchAsync(CancellationToken c
             }
         }
 
-        private void KillCursor(long cursorId, CancellationToken cancellationToken)
+        private void KillCursors(CancellationToken cancellationToken)
         {
-            try
+            using (EventContext.BeginOperation(_operationId))
+            using (EventContext.BeginKillCursors(_collectionNamespace))
+            using (var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10)))
+            using (var channel = _channelSource.GetChannel(cancellationTokenSource.Token))
             {
-                using (EventContext.BeginOperation(_operationId))
-                using (EventContext.BeginKillCursors(_collectionNamespace))
-                using (var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10)))
-                using (var channel = _channelSource.GetChannel(cancellationTokenSource.Token))
+                if (Feature.KillCursorsCommand.IsSupported(channel.ConnectionDescription.ServerVersion))
+                {
+                    ExecuteKillCursorsCommand(channel, cancellationToken);
+                }
+                else
                 {
                     ExecuteKillCursorsProtocol(channel, cancellationToken);
                 }
             }
-            catch
+        }
+
+        private async Task KillCursorsAsync(CancellationToken cancellationToken)
+        {
+            using (EventContext.BeginOperation(_operationId))
+            using (EventContext.BeginKillCursors(_collectionNamespace))
+            using (var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(10)))
+            using (var channel = await _channelSource.GetChannelAsync(cancellationTokenSource.Token).ConfigureAwait(false))
             {
-                // ignore exceptions
+                if (Feature.KillCursorsCommand.IsSupported(channel.ConnectionDescription.ServerVersion))
+                {
+                    await ExecuteKillCursorsCommandAsync(channel, cancellationToken).ConfigureAwait(false);
+                }
+                else
+                {
+                    await ExecuteKillCursorsProtocolAsync(channel, cancellationToken).ConfigureAwait(false);
+                }
             }
         }
 
@@ -393,6 +545,28 @@ private void ThrowIfDisposed()
             }
         }
 
+        private void ThrowIfKillCursorsCommandFailed(BsonDocument commandResult, ConnectionId connectionId)
+        {
+            if (!commandResult.GetValue("ok", false).ToBoolean())
+            {
+                throw new MongoCommandException(connectionId, "Kill cursors command returned an error.", commandResult);
+            }
+            else
+            {
+                var notFoundCursors = commandResult["cursorsNotFound"].AsBsonArray;
+                if (notFoundCursors.Count > 0)
+                {
+                    throw new MongoCursorNotFoundException(connectionId, _cursorId, commandResult);
+                }
+
+                var killedCursors = commandResult["cursorsKilled"].AsBsonArray.Select(c => c.ToInt64());
+                if (!killedCursors.Contains(_cursorId))
+                {
+                    throw new MongoCommandException(connectionId, "Kill cursors command failed.", commandResult);
+                }
+            }
+        }
+
         private bool TryMoveNext(out bool hasMore)
         {
             hasMore = false;
@@ -419,4 +593,4 @@ private bool TryMoveNext(out bool hasMore)
             return false;
         }
     }
-}
+}