natsdotnet/tests/NATS.Server.Tests/JetStream/Cluster/TopologyPlacementTests.cs

// Go parity: golang/nats-server/server/jetstream_cluster.go:7212 selectPeerGroup
// Covers: UniqueTag enforcement, HA asset limits, weighted scoring by available resources.
using NATS.Server.JetStream.Cluster;

namespace NATS.Server.Tests.JetStream.Cluster;

/// <summary>
/// Tests for topology-aware placement: JetStreamUniqueTag enforcement,
/// MaxAssetsPerPeer HA limits, and weighted scoring.
/// Go reference: jetstream_cluster.go:7212 selectPeerGroup (uniqueTagPrefix, maxHaAssets, weighted sort).
/// </summary>
public class TopologyPlacementTests
{
    // ---------------------------------------------------------------
    // UniqueTag enforcement
    // Go reference: jetstream_cluster.go:7251 uniqueTagPrefix / checkUniqueTag
    // ---------------------------------------------------------------

    [Fact]
    public void UniqueTag_prevents_same_tag_value_replicas()
    {
        // 3 peers: p1 and p2 in az:us-east-1a, p3 in az:us-east-1b.
        // R=2 with UniqueTag="az" must pick one from each AZ.
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", Tags = ["az:us-east-1a"], AvailableStorage = 1000 },
            new() { PeerId = "p2", Tags = ["az:us-east-1a"], AvailableStorage = 2000 },
            new() { PeerId = "p3", Tags = ["az:us-east-1b"], AvailableStorage = 900 },
        };
        var policy = new PlacementPolicy { UniqueTag = "az" };

        var group = PlacementEngine.SelectPeerGroup("az-group", 2, peers, policy);

        group.Peers.Count.ShouldBe(2);
        // One peer must be from az:us-east-1a and one from az:us-east-1b.
        var selectedPeers = peers.Where(p => group.Peers.Contains(p.PeerId)).ToList();
        var azValues = selectedPeers
            .SelectMany(p => p.Tags)
            .Where(t => t.StartsWith("az:", StringComparison.OrdinalIgnoreCase))
            .ToList();
        azValues.Distinct(StringComparer.OrdinalIgnoreCase).Count().ShouldBe(2);
    }

    [Fact]
    public void UniqueTag_throws_when_not_enough_unique_values()
    {
        // All 3 peers share the same AZ tag; R=2 requires 2 unique AZ values → impossible.
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", Tags = ["az:us-east-1a"] },
            new() { PeerId = "p2", Tags = ["az:us-east-1a"] },
            new() { PeerId = "p3", Tags = ["az:us-east-1a"] },
        };
        var policy = new PlacementPolicy { UniqueTag = "az" };

        Should.Throw<InvalidOperationException>(
            () => PlacementEngine.SelectPeerGroup("fail", 2, peers, policy));
    }

    [Fact]
    public void Tag_prefix_matching_for_unique_constraint()
    {
        // UniqueTag="az" should match tags like "az:us-east-1a", "az:us-west-2b", etc.
        // Go reference: jetstream_cluster.go:7265 strings.HasPrefix(t, uniqueTagPrefix)
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", Tags = ["az:us-east-1a", "ssd"] },
            new() { PeerId = "p2", Tags = ["az:us-west-2b", "ssd"] },
            new() { PeerId = "p3", Tags = ["az:eu-central-1a", "ssd"] },
        };
        var policy = new PlacementPolicy { UniqueTag = "az" };

        var group = PlacementEngine.SelectPeerGroup("prefix", 3, peers, policy);

        group.Peers.Count.ShouldBe(3);
        group.Peers.ShouldContain("p1");
        group.Peers.ShouldContain("p2");
        group.Peers.ShouldContain("p3");
    }

    [Fact]
    public void Empty_unique_tag_ignored()
    {
        // UniqueTag="" or null → no unique constraint applied, normal selection.
        // Go reference: jetstream_cluster.go:7252 if uniqueTagPrefix != _EMPTY_
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", Tags = ["az:us-east-1a"] },
            new() { PeerId = "p2", Tags = ["az:us-east-1a"] },
            new() { PeerId = "p3", Tags = ["az:us-east-1a"] },
        };

        // No UniqueTag policy — all 3 peers are valid, R=3 should succeed.
        var groupNull = PlacementEngine.SelectPeerGroup("no-unique-null", 3, peers, policy: null);
        groupNull.Peers.Count.ShouldBe(3);

        // Empty string UniqueTag → treated as disabled.
        var policy = new PlacementPolicy { UniqueTag = "" };
        var groupEmpty = PlacementEngine.SelectPeerGroup("no-unique-empty", 3, peers, policy);
        groupEmpty.Peers.Count.ShouldBe(3);
    }

    [Fact]
    public void UniqueTag_combined_with_cluster_filter()
    {
        // Both cluster filter and UniqueTag must be applied together.
        // Go reference: jetstream_cluster.go:7346 cluster check before uniqueTag check
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", Cluster = "us-east", Tags = ["az:us-east-1a"] },
            new() { PeerId = "p2", Cluster = "us-east", Tags = ["az:us-east-1a"] },
            new() { PeerId = "p3", Cluster = "us-east", Tags = ["az:us-east-1b"] },
            new() { PeerId = "p4", Cluster = "us-west", Tags = ["az:us-west-2a"] },
        };
        var policy = new PlacementPolicy { Cluster = "us-east", UniqueTag = "az" };

        // Only p1/p2/p3 are in us-east; UniqueTag="az" → picks one from 1a and one from 1b.
        var group = PlacementEngine.SelectPeerGroup("combo", 2, peers, policy);

        group.Peers.Count.ShouldBe(2);
        group.Peers.ShouldNotContain("p4");
        var selectedPeers = peers.Where(p => group.Peers.Contains(p.PeerId)).ToList();
        var azValues = selectedPeers
            .SelectMany(p => p.Tags)
            .Where(t => t.StartsWith("az:", StringComparison.OrdinalIgnoreCase))
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToList();
        azValues.Count.ShouldBe(2);
    }

    // ---------------------------------------------------------------
    // MaxAssetsPerPeer HA limit deprioritization
    // Go reference: jetstream_cluster.go:7428 maxHaAssets check (deprioritize vs hard exclude)
    // ---------------------------------------------------------------

    [Fact]
    public void MaxAssetsPerPeer_deprioritizes_overloaded_peers()
    {
        // p1 is at its asset limit but p2 and p3 are not.
        // With enough non-overloaded candidates, overloaded peer should not be selected.
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", AvailableStorage = 10_000, CurrentAssets = 5, MaxAssetsPerPeer = 5 },
            new() { PeerId = "p2", AvailableStorage = 8_000,  CurrentAssets = 1, MaxAssetsPerPeer = 5 },
            new() { PeerId = "p3", AvailableStorage = 6_000,  CurrentAssets = 0, MaxAssetsPerPeer = 5 },
        };

        var group = PlacementEngine.SelectPeerGroup("ha-limit", 2, peers);

        // p1 is deprioritized (at max), so p2 and p3 should be selected over p1.
        group.Peers.Count.ShouldBe(2);
        group.Peers.ShouldContain("p2");
        group.Peers.ShouldContain("p3");
        group.Peers.ShouldNotContain("p1");
    }

    [Fact]
    public void MaxAssetsPerPeer_still_used_when_no_alternatives()
    {
        // All peers are at their HA asset limit, but we must still select from them.
        // Go reference: jetstream_cluster.go — deprioritize (move to end), not hard exclude.
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", AvailableStorage = 1000, CurrentAssets = 3, MaxAssetsPerPeer = 3 },
            new() { PeerId = "p2", AvailableStorage = 900,  CurrentAssets = 3, MaxAssetsPerPeer = 3 },
        };

        // Should succeed even though both peers are at max.
        var group = PlacementEngine.SelectPeerGroup("ha-fallback", 2, peers);

        group.Peers.Count.ShouldBe(2);
        group.Peers.ShouldContain("p1");
        group.Peers.ShouldContain("p2");
    }

    [Fact]
    public void Zero_MaxAssets_means_unlimited()
    {
        // MaxAssetsPerPeer=0 → no asset limit, peer treated as not overloaded regardless of CurrentAssets.
        var peers = new List<PeerInfo>
        {
            new() { PeerId = "p1", AvailableStorage = 5000, CurrentAssets = 100, MaxAssetsPerPeer = 0 },
            new() { PeerId = "p2", AvailableStorage = 4000, CurrentAssets = 200, MaxAssetsPerPeer = 0 },
        };

        var group = PlacementEngine.SelectPeerGroup("unlimited", 2, peers);

        group.Peers.Count.ShouldBe(2);
        group.Peers.ShouldContain("p1");
        group.Peers.ShouldContain("p2");
    }

    // ---------------------------------------------------------------
    // Weighted score = AvailableStorage - (CurrentAssets * AssetCostWeight)
    // Go reference: jetstream_cluster.go:7469 sort by avail then ns (stream count)
    // ---------------------------------------------------------------

    [Fact]
    public void Weighted_score_prefers_less_loaded_peers()
    {
        // p1: more storage but many assets → lower score
        // p2: less storage but few assets → higher score
        // With DefaultAssetCostWeight = 1GB, even a small difference in assets
        // can overcome a moderate storage advantage.
        const long gb = PlacementEngine.DefaultAssetCostWeight; // 1_073_741_824L
        var peers = new List<PeerInfo>
        {
            // p1: score = 10*GB - 5*GB = 5*GB
            new() { PeerId = "p1", AvailableStorage = 10 * gb, CurrentAssets = 5 },
            // p2: score = 9*GB - 1*GB = 8*GB  (wins despite less raw storage)
            new() { PeerId = "p2", AvailableStorage = 9 * gb,  CurrentAssets = 1 },
            // p3: score = 3*GB - 0 = 3*GB
            new() { PeerId = "p3", AvailableStorage = 3 * gb,  CurrentAssets = 0 },
        };

        var group = PlacementEngine.SelectPeerGroup("weighted", 2, peers);

        // p2 has the highest score (8*GB), p1 has second (5*GB).
        group.Peers.Count.ShouldBe(2);
        group.Peers[0].ShouldBe("p2");
        group.Peers[1].ShouldBe("p1");
    }

    [Fact]
    public void Weighted_score_with_custom_cost_weight()
    {
        // Verify score formula: score = AvailableStorage - (CurrentAssets * AssetCostWeight)
        // Use a fixed, small cost weight to make the math obvious.
        const long costWeight = 1000L;
        var peers = new List<PeerInfo>
        {
            // score = 5000 - (3 * 1000) = 2000
            new() { PeerId = "p1", AvailableStorage = 5000, CurrentAssets = 3 },
            // score = 4000 - (0 * 1000) = 4000  (wins)
            new() { PeerId = "p2", AvailableStorage = 4000, CurrentAssets = 0 },
            // score = 6000 - (5 * 1000) = 1000  (loses)
            new() { PeerId = "p3", AvailableStorage = 6000, CurrentAssets = 5 },
        };

        var group = PlacementEngine.SelectPeerGroup("custom-weight", 2, peers, assetCostWeight: costWeight);

        group.Peers.Count.ShouldBe(2);
        group.Peers[0].ShouldBe("p2"); // score 4000
        group.Peers[1].ShouldBe("p1"); // score 2000
    }
}