From a3bf1af10efe754bac9313987cca331454dd5eb8 Mon Sep 17 00:00:00 2001 From: Qing Wang Date: Sat, 11 Dec 2021 07:59:16 +0800 Subject: [PATCH] [core] Fix the risk of iterator invalidation issue. (#20989) We erase the elements from object_id_refs_ in the method `RemoveLocalReferenceInternal()` which may cause iterator invalidation issue. Note that, normally flatmap will not trigger any iterator invalidation except triggering `rehash()`. But in this case, we may remove other elements(not only the current iterator), so there is still a risk of it. --- src/ray/core_worker/reference_count.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/ray/core_worker/reference_count.cc b/src/ray/core_worker/reference_count.cc index 8ee731473..561dca4ac 100644 --- a/src/ray/core_worker/reference_count.cc +++ b/src/ray/core_worker/reference_count.cc @@ -256,11 +256,15 @@ void ReferenceCounter::SetNestedRefInUseRecursive(ReferenceTable::iterator inner void ReferenceCounter::ReleaseAllLocalReferences() { absl::MutexLock lock(&mutex_); + std::vector refs_to_remove; for (auto &ref : object_id_refs_) { for (int i = ref.second.local_ref_count; i > 0; --i) { - RemoveLocalReferenceInternal(ref.first, nullptr); + refs_to_remove.push_back(ref.first); } } + for (const auto &object_id_to_remove : refs_to_remove) { + RemoveLocalReferenceInternal(object_id_to_remove, nullptr); + } } void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id,