Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 58 additions & 41 deletions src/SignalR/server/Core/src/Internal/DefaultHubDispatcher.cs
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,7 @@ private async Task<bool> Invoke(HubMethodDescriptor descriptor, HubConnectionCon
var scope = _serviceScopeFactory.CreateAsyncScope();
IHubActivator<THub>? hubActivator = null;
THub? hub = null;
long? streamOwner = null;
try
{
hubActivator = scope.ServiceProvider.GetRequiredService<IHubActivator<THub>>();
Expand Down Expand Up @@ -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)
Expand All @@ -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
{
Expand All @@ -439,7 +440,8 @@ static async Task ExecuteInvocation(DefaultHubDispatcher<THub> dispatcher,
HubCallerContext hubCallerContext,
HubMethodInvocationMessage hubMethodInvocationMessage,
bool isStreamCall,
CancellationTokenSource? cts)
CancellationTokenSource? cts,
long? streamOwner)
{
var logger = dispatcher._logger;
var enableDetailedErrors = dispatcher._enableDetailedErrors;
Expand Down Expand Up @@ -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);
}
}

Expand All @@ -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)
Expand Down Expand Up @@ -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<THub>? hubActivator,
private static ValueTask CleanupInvocation(HubConnectionContext connection, HubMethodInvocationMessage hubMessage, long? streamOwner, IHubActivator<THub>? 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);
}
}

Expand All @@ -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<THub> hubActivator, THub hub, CancellationTokenSource? streamCts, HubMethodInvocationMessage hubMethodInvocationMessage, HubMethodDescriptor descriptor)
IHubActivator<THub> hubActivator, THub hub, CancellationTokenSource? streamCts, HubMethodInvocationMessage hubMethodInvocationMessage,
HubMethodDescriptor descriptor, long? streamOwner)
{
string? error = null;

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -817,47 +820,61 @@ 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)
arguments = new object?[descriptor.OriginalParameterTypes!.Count];

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();
}
}
}
Expand Down
51 changes: 38 additions & 13 deletions src/SignalR/server/Core/src/StreamTracker.cs
Original file line number Diff line number Diff line change
Expand Up @@ -16,32 +16,42 @@ 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<string, IStreamConverter> _lookup = new ConcurrentDictionary<string, IStreamConverter>();
private readonly ConcurrentDictionary<string, (long Owner, IStreamConverter Converter)> _lookup = new();
private long _nextStreamOwner;

public StreamTracker(int streamBufferCapacity)
{
_streamConverterArgs = new object[] { streamBufferCapacity };
}

public long GetNextStreamOwner()
{
return Interlocked.Increment(ref _nextStreamOwner);
}

/// <summary>
/// Creates a new stream and returns the ChannelReader for it as an object.
/// </summary>
[UnconditionalSuppressMessage("Trimming", "IL2060:MakeGenericMethod",
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.");
Comment thread
javiercn marked this conversation as resolved.
}

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;
}
Expand All @@ -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;
}

Expand All @@ -63,30 +73,45 @@ 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.");
}

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;
}

public void CompleteAll(Exception ex)
{
foreach (var converter in _lookup)
{
converter.Value.TryComplete(ex);
converter.Value.Converter.TryComplete(ex);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,11 @@ public async Task<string> StreamingConcat(ChannelReader<string> source)
return sb.ToString();
}

public async Task<string> StreamingConcatTwoStreams(ChannelReader<string> first, ChannelReader<string> second)
{
return await StreamingConcat(first) + await StreamingConcat(second);
}
Comment thread
javiercn marked this conversation as resolved.

public async Task StreamDontRead(ChannelReader<string> source)
{
while (await source.WaitToReadAsync())
Expand Down
Loading
Loading