natsnet/dotnet/tests/ZB.MOM.NatsNet.Server.Tests/Internal/DataStructures/HashWheelTests.cs

using Shouldly;
using ZB.MOM.NatsNet.Server.Internal.DataStructures;

namespace ZB.MOM.NatsNet.Server.Tests.Internal.DataStructures;

/// <summary>
/// Tests for <see cref="HashWheel"/>, mirroring thw_test.go (functional tests only;
/// benchmarks are omitted as they require BenchmarkDotNet).
/// </summary>
public sealed class HashWheelTests
{
    private static readonly long Second = 1_000_000_000L; // nanoseconds

    [Fact]
    public void HashWheelBasics_ShouldSucceed()
    {
        // Mirror: TestHashWheelBasics
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
        var seq = 1UL;
        var expires = now + 5 * Second;

        hw.Add(seq, expires);
        hw.Count.ShouldBe(1UL);

        // Remove non-existent sequence.
        Should.Throw<InvalidOperationException>(() => hw.Remove(999, expires));
        hw.Count.ShouldBe(1UL);

        // Remove properly.
        hw.Remove(seq, expires);
        hw.Count.ShouldBe(0UL);

        // Already gone.
        Should.Throw<InvalidOperationException>(() => hw.Remove(seq, expires));
        hw.Count.ShouldBe(0UL);
    }

    [Fact]
    public void HashWheelUpdate_ShouldSucceed()
    {
        // Mirror: TestHashWheelUpdate
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
        var oldExpires = now + 5 * Second;
        var newExpires = now + 10 * Second;

        hw.Add(1, oldExpires);
        hw.Count.ShouldBe(1UL);

        hw.Update(1, oldExpires, newExpires);
        hw.Count.ShouldBe(1UL);

        // Old position gone.
        Should.Throw<InvalidOperationException>(() => hw.Remove(1, oldExpires));
        hw.Count.ShouldBe(1UL);

        // New position exists.
        hw.Remove(1, newExpires);
        hw.Count.ShouldBe(0UL);
    }

    [Fact]
    public void HashWheelExpiration_ShouldExpireOnly_AlreadyExpired()
    {
        // Mirror: TestHashWheelExpiration
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;

        var seqs = new Dictionary<ulong, long>
        {
            [1] = now - 1 * Second,  // already expired
            [2] = now + 1 * Second,
            [3] = now + 10 * Second,
            [4] = now + 60 * Second,
        };
        foreach (var (s, exp) in seqs)
            hw.Add(s, exp);
        hw.Count.ShouldBe((ulong)seqs.Count);

        var expired = new HashSet<ulong>();
        hw.ExpireTasksInternal(now, (s, _) => { expired.Add(s); return true; });

        expired.Count.ShouldBe(1);
        expired.ShouldContain(1UL);
        hw.Count.ShouldBe(3UL);
    }

    [Fact]
    public void HashWheelManualExpiration_ShouldRespectCallbackReturn()
    {
        // Mirror: TestHashWheelManualExpiration
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;

        for (var s = 1UL; s <= 4; s++)
            hw.Add(s, now);
        hw.Count.ShouldBe(4UL);

        // Iterate without removing.
        var expired = new Dictionary<ulong, ulong>();
        for (var i = 0UL; i <= 1; i++)
        {
            hw.ExpireTasksInternal(now, (s, _) => { expired.TryAdd(s, 0); expired[s]++; return false; });

            expired.Count.ShouldBe(4);
            expired[1].ShouldBe(1 + i);
            expired[2].ShouldBe(1 + i);
            expired[3].ShouldBe(1 + i);
            expired[4].ShouldBe(1 + i);
            hw.Count.ShouldBe(4UL);
        }

        // Remove only even sequences.
        for (var i = 0UL; i <= 1; i++)
        {
            hw.ExpireTasksInternal(now, (s, _) => { expired.TryAdd(s, 0); expired[s]++; return s % 2 == 0; });

            expired[1].ShouldBe(3 + i);
            expired[2].ShouldBe(3UL);
            expired[3].ShouldBe(3 + i);
            expired[4].ShouldBe(3UL);
            hw.Count.ShouldBe(2UL);
        }

        // Manually remove remaining.
        hw.Remove(1, now);
        hw.Remove(3, now);
        hw.Count.ShouldBe(0UL);
    }

    [Fact]
    public void HashWheelExpirationLargerThanWheel_ShouldExpireAll()
    {
        // Mirror: TestHashWheelExpirationLargerThanWheel
        const int WheelMask = (1 << 12) - 1;
        var hw = HashWheel.NewHashWheel();

        hw.Add(1, 0);
        hw.Add(2, Second);
        hw.Count.ShouldBe(2UL);

        // Timestamp large enough to wrap the entire wheel.
        var nowWrapped = Second * WheelMask;

        var expired = new HashSet<ulong>();
        hw.ExpireTasksInternal(nowWrapped, (s, _) => { expired.Add(s); return true; });

        expired.Count.ShouldBe(2);
        hw.Count.ShouldBe(0UL);
    }

    [Fact]
    public void HashWheelNextExpiration_ShouldReturnEarliest()
    {
        // Mirror: TestHashWheelNextExpiration
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;

        var seqs = new Dictionary<ulong, long>
        {
            [1] = now + 5 * Second,
            [2] = now + 3 * Second, // earliest
            [3] = now + 10 * Second,
        };
        foreach (var (s, exp) in seqs)
            hw.Add(s, exp);

        var tick = now + 6 * Second;
        hw.GetNextExpiration(tick).ShouldBe(seqs[2]);

        var empty = HashWheel.NewHashWheel();
        empty.GetNextExpiration(now + Second).ShouldBe(long.MaxValue);
    }

    [Fact]
    public void HashWheelStress_ShouldHandleLargeScale()
    {
        // Mirror: TestHashWheelStress
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
        const int numSeqs = 100_000;

        for (var seq = 0; seq < numSeqs; seq++)
        {
            var exp = now + (long)seq * Second;
            hw.Add((ulong)seq, exp);
        }

        // Update even sequences.
        for (var seq = 0; seq < numSeqs; seq += 2)
        {
            var oldExp = now + (long)seq * Second;
            var newExp = now + (long)(seq + numSeqs) * Second;
            hw.Update((ulong)seq, oldExp, newExp);
        }

        // Remove odd sequences.
        for (var seq = 1; seq < numSeqs; seq += 2)
        {
            var exp = now + (long)seq * Second;
            hw.Remove((ulong)seq, exp);
        }
    }

    [Fact]
    public void HashWheelEncodeDecode_ShouldRoundTrip()
    {
        // Mirror: TestHashWheelEncodeDecode
        var hw = HashWheel.NewHashWheel();
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
        const int numSeqs = 100_000;

        for (var seq = 0; seq < numSeqs; seq++)
        {
            var exp = now + (long)seq * Second;
            hw.Add((ulong)seq, exp);
        }

        var b = hw.Encode(12345);
        b.Length.ShouldBeGreaterThan(17);

        var nhw = HashWheel.NewHashWheel();
        var stamp = nhw.Decode(b);
        stamp.ShouldBe(12345UL);

        // Lowest expiry should match.
        hw.GetNextExpiration(long.MaxValue).ShouldBe(nhw.GetNextExpiration(long.MaxValue));

        // Verify all entries transferred by removing them from nhw.
        for (var seq = 0; seq < numSeqs; seq++)
        {
            var exp = now + (long)seq * Second;
            nhw.Remove((ulong)seq, exp); // throws if missing
        }
        nhw.Count.ShouldBe(0UL);
    }
}