Fix jobs to cancel w/ HostingBackgroundService shutdown

src/Microsoft.Health.Fhir.Core/Features/Operations/Export/ExportJobTask.cs

@@ -293,8 +293,8 @@ public async Task ExecuteAsync(ExportJobRecord exportJobRecord, WeakETag weakETa
             }
             catch (Exception ex) when ((ex is OperationCanceledException || ex is TaskCanceledException) && cancellationToken.IsCancellationRequested)
             {
-                _logger.LogInformation(ex, "[JobId:{JobId}] The job was canceled.", _exportJobRecord.Id);
-                await CompleteJobAsync(OperationStatus.Canceled, CancellationToken.None);
+                // Do nothing - JobHosting will cancel the job if needed. Job will be rescheduled if cancellation due to HostingBackgroudService shutdown.
+                _logger.LogInformation("[JobId:{JobId}] The job is being abandoned by ExportJobTask due to cancellation.", _exportJobRecord.Id);
             }
             catch (Exception ex)
             {

src/Microsoft.Health.TaskManagement/JobHosting.cs

@@ -97,7 +97,7 @@
                                 activity?.SetTag("Version", nextJob.Version);
 
                                 _logger.LogJobInformation(nextJob, "Job dequeued.");
                                 await ExecuteJobAsync(nextJob);
                             }
                         }
                         else
@@ -111,18 +111,27 @@
 
             try
             {
-                await Task.WhenAny(workers.ToArray()); // If any worker crashes exit.
+                 // If any worker crashes or complete after cancellation due to shutdown,
+                 // cancel all workers and wait for completion so they don't crash unnecessarily.
+                await Task.WhenAny(workers.ToArray());
+#if NET6_0
+                cancellationTokenSource.Cancel();
+#else
+                await cancellationTokenSource.CancelAsync();
+#endif
+                await Task.WhenAll(workers.ToArray());
             }
             catch (Exception ex)
             {
                 _logger.LogError(ex, "Job failed to execute. Queue type: {QueueType}", queueType);
             }
         }
 
-        private async Task ExecuteJobAsync(JobInfo jobInfo)
+        private async Task ExecuteJobAsync(JobInfo jobInfo, CancellationToken jobQueueCancellationToken)
         {
             EnsureArg.IsNotNull(jobInfo, nameof(jobInfo));
             using var jobCancellationToken = new CancellationTokenSource();
+            using var combinedCancellationTokens = CancellationTokenSource.CreateLinkedTokenSource(jobQueueCancellationToken, jobCancellationToken.Token);
 
             using IScoped<IJob> job = _jobFactory.Create(jobInfo);
 
@@ -153,7 +162,7 @@
                                     jobInfo.Version,
                                     cancellationSource => job.Value.ExecuteAsync(jobInfo, cancellationSource.Token),
                                     TimeSpan.FromSeconds(JobHeartbeatIntervalInSeconds),
-                                    jobCancellationToken);
+                                    combinedCancellationTokens);
 
                 jobInfo.Result = await runningJob;
             }
@@ -176,7 +185,14 @@
             }
             catch (OperationCanceledException ex)
             {
-                _logger.LogWarning(ex, "Job with id: {JobId} and group id: {GroupId} of type: {JobType} canceled.", jobInfo.Id, jobInfo.GroupId, jobInfo.QueueType);
+                // If cancellation not due to job cancellation, exit. Another worker will pickup/restart the job.
+                if (!jobInfo.CancelRequested)
+                {
+                    return;
+                }
+
+                // If cancellation due to job cancellation, update the status to cancelled.
+                _logger.LogJobInformation(ex, jobInfo, "Job with id: {JobId} and group id: {GroupId} of type: {JobType} canceled due to cancellation.", jobInfo.Id, jobInfo.GroupId, jobInfo.QueueType);
                 jobInfo.Status = JobStatus.Cancelled;
 
                 try
@@ -185,7 +201,7 @@
                 }
                 catch (Exception completeEx)
                 {
-                    _logger.LogError(completeEx, "Job with id: {JobId} and group id: {GroupId} of type: {JobType} failed to complete.", jobInfo.Id, jobInfo.GroupId, jobInfo.QueueType);
+                    _logger.LogJobError(completeEx, jobInfo, "Job with id: {JobId} and group id: {GroupId} of type: {JobType} failed to complete.", jobInfo.Id, jobInfo.GroupId, jobInfo.QueueType);
                 }
 
                 return;