async first take

1 files changed, 80 insertions, 2 deletions

diff --git a/src/cpp/server/server_cc.cc b/src/cpp/server/server_cc.cc
index c50c076bdc..20641aeea8 100644
--- a/src/cpp/server/server_cc.cc
+++ b/src/cpp/server/server_cc.cc
@@ -37,6 +37,7 @@
 
 #include <grpc++/completion_queue.h>
 #include <grpc++/generic/async_generic_service.h>
+#include <grpc++/impl/codegen/async_unary_call.h>
 #include <grpc++/impl/codegen/completion_queue_tag.h>
 #include <grpc++/impl/grpc_library.h>
 #include <grpc++/impl/method_handler_impl.h>
@@ -118,6 +119,67 @@ class Server::UnimplementedAsyncResponse final
   UnimplementedAsyncRequest* const request_;
 };
 
+class Server::HealthCheckAsyncRequestContext {
+ protected:
+  HealthCheckAsyncRequestContext() : rpc_(&server_context_) {}
+  ServerContext server_context_;
+  ServerAsyncResponseWriter<ByteBuffer> rpc_;
+};
+
+class Server::HealthCheckAsyncRequest final
+    : public HealthCheckAsyncRequestContext,
+      public RegisteredAsyncRequest {
+ public:
+  HealthCheckAsyncRequest(
+      DefaultHealthCheckService::AsyncHealthCheckServiceImpl* service,
+      Server* server, ServerCompletionQueue* cq)
+      : RegisteredAsyncRequest(server, &server_context_, &rpc_, cq, this,
+                               false),
+        service_(service),
+        server_(server),
+        cq_(cq),
+        had_request_(false) {
+    IssueRequest(service->method()->server_tag(), &payload_, cq);
+  }
+
+  bool FinalizeResult(void** tag, bool* status) override;
+
+ private:
+  DefaultHealthCheckService::AsyncHealthCheckServiceImpl* service_;
+  Server* const server_;
+  ServerCompletionQueue* const cq_;
+  grpc_byte_buffer* payload_;
+  bool had_request_;
+  ByteBuffer request_;
+  ByteBuffer response_;
+};
+
+bool Server::HealthCheckAsyncRequest::FinalizeResult(void** tag, bool* status) {
+  if (!had_request_) {
+    had_request_ = true;
+    bool serialization_status =
+        *status && payload_ &&
+        SerializationTraits<ByteBuffer>::Deserialize(
+            payload_, &request_, server_->max_receive_message_size())
+            .ok();
+    RegisteredAsyncRequest::FinalizeResult(tag, status);
+    *status = serialization_status && *status;
+    if (*status) {
+      new HealthCheckAsyncRequest(service_, server_, cq_);
+      Status s = service_->Check(&server_context_, &request_, &response_);
+      rpc_.Finish(response_, s, this);
+      return false;
+    } else {
+      // TODO what to do here
+      delete this;
+      return false;
+    }
+  } else {
+    delete this;
+    return false;
+  }
+}
+
 class ShutdownTag : public CompletionQueueTag {
  public:
   bool FinalizeResult(void** tag, bool* status) { return false; }
@@ -498,6 +560,8 @@ bool Server::Start(ServerCompletionQueue** cqs, size_t num_cqs) {
 
   // Only create default health check service when user did not provide an
   // explicit one.
+  DefaultHealthCheckService::AsyncHealthCheckServiceImpl* async_health_service =
+      nullptr;
   if (health_check_service_ == nullptr && !health_check_service_disabled_ &&
       DefaultHealthCheckServiceEnabled()) {
     auto* default_hc_service = new DefaultHealthCheckService;
@@ -505,6 +569,10 @@ bool Server::Start(ServerCompletionQueue** cqs, size_t num_cqs) {
     if (!sync_server_cqs_->empty()) {  // Has sync methods.
       RegisterService(nullptr, default_hc_service->GetSyncHealthCheckService());
     }
+    if (sync_server_cqs_->empty()) {  // No sync methods.
+      async_health_service = default_hc_service->GetAsyncHealthCheckService();
+      RegisterService(nullptr, async_health_service);
+    }
   }
 
   grpc_server_start(server_);
@@ -521,6 +589,14 @@ bool Server::Start(ServerCompletionQueue** cqs, size_t num_cqs) {
     }
   }
 
+  if (async_health_service) {
+    for (size_t i = 0; i < num_cqs; i++) {
+      if (cqs[i]->IsFrequentlyPolled()) {
+        new HealthCheckAsyncRequest(async_health_service, this, cqs[i]);
+      }
+    }
+  }
+
   for (auto it = sync_req_mgrs_.begin(); it != sync_req_mgrs_.end(); it++) {
     (*it)->Start();
   }
@@ -641,8 +717,10 @@ bool ServerInterface::BaseAsyncRequest::FinalizeResult(void** tag,
 
 ServerInterface::RegisteredAsyncRequest::RegisteredAsyncRequest(
     ServerInterface* server, ServerContext* context,
-    ServerAsyncStreamingInterface* stream, CompletionQueue* call_cq, void* tag)
-    : BaseAsyncRequest(server, context, stream, call_cq, tag, true) {}
+    ServerAsyncStreamingInterface* stream, CompletionQueue* call_cq, void* tag,
+    bool delete_on_finalize)
+    : BaseAsyncRequest(server, context, stream, call_cq, tag,
+                       delete_on_finalize) {}
 
 void ServerInterface::RegisteredAsyncRequest::IssueRequest(
     void* registered_method, grpc_byte_buffer** payload,