[Java worker] Refactor object store and worker context on top of core worker (#5079)

src/ray/common/ray_config_def.h

@@ -151,3 +151,10 @@ RAY_CONFIG(uint32_t, num_actor_checkpoints_to_keep, 20)
 
 /// Maximum number of ids in one batch to send to GCS to delete keys.
 RAY_CONFIG(uint32_t, maximum_gcs_deletion_batch_size, 1000)
+
+/// When getting objects from object store, print a warning every this number of attempts.
+RAY_CONFIG(uint32_t, object_store_get_warn_per_num_attempts, 50)
+
+/// When getting objects from object store, max number of ids to print in the warning
+/// message.
+RAY_CONFIG(uint32_t, object_store_get_max_ids_to_print_in_warning, 20)

src/ray/core_worker/common.h

@@ -9,9 +9,7 @@
 #include "ray/raylet/raylet_client.h"
 
 namespace ray {
-
-/// Type of this worker.
-enum class WorkerType { WORKER, DRIVER };
+using WorkerType = rpc::WorkerType;
 
 /// Information about a remote function.
 struct RayFunction {

src/ray/core_worker/context.cc

@@ -6,69 +6,81 @@ namespace ray {
 /// per-thread context for core worker.
 struct WorkerThreadContext {
   WorkerThreadContext()
-      : current_task_id(TaskID::FromRandom()), task_index(0), put_index(0) {}
+      : current_task_id_(TaskID::FromRandom()), task_index_(0), put_index_(0) {}
 
-  int GetNextTaskIndex() { return ++task_index; }
+  int GetNextTaskIndex() { return ++task_index_; }
 
-  int GetNextPutIndex() { return ++put_index; }
+  int GetNextPutIndex() { return ++put_index_; }
 
-  const TaskID &GetCurrentTaskID() const { return current_task_id; }
+  const TaskID &GetCurrentTaskID() const { return current_task_id_; }
 
-  void SetCurrentTask(const TaskID &task_id) {
-    current_task_id = task_id;
-    task_index = 0;
-    put_index = 0;
+  std::shared_ptr<const TaskSpecification> GetCurrentTask() const {
+    return current_task_;
+  }
+
+  void SetCurrentTaskId(const TaskID &task_id) {
+    current_task_id_ = task_id;
+    task_index_ = 0;
+    put_index_ = 0;
   }
 
   void SetCurrentTask(const TaskSpecification &task_spec) {
-    SetCurrentTask(task_spec.TaskId());
+    SetCurrentTaskId(task_spec.TaskId());
+    current_task_ = std::make_shared<const TaskSpecification>(task_spec);
   }
 
  private:
   /// The task ID for current task.
-  TaskID current_task_id;
+  TaskID current_task_id_;
+
+  /// The current task.
+  std::shared_ptr<const TaskSpecification> current_task_;
 
   /// Number of tasks that have been submitted from current task.
-  int task_index;
+  int task_index_;
 
   /// Number of objects that have been put from current task.
-  int put_index;
+  int put_index_;
 };
 
 thread_local std::unique_ptr<WorkerThreadContext> WorkerContext::thread_context_ =
     nullptr;
 
 WorkerContext::WorkerContext(WorkerType worker_type, const JobID &job_id)
-    : worker_type(worker_type),
-      worker_id(worker_type == WorkerType::DRIVER ? ComputeDriverIdFromJob(job_id)
-                                                  : WorkerID::FromRandom()),
-      current_job_id(worker_type == WorkerType::DRIVER ? job_id : JobID::Nil()) {
+    : worker_type_(worker_type),
+      worker_id_(worker_type_ == WorkerType::DRIVER ? ComputeDriverIdFromJob(job_id)
+                                                    : WorkerID::FromRandom()),
+      current_job_id_(worker_type_ == WorkerType::DRIVER ? job_id : JobID::Nil()) {
   // For worker main thread which initializes the WorkerContext,
   // set task_id according to whether current worker is a driver.
   // (For other threads it's set to random ID via GetThreadContext).
-  GetThreadContext().SetCurrentTask(
-      (worker_type == WorkerType::DRIVER) ? TaskID::FromRandom() : TaskID::Nil());
+  GetThreadContext().SetCurrentTaskId(
+      (worker_type_ == WorkerType::DRIVER) ? TaskID::FromRandom() : TaskID::Nil());
 }
 
-const WorkerType WorkerContext::GetWorkerType() const { return worker_type; }
+const WorkerType WorkerContext::GetWorkerType() const { return worker_type_; }
 
-const WorkerID &WorkerContext::GetWorkerID() const { return worker_id; }
+const WorkerID &WorkerContext::GetWorkerID() const { return worker_id_; }
 
 int WorkerContext::GetNextTaskIndex() { return GetThreadContext().GetNextTaskIndex(); }
 
 int WorkerContext::GetNextPutIndex() { return GetThreadContext().GetNextPutIndex(); }
 
-const JobID &WorkerContext::GetCurrentJobID() const { return current_job_id; }
+const JobID &WorkerContext::GetCurrentJobID() const { return current_job_id_; }
 
 const TaskID &WorkerContext::GetCurrentTaskID() const {
   return GetThreadContext().GetCurrentTaskID();
 }
 
 void WorkerContext::SetCurrentTask(const TaskSpecification &task_spec) {
-  current_job_id = task_spec.JobId();
+  current_job_id_ = task_spec.JobId();
   GetThreadContext().SetCurrentTask(task_spec);
 }
 
+std::shared_ptr<const TaskSpecification> WorkerContext::GetCurrentTask() const {
+  return GetThreadContext().GetCurrentTask();
+}
+
 WorkerThreadContext &WorkerContext::GetThreadContext() {
   if (thread_context_ == nullptr) {
     thread_context_ = std::unique_ptr<WorkerThreadContext>(new WorkerThreadContext());

src/ray/core_worker/context.h

@@ -22,19 +22,21 @@ class WorkerContext {
 
   void SetCurrentTask(const TaskSpecification &task_spec);
 
+  std::shared_ptr<const TaskSpecification> GetCurrentTask() const;
+
   int GetNextTaskIndex();
 
   int GetNextPutIndex();
 
  private:
   /// Type of the worker.
-  const WorkerType worker_type;
+  const WorkerType worker_type_;
 
   /// ID for this worker.
-  const WorkerID worker_id;
+  const WorkerID worker_id_;
 
   /// Job ID for this worker.
-  JobID current_job_id;
+  JobID current_job_id_;
 
  private:
   static WorkerThreadContext &GetThreadContext();

src/ray/core_worker/core_worker.cc

@@ -15,7 +15,7 @@ CoreWorker::CoreWorker(
       task_interface_(worker_context_, raylet_client_),
       object_interface_(worker_context_, raylet_client_, store_socket) {
   int rpc_server_port = 0;
-  if (worker_type_ == ray::WorkerType::WORKER) {
+  if (worker_type_ == WorkerType::WORKER) {
     RAY_CHECK(execution_callback != nullptr);
     task_execution_interface_ = std::unique_ptr<CoreWorkerTaskExecutionInterface>(
         new CoreWorkerTaskExecutionInterface(worker_context_, raylet_client_,
@@ -28,8 +28,8 @@ CoreWorker::CoreWorker(
   // instead of crashing.
   raylet_client_ = std::unique_ptr<RayletClient>(new RayletClient(
       raylet_socket_, ClientID::FromBinary(worker_context_.GetWorkerID().Binary()),
-      (worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(),
-      language_, rpc_server_port));
+      (worker_type_ == WorkerType::WORKER), worker_context_.GetCurrentJobID(), language_,
+      rpc_server_port));
 }
 
 }  // namespace ray

src/ray/core_worker/lib/java/jni_init.cc

@@ -0,0 +1,75 @@
+#include "ray/core_worker/lib/java/jni_utils.h"
+
+jclass java_boolean_class;
+jmethodID java_boolean_init;
+
+jclass java_list_class;
+jmethodID java_list_size;
+jmethodID java_list_get;
+jmethodID java_list_add;
+
+jclass java_array_list_class;
+jmethodID java_array_list_init;
+jmethodID java_array_list_init_with_capacity;
+
+jclass java_ray_exception_class;
+
+jclass java_native_ray_object_class;
+jmethodID java_native_ray_object_init;
+jfieldID java_native_ray_object_data;
+jfieldID java_native_ray_object_metadata;
+
+jint JNI_VERSION = JNI_VERSION_1_8;
+
+inline jclass LoadClass(JNIEnv *env, const char *class_name) {
+  jclass tempLocalClassRef = env->FindClass(class_name);
+  jclass ret = (jclass)env->NewGlobalRef(tempLocalClassRef);
+  env->DeleteLocalRef(tempLocalClassRef);
+  return ret;
+}
+
+/// Load and cache frequently-used Java classes and methods
+jint JNI_OnLoad(JavaVM *vm, void *reserved) {
+  JNIEnv *env;
+  if (vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION) != JNI_OK) {
+    return JNI_ERR;
+  }
+
+  java_boolean_class = LoadClass(env, "java/lang/Boolean");
+  java_boolean_init = env->GetMethodID(java_boolean_class, "<init>", "(Z)V");
+
+  java_list_class = LoadClass(env, "java/util/List");
+  java_list_size = env->GetMethodID(java_list_class, "size", "()I");
+  java_list_get = env->GetMethodID(java_list_class, "get", "(I)Ljava/lang/Object;");
+  java_list_add = env->GetMethodID(java_list_class, "add", "(Ljava/lang/Object;)Z");
+
+  java_array_list_class = LoadClass(env, "java/util/ArrayList");
+  java_array_list_init = env->GetMethodID(java_array_list_class, "<init>", "()V");
+  java_array_list_init_with_capacity =
+      env->GetMethodID(java_array_list_class, "<init>", "(I)V");
+
+  java_ray_exception_class = LoadClass(env, "org/ray/api/exception/RayException");
+
+  java_native_ray_object_class =
+      LoadClass(env, "org/ray/runtime/objectstore/NativeRayObject");
+  java_native_ray_object_init =
+      env->GetMethodID(java_native_ray_object_class, "<init>", "([B[B)V");
+  java_native_ray_object_data =
+      env->GetFieldID(java_native_ray_object_class, "data", "[B");
+  java_native_ray_object_metadata =
+      env->GetFieldID(java_native_ray_object_class, "metadata", "[B");
+
+  return JNI_VERSION;
+}
+
+/// Unload java classes
+void JNI_OnUnload(JavaVM *vm, void *reserved) {
+  JNIEnv *env;
+  vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION);
+
+  env->DeleteGlobalRef(java_boolean_class);
+  env->DeleteGlobalRef(java_list_class);
+  env->DeleteGlobalRef(java_array_list_class);
+  env->DeleteGlobalRef(java_ray_exception_class);
+  env->DeleteGlobalRef(java_native_ray_object_class);
+}

src/ray/core_worker/lib/java/jni_utils.h

@@ -0,0 +1,180 @@
+#ifndef RAY_COMMON_JAVA_JNI_HELPER_H
+#define RAY_COMMON_JAVA_JNI_HELPER_H
+
+#include <jni.h>
+#include "ray/common/buffer.h"
+#include "ray/common/id.h"
+#include "ray/common/status.h"
+#include "ray/core_worker/store_provider/store_provider.h"
+
+/// Boolean class
+extern jclass java_boolean_class;
+/// Constructor of Boolean class
+extern jmethodID java_boolean_init;
+
+/// List class
+extern jclass java_list_class;
+/// size method of List class
+extern jmethodID java_list_size;
+/// get method of List class
+extern jmethodID java_list_get;
+/// add method of List class
+extern jmethodID java_list_add;
+
+/// ArrayList class
+extern jclass java_array_list_class;
+/// Constructor of ArrayList class
+extern jmethodID java_array_list_init;
+/// Constructor of ArrayList class with single parameter capacity
+extern jmethodID java_array_list_init_with_capacity;
+
+/// RayException class
+extern jclass java_ray_exception_class;
+
+/// NativeRayObject class
+extern jclass java_native_ray_object_class;
+/// Constructor of NativeRayObject class
+extern jmethodID java_native_ray_object_init;
+/// data field of NativeRayObject class
+extern jfieldID java_native_ray_object_data;
+/// metadata field of NativeRayObject class
+extern jfieldID java_native_ray_object_metadata;
+
+/// Throws a Java RayException if the status is not OK.
+#define THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, ret)               \
+  {                                                                          \
+    if (!(status).ok()) {                                                    \
+      (env)->ThrowNew(java_ray_exception_class, (status).message().c_str()); \
+      return (ret);                                                          \
+    }                                                                        \
+  }
+
+/// Convert a Java byte array to a C++ UniqueID.
+template <typename ID>
+inline ID JavaByteArrayToId(JNIEnv *env, const jbyteArray &bytes) {
+  std::string id_str(ID::Size(), 0);
+  env->GetByteArrayRegion(bytes, 0, ID::Size(),
+                          reinterpret_cast<jbyte *>(&id_str.front()));
+  return ID::FromBinary(id_str);
+}
+
+/// Convert C++ UniqueID to a Java byte array.
+template <typename ID>
+inline jbyteArray IdToJavaByteArray(JNIEnv *env, const ID &id) {
+  jbyteArray array = env->NewByteArray(ID::Size());
+  env->SetByteArrayRegion(array, 0, ID::Size(),
+                          reinterpret_cast<const jbyte *>(id.Data()));
+  return array;
+}
+
+/// Convert C++ UniqueID to a Java ByteBuffer.
+template <typename ID>
+inline jobject IdToJavaByteBuffer(JNIEnv *env, const ID &id) {
+  return env->NewDirectByteBuffer(
+      reinterpret_cast<void *>(const_cast<uint8_t *>(id.Data())), id.Size());
+}
+
+/// Convert a Java String to C++ std::string.
+inline std::string JavaStringToNativeString(JNIEnv *env, jstring jstr) {
+  const char *c_str = env->GetStringUTFChars(jstr, nullptr);
+  std::string result(c_str);
+  env->ReleaseStringUTFChars(static_cast<jstring>(jstr), c_str);
+  return result;
+}
+
+/// Convert a Java List to C++ std::vector.
+template <typename NativeT>
+inline void JavaListToNativeVector(
+    JNIEnv *env, jobject java_list, std::vector<NativeT> *native_vector,
+    std::function<NativeT(JNIEnv *, jobject)> element_converter) {
+  int size = env->CallIntMethod(java_list, java_list_size);
+  native_vector->clear();
+  for (int i = 0; i < size; i++) {
+    native_vector->emplace_back(
+        element_converter(env, env->CallObjectMethod(java_list, java_list_get, (jint)i)));
+  }
+}
+
+/// Convert a C++ std::vector to a Java List.
+template <typename NativeT>
+inline jobject NativeVectorToJavaList(
+    JNIEnv *env, const std::vector<NativeT> &native_vector,
+    std::function<jobject(JNIEnv *, const NativeT &)> element_converter) {
+  jobject java_list =
+      env->NewObject(java_array_list_class, java_array_list_init_with_capacity,
+                     (jint)native_vector.size());
+  for (const auto &item : native_vector) {
+    env->CallVoidMethod(java_list, java_list_add, element_converter(env, item));
+  }
+  return java_list;
+}
+
+/// Convert a C++ ray::Buffer to a Java byte array.
+inline jbyteArray NativeBufferToJavaByteArray(JNIEnv *env,
+                                              const std::shared_ptr<ray::Buffer> buffer) {
+  if (!buffer) {
+    return nullptr;
+  }
+  jbyteArray java_byte_array = env->NewByteArray(buffer->Size());
+  if (buffer->Size() > 0) {
+    env->SetByteArrayRegion(java_byte_array, 0, buffer->Size(),
+                            reinterpret_cast<const jbyte *>(buffer->Data()));
+  }
+  return java_byte_array;
+}
+
+/// A helper method to help access a Java NativeRayObject instance and ensure memory
+/// safety.
+///
+/// \param[in] java_obj The Java NativeRayObject object.
+/// \param[in] reader The callback function to access a C++ ray::RayObject instance.
+/// \return The return value of callback function.
+template <typename ReturnT>
+inline ReturnT ReadJavaNativeRayObject(
+    JNIEnv *env, const jobject &java_obj,
+    std::function<ReturnT(const std::shared_ptr<ray::RayObject> &)> reader) {
+  if (!java_obj) {
+    return reader(nullptr);
+  }
+  auto java_data = (jbyteArray)env->GetObjectField(java_obj, java_native_ray_object_data);
+  auto java_metadata =
+      (jbyteArray)env->GetObjectField(java_obj, java_native_ray_object_metadata);
+  auto data_size = env->GetArrayLength(java_data);
+  jbyte *data = data_size > 0 ? env->GetByteArrayElements(java_data, nullptr) : nullptr;
+  auto metadata_size = java_metadata ? env->GetArrayLength(java_metadata) : 0;
+  jbyte *metadata =
+      metadata_size > 0 ? env->GetByteArrayElements(java_metadata, nullptr) : nullptr;
+  auto data_buffer = std::make_shared<ray::LocalMemoryBuffer>(
+      reinterpret_cast<uint8_t *>(data), data_size);
+  auto metadata_buffer = java_metadata
+                             ? std::make_shared<ray::LocalMemoryBuffer>(
+                                   reinterpret_cast<uint8_t *>(metadata), metadata_size)
+                             : nullptr;
+
+  auto native_obj = std::make_shared<ray::RayObject>(data_buffer, metadata_buffer);
+  auto result = reader(native_obj);
+
+  if (data) {
+    env->ReleaseByteArrayElements(java_data, data, JNI_ABORT);
+  }
+  if (metadata) {
+    env->ReleaseByteArrayElements(java_metadata, metadata, JNI_ABORT);
+  }
+
+  return result;
+}
+
+/// Convert a C++ ray::RayObject to a Java NativeRayObject.
+inline jobject ToJavaNativeRayObject(JNIEnv *env,
+                                     const std::shared_ptr<ray::RayObject> &rayObject) {
+  if (!rayObject) {
+    return nullptr;
+  }
+  auto java_data = NativeBufferToJavaByteArray(env, rayObject->GetData());
+  auto java_metadata = NativeBufferToJavaByteArray(env, rayObject->GetMetadata());
+  auto java_obj = env->NewObject(java_native_ray_object_class,
+                                 java_native_ray_object_init, java_data, java_metadata);
+  return java_obj;
+}
+
+#endif  // RAY_COMMON_JAVA_JNI_HELPER_H

src/ray/core_worker/lib/java/org_ray_runtime_WorkerContext.cc

@@ -0,0 +1,134 @@
+#include "ray/core_worker/lib/java/org_ray_runtime_WorkerContext.h"
+#include <jni.h>
+#include "ray/common/id.h"
+#include "ray/core_worker/context.h"
+#include "ray/core_worker/lib/java/jni_utils.h"
+
+inline ray::WorkerContext *GetWorkerContextFromPointer(
+    jlong nativeWorkerContextFromPointer) {
+  return reinterpret_cast<ray::WorkerContext *>(nativeWorkerContextFromPointer);
+}
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeCreateWorkerContext
+ * Signature: (I[B)J
+ */
+JNIEXPORT jlong JNICALL Java_org_ray_runtime_WorkerContext_nativeCreateWorkerContext(
+    JNIEnv *env, jclass, jint workerType, jbyteArray jobId) {
+  return reinterpret_cast<jlong>(
+      new ray::WorkerContext(static_cast<ray::rpc::WorkerType>(workerType),
+                             JavaByteArrayToId<ray::JobID>(env, jobId)));
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentTaskId
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_ray_runtime_WorkerContext_nativeGetCurrentTaskId(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  auto task_id =
+      GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->GetCurrentTaskID();
+  return IdToJavaByteArray<ray::TaskID>(env, task_id);
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeSetCurrentTask
+ * Signature: (J[B)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_WorkerContext_nativeSetCurrentTask(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer, jbyteArray taskSpec) {
+  jbyte *data = env->GetByteArrayElements(taskSpec, NULL);
+  jsize size = env->GetArrayLength(taskSpec);
+  ray::rpc::TaskSpec task_spec_message;
+  task_spec_message.ParseFromArray(data, size);
+  env->ReleaseByteArrayElements(taskSpec, data, JNI_ABORT);
+
+  ray::TaskSpecification spec(task_spec_message);
+  GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->SetCurrentTask(spec);
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentTask
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_ray_runtime_WorkerContext_nativeGetCurrentTask(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  auto spec =
+      GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->GetCurrentTask();
+  if (!spec) {
+    return nullptr;
+  }
+
+  auto task_message = spec->Serialize();
+  jbyteArray result = env->NewByteArray(task_message.size());
+  env->SetByteArrayRegion(
+      result, 0, task_message.size(),
+      reinterpret_cast<jbyte *>(const_cast<char *>(task_message.data())));
+  return result;
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentJobId
+ * Signature: (J)Ljava/nio/ByteBuffer;
+ */
+JNIEXPORT jobject JNICALL Java_org_ray_runtime_WorkerContext_nativeGetCurrentJobId(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  const auto &job_id =
+      GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->GetCurrentJobID();
+  return IdToJavaByteBuffer<ray::JobID>(env, job_id);
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentWorkerId
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_ray_runtime_WorkerContext_nativeGetCurrentWorkerId(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  auto worker_id =
+      GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->GetWorkerID();
+  return IdToJavaByteArray<ray::WorkerID>(env, worker_id);
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetNextTaskIndex
+ * Signature: (J)I
+ */
+JNIEXPORT jint JNICALL Java_org_ray_runtime_WorkerContext_nativeGetNextTaskIndex(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  return GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->GetNextTaskIndex();
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetNextPutIndex
+ * Signature: (J)I
+ */
+JNIEXPORT jint JNICALL Java_org_ray_runtime_WorkerContext_nativeGetNextPutIndex(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  return GetWorkerContextFromPointer(nativeWorkerContextFromPointer)->GetNextPutIndex();
+}
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeDestroy
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_WorkerContext_nativeDestroy(
+    JNIEnv *env, jclass, jlong nativeWorkerContextFromPointer) {
+  delete GetWorkerContextFromPointer(nativeWorkerContextFromPointer);
+}
+
+#ifdef __cplusplus
+}
+#endif

src/ray/core_worker/lib/java/org_ray_runtime_WorkerContext.h

@@ -0,0 +1,87 @@
+/* DO NOT EDIT THIS FILE - it is machine generated */
+#include <jni.h>
+/* Header for class org_ray_runtime_WorkerContext */
+
+#ifndef _Included_org_ray_runtime_WorkerContext
+#define _Included_org_ray_runtime_WorkerContext
+#ifdef __cplusplus
+extern "C" {
+#endif
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeCreateWorkerContext
+ * Signature: (I[B)J
+ */
+JNIEXPORT jlong JNICALL Java_org_ray_runtime_WorkerContext_nativeCreateWorkerContext(
+    JNIEnv *, jclass, jint, jbyteArray);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentTaskId
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_runtime_WorkerContext_nativeGetCurrentTaskId(JNIEnv *, jclass, jlong);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeSetCurrentTask
+ * Signature: (J[B)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_WorkerContext_nativeSetCurrentTask(
+    JNIEnv *, jclass, jlong, jbyteArray);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentTask
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_runtime_WorkerContext_nativeGetCurrentTask(JNIEnv *, jclass, jlong);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentJobId
+ * Signature: (J)Ljava/nio/ByteBuffer;
+ */
+JNIEXPORT jobject JNICALL
+Java_org_ray_runtime_WorkerContext_nativeGetCurrentJobId(JNIEnv *, jclass, jlong);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetCurrentWorkerId
+ * Signature: (J)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_runtime_WorkerContext_nativeGetCurrentWorkerId(JNIEnv *, jclass, jlong);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetNextTaskIndex
+ * Signature: (J)I
+ */
+JNIEXPORT jint JNICALL Java_org_ray_runtime_WorkerContext_nativeGetNextTaskIndex(JNIEnv *,
+                                                                                 jclass,
+                                                                                 jlong);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeGetNextPutIndex
+ * Signature: (J)I
+ */
+JNIEXPORT jint JNICALL Java_org_ray_runtime_WorkerContext_nativeGetNextPutIndex(JNIEnv *,
+                                                                                jclass,
+                                                                                jlong);
+
+/*
+ * Class:     org_ray_runtime_WorkerContext
+ * Method:    nativeDestroy
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_WorkerContext_nativeDestroy(JNIEnv *, jclass,
+                                                                        jlong);
+
+#ifdef __cplusplus
+}
+#endif
+#endif

src/ray/core_worker/lib/java/org_ray_runtime_objectstore_ObjectInterfaceImpl.cc

@@ -0,0 +1,149 @@
+#include "ray/core_worker/lib/java/org_ray_runtime_objectstore_ObjectInterfaceImpl.h"
+#include <jni.h>
+#include "ray/common/id.h"
+#include "ray/core_worker/common.h"
+#include "ray/core_worker/lib/java/jni_utils.h"
+#include "ray/core_worker/object_interface.h"
+
+inline ray::CoreWorkerObjectInterface *GetObjectInterfaceFromPointer(
+    jlong nativeObjectInterfacePointer) {
+  return reinterpret_cast<ray::CoreWorkerObjectInterface *>(nativeObjectInterfacePointer);
+}
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeCreateObjectInterface
+ * Signature: (JJLjava/lang/String;)J
+ */
+JNIEXPORT jlong JNICALL
+Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeCreateObjectInterface(
+    JNIEnv *env, jclass, jlong nativeWorkerContext, jlong nativeRayletClient,
+    jstring storeSocketName) {
+  return reinterpret_cast<jlong>(new ray::CoreWorkerObjectInterface(
+      *reinterpret_cast<ray::WorkerContext *>(nativeWorkerContext),
+      *reinterpret_cast<std::unique_ptr<RayletClient> *>(nativeRayletClient),
+      JavaStringToNativeString(env, storeSocketName)));
+}
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativePut
+ * Signature: (JLorg/ray/runtime/objectstore/NativeRayObject;)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativePut__JLorg_ray_runtime_objectstore_NativeRayObject_2(
+    JNIEnv *env, jclass, jlong nativeObjectInterfacePointer, jobject obj) {
+  ray::Status status;
+  ray::ObjectID object_id = ReadJavaNativeRayObject<ray::ObjectID>(
+      env, obj,
+      [nativeObjectInterfacePointer,
+       &status](const std::shared_ptr<ray::RayObject> &rayObject) {
+        RAY_CHECK(rayObject != nullptr);
+        ray::ObjectID object_id;
+        status = GetObjectInterfaceFromPointer(nativeObjectInterfacePointer)
+                     ->Put(*rayObject, &object_id);
+        return object_id;
+      });
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
+  return IdToJavaByteArray<ray::ObjectID>(env, object_id);
+}
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativePut
+ * Signature: (J[BLorg/ray/runtime/objectstore/NativeRayObject;)V
+ */
+JNIEXPORT void JNICALL
+Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativePut__J_3BLorg_ray_runtime_objectstore_NativeRayObject_2(
+    JNIEnv *env, jclass, jlong nativeObjectInterfacePointer, jbyteArray objectId,
+    jobject obj) {
+  auto object_id = JavaByteArrayToId<ray::ObjectID>(env, objectId);
+  auto status = ReadJavaNativeRayObject<ray::Status>(
+      env, obj,
+      [nativeObjectInterfacePointer,
+       &object_id](const std::shared_ptr<ray::RayObject> &rayObject) {
+        RAY_CHECK(rayObject != nullptr);
+        return GetObjectInterfaceFromPointer(nativeObjectInterfacePointer)
+            ->Put(*rayObject, object_id);
+      });
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
+}
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeGet
+ * Signature: (JLjava/util/List;J)Ljava/util/List;
+ */
+JNIEXPORT jobject JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeGet(
+    JNIEnv *env, jclass, jlong nativeObjectInterfacePointer, jobject ids,
+    jlong timeoutMs) {
+  std::vector<ray::ObjectID> object_ids;
+  JavaListToNativeVector<ray::ObjectID>(
+      env, ids, &object_ids, [](JNIEnv *env, jobject id) {
+        return JavaByteArrayToId<ray::ObjectID>(env, static_cast<jbyteArray>(id));
+      });
+  std::vector<std::shared_ptr<ray::RayObject>> results;
+  auto status = GetObjectInterfaceFromPointer(nativeObjectInterfacePointer)
+                    ->Get(object_ids, (int64_t)timeoutMs, &results);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
+  return NativeVectorToJavaList<std::shared_ptr<ray::RayObject>>(env, results,
+                                                                 ToJavaNativeRayObject);
+}
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeWait
+ * Signature: (JLjava/util/List;IJ)Ljava/util/List;
+ */
+JNIEXPORT jobject JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeWait(
+    JNIEnv *env, jclass, jlong nativeObjectInterfacePointer, jobject objectIds,
+    jint numObjects, jlong timeoutMs) {
+  std::vector<ray::ObjectID> object_ids;
+  JavaListToNativeVector<ray::ObjectID>(
+      env, objectIds, &object_ids, [](JNIEnv *env, jobject id) {
+        return JavaByteArrayToId<ray::ObjectID>(env, static_cast<jbyteArray>(id));
+      });
+  std::vector<bool> results;
+  auto status = GetObjectInterfaceFromPointer(nativeObjectInterfacePointer)
+                    ->Wait(object_ids, (int)numObjects, (int64_t)timeoutMs, &results);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
+  return NativeVectorToJavaList<bool>(env, results, [](JNIEnv *env, const bool &item) {
+    return env->NewObject(java_boolean_class, java_boolean_init, (jboolean)item);
+  });
+}
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeDelete
+ * Signature: (JLjava/util/List;ZZ)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeDelete(
+    JNIEnv *env, jclass, jlong nativeObjectInterfacePointer, jobject objectIds,
+    jboolean localOnly, jboolean deleteCreatingTasks) {
+  std::vector<ray::ObjectID> object_ids;
+  JavaListToNativeVector<ray::ObjectID>(
+      env, objectIds, &object_ids, [](JNIEnv *env, jobject id) {
+        return JavaByteArrayToId<ray::ObjectID>(env, static_cast<jbyteArray>(id));
+      });
+  auto status = GetObjectInterfaceFromPointer(nativeObjectInterfacePointer)
+                    ->Delete(object_ids, (bool)localOnly, (bool)deleteCreatingTasks);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
+}
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeDestroy
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeDestroy(
+    JNIEnv *env, jclass, jlong nativeObjectInterfacePointer) {
+  delete GetObjectInterfaceFromPointer(nativeObjectInterfacePointer);
+}
+
+#ifdef __cplusplus
+}
+#endif

src/ray/core_worker/lib/java/org_ray_runtime_objectstore_ObjectInterfaceImpl.h

@@ -0,0 +1,72 @@
+/* DO NOT EDIT THIS FILE - it is machine generated */
+#include <jni.h>
+/* Header for class org_ray_runtime_objectstore_ObjectInterfaceImpl */
+
+#ifndef _Included_org_ray_runtime_objectstore_ObjectInterfaceImpl
+#define _Included_org_ray_runtime_objectstore_ObjectInterfaceImpl
+#ifdef __cplusplus
+extern "C" {
+#endif
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeCreateObjectInterface
+ * Signature: (JJLjava/lang/String;)J
+ */
+JNIEXPORT jlong JNICALL
+Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeCreateObjectInterface(
+    JNIEnv *, jclass, jlong, jlong, jstring);
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativePut
+ * Signature: (JLorg/ray/runtime/objectstore/NativeRayObject;)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativePut__JLorg_ray_runtime_objectstore_NativeRayObject_2(
+    JNIEnv *, jclass, jlong, jobject);
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativePut
+ * Signature: (J[BLorg/ray/runtime/objectstore/NativeRayObject;)V
+ */
+JNIEXPORT void JNICALL
+Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativePut__J_3BLorg_ray_runtime_objectstore_NativeRayObject_2(
+    JNIEnv *, jclass, jlong, jbyteArray, jobject);
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeGet
+ * Signature: (JLjava/util/List;J)Ljava/util/List;
+ */
+JNIEXPORT jobject JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeGet(
+    JNIEnv *, jclass, jlong, jobject, jlong);
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeWait
+ * Signature: (JLjava/util/List;IJ)Ljava/util/List;
+ */
+JNIEXPORT jobject JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeWait(
+    JNIEnv *, jclass, jlong, jobject, jint, jlong);
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeDelete
+ * Signature: (JLjava/util/List;ZZ)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeDelete(
+    JNIEnv *, jclass, jlong, jobject, jboolean, jboolean);
+
+/*
+ * Class:     org_ray_runtime_objectstore_ObjectInterfaceImpl
+ * Method:    nativeDestroy
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_ray_runtime_objectstore_ObjectInterfaceImpl_nativeDestroy(
+    JNIEnv *, jclass, jlong);
+
+#ifdef __cplusplus
+}
+#endif
+#endif

src/ray/core_worker/store_provider/plasma_store_provider.cc

@@ -3,6 +3,7 @@
 #include "ray/core_worker/context.h"
 #include "ray/core_worker/core_worker.h"
 #include "ray/core_worker/object_interface.h"
+#include "ray/protobuf/gcs.pb.h"
 
 namespace ray {
 
@@ -101,11 +102,14 @@ Status CoreWorkerPlasmaStoreProvider::Get(
             std::make_shared<PlasmaBuffer>(object_buffers[i].data),
             std::make_shared<PlasmaBuffer>(object_buffers[i].metadata));
         unready.erase(object_id);
+        if (IsException(object_buffers[i])) {
+          should_break = true;
+        }
       }
     }
 
     num_attempts += 1;
-    // TODO(zhijunfu): log a message if attempted too many times.
+    WarnIfAttemptedTooManyTimes(num_attempts, unready);
   }
 
   if (was_blocked) {
@@ -144,4 +148,45 @@ Status CoreWorkerPlasmaStoreProvider::Delete(const std::vector<ObjectID> &object
   return raylet_client_->FreeObjects(object_ids, local_only, delete_creating_tasks);
 }
 
+bool CoreWorkerPlasmaStoreProvider::IsException(const plasma::ObjectBuffer &buffer) {
+  // TODO (kfstorm): metadata should be structured.
+  const std::string metadata = buffer.metadata->ToString();
+  const auto error_type_descriptor = ray::rpc::ErrorType_descriptor();
+  for (int i = 0; i < error_type_descriptor->value_count(); i++) {
+    const auto error_type_number = error_type_descriptor->value(i)->number();
+    if (metadata == std::to_string(error_type_number)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+void CoreWorkerPlasmaStoreProvider::WarnIfAttemptedTooManyTimes(
+    int num_attempts, const std::unordered_map<ObjectID, int> &unready) {
+  if (num_attempts % RayConfig::instance().object_store_get_warn_per_num_attempts() ==
+      0) {
+    std::ostringstream oss;
+    size_t printed = 0;
+    for (auto &entry : unready) {
+      if (printed >=
+          RayConfig::instance().object_store_get_max_ids_to_print_in_warning()) {
+        break;
+      }
+      if (printed > 0) {
+        oss << ", ";
+      }
+      oss << entry.first.Hex();
+    }
+    if (printed < unready.size()) {
+      oss << ", etc";
+    }
+    RAY_LOG(WARNING)
+        << "Attempted " << num_attempts << " times to reconstruct objects, but "
+        << "some objects are still unavailable. If this message continues to print,"
+        << " it may indicate that object's creating task is hanging, or something wrong"
+        << " happened in raylet backend. " << unready.size()
+        << " object(s) pending: " << oss.str() << ".";
+  }
+}
+
 }  // namespace ray

src/ray/core_worker/store_provider/plasma_store_provider.h

@@ -60,6 +60,20 @@ class CoreWorkerPlasmaStoreProvider : public CoreWorkerStoreProvider {
                 bool delete_creating_tasks) override;
 
  private:
+  /// Whether the buffer represents an exception object.
+  ///
+  /// \param[in] buffer the object buffer.
+  /// \return Whether it represents an exception object.
+  static bool IsException(const plasma::ObjectBuffer &buffer);
+
+  /// Print a warning if we've attempted too many times, but some objects are still
+  /// unavailable.
+  ///
+  /// \param[in] num_attemps The number of attempted times.
+  /// \param[in] unready The unready objects.
+  static void WarnIfAttemptedTooManyTimes(
+      int num_attempts, const std::unordered_map<ObjectID, int> &unready);
+
   /// Plasma store client.
   plasma::PlasmaClient store_client_;
 

src/ray/protobuf/common.proto

@@ -11,6 +11,12 @@ enum Language {
   CPP = 2;
 }
 
+// Type of a worker.
+enum WorkerType {
+  WORKER = 0;
+  DRIVER = 1;
+}
+
 // Type of a task.
 enum TaskType {
   // Normal task.

src/ray/protobuf/gcs.proto

@@ -267,4 +267,6 @@ enum ErrorType {
   // 2) The object's creating task is already cleaned up from GCS (this currently
   //    crashes raylet).
   OBJECT_UNRECONSTRUCTABLE = 2;
+  // Indicates that a task failed due to user code failure.
+  TASK_EXECUTION_EXCEPTION = 3;
 }

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

@@ -3,39 +3,14 @@
 #include <jni.h>
 
 #include "ray/common/id.h"
+#include "ray/core_worker/lib/java/jni_utils.h"
 #include "ray/raylet/raylet_client.h"
 #include "ray/util/logging.h"
 
-template <typename ID>
-class UniqueIdFromJByteArray {
- public:
-  const ID &GetId() const { return id; }
-
-  UniqueIdFromJByteArray(JNIEnv *env, const jbyteArray &bytes) {
-    std::string id_str(ID::Size(), 0);
-    env->GetByteArrayRegion(bytes, 0, ID::Size(),
-                            reinterpret_cast<jbyte *>(&id_str.front()));
-    id = ID::FromBinary(id_str);
-  }
-
- private:
-  ID id;
-};
-
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-inline bool ThrowRayExceptionIfNotOK(JNIEnv *env, const ray::Status &status) {
-  if (!status.ok()) {
-    jclass exception_class = env->FindClass("org/ray/api/exception/RayException");
-    env->ThrowNew(exception_class, status.message().c_str());
-    return true;
-  } else {
-    return false;
-  }
-}
-
 /*
  * Class:     org_ray_runtime_raylet_RayletClientImpl
  * Method:    nativeInit
@@ -44,11 +19,11 @@ inline bool ThrowRayExceptionIfNotOK(JNIEnv *env, const ray::Status &status) {
 JNIEXPORT jlong JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeInit(
     JNIEnv *env, jclass, jstring sockName, jbyteArray workerId, jboolean isWorker,
     jbyteArray jobId) {
-  UniqueIdFromJByteArray<ClientID> worker_id(env, workerId);
-  UniqueIdFromJByteArray<JobID> job_id(env, jobId);
+  const auto worker_id = JavaByteArrayToId<ClientID>(env, workerId);
+  const auto job_id = JavaByteArrayToId<JobID>(env, jobId);
   const char *nativeString = env->GetStringUTFChars(sockName, JNI_FALSE);
-  auto raylet_client = new RayletClient(nativeString, worker_id.GetId(), isWorker,
-                                        job_id.GetId(), Language::JAVA);
+  auto raylet_client = new std::unique_ptr<RayletClient>(
+      new RayletClient(nativeString, worker_id, isWorker, job_id, Language::JAVA));
   env->ReleaseStringUTFChars(sockName, nativeString);
   return reinterpret_cast<jlong>(raylet_client);
 }
@@ -60,7 +35,7 @@ JNIEXPORT jlong JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeInit(
  */
 JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeSubmitTask(
     JNIEnv *env, jclass, jlong client, jbyteArray taskSpec) {
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
 
   jbyte *data = env->GetByteArrayElements(taskSpec, NULL);
   jsize size = env->GetArrayLength(taskSpec);
@@ -70,7 +45,7 @@ JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeSubmit
 
   ray::TaskSpecification task_spec(task_spec_message);
   auto status = raylet_client->SubmitTask(task_spec);
-  ThrowRayExceptionIfNotOK(env, status);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 
 /*
@@ -80,13 +55,11 @@ JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeSubmit
  */
 JNIEXPORT jbyteArray JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeGetTask(
     JNIEnv *env, jclass, jlong client) {
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
 
   std::unique_ptr<ray::TaskSpecification> spec;
   auto status = raylet_client->GetTask(&spec);
-  if (ThrowRayExceptionIfNotOK(env, status)) {
-    return nullptr;
-  }
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
 
   // Serialize the task spec and copy to Java byte array.
   auto task_data = spec->Serialize();
@@ -109,8 +82,9 @@ JNIEXPORT jbyteArray JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_native
  */
 JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeDestroy(
     JNIEnv *env, jclass, jlong client) {
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  ThrowRayExceptionIfNotOK(env, raylet_client->Disconnect());
+  auto raylet_client = reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
+  auto status = (*raylet_client)->Disconnect();
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
   delete raylet_client;
 }
 
@@ -128,15 +102,14 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeFetchOrReconstruct(
   for (int i = 0; i < len; i++) {
     jbyteArray object_id_bytes =
         static_cast<jbyteArray>(env->GetObjectArrayElement(objectIds, i));
-    UniqueIdFromJByteArray<ObjectID> object_id(env, object_id_bytes);
-    object_ids.push_back(object_id.GetId());
+    const auto object_id = JavaByteArrayToId<ObjectID>(env, object_id_bytes);
+    object_ids.push_back(object_id);
     env->DeleteLocalRef(object_id_bytes);
   }
-  UniqueIdFromJByteArray<TaskID> current_task_id(env, currentTaskId);
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  auto status =
-      raylet_client->FetchOrReconstruct(object_ids, fetchOnly, current_task_id.GetId());
-  ThrowRayExceptionIfNotOK(env, status);
+  const auto current_task_id = JavaByteArrayToId<TaskID>(env, currentTaskId);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
+  auto status = raylet_client->FetchOrReconstruct(object_ids, fetchOnly, current_task_id);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 
 /*
@@ -146,10 +119,10 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeFetchOrReconstruct(
  */
 JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeNotifyUnblocked(
     JNIEnv *env, jclass, jlong client, jbyteArray currentTaskId) {
-  UniqueIdFromJByteArray<TaskID> current_task_id(env, currentTaskId);
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  auto status = raylet_client->NotifyUnblocked(current_task_id.GetId());
-  ThrowRayExceptionIfNotOK(env, status);
+  const auto current_task_id = JavaByteArrayToId<TaskID>(env, currentTaskId);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
+  auto status = raylet_client->NotifyUnblocked(current_task_id);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 
 /*
@@ -166,22 +139,20 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeWaitObject(
   for (int i = 0; i < len; i++) {
     jbyteArray object_id_bytes =
         static_cast<jbyteArray>(env->GetObjectArrayElement(objectIds, i));
-    UniqueIdFromJByteArray<ObjectID> object_id(env, object_id_bytes);
-    object_ids.push_back(object_id.GetId());
+    const auto object_id = JavaByteArrayToId<ObjectID>(env, object_id_bytes);
+    object_ids.push_back(object_id);
     env->DeleteLocalRef(object_id_bytes);
   }
-  UniqueIdFromJByteArray<TaskID> current_task_id(env, currentTaskId);
+  const auto current_task_id = JavaByteArrayToId<TaskID>(env, currentTaskId);
 
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
 
   // Invoke wait.
   WaitResultPair result;
-  auto status = raylet_client->Wait(object_ids, numReturns, timeoutMillis,
-                                    static_cast<bool>(isWaitLocal),
-                                    current_task_id.GetId(), &result);
-  if (ThrowRayExceptionIfNotOK(env, status)) {
-    return nullptr;
-  }
+  auto status =
+      raylet_client->Wait(object_ids, numReturns, timeoutMillis,
+                          static_cast<bool>(isWaitLocal), current_task_id, &result);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
 
   // Convert result to java object.
   jboolean put_value = true;
@@ -216,11 +187,10 @@ JNIEXPORT jbyteArray JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativeGenerateTaskId(
     JNIEnv *env, jclass, jbyteArray jobId, jbyteArray parentTaskId,
     jint parent_task_counter) {
-  UniqueIdFromJByteArray<JobID> job_id(env, jobId);
-  UniqueIdFromJByteArray<TaskID> parent_task_id(env, parentTaskId);
+  const auto job_id = JavaByteArrayToId<JobID>(env, jobId);
+  const auto parent_task_id = JavaByteArrayToId<TaskID>(env, parentTaskId);
 
-  TaskID task_id =
-      ray::GenerateTaskId(job_id.GetId(), parent_task_id.GetId(), parent_task_counter);
+  TaskID task_id = ray::GenerateTaskId(job_id, parent_task_id, parent_task_counter);
   jbyteArray result = env->NewByteArray(task_id.Size());
   if (nullptr == result) {
     return nullptr;
@@ -245,13 +215,13 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeFreePlasmaObjects(
   for (int i = 0; i < len; i++) {
     jbyteArray object_id_bytes =
         static_cast<jbyteArray>(env->GetObjectArrayElement(objectIds, i));
-    UniqueIdFromJByteArray<ObjectID> object_id(env, object_id_bytes);
-    object_ids.push_back(object_id.GetId());
+    const auto object_id = JavaByteArrayToId<ObjectID>(env, object_id_bytes);
+    object_ids.push_back(object_id);
     env->DeleteLocalRef(object_id_bytes);
   }
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
   auto status = raylet_client->FreeObjects(object_ids, localOnly, deleteCreatingTasks);
-  ThrowRayExceptionIfNotOK(env, status);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 
 /*
@@ -263,13 +233,11 @@ JNIEXPORT jbyteArray JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativePrepareCheckpoint(JNIEnv *env, jclass,
                                                                      jlong client,
                                                                      jbyteArray actorId) {
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  UniqueIdFromJByteArray<ActorID> actor_id(env, actorId);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
+  const auto actor_id = JavaByteArrayToId<ActorID>(env, actorId);
   ActorCheckpointID checkpoint_id;
-  auto status = raylet_client->PrepareActorCheckpoint(actor_id.GetId(), checkpoint_id);
-  if (ThrowRayExceptionIfNotOK(env, status)) {
-    return nullptr;
-  }
+  auto status = raylet_client->PrepareActorCheckpoint(actor_id, checkpoint_id);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
   jbyteArray result = env->NewByteArray(checkpoint_id.Size());
   env->SetByteArrayRegion(result, 0, checkpoint_id.Size(),
                           reinterpret_cast<const jbyte *>(checkpoint_id.Data()));
@@ -284,12 +252,11 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativePrepareCheckpoint(JNIEnv *env
 JNIEXPORT void JNICALL
 Java_org_ray_runtime_raylet_RayletClientImpl_nativeNotifyActorResumedFromCheckpoint(
     JNIEnv *env, jclass, jlong client, jbyteArray actorId, jbyteArray checkpointId) {
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  UniqueIdFromJByteArray<ActorID> actor_id(env, actorId);
-  UniqueIdFromJByteArray<ActorCheckpointID> checkpoint_id(env, checkpointId);
-  auto status = raylet_client->NotifyActorResumedFromCheckpoint(actor_id.GetId(),
-                                                                checkpoint_id.GetId());
-  ThrowRayExceptionIfNotOK(env, status);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
+  const auto actor_id = JavaByteArrayToId<ActorID>(env, actorId);
+  const auto checkpoint_id = JavaByteArrayToId<ActorCheckpointID>(env, checkpointId);
+  auto status = raylet_client->NotifyActorResumedFromCheckpoint(actor_id, checkpoint_id);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 
 /*
@@ -300,14 +267,14 @@ Java_org_ray_runtime_raylet_RayletClientImpl_nativeNotifyActorResumedFromCheckpo
 JNIEXPORT void JNICALL Java_org_ray_runtime_raylet_RayletClientImpl_nativeSetResource(
     JNIEnv *env, jclass, jlong client, jstring resourceName, jdouble capacity,
     jbyteArray nodeId) {
-  auto raylet_client = reinterpret_cast<RayletClient *>(client);
-  UniqueIdFromJByteArray<ClientID> node_id(env, nodeId);
+  auto &raylet_client = *reinterpret_cast<std::unique_ptr<RayletClient> *>(client);
+  const auto node_id = JavaByteArrayToId<ClientID>(env, nodeId);
   const char *native_resource_name = env->GetStringUTFChars(resourceName, JNI_FALSE);
 
-  auto status = raylet_client->SetResource(
-      native_resource_name, static_cast<double>(capacity), node_id.GetId());
+  auto status = raylet_client->SetResource(native_resource_name,
+                                           static_cast<double>(capacity), node_id);
   env->ReleaseStringUTFChars(resourceName, native_resource_name);
-  ThrowRayExceptionIfNotOK(env, status);
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, (void)0);
 }
 
 #ifdef __cplusplus