mxaccessgw/src/MxGateway.Server/Sessions/GatewaySession.cs

using MxGateway.Contracts.Proto;
using MxGateway.Server.Workers;

namespace MxGateway.Server.Sessions;

public sealed class GatewaySession
{
    private readonly object _syncRoot = new();
    private readonly SemaphoreSlim _closeLock = new(1, 1);
    private IWorkerClient? _workerClient;
    private SessionState _state = SessionState.Creating;
    private string? _finalFault;
    private DateTimeOffset _lastClientActivityAt;
    private DateTimeOffset? _leaseExpiresAt;
    private bool _closeStarted;
    private int _activeEventSubscriberCount;

    public GatewaySession(
        string sessionId,
        string backendName,
        string pipeName,
        string nonce,
        string? clientIdentity,
        string? clientSessionName,
        string? clientCorrelationId,
        TimeSpan commandTimeout,
        TimeSpan startupTimeout,
        TimeSpan shutdownTimeout,
        DateTimeOffset openedAt)
    {
        if (string.IsNullOrWhiteSpace(sessionId))
        {
            throw new ArgumentException("Session id is required.", nameof(sessionId));
        }

        if (string.IsNullOrWhiteSpace(backendName))
        {
            throw new ArgumentException("Backend name is required.", nameof(backendName));
        }

        if (string.IsNullOrWhiteSpace(pipeName))
        {
            throw new ArgumentException("Pipe name is required.", nameof(pipeName));
        }

        if (string.IsNullOrWhiteSpace(nonce))
        {
            throw new ArgumentException("Nonce is required.", nameof(nonce));
        }

        SessionId = sessionId;
        BackendName = backendName;
        PipeName = pipeName;
        Nonce = nonce;
        ClientIdentity = clientIdentity;
        ClientSessionName = clientSessionName;
        ClientCorrelationId = clientCorrelationId;
        CommandTimeout = commandTimeout;
        StartupTimeout = startupTimeout;
        ShutdownTimeout = shutdownTimeout;
        OpenedAt = openedAt;
        _lastClientActivityAt = openedAt;
    }

    public string SessionId { get; }

    public string BackendName { get; }

    public string PipeName { get; }

    public string Nonce { get; }

    public string? ClientIdentity { get; }

    public string? ClientSessionName { get; }

    public string? ClientCorrelationId { get; }

    public TimeSpan CommandTimeout { get; }

    public TimeSpan StartupTimeout { get; }

    public TimeSpan ShutdownTimeout { get; }

    public DateTimeOffset OpenedAt { get; }

    public int? WorkerProcessId => _workerClient?.ProcessId;

    public IWorkerClient? WorkerClient => _workerClient;

    public SessionState State
    {
        get
        {
            lock (_syncRoot)
            {
                return _state;
            }
        }
    }

    public DateTimeOffset LastClientActivityAt
    {
        get
        {
            lock (_syncRoot)
            {
                return _lastClientActivityAt;
            }
        }
    }

    public DateTimeOffset? LeaseExpiresAt
    {
        get
        {
            lock (_syncRoot)
            {
                return _leaseExpiresAt;
            }
        }
    }

    public string? FinalFault
    {
        get
        {
            lock (_syncRoot)
            {
                return _finalFault;
            }
        }
    }

    public int ActiveEventSubscriberCount
    {
        get
        {
            lock (_syncRoot)
            {
                return _activeEventSubscriberCount;
            }
        }
    }

    public void AttachWorkerClient(IWorkerClient workerClient)
    {
        ArgumentNullException.ThrowIfNull(workerClient);

        lock (_syncRoot)
        {
            _workerClient = workerClient;
        }
    }

    public void TransitionTo(SessionState nextState)
    {
        lock (_syncRoot)
        {
            if (_state is SessionState.Closed)
            {
                return;
            }

            if (_state is SessionState.Faulted && nextState is not SessionState.Closed)
            {
                return;
            }

            _state = nextState;
        }
    }

    public void MarkReady()
    {
        TransitionTo(SessionState.Ready);
    }

    public void MarkFaulted(string reason)
    {
        lock (_syncRoot)
        {
            if (_state is SessionState.Closed)
            {
                return;
            }

            _finalFault = reason;
            _state = SessionState.Faulted;
        }
    }

    public void TouchClientActivity(DateTimeOffset activityAt)
    {
        lock (_syncRoot)
        {
            _lastClientActivityAt = activityAt;
        }
    }

    public void ExtendLease(DateTimeOffset leaseExpiresAt)
    {
        lock (_syncRoot)
        {
            _leaseExpiresAt = leaseExpiresAt;
        }
    }

    public bool IsLeaseExpired(DateTimeOffset now)
    {
        lock (_syncRoot)
        {
            return _leaseExpiresAt is not null && _leaseExpiresAt <= now;
        }
    }

    public IDisposable AttachEventSubscriber(bool allowMultipleSubscribers)
    {
        lock (_syncRoot)
        {
            if (_state != SessionState.Ready || _workerClient?.State != WorkerClientState.Ready)
            {
                throw new SessionManagerException(
                    SessionManagerErrorCode.SessionNotReady,
                    $"Session {SessionId} is not ready for event streaming. Current state is {_state}.");
            }

            if (!allowMultipleSubscribers && _activeEventSubscriberCount > 0)
            {
                throw new SessionManagerException(
                    SessionManagerErrorCode.EventSubscriberAlreadyActive,
                    $"Session {SessionId} already has an active event stream subscriber.");
            }

            _activeEventSubscriberCount++;
            return new EventSubscriberLease(this);
        }
    }

    public async Task<WorkerCommandReply> InvokeAsync(
        WorkerCommand command,
        CancellationToken cancellationToken)
    {
        IWorkerClient workerClient = GetReadyWorkerClient();
        TouchClientActivity(DateTimeOffset.UtcNow);

        return await workerClient.InvokeAsync(command, CommandTimeout, cancellationToken).ConfigureAwait(false);
    }

    public IAsyncEnumerable<WorkerEvent> ReadEventsAsync(CancellationToken cancellationToken)
    {
        IWorkerClient workerClient = GetReadyWorkerClient();
        TouchClientActivity(DateTimeOffset.UtcNow);

        return workerClient.ReadEventsAsync(cancellationToken);
    }

    public async Task<SessionCloseResult> CloseAsync(
        string reason,
        CancellationToken cancellationToken)
    {
        await _closeLock.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            if (_state is SessionState.Closed)
            {
                return new SessionCloseResult(SessionId, SessionState.Closed, AlreadyClosed: true);
            }

            bool alreadyClosing = _closeStarted;
            _closeStarted = true;
            _state = SessionState.Closing;

            if (_workerClient is not null)
            {
                try
                {
                    await _workerClient.ShutdownAsync(ShutdownTimeout, cancellationToken).ConfigureAwait(false);
                }
                catch
                {
                    _workerClient.Kill(reason);
                    throw;
                }
            }

            _state = SessionState.Closed;
            return new SessionCloseResult(SessionId, SessionState.Closed, alreadyClosing);
        }
        finally
        {
            _closeLock.Release();
        }
    }

    public void KillWorker(string reason)
    {
        _workerClient?.Kill(reason);
        TransitionTo(SessionState.Closed);
    }

    public async ValueTask DisposeAsync()
    {
        _closeLock.Dispose();
        if (_workerClient is not null)
        {
            await _workerClient.DisposeAsync().ConfigureAwait(false);
        }
    }

    private IWorkerClient GetReadyWorkerClient()
    {
        lock (_syncRoot)
        {
            if (_state != SessionState.Ready || _workerClient?.State != WorkerClientState.Ready)
            {
                throw new SessionManagerException(
                    SessionManagerErrorCode.SessionNotReady,
                    $"Session {SessionId} is not ready. Current state is {_state}.");
            }

            return _workerClient;
        }
    }

    private void DetachEventSubscriber()
    {
        lock (_syncRoot)
        {
            if (_activeEventSubscriberCount > 0)
            {
                _activeEventSubscriberCount--;
            }
        }
    }

    private sealed class EventSubscriberLease(GatewaySession session) : IDisposable
    {
        private bool _disposed;

        public void Dispose()
        {
            if (_disposed)
            {
                return;
            }

            session.DetachEventSubscriber();
            _disposed = true;
        }
    }
}