Prevent the receive completion operation queued if unnecessary

src/core/api.c

@@ -1321,11 +1321,13 @@ MsQuicStreamReceiveComplete(
         Stream,
         BufferLength);
 
+    CxPlatRwLockAcquireShared(&Stream->ReceiveCompleteInlineLock);
+    BOOLEAN ReceiveCompleteInline = Stream->Flags.ReceiveCallActive;
     InterlockedExchangeAdd64(
         (int64_t*)&Stream->RecvCompletionLength, (int64_t)BufferLength);
+    CxPlatRwLockReleaseShared(&Stream->ReceiveCompleteInlineLock);
 
-    if (Connection->WorkerThreadID == CxPlatCurThreadID() &&
-        Stream->Flags.ReceiveCallActive) {
+    if (ReceiveCompleteInline) {
         goto Exit; // No need to queue a completion operation when run inline
     }
 

src/core/stream.c

@@ -76,6 +76,7 @@ QuicStreamInitialize(
     QuicRangeInitialize(
         QUIC_MAX_RANGE_ALLOC_SIZE,
         &Stream->SparseAckRanges);
+    CxPlatRwLockInitialize(&Stream->ReceiveCompleteInlineLock);
     Stream->ReceiveCompleteOperation = &Stream->ReceiveCompleteOperationStorage;
     Stream->ReceiveCompleteOperationStorage.API_CALL.Context = &Stream->ReceiveCompleteApiCtxStorage;
     Stream->ReceiveCompleteOperation->Type = QUIC_OPER_TYPE_API_CALL;
@@ -159,6 +160,7 @@ QuicStreamInitialize(
 #endif
         QuicPerfCounterDecrement(QUIC_PERF_COUNTER_STRM_ACTIVE);
         CxPlatDispatchLockUninitialize(&Stream->ApiSendRequestLock);
+        CxPlatRwLockUninitialize(&Stream->ReceiveCompleteInlineLock);
         Stream->Flags.Freed = TRUE;
         CxPlatPoolFree(&Worker->StreamPool, Stream);
     }
@@ -201,6 +203,7 @@ QuicStreamFree(
     QuicRecvBufferUninitialize(&Stream->RecvBuffer);
     QuicRangeUninitialize(&Stream->SparseAckRanges);
     CxPlatDispatchLockUninitialize(&Stream->ApiSendRequestLock);
+    CxPlatRwLockUninitialize(&Stream->ReceiveCompleteInlineLock);
     CxPlatRefUninitialize(&Stream->RefCount);
 
     if (Stream->RecvBuffer.PreallocatedChunk) {

src/core/stream.h

@@ -418,6 +418,11 @@ typedef struct QUIC_STREAM {
     //
     volatile uint64_t RecvCompletionLength;
 
+    //
+    // This RW lock protects `Flags.ReceiveCallActive`.
+    //
+    CXPLAT_RW_LOCK ReceiveCompleteInlineLock;
+
     //
     // The error code for why the receive path was shutdown.
     //

src/core/stream_recv.c

@@ -873,6 +873,15 @@ QuicStreamRecvFlush(
     while (FlushRecv) {
         CXPLAT_DBG_ASSERT(!Stream->Flags.SentStopSending);
 
+        CxPlatRwLockAcquireExclusive(&Stream->ReceiveCompleteInlineLock);
+        Stream->Flags.ReceiveCallActive = TRUE;
+        //
+        // Save the current value of RecvCompletionLength to check if it has
+        // changed during the app callback.
+        //
+        uint64_t RecvCompletionLength = Stream->RecvCompletionLength;
+        CxPlatRwLockReleaseExclusive(&Stream->ReceiveCompleteInlineLock);
+
         QUIC_BUFFER RecvBuffers[3];
         QUIC_STREAM_EVENT Event = {0};
         Event.Type = QUIC_STREAM_EVENT_RECEIVE;
@@ -922,7 +931,6 @@ QuicStreamRecvFlush(
         }
 
         Stream->Flags.ReceiveEnabled = Stream->Flags.ReceiveMultiple;
-        Stream->Flags.ReceiveCallActive = TRUE;
         Stream->RecvPendingLength += Event.RECEIVE.TotalBufferLength;
         CXPLAT_DBG_ASSERT(Stream->RecvPendingLength <= Stream->RecvBuffer.ReadPendingLength);
 
@@ -936,12 +944,24 @@ QuicStreamRecvFlush(
 
         QUIC_STATUS Status = QuicStreamIndicateEvent(Stream, &Event);
 
+        CxPlatRwLockAcquireExclusive(&Stream->ReceiveCompleteInlineLock);
         Stream->Flags.ReceiveCallActive = FALSE;
+        //
+        // Check if the app has called the receive complete API.
+        //
+        BOOLEAN ReceiveCompleteCalled = (Stream->RecvCompletionLength > RecvCompletionLength);
+        //
+        // Save the current value of RecvCompletionLength to determine the length of
+        // data which will be processed in current iteration.
+        //
+        RecvCompletionLength = Stream->RecvCompletionLength;
+
+        CxPlatRwLockReleaseExclusive(&Stream->ReceiveCompleteInlineLock);
 
         if (Status == QUIC_STATUS_CONTINUE) {
             CXPLAT_DBG_ASSERT(!Stream->Flags.SentStopSending);
             InterlockedExchangeAdd64(
-                (int64_t*)&Stream->RecvCompletionLength,
+                (int64_t*)&RecvCompletionLength,
                 (int64_t)Event.RECEIVE.TotalBufferLength);
             FlushRecv = TRUE;
             //
@@ -951,12 +971,7 @@ QuicStreamRecvFlush(
             Stream->Flags.ReceiveEnabled = TRUE;
 
         } else if (Status == QUIC_STATUS_PENDING) {
-            //
-            // The app called the receive complete API inline if
-            // RecvCompletionLength is non-zero.
-            //
-            FlushRecv = (Stream->RecvCompletionLength != 0);
-
+            FlushRecv = ReceiveCompleteCalled;
         } else {
             //
             // All failure status returns shouldn't be used by the app are
@@ -969,13 +984,13 @@ QuicStreamRecvFlush(
                 Status, 0);
 
             InterlockedExchangeAdd64(
-                (int64_t*)&Stream->RecvCompletionLength,
+                (int64_t*)&RecvCompletionLength,
                 (int64_t)Event.RECEIVE.TotalBufferLength);
             FlushRecv = TRUE;
         }
 
         if (FlushRecv) {
-            uint64_t BufferLength = Stream->RecvCompletionLength;
+            uint64_t BufferLength = RecvCompletionLength;
             InterlockedExchangeAdd64(
                 (int64_t*)&Stream->RecvCompletionLength,
                 -(int64_t)BufferLength);