How to make Dispose await for all async methods?

Question

I have disposable class with async methods.

class Gateway : IDisposable {
  public Gateway() {}
  public void Dispose() {}

  public async Task<Data> Request1 () {...}
  public async Task<Data> Request2 () {...}
  public async Task<Data> Request3 () {...}
}

I need Dispose to await until all running requests are completed.

So, either I need to track of all running tasks, or use AsyncLock from AsyncEx, or something else?

Updated

As I can see someone is afraid of blocking Dispose. Then we could make Task WaitForCompletionAsync() or Task CancelAllAsync() methods.

Answer 1

For the time being, you'll have to add a CloseAsync method that your users have to invoke.

Once C# 8.0 is released, you can rely on the IAsyncDisposable Interface and its language support:

await using (var asyncDisposable = GetAsyncDisposable())
{
    // ...
} // await asyncDisposable.DisposeAsync()

Answer 2

Here is a solution for reusable async disposal support. Due to that .NET Core 3.0 is not yet released, I will provide code for both current C# version (7.3) and beta (8.0).

Once IDisposable.Dispose() is called on the object, it won't block and will ensure the disposal at once all tasks are completed.

Source Code (Current C# version, without IAsyncDisposable )

Usings related:

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

The interface that can dispose after all tracking tasks are complete:

public interface ITrackingDisposable : IDisposable
{
    //The implementation of the actual disposings
    Task FinishDisposeAsync();
}

The disposer that tracks all running tasks and call the deferred disposal at the appropriate timing:

public class TrackingDisposer : IDisposable
{
    private readonly LinkedList<Task> _tasks = new LinkedList<Task>();

    private readonly ITrackingDisposable _target;

    public bool IsDisposed { get; private set; } = false;

    //The supported class must implement ITrackingDisposable
    public TrackingDisposer(ITrackingDisposable target)
    => _target = target ?? throw new ArgumentNullException();

    //Add a task to the tracking list, returns false if disposed
    //Without return value
    public bool Track(Func<Task> func, out Task result)
    {
        lock (_tasks)
        {
            if (IsDisposed)
            {
                result = null;
                return false;
            }

            var task = func();
            var node = _tasks.AddFirst(task);

            async Task ending()
            {
                await task;
                var dispose = false;
                lock (_tasks)
                {
                    _tasks.Remove(node);
                    dispose = IsDisposed && _tasks.Count == 0;
                }
                if (dispose)
                {
                    await _target.FinishDisposeAsync();
                }
            }

            result = ending();
        }
        return true;
    }

    //With return value
    public bool Track<TResult>(Func<Task<TResult>> func, out Task<TResult> result)
    {
        lock (_tasks)
        {
            if (IsDisposed)
            {
                result = null;
                return false;
            }

            var task = func();
            var node = _tasks.AddFirst(task);

            async Task<TResult> ending()
            {
                var result = await task;
                var dispose = false;
                lock (_tasks)
                {
                    _tasks.Remove(node);
                    dispose = IsDisposed && _tasks.Count == 0;
                }
                if (dispose)
                {
                    await _target.FinishDisposeAsync();
                }
                return result;
            }

            result = ending();
        }
        return true;
    }

    //The entry of applying for dispose
    public void Dispose()
    {
        var dispose = false;

        lock (_tasks)
        {
            if (IsDisposed)
            {
                return;
            }

            IsDisposed = true;
            dispose = _tasks.Count == 0;
        }

        if (dispose)
        {
            _target.FinishDisposeAsync();
        }
    }
}

A base class simplifying the implementation:

public abstract class TrackingDisposable : ITrackingDisposable
{
    private readonly TrackingDisposer _disposer;

    public TrackingDisposable()
    => _disposer = new TrackingDisposer(this);

    protected virtual void FinishDispose() { }

    protected virtual Task FinishDisposeAsync()
    => Task.CompletedTask;

    Task ITrackingDisposable.FinishDisposeAsync()
    {
        FinishDispose();
        return FinishDisposeAsync();
    }

    public void Dispose()
    => _disposer.Dispose();

    protected Task Track(Func<Task> func)
    => _disposer.Track(func, out var result)
        ? result
        : throw new ObjectDisposedException(nameof(TrackingDisposable));

    protected Task<TResult> Track<TResult>(Func<Task<TResult>> func)
    => _disposer.Track(func, out var result)
        ? result
        : throw new ObjectDisposedException(nameof(TrackingDisposable));
}

Demo & Test Output

Testing class:

internal sealed class TestDisposingObject : TrackingDisposable
{
    public Task Job0Async() => Track(async () =>
    {
        await Task.Delay(200);
        Console.WriteLine("Job0 done.");
    });

    public Task<string> Job1Async(int ms) => Track(async () =>
    {
        await Task.Delay(ms);
        return "Job1 done.";
    });

    protected override void FinishDispose()
    => Console.WriteLine("Disposed.");
}

Main:

internal static class Program
{
    private static async Task Main()
    {
        var result0 = default(Task);
        var result1 = default(Task);
        var obj = new TestDisposingObject();
        result0 = obj.Job0Async();
        result1 = obj.Job1Async(100).ContinueWith(r => Console.WriteLine(r.Result));
        obj.Dispose();
        Console.WriteLine("Waiting For jobs done...");
        await Task.WhenAll(result0, result1);
    }
}

Output:

Waiting For jobs done...
Job1 done.
Job0 done.
Disposed.

Additional, C# 8.0 (with IAsyncDisposable )

Replace the type definition with follows:

public interface ITrackingDisposable : IDisposable, IAsyncDisposable
{
    Task FinishDisposeAsync();
}

public class TrackingDisposer : IDisposable, IAsyncDisposable
{
    private readonly LinkedList<Task> _tasks = new LinkedList<Task>();

    private readonly ITrackingDisposable _target;

    private readonly TaskCompletionSource<object> _disposing = new TaskCompletionSource<object>();

    public bool IsDisposed { get; private set; } = false;

    public TrackingDisposer(ITrackingDisposable target)
    => _target = target ?? throw new ArgumentNullException();

    public bool Track(Func<Task> func, out Task result)
    {
        lock (_tasks)
        {
            if (IsDisposed)
            {
                result = null;
                return false;
            }

            var task = func();
            var node = _tasks.AddFirst(task);

            async Task ending()
            {
                await task;
                var dispose = false;
                lock (_tasks)
                {
                    _tasks.Remove(node);
                    dispose = IsDisposed && _tasks.Count == 0;
                }
                if (dispose)
                {
                    await _target.FinishDisposeAsync();
                    _disposing.SetResult(null);
                }
            }

            result = ending();
        }
        return true;
    }

    public bool Track<TResult>(Func<Task<TResult>> func, out Task<TResult> result)
    {
        lock (_tasks)
        {
            if (IsDisposed)
            {
                result = null;
                return false;
            }

            var task = func();
            var node = _tasks.AddFirst(task);

            async Task<TResult> ending()
            {
                var result = await task;
                var dispose = false;
                lock (_tasks)
                {
                    _tasks.Remove(node);
                    dispose = IsDisposed && _tasks.Count == 0;
                }
                if (dispose)
                {
                    await _target.FinishDisposeAsync();
                    _disposing.SetResult(null);
                }
                return result;
            }

            result = ending();
        }
        return true;
    }

    public void Dispose()
    {
        var dispose = false;

        lock (_tasks)
        {
            if (IsDisposed)
            {
                return;
            }

            IsDisposed = true;
            dispose = _tasks.Count == 0;
        }

        if (dispose)
        {
            _target.FinishDisposeAsync();
            _disposing.SetResult(null);
        }
    }

    public ValueTask DisposeAsync()
    {
        Dispose();
        return new ValueTask(_disposing.Task);
    }
}

public abstract class TrackingDisposable : ITrackingDisposable
{
    private readonly TrackingDisposer _disposer;

    public TrackingDisposable()
    => _disposer = new TrackingDisposer(this);

    protected virtual void FinishDispose() { }

    protected virtual Task FinishDisposeAsync()
    => Task.CompletedTask;

    Task ITrackingDisposable.FinishDisposeAsync()
    {
        FinishDispose();
        return FinishDisposeAsync();
    }

    public void Dispose()
    => _disposer.Dispose();

    public ValueTask DisposeAsync() => _disposer.DisposeAsync();

    protected Task Track(Func<Task> func)
    => _disposer.Track(func, out var result)
        ? result
        : throw new ObjectDisposedException(nameof(TrackingDisposable));

    protected Task<TResult> Track<TResult>(Func<Task<TResult>> func)
    => _disposer.Track(func, out var result)
        ? result
        : throw new ObjectDisposedException(nameof(TrackingDisposable));
}

Test Main:

internal static class Program
{
    private static async Task Main()
    {
        await using var obj = new TestDisposingObject();
        _ = obj.Job0Async();
        _ = obj.Job1Async(100).ContinueWith(r => Console.WriteLine(r.Result));
        Console.WriteLine("Waiting For jobs done...");
    }
}

Answer 3

The issue here is that there is no async version of Dispose() (yet). So you have to ask yourself-- what do you expect to happen when you call Dispose() , or when a using block ends....? In other words, what is the requirement?

You could require Dispose to await all outstanding tasks and then do its work. But Dispose can't use await (it's not async). The best it can do is call Result to force the task to complete, but that would be a blocking call, and if any of the async tasks are awaiting anything else, it could easily deadlock.

Instead, I suggest the following requirement: When the caller calls Dispose() , the call will flag the Gateway to be disposed and then return immediately, secure in the knowledge that the disposal mechanism will activate itself when the last task has completed.

If that requirement is adequate, it is possible, but a bit messy. Here's how:

1. Each time a method (such as Request ) is called, "wrap" the returned Task in another Task that includes a check to see if the caller has requested the Gateway to be disposed.

2. If a disposal has been requested, go ahead and dispose right then and there before flagging the task as completed. Thus when the caller awaits the task, it will force the disposal.

Here's my implementation. I told you it was ugly.

class Gateway : IDisposable 
{
    protected readonly HttpClient _client = new HttpClient();  //an inner class that must be disposed when Gateway disposes
    protected bool _disposalRequested = false;
    protected bool _disposalCompleted = false;
    protected int _tasksRunning = 0;


    public void Dispose()
    {
        Console.WriteLine("Dispose() called.");
        _disposalRequested = true;  
        if (_tasksRunning == 0)
        {
            Console.WriteLine("No running tasks, so disposing immediately.");
            DisposeInternal();
        }
        else
        {
            Console.WriteLine("There are running tasks, so disposal shall be deferred.");
        }
    }

    protected void DisposeInternal()
    {
        if (!_disposalCompleted)
        {
            Console.WriteLine("Disposing");
            _client.Dispose();
            _disposalCompleted = true;
        }
    }

    protected async Task<T> AddDisposeWrapper<T>(Func<Task<T>> func)
    {
        if (_disposalRequested) throw new ObjectDisposedException("Disposal has already been requested. No new requests can be handled at this point.");

        _tasksRunning++;
        var result = await func();
        _tasksRunning--;
        await DisposalCheck();
        return result;
    }

    protected async Task DisposalCheck()
    {
        if (_disposalRequested) DisposeInternal();
    }

    public Task<Data> Request1()
    {
        return AddDisposeWrapper
        (
            Request1Internal
        );
    }

    public Task<Data> Request2()
    {
        return AddDisposeWrapper
        (
            Request2Internal
        );
    }

    protected async Task<Data> Request1Internal()
    {
        Console.WriteLine("Performing Request1 (slow)");
        await Task.Delay(3000);
        Console.WriteLine("Request1 has finished. Returning new Data.");
        return new Data();
    }

    protected async Task<Data> Request2Internal()
    {
        Console.WriteLine("Performing Request2 (fast)");
        await Task.Delay(1);
        Console.WriteLine("Request2 has finished. Returning new Data.");
        return new Data();
    }
}

Here's some test code:

public class Program
{
    public static async Task Test1()
    {
        Task<Data> task;
        using (var gateway = new Gateway())
        {
            task = gateway.Request1();
            await Task.Delay(1000);
        }
        var data = await task;
        Console.WriteLine("Test 1 is complete.");
    }

    public static async Task Test2()
    {
        Task<Data> task;
        using (var gateway = new Gateway())
        {
            task = gateway.Request2();
            await Task.Delay(1000);
        }
        var data = await task;
        Console.WriteLine("Test 2 is complete.");
    }

    public static async Task MainAsync()
    {
        await Test1();
        await Test2();
    }

    public static void Main()
    {
        MainAsync().GetAwaiter().GetResult();
        Console.WriteLine("Run completed at {0:yyyy-MM-dd HH:mm:ss}", DateTime.Now);
    }
}

This is the output:

Performing Request1 (slow)
Dispose() called.
There are running tasks, so disposal shall be deferred.
Request1 has finished. Returning new Data.
Disposing
Test 1 is complete.
Performing Request2 (fast)
Request2 has finished. Returning new Data.
Dispose() called.
No running tasks, so disposing immediately.
Disposing
Test 2 is complete.
Run completed at 2019-05-15 00:34:46

And here is my Fiddle, in case you wish to try it out: Link

I don't really recommend this (if something is going to be disposed, you should have better control over its lifespan), but it was fun writing this code for you.

Note: Due to the use of reference counting, additional work would be needed to make this solution thread-safe or to make it resilient to the case where one of Gateway's request methods throws an exception.

Answer 4

Disposing and awaiting for completion are different things. So, I'd rather will throw exception when tasks are still running.

I wrote example with Nito.AsyncEx.AsyncConditionVariable . I didn't tested it, but I think it should work. Just use Completion.WaitAsync() .

Also I recommend this article: https://blog.stephencleary.com/2013/03/async-oop-6-disposal.html

class Gateway : IDisposable {

  private int runningTaskCount;
  public AsyncConditionVariable Completion { get; } = new AsyncConditionVariable( new AsyncLock() );

  public Gateway() {
  }
  public void Dispose() {
    if (runningTaskCount != 0) throw new InvalidOperationException( "You can not call this method when tasks are running" );
  }

  public async Task<Data> Request1 () {
    BeginTask();
    ...
    EndTask();
  }

  private void BeginTask() {
    Interlocked.Increment( ref runningTaskCount );
  }
  private void EndTask() {
    var result = Interlocked.Decrement( ref runningTaskCount );
    if (result == 0) Completion.NotifyAll();
  }

}