This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 933315982d [GLUTEN-13225][VL] Fix use-after-free of memory pool when
dumping inputs for micro benchmark (#13227)
933315982d is described below
commit 933315982dc3f883acf890c77414956795c45636
Author: Rong Ma <[email protected]>
AuthorDate: Fri Oct 9 17:22:58 2026 +0100
[GLUTEN-13225][VL] Fix use-after-free of memory pool when dumping inputs
for micro benchmark (#13227)
---
cpp/velox/compute/VeloxRuntime.cc | 5 +-
cpp/velox/tests/CMakeLists.txt | 2 +
cpp/velox/tests/VeloxWholeStageDumperTest.cc | 136 +++++++++++++++++++++++++++
cpp/velox/utils/VeloxWholeStageDumper.cc | 18 ++--
cpp/velox/utils/VeloxWholeStageDumper.h | 7 +-
5 files changed, 155 insertions(+), 13 deletions(-)
diff --git a/cpp/velox/compute/VeloxRuntime.cc
b/cpp/velox/compute/VeloxRuntime.cc
index 45649c05f8..5c4621a209 100644
--- a/cpp/velox/compute/VeloxRuntime.cc
+++ b/cpp/velox/compute/VeloxRuntime.cc
@@ -637,10 +637,7 @@ void VeloxRuntime::enableDumping() {
GLUTEN_CHECK(taskInfo.has_value(), "Task info is not set. Please set task
info before enabling dumping.");
dumper_ = std::make_shared<VeloxWholeStageDumper>(
- taskInfo.value(),
- saveDir.value(),
- veloxCfg_->get<int64_t>(kSparkBatchSize, 4096),
- memoryManager()->getAggregateMemoryPool().get());
+ taskInfo.value(), saveDir.value(),
veloxCfg_->get<int64_t>(kSparkBatchSize, 4096), memoryManager());
dumper_->dumpConf(getConfMap());
}
diff --git a/cpp/velox/tests/CMakeLists.txt b/cpp/velox/tests/CMakeLists.txt
index b3b368703a..6a10641fbf 100644
--- a/cpp/velox/tests/CMakeLists.txt
+++ b/cpp/velox/tests/CMakeLists.txt
@@ -142,6 +142,8 @@ add_velox_test(spark_functions_test SOURCES
SparkFunctionTest.cc
add_velox_test(runtime_test SOURCES RuntimeTest.cc)
add_velox_test(velox_memory_test SOURCES MemoryManagerTest.cc)
add_velox_test(buffer_outputstream_test SOURCES BufferOutputStreamTest.cc)
+add_velox_test(velox_whole_stage_dumper_test SOURCES
+ VeloxWholeStageDumperTest.cc)
if(ENABLE_S3)
add_velox_test(gluten_s3_file_system_test SOURCES GlutenS3FileSystemTest.cc)
endif()
diff --git a/cpp/velox/tests/VeloxWholeStageDumperTest.cc
b/cpp/velox/tests/VeloxWholeStageDumperTest.cc
new file mode 100644
index 0000000000..1ecb5f376a
--- /dev/null
+++ b/cpp/velox/tests/VeloxWholeStageDumperTest.cc
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "utils/VeloxWholeStageDumper.h"
+#include "compute/VeloxBackend.h"
+#include "memory/VeloxColumnarBatch.h"
+#include "memory/VeloxMemoryManager.h"
+#include "velox/vector/tests/utils/VectorTestBase.h"
+
+#include <gtest/gtest.h>
+
+#include <filesystem>
+
+using namespace facebook::velox;
+
+namespace gluten {
+namespace {
+
+// A ColumnarBatchIterator that serves a fixed list of row vectors.
+class RowVectorBatchIterator final : public ColumnarBatchIterator {
+ public:
+ explicit RowVectorBatchIterator(std::vector<RowVectorPtr> vectors) :
vectors_(std::move(vectors)) {}
+
+ std::shared_ptr<ColumnarBatch> next() override {
+ if (idx_ >= vectors_.size()) {
+ return nullptr;
+ }
+ return std::make_shared<VeloxColumnarBatch>(vectors_[idx_++]);
+ }
+
+ private:
+ std::vector<RowVectorPtr> vectors_;
+ size_t idx_{0};
+};
+
+} // namespace
+
+class VeloxWholeStageDumperTest : public ::testing::Test, public
test::VectorTestBase {
+ protected:
+ static void SetUpTestSuite() {
+ VeloxBackend::create(AllocationListener::noop(), {});
+
memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{});
+ }
+
+ static void TearDownTestSuite() {
+ VeloxBackend::get()->tearDown();
+ }
+
+ void SetUp() override {
+ vmm_ = std::make_unique<VeloxMemoryManager>(
+ kVeloxBackendKind, AllocationListener::noop(),
*VeloxBackend::get()->getBackendConf());
+ saveDir_ = std::filesystem::path(::testing::TempDir()) /
+ fmt::format("gluten-whole-stage-dumper-test-{}",
::testing::UnitTest::GetInstance()->random_seed());
+ std::filesystem::remove_all(saveDir_);
+ }
+
+ void TearDown() override {
+ vmm_.reset();
+ std::filesystem::remove_all(saveDir_);
+ }
+
+ std::unique_ptr<VeloxMemoryManager> vmm_;
+ std::filesystem::path saveDir_;
+};
+
+// Regression test: batches read back from the dumped file must stay valid
after the reader iterator is gone.
+// The reader iterator is dropped by the value stream as soon as the input is
exhausted, while vectors it produced
+// may still be referenced by downstream operators or by batches handed to
Java. If they were allocated from a pool
+// owned by the reader, freeing them later would hit a destroyed pool and
crash in AlignedBuffer::freeToPool.
+TEST_F(VeloxWholeStageDumperTest, dumpedBatchesOutliveReaderIterator) {
+ constexpr int32_t kNumBatches = 3;
+ constexpr vector_size_t kRowsPerBatch = 1000;
+
+ std::vector<RowVectorPtr> inputs;
+ for (int32_t i = 0; i < kNumBatches; ++i) {
+ inputs.push_back(makeRowVector(
+ {"c0", "c1", "c2"},
+ {makeFlatVector<int64_t>(kRowsPerBatch, [i](auto row) { return i *
kRowsPerBatch + row; }),
+ makeFlatVector<double>(kRowsPerBatch, [](auto row) { return row *
0.5; }),
+ makeFlatVector<StringView>(kRowsPerBatch, [](auto row) {
+ return row % 2 == 0 ? StringView("short") : StringView("a string
that is not inlined");
+ })}));
+ }
+ auto input = std::make_shared<RowVectorBatchIterator>(inputs);
+
+ const SparkTaskInfo taskInfo{/*stageId=*/1, /*partitionId=*/2, /*taskId=*/3,
/*vId=*/4};
+ VeloxWholeStageDumper dumper(taskInfo, saveDir_.string(), kRowsPerBatch,
vmm_.get());
+
+ // Capture the baseline before the reader exists: the reader allocates its
own state from the leaf pool as soon
+ // as it is created, and all of it must be gone by the end of the test.
+ const auto leafPool = vmm_->getLeafMemoryPool();
+ const auto baselineBytes = leafPool->usedBytes();
+
+ auto reader = dumper.dumpInputIterator(0, input);
+ ASSERT_NE(reader, nullptr);
+ ASSERT_TRUE(std::filesystem::exists(saveDir_ / "data_1_2_4_0.parquet"));
+
+ std::vector<std::shared_ptr<ColumnarBatch>> batches;
+ int64_t numRows = 0;
+ while (auto cb = reader->next()) {
+ numRows += cb->numRows();
+ // Vectors must be backed by the Runtime-lifetime leaf pool rather than a
pool owned by the reader.
+ ASSERT_EQ(VeloxColumnarBatch::from(leafPool.get(),
cb)->getRowVector()->pool(), leafPool.get());
+ batches.push_back(std::move(cb));
+ }
+ ASSERT_EQ(numRows, kNumBatches * kRowsPerBatch);
+ ASSERT_GT(leafPool->usedBytes(), baselineBytes);
+
+ // Destroying the reader frees the reader's own buffers, but must not
destroy the pool backing the batches,
+ // so the memory held by the batches is still accounted for.
+ reader.reset();
+ const auto bytesHeldByBatches = leafPool->usedBytes();
+ ASSERT_GT(bytesHeldByBatches, baselineBytes);
+
+ // Releasing the batches after the reader is gone must be safe,
+ // and must give the memory back to the pool, which is still alive.
+ batches.clear();
+ ASSERT_LT(leafPool->usedBytes(), bytesHeldByBatches);
+ ASSERT_EQ(leafPool->usedBytes(), baselineBytes);
+}
+
+} // namespace gluten
diff --git a/cpp/velox/utils/VeloxWholeStageDumper.cc
b/cpp/velox/utils/VeloxWholeStageDumper.cc
index d8a451dac7..3f0bb556f8 100644
--- a/cpp/velox/utils/VeloxWholeStageDumper.cc
+++ b/cpp/velox/utils/VeloxWholeStageDumper.cc
@@ -58,8 +58,8 @@ VeloxWholeStageDumper::VeloxWholeStageDumper(
const SparkTaskInfo& taskInfo,
const std::string& saveDir,
int64_t batchSize,
- facebook::velox::memory::MemoryPool* aggregatePool)
- : taskInfo_(taskInfo), saveDir_(saveDir), batchSize_(batchSize),
pool_(aggregatePool) {}
+ VeloxMemoryManager* memoryManager)
+ : taskInfo_(taskInfo), saveDir_(saveDir), batchSize_(batchSize),
memoryManager_(memoryManager) {}
void VeloxWholeStageDumper::dumpConf(const std::unordered_map<std::string,
std::string>& confMap) {
const auto& backendConfMap =
VeloxBackend::get()->getBackendConf()->rawConfigs();
@@ -113,18 +113,22 @@ std::shared_ptr<ColumnarBatchIterator>
VeloxWholeStageDumper::dumpInputIterator(
fmt::format("data_{}_{}_{}_{}.parquet", taskInfo_.stageId,
taskInfo_.partitionId, taskInfo_.vId, iteratorIndex);
const auto dumpPath = checkAndGetDumpPath(saveDir_, fileName);
- // Velox parquet writer requires aggregate memory pool.
+ // Velox parquet writer requires aggregate memory pool. The writer is
short-lived so a dedicated child pool is fine.
auto writer = std::make_shared<VeloxColumnarBatchWriter>(
- dumpPath, batchSize_,
pool_->addAggregateChild(fmt::format("dump_iterator.{}", iteratorIndex)));
+ dumpPath,
+ batchSize_,
+
memoryManager_->getAggregateMemoryPool()->addAggregateChild(fmt::format("dump_iterator.{}",
iteratorIndex)));
while (auto cb = inputIterator->next()) {
GLUTEN_THROW_NOT_OK(writer->write(cb));
}
GLUTEN_THROW_NOT_OK(writer->close());
- // Velox parquet reader requires leaf memory pool.
- return std::make_shared<ParquetStreamReaderIterator>(
- dumpPath, batchSize_,
pool_->addLeafChild(fmt::format("retrieve_iterator.{}", iteratorIndex)));
+ // Velox parquet reader requires leaf memory pool. Do not create a dedicated
child pool here: the reader
+ // iterator is destroyed as soon as the input stream is exhausted, while
vectors it produced may still be
+ // referenced by downstream operators or by batches handed to Java. Freeing
them against a destroyed pool
+ // crashes in AlignedBuffer::freeToPool, so use the Runtime's leaf pool
which outlives all of them.
+ return std::make_shared<ParquetStreamReaderIterator>(dumpPath, batchSize_,
memoryManager_->getLeafMemoryPool());
}
} // namespace gluten
diff --git a/cpp/velox/utils/VeloxWholeStageDumper.h
b/cpp/velox/utils/VeloxWholeStageDumper.h
index dc6c9da90a..9327f40474 100644
--- a/cpp/velox/utils/VeloxWholeStageDumper.h
+++ b/cpp/velox/utils/VeloxWholeStageDumper.h
@@ -18,6 +18,7 @@
#pragma once
#include "compute/VeloxRuntime.h"
+#include "memory/VeloxMemoryManager.h"
#include "utils/WholeStageDumper.h"
namespace gluten {
@@ -28,7 +29,7 @@ class VeloxWholeStageDumper final : public WholeStageDumper {
const SparkTaskInfo& taskInfo,
const std::string& saveDir,
int64_t batchSize,
- facebook::velox::memory::MemoryPool* aggregatePool);
+ VeloxMemoryManager* memoryManager);
void dumpConf(const std::unordered_map<std::string, std::string>& confMap)
override;
@@ -45,7 +46,9 @@ class VeloxWholeStageDumper final : public WholeStageDumper {
std::string saveDir_;
int64_t batchSize_;
- facebook::velox::memory::MemoryPool* pool_;
+ // Owned by the Runtime / JNI wrapper and released after the Runtime, so it
outlives the dumper and
+ // every iterator / batch created by the Runtime.
+ VeloxMemoryManager* memoryManager_;
};
} // namespace gluten
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
