From 14eca7cb1e21908e3f77307b0e6c404ea1241f10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marie=20P=C3=ADchov=C3=A1?=
 <11718369+ManickaP@users.noreply.github.com>
Date: Wed, 21 Aug 2024 17:00:54 +0200
Subject: [PATCH] [QUIC] Add synchronize for streams available callback
 (#106678)

* Add synchronize for streams available callback

* Feedback
---
 .../tests/FunctionalTests/QuicConnectionTests.cs    | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs b/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs
index 79874b31d4895..f059f28d57719 100644
--- a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs
+++ b/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs
@@ -455,11 +455,13 @@ public async Task GetStreamCapacity_SumInvariant()
         {
             int maxStreamIndex = 0;
             const int Limit = 5;
+            SemaphoreSlim streamsAvailableFired = new SemaphoreSlim(0);
 
             var clientOptions = CreateQuicClientOptions(new IPEndPoint(0, 0));
             clientOptions.StreamCapacityCallback = (connection, args) =>
             {
                 Interlocked.Add(ref maxStreamIndex, args.BidirectionalIncrement);
+                streamsAvailableFired.Release();
             };
 
             var listenerOptions = CreateQuicListenerOptions();
@@ -472,6 +474,10 @@ public async Task GetStreamCapacity_SumInvariant()
 
             (QuicConnection clientConnection, QuicConnection serverConnection) = await CreateConnectedQuicConnection(clientOptions, listenerOptions);
 
+            while (maxStreamIndex < Limit)
+            {
+                await streamsAvailableFired.WaitAsync().WaitAsync(PassingTestTimeout);
+            }
             Assert.Equal(Limit, maxStreamIndex);
 
             Queue<(QuicStream client, QuicStream server)> streams = new();
@@ -509,8 +515,11 @@ public async Task GetStreamCapacity_SumInvariant()
                 }
             }
 
-            // give time to update the count
-            await Task.Delay(1000);
+            // wait for the callback
+            while (maxStreamIndex < 3 * Limit)
+            {
+                await streamsAvailableFired.WaitAsync().WaitAsync(PassingTestTimeout);
+            }
 
             // by now, we opened and closed 2 * Limit, and expect a budget of 'Limit' more
             Assert.Equal(3 * Limit, maxStreamIndex);