Copilot commented on code in PR #13227: URL: https://github.com/apache/gluten/pull/13227#discussion_r4222492731
########## cpp/velox/tests/VeloxWholeStageDumperTest.cc: ########## @@ -0,0 +1,124 @@ +/* + * 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 SetUpTestCase() { + VeloxBackend::create(AllocationListener::noop(), {}); + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + } + Review Comment: `memory::MemoryManager::testingSetInstance(...)` (and potentially `VeloxBackend::create(...)`) modifies global process state, but the test doesn't restore it. This can introduce order-dependence and flakiness when the full test suite runs. Consider adding `TearDownTestCase()` to reset/restore the global MemoryManager (and backend instance if the API supports it), or use an existing RAII/guard utility used by other tests in this repo that temporarily swaps the singleton. ########## cpp/velox/utils/VeloxWholeStageDumper.h: ########## @@ -45,7 +45,9 @@ class VeloxWholeStageDumper final : public WholeStageDumper { std::string saveDir_; int64_t batchSize_; - facebook::velox::memory::MemoryPool* pool_; + // Owned by the Runtime's Java-side counterpart and released after the Runtime, so it outlives the dumper and Review Comment: The comment states the memory manager is owned by the Runtime's Java-side counterpart, but the updated call site passes `VeloxRuntime::memoryManager()` (native). If the ownership is actually native (and merely *coordinated* by Java/JNI), it would be clearer/less misleading to rephrase this to reflect the actual ownership/lifetime contract used here (e.g., owned by the Runtime/JNI wrapper and guaranteed to outlive dump iterators/batches). ########## cpp/velox/tests/VeloxWholeStageDumperTest.cc: ########## @@ -0,0 +1,124 @@ +/* + * 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 SetUpTestCase() { + VeloxBackend::create(AllocationListener::noop(), {}); + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + } + + 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(); + 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(), 0); + + // Destroying the reader must not destroy the pool backing the batches. + reader.reset(); + ASSERT_GT(leafPool->usedBytes(), 0); + + // Releasing the batches after the reader is gone must be safe and must return all memory to the pool. + batches.clear(); + ASSERT_EQ(leafPool->usedBytes(), 0); Review Comment: The regression intent (no crash / no UAF) is solid, but asserting exact `usedBytes()` values can be brittle because some Velox components may retain pool allocations (e.g., internal caches/metadata) beyond vector lifetimes. To make the test more robust while still validating the bugfix, consider focusing assertions on safety and relative changes (e.g., `usedBytes()` decreases after `batches.clear()`), or capture a baseline `usedBytes()` before reading and assert you return to that baseline after cleanup. -- 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]
