[Java] Release actor instance reference when Ray.exitActor() is invoked (#11324)

src/ray/core_worker/core_worker.cc

@@ -541,6 +541,9 @@ void CoreWorker::Shutdown() {
   if (options_.worker_type == WorkerType::WORKER) {
     task_execution_service_.stop();
   }
+  if (options_.on_worker_shutdown) {
+    options_.on_worker_shutdown(GetWorkerID());
+  }
 }
 
 void CoreWorker::Disconnect() {

src/ray/core_worker/core_worker.h

@@ -123,6 +123,8 @@ struct CoreWorkerOptions {
   std::string stderr_file;
   /// Language worker callback to execute tasks.
   TaskExecutionCallback task_execution_callback;
+  /// The callback to be called when shutting down a `CoreWorker` instance.
+  std::function<void(const WorkerID &)> on_worker_shutdown;
   /// Application-language callback to check for signals that have been received
   /// since calling into C++. This will be called periodically (at least every
   /// 1s) during long-running operations. If the function returns anything but StatusOK,

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

@@ -150,6 +150,7 @@ JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeInitialize(
         if (throwable &&
             env->IsInstanceOf(throwable,
                               java_ray_intentional_system_exit_exception_class)) {
+          env->ExceptionClear();
           return ray::Status::IntentionalSystemExit();
         }
         RAY_CHECK_JAVA_EXCEPTION(env);
@@ -211,6 +212,16 @@ JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeInitialize(
     }
   };
 
+  auto on_worker_shutdown = [](const ray::WorkerID &worker_id) {
+    JNIEnv *env = GetJNIEnv();
+    auto worker_id_bytes = IdToJavaByteArray<ray::WorkerID>(env, worker_id);
+    if (java_task_executor) {
+      env->CallVoidMethod(java_task_executor,
+                          java_native_task_executor_on_worker_shutdown, worker_id_bytes);
+      RAY_CHECK_JAVA_EXCEPTION(env);
+    }
+  };
+
   std::string serialized_job_config =
       (jobConfig == nullptr ? "" : JavaByteArrayToNativeString(env, jobConfig));
   ray::CoreWorkerOptions options;
@@ -229,6 +240,7 @@ JNIEXPORT void JNICALL Java_io_ray_runtime_RayNativeRuntime_nativeInitialize(
   options.raylet_ip_address = JavaStringToNativeString(env, nodeIpAddress);
   options.driver_name = JavaStringToNativeString(env, driverName);
   options.task_execution_callback = task_execution_callback;
+  options.on_worker_shutdown = on_worker_shutdown;
   options.gc_collect = gc_collect;
   options.ref_counting_enabled = true;
   options.num_workers = static_cast<int>(numWorkersPerProcess);

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

@@ -107,6 +107,9 @@ jclass java_task_executor_class;
 jmethodID java_task_executor_parse_function_arguments;
 jmethodID java_task_executor_execute;
 
+jclass java_native_task_executor_class;
+jmethodID java_native_task_executor_on_worker_shutdown;
+
 jclass java_placement_group_class;
 jfieldID java_placement_group_id;
 
@@ -267,6 +270,10 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved) {
   java_task_executor_execute =
       env->GetMethodID(java_task_executor_class, "execute",
                        "(Ljava/util/List;Ljava/util/List;)Ljava/util/List;");
+  java_native_task_executor_class =
+      LoadClass(env, "io/ray/runtime/task/NativeTaskExecutor");
+  java_native_task_executor_on_worker_shutdown =
+      env->GetMethodID(java_native_task_executor_class, "onWorkerShutdown", "([B)V");
   return CURRENT_JNI_VERSION;
 }
 
@@ -298,4 +305,5 @@ void JNI_OnUnload(JavaVM *vm, void *reserved) {
   env->DeleteGlobalRef(java_actor_creation_options_class);
   env->DeleteGlobalRef(java_native_ray_object_class);
   env->DeleteGlobalRef(java_task_executor_class);
+  env->DeleteGlobalRef(java_native_task_executor_class);
 }

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

@@ -187,6 +187,11 @@ extern jmethodID java_task_executor_parse_function_arguments;
 /// execute method of TaskExecutor class
 extern jmethodID java_task_executor_execute;
 
+/// NativeTaskExecutor class
+extern jclass java_native_task_executor_class;
+/// onWorkerShutdown method of NativeTaskExecutor class
+extern jmethodID java_native_task_executor_on_worker_shutdown;
+
 /// PlacementGroup class
 extern jclass java_placement_group_class;
 /// id field of PlacementGroup class