[QUIC] Root listener/connection while waiting on new connection/stream event (#74450)

ManickaP · web-flow · commit 37eff15f3cef · 2022-08-26T20:25:06.000+02:00
* Fixed GC collecting listener and/or connection while waiting on new connection/stream event

* Minor fixes and cleanups
diff --git a/src/libraries/System.Net.Http/src/System.Net.Http.csproj b/src/libraries/System.Net.Http/src/System.Net.Http.csproj
@@ -446,7 +446,6 @@
     <Reference Include="System.Diagnostics.DiagnosticSource" />
     <Reference Include="System.Diagnostics.Tracing" />
     <Reference Include="System.IO.Compression" />
-    <Reference Include="System.Linq" />
     <Reference Include="System.Memory" />
     <Reference Include="System.Net.NameResolution" />
     <Reference Include="System.Net.NetworkInformation" />
diff --git a/src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/Http3Connection.cs b/src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/Http3Connection.cs
@@ -7,7 +7,6 @@
 using System.Runtime.Versioning;
 using System.Net.Quic;
 using System.IO;
-using System.Linq;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Globalization;
@@ -155,7 +154,7 @@ private void CheckForShutdown()
 
                     if (_clientControl != null)
                     {
-                        _clientControl.Dispose();
+                        await _clientControl.DisposeAsync().ConfigureAwait(false);
                         _clientControl = null;
                     }
 
@@ -245,7 +244,10 @@ public async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, lon
             }
             finally
             {
-                requestStream?.Dispose();
+                if (requestStream is not null)
+                {
+                    await requestStream.DisposeAsync().ConfigureAwait(false);
+                }
             }
         }
 
@@ -562,7 +564,6 @@ private async Task ProcessServerStreamAsync(QuicStream stream)
                             }
 
                             stream.Abort(QuicAbortDirection.Read, (long)Http3ErrorCode.StreamCreationError);
-                            stream.Dispose();
                             return;
                     }
                 }
diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs b/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>
diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicListener.cs b/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)
diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs
@@ -398,7 +398,6 @@ public void Abort(QuicAbortDirection abortDirection, long errorCode)
         QUIC_STREAM_SHUTDOWN_FLAGS flags = QUIC_STREAM_SHUTDOWN_FLAGS.NONE;
         if (abortDirection.HasFlag(QuicAbortDirection.Read))
         {
-            flags |= QUIC_STREAM_SHUTDOWN_FLAGS.ABORT_RECEIVE;
             if (_receiveTcs.TrySetException(ThrowHelper.GetOperationAbortedException(SR.net_quic_reading_aborted), final: true))
             {
                 flags |= QUIC_STREAM_SHUTDOWN_FLAGS.ABORT_RECEIVE;
@@ -537,12 +536,12 @@ private unsafe int HandleEventShutdownComplete(ref SHUTDOWN_COMPLETE data)
                 (shutdownByApp: true, closedRemotely: true) => ThrowHelper.GetConnectionAbortedException((long)data.ConnectionErrorCode),
                 // It's local shutdown by app, this side called QuicConnection.CloseAsync, throw QuicError.OperationAborted.
                 (shutdownByApp: true, closedRemotely: false) => ThrowHelper.GetOperationAbortedException(),
-                // It's remote shutdown by transport, we received a CONNECTION_CLOSE frame with a QUIC transport error code
+                // It's remote shutdown by transport, we received a CONNECTION_CLOSE frame with a QUIC transport error code, throw error based on the status.
                 // TODO: we should propagate the transport error code
                 // https://github.com/dotnet/runtime/issues/72666
                 (shutdownByApp: false, closedRemotely: true) => ThrowHelper.GetExceptionForMsQuicStatus(data.ConnectionCloseStatus, $"Shutdown by transport {data.ConnectionErrorCode}"),
-                // It's local shutdown by transport, due to some timeout
-                // TODO: we should propagate transport error code
+                // It's local shutdown by transport, most likely due to a timeout, throw error based on the status.
+                // TODO: we should propagate the transport error code
                 // https://github.com/dotnet/runtime/issues/72666
                 (shutdownByApp: false, closedRemotely: false) => ThrowHelper.GetExceptionForMsQuicStatus(data.ConnectionCloseStatus),
             };
diff --git a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs b/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)));
+        }
     }
 }
diff --git a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicListenerTests.cs b/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)));
+        }
     }
 }
diff --git a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicStreamConnectedStreamConformanceTests.cs b/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;
-        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,6 @@`
`7`	`7`	`using System.Runtime.Versioning;`
`8`	`8`	`using System.Net.Quic;`
`9`	`9`	`using System.IO;`
`10`		`-using System.Linq;`
`11`	`10`	`using System.Collections.Generic;`
`12`	`11`	`using System.Diagnostics;`
`13`	`12`	`using System.Globalization;`
`@@ -155,7 +154,7 @@ private void CheckForShutdown()`
`155`	`154`
`156`	`155`	`if (_clientControl != null)`
`157`	`156`	`{`
`158`		`- _clientControl.Dispose();`
	`157`	`+ await _clientControl.DisposeAsync().ConfigureAwait(false);`
`159`	`158`	`_clientControl = null;`
`160`	`159`	`}`
`161`	`160`
`@@ -245,7 +244,10 @@ public async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, lon`
`245`	`244`	`}`
`246`	`245`	`finally`
`247`	`246`	`{`
`248`		`- requestStream?.Dispose();`
	`247`	`+ if (requestStream is not null)`
	`248`	`+ {`
	`249`	`+ await requestStream.DisposeAsync().ConfigureAwait(false);`
	`250`	`+ }`
`249`	`251`	`}`
`250`	`252`	`}`
`251`	`253`
`@@ -562,7 +564,6 @@ private async Task ProcessServerStreamAsync(QuicStream stream)`
`562`	`564`	`}`
`563`	`565`
`564`	`566`	`stream.Abort(QuicAbortDirection.Read, (long)Http3ErrorCode.StreamCreationError);`
`565`		`- stream.Dispose();`
`566`	`567`	`return;`
`567`	`568`	`}`
`568`	`569`	`}`
Original file line number	Diff line number	Diff line change
`@@ -392,6 +392,7 @@ public async ValueTask<QuicStream> AcceptInboundStreamAsync(CancellationToken ca`
`392`	`392`	`throw new InvalidOperationException(SR.net_quic_accept_not_allowed);`
`393`	`393`	`}`
`394`	`394`
	`395`	`+ GCHandle keepObject = GCHandle.Alloc(this);`
`395`	`396`	`try`
`396`	`397`	`{`
`397`	`398`	`return await _acceptQueue.Reader.ReadAsync(cancellationToken).ConfigureAwait(false);`
`@@ -401,6 +402,10 @@ public async ValueTask<QuicStream> AcceptInboundStreamAsync(CancellationToken ca`
`401`	`402`	`ExceptionDispatchInfo.Capture(ex.InnerException).Throw();`
`402`	`403`	`throw;`
`403`	`404`	`}`
	`405`	`+ finally`
	`406`	`+ {`
	`407`	`+ keepObject.Free();`
	`408`	`+ }`
`404`	`409`	`}`
`405`	`410`
`406`	`411`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -162,6 +162,7 @@ public async ValueTask<QuicConnection> AcceptConnectionAsync(CancellationToken c`
`162`	`162`	`{`
`163`	`163`	`ObjectDisposedException.ThrowIf(_disposed == 1, this);`
`164`	`164`
	`165`	`+ GCHandle keepObject = GCHandle.Alloc(this);`
`165`	`166`	`try`
`166`	`167`	`{`
`167`	`168`	`PendingConnection pendingConnection = await _acceptQueue.Reader.ReadAsync(cancellationToken).ConfigureAwait(false);`
`@@ -175,6 +176,10 @@ public async ValueTask<QuicConnection> AcceptConnectionAsync(CancellationToken c`
`175`	`176`	`ExceptionDispatchInfo.Capture(ex.InnerException).Throw();`
`176`	`177`	`throw;`
`177`	`178`	`}`
	`179`	`+ finally`
	`180`	`+ {`
	`181`	`+ keepObject.Free();`
	`182`	`+ }`
`178`	`183`	`}`
`179`	`184`
`180`	`185`	`private unsafe int HandleEventNewConnection(ref NEW_CONNECTION_DATA data)`
Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ await WhenAllOrAnyFailed(`
`102`	`102`	`}`
`103`	`103`	`catch (Exception ex)`
`104`	`104`	`{`
`105`		`- _output?.WriteLine($"Failed to {ex.Message}");`
	`105`	`+ _output?.WriteLine($"Failed to connect: {ex.Message}");`
`106`	`106`	`throw;`
`107`	`107`	`}`
`108`	`108`	`}));`
`@@ -153,14 +153,5 @@ public override void Dispose()`
`153`	`153`	`}`
`154`	`154`	`}`
`155`	`155`	`}`
`156`		`-`
`157`		`- [OuterLoop("May take several seconds")]`
`158`		`- [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]`
`159`		`- [SkipOnPlatform(TestPlatforms.LinuxBionic, "SElinux blocks UNIX sockets in our CI environment")]`
`160`		`- [ActiveIssue("https://github.com/dotnet/runtime/issues/73377")]`
`161`		`- public override Task Parallel_ReadWriteMultipleStreamsConcurrently()`
`162`		`- {`
`163`		`- return Task.CompletedTask;`
`164`		`- }`
`165`	`156`	`}`
`166`	`157`	`}`