Add option to limit # of messages per stream with tests

src/proto/grpc/testing/control.proto

@@ -113,6 +113,9 @@ message ClientConfig {
   string other_client_api = 15;
 
   repeated ChannelArg channel_args = 16;
+
+  // Number of messages on a stream before it gets finished/restarted
+  int32 messages_per_stream = 18;
 }
 
 message ClientStatus { ClientStats stats = 1; }

test/cpp/qps/client_async.cc

@@ -63,13 +63,13 @@ class ClientRpcContext {
   virtual ~ClientRpcContext() {}
   // next state, return false if done. Collect stats when appropriate
   virtual bool RunNextState(bool, HistogramEntry* entry) = 0;
-  virtual ClientRpcContext* StartNewClone() = 0;
+  virtual void StartNewClone(CompletionQueue* cq) = 0;
   static void* tag(ClientRpcContext* c) { return reinterpret_cast<void*>(c); }
   static ClientRpcContext* detag(void* t) {
     return reinterpret_cast<ClientRpcContext*>(t);
   }
 
-  virtual void Start(CompletionQueue* cq) = 0;
+  virtual void Start(CompletionQueue* cq, const ClientConfig& config) = 0;
 };
 
 template <class RequestType, class ResponseType>
@@ -94,22 +94,17 @@ class ClientRpcContextUnaryImpl : public ClientRpcContext {
         next_issue_(next_issue),
         start_req_(start_req) {}
   ~ClientRpcContextUnaryImpl() override {}
-  void Start(CompletionQueue* cq) override {
-    cq_ = cq;
-    if (!next_issue_) {  // ready to issue
-      RunNextState(true, nullptr);
-    } else {  // wait for the issue time
-      alarm_.reset(new Alarm(cq_, next_issue_(), ClientRpcContext::tag(this)));
-    }
+  void Start(CompletionQueue* cq, const ClientConfig& config) override {
+    StartInternal(cq);
   }
   bool RunNextState(bool ok, HistogramEntry* entry) override {
     switch (next_state_) {
       case State::READY:
         start_ = UsageTimer::Now();
         response_reader_ = start_req_(stub_, &context_, req_, cq_);
+        next_state_ = State::RESP_DONE;
         response_reader_->Finish(&response_, &status_,
                                  ClientRpcContext::tag(this));
-        next_state_ = State::RESP_DONE;
         return true;
       case State::RESP_DONE:
         if (status_.ok()) {
@@ -123,9 +118,10 @@ class ClientRpcContextUnaryImpl : public ClientRpcContext {
         return false;
     }
   }
-  ClientRpcContext* StartNewClone() override {
-    return new ClientRpcContextUnaryImpl(stub_, req_, next_issue_, start_req_,
-                                         callback_);
+  void StartNewClone(CompletionQueue* cq) override {
+    auto* clone = new ClientRpcContextUnaryImpl(stub_, req_, next_issue_,
+                                                start_req_, callback_);
+    clone->StartInternal(cq);
   }
 
  private:
@@ -147,6 +143,15 @@ class ClientRpcContextUnaryImpl : public ClientRpcContext {
   double start_;
   std::unique_ptr<grpc::ClientAsyncResponseReader<ResponseType>>
       response_reader_;
+
+  void StartInternal(CompletionQueue* cq) {
+    cq_ = cq;
+    if (!next_issue_) {  // ready to issue
+      RunNextState(true, nullptr);
+    } else {  // wait for the issue time
+      alarm_.reset(new Alarm(cq_, next_issue_(), ClientRpcContext::tag(this)));
+    }
+  }
 };
 
 typedef std::forward_list<ClientRpcContext*> context_list;
@@ -185,7 +190,7 @@ class AsyncClient : public ClientImpl<StubType, RequestType> {
         auto* cq = cli_cqs_[t].get();
         auto ctx =
             setup_ctx(channels_[ch].get_stub(), next_issuers_[t], request_);
-        ctx->Start(cq);
+        ctx->Start(cq, config);
       }
       t = (t + 1) % cli_cqs_.size();
     }
@@ -248,8 +253,7 @@ class AsyncClient : public ClientImpl<StubType, RequestType> {
         } else if (!ctx->RunNextState(ok, entry)) {
           // The RPC and callback are done, so clone the ctx
           // and kickstart the new one
-          auto clone = ctx->StartNewClone();
-          clone->Start(cli_cqs_[thread_idx].get());
+          ctx->StartNewClone(cli_cqs_[thread_idx].get());
           // delete the old version
           delete ctx;
         }
@@ -330,10 +334,8 @@ class ClientRpcContextStreamingImpl : public ClientRpcContext {
         next_issue_(next_issue),
         start_req_(start_req) {}
   ~ClientRpcContextStreamingImpl() override {}
-  void Start(CompletionQueue* cq) override {
-    cq_ = cq;
-    stream_ = start_req_(stub_, &context_, cq, ClientRpcContext::tag(this));
-    next_state_ = State::STREAM_IDLE;
+  void Start(CompletionQueue* cq, const ClientConfig& config) override {
+    StartInternal(cq, config.messages_per_stream());
   }
   bool RunNextState(bool ok, HistogramEntry* entry) override {
     while (true) {
@@ -346,9 +348,9 @@ class ClientRpcContextStreamingImpl : public ClientRpcContext {
           }
           break;  // loop around, don't return
         case State::WAIT:
+          next_state_ = State::READY_TO_WRITE;
           alarm_.reset(
               new Alarm(cq_, next_issue_(), ClientRpcContext::tag(this)));
-          next_state_ = State::READY_TO_WRITE;
           return true;
         case State::READY_TO_WRITE:
           if (!ok) {
@@ -369,17 +371,32 @@ class ClientRpcContextStreamingImpl : public ClientRpcContext {
         case State::READ_DONE:
           entry->set_value((UsageTimer::Now() - start_) * 1e9);
           callback_(status_, &response_);
+          if ((messages_per_stream_ != 0) &&
+              (++messages_issued_ >= messages_per_stream_)) {
+            next_state_ = State::WRITES_DONE_DONE;
+            stream_->WritesDone(ClientRpcContext::tag(this));
+            return true;
+          }
           next_state_ = State::STREAM_IDLE;
           break;  // loop around
+        case State::WRITES_DONE_DONE:
+          next_state_ = State::FINISH_DONE;
+          stream_->Finish(&status_, ClientRpcContext::tag(this));
+          return true;
+        case State::FINISH_DONE:
+          next_state_ = State::INVALID;
+          return false;
+          break;
         default:
           GPR_ASSERT(false);
           return false;
       }
     }
   }
-  ClientRpcContext* StartNewClone() override {
-    return new ClientRpcContextStreamingImpl(stub_, req_, next_issue_,
-                                             start_req_, callback_);
+  void StartNewClone(CompletionQueue* cq) override {
+    auto* clone = new ClientRpcContextStreamingImpl(stub_, req_, next_issue_,
+                                                    start_req_, callback_);
+    clone->StartInternal(cq, messages_per_stream_);
   }
 
  private:
@@ -395,7 +412,9 @@ class ClientRpcContextStreamingImpl : public ClientRpcContext {
     WAIT,
     READY_TO_WRITE,
     WRITE_DONE,
-    READ_DONE
+    READ_DONE,
+    WRITES_DONE_DONE,
+    FINISH_DONE
   };
   State next_state_;
   std::function<void(grpc::Status, ResponseType*)> callback_;
@@ -408,6 +427,18 @@ class ClientRpcContextStreamingImpl : public ClientRpcContext {
   double start_;
   std::unique_ptr<grpc::ClientAsyncReaderWriter<RequestType, ResponseType>>
       stream_;
+
+  // Allow a limit on number of messages in a stream
+  int messages_per_stream_;
+  int messages_issued_;
+
+  void StartInternal(CompletionQueue* cq, int messages_per_stream) {
+    cq_ = cq;
+    next_state_ = State::STREAM_IDLE;
+    stream_ = start_req_(stub_, &context_, cq, ClientRpcContext::tag(this));
+    messages_per_stream_ = messages_per_stream;
+    messages_issued_ = 0;
+  }
 };
 
 class AsyncStreamingClient final
@@ -459,13 +490,8 @@ class ClientRpcContextGenericStreamingImpl : public ClientRpcContext {
         next_issue_(next_issue),
         start_req_(start_req) {}
   ~ClientRpcContextGenericStreamingImpl() override {}
-  void Start(CompletionQueue* cq) override {
-    cq_ = cq;
-    const grpc::string kMethodName(
-        "/grpc.testing.BenchmarkService/StreamingCall");
-    stream_ = start_req_(stub_, &context_, kMethodName, cq,
-                         ClientRpcContext::tag(this));
-    next_state_ = State::STREAM_IDLE;
+  void Start(CompletionQueue* cq, const ClientConfig& config) override {
+    StartInternal(cq, config.messages_per_stream());
   }
   bool RunNextState(bool ok, HistogramEntry* entry) override {
     while (true) {
@@ -478,9 +504,9 @@ class ClientRpcContextGenericStreamingImpl : public ClientRpcContext {
           }
           break;  // loop around, don't return
         case State::WAIT:
+          next_state_ = State::READY_TO_WRITE;
           alarm_.reset(
               new Alarm(cq_, next_issue_(), ClientRpcContext::tag(this)));
-          next_state_ = State::READY_TO_WRITE;
           return true;
         case State::READY_TO_WRITE:
           if (!ok) {
@@ -501,17 +527,32 @@ class ClientRpcContextGenericStreamingImpl : public ClientRpcContext {
         case State::READ_DONE:
           entry->set_value((UsageTimer::Now() - start_) * 1e9);
           callback_(status_, &response_);
+          if ((messages_per_stream_ != 0) &&
+              (++messages_issued_ >= messages_per_stream_)) {
+            next_state_ = State::WRITES_DONE_DONE;
+            stream_->WritesDone(ClientRpcContext::tag(this));
+            return true;
+          }
           next_state_ = State::STREAM_IDLE;
           break;  // loop around
+        case State::WRITES_DONE_DONE:
+          next_state_ = State::FINISH_DONE;
+          stream_->Finish(&status_, ClientRpcContext::tag(this));
+          return true;
+        case State::FINISH_DONE:
+          next_state_ = State::INVALID;
+          return false;
+          break;
         default:
           GPR_ASSERT(false);
           return false;
       }
     }
   }
-  ClientRpcContext* StartNewClone() override {
-    return new ClientRpcContextGenericStreamingImpl(stub_, req_, next_issue_,
-                                                    start_req_, callback_);
+  void StartNewClone(CompletionQueue* cq) override {
+    auto* clone = new ClientRpcContextGenericStreamingImpl(
+        stub_, req_, next_issue_, start_req_, callback_);
+    clone->StartInternal(cq, messages_per_stream_);
   }
 
  private:
@@ -527,7 +568,9 @@ class ClientRpcContextGenericStreamingImpl : public ClientRpcContext {
     WAIT,
     READY_TO_WRITE,
     WRITE_DONE,
-    READ_DONE
+    READ_DONE,
+    WRITES_DONE_DONE,
+    FINISH_DONE
   };
   State next_state_;
   std::function<void(grpc::Status, ByteBuffer*)> callback_;
@@ -539,6 +582,21 @@ class ClientRpcContextGenericStreamingImpl : public ClientRpcContext {
   grpc::Status status_;
   double start_;
   std::unique_ptr<grpc::GenericClientAsyncReaderWriter> stream_;
+
+  // Allow a limit on number of messages in a stream
+  int messages_per_stream_;
+  int messages_issued_;
+
+  void StartInternal(CompletionQueue* cq, int messages_per_stream) {
+    cq_ = cq;
+    const grpc::string kMethodName(
+        "/grpc.testing.BenchmarkService/StreamingCall");
+    next_state_ = State::STREAM_IDLE;
+    stream_ = start_req_(stub_, &context_, kMethodName, cq,
+                         ClientRpcContext::tag(this));
+    messages_per_stream_ = messages_per_stream;
+    messages_issued_ = 0;
+  }
 };
 
 static std::unique_ptr<grpc::GenericStub> GenericStubCreator(

test/cpp/qps/client_sync.cc

@@ -142,10 +142,13 @@ class SynchronousStreamingClient final : public SynchronousClient {
   SynchronousStreamingClient(const ClientConfig& config)
       : SynchronousClient(config),
         context_(num_threads_),
-        stream_(num_threads_) {
+        stream_(num_threads_),
+        messages_per_stream_(config.messages_per_stream()),
+        messages_issued_(num_threads_) {
     for (size_t thread_idx = 0; thread_idx < num_threads_; thread_idx++) {
       auto* stub = channels_[thread_idx % channels_.size()].get_stub();
       stream_[thread_idx] = stub->StreamingCall(&context_[thread_idx]);
+      messages_issued_[thread_idx] = 0;
     }
     StartThreads(num_threads_);
   }
@@ -173,11 +176,17 @@ class SynchronousStreamingClient final : public SynchronousClient {
         stream_[thread_idx]->Read(&responses_[thread_idx])) {
       entry->set_value((UsageTimer::Now() - start) * 1e9);
       // don't set the status since there isn't one yet
-      return true;
+      if ((messages_per_stream_ != 0) &&
+          (++messages_issued_[thread_idx] < messages_per_stream_)) {
+        return true;
+      } else {
+        // Fall through to the below resetting code after finish
+      }
     }
     stream_[thread_idx]->WritesDone();
     Status s = stream_[thread_idx]->Finish();
-    // don't set the value since the stream is failed and shouldn't be timed
+    // don't set the value since this is either a failure (shouldn't be timed)
+    // or a stream-end (already has been timed)
     entry->set_status(s.error_code());
     if (!s.ok()) {
       gpr_log(GPR_ERROR, "Stream %" PRIuPTR " received an error %s", thread_idx,
@@ -187,6 +196,7 @@ class SynchronousStreamingClient final : public SynchronousClient {
     context_[thread_idx].~ClientContext();
     new (&context_[thread_idx]) ClientContext();
     stream_[thread_idx] = stub->StreamingCall(&context_[thread_idx]);
+    messages_issued_[thread_idx] = 0;
     return true;
   }
 
@@ -197,6 +207,8 @@ class SynchronousStreamingClient final : public SynchronousClient {
   std::vector<
       std::unique_ptr<grpc::ClientReaderWriter<SimpleRequest, SimpleResponse>>>
       stream_;
+  const int messages_per_stream_;
+  std::vector<int> messages_issued_;
 };
 
 std::unique_ptr<Client> CreateSynchronousUnaryClient(

tools/run_tests/performance/scenario_config.py

@@ -112,6 +112,7 @@ def _ping_pong_scenario(name, rpc_type,
                         channels=None,
                         outstanding=None,
                         resource_quota_size=None,
+                        messages_per_stream=None,
                         excluded_poll_engines=[]):
   """Creates a basic ping pong scenario."""
   scenario = {
@@ -165,6 +166,8 @@ def _ping_pong_scenario(name, rpc_type,
     scenario['client_config']['client_channels'] = 1
     scenario['client_config']['async_client_threads'] = 1
 
+  if messages_per_stream:
+    scenario['client_config']['messages_per_stream'] = messages_per_stream
   if client_language:
     # the CLIENT_LANGUAGE field is recognized by run_performance_tests.py
     scenario['CLIENT_LANGUAGE'] = client_language
@@ -214,6 +217,26 @@ class CXXLanguage:
           secure=secure,
           categories=smoketest_categories+[SCALABLE])
 
+      for mps in geometric_progression(1, 20, 10):
+        yield _ping_pong_scenario(
+            'cpp_generic_async_streaming_qps_unconstrained_%smps_%s' % (mps, secstr),
+            rpc_type='STREAMING',
+            client_type='ASYNC_CLIENT',
+            server_type='ASYNC_GENERIC_SERVER',
+            unconstrained_client='async', use_generic_payload=True,
+            secure=secure, messages_per_stream=mps,
+            categories=smoketest_categories+[SCALABLE])
+
+      for mps in geometric_progression(1, 200, math.sqrt(10)):
+        yield _ping_pong_scenario(
+            'cpp_generic_async_streaming_qps_unconstrained_%smps_%s' % (mps, secstr),
+            rpc_type='STREAMING',
+            client_type='ASYNC_CLIENT',
+            server_type='ASYNC_GENERIC_SERVER',
+            unconstrained_client='async', use_generic_payload=True,
+            secure=secure, messages_per_stream=mps,
+            categories=[SWEEP])
+
       yield _ping_pong_scenario(
           'cpp_generic_async_streaming_qps_1channel_1MBmsg_%s' % secstr,
           rpc_type='STREAMING',
@@ -331,6 +354,27 @@ class CXXLanguage:
           #     categories=smoketest_categories+[SCALABLE],
           #     resource_quota_size=500*1024)
 
+          if rpc_type == 'streaming':
+            for mps in geometric_progression(1, 20, 10):
+              yield _ping_pong_scenario(
+                  'cpp_protobuf_%s_%s_qps_unconstrained_%smps_%s' % (synchronicity, rpc_type, mps, secstr),
+                  rpc_type=rpc_type.upper(),
+                  client_type='%s_CLIENT' % synchronicity.upper(),
+                  server_type='%s_SERVER' % synchronicity.upper(),
+                  unconstrained_client=synchronicity,
+                  secure=secure, messages_per_stream=mps,
+                  categories=smoketest_categories+[SCALABLE])
+
+            for mps in geometric_progression(1, 200, math.sqrt(10)):
+              yield _ping_pong_scenario(
+                  'cpp_protobuf_%s_%s_qps_unconstrained_%smps_%s' % (synchronicity, rpc_type, mps, secstr),
+                  rpc_type=rpc_type.upper(),
+                  client_type='%s_CLIENT' % synchronicity.upper(),
+                  server_type='%s_SERVER' % synchronicity.upper(),
+                  unconstrained_client=synchronicity,
+                  secure=secure, messages_per_stream=mps,
+                  categories=[SWEEP])
+
           for channels in geometric_progression(1, 20000, math.sqrt(10)):
             for outstanding in geometric_progression(1, 200000, math.sqrt(10)):
                 if synchronicity == 'sync' and outstanding > 1200: continue