mxaccessgw/src/ZB.MOM.WW.MxGateway.Worker/Sta/StaRuntime.cs

using System;
using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;

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

public sealed class StaRuntime : IDisposable
{
    private readonly IStaComApartmentInitializer comApartmentInitializer;
    private readonly IStaMessagePump messagePump;
    private readonly ConcurrentQueue<IStaWorkItem> commandQueue = new();
    private readonly AutoResetEvent commandWakeEvent = new(false);
    private readonly ManualResetEventSlim startedEvent = new(false);
    private readonly ManualResetEventSlim stoppedEvent = new(false);
    private readonly object gate = new();
    private readonly Thread staThread;
    private readonly TimeSpan idlePumpInterval;
    private bool disposed;
    private bool startRequested;
    private bool shutdownRequested;
    private Exception? startupException;
    private Exception? loopFault;
    private Action<Exception>? loopFaultHandler;
    private long lastActivityUtcTicks;
    private bool comInitialized;

    /// <summary>
    /// Initializes a new instance of <see cref="StaRuntime"/> with default dependencies.
    /// </summary>
    public StaRuntime()
        : this(new StaComApartmentInitializer(), new StaMessagePump(), TimeSpan.FromMilliseconds(50))
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="StaRuntime"/> with custom dependencies.
    /// </summary>
    /// <param name="comApartmentInitializer">COM apartment initializer.</param>
    /// <param name="messagePump">Message pump for the STA thread.</param>
    /// <param name="idlePumpInterval">Interval for idle message pump waits.</param>
    public StaRuntime(
        IStaComApartmentInitializer comApartmentInitializer,
        StaMessagePump messagePump,
        TimeSpan idlePumpInterval)
        : this(comApartmentInitializer, (IStaMessagePump)messagePump, idlePumpInterval)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="StaRuntime"/> over any
    /// <see cref="IStaMessagePump"/> (test seam).
    /// </summary>
    /// <param name="comApartmentInitializer">COM apartment initializer.</param>
    /// <param name="messagePump">Message pump for the STA thread.</param>
    /// <param name="idlePumpInterval">Interval for idle message pump waits.</param>
    internal StaRuntime(
        IStaComApartmentInitializer comApartmentInitializer,
        IStaMessagePump messagePump,
        TimeSpan idlePumpInterval)
    {
        this.comApartmentInitializer = comApartmentInitializer
            ?? throw new ArgumentNullException(nameof(comApartmentInitializer));
        this.messagePump = messagePump ?? throw new ArgumentNullException(nameof(messagePump));

        if (idlePumpInterval <= TimeSpan.Zero)
        {
            throw new ArgumentOutOfRangeException(
                nameof(idlePumpInterval),
                "The idle pump interval must be greater than zero.");
        }

        this.idlePumpInterval = idlePumpInterval;
        lastActivityUtcTicks = DateTimeOffset.UtcNow.UtcTicks;
        staThread = new Thread(ThreadMain)
        {
            IsBackground = true,
            Name = "MxGateway.Worker.STA"
        };
        staThread.SetApartmentState(ApartmentState.STA);
    }

    /// <summary>
    /// Gets the managed thread ID of the STA thread.
    /// </summary>
    public int? StaThreadId { get; private set; }

    /// <summary>
    /// Gets the timestamp of the last STA thread activity.
    /// </summary>
    public DateTimeOffset LastActivityUtc =>
        new(new DateTime(Volatile.Read(ref lastActivityUtcTicks), DateTimeKind.Utc));

    /// <summary>
    /// Gets a value indicating whether the STA runtime is currently running.
    /// </summary>
    public bool IsRunning => startedEvent.IsSet && !stoppedEvent.IsSet;

    /// <summary>
    /// Gets the exception that terminated the STA loop after a successful
    /// <see cref="Start"/>, or <see langword="null"/> while the loop is healthy
    /// (or stopped by a requested shutdown). Startup failures are not reported
    /// here — <see cref="Start"/> throws them.
    /// </summary>
    public Exception? LoopFault
    {
        get
        {
            lock (gate)
            {
                return loopFault;
            }
        }
    }

    /// <summary>
    /// Registers the callback invoked once, on the dying STA thread, when the
    /// STA loop terminates with an exception after a successful
    /// <see cref="Start"/>. When the loop has already died the callback runs
    /// immediately on the calling thread, so a fault that lands between
    /// <see cref="Start"/> and registration is never lost. Only one handler is
    /// kept; a later registration replaces an earlier one.
    /// </summary>
    /// <remarks>
    /// Without this, a post-start loop death (a failed
    /// <c>MsgWaitForMultipleObjectsEx</c>, or an exception escaping a dispatched
    /// window message) left the worker process alive with no STA: every later
    /// command queued forever and the gateway only learned of it through the
    /// heartbeat watchdog, if at all.
    /// </remarks>
    /// <param name="handler">Callback receiving the terminating exception.</param>
    public void RegisterLoopFaultHandler(Action<Exception> handler)
    {
        if (handler is null)
        {
            throw new ArgumentNullException(nameof(handler));
        }

        Exception? alreadyFaulted;
        lock (gate)
        {
            loopFaultHandler = handler;
            alreadyFaulted = loopFault;
        }

        if (alreadyFaulted is not null)
        {
            handler(alreadyFaulted);
        }
    }

    /// <summary>
    ///     Pumps any pending Windows messages on the calling thread and refreshes
    ///     the STA activity timestamp. Intended for commands that synchronously
    ///     hold the STA (e.g. ReadBulk) and must allow inbound MXAccess COM events
    ///     to dispatch while they wait. Because a long-running command invokes this
    ///     on every wait iteration, refreshing activity here keeps a busy STA from
    ///     being mistaken for a hung one: a healthy command that keeps pumping stays
    ///     fresh past <c>HeartbeatStuckCeiling</c>, while a genuinely stuck STA (no
    ///     pumping) still accrues staleness and faults correctly. Callers must
    ///     already be on the STA; the method is otherwise safe (PeekMessage simply
    ///     finds no messages).
    /// </summary>
    /// <returns>The number of messages pumped.</returns>
    public int PumpPendingMessages()
    {
        int pumpedMessages = messagePump.PumpPendingMessages();
        MarkActivity();
        return pumpedMessages;
    }

    /// <summary>
    /// Starts the STA thread.
    /// </summary>
    public void Start()
    {
        ThrowIfDisposed();

        lock (gate)
        {
            if (shutdownRequested)
            {
                throw new StaRuntimeShutdownException();
            }

            if (!startRequested)
            {
                startRequested = true;
                staThread.Start();
            }
        }

        startedEvent.Wait();
        if (startupException is not null)
        {
            throw new InvalidOperationException(
                "The worker STA runtime failed to initialize.",
                startupException);
        }
    }

    /// <summary>
    /// Invokes an action on the STA thread asynchronously.
    /// </summary>
    /// <param name="command">Action to invoke on the STA thread.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Task that completes when the action executes.</returns>
    public Task InvokeAsync(Action command, CancellationToken cancellationToken = default)
    {
        if (command is null)
        {
            throw new ArgumentNullException(nameof(command));
        }

        return InvokeAsync(
            () =>
            {
                command();
                return true;
            },
            cancellationToken);
    }

    /// <summary>
    /// Invokes a function on the STA thread asynchronously.
    /// </summary>
    /// <typeparam name="T">Return type of the function.</typeparam>
    /// <param name="command">Function to invoke on the STA thread.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Task that returns the function result.</returns>
    public Task<T> InvokeAsync<T>(Func<T> command, CancellationToken cancellationToken = default)
    {
        if (command is null)
        {
            throw new ArgumentNullException(nameof(command));
        }

        ThrowIfDisposed();

        if (cancellationToken.IsCancellationRequested)
        {
            return Task.FromCanceled<T>(cancellationToken);
        }

        StaWorkItem<T> workItem = new(command, cancellationToken);

        lock (gate)
        {
            if (loopFault is not null)
            {
                // The loop is gone; nothing would ever dequeue this item.
                return Task.FromException<T>(new InvalidOperationException(
                    "The worker STA loop terminated unexpectedly.",
                    loopFault));
            }

            if (shutdownRequested)
            {
                return Task.FromException<T>(new StaRuntimeShutdownException());
            }

            commandQueue.Enqueue(workItem);
        }

        commandWakeEvent.Set();
        return workItem.Task;
    }

    /// <summary>
    /// Requests graceful shutdown of the STA runtime within a timeout.
    /// </summary>
    /// <param name="timeout">Maximum time to wait for shutdown.</param>
    /// <returns>True if shutdown completed; otherwise false.</returns>
    public bool Shutdown(TimeSpan timeout)
    {
        if (timeout < TimeSpan.Zero && timeout != Timeout.InfiniteTimeSpan)
        {
            throw new ArgumentOutOfRangeException(nameof(timeout));
        }

        lock (gate)
        {
            shutdownRequested = true;
        }

        commandWakeEvent.Set();

        if (!startedEvent.IsSet && !staThread.IsAlive)
        {
            CancelQueuedCommands();
            stoppedEvent.Set();
            return true;
        }

        bool stopped = stoppedEvent.Wait(timeout);
        if (stopped)
        {
            CancelQueuedCommands();
        }

        return stopped;
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (disposed)
        {
            return;
        }

        bool stopped = Shutdown(TimeSpan.FromSeconds(5));
        if (stopped)
        {
            commandWakeEvent.Dispose();
            startedEvent.Dispose();
            stoppedEvent.Dispose();
        }

        disposed = true;
    }

    private void ThreadMain()
    {
        bool started = false;
        try
        {
            StaThreadId = Thread.CurrentThread.ManagedThreadId;
            comApartmentInitializer.Initialize();
            comInitialized = true;
            MarkActivity();
            started = true;
            startedEvent.Set();

            while (!IsShutdownRequested())
            {
                ProcessQueuedCommands();
                messagePump.WaitForWorkOrMessages(commandWakeEvent, idlePumpInterval);
                messagePump.PumpPendingMessages();
                MarkActivity();
            }

            ProcessQueuedCommands();
        }
        catch (Exception exception)
        {
            if (!started)
            {
                // Start() is blocked on startedEvent and rethrows this.
                startupException = exception;
                startedEvent.Set();
            }
            else
            {
                // Post-start loop death. Start() returned long ago, so nothing
                // would ever read startupException: record the fault where
                // InvokeAsync and the registered handler see it. Set under the
                // gate so every item enqueued before it is still in the queue
                // for the finally block's CancelQueuedCommands, and every
                // InvokeAsync after it fails fast instead of queueing forever.
                Action<Exception>? handler;
                lock (gate)
                {
                    loopFault = exception;
                    handler = loopFaultHandler;
                }

                if (handler is not null)
                {
                    try
                    {
                        handler(exception);
                    }
                    catch
                    {
                        // The handler reports the fault; a failure inside it
                        // must not skip the apartment cleanup below.
                    }
                }
            }
        }
        finally
        {
            CancelQueuedCommands();
            try
            {
                if (comInitialized)
                {
                    comApartmentInitializer.Uninitialize();
                }
            }
            finally
            {
                MarkActivity();
                stoppedEvent.Set();
            }
        }
    }

    private void ProcessQueuedCommands()
    {
        while (commandQueue.TryDequeue(out IStaWorkItem? workItem))
        {
            MarkActivity();
            workItem.Execute();
            MarkActivity();

            // Pump between work items, not only once the whole queue is
            // drained: a burst of queued commands (bulk callers, the alarm
            // poll, a reconnecting client replaying subscriptions) would
            // otherwise hold every inbound COM callback — OnDataChange,
            // OnWriteComplete — behind the entire backlog. PeekMessage with an
            // empty queue is cheap, so the per-item cost is negligible.
            if (!commandQueue.IsEmpty)
            {
                messagePump.PumpPendingMessages();
            }
        }
    }

    private void CancelQueuedCommands()
    {
        while (commandQueue.TryDequeue(out IStaWorkItem? workItem))
        {
            workItem.CancelBeforeExecution();
        }
    }

    private bool IsShutdownRequested()
    {
        lock (gate)
        {
            return shutdownRequested;
        }
    }

    private void MarkActivity()
    {
        Volatile.Write(ref lastActivityUtcTicks, DateTimeOffset.UtcNow.UtcTicks);
    }

    private void ThrowIfDisposed()
    {
        if (disposed)
        {
            throw new ObjectDisposedException(nameof(StaRuntime));
        }
    }
}