diff --git a/java/test/src/main/java/org/ray/api/test/MultiThreadingTest.java b/java/test/src/main/java/org/ray/api/test/MultiThreadingTest.java index b9e76add1..86fdc1b94 100644 --- a/java/test/src/main/java/org/ray/api/test/MultiThreadingTest.java +++ b/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 = 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); diff --git a/src/ray/core_worker/context.cc b/src/ray/core_worker/context.cc index 896f8a70e..663305e4c 100644 --- a/src/ray/core_worker/context.cc +++ b/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(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 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 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) { diff --git a/src/ray/core_worker/context.h b/src/ray/core_worker/context.h index 77e9e2814..3a02c6915 100644 --- a/src/ray/core_worker/context.h +++ b/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();