[Java] Attach owner address for pass-by-reference task arguments (#9634)

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

@@ -276,10 +276,7 @@ Java_io_ray_runtime_RayNativeRuntime_nativeGetActorIdOfNamedActor(JNIEnv *env, j
   } else {
     actor_id = ray::ActorID::Nil();
   }
-  jbyteArray bytes = env->NewByteArray(actor_id.Size());
-  env->SetByteArrayRegion(bytes, 0, actor_id.Size(),
-                          reinterpret_cast<const jbyte *>(actor_id.Data()));
-  return bytes;
+  return IdToJavaByteArray<ray::ActorID>(env, actor_id);
 }
 
 JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeKillActor(

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

@@ -51,10 +51,8 @@ JNIEXPORT jbyteArray JNICALL Java_io_ray_runtime_actor_NativeActorHandle_nativeS
   ObjectID actor_handle_id;
   ray::Status status = ray::CoreWorkerProcess::GetCoreWorker().SerializeActorHandle(
       actor_id, &output, &actor_handle_id);
-  jbyteArray bytes = env->NewByteArray(output.size());
-  env->SetByteArrayRegion(bytes, 0, output.size(),
-                          reinterpret_cast<const jbyte *>(output.c_str()));
-  return bytes;
+  THROW_EXCEPTION_AND_RETURN_IF_NOT_OK(env, status, nullptr);
+  return NativeStringToJavaByteArray(env, output);
 }
 
 JNIEXPORT jbyteArray JNICALL

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

@@ -66,6 +66,12 @@ Java_io_ray_runtime_context_NativeWorkerContext_nativeGetCurrentActorId(JNIEnv *
   return IdToJavaByteBuffer<ray::ActorID>(env, actor_id);
 }
 
+JNIEXPORT jbyteArray JNICALL
+Java_io_ray_runtime_context_NativeWorkerContext_nativeGetRpcAddress(JNIEnv *env, jclass) {
+  const auto &rpc_address = ray::CoreWorkerProcess::GetCoreWorker().GetRpcAddress();
+  return NativeStringToJavaByteArray(env, rpc_address.SerializeAsString());
+}
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -63,6 +63,14 @@ Java_io_ray_runtime_context_NativeWorkerContext_nativeGetCurrentWorkerId(JNIEnv
 JNIEXPORT jobject JNICALL
 Java_io_ray_runtime_context_NativeWorkerContext_nativeGetCurrentActorId(JNIEnv *, jclass);
 
+/*
+ * Class:     io_ray_runtime_context_NativeWorkerContext
+ * Method:    nativeGetRpcAddress
+ * Signature: ()[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_io_ray_runtime_context_NativeWorkerContext_nativeGetRpcAddress(JNIEnv *, jclass);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -1,3 +1,17 @@
+// Copyright 2017 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 /* DO NOT EDIT THIS FILE - it is machine generated */
 #include <jni.h>
 /* Header for class io_ray_runtime_metric_NativeMetric */

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

@@ -174,6 +174,15 @@ Java_io_ray_runtime_object_NativeObjectStore_nativeGetAllReferenceCounts(JNIEnv
       });
 }
 
+JNIEXPORT jbyteArray JNICALL
+Java_io_ray_runtime_object_NativeObjectStore_nativeGetOwnerAddress(JNIEnv *env, jclass,
+                                                                   jbyteArray objectId) {
+  auto object_id = JavaByteArrayToId<ray::ObjectID>(env, objectId);
+  const auto &rpc_address =
+      ray::CoreWorkerProcess::GetCoreWorker().GetOwnerAddress(object_id);
+  return NativeStringToJavaByteArray(env, rpc_address.SerializeAsString());
+}
+
 JNIEXPORT jbyteArray JNICALL
 Java_io_ray_runtime_object_NativeObjectStore_nativePromoteAndGetOwnershipInfo(
     JNIEnv *env, jclass, jbyteArray objectId) {

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

@@ -94,6 +94,15 @@ JNIEXPORT jobject JNICALL
 Java_io_ray_runtime_object_NativeObjectStore_nativeGetAllReferenceCounts(JNIEnv *,
                                                                          jclass);
 
+/*
+ * Class:     io_ray_runtime_object_NativeObjectStore
+ * Method:    nativeGetOwnerAddress
+ * Signature: ([B)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_io_ray_runtime_object_NativeObjectStore_nativeGetOwnerAddress(JNIEnv *, jclass,
+                                                                   jbyteArray);
+
 /*
  * Class:     io_ray_runtime_object_NativeObjectStore
  * Method:    nativePromoteAndGetOwnershipInfo

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

@@ -37,10 +37,7 @@ Java_io_ray_runtime_task_NativeTaskExecutor_nativePrepareCheckpoint(JNIEnv *env,
   ActorCheckpointID checkpoint_id;
   auto status = core_worker.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()));
-  return result;
+  return IdToJavaByteArray<ActorCheckpointID>(env, checkpoint_id);
 }
 
 JNIEXPORT void JNICALL

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

@@ -62,8 +62,14 @@ inline std::vector<std::unique_ptr<ray::TaskArg>> ToTaskArgs(JNIEnv *env, jobjec
               env->CallObjectMethod(java_id, java_base_id_get_bytes));
           RAY_CHECK_JAVA_EXCEPTION(env);
           auto id = JavaByteArrayToId<ray::ObjectID>(env, java_id_bytes);
-          return std::unique_ptr<ray::TaskArg>(new ray::TaskArgByReference(
-              id, ray::CoreWorkerProcess::GetCoreWorker().GetOwnerAddress(id)));
+          auto java_owner_address =
+              env->GetObjectField(arg, java_function_arg_owner_address);
+          RAY_CHECK(java_owner_address);
+          auto owner_address =
+              JavaProtobufObjectToNativeProtobufObject<ray::rpc::Address>(
+                  env, java_owner_address);
+          return std::unique_ptr<ray::TaskArg>(
+              new ray::TaskArgByReference(id, owner_address));
         }
         auto java_value =
             static_cast<jbyteArray>(env->GetObjectField(arg, java_function_arg_value));

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

@@ -62,6 +62,9 @@ jmethodID java_jni_exception_util_get_stack_trace;
 jclass java_base_id_class;
 jmethodID java_base_id_get_bytes;
 
+jclass java_abstract_message_lite_class;
+jmethodID java_abstract_message_lite_to_byte_array;
+
 jclass java_function_descriptor_class;
 jmethodID java_function_descriptor_get_language;
 jmethodID java_function_descriptor_to_list;
@@ -71,6 +74,7 @@ jmethodID java_language_get_number;
 
 jclass java_function_arg_class;
 jfieldID java_function_arg_id;
+jfieldID java_function_arg_owner_address;
 jfieldID java_function_arg_value;
 
 jclass java_base_task_options_class;
@@ -183,6 +187,11 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved) {
   java_base_id_class = LoadClass(env, "io/ray/api/id/BaseId");
   java_base_id_get_bytes = env->GetMethodID(java_base_id_class, "getBytes", "()[B");
 
+  java_abstract_message_lite_class =
+      LoadClass(env, "com/google/protobuf/AbstractMessage");
+  java_abstract_message_lite_to_byte_array =
+      env->GetMethodID(java_abstract_message_lite_class, "toByteArray", "()[B");
+
   java_function_descriptor_class =
       LoadClass(env, "io/ray/runtime/functionmanager/FunctionDescriptor");
   java_function_descriptor_get_language =
@@ -197,6 +206,9 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved) {
   java_function_arg_class = LoadClass(env, "io/ray/runtime/task/FunctionArg");
   java_function_arg_id =
       env->GetFieldID(java_function_arg_class, "id", "Lio/ray/api/id/ObjectId;");
+  java_function_arg_owner_address =
+      env->GetFieldID(java_function_arg_class, "ownerAddress",
+                      "Lio/ray/runtime/generated/Common$Address;");
   java_function_arg_value = env->GetFieldID(java_function_arg_class, "value",
                                             "Lio/ray/runtime/object/NativeRayObject;");
 
@@ -278,6 +290,7 @@ void JNI_OnUnload(JavaVM *vm, void *reserved) {
   env->DeleteGlobalRef(java_ray_intentional_system_exit_exception_class);
   env->DeleteGlobalRef(java_jni_exception_util_class);
   env->DeleteGlobalRef(java_base_id_class);
+  env->DeleteGlobalRef(java_abstract_message_lite_class);
   env->DeleteGlobalRef(java_function_descriptor_class);
   env->DeleteGlobalRef(java_language_class);
   env->DeleteGlobalRef(java_function_arg_class);

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

@@ -107,6 +107,11 @@ extern jclass java_base_id_class;
 /// getBytes method of BaseId class
 extern jmethodID java_base_id_get_bytes;
 
+/// AbstractMessageLite class
+extern jclass java_abstract_message_lite_class;
+/// toByteArray method of AbstractMessageLite class
+extern jmethodID java_abstract_message_lite_to_byte_array;
+
 /// FunctionDescriptor interface
 extern jclass java_function_descriptor_class;
 /// getLanguage method of FunctionDescriptor interface
@@ -123,6 +128,8 @@ extern jmethodID java_language_get_number;
 extern jclass java_function_arg_class;
 /// id field of FunctionArg class
 extern jfieldID java_function_arg_id;
+/// ownerAddress field of FunctionArg class
+extern jfieldID java_function_arg_owner_address;
 /// value field of FunctionArg class
 extern jfieldID java_function_arg_value;
 
@@ -528,6 +535,27 @@ inline jobject NativeRayFunctionDescriptorToJavaStringList(
   return NativeStringVectorToJavaStringList(env, std::vector<std::string>());
 }
 
+/// Convert a Java protobuf object to a C++ protobuf object
+template <typename NativeT>
+inline NativeT JavaProtobufObjectToNativeProtobufObject(JNIEnv *env, jobject java_obj) {
+  NativeT native_obj;
+  if (java_obj) {
+    jbyteArray bytes = static_cast<jbyteArray>(
+        env->CallObjectMethod(java_obj, java_abstract_message_lite_to_byte_array));
+    RAY_CHECK_JAVA_EXCEPTION(env);
+    RAY_CHECK(bytes != nullptr);
+    auto buffer = JavaByteArrayToNativeBuffer(env, bytes);
+    RAY_CHECK(buffer);
+    native_obj.ParseFromArray(buffer->Data(), buffer->Size());
+    // Destroy the buffer before deleting the local ref of `bytes`. We need to make sure
+    // that `bytes` is still available when invoking the destructor of
+    // `JavaByteArrayBuffer`.
+    buffer.reset();
+    env->DeleteLocalRef(bytes);
+  }
+  return native_obj;
+}
+
 // Return an actor fullname with job id prepended if this tis a global actor.
 inline std::string GetActorFullName(bool global, std::string name) {
   if (name.empty()) {