[Java] Fix getCurrentActorId in multi-threading scenario. (#5506)

java/test/src/main/java/org/ray/api/test/MultiThreadingTest.java

@@ -15,6 +15,7 @@ import org.ray.api.RayObject;
 import org.ray.api.TestUtils;
 import org.ray.api.WaitResult;
 import org.ray.api.annotation.RayRemote;
+import org.ray.api.id.ActorId;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
@@ -42,6 +43,33 @@ public class MultiThreadingTest extends BaseTest {
     }
   }
 
+  @RayRemote
+  public static class ActorIdTester {
+
+    private final ActorId actorId;
+
+    public ActorIdTester() {
+      actorId = Ray.getRuntimeContext().getCurrentActorId();
+      Assert.assertNotEquals(actorId, ActorId.NIL);
+    }
+
+    @RayRemote
+    public ActorId getCurrentActorId() {
+      final ActorId[] result = new ActorId[1];
+      Thread thread = new Thread(() -> {
+        result[0] = Ray.getRuntimeContext().getCurrentActorId();
+      });
+      thread.start();
+      try {
+        thread.join();
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+      Assert.assertEquals(result[0], actorId);
+      return result[0];
+    }
+  }
+
   public static String testMultiThreading() {
     Random random = new Random();
     // Test calling normal functions.
@@ -105,6 +133,14 @@ public class MultiThreadingTest extends BaseTest {
     Assert.assertEquals("ok", obj.get());
   }
 
+  @Test
+  public void testGetCurrentActorId() {
+    TestUtils.skipTestUnderSingleProcess();
+    RayActor<ActorIdTester> actorIdTester = Ray.createActor(ActorIdTester::new);
+    ActorId actorId = Ray.call(ActorIdTester::getCurrentActorId, actorIdTester).get();
+    Assert.assertEquals(actorId, actorIdTester.getId());
+  }
+
   private static void runTestCaseInMultipleThreads(Runnable testCase, int numRepeats) {
     ExecutorService service = Executors.newFixedThreadPool(NUM_THREADS);
 

src/ray/core_worker/context.cc

@@ -6,10 +6,7 @@ namespace ray {
 /// per-thread context for core worker.
 struct WorkerThreadContext {
   WorkerThreadContext()
-      : current_task_id_(TaskID::ForFakeTask()),
-        current_actor_id_(ActorID::Nil()),
-        task_index_(0),
-        put_index_(0) {}
+      : current_task_id_(TaskID::ForFakeTask()), task_index_(0), put_index_(0) {}
 
   int GetNextTaskIndex() { return ++task_index_; }
 
@@ -21,8 +18,6 @@ struct WorkerThreadContext {
     return current_task_;
   }
 
-  const ActorID &GetCurrentActorID() const { return current_actor_id_; }
-
   void SetCurrentTaskId(const TaskID &task_id) {
     current_task_id_ = task_id;
     task_index_ = 0;
@@ -32,22 +27,12 @@ struct WorkerThreadContext {
   void SetCurrentTask(const TaskSpecification &task_spec) {
     SetCurrentTaskId(task_spec.TaskId());
     current_task_ = std::make_shared<const TaskSpecification>(task_spec);
-    if (task_spec.IsActorCreationTask()) {
-      RAY_CHECK(current_actor_id_.IsNil());
-      current_actor_id_ = task_spec.ActorCreationId();
-    }
-    if (task_spec.IsActorTask()) {
-      RAY_CHECK(current_actor_id_ == task_spec.ActorId());
-    }
   }
 
  private:
   /// The task ID for current task.
   TaskID current_task_id_;
 
-  /// ID of current actor.
-  ActorID current_actor_id_;
-
   /// The current task.
   std::shared_ptr<const TaskSpecification> current_task_;
 
@@ -65,7 +50,8 @@ 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()) {
+      current_job_id_(worker_type_ == WorkerType::DRIVER ? job_id : JobID::Nil()),
+      current_actor_id_(ActorID::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).
@@ -91,14 +77,19 @@ const TaskID &WorkerContext::GetCurrentTaskID() const {
 void WorkerContext::SetCurrentTask(const TaskSpecification &task_spec) {
   current_job_id_ = task_spec.JobId();
   GetThreadContext().SetCurrentTask(task_spec);
+  if (task_spec.IsActorCreationTask()) {
+    RAY_CHECK(current_actor_id_.IsNil());
+    current_actor_id_ = task_spec.ActorCreationId();
+  }
+  if (task_spec.IsActorTask()) {
+    RAY_CHECK(current_actor_id_ == task_spec.ActorId());
+  }
 }
 std::shared_ptr<const TaskSpecification> WorkerContext::GetCurrentTask() const {
   return GetThreadContext().GetCurrentTask();
 }
 
-const ActorID &WorkerContext::GetCurrentActorID() const {
-  return GetThreadContext().GetCurrentActorID();
-}
+const ActorID &WorkerContext::GetCurrentActorID() const { return current_actor_id_; }
 
 WorkerThreadContext &WorkerContext::GetThreadContext() {
   if (thread_context_ == nullptr) {

src/ray/core_worker/context.h

@@ -40,6 +40,9 @@ class WorkerContext {
   /// Job ID for this worker.
   JobID current_job_id_;
 
+  /// ID of current actor.
+  ActorID current_actor_id_;
+
  private:
   static WorkerThreadContext &GetThreadContext();