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/AsyncIOOperation.cs

@@ -134,6 +134,8 @@ protected virtual void ResetOperation()
         _continuation = null;
     }
 
+    public bool InUse() => _continuation is not null;
+
     public readonly struct AsyncContinuation
     {
         public Action<object?> Continuation { get; }

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.Initialize.cs

@@ -36,7 +36,6 @@ protected override void ResetOperation()
             base.ResetOperation();
 
             _requestHandler = default;
-            _engine.ReturnOperation(this);
         }
     }
 }

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,16 +70,17 @@ protected override void ResetOperation()
         {
             base.ResetOperation();
 
-            _thisHandle.Free();
-
             _memory = default;
             _inputHandle.Dispose();
             _inputHandle = default;
             _requestHandler = default;
-
-            _engine.ReturnOperation(this);
         }
 
         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,27 @@ 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();
+        }
 
+        public void Dispose()
+        {
             _thisHandle.Free();
-
-            _engine.ReturnOperation(this);
         }
     }
 }

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Buffers;
+using System.Diagnostics;
 
 namespace Microsoft.AspNetCore.Server.IIS.Core.IO;
 
@@ -15,11 +16,9 @@ internal sealed partial class WebSocketsAsyncIOEngine : IAsyncIOEngine
 
     private AsyncInitializeOperation? _initializationFlush;
 
-    private WebSocketWriteOperation? _cachedWebSocketWriteOperation;
+    private WebSocketWriteOperation? _webSocketWriteOperation;
 
-    private WebSocketReadOperation? _cachedWebSocketReadOperation;
-
-    private AsyncInitializeOperation? _cachedAsyncInitializeOperation;
+    private WebSocketReadOperation? _webSocketReadOperation;
 
     public WebSocketsAsyncIOEngine(IISHttpContext context, NativeSafeHandle handler)
     {
@@ -33,7 +32,10 @@ public ValueTask<int> ReadAsync(Memory<byte> memory)
         {
             ThrowIfNotInitialized();
 
-            var read = GetReadOperation();
+            var read = _webSocketReadOperation ??= new WebSocketReadOperation(this);
+
+            Debug.Assert(!read.InUse());
+
             read.Initialize(_handler, memory);
             read.Invoke();
             return new ValueTask<int>(read, 0);
@@ -46,7 +48,10 @@ public ValueTask<int> WriteAsync(ReadOnlySequence<byte> data)
         {
             ThrowIfNotInitialized();
 
-            var write = GetWriteOperation();
+            var write = _webSocketWriteOperation ??= new WebSocketWriteOperation(this);
+
+            Debug.Assert(!write.InUse());
+
             write.Initialize(_handler, data);
             write.Invoke();
             return new ValueTask<int>(write, 0);
@@ -64,7 +69,7 @@ public ValueTask FlushAsync(bool moreData)
 
             NativeMethods.HttpEnableWebsockets(_handler);
 
-            var init = GetInitializeOperation();
+            var init = new AsyncInitializeOperation(this);
             init.Initialize(_handler);
 
             var continuation = init.Invoke();
@@ -119,24 +124,9 @@ public void Complete()
         }
     }
 
-    private WebSocketReadOperation GetReadOperation() =>
-        Interlocked.Exchange(ref _cachedWebSocketReadOperation, null) ??
-        new WebSocketReadOperation(this);
-
-    private WebSocketWriteOperation GetWriteOperation() =>
-        Interlocked.Exchange(ref _cachedWebSocketWriteOperation, null) ??
-        new WebSocketWriteOperation(this);
-
-    private AsyncInitializeOperation GetInitializeOperation() =>
-        Interlocked.Exchange(ref _cachedAsyncInitializeOperation, null) ??
-        new AsyncInitializeOperation(this);
-
-    private void ReturnOperation(AsyncInitializeOperation operation) =>
-        Volatile.Write(ref _cachedAsyncInitializeOperation, operation);
-
-    private void ReturnOperation(WebSocketWriteOperation operation) =>
-        Volatile.Write(ref _cachedWebSocketWriteOperation, operation);
-
-    private void ReturnOperation(WebSocketReadOperation operation) =>
-        Volatile.Write(ref _cachedWebSocketReadOperation, operation);
+    public void Dispose()
+    {
+        _webSocketWriteOperation?.Dispose();
+        _webSocketReadOperation?.Dispose();
+    }
 }