Re-design Azure IoTHub health checks to allow for managed identity an…

…d other features (#2216)

Co-authored-by: Adam Sitnik <[email protected]>

src/HealthChecks.Azure.IoTHub/DependencyInjection/IoTHubHealthChecksBuilderExtensions.cs

@@ -1,20 +1,27 @@
 using HealthChecks.Azure.IoTHub;
+using Microsoft.Azure.Devices;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace Microsoft.Extensions.DependencyInjection;
 
 /// <summary>
-/// Extension methods to configure <see cref="IoTHubHealthCheck"/>.
+/// Extension methods to configure <see cref="IoTHubRegistryManagerHealthCheck"/>.
 /// </summary>
 public static class IoTHubHealthChecksBuilderExtensions
 {
-    private const string NAME = "iothub";
+    private const string NAME_REGISTRY_MANAGER_READ = "iothub_registrymanager_read";
+    private const string NAME_REGISTRY_MANAGER_WRITE = "iothub_registrymanager_write";
+    private const string NAME_SERVICE_CLIENT = "iothub_serviceclient";
 
     /// <summary>
-    /// Add a health check for Azure IoT Hub.
+    /// Adds a read health check for Azure IoT Hub registry manager.
     /// </summary>
     /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
-    /// <param name="optionsFactory">A action to configure the Azure IoT Hub connection to use.</param>
+    /// <param name="registryManagerFactory">
+    /// An optional factory to obtain <see cref="RegistryManager" /> instance.
+    /// When not provided, <see cref="RegistryManager" /> is simply resolved from <see cref="IServiceProvider"/>.
+    /// </param>
+    /// <param name="query">The query to perform.</param>
     /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'iothub' will be used for the name.</param>
     /// <param name="failureStatus">
     /// The <see cref="HealthStatus"/> that should be reported when the health check fails. Optional. If <c>null</c> then
@@ -23,22 +30,92 @@ public static class IoTHubHealthChecksBuilderExtensions
     /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
     /// <param name="timeout">An optional <see cref="TimeSpan"/> representing the timeout of the check.</param>
     /// <returns>The specified <paramref name="builder"/>.</returns>
-    public static IHealthChecksBuilder AddAzureIoTHub(
+    public static IHealthChecksBuilder AddAzureIoTHubRegistryReadCheck(
         this IHealthChecksBuilder builder,
-        Action<IoTHubOptions>? optionsFactory,
-        string? name = default,
+        Func<IServiceProvider, RegistryManager>? registryManagerFactory = default,
+        string query = "SELECT deviceId FROM devices",
+        string? name = NAME_REGISTRY_MANAGER_READ,
         HealthStatus? failureStatus = default,
         IEnumerable<string>? tags = default,
         TimeSpan? timeout = default)
     {
-        var options = new IoTHubOptions();
-        optionsFactory?.Invoke(options);
+        Guard.ThrowIfNull(query);
+
+        return builder.Add(new HealthCheckRegistration(
+           name ?? NAME_REGISTRY_MANAGER_READ,
+           sp => new IoTHubRegistryManagerHealthCheck(
+                registryManager: registryManagerFactory?.Invoke(sp) ?? sp.GetRequiredService<RegistryManager>(),
+                readQuery: query),
+           failureStatus,
+           tags,
+           timeout));
+    }
 
-        var registrationName = name ?? NAME;
+    /// <summary>
+    /// Add a write health check for Azure IoT Hub registry manager.
+    /// </summary>
+    /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
+    /// <param name="registryManagerFactory">
+    /// An optional factory to obtain <see cref="RegistryManager" /> instance.
+    /// When not provided, <see cref="RegistryManager" /> is simply resolved from <see cref="IServiceProvider"/>.
+    /// </param>
+    /// <param name="deviceId">The id of the device to add and remove.</param>
+    /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'iothub' will be used for the name.</param>
+    /// <param name="failureStatus">
+    /// The <see cref="HealthStatus"/> that should be reported when the health check fails. Optional. If <c>null</c> then
+    /// the default status of <see cref="HealthStatus.Unhealthy"/> will be reported.
+    /// </param>
+    /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
+    /// <param name="timeout">An optional <see cref="TimeSpan"/> representing the timeout of the check.</param>
+    /// <returns>The specified <paramref name="builder"/>.</returns>
+    public static IHealthChecksBuilder AddAzureIoTHubRegistryWriteCheck(
+        this IHealthChecksBuilder builder,
+        Func<IServiceProvider, RegistryManager>? registryManagerFactory = default,
+        string deviceId = "health-check-registry-write-device-id",
+        string? name = NAME_REGISTRY_MANAGER_WRITE,
+        HealthStatus? failureStatus = default,
+        IEnumerable<string>? tags = default,
+        TimeSpan? timeout = default)
+    {
+        Guard.ThrowIfNull(deviceId);
 
         return builder.Add(new HealthCheckRegistration(
-           registrationName,
-           sp => new IoTHubHealthCheck(options),
+           name ?? NAME_REGISTRY_MANAGER_WRITE,
+           sp => new IoTHubRegistryManagerHealthCheck(
+                registryManager: registryManagerFactory?.Invoke(sp) ?? sp.GetRequiredService<RegistryManager>(),
+                writeDeviceId: deviceId),
+           failureStatus,
+           tags,
+           timeout));
+    }
+
+    /// <summary>
+    /// Add a health check for Azure IoT Hub service client.
+    /// </summary>
+    /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
+    /// <param name="serviceClientFactory">
+    /// An optional factory to obtain <see cref="ServiceClient" /> instance.
+    /// When not provided, <see cref="ServiceClient" /> is simply resolved from <see cref="IServiceProvider"/>.
+    /// </param>
+    /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'iothub' will be used for the name.</param>
+    /// <param name="failureStatus">
+    /// The <see cref="HealthStatus"/> that should be reported when the health check fails. Optional. If <c>null</c> then
+    /// the default status of <see cref="HealthStatus.Unhealthy"/> will be reported.
+    /// </param>
+    /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
+    /// <param name="timeout">An optional <see cref="TimeSpan"/> representing the timeout of the check.</param>
+    /// <returns>The specified <paramref name="builder"/>.</returns>
+    public static IHealthChecksBuilder AddAzureIoTHubServiceClient(
+        this IHealthChecksBuilder builder,
+        Func<IServiceProvider, ServiceClient>? serviceClientFactory = default,
+        string? name = NAME_SERVICE_CLIENT,
+        HealthStatus? failureStatus = default,
+        IEnumerable<string>? tags = default,
+        TimeSpan? timeout = default)
+    {
+        return builder.Add(new HealthCheckRegistration(
+           name ?? NAME_SERVICE_CLIENT,
+           sp => new IoTHubServiceClientHealthCheck(serviceClient: serviceClientFactory?.Invoke(sp) ?? sp.GetRequiredService<ServiceClient>()),
            failureStatus,
            tags,
            timeout));

src/HealthChecks.Azure.IoTHub/IoTHubRegistryManagerHealthCheck.cs

@@ -0,0 +1,71 @@
+using Microsoft.Azure.Devices;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace HealthChecks.Azure.IoTHub;
+
+public sealed class IoTHubRegistryManagerHealthCheck : IHealthCheck
+{
+    private readonly RegistryManager _registryManager;
+    private readonly string? _readQuery;
+    private readonly string? _writeDeviceId;
+
+    public IoTHubRegistryManagerHealthCheck(RegistryManager registryManager, string? readQuery = default, string? writeDeviceId = default)
+    {
+        _registryManager = Guard.ThrowIfNull(registryManager);
+
+        if (string.IsNullOrEmpty(readQuery) && string.IsNullOrEmpty(writeDeviceId))
+        {
+            throw new ArgumentException("Either readQuery or writeDeviceId has to be provided");
+        }
+
+        _readQuery = readQuery;
+        _writeDeviceId = writeDeviceId;
+    }
+
+    /// <inheritdoc />
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            if (!string.IsNullOrEmpty(_writeDeviceId))
+            {
+                await ExecuteRegistryWriteCheckAsync(cancellationToken).ConfigureAwait(false);
+            }
+            else
+            {
+                await ExecuteRegistryReadCheckAsync().ConfigureAwait(false);
+            }
+
+            return HealthCheckResult.Healthy();
+        }
+        catch (Exception ex)
+        {
+            return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
+        }
+    }
+
+    private async Task ExecuteRegistryReadCheckAsync()
+    {
+        var query = _registryManager.CreateQuery(_readQuery!, 1);
+        await query.GetNextAsJsonAsync().ConfigureAwait(false);
+    }
+
+    private async Task ExecuteRegistryWriteCheckAsync(CancellationToken cancellationToken)
+    {
+        string deviceId = _writeDeviceId!;
+        var device = await _registryManager.GetDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
+
+        // in default implementation of configuration deviceId equals "health-check-registry-write-device-id"
+        // if in previous health check device were not removed -- try remove it
+        // if in previous health check device were added and removed -- try create and remove it
+        if (device != null)
+        {
+            await _registryManager.RemoveDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
+        }
+        else
+        {
+            await _registryManager.AddDeviceAsync(new Device(deviceId), cancellationToken).ConfigureAwait(false);
+            await _registryManager.RemoveDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
+        }
+    }
+}

src/HealthChecks.Azure.IoTHub/IoTHubServiceClientHealthCheck.cs

@@ -0,0 +1,29 @@
+using Microsoft.Azure.Devices;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace HealthChecks.Azure.IoTHub;
+
+public sealed class IoTHubServiceClientHealthCheck : IHealthCheck
+{
+    private readonly ServiceClient _serviceClient;
+
+    public IoTHubServiceClientHealthCheck(ServiceClient serviceClient)
+    {
+        _serviceClient = Guard.ThrowIfNull(serviceClient);
+    }
+
+    /// <inheritdoc />
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            await _serviceClient.GetServiceStatisticsAsync(cancellationToken).ConfigureAwait(false);
+
+            return HealthCheckResult.Healthy();
+        }
+        catch (Exception ex)
+        {
+            return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
+        }
+    }
+}