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

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;
using ZB.MOM.WW.MxGateway.Contracts.Proto;
using ZB.MOM.WW.MxGateway.Worker.Conversion;
using ZB.MOM.WW.MxGateway.Worker.Sta;

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

public sealed class MxAccessStaSession : IWorkerRuntimeSession
{
    private static readonly TimeSpan AlarmPollInterval = TimeSpan.FromMilliseconds(500);

    private readonly IMxAccessComObjectFactory factory;
    private readonly IMxAccessEventSink eventSink;
    private readonly MxAccessEventQueue eventQueue;
    private readonly StaRuntime staRuntime;
    // Worker-024: the factory takes an Action so MxAccessStaSession can hand
    // the alarm handler its STA-affinity guard (a closure over
    // alarmConsumerThreadId captured at the factory call site). The handler
    // then invokes the guard at the entry of every method that touches the
    // wnwrap consumer, matching the STA-affinity invariant already enforced
    // for the poll path via EnsureOnAlarmConsumerThread.
    private readonly Func<MxAccessEventQueue, Action, IAlarmCommandHandler>? alarmCommandHandlerFactory;
    private StaCommandDispatcher? commandDispatcher;
    private MxAccessSession? session;
    private IAlarmCommandHandler? alarmCommandHandler;
    private CancellationTokenSource? alarmPollCts;
    private Task? alarmPollTask;
    private int? alarmConsumerThreadId;
    private bool disposed;

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with default dependencies.
    /// </summary>
    public MxAccessStaSession()
        : this(
            new StaRuntime(),
            new MxAccessComObjectFactory(),
            new MxAccessEventQueue())
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with default STA runtime,
    /// factory, and event queue, but with a custom alarm-command handler factory. The factory is
    /// invoked on the STA thread during
    /// <see cref="StartAsync(string, int, CancellationToken)"/>; pass <c>null</c> to opt out
    /// of alarm-side commands.
    /// </summary>
    internal MxAccessStaSession(Func<MxAccessEventQueue, Action, IAlarmCommandHandler>? alarmCommandHandlerFactory)
        : this(
            new StaRuntime(),
            new MxAccessComObjectFactory(),
            new MxAccessEventQueue(),
            alarmCommandHandlerFactory)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with custom STA runtime and factory.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventSink">Event sink for MXAccess events.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        IMxAccessEventSink eventSink)
        : this(staRuntime, factory, eventSink, new MxAccessEventQueue())
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with custom event queue.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        MxAccessEventQueue eventQueue)
        : this(staRuntime, factory, new MxAccessBaseEventSink(eventQueue), eventQueue)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with custom event queue
    /// and an alarm-command handler factory.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    /// <param name="alarmCommandHandlerFactory">
    ///     Factory that constructs the alarm-command handler from the event queue.
    ///     Pass <c>null</c> to opt out of alarm-side commands.
    /// </param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        MxAccessEventQueue eventQueue,
        Func<MxAccessEventQueue, Action, IAlarmCommandHandler>? alarmCommandHandlerFactory)
        : this(staRuntime, factory, new MxAccessBaseEventSink(eventQueue), eventQueue, alarmCommandHandlerFactory)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with all dependencies.
    /// </summary>
    /// <param name="staRuntime">STA thread runtime.</param>
    /// <param name="factory">MXAccess COM object factory.</param>
    /// <param name="eventSink">Event sink for MXAccess events.</param>
    /// <param name="eventQueue">Event queue for buffering MXAccess events.</param>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        IMxAccessEventSink eventSink,
        MxAccessEventQueue eventQueue)
        : this(staRuntime, factory, eventSink, eventQueue, alarmCommandHandlerFactory: null)
    {
    }

    /// <summary>
    /// Initializes a new instance of <see cref="MxAccessStaSession"/> with all
    /// dependencies including an alarm-command handler factory. The factory is
    /// invoked on the STA thread during <see cref="StartAsync(string, int, CancellationToken)"/>;
    /// pass <c>null</c> to opt out of alarm-side commands (the worker rejects
    /// them with an "alarm consumer not configured" diagnostic).
    /// </summary>
    public MxAccessStaSession(
        StaRuntime staRuntime,
        IMxAccessComObjectFactory factory,
        IMxAccessEventSink eventSink,
        MxAccessEventQueue eventQueue,
        Func<MxAccessEventQueue, Action, IAlarmCommandHandler>? alarmCommandHandlerFactory)
    {
        this.staRuntime = staRuntime ?? throw new ArgumentNullException(nameof(staRuntime));
        this.factory = factory ?? throw new ArgumentNullException(nameof(factory));
        this.eventSink = eventSink ?? throw new ArgumentNullException(nameof(eventSink));
        this.eventQueue = eventQueue ?? throw new ArgumentNullException(nameof(eventQueue));
        this.alarmCommandHandlerFactory = alarmCommandHandlerFactory;
    }

    /// <summary>
    /// Gets the event queue for this session.
    /// </summary>
    public MxAccessEventQueue EventQueue => eventQueue;

    /// <summary>
    /// Starts the MXAccess COM session asynchronously.
    /// </summary>
    /// <param name="workerProcessId">Worker process identifier.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Worker ready message.</returns>
    public Task<WorkerReady> StartAsync(
        int workerProcessId,
        CancellationToken cancellationToken = default)
    {
        return StartAsync(string.Empty, workerProcessId, cancellationToken);
    }

    /// <summary>
    /// Starts the MXAccess COM session with a session ID asynchronously.
    /// </summary>
    /// <param name="sessionId">Session identifier.</param>
    /// <param name="workerProcessId">Worker process identifier.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Worker ready message.</returns>
    public async Task<WorkerReady> StartAsync(
        string sessionId,
        int workerProcessId,
        CancellationToken cancellationToken = default)
    {
        staRuntime.Start();

        WorkerReady ready = await staRuntime.InvokeAsync(
            () =>
            {
                if (session is not null)
                {
                    throw new InvalidOperationException("MXAccess COM session has already been created.");
                }

                session = MxAccessSession.Create(factory, eventSink, sessionId);
                if (alarmCommandHandlerFactory is not null)
                {
                    // STA-affinity invariant: the alarm consumer factory and
                    // every IMxAccessAlarmConsumer call must run on the STA
                    // thread, because the production wnwrap consumer holds an
                    // Apartment-threaded COM object. The factory runs here
                    // inside staRuntime.InvokeAsync, so this records the STA
                    // thread id; RunAlarmPollLoopAsync then asserts each
                    // PollOnce executes on the same thread.
                    alarmConsumerThreadId = Environment.CurrentManagedThreadId;
                    // Worker-024: hand the handler an affinity guard so each
                    // of its command-path entries (Subscribe / Acknowledge /
                    // AcknowledgeByName / QueryActive / Unsubscribe / PollOnce)
                    // asserts the same STA-affinity invariant the poll path
                    // already enforced. Without this the command path relied
                    // on convention alone; a future refactor that let a
                    // command run off-STA would silently deadlock on
                    // cross-apartment marshaling against the wnwrap consumer.
                    alarmCommandHandler = alarmCommandHandlerFactory(
                        eventQueue,
                        EnsureOnAlarmConsumerThread);
                }
                commandDispatcher = new StaCommandDispatcher(
                    staRuntime,
                    new MxAccessCommandExecutor(
                        session,
                        new VariantConverter(),
                        alarmCommandHandler,
                        // ReadBulk needs to pump Windows messages while it waits
                        // for the first OnDataChange callback so the inbound COM
                        // event can dispatch on this same STA thread. The pump
                        // step closes over staRuntime so it always pumps the
                        // pump tied to the apartment that owns this session.
                        pumpStep: () => staRuntime.PumpPendingMessages()));

                return session.CreateWorkerReady(workerProcessId);
            },
            cancellationToken).ConfigureAwait(false);

        if (alarmCommandHandler is not null)
        {
            alarmPollCts = new CancellationTokenSource();
            alarmPollTask = RunAlarmPollLoopAsync(alarmCommandHandler, alarmPollCts.Token);
        }

        return ready;
    }

    private Task RunAlarmPollLoopAsync(
        IAlarmCommandHandler handler,
        CancellationToken cancellationToken)
    {
        return Task.Run(async () =>
        {
            while (!cancellationToken.IsCancellationRequested)
            {
                try
                {
                    await Task.Delay(AlarmPollInterval, cancellationToken).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    return;
                }

                if (cancellationToken.IsCancellationRequested)
                {
                    return;
                }

                try
                {
                    await staRuntime.InvokeAsync(
                        () =>
                        {
                            EnsureOnAlarmConsumerThread();
                            handler.PollOnce();
                        },
                        cancellationToken).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    return;
                }
                catch (ObjectDisposedException)
                {
                    // STA runtime or alarm handler disposed — stop the loop gracefully.
                    return;
                }
                catch (StaRuntimeShutdownException)
                {
                    // STA runtime shutting down — stop the loop gracefully.
                    // The dedicated shutdown type lets us distinguish this
                    // graceful-stop signal from the STA-affinity assertion
                    // raised by EnsureOnAlarmConsumerThread (Worker-008),
                    // which is also an InvalidOperationException but signals
                    // a programming-error regression — that case falls through
                    // to the generic Exception arm below and is recorded as a
                    // fault on the event queue, so an affinity regression
                    // becomes observable on the IPC fault path instead of
                    // silently stopping alarm delivery.
                    return;
                }
                catch (Exception exception)
                {
                    // A real alarm-poll failure (COMException from
                    // GetXmlCurrentAlarms2, malformed-XML parse failure, an
                    // STA-affinity InvalidOperationException from
                    // EnsureOnAlarmConsumerThread, etc.). Record it as a
                    // fault on the event queue so a broken alarm subscription
                    // — or an affinity-invariant regression — becomes
                    // observable on the IPC fault path instead of silently
                    // faulting this never-awaited task. The loop then stops —
                    // the subscription is dead.
                    eventQueue.RecordFault(CreateAlarmPollFault(exception));
                    return;
                }
            }
        }, CancellationToken.None);
    }

    private void EnsureOnAlarmConsumerThread()
    {
        AssertOnAlarmConsumerThread(alarmConsumerThreadId, Environment.CurrentManagedThreadId);
    }

    /// <summary>
    /// Enforces the STA-affinity invariant for the alarm consumer: every
    /// <see cref="IMxAccessAlarmConsumer"/> call (and the consumer factory)
    /// must run on the same thread the consumer was created on (the worker's
    /// STA). Throws <see cref="InvalidOperationException"/> when a caller
    /// breaks affinity — a programming error that would otherwise risk a
    /// cross-apartment COM deadlock in the production wnwrap consumer, since
    /// its CLSID is registered <c>ThreadingModel=Apartment</c>. The check is
    /// a no-op until the consumer thread has been recorded (no alarm handler
    /// configured, or session not yet started).
    /// </summary>
    /// <param name="expectedThreadId">
    /// The managed thread id the alarm consumer was created on, or
    /// <c>null</c> if no alarm consumer is configured.
    /// </param>
    /// <param name="actualThreadId">The current managed thread id.</param>
    internal static void AssertOnAlarmConsumerThread(int? expectedThreadId, int actualThreadId)
    {
        if (expectedThreadId is not null && actualThreadId != expectedThreadId.Value)
        {
            throw new InvalidOperationException(
                $"Alarm consumer accessed off its owning STA thread. Expected thread {expectedThreadId.Value}, "
                + $"actual {actualThreadId}. All IMxAccessAlarmConsumer calls must run on the STA that "
                + "created the consumer.");
        }
    }

    private static WorkerFault CreateAlarmPollFault(Exception exception)
    {
        string message =
            $"MXAccess alarm poll failed: {exception.Message}";
        WorkerFault fault = new()
        {
            Category = WorkerFaultCategory.MxaccessEventConversionFailed,
            ExceptionType = exception.GetType().FullName ?? string.Empty,
            DiagnosticMessage = message,
            ProtocolStatus = new ProtocolStatus
            {
                Code = ProtocolStatusCode.WorkerUnavailable,
                Message = message,
            },
        };

        if (exception is System.Runtime.InteropServices.COMException comException)
        {
            fault.Hresult = comException.HResult;
        }

        return fault;
    }

    /// <summary>
    /// Dispatches a command to the STA thread for execution asynchronously.
    /// </summary>
    /// <param name="command">The command to dispatch.</param>
    /// <returns>Command reply.</returns>
    public Task<MxCommandReply> DispatchAsync(StaCommand command)
    {
        if (commandDispatcher is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return commandDispatcher.DispatchAsync(command);
    }

    /// <summary>
    /// Captures a heartbeat snapshot of the session's runtime state.
    /// </summary>
    /// <returns>Heartbeat snapshot.</returns>
    public WorkerRuntimeHeartbeatSnapshot CaptureHeartbeat()
    {
        uint pendingCommandCount = 0;
        string currentCommandCorrelationId = string.Empty;

        if (commandDispatcher is not null)
        {
            pendingCommandCount = (uint)commandDispatcher.PendingCommandCount;
            currentCommandCorrelationId = commandDispatcher.CurrentCommandCorrelationId;
        }

        return new WorkerRuntimeHeartbeatSnapshot(
            staRuntime.LastActivityUtc,
            pendingCommandCount,
            (uint)eventQueue.Count,
            eventQueue.LastEventSequence,
            currentCommandCorrelationId);
    }

    /// <summary>
    /// Requests graceful shutdown of the command dispatcher.
    /// </summary>
    public void RequestShutdown()
    {
        commandDispatcher?.RequestShutdown();
    }

    /// <summary>
    /// Drains up to the specified number of events from the queue.
    /// </summary>
    /// <param name="maxEvents">Maximum events to drain.</param>
    /// <returns>Drained events.</returns>
    public IReadOnlyList<WorkerEvent> DrainEvents(uint maxEvents)
    {
        return eventQueue.Drain(maxEvents);
    }

    /// <summary>
    /// Drains a fault from the queue if present.
    /// </summary>
    /// <returns>Drained fault or null.</returns>
    public WorkerFault? DrainFault()
    {
        return eventQueue.DrainFault();
    }

    /// <summary>
    /// Cancels a queued command by correlation ID.
    /// </summary>
    /// <param name="correlationId">Correlation ID of the command to cancel.</param>
    /// <returns>True if cancelled; otherwise false.</returns>
    public bool CancelCommand(string correlationId)
    {
        return commandDispatcher?.CancelQueuedCommand(correlationId) ?? false;
    }

    /// <summary>
    /// Gets the registered server handles asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Registered server handles.</returns>
    public Task<IReadOnlyList<RegisteredServerHandle>> GetRegisteredServerHandlesAsync(
        CancellationToken cancellationToken = default)
    {
        if (session is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return staRuntime.InvokeAsync(
            () => session.HandleRegistry.ServerHandles,
            cancellationToken);
    }

    /// <summary>
    /// Gets the registered item handles asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Registered item handles.</returns>
    public Task<IReadOnlyList<RegisteredItemHandle>> GetRegisteredItemHandlesAsync(
        CancellationToken cancellationToken = default)
    {
        if (session is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return staRuntime.InvokeAsync(
            () => session.HandleRegistry.ItemHandles,
            cancellationToken);
    }

    /// <summary>
    /// Gets the registered advice handles asynchronously.
    /// </summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Registered advice handles.</returns>
    public Task<IReadOnlyList<RegisteredAdviceHandle>> GetRegisteredAdviceHandlesAsync(
        CancellationToken cancellationToken = default)
    {
        if (session is null)
        {
            throw new InvalidOperationException("MXAccess COM session has not been started.");
        }

        return staRuntime.InvokeAsync(
            () => session.HandleRegistry.AdviceHandles,
            cancellationToken);
    }

    /// <summary>
    /// Performs graceful shutdown of the MXAccess session within a timeout.
    /// </summary>
    /// <param name="timeout">Maximum time allowed for shutdown.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <returns>Shutdown result with any cleanup failures.</returns>
    public async Task<MxAccessShutdownResult> ShutdownGracefullyAsync(
        TimeSpan timeout,
        CancellationToken cancellationToken = default)
    {
        if (timeout <= TimeSpan.Zero)
        {
            throw new ArgumentOutOfRangeException(
                nameof(timeout),
                "MXAccess graceful shutdown timeout must be greater than zero.");
        }

        if (disposed)
        {
            return new MxAccessShutdownResult(Array.Empty<MxAccessShutdownFailure>());
        }

        commandDispatcher?.RequestShutdown();

        // Cancel the STA poll loop before disposing the alarm handler.
        // The loop references the alarm handler and must be stopped first
        // so that no further PollOnce calls race with disposal.
        CancellationTokenSource? pollCtsToDispose = alarmPollCts;
        Task? pollTaskToAwait = alarmPollTask;
        alarmPollCts = null;
        alarmPollTask = null;
        if (pollCtsToDispose is not null)
        {
            pollCtsToDispose.Cancel();
            if (pollTaskToAwait is not null)
            {
                try
                {
                    await pollTaskToAwait.ConfigureAwait(false);
                }
                catch
                {
                    // Swallow — poll loop cancellation must not block data shutdown.
                }
            }
            pollCtsToDispose.Dispose();
        }

        // Stop the alarm consumer's polling timer and tear down the
        // dispatcher BEFORE the data-side cleanup begins. The alarm
        // consumer holds a wnwrap COM RCW that needs the STA pump to
        // unwind cleanly; doing it here gives it the opportunity while
        // the STA is still alive.
        IAlarmCommandHandler? alarmHandlerToDispose = alarmCommandHandler;
        alarmCommandHandler = null;
        if (alarmHandlerToDispose is not null)
        {
            try
            {
                await staRuntime.InvokeAsync(
                    () => alarmHandlerToDispose.Dispose(),
                    cancellationToken).ConfigureAwait(false);
            }
            catch
            {
                // Swallow — alarm cleanup must not block data shutdown.
            }
        }

        Stopwatch stopwatch = Stopwatch.StartNew();
        MxAccessShutdownResult result;
        if (session is null)
        {
            result = new MxAccessShutdownResult(Array.Empty<MxAccessShutdownFailure>());
        }
        else
        {
            using CancellationTokenSource shutdownCancellation =
                CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            shutdownCancellation.CancelAfter(timeout);

            Task<MxAccessShutdownResult> cleanupTask = staRuntime.InvokeAsync(
                () => session.ShutdownGracefully(),
                shutdownCancellation.Token);
            Task delayTask = Task.Delay(timeout, cancellationToken);
            Task completedTask = await Task.WhenAny(cleanupTask, delayTask).ConfigureAwait(false);
            if (completedTask != cleanupTask)
            {
                cancellationToken.ThrowIfCancellationRequested();
                throw new TimeoutException($"MXAccess graceful shutdown exceeded {timeout}.");
            }

            try
            {
                result = await cleanupTask.ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
            {
                throw new TimeoutException($"MXAccess graceful shutdown exceeded {timeout}.");
            }
        }

        TimeSpan remaining = timeout - stopwatch.Elapsed;
        if (remaining <= TimeSpan.Zero || !staRuntime.Shutdown(remaining))
        {
            throw new TimeoutException($"MXAccess graceful shutdown exceeded {timeout}.");
        }

        staRuntime.Dispose();
        disposed = true;
        return result;
    }

    /// <summary>Releases resources and shuts down the session.</summary>
    public void Dispose()
    {
        if (disposed)
        {
            return;
        }

        RequestShutdown();

        // Cancel the STA poll loop and join it before disposing the alarm
        // handler. Joining (rather than discarding alarmPollTask) makes the
        // stop deterministic: once Dispose returns, no further PollOnce calls
        // can be in flight, so callers and tests can rely on a frozen poll
        // count instead of an elapsed-time "no further polls" window.
        CancellationTokenSource? pollCtsToDispose = alarmPollCts;
        Task? pollTaskToJoin = alarmPollTask;
        alarmPollCts = null;
        alarmPollTask = null;
        if (pollCtsToDispose is not null)
        {
            try { pollCtsToDispose.Cancel(); } catch { }
            if (pollTaskToJoin is not null)
            {
                try
                {
                    pollTaskToJoin.Wait(TimeSpan.FromSeconds(5));
                }
                catch (AggregateException) { }
                catch (ObjectDisposedException) { }
            }
            try { pollCtsToDispose.Dispose(); } catch { }
        }

        IAlarmCommandHandler? alarmHandlerToDispose = alarmCommandHandler;
        alarmCommandHandler = null;
        if (alarmHandlerToDispose is not null)
        {
            try
            {
                staRuntime.InvokeAsync(() => alarmHandlerToDispose.Dispose())
                    .Wait(TimeSpan.FromSeconds(2));
            }
            catch (AggregateException) { }
            catch (ObjectDisposedException) { }
        }

        if (session is not null)
        {
            try
            {
                staRuntime.InvokeAsync(() => session.Dispose())
                    .Wait(TimeSpan.FromSeconds(2));
            }
            catch (AggregateException)
            {
            }
            catch (ObjectDisposedException)
            {
            }
        }

        staRuntime.Dispose();
        disposed = true;
    }
}