🐛 Give the remotes health probe its own budget and make its timeout count

The probe used the remote's HttpClient, whose timeout is now the query time
limit, so a short limit failed the health check of a remote that was fine.
The probe's own 10 second budget never applied: its cancellation was caught
by a no-op branch, so a hanging remote passed the check and was never
disabled, and a shutdown waited out the full budget.

The probe now runs on its own budget regardless of the query time limit; a
remote that does not answer in time is disabled and fails the check, and a
shutdown aborts the probe without a verdict on the remote.

Author: ramonsmits

src/ServiceControl.UnitTests/ScatterGather/CheckRemotesTests.cs

@@ -0,0 +1,110 @@
+namespace ServiceControl.UnitTests.ScatterGather;
+
+using System;
+using System.Diagnostics;
+using System.Net;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using NUnit.Framework;
+using ServiceBus.Management.Infrastructure.Settings;
+using ServiceControl.Operations;
+
+/// <summary>
+/// The remotes health check probes each remote on its own budget: a remote that does not answer the probe is
+/// disabled until it answers again, and neither the query time limit nor a shutdown decides that.
+/// </summary>
+[TestFixture]
+class CheckRemotesTests
+{
+    const string RemoteAddress = "http://audit-1/api";
+
+    [Test]
+    public async Task A_remote_that_does_not_answer_the_probe_in_time_is_disabled_and_fails_the_check()
+    {
+        var settings = Settings();
+        var factory = new FakeHttpClientFactory();
+        factory.Register(settings.RemoteInstances[0], Hanging);
+
+        var check = new CheckRemotes(settings, factory, probeTimeout: TimeSpan.FromMilliseconds(200));
+
+        var result = await check.PerformCheck();
+
+        Assert.That(result.HasFailed, Is.True, "a hanging remote must fail the check");
+        Assert.That(result.FailureReason, Does.Contain("did not respond"));
+        Assert.That(settings.RemoteInstances[0].TemporarilyUnavailable, Is.True, "a hanging remote must be skipped by the queries until it answers again");
+    }
+
+    [Test]
+    public async Task The_probe_has_its_own_budget_and_is_not_bound_by_the_query_time_limit()
+    {
+        var settings = Settings();
+        var factory = new FakeHttpClientFactory();
+        factory.Register(settings.RemoteInstances[0], AnsweringAfter(TimeSpan.FromMilliseconds(300)), queryTimeLimit: TimeSpan.FromMilliseconds(50));
+
+        var check = new CheckRemotes(settings, factory, probeTimeout: TimeSpan.FromSeconds(5));
+
+        var result = await check.PerformCheck();
+
+        Assert.That(result.HasFailed, Is.False, "a short query time limit is not a verdict on the remote's health");
+        Assert.That(settings.RemoteInstances[0].TemporarilyUnavailable, Is.False);
+    }
+
+    [Test]
+    public async Task A_probe_aborted_by_shutdown_does_not_disable_the_remote()
+    {
+        var settings = Settings();
+        var factory = new FakeHttpClientFactory();
+        factory.Register(settings.RemoteInstances[0], Hanging);
+
+        var check = new CheckRemotes(settings, factory, probeTimeout: TimeSpan.FromSeconds(10));
+        using var shutdown = new CancellationTokenSource(TimeSpan.FromMilliseconds(100));
+        var stopwatch = Stopwatch.StartNew();
+
+        await check.PerformCheck(shutdown.Token);
+
+        Assert.That(stopwatch.Elapsed, Is.LessThan(TimeSpan.FromSeconds(5)), "shutdown must abort the probe instead of waiting out its budget");
+        Assert.That(settings.RemoteInstances[0].TemporarilyUnavailable, Is.False, "an aborted probe says nothing about the remote's health");
+    }
+
+    [Test]
+    public async Task A_remote_that_refuses_the_connection_is_disabled_and_fails_the_check()
+    {
+        var settings = Settings();
+        var factory = new FakeHttpClientFactory();
+        factory.Register(settings.RemoteInstances[0], (_, _) => throw new HttpRequestException("Connection refused"));
+
+        var result = await new CheckRemotes(settings, factory, probeTimeout: TimeSpan.FromSeconds(5)).PerformCheck();
+
+        Assert.That(result.HasFailed, Is.True);
+        Assert.That(settings.RemoteInstances[0].TemporarilyUnavailable, Is.True);
+    }
+
+    [Test]
+    public async Task A_remote_that_answers_the_probe_is_enabled_again()
+    {
+        var settings = Settings();
+        settings.RemoteInstances[0].TemporarilyUnavailable = true;
+        var factory = new FakeHttpClientFactory();
+        factory.Register(settings.RemoteInstances[0], AnsweringAfter(TimeSpan.Zero));
+
+        var result = await new CheckRemotes(settings, factory, probeTimeout: TimeSpan.FromSeconds(5)).PerformCheck();
+
+        Assert.That(result.HasFailed, Is.False);
+        Assert.That(settings.RemoteInstances[0].TemporarilyUnavailable, Is.False);
+    }
+
+    static Settings Settings() => new() { RemoteInstances = [new RemoteInstanceSetting(RemoteAddress)] };
+
+    static readonly Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> Hanging = async (_, token) =>
+    {
+        await Task.Delay(Timeout.Infinite, token);
+        return null;
+    };
+
+    static Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> AnsweringAfter(TimeSpan delay) => async (_, token) =>
+    {
+        await Task.Delay(delay, token);
+        return new HttpResponseMessage(HttpStatusCode.OK);
+    };
+}

src/ServiceControl.UnitTests/ScatterGather/FakeHttpClientFactory.cs

@@ -0,0 +1,33 @@
+#pragma warning disable PS0003 // Make the CancellationToken parameter optional — HttpMessageHandler.SendAsync override signature is fixed
+
+namespace ServiceControl.UnitTests.ScatterGather;
+
+using System;
+using System.Collections.Concurrent;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using ServiceBus.Management.Infrastructure.Settings;
+
+/// <summary>
+/// Hands out a client per remote instance whose answers are scripted, with the timeout the real
+/// remote clients get from <see cref="ServiceControl.Infrastructure.WebApi.RemoteInstanceServiceCollectionExtensions"/>: the query time limit.
+/// </summary>
+class FakeHttpClientFactory : IHttpClientFactory
+{
+    readonly ConcurrentDictionary<string, (HttpMessageHandler Handler, string BaseAddress, TimeSpan Timeout)> handlers = new();
+
+    public void Register(RemoteInstanceSetting remote, Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> responder, TimeSpan? queryTimeLimit = null) =>
+        handlers[remote.InstanceId] = (new StubHandler(responder), remote.BaseAddress, queryTimeLimit ?? TimeSpan.FromSeconds(100));
+
+    public HttpClient CreateClient(string name)
+    {
+        var (handler, baseAddress, timeout) = handlers[name];
+        return new HttpClient(handler, disposeHandler: false) { BaseAddress = new Uri(baseAddress), Timeout = timeout };
+    }
+
+    class StubHandler(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> responder) : HttpMessageHandler
+    {
+        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) => responder(request, cancellationToken);
+    }
+}

src/ServiceControl.UnitTests/ScatterGather/IncompleteResultsTests.cs

@@ -1,9 +1,6 @@
-#pragma warning disable PS0003 // Make the CancellationToken parameter optional — HttpMessageHandler.SendAsync override signature is fixed
-
 namespace ServiceControl.UnitTests.ScatterGather;
 
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net;
@@ -239,23 +236,4 @@ class TestApi(Settings settings, IHttpClientFactory factory, Func<CancellationTo
     {
         protected override Task<QueryResult<IList<MessagesView>>> LocalQuery(ScatterGatherApiMessageViewContext input, CancellationToken cancellationToken = default) => local(cancellationToken);
     }
-
-    class FakeHttpClientFactory : IHttpClientFactory
-    {
-        readonly ConcurrentDictionary<string, (HttpMessageHandler Handler, string BaseAddress)> handlers = new();
-
-        public void Register(RemoteInstanceSetting remote, Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> responder) =>
-            handlers[remote.InstanceId] = (new StubHandler(responder), remote.BaseAddress);
-
-        public HttpClient CreateClient(string name)
-        {
-            var (handler, baseAddress) = handlers[name];
-            return new HttpClient(handler, disposeHandler: false) { BaseAddress = new Uri(baseAddress) };
-        }
-    }
-
-    class StubHandler(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> responder) : HttpMessageHandler
-    {
-        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) => responder(request, cancellationToken);
-    }
 }

src/ServiceControl/Operations/CheckRemotes.cs

@@ -11,9 +11,14 @@
 
     class CheckRemotes : CustomCheck
     {
-        public CheckRemotes(Settings settings, IHttpClientFactory httpClientFactory) : base("ServiceControl Remotes", "Health", TimeSpan.FromSeconds(30))
+        public CheckRemotes(Settings settings, IHttpClientFactory httpClientFactory) : this(settings, httpClientFactory, TimeSpan.FromSeconds(10))
+        {
+        }
+
+        internal CheckRemotes(Settings settings, IHttpClientFactory httpClientFactory, TimeSpan probeTimeout) : base("ServiceControl Remotes", "Health", TimeSpan.FromSeconds(30))
         {
             this.httpClientFactory = httpClientFactory;
+            this.probeTimeout = probeTimeout;
             remoteInstanceSetting = settings.RemoteInstances;
             remoteQueryTasks = new List<Task>(remoteInstanceSetting.Length);
         }
@@ -22,11 +27,12 @@ public override async Task<CheckResult> PerformCheck(CancellationToken cancellat
         {
             try
             {
-                var queryTimeout = TimeSpan.FromSeconds(10);
-                using var cancellationTokenSource = new CancellationTokenSource(queryTimeout);
+                using var probe = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+                probe.CancelAfter(probeTimeout);
+
                 foreach (var remote in remoteInstanceSetting)
                 {
-                    remoteQueryTasks.Add(CheckSuccessStatusCode(remote, queryTimeout, cancellationTokenSource.Token));
+                    remoteQueryTasks.Add(CheckSuccessStatusCode(remote, probe, cancellationToken));
                 }
 
                 try
@@ -59,12 +65,15 @@ public override async Task<CheckResult> PerformCheck(CancellationToken cancellat
             }
         }
 
-        async Task CheckSuccessStatusCode(RemoteInstanceSetting remoteSettings, TimeSpan queryTimeout, CancellationToken cancellationToken)
+        async Task CheckSuccessStatusCode(RemoteInstanceSetting remoteSettings, CancellationTokenSource probe, CancellationToken cancellationToken)
         {
             try
             {
                 var client = httpClientFactory.CreateClient(remoteSettings.InstanceId);
-                var response = await client.GetAsync("/api", cancellationToken);
+                // The remote's client is configured with the query time limit; the probe runs on its own budget.
+                client.Timeout = Timeout.InfiniteTimeSpan;
+
+                var response = await client.GetAsync("/api", probe.Token);
                 response.EnsureSuccessStatusCode();
                 remoteSettings.TemporarilyUnavailable = false;
             }
@@ -73,18 +82,21 @@ async Task CheckSuccessStatusCode(RemoteInstanceSetting remoteSettings, TimeSpan
                 remoteSettings.TemporarilyUnavailable = true;
                 throw new TimeoutException($"The remote instance at '{remoteSettings.BaseAddress}' doesn't seem to be available. It will be temporarily disabled. Reason: {e.Message}", e);
             }
-            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
-            {
-                // Cancelled, noop
-            }
-            catch (OperationCanceledException e) // Intentional, OCE gracefully handled by other catch
+            catch (OperationCanceledException e) when (probe.Token.IsCancellationRequested)
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    // Shutting down; an aborted probe says nothing about the remote's health
+                    return;
+                }
+
                 remoteSettings.TemporarilyUnavailable = true;
-                throw new TimeoutException($"The remote at '{remoteSettings.BaseAddress}' did not respond within the allotted time of '{queryTimeout}'. It will be temporarily disabled.", e);
+                throw new TimeoutException($"The remote at '{remoteSettings.BaseAddress}' did not respond within the allotted time of '{probeTimeout}'. It will be temporarily disabled.", e);
             }
         }
 
         readonly IHttpClientFactory httpClientFactory;
+        readonly TimeSpan probeTimeout;
         RemoteInstanceSetting[] remoteInstanceSetting;
         List<Task> remoteQueryTasks;
     }