diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp index f5bc90391369d..ea40e5cb3bb4b 100644 --- a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp +++ b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp @@ -131,8 +131,15 @@ std::shared_ptr QueryContextManager::findOrCreateQueryCtx( } velox::core::QueryConfig queryConfig{std::move(configStrings)}; + // NOTE: the monotonically increasing 'poolId' is appended to 'queryId' to + // ensure that the name of root memory pool instance is always unique. In some + // edge case, we found some background activities such as the long-running + // memory arbitration process will still hold the query root memory pool even + // though the query ctx has been evicted out of the cache. The query ctx cache + // is still indexed by the query id. + static std::atomic_uint64_t poolId{0}; auto pool = memory::defaultMemoryManager().addRootPool( - queryId, + fmt::format("{}_{}", queryId, poolId++), queryConfig.queryMaxMemoryPerNode() != 0 ? queryConfig.queryMaxMemoryPerNode() : SystemConfig::instance()->queryMaxMemoryPerNode(), @@ -163,4 +170,13 @@ void QueryContextManager::visitAllContexts( } } +void QueryContextManager::testingClearCache() { + queryContextCache_.wlock()->testingClear(); +} + +void QueryContextCache::testingClear() { + queryCtxs_.clear(); + queryIds_.clear(); +} + } // namespace facebook::presto diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.h b/presto-native-execution/presto_cpp/main/QueryContextManager.h index a50dc37841a7d..12ac3fb29e30d 100644 --- a/presto-native-execution/presto_cpp/main/QueryContextManager.h +++ b/presto-native-execution/presto_cpp/main/QueryContextManager.h @@ -87,6 +87,8 @@ class QueryContextCache { return queryCtxs_; } + void testingClear(); + private: size_t capacity_; @@ -106,11 +108,14 @@ class QueryContextManager { const protocol::TaskId& taskId, const protocol::SessionRepresentation& session); - // Calls the given functor for every present query context. + /// Calls the given functor for every present query context. void visitAllContexts(std::function visitor) const; + /// Test method to clear the query context cache. + void testingClearCache(); + private: std::shared_ptr findOrCreateQueryCtx( const protocol::TaskId& taskId, diff --git a/presto-native-execution/presto_cpp/main/tests/QueryContextManagerTest.cpp b/presto-native-execution/presto_cpp/main/tests/QueryContextManagerTest.cpp index 5245f4649b276..9a1475358bc23 100644 --- a/presto-native-execution/presto_cpp/main/tests/QueryContextManagerTest.cpp +++ b/presto-native-execution/presto_cpp/main/tests/QueryContextManagerTest.cpp @@ -12,6 +12,7 @@ * limitations under the License. */ #include "presto_cpp/main/QueryContextManager.h" +#include #include #include "presto_cpp/main/TaskManager.h" @@ -68,4 +69,50 @@ TEST_F(QueryContextManagerTest, defaultSessionProperties) { EXPECT_EQ(queryCtx->queryConfig().spillWriteBufferSize(), 1L << 20); } +TEST_F(QueryContextManagerTest, duplicateQueryRootPoolName) { + const protocol::TaskId fakeTaskId = "scan.0.0.1.0"; + const protocol::SessionRepresentation fakeSession{.systemProperties = {}}; + auto* queryCtxManager = taskManager_->getQueryContextManager(); + struct { + bool hasPendingReference; + bool clearCache; + bool expectedNewPoolName; + + std::string debugString() const { + return fmt::format( + "hasPendingReference: {}, clearCache: {}, expectedNewPoolName: {}", + hasPendingReference, + clearCache, + expectedNewPoolName); + } + } testSettings[] = { + {true, true, true}, + {true, false, false}, + {false, true, true}, + {false, false, true}}; + for (const auto& testData : testSettings) { + SCOPED_TRACE(testData.debugString()); + queryCtxManager->testingClearCache(); + + auto queryCtx = + queryCtxManager->findOrCreateQueryCtx(fakeTaskId, fakeSession); + const auto poolName = queryCtx->pool()->name(); + ASSERT_THAT(poolName, testing::HasSubstr("scan_")); + if (!testData.hasPendingReference) { + queryCtx.reset(); + } + if (testData.clearCache) { + queryCtxManager->testingClearCache(); + } + auto newQueryCtx = + queryCtxManager->findOrCreateQueryCtx(fakeTaskId, fakeSession); + const auto newPoolName = newQueryCtx->pool()->name(); + ASSERT_THAT(newPoolName, testing::HasSubstr("scan_")); + if (testData.expectedNewPoolName) { + ASSERT_NE(poolName, newPoolName); + } else { + ASSERT_EQ(poolName, newPoolName); + } + } +} } // namespace facebook::presto