Copilot commented on code in PR #13227:
URL: https://github.com/apache/gluten/pull/13227#discussion_r4229369391
##########
cpp/velox/utils/VeloxWholeStageDumper.cc:
##########
@@ -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());
Review Comment:
`getLeafMemoryPool()` appears to return a smart pointer in this PR (the new
test calls `.get()` on its result), but here it’s passed directly into
`ParquetStreamReaderIterator`. If `ParquetStreamReaderIterator` expects a
`facebook::velox::memory::MemoryPool*` (as suggested by the previous
`addLeafChild(...)` usage), this will be a type mismatch or (worse) lead to
unintended overload resolution. Prefer passing the raw pointer
(`leafPool.get()`) and, if lifetime extension is required, ensure the iterator
(or dumper) retains a shared ownership reference to the leaf pool (e.g., store
the shared_ptr) rather than only a raw pointer.
##########
cpp/velox/tests/VeloxWholeStageDumperTest.cc:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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();
Review Comment:
`memory::MemoryManager::testingSetInstance(...)` mutates global process
state; the suite doesn’t restore/reset it in `TearDownTestSuite()`. This can
make other tests order-dependent or flaky if they rely on the default
`MemoryManager` instance/configuration. Consider resetting/restoring the
previous instance in `TearDownTestSuite()` (using the project’s supported reset
API or by saving/restoring the prior singleton) so this test suite is hermetic.
##########
cpp/velox/tests/VeloxWholeStageDumperTest.cc:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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());
+
+ auto reader = dumper.dumpInputIterator(0, input);
+ ASSERT_NE(reader, nullptr);
+ ASSERT_TRUE(std::filesystem::exists(saveDir_ / "data_1_2_4_0.parquet"));
+
+ const auto leafPool = vmm_->getLeafMemoryPool();
+ const auto baselineBytes = leafPool->usedBytes();
+ 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);
+ const auto bytesHeldByBatches = leafPool->usedBytes();
+ ASSERT_GT(bytesHeldByBatches, baselineBytes);
+
+ // Destroying the reader must not destroy the pool backing the batches, nor
free the memory they hold.
+ reader.reset();
+ ASSERT_GE(leafPool->usedBytes(), bytesHeldByBatches);
+
+ // 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);
Review Comment:
The final `ASSERT_EQ(leafPool->usedBytes(), baselineBytes)` can be brittle
if non-batch allocations occur during parquet read/write (e.g., internal
buffers, lazy init, metadata) that don’t fully return to the exact baseline,
leading to flakes unrelated to the UAF regression. To keep the test focused on
the lifetime guarantee, consider relaxing the assertion (e.g., assert it drops
back near/below a threshold) while still validating that clearing `batches`
releases the bulk of memory and (critically) does not crash after
`reader.reset()`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
