Search code examples
c#tasktask-parallel-librarycancellation

How to Create a Hierarchical CancellationTokenSource?


In our project, we have decided to provide a cancellation mechanism for users with the help of CancellationToken.

Because of the structure of the works in the project, I need a hierarchical cancellation mechanism. By hierarchical, I mean that parent source cancellation causes all child sources to be recursively canceled but child sources cancellations are not propagated to the parent.

Is there such an option available in .NET out of the box? If not, I'm not sure whether registering a delegate to the parent token is enough or further considerations should be given.


Solution

  • Based on the implementation in the source code for Linked2CancellationTokenSource, I came to this implementation:

    public class HierarchicalCancellationTokenSource : CancellationTokenSource
    {
        private readonly CancellationTokenRegistration _parentReg;
    
        public HierarchicalCancellationTokenSource(CancellationToken parentToken)
        {
            this._parentReg = parentToken.Register(
                static s => ((CancellationTokenSource)s).Cancel(false),
                this,
                useSynchronizationContext: false);
        }
    
        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                this._parentReg.Dispose();
            }
    
            base.Dispose(disposing);
        }
    }
    

    And a demo:

    CancellationTokenSource[] CreateChildSources(CancellationTokenSource parentSource) =>
        Enumerable.Range(0, 2)
            .Select(_ => new HierarchicalCancellationTokenSource(parentSource.Token))
            .ToArray();
    
    var rootSource = new CancellationTokenSource();
    var childSources = CreateChildSources(rootSource);
    var grandChildSources = childSources.SelectMany(CreateChildSources).ToArray();
    
    var allTokens = new[] { rootSource.Token }
        .Concat(childSources.Select(s => s.Token))
        .Concat(grandChildSources.Select(s => s.Token))
        .ToArray();
    
    for (int i = 0; i < allTokens.Length; i++)
    {
        allTokens[i].Register(
            i => Console.WriteLine(
                $"{new string('+', (int)Math.Log2((int)i))}{i} canceled."),
            i + 1);
    }
    
    rootSource.Cancel();
    
    /* Output:
    1 canceled.
    +3 canceled.
    ++7 canceled.
    ++6 canceled.
    +2 canceled.
    ++5 canceled.
    ++4 canceled.
    */