Skip to content

Avoid repetitive GCHandles in IIS WebSocket operations #48462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 8, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/Servers/IIS/IIS/src/Core/IISHttpContext.cs
Original file line number Diff line number Diff line change
Expand Up @@ -763,6 +763,8 @@ protected virtual void Dispose(bool disposing)
localAbortCts?.Dispose();

disposedValue = true;

AsyncIO?.Dispose();
}
}

Expand Down
7 changes: 6 additions & 1 deletion src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -244,4 +244,9 @@ private void ReturnOperation(AsyncFlushOperation operation)
{
Volatile.Write(ref _cachedAsyncFlushOperation, operation);
}

public void Dispose()
{
_stopped = true;
}
}
2 changes: 2 additions & 0 deletions src/Servers/IIS/IIS/src/Core/IO/AsyncIOOperation.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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; }
Expand Down
2 changes: 1 addition & 1 deletion src/Servers/IIS/IIS/src/Core/IO/IAsyncIOEngine.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ protected override void ResetOperation()
base.ResetOperation();

_requestHandler = default;
_engine.ReturnOperation(this);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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(
Expand Down Expand Up @@ -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();
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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);
}
}
}
44 changes: 17 additions & 27 deletions src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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)
{
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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();
Expand Down Expand Up @@ -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();
}
}