Add CancellationTokenRegistration.OnUnregisteredWithoutInvoke(Action) API

[timcassell]

Currently, I have no way of knowing if/when my callback registered via CancellationToken.Register(Action) was removed from the token without it being invoked without me explicitly removing it myself via CancellationTokenRegistration.Unregister. Like when the source is disposed without Cancel having ever been called, or if the new TryReset is called.

This would be useful for cleaning up resources that have no definite expiration aside from the life of the registration.

Adjusting the MSDN example, I could do this:

bool wasCanceled = false;
CancellationTokenRegistration ctr = token.Register(() => wasCanceled = true);
using (ctr)
{
    await SomeAsyncTask();
}
if (!wasCanceled)
{
    // cleanup
}

But this won't work if the object I need to clean up is relying entirely on the lifetime of the registered action, not some other async action.

public class MyCustomToken
{
    private CancellationTokenRegistration registration;
    
    public class MyCustomToken(CancellationToken token)
    {
        registration = token.Register(OnCanceled);
        registration.OnUnregisteredWithoutInvoke(OnUnregisteredWithoutInvoke);
    }
    
    private void OnCanceled()
    {
        DoCustomWork();
        Cleanup();
    }
    
    private void DoCustomWork()
    {
        // Do my custom thing.
    }
    
    private void OnUnregisteredWithoutInvoke()
    {
        Cleanup();
    }
    
    private void Cleanup()
    {
        // Custom Cleanup
        registration.Dispose();
    }
}

[huoyaoyuan]

You can open an issue with api-proposal category to increase visibility.

[timcassell]

@stephentoub #60843 was locked, so I thought I should post here rather than create a new issue (this is what I was trying to do after all).

Trying to hook up my custom cancelation token to BCL token, I'm stuck on 2 parts.

1. How can I cache the token to reregister the cancelation while also allowing the underlying source to be collected? Like a WeakReference, but I can't use it because it's a struct. And I'm pretty sure caching it directly on a field in a resurrecting object won't work like I did here, but I don't know enough about the GC to say if it does or not.

2. How can I tell if the token's source has been disposed or finalized so I can skip the reregistration and resurrection? I only want to reregister the cancelation if TryReset was the cause of the finalizer being invoked.

For both of these, I had the thought to use reflection to extract the wrapped source from the token to be able to use it in a WeakReference, and its Token property throws ObjectDisposedException which I could catch to detect if it's disposed. But that's relying on an implementation detail, and I need a portable solution.

partial class Extensions
{
    /// <summary>
    /// Convert <paramref name="token"/> to a <see cref="CancelationToken"/>.
    /// </summary>
    /// <param name="token">The cancellation token to convert</param>
    /// <returns>A <see cref="CancelationToken"/> that will be canceled when <paramref name="token"/> is canceled.</returns>
    public static CancelationToken ToCancelationToken(this System.Threading.CancellationToken token)
    {
        if (!token.CanBeCanceled)
        {
            return default(CancelationToken);
        }
        if (token.IsCancellationRequested)
        {
            return CancelationToken.Canceled();
        }

        return new TokenWrapper(Internal.CancelationRef.GetOrCreate(), token).Token;
    }

    private sealed class TokenWrapper
    {
        // Is it even possible for the internal CancellationTokenSource to be finalized if we have a field here and this is resurrected?
        // We can't use WeakReference because it's a value type.
        internal readonly System.Threading.CancellationToken _bclToken;

        private readonly Internal.CancelationRef _cancelationRef;
        private readonly int _tokenId;

        internal CancelationToken Token { get { return new CancelationToken(_cancelationRef, _tokenId); } }

        internal TokenWrapper(Internal.CancelationRef cancelationRef, System.Threading.CancellationToken token)
        {
            _bclToken = token;
            _cancelationRef = cancelationRef;
            _tokenId = cancelationRef.TokenId;
            RegisterCancelation();
        }

        ~TokenWrapper()
        {
            // How can we tell if the token was disposed/finalized vs TryReset?
            // CanBeCanceled obviously doesn't work.
            if (!_bclToken.CanBeCanceled)
            {
                _cancelationRef.Dispose();
                return;
            }

            GC.ReRegisterForFinalize(this);
            _cancelationRef.ClearListeners();
            RegisterCancelation();
        }

        private void RegisterCancelation()
        {
            _bclToken.Register(state =>
            {
                GC.SuppressFinalize(state);
                var wrapper = Unsafe.As<TokenWrapper>(state);
                wrapper._cancelationRef.Cancel();
            }, this, false);
        }
    }
}

[timcassell]

And even if I do extract the source and wrap it in a WeakReference, I have no way of detecting the case where the source has TryReset called, then Cancel, then is collected before my finalizer is called. I would try to get the source from the WeakReference, but it would no longer exist, so I couldn't tell if it was canceled to cancel my token.

~TokenWrapper()
{
    System.Threading.CancellationTokenSource source;
    if (!_weakBclTokenSource.TryGetTarget(out source))
    {
        // It is possible the source had TryReset() then Cancel() called before being collected, and we have no way of knowing at this point. :(
        _cancelationRef.Dispose();
        return;
    }

    if (source.IsCancellationRequested)
    {
        _cancelationRef.Cancel();
        return;
    }

    System.Threading.CancellationToken token;
    // If the source was disposed, the Token property will throw ObjectDisposedException.
    try
    {
        token = source.Token;
    }
    catch (ObjectDisposedException)
    {
        _cancelationRef.Dispose();
        return;
    }

    GC.ReRegisterForFinalize(this);
    RegisterCancelation(token);
}

[timcassell]

I actually think I may have figured it out. I forgot the GC still works with circular references, so a strong reference inside the wrapper should still allow the source to be collected, as long as the wrapper is never referenced. So I can just use the wrapper to reregister if TryReset is called, and the CancelationRef's own finalizer to handle cleanup.

Are there any logic holes here that I'm overlooking?

partial class Extensions
{
    private static readonly bool s_canExtractSource = GetCanExtractSource();

    private static bool GetCanExtractSource()
    {
        // This assumes the CancellationToken is implemented like this, and will return false if it's different.
        // public struct CancellationToken
        // {
        //     private CancellationTokenSource m_source;
        //     ...
        // }
        var fields = typeof(System.Threading.CancellationToken).GetFields(System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Public | System.Reflection.BindingFlags.Instance);
        return fields.Length == 1 && typeof(System.Threading.CancellationTokenSource).IsAssignableFrom(fields[0].FieldType);
    }

    // Implementation detail, the token wraps the source, so we can retrieve it by placing it in this explicit layout struct and reading the source.
    // This is equivalent to `Unsafe.As`, but also works in older runtimes that don't support Unsafe. It's also more efficient than using reflection (and some runtimes don't support TypedReference).
    // I think it is very unlikely, but the internal implementation of CancellationToken could change in the future (or different runtime) to break this code, which is why we have the s_canExtractSource check.
    [System.Runtime.InteropServices.StructLayout(System.Runtime.InteropServices.LayoutKind.Explicit)]
    private struct TokenSourceExtractor
    {
        [System.Runtime.InteropServices.FieldOffset(0)]
        internal System.Threading.CancellationToken _token;
        [System.Runtime.InteropServices.FieldOffset(0)]
        internal System.Threading.CancellationTokenSource _source;
    }

    /// <summary>
    /// Convert <paramref name="token"/> to a <see cref="CancelationToken"/>.
    /// </summary>
    /// <param name="token">The cancellation token to convert</param>
    /// <returns>A <see cref="CancelationToken"/> that will be canceled when <paramref name="token"/> is canceled.</returns>
    public static CancelationToken ToCancelationToken(this System.Threading.CancellationToken token)
    {
        if (!s_canExtractSource)
        {
            throw new System.Reflection.TargetException("Cannot convert System.Threading.CancellationToken to Proto.Promises.CancelationToken due to an implementation change. Please notify the developer.");
        }

        if (!token.CanBeCanceled)
        {
            return default(CancelationToken);
        }
        if (token.IsCancellationRequested)
        {
            return CancelationToken.Canceled();
        }

        // This relies on internal implementation details. If the implementation changes, the s_canExtractSource check should catch it.
        var source = new TokenSourceExtractor() { _token = token }._source;

        if (source == null)
        {
            // Source should never be null if token.CanBeCanceled returned true.
            throw new System.Reflection.TargetException("The token's internal source was null.");
        }

        var cancelationRef = Internal.CancelationRef.GetOrCreateForBclTokenConvert(source);
        var tokenId = cancelationRef.TokenId;
        HookupCancelation(cancelationRef, source, token);
        return new CancelationToken(cancelationRef, tokenId);
    }

    static void HookupCancelation(Internal.CancelationRef cancelationRef, System.Threading.CancellationTokenSource source, System.Threading.CancellationToken token)
    {
        _ = new Reregister(cancelationRef, source, token);
    }

    // The purpose of this class is to re-register the cancel callback if `CancellationTokenSource.TryReset()` is called, but not if the source is disposed or garbage collected.
    private sealed class Reregister
    {
        // A strong reference to the source forces it to stay alive until this is collected so that we can check its canceled/disposed state inside our finalizer.
        // This is necessary in case the source had TryReset() then Cancel() and all references except this dropped before our finalizer is called.
        private readonly System.Threading.CancellationTokenSource _source;
        private readonly Internal.CancelationRef _cancelationRef;

        private System.Threading.CancellationTokenRegistration _registration;

        internal Reregister(Internal.CancelationRef cancelationRef, System.Threading.CancellationTokenSource source, System.Threading.CancellationToken token)
        {
            _cancelationRef = cancelationRef;
            _source = source;
            RegisterCancelation(token);
        }

        ~Reregister()
        {
            if (_source.IsCancellationRequested)
            {
                _cancelationRef.Cancel();
                return;
            }

            System.Threading.CancellationToken token;
            // If the source was disposed, the Token property will throw ObjectDisposedException.
            try
            {
                token = _source.Token;
            }
            catch (ObjectDisposedException)
            {
                return;
            }

            // If we are able to unregister, it means the source is ready for garbage collection without canceling, so we shouldn't resurrect this.
            if (!_registration.Unregister())
            {
                GC.ReRegisterForFinalize(this);
                RegisterCancelation(token);
            }
        }

        private void OnCancel()
        {
            GC.SuppressFinalize(this);
            _cancelationRef.Cancel();
        }

        private void RegisterCancelation(System.Threading.CancellationToken token)
        {
            _registration = token.Register(state => state.UnsafeAs<Reregister>().OnCancel(), this, false);
        }
    }
}

[timcassell]

Well, after some testing, that seems to mostly work. The only thing I missed was that my CancelationRef's finalizer can run before the wrapper's finalizer randomly. So to fix that I just place it in a static HashSet to establish a real strong reference, then remove it when I no longer need the wrapper.

[timcassell]

Omg, the realization just hit me like a ton of bricks... When a new registration is added to my own token, I can just call _registration.Unregister() and immediately re-register the cancelation (and the return value can be used to see if I need to remove my registrations). So simple! No need for an extra finalizer at all! Unfortunately, it still doesn't catch the disposed case without catching the exception from the source, and that makes each of my registrations more expensive, but it's better than relying on the timing of the GC.

Of course, an IsRegistered property would be even better to do less work.