1 files changed, 24 insertions, 4 deletions
diff --git a/grpc/src/cpp/server/server_builder.cc b/grpc/src/cpp/server/server_builder.cc
index dc381257..695aa541 100644
--- a/grpc/src/cpp/server/server_builder.cc
+++ b/grpc/src/cpp/server/server_builder.cc
@@ -130,6 +130,12 @@ ServerBuilder& ServerBuilder::experimental_type::RegisterCallbackGenericService(
 }
 #endif
 
+ServerBuilder& ServerBuilder::experimental_type::SetContextAllocator(
+    std::unique_ptr<grpc::ContextAllocator> context_allocator) {
+  builder_->context_allocator_ = std::move(context_allocator);
+  return *builder_;
+}
+
 std::unique_ptr<grpc::experimental::ExternalConnectionAcceptor>
 ServerBuilder::experimental_type::AddExternalConnectionAcceptor(
     experimental_type::ExternalConnectionType type,
@@ -217,8 +223,8 @@ ServerBuilder& ServerBuilder::AddListeningPort(
   return *this;
 }
 
-std::unique_ptr<grpc::Server> ServerBuilder::BuildAndStart() {
-  grpc::ChannelArguments args;
+ChannelArguments ServerBuilder::BuildChannelArgs() {
+  ChannelArguments args;
   if (max_receive_message_size_ >= -1) {
     args.SetInt(GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH, max_receive_message_size_);
   }
@@ -239,16 +245,19 @@ std::unique_ptr<grpc::Server> ServerBuilder::BuildAndStart() {
     args.SetInt(GRPC_COMPRESSION_CHANNEL_DEFAULT_ALGORITHM,
                 maybe_default_compression_algorithm_.algorithm);
   }
-
   if (resource_quota_ != nullptr) {
     args.SetPointerWithVtable(GRPC_ARG_RESOURCE_QUOTA, resource_quota_,
                               grpc_resource_quota_arg_vtable());
   }
-
   for (const auto& plugin : plugins_) {
     plugin->UpdateServerBuilder(this);
     plugin->UpdateChannelArguments(&args);
   }
+  return args;
+}
+
+std::unique_ptr<grpc::Server> ServerBuilder::BuildAndStart() {
+  ChannelArguments args = BuildChannelArgs();
 
   // == Determine if the server has any syncrhonous methods ==
   bool has_sync_methods = false;
@@ -298,6 +307,10 @@ std::unique_ptr<grpc::Server> ServerBuilder::BuildAndStart() {
     }
   }
 
+  if (callback_generic_service_ != nullptr) {
+    has_frequently_polled_cqs = true;
+  }
+
   const bool is_hybrid_server = has_sync_methods && has_frequently_polled_cqs;
 
   if (has_sync_methods) {
@@ -369,6 +382,13 @@ std::unique_ptr<grpc::Server> ServerBuilder::BuildAndStart() {
     return nullptr;
   }
 
+#ifdef GRPC_CALLBACK_API_NONEXPERIMENTAL
+  server->RegisterContextAllocator(std::move(context_allocator_));
+#else
+  server->experimental_registration()->RegisterContextAllocator(
+      std::move(context_allocator_));
+#endif
+
   for (const auto& value : services_) {
     if (!server->RegisterService(value->host.get(), value->service)) {
       return nullptr;