Make CreateThreadPool settable

sheenaqotj · Jul 19, 2017 · 30939f8 · 30939f8
1 parent 10bed24
commit 30939f8
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/src/cpp/client/secure_credentials.cc b/src/cpp/client/secure_credentials.cc
@@ -207,7 +207,7 @@ void MetadataCredentialsPluginWrapper::InvokePlugin(
 
 MetadataCredentialsPluginWrapper::MetadataCredentialsPluginWrapper(
     std::unique_ptr<MetadataCredentialsPlugin> plugin)
-    : thread_pool_(CreateDefaultThreadPool()), plugin_(std::move(plugin)) {}
+    : thread_pool_(CreateThreadPool()), plugin_(std::move(plugin)) {}
 
 std::shared_ptr<CallCredentials> MetadataCredentialsFromPlugin(
     std::unique_ptr<MetadataCredentialsPlugin> plugin) {

diff --git a/src/cpp/server/create_default_thread_pool.cc b/src/cpp/server/create_default_thread_pool.cc
@@ -24,12 +24,22 @@
 
 namespace grpc {
 
-ThreadPoolInterface* CreateDefaultThreadPool() {
+static ThreadPoolInterface* CreateDefaultThreadPool() {
   int cores = gpr_cpu_num_cores();
   if (!cores) cores = 4;
   return new DynamicThreadPool(cores);
 }
 
+static CreateThreadPoolFunc g_ctp_impl = CreateDefaultThreadPool;
+
+ThreadPoolInterface* CreateThreadPool() {
+  return g_ctp_impl();
+}
+
+void SetCreateThreadPool(CreateThreadPoolFunc func) {
+  g_ctp_impl = func;
+}
+
 }  // namespace grpc
 
 #endif  // !GRPC_CUSTOM_DEFAULT_THREAD_POOL
diff --git a/src/cpp/server/secure_server_credentials.h b/src/cpp/server/secure_server_credentials.h
@@ -39,7 +39,7 @@ class AuthMetadataProcessorAyncWrapper final {
 
   AuthMetadataProcessorAyncWrapper(
       const std::shared_ptr<AuthMetadataProcessor>& processor)
-      : thread_pool_(CreateDefaultThreadPool()), processor_(processor) {}
+      : thread_pool_(CreateThreadPool()), processor_(processor) {}
 
  private:
   void InvokeProcessor(grpc_auth_context* context, const grpc_metadata* md,

diff --git a/src/cpp/server/thread_pool_interface.h b/src/cpp/server/thread_pool_interface.h
@@ -32,7 +32,11 @@ class ThreadPoolInterface {
   virtual void Add(const std::function<void()>& callback) = 0;
 };
 
-ThreadPoolInterface* CreateDefaultThreadPool();
+// Allows different codebases to use their own thread pool impls
+typedef ThreadPoolInterface* (*CreateThreadPoolFunc)(void);
+void SetCreateThreadPool(CreateThreadPoolFunc func);
+
+ThreadPoolInterface* CreateThreadPool();
 
 }  // namespace grpc