From bab9d81aa0a5a74d2d479ae3c06f06bb9d02ee69 Mon Sep 17 00:00:00 2001
From: Dominion <Cyberboss@users.noreply.github.com>
Date: Mon, 5 Jun 2023 21:24:12 -0400
Subject: [PATCH] Fix attempting to wait on a job after the JobService has
 stopped

---
 .../Components/Instance.cs                    |  6 ++---
 src/Tgstation.Server.Host/Jobs/JobHandler.cs  | 25 +++++++++++--------
 src/Tgstation.Server.Host/Jobs/JobService.cs  | 10 ++++++++
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/src/Tgstation.Server.Host/Components/Instance.cs b/src/Tgstation.Server.Host/Components/Instance.cs
index 95ca246c42..0e441edb15 100644
--- a/src/Tgstation.Server.Host/Components/Instance.cs
+++ b/src/Tgstation.Server.Host/Components/Instance.cs
@@ -514,11 +514,9 @@ namespace Tgstation.Server.Host.Components
 						await jobManager.RegisterOperation(
 							repositoryUpdateJob,
 							RepositoryAutoUpdateJob,
-							cancellationToken)
-							;
+							cancellationToken);
 
-						// DCT: First token will cancel the job, second is for cancelling the cancellation, unwanted
-						await jobManager.WaitForJobCompletion(repositoryUpdateJob, null, cancellationToken, default);
+						await jobManager.WaitForJobCompletion(repositoryUpdateJob, null, cancellationToken, cancellationToken);
 
 						Job compileProcessJob;
 						using (var repo = await RepositoryManager.LoadRepository(cancellationToken))
diff --git a/src/Tgstation.Server.Host/Jobs/JobHandler.cs b/src/Tgstation.Server.Host/Jobs/JobHandler.cs
index 19941fbefd..98cbe88651 100644
--- a/src/Tgstation.Server.Host/Jobs/JobHandler.cs
+++ b/src/Tgstation.Server.Host/Jobs/JobHandler.cs
@@ -11,6 +11,21 @@ namespace Tgstation.Server.Host.Jobs
 	/// </summary>
 	sealed class JobHandler : IDisposable
 	{
+		/// <summary>
+		/// If the job has started.
+		/// </summary>
+		public bool Started => task != null;
+
+		/// <summary>
+		/// The progress of the job.
+		/// </summary>
+		public int? Progress { get; set; }
+
+		/// <summary>
+		/// The stage of the job.
+		/// </summary>
+		public string Stage { get; set; }
+
 		/// <summary>
 		/// The <see cref="CancellationTokenSource"/> for <see cref="task"/>.
 		/// </summary>
@@ -39,16 +54,6 @@ namespace Tgstation.Server.Host.Jobs
 		/// <inheritdoc />
 		public void Dispose() => cancellationTokenSource.Dispose();
 
-		/// <summary>
-		/// The progress of the job.
-		/// </summary>
-		public int? Progress { get; set; }
-
-		/// <summary>
-		/// The stage of the job.
-		/// </summary>
-		public string Stage { get; set; }
-
 		/// <summary>
 		/// Wait for <see cref="task"/> to complete.
 		/// </summary>
diff --git a/src/Tgstation.Server.Host/Jobs/JobService.cs b/src/Tgstation.Server.Host/Jobs/JobService.cs
index c0d01f2dcf..09350a4c4c 100644
--- a/src/Tgstation.Server.Host/Jobs/JobService.cs
+++ b/src/Tgstation.Server.Host/Jobs/JobService.cs
@@ -256,13 +256,23 @@ namespace Tgstation.Server.Host.Jobs
 		{
 			if (job == null)
 				throw new ArgumentNullException(nameof(job));
+
+			if (!cancellationToken.CanBeCanceled)
+				throw new ArgumentException("A cancellable CancellationToken should be provided!", nameof(cancellationToken));
+
 			JobHandler handler;
+			bool noMoreJobsShouldStart;
 			lock (synchronizationLock)
 			{
 				if (!jobs.TryGetValue(job.Id.Value, out handler))
 					return;
+
+				noMoreJobsShouldStart = this.noMoreJobsShouldStart;
 			}
 
+			if (noMoreJobsShouldStart && !handler.Started)
+				await Extensions.TaskExtensions.InfiniteTask().WithToken(cancellationToken);
+
 			Task cancelTask = null;
 			using (jobCancellationToken.Register(() => cancelTask = CancelJob(job, canceller, true, cancellationToken)))
 				await handler.Wait(cancellationToken);