From 0dca581ef26edf33b7b17e1000a6098e9f02de7a Mon Sep 17 00:00:00 2001
From: Jordan <Cyberboss@users.noreply.github.com>
Date: Mon, 19 Jun 2023 00:52:56 -0400
Subject: [PATCH] Pass in CancellationToken to (Try)SetCanceled

---
 .../Components/Chat/Providers/DiscordProvider.cs             | 2 +-
 .../Components/Chat/Providers/IrcProvider.cs                 | 2 +-
 .../Components/Watchdog/WatchdogBase.cs                      | 5 +----
 src/Tgstation.Server.Host/Extensions/TaskExtensions.cs       | 3 +--
 src/Tgstation.Server.Host/Transfer/FileUploadProvider.cs     | 2 +-
 .../IO/TestRequestFileStreamProvider.cs                      | 2 +-
 6 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/src/Tgstation.Server.Host/Components/Chat/Providers/DiscordProvider.cs b/src/Tgstation.Server.Host/Components/Chat/Providers/DiscordProvider.cs
index 81b652cf49..b485661d3f 100644
--- a/src/Tgstation.Server.Host/Components/Chat/Providers/DiscordProvider.cs
+++ b/src/Tgstation.Server.Host/Components/Chat/Providers/DiscordProvider.cs
@@ -646,7 +646,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers
 				Task<Result> localGatewayTask;
 				gatewayReadyTcs = new TaskCompletionSource();
 
-				using var gatewayConnectionAbortRegistration = cancellationToken.Register(() => gatewayReadyTcs.TrySetCanceled());
+				using var gatewayConnectionAbortRegistration = cancellationToken.Register(() => gatewayReadyTcs.TrySetCanceled(cancellationToken));
 				gatewayCancellationToken.Register(() => Logger.LogTrace("Stopping gateway client..."));
 
 				// reconnects keep happening until we stop or it faults, our auto-reconnector will handle the latter
diff --git a/src/Tgstation.Server.Host/Components/Chat/Providers/IrcProvider.cs b/src/Tgstation.Server.Host/Components/Chat/Providers/IrcProvider.cs
index 3312e4b8b4..5e496dba40 100644
--- a/src/Tgstation.Server.Host/Components/Chat/Providers/IrcProvider.cs
+++ b/src/Tgstation.Server.Host/Components/Chat/Providers/IrcProvider.cs
@@ -407,7 +407,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers
 				await NonBlockingListen(cancellationToken);
 
 				var nickCheckCompleteTcs = new TaskCompletionSource();
-				using (cancellationToken.Register(() => nickCheckCompleteTcs.TrySetCanceled()))
+				using (cancellationToken.Register(() => nickCheckCompleteTcs.TrySetCanceled(cancellationToken)))
 				{
 					listenTask = Task.Factory.StartNew(
 					async () =>
diff --git a/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs b/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs
index 6df80546ca..0d2f0597d3 100644
--- a/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs
+++ b/src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs
@@ -862,20 +862,17 @@ namespace Tgstation.Server.Host.Components.Watchdog
 									cancellationToken);
 
 							// cancel waiting if requested
-							var cancelTcs = new TaskCompletionSource();
 							var toWaitOn = Task.WhenAny(
 								activeServerLifetime,
 								activeServerReboot,
 								activeServerStartup,
 								healthCheck,
 								newDmbAvailable,
-								cancelTcs.Task,
 								activeLaunchParametersChanged,
 								serverPrimed);
 
 							// wait for something to happen
-							using (cancellationToken.Register(() => cancelTcs.SetCanceled()))
-								await toWaitOn;
+							await toWaitOn.WithToken(cancellationToken);
 
 							cancellationToken.ThrowIfCancellationRequested();
 							Logger.LogTrace("Monitor activated");
diff --git a/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs b/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs
index 48644e0cd5..aeac45a0cc 100644
--- a/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs
+++ b/src/Tgstation.Server.Host/Extensions/TaskExtensions.cs
@@ -50,9 +50,8 @@ namespace Tgstation.Server.Host.Extensions
 			ArgumentNullException.ThrowIfNull(task);
 
 			var cancelTcs = new TaskCompletionSource();
-			using (cancellationToken.Register(() => cancelTcs.SetCanceled()))
+			using (cancellationToken.Register(() => cancelTcs.SetCanceled(cancellationToken)))
 				await Task.WhenAny(task, cancelTcs.Task);
-			cancellationToken.ThrowIfCancellationRequested();
 
 			return await task;
 		}
diff --git a/src/Tgstation.Server.Host/Transfer/FileUploadProvider.cs b/src/Tgstation.Server.Host/Transfer/FileUploadProvider.cs
index f81208601c..402cd6b951 100644
--- a/src/Tgstation.Server.Host/Transfer/FileUploadProvider.cs
+++ b/src/Tgstation.Server.Host/Transfer/FileUploadProvider.cs
@@ -69,7 +69,7 @@ namespace Tgstation.Server.Host.Transfer
 		/// <inheritdoc />
 		public async Task<Stream> GetResult(CancellationToken cancellationToken)
 		{
-			using (cancellationToken.Register(() => streamTcs.TrySetCanceled()))
+			using (cancellationToken.Register(() => streamTcs.TrySetCanceled(cancellationToken)))
 			using (ticketExpiryCts.Token.Register(() => streamTcs.TrySetResult(null)))
 				return await streamTcs.Task;
 		}
diff --git a/tests/Tgstation.Server.Host.Tests/IO/TestRequestFileStreamProvider.cs b/tests/Tgstation.Server.Host.Tests/IO/TestRequestFileStreamProvider.cs
index d9e0bdde91..638fa965c5 100644
--- a/tests/Tgstation.Server.Host.Tests/IO/TestRequestFileStreamProvider.cs
+++ b/tests/Tgstation.Server.Host.Tests/IO/TestRequestFileStreamProvider.cs
@@ -115,7 +115,7 @@ namespace Tgstation.Server.Host.IO.Tests
 				.Setup(x => x.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, It.IsAny<CancellationToken>()))
 				.Returns<HttpRequestMessage, HttpCompletionOption, CancellationToken>((request, option, cancellationToken) =>
 				{
-					cancellationToken.Register(() => tcs.TrySetCanceled());
+					cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken));
 					return tcs.Task;
 				})
 				.Verifiable();