Progress

2 files changed, 36 insertions, 22 deletions

diff --git a/src/cpp/server/server.cc b/src/cpp/server/server.cc
index 1abdf702e2..ac1b9ddc50 100644
--- a/src/cpp/server/server.cc
+++ b/src/cpp/server/server.cc
@@ -38,24 +38,20 @@
 #include <grpc/grpc_security.h>
 #include <grpc/support/log.h>
 #include "src/cpp/server/server_rpc_handler.h"
-#include "src/cpp/server/thread_pool.h"
 #include <grpc++/async_server_context.h>
 #include <grpc++/completion_queue.h>
 #include <grpc++/impl/rpc_service_method.h>
 #include <grpc++/server_credentials.h>
+#include <grpc++/thread_pool_interface.h>
 
 namespace grpc {
 
-// TODO(rocking): consider a better default value like num of cores.
-static const int kNumThreads = 4;
-
-Server::Server(ThreadPoolInterface *thread_pool, ServerCredentials *creds)
+Server::Server(ThreadPoolInterface *thread_pool, bool thread_pool_owned, ServerCredentials *creds)
     : started_(false),
       shutdown_(false),
       num_running_cb_(0),
-      thread_pool_(thread_pool == nullptr ? new ThreadPool(kNumThreads)
-                                          : thread_pool),
-      thread_pool_owned_(thread_pool == nullptr),
+      thread_pool_(thread_pool),
+      thread_pool_owned_(thread_pool_owned),
       secure_(creds != nullptr) {
   if (creds) {
     server_ =
@@ -82,31 +78,35 @@ Server::~Server() {
   }
 }
 
-void Server::RegisterService(RpcService *service) {
+bool Server::RegisterService(RpcService *service) {
   for (int i = 0; i < service->GetMethodCount(); ++i) {
     RpcServiceMethod *method = service->GetMethod(i);
+    if (method_map_.find(method->name()) != method_map_.end()) {
+      gpr_log(GPR_DEBUG, "Attempt to register %s multiple times", method->name());
+      return false;
+    }
     method_map_.insert(std::make_pair(method->name(), method));
   }
 }
 
-void Server::AddPort(const grpc::string &addr) {
+int Server::AddPort(const grpc::string &addr) {
   GPR_ASSERT(!started_);
-  int success;
   if (secure_) {
-    success = grpc_server_add_secure_http2_port(server_, addr.c_str());
+    return grpc_server_add_secure_http2_port(server_, addr.c_str());
   } else {
-    success = grpc_server_add_http2_port(server_, addr.c_str());
+    return grpc_server_add_http2_port(server_, addr.c_str());
   }
-  GPR_ASSERT(success);
 }
 
-void Server::Start() {
+bool Server::Start() {
   GPR_ASSERT(!started_);
   started_ = true;
   grpc_server_start(server_);
 
   // Start processing rpcs.
   ScheduleCallback();
+
+  return true;
 }
 
 void Server::AllowOneRpc() {
@@ -141,8 +141,7 @@ void Server::ScheduleCallback() {
     std::unique_lock<std::mutex> lock(mu_);
     num_running_cb_++;
   }
-  std::function<void()> callback = std::bind(&Server::RunRpc, this);
-  thread_pool_->ScheduleCallback(callback);
+  thread_pool_->ScheduleCallback(std::bind(&Server::RunRpc, this));
 }
 
 void Server::RunRpc() {
diff --git a/src/cpp/server/server_builder.cc b/src/cpp/server/server_builder.cc
index 66e2055af0..8d8276ca00 100644
--- a/src/cpp/server/server_builder.cc
+++ b/src/cpp/server/server_builder.cc
@@ -33,13 +33,15 @@
 
 #include <grpc++/server_builder.h>
 
+#include <grpc/support/cpu.h>
 #include <grpc/support/log.h>
 #include <grpc++/impl/service_type.h>
 #include <grpc++/server.h>
+#include "src/cpp/server/thread_pool.h"
 
 namespace grpc {
 
-ServerBuilder::ServerBuilder() : thread_pool_(nullptr) {}
+ServerBuilder::ServerBuilder() {}
 
 void ServerBuilder::RegisterService(SynchronousService *service) {
   services_.push_back(service->service());
@@ -64,14 +66,27 @@ void ServerBuilder::SetThreadPool(ThreadPoolInterface *thread_pool) {
 }
 
 std::unique_ptr<Server> ServerBuilder::BuildAndStart() {
-  std::unique_ptr<Server> server(new Server(thread_pool_, creds_.get()));
+  bool thread_pool_owned = false;
+  if (!thread_pool_ && services_.size()) {
+    int cores = gpr_cpu_num_cores();
+    if (!cores) cores = 4;
+    thread_pool_ = new ThreadPool(cores);
+    thread_pool_owned = true;
+  }
+  std::unique_ptr<Server> server(new Server(thread_pool_, thread_pool_owned, creds_.get()));
   for (auto *service : services_) {
-    server->RegisterService(service);
+    if (!server->RegisterService(service)) {
+      return nullptr;
+    }
   }
   for (auto &port : ports_) {
-    server->AddPort(port);
+    if (!server->AddPort(port)) {
+      return nullptr;
+    }
+  }
+  if (!server->Start()) {
+    return nullptr;
   }
-  server->Start();
   return server;
 }