'functionCompletionTask' is being set when already complete in Extensions.Http.AspNetCore

[liliankasem]

Description

There appears to be a bug in Extensions.Http.AspNetCore during function completion. When a Http invocation gets cancelled, an InvalidOperationException is thrown instead of expected TaskCanceledException.

This occurs because during cancellation we cancel the CancellationTokenSource in the invocation handler:

azure-functions-dotnet-worker/src/DotNetWorker.Grpc/Handlers/InvocationHandler.cs

Line 147 in ca7afd2

cancellationTokenSource?.Cancel();

This completes the _functionCompletionTask:

azure-functions-dotnet-worker/extensions/Worker.Extensions.Http.AspNetCore/src/Coordinator/ContextReference.cs

Lines 35 to 45 in ca7afd2

	internal void SetCancellationToken(CancellationToken token)
	{
	_token = token;
	_tokenRegistration = _token.Register(() =>
	{
	_functionStartTask.TrySetCanceled();
	_functionCompletionTask.TrySetCanceled();
	_functionContextValueSource.TrySetCanceled();
	_httpContextValueSource.TrySetCanceled();
	});
	}

When CompleteFunction is called the _functionCompletionTask task is already complete, but we try to set the result here anyway:

azure-functions-dotnet-worker/extensions/Worker.Extensions.Http.AspNetCore/src/Coordinator/ContextReference.cs

Lines 53 to 60 in ca7afd2

	internal void CompleteFunction()
	{
	if (_httpContextValueSource.Task.IsCompleted)
	{
	if (_httpContextValueSource.Task.IsCanceled || _token.IsCancellationRequested)
	{
	_functionCompletionTask.SetCanceled();
	}

Which leads to the InvalidOperationException exception.

Exception: System.InvalidOperationException: An attempt was made to transition a task to a final state when it had already completed.

Repro

Repro can be found in this GH issue:

• InvalidOperationException when canceling request #2376

Fix

First we can implement a check at the start of CompleteFunction to see if the task is complete before trying to set a result i.e.

if (_functionCompletionTask.Task.IsCompleted)
{
    return;
}

This still leads to a different exception FunctionInvocationException with a TaskCanceledException inner exception.

Executed 'Functions.BlobStream' (Failed, Id=c4cb3e5e-a039-4a44-a4d2-172ac6be0668, Duration=25692ms)
      Microsoft.Azure.WebJobs.Host.FunctionInvocationException: Exception while executing function: Functions.BlobStream
       ---> System.InvalidOperationException: Failed to proxy request with ForwarderError: ResponseBodyCanceled
       ---> System.Threading.Tasks.TaskCanceledException: The operation was canceled.
         at System.Net.Http.CancellationHelper.ThrowOperationCanceledException(Exception innerException, CancellationToken cancellationToken)
         at System.Net.Http.HttpConnection.ContentLengthReadStream.ReadAsync(Memory`1 buffer, CancellationToken cancellationToken)
         at Yarp.ReverseProxy.Forwarder.StreamCopier.CopyAsync(Stream input, Stream output, Int64 promisedContentLength, StreamCopierTelemetry telemetry, ActivityCancellationTokenSource activityToken, Boolean autoFlush, CancellationToken cancellation)

If there are other enhancements to make here, they should be considered.