Make QPS server builder customizeabled

6 files changed, 117 insertions, 17 deletions

diff --git a/test/cpp/qps/BUILD b/test/cpp/qps/BUILD
index e7d093c71a..b958c75fc7 100644
--- a/test/cpp/qps/BUILD
+++ b/test/cpp/qps/BUILD
@@ -34,11 +34,13 @@ grpc_cc_library(
         "qps_worker.cc",
         "server_async.cc",
         "server_sync.cc",
+        "qps_server_builder.cc",
     ],
     hdrs = [
         "client.h",
         "qps_worker.h",
         "server.h",
+        "qps_server_builder.h",
     ],
     deps = [
         ":histogram",
@@ -55,6 +57,10 @@ grpc_cc_library(
         "//test/core/util:gpr_test_util",
         "//test/core/util:grpc_test_util",
         "//test/cpp/util:test_util",
+        "//test/cpp/util:test_config",
+    ],
+    external_deps = [
+        "gflags",
     ],
 )
 
diff --git a/test/cpp/qps/qps_server_builder.cc b/test/cpp/qps/qps_server_builder.cc
new file mode 100644
index 0000000000..5fbc682b75
--- /dev/null
+++ b/test/cpp/qps/qps_server_builder.cc
@@ -0,0 +1,45 @@
+/*
+ *
+ * Copyright 2016 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+#include "qps_server_builder.h"
+
+using grpc::ServerBuilder;
+
+namespace grpc {
+namespace testing {
+
+namespace {
+std::unique_ptr<ServerBuilder> DefaultCreateQpsServerBuilder() {
+  return std::unique_ptr<ServerBuilder>(new ServerBuilder());
+}
+
+std::function<std::unique_ptr<ServerBuilder>()> g_create_qps_server_builder =
+    DefaultCreateQpsServerBuilder;
+}  // namespace
+
+std::unique_ptr<ServerBuilder> CreateQpsServerBuilder() {
+  return g_create_qps_server_builder();
+}
+
+void SetCreateQpsServerBuilderFunc(
+    std::function<std::unique_ptr<ServerBuilder>()> create_qps_server_builder) {
+  g_create_qps_server_builder = std::move(create_qps_server_builder);
+}
+
+}  // namespace testing
+}  // namespace grpc
diff --git a/test/cpp/qps/qps_server_builder.h b/test/cpp/qps/qps_server_builder.h
new file mode 100644
index 0000000000..98f9fa72e9
--- /dev/null
+++ b/test/cpp/qps/qps_server_builder.h
@@ -0,0 +1,46 @@
+/*
+ *
+ * Copyright 2016 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+#ifndef GRPC_QPS_SERVER_BUILDER_H
+#define GRPC_QPS_SERVER_BUILDER_H
+
+#include <functional>
+#include <memory>
+
+#include <grpcpp/server_builder.h>
+
+namespace grpc {
+namespace testing {
+
+// CreateQpsServerBuilder creates a new ServerBuilder.
+// This uses the "create ServerBuilder" func that was set
+// in SetCreateQpsServerBuilderFunc if one has been set,
+// otherwise, this defaults to creating a new ServerBuilder
+// with only its default constructor.
+std::unique_ptr<ServerBuilder> CreateQpsServerBuilder();
+
+// SetCreateQpsServerBuilderFunc sets a function to use to create new
+// ServerBuilders in "CreateQpsServerBuilder". It can be used to modify options
+// that the server is built with.
+void SetCreateQpsServerBuilderFunc(
+    std::function<std::unique_ptr<ServerBuilder>()>);
+
+}  // namespace testing
+}  // namespace grpc
+
+#endif  // GRPC_QPS_SERVER_BUILDER_H
diff --git a/test/cpp/qps/qps_worker.cc b/test/cpp/qps/qps_worker.cc
index d3f0380474..7ddf3c1cf3 100644
--- a/test/cpp/qps/qps_worker.cc
+++ b/test/cpp/qps/qps_worker.cc
@@ -39,6 +39,7 @@
 #include "test/core/util/grpc_profiler.h"
 #include "test/core/util/histogram.h"
 #include "test/cpp/qps/client.h"
+#include "test/cpp/qps/qps_server_builder.h"
 #include "test/cpp/qps/server.h"
 #include "test/cpp/util/create_test_channel.h"
 #include "test/cpp/util/test_credentials_provider.h"
@@ -272,18 +273,18 @@ QpsWorker::QpsWorker(int driver_port, int server_port,
   impl_.reset(new WorkerServiceImpl(server_port, this));
   gpr_atm_rel_store(&done_, static_cast<gpr_atm>(0));
 
-  ServerBuilder builder;
+  std::unique_ptr<ServerBuilder> builder = CreateQpsServerBuilder();
   if (driver_port >= 0) {
     char* server_address = nullptr;
     gpr_join_host_port(&server_address, "::", driver_port);
-    builder.AddListeningPort(
+    builder->AddListeningPort(
         server_address,
         GetCredentialsProvider()->GetServerCredentials(credential_type));
     gpr_free(server_address);
   }
-  builder.RegisterService(impl_.get());
+  builder->RegisterService(impl_.get());
 
-  server_ = builder.BuildAndStart();
+  server_ = builder->BuildAndStart();
 }
 
 QpsWorker::~QpsWorker() {}
diff --git a/test/cpp/qps/server_async.cc b/test/cpp/qps/server_async.cc
index 1dfef6cfc1..5cd975cf74 100644
--- a/test/cpp/qps/server_async.cc
+++ b/test/cpp/qps/server_async.cc
@@ -38,6 +38,7 @@
 #include "src/core/lib/surface/completion_queue.h"
 #include "src/proto/grpc/testing/benchmark_service.grpc.pb.h"
 #include "test/core/util/test_config.h"
+#include "test/cpp/qps/qps_server_builder.h"
 #include "test/cpp/qps/server.h"
 
 namespace grpc {
@@ -74,19 +75,19 @@ class AsyncQpsServerTest final : public grpc::testing::Server {
                                  ResponseType*)>
           process_rpc)
       : Server(config) {
-    ServerBuilder builder;
+    std::unique_ptr<ServerBuilder> builder = CreateQpsServerBuilder();
 
     auto port_num = port();
     // Negative port number means inproc server, so no listen port needed
     if (port_num >= 0) {
       char* server_address = nullptr;
       gpr_join_host_port(&server_address, "::", port_num);
-      builder.AddListeningPort(server_address,
-                               Server::CreateServerCredentials(config));
+      builder->AddListeningPort(server_address,
+                                Server::CreateServerCredentials(config));
       gpr_free(server_address);
     }
 
-    register_service(&builder, &async_service_);
+    register_service(builder.get(), &async_service_);
 
     int num_threads = config.async_server_threads();
     if (num_threads <= 0) {  // dynamic sizing
@@ -97,15 +98,15 @@ class AsyncQpsServerTest final : public grpc::testing::Server {
     int tpc = std::max(1, config.threads_per_cq());  // 1 if unspecified
     int num_cqs = (num_threads + tpc - 1) / tpc;     // ceiling operator
     for (int i = 0; i < num_cqs; i++) {
-      srv_cqs_.emplace_back(builder.AddCompletionQueue());
+      srv_cqs_.emplace_back(builder->AddCompletionQueue());
     }
     for (int i = 0; i < num_threads; i++) {
       cq_.emplace_back(i % srv_cqs_.size());
     }
 
-    ApplyConfigToBuilder(config, &builder);
+    ApplyConfigToBuilder(config, builder.get());
 
-    server_ = builder.BuildAndStart();
+    server_ = builder->BuildAndStart();
 
     auto process_rpc_bound =
         std::bind(process_rpc, config.payload_config(), std::placeholders::_1,
diff --git a/test/cpp/qps/server_sync.cc b/test/cpp/qps/server_sync.cc
index b8facf9b56..2e63f5ec86 100644
--- a/test/cpp/qps/server_sync.cc
+++ b/test/cpp/qps/server_sync.cc
@@ -27,6 +27,7 @@
 
 #include "src/core/lib/gpr/host_port.h"
 #include "src/proto/grpc/testing/benchmark_service.grpc.pb.h"
+#include "test/cpp/qps/qps_server_builder.h"
 #include "test/cpp/qps/server.h"
 #include "test/cpp/qps/usage_timer.h"
 
@@ -154,23 +155,23 @@ class BenchmarkServiceImpl final : public BenchmarkService::Service {
 class SynchronousServer final : public grpc::testing::Server {
  public:
   explicit SynchronousServer(const ServerConfig& config) : Server(config) {
-    ServerBuilder builder;
+    std::unique_ptr<ServerBuilder> builder = CreateQpsServerBuilder();
 
     auto port_num = port();
     // Negative port number means inproc server, so no listen port needed
     if (port_num >= 0) {
       char* server_address = nullptr;
       gpr_join_host_port(&server_address, "::", port_num);
-      builder.AddListeningPort(server_address,
-                               Server::CreateServerCredentials(config));
+      builder->AddListeningPort(server_address,
+                                Server::CreateServerCredentials(config));
       gpr_free(server_address);
     }
 
-    ApplyConfigToBuilder(config, &builder);
+    ApplyConfigToBuilder(config, builder.get());
 
-    builder.RegisterService(&service_);
+    builder->RegisterService(&service_);
 
-    impl_ = builder.BuildAndStart();
+    impl_ = builder->BuildAndStart();
   }
 
   std::shared_ptr<Channel> InProcessChannel(