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

Code before:
  // Create a shared owner
  rpc::Address owner_address;
  auto owner_node_id = NodeID::FromRandom();
  auto owner_worker_id = WorkerID::FromRandom();
  owner_address.set_node_id(owner_node_id.Binary());
  owner_address.set_ip_address("1234");
  owner_address.set_port(5678);
  owner_address.set_worker_id(owner_worker_id.Binary());

  // Register 3 actors all owned by the same worker (non-detached)
  std::vector<std::shared_ptr<gcs::GcsActor>> actors;
  for (int i = 0; i < 3; i++) {
    auto actor_creation_task_spec = GenActorCreationTask(job_id,
                                                         /*max_restarts=*/0,
                                                         /*detached=*/false,
                                                         /*name=*/"",
                                                         /*ray_namespace=*/"test",
                                                         owner_address);
    rpc::RegisterActorRequest register_request;
    register_request.mutable_task_spec()->CopyFrom(actor_creation_task_spec.GetMessage());

    Status status =
        gcs_actor_manager_->RegisterActor(register_request, [](const Status &) {});
    ASSERT_TRUE(status.ok());
    io_service_.run_one();
    io_service_.run_one();

    auto actor_id = ActorID::FromBinary(
        register_request.task_spec().actor_creation_task_spec().actor_id());
    ASSERT_TRUE(gcs_actor_manager_->registered_actors_.contains(actor_id));
    auto actor = gcs_actor_manager_->registered_actors_[actor_id];

    rpc::CreateActorRequest create_request;
    create_request.mutable_task_spec()->CopyFrom(actor_creation_task_spec.GetMessage());
    RAY_CHECK_OK(gcs_actor_manager_->CreateActor(
        create_request,
        [](std::shared_ptr<gcs::GcsActor>, const rpc::PushTaskReply &, const Status &) {
        }));

    ASSERT_EQ(mock_actor_scheduler_->actors.size(), 1);
    auto scheduled_actor = mock_actor_scheduler_->actors.back();
    mock_actor_scheduler_->actors.pop_back();

    // Make the actor alive on a different node (not the owner's node)
    auto actor_address = RandomAddress();
    scheduled_actor->UpdateAddress(actor_address);
    gcs_actor_manager_->OnActorCreationSuccess(scheduled_actor, rpc::PushTaskReply());
    io_service_.run_one();

    ASSERT_EQ(scheduled_actor->GetState(), rpc::ActorTableData::ALIVE);
    actors.push_back(scheduled_actor);

Review comment (refactor):
Hardcoding the number of times you have to run the io_context is a little fragile. I think @Kunchd introduced a utility function that's less error-prone and more readable. It's conveniently in the same file! 

https://github.com/ray-project/ray/pull/60652/changes#diff-855b531c5bb189015300ee399fdc4928c639a817774087f7a669f13edab595e3R287

Code after:
  // triggering another restart of the actor.
  rpc::RestartActorForLineageReconstructionRequest request3;
  request3.set_actor_id(actor->GetActorID().Binary());
  request3.set_num_restarts_due_to_lineage_reconstruction(1);
  rpc::RestartActorForLineageReconstructionReply reply3;
  gcs_actor_manager_->HandleRestartActorForLineageReconstruction(
      request3, &reply3, [](auto, auto, auto) {});
  ASSERT_EQ(reply3.status().code(), static_cast<int>(StatusCode::OK));
  // Make sure the actor is not restarted again.
  ASSERT_EQ(actor->GetState(), rpc::ActorTableData::ALIVE);
  ASSERT_EQ(actor->GetActorTableData().num_restarts(), 1);
  ASSERT_EQ(actor->GetActorTableData().num_restarts_due_to_lineage_reconstruction(), 1);
}

TEST_F(GcsActorManagerTest, TestDestroyActorWhenActorIsCreating) {
  auto job_id = JobID::FromInt(1);
  auto registered_actor = RegisterActor(job_id, /*max_restarts*/ -1);
  rpc::CreateActorRequest create_actor_request;
  create_actor_request.mutable_task_spec()->CopyFrom(
      registered_actor->GetCreationTaskSpecification().GetMessage());

  std::vector<std::shared_ptr<gcs::GcsActor>> finished_actors;
  Status status = gcs_actor_manager_->CreateActor(
      create_actor_request,
      [&finished_actors](const std::shared_ptr<gcs::GcsActor> &result_actor,
                         const rpc::PushTaskReply &,
                         const Status &) { finished_actors.emplace_back(result_actor); });
  RAY_CHECK_OK(status);

  ASSERT_EQ(finished_actors.size(), 0);
  ASSERT_EQ(mock_actor_scheduler_->actors.size(), 1);
  auto actor = mock_actor_scheduler_->actors.back();
  mock_actor_scheduler_->actors.pop_back();

  // Make sure actor in the phase of creating, that is the worker id is not nil and the
  // actor state is not alive yet.
  actor->UpdateAddress(RandomAddress());
  actor->UpdateLocalRayletAddress(RandomAddress());
  const auto &worker_id = actor->GetWorkerID();
  ASSERT_TRUE(!worker_id.IsNil());
  const auto &local_raylet_address = actor->LocalRayletAddress();
  ASSERT_TRUE(local_raylet_address.has_value());
  ASSERT_NE(actor->GetState(), rpc::ActorTableData::ALIVE);

  // Then handle the kill actor requst (no restart).
  rpc::KillActorViaGcsReply reply;
  rpc::KillActorViaGcsRequest request;
  request.set_actor_id(actor->GetActorID().Binary());
  request.set_force_kill(true);
  // Set the `no_restart` flag to true so that the actor will be destoryed.
  request.set_no_restart(true);