diff --git a/Adaptors/MongoDB/src/AuthenticationTable.cs b/Adaptors/MongoDB/src/AuthenticationTable.cs index 5a3206d6a..6b56c0e86 100644 --- a/Adaptors/MongoDB/src/AuthenticationTable.cs +++ b/Adaptors/MongoDB/src/AuthenticationTable.cs @@ -27,6 +27,7 @@ using ArmoniK.Core.Adapters.MongoDB.Table.DataModel.Auth; using ArmoniK.Core.Base.DataStructures; using ArmoniK.Core.Common.Auth.Authentication; +using ArmoniK.Core.Utils; using JetBrains.Annotations; @@ -137,10 +138,20 @@ public void AddCertificates(IEnumerable certificates) } /// - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(AuthenticationTable)} is not initialized", + sessionProvider_, + userCollectionProvider_, + authCollectionProvider_, + roleCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } /// public async Task Init(CancellationToken cancellationToken) diff --git a/Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs b/Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs index 99a4ed1f8..a99030546 100644 --- a/Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs +++ b/Adaptors/MongoDB/src/Common/MongoCollectionProvider.cs @@ -67,10 +67,11 @@ public Task Check(HealthCheckTag tag) { HealthCheckTag.Startup or HealthCheckTag.Readiness => Task.FromResult(isInitialized_ ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy("MongoCollection not initialized yet.")), - HealthCheckTag.Liveness => Task.FromResult(isInitialized_ && mongoCollection_ is null + : HealthCheckResult + .Unhealthy($"Mongo Collection<{typeof(TData)}> not initialized yet.")), + HealthCheckTag.Liveness => Task.FromResult(isInitialized_ && mongoCollection_ is not null ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy("MongoCollection not initialized yet.")), + : HealthCheckResult.Unhealthy($"Mongo Collection<{typeof(TData)}> not initialized yet.")), _ => throw new ArgumentOutOfRangeException(nameof(tag), tag, null), diff --git a/Adaptors/MongoDB/src/PartitionTable.cs b/Adaptors/MongoDB/src/PartitionTable.cs index d6371bbb4..69e7cf9cf 100644 --- a/Adaptors/MongoDB/src/PartitionTable.cs +++ b/Adaptors/MongoDB/src/PartitionTable.cs @@ -30,6 +30,7 @@ using ArmoniK.Core.Base.DataStructures; using ArmoniK.Core.Common.Exceptions; using ArmoniK.Core.Common.Storage; +using ArmoniK.Core.Utils; using Microsoft.Extensions.Diagnostics.HealthChecks; using Microsoft.Extensions.Logging; @@ -59,11 +60,21 @@ public PartitionTable(SessionProvider activitySource_ = activitySource; } - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + /// + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(PartitionTable)} is not initialized", + sessionProvider_, + partitionCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } + /// public async Task Init(CancellationToken cancellationToken) { if (!isInitialized_) diff --git a/Adaptors/MongoDB/src/ResultTable.cs b/Adaptors/MongoDB/src/ResultTable.cs index 98dfa4187..3a951b602 100644 --- a/Adaptors/MongoDB/src/ResultTable.cs +++ b/Adaptors/MongoDB/src/ResultTable.cs @@ -30,6 +30,7 @@ using ArmoniK.Core.Base.Exceptions; using ArmoniK.Core.Common.Exceptions; using ArmoniK.Core.Common.Storage; +using ArmoniK.Core.Utils; using ArmoniK.Utils; using Microsoft.Extensions.Diagnostics.HealthChecks; @@ -391,8 +392,16 @@ await resultCollectionProvider_.Init(cancellationToken) public ILogger Logger { get; } /// - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(ResultTable)} is not initialized", + sessionProvider_, + resultCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } } diff --git a/Adaptors/MongoDB/src/ResultWatcher.cs b/Adaptors/MongoDB/src/ResultWatcher.cs index ac974ce88..aa555e133 100644 --- a/Adaptors/MongoDB/src/ResultWatcher.cs +++ b/Adaptors/MongoDB/src/ResultWatcher.cs @@ -27,6 +27,7 @@ using ArmoniK.Core.Base.DataStructures; using ArmoniK.Core.Common.Storage; using ArmoniK.Core.Common.Storage.Events; +using ArmoniK.Core.Utils; using Microsoft.Extensions.Diagnostics.HealthChecks; using Microsoft.Extensions.Logging; @@ -65,10 +66,18 @@ public ResultWatcher(SessionProvider ses } /// - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(ResultWatcher)} is not initialized", + sessionProvider_, + resultCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } /// public async Task Init(CancellationToken cancellationToken) diff --git a/Adaptors/MongoDB/src/SessionTable.cs b/Adaptors/MongoDB/src/SessionTable.cs index 413669db7..c6b03a58f 100644 --- a/Adaptors/MongoDB/src/SessionTable.cs +++ b/Adaptors/MongoDB/src/SessionTable.cs @@ -27,6 +27,7 @@ using ArmoniK.Core.Adapters.MongoDB.Table.DataModel; using ArmoniK.Core.Base.DataStructures; using ArmoniK.Core.Common.Storage; +using ArmoniK.Core.Utils; using ArmoniK.Utils; using JetBrains.Annotations; @@ -179,10 +180,18 @@ await sessionCollectionProvider_.Init(cancellationToken) } /// - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(SessionTable)} is not initialized", + sessionProvider_, + sessionCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } /// public async Task UpdateOneSessionAsync(string sessionId, diff --git a/Adaptors/MongoDB/src/TaskTable.cs b/Adaptors/MongoDB/src/TaskTable.cs index 1a368f19d..3d554efbf 100644 --- a/Adaptors/MongoDB/src/TaskTable.cs +++ b/Adaptors/MongoDB/src/TaskTable.cs @@ -30,6 +30,7 @@ using ArmoniK.Core.Base.Exceptions; using ArmoniK.Core.Common.Exceptions; using ArmoniK.Core.Common.Storage; +using ArmoniK.Core.Utils; using ArmoniK.Utils; using Microsoft.Extensions.Diagnostics.HealthChecks; @@ -435,10 +436,18 @@ await taskCollection.UpdateManyAsync(data => taskIds.Contains(data.TaskId), public ILogger Logger { get; } /// - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(TaskTable)} is not initialized", + sessionProvider_, + taskCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } /// public async Task Init(CancellationToken cancellationToken) diff --git a/Adaptors/MongoDB/src/TaskWatcher.cs b/Adaptors/MongoDB/src/TaskWatcher.cs index d5686c1bb..b96939799 100644 --- a/Adaptors/MongoDB/src/TaskWatcher.cs +++ b/Adaptors/MongoDB/src/TaskWatcher.cs @@ -27,6 +27,7 @@ using ArmoniK.Core.Base.DataStructures; using ArmoniK.Core.Common.Storage; using ArmoniK.Core.Common.Storage.Events; +using ArmoniK.Core.Utils; using Microsoft.Extensions.Diagnostics.HealthChecks; using Microsoft.Extensions.Logging; @@ -65,10 +66,18 @@ public TaskWatcher(SessionProvider sessi } /// - public Task Check(HealthCheckTag tag) - => Task.FromResult(isInitialized_ - ? HealthCheckResult.Healthy() - : HealthCheckResult.Unhealthy()); + public async Task Check(HealthCheckTag tag) + { + var result = await HealthCheckResultCombiner.Combine(tag, + $"{nameof(TaskWatcher)} is not initialized", + sessionProvider_, + taskCollectionProvider_) + .ConfigureAwait(false); + + return isInitialized_ && result.Status == HealthStatus.Healthy + ? HealthCheckResult.Healthy() + : HealthCheckResult.Unhealthy(result.Description); + } /// public async Task Init(CancellationToken cancellationToken) diff --git a/Common/src/Pollster/Pollster.cs b/Common/src/Pollster/Pollster.cs index 348ddd778..8522cc15b 100644 --- a/Common/src/Pollster/Pollster.cs +++ b/Common/src/Pollster/Pollster.cs @@ -22,7 +22,6 @@ using System.Diagnostics.CodeAnalysis; using System.Diagnostics.Metrics; using System.Linq; -using System.Text; using System.Threading; using System.Threading.Tasks; @@ -36,6 +35,7 @@ using ArmoniK.Core.Common.Storage; using ArmoniK.Core.Common.Stream.Worker; using ArmoniK.Core.Common.Utils; +using ArmoniK.Core.Utils; using ArmoniK.Utils; using Grpc.Core; @@ -180,56 +180,19 @@ public async Task Check(HealthCheckTag tag) return HealthCheckResult.Unhealthy("End of main loop reached, no more tasks will be executed."); } - var checks = new List> - { - pullQueueStorage_.Check(tag), - dataPrefetcher_.Check(tag), - workerStreamHandler_.Check(tag), - objectStorage_.Check(tag), - resultTable_.Check(tag), - sessionTable_.Check(tag), - taskTable_.Check(tag), - }; - - var exceptions = new List(); - var data = new Dictionary(); - var description = new StringBuilder(); - var worstStatus = HealthStatus.Healthy; - - foreach (var healthCheckResult in await checks.WhenAll() - .ConfigureAwait(false)) - { - if (healthCheckResult.Status == HealthStatus.Healthy) - { - continue; - } - - if (healthCheckResult.Exception is not null) - { - exceptions.Add(healthCheckResult.Exception); - } - - foreach (var (key, value) in healthCheckResult.Data) - { - data[key] = value; - } - - if (healthCheckResult.Description is not null) - { - description.AppendLine(healthCheckResult.Description); - } - - worstStatus = worstStatus < healthCheckResult.Status - ? worstStatus - : healthCheckResult.Status; - } - - var result = new HealthCheckResult(worstStatus, - description.ToString(), - new AggregateException(exceptions), - data); - - if (worstStatus == HealthStatus.Unhealthy && tag == HealthCheckTag.Liveness) + // no need for description because this check is registered as the agent health check and it will add proper metadata. + var result = await HealthCheckResultCombiner.Combine(tag, + string.Empty, + pullQueueStorage_, + dataPrefetcher_, + workerStreamHandler_, + objectStorage_, + resultTable_, + sessionTable_, + taskTable_) + .ConfigureAwait(false); + + if (result.Status == HealthStatus.Unhealthy && tag == HealthCheckTag.Liveness) { healthCheckFailedResult_ = result; } diff --git a/Common/tests/Pollster/PollsterTest.cs b/Common/tests/Pollster/PollsterTest.cs index 539e970b2..84407c114 100644 --- a/Common/tests/Pollster/PollsterTest.cs +++ b/Common/tests/Pollster/PollsterTest.cs @@ -356,9 +356,8 @@ await testServiceProvider.Pollster.Init(CancellationToken.None) Console.WriteLine(res.Description); - Assert.AreEqual(new StringBuilder().AppendLine(desc) - .ToString(), - healthResult.Description); + Assert.AreEqual(desc, + healthResult.Description?.Trim()); Assert.AreEqual(new AggregateException(ex).Message, healthResult.Exception?.Message); Assert.AreEqual(HealthStatus.Unhealthy, diff --git a/Utils/src/HealthCheckResultCombiner.cs b/Utils/src/HealthCheckResultCombiner.cs new file mode 100644 index 000000000..cd4bbfb7d --- /dev/null +++ b/Utils/src/HealthCheckResultCombiner.cs @@ -0,0 +1,90 @@ +// This file is part of the ArmoniK project +// +// Copyright (C) ANEO, 2021-2025. All rights reserved. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published +// by the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY, without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text; +using System.Threading.Tasks; + +using ArmoniK.Core.Base; +using ArmoniK.Core.Base.DataStructures; +using ArmoniK.Utils; + +using Microsoft.Extensions.Diagnostics.HealthChecks; + +namespace ArmoniK.Core.Utils; + +/// +/// Helper class to combine the results of multiple +/// +public static class HealthCheckResultCombiner +{ + /// + /// Combine the results of multiple for a given tag and description + /// + /// The tag on which to combine the results + /// The description to add to the results + /// The sources from which to get the results + /// + /// The combined result + /// + public static async Task Combine(HealthCheckTag tag, + string desc, + params IHealthCheckProvider[] providers) + { + var exceptions = new List(); + var data = new Dictionary(); + var description = new StringBuilder(); + var worstStatus = HealthStatus.Healthy; + description.AppendLine(desc); + + foreach (var healthCheckResult in await providers.Select(p => p.Check(tag)) + .WhenAll() + .ConfigureAwait(false)) + { + if (healthCheckResult.Status == HealthStatus.Healthy) + { + continue; + } + + if (healthCheckResult.Exception is not null) + { + exceptions.Add(healthCheckResult.Exception); + } + + foreach (var (key, value) in healthCheckResult.Data) + { + data[key] = value; + } + + if (healthCheckResult.Description is not null) + { + description.AppendLine(healthCheckResult.Description); + } + + worstStatus = worstStatus < healthCheckResult.Status + ? worstStatus + : healthCheckResult.Status; + } + + return new HealthCheckResult(worstStatus, + description.ToString(), + new AggregateException(exceptions), + data); + } +}