[core] Use async next for grpc server (#51378)

It's recommended to use AsyncNext if you're bringing your own threading
model, which we are here by calling poll from our polling threads.
Otherwise Next will use its own internal threading model.

---------

Signed-off-by: dayshah <dhyey2019@gmail.com>

Author: dayshah

src/ray/rpc/client_call.h

@@ -308,7 +308,7 @@ class ClientCallManager {
       } else if (status != grpc::CompletionQueue::TIMEOUT) {
         // NOTE: CompletionQueue::TIMEOUT and gRPC deadline exceeded are different.
         // If the client deadline is exceeded, event is obtained at this block.
-        auto tag = reinterpret_cast<ClientCallTag *>(got_tag);
+        auto tag = static_cast<ClientCallTag *>(got_tag);
         // Refresh the tag.
         got_tag = nullptr;
         tag->GetCall()->SetReturnStatus();

src/ray/rpc/grpc_server.cc

@@ -43,6 +43,7 @@ void GrpcServer::Init() {
 
 void GrpcServer::Shutdown() {
   if (!is_closed_) {
+    shutdown_ = true;
     // Drain the executor threads.
     // Shutdown the server with an immediate deadline.
     // TODO(edoakes): do we want to do this in all cases?
@@ -189,7 +190,19 @@ void GrpcServer::PollEventsFromCompletionQueue(int index) {
   bool ok;
 
   // Keep reading events from the `CompletionQueue` until it's shutdown.
-  while (cqs_[index]->Next(&tag, &ok)) {
+  while (true) {
+    auto deadline = gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
+                                 gpr_time_from_millis(250, GPR_TIMESPAN));
+    auto status = cqs_[index]->AsyncNext(&tag, &ok, deadline);
+    if (status == grpc::CompletionQueue::SHUTDOWN ||
+        (status == grpc::CompletionQueue::TIMEOUT && shutdown_)) {
+      // If we timed out and shutdown, then exit immediately. This should not
+      // be needed, but gRPC seems to not return SHUTDOWN correctly in these
+      // cases (e.g., test_wait will hang on shutdown without this check).
+      break;
+    } else if (status == grpc::CompletionQueue::TIMEOUT) {
+      continue;
+    }
     auto *server_call = static_cast<ServerCall *>(tag);
     bool delete_call = false;
     // A new call is needed after the server sends a reply, no matter the reply is
@@ -241,6 +254,5 @@ void GrpcServer::PollEventsFromCompletionQueue(int index) {
     }
   }
 }
-
 }  // namespace rpc
 }  // namespace ray

src/ray/rpc/grpc_server.h

@@ -105,7 +105,8 @@ class GrpcServer {
         cluster_id_(ClusterID::Nil()),
         is_closed_(true),
         num_threads_(num_threads),
-        keepalive_time_ms_(keepalive_time_ms) {
+        keepalive_time_ms_(keepalive_time_ms),
+        shutdown_(false) {
     Init();
   }
 
@@ -182,6 +183,8 @@ class GrpcServer {
   /// gRPC server cannot get the ping response within the time, it triggers
   /// the watchdog timer fired error, which will close the connection.
   const int64_t keepalive_time_ms_;
+
+  std::atomic_bool shutdown_;
 };
 
 /// Base class that represents an abstract gRPC service.