Use async iterators in Finally.

dotnet · Jan 29, 2019 · b48137d · b48137d
1 parent c66fdda
commit b48137d
Showing 1 changed file with 41 additions and 0 deletions.
diff --git a/Ix.NET/Source/System.Interactive.Async/System/Linq/Operators/Finally.cs b/Ix.NET/Source/System.Interactive.Async/System/Linq/Operators/Finally.cs
@@ -4,6 +4,7 @@
 
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Linq
@@ -17,7 +18,26 @@ public static IAsyncEnumerable<TSource> Finally<TSource>(this IAsyncEnumerable<T
             if (finallyAction == null)
                 throw Error.ArgumentNull(nameof(finallyAction));
 
+#if USE_ASYNC_ITERATOR
+            return Create(Core);
+
+            async IAsyncEnumerator<TSource> Core(CancellationToken cancellationToken)
+            {
+                try
+                {
+                    await foreach (TSource item in source.WithCancellation(cancellationToken).ConfigureAwait(false))
+                    {
+                        yield return item;
+                    }
+                }
+                finally
+                {
+                    finallyAction();
+                }
+            }
+#else
             return new FinallyAsyncIterator<TSource>(source, finallyAction);
+#endif
         }
 
         public static IAsyncEnumerable<TSource> Finally<TSource>(this IAsyncEnumerable<TSource> source, Func<Task> finallyAction)
@@ -27,11 +47,31 @@ public static IAsyncEnumerable<TSource> Finally<TSource>(this IAsyncEnumerable<T
             if (finallyAction == null)
                 throw Error.ArgumentNull(nameof(finallyAction));
 
+#if USE_ASYNC_ITERATOR
+            return Create(Core);
+
+            async IAsyncEnumerator<TSource> Core(CancellationToken cancellationToken)
+            {
+                try
+                {
+                    await foreach (TSource item in source.WithCancellation(cancellationToken).ConfigureAwait(false))
+                    {
+                        yield return item;
+                    }
+                }
+                finally
+                {
+                    await finallyAction().ConfigureAwait(false);
+                }
+            }
+#else
             return new FinallyAsyncIteratorWithTask<TSource>(source, finallyAction);
+#endif
         }
 
         // REVIEW: No cancellation support for finally action.
 
+#if !USE_ASYNC_ITERATOR
         private sealed class FinallyAsyncIterator<TSource> : AsyncIterator<TSource>
         {
             private readonly Action _finallyAction;
@@ -148,4 +188,5 @@ protected override async ValueTask<bool> MoveNextCore()
             }
         }
     }
+#endif
 }