using System.Collections.Concurrent; using System.Runtime.CompilerServices; using System.Threading.Channels; using Google.Protobuf.WellKnownTypes; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging.Abstractions; using ZB.MOM.WW.MxGateway.Contracts; using ZB.MOM.WW.MxGateway.Contracts.Proto; using ZB.MOM.WW.MxGateway.Server.Metrics; namespace ZB.MOM.WW.MxGateway.Server.Workers; public sealed class WorkerClient : IWorkerClient { private const string GatewayVersionFallback = "unknown"; private static readonly TimeSpan DisposeTaskTimeout = TimeSpan.FromSeconds(5); private readonly object _syncRoot = new(); private readonly WorkerClientConnection _connection; private readonly WorkerClientOptions _options; private readonly GatewayMetrics? _metrics; private readonly TimeProvider _timeProvider; private readonly ILogger _logger; private readonly WorkerFrameReader _reader; private readonly WorkerFrameWriter _writer; private readonly Channel _outboundEnvelopes; private readonly Channel _events; // Staging hand-off between the read loop and the dedicated event writer. The read loop writes // here with a non-blocking TryWrite so a full consumer channel (_events) can never stall the read // loop behind an event — replies and heartbeats keep flowing. Bounded at 2 × EventChannelCapacity // (GWC-24): an unbounded staging channel let a consumer that drains slower than the worker // produces grow gateway memory without limit and without a fault, because each individual timed // write into _events still completed inside EventChannelFullModeTimeout. A rejected TryWrite here // is the sustained-slow-drain signal and faults the client immediately (ProtocolViolation), which // is the fail-fast backpressure policy in docs/DesignDecisions.md. Total gateway-side buffering // per session is therefore 3 × MxGateway:Events:QueueCapacity, and the single _eventQueueDepth // gauge covers staged + queued events so the whole backlog is observable. private readonly Channel _eventStaging; private readonly int _eventStagingCapacity; private readonly ConcurrentDictionary _pendingCommands = new(StringComparer.Ordinal); private readonly SemaphoreSlim _pendingCommandSlots; private readonly CancellationTokenSource _stopCts = new(); // GWC-30: this client's contribution to the gateway-wide worker queue-depth gauge. Registered // once here and read only when the gauge is scraped, so staging and consuming an event cost an // Interlocked on _eventQueueDepth and nothing else — previously each of those two hot-path steps // called into GatewayMetrics and took its process-wide lock. Null when metrics are disabled. private readonly IDisposable? _eventQueueDepthRegistration; // Touched only by WriteLoopAsync — the single consumer of _outboundEnvelopes — so it needs no // interlocking. See WriteLoopAsync for why the stamp happens there rather than at construction. private ulong _nextSequence; private WorkerClientState _state; private DateTimeOffset _lastHeartbeatAt; private int? _processId; private int _eventQueueDepth; private int _eventsReaderClaimed; private Task? _readLoopTask; private Task? _writeLoopTask; private Task? _eventWriteLoopTask; private Task? _heartbeatLoopTask; private bool _workerStartRecorded; private bool _disposed; /// Initializes a client for communicating with a worker process over a named pipe. /// Named pipe connection to the worker. /// Worker client configuration; defaults to WorkerClientOptions if null. /// Gateway metrics sink; null disables metrics recording. /// Time provider for timestamps; defaults to system time if null. /// Logger instance; defaults to NullLogger if null. public WorkerClient( WorkerClientConnection connection, WorkerClientOptions? options = null, GatewayMetrics? metrics = null, TimeProvider? timeProvider = null, ILogger? logger = null) { _connection = connection ?? throw new ArgumentNullException(nameof(connection)); _options = options ?? new WorkerClientOptions(); _metrics = metrics; _timeProvider = timeProvider ?? TimeProvider.System; _logger = logger ?? NullLogger.Instance; _reader = new WorkerFrameReader(connection.Stream, connection.FrameOptions); _writer = new WorkerFrameWriter(connection.Stream, connection.FrameOptions); _pendingCommandSlots = new SemaphoreSlim(_options.MaxPendingCommands, _options.MaxPendingCommands); _outboundEnvelopes = Channel.CreateBounded( new BoundedChannelOptions(_options.MaxPendingCommands + 4) { SingleReader = true, SingleWriter = false, FullMode = BoundedChannelFullMode.Wait, AllowSynchronousContinuations = false, }); _events = Channel.CreateBounded( new BoundedChannelOptions(_options.EventChannelCapacity) { // The worker event channel has exactly ONE consumer: the per-session // SessionEventDistributor pump. The alarm monitor and dashboard mirror both // attach to the distributor rather than draining this channel directly, so a // second concurrent reader would silently split events between the two // enumerators. SingleReader=true asserts that invariant; ReadEventsAsync adds a // claimed-once guard so a regression fails loudly instead of losing events. SingleReader = true, SingleWriter = true, FullMode = BoundedChannelFullMode.Wait, AllowSynchronousContinuations = false, }); _eventStagingCapacity = checked(2 * _options.EventChannelCapacity); _eventStaging = Channel.CreateBounded( new BoundedChannelOptions(_eventStagingCapacity) { // The read loop is the only writer; EventWriteLoopAsync is the only reader. SingleReader = true, SingleWriter = true, // Wait (not Drop*) so the read loop's non-blocking TryWrite returns false exactly // when the bound is reached — the same Wait+TryWrite overflow-detection idiom the // session event distributor uses. The read loop never awaits this channel. FullMode = BoundedChannelFullMode.Wait, AllowSynchronousContinuations = false, }); _eventQueueDepthRegistration = _metrics?.RegisterWorkerEventQueueDepthSource( () => Volatile.Read(ref _eventQueueDepth)); _lastHeartbeatAt = _timeProvider.GetUtcNow(); } /// public string SessionId => _connection.SessionId; /// public int? ProcessId { get { lock (_syncRoot) { return _processId; } } } /// public WorkerClientState State { get { lock (_syncRoot) { return _state; } } } /// public DateTimeOffset LastHeartbeatAt { get { lock (_syncRoot) { return _lastHeartbeatAt; } } } /// public async Task StartAsync(CancellationToken cancellationToken) { ThrowIfDisposed(); TransitionFromCreatedToHandshaking(); _writeLoopTask = Task.Run(WriteLoopAsync); await EnqueueAsync(CreateGatewayHelloEnvelope(), cancellationToken).ConfigureAwait(false); WorkerEnvelope helloEnvelope = await ReadHandshakeEnvelopeAsync( WorkerEnvelope.BodyOneofCase.WorkerHello, cancellationToken).ConfigureAwait(false); ValidateWorkerHello(helloEnvelope.WorkerHello); WorkerEnvelope readyEnvelope = await ReadHandshakeEnvelopeAsync( WorkerEnvelope.BodyOneofCase.WorkerReady, cancellationToken).ConfigureAwait(false); MarkReady(readyEnvelope.WorkerReady); _readLoopTask = Task.Run(ReadLoopAsync); _eventWriteLoopTask = Task.Run(EventWriteLoopAsync); _heartbeatLoopTask = Task.Run(HeartbeatLoopAsync); } /// public async Task InvokeAsync( WorkerCommand command, TimeSpan timeout, CancellationToken cancellationToken) { ArgumentNullException.ThrowIfNull(command); ThrowIfDisposed(); EnsureReady(); if (timeout <= TimeSpan.Zero) { throw new ArgumentOutOfRangeException(nameof(timeout), timeout, "Command timeout must be greater than zero."); } string correlationId = Guid.NewGuid().ToString("N"); string method = GetCommandMethod(command); if (!_pendingCommandSlots.Wait(0)) { _metrics?.QueueOverflow("worker-pending-commands"); throw new WorkerClientException( WorkerClientErrorCode.PendingCommandLimitExceeded, $"Worker session {SessionId} already has {_options.MaxPendingCommands} pending command(s)."); } PendingCommand pendingCommand = new( correlationId, method, _timeProvider.GetTimestamp()); if (!_pendingCommands.TryAdd(correlationId, pendingCommand)) { ReleasePendingCommandSlot(); throw new InvalidOperationException("Generated a duplicate command correlation id."); } _metrics?.CommandStarted(method); try { WorkerEnvelope commandEnvelope = CreateCommandEnvelope(correlationId, command); // Reject an oversized command at the enqueue boundary so only this correlation fails // (ResourceExhausted) rather than the frame reaching the write loop and faulting the whole // session. Command envelopes are the only gateway-authored outbound payload whose // size the caller controls; checking here keeps a MessageTooLarge in the write loop a // genuine desync signal. // // PERF(GWC-31): this size cannot be handed to WorkerFrameWriter to spare its own // CalculateSize. WriteLoopAsync stamps envelope.Sequence immediately before the write // (GWC-28), and a non-zero varint field grows the encoding — so the number computed here // is a lower bound on the frame the writer actually emits, never the frame length. Passing // it as a knownSize would under-length the prefix and desync the worker's framing. The // pre-check stays a pre-check: it is conservative in the right direction. int envelopeSize = commandEnvelope.CalculateSize(); if (envelopeSize > _connection.FrameOptions.MaxMessageBytes) { throw new WorkerClientException( WorkerClientErrorCode.CommandTooLarge, $"Worker command {method} serializes to {envelopeSize} bytes, exceeding the negotiated " + $"worker-frame maximum of {_connection.FrameOptions.MaxMessageBytes} bytes."); } await EnqueueAsync(commandEnvelope, cancellationToken).ConfigureAwait(false); // GWC-31: one pooled timer instead of a linked CTS + Task.Delay + WhenAny per command. // Task.WaitAsync arms a TimerQueueTimer on the shared timer queue and cancels it when the // reply lands, so the steady-state cost of a command that replies in time is a single // continuation — the old shape allocated a linked CancellationTokenSource, its // registration, a delay Task, and the WhenAny Task on every invoke, and left the delay // Task rooted until the cancel completed. WaitAsync raises TimeoutException for the // deadline and OperationCanceledException for the caller's token — but unlike the old // wait it races the two and reports whichever fired first, whereas the old code inspected // cancellationToken.IsCancellationRequested BEFORE classifying a won delay as a timeout. // The filter on the CommandTimeout clause restores that priority: a token canceled around // the deadline is still classified as cancellation, never as CommandTimeout. Error codes // and messages are unchanged. try { return await pendingCommand.Task.WaitAsync(timeout, cancellationToken).ConfigureAwait(false); } catch (TimeoutException) when (!cancellationToken.IsCancellationRequested) { string timeoutMessage = $"Worker command {method} timed out after {timeout}."; bool removed = RemovePendingCommandAsFailed( correlationId, pendingCommand, WorkerClientErrorCode.CommandTimeout, timeoutMessage); // The gateway has stopped waiting, but the worker has not stopped working: the // correlation is still on its single STA queue and would execute (or keep executing) // regardless. Tell it, so a queued-but-not-started command is dropped instead of // occupying the STA behind a caller that is already gone. Gated on the removal so a // reply that won the race — the pending entry is already gone and the caller is about // to see it — never has a cancel chase it. Best-effort by design; the send cannot // throw, so it can never replace the timeout the caller is owed. if (removed) { TrySendCancelForTimedOutCommand(correlationId, method, timeout); } throw new WorkerClientException( WorkerClientErrorCode.CommandTimeout, timeoutMessage); } catch (OperationCanceledException) { RemovePendingCommandAsFailed( correlationId, pendingCommand, WorkerClientErrorCode.GatewayShutdown, "Command wait was canceled."); // WaitAsync surfaces TaskCanceledException; throwing through the token keeps the // exception the caller observes exactly what the hand-rolled wait produced. cancellationToken.ThrowIfCancellationRequested(); throw; } catch (TimeoutException) { // The deadline and the caller's cancellation raced and WaitAsync picked the timer. // The old wait classified this as cancellation, so this clause — reached only when // the filter above saw a canceled token — reproduces that treatment exactly. RemovePendingCommandAsFailed( correlationId, pendingCommand, WorkerClientErrorCode.GatewayShutdown, "Command wait was canceled."); cancellationToken.ThrowIfCancellationRequested(); throw; } } catch { if (_pendingCommands.TryRemove(correlationId, out _)) { ReleasePendingCommandSlot(); } throw; } } /// public IAsyncEnumerable ReadEventsAsync(CancellationToken cancellationToken) { // The event channel is SingleReader: only one enumerator may ever drain it, otherwise // the two readers would each receive a random subset of events. Claim the reader at CALL // time (not lazily on first MoveNext) and fail loudly on a second consumer rather than // silently splitting the stream. The distributor pump is the only intended // caller; the alarm monitor and dashboard mirror attach to the distributor instead. if (Interlocked.CompareExchange(ref _eventsReaderClaimed, 1, 0) != 0) { throw new InvalidOperationException( "WorkerClient.ReadEventsAsync was already claimed by another consumer. The worker event " + "channel is single-reader; attach to the SessionEventDistributor instead of draining it twice."); } return ReadEventsCoreAsync(cancellationToken); } private async IAsyncEnumerable ReadEventsCoreAsync( [EnumeratorCancellation] CancellationToken cancellationToken) { await foreach (WorkerEvent workerEvent in _events.Reader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { // No metrics call on the hot path: the gauge pulls _eventQueueDepth when scraped (GWC-30). Interlocked.Decrement(ref _eventQueueDepth); yield return workerEvent; } } /// public async Task ShutdownAsync(TimeSpan timeout, CancellationToken cancellationToken) { ThrowIfDisposed(); if (timeout <= TimeSpan.Zero) { throw new ArgumentOutOfRangeException(nameof(timeout), timeout, "Shutdown timeout must be greater than zero."); } WorkerClientState state = State; if (state == WorkerClientState.Closed) { return; } if (state == WorkerClientState.Faulted) { KillOwnedProcess("ShutdownFaulted"); return; } MarkClosing(); await EnqueueAsync(CreateShutdownEnvelope(timeout, "gateway-shutdown"), cancellationToken).ConfigureAwait(false); _outboundEnvelopes.Writer.TryComplete(); using CancellationTokenSource timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); timeoutCts.CancelAfter(timeout); try { await WaitForBackgroundTasksAsync(timeoutCts.Token).ConfigureAwait(false); await WaitForProcessExitAsync(timeoutCts.Token).ConfigureAwait(false); MarkClosed("shutdown"); } catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) { SetFaulted( WorkerClientErrorCode.ShutdownTimeout, "Worker shutdown timed out.", null); throw new WorkerClientException( WorkerClientErrorCode.ShutdownTimeout, $"Worker shutdown timed out after {timeout}."); } } /// public void Kill(string reason) { ThrowIfDisposed(); KillOwnedProcess(reason); SetFaulted( WorkerClientErrorCode.WorkerFaulted, $"Worker was killed by the gateway: {reason}.", null); } /// Disposes the worker client and releases resources. /// A task that represents the asynchronous operation. public async ValueTask DisposeAsync() { if (_disposed) { return; } _disposed = true; // Drop out of the worker queue-depth gauge before teardown: whatever this client still holds // is about to be discarded, and the sum must not keep counting a client that is going away. _eventQueueDepthRegistration?.Dispose(); KillOwnedProcess("Dispose"); _stopCts.Cancel(); _outboundEnvelopes.Writer.TryComplete(); _eventStaging.Writer.TryComplete(); _events.Writer.TryComplete(); CompletePendingCommands( new WorkerClientException( WorkerClientErrorCode.GatewayShutdown, "Worker client was disposed.")); await _connection.Stream.DisposeAsync().ConfigureAwait(false); using CancellationTokenSource disposeTimeout = new(DisposeTaskTimeout); try { await WaitForBackgroundTasksAsync(disposeTimeout.Token).ConfigureAwait(false); } catch (OperationCanceledException) { _logger.LogWarning( "Timed out waiting for worker client background tasks to stop for session {SessionId}.", SessionId); } _connection.ProcessHandle?.Dispose(); _pendingCommandSlots.Dispose(); _stopCts.Dispose(); } /// Manages writing envelopes to the worker pipe. private async Task WriteLoopAsync() { try { await foreach (WorkerEnvelope envelope in _outboundEnvelopes.Reader.ReadAllAsync(_stopCts.Token).ConfigureAwait(false)) { // GWC-28: stamp the sequence at the point of writing, not when the envelope is built. // Stamping at construction let two concurrent InvokeAsync callers take 1 and 2 and then // enqueue in the order 2, 1 — non-monotonic on the wire, breaking gateway.md's // "monotonic per sender" contract. This loop is the channel's single consumer // (SingleReader = true), so wire order and stamp order are the same thing here and // _nextSequence needs no interlocking. Mirrors the worker's WRK-04 fix. envelope.Sequence = unchecked(++_nextSequence); await _writer.WriteAsync(envelope, _stopCts.Token).ConfigureAwait(false); } } catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState()) { } catch (Exception exception) { SetFaulted( WorkerClientErrorCode.WriteFailed, "Worker pipe write failed.", exception); } } /// Manages reading envelopes from the worker pipe. private async Task ReadLoopAsync() { try { while (!_stopCts.IsCancellationRequested) { WorkerEnvelope envelope = await _reader.ReadAsync(_stopCts.Token).ConfigureAwait(false); DispatchEnvelope(envelope); } } catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState()) { } catch (WorkerFrameProtocolException exception) when (exception.ErrorCode == WorkerFrameProtocolErrorCode.EndOfStream) { SetFaulted( WorkerClientErrorCode.PipeDisconnected, "Worker pipe disconnected.", exception); } catch (Exception exception) { SetFaulted( WorkerClientErrorCode.ProtocolViolation, "Worker read loop failed.", exception); } } /// /// Monitors worker heartbeat and detects stale sessions. Mirrors the /// worker-side watchdog: while a command is in flight on the /// gateway↔worker pipe, the heartbeat watchdog is suppressed up to /// — the worker /// may be busy executing a slow STA command and the heartbeat write may /// be queued behind a long event-drain burst, neither of which /// indicates the worker is actually hung. Once the oldest pending /// command exceeds the ceiling, the fault fires anyway so a truly stuck /// COM call doesn't hide the worker forever. /// private async Task HeartbeatLoopAsync() { try { while (!_stopCts.IsCancellationRequested) { await Task.Delay(_options.HeartbeatCheckInterval, _stopCts.Token).ConfigureAwait(false); if (State != WorkerClientState.Ready) { continue; } DateTimeOffset lastHeartbeatAt = LastHeartbeatAt; DateTimeOffset now = _timeProvider.GetUtcNow(); if (now - lastHeartbeatAt <= _options.HeartbeatGrace) { continue; } if (TryGetOldestPendingCommandAge(out TimeSpan oldestCommandAge) && oldestCommandAge <= _options.HeartbeatStuckCeiling) { continue; } _metrics?.HeartbeatFailed(SessionId); SetFaulted( WorkerClientErrorCode.HeartbeatExpired, $"Worker heartbeat expired. Last heartbeat was at {lastHeartbeatAt:O}.", null); } } catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState()) { } } /// /// Returns the age of the oldest pending command on the worker pipe, /// measured via against /// , or false when no /// commands are pending. Used by the heartbeat watchdog /// to decide whether a heartbeat gap is plausibly the result of /// pipe-write contention rather than a hung worker. /// private bool TryGetOldestPendingCommandAge(out TimeSpan oldestAge) { long oldestStart = long.MaxValue; foreach (PendingCommand pending in _pendingCommands.Values) { if (pending.StartTimestamp < oldestStart) { oldestStart = pending.StartTimestamp; } } if (oldestStart == long.MaxValue) { oldestAge = TimeSpan.Zero; return false; } oldestAge = _timeProvider.GetElapsedTime(oldestStart); return true; } /// /// Routes a received envelope to its handler. Every branch dispatches synchronously and /// immediately — the event branch only stages the event for the dedicated writer — so a full /// event channel can never delay a command reply, heartbeat, fault, or shutdown ack behind an /// event backlog. /// /// The envelope to dispatch. private void DispatchEnvelope(WorkerEnvelope envelope) { switch (envelope.BodyCase) { case WorkerEnvelope.BodyOneofCase.WorkerCommandReply: CompleteCommand(envelope); break; case WorkerEnvelope.BodyOneofCase.WorkerEvent: StageWorkerEvent(envelope.WorkerEvent); break; case WorkerEnvelope.BodyOneofCase.WorkerHeartbeat: MarkHeartbeat(envelope.WorkerHeartbeat); break; case WorkerEnvelope.BodyOneofCase.WorkerFault: SetFaulted( WorkerClientErrorCode.WorkerFaulted, CreateWorkerFaultMessage(envelope.WorkerFault), null); break; case WorkerEnvelope.BodyOneofCase.WorkerShutdownAck: MarkClosed("worker-shutdown-ack"); break; default: SetFaulted( WorkerClientErrorCode.ProtocolViolation, $"Worker sent unexpected envelope body {envelope.BodyCase}.", null); break; } } /// /// Hands a received worker event to the dedicated event writer without blocking the read loop. /// TryWrite is non-blocking, so the read loop keeps dispatching replies, heartbeats and /// faults regardless of how backed up the event path is. It returns false in two cases: /// the staging channel has been completed during shutdown (the event is safely dropped because /// the client is already terminal), or the channel is full at its /// 2 × bound. The latter means the /// consumer has been draining slower than the worker produces for the whole time it took the /// worker to emit that many further events, so the client is faulted immediately (GWC-24) — /// is non-blocking, so the read loop still never awaits here. The /// complementary full-stall case (a consumer that stops entirely) is caught earlier by the /// timed write in . /// /// The event received from the worker. private void StageWorkerEvent(WorkerEvent workerEvent) { if (workerEvent.Event is not null) { _metrics?.EventReceived(SessionId, workerEvent.Event.Family.ToString()); } if (_eventStaging.Writer.TryWrite(workerEvent)) { // Counted here rather than at the _events write so the single gauge reports total // undelivered events (staged + queued). ReadEventsCoreAsync decrements on consumer read. Interlocked.Increment(ref _eventQueueDepth); return; } if (IsTerminalState()) { // Shutdown completed the staging channel; dropping the event is the documented behavior. return; } _metrics?.QueueOverflow("worker-event-staging"); int depthAtOverflow = Volatile.Read(ref _eventQueueDepth); SetFaulted( WorkerClientErrorCode.ProtocolViolation, $"Worker event staging channel is full at its {_eventStagingCapacity}-event bound " + $"(2 x EventChannelCapacity {_options.EventChannelCapacity}); undelivered depth is " + $"{depthAtOverflow}. The event consumer is draining slower than the worker produces. " + $"Attach or unblock the StreamEvents consumer or raise MxGateway:Events:QueueCapacity.", null); } /// /// Drains staged worker events and applies the bounded-channel backpressure (and /// sustained-overflow fault) on a dedicated task, so the timed write /// never runs on the read loop. Mirrors for events. /// private async Task EventWriteLoopAsync() { try { await foreach (WorkerEvent workerEvent in _eventStaging.Reader.ReadAllAsync(_stopCts.Token).ConfigureAwait(false)) { await EnqueueWorkerEventAsync(workerEvent, _stopCts.Token).ConfigureAwait(false); } } catch (OperationCanceledException) when (_stopCts.IsCancellationRequested || IsTerminalState()) { } } /// /// Enqueues a worker event for client consumption. The channel is /// configured with /// and a brief consumer hiccup is tolerated for up to /// /// (default 5s) before the worker is faulted. The channel previously /// used TryWrite (non-blocking), which never honored the /// configured FullModeTimeout — the worker faulted on the first /// missed slot even though the wait-mode channel would have absorbed /// the burst. The diagnostic now names the capacity, current depth, and /// the actionable fix (attach StreamEvents or raise /// MxGateway:Events:QueueCapacity). This is the full-stall half of /// the backpressure policy; a consumer that merely drains too slowly is /// caught by the staging bound in (GWC-24). /// Queue depth is not counted here — the event was already counted when it /// was staged, so moving it between the two channels changes nothing. /// /// The event to enqueue. /// Cancellation token. private async Task EnqueueWorkerEventAsync( WorkerEvent workerEvent, CancellationToken cancellationToken) { if (_events.Writer.TryWrite(workerEvent)) { return; } using CancellationTokenSource fullModeCts = CancellationTokenSource .CreateLinkedTokenSource(cancellationToken); fullModeCts.CancelAfter(_options.EventChannelFullModeTimeout); try { await _events.Writer.WriteAsync(workerEvent, fullModeCts.Token).ConfigureAwait(false); return; } catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested) { // Only the full-mode timeout fired — the outer cancellation is // a different concern and is rethrown by the await above when it // triggers. } _metrics?.QueueOverflow("worker-events"); int depthAtOverflow = Volatile.Read(ref _eventQueueDepth); SetFaulted( WorkerClientErrorCode.ProtocolViolation, $"Worker event channel rejected an event after waiting " + $"{_options.EventChannelFullModeTimeout.TotalMilliseconds:F0} ms; " + $"undelivered depth is {depthAtOverflow} against a consumer channel " + $"of {_options.EventChannelCapacity} capacity. " + $"Attach a StreamEvents consumer or raise MxGateway:Events:QueueCapacity.", null); } /// Completes pending command with worker reply. /// Envelope containing command reply. private void CompleteCommand(WorkerEnvelope envelope) { string correlationId = envelope.CorrelationId; if (string.IsNullOrWhiteSpace(correlationId)) { correlationId = envelope.WorkerCommandReply.Reply?.CorrelationId ?? string.Empty; } if (!_pendingCommands.TryRemove(correlationId, out PendingCommand? pendingCommand)) { _logger.LogDebug( "Ignoring late or unknown worker command reply for session {SessionId} and correlation {CorrelationId}.", SessionId, correlationId); return; } ReleasePendingCommandSlot(); TimeSpan duration = _timeProvider.GetElapsedTime(pendingCommand.StartTimestamp); _metrics?.CommandSucceeded(pendingCommand.Method, duration); pendingCommand.SetResult(envelope.WorkerCommandReply); } /// Fails a pending command with an error. /// Command correlation ID. /// The pending command. /// Error code. /// Error message. /// /// true when this call removed the pending entry and owns the failure; false when /// the entry was already gone — a reply, fault, or shutdown got there first, so the caller must /// not take any further action on behalf of that correlation. /// private bool RemovePendingCommandAsFailed( string correlationId, PendingCommand pendingCommand, WorkerClientErrorCode errorCode, string message) { if (!_pendingCommands.TryRemove(correlationId, out _)) { return false; } ReleasePendingCommandSlot(); TimeSpan duration = _timeProvider.GetElapsedTime(pendingCommand.StartTimestamp); _metrics?.CommandFailed(pendingCommand.Method, errorCode.ToString(), duration); pendingCommand.SetException(new WorkerClientException(errorCode, message)); return true; } /// /// Forwards a WorkerCancel for a correlation the gateway has given up waiting for, so the /// worker can drop it from its STA queue (WorkerPipeSession routes the envelope to /// CancelCommand). A cancel that arrives after the command already reached the COM call /// is a no-op — MXAccess offers no way to abort an in-flight call — so this shortens the STA /// backlog rather than freeing a call already running on it. /// /// A command whose envelope has not yet left _outboundEnvelopes is handled by the same /// path rather than by pulling it back out: exposes no removal, and /// the queue is FIFO, so the worker simply reads the command and then its cancel and drops the /// correlation before it ever reaches the STA. Nothing is gained by dequeuing it here. /// /// /// /// The whole body sits under one catch-all that debug-logs, because the guarantee this method /// owes its caller is structural, not incidental: the caller is on the throw path for the /// timeout, so anything escaping here — an envelope that failed to build, a TryWrite /// against a disposed channel, a scheduler refusing the detached task — would replace the /// the caller is owed with an unrelated /// exception. Losing the cancel costs the worker one wasted command; losing the timeout /// misreports why the call failed. The detached task carries its own handler for the same /// reason: its failures (including the from /// _stopCts if the client is disposed underneath it) happen after this method returns /// and would otherwise be unobserved. It is deliberately not tracked or awaited — it holds no /// resource the shutdown path needs back, and the outbound channel is completed on close. /// /// Correlation id of the command that timed out. /// Command method name, for the cancel reason and diagnostics. /// The elapsed command timeout, for the cancel reason. private void TrySendCancelForTimedOutCommand( string correlationId, string method, TimeSpan timeout) { try { WorkerEnvelope cancelEnvelope = CreateEnvelope( correlationId, envelope => envelope.WorkerCancel = new WorkerCancel { Reason = $"gateway command timeout after {timeout}", }); if (_outboundEnvelopes.Writer.TryWrite(cancelEnvelope)) { return; } _ = Task.Run(async () => { try { await EnqueueAsync(cancelEnvelope, _stopCts.Token).ConfigureAwait(false); } catch (Exception exception) { LogCancelNotForwarded(exception, method, correlationId); } }); } catch (Exception exception) { LogCancelNotForwarded(exception, method, correlationId); } } /// Records a cancel that could not be forwarded for a timed-out command. /// The failure that stopped the cancel from being sent. /// Command method name of the timed-out command. /// Correlation id of the timed-out command. private void LogCancelNotForwarded( Exception exception, string method, string correlationId) { _logger.LogDebug( exception, "Could not forward a cancel for timed-out worker command {Method} on session {SessionId} " + "and correlation {CorrelationId}.", method, SessionId, correlationId); } /// Reads and validates a handshake envelope. /// Expected envelope body type. /// Cancellation token. /// The read envelope. private async Task ReadHandshakeEnvelopeAsync( WorkerEnvelope.BodyOneofCase expectedBody, CancellationToken cancellationToken) { WorkerEnvelope envelope = await _reader.ReadAsync(cancellationToken).ConfigureAwait(false); if (envelope.BodyCase != expectedBody) { throw new WorkerClientException( WorkerClientErrorCode.ProtocolViolation, $"Worker handshake expected {expectedBody} but received {envelope.BodyCase}."); } return envelope; } /// Validates worker hello message protocol and nonce. /// The hello message to validate. private void ValidateWorkerHello(WorkerHello workerHello) { if (workerHello.ProtocolVersion != _connection.FrameOptions.ProtocolVersion) { throw new WorkerClientException( WorkerClientErrorCode.ProtocolViolation, "Worker hello protocol version does not match the gateway protocol version."); } if (!string.Equals(workerHello.Nonce, _connection.Nonce, StringComparison.Ordinal)) { throw new WorkerClientException( WorkerClientErrorCode.ProtocolViolation, "Worker hello nonce does not match the gateway nonce."); } lock (_syncRoot) { _processId = workerHello.WorkerProcessId == 0 ? _connection.ProcessHandle?.ProcessId : workerHello.WorkerProcessId; } } /// Marks the worker as ready and records startup metrics. /// The ready message. private void MarkReady(WorkerReady ready) { lock (_syncRoot) { _processId = ready.WorkerProcessId == 0 ? _processId ?? _connection.ProcessHandle?.ProcessId : ready.WorkerProcessId; _lastHeartbeatAt = _timeProvider.GetUtcNow(); _state = WorkerClientState.Ready; _workerStartRecorded = true; } DateTimeOffset readyAt = _timeProvider.GetUtcNow(); DateTimeOffset launchedAt = _connection.ProcessHandle?.LaunchedAt ?? readyAt; _metrics?.WorkerStarted(readyAt - launchedAt); } /// Updates the heartbeat timestamp and process ID. /// The heartbeat message. private void MarkHeartbeat(WorkerHeartbeat heartbeat) { lock (_syncRoot) { _lastHeartbeatAt = _timeProvider.GetUtcNow(); if (heartbeat.WorkerProcessId != 0) { _processId = heartbeat.WorkerProcessId; } } } /// Transitions to closing state. private void MarkClosing() { lock (_syncRoot) { if (_state is WorkerClientState.Closed or WorkerClientState.Faulted) { return; } _state = WorkerClientState.Closing; } } /// Marks client as closed and cleans up resources. /// Reason for closure. private void MarkClosed(string reason) { lock (_syncRoot) { if (_state == WorkerClientState.Closed) { return; } _state = WorkerClientState.Closed; } _stopCts.Cancel(); _outboundEnvelopes.Writer.TryComplete(); _eventStaging.Writer.TryComplete(); _events.Writer.TryComplete(); CompletePendingCommands( new WorkerClientException( WorkerClientErrorCode.GatewayShutdown, $"Worker client closed because {reason}.")); RecordWorkerStoppedOnce(reason); } /// Marks client as faulted and propagates the error. /// Error code. /// Error message. /// Optional inner exception. private void SetFaulted( WorkerClientErrorCode errorCode, string message, Exception? exception) { WorkerClientException fault = exception is null ? new WorkerClientException(errorCode, message) : new WorkerClientException(errorCode, message, exception); lock (_syncRoot) { if (_state is WorkerClientState.Faulted or WorkerClientState.Closed) { return; } _state = WorkerClientState.Faulted; } _stopCts.Cancel(); _outboundEnvelopes.Writer.TryComplete(fault); _eventStaging.Writer.TryComplete(fault); _events.Writer.TryComplete(fault); KillOwnedProcess(errorCode.ToString()); CompletePendingCommands(fault); RecordWorkerStoppedOnce(errorCode.ToString()); _metrics?.Fault(errorCode.ToString()); _logger.LogWarning(exception, "Worker client faulted for session {SessionId}: {Message}", SessionId, message); } private void KillOwnedProcess(string reason) { WorkerProcessHandle? processHandle = _connection.ProcessHandle; if (processHandle is null) { return; } try { if (!processHandle.Process.HasExited) { processHandle.Process.Kill(entireProcessTree: true); _metrics?.WorkerKilled(reason); } } catch (Exception exception) { _logger.LogWarning( exception, "Failed to kill worker process {ProcessId} for session {SessionId}.", processHandle.ProcessId, SessionId); } } /// Records worker stopped metric only once. /// Reason for stopping. private void RecordWorkerStoppedOnce(string reason) { bool shouldRecord; lock (_syncRoot) { shouldRecord = _workerStartRecorded; _workerStartRecorded = false; } if (shouldRecord) { _metrics?.WorkerStopped(reason); } } /// Fails all pending commands with the given exception. /// Exception to apply to all commands. private void CompletePendingCommands(Exception exception) { foreach (KeyValuePair item in _pendingCommands.ToArray()) { if (_pendingCommands.TryRemove(item.Key, out PendingCommand? pendingCommand)) { ReleasePendingCommandSlot(); TimeSpan duration = _timeProvider.GetElapsedTime(pendingCommand.StartTimestamp); _metrics?.CommandFailed(pendingCommand.Method, exception.GetType().Name, duration); pendingCommand.SetException(exception); } } } /// Releases a pending command slot. private void ReleasePendingCommandSlot() { try { _pendingCommandSlots.Release(); } catch (SemaphoreFullException) { } } /// Transitions from created to handshaking state. private void TransitionFromCreatedToHandshaking() { lock (_syncRoot) { if (_state != WorkerClientState.Created) { throw new WorkerClientException( WorkerClientErrorCode.InvalidState, $"Worker client cannot start from state {_state}."); } _state = WorkerClientState.Handshaking; } } /// Throws if client is not in ready state. private void EnsureReady() { WorkerClientState state = State; if (state != WorkerClientState.Ready) { throw new WorkerClientException( WorkerClientErrorCode.InvalidState, $"Worker client is not ready. Current state is {state}."); } } /// Checks if current state is terminal. /// True if closed, closing, or faulted. private bool IsTerminalState() { WorkerClientState state = State; return state is WorkerClientState.Closing or WorkerClientState.Closed or WorkerClientState.Faulted; } /// Enqueues an envelope for writing to the worker. /// Envelope to enqueue. /// Cancellation token. private async Task EnqueueAsync( WorkerEnvelope envelope, CancellationToken cancellationToken) { try { await _outboundEnvelopes.Writer.WriteAsync(envelope, cancellationToken).ConfigureAwait(false); } catch (ChannelClosedException exception) { throw new WorkerClientException( WorkerClientErrorCode.WriteFailed, "Worker outbound channel is closed.", exception); } } /// Creates gateway hello envelope. /// The hello envelope. private WorkerEnvelope CreateGatewayHelloEnvelope() { return CreateEnvelope( correlationId: string.Empty, envelope => envelope.GatewayHello = new GatewayHello { SupportedProtocolVersion = _connection.FrameOptions.ProtocolVersion, Nonce = _connection.Nonce, GatewayVersion = typeof(GatewayContractInfo).Assembly.GetName().Version?.ToString() ?? GatewayVersionFallback, // Convey the negotiated worker-frame maximum so the worker adopts it instead of a // hard-coded default. Sits above the public gRPC cap by the envelope reserve. MaxFrameBytes = (uint)_connection.FrameOptions.MaxMessageBytes, }); } /// Creates command envelope. /// Command correlation ID. /// The command to wrap. /// The command envelope. private WorkerEnvelope CreateCommandEnvelope( string correlationId, WorkerCommand command) { // IPC-05: no second clone. MxAccessGrpcMapper.MapCommand already deep-cloned the command // out of the caller-owned gRPC MxCommandRequest, so this WorkerCommand is a fresh graph // owned by the invoke pipeline and referenced by no other consumer. The envelope is built // and owned entirely inside WorkerClient and the command is not touched again after this // point, so transferring it into the envelope introduces no aliasing hazard. return CreateEnvelope( correlationId, envelope => envelope.WorkerCommand = command); } /// Creates shutdown envelope. /// Shutdown timeout. /// Shutdown reason. /// The shutdown envelope. private WorkerEnvelope CreateShutdownEnvelope( TimeSpan timeout, string reason) { return CreateEnvelope( correlationId: string.Empty, envelope => envelope.WorkerShutdown = new WorkerShutdown { GracePeriod = Duration.FromTimeSpan(timeout), Reason = reason, }); } /// Creates a new worker envelope with common fields. /// Correlation ID for the envelope. /// Action to set envelope body. /// The created envelope. private WorkerEnvelope CreateEnvelope( string correlationId, Action setBody) { // Sequence is deliberately left unset here: WriteLoopAsync stamps it immediately before the // frame goes out, so the numbers are monotonic in wire order however the callers interleave // between construction and enqueue (GWC-28, mirroring the worker's WRK-04 fix). WorkerEnvelope envelope = new() { ProtocolVersion = _connection.FrameOptions.ProtocolVersion, SessionId = SessionId, CorrelationId = correlationId, }; setBody(envelope); return envelope; } /// Gets the human-readable command method name. /// The command to get method name from. /// Command method name. private static string GetCommandMethod(WorkerCommand command) { return command.Command?.Kind.ToString() ?? MxCommandKind.Unspecified.ToString(); } /// Creates a fault message from worker fault data. /// The worker fault. /// Formatted fault message. private static string CreateWorkerFaultMessage(WorkerFault fault) { return string.IsNullOrWhiteSpace(fault.DiagnosticMessage) ? $"Worker faulted with category {fault.Category}." : $"Worker faulted with category {fault.Category}: {fault.DiagnosticMessage}"; } /// Waits for all background tasks to complete. /// Cancellation token. private async Task WaitForBackgroundTasksAsync(CancellationToken cancellationToken) { Task[] tasks = new[] { _readLoopTask, _writeLoopTask, _eventWriteLoopTask, _heartbeatLoopTask } .Where(task => task is not null) .Cast() .ToArray(); if (tasks.Length == 0) { return; } await Task.WhenAll(tasks).WaitAsync(cancellationToken).ConfigureAwait(false); } /// Waits for the worker process to exit. /// Cancellation token. private async Task WaitForProcessExitAsync(CancellationToken cancellationToken) { WorkerProcessHandle? processHandle = _connection.ProcessHandle; if (processHandle is null || processHandle.Process.HasExited) { return; } await processHandle.Process.WaitForExitAsync(cancellationToken).ConfigureAwait(false); } /// Throws if the client has been disposed. private void ThrowIfDisposed() { ObjectDisposedException.ThrowIf(_disposed, this); } private sealed class PendingCommand { private readonly TaskCompletionSource _completion = new(TaskCreationOptions.RunContinuationsAsynchronously); /// Initializes a pending command awaiting a worker reply. /// Command correlation ID for reply matching. /// Command method name. /// Start time in milliseconds for duration tracking. public PendingCommand( string correlationId, string method, long startTimestamp) { CorrelationId = correlationId; Method = method; StartTimestamp = startTimestamp; } /// Gets the command correlation ID. public string CorrelationId { get; } /// Gets the command method name. public string Method { get; } /// Gets the command start timestamp. public long StartTimestamp { get; } /// Gets the task that completes when reply arrives. public Task Task => _completion.Task; /// Completes the command with a reply. /// The command reply. public void SetResult(WorkerCommandReply reply) { _completion.TrySetResult(reply); } /// Completes the command with an exception. /// The exception. public void SetException(Exception exception) { _completion.TrySetException(exception); } } }