mxaccessgw/src/ZB.MOM.WW.MxGateway.Worker/MxAccess/MxAccessWriteCompletionCache.cs

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

namespace ZB.MOM.WW.MxGateway.Worker.MxAccess;

/// <summary>
///     Per-session cache of the most recent <c>OnWriteComplete</c> status rows
///     for each (server handle, item handle) pair. Written by the MXAccess
///     event sink as completion callbacks arrive; read by the write command
///     executor so a WriteSecured/WriteSecured2 reply can carry the correlated
///     MXAccess outcome instead of proving command acceptance only.
/// </summary>
/// <remarks>
///     Same threading posture as <see cref="MxAccessValueCache"/>: writers and
///     readers run on the worker's STA thread (COM dispatches events on the
///     apartment thread; commands also execute on the STA), so no internal
///     locking is required. A single sync root keeps it nominally thread-safe
///     for tests that drive it from a non-STA thread.
/// </remarks>
public sealed class MxAccessWriteCompletionCache
{
    private readonly Dictionary<long, CompletionEntry> entries = new();
    private readonly object syncRoot = new();

    // Set by Record so a waiter blocked in TryWaitForCompletion wakes the moment
    // a completion lands instead of after the next blind poll tick. In the live
    // worker the completion is recorded from inside the waiter's own pumpStep
    // (COM dispatches OnWriteComplete on the STA), so the message wake is what
    // does the work there; this handle covers the cross-thread producers — the
    // tests, and any future off-STA recorder.
    //
    // Single-waiter invariant: an AutoResetEvent releases exactly one waiter per
    // Set. Today the sole waiter is the STA (one worker per session, and the
    // write path holds the apartment while it waits), so no wake is ever lost. A
    // second concurrent waiter would not hang — the loser still re-checks the
    // entry every fallback tick and still honors its deadline — it would just
    // fall back to 50 ms latency for that iteration.
    //
    // NOT IDisposable, and that is one invariant with the wait's use of
    // DangerousGetHandle: the cache is co-owned by MxAccessBaseEventSink and
    // MxAccessSession with no single teardown owner, so a late COM callback
    // could otherwise Set a disposed handle. StaWaitHelper deliberately does not
    // DangerousAddRef the handle; GC.KeepAlive there closes the *finalization*
    // race but not a dispose race — if this field were ever disposed while the
    // STA sat inside the wait, the OS could recycle that handle value and the
    // wait would silently attach to an unrelated object. Making these caches
    // IDisposable is therefore not a free refactor: it must add DangerousAddRef/
    // DangerousRelease around the wait in the same change. One kernel event per
    // session (one session per worker process), released by the finalizer.
    private readonly AutoResetEvent completionSignal = new(false);

    /// <summary>Records the status rows of a fresh OnWriteComplete callback for the given handle pair.</summary>
    /// <param name="serverHandle">MXAccess server handle.</param>
    /// <param name="itemHandle">MXAccess item handle.</param>
    /// <param name="statuses">Status rows from the mapped OnWriteComplete event; cloned before storing.</param>
    public void Record(
        int serverHandle,
        int itemHandle,
        RepeatedField<MxStatusProxy> statuses)
    {
        if (statuses is null)
        {
            throw new ArgumentNullException(nameof(statuses));
        }

        lock (syncRoot)
        {
            long key = CreateItemKey(serverHandle, itemHandle);
            ulong version = entries.TryGetValue(key, out CompletionEntry existing)
                ? existing.Version + 1
                : 1UL;

            // Still a defensive copy, deliberately — this is NOT an oversight
            // left behind by the borrow that MxAccessValueCache.Set adopted (see
            // the ownership contract there). The value cache is handed the whole
            // enqueued MxEvent, so the queue's write-once ownership invariant
            // covers everything it retains. This method is handed a bare
            // RepeatedField whose provenance its signature cannot constrain:
            // the production sink does pass an enqueued event's Statuses, but
            // callers that build and keep their own rows are equally valid
            // against this API, and a borrowed alias would then let a later
            // mutation rewrite an already-recorded completion. The write path is
            // command-rate, not the per-OnDataChange streaming hot path, so the
            // clone costs nothing worth reclaiming.
            entries[key] = new CompletionEntry(version, statuses.Clone());
        }

        // Signaled outside the lock: the waiter re-takes syncRoot the instant it
        // wakes, so signaling while holding it would hand it guaranteed contention.
        completionSignal.Set();
    }

    /// <summary>Returns the current completion version for a handle pair, or 0 if none was recorded.</summary>
    /// <param name="serverHandle">MXAccess server handle.</param>
    /// <param name="itemHandle">MXAccess item handle.</param>
    /// <returns>The current completion version, or 0 if no completion was recorded.</returns>
    public ulong CurrentVersion(
        int serverHandle,
        int itemHandle)
    {
        lock (syncRoot)
        {
            return entries.TryGetValue(CreateItemKey(serverHandle, itemHandle), out CompletionEntry existing)
                ? existing.Version
                : 0UL;
        }
    }

    /// <summary>
    ///     Waits for a completion newer than <paramref name="sinceVersion"/> until it
    ///     arrives or the deadline elapses, calling <paramref name="pumpStep"/> on every
    ///     iteration so the worker's STA can dispatch the inbound MXAccess
    ///     OnWriteComplete message. Same loop shape as
    ///     <see cref="MxAccessValueCache.TryWaitForUpdate"/>.
    /// </summary>
    /// <remarks>
    ///     The idle part of each iteration blocks in
    ///     <see cref="StaWaitHelper.WaitForSignalOrMessages"/> rather than
    ///     <c>Thread.Sleep</c>: a sleeping STA pumps no Windows messages, so the
    ///     inbound OnWriteComplete could only ever be dispatched at poll-tick
    ///     granularity. The Win32 wait returns the instant a message needs
    ///     pumping (or the completion is recorded from another thread), so the
    ///     apartment keeps dispatching COM events continuously for the whole
    ///     wait. Timeouts, deadline math, and both return shapes are unchanged.
    /// </remarks>
    /// <param name="serverHandle">MXAccess server handle.</param>
    /// <param name="itemHandle">MXAccess item handle.</param>
    /// <param name="sinceVersion">Version snapshot captured before the write COM call.</param>
    /// <param name="deadlineUtc">Absolute UTC deadline.</param>
    /// <param name="pumpStep">Action that pumps any pending Windows messages.</param>
    /// <param name="statuses">The recorded status rows if a completion arrived before the deadline; empty otherwise.</param>
    /// <param name="pollIntervalMs">Upper bound on one idle wait slice — the cadence at which <paramref name="pumpStep"/> runs when nothing wakes the wait. Default 5 ms, itself capped at <see cref="StaWaitHelper.MaxFallbackTickMilliseconds"/>.</param>
    /// <returns><see langword="true"/> if a completion newer than <paramref name="sinceVersion"/> arrived before the deadline; otherwise <see langword="false"/>.</returns>
    public bool TryWaitForCompletion(
        int serverHandle,
        int itemHandle,
        ulong sinceVersion,
        DateTime deadlineUtc,
        Action pumpStep,
        out RepeatedField<MxStatusProxy> statuses,
        int pollIntervalMs = 5)
    {
        if (pumpStep is null)
        {
            throw new ArgumentNullException(nameof(pumpStep));
        }

        while (true)
        {
            // Pumped unconditionally at the top of every iteration — including
            // the one that immediately follows a signaled wait, and the very
            // first before any waiting. Dispatching the message is what fills
            // the cache in the live worker, and unit tests stand in for the STA
            // by recording the completion from inside their fake pumpStep, so
            // this call must never be conditional on a real message queue.
            pumpStep();

            lock (syncRoot)
            {
                if (entries.TryGetValue(CreateItemKey(serverHandle, itemHandle), out CompletionEntry entry)
                    && entry.Version > sinceVersion)
                {
                    statuses = entry.Statuses;
                    return true;
                }
            }

            // Same expiry test as the sleep-polled version: remaining <= zero is
            // exactly DateTime.UtcNow >= deadlineUtc.
            TimeSpan remaining = deadlineUtc - DateTime.UtcNow;
            if (remaining <= TimeSpan.Zero)
            {
                statuses = new RepeatedField<MxStatusProxy>();
                return false;
            }

            StaWaitHelper.WaitForSignalOrMessages(
                completionSignal,
                StaWaitHelper.ClampWaitMilliseconds(remaining, pollIntervalMs));
        }
    }

    private static long CreateItemKey(
        int serverHandle,
        int itemHandle)
    {
        return ((long)serverHandle << 32) | (uint)itemHandle;
    }

    /// <summary>
    ///     Snapshot of the most recent OnWriteComplete status rows for a handle
    ///     pair. <see cref="Version"/> increments by one on every
    ///     <see cref="Record"/> call so the write executor can detect "a new
    ///     completion arrived since I captured my baseline".
    /// </summary>
    /// <remarks>
    ///     Plain readonly struct (not a record) so this compiles under the
    ///     worker's net48 target, which lacks <c>IsExternalInit</c>.
    /// </remarks>
    private readonly struct CompletionEntry
    {
        public CompletionEntry(
            ulong version,
            RepeatedField<MxStatusProxy> statuses)
        {
            Version = version;
            Statuses = statuses;
        }

        public ulong Version { get; }

        public RepeatedField<MxStatusProxy> Statuses { get; }
    }
}