Skip to content

Commit

Permalink
Add ICancellationStrategy extensibility point
Browse files Browse the repository at this point in the history
This allows specific JsonRpc instances to be configured to implement `CancellationToken` support in a totally different way. They are *not* given the freedom to mutate the the request object itself, but they may trigger and respond to any other JSON-RPC message or do something entirely out of band (e.g. creating a file and/or testing for existing of the file).

As part of this change, we had to move `$/cancelRequest` handling out of its very special place in `JsonRpc.HandleRpcAsync` into a standard method handler. But that would have changed the method handler to be invoked on the user-configurable `SynchronizationContext` (which is now by default order-preserving). But this special method *shouldn't* be subject to that order-preserving `SynchronizationContext` because in fact we want these cancellation requests handled as soon as they come in so they can cancel an ongoing RPC server method that has not yet yielded. In order to support this, I modified how we record `TargetMethod` to include an optional override `SynchronizationContext`.
  • Loading branch information
AArnott committed Sep 3, 2020
1 parent fde90ef commit 28cfabd
Show file tree
Hide file tree
Showing 9 changed files with 436 additions and 149 deletions.
128 changes: 128 additions & 0 deletions src/StreamJsonRpc.Tests/CustomCancellationStrategyTests.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.VisualStudio.Threading;
using Nerdbank.Streams;
using StreamJsonRpc;
using Xunit;
using Xunit.Abstractions;

public class CustomCancellationStrategyTests : TestBase
{
private readonly JsonRpc clientRpc;
private readonly Server server;
private readonly JsonRpc serverRpc;
private readonly MockCancellationStrategy mockStrategy;

public CustomCancellationStrategyTests(ITestOutputHelper logger)
: base(logger)
{
this.mockStrategy = new MockCancellationStrategy(logger);

var streams = FullDuplexStream.CreatePair();
this.clientRpc = new JsonRpc(streams.Item1)
{
CancellationStrategy = this.mockStrategy,
TraceSource = new TraceSource("Client", SourceLevels.Verbose)
{
Listeners = { new XunitTraceListener(logger) },
},
};
this.clientRpc.StartListening();

this.server = new Server();
this.serverRpc = new JsonRpc(streams.Item2)
{
CancellationStrategy = this.mockStrategy,
TraceSource = new TraceSource("Server", SourceLevels.Verbose)
{
Listeners = { new XunitTraceListener(logger) },
},
};
this.serverRpc.AddLocalRpcTarget(this.server);
this.serverRpc.StartListening();
}

/// <summary>
/// Verifies that cancellation can occur through a custom strategy.
/// </summary>
[Fact]
public async Task CancelRequest()
{
using var cts = new CancellationTokenSource();
Task invokeTask = this.clientRpc.InvokeWithCancellationAsync(nameof(Server.NoticeCancellationAsync), cancellationToken: cts.Token);
var completingTask = await Task.WhenAny(invokeTask, this.server.MethodEntered.WaitAsync()).WithCancellation(this.TimeoutToken);
await completingTask; // rethrow an exception if there is one.

cts.Cancel();
await invokeTask.WithCancellation(this.TimeoutToken);
Assert.True(this.mockStrategy.CancelRequestMade);
}

private class Server
{
internal AsyncManualResetEvent MethodEntered { get; } = new AsyncManualResetEvent();

public async Task NoticeCancellationAsync(CancellationToken cancellationToken)
{
this.MethodEntered.Set();
var canceled = new AsyncManualResetEvent();
using (cancellationToken.Register(canceled.Set))
{
await canceled.WaitAsync();
}
}
}

private class MockCancellationStrategy : ICancellationStrategy
{
private readonly Dictionary<RequestId, CancellationTokenSource> cancelableRequests = new Dictionary<RequestId, CancellationTokenSource>();
private readonly ITestOutputHelper logger;

internal MockCancellationStrategy(ITestOutputHelper logger)
{
this.logger = logger;
}

internal bool CancelRequestMade { get; private set; }

public void CancelOutboundRequest(RequestId requestId)
{
this.logger.WriteLine("Canceling outbound request: {0}", requestId);
CancellationTokenSource? cts;
lock (this.cancelableRequests)
{
this.cancelableRequests.TryGetValue(requestId, out cts);
}

cts?.Cancel();
this.CancelRequestMade = true;
}

public void IncomingRequestStarted(RequestId requestId, CancellationTokenSource cancellationTokenSource)
{
this.logger.WriteLine("Recognizing incoming request start: {0}", requestId);
lock (this.cancelableRequests)
{
this.cancelableRequests.Add(requestId, cancellationTokenSource);
}
}

public void IncomingRequestEnded(RequestId requestId)
{
this.logger.WriteLine("Recognizing incoming request end: {0}", requestId);
lock (this.cancelableRequests)
{
this.cancelableRequests.Remove(requestId);
}
}
}
}
9 changes: 9 additions & 0 deletions src/StreamJsonRpc.Tests/JsonRpcTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2181,6 +2181,15 @@ public async Task SerializableExceptions_Null()
Assert.Null(this.server.ReceivedException);
}

[Fact]
public void CancellationStrategy_ConfigurationLocked()
{
Assert.Throws<InvalidOperationException>(() => this.clientRpc.CancellationStrategy = null);
Assert.NotNull(this.clientRpc.CancellationStrategy);
this.clientRpc.AllowModificationWhileListening = true;
this.clientRpc.CancellationStrategy = null;
}

protected static Exception CreateExceptionToBeThrownByDeserializer() => new Exception("This exception is meant to be thrown.");

protected override void Dispose(bool disposing)
Expand Down
46 changes: 46 additions & 0 deletions src/StreamJsonRpc/ICancellationStrategy.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

namespace StreamJsonRpc
{
using System.Threading;

/// <summary>
/// Defines an extensibility point by which RPC methods may be canceled using <see cref="CancellationToken"/>.
/// </summary>
/// <remarks>
/// <para>A cancellation strategy can be set on the <see cref="JsonRpc.CancellationStrategy"/> property.</para>
/// <para>The default implementation is defined by <see cref="StandardCancellationStrategy"/>.</para>
/// <para>Implementations must be thread-safe.</para>
/// </remarks>
public interface ICancellationStrategy
{
/// <summary>
/// Translates a canceled <see cref="CancellationToken"/> that was used in an outbound RPC request into terms that
/// the RPC server can understand.
/// </summary>
/// <param name="requestId">The ID of the canceled request.</param>
void CancelOutboundRequest(RequestId requestId);

/// <summary>
/// Reports an incoming request and the <see cref="CancellationTokenSource"/> that is assigned to it.
/// </summary>
/// <param name="requestId">The ID of the incoming request.</param>
/// <param name="cancellationTokenSource">A means to cancel the <see cref="CancellationToken"/> that will be used when invoking the RPC server method.</param>
/// <remarks>
/// Implementations are expected to store the arguments in a dictionary so the implementing strategy can cancel it when the trigger occurs.
/// The trigger is outside the scope of this interface and will vary by implementation.
/// </remarks>
void IncomingRequestStarted(RequestId requestId, CancellationTokenSource cancellationTokenSource);

/// <summary>
/// Reports that an incoming request is no longer a candidate for cancellation.
/// </summary>
/// <param name="requestId">The ID of the request that has been fulfilled.</param>
/// <remarks>
/// Implementations are expected to release memory allocated by a prior call to <see cref="IncomingRequestStarted(RequestId, CancellationTokenSource)"/>.
/// This method should *not* dispose of the <see cref="CancellationTokenSource"/> received previously as <see cref="JsonRpc"/> owns its lifetime.
/// </remarks>
void IncomingRequestEnded(RequestId requestId);
}
}
Loading

0 comments on commit 28cfabd

Please sign in to comment.