From 40383cf1f9417b7018b51d02b72afdf1cfa1705b Mon Sep 17 00:00:00 2001
From: Brennan Conroy <brecon@microsoft.com>
Date: Mon, 20 Jul 2020 19:11:17 +0000
Subject: [PATCH] Merged PR 9189: Modify WebSockets ValueTaskSource

Create the GCHandle per operation and Free it when resetting the ValueTaskSource for re-use.
---
 .../IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs     | 6 ++++--
 .../IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs    | 6 ++++--
 src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs  | 2 +-
 3 files changed, 9 insertions(+), 5 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);
             }
         }
diff --git a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs b/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs
index a796af23b3..2c1ac02e20 100644
--- a/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs
+++ b/src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs
@@ -39,7 +39,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 var read = GetReadOperation();
                 read.Initialize(_handler, memory);
                 read.Invoke();
-                return new ValueTask<int>(read, 0);   
+                return new ValueTask<int>(read, 0);
             }
         }