Merged PR 9194: Modify WebSockets ValueTaskSourcev2.2.13

Modify WebSockets ValueTaskSource

2 files changed, 8 insertions, 4 deletions

diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs b/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs
index 2dac1a234e..db4068c52c 100644
--- a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs
+++ b/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs
@@ -25,7 +25,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             };
 
             private readonly WebSocketsAsyncIOEngine _engine;
-            private readonly GCHandle _thisHandle;
+            private GCHandle _thisHandle;
             private MemoryHandle _inputHandle;
             private IntPtr _requestHandler;
             private Memory<byte> _memory;
@@ -33,11 +33,11 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             public WebSocketReadOperation(WebSocketsAsyncIOEngine engine)
             {
                 _engine = engine;
-                _thisHandle = GCHandle.Alloc(this);
             }
 
             protected override unsafe bool InvokeOperation(out int hr, out int bytes)
             {
+                _thisHandle = GCHandle.Alloc(this);
                 _inputHandle = _memory.Pin();
 
                 hr = NativeMethods.HttpWebsocketsReadBytes(
@@ -67,6 +67,8 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             {
                 base.ResetOperation();
 
+                _thisHandle.Free();
+
                 _memory = default;
                 _inputHandle.Dispose();
                 _inputHandle = default;
diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs b/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs
index 3eff3bba46..5e6630aae5 100644
--- a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs
+++ b/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs
@@ -25,16 +25,16 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             };
 
             private readonly WebSocketsAsyncIOEngine _engine;
-            private readonly GCHandle _thisHandle;
+            private GCHandle _thisHandle;
 
             public WebSocketWriteOperation(WebSocketsAsyncIOEngine engine)
             {
                 _engine = engine;
-                _thisHandle = GCHandle.Alloc(this);
             }
 
             protected override unsafe int WriteChunks(IntPtr 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);
             }
 
@@ -42,6 +42,8 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             {
                 base.ResetOperation();
 
+                _thisHandle.Free();
+
                 _engine.ReturnOperation(this);
             }
         }