diff --git a/src/Microsoft.AspNetCore.SignalR.Core/Internal/DefaultHubDispatcher.cs b/src/Microsoft.AspNetCore.SignalR.Core/Internal/DefaultHubDispatcher.cs index d26d10ac89..70ef9734bd 100644 --- a/src/Microsoft.AspNetCore.SignalR.Core/Internal/DefaultHubDispatcher.cs +++ b/src/Microsoft.AspNetCore.SignalR.Core/Internal/DefaultHubDispatcher.cs @@ -188,11 +188,45 @@ namespace Microsoft.AspNetCore.SignalR.Internal { InitializeHub(hub, connection); - var result = await ExecuteHubMethod(methodExecutor, hub, hubMethodInvocationMessage.Arguments); + CancellationTokenSource cts = null; + var arguments = hubMethodInvocationMessage.Arguments; + if (descriptor.HasSyntheticArguments) + { + // In order to add the synthetic arguments we need a new array because the invocation array is too small (it doesn't know about synthetic arguments) + arguments = new object[descriptor.OriginalParameterTypes.Count]; + + var hubInvocationArgumentPointer = 0; + for (var parameterPointer = 0; parameterPointer < arguments.Length; parameterPointer++) + { + if (hubMethodInvocationMessage.Arguments.Length > hubInvocationArgumentPointer && + hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer].GetType() == descriptor.OriginalParameterTypes[parameterPointer]) + { + // The types match so it isn't a synthetic argument, just copy it into the arguments array + arguments[parameterPointer] = hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer]; + hubInvocationArgumentPointer++; + } + else + { + // This is the only synthetic argument type we currently support + if (descriptor.OriginalParameterTypes[parameterPointer] == typeof(CancellationToken)) + { + cts = CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted); + arguments[parameterPointer] = cts.Token; + } + else + { + // This should never happen + Debug.Assert(false, $"Failed to bind argument of type '{descriptor.OriginalParameterTypes[parameterPointer].Name}' for hub method '{methodExecutor.MethodInfo.Name}'."); + } + } + } + } + + var result = await ExecuteHubMethod(methodExecutor, hub, arguments); if (isStreamedInvocation) { - if (!TryGetStreamingEnumerator(connection, hubMethodInvocationMessage.InvocationId, descriptor, result, out var enumerator, out var streamCts)) + if (!TryGetStreamingEnumerator(connection, hubMethodInvocationMessage.InvocationId, descriptor, result, out var enumerator, ref cts)) { Log.InvalidReturnValueFromStreamingMethod(_logger, methodExecutor.MethodInfo.Name); @@ -204,7 +238,7 @@ namespace Microsoft.AspNetCore.SignalR.Internal disposeScope = false; Log.StreamingResult(_logger, hubMethodInvocationMessage.InvocationId, methodExecutor); // Fire-and-forget stream invocations, otherwise they would block other hub invocations from being able to run - _ = StreamResultsAsync(hubMethodInvocationMessage.InvocationId, connection, enumerator, scope, hubActivator, hub, streamCts); + _ = StreamResultsAsync(hubMethodInvocationMessage.InvocationId, connection, enumerator, scope, hubActivator, hub, cts); } // Non-empty/null InvocationId ==> Blocking invocation that needs a response else if (!string.IsNullOrEmpty(hubMethodInvocationMessage.InvocationId)) @@ -375,29 +409,24 @@ namespace Microsoft.AspNetCore.SignalR.Internal return true; } - private bool TryGetStreamingEnumerator(HubConnectionContext connection, string invocationId, HubMethodDescriptor hubMethodDescriptor, object result, out IAsyncEnumerator enumerator, out CancellationTokenSource streamCts) + private bool TryGetStreamingEnumerator(HubConnectionContext connection, string invocationId, HubMethodDescriptor hubMethodDescriptor, object result, out IAsyncEnumerator enumerator, ref CancellationTokenSource streamCts) { if (result != null) { if (hubMethodDescriptor.IsChannel) { - streamCts = CreateCancellation(); + if (streamCts == null) + { + streamCts = CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted); + } + connection.ActiveRequestCancellationSources.TryAdd(invocationId, streamCts); enumerator = hubMethodDescriptor.FromChannel(result, streamCts.Token); return true; } } - streamCts = null; enumerator = null; return false; - - CancellationTokenSource CreateCancellation() - { - var userCts = new CancellationTokenSource(); - connection.ActiveRequestCancellationSources.TryAdd(invocationId, userCts); - - return CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted, userCts.Token); - } } private void DiscoverHubMethods() diff --git a/src/Microsoft.AspNetCore.SignalR.Core/Internal/HubMethodDescriptor.cs b/src/Microsoft.AspNetCore.SignalR.Core/Internal/HubMethodDescriptor.cs index a15dce772e..6c7c895659 100644 --- a/src/Microsoft.AspNetCore.SignalR.Core/Internal/HubMethodDescriptor.cs +++ b/src/Microsoft.AspNetCore.SignalR.Core/Internal/HubMethodDescriptor.cs @@ -22,8 +22,6 @@ namespace Microsoft.AspNetCore.SignalR.Internal public HubMethodDescriptor(ObjectMethodExecutor methodExecutor, IEnumerable policies) { MethodExecutor = methodExecutor; - ParameterTypes = methodExecutor.MethodParameters.Select(p => p.ParameterType).ToArray(); - Policies = policies.ToArray(); NonAsyncReturnType = (MethodExecutor.IsMethodAsync) ? MethodExecutor.AsyncResultType @@ -34,6 +32,25 @@ namespace Microsoft.AspNetCore.SignalR.Internal IsChannel = true; StreamReturnType = channelItemType; } + + // Take out synthetic arguments that will be provided by the server, this list will be given to the protocol parsers + ParameterTypes = methodExecutor.MethodParameters.Where(p => + { + // Only streams can take CancellationTokens currently + if (IsStreamable && p.ParameterType == typeof(CancellationToken)) + { + HasSyntheticArguments = true; + return false; + } + return true; + }).Select(p => p.ParameterType).ToArray(); + + if (HasSyntheticArguments) + { + OriginalParameterTypes = methodExecutor.MethodParameters.Select(p => p.ParameterType).ToArray(); + } + + Policies = policies.ToArray(); } private Func> _convertToEnumerator; @@ -42,6 +59,8 @@ namespace Microsoft.AspNetCore.SignalR.Internal public IReadOnlyList ParameterTypes { get; } + public IReadOnlyList OriginalParameterTypes { get; } + public Type NonAsyncReturnType { get; } public bool IsChannel { get; } @@ -52,6 +71,8 @@ namespace Microsoft.AspNetCore.SignalR.Internal public IList Policies { get; } + public bool HasSyntheticArguments { get; private set; } + private static bool IsChannelType(Type type, out Type payloadType) { var channelType = type.AllBaseTypes().FirstOrDefault(t => t.IsGenericType && t.GetGenericTypeDefinition() == typeof(ChannelReader<>)); diff --git a/test/Microsoft.AspNetCore.SignalR.Tests.Utils/CancellationTokenExtensions.cs b/test/Microsoft.AspNetCore.SignalR.Tests.Utils/CancellationTokenExtensions.cs new file mode 100644 index 0000000000..fadb9626bf --- /dev/null +++ b/test/Microsoft.AspNetCore.SignalR.Tests.Utils/CancellationTokenExtensions.cs @@ -0,0 +1,21 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information. + +using System.Threading; +using System.Threading.Tasks; + +namespace Microsoft.AspNetCore.SignalR.Tests +{ + public static class CancellationTokenExtensions + { + public static Task WaitForCancellationAsync(this CancellationToken token) + { + var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + token.Register((t) => + { + ((TaskCompletionSource)t).SetResult(null); + }, tcs); + return tcs.Task; + } + } +} diff --git a/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs b/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs index 7eb0aec498..37fc164c14 100644 --- a/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs +++ b/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; +using System.Threading; using System.Threading.Channels; using System.Threading.Tasks; using Microsoft.AspNetCore.Authorization; @@ -165,6 +166,10 @@ namespace Microsoft.AspNetCore.SignalR.Tests return Clients.Caller.SendAsync("Send", new string('x', 3000), new SelfRef()); } + public void InvalidArgument(CancellationToken token) + { + } + private class SelfRef { public SelfRef() @@ -547,6 +552,51 @@ namespace Microsoft.AspNetCore.SignalR.Tests return Channel.CreateUnbounded().Reader; } + public ChannelReader CancelableStream(CancellationToken token) + { + var channel = Channel.CreateBounded(10); + + Task.Run(async () => + { + _tcsService.StartedMethod.SetResult(null); + await token.WaitForCancellationAsync(); + channel.Writer.TryComplete(); + _tcsService.EndMethod.SetResult(null); + }); + + return channel.Reader; + } + + public ChannelReader CancelableStream2(int ignore, int ignore2, CancellationToken token) + { + var channel = Channel.CreateBounded(10); + + Task.Run(async () => + { + _tcsService.StartedMethod.SetResult(null); + await token.WaitForCancellationAsync(); + channel.Writer.TryComplete(); + _tcsService.EndMethod.SetResult(null); + }); + + return channel.Reader; + } + + public ChannelReader CancelableStreamMiddle(int ignore, CancellationToken token, int ignore2) + { + var channel = Channel.CreateBounded(10); + + Task.Run(async () => + { + _tcsService.StartedMethod.SetResult(null); + await token.WaitForCancellationAsync(); + channel.Writer.TryComplete(); + _tcsService.EndMethod.SetResult(null); + }); + + return channel.Reader; + } + public int SimpleMethod() { return 21; diff --git a/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs b/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs index c637f9e131..c8d884ec5f 100644 --- a/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs +++ b/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs @@ -2381,6 +2381,95 @@ namespace Microsoft.AspNetCore.SignalR.Tests } } + [Theory] + [InlineData(nameof(LongRunningHub.CancelableStream))] + [InlineData(nameof(LongRunningHub.CancelableStream2), 1, 2)] + [InlineData(nameof(LongRunningHub.CancelableStreamMiddle), 1, 2)] + public async Task StreamHubMethodCanAcceptCancellationTokenAsArgumentAndBeTriggeredOnCancellation(string methodName, params object[] args) + { + var tcsService = new TcsService(); + var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(builder => + { + builder.AddSingleton(tcsService); + }); + var connectionHandler = serviceProvider.GetService>(); + + using (var client = new TestClient()) + { + var connectionHandlerTask = await client.ConnectAsync(connectionHandler).OrTimeout(); + + var streamInvocationId = await client.SendStreamInvocationAsync(methodName, args).OrTimeout(); + // Wait for the stream method to start + await tcsService.StartedMethod.Task.OrTimeout(); + + // Cancel the stream which should trigger the CancellationToken in the hub method + await client.SendHubMessageAsync(new CancelInvocationMessage(streamInvocationId)).OrTimeout(); + + var result = await client.ReadAsync().OrTimeout(); + + var simpleCompletion = Assert.IsType(result); + Assert.Null(simpleCompletion.Result); + + // CancellationToken passed to hub method will allow EndMethod to be triggered if it is canceled. + await tcsService.EndMethod.Task.OrTimeout(); + + // Shut down + client.Dispose(); + + await connectionHandlerTask.OrTimeout(); + } + } + + [Fact] + public async Task StreamHubMethodCanAcceptCancellationTokenAsArgumentAndBeTriggeredOnConnectionAborted() + { + var tcsService = new TcsService(); + var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(builder => + { + builder.AddSingleton(tcsService); + }); + var connectionHandler = serviceProvider.GetService>(); + + using (var client = new TestClient()) + { + var connectionHandlerTask = await client.ConnectAsync(connectionHandler).OrTimeout(); + + var streamInvocationId = await client.SendStreamInvocationAsync(nameof(LongRunningHub.CancelableStream)).OrTimeout(); + // Wait for the stream method to start + await tcsService.StartedMethod.Task.OrTimeout(); + + // Shut down the client which should trigger the CancellationToken in the hub method + client.Dispose(); + + // CancellationToken passed to hub method will allow EndMethod to be triggered if it is canceled. + await tcsService.EndMethod.Task.OrTimeout(); + + await connectionHandlerTask.OrTimeout(); + } + } + + [Fact] + public async Task InvokeHubMethodCannotAcceptCancellationTokenAsArgument() + { + var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(); + var connectionHandler = serviceProvider.GetService>(); + + using (var client = new TestClient()) + { + var connectionHandlerTask = await client.ConnectAsync(connectionHandler).OrTimeout(); + + var invocationId = await client.SendInvocationAsync(nameof(MethodHub.InvalidArgument)).OrTimeout(); + + var completion = Assert.IsType(await client.ReadAsync().OrTimeout()); + + Assert.Equal("Failed to invoke 'InvalidArgument' due to an error on the server.", completion.Error); + + client.Dispose(); + + await connectionHandlerTask.OrTimeout(); + } + } + private class CustomHubActivator : IHubActivator where THub : Hub { public int ReleaseCount;