Add return value for recontruction RPC. (#3493)

* Add return value for recontruct RPC.

* Fix comment function name

src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.cc

@@ -119,9 +119,9 @@ JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeDestro
 /*
  * Class:     org_ray_runtime_raylet_RayletClientImpl
  * Method:    nativeFetchOrReconstruct
- * Signature: (J[[BZ)V
+ * Signature: (J[[BZ[B)I
  */
-JNIEXPORT void JNICALL
+JNIEXPORT jint JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativeFetchOrReconstruct(
     JNIEnv *env, jclass, jlong client, jobjectArray objectIds, jboolean fetchOnly,
     jbyteArray currentTaskId) {
@@ -136,7 +136,8 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeFetchOrReconstruct(
   }
   UniqueIdFromJByteArray current_task_id(env, currentTaskId);
   auto conn = reinterpret_cast<LocalSchedulerConnection *>(client);
-  local_scheduler_fetch_or_reconstruct(conn, object_ids, fetchOnly, *current_task_id.PID);
+  return local_scheduler_fetch_or_reconstruct(conn, object_ids, fetchOnly,
+                                              *current_task_id.PID);
 }
 
 /*

src/ray/raylet/lib/java/org_ray_runtime_raylet_RayletClientImpl.h

@@ -42,9 +42,9 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeDestroy(JNIEnv *, jclass, jlo
 /*
  * Class:     org_ray_runtime_raylet_RayletClientImpl
  * Method:    nativeFetchOrReconstruct
- * Signature: (J[[BZ)V
+ * Signature: (J[[BZ[B)I
  */
-JNIEXPORT void JNICALL
+JNIEXPORT jint JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativeFetchOrReconstruct(JNIEnv *, jclass,
                                                                       jlong, jobjectArray,
                                                                       jboolean,

src/ray/raylet/lib/python/local_scheduler_extension.cc

@@ -1,4 +1,5 @@
 #include <Python.h>
+#include <sstream>
 
 #include "common_extension.h"
 #include "config_extension.h"
@@ -92,10 +93,19 @@ static PyObject *PyLocalSchedulerClient_fetch_or_reconstruct(PyObject *self,
     }
     object_ids.push_back(object_id);
   }
-  local_scheduler_fetch_or_reconstruct(
+  int ret = local_scheduler_fetch_or_reconstruct(
       reinterpret_cast<PyLocalSchedulerClient *>(self)->local_scheduler_connection,
       object_ids, fetch_only, current_task_id);
-  Py_RETURN_NONE;
+  if (ret == 0) {
+    Py_RETURN_NONE;
+  } else {
+    std::ostringstream stream;
+    stream << "local_scheduler_fetch_or_reconstruct failed: "
+           << "local scheduler connection may be closed, "
+           << "check raylet status. return value: " << ret;
+    PyErr_SetString(CommonError, stream.str().c_str());
+    Py_RETURN_NONE;
+  }
 }
 
 static PyObject *PyLocalSchedulerClient_notify_unblocked(PyObject *self, PyObject *args) {

src/ray/raylet/local_scheduler_client.cc

@@ -306,18 +306,16 @@ void local_scheduler_task_done(LocalSchedulerConnection *conn) {
                 &conn->write_mutex);
 }
 
-void local_scheduler_fetch_or_reconstruct(LocalSchedulerConnection *conn,
-                                          const std::vector<ObjectID> &object_ids,
-                                          bool fetch_only,
-                                          const TaskID &current_task_id) {
+int local_scheduler_fetch_or_reconstruct(LocalSchedulerConnection *conn,
+                                         const std::vector<ObjectID> &object_ids,
+                                         bool fetch_only, const TaskID &current_task_id) {
   flatbuffers::FlatBufferBuilder fbb;
   auto object_ids_message = to_flatbuf(fbb, object_ids);
   auto message = ray::protocol::CreateFetchOrReconstruct(
       fbb, object_ids_message, fetch_only, to_flatbuf(fbb, current_task_id));
   fbb.Finish(message);
-  write_message(conn->conn, static_cast<int64_t>(MessageType::FetchOrReconstruct),
-                fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex);
-  /* TODO(swang): Propagate the error. */
+  return write_message(conn->conn, static_cast<int64_t>(MessageType::FetchOrReconstruct),
+                       fbb.GetSize(), fbb.GetBufferPointer(), &conn->write_mutex);
 }
 
 void local_scheduler_notify_unblocked(LocalSchedulerConnection *conn,

src/ray/raylet/local_scheduler_client.h

@@ -97,11 +97,11 @@ void local_scheduler_task_done(LocalSchedulerConnection *conn);
  * @param object_ids The IDs of the objects to reconstruct.
  * @param fetch_only Only fetch objects, do not reconstruct them.
  * @param current_task_id The task that needs the objects.
- * @return Void.
+ * @return int 0 means correct, other numbers mean error.
  */
-void local_scheduler_fetch_or_reconstruct(LocalSchedulerConnection *conn,
-                                          const std::vector<ObjectID> &object_ids,
-                                          bool fetch_only, const TaskID &current_task_id);
+int local_scheduler_fetch_or_reconstruct(LocalSchedulerConnection *conn,
+                                         const std::vector<ObjectID> &object_ids,
+                                         bool fetch_only, const TaskID &current_task_id);
 
 /**
  * Notify the local scheduler that this client (worker) is no longer blocked.