lmxopcua/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport/DeviceSessionPool.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport.Cip;
namespace ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport;

/// <summary>
/// Process-wide owner of device connections: exactly one <see cref="CipConnection"/> and one
/// <see cref="RequestScheduler"/> per <see cref="DeviceKey"/> (review B1), handed out as ref-counted
/// <see cref="ConnectionLease"/>s. The last lease's dispose closes the connection; a later acquire builds
/// fresh instances, which do not open until the previous connection has finished closing.
/// </summary>
/// <remarks>Forward Open throttles are kept per key for the process lifetime and never pruned, so a
/// reconnect after the last release is still rate-limited against the previous attempt.</remarks>
public sealed class DeviceSessionPool
{
    /// <summary>The process-wide pool the drivers use. Tests build their own.</summary>
    public static DeviceSessionPool Shared { get; } = new();

    internal sealed class Entry(CipConnection connection, RequestScheduler scheduler, ConnectionOptions options,
        Task predecessorClosed)
    {
        /// <summary>The shared CIP connection for the device.</summary>
        public CipConnection Connection { get; } = connection;
        /// <summary>Completes when every earlier connection for this key has finished closing.</summary>
        public Task PredecessorClosed { get; } = predecessorClosed;
        /// <summary>The request scheduler serializing and packing traffic on <see cref="Connection"/>.</summary>
        public RequestScheduler Scheduler { get; } = scheduler;
        /// <summary>The options the connection was created with.</summary>
        public ConnectionOptions Options { get; } = options;
        public int RefCount;            // under the pool lock
        public bool MismatchLogged;     // under the pool lock
    }

    private readonly object _lock = new();
    private readonly Dictionary<DeviceKey, Entry> _entries = new();
    private readonly Dictionary<DeviceKey, ForwardOpenThrottle> _throttles = new();
    private readonly Dictionary<DeviceKey, Task> _closing = new();
    private readonly TimeProvider _time;
    private volatile ILogger _logger;

    /// <summary>Initializes a new instance of the <see cref="DeviceSessionPool"/> class.</summary>
    /// <param name="time">Time source; defaults to <see cref="TimeProvider.System"/>.</param>
    /// <param name="logger">Logger for mismatch and close-failure warnings; defaults to a no-op logger.</param>
    public DeviceSessionPool(TimeProvider? time = null, ILogger? logger = null)
    {
        _time = time ?? TimeProvider.System;
        _logger = logger ?? NullLogger.Instance;
    }

    /// <summary>Give <see cref="Shared"/> a logger (it is built before any host exists), so its
    /// options-mismatch and close-failure warnings reach the service log. Call once at startup.</summary>
    /// <param name="loggerFactory">The factory used to create the shared pool's logger.</param>
    public static void ConfigureSharedLogger(ILoggerFactory loggerFactory)
    {
        ArgumentNullException.ThrowIfNull(loggerFactory);
        Shared._logger = loggerFactory.CreateLogger<DeviceSessionPool>();
    }

    /// <summary>Wiring-test hook: the logger this pool currently writes to (final review I2).</summary>
    internal ILogger Logger => _logger;

    /// <summary>Counters and open-attempt limits for this pool's connections (gate and harness diagnostics).</summary>
    internal ConnectionDiagnostics Diagnostics { get; } = new();

    /// <summary>Test hook: keys with a live entry.</summary>
    internal int EntryCount { get { lock (_lock) return _entries.Count; } }

    /// <summary>Lease the shared connection for <paramref name="key"/>, creating it on first use. Performs
    /// no I/O: the connection opens on its first request. When the key already has an entry, that entry's
    /// options win; a differing <paramref name="options"/> is logged once per entry.</summary>
    /// <param name="key">The device to lease a connection for.</param>
    /// <param name="options">The options to open the connection with if the key has no entry yet.</param>
    /// <returns>A lease on the device's shared connection.</returns>
    public ConnectionLease Acquire(DeviceKey key, ConnectionOptions options)
    {
        ArgumentNullException.ThrowIfNull(key);
        ArgumentNullException.ThrowIfNull(options);
        lock (_lock)
        {
            if (!_entries.TryGetValue(key, out var entry))
            {
                if (!_throttles.TryGetValue(key, out var throttle))
                    _throttles[key] = throttle = new ForwardOpenThrottle(options.MinForwardOpenInterval, _time);
                var predecessor = _closing.GetValueOrDefault(key) ?? Task.CompletedTask;
                var connection = new CipConnection(key, options, throttle, _time, predecessor, Diagnostics);
                entry = new Entry(connection, new RequestScheduler(connection, options.AllowPacking), options, predecessor);
                _entries[key] = entry;
            }
            else if (!entry.Options.Equals(options) && !entry.MismatchLogged)
            {
                entry.MismatchLogged = true;
                _logger.LogWarning(
                    "Device {Host}:{Port} path '{CipPath}' is already connected with different connection options; " +
                    "the existing options ({Existing}) stay in force and the requested ({Requested}) are ignored.",
                    key.Host, key.Port, key.CipPath, entry.Options, options);
            }
            entry.RefCount++;
            return new ConnectionLease(this, key, entry);
        }
    }

    /// <summary>Drops one reference on <paramref name="entry"/>; the last release removes it and closes the connection.</summary>
    /// <param name="key">The device key the entry belongs to.</param>
    /// <param name="entry">The entry whose reference is released.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    internal async ValueTask ReleaseAsync(DeviceKey key, Entry entry)
    {
        TaskCompletionSource closed;
        Task chain;
        lock (_lock)
        {
            if (--entry.RefCount > 0) return;
            _entries.Remove(key);
            closed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
            // The successor must wait for THIS close and for every earlier one still running. An entry
            // that never opened closes instantly; without the chain it would hide a slow predecessor
            // (A closing, B acquired+released unopened, C opens while A's Forward Close is still pending).
            chain = Task.WhenAll(entry.PredecessorClosed, closed.Task);
            _closing[key] = chain;
        }
        _ = chain.ContinueWith(_ =>
        {
            lock (_lock)
            {
                if (_closing.TryGetValue(key, out var t) && t == chain) _closing.Remove(key);
            }
        }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);

        // Outside the lock.
        try
        {
            await CloseAsync(key, entry.Scheduler, entry.Connection, _logger).ConfigureAwait(false);
        }
        finally
        {
            closed.TrySetResult();   // never faults: a successor only needs to know the close is over
        }
    }

    /// <summary>
    /// Scheduler first (fails its waiters, stops sending), then Forward Close + UnRegisterSession via the
    /// connection. The connection is disposed even if the scheduler's dispose throws: otherwise its keep-alive
    /// would hold the press connection open forever. Never throws; each failure is logged.
    /// </summary>
    /// <param name="key">The device being closed (for logging).</param>
    /// <param name="scheduler">The request scheduler to dispose first.</param>
    /// <param name="connection">The connection to dispose after the scheduler.</param>
    /// <param name="logger">The logger that receives close failures.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    internal static async Task CloseAsync(DeviceKey key, IAsyncDisposable scheduler, IAsyncDisposable connection, ILogger logger)
    {
        try { await scheduler.DisposeAsync().ConfigureAwait(false); }
        catch (Exception ex)
        {
            logger.LogWarning(ex, "Disposing the request scheduler for {Host}:{Port} failed; closing the connection anyway.",
                key.Host, key.Port);
        }
        try { await connection.DisposeAsync().ConfigureAwait(false); }
        catch (Exception ex)
        {
            logger.LogWarning(ex, "Closing the connection to {Host}:{Port} failed.", key.Host, key.Port);
        }
    }
}

/// <summary>A reference on a device's shared connection. Dispose exactly once (further disposes are no-ops).</summary>
public sealed class ConnectionLease : IAsyncDisposable
{
    private readonly DeviceSessionPool _pool;
    private readonly DeviceKey _key;
    private readonly DeviceSessionPool.Entry _entry;
    private int _disposed;

    /// <summary>Initializes a new instance of the <see cref="ConnectionLease"/> class.</summary>
    /// <param name="pool">The owning pool.</param>
    /// <param name="key">The device key leased.</param>
    /// <param name="entry">The pool entry this lease references.</param>
    internal ConnectionLease(DeviceSessionPool pool, DeviceKey key, DeviceSessionPool.Entry entry)
    {
        _pool = pool;
        _key = key;
        _entry = entry;
    }

    /// <summary>The device this lease is on.</summary>
    public DeviceKey Key => _key;

    /// <summary>The device's channel. Prefer <see cref="Scheduler"/>: it serializes and packs requests.</summary>
    public ICipChannel Channel => _entry.Connection;

    /// <summary>The device's single request scheduler, shared by every lease on the key.</summary>
    public RequestScheduler Scheduler => _entry.Scheduler;

    /// <summary>Releases the lease; further calls are no-ops.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public ValueTask DisposeAsync() =>
        Interlocked.Exchange(ref _disposed, 1) == 0 ? _pool.ReleaseAsync(_key, _entry) : ValueTask.CompletedTask;

    /// <summary>
    /// Releases this lease without waiting for the teardown — the one release shape both native AB tag
    /// runtimes use from their synchronous <c>Dispose</c>, which the drivers call under a per-tag lock.
    /// Idempotent, like <see cref="DisposeAsync"/>.
    /// </summary>
    /// <remarks>
    /// <para>Nothing is lost by not waiting. The release's bookkeeping — the ref-count drop and, for the
    /// last lease, the entry's removal and the registration of its close in the pool's closing chain — runs
    /// synchronously, before this returns. So an acquire made the very next instant builds a connection that
    /// does not open until this one's Forward Close + UnRegisterSession have finished (the guarantee the old
    /// blocking 2 s wait in the AbCip runtime existed for), and the close itself is bounded by
    /// <see cref="CipConnection"/>'s own budgets (the open budget for an in-flight open, 1 s for the Forward
    /// Close), so it always completes and the socket is always released.</para>
    /// <para>The returned task is the teardown; it never faults (the pool logs its own close failures), so
    /// a caller may drop it.</para>
    /// </remarks>
    /// <returns>The teardown, for tests and callers that want to observe it.</returns>
    public Task ReleaseInBackground() => Observe(DisposeAsync());

    private static async Task Observe(ValueTask release)
    {
        try { await release.ConfigureAwait(false); }
        catch (Exception) { /* ReleaseAsync logs and swallows its own failures; never fault the observer */ }
    }
}