CSHARP-3458: Extend IAsyncCursor and IAsyncCursorSource to support IAsyncEnumerable

Author: adelinowona

src/MongoDB.Driver/Core/IAsyncCursor.cs

@@ -360,6 +360,17 @@ public static class IAsyncCursorExtensions
             return new AsyncCursorEnumerableOneTimeAdapter<TDocument>(cursor, cancellationToken);
         }
 
+        /// <summary>
+        /// Wraps a cursor in an IAsyncEnumerable that can be enumerated one time.
+        /// </summary>
+        /// <typeparam name="TDocument">The type of the document.</typeparam>
+        /// <param name="cursor">The cursor.</param>
+        /// <returns>An IAsyncEnumerable.</returns>
+        public static IAsyncEnumerable<TDocument> ToAsyncEnumerable<TDocument>(this IAsyncCursor<TDocument> cursor)
+        {
+            return new AsyncCursorEnumerableOneTimeAdapter<TDocument>(cursor);
+        }
+
         /// <summary>
         /// Returns a list containing all the documents returned by a cursor.
         /// </summary>

src/MongoDB.Driver/Core/IAsyncCursorSource.cs

@@ -336,6 +336,17 @@ public static class IAsyncCursorSourceExtensions
             return new AsyncCursorSourceEnumerableAdapter<TDocument>(source, cancellationToken);
         }
 
+        /// <summary>
+        /// Wraps a cursor source in an IAsyncEnumerable. Each time GetAsyncEnumerator is called a new cursor is fetched from the cursor source.
+        /// </summary>
+        /// <typeparam name="TDocument">The type of the document.</typeparam>
+        /// <param name="source">The source.</param>
+        /// <returns>An IAsyncEnumerable.</returns>
+        public static IAsyncEnumerable<TDocument> ToAsyncEnumerable<TDocument>(this IAsyncCursorSource<TDocument> source)
+        {
+            return new AsyncCursorSourceEnumerableAdapter<TDocument>(source);
+        }
+
         /// <summary>
         /// Returns a list containing all the documents returned by the cursor returned by a cursor source.
         /// </summary>

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

@@ -21,18 +21,33 @@
 
 namespace MongoDB.Driver.Core.Operations
 {
-    internal sealed class AsyncCursorEnumerableOneTimeAdapter<TDocument> : IEnumerable<TDocument>
+    internal sealed class AsyncCursorEnumerableOneTimeAdapter<TDocument> : IEnumerable<TDocument>, IAsyncEnumerable<TDocument>
     {
         private readonly CancellationToken _cancellationToken;
         private readonly IAsyncCursor<TDocument> _cursor;
         private bool _hasBeenEnumerated;
 
+        public AsyncCursorEnumerableOneTimeAdapter(IAsyncCursor<TDocument> cursor)
+        {
+            _cursor = Ensure.IsNotNull(cursor, nameof(cursor));
+        }
+
         public AsyncCursorEnumerableOneTimeAdapter(IAsyncCursor<TDocument> cursor, CancellationToken cancellationToken)
         {
             _cursor = Ensure.IsNotNull(cursor, nameof(cursor));
             _cancellationToken = cancellationToken;
         }
 
+        public IAsyncEnumerator<TDocument> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            if (_hasBeenEnumerated)
+            {
+                throw new InvalidOperationException("An IAsyncCursor can only be enumerated once.");
+            }
+            _hasBeenEnumerated = true;
+            return new AsyncCursorEnumerator<TDocument>(_cursor, cancellationToken);
+        }
+
         public IEnumerator<TDocument> GetEnumerator()
         {
             if (_hasBeenEnumerated)

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

@@ -17,11 +17,12 @@
 using System.Collections;
 using System.Collections.Generic;
 using System.Threading;
+using System.Threading.Tasks;
 using MongoDB.Driver.Core.Misc;
 
 namespace MongoDB.Driver.Core.Operations
 {
-    internal class AsyncCursorEnumerator<TDocument> : IEnumerator<TDocument>
+    internal class AsyncCursorEnumerator<TDocument> : IEnumerator<TDocument>, IAsyncEnumerator<TDocument>
     {
         // private fields
         private IEnumerator<TDocument> _batchEnumerator;
@@ -72,6 +73,12 @@ public void Dispose()
             }
         }
 
+        public ValueTask DisposeAsync()
+        {
+            Dispose();
+            return new ValueTask(Task.CompletedTask);
+        }
+
         public bool MoveNext()
         {
             ThrowIfDisposed();
@@ -102,6 +109,28 @@ public bool MoveNext()
             }
         }
 
+        public async ValueTask<bool> MoveNextAsync()
+        {
+            ThrowIfDisposed();
+            _started = true;
+
+            if (_batchEnumerator != null && _batchEnumerator.MoveNext())
+            {
+                return true;
+            }
+
+            while (await _cursor.MoveNextAsync(_cancellationToken).ConfigureAwait(false))
+            {
+                _batchEnumerator?.Dispose();
+                _batchEnumerator = _cursor.Current.GetEnumerator();
+                return _batchEnumerator.MoveNext();
+            }
+
+            _batchEnumerator = null;
+            _finished = true;
+            return false;
+        }
+
         public void Reset()
         {
             ThrowIfDisposed();

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

@@ -13,27 +13,37 @@
 * limitations under the License.
 */
 
-using System;
 using System.Collections;
 using System.Collections.Generic;
 using System.Threading;
 using MongoDB.Driver.Core.Misc;
 
 namespace MongoDB.Driver.Core.Operations
 {
-    internal class AsyncCursorSourceEnumerableAdapter<TDocument> : IEnumerable<TDocument>
+    internal class AsyncCursorSourceEnumerableAdapter<TDocument> : IEnumerable<TDocument>, IAsyncEnumerable<TDocument>
     {
         // private fields
         private readonly CancellationToken _cancellationToken;
         private readonly IAsyncCursorSource<TDocument> _source;
 
         // constructors
+        public AsyncCursorSourceEnumerableAdapter(IAsyncCursorSource<TDocument> source)
+        {
+            _source = Ensure.IsNotNull(source, nameof(source));
+        }
+
         public AsyncCursorSourceEnumerableAdapter(IAsyncCursorSource<TDocument> source, CancellationToken cancellationToken)
         {
             _source = Ensure.IsNotNull(source, nameof(source));
             _cancellationToken = cancellationToken;
         }
 
+        public IAsyncEnumerator<TDocument> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+        {
+            var cursor = _source.ToCursor(cancellationToken);
+            return new AsyncCursorEnumerator<TDocument>(cursor, cancellationToken);
+        }
+
         // public methods
         public IEnumerator<TDocument> GetEnumerator()
         {

tests/MongoDB.Driver.Tests/Core/IAsyncCursorExtensionsTests.cs

@@ -16,6 +16,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using FluentAssertions;
 using MongoDB.Bson;
 using MongoDB.Bson.Serialization.Serializers;
@@ -201,6 +202,37 @@ public void SingleOrDefault_should_throw_when_cursor_has_wrong_number_of_documen
             action.ShouldThrow<InvalidOperationException>();
         }
 
+        [Fact]
+        public void ToAsyncEnumerable_result_should_only_be_enumerable_one_time()
+        {
+            var cursor = CreateCursor(2);
+            var enumerable = cursor.ToAsyncEnumerable();
+            enumerable.GetAsyncEnumerator();
+
+            Action action = () => enumerable.GetAsyncEnumerator();
+
+            action.ShouldThrow<InvalidOperationException>();
+        }
+
+        [Fact]
+        public async Task ToAsyncEnumerable_should_return_expected_result()
+        {
+            var cursor = CreateCursor(2);
+            var expectedDocuments = new[]
+            {
+                new BsonDocument("_id", 0),
+                new BsonDocument("_id", 1)
+            };
+
+            var result = new List<BsonDocument>();
+            await foreach (var doc in cursor.ToAsyncEnumerable())
+            {
+                result.Add(doc);
+            }
+
+            result.Should().Equal(expectedDocuments);
+        }
+
         [Fact]
         public void ToEnumerable_result_should_only_be_enumerable_one_time()
         {

tests/MongoDB.Driver.Tests/Core/IAsyncCursorSourceExtensionsTests.cs

@@ -203,6 +203,31 @@ public void SingleOrDefault_should_throw_when_cursor_has_wrong_number_of_documen
             action.ShouldThrow<InvalidOperationException>();
         }
 
+        [Theory]
+        [ParameterAttributeData]
+        public async Task ToAsyncEnumerable_result_should_be_enumerable_multiple_times(
+            [Values(1, 2)] int times)
+        {
+            var source = CreateCursorSource(2);
+            var expectedDocuments = new[]
+            {
+                new BsonDocument("_id", 0),
+                new BsonDocument("_id", 1)
+            };
+
+            var result = new List<BsonDocument>();
+            for (var i = 0; i < times; i++)
+            {
+                await foreach (var doc in source.ToAsyncEnumerable())
+                {
+                    result.Add(doc);
+                }
+
+                result.Should().Equal(expectedDocuments);
+                result.Clear();
+            }
+        }
+
         [Theory]
         [ParameterAttributeData]
         public void ToEnumerable_result_should_be_enumerable_multiple_times(