[QUIC] Fixed disabled receives if buffer gets drained before RECEIVE …

…returns (#72439) * Fixed disabled receives if buffer gets drained before RECEIVE returns * Feedback
dotnet · Jul 19, 2022 · 518d6cb · 518d6cb
1 parent c22fa85
commit 518d6cb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ReceiveBuffers.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ReceiveBuffers.cs
@@ -28,6 +28,14 @@ public void SetFinal()
         }
     }
 
+    public bool HasCapacity()
+    {
+        lock (_syncRoot)
+        {
+            return _buffer.ActiveMemory.Length < MaxBufferedBytes;
+        }
+    }
+
     public int CopyFrom(ReadOnlySpan<QUIC_BUFFER> quicBuffers, int totalLength, bool final)
     {
         lock (_syncRoot)

diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs
@@ -509,7 +509,7 @@ private unsafe int HandleEventReceive(ref RECEIVE data)
         _receiveTcs.TrySetResult();
 
         data.TotalBufferLength = totalCopied;
-        return QUIC_STATUS_SUCCESS;
+        return (_receiveBuffers.HasCapacity() && Interlocked.CompareExchange(ref _receivedNeedsEnable, 0, 1) == 1) ? QUIC_STATUS_CONTINUE : QUIC_STATUS_SUCCESS;
     }
     private unsafe int HandleEventSendComplete(ref SEND_COMPLETE data)
     {