From b07fcc4abf1198eb52940860fc9fe9f4acb3fb54 Mon Sep 17 00:00:00 2001 From: Kerry Jiang Date: Sat, 13 Apr 2024 21:28:49 -0700 Subject: [PATCH] made the unit test of ping handling work --- src/WebSocket4Net/PingPongStatus.cs | 58 ++++++---------------------- src/WebSocket4Net/WebSocket.cs | 46 +++++++++++++++------- test/WebSocket4Net.Tests/MainTest.cs | 5 ++- 3 files changed, 47 insertions(+), 62 deletions(-) diff --git a/src/WebSocket4Net/PingPongStatus.cs b/src/WebSocket4Net/PingPongStatus.cs index e529861..a30503e 100644 --- a/src/WebSocket4Net/PingPongStatus.cs +++ b/src/WebSocket4Net/PingPongStatus.cs @@ -7,87 +7,53 @@ namespace WebSocket4Net { public class PingPongStatus { - public bool AutoPingEnabled { get; private set; } - - /// - /// The interval of the auto ping - /// - /// in seconds - public int AutoPingInterval { get; private set; } - - /// - /// How long we expect receive pong after ping is sent - /// - /// in seconds - public int ExpectedPongDelay { get; private set; } - private TaskCompletionSource _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(); - 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 pongTask) + if (task is Task) { - 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; } } diff --git a/src/WebSocket4Net/WebSocket.cs b/src/WebSocket4Net/WebSocket.cs index e8cb04d..966f251 100644 --- a/src/WebSocket4Net/WebSocket.cs +++ b/src/WebSocket4Net/WebSocket.cs @@ -26,6 +26,8 @@ public class WebSocket : EasyClient public Uri Uri { get; private set; } + public bool AutoPingEnabled { get; set; } + public CloseStatus CloseStatus { get; private set; } public PingPongStatus PingPongStatus { get; private set; } @@ -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) @@ -152,7 +154,13 @@ public async ValueTask 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; } @@ -196,10 +204,14 @@ public async ValueTask 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; @@ -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) diff --git a/test/WebSocket4Net.Tests/MainTest.cs b/test/WebSocket4Net.Tests/MainTest.cs index 34b5719..27f2d93 100644 --- a/test/WebSocket4Net.Tests/MainTest.cs +++ b/test/WebSocket4Net.Tests/MainTest.cs @@ -197,8 +197,11 @@ await session.SendAsync(new WebSocketPackage OpCode = OpCode.Ping, Data = new ReadOnlySequence(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);