lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests/DeviceSessionPoolTests.cs

using Microsoft.Extensions.Logging;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Cip;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Eip;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests.Cip;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests.Fakes;
namespace ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Tests;

[Collection(CipConnectionStatics.Name)]
public sealed class DeviceSessionPoolTests
{
    private static readonly TimeSpan Guard = TimeSpan.FromSeconds(10);
    private static readonly byte[] ReadTag = [0x4C, 0x02, 0x91, 0x03, 0x54, 0x61, 0x67, 0x00, 0x01, 0x00];
    private static CancellationToken TestCt => TestContext.Current.CancellationToken;

    private static ConnectionOptions Options(int size = 4002) => new() { ConnectionSize = size, Timeout = TimeSpan.FromSeconds(2) };

    [Fact]
    public void Device_keys_compare_host_and_plc_attribute_case_insensitively()
    {
        var a = new DeviceKey("Press-07.plant.local", 44818, "1,0", "ControlLogix");
        var b = new DeviceKey("PRESS-07.PLANT.LOCAL", 44818, "1,0", "controllogix");
        a.ShouldBe(b);
        a.GetHashCode().ShouldBe(b.GetHashCode());
        a.ShouldNotBe(a with { CipPath = "1,1" });
        a.ShouldNotBe(a with { Port = 2222 });
    }

    [Fact]
    public async Task Leases_for_one_key_share_one_scheduler()
    {
        var pool = new DeviceSessionPool(new ManualTimeProvider());
        var key = new DeviceKey("plc1", 44818, "1,0", "ControlLogix");

        await using var a = pool.Acquire(key, Options());
        await using var b = pool.Acquire(key with { Host = "PLC1" }, Options());
        await using var other = pool.Acquire(key with { CipPath = "1,2" }, Options());

        b.Scheduler.ShouldBeSameAs(a.Scheduler);
        b.Channel.ShouldBeSameAs(a.Channel);
        other.Scheduler.ShouldNotBeSameAs(a.Scheduler);
        other.Channel.ShouldNotBeSameAs(a.Channel);
        pool.EntryCount.ShouldBe(2);
    }

    [Fact]
    public async Task A_second_acquire_with_different_options_reuses_the_first_and_logs_once()
    {
        var logger = new CapturingLogger();
        var pool = new DeviceSessionPool(new ManualTimeProvider(), logger);
        var key = new DeviceKey("plc1", 44818, "1,0", "ControlLogix");
        var first = Options(4002);
        var second = Options(504) with { AllowPacking = true };

        await using var a = pool.Acquire(key, first);
        await using var b = pool.Acquire(key, second);
        await using var c = pool.Acquire(key, second);
        await using var d = pool.Acquire(key, first);

        ((CipConnection)b.Channel).Options.ShouldBe(first);
        ((CipConnection)c.Channel).Options.ShouldBe(first);
        logger.Warnings.ShouldBe(1);
    }

    [Fact]
    public async Task Last_lease_dispose_sends_forward_close_then_unregister()
    {
        await using var target = new FakeCipTarget();
        var pool = new DeviceSessionPool(new ManualTimeProvider());
        var key = new DeviceKey("127.0.0.1", target.Port, "1,0", "ControlLogix");
        var a = pool.Acquire(key, Options());
        var b = pool.Acquire(key, Options());
        (await a.Scheduler.SendAsync(ReadTag, TestCt)).GeneralStatus.ShouldBe((byte)0);

        await a.DisposeAsync();
        await a.DisposeAsync();   // idempotent: must not release b's reference
        target.ForwardCloses.ShouldBe(0);
        (await b.Scheduler.SendAsync(ReadTag, TestCt)).GeneralStatus.ShouldBe((byte)0);

        await b.DisposeAsync();
        await target.Server.WaitForFrameAsync(f => FakeEipServer.Command(f) == EncapCommand.UnRegisterSession, Guard);

        target.ForwardCloses.ShouldBe(1);
        var frames = target.Server.Received.ToList();
        var close = frames.FindIndex(f => FakeEipServer.Command(f) == EncapCommand.SendRRData
                                          && FakeCipTarget.IsForwardClose(FakeEipServer.RequestCip(f)));
        var unregister = frames.FindIndex(f => FakeEipServer.Command(f) == EncapCommand.UnRegisterSession);
        close.ShouldBeGreaterThanOrEqualTo(0);
        close.ShouldBeLessThan(unregister);
        // The Forward Close names the connection that was opened.
        target.ForwardCloseRequests.Single().AsSpan(8, 8).ToArray()
            .ShouldBe(target.ForwardOpenRequests.Single().AsSpan(16, 8).ToArray());
        pool.EntryCount.ShouldBe(0);
        await Should.ThrowAsync<ObjectDisposedException>(() => b.Scheduler.SendAsync(ReadTag, TestCt));
    }

    [Fact]
    public async Task A_later_acquire_builds_new_instances()
    {
        var pool = new DeviceSessionPool(new ManualTimeProvider());
        var key = new DeviceKey("plc1", 44818, "1,0", "ControlLogix");
        var first = pool.Acquire(key, Options());
        var scheduler = first.Scheduler;
        var channel = first.Channel;
        await first.DisposeAsync();

        await using var second = pool.Acquire(key, Options());

        second.Scheduler.ShouldNotBeSameAs(scheduler);
        second.Channel.ShouldNotBeSameAs(channel);
    }

    [Fact]
    public async Task A_reconnect_after_the_last_release_is_still_throttled()
    {
        await using var target = new FakeCipTarget();
        var time = new ManualTimeProvider();
        var pool = new DeviceSessionPool(time);
        var key = new DeviceKey("127.0.0.1", target.Port, "1,0", "ControlLogix");

        await using (var first = pool.Acquire(key, Options()))
            await first.Scheduler.SendAsync(ReadTag, TestCt);                    // Forward Open at fake t0

        await using var second = pool.Acquire(key, Options());
        var send = second.Scheduler.SendAsync(ReadTag, TestCt);
        await time.WaitForTimerDueAtAsync(time.Start + TimeSpan.FromSeconds(1), Guard);
        target.ForwardOpens.ShouldBe(1);

        time.Advance(TimeSpan.FromSeconds(1));
        await send.WaitAsync(Guard, TestCt);
        target.ForwardOpens.ShouldBe(2);
    }

    [Fact]
    public async Task Concurrent_acquire_and_release_is_safe()
    {
        var pool = new DeviceSessionPool(new ManualTimeProvider());
        var keys = Enumerable.Range(0, 3).Select(i => new DeviceKey($"plc{i}", 44818, "1,0", "ControlLogix")).ToArray();
        var schedulers = new System.Collections.Concurrent.ConcurrentDictionary<(int, RequestScheduler), byte>();

        var tasks = Enumerable.Range(0, 100).Select(i => Task.Run(async () =>
        {
            var lease = pool.Acquire(keys[i % 3], Options());
            schedulers.TryAdd((i % 3, lease.Scheduler), 0);
            await Task.Yield();
            await lease.DisposeAsync();
        }, TestCt)).ToArray();
        await Task.WhenAll(tasks).WaitAsync(Guard, TestCt);

        pool.EntryCount.ShouldBe(0);
        // Every lease got a scheduler for its own key only, never another key's.
        schedulers.Keys.GroupBy(k => k.Item2).ShouldAllBe(g => g.Select(k => k.Item1).Distinct().Count() == 1);
    }

    [Fact]
    public async Task Device_keys_ignore_whitespace_in_the_cip_path()
    {
        var a = new DeviceKey("plc1", 44818, "1,0", "ControlLogix");
        var b = new DeviceKey("plc1", 44818, " 1, 0 ", "ControlLogix");
        b.CipPath.ShouldBe("1,0");
        b.ShouldBe(a);
        b.GetHashCode().ShouldBe(a.GetHashCode());
        (a with { CipPath = "1, 2" }).CipPath.ShouldBe("1,2");
        new DeviceKey("plc1", 44818, "   ", "ControlLogix").CipPath.ShouldBe("");

        var pool = new DeviceSessionPool(new ManualTimeProvider());
        await using var la = pool.Acquire(a, Options());
        await using var lb = pool.Acquire(b, Options());
        lb.Scheduler.ShouldBeSameAs(la.Scheduler);   // one connection, one throttle
        pool.EntryCount.ShouldBe(1);
    }

    [Fact]
    public void Whitespace_inside_a_cip_path_part_is_kept_so_a_typo_fails_loudly()
    {
        var typo = new DeviceKey("plc1", 44818, "1,2 3", "ControlLogix");
        typo.CipPath.ShouldNotBe("1,23");                    // never silently slot 23 on another controller
        typo.CipPath.ShouldBe("1,2 3");
        Should.Throw<ArgumentException>(() => CipPath.EncodeRoute(typo.CipPath));
        CipPath.EncodeRoute(new DeviceKey("plc1", 44818, " 1, 0 ", "ControlLogix").CipPath).ShouldBe(new byte[] { 0x01, 0x00 });
    }

    [Fact]
    public async Task A_new_connection_waits_for_every_earlier_close_not_just_the_latest()
    {
        // A is closing (its Forward Close is held); B is acquired and released without ever opening, so
        // its own close is instant; C must still not open until A's close has finished.
        await using var target = new FakeCipTarget();
        var closeHeld = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        target.ForwardCloseGate = () => closeHeld.Task;
        var time = new ManualTimeProvider();
        var pool = new DeviceSessionPool(time);
        var key = new DeviceKey("127.0.0.1", target.Port, "1,0", "ControlLogix");

        var a = pool.Acquire(key, Options());
        await a.Scheduler.SendAsync(ReadTag, TestCt);
        var aClosing = a.DisposeAsync().AsTask();
        await FakeCipTarget.Until(() => target.ForwardCloses == 1, Guard, "A's Forward Close");

        await pool.Acquire(key, Options()).DisposeAsync();            // B: never opened
        time.Advance(TimeSpan.FromSeconds(5));                        // the throttle is no obstacle for C

        await using var c = pool.Acquire(key, Options());
        var cSend = c.Scheduler.SendAsync(ReadTag, TestCt);
        // Positive control (#583): C's open has reached the predecessor wait and parked there, so the absence of
        // its Forward Open below is the wait holding it, not C merely not having got that far yet.
        await FakeCipTarget.Until(() => pool.Diagnostics.PredecessorWaits == 1, Guard, "C to park behind A's close");
        target.ForwardOpens.ShouldBe(1);                              // C is held back behind A
        cSend.IsCompleted.ShouldBeFalse();

        closeHeld.SetResult();
        await aClosing.WaitAsync(Guard, TestCt);
        (await cSend.WaitAsync(Guard, TestCt)).GeneralStatus.ShouldBe((byte)0);
        target.ForwardOpens.ShouldBe(2);
    }

    [Fact]
    public async Task Disposing_during_an_open_still_forward_closes_what_the_target_accepted()
    {
        // The scheduler's stop token (and any caller's) ends only the WAIT; the Forward Open on the wire
        // completes, and the connection it created is closed.
        await using var target = new FakeCipTarget();
        var openHeld = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        target.ForwardOpenGate = () => openHeld.Task;
        var pool = new DeviceSessionPool(new ManualTimeProvider());
        var key = new DeviceKey("127.0.0.1", target.Port, "1,0", "ControlLogix");

        var lease = pool.Acquire(key, Options());
        var send = lease.Scheduler.SendAsync(ReadTag, TestCt);
        await FakeCipTarget.Until(() => target.ForwardOpens == 1, Guard, "the Forward Open");
        var disposing = lease.DisposeAsync().AsTask();

        var failed = await Should.ThrowAsync<AbTransportException>(() => send.WaitAsync(Guard, TestCt));
        failed.Status.ShouldBe(AbStatus.BadConnection);               // the scheduler's waiter is released at once
        target.ForwardCloses.ShouldBe(0);

        openHeld.SetResult();                                         // the target accepts, late
        await disposing.WaitAsync(Guard, TestCt);

        target.ForwardCloses.ShouldBe(1);
        target.ForwardCloseRequests.Single().AsSpan(8, 8).ToArray()
            .ShouldBe(target.ForwardOpenRequests.Single().AsSpan(16, 8).ToArray());
        target.ConnectedRequests.ShouldBe(0);
    }

    [Fact]
    public async Task The_shared_pool_logger_can_be_configured()
    {
        var logger = new CapturingLogger();
        DeviceSessionPool.ConfigureSharedLogger(new SingleLoggerFactory(logger));
        try
        {
            var key = new DeviceKey($"shared-logger-{Guid.NewGuid():N}", 44818, "1,0", "ControlLogix");
            await using var a = DeviceSessionPool.Shared.Acquire(key, Options(4002));
            await using var b = DeviceSessionPool.Shared.Acquire(key, Options(504));
            logger.Warnings.ShouldBe(1);
        }
        finally
        {
            DeviceSessionPool.ConfigureSharedLogger(Microsoft.Extensions.Logging.Abstractions.NullLoggerFactory.Instance);
        }
    }

    private sealed class SingleLoggerFactory(ILogger logger) : ILoggerFactory
    {
        public ILogger CreateLogger(string categoryName) => logger;
        public void AddProvider(ILoggerProvider provider) { }
        public void Dispose() { }
    }

    /// <summary>
    ///     Gitea #579: <see cref="ConnectionLease.ReleaseInBackground"/> returns before the last release's Forward
    ///     Close completes, yet a successor acquired that same instant still cannot open until the close is over —
    ///     the closing chain is registered synchronously, inside the call.
    /// </summary>
    [Fact]
    public async Task Release_in_background_returns_before_the_close_and_a_successor_still_waits_for_it()
    {
        await using var target = new FakeCipTarget();
        var closeHeld = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        target.ForwardCloseGate = () => closeHeld.Task;
        var time = new ManualTimeProvider();
        var pool = new DeviceSessionPool(time);
        var key = new DeviceKey("127.0.0.1", target.Port, "1,0", "ControlLogix");

        var a = pool.Acquire(key, Options());
        await a.Scheduler.SendAsync(ReadTag, TestCt);
        var teardown = a.ReleaseInBackground();
        pool.EntryCount.ShouldBe(0);                                  // bookkeeping done before it returned
        await FakeCipTarget.Until(() => target.ForwardCloses == 1, Guard, "the Forward Close");
        teardown.IsCompleted.ShouldBeFalse();                         // ...but the close is still pending

        time.Advance(TimeSpan.FromSeconds(5));                        // the throttle is no obstacle
        await using var b = pool.Acquire(key, Options());
        var bSend = b.Scheduler.SendAsync(ReadTag, TestCt);
        await FakeCipTarget.Until(() => pool.Diagnostics.PredecessorWaits == 1, Guard, "B to park behind A's close");
        target.ForwardOpens.ShouldBe(1);                              // held behind A's close
        bSend.IsCompleted.ShouldBeFalse();

        closeHeld.SetResult();
        await teardown.WaitAsync(Guard, TestCt);
        (await bSend.WaitAsync(Guard, TestCt)).GeneralStatus.ShouldBe((byte)0);
        target.ForwardOpens.ShouldBe(2);
        await a.ReleaseInBackground().WaitAsync(Guard, TestCt);      // idempotent: b's reference is untouched
        pool.EntryCount.ShouldBe(1);
    }

    /// <summary>
    ///     Gitea #566: releasing a lease while the device's shared open is still on the wire does not wait for that
    ///     open (up to its whole open budget) — <see cref="ConnectionLease.ReleaseInBackground"/> (#579) returns at
    ///     once, and the teardown it hands back still Forward Closes what the open produced once it lands.
    /// </summary>
    [Fact]
    public async Task Release_in_background_does_not_wait_behind_an_in_flight_shared_open()
    {
        await using var target = new FakeCipTarget();
        var openHeld = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        target.ForwardOpenGate = () => openHeld.Task;
        var pool = new DeviceSessionPool(new ManualTimeProvider());
        var key = new DeviceKey("127.0.0.1", target.Port, "1,0", "ControlLogix");

        var lease = pool.Acquire(key, Options());
        var send = lease.Scheduler.SendAsync(ReadTag, TestCt);
        await FakeCipTarget.Until(() => target.ForwardOpens == 1, Guard, "the Forward Open");
        ((CipConnection)lease.Channel).OpenInProgress.ShouldBeTrue();

        var teardown = lease.ReleaseInBackground();   // returns synchronously: the open is still held

        pool.EntryCount.ShouldBe(0);
        teardown.IsCompleted.ShouldBeFalse("the teardown waits for the open on the wire; the release call did not");
        (await Should.ThrowAsync<AbTransportException>(() => send.WaitAsync(Guard, TestCt))).Status.ShouldBe(AbStatus.BadConnection);

        openHeld.SetResult();
        await teardown.WaitAsync(Guard, TestCt);
        target.ForwardCloses.ShouldBe(1);             // what the late-accepted open created is still closed
    }

    /// <summary>Gitea #566: the close path disposes the connection even when the scheduler's dispose throws —
    ///     otherwise the connection's keep-alive would hold the press connection open forever.</summary>
    [Fact]
    public async Task Close_disposes_the_connection_even_when_the_scheduler_dispose_throws()
    {
        var logger = new CapturingLogger();
        var scheduler = new Disposable(() => throw new InvalidOperationException("scheduler dispose failed"));
        var connection = new Disposable(() => { });
        var key = new DeviceKey("plc1", 44818, "1,0", "ControlLogix");

        await DeviceSessionPool.CloseAsync(key, scheduler, connection, logger);

        scheduler.Disposed.ShouldBe(1);
        connection.Disposed.ShouldBe(1);
        logger.Warnings.ShouldBe(1);

        // And a throwing connection dispose is logged, never propagated to the release.
        var failing = new Disposable(() => throw new InvalidOperationException("close failed"));
        await DeviceSessionPool.CloseAsync(key, new Disposable(() => { }), failing, logger);
        failing.Disposed.ShouldBe(1);
        logger.Warnings.ShouldBe(2);
    }

    private sealed class Disposable(Action onDispose) : IAsyncDisposable
    {
        private int _disposed;
        public int Disposed => Volatile.Read(ref _disposed);

        public ValueTask DisposeAsync()
        {
            Interlocked.Increment(ref _disposed);
            onDispose();
            return ValueTask.CompletedTask;
        }
    }

    [Fact]
    public void Shared_is_one_process_wide_instance() =>
        DeviceSessionPool.Shared.ShouldBeSameAs(DeviceSessionPool.Shared);

    private sealed class CapturingLogger : ILogger
    {
        private int _warnings;
        public int Warnings => Volatile.Read(ref _warnings);
        public IDisposable? BeginScope<TState>(TState state) where TState : notnull => null;
        public bool IsEnabled(LogLevel logLevel) => true;

        public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception,
            Func<TState, Exception?, string> formatter)
        {
            if (logLevel == LogLevel.Warning) Interlocked.Increment(ref _warnings);
        }
    }
}