Switch from .WithToken to .WaitAsync

src/Tgstation.Server.Host/Components/Byond/ByondManager.cs

@@ -235,7 +235,7 @@ namespace Tgstation.Server.Host.Components.Byond
 				await Task.WhenAny(
 					containerTask,
 					activeVersionUpdate)
-					.WithToken(cancellationToken);
+					.WaitAsync(cancellationToken);
 
 				if (containerTask.IsCompleted)
 					logger.LogTrace("All BYOND locks for {version} are gone", version);
@@ -455,7 +455,7 @@ namespace Tgstation.Server.Host.Components.Byond
 					if (neededForLock && !installation.InstallationTask.IsCompleted)
 						logger.LogWarning("The required BYOND version ({version}) is not readily available! We will have to wait for it to install.", version);
 
-					await installation.InstallationTask.WithToken(cancellationToken);
+					await installation.InstallationTask.WaitAsync(cancellationToken);
 					return installLock;
 				}
 

src/Tgstation.Server.Host/Components/Chat/ChatManager.cs

@@ -482,7 +482,7 @@ namespace Tgstation.Server.Host.Components.Chat
 			if (waitingForInitialConnection)
 			{
 				logger.LogTrace("Waiting for initial chat bot connections before updating tracking contexts...");
-				await initialProviderConnectionsTask.WithToken(cancellationToken);
+				await initialProviderConnectionsTask.WaitAsync(cancellationToken);
 			}
 
 			List<Task> tasks;
@@ -932,7 +932,7 @@ namespace Tgstation.Server.Host.Components.Chat
 					if (messageTasks.Count == 0)
 					{
 						logger.LogTrace("No providers active, pausing messsage monitoring...");
-						await updatedTask.WithToken(cancellationToken);
+						await updatedTask.WaitAsync(cancellationToken);
 						logger.LogTrace("Resuming message monitoring...");
 						continue;
 					}
@@ -1073,7 +1073,7 @@ namespace Tgstation.Server.Host.Components.Chat
 			{
 				var cancellationToken = handlerCts.Token;
 				if (waitForConnections)
-					await initialProviderConnectionsTask.WithToken(cancellationToken);
+					await initialProviderConnectionsTask.WaitAsync(cancellationToken);
 
 				await SendMessage(
 					channelIdsFactory(),

src/Tgstation.Server.Host/Components/Chat/Providers/IrcProvider.cs

@@ -377,7 +377,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers
 					cancellationToken,
 					DefaultIOManager.BlockingTaskCreationOptions,
 					TaskScheduler.Current)
-					.WithToken(cancellationToken);
+					.WaitAsync(cancellationToken);
 
 				cancellationToken.ThrowIfCancellationRequested();
 
@@ -584,7 +584,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers
 			cancellationToken,
 			TaskCreationOptions.None,
 			TaskScheduler.Current)
-			.WithToken(cancellationToken);
+			.WaitAsync(cancellationToken);
 
 		/// <summary>
 		/// Run SASL authentication on <see cref="client"/>.

src/Tgstation.Server.Host/Components/Chat/Providers/Provider.cs

@@ -149,7 +149,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers
 		{
 			while (true)
 			{
-				await nextMessage.Task.WithToken(cancellationToken);
+				await nextMessage.Task.WaitAsync(cancellationToken);
 				lock (messageQueue)
 					if (messageQueue.Count > 0)
 					{

src/Tgstation.Server.Host/Components/Deployment/DreamMaker.cs

@@ -806,10 +806,10 @@ namespace Tgstation.Server.Host.Components.Deployment
 			using (var provider = new TemporaryDmbProvider(ioManager.ResolvePath(job.DirectoryName.ToString()), String.Concat(job.DmeName, DmbExtension), job))
 			await using (var controller = await sessionControllerFactory.LaunchNew(provider, byondLock, launchParameters, true, cancellationToken))
 			{
-				var launchResult = await controller.LaunchResult.WithToken(cancellationToken);
+				var launchResult = await controller.LaunchResult.WaitAsync(cancellationToken);
 
 				if (launchResult.StartupTime.HasValue)
-					await controller.Lifetime.WithToken(cancellationToken);
+					await controller.Lifetime.WaitAsync(cancellationToken);
 
 				if (!controller.Lifetime.IsCompleted)
 					await controller.DisposeAsync();

src/Tgstation.Server.Host/Components/Deployment/Remote/GitLabRemoteDeploymentManager.cs

@@ -59,7 +59,7 @@ namespace Tgstation.Server.Host.Components.Deployment.Remote
 					.GetAsync(
 						$"{repository.RemoteRepositoryOwner}/{repository.RemoteRepositoryName}",
 						x.TestMerge.Number)
-					.WithToken(cancellationToken));
+					.WaitAsync(cancellationToken));
 			try
 			{
 				await Task.WhenAll(tasks);
@@ -138,7 +138,7 @@ namespace Tgstation.Server.Host.Components.Deployment.Remote
 					$"{remoteRepositoryOwner}/{remoteRepositoryName}",
 					testMergeNumber,
 					new CreateMergeRequestNoteRequest(comment))
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 		}
 
 		/// <inheritdoc />

src/Tgstation.Server.Host/Components/InstanceManager.cs

@@ -323,7 +323,7 @@ namespace Tgstation.Server.Host.Components
 
 				try
 				{
-					await container.OnZeroReferences.WithToken(cancellationToken);
+					await container.OnZeroReferences.WaitAsync(cancellationToken);
 
 					// we are the one responsible for cancelling his jobs
 					var tasks = new List<Task>();

src/Tgstation.Server.Host/Components/Repository/GitLabRemoteFeatures.cs

@@ -66,7 +66,7 @@ namespace Tgstation.Server.Host.Components.Repository
 				var mr = await client
 					.MergeRequests
 					.GetAsync($"{RemoteRepositoryOwner}/{RemoteRepositoryName}", parameters.Number)
-					.WithToken(cancellationToken);
+					.WaitAsync(cancellationToken);
 
 				var revisionToUse = parameters.TargetCommitSha == null
 					|| mr.Sha.StartsWith(parameters.TargetCommitSha, StringComparison.OrdinalIgnoreCase)

src/Tgstation.Server.Host/Components/Session/SessionController.cs

@@ -779,7 +779,7 @@ namespace Tgstation.Server.Host.Components.Session
 						}
 
 						Interlocked.Exchange(ref rebootTcs, new TaskCompletionSource()).SetResult();
-						await RebootGate.WithToken(cancellationToken);
+						await RebootGate.WaitAsync(cancellationToken);
 					}
 					finally
 					{

src/Tgstation.Server.Host/Components/Watchdog/WatchdogBase.cs

@@ -411,7 +411,7 @@ namespace Tgstation.Server.Host.Components.Watchdog
 				if (Status != WatchdogStatus.Offline)
 				{
 					Logger.LogDebug("Waiting for server to gracefully shut down.");
-					await monitorTask.WithToken(cancellationToken);
+					await monitorTask.WaitAsync(cancellationToken);
 				}
 				else
 					Logger.LogTrace("Graceful shutdown requested but server is already offline.");
@@ -589,7 +589,7 @@ namespace Tgstation.Server.Host.Components.Watchdog
 		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
 		protected async Task CheckLaunchResult(ISessionController controller, string serverName, CancellationToken cancellationToken)
 		{
-			var launchResult = await controller.LaunchResult.WithToken(cancellationToken);
+			var launchResult = await controller.LaunchResult.WaitAsync(cancellationToken);
 
 			// Dead sessions won't trigger this
 			if (launchResult.ExitCode.HasValue) // you killed us ray...
@@ -872,7 +872,7 @@ namespace Tgstation.Server.Host.Components.Watchdog
 								serverPrimed);
 
 							// wait for something to happen
-							await toWaitOn.WithToken(cancellationToken);
+							await toWaitOn.WaitAsync(cancellationToken);
 
 							cancellationToken.ThrowIfCancellationRequested();
 							Logger.LogTrace("Monitor activated");

src/Tgstation.Server.Host/Core/Application.cs

@@ -456,7 +456,7 @@ namespace Tgstation.Server.Host.Core
 
 			// 503 requests made while the application is starting
 			applicationBuilder.UseAsyncInitialization<IInstanceManager>(
-				(instanceManager, cancellationToken) => instanceManager.Ready.WithToken(cancellationToken));
+				(instanceManager, cancellationToken) => instanceManager.Ready.WaitAsync(cancellationToken));
 
 			if (generalConfiguration.HostApiDocumentation)
 			{

src/Tgstation.Server.Host/Core/CommandPipeReader.cs

@@ -69,7 +69,7 @@ namespace Tgstation.Server.Host.Core
 				while (!cancellationToken.IsCancellationRequested)
 				{
 					logger.LogTrace("Waiting to read command line...");
-					var line = await streamReader.ReadLineAsync().WithToken(cancellationToken);
+					var line = await streamReader.ReadLineAsync().WaitAsync(cancellationToken);
 
 					logger?.LogInformation("Received pipe command: {command}", line);
 					switch (line)

src/Tgstation.Server.Host/Extensions/TaskExtensions.cs

@@ -1,6 +1,4 @@
-using System;
-using System.Threading;
-using System.Threading.Tasks;
+using System.Threading.Tasks;
 
 namespace Tgstation.Server.Host.Extensions
 {
@@ -18,46 +16,5 @@ namespace Tgstation.Server.Host.Extensions
 		/// Gets a <see cref="Task"/> that never completes.
 		/// </summary>
 		public static Task InfiniteTask => InfiniteTaskCompletionSource.Task;
-
-		/// <summary>
-		/// Create a <see cref="Task"/> that can be awaited while respecting a given <paramref name="cancellationToken"/>.
-		/// </summary>
-		/// <param name="task">The <see cref="Task"/> to add cancel support to.</param>
-		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		public static Task WithToken(this Task task, CancellationToken cancellationToken)
-		{
-			ArgumentNullException.ThrowIfNull(task);
-
-			async Task<object> Wrap()
-			{
-				await task;
-				return null;
-			}
-
-			return Wrap().WithToken(cancellationToken);
-		}
-
-		/// <summary>
-		/// Create a <see cref="Task{TResult}"/> that can be awaited while respecting a given <paramref name="cancellationToken"/>.
-		/// </summary>
-		/// <typeparam name="T">The result <see cref="Type"/> of the <paramref name="task"/>.</typeparam>
-		/// <param name="task">The <see cref="Task{TResult}"/> to add cancel support to.</param>
-		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task{TResult}"/> resulting in the result of <paramref name="task"/>.</returns>
-		public static async Task<T> WithToken<T>(this Task<T> task, CancellationToken cancellationToken)
-		{
-			ArgumentNullException.ThrowIfNull(task);
-
-			var cancelTcs = new TaskCompletionSource();
-			Task completedTask;
-			using (cancellationToken.Register(() => cancelTcs.SetCanceled(cancellationToken)))
-				completedTask = await Task.WhenAny(task, cancelTcs.Task);
-
-			if (completedTask != task)
-				await cancelTcs.Task;
-
-			return await task;
-		}
 	}
 }

src/Tgstation.Server.Host/IO/Console.cs

@@ -102,7 +102,7 @@ namespace Tgstation.Server.Host.IO
 			cancellationToken,
 			DefaultIOManager.BlockingTaskCreationOptions,
 			TaskScheduler.Current)
-			.WithToken(cancellationToken);
+			.WaitAsync(cancellationToken);
 
 		/// <inheritdoc />
 		public Task WriteAsync(string text, bool newLine, CancellationToken cancellationToken) => Task.Factory.StartNew(

src/Tgstation.Server.Host/IO/DefaultIOManager.cs

@@ -389,7 +389,7 @@ namespace Tgstation.Server.Host.IO
 
 				async Task CopyThisFile()
 				{
-					await subdirCreationTask.WithToken(cancellationToken);
+					await subdirCreationTask.WaitAsync(cancellationToken);
 					using var lockContext = semaphore != null
 						? await SemaphoreSlimContext.Lock(semaphore, cancellationToken)
 						: null;

src/Tgstation.Server.Host/Jobs/JobHandler.cs

@@ -64,7 +64,7 @@ namespace Tgstation.Server.Host.Jobs
 			if (task == null)
 				throw new InvalidOperationException("Job not started!");
 
-			return task.WithToken(cancellationToken);
+			return task.WaitAsync(cancellationToken);
 		}
 
 		/// <summary>

src/Tgstation.Server.Host/Jobs/JobService.cs

@@ -265,7 +265,7 @@ namespace Tgstation.Server.Host.Jobs
 			}
 
 			if (noMoreJobsShouldStart && !handler.Started)
-				await Extensions.TaskExtensions.InfiniteTask.WithToken(cancellationToken);
+				await Extensions.TaskExtensions.InfiniteTask.WaitAsync(cancellationToken);
 
 			Task cancelTask = null;
 			using (jobCancellationToken.Register(() => cancelTask = CancelJob(job, canceller, true, cancellationToken)))
@@ -320,7 +320,7 @@ namespace Tgstation.Server.Host.Jobs
 								}
 						}
 
-						var instanceCoreProvider = await activationTcs.Task.WithToken(cancellationToken);
+						var instanceCoreProvider = await activationTcs.Task.WaitAsync(cancellationToken);
 
 						logger.LogTrace("Starting job...");
 						await operation(

src/Tgstation.Server.Host/Swarm/SwarmService.cs

@@ -282,7 +282,7 @@ namespace Tgstation.Server.Host.Swarm
 				? asyncDelayer.Delay(
 					TimeSpan.FromMinutes(SwarmConstants.UpdateCommitTimeoutMinutes),
 					cancellationToken)
-				: Extensions.TaskExtensions.InfiniteTask.WithToken(cancellationToken);
+				: Extensions.TaskExtensions.InfiniteTask.WaitAsync(cancellationToken);
 
 			var commitTask = Task.WhenAny(localUpdateOperation.CommitGate, timeoutTask);
 

src/Tgstation.Server.Host/System/ProcessExecutor.cs

@@ -262,7 +262,7 @@ namespace Tgstation.Server.Host.System
 				if (errorOpen && errorReadTask == null)
 					errorReadTask = stdErrHandle.ReadLineAsync();
 
-				var completedTask = await Task.WhenAny(outputReadTask ?? errorReadTask, errorReadTask ?? outputReadTask).WithToken(disposeToken);
+				var completedTask = await Task.WhenAny(outputReadTask ?? errorReadTask, errorReadTask ?? outputReadTask).WaitAsync(disposeToken);
 				var line = await completedTask;
 				if (completedTask == outputReadTask)
 				{

src/Tgstation.Server.Host/System/SystemDManager.cs

@@ -128,7 +128,7 @@ namespace Tgstation.Server.Host.System
 
 			try
 			{
-				await instanceManager.Ready.WithToken(cancellationToken);
+				await instanceManager.Ready.WaitAsync(cancellationToken);
 				CheckReady();
 
 				var watchdogUsec = Environment.GetEnvironmentVariable("WATCHDOG_USEC");

src/Tgstation.Server.Host/Transfer/FileTransferService.cs

@@ -258,7 +258,7 @@ namespace Tgstation.Server.Host.Transfer
 				var expireAt = DateTimeOffset.UtcNow + TimeSpan.FromMinutes(TicketValidityMinutes);
 				try
 				{
-					await oldExpireTask.WithToken(disposeCts.Token);
+					await oldExpireTask.WaitAsync(disposeCts.Token);
 
 					var now = DateTimeOffset.UtcNow;
 					if (now < expireAt)

src/Tgstation.Server.Host/Transfer/FileUploadProvider.cs

@@ -126,7 +126,7 @@ namespace Tgstation.Server.Host.Transfer
 			{
 				streamTcs.TrySetResult(bufferedStream ?? stream);
 
-				await completionTcs.Task.WithToken(cancellationToken);
+				await completionTcs.Task.WaitAsync(cancellationToken);
 				return errorMessage;
 			}
 		}

src/Tgstation.Server.Host/Utils/GitHub/GitHubService.cs

@@ -65,7 +65,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 					{
 						RedirectUri = oAuthConfiguration.RedirectUrl,
 					})
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 
 			var token = response.AccessToken;
 			return token;
@@ -79,7 +79,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 				.Repository
 				.Release
 					.GetAll(updatesConfiguration.GitHubRepositoryId)
-					.WithToken(cancellationToken);
+					.WaitAsync(cancellationToken);
 
 			logger.LogTrace("{totalReleases} total releases", allReleases.Count);
 			var releases = allReleases
@@ -123,7 +123,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 			var repository = await gitHubClient
 				.Repository
 					.Get(updatesConfiguration.GitHubRepositoryId)
-					.WithToken(cancellationToken);
+					.WaitAsync(cancellationToken);
 
 			var repoUrl = new Uri(repository.HtmlUrl);
 			logger.LogTrace("Maps to {repostioryUrl}", repoUrl);
@@ -136,7 +136,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 		{
 			logger.LogTrace("CreateOAuthAccessToken");
 
-			var userDetails = await gitHubClient.User.Current().WithToken(cancellationToken);
+			var userDetails = await gitHubClient.User.Current().WaitAsync(cancellationToken);
 			return userDetails.Id;
 		}
 
@@ -159,7 +159,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 					repoName,
 					issueNumber,
 					comment)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 		}
 
 		/// <inheritdoc />
@@ -176,7 +176,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 				.Get(
 					repoOwner,
 					repoName)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 
 			return repo.Id;
 		}
@@ -199,7 +199,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 					repoOwner,
 					repoName,
 					newDeployment)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 
 			return deployment.Id;
 		}
@@ -223,7 +223,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 					repoName,
 					deploymentId,
 					newDeploymentStatus)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 		}
 
 		/// <inheritdoc />
@@ -240,7 +240,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 					repoId,
 					deploymentId,
 					newDeploymentStatus)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 		}
 
 		/// <inheritdoc />
@@ -258,7 +258,7 @@ namespace Tgstation.Server.Host.Utils.GitHub
 					repoOwner,
 					repoName,
 					pullRequestNumber)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 		}
 	}
 }

tests/Tgstation.Server.Host.Tests.Signals/Program.cs

@@ -40,17 +40,17 @@ namespace Tgstation.Server.Host.Tests.Signals
 			Assert.IsFalse(tcs.Task.IsCompleted);
 
 			using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
-			await signalHandler.StopAsync(default).WithToken(cts.Token);
+			await signalHandler.StopAsync(default).WaitAsync(cts.Token);
 			Assert.IsFalse(tcs.Task.IsCompleted);
 
 			using var signalHandler2 = new PosixSignalHandler(mockServerControl.Object, mockAsyncDelayer.Object, Mock.Of<ILogger<PosixSignalHandler>>());
 			await signalHandler2.StartAsync(default);
 
 			using var cts2 = new CancellationTokenSource(TimeSpan.FromSeconds(20));
-			await tcs.Task.WithToken(cts2.Token);
+			await tcs.Task.WaitAsync(cts2.Token);
 
 			using var cts3 = new CancellationTokenSource(TimeSpan.FromSeconds(5));
-			await signalHandler2.StopAsync(default).WithToken(cts3.Token);
+			await signalHandler2.StopAsync(default).WaitAsync(cts3.Token);
 		}
 	}
 }

tests/Tgstation.Server.Tests/Live/DummyGitHubService.cs

@@ -45,7 +45,7 @@ namespace Tgstation.Server.Tests.Live
 					.Repository
 					.Release
 					.GetAll("tgstation", "tgstation-server")
-					.WithToken(cancellationToken);
+					.WaitAsync(cancellationToken);
 
 				targetRelease = releases.FirstOrDefault(release => release.TagName == $"{new UpdatesConfiguration().GitTagPrefix}{TestLiveServer.TestUpdateVersion}");
 			}
@@ -59,7 +59,7 @@ namespace Tgstation.Server.Tests.Live
 			testPr = await gitHubClient
 				.PullRequest
 				.Get("Cyberboss", "common_core", 2)
-				.WithToken(cancellationToken);
+				.WaitAsync(cancellationToken);
 
 			ServiceCollectionExtensions.UseGitHubServiceFactory<DummyGitHubServiceFactory>();
 		}

tests/Tgstation.Server.Tests/Live/Instance/WatchdogTest.cs

@@ -590,7 +590,7 @@ namespace Tgstation.Server.Tests.Live.Instance
 				var bridgeTestTopicResult = await TopicClient.SendTopic(IPAddress.Loopback, "tgs_integration_test_tactics2=1", TestLiveServer.DDPort, cancellationToken);
 				Assert.AreEqual("ack2", bridgeTestTopicResult.StringData);
 
-				await bridgeTestsTcs.Task.WithToken(cancellationToken);
+				await bridgeTestsTcs.Task.WaitAsync(cancellationToken);
 			}
 
 			BridgeController.LogContent = true;

tests/Tgstation.Server.Tests/Live/TestLiveServer.cs

@@ -1303,7 +1303,7 @@ namespace Tgstation.Server.Tests.Live
 				serverCts.Cancel();
 				try
 				{
-					await serverTask.WithToken(hardCancellationToken);
+					await serverTask.WaitAsync(hardCancellationToken);
 				}
 				catch (OperationCanceledException) { }
 

tests/Tgstation.Server.Tests/TestSystemInteraction.cs

@@ -30,7 +30,7 @@ namespace Tgstation.Server.Tests
 			await using var process = await processExecutor.LaunchProcess("test." + platformIdentifier.ScriptFileExtension, ".", string.Empty, null, true, true);
 			using var cts = new CancellationTokenSource();
 			cts.CancelAfter(3000);
-			var exitCode = await process.Lifetime.WithToken(cts.Token);
+			var exitCode = await process.Lifetime.WaitAsync(cts.Token);
 
 			Assert.AreEqual(0, exitCode);
 			var result = (await process.GetCombinedOutput(default)).Trim();
@@ -59,7 +59,7 @@ namespace Tgstation.Server.Tests
 				{
 					using var cts = new CancellationTokenSource();
 					cts.CancelAfter(3000);
-					var exitCode = await process.Lifetime.WithToken(cts.Token);
+					var exitCode = await process.Lifetime.WaitAsync(cts.Token);
 
 					await process.GetCombinedOutput(cts.Token);