Structured concurrency groups for disposable resources

[timcassell]

Structured concurrency groups are super powerful for loading resources concurrently and canceling operations if an exception occurs, but they lack the native ability to dispose successfully loaded resources when an exception or cancelation occurs. I have implemented here 2 types that wrap the built-in structured concurrency groups to support this behavior for merge and race scenarios.

PromiseAllDisposableGroup

public readonly struct PromiseAllDisposableGroup<TDisposable> where TDisposable : IDisposable
{
    private readonly PromiseAllResultsGroup<TDisposable> _promiseGroup;

    private PromiseAllDisposableGroup(PromiseAllResultsGroup<TDisposable> promiseGroup)
    {
        _promiseGroup = promiseGroup;
    }

    public static PromiseAllDisposableGroup<TDisposable> New(out CancelationToken groupCancelationToken)
        => New(CancelationToken.None, out groupCancelationToken);

    public static PromiseAllDisposableGroup<TDisposable> New(CancelationToken sourceCancelationToken, out CancelationToken groupCancelationToken)
        => new(PromiseAllResultsGroup<TDisposable>.New(sourceCancelationToken, out groupCancelationToken));

    public PromiseAllDisposableGroup<TDisposable> Add(Promise<TDisposable> promise)
        => new(_promiseGroup.Add(promise));

    public async Promise<IList<TDisposable>> WaitAsync()
    {
        // If an exception occurred, or the token was canceled before promises completed successfully,
        // we need to dispose the items that were retrieved before re-throwing.
        var results = await _promiseGroup.WaitAsync();

        var items = new List<TDisposable>(results.Count);
        bool canceled = false;
        List<Exception> exceptions = null;
        foreach (var result in results)
        {
            switch (result.State)
            {
                case Promise.State.Resolved:
                    items.Add(result.Value);
                    break;
                case Promise.State.Canceled:
                    canceled = true;
                    break;
                case Promise.State.Rejected:
                    exceptions ??= new List<Exception>(results.Count);
                    exceptions.Add(result.Reason as Exception ?? Promise.RejectException(result.Reason));
                    break;
            }
        }
        if (canceled || exceptions != null)
        {
            foreach (var item in items)
            {
                item.Dispose();
            }
            throw exceptions == null ? Promise.CancelException() : new AggregateException(exceptions);
        }

        return items;
    }
}

PromiseRaceDisposableGroup

public readonly struct PromiseRaceDisposableGroup<TDisposable> where TDisposable : IDisposable
{
    private readonly PromiseEachGroup<TDisposable> _promiseGroup;
    private readonly CancelationSource _cancelationSource;
    private readonly bool _cancelOnNonResolved;

    private PromiseRaceDisposableGroup(PromiseEachGroup<TDisposable> promiseGroup, CancelationSource cancelationSource, bool cancelOnNonResolved)
    {
        _promiseGroup = promiseGroup;
        _cancelationSource = cancelationSource;
        _cancelOnNonResolved = cancelOnNonResolved;
    }

    public static PromiseRaceDisposableGroup<TDisposable> New(out CancelationToken groupCancelationToken, bool cancelOnNonResolved = true)
        => New(CancelationToken.None, out groupCancelationToken);

    public static PromiseRaceDisposableGroup<TDisposable> New(CancelationToken sourceCancelationToken, out CancelationToken groupCancelationToken, bool cancelOnNonResolved = true)
    {
        var cancelationSource = CancelationSource.New(sourceCancelationToken);
        return new(PromiseEachGroup<TDisposable>.New(cancelationSource.Token, out groupCancelationToken), cancelationSource, cancelOnNonResolved);
    }

    public PromiseRaceDisposableGroup<TDisposable> Add(Promise<TDisposable> promise)
        => new(_promiseGroup.Add(promise), _cancelationSource, _cancelOnNonResolved);

    private void CancelGroup(ref List<Exception> exceptions, ref bool cancelationThrew)
    {
        try
        {
            _cancelationSource.Cancel();
        }
        catch (Exception e)
        {
            cancelationThrew = true;
            exceptions ??= new List<Exception>();
            exceptions.Add(e);
        }
    }

    public async Promise<TDisposable> WaitAsync()
    {
        var enumerator = _promiseGroup.GetAsyncEnumerable().GetAsyncEnumerator();

        if (!await enumerator.MoveNextAsync())
        {
            throw new System.InvalidOperationException("The promise race disposable group must have at least one promise added to it.");
        }

        bool resolved = false;
        TDisposable value = default;
        bool cancelationThrew = false;
        List<Exception> exceptions = null;
        // Iterate over the results.
        do
        {
            var result = enumerator.Current;
            if (result.State == Promise.State.Resolved)
            {
                // We have obtained the first successful result, cancel the rest.
                CancelGroup(ref exceptions, ref cancelationThrew);
                if (resolved)
                {
                    // Dispose all successful results except the first.
                    result.Value.Dispose();
                }
                else
                {
                    resolved = true;
                    value = result.Value;
                }
            }
            else
            {
                if (_cancelOnNonResolved)
                {
                    CancelGroup(ref exceptions, ref cancelationThrew);
                }
                if (result.State == Promise.State.Rejected)
                {
                    exceptions ??= new List<Exception>();
                    exceptions.Add(result.Reason as Exception ?? Promise.RejectException(result.Reason));
                }
            }
        } while (await enumerator.MoveNextAsync());

        _cancelationSource.Dispose();

        // Race promise group ignores rejections if any promise was resolved,
        // unless a cancelation token callback threw.
        if ((!resolved | cancelationThrew) & exceptions != null)
        {
            if (resolved)
            {
                value.Dispose();
            }
            throw new AggregateException(exceptions);
        }
        if (resolved)
        {
            return value;
        }
        throw Promise.CancelException();
    }
}

An example usage for such types is to wait on multiple AsyncSemaphore instances.

// Wait all
var promiseGroup = PromiseAllDisposableGroup<AsyncSemaphore.Scope>.New(cancelationToken, out var groupCancelationToken);
foreach (var sem in semaphores)
{
    promiseGroup = promiseGroup.Add(sem.EnterScopeAsync(groupCancelationToken));
}
var scopes = await promiseGroup.WaitAsync();
try
{
    // Do work
}
finally
{
    foreach (var scope in scopes)
    {
        scope.Dispose();
    }
}

// Wait any
var promiseGroup = PromiseRaceDisposableGroup<AsyncSemaphore.Scope>.New(cancelationToken, out var groupCancelationToken);
foreach (var sem in semaphores)
{
    promiseGroup = promiseGroup.Add(sem.EnterScopeAsync(groupCancelationToken));
}
using (await promiseGroup.WaitAsync())
{
    // Do work
}

[timcassell]

In v3.5.0, this is much simpler by just passing a cleanup delegate to the promise group.

var promiseGroup = PromiseAllGroup<AsyncSemaphore.Scope>
    .New(cancelationToken, out var groupCancelationToken, scope => scope.dispose());

var promiseGroup = PromiseRaceGroup<AsyncSemaphore.Scope>
    .New(cancelationToken, out var groupCancelationToken, scope => scope.dispose());