CSHARP-6220: Prevent cancellation-token retention in WaitAsync - #2124
mohammadshahsavari wants to merge 3 commits into
Conversation
|
Assigned |
| { | ||
| var timeoutTask = Task.Delay(timeout, cancellationToken); | ||
| await Task.WhenAny(task, timeoutTask).ConfigureAwait(false); | ||
| using (var timeoutCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken)) |
There was a problem hiding this comment.
Should we check cancellationToken.CanBeCanceled before creating the CancellationTokenSource to optimize the case when the CancellationToken.None is being used?
There was a problem hiding this comment.
Thanks, good point. I updated both overloads to use a regular CTS when the token can’t be canceled.
There was a problem hiding this comment.
Do we really need the CancellationTokenSource at all in case when CancellationToken is not cancellable? As far as i understood for non-cancellable tokens registration for callback is no-op, so it will not hold a reference to anything. I suppose we can save on creating a CancellationTokenSource in such cases.
There was a problem hiding this comment.
You’re right. The registration issue only applies to cancellable tokens so I’ve removed the CTS allocation for non-cancellable tokens.
| { | ||
| var timeoutTask = Task.Delay(timeout, cancellationToken); | ||
| await Task.WhenAny(task, timeoutTask).ConfigureAwait(false); | ||
| using (var timeoutCancellationTokenSource = cancellationToken.CanBeCanceled ? |
There was a problem hiding this comment.
Do we need to have a special case from non-cancelable token sources?
There was a problem hiding this comment.
The previous special case created a regular CTS, which wasn’t necessary. I’ve updated it so non-cancelable tokens don’t create a CTS at all. the conditional now only avoids that unnecessary allocation.
| new CancellationTokenSource()) | ||
| { | ||
| var timeoutTask = Task.Delay(timeout, timeoutCancellationTokenSource.Token); | ||
| try |
There was a problem hiding this comment.
I don't think we should have the try-finally here. Task.WhenAny does not throw when the tasks themselves throw (this is why we have the pre-existing checks afterwards).
There was a problem hiding this comment.
I’ve removed the try-finally and cancel the timeout task directly after Task.WhenAny.
Summary
TaskExtensions.WaitAsyncimplementation from retaining registrations on long-lived caller cancellation tokens.Task.WhenAnycompletes.Testing
dotnet build CSharpDriver.slndotnet test tests/MongoDB.Driver.Tests/MongoDB.Driver.Tests.csproj -f net472 --filter "FullyQualifiedName~TaskExtensionsTests"JIRA: https://jira.mongodb.org/browse/CSHARP-6220