Fix DistributeByGroupAffinity assigning a bumped projection version to nodes that cannot build it

src/Persistence/MartenTests/MultiTenancy/blue_green_version_bump_assignment.cs

@@ -0,0 +1,258 @@
+using IntegrationTests;
+using JasperFx.Events;
+using JasperFx.Events.Projections;
+using Marten;
+using Marten.Events.Aggregation;
+using Npgsql;
+using Shouldly;
+using JasperFx;
+using Weasel.Postgresql;
+using Weasel.Postgresql.Migrations;
+using Wolverine.Runtime.Agents;
+
+namespace MartenTests.MultiTenancy;
+
+/// <summary>
+/// End-to-end cover for a projection version bump on a <b>multi-database</b> store, through the real
+/// <see cref="EventSubscriptionAgentFamily.EvaluateAssignmentsAsync"/> rather than
+/// <see cref="AssignmentGrid.DistributeByGroupAffinity(string, Func{Uri, string}, Func{Uri, bool})"/>
+/// directly — so it also covers the store-cardinality pass selection and
+/// <c>RetireSupersededAgentsAsync</c>, both of which sit between a leader and the placement.
+///
+/// <para>The scenario is a blue/green rollout: the "blue" fleet runs the store as it is deployed today
+/// and is <i>already running</i> its agents; the "green" fleet runs a build where one projection's
+/// <c>Version</c> is one higher. Both fleets are the same application over the same tenant databases,
+/// and the leader evaluating the assignments is a blue node — which is what makes the green fleet's
+/// agents reachable only through its persisted node capabilities.</para>
+/// </summary>
+public class blue_green_version_bump_assignment : IAsyncLifetime
+{
+    private const uint PreviousVersion = 2;
+    private const uint BumpedVersion = 3;
+
+    // Our shape, and the one that makes a group big enough to matter: several tenants share a shard
+    // database, and events are tenant-partitioned, so the daemon runs an agent per
+    // (database, tenant, projection version) rather than one per (database, version).
+    private static readonly string[] _databases = ["bgshard1", "bgshard2", "bgshard3"];
+    private static readonly string[] _tenantsPerDatabase = ["alpha", "beta", "gamma"];
+
+    private DocumentStore _blue = null!;
+    private DocumentStore _green = null!;
+    private IReadOnlyList<Uri> _blueAgents = null!;
+    private IReadOnlyList<Uri> _greenAgents = null!;
+
+    public async ValueTask InitializeAsync()
+    {
+        await using var conn = new NpgsqlConnection(Servers.PostgresConnectionString);
+        await conn.OpenAsync(TestContext.Current.CancellationToken);
+
+        var connectionStrings = new List<string>();
+        foreach (var database in _databases)
+        {
+            connectionStrings.Add(await CreateDatabaseIfNotExists(conn, database));
+        }
+
+        await conn.CloseAsync();
+
+        _blue = StoreFor(PreviousVersion, connectionStrings);
+        _green = StoreFor(BumpedVersion, connectionStrings);
+
+        _blueAgents = await AdvertisedAgentsAsync(_blue);
+        _greenAgents = await AdvertisedAgentsAsync(_green);
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _blue.DisposeAsync();
+        await _green.DisposeAsync();
+    }
+
+    [Fact]
+    public async Task the_two_fleets_advertise_the_bumped_projection_under_disjoint_identities()
+    {
+        // Guard the shape this test exists for: agents are per (database, tenant, projection), so each
+        // shard database's group is tenants x projections deep. Without this, a change that stopped
+        // distributing per tenant would silently shrink every group to one agent per projection and this
+        // whole file would go on passing while testing something much weaker.
+        foreach (var database in _blueAgents.GroupBy(DatabaseKeyOf))
+        {
+            database.Count().ShouldBe(_tenantsPerDatabase.Length * 2,
+                $"{database.Key} must carry an agent per tenant for each of the two projections");
+        }
+
+        _blueAgents.Count.ShouldBe(_databases.Length * _tenantsPerDatabase.Length * 2);
+        _greenAgents.Count.ShouldBe(_blueAgents.Count);
+
+        var blueBumped = AgentsFor(_blueAgents, "Trip");
+        var greenBumped = AgentsFor(_greenAgents, "Trip");
+
+        blueBumped.ShouldNotBeEmpty();
+        greenBumped.Count.ShouldBe(blueBumped.Count);
+        blueBumped.Intersect(greenBumped).ShouldBeEmpty(
+            "a version bump has to change the agent identity, or the leader has no way to tell the two fleets apart");
+
+        // The projection that was NOT bumped is one and the same agent on both fleets, which is what
+        // makes it declared by every node and therefore free to land anywhere.
+        AgentsFor(_greenAgents, "Passenger")
+            .OrderBy(x => x.ToString())
+            .ShouldBe(AgentsFor(_blueAgents, "Passenger").OrderBy(x => x.ToString()));
+    }
+
+    [Fact]
+    public async Task the_bumped_version_is_assigned_to_the_fleet_that_declares_it()
+    {
+        var (grid, blue, green) = BuildGrid();
+
+        // The leader is a blue node: its own family enumerates only the previous version, exactly as in
+        // production, so the green agents reach the grid solely through node capabilities.
+        await using var leader = new EventSubscriptionAgentFamily([_blue], []);
+        await leader.EvaluateAssignmentsAsync(grid);
+
+        foreach (var uri in AgentsFor(_greenAgents, "Trip"))
+        {
+            var host = grid.AgentFor(uri).AssignedNode;
+            host.ShouldNotBeNull($"{uri} was left unassigned, so nothing would build the new version");
+            green.ShouldContain(host!,
+                $"{uri} may only run on a node that declares it — BuildAgentAsync throws 'Unknown event projection or subscription' on the other fleet");
+        }
+
+        foreach (var uri in AgentsFor(_blueAgents, "Trip"))
+        {
+            var host = grid.AgentFor(uri).AssignedNode;
+            host.ShouldNotBeNull($"{uri} was left unassigned, so the fleet still serving would stop projecting");
+            blue.ShouldContain(host!, $"{uri} is the version the blue fleet serves and must stay there");
+        }
+    }
+
+    // The "one owner per version per database" bound that the sibling-partition preference exists for is
+    // asserted in CoreTests.Runtime.Agents.distribute_by_group_affinity instead. It needs agents to
+    // outnumber nodes the way they do in a real cluster (~5k agents over ~10 nodes); with three databases
+    // over four nodes the per-node ceiling binds first and legitimately spreads a group's third partition
+    // to balance load, which is the intended trade and not what this test is about.
+
+    private (AssignmentGrid Grid, List<AssignmentGrid.Node> Blue, List<AssignmentGrid.Node> Green) BuildGrid()
+    {
+        var grid = new AssignmentGrid();
+
+        var blue = new List<AssignmentGrid.Node>
+        {
+            grid.WithNode(1, Guid.NewGuid()).HasCapabilities(_blueAgents),
+            grid.WithNode(2, Guid.NewGuid()).HasCapabilities(_blueAgents)
+        };
+
+        var green = new List<AssignmentGrid.Node>
+        {
+            grid.WithNode(3, Guid.NewGuid()).HasCapabilities(_greenAgents),
+            grid.WithNode(4, Guid.NewGuid()).HasCapabilities(_greenAgents)
+        };
+
+        // Blue is already running what it advertises. That incumbency is the whole point: it is what
+        // keeps a blue node in the candidate set of a group it can only partly run.
+        for (var i = 0; i < blue.Count; i++)
+        {
+            blue[i].Running(_blueAgents.Where((_, index) => index % blue.Count == i).ToArray());
+        }
+
+        // Mirrors NodeAgentController.EvaluateAssignmentsAsync, which seeds the grid from the union of
+        // every node's persisted capabilities before handing it to the families.
+        grid.WithAgents(_blueAgents.Concat(_greenAgents).Distinct().ToArray());
+
+        return (grid, blue, green);
+    }
+
+    private static async Task<IReadOnlyList<Uri>> AdvertisedAgentsAsync(IEventStore store)
+    {
+        await using var family = new EventSubscriptionAgentFamily([store], []);
+        return await family.SupportedAgentsAsync();
+    }
+
+    private static DocumentStore StoreFor(uint tripVersion, IReadOnlyList<string> connectionStrings)
+    {
+        return DocumentStore.For(opts =>
+        {
+            opts.DatabaseSchemaName = "bluegreen";
+
+            // Nothing here writes or projects; the test only asks the store what it would advertise.
+            opts.AutoCreateSchemaObjects = AutoCreate.None;
+
+            // Sharded shape: every database holds several tenants, and the event tables are partitioned
+            // per tenant, which is what turns IEventStore.DistributesAgentsPerTenant on and produces the
+            // per-(database, tenant) agents a real shard database's group is made of.
+            opts.Events.TenancyStyle = JasperFx.MultiTenancy.TenancyStyle.Conjoined;
+            opts.Events.UseTenantPartitionedEvents = true;
+            opts.Events.UseArchivedStreamPartitioning = false;
+
+            opts.MultiTenantedDatabases(tenancy =>
+            {
+                for (var i = 0; i < connectionStrings.Count; i++)
+                {
+                    tenancy.AddMultipleTenantDatabase(connectionStrings[i], _databases[i])
+                        .ForTenants(_tenantsPerDatabase.Select(t => $"{_databases[i]}-{t}").ToArray());
+                }
+            });
+
+            // Conjoined events require conjoined read models.
+            opts.Schema.For<BlueGreenTrip>().MultiTenanted();
+            opts.Schema.For<BlueGreenPassengerCount>().MultiTenanted();
+
+            opts.Projections.Add(new TripProjection { Version = tripVersion }, ProjectionLifecycle.Async);
+            opts.Projections.Add(new PassengerProjection(), ProjectionLifecycle.Async);
+        });
+    }
+
+    private static async Task<string> CreateDatabaseIfNotExists(NpgsqlConnection conn, string databaseName)
+    {
+        var builder = new NpgsqlConnectionStringBuilder(Servers.PostgresConnectionString);
+
+        if (!await conn.DatabaseExists(databaseName))
+        {
+            await new DatabaseSpecification().BuildDatabase(conn, databaseName);
+        }
+
+        builder.Database = databaseName;
+        return builder.ConnectionString;
+    }
+
+    private static List<Uri> AgentsFor(IEnumerable<Uri> agents, string projectionName) =>
+        agents.Where(uri => uri.Segments.Any(segment =>
+            segment.Trim('/').Equals(projectionName, StringComparison.OrdinalIgnoreCase))).ToList();
+
+    // The (type, name, databaseId) prefix of an agent URI, mirroring the internal
+    // EventSubscriptionAgentFamily.DatabaseKeyOf that group affinity keys on.
+    private static string DatabaseKeyOf(Uri uri) =>
+        uri.Segments.Length >= 3
+            ? $"{uri.Host}/{uri.Segments[1].Trim('/')}/{uri.Segments[2].Trim('/')}"
+            : uri.AbsoluteUri;
+}
+
+public record TripStarted(Guid Id, string Description);
+
+public record PassengerBoarded(Guid TripId, string Name);
+
+public class BlueGreenTrip
+{
+    public Guid Id { get; set; }
+    public string Description { get; set; } = string.Empty;
+}
+
+public class BlueGreenPassengerCount
+{
+    public Guid Id { get; set; }
+    public int Count { get; set; }
+}
+
+public partial class TripProjection : SingleStreamProjection<BlueGreenTrip, Guid>
+{
+    public TripProjection() => Name = "Trip";
+
+    public static BlueGreenTrip Create(TripStarted started) =>
+        new() { Id = started.Id, Description = started.Description };
+}
+
+public partial class PassengerProjection : SingleStreamProjection<BlueGreenPassengerCount, Guid>
+{
+    public PassengerProjection() => Name = "Passenger";
+
+    public static BlueGreenPassengerCount Create(PassengerBoarded boarded) =>
+        new() { Id = boarded.TripId, Count = 1 };
+}

src/Testing/CoreTests/Runtime/Agents/distribute_by_group_affinity.cs

@@ -235,5 +235,85 @@ public void an_incumbent_node_keeps_its_group_up_to_the_ceiling()
         grid.AgentFor(g2).AssignedNode.ShouldBe(node2, "an under-ceiling incumbent keeps its group");
         grid.AgentFor(g3).AssignedNode.ShouldBe(node3, "only the over-ceiling group moves, to the empty node");
     }
+
+    // event-subscriptions://{type}/{name}/{databaseId}/{projection}/{shardKey}/v{version}/{tenant}
+    // — the real EventSubscriptionAgentFamily.UriFor grammar, so the version sits in its own segment.
+    private static Uri VersionedAgent(string db, uint version, string tenant) =>
+        new($"event-subscriptions://marten/main/{db}/Proj/All/v{version}/{tenant}");
+
+    [Fact]
+    public void a_version_bump_splits_a_group_between_the_old_and_new_version_nodes()
+    {
+        // A projection version bump on a sharded store: one shard database's group spans the previous
+        // version's agent — declared by, and RUNNING on, the blue node — and the new version's agent,
+        // declared only by the green node. No node is capable of the whole group, so the members must
+        // fall back individually: the old version stays on blue, the new version goes to green.
+        //
+        // This is the intersection of the two cases above, and it is what a blue/green deployment of a
+        // sharded store looks like at every evaluation for the whole rollout, not just transiently.
+        var previous = VersionedAgent("db1", 22, "t1");
+        var bumped = VersionedAgent("db1", 23, "t1");
+
+        var grid = new AssignmentGrid();
+        var blue = grid.WithNode(1, Guid.NewGuid()).HasCapabilities(new[] { previous });
+        blue.Running(previous);
+        var green = grid.WithNode(2, Guid.NewGuid()).HasCapabilities(new[] { bumped });
+
+        grid.WithAgents(previous, bumped);
+
+        grid.DistributeByGroupAffinity("event-subscriptions", DatabaseKey);
+
+        grid.AgentFor(previous).AssignedNode.ShouldBe(blue,
+            "the previous version keeps running where it is");
+        grid.AgentFor(bumped).AssignedNode.ShouldBe(green,
+            "the new version's agent may only run on the node that declares it — the blue node cannot build it");
+    }
+
+    [Fact]
+    public void a_split_group_costs_only_as_many_nodes_as_the_capability_split_forces()
+    {
+        // The same version bump, now with the projections whose version did NOT change also in the group.
+        // Every node declares those, so nothing stops them landing on a third node — and what a shard
+        // database costs in connection pools is the number of DISTINCT nodes holding any of its agents, so
+        // a third host is a third pool set on that database. A split group must still occupy only as many
+        // nodes as the capability split forces: one per version.
+        var databases = new[] { "db1", "db2", "db3", "db4" };
+        var tenants = new[] { "t1", "t2" };
+
+        Uri Unchanged(string db, string tenant) =>
+            new($"event-subscriptions://marten/main/{db}/Other/All/v7/{tenant}");
+
+        Uri[] Across(Func<string, string, Uri> agent) =>
+            databases.SelectMany(db => tenants.Select(t => agent(db, t))).ToArray();
+
+        var previous = Across((db, t) => VersionedAgent(db, 22, t));
+        var bumped = Across((db, t) => VersionedAgent(db, 23, t));
+        var unchanged = Across(Unchanged);
+
+        var grid = new AssignmentGrid();
+        var blues = new[] { grid.WithNode(1, Guid.NewGuid()), grid.WithNode(2, Guid.NewGuid()) };
+        var greens = new[] { grid.WithNode(3, Guid.NewGuid()), grid.WithNode(4, Guid.NewGuid()) };
+
+        foreach (var blue in blues) blue.HasCapabilities(previous.Concat(unchanged));
+        foreach (var green in greens) green.HasCapabilities(bumped.Concat(unchanged));
+
+        grid.WithAgents(previous.Concat(bumped).Concat(unchanged).ToArray());
+
+        grid.DistributeByGroupAffinity("event-subscriptions", DatabaseKey);
+
+        foreach (var db in databases)
+        {
+            var hosts = tenants
+                .SelectMany(t => new[] { VersionedAgent(db, 22, t), VersionedAgent(db, 23, t), Unchanged(db, t) })
+                .Select(uri => grid.AgentFor(uri).AssignedNode)
+                .Distinct()
+                .ToList();
+
+            hosts.Count.ShouldBe(2,
+                $"{db} must sit on exactly two nodes — one per version — so it attracts two pool sets, not three");
+            hosts.ShouldContain(host => blues.Contains(host!), $"{db}'s previous version must be on a blue node");
+            hosts.ShouldContain(host => greens.Contains(host!), $"{db}'s new version must be on a green node");
+        }
+    }
 }