reactive/AsyncRx.NET/System.Reactive.Async.Linq/System/Reactive/Internal/EventPatternSourceBase.cs

// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT License.
// See the LICENSE file in the project root for more information. 

using System.Collections.Generic;
using System.Threading.Tasks;

namespace System.Reactive
{
    /// <summary>
    /// Base class for classes that expose an observable sequence as a well-known event pattern (sender, event arguments).
    /// Contains functionality to maintain a map of event handler delegates to observable sequence subscriptions. Subclasses
    /// should only add an event with custom add and remove methods calling into the base class's operations.
    /// </summary>
    /// <typeparam name="TSender">The type of the sender that raises the event.</typeparam>
    /// <typeparam name="TEventArgs">The type of the event data generated by the event.</typeparam>
    internal abstract class EventPatternSourceBase<TSender, TEventArgs>
    {
        private readonly IAsyncObservable<EventPattern<TSender, TEventArgs>> _source;
        private readonly Dictionary<Delegate, Stack<IAsyncDisposable>> _subscriptions;
        private readonly Action<Action<TSender, TEventArgs>, /*object,*/ EventPattern<TSender, TEventArgs>> _invokeHandler;

        /// <summary>
        /// Creates a new event pattern source.
        /// </summary>
        /// <param name="source">Source sequence to expose as an event.</param>
        /// <param name="invokeHandler">Delegate used to invoke the event for each element of the sequence.</param>
        /// <exception cref="ArgumentNullException"><paramref name="source"/> or <paramref name="invokeHandler"/> is null.</exception>
        protected EventPatternSourceBase(IAsyncObservable<EventPattern<TSender, TEventArgs>> source, Action<Action<TSender, TEventArgs>, /*object,*/ EventPattern<TSender, TEventArgs>> invokeHandler)
        {
            _source = source ?? throw new ArgumentNullException(nameof(source));
            _invokeHandler = invokeHandler ?? throw new ArgumentNullException(nameof(invokeHandler));
            _subscriptions = new Dictionary<Delegate, Stack<IAsyncDisposable>>();
        }

        /// <summary>
        /// Adds the specified event handler, causing a subscription to the underlying source.
        /// </summary>
        /// <param name="handler">Event handler to add. The same delegate should be passed to the Remove operation in order to remove the event handler.</param>
        /// <param name="invoke">Invocation delegate to raise the event in the derived class.</param>
        /// <exception cref="ArgumentNullException"><paramref name="handler"/> or <paramref name="invoke"/> is null.</exception>
        protected void Add(Delegate handler, Action<TSender, TEventArgs> invoke)
        {
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));
            if (invoke == null)
                throw new ArgumentNullException(nameof(invoke));

            var gate = new object();
            var isAdded = false;
            var isDone = false;

            var remove = new Action(() =>
            {
                lock (gate)
                {
                    if (isAdded)
                        Remove(handler);
                    else
                        isDone = true;
                }
            });

            //
            // [OK] Use of unsafe SubscribeAsync: non-pretentious wrapper of an observable in an event; exceptions can occur during +=.
            //
            var d = _source.SubscribeAsync(
                x => { _invokeHandler(invoke, /*this,*/ x); return default; },
                ex => { remove(); return new ValueTask(Task.FromException(ex)); },
                () => { remove(); return default; }
            ).GetAwaiter().GetResult();

            lock (gate)
            {
                if (!isDone)
                {
                    Add(handler, d);
                    isAdded = true;
                }
            }
        }

        private void Add(Delegate handler, IAsyncDisposable disposable)
        {
            lock (_subscriptions)
            {
                var l = new Stack<IAsyncDisposable>();
                if (!_subscriptions.TryGetValue(handler, out l))
                    _subscriptions[handler] = l = new Stack<IAsyncDisposable>();

                l.Push(disposable);
            }
        }

        /// <summary>
        /// Removes the specified event handler, causing a disposal of the corresponding subscription to the underlying source that was created during the Add operation.
        /// </summary>
        /// <param name="handler">Event handler to remove. This should be the same delegate as one that was passed to the Add operation.</param>
        /// <exception cref="ArgumentNullException"><paramref name="handler"/> is null.</exception>
        protected void Remove(Delegate handler)
        {
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));

            var d = default(IAsyncDisposable);

            lock (_subscriptions)
            {
                var l = new Stack<IAsyncDisposable>();
                if (_subscriptions.TryGetValue(handler, out l))
                {
                    d = l.Pop();
                    if (l.Count == 0)
                        _subscriptions.Remove(handler);
                }
            }

            if (d != null)
            {
                d.DisposeAsync().GetAwaiter().GetResult();
            }
        }
    }
}