Fix error handling in Create()

akarnokd · Oct 30, 2018 · f340828 · f340828
1 parent 0125fe0
commit f340828
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/async-enumerable-dotnet-test/CreateTest.cs b/async-enumerable-dotnet-test/CreateTest.cs
@@ -5,6 +5,7 @@
 using Xunit;
 using async_enumerable_dotnet;
 using System.Threading.Tasks;
+using System;
 
 namespace async_enumerable_dotnet_test
 {
@@ -43,5 +44,43 @@ public async void Range_Loop()
                 await Range();
             }
         }
+
+        [Fact]
+        public async void Items_And_Error()
+        {
+            var result = AsyncEnumerable.Create<int>(async e =>
+            {
+                await e.Next(1);
+
+                await e.Next(2);
+
+                throw new InvalidOperationException();
+            });
+
+            await result.AssertFailure(typeof(InvalidOperationException), 1, 2);
+        }
+
+        [Fact]
+        public async ValueTask Take()
+        {
+            await AsyncEnumerable.Create<int>(async e =>
+            {
+                for (var i = 0; i < 10 && !e.DisposeAsyncRequested; i++)
+                {
+                    await e.Next(i);
+                }
+            })
+            .Take(5)
+            .AssertResult(0, 1, 2, 3, 4);
+        }
+
+        [Fact]
+        public async void Take_Loop()
+        {
+            for (int j = 0; j < 1000; j++)
+            {
+                await Take();
+            }
+        }
     }
 }
diff --git a/async-enumerable-dotnet/impl/CreateEmitter.cs b/async-enumerable-dotnet/impl/CreateEmitter.cs
@@ -32,6 +32,7 @@ private sealed class CreateEmitterEnumerator : IAsyncEnumerator<T>, IAsyncEmitte
             public bool DisposeAsyncRequested => _disposeRequested;
 
             private bool _hasValue;
+            private Exception _error;
 
             public T Current { get; private set; }
 
@@ -41,8 +42,7 @@ private sealed class CreateEmitterEnumerator : IAsyncEnumerator<T>, IAsyncEmitte
 
             internal void SetTask(Task task)
             {
-                _task = task;
-                task.ContinueWith(async t =>
+                _task = task.ContinueWith(async t =>
                 {
                     if (_disposeRequested)
                     {
@@ -55,6 +55,8 @@ internal void SetTask(Task task)
                         return;
                     }
 
+                    _error = ExceptionHelper.Extract(t.Exception);
+
                     ResumeHelper.Resume(ref _valueReady);
                 });
             }
@@ -78,6 +80,13 @@ public async ValueTask<bool> MoveNextAsync()
                     return true;
                 }
                 Current = default;
+
+                var ex = _error;
+                if (ex != null)
+                {
+                    _error = null;
+                    throw ex;
+                }
                 return false;
             }