[Dist] re-try to receive rpc ndarray msg (#4142)

src/rpc/network/socket_communicator.cc

@@ -291,14 +291,17 @@ rpc::RPCStatus SocketReceiver::Recv(rpc::RPCMessage* msg, int timeout) {
   std::vector<void*> buffer_list(nonempty_ndarray_count);
   for (int i = 0; i < nonempty_ndarray_count; ++i) {
     Message ndarray_data_msg;
-    status = RecvFrom(&ndarray_data_msg, send_id, timeout);
-    if (status == QUEUE_EMPTY) {
-      // As we cannot handle this timeout for now, let's treat it as fatal
-      // error.
-      LOG(FATAL) << "Timed out when trying to receive rpc ndarray data after "
-                 << timeout << " milliseconds.";
-      return rpc::kRPCTimeOut;
-    }
+    // As meta message has been received, data message is always expected unless
+    // connection is closed.
+    STATUS status;
+    do {
+      status = RecvFrom(&ndarray_data_msg, send_id, timeout);
+      if (status == QUEUE_EMPTY) {
+        DLOG(WARNING)
+            << "Timed out when trying to receive rpc ndarray data after "
+            << timeout << " milliseconds.";
+      }
+    } while (status == QUEUE_EMPTY);
     CHECK_EQ(status, REMOVE_SUCCESS);
     buffer_list[i] = ndarray_data_msg.data;
   }