Skip to content

Commit ae1464c

Browse files
YunchuWangYunchuWang
authored
Release cancellation registrations after external-event waits complete (#838)
1 parent 3612c73 commit ae1464c

3 files changed

Lines changed: 167 additions & 2 deletions

File tree

‎src/Worker/Core/RELEASENOTES.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
- Add new `IDurableTaskWorkerBuilder AddDurableTaskWorker(IServiceCollection, string?)` API
2+
- Release cancellation registrations after external-event waits complete so live token sources do not retain completed waits and their payloads.

‎src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -336,8 +336,19 @@ public override Task<T> WaitForExternalEvent<T>(string eventName, CancellationTo
336336
this.CompleteExternalEvent(eventName, bufferedEvent);
337337
}
338338

339-
// TODO: this needs to be tracked and disposed appropriately.
340-
cancellationToken.Register(() => eventSource.TrySetCanceled(cancellationToken));
339+
if (cancellationToken.CanBeCanceled && !eventSource.Task.IsCompleted)
340+
{
341+
CancellationTokenRegistration registration = cancellationToken.Register(() => eventSource.TrySetCanceled(cancellationToken));
342+
343+
// Cleanup must not depend on the orchestrator's task scheduler.
344+
_ = eventSource.Task.ContinueWith(
345+
static (_, state) => ((CancellationTokenRegistration)state!).Dispose(),
346+
registration,
347+
CancellationToken.None,
348+
TaskContinuationOptions.ExecuteSynchronously,
349+
TaskScheduler.Default);
350+
}
351+
341352
return eventSource.Task;
342353
}
343354

‎test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -391,6 +391,159 @@ await wrapper.CallSubOrchestratorAsync<string>(
391391
innerContext.LastSubOrchestrationVersion.Should().Be(string.Empty);
392392
}
393393

394+
[Fact]
395+
public async Task WaitForExternalEvent_PreCanceledToken_PreservesCancellationTokenAsync()
396+
{
397+
// Arrange
398+
TrackingOrchestrationContext innerContext = new();
399+
OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
400+
TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
401+
using CancellationTokenSource source = new();
402+
source.Cancel();
403+
404+
// Act
405+
Task<string> wait = wrapper.WaitForExternalEvent<string>("Event", source.Token);
406+
407+
// Assert
408+
wait.IsCanceled.Should().BeTrue();
409+
TaskCanceledException exception = await Assert.ThrowsAsync<TaskCanceledException>(() => wait);
410+
exception.CancellationToken.Should().Be(source.Token);
411+
wrapper.CompleteExternalEvent("Event", "\"payload\"");
412+
Task<string> nextWait = wrapper.WaitForExternalEvent<string>("Event");
413+
nextWait.IsCompletedSuccessfully.Should().BeTrue();
414+
(await nextWait).Should().Be("payload");
415+
}
416+
417+
[Theory]
418+
[InlineData(false)]
419+
[InlineData(true)]
420+
public async Task WaitForExternalEvent_BufferedEvent_CompletesSynchronouslyBeforeCancellationAsync(bool preCanceled)
421+
{
422+
// Arrange
423+
TrackingOrchestrationContext innerContext = new();
424+
OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
425+
TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
426+
using CancellationTokenSource source = new();
427+
wrapper.CompleteExternalEvent("Event", "\"payload\"");
428+
if (preCanceled)
429+
{
430+
source.Cancel();
431+
}
432+
433+
// Act
434+
Task<string> wait = wrapper.WaitForExternalEvent<string>("Event", source.Token);
435+
source.Cancel();
436+
437+
// Assert
438+
wait.IsCompletedSuccessfully.Should().BeTrue();
439+
(await wait).Should().Be("payload");
440+
}
441+
442+
[Fact]
443+
public async Task CompleteExternalEvent_CancelableWait_PreservesSynchronousContinuationsAsync()
444+
{
445+
// Arrange
446+
TrackingOrchestrationContext innerContext = new();
447+
OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
448+
TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
449+
using CancellationTokenSource source = new();
450+
Task<string> wait = wrapper.WaitForExternalEvent<string>("Event", source.Token);
451+
Task<string> canceledWait = wrapper.WaitForExternalEvent<string>("OtherEvent", source.Token);
452+
bool continuationRan = false;
453+
Task continuation = wait.ContinueWith(
454+
_ =>
455+
{
456+
continuationRan = true;
457+
source.Cancel();
458+
},
459+
CancellationToken.None,
460+
TaskContinuationOptions.ExecuteSynchronously,
461+
TaskScheduler.Default);
462+
463+
// Act
464+
wrapper.CompleteExternalEvent("Event", "\"payload\"");
465+
466+
// Assert
467+
continuationRan.Should().BeTrue();
468+
canceledWait.IsCanceled.Should().BeTrue();
469+
wait.IsCompletedSuccessfully.Should().BeTrue();
470+
(await wait).Should().Be("payload");
471+
await continuation;
472+
}
473+
474+
[Fact]
475+
public async Task CompleteExternalEvent_CancelableWaiters_PreservesLifoOrderAsync()
476+
{
477+
// Arrange
478+
TrackingOrchestrationContext innerContext = new();
479+
OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
480+
TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
481+
using CancellationTokenSource liveSource = new();
482+
using CancellationTokenSource canceledSource = new();
483+
Task<string> oldest = wrapper.WaitForExternalEvent<string>("Event", liveSource.Token);
484+
Task<string> canceled = wrapper.WaitForExternalEvent<string>("Event", canceledSource.Token);
485+
Task<string> newest = wrapper.WaitForExternalEvent<string>("Event", liveSource.Token);
486+
canceledSource.Cancel();
487+
488+
// Act
489+
wrapper.CompleteExternalEvent("Event", "\"first\"");
490+
wrapper.CompleteExternalEvent("Event", "\"second\"");
491+
liveSource.Cancel();
492+
493+
// Assert
494+
canceled.IsCanceled.Should().BeTrue();
495+
(await newest).Should().Be("first");
496+
(await oldest).Should().Be("second");
497+
}
498+
499+
[Fact]
500+
public async Task CompleteExternalEvent_CancellationRace_DeliversOrBuffersExactlyOnceAsync()
501+
{
502+
for (int i = 0; i < 100; i++)
503+
{
504+
// Arrange
505+
TrackingOrchestrationContext innerContext = new();
506+
OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
507+
TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
508+
using CancellationTokenSource source = new();
509+
using Barrier start = new(2);
510+
Task<string> wait = wrapper.WaitForExternalEvent<string>("Event", source.Token);
511+
512+
// Act
513+
Task cancellation = Task.Run(() =>
514+
{
515+
start.SignalAndWait();
516+
source.Cancel();
517+
});
518+
Task completion = Task.Run(() =>
519+
{
520+
start.SignalAndWait();
521+
wrapper.CompleteExternalEvent("Event", "\"payload\"");
522+
});
523+
await Task.WhenAll(cancellation, completion).WaitAsync(TimeSpan.FromSeconds(30));
524+
525+
// Assert
526+
if (wait.IsCanceled)
527+
{
528+
TaskCanceledException exception = await Assert.ThrowsAsync<TaskCanceledException>(() => wait);
529+
exception.CancellationToken.Should().Be(source.Token);
530+
Task<string> bufferedWait = wrapper.WaitForExternalEvent<string>("Event");
531+
bufferedWait.IsCompletedSuccessfully.Should().BeTrue();
532+
(await bufferedWait).Should().Be("payload");
533+
}
534+
else
535+
{
536+
wait.IsCompletedSuccessfully.Should().BeTrue();
537+
(await wait).Should().Be("payload");
538+
}
539+
540+
Task<string> nextWait = wrapper.WaitForExternalEvent<string>("Event");
541+
nextWait.IsCompleted.Should().BeFalse();
542+
wrapper.CompleteExternalEvent("Event", "\"next\"");
543+
(await nextWait).Should().Be("next");
544+
}
545+
}
546+
394547
[Fact]
395548
public void CompleteExternalEvent_CanceledWaiter_BuffersEventForNextWaiter()
396549
{

0 commit comments

Comments
 (0)