protected override Task OnClientConnectedAsync(OpenConnectionMessage openConnectionMessage)
{
// Create empty transport with only channel for async processing messages
var connectionId = openConnectionMessage.ConnectionId;
var clientContext = new ClientConnectionContext(this, connectionId, GetInstanceId(openConnectionMessage.Headers));
bool isDiagnosticClient = false;
openConnectionMessage.Headers.TryGetValue(Constants.AsrsIsDiagnosticClient, out var isDiagnosticClientValue);
if (!StringValues.IsNullOrEmpty(isDiagnosticClientValue))
{
isDiagnosticClient = Convert.ToBoolean(isDiagnosticClientValue.FirstOrDefault());
}
// todo: ignore asp.net for now
using (new ClientConnectionScope(outboundConnection: this, isDiagnosticClient: isDiagnosticClient))
{
if (_clientConnectionManager.TryAddClientConnection(clientContext))
{
_clientConnections.TryAdd(connectionId, clientContext);
clientContext.ApplicationTask = ProcessMessageAsync(clientContext, clientContext.CancellationToken);
return(ForwardMessageToApplication(connectionId, openConnectionMessage));
}
else
{
// the manager still contains this connectionId, probably this connection is not yet cleaned up
Log.DuplicateConnectionId(Logger, connectionId, null);
return(SafeWriteAsync(
new CloseConnectionMessage(connectionId, $"Duplicate connection ID {connectionId}")));
}
}
}
private async Task ProcessMessageAsync(ClientConnectionContext clientContext, CancellationToken cancellation)
{
var connectionId = clientContext.ConnectionId;
try
{
// Check if channel is closed.
while (await clientContext.Input.WaitToReadAsync(cancellation))
{
while (clientContext.Input.TryRead(out var serviceMessage))
{
cancellation.ThrowIfCancellationRequested();
switch (serviceMessage)
{
case OpenConnectionMessage openConnectionMessage:
await OnConnectedAsyncCore(clientContext, openConnectionMessage);
break;
case CloseConnectionMessage closeConnectionMessage:
// should not wait for application task when inside the application task
// As the messages are in a queue, close message should be after all the other messages
await PerformDisconnectCore(closeConnectionMessage.ConnectionId, false);
return;
case ConnectionDataMessage connectionDataMessage:
ProcessOutgoingMessages(clientContext, connectionDataMessage);
break;
default:
break;
}
}
}
}
catch (OperationCanceledException e)
{
Log.SendLoopStopped(Logger, connectionId, e);
}
catch (Exception e)
{
// Internal exception is already caught and here only for channel exception.
// Notify client to disconnect.
Log.SendLoopStopped(Logger, connectionId, e);
_ = PerformDisconnectCore(connectionId, false);
_ = SafeWriteAsync(new CloseConnectionMessage(connectionId, e.Message));
}
}
private async Task OnConnectedAsyncCore(ClientConnectionContext clientContext, OpenConnectionMessage message)
{
var connectionId = message.ConnectionId;
try
{
clientContext.Transport = await _clientConnectionManager.CreateConnection(message);
Log.ConnectedStarting(Logger, connectionId);
}
catch (Exception e)
{
Log.ConnectedStartingFailed(Logger, connectionId, e);
// Should not wait for application task inside the application task
_ = PerformDisconnectCore(connectionId, false);
_ = SafeWriteAsync(new CloseConnectionMessage(connectionId, e.Message));
}
}
public static Task WriteAsync(
this ClientConnectionContext connection,
string connectionId,
object value,
IServiceProtocol protocol,
JsonSerializer serializer,
IMemoryPool pool)
{
using var writer = new MemoryPoolTextWriter(pool);
serializer.Serialize(writer, value);
writer.Flush();
// Reuse ConnectionDataMessage to wrap the payload
var wrapped = new ConnectionDataMessage(string.Empty, writer.Buffer);
var message = new ConnectionDataMessage(connectionId, protocol.GetMessageBytes(wrapped));
return(connection.WriteMessageAsync(message));
}
private void ProcessOutgoingMessages(ClientConnectionContext clientContext, ConnectionDataMessage connectionDataMessage)
{
var connectionId = connectionDataMessage.ConnectionId;
try
{
var payload = connectionDataMessage.Payload;
Log.WriteMessageToApplication(Logger, payload.Length, connectionId);
var message = GetString(payload);
if (message == ReconnectMessage)
{
clientContext.Transport?.Reconnected?.Invoke();
}
else
{
clientContext.Transport?.OnReceived(message);
}
}
catch (Exception e)
{
Log.FailToWriteMessageToApplication(Logger, nameof(ConnectionDataMessage), connectionDataMessage.ConnectionId, connectionDataMessage.TracingId, e);
}
}
private async Task WaitForApplicationTask(ClientConnectionContext clientContext, bool closeGracefully)
{
clientContext.Output.TryComplete();
var app = clientContext.ApplicationTask;
if (!app.IsCompleted)
{
try
{
if (!closeGracefully)
{
clientContext.CancelPendingRead();
}
using (var delayCts = new CancellationTokenSource())
{
var resultTask =
await Task.WhenAny(app, Task.Delay(CloseApplicationTimeout, delayCts.Token));
if (resultTask != app)
{
// Application task timed out and it might never end writing to Transport.Output, cancel reading the pipe so that our ProcessOutgoing ends
clientContext.CancelPendingRead();
Log.ApplicationTaskTimedOut(Logger);
}
else
{
delayCts.Cancel();
}
}
}
catch (Exception ex)
{
Log.ApplicationTaskFailed(Logger, ex);
}
}
}
public bool TryGetClientConnection(string connectionId, out ClientConnectionContext connection)
{
return(_clientConnections.TryGetValue(connectionId, out connection));
}
public bool TryRemoveClientConnection(string connectionId, out ClientConnectionContext connection)
{
return(_clientConnections.TryRemove(connectionId, out connection));
}
public bool TryAddClientConnection(ClientConnectionContext connection)
{
return(_clientConnections.TryAdd(connection.ConnectionId, connection));
}