diff --git a/docs/source/user_guide/read.rst b/docs/source/user_guide/read.rst index a94776353..61286f045 100644 --- a/docs/source/user_guide/read.rst +++ b/docs/source/user_guide/read.rst @@ -41,6 +41,57 @@ cross-language task scheduling and interaction (e.g., Java and C++), substantial across the two language ecosystems. +Planning with a custom memory pool +---------------------------------- + +Use ``ScanContextBuilder::WithMemoryPool(pool)`` with ordinary +``TableScan::Create`` and ``CreatePlan`` to supply an application-owned pool. +Planning continues to use the normal metadata decoders, filters and caches. +A pool that refuses an allocation can throw ``std::bad_alloc``. The caller must +handle both an exception from the planning call and a failed ``Result``: adapters +such as the Arrow memory pool can translate allocation refusal into an error. +A failed call must not be treated as an empty or partial successful plan. + +The budget covers only requests routed through the supplied pool. STL containers, +metadata parser storage, shared services and other allocations may bypass it. +It is not a limit on total planning memory, data-reader memory or process RSS. +Unaccounted allocations are not guaranteed to be small merely because the pool +is bounded. Applications requiring a stronger bound need a separate policy and +validation of the uncovered allocations. + +The application chooses the accounting policy in its pool. A **live-byte** limit +charges outstanding requested bytes and returns credit on ``Free``. A growing +``Realloc`` must reserve the increase before requesting storage; on failure it +must preserve the original allocation and its charge. A **cumulative** limit +instead charges allocation work without refunding frees. These are different +contracts: freeing temporary buffers permits another allocation under a live-byte +limit, but not necessarily under a cumulative limit. Neither policy accounts for +allocations that do not reach that pool. ``MemoryPool::CurrentUsage`` describes +live usage; the scan API does not install either budget policy for the caller. + +Keep a custom pool alive until all objects using its allocations have been +released, including retained plans, separately retained splits, cache entries +and executor task closures. A plan or split does not automatically retain a +shared owner of every custom pool. It can outlive its scan when the caller keeps +the pool alive. Use a thread-safe pool, or synchronize all its operations including +cleanup, when planning tasks or consumers can run concurrently. A scan itself +must not be called concurrently. + +Cache entries remain charged to the pool that allocated them until they are +released. A shared cache can retain allocations from multiple pools; keep each +such pool alive until all entries using its allocations have been released. + +When ordinary manifest planning fails, it waits for already accepted metadata +tasks before unwinding the task submission scope. For this guarantee, if +``Executor::Add`` throws, it must not have accepted or run the submitted task. +The default executor satisfies this condition. + +Waiting for a future to become ready does not guarantee that its task closure +has been destroyed. Keep the pool alive until task closures and their retained +results have been released, including when a shared or custom executor outlives +the scan. +Waiting for tasks does not cancel native I/O. + Schema Evolution ----------------------- Scope and Compatibility diff --git a/include/paimon/scan_context.h b/include/paimon/scan_context.h index 6fbb006f0..efee52743 100644 --- a/include/paimon/scan_context.h +++ b/include/paimon/scan_context.h @@ -225,6 +225,9 @@ class PAIMON_EXPORT ScanContextBuilder { /// @param memory_pool The memory pool to use. /// @return Reference to this builder for method chaining. /// @note if not set, memory_pool is default pool + /// A custom pool controls only allocations routed through it, not all planning memory. + /// Keep it alive while any plan, split, cache entry, or executor task can retain its + /// allocations; destroying the scan alone does not release all such objects. ScanContextBuilder& WithMemoryPool(const std::shared_ptr& memory_pool); /// Set custom executor for task execution. /// @param executor The executor to use. diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index ec0c6ea8c..7726b7bad 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -1040,6 +1040,7 @@ if(PAIMON_BUILD_TESTS) core/table/source/snapshot/snapshot_reader_test.cpp core/table/source/startup_mode_test.cpp core/table/source/table_scan_test.cpp + core/table/source/planning_memory_pool_test.cpp core/table/system/system_table_test.cpp core/tag/tag_test.cpp core/utils/blob_view_lookup_test.cpp diff --git a/src/paimon/core/operation/file_store_scan.cpp b/src/paimon/core/operation/file_store_scan.cpp index ca2de308a..b6fbd1ced 100644 --- a/src/paimon/core/operation/file_store_scan.cpp +++ b/src/paimon/core/operation/file_store_scan.cpp @@ -37,6 +37,7 @@ #include "paimon/common/predicate/literal_converter.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/field_type_utils.h" +#include "paimon/common/utils/scope_guard.h" #include "paimon/core/io/data_file_meta.h" #include "paimon/core/io/file_index_evaluator.h" #include "paimon/core/io/file_index_options.h" @@ -323,6 +324,11 @@ Status FileStoreScan::ReadManifestsWithSnapshot(const Snapshot& snapshot, return metas; }; std::vector>>> futures; + ScopeGuard drain([&futures] { + for (auto& future : futures) { + if (future.valid()) future.wait(); + } + }); futures.reserve(2); futures.push_back( Via(executor_.get(), [read_list]() { return read_list(/*base=*/true); })); @@ -352,6 +358,12 @@ Status FileStoreScan::ReadFileEntries(const std::vector& manif bool apply_scan_filter, std::vector* manifest_entries) const { std::vector>>> futures; + ScopeGuard drain([&futures] { + for (auto& future : futures) { + if (future.valid()) future.wait(); + } + }); + futures.reserve(manifest_metas.size()); for (const auto& meta : manifest_metas) { auto read_meta_task = [this, meta, apply_scan_filter]() -> Result> { @@ -511,6 +523,11 @@ Status FileStoreScan::ReadAndMergeBucketFileEntries( const bool inferred_bucket = !bucket_filter_ && bucket_selector_ != nullptr; if (core_options_.ScanManifestEntryLazyDecodeEnabled()) { std::vector>>> futures; + ScopeGuard drain([&futures] { + for (auto& future : futures) { + if (future.valid()) future.wait(); + } + }); futures.reserve(manifest_metas.size()); for (const auto& meta : manifest_metas) { auto read_meta_task = [this, meta, bucket, diff --git a/src/paimon/core/table/source/planning_memory_pool_test.cpp b/src/paimon/core/table/source/planning_memory_pool_test.cpp new file mode 100644 index 000000000..c8b2fed50 --- /dev/null +++ b/src/paimon/core/table/source/planning_memory_pool_test.cpp @@ -0,0 +1,402 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/memory_pool.h" +#include "gtest/gtest.h" +#include "paimon/common/io/cache/lru_cache.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/executor.h" +#include "paimon/fs/local/local_file_system.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/scan_context.h" +#include "paimon/table/source/table_scan.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +namespace { +class InlinePlanningExecutor final : public Executor { + public: + void Add(std::function task) override { + task(); + } + void ShutdownNow() override {} + uint32_t GetThreadNum() const override { + return 1; + } +}; + +class PlanningFaultPool final : public MemoryPool { + public: + void* Malloc(uint64_t size, uint64_t alignment) override { + std::lock_guard lock(mutex_); + Check(); + ++malloc_calls; + return upstream_->Malloc(size, alignment); + } + void* Realloc(void* pointer, size_t old_size, size_t new_size, uint64_t alignment) override { + std::lock_guard lock(mutex_); + Check(); + ++realloc_calls; + return upstream_->Realloc(pointer, old_size, new_size, alignment); + } + void Free(void* pointer, uint64_t size) override { + std::lock_guard lock(mutex_); + upstream_->Free(pointer, size); + } + void Free(void* pointer, uint64_t size, uint64_t alignment) override { + std::lock_guard lock(mutex_); + upstream_->Free(pointer, size, alignment); + } + uint64_t CurrentUsage() const override { + return upstream_->CurrentUsage(); + } + uint64_t MaxMemoryUsage() const override { + return upstream_->MaxMemoryUsage(); + } + std::atomic calls{0}; + std::atomic fail_at{std::numeric_limits::max()}; + std::atomic malloc_calls{0}; + std::atomic realloc_calls{0}; + std::atomic rejected{false}; + + private: + void Check() { + if (calls.fetch_add(1) == fail_at.load()) { + rejected = true; + throw std::bad_alloc(); + } + } + std::unique_ptr upstream_ = GetMemoryPool(); + std::mutex mutex_; +}; + +std::string PlanningFixture() { + return GetDataDir() + "/parquet/append_types_compatibility.db/rust_time_types"; +} + +struct PlanningInput { + std::string path = PlanningFixture(); + std::string snapshot = "1"; + std::shared_ptr cache; + std::optional bucket; + std::shared_ptr file_system; +}; + +Result> OrdinaryScan( + const std::shared_ptr& pool, + const std::shared_ptr& executor = std::make_shared(), + const PlanningInput& input = {}) { + std::ifstream schema_file(input.path + "/schema/schema-0"); + std::string schema{std::istreambuf_iterator(schema_file), + std::istreambuf_iterator()}; + ScanContextBuilder builder(input.path); + builder.SetTableSchema(schema).WithMemoryPool(pool).WithExecutor(executor).AddOption( + "scan.snapshot-id", input.snapshot); + if (input.cache) { + builder.WithCache(input.cache) + .AddOption("scan.manifest-entry-cache.max-snapshots", "8") + .AddOption("scan.manifest-entry.lazy-decode.enabled", "true"); + } + if (input.bucket) builder.SetBucketFilter(*input.bucket); + if (input.file_system) builder.WithFileSystem(input.file_system); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr context, builder.Finish()); + return TableScan::Create(std::move(context)); +} + +class RejectingPlanningExecutor final : public Executor { + public: + explicit RejectingPlanningExecutor(uint32_t rejected_submission = 2, + uint32_t blocked_submission = 0) + : rejected_submission_(rejected_submission), + blocked_submission_(blocked_submission ? blocked_submission : rejected_submission - 1) {} + ~RejectingPlanningExecutor() override { + Release(false); + worker_.reset(); + } + void Add(std::function task) override { + uint32_t submission = ++submissions_; + if (submission == rejected_submission_) { + rejected_ = true; + throw std::bad_alloc(); + } + if (submission != blocked_submission_) { + worker_->Add([this, submission, task = std::move(task)] { + task(); + if (submission == 1) { + std::lock_guard lock(mutex_); + first_completed_ = true; + changed_.notify_all(); + } + }); + return; + } + worker_->Add([this, task = std::move(task)] { + bool execute; + { + std::unique_lock lock(mutex_); + blocked_ = true; + changed_.notify_all(); + changed_.wait(lock, [this] { return released_; }); + execute = execute_; + } + if (execute) task(); + }); + } + void ShutdownNow() override { + worker_->ShutdownNow(); + } + uint32_t GetThreadNum() const override { + return 2; + } + bool WaitUntilBlocked() { + std::unique_lock lock(mutex_); + return changed_.wait_for(lock, std::chrono::seconds(5), [this] { return blocked_; }); + } + void Release(bool execute) { + std::lock_guard lock(mutex_); + execute_ = execute; + released_ = true; + changed_.notify_all(); + } + bool Rejected() const { + return rejected_; + } + bool WaitUntilFirstCompleted() { + std::unique_lock lock(mutex_); + return changed_.wait_for(lock, std::chrono::seconds(5), + [this] { return first_completed_; }); + } + + private: + std::unique_ptr worker_ = CreateDefaultExecutor(2).value(); + std::atomic submissions_{0}; + uint32_t rejected_submission_; + uint32_t blocked_submission_; + std::atomic rejected_{false}; + std::mutex mutex_; + std::condition_variable changed_; + bool blocked_ = false; + bool released_ = false; + bool execute_ = false; + bool first_completed_ = false; +}; + +class PoolRefusingManifestFileSystem final : public LocalFileSystem { + public: + explicit PoolRefusingManifestFileSystem(std::shared_ptr pool) + : pool_(std::move(pool)) {} + using LocalFileSystem::Open; + Result> Open(const std::string& path) const override { + if (path.find("/manifest/") != std::string::npos && !refused_.exchange(true)) { + pool_->fail_at = pool_->calls.load(); + // A supplied file system may allocate through the caller's pool too. + // Force refusal inside the real manifest worker, before opening a file. + void* allocation = pool_->Malloc(1, 64); + pool_->Free(allocation, 1); + } + return LocalFileSystem::Open(path); + } + + private: + std::shared_ptr pool_; + mutable std::atomic refused_{false}; +}; +} // namespace + +TEST(OrdinaryPlanningPoolTest, RejectEachPlanningAllocationWithoutLeaking) { + auto baseline_pool = std::make_shared(); + auto scan_result = OrdinaryScan(baseline_pool); + ASSERT_TRUE(scan_result.ok()) << scan_result.status().ToString(); + auto baseline_scan = std::move(scan_result).value(); + RecordProperty("factory_pool_calls", baseline_pool->calls.load()); + baseline_pool->calls = 0; + baseline_pool->malloc_calls = 0; + baseline_pool->realloc_calls = 0; + auto baseline_result = baseline_scan->CreatePlan(); + ASSERT_TRUE(baseline_result.ok()) << baseline_result.status().ToString(); + auto baseline = std::move(baseline_result).value(); + ASSERT_FALSE(baseline->Splits().empty()); + const uint64_t allocation_count = baseline_pool->calls; + ASSERT_GT(allocation_count, 0); + RecordProperty("planning_malloc_calls", baseline_pool->malloc_calls.load()); + RecordProperty("planning_realloc_calls", baseline_pool->realloc_calls.load()); + baseline.reset(); + baseline_scan.reset(); + ASSERT_EQ(baseline_pool->CurrentUsage(), 0); + + for (uint64_t ordinal = 0; ordinal < allocation_count; ++ordinal) { + SCOPED_TRACE(ordinal); + auto pool = std::make_shared(); + auto result = OrdinaryScan(pool); + ASSERT_TRUE(result.ok()) << result.status().ToString(); + auto scan = std::move(result).value(); + pool->calls = 0; + pool->fail_at = ordinal; + bool failed = false; + try { + auto plan = scan->CreatePlan(); + failed = !plan.ok(); + } catch (const std::bad_alloc&) { + failed = true; + } + ASSERT_TRUE(pool->rejected); + ASSERT_TRUE(failed) << "allocation refusal returned a successful plan"; + scan.reset(); + EXPECT_EQ(pool->CurrentUsage(), 0) << "pool storage survived failed planning"; + pool->fail_at = std::numeric_limits::max(); + { + auto recovery_result = OrdinaryScan(pool); + ASSERT_TRUE(recovery_result.ok()) << recovery_result.status().ToString(); + auto recovery = std::move(recovery_result).value(); + auto plan = recovery->CreatePlan(); + ASSERT_TRUE(plan.ok()) << plan.status().ToString(); + EXPECT_FALSE(plan.value()->Splits().empty()); + } + EXPECT_EQ(pool->CurrentUsage(), 0); + } +} + +TEST(OrdinaryPlanningPoolTest, RetainedSplitRemainsUsableWithCallerOwnedPool) { + auto pool = std::make_shared(); + auto scan_result = OrdinaryScan(pool); + ASSERT_TRUE(scan_result.ok()) << scan_result.status().ToString(); + auto scan = std::move(scan_result).value(); + auto plan_result = scan->CreatePlan(); + ASSERT_TRUE(plan_result.ok()) << plan_result.status().ToString(); + auto plan = std::move(plan_result).value(); + ASSERT_FALSE(plan->Splits().empty()); + auto split = plan->Splits().front(); + auto original = Split::Serialize(split, pool); + ASSERT_TRUE(original.ok()) << original.status().ToString(); + scan.reset(); + EXPECT_EQ(plan->SnapshotId(), 1); + plan.reset(); + auto retained = Split::Serialize(split, pool); + ASSERT_TRUE(retained.ok()) << retained.status().ToString(); + EXPECT_EQ(original.value(), retained.value()); + split.reset(); + EXPECT_EQ(pool->CurrentUsage(), 0); +} + +TEST(OrdinaryPlanningPoolTest, RefusedArrowReallocationPreservesExistingBuffer) { + auto pool = std::make_shared(); + auto arrow_pool = GetArrowPool(pool); + uint8_t* data = nullptr; + ASSERT_TRUE(arrow_pool->Allocate(16, 64, &data).ok()); + ASSERT_NE(data, nullptr); + data[0] = 42; + uint8_t* original = data; + pool->fail_at = pool->calls.load(); + EXPECT_TRUE(arrow_pool->Reallocate(16, 32, 64, &data).IsOutOfMemory()); + EXPECT_TRUE(pool->rejected); + EXPECT_EQ(data, original); + EXPECT_EQ(data[0], 42); + EXPECT_EQ(pool->CurrentUsage(), 16); + arrow_pool->Free(data, 16, 64); + EXPECT_EQ(pool->CurrentUsage(), 0); +} + +TEST(OrdinaryPlanningPoolTest, SubmissionFailureDrainsAcceptedManifestTask) { + for (uint32_t phase = 0; phase < 3; ++phase) { + SCOPED_TRACE(phase); + PlanningInput input; + uint32_t rejected_submission = 2; + if (phase > 0) { + input.path = GetDataDir() + "/orc/append_09.db/append_09"; + input.snapshot = "5"; + rejected_submission = 4; + } + if (phase == 2) { + input.cache = std::make_shared(16 * 1024 * 1024); + input.bucket = 0; + } + auto pool = std::make_shared(); + auto executor = std::make_shared(rejected_submission); + auto result = OrdinaryScan(pool, executor, input); + ASSERT_TRUE(result.ok()) << result.status().ToString(); + auto scan = std::move(result).value(); + auto planning = std::async(std::launch::async, [&scan] { + try { + auto plan = scan->CreatePlan(); + return !plan.ok(); + } catch (const std::bad_alloc&) { + return true; + } + }); + bool blocked = executor->WaitUntilBlocked(); + bool returned_early = + planning.wait_for(std::chrono::milliseconds(200)) == std::future_status::ready; + // On the faulty path the task borrows an already unwound snapshot. Drop it + // without dereferencing that snapshot so the regression reports the ordering. + executor->Release(!returned_early); + EXPECT_TRUE(blocked); + EXPECT_TRUE(executor->Rejected()); + EXPECT_FALSE(returned_early) + << "CreatePlan returned with accepted manifest work still pending"; + EXPECT_TRUE(planning.get()); + scan.reset(); + executor.reset(); + input.cache.reset(); + EXPECT_EQ(pool->CurrentUsage(), 0); + } +} + +TEST(OrdinaryPlanningPoolTest, WorkerPoolExceptionDrainsAcceptedManifestTask) { + auto pool = std::make_shared(); + auto executor = std::make_shared(0, 2); + PlanningInput input; + input.file_system = std::make_shared(pool); + auto result = OrdinaryScan(pool, executor, input); + ASSERT_TRUE(result.ok()) << result.status().ToString(); + auto scan = std::move(result).value(); + auto planning = std::async(std::launch::async, [&scan] { + try { + (void)scan->CreatePlan(); + return false; + } catch (const std::bad_alloc&) { + return true; + } + }); + bool blocked = executor->WaitUntilBlocked(); + bool first_completed = executor->WaitUntilFirstCompleted(); + bool returned_early = + planning.wait_for(std::chrono::milliseconds(200)) == std::future_status::ready; + executor->Release(!returned_early); + EXPECT_TRUE(blocked); + EXPECT_TRUE(first_completed); + EXPECT_TRUE(pool->rejected); + EXPECT_FALSE(returned_early) << "worker exception unwound before the other manifest task"; + EXPECT_TRUE(planning.get()); + scan.reset(); + executor.reset(); + input.file_system.reset(); + EXPECT_EQ(pool->CurrentUsage(), 0); +} +} // namespace paimon::test