using System.Diagnostics.CodeAnalysis; using System.Runtime.ExceptionServices; using System.Security.Cryptography; using Google.Protobuf.WellKnownTypes; using Microsoft.Extensions.Hosting; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Logging.Abstractions; using Microsoft.Extensions.Options; using ZB.MOM.WW.MxGateway.Contracts; using ZB.MOM.WW.MxGateway.Contracts.Proto; using ZB.MOM.WW.MxGateway.Server.Configuration; using ZB.MOM.WW.MxGateway.Server.Metrics; using ZB.MOM.WW.MxGateway.Server.Workers; namespace ZB.MOM.WW.MxGateway.Server.Sessions; public sealed class SessionManager : ISessionManager { public const string DefaultCloseReason = "client-close"; public const string GatewayShutdownReason = "gateway-shutdown"; public const string LeaseExpiredReason = "lease-expired"; public const string DetachGraceExpiredReason = "detach-grace-expired"; public const string FaultedReason = "faulted-reaped"; // Bounded so a mass expiry (or a host stop with a full registry) cannot stampede // worker-process teardown: every concurrent close is one x86 worker being shut down or // killed, and the point of the fan-out is to hide a few hung workers, not to tear the whole // registry down at once. private const int MaxParallelSessionCloses = 4; // ShutdownAsync re-snapshots the registry until it is empty so an open that registered its // session just after a snapshot is still drained. The stopping gate in OpenSessionAsync is // what actually closes that window (it is re-checked after registration), so the loop is a // backstop and a small bound is enough; anything still registered after the last pass is // logged and left to the orphan terminator on the next start. private const int MaxShutdownDrainPasses = 3; private readonly ISessionRegistry _registry; private readonly ISessionWorkerClientFactory _workerClientFactory; private readonly GatewayMetrics _metrics; private readonly TimeProvider _timeProvider; private readonly ILogger _logger; private readonly GatewayOptions _options; private readonly SemaphoreSlim _sessionSlots; private readonly Grpc.MxAccessGrpcMapper _eventMapper; private readonly ILogger _distributorLogger; private readonly Dashboard.Hubs.IDashboardEventBroadcaster? _dashboardEventBroadcaster; private readonly ArrayAddressNormalizer? _addressNormalizer; // 0 = accepting opens, 1 = stopping. Set by the host's ApplicationStopping signal (which fires // before hosted services stop, while Kestrel still accepts gRPC calls) and, as a fallback for // hosts without a lifetime, at the start of ShutdownAsync. Never reset: a manager that has // begun draining does not come back. private int _stopping; /// /// Initializes a new instance of . /// /// Session registry. /// Worker client factory. /// Gateway options. /// Gateway metrics. /// Time provider for timestamps. /// Logger. /// Mapper used by each session's event distributor to map worker events to public events. /// Logger passed to each session's event distributor pump. /// /// Dashboard SignalR fan-out sink. Each session registers an internal distributor /// subscriber that mirrors raw session events to this broadcaster, so the /// dashboard receives events regardless of whether a gRPC client is streaming. Null in /// unit tests that do not exercise the dashboard mirror. /// /// /// Rewrites bare array AddItem addresses to their writable [] form using Galaxy /// metadata; handed to each session so the normalization runs at the outbound choke point. /// Null in unit tests that do not exercise array-write ergonomics. /// /// /// Host lifetime whose ApplicationStopping closes the open gate as soon as the host /// begins stopping — before hosted services stop and while Kestrel is still accepting gRPC /// calls — so no session is opened behind the shutdown drain. Null in unit tests; the gate /// then closes when starts. /// public SessionManager( ISessionRegistry registry, ISessionWorkerClientFactory workerClientFactory, IOptions options, GatewayMetrics metrics, TimeProvider? timeProvider = null, ILogger? logger = null, Grpc.MxAccessGrpcMapper? eventMapper = null, ILogger? distributorLogger = null, Dashboard.Hubs.IDashboardEventBroadcaster? dashboardEventBroadcaster = null, ArrayAddressNormalizer? addressNormalizer = null, IHostApplicationLifetime? applicationLifetime = null) { _registry = registry ?? throw new ArgumentNullException(nameof(registry)); _workerClientFactory = workerClientFactory ?? throw new ArgumentNullException(nameof(workerClientFactory)); ArgumentNullException.ThrowIfNull(options); _metrics = metrics ?? throw new ArgumentNullException(nameof(metrics)); _timeProvider = timeProvider ?? TimeProvider.System; _logger = logger ?? NullLogger.Instance; _eventMapper = eventMapper ?? new Grpc.MxAccessGrpcMapper(); _distributorLogger = distributorLogger ?? NullLogger.Instance; _dashboardEventBroadcaster = dashboardEventBroadcaster; _addressNormalizer = addressNormalizer; _options = options.Value; _sessionSlots = new SemaphoreSlim(_options.Sessions.MaxSessions, _options.Sessions.MaxSessions); // The registration is never disposed: the manager is a host singleton that lives exactly // as long as the lifetime it observes. applicationLifetime?.ApplicationStopping.Register(MarkStopping); } /// public Task OpenSessionAsync( SessionOpenRequest request, string? clientIdentity, string? ownerKeyId, CancellationToken cancellationToken) => OpenSessionAsync(request, clientIdentity, ownerKeyId, ownerDashboardTags: null, cancellationToken); /// public async Task OpenSessionAsync( SessionOpenRequest request, string? clientIdentity, string? ownerKeyId, IReadOnlyList? ownerDashboardTags, CancellationToken cancellationToken) { ArgumentNullException.ThrowIfNull(request); ThrowIfStopping(); EnsureSessionCapacity(); GatewaySession? session = null; bool registered = false; try { session = CreateSession(request, clientIdentity, ownerKeyId, ownerDashboardTags); if (!_registry.TryAdd(session)) { throw new SessionManagerException( SessionManagerErrorCode.OpenFailed, $"Session id collision while opening session {session.SessionId}."); } registered = true; // Re-check after registration. The stopping flag is written before ShutdownAsync takes // its registry snapshot, so an open that does not observe it here registered its // session before that snapshot and will be drained; one that does observe it rolls // back itself. Either way no session is created behind the drain. ThrowIfStopping(); session.TransitionTo(SessionState.StartingWorker); IWorkerClient workerClient = await _workerClientFactory .CreateAsync(session, cancellationToken) .ConfigureAwait(false); // The session has been registered (and killable / drainable) for the whole worker // startup. If a kill or the host drain closed it meanwhile, the attach is refused and // the late worker is discarded here — it must not be returned on a dead session, and // nothing else would ever stop it. bool ready = await session .TryAttachWorkerClientAndMarkReadyAsync(workerClient, _metrics.SessionOpened, cancellationToken) .ConfigureAwait(false); if (!ready) { await DiscardWorkerClientAsync(session, workerClient).ConfigureAwait(false); throw new InvalidOperationException( $"Session {session.SessionId} was closed while its worker was starting."); } return session; } catch (SessionManagerException exception) when (exception.ErrorCode == SessionManagerErrorCode.GatewayStopping) { // Not an open failure: the gateway is going away and the caller should retry against // the restarted one, so no fault is recorded and the code is surfaced as-is. await RollBackOpenAsync(session, registered, exception).ConfigureAwait(false); throw; } catch (Exception exception) { await RollBackOpenAsync(session, registered, exception).ConfigureAwait(false); _metrics.Fault(SessionManagerErrorCode.OpenFailed.ToString()); _logger.LogWarning( exception, "Failed to open gateway session {SessionId}.", session?.SessionId ?? ""); throw new SessionManagerException( SessionManagerErrorCode.OpenFailed, session is null ? "Failed to create session." : $"Failed to open session {session.SessionId}.", exception); } } // Order — fault, deregister (dropping per-session event metrics), release the slot, dispose — // matters: releasing the slot before disposal would let the next open race this worker's // teardown on the same machine. A registered session goes through RemoveSessionAsync, the // single removal path, which is a no-op when a concurrent kill or drain already removed it // (that teardown released the slot and disposed the session). The open-session gauge is never // touched here: SessionOpened is recorded inside the attach gate, after which nothing in the // open sequence can fail. private async Task RollBackOpenAsync(GatewaySession? session, bool registered, Exception cause) { if (session is null) { ReleaseSessionSlot(); return; } session.MarkFaulted(cause.Message); if (registered) { await RemoveSessionAsync(session).ConfigureAwait(false); return; } ReleaseSessionSlot(); await session.DisposeAsync().ConfigureAwait(false); } // Kills and disposes a worker client that arrived after its session was closed. Best-effort: // the client is already orphaned from the session, so a failure here is logged rather than // allowed to mask the open failure the caller is about to receive. private async Task DiscardWorkerClientAsync(GatewaySession session, IWorkerClient workerClient) { try { workerClient.Kill("session-closed-during-startup"); } catch (Exception exception) { _logger.LogWarning( exception, "Failed to kill the late worker for closed session {SessionId}.", session.SessionId); } try { await workerClient.DisposeAsync().ConfigureAwait(false); } catch (Exception exception) { _logger.LogWarning( exception, "Failed to dispose the late worker client for closed session {SessionId}.", session.SessionId); } } /// public bool TryGetSession( string sessionId, [MaybeNullWhen(false)] out GatewaySession session) { return _registry.TryGet(sessionId, out session); } /// public async Task InvokeAsync( string sessionId, WorkerCommand command, CancellationToken cancellationToken) { GatewaySession session = GetRequiredSession(sessionId); try { return await session.InvokeAsync(command, cancellationToken).ConfigureAwait(false); } catch (SessionManagerException) { throw; } catch (Exception exception) { if (session.WorkerClient?.State == WorkerClientState.Faulted) { session.MarkFaulted(exception.Message); } throw; } } /// public async Task CloseSessionAsync( string sessionId, CancellationToken cancellationToken) { GatewaySession session = GetRequiredSession(sessionId); SessionCloseResult result = await CloseSessionCoreAsync( session, DefaultCloseReason, cancellationToken).ConfigureAwait(false); return result; } /// public async Task KillWorkerAsync( string sessionId, string reason, CancellationToken cancellationToken) { ArgumentException.ThrowIfNullOrWhiteSpace(reason); cancellationToken.ThrowIfCancellationRequested(); GatewaySession session = GetRequiredSession(sessionId); bool wasClosed; try { wasClosed = await session.KillWorkerWithCloseGateAsync(reason, cancellationToken).ConfigureAwait(false); } catch (Exception exception) { session.MarkFaulted(exception.Message); _metrics.Fault(SessionManagerErrorCode.CloseFailed.ToString()); // If the open-session gauge was incremented for this session (it was not when the // kill landed while the worker was still starting), decrement it here so // mxgateway.sessions.open does not leak when the kill path throws. if (session.OpenRecorded) { _metrics.SessionRemoved(); } await RemoveSessionAsync(session).ConfigureAwait(false); throw new SessionManagerException( SessionManagerErrorCode.CloseFailed, $"Failed to kill worker for session {sessionId}.", exception); } if (!wasClosed && session.OpenRecorded) { _metrics.SessionClosed(); } await RemoveSessionAsync(session).ConfigureAwait(false); _logger.LogInformation( "Worker for session {SessionId} killed; reason={Reason}.", sessionId, reason); return new SessionCloseResult(sessionId, SessionState.Closed, AlreadyClosed: wasClosed); } /// public async Task CloseExpiredLeasesAsync( DateTimeOffset now, CancellationToken cancellationToken) { // Selection phase — deliberately sequential. Only the close calls below run in parallel: // deciding WHICH sessions to close must stay a single ordered pass so the sweep-precedence // rule and the TOCTOU re-check keep their meaning. List<(GatewaySession Session, string Reason)> selected = []; foreach (GatewaySession session in _registry.Snapshot()) { // A session is swept when its normal lease has expired, it has FAULTED (a faulted // session is permanently unusable yet otherwise pins a session slot and a live x86 // worker until its DefaultLeaseSeconds lease expires), OR its detach-grace retention // window has elapsed (last external subscriber dropped and no client reconnected // within DetachGraceSeconds). All three are the same teardown; only the reason differs // so operators can tell a fault reap from a short reconnect-window expiry from a long // idle-lease expiry in logs/metrics. Precedence when several fire simultaneously: // lease-expired, then faulted, then detach-grace. // // TOCTOU note: eligibility is re-verified atomically inside TryBeginCloseIfExpired // under _syncRoot, so a client that reattaches a subscriber between the check above // and the close call wins the race and the session is left open and usable. string? reason = session.IsLeaseExpired(now) ? LeaseExpiredReason : session.IsFaultedReapable(now) ? FaultedReason : session.IsDetachGraceExpired(now) ? DetachGraceExpiredReason : null; if (reason is null) { continue; } // Re-verify eligibility atomically and begin the Closing transition before // delegating to CloseSessionCoreAsync. If a subscriber reattached between the // IsLeaseExpired/IsDetachGraceExpired check above and here, TryBeginCloseIfExpired // returns false and we skip this session (it is no longer expired). if (!session.TryBeginCloseIfExpired(now, out bool alreadyClosing) && !alreadyClosing) { continue; } selected.Add((session, reason)); } if (selected.Count == 0) { return 0; } int closedCount = 0; object failureSyncRoot = new(); ExceptionDispatchInfo? firstFailure = null; // Close phase. Each close is bounded by the worker shutdown timeout (default 10 s), so a // mass expiry with a few hung workers would serialize reaping and starve session slots. // Parallel close is safe because TryBeginCloseIfExpired above already flipped every // selected session to Closing under its own lock — that idempotent begin-close is the // per-session exclusivity invariant, so no two teardowns can run against one session and // a session selected here cannot be re-selected by a concurrent sweep. await Parallel.ForEachAsync( selected, new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses, CancellationToken = cancellationToken, }, async (candidate, closeToken) => { try { await CloseSessionCoreAsync(candidate.Session, candidate.Reason, closeToken).ConfigureAwait(false); Interlocked.Increment(ref closedCount); } catch (Exception exception) { // The sequential sweep let a close failure propagate to the lease monitor, // which logs it; that signal is preserved by rethrowing the first failure // below. It is captured rather than thrown here so one failed (or hung) // teardown does not abandon the rest of the already-selected set. lock (failureSyncRoot) { firstFailure ??= ExceptionDispatchInfo.Capture(exception); } } }).ConfigureAwait(false); firstFailure?.Throw(); return closedCount; } /// public async Task ShutdownAsync(CancellationToken cancellationToken) { // Sessions are drained in parallel: at a worst-case worker shutdown timeout each, a // one-at-a-time drain of a full registry outruns any host stop-timeout and leaves the // tail to the orphan killer. Per-session exclusivity comes from GatewaySession.CloseAsync's // own close gate, and each iteration touches only its own session plus thread-safe // registry/metrics state. // // The body must be exception-TOTAL. Parallel.ForEachAsync cancels the token it hands the // sibling bodies as soon as one body throws, so a single escaping exception would abort up // to MaxParallelSessionCloses - 1 in-flight graceful shutdowns AND make their kill fallback // throw immediately on the freshly cancelled token — sessions neither closed nor killed, // i.e. leaked x86 workers that nothing reattaches to (a gateway restart terminates orphans // rather than adopting them). // // For the same reason the loop itself is NOT bound to cancellationToken: a cancelled // ParallelOptions token stops dispatching the remaining sessions entirely, so a stop // deadline would leave the untried tail neither closed nor killed. Note this FIXES a leak // the sequential drain also had rather than restoring its behavior: there the kill fallback // ran on the caller's cancelled token, and KillWorkerAsync's entry // ThrowIfCancellationRequested threw out of the loop on the first session — zero kills, not // "fail fast and still kill". The token is passed to the graceful close only, and the kill // runs on CancellationToken.None, so a host stop deadline turns the drain into a kill sweep // rather than into a leak. // // The asymmetry with CloseExpiredLeasesAsync (whose ParallelOptions IS token-bound) is // intentional: that sweep is periodic maintenance whose missed sessions are picked up by // the next pass and, ultimately, by this drain. This drain is terminal — nothing runs after // it — so it must not be abandoned partway. // // Close the open gate BEFORE the first snapshot (see the re-check in OpenSessionAsync), then // re-snapshot until the registry is empty so an open that registered just before the gate // closed is drained too. MarkStopping(); for (int pass = 0; pass < MaxShutdownDrainPasses; pass++) { if (_registry.Count == 0) { return; } await Parallel.ForEachAsync( _registry.Snapshot(), new ParallelOptions { MaxDegreeOfParallelism = MaxParallelSessionCloses }, async (session, _) => { try { await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false); } catch (Exception exception) { _logger.LogWarning( exception, "Graceful shutdown failed for session {SessionId}; killing worker.", session.SessionId); // Defensive fallback: CloseSessionCoreAsync's inner // SessionCloseStartedException catch normally removes the session // and accounts the close. This outer fallback only fires for // sessions still in the registry — route through KillWorkerAsync // so the bookkeeping is identical to the dashboard kill path. if (_registry.TryGet(session.SessionId, out GatewaySession? registeredSession) && registeredSession is not null) { try { // Deliberately NOT the caller's token: the kill is the last-resort orphan // preventer, so it must still run when the host stop deadline (or a // sibling body's failure) has already cancelled the drain. It is a // synchronous Kill plus registry/dispose bookkeeping, not a wait on // the worker, so it cannot extend the drain meaningfully. await KillWorkerAsync(session.SessionId, GatewayShutdownReason, CancellationToken.None).ConfigureAwait(false); } catch (Exception killException) { _logger.LogWarning( killException, "Worker kill fallback failed for session {SessionId}.", session.SessionId); } } } }).ConfigureAwait(false); } int remaining = _registry.Count; if (remaining > 0) { _logger.LogWarning( "Session drain finished with {Remaining} session(s) still registered after {Passes} passes; their workers are left to the orphan terminator on the next start.", remaining, MaxShutdownDrainPasses); } } private void MarkStopping() { Volatile.Write(ref _stopping, 1); } private void ThrowIfStopping() { if (Volatile.Read(ref _stopping) != 0) { throw new SessionManagerException( SessionManagerErrorCode.GatewayStopping, "The gateway is stopping and is not accepting new sessions."); } } private async Task CloseSessionCoreAsync( GatewaySession session, string reason, CancellationToken cancellationToken) { bool wasClosed = session.State == SessionState.Closed; try { SessionCloseResult result = await session.CloseAsync(reason, cancellationToken).ConfigureAwait(false); if (result.WorkerShutdownFailure is not null) { // A clean close from the session's point of view (the worker was killed and the // session is Closed), but the operator should see that the worker would not stop // gracefully — a run of these is the signal that workers are hanging on shutdown. _logger.LogWarning( result.WorkerShutdownFailure, "Worker for session {SessionId} did not shut down gracefully and was killed; reason={Reason}.", session.SessionId, reason); } if (!wasClosed && !result.AlreadyClosed && session.OpenRecorded) { _metrics.SessionClosed(); } await RemoveSessionAsync(session).ConfigureAwait(false); return result; } catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { throw; } catch (SessionCloseStartedException exception) { session.MarkFaulted(exception.Message); if (!wasClosed && session.OpenRecorded) { // Account the close as a SessionClosed (decrements the open-session // gauge AND increments the sessions.closed counter), not just // SessionRemoved. The session is being removed from the registry // below; treating this as a half-finished close that only // decremented the gauge would under-count the closed counter. _metrics.SessionClosed(); } _metrics.Fault(SessionManagerErrorCode.CloseFailed.ToString()); await RemoveSessionAsync(session).ConfigureAwait(false); throw new SessionManagerException( SessionManagerErrorCode.CloseFailed, $"Failed to close session {session.SessionId}.", exception); } } private GatewaySession GetRequiredSession(string sessionId) { if (!_registry.TryGet(sessionId, out GatewaySession? session) || session is null) { throw new SessionManagerException( SessionManagerErrorCode.SessionNotFound, $"Session {sessionId} was not found."); } return session; } private void EnsureSessionCapacity() { if (!_sessionSlots.Wait(0)) { throw new SessionManagerException( SessionManagerErrorCode.SessionLimitExceeded, $"Gateway session limit {_options.Sessions.MaxSessions} has been reached."); } } private async Task RemoveSessionAsync(GatewaySession session) { if (!_registry.TryRemove(session.SessionId, out GatewaySession? removedSession)) { return; } _metrics.RemoveSessionEvents(session.SessionId); ReleaseSessionSlot(); await removedSession.DisposeAsync().ConfigureAwait(false); } private void ReleaseSessionSlot() { try { _sessionSlots.Release(); } catch (SemaphoreFullException) { } } private GatewaySession CreateSession( SessionOpenRequest request, string? clientIdentity, string? ownerKeyId, IReadOnlyList? ownerDashboardTags) { string sessionUid = Guid.NewGuid().ToString("N"); string sessionId = $"session-{sessionUid}"; string backendName = string.IsNullOrWhiteSpace(request.RequestedBackend) ? GatewayContractInfo.DefaultBackendName : request.RequestedBackend!; TimeSpan commandTimeout = ResolveCommandTimeout(request.CommandTimeout); TimeSpan startupTimeout = TimeSpan.FromSeconds(_options.Worker.StartupTimeoutSeconds); TimeSpan shutdownTimeout = TimeSpan.FromSeconds(_options.Worker.ShutdownTimeoutSeconds); TimeSpan leaseDuration = TimeSpan.FromSeconds(_options.Sessions.DefaultLeaseSeconds); // The short prefix and bare guid keep the pipe's Unix-domain-socket path // (TMPDIR + "CoreFxPipe_" + name) inside the 104-byte sun_path limit on // macOS, whose default per-user TMPDIR is ~49 chars; the gateway PID keeps // the name collision-free across gateway restarts. string pipeName = $"mxgw-{Environment.ProcessId}-{sessionUid}"; string nonce = CreateNonce(); DateTimeOffset openedAt = _timeProvider.GetUtcNow(); string clientCorrelationId = CreateClientCorrelationId(request.ClientSessionName, sessionId); SessionEventStreaming eventStreaming = new( _eventMapper, _options.Events, _distributorLogger, _timeProvider, _metrics, _dashboardEventBroadcaster, _options.Sessions.AllowMultipleEventSubscribers); return new GatewaySession( sessionId, backendName, pipeName, nonce, clientIdentity, ownerKeyId, request.ClientSessionName, clientCorrelationId, commandTimeout, startupTimeout, shutdownTimeout, leaseDuration, openedAt, eventStreaming, TimeSpan.FromSeconds(Math.Max(0, _options.Sessions.DetachGraceSeconds)), TimeSpan.FromMilliseconds(Math.Max(0, _options.Sessions.WorkerReadyWaitTimeoutMs)), _addressNormalizer, TimeSpan.FromSeconds(Math.Max(0, _options.Sessions.FaultedGraceSeconds)), ownerDashboardTags); } private static string CreateClientCorrelationId( string? clientSessionName, string sessionId) { string clientName = string.IsNullOrWhiteSpace(clientSessionName) ? "client" : clientSessionName!; return $"{clientName}-{sessionId}"; } private TimeSpan ResolveCommandTimeout(Duration? requestedTimeout) { if (requestedTimeout is null) { return TimeSpan.FromSeconds(_options.Sessions.DefaultCommandTimeoutSeconds); } TimeSpan timeout = requestedTimeout.ToTimeSpan(); return timeout <= TimeSpan.Zero ? TimeSpan.FromSeconds(_options.Sessions.DefaultCommandTimeoutSeconds) : timeout; } private static string CreateNonce() { Span bytes = stackalloc byte[32]; RandomNumberGenerator.Fill(bytes); return Convert.ToBase64String(bytes); } }