chore: fix benchmark cancellation for ServerStreaming Scenario

Cysharp · Nov 2, 2024 · f81c513 · f81c513
1 parent 973e5e7
commit f81c513
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/perf/BenchmarkApp/PerformanceTest.Client/ServerStreamingScenario.cs b/perf/BenchmarkApp/PerformanceTest.Client/ServerStreamingScenario.cs
@@ -1,3 +1,4 @@
+using Grpc.Core;
 using Grpc.Net.Client;
 using MagicOnion.Client;
 using PerformanceTest.Shared;
@@ -19,13 +20,21 @@ public ValueTask PrepareAsync(GrpcChannel channel)
     // So most times MoveNext won't wait at all, and it may wait occasionally.
     public async ValueTask RunAsync(int connectionId, PerformanceTestRunningContext ctx, CancellationToken cancellationToken)
     {
-        var stream = await client.ServerStreamingAsync();
-        while(!cancellationToken.IsCancellationRequested)
+        using var stream = await client.ServerStreamingAsync();
+        while (!cancellationToken.IsCancellationRequested)
         {
             ctx.Increment();
             var begin = timeProvider.GetTimestamp();
-            await stream.ResponseStream.MoveNext(cancellationToken);
-            ctx.Latency(connectionId, timeProvider.GetElapsedTime(begin));
+            try
+            {
+                await stream.ResponseStream.MoveNext(cancellationToken);
+                ctx.Latency(connectionId, timeProvider.GetElapsedTime(begin));
+            }
+            catch (RpcException ex) when (ex.StatusCode == StatusCode.Cancelled && cancellationToken.IsCancellationRequested)
+            {
+                // canceling call is expected behavior.
+                break;
+            }
         }
     }