Writing your own CancellationToken from scratch

 
 
  • Gérald Barré

This post is part of the series 'Async/await in C#'. Be sure to check out the rest of the blog posts of the series!

The previous post built a MyTask that can start work and report that it finished. It cannot stop anything. That gap is not an oversight in the implementation: stopping is a separate concern, and .NET models it as a separate pair of types.

This post builds those too. MyCancellationTokenSource and MyCancellationToken, from scratch, and then wires them into MyTask. The implementation is short. The design decisions inside it are the interesting part, because every one of them answers a question people ask about CancellationToken.

#Why it has to be cooperative

The obvious API is the one that does not exist: task.Kill().

.NET used to have something close. Thread.Abort injected an exception into another thread at whatever instruction it happened to be executing. That point could be in the middle of updating two related fields, between acquiring a lock and entering the try that would release it, or inside a finally block that was halfway through cleaning up. The result was a process in a state no code was written to handle, with no way to reason about what had and had not happened.

Thread.Abort throws PlatformNotSupportedException on .NET Core and later. There is no supported way to stop code that does not want to stop.

So cancellation is cooperative: the operation is told that it should stop, and it stops at a point it chose. Everything below follows from that. All the machinery does is deliver a one-way flag and, optionally, notify people when it flips.

#Two types, on purpose

The split between source and token is the first thing people notice and the last thing they appreciate.

  • MyCancellationTokenSource requests cancellation. It has Cancel().
  • MyCancellationToken observes it. It has no way to cancel anything.

It is a capability boundary. A method that takes a CancellationToken is handed the right to know, not the right to decide. If the two were one type, every method you passed a token to could cancel the whole operation, including operations it knows nothing about.

The token is a struct wrapping a reference:

C#
public readonly struct MyCancellationToken(MyCancellationTokenSource? source)
{
    public static MyCancellationToken None => default;

    public bool IsCancellationRequested => source is { IsCancellationRequested: true };

    public bool CanBeCanceled => source is not null;
}

Two consequences fall out of that one field.

default is a valid token that never cancels. No null checks, no CancellationToken?, no special case. That is why CancellationToken token = default is the idiomatic optional parameter.

CanBeCanceled is worth checking. It is false for default, which lets a hot path skip building registration machinery that could never fire. The real CancellationToken exposes it for exactly that reason.

#One volatile bool

The source's state is one field:

C#
public sealed class MyCancellationTokenSource : IDisposable
{
    private volatile bool _cancellationRequested;

    public bool IsCancellationRequested => _cancellationRequested;
}

volatile is not decoration here. Consider the polling loop this exists to support:

C#
while (!token.IsCancellationRequested)
{
    DoOneChunkOfWork();
}

Nothing inside that loop writes the flag, and nothing the JIT can see writes it. Without volatile the JIT is entitled to read it once, cache it in a register, and turn the loop into while (true). That is a legal optimisation of a single-threaded program, and it would make cancellation silently never work. volatile forces the read to happen every iteration.

The flag is one-way. A cancelled source never becomes uncancelled, which means an operation that has observed cancellation never has to wonder whether the answer changed behind it.

ThrowIfCancellationRequested is the convenience on top:

C#
public void ThrowIfCancellationRequested()
{
    if (IsCancellationRequested)
        throw new MyOperationCanceledException(this);
}

The exception carries the token. That matters at the catch site: "the thing I asked to stop was stopped" and "something else was cancelled and its exception passed through me" look identical otherwise. OperationCanceledException.CancellationToken exists so you can tell them apart, and so a Task can decide whether to enter the Canceled state or the Faulted one. The dedicated post on cancellation goes into what that distinction buys you.

#Callbacks, and the race in the middle

Polling only works when the operation is running code. A socket read that is parked waiting for bytes is not polling anything, and it is exactly the operation you most want to cancel. For those, cancellation has to be an event.

C#
internal MyCancellationRegistration Register(Action callback)
{
    lock (_lock)
    {
        if (!_cancellationRequested)
        {
            var registration = new Registration(callback);
            _registrations!.Add(registration);
            return new MyCancellationRegistration(this, registration);
        }
    }

    // Registering on an already cancelled source runs the callback immediately, outside the lock.
    callback();
    return default;
}

That last branch is the same race as the one in MyTask, in a new costume. Cancel and Register can run at the same instant on two threads, and the callback must run exactly once: either because Cancel found it in the list, or because Register found the flag already set. Deciding under the lock and acting outside it is what makes that true.

Cancel takes the list out rather than iterating it in place:

C#
public void Cancel()
{
    List<Registration>? registrations;

    lock (_lock)
    {
        if (_cancellationRequested)
            return;

        _cancellationRequested = true;

        registrations = _registrations;
        _registrations = null;
    }

    if (registrations is null)
        return;

    List<Exception>? exceptions = null;

    foreach (var registration in registrations)
    {
        try
        {
            registration.Callback();
        }
        catch (Exception ex)
        {
            exceptions ??= [];
            exceptions.Add(ex);
        }
    }

    if (exceptions is not null)
        throw new AggregateException(exceptions);
}

Three deliberate details:

  • Cancel on an already cancelled source does nothing. It is idempotent, not an error.
  • Callbacks run outside the lock. A callback that takes another lock, or registers another callback, must not deadlock against this one.
  • One failing callback does not stop the others. They are unrelated pieces of cleanup; the exceptions are collected and thrown together.

A registration can be undone, because a callback that outlives the operation that needed it is a leak:

C#
public readonly struct MyCancellationRegistration : IDisposable
{
    public void Dispose()
    {
        if (_source is not null && _registration is not null)
            _source.Unregister(_registration);
    }
}

A timeout is not a new concept, just a timer that calls Cancel:

C#
public void CancelAfter(TimeSpan delay)
{
    lock (_lock)
    {
        if (_cancellationRequested)
            return;

        _timer ??= new Timer(static state => ((MyCancellationTokenSource)state!).Cancel(), this, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
        _timer.Change(delay, Timeout.InfiniteTimeSpan);
    }
}

The source owns the timer and disposes it, which is half the reason CancellationTokenSource is IDisposable at all. A source with a pending timer that nobody disposes keeps a timer registration alive.

Linking is the other composition primitive: cancel when the caller gives up or when the timeout expires. It is Register applied to the sources you want to follow:

C#
public static MyCancellationTokenSource CreateLinked(params ReadOnlySpan<MyCancellationToken> tokens)
{
    var linked = new MyCancellationTokenSource();
    var registrations = new List<MyCancellationRegistration>(tokens.Length);

    foreach (var token in tokens)
    {
        registrations.Add(token.Register(linked.Cancel));
    }

    linked._linkedRegistrations = registrations;
    return linked;
}

Look at what that leaves behind: a callback on each parent source, holding a reference to the child. If the parent is a long-lived token (a host lifetime, an application-wide shutdown token) and the child is per-request, then failing to dispose the child means every request permanently adds an entry to a list that lives as long as the process. That is the linked-token-source leak, and writing the seven lines above is the clearest way to see why it happens.

#Wiring it into MyTask

MyTask needs a third terminal state, because cancelled is not failed:

C#
public void SetCanceled(MyCancellationToken token)
    => Complete(new MyOperationCanceledException(token), canceled: true);

Then the leaf operation honours the token. Delay is the whole idea in twenty lines:

C#
public static MyTask Delay(TimeSpan delay, MyCancellationToken token = default)
{
    var task = new MyTask();

    if (token.IsCancellationRequested)
    {
        task.SetCanceled(token);
        return task;
    }

    Timer? timer = null;
    MyCancellationRegistration registration = default;

    timer = new Timer(_ =>
    {
        registration.Dispose();
        timer!.Dispose();
        task.SetResult();
    });

    registration = token.Register(() =>
    {
        timer!.Dispose();
        task.SetCanceled(token);
    });

    timer.Change(delay, Timeout.InfiniteTimeSpan);

    return task;
}

Two events race to complete one task, and Complete ignores the loser. Nothing polls, and no thread waits. That is the shape of every real cancellable I/O operation: the registration tells the driver to abandon the pending read, and whichever of "data arrived" and "give up" happens first is the one that completes the task.

Note what MyTask itself did not gain: a token. Cancellation is not a property of the task, it is a parameter threaded into the operations that produce one. That is why Task has no Cancel() method, and why a CancellationToken you did not pass down cannot cancel anything, no matter how many layers above it was created.

#Running it

The sample exercises every path:

polling loop stopped after 10 iterations
task IsCanceled = True

callback registered before Cancel
callback registered after Cancel, invoked immediately

default token: CanBeCanceled = False, IsCancellationRequested = False

10s delay canceled after 201ms

linked token canceled after 151ms
caller canceled: False, timeout canceled: True

The second block is the one to read twice. Three callbacks were registered; one was unregistered before Cancel and never ran, one ran during Cancel, and one registered after Cancel and ran immediately on the registering thread. All three are the correct behaviour, and all three come from the same eight lines of Register.

The last block shows a ten second delay abandoned after 151 milliseconds because one of two linked sources timed out, with the caller's own source untouched.

#What the real implementation adds

CancellationTokenSource is a few hundred lines rather than a hundred, and the difference is almost entirely contention and lifetime:

  • The registration list is partitioned and lock-free. A single List<T> under a lock is a bottleneck when thousands of concurrent requests each register on a shared token. The real one spreads registrations across per-thread partitions.
  • CancellationTokenRegistration is a struct with a version number, so that disposing a stale registration cannot remove a newer one that reused the slot.
  • Callbacks flow ExecutionContext, captured at registration, and can be posted to the SynchronizationContext captured at registration. MyCancellationTokenSource runs them on whichever thread called Cancel.
  • Dispose waits for a running callback instead of returning immediately, and detects the self-disposal case rather than deadlocking.
  • Token throws ObjectDisposedException after disposal, so handing out a token that can never fire is an error rather than a silent bug.
  • CancellationTokenSource.CreateLinkedTokenSource has a fast path for one and two tokens, because those are almost all the real calls.

The model does not change. A flag, a list of callbacks, and a rule about who runs them.

#What to take away

  • Cancellation is cooperative because the alternative, injecting an exception into arbitrary code, left processes in unreasonable states. Thread.Abort is gone.
  • The source requests and the token observes. Passing a token grants the right to know, not the right to cancel.
  • The token is a struct over one reference, which is why default is a valid never-cancelled token and why CanBeCanceled is worth checking.
  • The flag is volatile so a polling loop cannot cache it, and one-way so an observed cancellation stays true.
  • Register covers the operations that cannot poll. Registering after cancellation must run the callback inline, and that branch is the whole race.
  • Callbacks run outside the lock, all of them run even if one throws, and registrations must be disposed.
  • A linked source leaves a callback on every parent. Not disposing it is the classic leak.
  • Cancellation is a parameter threaded through operations, not a feature of the task. That is why there is no Task.Cancel().

The next post returns to the compiler's machinery: the async method builder, the piece that would let MyTask be returned from an async method, and where the allocations in a real async method actually come from.

#Additional resources

Do you have a question or a suggestion about this post? Contact me!

Follow me:
Enjoy this blog?