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

Code before:
      if (inner_it == object_id_refs_.end()) {
        inner_it = object_id_refs_.emplace(inner_id, Reference()).first;
      }
      // Add the task's caller as a borrower.
      if (inner_it->second.owned_by_us_) {
        auto inserted =
            inner_it->second.mutable_borrow()->borrowers.insert(owner_address).second;
        if (inserted) {
          // Wait for it to remove its reference.
          WaitForRefRemoved(inner_it, owner_address, object_id);
        }
      } else {
        inner_it->second.mutable_borrow()->stored_in_objects.emplace(object_id,
                                                                     owner_address);
      }
      PRINT_REF_COUNT(inner_it);
    }
  }
}

void ReferenceCounter::PublishRefRemoved(const ObjectID &object_id) {
  absl::MutexLock lock(&mutex_);
  PublishRefRemovedInternal(object_id);
}

void ReferenceCounter::PublishRefRemovedInternal(const ObjectID &object_id) {
  RAY_LOG(DEBUG).WithField(object_id) << "PublishRefRemoved ";
  auto it = object_id_refs_.find(object_id);
  if (it != object_id_refs_.end()) {
    PRINT_REF_COUNT(it);
  }

  rpc::PubMessage pub_message;
  pub_message.set_key_id(object_id.Binary());
  pub_message.set_channel_type(rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL);
  auto *worker_ref_removed_message = pub_message.mutable_worker_ref_removed_message();

  RAY_UNUSED(GetAndClearLocalBorrowersInternal(
      object_id,
      /*for_ref_removed=*/true,
      /*deduct_local_ref=*/false,
      worker_ref_removed_message->mutable_borrowed_refs()));

  // Send the owner information about any new borrowers.
  RAY_LOG(DEBUG).WithField(object_id)
      << "Publishing WaitForRefRemoved message for object, message has "
      << worker_ref_removed_message->borrowed_refs().size() << " borrowed references.";
  object_info_publisher_->Publish(std::move(pub_message));
}

void ReferenceCounter::SubscribeRefRemoved(const ObjectID &object_id,

Review comment (suggestion):
this was moved inside `GetAndClearLocalBorrowersInternal` for implementation convenience

Code after:
          << BuildAddress(owner_address.ip_address(), owner_address.port())
          << " to object, borrower owns outer ID " << object_id;
      auto inner_it = object_id_refs_.find(inner_id);
      if (inner_it == object_id_refs_.end()) {
        inner_it = object_id_refs_.emplace(inner_id, Reference()).first;
      }
      // Add the task's caller as a borrower.
      if (inner_it->second.owned_by_us_) {
        auto inserted =
            inner_it->second.mutable_borrow()->borrowers.insert(owner_address).second;
        if (inserted) {
          // Wait for it to remove its reference.
          WaitForRefRemoved(inner_it, owner_address, object_id);
        }
      } else {
        inner_it->second.mutable_borrow()->stored_in_objects.emplace(object_id,
                                                                     owner_address);
      }
      PRINT_REF_COUNT(inner_it);
    }
  }
}

void ReferenceCounter::PublishRefRemoved(const ObjectID &object_id) {
  absl::MutexLock lock(&mutex_);
  PublishRefRemovedInternal(object_id);
}

void ReferenceCounter::PublishRefRemovedInternal(const ObjectID &object_id) {
  RAY_LOG(DEBUG).WithField(object_id) << "PublishRefRemoved ";
  auto it = object_id_refs_.find(object_id);
  if (it != object_id_refs_.end()) {
    PRINT_REF_COUNT(it);
  }

  rpc::PubMessage pub_message;
  pub_message.set_key_id(object_id.Binary());
  pub_message.set_channel_type(rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL);
  auto *worker_ref_removed_message = pub_message.mutable_worker_ref_removed_message();

  RAY_UNUSED(GetAndClearLocalBorrowersInternal(
      object_id,
      /*for_ref_removed=*/true,
      /*deduct_local_ref=*/false,
      worker_ref_removed_message->mutable_borrowed_refs()));

  // Send the owner information about any new borrowers.
  RAY_LOG(DEBUG).WithField(object_id)
      << "Publishing WaitForRefRemoved message for object, message has "
      << worker_ref_removed_message->borrowed_refs().size() << " borrowed references.";
  object_info_publisher_->Publish(std::move(pub_message));