Why Bedrock ProtocolReader ignores CancellationToken?

Viewed 259

It seems that my Custom ProtocolReader ignores CancellationToken, also the result.Cancelled is always false.

I took most of the code from: https://github.com/davidfowl/BedrockFramework/blob/master/samples/ServerApplication/MyCustomProtocol.cs

Am is missing something?

Here is my ProtocolReader:

public class CustomProtocolReader : IMessageReader<SocketObject>
{
    public bool TryParseMessage(in ReadOnlySequence<byte> input,
        ref SequencePosition consumed, 
        ref SequencePosition examined, 
        out SocketObject message)
    {
        var reader = new SequenceReader<byte>(input);

        var payload = input.Slice(reader.Position, input.Length);
        message = new SocketObject { Buffer = payload.ToArray() };

        consumed = payload.End;
        examined = consumed;

        return true;
    }
}

and here is the ConnectionHandler:

while (true)
{
    try
    {
        var cts = new CancellationTokenSource();
        cts.CancelAfter(2000);

        // this should throw if nothing is received after 2 seconds.
        var result = await reader.ReadAsync(protocol, cts.Token); 

        if (result.IsCompleted || result.IsCanceled)  // this never hits
        {
            Console.WriteLine("Broke or cancelled");
            break;
        }
    }
    catch (OperationCanceledException) // this never gets fired
    {
        Console.WriteLine("Cancelled");
        break;
    }
    finally
    {
        reader.Advance();
    }
}

Update 1:

default connection.Transport.Input.ReadAsync works as expected and respects CancellationToken.

1 Answers

Update: Ohh well.. I've just overlooked this part. So yes, you should expect the cancellation after 2 seconds and even ValueTask should print IsCancled = true. :/ I'm afraid, but I don't know about that. I tried to understand the IValueTaskSource.OnCompleted implementation of Pipe/PipeAwaitable but it is very complicated. I don't have the time right now.

I will let the answer here because of historical reasons.

It seems that my Custom ProtocolReader ignores CancellationToken

It does not.

First of all you have to know that CancellationToken is not actively doing something. When using CancellationToken.CancelAfter it just notifies its origin CancellationTokenSource and their children (the cancellation token sources that got created by CancellationTokenSource.CreateLinkedTokenSource) to set their IsCancellationRequested to true and fires all handlers (registered by CancellationToken.Register) in all involved CancellationTokenSource's. It is in your response to listen to a cancellation request by continously looking at IsCancellationRequested or by register a handler with CancellationToken.Register (see example implementation below).

The deepest subroutine when calling reader.ReadAsync is checking one time if the token you passed to reader.ReadAsync has been cancellation requested. Let me elaborate this in detail:

I assume that your reader of type ProtocolReader has been created from connectionContext.CreateReader (like in MyCustomProtocol) which internally passes Input of the Transport of ConnectionContext to constructor of ProtocolReader. The implementation of ConnectionContext should be SocketConnection. The property Transport is of type IDuplexPipe and is created in SocketConnection.StartAsync which creates a DuplexPipePair and saves its property Transport to SocketConnection.Transport (see property Transport of SocketConnection). The type of Transport has Input and Output and are both of type Pipe. This class uses PipeAwaitable internally in implementation of System.IO.Pipelines.Pipe.ReadAsync (see method ReadAsync of type Pipe) by calling BeginOperation of PipeAwaitable. Finally, this method calls first time CancellationToken.ThrowIfCancellationRequested() of the token you passed initially by var result = await reader.ReadAsync(protocol, cts.Token);.

So your cancellation token is only get ignored when your cancellation has not been cancellation requested in the moment the very one code passage CancellationToken.ThrowIfCancellationRequested() could pass.

also the result.Cancelled is always false.

Your result is not a Task. When you create a Task manually (Task.Factory.StartNew) you can pass a CancellationToken. When this token gets cancellation requested before or after task creation then Task.IsCanceled will become true unless already finished (see Task.AssignCancellationToken).

But you have maybe recognized that you get not Task but ValueTask. This a lightweight task implementation to reduce Task creation overhead. This does not accept any CancellationToken. So your initial passed CancellationToken won't be used to notify ValueTask being canceled. In other words. As soon as the ValueTask gets created, the cancellation from CancellationTokenSource won't turn ValueTask.IsCanceled to true anymore. Only its internal held variable and if of type Task or IValueTaskSource and its implementation of Task.IsCanceled or IValueTaskSource.GetStatus can change this, but this is not the case (see next paragraph).

ValueTask accepts either any "precalculated" value, a Task or a "task"/promise that implements IValueTaskSource. In case of a value that implements IValueTaskSource, the only way to indicate ValueTask.IsCanceled = true is when IValueTaskSource.GetStatus returns ValueTaskSourceStatus.Canceled. So ReadAsync from Pipe returns either the synchronous result wrapped by ValueTask or still pending/promised result also wrapped by ValueTask and this is important, this pending/promised result will be an internal variable of type Pipe that has been initialized in constructor of the public class Pipe. But your initially passed cancellation token in context of Pipe.ReadAsync is not interacting with the internal variable of type Pipe (implements IValueTaskSource) to get its implemented method IValueTaskSource.GetStatus to return ValueTaskSourceStatus.Canceled.

So the reason that your task won't never become true is that in no branch a Task with your initial cancellation token is ever created nor a IValueTaskSource implementation is implemented that is interacting with your initial passed cancellation token.

So what are the consequences?

When calling ReadAsync after your token has been cancellation requested, you are fine - you get your desired OperationCanceledException. After that you're are not covered by expected cancellation behaviour. So that's where you have to start to implement a solution by yourself:

        public Task ReadAsync() =>
            // This is just an example.
            Task.Delay(int.MaxValue);

        public async Task DoTaskNotForever(CancellationToken token = default)
        {
            using var cts = CancellationTokenSource.CreateLinkedTokenSource(token);
            cts.CancelAfter(2);
            var tcs = new TaskCompletionSource<object>();
            using var registration = cts.Token.Register(() => tcs.SetCanceled());

            while (true) {
                try {
                    await Task.WhenAny(ReadAsync(), tcs.Task).ConfigureAwait(false);
                } catch (OperationCanceledException error) { 
                    // Ensure to flush/close/reconnect socket !!! (see below why)
                }
            }
        }

Before you use the above... I can only advice you to use the library https://github.com/StephenCleary/AsyncEx. Especially Nito.AsyncEx.Tasks and the Task extension method WaitAsync.

It is something like: Tasks.WhenAny(theActualTask, throwOnCancellationRequestedTask) where throwOnCancellationRequestedTask looks continously for cancellation request and throws if so:

using using Nito.AsyncEx;

...

while (true)
{
    try
    {
        var cts = new CancellationTokenSource();
        cts.CancelAfter(2000);

        // this should throw if nothing is received after 2 seconds.
        var result = await reader.ReadAsync(protocol, cts.Token).WaitAsync(cts.Token); 

        if (result.IsCompleted || result.IsCanceled)  // this should hit
        {
            Console.WriteLine("Broke or cancelled");
            break;
        }
    }
    catch (OperationCanceledException) // should get fired after 2 seconds unless ReadAsync has not yet finished
    {
        Console.WriteLine("Cancelled");
        break;
    }
    finally
    {
        reader.Advance();
    }
}

"Ensure to flush/close/reconnect socket !!!"

After timeout cancellation by WaitAsync (that from Nito.AsyncEx.Tasks) or your own implementation, you should recreate the socket, or implement a safe prodecure that ensures that the finished or still running ValueTask from ReadAsync gets "finished". In SerialPort you can flush the SerialPort.BaseStream so the Task you got from SerialPort.ReadAsync gets "finished".

Related