ResilientMessageTransport.cs

이 파일의 문서화 페이지로 가기

using Dreamine.Communication.Abstractions.Enums;
using Dreamine.Communication.Abstractions.Interfaces;
using Dreamine.Communication.Abstractions.Models;
using Dreamine.Communication.Abstractions.Options;
using Dreamine.Communication.Core.Queues;
 
namespace Dreamine.Communication.Core.Resilience;

public sealed class ResilientMessageTransport : IMessageTransport
{
    private readonly IMessageTransport _innerTransport;
    private readonly ReconnectPolicy _reconnectPolicy;
    private readonly OutboundQueueOptions _queueOptions;
    private readonly IOutboundMessageQueue _outboundQueue;
    private readonly ReconnectDelayCalculator _delayCalculator;
    private readonly SemaphoreSlim _flushLock = new(1, 1);
    private readonly object _loopSync = new();
    private readonly object _notifySync = new();

    private CancellationTokenSource _lifetimeCts = new();
    private Task? _resilienceLoopTask;
    private int _manualDisconnectRequested = 1;
    private int _resilienceLoopRequested;
    private int _disposeGuard;   // Interlocked: 0 = live, 1 = disposed
    private bool _disposed;
    private ConnectionState _lastNotifiedState = ConnectionState.Disconnected;

    public ResilientMessageTransport(
        IMessageTransport innerTransport,
        ReconnectPolicy? reconnectPolicy = null,
        OutboundQueueOptions? queueOptions = null,
        IOutboundMessageQueue? outboundQueue = null)
    {
        _innerTransport = innerTransport ?? throw new ArgumentNullException(nameof(innerTransport));
        _reconnectPolicy = reconnectPolicy ?? new ReconnectPolicy();
        _queueOptions = queueOptions ?? new OutboundQueueOptions();
        _outboundQueue = outboundQueue ?? new InMemoryOutboundMessageQueue(_queueOptions);
        _delayCalculator = new ReconnectDelayCalculator(_reconnectPolicy);
 
        _lastNotifiedState = State;
        _innerTransport.MessageReceived += OnInnerMessageReceived;
    }

    public event EventHandler<MessageEnvelope>? MessageReceived;

    public event EventHandler<ConnectionState>? StateChanged;

    public event EventHandler<int>? QueueCountChanged;

    public TransportKind Kind => _innerTransport.Kind;

    public ConnectionState State
    {
        get
        {
            if (IsOperationalState(_innerTransport.State))
            {
                return _innerTransport.State;
            }
 
            if (_reconnectPolicy.Enabled &&
                !IsManualDisconnectRequested &&
                IsResilienceLoopRequested)
            {
                return ConnectionState.Connecting;
            }
 
            return _innerTransport.State;
        }
    }

    public int QueuedMessageCount => _outboundQueue.Count;

    public async Task ConnectAsync(CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
 
        SetManualDisconnectRequested(false);
        EnsureResilienceLoopStarted();
        NotifyStateChangedIfNeeded();
 
        try
        {
            await _innerTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
            NotifyStateChangedIfNeeded();
 
            if (_queueOptions.FlushOnReconnect)
            {
                await FlushQueueAsync(cancellationToken).ConfigureAwait(false);
            }
        }
        catch
        {
            NotifyStateChangedIfNeeded();
 
            if (!_reconnectPolicy.Enabled)
            {
                throw;
            }
 
            // 자동 재접속 모드에서는 최초 연결 실패를 치명 오류로 보지 않습니다.
            // 이후 백그라운드 재접속 루프가 계속 연결을 시도합니다.
        }
    }

    public async Task DisconnectAsync(CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
 
        SetManualDisconnectRequested(true);
        SetResilienceLoopRequested(false);
        NotifyStateChangedIfNeeded();
 
        await _innerTransport.DisconnectAsync(cancellationToken).ConfigureAwait(false);
        NotifyStateChangedIfNeeded();
    }

    public async Task SendAsync(
        MessageEnvelope message,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(message);
        ThrowIfDisposed();
 
        SetManualDisconnectRequested(false);
 
        if (IsOperationalState(_innerTransport.State))
        {
            try
            {
                await _innerTransport.SendAsync(message, cancellationToken).ConfigureAwait(false);
                NotifyStateChangedIfNeeded();
                return;
            }
            catch
            {
                await HandleDisconnectedSendAsync(message, cancellationToken).ConfigureAwait(false);
                EnsureResilienceLoopStarted();
                NotifyStateChangedIfNeeded();
                return;
            }
        }
 
        await HandleDisconnectedSendAsync(message, cancellationToken).ConfigureAwait(false);
        EnsureResilienceLoopStarted();
        NotifyStateChangedIfNeeded();
    }

    public async Task FlushQueueAsync(CancellationToken cancellationToken = default)
    {
        ThrowIfDisposed();
 
        if (!_queueOptions.FlushOnReconnect)
        {
            return;
        }
 
        await _flushLock.WaitAsync(cancellationToken).ConfigureAwait(false);
 
        try
        {
            while (IsOperationalState(_innerTransport.State))
            {
                cancellationToken.ThrowIfCancellationRequested();
 
                var queuedMessage = await _outboundQueue
                    .TryDequeueAsync(cancellationToken)
                    .ConfigureAwait(false);
 
                if (queuedMessage is null)
                {
                    NotifyQueueCountChanged();
                    NotifyStateChangedIfNeeded();
                    return;
                }
 
                NotifyQueueCountChanged();
 
                if (queuedMessage.IsExpired(
                        _queueOptions.MaxMessageAge,
                        DateTimeOffset.UtcNow))
                {
                    continue;
                }
 
                try
                {
                    await _innerTransport
                        .SendAsync(queuedMessage.Message, cancellationToken)
                        .ConfigureAwait(false);
 
                    NotifyQueueCountChanged();
                    NotifyStateChangedIfNeeded();
                }
                catch (Exception ex)
                {
                    await _outboundQueue
                        .EnqueueFrontAsync(queuedMessage.WithFailure(ex), cancellationToken)
                        .ConfigureAwait(false);
 
                    NotifyQueueCountChanged();
                    EnsureResilienceLoopStarted();
                    NotifyStateChangedIfNeeded();
                    return;
                }
            }
        }
        finally
        {
            _flushLock.Release();
        }
    }

    public async ValueTask DisposeAsync()
    {
        // Interlocked guard prevents concurrent double-dispose.
        if (Interlocked.CompareExchange(ref _disposeGuard, 1, 0) != 0)
        {
            return;
        }
 
        _disposed = true;
        SetManualDisconnectRequested(true);
        SetResilienceLoopRequested(false);
 
        _innerTransport.MessageReceived -= OnInnerMessageReceived;
 
        try
        {
            await _lifetimeCts.CancelAsync().ConfigureAwait(false);
        }
        catch
        {
            // Dispose 경로에서는 취소 실패를 무시합니다.
        }
 
        if (_resilienceLoopTask is not null)
        {
            try
            {
                await _resilienceLoopTask.ConfigureAwait(false);
            }
            catch
            {
                // Dispose 경로에서는 백그라운드 루프 예외를 무시합니다.
            }
        }
 
        _lifetimeCts.Dispose();
        _flushLock.Dispose();
 
        await _innerTransport.DisposeAsync().ConfigureAwait(false);
    }

    private async Task HandleDisconnectedSendAsync(
        MessageEnvelope message,
        CancellationToken cancellationToken)
    {
        switch (_queueOptions.DisconnectedSendPolicy)
        {
            case DisconnectedSendPolicy.Fail:
                throw new InvalidOperationException("Transport is not connected.");
 
            case DisconnectedSendPolicy.Queue:
                await _outboundQueue
                    .EnqueueAsync(new QueuedOutboundMessage(message), cancellationToken)
                    .ConfigureAwait(false);
 
                NotifyQueueCountChanged();
                return;
 
            case DisconnectedSendPolicy.WaitUntilConnected:
                await WaitUntilConnectedAsync(cancellationToken).ConfigureAwait(false);
                await _innerTransport.SendAsync(message, cancellationToken).ConfigureAwait(false);
                NotifyStateChangedIfNeeded();
                return;
 
            default:
                throw new InvalidOperationException(
                    $"Unsupported disconnected send policy: {_queueOptions.DisconnectedSendPolicy}");
        }
    }

    private async Task WaitUntilConnectedAsync(CancellationToken cancellationToken)
    {
        EnsureResilienceLoopStarted();
        NotifyStateChangedIfNeeded();
 
        while (!IsOperationalState(_innerTransport.State))
        {
            cancellationToken.ThrowIfCancellationRequested();
            await Task.Delay(_reconnectPolicy.WatchInterval, cancellationToken)
                .ConfigureAwait(false);
        }
    }

    private void EnsureResilienceLoopStarted()
    {
        if (!_reconnectPolicy.Enabled)
        {
            return;
        }
 
        lock (_loopSync)
        {
            if (_disposed)
            {
                return;
            }
 
            SetResilienceLoopRequested(true);
 
            if (_resilienceLoopTask is { IsCompleted: false })
            {
                return;
            }
 
            if (_lifetimeCts.IsCancellationRequested)
            {
                _lifetimeCts.Dispose();
                _lifetimeCts = new CancellationTokenSource();
            }
 
            _resilienceLoopTask = Task.Run(
                () => RunResilienceLoopAsync(_lifetimeCts.Token),
                CancellationToken.None);
        }
    }

    private async Task RunResilienceLoopAsync(CancellationToken cancellationToken)
    {
        var reconnectAttempt = 0;
 
        while (!cancellationToken.IsCancellationRequested)
        {
            NotifyStateChangedIfNeeded();
 
            if (IsManualDisconnectRequested)
            {
                await DelaySafelyAsync(
                    _reconnectPolicy.WatchInterval,
                    cancellationToken)
                    .ConfigureAwait(false);
 
                continue;
            }
 
            if (IsOperationalState(_innerTransport.State))
            {
                reconnectAttempt = 0;
                NotifyStateChangedIfNeeded();
 
                if (_queueOptions.FlushOnReconnect && _outboundQueue.Count > 0)
                {
                    await FlushQueueAsync(cancellationToken).ConfigureAwait(false);
                }
 
                await DelaySafelyAsync(
                    _reconnectPolicy.WatchInterval,
                    cancellationToken)
                    .ConfigureAwait(false);
 
                continue;
            }
 
            if (_innerTransport.State == ConnectionState.Disconnecting)
            {
                await DelaySafelyAsync(
                    _reconnectPolicy.WatchInterval,
                    cancellationToken)
                    .ConfigureAwait(false);
 
                continue;
            }
 
            if (_reconnectPolicy.MaxRetryCount is not null &&
                reconnectAttempt >= _reconnectPolicy.MaxRetryCount.Value)
            {
                await DelaySafelyAsync(
                    _reconnectPolicy.MaxDelay,
                    cancellationToken)
                    .ConfigureAwait(false);
 
                continue;
            }
 
            reconnectAttempt++;
            NotifyStateChangedIfNeeded();
 
            try
            {
                await _innerTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
                reconnectAttempt = 0;
                NotifyStateChangedIfNeeded();
 
                if (_queueOptions.FlushOnReconnect)
                {
                    await FlushQueueAsync(cancellationToken).ConfigureAwait(false);
                }
            }
            catch
            {
                NotifyStateChangedIfNeeded();
 
                var delay = _delayCalculator.Calculate(reconnectAttempt);
 
                await DelaySafelyAsync(delay, cancellationToken)
                    .ConfigureAwait(false);
            }
        }
    }

    private static bool IsOperationalState(ConnectionState state)
    {
        return state is ConnectionState.Connected or ConnectionState.Listening;
    }

    private bool IsManualDisconnectRequested =>
        Volatile.Read(ref _manualDisconnectRequested) != 0;

    private bool IsResilienceLoopRequested =>
        Volatile.Read(ref _resilienceLoopRequested) != 0;

    private void SetManualDisconnectRequested(bool value)
    {
        Interlocked.Exchange(ref _manualDisconnectRequested, value ? 1 : 0);
    }

    private void SetResilienceLoopRequested(bool value)
    {
        Interlocked.Exchange(ref _resilienceLoopRequested, value ? 1 : 0);
    }

    private static async Task DelaySafelyAsync(
        TimeSpan delay,
        CancellationToken cancellationToken)
    {
        if (delay <= TimeSpan.Zero)
        {
            delay = TimeSpan.FromMilliseconds(1);
        }
 
        try
        {
            await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
        }
        catch (OperationCanceledException)
        {
            // 정상적인 종료 경로입니다.
        }
    }

    private void OnInnerMessageReceived(object? sender, MessageEnvelope message)
    {
        MessageReceived?.Invoke(this, message);
    }

    private void NotifyStateChangedIfNeeded()
    {
        var currentState = State;
        EventHandler<ConnectionState>? handler;
 
        lock (_notifySync)
        {
            if (_lastNotifiedState == currentState)
            {
                return;
            }
 
            _lastNotifiedState = currentState;
            handler = StateChanged;
        }
 
        handler?.Invoke(this, currentState);
    }

    private void NotifyQueueCountChanged()
    {
        QueueCountChanged?.Invoke(this, _outboundQueue.Count);
    }

    private void ThrowIfDisposed()
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
    }
}