Address doc related PR comments

davik · davik · commit 84305355d5f9 · 2025-10-06T23:36:48.000Z
Signed-off-by: davik &lt;davik@anyscale.com&gt;
diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc
@@ -2707,7 +2707,7 @@ Status CoreWorker::ExecuteTask(
     std::vector<std::pair<ObjectID, std::shared_ptr<RayObject>>> *return_objects,
     std::vector<std::pair<ObjectID, std::shared_ptr<RayObject>>> *dynamic_return_objects,
     std::vector<std::pair<ObjectID, bool>> *streaming_generator_returns,
-    ReferenceCounter::ReferenceTableProto *borrowed_refs,
+    ReferenceCounterInterface::ReferenceTableProto *borrowed_refs,
     bool *is_retryable_error,
     std::string *application_error) {
   RAY_LOG(DEBUG) << "Executing task, task info = " << task_spec.DebugString();
@@ -3082,7 +3082,7 @@ Status CoreWorker::ReportGeneratorItemReturns(
     // we borrow the object. When the object value is allocatd, the
     // memory store is updated. We should clear borrowers and memory store
     // here.
-    ReferenceCounter::ReferenceTableProto borrowed_refs;
+    ReferenceCounterInterface::ReferenceTableProto borrowed_refs;
     reference_counter_->PopAndClearLocalBorrowers(
         {dynamic_return_object.first}, &borrowed_refs, &deleted);
     memory_store_->Delete(deleted);
@@ -3152,7 +3152,7 @@ void CoreWorker::HandleReportGeneratorItemReturns(
 std::vector<rpc::ObjectReference> CoreWorker::ExecuteTaskLocalMode(
     const TaskSpecification &task_spec, const ActorID &actor_id) {
   auto return_objects = std::vector<std::pair<ObjectID, std::shared_ptr<RayObject>>>();
-  auto borrowed_refs = ReferenceCounter::ReferenceTableProto();
+  auto borrowed_refs = ReferenceCounterInterface::ReferenceTableProto();
 
   std::vector<rpc::ObjectReference> returned_refs;
   size_t num_returns = task_spec.NumReturns();
diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h
@@ -1483,7 +1483,7 @@ class CoreWorker {
       std::vector<std::pair<ObjectID, std::shared_ptr<RayObject>>>
           *dynamic_return_objects,
       std::vector<std::pair<ObjectID, bool>> *streaming_generator_returns,
-      ReferenceCounter::ReferenceTableProto *borrowed_refs,
+      ReferenceCounterInterface::ReferenceTableProto *borrowed_refs,
       bool *is_retryable_error,
       std::string *application_error);
 
diff --git a/src/ray/core_worker/reference_counter.cc b/src/ray/core_worker/reference_counter.cc
@@ -960,7 +960,7 @@ ReferenceCounter::GetAllReferenceCounts() const {
 
 void ReferenceCounter::PopAndClearLocalBorrowers(
     const std::vector<ObjectID> &borrowed_ids,
-    ReferenceCounter::ReferenceTableProto *proto,
+    ReferenceTableProto *proto,
     std::vector<ObjectID> *deleted) {
   absl::MutexLock lock(&mutex_);
   ReferenceProtoTable borrowed_refs;
diff --git a/src/ray/core_worker/reference_counter.h b/src/ray/core_worker/reference_counter.h
@@ -43,12 +43,6 @@ namespace core {
 class ReferenceCounter : public ReferenceCounterInterface,
                          public LocalityDataProviderInterface {
  public:
-  using ReferenceTableProto =
-      ::google::protobuf::RepeatedPtrField<rpc::ObjectReferenceCount>;
-  using ReferenceRemovedCallback = std::function<void(const ObjectID &)>;
-  using LineageReleasedCallback =
-      std::function<int64_t(const ObjectID &, std::vector<ObjectID> *)>;
-
   ReferenceCounter(rpc::Address rpc_address,
                    pubsub::PublisherInterface *object_info_publisher,
                    pubsub::SubscriberInterface *object_info_subscriber,
diff --git a/src/ray/core_worker/reference_counter_interface.h b/src/ray/core_worker/reference_counter_interface.h
@@ -30,9 +30,10 @@
 namespace ray {
 namespace core {
 
-// Interface for mocking.
 class ReferenceCounterInterface {
  public:
+  using ReferenceTableProto =
+      ::google::protobuf::RepeatedPtrField<rpc::ObjectReferenceCount>;
   using LineageReleasedCallback =
       std::function<int64_t(const ObjectID &, std::vector<ObjectID> *)>;
 
@@ -213,7 +214,6 @@ class ReferenceCounterInterface {
   /// if one exists. An outer_id may not exist if object_id was inlined
   /// directly in a task spec, or if it was passed in the application
   /// out-of-band.
-  /// task ID (for non-actors) or the actor ID of the owner.
   /// \param[in] owner_address The owner's address.
   virtual bool AddBorrowedObject(const ObjectID &object_id,
                                  const ObjectID &outer_id,
@@ -241,7 +241,7 @@ class ReferenceCounterInterface {
 
   //// Checks to see if objects have an owner.
   ///
-  /// \param[in] object_id The ID of the object.
+  /// \param[in] object_ids The IDs of the objects.
   /// \return StatusT::OK if all objects have owners.
   /// \return StatusT::NotFound if any object does not have an owner. The error message
   /// contains objects without owners.
diff --git a/src/ray/core_worker/task_manager.cc b/src/ray/core_worker/task_manager.cc
@@ -1329,7 +1329,7 @@ void TaskManager::FailPendingTask(const TaskID &task_id,
   RemoveFinishedTaskReferences(spec,
                                /*release_lineage=*/true,
                                rpc::Address(),
-                               ReferenceCounter::ReferenceTableProto());
+                               ReferenceCounterInterface::ReferenceTableProto());
 
   MarkTaskReturnObjectsFailed(spec, error_type, ray_error_info, store_in_plasma_ids);
 
@@ -1399,7 +1399,7 @@ void TaskManager::RemoveFinishedTaskReferences(
     TaskSpecification &spec,
     bool release_lineage,
     const rpc::Address &borrower_addr,
-    const ReferenceCounter::ReferenceTableProto &borrowed_refs) {
+    const ReferenceCounterInterface::ReferenceTableProto &borrowed_refs) {
   std::vector<ObjectID> plasma_dependencies = ExtractPlasmaDependencies(spec);
 
   std::vector<ObjectID> return_ids;
diff --git a/src/ray/core_worker/task_manager.h b/src/ray/core_worker/task_manager.h
@@ -645,7 +645,7 @@ class TaskManager : public TaskManagerInterface {
       TaskSpecification &spec,
       bool release_lineage,
       const rpc::Address &worker_addr,
-      const ReferenceCounter::ReferenceTableProto &borrowed_refs);
+      const ReferenceCounterInterface::ReferenceTableProto &borrowed_refs);
 
   /// Get the objects that were stored in plasma upon the first successful
   /// execution of this task. If the task is re-executed, these objects should
diff --git a/src/ray/core_worker/tests/reference_counter_test.cc b/src/ray/core_worker/tests/reference_counter_test.cc
@@ -37,7 +37,7 @@ namespace ray {
 namespace core {
 
 static const rpc::Address empty_borrower;
-static const ReferenceCounter::ReferenceTableProto empty_refs;
+static const ReferenceCounterInterface::ReferenceTableProto empty_refs;
 
 class ReferenceCountTest : public ::testing::Test {
  protected:
@@ -406,7 +406,7 @@ class MockWorkerClient : public MockCoreWorkerClientInterface {
     return return_id;
   }
 
-  ReferenceCounter::ReferenceTableProto FinishExecutingTask(
+  ReferenceCounterInterface::ReferenceTableProto FinishExecutingTask(
       const ObjectID &arg_id,
       const ObjectID &return_id,
       const ObjectID *return_wrapped_id = nullptr,
@@ -415,7 +415,7 @@ class MockWorkerClient : public MockCoreWorkerClientInterface {
       rc_.AddNestedObjectIds(return_id, {*return_wrapped_id}, *owner_address);
     }
 
-    ReferenceCounter::ReferenceTableProto refs;
+    ReferenceCounterInterface::ReferenceTableProto refs;
     if (!arg_id.IsNil()) {
       rc_.PopAndClearLocalBorrowers({arg_id}, &refs, nullptr);
     }
@@ -427,7 +427,7 @@ class MockWorkerClient : public MockCoreWorkerClientInterface {
       const ObjectID &arg_id,
       const absl::flat_hash_map<ObjectID, std::vector<ObjectID>> &nested_return_ids = {},
       const rpc::Address &borrower_address = empty_borrower,
-      const ReferenceCounter::ReferenceTableProto &borrower_refs = empty_refs) {
+      const ReferenceCounterInterface::ReferenceTableProto &borrower_refs = empty_refs) {
     std::vector<ObjectID> arguments;
     for (const auto &pair : nested_return_ids) {
       // NOTE(swang): https://github.com/ray-project/ray/issues/17553.
@@ -2735,7 +2735,7 @@ TEST_F(ReferenceCountTest, TestGetObjectStatusReplyDelayed) {
   ASSERT_TRUE(rc->HasReference(outer_id));
   // Task finishes and our local ref to the outer ObjectRef is deleted. We
   // return borrower information to the owner.
-  ReferenceCounter::ReferenceTableProto refs_proto;
+  ReferenceCounterInterface::ReferenceTableProto refs_proto;
   rc->PopAndClearLocalBorrowers({outer_id}, &refs_proto, nullptr);
   ASSERT_FALSE(rc->HasReference(outer_id));
   // Future resolution is async, so we may receive information about the inner
