Task cancellation + remote target may result in a hung RPC call.

Andrew, this is Ricardo from Dell. I have an issue for you.

Use Case:

• Using dynamic proxies.
• Use the RemoteRpcTarget functionality and have one 'client' invoke a different 'client'
• Implement different types of methods with CancellationToken support. (Types of methods discussed below).

Target Code

The bug I'm discussing is in the following code segment in JsonRpc.DispatchIncomingRequestAsync().

	// Before we forward the request to the remote targets, we need to change the request ID so it gets a new ID in case we run into collisions.  For example,
	// if origin issues a request destined for the remote target at the same time as a request issued by the relay to the remote target, their IDs could be mixed up.
	// See InvokeRemoteTargetWithExistingId unit test for an example.
	RequestId previousId = request.RequestId;
	
	JsonRpcMessage? remoteResponse = null;
	foreach (JsonRpc remoteTarget in remoteRpcTargets)
	{
	    if (request.IsResponseExpected)
	    {
	        request.RequestId = remoteTarget.CreateNewRequestId();
	    }
	
	    CancellationToken token = request.IsResponseExpected ? localMethodCancellationSource!.Token : CancellationToken.None;
	    remoteResponse = await remoteTarget.InvokeCoreAsync(request, null, token).ConfigureAwait(false);
	
	    if (remoteResponse is JsonRpcError error && error.Error != null)
	    {
	        if (error.Error.Code == JsonRpcErrorCode.MethodNotFound || error.Error.Code == JsonRpcErrorCode.InvalidParams)
	        {
	            // If the result is an error and that error is method not found or invalid parameters on the remote target, then we continue on to the next target.
	            continue;
	        }
	    }
	
	    // Otherwise, we simply return the json response;
	    break;
	}
	
	if (remoteResponse != null)
	{
	    if (remoteResponse is IJsonRpcMessageWithId messageWithId)
	    {
	        messageWithId.RequestId = previousId;
	    }
	
	    return remoteResponse;
	}

Working Scenario

There is a particular type of "cancellation" which works fine. If the target method is defined as:

        public Task CancellableAsync(int secondsToDelay, CancellationToken cancellationToken)
        {
            return Task.Run(() =>
            {
                for (int i = 0; i < secondsToDelay; i++)
                {
                    Thread.Sleep(1000);
                    cancellationToken.ThrowIfCancellationRequested();
                }
            });
        }

And the cancellation token is cancelled during the test. In JsonRpc.DispatchIncomingRequestAsync() this results in a remoteResponse which is a JsonRpcError. This gets handled correctly and the last bit of the code replaces the RequestId with the original one to return the reply to to the calling client. All works well.

Broken Scenario

Now, if the target method is defined as:

        public Task CancellableAsync(int secondsToDelay, CancellationToken cancellationToken)
        {
            return Task.Run(() =>
            {
                for (int i = 0; i < secondsToDelay; i++)
                {
                    Thread.Sleep(1000);
                    cancellationToken.ThrowIfCancellationRequested();
                }
            }, cancellationToken);
        }

or

        public async Task<int> CancellableWithAwaitAsync(int secondsToDelay, CancellationToken cancellationToken)
        {
            await Task.Delay(secondsToDelay * 1000, cancellationToken);
            return secondsToDelay;
        }

If these methods are called and cancelled, the line in JsonRpc.DispatchIncomingRequestAsync()

	    remoteResponse = await remoteTarget.InvokeCoreAsync(request, null, token).ConfigureAwait(false);

raises a TaskCancelled exception which is not handled here. It is caught by the "catch all" at the end and turned into a JsonRpcError which is returned. BUT, the ID is never set to the previous ID. So the message returned to the original client has the incorrect ID. This leads to a hung method that never returns.
Note that I believe the big difference between the two scenarios is that the last one results in a cancelled Task while the first, working, scenario is just a Task that threw an exception.

My Mitigation

I mitigated the problem on my end with this change:

	// Before we forward the request to the remote targets, we need to change the request ID so it gets a new ID in case we run into collisions.  For example,
	// if origin issues a request destined for the remote target at the same time as a request issued by the relay to the remote target, their IDs could be mixed up.
	// See InvokeRemoteTargetWithExistingId unit test for an example.
	RequestId previousId = request.RequestId;
	
	JsonRpcMessage? remoteResponse = null;
	foreach (JsonRpc remoteTarget in remoteRpcTargets)
	{
	    if (request.IsResponseExpected)
	    {
	        request.RequestId = remoteTarget.CreateNewRequestId();
	    }
	
	    try
	    {
	        CancellationToken token = request.IsResponseExpected ? localMethodCancellationSource!.Token : CancellationToken.None;
	        remoteResponse = await remoteTarget.InvokeCoreAsync(request, null, token).ConfigureAwait(false);
	
	        if (remoteResponse is JsonRpcError error && error.Error != null)
	        {
	            if (error.Error.Code == JsonRpcErrorCode.MethodNotFound || error.Error.Code == JsonRpcErrorCode.InvalidParams)
	            {
	                // If the result is an error and that error is method not found or invalid parameters on the remote target, then we continue on to the next target.
	                continue;
	            }
	        }
	    }
	    catch (Exception ex)
	    {
	        JsonRpcError error = this.CreateError(request, ex);
	
	        if (error.Error != null && JsonRpcEventSource.Instance.IsEnabled(System.Diagnostics.Tracing.EventLevel.Warning, System.Diagnostics.Tracing.EventKeywords.None))
	        {
	            JsonRpcEventSource.Instance.SendingError(request.RequestId.NumberIfPossibleForEvent, error.Error.Code);
	        }
	
	        remoteResponse = error;
	    }
	
	    // Otherwise, we simply return the json response;
	    break;
	}
	
	if (remoteResponse != null)
	{
	    if (remoteResponse is IJsonRpcMessageWithId messageWithId)
	    {
	        messageWithId.RequestId = previousId;
	    }
	
	    return remoteResponse;
	}

But I figured you guys might want to deal with the fact that one version results in a JsonRpcError while the other version results in an exception.

Start at JsonRpc.DispatchIncomingRequestAsync, especially the remoteTarget.InvokeCoreAsync call and the existing InvokeRemoteTargetWithExistingId unit test. Reproduce cancellation for both Task-cancellation forms described in the issue and trace how the response ID is handled. Done means a cancelled remote call returns to the original client without hanging and preserves the original request ID.