Avoid repetitive GCHandles in IIS WebSocket operations (#48462)

BrennanConroy · web-flow · commit d868553a675f · 2023-06-08T13:28:10.000-07:00
diff --git a/src/Servers/IIS/IIS/src/Core/IISHttpContext.cs b/src/Servers/IIS/IIS/src/Core/IISHttpContext.cs
@@ -763,6 +763,8 @@ protected virtual void Dispose(bool disposing)
             localAbortCts?.Dispose();
 
             disposedValue = true;
+
+            AsyncIO?.Dispose();
         }
     }
 
diff --git a/src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.cs b/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;
+    }
 }
diff --git a/src/Servers/IIS/IIS/src/Core/IO/AsyncIOOperation.cs b/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; }
diff --git a/src/Servers/IIS/IIS/src/Core/IO/IAsyncIOEngine.cs b/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);
diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Initialize.cs b/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);
         }
     }
 }
diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs b/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();
+        }
     }
 }
diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs b/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);
         }
     }
 }
diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs b/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();
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -763,6 +763,8 @@ protected virtual void Dispose(bool disposing)`
`763`	`763`	`localAbortCts?.Dispose();`
`764`	`764`
`765`	`765`	`disposedValue = true;`
	`766`	`+`
	`767`	`+ AsyncIO?.Dispose();`
`766`	`768`	`}`
`767`	`769`	`}`
`768`	`770`
Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@`
`6`	`6`
`7`	`7`	`namespace Microsoft.AspNetCore.Server.IIS.Core.IO;`
`8`	`8`
`9`		`-internal sealed partial class AsyncIOEngine : IAsyncIOEngine`
	`9`	`+internal sealed partial class AsyncIOEngine : IAsyncIOEngine, IDisposable`
`10`	`10`	`{`
`11`	`11`	`private const ushort ResponseMaxChunks = 65533;`
`12`	`12`
`@@ -244,4 +244,9 @@ private void ReturnOperation(AsyncFlushOperation operation)`
`244`	`244`	`{`
`245`	`245`	`Volatile.Write(ref _cachedAsyncFlushOperation, operation);`
`246`	`246`	`}`
	`247`	`+`
	`248`	`+ public void Dispose()`
	`249`	`+ {`
	`250`	`+ _stopped = true;`
	`251`	`+ }`
`247`	`252`	`}`
Original file line number	Diff line number	Diff line change
`@@ -134,6 +134,8 @@ protected virtual void ResetOperation()`
`134`	`134`	`_continuation = null;`
`135`	`135`	`}`
`136`	`136`
	`137`	`+ public bool InUse() => _continuation is not null;`
	`138`	`+`
`137`	`139`	`public readonly struct AsyncContinuation`
`138`	`140`	`{`
`139`	`141`	`public Action<object?> Continuation { get; }`
Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@`
`5`	`5`
`6`	`6`	`namespace Microsoft.AspNetCore.Server.IIS.Core.IO;`
`7`	`7`
`8`		`-internal interface IAsyncIOEngine`
	`8`	`+internal interface IAsyncIOEngine : IDisposable`
`9`	`9`	`{`
`10`	`10`	`ValueTask<int> ReadAsync(Memory<byte> memory);`
`11`	`11`	`ValueTask<int> WriteAsync(ReadOnlySequence<byte> data);`
Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,6 @@ protected override void ResetOperation()`
`36`	`36`	`base.ResetOperation();`
`37`	`37`
`38`	`38`	`_requestHandler = default;`
`39`		`- _engine.ReturnOperation(this);`
`40`	`39`	`}`
`41`	`40`	`}`
`42`	`41`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO;`
`9`	`9`
`10`	`10`	`internal partial class WebSocketsAsyncIOEngine`
`11`	`11`	`{`
`12`		`- internal sealed class WebSocketReadOperation : AsyncIOOperation`
	`12`	`+ internal sealed class WebSocketReadOperation : AsyncIOOperation, IDisposable`
`13`	`13`	`{`
`14`	`14`	`[UnmanagedCallersOnly]`
`15`	`15`	`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`
`26`	`26`	`}`
`27`	`27`
`28`	`28`	`private readonly WebSocketsAsyncIOEngine _engine;`
`29`		`- private GCHandle _thisHandle;`
	`29`	`+ private readonly GCHandle _thisHandle;`
`30`	`30`	`private MemoryHandle _inputHandle;`
`31`	`31`	`private NativeSafeHandle? _requestHandler;`
`32`	`32`	`private Memory<byte> _memory;`
`33`	`33`
`34`	`34`	`public WebSocketReadOperation(WebSocketsAsyncIOEngine engine)`
`35`	`35`	`{`
	`36`	`+ _thisHandle = GCHandle.Alloc(this);`
`36`	`37`	`_engine = engine;`
`37`	`38`	`}`
`38`	`39`
`39`	`40`	`protected override unsafe bool InvokeOperation(out int hr, out int bytes)`
`40`	`41`	`{`
`41`	`42`	`Debug.Assert(_requestHandler != null, "Must initialize first.");`
`42`	`43`
`43`		`- _thisHandle = GCHandle.Alloc(this);`
`44`	`44`	`_inputHandle = _memory.Pin();`
`45`	`45`
`46`	`46`	`hr = NativeMethods.HttpWebsocketsReadBytes(`
`@@ -70,16 +70,17 @@ protected override void ResetOperation()`
`70`	`70`	`{`
`71`	`71`	`base.ResetOperation();`
`72`	`72`
`73`		`- _thisHandle.Free();`
`74`		`-`
`75`	`73`	`_memory = default;`
`76`	`74`	`_inputHandle.Dispose();`
`77`	`75`	`_inputHandle = default;`
`78`	`76`	`_requestHandler = default;`
`79`		`-`
`80`		`- _engine.ReturnOperation(this);`
`81`	`77`	`}`
`82`	`78`
`83`	`79`	`protected override bool IsSuccessfulResult(int hr) => hr == NativeMethods.ERROR_HANDLE_EOF;`
	`80`	`+`
	`81`	`+ public void Dispose()`
	`82`	`+ {`
	`83`	`+ _thisHandle.Free();`
	`84`	`+ }`
`84`	`85`	`}`
`85`	`86`	`}`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO;`
`8`	`8`
`9`	`9`	`internal partial class WebSocketsAsyncIOEngine`
`10`	`10`	`{`
`11`		`- internal sealed class WebSocketWriteOperation : AsyncWriteOperationBase`
	`11`	`+ internal sealed class WebSocketWriteOperation : AsyncWriteOperationBase, IDisposable`
`12`	`12`	`{`
`13`	`13`	`[UnmanagedCallersOnly]`
`14`	`14`	`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`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`private readonly WebSocketsAsyncIOEngine _engine;`
`27`		`- private GCHandle _thisHandle;`
	`27`	`+ private readonly GCHandle _thisHandle;`
`28`	`28`
`29`	`29`	`public WebSocketWriteOperation(WebSocketsAsyncIOEngine engine)`
`30`	`30`	`{`
	`31`	`+ _thisHandle = GCHandle.Alloc(this);`
`31`	`32`	`_engine = engine;`
`32`	`33`	`}`
`33`	`34`
`34`	`35`	`protected override unsafe int WriteChunks(NativeSafeHandle requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks, out bool completionExpected)`
`35`	`36`	`{`
`36`		`- _thisHandle = GCHandle.Alloc(this);`
`37`	`37`	`return NativeMethods.HttpWebsocketsWriteBytes(requestHandler, dataChunks, chunkCount, &WriteCallback, (IntPtr)_thisHandle, out completionExpected);`
`38`	`38`	`}`
`39`	`39`
`40`	`40`	`protected override void ResetOperation()`
`41`	`41`	`{`
`42`	`42`	`base.ResetOperation();`
	`43`	`+ }`
`43`	`44`
	`45`	`+ public void Dispose()`
	`46`	`+ {`
`44`	`47`	`_thisHandle.Free();`
`45`		`-`
`46`		`- _engine.ReturnOperation(this);`
`47`	`48`	`}`
`48`	`49`	`}`
`49`	`50`	`}`