1 files changed, 116 insertions, 15 deletions

diff --git a/test/cpp/qps/driver.cc b/test/cpp/qps/driver.cc
index 490156aec2..c70b0303b8 100644
--- a/test/cpp/qps/driver.cc
+++ b/test/cpp/qps/driver.cc
@@ -34,6 +34,7 @@
 #include <deque>
 #include <list>
 #include <thread>
+#include <unordered_map>
 #include <vector>
 
 #include <grpc++/channel.h>
@@ -59,7 +60,42 @@ using std::vector;
 
 namespace grpc {
 namespace testing {
-static deque<string> get_hosts(const string& name) {
+static std::string get_host(const std::string& worker) {
+  char* host;
+  char* port;
+
+  gpr_split_host_port(worker.c_str(), &host, &port);
+  const string s(host);
+
+  gpr_free(host);
+  gpr_free(port);
+  return s;
+}
+
+static std::unordered_map<string, std::deque<int>> get_hosts_and_cores(
+    const deque<string>& workers) {
+  std::unordered_map<string, std::deque<int>> hosts;
+  for (auto it = workers.begin(); it != workers.end(); it++) {
+    const string host = get_host(*it);
+    if (hosts.find(host) == hosts.end()) {
+      auto stub = WorkerService::NewStub(
+          CreateChannel(*it, InsecureChannelCredentials()));
+      grpc::ClientContext ctx;
+      CoreRequest dummy;
+      CoreResponse cores;
+      grpc::Status s = stub->CoreCount(&ctx, dummy, &cores);
+      assert(s.ok());
+      std::deque<int> dq;
+      for (int i = 0; i < cores.cores(); i++) {
+        dq.push_back(i);
+      }
+      hosts[host] = dq;
+    }
+  }
+  return hosts;
+}
+
+static deque<string> get_workers(const string& name) {
   char* env = gpr_getenv(name.c_str());
   if (!env) return deque<string>();
 
@@ -105,18 +141,18 @@ struct ClientData {
 
 std::unique_ptr<ScenarioResult> RunScenario(
     const ClientConfig& initial_client_config, size_t num_clients,
-    const ServerConfig& server_config, size_t num_servers, int warmup_seconds,
-    int benchmark_seconds, int spawn_local_worker_count) {
+    const ServerConfig& initial_server_config, size_t num_servers,
+    int warmup_seconds, int benchmark_seconds, int spawn_local_worker_count) {
   // ClientContext allocations (all are destroyed at scope exit)
   list<ClientContext> contexts;
 
   // To be added to the result, containing the final configuration used for
   // client and config (including host, etc.)
   ClientConfig result_client_config;
-  ServerConfig result_server_config;
+  const ServerConfig result_server_config = initial_server_config;
 
   // Get client, server lists
-  auto workers = get_hosts("QPS_WORKERS");
+  auto workers = get_workers("QPS_WORKERS");
   ClientConfig client_config = initial_client_config;
 
   // Spawn some local workers if desired
@@ -143,6 +179,9 @@ std::unique_ptr<ScenarioResult> RunScenario(
     }
   }
 
+  // Setup the hosts and core counts
+  auto hosts_cores = get_hosts_and_cores(workers);
+
   // if num_clients is set to <=0, do dynamic sizing: all workers
   // except for servers are clients
   if (num_clients <= 0) {
@@ -172,18 +211,49 @@ std::unique_ptr<ScenarioResult> RunScenario(
             i);
     servers[i].stub = WorkerService::NewStub(
         CreateChannel(workers[i], InsecureChannelCredentials()));
+
+    ServerConfig server_config = initial_server_config;
+    char* host;
+    char* driver_port;
+    char* cli_target;
+    gpr_split_host_port(workers[i].c_str(), &host, &driver_port);
+    string host_str(host);
+    int server_core_limit = initial_server_config.core_limit();
+    int client_core_limit = initial_client_config.core_limit();
+
+    if (server_core_limit == 0 && client_core_limit > 0) {
+      // In this case, limit the server cores if it matches the
+      // same host as one or more clients
+      const auto& dq = hosts_cores.at(host_str);
+      bool match = false;
+      int limit = dq.size();
+      for (size_t cli = 0; cli < num_clients; cli++) {
+        if (host_str == get_host(workers[cli + num_servers])) {
+          limit -= client_core_limit;
+          match = true;
+        }
+      }
+      if (match) {
+        GPR_ASSERT(limit > 0);
+        server_core_limit = limit;
+      }
+    }
+    if (server_core_limit > 0) {
+      auto& dq = hosts_cores.at(host_str);
+      GPR_ASSERT(dq.size() >= static_cast<size_t>(server_core_limit));
+      for (int core = 0; core < server_core_limit; core++) {
+        server_config.add_core_list(dq.front());
+        dq.pop_front();
+      }
+    }
+
     ServerArgs args;
-    result_server_config = server_config;
     *args.mutable_setup() = server_config;
     servers[i].stream =
         servers[i].stub->RunServer(runsc::AllocContext(&contexts, deadline));
     GPR_ASSERT(servers[i].stream->Write(args));
     ServerStatus init_status;
     GPR_ASSERT(servers[i].stream->Read(&init_status));
-    char* host;
-    char* driver_port;
-    char* cli_target;
-    gpr_split_host_port(workers[i].c_str(), &host, &driver_port);
     gpr_join_host_port(&cli_target, host, init_status.port());
     client_config.add_server_targets(cli_target);
     gpr_free(host);
@@ -191,19 +261,50 @@ std::unique_ptr<ScenarioResult> RunScenario(
     gpr_free(cli_target);
   }
 
+  // Targets are all set by now
+  result_client_config = client_config;
   // Start clients
   using runsc::ClientData;
   // clients is array rather than std::vector to avoid gcc-4.4 issues
   // where class contained in std::vector must have a copy constructor
   auto* clients = new ClientData[num_clients];
   for (size_t i = 0; i < num_clients; i++) {
-    gpr_log(GPR_INFO, "Starting client on %s (worker #%d)",
-            workers[i + num_servers].c_str(), i + num_servers);
+    const auto& worker = workers[i + num_servers];
+    gpr_log(GPR_INFO, "Starting client on %s (worker #%d)", worker.c_str(),
+            i + num_servers);
     clients[i].stub = WorkerService::NewStub(
-        CreateChannel(workers[i + num_servers], InsecureChannelCredentials()));
+        CreateChannel(worker, InsecureChannelCredentials()));
+    ClientConfig per_client_config = client_config;
+
+    int server_core_limit = initial_server_config.core_limit();
+    int client_core_limit = initial_client_config.core_limit();
+    if ((server_core_limit > 0) || (client_core_limit > 0)) {
+      auto& dq = hosts_cores.at(get_host(worker));
+      if (client_core_limit == 0) {
+        // limit client cores if it matches a server host
+        bool match = false;
+        int limit = dq.size();
+        for (size_t srv = 0; srv < num_servers; srv++) {
+          if (get_host(worker) == get_host(workers[srv])) {
+            match = true;
+          }
+        }
+        if (match) {
+          GPR_ASSERT(limit > 0);
+          client_core_limit = limit;
+        }
+      }
+      if (client_core_limit > 0) {
+        GPR_ASSERT(dq.size() >= static_cast<size_t>(client_core_limit));
+        for (int core = 0; core < client_core_limit; core++) {
+          per_client_config.add_core_list(dq.front());
+          dq.pop_front();
+        }
+      }
+    }
+
     ClientArgs args;
-    result_client_config = client_config;
-    *args.mutable_setup() = client_config;
+    *args.mutable_setup() = per_client_config;
     clients[i].stream =
         clients[i].stub->RunClient(runsc::AllocContext(&contexts, deadline));
     GPR_ASSERT(clients[i].stream->Write(args));