Adding Timeout.

AsyncRx.NET/System.Reactive.Async.Linq/System/Reactive/Linq/Operators/Timeout.cs

@@ -0,0 +1,238 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information. 
+
+using System.Reactive.Concurrency;
+using System.Reactive.Disposables;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace System.Reactive.Linq
+{
+    partial class AsyncObservable
+    {
+        public static IAsyncObservable<TSource> Timeout<TSource>(this IAsyncObservable<TSource> source, TimeSpan dueTime)
+        {
+            if (source == null)
+                throw new ArgumentNullException(nameof(source));
+
+            return Create<TSource>(async observer =>
+            {
+                var sourceSubscription = new SingleAssignmentAsyncDisposable();
+
+                var (sink, disposable) = await AsyncObserver.Timeout(observer, sourceSubscription, dueTime).ConfigureAwait(false);
+
+                var sourceSubscriptionInner = await source.SubscribeAsync(sink).ConfigureAwait(false);
+
+                await sourceSubscription.AssignAsync(sourceSubscriptionInner).ConfigureAwait(false);
+
+                return disposable;
+            });
+        }
+
+        public static IAsyncObservable<TSource> Timeout<TSource>(this IAsyncObservable<TSource> source, TimeSpan dueTime, IAsyncScheduler scheduler)
+        {
+            if (source == null)
+                throw new ArgumentNullException(nameof(source));
+            if (scheduler == null)
+                throw new ArgumentNullException(nameof(scheduler));
+
+            return Create<TSource>(async observer =>
+            {
+                var sourceSubscription = new SingleAssignmentAsyncDisposable();
+
+                var (sink, disposable) = await AsyncObserver.Timeout(observer, sourceSubscription, dueTime, scheduler).ConfigureAwait(false);
+
+                var sourceSubscriptionInner = await source.SubscribeAsync(sink).ConfigureAwait(false);
+
+                await sourceSubscription.AssignAsync(sourceSubscriptionInner).ConfigureAwait(false);
+
+                return disposable;
+            });
+        }
+
+        public static IAsyncObservable<TSource> Timeout<TSource>(this IAsyncObservable<TSource> source, TimeSpan dueTime, IAsyncObservable<TSource> other)
+        {
+            if (source == null)
+                throw new ArgumentNullException(nameof(source));
+            if (other == null)
+                throw new ArgumentNullException(nameof(other));
+
+            return Create<TSource>(async observer =>
+            {
+                var sourceSubscription = new SingleAssignmentAsyncDisposable();
+
+                var (sink, disposable) = await AsyncObserver.Timeout(observer, sourceSubscription, dueTime, other).ConfigureAwait(false);
+
+                var sourceSubscriptionInner = await source.SubscribeAsync(sink).ConfigureAwait(false);
+
+                await sourceSubscription.AssignAsync(sourceSubscriptionInner).ConfigureAwait(false);
+
+                return disposable;
+            });
+        }
+
+        public static IAsyncObservable<TSource> Timeout<TSource>(this IAsyncObservable<TSource> source, TimeSpan dueTime, IAsyncObservable<TSource> other, IAsyncScheduler scheduler)
+        {
+            if (source == null)
+                throw new ArgumentNullException(nameof(source));
+            if (other == null)
+                throw new ArgumentNullException(nameof(other));
+            if (scheduler == null)
+                throw new ArgumentNullException(nameof(scheduler));
+
+            return Create<TSource>(async observer =>
+            {
+                var sourceSubscription = new SingleAssignmentAsyncDisposable();
+
+                var (sink, disposable) = await AsyncObserver.Timeout(observer, sourceSubscription, dueTime, other, scheduler).ConfigureAwait(false);
+
+                var sourceSubscriptionInner = await source.SubscribeAsync(sink).ConfigureAwait(false);
+
+                await sourceSubscription.AssignAsync(sourceSubscriptionInner).ConfigureAwait(false);
+
+                return disposable;
+            });
+        }
+    }
+
+    partial class AsyncObserver
+    {
+        public static Task<(IAsyncObserver<TSource>, IAsyncDisposable)> Timeout<TSource>(IAsyncObserver<TSource> observer, IAsyncDisposable sourceSubscription, TimeSpan dueTime)
+        {
+            if (observer == null)
+                throw new ArgumentNullException(nameof(observer));
+            if (sourceSubscription == null)
+                throw new ArgumentNullException(nameof(sourceSubscription));
+
+            return Timeout(observer, sourceSubscription, dueTime, AsyncObservable.Throw<TSource>(new TimeoutException()), TaskPoolAsyncScheduler.Default);
+        }
+
+        public static Task<(IAsyncObserver<TSource>, IAsyncDisposable)> Timeout<TSource>(IAsyncObserver<TSource> observer, IAsyncDisposable sourceSubscription, TimeSpan dueTime, IAsyncScheduler scheduler)
+        {
+            if (observer == null)
+                throw new ArgumentNullException(nameof(observer));
+            if (sourceSubscription == null)
+                throw new ArgumentNullException(nameof(sourceSubscription));
+            if (scheduler == null)
+                throw new ArgumentNullException(nameof(scheduler));
+
+            return Timeout(observer, sourceSubscription, dueTime, AsyncObservable.Throw<TSource>(new TimeoutException()), scheduler);
+        }
+
+        public static Task<(IAsyncObserver<TSource>, IAsyncDisposable)> Timeout<TSource>(IAsyncObserver<TSource> observer, IAsyncDisposable sourceSubscription, TimeSpan dueTime, IAsyncObservable<TSource> other)
+        {
+            if (observer == null)
+                throw new ArgumentNullException(nameof(observer));
+            if (sourceSubscription == null)
+                throw new ArgumentNullException(nameof(sourceSubscription));
+            if (other == null)
+                throw new ArgumentNullException(nameof(other));
+
+            return Timeout(observer, sourceSubscription, dueTime, other, TaskPoolAsyncScheduler.Default);
+        }
+
+        public static Task<(IAsyncObserver<TSource>, IAsyncDisposable)> Timeout<TSource>(IAsyncObserver<TSource> observer, IAsyncDisposable sourceSubscription, TimeSpan dueTime, IAsyncObservable<TSource> other, IAsyncScheduler scheduler)
+        {
+            if (observer == null)
+                throw new ArgumentNullException(nameof(observer));
+            if (sourceSubscription == null)
+                throw new ArgumentNullException(nameof(sourceSubscription));
+            if (other == null)
+                throw new ArgumentNullException(nameof(other));
+            if (scheduler == null)
+                throw new ArgumentNullException(nameof(scheduler));
+
+            async Task<(IAsyncObserver<TSource>, IAsyncDisposable)> CoreAsync()
+            {
+                var gate = new AsyncLock();
+
+                var switched = false;
+                var id = 0UL;
+
+                var timer = new SerialAsyncDisposable();
+                var subscription = new SerialAsyncDisposable();
+
+                var d = StableCompositeAsyncDisposable.Create(timer, subscription);
+
+                await subscription.AssignAsync(sourceSubscription).ConfigureAwait(false);
+
+                async Task<bool> OnAsync()
+                {
+                    var hasWon = false;
+
+                    using (await gate.LockAsync().ConfigureAwait(false))
+                    {
+                        if (!switched)
+                        {
+                            unchecked
+                            {
+                                ++id;
+                            }
+
+                            hasWon = true;
+                        }
+                    }
+
+                    return hasWon;
+                }
+
+                async Task CreateTimerAsync()
+                {
+                    var timerId = id;
+
+                    var timeout = await scheduler.ScheduleAsync(async ct =>
+                    {
+                        var hasWon = false;
+
+                        using (await gate.LockAsync().ConfigureAwait(false))
+                        {
+                            hasWon = switched = timerId == id;
+                        }
+
+                        if (hasWon)
+                        {
+                            var otherSubscription = await other.SubscribeAsync(observer).RendezVous(scheduler);
+
+                            await subscription.AssignAsync(otherSubscription).RendezVous(scheduler);
+                        }
+                    }, dueTime).ConfigureAwait(false);
+
+                    await timer.AssignAsync(timeout).ConfigureAwait(false);
+                }
+
+                var sink = Create<TSource>(
+                    async x =>
+                    {
+                        if (await OnAsync().ConfigureAwait(false))
+                        {
+                            await observer.OnNextAsync(x).ConfigureAwait(false);
+
+                            await CreateTimerAsync().ConfigureAwait(false);
+                        }
+                    },
+                    async ex =>
+                    {
+                        if (await OnAsync().ConfigureAwait(false))
+                        {
+                            await observer.OnErrorAsync(ex).ConfigureAwait(false);
+                        }
+                    },
+                    async () =>
+                    {
+                        if (await OnAsync().ConfigureAwait(false))
+                        {
+                            await observer.OnCompletedAsync().ConfigureAwait(false);
+                        }
+                    }
+                );
+
+                await CreateTimerAsync().ConfigureAwait(false);
+
+                return (sink, d);
+            }
+
+            return CoreAsync();
+        }
+    }
+}