88
99namespace AsyncEnumerableExtensions
1010{
11- public class TaskAsyncEnumerable < T > : IAsyncEnumerable < T >
11+ internal class TaskAsyncEnumerable < T > : IAsyncEnumerable < T >
1212 {
1313
1414 private readonly Func < IAsyncEnumerableSink < T > , CancellationToken , Task > generator ;
15- private readonly bool acceptsCancellationToken ;
1615
1716 public TaskAsyncEnumerable ( Func < IAsyncEnumerableSink < T > , Task > sourceTask )
1817 {
1918 if ( sourceTask == null ) throw new ArgumentNullException ( nameof ( sourceTask ) ) ;
2019 generator = ( sink , ct ) => sourceTask ( sink ) ;
21- acceptsCancellationToken = false ;
2220 }
2321
2422 public TaskAsyncEnumerable ( Func < IAsyncEnumerableSink < T > , CancellationToken , Task > sourceTask )
2523 {
2624 this . generator = sourceTask ?? throw new ArgumentNullException ( nameof ( sourceTask ) ) ;
27- acceptsCancellationToken = true ;
2825 }
2926
3027 /// <inheritdoc />
3128 public IAsyncEnumerator < T > GetEnumerator ( )
3229 {
33- return new Enumerator ( generator , acceptsCancellationToken ) ;
30+ return new Enumerator ( generator ) ;
3431 }
3532
3633 private class Enumerator : IAsyncEnumerator < T >
@@ -43,7 +40,7 @@ private class Enumerator : IAsyncEnumerator<T>
4340 private CancellationToken lastMoveNextCancellationToken = CancellationToken . None ;
4441 private CancellationTokenSource lastCombinedCancellationTokenSource ;
4542
46- public Enumerator ( Func < IAsyncEnumerableSink < T > , CancellationToken , Task > generator , bool acceptsCancellationToken )
43+ public Enumerator ( Func < IAsyncEnumerableSink < T > , CancellationToken , Task > generator )
4744 {
4845 Debug . Assert ( generator != null ) ;
4946 this . generator = generator ;
@@ -188,10 +185,11 @@ public bool Yield(IEnumerable<T> items)
188185
189186 public Task Wait ( CancellationToken cancellationToken )
190187 {
191- cancellationToken . ThrowIfCancellationRequested ( ) ;
188+ if ( cancellationToken . IsCancellationRequested )
189+ return Task . Delay ( - 1 , cancellationToken ) ;
192190 lock ( syncLock )
193191 {
194- if ( queue == null ) throw new ObjectDisposedException ( nameof ( AsyncEnumerableBuffer < T > ) ) ;
192+ if ( queue == null ) return ObjectDisposedTask ;
195193 if ( queue . Count == 0 ) return CompletedTask ;
196194 if ( onQueueExhaustedTcs == null )
197195 {
@@ -249,8 +247,8 @@ internal void Terminate()
249247 {
250248 lock ( syncLock )
251249 {
252- onQueueExhaustedTcs ? . TrySetCanceled ( ) ;
253- onYieldTcs ? . TrySetCanceled ( ) ;
250+ onQueueExhaustedTcs ? . TrySetException ( new ObjectDisposedException ( nameof ( AsyncEnumerableBuffer < T > ) ) ) ;
251+ onYieldTcs ? . TrySetException ( new ObjectDisposedException ( nameof ( AsyncEnumerableBuffer < T > ) ) ) ;
254252 onQueueExhaustedTcs = null ;
255253 onYieldTcs = null ;
256254 queue = null ;
0 commit comments