Feature request: Extended ICancellationStrategy cancellation strategy

[simon-curtis]

Reason for Change Proposal

We are using this library to remotely control our test machines. However, the current implementation of ICancellationStrategy is quite restrictive.

We would like to bind the IHostApplicationLifetime.ApplicationStopping cancellation token to requests in order to perform a graceful shutdown, rather than abruptly closing the WebSocket connection.

After experimenting extensively with different configurations, I have been unable to find a viable workaround within the current implementation.

At present, the JsonRpc constructor explicitly creates a StandardCancellationStrategy, which then registers the $/cancelRequest method endpoint. This design prevents a custom implementation of ICancellationStrategy from being used while maintaining the same cancellation logic.

Attempted Workaround

I built a wrapper around the existing ICancellationStrategy to track request IDs independently, so that requests could be cancelled gracefully when the host application stops:

public class CancellationStrategyWrapper : ICancellationStrategy
{
    private readonly ICancellationStrategy _currentStrategy;
    private readonly ConcurrentHashSet<RequestId> _requestIds = [];

    public CancellationStrategyWrapper(
        ICancellationStrategy currentStrategy,
        JsonRpc rpc,
        CancellationToken stoppingToken)
    {
        _currentStrategy = currentStrategy;
        stoppingToken.Register(() =>
        {
            foreach (var requestId in _requestIds)
            {
                _ = rpc.NotifyWithParameterObjectAsync("$/cancelRequest", new { id = requestId });
            }
        });
    }

    public void CancelOutboundRequest(RequestId requestId) =>
        _currentStrategy.CancelOutboundRequest(requestId);

    public void OutboundRequestEnded(RequestId requestId) =>
        _currentStrategy.OutboundRequestEnded(requestId);

    public void IncomingRequestStarted(RequestId requestId, CancellationTokenSource cancellationTokenSource)
    {
        _requestIds.Add(requestId);
        _currentStrategy.IncomingRequestStarted(requestId, cancellationTokenSource);
    }

    public void IncomingRequestEnded(RequestId requestId)
    {
        _requestIds.Remove(requestId);
        _currentStrategy.IncomingRequestEnded(requestId);
    }
}

Unexpectedly, this caused StandardCancellationStrategy.inboundCancellationSources to be populated correctly during IncomingRequestStarted, but to appear empty in the CancelInboundRequest method.

Suggested Fix

Move the construction of StandardCancellationStrategy into the JsonRpc.CancellationStrategy property as a nullable assignment.

This would allow developers to override the default behaviour before the method is reserved, without breaking existing functionality:

// JsonRpc.cs

public class JsonRpc
{
    public JsonRpc(IJsonRpcMessageHandler messageHandler)
    {
        // ...
-       this.CancellationStrategy = new StandardCancellationStrategy(this);
    }
    
    public ICancellationStrategy? CancellationStrategy
    {
-       get => this.cancellationStrategy;
+       get => this.cancellationStrategy ??= new StandardCancellationStrategy(this);
        set
        {
            this.ThrowIfConfigurationLocked();
            this.cancellationStrategy = value;
        }
    }
}

Possible Alternative Solutions

1. Adding a CancelAllRequests method to ICancellationStrategy, allowing the user to cancel from outside of the strategy. (Doesn't solve the method name clash for own implementation)
2. Adding a RegisterExternalCancellationToken to StandardCancellationStrategy, registering the cancellation for the user with a convenience method. (Doesn't solve the method name clash for own implementation)
3. Converting to a builder pattern allowing for builder.WithCancellationStrategy<TStrategy>(), this would allow for much more expressive configuration.

Current Impact

This limitation is currently blocking a pending release.
As a temporary workaround, I’m passing the CancellationToken into each concrete type’s constructor and creating linked sources for every method, which not ideal but will do for the time.

I am happy to contribute more to this should you want to go down the alternative solutions route.

[AArnott]

Seems minimal and safe. Thanks for finding such a change that unblocks you.

[AArnott]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[AArnott]

@simon-curtis I am not sure when 2.24 will release to nuget.org. If you'd like to get this fix sooner, would you care to retarget your PR to our v2.23 branch?

[AArnott]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).