Creating a Small Helper Class to Safely Handle CancellationTokenSource Cancel and Dispose

Introduction

When writing async code in Unity, you often end up working with CancellationToken for loading, networking, Addressables, screen transitions, timeouts, and so on.

The object you usually create for that is CancellationTokenSource, but it is surprisingly easy to handle it incorrectly.

For example:

private CancellationTokenSource? _cts;

private void CancelAndDispose()
{
    _cts?.Cancel();
    _cts?.Dispose();
}

At first glance, this looks fine. However, the disposed CancellationTokenSource remains in the field.

If CancelAndDispose() is called again, the second Cancel() may throw an ObjectDisposedException.

Also, when using CancelAfter() or CreateLinkedTokenSource(), disposing the CTS after use becomes even more important.

This article explains the difference between Cancel() and Dispose(), then introduces a small helper class that safely handles CancelAfter() and CreateLinkedTokenSource() as well.

Cancel and Dispose Are Different Things

First, Cancel() and Dispose() do different jobs.

Cancel() sends a cancellation request to code that has received a CancellationToken.

var cts = new CancellationTokenSource();

var task = LoadAsync(cts.Token);

cts.Cancel();

This is only a cancellation request.

Calling Cancel() does not forcibly stop a running operation. The callee must either check the CancellationToken or pass it to an API that supports cancellation.

On the other hand, Dispose() is not a cancellation request.

cts.Dispose();

Dispose() releases resources held internally by CancellationTokenSource.

In short:

Cancel()
    Sends a cancellation request

Dispose()
    Releases resources

Calling Cancel() does not automatically call Dispose(). Depending on the situation, you may need both.

cts.Cancel();
cts.Dispose();

However, that does not mean you should blindly write those two lines everywhere.

Disposed Does Not Mean Null

Let's look at this code again.

private CancellationTokenSource? _cts;

private void CancelAndDispose()
{
    _cts?.Cancel();
    _cts?.Dispose();
}

Dispose() does not set _cts to null.

That means the following call order can be dangerous:

CancelAndDispose();
CancelAndDispose();

Even after _cts.Dispose() is called the first time, _cts is still not null.

So the second call still executes _cts?.Cancel().

CancellationTokenSource.Cancel() may throw ObjectDisposedException if the target CTS has already been disposed.

In Unity, cleanup code can be called through multiple paths.

public void Close()
{
    CancelAndDispose();
}

private void OnDestroy()
{
    CancelAndDispose();
}

You could catch and ignore ObjectDisposedException, but that is not the real fix.

The real problem is that a disposed CTS remains in the field.

Instead of hiding the exception, it is safer to structure the code so that a disposed CTS is never touched again.

Dispose Is Important with CancelAfter and CreateLinkedTokenSource

CancelAfter() schedules a cancellation request after a specified delay.

using var cts = new CancellationTokenSource();
cts.CancelAfter(TimeSpan.FromSeconds(10));

await LoadAsync(cts.Token);

This is commonly used for timeouts.

CancelAfter() is convenient, but the CTS needs internal machinery to cancel itself later. Because of that, it is important to dispose the CTS when you are done with it.

The same applies to CreateLinkedTokenSource().

For example, you may want to combine multiple cancellation conditions, such as "the MonoBehaviour was destroyed" and "the operation timed out".

using var timeoutCts = new CancellationTokenSource();
timeoutCts.CancelAfter(TimeSpan.FromSeconds(10));

using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(
    destroyCancellationToken,
    timeoutCts.Token);

await LoadAsync(linkedCts.Token);

In this example, both timeoutCts and linkedCts should be disposed.

CreateLinkedTokenSource() creates a new CancellationTokenSource that becomes canceled when any of the provided tokens is canceled.

So this newly created CTS is also something you should dispose after use.

If you leave a CTS created with CancelAfter() or CreateLinkedTokenSource() undisposed, the release of internal resources and registrations may be delayed.

Instead of saying this always causes an immediate memory leak, it is more accurate to say that resources owned by the finished CTS may remain alive longer than necessary.

Managing CTS Fields Directly Gets Messy

If everything is contained inside a single method, using is often enough.

using var cts = new CancellationTokenSource();
cts.CancelAfter(TimeSpan.FromSeconds(10));

await LoadAsync(cts.Token);

However, in game code, you often want to keep a CTS in a field.

For example, when restarting a load operation, you may want to cancel the previous load first.

private CancellationTokenSource? _loadCts;

public async void Reload()
{
    _loadCts?.Cancel();
    _loadCts?.Dispose();

    _loadCts = new CancellationTokenSource();
    var token = _loadCts.Token;

    try
    {
        await LoadAsync(token);
    }
    catch (OperationCanceledException) when (token.IsCancellationRequested)
    {
        // Cancellation is treated as a normal flow.
    }
}

This code mixes together cancellation of the previous CTS, disposal, creation of a new CTS, timeout setup, linked token creation, and protection against double cleanup.

Writing that every time is annoying.

And annoying code tends to get copied incorrectly sooner or later.

So instead, we can put ownership of the CancellationTokenSource into a dedicated helper class.

Creating a Helper Class

The policy for this class is:

- Return only CancellationToken to the caller
- Let the helper class create, replace, cancel, and dispose CTS instances
- Do not leave a disposed CTS in the field
- End the old CTS when Renew / RenewOwned is called
- Dispose CTS instances created by CancelAfter and CreateLinkedTokenSource
- By default, do not throw ObjectDisposedException / AggregateException from Cancel()
- If onCancelException is specified, report the exception there before swallowing it
- Allow a strict mode that propagates Cancel exceptions when needed

Cancel() may throw AggregateException if a registered cancellation callback throws.

If an old CTS throws during Cancel() inside RenewOwned(), the caller sees RenewOwned() as failed.

However, depending on the implementation order, the new CTS may already have become active internally.

For this reason, this helper class does not propagate exceptions caused by Cancel() by default. If onCancelException is specified, it reports the exception there.

If onCancelException is not specified, the exception is simply swallowed. In a real project, it is recommended to pass a logging function such as Debug.LogException.

Implementation

#nullable enable

using System;
using System.Threading;

public sealed class CancellationTokenSourceOwner : IDisposable
{
    private readonly object _gate = new();
    private readonly bool _throwOnCancelException;
    private readonly Action<Exception>? _onCancelException;

    private CancellationTokenSource? _cts;
    private bool _disposed;

    public CancellationTokenSourceOwner(
        bool throwOnCancelException = false,
        Action<Exception>? onCancelException = null)
    {
        _throwOnCancelException = throwOnCancelException;
        _onCancelException = onCancelException;
    }

    public CancellationToken Renew()
    {
        return RenewOwned(new CancellationTokenSource());
    }

    public CancellationToken RenewAfter(TimeSpan delay)
    {
        var next = new CancellationTokenSource();

        try
        {
            next.CancelAfter(delay);
        }
        catch
        {
            next.Dispose();
            throw;
        }

        return RenewOwned(next);
    }

    public CancellationToken RenewLinked(params CancellationToken[] tokens)
    {
        var next = CancellationTokenSource.CreateLinkedTokenSource(tokens);
        return RenewOwned(next);
    }

    public CancellationToken RenewLinkedAfter(
        TimeSpan delay,
        params CancellationToken[] tokens)
    {
        var next = CancellationTokenSource.CreateLinkedTokenSource(tokens);

        try
        {
            next.CancelAfter(delay);
        }
        catch
        {
            next.Dispose();
            throw;
        }

        return RenewOwned(next);
    }

    /// <summary>
    /// Takes ownership of next.
    /// After calling this method, the caller must not dispose next.
    /// </summary>
    public CancellationToken RenewOwned(CancellationTokenSource next)
    {
        if (next == null)
            throw new ArgumentNullException(nameof(next));

        CancellationTokenSource? old;
        CancellationToken token;

        lock (_gate)
        {
            if (_disposed)
            {
                next.Dispose();
                throw new ObjectDisposedException(nameof(CancellationTokenSourceOwner));
            }

            token = next.Token;
            old = _cts;
            _cts = next;
        }

        CancelAndDisposeCore(old);

        return token;
    }

    public void CancelAndDispose()
    {
        CancellationTokenSource? old;

        lock (_gate)
        {
            old = _cts;
            _cts = null;
        }

        CancelAndDisposeCore(old);
    }

    public void Dispose()
    {
        CancellationTokenSource? old;

        lock (_gate)
        {
            if (_disposed)
                return;

            _disposed = true;
            old = _cts;
            _cts = null;
        }

        CancelAndDisposeCore(old);
    }

    private void CancelAndDisposeCore(CancellationTokenSource? cts)
    {
        if (cts == null)
            return;

        try
        {
            try
            {
                cts.Cancel();
            }
            catch (Exception ex) when (ShouldHandleCancelException(ex))
            {
                ReportCancelException(ex);
            }
        }
        finally
        {
            cts.Dispose();
        }
    }

    private bool ShouldHandleCancelException(Exception ex)
    {
        if (_throwOnCancelException)
            return false;

        return ex is ObjectDisposedException || ex is AggregateException;
    }

    private void ReportCancelException(Exception ex)
    {
        var handler = _onCancelException;
        if (handler == null)
            return;

        try
        {
            handler(ex);
        }
        catch
        {
            // Do not let logging failures break cleanup.
        }
    }
}

Implementation Notes

RenewOwned Transfers Ownership

The actual CTS replacement logic is centralized in RenewOwned().

Renew(), RenewAfter(), RenewLinked(), and RenewLinkedAfter() are convenience methods for different use cases, and they all call RenewOwned() internally.

var next = new CancellationTokenSource();
next.CancelAfter(TimeSpan.FromSeconds(5));

var token = _owner.RenewOwned(next);

At this point, ownership of next moves to _owner.

The caller must not call next.Dispose() after passing it to RenewOwned().

Change State Inside the Lock, Cancel Outside the Lock

In RenewOwned(), the class removes the old CTS and activates the new CTS inside the lock.

However, Cancel() itself is called outside the lock.

The reason is that Cancel() may execute registered cancellation callbacks synchronously.

If the owner lock is held while calling Cancel(), the design can become complicated depending on what those callbacks do.

The order is:

RenewOwned:
    Activate the new CTS inside the lock
    Exit the lock
    Cancel and dispose the old CTS

CancelAndDispose / Dispose:
    Remove the active CTS inside the lock
    Exit the lock
    Cancel and dispose the removed CTS

There is one important caveat.

RenewOwned() activates the new CTS before canceling the old CTS.

That means this class does not support reentrant usage where a cancellation callback from the old token calls Renew(), CancelAndDispose(), or Dispose() on the same owner.

Cancellation callbacks should avoid manipulating the owner itself and should be kept to minimal notifications or logging.

Cancel Exceptions Are Not Propagated by Default

CancelAndDisposeCore() always calls Dispose() even if Cancel() throws.

try
{
    try
    {
        cts.Cancel();
    }
    catch (Exception ex) when (ShouldHandleCancelException(ex))
    {
        ReportCancelException(ex);
    }
}
finally
{
    cts.Dispose();
}

By default, ObjectDisposedException and AggregateException caused by Cancel() are not propagated.

If onCancelException is specified, the exception is reported there and then swallowed.

In Unity, this lets you avoid throwing from cleanup code while still keeping a log.

private readonly CancellationTokenSourceOwner _loadToken = new CancellationTokenSourceOwner(
    onCancelException: ex => Debug.LogException(ex));

If ObjectDisposedException is reported, check whether the CTS passed to RenewOwned(next) was disposed by the caller, or whether the same CTS is being shared outside the owner.

If you want to treat cancellation callback exceptions as bugs, enable strict mode.

private readonly CancellationTokenSourceOwner _loadToken = new CancellationTokenSourceOwner(
    throwOnCancelException: true);

In strict mode, RenewOwned() may throw even though the new CTS has already become active internally.

For Unity loading and screen transition code, passing a logging function to onCancelException and using the non-throwing default is usually easier to handle.

Usage

The basic method is Renew().

private readonly CancellationTokenSourceOwner _loadToken = new CancellationTokenSourceOwner(
    onCancelException: ex => Debug.LogException(ex));

public async void Reload()
{
    var token = _loadToken.Renew();

    try
    {
        await LoadAsync(token);
    }
    catch (OperationCanceledException) when (token.IsCancellationRequested)
    {
        // Cancellation is treated as a normal flow.
    }
    catch (Exception ex)
    {
        Debug.LogException(ex);
    }
}

Calling Renew() creates a new CTS. If an old CTS exists, it is canceled and then disposed.

Use RenewAfter() when you need a timeout.

var token = _loadToken.RenewAfter(TimeSpan.FromSeconds(10));

Use RenewLinked() when you need to link with an external token.

var token = _loadToken.RenewLinked(destroyCancellationToken);

Use RenewLinkedAfter() when you need both a timeout and an external token.

var token = _loadToken.RenewLinkedAfter(
    TimeSpan.FromSeconds(10),
    destroyCancellationToken);

This cancels when either the MonoBehaviour is destroyed or 10 seconds pass.

Internally, this uses both CreateLinkedTokenSource() and CancelAfter(), but disposal of the CTS remains inside CancellationTokenSourceOwner.

For cleanup:

public void Close()
{
    _loadToken.CancelAndDispose();
}

private void OnDestroy()
{
    _loadToken.Dispose();
}

CancelAndDispose() removes the active CTS first, then calls Cancel() and Dispose().

So even if it is called multiple times, the second and later calls have no active CTS to touch.

Dispose() ends the lifetime of the owner itself. Calling Renew() after Dispose() throws ObjectDisposedException.

Use CancelAndDispose() when you want to stop the current operation but keep the owner reusable. Use Dispose() when the owner itself is done.

Combining with destroyCancellationToken

Unity 2022.2 and later provides MonoBehaviour.destroyCancellationToken.

destroyCancellationToken is managed by Unity, so you do not dispose it yourself.

However, a linked CTS created by CreateLinkedTokenSource() is created by your code, so it must be disposed.

When using this helper class, the mental model becomes:

destroyCancellationToken
    Token managed by Unity
    Do not dispose it yourself

RenewLinked(destroyCancellationToken)
    Owner creates a linked CTS
    Owner disposes it

RenewLinkedAfter(timeout, destroyCancellationToken)
    Owner creates a linked CTS
    Owner sets CancelAfter
    Owner disposes it

Caveats

CancellationTokenSourceOwner is not a universal cancellation management system.

Use it with the following assumptions in mind.

Do Not Reenter the Owner from Cancellation Callbacks

This class does not support calling Renew(), CancelAndDispose(), or Dispose() on the same owner from a cancellation callback of the old token.

RenewOwned() activates the new CTS before canceling the old CTS.

If a callback from the old token manipulates the owner, it may affect the new CTS.

Renew / CancelAndDispose May Be Synchronously Expensive

Cancel() may run cancellation callbacks synchronously.

That means Renew(), CancelAndDispose(), and Dispose() may not return until callbacks of the old CTS have finished.

If you call these methods on the Unity main thread, heavy cancellation callbacks can cause frame drops.

The Old CTS Is Disposed Immediately After Cancel

This class disposes the old CTS immediately after calling Cancel().

For normal async / await code and APIs that accept CancellationToken, this is usually convenient.

However, it does not work well with code that touches token.WaitHandle after cancellation.

var token = _loadToken.Renew();

_loadToken.CancelAndDispose();

// Not recommended:
// The old CTS has already been disposed, so this design does not work well with WaitHandle.
var handle = token.WaitHandle;

If you pass tokens to old APIs or custom implementations that use WaitHandle, consider a design where the CTS is disposed only after the operation has fully completed.

Always Handle Exceptions from async void

Unity event handlers sometimes use async void.

However, exceptions escaping from async void are hard to handle, so it is safer to explicitly catch both OperationCanceledException and normal exceptions.

try
{
    await LoadAsync(token);
}
catch (OperationCanceledException) when (token.IsCancellationRequested)
{
    // Cancellation is normal flow.
}
catch (Exception ex)
{
    Debug.LogException(ex);
}

Outside Unity event handlers, prefer async Task when possible so the caller can await the operation.

Summary

CancellationTokenSource is not finished just because you called Cancel().

Cancel()
    Sends a cancellation request

Dispose()
    Releases resources

These are separate responsibilities.

Especially when using CancelAfter() or CreateLinkedTokenSource(), failing to dispose the CTS after use may delay the release of internal resources and registrations.

Also, when keeping CTS directly in fields, it is easy to leave a disposed CTS behind, call Cancel() on an already disposed CTS, or forget to dispose CTS instances created for timeouts or linked tokens.

Putting ownership of the CTS into a dedicated class makes the code easier to reason about.

The caller receives only a CancellationToken, while creation, cancellation, disposal, renewal, linked tokens, and timeouts are handled by the owner class.

If you create a CancellationTokenSource yourself, you are responsible for disposing it.
But instead of repeating that logic in every class, put ownership in one place.

With this structure, Unity loading and networking code becomes easier to handle safely.

推荐订阅源

DEV Community