keep MsQuicConnection alive when streams are pending

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicConnection.cs

@@ -54,7 +54,7 @@ internal sealed class State
             public bool Connected;
             public long AbortErrorCode = -1;
             public int StreamCount;
-            public bool Closing;
+            private IntPtr _closingGCHandle;
 
             // Queue for accepted streams.
             // Backlog limit is managed by MsQuic so it can be unbounded here.
@@ -66,23 +66,34 @@ internal sealed class State
 
             public void RemoveStream(MsQuicStream stream)
             {
+                bool closing;
                 lock (this)
                 {
                     StreamCount--;
+                    closing = _closingGCHandle != IntPtr.Zero && StreamCount == 0;
                 }
 
-                if (Closing && StreamCount == 0)
+                if (closing)
                 {
                     Handle?.Dispose();
-                    //if (_stateHandle.IsAllocated) _stateHandle.Free();
+                    try
+                    {
+                        GCHandle gcHandle = GCHandle.FromIntPtr(_closingGCHandle);
+                        Debug.Assert(gcHandle.IsAllocated);
+                        if (gcHandle.IsAllocated) gcHandle.Free();
+                    }
+                    catch (Exception ex)
+                    {
+                        if (NetEventSource.Log.IsEnabled()) NetEventSource.Error(this, $"Failed to restore GCHandle from ptr: {ex.Message}");
+                    }
                 }
             }
 
             public bool TryAddStream(SafeMsQuicStreamHandle streamHandle, QUIC_STREAM_OPEN_FLAGS flags)
             {
                 lock (this)
                 {
-                    if (Closing)
+                    if (_closingGCHandle != IntPtr.Zero)
                     {
                         return false;
                     }
@@ -101,6 +112,12 @@ public bool TryAddStream(SafeMsQuicStreamHandle streamHandle, QUIC_STREAM_OPEN_F
                     }
                 }
             }
+
+            // This is called under lock from connection dispose
+            public void SetClosing(GCHandle handle)
+            {
+                _closingGCHandle = GCHandle.ToIntPtr(handle);
+            }
         }
 
         // constructor for inbound connections
@@ -568,15 +585,23 @@ private void Dispose(bool disposing)
             lock (_state)
             {
                 _state.Connection = null;
-                _state.Closing = true;
                 if (_state.StreamCount == 0)
                 {
-                    _state?.Handle?.Dispose();
+                    _state!.Handle?.Dispose();
+                    if (_stateHandle.IsAllocated) _stateHandle.Free();
+                }
+                else
+                {
+                    // normally, _state would be rooted by 'this' and we would hold _stateHandle to prevent
+                    // GC from moving _state as we handed pointer to it by msquic. It was handed to msquic and
+                    // we may try to get _state from it in NativeCallbackHandler()
+                    // At this point we are being disposed either explicitly or by finalizer but we still have active streams.
+                    // To prevent issue, the handle will be transferred to state and it will be released when last stream is gone
+                    _state.SetClosing(_stateHandle);
                 }
             }
 
             FlushAcceptQueue().GetAwaiter().GetResult();
-            if (_stateHandle.IsAllocated) _stateHandle.Free();
         }
 
         // TODO: this appears abortive and will cause prior successfully shutdown and closed streams to drop data.

src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicStream.cs

@@ -554,6 +554,7 @@ private void Dispose(bool disposing)
             Marshal.FreeHGlobal(_state.SendQuicBuffers);
             if (_stateHandle.IsAllocated) _stateHandle.Free();
             CleanupSendState(_state);
+            Debug.Assert(_state.ConnectionState != null);
             _state.ConnectionState?.RemoveStream(this);
 
             if (NetEventSource.Log.IsEnabled())