0
\$\begingroup\$

I made a web socket client wrapper using SignalR's implementation, but there are some downsides that I want to improve.

  • CancellationToken's implementation is wrong/incomplete
  • ReadMultipleMessagesAsync is parsing messages ending with (byte)'}', which is kinda wrong, because it removes the preceding symbol from the parsed JSON and then I'm unable to deserialize it.
  • SignalR's implementation doesn't implement IDisposable, so do I. However, I think my BinanceMessageTransport class should do and call the StopAsync() on Dispose(). Let me know about your opinion.
  • Perhaps, my reconnection logic can be improved
  • the exceptional handling could probably be improved
  • anything else?

The project can be found on GitHub and can also be executed from the Sample App in the tests folder.

Usage

public interface IMessageTransport
{
    /// <summary>
    ///     The received messages from the web socket stream.
    /// </summary>
    IObservable<string> MessageReceived { get; }

    /// <summary>
    ///     Starts the message transport service.
    /// </summary>
    /// <returns></returns>
    Task StartAsync();

    /// <summary>
    ///     Stops the message transport service.
    /// </summary>
    /// <returns></returns>
    Task StopAsync();
}

public sealed class BinanceMessageTransport : IMessageTransport
{
    private const string Endpoint = "wss://stream.binance.com:9443/ws/bnbusdt@trade";
    private readonly ILogger<BinanceMessageTransport> _logger;

    private readonly ISubject<string> _messageReceivedSubject;
    private readonly IWebSocketPipe _webSocketPipe;

    /// <summary>
    ///     Creates a new <see cref="BinanceMessageTransport" /> instance.
    /// </summary>
    public BinanceMessageTransport() : this(NullLoggerFactory.Instance)
    {
    }

    /// <summary>
    ///     Creates a new <see cref="BinanceMessageTransport" /> instance.
    /// </summary>
    /// <param name="loggerFactory">The <see cref="ILoggerFactory" /> instance used for internal logging.</param>
    public BinanceMessageTransport(ILoggerFactory loggerFactory)
    {
        _webSocketPipe = WebSocketPipe.Create(loggerFactory);
        _messageReceivedSubject = new Subject<string>();
        _logger = loggerFactory.CreateLogger<BinanceMessageTransport>();
    }

    /// <summary>
    ///     How long to wait before reconnecting if no message comes from server.
    ///     Set null to disable this feature.
    ///     Default: 1 minute
    /// </summary>
    public TimeSpan? ReconnectTimeout { get; set; } = TimeSpan.FromSeconds(5);

    /// <inheritdoc />
    public IObservable<string> MessageReceived => _messageReceivedSubject.AsObservable();

    /// <inheritdoc />
    public async Task StartAsync()
    {
        try
        {
            await _webSocketPipe.StartAsync(new Uri(Endpoint), CancellationToken.None).ConfigureAwait(false);
        }
        catch (WebSocketException)
        {
            // Reconnect logic
            if (ReconnectTimeout is null)
            {
                throw;
            }

            _logger.LogInformation("About to reconnect in {Time} ...", ReconnectTimeout.Value.ToReadableString());

            await Task.Delay(ReconnectTimeout.Value, CancellationToken.None).ConfigureAwait(false);
            await StartAsync().ConfigureAwait(false);
        }

        await ReadMultipleMessagesAsync(_webSocketPipe.Input, CancellationToken.None).ConfigureAwait(false);
    }

    /// <inheritdoc />
    public async Task StopAsync()
    {
        if (_webSocketPipe.IsRunning)
        {
            await _webSocketPipe.StopAsync().ConfigureAwait(false);
        }
    }

    private static bool TryParseMessage(ref ReadOnlySequence<byte> buffer, out ReadOnlySequence<byte> payload)
    {
        var position = buffer.PositionOf((byte)'}');
        if (position == null)
        {
            payload = default;
            return false;
        }

        payload = buffer.Slice(0, position.Value);

        // Skip record separator
        buffer = buffer.Slice(buffer.GetPosition(1, position.Value));

        return true;
    }

    private static bool TryParseMessage(ref ReadOnlySequence<byte> input, out string message)
    {
        if (!TryParseMessage(ref input, out ReadOnlySequence<byte> payload))
        {
            message = string.Empty;
            return false;
        }

        // Handle "new" call
        var json = Encoding.UTF8.GetString(payload.ToArray());

        message = json;
        return true;
    }

    // Source: https://github.com/davidfowl/DocsStaging/blob/master/Pipelines.md#reading-multiple-messages
    private async Task ReadMultipleMessagesAsync(PipeReader reader, CancellationToken cancellationToken = default)
    {
        try
        {
            while (true)
            {
                var result = await reader.ReadAsync(cancellationToken);
                var buffer = result.Buffer;

                try
                {
                    // Process all messages from the buffer, modifying the input buffer on each iteration
                    while (TryParseMessage(ref buffer, out string message))
                    {
                        _messageReceivedSubject.OnNext(message);
                    }

                    // There's no more data to be processed
                    if (result.IsCompleted)
                    {
                        if (buffer.Length > 0)
                        {
                            // We have an incomplete message and there's no more data to process
                            throw new InvalidDataException("Incomplete message!");
                        }

                        break;
                    }
                }
                finally
                {
                    // Since we're processing all messages in the buffer, we can use the remaining buffer's Start and End
                    // position to determine consumed and examined
                    reader.AdvanceTo(buffer.Start, buffer.End);
                }
            }
        }
        finally
        {
            await reader.CompleteAsync();
        }
    }
}

public class Program
{
    private static readonly ManualResetEvent ExitEvent = new(false);

    private static async Task Main()
    {
        Log.Logger = new LoggerConfiguration()
            .MinimumLevel.Verbose()
            .Enrich.FromLogContext()
            .WriteTo.Console(LogEventLevel.Information, theme: SystemConsoleTheme.Colored)
            .CreateLogger();

        AppDomain.CurrentDomain.ProcessExit += CurrentDomainOnProcessExit;
        AssemblyLoadContext.Default.Unloading += DefaultOnUnloading;
        Console.CancelKeyPress += ConsoleOnCancelKeyPress;

        var disposable = new CompositeDisposable();

        using var loggerFactory = LoggerFactory.Create(builder => { builder.AddSerilog(Log.Logger); });

        IMessageTransport messageTransport = new BinanceMessageTransport(loggerFactory);

        messageTransport.MessageReceived.Subscribe(message => {
            Log.Information("Message received: {Message}", message);
        }).DisposeWith(disposable);

        try
        {
            await messageTransport.StartAsync().ConfigureAwait(false);
        }
        catch (Exception ex) when (ex is WebSocketException or InvalidOperationException or InvalidDataException
                                       or ArgumentNullException)
        {
            Log.Error(ex, "Exception: {Message}", ex.Message);
        }
        finally
        {
            await messageTransport.StopAsync().ConfigureAwait(false);
            disposable.Dispose();
        }

        ExitEvent.WaitOne();

        Log.CloseAndFlush();
    }

    private static void CurrentDomainOnProcessExit(object? sender, EventArgs eventArgs)
    {
        Log.Warning("Exiting process");
        ExitEvent.Set();
    }

    private static void DefaultOnUnloading(AssemblyLoadContext assemblyLoadContext)
    {
        Log.Warning("Unloading process");
        ExitEvent.Set();
    }

    private static void ConsoleOnCancelKeyPress(object? sender, ConsoleCancelEventArgs e)
    {
        Log.Warning("Canceling process");
        e.Cancel = true;
        ExitEvent.Set();
    }
}

SignalR's implementation

public interface IWebSocketPipe : IDuplexPipe
{
    /// <summary>
    ///     Whether the connection is running.
    /// </summary>
    bool IsRunning { get; }

    /// <summary>
    ///     Initializes and starts the web socket client connection. It can be stopped using <see cref="StopAsync" />.
    /// </summary>
    /// <param name="uri">The web socket URI.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken" /> instance.</param>
    /// <returns></returns>
    Task StartAsync(Uri uri, CancellationToken cancellationToken);

    /// <summary>
    ///     Stops the web socket client connection.
    /// </summary>
    /// <returns></returns>
    Task StopAsync();
}

internal sealed class SimpleWebSocketPipe : IWebSocketPipe
{
    private readonly TimeSpan _closeTimeout;
    private readonly ILogger _logger;
    private volatile bool _aborted;
    private IDuplexPipe? _application;
    private IDuplexPipe? _transport;
    private WebSocket? _webSocket;

    public SimpleWebSocketPipe(ILoggerFactory loggerFactory)
    {
        _logger = loggerFactory.CreateLogger<SimpleWebSocketPipe>();

        _closeTimeout = TimeSpan.FromMilliseconds(10_000);
    }

    internal Task Running { get; private set; } = Task.CompletedTask;

    public bool IsRunning => _webSocket?.State == WebSocketState.Open;

    public PipeReader Input => _transport!.Input;

    public PipeWriter Output => _transport!.Output;

    public async Task StartAsync(Uri uri, CancellationToken cancellationToken)
    {
        if (!Running.IsCompleted)
        {
            throw new InvalidOperationException("The web socket has already started.");
        }

        if (uri is null)
        {
            throw new ArgumentNullException(nameof(uri));
        }

        var factory = DefaultWebSocketFactory;
        _webSocket = await factory(uri, cancellationToken).ConfigureAwait(false);

        if (_webSocket == null)
        {
            throw new InvalidOperationException("Configured WebSocketFactory did not return a value.");
        }

        _logger.LogDebug("Started transport.");

        var pipeOptions = new PipeOptions(
            readerScheduler: PipeScheduler.ThreadPool,
            writerScheduler: PipeScheduler.ThreadPool,
            useSynchronizationContext: false);

        var pair = DuplexPipe.CreateConnectionPair(pipeOptions, pipeOptions);

        _transport = pair.Transport;
        _application = pair.Application;

        Running = ProcessAsync(_webSocket);
    }

    public async Task StopAsync()
    {
        if (Running.IsCompleted)
        {
            throw new InvalidOperationException("The web socket hasn't started yet.");
        }

        _logger.LogInformation("Transport is stopping.");

        if (_application is null)
        {
            // We never started
            return;
        }

        await _transport!.Input.CompleteAsync().ConfigureAwait(false);
        await _transport!.Output.CompleteAsync().ConfigureAwait(false);

        // Cancel any pending reads from the application, this should start the entire shutdown process
        _application.Input.CancelPendingRead();

        try
        {
            await Running.ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Transport stopped.");
            // exceptions have been handled in the Running task continuation by closing the channel with the exception
            return;
        }
        finally
        {
            _webSocket?.Dispose();
        }

        _logger.LogDebug("Transport stopped.");
    }

    private static async ValueTask<WebSocket> DefaultWebSocketFactory(Uri uri, CancellationToken cancellationToken)
    {
        var webSocket = new ClientWebSocket();

        try
        {
            await webSocket.ConnectAsync(uri, cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            webSocket.Dispose();
            throw;
        }

        return webSocket;
    }

    private async Task ProcessAsync(WebSocket socket)
    {
        Debug.Assert(_application != null);

        using (socket)
        {
            // Begin sending and receiving.
            var receiving = StartReceiving(socket);
            var sending = StartSending(socket);

            // Wait for send or receive to complete
            var trigger = await Task.WhenAny(receiving, sending).ConfigureAwait(false);

            if (trigger == receiving)
            {
                // We're waiting for the application to finish and there are 2 things it could be doing
                // 1. Waiting for application data
                // 2. Waiting for a websocket send to complete

                // Cancel the application so that ReadAsync yields
                _application.Input.CancelPendingRead();

                using var delayCts = new CancellationTokenSource();

                var resultTask = await Task.WhenAny(sending, Task.Delay(_closeTimeout, delayCts.Token))
                    .ConfigureAwait(false);

                if (resultTask != sending)
                {
                    _aborted = true;

                    // Abort the websocket if we're stuck in a pending send to the client
                    socket.Abort();
                }
                else
                {
                    // Cancel the timeout
                    delayCts.Cancel();
                }
            }
            else
            {
                // We're waiting on the websocket to close and there are 2 things it could be doing
                // 1. Waiting for websocket data
                // 2. Waiting on a flush to complete (backpressure being applied)

                _aborted = true;

                // Abort the websocket if we're stuck in a pending receive from the client
                socket.Abort();

                // Cancel any pending flush so that we can quit
                _application.Output.CancelPendingFlush();
            }
        }
    }

    private async Task StartReceiving(WebSocket webSocket)
    {
        Debug.Assert(_application != null);

        try
        {
            while (true)
            {
                // Do a 0 byte read so that idle connections don't allocate a buffer when waiting for a read
                var result = await webSocket.ReceiveAsync(Memory<byte>.Empty, CancellationToken.None)
                    .ConfigureAwait(false);

                if (result.MessageType == WebSocketMessageType.Close)
                {
                    _logger.LogInformation("WebSocket closed by the server. Close status {CloseStatus}.",
                        webSocket.CloseStatus);

                    if (webSocket.CloseStatus != WebSocketCloseStatus.NormalClosure)
                    {
                        throw new InvalidOperationException($"Websocket closed with error: {webSocket.CloseStatus}.");
                    }

                    return;
                }

                var memory = _application.Output.GetMemory();

                // Because we checked the CloseStatus from the 0 byte read above, we don't need to check again after reading
                var receiveResult = await webSocket.ReceiveAsync(memory, CancellationToken.None).ConfigureAwait(false);

                // Need to check again for netstandard2.1 because a close can happen between a 0-byte read and the actual read
                if (receiveResult.MessageType == WebSocketMessageType.Close)
                {
                    _logger.LogInformation("WebSocket closed by the server. Close status {CloseStatus}.",
                        webSocket.CloseStatus);

                    if (webSocket.CloseStatus != WebSocketCloseStatus.NormalClosure)
                    {
                        throw new InvalidOperationException($"Websocket closed with error: {webSocket.CloseStatus}.");
                    }

                    return;
                }

                _logger.LogDebug("Message received. Type: {MessageType}, size: {Count}, EndOfMessage: {EndOfMessage}.",
                    receiveResult.MessageType, receiveResult.Count, receiveResult.EndOfMessage);

                _application.Output.Advance(receiveResult.Count);

                var flushResult = await _application.Output.FlushAsync().ConfigureAwait(false);

                // We canceled in the middle of applying back pressure
                // or if the consumer is done
                if (flushResult.IsCanceled || flushResult.IsCompleted)
                {
                    break;
                }
            }
        }
        catch (OperationCanceledException)
        {
            _logger.LogDebug("Receive loop canceled.");
        }
        catch (Exception ex)
        {
            if (!_aborted)
            {
                await _application.Output.CompleteAsync(ex).ConfigureAwait(false);
            }
        }
        finally
        {
            // We're done writing
            await _application.Output.CompleteAsync().ConfigureAwait(false);

            _logger.LogDebug("Receive loop stopped.");
        }
    }

    private async Task StartSending(WebSocket webSocket)
    {
        Debug.Assert(_application != null);

        Exception? error = null;

        try
        {
            while (true)
            {
                var result = await _application.Input.ReadAsync();
                var buffer = result.Buffer;

                // Get a frame from the application
                try
                {
                    if (result.IsCanceled)
                    {
                        break;
                    }

                    if (!buffer.IsEmpty)
                    {
                        try
                        {
                            _logger.LogDebug("Received message from application. Payload size: {Count}.",
                                buffer.Length);

                            if (WebSocketCanSend(webSocket))
                            {
                                await webSocket.SendAsync(buffer, WebSocketMessageType.Text);
                            }
                            else
                            {
                                break;
                            }
                        }
                        catch (Exception ex)
                        {
                            if (!_aborted)
                            {
                                _logger.LogError(ex, "Error while sending a message.");
                            }

                            break;
                        }
                    }
                    else if (result.IsCompleted)
                    {
                        break;
                    }
                }
                finally
                {
                    _application.Input.AdvanceTo(buffer.End);
                }
            }
        }
        catch (Exception ex)
        {
            error = ex;
        }
        finally
        {
            if (WebSocketCanSend(webSocket))
            {
                try
                {
                    // We're done sending, send the close frame to the client if the websocket is still open
                    await webSocket.CloseOutputAsync(
                        error != null ? WebSocketCloseStatus.InternalServerError : WebSocketCloseStatus.NormalClosure,
                        "", CancellationToken.None);
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Closing webSocket failed.");
                }
            }

            await _application.Input.CompleteAsync().ConfigureAwait(false);

            _logger.LogDebug("Send loop stopped.");
        }

        static bool WebSocketCanSend(WebSocket ws)
        {
            return ws.State is not (WebSocketState.Aborted or WebSocketState.Closed or WebSocketState.CloseSent);
        }
    }
}

public static class WebSocketPipe
{
    /// <summary>
    ///     Creates a web socket pipe.
    /// </summary>
    /// <returns>The <see cref="IWebSocketPipe" /> that was created.</returns>
    public static IWebSocketPipe Create()
    {
        return Create(NullLoggerFactory.Instance);
    }

    /// <summary>
    ///     Creates a web socket pipe from the specified <paramref name="loggerFactory" />.
    /// </summary>
    /// <param name="loggerFactory">The <see cref="ILoggerFactory" /> instance used for logging.</param>
    /// <returns>The <see cref="IWebSocketPipe" /> that was created.</returns>
    public static IWebSocketPipe Create(ILoggerFactory loggerFactory)
    {
        if (loggerFactory is null)
        {
            throw new ArgumentNullException(nameof(loggerFactory));
        }

        return new SimpleWebSocketPipe(loggerFactory);
    }
}

internal class DuplexPipe : IDuplexPipe
{
    public DuplexPipe(PipeReader reader, PipeWriter writer)
    {
        Input = reader;
        Output = writer;
    }

    public PipeReader Input { get; }

    public PipeWriter Output { get; }

    public static DuplexPipePair CreateConnectionPair(PipeOptions inputOptions, PipeOptions outputOptions)
    {
        var input = new Pipe(inputOptions);
        var output = new Pipe(outputOptions);

        var transportToApplication = new DuplexPipe(output.Reader, input.Writer);
        var applicationToTransport = new DuplexPipe(input.Reader, output.Writer);

        return new DuplexPipePair(applicationToTransport, transportToApplication);
    }

    // This class exists to work around issues with value tuple on .NET Framework
    public readonly struct DuplexPipePair
    {
        public IDuplexPipe Transport { get; }
        public IDuplexPipe Application { get; }

        public DuplexPipePair(IDuplexPipe transport, IDuplexPipe application)
        {
            Transport = transport;
            Application = application;
        }
    }
}

internal static class WebSocketExtensions
{
    public static ValueTask SendAsync(this WebSocket webSocket, ReadOnlySequence<byte> buffer,
        WebSocketMessageType webSocketMessageType, CancellationToken cancellationToken = default)
    {
        return buffer.IsSingleSegment
            ? webSocket.SendAsync(buffer.First, webSocketMessageType, true, cancellationToken)
            : SendMultiSegmentAsync(webSocket, buffer, webSocketMessageType, cancellationToken);
    }

    private static async ValueTask SendMultiSegmentAsync(WebSocket webSocket, ReadOnlySequence<byte> buffer,
        WebSocketMessageType webSocketMessageType, CancellationToken cancellationToken = default)
    {
        var position = buffer.Start;
        // Get a segment before the loop so we can be one segment behind while writing
        // This allows us to do a non-zero byte write for the endOfMessage = true send
        buffer.TryGet(ref position, out var prevSegment);
        while (buffer.TryGet(ref position, out var segment))
        {
            await webSocket.SendAsync(prevSegment, webSocketMessageType, false, cancellationToken);

            prevSegment = segment;
        }

        // End of message frame
        await webSocket.SendAsync(prevSegment, webSocketMessageType, true, cancellationToken);
    }
}

\$\endgroup\$
6
  • \$\begingroup\$ public BinanceMessageTransport() : this(NullLoggerFactory.Instance) << This looks really strange to me. Did you define a parameterless ctor for ease testing? \$\endgroup\$ Dec 10 '21 at 9:48
  • \$\begingroup\$ @PeterCsala, the purpose of that constructor is to make ILoggerFactory loggerFactory optional. I could probably do ILoggerFactory? loggerFactory, but I checked Microsoft's implementation of ILoggerFactory and other classes, and they seem to be using it that way. In other words, I want to be able to do new BinanceMessageTransport() and new BinanceMessageTransport(loggerFactory);. \$\endgroup\$
    – nop
    Dec 10 '21 at 11:03
  • \$\begingroup\$ I quite not sure how will it work with the DI? If you have two constructors which one will be used whenever the DI instantiates the object? \$\endgroup\$ Dec 10 '21 at 11:13
  • 1
    \$\begingroup\$ @PeterCsala, I believe it's okay, because BinanceMessageTransport won't be DI'd. If it has to be, I will remove the first constructor, because I believe multiple constructors during DI is an anti-pattern. \$\endgroup\$
    – nop
    Dec 10 '21 at 12:08
  • \$\begingroup\$ Cleaner code would be to remove the multiple constructors and have a single one. Just do a check for null on your ILoggerFactory, and if so - set the NullLoggerFactory instead. That way you keep your ability to mock for unit testing, and you can use it with DI because it only has the single constructor. Multiple constructors will NOT work with Dependency Injection because it wont know which one to use \$\endgroup\$ Jan 6 at 13:58

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.