Search code examples
c#.netcancellationtokensource

Cancellation doesn't go as expected


The while loop currently blocks the thread as it retries forever until the connection is established. I expect it to retry forever but it shouldn't block the thread (just like if we call StartAsync without awaiting it) and make it possible for us to call StopAsync during StartAsync's execution and cancel that process of retrying forever. It would require passing a CancellationToken to ConnectAsync too.

await client.StartAsync(); // do not block

await Task.Delay(5000);

await client.StopAsync();

I was thinking of moving the CancellationToken before the while loop and pass it to the while loop as well as loop until while (!await ConnectAsync().ConfigureAwait(false) && !_tokenSource.IsCancellationRequested) and then wrap that logic into a Task.Run to prevent blocking. What do you think?

Full code (@GeneralClient2 class)

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()
{
    if (!IsRunning)
    {
        return;
    }

    _logger.LogDebug("Stopping");

    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
    {
    }

    await _processingReceive.ConfigureAwait(false);

    _logger.LogDebug("Stopped");
}

private async ValueTask<bool> ConnectAsync()
{
    _logger.LogDebug("Connecting");

    var ws = new ClientWebSocket();

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

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

    _clientWebSocket = ws;

    _logger.LogDebug("Connected");

    return true;
}

Solution

  • The while loop currently blocks the thread

    Are you sure? because ClientWebSocket.ConnectAsync explicitly states

    This operation will not block. The returned Task object will complete after the connect request on the ClientWebSocket instance has completed.

    So ConnectAsync should not block, and therefore neither the while loop. But even if it is non-blocking it might still consume significant CPU usage. Or something might throw before you even call ClientWebSocket.ConnectAsync, and since you just eat all the exceptions you will never know.

    make it possible for us to call StopAsync during StartAsync's execution and cancel that process of retrying forever

    You should be careful when stopping some service in a threaded or async context. Since some other task might find that the needed resource have been disposed.

    while (!await ConnectAsync().ConfigureAwait(false) && !_tokenSource.IsCancellationRequested)

    The problem with this is that the loop will exit and continue running the method if no connection has been established. Because of this it is recommended to use ThrowIfCancellationRequested, even if it is somewhat painful to use exceptions for flow control.

    My recommendations would be to make StartAsync take a cancellationToken that aborts the connection process. This method should return an object representing the connection, i.e. Task<MyConnection>. This connection object should be disposable. Stopping the connection could be done like

    // Create connection
    var cts = new CancellationTokenSource();
    var myStartAsyncConnectionTask = StartAsync(cts.Token);
    
    // Close connection
    cts.Cancel();
    try{
        var myConnection = await myStartAsyncConnectionTask;
        myConnection.Dispose();
    }
    catch(OperationCancelledException)
    {
    ...
    }
    

    That should work regardless of what state the connection is in. If a connection has been established the cancellation will do nothing, and the object will be disposed. If it has failed to connect, awaiting the task should throw. Note that StartAsync need to be written so that it cleans up any created resource in case the method throws any exception at any stage.