diff --git a/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs b/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs index aeac45a0cc..e2dbb8f8ff 100644 --- a/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs +++ b/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs @@ -50,8 +50,12 @@ namespace Tgstation.Server.Host.Extensions ArgumentNullException.ThrowIfNull(task); var cancelTcs = new TaskCompletionSource(); + Task completedTask; using (cancellationToken.Register(() => cancelTcs.SetCanceled(cancellationToken))) - await Task.WhenAny(task, cancelTcs.Task); + completedTask = await Task.WhenAny(task, cancelTcs.Task); + + if (completedTask != task) + await cancelTcs.Task; return await task; } diff --git a/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs b/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs index 3f8e8663d8..ed8660ec0e 100644 --- a/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs +++ b/tests/Tgstation.Server.Host.Tests/Jobs/TestJobHandler.cs @@ -41,7 +41,7 @@ namespace Tgstation.Server.Host.Jobs.Tests using var handler = new JobHandler(TestJob); await Assert.ThrowsExceptionAsync(() => handler.Wait(cts.Token)); handler.Start(); - await Assert.ThrowsExceptionAsync(() => handler.Wait(cts.Token)); + await Assert.ThrowsExceptionAsync(() => handler.Wait(cts.Token)); tcs.SetResult(); await handler.Wait(default); }