Clean-up core list usage and make it possible to reset the core list

6 files changed, 30 insertions, 30 deletions

diff --git a/src/proto/grpc/testing/control.proto b/src/proto/grpc/testing/control.proto
index b2b83af3df..40e8a97471 100644
--- a/src/proto/grpc/testing/control.proto
+++ b/src/proto/grpc/testing/control.proto
@@ -134,8 +134,7 @@ message ServerConfig {
   int32 port = 4;
   // Only for async server. Number of threads used to serve the requests.
   int32 async_server_threads = 7;
-  // restrict core usage, currently unused
-  int32 core_limit = 8;
+  // payload config, used in generic server
   PayloadConfig payload_config = 9;
 
   // Specify the cores we should run the server on, if desired
diff --git a/test/cpp/qps/client.h b/test/cpp/qps/client.h
index 576adeb256..7d5f6466f9 100644
--- a/test/cpp/qps/client.h
+++ b/test/cpp/qps/client.h
@@ -324,6 +324,8 @@ class ClientImpl : public Client {
              std::function<std::unique_ptr<StubType>(std::shared_ptr<Channel>)>
                  create_stub)
       : channels_(config.client_channels()), create_stub_(create_stub) {
+    LimitCores(config.core_list().data(), config.core_list_size());
+
     for (int i = 0; i < config.client_channels(); i++) {
       channels_[i].init(config.server_targets(i % config.server_targets_size()),
                         config, create_stub_);
diff --git a/test/cpp/qps/limit_cores.cc b/test/cpp/qps/limit_cores.cc
index 0ba46d3d0a..5fd8d555a5 100644
--- a/test/cpp/qps/limit_cores.cc
+++ b/test/cpp/qps/limit_cores.cc
@@ -46,23 +46,31 @@ namespace testing {
 #define _GNU_SOURCE
 #endif
 #include <sched.h>
-int LimitCores(std::vector<int> cores) {
-  size_t num_cores = static_cast<size_t>(gpr_cpu_num_cores());
-  if (num_cores > cores.size()) {
-    cpu_set_t *cpup = CPU_ALLOC(num_cores);
-    GPR_ASSERT(cpup);
-    size_t size = CPU_ALLOC_SIZE(num_cores);
-    CPU_ZERO_S(size, cpup);
+int LimitCores(const int *cores, int cores_size) {
+  int num_cores = gpr_cpu_num_cores();
+  int cores_set = 0;
 
-    for (size_t i = 0; i < cores.size(); i++) {
-      CPU_SET_S(cores[i], size, cpup);
+  cpu_set_t *cpup = CPU_ALLOC(num_cores);
+  GPR_ASSERT(cpup);
+  size_t size = CPU_ALLOC_SIZE(num_cores);
+  CPU_ZERO_S(size, cpup);
+
+  if (cores_size > 0) {
+    for (int i = 0; i < cores_size; i++) {
+      if (cores[i] < num_cores) {
+        CPU_SET_S(cores[i], size, cpup);
+        cores_set++;
+      }
     }
-    GPR_ASSERT(sched_setaffinity(0, size, cpup) == 0);
-    CPU_FREE(cpup);
-    return cores.size();
   } else {
-    return num_cores;
+    for (int i = 0; i < num_cores; i++) {
+      CPU_SET_S(i, size, cpup);
+      cores_set++;
+    }
   }
+  GPR_ASSERT(sched_setaffinity(0, size, cpup) == 0);
+  CPU_FREE(cpup);
+  return cores_set;
 }
 #else
 // LimitCores is not currently supported for non-Linux platforms
diff --git a/test/cpp/qps/limit_cores.h b/test/cpp/qps/limit_cores.h
index 54c805216c..5467f3b881 100644
--- a/test/cpp/qps/limit_cores.h
+++ b/test/cpp/qps/limit_cores.h
@@ -38,7 +38,10 @@
 
 namespace grpc {
 namespace testing {
-int LimitCores(std::vector<int> core_vec);
+// LimitCores takes array and size arguments (instead of vector) for more direct
+// conversion from repeated field of protobuf. Use a cores_size of 0 to remove
+// existing limits (from an empty repeated field)
+int LimitCores(const int *cores, int cores_size);
 }  // namespace testing
 }  // namespace grpc
 
diff --git a/test/cpp/qps/qps_worker.cc b/test/cpp/qps/qps_worker.cc
index 5cb5850fd4..6289c1a843 100644
--- a/test/cpp/qps/qps_worker.cc
+++ b/test/cpp/qps/qps_worker.cc
@@ -87,10 +87,6 @@ static std::unique_ptr<Server> CreateServer(const ServerConfig& config) {
   gpr_log(GPR_INFO, "Starting server of type %s",
           ServerType_Name(config.server_type()).c_str());
 
-  if (config.core_limit() > 0) {
-    LimitCores(config.core_limit());
-  }
-
   switch (config.server_type()) {
     case ServerType::SYNC_SERVER:
       return CreateSynchronousServer(config);
diff --git a/test/cpp/qps/server.h b/test/cpp/qps/server.h
index bc6f9f99e3..94a6f8acfa 100644
--- a/test/cpp/qps/server.h
+++ b/test/cpp/qps/server.h
@@ -51,18 +51,10 @@ namespace testing {
 class Server {
  public:
   explicit Server(const ServerConfig& config) : timer_(new Timer) {
-    int clsize = config.core_list_size();
-    if (clsize > 0) {
-      std::vector<int> core_list;
-      for (int i = 0; i < clsize; i++) {
-        core_list.push_back(config.core_list(i));
-      }
-      cores_ = LimitCores(core_list);
-    } else {
-      cores_ = gpr_cpu_num_cores();
-    }
+    cores_ = LimitCores(config.core_list().data(), config.core_list_size());
     if (config.port()) {
       port_ = config.port();
+
     } else {
       port_ = grpc_pick_unused_port_or_die();
     }