lmxopcua/tests/Server/ZB.MOM.WW.OtOpcUa.Runtime.Tests/Drivers/DriverInstanceActorConcurrentWriteTests.cs

using System.Collections.Concurrent;
using System.Diagnostics;
using Akka.Actor;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Commons.OpcUa;
using ZB.MOM.WW.OtOpcUa.Commons.Types;
using ZB.MOM.WW.OtOpcUa.Core.Abstractions;
using ZB.MOM.WW.OtOpcUa.Runtime.Drivers;
using ZB.MOM.WW.OtOpcUa.Runtime.Tests.Harness;

namespace ZB.MOM.WW.OtOpcUa.Runtime.Tests.Drivers;

/// <summary>
///     Review 02 P2: <c>WriteAttribute</c> was a <c>ReceiveAsync</c>, which suspends the mailbox for the whole
///     device round-trip (up to the 5 s bound). Writes therefore serialised — the SECOND of two operator writes to
///     a ~3 s device missed the host's 8 s Ask and was reverted although the device took it — and every data
///     change queued behind the write. Writes now run off the mailbox under a per-driver in-flight cap; the
///     subscribe/unsubscribe pair that must serialise against the subscription handle stays <c>ReceiveAsync</c>.
/// </summary>
public sealed class DriverInstanceActorConcurrentWriteTests : RuntimeActorTestBase
{
    private static readonly TimeSpan WriteLatency = TimeSpan.FromMilliseconds(1200);

    private (IActorRef Actor, Akka.TestKit.TestProbe Parent) Connect(IDriver driver, int? maxWritesInFlight = null)
    {
        var parent = CreateTestProbe();
        var actor = parent.ChildActorOf(DriverInstanceActor.Props(
            driver, commsLossTimeout: TimeSpan.Zero, maxWritesInFlight: maxWritesInFlight));
        actor.Tell(new DriverInstanceActor.InitializeRequested("{}"));
        parent.ExpectMsg<DriverInstanceActor.ConnectivityChanged>(m => m.Connected, PresenceBudget);
        return (actor, parent);
    }

    [Fact]
    public async Task Two_concurrent_slow_writes_both_complete_within_one_device_round_trip()
    {
        var driver = new SlowWritableDriver { Latency = WriteLatency };
        var (actor, _) = Connect(driver);

        var sw = Stopwatch.StartNew();
        var askBudget = TimeSpan.FromSeconds(8); // the host's Ask
        var first = actor.Ask<DriverInstanceActor.WriteAttributeResult>(
            new DriverInstanceActor.WriteAttribute("tag-1", 1), askBudget);
        var second = actor.Ask<DriverInstanceActor.WriteAttributeResult>(
            new DriverInstanceActor.WriteAttribute("tag-2", 2), askBudget);
        var results = await Task.WhenAll(first, second);
        sw.Stop();

        results.ShouldAllBe(r => r.Success);
        driver.Writes.Select(w => w.FullReference).ShouldBe(new[] { "tag-1", "tag-2" }, ignoreOrder: true);
        // Serialised writes take two latencies; concurrent ones take one. The margin covers a loaded machine.
        sw.Elapsed.ShouldBeLessThan(WriteLatency + WriteLatency / 2,
            "the second write waited for the first — writes are still serialised on the mailbox");
        driver.MaxConcurrentWrites.ShouldBe(2, "positive control: both writes were inside the driver at once");
    }

    [Fact]
    public async Task Data_changes_are_delivered_while_a_write_is_pending()
    {
        var driver = new SlowWritableDriver { HoldWrites = true };
        var (actor, parent) = Connect(driver);
        await actor.Ask<DriverInstanceActor.SubscriptionEstablished>(
            new DriverInstanceActor.Subscribe(new[] { "tag-a" }, TimeSpan.FromMilliseconds(100)), PresenceBudget);

        var pending = actor.Ask<DriverInstanceActor.WriteAttributeResult>(
            new DriverInstanceActor.WriteAttribute("tag-a", 7), PresenceBudget);
        AwaitCondition(() => driver.WritesEntered == 1, PresenceBudget);

        driver.FireDataChange("tag-a", 3.5);
        var published = parent.ExpectMsg<DriverInstanceActor.AttributeValueBatchPublished>(PresenceBudget);
        published.Items.Single().Value.ShouldBe(3.5);
        pending.IsCompleted.ShouldBeFalse("the write is still held — the data change did not wait behind it");

        // Positive control: releasing the device completes the write with its real outcome.
        driver.ReleaseWrites();
        (await pending).Success.ShouldBeTrue();
    }

    [Fact]
    public async Task Writes_beyond_the_in_flight_cap_queue_and_drain_in_order()
    {
        var driver = new SlowWritableDriver { HoldWrites = true };
        var (actor, _) = Connect(driver, maxWritesInFlight: 2);

        var asks = Enumerable.Range(0, 4)
            .Select(i => actor.Ask<DriverInstanceActor.WriteAttributeResult>(
                new DriverInstanceActor.WriteAttribute($"tag-{i}", i), PresenceBudget))
            .ToArray();

        AwaitCondition(() => driver.WritesEntered == 2, PresenceBudget);
        // Settle: a third write entering the driver would show up here.
        await Task.Delay(300);
        driver.WritesEntered.ShouldBe(2, "the cap let a third write into the driver");

        driver.ReleaseWrites();
        var results = await Task.WhenAll(asks);
        results.ShouldAllBe(r => r.Success);
        driver.Writes.Select(w => w.FullReference).ShouldBe(new[] { "tag-0", "tag-1", "tag-2", "tag-3" });
    }

    [Fact]
    public async Task Queued_writes_fail_fast_when_the_driver_leaves_Connected()
    {
        var driver = new SlowWritableDriver { HoldWrites = true };
        var (actor, parent) = Connect(driver, maxWritesInFlight: 1);

        var inFlight = actor.Ask<DriverInstanceActor.WriteAttributeResult>(
            new DriverInstanceActor.WriteAttribute("tag-0", 0), PresenceBudget);
        AwaitCondition(() => driver.WritesEntered == 1, PresenceBudget);
        var queued = actor.Ask<DriverInstanceActor.WriteAttributeResult>(
            new DriverInstanceActor.WriteAttribute("tag-1", 1), PresenceBudget);

        actor.Tell(new DriverInstanceActor.ForceReconnect());
        parent.ExpectMsg<DriverInstanceActor.ConnectivityChanged>(m => !m.Connected, PresenceBudget);

        var queuedResult = await queued;
        queuedResult.Success.ShouldBeFalse();
        queuedResult.Reason.ShouldBe("driver not connected");
        driver.WritesEntered.ShouldBe(1, "a queued write must never reach a driver that is reconnecting");

        // The in-flight one still answers with the device's real outcome once it is released.
        driver.ReleaseWrites();
        (await inFlight).Success.ShouldBeTrue();
    }

    /// <summary>A writable + subscribable driver whose writes either take <see cref="Latency"/> or hold until released.</summary>
    private sealed class SlowWritableDriver : IDriver, IWritable, ISubscribable
    {
        private readonly StubHandle _handle = new();
        private readonly ConcurrentQueue<WriteRequest> _writes = new();
        private TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously);
        private int _entered;
        private int _inside;
        private int _maxInside;

        public TimeSpan Latency { get; init; }
        public bool HoldWrites { get; init; }

        public IReadOnlyList<WriteRequest> Writes => _writes.ToArray();
        public int WritesEntered => Volatile.Read(ref _entered);
        public int MaxConcurrentWrites => Volatile.Read(ref _maxInside);

        public string DriverInstanceId => "slow-writer-1";
        public string DriverType => "SlowWriter";

        public void ReleaseWrites() => _release.TrySetResult();

        public Task InitializeAsync(string driverConfigJson, CancellationToken cancellationToken) => Task.CompletedTask;
        public Task ReinitializeAsync(string driverConfigJson, CancellationToken cancellationToken) => Task.CompletedTask;
        public Task ShutdownAsync(CancellationToken cancellationToken) => Task.CompletedTask;
        public DriverHealth GetHealth() => new(DriverState.Healthy, DateTime.UtcNow, null);

        public async Task<IReadOnlyList<WriteResult>> WriteAsync(IReadOnlyList<WriteRequest> writes, CancellationToken cancellationToken)
        {
            Interlocked.Increment(ref _entered);
            var inside = Interlocked.Increment(ref _inside);
            int seen;
            while (inside > (seen = Volatile.Read(ref _maxInside)))
                Interlocked.CompareExchange(ref _maxInside, inside, seen);
            try
            {
                if (HoldWrites) await _release.Task.WaitAsync(cancellationToken);
                else if (Latency > TimeSpan.Zero) await Task.Delay(Latency, cancellationToken);
                foreach (var w in writes) _writes.Enqueue(w);
                return writes.Select(_ => new WriteResult(OpcUaStatus.Good)).ToList();
            }
            finally
            {
                Interlocked.Decrement(ref _inside);
            }
        }

        public event EventHandler<DataChangeEventArgs>? OnDataChange;

        public Task<ISubscriptionHandle> SubscribeAsync(
            IReadOnlyList<string> fullReferences, TimeSpan publishingInterval, CancellationToken cancellationToken) =>
            Task.FromResult<ISubscriptionHandle>(_handle);

        public Task UnsubscribeAsync(ISubscriptionHandle handle, CancellationToken cancellationToken) => Task.CompletedTask;

        public void FireDataChange(string fullRef, object? value) =>
            OnDataChange?.Invoke(this, new DataChangeEventArgs(
                _handle, fullRef, new DataValueSnapshot(value, OpcUaStatus.Good, DateTime.UtcNow, DateTime.UtcNow)));
    }
}