Skip to content

Commit

Permalink
Fix memory pool lock order reversion detected by Meta internal tsan t…
Browse files Browse the repository at this point in the history 
…est (facebookincubator#10051)

Summary:
We tsan lock to access memory reclaimer which is no-op in release build. Since we now
report the actual used memory in memory pool stats which might grab the child lock while
holding pool lock. There is another code path in child pool creation which grab the pool
lock to access parent memory reclaimer (which is no-op in release build) while holding the
child lock on child pool creation. This PR un-break this deadlock in tsan build by moving
the memory reclaimer check in parent pool without holding any lock.

Also fix tsan failures in MockSharedArbitratorTest

Pull Request resolved: facebookincubator#10051

Reviewed By: tanjialiang

Differential Revision: D58171780

Pulled By: xiaoxmeng

fbshipit-source-id: a2872cfd3ccf421a00d2cafa5fdaed2cdc03b109
  • Loading branch information
@xiaoxmeng @facebook-github-bot
xiaoxmeng authored and facebook-github-bot committed Jun 5, 2024
1 parent 3a7f8a8 commit a6bb4d8
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 47 deletions.
30 changes: 18 additions & 12 deletions velox/common/memory/MemoryPool.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -307,8 +307,15 @@ void MemoryPool::visitChildren(
std::shared_ptr<MemoryPool> MemoryPool::addLeafChild(
const std::string& name,
bool threadSafe,
std::unique_ptr<MemoryReclaimer> reclaimer) {
std::unique_ptr<MemoryReclaimer> _reclaimer) {
CHECK_POOL_MANAGEMENT_OP(addLeafChild);
// NOTE: we shall only set reclaimer in a child pool if its parent has also
// set. Otherwise it should be mis-configured.
VELOX_CHECK(
reclaimer() != nullptr || _reclaimer == nullptr,
"Child memory pool {} shall only set memory reclaimer if its parent {} has also set",
name,
name_);

std::unique_lock guard{poolMutex_};
VELOX_CHECK_EQ(
Expand All @@ -322,15 +329,22 @@ std::shared_ptr<MemoryPool> MemoryPool::addLeafChild(
name,
MemoryPool::Kind::kLeaf,
threadSafe,
std::move(reclaimer));
std::move(_reclaimer));
children_.emplace(name, child);
return child;
}

std::shared_ptr<MemoryPool> MemoryPool::addAggregateChild(
const std::string& name,
std::unique_ptr<MemoryReclaimer> reclaimer) {
std::unique_ptr<MemoryReclaimer> _reclaimer) {
CHECK_POOL_MANAGEMENT_OP(addAggregateChild);
// NOTE: we shall only set reclaimer in a child pool if its parent has also
// set. Otherwise it should be mis-configured.
VELOX_CHECK(
reclaimer() != nullptr || _reclaimer == nullptr,
"Child memory pool {} shall only set memory reclaimer if its parent {} has also set",
name,
name_);

std::unique_lock guard{poolMutex_};
VELOX_CHECK_EQ(
Expand All @@ -344,7 +358,7 @@ std::shared_ptr<MemoryPool> MemoryPool::addAggregateChild(
name,
MemoryPool::Kind::kAggregate,
true,
std::move(reclaimer));
std::move(_reclaimer));
children_.emplace(name, child);
return child;
}
Expand Down Expand Up @@ -413,14 +427,6 @@ MemoryPoolImpl::MemoryPoolImpl(
// actually used memory arbitration policy.
capacity_(parent_ != nullptr ? kMaxMemory : 0) {
VELOX_CHECK(options.threadSafe || isLeaf());
// NOTE: we shall only set reclaimer in a child pool if its parent has also
// set. Otherwise. it should be mis-configured.
VELOX_CHECK(
parent_ == nullptr || parent_->reclaimer() != nullptr ||
reclaimer_ == nullptr,
"Child memory pool {} shall only set memory reclaimer if its parent {} has also set",
name_,
parent_->name());
VELOX_CHECK(
isRoot() || (destructionCb_ == nullptr && growCapacityCb_ == nullptr),
"Only root memory pool allows to set destruction and capacity grow callbacks: {}",
Expand Down
79 changes: 44 additions & 35 deletions velox/common/memory/tests/MockSharedArbitratorTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -272,27 +272,32 @@ class MockMemoryOperator {
}

void freeAll() {
std::lock_guard<std::mutex> l(mu_);
for (auto entry : allocations_) {
totalBytes_ -= entry.second;
std::unordered_map<void*, size_t> allocationsToFree;
{
std::lock_guard<std::mutex> l(mu_);
for (auto entry : allocations_) {
totalBytes_ -= entry.second;
}
VELOX_CHECK_EQ(totalBytes_, 0);
allocationsToFree.swap(allocations_);
}
VELOX_CHECK_EQ(totalBytes_, 0);
for (auto entry : allocations_) {
for (auto entry : allocationsToFree) {
pool_->free(entry.first, entry.second);
}
allocations_.clear();
}

void free() {
Allocation allocation;
std::lock_guard<std::mutex> l(mu_);
if (allocations_.empty()) {
return;
{
std::lock_guard<std::mutex> l(mu_);
if (allocations_.empty()) {
return;
}
allocation.buffer = allocations_.begin()->first;
allocation.size = allocations_.begin()->second;
totalBytes_ -= allocation.size;
allocations_.erase(allocations_.begin());
}
allocation.buffer = allocations_.begin()->first;
allocation.size = allocations_.begin()->second;
totalBytes_ -= allocation.size;
allocations_.erase(allocations_.begin());
pool_->free(allocation.buffer, allocation.size);
}

Expand All @@ -312,35 +317,39 @@ class MockMemoryOperator {
VELOX_CHECK_GT(targetBytes, 0);
uint64_t bytesReclaimed{0};
std::vector<Allocation> allocationsToFree;
std::lock_guard<std::mutex> l(mu_);
VELOX_CHECK_NOT_NULL(pool_);
VELOX_CHECK_EQ(pool->name(), pool_->name());
auto allocIt = allocations_.begin();
while (allocIt != allocations_.end() &&
((targetBytes != 0) && (bytesReclaimed < targetBytes))) {
allocationsToFree.push_back({allocIt->first, allocIt->second});
bytesReclaimed += allocIt->second;
allocIt = allocations_.erase(allocIt);
}
totalBytes_ -= bytesReclaimed;
const auto oldReservedBytes = pool_->reservedBytes();
{
std::lock_guard<std::mutex> l(mu_);
VELOX_CHECK_NOT_NULL(pool_);
VELOX_CHECK_EQ(pool->name(), pool_->name());
auto allocIt = allocations_.begin();
while (allocIt != allocations_.end() &&
((targetBytes != 0) && (bytesReclaimed < targetBytes))) {
allocationsToFree.push_back({allocIt->first, allocIt->second});
bytesReclaimed += allocIt->second;
allocIt = allocations_.erase(allocIt);
}
totalBytes_ -= bytesReclaimed;
}
for (const auto& allocation : allocationsToFree) {
pool_->free(allocation.buffer, allocation.size);
}
const auto newReservedBytes = pool_->reservedBytes();
VELOX_CHECK_GE(oldReservedBytes, newReservedBytes);
return newReservedBytes - oldReservedBytes;
return bytesReclaimed;
}

void abort(MemoryPool* pool) {
std::lock_guard<std::mutex> l(mu_);
VELOX_CHECK_NOT_NULL(pool_);
VELOX_CHECK_EQ(pool->name(), pool_->name());
for (const auto& allocation : allocations_) {
totalBytes_ -= allocation.second;
pool_->free(allocation.first, allocation.second);
std::unordered_map<void*, size_t> allocationsToFree;
{
std::lock_guard<std::mutex> l(mu_);
VELOX_CHECK_NOT_NULL(pool_);
VELOX_CHECK_EQ(pool->name(), pool_->name());
for (const auto& allocation : allocations_) {
totalBytes_ -= allocation.second;
}
allocationsToFree.swap(allocations_);
}
for (auto entry : allocationsToFree) {
pool_->free(entry.first, entry.second);
}
allocations_.clear();
}

void setPool(MemoryPool* pool) {
Expand Down

0 comments on commit a6bb4d8

Please sign in to comment.