Repository: ray-project/ray
File: src/ray/core_worker/core_worker.cc (C++)

Code before:
           << placement_group_id
           << ". It is probably "
              "because GCS server is dead or there's a high load there.";
    return Status::TimedOut(stream.str());
  }
  return status;
}

Status CoreWorker::WaitPlacementGroupReady(const PlacementGroupID &placement_group_id,
                                           int64_t timeout_seconds) {
  auto status = gcs_client_->PlacementGroups().SyncWaitUntilReady(placement_group_id,
                                                                  timeout_seconds);
  if (status.IsTimedOut()) {
    std::ostringstream stream;
    stream << "There was timeout in waiting for placement group " << placement_group_id
           << " creation.";
    return Status::TimedOut(stream.str());
  }
  return status;
}

ObjectID CoreWorker::AsyncWaitPlacementGroupReady(
    const PlacementGroupID &placement_group_id,
    const std::string &serialized_object_data,
    const std::string &serialized_object_metadata) {
  // Generate ObjectID and register ownership.
  // The object will be stored directly in memory_store_ and fate-shares with the owner,
  // so we set pinned_at_node_id to nullopt (same as small task returns).
  ObjectID object_id = ObjectID::FromIndex(worker_context_->GetCurrentInternalTaskId(),
                                           worker_context_->GetNextPutIndex());
  reference_counter_->AddOwnedObject(object_id,
                                     /*contained_object_ids=*/{},
                                     rpc_address_,
                                     CurrentCallSite(),
                                     /*object_size=*/-1,
                                     LineageReconstructionEligibility::INELIGIBLE_PUT,
                                     /*add_local_ref=*/true,
                                     /*pinned_at_node_id=*/std::nullopt);

  // Async RPC to GCS that returns when the placement group is ready (or removed).
  // The callback puts the result into memory store, completing ray.get()/wait()/await.
  rpc::WaitPlacementGroupUntilReadyRequest request;
  request.set_placement_group_id(placement_group_id.Binary());

  gcs_client_->GetGcsRpcClient().WaitPlacementGroupUntilReady(
      std::move(request),
      [this, object_id, serialized_object_data, serialized_object_metadata](
          const Status &status, const rpc::WaitPlacementGroupUntilReadyReply &reply) {
        std::shared_ptr<RayObject> result;
        if (status.ok()) {
          auto data = std::make_shared<LocalMemoryBuffer>(serialized_object_data.size());

Review comment (refactor):
We control both the client and server implementations here, so we can `RAY_CHECK` here if we get an unexpected status.

Even better, this is a good use case for the `StatusSet` utility that @dayshah added awhile back. Take a look at the usage example at the top of `status.h` and try using it here. Ideally the RPC client would define the set of statuses that can be returned for this method, and the caller is forced to handle them all.

Code after:

  gcs_client_->GetGcsRpcClient().WaitPlacementGroupUntilReady(
      std::move(request),
      [this, object_id, serialized_object_data, serialized_object_metadata](
          const Status &status, const rpc::WaitPlacementGroupUntilReadyReply &reply) {
        // timeout_ms=-1 retries transient gRPC failures, so any other error
        // here indicates an unexpected GCS server-side failure.
        RAY_CHECK(status.ok() || status.IsNotFound())
            << "Unexpected status from WaitPlacementGroupUntilReady: " << status;

        std::shared_ptr<RayObject> result;
        if (status.ok()) {
          auto data = std::make_shared<LocalMemoryBuffer>(serialized_object_data.size());
          memcpy(
              data->Data(), serialized_object_data.data(), serialized_object_data.size());
          auto metadata =
              std::make_shared<LocalMemoryBuffer>(serialized_object_metadata.size());
          memcpy(metadata->Data(),
                 serialized_object_metadata.data(),
                 serialized_object_metadata.size());
          result = std::make_shared<RayObject>(
              data, metadata, std::vector<rpc::ObjectReference>());
        } else {
          result =
              std::make_shared<RayObject>(rpc::ErrorType::TASK_PLACEMENT_GROUP_REMOVED);
        }
        memory_store_->Put(*result, object_id, /*has_reference=*/true);
      },
      // timeout_ms=-1 means infinite wait with automatic retry.
      // Users can still set their own timeout via ray.get(ref, timeout=...).
      /*timeout_ms=*/-1);

  return object_id;
}

Status CoreWorker::SubmitActorTask(
    const ActorID &actor_id,
    const RayFunction &function,
    const std::vector<std::unique_ptr<TaskArg>> &args,
    const TaskOptions &task_options,
    int max_retries,
    bool retry_exceptions,
    const std::string &serialized_retry_exception_allowlist,
    const std::string &call_site,
    std::vector<rpc::ObjectReference> &task_returns,
    const TaskID current_task_id) {
  SubscribeToNodeChanges();
  absl::ReleasableMutexLock lock(&actor_task_mutex_);
  task_returns.clear();
  if (!actor_task_submitter_->CheckActorExists(actor_id)) {
    std::string err_msg = absl::StrFormat(