Avoid repetitive GCHandles in IIS WebSocket operations

src/Servers/IIS/IIS/src/Core/IISHttpContext.cs

@@ -763,6 +763,8 @@ protected virtual void Dispose(bool disposing)
             localAbortCts?.Dispose();
 
             disposedValue = true;
+
+            AsyncIO?.Dispose();
         }
     }
 

src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.cs

@@ -6,7 +6,7 @@
 
 namespace Microsoft.AspNetCore.Server.IIS.Core.IO;
 
-internal sealed partial class AsyncIOEngine : IAsyncIOEngine
+internal sealed partial class AsyncIOEngine : IAsyncIOEngine, IDisposable
 {
     private const ushort ResponseMaxChunks = 65533;
 
@@ -244,4 +244,9 @@ private void ReturnOperation(AsyncFlushOperation operation)
     {
         Volatile.Write(ref _cachedAsyncFlushOperation, operation);
     }
+
+    public void Dispose()
+    {
+        _stopped = true;
+    }
 }

src/Servers/IIS/IIS/src/Core/IO/IAsyncIOEngine.cs

@@ -5,7 +5,7 @@
 
 namespace Microsoft.AspNetCore.Server.IIS.Core.IO;
 
-internal interface IAsyncIOEngine
+internal interface IAsyncIOEngine : IDisposable
 {
     ValueTask<int> ReadAsync(Memory<byte> memory);
     ValueTask<int> WriteAsync(ReadOnlySequence<byte> data);

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs

@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO;
 
 internal partial class WebSocketsAsyncIOEngine
 {
-    internal sealed class WebSocketReadOperation : AsyncIOOperation
+    internal sealed class WebSocketReadOperation : AsyncIOOperation, IDisposable
     {
         [UnmanagedCallersOnly]
         public static NativeMethods.REQUEST_NOTIFICATION_STATUS ReadCallback(IntPtr httpContext, IntPtr completionInfo, IntPtr completionContext)
@@ -26,21 +26,21 @@ public static NativeMethods.REQUEST_NOTIFICATION_STATUS ReadCallback(IntPtr http
         }
 
         private readonly WebSocketsAsyncIOEngine _engine;
-        private GCHandle _thisHandle;
+        private readonly GCHandle _thisHandle;
         private MemoryHandle _inputHandle;
         private NativeSafeHandle? _requestHandler;
         private Memory<byte> _memory;
 
         public WebSocketReadOperation(WebSocketsAsyncIOEngine engine)
         {
+            _thisHandle = GCHandle.Alloc(this);
             _engine = engine;
         }
 
         protected override unsafe bool InvokeOperation(out int hr, out int bytes)
         {
             Debug.Assert(_requestHandler != null, "Must initialize first.");
 
-            _thisHandle = GCHandle.Alloc(this);
             _inputHandle = _memory.Pin();
 
             hr = NativeMethods.HttpWebsocketsReadBytes(
@@ -70,8 +70,6 @@ protected override void ResetOperation()
         {
             base.ResetOperation();
 
-            _thisHandle.Free();
-
             _memory = default;
             _inputHandle.Dispose();
             _inputHandle = default;
@@ -81,5 +79,10 @@ protected override void ResetOperation()
         }
 
         protected override bool IsSuccessfulResult(int hr) => hr == NativeMethods.ERROR_HANDLE_EOF;
+
+        public void Dispose()
+        {
+            _thisHandle.Free();
+        }
     }
 }

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs

@@ -8,7 +8,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO;
 
 internal partial class WebSocketsAsyncIOEngine
 {
-    internal sealed class WebSocketWriteOperation : AsyncWriteOperationBase
+    internal sealed class WebSocketWriteOperation : AsyncWriteOperationBase, IDisposable
     {
         [UnmanagedCallersOnly]
         private static NativeMethods.REQUEST_NOTIFICATION_STATUS WriteCallback(IntPtr httpContext, IntPtr completionInfo, IntPtr completionContext)
@@ -24,26 +24,29 @@ private static NativeMethods.REQUEST_NOTIFICATION_STATUS WriteCallback(IntPtr ht
         }
 
         private readonly WebSocketsAsyncIOEngine _engine;
-        private GCHandle _thisHandle;
+        private readonly GCHandle _thisHandle;
 
         public WebSocketWriteOperation(WebSocketsAsyncIOEngine engine)
         {
+            _thisHandle = GCHandle.Alloc(this);
             _engine = engine;
         }
 
         protected override unsafe int WriteChunks(NativeSafeHandle requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks, out bool completionExpected)
         {
-            _thisHandle = GCHandle.Alloc(this);
             return NativeMethods.HttpWebsocketsWriteBytes(requestHandler, dataChunks, chunkCount, &WriteCallback, (IntPtr)_thisHandle, out completionExpected);
         }
 
         protected override void ResetOperation()
         {
             base.ResetOperation();
 
-            _thisHandle.Free();
-
             _engine.ReturnOperation(this);
         }
+
+        public void Dispose()
+        {
+            _thisHandle.Free();
+        }
     }
 }

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs

@@ -139,4 +139,10 @@ private void ReturnOperation(WebSocketWriteOperation operation) =>
 
     private void ReturnOperation(WebSocketReadOperation operation) =>
         Volatile.Write(ref _cachedWebSocketReadOperation, operation);
+
+    public void Dispose()
+    {
+        _cachedWebSocketWriteOperation?.Dispose();
+        _cachedWebSocketReadOperation?.Dispose();
+    }
 }