2
\$\begingroup\$

I'm working on a web socket client wrapper around ClientWebSocket and this is kinda my first public library project, which is why I'm asking this question here. I was inspired by this web socket client wrapper, because it uses Rx .NET instead of regular .NET events and I also wanted to give System.IO.Pipelines a try.

There are a few things I'm not quite certain about:

  • I should probably move ClientWebSocket and CancellationTokenSource's instantiation into the StartAsync method. That way people will be able to start and stop the web socket stream anytime they wish. I think it's even necessary, because with the current design, I'm not able to implement reconnect logic.

  • The call to ListenAsync in StartAsync is currently blocking the UI thread. I'm not sure if that's good by design, because I saw how people did it here. They split the write and read logic, just like Microsoft's example here. Microsoft did the same in SignalR's implementation.

  • I think ClientWebSocket has to be instatiated just like they did here (a factory), because that way the users can take advantage of the options that ClientWebSocket provides. For ex. the proxy settings, etc.

  • Perhaps a SemaphoreSlim would be useful between the calls in the methods. Example usage here.

  • I'm not quite sure if I really need the second while loop in ListenAsync

  • Suggestions about the logging and what I did with the LogManager class?

The idea of the library is high-performance (System.IO.Pipelines), Rx Events (System.Reactive), proper logging and error handling, connection lost and connection restored events (not implemented yet) and reconnection based on a specific TimeSpan (also missing). I would like to know how you guys would've done it.

public interface IWebSocketClient : IDisposable
{
    /// <summary>
    ///     Client identification.
    /// </summary>
    string ClientId { get; }

    /// <summary>
    ///     Web socket state.
    /// </summary>
    WebSocketState State { get; }

    /// <summary>
    ///     Keep alive interval (heartbeat).
    /// </summary>
    TimeSpan KeepAliveInterval { get; set; }

    /// <summary>
    ///     Stream for received message (raw format).
    /// </summary>
    IObservable<string> MessageReceived { get; }

    /// <summary>
    ///     Connects to the web socket server and starts listening for messages.
    /// </summary>
    /// <param name="cancellationToken">The cancellation token.</param>
    /// <returns></returns>
    Task StartAsync(CancellationToken cancellationToken);
}

public class WebSocketClient : IWebSocketClient
{
    private const int ReceiveBufferSize = 2048;

    private readonly ISubject<string> _messageReceivedSubject = new Subject<string>();
    private readonly WebSocketClientOptions _options;
    private readonly Pipe _pipe;
    private readonly ClientWebSocket _webSocket;

    /// <summary>
    ///     Whether this instance is disposed.
    /// </summary>
    private bool _isDisposed;

    /// <summary>
    ///     The log manager object.
    /// </summary>
    protected internal LogManager LogManager;

    /// <summary>
    ///     Initializes a new instance of the <see cref="WebSocketClient" /> class.
    /// </summary>
    /// <param name="options"></param>
    protected WebSocketClient(WebSocketClientOptions options)
    {
        Guard.Against.Null(options, nameof(options));

        _options = options;
        LogManager = new LogManager
        {
            Level = options.LogLevel ?? LogLevel.Information,
            Writers = options.LogWriters ?? new List<ILogger>()
        };

        _webSocket = new ClientWebSocket();
        _pipe = new Pipe(new PipeOptions(useSynchronizationContext: false));
    }

    /// <summary>
    ///     The pipe writer.
    /// </summary>
    private PipeWriter Writer => _pipe.Writer;

    /// <summary>
    ///     The pipe reader.
    /// </summary>
    private PipeReader Reader => _pipe.Reader;

    /// <inheritdoc />
    public string ClientId { get; } = Guid.NewGuid().ToString();

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

    /// <inheritdoc />
    public TimeSpan KeepAliveInterval
    {
        get => _webSocket.Options.KeepAliveInterval;
        set => _webSocket.Options.KeepAliveInterval = value;
    }

    /// <inheritdoc />
    public WebSocketState State => _webSocket.State;

    /// <summary>
    ///     Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
    /// </summary>
    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    async Task IWebSocketClient.StartAsync(CancellationToken cancellationToken)
    {
        ThrowIfDisposed();

        Guard.Against.Null(_options.BaseAddress, nameof(_options.BaseAddress));

        if (State == WebSocketState.Open)
            throw new InvalidOperationException("Client was already connected");

        try
        {
            await _webSocket.ConnectAsync(_options.BaseAddress, cancellationToken).ConfigureAwait(false);

            LogManager.Write(LogLevel.Debug, $"Socket {ClientId} connected");
        }
        catch (Exception ex) when (ex is WebSocketException)
        {
            LogManager.Write(LogLevel.Error, ex.Message);
        }

        await ListenAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <summary>
    ///     Factory method.
    /// </summary>
    /// <returns></returns>
    public static IWebSocketClient Create(WebSocketClientOptions options) => new WebSocketClient(options);

    private async Task ListenAsync(CancellationToken cancellationToken)
    {
        while (State == WebSocketState.Open && !cancellationToken.IsCancellationRequested && !_isDisposed)
        {
            try
            {
                var bytesRead = await _webSocket
                    .ReceiveAsync(Writer.GetMemory(ReceiveBufferSize), cancellationToken)
                    .ConfigureAwait(false);

                // TODO: I don't think this loop is needed at all
                while (!cancellationToken.IsCancellationRequested && !bytesRead.EndOfMessage &&
                       bytesRead.MessageType != WebSocketMessageType.Close && !_isDisposed)
                {
                    if (bytesRead.Count == 0)
                        break;

                    Writer.Advance(bytesRead.Count);

                    bytesRead = await _webSocket
                        .ReceiveAsync(Writer.GetMemory(ReceiveBufferSize), cancellationToken)
                        .ConfigureAwait(false);
                }

                // We didn't get a complete message, we can't flush partial message.
                if (cancellationToken.IsCancellationRequested || !bytesRead.EndOfMessage ||
                    bytesRead.MessageType == WebSocketMessageType.Close || _isDisposed)
                    break;

                // Advance the EndOfMessage bytes before flushing
                Writer.Advance(bytesRead.Count);
                var flushResult = await Writer.FlushAsync(cancellationToken).ConfigureAwait(false);
                if (flushResult.IsCompleted)
                    break;

                // TODO: Should probably be using the following template: https://docs.microsoft.com/en-us/dotnet/standard/io/pipelines#pipe-basic-usage оr https://github.com/teraa/Tw.Net/blob/b5fd6f1ee184316a9a6a2f1eec46698ea08bf682/src/Common/Clients/WebSocketClient.cs#L211
                if (bytesRead.MessageType != WebSocketMessageType.Close)
                    ReadPipeAsync();
            }
            catch (Exception ex) when (ex is OperationCanceledException or WebSocketException
                                           or InvalidOperationException)
            {
                LogManager.Write(LogLevel.Error, ex.Message);

                break;
            }
        }

        await CompleteAsync(cancellationToken, _webSocket.CloseStatus, _webSocket.CloseStatusDescription);
    }

    private void ReadPipeAsync()
    {
        if (Reader.TryRead(out var readResult) && readResult.Buffer.Length > 0)
        {
            var json = Encoding.UTF8.GetString(readResult.Buffer);

            Reader.AdvanceTo(readResult.Buffer.End);

            _messageReceivedSubject.OnNext(json);
        }
    }

    private async Task CompleteAsync(
        CancellationToken cancellationToken,
        WebSocketCloseStatus? closeStatus = default,
        string? closeStatusDescription = default)
    {
        await Reader.CompleteAsync().ConfigureAwait(false);
        await Writer.CompleteAsync().ConfigureAwait(false);

        if (_options.CloseWhenCompleted || closeStatus != null)
        {
            await CloseAsync(closeStatus ?? WebSocketCloseStatus.NormalClosure, closeStatusDescription ?? string.Empty,
                cancellationToken);
        }
    }

    private async Task CloseAsync(
        WebSocketCloseStatus closeStatus,
        string closeStatusDescription,
        CancellationToken cancellationToken)
    {
        if (State is WebSocketState.Closed or WebSocketState.CloseSent or WebSocketState.Aborted)
            return;

        await _webSocket.CloseAsync(closeStatus, closeStatusDescription, cancellationToken);
    }

    /// <summary>
    ///     Throw if disposed.
    /// </summary>
    /// <exception cref="ObjectDisposedException"></exception>
    private void ThrowIfDisposed()
    {
        if (_isDisposed)
        {
            throw new ObjectDisposedException($"{nameof(WebSocketClient)} has been disposed");
        }
    }

    /// <summary>
    ///     The standard dispose destructor.
    /// </summary>
    ~WebSocketClient()
    {
        Dispose(false);
    }

    /// <summary>
    ///     Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
    /// </summary>
    /// <param name="disposing">If this method is called by a user's code.</param>
    protected virtual void Dispose(bool disposing)
    {
        if (_isDisposed)
        {
            return;
        }

        if (disposing)
        {
            _webSocket.Abort();
            _webSocket.Dispose();

            _messageReceivedSubject.OnCompleted();
        }

        _isDisposed = true;
    }
}

/// <summary>
///     Configures the behavior of <see cref="IWebSocketClient" />.
/// </summary>
public class WebSocketClientOptions
{
    /// <summary>
    ///     The base address.
    /// </summary>
    public Uri? BaseAddress { get; set; }

    /// <summary>
    ///     Whether to close the underlying <see cref="System.Net.WebSockets.WebSocket" /> when
    ///     either the <see cref="IDuplexPipe.Input" /> or <see cref="IDuplexPipe.Output" />
    ///     is marked as completed. Defaults to <see langword="false" />.
    /// </summary>
    public bool CloseWhenCompleted { get; set; }

    /// <summary>
    ///     The log level.
    /// </summary>
    public LogLevel? LogLevel { get; set; }

    /// <summary>
    ///     Collection of ILogger implementations to forward the message to.
    /// </summary>
    public ICollection<ILogger>? LogWriters { get; set; }
}

/// <summary>
///     Log manager implementation.
/// </summary>
public class LogManager
{
    /// <summary>
    ///     The log level.
    /// </summary>
    public LogLevel Level { get; set; }

    /// <summary>
    ///     Collection of ILogger implementations to forward the message to.
    /// </summary>
    public ICollection<ILogger> Writers { get; set; } = null!;

    /// <summary>
    ///     Write a log entry
    /// </summary>
    /// <param name="logLevel">The verbosity of the message</param>
    /// <param name="message">The message to log</param>
    public void Write(LogLevel logLevel, string message)
    {
        if ((int) logLevel < (int) Level)
            return;

        foreach (var writer in Writers)
        {
            try
            {
                writer.Log(logLevel, message);
            }
            catch (Exception ex)
            {
                Debug.WriteLine($"Failed to write log to writer {writer.GetType()}: {ex.Message}");
            }
        }
    }
}

/// <summary>
///     A helper class for disposables.
/// </summary>
public static class DisposableExtensions
{
    /// <summary>
    ///     Ensures the provided disposable is disposed with the specified
    ///     <paramref name="collection" />
    ///     .
    /// </summary>
    /// <typeparam name="T">The type of the disposable.</typeparam>
    /// <param name="disposable">The disposable.</param>
    /// <param name="collection">
    ///     The collection to which
    ///     <paramref name="disposable" />
    ///     will be added.
    /// </param>
    /// <returns>The disposable.</returns>
    public static T DisposeWith<T>(this T disposable, ICollection<IDisposable> collection) where T : IDisposable
    {
        Guard.Against.Null(disposable, nameof(disposable));
        Guard.Against.Null(collection, nameof(collection));

        collection.Add(disposable);
        return disposable;
    }
}

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.Debug, theme: SystemConsoleTheme.Colored)
            .CreateLogger();

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

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

        var disposable = new CompositeDisposable();

        var uri = new Uri("wss://stream.binance.com:9443/ws/bnbusdt@trade");

        var client = WebSocketClient.Create(new WebSocketClientOptions
        {
            BaseAddress = uri,
            CloseWhenCompleted = true,
            LogLevel = LogLevel.Debug,
            LogWriters = new List<ILogger>
            {
                loggerFactory.CreateLogger<WebSocketClient>()
            }
        });

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

        await client.StartAsync(CancellationToken.None).ConfigureAwait(false);

        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();
    }
}
\$\endgroup\$
3
  • 1
    \$\begingroup\$ What's with catch (Exception ex) when (ex is WebSocketException)? Can't this just be catch (WebSocketException ex)? \$\endgroup\$
    – Reinderien
    Dec 8 '21 at 15:15
  • \$\begingroup\$ @Reinderien, it's identical \$\endgroup\$
    – nop
    Dec 8 '21 at 20:09
  • \$\begingroup\$ @nop it is identical neither in C#, nor in Runtime. \$\endgroup\$
    – aepot
    Dec 18 '21 at 15:35

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.