General web socket client with auto reconnect capabilities

Question

This is a web socket client wrapper with auto reconnect capabilities.

First of all, I know I should avoid working with strings to reduce allocation, that's on test purpose. So, ignore the fact that it is an unbounded Channel<string>, which is not directly parsing to JSON from ReadOnlyMemory<byte> (stringifies instead) and doesn't have a max message size.

I want to have a code review on everything else, i.e. the way it starts/stops/reconnects the web socket, the way I use both classes together, etc.

There are two general clients in the code below:

1. GeneralClient2. The initial idea was to signal exit by completing the writer as it's preferred by Microsoft and Marc Gravell. In other words, calling .Complete/.TryComplete on the writer, completes the reading loop (ProcessSendAsync/ProcessDataAsync).

Edit:

I realized completing the writer is not a good idea in this case because:

• It requires the Channel<T> to be nullable because a completed writer cannot be reused. This led to nullability checks in SendAsync, which completely killed the point of the channels as I was using them in oppose to ConcurrentQueue + AsyncResetEvent (similar to what's been done here). In other words, I want to be able to enqueue messages anytime I want and not to be getting ChannelClosedException because of some bad timing.

• Channels have the best performance (producer/consumer pattern) compared to DataFlow, BlockingCollection and basically anything else. If SendAsync is called from multiple threads at the same time, it will throw an exception. It can be avoided using AsyncResetEvent/SemaphoreSlim, etc. Channels allow us to specify SingleReader = true, which acts as a lock. We can make it bounded and restrict it more, so the clien't can't DOS us.

2. GeneralClient uses a CancellationToken to abort the websocket - Similar to this example, which is well documented, even though, that example first tries to gracefully close the socket and if something went side ways, it then aborts it. You can express your opinion on this one, but I would rather have a review on GeneralClient2 as it's the preferred way.

using System.Buffers;
using System.Diagnostics;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using System.Threading.Channels;
using Nito.AsyncEx;

namespace CodeReview;

public sealed class GeneralClient2 : IDisposable
{
    private readonly string _url;
    private readonly SemaphoreSlim _semaphore = new(1, 1);

    private readonly Channel<string> _incomingMessages = Channel.CreateUnbounded<string>(new UnboundedChannelOptions
    {
        SingleReader = false,
        SingleWriter = true
    });

    private readonly Channel<string> _outgoingMessages = Channel.CreateUnbounded<string>(new UnboundedChannelOptions
    {
        SingleReader = true,
        SingleWriter = false
    });

    private ClientWebSocket? _clientWebSocket;
    private CancellationTokenSource? _tokenSource;
    private Task _processingSend = Task.CompletedTask;
    private Task _processingData = Task.CompletedTask;
    private Task _processingReceive = Task.CompletedTask;

    public GeneralClient2(string url)
    {
        if (string.IsNullOrWhiteSpace(url))
        {
            throw new ArgumentNullException(nameof(url));
        }

        _url = url;
    }

    public bool IsRunning { get; private set; }

    public event EventHandler? Connected;
    public event EventHandler? Disconnected;
    public event EventHandler<MessageReceivedEventArgs>? MessageReceived;

    public void Dispose()
    {
        _semaphore.Dispose();

        _incomingMessages.Writer.TryComplete();
        _outgoingMessages.Writer.TryComplete();
    }

    public async Task StartAsync()
    {
        // Prevent a race condition
        await _semaphore.WaitAsync().ConfigureAwait(false);

        try
        {
            if (IsRunning)
            {
                return;
            }

            while (!await ConnectAsync().ConfigureAwait(false))
            {
            }

            IsRunning = true;

            Debug.Assert(_clientWebSocket != null);

            _tokenSource = new CancellationTokenSource();

            _processingSend = ProcessSendAsync(_clientWebSocket, _tokenSource.Token);
            _processingData = ProcessDataAsync(_tokenSource.Token);
            _processingReceive = ProcessReceiveAsync(_clientWebSocket);
        }
        finally
        {
            _semaphore.Release();
        }
    }

    public async Task StopAsync()
    {
        Console.WriteLine("Stopping");

        if (!IsRunning)
        {
            return;
        }

        try
        {
            if (_clientWebSocket is { State: not (WebSocketState.Aborted or WebSocketState.Closed or WebSocketState.CloseSent) })
            {
                await _clientWebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false);
            }
        }
        catch
        {
            // Any exception thrown here will be caused by the socket already being closed,
            // which is the state we want to put it in by calling this method, which
            // means we don't care if it was already closed and threw an exception
            // when we tried to close it again.
        }

        await _processingReceive.ConfigureAwait(false);

        Console.WriteLine("Stopped");
    }

    public ValueTask SendAsync(string message)
    {
        return _outgoingMessages.Writer.WriteAsync(message);
    }

    private async ValueTask<bool> ConnectAsync()
    {
        Console.WriteLine("Connecting");

        var ws = new ClientWebSocket();

        try
        {
            await ws.ConnectAsync(new Uri(_url), CancellationToken.None).ConfigureAwait(false);

            Connected?.Invoke(this, EventArgs.Empty);
        }
        catch (Exception) // TaskCanceledException & WebSocketException
        {
            ws.Dispose();
            return false;
        }

        _clientWebSocket = ws;

        Console.WriteLine("Connected");

        return true;
    }

    private async Task ProcessSendAsync(WebSocket webSocket, CancellationToken cancellationToken)
    {
        try
        {
            while (await _outgoingMessages.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
            {
                while (_outgoingMessages.Reader.TryRead(out var message))
                {
                    // "SingleReader = true" acts as a lock.
                    // The lock is required because the client will throw an exception if SendAsync is 
                    // called from multiple threads at the same time. But this issue only happens with several
                    // framework versions.
                    var data = new ArraySegment<byte>(Encoding.UTF8.GetBytes(message));
                    await webSocket.SendAsync(data, WebSocketMessageType.Text, true, CancellationToken.None).ConfigureAwait(false);
                }
            }
        }
        catch (OperationCanceledException)
        {
            // normal upon task/token cancellation, disregard
        }

        Console.WriteLine("Send loop end");
    }

    private async Task ProcessDataAsync(CancellationToken cancellationToken)
    {
        try
        {
            while (await _incomingMessages.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
            {
                while (_incomingMessages.Reader.TryRead(out var message))
                {
                    await ProcessMessageAsync(message).ConfigureAwait(false);
                }
            }
        }
        catch (OperationCanceledException)
        {
            // normal upon task/token cancellation, disregard
        }

        Console.WriteLine("Data loop end");
    }

    private Task ProcessMessageAsync(string message)
    {
        MessageReceived?.Invoke(this, new MessageReceivedEventArgs(message));
        return Task.CompletedTask;
    }

    private async Task ProcessReceiveAsync(WebSocket webSocket)
    {
        Debug.Assert(_incomingMessages != null && _outgoingMessages != null);

        try
        {
            while (true)
            {
                ValueWebSocketReceiveResult receiveResult;

                using var buffer = MemoryPool<byte>.Shared.Rent(4096);
                await using var ms = new MemoryStream(buffer.Memory.Length);
                do
                {
                    receiveResult = await webSocket.ReceiveAsync(buffer.Memory, CancellationToken.None).ConfigureAwait(false);

                    if (receiveResult.MessageType == WebSocketMessageType.Close)
                    {
                        break;
                    }

                    await ms.WriteAsync(buffer.Memory[..receiveResult.Count], CancellationToken.None).ConfigureAwait(false);
                } while (!receiveResult.EndOfMessage);

                ms.Seek(0, SeekOrigin.Begin);

                if (receiveResult.MessageType == WebSocketMessageType.Text)
                {
                    using var reader = new StreamReader(ms, Encoding.UTF8);
                    var message = await reader.ReadToEndAsync().ConfigureAwait(false);

                    await _incomingMessages.Writer.WriteAsync(message, CancellationToken.None).ConfigureAwait(false);
                }
                else if (receiveResult.MessageType == WebSocketMessageType.Close)
                {
                    await CloseAsync().ConfigureAwait(false);

                    return;
                }
            }
        }
        catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
        {
            Console.WriteLine("WebSocketException prematurely");

            await CloseAsync().ConfigureAwait(false);

            await StartAsync().ConfigureAwait(false);
        }
    }

    private async Task CloseAsync()
    {
        // Cancel loops
        _tokenSource?.Cancel();
        _tokenSource?.Dispose();

        // Wait for the tasks to finish
        await Task.WhenAll(_processingSend, _processingData).ConfigureAwait(false);

        // Prevent a leak
        _clientWebSocket?.Dispose();

        Disconnected?.Invoke(this, EventArgs.Empty);

        IsRunning = false;
    }
}

public sealed class GeneralClient : IDisposable
{
    private readonly string _url;
    private readonly Channel<string> _incomingMessages;
    private readonly Channel<string> _outgoingMessages;
    private readonly AsyncManualResetEvent _connectionResetEvent = new(false);

    private ClientWebSocket? _clientWebSocket;
    private CancellationTokenSource? _tokenSource;
    private Task _processingSend = Task.CompletedTask;
    private Task _processingData = Task.CompletedTask;
    private Task _processingLoop = Task.CompletedTask;

    public GeneralClient(string url)
    {
        if (string.IsNullOrWhiteSpace(url))
        {
            throw new ArgumentNullException(nameof(url));
        }

        _url = url;

        _incomingMessages = Channel.CreateUnbounded<string>(new UnboundedChannelOptions
        {
            SingleReader = false,
            SingleWriter = true
        });

        _outgoingMessages = Channel.CreateUnbounded<string>(new UnboundedChannelOptions
        {
            SingleReader = true,
            SingleWriter = false
        });
    }

    public event EventHandler? Connected;
    public event EventHandler? Disconnected;
    public event EventHandler<MessageReceivedEventArgs>? MessageReceived;

    public void Dispose()
    {
        _incomingMessages.Writer.TryComplete();
        _outgoingMessages.Writer.TryComplete();
    }

    public Task StartAsync()
    {
        // avoid race conditions
        if (!_processingLoop.IsCompleted)
        {
            return Task.CompletedTask;
        }

        _tokenSource = new CancellationTokenSource();

        _processingSend = ProcessSendAsync(_tokenSource.Token);
        _processingData = ProcessDataAsync(_tokenSource.Token);

        _processingLoop = Task.Run(async () =>
        {
            while (!await ConnectAsync().ConfigureAwait(false))
            {
            }
        });

        return Task.CompletedTask;
    }

    public async Task StopAsync()
    {
        if (_processingLoop.IsCompleted)
        {
            return;
        }

        Console.WriteLine("Stopping");

        _tokenSource?.Cancel();
        _tokenSource?.Dispose();

        await Task.WhenAll(_processingSend, _processingData, _processingLoop).ConfigureAwait(false);

        Disconnected?.Invoke(this, EventArgs.Empty);

        Console.WriteLine("Stopped");
    }

    private async ValueTask<bool> ConnectAsync()
    {
        Console.WriteLine("Connecting");

        Debug.Assert(_tokenSource != null);

        using var ws = new ClientWebSocket();

        try
        {
            // This call timeouts in 20 seconds with WebSocketException by default
            await ws.ConnectAsync(new Uri(_url), CancellationToken.None).ConfigureAwait(false);

            Connected?.Invoke(this, EventArgs.Empty);
        }
        catch (Exception)
        {
            return false;
        }

        _clientWebSocket = ws;

        Console.WriteLine("Connected");

        _connectionResetEvent.Set();

        try
        {
            await ProcessReceiveAsync(ws, _tokenSource.Token).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            //  normal upon task/token cancellation, disregard
        }
        catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
        {
            return false;
        }
        catch (Exception)
        {
            return false;
        }

        _connectionResetEvent.Reset();

        Console.WriteLine("End");

        return true;
    }

    public void Send(string message)
    {
        _outgoingMessages.Writer.TryWrite(message);
    }

    public ValueTask SendAsync(string message)
    {
        return _outgoingMessages.Writer.WriteAsync(message);
    }

    private async Task ProcessSendAsync(CancellationToken cancellationToken)
    {
        try
        {
            while (await _outgoingMessages.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
            {
                while (_outgoingMessages.Reader.TryRead(out var message))
                {
                    await _connectionResetEvent.WaitAsync(cancellationToken).ConfigureAwait(false);

                    Debug.Assert(_clientWebSocket != null);

                    var data = new ArraySegment<byte>(Encoding.UTF8.GetBytes(message));
                    await _clientWebSocket.SendAsync(data, WebSocketMessageType.Text, true, CancellationToken.None).ConfigureAwait(false);
                }
            }
        }
        catch (OperationCanceledException)
        {
            //  normal upon task/token cancellation, disregard
        }

        Console.WriteLine("Send loop end");
    }

    private async Task ProcessDataAsync(CancellationToken cancellationToken)
    {
        try
        {
            while (await _incomingMessages.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false))
            {
                while (_incomingMessages.Reader.TryRead(out var message))
                {
                    await ProcessMessageAsync(message).ConfigureAwait(false);
                }
            }
        }
        catch (OperationCanceledException)
        {
            //  normal upon task/token cancellation, disregard
        }

        Console.WriteLine("Data loop end");
    }

    private Task ProcessMessageAsync(string message)
    {
        MessageReceived?.Invoke(this, new MessageReceivedEventArgs(message));
        return Task.CompletedTask;
    }

    private async Task ProcessReceiveAsync(WebSocket webSocket, CancellationToken cancellationToken)
    {
        while (true)
        {
            ValueWebSocketReceiveResult receiveResult;

            using var buffer = MemoryPool<byte>.Shared.Rent(4096);
            await using var ms = new MemoryStream(buffer.Memory.Length);
            do
            {
                receiveResult = await webSocket.ReceiveAsync(buffer.Memory, cancellationToken).ConfigureAwait(false);

                if (receiveResult.MessageType == WebSocketMessageType.Close)
                {
                    break;
                }

                await ms.WriteAsync(buffer.Memory[..receiveResult.Count], cancellationToken).ConfigureAwait(false);
            } while (!receiveResult.EndOfMessage);

            ms.Seek(0, SeekOrigin.Begin);

            if (receiveResult.MessageType == WebSocketMessageType.Text)
            {
                using var reader = new StreamReader(ms, Encoding.UTF8);
                var message = await reader.ReadToEndAsync().ConfigureAwait(false);

                await _incomingMessages.Writer.WriteAsync(message, CancellationToken.None).ConfigureAwait(false);
            }
            else if (receiveResult.MessageType == WebSocketMessageType.Close)
            {
                break;
            }
        }
    }
}

public sealed class DeribitClient : IDisposable
{
    private readonly GeneralClient2 _client;

    public DeribitClient()
    {
        _client = new GeneralClient2("wss://www.deribit.com/ws/api/v2");

        _client.Connected += OnConnected;
        _client.Disconnected += OnDisconnected;
        _client.MessageReceived += OnMessageReceived;
    }

    public void Dispose()
    {
        _client.Connected -= OnConnected;
        _client.Disconnected -= OnDisconnected;
        _client.MessageReceived -= OnMessageReceived;

        _client.Dispose();
    }

    public Task StartAsync()
    {
        return _client.StartAsync();
    }

    public Task StopAsync()
    {
        return _client.StopAsync();
    }

    public ValueTask SubscribeToDeribitPriceIndexAsync()
    {
        string[] subscriptions = { "deribit_price_index.btc_usd" };
        var @params = new Dictionary<string, dynamic?>
        {
            { "channels", subscriptions }
        };
        var request = new JsonRpcRequest("2.0", Guid.NewGuid(), "public/subscribe", @params);
        var message = JsonSerializer.Serialize(request);

        return _client.SendAsync(message);
    }

    private void OnConnected(object? sender, EventArgs e)
    {
    }

    private void OnDisconnected(object? sender, EventArgs e)
    {
    }

    private void OnMessageReceived(object? sender, MessageReceivedEventArgs e)
    {
        Console.WriteLine($"Message received: {e.Message}{Environment.NewLine}");
    }
}

public class MessageReceivedEventArgs : EventArgs
{
    public MessageReceivedEventArgs(string message)
    {
        Message = message;
    }

    public string Message { get; }
}

public record JsonRpcRequest(
    [property: JsonPropertyName("jsonrpc")]
    string JsonRpc,
    [property: JsonPropertyName("id")] Guid Id,
    [property: JsonPropertyName("method")] string Method,
    [property: JsonPropertyName("params")] object Params);

It can be executed through the following:

using CodeReview;

using var client = new DeribitClient();
await client.StartAsync().ConfigureAwait(false);
await client.SubscribeToDeribitPriceIndexAsync().ConfigureAwait(false);

//await Task.Delay(5000).ConfigureAwait(false);

//await client.StopAsync().ConfigureAwait(false);

//await Task.Delay(5000).ConfigureAwait(false);

//await client.StartAsync().ConfigureAwait(false);

Console.ReadLine();

Answer 1

Here are my observations

DeribitClient

SubscribeToDeribitPriceIndexAsync

• The re-creation of the @params collection for each function call is unnecessary
  • You can create a static private field for that
  • I would also suggest to replace dynamic? type to object

static Dictionary<string, object> rpcRequestParameter = new() 
{
    { "channels", new [] { "deribit_price_index.btc_usd" } }
};

public ValueTask SubscribeToDeribitPriceIndexAsync()
{
    var request = new JsonRpcRequest("2.0", Guid.NewGuid(), "public/subscribe", rpcRequestParameter);
    var message = JsonSerializer.Serialize(request);
    return _client.SendAsync(message);
}

GeneralClient2

_processingXYZ

• Maybe it is just me but it seems really weird to initialize a Task field with Task.CompletedTask

constructor

• I would suggest to check that the provided string is a valid url
  • Either use Uri.IsWellFormedUriString
  • Or use Uri.TryCreate

StartAsync

_semaphore

• I'm not sure why did you protect only StartAsync method from race-condition
  • Don't you need to protect StopAsync as well?
• Please also spend sometime to find some better name for this field, like startExclusiveLock

ClientWebSocket

• Why don't you make ws a class level variable?
  • Passing this object around different methods feels a bit unnatural

await StartAsync() inside catch

• It does seems like it can be an infinite loop, since the method calls recursive itself without any exit condition

Console.WriteLine

• I guess this class is intended to be released as a part of a library based on the ConfigurateAwait(false) calls
  • You don't know whether this lib will be used inside a WPF application or in an ASP.NET Core app
• Please try to avoid using Console class inside a library class

---

UPDATE #1

StopAsync

IsRunning

• I'm not sure whether or not it is a good idea to expose IsRunning
  • It does feel like it is part of the internal state
  • You are just making shortcuts based on its value
• DeribitClient does not use that property at all

State: not (WebSocketState.Aborted or ...

• Turn your guard expression into an early exit

if (_clientWebSocket is { State: WebSocketState.Aborted or WebSocketState.Closed or WebSocketState.CloseSent })
     return;
            
await _clientWebSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None).ConfigureAwait(false);

• Based on the documentation only the Aborted state could cause exception
  • If the documentation is correct then you don't need the try-catch block

SendAsync

• I would suggest to pass a CancellationToken with time constraint to the WriteAsync in order to avoid infinite waiting

ConnectAsync

• The new Uri(_url) is error-prone
  • Please validate your input as soon as you receive it
• catch (Exception) // TaskCanceledException & WebSocketException
  • Either have a simple catch without any Exception type
  • Or use a when statement

catch (Exception ex) when (ex is OperationCanceledException or WebSocketException)

ProcessSendAsync

• Here only the WaitToReadAsync call could throw OperationCanceledException
  • So, you should wrap only that call with the try-catch

bool shouldContinue;
do
{
    try
    {
        shouldContinue = await _outgoingMessages.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false);
    }
    catch (OperationCanceledException)
    {
        shouldContinue = false;
    }

    while (_outgoingMessages.Reader.TryRead(out var message))
    {
        var data = new ArraySegment<byte>(Encoding.UTF8.GetBytes(message));
        await webSocket.SendAsync(data, WebSocketMessageType.Text, true, CancellationToken.None).ConfigureAwait(false);
    }
} while (shouldContinue);

ProcessDataAsync

• Same as with the previous + the ProcessMessageAsync could be inlined

bool shouldContinue;
do
{
    try
    {
        shouldContinue = await _incomingMessages.Reader.WaitToReadAsync(cancellationToken).ConfigureAwait(false);
    }
    catch (OperationCanceledException)
    {
        shouldContinue = false;
    }

    while (_incomingMessages.Reader.TryRead(out var message))
    {
        MessageReceived?.Invoke(this, new MessageReceivedEventArgs(message));
    }
}
while (shouldContinue);

ProcessReceiveAsync

• Debug.Assert is unnecessary since the Channels are declared as readonly
• It is strange that this method does not receive a CancellationToken as a parameter
• I would suggest to split this giant infinite loop into smaller chunks

private async Task ProcessReceiveAsync(WebSocket webSocket)
{
    while (true)
    {
        try
        {
            var (messageStream, receiveResult) = await ReadFromSocket(webSocket);
            await ForwardIncommingMessage(messageStream, receiveResult);
            if (receiveResult.MessageType == WebSocketMessageType.Close)
            {
                await CloseAsync().ConfigureAwait(false);
                return;
            }
        }
        catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
        {
            await CloseAsync().ConfigureAwait(false);
            await StartAsync().ConfigureAwait(false);
        }
    }
}

private async Task<(MemoryStream, ValueWebSocketReceiveResult)> ReadFromSocket(WebSocket webSocket)
{
    ValueWebSocketReceiveResult receiveResult;

    using var buffer = MemoryPool<byte>.Shared.Rent(4096);
    await using var ms = new MemoryStream(buffer.Memory.Length);
    do
    {
        receiveResult = await webSocket.ReceiveAsync(buffer.Memory, CancellationToken.None).ConfigureAwait(false);
        if (receiveResult.MessageType == WebSocketMessageType.Close)
            break;
            
        await ms.WriteAsync(buffer.Memory[..receiveResult.Count], CancellationToken.None).ConfigureAwait(false);
    } while (!receiveResult.EndOfMessage);

    ms.Seek(0, SeekOrigin.Begin);
    return (ms, receiveResult);
}

private async Task ForwardIncommingMessage(MemoryStream messageStream, ValueWebSocketReceiveResult receiveResult)
{
    if (receiveResult.MessageType != WebSocketMessageType.Text)
        return;
        
    using var reader = new StreamReader(messageStream, Encoding.UTF8);
    var message = await reader.ReadToEndAsync().ConfigureAwait(false);
    await _incomingMessages.Writer.WriteAsync(message, CancellationToken.None).ConfigureAwait(false);
}

---

Disclaimer: I haven't reviewed GeneralClient only GeneralClient2