Fix threading issues?

LobbyClient/TcpClient.cs

@@ -1,4 +1,5 @@
-using System;
+using LobbyServerDto;
+using System;
 using System.IO;
 using System.Net.Sockets;
 using System.Threading;
@@ -25,7 +26,20 @@ namespace Lobbies
         private NetworkStream? networkStream;
         private CancellationTokenSource? cancellationTokenSource = new CancellationTokenSource();
         private bool running = false;
-        private readonly SemaphoreSlim sendLock = new SemaphoreSlim(1, 1);
+        private readonly BufferRental bufferRental = new BufferRental(MaxMessageSize);
+        private async Task ReadExact(NetworkStream stream, Memory<byte> buffer, int length, CancellationToken token)
+        {
+            int offset = 0;
+
+            while (offset < length)
+            {
+                int read = await stream.ReadAsync(buffer.Slice(offset), token);
+                if (read == 0)
+                    throw new EndOfStreamException();
+
+                offset += read;
+            }
+        }
 
         internal async Task Connect(string host, int port)
         {
@@ -60,81 +74,26 @@ namespace Lobbies
                 pingTask = Task.Run(() => PingLoop(token), token);
 
                 Memory<byte> buffer = new byte[MaxMessageSize];
-                Memory<byte> target = new byte[MaxMessageSize];
-
-                int bufferedBytes = 0;
-                int currentMessageLength = -1;
-                int currentMessageOffset = 0;
 
                 while (running && !token.IsCancellationRequested)
                 {
-                    if (bufferedBytes == buffer.Length)
-                        throw new InvalidDataException("Receive buffer overflow.");
+                    await ReadExact(networkStream, buffer, HeaderSize, token);                    
 
-                    int bytesRead = await networkStream.ReadAsync(buffer.Slice(bufferedBytes), token);
-                    if (bytesRead == 0)
-                        break;
+                    int length = BitConverter.ToInt32(buffer.Span.Slice(0, HeaderSize));
 
-                    bufferedBytes += bytesRead;
+                    if (length < 0)
+                        throw new InvalidDataException("Negative message length received.");
 
-                    while (true)
-                    {
-                        if (currentMessageLength < 0)
-                        {
-                            if (bufferedBytes < HeaderSize)
-                                break;
+                    if (length > MaxMessageSize)
+                        throw new InvalidDataException($"Message too large: {length} > {MaxMessageSize}.");
 
-                            currentMessageLength = BitConverter.ToInt32(buffer.Span.Slice(0, HeaderSize));
+                    if (length == 0)
+                        continue;
 
-                            if (currentMessageLength < 0)
-                                throw new InvalidDataException("Negative message length received.");
+                    await ReadExact(networkStream, buffer, length, token);
 
-                            if (currentMessageLength > MaxMessageSize)
-                                throw new InvalidDataException($"Message too large: {currentMessageLength} > {MaxMessageSize}.");
-
-                            if (bufferedBytes > HeaderSize)
-                                buffer.Slice(HeaderSize, bufferedBytes - HeaderSize).CopyTo(buffer);
-
-                            bufferedBytes -= HeaderSize;
-                            currentMessageOffset = 0;
-
-                            if (currentMessageLength == 0)
-                            {
-                                currentMessageLength = -1;
-                                continue;
-                            }
-                        }
-
-                        int remainingMessageBytes = currentMessageLength - currentMessageOffset;
-                        if (remainingMessageBytes <= 0)
-                        {
-                            DataReceived?.Invoke(currentMessageLength, target.Slice(0, currentMessageLength));
-                            currentMessageLength = -1;
-                            currentMessageOffset = 0;
-                            continue;
-                        }
-
-                        if (bufferedBytes == 0)
-                            break;
-
-                        int chunkSize = Math.Min(bufferedBytes, remainingMessageBytes);
-
-                        buffer.Slice(0, chunkSize).CopyTo(target.Slice(currentMessageOffset));
-                        currentMessageOffset += chunkSize;
-
-                        if (bufferedBytes > chunkSize)
-                            buffer.Slice(chunkSize, bufferedBytes - chunkSize).CopyTo(buffer);
-
-                        bufferedBytes -= chunkSize;
-
-                        if (currentMessageOffset == currentMessageLength)
-                        {
-                            DataReceived?.Invoke(currentMessageLength, target.Slice(0, currentMessageLength));
-                            currentMessageLength = -1;
-                            currentMessageOffset = 0;
-                        }
-                    }
-                }
+                    DataReceived?.Invoke(length, buffer.Slice(0, length));
+                }                         
             }
             catch (OperationCanceledException)
             {
@@ -185,37 +144,42 @@ namespace Lobbies
             if (!running || networkStream == null)
                 return;
 
-            await sendLock.WaitAsync(token);
             try
             {
                 await networkStream.WriteAsync(BitConverter.GetBytes(0), 0, 4, token);
             }
-            finally
+            catch
             {
-                sendLock.Release();
             }
         }
 
         internal async Task Send(byte[] buffer, int offset, int count)
         {
+            byte[] frame = bufferRental.Rent();
             try
             {
                 if (!running || networkStream == null || cancellationTokenSource == null)
                     return;
 
-                await sendLock.WaitAsync(cancellationTokenSource.Token);
-                try
-                {
-                    await networkStream.WriteAsync(BitConverter.GetBytes(count), 0, 4, cancellationTokenSource.Token);
-                    await networkStream.WriteAsync(buffer, offset, count, cancellationTokenSource.Token);
-                }
-                finally
-                {
-                    sendLock.Release();
-                }
+                if (count < 0 || count > MaxMessageSize)
+                    throw new ArgumentOutOfRangeException(nameof(count));
+
+                if (offset < 0 || offset + count > buffer.Length)
+                    throw new ArgumentOutOfRangeException(nameof(offset));
+
+                BitConverter.GetBytes(count).CopyTo(frame, 0);
+                Buffer.BlockCopy(buffer, offset, frame, HeaderSize, count);
+
+                await networkStream.WriteAsync(frame, 0, HeaderSize + count, cancellationTokenSource.Token);
             }
-            catch { }
-        }
+            catch
+            {
+            }
+            finally
+            {
+                bufferRental.Return(frame);
+            }
+        }      
 
         internal void Stop()
         {