SpinQuic: Pend some receives and try resuming them. (microsoft#1613)

boscorona · May 18, 2021 · 7977dda · 7977dda
1 parent 3163953
commit 7977dda
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/src/tools/spin/spinquic.cpp b/src/tools/spin/spinquic.cpp
@@ -113,6 +113,7 @@ typedef enum {
     SpinQuicAPICallGetParamConnection,
     SpinQuicAPICallGetParamStream,
     SpinQuicAPICallDatagramSend,
+    SpinQuicAPICallReceiveSetEnabled,
     SpinQuicAPICallCount    // Always the last element
 } SpinQuicAPICall;
 
@@ -200,6 +201,11 @@ QUIC_STATUS QUIC_API SpinQuicHandleStreamEvent(HQUIC Stream, void * /* Context *
     case QUIC_STREAM_EVENT_PEER_SEND_SHUTDOWN:
         MsQuic->StreamShutdown(Stream, (QUIC_STREAM_SHUTDOWN_FLAGS)GetRandom(16), 0);
         break;
+    case QUIC_STREAM_EVENT_RECEIVE:
+        if (GetRandom(5) == 0) {
+            return QUIC_STATUS_PENDING;
+        }
+        break;
     default:
         break;
     }
@@ -519,6 +525,18 @@ void Spin(LockableVector<HQUIC>& Connections, std::vector<HQUIC>* Listeners = nu
             }
             break;
         }
+        case SpinQuicAPICallReceiveSetEnabled: {
+            auto Connection = Connections.TryGetRandom();
+            BAIL_ON_NULL_CONNECTION(Connection);
+            auto ctx = SpinQuicConnection::Get(Connection);
+            {
+                std::lock_guard<std::mutex> Lock(ctx->Lock);
+                auto Stream = ctx->TryGetStream();
+                if (Stream == nullptr) continue;
+                MsQuic->StreamReceiveSetEnabled(Stream, GetRandom(2) == 0);
+            }
+            break;
+        }
         case SpinQuicAPICallStreamShutdown: {
             auto Connection = Connections.TryGetRandom();
             BAIL_ON_NULL_CONNECTION(Connection);