Report database availability

src/Microsoft.Health.Fhir.Api/Features/Health/StorageInitializedHealthCheck.cs

@@ -9,10 +9,9 @@
 using EnsureThat;
 using MediatR;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
-using Microsoft.Health.Core;
+using Microsoft.Health.Fhir.Core.Data;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Messages.Search;
-using Microsoft.Health.Fhir.Core.Messages.Storage;
 
 namespace Microsoft.Health.Fhir.Api.Features.Health
 {
@@ -45,8 +44,9 @@ public Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, Canc
             }
 
             // Check if customer-managed key (CMK) is properly set.
-            var isCMKProperlySet = _databaseStatusReporter.IsCustomerManagerKeyProperlySetAsync(cancellationToken).GetAwaiter().GetResult();
-            if (!isCMKProperlySet)
+            var isCMKProperlySet = _databaseStatusReporter.IsCustomerManagedKeyProperlySetAsync(cancellationToken).GetAwaiter().GetResult();
+            DatabaseAvailability databaseAvailability = _databaseStatusReporter.GetDatabaseAvailability();
+            if (!isCMKProperlySet || databaseAvailability == DatabaseAvailability.DegradedByCustomerManagedKey)
             {
                 return Task.FromResult(new HealthCheckResult(HealthStatus.Degraded, DegradedCMKMessage));
             }

src/Microsoft.Health.Fhir.Core/Data/BaseDatabaseStatusReporter.cs

@@ -0,0 +1,37 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Health.Fhir.Core.Data
+{
+    public abstract class BaseDatabaseStatusReporter : IDatabaseStatusReporter
+    {
+        private long _databaseAvailability;
+
+        protected BaseDatabaseStatusReporter()
+        {
+            _databaseAvailability = (long)DatabaseAvailability.Available;
+        }
+
+        public abstract bool IsCustomerManagedKeyException(Exception exception);
+
+        public abstract Task<bool> IsCustomerManagedKeyProperlySetAsync(CancellationToken cancellationToken);
+
+#pragma warning disable CA1024 // Use properties where appropriate. Justification: Prefer to keep this as a method for consistency with the other methods in this interface.
+        public DatabaseAvailability GetDatabaseAvailability()
+        {
+            return (DatabaseAvailability)Interlocked.Read(ref _databaseAvailability);
+        }
+#pragma warning restore CA1024
+
+        public void ReportDatabaseAvailability(DatabaseAvailability availability)
+        {
+            Interlocked.Exchange(ref _databaseAvailability, (long)availability);
+        }
+    }
+}

src/Microsoft.Health.Fhir.Core/Data/DatabaseAvailability.cs

@@ -0,0 +1,13 @@
+// -------------------------------------------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+// -------------------------------------------------------------------------------------------------
+
+namespace Microsoft.Health.Fhir.Core.Data
+{
+    public enum DatabaseAvailability
+    {
+        Available,
+        DegradedByCustomerManagedKey,
+    }
+}

src/Microsoft.Health.Fhir.Api/Features/Health/IDatabaseStatusReporter.cs → src/Microsoft.Health.Fhir.Core/Data/IDatabaseStatusReporter.cs

@@ -3,10 +3,11 @@
 // Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
 // -------------------------------------------------------------------------------------------------
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace Microsoft.Health.Fhir.Api.Features.Health
+namespace Microsoft.Health.Fhir.Core.Data
 {
     /// <summary>
     /// Provides functionality to report the status of Data Store.
@@ -18,6 +19,12 @@ public interface IDatabaseStatusReporter
         /// </summary>
         /// <param name="cancellationToken">A token to cancel the operation.</param>
         /// <returns>A task that returns true if the key is healthy; otherwise, false.</returns>
-        Task<bool> IsCustomerManagerKeyProperlySetAsync(CancellationToken cancellationToken);
+        Task<bool> IsCustomerManagedKeyProperlySetAsync(CancellationToken cancellationToken);
+
+        bool IsCustomerManagedKeyException(Exception exception);
+
+        DatabaseAvailability GetDatabaseAvailability();
+
+        void ReportDatabaseAvailability(DatabaseAvailability availability);
     }
 }

src/Microsoft.Health.Fhir.CosmosDb.UnitTests/Features/Health/CosmosDbStatusReporterTests.cs

@@ -25,7 +25,7 @@ public async Task GivenHealthyCustomerKeyHealth_WhenIsCustomerManagerKeyProperly
             var reporter = new CosmosDbStatusReporter();
 
             // Act
-            bool result = await reporter.IsCustomerManagerKeyProperlySetAsync(CancellationToken.None);
+            bool result = await reporter.IsCustomerManagedKeyProperlySetAsync(CancellationToken.None);
 
             // Assert
             Assert.True(result);

src/Microsoft.Health.Fhir.CosmosDb/Features/Health/CosmosDbHealthCheck.cs

@@ -151,7 +151,7 @@ void LogWithDetails(LogLevel logLevel, Exception ex, string message, IEnumerable
                             [attempt, ex.GetType().Name]);
                     }
                 }
-                catch (CosmosException ex) when (ex.IsCmkClientError())
+                catch (CosmosException ex) when (ex.IsCustomerManagedKeyException())
                 {
                     // Handling CMK errors.
                     LogWithDetails(

src/Microsoft.Health.Fhir.CosmosDb/Features/Health/CosmosDbStatusReporter.cs

@@ -3,23 +3,33 @@
 // Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
 // -------------------------------------------------------------------------------------------------
 
+using System;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Health.Fhir.Api.Features.Health;
+using Microsoft.Azure.Cosmos;
+using Microsoft.Health.Fhir.Core.Data;
+using Microsoft.Health.Fhir.CosmosDb.Features.Storage;
 
 namespace Microsoft.Health.Fhir.CosmosDb.Features.Health
 {
     /// <summary>
-    /// Cosmos DB implementation of <see cref="IDatabaseStatusReporter"/>.
-    /// Always returns healthy status without performing any checks.
+    /// Cosmos DB implementation of <see cref="BaseDatabaseStatusReporter"/>.
     /// </summary>
-    public class CosmosDbStatusReporter : IDatabaseStatusReporter
+    public class CosmosDbStatusReporter : BaseDatabaseStatusReporter
     {
-        /// <inheritdoc />
-        public Task<bool> IsCustomerManagerKeyProperlySetAsync(CancellationToken cancellationToken = default)
+        public CosmosDbStatusReporter()
+            : base()
         {
-            // [WI] to implement: https://microsofthealth.visualstudio.com/Health/_workitems/edit/166817
-            return Task.FromResult(true);
+        }
+
+        public override bool IsCustomerManagedKeyException(Exception exception)
+        {
+            return exception is CosmosException cdbe && cdbe.IsCustomerManagedKeyException();
+        }
+
+        public override Task<bool> IsCustomerManagedKeyProperlySetAsync(CancellationToken cancellationToken = default)
+        {
+            return Task.FromResult(GetDatabaseAvailability() != DatabaseAvailability.DegradedByCustomerManagedKey);
         }
     }
 }

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosContainerProvider.cs

@@ -17,6 +17,7 @@
 using Microsoft.Health.Abstractions.Exceptions;
 using Microsoft.Health.Core;
 using Microsoft.Health.Extensions.DependencyInjection;
+using Microsoft.Health.Fhir.Core.Data;
 using Microsoft.Health.Fhir.Core.Messages.Storage;
 using Microsoft.Health.Fhir.CosmosDb.Core.Configs;
 using Microsoft.Health.Fhir.CosmosDb.Core.Features.Storage;
@@ -33,6 +34,7 @@ public class CosmosContainerProvider : IHostedService, IRequireInitializationOnF
         private const int CollectionSettingsVersion = 3;
         private readonly ILogger<CosmosContainerProvider> _logger;
         private readonly IMediator _mediator;
+        private readonly IDatabaseStatusReporter _databaseStatusReporter;
         private readonly ICosmosDbDistributedLockFactory _distributedLockFactory;
         private Lazy<Container> _container;
         private readonly RetryableInitializationOperation _initializationOperation;
@@ -50,6 +52,7 @@ public CosmosContainerProvider(
             IMediator mediator,
             IEnumerable<ICollectionInitializer> collectionInitializers,
             ICosmosDbDistributedLockFactory distributedLockFactory,
+            IDatabaseStatusReporter databaseStatusReporter,
             ICosmosClientTestProvider cosmosClientTestProvider)
         {
             EnsureArg.IsNotNull(cosmosDataStoreConfiguration, nameof(cosmosDataStoreConfiguration));
@@ -61,11 +64,13 @@ public CosmosContainerProvider(
             EnsureArg.IsNotNull(collectionDataUpdater, nameof(collectionDataUpdater));
             EnsureArg.IsNotNull(retryPolicyFactory, nameof(retryPolicyFactory));
             EnsureArg.IsNotNull(mediator, nameof(mediator));
+            EnsureArg.IsNotNull(databaseStatusReporter, nameof(databaseStatusReporter));
             EnsureArg.IsNotNull(distributedLockFactory, nameof(distributedLockFactory));
 
             _logger = logger;
             _mediator = mediator;
             _distributedLockFactory = distributedLockFactory;
+            _databaseStatusReporter = databaseStatusReporter;
 
             string collectionId = collectionConfiguration.Get(Constants.CollectionConfigurationName).CollectionId;
             _client = cosmosClientInitializer.CreateCosmosClient(cosmosDataStoreConfiguration);
@@ -163,6 +168,12 @@ private async Task InitializeDataStoreAsync(
             {
                 LogLevel logLevel = LogLevel.Critical;
                 _logger.Log(logLevel, ex, "Cosmos DB Database {DatabaseId} Initialization has failed.", cosmosDataStoreConfiguration.DatabaseId);
+
+                if (_databaseStatusReporter.IsCustomerManagedKeyException(ex))
+                {
+                    _databaseStatusReporter.ReportDatabaseAvailability(DatabaseAvailability.DegradedByCustomerManagedKey);
+                }
+
                 throw;
             }
         }

src/Microsoft.Health.Fhir.CosmosDb/Features/Storage/CosmosExceptionExtensions.cs

@@ -61,7 +61,7 @@ public static class CosmosExceptionExtensions
         /// </summary>
         /// <param name="exception">The exception object</param>
         /// <returns>True iff the error is due to client CMK setting.</returns>
-        public static bool IsCmkClientError(this CosmosException exception)
+        public static bool IsCustomerManagedKeyException(this CosmosException exception)
         {
             return exception.StatusCode == HttpStatusCode.Forbidden && Enum.IsDefined(typeof(KnownCosmosDbCmkSubStatusValueClientIssue), exception.SubStatusCode);
         }

src/Microsoft.Health.Fhir.CosmosDb/Registration/FhirServerBuilderCosmosDbRegistrationExtensions.cs

@@ -14,7 +14,6 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using Microsoft.Health.Extensions.DependencyInjection;
-using Microsoft.Health.Fhir.Api.Features.Health;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Features.Operations;
 using Microsoft.Health.Fhir.Core.Features.Operations.Export;

src/Microsoft.Health.Fhir.R4.Core.UnitTests/Features/Validation/ServerProvideProfileValidationTests.cs

@@ -15,6 +15,7 @@
 using Microsoft.Extensions.Options;
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Core.Configs;
+using Microsoft.Health.Fhir.Core.Data;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Features.Persistence;
 using Microsoft.Health.Fhir.Core.Features.Search;
@@ -61,6 +62,7 @@ public ServerProvideProfileValidationTests()
                 _options,
                 _mediator,
                 _hostApplicationLifetime,
+                Substitute.For<IDatabaseStatusReporter>(),
                 NullLogger<ServerProvideProfileValidation>.Instance);
         }
 

src/Microsoft.Health.Fhir.Shared.Api.UnitTests/Features/Health/StorageInitializedHealthCheckTests.cs

@@ -7,9 +7,9 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Health.Fhir.Core.Data;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Messages.Search;
-using Microsoft.Health.Fhir.Core.Messages.Storage;
 using Microsoft.Health.Fhir.Tests.Common;
 using Microsoft.Health.Test.Utilities;
 using NSubstitute;
@@ -27,7 +27,8 @@ public class StorageInitializedHealthCheckTests
     public StorageInitializedHealthCheckTests()
     {
         _databaseStatusReporter = Substitute.For<IDatabaseStatusReporter>();
-        _databaseStatusReporter.IsCustomerManagerKeyProperlySetAsync(Arg.Any<CancellationToken>()).Returns(Task.FromResult(true));
+        _databaseStatusReporter.IsCustomerManagedKeyProperlySetAsync(Arg.Any<CancellationToken>()).Returns(Task.FromResult(true));
+        _databaseStatusReporter.GetDatabaseAvailability().Returns(DatabaseAvailability.Available);
         _sut = new StorageInitializedHealthCheck(_databaseStatusReporter);
     }
 
@@ -53,6 +54,10 @@ public async Task GivenStorageInitializedHealthCheck_WhenCheckHealthAsync_ThenCh
     {
         using (Mock.Property(() => ClockResolver.TimeProvider, new Microsoft.Extensions.Time.Testing.FakeTimeProvider(DateTimeOffset.Now.AddMinutes(5).AddSeconds(1))))
         {
+            // Arrange
+            _databaseStatusReporter.IsCustomerManagedKeyProperlySetAsync(Arg.Any<CancellationToken>()).Returns(Task.FromResult(true));
+            _databaseStatusReporter.GetDatabaseAvailability().Returns(DatabaseAvailability.Available);
+
             HealthCheckResult result = await _sut.CheckHealthAsync(new HealthCheckContext(), CancellationToken.None);
 
             Assert.Equal(HealthStatus.Unhealthy, result.Status);
@@ -66,11 +71,30 @@ public async Task GivenUnhealthyCMK_WhenCheckHealthAsyncAfter5Minutes_ThenChange
         using (Mock.Property(() => ClockResolver.TimeProvider, new Microsoft.Extensions.Time.Testing.FakeTimeProvider(DateTimeOffset.Now.AddMinutes(5).AddSeconds(1))))
         {
             // Arrange
-            _databaseStatusReporter.IsCustomerManagerKeyProperlySetAsync(Arg.Any<CancellationToken>()).Returns(Task.FromResult(false));
+            _databaseStatusReporter.IsCustomerManagedKeyProperlySetAsync(Arg.Any<CancellationToken>()).Returns(Task.FromResult(false));
+            _databaseStatusReporter.GetDatabaseAvailability().Returns(DatabaseAvailability.Available);
 
             HealthCheckResult result = await _sut.CheckHealthAsync(new HealthCheckContext(), CancellationToken.None);
             Assert.Equal(HealthStatus.Degraded, result.Status);
             Assert.Contains("The health of the store has degraded. Customer-managed key is not properly set.", result.Description);
         }
     }
+
+    [Fact]
+    public async Task GivenDatabaseNotAvailable_WhenCheckHealthAsync_ThenReturnsDegraded()
+    {
+        using (Mock.Property(() => ClockResolver.TimeProvider, new Microsoft.Extensions.Time.Testing.FakeTimeProvider(DateTimeOffset.Now.AddMinutes(5).AddSeconds(1))))
+        {
+            // Arrange
+            _databaseStatusReporter.IsCustomerManagedKeyProperlySetAsync(Arg.Any<CancellationToken>()).Returns(Task.FromResult(true));
+            _databaseStatusReporter.GetDatabaseAvailability().Returns(DatabaseAvailability.DegradedByCustomerManagedKey);
+
+            // Act
+            HealthCheckResult result = await _sut.CheckHealthAsync(new HealthCheckContext(), CancellationToken.None);
+
+            // Assert
+            Assert.Equal(HealthStatus.Degraded, result.Status);
+            Assert.Contains("The health of the store has degraded. Customer-managed key is not properly set.", result.Description);
+        }
+    }
 }

src/Microsoft.Health.Fhir.Shared.Core.UnitTests/Features/Validation/ServerProvideProfileValidationTests.cs

@@ -15,6 +15,7 @@
 using Microsoft.Extensions.Options;
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Core.Configs;
+using Microsoft.Health.Fhir.Core.Data;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Features.Persistence;
 using Microsoft.Health.Fhir.Core.Features.Search;
@@ -63,6 +64,7 @@ public ServerProvideProfileValidationTests()
                 _options,
                 _mediator,
                 _hostApplicationLifetime,
+                Substitute.For<IDatabaseStatusReporter>(),
                 NullLogger<ServerProvideProfileValidation>.Instance);
         }
 

src/Microsoft.Health.Fhir.Shared.Core/Features/Validation/ServerProvideProfileValidation.cs

@@ -24,6 +24,8 @@
 using Microsoft.Health.Core.Extensions;
 using Microsoft.Health.Extensions.DependencyInjection;
 using Microsoft.Health.Fhir.Core.Configs;
+using Microsoft.Health.Fhir.Core.Data;
+using Microsoft.Health.Fhir.Core.Exceptions;
 using Microsoft.Health.Fhir.Core.Extensions;
 using Microsoft.Health.Fhir.Core.Features.Search;
 using Microsoft.Health.Fhir.Core.Features.Storage;
@@ -48,6 +50,7 @@ public sealed class ServerProvideProfileValidation : IProvideProfilesForValidati
         private readonly ValidateOperationConfiguration _validateOperationConfig;
         private readonly FhirMemoryCache<Resource> _resourcesByUri;
         private readonly IMediator _mediator;
+        private readonly IDatabaseStatusReporter _databaseStatusReporter;
 
         private bool _isDisposed = false;
         private string _mostRecentProfileHash = null;
@@ -61,18 +64,21 @@ public ServerProvideProfileValidation(
             IOptions<ValidateOperationConfiguration> options,
             IMediator mediator,
             IHostApplicationLifetime hostApplicationLifetime,
+            IDatabaseStatusReporter databaseStatusReporter,
             ILogger<ServerProvideProfileValidation> logger)
         {
             EnsureArg.IsNotNull(searchServiceFactory, nameof(searchServiceFactory));
             EnsureArg.IsNotNull(options?.Value, nameof(options));
             EnsureArg.IsNotNull(mediator, nameof(mediator));
             EnsureArg.IsNotNull(logger, nameof(logger));
             EnsureArg.IsNotNull(hostApplicationLifetime, nameof(hostApplicationLifetime));
+            EnsureArg.IsNotNull(databaseStatusReporter, nameof(databaseStatusReporter));
 
             _searchServiceFactory = searchServiceFactory;
             _expirationTime = DateTime.UtcNow;
             _validateOperationConfig = options.Value;
             _mediator = mediator;
+            _databaseStatusReporter = databaseStatusReporter;
             _logger = logger;
 
             _resourcesByUri = new FhirMemoryCache<Resource>(
@@ -346,7 +352,15 @@ private async Task BackgroundLoop(CancellationToken cancellationToken)
                 }
                 catch (Exception ex)
                 {
-                    _logger.LogError(ex, "Profiles: Background profile status task failed. Elapsed time: {ElapsedTime}ms", stopwatch.ElapsedMilliseconds);
+                    if (_databaseStatusReporter.IsCustomerManagedKeyException(ex))
+                    {
+                        _databaseStatusReporter.ReportDatabaseAvailability(DatabaseAvailability.DegradedByCustomerManagedKey);
+                        _logger.LogWarning(ex, "Profiles: Background profile status task failed due to customer managed key issue. Database availability set as degraded. Elapsed time: {ElapsedTime}ms", stopwatch.ElapsedMilliseconds);
+                    }
+                    else
+                    {
+                        _logger.LogError(ex, "Profiles: Background profile status task failed. Elapsed time: {ElapsedTime}ms", stopwatch.ElapsedMilliseconds);
+                    }
                 }
 
                 await Task.Delay(TimeSpan.FromSeconds(_validateOperationConfig.BackgroundProfileStatusCheckIntervalInSeconds), cancellationToken);