natsdotnet/tests/NATS.Server.Core.Tests/SlowConsumerStallGateTests.cs

using NATS.Server;
using NATS.Server.Auth;
using Shouldly;

namespace NATS.Server.Core.Tests;

/// <summary>
/// Tests for <see cref="SlowConsumerTracker"/> and <see cref="Account"/> slow-consumer counters.
/// Go reference: server/client.go — handleSlowConsumer, markConnAsSlow.
/// </summary>
public class SlowConsumerStallGateTests
{
    // ── SlowConsumerTracker ────────────────────────────────────────────────────

    [Fact]
    public void RecordSlowConsumer_increments_total_count()
    {
        var tracker = new SlowConsumerTracker();

        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Router);

        tracker.TotalCount.ShouldBe(3L);
    }

    [Fact]
    public void RecordSlowConsumer_increments_per_kind_count()
    {
        var tracker = new SlowConsumerTracker();

        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Gateway);

        tracker.GetCount(ClientKind.Client).ShouldBe(2L);
        tracker.GetCount(ClientKind.Gateway).ShouldBe(1L);
    }

    [Fact]
    public void GetCount_returns_zero_for_unrecorded_kind()
    {
        var tracker = new SlowConsumerTracker();

        tracker.GetCount(ClientKind.Leaf).ShouldBe(0L);
    }

    [Fact]
    public void Multiple_kinds_tracked_independently()
    {
        var tracker = new SlowConsumerTracker();

        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Router);
        tracker.RecordSlowConsumer(ClientKind.Gateway);
        tracker.RecordSlowConsumer(ClientKind.Leaf);

        tracker.GetCount(ClientKind.Client).ShouldBe(1L);
        tracker.GetCount(ClientKind.Router).ShouldBe(1L);
        tracker.GetCount(ClientKind.Gateway).ShouldBe(1L);
        tracker.GetCount(ClientKind.Leaf).ShouldBe(1L);
        tracker.GetCount(ClientKind.System).ShouldBe(0L);
        tracker.TotalCount.ShouldBe(4L);
    }

    [Fact]
    public void OnThresholdExceeded_fires_when_threshold_reached()
    {
        var tracker = new SlowConsumerTracker(threshold: 3);
        ClientKind? firedKind = null;
        tracker.OnThresholdExceeded(k => firedKind = k);

        tracker.RecordSlowConsumer(ClientKind.Client);
        firedKind.ShouldBeNull();

        tracker.RecordSlowConsumer(ClientKind.Client);
        firedKind.ShouldBeNull();

        tracker.RecordSlowConsumer(ClientKind.Router);
        firedKind.ShouldBe(ClientKind.Router);
    }

    [Fact]
    public void OnThresholdExceeded_not_fired_below_threshold()
    {
        var tracker = new SlowConsumerTracker(threshold: 5);
        var fired = false;
        tracker.OnThresholdExceeded(_ => fired = true);

        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Client);

        fired.ShouldBeFalse();
    }

    [Fact]
    public void Reset_clears_all_counters()
    {
        var tracker = new SlowConsumerTracker();

        tracker.RecordSlowConsumer(ClientKind.Client);
        tracker.RecordSlowConsumer(ClientKind.Router);
        tracker.TotalCount.ShouldBe(2L);

        tracker.Reset();

        tracker.TotalCount.ShouldBe(0L);
        tracker.GetCount(ClientKind.Client).ShouldBe(0L);
        tracker.GetCount(ClientKind.Router).ShouldBe(0L);
    }

    // ── Account slow-consumer counters ────────────────────────────────────────

    [Fact]
    public void Account_IncrementSlowConsumers_tracks_count()
    {
        var account = new Account("test-account");

        account.SlowConsumerCount.ShouldBe(0L);

        account.IncrementSlowConsumers();
        account.IncrementSlowConsumers();
        account.IncrementSlowConsumers();

        account.SlowConsumerCount.ShouldBe(3L);
    }

    [Fact]
    public void Account_ResetSlowConsumerCount_clears()
    {
        var account = new Account("test-account");

        account.IncrementSlowConsumers();
        account.IncrementSlowConsumers();
        account.SlowConsumerCount.ShouldBe(2L);

        account.ResetSlowConsumerCount();

        account.SlowConsumerCount.ShouldBe(0L);
    }

    [Fact]
    public void Thread_safety_concurrent_increments()
    {
        var tracker = new SlowConsumerTracker();

        Parallel.For(0, 1000, _ => tracker.RecordSlowConsumer(ClientKind.Client));

        tracker.TotalCount.ShouldBe(1000L);
        tracker.GetCount(ClientKind.Client).ShouldBe(1000L);
    }
}