Skip to content

Commit

Permalink
made the unit test of ping handling work
Browse files Browse the repository at this point in the history
  • Loading branch information
@kerryjiang
kerryjiang committed Apr 14, 2024
1 parent 4e881b3 commit b07fcc4
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 62 deletions.
58 changes: 12 additions & 46 deletions src/WebSocket4Net/PingPongStatus.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,87 +7,53 @@ namespace WebSocket4Net
{
public class PingPongStatus
{
public bool AutoPingEnabled { get; private set; }

/// <summary>
/// The interval of the auto ping
/// </summary>
/// <value>in seconds</value>
public int AutoPingInterval { get; private set; }

/// <summary>
/// How long we expect receive pong after ping is sent
/// </summary>
/// <value>in seconds</value>
public int ExpectedPongDelay { get; private set; }

private TaskCompletionSource<WebSocketPackage> _pongReceivedTaskSource;

public DateTimeOffset LastPingReceived { get; internal set; }

public DateTimeOffset LastPongReceived { get; internal set; }

internal WebSocket WebSocket { get; set; }

internal PingPongStatus(AutoPingOptions options)
{
if (options != null)
{
AutoPingEnabled = true;
AutoPingInterval = options.AutoPingInterval;
ExpectedPongDelay = options.ExpectedPongDelay;
}
}

internal void Start()
internal PingPongStatus()
{
RunAutoPing();
}

private async void RunAutoPing()
internal async Task RunAutoPing(WebSocket webSocket, AutoPingOptions options)
{
while (AutoPingEnabled)
while (webSocket.State == WebSocketState.Open)
{
var autoPingInterval = Math.Max(AutoPingInterval, 60) * 1000;
var autoPingInterval = Math.Max(options.AutoPingInterval, 60) * 1000;

await Task.Delay(autoPingInterval);

_pongReceivedTaskSource = new TaskCompletionSource<WebSocketPackage>();

await WebSocket.SendAsync(new WebSocketPackage
await webSocket.SendAsync(new WebSocketPackage
{
OpCode = OpCode.Ping
});

var pongExpectAfterPing = Math.Max(ExpectedPongDelay, autoPingInterval * 3);
var pongExpectAfterPing = Math.Max(options.ExpectedPongDelay, autoPingInterval * 3);
var task = await Task.WhenAny(_pongReceivedTaskSource.Task, Task.Delay(pongExpectAfterPing));

if (task is Task<WebSocketPackage> pongTask)
if (task is Task<WebSocketPackage>)
{
LastPongReceived = DateTimeOffset.Now;
continue;
}

// Pong doesn't arrive on time
await WebSocket.CloseAsync(CloseReason.UnexpectedCondition, "Pong is not received on time.");
await webSocket.CloseAsync(CloseReason.UnexpectedCondition, "Pong is not received on time.");
break;
}
}

internal void Stop()
{
AutoPingEnabled = false;
}

internal ValueTask OnPongReceived(WebSocketPackage pong)
internal void OnPongReceived(WebSocketPackage pong)
{
LastPongReceived = DateTimeOffset.Now;
_pongReceivedTaskSource.SetResult(pong);
return new ValueTask();
}

internal async ValueTask OnPingReceived(WebSocketPackage ping)
internal void OnPingReceived(WebSocketPackage ping)
{
ping.OpCode = OpCode.Pong;
await WebSocket.SendAsync(ping);
LastPingReceived = DateTimeOffset.Now;
}
}
Expand Down
46 changes: 31 additions & 15 deletions src/WebSocket4Net/WebSocket.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ public class WebSocket : EasyClient<WebSocketPackage>

public Uri Uri { get; private set; }

public bool AutoPingEnabled { get; set; }

public CloseStatus CloseStatus { get; private set; }

public PingPongStatus PingPongStatus { get; private set; }
Expand Down Expand Up @@ -77,7 +79,7 @@ public WebSocket(string url, ConnectionOptions connectionOptions)
throw new ArgumentException("Unexpected url schema.", nameof(url));
}

PingPongStatus = new PingPongStatus(new AutoPingOptions(60 * 5, 5));
PingPongStatus = new PingPongStatus();
}

private EndPoint ResolveUri(Uri uri, int defaultPort)
Expand Down Expand Up @@ -152,7 +154,13 @@ public async ValueTask<bool> OpenAsync(CancellationToken cancellationToken = def
}

State = WebSocketState.Open;
PingPongStatus.Start();

if (AutoPingEnabled)
{
var autoPingTask = PingPongStatus.RunAutoPing(this, new AutoPingOptions(60 * 5, 5));
_ = autoPingTask.ContinueWith(t => OnError("AutoPing failed", t.Exception), TaskContinuationOptions.OnlyOnFaulted);
}

return true;
}

Expand Down Expand Up @@ -196,10 +204,14 @@ public async ValueTask<WebSocketPackage> ReceiveAsync(bool returnControlPackage)
{
var package = await base.ReceiveAsync();

if (!returnControlPackage && package.OpCode != OpCode.Binary && package.OpCode != OpCode.Text && package.OpCode != OpCode.Handshake)
if (package.OpCode != OpCode.Binary && package.OpCode != OpCode.Text && package.OpCode != OpCode.Handshake)
{
await HandleControlPackage(package);
return await ReceiveAsync(returnControlPackage);

if (!returnControlPackage)
{
return await ReceiveAsync(returnControlPackage);
}
}

return package;
Expand All @@ -216,20 +228,24 @@ protected override async ValueTask OnPackageReceived(WebSocketPackage package)
await base.OnPackageReceived(package);
}

private ValueTask HandleControlPackage(WebSocketPackage package)
private async ValueTask HandleControlPackage(WebSocketPackage package)
{
if (package.OpCode == OpCode.Close)
switch (package.OpCode)
{
HandleCloseHandshake(package);
return new ValueTask();
case (OpCode.Close):
HandleCloseHandshake(package);
break;

case (OpCode.Ping):
PingPongStatus.OnPingReceived(package);
package.OpCode = OpCode.Pong;
await this.SendAsync(package);
break;

case (OpCode.Pong):
PingPongStatus.OnPongReceived(package);
break;
}

if (package.OpCode == OpCode.Ping)
return PingPongStatus.OnPingReceived(package);
else if (package.OpCode == OpCode.Pong)
return PingPongStatus.OnPongReceived(package);

return new ValueTask();
}

public async ValueTask SendAsync(string message)
Expand Down
5 changes: 4 additions & 1 deletion test/WebSocket4Net.Tests/MainTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,8 +197,11 @@ await session.SendAsync(new WebSocketPackage
OpCode = OpCode.Ping,
Data = new ReadOnlySequence<byte>(Utf8Encoding.GetBytes("Hello"))
});

var package = await websocket.ReceiveAsync(true);

Assert.NotNull(package);

await Task.Delay(1 * 1000);
var lastPingReceivedNow = websocket.PingPongStatus.LastPingReceived;

Assert.NotEqual(lastPingReceived, lastPingReceivedNow);
Expand Down

0 comments on commit b07fcc4

Please sign in to comment.