[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;
+    // 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) {
-      // 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 "
+        DLOG(WARNING)
+            << "Timed out when trying to receive rpc ndarray data after "
             << timeout << " milliseconds.";
-      return rpc::kRPCTimeOut;
       }
+    } while (status == QUEUE_EMPTY);
     CHECK_EQ(status, REMOVE_SUCCESS);
     buffer_list[i] = ndarray_data_msg.data;
   }