Skip to content

Commit

Permalink
Avoid allocations in ClientWebSocket.ConnectAsync in the common case (#…
Browse files Browse the repository at this point in the history
…75025)

* Avoid allocations in ClientWebSocket.ConnectAsync in the common case

* Remove outdated ActiveIssue

* Update comment typo
  • Loading branch information
MihaZupan committed Sep 5, 2022
1 parent 8fe12bc commit db4a954
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 27 deletions.
Expand Up @@ -18,8 +18,9 @@ namespace System.Net.WebSockets
{
internal sealed class WebSocketHandle
{
/// <summary>Shared, lazily-initialized handler for when using default options.</summary>
private static SocketsHttpHandler? s_defaultHandler;
// Shared, lazily-initialized invokers used to avoid some allocations when using default options.
private static HttpMessageInvoker? s_defaultInvokerDefaultProxy;
private static HttpMessageInvoker? s_defaultInvokerNoProxy;

private readonly CancellationTokenSource _abortSource = new CancellationTokenSource();
private WebSocketState _state = WebSocketState.Connecting;
Expand Down Expand Up @@ -47,15 +48,15 @@ public void Abort()

public async Task ConnectAsync(Uri uri, HttpMessageInvoker? invoker, CancellationToken cancellationToken, ClientWebSocketOptions options)
{
bool disposeHandler = false;
bool disposeInvoker = false;
if (invoker is null)
{
if (options.HttpVersion.Major >= 2 || options.HttpVersionPolicy == HttpVersionPolicy.RequestVersionOrHigher)
{
throw new ArgumentException(SR.net_WebSockets_CustomInvokerRequiredForHttp2, nameof(options));
}

invoker = new HttpMessageInvoker(SetupHandler(options, out disposeHandler));
invoker = SetupInvoker(options, out disposeInvoker);
}
else if (!options.AreCompatibleWithCustomInvoker())
{
Expand Down Expand Up @@ -235,44 +236,47 @@ public async Task ConnectAsync(Uri uri, HttpMessageInvoker? invoker, Cancellatio
}
}

// Disposing the handler will not affect any active stream wrapped in the WebSocket.
if (disposeHandler)
// Disposing the invoker will not affect any active stream wrapped in the WebSocket.
if (disposeInvoker)
{
invoker?.Dispose();
}
}
}

private static SocketsHttpHandler SetupHandler(ClientWebSocketOptions options, out bool disposeHandler)
private static HttpMessageInvoker SetupInvoker(ClientWebSocketOptions options, out bool disposeInvoker)
{
SocketsHttpHandler? handler;

// Create the handler for this request and populate it with all of the options.
// Try to use a shared handler rather than creating a new one just for this request, if
// the options are compatible.
if (options.AreCompatibleWithCustomInvoker() && options.Proxy is null)
// Create the invoker for this request and populate it with all of the options.
// If the options are compatible, reuse a shared invoker.
if (options.AreCompatibleWithCustomInvoker())
{
disposeHandler = false;
handler = s_defaultHandler;
if (handler == null)
disposeInvoker = false;

bool useDefaultProxy = options.Proxy is not null;

ref HttpMessageInvoker? invokerRef = ref useDefaultProxy ? ref s_defaultInvokerDefaultProxy : ref s_defaultInvokerNoProxy;

if (invokerRef is null)
{
handler = new SocketsHttpHandler()
var invoker = new HttpMessageInvoker(new SocketsHttpHandler()
{
PooledConnectionLifetime = TimeSpan.Zero,
UseProxy = false,
UseProxy = useDefaultProxy,
UseCookies = false,
};
if (Interlocked.CompareExchange(ref s_defaultHandler, handler, null) != null)
});

if (Interlocked.CompareExchange(ref invokerRef, invoker, null) is not null)
{
handler.Dispose();
handler = s_defaultHandler;
invoker.Dispose();
}
}

return invokerRef;
}
else
{
disposeHandler = true;
handler = new SocketsHttpHandler();
disposeInvoker = true;
var handler = new SocketsHttpHandler();
handler.PooledConnectionLifetime = TimeSpan.Zero;
handler.CookieContainer = options.Cookies;
handler.UseCookies = options.Cookies != null;
Expand All @@ -297,9 +301,9 @@ private static SocketsHttpHandler SetupHandler(ClientWebSocketOptions options, o
handler.SslOptions.ClientCertificates = new X509Certificate2Collection();
handler.SslOptions.ClientCertificates.AddRange(options.ClientCertificates);
}
}

return handler;
return new HttpMessageInvoker(handler);
}
}

private static WebSocketDeflateOptions ParseDeflateOptions(ReadOnlySpan<char> extension, WebSocketDeflateOptions original)
Expand Down
Expand Up @@ -48,7 +48,6 @@ public static void Proxy_Roundtrips()

[OuterLoop]
[ConditionalTheory(nameof(WebSocketsSupported)), MemberData(nameof(EchoServers))]
[ActiveIssue("https://github.com/dotnet/runtime/issues/43751")]
public async Task Proxy_SetNull_ConnectsSuccessfully(Uri server)
{
for (int i = 0; i < 3; i++) // Connect and disconnect multiple times to exercise shared handler on netcoreapp
Expand Down

0 comments on commit db4a954

Please sign in to comment.