integrate health-check

Author: mgravell

src/StackExchange.Redis/ConfigurationOptions.cs

@@ -761,6 +761,7 @@ public string TieBreaker
         public int WriteBuffer
         {
             get => 0;
+            // ReSharper disable once ValueParameterNotUsed
             set { }
         }
 

src/StackExchange.Redis/FrameworkShims.cs

@@ -28,7 +28,6 @@ internal sealed class OverloadResolutionPriorityAttribute(int priority) : Attrib
 #endif
 
 #if !NET
-
 namespace System.Text
 {
     internal static class EncodingExtensions

src/StackExchange.Redis/HealthCheck.ConnectedProbe.cs

@@ -0,0 +1,24 @@
+using System.Net;
+using System.Threading.Tasks;
+
+namespace StackExchange.Redis;
+
+public sealed partial class HealthCheck
+{
+    public partial class HealthCheckProbe
+    {
+        /// <summary>
+        /// Report health using the <see cref="IServer.IsConnected"/> property, without any additional tests.
+        /// </summary>
+        public static HealthCheckProbe IsConnected => ConnectedProbe.Instance;
+    }
+
+    private sealed class ConnectedProbe : HealthCheckProbe
+    {
+        public static ConnectedProbe Instance { get; } = new();
+        private ConnectedProbe() { }
+
+        public override Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IServer server)
+            => server.IsConnected ? HealthyTask : UnhealthyTask;
+    }
+}

src/StackExchange.Redis/HealthCheck.Execute.cs

@@ -1,16 +1,140 @@
 using System;
+using System.Buffers;
 using System.Diagnostics;
 using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace StackExchange.Redis;
 
 public sealed partial class HealthCheck
 {
+    /// <summary>
+    /// Evaluate the health of the specified multiplexer, by evaluating all endpoints.
+    /// </summary>
+    public Task<HealthCheckResult> CheckHealthAsync(IConnectionMultiplexer multiplexer)
+        => multiplexer.IsConnected ? CheckHealthCoreAsync(multiplexer) : HealthCheckProbe.UnhealthyTask;
+
+    private async Task<HealthCheckResult> CheckHealthCoreAsync(IConnectionMultiplexer multiplexer)
+    {
+        try
+        {
+            Task<HealthCheckResult>[] pending;
+            if (multiplexer is IInternalConnectionMultiplexer internalMultiplexer)
+            {
+                var snapshot = internalMultiplexer.GetServerSnapshot();
+                pending = GetReusablePending(ref _reusablePending, snapshot.Length);
+                for (int i = 0; i < pending.Length; i++)
+                {
+                    pending[i] = CheckHealthAsync(snapshot[i].GetRedisServer(null));
+                }
+            }
+            else
+            {
+                var servers = multiplexer.GetServers();
+                pending = GetReusablePending(ref _reusablePending, servers.Length);
+                for (int i = 0; i < pending.Length; i++)
+                {
+                    pending[i] = CheckHealthAsync(servers[i]);
+                }
+            }
+            var result = await CollateAsync(pending, TotalTimeoutMillis()).ForAwait();
+
+            // on successful completion (regardless of outcome), we can reuse the pending array
+            PutReusablePending(ref _reusablePending, ref pending);
+            return result;
+        }
+        catch
+        {
+            // definitely unhappy
+            return HealthCheckResult.Unhealthy;
+        }
+    }
+
+    internal int TotalTimeoutMillis()
+    {
+        int count = ProbeCount;
+        if (count <= 0)
+        {
+            Debug.Fail("We shouldn't get as far as calculating timeouts with a non-positive probe count.");
+            return 0;
+        }
+
+        TimeSpan probeTimeout = ProbeTimeout, probeInterval = ProbeInterval;
+
+        // the first probe doesn't have an interval before it, the rest do
+        var totalTicks = probeTimeout.Ticks
+            + ((probeTimeout.Ticks + probeInterval.Ticks) * (count - 1));
+        var millis = (int)TimeSpan.FromTicks(totalTicks).TotalMilliseconds;
+        Debug.Assert(millis > 0, "Total timeout should be positive");
+        return millis;
+    }
+
+    // apply timeout and collation logic to a group of probes
+    internal static async Task<HealthCheckResult> CollateAsync(Task<HealthCheckResult>[] probes, int timeoutMilliseconds)
+    {
+        var pendingAll = Task.WhenAll(probes).ObserveErrors();
+        int success = 0, failure = 0;
+
+        if (await pendingAll.TimeoutAfter(timeoutMilliseconds).ForAwait())
+        {
+            // all completed inside timeout; all results should now be available
+            for (int i = 0; i < probes.Length; i++)
+            {
+                var individualResult = await probes[i].ForAwait();
+                switch (individualResult)
+                {
+                    case HealthCheckResult.Healthy: success++; break;
+                    case HealthCheckResult.Unhealthy: failure++; break;
+                }
+            }
+        }
+        else
+        {
+            // timeout
+            for (int i = 0; i < probes.Length; i++)
+            {
+                _ = probes[i].ObserveErrors();
+            }
+            throw new TimeoutException();
+        }
+
+        if (failure > 0) return HealthCheckResult.Unhealthy;
+        if (success > 0) return HealthCheckResult.Healthy;
+        return HealthCheckResult.Inconclusive;
+    }
+
+    private Task<HealthCheckResult>[]? _reusablePending;
+
+    // The number of pending tasks is determined by the number of endpoints, which doesn't change frequently
+    // (if at all); consequently, we can often re-use this buffer between health-checks, as long as we're careful.
+    internal static Task<HealthCheckResult>[] GetReusablePending(ref Task<HealthCheckResult>[]? field, int count)
+    {
+        var result = Interlocked.Exchange(ref field, null);
+        if (result is null || result.Length != count)
+        {
+            result = count == 0 ? [] : new Task<HealthCheckResult>[count];
+        }
+        return result;
+    }
+
+    internal static void PutReusablePending(ref Task<HealthCheckResult>[]? field, ref Task<HealthCheckResult>[] value)
+    {
+        if (value is { Length: > 0 })
+        {
+            Array.Clear(value, 0, value.Length);
+            Interlocked.Exchange(ref field, value);
+            value = [];
+        }
+    }
+
     /// <summary>
     /// Evaluate the health of an endpoint.
     /// </summary>
-    public async Task<HealthCheckResult> CheckHealthAsync(IConnectionMultiplexer multiplexer, EndPoint endpoint)
+    public Task<HealthCheckResult> CheckHealthAsync(IServer server)
+        => server.IsConnected ? CheckHealthCoreAsync(server) : HealthCheckProbe.UnhealthyTask;
+
+    private async Task<HealthCheckResult> CheckHealthCoreAsync(IServer server)
     {
         try
         {
@@ -20,7 +144,7 @@ public async Task<HealthCheckResult> CheckHealthAsync(IConnectionMultiplexer mul
                 HealthCheckResult probeResult;
                 try
                 {
-                    var pendingProbe = Probe.CheckHealthAsync(this, multiplexer, endpoint);
+                    var pendingProbe = Probe.CheckHealthAsync(this, server);
                     probeResult = await pendingProbe.TimeoutAfter(timeout).ForAwait()
                         ? await pendingProbe.ForAwait() // completed
                         : HealthCheckResult.Unhealthy; // timeout

src/StackExchange.Redis/HealthCheck.HealthCheckProbe.cs

@@ -14,14 +14,24 @@ public abstract partial class HealthCheckProbe
         /// <summary>
         /// Check the health of the specified endpoint.
         /// </summary>
-        public abstract Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IConnectionMultiplexer multiplexer, EndPoint endpoint);
+        public abstract Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IServer server);
 
-        private static Task<HealthCheckResult>? _inconclusive;
+        private static Task<HealthCheckResult>? _inconclusive, _healthy, _unhealthy;
 
         /// <summary>
-        /// Reports a probe that was skipped without being evaluated.
+        /// Reports a memoized probe that was skipped without being evaluated.
         /// </summary>
-        protected static Task<HealthCheckResult> Inconclusive => _inconclusive ??= Task.FromResult(HealthCheckResult.Inconclusive);
+        protected internal static Task<HealthCheckResult> InconclusiveTask => _inconclusive ??= Task.FromResult(HealthCheckResult.Inconclusive);
+
+        /// <summary>
+        /// Reports a memoized probe that was healthy.
+        /// </summary>
+        protected internal static Task<HealthCheckResult> HealthyTask => _healthy ??= Task.FromResult(HealthCheckResult.Healthy);
+
+        /// <summary>
+        /// Reports a memoized probe that was unhealthy.
+        /// </summary>
+        protected internal static Task<HealthCheckResult> UnhealthyTask => _unhealthy ??= Task.FromResult(HealthCheckResult.Unhealthy);
     }
 
     /// <summary>
@@ -30,15 +40,14 @@ public abstract partial class HealthCheckProbe
     public abstract class KeyWriteHealthCheckProbe : HealthCheckProbe
     {
         /// <inheritdoc/>
-        public override Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IConnectionMultiplexer multiplexer, EndPoint endpoint)
+        public override Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IServer server)
         {
-            var server = multiplexer.GetServer(endpoint);
-            if (server.IsReplica) return Inconclusive;
+            if (server.IsReplica) return InconclusiveTask;
 
-            RedisKey key = server.InventKey("health-check/"u8);
-            if (key.IsNull) return Inconclusive;
-            Debug.Assert(multiplexer.GetServer(key).EndPoint == endpoint, "Key was not routed to the correct endpoint");
-            return CheckHealthAsync(healthCheck, multiplexer.GetDatabase(), key);
+            RedisKey key = server.InventKey("health-check/");
+            if (key.IsNull) return InconclusiveTask;
+            Debug.Assert(server.Multiplexer.GetServer(key).EndPoint == server.EndPoint, "Key was not routed to the correct endpoint");
+            return CheckHealthAsync(healthCheck, server.Multiplexer.GetDatabase(), key);
         }
 
         /// <summary>

src/StackExchange.Redis/HealthCheck.PingProbe.cs

@@ -18,9 +18,8 @@ private sealed class PingProbe : HealthCheckProbe
         public static PingProbe Instance { get; } = new();
         private PingProbe() { }
 
-        public override async Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IConnectionMultiplexer multiplexer, EndPoint endpoint)
+        public override async Task<HealthCheckResult> CheckHealthAsync(HealthCheck healthCheck, IServer server)
         {
-            var server = multiplexer.GetServer(endpoint);
             await server.PingAsync();
             return HealthCheckResult.Healthy;
         }