ASP.NET Core websocket manager - websockets are always disposed

Question

I'm handling some websockets in ASP.NET Core 2.0. I want to accept incoming requests and hand them off to a singleton service so that messages can easily be sent to them from an API controller. My server is not reading anything at all, it's only accepting connections and sending messages to them.

Whenever I try to access the socket to send a message, however, it's always disposed (ObjectDisposedException). I'm not sure how to properly pass the connection from the middleware to the service in a way that doesn't dispose of it. How can I do this?

The websocket manager (singleton):

public class WebSocketManager
{
    private readonly ConcurrentDictionary<Guid, WebSocket> _sockets;

    protected WebSocketManager()
    {
        _sockets = new ConcurrentDictionary<Guid, WebSocket>();
    }

    public Guid AddWebSocket(WebSocket socket)
    {
        var guid = Guid.NewGuid();
        _sockets.TryAdd(guid, socket);

        return guid;
    }

    public async Task SendAsync(Guid guid, string message)
    {
        if (!_sockets.TryGetValue(guid, out var socket))
            return;

        await SendMessageAsync(socket, message);
    }

    public async Task SendAllAsync(string message)
    {
        foreach (var socket in _sockets)
        {
            if (socket.Value.State == WebSocketState.Open)
                await SendMessageAsync(socket.Value, message);
            else
                await RemoveWebSocketAsync(socket.Key);
        }
    }

    private async Task SendMessageAsync(WebSocket socket, string message)
    {
        await socket.SendAsync(
            new ArraySegment<byte>(Encoding.UTF8.GetBytes(message)),
            WebSocketMessageType.Text,
            true,
            CancellationToken.None);
    }

    private async Task RemoveWebSocketAsync(Guid guid)
    {
        if (!_sockets.TryRemove(guid, out var socket))
            return;

        if (socket?.State == WebSocketState.Open)
            await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closing", CancellationToken.None);

        socket.Dispose();
    }
}

Middleware that accepts the connections:

public class WebSocketMiddleware
{
    private readonly RequestDelegate _next;
    private readonly WebSocketManager _wsMgr;

    public WebSocketMiddleware(
        RequestDelegate next,
        WebSocketManager wsMgr)
    {
        _next = next;
        _wsMgr = wsMgr;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        if (context.Request.Path != "/ws")
        {
            await _next(context);
            return;
        }

        if (!context.WebSockets.IsWebSocketRequest)
        {
            context.Response.StatusCode = 400;
            await context.Response.WriteAsync("This endpoint accepts only websocket connections.");
            return;
        }

        var guid = _wsMgr.AddWebSocket(await context.WebSockets.AcceptWebSocketAsync());

        // this call succeeds
        await _wsMgr.SendAsync(guid, "test message");
    }
}

Example controller call:

public class WinLossController : Controller
{
    private readonly WebSocketManager _wsMgr;

    public WinLossController(WebSocketManager wsMgr)
    {
        _wsMgr = wsMgr;
    }

    [HttpPost]
    public async Task<IActionResult> Update()
    {
        await _wsMgr.SendAllAsync("test controller update");

        return NoContent();
    }
}

Answer 1

As soon as you leave InvokeAsync of your middleware - websocket connection is closed, that's an expected flow.

Correct way to prevent that is to listen for incoming messages on socket. Even if you don't expect any data from client - you still need to do this, because messages client might send to you include some control messages, such as ping-pong message (heartbeat) and "close" message, indicating that client is closing communication.

So after doing:

var guid = _wsMgr.AddWebSocket(await context.WebSockets.AcceptWebSocketAsync());

Implement a loop in your manager which at least handles "close" message and await that:

await _wsMgs.ReceiveAsync(guid);

If you need to send something to client right after connecting - do that before awaiting receive loop, because this loop should only complete when you are done with this websocket.