wwtools/mbproxy/tests/Mbproxy.Tests/Proxy/Multiplexing/TxIdAllocatorTests.cs

using Mbproxy.Proxy.Multiplexing;
using Shouldly;
using Xunit;

namespace Mbproxy.Tests.Proxy.Multiplexing;

/// <summary>
/// Unit tests for <see cref="TxIdAllocator"/>. Pure logic — no I/O.
/// </summary>
[Trait("Category", "Unit")]
public sealed class TxIdAllocatorTests
{
    [Fact]
    public void Allocate_FromEmpty_Returns_NextSequential()
    {
        var alloc = new TxIdAllocator();

        alloc.TryAllocate(out ushort a).ShouldBeTrue();
        alloc.TryAllocate(out ushort b).ShouldBeTrue();
        alloc.TryAllocate(out ushort c).ShouldBeTrue();

        a.ShouldBe((ushort)0);
        b.ShouldBe((ushort)1);
        c.ShouldBe((ushort)2);
        alloc.InFlightCount.ShouldBe(3);
    }

    [Fact]
    public void Allocate_AfterRelease_Reuses_FreedId()
    {
        var alloc = new TxIdAllocator();

        alloc.TryAllocate(out ushort a).ShouldBeTrue();
        alloc.TryAllocate(out ushort b).ShouldBeTrue();
        alloc.TryAllocate(out ushort c).ShouldBeTrue();

        // Release the middle slot and allocate again. The next allocation should advance
        // forward from the cursor (3) and not re-use 1 until the cursor wraps and finds it free.
        alloc.Release(b);
        alloc.InFlightCount.ShouldBe(2);

        alloc.TryAllocate(out ushort d).ShouldBeTrue();
        d.ShouldBe((ushort)3, "allocator advances the cursor; freed slot 1 reuses only after wrap");
    }

    [Fact]
    public void Allocate_AllocatesEveryUshort_BeforeWrapping()
    {
        var alloc = new TxIdAllocator();
        var seen = new HashSet<ushort>();

        for (int i = 0; i < 65536; i++)
        {
            alloc.TryAllocate(out ushort id).ShouldBeTrue($"allocation {i} should succeed");
            seen.Add(id).ShouldBeTrue($"id {id} should be unique across the full 0..65535 sweep");
        }

        seen.Count.ShouldBe(65536);
        alloc.InFlightCount.ShouldBe(65536);
    }

    [Fact]
    public void Allocate_WrapsCorrectly_After0xFFFF()
    {
        var alloc = new TxIdAllocator();

        // Allocate every slot then release slot 5.
        for (int i = 0; i < 65536; i++)
            alloc.TryAllocate(out _).ShouldBeTrue();

        alloc.Release(5);

        // Next allocation should find slot 5 after the cursor wraps.
        alloc.TryAllocate(out ushort id).ShouldBeTrue();
        id.ShouldBe((ushort)5);
    }

    [Fact]
    public void Allocate_WhenSaturated_ReturnsFalse_DoesNotThrow()
    {
        var alloc = new TxIdAllocator();
        for (int i = 0; i < 65536; i++)
            alloc.TryAllocate(out _).ShouldBeTrue();

        alloc.TryAllocate(out ushort id).ShouldBeFalse("saturated allocator must refuse cleanly");
        id.ShouldBe((ushort)0);
    }

    [Fact]
    public void Release_OfNonAllocated_IsNoOp()
    {
        var alloc = new TxIdAllocator();

        alloc.TryAllocate(out ushort a).ShouldBeTrue();
        // a == 0. Release a slot that was never allocated.
        alloc.Release(42);
        alloc.InFlightCount.ShouldBe(1, "releasing a non-allocated id must not decrement the count");
    }

    [Fact]
    public async Task Concurrent_AllocateRelease_NoDuplicateIds_Under_Parallel_Stress()
    {
        var alloc = new TxIdAllocator();
        const int taskCount = 100;
        const int opsPerTask = 1000;

        // Each task allocates and immediately releases its id, hammering the lock.
        // If allocate ever hands out a duplicate, two tasks would see the same id.
        var observed = new System.Collections.Concurrent.ConcurrentDictionary<int, byte>();

        await Task.WhenAll(Enumerable.Range(0, taskCount).Select(_ => Task.Run(() =>
        {
            for (int i = 0; i < opsPerTask; i++)
            {
                if (!alloc.TryAllocate(out ushort id))
                    continue;
                // Add a unique tag to detect a duplicate live id.
                observed.TryAdd(id, 1).ShouldBeTrue();
                observed.TryRemove(id, out byte _);
                alloc.Release(id);
            }
        })));

        alloc.InFlightCount.ShouldBe(0, "every allocation was released; count must be back to 0");
    }

    [Fact]
    public void WrapCount_IncrementsOnEachFullWrap()
    {
        var alloc = new TxIdAllocator();
        alloc.WrapCount.ShouldBe(0);

        // First sweep: 65536 allocations bring the cursor from 0 back to 0 → one wrap.
        for (int i = 0; i < 65536; i++)
            alloc.TryAllocate(out _).ShouldBeTrue();

        alloc.WrapCount.ShouldBe(1);

        // Release everything, then sweep again: should bump WrapCount to 2.
        for (ushort i = 0; ; i++)
        {
            alloc.Release(i);
            if (i == 65535) break;
        }
        for (int i = 0; i < 65536; i++)
            alloc.TryAllocate(out _).ShouldBeTrue();
        alloc.WrapCount.ShouldBe(2);
    }
}