Potential race condition in RefCount when using it with source that supports resubscription after error

[ArXen42]

Bug

Which library version?
6.1.0, 7.0.0-preview.1

What are the platform(s), environment(s) and related component version(s)?
I've tested on Linux x64 and arm64

What is the use case or problem?

There appears to be race condition within RefCount implementation when using it like this:

src.Multicast<StatelessSubject>().RefCount().ObserveOn(TaskPoolScheduler.Default).Retry()

When error occurs, retry attempts will race with each other and potentially lead to a state when subscriptions exist but connection to source observable is not recreated. More complete example:

// Simulate polling some remote device
var src = Observable.Create<Int32>(async (observer, ct) =>
    {
        Console.WriteLine("START");
        Int32 i = 0;
        while (!ct.IsCancellationRequested)
        {
            Console.WriteLine(++i);
            observer.OnNext(i);
            await Task.Delay(500);

            // Simulate connection failure
            if (i > 5)
                throw new InvalidOperationException();
        }
    })
    .SubscribeOn(TaskPoolScheduler.Default);

src = src.Multicast(new StatelessSubject<Int32>())
    .RefCount();

const Int32 parallelCount = 12;
for (Int32 i = 0; i < parallelCount; i++)
{
    Int32 i1 = i;
    src
        .ObserveOn(TaskPoolScheduler.Default)
        .Do(_ => {},
            ex => { Console.WriteLine($"Error {i1}"); }, // Simulate handling error in downstream consumer
            () => Console.WriteLine("Completed"))
        .Retry() // This retry can't just be moved before .Multicast since downstream consumers might need to know when errors occur (to reset it's own internal business logic, etc) 
        .Subscribe();
}

Thread.Sleep(-1);

/// "Dumb" subject that doesn't remember error state that I've taken from https://stackoverflow.com/a/64991229
public class StatelessSubject<T> : ISubject<T>, ISubject<T, T>, IObserver<T>, IObservable<T>
{
    private IImmutableList<IObserver<T>> _observers = ImmutableArray<IObserver<T>>.Empty;

    public IDisposable Subscribe(IObserver<T> observer)
    {
        ImmutableInterlocked.Update(ref _observers, x => x.Add(observer));
        return Disposable.Create(() => ImmutableInterlocked.Update(ref _observers, (Func<IImmutableList<IObserver<T>>, IImmutableList<IObserver<T>>>)(x => x.Remove(observer))));
    }

    public void OnNext(T value)
    {
        foreach (var observer in Volatile.Read(ref _observers))
            observer.OnNext(value);
    }

    public void OnError(Exception error)
    {
        foreach (var observer in Interlocked.Exchange(ref _observers, []))
            observer.OnError(error);
    }

    public void OnCompleted()
    {
        foreach (var observer in Interlocked.Exchange(ref _observers, []))
            observer.OnCompleted();
    }
}

With large enough parallelCount it will stuck after a few iterations.

What is the expected outcome?

When error occurs in the source, I'd expect RefCount to atomically reset it's state (clear internal subscription to source and all observers), propagate error to "old" observers and be ready to immediately accept subscriptions from new observers, restarting source subscription anew.

What is the actual outcome?

RefCount operator entering incorrect state, when it has active downstream subscriptions, but doesn't create upstream subscription.

[ArXen42]

I'm not entirely sure if this is a bug or simply unsupported scenario for RefCount, though it doesn't look like my .Multicast<StatelessSubject> violates Rx grammar, so it feels like it should be valid to use with RefCount.

Since the internals of the operator are rather complex, I don't know how to properly fix this. For now I've made simple implementation from scratch that seems to work well for my scenario:

public class RefCountThreadSafe<T>(IObservable<T> src) : IObservable<T>
{
    private readonly Object      _gate = new();
    private          Connection? _currentConnection;

    public IDisposable Subscribe(IObserver<T> observer)
    {
        Boolean    shouldConnect;
        Connection connection;

        lock (_gate)
        {
            if (_currentConnection is null || _currentConnection.IsTerminated)
            {
                _currentConnection = new(observer);
                shouldConnect      = true;
            }
            else
            {
                _currentConnection.Observers = _currentConnection.Observers.Add(observer);
                shouldConnect                = false;
            }

            connection = _currentConnection;
        }

        if (shouldConnect)
        {
            var sourceSubscription = src.SubscribeSafe(new AnonymousObserver<T>(v => ForwardOnNext(connection, v),
                ex => ForwardTerminal(connection, obs => obs.OnError(ex)),
                () => ForwardTerminal(connection, obs => obs.OnCompleted())));

            // OnError might be called here in case of exception in source Subscribe
            IDisposable? toDispose = null;
            lock (_gate)
            {
                if (connection.IsTerminated)
                    toDispose = sourceSubscription;
                else
                    connection.SourceSubscription = sourceSubscription;
            }

            if (toDispose != null)
            {
                toDispose.Dispose();
            }
        }

        return Disposable.Create(() => Unsubscribe(connection, observer));
    }

    private void Unsubscribe(Connection connection, IObserver<T> observer)
    {
        IDisposable? toDispose = null;

        lock (_gate)
        {
            if (connection.IsTerminated)
                return;

            var removed = connection.Observers.Remove(observer);
            if (removed == connection.Observers)
            {
                Debug.WriteLine("Observer was already removed");
                return;
            }

            connection.Observers = removed;
            if (connection.Observers.Length == 0)
            {
                connection.IsTerminated       = true;
                toDispose                     = connection.SourceSubscription;
                connection.SourceSubscription = null;
            }
        }

        toDispose?.Dispose();
    }

    private void ForwardOnNext(Connection connection, T value)
    {
        ImmutableArray<IObserver<T>> observers;
        lock (_gate)
        {
            observers = connection.Observers;
        }

        foreach (var observer in observers)
        {
            if (Volatile.Read(ref connection.IsTerminated))
                return;
            observer.OnNext(value);
        }
    }

    private void ForwardTerminal(Connection connection, Action<IObserver<T>> action)
    {
        ImmutableArray<IObserver<T>> observers;
        IDisposable?                 toDispose;

        lock (_gate)
        {
            if (connection.IsTerminated)
                return;

            connection.IsTerminated       = true;
            observers                     = connection.Observers;
            toDispose                     = connection.SourceSubscription;
            connection.SourceSubscription = null;
        }

        try
        {
            foreach (var observer in observers)
                action(observer);
        }
        finally
        {
            toDispose?.Dispose();
        }
    }

    private sealed class Connection(IObserver<T> firstObserver)
    {
        public ImmutableArray<IObserver<T>> Observers = [firstObserver];
        public IDisposable?                 SourceSubscription;
        public Boolean                      IsTerminated;
    }
}