Ensure async task completes before sending packets (#5765)

## Description

Spurious failures of
Misc/WithReceiveResumeNoDataArgs.ReceiveResumeNoData/2 have been
observed in the OS repository.
From traces, this is due to a race between `StreamReceiveSetEnabled`
(processed async) and receiving the shutdown frame sent by the client:
it was possible the shutdown frame was indicated by the server before it
blocks receives.

## Testing

CI. This is a rare race, so validation is mostly for non-regression. 

## Documentation

N/A

---------

Co-authored-by: Michael Friesen <3517159+mtfriesen@users.noreply.github.com>

Author: guhetier

src/test/lib/DataTest.cpp

@@ -1989,6 +1989,11 @@ QuicTestReceiveResume(
                 TEST_FAILURE("PauseFirst MsQuic->StreamReceiveSetEnabled(FALSE) failed, 0x%x", Status);
                 return;
             }
+
+            //
+            // Ensure the receive pause task is processed.
+            //
+            DrainConnectionWorkQueue(ServerContext.Conn.Handle);
         }
 
         Status =
@@ -2206,6 +2211,9 @@ QuicTestReceiveResumeNoData(
             return;
         }
 
+        //
+        // Disable receive notifications on the server.
+        //
         Status =
             MsQuic->StreamReceiveSetEnabled(
                 ServerContext.Stream.Handle,
@@ -2215,6 +2223,12 @@ QuicTestReceiveResumeNoData(
             return;
         }
 
+        //
+        // Ensure the StreamReceiveSetEnabled task is processed so it doesn't race with receiving
+        // the shutdown.
+        //
+        DrainConnectionWorkQueue(ServerContext.Conn.Handle);
+
         Status =
             MsQuic->StreamShutdown(
                 ClientContext.Stream.Handle,
@@ -2227,10 +2241,18 @@ QuicTestReceiveResumeNoData(
         }
 
         if (ShutdownType == GracefulShutdown) {
+            //
+            // On a graceful shutdown, the shutdown notification is emitted only after all data
+            // was delivered (in this case, zero bytes and a FIN).
+            //
             if (CxPlatEventWaitWithTimeout(ServerContext.TestEvent.Handle, TimeoutMs)) {
                 TEST_FAILURE("Server got shutdown event when it shouldn't have!");
                 return;
             }
+
+            //
+            // Re-enable receive notifications.
+            //
             Status =
                 MsQuic->StreamReceiveSetEnabled(
                     ServerContext.Stream.Handle,
@@ -2242,7 +2264,7 @@ QuicTestReceiveResumeNoData(
         }
 
         //
-        // Validate the test was shutdown as expected.
+        // Validate the stream was shutdown as expected.
         //
         if (!CxPlatEventWaitWithTimeout(ServerContext.TestEvent.Handle, TimeoutMs)) {
             TEST_FAILURE("Server failed to get shutdown before timeout!");

src/test/lib/TestHelpers.h

@@ -1063,3 +1063,16 @@ WaitForMsQuicInUse() {
 
     return MsQuicInUse && Status == QUIC_STATUS_SUCCESS;
 }
+
+_IRQL_requires_max_(PASSIVE_LEVEL)
+QUIC_INLINE
+void
+DrainConnectionWorkQueue(HQUIC Connection) {
+    //
+    // Call the GetParam API: it schedules a work item on the connection and waits for its completion,
+    // ensuring all prior work items have completed.
+    //
+    QUIC_STATISTICS_V2 Stats{};
+    uint32_t StatsSize = sizeof(Stats);
+    (void)MsQuic->GetParam(Connection, QUIC_PARAM_CONN_STATISTICS_V2, &StatsSize, &Stats);
+}