Skip to content

Commit

Permalink
Switch to ValueTask for async methods
Browse files Browse the repository at this point in the history
We only have CompleteAsync and RunAsync on IWebSocketPipe, but since the underlying primitives we invoke on either the pipe reader/writer or websocket, all return ValueTask to avoid allocations, we might as well follow suit.

Background: https://devblogs.microsoft.com/dotnet/understanding-the-whys-whats-and-whens-of-valuetask/

Fixes #3
  • Loading branch information
kzu committed Sep 22, 2021
1 parent 1f20632 commit 269dfdc
Show file tree
Hide file tree
Showing 4 changed files with 17 additions and 13 deletions.
8 changes: 5 additions & 3 deletions src/Tests/SimpleWebSocketPipeTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ public record SimpleWebSocketPipeTests(ITestOutputHelper Output)
public async Task WhenWebSocketNotOpen_ThenThrowsAsync()
{
IWebSocketPipe pipe = WebSocketPipe.Create(new ClientWebSocket());
await Assert.ThrowsAsync<InvalidOperationException>(() => pipe.RunAsync());
await Assert.ThrowsAsync<InvalidOperationException>(() => pipe.RunAsync().AsTask());
}

[Fact]
Expand All @@ -26,7 +26,7 @@ public async Task WhenConnected_ThenRuns()
using var pipe = WebSocketPipe.Create(socket);

await Task.WhenAll(
pipe.RunAsync(server.Cancellation.Token),
pipe.RunAsync(server.Cancellation.Token).AsTask(),
Task.Delay(100).ContinueWith(_ => server.Cancellation.Cancel()));
}

Expand All @@ -41,7 +41,9 @@ public async Task WhenServerClosesWebSocket_ThenClientCompletesGracefully()

await server.DisposeAsync();

Task.WaitAny(run, Task.Delay(100).ContinueWith(_ => throw new TimeoutException()));
Task.WaitAny(
run.AsTask(),
Task.Delay(100).ContinueWith(_ => throw new TimeoutException()));
}

[Fact]
Expand Down
5 changes: 2 additions & 3 deletions src/Tests/WebSocketServer.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
using System.IO.Pipelines;
using System.Net;
using System.Net;
using System.Net.WebSockets;
using Xunit.Abstractions;

Expand Down Expand Up @@ -52,7 +51,7 @@ static WebSocketServer Create(Func<IWebSocketPipe, Task>? pipeBehavior, Func<Web
if (pipeBehavior != null)
{
using var pipe = WebSocketPipe.Create(websocket, options);
await Task.WhenAll(pipeBehavior(pipe), pipe.RunAsync(cts.Token));
await Task.WhenAll(pipeBehavior(pipe), pipe.RunAsync(cts.Token).AsTask());
}
else if (socketBehavior != null)
{
Expand Down
4 changes: 2 additions & 2 deletions src/WebSocketPipe/IWebSocketPipe.cs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ public interface IWebSocketPipe : IDuplexPipe, IDisposable
/// <param name="closeStatusDescription">Optional close status description to use if the underlying
/// <see cref="WebSocket"/> is closed.</param>
/// <returns></returns>
public Task CompleteAsync(WebSocketCloseStatus? closeStatus = null, string? closeStatusDescription = null);
public ValueTask CompleteAsync(WebSocketCloseStatus? closeStatus = null, string? closeStatusDescription = null);

/// <summary>
/// Starts populating the <see cref="IDuplexPipe.Input"/> with incoming data from the underlying
Expand All @@ -54,6 +54,6 @@ public interface IWebSocketPipe : IDuplexPipe, IDisposable
/// <see cref="IDuplexPipe.Output"/> are completed, or an explicit invocation of <see cref="CompleteAsync"/>
/// is executed.
/// </returns>
public Task RunAsync(CancellationToken cancellation = default);
public ValueTask RunAsync(CancellationToken cancellation = default);
}
}
13 changes: 8 additions & 5 deletions src/WebSocketPipe/SimpleWebSocketPipe.cs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ public SimpleWebSocketPipe(WebSocket webSocket, WebSocketPipeOptions options)

public string? SubProtocol => webSocket.SubProtocol;

public async Task RunAsync(CancellationToken cancellation = default)
public async ValueTask RunAsync(CancellationToken cancellation = default)
{
if (webSocket.State != WebSocketState.Open)
throw new InvalidOperationException($"WebSocket must be opened. State was {webSocket.State}");
Expand All @@ -51,10 +51,13 @@ public async Task RunAsync(CancellationToken cancellation = default)

// NOTE: when both are completed, the CompleteAsync will be called automatically
// by both writing and reading, so we ensure CloseWhenCompleted is performed.
await Task.WhenAll(reading, writing);

// TODO: replace with ValueTask.WhenAll if/when it ships.
// See https://github.com/dotnet/runtime/issues/23625
await Task.WhenAll(reading.AsTask(), writing.AsTask());
}

public async Task CompleteAsync(WebSocketCloseStatus? closeStatus = null, string? closeStatusDescription = null)
public async ValueTask CompleteAsync(WebSocketCloseStatus? closeStatus = null, string? closeStatusDescription = null)
{
if (completed)
return;
Expand Down Expand Up @@ -87,7 +90,7 @@ async ValueTask CloseAsync(WebSocketCloseStatus closeStatus, string closeStatusD
await Task.WhenAny(closeTask, Task.Delay(closeTimeout));
}

async Task FillInputAsync(CancellationToken cancellation)
async ValueTask FillInputAsync(CancellationToken cancellation)
{
while (webSocket.State == WebSocketState.Open && !cancellation.IsCancellationRequested)
{
Expand Down Expand Up @@ -126,7 +129,7 @@ ex is WebSocketException ||
await CompleteAsync(webSocket.CloseStatus, webSocket.CloseStatusDescription);
}

async Task SendOutputAsync(CancellationToken cancellation)
async ValueTask SendOutputAsync(CancellationToken cancellation)
{
while (webSocket.State == WebSocketState.Open && !cancellation.IsCancellationRequested)
{
Expand Down

0 comments on commit 269dfdc

Please sign in to comment.