Initial interface rework to allow knowing whether to pull payload at registration, not at request time

1 files changed, 16 insertions, 2 deletions

diff --git a/src/cpp/server/server.cc b/src/cpp/server/server.cc
index 3e44c502fc..fafe31e84c 100644
--- a/src/cpp/server/server.cc
+++ b/src/cpp/server/server.cc
@@ -321,6 +321,19 @@ void Server::SetGlobalCallbacks(GlobalCallbacks* callbacks) {
   g_callbacks.reset(callbacks);
 }
 
+static grpc_server_register_method_payload_handling PayloadHandlingForMethod(
+    RpcServiceMethod* method) {
+  switch (method->method_type()) {
+    case RpcMethod::NORMAL_RPC:
+    case RpcMethod::SERVER_STREAMING:
+      return GRPC_SRM_PAYLOAD_READ_INITIAL_BYTE_BUFFER;
+    case RpcMethod::CLIENT_STREAMING:
+    case RpcMethod::BIDI_STREAMING:
+      return GRPC_SRM_PAYLOAD_NONE;
+  }
+  GPR_UNREACHABLE_CODE(return GRPC_SRM_PAYLOAD_NONE;);
+}
+
 bool Server::RegisterService(const grpc::string* host, Service* service) {
   bool has_async_methods = service->has_async_methods();
   if (has_async_methods) {
@@ -334,8 +347,9 @@ bool Server::RegisterService(const grpc::string* host, Service* service) {
       continue;
     }
     RpcServiceMethod* method = it->get();
-    void* tag = grpc_server_register_method(server_, method->name(),
-                                            host ? host->c_str() : nullptr, 0);
+    void* tag = grpc_server_register_method(
+        server_, method->name(), host ? host->c_str() : nullptr,
+        PayloadHandlingForMethod(method), 0);
     if (tag == nullptr) {
       gpr_log(GPR_DEBUG, "Attempt to register %s multiple times",
               method->name());