mxaccessgw/src/ZB.MOM.WW.MxGateway.Worker.Tests/TestSupport/FakeRuntimeSession.cs

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Google.Protobuf.WellKnownTypes;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Worker.Ipc;
using ZB.MOM.WW.MxGateway.Worker.MxAccess;
using ZB.MOM.WW.MxGateway.Worker.Sta;

namespace ZB.MOM.WW.MxGateway.Worker.Tests.TestSupport;

/// <summary>
/// Single configurable <see cref="IWorkerRuntimeSession"/> test double shared by
/// the IPC tests. Replaces the two independent (and previously diverged)
/// <c>FakeRuntimeSession</c> copies in WorkerPipeSessionTests and
/// WorkerPipeClientTests: one supported dispatch blocking and event enqueue, the
/// other did not. This consolidated double supports every configuration both
/// call sites needed, so a minimal caller simply leaves the options unset.
/// </summary>
internal sealed class FakeRuntimeSession : IWorkerRuntimeSession
{
    /// <summary>
    ///     Backstop on the <see cref="BlockDispatch"/> wait so a test that never releases leaves no
    ///     thread parked forever. It is a safety net, never a scenario's timing budget: nothing
    ///     asserts on it firing, and a test whose blocked window outruns it silently gets its reply
    ///     mid-window, which then fails as an opaque cancellation somewhere later. Kept far above
    ///     any test's window — and above the 20 s cancellation those tests arm — so the test's own
    ///     token always fails first, with its own message. <see cref="Dispose"/> releases the wait
    ///     regardless, so teardown never depends on this elapsing.
    /// </summary>
    private static readonly TimeSpan BlockedDispatchSafetyNet = TimeSpan.FromSeconds(30);

    private readonly ManualResetEventSlim releaseDispatch = new(false);
    private readonly object gate = new();
    private readonly Queue<WorkerEvent> events = new();
    private readonly List<string> cancelledCorrelationIds = new();

    // Mirrors MxAccessEventQueue's coalesced wake signal so the drain loop under test is driven the
    // same way it is in production: EnqueueEvent(s) releases one permit, WaitForEventsAsync consumes
    // it. Never disposed — the drain loop can still be parked on it while Dispose runs, and a
    // disposed SemaphoreSlim would turn that shutdown into an ObjectDisposedException.
    private readonly SemaphoreSlim eventSignal = new(0, 1);
    private TimeSpan? lastWaitForEventsTimeout;
    private bool refreshStaActivityOnCapture;
    private WorkerRuntimeHeartbeatSnapshot snapshot = new(
        DateTimeOffset.UtcNow,
        pendingCommandCount: 0,
        outboundEventQueueDepth: 0,
        lastEventSequence: 0,
        currentCommandCorrelationId: string.Empty);

    /// <summary>Gets the event signaled when dispatch begins.</summary>
    public ManualResetEventSlim DispatchStarted { get; } = new(false);

    /// <summary>Blocks dispatch execution until explicitly released.</summary>
    public bool BlockDispatch { get; set; }

    /// <summary>Gets or sets whether to throw an exception after dispatch is released.</summary>
    public bool ThrowAfterDispatchReleased { get; set; }

    /// <summary>Gets or sets whether ShutdownGracefullyAsync throws a TimeoutException.</summary>
    public bool ThrowTimeoutOnShutdown { get; set; }

    /// <summary>
    ///     Optional diagnostic message stuffed into every dispatched command reply. A long value
    ///     pushes the STA command reply past a small negotiated frame maximum, which is how a test
    ///     drives the <c>ProcessCommandAsync</c> reply-size backstop.
    /// </summary>
    public string? DispatchReplyDiagnosticMessage { get; set; }

    /// <summary>Gets a value indicating whether Dispose was called.</summary>
    public bool Disposed { get; private set; }

    /// <inheritdoc />
    public Task<WorkerReady> StartAsync(
        string sessionId,
        int workerProcessId,
        CancellationToken cancellationToken = default)
    {
        return Task.FromResult(new WorkerReady
        {
            WorkerProcessId = workerProcessId,
            MxaccessProgid = MxAccessInteropInfo.ProgId,
            MxaccessClsid = MxAccessInteropInfo.Clsid,
            ReadyTimestamp = Timestamp.FromDateTimeOffset(DateTimeOffset.UtcNow),
        });
    }

    /// <inheritdoc />
    public Task<MxCommandReply> DispatchAsync(StaCommand command)
    {
        return Task.Run(
            () =>
            {
                SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
                    DateTimeOffset.UtcNow,
                    pendingCommandCount: 0,
                    outboundEventQueueDepth: 0,
                    lastEventSequence: 0,
                    command.CorrelationId));
                DispatchStarted.Set();

                if (BlockDispatch)
                {
                    releaseDispatch.Wait(BlockedDispatchSafetyNet);
                }

                SetSnapshot(new WorkerRuntimeHeartbeatSnapshot(
                    DateTimeOffset.UtcNow,
                    pendingCommandCount: 0,
                    outboundEventQueueDepth: 0,
                    lastEventSequence: 0,
                    currentCommandCorrelationId: string.Empty));

                if (ThrowAfterDispatchReleased)
                {
                    throw new InvalidOperationException("Command failed after shutdown started.");
                }

                MxCommandReply reply = new()
                {
                    SessionId = command.SessionId,
                    CorrelationId = command.CorrelationId,
                    Kind = command.Kind,
                    ProtocolStatus = new ProtocolStatus
                    {
                        Code = ProtocolStatusCode.Ok,
                        Message = "OK",
                    },
                };

                if (DispatchReplyDiagnosticMessage is not null)
                {
                    reply.DiagnosticMessage = DispatchReplyDiagnosticMessage;
                }

                return reply;
            });
    }

    /// <summary>
    ///     When set, <see cref="CaptureHeartbeat"/> stamps the snapshot's
    ///     <c>LastStaActivityUtc</c> with the capture time and leaves every other field as the last
    ///     <see cref="SetSnapshot"/> left it. Models a live STA whose pump calls
    ///     <c>MarkActivity()</c> on each wait iteration (<c>StaRuntime.ThreadMain</c>), so a healthy
    ///     worker is never captured stale — which a watchdog test needs to hold for the <em>whole</em>
    ///     session, including the handshake window before any command exists for the watchdog to
    ///     suppress on. A test-owned refresh loop cannot hold it: it is a thread-pool continuation
    ///     racing a compressed grace, and the gap between this fake being constructed and that loop's
    ///     first tick is already enough to look hung.
    /// </summary>
    public bool RefreshStaActivityOnCapture
    {
        get
        {
            lock (gate)
            {
                return refreshStaActivityOnCapture;
            }
        }

        set
        {
            lock (gate)
            {
                refreshStaActivityOnCapture = value;
            }
        }
    }

    /// <inheritdoc />
    public WorkerRuntimeHeartbeatSnapshot CaptureHeartbeat()
    {
        lock (gate)
        {
            if (refreshStaActivityOnCapture)
            {
                snapshot = new WorkerRuntimeHeartbeatSnapshot(
                    DateTimeOffset.UtcNow,
                    snapshot.PendingCommandCount,
                    snapshot.OutboundEventQueueDepth,
                    snapshot.LastEventSequence,
                    snapshot.CurrentCommandCorrelationId,
                    snapshot.StaCallInProgress);
            }

            return snapshot;
        }
    }

    /// <summary>
    ///     When set, <see cref="DrainEvents"/> returns no events for the
    ///     WorkerPipeSession background drain loop's fixed batch size, so an
    ///     explicit DrainEvents control command (which drains all via
    ///     <c>maxEvents == 0</c>) can claim the queued events deterministically
    ///     without racing the 25 ms background loop. Mirrors
    ///     <c>WorkerPipeSession.EventDrainBatchSize</c>.
    /// </summary>
    public uint? SuppressDrainForBatchSize { get; set; }

    /// <summary>
    ///     Records the <c>maxEvents</c> argument of the most recent non-suppressed
    ///     <see cref="DrainEvents"/> call — i.e. the effective cap the session passed for an explicit
    ///     DrainEvents control command. Lets a test assert the worker bounds the drain rather
    ///     than forwarding the client's raw <c>max_events = 0</c>.
    /// </summary>
    public uint? LastDrainMaxEvents { get; private set; }

    /// <summary>
    ///     Optional real event queue backing the drain paths. When set, both
    ///     <see cref="DrainEvents(uint)"/> and <see cref="DrainEvents(uint, int)"/> delegate to it
    ///     so a test can exercise the production byte-budgeting logic behind the fake session.
    /// </summary>
    public MxAccessEventQueue? BackingQueue { get; set; }

    /// <summary>
    ///     When set, <see cref="DrainEvents(uint, int)"/> ignores the byte budget and drains purely
    ///     by count. Simulates the "sizing bug or future command" case the control-reply size
    ///     backstop exists for, so a test can drive an oversized reply without a real budgeting
    ///     defect.
    /// </summary>
    public bool IgnoreDrainByteBudget { get; set; }

    /// <summary>
    ///     Records the <c>maxTotalBytes</c> argument of the most recent byte-budgeted
    ///     <see cref="DrainEvents(uint, int)"/> call.
    /// </summary>
    public int? LastDrainMaxTotalBytes { get; private set; }

    /// <inheritdoc />
    public IReadOnlyList<WorkerEvent> DrainEvents(uint maxEvents)
    {
        if (SuppressDrainForBatchSize is uint suppressed && maxEvents == suppressed)
        {
            return Array.Empty<WorkerEvent>();
        }

        LastDrainMaxEvents = maxEvents;

        if (BackingQueue is not null)
        {
            return BackingQueue.Drain(maxEvents);
        }

        lock (gate)
        {
            int drainCount = maxEvents == 0
                ? events.Count
                : Math.Min(events.Count, checked((int)Math.Min(maxEvents, int.MaxValue)));
            List<WorkerEvent> drained = new(drainCount);
            for (int index = 0; index < drainCount; index++)
            {
                drained.Add(events.Dequeue());
            }

            return drained;
        }
    }

    /// <inheritdoc />
    public WorkerEventDrainResult DrainEvents(uint maxEvents, int maxTotalBytes)
    {
        if (SuppressDrainForBatchSize is uint suppressed && maxEvents == suppressed)
        {
            return new WorkerEventDrainResult(
                Array.Empty<WorkerEvent>(),
                truncatedBySize: false,
                remainingCount: PendingEventCount,
                oversizedHeadSequence: 0);
        }

        LastDrainMaxEvents = maxEvents;
        LastDrainMaxTotalBytes = maxTotalBytes;

        if (BackingQueue is not null && !IgnoreDrainByteBudget)
        {
            return BackingQueue.Drain(maxEvents, maxTotalBytes);
        }

        // Count-only drain: either no backing queue (the simple fakes) or a deliberately
        // budget-blind drain used to exercise the reply-size backstop.
        IReadOnlyList<WorkerEvent> drained = BackingQueue is not null
            ? BackingQueue.Drain(maxEvents)
            : DrainByCount(maxEvents);
        return new WorkerEventDrainResult(
            drained,
            truncatedBySize: false,
            remainingCount: PendingEventCount,
            oversizedHeadSequence: 0);
    }

    private int PendingEventCount
    {
        get
        {
            if (BackingQueue is not null)
            {
                return BackingQueue.Count;
            }

            lock (gate)
            {
                return events.Count;
            }
        }
    }

    private IReadOnlyList<WorkerEvent> DrainByCount(uint maxEvents)
    {
        lock (gate)
        {
            int drainCount = maxEvents == 0
                ? events.Count
                : Math.Min(events.Count, checked((int)Math.Min(maxEvents, int.MaxValue)));
            List<WorkerEvent> drained = new(drainCount);
            for (int index = 0; index < drainCount; index++)
            {
                drained.Add(events.Dequeue());
            }

            return drained;
        }
    }

    /// <summary>
    ///     When set, <see cref="WaitForEventsAsync"/> honours only the wake signal and cancellation,
    ///     never the fallback timeout. A drain loop that ships an event while this is set can only
    ///     have been woken by the enqueue signal, which is what makes "the drain is signal-driven,
    ///     not poll-driven" assertable without racing the 25 ms fallback tick.
    /// </summary>
    public bool WaitForEventsOnSignalOnly { get; set; }

    /// <summary>
    ///     The <c>timeout</c> argument of the most recent <see cref="WaitForEventsAsync"/> call, so
    ///     a test can assert the drain loop still passes its fallback ceiling.
    /// </summary>
    public TimeSpan? LastWaitForEventsTimeout
    {
        get
        {
            lock (gate)
            {
                return lastWaitForEventsTimeout;
            }
        }
    }

    /// <inheritdoc />
    public Task WaitForEventsAsync(TimeSpan timeout, CancellationToken cancellationToken)
    {
        lock (gate)
        {
            lastWaitForEventsTimeout = timeout;
        }

        if (BackingQueue is not null)
        {
            if (WaitForEventsOnSignalOnly)
            {
                // The two are mutually exclusive: a real backing queue owns its own signal and
                // always honours the timeout, so silently letting it win would leave a
                // signal-only test passing on a fallback tick — exactly the false green the
                // option exists to rule out.
                throw new InvalidOperationException(
                    "FakeRuntimeSession cannot combine BackingQueue with WaitForEventsOnSignalOnly: "
                    + "the backing queue honours the fallback timeout, which defeats the signal-only wait.");
            }

            // Tests that drive a real queue enqueue into it directly, so the real queue owns the
            // wake signal too.
            return BackingQueue.WaitForEventsAsync(timeout, cancellationToken);
        }

        if (WaitForEventsOnSignalOnly)
        {
            return eventSignal.WaitAsync(cancellationToken);
        }

        return eventSignal.WaitAsync(timeout, cancellationToken);
    }

    /// <inheritdoc />
    public WorkerFault? DrainFault()
    {
        return null;
    }

    /// <summary>
    /// Gets a snapshot of every correlation id passed to
    /// <see cref="CancelCommand"/>. Recording lets the IPC tests
    /// assert that a <c>WorkerCancel</c> envelope dispatched on the
    /// gateway side reaches the runtime session.
    /// </summary>
    public IReadOnlyList<string> CancelledCorrelationIds
    {
        get
        {
            lock (gate)
            {
                return new List<string>(cancelledCorrelationIds);
            }
        }
    }

    private bool cancelCommandReturnValue;

    /// <summary>
    /// Optional return value yielded by <see cref="CancelCommand"/>.
    /// Defaults to <c>false</c> (the runtime had no matching in-flight
    /// command), matching the previous test-double behaviour. Mutated
    /// and read under <c>lock(gate)</c> to match the locking convention
    /// the rest of this fake uses for <c>cancelledCorrelationIds</c>,
    /// <c>snapshot</c>, and <c>events</c>.
    /// </summary>
    public bool CancelCommandReturnValue
    {
        get
        {
            lock (gate)
            {
                return cancelCommandReturnValue;
            }
        }

        set
        {
            lock (gate)
            {
                cancelCommandReturnValue = value;
            }
        }
    }

    /// <inheritdoc />
    public bool CancelCommand(string correlationId)
    {
        lock (gate)
        {
            cancelledCorrelationIds.Add(correlationId);
            return cancelCommandReturnValue;
        }
    }

    /// <inheritdoc />
    public void RequestShutdown()
    {
        releaseDispatch.Set();
    }

    /// <inheritdoc />
    public Task<MxAccessShutdownResult> ShutdownGracefullyAsync(
        TimeSpan timeout,
        CancellationToken cancellationToken = default)
    {
        releaseDispatch.Set();
        if (ThrowTimeoutOnShutdown)
        {
            return Task.FromException<MxAccessShutdownResult>(
                new TimeoutException("Simulated graceful shutdown timeout."));
        }

        return Task.FromResult(new MxAccessShutdownResult(Array.Empty<MxAccessShutdownFailure>()));
    }

    /// <summary>Releases a blocked dispatch.</summary>
    public void ReleaseDispatch()
    {
        releaseDispatch.Set();
    }

    /// <summary>Sets the current heartbeat snapshot.</summary>
    /// <param name="value">The snapshot to set.</param>
    public void SetSnapshot(WorkerRuntimeHeartbeatSnapshot value)
    {
        lock (gate)
        {
            snapshot = value;
        }
    }

    /// <summary>Enqueues a worker event to be drained.</summary>
    /// <param name="workerEvent">The event to enqueue.</param>
    public void EnqueueEvent(WorkerEvent workerEvent)
    {
        lock (gate)
        {
            events.Enqueue(workerEvent);
        }

        SignalWake();
    }

    /// <summary>
    ///     Enqueues a batch of worker events atomically under one lock so the drain loop cannot
    ///     observe a partial batch. Lets a test assert the drain loop coalesces a whole batch into one
    ///     flush (WRK-25) without racing a mid-enqueue drain that would split the batch.
    /// </summary>
    /// <param name="workerEvents">The events to enqueue in order.</param>
    public void EnqueueEvents(IEnumerable<WorkerEvent> workerEvents)
    {
        lock (gate)
        {
            foreach (WorkerEvent workerEvent in workerEvents)
            {
                events.Enqueue(workerEvent);
            }
        }

        SignalWake();
    }

    // Coalesced wake, released outside the gate exactly as MxAccessEventQueue does.
    private void SignalWake()
    {
        if (eventSignal.CurrentCount > 0)
        {
            return;
        }

        try
        {
            eventSignal.Release();
        }
        catch (SemaphoreFullException)
        {
            // A concurrent enqueue already published the pending wake this call wanted.
        }
    }

    /// <inheritdoc />
    public void Dispose()
    {
        Disposed = true;
        releaseDispatch.Set();
        releaseDispatch.Dispose();
        DispatchStarted.Dispose();
    }
}