<commit_before>﻿using System;
using System.Buffers;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Lime.Protocol;
using Lime.Protocol.Network;
using Lime.Protocol.Serialization;

namespace Lime.Transport.WebSocket
{
    public abstract class WebSocketTransport : TransportBase, IDisposable
    {
        public const int DEFAULT_BUFFER_SIZE = 8192;
        public const int DEFAULT_MAX_BUFFER_COUNT = 1024;

        private static readonly TimeSpan CloseTimeout = TimeSpan.FromSeconds(5);

        private readonly IEnvelopeSerializer _envelopeSerializer;
        private readonly ITraceWriter _traceWriter;
        private readonly ArrayPool<byte> _arrayPool;
        private readonly int _bufferSize;
        private readonly SemaphoreSlim _sendSemaphore;
        private readonly SemaphoreSlim _receiveSemaphore;
        private readonly SemaphoreSlim _closeSemaphore;
        private readonly WebSocketMessageType _webSocketMessageType;
        private readonly CancellationTokenSource _receiveCts;

        private WebSocketReceiveResult _closeFrame;        
        protected WebSocketCloseStatus CloseStatus;
        protected string CloseStatusDescription;

        protected WebSocketTransport(
            System.Net.WebSockets.WebSocket webSocket,
            IEnvelopeSerializer envelopeSerializer,
            ITraceWriter traceWriter = null,
            int bufferSize = DEFAULT_BUFFER_SIZE,
            WebSocketMessageType webSocketMessageType = WebSocketMessageType.Text,
            ArrayPool<byte> arrayPool = null)
        {
            WebSocket = webSocket;
            _envelopeSerializer = envelopeSerializer;
            _traceWriter = traceWriter;
            _arrayPool = arrayPool ?? ArrayPool<byte>.Shared;
            _bufferSize = bufferSize;
            _webSocketMessageType = webSocketMessageType;
            _sendSemaphore = new SemaphoreSlim(1);
            _receiveSemaphore = new SemaphoreSlim(1);
            _closeSemaphore = new SemaphoreSlim(1);
            _receiveCts = new CancellationTokenSource();
            CloseStatus = WebSocketCloseStatus.NormalClosure;
            CloseStatusDescription = string.Empty;
        }

        protected System.Net.WebSockets.WebSocket WebSocket { get; }

        public override async Task SendAsync(Envelope envelope, CancellationToken cancellationToken)
        {
            if (envelope == null) throw new ArgumentNullException(nameof(envelope));
            EnsureOpen("send");

            var serializedEnvelope = _envelopeSerializer.Serialize(envelope);

            if (_traceWriter != null &&
                _traceWriter.IsEnabled)
            {
                await _traceWriter.TraceAsync(serializedEnvelope, DataOperation.Send).ConfigureAwait(false);
            }

            var buffer = Encoding.UTF8.GetBytes(serializedEnvelope);
            

            await _sendSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);

            try
            {
                EnsureOpen("send");

                await WebSocket.SendAsync(new ArraySegment<byte>(buffer), _webSocketMessageType, true,
                    cancellationToken).ConfigureAwait(false);
            }
            catch (WebSocketException)
            {
                await CloseWithTimeoutAsync().ConfigureAwait(false);
                throw;
            }
            finally
            {
                _sendSemaphore.Release();
            }
        }

        public override async Task<Envelope> ReceiveAsync(CancellationToken cancellationToken)
        {
            EnsureOpen("receive");

            var segments = new List<BufferSegment>();

            await _receiveSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);

            try
            {
                EnsureOpen("receive");

                while (!_receiveCts.IsCancellationRequested)
                {
                    cancellationToken.ThrowIfCancellationRequested();

                    var segment = new BufferSegment
                    {
                        Buffer = _arrayPool.Rent(_bufferSize)
                    };
                    segments.Add(segment);

                    // The websocket class go to the 'Aborted' state if the receiveasync operation is cancelled.
                    // In this case, we are unable to close the connection clearly when required.
                    var receiveTask = WebSocket
                        .ReceiveAsync(new ArraySegment<byte>(segment.Buffer), _receiveCts.Token)
                        .ContinueWith(t =>
                        {
                            if (t.IsCanceled) return null;
                            return t.Result;
                        });
                    var cancellationTask = cancellationToken.AsTask();

                    // If the token is cancelled
                    var completedTask = await Task.WhenAny(receiveTask, cancellationTask).ConfigureAwait(false);
                    if (completedTask != receiveTask)
                    {
                        // The task above will thrown a TaskCancelledException, but just in case...
                        cancellationToken.ThrowIfCancellationRequested();
                    }

                    var receiveResult = receiveTask.Result;
                    if (receiveResult == null) continue;

                    segment.Count = receiveResult.Count;

                    if (receiveResult.MessageType == WebSocketMessageType.Close)
                    {
                        HandleCloseMessage(receiveResult);
                        break;
                    }

                    if (receiveResult.MessageType != _webSocketMessageType)
                    {
                        CloseStatus = WebSocketCloseStatus.InvalidMessageType;
                        CloseStatusDescription = "An unsupported message type was received";<commit_msg>Correcly handling exceptions on WS receive operation<commit_after>﻿using System;
using System.Buffers;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Lime.Protocol;
using Lime.Protocol.Network;
using Lime.Protocol.Serialization;

namespace Lime.Transport.WebSocket
{
    public abstract class WebSocketTransport : TransportBase, IDisposable
    {
        public const int DEFAULT_BUFFER_SIZE = 8192;
        public const int DEFAULT_MAX_BUFFER_COUNT = 1024;

        private static readonly TimeSpan CloseTimeout = TimeSpan.FromSeconds(5);

        private readonly IEnvelopeSerializer _envelopeSerializer;
        private readonly ITraceWriter _traceWriter;
        private readonly ArrayPool<byte> _arrayPool;
        private readonly int _bufferSize;
        private readonly SemaphoreSlim _sendSemaphore;
        private readonly SemaphoreSlim _receiveSemaphore;
        private readonly SemaphoreSlim _closeSemaphore;
        private readonly WebSocketMessageType _webSocketMessageType;
        private readonly CancellationTokenSource _receiveCts;

        private WebSocketReceiveResult _closeFrame;        
        protected WebSocketCloseStatus CloseStatus;
        protected string CloseStatusDescription;

        protected WebSocketTransport(
            System.Net.WebSockets.WebSocket webSocket,
            IEnvelopeSerializer envelopeSerializer,
            ITraceWriter traceWriter = null,
            int bufferSize = DEFAULT_BUFFER_SIZE,
            WebSocketMessageType webSocketMessageType = WebSocketMessageType.Text,
            ArrayPool<byte> arrayPool = null)
        {
            WebSocket = webSocket;
            _envelopeSerializer = envelopeSerializer;
            _traceWriter = traceWriter;
            _arrayPool = arrayPool ?? ArrayPool<byte>.Shared;
            _bufferSize = bufferSize;
            _webSocketMessageType = webSocketMessageType;
            _sendSemaphore = new SemaphoreSlim(1);
            _receiveSemaphore = new SemaphoreSlim(1);
            _closeSemaphore = new SemaphoreSlim(1);
            _receiveCts = new CancellationTokenSource();
            CloseStatus = WebSocketCloseStatus.NormalClosure;
            CloseStatusDescription = string.Empty;
        }

        protected System.Net.WebSockets.WebSocket WebSocket { get; }

        public override async Task SendAsync(Envelope envelope, CancellationToken cancellationToken)
        {
            if (envelope == null) throw new ArgumentNullException(nameof(envelope));
            EnsureOpen("send");

            var serializedEnvelope = _envelopeSerializer.Serialize(envelope);

            if (_traceWriter != null &&
                _traceWriter.IsEnabled)
            {
                await _traceWriter.TraceAsync(serializedEnvelope, DataOperation.Send).ConfigureAwait(false);
            }

            var buffer = Encoding.UTF8.GetBytes(serializedEnvelope);
            

            await _sendSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);

            try
            {
                EnsureOpen("send");

                await WebSocket.SendAsync(new ArraySegment<byte>(buffer), _webSocketMessageType, true,
                    cancellationToken).ConfigureAwait(false);
            }
            catch (WebSocketException)
            {
                await CloseWithTimeoutAsync().ConfigureAwait(false);
                throw;
            }
            finally
            {
                _sendSemaphore.Release();
            }
        }

        public override async Task<Envelope> ReceiveAsync(CancellationToken cancellationToken)
        {
            EnsureOpen("receive");

            var segments = new List<BufferSegment>();

            await _receiveSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);

            try
            {
                EnsureOpen("receive");

                while (!_receiveCts.IsCancellationRequested)
                {
                    cancellationToken.ThrowIfCancellationRequested();

                    var segment = new BufferSegment
                    {
                        Buffer = _arrayPool.Rent(_bufferSize)
                    };
                    segments.Add(segment);

                    // The websocket class go to the 'Aborted' state if the receiveasync operation is cancelled.
                    // In this case, we are unable to close the connection clearly when required.
                    var receiveTask = WebSocket.ReceiveAsync(new ArraySegment<byte>(segment.Buffer), _receiveCts.Token);
                    var cancellationTask = cancellationToken.AsTask();

                    // If the token is cancelled
                    var completedTask = await Task.WhenAny(receiveTask, cancellationTask).ConfigureAwait(false);
                    if (completedTask != receiveTask)
                    {
                        // The task above will thrown a TaskCancelledException, but just in case...
                        cancellationToken.ThrowIfCancellationRequested();
                    }

                    var receiveResult = await receiveTask;
                    if (receiveResult == null) continue;

                    segment.Count = receiveResult.Count;

                    if (receiveResult.MessageType == WebSocketMessageType.Close)
                    {
                        HandleCloseMessage(receiveResult);
                        break;
                    }

                    if (receiveResult.MessageType != _webSocketMessageType)
                    {
                        CloseStatus = WebSocketCloseStatus.InvalidMessageType;
                        CloseStatusDescription = "An unsupported message type was received";
{'fraction_non_alphanumeric': 0.061146496815286625, 'fraction_numerical': 0.002208067940552017, 'mean_word_length': 1.8883983321069413, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4734121', 'n_tokens_mistral': 2748, 'n_tokens_neox': 2634, 'n_words': 606}