-
Notifications
You must be signed in to change notification settings - Fork 1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Allow Streams to be consumed as IAsyncEnumerable (#4742)
* Add `.RunAsAsyncEnumerable` and `.RunAsAsyncEnumerableBuffer` to Akka Streams Source DSL * Add tests for CancellationToken support, use Cancellationtoken in rest of stream flow Co-authored-by: Aaron Stannard <aaron@petabridge.com>
- Loading branch information
1 parent
a0f4705
commit b77ed9a
Showing
5 changed files
with
295 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,151 @@ | ||
| //----------------------------------------------------------------------- | ||
| // <copyright file="AsyncEnumerableSpec.cs" company="Akka.NET Project"> | ||
| // Copyright (C) 2009-2021 Lightbend Inc. <http://www.lightbend.com> | ||
| // Copyright (C) 2013-2021 .NET Foundation <https://github.com/akkadotnet/akka.net> | ||
| // </copyright> | ||
| //----------------------------------------------------------------------- | ||
|
|
||
| using System; | ||
| using System.Linq; | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
| using Akka.Pattern; | ||
| using Akka.Routing; | ||
| using Akka.Streams.Dsl; | ||
| using Akka.Streams.TestKit; | ||
| using Akka.TestKit; | ||
| using FluentAssertions; | ||
| using Nito.AsyncEx.Synchronous; | ||
| using Xunit; | ||
| using Xunit.Abstractions; | ||
|
|
||
| namespace Akka.Streams.Tests.Dsl | ||
| { | ||
| #if NETCOREAPP | ||
| public class AsyncEnumerableSpec : AkkaSpec | ||
| { | ||
| private ActorMaterializer Materializer { get; } | ||
|
|
||
| public AsyncEnumerableSpec(ITestOutputHelper helper) : base(helper) | ||
| { | ||
| var settings = ActorMaterializerSettings.Create(Sys).WithInputBuffer(2, 16); | ||
| Materializer = ActorMaterializer.Create(Sys, settings); | ||
| } | ||
|
|
||
| [Fact] public async Task RunAsAsyncEnumerable_Uses_CancellationToken() | ||
| { | ||
| var input = Enumerable.Range(1, 6).ToList(); | ||
|
|
||
| var cts = new CancellationTokenSource(); | ||
| var token = cts.Token; | ||
|
|
||
| var asyncEnumerable = Source.From(input).RunAsAsyncEnumerable(Materializer); | ||
| var output = input.ToArray(); | ||
| bool caught = false; | ||
| try | ||
| { | ||
| await foreach (var a in asyncEnumerable.WithCancellation(token)) | ||
| { | ||
| cts.Cancel(); | ||
| } | ||
| } | ||
| catch (OperationCanceledException e) | ||
| { | ||
| caught = true; | ||
| } | ||
|
|
||
| caught.ShouldBeTrue(); | ||
| } | ||
|
|
||
| [Fact] | ||
| public async Task RunAsAsyncEnumerable_must_return_an_IAsyncEnumerableT_from_a_Source() | ||
| { | ||
| var input = Enumerable.Range(1, 6).ToList(); | ||
| var asyncEnumerable = Source.From(input).RunAsAsyncEnumerable(Materializer); | ||
| var output = input.ToArray(); | ||
| await foreach (var a in asyncEnumerable) | ||
| { | ||
| (output[0] == a).ShouldBeTrue("Did not get elements in order!"); | ||
| output = output.Skip(1).ToArray(); | ||
| } | ||
| output.Length.ShouldBe(0,"Did not receive all elements!"); | ||
| } | ||
|
|
||
| [Fact] | ||
| public async Task RunAsAsyncEnumerable_must_allow_multiple_enumerations() | ||
| { | ||
| var input = Enumerable.Range(1, 6).ToList(); | ||
| var asyncEnumerable = Source.From(input).RunAsAsyncEnumerable(Materializer); | ||
| var output = input.ToArray(); | ||
| await foreach (var a in asyncEnumerable) | ||
| { | ||
| (output[0] == a).ShouldBeTrue("Did not get elements in order!"); | ||
| output = output.Skip(1).ToArray(); | ||
| } | ||
| output.Length.ShouldBe(0,"Did not receive all elements!"); | ||
|
|
||
| output = input.ToArray(); | ||
| await foreach (var a in asyncEnumerable) | ||
| { | ||
| (output[0] == a).ShouldBeTrue("Did not get elements in order!"); | ||
| output = output.Skip(1).ToArray(); | ||
| } | ||
| output.Length.ShouldBe(0,"Did not receive all elements in second enumeration!!"); | ||
| } | ||
|
|
||
|
|
||
| [Fact] | ||
| public async Task RunAsAsyncEnumerable_Throws_on_Abrupt_Stream_termination() | ||
| { | ||
| var materializer = ActorMaterializer.Create(Sys); | ||
| var probe = this.CreatePublisherProbe<int>(); | ||
| var task = Source.FromPublisher(probe).RunAsAsyncEnumerable(materializer); | ||
|
|
||
| var a = Task.Run( async () => | ||
| { | ||
| await foreach (var notused in task) | ||
| { | ||
| materializer.Shutdown(); | ||
| } | ||
| }); | ||
| //since we are collapsing the stream inside the read | ||
| //we want to send messages so we aren't just waiting forever. | ||
| probe.SendNext(1); | ||
| probe.SendNext(2); | ||
| bool thrown = false; | ||
| try | ||
| { | ||
| await a; | ||
| } | ||
| catch (AbruptTerminationException e) | ||
| { | ||
| thrown = true; | ||
| } | ||
| thrown.ShouldBeTrue(); | ||
| } | ||
|
|
||
| [Fact] | ||
| public async Task RunAsAsyncEnumerable_Throws_if_materializer_gone_before_Enumeration() | ||
| { | ||
| var materializer = ActorMaterializer.Create(Sys); | ||
| var probe = this.CreatePublisherProbe<int>(); | ||
| var task = Source.FromPublisher(probe).RunAsAsyncEnumerable(materializer); | ||
| materializer.Shutdown(); | ||
|
|
||
| async Task ShouldThrow() | ||
| { | ||
| await foreach (var a in task) | ||
| { | ||
|
|
||
| } | ||
| } | ||
|
|
||
| await Assert.ThrowsAsync<IllegalStateException>(ShouldThrow); | ||
| } | ||
|
|
||
|
|
||
| } | ||
|
|
||
| #else | ||
| #endif | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| // //----------------------------------------------------------------------- | ||
| // // <copyright file="AsyncEnumerable.cs" company="Akka.NET Project"> | ||
| // // Copyright (C) 2009-2021 Lightbend Inc. <http://www.lightbend.com> | ||
| // // Copyright (C) 2013-2021 .NET Foundation <https://github.com/akkadotnet/akka.net> | ||
| // // </copyright> | ||
| // //----------------------------------------------------------------------- | ||
|
|
||
| using System.Collections.Generic; | ||
| using System.Threading; | ||
| using System.Threading.Tasks; | ||
| using Akka.Configuration.Hocon; | ||
|
|
||
| namespace Akka.Streams.Dsl | ||
| { | ||
| /// <summary> | ||
| /// Used to treat an <see cref="IRunnableGraph{TMat}"/> of <see cref="ISinkQueue{T}"/> | ||
| /// as an <see cref="IAsyncEnumerable{T}"/> | ||
| /// </summary> | ||
| /// <typeparam name="T"></typeparam> | ||
| public sealed class StreamsAsyncEnumerableRerunnable<T,TMat> : IAsyncEnumerable<T> | ||
| { | ||
| private static readonly Sink<T, ISinkQueue<T>> defaultSinkqueue = | ||
| Sink.Queue<T>(); | ||
| private readonly Source<T, TMat> _source; | ||
| private readonly IMaterializer _materializer; | ||
|
|
||
| private readonly Sink<T, ISinkQueue<T>> thisSinkQueue; | ||
| //private readonly IRunnableGraph<(UniqueKillSwitch, ISinkQueue<T>)> _graph; | ||
| public StreamsAsyncEnumerableRerunnable(Source<T,TMat> source, IMaterializer materializer) | ||
| { | ||
| _source = source; | ||
| _materializer = materializer; | ||
| thisSinkQueue = defaultSinkqueue; | ||
| } | ||
|
|
||
| public StreamsAsyncEnumerableRerunnable(Source<T, TMat> source, | ||
| IMaterializer materializer, int minBuf, int maxBuf):this(source, materializer) | ||
| { | ||
| thisSinkQueue = | ||
| defaultSinkqueue.WithAttributes( | ||
| Attributes.CreateInputBuffer(minBuf, maxBuf)); | ||
| } | ||
|
|
||
| public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken) | ||
| { | ||
| cancellationToken.ThrowIfCancellationRequested(); | ||
|
|
||
| return new SinkQueueAsyncEnumerator<T>(_source | ||
| .Via(cancellationToken.AsFlow<T>(cancelGracefully: true)) | ||
| .ViaMaterialized(KillSwitches.Single<T>(), Keep.Right) | ||
| .ToMaterialized(thisSinkQueue, Keep.Both) | ||
| .Run(_materializer), | ||
| cancellationToken); | ||
| } | ||
| } | ||
| /// <summary> | ||
| /// Wraps a Sink Queue and Killswitch around <see cref="IAsyncEnumerator{T}"/> | ||
| /// </summary> | ||
| /// <typeparam name="T"></typeparam> | ||
| public sealed class SinkQueueAsyncEnumerator<T> : IAsyncEnumerator<T> | ||
| { | ||
| private ISinkQueue<T> _sinkQueue; | ||
| private IKillSwitch _killSwitch; | ||
| private CancellationToken _token; | ||
| public SinkQueueAsyncEnumerator((UniqueKillSwitch killSwitch,ISinkQueue<T> sinkQueue) queueAndSwitch, CancellationToken token) | ||
| { | ||
| _sinkQueue = queueAndSwitch.sinkQueue; | ||
| _killSwitch = queueAndSwitch.killSwitch; | ||
| _token = token; | ||
| } | ||
| public async ValueTask DisposeAsync() | ||
| { | ||
| //If we are disposing, let's shut down the stream | ||
| //so that we don't have data hanging around. | ||
| _killSwitch.Shutdown(); | ||
| _killSwitch = null; | ||
| _sinkQueue = null; | ||
| } | ||
|
|
||
| public async ValueTask<bool> MoveNextAsync() | ||
| { | ||
| _token.ThrowIfCancellationRequested(); | ||
| var opt = await _sinkQueue.PullAsync(); | ||
| if (opt.HasValue) | ||
| { | ||
| Current = opt.Value; | ||
| return true; | ||
| } | ||
| else | ||
| { | ||
| return false; | ||
| } | ||
| } | ||
|
|
||
| public T Current { get; private set; } | ||
| } | ||
| } |