mxaccessgw/src/ZB.MOM.WW.MxGateway.Server/Sessions/SessionManager.cs

using System.Diagnostics.CodeAnalysis;
using System.Security.Cryptography;
using Google.Protobuf.WellKnownTypes;
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";

    private readonly ISessionRegistry _registry;
    private readonly ISessionWorkerClientFactory _workerClientFactory;
    private readonly GatewayMetrics _metrics;
    private readonly TimeProvider _timeProvider;
    private readonly ILogger<SessionManager> _logger;
    private readonly GatewayOptions _options;
    private readonly SemaphoreSlim _sessionSlots;
    private readonly Grpc.MxAccessGrpcMapper _eventMapper;
    private readonly ILogger<SessionEventDistributor> _distributorLogger;
    private readonly Dashboard.Hubs.IDashboardEventBroadcaster? _dashboardEventBroadcaster;
    private readonly ArrayAddressNormalizer? _addressNormalizer;

    /// <summary>
    /// Initializes a new instance of <see cref="SessionManager"/>.
    /// </summary>
    /// <param name="registry">Session registry.</param>
    /// <param name="workerClientFactory">Worker client factory.</param>
    /// <param name="options">Gateway options.</param>
    /// <param name="metrics">Gateway metrics.</param>
    /// <param name="timeProvider">Time provider for timestamps.</param>
    /// <param name="logger">Logger.</param>
    /// <param name="eventMapper">Mapper used by each session's event distributor to map worker events to public events.</param>
    /// <param name="distributorLogger">Logger passed to each session's event distributor pump.</param>
    /// <param name="dashboardEventBroadcaster">
    ///     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.
    /// </param>
    /// <param name="addressNormalizer">
    ///     Rewrites bare array AddItem addresses to their writable <c>[]</c> 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.
    /// </param>
    public SessionManager(
        ISessionRegistry registry,
        ISessionWorkerClientFactory workerClientFactory,
        IOptions<GatewayOptions> options,
        GatewayMetrics metrics,
        TimeProvider? timeProvider = null,
        ILogger<SessionManager>? logger = null,
        Grpc.MxAccessGrpcMapper? eventMapper = null,
        ILogger<SessionEventDistributor>? distributorLogger = null,
        Dashboard.Hubs.IDashboardEventBroadcaster? dashboardEventBroadcaster = null,
        ArrayAddressNormalizer? addressNormalizer = 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<SessionManager>.Instance;
        _eventMapper = eventMapper ?? new Grpc.MxAccessGrpcMapper();
        _distributorLogger = distributorLogger ?? NullLogger<SessionEventDistributor>.Instance;
        _dashboardEventBroadcaster = dashboardEventBroadcaster;
        _addressNormalizer = addressNormalizer;
        _options = options.Value;
        _sessionSlots = new SemaphoreSlim(_options.Sessions.MaxSessions, _options.Sessions.MaxSessions);
    }

    /// <inheritdoc />
    public async Task<GatewaySession> OpenSessionAsync(
        SessionOpenRequest request,
        string? clientIdentity,
        string? ownerKeyId,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);
        EnsureSessionCapacity();

        GatewaySession? session = null;
        bool sessionOpenedRecorded = false;
        try
        {
            session = CreateSession(request, clientIdentity, ownerKeyId);
            if (!_registry.TryAdd(session))
            {
                throw new SessionManagerException(
                    SessionManagerErrorCode.OpenFailed,
                    $"Session id collision while opening session {session.SessionId}.");
            }

            session.TransitionTo(SessionState.StartingWorker);
            IWorkerClient workerClient = await _workerClientFactory
                .CreateAsync(session, cancellationToken)
                .ConfigureAwait(false);

            session.AttachWorkerClient(workerClient);
            session.MarkReady();
            _metrics.SessionOpened();
            sessionOpenedRecorded = true;

            return session;
        }
        catch (Exception exception)
        {
            session?.MarkFaulted(exception.Message);
            if (session is not null)
            {
                _registry.TryRemove(session.SessionId, out _);
                await session.DisposeAsync().ConfigureAwait(false);
            }

            // If SessionOpened() already incremented the open-session gauge,
            // a failure after that point (e.g. auto-subscribe rejection) must
            // decrement it again so mxgateway.sessions.open does not leak.
            if (sessionOpenedRecorded)
            {
                _metrics.SessionRemoved();
            }

            ReleaseSessionSlot();
            _metrics.Fault(SessionManagerErrorCode.OpenFailed.ToString());
            _logger.LogWarning(
                exception,
                "Failed to open gateway session {SessionId}.",
                session?.SessionId ?? "<not-created>");

            throw new SessionManagerException(
                SessionManagerErrorCode.OpenFailed,
                session is null ? "Failed to create session." : $"Failed to open session {session.SessionId}.",
                exception);
        }
    }

    /// <inheritdoc />
    public bool TryGetSession(
        string sessionId,
        [MaybeNullWhen(false)] out GatewaySession session)
    {
        return _registry.TryGet(sessionId, out session);
    }

    /// <inheritdoc />
    public async Task<WorkerCommandReply> 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;
        }
    }

    /// <inheritdoc />
    public IAsyncEnumerable<WorkerEvent> ReadEventsAsync(
        string sessionId,
        CancellationToken cancellationToken)
    {
        GatewaySession session = GetRequiredSession(sessionId);

        return session.ReadEventsAsync(cancellationToken);
    }

    /// <inheritdoc />
    public async Task<SessionCloseResult> CloseSessionAsync(
        string sessionId,
        CancellationToken cancellationToken)
    {
        GatewaySession session = GetRequiredSession(sessionId);
        SessionCloseResult result = await CloseSessionCoreAsync(
            session,
            DefaultCloseReason,
            cancellationToken).ConfigureAwait(false);

        return result;
    }

    /// <inheritdoc />
    public async Task<SessionCloseResult> 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());

            _metrics.SessionRemoved();
            await RemoveSessionAsync(session).ConfigureAwait(false);
            throw new SessionManagerException(
                SessionManagerErrorCode.CloseFailed,
                $"Failed to kill worker for session {sessionId}.",
                exception);
        }

        if (!wasClosed)
        {
            _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);
    }

    /// <inheritdoc />
    public async Task<int> CloseExpiredLeasesAsync(
        DateTimeOffset now,
        CancellationToken cancellationToken)
    {
        int closedCount = 0;
        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;
            }

            await CloseSessionCoreAsync(session, reason, cancellationToken).ConfigureAwait(false);
            closedCount++;
        }

        return closedCount;
    }

    /// <inheritdoc />
    public async Task ShutdownAsync(CancellationToken cancellationToken)
    {
        foreach (GatewaySession session in _registry.Snapshot())
        {
            try
            {
                await CloseSessionCoreAsync(session, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
            }
            catch (Exception exception)
            {
                _logger.LogWarning(
                    exception,
                    "Graceful shutdown failed for session {SessionId}; killing worker.",
                    session.SessionId);

                if (_registry.TryGet(session.SessionId, out _))
                {
                    try
                    {
                        await KillWorkerAsync(session.SessionId, GatewayShutdownReason, cancellationToken).ConfigureAwait(false);
                    }
                    catch (SessionManagerException killException)
                    {
                        _logger.LogWarning(
                            killException,
                            "Worker kill fallback failed for session {SessionId}.",
                            session.SessionId);
                    }
                }
            }
        }
    }

    private async Task<SessionCloseResult> CloseSessionCoreAsync(
        GatewaySession session,
        string reason,
        CancellationToken cancellationToken)
    {
        bool wasClosed = session.State == SessionState.Closed;
        try
        {
            SessionCloseResult result = await session.CloseAsync(reason, cancellationToken).ConfigureAwait(false);
            if (!wasClosed && !result.AlreadyClosed)
            {
                _metrics.SessionClosed();
            }

            await RemoveSessionAsync(session).ConfigureAwait(false);
            return result;
        }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
        {
            throw;
        }
        catch (SessionCloseStartedException exception)
        {
            session.MarkFaulted(exception.Message);
            if (!wasClosed)
            {
                _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)
    {
        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 (NEXT-01).
        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)));
    }

    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<byte> bytes = stackalloc byte[32];
        RandomNumberGenerator.Fill(bytes);

        return Convert.ToBase64String(bytes);
    }

}