natsdotnet/tests/NATS.Server.Tests/Raft/RaftElectionJitterTests.cs

using NATS.Server.Raft;

namespace NATS.Server.Tests.Raft;

/// <summary>
/// Tests for randomized election timeout jitter in RaftNode.
/// Jitter prevents synchronized elections after network partitions (split-vote avoidance).
/// Go reference: raft.go resetElectionTimeout — uses rand.Int63n to jitter election timeout.
/// </summary>
public class RaftElectionJitterTests : IDisposable
{
    private readonly List<RaftNode> _nodesToDispose = [];

    public void Dispose()
    {
        foreach (var node in _nodesToDispose)
            node.Dispose();
    }

    private RaftNode CreateTrackedNode(string id, Random? random = null)
    {
        var node = new RaftNode(id, random: random);
        _nodesToDispose.Add(node);
        return node;
    }

    [Fact]
    public void RandomizedElectionTimeout_within_range()
    {
        // Go reference: raft.go resetElectionTimeout — timeout is always within [min, max).
        var node = CreateTrackedNode("n1");
        node.ElectionTimeoutMinMs = 150;
        node.ElectionTimeoutMaxMs = 300;

        for (var i = 0; i < 100; i++)
        {
            var timeout = node.RandomizedElectionTimeout();
            timeout.TotalMilliseconds.ShouldBeGreaterThanOrEqualTo(150);
            timeout.TotalMilliseconds.ShouldBeLessThan(300);
        }
    }

    [Fact]
    public void RandomizedElectionTimeout_varies()
    {
        // Multiple calls must produce at least 2 distinct values in 10 samples,
        // confirming that the timeout is not fixed.
        // Go reference: raft.go resetElectionTimeout — jitter ensures nodes don't all
        // timeout simultaneously after a partition.
        var node = CreateTrackedNode("n1");
        node.ElectionTimeoutMinMs = 150;
        node.ElectionTimeoutMaxMs = 300;

        var samples = Enumerable.Range(0, 10)
            .Select(_ => node.RandomizedElectionTimeout().TotalMilliseconds)
            .ToList();

        var distinct = samples.Distinct().Count();
        distinct.ShouldBeGreaterThanOrEqualTo(2);
    }

    [Fact]
    public void RandomizedElectionTimeout_uses_total_milliseconds()
    {
        // Verifies that TotalMilliseconds (not .Milliseconds component) gives the full
        // value. For timeouts >= 1000 ms, .Milliseconds would return only the sub-second
        // component (0-999), while TotalMilliseconds returns the complete value.
        // This test uses a range that straddles 1000 ms to expose the bug if present.
        // Go reference: raft.go resetElectionTimeout uses full duration, not sub-second part.
        var node = CreateTrackedNode("n1");
        node.ElectionTimeoutMinMs = 1200;
        node.ElectionTimeoutMaxMs = 1500;

        for (var i = 0; i < 50; i++)
        {
            var timeout = node.RandomizedElectionTimeout();

            // TotalMilliseconds must be in [1200, 1500)
            timeout.TotalMilliseconds.ShouldBeGreaterThanOrEqualTo(1200);
            timeout.TotalMilliseconds.ShouldBeLessThan(1500);

            // The .Milliseconds property would return only the 0-999 sub-second
            // component. If the implementation incorrectly used .Milliseconds,
            // these values would be wrong (< 1200). Verify TotalMilliseconds is >= 1200.
            ((int)timeout.TotalMilliseconds).ShouldBeGreaterThanOrEqualTo(1200);
        }
    }

    [Fact]
    public void ResetElectionTimeout_randomizes_each_time()
    {
        // Successive calls to ResetElectionTimeout should obtain fresh random intervals.
        // We observe the timer change indirectly by injecting a deterministic Random
        // and confirming it gets called on each reset.
        // Go reference: raft.go resetElectionTimeout — called on every heartbeat and
        // leader append to re-arm with a fresh random deadline.
        var callCount = 0;
        var deterministicRandom = new CountingRandom(() => callCount++);

        var node = CreateTrackedNode("n1", deterministicRandom);
        node.ElectionTimeoutMinMs = 150;
        node.ElectionTimeoutMaxMs = 300;
        node.StartElectionTimer();

        var countAfterStart = callCount;

        node.ResetElectionTimeout();
        node.ResetElectionTimeout();
        node.ResetElectionTimeout();

        // Each ResetElectionTimeout + StartElectionTimer must have called Next() at least once each
        callCount.ShouldBeGreaterThan(countAfterStart);

        node.StopElectionTimer();
    }

    [Fact]
    public void Different_nodes_get_different_timeouts()
    {
        // Two nodes created with the default shared Random should produce at least some
        // distinct timeout values across multiple samples (probabilistic).
        // Go reference: raft.go — each server picks an independent random timeout so
        // they do not all call elections at exactly the same moment.
        var node1 = CreateTrackedNode("n1");
        var node2 = CreateTrackedNode("n2");

        node1.ElectionTimeoutMinMs = node2.ElectionTimeoutMinMs = 150;
        node1.ElectionTimeoutMaxMs = node2.ElectionTimeoutMaxMs = 300;

        var samples1 = Enumerable.Range(0, 20)
            .Select(_ => node1.RandomizedElectionTimeout().TotalMilliseconds)
            .ToList();
        var samples2 = Enumerable.Range(0, 20)
            .Select(_ => node2.RandomizedElectionTimeout().TotalMilliseconds)
            .ToList();

        // The combined set of 40 values must contain at least 2 distinct values
        // (if both nodes returned the exact same value every time, that would
        // indicate no jitter at all).
        var allValues = samples1.Concat(samples2).Distinct().Count();
        allValues.ShouldBeGreaterThanOrEqualTo(2);
    }

    /// <summary>
    /// A <see cref="Random"/> subclass that invokes a callback on every call to
    /// <see cref="Next(int, int)"/>, allowing tests to count how many times the
    /// RaftNode requests a new random value.
    /// </summary>
    private sealed class CountingRandom(Action onNext) : Random
    {
        public override int Next(int minValue, int maxValue)
        {
            onNext();
            return base.Next(minValue, maxValue);
        }
    }
}