diff --git a/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs b/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs index f049147911f7..279fd103e2e4 100644 --- a/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs +++ b/src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs @@ -370,6 +370,7 @@ private async Task Invoke(HubMethodDescriptor descriptor, HubConnectionCon var scope = _serviceScopeFactory.CreateAsyncScope(); IHubActivator? hubActivator = null; THub? hub = null; + long? streamOwner = null; try { hubActivator = scope.ServiceProvider.GetRequiredService>(); @@ -409,7 +410,7 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection, CancellationTokenSource? cts = null; if (descriptor.HasSyntheticArguments) { - ReplaceArguments(descriptor, hubMethodInvocationMessage, isStreamCall, connection, scope, ref arguments, out cts); + ReplaceArguments(descriptor, hubMethodInvocationMessage, isStreamCall, connection, scope, ref arguments, ref streamOwner, out cts); } if (isStreamCall || isStreamResponse) @@ -424,7 +425,7 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection, if (isStreamResponse) { _ = StreamAsync(hubMethodInvocationMessage.InvocationId!, connection, hubCallerContext, - arguments, scope, hubActivator, hub, cts, hubMethodInvocationMessage, descriptor); + arguments, scope, hubActivator, hub, cts, hubMethodInvocationMessage, descriptor, streamOwner); } else { @@ -439,7 +440,8 @@ static async Task ExecuteInvocation(DefaultHubDispatcher dispatcher, HubCallerContext hubCallerContext, HubMethodInvocationMessage hubMethodInvocationMessage, bool isStreamCall, - CancellationTokenSource? cts) + CancellationTokenSource? cts, + long? streamOwner) { var logger = dispatcher._logger; var enableDetailedErrors = dispatcher._enableDetailedErrors; @@ -509,7 +511,7 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection, // And normal invocations handle cleanup below in the finally if (isStreamCall) { - await CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope); + await CleanupInvocation(connection, hubMethodInvocationMessage, streamOwner, hubActivator, hub, scope); } } @@ -521,7 +523,7 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection, } } - invocation = ExecuteInvocation(this, methodExecutor, hub, arguments, scope, hubActivator, connection, hubCallerContext, hubMethodInvocationMessage, isStreamCall, cts); + invocation = ExecuteInvocation(this, methodExecutor, hub, arguments, scope, hubActivator, connection, hubCallerContext, hubMethodInvocationMessage, isStreamCall, cts, streamOwner); } if (isStreamCall || isStreamResponse) @@ -557,21 +559,21 @@ await SendInvocationError(hubMethodInvocationMessage.InvocationId, connection, { wasSemaphoreReleased = !hubCallerClients.TrySetSemaphoreReleased(); } - await CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope); + await CleanupInvocation(connection, hubMethodInvocationMessage, streamOwner, hubActivator, hub, scope); } } return !wasSemaphoreReleased; } - private static ValueTask CleanupInvocation(HubConnectionContext connection, HubMethodInvocationMessage hubMessage, IHubActivator? hubActivator, + private static ValueTask CleanupInvocation(HubConnectionContext connection, HubMethodInvocationMessage hubMessage, long? streamOwner, IHubActivator? hubActivator, THub? hub, AsyncServiceScope scope) { - if (hubMessage.StreamIds != null) + if (streamOwner is not null) { - foreach (var stream in hubMessage.StreamIds) + foreach (var streamId in hubMessage.StreamIds!) { - connection.StreamTracker.TryComplete(CompletionMessage.Empty(stream)); + connection.StreamTracker.TryComplete(streamId, streamOwner.Value); } } @@ -584,7 +586,8 @@ private static ValueTask CleanupInvocation(HubConnectionContext connection, HubM } private async Task StreamAsync(string invocationId, HubConnectionContext connection, HubCallerContext hubCallerContext, object?[] arguments, AsyncServiceScope scope, - IHubActivator hubActivator, THub hub, CancellationTokenSource? streamCts, HubMethodInvocationMessage hubMethodInvocationMessage, HubMethodDescriptor descriptor) + IHubActivator hubActivator, THub hub, CancellationTokenSource? streamCts, HubMethodInvocationMessage hubMethodInvocationMessage, + HubMethodDescriptor descriptor, long? streamOwner) { string? error = null; @@ -671,7 +674,7 @@ private async Task StreamAsync(string invocationId, HubConnectionContext connect Activity.Current = previousActivity; } - await CleanupInvocation(connection, hubMethodInvocationMessage, hubActivator, hub, scope); + await CleanupInvocation(connection, hubMethodInvocationMessage, streamOwner, hubActivator, hub, scope); // Only remove/dispose the CTS if we successfully registered it, otherwise we'd evict // another invocation's CTS on ID collision. @@ -817,7 +820,8 @@ await connection.WriteAsync(CompletionMessage.WithError(hubMethodInvocationMessa } private void ReplaceArguments(HubMethodDescriptor descriptor, HubMethodInvocationMessage hubMethodInvocationMessage, bool isStreamCall, - HubConnectionContext connection, AsyncServiceScope scope, ref object?[] arguments, out CancellationTokenSource? cts) + HubConnectionContext connection, AsyncServiceScope scope, ref object?[] arguments, + ref long? streamOwner, out CancellationTokenSource? cts) { cts = null; // 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) @@ -825,39 +829,52 @@ private void ReplaceArguments(HubMethodDescriptor descriptor, HubMethodInvocatio var streamPointer = 0; var hubInvocationArgumentPointer = 0; - for (var parameterPointer = 0; parameterPointer < arguments.Length; parameterPointer++) + var argumentsReplaced = false; + try { - // populate the synthetic arguments first - if (descriptor.IsServiceArgument(parameterPointer)) - { - arguments[parameterPointer] = descriptor.GetService(scope.ServiceProvider, parameterPointer, descriptor.OriginalParameterTypes[parameterPointer]); - } - else if (descriptor.OriginalParameterTypes[parameterPointer] == typeof(CancellationToken)) + for (var parameterPointer = 0; parameterPointer < arguments.Length; parameterPointer++) { - cts = CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted); - arguments[parameterPointer] = cts.Token; - } - else if (isStreamCall && ReflectionHelper.IsStreamingType(descriptor.OriginalParameterTypes[parameterPointer], mustBeDirectType: true)) - { - Log.StartingParameterStream(_logger, hubMethodInvocationMessage.StreamIds![streamPointer]); - var itemType = descriptor.StreamingParameters![streamPointer]; - arguments[parameterPointer] = connection.StreamTracker.AddStream(hubMethodInvocationMessage.StreamIds[streamPointer], - itemType, descriptor.OriginalParameterTypes[parameterPointer]); + // populate the synthetic arguments first + if (descriptor.IsServiceArgument(parameterPointer)) + { + arguments[parameterPointer] = descriptor.GetService(scope.ServiceProvider, parameterPointer, descriptor.OriginalParameterTypes[parameterPointer]); + } + else if (descriptor.OriginalParameterTypes[parameterPointer] == typeof(CancellationToken)) + { + cts = CancellationTokenSource.CreateLinkedTokenSource(connection.ConnectionAborted); + arguments[parameterPointer] = cts.Token; + } + else if (isStreamCall && ReflectionHelper.IsStreamingType(descriptor.OriginalParameterTypes[parameterPointer], mustBeDirectType: true)) + { + Log.StartingParameterStream(_logger, hubMethodInvocationMessage.StreamIds![streamPointer]); + var itemType = descriptor.StreamingParameters![streamPointer]; + arguments[parameterPointer] = connection.StreamTracker.AddStream(hubMethodInvocationMessage.StreamIds[streamPointer], + itemType, descriptor.OriginalParameterTypes[parameterPointer], streamOwner ??= connection.StreamTracker.GetNextStreamOwner()); - streamPointer++; - } - else if (hubMethodInvocationMessage.Arguments?.Length > hubInvocationArgumentPointer && - (hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer] == null || - descriptor.OriginalParameterTypes[parameterPointer].IsAssignableFrom(hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer]?.GetType()))) - { - // The types match so it isn't a synthetic argument, just copy it into the arguments array - arguments[parameterPointer] = hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer]; - hubInvocationArgumentPointer++; + streamPointer++; + } + else if (hubMethodInvocationMessage.Arguments?.Length > hubInvocationArgumentPointer && + (hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer] == null || + descriptor.OriginalParameterTypes[parameterPointer].IsAssignableFrom(hubMethodInvocationMessage.Arguments[hubInvocationArgumentPointer]?.GetType()))) + { + // 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 should never happen + Debug.Assert(false, $"Failed to bind argument of type '{descriptor.OriginalParameterTypes[parameterPointer].Name}' for hub method '{descriptor.MethodExecutor.MethodInfo.Name}'."); + } } - else + + argumentsReplaced = true; + } + finally + { + if (!argumentsReplaced) { - // This should never happen - Debug.Assert(false, $"Failed to bind argument of type '{descriptor.OriginalParameterTypes[parameterPointer].Name}' for hub method '{descriptor.MethodExecutor.MethodInfo.Name}'."); + cts?.Dispose(); } } } diff --git a/src/SignalR/server/Core/src/StreamTracker.cs b/src/SignalR/server/Core/src/StreamTracker.cs index 81fea34566de..7666cd420b06 100644 --- a/src/SignalR/server/Core/src/StreamTracker.cs +++ b/src/SignalR/server/Core/src/StreamTracker.cs @@ -16,13 +16,19 @@ internal sealed class StreamTracker { private static readonly MethodInfo _buildConverterMethod = typeof(StreamTracker).GetMethods(BindingFlags.NonPublic | BindingFlags.Static).Single(m => m.Name.Equals(nameof(BuildStream))); private readonly object[] _streamConverterArgs; - private readonly ConcurrentDictionary _lookup = new ConcurrentDictionary(); + private readonly ConcurrentDictionary _lookup = new(); + private long _nextStreamOwner; public StreamTracker(int streamBufferCapacity) { _streamConverterArgs = new object[] { streamBufferCapacity }; } + public long GetNextStreamOwner() + { + return Interlocked.Increment(ref _nextStreamOwner); + } + /// /// Creates a new stream and returns the ChannelReader for it as an object. /// @@ -30,18 +36,22 @@ public StreamTracker(int streamBufferCapacity) Justification = "BuildStream doesn't have trimming annotations.")] [UnconditionalSuppressMessage("AOT", "IL3050:RequiresDynamicCode", Justification = "HubMethodDescriptor checks for ValueType streaming item types when PublishAot=true. Developers will get an exception in this situation before publishing.")] - public object AddStream(string streamId, Type itemType, Type targetType) + public object AddStream(string streamId, Type itemType, Type targetType, long streamOwner) { Debug.Assert(RuntimeFeature.IsDynamicCodeSupported || !itemType.IsValueType, "HubMethodDescriptor ensures itemType is not a ValueType when PublishAot=true."); var newConverter = (IStreamConverter)_buildConverterMethod.MakeGenericMethod(itemType).Invoke(null, _streamConverterArgs)!; - _lookup[streamId] = newConverter; + if (!_lookup.TryAdd(streamId, (streamOwner, newConverter))) + { + throw new HubException($"Stream ID '{streamId}' is already in use."); + } + return newConverter.GetReaderAsObject(targetType); } - private bool TryGetConverter(string streamId, [NotNullWhen(true)] out IStreamConverter? converter) + private bool TryGetRegistration(string streamId, out (long Owner, IStreamConverter Converter) registration) { - if (_lookup.TryGetValue(streamId, out converter)) + if (_lookup.TryGetValue(streamId, out registration)) { return true; } @@ -51,9 +61,9 @@ private bool TryGetConverter(string streamId, [NotNullWhen(true)] out IStreamCon public bool TryProcessItem(StreamItemMessage message, [NotNullWhen(true)] out Task? task) { - if (TryGetConverter(message.InvocationId!, out var converter)) + if (TryGetRegistration(message.InvocationId!, out var registration)) { - task = converter.WriteToStream(message.Item); + task = registration.Converter.WriteToStream(message.Item); return true; } @@ -63,9 +73,9 @@ public bool TryProcessItem(StreamItemMessage message, [NotNullWhen(true)] out Ta public Type GetStreamItemType(string streamId) { - if (TryGetConverter(streamId, out var converter)) + if (TryGetRegistration(streamId, out var registration)) { - return converter.GetItemType(); + return registration.Converter.GetItemType(); } throw new KeyNotFoundException($"No stream with id '{streamId}' could be found."); @@ -73,12 +83,27 @@ public Type GetStreamItemType(string streamId) public bool TryComplete(CompletionMessage message) { - _lookup.TryRemove(message.InvocationId!, out var converter); - if (converter == null) + if (!_lookup.TryRemove(message.InvocationId!, out var registration)) + { + return false; + } + registration.Converter.TryComplete(message.HasResult || message.Error == null ? null : new HubException(message.Error)); + return true; + } + + public bool TryComplete(string streamId, long streamOwner) + { + if (!_lookup.TryGetValue(streamId, out var registration) || registration.Owner != streamOwner) { return false; } - converter.TryComplete(message.HasResult || message.Error == null ? null : new HubException(message.Error)); + + if (!_lookup.TryRemove(KeyValuePair.Create(streamId, registration))) + { + return false; + } + + registration.Converter.TryComplete(null); return true; } @@ -86,7 +111,7 @@ public void CompleteAll(Exception ex) { foreach (var converter in _lookup) { - converter.Value.TryComplete(ex); + converter.Value.Converter.TryComplete(ex); } } diff --git a/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs b/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs index afb56aa755a0..03328940668d 100644 --- a/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs +++ b/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTestUtils/Hubs.cs @@ -223,6 +223,11 @@ public async Task StreamingConcat(ChannelReader source) return sb.ToString(); } + public async Task StreamingConcatTwoStreams(ChannelReader first, ChannelReader second) + { + return await StreamingConcat(first) + await StreamingConcat(second); + } + public async Task StreamDontRead(ChannelReader source) { while (await source.WaitToReadAsync()) diff --git a/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs b/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs index 051ae6409ac0..6fe265df7af5 100644 --- a/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs +++ b/src/SignalR/server/SignalR/test/Microsoft.AspNetCore.SignalR.Tests/HubConnectionHandlerTests.cs @@ -3696,6 +3696,88 @@ public async Task UploadStringsToConcat() } } + [Fact] + public async Task UploadStreamWithDuplicateIdsFailsAndConnectionContinues() + { + var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(services => + { + services.AddSignalR(options => + { + options.EnableDetailedErrors = true; + options.StreamBufferCapacity = 1; + }); + }); + var connectionHandler = serviceProvider.GetService>(); + + using (var client = new TestClient()) + { + var connectionHandlerTask = await client.ConnectAsync(connectionHandler).DefaultTimeout(); + await client.BeginUploadStreamAsync("duplicate", nameof(MethodHub.StreamingConcatTwoStreams), new[] { "id", "id" }, Array.Empty()).DefaultTimeout(); + await client.SendHubMessageAsync(new StreamItemMessage("id", "first")).DefaultTimeout(); + await client.SendHubMessageAsync(new StreamItemMessage("id", "second")).DefaultTimeout(); + await client.SendInvocationAsync(nameof(MethodHub.Echo), "test").DefaultTimeout(); + + var duplicateCompletion = Assert.IsType(await client.ReadAsync().DefaultTimeout()); + Assert.Equal("An unexpected error occurred invoking 'StreamingConcatTwoStreams' on the server. HubException: Stream ID 'id' is already in use.", duplicateCompletion.Error); + + var echoCompletion = Assert.IsType(await client.ReadAsync().DefaultTimeout()); + Assert.Equal("test", echoCompletion.Result); + } + } + + [Fact] + public async Task ActiveUploadStreamCannotBeReplacedAndIdCanBeReusedAfterCompletion() + { + var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(services => + { + services.AddSignalR(options => options.EnableDetailedErrors = true); + }); + var connectionHandler = serviceProvider.GetService>(); + + using (var client = new TestClient()) + { + var connectionHandlerTask = await client.ConnectAsync(connectionHandler).DefaultTimeout(); + await client.BeginUploadStreamAsync("original", nameof(MethodHub.StreamingConcat), new[] { "id" }, Array.Empty()).DefaultTimeout(); + await client.BeginUploadStreamAsync("duplicate", nameof(MethodHub.StreamingConcat), new[] { "id" }, Array.Empty()).DefaultTimeout(); + + var duplicateCompletion = Assert.IsType(await client.ReadAsync().DefaultTimeout()); + Assert.Equal("An unexpected error occurred invoking 'StreamingConcat' on the server. HubException: Stream ID 'id' is already in use.", duplicateCompletion.Error); + + await client.SendHubMessageAsync(new StreamItemMessage("id", "original")).DefaultTimeout(); + await client.SendHubMessageAsync(CompletionMessage.Empty("id")).DefaultTimeout(); + + var originalCompletion = Assert.IsType(await client.ReadAsync().DefaultTimeout()); + Assert.Equal("original", originalCompletion.Result); + + await client.BeginUploadStreamAsync("reused", nameof(MethodHub.StreamingConcat), new[] { "id" }, Array.Empty()).DefaultTimeout(); + await client.SendHubMessageAsync(new StreamItemMessage("id", "reused")).DefaultTimeout(); + await client.SendHubMessageAsync(CompletionMessage.Empty("id")).DefaultTimeout(); + + var reusedCompletion = Assert.IsType(await client.ReadAsync().DefaultTimeout()); + Assert.Equal("reused", reusedCompletion.Result); + } + } + + [Fact] + public async Task UploadMultipleStreamsWithUniqueIds() + { + var serviceProvider = HubConnectionHandlerTestUtils.CreateServiceProvider(); + var connectionHandler = serviceProvider.GetService>(); + + using (var client = new TestClient()) + { + var connectionHandlerTask = await client.ConnectAsync(connectionHandler).DefaultTimeout(); + await client.BeginUploadStreamAsync("invocation", nameof(MethodHub.StreamingConcatTwoStreams), new[] { "first", "second" }, Array.Empty()).DefaultTimeout(); + await client.SendHubMessageAsync(new StreamItemMessage("first", "hello ")).DefaultTimeout(); + await client.SendHubMessageAsync(CompletionMessage.Empty("first")).DefaultTimeout(); + await client.SendHubMessageAsync(new StreamItemMessage("second", "world")).DefaultTimeout(); + await client.SendHubMessageAsync(CompletionMessage.Empty("second")).DefaultTimeout(); + + var completion = Assert.IsType(await client.ReadAsync().DefaultTimeout()); + Assert.Equal("hello world", completion.Result); + } + } + [Fact] public async Task UploadStreamedObjects() {