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

Code before:
  for (const auto &id : ref_count.contains()) {
    ref.mutable_nested()->contains.insert(ObjectID::FromBinary(id));
  }
  const auto contained_in_borrowed_ids =
      IdVectorFromProtobuf<ObjectID>(ref_count.contained_in_borrowed_ids());
  ref.mutable_nested()->contained_in_borrowed_ids.insert(
      contained_in_borrowed_ids.begin(), contained_in_borrowed_ids.end());
  return ref;
}

void ReferenceCounter::Reference::ToProto(rpc::ObjectReferenceCount *ref,
                                          const ObjectID &object_id,
                                          bool has_local_ref) const {
  ref->mutable_reference()->set_object_id(object_id.Binary());
  if (owner_address_) {
    ref->mutable_reference()->mutable_owner_address()->CopyFrom(*owner_address_);
  }
  ref->set_has_local_ref(has_local_ref);
  for (const auto &borrower : borrow().borrowers) {
    ref->add_borrowers()->CopyFrom(borrower);
  }
  for (const auto &object : borrow().stored_in_objects) {
    auto ref_object = ref->add_stored_in_objects();
    ref_object->set_object_id(object.first.Binary());
    ref_object->mutable_owner_address()->CopyFrom(object.second);
  }
  for (const auto &contained_in_borrowed_id : nested().contained_in_borrowed_ids) {
    ref->add_contained_in_borrowed_ids(contained_in_borrowed_id.Binary());
  }
  for (const auto &contains_id : nested().contains) {
    ref->add_contains(contains_id.Binary());
  }
}

std::optional<std::string> ReferenceCounter::GetTensorTransport(
    const ObjectID &object_id) const {
  absl::MutexLock lock(&mutex_);
  auto it = object_id_refs_.find(object_id);
  if (it == object_id_refs_.end()) {
    return std::nullopt;
  }
  return it->second.tensor_transport_;
}

}  // namespace core
}  // namespace ray


Review comment (bug):
This is a follow up that's out of the scope of this change since it's mean for private core debugging only, but as we improve upon the debugging experience of the system, we should also add more documentations for what the information we are emitting represent.

Code after:
    return;
  }

  // Always publish the location when subscribed for the first time.
  // This will ensure that the subscriber will get the first snapshot of the
  // object location.
  PushToLocationSubscribers(it);
}

std::string ReferenceCounter::DebugString() const {
  absl::MutexLock lock(&mutex_);
  std::stringstream ss;
  ss << "ReferenceTable{size: " << object_id_refs_.size();
  if (!object_id_refs_.empty()) {
    ss << " sample: " << object_id_refs_.begin()->first << ":"
       << object_id_refs_.begin()->second.DebugString();
  }
  ss << "}";
  return ss.str();
}

json ReferenceCounter::NestedReferenceCount::ToJson() const {
  return {
      {"contained_in_owned", IdContainerToJsonArray(contained_in_owned)},
      {"contained_in_borrowed_ids", IdContainerToJsonArray(contained_in_borrowed_ids)},
      {"contains", IdContainerToJsonArray(contains)}};
}

json ReferenceCounter::BorrowInfo::ToJson() const {
  json stored_in_objects_json = json::array();
  for (const auto &[object_id, addr] : stored_in_objects) {
    stored_in_objects_json.push_back(
        {{"object_id", object_id.Hex()}, {"address", AddressToJson(addr)}});
  }
  json borrowers_json = json::array();
  for (const auto &address : borrowers) {
    borrowers_json.push_back(AddressToJson(address));
  }
  return {{"stored_in_objects", stored_in_objects_json}, {"borrowers", borrowers_json}};
}

std::string ReferenceCounter::ToJsonString() const {
  absl::MutexLock lock(&mutex_);
  json ref_table_json = json::array();
  for (const auto &[obj_id, reference] : object_id_refs_) {
    ref_table_json.push_back({
        {"object_id", obj_id.Hex()},
        {"reference", reference.ToJson()},
    });
  }
  json output = {{"rpc_address", AddressToJson(rpc_address_)},