diff --git a/src/Servers/IIS/IIS/src/Core/IISHttpContext.cs b/src/Servers/IIS/IIS/src/Core/IISHttpContext.cs index 504f9ae2339f..a8a7c034d5c0 100644 --- 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 index c480a1d1bb01..92735d8cfbb7 100644 --- 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 index 8aca9e770f1b..ee7a74b8c72c 100644 --- 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 Continuation { get; } diff --git a/src/Servers/IIS/IIS/src/Core/IO/IAsyncIOEngine.cs b/src/Servers/IIS/IIS/src/Core/IO/IAsyncIOEngine.cs index 78cbd0a095cd..43034f7119da 100644 --- 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 ReadAsync(Memory memory); ValueTask WriteAsync(ReadOnlySequence 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 index 10018a0512aa..347a290220c3 100644 --- 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 index bc29203b5264..ea78fcf3dac7 100644 --- 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,13 +26,14 @@ 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 _memory; public WebSocketReadOperation(WebSocketsAsyncIOEngine engine) { + _thisHandle = GCHandle.Alloc(this); _engine = engine; } @@ -40,7 +41,6 @@ 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 index 6b7d1bf8b3ec..d38c541d5960 100644 --- 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 index 7748802be688..3c38a4bca06a 100644 --- 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 ReadAsync(Memory memory) { ThrowIfNotInitialized(); - var read = GetReadOperation(); + var read = _webSocketReadOperation ??= new WebSocketReadOperation(this); + + Debug.Assert(!read.InUse()); + read.Initialize(_handler, memory); read.Invoke(); return new ValueTask(read, 0); @@ -46,7 +48,10 @@ public ValueTask WriteAsync(ReadOnlySequence data) { ThrowIfNotInitialized(); - var write = GetWriteOperation(); + var write = _webSocketWriteOperation ??= new WebSocketWriteOperation(this); + + Debug.Assert(!write.InUse()); + write.Initialize(_handler, data); write.Invoke(); return new ValueTask(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(); + } }