From f875b0dedd5809e2567c517ecf513f1a560318bd Mon Sep 17 00:00:00 2001 From: Shaojie Li Date: Sun, 7 Jun 2026 21:04:19 +0000 Subject: [PATCH] fix: Restore thread-local stat writer in MockSharedArbitratorTest Three tests (globalArbitrationWithoutSpill, globalArbitrationSmallParticipantLargeGrow, and a DEBUG_ONLY test) set the thread-local runtime stat writer to a stack-owned TestRuntimeStatWriter on the main thread via a bare setThreadLocalRunTimeStatWriter() call, but never restore it. localRuntimeStatWriter is a thread_local raw pointer, so once the TestRuntimeStatWriter is destroyed at test exit the pointer dangles. The next test that runs on the main thread and reaches addThreadLocalRuntimeStat (e.g. any global arbitration that calls finishArbitration during growCapacity) dereferences the freed writer and segfaults. The tests pass in isolation but crash when run after one of these, depending on execution order. Use the existing RAII RuntimeStatWriterScopeGuard, which restores the previous writer on destruction, so the thread-local pointer is never left dangling. --- velox/common/memory/tests/MockSharedArbitratorTest.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp index 830aa1e699b..d539bba0441 100644 --- a/velox/common/memory/tests/MockSharedArbitratorTest.cpp +++ b/velox/common/memory/tests/MockSharedArbitratorTest.cpp @@ -1977,7 +1977,7 @@ DEBUG_ONLY_TEST_F( std::unordered_map runtimeStats; auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + RuntimeStatWriterScopeGuard statsWriterGuard(statsWriter.get()); localArbitrationOp->allocate(memoryPoolReservedCapacity); // Inject some delay for global arbitration. @@ -2111,7 +2111,7 @@ TEST_F(MockSharedArbitrationTest, globalArbitrationWithoutSpill) { std::unordered_map runtimeStats; auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + RuntimeStatWriterScopeGuard statsWriterGuard(statsWriter.get()); triggerOp->allocate(memoryCapacity / 2); ASSERT_EQ( @@ -2166,7 +2166,7 @@ TEST_F(MockSharedArbitrationTest, globalArbitrationSmallParticipantLargeGrow) { std::unordered_map runtimeStats; auto statsWriter = std::make_unique(runtimeStats); - setThreadLocalRunTimeStatWriter(statsWriter.get()); + RuntimeStatWriterScopeGuard statsWriterGuard(statsWriter.get()); // task0 has 256MB + 256MB (attempt) = 512MB in top abort capacity limit // bucket, which shall be evaluated first, and hence killed by global