Add a semaphore for receiving data through WebSocket (#64)

occluder · Dec 14, 2023 · c16f570 · c16f570
1 parent 5077248
commit c16f570
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 6 deletions.
diff --git a/MiniTwitch.Common/CHANGELOG.md b/MiniTwitch.Common/CHANGELOG.md
@@ -0,0 +1,5 @@
+# MiniTwitch.Common Changelog
+
+## Upcoming version
+
+- Minor: Improved handling of WebSocket reads
diff --git a/MiniTwitch.Common/Internal/Models/ByteBucket.cs b/MiniTwitch.Common/Internal/Models/ByteBucket.cs
@@ -23,9 +23,15 @@ public ByteBucket(short tempBufferSize = 2048)
     /// <returns><see langword="true"/> if the bucket should be drained</returns>
     public async ValueTask<bool> FillFrom(ClientWebSocket source, CancellationToken cToken)
     {
+        ValueTask<ValueWebSocketReceiveResult> valueTask = source.ReceiveAsync(_temp, cToken);
+        if (valueTask.IsCanceled || cToken.IsCancellationRequested)
+        {
+            return false;
+        }
+
         // Receive bytes into _temp
         // This overwrites previous data
-        ValueWebSocketReceiveResult result = await source.ReceiveAsync(_temp, cToken);
+        ValueWebSocketReceiveResult result = await valueTask.ConfigureAwait(false);
         // Copy the received bytes into the bucket
         _temp[..result.Count].CopyTo(_bucket[_reached..(result.Count + _reached)]);
         _reached += result.Count;

diff --git a/MiniTwitch.Common/WebSocketClient.cs b/MiniTwitch.Common/WebSocketClient.cs
@@ -22,8 +22,9 @@ public sealed class WebSocketClient : IAsyncDisposable
     #endregion
 
     #region Fields
-    private readonly SemaphoreSlim _sendLock = new(1);
     private readonly SemaphoreSlim _reconnectionLock = new(0);
+    private readonly SemaphoreSlim _receiveLock = new(1);
+    private readonly SemaphoreSlim _sendLock = new(1);
     private readonly TimeSpan _reconnectDelay;
     private readonly ByteBucket _bucket;
     private CancellationTokenSource _cts;
@@ -148,6 +149,7 @@ private async Task Receive()
         {
             try
             {
+                await _receiveLock.WaitAsync(_cts.Token);
                 // Continue if not at the end of message
                 if (!await _bucket.FillFrom(_client, _cts.Token))
                     continue;
@@ -157,7 +159,8 @@ private async Task Receive()
             }
             catch (WebSocketException wse)
             {
-                Log(LogLevel.Critical, "An error occurred while receiving data from the WebSocket connection: {msg}", wse.Message);
+                Log(LogLevel.Critical, "An error occurred while receiving data from the WebSocket connection: {msg}",
+                    wse.Message);
                 break;
             }
             catch (InvalidOperationException)
@@ -174,6 +177,10 @@ private async Task Receive()
             {
                 LogException(ex, "Exception caught in data receiver: ");
             }
+            finally
+            {
+                _ = _receiveLock.Release();
+            }
         }
 
         // Don't restart if it's a user disconnect
@@ -186,16 +193,16 @@ public async ValueTask SendAsync(string data, bool sensitive = false, Cancellati
         var (written, bytes) = StringToBytes(data);
         if (!await _sendLock.WaitAsync(TimeSpan.FromSeconds(10), cancellationToken).ConfigureAwait(false))
         {
-            Log(LogLevel.Warning, "{method} timed out after 10 seconds.", nameof(SendAsync));
+            Log(LogLevel.Warning, "{Method} timed out after 10 seconds.", nameof(SendAsync));
             return;
         }
         else if (!this.IsConnected)
         {
-            Log(LogLevel.Warning, "Cannot send data in non-connect state. ({state})", _client.State);
+            Log(LogLevel.Warning, "Cannot send data in non-connect state. ({State})", _client.State);
         }
 
         if (!sensitive)
-            Log(LogLevel.Debug, "Sending data: {msg}", data);
+            Log(LogLevel.Debug, "Sending data: {Message}", data);
 
         try
         {