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

java/runtime/src/main/java/io/ray/runtime/AbstractRayRuntime.java

@@ -282,6 +282,11 @@ public abstract class AbstractRayRuntime implements RayRuntimeInternal {
     return objectStore;
   }
 
+  @Override
+  public TaskExecutor getTaskExecutor() {
+    return taskExecutor;
+  }
+
   @Override
   public FunctionManager getFunctionManager() {
     return functionManager;

java/runtime/src/main/java/io/ray/runtime/RayRuntimeInternal.java

@@ -6,6 +6,7 @@ import io.ray.runtime.context.WorkerContext;
 import io.ray.runtime.functionmanager.FunctionManager;
 import io.ray.runtime.gcs.GcsClient;
 import io.ray.runtime.object.ObjectStore;
+import io.ray.runtime.task.TaskExecutor;
 
 /**
  * This interface is required to make {@link RayRuntimeProxy} work.
@@ -21,6 +22,8 @@ public interface RayRuntimeInternal extends RayRuntime {
 
   ObjectStore getObjectStore();
 
+  TaskExecutor getTaskExecutor();
+
   FunctionManager getFunctionManager();
 
   RayConfig getRayConfig();

java/runtime/src/main/java/io/ray/runtime/task/NativeTaskExecutor.java

@@ -45,6 +45,11 @@ public class NativeTaskExecutor extends TaskExecutor<NativeTaskExecutor.NativeAc
     return new NativeActorContext();
   }
 
+  public void onWorkerShutdown(byte[] workerIdBytes) {
+    // This is to make sure no memory leak when `Ray.exitActor()` is called.
+    removeActorContext(new UniqueId(workerIdBytes));
+  }
+
   @Override
   protected void maybeSaveCheckpoint(Object actor, ActorId actorId) {
     if (!(actor instanceof Checkpointable)) {

java/runtime/src/main/java/io/ray/runtime/task/TaskExecutor.java

@@ -65,6 +65,10 @@ public abstract class TaskExecutor<T extends TaskExecutor.ActorContext> {
     this.actorContextMap.put(runtime.getWorkerContext().getCurrentWorkerId(), actorContext);
   }
 
+  protected void removeActorContext(UniqueId workerId) {
+    this.actorContextMap.remove(workerId);
+  }
+
   private RayFunction getRayFunction(List<String> rayFunctionInfo) {
     JobId jobId = runtime.getWorkerContext().getCurrentJobId();
     JavaFunctionDescriptor functionDescriptor = parseFunctionDescriptor(rayFunctionInfo);