lmxopcua/tests/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Modbus.Tests/ModbusBackgroundLoopTokenLifetimeTests.cs

using System.Diagnostics.Tracing;
using System.Reflection;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Driver.Modbus;

namespace ZB.MOM.WW.OtOpcUa.Driver.Modbus.Tests;

/// <summary>
///     #729 — the probe worker delegate must not obtain its <see cref="CancellationToken"/> after it is
///     already running. The scheduling token passed to <c>Task.Run</c> only suppresses a task that is still
///     QUEUED; once the pool has passed the task's cancellation/start check, lifecycle teardown can finish
///     underneath it (the bounded join gives up and the CTS is disposed). A delegate that then reads
///     <c>cts.Token</c> throws <see cref="ObjectDisposedException"/>; one that reads the <c>_probeCts</c>
///     field sees null or the next generation's token.
/// </summary>
/// <remarks>
///     The seam is TPL's own <c>TaskStarted</c> event: <c>TplEventSource</c> raises it synchronously on the
///     worker thread AFTER the start check and BEFORE the delegate is invoked, so holding the listener
///     callback parks the probe worker in exactly the window the issue describes — deterministically, no
///     timing stress. The target task is identified by the <c>TaskScheduled</c> event raised on the test's
///     own (AsyncLocal-armed) flow during <c>InitializeAsync</c>; more than one armed schedule fails closed.
///     Queued-task cancellation (shutdown before the pool reaches the start check) is not exercised here:
///     there is no deterministic seam ahead of the start check, and that path is the BCL's own behaviour
///     for the scheduling token, which the fix leaves unchanged.
/// </remarks>
[Trait("Category", "Unit")]
public sealed class ModbusBackgroundLoopTokenLifetimeTests
{
    private static readonly TimeSpan Bound = TimeSpan.FromSeconds(15);

    private sealed class CountingTransport : IModbusTransport
    {
        public int SendCount;
        public int DisposeCount;

        public Task ConnectAsync(CancellationToken ct) => Task.CompletedTask;

        public Task<byte[]> SendAsync(byte unitId, byte[] pdu, CancellationToken ct)
        {
            Interlocked.Increment(ref SendCount);
            return Task.FromResult(new byte[] { pdu[0], 2, 0, 0 });
        }

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

    /// <summary>Parks the one task scheduled on the armed flow inside its TaskStarted callback.</summary>
    private sealed class TaskStartGate : EventListener
    {
        private const string TplSourceName = "System.Threading.Tasks.TplEventSource";
        private const EventKeywords TaskTransfer = (EventKeywords)0x1;
        private const EventKeywords Tasks = (EventKeywords)0x2;

        public static readonly AsyncLocal<bool> Armed = new();

        private readonly ManualResetEventSlim _release = new(false);
        private int _targetTaskId = -1;
        private int _armedSchedules;
        private volatile bool _closed;

        public TaskCompletionSource Started { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
        public int ArmedSchedules => Volatile.Read(ref _armedSchedules);
        public bool ReleaseTimedOut { get; private set; }

        /// <summary>Stops correlating new schedules once the arming operation has returned.</summary>
        public void Close() => _closed = true;

        public void Release() => _release.Set();

        protected override void OnEventSourceCreated(EventSource eventSource)
        {
            if (eventSource.Name == TplSourceName)
                EnableEvents(eventSource, EventLevel.Informational, TaskTransfer | Tasks);
        }

        protected override void OnEventWritten(EventWrittenEventArgs e)
        {
            if (e.EventName == "TaskScheduled")
            {
                if (_closed || !Armed.Value) return;
                Interlocked.Increment(ref _armedSchedules);
                Interlocked.CompareExchange(ref _targetTaskId, TaskId(e), -1);
            }
            else if (e.EventName == "TaskStarted")
            {
                var target = Volatile.Read(ref _targetTaskId);
                if (target == -1 || TaskId(e) != target) return;
                Started.TrySetResult();
                // Bounded: a test that never releases must not wedge a pool thread forever.
                if (!_release.Wait(Bound)) ReleaseTimedOut = true;
            }
        }

        private static int TaskId(EventWrittenEventArgs e)
        {
            var i = e.PayloadNames!.IndexOf("TaskID");
            return i < 0 ? -2 : Convert.ToInt32(e.Payload![i]);
        }

        public override void Dispose()
        {
            _release.Set();
            base.Dispose();
        }
    }

    private static ModbusDriverOptions ProbeOnlyOptions() => new()
    {
        Host = "fake",
        Probe = new ModbusProbeOptions
        {
            Enabled = true,
            Interval = TimeSpan.FromMilliseconds(50),
            Timeout = TimeSpan.FromSeconds(1),
        },
        AutoProhibitReprobeInterval = null,
        RawTags = ModbusRawTags.Entries([new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)]),
    };

    private static T? Field<T>(ModbusDriver drv, string name) where T : class =>
        (T?)typeof(ModbusDriver).GetField(name, BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(drv);

    /// <summary>Initializes on an armed flow and returns the parked background-loop task.</summary>
    private static async Task<Task> InitializeAndParkProbeAsync(ModbusDriver drv, TaskStartGate gate, string taskField = "_probeTask")
    {
        TaskStartGate.Armed.Value = true;
        try { await drv.InitializeAsync("{}", CancellationToken.None); }
        finally
        {
            gate.Close();
            TaskStartGate.Armed.Value = false;
        }

        gate.ArmedSchedules.ShouldBe(1, "exactly one task (the background loop) must be scheduled by initialize; anything else is ambiguous");
        var probeTask = Field<Task>(drv, taskField);
        Assert.NotNull(probeTask); // not Shouldly: its Task overload would be an un-awaited async assertion

        await gate.Started.Task.WaitAsync(Bound, TestContext.Current.CancellationToken);
        probeTask.IsCompleted.ShouldBeFalse("the worker must be parked after its start check, before its delegate runs");
        return probeTask;
    }

    private static async Task SettleAsync(Task task)
    {
        try { await task.WaitAsync(Bound, TestContext.Current.CancellationToken); }
        catch (TimeoutException) { throw; }
        catch { /* outcome asserted by the caller from task.Status */ }
    }

    /// <summary>
    ///     A probe worker already past its start check when ShutdownAsync tears the driver down must still
    ///     complete cleanly on its own cancelled token once it runs — not fault on a disposed CTS, and not
    ///     touch the disposed transport.
    /// </summary>
    [Fact]
    public async Task Started_probe_worker_released_after_shutdown_exits_cleanly_on_its_own_cancelled_token()
    {
        var fake = new CountingTransport();
        var drv = new ModbusDriver(ProbeOnlyOptions(), "modbus-729", _ => fake);
        using var gate = new TaskStartGate();
        try
        {
            var probeTask = await InitializeAndParkProbeAsync(drv, gate);

            // Teardown runs to completion while the worker is parked: the bounded join times out and the
            // CTS is disposed and the field cleared before the delegate has read anything.
            await drv.ShutdownAsync(CancellationToken.None);
            Field<CancellationTokenSource>(drv, "_probeCts").ShouldBeNull();
            fake.DisposeCount.ShouldBe(1);
            probeTask.IsCompleted.ShouldBeFalse();

            gate.Release();
            await SettleAsync(probeTask);

            gate.ReleaseTimedOut.ShouldBeFalse();
            probeTask.Exception?.Flatten().InnerException.ShouldBeNull(
                "the late worker must not fault obtaining its token after teardown disposed the CTS");
            probeTask.Status.ShouldBe(TaskStatus.RanToCompletion);
            Volatile.Read(ref fake.SendCount).ShouldBe(0, "a worker released after teardown must not probe the disposed transport");
        }
        finally
        {
            gate.Release();
            await drv.DisposeAsync();
        }
    }

    /// <summary>
    ///     A probe worker of generation N released after a replacement initialize must run on generation
    ///     N's (cancelled) token — not adopt generation N+1's live token and keep probing the N transport —
    ///     while generation N+1's own probe loop stays live.
    /// </summary>
    [Fact]
    public async Task Started_probe_worker_released_after_replacement_keeps_its_own_generation_token()
    {
        var transports = new List<CountingTransport>();
        var drv = new ModbusDriver(ProbeOnlyOptions(), "modbus-729", _ =>
        {
            var t = new CountingTransport();
            lock (transports) transports.Add(t);
            return t;
        });
        using var gate = new TaskStartGate();
        try
        {
            var firstProbe = await InitializeAndParkProbeAsync(drv, gate);
            var first = transports[0];

            await drv.ReinitializeAsync("{}", CancellationToken.None);
            transports.Count.ShouldBe(2);
            var second = transports[1];
            first.DisposeCount.ShouldBe(1);
            var liveCts = Field<CancellationTokenSource>(drv, "_probeCts");
            liveCts.ShouldNotBeNull();
            firstProbe.IsCompleted.ShouldBeFalse();

            gate.Release();
            await SettleAsync(firstProbe);

            gate.ReleaseTimedOut.ShouldBeFalse();
            firstProbe.Exception?.Flatten().InnerException.ShouldBeNull();
            firstProbe.Status.ShouldBe(TaskStatus.RanToCompletion,
                "the generation-1 worker must exit on its own cancelled token");
            Volatile.Read(ref first.SendCount).ShouldBe(0, "the generation-1 worker must not probe its disposed transport");

            // Generation 2 is unaffected: its loop is live and its CTS is still the published, uncancelled one.
            var deadline = DateTime.UtcNow + Bound;
            while (Volatile.Read(ref second.SendCount) == 0 && DateTime.UtcNow < deadline)
                await Task.Delay(10, TestContext.Current.CancellationToken);
            Volatile.Read(ref second.SendCount).ShouldBeGreaterThan(0);
            Field<CancellationTokenSource>(drv, "_probeCts").ShouldBeSameAs(liveCts);
            liveCts!.IsCancellationRequested.ShouldBeFalse();
        }
        finally
        {
            gate.Release();
            await drv.DisposeAsync();
        }
    }

    /// <summary>The re-probe loop gets the same guarantee as the probe loop.</summary>
    [Fact]
    public async Task Started_reprobe_worker_released_after_shutdown_exits_cleanly_on_its_own_cancelled_token()
    {
        var fake = new CountingTransport();
        var opts = new ModbusDriverOptions
        {
            Host = "fake",
            Probe = new ModbusProbeOptions { Enabled = false },
            AutoProhibitReprobeInterval = TimeSpan.FromMilliseconds(50),
            RawTags = ModbusRawTags.Entries([new ModbusTagDefinition("A", ModbusRegion.HoldingRegisters, 0, ModbusDataType.Int16)]),
        };
        var drv = new ModbusDriver(opts, "modbus-729", _ => fake);
        using var gate = new TaskStartGate();
        try
        {
            var reprobeTask = await InitializeAndParkProbeAsync(drv, gate, "_reprobeTask");

            await drv.ShutdownAsync(CancellationToken.None);
            Field<CancellationTokenSource>(drv, "_reprobeCts").ShouldBeNull();
            reprobeTask.IsCompleted.ShouldBeFalse();

            gate.Release();
            await SettleAsync(reprobeTask);

            gate.ReleaseTimedOut.ShouldBeFalse();
            reprobeTask.Exception?.Flatten().InnerException.ShouldBeNull(
                "the late re-probe worker must not fault obtaining its token after teardown disposed the CTS");
            reprobeTask.Status.ShouldBe(TaskStatus.RanToCompletion);
        }
        finally
        {
            gate.Release();
            await drv.DisposeAsync();
        }
    }
}