How to gracefully close a two-way WebSocket in .Net

Question

I have a WebSocket server that accepts a stream of binary data from a client and responds with another stream of text data for every 4MB read. The server uses IIS 8 and asp.net web api.

Server

public class WebSocketController : ApiController
{
    public HttpResponseMessage Get()
    {
        if (!HttpContext.Current.IsWebSocketRequest)
        {
            return new HttpResponseMessage(HttpStatusCode.BadRequest);
        }

        HttpContext.Current.AcceptWebSocketRequest(async (context) =>
        {
            try
            {
                WebSocket socket = context.WebSocket;

                byte[] requestBuffer = new byte[4194304];
                int offset = 0;

                while (socket.State == WebSocketState.Open)
                {
                    var requestSegment = new ArraySegment<byte>(requestBuffer, offset, requestBuffer.Length - offset);
                    WebSocketReceiveResult result = await socket.ReceiveAsync(requestSegment, CancellationToken.None);

                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        // Send one last response before closing
                        var response = new ArraySegment<byte>(Encoding.UTF8.GetBytes("Server got " + offset + " bytes\n"));
                        await socket.SendAsync(response, WebSocketMessageType.Text, true, CancellationToken.None);

                        // Close
                        await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);
                        break;
                    }

                    offset += result.Count;
                    if (offset == requestBuffer.Length)
                    {
                        // Regular response
                        var response = new ArraySegment<byte>(Encoding.UTF8.GetBytes("Server got 4194304 bytes\n"));
                        await socket.SendAsync(response, WebSocketMessageType.Text, true, CancellationToken.None);
                        offset = 0;
                    }
                }
            }
            catch (Exception ex)
            {
                // Log and continue
            }
        });

        return new HttpResponseMessage(HttpStatusCode.SwitchingProtocols);
    }
}

The c# client uses the ClientWebSocket class to connect to the server and send requests. It creates a task for receiving responses from the server that runs in parallel with the request sending. When it is done sending the requests it calls CloseAsync on the socket and then waits for the Receive task to complete.

Client

using System;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace WebSocketClient
{
    class Program
    {
        static void Main(string[] args)
        {
            try
            {
                CallWebSocketServer().Wait();
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
            }
        }

        static async Task CallWebSocketServer()
        {
            using (ClientWebSocket socket = new ClientWebSocket())
            {
                await socket.ConnectAsync(new Uri("ws://localhost/RestWebController"), CancellationToken.None);
                byte[] buffer = new byte[128 * 1024];

                Task receiveTask = Receive(socket);
                for (int i = 0; i < 1024; ++i)
                {
                    await socket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Binary, true, CancellationToken.None);
                }
               
                await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);
                receiveTask.Wait();

                Console.WriteLine("All done");
            }
        }

        static async Task Receive(ClientWebSocket socket)
        {
            try
            {
                byte[] recvBuffer = new byte[64 * 1024];
                while (socket.State == WebSocketState.Open)
                {
                    var result = await socket.ReceiveAsync(new ArraySegment<byte>(recvBuffer), CancellationToken.None);
                    Console.WriteLine("Client got {0} bytes", result.Count);
                    Console.WriteLine(Encoding.UTF8.GetString(recvBuffer, 0, result.Count));
                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        Console.WriteLine("Close loop complete");
                        break;
                    }
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine("Exception in receive - {0}", ex.Message);
            }
        }
    }
}

The problem is that the client blocks at the CloseAsync call.

What would be the correct way of gracefully closing the WebSocket in this scenario?

Answer 1

I did some research using various combinations of the following code. The scenario is that the client sends data until the server closes the WebSocket.

Server

public class HomeController : Controller
{
    public async Task Index()
    {
        using (WebSocket webSocket = await HttpContext.WebSockets.AcceptWebSocketAsync())
        {
            await Task.Delay(3000);
            await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "ServerClose", new CancellationTokenSource(20_000).Token);
        }
    }
}

Client

static async Task SendDataToWebSocket(string webSocketUrl)
{
    ClientWebSocket webSocket = new ClientWebSocket();
    try
    {
        await webSocket.ConnectAsync(new Uri(webSocketUrl), new CancellationTokenSource(5000).Token);
        ProcessCloseEvent(webSocket);

        while (true)
        {
            ArraySegment<byte> data = new ArraySegment<byte>(Encoding.UTF8.GetBytes("something"));
            await webSocket.SendAsync(data, WebSocketMessageType.Text, true, CancellationToken.None);
            await Task.Delay(1000);
        }
    }
    catch (Exception e)
    {
        Console.WriteLine("Error: " + e.Message);
    }
}

private static async void ProcessCloseEvent(ClientWebSocket webSocket)
{
    ArraySegment<byte> data = new ArraySegment<byte>(new byte[1000]);
    var result = await webSocket.ReceiveAsync(data, CancellationToken.None);
    Console.WriteLine("Count: {0}, CloseStatus: {1}, Data: {2}", result.Count, result.CloseStatus, Encoding.UTF8.GetString(data));
    await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "ClientClose", CancellationToken.None);
}

I observed the following:

• CloseAsync works in pairs, meaning that the first side calls CloseAsync and waits until the other side calls it too.
• Bacause CloseAsync with CancelationToken.None will wait indefinitely, people prefer the use of CloseOutputAsync that seems not to be waiting for the other side to respond.
• Even though the WebSocket close event can be observed in the result returned by ReceiveAsync call, the WebSocket.State and WebSocket.CloseStatus only gets updated after ReceiveAsync.
• CloseAsync and CloseOutputAsync can be paired in any combination to respond to the WebSocket close event. Since it shouldn't be necessary to wait for a response when closing the WebSocket or know if the other side handled the close event, CloseOutputAsync should be better to use.
• Once the client responds with CloseAsync to the server's request, the WebSocket.State is changed to Closed and subsequent SendAsync and ReceiveAsync will throw an exception.
• If the client does not react to WebSocket close event and server disposes the WebSocket, the SendAsync at client side with throw an IOException with this message: "Unable to write data to the transport connection: An established connection was aborted by the software in your host machine.."
• If the client does not react to WebSocket close event and server disposes the WebSocket, the ReceiveAsync at client side will throw an WebSocketException with the message: "The remote party closed the WebSocket connection without completing the close handshake."
• After the server disposes the WebSocket, the client is able to send/receive a few messages without any exception, so there should be no concurrency problems between checking for close and exchanging data.
• Since CloseAsync waits for the other side, it can be used to determine if all sent data has been successfully received. This does not apply if both sides send and close at the same time.

NOTE: I have used .NET 6 for testing. The above observations may not apply in combinations with clients/servers written in other languages.

Answer 2

Figured this out.

Server

Basically, I had to call the ClientWebSocket.CloseOutputAsync (instead of the CloseAsync) method to tell the framework no more output is going to be sent from the client.

await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, string.Empty, CancellationToken.None);

Client

Then in the Receive function, I had to allow for socket state WebSocketState.CloseSent to receive the Close response from the server

static async Task Receive(ClientWebSocket socket)
{
    try
    {
        byte[] recvBuffer = new byte[64 * 1024];
        while (socket.State == WebSocketState.Open || socket.State == WebSocketState.CloseSent)
        {
            var result = await socket.ReceiveAsync(new ArraySegment<byte>(recvBuffer), CancellationToken.None);
            Console.WriteLine("Client got {0} bytes", result.Count);
            Console.WriteLine(Encoding.UTF8.GetString(recvBuffer, 0, result.Count));
            if (result.MessageType == WebSocketMessageType.Close)
            {
                Console.WriteLine("Close loop complete");
                break;
            }
        }
    }
    catch (Exception ex)
    {
        Console.WriteLine("Exception in receive - {0}", ex.Message);
    }
}

Answer 3

i suggest you to look at these links:

asynchronous server: https://msdn.microsoft.com/en-us/library/fx6588te%28v=vs.110%29.aspx

asynchronous client: https://msdn.microsoft.com/en-us/library/bew39x2a(v=vs.110).aspx

recently i implementled something similar with these links as an example. the methods "BeginReceive" (for server) and "BeginConnect" (for client) start each a new thread. so there won't be anything that blocks