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 StaMessagePump 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 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 = 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>
    ///     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 (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()
    {
        try
        {
            StaThreadId = Thread.CurrentThread.ManagedThreadId;
            comApartmentInitializer.Initialize();
            comInitialized = true;
            MarkActivity();
            startedEvent.Set();

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

            ProcessQueuedCommands();
        }
        catch (Exception exception)
        {
            startupException = exception;
            startedEvent.Set();
        }
        finally
        {
            CancelQueuedCommands();
            try
            {
                if (comInitialized)
                {
                    comApartmentInitializer.Uninitialize();
                }
            }
            finally
            {
                MarkActivity();
                stoppedEvent.Set();
            }
        }
    }

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

    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));
        }
    }
}