Managed distribution: per-tenant agents + database-affine assignment for sharded stores (#3280, marten#4806)

docs/guide/durability/marten/distribution.md

@@ -59,6 +59,45 @@ Some other facts about this integration:
   new versions of a projection or subscription on some application nodes in order to do try ["blue/green deployment."](https://en.wikipedia.org/wiki/Blue%E2%80%93green_deployment)
 * This capability does depend on Wolverine's built-in [leadership election](https://en.wikipedia.org/wiki/Leader_election) -- which fortunately got a _lot_ better in Wolverine 3.0
 
+## Database-Affine Distribution for Sharded Stores <Badge type="tip" text="6.x" />
+
+By default Wolverine spreads subscription and projection agents *evenly* across the cluster (with blue/green
+capability matching). That is the right choice for a store-global event store or database-per-tenant multi-tenancy,
+where there is one agent per database anyway.
+
+It is *not* the right choice for a Marten store that combines [sharded multi-tenancy](https://martendb.io/configuration/multitenancy.html#sharded-multi-tenancy-with-database-pooling)
+with [per-tenant event partitioning](https://martendb.io/events/multitenancy.html#per-tenant-event-partitioning). There,
+many tenants are co-located in one shard database and each draws its own event sequence, so Wolverine fans agents out
+one-per-`(shard, tenant)` rather than one-per-database. With hundreds of tenants scattered across many shard databases,
+an even per-agent spread makes *every* node open a connection pool to *nearly every* shard database — so the pool count
+grows as `nodes × databases` and quickly exhausts a shared server's `max_connections`.
+
+To bound that, opt into **database-affine assignment** on the Marten side. Wolverine then keeps every agent for a given
+shard database together on a single node, so a node only opens pools to the databases it actually owns and the pool
+count scales with the number of shard databases, not `nodes × databases`:
+
+```cs
+services.AddMarten(opts =>
+{
+    // ... sharded tenancy + per-tenant event partitioning ...
+    opts.Events.UseTenantPartitionedEvents = true;
+})
+.IntegrateWithWolverine(x =>
+{
+    x.UseWolverineManagedEventSubscriptionDistribution = true;
+
+    // Keep a shard database's per-tenant agents together on one node.
+    x.UseDatabaseAffineAgentAssignment = true;
+});
+```
+
+This setting on `IntegrateWithWolverine` flows through to `StoreOptions.Events.UseDatabaseAffineAgentAssignment`
+(you can also set it there directly). Wolverine reads it off the store through
+`IEventStore.GroupAgentAssignmentsByDatabase`, so the behavior is entirely opt-in and store-driven — a store that
+does not set it keeps the default even distribution. The grouping key is the
+`[event store type]/[event store name]/[database]` prefix of the agent `Uri` (see below), so all of a shard database's
+per-tenant agents share one group.
+
 ## Uri Structure
 
 The `Uri` structure for event subscriptions or projections is:

src/Persistence/MartenTests/Distribution/integrate_with_wolverine_affine_settings.cs

@@ -0,0 +1,63 @@
+using IntegrationTests;
+using JasperFx.CodeGeneration;
+using Marten;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Shouldly;
+using Wolverine;
+using Wolverine.Marten;
+using Xunit;
+
+namespace MartenTests.Distribution;
+
+/// <summary>
+/// JasperFx/marten#4806: the database-affine agent-assignment setting can be configured on
+/// IntegrateWithWolverine (next to UseWolverineManagedEventSubscriptionDistribution), and it must flow
+/// through to the Marten event store — which is where Wolverine's distribution reads it via
+/// IEventStore.GroupAgentAssignmentsByDatabase. This guards the bridge (MartenOverrides resolving the
+/// MartenIntegration and applying the setting) against a DI-resolution regression that would silently
+/// leave the app on the default even distribution.
+/// </summary>
+public class integrate_with_wolverine_affine_settings
+{
+    private static async Task<IHost> StartHostAsync(Action<MartenIntegration> configure, string schema)
+    {
+        return await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Services.AddMarten(m =>
+                    {
+                        m.DisableNpgsqlLogging = true;
+                        m.Connection(Servers.PostgresConnectionString);
+                        m.DatabaseSchemaName = schema;
+                    })
+                    .IntegrateWithWolverine(configure);
+
+                opts.Discovery.DisableConventionalDiscovery();
+                opts.CodeGeneration.TypeLoadMode = TypeLoadMode.Auto;
+            }).StartAsync();
+    }
+
+    [Fact]
+    public async Task affine_setting_flows_to_the_event_store()
+    {
+        using var host = await StartHostAsync(m =>
+        {
+            m.UseWolverineManagedEventSubscriptionDistribution = true;
+            m.UseDatabaseAffineAgentAssignment = true;
+        }, "affine_on");
+
+        var store = (DocumentStore)host.Services.GetRequiredService<IDocumentStore>();
+        store.Options.Events.UseDatabaseAffineAgentAssignment.ShouldBeTrue();
+    }
+
+    [Fact]
+    public async Task default_is_unchanged_when_not_configured()
+    {
+        using var host = await StartHostAsync(
+            m => m.UseWolverineManagedEventSubscriptionDistribution = true, "affine_off");
+
+        var store = (DocumentStore)host.Services.GetRequiredService<IDocumentStore>();
+        store.Options.Events.UseDatabaseAffineAgentAssignment.ShouldBeFalse();
+    }
+}

src/Persistence/MartenTests/Distribution/sharded_tenant_partitioned_distribution.cs

@@ -0,0 +1,190 @@
+using IntegrationTests;
+using JasperFx.CodeGeneration;
+using JasperFx.Core;
+using JasperFx.Events;
+using JasperFx.Events.Projections;
+using JasperFx.MultiTenancy;
+using Marten;
+using Marten.Events;
+using MartenTests.Distribution.Support;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
+using Npgsql;
+using Shouldly;
+using Weasel.Postgresql;
+using Weasel.Postgresql.Migrations;
+using Wolverine;
+using Wolverine.Marten;
+using Wolverine.Runtime.Agents;
+using Wolverine.Tracking;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace MartenTests.Distribution;
+
+// Phase 3 of the per-tenant-partitioned-events matrix — the #3280 regression itself: under
+// MultiTenantedWithShardedDatabases + UseTenantPartitionedEvents, multiple tenants are CO-LOCATED in one
+// shard database, each drawing its own event sequence. A single store-global agent per (shard, database)
+// cannot track them (the lagging tenant's appends fall below the shared high-water mark and are skipped),
+// so managed distribution must fan out one agent per (shard, tenant) — asserted here end-to-end: two
+// ":all/<tenant>" agents for the co-located database, and both tenants' events actually projected.
+//
+// NOTE: this requires the companion JasperFx.Events (jasperfx#482) + Marten (marten#4799) changes — it
+// goes green in CI once those merge, per the PR's stated merge order.
+public class sharded_tenant_partitioned_distribution(ITestOutputHelper output) : PostgresqlContext, IAsyncLifetime
+{
+    private const string TenantA = "tenant-a";
+    private const string TenantB = "tenant-b";
+
+    private IHost theHost = null!;
+    private IDocumentStore theStore = null!;
+
+    // Unique names per run: managed partitions/databases from one run otherwise break the next run's
+    // resource-setup DDL reconciliation.
+    private readonly string theSuffix = Guid.NewGuid().ToString("N")[..8];
+    private string ShardDatabase => $"w3281_shard_{theSuffix}";
+    private string MasterSchema => $"csp_sharded_{theSuffix}";
+
+    public async Task InitializeAsync()
+    {
+        // Provision the physical shard database (the master/pool lives in the default test database
+        // under its own schema).
+        await using (var conn = new NpgsqlConnection(Servers.PostgresConnectionString))
+        {
+            await conn.OpenAsync();
+            if (!await conn.DatabaseExists(ShardDatabase))
+            {
+                await new DatabaseSpecification().BuildDatabase(conn, ShardDatabase);
+            }
+        }
+
+        var shardConnectionString = new NpgsqlConnectionStringBuilder(Servers.PostgresConnectionString)
+        {
+            Database = ShardDatabase
+        }.ConnectionString;
+
+        void ConfigureShardedStore(StoreOptions m)
+        {
+            m.DisableNpgsqlLogging = true;
+            m.MultiTenantedWithShardedDatabases(x =>
+            {
+                x.ConnectionString = Servers.PostgresConnectionString;
+                x.SchemaName = MasterSchema;
+                x.PartitionSchemaName = MasterSchema + "_tenants";
+                x.AddDatabase("shard-1", shardConnectionString);
+            });
+
+            m.Events.StreamIdentity = StreamIdentity.AsString;
+            m.Events.TenancyStyle = TenancyStyle.Conjoined;
+            m.Events.AppendMode = EventAppendMode.Quick;
+            m.Events.UseTenantPartitionedEvents = true;
+            m.Events.UseIdentityMapForAggregates = false;
+
+            m.Schema.For<PartitionedCounter>().MultiTenanted();
+            m.Projections.Snapshot<PartitionedCounter>(SnapshotLifecycle.Async);
+        }
+
+        // Provision the two CO-LOCATED tenants BEFORE the Wolverine host starts, so the first agent
+        // enumeration already sees them and fans out per tenant. (Adding a database's FIRST tenants while
+        // the host runs leaves the store-global agent from the empty-database phase running alongside the
+        // new per-tenant agents until assignments reconcile — a transition edge outside this test's scope.)
+        await using (var provisioning = DocumentStore.For(ConfigureShardedStore))
+        {
+            // Apply the event schema (parent partitioned tables) to the shard database first, so the
+            // per-tenant LIST partitions + sequences created by the explicit assignment have a parent.
+            await provisioning.Storage.ApplyAllConfiguredChangesToDatabaseAsync();
+            await provisioning.Advanced.AddTenantToShardAsync(TenantA, "shard-1", default);
+            await provisioning.Advanced.AddTenantToShardAsync(TenantB, "shard-1", default);
+        }
+
+        theHost = await Host.CreateDefaultBuilder()
+            .UseWolverine(opts =>
+            {
+                opts.Durability.HealthCheckPollingTime = 1.Seconds();
+                opts.Durability.CheckAssignmentPeriod = 1.Seconds();
+
+                opts.Services.AddMarten(ConfigureShardedStore)
+                    .IntegrateWithWolverine(m =>
+                    {
+                        m.UseWolverineManagedEventSubscriptionDistribution = true;
+                        // Sharded tenancy has no single "main" database, so Wolverine's message store needs
+                        // an explicit master (mirrors a real deployment pointing this at a references DB).
+                        m.MainDatabaseConnectionString = Servers.PostgresConnectionString;
+                    });
+
+                opts.Services.AddSingleton<ILoggerProvider>(new OutputLoggerProvider(output));
+                opts.Discovery.DisableConventionalDiscovery();
+                opts.CodeGeneration.TypeLoadMode = TypeLoadMode.Auto;
+            }).StartAsync();
+
+        theStore = theHost.Services.GetRequiredService<IDocumentStore>();
+    }
+
+    public async Task DisposeAsync()
+    {
+        theHost.GetRuntime().Agents.DisableHealthChecks();
+        await theHost.StopAsync();
+        theHost.Dispose();
+    }
+
+    [Fact]
+    public async Task co_located_tenants_get_their_own_agents_and_both_project()
+    {
+        await theHost.WaitUntilAssumesLeadershipAsync(10.Seconds());
+
+        // One async projection over one shard database with two co-located tenants => one agent PER TENANT.
+        await theHost.WaitUntilAssignmentsChangeTo(w =>
+        {
+            w.AgentScheme = EventSubscriptionAgentFamily.SchemeName;
+            w.ExpectRunningAgents(theHost, 2);
+        }, 60.Seconds());
+
+        var uris = await GetAgentUrisAsync(theHost);
+        uris.Length.ShouldBe(2);
+        uris.ShouldContain(u => u.TrimEnd('/').EndsWith("/" + TenantA, StringComparison.OrdinalIgnoreCase));
+        uris.ShouldContain(u => u.TrimEnd('/').EndsWith("/" + TenantB, StringComparison.OrdinalIgnoreCase));
+
+        // Both tenants' events are projected — each tenant advances against its OWN high-water mark, so the
+        // second tenant's appends are not skipped below a shared mark (the #3280 failure mode).
+        var id = "pc-" + Guid.NewGuid().ToString("N");
+        await AppendAsync(TenantA, id, 5);
+        await AppendAsync(TenantB, id, 11); // same stream id, other tenant partition
+
+        (await WaitForProjectionAsync(TenantA, id, 60.Seconds()))!.Total.ShouldBe(5);
+        (await WaitForProjectionAsync(TenantB, id, 60.Seconds()))!.Total.ShouldBe(11);
+    }
+
+    private async Task AppendAsync(string tenant, string id, int amount)
+    {
+        await using var session = theStore.LightweightSession(tenant);
+        session.Events.StartStream<PartitionedCounter>(id, new CounterChanged(amount));
+        await session.SaveChangesAsync();
+    }
+
+    private async Task<PartitionedCounter?> WaitForProjectionAsync(string tenant, string id, TimeSpan timeout)
+    {
+        var deadline = DateTimeOffset.UtcNow + timeout;
+        while (DateTimeOffset.UtcNow < deadline)
+        {
+            await using var session = theStore.QuerySession(tenant);
+            var doc = await session.LoadAsync<PartitionedCounter>(id);
+            if (doc != null)
+            {
+                return doc;
+            }
+
+            await Task.Delay(250.Milliseconds());
+        }
+
+        return null;
+    }
+
+    private static async Task<string[]> GetAgentUrisAsync(IHost host)
+    {
+        var family = host.Services.GetServices<IAgentFamily>()
+            .OfType<EventSubscriptionAgentFamily>().Single();
+        var agents = await family.AllKnownAgentsAsync();
+        return [.. agents.Select(x => x.AbsoluteUri)];
+    }
+}

src/Persistence/Wolverine.Marten/MartenIntegration.cs

@@ -40,12 +40,24 @@ public class MartenIntegration : IWolverineExtension, IEventForwarding
     public bool UseFastEventForwarding { get; set; }
 
     /// <summary>
-    /// Use this when using Wolverine to evenly distribute event projection and subscription 
-    /// work of Marten asynchronous projections. This replaces Marten's <c>AddAsyncDaemon(HotCold)</c> 
+    /// Use this when using Wolverine to evenly distribute event projection and subscription
+    /// work of Marten asynchronous projections. This replaces Marten's <c>AddAsyncDaemon(HotCold)</c>
     /// option and should not be used in combination with Marten's own load distribution.
     /// </summary>
     public bool UseWolverineManagedEventSubscriptionDistribution { get; set; }
 
+    /// <summary>
+    /// Opt-in (default false): when <see cref="UseWolverineManagedEventSubscriptionDistribution"/> is on and
+    /// the Marten store uses sharded databases with per-tenant event partitioning, assign each shard
+    /// database's per-(shard, tenant) agents with <b>database affinity</b> — all of a database's agents run
+    /// on one node — instead of spreading them evenly by count. This bounds each node to the shard databases
+    /// it owns, so connection pools scale with the number of shard databases rather than nodes × databases
+    /// (which otherwise exhausts a shared server's max_connections). Flows through to
+    /// <c>StoreOptions.Events.UseDatabaseAffineAgentAssignment</c>; only takes effect for a sharded,
+    /// per-tenant-partitioned store. See JasperFx/marten#4806.
+    /// </summary>
+    public bool UseDatabaseAffineAgentAssignment { get; set; }
+
     public void Configure(WolverineOptions options)
     {
         // Duplicate incoming messages
@@ -201,6 +213,17 @@ public void Configure(IServiceProvider services, StoreOptions options)
                     nameof(Saga.Version), typeof(int))!;
             }
         });
+
+        // Bridge the Wolverine-managed distribution's database-affine agent-assignment setting — configured
+        // on IntegrateWithWolverine — onto the Marten event store, because Wolverine's distribution is what
+        // consumes it (via IEventStore.GroupAgentAssignmentsByDatabase). Only the main store
+        // (StoreType == null): this is a sharded per-tenant-store concern, not ancillary.
+        // See JasperFx/marten#4806.
+        if (StoreType == null &&
+            services.GetService<MartenIntegration>() is { UseDatabaseAffineAgentAssignment: true })
+        {
+            options.Events.UseDatabaseAffineAgentAssignment = true;
+        }
     }
 }