Fixed GC collecting listener and/or connection while waiting on new connection/stream event (#74740)

Author: ManickaP

src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs

@@ -392,6 +392,7 @@ public async ValueTask<QuicStream> AcceptInboundStreamAsync(CancellationToken ca
             throw new InvalidOperationException(SR.net_quic_accept_not_allowed);
         }
 
+        GCHandle keepObject = GCHandle.Alloc(this);
         try
         {
             return await _acceptQueue.Reader.ReadAsync(cancellationToken).ConfigureAwait(false);
@@ -401,6 +402,10 @@ public async ValueTask<QuicStream> AcceptInboundStreamAsync(CancellationToken ca
             ExceptionDispatchInfo.Capture(ex.InnerException).Throw();
             throw;
         }
+        finally
+        {
+            keepObject.Free();
+        }
     }
 
     /// <summary>

src/libraries/System.Net.Quic/src/System/Net/Quic/QuicListener.cs

@@ -162,6 +162,7 @@ public async ValueTask<QuicConnection> AcceptConnectionAsync(CancellationToken c
     {
         ObjectDisposedException.ThrowIf(_disposed == 1, this);
 
+        GCHandle keepObject = GCHandle.Alloc(this);
         try
         {
             PendingConnection pendingConnection = await _acceptQueue.Reader.ReadAsync(cancellationToken).ConfigureAwait(false);
@@ -175,6 +176,10 @@ public async ValueTask<QuicConnection> AcceptConnectionAsync(CancellationToken c
             ExceptionDispatchInfo.Capture(ex.InnerException).Throw();
             throw;
         }
+        finally
+        {
+            keepObject.Free();
+        }
     }
 
     private unsafe int HandleEventNewConnection(ref NEW_CONNECTION_DATA data)

src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs

@@ -336,5 +336,33 @@ public async Task Connect_PeerCertificateDisposed(bool useGetter)
             }
             peerCertificate.Dispose();
         }
+
+        [Fact]
+        public async Task Connection_AwaitsStream_ConnectionSurvivesGC()
+        {
+            const byte data = 0xDC;
+
+            TaskCompletionSource<IPEndPoint> listenerEndpointTcs = new TaskCompletionSource<IPEndPoint>();
+            await Task.WhenAll(
+                Task.Run(async () =>
+                {
+                    await using var listener = await CreateQuicListener();
+                    listenerEndpointTcs.SetResult(listener.LocalEndPoint);
+                    await using var connection = await listener.AcceptConnectionAsync();
+                    await using var stream = await connection.AcceptInboundStreamAsync();
+                    var buffer = new byte[1];
+                    Assert.Equal(1, await stream.ReadAsync(buffer));
+                    Assert.Equal(data, buffer[0]);
+                }).WaitAsync(TimeSpan.FromSeconds(5)),
+                Task.Run(async () =>
+                {
+                    var endpoint = await listenerEndpointTcs.Task;
+                    await using var connection = await CreateQuicConnection(endpoint);
+                    await Task.Delay(TimeSpan.FromSeconds(0.5));
+                    GC.Collect();
+                    await using var stream = await connection.OpenOutboundStreamAsync(QuicStreamType.Unidirectional);
+                    await stream.WriteAsync(new byte[1] { data }, completeWrites: true);
+                }).WaitAsync(TimeSpan.FromSeconds(5)));
+        }
     }
 }

src/libraries/System.Net.Quic/tests/FunctionalTests/QuicListenerTests.cs

@@ -106,7 +106,7 @@ public async Task TwoListenersOnSamePort_DisjointAlpn_Success()
             QuicListenerOptions listenerOptions = CreateQuicListenerOptions();
             listenerOptions.ListenEndPoint = listener1.LocalEndPoint;
             listenerOptions.ApplicationProtocols[0] = new SslApplicationProtocol("someprotocol");
-            listenerOptions.ConnectionOptionsCallback = (_, _, _) => 
+            listenerOptions.ConnectionOptionsCallback = (_, _, _) =>
             {
                 var options = CreateQuicServerOptions();
                 options.ServerAuthenticationOptions.ApplicationProtocols[0] = listenerOptions.ApplicationProtocols[0];
@@ -144,5 +144,27 @@ public async Task TwoListenersOnSamePort_SameAlpn_Throws()
             //
             await AssertThrowsQuicExceptionAsync(QuicError.InternalError, async () => await CreateQuicListener(listener.LocalEndPoint));
         }
+
+        [Fact]
+        public async Task Listener_AwaitsConnection_ListenerSurvivesGC()
+        {
+            TaskCompletionSource<IPEndPoint> listenerEndpointTcs = new TaskCompletionSource<IPEndPoint>();
+            await Task.WhenAll(
+                Task.Run(async () =>
+                {
+                    await using var listener = await CreateQuicListener();
+                    listenerEndpointTcs.SetResult(listener.LocalEndPoint);
+                    var connection = await listener.AcceptConnectionAsync();
+                    await connection.DisposeAsync();
+                }).WaitAsync(TimeSpan.FromSeconds(5)),
+                Task.Run(async () =>
+                {
+                    var endpoint = await listenerEndpointTcs.Task;
+                    await Task.Delay(TimeSpan.FromSeconds(0.5));
+                    GC.Collect();
+                    var connection = await CreateQuicConnection(endpoint);
+                    await connection.DisposeAsync();
+                }).WaitAsync(TimeSpan.FromSeconds(5)));
+        }
     }
 }

src/libraries/System.Net.Quic/tests/FunctionalTests/QuicStreamConnectedStreamConformanceTests.cs

@@ -102,7 +102,7 @@ await WhenAllOrAnyFailed(
                         }
                         catch (Exception ex)
                         {
-                            _output?.WriteLine($"Failed to {ex.Message}");
+                            _output?.WriteLine($"Failed to connect: {ex.Message}");
                             throw;
                         }
                     }));
@@ -153,14 +153,5 @@ public override void Dispose()
                 }
             }
         }
-
-        [OuterLoop("May take several seconds")]
-        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
-        [SkipOnPlatform(TestPlatforms.LinuxBionic, "SElinux blocks UNIX sockets in our CI environment")]
-        [ActiveIssue("https://github.com/dotnet/runtime/issues/73377")]
-        public override Task Parallel_ReadWriteMultipleStreamsConcurrently()
-        {
-            return Task.CompletedTask;
-        }
     }
 }