Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 96 additions & 12 deletions TUnit.AspNetCore.Core/FlowSuppressingHostedService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,31 @@ namespace TUnit.AspNetCore;
/// Without this, wrapped services that own unmanaged resources leak silently because
/// the container only sees the non-disposable wrapper.
/// </para>
/// <para>
/// The stop-phase methods are additionally guarded to run the inner service's stop
/// exactly once. Minimal-hosting SUTs park their own <c>app.Run()</c> in
/// <c>WaitForShutdownAsync</c>, which calls <c>Host.StopAsync</c> when
/// <c>ApplicationStopping</c> fires — concurrently with the <c>Host.StopAsync</c> that
/// <c>WebApplicationFactory.DisposeAsync</c> already has in flight (the trigger of that
/// very <c>ApplicationStopping</c>). <c>Host.StopAsync</c> has no concurrency guard, so
/// every hosted service's <c>StopAsync</c> runs twice in parallel, breaking services
/// with non-thread-safe shutdown (e.g. Rebus' bus dispose throws
/// <see cref="ObjectDisposedException"/>). See https://github.com/thomhurst/TUnit/issues/6339.
/// All concurrent and repeated callers observe the single underlying stop task.
/// </para>
/// </remarks>
internal sealed class FlowSuppressingHostedService(IHostedService inner) : IHostedLifecycleService, IAsyncDisposable, IDisposable
{
private readonly object _stopGate = new();
private Task? _stopTask;
private Task? _stoppingTask;
private Task? _stoppedTask;

public Task StartAsync(CancellationToken cancellationToken) =>
RunOnCleanContext(inner.StartAsync, cancellationToken);

public Task StopAsync(CancellationToken cancellationToken) =>
inner.StopAsync(cancellationToken);
InvokeOnce(ref _stopTask, inner.StopAsync, cancellationToken);

public Task StartingAsync(CancellationToken cancellationToken) =>
inner is IHostedLifecycleService lifecycle
Expand All @@ -42,18 +59,30 @@ inner is IHostedLifecycleService lifecycle
? RunOnCleanContext(lifecycle.StartedAsync, cancellationToken)
: Task.CompletedTask;

// Stop lifecycle is intentionally not wrapped: stop methods typically signal
// cancellation and await shutdown rather than spawning new long-running background
// work, so context capture during Stop is not the span-leak vector that Start is.
public Task StoppingAsync(CancellationToken cancellationToken) =>
inner is IHostedLifecycleService lifecycle
? lifecycle.StoppingAsync(cancellationToken)
: Task.CompletedTask;
// Stop lifecycle is not flow-suppressed: stop methods typically signal cancellation
// and await shutdown rather than spawning new long-running background work, so
// context capture during Stop is not the span-leak vector that Start is. They are
// once-guarded, though — concurrent Host.StopAsync calls invoke these hooks twice
// (see the class remarks).
public Task StoppingAsync(CancellationToken cancellationToken)
{
if (inner is not IHostedLifecycleService lifecycle)
{
return Task.CompletedTask;
}

public Task StoppedAsync(CancellationToken cancellationToken) =>
inner is IHostedLifecycleService lifecycle
? lifecycle.StoppedAsync(cancellationToken)
: Task.CompletedTask;
return InvokeOnce(ref _stoppingTask, lifecycle.StoppingAsync, cancellationToken);
}

public Task StoppedAsync(CancellationToken cancellationToken)
{
if (inner is not IHostedLifecycleService lifecycle)
{
return Task.CompletedTask;
}

return InvokeOnce(ref _stoppedTask, lifecycle.StoppedAsync, cancellationToken);
}

/// <inheritdoc />
public async ValueTask DisposeAsync()
Expand Down Expand Up @@ -83,6 +112,61 @@ public void Dispose()
}
}

// Publish a placeholder before invoking the operation outside the lock. Duplicate callers
// can immediately wait with their own cancellation token, even when the operation blocks
// synchronously before returning its task.
private Task InvokeOnce(
ref Task? cachedTask,
Func<CancellationToken, Task> operation,
CancellationToken cancellationToken)
{
TaskCompletionSource? completionSource = null;
Task task;
var ownsInvocation = false;

lock (_stopGate)
{
if (cachedTask is null)
{
completionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
cachedTask = completionSource.Task;
ownsInvocation = true;
}

task = cachedTask;
}

if (completionSource is not null)
{
_ = ExecuteAndCompleteAsync(completionSource, operation, cancellationToken);
}

return ownsInvocation ? task : WaitWithCancellation(task, cancellationToken);
}

private static async Task ExecuteAndCompleteAsync(
TaskCompletionSource completionSource,
Func<CancellationToken, Task> operation,
CancellationToken cancellationToken)
{
try
{
await operation(cancellationToken).ConfigureAwait(false);
completionSource.SetResult();
}
catch (OperationCanceledException exception)
{
completionSource.SetCanceled(exception.CancellationToken);
}
catch (Exception exception)
{
completionSource.SetException(exception);
}
}

private static Task WaitWithCancellation(Task task, CancellationToken cancellationToken) =>
cancellationToken.CanBeCanceled ? task.WaitAsync(cancellationToken) : task;

// Dispatch onto a thread-pool worker with a clean captured ExecutionContext by
// combining SuppressFlow + Task.Run. Unlike wrapping `using (SuppressFlow()) return op(ct);`
// which only suppresses during the synchronous body, this keeps the inner operation
Expand Down
72 changes: 72 additions & 0 deletions TUnit.AspNetCore.Tests/HostStopExactlyOnceProbeTests.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using TUnit.AspNetCore;

namespace TUnit.AspNetCore.Tests;

/// <summary>
/// End-to-end regression test for https://github.com/thomhurst/TUnit/issues/6339.
/// <para>
/// With a minimal-hosting SUT, disposing the per-test factory races the SUT's own
/// <c>app.Run()</c> shutdown path: <c>WebApplicationFactory.DisposeAsync</c> calls
/// <c>Host.StopAsync</c>, whose <c>ApplicationStopping</c> signal wakes the app's parked
/// <c>WaitForShutdownAsync</c>, which calls <c>Host.StopAsync</c> again — concurrently.
/// Before the stop-once guard in <see cref="FlowSuppressingHostedService"/>, each hosted
/// service's <c>StopAsync</c> ran twice in parallel (observed here as 4–8 failures per
/// 50 tests), which is how Rebus' non-thread-safe bus dispose produced the reporter's
/// <see cref="ObjectDisposedException"/>.
/// </para>
/// <para>
/// The probe throws — with both stack traces — if it ever observes a second
/// <c>StopAsync</c> entry, so a regression surfaces as loud test failures.
/// </para>
/// </summary>
public class HostStopExactlyOnceProbeTests : WebApplicationTest<TestWebAppFactory, Program>
{
protected override void ConfigureTestServices(IServiceCollection services)
{
services.AddSingleton<IHostedService, StopProbeHostedService>();
}

[Test]
[Repeat(50)]
public async Task Host_Hosted_Service_Stops_Exactly_Once_Under_Parallel_Churn()
{
var client = Factory.CreateClient();
using var response = await client.GetAsync("/");

// Small jitter so test completions (and After-hook disposals) overlap heavily.
await Task.Delay(Random.Shared.Next(0, 20));
}
}

/// <summary>
/// Throws on a second StopAsync entry — sequential re-stop or concurrent overlap — so the
/// offending call site's stack appears in the test output alongside the first stop's stack.
/// </summary>
internal sealed class StopProbeHostedService : IHostedService
{
private int _stopEntries;
private volatile string? _firstStopStack;

public Task StartAsync(CancellationToken cancellationToken) => Task.CompletedTask;

public async Task StopAsync(CancellationToken cancellationToken)
{
var entry = Interlocked.Increment(ref _stopEntries);

if (entry > 1)
{
throw new InvalidOperationException(
$"StopAsync entered {entry} times on probe {GetHashCode():x8}.\n" +
$"--- first stop stack ---\n{_firstStopStack}\n" +
$"--- this stop stack ---\n{Environment.StackTrace}");
}

_firstStopStack = Environment.StackTrace;

// Widen the window like a real bus shutdown (Rebus stops workers, waits on its
// cleanup task for up to 5s). A concurrent second stop lands inside this delay.
await Task.Delay(100, CancellationToken.None);
}
}
Loading
Loading