fix: health check waits for tasks to finish before becoming unhealthy

aneoconsulting · Oct 30, 2024 · 7c5a3e5 · 7c5a3e5
1 parent 8f6a2f9
commit 7c5a3e5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 12 deletions.
diff --git a/Common/src/Pollster/Pollster.cs b/Common/src/Pollster/Pollster.cs
@@ -43,6 +43,8 @@
 using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Microsoft.Extensions.Logging;
 
+using TaskStatus = ArmoniK.Core.Common.Storage.TaskStatus;
+
 namespace ArmoniK.Core.Common.Pollster;
 
 public class Pollster : IInitializable
@@ -170,12 +172,7 @@ public async Task Init(CancellationToken cancellationToken)
 
   public async Task<HealthCheckResult> Check(HealthCheckTag tag)
   {
-    if (healthCheckFailedResult_ is not null)
-    {
-      return healthCheckFailedResult_ ?? HealthCheckResult.Unhealthy("Health Check failed previously so this polling agent should be destroyed.");
-    }
-
-    if (endLoopReached_)
+    if (endLoopReached_ && taskProcessingDict_.IsEmpty)
     {
       return HealthCheckResult.Unhealthy("End of main loop reached, no more tasks will be executed.");
     }
@@ -234,6 +231,12 @@ public async Task<HealthCheckResult> Check(HealthCheckTag tag)
       healthCheckFailedResult_ = result;
     }
 
+    if (tag == HealthCheckTag.Liveness && taskProcessingDict_.Any(pair => pair.Value.GetAcquiredTaskInfo()
+                                                                              .TaskStatus != TaskStatus.Dispatched))
+    {
+      return HealthCheckResult.Healthy();
+    }
+
     if (tag == HealthCheckTag.Readiness && taskProcessingDict_.IsEmpty)
     {
       return HealthCheckResult.Unhealthy("No tasks to process");
@@ -264,11 +267,10 @@ await Init(exceptionManager_.EarlyCancellationToken)
         if (healthCheckFailedResult_ is not null)
         {
           var hcr = healthCheckFailedResult_.Value;
-          exceptionManager_.FatalError(logger_,
-                                       hcr.Exception,
-                                       "Health Check failed with status {Status} thus no more tasks will be executed:\n{Description}",
-                                       hcr.Status,
-                                       hcr.Description);
+          logger_.LogError(hcr.Exception,
+                           "Health Check failed with status {Status} thus no more tasks will be acquired (tasks already acquired will be executed to completion if possible):\n{Description}",
+                           hcr.Status,
+                           hcr.Description);
           return;
         }
 

diff --git a/Common/tests/Pollster/PollsterTest.cs b/Common/tests/Pollster/PollsterTest.cs
@@ -376,7 +376,6 @@ await testServiceProvider.Pollster.Init(CancellationToken.None)
     // This test that we return from the mainloop after the health check is unhealthy
     await testServiceProvider.Pollster.MainLoop()
                              .ConfigureAwait(false);
-    Assert.True(testServiceProvider.ExceptionManager.Failed);
   }
 
   [Test]