natsdotnet/tests/NATS.Server.Tests/Internal/TimeHashWheel/HashWheelTests.cs

// Go reference: server/thw/thw_test.go

using NATS.Server.Internal.TimeHashWheel;

namespace NATS.Server.Tests.Internal.TimeHashWheel;

public class HashWheelTests
{
    /// <summary>
    /// Helper to produce nanosecond timestamps relative to a base, matching
    /// the Go test pattern of now.Add(N * time.Second).UnixNano().
    /// </summary>
    private static long NowNanos() => DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000;

    private static long SecondsToNanos(long seconds) => seconds * 1_000_000_000;

    // Go: TestHashWheelBasics server/thw/thw_test.go:22
    [Fact]
    public void Basics_AddRemoveCount()
    {
        var hw = new HashWheel();
        var now = NowNanos();

        // Add a sequence.
        ulong seq = 1;
        var expires = now + SecondsToNanos(5);
        hw.Add(seq, expires);
        hw.Count.ShouldBe(1UL);

        // Try to remove non-existent sequence.
        hw.Remove(999, expires).ShouldBeFalse();
        hw.Count.ShouldBe(1UL);

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

        // Verify it's gone.
        hw.Remove(seq, expires).ShouldBeFalse();
        hw.Count.ShouldBe(0UL);
    }

    // Go: TestHashWheelUpdate server/thw/thw_test.go:44
    [Fact]
    public void Update_ChangesExpiration()
    {
        var hw = new HashWheel();
        var now = NowNanos();
        var oldExpires = now + SecondsToNanos(5);
        var newExpires = now + SecondsToNanos(10);

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

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

        // Verify old expiration is gone.
        hw.Remove(1, oldExpires).ShouldBeFalse();
        hw.Count.ShouldBe(1UL);

        // Verify new expiration exists.
        hw.Remove(1, newExpires).ShouldBeTrue();
        hw.Count.ShouldBe(0UL);
    }

    // Go: TestHashWheelExpiration server/thw/thw_test.go:67
    [Fact]
    public void Expiration_FiresCallbackForExpired()
    {
        var hw = new HashWheel();
        var now = NowNanos();

        // Add sequences with different expiration times.
        var seqs = new Dictionary<ulong, long>
        {
            [1] = now - SecondsToNanos(1), // Already expired
            [2] = now + SecondsToNanos(1), // Expires soon
            [3] = now + SecondsToNanos(10), // Expires later
            [4] = now + SecondsToNanos(60), // Expires much later
        };

        foreach (var (seq, expires) in seqs)
        {
            hw.Add(seq, expires);
        }

        hw.Count.ShouldBe((ulong)seqs.Count);

        // Process expired tasks using internal method with explicit "now" timestamp.
        var expired = new Dictionary<ulong, bool>();
        hw.ExpireTasksInternal(now, (seq, _) =>
        {
            expired[seq] = true;
            return true;
        });

        // Verify only sequence 1 expired.
        expired.Count.ShouldBe(1);
        expired.ShouldContainKey(1UL);
        hw.Count.ShouldBe(3UL);
    }

    // Go: TestHashWheelManualExpiration server/thw/thw_test.go:97
    [Fact]
    public void ManualExpiration_SpecificTime()
    {
        var hw = new HashWheel();
        var now = NowNanos();

        for (ulong seq = 1; seq <= 4; seq++)
        {
            hw.Add(seq, now);
        }

        hw.Count.ShouldBe(4UL);

        // Loop over expired multiple times, but without removing them.
        var expired = new Dictionary<ulong, ulong>();
        for (ulong i = 0; i <= 1; i++)
        {
            hw.ExpireTasksInternal(now, (seq, _) =>
            {
                if (!expired.TryGetValue(seq, out var count))
                {
                    count = 0;
                }

                expired[seq] = count + 1;
                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);
        }

        // Only remove even sequences.
        for (ulong i = 0; i <= 1; i++)
        {
            hw.ExpireTasksInternal(now, (seq, _) =>
            {
                if (!expired.TryGetValue(seq, out var count))
                {
                    count = 0;
                }

                expired[seq] = count + 1;
                return seq % 2 == 0;
            });

            // Verify even sequences are removed.
            expired[1].ShouldBe(3 + i);
            expired[2].ShouldBe(3UL);
            expired[3].ShouldBe(3 + i);
            expired[4].ShouldBe(3UL);
            hw.Count.ShouldBe(2UL);
        }

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

    // Go: TestHashWheelExpirationLargerThanWheel server/thw/thw_test.go:143
    [Fact]
    public void LargerThanWheel_HandlesWrapAround()
    {
        var hw = new HashWheel();

        // Add sequences such that they can be expired immediately.
        var seqs = new Dictionary<ulong, long>
        {
            [1] = 0,
            [2] = SecondsToNanos(1),
        };

        foreach (var (seq, expires) in seqs)
        {
            hw.Add(seq, expires);
        }

        hw.Count.ShouldBe(2UL);

        // Pick a timestamp such that the expiration needs to wrap around the whole wheel.
        // Go: now := int64(time.Second) * wheelMask
        var now = SecondsToNanos(1) * HashWheel.WheelSize - SecondsToNanos(1);

        // Process expired tasks.
        var expired = new Dictionary<ulong, bool>();
        hw.ExpireTasksInternal(now, (seq, _) =>
        {
            expired[seq] = true;
            return true;
        });

        // Verify both sequences are expired.
        expired.Count.ShouldBe(2);
        hw.Count.ShouldBe(0UL);
    }

    // Go: TestHashWheelNextExpiration server/thw/thw_test.go:171
    [Fact]
    public void NextExpiration_FindsEarliest()
    {
        var hw = new HashWheel();
        var now = NowNanos();

        // Add sequences with different expiration times.
        var seqs = new Dictionary<ulong, long>
        {
            [1] = now + SecondsToNanos(5),
            [2] = now + SecondsToNanos(3), // Earliest
            [3] = now + SecondsToNanos(10),
        };

        foreach (var (seq, expires) in seqs)
        {
            hw.Add(seq, expires);
        }

        hw.Count.ShouldBe((ulong)seqs.Count);

        // Test GetNextExpiration.
        var nextExternalTick = now + SecondsToNanos(6);
        // Should return sequence 2's expiration.
        hw.GetNextExpiration(nextExternalTick).ShouldBe(seqs[2]);

        // Test with empty wheel.
        var empty = new HashWheel();
        empty.GetNextExpiration(now + SecondsToNanos(1)).ShouldBe(long.MaxValue);
    }

    // Go: TestHashWheelStress server/thw/thw_test.go:197
    [Fact]
    public void Stress_ConcurrentAddRemove()
    {
        var hw = new HashWheel();
        var now = NowNanos();
        const int numSequences = 100_000;

        // Add many sequences.
        for (var seq = 0; seq < numSequences; seq++)
        {
            var expires = now + SecondsToNanos(seq);
            hw.Add((ulong)seq, expires);
        }

        // Update many sequences (every other one).
        for (var seq = 0; seq < numSequences; seq += 2)
        {
            var oldExpires = now + SecondsToNanos(seq);
            var newExpires = now + SecondsToNanos(seq + numSequences);
            hw.Update((ulong)seq, oldExpires, newExpires);
        }

        // Remove odd-numbered sequences.
        for (var seq = 1; seq < numSequences; seq += 2)
        {
            var expires = now + SecondsToNanos(seq);
            hw.Remove((ulong)seq, expires).ShouldBeTrue();
        }

        // After updates and removals, only half remain (the even ones with updated expiration).
        hw.Count.ShouldBe((ulong)(numSequences / 2));
    }

    // Go: TestHashWheelEncodeDecode server/thw/thw_test.go:222
    [Fact]
    public void EncodeDecode_RoundTrips()
    {
        var hw = new HashWheel();
        var now = NowNanos();
        const int numSequences = 100_000;

        // Add many sequences.
        for (var seq = 0; seq < numSequences; seq++)
        {
            var expires = now + SecondsToNanos(seq);
            hw.Add((ulong)seq, expires);
        }

        var encoded = hw.Encode(12345);
        encoded.Length.ShouldBeGreaterThan(17); // Bigger than just the header.

        var nhw = new HashWheel();
        var (highSeq, bytesRead) = nhw.Decode(encoded);
        highSeq.ShouldBe(12345UL);
        bytesRead.ShouldBe(encoded.Length);
        hw.GetNextExpiration(long.MaxValue).ShouldBe(nhw.GetNextExpiration(long.MaxValue));

        // Verify all slots match.
        for (var s = 0; s < HashWheel.WheelSize; s++)
        {
            var slot = hw.Wheel[s];
            var nslot = nhw.Wheel[s];

            if (slot is null)
            {
                nslot.ShouldBeNull();
                continue;
            }

            nslot.ShouldNotBeNull();
            slot.Lowest.ShouldBe(nslot!.Lowest);
            slot.Entries.Count.ShouldBe(nslot.Entries.Count);

            foreach (var (seq, ts) in slot.Entries)
            {
                nslot.Entries.ShouldContainKey(seq);
                nslot.Entries[seq].ShouldBe(ts);
            }
        }
    }
}