Experimental C++ server callback unary API

3 files changed, 298 insertions, 42 deletions

diff --git a/src/cpp/server/server_builder.cc b/src/cpp/server/server_builder.cc
index fc42b6c886..0dc03b6876 100644
--- a/src/cpp/server/server_builder.cc
+++ b/src/cpp/server/server_builder.cc
@@ -71,7 +71,9 @@ ServerBuilder::~ServerBuilder() {
 std::unique_ptr<ServerCompletionQueue> ServerBuilder::AddCompletionQueue(
     bool is_frequently_polled) {
   ServerCompletionQueue* cq = new ServerCompletionQueue(
-      is_frequently_polled ? GRPC_CQ_DEFAULT_POLLING : GRPC_CQ_NON_LISTENING);
+      GRPC_CQ_NEXT,
+      is_frequently_polled ? GRPC_CQ_DEFAULT_POLLING : GRPC_CQ_NON_LISTENING,
+      nullptr);
   cqs_.push_back(cq);
   return std::unique_ptr<ServerCompletionQueue>(cq);
 }
@@ -256,15 +258,22 @@ std::unique_ptr<Server> ServerBuilder::BuildAndStart() {
 
     // Create completion queues to listen to incoming rpc requests
     for (int i = 0; i < sync_server_settings_.num_cqs; i++) {
-      sync_server_cqs->emplace_back(new ServerCompletionQueue(polling_type));
+      sync_server_cqs->emplace_back(
+          new ServerCompletionQueue(GRPC_CQ_NEXT, polling_type, nullptr));
     }
   }
 
-  std::unique_ptr<Server> server(new Server(
-      max_receive_message_size_, &args, sync_server_cqs,
-      sync_server_settings_.min_pollers, sync_server_settings_.max_pollers,
-      sync_server_settings_.cq_timeout_msec, resource_quota_,
-      std::move(interceptor_creators_)));
+  // == Determine if the server has any callback methods ==
+  bool has_callback_methods = false;
+  for (auto it = services_.begin(); it != services_.end(); ++it) {
+    if ((*it)->service->has_callback_methods()) {
+      has_callback_methods = true;
+      break;
+    }
+  }
+
+  // TODO(vjpai): Add a section here for plugins once they can support callback
+  // methods
 
   if (has_sync_methods) {
     // This is a Sync server
@@ -276,6 +285,16 @@ std::unique_ptr<Server> ServerBuilder::BuildAndStart() {
             sync_server_settings_.cq_timeout_msec);
   }
 
+  if (has_callback_methods) {
+    gpr_log(GPR_INFO, "Callback server.");
+  }
+
+  std::unique_ptr<Server> server(new Server(
+      max_receive_message_size_, &args, sync_server_cqs,
+      sync_server_settings_.min_pollers, sync_server_settings_.max_pollers,
+      sync_server_settings_.cq_timeout_msec, resource_quota_,
+      std::move(interceptor_creators_)));
+
   ServerInitializer* initializer = server->initializer();
 
   // Register all the completion queues with the server. i.e
@@ -289,6 +308,12 @@ std::unique_ptr<Server> ServerBuilder::BuildAndStart() {
     num_frequently_polled_cqs++;
   }
 
+  if (has_callback_methods) {
+    auto* cq = server->CallbackCQ();
+    grpc_server_register_completion_queue(server->server_, cq->cq(), nullptr);
+    num_frequently_polled_cqs++;
+  }
+
   // cqs_ contains the completion queue added by calling the ServerBuilder's
   // AddCompletionQueue() API. Some of them may not be frequently polled (i.e by
   // calling Next() or AsyncNext()) and hence are not safe to be used for
diff --git a/src/cpp/server/server_cc.cc b/src/cpp/server/server_cc.cc
index 82a9d719fa..5d24e01a47 100644
--- a/src/cpp/server/server_cc.cc
+++ b/src/cpp/server/server_cc.cc
@@ -147,9 +147,9 @@ class Server::UnimplementedAsyncResponse final
 
 class Server::SyncRequest final : public internal::CompletionQueueTag {
  public:
-  SyncRequest(internal::RpcServiceMethod* method, void* tag)
+  SyncRequest(internal::RpcServiceMethod* method, void* method_tag)
       : method_(method),
-        tag_(tag),
+        method_tag_(method_tag),
         in_flight_(false),
         has_request_payload_(
             method->method_type() == internal::RpcMethod::NORMAL_RPC ||
@@ -176,10 +176,10 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
   void Request(grpc_server* server, grpc_completion_queue* notify_cq) {
     GPR_ASSERT(cq_ && !in_flight_);
     in_flight_ = true;
-    if (tag_) {
+    if (method_tag_) {
       if (GRPC_CALL_OK !=
           grpc_server_request_registered_call(
-              server, tag_, &call_, &deadline_, &request_metadata_,
+              server, method_tag_, &call_, &deadline_, &request_metadata_,
               has_request_payload_ ? &request_payload_ : nullptr, cq_,
               notify_cq, this)) {
         TeardownRequest();
@@ -211,6 +211,9 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
     return true;
   }
 
+  // The CallData class represents a call that is "active" as opposed
+  // to just being requested. It wraps and takes ownership of the cq from
+  // the call request
   class CallData final {
    public:
     explicit CallData(Server* server, SyncRequest* mrd)
@@ -281,7 +284,7 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
         auto* handler = resources_ ? method_->handler()
                                    : server_->resource_exhausted_handler_.get();
         handler->RunHandler(internal::MethodHandler::HandlerParameter(
-            &call_, &ctx_, request_, request_status_));
+            &call_, &ctx_, request_, request_status_, nullptr));
         request_ = nullptr;
         global_callbacks_->PostSynchronousRequest(&ctx_);
 
@@ -314,7 +317,7 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
 
  private:
   internal::RpcServiceMethod* const method_;
-  void* const tag_;
+  void* const method_tag_;
   bool in_flight_;
   const bool has_request_payload_;
   grpc_call* call_;
@@ -325,6 +328,176 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
   grpc_completion_queue* cq_;
 };
 
+class Server::CallbackRequest final : public internal::CompletionQueueTag {
+ public:
+  CallbackRequest(Server* server, internal::RpcServiceMethod* method,
+                  void* method_tag)
+      : server_(server),
+        method_(method),
+        method_tag_(method_tag),
+        has_request_payload_(
+            method->method_type() == internal::RpcMethod::NORMAL_RPC ||
+            method->method_type() == internal::RpcMethod::SERVER_STREAMING),
+        cq_(server->CallbackCQ()),
+        tag_(this) {
+    Setup();
+  }
+
+  ~CallbackRequest() { Clear(); }
+
+  void Request() {
+    if (method_tag_) {
+      if (GRPC_CALL_OK !=
+          grpc_server_request_registered_call(
+              server_->c_server(), method_tag_, &call_, &deadline_,
+              &request_metadata_,
+              has_request_payload_ ? &request_payload_ : nullptr, cq_->cq(),
+              cq_->cq(), static_cast<void*>(&tag_))) {
+        return;
+      }
+    } else {
+      if (!call_details_) {
+        call_details_ = new grpc_call_details;
+        grpc_call_details_init(call_details_);
+      }
+      if (grpc_server_request_call(server_->c_server(), &call_, call_details_,
+                                   &request_metadata_, cq_->cq(), cq_->cq(),
+                                   static_cast<void*>(&tag_)) != GRPC_CALL_OK) {
+        return;
+      }
+    }
+  }
+
+  bool FinalizeResult(void** tag, bool* status) override { return false; }
+
+ private:
+  class CallbackCallTag : public grpc_experimental_completion_queue_functor {
+   public:
+    CallbackCallTag(Server::CallbackRequest* req) : req_(req) {
+      functor_run = &CallbackCallTag::StaticRun;
+    }
+
+    // force_run can not be performed on a tag if operations using this tag
+    // have been sent to PerformOpsOnCall. It is intended for error conditions
+    // that are detected before the operations are internally processed.
+    void force_run(bool ok) { Run(ok); }
+
+   private:
+    Server::CallbackRequest* req_;
+    internal::Call* call_;
+
+    static void StaticRun(grpc_experimental_completion_queue_functor* cb,
+                          int ok) {
+      static_cast<CallbackCallTag*>(cb)->Run(static_cast<bool>(ok));
+    }
+    void Run(bool ok) {
+      void* ignored = req_;
+      bool new_ok = ok;
+      GPR_ASSERT(!req_->FinalizeResult(&ignored, &new_ok));
+      GPR_ASSERT(ignored == req_);
+
+      if (!ok) {
+        // The call has been shutdown
+        req_->Clear();
+        return;
+      }
+
+      // Bind the call, deadline, and metadata from what we got
+      req_->ctx_.set_call(req_->call_);
+      req_->ctx_.cq_ = req_->cq_;
+      req_->ctx_.BindDeadlineAndMetadata(req_->deadline_,
+                                         &req_->request_metadata_);
+      req_->request_metadata_.count = 0;
+
+      // Create a C++ Call to control the underlying core call
+      call_ = new (grpc_call_arena_alloc(req_->call_, sizeof(internal::Call)))
+          internal::Call(
+              req_->call_, req_->server_, req_->cq_,
+              req_->server_->max_receive_message_size(),
+              req_->ctx_.set_server_rpc_info(
+                  req_->method_->name(), req_->server_->interceptor_creators_));
+
+      req_->interceptor_methods_.SetCall(call_);
+      req_->interceptor_methods_.SetReverse();
+      // Set interception point for RECV INITIAL METADATA
+      req_->interceptor_methods_.AddInterceptionHookPoint(
+          experimental::InterceptionHookPoints::POST_RECV_INITIAL_METADATA);
+      req_->interceptor_methods_.SetRecvInitialMetadata(
+          &req_->ctx_.client_metadata_);
+
+      if (req_->has_request_payload_) {
+        // Set interception point for RECV MESSAGE
+        req_->request_ = req_->method_->handler()->Deserialize(
+            req_->call_, req_->request_payload_, &req_->request_status_);
+        req_->request_payload_ = nullptr;
+        req_->interceptor_methods_.AddInterceptionHookPoint(
+            experimental::InterceptionHookPoints::POST_RECV_MESSAGE);
+        req_->interceptor_methods_.SetRecvMessage(req_->request_);
+      }
+
+      if (req_->interceptor_methods_.RunInterceptors(
+              [this] { ContinueRunAfterInterception(); })) {
+        ContinueRunAfterInterception();
+      } else {
+        // There were interceptors to be run, so ContinueRunAfterInterception
+        // will be run when interceptors are done.
+      }
+    }
+    void ContinueRunAfterInterception() {
+      // req_->ctx_.BeginCompletionOp(call_);
+      req_->method_->handler()->RunHandler(
+          internal::MethodHandler::HandlerParameter(
+              call_, &req_->ctx_, req_->request_, req_->request_status_,
+              [this] {
+                req_->Reset();
+                req_->Request();
+              }));
+    }
+  };
+
+  void Reset() {
+    Clear();
+    Setup();
+  }
+
+  void Clear() {
+    if (call_details_) {
+      delete call_details_;
+      call_details_ = nullptr;
+    }
+    grpc_metadata_array_destroy(&request_metadata_);
+    if (has_request_payload_ && request_payload_) {
+      grpc_byte_buffer_destroy(request_payload_);
+    }
+    ctx_.Clear();
+    interceptor_methods_.ClearState();
+  }
+
+  void Setup() {
+    grpc_metadata_array_init(&request_metadata_);
+    ctx_.Setup(gpr_inf_future(GPR_CLOCK_REALTIME));
+    request_payload_ = nullptr;
+    request_ = nullptr;
+    request_status_ = Status();
+  }
+
+  Server* const server_;
+  internal::RpcServiceMethod* const method_;
+  void* const method_tag_;
+  const bool has_request_payload_;
+  grpc_byte_buffer* request_payload_;
+  void* request_;
+  Status request_status_;
+  grpc_call_details* call_details_ = nullptr;
+  grpc_call* call_;
+  gpr_timespec deadline_;
+  grpc_metadata_array request_metadata_;
+  CompletionQueue* cq_;
+  CallbackCallTag tag_;
+  ServerContext ctx_;
+  internal::InterceptorBatchMethodsImpl interceptor_methods_;
+};
+
 // Implementation of ThreadManager. Each instance of SyncRequestThreadManager
 // manages a pool of threads that poll for incoming Sync RPCs and call the
 // appropriate RPC handlers
@@ -504,6 +677,9 @@ Server::Server(
 Server::~Server() {
   {
     std::unique_lock<std::mutex> lock(mu_);
+    if (callback_cq_ != nullptr) {
+      callback_cq_->Shutdown();
+    }
     if (started_ && !shutdown_) {
       lock.unlock();
       Shutdown();
@@ -576,21 +752,28 @@ bool Server::RegisterService(const grpc::string* host, Service* service) {
     }
 
     internal::RpcServiceMethod* method = it->get();
-    void* tag = grpc_server_register_method(
+    void* method_registration_tag = grpc_server_register_method(
         server_, method->name(), host ? host->c_str() : nullptr,
         PayloadHandlingForMethod(method), 0);
-    if (tag == nullptr) {
+    if (method_registration_tag == nullptr) {
       gpr_log(GPR_DEBUG, "Attempt to register %s multiple times",
               method->name());
       return false;
     }
 
-    if (method->handler() == nullptr) {  // Async method
-      method->set_server_tag(tag);
-    } else {
+    if (method->handler() == nullptr) {  // Async method without handler
+      method->set_server_tag(method_registration_tag);
+    } else if (method->api_type() ==
+               internal::RpcServiceMethod::ApiType::SYNC) {
       for (auto it = sync_req_mgrs_.begin(); it != sync_req_mgrs_.end(); it++) {
-        (*it)->AddSyncMethod(method, tag);
+        (*it)->AddSyncMethod(method, method_registration_tag);
       }
+    } else {
+      // a callback method
+      auto* req = new CallbackRequest(this, method, method_registration_tag);
+      callback_reqs_.emplace_back(req);
+      // Enqueue it so that it will be Request'ed later once
+      // all request matchers are created at core server startup
     }
 
     method_name = method->name();
@@ -641,7 +824,8 @@ void Server::Start(ServerCompletionQueue** cqs, size_t num_cqs) {
     // performance.  This ensures that we don't introduce thread hops
     // for application requests that wind up on this CQ, which is polled
     // in its own thread.
-    health_check_cq = new ServerCompletionQueue(GRPC_CQ_NON_POLLING);
+    health_check_cq =
+        new ServerCompletionQueue(GRPC_CQ_NEXT, GRPC_CQ_NON_POLLING, nullptr);
     grpc_server_register_completion_queue(server_, health_check_cq->cq(),
                                           nullptr);
     default_health_check_service_impl =
@@ -678,6 +862,10 @@ void Server::Start(ServerCompletionQueue** cqs, size_t num_cqs) {
     (*it)->Start();
   }
 
+  for (auto& cbreq : callback_reqs_) {
+    cbreq->Request();
+  }
+
   if (default_health_check_service_impl != nullptr) {
     default_health_check_service_impl->StartServingThread();
   }
@@ -910,4 +1098,41 @@ Server::UnimplementedAsyncResponse::UnimplementedAsyncResponse(
 
 ServerInitializer* Server::initializer() { return server_initializer_.get(); }
 
+namespace {
+class ShutdownCallback : public grpc_experimental_completion_queue_functor {
+ public:
+  ShutdownCallback() { functor_run = &ShutdownCallback::Run; }
+  // TakeCQ takes ownership of the cq into the shutdown callback
+  // so that the shutdown callback will be responsible for destroying it
+  void TakeCQ(CompletionQueue* cq) { cq_ = cq; }
+
+  // The Run function will get invoked by the completion queue library
+  // when the shutdown is actually complete
+  static void Run(grpc_experimental_completion_queue_functor* cb, int) {
+    auto* callback = static_cast<ShutdownCallback*>(cb);
+    delete callback->cq_;
+    grpc_core::Delete(callback);
+  }
+
+ private:
+  CompletionQueue* cq_ = nullptr;
+};
+}  // namespace
+
+CompletionQueue* Server::CallbackCQ() {
+  // TODO(vjpai): Consider using a single global CQ for the default CQ
+  // if there is no explicit per-server CQ registered
+  std::lock_guard<std::mutex> l(mu_);
+  if (callback_cq_ == nullptr) {
+    auto* shutdown_callback = grpc_core::New<ShutdownCallback>();
+    callback_cq_ = new CompletionQueue(grpc_completion_queue_attributes{
+        GRPC_CQ_CURRENT_VERSION, GRPC_CQ_CALLBACK, GRPC_CQ_DEFAULT_POLLING,
+        shutdown_callback});
+
+    // Transfer ownership of the new cq to its own shutdown callback
+    shutdown_callback->TakeCQ(callback_cq_);
+  }
+  return callback_cq_;
+};
+
 }  // namespace grpc
diff --git a/src/cpp/server/server_context.cc b/src/cpp/server/server_context.cc
index 995e787785..51a2689c6a 100644
--- a/src/cpp/server/server_context.cc
+++ b/src/cpp/server/server_context.cc
@@ -209,31 +209,35 @@ bool ServerContext::CompletionOp::FinalizeResult(void** tag, bool* status) {
 
 // ServerContext body
 
-ServerContext::ServerContext()
-    : completion_op_(nullptr),
-      has_notify_when_done_tag_(false),
-      async_notify_when_done_tag_(nullptr),
-      deadline_(gpr_inf_future(GPR_CLOCK_REALTIME)),
-      call_(nullptr),
-      cq_(nullptr),
-      sent_initial_metadata_(false),
-      compression_level_set_(false),
-      has_pending_ops_(false) {}
-
-ServerContext::ServerContext(gpr_timespec deadline, grpc_metadata_array* arr)
-    : completion_op_(nullptr),
-      has_notify_when_done_tag_(false),
-      async_notify_when_done_tag_(nullptr),
-      deadline_(deadline),
-      call_(nullptr),
-      cq_(nullptr),
-      sent_initial_metadata_(false),
-      compression_level_set_(false),
-      has_pending_ops_(false) {
+ServerContext::ServerContext() { Setup(gpr_inf_future(GPR_CLOCK_REALTIME)); }
+
+ServerContext::ServerContext(gpr_timespec deadline, grpc_metadata_array* arr) {
+  Setup(deadline);
   std::swap(*client_metadata_.arr(), *arr);
 }
 
-ServerContext::~ServerContext() {
+void ServerContext::Setup(gpr_timespec deadline) {
+  completion_op_ = nullptr;
+  has_notify_when_done_tag_ = false;
+  async_notify_when_done_tag_ = nullptr;
+  deadline_ = deadline;
+  call_ = nullptr;
+  cq_ = nullptr;
+  sent_initial_metadata_ = false;
+  compression_level_set_ = false;
+  has_pending_ops_ = false;
+  rpc_info_ = nullptr;
+}
+
+void ServerContext::BindDeadlineAndMetadata(gpr_timespec deadline,
+                                            grpc_metadata_array* arr) {
+  deadline_ = deadline;
+  std::swap(*client_metadata_.arr(), *arr);
+}
+
+ServerContext::~ServerContext() { Clear(); }
+
+void ServerContext::Clear() {
   if (call_) {
     grpc_call_unref(call_);
   }
@@ -243,6 +247,8 @@ ServerContext::~ServerContext() {
   if (rpc_info_) {
     rpc_info_->Unref();
   }
+  // Don't need to clear out call_, completion_op_, or rpc_info_ because this is
+  // either called from destructor or just before Setup
 }
 
 void ServerContext::BeginCompletionOp(internal::Call* call) {