Merge pull request #17465 from mhaidrygoog/server_streaming_callback_benchmark

Benchmark Streaming support for callback server

1 files changed, 43 insertions, 3 deletions

diff --git a/test/cpp/qps/server_callback.cc b/test/cpp/qps/server_callback.cc
index 8bedd44739..4a346dd017 100644
--- a/test/cpp/qps/server_callback.cc
+++ b/test/cpp/qps/server_callback.cc
@@ -34,13 +34,53 @@ class BenchmarkCallbackServiceImpl final
     : public BenchmarkService::ExperimentalCallbackService {
  public:
   void UnaryCall(
-      ServerContext* context, const SimpleRequest* request,
-      SimpleResponse* response,
-      experimental::ServerCallbackRpcController* controller) override {
+      ServerContext* context, const ::grpc::testing::SimpleRequest* request,
+      ::grpc::testing::SimpleResponse* response,
+      ::grpc::experimental::ServerCallbackRpcController* controller) override {
     auto s = SetResponse(request, response);
     controller->Finish(s);
   }
 
+  ::grpc::experimental::ServerBidiReactor<::grpc::testing::SimpleRequest,
+                                          ::grpc::testing::SimpleResponse>*
+  StreamingCall() override {
+    class Reactor
+        : public ::grpc::experimental::ServerBidiReactor<
+              ::grpc::testing::SimpleRequest, ::grpc::testing::SimpleResponse> {
+     public:
+      Reactor() {}
+      void OnStarted(ServerContext* context) override { StartRead(&request_); }
+
+      void OnReadDone(bool ok) override {
+        if (!ok) {
+          Finish(::grpc::Status::OK);
+          return;
+        }
+        auto s = SetResponse(&request_, &response_);
+        if (!s.ok()) {
+          Finish(s);
+          return;
+        }
+        StartWrite(&response_);
+      }
+
+      void OnWriteDone(bool ok) override {
+        if (!ok) {
+          Finish(::grpc::Status::OK);
+          return;
+        }
+        StartRead(&request_);
+      }
+
+      void OnDone() override { delete (this); }
+
+     private:
+      SimpleRequest request_;
+      SimpleResponse response_;
+    };
+    return new Reactor;
+  }
+
  private:
   static Status SetResponse(const SimpleRequest* request,
                             SimpleResponse* response) {