Skip to content

CSHARP-6220: Prevent cancellation-token retention in WaitAsync - #2124

Open
mohammadshahsavari wants to merge 3 commits into
mongodb:mainfrom
mohammadshahsavari:CSHARP-6220
Open

mohammadshahsavari wants to merge 3 commits into
mongodb:mainfrom
mohammadshahsavari:CSHARP-6220

Conversation

@mohammadshahsavari

Copy link
Copy Markdown

Summary

  • Prevent the pre-.NET 6 TaskExtensions.WaitAsync implementation from retaining registrations on long-lived caller cancellation tokens.
  • Cancel and dispose a linked timeout cancellation source after Task.WhenAny completes.
  • Add regression coverage for both generic and non-generic tasks.

Testing

  • dotnet build CSharpDriver.sln
  • dotnet test tests/MongoDB.Driver.Tests/MongoDB.Driver.Tests.csproj -f net472 --filter "FullyQualifiedName~TaskExtensionsTests"

JIRA: https://jira.mongodb.org/browse/CSHARP-6220

@codeowners-service-app

Copy link
Copy Markdown

Assigned sanych-sun for team dbx-csharp-dotnet because ajcvickers is out of office.

{
var timeoutTask = Task.Delay(timeout, cancellationToken);
await Task.WhenAny(task, timeoutTask).ConfigureAwait(false);
using (var timeoutCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we check cancellationToken.CanBeCanceled before creating the CancellationTokenSource to optimize the case when the CancellationToken.None is being used?

@mohammadshahsavari mohammadshahsavari Sep 16, 2026 •

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, good point. I updated both overloads to use a regular CTS when the token can’t be canceled.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to have a special case from non-cancelable token sources?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I’ve removed the try-finally and cancel the timeout task directly after Task.WhenAny.

@sanych-sun sanych-sun added the improvement Optimizations or refactoring (no new features or fixes). label Sep 16, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

improvement Optimizations or refactoring (no new features or fixes).

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants