Cooperatively pausing async methods

Recently I was writing an app that processed a bunch of files asynchronously.  As with the Windows copy file dialog, I wanted to be able to provide the user with a button that would pause the processing operation.

To achieve that, I implemented a simple mechanism that would allow me to pass a “pause token” into the async method, which the async method could asynchronous wait on at appropriate points. 

public async Task ProcessFiles(
    IEnumerable<StorageFile> files, PauseToken pauseToken)
{
    foreach(var file in files)
    {
        await pauseToken.WaitWhilePausedAsync();
        await ProcessAsync(file);
    }
}

My pause token follows a similar design to that of CancellationToken and CancellationTokenSource.  I have a PauseToken instance that I can pass to any number of operations (synchronous or asynchronous), and those operations can monitor that token to be alerted to pause requests.  Separately, a PauseTokenSource is responsible for creating the PauseToken to be handed out and for issuing the pause requests.

public class PauseTokenSource
{
    public bool IsPaused { get; set; }
    public PauseToken Token { get; }
}

public struct PauseToken
{
    public bool IsPaused { get; }
    public Task WaitWhilePausedAsync();
}

We’ll start by implementing PauseTokenSource, which is the meat of the implementation; as with CancellationToken and CancellationTokenSource, PauseToken is just a thin value-type veneer on top of PauseTokenSource that just delegates most calls to the underlying reference type.  PauseTokenSource has one instance field:

private volatile TaskCompletionSource<bool> m_paused;

The m_paused field is the TaskCompletionSource<bool> that can be used to complete the Task we’ll hand out to waiters when the instance is paused (such that when we’re un-paused, we’ll set the Task to wake up all the waiters): if m_paused is null, we’re not paused, and if it’s non-null, we’re currently paused.

The bulk of the implementation is then in PauseTokenSource.IsPaused.  Its getter just returns whether m_paused is not null, but its setter is more complicated:

public bool IsPaused
{
    get { return m_paused != null; }
    set
    {
        if (value)
        {
            Interlocked.CompareExchange(
                ref m_paused, new TaskCompletionSource<bool>(), null);
        }
        else
        {
            while (true)
            {
                var tcs = m_paused;
                if (tcs == null) return;
                if (Interlocked.CompareExchange(ref m_paused, null, tcs) == tcs)
                {
                    tcs.SetResult(true);
                    break;
                }
            }
        }
    }
}

If IsPaused is being set to true, then we simply need to transition m_paused from null to a new TaskCompletionSource<bool>; we do this with an interlocked compare-exchange so that we only do the transition if m_paused is null, regardless of what other threads we might be competing with.

If IsPaused is being set to false, we need to do two things: transition m_paused from non-null to null, and complete the Task from the TaskCompletionSource<bool> that was stored in m_paused.  We do this with another Interlocked.CompareExchange, and as we need to tell the CompareExchange operation what exact value we expect to find in m_paused, and as another thread could be changing it out from under us, we need to do this in a standard compare-exchange loop: grab the current value, do the compare exchange assuming that value, and if the value actually changed between the time we grabbed it and the time we did the compare-exchange, repeat.

To shield these implementation details from PauseToken, we’ll add an internal WaitWhilePauseAsync method to PauseTokenSource that PauseToken can then access.

internal Task WaitWhilePausedAsync()
{
    var cur = m_paused;
    return cur != null ? cur.Task : s_completedTask;
}

This method just grabs m_paused, and if it’s non-null returns its Task.  If it is null, then we’re not paused, so we can hand back an already completed Task in order to avoid unnecessary allocations (since we should expect it to be very common that WaitWhilePauseAsync is called when not actually paused):

internal static readonly Task s_completedTask = Task.FromResult(true);

The last member we need on PauseTokenSource is the Token property that will return the associated PauseToken:

public PauseToken Token { get { return new PauseToken(this); } }

Now for implementing PauseToken.  Its implementation is very simple, as it’s just a wrapper over the PauseTokenSource from which its constructed:

public struct PauseToken
{
    private readonly PauseTokenSource m_source;
    internal PauseToken(PauseTokenSource source) { m_source = source; }

    public bool IsPaused { get { return m_source != null && m_source.IsPaused; } }

    public Task WaitWhilePausedAsync()
    {
        return IsPaused ?
            m_source.WaitWhilePausedAsync() :
            PauseTokenSource.s_completedTask;
    }
}

PauseToken’s IsPaused property only has a getter and not a setter, since our design requires that all transitioning from un-paused to paused, and vice versa, is done via the PauseTokenSource (that way, only someone with access to the source can cause the transition).  PauseToken’s IsPaused getter just delegates to the source’s IsPaused; of course, as this PauseToken is a struct, it’s possible it could have been default initialized such that m_source would be null… in that case, we’ll just return false from IsPaused.

Finally, we have our PauseToken’s WaitWhilePauseAsync method.  If we’re paused, we simply delegate to the source’s WaitWhilePausedAsync implementation we already saw.  If we’re not paused (which could include not having a source), we just return our cached already-completed Task.

That’s it: our implementation is now complete, and we can start using it to pause asynchronous operations.  Here’s a basic console-based example of using our new PauseToken type:

class Program
{
    static void Main()
    {
        var pts = new PauseTokenSource();
        Task.Run(() =>
        {
            while (true)
            {
                Console.ReadLine();
                pts.IsPaused = !pts.IsPaused;
            }
        });
        SomeMethodAsync(pts.Token).Wait();
    }

    public static async Task SomeMethodAsync(PauseToken pause)
    {
        for (int i = 0; i < 100; i++)
        {
            Console.WriteLine(i);
            await Task.Delay(100);
            await pause.WaitWhilePausedAsync();
        }
    }
}

As a final thought, for those of you familiar with various kinds of synchronization primitives, PauseTokenSource might remind you of one in particular: manual reset events.  In fact, that’s basically what it is, just with a different API set (for comparison, see this blog post on building an AsyncManualResetEvent).  Setting IsPaused to false is like setting/signaling a manual reset event, and setting it to true is like resetting one.

Enjoy!