Repository: ray-project/ray
File: src/ray/gcs/actor/tests/gcs_actor_manager_test.cc (C++)

Code before:
  // Make the actor alive on another node again
  auto new_address_2 = RandomAddress();
  auto new_node_id_2 = NodeID::FromBinary(new_address_2.node_id());
  auto new_worker_id_2 = WorkerID::FromBinary(new_address_2.worker_id());
  actor->UpdateAddress(new_address_2);
  gcs_actor_manager_->OnActorCreationSuccess(actor, rpc::PushTaskReply());
  io_service_.run_one();
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::ALIVE);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 2);
  ASSERT_EQ(actor->GetActorTableData().num_restarts_due_to_node_preemption(), 1);
  ASSERT_FALSE(actor->GetActorTableData().preempted());  // Turn preempted back

  // Third restart: actor reaches max_restarts, so num_restarts and
  // num_restarts_due_to_node_preemption should not increment
  gcs_actor_manager_->OnWorkerDead(new_node_id_2, new_worker_id_2);
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::DEAD);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 2);
  ASSERT_EQ(actor->GetActorTableData().num_restarts_due_to_node_preemption(), 1);
  ASSERT_FALSE(actor->GetActorTableData().preempted());
}

// This test simulates the following get_if_exists=True flow:
// 1. Creates a named actor
// 2. Actor goes out of scope, but ref is not deleted yet
// 3. get_if_exists=True is called (i.e. get, create, get)
// 4. First get should return the existing actor
TEST_F(GcsActorManagerTest, TestGetIfExistsOnOutOfScopeActorReturnsSameActor) {
  auto send_reply_callback = [](Status, std::function<void()>, std::function<void()>) {};
  const JobID job_id = JobID::FromInt(1);
  const std::string actor_name = "test_get_if_exists_actor";
  const std::string ray_namespace = "test_namespace";

  std::shared_ptr<gcs::GcsActor> actor =
      RegisterAndCreateNamedActor(job_id, actor_name, ray_namespace, send_reply_callback);

  // Simulate actor going out of scope
  rpc::ReportActorOutOfScopeRequest out_of_scope_request;
  out_of_scope_request.set_actor_id(actor->GetActorID().Binary());
  out_of_scope_request.set_num_restarts_due_to_lineage_reconstruction(0);
  rpc::ReportActorOutOfScopeReply out_of_scope_reply;
  gcs_actor_manager_->HandleReportActorOutOfScope(
      out_of_scope_request,
      &out_of_scope_reply,
      [](auto status, auto success_callback, auto failure_callback) {});
  while (io_service_.poll_one()) {
  }
  // Actor should now be DEAD
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::DEAD);

  // Simulate get_if_exists=True flow - first get call
  rpc::GetNamedActorInfoRequest get_request2;

Review comment (bug):
Often we write unit tests in response to bugs we find and their fixes. This is good. However, we shouldn't over-specify the test description or the comments to the specific bug but instead to the API of the class we're testing. 

I would remove all references to `get_if_exists` and just say that HandleGetNamedActorInfo returns the Actor even if it's gone out of scope.

Code after:
  io_service_.run_one();
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::ALIVE);

  // Initially num_restarts should be 0
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 0);
  ASSERT_FALSE(actor->GetActorTableData().preempted());

  // First restart: actor is NOT preempted, so num_restarts should increment
  gcs_actor_manager_->OnWorkerDead(node_id, worker_id);
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::RESTARTING);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 1);  // Should increment
  ASSERT_FALSE(actor->GetActorTableData().preempted());

  // Make the actor alive on a specific node again.
  auto new_address = RandomAddress();
  auto new_node_id = NodeID::FromBinary(new_address.node_id());
  auto new_worker_id = WorkerID::FromBinary(new_address.worker_id());
  actor->UpdateAddress(new_address);
  gcs_actor_manager_->OnActorCreationSuccess(actor, rpc::PushTaskReply());
  io_service_.run_one();
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::ALIVE);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 1);
  ASSERT_EQ(actor->GetActorTableData().num_restarts_due_to_node_preemption(), 0);

  // Now set the actor as preempted using SetPreemptedAndPublish
  gcs_actor_manager_->SetPreemptedAndPublish(new_node_id);
  io_service_.run_one();
  ASSERT_TRUE(actor->GetActorTableData().preempted());

  // Second restart: actor is preempted, so num_restarts and
  // num_restarts_due_to_node_preemption should increment
  gcs_actor_manager_->OnWorkerDead(new_node_id, new_worker_id);
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::RESTARTING);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 2);  // Should increment
  ASSERT_EQ(actor->GetActorTableData().num_restarts_due_to_node_preemption(), 1);

  // Make the actor alive on another node again
  auto new_address_2 = RandomAddress();
  auto new_node_id_2 = NodeID::FromBinary(new_address_2.node_id());
  auto new_worker_id_2 = WorkerID::FromBinary(new_address_2.worker_id());
  actor->UpdateAddress(new_address_2);
  gcs_actor_manager_->OnActorCreationSuccess(actor, rpc::PushTaskReply());
  io_service_.run_one();
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::ALIVE);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 2);
  ASSERT_EQ(actor->GetActorTableData().num_restarts_due_to_node_preemption(), 1);
  ASSERT_FALSE(actor->GetActorTableData().preempted());  // Turn preempted back

  // Third restart: actor reaches max_restarts, so num_restarts and
  // num_restarts_due_to_node_preemption should not increment
  gcs_actor_manager_->OnWorkerDead(new_node_id_2, new_worker_id_2);