UdpTransport.cs

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

using System.Net;
using System.Net.Sockets;
using Dreamine.Communication.Abstractions.Enums;
using Dreamine.Communication.Abstractions.Interfaces;
using Dreamine.Communication.Abstractions.Models;
using Dreamine.Communication.Core.Protocols;
using Dreamine.Communication.Sockets.Options;
 
namespace Dreamine.Communication.Sockets.Udp;

public sealed class UdpTransport : IMessageTransport
{
    private readonly UdpTransportOptions _options;
    private readonly IMessageProtocolAdapter _protocolAdapter;

    private UdpClient? _client;
    private IPEndPoint? _remoteEndPoint;
    private CancellationTokenSource? _receiveLoopCts;
    private Task? _receiveLoopTask;
    private int _state = (int)ConnectionState.Disconnected;

    public UdpTransport(UdpTransportOptions options)
        : this(options, new DreamineEnvelopeProtocolAdapter())
    {
    }

    public UdpTransport(
        UdpTransportOptions options,
        IMessageProtocolAdapter protocolAdapter)
    {
        _options = options ?? throw new ArgumentNullException(nameof(options));
        _protocolAdapter = protocolAdapter ?? throw new ArgumentNullException(nameof(protocolAdapter));
 
        ValidateOptions(_options);
    }

    public ConnectionState State => (ConnectionState)Volatile.Read(ref _state);

    public TransportKind Kind => TransportKind.Udp;

    public event EventHandler<MessageEnvelope>? MessageReceived;

    public Task ConnectAsync(CancellationToken cancellationToken = default)
    {
        if (State is ConnectionState.Connected or ConnectionState.Connecting)
        {
            return Task.CompletedTask;
        }
 
        cancellationToken.ThrowIfCancellationRequested();
 
        SetState(ConnectionState.Connecting);
 
        try
        {
            var localEndPoint = _options.CreateLocalEndPoint();
            _remoteEndPoint = _options.CreateRemoteEndPoint();
 
            _client = new UdpClient(AddressFamily.InterNetwork)
            {
                EnableBroadcast = _options.EnableBroadcast
            };
 
            _client.Client.ReceiveBufferSize = _options.ReceiveBufferSize;
            _client.Client.SendBufferSize = _options.SendBufferSize;
 
            if (_options.ReuseAddress)
            {
                _client.Client.SetSocketOption(
                    SocketOptionLevel.Socket,
                    SocketOptionName.ReuseAddress,
                    true);
            }
 
            _client.Client.Bind(localEndPoint);
 
            _receiveLoopCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 
            SetState(ConnectionState.Connected);
 
            _receiveLoopTask = Task.Run(
                () => ReceiveLoopAsync(_receiveLoopCts.Token),
                _receiveLoopCts.Token);
 
            return Task.CompletedTask;
        }
        catch
        {
            SetState(ConnectionState.Faulted);
 
            _client?.Dispose();
            _client = null;
            _remoteEndPoint = null;
 
            _receiveLoopCts?.Dispose();
            _receiveLoopCts = null;
 
            throw;
        }
    }

    public async Task DisconnectAsync(CancellationToken cancellationToken = default)
    {
        if (State == ConnectionState.Disconnected)
        {
            return;
        }
 
        SetState(ConnectionState.Disconnecting);
 
        _receiveLoopCts?.Cancel();
 
        _client?.Close();
        _client?.Dispose();
        _client = null;
        _remoteEndPoint = null;
 
        if (_receiveLoopTask is not null)
        {
            try
            {
                await _receiveLoopTask.ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
            }
            catch (ObjectDisposedException)
            {
            }
            catch (SocketException)
            {
            }
        }
 
        _receiveLoopCts?.Dispose();
        _receiveLoopCts = null;
        _receiveLoopTask = null;
 
        cancellationToken.ThrowIfCancellationRequested();
 
        SetState(ConnectionState.Disconnected);
    }

    public async Task SendAsync(
        MessageEnvelope message,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(message);
 
        if (_client is null ||
            _remoteEndPoint is null ||
            State != ConnectionState.Connected)
        {
            throw new InvalidOperationException("UDP transport is not connected.");
        }
 
        cancellationToken.ThrowIfCancellationRequested();
 
        var payload = _protocolAdapter.Encode(message);
 
        await _client.SendAsync(payload, _remoteEndPoint, cancellationToken)
            .ConfigureAwait(false);
    }

    public async ValueTask DisposeAsync()
    {
        await DisconnectAsync().ConfigureAwait(false);
    }

    private async Task ReceiveLoopAsync(CancellationToken cancellationToken)
    {
        if (_client is null)
        {
            return;
        }
 
        try
        {
            while (!cancellationToken.IsCancellationRequested &&
                   State == ConnectionState.Connected)
            {
                var result = await _client.ReceiveAsync(cancellationToken)
                    .ConfigureAwait(false);
 
                var message = _protocolAdapter.Decode(result.Buffer);
                MessageReceived?.Invoke(this, message);
            }
        }
        catch (OperationCanceledException)
        {
        }
        catch (ObjectDisposedException)
        {
        }
        catch (SocketException)
        {
            if (!cancellationToken.IsCancellationRequested)
            {
                SetState(ConnectionState.Faulted);
            }
        }
        catch
        {
            if (!cancellationToken.IsCancellationRequested)
            {
                SetState(ConnectionState.Faulted);
            }
        }
    }

    private static void ValidateOptions(UdpTransportOptions options)
    {
        ArgumentNullException.ThrowIfNull(options);
        ArgumentException.ThrowIfNullOrWhiteSpace(options.LocalHost);
        ArgumentException.ThrowIfNullOrWhiteSpace(options.RemoteHost);
 
        ValidatePort(options.LocalPort, nameof(options.LocalPort));
        ValidatePort(options.RemotePort, nameof(options.RemotePort));
 
        if (options.ReceiveBufferSize <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(options.ReceiveBufferSize));
        }
 
        if (options.SendBufferSize <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(options.SendBufferSize));
        }
    }

    private static void ValidatePort(int port, string parameterName)
    {
        if (port <= 0 || port > 65535)
        {
            throw new ArgumentOutOfRangeException(parameterName);
        }
    }

    private void SetState(ConnectionState state)
    {
        Interlocked.Exchange(ref _state, (int)state);
    }
}