How do I properly close a persistent System.Net.WebSockets.ClientWebSocket connection?

Viewed 2893

I am using System.Net.WebSockets.ClientWebSocket to create a C# WebSocket client to connect to a server, to which I do not own the source code. Everything works well so far, but I wish to "correctly" disconnect my client. Here's an abbreviated version of my source code:

public class Client : IDisposable
{
    private ClientWebSocket socket;
    private string endpoint;
    private Task receiveTask;

    public Client(string endpoint)
    {
        this.endpoint = endpoint;
        this.socket = new ClientWebSocket();
    }

    public async Task Initialize()
    {
        byte[] contentBuffer = Encoding.UTF8.GetBytes("notify message");

        // Connect to the server, and send a message to notify
        // it of the client's availability to receive data.
        await OpenConnection();
        await socket.SendAsync(new ArraySegment<byte>(contentBuffer), WebSocketMessageType.Text, true, CancellationToken.None);
    }

    private async Task OpenConnection()
    {
        if (socket.State != WebSocketState.Open)
        {
            await socket.ConnectAsync(new Uri(endpoint), CancellationToken.None);
            receiveTask = Task.Run(async () => await Receive());
        }
    }

    private async Task Receive()
    {
        while (socket.State == WebSocketState.Open)
        {
            byte[] buffer = new byte[1024];
            var result = await m_sessionSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);

            if (result.MessageType == WebSocketMessageType.Close)
            {
                break;
            }
            else
            {
                using (var stream = new MemoryStream())
                {
                    stream.Write(buffer, 0, result.Count);
                    while (!result.EndOfMessage)
                    {
                        result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
                        stream.Write(buffer, 0, result.Count);
                    }

                    stream.Seek(0, SeekOrigin.Begin);
                    using (var reader = new StreamReader(stream, Encoding.UTF8))
                    {
                        string message = reader.ReadToEnd();
                        // Do stuff with received message
                    }
                }
            }
        }
    }

    public void Dispose()
    {
        // NOTE: This is a gross oversimplification. Assume in the
        // actual project that a proper implementation of the Dispose
        // pattern has been created.
        if (socket.State == WebSocketState.Open)
        {
            // How do I notify the server of disconnection?
            // The below has some sort of race condition, whereby
            // it hangs the client.
            Task.Run(async () => await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None)).Wait();
        }

        if (receiveTask != null)
        {
            receiveTask.Dispose();
        }

        socket.Dispose()
    }
}

Now for the challenge I'm having. I can't dispose of the Task running the receive of data until it's finished. That task is blocked awaiting receive from the server. Attempting to call CloseAsync from within Dispose seems to encounter a deadlock of some kind, so that might not be the right approach.

Obviously, I could use a CancellationTokenSource to pass a CancellationToken to the Receive method, passing it, in turn, to the ReceiveAsync method of the WebSocket. If I do that and tell the token source to flag the token, it will release the block, and I can check the token for the IsCancellationRequested flag inside the receive method, exiting when that's seen. However, if I do that, it throws an exception. Do I really need to catch an exception (yuck!) in my Dispose method just for wanting to cancel the task? Also, if I cancel the task, does that properly notify the server of the disconnection? Or am I just doing this wrong in the first place by having the long-running (non-awaited) Task running the ReceiveAsync method?

1 Answers

You should not use IDisposable in this case. Gracefully closing a websocket connection involves async I/O. It needs a clean "close handshake" which is when the sender sends a CLOSE frame and the receives sends back a CLOSE frame. Only then will the connection transition to a clean socket close (with a FIN and not RST).

If you need "IDisposable" semantics, you could look into the new IAsyncDisposable pattern introduced in latest .NET Core. But in general, using synchronous Dispose() pattern is not a good idea.

Related