private async Task ProcessIncomingMessageAsync(ClientConnectionContext connection)
{
Exception exception = null;
try
{
await connection.ApplicationTask;
}
catch (Exception e)
{
exception = e;
}
finally
{
// If we aren't already aborted, we send the abort message to the service
if (connection.AbortOnClose)
{
// Inform the Service that we will remove the client because SignalR told us it is disconnected.
var serviceMessage = new CloseConnectionMessage(connection.ConnectionId, errorMessage: exception?.Message);
await SafeWriteAsync(serviceMessage);
Log.CloseConnection(Logger, connection.ConnectionId);
}
}
}
public async Task ClosingConnectionSendsCloseMessage()
{
var proxy = new ServiceConnectionProxy(context =>
{
// Just let the connection end immediately
return(Task.CompletedTask);
});
var serverTask = proxy.WaitForServerConnectionAsync(1);
_ = proxy.StartAsync();
await serverTask.OrTimeout();
var task = proxy.WaitForConnectionAsync("1");
var closeMessageTask = proxy.WaitForApplicationMessageAsync(typeof(CloseConnectionMessage));
await proxy.WriteMessageAsync(new OpenConnectionMessage("1", null));
var connection = await task.OrTimeout();
CloseConnectionMessage message = (CloseConnectionMessage)await closeMessageTask.OrTimeout();
Assert.Equal(message.ConnectionId, connection.ConnectionId);
proxy.Stop();
}
private async Task WaitOnApplicationTask(ServiceConnectionContext connection)
{
Exception exception = null;
try
{
// Wait for the application task to complete
await connection.ApplicationTask;
}
catch (Exception ex)
{
// Capture the exception to communicate it to the transport (this isn't strictly required)
exception = ex;
}
finally
{
// Close the transport side since the application is no longer running
connection.Transport.Output.Complete(exception);
connection.Transport.Input.Complete();
}
// If we aren't already aborted, we send the abort message to the service
if (connection.AbortOnClose)
{
// Inform the Service that we will remove the client because SignalR told us it is disconnected.
var serviceMessage = new CloseConnectionMessage(connection.ConnectionId, errorMessage: "");
await WriteAsync(serviceMessage);
Log.CloseConnection(_logger, connection.ConnectionId);
}
}
public async Task TestServiceConnectionForMigratedOut()
{
var factory = new TestClientConnectionFactory();
CancellationTokenSource cts = new();
var connection = CreateServiceConnection(clientConnectionFactory: factory, mode: GracefulShutdownMode.MigrateClients);
// create a connection with migration header.
await connection.OnClientConnectedAsyncForTest(new OpenConnectionMessage("foo", new Claim[0]));
var context = factory.Connections[0];
var closeMessage = new CloseConnectionMessage(context.ConnectionId);
closeMessage.Headers.Add(Constants.AsrsMigrateTo, "another-server");
var disconnect = connection.OnClientDisconnectedAsyncForTest(closeMessage);
var feature = context.Features.Get <IConnectionMigrationFeature>();
Assert.NotNull(feature);
Assert.Equal("another-server", feature.MigrateTo);
cts.Cancel();
await disconnect.OrTimeout();
}
protected override async Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
await base.OnDisconnectedAsync(closeConnectionMessage);
var tcs = _waitForConnectionClose.GetOrAdd(closeConnectionMessage.ConnectionId, i => new TaskCompletionSource <object>());
tcs.TrySetResult(null);
}
protected override Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
var connectionId = closeConnectionMessage.ConnectionId;
PerformDisconnectCore(connectionId);
return(Task.CompletedTask);
}
protected override Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
if (_connections.TryRemove(closeConnectionMessage.ConnectionId, out var connection))
{
connection.Disconnect();
}
return(Task.CompletedTask);
}
public async Task CloseConnection()
{
var closeClientMessage = new CloseConnectionMessage(ConnectionId, "bbb");
_servicePro.WriteMessage(closeClientMessage, ServiceSideConnection.MockServicePipe.Output);
var flushResult = await ServiceSideConnection.MockServicePipe.Output.FlushAsync();
if (flushResult.IsCanceled || flushResult.IsCompleted)
{
throw new InvalidOperationException($"CloseConnectionMessage for client connection id {ConnectionId} resulted in FlushResult IsCanceled {flushResult.IsCanceled} IsCompleted {flushResult.IsCompleted}");
}
}
public async Task TestCloseConnectionMessageWithMigrateOut()
{
var clientConnectionFactory = new TestClientConnectionFactory();
var connection = CreateServiceConnection(clientConnectionFactory: clientConnectionFactory, handler: new TestConnectionHandler(3000, "foobar"));
_ = connection.StartAsync();
await connection.ConnectionInitializedTask.OrTimeout(1000);
var openConnectionMessage = new OpenConnectionMessage("foo", Array.Empty <Claim>());
_ = connection.WriteFromServiceAsync(openConnectionMessage);
await connection.ClientConnectedTask;
Assert.Equal(1, clientConnectionFactory.Connections.Count);
var clientConnection = clientConnectionFactory.Connections[0];
Assert.False(clientConnection.IsMigrated);
// write a signalr handshake response
var message = new SignalRProtocol.HandshakeResponseMessage("");
SignalRProtocol.HandshakeProtocol.WriteResponseMessage(message, clientConnection.Transport.Output);
await clientConnection.Transport.Output.FlushAsync();
// write a close connection message with migration header
var closeMessage = new CloseConnectionMessage(clientConnection.ConnectionId);
closeMessage.Headers.Add(Constants.AsrsMigrateTo, "another-server");
await connection.WriteFromServiceAsync(closeMessage);
// wait until app task completed.
await Assert.ThrowsAsync <TimeoutException>(async() => await clientConnection.LifetimeTask.OrTimeout(1000));
await clientConnection.LifetimeTask.OrTimeout(3000);
// expect a handshake response message.
await connection.ExpectSignalRMessage(SignalRProtocol.HandshakeResponseMessage.Empty).OrTimeout(3000);
// signalr close message should be skipped.
await Assert.ThrowsAsync <TimeoutException>(async() => await connection.ExpectSignalRMessage(SignalRProtocol.CloseMessage.Empty).OrTimeout(1000));
var feature = clientConnection.Features.Get <IConnectionMigrationFeature>();
Assert.NotNull(feature);
Assert.Equal("another-server", feature.MigrateTo);
await connection.StopAsync();
}
protected override async Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
var connectionId = closeConnectionMessage.ConnectionId;
if (_clientConnections.TryGetValue(connectionId, out var clientContext))
{
try
{
await clientContext.Output.WriteAsync(closeConnectionMessage);
}
catch (Exception e)
{
Log.FailToWriteMessageToApplication(_logger, connectionId, e);
}
}
}
protected override Task OnClientDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
var connectionId = closeConnectionMessage.ConnectionId;
if (_enableConnectionMigration && _clientConnectionManager.ClientConnections.TryGetValue(connectionId, out var context))
{
if (!context.HttpContext.Request.Headers.ContainsKey(Constants.AsrsMigrateOut))
{
context.HttpContext.Request.Headers.Add(Constants.AsrsMigrateOut, "");
}
// We have to prevent SignalR `{type: 7}` (close message) from reaching our client while doing migration.
// Since all user-created messages will be sent to `ServiceConnection` directly.
// We can simply ignore all messages came from the application pipe.
context.Application.Input.CancelPendingRead();
}
return(PerformDisconnectAsyncCore(connectionId, false));
}
protected override Task OnClientDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
var connectionId = closeConnectionMessage.ConnectionId;
if (_clientConnectionManager.ClientConnections.TryGetValue(connectionId, out var context))
{
if (closeConnectionMessage.Headers.TryGetValue(Constants.AsrsMigrateTo, out var to))
{
context.Features.Set <IConnectionMigrationFeature>(new ConnectionMigrationFeature(ServerId, to));
// We have to prevent SignalR `{type: 7}` (close message) from reaching our client while doing migration.
// Since all user-created messages will be sent to `ServiceConnection` directly.
// We can simply ignore all messages came from the application pipe.
}
}
return(PerformDisconnectAsyncCore(connectionId));
}
protected override Task OnClientDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
var connectionId = closeConnectionMessage.ConnectionId;
if (_clientConnectionManager.ClientConnections.TryGetValue(connectionId, out var context))
{
if (closeConnectionMessage.Headers.TryGetValue(Constants.AsrsMigrateTo, out var to))
{
context.AbortOnClose = false;
context.Features.Set <IConnectionMigrationFeature>(new ConnectionMigrationFeature(ServerId, to));
// We have to prevent SignalR `{type: 7}` (close message) from reaching our client while doing migration.
// Since all data messages will be sent to `ServiceConnection` directly.
// We can simply ignore all messages came from the application.
context.CancelOutgoing();
// The close connection message must be the last message, so we could complete the pipe.
context.CompleteIncoming();
}
}
return(PerformDisconnectAsyncCore(connectionId));
}
public async Task TestServiceConnectionForMigratedOut()
{
var factory = new TestClientConnectionFactory();
var connection = MockServiceConnection(null, factory);
// create a connection with migration header.
await connection.OnClientConnectedAsyncForTest(new OpenConnectionMessage("foo", new Claim[0]));
var context = factory.Connections[0];
var closeMessage = new CloseConnectionMessage(context.ConnectionId);
closeMessage.Headers.Add(Constants.AsrsMigrateTo, "another-server");
await connection.OnClientDisconnectedAsyncForTest(closeMessage);
var feature = context.Features.Get <IConnectionMigrationFeature>();
Assert.NotNull(feature);
Assert.Equal("another-server", feature.MigrateTo);
}
protected override Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
var connectionId = closeConnectionMessage.ConnectionId;
return(PerformDisconnectAsyncCore(connectionId, false));
}
public static void StartToCloseConnection(ILogger logger, CloseConnectionMessage message)
{
_startToCloseConnection(logger, message.TracingId, message.ConnectionId, message.ErrorMessage, null);
}
protected override Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
return(Task.CompletedTask);
}
protected override Task OnClientDisconnectedAsync(CloseConnectionMessage closeConnectionMessage)
{
return(ForwardMessageToApplication(closeConnectionMessage.ConnectionId, closeConnectionMessage));
}
private async Task ProcessClientConnectionAsync(ClientConnectionContext connection)
{
try
{
// Writing from the application to the service
var transport = ProcessOutgoingMessagesAsync(connection, connection.OutgoingAborted);
// Waiting for the application to shutdown so we can clean up the connection
var app = ProcessApplicationTaskAsyncCore(connection);
var task = await Task.WhenAny(app, transport);
// remove it from the connection list
RemoveClientConnection(connection.ConnectionId);
// This is the exception from application
Exception exception = null;
if (task == app)
{
exception = app.Exception?.GetBaseException();
// there is no need to write to the transport as application is no longer running
Log.WaitingForTransport(Logger);
// app task completes connection.Transport.Output, which will completes connection.Application.Input and ends the transport
// Transports are written by us and are well behaved, wait for them to drain
connection.CancelOutgoing(_closeTimeOutMilliseconds);
// transport never throws
await transport;
}
else
{
// transport task ends first, no data will be dispatched out
Log.WaitingForApplication(Logger);
try
{
// always wait for the application to complete
await app;
}
catch (Exception e)
{
exception = e;
}
}
if (exception != null)
{
Log.ApplicationTaskFailed(Logger, exception);
}
// If we aren't already aborted, we send the abort message to the service
if (connection.AbortOnClose)
{
// Inform the Service that we will remove the client because SignalR told us it is disconnected.
var serviceMessage =
new CloseConnectionMessage(connection.ConnectionId, errorMessage: exception?.Message);
// when it fails, it means the underlying connection is dropped
// service is responsible for closing the client connections in this case and there is no need to throw
await SafeWriteAsync(serviceMessage);
Log.CloseConnection(Logger, connection.ConnectionId);
}
Log.ConnectedEnding(Logger, connection.ConnectionId);
}
catch (Exception e)
{
// When it throws, there must be something wrong
Log.ProcessConnectionFailed(Logger, connection.ConnectionId, e);
}
finally
{
connection.OnCompleted();
}
}
protected abstract Task OnDisconnectedAsync(CloseConnectionMessage closeConnectionMessage);
public Task OnClientDisconnectedAsyncForTest(CloseConnectionMessage message)
{
return(base.OnClientDisconnectedAsync(message));
}
private bool CloseConnectionMessagesEqual(CloseConnectionMessage x, CloseConnectionMessage y)
{
return(StringEqual(x.ConnectionId, y.ConnectionId) && StringEqual(x.ErrorMessage, y.ErrorMessage));
}
protected override async Task OnClientDisconnectedAsync(CloseConnectionMessage message)
{
await base.OnClientDisconnectedAsync(message);
_clientDisconnectedTcs.TrySetResult();
}
