From f62d68d51371af260be14b6dae029757e354f773 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Mon, 18 May 2020 21:39:57 +0100
Subject: [PATCH] benchmark: exit from UnconstrainedStreamingCall() when stream
 is canceled (#3623)

---
 benchmark/benchmark.go | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/benchmark/benchmark.go b/benchmark/benchmark.go
index 8d9b20f2..e844bb1e 100644
--- a/benchmark/benchmark.go
+++ b/benchmark/benchmark.go
@@ -113,6 +113,7 @@ func (s *testServer) UnconstrainedStreamingCall(stream testpb.BenchmarkService_U
 			err := stream.RecvMsg(in)
 			switch status.Code(err) {
 			case codes.Canceled:
+				return
 			case codes.OK:
 			default:
 				log.Fatalf("server recv error: %v", err)
@@ -125,6 +126,7 @@ func (s *testServer) UnconstrainedStreamingCall(stream testpb.BenchmarkService_U
 			err := stream.Send(response)
 			switch status.Code(err) {
 			case codes.Unavailable:
+				return
 			case codes.OK:
 			default:
 				log.Fatalf("server send error: %v", err)