Fix #5061: provider-aware databaseFilter overload for AddMartenHighWaterHealthCheck

docs/cSpell.json

@@ -10,6 +10,7 @@
     "hypertables",
     "idempotently",
     "rollups",
+    "rollouts",
     "Polecat",
     "Oakton",
     "Serilog",

docs/events/projections/healthchecks.md

@@ -119,6 +119,34 @@ Services.AddHealthChecks().AddMartenHighWaterHealthCheck(
     includeExternallyManaged: true);
 ```
 
+#### When ownership is runtime state
+
+That `databaseFilter` closure is captured at registration time, so it can only read what the
+application already has in hand — an ambient collection the app keeps up to date, a config value, a
+static. It cannot resolve services.
+
+Under **Wolverine-managed daemon distribution** that is not enough: which databases a node owns is
+decided by Wolverine's agent assignments and moves over the node's lifetime (rollouts, node loss,
+leadership changes), and the authoritative answer has to be read on *each* probe. Use the overload
+whose filter takes the `IServiceProvider` — it is invoked once per database per probe
+(see [marten#5061](https://github.com/JasperFx/marten/issues/5061)):
+
+```cs
+Services.AddHealthChecks().AddMartenHighWaterHealthCheck(
+    // Re-evaluated on every probe, against the provider the health check itself was resolved
+    // from — so scoped services are legal too.
+    (services, database) => services.GetRequiredService<IWolverineRuntime>()
+        .Agents.AllLocallyOwnedDatabaseIds()
+        .Any(id => id.Name.EqualsIgnoreCase(database.Identifier)),
+
+    staleThreshold: TimeSpan.FromSeconds(30),
+    includeExternallyManaged: true);
+```
+
+Both overloads register the settings through a factory, so replacing
+`HighWaterHealthCheckSettings` in the container is also a supported way to override the filter if
+you need to reach further into DI than the predicate allows.
+
 Under `UseTenantPartitionedEvents` the high-water mark is tracked **per tenant** as
 `HighWaterMark:<tenant>` progression rows rather than a single store-global `HighWaterMark`. The
 check evaluates those per-tenant rows too, using the liveness heartbeat (the sequence-gap

src/DaemonTests.ManualOnly/HealthChecks/HighWaterHealthCheckTests.cs

@@ -43,17 +43,21 @@ public HighWaterHealthCheckTests(ITestOutputHelper output): base(output)
 
     private HighWaterHealthCheck buildCheck(TimeSpan? staleThreshold = null, long minimumGap = 1,
         bool autoRestart = false, IProjectionCoordinator? coordinator = null,
-        Func<IMartenDatabase, bool>? databaseFilter = null, bool includeExternallyManaged = false)
+        Func<IMartenDatabase, bool>? databaseFilter = null, bool includeExternallyManaged = false,
+        Func<IServiceProvider, IMartenDatabase, bool>? scopedDatabaseFilter = null,
+        Action<ServiceCollection>? configure = null)
     {
         var services = new ServiceCollection();
         if (coordinator != null)
         {
             services.AddSingleton(coordinator);
         }
 
+        configure?.Invoke(services);
+
         return new(theStore,
             new HighWaterHealthCheckSettings(staleThreshold ?? 30.Seconds(), minimumGap, autoRestart, databaseFilter,
-                includeExternallyManaged), _timeProvider,
+                includeExternallyManaged, scopedDatabaseFilter), _timeProvider,
             _tracker, services.BuildServiceProvider());
     }
 
@@ -374,6 +378,100 @@ public async Task database_filter_including_the_database_still_detects_stuck_mar
         (await check.CheckHealthAsync(new HealthCheckContext(), TestContext.Current.CancellationToken)).Status.ShouldBe(HealthStatus.Unhealthy);
     }
 
+    // ---- provider-aware database scoping (marten#5061) -----------------------------------
+
+    // Stands in for the runtime-owned ownership state that the plain databaseFilter closure cannot
+    // reach: on the reporter's system this is IWolverineRuntime.Agents.AllLocallyOwnedDatabaseIds(),
+    // which changes as Wolverine rebalances agents over the node's lifetime.
+    private sealed class FakeOwnershipRegistry
+    {
+        public bool OwnsEverything { get; set; }
+    }
+
+    [Fact]
+    public async Task scoped_database_filter_excluding_the_database_reports_healthy_despite_stuck_mark()
+    {
+        StoreOptions(x =>
+        {
+            x.Projections.Add(new HwFakeProjection(), ProjectionLifecycle.Async);
+            x.Projections.AsyncMode = DaemonMode.Solo;
+        });
+        await appendEventsAsync(20);
+        await seedHighWaterMarkAsync(1);
+
+        var check = buildCheck(30.Seconds(),
+            scopedDatabaseFilter: (services, _) =>
+                services.GetRequiredService<FakeOwnershipRegistry>().OwnsEverything,
+            configure: services => services.AddSingleton(new FakeOwnershipRegistry { OwnsEverything = false }));
+
+        (await check.CheckHealthAsync(new HealthCheckContext(), TestContext.Current.CancellationToken)).Status
+            .ShouldBe(HealthStatus.Healthy);
+
+        _timeProvider.GetUtcNow().Returns(_now.AddSeconds(60));
+        (await check.CheckHealthAsync(new HealthCheckContext(), TestContext.Current.CancellationToken)).Status
+            .ShouldBe(HealthStatus.Healthy);
+    }
+
+    [Fact]
+    public async Task scoped_database_filter_is_re_evaluated_on_every_probe()
+    {
+        // The whole point of marten#5061: ownership is runtime state. A node that does not own the
+        // database at probe 1 but does at probe 2 must start asserting on it without re-registration.
+        StoreOptions(x =>
+        {
+            x.Projections.Add(new HwFakeProjection(), ProjectionLifecycle.Async);
+            x.Projections.AsyncMode = DaemonMode.Solo;
+        });
+        await appendEventsAsync(20);
+        await seedHighWaterMarkAsync(1);
+
+        var ownership = new FakeOwnershipRegistry { OwnsEverything = false };
+        var check = buildCheck(30.Seconds(),
+            scopedDatabaseFilter: (services, _) =>
+                services.GetRequiredService<FakeOwnershipRegistry>().OwnsEverything,
+            configure: services => services.AddSingleton(ownership));
+
+        // Not owned yet, and well past the staleness threshold: not our problem.
+        _timeProvider.GetUtcNow().Returns(_now.AddSeconds(60));
+        (await check.CheckHealthAsync(new HealthCheckContext(), TestContext.Current.CancellationToken)).Status
+            .ShouldBe(HealthStatus.Healthy);
+
+        // The agent gets assigned here. First probe under ownership establishes the stuck-mark
+        // reading; the gap heuristic then needs a full staleness window at the same mark.
+        ownership.OwnsEverything = true;
+        (await check.CheckHealthAsync(new HealthCheckContext(), TestContext.Current.CancellationToken)).Status
+            .ShouldBe(HealthStatus.Healthy);
+
+        _timeProvider.GetUtcNow().Returns(_now.AddSeconds(120));
+        (await check.CheckHealthAsync(new HealthCheckContext(), TestContext.Current.CancellationToken)).Status
+            .ShouldBe(HealthStatus.Unhealthy);
+    }
+
+    [Fact]
+    public void scoped_overload_registers_the_settings_through_a_factory_carrying_the_filter()
+    {
+        _builder = new();
+        _builder.AddMartenHighWaterHealthCheck(
+            (_, database) => database.Identifier == "owned",
+            staleThreshold: 30.Seconds(),
+            includeExternallyManaged: true);
+
+        var services = _builder.Services.BuildServiceProvider();
+        var settings = services.GetRequiredService<HighWaterHealthCheckSettings>();
+
+        settings.DatabaseFilter.ShouldBeNull();
+        settings.ScopedDatabaseFilter.ShouldNotBeNull();
+        settings.IncludeExternallyManaged.ShouldBeTrue();
+
+        // The registration must be a factory, not a pre-built instance — that is what makes
+        // overriding the settings a supported extension point (marten#5061 suggestion 3).
+        _builder.Services.ShouldContain(x =>
+            x.ServiceType == typeof(HighWaterHealthCheckSettings) && x.ImplementationFactory != null);
+
+        services.GetServices<HealthCheckRegistration>()
+            .ShouldContain(reg => reg.Name == nameof(HighWaterHealthCheck));
+    }
+
     // ---- ExternallyManaged gate (marten#4991) --------------------------------------------
 
     [Fact]

src/Marten.AspNetCore/Daemon/HighWaterHealthCheckExtensions.cs

@@ -83,6 +83,12 @@ public static class HighWaterHealthCheckExtensions
     ///     nodes (e.g. Wolverine-managed), scope this to the databases whose async agents run on
     ///     the local node so the probe does not fan out to (or auto-restart agents on) databases
     ///     this node does not host. Defaults to <c>null</c> (every database — today's behavior).
+    ///     <para>
+    ///         This predicate is captured at registration time, so it cannot resolve services. When
+    ///         ownership is runtime state that has to be re-read on each probe — Wolverine-managed
+    ///         daemon distribution being the usual case — use the overload taking a
+    ///         <c>Func&lt;IServiceProvider, IMartenDatabase, bool&gt;</c> instead (marten#5061).
+    ///     </para>
     /// </param>
     /// <param name="includeExternallyManaged">
     ///     marten#4991: by default the check only runs under <see cref="DaemonMode.Solo" /> /
@@ -103,9 +109,64 @@ public static IHealthChecksBuilder AddMartenHighWaterHealthCheck(
         bool includeExternallyManaged = false
     )
     {
-        builder.Services.AddSingleton(new HighWaterHealthCheckSettings(
+        return builder.registerHighWaterHealthCheck(new HighWaterHealthCheckSettings(
             staleThreshold ?? TimeSpan.FromSeconds(30), minimumGap, autoRestart, databaseFilter,
             includeExternallyManaged));
+    }
+
+    /// <summary>
+    ///     marten#5061: same check, but with a <paramref name="databaseFilter" /> that is handed the
+    ///     <see cref="IServiceProvider" /> and re-evaluated on <em>every</em> probe. Use this
+    ///     overload when "the databases this node owns" is runtime state rather than something known
+    ///     at registration time — the common case under Wolverine-managed daemon distribution, where
+    ///     agent assignments move between nodes over the process's lifetime:
+    ///     <code>
+    ///     Services.AddHealthChecks().AddMartenHighWaterHealthCheck(
+    ///         (services, database) => services.GetRequiredService&lt;IWolverineRuntime&gt;()
+    ///             .Agents.AllLocallyOwnedDatabaseIds()
+    ///             .Any(id =&gt; id.Name.EqualsIgnoreCase(database.Identifier)),
+    ///         staleThreshold: TimeSpan.FromSeconds(30),
+    ///         includeExternallyManaged: true);
+    ///     </code>
+    ///     The <see cref="IServiceProvider" /> passed in is the one the health check itself was
+    ///     resolved from, so scoped services are legal.
+    /// </summary>
+    /// <param name="builder"><see cref="IHealthChecksBuilder" /></param>
+    /// <param name="databaseFilter">
+    ///     Provider-aware predicate, invoked once per database per probe. Return <c>true</c> for the
+    ///     databases this node should assert on.
+    /// </param>
+    /// <param name="staleThreshold">See the other overload. Defaults to 30 seconds.</param>
+    /// <param name="minimumGap">See the other overload. Defaults to 1.</param>
+    /// <param name="autoRestart">See the other overload. Defaults to <c>false</c>.</param>
+    /// <param name="includeExternallyManaged">
+    ///     See the other overload. Usually <c>true</c> when this overload is what you need, since
+    ///     runtime-assigned ownership implies <see cref="DaemonMode.ExternallyManaged" />. Defaults
+    ///     to <c>false</c>.
+    /// </param>
+    public static IHealthChecksBuilder AddMartenHighWaterHealthCheck(
+        this IHealthChecksBuilder builder,
+        Func<IServiceProvider, IMartenDatabase, bool> databaseFilter,
+        TimeSpan? staleThreshold = null,
+        long minimumGap = 1,
+        bool autoRestart = false,
+        bool includeExternallyManaged = false
+    )
+    {
+        ArgumentNullException.ThrowIfNull(databaseFilter);
+
+        return builder.registerHighWaterHealthCheck(new HighWaterHealthCheckSettings(
+            staleThreshold ?? TimeSpan.FromSeconds(30), minimumGap, autoRestart, DatabaseFilter: null,
+            includeExternallyManaged, databaseFilter));
+    }
+
+    private static IHealthChecksBuilder registerHighWaterHealthCheck(this IHealthChecksBuilder builder,
+        HighWaterHealthCheckSettings settings)
+    {
+        // marten#5061: registered as a factory rather than a bare instance so an application that
+        // needs to reach further into DI than ScopedDatabaseFilter allows still has a supported
+        // override point (Replace a Singleton factory, not a Singleton instance).
+        builder.Services.Replace(ServiceDescriptor.Singleton(_ => settings));
         builder.Services.TryAddSingleton(TimeProvider.System);
         builder.Services.TryAddSingleton<HighWaterStateTracker>();
         return builder.AddCheck<HighWaterHealthCheck>(
@@ -117,12 +178,18 @@ public static IHealthChecksBuilder AddMartenHighWaterHealthCheck(
     /// <summary>
     ///     DI-injected settings for <see cref="HighWaterHealthCheck" />.
     /// </summary>
+    /// <param name="ScopedDatabaseFilter">
+    ///     marten#5061: provider-aware alternative to <paramref name="DatabaseFilter" />, resolved
+    ///     against the health check's own <see cref="IServiceProvider" /> on every probe. When both
+    ///     are set a database must satisfy both.
+    /// </param>
     public record HighWaterHealthCheckSettings(
         TimeSpan StaleThreshold,
         long MinimumGap,
         bool AutoRestart = false,
         Func<IMartenDatabase, bool>? DatabaseFilter = null,
-        bool IncludeExternallyManaged = false);
+        bool IncludeExternallyManaged = false,
+        Func<IServiceProvider, IMartenDatabase, bool>? ScopedDatabaseFilter = null);
 
     /// <summary>
     ///     Tracks the fallback gap heuristic's "first observed a stuck mark" reading (keyed per
@@ -155,6 +222,7 @@ public class HighWaterHealthCheck: IHealthCheck
         private readonly long _minimumGap;
         private readonly bool _autoRestart;
         private readonly Func<IMartenDatabase, bool>? _databaseFilter;
+        private readonly Func<IServiceProvider, IMartenDatabase, bool>? _scopedDatabaseFilter;
         private readonly bool _includeExternallyManaged;
         private readonly HighWaterStateTracker _tracker;
         private readonly IServiceProvider _services;
@@ -168,6 +236,7 @@ public HighWaterHealthCheck(IDocumentStore store, HighWaterHealthCheckSettings s
             _minimumGap = settings.MinimumGap;
             _autoRestart = settings.AutoRestart;
             _databaseFilter = settings.DatabaseFilter;
+            _scopedDatabaseFilter = settings.ScopedDatabaseFilter;
             _includeExternallyManaged = settings.IncludeExternallyManaged;
             _tracker = tracker;
             _services = services;
@@ -227,10 +296,19 @@ public async Task<HealthCheckResult> CheckHealthAsync(
 
                 // marten#4991: scope to the databases this node owns so the probe does not fan out
                 // to (or auto-restart) databases the local node does not host the daemon for.
+                // marten#5061: ScopedDatabaseFilter is evaluated here, per probe, against the
+                // provider the check was resolved from — that is the only way to express ownership
+                // that is runtime state (Wolverine reassigns agents over a node's lifetime) rather
+                // than something fixed at registration. When both filters are set, both must pass.
                 IEnumerable<IMartenDatabase> scoped = databases;
                 if (_databaseFilter != null)
                 {
-                    scoped = databases.Where(_databaseFilter);
+                    scoped = scoped.Where(_databaseFilter);
+                }
+
+                if (_scopedDatabaseFilter != null)
+                {
+                    scoped = scoped.Where(database => _scopedDatabaseFilter(_services, database));
                 }
 
                 foreach (var database in scoped)