diff --git a/docs/design/allocation_accounted_memory_admission_impl.md b/docs/design/allocation_accounted_memory_admission_impl.md new file mode 100644 index 0000000000000..315b968459b22 --- /dev/null +++ b/docs/design/allocation_accounted_memory_admission_impl.md @@ -0,0 +1,237 @@ +# Allocation-accounted HashBuild memory admission + +Issue: #26459 + +## Goal + +Within the controlled domain, HashBuild and its join consumers must admit the +bytes they physically retain, not a predicted multiple of logical payload size. +A query may spill or degrade an optional optimization when the retained +allocation cannot be admitted, but must not fail because an estimator guessed +a larger, non-existent allocation. + +The implementation has one production path: + +1. the MPool performs the physical allocation; +2. the allocation carries immutable account, owner, and site provenance; +3. the account charges the query generation and the CN aggregate controller; +4. the same allocation releases the charge when MPool frees it. + +There is no feature switch, activation gate, or parallel compatibility ledger. +Physical storage is charged only by its allocation. A shuffle HashBuild may +additionally reserve a conservative recovery floor before accepting spillable +retained state; that floor proves the retained state can be drained under a +shared budget. A capacity failure while growing the floor changes the decision +from retain to direct spill; lifecycle and invariant failures remain terminal. +The projection is never reapplied to an upstream-owned direct source as a +query-fatal estimate. Cardinality estimates may likewise select a hash-table +capacity without creating a separately releasable physical allocation charge. + +## Scope + +The account covers retained physical storage owned by the HashBuild execution +family: + +- hash table cells, descriptors, iterator keys, and selection lists; +- copied build batches and retained unique keys; +- JoinMap-owned batches and grouping metadata; +- HashJoin, LoopJoin, DedupJoin, and RightDedupJoin result/finalize state; +- join matched/capture state; +- Product result state; +- runtime-filter payloads until ownership transfer; +- spill encode/decode/scatter buffers and rebuilt retained state. + +HashBuild and join expression trees receive the same account before +construction. Their fixed/variable values, nested function results, selected +row buffers, list results, and other owned MPool vectors allocate directly +through it. Function-result wrappers retain the immutable selection even when +the current vector is transferred, so later reuse cannot fall back to an +unaccounted vector. Borrowed input and serialized-plan vectors keep their +source ownership. + +Library-internal Go heap objects remain outside the controlled domain. This +explicit boundary avoids pretending that regexp, JSON, JQ, and other libraries +with no allocator/free hooks are exactly charged. + +The account is therefore not advertised as total query RSS. Unobservable +library allocations must not be represented by an estimated charge inside +this exact, terminal-zero ledger. + +ProductL2 scratch and native CPU/GPU index storage are also outside this first +controlled domain. ProductL2 still consumes an accounted JoinMap: those source +allocations keep their original provenance through transfer and are released +by JoinMap `Free`. Its additional search index and scratch remain under the +existing implementation until their CPU and GPU allocators expose one common +physical capacity contract. Adding only the visible Go buffers would claim +false exactness and create a partial second path. + +## Ownership model + +Each physical allocation has one owner and one release path. Provenance is +attached before the first owned allocation and cannot change while storage is +live. Views borrow storage and do not create another charge. Copies allocate +new storage in their destination account. + +The main transfer boundaries are: + +| Storage | Initial owner | Transfer | Terminal release | +| --- | --- | --- | --- | +| retained build batch | HashmapBuilder | JoinMap | JoinMap `Free` | +| hash cells/descriptors | HashmapBuilder | JoinMap | JoinMap `Free` | +| grouping selections | HashmapBuilder | JoinMap | JoinMap `Free` | +| spill file + disk/FD tokens | HashBuild | SpillBuildPayload | SpillEngine/file close | +| matched bitmap | parallel worker | BitmapMailbox/merger | merger or mailbox drain | +| Product build batches | producer JoinMap | Product | Product reset/free | +| runtime-filter payload | HashBuild | message board | message destruction | + +Transfers are move-only. A successful send clears the sender's ownership; a +failed send leaves ownership with the sender. Cancellation seals mailboxes and +drains queued accounted objects before terminal validation. + +## Execution lifecycle + +Compile opens one allocation generation for each local statement attempt. +Every operator implementing `SetAllocationAccount` / `ClearAllocationAccount` +is an owner in that generation. + +The sequence is: + +1. collect owners from physical scope templates; +2. open the account with the live HashBuild capacity controller; +3. configure all owners atomically, rolling back in reverse order on failure; +4. attach parallel scan/load clones created during `runOnce` to the same + generation before worker `Prepare` runs; +5. execute and drain the message board; +6. clear owners in reverse order; +7. seal and finalize the account; +8. export exactly one terminal snapshot. + +For a coordinator-local attempt, that sequence owns one MessageBoard. Remote +PipelineMessage handlers are different: all fragments of the same statement on +one CN share the board, and accounted JoinMap/result storage may remain owned by +a producer account until a sibling consumes it. The coordinator therefore +computes the complete physical scope graph's RPC count per CN, creates a unique +physical execution ID, and carries both in ProcessInfo. Each remote handler +registers before decoding/execution, stages its account and MPool after its +operators quiesce, and detaches the shared board before Compile cleanup. A +quiesced handler returns without waiting for siblings and adds one counted +pending terminal-memory-domain signal. This avoids a B-to-C-to-B response +dependency cycle. +The final fragment drains the board, completes every staged account, samples +every staged MPool, and publishes the aggregate plus a completion marker exactly +once. The coordinator reduces pending/completed markers independent of response +order; unresolved counts preserve every suppressed fragment MPool domain and +make the attempt explicitly partial instead of silently omitting memory. + +If dispatch fails before every planned handler arrives, the first quiesced +fragment starts a bounded orphan-registration timer. Registration of the full +set cancels it. Expiry closes only that board generation; active registered +fragments retain their accounts until they quiesce, after which the final +registered fragment completes the group. If no response remains to carry a +completion marker, the already-returned pending marker preserves the missing +domain as an explicit partial result. A fragment failure aborts an incomplete +group without waiting for the timer and cancels active registered siblings. +This bounds the MessageCenter/group maps without sealing live +allocations or imposing an execution timeout on a fully registered query. +Prepared executions and retries use fresh execution IDs, so a stale group and +board cannot capture the next generation. + +ProcessInfo without both topology and execution-ID metadata remains decodable. +Remote plans outside the accounted owner domain keep their legacy lifecycle; +plans containing an allocation-account owner are rejected explicitly. The +server never guesses a group size and never silently falls back to unaccounted +HashBuild execution. In the opposite upgrade direction, query-candidate +discovery already excludes CNs whose CommitID differs from the coordinator +binary, so a new coordinator cannot dispatch this lifecycle to an old handler. + +Prepared statements and retries create a new generation. Reset frees all +generation-bound state; it does not carry an executor, bitmap, mailbox payload, +or allocation selection into the next attempt. Runtime parallel clones are +also registered as owners and are cleared before their reuse-pool release. + +A valid terminal snapshot requires zero live bytes and zero live allocation +metadata. A mismatch suspends new admission until late physical frees drain the +tombstone; it is never converted into retryable capacity pressure. + +## Capacity and pressure + +The controller enforces both the query-generation cap and the CN aggregate +cap. Admission uses checked arithmetic and charges the physical MPool capacity +requested by the allocator. + +Pressure reasons are typed and disjoint: + +- memory capacity: reclaim, spill, reduce an unpublished input unit, or degrade + an optional runtime filter; +- account sealed/suspended: terminal lifecycle error; +- owner/site mismatch or allocator invariant: terminal correctness error; +- spill disk cap: spill-resource error, never a memory-reduction retry; +- spill FD cap: spill-resource error, never a memory-reduction retry; +- minimum input unit: terminal capacity error after monotonic progress is no + longer possible. + +Retry is allowed only when progress is observable: retained bytes decrease, +spill epoch advances, input units shrink, or optional state is disabled. This +prevents a capacity loop from replaying the same publication or I/O. + +Runtime filters are optional. If their retained payload cannot be admitted, +HashBuild publishes PASS and releases unpublished scratch. Required hash/join +state does not silently bypass admission. + +## Spill resources + +SpillEngine requires the same live budget generation as its producer. There is +no nil-budget file path. + +Memory, disk bytes, and open file descriptors are separate physical resources: + +- memory is charged by MPool allocations; +- each spill file owns one growable disk token; +- each open spill file owns one FD token; +- a file handoff moves both tokens with the file; +- close releases all three exactly once. + +Recursive spill validates schema, framing, row conservation, queue bounds, and +file metadata. Repartitioning keeps only bounded control arrays plus admitted +scatter buffers. Test fixtures use the same builder copy and budget paths as +production. + +## Grouping semantics + +Grouping sentinels are a distinct key domain from ordinary zero/empty values. +HashBuild selects grouping-aware key encoding whenever input contains grouping +bits. Hashing preserves this distinction for partitioning, equality preserves +it in resident maps, copies preserve the bitmap, and ordering treats the +sentinel as SQL NULL for NULLS FIRST/LAST behavior. + +Sample keeps its ordinary batched hash path unchanged. If grouping bits appear +after ordinary groups were already inserted, it lazily opens a grouping-aware +key domain and translates both maps' local IDs into one sample-pool ID space. +Its iterators are reused per batch, so alternating grouping bits do not create +one iterator allocation per row. + +## Non-goals + +- estimating or limiting total query RSS; +- charging regexp/JQ/JSON/library-internal Go heap as if it were exact; +- changing optimizer join selection; +- making spill as fast as a sufficient-memory in-memory join; +- using remote benchmark workflows as a correctness oracle. + +## Completion criteria + +The implementation is complete only when: + +- no estimated HashBuild value is a query-fatal admission gate for unretained + work; the recovery projection is limited to the retain-versus-spill decision; +- every retained HashBuild/join allocation in the controlled domain has + immutable provenance; +- runtime parallel clones join the current attempt before `Prepare`; +- every transfer has exactly one owner after success and on cancellation; +- memory, disk, and FD rejection remain distinct; +- prepared/retry generations terminate independently at zero; +- remote fragments sharing one MessageBoard terminate at the per-CN statement + boundary, with incomplete dispatch bounded and old board generations unable + to close a replacement; +- local unit, race, build, vet, lifecycle, spill, and performance checks pass; +- independent reviews report no blocker or major correctness/performance issue. diff --git a/docs/design/evidence/26459_allocation_accounting_bench.txt b/docs/design/evidence/26459_allocation_accounting_bench.txt new file mode 100644 index 0000000000000..c7f55f43bbc25 --- /dev/null +++ b/docs/design/evidence/26459_allocation_accounting_bench.txt @@ -0,0 +1,213 @@ +Artifact +======== + +Date: 2026-08-01 +Branch: feature/26459-statement-lifecycle +Base: origin/main 49f9e33cae48 +Code under test: ea5b314cd247 +Go: go1.26.4 linux/amd64 +CPU: 11th Gen Intel Core i7-11700, GOMAXPROCS=16 + +All commands used the same CGo environment as the MatrixOne test wrapper: + + .agents/skills/mo-dev/scripts/mo-cgo-test + +No remote workflow or container build was used. Values below are medians of +the displayed count unless noted otherwise. + + +Physical MPool admission +======================== + +Command: + + mo-cgo-test ./pkg/common/mpool -run '^$' \ + -bench '^BenchmarkMPoolAccountedAllocation$' \ + -benchmem -benchtime=300ms -count=3 + +Operation Unaccounted Accounted Delta +alloc/free 64 B 252.1 ns 356.5 ns +41.4% +alloc/free 4 KiB 302.9 ns 405.4 ns +33.8% +alloc/free 16 KiB 343.4 ns 447.1 ns +30.2% +alloc/free 64 KiB 1024 ns 1123 ns +9.7% +grow replacement 1307 ns 1495 ns +14.4% +parallel accounted alloc/free 64K 280.7 ns + +Every sample reported 0 B/op and 0 allocs/op. The fixed transaction cost is +largest for tiny allocations; retained HashBuild cells and batch buffers use +larger capacity changes rather than one admission per row. Fault-injection +branches that existed only for tests were removed before this measurement. + + +Production shared-budget admission +================================== + +Command: + + mo-cgo-test ./pkg/vm/process -run '^$' \ + -bench '^BenchmarkHashBuildBudgetAllocationAccount$' \ + -benchmem -benchtime=300ms -count=3 + +Operation Unaccounted Budget-controlled Delta +alloc/free 64 B 253.1 ns 503.5 ns +98.9% +alloc/free 4 KiB 301.4 ns 549.7 ns +82.4% +alloc/free 64 KiB 1029 ns 1304 ns +26.7% + +This deliberately measures the full production controller mutex and aggregate +metrics, not only a local AllocationAccount. Every sample remains 0 B/op and +0 allocs/op. This primitive is the worst-case fixed cost; retained-closure +results below show the operator-level effect after allocation reuse. + + +Vector growth and reuse +======================= + +Command: + + mo-cgo-test ./pkg/container/vector -run '^$' \ + -bench '^(BenchmarkVectorAllocationAccount|BenchmarkVectorElementAccounting)$' \ + -benchmem -benchtime=1s -count=3 + +Operation Unaccounted Accounted Allocation result +fixed preextend/free 1066 ns 1172 ns 0 B/op, 0 allocs/op +varlen preextend/free 79088 ns 78729 ns 48 B/op, 2 allocs/op +accounted fixed Reset reuse 2.65 ns 0 B/op, 0 allocs/op + +Accounting adds no Go object to vector allocation or Reset reuse. The varlen +difference was -0.5%, inside local run noise. + +The per-row UnionOne and Copy hot paths were measured separately after +eliminating duplicate bitmap probes and skipping accounted-capacity checks for +ordinary vectors: + +Operation Unaccounted Accounted Delta +UnionOne, retained capacity 11.58 ns 11.83 ns +2.2% +Copy, retained capacity 8.96 ns 8.95 ns -0.1% + +All samples were 0 B/op and 0 allocs/op. + + +Caller-owned Packer storage +=========================== + +Command: + + # origin/main 49f9e33cae48 + taskset -c 4 env GOMAXPROCS=1 mo-cgo-test \ + ./pkg/container/types -run '^$' \ + -bench '^BenchmarkPacker(Encode)?$' \ + -benchmem -cpu=1 -benchtime=500ms -count=8 + + # code commit ea5b314cd247 + taskset -c 4 env GOMAXPROCS=1 mo-cgo-test \ + ./pkg/container/types -run '^$' \ + -bench '^BenchmarkPacker(Encode)?$' \ + -benchmem -cpu=1 -benchtime=500ms -count=8 + +Operation origin/main Code commit Result +allocator-backed retained encode 9.58 ns 9.23 ns -3.7% +fixed-buffer retained encode n/a 9.15 ns +Packer construction/close 48 B/op 48 B/op unchanged + +All retained-encode samples report 0 B/op and 0 allocs/op. The fixed-buffer +path used by serialized runtime filters therefore adds no retained-encode +allocation or CPU regression. Keeping fixed/overflow as one-byte state also +preserves the baseline 48-byte Packer allocation class; an error interface in +the object would have raised it to 64 B/op. Runtime-filter scratch storage is +admitted once through the statement physical account. + + +Retained HashBuild closure +========================== + +Command: + + mo-cgo-test ./pkg/sql/colexec/hashbuild -run '^$' \ + -bench '^BenchmarkResidentHashBuildAccounting$' \ + -benchmem -benchtime=300ms -count=3 + +Benchmark Local account Budget-controlled Delta +resident int key, 32 rows 4034 ns 4400 ns +9.1% +resident varchar key, 32 rows 6953 ns 7649 ns +10.0% +resident int key, 8192 rows 137067 ns 139062 ns +1.5% +resident varchar key, 8192 rows 422190 ns 423558 ns +0.3% + +Both modes use the one production physical-account ownership path; the only +difference is whether that account also uses the shared query/CN controller. +B/op and allocs/op are identical within each pair. The full-batch delta is +0.3-1.5%; the 32-row case exposes the fixed cost rather than a row-scaled cost. + + +Pipeline spool ownership transfer +================================= + +Command: + + mo-cgo-test ./pkg/container/pSpool -run '^$' \ + -bench '^BenchmarkCachedBatchReuse$' \ + -benchmem -benchtime=2s -count=5 + +Median: 1603 ns/op, 224 B/op, 3 allocs/op. + +The same DetachedBuffer path handles both storage outside the controlled +HashBuild domain and account-provenance storage. There is no second raw-byte +cache implementation. + + +Grouping-aware Sample correctness path +====================================== + +Command: + + mo-cgo-test ./pkg/sql/colexec/sample -run '^$' \ + -bench '^(BenchmarkSampleGroupedHashFastPath|BenchmarkSampleAlternatingGrouping)$' \ + -benchmem -benchtime=1s -count=5 + +Ordinary 256-row grouped path: 6075 ns/op, 9344 B/op, 6 allocs/op. +Alternating ordinary/GROUPING rows: 26869 ns/op, 19256 B/op, 12 allocs/op. + +A paired temporary reference containing the pre-fix ordinary hash loop was +measured earlier on this branch in the same binary and then removed (the +Sample code is unchanged by the final review fixes): reference median 5989 ns/op, +final median 6064 ns/op (+1.25%); both were 9344 B/op and 6 allocs/op. The +committed alternating-domain test also applies AllocsPerRun and proves the +iterator count remains constant per batch rather than per grouping run. + + +Spill scatter and serialization +=============================== + +Command: + + mo-cgo-test ./pkg/sql/colexec/spillutil -run '^$' \ + -bench '^BenchmarkSpillScatterAccounting$' \ + -benchmem -benchtime=1s -count=5 + +Final median: 62709 ns/op, 344 B/op, 4 allocs/op, about 523 MB/s. + +An allocation profile found that primitive marshal fallbacks took the address +of integer parameters even when AccountedBuffer's typed writer fast path was +selected. Before the fix the median was 67751 ns/op, 2904 B/op, and 516 +allocs/op. Encoding fallback integers into branch-local fixed arrays reduced +time by 7.8% and removed 512 Go allocations per scatter operation without +changing the wire format. + + +Acceptance +========== + +- physical admission/release adds zero Go allocations; +- the production shared controller adds no Go allocations and no per-row + accounting object; +- vector and spool reuse retain bounded allocation counts; +- UnionOne and Copy retained-capacity paths add no allocations and show at + most a 2.2% local accounting delta; +- caller-owned Packer encoding matches allocator-backed retained encoding and + adds no Go allocations; +- ordinary Sample grouping has unchanged allocation shape and only 1.25% + measured CPU delta; +- mixed GROUPING input uses a fixed number of iterators per batch; +- spill serialization has four bounded Go allocations per 4096-row scatter; +- full-batch resident HashBuild controller overhead is 0.3-1.5%, while the + intentionally worst-case 32-row microbenchmark is 9.1-10.0%; +- no benchmark indicates a compatibility ledger or unaccounted builder path. diff --git a/docs/design/evidence/26459_local_validation.md b/docs/design/evidence/26459_local_validation.md new file mode 100644 index 0000000000000..9b39451504576 --- /dev/null +++ b/docs/design/evidence/26459_local_validation.md @@ -0,0 +1,184 @@ +# #26459 local validation evidence + +This file records evidence for the single production path described in +`../allocation_accounted_memory_admission_impl.md`. Removed implementations +are not retained as validation dimensions. + +## Static closure checks + +- no production allocation-account enable switch; +- no logical-size estimate is a query-fatal gate for unretained HashBuild work; + the shuffle recovery floor is used only before retaining spillable state; +- every join/HashBuild expression-owned MPool vector is constructed with the + attempt account, while opaque library Go heap remains an explicit boundary; +- SpillEngine construction rejects a missing or closed budget generation; +- runtime scan/load clones are attached to the current attempt before worker + `Prepare`; +- terminal lifecycle validates zero bytes and zero live metadata; +- memory, spill disk, and spill FD admission errors have distinct components. + +## Required local test matrix + +All Go tests use the repository CGo wrapper so `usearch` headers, libraries, +link flags, and runtime paths match the MatrixOne build contract. + +```text +.agents/skills/mo-dev/scripts/mo-cgo-test -count=1 -timeout=240s \ + ./pkg/common/mpool ./pkg/common/bitmap ./pkg/common/hashmap/... \ + ./pkg/container/vector ./pkg/container/batch ./pkg/vm/message \ + ./pkg/vm/process ./pkg/sql/util ./pkg/sql/colexec \ + ./pkg/sql/colexec/hashbuild \ + ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/dedupjoin \ + ./pkg/sql/colexec/rightdedupjoin ./pkg/sql/colexec/loopjoin \ + ./pkg/sql/colexec/product ./pkg/sql/colexec/productl2 \ + ./pkg/sql/colexec/spillutil ./pkg/sql/compile +``` + +Selected race coverage: + +```text +.agents/skills/mo-dev/scripts/mo-cgo-test -race -count=1 -timeout=300s \ + ./pkg/common/mpool ./pkg/container/vector ./pkg/container/pSpool \ + ./pkg/sql/util ./pkg/sql/colexec ./pkg/sql/colexec/hashbuild \ + ./pkg/sql/colexec/hashjoin ./pkg/sql/colexec/loopjoin \ + ./pkg/sql/colexec/dedupjoin ./pkg/sql/colexec/rightdedupjoin \ + ./pkg/sql/colexec/productl2 \ + ./pkg/sql/colexec/spillutil ./pkg/sql/colexec/sample ./pkg/sql/compile +``` + +Static checks: + +```text +go vet +go build +``` + +## Fresh local result + +The final semantic edit was followed by a clean local run on 2026-08-01. + +- the complete package matrix above passed, with Sample, pSpool, compare, + hashtable, nulls, shuffle, SQL util, vector-index CPU packages, and disttae + added to the command; +- `-race -p=2 -count=1` passed for mpool, hashjoin, spillutil, Sample, and + compile; +- `go vet -mod=readonly` passed for every modified production package; +- `go build -mod=readonly` passed for every modified production package; +- `git diff --check` and `gofmt` were clean; +- allocation and performance results are recorded in + `26459_allocation_accounting_bench.txt`. + +No partial or still-running session is counted as a pass. + +After the distributed q7 failure exposed the shared remote-MessageBoard +boundary, the final remote-lifecycle amendment was validated separately: + +- the complete `pkg/sql/compile` suite passed in normal and race modes; +- `pkg/pb/pipeline`, `pkg/vm/message`, and `pkg/vm/process` passed; +- protobuf regeneration was clean and reproducible; +- vet passed for all four affected packages; +- an independent lifecycle/concurrency review converged with no blocker or + major finding after its pending-domain cardinality and mixed-version concerns + were resolved in code or by existing scheduling evidence. + +## Behavioral coverage + +The local suite covers: + +- exact physical allocation/release and capacity rollback; +- owner/site mismatch and sealed-generation terminal errors; +- prepared and retry generation reuse; +- runtime parallel clone attachment; +- JoinMap and spill payload move-only ownership; +- cancellation-safe bitmap mailbox seal/drain; +- spill disk/FD admission and release; +- recursive spill row/schema/file validation; +- minimum-unit and monotonic pressure termination; +- optional runtime-filter degradation; +- nested/selected expression result accounting, capacity rejection, transfer, + and terminal release; +- HashJoin, LoopJoin, DedupJoin multi-batch finalize, and RightDedupJoin result + accounting, batch-local reuse, capacity rejection, and prepared + `Reset -> ClearAllocationAccount` terminal release; +- grouping-aware copy, hash, equality, and ordering; +- late and alternating Sample grouping domains across row, percent, and merge + modes; +- Product cleanup and account terminal zero; +- accounted JoinMap release after an unaccounted ProductL2 consumer frees it. +- remote scope-graph fragment counting across nested CN execution addresses; +- ProcessInfo topology-map and execution-ID wire round trips, with distinct + MessageBoard generations across retries; +- remote statement-group board drain after a producer finishes before a later + sibling registers, including accounted queued-message destruction, exact + account terminal zero, and aggregate MPool terminal sampling; +- incomplete remote dispatch expiry releases staged account and MPool domains, + while an unresolved pending terminal marker makes the coordinator summary + explicitly partial; this includes the case where another registered fragment + is still active when the old board generation closes; +- fragment failure aborts an incomplete group immediately, and a legacy remote + plan containing an accounted owner is rejected instead of running without an + account; +- counted pending/completed group markers resolve independent of + terminal-response order; a four-fragment lost-final-response case preserves + all three suppressed reported domains plus the directly missing domain. + +## Performance evidence + +Performance validation is local and allocation-focused. Benchmarks record +`ns/op`, `B/op`, and `allocs/op` for account acquire/release, vector growth, +hash-map build/lookup, and spill scatter. The acceptance rule is no +new per-row or per-allocation Go object in steady state and no material +regression outside measurement noise. + +The local measurements are complemented by the distributed validation below. + +## Distributed workload evidence + +The final semantic head before the two review counterexample fixes, +`d13b9103c8`, completed the TPCH 100G and 1T TKE run +[`30758186183`](https://github.com/matrixorigin/mo-auto-test/actions/runs/30758186183). +The workflow built that commit, loaded the native fixtures in 11 seconds and +49 seconds, and compared every Q1-Q22 result with its golden result. No query +failed and the run reported no OOM or budget-admission error. + +The measured query-only totals were: + +| Workload | Candidate turns | Candidate average | Recent main average | Delta | +| --- | --- | ---: | ---: | ---: | +| TPCH 100G | 97.739 / 95.005 / 94.982 / 97.469 s | 96.298 s | 98.252 s | -1.99% | +| TPCH 1T | 1045.573 / 1039.371 s | 1042.472 s | 1027.439 s | +1.46% | + +The cited main result is job +[`91396824792`](https://github.com/matrixorigin/mo-nightly-regression/actions/runs/30708854656/job/91396824792): +100G turns were 103.550 / 95.127 / 97.321 / 97.010 seconds and 1T turns +were 1027.322 / 1027.557 seconds. These are adjacent runs of the same TKE +benchmark shape, not a simultaneous same-base A/B; the deltas establish that +the stabilized candidate is within normal workload variance, not a stronger +causal performance claim. Compared with the earlier regressed candidate run +`30738374292` (135.093 seconds for 100G and 1591.049 seconds for 1T), this head +recovered 28.7% and 34.5% respectively. + +The current review fixes after `d13b9103c8` are allocation-boundary and +late-RPC-lifetime corrections. They add no per-row work: existing-buffer grow +now passes the logical requirement to the allocator's single capacity policy, +and aborted remote generations retain only a key and timer for the maximum +possible RPC lifetime. Their focused and package validation is recorded in the +PR review response after the final commit. + +## Incident acceptance matrix + +This matrix separates durable mechanism regressions from workload executions; +one is not presented as a substitute for the other. + +| Incident | Durable regression retained on this branch | Workload evidence | Current-head gap | +| --- | --- | --- | --- | +| #26174 | HashBuild build/hashmap/spill regressions introduced by #26178, plus exact physical batch/vector allocation boundaries in this PR | #26178 TKE BVT: all three 3,840,001-row fulltext inserts succeeded with zero HashBuild rejection | full fulltext workload has not been rerun at the final head | +| #26192 | exact accounted runtime-filter payload, one-byte-short PASS degradation, varlena/null coverage, and spill decode/reuse lifecycle tests | historical LOAD failure shape is covered by #26231/#26318; the current TPCH fixture LOAD path succeeds | the original `ca_comprehensive_dataset` workload has not been rerun at the final head | +| #26413 | segmented `CopyIntoBatches` and accounted hash-map growth/rollback regressions, including large external-batch shapes | #26438 verified the real Parquet self-join with both expected 50,000-row results | the Hive fixture has not been rerun at the final head | +| #26454 | `TestIssue26454ExpressionKeyBuildUsesActualCapacity` exercises the CONCAT/CAST and CASE key shapes under a 16 MiB physical account and validates terminal zero; `TestShuffleHashBuildDirectSpillUsesActualAllocation` proves the conservative recovery projection cannot reject an upstream-owned direct source | the exact jinpan SQL/data is not available in this repository | full jinpan workload remains external evidence | +| #26586 | retained shuffle batches reserve recovery ownership before copy, while direct spill uses allocation-led admission; the focused direct-spill regression and full HashBuild race suite pass | `TestHashBuildSharedBudgetRecoverySQL` passes locally with the 28 MiB shared-budget reproducer | the external TPCH Q9 workload is not rerun for this control-flow-only correction | +| #25782 | `TestShuffleHashBuildAccountedSpillLifecycle`, `TestHashTableAccountedHighCardinalityResizeReturnsToZero`, broadcast error propagation, recursive spill, and terminal-zero tests | the two-CN 132,096-row harness at `f5cc97efe7` returned the exact count with positive spill and zero OOM; current-head TPCH 1T also completed without OOM/query failure | the private original high-cardinality SQL harness has not been rerun at the final head | + +Accordingly, the current TKE TPCH acceptance is complete, while the unavailable +external-data workloads and the original private #25782 harness remain explicit +follow-up evidence rather than being silently marked complete. diff --git a/docs/rfcs/00000000_allocation_accounted_memory_admission.md b/docs/rfcs/00000000_allocation_accounted_memory_admission.md new file mode 100644 index 0000000000000..58eaaf9b25244 --- /dev/null +++ b/docs/rfcs/00000000_allocation_accounted_memory_admission.md @@ -0,0 +1,289 @@ +- Status: implementation validation +- Start Date: 2026-07-30 +- Authors: aptend +- Issue: [#26459](https://github.com/matrixorigin/matrixone/issues/26459) +- Implementation plan: + [Allocation-accounted HashBuild memory admission](../design/allocation_accounted_memory_admission_impl.md) + +# Allocation-accounted memory admission for spillable SQL execution + +## Summary + +HashBuild needs a finite query/CN memory boundary, but a predicted future size +is not a reliable hard-admission fact. SQL maximum widths, payload multipliers, +recursive expression sums, and logical batch sizes can reject valid work by +orders of magnitude. Relaxing those estimates can instead miss a real +allocate-copy-free overlap and allow OOM. + +This RFC makes the capacity of live physical allocations the source of truth: + +```text +hard HashBuild admission = live retained physical allocation capacity +prediction = planning and early-spill hint only +``` + +MPool admits an accounted allocation before allocating it. Immutable account, +owner, and site provenance follows the allocation across reuse and ownership +handoff. The same physical Free releases the charge. Growth admits the actual +replacement capacity while the old allocation is still charged, so overlap is +represented without a multiplier. + +The first consumer is retained HashBuild and join storage. The design has one +production path: no activation switch, estimated-memory reservation, or +compatibility ledger remains beside physical accounting. + +## Motivation + +The related incidents show both sides of the same modeling error: + +| Incidents | Failure | Mismatch | +| --- | --- | --- | +| #25782 / #25837 | CN OOM instead of bounded spill | hash-table growth was not admitted at allocation time | +| #26174 / #26178 | false fulltext INSERT rejection | non-simultaneous ingress/tail/drain states were summed | +| #26192 / #26231 / #26318 | false LOAD DATA rejection | runtime-filter/payload multipliers duplicated live owners | +| #26413 / #26438 | false external-table join rejection | logical ingress did not match physical batch segmentation | +| #26454 | request reported hundreds of GiB with single-digit-GiB usage | TEXT maximum width was multiplied through CAST/CONCAT | +| #26186 | spill overlap could be under-counted | independent estimates omitted a live transition state | + +Fixing individual multipliers cannot close this class. A stricter estimate +rejects more valid queries; a looser estimate misses a different overlap. + +Three concerns must remain separate: + +1. ownership accounting: which retained allocations are live; +2. prediction: which strategy or capacity may be useful; +3. pressure response: spill, reclaim, reduce, degrade, or error. + +Only ownership accounting can support a hard live-memory invariant. + +## Goals + +1. Make each HashBuild hard memory charge correspond to live physical retained + storage. +2. Admit before allocation and leave allocator/account state unchanged on + rejection. +3. Represent real replacement overlap during growth. +4. Preserve provenance through vector reuse, copies, and owner transfer. +5. Give Reset, Free, cancellation, retry, and prepared reuse one lifecycle. +6. Treat only typed memory-capacity rejection as reclaimable pressure. +7. Prevent a prediction alone from rejecting a statement. +8. Add no per-row reservation or Go allocation to steady-state reuse paths. +9. Keep spill disk and FD accounting distinct from memory. +10. Preserve the no-OOM objective introduced for #25782. + +## Non-goals + +This RFC does not: + +- equate the HashBuild account with total query RSS; +- account all Go runtime/library allocations; +- guarantee completion under every finite cap; +- change optimizer join selection; +- make disk spill as fast as sufficient-memory execution; +- add type- or workload-specific exemptions; +- raise a cap to hide a false estimate; +- introduce a mutable process-wide current account. + +## Controlled domain + +The first implementation accounts physical storage retained by HashBuild and +its join consumers: hash tables, copied build batches, JoinMap state, retained +keys, join bitmaps/capture/result state, Product result state, +runtime-filter payloads, and spill encode/decode/rebuild buffers. + +Every MPool-backed vector owned by a HashBuild/join ExpressionExecutor tree is +inside the account: fixed and variable values, reusable function results, +nested child results, selected-row buffers, and list results. Allocation +provenance is installed when the tree is constructed, before the first owned +capacity allocation. Borrowed input and serialized-plan vectors keep their +source ownership. + +Library-internal Go heap used by regexp, JSON, JQ, and similar implementations +remains outside this controlled domain because those libraries do not expose +allocator/free hooks. The account is exact for its stated allocator-visible +domain; it is not advertised as total expression memory or total query RSS. + +This is a static ownership boundary, not a runtime fallback. A future general +expression-memory design must either use allocator-aware implementations or a +separate explicitly non-exact transient policy. Estimates must not be inserted +into this exact ledger. + +ProductL2's accounted input remains the producer-owned JoinMap. Its additional +CPU/GPU index and scratch storage is outside the first controlled domain because +the native GPU allocator does not expose an admission/free capacity contract. +Partially charging only Go-visible buffers would be another ledger, not exact +physical ownership. ProductL2 can join this domain only when both platform +implementations provide the same allocation contract. + +## Allocation contract + +An accounted allocation records: + +- actual allocator capacity; +- allocation account and generation; +- bounded owner class and allocation site. + +The contract is transactional: + +1. compute the allocator capacity using checked arithmetic; +2. acquire account/controller capacity; +3. allocate physical storage; +4. publish metadata and ownership; +5. on failure, undo the acquisition before returning; +6. on Free, deallocate and release exactly once. + +For growth, old and replacement allocations are both charged until the copy +succeeds and the old allocation is physically freed. Views borrow storage and +do not charge again. Copies charge their destination allocation. + +The query generation and CN aggregate are checked at the same controller +boundary. There is no operator-owned memory reservation token parallel to the +allocator lease. + +## Lifecycle + +Compile opens one generation for each local statement attempt, configures all +physical-plan owners, and records them for reverse-order cleanup. Parallel +scan/load workers created during `runOnce` join the same generation before +worker `Prepare` starts. Configuration is atomic: a failure clears newly +configured owners in reverse order. + +At completion, the message board is drained, owners are cleared, and the +account is sealed. A valid terminal state requires zero live bytes and zero +live metadata. A late allocation or release mismatch is a lifecycle invariant, +not capacity pressure. Prepared statements and retries use new generations; +generation-bound state cannot survive Reset. + +Remote pipeline RPCs for one statement on one CN share a MessageBoard, so an +individual RPC is not a terminal ownership boundary. The coordinator carries +the exact number of planned fragments per target CN plus a unique physical +execution ID. Every received fragment joins the execution-and-board-keyed +statement group, clears its reachable operators when it quiesces, and transfers +its account and MPool terminal ownership to that group. Quiesced handlers do not +wait for siblings, because nested B-to-C-to-B execution can otherwise create a +response dependency cycle. Their terminal responses carry a counted pending +memory-domain signal. The finalizer drains the board, completes all accounts, samples all +fragment MPool domains, and lets exactly one terminal response publish the +aggregate plus a completion marker; the other responses publish no duplicate +terminal facts. Pending/completed signals reduce independent of response order; +unresolved counts preserve the cardinality of every suppressed fragment MPool +domain and make the root summary explicitly partial. +A missing planned RPC is bounded by the MessageBoard receive interval: expiry +closes the old board generation and cancels active registered fragments; their +accounts remain live until those fragments actually quiesce. A fragment failure +aborts the incomplete group immediately. The timer is canceled once all planned +fragments register, so it never limits a fully dispatched statement. Prepared +executions and retries use new execution IDs, so an incomplete old group cannot +capture a replacement. + +The topology and execution-ID fields are a required capability for remote +plans containing allocation-account owners. A legacy ProcessInfo remains +decodable, and remote plans outside the accounted domain keep their legacy +lifecycle. An accounted legacy plan is rejected explicitly rather than running +with a guessed fragment count or an unaccounted fallback. Query-candidate +discovery admits only CNs with the coordinator binary's CommitID, which prevents +the inverse new-coordinator-to-old-handler execution during a rolling upgrade. + +## Pressure protocol + +Typed reasons keep control flow honest: + +- memory capacity may reclaim, spill, reduce an unpublished input, or disable + an optional optimization; +- sealed, suspended, owner/site mismatch, or allocator invariant is terminal; +- spill disk and spill FD rejection are resource-specific terminal results; +- minimum-unit pressure ends retries when no smaller valid input exists. + +A retry requires monotonic evidence: fewer live bytes, a new spill epoch, a +smaller input unit, or an optional structure disabled. Publication and spill +I/O cannot be replayed merely because the same capacity error recurred. + +Runtime filters are optional and degrade to PASS if their retained payload +cannot be admitted. Required join state never bypasses admission. + +## Spill ownership + +SpillEngine must receive the producer's live budget generation. Memory, disk, +and FD tokens describe different physical resources. One file owns one +growable disk token and one FD token; file handoff moves both. Close releases +them once. Recursive spill validates framing, schema, row conservation, file +metadata, and bounded queue progress. + +## Relationship to SQL Resource Accounting + +[`SQL Resource Accounting`](./00000000_sql_resource_accounting.md) observes and +persists statement resource facts. This RFC controls an allocation before it +happens. + +| Concern | SQL Resource Accounting | This RFC | +| --- | --- | --- | +| Purpose | observation and diagnostics | admission and pressure control | +| Time | during/after execution | immediately before allocation | +| Unit | domain usage/peak | live physical retained capacity | +| Missing data | quality flag | storage is outside this controlled domain | +| Failure action | report | reclaim/spill/reduce/degrade/error | + +The systems may cross-check terminal facts, but neither authorizes the other. + +## Performance constraints + +- no per-row account object; +- no per-allocation Go object in steady state; +- fixed-size provenance in allocator metadata; +- controller calls only when physical capacity changes; +- views and Reset reuse do not re-admit unchanged storage; +- spill tokens scale with open files, not records; +- no extra expression-account bookkeeping in hot predicates. + +Benchmarks must compare unaccounted MPool behavior, accounted acquire/release, +vector growth/reuse, hash build/lookup, and spill scatter. Correctness +requires exact terminal zero; performance acceptance requires no material +regression outside measurement noise for unaccounted paths and no new +data-scaled Go allocation in accounted paths. + +## Alternatives rejected + +### Patch each multiplier + +It cannot close both false-positive and false-negative estimate errors. + +### Reserve an estimated amount, reconcile later + +The estimate can reject before real allocation or under-count live overlap. It +also creates a second release owner beside MPool. + +### Sample RSS or heap counters + +Sampling is observational, process-wide, and too late to authorize a specific +allocation. + +### Estimate or manually charge selected expression buffers + +This creates another ledger and still reports false exactness. The +implementation instead propagates one immutable allocation selection through +the complete expression tree and charges every physical MPool capacity change; +opaque library Go heap remains an explicit boundary rather than an estimate. + +### Keep old and new production paths behind a switch + +Two admission semantics double ownership states, tests, and failure modes. The +final implementation removes the old path instead. + +## Rollout and validation + +The merge unit is one PR with reviewable commits, but the final branch must +present a single production path. Local validation includes unit tests, +prepared/retry and parallel-clone lifecycle tests, cancellation/transfer tests, +spill resource tests, race tests, build, vet, and allocation benchmarks. + +Remote benchmark workflows are not required for correctness convergence in +this implementation cycle. + +## Completion invariant + +The RFC is complete when every retained HashBuild/join allocation in scope has +one physical owner, one admission path, and one terminal release; estimated +hard memory reservations and activation gates are absent; runtime clones share +the statement generation; all pressure/resource types remain disjoint; and +independent review plus fresh local validation reports no blocker or major +regression. diff --git a/pkg/common/bitmap/bitmap.go b/pkg/common/bitmap/bitmap.go index 80ed391dff5a0..cc9a10c9f8150 100644 --- a/pkg/common/bitmap/bitmap.go +++ b/pkg/common/bitmap/bitmap.go @@ -18,9 +18,12 @@ import ( "bytes" "encoding" "fmt" + "io" + "math" "math/bits" "unsafe" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/container/types" ) @@ -32,6 +35,8 @@ import ( type bitmask = uint64 +const MarshalHeaderSize = 24 + /* * Array giving the position of the right-most set bit for each possible * byte value. count the right-most position as the 0th bit, and the @@ -57,16 +62,62 @@ var rightmost_one_pos_8 = [256]uint8{ 4, 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, } +func encodeTaggedLen(length int64, external bool) int64 { + if length < 0 { + panic("negative bitmap length") + } + if external { + return ^length + } + return length +} + +func (n *Bitmap) logicalLen() int64 { + if n.taggedLen < 0 { + return ^n.taggedLen + } + return n.taggedLen +} + +func (n *Bitmap) setLogicalLen(length int64) { + n.taggedLen = encodeTaggedLen(length, n.HasExternalStorage()) +} + func (n *Bitmap) InitWith(m *Bitmap) { - n.len = m.len + if n == m { + return + } + n.setLogicalLen(m.logicalLen()) n.count = m.count + if n.HasExternalStorage() { + if len(m.data) > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + previousLength := len(n.data) + storage := n.data[:cap(n.data)] + clear(storage[len(m.data):max(previousLength, len(m.data))]) + n.data = storage[:len(m.data)] + copy(n.data, m.data) + return + } n.data = append([]uint64(nil), m.data...) } -func (n *Bitmap) InitWithSize(len int64) { - n.len = len +func (n *Bitmap) InitWithSize(length int64) { + n.setLogicalLen(length) n.count = 0 - n.data = make([]uint64, (len+63)/64) + words := int((length + 63) / 64) + if n.HasExternalStorage() { + if words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + previousLength := len(n.data) + storage := n.data[:cap(n.data)] + clear(storage[:max(previousLength, words)]) + n.data = storage[:words] + return + } + n.data = make([]uint64, words) } func (n *Bitmap) Clone() *Bitmap { @@ -113,7 +164,7 @@ func (itr *BitmapIterator) hasNext(i uint64) (uint64, bool) { // if the uint64 is not 0, then calculate the rightest_one position in a word, add up prev result and return. // when there is 1 in Bitmap, return true, otherwise Bitmap is empty and return false. // either case loop over words not bits - nwords := (itr.bm.len + 63) / 64 + nwords := (itr.bm.logicalLen() + 63) / 64 current_word := i >> 6 mask := (^(bitmask)(0)) << (i & 0x3F) // ignore bits check before var result uint64 @@ -158,14 +209,69 @@ func (itr *BitmapIterator) Next() uint64 { // Reset set n.data to nil func (n *Bitmap) Reset() { - n.len = 0 + n.setLogicalLen(0) n.count = 0 + if n.HasExternalStorage() { + clear(n.data) + storage := n.data[:cap(n.data)] + n.data = storage[:0] + return + } + n.data = nil +} + +// InstallExternalStorage replaces the bitmap backing with caller-owned +// storage while preserving the logical bitmap. The caller remains responsible +// for releasing the returned previous external storage, if any. +func (n *Bitmap) InstallExternalStorage(storage []uint64) []uint64 { + required := len(n.data) + if required > cap(storage) { + panic("bitmap external storage capacity exceeded") + } + var previous []uint64 + if n.HasExternalStorage() && cap(n.data) > 0 { + previous = n.data[:cap(n.data)] + } + target := storage[:cap(storage)] + if len(target) > required { + clear(target[required:]) + } + copy(target[:required], n.data) + n.data = target[:required] + n.taggedLen = encodeTaggedLen(n.logicalLen(), true) + return previous +} + +// ReleaseExternalStorage detaches caller-owned storage and clears the bitmap. +// It returns nil for a bitmap that owns its Go-allocated backing. +func (n *Bitmap) ReleaseExternalStorage() []uint64 { + if !n.HasExternalStorage() { + return nil + } + var storage []uint64 + if cap(n.data) > 0 { + storage = n.data[:cap(n.data)] + } + n.count = 0 + n.taggedLen = 0 n.data = nil + return storage +} + +func (n *Bitmap) ExternalStorageCapacity() int { + if n == nil || !n.HasExternalStorage() { + return 0 + } + return cap(n.data) +} + +func (n *Bitmap) HasExternalStorage() bool { + return n != nil && n.taggedLen < 0 } // Len returns the number of bits in the Bitmap. func (n *Bitmap) Len() int64 { - return n.len + return n.logicalLen() } // Size return number of bytes in n.data @@ -210,7 +316,7 @@ func (n *Bitmap) AddMany(rows []uint64) { } func (n *Bitmap) Remove(row uint64) { - if row >= uint64(n.len) { + if row >= uint64(n.logicalLen()) { return } if n.data[row>>6]&(1<<(row&0x3F)) != 0 { @@ -221,7 +327,7 @@ func (n *Bitmap) Remove(row uint64) { // Contains returns true if the row is contained in the Bitmap func (n *Bitmap) Contains(row uint64) bool { - if row >= uint64(n.len) { + if row >= uint64(n.logicalLen()) { return false } idx := row >> 6 @@ -255,8 +361,8 @@ func (n *Bitmap) AddRange(start, end uint64) { } func (n *Bitmap) RemoveRange(start, end uint64) { - if end > uint64(n.len) { - end = uint64(n.len) + if end > uint64(n.logicalLen()) { + end = uint64(n.logicalLen()) } if start >= end { return @@ -297,7 +403,7 @@ func (n *Bitmap) IsSame(b *Bitmap) bool { func (n *Bitmap) Or(b *Bitmap) { n.TryExpand(b) - size := (int(b.len) + 63) / 64 + size := (int(b.logicalLen()) + 63) / 64 for i := range size { cnt := bits.OnesCount64(n.data[i]) n.data[i] |= b.data[i] @@ -308,7 +414,7 @@ func (n *Bitmap) Or(b *Bitmap) { func (n *Bitmap) And(b *Bitmap) { n.TryExpand(b) n.count = 0 - size := (int(b.len) + 63) / 64 + size := (int(b.logicalLen()) + 63) / 64 for i := range size { n.data[i] &= b.data[i] n.count += int64(bits.OnesCount64(n.data[i])) @@ -319,7 +425,7 @@ func (n *Bitmap) And(b *Bitmap) { } func (n *Bitmap) Negate() { - nBlock, nTail := int(n.len)/64, int(n.len)%64 + nBlock, nTail := int(n.logicalLen())/64, int(n.logicalLen())%64 n.count = 0 for i := range nBlock { n.data[i] = ^n.data[i] @@ -333,16 +439,19 @@ func (n *Bitmap) Negate() { } func (n *Bitmap) TryExpand(m *Bitmap) { - n.TryExpandWithSize(int(m.len)) + n.TryExpandWithSize(int(m.logicalLen())) } func (n *Bitmap) TryExpandWithSize(size int) { - if int(n.len) >= size { + if int(n.logicalLen()) >= size { return } newCap := (size + 63) / 64 - n.len = int64(size) + n.setLogicalLen(int64(size)) if newCap > cap(n.data) { + if n.HasExternalStorage() { + panic("bitmap external storage capacity exceeded") + } data := make([]uint64, newCap) copy(data, n.data) n.data = data @@ -353,9 +462,176 @@ func (n *Bitmap) TryExpandWithSize(size int) { } } +// RemapOrdered rewrites the bitmap in place for an ordered row selection. +// When negate is false, output row i comes from sels[i]. When negate is true, +// sels identifies rows to remove. The caller must provide strictly increasing, +// non-negative row indexes. Because every destination row is at or before its +// source row, one cached source word is sufficient to avoid allocating a +// second data-scaled bitmap. Selection rows beyond the bitmap's logical +// length are valid and read as clear: a null bitmap may be shorter than its +// owning vector when the vector's trailing rows are all non-null. +func (n *Bitmap) RemapOrdered(sels []int64, negate bool) { + if n == nil { + return + } + oldLength := n.logicalLen() + previous := int64(-1) + for _, sel := range sels { + if sel <= previous || sel < 0 { + panic("bitmap ordered remap requires strictly increasing non-negative rows") + } + previous = sel + } + logicalLength := int64(len(sels)) + if negate { + logicalLength = oldLength + } + n.prepareOrderedRemap(logicalLength) + + sourceWordIndex := int64(-1) + var sourceWord uint64 + readSource := func(row int64) bool { + if row >= oldLength { + return false + } + wordIndex := row >> 6 + if wordIndex != sourceWordIndex { + sourceWordIndex = wordIndex + sourceWord = n.data[wordIndex] + } + return sourceWord&(uint64(1)<> 6 + mask := uint64(1) << uint(row&63) + if value { + n.data[wordIndex] |= mask + } else { + n.data[wordIndex] &^= mask + } + } + + output := int64(0) + if !negate { + for _, sel := range sels { + writeDestination(output, readSource(sel)) + output++ + } + } else { + selIndex := 0 + for source := int64(0); source < oldLength; source++ { + if selIndex < len(sels) && source == sels[selIndex] { + selIndex++ + continue + } + writeDestination(output, readSource(source)) + output++ + } + } + n.finishOrderedRemap(output, logicalLength) +} + +// RemapMaskOrdered is RemapOrdered for an ordered bitmap selection. Selection +// bitmap iteration is monotonic, so the rewrite uses no row-scaled scratch. +func (n *Bitmap) RemapMaskOrdered(sels *Bitmap, negate bool) { + if n == nil || sels == nil { + return + } + oldLength := n.logicalLen() + logicalLength := int64(sels.Count()) + if negate { + logicalLength = oldLength + } + n.prepareOrderedRemap(logicalLength) + sourceWordIndex := int64(-1) + var sourceWord uint64 + readSource := func(row int64) bool { + if row >= oldLength { + return false + } + wordIndex := row >> 6 + if wordIndex != sourceWordIndex { + sourceWordIndex = wordIndex + sourceWord = n.data[wordIndex] + } + return sourceWord&(uint64(1)<> 6 + mask := uint64(1) << uint(row&63) + if value { + n.data[wordIndex] |= mask + } else { + n.data[wordIndex] &^= mask + } + } + + output := int64(0) + iterator := sels.Iterator() + if !negate { + for iterator.HasNext() { + source := int64(iterator.Next()) + writeDestination(output, readSource(source)) + output++ + } + } else { + var selected int64 = -1 + if iterator.HasNext() { + selected = int64(iterator.Next()) + } + for source := int64(0); source < oldLength; source++ { + if source == selected { + if iterator.HasNext() { + selected = int64(iterator.Next()) + } else { + selected = -1 + } + continue + } + writeDestination(output, readSource(source)) + output++ + } + } + n.finishOrderedRemap(output, logicalLength) +} + +func (n *Bitmap) prepareOrderedRemap(logicalLength int64) { + words := int((logicalLength + 63) / 64) + if words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + if words > len(n.data) { + storage := n.data[:cap(n.data)] + clear(storage[len(n.data):words]) + n.data = storage[:words] + } +} + +func (n *Bitmap) finishOrderedRemap(written, logicalLength int64) { + words := int((logicalLength + 63) / 64) + if written < logicalLength { + word := int(written >> 6) + if tail := uint(written & 63); tail != 0 { + n.data[word] &= (uint64(1) << tail) - 1 + word++ + } + clear(n.data[word:words]) + } + if words > 0 && logicalLength&63 != 0 { + n.data[words-1] &= (uint64(1) << uint(logicalLength&63)) - 1 + } + clear(n.data[words:]) + n.data = n.data[:words] + n.setLogicalLen(logicalLength) + n.count = 0 + for _, word := range n.data { + n.count += int64(bits.OnesCount64(word)) + } +} + func (n *Bitmap) Filter(sels []int64) *Bitmap { var b Bitmap - b.InitWithSize(n.len) + b.InitWithSize(n.logicalLen()) for i, sel := range sels { if n.Contains(uint64(sel)) { b.Add(uint64(i)) @@ -368,6 +644,61 @@ func (n *Bitmap) Count() int { return int(n.count) } +// CountRange returns the number of set bits in [start, end). It never scans +// outside the bitmap's logical coverage and does not allocate. +func (n *Bitmap) CountRange(start, end uint64) int { + if n == nil || start >= end || start >= uint64(n.logicalLen()) { + return 0 + } + if end > uint64(n.logicalLen()) { + end = uint64(n.logicalLen()) + } + first := start >> 6 + last := (end - 1) >> 6 + if first == last { + mask := (^uint64(0) << (start & 63)) & + (^uint64(0) >> ((-end) & 63)) + return bits.OnesCount64(n.data[first] & mask) + } + count := bits.OnesCount64(n.data[first] & (^uint64(0) << (start & 63))) + for word := first + 1; word < last; word++ { + count += bits.OnesCount64(n.data[word]) + } + count += bits.OnesCount64(n.data[last] & (^uint64(0) >> ((-end) & 63))) + return count +} + +// AnySetNotIn reports whether [start, end) contains a bit set in n and not in +// other. It is used when one provenance bitmap (GROUPING) overrides another +// (SQL NULL) without expanding either bitmap row by row. +func (n *Bitmap) AnySetNotIn(other *Bitmap, start, end uint64) bool { + if n == nil || start >= end || start >= uint64(n.logicalLen()) { + return false + } + if end > uint64(n.logicalLen()) { + end = uint64(n.logicalLen()) + } + first := start >> 6 + last := (end - 1) >> 6 + for word := first; word <= last; word++ { + mask := ^uint64(0) + if word == first { + mask &= ^uint64(0) << (start & 63) + } + if word == last { + mask &= ^uint64(0) >> ((-end) & 63) + } + value := n.data[word] & mask + if other != nil && word < uint64(len(other.data)) { + value &^= other.data[word] + } + if value != 0 { + return true + } + } + return false +} + func (n *Bitmap) ToArray() []uint64 { rows := make([]uint64, 0, n.Count()) ToArray(n, &rows) @@ -394,20 +725,143 @@ func (n *Bitmap) ToI64Array(out *[]int64) []int64 { func (n *Bitmap) Marshal() []byte { var buf bytes.Buffer - u1 := uint64(n.len) - u2 := uint64(len(n.data) * 8) - buf.Write(types.EncodeInt64(&n.count)) - buf.Write(types.EncodeUint64(&u1)) - buf.Write(types.EncodeUint64(&u2)) - buf.Write(types.EncodeSlice(n.data)) + _ = n.MarshalTo(&buf) return buf.Bytes() } +func (n *Bitmap) MarshalSize() int { + if n == nil { + return 0 + } + return MarshalHeaderSize + len(n.data)*8 +} + +// Validate checks the in-memory representation after streaming decode. +func (n *Bitmap) Validate() error { + if n == nil || n.logicalLen() < 0 || n.count < 0 || + n.count > n.logicalLen() || + len(n.data) != int((n.logicalLen()+63)/64) { + return moerr.NewInvalidInputNoCtx("invalid bitmap representation") + } + actual := int64(0) + for i, word := range n.data { + if i == len(n.data)-1 && n.logicalLen()%64 != 0 && + word>>uint(n.logicalLen()%64) != 0 { + return moerr.NewInvalidInputNoCtx("invalid bitmap trailing bits") + } + actual += int64(bits.OnesCount64(word)) + } + if actual != n.count { + return moerr.NewInvalidInputNoCtx("invalid bitmap count") + } + return nil +} + +// DecodeMarshalHeader validates the fixed bitmap wire header. +func DecodeMarshalHeader(data []byte) ( + count int64, + bitLength int64, + dataSize int, + err error, +) { + if len(data) < MarshalHeaderSize { + return 0, 0, 0, io.ErrUnexpectedEOF + } + count = types.DecodeInt64(data[:8]) + rawBitLength := types.DecodeUint64(data[8:16]) + rawDataSize := types.DecodeUint64(data[16:24]) + if count < 0 || + rawBitLength > math.MaxInt64 || + rawDataSize > math.MaxInt || + rawDataSize%8 != 0 { + return 0, 0, 0, moerr.NewInvalidInputNoCtx("invalid bitmap wire header") + } + bitLength = int64(rawBitLength) + dataSize = int(rawDataSize) + if count > bitLength || + uint64(dataSize/8) != (rawBitLength+63)/64 { + return 0, 0, 0, moerr.NewInvalidInputNoCtx("invalid bitmap wire header") + } + return count, bitLength, dataSize, nil +} + +// PrepareExternalUnmarshal publishes a validated bitmap header into existing +// caller-owned storage and returns the payload bytes to fill. +func (n *Bitmap) PrepareExternalUnmarshal( + header []byte, + totalSize int, +) ([]byte, error) { + if !n.HasExternalStorage() { + return nil, moerr.NewInvalidInputNoCtx("bitmap does not use external storage") + } + count, bitLength, dataSize, err := DecodeMarshalHeader(header) + if err != nil { + return nil, err + } + if totalSize != MarshalHeaderSize+dataSize || + dataSize/8 > cap(n.data) { + return nil, moerr.NewInvalidInputNoCtx("invalid bitmap external storage capacity") + } + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:dataSize/8] + n.count = count + n.setLogicalLen(bitLength) + return types.EncodeSlice(n.data), nil +} + +func (n *Bitmap) MarshalTo(w io.Writer) error { + if n == nil { + return nil + } + if w == nil { + return io.ErrClosedPipe + } + bitLength := uint64(n.logicalLen()) + dataLength := uint64(len(n.data) * 8) + if typed, ok := w.(interface { + WriteInt64(int64) error + WriteUint64(uint64) error + }); ok { + if err := typed.WriteInt64(n.count); err != nil { + return err + } + if err := typed.WriteUint64(bitLength); err != nil { + return err + } + if err := typed.WriteUint64(dataLength); err != nil { + return err + } + return writeBitmapMarshalBytes(w, types.EncodeSlice(n.data)) + } + if err := writeBitmapMarshalBytes(w, types.EncodeInt64(&n.count)); err != nil { + return err + } + if err := writeBitmapMarshalBytes(w, types.EncodeUint64(&bitLength)); err != nil { + return err + } + if err := writeBitmapMarshalBytes(w, types.EncodeUint64(&dataLength)); err != nil { + return err + } + return writeBitmapMarshalBytes(w, types.EncodeSlice(n.data)) +} + +func writeBitmapMarshalBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } + return nil +} + // MarshalV1 in version 1, Bitmap.emptyFlag is type int32, now we use Bitmap.count replace it func (n *Bitmap) MarshalV1() []byte { var buf bytes.Buffer empty := int32(0) - u1 := uint64(n.len) + u1 := uint64(n.logicalLen()) u2 := uint64(len(n.data) * 8) buf.Write(types.EncodeInt32(&empty)) buf.Write(types.EncodeUint64(&u1)) @@ -419,21 +873,41 @@ func (n *Bitmap) MarshalV1() []byte { func (n *Bitmap) Unmarshal(data []byte) { n.count = types.DecodeInt64(data[:8]) data = data[8:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] if size == 0 { - n.data = nil + if n.HasExternalStorage() { + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:0] + } else { + n.data = nil + } } else { + if n.HasExternalStorage() { + words := size / 8 + if size%8 != 0 || words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:words] + copy(n.data, types.DecodeSlice[uint64](data[:size])) + return + } n.data = types.DecodeSlice[uint64](data[:size]) } } func (n *Bitmap) UnmarshalNoCopy(data []byte) { + if n.HasExternalStorage() { + panic("cannot install alias into bitmap external storage") + } n.count = types.DecodeInt64(data[:8]) data = data[8:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] @@ -447,14 +921,31 @@ func (n *Bitmap) UnmarshalNoCopy(data []byte) { // UnmarshalV1 in version 1, Bitmap.emptyFlag is type int32, now we use Bitmap.count replace it func (n *Bitmap) UnmarshalV1(data []byte) { data = data[4:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] if size == 0 { - n.data = nil + if n.HasExternalStorage() { + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:0] + } else { + n.data = nil + } } else { - n.data = types.DecodeSlice[uint64](data[:size]) + if n.HasExternalStorage() { + words := size / 8 + if size%8 != 0 || words > cap(n.data) { + panic("bitmap external storage capacity exceeded") + } + storage := n.data[:cap(n.data)] + clear(storage) + n.data = storage[:words] + copy(n.data, types.DecodeSlice[uint64](data[:size])) + } else { + n.data = types.DecodeSlice[uint64](data[:size]) + } } n.count = 0 for i := 0; i < len(n.data); i++ { @@ -463,8 +954,11 @@ func (n *Bitmap) UnmarshalV1(data []byte) { } func (n *Bitmap) UnmarshalNoCopyV1(data []byte) { + if n.HasExternalStorage() { + panic("cannot install alias into bitmap external storage") + } data = data[4:] - n.len = int64(types.DecodeUint64(data[:8])) + n.setLogicalLen(int64(types.DecodeUint64(data[:8]))) data = data[8:] size := int(types.DecodeUint64(data[:8])) data = data[8:] diff --git a/pkg/common/bitmap/bitmap_test.go b/pkg/common/bitmap/bitmap_test.go index 35bd38cb43b0d..131e764d258dc 100644 --- a/pkg/common/bitmap/bitmap_test.go +++ b/pkg/common/bitmap/bitmap_test.go @@ -15,8 +15,11 @@ package bitmap import ( + "bytes" "fmt" + "io" "testing" + "unsafe" "github.com/stretchr/testify/require" ) @@ -32,6 +35,12 @@ func newBm(n int) *Bitmap { return &bm } +type shortMarshalWriter struct{} + +func (shortMarshalWriter) Write(value []byte) (int, error) { + return len(value) - 1, nil +} + func TestNulls(t *testing.T) { np := newBm(Rows) np.AddRange(0, 0) @@ -68,7 +77,13 @@ func TestNulls(t *testing.T) { fmt.Printf("numbers: %v\n", np.Count()) nq := newBm(Rows) - nq.Unmarshal(np.Marshal()) + encoded := np.Marshal() + var streamed bytes.Buffer + require.NoError(t, np.MarshalTo(&streamed)) + require.Equal(t, encoded, streamed.Bytes()) + require.Equal(t, len(encoded), np.MarshalSize()) + require.ErrorIs(t, np.MarshalTo(shortMarshalWriter{}), io.ErrShortWrite) + nq.Unmarshal(encoded) require.Equal(t, np.ToArray(), nq.ToArray()) @@ -159,8 +174,8 @@ func TestBitmap_Compatibility(t *testing.T) { np.AddMany(rows) npV1 := &Bitmap{ - len: np.len, - data: np.data, + taggedLen: np.taggedLen, + data: np.data, } data := npV1.MarshalV1() @@ -218,3 +233,124 @@ func TestBitmap_And2(t *testing.T) { np.And(np2) require.Equal(t, 100, np.Count()) } + +func TestBitmapExternalStorageLifecycle(t *testing.T) { + require.Equal( + t, + 2*unsafe.Sizeof(int64(0))+unsafe.Sizeof([]uint64(nil)), + unsafe.Sizeof(Bitmap{}), + ) + storage := []uint64{^uint64(0), ^uint64(0)} + var value Bitmap + require.Nil(t, value.InstallExternalStorage(storage)) + require.True(t, value.HasExternalStorage()) + require.Equal(t, 2, value.ExternalStorageCapacity()) + + value.InitWithSize(65) + require.Equal(t, []uint64{0, 0}, storage) + value.Add(64) + require.True(t, value.Contains(64)) + + value.Reset() + require.True(t, value.HasExternalStorage()) + require.Equal(t, 2, value.ExternalStorageCapacity()) + require.Equal(t, []uint64{0, 0}, storage) + value.TryExpandWithSize(128) + value.Add(127) + require.True(t, value.Contains(127)) + require.Panics(t, func() { + value.TryExpandWithSize(129) + }) + + released := value.ReleaseExternalStorage() + require.Equal(t, storage, released) + require.False(t, value.HasExternalStorage()) + value.TryExpandWithSize(129) + value.Add(128) + require.True(t, value.Contains(128)) +} + +func TestBitmapRemapOrdered(t *testing.T) { + for _, external := range []bool{false, true} { + t.Run(fmt.Sprintf("external=%t", external), func(t *testing.T) { + var value Bitmap + if external { + value.InstallExternalStorage(make([]uint64, 3)) + } + value.InitWithSize(130) + value.AddMany([]uint64{0, 2, 63, 64, 65, 128, 129}) + + value.RemapOrdered([]int64{0, 2, 64, 65, 129}, false) + require.Equal(t, int64(5), value.Len()) + require.Equal(t, 5, value.Count()) + for row := uint64(0); row < 5; row++ { + require.True(t, value.Contains(row)) + } + + value.InitWithSize(130) + value.AddMany([]uint64{0, 2, 63, 64, 65, 128, 129}) + value.RemapOrdered([]int64{1, 63, 128}, true) + require.Equal(t, int64(130), value.Len()) + require.Equal(t, 5, value.Count()) + require.True(t, value.Contains(0)) + require.True(t, value.Contains(1)) + require.True(t, value.Contains(62)) + require.True(t, value.Contains(63)) + require.True(t, value.Contains(126)) + }) + } +} + +func TestBitmapRemapMaskOrdered(t *testing.T) { + var value Bitmap + value.InitWithSize(130) + value.AddMany([]uint64{1, 63, 64, 127, 129}) + + selection := newBm(130) + selection.AddMany([]uint64{1, 64, 129}) + value.RemapMaskOrdered(selection, false) + require.Equal(t, int64(3), value.Len()) + require.Equal(t, 3, value.Count()) + require.True(t, value.Contains(0)) + require.True(t, value.Contains(1)) + require.True(t, value.Contains(2)) + + value.InitWithSize(130) + value.AddMany([]uint64{1, 63, 64, 127, 129}) + value.RemapMaskOrdered(selection, true) + require.Equal(t, int64(130), value.Len()) + require.Equal(t, 2, value.Count()) + require.True(t, value.Contains(62)) + require.True(t, value.Contains(125)) +} + +func TestBitmapExternalStorageUnmarshal(t *testing.T) { + source := newBm(128) + source.Add(1) + source.Add(127) + encoded := source.Marshal() + + storage := make([]uint64, 2) + var copied Bitmap + copied.InstallExternalStorage(storage) + copied.Unmarshal(encoded) + require.True(t, copied.IsSame(source)) + require.True(t, copied.HasExternalStorage()) + + var streamed Bitmap + streamStorage := make([]uint64, 2) + streamed.InstallExternalStorage(streamStorage) + payload, err := streamed.PrepareExternalUnmarshal( + encoded[:MarshalHeaderSize], + len(encoded), + ) + require.NoError(t, err) + copy(payload, encoded[MarshalHeaderSize:]) + require.True(t, streamed.IsSame(source)) + + _, err = streamed.PrepareExternalUnmarshal( + encoded[:MarshalHeaderSize], + len(encoded)-1, + ) + require.Error(t, err) +} diff --git a/pkg/common/bitmap/cbitmap.go b/pkg/common/bitmap/cbitmap.go index 1edc7ffb8d595..2576718f2a2fb 100644 --- a/pkg/common/bitmap/cbitmap.go +++ b/pkg/common/bitmap/cbitmap.go @@ -27,7 +27,7 @@ func (n *Bitmap) cPtr() *C.uint64_t { return (*C.uint64_t)(unsafe.Pointer(&n.data[0])) } func (n *Bitmap) cLen() C.uint64_t { - return C.uint64_t(n.len) + return C.uint64_t(n.logicalLen()) } func (n *Bitmap) C_IsEmpty() bool { @@ -53,5 +53,5 @@ func (n *Bitmap) RawPtrLen() (uintptr, uintptr) { if n == nil || len(n.data) == 0 { return 0, 0 } - return uintptr(unsafe.Pointer(&n.data[0])), uintptr(n.len) + return uintptr(unsafe.Pointer(&n.data[0])), uintptr(n.logicalLen()) } diff --git a/pkg/common/bitmap/types.go b/pkg/common/bitmap/types.go index bf1141d9fc19d..0e12c750d66ad 100644 --- a/pkg/common/bitmap/types.go +++ b/pkg/common/bitmap/types.go @@ -22,8 +22,11 @@ type Iterator interface { type Bitmap struct { count int64 //in version 1, we use emptyFlag with type int32 to indicate whether it is empty - len int64 - data []uint64 + // taggedLen stores the logical length directly for Go-owned backing and its + // bitwise complement for caller-owned backing. This keeps Bitmap's footprint + // unchanged while making backing ownership explicit. + taggedLen int64 + data []uint64 } type BitmapIterator struct { diff --git a/pkg/common/hashmap/inthashmap.go b/pkg/common/hashmap/inthashmap.go index 22de287951b43..10e7d26fec42f 100644 --- a/pkg/common/hashmap/inthashmap.go +++ b/pkg/common/hashmap/inthashmap.go @@ -32,8 +32,16 @@ func init() { } func NewIntHashMap(hasNull bool, memPool *mpool.MPool) (*IntHashMap, error) { + return NewIntHashMapWithAllocation(hasNull, memPool, nil) +} + +func NewIntHashMapWithAllocation( + hasNull bool, + memPool *mpool.MPool, + allocation *hashtable.AllocationAccountSelection, +) (*IntHashMap, error) { mp := &hashtable.Int64HashMap{} - if err := mp.Init(memPool); err != nil { + if err := mp.InitWithAllocation(memPool, allocation); err != nil { return nil, err } return &IntHashMap{ diff --git a/pkg/common/hashmap/inthashmap_lazy_test.go b/pkg/common/hashmap/inthashmap_lazy_test.go index be4131421bbd6..665f964abd6c4 100644 --- a/pkg/common/hashmap/inthashmap_lazy_test.go +++ b/pkg/common/hashmap/inthashmap_lazy_test.go @@ -52,7 +52,8 @@ func TestIntHashMapIteratorLazyBuffers(t *testing.T) { insertedVs := append([]uint64(nil), vs...) insertedZvs := append([]int64(nil), zvs...) - foundVs, foundZvs := itr.Find(0, count, vecs) + foundVs, foundZvs, err := itr.Find(0, count, vecs) + require.NoError(t, err) if count > 0 { require.Equal(t, insertedVs, foundVs) require.Equal(t, insertedZvs, foundZvs) @@ -224,7 +225,10 @@ func BenchmarkIntHashMapFindFloat32(b *testing.B) { b.SetBytes(int64(count * types.T_float32.TypeLen())) b.ResetTimer() for i := 0; i < b.N; i++ { - benchmarkIntValues, benchmarkIntZValues = itr.Find(0, count, vecs) + benchmarkIntValues, benchmarkIntZValues, err = itr.Find(0, count, vecs) + if err != nil { + b.Fatal(err) + } } }) } diff --git a/pkg/common/hashmap/inthashmap_test.go b/pkg/common/hashmap/inthashmap_test.go index 8f28a8e63cba2..d5a746e1ae229 100644 --- a/pkg/common/hashmap/inthashmap_test.go +++ b/pkg/common/hashmap/inthashmap_test.go @@ -24,6 +24,48 @@ import ( "github.com/stretchr/testify/require" ) +func TestIntHashMapProbeGroupingDoesNotMatchRawKey(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewIntHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(0), false, mp)) + grouping := vector.NewRollupConst(types.T_uint8.ToType(), 1, mp) + defer raw.Free(mp) + defer grouping.Free(mp) + + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 1, []*vector.Vector{grouping}) + require.NoError(t, err) + require.Equal(t, []uint64{0}, values) + require.Equal(t, []int64{0}, zValues) +} + +func TestIntHashMapPartialGroupingRowsDoNotMatchRawKeys(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewIntHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + build := vector.NewVec(types.T_int32.ToType()) + probe := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList(build, []int32{7, 8}, nil, mp)) + require.NoError(t, vector.AppendFixedList(probe, []int32{7, 8}, nil, mp)) + probe.GetGrouping().Add(0) + defer build.Free(mp) + defer probe.Free(mp) + + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 2, []*vector.Vector{build}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 2, []*vector.Vector{probe}) + require.NoError(t, err) + require.Equal(t, []uint64{0, 2}, values) + require.Equal(t, []int64{0, 1}, zValues) +} + func TestIntHashMap_Iterator(t *testing.T) { { m := mpool.MustNewZero() @@ -42,7 +84,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, rowCount, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 1, 1, 2, 2, 2, 3, 3, 3, 4}, vs) - vs, _ = itr.Find(0, rowCount, vecs) + vs, _, err = itr.Find(0, rowCount, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 1, 1, 2, 2, 2, 3, 3, 3, 4}, vs) for _, vec := range vecs { vec.Free(m) @@ -63,7 +106,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -83,7 +127,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -103,7 +148,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -123,7 +169,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -143,7 +190,8 @@ func TestIntHashMap_Iterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{0, 1, 0, 2, 0, 3, 0, 4, 0, 5}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{0, 1, 0, 2, 0, 3, 0, 4, 0, 5}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -208,7 +256,8 @@ func TestIntHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, _ := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + foundVs, _, err := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + require.NoError(t, err) require.Equal(t, expectedMappedValue, foundVs) }) @@ -245,7 +294,8 @@ func TestIntHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, foundZvs := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + foundVs, foundZvs, err := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + require.NoError(t, err) for i := 0; i < numElements; i++ { require.Equal(t, originalVs[i], foundVs[i], "Mismatch at index %d for mapped value", i) require.Equal(t, originalZvs[i], foundZvs[i], "Mismatch at index %d for zValue", i) diff --git a/pkg/common/hashmap/iterator.go b/pkg/common/hashmap/iterator.go index 8b99d6405551f..cac9b7eb2e299 100644 --- a/pkg/common/hashmap/iterator.go +++ b/pkg/common/hashmap/iterator.go @@ -17,9 +17,47 @@ package hashmap import ( "unsafe" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/vector" ) +func validateIteratorVectors( + vecs []*vector.Vector, + start int, + count int, +) error { + if len(vecs) == 0 || start < 0 || count < 0 || count > UnitLimit { + return mpool.ErrAllocationAccountInvalid + } + for _, vec := range vecs { + if vec == nil || start > vec.Length() || count > vec.Length()-start { + return mpool.ErrAllocationAccountInvalid + } + } + return nil +} + +func hasGroupingInRange(vecs []*vector.Vector, start, count int) bool { + end := uint64(start + count) + for _, vec := range vecs { + if vec != nil && vec.GetGrouping().GetBitmap().CountRange( + uint64(start), end, + ) > 0 { + return true + } + } + return false +} + +func rowHasGrouping(vecs []*vector.Vector, row int) bool { + for _, vec := range vecs { + if vec.GetGrouping().Contains(uint64(row)) { + return true + } + } + return false +} + // MaxStrIteratorCapacity limits how many bytes of backing storage we keep when // reusing a string iterator. Avoids retaining oversized buffers after handling // very large strings. @@ -31,7 +69,12 @@ func IteratorChangeOwner(itr Iterator, m HashMap) { return } it := itr.(*strHashmapIterator) - it.mp = m.(*StrHashMap) + next := m.(*StrHashMap) + if it.mp != nil && + it.mp.iteratorAllocation != next.iteratorAllocation { + it.releaseScratch() + } + it.mp = next } // IteratorClearOwner detaches the iterator from its hashmap to allow the old @@ -41,6 +84,7 @@ func IteratorClearOwner(itr Iterator) { case *intHashMapIterator: it.mp = nil case *strHashmapIterator: + it.releaseAccountedScratch() it.mp = nil } } @@ -52,28 +96,58 @@ func StrIteratorCapacity(itr Iterator) int { if !ok || it == nil { return 0 } - total := 0 - for i := range it.keys { - total += cap(it.keys[i]) + return cap(it.keyBuffer) +} + +func (itr *strHashmapIterator) releaseScratch() { + if itr == nil { + return + } + if cap(itr.keyBuffer) > 0 && itr.mp != nil && itr.mp.mp != nil && + itr.mp.iteratorAllocation != nil { + itr.mp.mp.Free(itr.keyBuffer) } - return total + itr.keyBuffer = nil + clear(itr.keys) } -func (itr *strHashmapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64) { - for i := 0; i < count; i++ { - itr.keys[i] = itr.keys[i][:0] +func (itr *strHashmapIterator) releaseAccountedScratch() { + if itr == nil || cap(itr.keyBuffer) == 0 || itr.mp == nil || + itr.mp.iteratorAllocation == nil { + return + } + itr.mp.mp.Free(itr.keyBuffer) + itr.keyBuffer = nil + clear(itr.keys) +} + +func (itr *strHashmapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { + if err := itr.prepareHashKeys(vecs, start, count); err != nil { + return nil, nil, err } copy(itr.zValues[:count], OneInt64s[:count]) copy(itr.values[:count], zeroUint64[:count]) itr.encodeHashKeys(vecs, start, count) itr.mp.hashMap.FindStringBatch(itr.strHashStates, itr.keys[:count], itr.values) - return itr.values[:count], itr.zValues[:count] + if !itr.mp.hasNull && !itr.mp.groupingAware && + hasGroupingInRange(vecs, start, count) { + for i := 0; i < count; i++ { + if rowHasGrouping(vecs, start+i) { + itr.values[i] = 0 + itr.zValues[i] = 0 + } + } + } + return itr.values[:count], itr.zValues[:count], nil } // Insert a row from multiple columns into the hashmap, return true if it is new, otherwise false func (itr *strHashmapIterator) DetectDup(vecs []*vector.Vector, row int) (bool, error) { keys := itr.keys defer func() { keys[0] = keys[0][:0] }() + if err := itr.prepareHashKeys(vecs, row, 1); err != nil { + return false, err + } itr.encodeHashKeys(vecs, row, 1) if err := itr.mp.hashMap.InsertStringBatch(itr.strHashStates, keys[:1], itr.values[:1]); err != nil { return false, err @@ -88,6 +162,9 @@ func (itr *strHashmapIterator) DetectDup(vecs []*vector.Vector, row int) (bool, func (itr *strHashmapIterator) Insert(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { var err error + if err = itr.prepareHashKeys(vecs, start, count); err != nil { + return nil, nil, err + } defer func() { for i := 0; i < count; i++ { itr.keys[i] = itr.keys[i][:0] @@ -110,10 +187,16 @@ func (itr *strHashmapIterator) Insert(start, count int, vecs []*vector.Vector) ( return vs, zvs, err } -func (itr *intHashMapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64) { +func (itr *intHashMapIterator) Find(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { + if itr == nil || itr.mp == nil { + return nil, nil, mpool.ErrAllocationAccountInvalid + } + if err := validateIteratorVectors(vecs, start, count); err != nil { + return nil, nil, err + } itr.ensureCapacity(count) if count == 0 { - return itr.values, itr.zValues + return itr.values, itr.zValues, nil } for i := 0; i < count; i++ { itr.keys[i] = 0 @@ -124,15 +207,37 @@ func (itr *intHashMapIterator) Find(start, count int, vecs []*vector.Vector) ([] itr.encodeHashKeys(vecs, start, count) copy(itr.hashes[:count], zeroUint64[:count]) itr.mp.hashMap.FindBatch(count, itr.hashes[:count], unsafe.Pointer(&itr.keys[0]), itr.values[:count]) - return itr.values[:count], itr.zValues[:count] + if hasGroupingInRange(vecs, start, count) { + for i := 0; i < count; i++ { + if rowHasGrouping(vecs, start+i) { + itr.values[i] = 0 + itr.zValues[i] = 0 + } + } + } + return itr.values[:count], itr.zValues[:count], nil } func (itr *intHashMapIterator) DetectDup(vecs []*vector.Vector, row int) (bool, error) { - panic("not implemented yet!!!") + if itr == nil || itr.mp == nil { + return false, mpool.ErrAllocationAccountInvalid + } + before := itr.mp.rows + values, zValues, err := itr.Insert(row, 1, vecs) + if err != nil { + return false, err + } + return zValues[0] != 0 && values[0] > before, nil } func (itr *intHashMapIterator) Insert(start, count int, vecs []*vector.Vector) ([]uint64, []int64, error) { var err error + if itr == nil || itr.mp == nil { + return nil, nil, mpool.ErrAllocationAccountInvalid + } + if err = validateIteratorVectors(vecs, start, count); err != nil { + return nil, nil, err + } itr.ensureCapacity(count) if count == 0 { return itr.values, itr.zValues, nil diff --git a/pkg/common/hashmap/keycodec/keycodec.go b/pkg/common/hashmap/keycodec/keycodec.go index b259784429efc..0f72f997e3e83 100644 --- a/pkg/common/hashmap/keycodec/keycodec.go +++ b/pkg/common/hashmap/keycodec/keycodec.go @@ -25,6 +25,23 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" ) +// ValidVectors verifies the row-shape contract required by both resident hash +// maps and spill partitioning. Hashing a short or nil key must never silently +// leave a suffix at its previous seed value. +func ValidVectors(vecs []*vector.Vector, rows int) bool { + if rows < 0 || len(vecs) == 0 { + return false + } + for _, vec := range vecs { + if vec == nil || vec.Length() != rows { + return false + } + } + return true +} + +var groupingColumnHash = xxhash.Sum64([]byte{2}) + // Float32Codec holds the SQL comparison normalization for one FLOAT32 type. // Construct it once per vector so scale processing is not repeated per row. type Float32Codec struct { @@ -194,6 +211,10 @@ func ComputeXXHash(keyVecs []*vector.Vector, hashValues []uint64, seed uint64) { } for _, vec := range keyVecs { + if vec.GetGrouping().GetBitmap().CountRange(0, uint64(rowCount)) > 0 { + computeGroupingXXHash(vec, hashValues) + continue + } switch vec.GetType().Oid { case types.T_float32: computeFloat32XXHash(vec, hashValues) @@ -273,6 +294,41 @@ func computeFloat32XXHash(vec *vector.Vector, hashValues []uint64) { } } +func computeGroupingXXHash(vec *vector.Vector, hashValues []uint64) { + rowCount := len(hashValues) + grouping := vec.GetGrouping() + nulls := vec.GetNulls() + for i := 0; i < rowCount; i++ { + if grouping.Contains(uint64(i)) { + hashValues[i] = HashCombine(hashValues[i], groupingColumnHash) + continue + } + if vec.IsConstNull() || nulls.Contains(uint64(i)) { + hashValues[i] = HashCombine(hashValues[i], 0) + continue + } + row := i + if vec.IsConst() { + row = 0 + } + switch vec.GetType().Oid { + case types.T_float32: + values := vector.MustFixedColNoTypeCheck[float32](vec) + value := NewFloat32Codec(vec.GetType().Scale).CanonicalBytes(values[row]) + hashValues[i] = HashCombine(hashValues[i], xxhash.Sum64(value[:])) + continue + case types.T_float64: + values := vector.MustFixedColNoTypeCheck[float64](vec) + value := CanonicalFloat64Bytes(values[row]) + hashValues[i] = HashCombine(hashValues[i], xxhash.Sum64(value[:])) + continue + } + hashValues[i] = HashCombine( + hashValues[i], xxhash.Sum64(vec.GetRawBytesAt(row)), + ) + } +} + func computeFloat64XXHash(vec *vector.Vector, hashValues []uint64) { rowCount := len(hashValues) if vec.IsConst() { diff --git a/pkg/common/hashmap/keycodec/keycodec_test.go b/pkg/common/hashmap/keycodec/keycodec_test.go index 8fd2fa380c822..79cdcbf784196 100644 --- a/pkg/common/hashmap/keycodec/keycodec_test.go +++ b/pkg/common/hashmap/keycodec/keycodec_test.go @@ -178,3 +178,83 @@ func TestComputeXXHashCompositeScaledFloat32Contract(t *testing.T) { require.NotEqual(t, hashes[0], hashes[2], "the FLOAT32 codec must preserve prior column hash state") require.NotEqual(t, hashes[0], hashes[3], "a distinct canonical FLOAT32 value must change the composite hash") } + +func TestComputeXXHashCanonicalizesGroupingRows(t *testing.T) { + mp := mpool.MustNewZero() + defer func() { require.Zero(t, mp.CurrNB()) }() + + for _, test := range []struct { + name string + left any + right any + newVec func(any) *vector.Vector + }{ + { + name: "fixed", + left: []int64{11, 22, 33}, + right: []int64{101, 22, 303}, + newVec: func(values any) *vector.Vector { + vec := vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + vec, values.([]int64), nil, mp, + )) + return vec + }, + }, + { + name: "float64", + left: []float64{11, 22, 33}, + right: []float64{101, 22, 303}, + newVec: func(values any) *vector.Vector { + vec := vector.NewVec(types.T_float64.ToType()) + require.NoError(t, vector.AppendFixedList( + vec, values.([]float64), nil, mp, + )) + return vec + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + left := test.newVec(test.left) + right := test.newVec(test.right) + ordinary := test.newVec(test.left) + defer left.Free(mp) + defer right.Free(mp) + defer ordinary.Free(mp) + left.GetGrouping().AddRange(0, 3) + right.GetGrouping().AddRange(0, 3) + + leftHashes := make([]uint64, 3) + rightHashes := make([]uint64, 3) + ComputeXXHash([]*vector.Vector{left}, leftHashes, 17) + ComputeXXHash([]*vector.Vector{right}, rightHashes, 17) + + require.Equal(t, leftHashes, rightHashes) + + left.GetGrouping().Reset() + left.GetGrouping().Add(0) + ComputeXXHash([]*vector.Vector{left}, leftHashes, 17) + ComputeXXHash([]*vector.Vector{ordinary}, rightHashes, 17) + require.NotEqual(t, rightHashes[0], leftHashes[0]) + require.Equal(t, rightHashes[1:], leftHashes[1:]) + }) + } +} + +func TestComputeXXHashDoesNotTreatStaleGroupingAsFull(t *testing.T) { + mp := mpool.MustNewZero() + left := vector.NewVec(types.T_int64.ToType()) + right := vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(left, int64(11), false, mp)) + require.NoError(t, vector.AppendFixed(right, int64(22), false, mp)) + left.GetGrouping().Add(5) + right.GetGrouping().Add(5) + leftHash := []uint64{0} + rightHash := []uint64{0} + ComputeXXHash([]*vector.Vector{left}, leftHash, 17) + ComputeXXHash([]*vector.Vector{right}, rightHash, 17) + require.NotEqual(t, leftHash, rightHash) + left.Free(mp) + right.Free(mp) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/common/hashmap/strhashmap.go b/pkg/common/hashmap/strhashmap.go index 98addc3841f0f..25c33cbe4b9b1 100644 --- a/pkg/common/hashmap/strhashmap.go +++ b/pkg/common/hashmap/strhashmap.go @@ -38,13 +38,37 @@ func init() { } func NewStrHashMap(hasNull bool, memPool *mpool.MPool) (*StrHashMap, error) { + return NewStrHashMapWithAllocation(hasNull, memPool, nil) +} + +func NewStrHashMapWithAllocation( + hasNull bool, + memPool *mpool.MPool, + allocation *hashtable.AllocationAccountSelection, +) (*StrHashMap, error) { + return NewStrHashMapWithAllocations( + hasNull, + memPool, + allocation, + nil, + ) +} + +func NewStrHashMapWithAllocations( + hasNull bool, + memPool *mpool.MPool, + allocation *hashtable.AllocationAccountSelection, + iteratorAllocation *IteratorAllocation, +) (*StrHashMap, error) { mp := &hashtable.StringHashMap{} - if err := mp.Init(memPool); err != nil { + if err := mp.InitWithAllocation(memPool, allocation); err != nil { return nil, err } return &StrHashMap{ - hashMap: mp, - hasNull: hasNull, + hashMap: mp, + hasNull: hasNull, + mp: memPool, + iteratorAllocation: iteratorAllocation, }, nil } @@ -58,10 +82,185 @@ func (m *StrHashMap) NewIterator() Iterator { } } +func (itr *strHashmapIterator) prepareHashKeys( + vecs []*vector.Vector, + start int, + count int, +) error { + if itr == nil || itr.mp == nil || start < 0 || count < 0 || + count > UnitLimit { + return mpool.ErrAllocationAccountInvalid + } + if err := validateIteratorVectors(vecs, start, count); err != nil { + return err + } + for i := 0; i < count; i++ { + itr.keyLengths[i] = 0 + } + const maxInt = int(^uint(0) >> 1) + add := func(row int, size int) error { + if size < 0 || itr.keyLengths[row] > maxInt-size { + return mpool.ErrAllocationAccountInvalid + } + itr.keyLengths[row] += size + return nil + } + for _, vec := range vecs { + withDomain := itr.mp.hasNull || itr.mp.groupingAware + prefix := 0 + if withDomain { + prefix = 1 + } + if vec.IsGrouping() { + for i := 0; i < count; i++ { + if err := add(i, 1); err != nil { + return err + } + } + continue + } + if vec.IsConstNull() { + if itr.mp.hasNull { + for i := 0; i < count; i++ { + if err := add(i, 1); err != nil { + return err + } + } + } + continue + } + + // Most join keys are flat and non-null. Size them from the physical + // representation directly, avoiding repeated type/null/const dispatch + // before the encoder's required value pass. + hasGrouping := withDomain && vec.HasGrouping() + if !hasGrouping && !vec.GetNulls().Any() { + if vec.GetType().IsFixedLen() { + size := prefix + vec.GetType().TypeSize() + for i := 0; i < count; i++ { + if err := add(i, size); err != nil { + return err + } + } + continue + } + if vec.IsConst() { + size := prefix + 4 + len(vec.GetBytesAt(0)) + for i := 0; i < count; i++ { + if err := add(i, size); err != nil { + return err + } + } + continue + } + values, area := vector.MustVarlenaRawData(vec) + for i := 0; i < count; i++ { + value := values[start+i].GetByteSlice(area) + if err := add(i, prefix+4+len(value)); err != nil { + return err + } + } + continue + } + + fixed := vec.GetType().IsFixedLen() + for i := 0; i < count; i++ { + row := start + i + if withDomain && vec.GetGrouping().Contains(uint64(row)) { + if err := add(i, 1); err != nil { + return err + } + continue + } + if vec.GetNulls().Contains(uint64(row)) { + if itr.mp.hasNull { + if err := add(i, 1); err != nil { + return err + } + } + continue + } + if fixed { + if err := add(i, prefix+vec.GetType().TypeSize()); err != nil { + return err + } + continue + } + valueRow := row + if vec.IsConst() { + valueRow = 0 + } + if err := add(i, prefix+4+len(vec.GetBytesAt(valueRow))); err != nil { + return err + } + } + } + + total := 0 + for i := 0; i < count; i++ { + if itr.keyLengths[i] < 16 { + itr.keyLengths[i] = 16 + } + if total > maxInt-itr.keyLengths[i] { + return mpool.ErrAllocationAccountInvalid + } + total += itr.keyLengths[i] + } + if cap(itr.keyBuffer) < total { + if allocation := itr.mp.iteratorAllocation; allocation != nil { + var next []byte + var err error + if cap(itr.keyBuffer) > 0 { + // Grow owns the capacity policy. Passing a pre-grown capacity + // would apply that policy twice and falsely inflate admission. + next, err = itr.mp.mp.Grow(itr.keyBuffer, total, true) + } else { + capacity, ok := mpool.GrowCapacity(0, int64(total)) + if !ok || int64(int(capacity)) != capacity { + return mpool.ErrAllocationAllocatorLimit + } + next, err = itr.mp.mp.AllocAccounted( + int(capacity), + allocation.account, + allocation.owner, + allocation.site, + ) + } + if err != nil { + return err + } + itr.keyBuffer = next + } else { + itr.keyBuffer = make([]byte, total) + } + } + itr.keyBuffer = itr.keyBuffer[:total] + storage := itr.keyBuffer + offset := 0 + for i := 0; i < count; i++ { + end := offset + itr.keyLengths[i] + itr.keys[i] = storage[offset:offset:end] + offset = end + } + return nil +} + func (m *StrHashMap) HasNull() bool { return m.hasNull } +// SetGroupingAware selects a collision-free key domain for maps that may see +// GROUPING rows. It must be set before the +// first insert. Ordinary columns receive a 0 domain byte and GROUPING columns +// receive 2, so no raw fixed-width value can alias the sentinel. +func (m *StrHashMap) SetGroupingAware() error { + if m == nil || m.rows != 0 { + return mpool.ErrAllocationAccountInvalid + } + m.groupingAware = true + return nil +} + func (m *StrHashMap) Free() { m.hashMap.Free() } @@ -96,6 +295,10 @@ func (m *StrHashMap) Size() int64 { func (itr *strHashmapIterator) encodeHashKeys(vecs []*vector.Vector, start, count int) { for _, vec := range vecs { + if itr.mp.groupingAware || itr.mp.hasNull { + fillGroupingAwareStr(itr, vec, count, start) + continue + } if vec.GetType().IsFixedLen() { switch vec.GetType().Oid { case types.T_float32: @@ -137,7 +340,6 @@ func fillFloat32GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start i } return } - values := vector.MustFixedColNoTypeCheck[float32](vec) codec := keycodec.NewFloat32Codec(vec.GetType().Scale) if vec.IsConst() { @@ -190,6 +392,78 @@ func fillFloat32GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start i } } +func fillGroupingAwareStr( + itr *strHashmapIterator, + vec *vector.Vector, + n int, + start int, +) { + keys := itr.keys + if vec.IsGrouping() { + for i := 0; i < n; i++ { + keys[i] = append(keys[i], byte(2)) + } + return + } + if vec.IsConstNull() { + for i := 0; i < n; i++ { + row := start + i + if vec.GetGrouping().Contains(uint64(row)) { + keys[i] = append(keys[i], byte(2)) + } else if itr.mp.hasNull { + keys[i] = append(keys[i], byte(1)) + } else { + itr.zValues[i] = 0 + } + } + return + } + float32Codec := keycodec.NewFloat32Codec(vec.GetType().Scale) + for i := 0; i < n; i++ { + row := start + i + if vec.GetGrouping().Contains(uint64(row)) { + keys[i] = append(keys[i], byte(2)) + continue + } + if vec.GetNulls().Contains(uint64(row)) { + if itr.mp.hasNull { + keys[i] = append(keys[i], byte(1)) + } else { + itr.zValues[i] = 0 + } + continue + } + keys[i] = append(keys[i], byte(0)) + valueRow := row + if vec.IsConst() { + valueRow = 0 + } + switch vec.GetType().Oid { + case types.T_float32: + values := vector.MustFixedColNoTypeCheck[float32](vec) + value := float32Codec.CanonicalBytes(values[valueRow]) + keys[i] = append(keys[i], value[:]...) + continue + case types.T_float64: + values := vector.MustFixedColNoTypeCheck[float64](vec) + value := keycodec.CanonicalFloat64Bytes(values[valueRow]) + keys[i] = append(keys[i], value[:]...) + continue + } + if vec.GetType().IsFixedLen() { + size := vec.GetType().TypeSize() + data := vec.GetData() + value := data[valueRow*size : (valueRow+1)*size] + keys[i] = append(keys[i], value...) + continue + } + value := vec.GetBytesAt(valueRow) + length := uint32(len(value)) + keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) + keys[i] = append(keys[i], value...) + } +} + func fillFloat64GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start int) { keys := itr.keys if vec.IsGrouping() { @@ -265,7 +539,7 @@ func fillFloat64GroupStr(itr *strHashmapIterator, vec *vector.Vector, n, start i func fillStringGroupStrForConstVec(itr *strHashmapIterator, vec *vector.Vector, n int, start int) { keys := itr.keys bytes := vec.GetBytesAt(start) - length := uint16(len(bytes)) + length := uint32(len(bytes)) // can't be const null if itr.mp.hasNull { gsp := vec.GetGrouping() @@ -342,7 +616,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -359,7 +633,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -372,7 +646,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].ByteSlice() // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -382,7 +656,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].GetByteSlice(area) // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -408,7 +682,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -421,7 +695,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].ByteSlice() // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -442,7 +716,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s // this is not null value keys[i] = append(keys[i], 0) // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -455,7 +729,7 @@ func fillStringGroupStr(itr *strHashmapIterator, vec *vector.Vector, lenV int, s bytes := va[i+start].GetByteSlice(area) // for "a"，"bc" and "ab","c", we need to distinct // give the length - length := uint16(len(bytes)) + length := uint32(len(bytes)) keys[i] = append(keys[i], util.UnsafeToBytes(&length)...) // append the pure value bytes keys[i] = append(keys[i], bytes...) @@ -558,15 +832,17 @@ func (m *StrHashMap) UnmarshalBinary(data []byte, mp *mpool.MPool) error { func (m *StrHashMap) WriteTo(w io.Writer) (int64, error) { var n int64 - // Serialize hasNull (1 byte) + // The low two bits retain the key grammar. Historical payloads used only + // bit zero, so 0/1 remain backward-compatible. + flags := byte(0) if m.hasNull { - if _, err := w.Write([]byte{1}); err != nil { - return 0, err - } - } else { - if _, err := w.Write([]byte{0}); err != nil { - return 0, err - } + flags |= 1 + } + if m.groupingAware { + flags |= 2 + } + if _, err := w.Write([]byte{flags}); err != nil { + return 0, err } n++ @@ -598,7 +874,11 @@ func (m *StrHashMap) UnmarshalFrom(r io.Reader, mp *mpool.MPool) (int64, error) return 0, err } n += int64(rn) - m.hasNull = b[0] == 1 + if b[0]&^byte(3) != 0 { + return 0, mpool.ErrAllocationAccountInvalid + } + m.hasNull = b[0]&1 != 0 + m.groupingAware = b[0]&2 != 0 // Deserialize rows rowsData := make([]byte, 8) @@ -607,6 +887,7 @@ func (m *StrHashMap) UnmarshalFrom(r io.Reader, mp *mpool.MPool) (int64, error) } n += int64(rn) m.rows = types.DecodeUint64(rowsData) + m.mp = mp // Deserialize the underlying StringHashMap m.hashMap = &hashtable.StringHashMap{} diff --git a/pkg/common/hashmap/strhashmap_test.go b/pkg/common/hashmap/strhashmap_test.go index 1d9609893878b..3359c71152b65 100644 --- a/pkg/common/hashmap/strhashmap_test.go +++ b/pkg/common/hashmap/strhashmap_test.go @@ -141,7 +141,8 @@ func runIntHashMapFloat32CompositeFloatLastCase( require.NoError(t, err) require.Equal(t, []uint64{1}, values) require.Equal(t, []int64{1}, zValues) - values, zValues = itr.Find(0, 1, probeKeys) + values, zValues, err = itr.Find(0, 1, probeKeys) + require.NoError(t, err) require.Equal(t, []uint64{1}, values) require.Equal(t, []int64{1}, zValues) } @@ -236,6 +237,325 @@ func runFloatHashMapContract(t *testing.T, makeVectors floatHashMapVectorFactory } } +func TestStringHashMapCanonicalizesFullyGroupedKeys(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + require.NoError(t, hashMap.SetGroupingAware()) + defer hashMap.Free() + + build := vector.NewVec(types.T_int32.ToType()) + probe := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixed(build, int32(111), false, mp)) + require.NoError(t, vector.AppendFixed(probe, int32(222), false, mp)) + build.GetGrouping().Add(0) + probe.GetGrouping().Add(0) + defer build.Free(mp) + defer probe.Free(mp) + + iterator := hashMap.NewIterator() + values, zValues, err := iterator.Insert(0, 1, []*vector.Vector{build}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + values, zValues, err = iterator.Find(0, 1, []*vector.Vector{probe}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) +} + +func TestStringHashIteratorAccountedGrowthCapacityBoundary(t *testing.T) { + const ( + oldCapacity = 10_240 + required = oldCapacity + 1 + payloadSize = required - 4 + ) + newCapacity, ok := mpool.GrowCapacity(oldCapacity, required) + require.True(t, ok) + exactLimit := uint64(oldCapacity) + uint64(newCapacity) + + for _, testCase := range []struct { + name string + limit uint64 + wantError bool + }{ + {name: "exact-old-plus-rounded-new", limit: exactLimit}, + {name: "one-byte-short", limit: exactLimit - 1, wantError: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 2) + require.NoError(t, err) + account, err := registry.Open(testCase.limit) + require.NoError(t, err) + allocation, err := NewIteratorAllocation( + account, + mpool.AllocationOwnerMin, + mpool.AllocationSiteMin, + ) + require.NoError(t, err) + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMapWithAllocations( + false, + mp, + nil, + allocation, + ) + require.NoError(t, err) + iterator := hashMap.NewIterator().(*strHashmapIterator) + iterator.keyBuffer, err = mp.AllocAccounted( + oldCapacity, + account, + mpool.AllocationOwnerMin, + mpool.AllocationSiteMin, + ) + require.NoError(t, err) + iterator.keyBuffer = iterator.keyBuffer[:0] + vec := vector.NewVec(types.T_varchar.ToType()) + require.NoError(t, vector.AppendBytes( + vec, + make([]byte, payloadSize), + false, + mp, + )) + + err = iterator.prepareHashKeys([]*vector.Vector{vec}, 0, 1) + if testCase.wantError { + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, oldCapacity, cap(iterator.keyBuffer)) + require.Equal(t, uint64(oldCapacity), account.Snapshot().Used) + } else { + require.NoError(t, err) + require.Equal(t, int(newCapacity), cap(iterator.keyBuffer)) + require.Equal(t, uint64(newCapacity), account.Snapshot().Used) + require.Equal(t, exactLimit, account.Snapshot().Peak) + } + + iterator.releaseScratch() + hashMap.Free() + vec.Free(mp) + require.Zero(t, mp.Stats().NumCurrBytes.Load()) + require.Zero(t, account.Seal().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + }) + } +} + +func TestGroupingAwareStringHashMapSeparatesRawSentinelBytes(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + require.NoError(t, hashMap.SetGroupingAware()) + defer hashMap.Free() + + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(2), false, mp)) + grouping := vector.NewRollupConst(types.T_uint8.ToType(), 1, mp) + defer raw.Free(mp) + defer grouping.Free(mp) + + iterator := hashMap.NewIterator() + values, zValues, err := iterator.Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + values, zValues, err = iterator.Insert(0, 1, []*vector.Vector{grouping}) + require.NoError(t, err) + require.Equal(t, []uint64{2}, values) + require.Equal(t, []int64{1}, zValues) + require.Equal(t, uint64(2), hashMap.GroupCount()) +} + +func TestNullableStringHashMapTreatsGroupingRowsAsSentinel(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(true, mp) + require.NoError(t, err) + defer hashMap.Free() + + partial := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList( + partial, + []int32{7, 7, 9}, + nil, + mp, + )) + partial.GetGrouping().Add(0) + partial.GetGrouping().Add(2) + partial.GetNulls().Add(2) + defer partial.Free(mp) + + iterator := hashMap.NewIterator() + values, zValues, err := iterator.Insert(0, 3, []*vector.Vector{partial}) + require.NoError(t, err) + require.Equal(t, []uint64{1, 2, 1}, values) + require.Equal(t, []int64{1, 1, 1}, zValues) + + fullGrouping := vector.NewRollupConst(types.T_int32.ToType(), 1, mp) + defer fullGrouping.Free(mp) + values, zValues, err = iterator.Insert(0, 1, []*vector.Vector{fullGrouping}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + require.Equal(t, uint64(2), hashMap.GroupCount()) +} + +func TestGroupingAwareStringHashMapConstNullUsesRowwiseGrouping(t *testing.T) { + mp := mpool.MustNewZero() + for _, hasNull := range []bool{false, true} { + hashMap, err := NewStrHashMap(hasNull, mp) + require.NoError(t, err) + require.NoError(t, hashMap.SetGroupingAware()) + vec := vector.NewConstNull(types.T_int32.ToType(), 2, mp) + vec.GetGrouping().Add(0) + + values, zValues, err := hashMap.NewIterator().Insert( + 0, 2, []*vector.Vector{vec}, + ) + require.NoError(t, err) + if hasNull { + require.Equal(t, []uint64{1, 2}, values) + require.Equal(t, []int64{1, 1}, zValues) + } else { + require.Equal(t, []uint64{1, 0}, values) + require.Equal(t, []int64{1, 0}, zValues) + } + + vec.Free(mp) + hashMap.Free() + } + require.Zero(t, mp.CurrNB()) +} + +func TestStringHashMapProbeGroupingDoesNotMatchRawKey(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(2), false, mp)) + grouping := vector.NewRollupConst(types.T_uint8.ToType(), 1, mp) + defer raw.Free(mp) + defer grouping.Free(mp) + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 1, []*vector.Vector{grouping}) + require.NoError(t, err) + require.Equal(t, []uint64{0}, values) + require.Equal(t, []int64{0}, zValues) +} + +func TestGroupingAwareStringHashMapRoundTripRetainsKeyGrammar(t *testing.T) { + mp := mpool.MustNewZero() + original, err := NewStrHashMap(false, mp) + require.NoError(t, err) + require.NoError(t, original.SetGroupingAware()) + raw := vector.NewVec(types.T_uint8.ToType()) + require.NoError(t, vector.AppendFixed(raw, uint8(2), false, mp)) + defer raw.Free(mp) + _, _, err = original.NewIterator().Insert(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + require.ErrorIs(t, original.SetGroupingAware(), mpool.ErrAllocationAccountInvalid) + + encoded, err := original.MarshalBinary() + require.NoError(t, err) + original.Free() + restored := &StrHashMap{} + require.NoError(t, restored.UnmarshalBinary(encoded, mp)) + values, zValues, err := restored.NewIterator().Find(0, 1, []*vector.Vector{raw}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + restored.Free() +} + +func TestStringHashMapWideVarlenaLengthsDoNotCollide(t *testing.T) { + mp := mpool.MustNewZero() + hashMap, err := NewStrHashMap(false, mp) + require.NoError(t, err) + defer hashMap.Free() + + large := make([]byte, 1<<16) + left := vector.NewVec(types.T_binary.ToType()) + right := vector.NewVec(types.T_binary.ToType()) + require.NoError(t, vector.AppendBytesList(left, [][]byte{large, nil}, nil, mp)) + require.NoError(t, vector.AppendBytesList(right, [][]byte{nil, large}, nil, mp)) + defer left.Free(mp) + defer right.Free(mp) + + values, zValues, err := hashMap.NewIterator().Insert( + 0, 2, []*vector.Vector{left, right}, + ) + require.NoError(t, err) + require.Equal(t, []uint64{1, 2}, values) + require.Equal(t, []int64{1, 1}, zValues) +} + +func TestStringHashMapConstAndFlatVarlenaUseSameKey(t *testing.T) { + mp := mpool.MustNewZero() + for _, buildConst := range []bool{false, true} { + hashMap, err := NewStrHashMap(true, mp) + require.NoError(t, err) + build := vector.NewVec(types.T_varchar.ToType()) + probe := vector.NewVec(types.T_varchar.ToType()) + if buildConst { + build, err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("abc"), 1, mp) + require.NoError(t, err) + require.NoError(t, vector.AppendBytes(probe, []byte("abc"), false, mp)) + } else { + require.NoError(t, vector.AppendBytes(build, []byte("abc"), false, mp)) + probe, err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("abc"), 1, mp) + require.NoError(t, err) + } + iterator := hashMap.NewIterator() + _, _, err = iterator.Insert(0, 1, []*vector.Vector{build}) + require.NoError(t, err) + values, zValues, err := iterator.Find(0, 1, []*vector.Vector{probe}) + require.NoError(t, err) + require.Equal(t, []uint64{1}, values) + require.Equal(t, []int64{1}, zValues) + build.Free(mp) + probe.Free(mp) + hashMap.Free() + } + require.Zero(t, mp.CurrNB()) +} + +func TestHashMapIteratorsRejectMalformedRowShapes(t *testing.T) { + mp := mpool.MustNewZero() + for _, makeIterator := range []func() (HashMap, Iterator){ + func() (HashMap, Iterator) { + m, err := NewIntHashMap(false, mp) + require.NoError(t, err) + return m, m.NewIterator() + }, + func() (HashMap, Iterator) { + m, err := NewStrHashMap(false, mp) + require.NoError(t, err) + return m, m.NewIterator() + }, + } { + m, iterator := makeIterator() + short, err := vector.NewConstFixed(types.T_int32.ToType(), int32(1), 1, mp) + require.NoError(t, err) + for _, vecs := range [][]*vector.Vector{ + nil, + {nil}, + {short}, + } { + _, _, err = iterator.Insert(0, 2, vecs) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + _, _, err = iterator.Find(1, 1, vecs) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + } + require.Zero(t, m.GroupCount()) + short.Free(mp) + m.Free() + } + require.Zero(t, mp.CurrNB()) +} + func runFloatHashMapShape( t *testing.T, composite bool, @@ -280,7 +600,8 @@ func runFloatHashMapShape( require.Equal(t, shape.wantValues, liveHashMapValues(values, zValues)) require.Equal(t, shape.wantGroups, hashMap.GroupCount()) - values, zValues = itr.Find(shape.start, shape.count, probe) + values, zValues, err = itr.Find(shape.start, shape.count, probe) + require.NoError(t, err) require.Equal(t, shape.wantZValues, zValues) require.Equal(t, shape.wantValues, liveHashMapValues(values, zValues)) require.Equal(t, shape.wantGroups, hashMap.GroupCount()) @@ -454,7 +775,8 @@ func TestIterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -479,7 +801,8 @@ func TestIterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -504,7 +827,8 @@ func TestIterator(t *testing.T) { vs, _, err := itr.Insert(0, Rows, vecs) require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) - vs, _ = itr.Find(0, Rows, vecs) + vs, _, err = itr.Find(0, Rows, vecs) + require.NoError(t, err) require.Equal(t, []uint64{1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, vs[:Rows]) for _, vec := range vecs { vec.Free(m) @@ -835,7 +1159,8 @@ func TestStrHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, _ := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + foundVs, _, err := unmarshaledMp.NewIterator().Find(0, rowCount, vecs) + require.NoError(t, err) require.Equal(t, expectedMappedValue, foundVs) }) @@ -872,7 +1197,8 @@ func TestStrHashMap_MarshalUnmarshal(t *testing.T) { require.Equal(t, expectedGroupCount, unmarshaledMp.GroupCount()) require.Equal(t, mp.HasNull(), unmarshaledMp.HasNull()) - foundVs, foundZvs := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + foundVs, foundZvs, err := unmarshaledMp.NewIterator().Find(0, numElements, vecs) + require.NoError(t, err) for i := 0; i < numElements; i++ { require.Equal(t, originalVs[i], foundVs[i], "Mismatch at index %d for mapped value", i) require.Equal(t, originalZvs[i], foundZvs[i], "Mismatch at index %d for zValue", i) diff --git a/pkg/common/hashmap/types.go b/pkg/common/hashmap/types.go index f84800bb74237..190eaaa17dff7 100644 --- a/pkg/common/hashmap/types.go +++ b/pkg/common/hashmap/types.go @@ -77,16 +77,46 @@ type Iterator interface { // Find vecs[start, start+count) in hashmap // vs : the number of rows corresponding to each value in the hash table (start with 1, and 0 means not found.) // zvs : if zvs[i] is 0 indicates the presence null, 1 indicates the absence of a null. - Find(start, count int, vecs []*vector.Vector) (vs []uint64, zvs []int64) + Find(start, count int, vecs []*vector.Vector) (vs []uint64, zvs []int64, err error) +} + +// IteratorAllocation selects exact physical provenance for data-scaled hash +// key encoding scratch. It is immutable and shared by iterators created from +// one map generation. +type IteratorAllocation struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + site mpool.AllocationSite +} + +func NewIteratorAllocation( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) (*IteratorAllocation, error) { + allocation := &IteratorAllocation{ + account: account, + owner: owner, + site: site, + } + if account == nil || account.Handle() == 0 || + owner < mpool.AllocationOwnerMin || owner > mpool.AllocationOwnerMax || + site < mpool.AllocationSiteMin { + return nil, mpool.ErrAllocationAccountInvalid + } + return allocation, nil } // StrHashMap key is []byte, value is an uint64 value (starting from 1) // // each time a new key is inserted, the hashtable returns a last-value+1 or, if the old key is inserted, the value corresponding to that key type StrHashMap struct { - hasNull bool - rows uint64 - hashMap *hashtable.StringHashMap + hasNull bool + groupingAware bool + rows uint64 + hashMap *hashtable.StringHashMap + mp *mpool.MPool + iteratorAllocation *IteratorAllocation } // IntHashMap key is int64, value is an uint64 (start from 1) @@ -99,9 +129,11 @@ type IntHashMap struct { } type strHashmapIterator struct { - mp *StrHashMap - keys [][]byte - values []uint64 + mp *StrHashMap + keys [][]byte + values []uint64 + keyBuffer []byte + keyLengths [UnitLimit]int // zValues, 0 indicates the presence null, 1 indicates the absence of a null zValues []int64 strHashStates [][3]uint64 diff --git a/pkg/common/malloc/profiler.go b/pkg/common/malloc/profiler.go index 148d599ee35b4..fad3f9d5123c7 100644 --- a/pkg/common/malloc/profiler.go +++ b/pkg/common/malloc/profiler.go @@ -103,6 +103,25 @@ func (p *Profiler[T, P]) Sample( return p.getSampleValueFromPCs(pcs, int64(fullStackFraction)) } +// SampleNamed returns a stable synthetic sample without collecting a runtime +// stack. It is intended for allocations that already carry explicit, bounded +// provenance supplied by their owner. +func (p *Profiler[T, P]) SampleNamed(name string) P { + locations := []*profile.Location{p.getMockLocation(name)} + locationsHashSum := hashLocations(locations) + if v, ok := p.locationsToSample.Load(locationsHashSum); ok { + return v.(*SampleInfo[P]).Values + } + + var value T + P(&value).Init() + v, _ := p.locationsToSample.LoadOrStore(locationsHashSum, &SampleInfo[P]{ + Values: &value, + Locations: locations, + }) + return v.(*SampleInfo[P]).Values +} + func (p *Profiler[T, P]) getLocation(frame runtime.Frame) *profile.Location { locationKey := LocationKey{ File: frame.File, diff --git a/pkg/common/moerr/error.go b/pkg/common/moerr/error.go index 06358bde0d09d..90b3c50199834 100644 --- a/pkg/common/moerr/error.go +++ b/pkg/common/moerr/error.go @@ -65,6 +65,7 @@ const ( ErrQueryInterrupted uint16 = 20104 ErrNotSupported uint16 = 20105 ErrRemoteDispatchNotRegistered uint16 = 20106 + ErrMPoolCapacity uint16 = 20107 // Group 2: numeric and functions ErrDivByZero uint16 = 20200 @@ -393,6 +394,7 @@ var errorMsgRefer = map[uint16]moErrorMsgItem{ ErrQueryInterrupted: {ER_QUERY_INTERRUPTED, []string{MySQLDefaultSqlState}, "query interrupted"}, ErrNotSupported: {ER_UNKNOWN_ERROR, []string{MySQLDefaultSqlState}, "not supported: %s"}, ErrRemoteDispatchNotRegistered: {ER_UNKNOWN_ERROR, []string{MySQLDefaultSqlState}, "remote dispatch receiver %s is not registered yet"}, + ErrMPoolCapacity: {ER_ENGINE_OUT_OF_MEMORY, []string{MySQLDefaultSqlState}, "mpool physical capacity exceeded: %s"}, // Group 2: numeric ErrDivByZero: {ER_DIVISION_BY_ZERO, []string{MySQLDefaultSqlState}, "division by zero"}, @@ -953,6 +955,13 @@ func NewOOM(ctx context.Context) *Error { return newError(ctx, ErrOOM) } +// NewMPoolCapacity reports a physical allocator or MPool capacity failure. +// Its dedicated wire code lets pressure recovery distinguish retryable +// physical capacity from unrelated OOMs without wrapping the MO error. +func NewMPoolCapacity(ctx context.Context, msg string) *Error { + return newError(ctx, ErrMPoolCapacity, msg) +} + // NewResourceExhaustedf preserves the existing resource-exhaustion wire code // while adding bounded, actionable context for guards that reject before the // allocator or operating system itself fails. The formatted message is diff --git a/pkg/common/moerr/error_no_ctx.go b/pkg/common/moerr/error_no_ctx.go index 750ab8b2e6d76..68b07dc7637b4 100644 --- a/pkg/common/moerr/error_no_ctx.go +++ b/pkg/common/moerr/error_no_ctx.go @@ -57,6 +57,10 @@ func NewOOMNoCtx() *Error { return newError(Context(), ErrOOM) } +func NewMPoolCapacityNoCtxf(format string, args ...any) *Error { + return NewMPoolCapacity(Context(), fmt.Sprintf(format, args...)) +} + func NewDivByZeroNoCtx() *Error { return newError(Context(), ErrDivByZero) } diff --git a/pkg/common/moerr/error_test.go b/pkg/common/moerr/error_test.go index c73e5c8674bad..17950ad04d8d6 100644 --- a/pkg/common/moerr/error_test.go +++ b/pkg/common/moerr/error_test.go @@ -181,6 +181,19 @@ func TestResourceExhaustedWithDetailsEncoding(t *testing.T) { require.Equal(t, err, decoded) } +func TestMPoolCapacityEncoding(t *testing.T) { + err := NewMPoolCapacityNoCtxf("alloc %d bytes, cap %d", 8, 4) + require.Equal(t, ErrMPoolCapacity, err.ErrorCode()) + require.Equal(t, ER_ENGINE_OUT_OF_MEMORY, err.MySQLCode()) + require.Contains(t, err.Error(), "alloc 8 bytes, cap 4") + + data, marshalErr := err.MarshalBinary() + require.NoError(t, marshalErr) + decoded := new(Error) + require.NoError(t, decoded.UnmarshalBinary(data)) + require.Equal(t, err, decoded) +} + func TestErrSubqueryNo1RowContract(t *testing.T) { err := NewErrSubqueryNo1Row(context.Background()) require.Equal(t, ErrSubqueryNo1Row, err.ErrorCode()) diff --git a/pkg/common/mpool/accounted_buffer.go b/pkg/common/mpool/accounted_buffer.go new file mode 100644 index 0000000000000..6eea972bae8b7 --- /dev/null +++ b/pkg/common/mpool/accounted_buffer.go @@ -0,0 +1,266 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package mpool + +import ( + "encoding/binary" + "math" +) + +// AccountedBuffer is a non-copyable, allocation-accounted off-heap byte +// buffer. Reset retains physical capacity; Free is its terminal release. +type AccountedBuffer struct { + data []byte + mp *MPool + + account *AllocationAccount + owner AllocationOwner + site AllocationSite + capacityClass AllocationCapacityClass +} + +func NewAccountedBuffer( + mp *MPool, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, +) (*AccountedBuffer, error) { + return NewAccountedBufferWithCapacityClass( + mp, + account, + owner, + site, + AllocationCapacityClassDefault, + ) +} + +// NewAccountedBufferWithCapacityClass creates a reusable buffer whose physical +// allocations borrow the selected execution-local capacity. +func NewAccountedBufferWithCapacityClass( + mp *MPool, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, + capacityClass AllocationCapacityClass, +) (*AccountedBuffer, error) { + if mp == nil { + return nil, ErrAllocationAccountInvalid + } + request := allocationAccountRequest{ + account: account, + owner: owner, + site: site, + capacityClass: capacityClass, + } + if err := request.validate(); err != nil { + return nil, err + } + return &AccountedBuffer{ + mp: mp, + account: account, + owner: owner, + site: site, + capacityClass: capacityClass, + }, nil +} + +func (b *AccountedBuffer) Bytes() []byte { + if b == nil { + return nil + } + return b.data +} + +func (b *AccountedBuffer) Len() int { + if b == nil { + return 0 + } + return len(b.data) +} + +func (b *AccountedBuffer) Cap() int { + if b == nil { + return 0 + } + return cap(b.data) +} + +// EnsureCapacity admits and allocates an absolute retained capacity. +func (b *AccountedBuffer) EnsureCapacity(required int) error { + if b == nil || b.mp == nil || b.account == nil || required < 0 { + return ErrAllocationAccountInvalid + } + if required <= cap(b.data) { + return nil + } + if int64(required) > maxAllocationSize() { + return ErrAllocationAllocatorLimit + } + + oldLength := len(b.data) + if cap(b.data) == 0 { + capacity, ok := GrowCapacity(0, int64(required)) + if !ok || capacity > int64(math.MaxInt) { + return ErrAllocationAllocatorLimit + } + data, err := b.mp.AllocAccountedWithCapacityClass( + int(capacity), + b.account, + b.owner, + b.site, + b.capacityClass, + ) + if err != nil { + return err + } + b.data = data[:oldLength] + return nil + } + + // Grow owns the capacity policy. Pass the caller's requirement instead of + // applying GrowCapacity a second time to the already rounded capacity. + data, err := b.mp.Grow(b.data, required, true) + if err != nil { + return err + } + b.data = data[:oldLength] + return nil +} + +// Resize changes the logical length after admitting any required retained +// capacity. Existing bytes are preserved. +func (b *AccountedBuffer) Resize(length int) error { + if b == nil || length < 0 { + return ErrAllocationAccountInvalid + } + if err := b.EnsureCapacity(length); err != nil { + return err + } + b.data = b.data[:length] + return nil +} + +func (b *AccountedBuffer) Write(value []byte) (int, error) { + if b == nil { + return 0, ErrAllocationAccountInvalid + } + if len(value) > math.MaxInt-len(b.data) { + return 0, ErrAllocationAccountInvalid + } + oldLength := len(b.data) + required := oldLength + len(value) + if err := b.EnsureCapacity(required); err != nil { + return 0, err + } + b.data = b.data[:required] + copy(b.data[oldLength:], value) + return len(value), nil +} + +func (b *AccountedBuffer) WriteString(value string) (int, error) { + if b == nil || len(value) > math.MaxInt-len(b.data) { + return 0, ErrAllocationAccountInvalid + } + oldLength := len(b.data) + required := oldLength + len(value) + if err := b.EnsureCapacity(required); err != nil { + return 0, err + } + b.data = b.data[:required] + copy(b.data[oldLength:], value) + return len(value), nil +} + +func (b *AccountedBuffer) appendSpace(length int) ([]byte, error) { + if b == nil || length < 0 || length > math.MaxInt-len(b.data) { + return nil, ErrAllocationAccountInvalid + } + start := len(b.data) + if err := b.Resize(start + length); err != nil { + return nil, err + } + return b.data[start:], nil +} + +func (b *AccountedBuffer) WriteByte(value byte) error { + dst, err := b.appendSpace(1) + if err != nil { + return err + } + dst[0] = value + return nil +} + +func (b *AccountedBuffer) WriteUint32(value uint32) error { + dst, err := b.appendSpace(4) + if err != nil { + return err + } + binary.NativeEndian.PutUint32(dst, value) + return nil +} + +func (b *AccountedBuffer) WriteInt32(value int32) error { + return b.WriteUint32(uint32(value)) +} + +func (b *AccountedBuffer) WriteUint64(value uint64) error { + dst, err := b.appendSpace(8) + if err != nil { + return err + } + binary.NativeEndian.PutUint64(dst, value) + return nil +} + +func (b *AccountedBuffer) WriteInt64(value int64) error { + return b.WriteUint64(uint64(value)) +} + +func (b *AccountedBuffer) SetUint32(offset int, value uint32) error { + if b == nil || offset < 0 || offset > len(b.data)-4 { + return ErrAllocationAccountInvalid + } + binary.NativeEndian.PutUint32(b.data[offset:offset+4], value) + return nil +} + +func (b *AccountedBuffer) SetInt64(offset int, value int64) error { + if b == nil || offset < 0 || offset > len(b.data)-8 { + return ErrAllocationAccountInvalid + } + binary.NativeEndian.PutUint64(b.data[offset:offset+8], uint64(value)) + return nil +} + +func (b *AccountedBuffer) Reset() { + if b != nil { + b.data = b.data[:0] + } +} + +func (b *AccountedBuffer) Free() { + if b == nil { + return + } + if cap(b.data) > 0 { + b.mp.Free(b.data) + } + b.data = nil + b.mp = nil + b.account = nil + b.owner = 0 + b.site = 0 +} diff --git a/pkg/common/mpool/accounted_buffer_test.go b/pkg/common/mpool/accounted_buffer_test.go new file mode 100644 index 0000000000000..c70765a9f48f5 --- /dev/null +++ b/pkg/common/mpool/accounted_buffer_test.go @@ -0,0 +1,181 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package mpool + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAccountedBufferLifecycle(t *testing.T) { + registry, account := newTestAllocationAccount(t, 1<<20, 4) + mp := MustNew("accounted-buffer") + defer DeleteMPool(mp) + buffer, err := NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + + require.NoError(t, buffer.EnsureCapacity(32)) + firstCapacity := buffer.Cap() + require.GreaterOrEqual(t, firstCapacity, 32) + require.Equal(t, uint64(firstCapacity), account.Snapshot().Used) + require.NoError(t, buffer.Resize(16)) + require.Equal(t, 16, buffer.Len()) + require.Equal(t, firstCapacity, buffer.Cap()) + copy(buffer.Bytes(), "retained") + require.NoError(t, buffer.Resize(8)) + require.Equal(t, "retained", string(buffer.Bytes())) + buffer.Reset() + _, err = buffer.WriteString("accounted") + require.NoError(t, err) + require.Equal(t, "accounted", string(buffer.Bytes())) + + buffer.Reset() + require.Zero(t, buffer.Len()) + require.Equal(t, firstCapacity, buffer.Cap()) + _, err = buffer.Write([]byte("reuse")) + require.NoError(t, err) + require.Equal(t, uint64(firstCapacity), account.Snapshot().Used) + + require.NoError(t, buffer.EnsureCapacity(firstCapacity+1)) + secondCapacity := buffer.Cap() + require.Greater(t, secondCapacity, firstCapacity) + snapshot := account.Snapshot() + require.Equal(t, uint64(secondCapacity), snapshot.Used) + require.GreaterOrEqual( + t, + snapshot.Peak, + uint64(firstCapacity+secondCapacity), + ) + require.Equal(t, "reuse", string(buffer.Bytes())) + + buffer.Free() + buffer.Free() + require.Zero(t, account.Snapshot().Used) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestAccountedBufferFailureRetainsPublishedData(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 2) + mp := MustNew("accounted-buffer-failure") + defer DeleteMPool(mp) + buffer, err := NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + _, err = buffer.Write([]byte("published")) + require.NoError(t, err) + before := append([]byte(nil), buffer.Bytes()...) + snapshot := account.Snapshot() + used := snapshot.Used + + err = buffer.EnsureCapacity(usedSizeToInt(t, snapshot.Limit)) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, before, buffer.Bytes()) + require.Equal(t, used, account.Snapshot().Used) + + buffer.Free() + finalizeTestAllocationAccount(t, registry, account) +} + +func TestAccountedBufferGrowthCapacityBoundary(t *testing.T) { + const ( + oldCapacity = 10_240 + required = oldCapacity + 1 + ) + newCapacity, ok := GrowCapacity(oldCapacity, required) + require.True(t, ok) + exactLimit := uint64(oldCapacity) + uint64(newCapacity) + + for _, testCase := range []struct { + name string + limit uint64 + wantError bool + }{ + {name: "exact-old-plus-rounded-new", limit: exactLimit}, + {name: "one-byte-short", limit: exactLimit - 1, wantError: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + registry, account := newTestAllocationAccount(t, testCase.limit, 2) + mp := MustNew("accounted-buffer-growth-boundary") + defer DeleteMPool(mp) + buffer, err := NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.NoError(t, buffer.EnsureCapacity(oldCapacity)) + + err = buffer.EnsureCapacity(required) + if testCase.wantError { + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, oldCapacity, buffer.Cap()) + require.Equal(t, uint64(oldCapacity), account.Snapshot().Used) + } else { + require.NoError(t, err) + require.Equal(t, int(newCapacity), buffer.Cap()) + require.Equal(t, uint64(newCapacity), account.Snapshot().Used) + require.Equal(t, exactLimit, account.Snapshot().Peak) + } + + buffer.Free() + finalizeTestAllocationAccount(t, registry, account) + }) + } +} + +func TestAccountedBufferConfiguration(t *testing.T) { + _, err := NewAccountedBuffer(nil, nil, 0, 0) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + var buffer *AccountedBuffer + require.Nil(t, buffer.Bytes()) + require.Zero(t, buffer.Len()) + require.Zero(t, buffer.Cap()) + require.ErrorIs(t, buffer.Resize(0), ErrAllocationAccountInvalid) + _, err = buffer.Write([]byte("x")) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + buffer.Reset() + buffer.Free() + + registry, account := newTestAllocationAccount(t, 1<<20, 1) + mp := MustNew("accounted-buffer-invalid-resize") + defer DeleteMPool(mp) + buffer, err = NewAccountedBuffer( + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.ErrorIs(t, buffer.Resize(-1), ErrAllocationAccountInvalid) + buffer.Free() + finalizeTestAllocationAccount(t, registry, account) +} + +func usedSizeToInt(t testing.TB, value uint64) int { + t.Helper() + require.LessOrEqual(t, value, uint64(^uint(0)>>1)) + return int(value) +} diff --git a/pkg/common/mpool/allocation_account.go b/pkg/common/mpool/allocation_account.go new file mode 100644 index 0000000000000..bf92475552987 --- /dev/null +++ b/pkg/common/mpool/allocation_account.go @@ -0,0 +1,986 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package mpool + +import ( + "errors" + "fmt" + "math" + "runtime" + "sync" + "sync/atomic" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" +) + +// AllocationOwner and AllocationSite are bounded diagnostic dimensions. +// Callers assign stable values in their own allocation-site ledger. Zero is +// reserved so an accounted allocation can never be published without an +// explicit owner and site. +type AllocationOwner uint8 +type AllocationSite uint8 + +const ( + AllocationOwnerMin AllocationOwner = 1 + AllocationOwnerMax AllocationOwner = 63 + AllocationSiteMin AllocationSite = 1 + AllocationSiteMax AllocationSite = math.MaxUint8 +) + +var ( + ErrAllocationAccountCapacity error = allocationAccountSentinel("allocation account capacity exceeded") + ErrAllocationAccountSealed error = allocationAccountSentinel("allocation account is sealed") + ErrAllocationAccountInvalid error = allocationAccountSentinel("invalid allocation account") + ErrAllocationAccountStale error = allocationAccountSentinel("stale allocation account handle") + ErrAllocationAccountMismatch error = allocationAccountSentinel("allocation account ownership mismatch") + ErrAllocationAllocatorLimit error = allocationAccountSentinel("allocation exceeds allocator size limit") + ErrAllocationAccountInvariant error = allocationAccountSentinel("allocation account invariant failure") + ErrAllocationAdmissionSuspended error = allocationAccountSentinel("allocation account admission is suspended") + ErrAllocationMetadataSlots error = allocationAccountSentinel("allocation metadata slots exhausted") + ErrAllocationGenerationSlots error = allocationAccountSentinel("allocation account generation slots exhausted") + ErrAllocationAccountLive error = allocationAccountSentinel("allocation account still owns memory") +) + +type allocationAccountSentinel string + +func (e allocationAccountSentinel) Error() string { return string(e) } + +type allocationAccountDetailError struct { + cause error + detail string + detailFirst bool +} + +func (e *allocationAccountDetailError) Error() string { + if e == nil || e.cause == nil { + return "allocation account error" + } + if e.detail == "" { + return e.cause.Error() + } + if e.detailFirst { + return e.detail + ": " + e.cause.Error() + } + return e.cause.Error() + ": " + e.detail +} + +func prefixAllocationAccountError( + cause error, + format string, + args ...any, +) error { + return &allocationAccountDetailError{ + cause: cause, + detail: fmt.Sprintf(format, args...), + detailFirst: true, + } +} + +func (e *allocationAccountDetailError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func wrapAllocationAccountError( + cause error, + format string, + args ...any, +) error { + return &allocationAccountDetailError{ + cause: cause, + detail: fmt.Sprintf(format, args...), + } +} + +const ( + allocationAccountSealedBit = uint64(1) << 63 + allocationAccountUsedMask = allocationAccountSealedBit - 1 +) + +// AllocationAccountHandle identifies one use of a reusable registry slot. +// The upper 32 bits are the slot generation and the lower 32 bits are the slot. +type AllocationAccountHandle uint64 + +func newAllocationAccountHandle(slot, generation uint32) AllocationAccountHandle { + return AllocationAccountHandle(uint64(generation)<<32 | uint64(slot)) +} + +func (h AllocationAccountHandle) slot() uint32 { + return uint32(h) +} + +func (h AllocationAccountHandle) generation() uint32 { + return uint32(uint64(h) >> 32) +} + +// AllocationAccountSnapshot is immutable observation state. A terminal owner +// may publish it after Seal and exact zero; taking a snapshot does not mutate +// account lifecycle. +type AllocationAccountSnapshot struct { + Handle AllocationAccountHandle + Limit uint64 + Used uint64 + Peak uint64 + Sealed bool +} + +// AllocationAccountTerminalState classifies the one immutable snapshot +// exported by an execution generation. A nonzero account at terminal cleanup +// is an ownership invariant failure, not recoverable capacity pressure. +type AllocationAccountTerminalState uint8 + +const ( + AllocationAccountTerminalValid AllocationAccountTerminalState = iota + 1 + AllocationAccountTerminalInvariantFailure +) + +// AllocationFailureReason is the non-overlapping control-flow reason exposed +// to later pressure handling. Only Capacity is eligible for reclaim/spill or +// a smaller operation retry; every other reason is terminal for that logical +// operation or generation. +type AllocationFailureReason uint8 + +const ( + AllocationFailureNone AllocationFailureReason = iota + AllocationFailureCapacity + AllocationFailureSealed + AllocationFailureMismatch + AllocationFailureAllocatorLimit + AllocationFailureInvariant + AllocationFailureSuspended +) + +func AllocationFailureReasonOf(err error) AllocationFailureReason { + switch { + case errors.Is(err, ErrAllocationAccountInvariant): + return AllocationFailureInvariant + case errors.Is(err, ErrAllocationAccountInvalid), + errors.Is(err, ErrAllocationAccountStale), + errors.Is(err, ErrAllocationAccountLive), + errors.Is(err, ErrAllocationGenerationSlots): + return AllocationFailureInvariant + case errors.Is(err, ErrAllocationAccountMismatch): + return AllocationFailureMismatch + case errors.Is(err, ErrAllocationAccountSealed): + return AllocationFailureSealed + case errors.Is(err, ErrAllocationAllocatorLimit): + return AllocationFailureAllocatorLimit + case errors.Is(err, ErrAllocationAdmissionSuspended): + return AllocationFailureSuspended + case errors.Is(err, ErrAllocationAccountCapacity), + errors.Is(err, ErrAllocationMetadataSlots), + IsMPoolCapacityFailure(err): + return AllocationFailureCapacity + default: + return AllocationFailureNone + } +} + +// IsMPoolCapacityFailure recognizes both a direct MO error and a contextual +// wrapper retained by an intermediate owner. +func IsMPoolCapacityFailure(err error) bool { + var moErr *moerr.Error + return errors.As(err, &moErr) && moErr.ErrorCode() == moerr.ErrMPoolCapacity +} + +func IsRetryableAllocationCapacity(err error) bool { + return AllocationFailureReasonOf(err) == AllocationFailureCapacity +} + +// AllocationAccountTerminalSnapshot is the immutable terminal observation of +// one generation. Failure snapshots retain the live-byte value observed at +// the terminal boundary even if a later physical Free drains the tombstone. +type AllocationAccountTerminalSnapshot struct { + AllocationAccountSnapshot + State AllocationAccountTerminalState + LiveOwner AllocationOwner + LiveSite AllocationSite + LiveAllocations uint64 +} + +// AllocationCapacityController lets an account share a higher-level aggregate +// cap. The controller owns cap policy only; physical MPool metadata remains +// the sole release owner. +type AllocationCapacityController interface { + AcquireAllocationCapacity(uint64) error + ReleaseAllocationCapacity(uint64) +} + +// AllocationCapacityClass selects a capacity controller for one physical +// allocation. Class zero uses the account's statement controller. Non-zero +// classes are registered by execution owners whose future recovery storage +// must borrow pre-admitted headroom without charging that headroom twice. +type AllocationCapacityClass uint8 + +// AllocationCapacityClassDefault uses the statement's ordinary controller. +const AllocationCapacityClassDefault AllocationCapacityClass = 0 + +// AllocationAccount owns physical allocation capacity for one execution +// generation. state packs the sealed bit and used bytes into one atomic word, +// so Acquire and Seal have one unambiguous linearization point. +type AllocationAccount struct { + registry *AllocationAccountRegistry + handle AllocationAccountHandle + limit uint64 + control AllocationCapacityController + + state atomic.Uint64 + peak atomic.Uint64 + inflight atomic.Int64 + + capacityMu sync.Mutex + capacityControllers map[AllocationCapacityClass]*allocationCapacityRegistration + nextCapacityClass AllocationCapacityClass +} + +type allocationCapacityRegistration struct { + control AllocationCapacityController + used uint64 +} + +func (a *AllocationAccount) Handle() AllocationAccountHandle { + if a == nil { + return 0 + } + return a.handle +} + +func (a *AllocationAccount) Snapshot() AllocationAccountSnapshot { + if a == nil { + return AllocationAccountSnapshot{} + } + state := a.state.Load() + return AllocationAccountSnapshot{ + Handle: a.handle, + Limit: a.limit, + Used: state & allocationAccountUsedMask, + Peak: a.peak.Load(), + Sealed: state&allocationAccountSealedBit != 0, + } +} + +// RegisterCapacityController installs one execution-local capacity class. The +// owner must unregister it only after every allocation in the class is freed. +func (a *AllocationAccount) RegisterCapacityController( + control AllocationCapacityController, +) (AllocationCapacityClass, error) { + if a == nil || control == nil || a.registry == nil || a.handle == 0 { + return AllocationCapacityClassDefault, ErrAllocationAccountInvalid + } + a.inflight.Add(1) + defer a.inflight.Add(-1) + if a.state.Load()&allocationAccountSealedBit != 0 { + return AllocationCapacityClassDefault, ErrAllocationAccountSealed + } + a.capacityMu.Lock() + defer a.capacityMu.Unlock() + if a.state.Load()&allocationAccountSealedBit != 0 { + return AllocationCapacityClassDefault, ErrAllocationAccountSealed + } + if a.capacityControllers == nil { + a.capacityControllers = make( + map[AllocationCapacityClass]*allocationCapacityRegistration, + ) + } + for range uint16(math.MaxUint8) { + a.nextCapacityClass++ + if a.nextCapacityClass == AllocationCapacityClassDefault { + a.nextCapacityClass++ + } + if _, exists := a.capacityControllers[a.nextCapacityClass]; !exists { + a.capacityControllers[a.nextCapacityClass] = + &allocationCapacityRegistration{control: control} + return a.nextCapacityClass, nil + } + } + return AllocationCapacityClassDefault, ErrAllocationAccountInvariant +} + +// UnregisterCapacityController removes a quiescent execution-local class. +func (a *AllocationAccount) UnregisterCapacityController( + class AllocationCapacityClass, + control AllocationCapacityController, +) error { + if a == nil || class == AllocationCapacityClassDefault || control == nil { + return ErrAllocationAccountInvalid + } + a.capacityMu.Lock() + defer a.capacityMu.Unlock() + registration := a.capacityControllers[class] + if registration == nil || registration.control != control { + return ErrAllocationAccountMismatch + } + if registration.used != 0 { + return ErrAllocationAccountLive + } + delete(a.capacityControllers, class) + if len(a.capacityControllers) == 0 { + a.capacityControllers = nil + } + return nil +} + +// capacityControllerLocked requires capacityMu to be held, keeping a +// non-default registration stable through the complete acquire/release. +func (a *AllocationAccount) capacityControllerLocked( + class AllocationCapacityClass, +) (*allocationCapacityRegistration, error) { + if class == AllocationCapacityClassDefault { + return nil, ErrAllocationAccountInvalid + } + registration := a.capacityControllers[class] + if registration == nil { + return nil, ErrAllocationAccountInvalid + } + return registration, nil +} + +func (a *AllocationAccount) acquire(capacity uint64) error { + return a.acquireWithCapacityClass(capacity, AllocationCapacityClassDefault) +} + +func (a *AllocationAccount) acquireWithCapacityClass( + capacity uint64, + class AllocationCapacityClass, +) error { + if a == nil || a.registry == nil || a.handle == 0 { + return ErrAllocationAccountInvalid + } + if capacity == 0 { + return nil + } + state := a.state.Load() + if state&allocationAccountSealedBit != 0 { + return ErrAllocationAccountSealed + } + used := state & allocationAccountUsedMask + if used > a.limit || capacity > a.limit-used { + return newAllocationAccountCapacityError(used, capacity, a.limit) + } + + // Register before consulting the shared controller. Once Seal publishes the + // sealed bit, it either observes this transaction or this transaction + // observes sealed before acquiring controller capacity. + a.inflight.Add(1) + defer a.inflight.Add(-1) + state = a.state.Load() + if state&allocationAccountSealedBit != 0 { + return ErrAllocationAccountSealed + } + used = state & allocationAccountUsedMask + if used > a.limit || capacity > a.limit-used { + return newAllocationAccountCapacityError(used, capacity, a.limit) + } + control := a.control + var registration *allocationCapacityRegistration + if class != AllocationCapacityClassDefault { + a.capacityMu.Lock() + var err error + registration, err = a.capacityControllerLocked(class) + if err != nil { + a.capacityMu.Unlock() + return err + } + control = registration.control + defer a.capacityMu.Unlock() + } + if control != nil { + if err := control.AcquireAllocationCapacity(capacity); err != nil { + return err + } + } + acquired := false + defer func() { + if !acquired && control != nil { + control.ReleaseAllocationCapacity(capacity) + } + }() + + for { + state = a.state.Load() + if state&allocationAccountSealedBit != 0 { + return ErrAllocationAccountSealed + } + used = state & allocationAccountUsedMask + if used > a.limit || capacity > a.limit-used { + return newAllocationAccountCapacityError(used, capacity, a.limit) + } + next := used + capacity + if a.state.CompareAndSwap(state, next) { + if registration != nil { + registration.used += capacity + } + for { + peak := a.peak.Load() + if next <= peak || a.peak.CompareAndSwap(peak, next) { + acquired = true + return nil + } + } + } + } +} + +func newAllocationAccountCapacityError( + used uint64, + requested uint64, + limit uint64, +) error { + return wrapAllocationAccountError( + ErrAllocationAccountCapacity, + "used=%d requested=%d limit=%d", + used, + requested, + limit, + ) +} + +func (a *AllocationAccount) release(capacity uint64) { + a.releaseWithCapacityClass(capacity, AllocationCapacityClassDefault) +} + +func (a *AllocationAccount) releaseWithCapacityClass( + capacity uint64, + class AllocationCapacityClass, +) { + if capacity == 0 { + return + } + // Keep the local charge until the higher-level policy charge is gone. With + // metadata released by allocationLease first, exact local zero is therefore + // also a complete-release boundary. + control := a.control + var registration *allocationCapacityRegistration + if class != AllocationCapacityClassDefault { + a.capacityMu.Lock() + var err error + registration, err = a.capacityControllerLocked(class) + if err != nil { + a.capacityMu.Unlock() + panic(err) + } + control = registration.control + defer a.capacityMu.Unlock() + } + if control != nil { + state := a.state.Load() + if capacity > state&allocationAccountUsedMask { + panic("allocation account release underflow") + } + control.ReleaseAllocationCapacity(capacity) + } + for { + state := a.state.Load() + used := state & allocationAccountUsedMask + if capacity > used { + panic("allocation account release underflow") + } + next := state - capacity + if a.state.CompareAndSwap(state, next) { + if registration != nil { + if capacity > registration.used { + panic("allocation capacity class release underflow") + } + registration.used -= capacity + } + if next&allocationAccountUsedMask == 0 { + a.registry.tryDrainTombstone(a) + } + return + } + } +} + +// Seal prevents every later acquisition. It waits only for acquisitions that +// linearized before the sealed bit was published to finish updating peak. +func (a *AllocationAccount) Seal() AllocationAccountSnapshot { + if a == nil { + return AllocationAccountSnapshot{} + } + for { + state := a.state.Load() + if state&allocationAccountSealedBit != 0 || + a.state.CompareAndSwap(state, state|allocationAccountSealedBit) { + break + } + } + for a.inflight.Load() != 0 { + runtime.Gosched() + } + return a.Snapshot() +} + +type allocationAccountRegistrySlot struct { + account atomic.Pointer[AllocationAccount] + // terminal and tombstone are protected by AllocationAccountRegistry.mu. + terminal *AllocationAccountTerminalSnapshot + tombstone bool +} + +// AllocationAccountRegistry bounds live generations and accounted-allocation +// metadata for one CN. Registry slots are reused only after Seal and exact +// zero. Their generation counters never wrap. +type AllocationAccountRegistry struct { + mu sync.Mutex + + slots []allocationAccountRegistrySlot + generations []uint32 + free []uint32 + suspended bool + tombstones uint32 + + maxAllocations uint64 + liveAllocations atomic.Uint64 + peakAllocations atomic.Uint64 +} + +func NewAllocationAccountRegistry( + generationSlots uint32, + allocationSlots uint64, +) (*AllocationAccountRegistry, error) { + if generationSlots == 0 || uint64(generationSlots) >= uint64(math.MaxInt) { + return nil, ErrAllocationAccountInvalid + } + registry := &AllocationAccountRegistry{ + slots: make([]allocationAccountRegistrySlot, uint64(generationSlots)+1), + generations: make([]uint32, uint64(generationSlots)+1), + free: make([]uint32, generationSlots), + maxAllocations: allocationSlots, + } + for i := uint32(0); i < generationSlots; i++ { + registry.free[i] = generationSlots - i + } + return registry, nil +} + +func (r *AllocationAccountRegistry) Open( + limit uint64, +) (*AllocationAccount, error) { + return r.OpenWithController(limit, nil) +} + +func (r *AllocationAccountRegistry) OpenWithController( + limit uint64, + control AllocationCapacityController, +) (*AllocationAccount, error) { + if r == nil || limit > allocationAccountUsedMask { + return nil, ErrAllocationAccountInvalid + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.suspended { + return nil, ErrAllocationAdmissionSuspended + } + for len(r.free) > 0 { + index := len(r.free) - 1 + slot := r.free[index] + r.free = r.free[:index] + generation := r.generations[slot] + if generation == math.MaxUint32 { + continue + } + generation++ + r.generations[slot] = generation + account := &AllocationAccount{ + registry: r, + handle: newAllocationAccountHandle(slot, generation), + limit: limit, + control: control, + } + r.slots[slot].account.Store(account) + return account, nil + } + return nil, ErrAllocationGenerationSlots +} + +// CompleteTerminal seals one generation and publishes its immutable terminal +// state exactly once. A nonzero terminal state remains resolvable as a +// release-capable tombstone and suspends new generations on this registry. +// The tombstone is removed automatically after the last physical Free. +// +// first is true only for the call that created the immutable snapshot. A +// repeated call while a tombstone is live returns the same snapshot with +// first=false. +func (r *AllocationAccountRegistry) CompleteTerminal( + account *AllocationAccount, +) (snapshot AllocationAccountTerminalSnapshot, first bool, err error) { + return r.CompleteTerminalWithError(account, nil) +} + +// CompleteTerminalWithError additionally records an owner-lifecycle invariant +// discovered after physical producers quiesced. A zero-live failure is removed +// immediately (there is no provenance to retain); a nonzero failure follows +// the same tombstone/suspension path as CompleteTerminal. +func (r *AllocationAccountRegistry) CompleteTerminalWithError( + account *AllocationAccount, + terminalCause error, +) (snapshot AllocationAccountTerminalSnapshot, first bool, err error) { + if r == nil || account == nil || account.registry != r { + return snapshot, false, ErrAllocationAccountInvalid + } + account.Seal() + + r.mu.Lock() + defer r.mu.Unlock() + slot := account.handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) || + r.slots[slot].account.Load() != account { + return snapshot, false, ErrAllocationAccountStale + } + entry := &r.slots[slot] + if entry.terminal != nil { + snapshot = *entry.terminal + if snapshot.State == AllocationAccountTerminalInvariantFailure { + return snapshot, false, newAllocationTerminalInvariantError(snapshot) + } + return snapshot, false, nil + } + + current := account.Snapshot() + if !current.Sealed || account.inflight.Load() != 0 { + return AllocationAccountTerminalSnapshot{ + AllocationAccountSnapshot: current, + State: AllocationAccountTerminalInvariantFailure, + }, false, wrapAllocationAccountError( + ErrAllocationAccountInvariant, + "terminal account is not quiescent", + ) + } + account.capacityMu.Lock() + liveCapacityControllers := len(account.capacityControllers) + account.capacityMu.Unlock() + if liveCapacityControllers != 0 { + terminalCause = errors.Join( + terminalCause, + wrapAllocationAccountError( + ErrAllocationAccountInvariant, + "terminal account retains %d capacity controllers", + liveCapacityControllers, + ), + ) + } + snapshot = AllocationAccountTerminalSnapshot{ + AllocationAccountSnapshot: current, + State: AllocationAccountTerminalValid, + } + if terminalCause != nil { + snapshot.State = AllocationAccountTerminalInvariantFailure + } + if current.Used == 0 { + entry.terminal = &snapshot + r.removeSlotLocked(slot, account) + if snapshot.State == AllocationAccountTerminalInvariantFailure { + return snapshot, true, errors.Join( + terminalCause, + newAllocationTerminalInvariantError(snapshot), + ) + } + return snapshot, true, nil + } + + snapshot.State = AllocationAccountTerminalInvariantFailure + snapshot.LiveOwner, snapshot.LiveSite, snapshot.LiveAllocations = + allocationAccountLiveDiagnostic(account) + entry.terminal = &snapshot + entry.tombstone = true + r.tombstones++ + r.suspended = true + // A physical Free may have raced the terminal observation. The immutable + // failure snapshot remains truthful at its linearization point, while a + // now-empty tombstone can be removed immediately. + if account.Snapshot().Used == 0 && account.inflight.Load() == 0 { + r.removeTombstoneLocked(slot, account) + } + return snapshot, true, errors.Join( + terminalCause, + newAllocationTerminalInvariantError(snapshot), + ) +} + +func newAllocationTerminalInvariantError( + snapshot AllocationAccountTerminalSnapshot, +) error { + return wrapAllocationAccountError( + ErrAllocationAccountInvariant, + "handle=%d used=%d peak=%d limit=%d owner=%d site=%d live-allocations=%d", + snapshot.Handle, + snapshot.Used, + snapshot.Peak, + snapshot.Limit, + snapshot.LiveOwner, + snapshot.LiveSite, + snapshot.LiveAllocations, + ) +} + +// allocationAccountLiveDiagnostic is a terminal-only scan. It does not add a +// per-allocation hot-path counter: provenance already lives in the pointer +// metadata required for physical Free. The first live owner/site plus the +// exact live allocation count makes a nonzero terminal snapshot actionable. +func allocationAccountLiveDiagnostic( + account *AllocationAccount, +) (AllocationOwner, AllocationSite, uint64) { + if account == nil { + return 0, 0, 0 + } + var owner AllocationOwner + var site AllocationSite + var count uint64 + record := func(lease allocationLease) { + if lease.account != account { + return + } + if count == 0 { + owner = lease.owner + site = lease.site + } + count++ + } + for i := range globalPtrShards { + shard := &globalPtrShards[i] + shard.mu.Lock() + for _, lease := range shard.leases { + record(lease) + } + shard.mu.Unlock() + } + // noLock pools intentionally provide no synchronization for their local + // maps. Do not race unrelated single-threaded pools merely to enrich a + // terminal error; production query pools use the sharded metadata above. + return owner, site, count +} + +func (r *AllocationAccountRegistry) removeSlotLocked( + slot uint32, + account *AllocationAccount, +) { + entry := &r.slots[slot] + if entry.account.Load() != account { + return + } + entry.account.Store(nil) + entry.terminal = nil + entry.tombstone = false + if r.generations[slot] != math.MaxUint32 { + r.free = append(r.free, slot) + } +} + +func (r *AllocationAccountRegistry) removeTombstoneLocked( + slot uint32, + account *AllocationAccount, +) { + entry := &r.slots[slot] + if entry.account.Load() != account || !entry.tombstone { + return + } + if r.tombstones == 0 { + panic("allocation account tombstone underflow") + } + r.tombstones-- + r.removeSlotLocked(slot, account) + r.suspended = r.tombstones != 0 +} + +func (r *AllocationAccountRegistry) tryDrainTombstone( + account *AllocationAccount, +) { + if r == nil || account == nil || account.Snapshot().Used != 0 || + account.inflight.Load() != 0 { + return + } + r.mu.Lock() + defer r.mu.Unlock() + slot := account.handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) { + return + } + if r.slots[slot].account.Load() == account && + r.slots[slot].tombstone && + account.Snapshot().Used == 0 && + account.inflight.Load() == 0 { + r.removeTombstoneLocked(slot, account) + } +} + +func (r *AllocationAccountRegistry) AdmissionSuspended() bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + return r.suspended +} + +func (r *AllocationAccountRegistry) LiveTombstones() uint32 { + if r == nil { + return 0 + } + r.mu.Lock() + defer r.mu.Unlock() + return r.tombstones +} + +func (r *AllocationAccountRegistry) Resolve( + handle AllocationAccountHandle, +) (*AllocationAccount, bool) { + if r == nil { + return nil, false + } + slot := handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) { + return nil, false + } + account := r.slots[slot].account.Load() + return account, account != nil && account.handle == handle +} + +// Finalize removes a sealed, empty account and makes its slot reusable. A live +// account remains resolvable so physical Free can still release its charge. +func (r *AllocationAccountRegistry) Finalize( + account *AllocationAccount, +) (AllocationAccountSnapshot, error) { + if r == nil || account == nil || account.registry != r { + return AllocationAccountSnapshot{}, ErrAllocationAccountInvalid + } + snapshot := account.Snapshot() + if !snapshot.Sealed || snapshot.Used != 0 || + account.inflight.Load() != 0 { + return snapshot, ErrAllocationAccountLive + } + + r.mu.Lock() + defer r.mu.Unlock() + slot := account.handle.slot() + if slot == 0 || uint64(slot) >= uint64(len(r.slots)) || + r.slots[slot].account.Load() != account { + return snapshot, ErrAllocationAccountStale + } + if current := account.Snapshot(); !current.Sealed || current.Used != 0 || + account.inflight.Load() != 0 { + return current, ErrAllocationAccountLive + } + r.removeSlotLocked(slot, account) + return account.Snapshot(), nil +} + +func (r *AllocationAccountRegistry) reserveMetadata() error { + if r == nil { + return ErrAllocationAccountInvalid + } + for { + live := r.liveAllocations.Load() + if live >= r.maxAllocations { + return ErrAllocationMetadataSlots + } + if r.liveAllocations.CompareAndSwap(live, live+1) { + next := live + 1 + for { + peak := r.peakAllocations.Load() + if next <= peak || + r.peakAllocations.CompareAndSwap(peak, next) { + break + } + } + return nil + } + } +} + +func (r *AllocationAccountRegistry) releaseMetadata() { + for { + live := r.liveAllocations.Load() + if live == 0 { + panic("allocation metadata slot release underflow") + } + if r.liveAllocations.CompareAndSwap(live, live-1) { + return + } + } +} + +// LiveAllocationMetadata returns published and currently in-flight metadata +// slots. A failed unpublished transaction returns its slot before returning. +func (r *AllocationAccountRegistry) LiveAllocationMetadata() uint64 { + if r == nil { + return 0 + } + return r.liveAllocations.Load() +} + +// PeakAllocationMetadata returns the exact high-water slot count. +func (r *AllocationAccountRegistry) PeakAllocationMetadata() uint64 { + if r == nil { + return 0 + } + return r.peakAllocations.Load() +} + +func (r *AllocationAccountRegistry) MaxAllocationMetadata() uint64 { + if r == nil { + return 0 + } + return r.maxAllocations +} + +func (r *AllocationAccountRegistry) GenerationCapacity() uint32 { + if r == nil || len(r.slots) == 0 { + return 0 + } + return uint32(len(r.slots) - 1) +} + +type allocationAccountRequest struct { + account *AllocationAccount + owner AllocationOwner + site AllocationSite + capacityClass AllocationCapacityClass +} + +func (r allocationAccountRequest) validate() error { + if r.account == nil || r.account.registry == nil || + r.owner < AllocationOwnerMin || r.owner > AllocationOwnerMax || + r.site < AllocationSiteMin { + return ErrAllocationAccountInvalid + } + resolved, ok := r.account.registry.Resolve(r.account.handle) + if !ok || resolved != r.account { + return ErrAllocationAccountStale + } + return nil +} + +type allocationLease struct { + account *AllocationAccount + owner AllocationOwner + site AllocationSite + profiled bool + capacityClass AllocationCapacityClass + _ [4]byte +} + +func (l allocationLease) release(capacity uint64) { + if l.account == nil || l.account.registry == nil { + panic("invalid allocation account lease") + } + // Return finite metadata first. account.release retains the local charge + // until controller cleanup completes, so exact zero is a complete-release + // boundary. + l.account.registry.releaseMetadata() + l.account.releaseWithCapacityClass(capacity, l.capacityClass) +} diff --git a/pkg/common/mpool/allocation_account_mpool_test.go b/pkg/common/mpool/allocation_account_mpool_test.go new file mode 100644 index 0000000000000..75dc286ccc0e8 --- /dev/null +++ b/pkg/common/mpool/allocation_account_mpool_test.go @@ -0,0 +1,730 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package mpool + +import ( + "bytes" + "fmt" + "sync" + "testing" + "unsafe" + + "github.com/stretchr/testify/require" +) + +const ( + testAllocationOwner AllocationOwner = 1 + testAllocationSite AllocationSite = 1 +) + +func newTestAllocationAccount( + t testing.TB, + limit uint64, + metadataSlots uint64, +) (*AllocationAccountRegistry, *AllocationAccount) { + t.Helper() + registry, err := NewAllocationAccountRegistry(4, metadataSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + return registry, account +} + +func finalizeTestAllocationAccount( + t testing.TB, + registry *AllocationAccountRegistry, + account *AllocationAccount, +) { + t.Helper() + account.Seal() + _, err := registry.Finalize(account) + require.NoError(t, err) +} + +func TestMPoolAccountedAllocGrowFree(t *testing.T) { + require.Equal(t, uintptr(kMemHdrSz), unsafe.Sizeof(memHdr{})) + require.Equal(t, uintptr(16), unsafe.Sizeof(allocationLease{})) + require.LessOrEqual(t, unsafe.Sizeof(AllocationAccount{}), uintptr(96)) + require.LessOrEqual(t, unsafe.Sizeof(allocationAccountRegistrySlot{}), uintptr(32)) + + registry, account := newTestAllocationAccount(t, 1024, 8) + mp := MustNew("accounted-alloc-grow") + defer DeleteMPool(mp) + + empty, err := mp.AllocAccounted( + 0, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Nil(t, empty) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + var lease allocationLease + hdr, ok := mp.getPtrMetadata( + unsafe.Pointer(unsafe.SliceData(buffer)), + &lease, + ) + require.True(t, ok) + require.True(t, hdr.isOffHeap()) + require.True(t, hdr.isAccounted()) + require.Same(t, account, lease.account) + require.Equal(t, testAllocationOwner, lease.owner) + require.Equal(t, testAllocationSite, lease.site) + + same, err := mp.Grow(buffer, 32, true) + require.NoError(t, err) + require.Equal( + t, + unsafe.Pointer(unsafe.SliceData(buffer)), + unsafe.Pointer(unsafe.SliceData(same)), + ) + require.Equal(t, uint64(64), account.Snapshot().Used) + + grown, err := mp.Grow(same, 128, true) + require.NoError(t, err) + require.Equal(t, uint64(cap(grown)), account.Snapshot().Used) + require.Equal( + t, + uint64(64+cap(grown)), + account.Snapshot().Peak, + ) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) + + mp.Free(grown) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestMPoolTerminalLeakDiagnosticUsesPublishedProvenance(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNew("accounted-terminal-diagnostic") + defer DeleteMPool(mp) + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + snapshot, first, err := registry.CompleteTerminal(account) + require.True(t, first) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.Equal(t, testAllocationOwner, snapshot.LiveOwner) + require.Equal(t, testAllocationSite, snapshot.LiveSite) + require.Equal(t, uint64(1), snapshot.LiveAllocations) + require.Contains(t, err.Error(), "owner=1 site=1 live-allocations=1") + mp.Free(buffer) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(snapshot.Handle) + require.False(t, ok) +} + +func TestMPoolMakeSliceAccounted(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 2) + mp := MustNew("accounted-typed-slice") + defer DeleteMPool(mp) + + values, err := MakeSliceAccounted[int64]( + 4, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Len(t, values, 4) + require.Equal(t, uint64(32), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + for i := range values { + values[i] = int64(i + 1) + } + require.Equal(t, []int64{1, 2, 3, 4}, values) + + _, err = MakeSliceAccounted[int64]( + 5, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, uint64(32), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + empty, err := MakeSliceAccounted[int64]( + 0, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Nil(t, empty) + _, err = MakeSliceAccounted[int64]( + -1, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + _, err = MakeSliceAccounted[struct{}]( + 1, + mp, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAllocatorLimit) + + FreeSlice(mp, values[:0]) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestMPoolAccountedRollback(t *testing.T) { + t.Run("account-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 63, 1) + mp := MustNew("accounted-capacity") + defer DeleteMPool(mp) + + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Contains(t, err.Error(), "owner=1 site=1") + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("metadata-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 0) + mp := MustNew("accounted-metadata-capacity") + defer DeleteMPool(mp) + + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Contains(t, err.Error(), "owner=1 site=1") + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("pool-capacity", func(t *testing.T) { + const allocationSize = 768 << 10 + registry, account := newTestAllocationAccount( + t, + 2<<20, + 2, + ) + mp, err := NewMPool("accounted-pool-capacity", 1<<20, NoFixed) + require.NoError(t, err) + defer DeleteMPool(mp) + + first, err := mp.AllocAccounted( + allocationSize, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + _, err = mp.AllocAccounted( + allocationSize, + account, + testAllocationOwner, + testAllocationSite, + ) + require.Error(t, err) + require.True(t, IsMPoolCapacityFailure(err)) + require.Equal(t, AllocationFailureCapacity, AllocationFailureReasonOf(err)) + require.True(t, IsRetryableAllocationCapacity(err)) + require.Equal(t, uint64(allocationSize), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) + + mp.Free(first) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("global-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNew("accounted-global-capacity") + defer DeleteMPool(mp) + + oldGlobalCap := globalCap.Load() + globalBefore := GlobalStats().NumCurrBytes.Load() + globalCap.Store(globalBefore + 63) + t.Cleanup(func() { + globalCap.Store(oldGlobalCap) + }) + + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.Error(t, err) + require.True(t, IsMPoolCapacityFailure(err)) + require.Equal(t, AllocationFailureCapacity, AllocationFailureReasonOf(err)) + require.True(t, IsRetryableAllocationCapacity(err)) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + require.Zero(t, mp.CurrNB()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("grow-metadata-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 1) + mp := MustNew("accounted-grow-metadata-capacity") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + copy(buffer, bytes.Repeat([]byte{0x5a}, len(buffer))) + + _, err = mp.Grow(buffer, 128, true) + require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Equal(t, bytes.Repeat([]byte{0x5a}, len(buffer)), buffer) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + mp.Free(buffer) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("realloc-zero-metadata-capacity", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 1) + mp := MustNew("accounted-realloc-zero-metadata-capacity") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + copy(buffer, bytes.Repeat([]byte{0x5a}, len(buffer))) + + _, err = mp.ReallocZero(buffer, 128, true) + require.ErrorIs(t, err, ErrAllocationMetadataSlots) + require.Equal(t, bytes.Repeat([]byte{0x5a}, len(buffer)), buffer) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + mp.Free(buffer) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("allocator-size-limit-precedes-admission", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 0, 1) + mp := MustNew("accounted-allocator-size-limit") + defer DeleteMPool(mp) + + _, err := mp.AllocAccounted( + CapLimit-kMemHdrSz, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + _, err = mp.AllocAccounted( + CapLimit-kMemHdrSz+1, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAllocatorLimit) + require.NotErrorIs(t, err, ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) +} + +func TestMPoolAccountedGrowCapacityBoundary(t *testing.T) { + const ( + oldCapacity = 64 + required = 65 + ) + newCapacity, ok := GrowCapacity(oldCapacity, required) + require.True(t, ok) + require.Greater(t, newCapacity, int64(required)) + + for _, testCase := range []struct { + name string + limit uint64 + wantError bool + }{ + { + name: "exact-old-plus-rounded-new", + limit: oldCapacity + uint64(newCapacity), + }, + { + name: "one-byte-short", + limit: oldCapacity + uint64(newCapacity) - 1, + wantError: true, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + registry, account := newTestAllocationAccount( + t, + testCase.limit, + 2, + ) + mp := MustNew("accounted-grow-capacity-boundary") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + oldCapacity, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + grown, err := mp.Grow(buffer, required, true) + if testCase.wantError { + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, uint64(oldCapacity), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + mp.Free(buffer) + } else { + require.NoError(t, err) + require.Equal(t, newCapacity, int64(cap(grown))) + require.Equal(t, uint64(newCapacity), account.Snapshot().Used) + require.Equal(t, testCase.limit, account.Snapshot().Peak) + mp.Free(grown) + } + finalizeTestAllocationAccount(t, registry, account) + }) + } +} + +func TestMPoolAccountedReallocZero(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 2) + mp := MustNew("accounted-realloc-zero") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + copy(buffer, bytes.Repeat([]byte{0x5a}, 64)) + + _, err = mp.ReallocZero(buffer, 128, false) + require.ErrorIs(t, err, ErrAllocationAccountMismatch) + require.Equal(t, uint64(64), account.Snapshot().Used) + + replacement, err := mp.ReallocZero(buffer, 128, true) + require.NoError(t, err) + require.Equal(t, bytes.Repeat([]byte{0x5a}, 64), replacement[:64]) + require.Equal(t, make([]byte, 64), replacement[64:]) + require.Equal(t, uint64(128), account.Snapshot().Used) + require.Equal(t, uint64(192), account.Snapshot().Peak) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + require.Equal(t, uint64(2), registry.PeakAllocationMetadata()) + + mp.Free(replacement) + finalizeTestAllocationAccount(t, registry, account) +} + +func TestMPoolAccountedSealRejectsGrowth(t *testing.T) { + registry, account := newTestAllocationAccount(t, 512, 2) + mp := MustNew("accounted-sealed-growth") + defer DeleteMPool(mp) + + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + account.Seal() + _, err = mp.Grow(buffer, 128, true) + require.ErrorIs(t, err, ErrAllocationAccountSealed) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + mp.Free(buffer) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestMPoolAccountedCrossPoolAndTeardown(t *testing.T) { + t.Run("cross-pool", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + owner := MustNew("accounted-cross-owner") + other := MustNew("accounted-cross-other") + defer DeleteMPool(owner) + defer DeleteMPool(other) + + buffer, err := owner.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + other.Free(buffer) + require.Panics(t, func() { + other.Free(buffer) + }) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("deleted-owner", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + owner := MustNew("accounted-deleted-owner") + other := MustNew("accounted-deleted-other") + defer DeleteMPool(other) + + globalBefore := GlobalStats().NumCurrBytes.Load() + buffer, err := owner.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + DeleteMPool(owner) + require.Equal(t, globalBefore+64, GlobalStats().NumCurrBytes.Load()) + require.Equal(t, uint64(64), account.Snapshot().Used) + require.Equal(t, uint64(1), registry.LiveAllocationMetadata()) + + other.Free(buffer) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) + + t.Run("no-lock-teardown", func(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNewNoLock("accounted-no-lock-teardown") + globalBefore := GlobalStats().NumCurrBytes.Load() + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + DeleteMPool(mp) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) + }) +} + +func TestMPoolAccountedConcurrentAllocFree(t *testing.T) { + const ( + workers = 32 + rounds = 200 + size = 64 + ) + registry, account := newTestAllocationAccount( + t, + workers*size, + workers, + ) + mp := MustNew("accounted-concurrent") + defer DeleteMPool(mp) + + var wait sync.WaitGroup + wait.Add(workers) + for range workers { + go func() { + defer wait.Done() + for range rounds { + buffer, err := mp.AllocAccounted( + size, + account, + testAllocationOwner, + testAllocationSite, + ) + if err != nil { + t.Errorf("accounted alloc: %v", err) + return + } + mp.Free(buffer) + } + }() + } + wait.Wait() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + finalizeTestAllocationAccount(t, registry, account) +} + +func BenchmarkMPoolAccountedAllocation(b *testing.B) { + for _, accounted := range []bool{false, true} { + mode := "unaccounted" + if accounted { + mode = "accounted" + } + for _, size := range []int{64, 4 << 10, 16 << 10, 64 << 10} { + b.Run(fmt.Sprintf("%s/alloc-free/%d", mode, size), func(b *testing.B) { + mp := MustNew("benchmark-allocation-account") + defer DeleteMPool(mp) + var registry *AllocationAccountRegistry + var account *AllocationAccount + if accounted { + var err error + registry, err = NewAllocationAccountRegistry(1, 1) + require.NoError(b, err) + account, err = registry.Open(1 << 60) + require.NoError(b, err) + } + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for range b.N { + var buffer []byte + var allocErr error + if accounted { + buffer, allocErr = mp.AllocAccounted( + size, + account, + testAllocationOwner, + testAllocationSite, + ) + } else { + buffer, allocErr = mp.Alloc(size, true) + } + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + b.StopTimer() + if accounted { + finalizeTestAllocationAccount(b, registry, account) + } + }) + } + + b.Run(mode+"/grow-replacement", func(b *testing.B) { + mp := MustNew("benchmark-allocation-account-grow") + defer DeleteMPool(mp) + var registry *AllocationAccountRegistry + var account *AllocationAccount + if accounted { + var err error + registry, err = NewAllocationAccountRegistry(1, 2) + require.NoError(b, err) + account, err = registry.Open(1 << 60) + require.NoError(b, err) + } + b.ReportAllocs() + b.SetBytes(64 << 10) + b.ResetTimer() + for range b.N { + var buffer []byte + var allocErr error + if accounted { + buffer, allocErr = mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + } else { + buffer, allocErr = mp.Alloc(64, true) + } + if allocErr != nil { + b.Fatal(allocErr) + } + buffer, allocErr = mp.Grow(buffer, 64<<10, true) + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + b.StopTimer() + if accounted { + finalizeTestAllocationAccount(b, registry, account) + } + }) + } + + b.Run("accounted/parallel-alloc-free/65536", func(b *testing.B) { + mp := MustNew("benchmark-allocation-account-parallel") + defer DeleteMPool(mp) + registry, err := NewAllocationAccountRegistry(1, 1024) + require.NoError(b, err) + account, err := registry.Open(1 << 60) + require.NoError(b, err) + b.ReportAllocs() + b.SetBytes(64 << 10) + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + buffer, allocErr := mp.AllocAccounted( + 64<<10, + account, + testAllocationOwner, + testAllocationSite, + ) + if allocErr != nil { + b.Fatal(allocErr) + } + mp.Free(buffer) + } + }) + b.StopTimer() + finalizeTestAllocationAccount(b, registry, account) + }) +} diff --git a/pkg/common/mpool/allocation_account_test.go b/pkg/common/mpool/allocation_account_test.go new file mode 100644 index 0000000000000..5bdd1ca4aef37 --- /dev/null +++ b/pkg/common/mpool/allocation_account_test.go @@ -0,0 +1,519 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package mpool + +import ( + "errors" + "math" + "runtime" + "sync" + "sync/atomic" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/stretchr/testify/require" +) + +type testAllocationCapacityController struct { + used atomic.Uint64 + reject atomic.Bool + acquireStarted chan struct{} + acquireProceed chan struct{} + releaseStarted chan struct{} + releaseProceed chan struct{} +} + +func (c *testAllocationCapacityController) AcquireAllocationCapacity( + capacity uint64, +) error { + if c.reject.Load() { + return ErrAllocationAccountCapacity + } + c.used.Add(capacity) + if c.acquireStarted != nil { + close(c.acquireStarted) + <-c.acquireProceed + } + return nil +} + +func (c *testAllocationCapacityController) ReleaseAllocationCapacity( + capacity uint64, +) { + c.used.Add(^uint64(capacity - 1)) + if c.releaseStarted != nil { + close(c.releaseStarted) + <-c.releaseProceed + } +} + +func TestAllocationAccountRegistryLifecycle(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 2) + require.NoError(t, err) + + account, err := registry.Open(128) + require.NoError(t, err) + firstHandle := account.Handle() + resolved, ok := registry.Resolve(firstHandle) + require.True(t, ok) + require.Same(t, account, resolved) + + _, err = registry.Open(128) + require.ErrorIs(t, err, ErrAllocationGenerationSlots) + _, err = registry.Finalize(account) + require.ErrorIs(t, err, ErrAllocationAccountLive) + + require.NoError(t, account.acquire(64)) + require.NoError(t, account.acquire(64)) + err = account.acquire(1) + require.ErrorIs(t, err, ErrAllocationAccountCapacity) + require.Equal(t, AllocationAccountSnapshot{ + Handle: firstHandle, + Limit: 128, + Used: 128, + Peak: 128, + }, account.Snapshot()) + + account.release(64) + sealed := account.Seal() + require.True(t, sealed.Sealed) + require.Equal(t, uint64(64), sealed.Used) + require.ErrorIs(t, account.acquire(1), ErrAllocationAccountSealed) + // A zero-byte request is always a no-op, including after Seal. + require.NoError(t, account.acquire(0)) + + account.release(64) + final, err := registry.Finalize(account) + require.NoError(t, err) + require.True(t, final.Sealed) + require.Zero(t, final.Used) + _, ok = registry.Resolve(firstHandle) + require.False(t, ok) + _, err = registry.Finalize(account) + require.ErrorIs(t, err, ErrAllocationAccountStale) + + reused, err := registry.Open(128) + require.NoError(t, err) + require.NotEqual(t, firstHandle, reused.Handle()) + require.Equal(t, firstHandle.slot(), reused.Handle().slot()) + require.Equal(t, firstHandle.generation()+1, reused.Handle().generation()) + reused.Seal() + _, err = registry.Finalize(reused) + require.NoError(t, err) +} + +func TestAllocationAccountControllerRollback(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + controller := &testAllocationCapacityController{} + account, err := registry.OpenWithController(8, controller) + require.NoError(t, err) + + controller.reject.Store(true) + require.ErrorIs(t, account.acquire(1), ErrAllocationAccountCapacity) + require.Zero(t, controller.used.Load()) + require.Zero(t, account.Snapshot().Used) + + controller.reject.Store(false) + controller.acquireStarted = make(chan struct{}) + controller.acquireProceed = make(chan struct{}) + acquireResult := make(chan error, 1) + go func() { + acquireResult <- account.acquire(1) + }() + <-controller.acquireStarted + sealed := make(chan struct{}) + go func() { + account.Seal() + close(sealed) + }() + for !account.Snapshot().Sealed { + runtime.Gosched() + } + close(controller.acquireProceed) + require.ErrorIs(t, <-acquireResult, ErrAllocationAccountSealed) + <-sealed + require.Zero(t, controller.used.Load()) + require.Zero(t, account.Snapshot().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountCapacityClassLifecycle(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(8) + require.NoError(t, err) + controller := &testAllocationCapacityController{} + class, err := account.RegisterCapacityController(controller) + require.NoError(t, err) + require.NotEqual(t, AllocationCapacityClassDefault, class) + + require.NoError(t, account.acquireWithCapacityClass(4, class)) + require.Equal(t, uint64(4), controller.used.Load()) + require.ErrorIs( + t, account.UnregisterCapacityController(class, controller), + ErrAllocationAccountLive, + ) + account.releaseWithCapacityClass(4, class) + require.Zero(t, controller.used.Load()) + require.NoError(t, account.UnregisterCapacityController(class, controller)) + require.ErrorIs( + t, account.acquireWithCapacityClass(1, class), + ErrAllocationAccountInvalid, + ) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestAllocationAccountTerminalRejectsRegisteredCapacityClass(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + controller := &testAllocationCapacityController{} + _, err = account.RegisterCapacityController(controller) + require.NoError(t, err) + + snapshot, first, err := registry.CompleteTerminal(account) + require.True(t, first) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.Contains(t, err.Error(), "retains 1 capacity controllers") + require.Equal(t, AllocationAccountTerminalInvariantFailure, snapshot.State) + require.Zero(t, snapshot.Used) + require.False(t, registry.AdmissionSuspended()) +} + +func TestAllocationAccountFinalizeWaitsForRelease(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + controller := &testAllocationCapacityController{ + releaseStarted: make(chan struct{}), + releaseProceed: make(chan struct{}), + } + t.Cleanup(func() { + select { + case <-controller.releaseProceed: + default: + close(controller.releaseProceed) + } + }) + account, err := registry.OpenWithController(1, controller) + require.NoError(t, err) + require.NoError(t, account.acquire(1)) + require.NoError(t, registry.reserveMetadata()) + account.Seal() + + released := make(chan struct{}) + go func() { + defer close(released) + allocationLease{account: account, owner: 1, site: 1}.release(1) + }() + <-controller.releaseStarted + require.Equal(t, uint64(1), account.Snapshot().Used) + _, err = registry.Finalize(account) + require.ErrorIs(t, err, ErrAllocationAccountLive) + + close(controller.releaseProceed) + <-released + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountRegistryBounds(t *testing.T) { + _, err := NewAllocationAccountRegistry(0, 1) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + + registry, err := NewAllocationAccountRegistry(1, 0) + require.NoError(t, err) + _, err = registry.Open(allocationAccountUsedMask + 1) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + + account, err := registry.Open(1) + require.NoError(t, err) + require.NoError(t, account.acquire(1)) + require.ErrorIs(t, registry.reserveMetadata(), ErrAllocationMetadataSlots) + account.release(1) + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) + + registry.mu.Lock() + registry.generations[1] = math.MaxUint32 + registry.mu.Unlock() + _, err = registry.Open(1) + require.ErrorIs(t, err, ErrAllocationGenerationSlots) +} + +func TestAllocationAccountRequestValidation(t *testing.T) { + registry, account := newTestAllocationAccount(t, 64, 1) + mp := MustNew("accounted-request-validation") + defer DeleteMPool(mp) + + for _, testCase := range []struct { + name string + account *AllocationAccount + owner AllocationOwner + site AllocationSite + }{ + {name: "nil-account", owner: 1, site: 1}, + {name: "zero-owner", account: account, site: 1}, + {name: "owner-out-of-range", account: account, owner: 64, site: 1}, + {name: "zero-site", account: account, owner: 1}, + } { + t.Run(testCase.name, func(t *testing.T) { + _, err := mp.AllocAccounted( + 64, + testCase.account, + testCase.owner, + testCase.site, + ) + require.ErrorIs(t, err, ErrAllocationAccountInvalid) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + }) + } + finalizeTestAllocationAccount(t, registry, account) + _, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.ErrorIs(t, err, ErrAllocationAccountStale) +} + +func TestAllocationAccountSealLinearization(t *testing.T) { + const contenders = 128 + + registry, err := NewAllocationAccountRegistry(1, contenders) + require.NoError(t, err) + account, err := registry.Open(contenders) + require.NoError(t, err) + + start := make(chan struct{}) + sealed := make(chan struct{}) + var wait sync.WaitGroup + var acquired atomic.Uint64 + wait.Add(contenders) + for range contenders { + go func() { + defer wait.Done() + <-start + err := account.acquire(1) + if err == nil { + acquired.Add(1) + <-sealed + account.release(1) + return + } + if !errors.Is(err, ErrAllocationAccountSealed) { + t.Errorf("unexpected acquire error: %v", err) + } + }() + } + + close(start) + snapshot := account.Seal() + close(sealed) + wait.Wait() + require.Equal(t, acquired.Load(), snapshot.Used) + require.Equal(t, acquired.Load(), snapshot.Peak) + require.Zero(t, account.Snapshot().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountReleaseUnderflow(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + require.Panics(t, func() { + account.release(1) + }) + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestAllocationAccountTerminalTombstoneSuspendsAdmission(t *testing.T) { + registry, err := NewAllocationAccountRegistry(3, 2) + require.NoError(t, err) + account, err := registry.Open(64) + require.NoError(t, err) + require.NoError(t, account.acquire(64)) + + snapshot, first, err := registry.CompleteTerminal(account) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, first) + require.Equal(t, AllocationAccountTerminalInvariantFailure, snapshot.State) + require.Equal(t, uint64(64), snapshot.Used) + require.True(t, snapshot.Sealed) + require.True(t, registry.AdmissionSuspended()) + require.Equal(t, uint32(1), registry.LiveTombstones()) + require.ErrorIs(t, account.acquire(1), ErrAllocationAccountSealed) + _, err = registry.Open(64) + require.ErrorIs(t, err, ErrAllocationAdmissionSuspended) + + repeated, first, err := registry.CompleteTerminal(account) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.False(t, first) + require.Equal(t, snapshot, repeated) + + account.release(64) + require.False(t, registry.AdmissionSuspended()) + require.Zero(t, registry.LiveTombstones()) + _, ok := registry.Resolve(snapshot.Handle) + require.False(t, ok) + + next, err := registry.Open(64) + require.NoError(t, err) + require.NotEqual(t, snapshot.Handle, next.Handle()) + valid, first, err := registry.CompleteTerminal(next) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, AllocationAccountTerminalValid, valid.State) + require.Zero(t, valid.Used) + _, ok = registry.Resolve(valid.Handle) + require.False(t, ok) +} + +func TestAllocationAccountZeroLiveOwnerInvariantExportsFailure(t *testing.T) { + registry, err := NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + ownerErr := errors.New("owner teardown failed") + + snapshot, first, err := registry.CompleteTerminalWithError(account, ownerErr) + require.True(t, first) + require.ErrorIs(t, err, ownerErr) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.Equal(t, AllocationAccountTerminalInvariantFailure, snapshot.State) + require.Zero(t, snapshot.Used) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(account.Handle()) + require.False(t, ok) +} + +func TestAllocationAccountMultipleTombstonesDrainBeforeResume(t *testing.T) { + registry, err := NewAllocationAccountRegistry(3, 2) + require.NoError(t, err) + first, err := registry.Open(64) + require.NoError(t, err) + second, err := registry.Open(64) + require.NoError(t, err) + require.NoError(t, first.acquire(1)) + require.NoError(t, second.acquire(1)) + + _, created, err := registry.CompleteTerminal(first) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, created) + _, created, err = registry.CompleteTerminal(second) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, created) + require.Equal(t, uint32(2), registry.LiveTombstones()) + + first.release(1) + require.True(t, registry.AdmissionSuspended()) + require.Equal(t, uint32(1), registry.LiveTombstones()) + second.release(1) + require.False(t, registry.AdmissionSuspended()) + require.Zero(t, registry.LiveTombstones()) +} + +func TestAllocationAccountOpenSuspendLinearization(t *testing.T) { + const contenders = 128 + + registry, err := NewAllocationAccountRegistry(contenders+1, 1) + require.NoError(t, err) + leaked, err := registry.Open(1) + require.NoError(t, err) + require.NoError(t, leaked.acquire(1)) + + start := make(chan struct{}) + opened := make(chan *AllocationAccount, contenders) + var wait sync.WaitGroup + wait.Add(contenders) + for range contenders { + go func() { + defer wait.Done() + <-start + account, openErr := registry.Open(1) + if openErr == nil { + opened <- account + return + } + if !errors.Is(openErr, ErrAllocationAdmissionSuspended) && + !errors.Is(openErr, ErrAllocationGenerationSlots) { + t.Errorf("unexpected open error: %v", openErr) + } + }() + } + close(start) + _, first, err := registry.CompleteTerminal(leaked) + require.ErrorIs(t, err, ErrAllocationAccountInvariant) + require.True(t, first) + wait.Wait() + close(opened) + + _, err = registry.Open(1) + require.ErrorIs(t, err, ErrAllocationAdmissionSuspended) + for account := range opened { + _, _, finishErr := registry.CompleteTerminal(account) + require.NoError(t, finishErr) + } + leaked.release(1) + require.False(t, registry.AdmissionSuspended()) +} + +func TestAllocationFailureReasonsAreNonOverlapping(t *testing.T) { + testCases := []struct { + err error + reason AllocationFailureReason + retryable bool + }{ + {ErrAllocationAccountCapacity, AllocationFailureCapacity, true}, + {ErrAllocationMetadataSlots, AllocationFailureCapacity, true}, + {moerr.NewMPoolCapacityNoCtxf("test"), AllocationFailureCapacity, true}, + {ErrAllocationAccountSealed, AllocationFailureSealed, false}, + {ErrAllocationAccountMismatch, AllocationFailureMismatch, false}, + {ErrAllocationAllocatorLimit, AllocationFailureAllocatorLimit, false}, + {ErrAllocationAccountInvariant, AllocationFailureInvariant, false}, + {ErrAllocationAccountInvalid, AllocationFailureInvariant, false}, + {ErrAllocationAccountStale, AllocationFailureInvariant, false}, + {ErrAllocationAccountLive, AllocationFailureInvariant, false}, + {ErrAllocationGenerationSlots, AllocationFailureInvariant, false}, + {errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), ErrAllocationAccountInvariant), AllocationFailureInvariant, false}, + {errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), ErrAllocationAccountInvalid), AllocationFailureInvariant, false}, + {errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), ErrAllocationAccountSealed), AllocationFailureSealed, false}, + {ErrAllocationAdmissionSuspended, AllocationFailureSuspended, false}, + {errors.New("unrelated"), AllocationFailureNone, false}, + } + for _, testCase := range testCases { + require.Equal(t, testCase.reason, AllocationFailureReasonOf(testCase.err)) + require.Equal(t, testCase.retryable, IsRetryableAllocationCapacity(testCase.err)) + } + + // A terminal invariant dominates a joined underlying capacity error, so a + // failed terminal cleanup can never enter the pressure retry loop. + joined := errors.Join( + ErrAllocationAccountCapacity, + ErrAllocationAccountInvariant, + ) + require.Equal(t, AllocationFailureInvariant, AllocationFailureReasonOf(joined)) + require.False(t, IsRetryableAllocationCapacity(joined)) +} diff --git a/pkg/common/mpool/mpool.go b/pkg/common/mpool/mpool.go index ce33dd8c2cf30..af71a823fe594 100644 --- a/pkg/common/mpool/mpool.go +++ b/pkg/common/mpool/mpool.go @@ -228,15 +228,29 @@ type memHdr struct { poolId int64 allocSz int32 guard [3]uint8 - offHeap bool + kind uint8 } +const ( + memKindOnHeap uint8 = iota + memKindOffHeap + memKindAccountedOffHeap +) + func init() { if unsafe.Sizeof(memHdr{}) != kMemHdrSz { panic("memory header size assertion failed") } } +func (pHdr memHdr) isOffHeap() bool { + return pHdr.kind != memKindOnHeap +} + +func (pHdr memHdr) isAccounted() bool { + return pHdr.kind == memKindAccountedOffHeap +} + func (pHdr *memHdr) SetGuard() { pHdr.guard[0] = 0xDE pHdr.guard[1] = 0xAD @@ -324,6 +338,7 @@ type MPool struct { noLock bool ptrs map[unsafe.Pointer]memHdr + leases map[unsafe.Pointer]allocationLease } const ( @@ -334,14 +349,38 @@ const ( func (mp *MPool) recordPtrHdr(ptr unsafe.Pointer, pHdr memHdr) error { if !mp.noLock { return gRecordPtr(ptr, pHdr) - } else { - _, ok := mp.ptrs[ptr] - if ok { - return moerr.NewInternalErrorNoCtx("ptr already recorded") - } - mp.ptrs[ptr] = pHdr - return nil } + if _, ok := mp.ptrs[ptr]; ok { + return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + mp.ptrs[ptr] = pHdr + return nil +} + +func (mp *MPool) recordAccountedPtrMetadata( + ptr unsafe.Pointer, + pHdr memHdr, + lease allocationLease, +) error { + if !mp.noLock { + return gRecordAccountedPtrMetadata( + ptr, + pHdr, + lease, + ) + } + if _, ok := mp.ptrs[ptr]; ok { + return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + if _, ok := mp.leases[ptr]; ok { + return moerr.NewInternalErrorNoCtx("account lease already recorded") + } + if mp.leases == nil { + mp.leases = make(map[unsafe.Pointer]allocationLease) + } + mp.ptrs[ptr] = pHdr + mp.leases[ptr] = lease + return nil } func (mp *MPool) getPtrHdr(ptr unsafe.Pointer) (memHdr, bool) { @@ -352,25 +391,81 @@ func (mp *MPool) getPtrHdr(ptr unsafe.Pointer) (memHdr, bool) { return hdr, ok } -func (mp *MPool) removePtrHdr(ptr unsafe.Pointer) (memHdr, bool) { +func (mp *MPool) getPtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { if !mp.noLock { - return gRemovePtr(ptr) - } else { - hdr, ok := mp.ptrs[ptr] + return gGetPtrMetadata(ptr, lease) + } + hdr, ok := mp.ptrs[ptr] + if !ok { + return memHdr{}, false + } + if !hdr.isAccounted() { + return hdr, true + } + accountedLease, hasLease := mp.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) + } + *lease = accountedLease + return hdr, true +} + +func (mp *MPool) removePtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { + if !mp.noLock { + return gRemovePtrMetadata(ptr, lease) + } + hdr, ok := mp.ptrs[ptr] + if !ok { + if _, hasLease := mp.leases[ptr]; hasLease { + panic(moerr.NewInternalErrorNoCtx( + "account lease exists without allocation header", + )) + } + return memHdr{}, false + } + if !hdr.isAccounted() { delete(mp.ptrs, ptr) - return hdr, ok + return hdr, true + } + accountedLease, hasLease := mp.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) } + delete(mp.ptrs, ptr) + delete(mp.leases, ptr) + *lease = accountedLease + return hdr, true } func (mp *MPool) deallocateAllPtrs() { for ptr, hdr := range mp.ptrs { - if hdr.offHeap { + lease, hasLease := mp.leases[ptr] + if hdr.isAccounted() != hasLease { + panic(moerr.NewInternalErrorNoCtx( + "allocation header and account lease disagree during teardown", + )) + } + if hdr.isOffHeap() { sz := int(hdr.allocSz) profileRecordFree(uintptr(ptr), int64(sz)) simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) + if hasLease { + lease.release(uint64(sz)) + } } } mp.ptrs = nil + mp.leases = nil } func (mp *MPool) EnableDetailRecording() { @@ -403,33 +498,23 @@ func (mp *MPool) Cap() int64 { return mp.cap } -const ( - xxxIWouldRatherUseAfterFreeCrashLaterThanLeak = true -) - func (mp *MPool) destroy() { - if mp.stats.NumAlloc.Load() < mp.stats.NumFree.Load() { - // this is a memory leak, + liveBytes := mp.stats.NumCurrBytes.Load() + if liveBytes != 0 { logutil.Errorf("mp error: %s", mp.stats.Report("")) - - // here we MUST free all the memories allocated by this mpool. - // otherwise it is a memory leak. Whoever still holds - // a pointer of this mpool is a bug (the cross pool case). - // - // We are so messed up because the cross pool free. - // If a pointer is handed out to someone else and we free here - // it will be a use after free. We risk a crash or a leak. - // Eitherway we are screwed. - if xxxIWouldRatherUseAfterFreeCrashLaterThanLeak { + if mp.noLock { + // A noLock pool exclusively owns its local pointer metadata, so its + // teardown is the physical deallocation event. + nfree := mp.stats.NumAlloc.Load() - mp.stats.NumFree.Load() mp.deallocateAllPtrs() + mp.stats.RecordManyFrees(mp.tag, nfree, liveBytes) + globalStats.RecordManyFrees(mp.tag, nfree, liveBytes) + mp.resource.recordFree(liveBytes) } + // A normal pool may have handed allocations to another owner. Its + // global pointer metadata and account lease remain authoritative until + // a later physical Free, including after this pool is unregistered. } - - // Here we just compensate whatever left over in mp.stats - // into globalStats. - globalStats.RecordManyFrees(mp.tag, - mp.stats.NumAlloc.Load()-mp.stats.NumFree.Load(), - mp.stats.NumCurrBytes.Load()) } // New a MPool. Tag is user supplied, used for debugging/diagnostics. @@ -628,8 +713,9 @@ var globalPools sync.Map const numPtrShards = 128 type ptrShard struct { - mu sync.Mutex - m map[unsafe.Pointer]memHdr + mu sync.Mutex + m map[unsafe.Pointer]memHdr + leases map[unsafe.Pointer]allocationLease } var globalPtrShards [numPtrShards]ptrShard @@ -667,14 +753,55 @@ func GlobalCap() int64 { var CapLimit = math.MaxInt32 // 2GB - 1 +func maxAllocationSize() int64 { + return int64(CapLimit) - kMemHdrSz +} + func (mp *MPool) Alloc(sz int, offHeap bool) ([]byte, error) { detailk := mp.getDetailK() return mp.allocWithDetailK(detailk, int64(sz), offHeap) } +// AllocAccounted allocates off-heap memory owned by account. Owner and site +// are bounded diagnostics; zero is invalid. Existing unaccounted callers keep +// using Alloc. +func (mp *MPool) AllocAccounted( + sz int, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, +) ([]byte, error) { + return mp.AllocAccountedWithCapacityClass( + sz, + account, + owner, + site, + AllocationCapacityClassDefault, + ) +} + +// AllocAccountedWithCapacityClass allocates with an execution-local capacity +// controller while preserving the same physical allocation provenance. +func (mp *MPool) AllocAccountedWithCapacityClass( + sz int, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, + capacityClass AllocationCapacityClass, +) ([]byte, error) { + detailk := mp.getDetailK() + request := allocationAccountRequest{ + account: account, + owner: owner, + site: site, + capacityClass: capacityClass, + } + return mp.allocAccountedWithDetailK(detailk, int64(sz), request) +} + func (mp *MPool) allocWithDetailK(detailk string, sz int64, offHeap bool) ([]byte, error) { // reject unexpected alloc size. - if sz < 0 || sz > int64(CapLimit)-kMemHdrSz { + if sz < 0 || sz > maxAllocationSize() { logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) return nil, moerr.NewInternalErrorNoCtxf("mpool memory allocation exceed limit with requested size %d", sz) } @@ -684,48 +811,211 @@ func (mp *MPool) allocWithDetailK(detailk string, sz int64, offHeap bool) ([]byt return mp.alloc(detailk, sz, offHeap) } -func (mp *MPool) alloc(detailk string, sz int64, offHeap bool) ([]byte, error) { +func (mp *MPool) allocAccountedWithDetailK( + detailk string, + sz int64, + request allocationAccountRequest, +) ([]byte, error) { + if err := request.validate(); err != nil { + return nil, allocationAccountSiteError(request, err) + } + // reject unexpected alloc size. + if sz < 0 || sz > maxAllocationSize() { + logutil.Errorf("mpool memory allocation exceed limit with requested size %d: %s", sz, string(debug.Stack())) + return nil, allocationAccountSiteError(request, wrapAllocationAccountError( + ErrAllocationAllocatorLimit, + "requested=%d maximum=%d", + sz, + maxAllocationSize(), + )) + } + if sz == 0 { + return nil, nil + } + result, err := mp.allocAccounted(detailk, sz, request) + if err != nil { + return nil, allocationAccountSiteError(request, err) + } + return result, nil +} + +func allocationAccountSiteError( + request allocationAccountRequest, + err error, +) error { + if IsMPoolCapacityFailure(err) { + // Keep a direct MO error at the operator boundary: remote pipeline + // encoding preserves only direct *moerr.Error values. The ownership + // dimensions are folded into the serialized message instead. + return moerr.NewMPoolCapacityNoCtxf( + "allocation owner=%d site=%d: %s", + request.owner, + request.site, + err.Error(), + ) + } + return prefixAllocationAccountError( + err, + "allocation owner=%d site=%d", + request.owner, + request.site, + ) +} + +func (mp *MPool) alloc( + detailk string, + sz int64, + offHeap bool, +) ([]byte, error) { var bs []byte var err error hdr := memHdr{ poolId: mp.id, allocSz: int32(sz), - offHeap: offHeap, + } + if offHeap { + hdr.kind = memKindOffHeap } hdr.SetGuard() if offHeap { gcurr := globalStats.RecordAlloc("global", sz) if gcurr > GlobalCap() { - // compensate global globalStats.RecordFree("global", sz) - return nil, moerr.NewOOMNoCtx() + return nil, moerr.NewMPoolCapacityNoCtxf( + "global cap exceeded while allocating %d bytes", sz) } mycurr := mp.stats.RecordAlloc(mp.tag, sz) if mycurr > mp.Cap() { - // compensate both global and my mp.stats.RecordFree(mp.tag, sz) globalStats.RecordFree("global", sz) - return nil, moerr.NewInternalErrorNoCtxf("mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) + return nil, moerr.NewMPoolCapacityNoCtxf( + "mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) } bs, err = simpleCAllocator().Allocate(uint64(sz)) if err != nil { - panic(err) + mp.stats.RecordFree(mp.tag, sz) + globalStats.RecordFree("global", sz) + return nil, moerr.NewMPoolCapacityNoCtxf( + "physical allocator rejected %d bytes: %v", sz, err) } + } else { + bs = make([]byte, sz) + } + + ptr := unsafe.Pointer(&bs[0]) + if err = mp.recordPtrHdr(ptr, hdr); err != nil { + if offHeap { + simpleCAllocator().Deallocate(bs, uint64(sz)) + mp.stats.RecordFree(mp.tag, sz) + globalStats.RecordFree("global", sz) + } + return nil, err + } + if offHeap { mp.recordResourcePeak(mp.resource.recordAlloc(sz)) if mp.details != nil { mp.details.recordAlloc(detailk, sz) } - } else { - bs = make([]byte, sz) + profileRecordAlloc(3, uintptr(ptr), sz) } + return bs, nil +} - // always record the ptr, offHeap or not. - mp.recordPtrHdr(unsafe.Pointer(&bs[0]), hdr) - if offHeap { - profileRecordAlloc(3, uintptr(unsafe.Pointer(&bs[0])), sz) +func (mp *MPool) allocAccounted( + detailk string, + sz int64, + request allocationAccountRequest, +) ([]byte, error) { + var bs []byte + var err error + accountHeld := false + metadataHeld := false + globalHeld := false + poolHeld := false + physicalHeld := false + published := false + defer func() { + if published { + return + } + if physicalHeld { + simpleCAllocator().Deallocate(bs, uint64(sz)) + } + if poolHeld { + mp.stats.RecordFree(mp.tag, sz) + } + if globalHeld { + globalStats.RecordFree("global", sz) + } + if metadataHeld { + request.account.registry.releaseMetadata() + } + if accountHeld { + request.account.releaseWithCapacityClass(uint64(sz), request.capacityClass) + } + }() + + if err = request.account.acquireWithCapacityClass( + uint64(sz), + request.capacityClass, + ); err != nil { + return nil, err + } + accountHeld = true + if err = request.account.registry.reserveMetadata(); err != nil { + return nil, err + } + metadataHeld = true + + hdr := memHdr{ + poolId: mp.id, + allocSz: int32(sz), + kind: memKindAccountedOffHeap, } + hdr.SetGuard() + lease := allocationLease{ + account: request.account, + owner: request.owner, + site: request.site, + profiled: ProfilingEnabled(), + capacityClass: request.capacityClass, + } + + gcurr := globalStats.RecordAlloc("global", sz) + globalHeld = true + if gcurr > GlobalCap() { + return nil, moerr.NewMPoolCapacityNoCtxf( + "global cap exceeded while allocating %d bytes", sz) + } + mycurr := mp.stats.RecordAlloc(mp.tag, sz) + poolHeld = true + if mycurr > mp.Cap() { + return nil, moerr.NewMPoolCapacityNoCtxf( + "mpool out of space, alloc %d bytes, cap %d", sz, mp.cap) + } + bs, err = simpleCAllocator().Allocate(uint64(sz)) + if err != nil { + return nil, moerr.NewMPoolCapacityNoCtxf( + "physical allocator rejected %d bytes: %v", sz, err) + } + physicalHeld = true + + ptr := unsafe.Pointer(&bs[0]) + if err = mp.recordAccountedPtrMetadata( + ptr, + hdr, + lease, + ); err != nil { + return nil, err + } + published = true + mp.recordResourcePeak(mp.resource.recordAlloc(sz)) + if mp.details != nil { + mp.details.recordAlloc(detailk, sz) + } + profileRecordAccountedAlloc(lease, sz) return bs, nil } @@ -744,7 +1034,8 @@ func (mp *MPool) freeWithDetailK(detailk string, bs []byte) { } func (mp *MPool) freePtr(detailk string, ptr unsafe.Pointer) { - hdr, ok := mp.removePtrHdr(ptr) + var lease allocationLease + hdr, ok := mp.removePtrMetadata(ptr, &lease) if !ok { // this is a double free. panic(moerr.NewInternalErrorNoCtx("invalid ptr, double free")) @@ -761,25 +1052,42 @@ func (mp *MPool) freePtr(detailk string, ptr unsafe.Pointer) { // Call profileRecordFree and the full globalStats.RecordFree // (not just NumCurrBytes.Add) so NumFree/NumFreeBytes stay // consistent with freePtrInternal. - if hdr.offHeap { + if hdr.isOffHeap() { sz := int64(hdr.allocSz) - profileRecordFree(uintptr(ptr), sz) + if hdr.isAccounted() { + profileRecordAccountedFree(lease, sz) + } else { + profileRecordFree(uintptr(ptr), sz) + } globalStats.RecordFree("global", sz) simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) + if hdr.isAccounted() { + lease.release(uint64(sz)) + } } } else { owner := otherPool.(*MPool) owner.resource.crossPoolFree.Add(1) - owner.freePtrInternal(detailk, ptr, hdr) + owner.freePtrInternal(detailk, ptr, hdr, lease) } return } - mp.freePtrInternal(detailk, ptr, hdr) + mp.freePtrInternal(detailk, ptr, hdr, lease) } -func (mp *MPool) freePtrInternal(detailk string, ptr unsafe.Pointer, hdr memHdr) { - if !hdr.offHeap { +func (mp *MPool) freePtrInternal( + detailk string, + ptr unsafe.Pointer, + hdr memHdr, + lease allocationLease, +) { + if !hdr.isOffHeap() { + if hdr.isAccounted() { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation is not off-heap", + )) + } return } sz := int64(hdr.allocSz) @@ -789,7 +1097,11 @@ func (mp *MPool) freePtrInternal(detailk string, ptr unsafe.Pointer, hdr memHdr) bs[i] = 0xDD } } - profileRecordFree(uintptr(ptr), sz) + if hdr.isAccounted() { + profileRecordAccountedFree(lease, sz) + } else { + profileRecordFree(uintptr(ptr), sz) + } mp.stats.RecordFree(mp.tag, sz) globalStats.RecordFree("global", sz) mp.resource.recordFree(sz) @@ -798,6 +1110,9 @@ func (mp *MPool) freePtrInternal(detailk string, ptr unsafe.Pointer, hdr memHdr) } simpleCAllocator().Deallocate(unsafe.Slice((*byte)(ptr), sz), uint64(sz)) + if hdr.isAccounted() { + lease.release(uint64(sz)) + } } func (mp *MPool) reAllocWithDetailK(detailk string, old []byte, sz int64, offHeap bool, bufferMore bool) ([]byte, error) { @@ -814,7 +1129,36 @@ func (mp *MPool) reAllocWithDetailK(detailk string, old []byte, sz int64, offHea } } - ret, err := mp.allocWithDetailK(detailk, int64(newSz), offHeap) + var request *allocationAccountRequest + if ptr := unsafe.Pointer(unsafe.SliceData(old)); ptr != nil { + var lease allocationLease + hdr, ok := mp.getPtrMetadata(ptr, &lease) + if !ok { + return nil, moerr.NewInternalErrorNoCtx( + "invalid grow pointer: allocation metadata not found", + ) + } + if hdr.isAccounted() { + if !offHeap { + return nil, ErrAllocationAccountMismatch + } + accounted := allocationAccountRequest{ + account: lease.account, + owner: lease.owner, + site: lease.site, + capacityClass: lease.capacityClass, + } + request = &accounted + } + } + + var ret []byte + var err error + if request != nil { + ret, err = mp.allocAccountedWithDetailK(detailk, newSz, *request) + } else { + ret, err = mp.allocWithDetailK(detailk, newSz, offHeap) + } if err != nil { return nil, err } @@ -847,7 +1191,7 @@ func (mp *MPool) Grow2(old []byte, old2 []byte, sz int, offHeap bool) ([]byte, e // ReallocZero is like Realloc, but it clears the memory. func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { detailk := mp.getDetailK() - if sz < 0 || sz > CapLimit-kMemHdrSz { + if int64(sz) < 0 || int64(sz) > maxAllocationSize() { return nil, moerr.NewInternalErrorNoCtxf( "mpool memory allocation exceed limit with requested size %d", sz, @@ -870,9 +1214,10 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { oldptr := unsafe.Pointer(unsafe.SliceData(old)) var hdr memHdr + var lease allocationLease var ok bool if oldptr != nil { - hdr, ok = mp.getPtrHdr(oldptr) + hdr, ok = mp.getPtrMetadata(oldptr, &lease) } if !ok { if len(old) != 0 || cap(old) != 0 { @@ -913,10 +1258,34 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { return resized, nil } + if hdr.isAccounted() { + if !offHeap { + return nil, ErrAllocationAccountMismatch + } + request := allocationAccountRequest{ + account: lease.account, + owner: lease.owner, + site: lease.site, + capacityClass: lease.capacityClass, + } + replacement, err := mp.allocAccountedWithDetailK( + detailk, + int64(sz), + request, + ) + if err != nil { + return nil, err + } + copy(replacement, fullAllocation[:oldLength]) + clear(replacement[oldLength:]) + mp.freeWithDetailK(detailk, fullAllocation) + return replacement, nil + } + // Only resize in place when the source and destination are off-heap and // owned by this pool. Other provenance/ownership transitions use the normal // allocate-copy-free path so accounting and cross-pool cleanup stay correct. - if !hdr.offHeap || !offHeap || hdr.poolId != mp.id { + if !hdr.isOffHeap() || !offHeap || hdr.poolId != mp.id { return mp.reAllocWithDetailK( detailk, fullAllocation[:oldLength], @@ -931,12 +1300,14 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { // retain the new-size charge and release only the old-size charge on success. if globalStats.RecordAlloc("global", int64(sz)) > GlobalCap() { globalStats.RecordFree("global", int64(sz)) - return nil, moerr.NewOOMNoCtx() + return nil, moerr.NewMPoolCapacityNoCtxf( + "global cap exceeded while reallocating %d bytes", sz) } if mp.stats.RecordAlloc(mp.tag, int64(sz)) > mp.Cap() { mp.stats.RecordFree(mp.tag, int64(sz)) globalStats.RecordFree("global", int64(sz)) - return nil, moerr.NewInternalErrorNoCtxf("mpool out of space, realloc %d bytes, cap %d", sz, mp.cap) + return nil, moerr.NewMPoolCapacityNoCtxf( + "mpool out of space, realloc %d bytes, cap %d", sz, mp.cap) } newbs, err := simpleCAllocator().ReallocZero( @@ -947,19 +1318,26 @@ func (mp *MPool) ReallocZero(old []byte, sz int, offHeap bool) ([]byte, error) { if err != nil { mp.stats.RecordFree(mp.tag, int64(sz)) globalStats.RecordFree("global", int64(sz)) - return nil, err + return nil, moerr.NewMPoolCapacityNoCtxf( + "physical allocator rejected realloc to %d bytes: %v", sz, err) } newptr := unsafe.Pointer(&newbs[0]) - removedHdr, removed := mp.removePtrHdr(oldptr) + var removedLease allocationLease + removedHdr, removed := mp.removePtrMetadata(oldptr, &removedLease) if !removed || removedHdr != hdr { panic(moerr.NewInternalErrorNoCtx( "allocation metadata changed during realloc", )) } + if removedHdr.isAccounted() || removedLease.account != nil { + panic(moerr.NewInternalErrorNoCtx( + "unaccounted realloc removed an account lease", + )) + } newHdr := memHdr{ poolId: hdr.poolId, allocSz: int32(sz), - offHeap: true, + kind: memKindOffHeap, } newHdr.SetGuard() if err := mp.recordPtrHdr(newptr, newHdr); err != nil { @@ -991,6 +1369,70 @@ func MakeSlice[T any](n int, mp *MPool, offHeap bool) ([]T, error) { return makeSliceWithCapWithDetailK[T](detailk, n, n, mp, offHeap) } +// MakeSliceAccounted allocates an off-heap typed slice whose physical +// allocation is owned by account. FreeSlice releases the resulting charge. +func MakeSliceAccounted[T any]( + n int, + mp *MPool, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, +) ([]T, error) { + return MakeSliceAccountedWithCapacityClass[T]( + n, + mp, + account, + owner, + site, + AllocationCapacityClassDefault, + ) +} + +// MakeSliceAccountedWithCapacityClass is the typed-slice form of +// AllocAccountedWithCapacityClass. +func MakeSliceAccountedWithCapacityClass[T any]( + n int, + mp *MPool, + account *AllocationAccount, + owner AllocationOwner, + site AllocationSite, + capacityClass AllocationCapacityClass, +) ([]T, error) { + if n < 0 { + return nil, ErrAllocationAccountInvalid + } + if n == 0 { + return nil, nil + } + var value T + elementSize := unsafe.Sizeof(value) + maxSize := maxAllocationSize() + if elementSize == 0 || + maxSize <= 0 || + uint64(n) > uint64(maxSize)/uint64(elementSize) { + return nil, wrapAllocationAccountError( + ErrAllocationAllocatorLimit, + "elements=%d element-size=%d maximum=%d", + n, + elementSize, + maxSize, + ) + } + size := int(uint64(n) * uint64(elementSize)) + bs, err := mp.AllocAccountedWithCapacityClass( + size, + account, + owner, + site, + capacityClass, + ) + if err != nil { + return nil, err + } + values := unsafe.Slice((*T)(unsafe.Pointer(&bs[0])), n) + return values[:n:n], nil +} + func MakeSliceArgs[T any](mp *MPool, offHeap bool, args ...T) ([]T, error) { detailk := mp.getDetailK() ret, err := makeSliceWithCapWithDetailK[T](detailk, len(args), len(args), mp, offHeap) @@ -1005,6 +1447,7 @@ func FreeSlice[T any](mp *MPool, bs []T) { if cap(bs) == 0 { return } + bs = bs[:1] detailk := mp.getDetailK() mp.freePtr(detailk, unsafe.Pointer(&bs[0])) } @@ -1083,15 +1526,39 @@ func init() { } } -func gRecordPtr(ptr unsafe.Pointer, hdr memHdr) error { +func gRecordPtr( + ptr unsafe.Pointer, + hdr memHdr, +) error { + shard := getPtrShard(ptr) + shard.mu.Lock() + defer shard.mu.Unlock() + if _, ok := shard.m[ptr]; ok { + return moerr.NewInternalErrorNoCtx("ptr already recorded") + } + shard.m[ptr] = hdr + return nil +} + +func gRecordAccountedPtrMetadata( + ptr unsafe.Pointer, + hdr memHdr, + lease allocationLease, +) error { shard := getPtrShard(ptr) shard.mu.Lock() defer shard.mu.Unlock() - _, ok := shard.m[ptr] - if ok { + if _, ok := shard.m[ptr]; ok { return moerr.NewInternalErrorNoCtx("ptr already recorded") } + if _, ok := shard.leases[ptr]; ok { + return moerr.NewInternalErrorNoCtx("account lease already recorded") + } + if shard.leases == nil { + shard.leases = make(map[unsafe.Pointer]allocationLease) + } shard.m[ptr] = hdr + shard.leases[ptr] = lease return nil } @@ -1103,13 +1570,60 @@ func gGetPtr(ptr unsafe.Pointer) (memHdr, bool) { return hdr, ok } -func gRemovePtr(ptr unsafe.Pointer) (memHdr, bool) { +func gGetPtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { shard := getPtrShard(ptr) shard.mu.Lock() defer shard.mu.Unlock() hdr, ok := shard.m[ptr] + if !ok { + return memHdr{}, false + } + if !hdr.isAccounted() { + return hdr, true + } + accountedLease, hasLease := shard.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) + } + *lease = accountedLease + return hdr, true +} + +func gRemovePtrMetadata( + ptr unsafe.Pointer, + lease *allocationLease, +) (memHdr, bool) { + shard := getPtrShard(ptr) + shard.mu.Lock() + defer shard.mu.Unlock() + hdr, ok := shard.m[ptr] + if !ok { + if _, hasLease := shard.leases[ptr]; hasLease { + panic(moerr.NewInternalErrorNoCtx( + "account lease exists without allocation header", + )) + } + return memHdr{}, false + } + if !hdr.isAccounted() { + delete(shard.m, ptr) + return hdr, true + } + accountedLease, hasLease := shard.leases[ptr] + if !hasLease { + panic(moerr.NewInternalErrorNoCtx( + "accounted allocation has no account lease", + )) + } delete(shard.m, ptr) - return hdr, ok + delete(shard.leases, ptr) + *lease = accountedLease + return hdr, true } // alignUp rounds n up to a multiple of a. a must be a power of 2. @@ -1144,7 +1658,9 @@ func roundupsize(size int64) int64 { // request. Callers which must reserve memory before growing a slice use this // helper so admission and allocation share the same growth calculation. func GrowCapacity(oldCap int64, requiredSize int64) (int64, bool) { - if oldCap < 0 || requiredSize < 0 { + maxCapacity := maxAllocationSize() + if oldCap < 0 || requiredSize < 0 || + oldCap > maxCapacity || requiredSize > maxCapacity { return 0, false } if requiredSize <= oldCap { @@ -1175,8 +1691,8 @@ func GrowCapacity(oldCap int64, requiredSize int64) (int64, bool) { if newcap < requiredSize { return 0, false } - if newcap > int64(CapLimit) && requiredSize <= int64(CapLimit) { - newcap = int64(CapLimit) + if newcap > maxCapacity { + newcap = maxCapacity } return newcap, true } diff --git a/pkg/common/mpool/mpool_profile.go b/pkg/common/mpool/mpool_profile.go index 8736697c458fc..fd2f93a2fa7d7 100644 --- a/pkg/common/mpool/mpool_profile.go +++ b/pkg/common/mpool/mpool_profile.go @@ -15,6 +15,7 @@ package mpool import ( + "fmt" "sync" "sync/atomic" @@ -23,8 +24,9 @@ import ( var profilingEnabled atomic.Bool -// EnableProfiling turns on per-allocation stack tracking for off-heap mpool -// allocations. Tracked allocations appear in the malloc profiler output. +// EnableProfiling turns on tracking for off-heap mpool allocations. Ordinary +// allocations are grouped by sampled call stack; accounted allocations are +// grouped by their explicit owner/site provenance. func EnableProfiling() { profilingEnabled.Store(true) } // DisableProfiling turns off per-allocation stack tracking. @@ -43,6 +45,11 @@ type profileShard struct { var globalProfileShards [numProfileShards]profileShard +// Accounted allocations already carry stable, bounded provenance. Reusing one +// synthetic sample per owner/site avoids collecting and hashing the same +// runtime stack for every vector growth in a hash build. +var accountedProfileSamples [AllocationOwnerMax + 1][256]atomic.Pointer[malloc.HeapSampleValues] + func init() { for i := range globalProfileShards { globalProfileShards[i].m = make(map[uintptr]*malloc.HeapSampleValues, 64) @@ -98,6 +105,45 @@ func profileRecordFree(ptr uintptr, sz int64) { } } +func accountedProfileSample( + owner AllocationOwner, + site AllocationSite, +) *malloc.HeapSampleValues { + slot := &accountedProfileSamples[owner][site] + if values := slot.Load(); values != nil { + return values + } + values := malloc.GlobalProfiler().SampleNamed(fmt.Sprintf( + "| mpool accounted owner=%d site=%d |", + owner, + site, + )) + if slot.CompareAndSwap(nil, values) { + return values + } + return slot.Load() +} + +func profileRecordAccountedAlloc(lease allocationLease, sz int64) { + if !lease.profiled { + return + } + values := accountedProfileSample(lease.owner, lease.site) + values.Bytes.Allocated.Add(uint64(sz)) + values.Objects.Allocated.Add(1) + values.Bytes.Inuse.Add(sz) + values.Objects.Inuse.Add(1) +} + +func profileRecordAccountedFree(lease allocationLease, sz int64) { + if !lease.profiled { + return + } + values := accountedProfileSample(lease.owner, lease.site) + values.Bytes.Inuse.Add(-sz) + values.Objects.Inuse.Add(-1) +} + func profileRecordRealloc(skip int, oldPtr, newPtr uintptr, oldSz, newSz int64) { if !profilingEnabled.Load() { return diff --git a/pkg/common/mpool/mpool_profile_test.go b/pkg/common/mpool/mpool_profile_test.go index 8867cda1e013d..c95750f6a0126 100644 --- a/pkg/common/mpool/mpool_profile_test.go +++ b/pkg/common/mpool/mpool_profile_test.go @@ -103,3 +103,85 @@ func TestProfileWritable(t *testing.T) { mp.Free(bs) } + +func TestAccountedProfileUsesProvenanceAcrossProfilingToggle(t *testing.T) { + DisableProfiling() + defer DisableProfiling() + registry, account := newTestAllocationAccount(t, 1024, 1) + mp := MustNew("accounted-profile") + defer DeleteMPool(mp) + values := accountedProfileSample(testAllocationOwner, testAllocationSite) + before := values.Values() + trackedBefore := ProfileTrackedCount() + + EnableProfiling() + buffer, err := mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + require.NoError(t, err) + require.Equal(t, trackedBefore, ProfileTrackedCount(), + "accounted provenance does not need a per-pointer stack entry") + afterAlloc := values.Values() + require.Equal(t, int64(1), afterAlloc[0]-before[0]) + require.Equal(t, int64(64), afterAlloc[1]-before[1]) + require.Equal(t, int64(1), afterAlloc[2]-before[2]) + require.Equal(t, int64(64), afterAlloc[3]-before[3]) + + // The allocation lease remembers whether it was profiled, so disabling + // collection cannot strand an existing in-use sample. + DisableProfiling() + mp.Free(buffer) + afterFree := values.Values() + require.Equal(t, before[2], afterFree[2]) + require.Equal(t, before[3], afterFree[3]) + finalizeTestAllocationAccount(t, registry, account) +} + +func BenchmarkProfileAllocFree(b *testing.B) { + EnableProfiling() + defer DisableProfiling() + for _, accounted := range []bool{false, true} { + name := "stack" + if accounted { + name = "accounted-provenance" + } + b.Run(name, func(b *testing.B) { + mp := MustNew("profile-benchmark") + defer DeleteMPool(mp) + var registry *AllocationAccountRegistry + var account *AllocationAccount + if accounted { + registry, account = newTestAllocationAccount(b, 1<<60, 1) + // Initialize the bounded owner/site sample outside the measured loop. + accountedProfileSample(testAllocationOwner, testAllocationSite) + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + var buffer []byte + var err error + if accounted { + buffer, err = mp.AllocAccounted( + 64, + account, + testAllocationOwner, + testAllocationSite, + ) + } else { + buffer, err = mp.Alloc(64, true) + } + if err != nil { + b.Fatal(err) + } + mp.Free(buffer) + } + b.StopTimer() + if accounted { + finalizeTestAllocationAccount(b, registry, account) + } + }) + } +} diff --git a/pkg/common/mpool/mpool_test.go b/pkg/common/mpool/mpool_test.go index 133eacece5833..565a2af9967cc 100644 --- a/pkg/common/mpool/mpool_test.go +++ b/pkg/common/mpool/mpool_test.go @@ -381,6 +381,13 @@ func TestGrowCapacityValidation(t *testing.T) { capacity, ok := GrowCapacity(128, 64) require.True(t, ok) require.Equal(t, int64(128), capacity) + + maxCapacity := maxAllocationSize() + capacity, ok = GrowCapacity(maxCapacity, maxCapacity) + require.True(t, ok) + require.Equal(t, maxCapacity, capacity) + _, ok = GrowCapacity(maxCapacity, maxCapacity+1) + require.False(t, ok) } func TestUseMalloc(t *testing.T) { @@ -534,6 +541,33 @@ func TestCrossPoolFreeOnHeap(t *testing.T) { DeleteMPool(mp2) } +func TestMPoolTeardownTracksPhysicalLifetime(t *testing.T) { + t.Run("normal-pool-late-free", func(t *testing.T) { + owner := MustNew("teardown-normal-owner") + other := MustNew("teardown-normal-other") + defer DeleteMPool(other) + + globalBefore := GlobalStats().NumCurrBytes.Load() + buffer, err := owner.Alloc(64, true) + require.NoError(t, err) + DeleteMPool(owner) + require.Equal(t, globalBefore+64, GlobalStats().NumCurrBytes.Load()) + + other.Free(buffer) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + }) + + t.Run("no-lock-pool-owns-teardown", func(t *testing.T) { + mp := MustNewNoLock("teardown-no-lock-owner") + globalBefore := GlobalStats().NumCurrBytes.Load() + _, err := mp.Alloc(64, true) + require.NoError(t, err) + + DeleteMPool(mp) + require.Equal(t, globalBefore, GlobalStats().NumCurrBytes.Load()) + }) +} + // TestDoubleFree tests that double free is detected and panics. func TestDoubleFree(t *testing.T) { mp := MustNew("double-free-test") @@ -812,7 +846,7 @@ func TestMPoolReallocZeroUsesRecordedSourceProvenance(t *testing.T) { hdr, ok := mp.getPtrHdr(unsafe.Pointer(unsafe.SliceData(resized))) require.True(t, ok) - require.Equal(t, testCase.targetOffHeap, hdr.offHeap) + require.Equal(t, testCase.targetOffHeap, hdr.isOffHeap()) require.Equal(t, int32(newSize), hdr.allocSz) mp.Free(resized) diff --git a/pkg/compare/arraycompare.go b/pkg/compare/arraycompare.go index 3ac0695297ba8..d6e9c4410ee2c 100644 --- a/pkg/compare/arraycompare.go +++ b/pkg/compare/arraycompare.go @@ -45,8 +45,10 @@ func (c arrayCompare) Copy(vecSrc, vecDst int, src, dst int64, proc *process.Pro } func (c arrayCompare) Compare(veci, vecj int, vi, vj int64) int { - n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) - n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) + n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) || + c.vs[veci].GetGrouping().Contains(uint64(vi)) + n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) || + c.vs[vecj].GetGrouping().Contains(uint64(vj)) cmp := nullsCompare(n0, n1, c.nullsLast) if cmp != 0 { return cmp - nullsCompareFlag diff --git a/pkg/compare/compare.go b/pkg/compare/compare.go index ce547f19da34b..e9c4539bbe57e 100644 --- a/pkg/compare/compare.go +++ b/pkg/compare/compare.go @@ -230,8 +230,10 @@ func (c *compare[T]) Set(idx int, vec *vector.Vector) { } func (c *compare[T]) Compare(veci, vecj int, vi, vj int64) int { - n0 := c.isConstNull[veci] || c.ns[veci].Contains(uint64(vi)) - n1 := c.isConstNull[vecj] || c.ns[vecj].Contains(uint64(vj)) + n0 := c.isConstNull[veci] || c.ns[veci].Contains(uint64(vi)) || + c.gs[veci].Contains(uint64(vi)) + n1 := c.isConstNull[vecj] || c.ns[vecj].Contains(uint64(vj)) || + c.gs[vecj].Contains(uint64(vj)) cmp := nullsCompare(n0, n1, c.nullsLast) if cmp != 0 { return cmp - nullsCompareFlag @@ -239,7 +241,11 @@ func (c *compare[T]) Compare(veci, vecj int, vi, vj int64) int { return c.cmp(c.xs[veci][vi], c.xs[vecj][vj]) } -func (c *compare[T]) Copy(vecSrc, vecDst int, src, dst int64, _ *process.Process) error { +func (c *compare[T]) Copy(vecSrc, vecDst int, src, dst int64, proc *process.Process) error { + if c.gs[vecSrc].Contains(uint64(src)) || + c.gs[vecDst].Contains(uint64(dst)) { + return c.vs[vecDst].Copy(c.vs[vecSrc], dst, src, proc.Mp()) + } if c.isConstNull[vecSrc] || c.ns[vecSrc].Contains(uint64(src)) { nulls.Add(c.ns[vecDst], uint64(dst)) } else { diff --git a/pkg/compare/grouping_compare_test.go b/pkg/compare/grouping_compare_test.go new file mode 100644 index 0000000000000..cd228246e4198 --- /dev/null +++ b/pkg/compare/grouping_compare_test.go @@ -0,0 +1,90 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package compare + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/stretchr/testify/require" +) + +func TestGroupingSentinelComparesAsNull(t *testing.T) { + mp := mpool.MustNewZero() + tests := []struct { + name string + typ types.Type + add func(*vector.Vector, bool) error + }{ + { + name: "fixed", + typ: types.T_int64.ToType(), + add: func(v *vector.Vector, isNull bool) error { + return vector.AppendFixed(v, int64(0), isNull, mp) + }, + }, + { + name: "varlen", + typ: types.T_varchar.ToType(), + add: func(v *vector.Vector, isNull bool) error { + return vector.AppendBytes(v, nil, isNull, mp) + }, + }, + { + name: "array", + typ: types.T_array_float32.ToType(), + add: func(v *vector.Vector, isNull bool) error { + return vector.AppendArray(v, []float32{0}, isNull, mp) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + grouping := vector.NewVec(test.typ) + ordinary := vector.NewVec(test.typ) + nullValue := vector.NewVec(test.typ) + t.Cleanup(func() { + grouping.Free(mp) + ordinary.Free(mp) + nullValue.Free(mp) + }) + require.NoError(t, test.add(grouping, false)) + require.NoError(t, test.add(ordinary, false)) + require.NoError(t, test.add(nullValue, true)) + grouping.GetGrouping().Add(0) + + for _, nullsLast := range []bool{false, true} { + for _, desc := range []bool{false, true} { + cmp := New(test.typ, desc, nullsLast) + cmp.Set(0, grouping) + cmp.Set(1, ordinary) + if nullsLast { + require.Positive(t, cmp.Compare(0, 1, 0, 0)) + require.Negative(t, cmp.Compare(1, 0, 0, 0)) + } else { + require.Negative(t, cmp.Compare(0, 1, 0, 0)) + require.Positive(t, cmp.Compare(1, 0, 0, 0)) + } + + cmp.Set(1, nullValue) + require.Zero(t, cmp.Compare(0, 1, 0, 0)) + } + } + }) + } +} diff --git a/pkg/compare/strcompare.go b/pkg/compare/strcompare.go index f74ddc2a6375c..ba155c435e2c6 100644 --- a/pkg/compare/strcompare.go +++ b/pkg/compare/strcompare.go @@ -46,8 +46,10 @@ func (c *strCompare) Copy(vecSrc, vecDst int, src, dst int64, proc *process.Proc } func (c *strCompare) Compare(veci, vecj int, vi, vj int64) int { - n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) - n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) + n0 := c.isConstNull[veci] || c.vs[veci].GetNulls().Contains(uint64(vi)) || + c.vs[veci].GetGrouping().Contains(uint64(vi)) + n1 := c.isConstNull[vecj] || c.vs[vecj].GetNulls().Contains(uint64(vj)) || + c.vs[vecj].GetGrouping().Contains(uint64(vj)) cmp := nullsCompare(n0, n1, c.nullsLast) if cmp != 0 { return cmp - nullsCompareFlag diff --git a/pkg/container/batch/allocation_account_test.go b/pkg/container/batch/allocation_account_test.go new file mode 100644 index 0000000000000..799605335da72 --- /dev/null +++ b/pkg/container/batch/allocation_account_test.go @@ -0,0 +1,804 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package batch + +import ( + "bytes" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/stretchr/testify/require" +) + +type testBatchAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + selection *vector.AllocationAccountSelection +} + +func newTestBatchAllocationAccount( + t *testing.T, + allocationSlots uint64, +) testBatchAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) + require.NoError(t, err) + account, err := registry.Open(16 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + return testBatchAllocationAccount{ + registry: registry, + account: account, + selection: selection, + } +} + +func finalizeTestBatchAllocationAccount( + t *testing.T, + state testBatchAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func newBatchAllocationTestSource( + t *testing.T, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) *Batch { + t.Helper() + bat := NewWithSchema( + true, + []string{"id", "value"}, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + if selection != nil { + require.NoError(t, bat.SetAllocationAccount(selection)) + } + for i := 0; i < 32; i++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(i), false, mp)) + require.NoError( + t, + vector.AppendBytes( + bat.Vecs[1], + []byte("batch allocation payload that is not inline"), + false, + mp, + ), + ) + } + bat.SetRowCount(32) + return bat +} + +func TestBatchAllocationAccountCloneDupAndWindow(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, state.selection) + sourceUsed := state.account.Snapshot().Used + require.NotZero(t, sourceUsed) + _, err := source.Clone(mp, false) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + + cloned, err := source.Clone(mp, true) + require.NoError(t, err) + require.Same(t, state.selection, cloned.AllocationAccountSelection()) + require.Greater(t, state.account.Snapshot().Used, sourceUsed) + cloned.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + dup, err := source.Dup(mp) + require.NoError(t, err) + require.Same(t, state.selection, dup.AllocationAccountSelection()) + dup.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + selectedColumns, err := source.CloneSelectedColumns( + []int{1}, + []string{"value"}, + mp, + ) + require.NoError(t, err) + require.Same(t, state.selection, selectedColumns.AllocationAccountSelection()) + selectedColumns.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + window, err := source.Window(4, 12) + require.NoError(t, err) + require.Same(t, state.selection, window.AllocationAccountSelection()) + for _, vec := range window.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + window.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + var encoded bytes.Buffer + data, err := source.MarshalBinaryWithBuffer(&encoded, true) + require.NoError(t, err) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NoError(t, decoded.UnmarshalFromReader(bytes.NewReader(data), mp)) + require.Same(t, state.selection, decoded.AllocationAccountSelection()) + for _, vec := range decoded.Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + decoded.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + aliasDecoded := NewWithSchema( + true, + source.Attrs, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, aliasDecoded.SetAllocationAccount(state.selection)) + require.NoError(t, aliasDecoded.UnmarshalBinaryWithAnyMp(data, mp)) + require.Same(t, state.selection, aliasDecoded.AllocationAccountSelection()) + for _, vec := range aliasDecoded.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + require.NoError( + t, + aliasDecoded.UnmarshalFromReader(bytes.NewReader(data), mp), + ) + for _, vec := range aliasDecoded.Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + require.Equal(t, source.RowCount(), aliasDecoded.RowCount()) + aliasDecoded.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchDupWithoutAllocationAccountCrossesStatementBoundary(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, state.selection) + sourceUsed := state.account.Snapshot().Used + require.Positive(t, sourceUsed) + + cloned, err := source.DupWithoutAllocationAccount(mp) + require.NoError(t, err) + require.Nil(t, cloned.AllocationAccountSelection()) + for _, vec := range cloned.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + require.Equal(t, source.RowCount(), cloned.RowCount()) + require.Equal( + t, + int64(0), + vector.GetFixedAtNoTypeCheck[int64](cloned.Vecs[0], 0), + ) + require.Equal( + t, + []byte("batch allocation payload that is not inline"), + cloned.Vecs[1].GetBytesAt(0), + ) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) + require.Equal( + t, + int64(0), + vector.GetFixedAtNoTypeCheck[int64](cloned.Vecs[0], 0), + ) + require.Equal( + t, + []byte("batch allocation payload that is not inline"), + cloned.Vecs[1].GetBytesAt(0), + ) + cloned.Clean(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestBatchAccountedReaderAcceptsBitmapCapacityBeyondLogicalRows(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + source.Vecs[0].GetNulls().Add(31) + source.Shrink([]int64{0, 1, 2, 3, 4}, false) + + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryTo(&encoded)) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NoError(t, decoded.UnmarshalFromReader(&encoded, mp)) + require.Equal(t, 5, decoded.RowCount()) + require.Equal( + t, + int64(0), + vector.GetFixedAtWithTypeCheck[int64](decoded.Vecs[0], 0), + ) + + decoded.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAccountedReaderPreservesRowsWhenVectorCountChanges(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryTo(&encoded)) + decoded := NewOffHeapWithSize(1) + decoded.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + decoded.SetRowCount(7) + require.NoError(t, decoded.UnmarshalFromReader(&encoded, mp)) + require.Equal(t, source.RowCount(), decoded.RowCount()) + require.Len(t, decoded.Vecs, 2) + require.Equal( + t, + int64(31), + vector.GetFixedAtWithTypeCheck[int64](decoded.Vecs[0], 31), + ) + + decoded.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchGroupingCodecRoundTrip(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + source.Vecs[0].GetGrouping().Add(1, 7, 31) + source.Vecs[1].GetGrouping().Add(2, 9) + source.ExtraBuf = bytes.Repeat([]byte("x"), 1<<20) + + var encoded bytes.Buffer + spillSize, err := source.MarshalBinaryWithGroupingSize() + require.NoError(t, err) + stableSize, err := source.MarshalBinarySize() + require.NoError(t, err) + require.Greater(t, stableSize-spillSize, len(source.ExtraBuf)/2) + require.NoError(t, source.MarshalBinaryWithGroupingTo(&encoded)) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NoError(t, decoded.UnmarshalFromReaderWithGrouping(&encoded, mp)) + require.Empty(t, decoded.Attrs) + require.Empty(t, decoded.ExtraBuf) + for i := range source.Vecs { + require.True(t, decoded.Vecs[i].GetGrouping().IsSame(source.Vecs[i].GetGrouping())) + } + withoutGrouping := newBatchAllocationTestSource(t, mp, nil) + encoded.Reset() + require.NoError(t, withoutGrouping.MarshalBinaryWithGroupingTo(&encoded)) + require.NoError(t, decoded.UnmarshalFromReaderWithGrouping(&encoded, mp)) + for _, vec := range decoded.Vecs { + require.True(t, vec.GetGrouping().IsEmpty()) + } + + decoded.Clean(mp) + withoutGrouping.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchGroupingCodecRejectsStableMetadataBeforePayloadAllocation(t *testing.T) { + state := newTestBatchAllocationAccount(t, 32) + mp := mpool.MustNewZero() + + for _, test := range []struct { + name string + attrs []string + extra []byte + want string + }{ + { + name: "attributes", + attrs: []string{string(bytes.Repeat([]byte("a"), 1<<20))}, + want: "attributes are not allowed", + }, + { + name: "extra buffer", + extra: bytes.Repeat([]byte("x"), 1<<20), + want: "extra buffer is not allowed", + }, + } { + t.Run(test.name, func(t *testing.T) { + source := NewWithSize(0) + source.Attrs = test.attrs + source.ExtraBuf = test.extra + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryTo(&encoded)) + + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.ErrorContains( + t, + decoded.UnmarshalFromReaderWithGrouping(&encoded, mp), + test.want, + ) + require.Zero(t, state.account.Snapshot().Used) + decoded.Clean(mp) + }) + } + + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchGroupingCodecRejectsMismatchedRowCount(t *testing.T) { + state := newTestBatchAllocationAccount(t, 32) + mp := mpool.MustNewZero() + source := NewWithSize(1) + source.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(1), false, mp)) + source.SetRowCount(2) + + var encoded bytes.Buffer + require.NoError(t, source.MarshalBinaryWithGroupingTo(&encoded)) + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.ErrorContains( + t, + decoded.UnmarshalFromReaderWithGrouping(&encoded, mp), + "vector length does not match row count", + ) + + decoded.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAccountedReaderRejectsInvalidLengthsBeforeAllocation(t *testing.T) { + state := newTestBatchAllocationAccount(t, 16) + mp := mpool.MustNewZero() + encode := func(values ...[]byte) []byte { + return bytes.Join(values, nil) + } + zeroRows := int64(0) + zeroCount := int32(0) + negative := int32(-1) + huge := int32(1<<20 + 1) + one := int32(1) + oversized := int32(1 << 30) + + tests := []struct { + name string + data []byte + }{ + { + name: "negative vector count", + data: encode(types.EncodeInt64(&zeroRows), types.EncodeInt32(&negative)), + }, + { + name: "huge vector count", + data: encode(types.EncodeInt64(&zeroRows), types.EncodeInt32(&huge)), + }, + { + name: "negative attribute count", + data: encode(types.EncodeInt64(&zeroRows), types.EncodeInt32(&zeroCount), types.EncodeInt32(&negative)), + }, + { + name: "oversized attribute payload", + data: encode( + types.EncodeInt64(&zeroRows), + types.EncodeInt32(&zeroCount), + types.EncodeInt32(&one), + types.EncodeInt32(&oversized), + ), + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + decoded := NewOffHeapEmpty() + require.NoError(t, decoded.SetAllocationAccount(state.selection)) + require.NotPanics(t, func() { + require.Error(t, decoded.UnmarshalFromReader(bytes.NewReader(test.data), mp)) + }) + decoded.Clean(mp) + require.Zero(t, state.account.Snapshot().Used) + }) + } + finalizeTestBatchAllocationAccount(t, state) +} + +func newMixedBatchAllocationSource( + t *testing.T, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, + rows int, +) *Batch { + t.Helper() + bat := NewOffHeapWithSize(2) + bat.Attrs = []string{"accounted", "unaccounted"} + bat.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, bat.Vecs[0].SetAllocationAccount(selection)) + bat.Vecs[1] = vector.NewOffHeapVecWithType(types.T_varchar.ToType()) + for i := 0; i < rows; i++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(i), false, mp)) + require.NoError(t, vector.AppendBytes(bat.Vecs[1], []byte("unaccounted"), false, mp)) + } + bat.SetRowCount(rows) + return bat +} + +func TestMixedBatchAllocationClonePreservesVectorProvenance(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + mp := mpool.MustNewZero() + source := newMixedBatchAllocationSource(t, mp, state.selection, 8) + sourceUsed := state.account.Snapshot().Used + require.NotZero(t, sourceUsed) + require.Nil(t, source.AllocationAccountSelection()) + + _, err := source.Clone(mp, false) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + for _, clone := range []func() (*Batch, error){ + func() (*Batch, error) { return source.Clone(mp, true) }, + func() (*Batch, error) { return source.Dup(mp) }, + } { + got, err := clone() + require.NoError(t, err) + require.Nil(t, got.AllocationAccountSelection()) + require.Same(t, state.selection, got.Vecs[0].AllocationAccountSelection()) + require.Nil(t, got.Vecs[1].AllocationAccountSelection()) + got.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + } + + accounted, err := source.CloneSelectedColumns([]int{0}, []string{"accounted"}, mp) + require.NoError(t, err) + require.Same(t, state.selection, accounted.Vecs[0].AllocationAccountSelection()) + accounted.Clean(mp) + unaccounted, err := source.CloneSelectedColumns([]int{1}, []string{"unaccounted"}, mp) + require.NoError(t, err) + require.Nil(t, unaccounted.Vecs[0].AllocationAccountSelection()) + unaccounted.Clean(mp) + + source.FreeColumns(mp) + require.Zero(t, state.account.Snapshot().Used) + require.Same(t, state.selection, source.Vecs[0].AllocationAccountSelection()) + require.Nil(t, source.Vecs[1].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(9), false, mp)) + require.NotZero(t, state.account.Snapshot().Used) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchImplicitClonesPromoteAccountedVectorsOffHeap(t *testing.T) { + state := newTestBatchAllocationAccount(t, 16) + mp := mpool.MustNewZero() + source := NewWithSize(2) + source.Attrs = []string{"accounted", "unaccounted"} + source.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, source.Vecs[0].SetAllocationAccount(state.selection)) + source.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(1), false, mp)) + require.NoError(t, vector.AppendFixed(source.Vecs[1], int64(2), false, mp)) + source.SetRowCount(1) + sourceUsed := state.account.Snapshot().Used + + dup, err := source.Dup(mp) + require.NoError(t, err) + require.True(t, dup.offHeap) + require.Same(t, state.selection, dup.Vecs[0].AllocationAccountSelection()) + require.Nil(t, dup.Vecs[1].AllocationAccountSelection()) + dup.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + selected, err := source.CloneSelectedColumns( + []int{0}, + []string{"accounted"}, + mp, + ) + require.NoError(t, err) + require.True(t, selected.offHeap) + require.Same(t, state.selection, selected.Vecs[0].AllocationAccountSelection()) + selected.Clean(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestMixedBatchAllocationBatchSetPreservesVectorProvenance(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + mp := mpool.MustNewZero() + set := NewBatchSet(4) + first := newMixedBatchAllocationSource(t, mp, state.selection, 2) + second := newMixedBatchAllocationSource(t, mp, state.selection, 6) + + consumed, err := set.Extend(mp, first, nil) + require.NoError(t, err) + require.False(t, consumed) + consumed, err = set.Extend(mp, second, nil) + require.NoError(t, err) + require.False(t, consumed) + require.Equal(t, 2, set.Length()) + require.Equal(t, 8, set.RowCount()) + for i := 0; i < set.Length(); i++ { + require.Same(t, state.selection, set.Get(i).Vecs[0].AllocationAccountSelection()) + require.Nil(t, set.Get(i).Vecs[1].AllocationAccountSelection()) + } + + first.Clean(mp) + second.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetFillsTailAcrossEquivalentAllocationSelections(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + equivalent, err := vector.NewAllocationAccountSelection( + state.account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + require.NotSame(t, state.selection, equivalent) + mp := mpool.MustNewZero() + set := NewBatchSet(4) + first := newMixedBatchAllocationSource(t, mp, state.selection, 2) + second := newMixedBatchAllocationSource(t, mp, equivalent, 2) + + _, err = set.Extend(mp, first, nil) + require.NoError(t, err) + require.Equal(t, 1, set.ReadyDeltaFor(second, second.RowCount())) + _, err = set.Extend(mp, second, nil) + require.NoError(t, err) + require.Equal(t, 1, set.Length()) + require.Equal(t, 1, set.ReadyCount()) + require.Equal(t, 4, set.Get(0).RowCount()) + require.Same(t, state.selection, set.Get(0).Vecs[0].AllocationAccountSelection()) + + first.Clean(mp) + second.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetCoalescesEquivalentParallelProducerChunks(t *testing.T) { + const ( + batchMaxRows = 8192 + chunkRows = 72 + chunkCount = 114 + ) + state := newTestBatchAllocationAccount(t, 256) + mp := mpool.MustNewZero() + set := NewBatchSet(batchMaxRows) + + for range chunkCount { + selection, err := vector.NewAllocationAccountSelection( + state.account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + source := newMixedBatchAllocationSource( + t, + mp, + selection, + chunkRows, + ) + _, err = set.Extend(mp, source, nil) + require.NoError(t, err) + source.Clean(mp) + } + + require.Equal(t, 2, set.Length()) + require.Equal(t, 1, set.ReadyCount()) + require.Equal(t, batchMaxRows, set.Get(0).RowCount()) + require.Equal(t, chunkRows*chunkCount-batchMaxRows, set.Get(1).RowCount()) + + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetStartsNewTailWhenVectorProvenanceChanges(t *testing.T) { + state := newTestBatchAllocationAccount(t, 128) + mp := mpool.MustNewZero() + set := NewBatchSet(4) + unaccounted := newBatchAllocationTestSource(t, mp, nil) + unaccounted.Shrink([]int64{0, 1}, false) + mixed := newMixedBatchAllocationSource(t, mp, state.selection, 3) + + _, err := set.Extend(mp, unaccounted, nil) + require.NoError(t, err) + ready := set.ReadyCount() + require.Equal(t, 1, set.ReadyDeltaFor(mixed, mixed.RowCount())) + _, err = set.Extend(mp, mixed, nil) + require.NoError(t, err) + require.Equal(t, 1, set.ReadyCount()-ready) + require.Equal(t, 2, set.Length()) + require.Equal(t, 2, set.Get(0).RowCount()) + require.Equal(t, 3, set.Get(1).RowCount()) + require.Nil(t, set.Get(0).Vecs[0].AllocationAccountSelection()) + require.Same(t, state.selection, set.Get(1).Vecs[0].AllocationAccountSelection()) + + unaccountedUnion := newBatchAllocationTestSource(t, mp, nil) + ready = set.ReadyCount() + require.Equal(t, 1, set.ReadyDeltaFor(unaccountedUnion, 1)) + _, err = set.Union(mp, unaccountedUnion, []int32{0}, nil) + require.NoError(t, err) + require.Equal(t, 1, set.ReadyCount()-ready) + require.Equal(t, 3, set.Length()) + require.Nil(t, set.Get(2).Vecs[0].AllocationAccountSelection()) + + pushed := newMixedBatchAllocationSource(t, mp, state.selection, 1) + require.NoError(t, set.Push(mp, pushed)) + require.Equal(t, 4, set.Length()) + require.Same(t, state.selection, set.Get(3).Vecs[0].AllocationAccountSelection()) + require.Equal(t, 1, set.Get(3).RowCount()) + + unaccounted.Clean(mp) + mixed.Clean(mp) + unaccountedUnion.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchSetPreservesUniformBatchAllocationContext(t *testing.T) { + state := newTestBatchAllocationAccount(t, 256) + mp := mpool.MustNewZero() + set := NewBatchSet(16) + first := newBatchAllocationTestSource(t, mp, state.selection) + first.Shrink([]int64{0, 1, 2, 3, 4, 5, 6, 7}, false) + second := newBatchAllocationTestSource(t, mp, state.selection) + second.Shrink([]int64{ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, + 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, + }, false) + + _, err := set.Extend(mp, first, nil) + require.NoError(t, err) + _, err = set.Extend(mp, second, nil) + require.NoError(t, err) + require.Equal(t, 2, set.Length()) + for i := 0; i < set.Length(); i++ { + require.Same(t, state.selection, set.Get(i).AllocationAccountSelection()) + for _, vec := range set.Get(i).Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + } + + reuse := NewWithSchema( + true, + first.Attrs, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, reuse.SetAllocationAccount(state.selection)) + third := newBatchAllocationTestSource(t, mp, state.selection) + third.Shrink([]int64{0, 1, 2, 3, 4, 5, 6, 7}, false) + consumed, err := set.Extend(mp, third, reuse) + require.NoError(t, err) + require.True(t, consumed) + require.Equal(t, 3, set.Length()) + require.Same(t, state.selection, set.Get(2).AllocationAccountSelection()) + + set.Get(2).FreeColumns(mp) + require.Same(t, state.selection, set.Get(2).AllocationAccountSelection()) + for _, vec := range set.Get(2).Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + + first.Clean(mp) + second.Clean(mp) + third.Clean(mp) + set.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAllocationAccountDestinationCloneUnionAndReuse(t *testing.T) { + state := newTestBatchAllocationAccount(t, 64) + mp := mpool.MustNewZero() + source := newBatchAllocationTestSource(t, mp, nil) + + destination := NewWithSchema( + true, + source.Attrs, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, destination.SetAllocationAccount(state.selection)) + require.NoError(t, source.CloneTo(destination, mp)) + require.NotZero(t, state.account.Snapshot().Used) + require.Equal(t, source.RowCount(), destination.RowCount()) + + destination.FreeColumns(mp) + require.Zero(t, state.account.Snapshot().Used) + for _, vec := range destination.Vecs { + require.Same(t, state.selection, vec.AllocationAccountSelection()) + } + + require.NoError(t, destination.Union(source, []int64{1, 3, 5, 7}, mp)) + require.Equal(t, 4, destination.RowCount()) + require.NotZero(t, state.account.Snapshot().Used) + + destination.Clean(mp) + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAllocationAccountCloneRollback(t *testing.T) { + state := newTestBatchAllocationAccount(t, 1) + mp := mpool.MustNewZero() + source := NewWithSchema( + true, + []string{"left", "right"}, + []types.Type{types.T_int64.ToType(), types.T_int64.ToType()}, + ) + for i := range source.Vecs { + require.NoError(t, vector.AppendFixed(source.Vecs[i], int64(i), false, mp)) + } + source.SetRowCount(1) + + destination := NewWithSchema( + true, + source.Attrs, + []types.Type{types.T_int64.ToType(), types.T_int64.ToType()}, + ) + require.NoError(t, destination.SetAllocationAccount(state.selection)) + err := source.CloneTo(destination, mp) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + require.Nil(t, destination.AllocationAccountSelection()) + + source.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} + +func TestBatchAllocationAccountConfigurationIsAtomic(t *testing.T) { + state := newTestBatchAllocationAccount(t, 8) + mp := mpool.MustNewZero() + + onHeap := NewWithSchema( + false, + nil, + []types.Type{types.T_int64.ToType()}, + ) + require.ErrorIs( + t, + onHeap.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + + offHeap := NewWithSchema( + true, + nil, + []types.Type{types.T_int64.ToType(), types.T_int64.ToType()}, + ) + require.NoError(t, vector.AppendFixed(offHeap.Vecs[1], int64(1), false, mp)) + require.ErrorIs( + t, + offHeap.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + require.Nil(t, offHeap.Vecs[0].AllocationAccountSelection()) + require.Nil(t, offHeap.AllocationAccountSelection()) + + offHeap.Clean(mp) + onHeap.Clean(mp) + finalizeTestBatchAllocationAccount(t, state) +} diff --git a/pkg/container/batch/batch.go b/pkg/container/batch/batch.go index ece832ab7d7de..9bce06e335050 100644 --- a/pkg/container/batch/batch.go +++ b/pkg/container/batch/batch.go @@ -17,8 +17,10 @@ package batch import ( "bytes" "context" + "encoding/binary" "fmt" "io" + "math" "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" @@ -112,50 +114,198 @@ func (bat *Batch) MarshalBinary() ([]byte, error) { } func (bat *Batch) MarshalBinaryWithBuffer(w *bytes.Buffer, reset bool) ([]byte, error) { - // reset the buffer if caller wants to. if reset { w.Reset() } + if err := bat.MarshalBinaryTo(w); err != nil { + return nil, err + } + return w.Bytes(), nil +} - // row count. - rl := int64(bat.rowCount) - w.Write(types.EncodeInt64(&rl)) +func (bat *Batch) MarshalBinarySize() (int, error) { + return bat.prepareMarshalBinary(nil) +} + +func (bat *Batch) prepareMarshalBinary( + plans []vector.MarshalBinaryPlan, +) (int, error) { + if bat == nil { + return 0, moerr.NewInvalidInputNoCtx("invalid batch for marshal") + } + const fixedSize = uint64(8 + 4 + 4 + 4 + 4 + 4) + total := fixedSize + add := func(value uint64) bool { + if value > uint64(^uint(0)>>1)-total { + return false + } + total += value + return true + } + if uint64(len(bat.Vecs)) > uint64(^uint32(0)>>1) || + uint64(len(bat.Attrs)) > uint64(^uint32(0)>>1) || + uint64(len(bat.ExtraBuf)) > uint64(^uint32(0)>>1) { + return 0, moerr.NewInvalidInputNoCtx( + "batch field exceeds marshal format", + ) + } + if plans != nil && len(plans) < len(bat.Vecs) { + return 0, moerr.NewInvalidInputNoCtx("short batch marshal plan") + } + for i, vec := range bat.Vecs { + if vec == nil { + return 0, moerr.NewInvalidInputNoCtx( + "cannot marshal a nil batch vector", + ) + } + plan, err := vec.PrepareMarshalBinary() + if err != nil { + return 0, err + } + size := plan.Size() + if uint64(size) > uint64(^uint32(0)) || + !add(4+uint64(size)) { + return 0, moerr.NewInvalidInputNoCtx( + "batch vector exceeds marshal format", + ) + } + if plans != nil { + plans[i] = plan + } + } + for _, attr := range bat.Attrs { + if uint64(len(attr)) > uint64(^uint32(0)>>1) || + !add(4+uint64(len(attr))) { + return 0, moerr.NewInvalidInputNoCtx( + "batch attribute exceeds marshal format", + ) + } + } + if !add(uint64(len(bat.ExtraBuf))) { + return 0, moerr.NewInvalidInputNoCtx( + "batch marshal size exceeds platform limit", + ) + } + return int(total), nil +} + +func (bat *Batch) MarshalBinaryTo(w io.Writer) error { + if bat == nil || w == nil { + return io.ErrClosedPipe + } + var inlinePlans [64]vector.MarshalBinaryPlan + var plans []vector.MarshalBinaryPlan + if len(bat.Vecs) <= len(inlinePlans) { + plans = inlinePlans[:len(bat.Vecs)] + } else { + plans = make([]vector.MarshalBinaryPlan, len(bat.Vecs)) + } + size, err := bat.prepareMarshalBinary(plans) + if err != nil { + return err + } + if sized, ok := w.(interface { + Len() int + EnsureCapacity(int) error + }); ok { + if sized.Len() > math.MaxInt-size { + return moerr.NewInvalidInputNoCtx("batch marshal size exceeds platform limit") + } + if err := sized.EnsureCapacity(sized.Len() + size); err != nil { + return err + } + } + if err := writeBatchMarshalInt64(w, int64(bat.rowCount)); err != nil { + return err + } - // Vecs l := int32(len(bat.Vecs)) - w.Write(types.EncodeInt32(&l)) + if err := writeBatchMarshalInt32(w, l); err != nil { + return err + } for i := 0; i < int(l); i++ { - var size uint32 - offset := w.Len() - w.Write(types.EncodeUint32(&size)) - err := bat.Vecs[i].MarshalBinaryWithBuffer(w) - if err != nil { - return nil, err + if err := writeBatchMarshalUint32(w, uint32(plans[i].Size())); err != nil { + return err + } + if err := plans[i].MarshalTo(w); err != nil { + return err } - size = uint32(w.Len() - offset - 4) - buf := w.Bytes() - copy(buf[offset:], types.EncodeUint32(&size)) } - // Attrs l = int32(len(bat.Attrs)) - w.Write(types.EncodeInt32(&l)) + if err := writeBatchMarshalInt32(w, l); err != nil { + return err + } for i := 0; i < int(l); i++ { size := int32(len(bat.Attrs[i])) - w.Write(types.EncodeInt32(&size)) - n, _ := w.WriteString(bat.Attrs[i]) + if err := writeBatchMarshalInt32(w, size); err != nil { + return err + } + n, err := io.WriteString(w, bat.Attrs[i]) + if err != nil { + return err + } if int32(n) != size { - panic("unexpected length for string") + return io.ErrShortWrite } } - // ExtraBuf - types.WriteSizeBytes(bat.ExtraBuf, w) + extraSize := int32(len(bat.ExtraBuf)) + if err := writeBatchMarshalInt32(w, extraSize); err != nil { + return err + } + if err := writeBatchMarshalBytes(w, bat.ExtraBuf); err != nil { + return err + } - w.Write(types.EncodeInt32(&bat.Recursive)) - w.Write(types.EncodeInt32(&bat.ShuffleIDX)) + if err := writeBatchMarshalInt32(w, bat.Recursive); err != nil { + return err + } + return writeBatchMarshalInt32(w, bat.ShuffleIDX) +} - return w.Bytes(), nil +type batchPrimitiveWriter interface { + WriteUint32(uint32) error + WriteInt32(int32) error + WriteInt64(int64) error +} + +func writeBatchMarshalUint32(w io.Writer, value uint32) error { + if typed, ok := w.(batchPrimitiveWriter); ok { + return typed.WriteUint32(value) + } + var data [4]byte + binary.NativeEndian.PutUint32(data[:], value) + return writeBatchMarshalBytes(w, data[:]) +} + +func writeBatchMarshalInt32(w io.Writer, value int32) error { + if typed, ok := w.(batchPrimitiveWriter); ok { + return typed.WriteInt32(value) + } + var data [4]byte + binary.NativeEndian.PutUint32(data[:], uint32(value)) + return writeBatchMarshalBytes(w, data[:]) +} + +func writeBatchMarshalInt64(w io.Writer, value int64) error { + if typed, ok := w.(batchPrimitiveWriter); ok { + return typed.WriteInt64(value) + } + var data [8]byte + binary.NativeEndian.PutUint64(data[:], uint64(value)) + return writeBatchMarshalBytes(w, data[:]) +} + +func writeBatchMarshalBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } + return nil } func (bat *Batch) UnmarshalBinary(data []byte) (err error) { @@ -200,6 +350,7 @@ func (c *batchUnmarshalCursor) readUint32() (uint32, error) { } func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err error) { + allocationAccount := bat.allocationAccount cursor := batchUnmarshalCursor{data: data} rowCount, err := cursor.readInt64() if err != nil { @@ -234,6 +385,7 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } } bat.Clean(mp) + bat.allocationAccount = allocationAccount } bat.Vecs = make([]*vector.Vector, vecsLen) } @@ -290,6 +442,14 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } vecs[i].Free(mp) } + // UnmarshalBinary installs aliases into vecData. An empty accounted + // receiver must explicitly drop its future-allocation selection first; + // the Batch retains the destination context for a later owned copy. + if vecs[i].AllocationAccountSelection() != nil { + if err := vecs[i].SetAllocationAccount(nil); err != nil { + return err + } + } if err := vecs[i].UnmarshalBinary(vecData); err != nil { return err } @@ -400,28 +560,35 @@ func (bat *Batch) UnmarshalBinaryWithAnyMp(data []byte, mp *mpool.MPool) (err er } func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) { + return bat.unmarshalFromReader(r, mp, true) +} + +func (bat *Batch) unmarshalFromReader( + r io.Reader, + mp *mpool.MPool, + allowMetadata bool, +) (err error) { + if bat == nil || r == nil { + return io.ErrClosedPipe + } i64, err := types.ReadInt64(r) if err != nil { return err } - bat.rowCount = int(i64) + if i64 < 0 || int64(int(i64)) != i64 { + return moerr.NewInvalidInputNoCtx("invalid batch row count") + } + decodedRowCount := int(i64) l, err := types.ReadInt32AsInt(r) if err != nil { return err } - if l != len(bat.Vecs) { - if len(bat.Vecs) > 0 { - bat.Clean(mp) - } - bat.Vecs = make([]*vector.Vector, l) - for i := range bat.Vecs { - if bat.offHeap { - bat.Vecs[i] = vector.NewOffHeapVec() - } else { - bat.Vecs[i] = vector.NewVecFromReuse() - } - } + if err = validateReaderElementCount(r, l, 4, "vector"); err != nil { + return err + } + if err = bat.prepareOwnedDecodeVectors(l, mp); err != nil { + return err } vecs := bat.Vecs @@ -446,21 +613,39 @@ func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) if err != nil { return err } + if err = validateReaderElementCount(r, l, 4, "attribute"); err != nil { + return err + } + if !allowMetadata && l != 0 { + return moerr.NewInvalidInputNoCtx("spill batch attributes are not allowed") + } if l != len(bat.Attrs) { bat.Attrs = make([]string, l) } for i := 0; i < int(l); i++ { - _, bs, err := types.ReadSizeBytes(r) + bs, err := readBatchSizedBytes(r) if err != nil { return err } bat.Attrs[i] = string(bs) } - // ExtraBuf - if _, bat.ExtraBuf, err = types.ReadSizeBytes(r); err != nil { - return err + // ExtraBuf is a data-scaled Go-heap field in the stable Batch codec. Spill + // records do not use it and reject it before allocating its payload. + if allowMetadata { + if bat.ExtraBuf, err = readBatchSizedBytes(r); err != nil { + return err + } + } else { + extraSize, readErr := types.ReadInt32AsInt(r) + if readErr != nil { + return readErr + } + if extraSize != 0 { + return moerr.NewInvalidInputNoCtx("spill batch extra buffer is not allowed") + } + bat.ExtraBuf = nil } if bat.Recursive, err = types.ReadInt32(r); err != nil { @@ -469,6 +654,131 @@ func (bat *Batch) UnmarshalFromReader(r io.Reader, mp *mpool.MPool) (err error) if bat.ShuffleIDX, err = types.ReadInt32(r); err != nil { return err } + bat.rowCount = decodedRowCount + return nil +} + +func readBatchSizedBytes(r io.Reader) ([]byte, error) { + size, err := types.ReadInt32AsInt(r) + if err != nil { + return nil, err + } + if size < 0 { + return nil, moerr.NewInvalidInputNoCtx("negative batch buffer size") + } + if limited, ok := r.(*io.LimitedReader); ok && int64(size) > limited.N { + return nil, io.ErrUnexpectedEOF + } + if lengthAware, ok := r.(interface{ Len() int }); ok && size > lengthAware.Len() { + return nil, io.ErrUnexpectedEOF + } + if size == 0 { + return nil, nil + } + value := make([]byte, size) + if _, err = io.ReadFull(r, value); err != nil { + return nil, err + } + return value, nil +} + +func validateReaderElementCount( + r io.Reader, + count int, + minimumWireBytes int64, + field string, +) error { + const maxBatchWireFields = 1 << 20 + if count < 0 || count > maxBatchWireFields || minimumWireBytes <= 0 { + return moerr.NewInvalidInputNoCtx("invalid batch " + field + " count") + } + var remaining int64 = -1 + switch reader := r.(type) { + case *io.LimitedReader: + remaining = reader.N + case interface{ Len() int }: + remaining = int64(reader.Len()) + } + if remaining >= 0 && int64(count) > remaining/minimumWireBytes { + return moerr.NewInvalidInputNoCtx("invalid batch " + field + " count") + } + return nil +} + +// prepareOwnedDecodeVectors makes every destination an independent owner. +// Alias decoding deliberately installs borrowed vector buffers; those buffers +// must never be grown or relabeled by the owned streaming decoder. +func (bat *Batch) prepareOwnedDecodeVectors(count int, mp *mpool.MPool) error { + if count < 0 { + return moerr.NewInvalidInputNoCtx("invalid batch vector count") + } + allocationAccount := bat.allocationAccount + if count != len(bat.Vecs) { + if len(bat.Vecs) > 0 { + bat.Clean(mp) + bat.allocationAccount = allocationAccount + } + bat.Vecs = make([]*vector.Vector, count) + } + + const inlineReceivers = 16 + var inline [inlineReceivers]*vector.Vector + var used map[*vector.Vector]struct{} + for i, vec := range bat.Vecs { + selection := allocationAccount + if selection == nil && vec != nil { + selection = vec.AllocationAccountSelection() + } + if vec != nil { + exists := false + if i < inlineReceivers { + for j := 0; j < i; j++ { + if inline[j] == vec { + exists = true + break + } + } + inline[i] = vec + } else { + if used == nil { + used = make(map[*vector.Vector]struct{}, count) + for _, prior := range inline { + if prior != nil { + used[prior] = struct{}{} + } + } + } + _, exists = used[vec] + used[vec] = struct{}{} + } + if exists { + vec = nil + } + } + if vec == nil { + if bat.offHeap { + vec = vector.NewOffHeapVec() + } else { + vec = vector.NewVecFromReuse() + } + } else if vec.NeedDup() { + vec.Free(mp) + } + vec.SetOffHeap(bat.offHeap) + if !vector.AllocationAccountSelectionsEqual( + vec.AllocationAccountSelection(), + selection, + ) { + if err := vec.CanSetAllocationAccount(selection); err != nil { + vec.Free(mp) + vec.SetOffHeap(bat.offHeap) + } + if err := vec.SetAllocationAccount(selection); err != nil { + return err + } + } + bat.Vecs[i] = vec + } return nil } @@ -544,27 +854,70 @@ func (bat *Batch) SetAttributes(attrs []string) { bat.Attrs = attrs } +// AllocationAccountSelection returns the immutable destination selection used +// by this batch's owned off-heap vectors. +func (bat *Batch) AllocationAccountSelection() *vector.AllocationAccountSelection { + if bat == nil { + return nil + } + return bat.allocationAccount +} + +// SetAllocationAccount configures every existing empty destination vector as +// one transaction. Existing physical allocations are never relabeled. +func (bat *Batch) SetAllocationAccount( + selection *vector.AllocationAccountSelection, +) error { + if bat == nil || (selection != nil && !bat.offHeap) { + return mpool.ErrAllocationAccountInvalid + } + for _, vec := range bat.Vecs { + if vec != nil { + if err := vec.CanSetAllocationAccount(selection); err != nil { + return err + } + } + } + for _, vec := range bat.Vecs { + if vec != nil { + if err := vec.SetAllocationAccount(selection); err != nil { + panic(err) + } + } + } + bat.allocationAccount = selection + return nil +} + +func (bat *Batch) configureOwnedVector(vec *vector.Vector) { + if vec == nil { + return + } + vec.SetOffHeap(bat.offHeap) + if bat.allocationAccount != nil { + if err := vec.SetAllocationAccount(bat.allocationAccount); err != nil { + panic(err) + } + } +} + func (bat *Batch) InsertVector( pos int32, attr string, vec *vector.Vector, ) { + bat.configureOwnedVector(vec) bat.Vecs = append(bat.Vecs, nil) copy(bat.Vecs[pos+1:], bat.Vecs[pos:]) bat.Vecs[pos] = vec - if vec != nil { - vec.SetOffHeap(bat.offHeap) - } bat.Attrs = append(bat.Attrs, "") copy(bat.Attrs[pos+1:], bat.Attrs[pos:]) bat.Attrs[pos] = attr } func (bat *Batch) SetVector(pos int32, vec *vector.Vector) { + bat.configureOwnedVector(vec) bat.Vecs[pos] = vec - if vec != nil { - vec.SetOffHeap(bat.offHeap) - } } func (bat *Batch) GetVector(pos int32) *vector.Vector { @@ -578,15 +931,19 @@ func (bat *Batch) CloneSelectedColumns( ) (cloned *Batch, err error) { cloned = NewWithSize(len(selectCols)) cloned.Attrs = selectAttrs - cloned.offHeap = bat.offHeap + cloned.offHeap = bat.offHeap || bat.selectedColumnsHaveAllocationAccount(selectCols) var typ types.Type for idx := range selectCols { - if bat.offHeap { + if cloned.offHeap { cloned.Vecs[idx] = vector.NewOffHeapVecWithType(typ) } else { cloned.Vecs[idx] = vector.NewVec(typ) } } + if err = configureCloneAllocation(bat, cloned, selectCols); err != nil { + cloned.Clean(mp) + return nil, err + } if err = bat.CloneSelectedColumnsTo(selectCols, cloned, mp); err != nil { cloned.Clean(mp) cloned = nil @@ -626,6 +983,7 @@ func (bat *Batch) SelectColumns(cols []int, attrs []string) *Batch { rbat := NewWithSize(len(cols)) rbat.Attrs = attrs rbat.offHeap = bat.offHeap + rbat.allocationAccount = bat.allocationAccount for i, col := range cols { rbat.Vecs[i] = bat.Vecs[col] } @@ -650,6 +1008,7 @@ func (bat *Batch) Clean(m *mpool.MPool) { bat.Attrs = nil bat.ExtraBuf = nil bat.SetRowCount(0) + bat.allocationAccount = nil } func (bat *Batch) Last() bool { @@ -680,7 +1039,14 @@ func (bat *Batch) CleanOnlyData() { func (bat *Batch) FreeColumns(m *mpool.MPool) { for _, vec := range bat.Vecs { if vec != nil { + selection := vec.AllocationAccountSelection() vec.Free(m) + if bat.allocationAccount != nil { + selection = bat.allocationAccount + } + if err := vec.SetAllocationAccount(selection); err != nil { + panic(err) + } } } } @@ -704,12 +1070,72 @@ func (bat *Batch) GetSchema() (attrs []string, attrTypes []types.Type) { return } +func vectorAllocationSelectionsMatch(left, right *Batch) bool { + if left == nil || right == nil || len(left.Vecs) != len(right.Vecs) { + return false + } + if !vector.AllocationAccountSelectionsEqual( + left.allocationAccount, + right.allocationAccount, + ) { + return false + } + for i := range left.Vecs { + if left.Vecs[i] == nil || right.Vecs[i] == nil { + if left.Vecs[i] != right.Vecs[i] { + return false + } + continue + } + if !vector.AllocationAccountSelectionsEqual( + left.Vecs[i].AllocationAccountSelection(), + right.Vecs[i].AllocationAccountSelection(), + ) { + return false + } + } + return true +} + +func configureCloneAllocation( + source, destination *Batch, + selectedColumns []int, +) error { + if source == nil || destination == nil { + return mpool.ErrAllocationAccountInvalid + } + if source.allocationAccount != nil { + return destination.SetAllocationAccount(source.allocationAccount) + } + for destinationIdx := range destination.Vecs { + sourceIdx := destinationIdx + if len(selectedColumns) > 0 { + sourceIdx = selectedColumns[destinationIdx] + } + selection := source.Vecs[sourceIdx].AllocationAccountSelection() + if selection == nil { + continue + } + if !destination.offHeap { + return mpool.ErrAllocationAccountInvalid + } + if err := destination.Vecs[destinationIdx].SetAllocationAccount(selection); err != nil { + return err + } + } + return nil +} + func (bat *Batch) Clone(mp *mpool.MPool, offHeap bool) (*Batch, error) { var ( cloned *Batch attrs, attrTypes = bat.GetSchema() ) cloned = NewWithSchema(offHeap, attrs, attrTypes) + if err := configureCloneAllocation(bat, cloned, nil); err != nil { + cloned.Clean(mp) + return nil, err + } cloned.Recursive = bat.Recursive err := bat.CloneTo(cloned, mp) if err != nil { @@ -737,7 +1163,56 @@ func (bat *Batch) CloneTo(toBat *Batch, mp *mpool.MPool) (err error) { // Dup used to copy a Batch object, this method will create a new batch // and copy all vectors (Vecs) of the current batch to the new batch. func (bat *Batch) Dup(mp *mpool.MPool) (*Batch, error) { - return bat.Clone(mp, bat.offHeap) + return bat.Clone(mp, bat.offHeap || bat.hasAllocationAccountVector()) +} + +// CloneWithoutAllocationAccount deep-copies bat without carrying an +// allocation account into the destination. Use it only at an ownership +// boundary whose destination can outlive the source account. +func (bat *Batch) CloneWithoutAllocationAccount( + mp *mpool.MPool, + offHeap bool, +) (*Batch, error) { + attrs, attrTypes := bat.GetSchema() + cloned := NewWithSchema(offHeap, attrs, attrTypes) + cloned.Recursive = bat.Recursive + if err := bat.CloneTo(cloned, mp); err != nil { + return nil, err + } + return cloned, nil +} + +// DupWithoutAllocationAccount is the ownership-boundary counterpart of Dup. +func (bat *Batch) DupWithoutAllocationAccount(mp *mpool.MPool) (*Batch, error) { + return bat.CloneWithoutAllocationAccount( + mp, + bat.offHeap || bat.hasAllocationAccountVector(), + ) +} + +func (bat *Batch) hasAllocationAccountVector() bool { + for _, vec := range bat.Vecs { + if vec != nil && vec.AllocationAccountSelection() != nil { + return true + } + } + return false +} + +// HasAllocationAccount reports whether the batch or one of its vectors owns +// memory charged to an execution-scoped allocation account. +func (bat *Batch) HasAllocationAccount() bool { + return bat != nil && (bat.allocationAccount != nil || bat.hasAllocationAccountVector()) +} + +func (bat *Batch) selectedColumnsHaveAllocationAccount(selectCols []int) bool { + for _, sourceIdx := range selectCols { + vec := bat.Vecs[sourceIdx] + if vec != nil && vec.AllocationAccountSelection() != nil { + return true + } + } + return false } func (bat *Batch) Union(bat2 *Batch, sels []int64, m *mpool.MPool) error { @@ -870,6 +1345,8 @@ func (bat *Batch) Window(start, end int) (*Batch, error) { b := NewWithSize(len(bat.Vecs)) var err error b.Attrs = bat.Attrs + b.offHeap = bat.offHeap + b.allocationAccount = bat.allocationAccount for i, vec := range bat.Vecs { b.Vecs[i], err = vec.Window(start, end) if err != nil { @@ -880,3 +1357,38 @@ func (bat *Batch) Window(start, end int) (*Batch, error) { b.rowCount = end - start return b, nil } + +// WindowWithAllocation is the allocation-accounted counterpart of Window. +// Vector data and area remain borrowed; null/grouping range bitmaps are owned +// by selection and are released when the returned batch is cleaned. +func (bat *Batch) WindowWithAllocation( + start int, + end int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*Batch, error) { + if bat == nil || mp == nil || selection == nil || + start < 0 || end < start || end > bat.RowCount() { + return nil, mpool.ErrAllocationAccountInvalid + } + b := NewOffHeapWithSize(len(bat.Vecs)) + b.Attrs = bat.Attrs + if err := b.SetAllocationAccount(selection); err != nil { + b.Clean(mp) + return nil, err + } + for i, vec := range bat.Vecs { + if vec == nil { + b.Clean(mp) + return nil, mpool.ErrAllocationAccountInvalid + } + var err error + b.Vecs[i], err = vec.WindowWithAllocation(start, end, mp, selection) + if err != nil { + b.Clean(mp) + return nil, err + } + } + b.rowCount = end - start + return b, nil +} diff --git a/pkg/container/batch/batch_set.go b/pkg/container/batch/batch_set.go index 7e6ad16ec55d6..6b39d95dbdb01 100644 --- a/pkg/container/batch/batch_set.go +++ b/pkg/container/batch/batch_set.go @@ -42,8 +42,9 @@ func (bs *BatchSet) Length() int { return len(bs.batches) } -// ReadyCount returns the number of full batches that can be consumed. The last -// batch remains a writable tail only while it is partial. +// ReadyCount returns the number of batches that can be consumed. All batches +// before the last are sealed, including a partial batch whose successor has a +// different allocation provenance. The last is ready only when it is full. func (bs *BatchSet) ReadyCount() int { if len(bs.batches) == 0 { return 0 @@ -69,6 +70,23 @@ func (bs *BatchSet) ReadyDelta(rowCount int) int { return (lastRows+rowCount)/bs.batchMaxRow - lastRows/bs.batchMaxRow } +// ReadyDeltaFor returns how many batches become consumable when rows copied +// from source are appended. A provenance change seals the existing partial +// tail because already allocated vectors cannot be relabeled. +func (bs *BatchSet) ReadyDeltaFor(source *Batch, rowCount int) int { + if rowCount <= 0 { + return 0 + } + if len(bs.batches) > 0 { + last := bs.batches[len(bs.batches)-1] + if last.RowCount() < bs.batchMaxRow && + !vectorAllocationSelectionsMatch(last, source) { + return 1 + rowCount/bs.batchMaxRow + } + } + return bs.ReadyDelta(rowCount) +} + func (bs *BatchSet) Get(idx int) *Batch { if idx >= len(bs.batches) { return nil @@ -120,9 +138,14 @@ func (bs *BatchSet) Push(mpool *mpool.MPool, inBatch *Batch) error { return nil } - defer func() { - inBatch.Clean(mpool) - }() + if !vectorAllocationSelectionsMatch(bs.batches[batLen-1], inBatch) { + // A Vector cannot change provenance after its first allocation. Seal the + // existing partial tail and preserve the incoming Batch as a new owner. + bs.batches = append(bs.batches, inBatch) + return nil + } + + defer inBatch.Clean(mpool) // fast path 2 if lastBatRowCount+inBatch.RowCount() <= bs.batchMaxRow { @@ -148,8 +171,12 @@ func (bs *BatchSet) Extend(mpool *mpool.MPool, inBatch *Batch, reuseBuf *Batch) // empty bats or last batch is full - can directly use fast path lastIdx := batLen - 1 + if batLen > 0 && bs.batches[lastIdx].rowCount < bs.batchMaxRow && + !vectorAllocationSelectionsMatch(bs.batches[lastIdx], inBatch) { + return bs.extendWithNewProvenance(mpool, inBatch, reuseBuf) + } if batLen == 0 || bs.batches[lastIdx].rowCount == bs.batchMaxRow { - if reuseBuf != nil && len(reuseBuf.Vecs) == len(inBatch.Vecs) { + if vectorAllocationSelectionsMatch(reuseBuf, inBatch) { reuseBuf.CleanOnlyData() reuseBuf, err = reuseBuf.AppendWithCopy(context.TODO(), mpool, inBatch) if err != nil { @@ -168,7 +195,7 @@ func (bs *BatchSet) Extend(mpool *mpool.MPool, inBatch *Batch, reuseBuf *Batch) // fast path 2: inBatch is full if inBatch.rowCount == bs.batchMaxRow { - if reuseBuf != nil && len(reuseBuf.Vecs) == len(inBatch.Vecs) { + if vectorAllocationSelectionsMatch(reuseBuf, inBatch) { reuseBuf.CleanOnlyData() reuseBuf, err = reuseBuf.AppendWithCopy(context.TODO(), mpool, inBatch) if err != nil { @@ -205,6 +232,13 @@ func (bs *BatchSet) Union(mpool *mpool.MPool, inBatch *Batch, sels []int32, reus if selsLen > inBatch.RowCount() { panic("sels len > inBatch.RowCount()") } + if bs.Length() > 0 { + last := bs.batches[bs.Length()-1] + if last.rowCount < bs.batchMaxRow && + !vectorAllocationSelectionsMatch(last, inBatch) { + return bs.unionWithNewProvenance(mpool, inBatch, sels, reuseBuf) + } + } consumed := false if bs.Length() == 0 { @@ -283,7 +317,7 @@ func (bs *BatchSet) Union(mpool *mpool.MPool, inBatch *Batch, sels []int32, reus } func (bs *BatchSet) getOrCreateBatch(inBatch *Batch, reuseBuf *Batch, mpool *mpool.MPool) (*Batch, error) { - if reuseBuf != nil && len(reuseBuf.Vecs) == len(inBatch.Vecs) { + if vectorAllocationSelectionsMatch(reuseBuf, inBatch) { reuseBuf.CleanOnlyData() return reuseBuf, nil } @@ -291,9 +325,80 @@ func (bs *BatchSet) getOrCreateBatch(inBatch *Batch, reuseBuf *Batch, mpool *mpo for i := range tmpBat.Vecs { tmpBat.Vecs[i] = vector.NewOffHeapVecWithType(*inBatch.Vecs[i].GetType()) } + if inBatch.allocationAccount != nil { + if err := tmpBat.SetAllocationAccount(inBatch.allocationAccount); err != nil { + tmpBat.Clean(mpool) + return nil, err + } + return tmpBat, nil + } + for i := range tmpBat.Vecs { + if selection := inBatch.Vecs[i].AllocationAccountSelection(); selection != nil { + if err := tmpBat.Vecs[i].SetAllocationAccount(selection); err != nil { + tmpBat.Clean(mpool) + return nil, err + } + } + } return tmpBat, nil } +func (bs *BatchSet) extendWithNewProvenance( + mp *mpool.MPool, + inBatch *Batch, + reuseBuf *Batch, +) (bool, error) { + consumed := false + for start := 0; start < inBatch.RowCount(); { + tmpBat, err := bs.getOrCreateBatch(inBatch, reuseBuf, mp) + if err != nil { + return consumed, err + } + if tmpBat == reuseBuf { + consumed = true + reuseBuf = nil + } + bs.batches = append(bs.batches, tmpBat) + count := min(bs.batchMaxRow, inBatch.RowCount()-start) + if err := tmpBat.UnionWindow(inBatch, start, count, mp); err != nil { + return consumed, err + } + start += count + } + return consumed, nil +} + +func (bs *BatchSet) unionWithNewProvenance( + mp *mpool.MPool, + inBatch *Batch, + sels []int32, + reuseBuf *Batch, +) (bool, error) { + consumed := false + for start := 0; start < len(sels); { + tmpBat, err := bs.getOrCreateBatch(inBatch, reuseBuf, mp) + if err != nil { + return consumed, err + } + if tmpBat == reuseBuf { + consumed = true + reuseBuf = nil + } + bs.batches = append(bs.batches, tmpBat) + count := min(bs.batchMaxRow, len(sels)-start) + for i := range tmpBat.Vecs { + if err := tmpBat.Vecs[i].UnionInt32( + inBatch.Vecs[i], sels[start:start+count], mp, + ); err != nil { + return consumed, err + } + } + tmpBat.rowCount = count + start += count + } + return consumed, nil +} + func (bs *BatchSet) RowCount() int { rowCount := 0 for _, bat := range bs.batches { diff --git a/pkg/container/batch/batch_test.go b/pkg/container/batch/batch_test.go index 0044033f8c1e6..673f03796960c 100644 --- a/pkg/container/batch/batch_test.go +++ b/pkg/container/batch/batch_test.go @@ -17,6 +17,7 @@ package batch import ( "bytes" "fmt" + "io" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" @@ -51,6 +52,17 @@ func TestBatchMarshalAndUnmarshal(t *testing.T) { for _, tc := range tcs { data, err := tc.bat.MarshalBinary() require.NoError(t, err) + size, err := tc.bat.MarshalBinarySize() + require.NoError(t, err) + require.Equal(t, len(data), size) + var streamed bytes.Buffer + require.NoError(t, tc.bat.MarshalBinaryTo(&streamed)) + require.Equal(t, data, streamed.Bytes()) + require.ErrorIs( + t, + tc.bat.MarshalBinaryTo(shortBatchMarshalWriter{}), + io.ErrShortWrite, + ) rbat := new(Batch) err = rbat.UnmarshalBinary(data) @@ -85,6 +97,22 @@ func TestBatchMarshalAndUnmarshal(t *testing.T) { } } +type shortBatchMarshalWriter struct{} + +func (shortBatchMarshalWriter) Write(value []byte) (int, error) { + return len(value) - 1, nil +} + +func TestMarshalBinarySizeRejectsInvalidBatch(t *testing.T) { + var nilBatch *Batch + _, err := nilBatch.MarshalBinarySize() + require.Error(t, err) + + invalid := NewWithSize(1) + _, err = invalid.MarshalBinarySize() + require.Error(t, err) +} + func TestBatch(t *testing.T) { for _, tc := range tcs { data, err := types.Encode(tc.bat) diff --git a/pkg/container/batch/grouping_codec.go b/pkg/container/batch/grouping_codec.go new file mode 100644 index 0000000000000..7d54d0a4c66a7 --- /dev/null +++ b/pkg/container/batch/grouping_codec.go @@ -0,0 +1,131 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package batch + +import ( + "io" + "math" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" +) + +// MarshalBinaryWithGroupingSize returns the spill-only Batch wire size. Spill +// records retain vectors and grouping provenance but deliberately omit Attrs +// and ExtraBuf: joins do not consume either field, and ExtraBuf would otherwise +// create an unaccounted data-scaled Go-heap owner while decoding. +func (bat *Batch) MarshalBinaryWithGroupingSize() (int, error) { + metadataFree := *bat + metadataFree.Attrs = nil + metadataFree.ExtraBuf = nil + size, err := metadataFree.MarshalBinarySize() + if err != nil { + return 0, err + } + if len(bat.Vecs) > math.MaxInt32 || size > math.MaxInt-4 { + return 0, moerr.NewInvalidInputNoCtx("batch grouping payload exceeds marshal format") + } + size += 4 + for _, vec := range bat.Vecs { + groupingSize := vec.GroupingMarshalBinarySize() + if groupingSize < 0 || groupingSize > math.MaxInt32 || + size > math.MaxInt-4-groupingSize { + return 0, moerr.NewInvalidInputNoCtx("batch grouping payload exceeds marshal format") + } + size += 4 + groupingSize + } + return size, nil +} + +func (bat *Batch) MarshalBinaryWithGroupingTo(w io.Writer) error { + if bat == nil || w == nil { + return io.ErrClosedPipe + } + metadataFree := *bat + metadataFree.Attrs = nil + metadataFree.ExtraBuf = nil + if err := metadataFree.MarshalBinaryTo(w); err != nil { + return err + } + return bat.marshalGroupingTo(w) +} + +func (bat *Batch) marshalGroupingTo(w io.Writer) error { + if err := writeBatchMarshalInt32(w, int32(len(bat.Vecs))); err != nil { + return err + } + for _, vec := range bat.Vecs { + size := vec.GroupingMarshalBinarySize() + if size > math.MaxInt32 { + return moerr.NewInvalidInputNoCtx("vector grouping payload exceeds marshal format") + } + if err := writeBatchMarshalInt32(w, int32(size)); err != nil { + return err + } + if size > 0 { + if err := vec.MarshalGroupingTo(w); err != nil { + return err + } + } + } + return nil +} + +func (bat *Batch) UnmarshalFromReaderWithGrouping( + r io.Reader, + mp *mpool.MPool, +) error { + if err := bat.unmarshalFromReader(r, mp, false); err != nil { + return err + } + return bat.unmarshalGroupingFromReader(r, mp) +} + +func (bat *Batch) unmarshalGroupingFromReader( + r io.Reader, + mp *mpool.MPool, +) error { + if err := bat.CheckLength(); err != nil { + return moerr.NewInvalidInputNoCtx("spill batch vector length does not match row count") + } + count, err := types.ReadInt32AsInt(r) + if err != nil { + return err + } + if count != len(bat.Vecs) { + return moerr.NewInvalidInputNoCtx("batch grouping vector count mismatch") + } + for _, vec := range bat.Vecs { + size, err := types.ReadInt32AsInt(r) + if err != nil { + return err + } + if size < 0 { + return moerr.NewInvalidInputNoCtx("invalid vector grouping payload size") + } + if remaining, ok := r.(*io.LimitedReader); ok && int64(size) > remaining.N { + return io.ErrUnexpectedEOF + } + limited := &io.LimitedReader{R: r, N: int64(size)} + if err = vec.UnmarshalGroupingFromReader(limited, size, mp); err != nil { + return err + } + if limited.N != 0 { + return moerr.NewInvalidInputNoCtx("vector grouping payload was not fully consumed") + } + } + return nil +} diff --git a/pkg/container/batch/types.go b/pkg/container/batch/types.go index 6229e3f98b22a..f3d1a708579b4 100644 --- a/pkg/container/batch/types.go +++ b/pkg/container/batch/types.go @@ -56,4 +56,8 @@ type Batch struct { // row count of batch, to instead of old len(Zs). rowCount int offHeap bool + + // allocationAccount is the destination selection for owned off-heap + // vectors created or reused by this batch. Alias vectors do not copy it. + allocationAccount *vector.AllocationAccountSelection } diff --git a/pkg/container/hashtable/allocation_account.go b/pkg/container/hashtable/allocation_account.go new file mode 100644 index 0000000000000..df40ba29f6cb5 --- /dev/null +++ b/pkg/container/hashtable/allocation_account.go @@ -0,0 +1,114 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashtable + +import ( + "unsafe" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +func HashMapBlockDescriptorBytes() uint64 { + return uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) +} + +// AllocationAccountSelection is the immutable provenance for one hash table's +// cell blocks and outer descriptor storage. The outer []slice-header backing +// is itself data-scaled ownership and therefore uses a distinct physical site. +type AllocationAccountSelection struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + cellSite mpool.AllocationSite + descriptorSite mpool.AllocationSite +} + +func NewAllocationAccountSelection( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + cellSite mpool.AllocationSite, + descriptorSite mpool.AllocationSite, +) (*AllocationAccountSelection, error) { + selection := &AllocationAccountSelection{ + account: account, + owner: owner, + cellSite: cellSite, + descriptorSite: descriptorSite, + } + if err := selection.validate(); err != nil { + return nil, err + } + return selection, nil +} + +func (s *AllocationAccountSelection) validate() error { + if s == nil || s.account == nil || s.account.Handle() == 0 || + s.owner < mpool.AllocationOwnerMin || + s.owner > mpool.AllocationOwnerMax || + s.cellSite < mpool.AllocationSiteMin || + s.descriptorSite < mpool.AllocationSiteMin || + s.cellSite == s.descriptorSite { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +func makeHashTableCellSlice[T any]( + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, + site mpool.AllocationSite, +) ([]T, error) { + if selection == nil { + return mpool.MakeSlice[T](length, mp, true) + } + if err := selection.validate(); err != nil { + return nil, err + } + return mpool.MakeSliceAccounted[T]( + length, + mp, + selection.account, + selection.owner, + site, + ) +} + +func makeHashTableDescriptorSlice[T any]( + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, + site mpool.AllocationSite, +) ([]T, error) { + if selection == nil { + return make([]T, length), nil + } + return makeHashTableCellSlice[T](length, mp, selection, site) +} + +func freeHashTableCellSlice[T any](mp *mpool.MPool, values []T) { + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} + +func freeHashTableDescriptorSlice[T any]( + mp *mpool.MPool, + values []T, + selection *AllocationAccountSelection, +) { + if selection != nil { + freeHashTableCellSlice(mp, values) + } +} diff --git a/pkg/container/hashtable/allocation_account_test.go b/pkg/container/hashtable/allocation_account_test.go new file mode 100644 index 0000000000000..f57f55aa6556e --- /dev/null +++ b/pkg/container/hashtable/allocation_account_test.go @@ -0,0 +1,365 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashtable + +import ( + "testing" + "unsafe" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +const ( + testHashTableOwner mpool.AllocationOwner = 1 + testHashTableCellSite mpool.AllocationSite = 24 + testHashTableDescriptorSite mpool.AllocationSite = 25 +) + +func newHashTableAllocation( + t testing.TB, + limit uint64, + metadataSlots uint64, +) (*mpool.AllocationAccountRegistry, *mpool.AllocationAccount, *AllocationAccountSelection) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := NewAllocationAccountSelection( + account, + testHashTableOwner, + testHashTableCellSite, + testHashTableDescriptorSite, + ) + require.NoError(t, err) + return registry, account, selection +} + +func completeHashTableAllocation( + t testing.TB, + registry *mpool.AllocationAccountRegistry, + account *mpool.AllocationAccount, +) { + t.Helper() + snapshot, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + require.Zero(t, snapshot.Used) +} + +func TestHashTableAllocationAccountsCellAndDescriptorStorage(t *testing.T) { + t.Run("int", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 64<<20, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + descriptorBytes := uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) + require.Equal( + t, + Int64HashMapInitialAllocationBytes()+descriptorBytes, + account.Snapshot().Used, + ) + table.Free() + completeHashTableAllocation(t, registry, account) + }) + + t.Run("string", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 64<<20, 64) + mp := mpool.MustNewZero() + var table StringHashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + descriptorBytes := uint64(unsafe.Sizeof([]StringHashMapCell(nil))) + require.Equal( + t, + StringHashMapInitialAllocationBytes()+descriptorBytes, + account.Snapshot().Used, + ) + table.Free() + completeHashTableAllocation(t, registry, account) + }) +} + +func TestIntHashTableAccountedReplacementPeakAndRollback(t *testing.T) { + const requestedRows = uint64(20_000) + probeMP := mpool.MustNewZero() + var probe Int64HashMap + require.NoError(t, probe.Init(probeMP)) + plan := probe.PlanResize(requestedRows) + require.False(t, plan.Noop) + require.False(t, plan.ReuseCurrentBlocks) + probe.Free() + mpool.DeleteMPool(probeMP) + + descriptorSize := uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) + initialUsed := Int64HashMapInitialAllocationBytes() + descriptorSize + targetDescriptors := plan.TargetBlockCount * descriptorSize + expectedPeak := initialUsed + plan.AdditionalBytes + targetDescriptors + + t.Run("commit", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, expectedPeak, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + require.NoError(t, table.ResizeWithPlan(table.PlanResize(requestedRows))) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+targetDescriptors, + account.Snapshot().Used, + ) + table.Free() + completeHashTableAllocation(t, registry, account) + }) + + t.Run("one byte short", func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, expectedPeak-1, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + beforeCells := table.cells + before := account.Snapshot() + err := table.ResizeWithPlan(table.PlanResize(requestedRows)) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Same(t, &beforeCells[0][0], &table.cells[0][0]) + require.Equal(t, before.Used, account.Snapshot().Used) + require.Equal(t, uint64(2), registry.LiveAllocationMetadata()) + table.Free() + completeHashTableAllocation(t, registry, account) + }) + + t.Run("metadata rollback", func(t *testing.T) { + // Initial descriptor+cell consume two slots. The replacement descriptor + // consumes the last one; the first replacement cell must reject and the + // complete private replacement rolls back before returning. + registry, account, selection := newHashTableAllocation(t, expectedPeak, 3) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + beforeCells := table.cells + before := account.Snapshot() + err := table.ResizeWithPlan(table.PlanResize(requestedRows)) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Same(t, &beforeCells[0][0], &table.cells[0][0]) + require.Equal(t, before.Used, account.Snapshot().Used) + require.Equal(t, uint64(2), registry.LiveAllocationMetadata()) + table.Free() + completeHashTableAllocation(t, registry, account) + }) +} + +func TestIntHashTableAccountedSegmentedGrowth(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 96<<20, 128) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + + firstTarget := maxElemCnt(maxIntCellCntPerBlock, intCellSize) + require.NoError(t, table.ResizeOnDemand(int(firstTarget))) + plan := table.PlanResize(firstTarget + 1) + require.True(t, plan.ReuseCurrentBlocks) + before := account.Snapshot().Used + descriptorSize := uint64(unsafe.Sizeof([]Int64HashMapCell(nil))) + expectedPeak := before + plan.AdditionalBytes + plan.TargetBlockCount*descriptorSize + require.NoError(t, table.ResizeWithPlan(plan)) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+plan.TargetBlockCount*descriptorSize, + account.Snapshot().Used, + ) + + table.Free() + completeHashTableAllocation(t, registry, account) +} + +func TestStringHashTableAccountedReplacementPeakAndRollback(t *testing.T) { + const requestedRows = uint64(20_000) + probeMP := mpool.MustNewZero() + var probe StringHashMap + require.NoError(t, probe.Init(probeMP)) + plan := probe.PlanResize(requestedRows) + require.False(t, plan.Noop) + require.False(t, plan.ReuseCurrentBlocks) + probe.Free() + mpool.DeleteMPool(probeMP) + + descriptorSize := uint64(unsafe.Sizeof([]StringHashMapCell(nil))) + initialUsed := StringHashMapInitialAllocationBytes() + descriptorSize + targetDescriptors := plan.TargetBlockCount * descriptorSize + expectedPeak := initialUsed + plan.AdditionalBytes + targetDescriptors + + for _, tc := range []struct { + name string + limit uint64 + ok bool + }{ + {name: "commit", limit: expectedPeak, ok: true}, + {name: "one byte short", limit: expectedPeak - 1}, + } { + t.Run(tc.name, func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, tc.limit, 64) + mp := mpool.MustNewZero() + var table StringHashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + beforeCells := table.cells + before := account.Snapshot() + err := table.ResizeWithPlan(table.PlanResize(requestedRows)) + if tc.ok { + require.NoError(t, err) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+targetDescriptors, + account.Snapshot().Used, + ) + } else { + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Same(t, &beforeCells[0][0], &table.cells[0][0]) + require.Equal(t, before.Used, account.Snapshot().Used) + require.Equal(t, uint64(2), registry.LiveAllocationMetadata()) + } + table.Free() + completeHashTableAllocation(t, registry, account) + }) + } +} + +func TestStringHashTableAccountedSegmentedGrowth(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 192<<20, 128) + mp := mpool.MustNewZero() + var table StringHashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + + firstTarget := maxElemCnt(maxStrCellCntPerBlock, strCellSize) + require.NoError(t, table.ResizeOnDemand(firstTarget)) + plan := table.PlanResize(firstTarget + 1) + require.True(t, plan.ReuseCurrentBlocks) + before := account.Snapshot().Used + descriptorSize := uint64(unsafe.Sizeof([]StringHashMapCell(nil))) + expectedPeak := before + plan.AdditionalBytes + plan.TargetBlockCount*descriptorSize + require.NoError(t, table.ResizeWithPlan(plan)) + require.Equal(t, expectedPeak, account.Snapshot().Peak) + require.Equal( + t, + plan.NewBytes+plan.TargetBlockCount*descriptorSize, + account.Snapshot().Used, + ) + + table.Free() + completeHashTableAllocation(t, registry, account) +} + +func TestAccountedHashTableNoopAndStalePlanDoNotChangeCharge(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 64<<20, 64) + mp := mpool.MustNewZero() + var table Int64HashMap + require.NoError(t, table.InitWithAllocation(mp, selection)) + + before := account.Snapshot() + noop := table.PlanResize(1) + require.True(t, noop.Noop) + require.NoError(t, table.ResizeWithPlan(noop)) + require.Equal(t, before, account.Snapshot()) + + stale := table.PlanResize(20_000) + require.NoError(t, table.ResizeOnDemand(2_000)) + before = account.Snapshot() + err := table.ResizeWithPlan(stale) + require.ErrorIs(t, err, ErrStaleResizePlan) + require.Equal(t, before, account.Snapshot()) + + table.Free() + completeHashTableAllocation(t, registry, account) +} + +func TestHashTableAccountedHighCardinalityResizeReturnsToZero(t *testing.T) { + const rows = 1_000_000 + for _, tc := range []struct { + name string + run func(*mpool.MPool, *AllocationAccountSelection) error + }{ + { + name: "int", + run: func(mp *mpool.MPool, selection *AllocationAccountSelection) error { + var table Int64HashMap + if err := table.InitWithAllocation(mp, selection); err != nil { + return err + } + defer table.Free() + return table.ResizeOnDemand(rows) + }, + }, + { + name: "string", + run: func(mp *mpool.MPool, selection *AllocationAccountSelection) error { + var table StringHashMap + if err := table.InitWithAllocation(mp, selection); err != nil { + return err + } + defer table.Free() + return table.ResizeOnDemand(rows) + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + registry, account, selection := newHashTableAllocation(t, 256<<20, 256) + mp := mpool.MustNewZero() + require.NoError(t, tc.run(mp, selection)) + require.Positive(t, account.Snapshot().Peak) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + completeHashTableAllocation(t, registry, account) + }) + } +} + +func BenchmarkHashTableResizeAccounting(b *testing.B) { + const rows = 100_000 + b.Run("unaccounted", func(b *testing.B) { + mp := mpool.MustNewZero() + b.ReportAllocs() + b.ResetTimer() + for range b.N { + var table Int64HashMap + if err := table.Init(mp); err != nil { + b.Fatal(err) + } + if err := table.ResizeOnDemand(rows); err != nil { + b.Fatal(err) + } + table.Free() + } + }) + b.Run("accounted", func(b *testing.B) { + registry, account, selection := newHashTableAllocation(b, 256<<20, 256) + mp := mpool.MustNewZero() + b.ReportAllocs() + b.ResetTimer() + for range b.N { + var table Int64HashMap + if err := table.InitWithAllocation(mp, selection); err != nil { + b.Fatal(err) + } + if err := table.ResizeOnDemand(rows); err != nil { + b.Fatal(err) + } + table.Free() + } + b.StopTimer() + completeHashTableAllocation(b, registry, account) + }) +} diff --git a/pkg/container/hashtable/int64_hash_map.go b/pkg/container/hashtable/int64_hash_map.go index 29b1b28bc319b..d9b791f52a326 100644 --- a/pkg/container/hashtable/int64_hash_map.go +++ b/pkg/container/hashtable/int64_hash_map.go @@ -38,6 +38,7 @@ type Int64HashMap struct { cellCnt uint64 elemCnt uint64 cells [][]Int64HashMapCell + account *AllocationAccountSelection version uint64 admit ResizeAdmission @@ -68,19 +69,34 @@ func (ht *Int64HashMap) cellAt(index uint64) *Int64HashMapCell { func (ht *Int64HashMap) Free() { ht.freeCells(ht.cells) ht.cells = nil + ht.account = nil } func (ht *Int64HashMap) freeCells(cells [][]Int64HashMapCell) { for i, block := range cells { - mpool.FreeSlice(ht.mp, block) + freeHashTableCellSlice(ht.mp, block) cells[i] = nil } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) } func (ht *Int64HashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][]Int64HashMapCell, error) { - cells := make([][]Int64HashMapCell, blockCount) + cells, err := makeHashTableDescriptorSlice[[]Int64HashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) + if err != nil { + return nil, err + } for i := range cells { - block, err := mpool.MakeSlice[Int64HashMapCell](int(blockCellCnt), ht.mp, true) + block, err := makeHashTableCellSlice[Int64HashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) if err != nil { ht.freeCells(cells) return nil, err @@ -90,36 +106,83 @@ func (ht *Int64HashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][] return cells, nil } -func (ht *Int64HashMap) allocate(index int, ncells int) error { - if ht.cells[index] != nil { - panic("overwriting") - } - - cell, err := mpool.MakeSlice[Int64HashMapCell](ncells, ht.mp, true) +func (ht *Int64HashMap) appendCells( + blockCount int, + blockCellCnt uint64, +) ([][]Int64HashMapCell, error) { + cells, err := makeHashTableDescriptorSlice[[]Int64HashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) if err != nil { - return err + return nil, err } - ht.cells[index] = cell - return nil + copy(cells, ht.cells) + for i := len(ht.cells); i < len(cells); i++ { + block, allocErr := makeHashTableCellSlice[Int64HashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) + if allocErr != nil { + for j := len(ht.cells); j < i; j++ { + freeHashTableCellSlice(ht.mp, cells[j]) + cells[j] = nil + } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) + return nil, allocErr + } + cells[i] = block + } + return cells, nil } func (ht *Int64HashMap) Init(mp *mpool.MPool) (err error) { + return ht.InitWithAllocation(mp, nil) +} + +func (ht *Int64HashMap) InitWithAllocation( + mp *mpool.MPool, + account *AllocationAccountSelection, +) (err error) { + if account != nil { + if err = account.validate(); err != nil { + return err + } + } ht.mp = mp + ht.account = account ht.blockCellCntBits = kInitialCellCntBits ht.cellCntMask = kInitialCellCnt - 1 ht.elemCnt = 0 ht.cellCnt = kInitialCellCnt ht.version = 0 - ht.cells = make([][]Int64HashMapCell, 1) - - if err = ht.allocate(0, int(ht.blockCellCnt())); err != nil { + if ht.cells, err = ht.allocateCells(1, ht.blockCellCnt()); err != nil { + ht.account = nil return err } return } +func (ht *Int64HashMap) cellSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.cellSite +} + +func (ht *Int64HashMap) descriptorSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.descriptorSite +} + func (ht *Int64HashMap) InsertBatch(n int, hashes []uint64, keysPtr unsafe.Pointer, values []uint64) error { if n <= 0 { return nil @@ -279,15 +342,17 @@ func (ht *Int64HashMap) ResizeWithPlan(plan ResizePlan) error { }() if plan.ReuseCurrentBlocks { - newBlocks, err := ht.allocateCells( - int(plan.TargetBlockCount-plan.CurrentBlockCount), + newCells, err := ht.appendCells( + int(plan.TargetBlockCount), plan.TargetBlockCellCount, ) if err != nil { return err } oldCellCnt := ht.cellCnt - ht.cells = append(ht.cells, newBlocks...) + oldDescriptors := ht.cells + ht.cells = newCells + freeHashTableDescriptorSlice(ht.mp, oldDescriptors, ht.account) ht.cellCnt = plan.TargetCellCount ht.cellCntMask = ht.cellCnt - 1 ht.version++ @@ -344,8 +409,11 @@ func (ht *Int64HashMap) Size() int64 { ret := int64(41) for i := range ht.cells { ret += int64(len(ht.cells[i]) * int(intCellSize)) - // 16 is the len of ht.cells[i] - ret += 16 + } + if ht.account != nil { + ret += int64(len(ht.cells)) * int64(unsafe.Sizeof([]Int64HashMapCell(nil))) + } else { + ret += int64(len(ht.cells)) * 16 } return ret } diff --git a/pkg/container/hashtable/string_hash_map.go b/pkg/container/hashtable/string_hash_map.go index db32e3b582a43..464b2d05f3134 100644 --- a/pkg/container/hashtable/string_hash_map.go +++ b/pkg/container/hashtable/string_hash_map.go @@ -45,6 +45,7 @@ type StringHashMap struct { cellCnt uint64 elemCnt uint64 cells [][]StringHashMapCell + account *AllocationAccountSelection version uint64 admit ResizeAdmission @@ -75,19 +76,34 @@ func (ht *StringHashMap) cellAt(index uint64) *StringHashMapCell { func (ht *StringHashMap) Free() { ht.freeCells(ht.cells) ht.cells = nil + ht.account = nil } func (ht *StringHashMap) freeCells(cells [][]StringHashMapCell) { for i, block := range cells { - mpool.FreeSlice(ht.mp, block) + freeHashTableCellSlice(ht.mp, block) cells[i] = nil } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) } func (ht *StringHashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][]StringHashMapCell, error) { - cells := make([][]StringHashMapCell, blockCount) + cells, err := makeHashTableDescriptorSlice[[]StringHashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) + if err != nil { + return nil, err + } for i := range cells { - block, err := mpool.MakeSlice[StringHashMapCell](int(blockCellCnt), ht.mp, true) + block, err := makeHashTableCellSlice[StringHashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) if err != nil { ht.freeCells(cells) return nil, err @@ -97,34 +113,83 @@ func (ht *StringHashMap) allocateCells(blockCount int, blockCellCnt uint64) ([][ return cells, nil } -func (ht *StringHashMap) allocate(index int, ncells int) error { - if ht.cells[index] != nil { - panic("overwriting") - } - c, err := mpool.MakeSlice[StringHashMapCell](ncells, ht.mp, true) +func (ht *StringHashMap) appendCells( + blockCount int, + blockCellCnt uint64, +) ([][]StringHashMapCell, error) { + cells, err := makeHashTableDescriptorSlice[[]StringHashMapCell]( + blockCount, + ht.mp, + ht.account, + ht.descriptorSite(), + ) if err != nil { - return err + return nil, err } - ht.cells[index] = c - return nil + copy(cells, ht.cells) + for i := len(ht.cells); i < len(cells); i++ { + block, allocErr := makeHashTableCellSlice[StringHashMapCell]( + int(blockCellCnt), + ht.mp, + ht.account, + ht.cellSite(), + ) + if allocErr != nil { + for j := len(ht.cells); j < i; j++ { + freeHashTableCellSlice(ht.mp, cells[j]) + cells[j] = nil + } + freeHashTableDescriptorSlice(ht.mp, cells, ht.account) + return nil, allocErr + } + cells[i] = block + } + return cells, nil } func (ht *StringHashMap) Init(mp *mpool.MPool) (err error) { + return ht.InitWithAllocation(mp, nil) +} + +func (ht *StringHashMap) InitWithAllocation( + mp *mpool.MPool, + account *AllocationAccountSelection, +) (err error) { + if account != nil { + if err = account.validate(); err != nil { + return err + } + } ht.mp = mp + ht.account = account ht.blockCellCntBits = kInitialCellCntBits ht.elemCnt = 0 ht.cellCnt = kInitialCellCnt ht.version = 0 ht.cellCntMask = kInitialCellCnt - 1 - ht.cells = make([][]StringHashMapCell, 1) - if err := ht.allocate(0, int(ht.blockCellCnt())); err != nil { + if ht.cells, err = ht.allocateCells(1, ht.blockCellCnt()); err != nil { + ht.account = nil return err } return } +func (ht *StringHashMap) cellSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.cellSite +} + +func (ht *StringHashMap) descriptorSite() mpool.AllocationSite { + if ht.account == nil { + return 0 + } + return ht.account.descriptorSite +} + func (ht *StringHashMap) InsertStringBatch(states [][3]uint64, keys [][]byte, values []uint64) error { if len(keys) == 0 { return nil @@ -278,15 +343,17 @@ func (ht *StringHashMap) ResizeWithPlan(plan ResizePlan) error { }() if plan.ReuseCurrentBlocks { - newBlocks, err := ht.allocateCells( - int(plan.TargetBlockCount-plan.CurrentBlockCount), + newCells, err := ht.appendCells( + int(plan.TargetBlockCount), plan.TargetBlockCellCount, ) if err != nil { return err } oldCellCnt := ht.cellCnt - ht.cells = append(ht.cells, newBlocks...) + oldDescriptors := ht.cells + ht.cells = newCells + freeHashTableDescriptorSlice(ht.mp, oldDescriptors, ht.account) ht.cellCnt = plan.TargetCellCount ht.cellCntMask = ht.cellCnt - 1 ht.version++ @@ -340,6 +407,9 @@ func (ht *StringHashMap) Size() int64 { for i := range ht.cells { ret += int64(int(strCellSize) * len(ht.cells[i])) } + if ht.account != nil { + ret += int64(len(ht.cells)) * int64(unsafe.Sizeof([]StringHashMapCell(nil))) + } return ret } diff --git a/pkg/container/nulls/nulls.go b/pkg/container/nulls/nulls.go index f6b11432c6d62..e01555d7bd531 100644 --- a/pkg/container/nulls/nulls.go +++ b/pkg/container/nulls/nulls.go @@ -19,6 +19,7 @@ package nulls import ( "fmt" + "io" "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/util" @@ -280,6 +281,19 @@ func Filter(nsp *Nulls, sels []int64, negate bool) { } } +// FilterInPlaceOrdered preserves Filter semantics for Vector.Shrink's ordered +// selection contract without allocating a second row-scaled bitmap. +func FilterInPlaceOrdered(nsp *Nulls, sels []int64, negate bool) { + if nsp.np.EmptyByFlag() { + return + } + if !nsp.np.HasExternalStorage() { + Filter(nsp, sels, negate) + return + } + nsp.np.RemapOrdered(sels, negate) +} + func FilterByMask(nsp *Nulls, sels *bitmap.Bitmap, negate bool) { if nsp.np.EmptyByFlag() { return @@ -331,6 +345,19 @@ func FilterByMask(nsp *Nulls, sels *bitmap.Bitmap, negate bool) { } } +// FilterByMaskInPlace rewrites a null bitmap using the selection bitmap's +// naturally ordered iterator and therefore requires no row-scaled scratch. +func FilterByMaskInPlace(nsp *Nulls, sels *bitmap.Bitmap, negate bool) { + if nsp.np.EmptyByFlag() { + return + } + if !nsp.np.HasExternalStorage() { + FilterByMask(nsp, sels, negate) + return + } + nsp.np.RemapMaskOrdered(sels, negate) +} + // XXX This emptyFlag thing is broken -- it simply cannot be used concurrently. // Make any an alias of EmptyByFlag, otherwise there will be hell lots of race conditions. func (nsp *Nulls) Any() bool { @@ -372,6 +399,20 @@ func (nsp *Nulls) Show() ([]byte, error) { return nsp.np.Marshal(), nil } +func (nsp *Nulls) MarshalSize() int { + if nsp == nil || nsp.np.EmptyByFlag() { + return 0 + } + return nsp.np.MarshalSize() +} + +func (nsp *Nulls) MarshalTo(w io.Writer) error { + if nsp == nil || nsp.np.EmptyByFlag() { + return nil + } + return nsp.np.MarshalTo(w) +} + // ShowV1 in version 1, bitmap is v1 func (nsp *Nulls) ShowV1() ([]byte, error) { if nsp.np.EmptyByFlag() { diff --git a/pkg/container/nulls/nulls_test.go b/pkg/container/nulls/nulls_test.go index 6c9d24cc916d6..b3ef3cb4bc53e 100644 --- a/pkg/container/nulls/nulls_test.go +++ b/pkg/container/nulls/nulls_test.go @@ -15,9 +15,13 @@ package nulls import ( + "bytes" + "fmt" "testing" + "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestOr(t *testing.T) { @@ -68,6 +72,24 @@ func TestAny(t *testing.T) { }) } +func TestMarshalTo(t *testing.T) { + var n Nulls + var empty bytes.Buffer + require.NoError(t, n.MarshalTo(&empty)) + require.Zero(t, n.MarshalSize()) + require.Zero(t, empty.Len()) + + n.InitWithSize(128) + n.Add(3) + n.Add(65) + encoded, err := n.Show() + require.NoError(t, err) + var streamed bytes.Buffer + require.NoError(t, n.MarshalTo(&streamed)) + require.Equal(t, encoded, streamed.Bytes()) + require.Equal(t, len(encoded), n.MarshalSize()) +} + func TestSize(t *testing.T) { t.Run("Size test", func(t *testing.T) { var n Nulls @@ -210,6 +232,57 @@ func TestFilter(t *testing.T) { }) } +func TestFilterInPlaceOrderedMatchesFilter(t *testing.T) { + source := Build(130, 0, 2, 63, 64, 65, 128, 129) + for _, test := range []struct { + name string + sels []int64 + negate bool + }{ + {"select-across-words", []int64{0, 2, 64, 65, 129}, false}, + {"select-trailing-clear", []int64{1, 63, 130, 131}, false}, + {"remove-across-words", []int64{1, 63, 128}, true}, + {"remove-after-bitmap", []int64{130, 131}, true}, + } { + t.Run(test.name, func(t *testing.T) { + legacy := source.Clone() + Filter(legacy, test.sels, test.negate) + + var inPlace Nulls + inPlace.GetBitmap().InstallExternalStorage(make([]uint64, 4)) + inPlace.InitWith(source) + FilterInPlaceOrdered(&inPlace, test.sels, test.negate) + + require.True(t, inPlace.GetBitmap().IsSame(legacy.GetBitmap())) + require.Equal(t, legacy.GetBitmap().Len(), inPlace.GetBitmap().Len()) + require.Equal(t, legacy.Count(), inPlace.Count()) + }) + } +} + +func TestFilterByMaskInPlaceMatchesFilter(t *testing.T) { + source := Build(130, 0, 2, 63, 64, 65, 128, 129) + for _, negate := range []bool{false, true} { + t.Run(fmt.Sprintf("negate=%t", negate), func(t *testing.T) { + var selection bitmap.Bitmap + selection.InitWithSize(132) + selection.AddMany([]uint64{1, 63, 128, 130, 131}) + + legacy := source.Clone() + FilterByMask(legacy, &selection, negate) + + var inPlace Nulls + inPlace.GetBitmap().InstallExternalStorage(make([]uint64, 4)) + inPlace.InitWith(source) + FilterByMaskInPlace(&inPlace, &selection, negate) + + require.True(t, inPlace.GetBitmap().IsSame(legacy.GetBitmap())) + require.Equal(t, legacy.GetBitmap().Len(), inPlace.GetBitmap().Len()) + require.Equal(t, legacy.Count(), inPlace.Count()) + }) + } +} + func TestMerge(t *testing.T) { t.Run("merge test", func(t *testing.T) { var n, m Nulls diff --git a/pkg/container/pSpool/buffer.go b/pkg/container/pSpool/buffer.go index 686eba17ca736..b3f7ae4767935 100644 --- a/pkg/container/pSpool/buffer.go +++ b/pkg/container/pSpool/buffer.go @@ -59,19 +59,26 @@ func (b *spoolBuffer) putCacheID(mp *mpool.MPool, id uint32, bat *batch.Batch) { // 1. const vector size was too small, // 2. vector doesn't own its data and area, // we don't need to cache it. - if vec.IsConst() || vec.NeedDup() { - vec.Free(mp) - } - - data := vector.GetAndClearVecData(vec) - area := vector.GetAndClearVecArea(vec) - - if data != nil { - b.bytesCache[id].bs = append(b.bytesCache[id].bs, data) - } - if area != nil { - b.bytesCache[id].bs = append(b.bytesCache[id].bs, area) + if !vec.IsConst() && !vec.NeedDup() { + data := vector.DetachVectorData(vec) + area := vector.DetachVectorArea(vec) + if data.Capacity() != 0 { + b.bytesCache[id].buffers = append( + b.bytesCache[id].buffers, + data, + ) + } + if area.Capacity() != 0 { + b.bytesCache[id].buffers = append( + b.bytesCache[id].buffers, + area, + ) + } } + // data/area ownership has moved to the cache. Release the remaining + // Vector-owned state, including allocation-accounted bitmap backing, + // before dropping the Vector pointer. + vec.Free(mp) bat.ReplaceVector(vec, nil, i) } @@ -99,9 +106,9 @@ func (b *spoolBuffer) getCacheID() (uint32, *batch.Batch) { func (b *spoolBuffer) clean(mp *mpool.MPool) { for i := range b.bytesCache { - for j := range b.bytesCache[i].bs { - mp.Free(b.bytesCache[i].bs[j]) + for j := range b.bytesCache[i].buffers { + b.bytesCache[i].buffers[j].Free(mp) } - b.bytesCache[i].bs = nil + b.bytesCache[i].buffers = nil } } diff --git a/pkg/container/pSpool/copy.go b/pkg/container/pSpool/copy.go index e3797eeabe75c..27a2a8d6e994b 100644 --- a/pkg/container/pSpool/copy.go +++ b/pkg/container/pSpool/copy.go @@ -36,8 +36,7 @@ type cachedBatch struct { } type oneBatchMemoryCache struct { - // bytes to copy vector's data and area to. - bs [][]byte + buffers []vector.DetachedBuffer } func initCachedBatch(mp *mpool.MPool, capacity uint32) *cachedBatch { @@ -73,6 +72,15 @@ func (cb *cachedBatch) GetCopiedBatch( cacheID, dst = cb.buffer.getCacheID() dst.Recursive = src.Recursive dst.ShuffleIDX = src.ShuffleIDX + if sourceSelection := src.AllocationAccountSelection(); !vector.AllocationAccountSelectionsEqual( + sourceSelection, + dst.AllocationAccountSelection(), + ) { + if err = dst.SetAllocationAccount(sourceSelection); err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } + } if cap(dst.Vecs) >= len(src.Vecs) { dst.Vecs = dst.Vecs[:len(src.Vecs)] @@ -101,7 +109,19 @@ func (cb *cachedBatch) GetCopiedBatch( } typ := *vec.GetType() - dst.Vecs[i] = vector.NewOffHeapVecWithType(typ) + selection := vec.AllocationAccountSelection() + if selection == nil { + dst.Vecs[i] = vector.NewOffHeapVecWithType(typ) + } else { + dst.Vecs[i], err = vector.NewOffHeapVecWithTypeAndAllocation( + typ, + selection, + ) + if err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } + } if vec.IsConst() { if err = vector.GetConstSetFunction(typ, cb.mp)(dst.Vecs[i], vec, 0, vec.Length()); err != nil { @@ -110,8 +130,15 @@ func (cb *cachedBatch) GetCopiedBatch( } } else { - cb.buffer.bytesCache[cacheID].setSuitableDataAreaToVector( - len(vec.GetData()), len(vec.GetArea()), dst.Vecs[i]) + if err = cb.buffer.bytesCache[cacheID]. + setSuitableDataAreaToVector( + len(vec.GetData()), + len(vec.GetArea()), + dst.Vecs[i], + ); err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } dst.Vecs[i].Reset(typ) if err = vector.GetUnionAllFunction(typ, cb.mp)( dst.Vecs[i], @@ -122,6 +149,21 @@ func (cb *cachedBatch) GetCopiedBatch( dst.Vecs[i].SetSorted(vec.GetSorted()) } + if vec.HasGrouping() { + groupingRows := vec.GetGrouping().GetBitmap().Len() + if groupingRows < 0 || groupingRows > int64(math.MaxInt) { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, mpool.ErrAllocationAccountInvalid + } + if err = dst.Vecs[i].PreExtendGrouping( + int(groupingRows), + cb.mp, + ); err != nil { + cb.CacheBatch(true, cacheID, dst) + return nil, false, 0, err + } + dst.Vecs[i].SetGrouping(vec.GetGrouping()) + } dst.Vecs[i].SetIsBin(vec.GetIsBin()) // range src and found the same vector. @@ -144,25 +186,33 @@ func (cb *cachedBatch) GetCopiedBatch( // setSuitableDataAreaToVector get two long-enough bytes slices from the cache, and set them to the vector. // if not found, set the last one to the vector. func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( - dataSize, areaSize int, vec *vector.Vector) { + dataSize, areaSize int, + vec *vector.Vector, +) error { // return directly once cache was empty. - if len(mc.bs) == 0 { - return + if len(mc.buffers) == 0 { + return nil } setDataFirst := dataSize >= areaSize first, second := dataSize, areaSize + firstKind := vector.DetachedDataBuffer + secondKind := vector.DetachedAreaBuffer if !setDataFirst { first, second = areaSize, dataSize + firstKind, secondKind = secondKind, firstKind } if first > 0 { suitIdx := -1 suitDifference := math.MaxInt - for i, bs := range mc.bs { - if difference := cap(bs) - first; difference > 0 { + for i := range mc.buffers { + if !mc.buffers[i].CanAttachTo(vec, firstKind) { + continue + } + if difference := mc.buffers[i].Capacity() - first; difference > 0 { if difference < suitDifference { suitIdx = i suitDifference = difference @@ -172,10 +222,9 @@ func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( if suitIdx != -1 { mem := mc.removeItemAndArrange(suitIdx) - if setDataFirst { - vector.SetVecData(vec, mem) - } else { - vector.SetVecArea(vec, mem) + if err := mem.AttachTo(vec, firstKind); err != nil { + mc.buffers = append(mc.buffers, mem) + return err } } } @@ -184,8 +233,11 @@ func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( suitIdx := -1 suitDifference := math.MaxInt - for i, bs := range mc.bs { - if difference := cap(bs) - second; difference > 0 { + for i := range mc.buffers { + if !mc.buffers[i].CanAttachTo(vec, secondKind) { + continue + } + if difference := mc.buffers[i].Capacity() - second; difference > 0 { if difference < suitDifference { suitIdx = i suitDifference = difference @@ -195,34 +247,70 @@ func (mc *oneBatchMemoryCache) setSuitableDataAreaToVector( if suitIdx != -1 { mem := mc.removeItemAndArrange(suitIdx) - if setDataFirst { - vector.SetVecArea(vec, mem) - } else { - vector.SetVecData(vec, mem) + if err := mem.AttachTo(vec, secondKind); err != nil { + mc.buffers = append(mc.buffers, mem) + return err } } } - if len(mc.bs) > 0 && cap(vec.GetData()) == 0 && dataSize > 0 { - vector.SetVecData(vec, mc.bs[len(mc.bs)-1]) - mc.bs = mc.bs[:len(mc.bs)-1] + if cap(vec.GetData()) == 0 && dataSize > 0 { + if idx := mc.lastAttachable( + vec, + vector.DetachedDataBuffer, + ); idx >= 0 { + mem := mc.removeItemAndArrange(idx) + if err := mem.AttachTo( + vec, + vector.DetachedDataBuffer, + ); err != nil { + mc.buffers = append(mc.buffers, mem) + return err + } + } } - if len(mc.bs) > 0 && cap(vec.GetArea()) == 0 && areaSize > 0 { - vector.SetVecArea(vec, mc.bs[len(mc.bs)-1]) - mc.bs = mc.bs[:len(mc.bs)-1] + if cap(vec.GetArea()) == 0 && areaSize > 0 { + if idx := mc.lastAttachable( + vec, + vector.DetachedAreaBuffer, + ); idx >= 0 { + mem := mc.removeItemAndArrange(idx) + if err := mem.AttachTo( + vec, + vector.DetachedAreaBuffer, + ); err != nil { + mc.buffers = append(mc.buffers, mem) + return err + } + } + } + return nil +} + +func (mc *oneBatchMemoryCache) lastAttachable( + vec *vector.Vector, + kind vector.DetachedBufferKind, +) int { + for i := len(mc.buffers) - 1; i >= 0; i-- { + if mc.buffers[i].CanAttachTo(vec, kind) { + return i + } } + return -1 } // removeItemAndArrange return and remove the idx item of cache. -func (mc *oneBatchMemoryCache) removeItemAndArrange(idx int) []byte { - last := len(mc.bs) - 1 - dst := mc.bs[idx] +func (mc *oneBatchMemoryCache) removeItemAndArrange( + idx int, +) vector.DetachedBuffer { + last := len(mc.buffers) - 1 + dst := mc.buffers[idx] if idx != last { - mc.bs[idx] = mc.bs[last] - mc.bs = mc.bs[:last] + mc.buffers[idx] = mc.buffers[last] } - mc.bs = mc.bs[:last] + mc.buffers[last] = vector.DetachedBuffer{} + mc.buffers = mc.buffers[:last] return dst } diff --git a/pkg/container/pSpool/copy_benchmark_test.go b/pkg/container/pSpool/copy_benchmark_test.go new file mode 100644 index 0000000000000..e7b4e82009b19 --- /dev/null +++ b/pkg/container/pSpool/copy_benchmark_test.go @@ -0,0 +1,58 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package pSpool + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" +) + +func BenchmarkCachedBatchReuse(b *testing.B) { + mp := mpool.MustNewZero() + source := batch.NewWithSize(1) + source.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + values := make([]int64, 8192) + for i := range values { + values[i] = int64(i) + } + if err := vector.AppendFixedList( + source.Vecs[0], + values, + nil, + mp, + ); err != nil { + b.Fatal(err) + } + source.SetRowCount(len(values)) + cache := initCachedBatch(mp, 1) + b.Cleanup(func() { + source.Clean(mp) + cache.free() + }) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + copied, useCache, cacheID, err := cache.GetCopiedBatch(source) + if err != nil { + b.Fatal(err) + } + cache.CacheBatch(useCache, cacheID, copied) + } +} diff --git a/pkg/container/pSpool/sender.go b/pkg/container/pSpool/sender.go index 20d1af7fb46bb..5d78c2001308f 100644 --- a/pkg/container/pSpool/sender.go +++ b/pkg/container/pSpool/sender.go @@ -251,6 +251,45 @@ func (ps *PipelineSpool) ForceCleanupAfterTerminalSignal() { ps.cleanupOnce.Do(ps.forceCleanup) } +// ReleaseReusableCacheAfterProducerQuiesced returns buffers from batches that +// receivers have already released. The producer must have stopped, so no new +// SendBatch call can race to reuse those buffers. Receivers may still return +// current batches afterward; FinalizeAfterConsumersQuiesced performs the final +// cache pass once those receivers have joined. +func (ps *PipelineSpool) ReleaseReusableCacheAfterProducerQuiesced() { + if ps == nil { + return + } + + ps.mu.Lock() + defer ps.mu.Unlock() + ps.cache.free() +} + +// FinalizeAfterConsumersQuiesced releases every batch still retained by the +// transport. The caller must have joined all producer and consumer scopes, so +// no goroutine can subsequently read a queued slot or use a current batch. +// +// A typed terminal signal bypasses the spool queue. Consequently a receiver +// may stop with older GetFromSpool signals still queued even though its scope +// has completed. Those slots are no longer observable and must be reclaimed +// here before the statement allocation account can reach terminal zero. +func (ps *PipelineSpool) FinalizeAfterConsumersQuiesced() { + if ps == nil { + return + } + + ps.mu.Lock() + defer ps.mu.Unlock() + for i := range ps.shardPool { + ps.cleanSlotLocked(uint32(i)) + } + for i := range ps.rs { + ps.rs[i].flagLastPopRelease() + } + ps.cleanupOnce.Do(ps.forceCleanup) +} + // Abort terminates the spool without waiting for receiver acknowledgement. // Pending, not-yet-consumed slots are released immediately. Slots already handed // to receivers stay valid until their receiver calls ReleaseCurrent. The first diff --git a/pkg/container/pSpool/sender_test.go b/pkg/container/pSpool/sender_test.go index cbb6ee8aa9d8a..6a2d05276d1d3 100644 --- a/pkg/container/pSpool/sender_test.go +++ b/pkg/container/pSpool/sender_test.go @@ -99,6 +99,187 @@ func TestPipelineSpoolForceCleanupRetainsUntilReceiversDrained(t *testing.T) { require.Equal(t, int64(0), mp.CurrNB()) } +func TestCachedBatchPreservesAllocationProvenance(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(2, 32) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + equivalentSelection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + require.NotSame(t, selection, equivalentSelection) + otherAccount, err := registry.Open(1 << 20) + require.NoError(t, err) + otherSelection, err := vector.NewAllocationAccountSelection( + otherAccount, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + newSource := func( + value string, + target *vector.AllocationAccountSelection, + ) *batch.Batch { + source := batch.NewOffHeapWithSize(1) + require.NoError(t, source.SetAllocationAccount(target)) + vec := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) + source.SetVector(0, vec) + require.NoError(t, vector.AppendBytes( + vec, + []byte(value), + false, + mp, + )) + require.NoError(t, vec.PreExtendGrouping(1, mp)) + vec.GetGrouping().Add(0) + source.SetRowCount(1) + return source + } + firstSource := newSource("first cached allocation payload", selection) + secondSource := newSource("second", equivalentSelection) + secondSource.Vecs[0].ToConst() + cache := initCachedBatch(mp, 1) + + first, useCache, cacheID, err := cache.GetCopiedBatch(firstSource) + require.NoError(t, err) + require.True(t, useCache) + require.Same(t, selection, first.AllocationAccountSelection()) + require.Same( + t, + selection, + first.Vecs[0].AllocationAccountSelection(), + ) + require.True(t, first.Vecs[0].GetGrouping().Contains(0)) + cache.CacheBatch(useCache, cacheID, first) + beforeReuse := account.Snapshot().Used + + second, useCache, cacheID, err := cache.GetCopiedBatch(secondSource) + require.NoError(t, err) + require.True(t, useCache) + require.Greater(t, account.Snapshot().Used, beforeReuse) + require.Same(t, selection, second.AllocationAccountSelection()) + require.Same( + t, + equivalentSelection, + second.Vecs[0].AllocationAccountSelection(), + ) + require.True(t, second.Vecs[0].GetGrouping().Contains(0)) + cache.CacheBatch(useCache, cacheID, second) + require.Equal(t, beforeReuse, account.Snapshot().Used) + + otherSource := newSource("other allocation account", otherSelection) + firstAccountBeforeOther := account.Snapshot().Used + otherBeforeCopy := otherAccount.Snapshot().Used + other, useCache, cacheID, err := cache.GetCopiedBatch(otherSource) + require.NoError(t, err) + require.Equal(t, firstAccountBeforeOther, account.Snapshot().Used) + require.Greater(t, otherAccount.Snapshot().Used, otherBeforeCopy) + require.Same( + t, + otherSelection, + other.Vecs[0].AllocationAccountSelection(), + ) + cache.CacheBatch(useCache, cacheID, other) + + firstSource.Clean(mp) + secondSource.Clean(mp) + otherSource.Clean(mp) + cache.free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, otherAccount.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + account.Seal() + otherAccount.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) + _, err = registry.Finalize(otherAccount) + require.NoError(t, err) +} + +func TestCachedBatchAllocationFailureReturnsCacheOwnership(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 4) + require.NoError(t, err) + account, err := registry.Open(64) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + source := batch.NewOffHeapWithSize(1) + require.NoError(t, source.SetAllocationAccount(selection)) + source.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, vector.AppendFixedList( + source.Vecs[0], + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + nil, + mp, + )) + source.SetRowCount(8) + require.Equal(t, uint64(64), account.Snapshot().Used) + cache := initCachedBatch(mp, 1) + + _, _, _, err = cache.GetCopiedBatch(source) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Len(t, cache.buffer.readyToUse, 1) + require.Equal(t, uint64(64), account.Snapshot().Used) + + source.Clean(mp) + cache.free() + require.Zero(t, account.Snapshot().Used) + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) +} + +func TestUnaccountedCacheNonLastSelectionRetainsOwnership(t *testing.T) { + mp := mpool.MustNewZero() + cache := oneBatchMemoryCache{} + for _, size := range []int{64, 128, 256} { + owner := vector.NewOffHeapVecWithType(types.T_int8.ToType()) + require.NoError(t, owner.PreExtend(size, mp)) + cache.buffers = append(cache.buffers, vector.DetachVectorData(owner)) + owner.Free(mp) + } + vec := vector.NewOffHeapVecWithType(types.T_int8.ToType()) + require.NoError(t, cache.setSuitableDataAreaToVector( + 100, + 0, + vec, + )) + require.Len(t, cache.buffers, 2) + + vec.Free(mp) + for i := range cache.buffers { + cache.buffers[i].Free(mp) + } + require.Zero(t, mp.CurrNB()) +} + func TestPipelineSpoolForceCleanupAfterTerminalSignalDoesNotNeedNilEndMessage(t *testing.T) { mp := mpool.MustNewZeroNoFixed() t.Cleanup(func() { @@ -132,6 +313,103 @@ func TestPipelineSpoolForceCleanupAfterTerminalSignalDoesNotNeedNilEndMessage(t require.Equal(t, int64(0), mp.CurrNB()) } +func TestPipelineSpoolReleaseReusableCacheBeforeFinalization(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := InitMyPipelineSpool(mp, 2) + for range 2 { + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + } + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + sp.ReleaseCurrent(0) + got, info = sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + beforeRelease := mp.CurrNB() + + sp.ReleaseReusableCacheAfterProducerQuiesced() + require.Positive(t, mp.CurrNB()) + require.Less(t, mp.CurrNB(), beforeRelease) + + sp.ReleaseCurrent(0) + require.Positive(t, mp.CurrNB()) + sp.FinalizeAfterConsumersQuiesced() + require.Zero(t, mp.CurrNB()) +} + +func TestPipelineSpoolFinalizeAfterConsumersQuiescedReleasesPendingBatch(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + require.Greater(t, mp.CurrNB(), int64(0)) + + sp.ForceCleanupAfterTerminalSignal() + require.Greater(t, mp.CurrNB(), int64(0)) + + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) +} + +func TestPipelineSpoolFinalizeAfterConsumersQuiescedReleasesBroadcastCurrentAndPending(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := InitMyPipelineSpool(mp, 2) + done, err := sp.SendBatch(context.Background(), SendToAllLocal, src, nil) + require.NoError(t, err) + require.False(t, done) + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + require.Greater(t, mp.CurrNB(), int64(0)) + + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) + sp.FinalizeAfterConsumersQuiesced() + require.Equal(t, int64(0), mp.CurrNB()) +} + func TestPipelineSpoolLateReleaseAfterTerminalCleanupFreesDirectly(t *testing.T) { mp := mpool.MustNewZeroNoFixed() t.Cleanup(func() { diff --git a/pkg/container/types/packer.go b/pkg/container/types/packer.go index 5aae6cd942c3b..3cc955aaeb729 100644 --- a/pkg/container/types/packer.go +++ b/pkg/container/types/packer.go @@ -23,15 +23,22 @@ import ( "unsafe" "github.com/matrixorigin/matrixone/pkg/common/malloc" + "github.com/matrixorigin/matrixone/pkg/common/moerr" ) type Packer struct { buffer []byte bufferDeallocator malloc.Deallocator + fixed bool + overflow bool } const defaultPackerSize = uint64(4096) +var ErrPackerCapacity = moerr.NewInternalErrorNoCtx( + "packer fixed buffer capacity exceeded", +) + var packerAllocator = malloc.NewShardedAllocator( runtime.GOMAXPROCS(0), func() malloc.Allocator { @@ -45,16 +52,13 @@ func NewPacker() *Packer { return NewPackerWithSize(defaultPackerSize) } -// PackerAllocationSize returns the backing size-class allocation made by -// NewPackerWithSize. It lets memory-governed callers reserve the actual -// allocation, including allocator rounding, before constructing a packer. +// PackerAllocationSize returns the backing size-class allocation made by a +// packer request of size bytes. func PackerAllocationSize(size uint64) (uint64, bool) { return malloc.ClassAllocationSize(size) } // DefaultPackerCapacity returns the backing capacity retained by NewPacker. -// Memory-governed callers use this value to admit construction before the -// allocator is entered. func DefaultPackerCapacity() uint64 { size, ok := PackerAllocationSize(defaultPackerSize) if !ok { @@ -66,12 +70,6 @@ func DefaultPackerCapacity() uint64 { // PackerCapacityUpperBound bounds the backing capacity retained after any // sequence of Reset and append operations whose logical buffer length never // exceeds maxLength and whose individual append never exceeds maxAppend. -// -// ensureSizeSlow requests cap(buffer)+append from the class allocator. At a -// growth point the old capacity is strictly less than maxLength, so the -// request is at most maxLength+maxAppend-1. If the old capacity already covers -// maxLength, no growth occurs. This mirrors the allocator contract without -// replaying input-dependent append sequences in an admission hot path. func PackerCapacityUpperBound(maxLength, maxAppend uint64) (uint64, bool) { initial := DefaultPackerCapacity() if maxLength <= initial { @@ -80,8 +78,7 @@ func PackerCapacityUpperBound(maxLength, maxAppend uint64) (uint64, bool) { if maxAppend == 0 || maxLength > math.MaxUint64-maxAppend+1 { return 0, false } - request := maxLength + maxAppend - 1 - capacity, ok := PackerAllocationSize(request) + capacity, ok := PackerAllocationSize(maxLength + maxAppend - 1) if !ok { return 0, false } @@ -102,6 +99,15 @@ func NewPackerWithSize(size uint64) *Packer { } } +// NewPackerWithFixedBuffer uses caller-owned storage and never allocates. +// Err reports an encoding that exceeded the supplied physical capacity. +func NewPackerWithFixedBuffer(buffer []byte) *Packer { + return &Packer{ + buffer: buffer[:0:cap(buffer)], + fixed: true, + } +} + func NewPackerArray(length int) []*Packer { return NewPackerArrayWithSize(length, defaultPackerSize) } @@ -123,6 +129,7 @@ func (p *Packer) Close() { func (p *Packer) Reset() { p.buffer = p.buffer[:0] + p.overflow = false } // Allocated returns the size-class capacity retained by the packer. @@ -133,14 +140,11 @@ func (p *Packer) Allocated() uint64 { return uint64(cap(p.buffer)) } -func (p *Packer) ensureSize(n int) { - if len(p.buffer)+n <= cap(p.buffer) { +func (p *Packer) ensureSizeSlow(n int) { + if p.fixed { + p.overflow = true return } - p.ensureSizeSlow(n) -} - -func (p *Packer) ensureSizeSlow(n int) { newBuffer, newDec, err := packerAllocator.Allocate(uint64(cap(p.buffer)+n), malloc.NoClear) if err != nil { panic(err) @@ -155,12 +159,26 @@ func (p *Packer) ensureSizeSlow(n int) { } func (p *Packer) putByte(b byte) { - p.ensureSize(1) + if len(p.buffer) < cap(p.buffer) { + p.buffer = append(p.buffer, b) + return + } + p.ensureSizeSlow(1) + if p.overflow { + return + } p.buffer = append(p.buffer, b) } func (p *Packer) putBytes(bs []byte) { - p.ensureSize(len(bs)) + if len(bs) <= cap(p.buffer)-len(p.buffer) { + p.buffer = append(p.buffer, bs...) + return + } + p.ensureSizeSlow(len(bs)) + if p.overflow { + return + } p.buffer = append(p.buffer, bs...) } @@ -380,6 +398,16 @@ func (p *Packer) GetBuf() []byte { return p.buffer } +func (p *Packer) Err() error { + if p == nil { + return ErrPackerCapacity + } + if p.overflow { + return ErrPackerCapacity + } + return nil +} + func (p *Packer) Bytes() []byte { return slices.Clone(p.buffer) } diff --git a/pkg/container/types/packer_test.go b/pkg/container/types/packer_test.go index ef95f4f57cbd3..ec18ba8120297 100644 --- a/pkg/container/types/packer_test.go +++ b/pkg/container/types/packer_test.go @@ -15,10 +15,8 @@ package types import ( - "math" + "errors" "testing" - - "github.com/stretchr/testify/require" ) func TestPacker(t *testing.T) { @@ -46,40 +44,29 @@ func TestClosedPackerIsOK(t *testing.T) { packer.Close() } -func TestPackerCapacityUpperBoundCoversGrowthAndReuse(t *testing.T) { - const ( - columns = 4 - width = 65535 - ) - component := uint64(2*width + 3) - maxLength := uint64(columns) * component - bound, ok := PackerCapacityUpperBound(maxLength, component) - require.True(t, ok) - require.Equal(t, uint64(1<<20), bound) - - packer := NewPacker() - require.Equal(t, DefaultPackerCapacity(), packer.Allocated()) - value := make([]byte, width) - for range 2 { - packer.Reset() - for range columns { - packer.EncodeStringType(value) - } - require.Equal(t, int(maxLength), len(packer.GetBuf())) - require.LessOrEqual(t, packer.Allocated(), bound) +func TestFixedBufferPackerNeverAllocatesPastCapacity(t *testing.T) { + storage := make([]byte, 0, 3) + packer := NewPackerWithFixedBuffer(storage) + packer.EncodeBool(true) + packer.EncodeNull() + if err := packer.Err(); err != nil { + t.Fatal(err) + } + if len(packer.GetBuf()) != 2 { + t.Fatalf("encoded length = %d", len(packer.GetBuf())) + } + packer.EncodeInt64(42) + if !errors.Is(packer.Err(), ErrPackerCapacity) { + t.Fatalf("overflow error = %v", packer.Err()) + } + if len(packer.GetBuf()) > cap(storage) { + t.Fatal("fixed packer exceeded caller-owned storage") + } + packer.Reset() + packer.EncodeBool(true) + if err := packer.Err(); err != nil { + t.Fatalf("reset fixed packer error = %v", err) } - packer.Close() - require.Zero(t, packer.Allocated()) - - initial := DefaultPackerCapacity() - got, ok := PackerCapacityUpperBound(initial, initial) - require.True(t, ok) - require.Equal(t, initial, got) - got, ok = PackerCapacityUpperBound(initial+1, 1) - require.True(t, ok) - require.Equal(t, 2*initial, got) - _, ok = PackerCapacityUpperBound(math.MaxUint64, 2) - require.False(t, ok) } func BenchmarkPacker(b *testing.B) { @@ -91,11 +78,29 @@ func BenchmarkPacker(b *testing.B) { } func BenchmarkPackerEncode(b *testing.B) { - packer := NewPacker() - defer packer.Close() - b.ResetTimer() - for i := 0; i < b.N; i++ { - packer.EncodeInt64(42) - packer.Reset() + for _, fixed := range []bool{false, true} { + mode := "allocator-backed" + if fixed { + mode = "fixed-buffer" + } + b.Run(mode, func(b *testing.B) { + var packer *Packer + if fixed { + packer = NewPackerWithFixedBuffer(make([]byte, 16)) + } else { + packer = NewPacker() + defer packer.Close() + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + packer.EncodeInt64(42) + packer.Reset() + } + b.StopTimer() + if err := packer.Err(); err != nil { + b.Fatal(err) + } + }) } } diff --git a/pkg/container/vector/allocation_account.go b/pkg/container/vector/allocation_account.go new file mode 100644 index 0000000000000..4561dff3f0523 --- /dev/null +++ b/pkg/container/vector/allocation_account.go @@ -0,0 +1,621 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package vector + +import ( + "errors" + "io" + "math" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" +) + +func allocationAccountInvalid(message string) error { + return errors.Join( + mpool.ErrAllocationAccountInvalid, + moerr.NewInternalErrorNoCtx(message), + ) +} + +// AllocationAccountSelection is an immutable choice for the first owned +// off-heap allocations of a Vector. The physical MPool allocation metadata +// remains the sole owner of the resulting charge. +// +// A selection may be shared by all vectors owned by one Batch. Views do not +// copy it: they share storage and therefore must not create a second charge. +type AllocationAccountSelection struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + dataSite mpool.AllocationSite + areaSite mpool.AllocationSite + nullsSite mpool.AllocationSite + groupingSite mpool.AllocationSite + capacityClass mpool.AllocationCapacityClass +} + +// AllocationAccountSelectionsEqual reports whether two immutable selections +// describe the same physical allocation provenance. Separately constructed +// selections are interchangeable only when they charge the same account, +// owner, and allocation sites. +func AllocationAccountSelectionsEqual( + left, right *AllocationAccountSelection, +) bool { + if left == right { + return true + } + return left != nil && right != nil && + left.account == right.account && + left.owner == right.owner && + left.dataSite == right.dataSite && + left.areaSite == right.areaSite && + left.nullsSite == right.nullsSite && + left.groupingSite == right.groupingSite && + left.capacityClass == right.capacityClass +} + +func NewAllocationAccountSelection( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + dataSite mpool.AllocationSite, + areaSite mpool.AllocationSite, + nullsSite mpool.AllocationSite, + groupingSite mpool.AllocationSite, +) (*AllocationAccountSelection, error) { + return NewAllocationAccountSelectionWithCapacityClass( + account, + owner, + dataSite, + areaSite, + nullsSite, + groupingSite, + mpool.AllocationCapacityClassDefault, + ) +} + +// NewAllocationAccountSelectionWithCapacityClass applies one execution-local +// capacity class to every physical vector backing allocation. +func NewAllocationAccountSelectionWithCapacityClass( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + dataSite mpool.AllocationSite, + areaSite mpool.AllocationSite, + nullsSite mpool.AllocationSite, + groupingSite mpool.AllocationSite, + capacityClass mpool.AllocationCapacityClass, +) (*AllocationAccountSelection, error) { + selection := &AllocationAccountSelection{ + account: account, + owner: owner, + dataSite: dataSite, + areaSite: areaSite, + nullsSite: nullsSite, + groupingSite: groupingSite, + capacityClass: capacityClass, + } + if err := selection.validate(); err != nil { + return nil, err + } + return selection, nil +} + +// NewOffHeapVecWithTypeAndAllocation constructs an empty owning Vector whose +// future allocations use selection. +func NewOffHeapVecWithTypeAndAllocation( + typ types.Type, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec := NewOffHeapVecWithType(typ) + if err := vec.SetAllocationAccount(selection); err != nil { + vec.Free(nil) + return nil, err + } + return vec, nil +} + +// NewConstNullWithAllocation constructs a constant NULL Vector with explicit +// allocation provenance for any future owned backing. +func NewConstNullWithAllocation( + typ types.Type, + length int, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + vec.length = length + return vec, nil +} + +// NewConstFixedWithAllocation constructs an off-heap constant fixed-width +// Vector and charges its physical backing to selection. +func NewConstFixedWithAllocation[T any]( + typ types.Type, + value T, + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + if length > 0 { + if err = SetConstFixed(vec, value, length, mp); err != nil { + vec.Free(mp) + return nil, err + } + } + return vec, nil +} + +// NewConstBytesWithAllocation constructs an off-heap constant varlen Vector +// and charges its backing independently through selection. +func NewConstBytesWithAllocation( + typ types.Type, + value []byte, + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + if length > 0 { + if err = SetConstBytes(vec, value, length, mp); err != nil { + vec.Free(mp) + return nil, err + } + } + return vec, nil +} + +// NewConstArrayWithAllocation constructs an off-heap constant array Vector +// and charges its backing independently through selection. +func NewConstArrayWithAllocation[T types.ArrayElement]( + typ types.Type, + value []T, + length int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + vec, err := NewOffHeapVecWithTypeAndAllocation(typ, selection) + if err != nil { + return nil, err + } + vec.class = CONSTANT + if length > 0 { + if err = SetConstArray(vec, value, length, mp); err != nil { + vec.Free(mp) + return nil, err + } + } + return vec, nil +} + +func (s *AllocationAccountSelection) validate() error { + if s == nil || s.account == nil || s.account.Handle() == 0 || + s.owner < mpool.AllocationOwnerMin || + s.owner > mpool.AllocationOwnerMax || + s.dataSite < mpool.AllocationSiteMin || + s.areaSite < mpool.AllocationSiteMin || + s.nullsSite < mpool.AllocationSiteMin || + s.groupingSite < mpool.AllocationSiteMin { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +// AllocationAccountSelection returns the immutable selection used by this +// vector's future owned allocations. It is nil for unaccounted vectors and +// views outside the retained HashBuild domain. +func (v *Vector) AllocationAccountSelection() *AllocationAccountSelection { + if v == nil { + return nil + } + return v.allocationAccount +} + +// CanSetAllocationAccount reports whether selection can be installed without +// converting or relabeling an existing physical allocation. +func (v *Vector) CanSetAllocationAccount( + selection *AllocationAccountSelection, +) error { + if v == nil { + return mpool.ErrAllocationAccountInvalid + } + if selection != nil { + if err := selection.validate(); err != nil { + return err + } + if !v.offHeap { + return allocationAccountInvalid( + "allocation-accounted vector must be off-heap", + ) + } + } + if AllocationAccountSelectionsEqual(v.allocationAccount, selection) { + return nil + } + if v.hasBackingStorage() { + return allocationAccountInvalid("vector already has backing storage") + } + return nil +} + +func (v *Vector) hasBackingStorage() bool { + return cap(v.data) != 0 || + cap(v.area) != 0 || + v.nsp.GetBitmap().Size() != 0 || + v.gsp.GetBitmap().Size() != 0 || + v.nsp.GetBitmap().ExternalStorageCapacity() != 0 || + v.gsp.GetBitmap().ExternalStorageCapacity() != 0 +} + +// hasOwnedBackingStorage reports storage that UnmarshalBinary cannot replace +// without losing an MPool-owned allocation. Data and area marked cantFree are +// borrowed aliases; ordinary bitmap backing is Go-owned and remains GC-visible +// after replacement. Accounted bitmap storage is explicit external storage. +func (v *Vector) hasOwnedBackingStorage() bool { + return cap(v.data) != 0 && !v.cantFreeData || + cap(v.area) != 0 && !v.cantFreeArea || + v.nsp.GetBitmap().ExternalStorageCapacity() != 0 || + v.gsp.GetBitmap().ExternalStorageCapacity() != 0 +} + +// SetAllocationAccount selects the account used by future owned allocations. +// It is intentionally explicit and is legal only before the first backing +// allocation. Reset retains the selection; Free clears it. +func (v *Vector) SetAllocationAccount( + selection *AllocationAccountSelection, +) error { + if err := v.CanSetAllocationAccount(selection); err != nil { + return err + } + if AllocationAccountSelectionsEqual(v.allocationAccount, selection) { + return nil + } + if v.allocationAccount != nil && selection == nil { + v.nsp.GetBitmap().ReleaseExternalStorage() + v.gsp.GetBitmap().ReleaseExternalStorage() + } + v.allocationAccount = selection + if selection != nil { + v.nsp.GetBitmap().InstallExternalStorage(nil) + v.gsp.GetBitmap().InstallExternalStorage(nil) + } + return nil +} + +func (v *Vector) ensureBitmapCapacity(rows int, mp *mpool.MPool) error { + if v.allocationAccount == nil { + return nil + } + if rows < 0 || rows > math.MaxInt-64 || mp == nil { + return mpool.ErrAllocationAccountInvalid + } + // Nulls.AddRange currently expands through end+1 even though end is + // exclusive. Keep one admitted sentinel bit so raw bitmap mutation cannot + // escape to a Go allocation at the vector's logical row boundary. + if rows > 0 { + rows++ + } + requiredWords := (rows + 63) / 64 + if requiredWords <= v.nsp.GetBitmap().ExternalStorageCapacity() && + requiredWords <= v.gsp.GetBitmap().ExternalStorageCapacity() { + return nil + } + nulls, err := v.allocateBitmapGrowth( + v.nsp.GetBitmap(), + rows, + mp, + v.allocationAccount.nullsSite, + ) + if err != nil { + return err + } + grouping, err := v.allocateBitmapGrowth( + v.gsp.GetBitmap(), + rows, + mp, + v.allocationAccount.groupingSite, + ) + if err != nil { + mpool.FreeSlice(mp, nulls) + return err + } + if cap(nulls) > 0 { + previous := v.nsp.GetBitmap().InstallExternalStorage(nulls) + mpool.FreeSlice(mp, previous) + } + if cap(grouping) > 0 { + previous := v.gsp.GetBitmap().InstallExternalStorage(grouping) + mpool.FreeSlice(mp, previous) + } + return nil +} + +func (v *Vector) ensureNullCapacity(rows int, mp *mpool.MPool) error { + if v.allocationAccount == nil { + return nil + } + return v.ensureSingleBitmapCapacity( + v.nsp.GetBitmap(), + rows, + mp, + v.allocationAccount.nullsSite, + ) +} + +func (v *Vector) ensureGroupingCapacity(rows int, mp *mpool.MPool) error { + if v.allocationAccount == nil { + return nil + } + return v.ensureSingleBitmapCapacity( + v.gsp.GetBitmap(), + rows, + mp, + v.allocationAccount.groupingSite, + ) +} + +func (v *Vector) ensureSingleBitmapCapacity( + value *bitmap.Bitmap, + rows int, + mp *mpool.MPool, + site mpool.AllocationSite, +) error { + if rows < 0 || rows > math.MaxInt-64 || mp == nil { + return mpool.ErrAllocationAccountInvalid + } + if rows > 0 { + rows++ + } + storage, err := v.allocateBitmapGrowth(value, rows, mp, site) + if err != nil { + return err + } + if cap(storage) > 0 { + previous := value.InstallExternalStorage(storage) + mpool.FreeSlice(mp, previous) + } + return nil +} + +func (v *Vector) allocateBitmapGrowth( + value *bitmap.Bitmap, + rows int, + mp *mpool.MPool, + site mpool.AllocationSite, +) ([]uint64, error) { + requiredWords := (rows + 63) / 64 + if requiredWords <= value.ExternalStorageCapacity() { + return nil, nil + } + requiredBytes := int64(requiredWords) * 8 + oldBytes := int64(value.ExternalStorageCapacity()) * 8 + newBytes, ok := mpool.GrowCapacity(oldBytes, requiredBytes) + if !ok || newBytes > int64(math.MaxInt) || newBytes%8 != 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + next, err := mpool.MakeSliceAccountedWithCapacityClass[uint64]( + int(newBytes/8), + mp, + v.allocationAccount.account, + v.allocationAccount.owner, + site, + v.allocationAccount.capacityClass, + ) + if err != nil { + return nil, err + } + clear(next) + return next, nil +} + +func (v *Vector) freeBitmapStorage(mp *mpool.MPool) { + for _, value := range []*bitmap.Bitmap{ + v.nsp.GetBitmap(), + v.gsp.GetBitmap(), + } { + storage := value.ReleaseExternalStorage() + if cap(storage) > 0 { + mpool.FreeSlice(mp, storage) + } + } +} + +func (v *Vector) allocData(mp *mpool.MPool, size int) ([]byte, error) { + return v.allocOwned(mp, size, v.offHeap, true) +} + +func (v *Vector) allocArea(mp *mpool.MPool, size int) ([]byte, error) { + return v.allocOwned(mp, size, v.offHeap, false) +} + +func (v *Vector) allocOwned( + mp *mpool.MPool, + size int, + offHeap bool, + data bool, +) ([]byte, error) { + if mp == nil { + return nil, moerr.NewInternalErrorNoCtx( + "vector allocation does not have a mpool", + ) + } + if v.allocationAccount == nil { + return mp.Alloc(size, offHeap) + } + if !offHeap { + return nil, allocationAccountInvalid( + "accounted allocation must be off-heap", + ) + } + site := v.allocationAccount.areaSite + if data { + site = v.allocationAccount.dataSite + } + return mp.AllocAccountedWithCapacityClass( + size, + v.allocationAccount.account, + v.allocationAccount.owner, + site, + v.allocationAccount.capacityClass, + ) +} + +func (v *Vector) growData(mp *mpool.MPool, size int) ([]byte, error) { + return v.growOwned(mp, v.data, size, true) +} + +func (v *Vector) growArea(mp *mpool.MPool, size int) ([]byte, error) { + return v.growOwned(mp, v.area, size, false) +} + +func (v *Vector) growOwned( + mp *mpool.MPool, + old []byte, + size int, + data bool, +) ([]byte, error) { + if size <= cap(old) { + return old[:size], nil + } + if mp == nil { + return nil, moerr.NewInternalErrorNoCtx( + "vector growth does not have a mpool", + ) + } + if cap(old) != 0 || v.allocationAccount == nil { + return mp.Grow(old, size, v.offHeap) + } + + capacity, ok := mpool.GrowCapacity(0, int64(size)) + if !ok { + return nil, moerr.NewInternalErrorNoCtxf( + "invalid mpool grow capacity, old %d, required %d", + cap(old), + size, + ) + } + buf, err := v.allocOwned(mp, int(capacity), true, data) + if err != nil { + return nil, err + } + return buf[:size], nil +} + +func (v *Vector) growArea2( + mp *mpool.MPool, + src []byte, + size int, +) ([]byte, error) { + oldLen := len(v.area) + if size < oldLen+len(src) { + return nil, moerr.NewInternalErrorNoCtxf( + "mpool grow2 actually shrinks, %d+%d, %d", + oldLen, + len(src), + size, + ) + } + grown, err := v.growArea(mp, size) + if err != nil { + return nil, err + } + copy(grown[oldLen:oldLen+len(src)], src) + return grown, nil +} + +func (v *Vector) readSizeBytes( + r io.Reader, + mp *mpool.MPool, + data bool, +) (int32, []byte, error) { + size, err := types.ReadInt32(r) + if err != nil { + return 0, nil, err + } + var old []byte + if data { + old = v.data + } else { + old = v.area + } + if size == 0 { + if old != nil { + old = old[:0] + } + if data { + v.data = old + } else { + v.area = old + } + return 0, old, nil + } + if size < 0 { + return size, nil, moerr.NewInvalidInputNoCtx( + "negative vector buffer size", + ) + } + if err := validateStreamingReadSize(r, int64(size)); err != nil { + return size, nil, err + } + var buf []byte + if data { + buf, err = v.growData(mp, int(size)) + } else { + buf, err = v.growArea(mp, int(size)) + } + if err != nil { + return 0, nil, err + } + // Grow may already have freed the old allocation. Publish its replacement + // before reading so a short reader still leaves one reachable cleanup owner. + if data { + v.data = buf + } else { + v.area = buf + } + if _, err = io.ReadFull(r, buf); err != nil { + return size, buf, err + } + return size, buf, nil +} + +func validateStreamingReadSize(r io.Reader, size int64) error { + if size < 0 { + return moerr.NewInvalidInputNoCtx("negative vector buffer size") + } + var remaining int64 = -1 + switch reader := r.(type) { + case *io.LimitedReader: + remaining = reader.N + case interface{ Len() int }: + remaining = int64(reader.Len()) + } + if remaining >= 0 && size > remaining { + return io.ErrUnexpectedEOF + } + return nil +} diff --git a/pkg/container/vector/allocation_account_test.go b/pkg/container/vector/allocation_account_test.go new file mode 100644 index 0000000000000..787e063a3e316 --- /dev/null +++ b/pkg/container/vector/allocation_account_test.go @@ -0,0 +1,1480 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package vector + +import ( + "bytes" + "errors" + "math" + "math/rand" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/stretchr/testify/require" +) + +const ( + testVectorAllocationOwner mpool.AllocationOwner = 1 + testVectorDataAllocationSite mpool.AllocationSite = 1 + testVectorAreaAllocationSite mpool.AllocationSite = 2 + testVectorNullAllocationSite mpool.AllocationSite = 3 + testVectorGroupAllocationSite mpool.AllocationSite = 4 +) + +type testVectorAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + selection *AllocationAccountSelection +} + +func newTestVectorAllocationAccount( + t testing.TB, + limit uint64, + allocationSlots uint64, +) testVectorAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, allocationSlots) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := NewAllocationAccountSelection( + account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + return testVectorAllocationAccount{ + registry: registry, + account: account, + selection: selection, + } +} + +func finalizeTestVectorAllocationAccount( + t testing.TB, + state testVectorAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func newAccountedTestVector( + t testing.TB, + typ types.Type, + selection *AllocationAccountSelection, +) *Vector { + t.Helper() + vec := NewOffHeapVecWithType(typ) + require.NoError(t, vec.SetAllocationAccount(selection)) + return vec +} + +func TestVectorAllocationAccountConfiguration(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + + _, err := NewAllocationAccountSelection( + nil, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + + onHeap := NewVec(types.T_int64.ToType()) + require.ErrorIs( + t, + onHeap.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.Same(t, state.selection, vec.AllocationAccountSelection()) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + require.ErrorIs( + t, + vec.SetAllocationAccount(nil), + mpool.ErrAllocationAccountInvalid, + ) + require.Panics(t, func() { + vec.SetOffHeap(false) + }) + sourceUsed := state.account.Snapshot().Used + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Same(t, state.selection, dup.AllocationAccountSelection()) + require.Greater(t, state.account.Snapshot().Used, sourceUsed) + dup.Free(mp) + require.Equal(t, sourceUsed, state.account.Snapshot().Used) + _, err = vec.CloneToFlatCompact(mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + + vec.Free(mp) + require.Nil(t, vec.AllocationAccountSelection()) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestAllocationAccountSelectionsEqual(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + equivalent, err := NewAllocationAccountSelection( + state.account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + differentSite, err := NewAllocationAccountSelection( + state.account, + testVectorAllocationOwner, + testVectorDataAllocationSite+1, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + + require.NotSame(t, state.selection, equivalent) + require.True(t, AllocationAccountSelectionsEqual(state.selection, equivalent)) + require.False(t, AllocationAccountSelectionsEqual(state.selection, differentSite)) + require.False(t, AllocationAccountSelectionsEqual(state.selection, nil)) + require.True(t, AllocationAccountSelectionsEqual(nil, nil)) + + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + require.NoError(t, vec.CanSetAllocationAccount(equivalent)) + require.NoError(t, vec.SetAllocationAccount(equivalent)) + // Equivalent provenance is a no-op: existing physical ownership remains + // attached to the original immutable selection. + require.Same(t, state.selection, vec.AllocationAccountSelection()) + vec.Free(mp) + + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountFixedResetReuseAndFree(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + + require.NoError(t, vec.PreExtend(128, mp)) + initial := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)), initial.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + for i := 0; i < 64; i++ { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.ResetWithSameType() + require.Equal(t, initial.Used, state.account.Snapshot().Used) + for i := 0; i < 64; i++ { + require.NoError(t, AppendFixed(vec, int64(i*2), false, mp)) + } + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + growAt := vec.Capacity() + 1 + for vec.Length() < growAt { + require.NoError(t, AppendFixed(vec, int64(vec.Length()), false, mp)) + } + grown := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)), grown.Used) + require.Greater(t, grown.Peak, grown.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + require.NoError(t, vec.Shuffle([]int64{0, 2, 4, 6}, mp)) + shuffled := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)), shuffled.Used) + require.Equal(t, uint64(1), state.registry.LiveAllocationMetadata()) + + vec.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountVarlenaDataAndArea(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + + require.NoError(t, vec.PreExtendWithArea(64, 4096, mp)) + initial := state.account.Snapshot() + require.Equal(t, uint64(cap(vec.data)+cap(vec.area)), initial.Used) + require.Equal(t, uint64(2), state.registry.LiveAllocationMetadata()) + + for i := 0; i < 32; i++ { + require.NoError(t, AppendBytes(vec, bytes.Repeat([]byte{byte(i)}, 64), false, mp)) + } + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.ResetWithSameType() + require.Equal(t, initial.Used, state.account.Snapshot().Used) + require.NoError(t, AppendBytes(vec, bytes.Repeat([]byte("r"), 128), false, mp)) + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapResetReuseAndFree(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + + require.NoError(t, vec.PreExtend(32*1024, mp)) + require.NoError(t, vec.PreExtendBitmap(32*1024, mp)) + vec.SetLength(32 * 1024) + vec.SetAllNulls(32 * 1024) + vec.GetGrouping().AddRange(0, 32*1024) + require.True(t, vec.nsp.GetBitmap().HasExternalStorage()) + require.True(t, vec.gsp.GetBitmap().HasExternalStorage()) + require.Equal(t, 32*1024, vec.GetNulls().Count()) + require.Equal(t, 32*1024, vec.GetGrouping().Count()) + + initial := state.account.Snapshot() + expected := cap(vec.data) + + 8*vec.nsp.GetBitmap().ExternalStorageCapacity() + + 8*vec.gsp.GetBitmap().ExternalStorageCapacity() + require.Equal(t, uint64(expected), initial.Used) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + + vec.ResetWithSameType() + require.True(t, vec.GetNulls().IsEmpty()) + require.True(t, vec.GetGrouping().IsEmpty()) + require.Equal(t, initial.Used, state.account.Snapshot().Used) + require.NoError(t, vec.PreExtend(32*1024, mp)) + vec.SetNull(32*1024 - 1) + require.True(t, vec.IsNull(32*1024-1)) + require.Equal(t, initial.Used, state.account.Snapshot().Used) + + vec.ResetWithSameType() + require.NoError(t, vec.PreExtend(64*1024, mp)) + require.NoError(t, vec.PreExtendBitmap(64*1024, mp)) + grown := state.account.Snapshot() + require.Greater(t, grown.Used, initial.Used) + require.Greater(t, grown.Peak, grown.Used) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapShrinkUsesNoScratch(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + for _, row := range []uint64{0, 2, 63, 64, 129} { + vec.SetNull(row) + } + for _, row := range []uint64{1, 65, 128} { + vec.GetGrouping().Add(row) + } + before := state.account.Snapshot() + + vec.Shrink([]int64{0, 2, 64, 65, 129}, false) + require.Equal(t, []int64{0, 2, 64, 65, 129}, MustFixedColWithTypeCheck[int64](vec)) + for _, row := range []uint64{0, 1, 2, 4} { + require.True(t, vec.IsNull(row)) + } + require.Equal(t, 4, vec.GetNulls().Count()) + require.True(t, vec.GetGrouping().Contains(3)) + require.Equal(t, 1, vec.GetGrouping().Count()) + after := state.account.Snapshot() + require.Equal(t, before.Used, after.Used) + require.Equal(t, before.Peak, after.Peak) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountDupPreservesSparseBitmapRowDomain(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(0), true, mp)) + for i := 1; i < 130; i++ { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + require.NoError(t, vec.ensureGroupingCapacity(1, mp)) + vec.GetGrouping().Add(0) + require.Equal(t, 1, vec.GetNulls().Count()) + require.Equal(t, 1, vec.GetGrouping().Count()) + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.GreaterOrEqual( + t, + dup.GetNulls().GetBitmap().ExternalStorageCapacity(), + 3, + ) + require.GreaterOrEqual( + t, + dup.GetGrouping().GetBitmap().ExternalStorageCapacity(), + 3, + ) + sels := make([]int64, 129) + for i := range sels { + sels[i] = int64(i + 1) + } + require.NotPanics(t, func() { + dup.Shrink(sels, false) + }) + require.Equal(t, 129, dup.Length()) + require.Zero(t, dup.GetNulls().Count()) + require.Zero(t, dup.GetGrouping().Count()) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountDupPreservesStaleBitmapExtent(t *testing.T) { + t.Run("flat null and grouping", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + vec.SetNull(129) + vec.GetGrouping().Add(128) + vec.SetLength(1) + + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Equal(t, 1, dup.Length()) + require.True(t, dup.GetNulls().Contains(129)) + require.True(t, dup.GetGrouping().Contains(128)) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + }) + + t.Run("constant grouping", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + vec.SetClass(CONSTANT) + require.NoError(t, vec.PreExtendGrouping(130, mp)) + vec.GetGrouping().Add(129) + vec.SetLength(1) + + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Equal(t, 1, dup.Length()) + require.True(t, dup.GetGrouping().Contains(129)) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + }) + + t.Run("empty stale extent", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, AppendFixed(vec, int64(1), false, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + vec.SetNull(129) + vec.GetGrouping().Add(129) + vec.UnsetNull(129) + vec.GetGrouping().Del(129) + require.Zero(t, vec.GetNulls().Count()) + require.Zero(t, vec.GetGrouping().Count()) + + dup, err := vec.Dup(mp) + require.NoError(t, err) + require.Equal(t, int64(130), dup.GetNulls().GetBitmap().Len()) + require.Equal(t, int64(130), dup.GetGrouping().GetBitmap().Len()) + + dup.Free(mp) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + }) +} + +func TestVectorAllocationAccountBitmapShuffleAccountsScratch(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 16) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + for _, row := range []uint64{1, 64, 129} { + vec.SetNull(row) + } + for _, row := range []uint64{2, 63, 128} { + vec.GetGrouping().Add(row) + } + before := state.account.Snapshot() + + require.NoError(t, vec.Shuffle([]int64{129, 1, 64, 1, 2, 128, 63}, mp)) + require.Equal(t, []int64{129, 1, 64, 1, 2, 128, 63}, MustFixedColWithTypeCheck[int64](vec)) + for _, row := range []uint64{0, 1, 2, 3} { + require.True(t, vec.IsNull(row)) + } + require.Equal(t, 4, vec.GetNulls().Count()) + for _, row := range []uint64{4, 5, 6} { + require.True(t, vec.GetGrouping().Contains(row)) + } + require.Equal(t, 3, vec.GetGrouping().Count()) + after := state.account.Snapshot() + require.Greater(t, after.Peak, before.Peak) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + + var goScratch []byte + require.NoError(t, vec.ShuffleWithBuf([]int64{6, 5, 4, 3, 2, 1, 0}, mp, &goScratch)) + require.Nil(t, goScratch) + require.Equal(t, []int64{63, 128, 2, 1, 64, 1, 129}, MustFixedColWithTypeCheck[int64](vec)) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapShuffleFailurePreservesVector(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 4) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(130, mp)) + require.NoError(t, vec.PreExtendBitmap(130, mp)) + for i := range 130 { + require.NoError(t, AppendFixed(vec, int64(i), false, mp)) + } + vec.SetNull(1) + vec.GetGrouping().Add(2) + before := state.account.Snapshot() + + err := vec.Shuffle([]int64{2, 1, 0}, mp) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Equal(t, before.Used, state.account.Snapshot().Used) + require.Equal(t, uint64(3), state.registry.LiveAllocationMetadata()) + require.Equal(t, 130, vec.Length()) + require.True(t, vec.IsNull(1)) + require.True(t, vec.GetGrouping().Contains(2)) + require.Equal(t, int64(0), MustFixedColWithTypeCheck[int64](vec)[0]) + require.Equal(t, int64(129), MustFixedColWithTypeCheck[int64](vec)[129]) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapGrowthFailurePreservesOwner(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1000, 8) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(64, mp)) + require.NoError(t, vec.PreExtendBitmap(64, mp)) + vec.SetLength(64) + vec.SetNull(7) + vec.GetGrouping().Add(9) + + used := state.account.Snapshot().Used + dataCapacity := cap(vec.data) + nullCapacity := vec.nsp.GetBitmap().ExternalStorageCapacity() + groupCapacity := vec.gsp.GetBitmap().ExternalStorageCapacity() + // The null replacement fits by itself, but admitting the grouping + // replacement would exceed the account. Neither replacement is published. + err := vec.PreExtendBitmap(2*1024, mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, used, state.account.Snapshot().Used) + require.Equal(t, dataCapacity, cap(vec.data)) + require.Equal(t, nullCapacity, vec.nsp.GetBitmap().ExternalStorageCapacity()) + require.Equal(t, groupCapacity, vec.gsp.GetBitmap().ExternalStorageCapacity()) + require.True(t, vec.IsNull(7)) + require.True(t, vec.GetGrouping().Contains(9)) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapRejectsUnadmittedRawGrowth(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + unaccounted := NewOffHeapVecWithType(types.T_int64.ToType()) + unaccounted.GetNulls().Add(0) + require.ErrorIs( + t, + unaccounted.SetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + unaccounted.Free(mp) + + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.Panics(t, func() { + vec.GetNulls().Add(0) + }) + require.Zero(t, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountBitmapCopyDecode(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 32) + mp := mpool.MustNewZero() + source := NewOffHeapVecWithType(types.T_int64.ToType()) + for i := 0; i < 128; i++ { + require.NoError(t, AppendFixed(source, int64(i), i%7 == 0, mp)) + } + source.GetGrouping().Add(2, 9, 64) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + + copied := newAccountedTestVector( + t, + types.T_int64.ToType(), + state.selection, + ) + require.NoError(t, copied.UnmarshalBinaryWithCopy(encoded, mp)) + require.Equal(t, source.Length(), copied.Length()) + require.True(t, copied.GetNulls().IsSame(source.GetNulls())) + require.True(t, copied.nsp.GetBitmap().HasExternalStorage()) + require.True(t, copied.gsp.GetBitmap().HasExternalStorage()) + copied.Free(mp) + + fromReader := newAccountedTestVector( + t, + types.T_int64.ToType(), + state.selection, + ) + require.NoError(t, fromReader.UnmarshalWithReader(bytes.NewReader(encoded), mp)) + require.Equal(t, source.Length(), fromReader.Length()) + require.True(t, fromReader.GetNulls().IsSame(source.GetNulls())) + require.True(t, fromReader.nsp.GetBitmap().HasExternalStorage()) + fromReader.Free(mp) + + duplicate, err := source.DupOffHeapWithAllocation(mp, state.selection) + require.NoError(t, err) + require.True(t, duplicate.GetNulls().IsSame(source.GetNulls())) + require.True(t, duplicate.GetGrouping().IsSame(source.GetGrouping())) + duplicate.Free(mp) + + window, err := source.CloneWindowWithAllocation( + 1, + 65, + mp, + state.selection, + ) + require.NoError(t, err) + require.Equal(t, 64, window.Length()) + require.True(t, window.IsNull(6)) + require.True(t, window.IsNull(13)) + window.Free(mp) + + rollup := NewOffHeapVecWithType(types.T_int64.ToType()) + rollup.SetLength(128) + rollup.GetGrouping().AddRange(0, 128) + rollup.ToConst() + rollupCopy, err := rollup.DupOffHeapWithAllocation(mp, state.selection) + require.NoError(t, err) + require.True(t, rollupCopy.IsConstNull()) + require.True(t, rollupCopy.GetGrouping().IsSame(rollup.GetGrouping())) + rollupCopy.Free(mp) + rollup.Free(mp) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountViewAndDeepCopy(t *testing.T) { + stateA := newTestVectorAllocationAccount(t, 1<<20, 32) + stateB := newTestVectorAllocationAccount(t, 1<<20, 32) + mp := mpool.MustNewZero() + + source := newAccountedTestVector(t, types.T_varchar.ToType(), stateA.selection) + for _, value := range [][]byte{ + []byte("first value that is not inline"), + []byte("second value that is not inline"), + []byte("third value that is not inline"), + } { + require.NoError(t, AppendBytes(source, value, false, mp)) + } + sourceUsed := stateA.account.Snapshot().Used + + view, err := source.Window(1, 3) + require.NoError(t, err) + require.Nil(t, view.AllocationAccountSelection()) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + view.Free(mp) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + + dup, err := source.DupOffHeap(mp) + require.NoError(t, err) + require.Same(t, stateA.selection, dup.AllocationAccountSelection()) + require.Greater(t, stateA.account.Snapshot().Used, sourceUsed) + dup.Free(mp) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + + crossOwner, err := source.DupOffHeapWithAllocation(mp, stateB.selection) + require.NoError(t, err) + require.Equal(t, sourceUsed, stateA.account.Snapshot().Used) + require.NotZero(t, stateB.account.Snapshot().Used) + crossOwner.Free(mp) + require.Zero(t, stateB.account.Snapshot().Used) + + window, err := source.CloneWindowWithAllocation(1, 3, mp, stateB.selection) + require.NoError(t, err) + require.Equal(t, source.GetBytesAt(1), window.GetBytesAt(0)) + require.Equal(t, source.GetBytesAt(2), window.GetBytesAt(1)) + window.Free(mp) + + compact, err := source.CloneToFlatCompactWithAllocation(mp, stateB.selection) + require.NoError(t, err) + require.Equal(t, source.Length(), compact.Length()) + compact.Free(mp) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, stateA) + finalizeTestVectorAllocationAccount(t, stateB) +} + +func TestWindowPreservesGroupingProvenance(t *testing.T) { + mp := mpool.MustNewZero() + for _, typ := range []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()} { + source := NewVec(typ) + if typ.IsVarlen() { + for _, value := range []string{"zero", "one", "two", "three"} { + require.NoError(t, AppendBytes(source, []byte(value), false, mp)) + } + } else { + require.NoError(t, AppendFixedList(source, []int64{0, 1, 2, 3}, nil, mp)) + } + source.GetGrouping().Add(1, 3) + + window, err := source.Window(1, 4) + require.NoError(t, err) + require.True(t, window.GetGrouping().Contains(0)) + require.False(t, window.GetGrouping().Contains(1)) + require.True(t, window.GetGrouping().Contains(2)) + window.Free(mp) + + clone, err := source.CloneWindow(1, 4, mp) + require.NoError(t, err) + require.True(t, clone.GetGrouping().Contains(0)) + require.False(t, clone.GetGrouping().Contains(1)) + require.True(t, clone.GetGrouping().Contains(2)) + clone.Free(mp) + source.Free(mp) + } + + rollup := NewRollupConst(types.T_int64.ToType(), 4, mp) + window, err := rollup.Window(1, 3) + require.NoError(t, err) + require.True(t, window.IsGrouping()) + window.Free(mp) + clone, err := rollup.CloneWindow(1, 3, mp) + require.NoError(t, err) + require.True(t, clone.IsGrouping()) + clone.Free(mp) + rollup.Free(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestAccountedWindowOwnsRangeBitmaps(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + source := NewVec(types.T_int64.ToType()) + require.NoError(t, AppendFixedList(source, []int64{0, 1, 2, 3}, []bool{false, true, false, true}, mp)) + source.GetGrouping().Add(1, 2) + + window, err := source.WindowWithAllocation(1, 4, mp, state.selection) + require.NoError(t, err) + require.True(t, window.GetNulls().GetBitmap().HasExternalStorage()) + require.True(t, window.GetGrouping().GetBitmap().HasExternalStorage()) + require.True(t, window.GetGrouping().Contains(0)) + require.True(t, window.GetGrouping().Contains(1)) + require.NotZero(t, state.account.Snapshot().Used) + window.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestIsGroupingRejectsOutOfRangeBits(t *testing.T) { + for _, typ := range []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()} { + vec := NewVec(typ) + vec.SetLength(1) + vec.GetGrouping().Add(5) + require.False(t, vec.IsGrouping()) + vec.Free(nil) + } +} + +func TestConstSetPreservesSelectedGrouping(t *testing.T) { + mp := mpool.MustNewZero() + for _, typ := range []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()} { + source := NewVec(typ) + if typ.IsVarlen() { + require.NoError(t, AppendBytesList(source, [][]byte{[]byte("ordinary"), []byte("grouping")}, nil, mp)) + } else { + require.NoError(t, AppendFixedList(source, []int64{1, 2}, nil, mp)) + } + source.GetGrouping().Add(1) + destination := NewVec(typ) + set := GetConstSetFunction(typ, mp) + + require.NoError(t, set(destination, source, 1, 4)) + require.True(t, destination.IsGrouping()) + require.NoError(t, set(destination, source, 0, 4)) + require.False(t, destination.HasGrouping()) + + destination.Free(mp) + source.Free(mp) + } + require.Zero(t, mp.CurrNB()) +} + +func TestVectorAllocationAccountCopyRollback(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 1) + mp := mpool.MustNewZero() + source := NewVec(types.T_varchar.ToType()) + require.NoError( + t, + AppendBytes( + source, + bytes.Repeat([]byte("payload"), 32), + false, + mp, + ), + ) + + _, err := source.DupOffHeapWithAllocation(mp, state.selection) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountCrossPoolFreeAndSeal(t *testing.T) { + t.Run("cross pool free", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + ownerPool := mpool.MustNewZero() + freeingPool := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NoError(t, vec.PreExtend(128, ownerPool)) + require.NotZero(t, state.account.Snapshot().Used) + + vec.Free(freeingPool) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestVectorAllocationAccount(t, state) + }) + + t.Run("sealed before allocation", func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + state.account.Seal() + + err := AppendFixed(vec, int64(1), false, mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountSealed) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + vec.Free(mp) + _, err = state.registry.Finalize(state.account) + require.NoError(t, err) + }) +} + +func TestVectorAllocationAccountRandomizedAppendAndSelection(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 64) + mp := mpool.MustNewZero() + rng := rand.New(rand.NewSource(26459)) + + fixed := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + var fixedExpected []int64 + for i := 0; i < 2_000; i++ { + value := rng.Int63() + fixedExpected = append(fixedExpected, value) + require.NoError(t, AppendFixed(fixed, value, false, mp)) + } + require.Equal(t, fixedExpected, MustFixedColNoTypeCheck[int64](fixed)) + + varlen := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + var expected [][]byte + for i := 0; i < 1_000; i++ { + size := rng.Intn(96) + value := make([]byte, size) + _, err := rng.Read(value) + require.NoError(t, err) + expected = append(expected, append([]byte(nil), value...)) + require.NoError(t, AppendBytes(varlen, value, false, mp)) + } + for i := range expected { + require.True(t, bytes.Equal(expected[i], varlen.GetBytesAt(i))) + } + + selected := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + sels := []int64{1, 3, 7, 11, 23, 101, 509, 999} + require.NoError(t, selected.Union(varlen, sels, mp)) + for i, sel := range sels { + require.True(t, bytes.Equal(expected[sel], selected.GetBytesAt(i))) + } + + fixed.Free(mp) + varlen.Free(mp) + selected.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAccountedUnionPreservesGroupingWithoutNulls(t *testing.T) { + state := newTestVectorAllocationAccount(t, 8<<20, 64) + mp := mpool.MustNewZero() + + for _, typ := range []types.Type{ + types.T_int32.ToType(), + types.T_varchar.ToType(), + } { + t.Run(typ.String(), func(t *testing.T) { + source := NewOffHeapVecWithType(typ) + for i := range 6 { + if typ.IsVarlen() { + require.NoError(t, AppendBytes(source, []byte{byte('a' + i)}, false, mp)) + } else { + require.NoError(t, AppendFixed(source, int32(i), false, mp)) + } + } + source.GetGrouping().Add(1, 4) + + tests := []struct { + name string + run func(*Vector) error + want []bool + }{ + { + name: "union", + run: func(dst *Vector) error { + return dst.Union(source, []int64{4, 0, 1}, mp) + }, + want: []bool{true, false, true}, + }, + { + name: "union int32", + run: func(dst *Vector) error { + return dst.UnionInt32(source, []int32{4, 0, 1}, mp) + }, + want: []bool{true, false, true}, + }, + { + name: "union batch", + run: func(dst *Vector) error { + return dst.UnionBatch(source, 1, 4, nil, mp) + }, + want: []bool{true, false, false, true}, + }, + { + name: "union batch flags", + run: func(dst *Vector) error { + return dst.UnionBatch(source, 1, 4, []uint8{1, 0, 1, 1}, mp) + }, + want: []bool{true, false, true}, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dst := newAccountedTestVector(t, typ, state.selection) + require.NoError(t, test.run(dst)) + for row, want := range test.want { + require.Equal(t, want, dst.GetGrouping().Contains(uint64(row))) + } + require.True(t, dst.GetGrouping().GetBitmap().HasExternalStorage()) + dst.Free(mp) + }) + } + source.Free(mp) + }) + } + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAllocationAccountDecodeCopyAndReader(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + source := NewVec(types.T_varchar.ToType()) + require.NoError( + t, + AppendBytes(source, []byte("decoded payload that is not inline"), false, mp), + ) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + + copied := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + require.NoError(t, copied.UnmarshalBinaryWithCopy(encoded, mp)) + require.Equal(t, source.GetBytesAt(0), copied.GetBytesAt(0)) + require.Equal(t, uint64(2), state.registry.LiveAllocationMetadata()) + require.ErrorIs( + t, + copied.UnmarshalBinaryWithCopy(encoded, mp), + mpool.ErrAllocationAccountInvalid, + ) + require.ErrorIs( + t, + copied.UnmarshalBinary(encoded), + mpool.ErrAllocationAccountInvalid, + ) + copied.Free(mp) + + aliased := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + require.ErrorIs( + t, + aliased.UnmarshalBinary(encoded), + mpool.ErrAllocationAccountInvalid, + ) + aliased.Free(mp) + + readerDecoded := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + require.NoError(t, readerDecoded.UnmarshalWithReader(bytes.NewReader(encoded), mp)) + require.Equal(t, source.GetBytesAt(0), readerDecoded.GetBytesAt(0)) + readerDecoded.Free(mp) + + short := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + dataHeader := 1 + types.TSize + 4 + 4 + require.Error( + t, + short.UnmarshalWithReader( + bytes.NewReader(encoded[:dataHeader+1]), + mp, + ), + ) + require.Zero(t, state.account.Snapshot().Used) + short.Free(mp) + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestVectorAccountedReaderRejectsMalformedWire(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 16) + mp := mpool.MustNewZero() + source := NewVec(types.T_int64.ToType()) + require.NoError(t, AppendFixed(source, int64(7), true, mp)) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + + lengthOffset := 1 + types.TSize + dataLengthOffset := lengthOffset + 4 + dataLength := int(types.DecodeUint32(encoded[dataLengthOffset : dataLengthOffset+4])) + areaLengthOffset := dataLengthOffset + 4 + dataLength + areaLength := int(types.DecodeUint32(encoded[areaLengthOffset : areaLengthOffset+4])) + nullLengthOffset := areaLengthOffset + 4 + areaLength + nullOffset := nullLengthOffset + 4 + + tests := []struct { + name string + mutate func([]byte) + }{ + { + name: "invalid class", + mutate: func(data []byte) { + data[0] = 0xff + }, + }, + { + name: "negative length", + mutate: func(data []byte) { + value := uint32(math.MaxUint32) + copy(data[lengthOffset:lengthOffset+4], types.EncodeUint32(&value)) + }, + }, + { + name: "mismatched data length", + mutate: func(data []byte) { + value := uint32(2) + copy(data[lengthOffset:lengthOffset+4], types.EncodeUint32(&value)) + }, + }, + { + name: "oversized data payload", + mutate: func(data []byte) { + value := uint32(1 << 30) + copy(data[dataLengthOffset:dataLengthOffset+4], types.EncodeUint32(&value)) + }, + }, + { + name: "invalid null bitmap count", + mutate: func(data []byte) { + value := int64(2) + copy(data[nullOffset:nullOffset+8], types.EncodeInt64(&value)) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + data := append([]byte(nil), encoded...) + test.mutate(data) + decoded := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NotPanics(t, func() { + require.Error(t, decoded.UnmarshalWithReader(bytes.NewReader(data), mp)) + }) + decoded.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + copied := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NotPanics(t, func() { + require.Error(t, copied.UnmarshalBinaryWithCopy(data, mp)) + }) + copied.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + }) + } + for end := range encoded { + accounted := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + require.NotPanics(t, func() { + require.Error(t, accounted.UnmarshalBinaryWithCopy(encoded[:end], mp)) + }) + accounted.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + + unaccounted := NewOffHeapVecWithType(types.T_int64.ToType()) + require.NotPanics(t, func() { + require.Error(t, unaccounted.UnmarshalBinaryWithCopy(encoded[:end], mp)) + }) + unaccounted.Free(mp) + } + + source.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func BenchmarkVectorAllocationAccount(b *testing.B) { + const rows = 8192 + mp := mpool.MustNewZero() + state := newTestVectorAllocationAccount(b, 1<<40, 64) + + b.Run("unaccounted-fixed-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_int64.ToType()) + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("accounted-fixed-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_int64.ToType()) + if err := vec.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("unaccounted-varlen-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_varchar.ToType()) + if err := vec.PreExtendWithArea(rows, 1<<20, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("accounted-varlen-preextend-free", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + vec := NewOffHeapVecWithType(types.T_varchar.ToType()) + if err := vec.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + if err := vec.PreExtendWithArea(rows, 1<<20, mp); err != nil { + b.Fatal(err) + } + vec.Free(mp) + } + }) + b.Run("accounted-fixed-reset-reuse", func(b *testing.B) { + vec := newAccountedTestVector(b, types.T_int64.ToType(), state.selection) + if err := vec.PreExtend(rows, mp); err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + vec.ResetWithSameType() + } + b.StopTimer() + vec.Free(mp) + }) + finalizeTestVectorAllocationAccount(b, state) +} + +func BenchmarkVectorElementAccounting(b *testing.B) { + const rows = 8192 + mp := mpool.MustNewZero() + state := newTestVectorAllocationAccount(b, 1<<40, 64) + source := NewOffHeapVecWithType(types.T_int64.ToType()) + for i := range rows { + if err := AppendFixed(source, int64(i), false, mp); err != nil { + b.Fatal(err) + } + } + b.Cleanup(func() { + source.Free(mp) + finalizeTestVectorAllocationAccount(b, state) + }) + + for _, accounted := range []bool{false, true} { + mode := "unaccounted" + if accounted { + mode = "accounted" + } + b.Run("union-one/"+mode, func(b *testing.B) { + destination := NewOffHeapVecWithType(types.T_int64.ToType()) + if accounted { + if err := destination.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + } + if err := destination.PreExtend(1, mp); err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + destination.ResetWithSameType() + if err := destination.UnionOne(source, int64(i%rows), mp); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + destination.Free(mp) + }) + + b.Run("copy/"+mode, func(b *testing.B) { + destination := NewOffHeapVecWithType(types.T_int64.ToType()) + if accounted { + if err := destination.SetAllocationAccount(state.selection); err != nil { + b.Fatal(err) + } + } + if err := destination.PreExtend(1, mp); err != nil { + b.Fatal(err) + } + destination.SetLength(1) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := destination.Copy(source, 0, int64(i%rows), mp); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + destination.Free(mp) + }) + } +} + +func TestVectorAllocationAccountErrorsAreTyped(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1, 1) + mp := mpool.MustNewZero() + vec := newAccountedTestVector(t, types.T_int64.ToType(), state.selection) + err := vec.PreExtend(1, mp) + require.True(t, errors.Is(err, mpool.ErrAllocationAccountCapacity)) + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestDetachedBufferPreservesAllocationProvenance(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + source := newAccountedTestVector( + t, + types.T_varchar.ToType(), + state.selection, + ) + require.NoError(t, AppendBytes( + source, + []byte("detached allocation payload that uses the vector area"), + false, + mp, + )) + used := state.account.Snapshot().Used + require.Positive(t, used) + + data := DetachVectorData(source) + area := DetachVectorArea(source) + require.Positive(t, data.Capacity()) + require.Positive(t, area.Capacity()) + source.Free(mp) + require.Equal(t, used, state.account.Snapshot().Used) + + equivalent, err := NewAllocationAccountSelection( + state.account, + testVectorAllocationOwner, + testVectorDataAllocationSite, + testVectorAreaAllocationSite, + testVectorNullAllocationSite, + testVectorGroupAllocationSite, + ) + require.NoError(t, err) + require.NotSame(t, state.selection, equivalent) + destination := newAccountedTestVector( + t, + types.T_varchar.ToType(), + equivalent, + ) + require.True(t, data.CanAttachTo(destination, DetachedDataBuffer)) + require.False(t, data.CanAttachTo(destination, DetachedAreaBuffer)) + require.NoError(t, data.AttachTo(destination, DetachedDataBuffer)) + require.NoError(t, area.AttachTo(destination, DetachedAreaBuffer)) + require.Zero(t, data.Capacity()) + require.Zero(t, area.Capacity()) + require.Equal(t, used, state.account.Snapshot().Used) + + destination.Free(mp) + require.Zero(t, state.account.Snapshot().Used) + data.Free(mp) + area.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestSetTypeAndFixDataAllocationFailureIsAtomic(t *testing.T) { + state := newTestVectorAllocationAccount(t, 512, 4) + mp := mpool.MustNewZero() + vec := newAccountedTestVector( + t, + types.T_date.ToType(), + state.selection, + ) + require.NoError(t, vec.PreExtend(128, mp)) + vec.SetLength(128) + used := state.account.Snapshot().Used + require.Equal(t, uint64(512), used) + + err := vec.SetTypeAndFixData(types.T_datetime.ToType(), mp) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, types.T_date, vec.GetType().Oid) + require.Equal(t, 128, vec.Length()) + require.Equal(t, used, state.account.Snapshot().Used) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestDetachedUnaccountedBufferAndTypeChange(t *testing.T) { + mp := mpool.MustNewZero() + source := NewOffHeapVecWithType(types.T_varchar.ToType()) + require.NoError(t, AppendBytes( + source, + []byte("unaccounted detached allocation payload"), + false, + mp, + )) + data := DetachVectorData(source) + area := DetachVectorArea(source) + source.Free(mp) + + destination := NewOffHeapVecWithType(types.T_varchar.ToType()) + // Unaccounted buffers may serve either backing because there is no + // allocation-site provenance to preserve. + require.True(t, data.CanAttachTo(destination, DetachedAreaBuffer)) + require.NoError(t, area.AttachTo(destination, DetachedDataBuffer)) + require.NoError(t, data.AttachTo(destination, DetachedAreaBuffer)) + require.Zero(t, data.Capacity()) + require.Zero(t, area.Capacity()) + destination.Free(mp) + require.Zero(t, mp.CurrNB()) + + fixed := NewOffHeapVecWithType(types.T_date.ToType()) + require.NoError(t, AppendFixed( + fixed, + types.Date(1), + false, + mp, + )) + require.NoError(t, fixed.SetTypeAndFixData( + types.T_datetime.ToType(), + mp, + )) + require.Equal(t, types.T_datetime, fixed.GetType().Oid) + require.Equal(t, 1, fixed.Length()) + require.Error(t, fixed.SetTypeAndFixData( + types.T_varchar.ToType(), + mp, + )) + fixed.Free(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestVectorAllocationAccountHelperBoundaries(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 8) + mp := mpool.MustNewZero() + var nilVec *Vector + require.Nil(t, nilVec.AllocationAccountSelection()) + require.ErrorIs( + t, + nilVec.CanSetAllocationAccount(state.selection), + mpool.ErrAllocationAccountInvalid, + ) + + vec := newAccountedTestVector(t, types.T_varchar.ToType(), state.selection) + _, err := vec.allocOwned(mp, 1, false, true) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + _, err = vec.growData(nil, 1) + require.Error(t, err) + _, err = vec.growData(mp, mpool.CapLimit) + require.Error(t, err) + _, err = vec.growArea2(mp, []byte{1}, 0) + require.Error(t, err) + + require.NoError(t, vec.PreExtendWithArea(1, 128, mp)) + used := state.account.Snapshot().Used + zero := int32(0) + size, area, err := vec.readSizeBytes( + bytes.NewReader(types.EncodeInt32(&zero)), + mp, + false, + ) + require.NoError(t, err) + require.Zero(t, size) + require.Empty(t, area) + require.Equal(t, used, state.account.Snapshot().Used) + + negative := int32(-1) + _, _, err = vec.readSizeBytes( + bytes.NewReader(types.EncodeInt32(&negative)), + mp, + false, + ) + require.Error(t, err) + + vec.Free(mp) + finalizeTestVectorAllocationAccount(t, state) +} + +func TestUnionAllPreservesConstGrouping(t *testing.T) { + for _, typ := range []types.Type{ + types.T_int64.ToType(), + types.T_varchar.ToType(), + } { + t.Run(typ.String(), func(t *testing.T) { + state := newTestVectorAllocationAccount(t, 1<<20, 32) + mp := mpool.MustNewZero() + destination := newAccountedTestVector(t, typ, state.selection) + if typ.IsVarlen() { + require.NoError(t, AppendBytes(destination, []byte("prefix"), false, mp)) + } else { + require.NoError(t, AppendFixed(destination, int64(1), false, mp)) + } + + rollup := NewRollupConst(typ, 3, mp) + require.NoError(t, GetUnionAllFunction(typ, mp)(destination, rollup)) + require.Equal(t, 4, destination.Length()) + for row := uint64(1); row < 4; row++ { + require.True(t, destination.GetGrouping().Contains(row)) + } + rollup.Free(mp) + + var ordinary *Vector + var err error + if typ.IsVarlen() { + ordinary, err = NewConstBytes(typ, []byte("value"), 3, mp) + } else { + ordinary, err = NewConstFixed(typ, int64(2), 3, mp) + } + require.NoError(t, err) + ordinary.GetGrouping().Add(1) + require.NoError(t, GetUnionAllFunction(typ, mp)(destination, ordinary)) + require.Equal(t, 7, destination.Length()) + require.True(t, destination.GetGrouping().Contains(5)) + require.False(t, destination.GetGrouping().Contains(4)) + require.False(t, destination.GetGrouping().Contains(6)) + ordinary.Free(mp) + + require.True(t, + destination.GetGrouping().GetBitmap().HasExternalStorage()) + destination.Free(mp) + finalizeTestVectorAllocationAccount(t, state) + require.Zero(t, mp.CurrNB()) + }) + } +} + +func TestUnmarshalBinaryRejectsOwnedDestinationWithoutLosingBacking(t *testing.T) { + mp := mpool.MustNewZero() + source := NewOffHeapVecWithType(types.T_varchar.ToType()) + target := NewOffHeapVecWithType(types.T_varchar.ToType()) + require.NoError(t, AppendBytes(source, bytes.Repeat([]byte("s"), 64), false, mp)) + targetValue := bytes.Repeat([]byte("t"), 64) + require.NoError(t, AppendBytes(target, targetValue, false, mp)) + encoded, err := source.MarshalBinary() + require.NoError(t, err) + before := mp.CurrNB() + require.Positive(t, before) + + err = target.UnmarshalBinary(encoded) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) + require.Equal(t, targetValue, target.GetBytesAt(0)) + require.Equal(t, before, mp.CurrNB()) + + source.Free(mp) + target.Free(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestUnmarshalBinaryReplacesBorrowedAliases(t *testing.T) { + mp := mpool.MustNewZero() + first := NewVec(types.T_int64.ToType()) + second := NewVec(types.T_int64.ToType()) + require.NoError(t, AppendFixedList(first, []int64{1, 2}, nil, mp)) + require.NoError(t, AppendFixedList(second, []int64{3, 4}, []bool{true, false}, mp)) + firstData, err := first.MarshalBinary() + require.NoError(t, err) + secondData, err := second.MarshalBinary() + require.NoError(t, err) + + var target Vector + require.NoError(t, target.UnmarshalBinary(firstData)) + require.Equal(t, []int64{1, 2}, MustFixedColWithTypeCheck[int64](&target)) + require.NoError(t, target.UnmarshalBinaryTrusted(secondData)) + require.Equal(t, []int64{0, 4}, MustFixedColWithTypeCheck[int64](&target)) + require.True(t, target.GetNulls().Contains(0)) + + first.Free(mp) + second.Free(mp) + target.Free(mp) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/container/vector/functionTool_test.go b/pkg/container/vector/functionTool_test.go index 6b3ba1f7e5cf4..73c66c9703134 100644 --- a/pkg/container/vector/functionTool_test.go +++ b/pkg/container/vector/functionTool_test.go @@ -26,6 +26,34 @@ import ( "github.com/stretchr/testify/require" ) +func TestFunctionResultAllocationSurvivesVectorTransfer(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + wrapper, err := NewFunctionResultWrapperWithAllocation( + types.T_int64.ToType(), mp, selection, + ) + require.NoError(t, err) + require.NoError(t, wrapper.PreExtendAndReset(64)) + transferred := wrapper.GetResultVector() + require.Same(t, selection, transferred.AllocationAccountSelection()) + firstUsed := account.Snapshot().Used + require.Positive(t, firstUsed) + + wrapper.SetResultVector(nil) + require.NoError(t, wrapper.PreExtendAndReset(64)) + require.Same(t, selection, wrapper.GetResultVector().AllocationAccountSelection()) + require.Greater(t, account.Snapshot().Used, firstUsed) + + transferred.Free(mp) + wrapper.Free() + require.Zero(t, account.Snapshot().Used) +} + func TestAppendByteJsonUsesStorageCompatibleTypeCodes(t *testing.T) { mp := mpool.MustNewZeroNoFixed() wrapper := NewFunctionResultWrapper(types.T_json.ToType(), mp) diff --git a/pkg/container/vector/functionTools.go b/pkg/container/vector/functionTools.go index 9605a068a0ad4..438613bfd6089 100644 --- a/pkg/container/vector/functionTools.go +++ b/pkg/container/vector/functionTools.go @@ -571,9 +571,10 @@ func OptGetBytesParamFromWrapper(wrapper FunctionResultWrapper, idx int, src *Ve var _ FunctionResultWrapper = &FunctionResult[int64]{} type FunctionResult[T types.FixedSizeT] struct { - typ types.Type - vec *Vector - mp *mpool.MPool + typ types.Type + vec *Vector + mp *mpool.MPool + allocation *AllocationAccountSelection isVarlena bool cols []T @@ -621,7 +622,11 @@ func (fr *FunctionResult[T]) getConvenientParamList() []reusableParameterWrapper func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { if fr.vec == nil { - fr.vec = NewOffHeapVecWithType(fr.typ) + var err error + fr.vec, err = NewOffHeapVecWithTypeAndAllocation(fr.typ, fr.allocation) + if err != nil { + return err + } } oldLength := fr.vec.Length() @@ -631,6 +636,9 @@ func (fr *FunctionResult[T]) PreExtendAndReset(targetSize int) error { return err } } + if err := fr.vec.PreExtendNulls(targetSize, fr.mp); err != nil { + return err + } fr.vec.ResetWithSameType() if !fr.isVarlena { @@ -760,6 +768,14 @@ func (fr *FunctionResult[T]) Free() { fr.convenientParam = nil } +func (fr *FunctionResult[T]) setAllocation(selection *AllocationAccountSelection) { + fr.allocation = selection +} + +type functionResultAllocationSetter interface { + setAllocation(*AllocationAccountSelection) +} + func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWrapper { if typ.IsVarlen() { return newResultFunc[types.Varlena](typ, mp) @@ -819,3 +835,25 @@ func NewFunctionResultWrapper(typ types.Type, mp *mpool.MPool) FunctionResultWra } panic(fmt.Sprintf("unexpected type %s for function result", typ)) } + +// NewFunctionResultWrapperWithAllocation constructs a reusable result whose +// current and future vectors allocate through selection. The selection stays +// with the wrapper when EvalWithoutResultReusing transfers its current vector. +func NewFunctionResultWrapperWithAllocation( + typ types.Type, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (FunctionResultWrapper, error) { + if selection != nil { + if err := selection.validate(); err != nil { + return nil, err + } + } + result := NewFunctionResultWrapper(typ, mp) + setter, ok := result.(functionResultAllocationSetter) + if !ok { + return nil, mpool.ErrAllocationAccountInvariant + } + setter.setAllocation(selection) + return result, nil +} diff --git a/pkg/container/vector/pSpoolTools.go b/pkg/container/vector/pSpoolTools.go index 72bbafb60d6f3..27dd7f90c5991 100644 --- a/pkg/container/vector/pSpoolTools.go +++ b/pkg/container/vector/pSpoolTools.go @@ -14,35 +14,119 @@ package vector -// SetVecData is dangerous and should be used with caution. -func SetVecData(v *Vector, data []byte) { +import "github.com/matrixorigin/matrixone/pkg/common/mpool" + +// DetachedBuffer transfers one owned Vector backing allocation through the +// pipeline spool without losing its immutable allocation provenance. +// A non-empty value must be attached or freed exactly once. +type DetachedBuffer struct { + data []byte + selection *AllocationAccountSelection + kind DetachedBufferKind +} + +type DetachedBufferKind uint8 + +const ( + DetachedDataBuffer DetachedBufferKind = iota + DetachedAreaBuffer +) + +func DetachVectorData(v *Vector) DetachedBuffer { + if v == nil { + return DetachedBuffer{} + } + buffer := DetachedBuffer{ + data: v.data, + selection: v.allocationAccount, + } if v.typ.IsVarlen() { v.areaDisjoint = false } - data = data[:cap(data)] - v.data = data + v.data = nil + return buffer } -// SetVecArea is dangerous and should be used with caution. -func SetVecArea(v *Vector, area []byte) { +func DetachVectorArea(v *Vector) DetachedBuffer { + if v == nil { + return DetachedBuffer{} + } + buffer := DetachedBuffer{ + data: v.area, + selection: v.allocationAccount, + kind: DetachedAreaBuffer, + } v.areaDisjoint = false - v.area = area + v.area = nil + return buffer } -// GetAndClearVecData is a dangerous function that may cause data leakage. -func GetAndClearVecData(v *Vector) []byte { - if v.typ.IsVarlen() { +func (b *DetachedBuffer) Capacity() int { + if b == nil { + return 0 + } + return cap(b.data) +} + +// CanAttachTo preserves data/area site provenance when an allocation has an +// account. Unaccounted storage has no site identity and can serve either +// backing without losing ownership information. +func (b *DetachedBuffer) CanAttachTo( + v *Vector, + kind DetachedBufferKind, +) bool { + if b == nil || v == nil || cap(b.data) == 0 || + !AllocationAccountSelectionsEqual(b.selection, v.allocationAccount) || + kind > DetachedAreaBuffer { + return false + } + return b.selection == nil || b.kind == kind +} + +func (b *DetachedBuffer) AttachTo( + v *Vector, + kind DetachedBufferKind, +) error { + if !b.CanAttachTo(v, kind) { + return allocationAccountInvalid( + "detached vector buffer provenance mismatch", + ) + } + if kind == DetachedAreaBuffer { + if cap(v.area) != 0 { + return allocationAccountInvalid( + "vector area already has backing storage", + ) + } + v.area = b.data v.areaDisjoint = false + } else { + if cap(v.data) != 0 { + return allocationAccountInvalid( + "vector data already has backing storage", + ) + } + v.data = b.data[:cap(b.data)] + if v.typ.IsVarlen() { + v.areaDisjoint = false + } } - s := v.data - v.data = nil - return s + b.clear() + return nil } -// GetAndClearVecArea is a dangerous function that may cause data leakage. -func GetAndClearVecArea(v *Vector) []byte { - v.areaDisjoint = false - s := v.area - v.area = nil - return s +func (b *DetachedBuffer) Free(mp *mpool.MPool) { + if b == nil { + return + } + if cap(b.data) != 0 { + mp.Free(b.data) + } + b.clear() +} + +func (b *DetachedBuffer) clear() { + b.data = nil + b.selection = nil + b.kind = DetachedDataBuffer } diff --git a/pkg/container/vector/tools.go b/pkg/container/vector/tools.go index 8920131cf729e..6c0b0c1c6cba3 100644 --- a/pkg/container/vector/tools.go +++ b/pkg/container/vector/tools.go @@ -194,13 +194,23 @@ func MustVarlenaToInt64Slice(v *Vector) [][3]int64 { } func MustVarlenaRawData(v *Vector) (data []types.Varlena, area []byte) { - data = MustFixedColNoTypeCheck[types.Varlena](v) + data = ToSliceNoTypeCheck2[types.Varlena](v) area = v.area return } // XXX extend will extend the vector's Data to accommodate rows more entry. func extend(v *Vector, rows int, m *mpool.MPool) error { + return extendWithBitmaps(v, rows, m, false, false) +} + +func extendWithBitmaps( + v *Vector, + rows int, + m *mpool.MPool, + needNulls bool, + needGrouping bool, +) error { if rows <= 0 { // we will at least extent by 1. // This is a pure hack to @@ -208,9 +218,23 @@ func extend(v *Vector, rows int, m *mpool.MPool) error { } tgtLen := v.length + rows + switch { + case needNulls && needGrouping: + if err := v.ensureBitmapCapacity(tgtLen, m); err != nil { + return err + } + case needNulls: + if err := v.ensureNullCapacity(tgtLen, m); err != nil { + return err + } + case needGrouping: + if err := v.ensureGroupingCapacity(tgtLen, m); err != nil { + return err + } + } tgtDataCap := tgtLen * v.typ.TypeSize() if tgtDataCap > cap(v.data) { - ndata, err := m.Grow(v.data, tgtDataCap, v.offHeap) + ndata, err := v.growData(m, tgtDataCap) if err != nil { return err } diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index e4b391bd46fea..c7dadb634f3db 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -16,8 +16,10 @@ package vector import ( "bytes" + "encoding/binary" "fmt" "io" + "math" "math/bits" "slices" "sort" @@ -68,11 +70,13 @@ type Vector struct { offHeap bool - // areaDisjoint is true only when every non-inline varlena descriptor in the - // vector's logical range references a valid byte range disjoint from every - // other descriptor in that range. The invariant deliberately includes null - // rows, so changing a null bitmap cannot make this proof stale. Operations - // that reuse descriptors clear the proof; ordinary value appends retain it. + // allocationAccount selects the account for this vector's first owned + // off-heap data and area allocations. Physical MPool metadata owns release. + allocationAccount *AllocationAccountSelection + + // areaDisjoint proves that every live non-inline varlena descriptor owns a + // distinct range in area. Spill projections use it to avoid scanning normal + // append-built vectors; operations that can introduce aliases clear it. areaDisjoint bool } @@ -194,7 +198,10 @@ func (v *Vector) Capacity() int { // Allocated returns the total allocated memory size of the vector. // it can be used to estimate the memory usage of the vector. func (v *Vector) Allocated() int { - return cap(v.data) + cap(v.area) + return cap(v.data) + + cap(v.area) + + 8*v.nsp.GetBitmap().ExternalStorageCapacity() + + 8*v.gsp.GetBitmap().ExternalStorageCapacity() } func (v *Vector) SetLength(n int) { @@ -204,6 +211,40 @@ func (v *Vector) SetLength(n int) { v.length = n } +// AppendCheckpoint captures the logical state changed by append operations. +// Capacity growth is deliberately not rolled back: it remains owned by the +// vector and can be reused by a later append. +type AppendCheckpoint struct { + length int + areaLength int + sorted bool +} + +func (v *Vector) MakeAppendCheckpoint() AppendCheckpoint { + return AppendCheckpoint{ + length: v.length, + areaLength: len(v.area), + sorted: v.sorted, + } +} + +// RollbackAppend restores the logical state captured before an attempted +// append. attemptedRows is needed because grouping bits can be published +// before a varlen copy fails and advances length. +func (v *Vector) RollbackAppend(checkpoint AppendCheckpoint, attemptedRows int) { + if checkpoint.length < 0 || checkpoint.length > v.length || + checkpoint.areaLength < 0 || checkpoint.areaLength > len(v.area) || + attemptedRows < 0 { + panic("invalid vector append checkpoint") + } + end := max(v.length, checkpoint.length+attemptedRows) + nulls.RemoveRange(&v.nsp, uint64(checkpoint.length), uint64(end)) + nulls.RemoveRange(&v.gsp, uint64(checkpoint.length), uint64(end)) + v.length = checkpoint.length + v.area = v.area[:checkpoint.areaLength] + v.sorted = checkpoint.sorted +} + // Size of data, I think this function is inherently broken. This // Size is not meaningful other than used in (approximate) memory accounting. func (v *Vector) Size() int { @@ -227,28 +268,41 @@ func (v *Vector) SetType(typ types.Type) { v.typ = typ } -// Bug #23240 -// Neither this function, nor the SetType function are good -// Maybe we should just disallow. -func (v *Vector) SetTypeAndFixData(typ types.Type, mp *mpool.MPool) { +// SetTypeAndFixData changes a fixed-width result type and grows its owned data +// before publishing the new type. A failed growth leaves the original vector +// type, length, and backing allocation intact. +func (v *Vector) SetTypeAndFixData( + typ types.Type, + mp *mpool.MPool, +) error { if v.typ.IsVarlen() && typ.IsVarlen() { v.typ = typ - return + return nil } if v.typ.IsVarlen() || typ.IsVarlen() { - // this is a weird thing to do, we should not allow it. - panic("SetTypeAndFixData is not allowed to change from/to varlen type") + return moerr.NewInvalidInputNoCtx( + "SetTypeAndFixData cannot change from or to a varlen type", + ) } + oldType := v.typ v.typ = typ oldLength := v.length v.length = 0 - extend(v, oldLength, mp) + if err := extend(v, oldLength, mp); err != nil { + v.typ = oldType + v.length = oldLength + return err + } v.length = oldLength + return nil } func (v *Vector) SetOffHeap(offHeap bool) { + if !offHeap && v.allocationAccount != nil { + panic("allocation-accounted vector must remain off-heap") + } v.offHeap = offHeap } @@ -545,7 +599,7 @@ func NewVecWithDataCopy( vec.areaDisjoint = !typ.IsVarlen() || length == 0 var err error if len(data) > 0 { - vec.data, err = mp.Alloc(len(data), false) + vec.data, err = vec.allocData(mp, len(data)) if err != nil { vec.Free(mp) return nil, err @@ -553,7 +607,7 @@ func NewVecWithDataCopy( copy(vec.data, data) } if len(area) > 0 { - vec.area, err = mp.Alloc(len(area), false) + vec.area, err = vec.allocArea(mp, len(area)) if err != nil { vec.Free(mp) return nil, err @@ -622,7 +676,9 @@ func (v *Vector) IsConst() bool { } func (v *Vector) IsGrouping() bool { - return v.length > 0 && v.length == v.gsp.Count() + return v.length > 0 && + v.length == v.gsp.Count() && + v.length == v.gsp.GetBitmap().CountRange(0, uint64(v.length)) } func (v *Vector) SetClass(class int) { @@ -769,6 +825,7 @@ func (v *Vector) Free(mp *mpool.MPool) { if !v.cantFreeArea { mp.Free(v.area) } + v.freeBitmapStorage(mp) v.class = FLAT v.data = nil v.area = nil @@ -780,6 +837,7 @@ func (v *Vector) Free(mp *mpool.MPool) { v.gsp.Reset() v.sorted = false v.isBin = false + v.allocationAccount = nil v.areaDisjoint = true // if !v.OnUsed || v.OnPut { @@ -805,50 +863,174 @@ func (v *Vector) MarshalBinary() ([]byte, error) { } func (v *Vector) MarshalBinaryWithBuffer(buf *bytes.Buffer) error { + return v.MarshalBinaryTo(buf) +} - // write class - buf.WriteByte(uint8(v.class)) - - // write type - data := types.EncodeType(&v.typ) - buf.Write(data) +// MarshalBinaryPlan is a validated, allocation-free snapshot of one Vector's +// wire lengths. It lets batch writers size once and encode once. +type MarshalBinaryPlan struct { + vector *Vector + size int + dataLength uint32 + areaLength uint32 + nullLength uint32 +} - // write length - length := uint32(v.length) - buf.Write(types.EncodeUint32(&length)) +func (p MarshalBinaryPlan) Size() int { + return p.size +} - // write dataLen, data - dataLen := uint32(v.typ.TypeSize()) +func (v *Vector) PrepareMarshalBinary() (MarshalBinaryPlan, error) { + if v == nil || v.length < 0 { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx("invalid vector for marshal") + } + const maxWireBuffer = uint64(^uint32(0)) + if uint64(v.length) > maxWireBuffer { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( + "vector length exceeds marshal format", + ) + } + typeSize := v.typ.TypeSize() + if typeSize < 0 { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( + "vector type has invalid marshal size", + ) + } + dataLength := uint64(typeSize) if !v.IsConst() { - dataLen *= uint32(v.length) + if v.length != 0 && + dataLength > ^uint64(0)/uint64(v.length) { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( + "vector data exceeds marshal format", + ) + } + dataLength *= uint64(v.length) } else if v.IsConstNull() { - dataLen = 0 + dataLength = 0 + } + areaLength := uint64(len(v.area)) + nullLength := uint64(v.nsp.MarshalSize()) + if dataLength > maxWireBuffer || + areaLength > maxWireBuffer || + nullLength > maxWireBuffer { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( + "vector buffer exceeds marshal format", + ) } - buf.Write(types.EncodeUint32(&dataLen)) - if dataLen > 0 { - buf.Write(v.data[:dataLen]) + if dataLength > uint64(len(v.data)) { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( + "vector data is shorter than its marshal length", + ) } - - // write areaLen, area - areaLen := uint32(len(v.area)) - buf.Write(types.EncodeUint32(&areaLen)) - if areaLen > 0 { - buf.Write(v.area) + total := uint64(1+types.TSize+4+4+4+4+1) + + dataLength + areaLength + nullLength + if total > uint64(^uint(0)>>1) { + return MarshalBinaryPlan{}, moerr.NewInvalidInputNoCtx( + "vector marshal size exceeds platform limit", + ) } + return MarshalBinaryPlan{ + vector: v, + size: int(total), + dataLength: uint32(dataLength), + areaLength: uint32(areaLength), + nullLength: uint32(nullLength), + }, nil +} - // write nspLen, nsp - nspData, err := v.nsp.Show() +func (v *Vector) MarshalBinarySize() (int, error) { + plan, err := v.PrepareMarshalBinary() + return plan.Size(), err +} + +func (v *Vector) MarshalBinaryTo(w io.Writer) error { + plan, err := v.PrepareMarshalBinary() if err != nil { return err } - nspLen := uint32(len(nspData)) - buf.Write(types.EncodeUint32(&nspLen)) - if nspLen > 0 { - buf.Write(nspData) + return plan.MarshalTo(w) +} + +func (p MarshalBinaryPlan) MarshalTo(w io.Writer) error { + v := p.vector + if v == nil || w == nil { + return io.ErrClosedPipe + } + if err := writeVectorMarshalByte(w, uint8(v.class)); err != nil { + return err + } + if err := writeVectorMarshalBytes(w, types.EncodeType(&v.typ)); err != nil { + return err } - buf.Write(types.EncodeBool(&v.sorted)) + if err := writeVectorMarshalUint32(w, uint32(v.length)); err != nil { + return err + } + + if err := writeVectorMarshalUint32(w, p.dataLength); err != nil { + return err + } + if p.dataLength > 0 { + if err := writeVectorMarshalBytes(w, v.data[:p.dataLength]); err != nil { + return err + } + } + if err := writeVectorMarshalUint32(w, p.areaLength); err != nil { + return err + } + if p.areaLength > 0 { + if err := writeVectorMarshalBytes(w, v.area); err != nil { + return err + } + } + + if err := writeVectorMarshalUint32(w, p.nullLength); err != nil { + return err + } + if p.nullLength > 0 { + if err := v.nsp.MarshalTo(w); err != nil { + return err + } + } + + if v.sorted { + return writeVectorMarshalByte(w, 1) + } + return writeVectorMarshalByte(w, 0) +} + +type vectorPrimitiveWriter interface { + WriteByte(byte) error + WriteUint32(uint32) error +} + +func writeVectorMarshalByte(w io.Writer, value byte) error { + if typed, ok := w.(vectorPrimitiveWriter); ok { + return typed.WriteByte(value) + } + var data [1]byte + data[0] = value + return writeVectorMarshalBytes(w, data[:]) +} + +func writeVectorMarshalUint32(w io.Writer, value uint32) error { + if typed, ok := w.(vectorPrimitiveWriter); ok { + return typed.WriteUint32(value) + } + var data [4]byte + binary.NativeEndian.PutUint32(data[:], value) + return writeVectorMarshalBytes(w, data[:]) +} + +func writeVectorMarshalBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { + return err + } + if written != len(value) { + return io.ErrShortWrite + } return nil } @@ -873,89 +1055,158 @@ func (v *Vector) UnmarshalBinaryTrusted(data []byte) error { return v.unmarshalBinary(data, false) } -func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { - v.areaDisjoint = false - read := func(size int) ([]byte, error) { - if size < 0 || size > len(data) { - return nil, io.ErrUnexpectedEOF - } - value := data[:size] - data = data[size:] - return value, nil - } - readUint32 := func() (uint32, error) { - value, err := read(4) - if err != nil { - return 0, err - } - return types.DecodeUint32(value), nil +type vectorBinaryLayout struct { + class byte + typ types.Type + length int + data []byte + area []byte + nulls []byte + sorted bool +} + +type vectorBinaryCursor struct { + data []byte + offset int +} + +func (c *vectorBinaryCursor) read(size int) ([]byte, error) { + if size < 0 || c.offset > len(c.data)-size { + return nil, io.ErrUnexpectedEOF } + value := c.data[c.offset : c.offset+size] + c.offset += size + return value, nil +} - class, err := read(1) +func (c *vectorBinaryCursor) readUint32() (uint32, error) { + value, err := c.read(4) if err != nil { - return err + return 0, err } - typ, err := read(types.TSize) + return types.DecodeUint32(value), nil +} + +func decodeVectorBinaryLayout( + data []byte, + validateValues bool, +) (vectorBinaryLayout, error) { + cursor := vectorBinaryCursor{data: data} + class, err := cursor.read(1) if err != nil { - return err + return vectorBinaryLayout{}, err } - length, err := readUint32() + typData, err := cursor.read(types.TSize) if err != nil { - return err + return vectorBinaryLayout{}, err } - dataLen, err := readUint32() - if err != nil { - return err + length, err := cursor.readUint32() + if err != nil || uint64(length) > uint64(math.MaxInt) { + if err != nil { + return vectorBinaryLayout{}, err + } + return vectorBinaryLayout{}, moerr.NewInvalidInputNoCtx("vector length exceeds platform limit") } - vecData, err := read(int(dataLen)) - if err != nil { - return err + readSized := func() ([]byte, error) { + size, err := cursor.readUint32() + if err != nil { + return nil, err + } + if uint64(size) > uint64(math.MaxInt) { + return nil, moerr.NewInvalidInputNoCtx("vector buffer exceeds platform limit") + } + return cursor.read(int(size)) } - areaLen, err := readUint32() + vectorData, err := readSized() if err != nil { - return err + return vectorBinaryLayout{}, err } - area, err := read(int(areaLen)) + area, err := readSized() if err != nil { - return err + return vectorBinaryLayout{}, err } - nspLen, err := readUint32() + nullData, err := readSized() if err != nil { - return err + return vectorBinaryLayout{}, err } - nspData, err := read(int(nspLen)) + sorted, err := cursor.read(1) if err != nil { - return err + return vectorBinaryLayout{}, err + } + if cursor.offset != len(cursor.data) { + return vectorBinaryLayout{}, moerr.NewInvalidInputNoCtx("trailing vector wire data") + } + if sorted[0] > 1 { + return vectorBinaryLayout{}, moerr.NewInvalidInputNoCtx("invalid vector sorted flag") + } + if err = validateVectorNullBitmap(nullData, validateValues); err != nil { + return vectorBinaryLayout{}, err + } + var decodedNulls nulls.Nulls + if len(nullData) > 0 { + if err = decodedNulls.ReadNoCopy(nullData); err != nil { + return vectorBinaryLayout{}, err + } + } + typ := types.DecodeType(typData) + if err = validateVectorBinary( + class[0], + typ, + length, + vectorData, + area, + &decodedNulls, + validateValues, + ); err != nil { + return vectorBinaryLayout{}, err + } + return vectorBinaryLayout{ + class: class[0], + typ: typ, + length: int(length), + data: vectorData, + area: area, + nulls: nullData, + sorted: sorted[0] != 0, + }, nil +} + +func (v *Vector) unmarshalBinary(data []byte, validateValues bool) error { + if v == nil { + return io.ErrClosedPipe } - sorted, err := read(1) + v.areaDisjoint = false + if v.allocationAccount != nil { + return allocationAccountInvalid( + "cannot install aliases in an accounted vector", + ) + } + layout, err := decodeVectorBinaryLayout(data, validateValues) if err != nil { return err } - - decodedType := types.DecodeType(typ) - if err := validateVectorNullBitmap(nspData, validateValues); err != nil { - return err + if v.hasOwnedBackingStorage() { + return allocationAccountInvalid( + "cannot replace owned vector storage with aliases", + ) } - var nsp nulls.Nulls - if len(nspData) > 0 { - if err := nsp.ReadNoCopy(nspData); err != nil { + var decodedNulls nulls.Nulls + if len(layout.nulls) > 0 { + if err = decodedNulls.ReadNoCopy(layout.nulls); err != nil { return err } } - if err := validateVectorBinary(class[0], decodedType, length, vecData, area, &nsp, validateValues); err != nil { - return err - } - v.class = int(class[0]) - v.typ = decodedType - v.length = int(length) - v.data = vecData - v.area = area - v.nsp = nsp - v.sorted = types.DecodeBool(sorted) - + v.class = int(layout.class) + v.typ = layout.typ + v.length = layout.length + v.data = layout.data + v.area = layout.area + v.nsp = decodedNulls + v.gsp.Reset() + v.sorted = layout.sorted v.cantFreeData = true v.cantFreeArea = true - + v.allocationAccount = nil return nil } @@ -1020,14 +1271,15 @@ func validateVectorNullBitmap(data []byte, validateValues bool) error { if len(data) == 0 { return nil } - if len(data) < 24 { + if len(data) < bitmap.MarshalHeaderSize { return io.ErrUnexpectedEOF } count := types.DecodeInt64(data[:8]) bitmapLen := types.DecodeUint64(data[8:16]) bitmapDataLen := types.DecodeUint64(data[16:24]) if count < 0 || bitmapLen > uint64(1<<63-1) || uint64(count) > bitmapLen || - bitmapDataLen%8 != 0 || bitmapDataLen != uint64(len(data)-24) { + bitmapDataLen%8 != 0 || + bitmapDataLen != uint64(len(data)-bitmap.MarshalHeaderSize) { return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") } if bitmapDataLen != ((bitmapLen+63)/64)*8 { @@ -1036,7 +1288,7 @@ func validateVectorNullBitmap(data []byte, validateValues bool) error { if !validateValues { return nil } - words := types.DecodeSlice[uint64](data[24:]) + words := types.DecodeSlice[uint64](data[bitmap.MarshalHeaderSize:]) actualCount := int64(0) for i, word := range words { if i == len(words)-1 && bitmapLen%64 != 0 && word>>uint(bitmapLen%64) != 0 { @@ -1078,65 +1330,74 @@ func canonicalVectorTypeSize(typ types.Type) (int, error) { } func (v *Vector) UnmarshalBinaryWithCopy(data []byte, mp *mpool.MPool) error { + if v == nil || mp == nil { + return io.ErrClosedPipe + } v.areaDisjoint = false - var err error - - // read class - v.class = int(data[0]) - data = data[1:] - - // read typ - v.typ = types.DecodeType(data[:types.TSize]) - data = data[types.TSize:] - - // read length - v.length = int(types.DecodeUint32(data[:4])) - data = data[4:] - - // read data - dataLen := int(types.DecodeUint32(data[:4])) - data = data[4:] - if dataLen > 0 { - v.data, err = mp.Alloc(dataLen, v.offHeap) + if v.hasBackingStorage() { + return allocationAccountInvalid( + "cannot replace vector storage without Free", + ) + } + layout, err := decodeVectorBinaryLayout(data, true) + if err != nil { + return err + } + decoded := NewVec(layout.typ) + decoded.offHeap = v.offHeap + if v.allocationAccount != nil { + if err = decoded.SetAllocationAccount(v.allocationAccount); err != nil { + return err + } + } + committed := false + defer func() { + if !committed { + decoded.Free(mp) + } + }() + decoded.class = int(layout.class) + decoded.length = layout.length + if len(layout.data) > 0 { + decoded.data, err = decoded.allocData(mp, len(layout.data)) if err != nil { return err } - copy(v.data, data[:dataLen]) - data = data[dataLen:] + copy(decoded.data, layout.data) } - - // read area - areaLen := int(types.DecodeUint32(data[:4])) - data = data[4:] - if areaLen > 0 { - v.area, err = mp.Alloc(areaLen, v.offHeap) + if len(layout.area) > 0 { + decoded.area, err = decoded.allocArea(mp, len(layout.area)) if err != nil { return err } - copy(v.area, data[:areaLen]) - data = data[areaLen:] + copy(decoded.area, layout.area) } - - // read nsp - nspLen := types.DecodeUint32(data[:4]) - data = data[4:] - if nspLen > 0 { - if err := v.nsp.Read(data[:nspLen]); err != nil { + if len(layout.nulls) > 0 { + if decoded.allocationAccount != nil { + _, bitLength, _, decodeErr := bitmap.DecodeMarshalHeader(layout.nulls) + if decodeErr != nil || bitLength > int64(math.MaxInt) { + return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") + } + if err = decoded.ensureNullCapacity(int(bitLength), mp); err != nil { + return err + } + } + if err = decoded.nsp.Read(layout.nulls); err != nil { return err } - data = data[nspLen:] - } else { - v.nsp.Reset() } - - v.sorted = types.DecodeBool(data[:1]) - //data = data[1:] - + decoded.sorted = layout.sorted + *v = *decoded + committed = true return nil } func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { + if v == nil || r == nil { + return io.ErrClosedPipe + } v.areaDisjoint = false + v.ResetWithSameType() var err error if v.class, err = types.ReadByteAsInt(r); err != nil { @@ -1150,9 +1411,15 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { if v.length, err = types.ReadInt32AsInt(r); err != nil { return err } + if v.length < 0 { + return moerr.NewInvalidInputNoCtx("negative vector length") + } + if v.length > math.MaxUint32 { + return moerr.NewInvalidInputNoCtx("vector length exceeds marshal format") + } // read data - dataLen, dataBuf, err := types.ReadSizeBytesMp(r, v.data, mp, v.offHeap) + dataLen, dataBuf, err := v.readSizeBytes(r, mp, true) if err != nil { return err } @@ -1161,7 +1428,7 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { } // read area - areaLen, areaBuf, err := types.ReadSizeBytesMp(r, v.area, mp, v.offHeap) + areaLen, areaBuf, err := v.readSizeBytes(r, mp, false) if err != nil { return err } @@ -1169,25 +1436,156 @@ func (v *Vector) UnmarshalWithReader(r io.Reader, mp *mpool.MPool) error { v.area = areaBuf } - // read nsp, do not use mpool. nspBuf is different because - // it is not managed by vector. In the following, it will - // be unmarshalled into v.nsp - nspLen, nspBuf, err := types.ReadSizeBytes(r) + if err = v.readNullsWithReader(r, mp); err != nil { + return err + } + + v.sorted, err = types.ReadBool(r) if err != nil { return err } - if nspLen > 0 { - v.nsp.Read(nspBuf) - } else { + return validateVectorBinary( + byte(v.class), + v.typ, + uint32(v.length), + v.data[:int(dataLen)], + v.area[:int(areaLen)], + &v.nsp, + true, + ) +} + +func (v *Vector) readNullsWithReader(r io.Reader, mp *mpool.MPool) error { + if v.allocationAccount == nil { + nspLen, err := types.ReadInt32(r) + if err != nil { + return err + } + if nspLen < 0 { + return moerr.NewInvalidInputNoCtx("negative vector null bitmap size") + } + if err = validateStreamingReadSize(r, int64(nspLen)); err != nil { + return err + } + if nspLen > 0 { + nspBuf := make([]byte, nspLen) + if _, err = io.ReadFull(r, nspBuf); err != nil { + return err + } + if err := validateVectorNullBitmap(nspBuf, true); err != nil { + return err + } + return v.nsp.Read(nspBuf) + } v.nsp.Reset() + return nil } - v.sorted, err = types.ReadBool(r) + size, err := types.ReadInt32(r) + if err != nil { + return err + } + if size == 0 { + v.nsp.Reset() + return nil + } + if size < bitmap.MarshalHeaderSize { + return moerr.NewInvalidInputNoCtx("invalid bitmap wire size") + } + var header [bitmap.MarshalHeaderSize]byte + if _, err = io.ReadFull(r, header[:]); err != nil { + return err + } + _, bitLength, _, err := bitmap.DecodeMarshalHeader(header[:]) + if err != nil { + return moerr.NewInvalidInputNoCtx("invalid vector null bitmap") + } + if bitLength > int64(math.MaxInt) { + return moerr.NewInvalidInputNoCtx("vector null bitmap exceeds platform limit") + } + if err = v.ensureNullCapacity(int(bitLength), mp); err != nil { + return err + } + payload, err := v.nsp.GetBitmap().PrepareExternalUnmarshal( + header[:], + int(size), + ) if err != nil { return err } + if _, err = io.ReadFull(r, payload); err != nil { + v.nsp.Reset() + return err + } + return v.nsp.GetBitmap().Validate() +} - return nil +// GroupingMarshalBinarySize returns the optional grouping bitmap wire size. +func (v *Vector) GroupingMarshalBinarySize() int { + if v == nil { + return 0 + } + return v.gsp.MarshalSize() +} + +// MarshalGroupingTo writes the optional grouping bitmap without changing the +// stable Vector wire format. +func (v *Vector) MarshalGroupingTo(w io.Writer) error { + if v == nil || w == nil { + return io.ErrClosedPipe + } + return v.gsp.MarshalTo(w) +} + +// UnmarshalGroupingFromReader restores a grouping bitmap whose size is framed +// by the caller. +func (v *Vector) UnmarshalGroupingFromReader( + r io.Reader, + size int, + mp *mpool.MPool, +) error { + if v == nil || r == nil || size < 0 { + return moerr.NewInvalidInputNoCtx("invalid vector grouping bitmap") + } + if size == 0 { + v.gsp.Reset() + return nil + } + if size < bitmap.MarshalHeaderSize { + return moerr.NewInvalidInputNoCtx("invalid vector grouping bitmap") + } + var header [bitmap.MarshalHeaderSize]byte + if _, err := io.ReadFull(r, header[:]); err != nil { + return err + } + _, bitLength, _, err := bitmap.DecodeMarshalHeader(header[:]) + if err != nil || bitLength > int64(math.MaxInt) { + return moerr.NewInvalidInputNoCtx("invalid vector grouping bitmap") + } + if v.allocationAccount == nil { + data := make([]byte, size) + copy(data, header[:]) + if _, err = io.ReadFull(r, data[len(header):]); err != nil { + return err + } + if err = validateVectorNullBitmap(data, true); err != nil { + return err + } + v.gsp.Reset() + return v.gsp.Read(data) + } + if err = v.ensureGroupingCapacity(int(bitLength), mp); err != nil { + return err + } + payload, err := v.gsp.GetBitmap().PrepareExternalUnmarshal(header[:], size) + if err != nil { + return err + } + if _, err = io.ReadFull(r, payload); err != nil { + v.gsp.Reset() + return err + } + return v.gsp.GetBitmap().Validate() } func (v *Vector) ToConst() { @@ -1208,6 +1606,25 @@ func (v *Vector) PreExtend(rows int, mp *mpool.MPool) error { return extend(v, rows, mp) } +// PreExtendBitmap ensures allocation-accounted null and grouping storage can +// represent rows without allocating vector data. Unaccounted vectors are +// unchanged. +func (v *Vector) PreExtendBitmap(rows int, mp *mpool.MPool) error { + return v.ensureBitmapCapacity(rows, mp) +} + +// PreExtendNulls ensures allocation-accounted null storage can represent rows. +// Unaccounted vectors are unchanged. +func (v *Vector) PreExtendNulls(rows int, mp *mpool.MPool) error { + return v.ensureNullCapacity(rows, mp) +} + +// PreExtendGrouping ensures allocation-accounted grouping storage can +// represent rows. Unaccounted vectors are unchanged. +func (v *Vector) PreExtendGrouping(rows int, mp *mpool.MPool) error { + return v.ensureGroupingCapacity(rows, mp) +} + // PreExtendArea use to expand the mpool and area of vector // extraAreaSize: the size of area to be extended // mp: mpool @@ -1231,7 +1648,7 @@ func (v *Vector) PreExtendWithArea(rows int, extraAreaSize int, mp *mpool.MPool) // grow area var err error oldSz := len(area1) - area1, err = mp.Grow(area1, voff+extraAreaSize, v.offHeap) + area1, err = v.growArea(mp, voff+extraAreaSize) if err != nil { return err } @@ -1245,29 +1662,59 @@ func (v *Vector) PreExtendWithArea(rows int, extraAreaSize int, mp *mpool.MPool) // Dup use to copy an identical vector func (v *Vector) Dup(mp *mpool.MPool) (*Vector, error) { - return v.dup(mp, false, v.offHeap) + if v.allocationAccount != nil { + return v.dup(mp, true, true, v.allocationAccount) + } + return v.dup(mp, false, v.offHeap, nil) } // DupOffHeap copies a vector with all owned backing data allocated off-heap. func (v *Vector) DupOffHeap(mp *mpool.MPool) (*Vector, error) { - return v.dup(mp, true, true) + return v.dup(mp, true, true, v.allocationAccount) } -func (v *Vector) dup(mp *mpool.MPool, offHeap, areaOffHeap bool) (*Vector, error) { - if v.IsConstNull() { - return NewConstNull(v.typ, v.Length(), mp), nil - } - - var err error +// DupOffHeapWithAllocation copies a vector into an explicitly selected +// destination account. Passing nil creates an unaccounted destination. +func (v *Vector) DupOffHeapWithAllocation( + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + return v.dup(mp, true, true, selection) +} +func (v *Vector) dup( + mp *mpool.MPool, + offHeap bool, + areaOffHeap bool, + selection *AllocationAccountSelection, +) (*Vector, error) { w := NewVecFromReuse() w.offHeap = offHeap + if selection != nil { + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } + } w.class = v.class w.typ = v.typ - w.length = v.length w.sorted = v.sorted - w.GetNulls().InitWith(v.GetNulls()) + if v.IsConstNull() { + w.length = v.length + if v.HasGrouping() { + if err := w.ensureGroupingCapacity( + max(v.length, int(v.GetGrouping().GetBitmap().Len())), + mp, + ); err != nil { + w.Free(mp) + return nil, err + } + w.GetGrouping().InitWith(v.GetGrouping()) + } + return w, nil + } + + var err error dataLen := v.typ.TypeSize() if v.IsConst() { if err := extend(w, 1, mp); err != nil { @@ -1281,10 +1728,34 @@ func (v *Vector) dup(mp *mpool.MPool, offHeap, areaOffHeap bool) (*Vector, error } dataLen *= v.length } + // A bitmap may be shorter than a sparse vector or longer than a reused vector + // that was shortened with SetLength. Preserve both the complete row domain + // and the source bitmap extent before InitWith copies its storage. + if v.GetNulls().GetBitmap().Len() > 0 { + if err := w.ensureNullCapacity( + max(v.length, int(v.GetNulls().GetBitmap().Len())), + mp, + ); err != nil { + w.Free(mp) + return nil, err + } + } + if v.GetGrouping().GetBitmap().Len() > 0 { + if err := w.ensureGroupingCapacity( + max(v.length, int(v.GetGrouping().GetBitmap().Len())), + mp, + ); err != nil { + w.Free(mp) + return nil, err + } + } + w.length = v.length + w.GetNulls().InitWith(v.GetNulls()) + w.GetGrouping().InitWith(v.GetGrouping()) copy(w.data, v.data[:dataLen]) if len(v.area) > 0 { - if w.area, err = mp.Alloc(len(v.area), areaOffHeap); err != nil { + if w.area, err = w.allocOwned(mp, len(v.area), areaOffHeap, false); err != nil { w.Free(mp) return nil, err } @@ -1298,7 +1769,37 @@ func (v *Vector) dup(mp *mpool.MPool, offHeap, areaOffHeap bool) (*Vector, error // retains varlen payload referenced by the vector's logical rows, so stale or // unreferenced bytes in area are not propagated into batch memory accounting. func (v *Vector) CloneToFlatCompact(mp *mpool.MPool) (*Vector, error) { - w := NewVec(v.typ) + if v.allocationAccount != nil { + return nil, allocationAccountInvalid( + "accounted compact clone requires a destination selection", + ) + } + return v.cloneToFlatCompact(mp, nil) +} + +// CloneToFlatCompactWithAllocation creates an off-heap compact copy under the +// explicit destination selection. Passing nil creates an unaccounted +// destination and is reserved for a deliberate ownership boundary. +func (v *Vector) CloneToFlatCompactWithAllocation( + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + return v.cloneToFlatCompact(mp, selection) +} + +func (v *Vector) cloneToFlatCompact( + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + var w *Vector + if selection == nil { + w = NewVec(v.typ) + } else { + w = NewOffHeapVecWithType(v.typ) + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } + } if v.class != FLAT || (!v.typ.IsFixedLen() && !v.typ.IsVarlen()) { if err := GetUnionAllFunction(v.typ, mp)(w, v); err != nil { w.Free(mp) @@ -1311,7 +1812,13 @@ func (v *Vector) CloneToFlatCompact(mp *mpool.MPool) (*Vector, error) { if v.length == 0 { return w, nil } - if err := extend(w, v.length, mp); err != nil { + if err := extendWithBitmaps( + w, + v.length, + mp, + !v.nsp.EmptyByFlag(), + !v.gsp.EmptyByFlag(), + ); err != nil { w.Free(mp) return nil, err } @@ -1338,7 +1845,7 @@ func (v *Vector) CloneToFlatCompact(mp *mpool.MPool) (*Vector, error) { } if totalArea > 0 { var err error - w.area, err = mp.Alloc(totalArea, w.offHeap) + w.area, err = w.allocArea(mp, totalArea) if err != nil { w.Free(mp) return nil, err @@ -1612,6 +2119,12 @@ func (v *Vector) ShuffleWithBuf(sels []int64, mp *mpool.MPool, buf *[]byte) (err if v.IsConst() { return nil } + // The reusable buffer is Go-heap storage and therefore has no physical + // allocation provenance. Allocation-accounted vectors must use Shuffle, + // whose replacement data and bitmap scratch are admitted to their owner. + if v.allocationAccount != nil { + return v.Shuffle(sels, mp) + } // Fall back to allocating Shuffle if the vector doesn't own its data // or the selection changes the element count. if v.cantFreeData || len(sels) != v.length { @@ -1682,6 +2195,24 @@ func (v *Vector) ShuffleWithBuf(sels []int64, mp *mpool.MPool, buf *[]byte) (err // Copy simply does v[vi] = w[wi] func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { disjoint := v.areaDisjoint + sourceGrouping := w.GetGrouping().Contains(uint64(wi)) + if sourceGrouping && v.allocationAccount != nil { + if err := v.ensureGroupingCapacity(int(vi)+1, mp); err != nil { + return err + } + } + if sourceGrouping { + v.GetGrouping().Set(uint64(vi)) + } else { + v.GetGrouping().Unset(uint64(vi)) + } + sourceNull := w.IsConstNull() || + (!w.IsConst() && w.GetNulls().Contains(uint64(wi))) + if sourceNull && v.allocationAccount != nil { + if err := v.ensureNullCapacity(int(vi)+1, mp); err != nil { + return err + } + } if w.class == CONSTANT { if w.IsConstNull() { if !v.typ.IsFixedLen() { @@ -1696,7 +2227,7 @@ func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { // Non-null constant vectors still share the regular null/data path below. wi = 0 } - if w.GetNulls().Contains(uint64(wi)) { + if sourceNull { if !v.typ.IsFixedLen() { vva := MustFixedColNoTypeCheck[types.Varlena](v) vva[vi] = types.Varlena{} @@ -1734,6 +2265,38 @@ func (v *Vector) Copy(w *Vector, vi, wi int64, mp *mpool.MPool) error { // GetUnionAllFunction: A more sensible function for copying vector, // which avoids having to do type conversions and type judgements every time you append. func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) error { + union := getUnionAllFunction(typ, mp) + return func(v, w *Vector) error { + oldLength := v.length + if w.gsp.Any() { + if err := v.ensureGroupingCapacity(oldLength+w.length, mp); err != nil { + return err + } + } + if err := union(v, w); err != nil { + return err + } + if w.gsp.Any() { + unionVectorBitmap(&v.gsp, &w.gsp, oldLength, w.length) + } + return nil + } +} + +func unionVectorBitmap( + destination *nulls.Nulls, + source *nulls.Nulls, + offset int, + length int, +) { + for row := 0; row < length; row++ { + if source.Contains(uint64(row)) { + destination.Set(uint64(offset + row)) + } + } +} + +func getUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) error { // a more simple and quickly union nsp but not good. unionNsp := func(dst *nulls.Nulls, more *nulls.Nulls, oldLength int, moreLength int) { u64offset := uint64(oldLength) @@ -1770,7 +2333,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1799,7 +2362,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1828,7 +2391,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1857,7 +2420,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1886,7 +2449,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1915,7 +2478,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1944,7 +2507,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -1973,7 +2536,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2002,7 +2565,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2031,7 +2594,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2060,7 +2623,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2089,7 +2652,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2118,7 +2681,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2147,7 +2710,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2176,7 +2739,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2205,7 +2768,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2234,7 +2797,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2263,7 +2826,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2292,7 +2855,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2321,7 +2884,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2350,7 +2913,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2379,7 +2942,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2408,7 +2971,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2437,7 +3000,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2468,11 +3031,11 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if sz := len(v.area) + len(w.area); sz > cap(v.area) { - area, err := mp.Grow(v.area, sz, v.offHeap) + area, err := v.growArea(mp, sz) if err != nil { return err } @@ -2481,13 +3044,13 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err var err error vs := toSliceOfLengthNoTypeCheck[types.Varlena](v, v.length+w.length) + if w.gsp.Any() { + unionNsp(&v.gsp, &w.gsp, v.length, w.length) + } bm := w.nsp.GetBitmap() if bm != nil && !bm.EmptyByFlag() { for i := range ws { - if w.gsp.Contains(uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if bm.Contains(uint64(i)) { vs[v.length] = types.Varlena{} nulls.Add(&v.nsp, uint64(v.length)) @@ -2525,7 +3088,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err } return nil } - if err := extend(v, w.length, mp); err != nil { + if err := extendWithBitmaps(v, w.length, mp, w.nsp.Any(), w.gsp.Any()); err != nil { return err } if w.nsp.Any() { @@ -2546,7 +3109,7 @@ func GetUnionAllFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector) err // GetConstSetFunction: A more sensible function for const vector set, // which avoids having to do type conversions and type judgements every time you append. -func GetConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel int64, length int) error { +func getConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel int64, length int) error { switch typ.Oid { case types.T_bool: return func(v, w *Vector, sel int64, length int) error { @@ -2841,6 +3404,29 @@ func GetConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel } } +func GetConstSetFunction(typ types.Type, mp *mpool.MPool) func(v, w *Vector, sel int64, length int) error { + set := getConstSetFunction(typ, mp) + return func(v, w *Vector, sel int64, length int) error { + if v == nil || w == nil || sel < 0 || sel >= int64(w.Length()) || length < 0 { + return moerr.NewInvalidInputNoCtx("invalid const vector selection") + } + grouping := w.gsp.Contains(uint64(sel)) + if grouping { + if err := v.ensureGroupingCapacity(length, mp); err != nil { + return err + } + } + if err := set(v, w, sel, length); err != nil { + return err + } + v.gsp.Reset() + if grouping && length > 0 { + v.gsp.AddRange(0, uint64(length)) + } + return nil + } +} + // fillSlice broadcasts val across s[start:end] using exponential copy doubling: // write one element, then double the filled region with copy() — O(log n) memmoves // instead of n scalar element stores. Used on the hot const-broadcast path. @@ -2878,7 +3464,7 @@ func pregrowVarlenaArea(vec *Vector, totalBytes int, mp *mpool.MPool) error { return nil } origLen := len(vec.area) - grown, err := mp.Grow(vec.area, need, vec.offHeap) + grown, err := vec.growArea(mp, need) if err != nil { return err } @@ -2895,26 +3481,34 @@ func (v *Vector) UnionOne(w *Vector, sel int64, mp *mpool.MPool) error { if v.typ.IsVarlen() { v.areaDisjoint = false } - if err := extend(v, 1, mp); err != nil { + sourceGrouping := nulls.Contains(&w.gsp, uint64(sel)) + sourceNull := w.IsConstNull() || + (!w.IsConst() && nulls.Contains(&w.nsp, uint64(sel))) + if err := extendWithBitmaps( + v, + 1, + mp, + sourceNull && v.allocationAccount != nil, + sourceGrouping && v.allocationAccount != nil, + ); err != nil { return err } oldLen := v.length v.length++ - if nulls.Contains(&w.gsp, uint64(sel)) { + if sourceGrouping { nulls.Add(&v.gsp, uint64(oldLen)) } if w.IsConst() { - if w.IsConstNull() { + if sourceNull { nulls.Add(&v.nsp, uint64(oldLen)) return nil } sel = 0 - } else if nulls.Contains(&w.nsp, uint64(sel)) { + } else if sourceNull { nulls.Add(&v.nsp, uint64(oldLen)) return nil } - if v.GetType().IsVarlen() { var vs, ws []types.Varlena ToSliceNoTypeCheck(v, &vs) @@ -2948,6 +3542,22 @@ func (v *Vector) UnionOne(w *Vector, sel int64, mp *mpool.MPool) error { return nil } +func appendSelectedGrouping[T int32 | int64]( + dst *Vector, + src *Vector, + oldLength int, + sels []T, +) { + if src.gsp.EmptyByFlag() { + return + } + for i, sel := range sels { + if src.gsp.Contains(uint64(sel)) { + nulls.Add(&dst.gsp, uint64(oldLength+i)) + } + } +} + // It is simply append. the purpose of retention is ease of use func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) error { if v.typ.IsVarlen() { @@ -2957,26 +3567,34 @@ func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) erro return nil } - if err := extend(v, cnt, mp); err != nil { + sourceGrouping := nulls.Contains(&w.gsp, uint64(sel)) + sourceNull := w.IsConstNull() || + (!w.IsConst() && nulls.Contains(&w.nsp, uint64(sel))) + if err := extendWithBitmaps( + v, + cnt, + mp, + sourceNull && v.allocationAccount != nil, + sourceGrouping && v.allocationAccount != nil, + ); err != nil { return err } oldLen := v.length v.length += cnt - if nulls.Contains(&w.gsp, uint64(sel)) { + if sourceGrouping { nulls.AddRange(&v.gsp, uint64(oldLen), uint64(oldLen+cnt)) } if w.IsConst() { - if w.IsConstNull() { + if sourceNull { nulls.AddRange(&v.nsp, uint64(oldLen), uint64(oldLen+cnt)) return nil } sel = 0 - } else if nulls.Contains(&w.nsp, uint64(sel)) { + } else if sourceNull { nulls.AddRange(&v.nsp, uint64(oldLen), uint64(oldLen+cnt)) return nil } - if v.GetType().IsVarlen() { var err error var va types.Varlena @@ -2998,6 +3616,37 @@ func (v *Vector) UnionMulti(w *Vector, sel int64, cnt int, mp *mpool.MPool) erro return nil } +func appendBatchGrouping( + dst *Vector, + src *Vector, + oldLength int, + offset int64, + cnt int, + flags []uint8, +) { + if src.gsp.EmptyByFlag() { + return + } + output := oldLength + if flags == nil { + for i := range cnt { + if src.gsp.Contains(uint64(offset) + uint64(i)) { + nulls.Add(&dst.gsp, uint64(output+i)) + } + } + return + } + for i, selected := range flags { + if selected == 0 { + continue + } + if src.gsp.Contains(uint64(offset) + uint64(i)) { + nulls.Add(&dst.gsp, uint64(output)) + } + output++ + } +} + func (v *Vector) Union(w *Vector, sels []int64, mp *mpool.MPool) error { return unionT[int64](v, w, sels, mp) } @@ -3013,7 +3662,13 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { return nil } - if err := extend(v, len(sels), mp); err != nil { + if err := extendWithBitmaps( + v, + len(sels), + mp, + w.IsConstNull() || !w.nsp.EmptyByFlag(), + w.IsGrouping() || !w.gsp.EmptyByFlag(), + ); err != nil { return err } @@ -3045,6 +3700,7 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { return nil } + appendSelectedGrouping(v, w, oldLen, sels) if v.GetType().IsVarlen() { var err error @@ -3072,9 +3728,6 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { } if !w.GetNulls().EmptyByFlag() { for i, sel := range sels { - if w.gsp.Contains(uint64(sel)) { - nulls.Add(&v.gsp, uint64(oldLen+i)) - } if w.nsp.Contains(uint64(sel)) { nulls.Add(&v.nsp, uint64(oldLen+i)) continue @@ -3097,9 +3750,6 @@ func unionT[T int32 | int64](v, w *Vector, sels []T, mp *mpool.MPool) error { tlen := v.GetType().TypeSize() if !w.nsp.EmptyByFlag() { for i, sel := range sels { - if w.gsp.Contains(uint64(sel)) { - nulls.Add(&v.gsp, uint64(oldLen+i)) - } if w.nsp.Contains(uint64(sel)) { nulls.Add(&v.nsp, uint64(oldLen+i)) continue @@ -3158,7 +3808,13 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp return nil } - if err := extend(v, addCnt, mp); err != nil { + if err := extendWithBitmaps( + v, + addCnt, + mp, + w.IsConstNull() || !w.nsp.EmptyByFlag(), + w.IsGrouping() || !w.gsp.EmptyByFlag(), + ); err != nil { return err } @@ -3190,6 +3846,7 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp return nil } + appendBatchGrouping(v, w, v.length, offset, cnt, flags) if v.GetType().IsVarlen() { var err error @@ -3214,10 +3871,20 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if len(w.area) > 0 { // preserve mpool semantics: append within cap, else mpool Grow2 (so // v.area stays mpool-tracked rather than escaping to the Go heap). - if baseOff+len(w.area) <= cap(v.area) || mp == nil { + if baseOff+len(w.area) <= cap(v.area) { v.area = append(v.area, w.area...) - } else if v.area, err = mp.Grow2(v.area, w.area, baseOff+len(w.area), v.offHeap); err != nil { - return err + } else if mp == nil { + if v.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector area growth does not have a mpool", + ) + } + v.area = append(v.area, w.area...) + } else { + v.area, err = v.growArea2(mp, w.area, baseOff+len(w.area)) + if err != nil { + return err + } } } // one memmove of the header array; inline varlenas carry their bytes here. @@ -3233,21 +3900,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp } } } - // propagate grouping bits (value is still real for these rows). - // Bound to [0,cnt): Foreach walks every set bit in the underlying - // bitmap, but w may carry stale bits at index >= w.length (SetLength - // shrinks length without clearing nsp/gsp, and vectors are reused). - // The per-row path only consults [0,cnt) via Contains, so we must skip - // stale bits here too — otherwise they pollute v.gsp / index past vCol. - if !w.gsp.EmptyByFlag() { - base, ucnt := uint64(oldLen), uint64(cnt) - w.gsp.Foreach(func(i uint64) bool { - if i < ucnt { - nulls.Add(&v.gsp, base+i) - } - return true - }) - } // propagate null bits and clear those (never-read) headers so a copied // big-header offset can't linger as a dangling reference into v.area. // Same [0,cnt) bound as gsp above: a stale nsp bit at i >= cnt would @@ -3307,9 +3959,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if !w.nsp.EmptyByFlag() { if flags == nil { for i := 0; i < cnt; i++ { - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3325,9 +3974,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3342,9 +3988,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp } else { if flags == nil { for i := 0; i < cnt; i++ { - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } err = BuildVarlenaFromVarlena(v, &vCol[v.length], &wCol[int(offset)+i], &w.area, mp) if err != nil { return err @@ -3356,9 +3999,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } err = BuildVarlenaFromVarlena(v, &vCol[v.length], &wCol[int(offset)+i], &w.area, mp) if err != nil { return err @@ -3372,9 +4012,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if !w.nsp.EmptyByFlag() { if flags == nil { for i := 0; i < cnt; i++ { - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3387,9 +4024,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } if w.nsp.Contains(uint64(offset) + uint64(i)) { nulls.Add(&v.nsp, uint64(v.length)) } else { @@ -3417,9 +4051,6 @@ func (v *Vector) UnionBatch(w *Vector, offset int64, cnt int, flags []uint8, mp if flags[i] == 0 { continue } - if w.gsp.Contains(uint64(offset) + uint64(i)) { - nulls.Add(&v.gsp, uint64(v.length)) - } copy(v.data[v.length*tlen:(v.length+1)*tlen], w.data[(int(offset)+i)*tlen:(int(offset)+i+1)*tlen]) v.length++ } @@ -4090,7 +4721,7 @@ func appendOneFixed[T any](vec *Vector, val T, isNull bool, mp *mpool.MPool) err return moerr.NewInternalErrorNoCtx("append to const vector") } - if err := extend(vec, 1, mp); err != nil { + if err := extendWithBitmaps(vec, 1, mp, isNull, false); err != nil { return err } length := vec.length @@ -4169,10 +4800,9 @@ func appendOneArray[T types.ArrayElement](vec *Vector, val []T, isNull bool, mp } } -// appendOneOwnedVarlena installs a descriptor built against vec.area by one of -// the helpers above. That construction either keeps the value inline or -// appends a fresh range, so it preserves an existing disjoint-area proof -// without the invalidate/restore stores required by generic descriptor writes. +// appendOneOwnedVarlena installs a descriptor built against vec.area. The +// value is either inline or references a freshly appended area range, so this +// preserves an existing disjoint-area proof. func appendOneOwnedVarlena( vec *Vector, value types.Varlena, @@ -4183,8 +4813,7 @@ func appendOneOwnedVarlena( } index := vec.length vec.length++ - values := toSliceOfLengthNoTypeCheck[types.Varlena](vec, vec.length) - values[index] = value + toSliceOfLengthNoTypeCheck[types.Varlena](vec, vec.length)[index] = value return nil } @@ -4192,7 +4821,7 @@ func appendMultiFixed[T any](vec *Vector, val T, isNull bool, cnt int, mp *mpool if vec.typ.IsVarlen() && !isNull { vec.areaDisjoint = false } - if err := extend(vec, cnt, mp); err != nil { + if err := extendWithBitmaps(vec, cnt, mp, isNull, false); err != nil { return err } length := vec.length @@ -4216,7 +4845,7 @@ func appendMultiBytes(vec *Vector, val []byte, isNull bool, cnt int, mp *mpool.M vec.areaDisjoint = false var err error var va types.Varlena - if err = extend(vec, cnt, mp); err != nil { + if err = extendWithBitmaps(vec, cnt, mp, isNull, false); err != nil { return err } length := vec.length @@ -4239,10 +4868,15 @@ func appendMultiBytes(vec *Vector, val []byte, isNull bool, cnt int, mp *mpool.M func appendList[T any](vec *Vector, vals []T, isNulls []bool, mp *mpool.MPool) error { if vec.typ.IsVarlen() { - // Generic lists can contain repeated or externally-owned descriptors. vec.areaDisjoint = false } - if err := extend(vec, len(vals), mp); err != nil { + if err := extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4262,7 +4896,13 @@ func appendBytesList(vec *Vector, vals [][]byte, isNulls []bool, mp *mpool.MPool var err error disjoint := vec.areaDisjoint vec.areaDisjoint = false - if err = extend(vec, len(vals), mp); err != nil { + if err = extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4290,7 +4930,13 @@ func appendStringList(vec *Vector, vals []string, isNulls []bool, mp *mpool.MPoo disjoint := vec.areaDisjoint vec.areaDisjoint = false - if err = extend(vec, len(vals), mp); err != nil { + if err = extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4320,7 +4966,13 @@ func appendArrayList[T types.ArrayElement](vec *Vector, vals [][]T, isNulls []bo disjoint := vec.areaDisjoint vec.areaDisjoint = false - if err = extend(vec, len(vals), mp); err != nil { + if err = extendWithBitmaps( + vec, + len(vals), + mp, + slices.Contains(isNulls, true), + false, + ); err != nil { return err } length := vec.length @@ -4350,8 +5002,8 @@ func shrinkFixed[T types.FixedSizeT](v *Vector, sels []int64, negate bool) { for i, sel := range sels { vs[i] = vs[sel] } - nulls.Filter(&v.gsp, sels, false) - nulls.Filter(&v.nsp, sels, false) + nulls.FilterInPlaceOrdered(&v.gsp, sels, false) + nulls.FilterInPlaceOrdered(&v.nsp, sels, false) v.length = len(sels) } else if len(sels) > 0 { for oldIdx, newIdx, selIdx, sel := 0, 0, 0, sels[0]; oldIdx < v.length; oldIdx++ { @@ -4370,8 +5022,8 @@ func shrinkFixed[T types.FixedSizeT](v *Vector, sels []int64, negate bool) { sel = sels[selIdx] } } - nulls.Filter(&v.gsp, sels, true) - nulls.Filter(&v.nsp, sels, true) + nulls.FilterInPlaceOrdered(&v.gsp, sels, true) + nulls.FilterInPlaceOrdered(&v.nsp, sels, true) v.length -= len(sels) } } @@ -4386,8 +5038,8 @@ func shrinkFixedByMask[T types.FixedSizeT](v *Vector, sels *bitmap.Bitmap, negat vs[idx] = vs[itr.Next()+offset] idx++ } - nulls.FilterByMask(&v.gsp, sels, false) - nulls.FilterByMask(&v.nsp, sels, false) + nulls.FilterByMaskInPlace(&v.gsp, sels, false) + nulls.FilterByMaskInPlace(&v.nsp, sels, false) v.length = length } else if length > 0 { sel := itr.Next() + offset @@ -4406,8 +5058,8 @@ func shrinkFixedByMask[T types.FixedSizeT](v *Vector, sels *bitmap.Bitmap, negat sel = itr.Next() + offset } } - nulls.FilterByMask(&v.gsp, sels, true) - nulls.FilterByMask(&v.nsp, sels, true) + nulls.FilterByMaskInPlace(&v.gsp, sels, true) + nulls.FilterByMaskInPlace(&v.nsp, sels, true) v.length -= length } } @@ -4419,16 +5071,18 @@ func shuffleFixedNoTypeCheck[T types.FixedSizeT](v *Vector, sels []int64, mp *mp ns := len(sels) var vs []T ToFixedColNoTypeCheck(v, &vs) - data, err := mp.Alloc(ns*v.GetType().TypeSize(), v.offHeap) + data, err := v.allocData(mp, ns*v.GetType().TypeSize()) if err != nil { return err } - v.data = data - ws := toSliceOfLengthNoTypeCheck[T](v, ns) + ws := util.UnsafeSliceCastToLength[T](data, ns) shuffle.FixedLengthShuffle(vs, ws, sels) - nulls.Filter(&v.gsp, sels, false) - nulls.Filter(&v.nsp, sels, false) + if err := v.remapShuffleBitmaps(sels, mp); err != nil { + mp.Free(data) + return err + } + v.data = data // XXX We should never allow "half-owned" vectors later. And unowned vector should be strictly read-only. if v.cantFreeData { v.cantFreeData = false @@ -4439,6 +5093,95 @@ func shuffleFixedNoTypeCheck[T types.FixedSizeT](v *Vector, sels []int64, mp *mp return nil } +type bitmapRemapScratch struct { + destination *bitmap.Bitmap + value bitmap.Bitmap + storage []uint64 +} + +func (s *bitmapRemapScratch) release(mp *mpool.MPool) { + if s == nil || cap(s.storage) == 0 { + return + } + s.value.ReleaseExternalStorage() + mpool.FreeSlice(mp, s.storage) + s.storage = nil +} + +// remapShuffleBitmaps preserves Shuffle's arbitrary-selection semantics. An +// allocation-accounted vector builds both results in admitted temporary +// storage before publishing either, so rejection cannot leave null and +// grouping ownership half-mutated. +func (v *Vector) remapShuffleBitmaps(sels []int64, mp *mpool.MPool) error { + if v.allocationAccount == nil { + nulls.Filter(&v.gsp, sels, false) + nulls.Filter(&v.nsp, sels, false) + return nil + } + + targets := [...]struct { + destination *bitmap.Bitmap + site mpool.AllocationSite + }{ + {v.gsp.GetBitmap(), v.allocationAccount.groupingSite}, + {v.nsp.GetBitmap(), v.allocationAccount.nullsSite}, + } + if targets[0].destination.EmptyByFlag() && + targets[1].destination.EmptyByFlag() { + return nil + } + if !targets[0].destination.EmptyByFlag() { + if err := v.ensureGroupingCapacity(len(sels), mp); err != nil { + return err + } + } + if !targets[1].destination.EmptyByFlag() { + if err := v.ensureNullCapacity(len(sels), mp); err != nil { + return err + } + } + + var scratch [2]bitmapRemapScratch + for i, target := range targets { + if target.destination.EmptyByFlag() { + continue + } + words := (len(sels) + 63) / 64 + storage, err := mpool.MakeSliceAccounted[uint64]( + words, + mp, + v.allocationAccount.account, + v.allocationAccount.owner, + target.site, + ) + if err != nil { + for j := range i { + scratch[j].release(mp) + } + return err + } + scratch[i].destination = target.destination + scratch[i].storage = storage + scratch[i].value.InstallExternalStorage(storage) + scratch[i].value.InitWithSize(int64(len(sels))) + for output, source := range sels { + if target.destination.Contains(uint64(source)) { + scratch[i].value.Add(uint64(output)) + } + } + } + + for i := range scratch { + if scratch[i].destination != nil { + scratch[i].destination.InitWith(&scratch[i].value) + } + } + for i := range scratch { + scratch[i].release(mp) + } + return nil +} + // shuffleFixedNoTypeCheckWithBuf permutes elements using a reusable scratch // buffer instead of allocating a new data buffer. Only valid when // len(sels) == v.length and !v.cantFreeData (caller checks). @@ -4486,29 +5229,60 @@ func vecToString[T types.FixedSizeT](v *Vector) string { // The returned object is NOT allowed to be modified ( // TODO: Nulls are deep copied. func (v *Vector) Window(start, end int) (*Vector, error) { - if v.IsConstNull() { - return NewConstNull(v.typ, end-start, nil), nil - } else if v.IsConst() { - vec := NewVec(v.typ) - vec.class = v.class - if v.typ.IsVarlen() { - vec.areaDisjoint = false - } - vec.data = v.data - vec.area = v.area - vec.length = end - start - vec.cantFreeArea = true - vec.cantFreeData = true - vec.sorted = v.sorted - return vec, nil + return v.window(start, end, nil, nil) +} + +// WindowWithAllocation returns a borrowed data window whose range bitmaps are +// physical allocations in selection. Accounted pressure paths must use this +// form so shrinking an operation cannot create invisible Go-heap owners. +func (v *Vector) WindowWithAllocation( + start int, + end int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + if mp == nil || selection == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + return v.window(start, end, mp, selection) +} + +func (v *Vector) window( + start int, + end int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { + if start < 0 || end < start || end > v.Length() { + return nil, moerr.NewInvalidInputNoCtx("invalid vector window") } w := NewVec(v.typ) - if start == end { - return w, nil + if selection != nil { + w.offHeap = true + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } } - nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) - w.data = v.data[start*v.typ.TypeSize() : end*v.typ.TypeSize()] + w.class = v.class w.length = end - start + w.sorted = v.sorted + if err := v.copyWindowBitmaps(w, start, end, mp); err != nil { + w.Free(mp) + return nil, err + } + if v.IsConst() { + if v.typ.IsVarlen() { + w.areaDisjoint = false + } + w.data = v.data + w.area = v.area + w.cantFreeArea = true + w.cantFreeData = true + return w, nil + } + if start != end { + w.data = v.data[start*v.typ.TypeSize() : end*v.typ.TypeSize()] + } if v.typ.IsVarlen() { w.area = v.area w.areaDisjoint = v.areaDisjoint @@ -4518,33 +5292,65 @@ func (v *Vector) Window(start, end int) (*Vector, error) { return w, nil } +func (v *Vector) copyWindowBitmaps(w *Vector, start, end int, mp *mpool.MPool) error { + length := end - start + hasNull := v.nsp.GetBitmap().CountRange(uint64(start), uint64(end)) > 0 + hasGrouping := v.gsp.GetBitmap().CountRange(uint64(start), uint64(end)) > 0 + if hasNull { + if err := w.PreExtendNulls(length, mp); err != nil { + return err + } + nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) + } + if hasGrouping { + if err := w.PreExtendGrouping(length, mp); err != nil { + return err + } + nulls.Range(&v.gsp, uint64(start), uint64(end), uint64(start), &w.gsp) + } + return nil +} + // CloneWindow Deep copies the content from start to end into another vector. Afterwise it's safe to destroy the original one. func (v *Vector) CloneWindow(start, end int, mp *mpool.MPool) (*Vector, error) { + return v.CloneWindowWithAllocation( + start, + end, + mp, + v.allocationAccount, + ) +} + +// CloneWindowWithAllocation deep-copies a window into an explicitly selected +// off-heap destination account. +func (v *Vector) CloneWindowWithAllocation( + start int, + end int, + mp *mpool.MPool, + selection *AllocationAccountSelection, +) (*Vector, error) { if start == end { - return NewOffHeapVecWithType(v.typ), nil + w := NewOffHeapVecWithType(v.typ) + if selection != nil { + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err + } + } + return w, nil } if end > v.Length() { panic(fmt.Sprintf("CloneWindow end %d >= length %d", end, v.Length())) } - if v.IsConstNull() { - return NewConstNull(v.typ, end-start, mp), nil - } else if v.IsConst() { - if v.typ.IsVarlen() { - return NewConstBytes(v.typ, v.GetBytesAt(0), end-start, mp) - } else { - vec := NewOffHeapVecWithType(v.typ) - vec.class = v.class - vec.data = make([]byte, len(v.data)) - copy(vec.data, v.data) - vec.length = end - start - vec.cantFreeArea = true - vec.cantFreeData = true - vec.sorted = v.sorted - return vec, nil + w := NewOffHeapVecWithType(v.typ) + if selection != nil { + if err := w.SetAllocationAccount(selection); err != nil { + return nil, err } } - w := NewOffHeapVecWithType(v.typ) if err := v.CloneWindowTo(w, start, end, mp); err != nil { + if mp != nil { + w.Free(mp) + } return nil, err } return w, nil @@ -4554,6 +5360,9 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error if start == end { return nil } + if err := v.copyWindowBitmaps(w, start, end, mp); err != nil { + return err + } if v.IsConstNull() { w.class = CONSTANT if v.typ.IsVarlen() { @@ -4564,21 +5373,36 @@ func (v *Vector) CloneWindowTo(w *Vector, start, end int, mp *mpool.MPool) error return nil } else if v.IsConst() { if v.typ.IsVarlen() { + w.class = CONSTANT return SetConstBytes(w, v.GetBytesAt(0), end-start, mp) } else { + if mp == nil { + if w.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector clone does not have a mpool", + ) + } + w.data = make([]byte, len(v.data)) + w.cantFreeData = true + } else { + if err := w.PreExtend(1, mp); err != nil { + return err + } + copy(w.data, v.data) + } w.class = v.class - w.data = make([]byte, len(v.data)) - copy(w.data, v.data) w.length = end - start - w.cantFreeArea = true - w.cantFreeData = true w.sorted = v.sorted return nil } } - nulls.Range(&v.nsp, uint64(start), uint64(end), uint64(start), &w.nsp) length := (end - start) * v.typ.TypeSize() if mp == nil { + if w.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector clone does not have a mpool", + ) + } w.data = make([]byte, length) copy(w.data, v.data[start*v.typ.TypeSize():end*v.typ.TypeSize()]) w.length = end - start @@ -5670,14 +6494,25 @@ func BuildVarlenaNoInline(vec *Vector, v1 *types.Varlena, bs *[]byte, m *mpool.M vlen := len(*bs) area1 := vec.GetArea() voff := len(area1) - if voff+vlen <= cap(area1) || m == nil { + if voff+vlen <= cap(area1) { + area1 = append(area1, *bs...) + v1.SetOffsetLen(uint32(voff), uint32(vlen)) + vec.area = area1 + return nil + } + if m == nil { + if vec.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector area growth does not have a mpool", + ) + } area1 = append(area1, *bs...) v1.SetOffsetLen(uint32(voff), uint32(vlen)) vec.area = area1 return nil } var err error - area1, err = m.Grow2(area1, *bs, voff+vlen, vec.offHeap) + area1, err = vec.growArea2(m, *bs, voff+vlen) if err != nil { return err } @@ -5695,13 +6530,18 @@ func BuildVarlenaNoInlineFromByteJson(vec *Vector, v1 *types.Varlena, bj bytejso if voff+vlen > cap(area1) && m != nil { // Pass nil to Grow2, we can grow area1 to voff+vlen without // copy bytejson data. - area1, err = m.Grow2(area1, nil, voff+vlen, vec.offHeap) + area1, err = vec.growArea2(m, nil, voff+vlen) if err != nil { return err } area1[voff] = byte(bj.Type) copy(area1[voff+1:voff+vlen], bj.Data) } else { + if voff+vlen > cap(area1) && vec.allocationAccount != nil { + return moerr.NewInternalErrorNoCtx( + "accounted vector area growth does not have a mpool", + ) + } area1 = append(area1, byte(bj.Type)) area1 = append(area1, bj.Data...) } @@ -5794,7 +6634,7 @@ func BuildVarlenaFromByteJsonEncoded( } if int(newAreaLen) > cap(vec.area) { - newArea, err := m.Grow2(vec.area, nil, int(newAreaLen), vec.offHeap) + newArea, err := vec.growArea2(m, nil, int(newAreaLen)) if err != nil { return err } diff --git a/pkg/container/vector/vector_test.go b/pkg/container/vector/vector_test.go index 848489c4f7dfa..a86e333fc0f37 100644 --- a/pkg/container/vector/vector_test.go +++ b/pkg/container/vector/vector_test.go @@ -15,7 +15,9 @@ package vector import ( + "bytes" "fmt" + "io" "slices" "strings" "testing" @@ -62,6 +64,34 @@ func TestLength(t *testing.T) { } } +func TestAppendCheckpointRollback(t *testing.T) { + mp := mpool.MustNewZero() + vec := NewVec(types.T_varchar.ToType()) + defer vec.Free(mp) + first := strings.Repeat("a", 64) + require.NoError(t, AppendBytes(vec, []byte(first), false, mp)) + vec.GetGrouping().Set(0) + vec.SetSorted(true) + checkpoint := vec.MakeAppendCheckpoint() + + require.NoError(t, AppendBytes(vec, []byte(strings.Repeat("b", 96)), false, mp)) + vec.GetNulls().Set(1) + vec.GetGrouping().Set(1) + // Grouping publication can precede a failed varlen copy and therefore can + // extend beyond the length reached by the copy itself. + vec.GetGrouping().Set(2) + vec.SetSorted(false) + vec.RollbackAppend(checkpoint, 2) + + require.Equal(t, 1, vec.Length()) + require.Equal(t, []string{first}, InefficientMustStrCol(vec)) + require.False(t, vec.GetNulls().Contains(1)) + require.True(t, vec.GetGrouping().Contains(0)) + require.False(t, vec.GetGrouping().Contains(1)) + require.False(t, vec.GetGrouping().Contains(2)) + require.True(t, vec.GetSorted()) +} + func TestCapacityForUntypedNull(t *testing.T) { vec := NewVec(types.T_any.ToType()) require.Equal(t, 0, vec.Capacity()) @@ -1506,6 +1536,21 @@ func TestShuffle(t *testing.T) { func TestCopy(t *testing.T) { mp := mpool.MustNewZero() + { // fixed grouping provenance + dst := NewVec(types.T_int32.ToType()) + src := NewVec(types.T_int32.ToType()) + require.NoError(t, AppendFixedList(dst, []int32{0, 0}, nil, mp)) + require.NoError(t, AppendFixedList(src, []int32{1, 2}, nil, mp)) + src.GetGrouping().Add(0) + dst.GetGrouping().Add(1) + require.NoError(t, dst.Copy(src, 0, 0, mp)) + require.NoError(t, dst.Copy(src, 1, 1, mp)) + require.True(t, dst.GetGrouping().Contains(0)) + require.False(t, dst.GetGrouping().Contains(1)) + dst.Free(mp) + src.Free(mp) + require.Equal(t, int64(0), mp.CurrNB()) + } { // fixed v := NewVec(types.T_int8.ToType()) AppendFixedList(v, []int8{0, 0, 1, 0}, nil, mp) @@ -1725,6 +1770,13 @@ func TestMarshalAndUnMarshal(t *testing.T) { require.NoError(t, err) data, err := v.MarshalBinary() require.NoError(t, err) + size, err := v.MarshalBinarySize() + require.NoError(t, err) + require.Equal(t, len(data), size) + var streamed bytes.Buffer + require.NoError(t, v.MarshalBinaryTo(&streamed)) + require.Equal(t, data, streamed.Bytes()) + require.ErrorIs(t, v.MarshalBinaryTo(shortVectorMarshalWriter{}), io.ErrShortWrite) w := NewVecFromReuse() err = w.UnmarshalBinary(data) require.NoError(t, err) @@ -1739,6 +1791,29 @@ func TestMarshalAndUnMarshal(t *testing.T) { require.Equal(t, int64(0), mp.CurrNB()) } +type shortVectorMarshalWriter struct{} + +func (shortVectorMarshalWriter) Write(value []byte) (int, error) { + return len(value) - 1, nil +} + +func TestMarshalBinarySizeRejectsInvalidVector(t *testing.T) { + var nilVector *Vector + _, err := nilVector.MarshalBinarySize() + require.Error(t, err) + + typ := types.T_int64.ToType() + typ.Size = -1 + invalidType := NewVec(typ) + _, err = invalidType.MarshalBinarySize() + require.Error(t, err) + + shortData := NewVec(types.T_int64.ToType()) + shortData.SetLength(1) + _, err = shortData.MarshalBinarySize() + require.Error(t, err) +} + func TestUnmarshalBinaryAcceptsNullBitmapCoveragePastLength(t *testing.T) { mp := mpool.MustNewZero() source := NewVec(types.T_int64.ToType()) @@ -3011,15 +3086,17 @@ func TestGetAny(t *testing.T) { func BenchmarkUnmarshal(b *testing.B) { mp := mpool.MustNewZero() - vec := NewVec(types.T_int8.ToType()) - AppendAny(vec, int8(42), false, mp) - data, err := vec.MarshalBinary() + source := NewVec(types.T_int8.ToType()) + AppendAny(source, int8(42), false, mp) + data, err := source.MarshalBinary() if err != nil { b.Fatal(err) } + source.Free(mp) + var target Vector b.ResetTimer() for i := 0; i < b.N; i++ { - err := vec.UnmarshalBinary(data) + err := target.UnmarshalBinary(data) if err != nil { b.Fatal(err) } diff --git a/pkg/container/vector/versions.go b/pkg/container/vector/versions.go index 347fc2c05eb97..30a01bf038eb3 100644 --- a/pkg/container/vector/versions.go +++ b/pkg/container/vector/versions.go @@ -71,6 +71,11 @@ func (v *Vector) MarshalBinaryWithBufferV1(buf *bytes.Buffer) error { func (v *Vector) UnmarshalBinaryV1(data []byte) error { v.areaDisjoint = false + if v.allocationAccount != nil { + return allocationAccountInvalid( + "cannot install aliases in an accounted vector", + ) + } // read class v.class = int(data[0]) data = data[1:] @@ -117,6 +122,7 @@ func (v *Vector) UnmarshalBinaryV1(data []byte) error { v.cantFreeData = true v.cantFreeArea = true + v.allocationAccount = nil return nil } diff --git a/pkg/frontend/export.go b/pkg/frontend/export.go index f2e46dfedbcdc..51216a78fa80c 100644 --- a/pkg/frontend/export.go +++ b/pkg/frontend/export.go @@ -377,6 +377,19 @@ func escapeJSONControlChars(s string) string { return builder.String() } +func sendExportBatchByte( + ctx context.Context, + byteChan chan *BatchByte, + value *BatchByte, +) bool { + select { + case byteChan <- value: + return true + case <-ctx.Done(): + return false + } +} + func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index int32, ByteChan chan *BatchByte, ep *ExportConfig) { var ( ok bool @@ -398,15 +411,6 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i return } - sendByte := func(bb *BatchByte) bool { - select { - case ByteChan <- bb: - return true - case <-ctx.Done(): - return false - } - } - symbol := ep.Symbol closeby := ep.userConfig.Fields.EnclosedBy.Value flag := ep.ColumnFlag @@ -482,7 +486,7 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i case types.T_geometry, types.T_geometry32: text, err := planfunction.GeometryPayloadToText(vec.GetBytesAt(i)) if err != nil { - sendByte(&BatchByte{err: err}) + sendExportBatchByte(ctx, ByteChan, &BatchByte{err: err}) bat.Clean(mp) return } @@ -574,7 +578,7 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i } // stop early if downstream already failed - sendByte(&BatchByte{ + sendExportBatchByte(ctx, ByteChan, &BatchByte{ err: moerr.NewInternalErrorf(ctx, "constructByte : unsupported type %d", vec.GetType().Oid), }) bat.Clean(mp) @@ -589,7 +593,7 @@ func constructByte(ctx context.Context, obj FeSession, bat *batch.Batch, index i copy(result, buffer.Bytes()) buffer = nil - if !sendByte(&BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ index: index, writeByte: result, err: nil, @@ -948,7 +952,10 @@ func (ec *ExportConfig) init() { func (ec *ExportConfig) Write(execCtx *ExecCtx, crs *perfcounter.CounterSet, bat *batch.Batch) error { ec.Index.Add(1) - copied, err := bat.Dup(execCtx.ses.GetMemPool()) + // CSV and JSON conversion runs asynchronously and can outlive runner.Run. + // The worker copy therefore belongs to the export pipeline, not to the + // producing statement's allocation account. + copied, err := cloneExportWorkerBatch(bat, execCtx.ses.GetMemPool()) if err != nil { return err } @@ -973,6 +980,13 @@ func (ec *ExportConfig) Write(execCtx *ExecCtx, crs *perfcounter.CounterSet, bat return nil } +func cloneExportWorkerBatch( + bat *batch.Batch, + mp *mpool.MPool, +) (*batch.Batch, error) { + return bat.DupWithoutAllocationAccount(mp) +} + // writeParquet writes a batch to the parquet writer func (ec *ExportConfig) writeParquet(execCtx *ExecCtx, bat *batch.Batch) error { defer bat.Clean(execCtx.ses.GetMemPool()) @@ -1156,8 +1170,11 @@ func constructJSONLine(ctx context.Context, obj FeSession, bat *batch.Batch, ind } val, err := vectorValueToJSON(vec, i, ss, backSes) if err != nil { - ByteChan <- &BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ err: err, + }) { + bat.Clean(mp) + return } bat.Clean(mp) return @@ -1166,8 +1183,11 @@ func constructJSONLine(ctx context.Context, obj FeSession, bat *batch.Batch, ind } jsonBytes, err := json.Marshal(row) if err != nil { - ByteChan <- &BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ err: moerr.NewInternalErrorf(ctx, "failed to marshal JSON: %v", err), + }) { + bat.Clean(mp) + return } bat.Clean(mp) return @@ -1181,10 +1201,13 @@ func constructJSONLine(ctx context.Context, obj FeSession, bat *batch.Batch, ind copy(result, buffer.Bytes()) buffer = nil - ByteChan <- &BatchByte{ + if !sendExportBatchByte(ctx, ByteChan, &BatchByte{ index: index, writeByte: result, err: nil, + }) { + bat.Clean(mp) + return } bat.Clean(mp) diff --git a/pkg/frontend/export_test.go b/pkg/frontend/export_test.go index f2a472fb2f8e0..ab006c0024eab 100644 --- a/pkg/frontend/export_test.go +++ b/pkg/frontend/export_test.go @@ -19,6 +19,7 @@ import ( "context" "strings" "testing" + "time" "github.com/prashantv/gostub" "github.com/smartystreets/goconvey/convey" @@ -27,9 +28,11 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/defines" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" ) var colName1, colName2 = "DATABASE()", "VARIABLE_VALUE" @@ -370,6 +373,74 @@ func TestConstructByteFormatsUnscaledFloat64WithFullPrecision(t *testing.T) { }) } +func TestExportWorkerBatchEndsStatementAllocationOwnership(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + source := batch.NewWithSchema( + true, + []string{"value"}, + []types.Type{types.T_int64.ToType()}, + ) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + source.SetRowCount(1) + require.Positive(t, account.Snapshot().Used) + + workerBatch, err := cloneExportWorkerBatch(source, mp) + require.NoError(t, err) + require.Nil(t, workerBatch.AllocationAccountSelection()) + require.Nil(t, workerBatch.Vecs[0].AllocationAccountSelection()) + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal(t, int64(42), vector.GetFixedAtNoTypeCheck[int64](workerBatch.Vecs[0], 0)) + workerBatch.Clean(mp) + require.Zero(t, mp.CurrNB()) +} + +func TestJSONExportCancellationCleansBlockedWorkerBatch(t *testing.T) { + mp := mpool.MustNewZero() + bat := batch.NewOffHeapWithSize(1) + bat.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(42), false, mp)) + bat.SetRowCount(1) + require.Positive(t, mp.CurrNB()) + + mrs := &MysqlResultSet{} + column := &MysqlColumn{} + column.SetName("value") + mrs.AddColumn(column) + ep := &ExportConfig{mrs: mrs} + bytesChan := make(chan *BatchByte, 1) + bytesChan <- &BatchByte{index: 0} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + ses := &backSession{feSessionImpl: feSessionImpl{pool: mp}} + + done := make(chan struct{}) + go func() { + defer close(done) + constructJSONLine(ctx, ses, bat, 1, bytesChan, ep) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("canceled JSON export worker remained blocked on a full channel") + } + require.Zero(t, mp.CurrNB()) +} + func Test_getExportFormat(t *testing.T) { convey.Convey("getExportFormat returns correct format", t, func() { // Test default format (empty string returns "csv") diff --git a/pkg/frontend/types.go b/pkg/frontend/types.go index 34db16b934b9e..308f03ba6d3cc 100644 --- a/pkg/frontend/types.go +++ b/pkg/frontend/types.go @@ -1359,7 +1359,9 @@ func (ses *feSessionImpl) GetResultBatches() []*batch.Batch { } func (ses *feSessionImpl) AppendResultBatch(bat *batch.Batch) error { - copied, err := bat.Dup(ses.pool) + // Result batches belong to the session and can remain reachable after the + // producing statement has sealed its allocation account. + copied, err := bat.DupWithoutAllocationAccount(ses.pool) if err != nil { return err } diff --git a/pkg/frontend/types_test.go b/pkg/frontend/types_test.go index 8d5c30c300386..395eb1bc89aaa 100644 --- a/pkg/frontend/types_test.go +++ b/pkg/frontend/types_test.go @@ -17,8 +17,11 @@ package frontend import ( "testing" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/config" + "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -392,6 +395,50 @@ func TestPrepareStmt_Close(t *testing.T) { ps.Close() } +func TestAppendResultBatchEndsStatementAllocationOwnership(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + + source := batch.NewOffHeapWithSize(1) + source.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + source.SetRowCount(1) + + ses := &feSessionImpl{pool: mp} + require.NoError(t, ses.AppendResultBatch(source)) + require.Len(t, ses.resultBatches, 1) + require.Nil(t, ses.resultBatches[0].AllocationAccountSelection()) + require.Nil(t, ses.resultBatches[0].Vecs[0].AllocationAccountSelection()) + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal( + t, + int64(42), + vector.GetFixedAtNoTypeCheck[int64](ses.resultBatches[0].Vecs[0], 0), + ) + + ses.ClearResultBatches() + require.Zero(t, mp.CurrNB()) +} + func BenchmarkSessionAllocator(b *testing.B) { allocator := NewSessionAllocator(&config.ParameterUnit{ SV: &config.FrontendParameters{ diff --git a/pkg/pb/pipeline/error_test.go b/pkg/pb/pipeline/error_test.go new file mode 100644 index 0000000000000..279d80ea7278a --- /dev/null +++ b/pkg/pb/pipeline/error_test.go @@ -0,0 +1,50 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package pipeline + +import ( + "context" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +func TestMPoolCapacityErrorPreservesWireIdentity(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 4) + require.NoError(t, err) + account, err := registry.Open(2 << 20) + require.NoError(t, err) + pool, err := mpool.NewMPool("pipeline-mpool-capacity", 1<<20, mpool.NoFixed) + require.NoError(t, err) + defer mpool.DeleteMPool(pool) + + first, err := pool.AllocAccounted(768<<10, account, 1, 1) + require.NoError(t, err) + defer pool.Free(first) + _, capacityErr := pool.AllocAccounted(768<<10, account, 1, 1) + require.Error(t, capacityErr) + require.IsType(t, new(moerr.Error), capacityErr) + require.True(t, mpool.IsMPoolCapacityFailure(capacityErr)) + + message := new(Message) + message.SetMoError(context.Background(), capacityErr) + + wireErr, ok := message.TryToGetMoErr() + require.True(t, ok) + require.True(t, moerr.IsMoErrCode(wireErr, moerr.ErrMPoolCapacity)) + require.Contains(t, wireErr.Error(), "allocation owner=1 site=1") +} diff --git a/pkg/pb/pipeline/pipeline.pb.go b/pkg/pb/pipeline/pipeline.pb.go index f06d2144764d6..9025de3bb79f6 100644 --- a/pkg/pb/pipeline/pipeline.pb.go +++ b/pkg/pb/pipeline/pipeline.pb.go @@ -5363,19 +5363,25 @@ func (m *PrepareParamInfo) GetIsBin() []bool { } type ProcessInfo struct { - Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` - Sql string `protobuf:"bytes,2,opt,name=sql,proto3" json:"sql,omitempty"` - Lim ProcessLimitation `protobuf:"bytes,3,opt,name=lim,proto3" json:"lim"` - UnixTime int64 `protobuf:"varint,4,opt,name=unix_time,json=unixTime,proto3" json:"unix_time,omitempty"` - AccountId uint32 `protobuf:"varint,5,opt,name=account_id,json=accountId,proto3" json:"account_id,omitempty"` - Snapshot txn.CNTxnSnapshot `protobuf:"bytes,6,opt,name=snapshot,proto3" json:"snapshot"` - SessionInfo SessionInfo `protobuf:"bytes,7,opt,name=session_info,json=sessionInfo,proto3" json:"session_info"` - SessionLogger SessionLoggerInfo `protobuf:"bytes,8,opt,name=session_logger,json=sessionLogger,proto3" json:"session_logger"` - PrepareParams PrepareParamInfo `protobuf:"bytes,9,opt,name=prepare_params,json=prepareParams,proto3" json:"prepare_params"` - AffectedRows int64 `protobuf:"varint,10,opt,name=affected_rows,json=affectedRows,proto3" json:"affected_rows,omitempty"` - XXX_NoUnkeyedLiteral struct{} `json:"-"` - XXX_unrecognized []byte `json:"-"` - XXX_sizecache int32 `json:"-"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + Sql string `protobuf:"bytes,2,opt,name=sql,proto3" json:"sql,omitempty"` + Lim ProcessLimitation `protobuf:"bytes,3,opt,name=lim,proto3" json:"lim"` + UnixTime int64 `protobuf:"varint,4,opt,name=unix_time,json=unixTime,proto3" json:"unix_time,omitempty"` + AccountId uint32 `protobuf:"varint,5,opt,name=account_id,json=accountId,proto3" json:"account_id,omitempty"` + Snapshot txn.CNTxnSnapshot `protobuf:"bytes,6,opt,name=snapshot,proto3" json:"snapshot"` + SessionInfo SessionInfo `protobuf:"bytes,7,opt,name=session_info,json=sessionInfo,proto3" json:"session_info"` + SessionLogger SessionLoggerInfo `protobuf:"bytes,8,opt,name=session_logger,json=sessionLogger,proto3" json:"session_logger"` + PrepareParams PrepareParamInfo `protobuf:"bytes,9,opt,name=prepare_params,json=prepareParams,proto3" json:"prepare_params"` + AffectedRows int64 `protobuf:"varint,10,opt,name=affected_rows,json=affectedRows,proto3" json:"affected_rows,omitempty"` + // Planned PipelineMessage RPCs per target CN for statement-level remote + // MessageBoard and resource-account terminal ownership. + RemoteFragmentCounts map[string]uint32 `protobuf:"bytes,11,rep,name=remote_fragment_counts,json=remoteFragmentCounts,proto3" json:"remote_fragment_counts,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"varint,2,opt,name=value,proto3"` + // Unique physical execution attempt. Unlike the SQL statement ID, this + // changes across retries and prepared-statement executions. + RemoteExecutionId []byte `protobuf:"bytes,12,opt,name=remote_execution_id,json=remoteExecutionId,proto3" json:"remote_execution_id,omitempty"` + XXX_NoUnkeyedLiteral struct{} `json:"-"` + XXX_unrecognized []byte `json:"-"` + XXX_sizecache int32 `json:"-"` } func (m *ProcessInfo) Reset() { *m = ProcessInfo{} } @@ -5481,6 +5487,20 @@ func (m *ProcessInfo) GetAffectedRows() int64 { return 0 } +func (m *ProcessInfo) GetRemoteFragmentCounts() map[string]uint32 { + if m != nil { + return m.RemoteFragmentCounts + } + return nil +} + +func (m *ProcessInfo) GetRemoteExecutionId() []byte { + if m != nil { + return m.RemoteExecutionId + } + return nil +} + type SessionInfo struct { User string `protobuf:"bytes,1,opt,name=user,proto3" json:"user,omitempty"` Host string `protobuf:"bytes,2,opt,name=host,proto3" json:"host,omitempty"` @@ -6103,6 +6123,7 @@ func init() { proto.RegisterType((*ProcessLimitation)(nil), "pipeline.ProcessLimitation") proto.RegisterType((*PrepareParamInfo)(nil), "pipeline.PrepareParamInfo") proto.RegisterType((*ProcessInfo)(nil), "pipeline.ProcessInfo") + proto.RegisterMapType((map[string]uint32)(nil), "pipeline.ProcessInfo.RemoteFragmentCountsEntry") proto.RegisterType((*SessionInfo)(nil), "pipeline.SessionInfo") proto.RegisterType((*SessionLoggerInfo)(nil), "pipeline.SessionLoggerInfo") proto.RegisterType((*Pipeline)(nil), "pipeline.Pipeline") @@ -6114,481 +6135,486 @@ func init() { func init() { proto.RegisterFile("pipeline.proto", fileDescriptor_7ac67a7adf3df9c7) } var fileDescriptor_7ac67a7adf3df9c7 = []byte{ - // 7584 bytes of a gzipped FileDescriptorProto + // 7656 bytes of a gzipped FileDescriptorProto 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0xec, 0x7c, 0x4d, 0x6f, 0x1c, 0xc9, - 0x92, 0xd8, 0x34, 0xfb, 0x3b, 0xfa, 0x83, 0xcd, 0x24, 0x45, 0xb5, 0xa4, 0x79, 0x23, 0x4d, 0xcf, - 0x48, 0xc3, 0xa7, 0xd1, 0x50, 0x12, 0x67, 0xe6, 0xbd, 0xd9, 0xf7, 0xf6, 0xed, 0x5b, 0x8a, 0x92, - 0xde, 0xf0, 0x3d, 0x51, 0xe2, 0x16, 0x29, 0x0f, 0x30, 0x80, 0x5d, 0x28, 0x56, 0x65, 0x77, 0xd7, - 0xb0, 0xba, 0xb2, 0x54, 0x99, 0x25, 0x91, 0xba, 0xd8, 0x07, 0x9f, 0x7c, 0xf1, 0x71, 0xf7, 0xb8, - 0x80, 0x7d, 0x58, 0xdb, 0x07, 0x1f, 0x8c, 0xe7, 0x9f, 0x60, 0x2c, 0x6c, 0xc3, 0x58, 0xf8, 0xe0, - 0xa3, 0x61, 0xbc, 0x3d, 0x1a, 0x30, 0x0c, 0x03, 0x36, 0x16, 0x30, 0x0c, 0x18, 0x11, 0x99, 0x59, - 0x55, 0xdd, 0x4d, 0x69, 0x3e, 0x6c, 0xec, 0x65, 0xf7, 0xd4, 0x95, 0x11, 0x91, 0x59, 0x59, 0x91, - 0x91, 0x11, 0x91, 0x11, 0x91, 0x0d, 0xfd, 0x24, 0x4c, 0x78, 0x14, 0xc6, 0x7c, 0x3b, 0x49, 0x85, - 0x12, 0xac, 0x65, 0xdb, 0x57, 0x3f, 0x99, 0x84, 0x6a, 0x9a, 0x9d, 0x6c, 0xfb, 0x62, 0x76, 0x77, - 0x22, 0x26, 0xe2, 0x2e, 0x11, 0x9c, 0x64, 0x63, 0x6a, 0x51, 0x83, 0x9e, 0x74, 0xc7, 0xab, 0x10, - 0x09, 0xff, 0xd4, 0x3e, 0x27, 0x91, 0x17, 0x9b, 0xe7, 0x55, 0x15, 0xce, 0xb8, 0x54, 0xde, 0x2c, - 0x31, 0x80, 0xb6, 0x3a, 0x33, 0xb8, 0xd1, 0xbf, 0xab, 0x42, 0xf3, 0x80, 0x4b, 0xe9, 0x4d, 0x38, - 0x1b, 0x41, 0x55, 0x86, 0xc1, 0xb0, 0x72, 0xa3, 0xb2, 0xd5, 0xdf, 0x19, 0x6c, 0xe7, 0xd3, 0x3a, - 0x52, 0x9e, 0xca, 0xa4, 0x83, 0x48, 0xa4, 0xf1, 0x67, 0xc1, 0x70, 0x65, 0x91, 0xe6, 0x80, 0xab, - 0xa9, 0x08, 0x1c, 0x44, 0xb2, 0x01, 0x54, 0x79, 0x9a, 0x0e, 0xab, 0x37, 0x2a, 0x5b, 0x5d, 0x07, - 0x1f, 0x19, 0x83, 0x5a, 0xe0, 0x29, 0x6f, 0x58, 0x23, 0x10, 0x3d, 0xb3, 0x0f, 0xa1, 0x9f, 0xa4, - 0xc2, 0x77, 0xc3, 0x78, 0x2c, 0x5c, 0xc2, 0xd6, 0x09, 0xdb, 0x45, 0xe8, 0x7e, 0x3c, 0x16, 0x0f, - 0x91, 0x6a, 0x08, 0x4d, 0x2f, 0xf6, 0xa2, 0x73, 0xc9, 0x87, 0x0d, 0x42, 0xdb, 0x26, 0xeb, 0xc3, - 0x4a, 0x18, 0x0c, 0x9b, 0x37, 0x2a, 0x5b, 0x35, 0x67, 0x25, 0x0c, 0xf0, 0x1d, 0x59, 0x16, 0x06, - 0xc3, 0x96, 0x7e, 0x07, 0x3e, 0xb3, 0x11, 0x74, 0x63, 0xce, 0x83, 0xa7, 0x42, 0x39, 0x3c, 0x89, - 0xce, 0x87, 0xed, 0x1b, 0x95, 0xad, 0x96, 0x33, 0x07, 0x63, 0x57, 0xa1, 0x15, 0xf0, 0x93, 0x6c, - 0x72, 0x20, 0x27, 0x43, 0xb8, 0x51, 0xd9, 0x6a, 0x3b, 0x79, 0x9b, 0x1d, 0xc3, 0xe5, 0x94, 0xbf, - 0xc8, 0xb8, 0x54, 0x3c, 0x70, 0x15, 0xf7, 0xd2, 0x40, 0xbc, 0x8a, 0xdd, 0x99, 0x08, 0xf8, 0xb0, - 0x43, 0x1c, 0x78, 0xb7, 0xcc, 0xa5, 0x94, 0x7b, 0xb3, 0x63, 0x43, 0x74, 0x20, 0x02, 0xee, 0x5c, - 0xca, 0x3b, 0x97, 0xc1, 0xcc, 0x81, 0x4d, 0xcf, 0xf7, 0x79, 0xb2, 0x3c, 0x68, 0xf7, 0x3b, 0x0c, - 0xba, 0x61, 0xfb, 0x96, 0xa1, 0x3f, 0xab, 0xfd, 0xc9, 0x9f, 0x5e, 0x7f, 0x67, 0xf4, 0x1c, 0xda, - 0x7b, 0x22, 0x8e, 0xb9, 0xaf, 0x44, 0xca, 0xae, 0x43, 0xc7, 0x8e, 0xe3, 0x9a, 0x65, 0xad, 0x3b, - 0x60, 0x41, 0xfb, 0x01, 0xfb, 0x08, 0x56, 0x7d, 0x4b, 0xed, 0x86, 0x71, 0xc0, 0xcf, 0x68, 0x5d, - 0xeb, 0x4e, 0x3f, 0x07, 0xef, 0x23, 0x74, 0xf4, 0x6f, 0xaa, 0xd0, 0x3c, 0x9a, 0x66, 0xe3, 0x71, - 0xc4, 0xd9, 0x87, 0xd0, 0x33, 0x8f, 0x7b, 0x22, 0xda, 0x0f, 0xce, 0xcc, 0xb8, 0xf3, 0x40, 0x76, - 0x03, 0x3a, 0x06, 0x70, 0x7c, 0x9e, 0x70, 0x33, 0x6c, 0x19, 0x34, 0x3f, 0xce, 0x41, 0x18, 0x93, - 0xb8, 0x54, 0x9d, 0x79, 0xe0, 0x02, 0x95, 0x77, 0x46, 0x12, 0x34, 0x4f, 0xe5, 0xd1, 0xdb, 0x76, - 0xa3, 0xf0, 0x25, 0x77, 0xf8, 0x64, 0x2f, 0x56, 0x24, 0x47, 0x75, 0xa7, 0x0c, 0x62, 0x3b, 0x70, - 0x49, 0xea, 0x2e, 0x6e, 0xea, 0xc5, 0x13, 0x2e, 0xdd, 0x2c, 0x8c, 0xd5, 0x4f, 0x3e, 0x1b, 0x36, - 0x6e, 0x54, 0xb7, 0x6a, 0xce, 0xba, 0x41, 0x3a, 0x84, 0x7b, 0x4e, 0x28, 0x76, 0x0f, 0x36, 0x16, - 0xfa, 0xe8, 0x2e, 0xcd, 0x1b, 0xd5, 0xad, 0xaa, 0xc3, 0xe6, 0xba, 0xec, 0x53, 0x8f, 0x47, 0xb0, - 0x96, 0x66, 0x31, 0xee, 0xb6, 0xc7, 0x61, 0xa4, 0x78, 0x7a, 0x94, 0x70, 0x9f, 0xe4, 0xb1, 0xb3, - 0x73, 0x79, 0x9b, 0x36, 0xa4, 0xb3, 0x88, 0x76, 0x96, 0x7b, 0xb0, 0x3b, 0x39, 0xf3, 0x1e, 0x9d, - 0x25, 0x29, 0x09, 0x6d, 0x67, 0x07, 0xf4, 0x00, 0x08, 0x71, 0xca, 0x68, 0x76, 0x1b, 0xd6, 0x82, - 0xd4, 0x0b, 0x63, 0xd7, 0x8b, 0x22, 0xf7, 0x24, 0xf3, 0x4f, 0xb9, 0x92, 0x24, 0xc8, 0x2d, 0x67, - 0x95, 0x10, 0xbb, 0x51, 0xf4, 0x40, 0x83, 0x47, 0x7f, 0xb5, 0x02, 0xad, 0x87, 0xa1, 0x4c, 0x3c, - 0xe5, 0x4f, 0xd9, 0x65, 0x68, 0x8e, 0xb3, 0xd8, 0x2f, 0x64, 0xa3, 0x81, 0xcd, 0xfd, 0x80, 0xfd, - 0x3e, 0xac, 0x46, 0xc2, 0xf7, 0x22, 0x37, 0x17, 0x83, 0xe1, 0xca, 0x8d, 0xea, 0x56, 0x67, 0x67, - 0xbd, 0x10, 0xcc, 0x5c, 0xcc, 0x9c, 0x3e, 0xd1, 0x16, 0x62, 0xf7, 0x0b, 0x18, 0xa4, 0x7c, 0x26, - 0x14, 0x2f, 0x75, 0xaf, 0x52, 0x77, 0x56, 0x74, 0xff, 0x2a, 0xf5, 0x92, 0xa7, 0x28, 0xcd, 0xab, - 0x9a, 0xb6, 0xe8, 0x7e, 0xbf, 0xb4, 0x52, 0x7c, 0xe2, 0x86, 0xc1, 0x99, 0x4b, 0x2f, 0x18, 0xd6, - 0x6e, 0x54, 0xb7, 0xea, 0x05, 0xdb, 0xf9, 0x64, 0x3f, 0x38, 0x7b, 0x82, 0x18, 0xf6, 0x29, 0x6c, - 0x2e, 0x76, 0xd1, 0xa3, 0x0e, 0xeb, 0xd4, 0x67, 0x7d, 0xae, 0x8f, 0x43, 0x28, 0xf6, 0x3e, 0x74, - 0x6d, 0x27, 0x85, 0x22, 0xda, 0xd0, 0x42, 0x23, 0x4b, 0x22, 0x7a, 0x19, 0x9a, 0xa1, 0x74, 0x65, - 0x18, 0x9f, 0x92, 0x9a, 0x69, 0x39, 0x8d, 0x50, 0x1e, 0x85, 0xf1, 0x29, 0xbb, 0x02, 0xad, 0x94, - 0xfb, 0x1a, 0xd3, 0x22, 0x4c, 0x33, 0xe5, 0x3e, 0xa1, 0x2e, 0x03, 0x3e, 0xba, 0xbe, 0xe2, 0x46, - 0xd9, 0x34, 0x52, 0xee, 0xef, 0x29, 0x3e, 0x92, 0x50, 0x3f, 0xe0, 0xe9, 0x84, 0xa3, 0xbe, 0xc1, - 0x8e, 0x47, 0xbe, 0x17, 0x13, 0xdf, 0x5b, 0x4e, 0xde, 0x46, 0x6d, 0x97, 0x78, 0xa9, 0x0a, 0xbd, - 0x88, 0xb6, 0x4c, 0xcb, 0xb1, 0x4d, 0x76, 0x0d, 0xda, 0x52, 0x79, 0xa9, 0xc2, 0xaf, 0xa3, 0xad, - 0x52, 0x77, 0x5a, 0x04, 0xc0, 0xdd, 0x76, 0x19, 0x9a, 0x3c, 0x0e, 0x08, 0x55, 0xd3, 0x2b, 0xc9, - 0xe3, 0x60, 0x3f, 0x38, 0x1b, 0xfd, 0xab, 0x0a, 0xf4, 0x0e, 0xb2, 0x48, 0x85, 0xbb, 0xe9, 0x24, - 0xe3, 0xb3, 0x58, 0xa1, 0x96, 0x7c, 0x18, 0x4a, 0x65, 0xde, 0x4c, 0xcf, 0x6c, 0x0b, 0xda, 0xbf, - 0x4a, 0x45, 0x96, 0x90, 0xb4, 0xe9, 0x95, 0x2e, 0x4b, 0x5b, 0x81, 0x44, 0xc9, 0x7c, 0x96, 0x06, - 0x3c, 0x7d, 0x70, 0x4e, 0xb4, 0xd5, 0x25, 0xda, 0x32, 0x9a, 0xbd, 0x0b, 0xed, 0x23, 0x9e, 0x78, - 0xa9, 0x87, 0x22, 0x50, 0x23, 0xd5, 0x5a, 0x00, 0xf0, 0x5b, 0x89, 0x78, 0x3f, 0x30, 0x1b, 0xd6, - 0x36, 0x47, 0xff, 0xa4, 0x02, 0xed, 0xdd, 0xc9, 0x24, 0xe5, 0x13, 0x4f, 0x91, 0x9e, 0x17, 0x09, - 0xcd, 0xb7, 0xea, 0xac, 0x88, 0x84, 0x6c, 0x09, 0x7e, 0x81, 0x66, 0x10, 0x3d, 0xb3, 0xf7, 0xa0, - 0xc6, 0x2f, 0x9e, 0x10, 0xc1, 0xd9, 0x26, 0x34, 0x7c, 0x11, 0x8f, 0xc3, 0x89, 0xb1, 0x40, 0xa6, - 0xc5, 0x7e, 0x06, 0x1d, 0xfd, 0xa4, 0x65, 0xa0, 0x4e, 0xea, 0xf7, 0x8a, 0xee, 0x9e, 0xcf, 0x60, - 0x8f, 0x28, 0x50, 0x22, 0x1c, 0xf0, 0xf3, 0xe7, 0xd1, 0x3f, 0xaf, 0x42, 0x9d, 0x38, 0x83, 0x6b, - 0x83, 0x16, 0xc5, 0xe5, 0x2f, 0xbd, 0xc8, 0x2e, 0x29, 0x02, 0x1e, 0xbd, 0xf4, 0x22, 0x76, 0x03, - 0xea, 0x38, 0x05, 0x79, 0x01, 0x63, 0x35, 0x82, 0xdd, 0x82, 0x3a, 0xbe, 0x5d, 0xce, 0xcf, 0x1e, - 0xdf, 0xf1, 0xa0, 0xf6, 0xe7, 0xff, 0xf9, 0xfa, 0x3b, 0x8e, 0x46, 0xb3, 0x8f, 0xa0, 0xe6, 0x4d, - 0x26, 0x92, 0x36, 0xc2, 0xdc, 0x5e, 0xcc, 0x67, 0xea, 0x10, 0x01, 0xfb, 0x1c, 0xda, 0x7a, 0xd1, - 0x91, 0xba, 0x4e, 0xd4, 0x97, 0x4b, 0x96, 0xba, 0x2c, 0x0f, 0x4e, 0x41, 0x89, 0xcb, 0x15, 0x4a, - 0xa3, 0x59, 0x68, 0x3b, 0xb4, 0x9c, 0x02, 0x80, 0xa6, 0x34, 0x49, 0xf9, 0x6e, 0x14, 0x09, 0xff, - 0x28, 0x7c, 0xcd, 0x8d, 0xe1, 0x9d, 0x83, 0xb1, 0x5b, 0xd0, 0x3f, 0xd4, 0xf2, 0xea, 0x70, 0x99, - 0x45, 0x4a, 0x1a, 0x63, 0xbc, 0x00, 0x65, 0xdb, 0xc0, 0xe6, 0x20, 0xc7, 0xf4, 0xf9, 0xed, 0x1b, - 0xd5, 0xad, 0x9e, 0x73, 0x01, 0x86, 0x7d, 0x00, 0xbd, 0x09, 0x72, 0x3a, 0x8c, 0x27, 0xee, 0x38, - 0xf2, 0xd0, 0x4e, 0x57, 0xd1, 0x8e, 0x5b, 0xe0, 0xe3, 0xc8, 0x9b, 0xd0, 0x0e, 0x49, 0xc2, 0x28, - 0x72, 0x67, 0x7c, 0x46, 0xd6, 0xb9, 0xea, 0xb4, 0x08, 0x70, 0xc0, 0x67, 0xa3, 0x7f, 0x51, 0x83, - 0xc6, 0x7e, 0x2c, 0x79, 0xaa, 0x70, 0xff, 0x79, 0xe3, 0x31, 0xf7, 0x15, 0xd7, 0x7a, 0xaf, 0xe6, - 0xe4, 0x6d, 0x64, 0xc1, 0xb1, 0xf8, 0x2a, 0x0d, 0x15, 0x3f, 0xfa, 0xd4, 0x08, 0x58, 0x01, 0x40, - 0x4d, 0xeb, 0x05, 0x81, 0x6b, 0xa9, 0xdd, 0x54, 0xbc, 0x92, 0xb4, 0x17, 0x5b, 0xce, 0xaa, 0x17, - 0x04, 0xbb, 0x06, 0xee, 0x88, 0x57, 0x92, 0xbd, 0x0f, 0xd5, 0x94, 0x8f, 0x49, 0xdc, 0x3a, 0x3b, - 0xab, 0x7a, 0x49, 0x9f, 0x9d, 0x7c, 0xc3, 0x7d, 0xe5, 0xf0, 0xb1, 0x83, 0x38, 0xb6, 0x01, 0x75, - 0x4f, 0xa9, 0x54, 0x2f, 0x51, 0xdb, 0xd1, 0x0d, 0xb6, 0x0d, 0xeb, 0xb4, 0xe7, 0x55, 0x28, 0x62, - 0x57, 0x79, 0x27, 0x11, 0x1a, 0x6f, 0x69, 0xec, 0xd4, 0x5a, 0x8e, 0x3a, 0x46, 0xcc, 0x7e, 0x20, - 0xd1, 0xb2, 0x2d, 0xd2, 0xc7, 0xde, 0x8c, 0x4b, 0x32, 0x53, 0x6d, 0x67, 0x7d, 0xbe, 0xc7, 0x53, - 0x44, 0x21, 0x3f, 0x8b, 0x3e, 0xa8, 0x35, 0x5a, 0xb4, 0x01, 0xbb, 0x39, 0x10, 0x95, 0xca, 0x25, - 0x68, 0x84, 0xd2, 0xe5, 0x71, 0x60, 0x14, 0x59, 0x3d, 0x94, 0x8f, 0xe2, 0x80, 0x7d, 0x0c, 0x6d, - 0xfd, 0x96, 0x80, 0x8f, 0xc9, 0xcc, 0x74, 0x76, 0xfa, 0x46, 0x62, 0x11, 0xfc, 0x90, 0x8f, 0x9d, - 0x96, 0x32, 0x4f, 0xe8, 0x82, 0x28, 0xe1, 0xf2, 0x33, 0xc5, 0xd3, 0xd8, 0x8b, 0x68, 0x55, 0x5a, - 0x0e, 0x28, 0xf1, 0xc8, 0x40, 0xd8, 0xe7, 0x70, 0xd9, 0x62, 0x5d, 0xa9, 0x66, 0xca, 0xcd, 0xe2, - 0xf0, 0xcc, 0x8d, 0xbd, 0x58, 0x90, 0x2f, 0x54, 0x75, 0x36, 0x2c, 0xfa, 0x48, 0xcd, 0xd4, 0xf3, - 0x38, 0x3c, 0x7b, 0xea, 0xc5, 0x82, 0x6d, 0xc1, 0x20, 0xef, 0xa6, 0x5e, 0xd3, 0x07, 0x0f, 0x7b, - 0xa4, 0x60, 0xfa, 0x16, 0x7e, 0xfc, 0x1a, 0xbf, 0x15, 0x6d, 0x43, 0x99, 0x52, 0x8c, 0xc7, 0x92, - 0x2b, 0x57, 0x72, 0x7f, 0xd8, 0xa7, 0x6f, 0x5e, 0x2f, 0xe8, 0x9f, 0x11, 0xee, 0x88, 0xfb, 0xa3, - 0xdf, 0x56, 0xa0, 0x43, 0xfb, 0xe2, 0x79, 0x12, 0xa0, 0x0a, 0xfa, 0x00, 0x7a, 0xf3, 0x8b, 0xae, - 0xe5, 0xa6, 0xeb, 0x95, 0x57, 0x7c, 0x13, 0x1a, 0xbb, 0x3e, 0x32, 0x8f, 0x04, 0xa7, 0xe7, 0x98, - 0x16, 0xfb, 0x29, 0xac, 0x66, 0x34, 0x8c, 0xeb, 0xab, 0x33, 0x37, 0x42, 0xd5, 0xa5, 0x37, 0xba, - 0x91, 0x0a, 0xfd, 0x8e, 0x3d, 0x75, 0xe6, 0xf4, 0x32, 0xfb, 0xf8, 0x04, 0x95, 0xda, 0x3d, 0xd8, - 0x48, 0x39, 0x4a, 0x8c, 0xfb, 0x9a, 0xa7, 0xc2, 0x55, 0x7c, 0x96, 0x88, 0x94, 0x0c, 0x21, 0x72, - 0x91, 0x69, 0xdc, 0xd7, 0x3c, 0x15, 0xc7, 0x06, 0x33, 0xfa, 0x11, 0xd4, 0x77, 0xd3, 0xd4, 0x3b, - 0x27, 0xd1, 0xc2, 0x87, 0x61, 0x85, 0x0c, 0xa0, 0x6e, 0x8c, 0x7c, 0xa8, 0x1e, 0x78, 0x09, 0xbb, - 0x09, 0x2b, 0xb3, 0x84, 0x30, 0x9d, 0x9d, 0x4b, 0x25, 0xbd, 0xe0, 0x25, 0xdb, 0x07, 0xc9, 0xa3, - 0x58, 0xa5, 0xe7, 0xce, 0xca, 0x2c, 0xb9, 0xfa, 0x39, 0x34, 0x4d, 0x13, 0x1d, 0xfa, 0x53, 0x7e, - 0x4e, 0x5f, 0xdd, 0x76, 0xf0, 0x11, 0x5f, 0xf0, 0xd2, 0x8b, 0x32, 0xeb, 0xd9, 0xe9, 0xc6, 0xcf, - 0x56, 0xbe, 0xa8, 0x8c, 0xfe, 0x67, 0x0d, 0x5a, 0x0f, 0x79, 0xc4, 0xe9, 0xdb, 0x47, 0xd0, 0x2d, - 0xef, 0x0a, 0xcb, 0xb7, 0xb9, 0x9d, 0x32, 0x82, 0xae, 0x36, 0xc9, 0xd4, 0x8b, 0x9b, 0x6d, 0x37, - 0x07, 0x43, 0x5b, 0xb1, 0xaf, 0x7d, 0x18, 0xda, 0x6f, 0x3d, 0xc7, 0x36, 0x11, 0xf3, 0xd4, 0x60, - 0x6a, 0x1a, 0x63, 0x9a, 0xec, 0x5d, 0x80, 0x54, 0xbc, 0x72, 0x43, 0x6d, 0x17, 0xb5, 0x89, 0x69, - 0xa5, 0xe2, 0xd5, 0x3e, 0x5a, 0xc6, 0xbf, 0x96, 0x6d, 0xf6, 0x53, 0x18, 0x96, 0xb6, 0x19, 0x7a, - 0xd2, 0x6e, 0x18, 0xbb, 0x27, 0xe8, 0x7c, 0x99, 0x1d, 0x57, 0x8c, 0x49, 0x8e, 0xf6, 0x7e, 0xfc, - 0x80, 0x3c, 0x33, 0xa3, 0x3c, 0xda, 0x6f, 0x51, 0x1e, 0x17, 0xea, 0x22, 0xb8, 0x58, 0x17, 0x3d, - 0x00, 0x38, 0xe2, 0x93, 0x19, 0x8f, 0xd5, 0x81, 0x97, 0x0c, 0x3b, 0xb4, 0xf0, 0xa3, 0x62, 0xe1, - 0xed, 0x6a, 0x6d, 0x17, 0x44, 0x5a, 0x0a, 0x4a, 0xbd, 0xd0, 0x5d, 0xf2, 0xbd, 0xd8, 0x55, 0x69, - 0x16, 0xfb, 0x9e, 0xd2, 0x27, 0x95, 0x96, 0xd3, 0xf1, 0xbd, 0xf8, 0xd8, 0x80, 0x4a, 0x0a, 0xa3, - 0x57, 0x56, 0x18, 0xb7, 0x60, 0x35, 0x49, 0xc3, 0x99, 0x97, 0x9e, 0xbb, 0xa7, 0xfc, 0x9c, 0x16, - 0x43, 0x6f, 0xbd, 0x9e, 0x01, 0xff, 0x86, 0x9f, 0xef, 0x07, 0x67, 0x57, 0x7f, 0x01, 0xab, 0x0b, - 0x13, 0xf8, 0x5e, 0x72, 0xf7, 0x1f, 0xab, 0xd0, 0x3e, 0x4c, 0xb9, 0x51, 0xf2, 0xd7, 0xa1, 0x23, - 0xfd, 0x29, 0x9f, 0x79, 0x5a, 0x37, 0xe8, 0x11, 0x40, 0x83, 0x48, 0x2f, 0xcc, 0xa9, 0xb1, 0x95, - 0x6f, 0x51, 0x63, 0x03, 0xa8, 0x6a, 0xb7, 0x0b, 0x37, 0x13, 0x3e, 0x16, 0xba, 0xbb, 0x56, 0xd6, - 0xdd, 0x37, 0xa0, 0x3b, 0xf5, 0xa4, 0xeb, 0x65, 0x4a, 0xb8, 0xbe, 0x88, 0x48, 0xe8, 0x5a, 0x0e, - 0x4c, 0x3d, 0xb9, 0x9b, 0x29, 0xb1, 0x27, 0x22, 0xf6, 0x23, 0x00, 0x5f, 0x44, 0x46, 0x0d, 0x19, - 0x9f, 0xb3, 0xed, 0x8b, 0x48, 0xeb, 0x1e, 0x94, 0x4a, 0x2e, 0x55, 0x38, 0xf3, 0xcc, 0x92, 0xba, - 0xbe, 0xc8, 0x62, 0x45, 0xb6, 0xb6, 0xea, 0xac, 0xe5, 0x28, 0x47, 0xbc, 0xda, 0x43, 0x04, 0xbb, - 0x07, 0x7d, 0x5f, 0xcc, 0x12, 0x37, 0x41, 0xce, 0x92, 0x07, 0xd4, 0x5a, 0x3a, 0x2c, 0x74, 0x91, - 0xe2, 0xf0, 0x94, 0x6b, 0x9f, 0x6c, 0x07, 0x56, 0xfd, 0x28, 0x93, 0x8a, 0xa7, 0xee, 0x89, 0xe9, - 0xb2, 0x7c, 0xbe, 0xe8, 0x19, 0x12, 0xe3, 0xc7, 0x8d, 0xa0, 0x17, 0x4a, 0x57, 0x44, 0x81, 0xab, - 0x15, 0x94, 0x91, 0xb3, 0x4e, 0x28, 0x9f, 0x45, 0x81, 0x51, 0x91, 0x9a, 0x26, 0xe6, 0xaf, 0x2c, - 0x4d, 0xc7, 0xd2, 0x3c, 0xe5, 0xaf, 0x0c, 0xcd, 0x9b, 0x14, 0x5a, 0xf7, 0x8d, 0x0a, 0xed, 0x3f, - 0xad, 0x40, 0xf3, 0x50, 0x48, 0xf5, 0x70, 0x16, 0xd9, 0x4d, 0x51, 0xf9, 0xbe, 0x9b, 0x62, 0xe5, - 0xe2, 0x4d, 0x71, 0x81, 0x58, 0x56, 0x2f, 0x10, 0x4b, 0x34, 0x35, 0x65, 0x3a, 0x12, 0x27, 0xed, - 0xcb, 0xf6, 0x0b, 0x42, 0x12, 0xa9, 0x6b, 0xe8, 0x3f, 0xb9, 0x81, 0xd6, 0x62, 0x7a, 0xe9, 0x5b, - 0xa1, 0x34, 0x1a, 0x4c, 0x23, 0x43, 0x92, 0x4e, 0xe3, 0x5c, 0xb5, 0x42, 0x69, 0xa4, 0xf5, 0xf7, - 0xe0, 0x4a, 0xde, 0xd3, 0x7d, 0x15, 0xaa, 0xa9, 0xc8, 0x94, 0x3b, 0xa6, 0x03, 0xa1, 0x34, 0x47, - 0x8f, 0x4d, 0x3b, 0xd2, 0x57, 0x1a, 0xad, 0x8f, 0x8b, 0xe4, 0xeb, 0x8d, 0xb3, 0x28, 0x72, 0x15, - 0x3f, 0x53, 0x66, 0xf1, 0x87, 0x9a, 0x37, 0x86, 0x6f, 0x8f, 0xb3, 0x28, 0x3a, 0xe6, 0x67, 0x0a, - 0x0d, 0x4c, 0x6b, 0x6c, 0x1a, 0xa3, 0x3f, 0xae, 0x01, 0x3c, 0x11, 0xfe, 0xe9, 0xb1, 0x97, 0x4e, - 0xb8, 0xc2, 0x03, 0x8d, 0xd5, 0x81, 0x46, 0x47, 0x37, 0x95, 0xd6, 0x7c, 0x6c, 0x07, 0x36, 0xed, - 0xf7, 0xa3, 0xe4, 0xe2, 0xe1, 0x4a, 0x2b, 0x31, 0xb3, 0x05, 0x99, 0xc1, 0xea, 0x83, 0x3f, 0x69, - 0x30, 0xf6, 0x45, 0xc1, 0x5b, 0xec, 0xa3, 0xce, 0x13, 0xe2, 0xed, 0x45, 0xbe, 0x6d, 0xaf, 0xe8, - 0x7e, 0x7c, 0x9e, 0xb0, 0x7b, 0x70, 0x29, 0xe5, 0xe3, 0x94, 0xcb, 0xa9, 0xab, 0x64, 0xf9, 0x65, - 0xfa, 0x5c, 0xb3, 0x66, 0x90, 0xc7, 0x32, 0x7f, 0xd7, 0x3d, 0xb8, 0xa4, 0x39, 0xb5, 0x38, 0x3d, - 0xad, 0xf1, 0xd7, 0x34, 0xb2, 0x3c, 0xbb, 0x1f, 0x01, 0x05, 0xca, 0xb4, 0x16, 0xb7, 0x8e, 0x6e, - 0x44, 0xcc, 0x38, 0x89, 0x38, 0xfa, 0x80, 0x7b, 0x53, 0x3c, 0xd4, 0x3f, 0xe4, 0x63, 0xc3, 0xfc, - 0x02, 0xc0, 0x46, 0x50, 0x3b, 0x10, 0x01, 0x27, 0x56, 0xf7, 0x77, 0xfa, 0xdb, 0x14, 0x72, 0x43, - 0x4e, 0x52, 0x6c, 0x86, 0x70, 0xec, 0x23, 0xa0, 0xe1, 0xb4, 0xf8, 0x2d, 0xef, 0xae, 0x16, 0x22, - 0x49, 0x06, 0xef, 0xc1, 0xa5, 0x62, 0x26, 0xae, 0xa7, 0x5c, 0x35, 0xe5, 0xa4, 0x40, 0xf5, 0x06, - 0x5b, 0xcb, 0x27, 0xb5, 0xab, 0x8e, 0xa7, 0x1c, 0x95, 0xe9, 0x16, 0x34, 0xc5, 0xc9, 0x37, 0x2e, - 0x6e, 0x84, 0xce, 0xc5, 0x1b, 0xa1, 0x21, 0x4e, 0xbe, 0x71, 0xf8, 0x98, 0xfd, 0xa4, 0x6c, 0x7c, - 0x16, 0x58, 0xd3, 0x25, 0xd6, 0x6c, 0xe4, 0xf8, 0x12, 0x77, 0x46, 0x5f, 0x40, 0x03, 0x3f, 0xe7, - 0x59, 0xc2, 0xb6, 0xa1, 0xa9, 0x48, 0x3c, 0xa4, 0x71, 0x16, 0x36, 0x0a, 0x9b, 0x51, 0xc8, 0x8e, - 0x63, 0x89, 0x46, 0x0e, 0xac, 0xe6, 0x0a, 0xf8, 0x79, 0x1c, 0xbe, 0xc8, 0x38, 0xfb, 0x25, 0xac, - 0x25, 0x29, 0x37, 0x62, 0xef, 0x66, 0xa7, 0xe8, 0x02, 0x99, 0x1d, 0xbc, 0x61, 0xa4, 0x34, 0xef, - 0x71, 0x8a, 0x12, 0xda, 0x4f, 0xe6, 0xda, 0xa3, 0xaf, 0xe1, 0x72, 0x4e, 0x71, 0xc4, 0x7d, 0x11, - 0x07, 0x5e, 0x7a, 0x4e, 0xb6, 0x72, 0x61, 0x6c, 0xf9, 0x7d, 0xc6, 0x3e, 0xa2, 0xb1, 0xff, 0x7b, - 0x05, 0x3a, 0x8f, 0xb3, 0xd7, 0xaf, 0xcf, 0xf5, 0x5e, 0x62, 0x5d, 0xa8, 0x3c, 0xa5, 0x01, 0x56, - 0x9c, 0xca, 0x53, 0x74, 0xe7, 0x0e, 0x4f, 0x71, 0x5f, 0x93, 0x9c, 0xb7, 0x1d, 0xd3, 0xc2, 0xd3, - 0xda, 0xe1, 0xe9, 0xf1, 0x5b, 0x24, 0x5a, 0xa3, 0xf1, 0x98, 0xf1, 0x20, 0x0b, 0x23, 0x74, 0x36, - 0x8c, 0xf0, 0xe6, 0x6d, 0x3c, 0xff, 0xec, 0x8f, 0xf5, 0x54, 0x1e, 0xa7, 0x62, 0xa6, 0x99, 0x65, - 0x54, 0xc6, 0x05, 0x18, 0xf6, 0x2b, 0x58, 0x37, 0x51, 0x22, 0xa3, 0x15, 0x5c, 0x99, 0x70, 0x9f, - 0x44, 0xf7, 0x7b, 0x45, 0x96, 0x46, 0x7f, 0x55, 0x83, 0xd6, 0x97, 0x9e, 0x9c, 0xfe, 0x5a, 0x84, - 0x31, 0xbb, 0x07, 0xed, 0x6f, 0x44, 0x18, 0xeb, 0xa3, 0xaf, 0x0e, 0xfa, 0xae, 0xeb, 0xb1, 0x9e, - 0x8a, 0x80, 0x6f, 0x23, 0x0d, 0x1d, 0x7a, 0x5b, 0xdf, 0x98, 0x27, 0xa3, 0xe4, 0xd3, 0x70, 0x32, - 0x55, 0x2e, 0x02, 0x8d, 0x6e, 0xed, 0x84, 0xd2, 0x41, 0x18, 0x8d, 0xfa, 0x2e, 0xa0, 0xbd, 0x9b, - 0xba, 0x22, 0x76, 0x93, 0x53, 0x73, 0x3a, 0x6a, 0x21, 0xe4, 0x59, 0x7c, 0x78, 0x8a, 0x7b, 0x2f, - 0x94, 0xae, 0x09, 0xb2, 0x18, 0x4f, 0xb6, 0x74, 0xc8, 0xfc, 0x10, 0xfa, 0xe8, 0x65, 0xc8, 0xd3, - 0x30, 0x71, 0x93, 0x54, 0x9c, 0x58, 0xa6, 0xa0, 0xef, 0x71, 0x74, 0x1a, 0x26, 0x87, 0x08, 0x23, - 0xe3, 0x6e, 0x42, 0x37, 0xa8, 0xb6, 0xb5, 0x15, 0x05, 0x03, 0x42, 0xfe, 0x52, 0x7c, 0x26, 0xd2, - 0xbe, 0x76, 0x93, 0x8c, 0x76, 0x33, 0xe5, 0x11, 0x39, 0xd5, 0x57, 0xa0, 0x85, 0x9b, 0x81, 0x50, - 0x2d, 0x8d, 0xf2, 0x85, 0x46, 0xfd, 0x18, 0x20, 0xe2, 0x63, 0xe5, 0xa2, 0x94, 0xe9, 0xd3, 0xe8, - 0x42, 0x1c, 0x04, 0xb1, 0x7b, 0x88, 0x64, 0x1f, 0x43, 0x47, 0x73, 0x41, 0xd3, 0xc2, 0x12, 0x2d, - 0x10, 0x5a, 0x13, 0xdf, 0x86, 0x4e, 0x2c, 0x62, 0x97, 0xbf, 0x20, 0x6a, 0xb3, 0x6f, 0xe7, 0x06, - 0x8e, 0x45, 0xfc, 0xe8, 0x05, 0x12, 0xb3, 0xbb, 0x66, 0x0e, 0x3a, 0x20, 0xd0, 0x7d, 0x43, 0x40, - 0x80, 0x66, 0xa2, 0x8f, 0xc6, 0xf7, 0xed, 0x4c, 0x74, 0x8f, 0xde, 0x1b, 0x7a, 0xe8, 0xf9, 0xe8, - 0x2e, 0x37, 0xa0, 0x4b, 0xeb, 0x3e, 0xf3, 0x12, 0x57, 0x79, 0x13, 0xe3, 0x8d, 0x01, 0xc2, 0x0e, - 0xbc, 0xe4, 0xd8, 0x9b, 0x30, 0x07, 0xae, 0x2c, 0xc8, 0xdb, 0x09, 0x8a, 0xae, 0xe6, 0xda, 0xaa, - 0x0d, 0x28, 0x5c, 0x2c, 0x75, 0x9b, 0x73, 0x52, 0x47, 0x22, 0x8f, 0xdc, 0x1d, 0xfd, 0xd3, 0x15, - 0x68, 0x3d, 0x11, 0x22, 0xf9, 0x81, 0xa2, 0x57, 0x5e, 0xd2, 0x95, 0x37, 0x2f, 0x69, 0x75, 0x7e, - 0x49, 0x17, 0x58, 0x5f, 0xfb, 0xee, 0xac, 0xaf, 0x7f, 0x6f, 0xd6, 0x37, 0x7e, 0x00, 0xeb, 0x9b, - 0x8b, 0xac, 0x1f, 0x35, 0xa1, 0x7e, 0xc4, 0xd5, 0xb3, 0x64, 0xf4, 0x2f, 0x5b, 0xd0, 0x7e, 0xc8, - 0x83, 0x4c, 0x33, 0xac, 0xfc, 0xf9, 0x95, 0x37, 0x7f, 0xfe, 0xca, 0xfc, 0xe7, 0xa3, 0x21, 0xb2, - 0x12, 0x7d, 0x41, 0x6c, 0xac, 0x65, 0x05, 0x1a, 0x45, 0xbf, 0x90, 0x67, 0x13, 0x60, 0x9a, 0x63, - 0x53, 0x2e, 0xce, 0x6f, 0x97, 0x8d, 0xfa, 0x0f, 0x92, 0x8d, 0x05, 0xad, 0xb0, 0x14, 0x7a, 0xfa, - 0x56, 0xae, 0x2d, 0x6a, 0x84, 0xd6, 0x92, 0x46, 0x78, 0x02, 0xeb, 0x22, 0x76, 0x83, 0x2c, 0x89, - 0x42, 0x3c, 0xab, 0xb8, 0x9e, 0x3e, 0xa9, 0xb7, 0x6d, 0xbe, 0x25, 0x17, 0xbd, 0x67, 0xf1, 0x43, - 0x4b, 0xa4, 0xcf, 0xef, 0xce, 0x9a, 0x58, 0x04, 0xa1, 0x9a, 0x0a, 0x70, 0x69, 0xc8, 0xae, 0x92, - 0x47, 0xa8, 0x13, 0x47, 0x5d, 0x82, 0xee, 0x89, 0x88, 0x2c, 0xc5, 0x17, 0xb0, 0x5a, 0x50, 0x69, - 0x19, 0xe9, 0xbc, 0x41, 0x46, 0x7a, 0xb6, 0xa3, 0x16, 0x93, 0xbf, 0x0e, 0x2d, 0xf0, 0x09, 0xac, - 0xdb, 0xb0, 0x84, 0x71, 0x0e, 0x68, 0x05, 0xfb, 0x24, 0x41, 0x03, 0x13, 0x89, 0x20, 0xbf, 0x80, - 0x96, 0xe8, 0xe7, 0xb0, 0x51, 0x22, 0xc7, 0x73, 0x43, 0x59, 0x1b, 0x94, 0x65, 0x65, 0x2d, 0xef, - 0x8b, 0xcd, 0x27, 0x3a, 0x3c, 0xdb, 0x09, 0x78, 0x64, 0x5f, 0x34, 0x1c, 0xe8, 0x63, 0x4f, 0xc0, - 0x23, 0x93, 0x2d, 0x3a, 0x80, 0x0f, 0xf1, 0x74, 0x81, 0x78, 0xdf, 0x4b, 0x54, 0x96, 0x72, 0x37, - 0x89, 0x3c, 0x9f, 0x4f, 0x45, 0x14, 0xf0, 0xb4, 0x98, 0xdc, 0x1a, 0x4d, 0xee, 0xba, 0x88, 0x82, - 0x3d, 0x11, 0xed, 0x69, 0xca, 0xc3, 0x82, 0xd0, 0xce, 0x75, 0x17, 0xde, 0x5b, 0x1a, 0x0e, 0x0d, - 0x47, 0x31, 0x10, 0xa3, 0x81, 0xae, 0xcc, 0x0f, 0x84, 0x24, 0x76, 0x88, 0xfb, 0x70, 0x49, 0xaf, - 0x9d, 0x16, 0xee, 0x53, 0xce, 0x13, 0x37, 0xf2, 0xa4, 0x1a, 0xae, 0x6b, 0x23, 0x4d, 0x48, 0x12, - 0xe0, 0xdf, 0x70, 0x9e, 0x3c, 0xf1, 0xf4, 0x5b, 0x75, 0x17, 0xe3, 0xc7, 0x53, 0x9f, 0x39, 0xde, - 0x6e, 0xe8, 0xb7, 0x12, 0x95, 0x76, 0xe6, 0xb1, 0x73, 0x89, 0xc9, 0xbf, 0x0f, 0xd7, 0xe6, 0x86, - 0x98, 0x79, 0xe9, 0x69, 0xe1, 0xd8, 0x0e, 0x2f, 0x11, 0xdf, 0x2e, 0x97, 0xfa, 0x1f, 0x10, 0x81, - 0x1e, 0x61, 0xf4, 0xdf, 0xea, 0xd0, 0x27, 0x3b, 0xfc, 0xb7, 0x6a, 0xe3, 0x6f, 0xd5, 0xc6, 0xdf, - 0x00, 0xb5, 0x31, 0xfa, 0x07, 0x15, 0x68, 0x1e, 0xa6, 0x22, 0xc8, 0x7c, 0xf5, 0x03, 0x25, 0x7d, - 0x5e, 0x82, 0xaa, 0xdf, 0x26, 0x41, 0xb5, 0x25, 0x73, 0xfd, 0xcf, 0x2a, 0xd0, 0x36, 0x53, 0x78, - 0xb2, 0xf3, 0x03, 0x27, 0x51, 0x24, 0xaf, 0x2a, 0x17, 0x26, 0xaf, 0xbe, 0x75, 0x16, 0x28, 0x58, - 0x2f, 0x75, 0x16, 0x5f, 0x24, 0x45, 0x26, 0xab, 0xed, 0x74, 0x35, 0xf4, 0x59, 0x42, 0x09, 0xab, - 0x57, 0xd0, 0xa6, 0x93, 0x13, 0x69, 0x86, 0x4d, 0x68, 0xa4, 0x94, 0x61, 0x31, 0x13, 0x35, 0xad, - 0xb7, 0xef, 0xd3, 0x95, 0x1f, 0xe6, 0xfa, 0xfd, 0xdb, 0x15, 0xe8, 0xd1, 0x31, 0xf6, 0x71, 0x16, - 0xeb, 0x9d, 0x90, 0x87, 0xcf, 0x2a, 0xf3, 0xe1, 0xb3, 0x5a, 0x8a, 0xa7, 0x4d, 0xfd, 0x9a, 0xae, - 0x7e, 0xcd, 0x9e, 0x88, 0x1e, 0xf2, 0xb1, 0x43, 0x18, 0x64, 0x95, 0x97, 0x4e, 0xe4, 0x45, 0x79, - 0x3e, 0x84, 0xe3, 0x57, 0x25, 0x5e, 0xea, 0xcd, 0xa4, 0xcd, 0xf3, 0xe9, 0x16, 0x63, 0x50, 0xa3, - 0xfd, 0xa6, 0xd9, 0x42, 0xcf, 0x26, 0x22, 0x23, 0xc3, 0x78, 0x92, 0x2b, 0x8f, 0x16, 0xe5, 0x77, - 0x27, 0x11, 0x67, 0x0f, 0x81, 0xe9, 0x80, 0x6d, 0xca, 0x3d, 0x34, 0x41, 0x34, 0x0e, 0x69, 0x90, - 0xce, 0xce, 0xa6, 0x7e, 0x2d, 0xf1, 0xd2, 0x21, 0xf4, 0x21, 0x62, 0x9d, 0x41, 0xb8, 0x00, 0xb9, - 0x80, 0x99, 0xda, 0x0e, 0xe5, 0xa7, 0x8f, 0xef, 0xcc, 0x4c, 0x32, 0x4e, 0xc4, 0xcc, 0x5d, 0xb8, - 0x64, 0xb3, 0x27, 0xa8, 0x2e, 0x76, 0x70, 0x2f, 0xd0, 0x79, 0xd8, 0x7e, 0x63, 0xa5, 0xf4, 0x8d, - 0x1b, 0x50, 0x2f, 0xd7, 0x75, 0xe8, 0xc6, 0xe8, 0x26, 0x74, 0xc6, 0x61, 0xc4, 0x4d, 0x14, 0x12, - 0x99, 0x66, 0xe2, 0x91, 0x15, 0xaa, 0x6c, 0x30, 0xad, 0xd1, 0x6f, 0x2b, 0x70, 0x39, 0xf1, 0xd2, - 0x17, 0x19, 0x57, 0x14, 0x8b, 0xa4, 0x6c, 0x9b, 0x2b, 0xa7, 0x5e, 0x1a, 0xe0, 0xc6, 0xa1, 0x21, - 0xf4, 0xe8, 0xba, 0x7c, 0xa0, 0x8d, 0x10, 0x3d, 0x97, 0x5b, 0xb0, 0x5a, 0xea, 0xa1, 0xbc, 0xd4, - 0x46, 0x8b, 0x7a, 0xa9, 0x78, 0x45, 0x49, 0xd3, 0x23, 0x04, 0xe2, 0x81, 0xb2, 0xa0, 0xe3, 0x64, - 0x6d, 0x28, 0x0b, 0x6f, 0xa9, 0x1e, 0xc5, 0x01, 0xee, 0x9c, 0x38, 0x9b, 0xe9, 0x60, 0x8a, 0xae, - 0xfe, 0x68, 0xc6, 0xd9, 0x8c, 0xe2, 0x27, 0x1b, 0x50, 0x3f, 0x39, 0x57, 0xe4, 0xad, 0x23, 0x5c, - 0x37, 0x46, 0x7f, 0x51, 0x87, 0xf5, 0x7d, 0x9f, 0x9f, 0xf0, 0x74, 0xf2, 0xd0, 0x53, 0xde, 0xe3, - 0x30, 0xe2, 0xc7, 0x9e, 0x3c, 0xc5, 0x05, 0xa7, 0x39, 0x27, 0x9e, 0x9a, 0x1a, 0x2e, 0xb5, 0x10, - 0x70, 0xe8, 0xa9, 0x29, 0x9a, 0x02, 0x42, 0x8e, 0x45, 0x3a, 0x33, 0xb1, 0xad, 0xb6, 0x43, 0xdf, - 0xf8, 0x98, 0x20, 0x79, 0x6f, 0x19, 0xbe, 0xe6, 0xa6, 0x56, 0x85, 0x7a, 0x53, 0xe2, 0xf3, 0x7d, - 0xe8, 0xa6, 0xdc, 0x17, 0x69, 0x60, 0x02, 0xb6, 0x7a, 0x9e, 0x1d, 0x0d, 0xd3, 0xa1, 0xda, 0xdb, - 0x50, 0x64, 0x15, 0xe8, 0xf8, 0xee, 0x86, 0x36, 0xf1, 0xbd, 0x9a, 0x23, 0x70, 0xe5, 0xf7, 0x03, - 0xf6, 0x77, 0x61, 0x50, 0xd0, 0x52, 0x88, 0xdb, 0x1e, 0x2f, 0x76, 0x8a, 0x10, 0xcc, 0x05, 0x9f, - 0xb8, 0x7d, 0x68, 0x7b, 0xfd, 0x1d, 0xea, 0xa4, 0xc3, 0xf8, 0xc5, 0xf0, 0x1a, 0xca, 0x3e, 0x80, - 0x9e, 0x4c, 0xa2, 0x50, 0x19, 0x01, 0x90, 0xa6, 0xa2, 0xa5, 0x4b, 0x40, 0x1d, 0x89, 0x96, 0x17, - 0x2d, 0x61, 0xeb, 0x3b, 0x2d, 0x61, 0x7b, 0x79, 0x09, 0x7f, 0x0c, 0x03, 0x3f, 0xe5, 0x01, 0x8f, - 0x55, 0xe8, 0x45, 0xae, 0xf4, 0x45, 0x62, 0x4d, 0xdf, 0x6a, 0x01, 0x3f, 0x42, 0x30, 0xfb, 0x09, - 0x5c, 0xf6, 0x45, 0xac, 0x78, 0xac, 0x5c, 0xc9, 0x5f, 0x64, 0x3c, 0xf6, 0xb9, 0x1b, 0x67, 0xb3, - 0x13, 0x9e, 0x9a, 0x9c, 0xee, 0x25, 0x83, 0x3e, 0x32, 0xd8, 0xa7, 0x84, 0x64, 0xf7, 0x60, 0x43, - 0x2f, 0xcf, 0x42, 0x27, 0x9d, 0x45, 0x64, 0xb4, 0x52, 0xf3, 0x3d, 0xb6, 0x61, 0x7d, 0xea, 0x49, - 0x37, 0xe5, 0x32, 0x0c, 0x32, 0x2f, 0x32, 0x3b, 0xd4, 0xe4, 0x2e, 0xd6, 0xa6, 0x9e, 0x74, 0x0c, - 0xc6, 0x84, 0x87, 0x28, 0x7a, 0x3d, 0x47, 0xeb, 0x4e, 0x3d, 0x39, 0xa5, 0xe3, 0x73, 0xdb, 0x61, - 0xe9, 0x1c, 0xf5, 0x97, 0x9e, 0x9c, 0x5e, 0x7d, 0x00, 0x1b, 0x17, 0x2d, 0xc8, 0xb7, 0xa5, 0x35, - 0xda, 0xa5, 0xb4, 0x86, 0xa9, 0xeb, 0xfa, 0x1f, 0x2b, 0x70, 0xc9, 0xae, 0x37, 0x39, 0x7e, 0xb9, - 0x50, 0x5f, 0x27, 0x1b, 0x89, 0xce, 0x62, 0x7e, 0x96, 0x6e, 0x3b, 0xa0, 0x41, 0x74, 0x70, 0xde, - 0x82, 0x81, 0x21, 0x28, 0x84, 0x5f, 0xbf, 0xa5, 0x1f, 0xe4, 0x43, 0xd1, 0x16, 0xa0, 0x0f, 0x1c, - 0xf3, 0x14, 0x79, 0x14, 0x50, 0x45, 0x1e, 0x75, 0x21, 0x61, 0xa7, 0x0f, 0xb4, 0x38, 0x2b, 0x72, - 0xec, 0x0e, 0x30, 0xfe, 0x22, 0xf3, 0xa2, 0x50, 0x9d, 0xbb, 0xe3, 0x90, 0x47, 0x01, 0xe5, 0xd0, - 0x74, 0xa1, 0xce, 0xc0, 0x62, 0x1e, 0x23, 0x62, 0x3f, 0x90, 0xa5, 0x99, 0x98, 0xd4, 0x4c, 0xbe, - 0x01, 0xcc, 0x4c, 0x8e, 0x08, 0xbc, 0x1f, 0x5c, 0xbc, 0x57, 0x1a, 0x17, 0xef, 0x95, 0x8f, 0x60, - 0x75, 0x71, 0xcd, 0x75, 0xba, 0xa4, 0x2f, 0xe7, 0xd7, 0xfb, 0x22, 0x21, 0x6c, 0x5d, 0x28, 0x84, - 0x86, 0xe9, 0xff, 0x6b, 0x05, 0x36, 0x0c, 0xd3, 0xf7, 0x44, 0x94, 0xcd, 0xd0, 0xda, 0x26, 0x61, - 0x3c, 0x41, 0x83, 0x3c, 0x13, 0xda, 0x2d, 0x29, 0xa9, 0x3f, 0x98, 0x89, 0x5c, 0x17, 0x6f, 0xc1, - 0x20, 0xd4, 0x3d, 0x73, 0xbe, 0xd8, 0xd2, 0x3a, 0x03, 0x37, 0x5c, 0x41, 0x29, 0x94, 0xb1, 0x97, - 0xc8, 0xa9, 0x50, 0x86, 0x94, 0x94, 0xb8, 0xe6, 0xf9, 0x9a, 0x45, 0x11, 0x35, 0x79, 0x87, 0x77, - 0x80, 0xf9, 0x59, 0x9a, 0xe2, 0xfe, 0x28, 0x91, 0xeb, 0x84, 0xc4, 0xc0, 0x60, 0x0a, 0xea, 0x0f, - 0xa0, 0x39, 0x13, 0x85, 0x47, 0x30, 0xe7, 0xdc, 0x39, 0x8d, 0x99, 0x20, 0x09, 0xb9, 0x8a, 0x5e, - 0xcb, 0x8b, 0x2c, 0x4c, 0x79, 0x60, 0xed, 0xa0, 0x6d, 0x1b, 0x23, 0x39, 0x0d, 0x83, 0x80, 0xc7, - 0x26, 0x18, 0xde, 0x0a, 0xe5, 0x97, 0xd4, 0xa6, 0xca, 0x33, 0x3e, 0xf6, 0xb2, 0x48, 0xb9, 0x71, - 0x16, 0xd1, 0xae, 0x88, 0x4c, 0x3d, 0xd4, 0xaa, 0x41, 0x3c, 0xcd, 0x22, 0xdc, 0x11, 0x91, 0x59, - 0x52, 0xb2, 0x25, 0x28, 0x82, 0xee, 0x34, 0x8c, 0x15, 0xa9, 0x8a, 0x36, 0x2d, 0x29, 0x22, 0x50, - 0x08, 0xbf, 0x0c, 0x63, 0x35, 0xfa, 0xb3, 0x15, 0xd8, 0x34, 0x8c, 0x3f, 0x32, 0x0c, 0x30, 0xf6, - 0x91, 0x3c, 0x76, 0xcb, 0x2e, 0x93, 0xab, 0xa8, 0x3a, 0x60, 0x41, 0xfb, 0x34, 0xe1, 0x42, 0xba, - 0x56, 0x4c, 0x9d, 0x94, 0x95, 0xab, 0x3b, 0xc0, 0x96, 0xe4, 0x4a, 0x9a, 0x98, 0xd1, 0x60, 0x41, - 0xb0, 0x24, 0xfb, 0x0c, 0x36, 0x67, 0x5c, 0x79, 0xb4, 0x11, 0x22, 0xe1, 0x7b, 0xd4, 0x8b, 0xb6, - 0xbc, 0x66, 0xf7, 0x86, 0xc5, 0x3e, 0x31, 0x48, 0xdc, 0xf4, 0xf8, 0x8e, 0x99, 0x17, 0x87, 0x63, - 0x2e, 0x15, 0xd9, 0x79, 0xdd, 0x43, 0x3b, 0x1e, 0x03, 0x8b, 0x41, 0x4b, 0x4e, 0xd4, 0xe4, 0x31, - 0x8e, 0xf5, 0x22, 0x36, 0x88, 0xa6, 0x99, 0xf2, 0xb1, 0x59, 0xbb, 0x1e, 0xae, 0x55, 0x1c, 0xc6, - 0x13, 0x5d, 0x1c, 0xda, 0xd4, 0x3e, 0x9d, 0x05, 0x1e, 0x88, 0x80, 0x8f, 0xfe, 0xa4, 0x96, 0xcb, - 0xe8, 0xa1, 0x81, 0x1f, 0x29, 0x4f, 0x49, 0x76, 0x13, 0xfa, 0xf9, 0xe4, 0xb5, 0x8d, 0xd4, 0xbc, - 0xea, 0x59, 0xe8, 0x03, 0x04, 0xa2, 0xf8, 0xcd, 0xcf, 0x56, 0xd3, 0xae, 0xe8, 0x84, 0x63, 0x79, - 0xba, 0x9a, 0x1e, 0x87, 0xb5, 0xf4, 0x9a, 0xd4, 0x94, 0x6d, 0x5a, 0xa8, 0x26, 0xfb, 0xa4, 0x60, - 0x82, 0x74, 0x25, 0x8f, 0x74, 0xb5, 0x4d, 0x6d, 0x7e, 0x54, 0x79, 0x64, 0x10, 0xb8, 0x35, 0x0b, - 0xf2, 0x24, 0xcd, 0x62, 0x1e, 0x18, 0x93, 0xbe, 0x9a, 0xc3, 0x0f, 0x09, 0x8c, 0x13, 0xce, 0x35, - 0x53, 0x69, 0xe8, 0x86, 0x1e, 0x3a, 0x30, 0x9a, 0xa9, 0x18, 0x1a, 0x65, 0xb4, 0xa0, 0x37, 0x63, - 0x6b, 0x05, 0xb1, 0x9a, 0x53, 0x9b, 0xb1, 0x7f, 0x0a, 0xc3, 0x9c, 0x56, 0x7f, 0x5d, 0xf1, 0x82, - 0x96, 0x36, 0x3e, 0xb6, 0x0b, 0x7d, 0x66, 0xfe, 0x92, 0x4f, 0x61, 0x73, 0xb1, 0xa3, 0x79, 0x53, - 0x9b, 0xba, 0xad, 0xcf, 0x75, 0x2b, 0xbe, 0x24, 0x5f, 0x5f, 0xdf, 0xf3, 0xa7, 0xdc, 0x9d, 0x86, - 0xa6, 0x72, 0xb3, 0xea, 0xac, 0x59, 0xd4, 0x1e, 0x62, 0xbe, 0x0c, 0x95, 0xbc, 0x80, 0x7e, 0x16, - 0x4a, 0x69, 0xac, 0xe2, 0x3c, 0xfd, 0x41, 0x28, 0xe5, 0xe8, 0x1f, 0x03, 0x74, 0xad, 0xa7, 0x48, - 0xc5, 0x85, 0x77, 0xca, 0x4e, 0x77, 0x67, 0x67, 0x60, 0xbd, 0x67, 0x24, 0xd9, 0x55, 0x2a, 0xb5, - 0xf9, 0x0b, 0xed, 0x8c, 0xcf, 0xf9, 0x3b, 0x2b, 0xe4, 0x20, 0x14, 0xfe, 0xce, 0x2e, 0xac, 0x95, - 0x3c, 0x48, 0x57, 0x09, 0xe5, 0x45, 0xc6, 0x29, 0x2f, 0x55, 0x94, 0x94, 0x48, 0x9c, 0x55, 0x6c, - 0x68, 0xdf, 0xe2, 0x18, 0xa9, 0xd1, 0xd9, 0xf7, 0x45, 0x64, 0xab, 0xd9, 0x16, 0x9c, 0x7d, 0xc4, - 0x50, 0xae, 0x3c, 0xe5, 0x78, 0x76, 0x94, 0x2f, 0x22, 0xb3, 0x83, 0xda, 0x1a, 0x72, 0xf4, 0x22, - 0xca, 0x27, 0x48, 0xce, 0x74, 0x83, 0xce, 0x11, 0x34, 0x41, 0x3a, 0x53, 0x7d, 0x02, 0x1d, 0x91, - 0x86, 0x93, 0x90, 0x52, 0x5f, 0xda, 0xc1, 0x59, 0x7c, 0x09, 0x68, 0x82, 0x3d, 0x7c, 0xd5, 0x08, - 0x1a, 0xc6, 0xfc, 0x2f, 0xe7, 0xcf, 0x0d, 0x06, 0x1d, 0x22, 0xa9, 0xd2, 0xd0, 0x57, 0x38, 0x1d, - 0xbd, 0x23, 0x75, 0x61, 0x54, 0x4f, 0x83, 0x8f, 0x5e, 0x44, 0x94, 0xfd, 0xbb, 0x05, 0xab, 0x3e, - 0x99, 0x0b, 0xbd, 0xa1, 0x22, 0x1e, 0xd3, 0x9a, 0xd6, 0x9d, 0x9e, 0x06, 0xe3, 0xfc, 0x9e, 0xf0, - 0xd8, 0x14, 0x61, 0x79, 0x51, 0x84, 0x27, 0x46, 0xe1, 0x05, 0x26, 0x63, 0xde, 0xb5, 0xc0, 0x27, - 0xc2, 0x0b, 0xd8, 0xcf, 0xe0, 0x2a, 0xe2, 0x5c, 0x3e, 0x4b, 0xd4, 0x39, 0xda, 0x37, 0x9e, 0x86, - 0xbe, 0xeb, 0x49, 0xca, 0xa0, 0x9b, 0xc4, 0xf9, 0x26, 0x52, 0x3c, 0x42, 0x82, 0xa7, 0x1a, 0xbf, - 0x2b, 0xbf, 0xe6, 0xa9, 0x60, 0x5f, 0x53, 0x06, 0xf0, 0x22, 0xf7, 0xdd, 0x1e, 0xf5, 0xdf, 0x2f, - 0xd6, 0xea, 0x0d, 0x94, 0x54, 0xa1, 0x82, 0x08, 0xc7, 0x3a, 0x7d, 0xd4, 0x9f, 0xfd, 0x06, 0x98, - 0x35, 0x70, 0x24, 0xf9, 0xca, 0x93, 0xa7, 0x92, 0xa2, 0x00, 0x9d, 0x9d, 0x1f, 0xbd, 0xd5, 0x47, - 0x75, 0xac, 0x65, 0x44, 0x20, 0x02, 0x24, 0xfb, 0x23, 0xd8, 0xc8, 0x07, 0x33, 0xbe, 0x0c, 0x0d, - 0xa7, 0x83, 0x04, 0xd7, 0x97, 0x87, 0x9b, 0x73, 0x81, 0x1c, 0x3b, 0x13, 0x0d, 0xd6, 0x43, 0xfe, - 0x0a, 0x56, 0xed, 0x90, 0x9a, 0xeb, 0x72, 0x38, 0xa0, 0xd1, 0xde, 0x5b, 0x1a, 0x6d, 0xce, 0xb6, - 0xe7, 0xf6, 0x59, 0x43, 0xf1, 0x43, 0x73, 0x4b, 0x6e, 0xad, 0xcc, 0x70, 0x8d, 0x64, 0xe4, 0xc6, - 0xd2, 0x48, 0x0b, 0xc6, 0xca, 0xb1, 0x53, 0xb0, 0x70, 0x76, 0x1f, 0x2e, 0xd9, 0xc1, 0x04, 0x25, - 0x6c, 0xdd, 0x50, 0x50, 0x2e, 0x97, 0x69, 0x17, 0xcb, 0x20, 0x75, 0x32, 0x77, 0x5f, 0x38, 0x7c, - 0xcc, 0x7e, 0x01, 0xd7, 0x6c, 0x17, 0x6d, 0x85, 0xe9, 0x44, 0x9a, 0x7f, 0xd4, 0x3a, 0xd9, 0xae, - 0xa1, 0x21, 0xd1, 0x76, 0x19, 0x4f, 0xa0, 0x76, 0xfa, 0x5b, 0x30, 0xa0, 0xd2, 0x54, 0x5c, 0x56, - 0x91, 0x06, 0x61, 0xec, 0x45, 0xc3, 0x0d, 0x92, 0x9a, 0x3e, 0xc2, 0x1d, 0xf1, 0xea, 0x99, 0x86, - 0xb2, 0x63, 0xd8, 0xb4, 0x2f, 0xca, 0xd5, 0x8c, 0x44, 0x53, 0x42, 0x61, 0xc7, 0x8b, 0x18, 0x37, - 0x67, 0x70, 0x1c, 0xbb, 0x84, 0xf3, 0x66, 0xe8, 0x21, 0x5c, 0x5f, 0x58, 0xda, 0x99, 0x77, 0xe6, - 0xce, 0xf8, 0x4c, 0xa4, 0xe7, 0xc6, 0x80, 0x6c, 0x92, 0x02, 0xbb, 0x36, 0xb7, 0x88, 0x07, 0xde, - 0xd9, 0x01, 0xd1, 0x68, 0x73, 0xf2, 0x4b, 0x78, 0x77, 0x61, 0x14, 0x5d, 0xe9, 0xc9, 0x63, 0xef, - 0x24, 0xe2, 0xc1, 0xf0, 0x32, 0x7d, 0xd1, 0x95, 0xb9, 0x21, 0x8e, 0x90, 0xe2, 0x91, 0x26, 0x30, - 0x0e, 0xdd, 0x09, 0xb4, 0x29, 0x0c, 0x41, 0xda, 0x30, 0xaf, 0xba, 0xad, 0xbc, 0xbd, 0xea, 0xf6, - 0x13, 0xe8, 0x1a, 0x6f, 0xff, 0x4d, 0x65, 0xbc, 0x1d, 0x8d, 0xc7, 0x67, 0x39, 0xba, 0x03, 0x6d, - 0x72, 0xf5, 0xe9, 0x1d, 0xd7, 0xa1, 0x43, 0xd5, 0x5e, 0xee, 0x49, 0x24, 0xfc, 0x53, 0xeb, 0x9c, - 0x13, 0xe8, 0x01, 0x42, 0x46, 0x00, 0xad, 0xe7, 0x71, 0x28, 0xe2, 0xdd, 0x28, 0x1a, 0xfd, 0x65, - 0x03, 0xda, 0xe8, 0x13, 0x50, 0xdc, 0x04, 0x8f, 0x55, 0xb4, 0x70, 0x94, 0x4b, 0x9d, 0x79, 0x89, - 0xa9, 0x2b, 0xee, 0x20, 0x10, 0xa9, 0x0e, 0xbc, 0x64, 0x21, 0xd5, 0xba, 0xb2, 0x90, 0x6a, 0x7d, - 0x5f, 0xdf, 0x7d, 0xd1, 0xf5, 0x66, 0xdc, 0x16, 0xaa, 0xd2, 0x00, 0x0f, 0x34, 0x08, 0x7d, 0x15, - 0x22, 0xf1, 0x22, 0xf2, 0x6f, 0xf0, 0xf4, 0x14, 0x49, 0x93, 0x95, 0x25, 0xb9, 0xd9, 0x35, 0x88, - 0x23, 0xae, 0xf5, 0x71, 0x29, 0x58, 0x56, 0x5f, 0x0c, 0x96, 0xdd, 0x06, 0xf0, 0x45, 0x1c, 0x90, - 0x0b, 0xb5, 0x90, 0x0d, 0xd3, 0x29, 0xd1, 0x02, 0xfb, 0x1d, 0x42, 0xb3, 0x1f, 0xc1, 0x20, 0xa7, - 0x40, 0x0f, 0xc9, 0x8f, 0xf3, 0xf3, 0xa7, 0xa1, 0x72, 0xf8, 0x78, 0x2f, 0x56, 0x8b, 0x31, 0xdc, - 0xf6, 0x52, 0x0c, 0xf7, 0x0d, 0xc9, 0x73, 0xf8, 0xde, 0xd7, 0x32, 0xae, 0x40, 0x8b, 0xaa, 0x74, - 0x82, 0x2c, 0x31, 0xba, 0xba, 0x19, 0x4a, 0x8a, 0xb5, 0xbf, 0x29, 0x4e, 0xdc, 0xfd, 0xff, 0x15, - 0x27, 0xee, 0x7d, 0xb7, 0x38, 0x71, 0xff, 0xbb, 0xc5, 0x89, 0x17, 0xe2, 0xaa, 0xab, 0x8b, 0xe9, - 0x98, 0x37, 0x26, 0x3f, 0x06, 0x6f, 0x4c, 0x7e, 0x7c, 0x4b, 0xe6, 0x62, 0xed, 0xad, 0x99, 0x8b, - 0xef, 0x90, 0x3a, 0x61, 0xdf, 0x96, 0x3a, 0xb9, 0x05, 0xab, 0x2a, 0xf5, 0xfc, 0x53, 0x7d, 0x12, - 0x39, 0xe5, 0xe7, 0xd2, 0xa4, 0x6a, 0x7a, 0x04, 0xc6, 0x73, 0xc8, 0x6f, 0xf8, 0xb9, 0x1c, 0x3d, - 0x07, 0xa0, 0x23, 0x1a, 0x7d, 0xda, 0x9b, 0x64, 0xa3, 0xf2, 0xbd, 0x0b, 0x2b, 0xfe, 0x4f, 0x05, - 0xe0, 0xc8, 0x9b, 0x25, 0x3a, 0xc6, 0xc9, 0xfe, 0x10, 0x3a, 0x92, 0x5a, 0xe5, 0x0c, 0x77, 0xc9, - 0x90, 0x15, 0xa4, 0xe6, 0x51, 0xdf, 0x2e, 0x90, 0xf9, 0x33, 0x89, 0xb5, 0x1e, 0x21, 0x2f, 0x62, - 0xab, 0x5b, 0x02, 0x8a, 0x7d, 0xdd, 0x84, 0xbe, 0x21, 0x48, 0x78, 0xea, 0xf3, 0x58, 0x57, 0xc6, - 0x56, 0x9c, 0x9e, 0x86, 0x1e, 0x6a, 0x20, 0xbb, 0x9f, 0x93, 0x59, 0x93, 0xb1, 0x9c, 0xa6, 0x31, - 0x5d, 0x8c, 0xcd, 0x18, 0xed, 0xd8, 0x4f, 0xa1, 0x89, 0xb4, 0xa0, 0x86, 0xef, 0x1b, 0xbc, 0xc3, - 0x3a, 0xd0, 0x34, 0xa3, 0x0e, 0x2a, 0xac, 0x07, 0x6d, 0xba, 0xce, 0x42, 0xb8, 0x95, 0xd1, 0x1f, - 0xaf, 0x41, 0x67, 0x3f, 0x96, 0x2a, 0xcd, 0xb4, 0x08, 0x17, 0x97, 0x36, 0xea, 0x74, 0x69, 0xc3, - 0x54, 0x50, 0xea, 0xcf, 0xa0, 0x0a, 0xca, 0x4f, 0xa0, 0x69, 0xee, 0x07, 0x99, 0xc0, 0xf7, 0x85, - 0x97, 0x8b, 0x2c, 0x0d, 0xdb, 0x86, 0x56, 0x60, 0x2e, 0x2e, 0x99, 0x34, 0x7e, 0xe9, 0x36, 0x91, - 0xbd, 0xd2, 0xe4, 0xe4, 0x34, 0xec, 0x7d, 0xa8, 0x7a, 0x93, 0x89, 0x39, 0xf5, 0xae, 0x16, 0xa4, - 0xe4, 0xc4, 0x38, 0x88, 0x63, 0x77, 0xa1, 0x4d, 0xea, 0x93, 0x2a, 0x59, 0x1a, 0x8b, 0x63, 0xda, - 0x32, 0x19, 0xad, 0x51, 0x29, 0x66, 0x7e, 0x17, 0xda, 0x91, 0x10, 0x89, 0xee, 0xd0, 0x5c, 0xec, + 0x92, 0xd8, 0x34, 0xfb, 0x3b, 0xba, 0x9b, 0x6c, 0x26, 0x29, 0xaa, 0x25, 0xcd, 0x8c, 0x34, 0x3d, + 0x23, 0x0d, 0x9f, 0x46, 0x43, 0x49, 0x9c, 0x99, 0xf7, 0x66, 0xdf, 0xdb, 0xb7, 0x6f, 0x29, 0x4a, + 0x9a, 0xe1, 0x7b, 0xa2, 0xc4, 0x2d, 0x52, 0x1e, 0x60, 0x00, 0xbb, 0x50, 0xac, 0xca, 0xee, 0xae, + 0x61, 0x75, 0x65, 0xa9, 0x32, 0x4b, 0x22, 0x75, 0xb1, 0x0f, 0x3e, 0xf9, 0xe2, 0xe3, 0xee, 0x71, + 0x01, 0xfb, 0xb0, 0xb6, 0x01, 0xfb, 0x60, 0x3c, 0xff, 0x04, 0x63, 0x61, 0xfb, 0xb0, 0xf0, 0xc1, + 0x47, 0xc3, 0x78, 0x7b, 0x34, 0x60, 0x18, 0x06, 0x6c, 0x2c, 0x60, 0x18, 0x30, 0x22, 0x32, 0xb3, + 0xaa, 0xba, 0x9b, 0x94, 0x66, 0xc6, 0xc6, 0x5e, 0x76, 0x4f, 0x5d, 0x19, 0x11, 0x99, 0x95, 0x15, + 0x19, 0x19, 0x11, 0x19, 0x11, 0xd9, 0xb0, 0x9c, 0x84, 0x09, 0x8f, 0xc2, 0x98, 0x6f, 0x25, 0xa9, + 0x50, 0x82, 0xb5, 0x6c, 0xfb, 0xea, 0xa7, 0xe3, 0x50, 0x4d, 0xb2, 0xe3, 0x2d, 0x5f, 0x4c, 0xef, + 0x8e, 0xc5, 0x58, 0xdc, 0x25, 0x82, 0xe3, 0x6c, 0x44, 0x2d, 0x6a, 0xd0, 0x93, 0xee, 0x78, 0x15, + 0x22, 0xe1, 0x9f, 0xd8, 0xe7, 0x24, 0xf2, 0x62, 0xf3, 0xbc, 0xa2, 0xc2, 0x29, 0x97, 0xca, 0x9b, + 0x26, 0x06, 0xd0, 0x56, 0xa7, 0x06, 0x37, 0xfc, 0xf7, 0x55, 0x68, 0xee, 0x73, 0x29, 0xbd, 0x31, + 0x67, 0x43, 0xa8, 0xca, 0x30, 0x18, 0x54, 0x6e, 0x54, 0x36, 0x97, 0xb7, 0xfb, 0x5b, 0xf9, 0xb4, + 0x0e, 0x95, 0xa7, 0x32, 0xe9, 0x20, 0x12, 0x69, 0xfc, 0x69, 0x30, 0x58, 0x9a, 0xa7, 0xd9, 0xe7, + 0x6a, 0x22, 0x02, 0x07, 0x91, 0xac, 0x0f, 0x55, 0x9e, 0xa6, 0x83, 0xea, 0x8d, 0xca, 0x66, 0xd7, + 0xc1, 0x47, 0xc6, 0xa0, 0x16, 0x78, 0xca, 0x1b, 0xd4, 0x08, 0x44, 0xcf, 0xec, 0x23, 0x58, 0x4e, + 0x52, 0xe1, 0xbb, 0x61, 0x3c, 0x12, 0x2e, 0x61, 0xeb, 0x84, 0xed, 0x22, 0x74, 0x2f, 0x1e, 0x89, + 0x87, 0x48, 0x35, 0x80, 0xa6, 0x17, 0x7b, 0xd1, 0x99, 0xe4, 0x83, 0x06, 0xa1, 0x6d, 0x93, 0x2d, + 0xc3, 0x52, 0x18, 0x0c, 0x9a, 0x37, 0x2a, 0x9b, 0x35, 0x67, 0x29, 0x0c, 0xf0, 0x1d, 0x59, 0x16, + 0x06, 0x83, 0x96, 0x7e, 0x07, 0x3e, 0xb3, 0x21, 0x74, 0x63, 0xce, 0x83, 0xa7, 0x42, 0x39, 0x3c, + 0x89, 0xce, 0x06, 0xed, 0x1b, 0x95, 0xcd, 0x96, 0x33, 0x03, 0x63, 0x57, 0xa1, 0x15, 0xf0, 0xe3, + 0x6c, 0xbc, 0x2f, 0xc7, 0x03, 0xb8, 0x51, 0xd9, 0x6c, 0x3b, 0x79, 0x9b, 0x1d, 0xc1, 0xe5, 0x94, + 0xbf, 0xc8, 0xb8, 0x54, 0x3c, 0x70, 0x15, 0xf7, 0xd2, 0x40, 0xbc, 0x8a, 0xdd, 0xa9, 0x08, 0xf8, + 0xa0, 0x43, 0x1c, 0x78, 0xb7, 0xcc, 0xa5, 0x94, 0x7b, 0xd3, 0x23, 0x43, 0xb4, 0x2f, 0x02, 0xee, + 0x5c, 0xca, 0x3b, 0x97, 0xc1, 0xcc, 0x81, 0x0d, 0xcf, 0xf7, 0x79, 0xb2, 0x38, 0x68, 0xf7, 0x7b, + 0x0c, 0xba, 0x6e, 0xfb, 0x96, 0xa1, 0x3f, 0xaf, 0xfd, 0xc9, 0x9f, 0x5e, 0x7f, 0x67, 0xf8, 0x1c, + 0xda, 0xbb, 0x22, 0x8e, 0xb9, 0xaf, 0x44, 0xca, 0xae, 0x43, 0xc7, 0x8e, 0xe3, 0x9a, 0x65, 0xad, + 0x3b, 0x60, 0x41, 0x7b, 0x01, 0xfb, 0x18, 0x56, 0x7c, 0x4b, 0xed, 0x86, 0x71, 0xc0, 0x4f, 0x69, + 0x5d, 0xeb, 0xce, 0x72, 0x0e, 0xde, 0x43, 0xe8, 0xf0, 0xdf, 0x56, 0xa1, 0x79, 0x38, 0xc9, 0x46, + 0xa3, 0x88, 0xb3, 0x8f, 0xa0, 0x67, 0x1e, 0x77, 0x45, 0xb4, 0x17, 0x9c, 0x9a, 0x71, 0x67, 0x81, + 0xec, 0x06, 0x74, 0x0c, 0xe0, 0xe8, 0x2c, 0xe1, 0x66, 0xd8, 0x32, 0x68, 0x76, 0x9c, 0xfd, 0x30, + 0x26, 0x71, 0xa9, 0x3a, 0xb3, 0xc0, 0x39, 0x2a, 0xef, 0x94, 0x24, 0x68, 0x96, 0xca, 0xa3, 0xb7, + 0xed, 0x44, 0xe1, 0x4b, 0xee, 0xf0, 0xf1, 0x6e, 0xac, 0x48, 0x8e, 0xea, 0x4e, 0x19, 0xc4, 0xb6, + 0xe1, 0x92, 0xd4, 0x5d, 0xdc, 0xd4, 0x8b, 0xc7, 0x5c, 0xba, 0x59, 0x18, 0xab, 0x9f, 0x7e, 0x3e, + 0x68, 0xdc, 0xa8, 0x6e, 0xd6, 0x9c, 0x35, 0x83, 0x74, 0x08, 0xf7, 0x9c, 0x50, 0xec, 0x1e, 0xac, + 0xcf, 0xf5, 0xd1, 0x5d, 0x9a, 0x37, 0xaa, 0x9b, 0x55, 0x87, 0xcd, 0x74, 0xd9, 0xa3, 0x1e, 0x8f, + 0x60, 0x35, 0xcd, 0x62, 0xdc, 0x6d, 0x8f, 0xc3, 0x48, 0xf1, 0xf4, 0x30, 0xe1, 0x3e, 0xc9, 0x63, + 0x67, 0xfb, 0xf2, 0x16, 0x6d, 0x48, 0x67, 0x1e, 0xed, 0x2c, 0xf6, 0x60, 0x77, 0x72, 0xe6, 0x3d, + 0x3a, 0x4d, 0x52, 0x12, 0xda, 0xce, 0x36, 0xe8, 0x01, 0x10, 0xe2, 0x94, 0xd1, 0xec, 0x36, 0xac, + 0x06, 0xa9, 0x17, 0xc6, 0xae, 0x17, 0x45, 0xee, 0x71, 0xe6, 0x9f, 0x70, 0x25, 0x49, 0x90, 0x5b, + 0xce, 0x0a, 0x21, 0x76, 0xa2, 0xe8, 0x81, 0x06, 0x0f, 0xff, 0x6a, 0x09, 0x5a, 0x0f, 0x43, 0x99, + 0x78, 0xca, 0x9f, 0xb0, 0xcb, 0xd0, 0x1c, 0x65, 0xb1, 0x5f, 0xc8, 0x46, 0x03, 0x9b, 0x7b, 0x01, + 0xfb, 0x7d, 0x58, 0x89, 0x84, 0xef, 0x45, 0x6e, 0x2e, 0x06, 0x83, 0xa5, 0x1b, 0xd5, 0xcd, 0xce, + 0xf6, 0x5a, 0x21, 0x98, 0xb9, 0x98, 0x39, 0xcb, 0x44, 0x5b, 0x88, 0xdd, 0x2f, 0xa1, 0x9f, 0xf2, + 0xa9, 0x50, 0xbc, 0xd4, 0xbd, 0x4a, 0xdd, 0x59, 0xd1, 0xfd, 0x9b, 0xd4, 0x4b, 0x9e, 0xa2, 0x34, + 0xaf, 0x68, 0xda, 0xa2, 0xfb, 0xfd, 0xd2, 0x4a, 0xf1, 0xb1, 0x1b, 0x06, 0xa7, 0x2e, 0xbd, 0x60, + 0x50, 0xbb, 0x51, 0xdd, 0xac, 0x17, 0x6c, 0xe7, 0xe3, 0xbd, 0xe0, 0xf4, 0x09, 0x62, 0xd8, 0x67, + 0xb0, 0x31, 0xdf, 0x45, 0x8f, 0x3a, 0xa8, 0x53, 0x9f, 0xb5, 0x99, 0x3e, 0x0e, 0xa1, 0xd8, 0x07, + 0xd0, 0xb5, 0x9d, 0x14, 0x8a, 0x68, 0x43, 0x0b, 0x8d, 0x2c, 0x89, 0xe8, 0x65, 0x68, 0x86, 0xd2, + 0x95, 0x61, 0x7c, 0x42, 0x6a, 0xa6, 0xe5, 0x34, 0x42, 0x79, 0x18, 0xc6, 0x27, 0xec, 0x0a, 0xb4, + 0x52, 0xee, 0x6b, 0x4c, 0x8b, 0x30, 0xcd, 0x94, 0xfb, 0x84, 0xba, 0x0c, 0xf8, 0xe8, 0xfa, 0x8a, + 0x1b, 0x65, 0xd3, 0x48, 0xb9, 0xbf, 0xab, 0xf8, 0x50, 0x42, 0x7d, 0x9f, 0xa7, 0x63, 0x8e, 0xfa, + 0x06, 0x3b, 0x1e, 0xfa, 0x5e, 0x4c, 0x7c, 0x6f, 0x39, 0x79, 0x1b, 0xb5, 0x5d, 0xe2, 0xa5, 0x2a, + 0xf4, 0x22, 0xda, 0x32, 0x2d, 0xc7, 0x36, 0xd9, 0x35, 0x68, 0x4b, 0xe5, 0xa5, 0x0a, 0xbf, 0x8e, + 0xb6, 0x4a, 0xdd, 0x69, 0x11, 0x00, 0x77, 0xdb, 0x65, 0x68, 0xf2, 0x38, 0x20, 0x54, 0x4d, 0xaf, + 0x24, 0x8f, 0x83, 0xbd, 0xe0, 0x74, 0xf8, 0xaf, 0x2b, 0xd0, 0xdb, 0xcf, 0x22, 0x15, 0xee, 0xa4, + 0xe3, 0x8c, 0x4f, 0x63, 0x85, 0x5a, 0xf2, 0x61, 0x28, 0x95, 0x79, 0x33, 0x3d, 0xb3, 0x4d, 0x68, + 0x7f, 0x95, 0x8a, 0x2c, 0x21, 0x69, 0xd3, 0x2b, 0x5d, 0x96, 0xb6, 0x02, 0x89, 0x92, 0xf9, 0x2c, + 0x0d, 0x78, 0xfa, 0xe0, 0x8c, 0x68, 0xab, 0x0b, 0xb4, 0x65, 0x34, 0x7b, 0x17, 0xda, 0x87, 0x3c, + 0xf1, 0x52, 0x0f, 0x45, 0xa0, 0x46, 0xaa, 0xb5, 0x00, 0xe0, 0xb7, 0x12, 0xf1, 0x5e, 0x60, 0x36, + 0xac, 0x6d, 0x0e, 0xff, 0x49, 0x05, 0xda, 0x3b, 0xe3, 0x71, 0xca, 0xc7, 0x9e, 0x22, 0x3d, 0x2f, + 0x12, 0x9a, 0x6f, 0xd5, 0x59, 0x12, 0x09, 0xd9, 0x12, 0xfc, 0x02, 0xcd, 0x20, 0x7a, 0x66, 0xef, + 0x43, 0x8d, 0x9f, 0x3f, 0x21, 0x82, 0xb3, 0x0d, 0x68, 0xf8, 0x22, 0x1e, 0x85, 0x63, 0x63, 0x81, + 0x4c, 0x8b, 0xfd, 0x1c, 0x3a, 0xfa, 0x49, 0xcb, 0x40, 0x9d, 0xd4, 0xef, 0x15, 0xdd, 0x3d, 0x9f, + 0xc1, 0x2e, 0x51, 0xa0, 0x44, 0x38, 0xe0, 0xe7, 0xcf, 0xc3, 0x7f, 0x5e, 0x85, 0x3a, 0x71, 0x06, + 0xd7, 0x06, 0x2d, 0x8a, 0xcb, 0x5f, 0x7a, 0x91, 0x5d, 0x52, 0x04, 0x3c, 0x7a, 0xe9, 0x45, 0xec, + 0x06, 0xd4, 0x71, 0x0a, 0xf2, 0x1c, 0xc6, 0x6a, 0x04, 0xbb, 0x05, 0x75, 0x7c, 0xbb, 0x9c, 0x9d, + 0x3d, 0xbe, 0xe3, 0x41, 0xed, 0xcf, 0xff, 0xf3, 0xf5, 0x77, 0x1c, 0x8d, 0x66, 0x1f, 0x43, 0xcd, + 0x1b, 0x8f, 0x25, 0x6d, 0x84, 0x99, 0xbd, 0x98, 0xcf, 0xd4, 0x21, 0x02, 0xf6, 0x05, 0xb4, 0xf5, + 0xa2, 0x23, 0x75, 0x9d, 0xa8, 0x2f, 0x97, 0x2c, 0x75, 0x59, 0x1e, 0x9c, 0x82, 0x12, 0x97, 0x2b, + 0x94, 0x46, 0xb3, 0xd0, 0x76, 0x68, 0x39, 0x05, 0x00, 0x4d, 0x69, 0x92, 0xf2, 0x9d, 0x28, 0x12, + 0xfe, 0x61, 0xf8, 0x9a, 0x1b, 0xc3, 0x3b, 0x03, 0x63, 0xb7, 0x60, 0xf9, 0x40, 0xcb, 0xab, 0xc3, + 0x65, 0x16, 0x29, 0x69, 0x8c, 0xf1, 0x1c, 0x94, 0x6d, 0x01, 0x9b, 0x81, 0x1c, 0xd1, 0xe7, 0xb7, + 0x6f, 0x54, 0x37, 0x7b, 0xce, 0x39, 0x18, 0xf6, 0x21, 0xf4, 0xc6, 0xc8, 0xe9, 0x30, 0x1e, 0xbb, + 0xa3, 0xc8, 0x43, 0x3b, 0x5d, 0x45, 0x3b, 0x6e, 0x81, 0x8f, 0x23, 0x6f, 0x4c, 0x3b, 0x24, 0x09, + 0xa3, 0xc8, 0x9d, 0xf2, 0x29, 0x59, 0xe7, 0xaa, 0xd3, 0x22, 0xc0, 0x3e, 0x9f, 0x0e, 0xff, 0x45, + 0x0d, 0x1a, 0x7b, 0xb1, 0xe4, 0xa9, 0xc2, 0xfd, 0xe7, 0x8d, 0x46, 0xdc, 0x57, 0x5c, 0xeb, 0xbd, + 0x9a, 0x93, 0xb7, 0x91, 0x05, 0x47, 0xe2, 0x9b, 0x34, 0x54, 0xfc, 0xf0, 0x33, 0x23, 0x60, 0x05, + 0x00, 0x35, 0xad, 0x17, 0x04, 0xae, 0xa5, 0x76, 0x53, 0xf1, 0x4a, 0xd2, 0x5e, 0x6c, 0x39, 0x2b, + 0x5e, 0x10, 0xec, 0x18, 0xb8, 0x23, 0x5e, 0x49, 0xf6, 0x01, 0x54, 0x53, 0x3e, 0x22, 0x71, 0xeb, + 0x6c, 0xaf, 0xe8, 0x25, 0x7d, 0x76, 0xfc, 0x1d, 0xf7, 0x95, 0xc3, 0x47, 0x0e, 0xe2, 0xd8, 0x3a, + 0xd4, 0x3d, 0xa5, 0x52, 0xbd, 0x44, 0x6d, 0x47, 0x37, 0xd8, 0x16, 0xac, 0xd1, 0x9e, 0x57, 0xa1, + 0x88, 0x5d, 0xe5, 0x1d, 0x47, 0x68, 0xbc, 0xa5, 0xb1, 0x53, 0xab, 0x39, 0xea, 0x08, 0x31, 0x7b, + 0x81, 0x44, 0xcb, 0x36, 0x4f, 0x1f, 0x7b, 0x53, 0x2e, 0xc9, 0x4c, 0xb5, 0x9d, 0xb5, 0xd9, 0x1e, + 0x4f, 0x11, 0x85, 0xfc, 0x2c, 0xfa, 0xa0, 0xd6, 0x68, 0xd1, 0x06, 0xec, 0xe6, 0x40, 0x54, 0x2a, + 0x97, 0xa0, 0x11, 0x4a, 0x97, 0xc7, 0x81, 0x51, 0x64, 0xf5, 0x50, 0x3e, 0x8a, 0x03, 0xf6, 0x09, + 0xb4, 0xf5, 0x5b, 0x02, 0x3e, 0x22, 0x33, 0xd3, 0xd9, 0x5e, 0x36, 0x12, 0x8b, 0xe0, 0x87, 0x7c, + 0xe4, 0xb4, 0x94, 0x79, 0x42, 0x17, 0x44, 0x09, 0x97, 0x9f, 0x2a, 0x9e, 0xc6, 0x5e, 0x44, 0xab, + 0xd2, 0x72, 0x40, 0x89, 0x47, 0x06, 0xc2, 0xbe, 0x80, 0xcb, 0x16, 0xeb, 0x4a, 0x35, 0x55, 0x6e, + 0x16, 0x87, 0xa7, 0x6e, 0xec, 0xc5, 0x82, 0x7c, 0xa1, 0xaa, 0xb3, 0x6e, 0xd1, 0x87, 0x6a, 0xaa, + 0x9e, 0xc7, 0xe1, 0xe9, 0x53, 0x2f, 0x16, 0x6c, 0x13, 0xfa, 0x79, 0x37, 0xf5, 0x9a, 0x3e, 0x78, + 0xd0, 0x23, 0x05, 0xb3, 0x6c, 0xe1, 0x47, 0xaf, 0xf1, 0x5b, 0xd1, 0x36, 0x94, 0x29, 0xc5, 0x68, + 0x24, 0xb9, 0x72, 0x25, 0xf7, 0x07, 0xcb, 0xf4, 0xcd, 0x6b, 0x05, 0xfd, 0x33, 0xc2, 0x1d, 0x72, + 0x7f, 0xf8, 0xdb, 0x0a, 0x74, 0x68, 0x5f, 0x3c, 0x4f, 0x02, 0x54, 0x41, 0x1f, 0x42, 0x6f, 0x76, + 0xd1, 0xb5, 0xdc, 0x74, 0xbd, 0xf2, 0x8a, 0x6f, 0x40, 0x63, 0xc7, 0x47, 0xe6, 0x91, 0xe0, 0xf4, + 0x1c, 0xd3, 0x62, 0x3f, 0x83, 0x95, 0x8c, 0x86, 0x71, 0x7d, 0x75, 0xea, 0x46, 0xa8, 0xba, 0xf4, + 0x46, 0x37, 0x52, 0xa1, 0xdf, 0xb1, 0xab, 0x4e, 0x9d, 0x5e, 0x66, 0x1f, 0x9f, 0xa0, 0x52, 0xbb, + 0x07, 0xeb, 0x29, 0x47, 0x89, 0x71, 0x5f, 0xf3, 0x54, 0xb8, 0x8a, 0x4f, 0x13, 0x91, 0x92, 0x21, + 0x44, 0x2e, 0x32, 0x8d, 0xfb, 0x96, 0xa7, 0xe2, 0xc8, 0x60, 0x86, 0xef, 0x41, 0x7d, 0x27, 0x4d, + 0xbd, 0x33, 0x12, 0x2d, 0x7c, 0x18, 0x54, 0xc8, 0x00, 0xea, 0xc6, 0xd0, 0x87, 0xea, 0xbe, 0x97, + 0xb0, 0x9b, 0xb0, 0x34, 0x4d, 0x08, 0xd3, 0xd9, 0xbe, 0x54, 0xd2, 0x0b, 0x5e, 0xb2, 0xb5, 0x9f, + 0x3c, 0x8a, 0x55, 0x7a, 0xe6, 0x2c, 0x4d, 0x93, 0xab, 0x5f, 0x40, 0xd3, 0x34, 0xd1, 0xa1, 0x3f, + 0xe1, 0x67, 0xf4, 0xd5, 0x6d, 0x07, 0x1f, 0xf1, 0x05, 0x2f, 0xbd, 0x28, 0xb3, 0x9e, 0x9d, 0x6e, + 0xfc, 0x7c, 0xe9, 0xcb, 0xca, 0xf0, 0x7f, 0xd6, 0xa0, 0xf5, 0x90, 0x47, 0x9c, 0xbe, 0x7d, 0x08, + 0xdd, 0xf2, 0xae, 0xb0, 0x7c, 0x9b, 0xd9, 0x29, 0x43, 0xe8, 0x6a, 0x93, 0x4c, 0xbd, 0xb8, 0xd9, + 0x76, 0x33, 0x30, 0xb4, 0x15, 0x7b, 0xda, 0x87, 0xa1, 0xfd, 0xd6, 0x73, 0x6c, 0x13, 0x31, 0x4f, + 0x0d, 0xa6, 0xa6, 0x31, 0xa6, 0xc9, 0xde, 0x05, 0x48, 0xc5, 0x2b, 0x37, 0xd4, 0x76, 0x51, 0x9b, + 0x98, 0x56, 0x2a, 0x5e, 0xed, 0xa1, 0x65, 0xfc, 0x6b, 0xd9, 0x66, 0x3f, 0x83, 0x41, 0x69, 0x9b, + 0xa1, 0x27, 0xed, 0x86, 0xb1, 0x7b, 0x8c, 0xce, 0x97, 0xd9, 0x71, 0xc5, 0x98, 0xe4, 0x68, 0xef, + 0xc5, 0x0f, 0xc8, 0x33, 0x33, 0xca, 0xa3, 0xfd, 0x06, 0xe5, 0x71, 0xae, 0x2e, 0x82, 0xf3, 0x75, + 0xd1, 0x03, 0x80, 0x43, 0x3e, 0x9e, 0xf2, 0x58, 0xed, 0x7b, 0xc9, 0xa0, 0x43, 0x0b, 0x3f, 0x2c, + 0x16, 0xde, 0xae, 0xd6, 0x56, 0x41, 0xa4, 0xa5, 0xa0, 0xd4, 0x0b, 0xdd, 0x25, 0xdf, 0x8b, 0x5d, + 0x95, 0x66, 0xb1, 0xef, 0x29, 0x7d, 0x52, 0x69, 0x39, 0x1d, 0xdf, 0x8b, 0x8f, 0x0c, 0xa8, 0xa4, + 0x30, 0x7a, 0x65, 0x85, 0x71, 0x0b, 0x56, 0x92, 0x34, 0x9c, 0x7a, 0xe9, 0x99, 0x7b, 0xc2, 0xcf, + 0x68, 0x31, 0xf4, 0xd6, 0xeb, 0x19, 0xf0, 0x6f, 0xf8, 0xd9, 0x5e, 0x70, 0x7a, 0xf5, 0x97, 0xb0, + 0x32, 0x37, 0x81, 0x1f, 0x24, 0x77, 0xff, 0xb1, 0x0a, 0xed, 0x83, 0x94, 0x1b, 0x25, 0x7f, 0x1d, + 0x3a, 0xd2, 0x9f, 0xf0, 0xa9, 0xa7, 0x75, 0x83, 0x1e, 0x01, 0x34, 0x88, 0xf4, 0xc2, 0x8c, 0x1a, + 0x5b, 0x7a, 0x8b, 0x1a, 0xeb, 0x43, 0x55, 0xbb, 0x5d, 0xb8, 0x99, 0xf0, 0xb1, 0xd0, 0xdd, 0xb5, + 0xb2, 0xee, 0xbe, 0x01, 0xdd, 0x89, 0x27, 0x5d, 0x2f, 0x53, 0xc2, 0xf5, 0x45, 0x44, 0x42, 0xd7, + 0x72, 0x60, 0xe2, 0xc9, 0x9d, 0x4c, 0x89, 0x5d, 0x11, 0xb1, 0xf7, 0x00, 0x7c, 0x11, 0x19, 0x35, + 0x64, 0x7c, 0xce, 0xb6, 0x2f, 0x22, 0xad, 0x7b, 0x50, 0x2a, 0xb9, 0x54, 0xe1, 0xd4, 0x33, 0x4b, + 0xea, 0xfa, 0x22, 0x8b, 0x15, 0xd9, 0xda, 0xaa, 0xb3, 0x9a, 0xa3, 0x1c, 0xf1, 0x6a, 0x17, 0x11, + 0xec, 0x1e, 0x2c, 0xfb, 0x62, 0x9a, 0xb8, 0x09, 0x72, 0x96, 0x3c, 0xa0, 0xd6, 0xc2, 0x61, 0xa1, + 0x8b, 0x14, 0x07, 0x27, 0x5c, 0xfb, 0x64, 0xdb, 0xb0, 0xe2, 0x47, 0x99, 0x54, 0x3c, 0x75, 0x8f, + 0x4d, 0x97, 0xc5, 0xf3, 0x45, 0xcf, 0x90, 0x18, 0x3f, 0x6e, 0x08, 0xbd, 0x50, 0xba, 0x22, 0x0a, + 0x5c, 0xad, 0xa0, 0x8c, 0x9c, 0x75, 0x42, 0xf9, 0x2c, 0x0a, 0x8c, 0x8a, 0xd4, 0x34, 0x31, 0x7f, + 0x65, 0x69, 0x3a, 0x96, 0xe6, 0x29, 0x7f, 0x65, 0x68, 0x2e, 0x52, 0x68, 0xdd, 0x0b, 0x15, 0xda, + 0x7f, 0x5a, 0x82, 0xe6, 0x81, 0x90, 0xea, 0xe1, 0x34, 0xb2, 0x9b, 0xa2, 0xf2, 0x43, 0x37, 0xc5, + 0xd2, 0xf9, 0x9b, 0xe2, 0x1c, 0xb1, 0xac, 0x9e, 0x23, 0x96, 0x68, 0x6a, 0xca, 0x74, 0x24, 0x4e, + 0xda, 0x97, 0x5d, 0x2e, 0x08, 0x49, 0xa4, 0xae, 0xa1, 0xff, 0xe4, 0x06, 0x5a, 0x8b, 0xe9, 0xa5, + 0x6f, 0x85, 0xd2, 0x68, 0x30, 0x8d, 0x0c, 0x49, 0x3a, 0x8d, 0x73, 0xd5, 0x0a, 0xa5, 0x91, 0xd6, + 0xdf, 0x83, 0x2b, 0x79, 0x4f, 0xf7, 0x55, 0xa8, 0x26, 0x22, 0x53, 0xee, 0x88, 0x0e, 0x84, 0xd2, + 0x1c, 0x3d, 0x36, 0xec, 0x48, 0xdf, 0x68, 0xb4, 0x3e, 0x2e, 0x92, 0xaf, 0x37, 0xca, 0xa2, 0xc8, + 0x55, 0xfc, 0x54, 0x99, 0xc5, 0x1f, 0x68, 0xde, 0x18, 0xbe, 0x3d, 0xce, 0xa2, 0xe8, 0x88, 0x9f, + 0x2a, 0x34, 0x30, 0xad, 0x91, 0x69, 0x0c, 0xff, 0xb8, 0x06, 0xf0, 0x44, 0xf8, 0x27, 0x47, 0x5e, + 0x3a, 0xe6, 0x0a, 0x0f, 0x34, 0x56, 0x07, 0x1a, 0x1d, 0xdd, 0x54, 0x5a, 0xf3, 0xb1, 0x6d, 0xd8, + 0xb0, 0xdf, 0x8f, 0x92, 0x8b, 0x87, 0x2b, 0xad, 0xc4, 0xcc, 0x16, 0x64, 0x06, 0xab, 0x0f, 0xfe, + 0xa4, 0xc1, 0xd8, 0x97, 0x05, 0x6f, 0xb1, 0x8f, 0x3a, 0x4b, 0x88, 0xb7, 0xe7, 0xf9, 0xb6, 0xbd, + 0xa2, 0xfb, 0xd1, 0x59, 0xc2, 0xee, 0xc1, 0xa5, 0x94, 0x8f, 0x52, 0x2e, 0x27, 0xae, 0x92, 0xe5, + 0x97, 0xe9, 0x73, 0xcd, 0xaa, 0x41, 0x1e, 0xc9, 0xfc, 0x5d, 0xf7, 0xe0, 0x92, 0xe6, 0xd4, 0xfc, + 0xf4, 0xb4, 0xc6, 0x5f, 0xd5, 0xc8, 0xf2, 0xec, 0xde, 0x03, 0x0a, 0x94, 0x69, 0x2d, 0x6e, 0x1d, + 0xdd, 0x88, 0x98, 0x71, 0x1c, 0x71, 0xf4, 0x01, 0x77, 0x27, 0x78, 0xa8, 0x7f, 0xc8, 0x47, 0x86, + 0xf9, 0x05, 0x80, 0x0d, 0xa1, 0xb6, 0x2f, 0x02, 0x4e, 0xac, 0x5e, 0xde, 0x5e, 0xde, 0xa2, 0x90, + 0x1b, 0x72, 0x92, 0x62, 0x33, 0x84, 0x63, 0x1f, 0x03, 0x0d, 0xa7, 0xc5, 0x6f, 0x71, 0x77, 0xb5, + 0x10, 0x49, 0x32, 0x78, 0x0f, 0x2e, 0x15, 0x33, 0x71, 0x3d, 0xe5, 0xaa, 0x09, 0x27, 0x05, 0xaa, + 0x37, 0xd8, 0x6a, 0x3e, 0xa9, 0x1d, 0x75, 0x34, 0xe1, 0xa8, 0x4c, 0x37, 0xa1, 0x29, 0x8e, 0xbf, + 0x73, 0x71, 0x23, 0x74, 0xce, 0xdf, 0x08, 0x0d, 0x71, 0xfc, 0x9d, 0xc3, 0x47, 0xec, 0xa7, 0x65, + 0xe3, 0x33, 0xc7, 0x9a, 0x2e, 0xb1, 0x66, 0x3d, 0xc7, 0x97, 0xb8, 0x33, 0xfc, 0x12, 0x1a, 0xf8, + 0x39, 0xcf, 0x12, 0xb6, 0x05, 0x4d, 0x45, 0xe2, 0x21, 0x8d, 0xb3, 0xb0, 0x5e, 0xd8, 0x8c, 0x42, + 0x76, 0x1c, 0x4b, 0x34, 0x74, 0x60, 0x25, 0x57, 0xc0, 0xcf, 0xe3, 0xf0, 0x45, 0xc6, 0xd9, 0xaf, + 0x60, 0x35, 0x49, 0xb9, 0x11, 0x7b, 0x37, 0x3b, 0x41, 0x17, 0xc8, 0xec, 0xe0, 0x75, 0x23, 0xa5, + 0x79, 0x8f, 0x13, 0x94, 0xd0, 0xe5, 0x64, 0xa6, 0x3d, 0xfc, 0x16, 0x2e, 0xe7, 0x14, 0x87, 0xdc, + 0x17, 0x71, 0xe0, 0xa5, 0x67, 0x64, 0x2b, 0xe7, 0xc6, 0x96, 0x3f, 0x64, 0xec, 0x43, 0x1a, 0xfb, + 0xbf, 0x57, 0xa0, 0xf3, 0x38, 0x7b, 0xfd, 0xfa, 0x4c, 0xef, 0x25, 0xd6, 0x85, 0xca, 0x53, 0x1a, + 0x60, 0xc9, 0xa9, 0x3c, 0x45, 0x77, 0xee, 0xe0, 0x04, 0xf7, 0x35, 0xc9, 0x79, 0xdb, 0x31, 0x2d, + 0x3c, 0xad, 0x1d, 0x9c, 0x1c, 0xbd, 0x41, 0xa2, 0x35, 0x1a, 0x8f, 0x19, 0x0f, 0xb2, 0x30, 0x42, + 0x67, 0xc3, 0x08, 0x6f, 0xde, 0xc6, 0xf3, 0xcf, 0xde, 0x48, 0x4f, 0xe5, 0x71, 0x2a, 0xa6, 0x9a, + 0x59, 0x46, 0x65, 0x9c, 0x83, 0x61, 0x5f, 0xc1, 0x9a, 0x89, 0x12, 0x19, 0xad, 0xe0, 0xca, 0x84, + 0xfb, 0x24, 0xba, 0x3f, 0x28, 0xb2, 0x34, 0xfc, 0xab, 0x1a, 0xb4, 0xbe, 0xf6, 0xe4, 0xe4, 0xd7, + 0x22, 0x8c, 0xd9, 0x3d, 0x68, 0x7f, 0x27, 0xc2, 0x58, 0x1f, 0x7d, 0x75, 0xd0, 0x77, 0x4d, 0x8f, + 0xf5, 0x54, 0x04, 0x7c, 0x0b, 0x69, 0xe8, 0xd0, 0xdb, 0xfa, 0xce, 0x3c, 0x19, 0x25, 0x9f, 0x86, + 0xe3, 0x89, 0x72, 0x11, 0x68, 0x74, 0x6b, 0x27, 0x94, 0x0e, 0xc2, 0x68, 0xd4, 0x77, 0x01, 0xed, + 0xdd, 0xc4, 0x15, 0xb1, 0x9b, 0x9c, 0x98, 0xd3, 0x51, 0x0b, 0x21, 0xcf, 0xe2, 0x83, 0x13, 0xdc, + 0x7b, 0xa1, 0x74, 0x4d, 0x90, 0xc5, 0x78, 0xb2, 0xa5, 0x43, 0xe6, 0x47, 0xb0, 0x8c, 0x5e, 0x86, + 0x3c, 0x09, 0x13, 0x37, 0x49, 0xc5, 0xb1, 0x65, 0x0a, 0xfa, 0x1e, 0x87, 0x27, 0x61, 0x72, 0x80, + 0x30, 0x32, 0xee, 0x26, 0x74, 0x83, 0x6a, 0x5b, 0x5b, 0x51, 0x30, 0x20, 0xe4, 0x2f, 0xc5, 0x67, + 0x22, 0xed, 0x6b, 0x37, 0xc9, 0x68, 0x37, 0x53, 0x1e, 0x91, 0x53, 0x7d, 0x05, 0x5a, 0xb8, 0x19, + 0x08, 0xd5, 0xd2, 0x28, 0x5f, 0x68, 0xd4, 0x4f, 0x00, 0x22, 0x3e, 0x52, 0x2e, 0x4a, 0x99, 0x3e, + 0x8d, 0xce, 0xc5, 0x41, 0x10, 0xbb, 0x8b, 0x48, 0xf6, 0x09, 0x74, 0x34, 0x17, 0x34, 0x2d, 0x2c, + 0xd0, 0x02, 0xa1, 0x35, 0xf1, 0x6d, 0xe8, 0xc4, 0x22, 0x76, 0xf9, 0x0b, 0xa2, 0x36, 0xfb, 0x76, + 0x66, 0xe0, 0x58, 0xc4, 0x8f, 0x5e, 0x20, 0x31, 0xbb, 0x6b, 0xe6, 0xa0, 0x03, 0x02, 0xdd, 0x0b, + 0x02, 0x02, 0x34, 0x13, 0x7d, 0x34, 0xbe, 0x6f, 0x67, 0xa2, 0x7b, 0xf4, 0x2e, 0xe8, 0xa1, 0xe7, + 0xa3, 0xbb, 0xdc, 0x80, 0x2e, 0xad, 0xfb, 0xd4, 0x4b, 0x5c, 0xe5, 0x8d, 0x8d, 0x37, 0x06, 0x08, + 0xdb, 0xf7, 0x92, 0x23, 0x6f, 0xcc, 0x1c, 0xb8, 0x32, 0x27, 0x6f, 0xc7, 0x28, 0xba, 0x9a, 0x6b, + 0x2b, 0x36, 0xa0, 0x70, 0xbe, 0xd4, 0x6d, 0xcc, 0x48, 0x1d, 0x89, 0x3c, 0x72, 0x77, 0xf8, 0x4f, + 0x97, 0xa0, 0xf5, 0x44, 0x88, 0xe4, 0x47, 0x8a, 0x5e, 0x79, 0x49, 0x97, 0x2e, 0x5e, 0xd2, 0xea, + 0xec, 0x92, 0xce, 0xb1, 0xbe, 0xf6, 0xfd, 0x59, 0x5f, 0xff, 0xc1, 0xac, 0x6f, 0xfc, 0x08, 0xd6, + 0x37, 0xe7, 0x59, 0x3f, 0x6c, 0x42, 0xfd, 0x90, 0xab, 0x67, 0xc9, 0xf0, 0x5f, 0xb5, 0xa0, 0xfd, + 0x90, 0x07, 0x99, 0x66, 0x58, 0xf9, 0xf3, 0x2b, 0x17, 0x7f, 0xfe, 0xd2, 0xec, 0xe7, 0xa3, 0x21, + 0xb2, 0x12, 0x7d, 0x4e, 0x6c, 0xac, 0x65, 0x05, 0x1a, 0x45, 0xbf, 0x90, 0x67, 0x13, 0x60, 0x9a, + 0x61, 0x53, 0x2e, 0xce, 0x6f, 0x96, 0x8d, 0xfa, 0x8f, 0x92, 0x8d, 0x39, 0xad, 0xb0, 0x10, 0x7a, + 0x7a, 0x2b, 0xd7, 0xe6, 0x35, 0x42, 0x6b, 0x41, 0x23, 0x3c, 0x81, 0x35, 0x11, 0xbb, 0x41, 0x96, + 0x44, 0x21, 0x9e, 0x55, 0x5c, 0x4f, 0x9f, 0xd4, 0xdb, 0x36, 0xdf, 0x92, 0x8b, 0xde, 0xb3, 0xf8, + 0xa1, 0x25, 0xd2, 0xe7, 0x77, 0x67, 0x55, 0xcc, 0x83, 0x50, 0x4d, 0x05, 0xb8, 0x34, 0x64, 0x57, + 0xc9, 0x23, 0xd4, 0x89, 0xa3, 0x2e, 0x41, 0x77, 0x45, 0x44, 0x96, 0xe2, 0x4b, 0x58, 0x29, 0xa8, + 0xb4, 0x8c, 0x74, 0x2e, 0x90, 0x91, 0x9e, 0xed, 0xa8, 0xc5, 0xe4, 0xaf, 0x43, 0x0b, 0x7c, 0x0a, + 0x6b, 0x36, 0x2c, 0x61, 0x9c, 0x03, 0x5a, 0xc1, 0x65, 0x92, 0xa0, 0xbe, 0x89, 0x44, 0x90, 0x5f, + 0x40, 0x4b, 0xf4, 0x0b, 0x58, 0x2f, 0x91, 0xe3, 0xb9, 0xa1, 0xac, 0x0d, 0xca, 0xb2, 0xb2, 0x9a, + 0xf7, 0xc5, 0xe6, 0x13, 0x1d, 0x9e, 0xed, 0x04, 0x3c, 0xb2, 0x2f, 0x1a, 0xf4, 0xf5, 0xb1, 0x27, + 0xe0, 0x91, 0xc9, 0x16, 0xed, 0xc3, 0x47, 0x78, 0xba, 0x40, 0xbc, 0xef, 0x25, 0x2a, 0x4b, 0xb9, + 0x9b, 0x44, 0x9e, 0xcf, 0x27, 0x22, 0x0a, 0x78, 0x5a, 0x4c, 0x6e, 0x95, 0x26, 0x77, 0x5d, 0x44, + 0xc1, 0xae, 0x88, 0x76, 0x35, 0xe5, 0x41, 0x41, 0x68, 0xe7, 0xba, 0x03, 0xef, 0x2f, 0x0c, 0x87, + 0x86, 0xa3, 0x18, 0x88, 0xd1, 0x40, 0x57, 0x66, 0x07, 0x42, 0x12, 0x3b, 0xc4, 0x7d, 0xb8, 0xa4, + 0xd7, 0x4e, 0x0b, 0xf7, 0x09, 0xe7, 0x89, 0x1b, 0x79, 0x52, 0x0d, 0xd6, 0xb4, 0x91, 0x26, 0x24, + 0x09, 0xf0, 0x6f, 0x38, 0x4f, 0x9e, 0x78, 0xfa, 0xad, 0xba, 0x8b, 0xf1, 0xe3, 0xa9, 0xcf, 0x0c, + 0x6f, 0xd7, 0xf5, 0x5b, 0x89, 0x4a, 0x3b, 0xf3, 0xd8, 0xb9, 0xc4, 0xe4, 0xdf, 0x87, 0x6b, 0x33, + 0x43, 0x4c, 0xbd, 0xf4, 0xa4, 0x70, 0x6c, 0x07, 0x97, 0x88, 0x6f, 0x97, 0x4b, 0xfd, 0xf7, 0x89, + 0x40, 0x8f, 0x30, 0xfc, 0x6f, 0x75, 0x58, 0x26, 0x3b, 0xfc, 0xb7, 0x6a, 0xe3, 0x6f, 0xd5, 0xc6, + 0xdf, 0x00, 0xb5, 0x31, 0xfc, 0x07, 0x15, 0x68, 0x1e, 0xa4, 0x22, 0xc8, 0x7c, 0xf5, 0x23, 0x25, + 0x7d, 0x56, 0x82, 0xaa, 0x6f, 0x93, 0xa0, 0xda, 0x82, 0xb9, 0xfe, 0x67, 0x15, 0x68, 0x9b, 0x29, + 0x3c, 0xd9, 0xfe, 0x91, 0x93, 0x28, 0x92, 0x57, 0x95, 0x73, 0x93, 0x57, 0x6f, 0x9d, 0x05, 0x0a, + 0xd6, 0x4b, 0x9d, 0xc5, 0x17, 0x49, 0x91, 0xc9, 0x6a, 0x3b, 0x5d, 0x0d, 0x7d, 0x96, 0x50, 0xc2, + 0xea, 0x15, 0xb4, 0xe9, 0xe4, 0x44, 0x9a, 0x61, 0x03, 0x1a, 0x29, 0x65, 0x58, 0xcc, 0x44, 0x4d, + 0xeb, 0xcd, 0xfb, 0x74, 0xe9, 0xc7, 0xb9, 0x7e, 0xff, 0x6e, 0x09, 0x7a, 0x74, 0x8c, 0x7d, 0x9c, + 0xc5, 0x7a, 0x27, 0xe4, 0xe1, 0xb3, 0xca, 0x6c, 0xf8, 0xac, 0x96, 0xe2, 0x69, 0x53, 0xbf, 0xa6, + 0xab, 0x5f, 0xb3, 0x2b, 0xa2, 0x87, 0x7c, 0xe4, 0x10, 0x06, 0x59, 0xe5, 0xa5, 0x63, 0x79, 0x5e, + 0x9e, 0x0f, 0xe1, 0xf8, 0x55, 0x89, 0x97, 0x7a, 0x53, 0x69, 0xf3, 0x7c, 0xba, 0xc5, 0x18, 0xd4, + 0x68, 0xbf, 0x69, 0xb6, 0xd0, 0xb3, 0x89, 0xc8, 0xc8, 0x30, 0x1e, 0xe7, 0xca, 0xa3, 0x45, 0xf9, + 0xdd, 0x71, 0xc4, 0xd9, 0x43, 0x60, 0x3a, 0x60, 0x9b, 0x72, 0x0f, 0x4d, 0x10, 0x8d, 0x43, 0x1a, + 0xa4, 0xb3, 0xbd, 0xa1, 0x5f, 0x4b, 0xbc, 0x74, 0x08, 0x7d, 0x80, 0x58, 0xa7, 0x1f, 0xce, 0x41, + 0xce, 0x61, 0xa6, 0xb6, 0x43, 0xf9, 0xe9, 0xe3, 0x7b, 0x33, 0x93, 0x8c, 0x13, 0x31, 0x73, 0x07, + 0x2e, 0xd9, 0xec, 0x09, 0xaa, 0x8b, 0x6d, 0xdc, 0x0b, 0x74, 0x1e, 0xb6, 0xdf, 0x58, 0x29, 0x7d, + 0xe3, 0x3a, 0xd4, 0xcb, 0x75, 0x1d, 0xba, 0x31, 0xbc, 0x09, 0x9d, 0x51, 0x18, 0x71, 0x13, 0x85, + 0x44, 0xa6, 0x99, 0x78, 0x64, 0x85, 0x2a, 0x1b, 0x4c, 0x6b, 0xf8, 0xdb, 0x0a, 0x5c, 0x4e, 0xbc, + 0xf4, 0x45, 0xc6, 0x15, 0xc5, 0x22, 0x29, 0xdb, 0xe6, 0xca, 0x89, 0x97, 0x06, 0xb8, 0x71, 0x68, + 0x08, 0x3d, 0xba, 0x2e, 0x1f, 0x68, 0x23, 0x44, 0xcf, 0xe5, 0x16, 0xac, 0x94, 0x7a, 0x28, 0x2f, + 0xb5, 0xd1, 0xa2, 0x5e, 0x2a, 0x5e, 0x51, 0xd2, 0xf4, 0x10, 0x81, 0x78, 0xa0, 0x2c, 0xe8, 0x38, + 0x59, 0x1b, 0xca, 0xc2, 0x5b, 0xaa, 0x47, 0x71, 0x80, 0x3b, 0x27, 0xce, 0xa6, 0x3a, 0x98, 0xa2, + 0xab, 0x3f, 0x9a, 0x71, 0x36, 0xa5, 0xf8, 0xc9, 0x3a, 0xd4, 0x8f, 0xcf, 0x14, 0x79, 0xeb, 0x08, + 0xd7, 0x8d, 0xe1, 0x5f, 0xd4, 0x61, 0x6d, 0xcf, 0xe7, 0xc7, 0x3c, 0x1d, 0x3f, 0xf4, 0x94, 0xf7, + 0x38, 0x8c, 0xf8, 0x91, 0x27, 0x4f, 0x70, 0xc1, 0x69, 0xce, 0x89, 0xa7, 0x26, 0x86, 0x4b, 0x2d, + 0x04, 0x1c, 0x78, 0x6a, 0x82, 0xa6, 0x80, 0x90, 0x23, 0x91, 0x4e, 0x4d, 0x6c, 0xab, 0xed, 0xd0, + 0x37, 0x3e, 0x26, 0x48, 0xde, 0x5b, 0x86, 0xaf, 0xb9, 0xa9, 0x55, 0xa1, 0xde, 0x94, 0xf8, 0xfc, + 0x00, 0xba, 0x29, 0xf7, 0x45, 0x1a, 0x98, 0x80, 0xad, 0x9e, 0x67, 0x47, 0xc3, 0x74, 0xa8, 0xf6, + 0x36, 0x14, 0x59, 0x05, 0x3a, 0xbe, 0xbb, 0xa1, 0x4d, 0x7c, 0xaf, 0xe4, 0x08, 0x5c, 0xf9, 0xbd, + 0x80, 0xfd, 0x5d, 0xe8, 0x17, 0xb4, 0x14, 0xe2, 0xb6, 0xc7, 0x8b, 0xed, 0x22, 0x04, 0x73, 0xce, + 0x27, 0x6e, 0x1d, 0xd8, 0x5e, 0x7f, 0x87, 0x3a, 0xe9, 0x30, 0x7e, 0x31, 0xbc, 0x86, 0xb2, 0x0f, + 0xa1, 0x27, 0x93, 0x28, 0x54, 0x46, 0x00, 0xa4, 0xa9, 0x68, 0xe9, 0x12, 0x50, 0x47, 0xa2, 0xe5, + 0x79, 0x4b, 0xd8, 0xfa, 0x5e, 0x4b, 0xd8, 0x5e, 0x5c, 0xc2, 0x9f, 0x40, 0xdf, 0x4f, 0x79, 0xc0, + 0x63, 0x15, 0x7a, 0x91, 0x2b, 0x7d, 0x91, 0x58, 0xd3, 0xb7, 0x52, 0xc0, 0x0f, 0x11, 0xcc, 0x7e, + 0x0a, 0x97, 0x7d, 0x11, 0x2b, 0x1e, 0x2b, 0x57, 0xf2, 0x17, 0x19, 0x8f, 0x7d, 0xee, 0xc6, 0xd9, + 0xf4, 0x98, 0xa7, 0x26, 0xa7, 0x7b, 0xc9, 0xa0, 0x0f, 0x0d, 0xf6, 0x29, 0x21, 0xd9, 0x3d, 0x58, + 0xd7, 0xcb, 0x33, 0xd7, 0x49, 0x67, 0x11, 0x19, 0xad, 0xd4, 0x6c, 0x8f, 0x2d, 0x58, 0x9b, 0x78, + 0xd2, 0x4d, 0xb9, 0x0c, 0x83, 0xcc, 0x8b, 0xcc, 0x0e, 0x35, 0xb9, 0x8b, 0xd5, 0x89, 0x27, 0x1d, + 0x83, 0x31, 0xe1, 0x21, 0x8a, 0x5e, 0xcf, 0xd0, 0xba, 0x13, 0x4f, 0x4e, 0xe8, 0xf8, 0xdc, 0x76, + 0x58, 0x3a, 0x43, 0xfd, 0xb5, 0x27, 0x27, 0x57, 0x1f, 0xc0, 0xfa, 0x79, 0x0b, 0xf2, 0xb6, 0xb4, + 0x46, 0xbb, 0x94, 0xd6, 0x30, 0x75, 0x5d, 0xff, 0x63, 0x09, 0x2e, 0xd9, 0xf5, 0x26, 0xc7, 0x2f, + 0x17, 0xea, 0xeb, 0x64, 0x23, 0xd1, 0x59, 0xcc, 0xcf, 0xd2, 0x6d, 0x07, 0x34, 0x88, 0x0e, 0xce, + 0x9b, 0xd0, 0x37, 0x04, 0x85, 0xf0, 0xeb, 0xb7, 0x2c, 0x07, 0xf9, 0x50, 0xb4, 0x05, 0xe8, 0x03, + 0x47, 0x3c, 0x45, 0x1e, 0x05, 0x54, 0x91, 0x47, 0x5d, 0x48, 0xd8, 0xe9, 0x03, 0x2d, 0xce, 0x8a, + 0x1c, 0xbb, 0x03, 0x8c, 0xbf, 0xc8, 0xbc, 0x28, 0x54, 0x67, 0xee, 0x28, 0xe4, 0x51, 0x40, 0x39, + 0x34, 0x5d, 0xa8, 0xd3, 0xb7, 0x98, 0xc7, 0x88, 0xd8, 0x0b, 0x64, 0x69, 0x26, 0x26, 0x35, 0x93, + 0x6f, 0x00, 0x33, 0x93, 0x43, 0x02, 0xef, 0x05, 0xe7, 0xef, 0x95, 0xc6, 0xf9, 0x7b, 0xe5, 0x63, + 0x58, 0x99, 0x5f, 0x73, 0x9d, 0x2e, 0x59, 0x96, 0xb3, 0xeb, 0x7d, 0x9e, 0x10, 0xb6, 0xce, 0x15, + 0x42, 0xc3, 0xf4, 0xff, 0xb5, 0x04, 0xeb, 0x86, 0xe9, 0xbb, 0x22, 0xca, 0xa6, 0x68, 0x6d, 0x93, + 0x30, 0x1e, 0xa3, 0x41, 0x9e, 0x0a, 0xed, 0x96, 0x94, 0xd4, 0x1f, 0x4c, 0x45, 0xae, 0x8b, 0x37, + 0xa1, 0x1f, 0xea, 0x9e, 0x39, 0x5f, 0x6c, 0x69, 0x9d, 0x81, 0x1b, 0xae, 0xa0, 0x14, 0xca, 0xd8, + 0x4b, 0xe4, 0x44, 0x28, 0x43, 0x4a, 0x4a, 0x5c, 0xf3, 0x7c, 0xd5, 0xa2, 0x88, 0x9a, 0xbc, 0xc3, + 0x3b, 0xc0, 0xfc, 0x2c, 0x4d, 0x71, 0x7f, 0x94, 0xc8, 0x75, 0x42, 0xa2, 0x6f, 0x30, 0x05, 0xf5, + 0x87, 0xd0, 0x9c, 0x8a, 0xc2, 0x23, 0x98, 0x71, 0xee, 0x9c, 0xc6, 0x54, 0x90, 0x84, 0x5c, 0x45, + 0xaf, 0xe5, 0x45, 0x16, 0xa6, 0x3c, 0xb0, 0x76, 0xd0, 0xb6, 0x8d, 0x91, 0x9c, 0x84, 0x41, 0xc0, + 0x63, 0x13, 0x0c, 0x6f, 0x85, 0xf2, 0x6b, 0x6a, 0x53, 0xe5, 0x19, 0x1f, 0x79, 0x59, 0xa4, 0xdc, + 0x38, 0x8b, 0x68, 0x57, 0x44, 0xa6, 0x1e, 0x6a, 0xc5, 0x20, 0x9e, 0x66, 0x11, 0xee, 0x88, 0xc8, + 0x2c, 0x29, 0xd9, 0x12, 0x14, 0x41, 0x77, 0x12, 0xc6, 0x8a, 0x54, 0x45, 0x9b, 0x96, 0x14, 0x11, + 0x28, 0x84, 0x5f, 0x87, 0xb1, 0x1a, 0xfe, 0xd9, 0x12, 0x6c, 0x18, 0xc6, 0x1f, 0x1a, 0x06, 0x18, + 0xfb, 0x48, 0x1e, 0xbb, 0x65, 0x97, 0xc9, 0x55, 0x54, 0x1d, 0xb0, 0xa0, 0x3d, 0x9a, 0x70, 0x21, + 0x5d, 0x4b, 0xa6, 0x4e, 0xca, 0xca, 0xd5, 0x1d, 0x60, 0x0b, 0x72, 0x25, 0x4d, 0xcc, 0xa8, 0x3f, + 0x27, 0x58, 0x92, 0x7d, 0x0e, 0x1b, 0x53, 0xae, 0x3c, 0xda, 0x08, 0x91, 0xf0, 0x3d, 0xea, 0x45, + 0x5b, 0x5e, 0xb3, 0x7b, 0xdd, 0x62, 0x9f, 0x18, 0x24, 0x6e, 0x7a, 0x7c, 0xc7, 0xd4, 0x8b, 0xc3, + 0x11, 0x97, 0x8a, 0xec, 0xbc, 0xee, 0xa1, 0x1d, 0x8f, 0xbe, 0xc5, 0xa0, 0x25, 0x27, 0x6a, 0xf2, + 0x18, 0x47, 0x7a, 0x11, 0x1b, 0x44, 0xd3, 0x4c, 0xf9, 0xc8, 0xac, 0x5d, 0x0f, 0xd7, 0x2a, 0x0e, + 0xe3, 0xb1, 0x2e, 0x0e, 0x6d, 0x6a, 0x9f, 0xce, 0x02, 0xf7, 0x45, 0xc0, 0x87, 0x7f, 0x52, 0xcb, + 0x65, 0xf4, 0xc0, 0xc0, 0x0f, 0x95, 0xa7, 0x24, 0xbb, 0x09, 0xcb, 0xf9, 0xe4, 0xb5, 0x8d, 0xd4, + 0xbc, 0xea, 0x59, 0xe8, 0x03, 0x04, 0xa2, 0xf8, 0xcd, 0xce, 0x56, 0xd3, 0x2e, 0xe9, 0x84, 0x63, + 0x79, 0xba, 0x9a, 0x1e, 0x87, 0xb5, 0xf4, 0x9a, 0xd4, 0x94, 0x6d, 0x5a, 0xa8, 0x26, 0xfb, 0xb4, + 0x60, 0x82, 0x74, 0x25, 0x8f, 0x74, 0xb5, 0x4d, 0x6d, 0x76, 0x54, 0x79, 0x68, 0x10, 0xb8, 0x35, + 0x0b, 0xf2, 0x24, 0xcd, 0x62, 0x1e, 0x18, 0x93, 0xbe, 0x92, 0xc3, 0x0f, 0x08, 0x8c, 0x13, 0xce, + 0x35, 0x53, 0x69, 0xe8, 0x86, 0x1e, 0x3a, 0x30, 0x9a, 0xa9, 0x18, 0x1a, 0x65, 0xb4, 0xa0, 0x37, + 0x63, 0x6b, 0x05, 0xb1, 0x92, 0x53, 0x9b, 0xb1, 0x7f, 0x06, 0x83, 0x9c, 0x56, 0x7f, 0x5d, 0xf1, + 0x82, 0x96, 0x36, 0x3e, 0xb6, 0x0b, 0x7d, 0x66, 0xfe, 0x92, 0xcf, 0x60, 0x63, 0xbe, 0xa3, 0x79, + 0x53, 0x9b, 0xba, 0xad, 0xcd, 0x74, 0x2b, 0xbe, 0x24, 0x5f, 0x5f, 0xdf, 0xf3, 0x27, 0xdc, 0x9d, + 0x84, 0xa6, 0x72, 0xb3, 0xea, 0xac, 0x5a, 0xd4, 0x2e, 0x62, 0xbe, 0x0e, 0x95, 0x3c, 0x87, 0x7e, + 0x1a, 0x4a, 0x69, 0xac, 0xe2, 0x2c, 0xfd, 0x7e, 0x28, 0xe5, 0xf0, 0x1f, 0x03, 0x74, 0xad, 0xa7, + 0x48, 0xc5, 0x85, 0x77, 0xca, 0x4e, 0x77, 0x67, 0xbb, 0x6f, 0xbd, 0x67, 0x24, 0xd9, 0x51, 0x2a, + 0xb5, 0xf9, 0x0b, 0xed, 0x8c, 0xcf, 0xf8, 0x3b, 0x4b, 0xe4, 0x20, 0x14, 0xfe, 0xce, 0x0e, 0xac, + 0x96, 0x3c, 0x48, 0x57, 0x09, 0xe5, 0x45, 0xc6, 0x29, 0x2f, 0x55, 0x94, 0x94, 0x48, 0x9c, 0x15, + 0x6c, 0x68, 0xdf, 0xe2, 0x08, 0xa9, 0xd1, 0xd9, 0xf7, 0x45, 0x64, 0xab, 0xd9, 0xe6, 0x9c, 0x7d, + 0xc4, 0x50, 0xae, 0x3c, 0xe5, 0x78, 0x76, 0x94, 0x2f, 0x22, 0xb3, 0x83, 0xda, 0x1a, 0x72, 0xf8, + 0x22, 0xca, 0x27, 0x48, 0xce, 0x74, 0x83, 0xce, 0x11, 0x34, 0x41, 0x3a, 0x53, 0x7d, 0x0a, 0x1d, + 0x91, 0x86, 0xe3, 0x90, 0x52, 0x5f, 0xda, 0xc1, 0x99, 0x7f, 0x09, 0x68, 0x82, 0x5d, 0x7c, 0xd5, + 0x10, 0x1a, 0xc6, 0xfc, 0x2f, 0xe6, 0xcf, 0x0d, 0x06, 0x1d, 0x22, 0xa9, 0xd2, 0xd0, 0x57, 0x38, + 0x1d, 0xbd, 0x23, 0x75, 0x61, 0x54, 0x4f, 0x83, 0x0f, 0x5f, 0x44, 0x94, 0xfd, 0xbb, 0x05, 0x2b, + 0x3e, 0x99, 0x0b, 0xbd, 0xa1, 0x22, 0x1e, 0xd3, 0x9a, 0xd6, 0x9d, 0x9e, 0x06, 0xe3, 0xfc, 0x9e, + 0xf0, 0xd8, 0x14, 0x61, 0x79, 0x51, 0x84, 0x27, 0x46, 0xe1, 0x05, 0x26, 0x63, 0xde, 0xb5, 0xc0, + 0x27, 0xc2, 0x0b, 0xd8, 0xcf, 0xe1, 0x2a, 0xe2, 0x5c, 0x3e, 0x4d, 0xd4, 0x19, 0xda, 0x37, 0x9e, + 0x86, 0xbe, 0xeb, 0x49, 0xca, 0xa0, 0x9b, 0xc4, 0xf9, 0x06, 0x52, 0x3c, 0x42, 0x82, 0xa7, 0x1a, + 0xbf, 0x23, 0xbf, 0xe5, 0xa9, 0x60, 0xdf, 0x52, 0x06, 0xf0, 0x3c, 0xf7, 0xdd, 0x1e, 0xf5, 0x3f, + 0x28, 0xd6, 0xea, 0x02, 0x4a, 0xaa, 0x50, 0x41, 0x84, 0x63, 0x9d, 0x3e, 0xea, 0xcf, 0x7e, 0x03, + 0xcc, 0x1a, 0x38, 0x92, 0x7c, 0xe5, 0xc9, 0x13, 0x49, 0x51, 0x80, 0xce, 0xf6, 0x7b, 0x6f, 0xf4, + 0x51, 0x1d, 0x6b, 0x19, 0x11, 0x88, 0x00, 0xc9, 0xfe, 0x08, 0xd6, 0xf3, 0xc1, 0x8c, 0x2f, 0x43, + 0xc3, 0xe9, 0x20, 0xc1, 0xf5, 0xc5, 0xe1, 0x66, 0x5c, 0x20, 0xc7, 0xce, 0x44, 0x83, 0xf5, 0x90, + 0x5f, 0xc1, 0x8a, 0x1d, 0x52, 0x73, 0x5d, 0x0e, 0xfa, 0x34, 0xda, 0xfb, 0x0b, 0xa3, 0xcd, 0xd8, + 0xf6, 0xdc, 0x3e, 0x6b, 0x28, 0x7e, 0x68, 0x6e, 0xc9, 0xad, 0x95, 0x19, 0xac, 0x92, 0x8c, 0xdc, + 0x58, 0x18, 0x69, 0xce, 0x58, 0x39, 0x76, 0x0a, 0x16, 0xce, 0xee, 0xc3, 0x25, 0x3b, 0x98, 0xa0, + 0x84, 0xad, 0x1b, 0x0a, 0xca, 0xe5, 0x32, 0xed, 0x62, 0x19, 0xa4, 0x4e, 0xe6, 0xee, 0x09, 0x87, + 0x8f, 0xd8, 0x2f, 0xe1, 0x9a, 0xed, 0xa2, 0xad, 0x30, 0x9d, 0x48, 0xf3, 0x8f, 0x5a, 0x23, 0xdb, + 0x35, 0x30, 0x24, 0xda, 0x2e, 0xe3, 0x09, 0xd4, 0x4e, 0x7f, 0x13, 0xfa, 0x54, 0x9a, 0x8a, 0xcb, + 0x2a, 0xd2, 0x20, 0x8c, 0xbd, 0x68, 0xb0, 0x4e, 0x52, 0xb3, 0x8c, 0x70, 0x47, 0xbc, 0x7a, 0xa6, + 0xa1, 0xec, 0x08, 0x36, 0xec, 0x8b, 0x72, 0x35, 0x23, 0xd1, 0x94, 0x50, 0xd8, 0xf1, 0x3c, 0xc6, + 0xcd, 0x18, 0x1c, 0xc7, 0x2e, 0xe1, 0xac, 0x19, 0x7a, 0x08, 0xd7, 0xe7, 0x96, 0x76, 0xea, 0x9d, + 0xba, 0x53, 0x3e, 0x15, 0xe9, 0x99, 0x31, 0x20, 0x1b, 0xa4, 0xc0, 0xae, 0xcd, 0x2c, 0xe2, 0xbe, + 0x77, 0xba, 0x4f, 0x34, 0xda, 0x9c, 0xfc, 0x0a, 0xde, 0x9d, 0x1b, 0x45, 0x57, 0x7a, 0xf2, 0xd8, + 0x3b, 0x8e, 0x78, 0x30, 0xb8, 0x4c, 0x5f, 0x74, 0x65, 0x66, 0x88, 0x43, 0xa4, 0x78, 0xa4, 0x09, + 0x8c, 0x43, 0x77, 0x0c, 0x6d, 0x0a, 0x43, 0x90, 0x36, 0xcc, 0xab, 0x6e, 0x2b, 0x6f, 0xae, 0xba, + 0xfd, 0x14, 0xba, 0xc6, 0xdb, 0xbf, 0xa8, 0x8c, 0xb7, 0xa3, 0xf1, 0xf8, 0x2c, 0x87, 0x77, 0xa0, + 0x4d, 0xae, 0x3e, 0xbd, 0xe3, 0x3a, 0x74, 0xa8, 0xda, 0xcb, 0x3d, 0x8e, 0x84, 0x7f, 0x62, 0x9d, + 0x73, 0x02, 0x3d, 0x40, 0xc8, 0x10, 0xa0, 0xf5, 0x3c, 0x0e, 0x45, 0xbc, 0x13, 0x45, 0xc3, 0xbf, + 0x6c, 0x40, 0x1b, 0x7d, 0x02, 0x8a, 0x9b, 0xe0, 0xb1, 0x8a, 0x16, 0x8e, 0x72, 0xa9, 0x53, 0x2f, + 0x31, 0x75, 0xc5, 0x1d, 0x04, 0x22, 0xd5, 0xbe, 0x97, 0xcc, 0xa5, 0x5a, 0x97, 0xe6, 0x52, 0xad, + 0x1f, 0xe8, 0xbb, 0x2f, 0xba, 0xde, 0x8c, 0xdb, 0x42, 0x55, 0x1a, 0xe0, 0x81, 0x06, 0xa1, 0xaf, + 0x42, 0x24, 0x5e, 0x44, 0xfe, 0x0d, 0x9e, 0x9e, 0x22, 0x69, 0xb2, 0xb2, 0x24, 0x37, 0x3b, 0x06, + 0x71, 0xc8, 0xb5, 0x3e, 0x2e, 0x05, 0xcb, 0xea, 0xf3, 0xc1, 0xb2, 0xdb, 0x00, 0xbe, 0x88, 0x03, + 0x72, 0xa1, 0xe6, 0xb2, 0x61, 0x3a, 0x25, 0x5a, 0x60, 0xbf, 0x47, 0x68, 0xf6, 0x63, 0xe8, 0xe7, + 0x14, 0xe8, 0x21, 0xf9, 0x71, 0x7e, 0xfe, 0x34, 0x54, 0x0e, 0x1f, 0xed, 0xc6, 0x6a, 0x3e, 0x86, + 0xdb, 0x5e, 0x88, 0xe1, 0x5e, 0x90, 0x3c, 0x87, 0x1f, 0x7c, 0x2d, 0xe3, 0x0a, 0xb4, 0xa8, 0x4a, + 0x27, 0xc8, 0x12, 0xa3, 0xab, 0x9b, 0xa1, 0xa4, 0x58, 0xfb, 0x45, 0x71, 0xe2, 0xee, 0xff, 0xaf, + 0x38, 0x71, 0xef, 0xfb, 0xc5, 0x89, 0x97, 0xbf, 0x5f, 0x9c, 0x78, 0x2e, 0xae, 0xba, 0x32, 0x9f, + 0x8e, 0xb9, 0x30, 0xf9, 0xd1, 0xbf, 0x30, 0xf9, 0xf1, 0x96, 0xcc, 0xc5, 0xea, 0x1b, 0x33, 0x17, + 0xdf, 0x23, 0x75, 0xc2, 0xde, 0x96, 0x3a, 0xb9, 0x05, 0x2b, 0x2a, 0xf5, 0xfc, 0x13, 0x7d, 0x12, + 0x39, 0xe1, 0x67, 0xd2, 0xa4, 0x6a, 0x7a, 0x04, 0xc6, 0x73, 0xc8, 0x6f, 0xf8, 0x99, 0x1c, 0x3e, + 0x07, 0xa0, 0x23, 0x1a, 0x7d, 0xda, 0x45, 0xb2, 0x51, 0xf9, 0xc1, 0x85, 0x15, 0xff, 0xa7, 0x02, + 0x70, 0xe8, 0x4d, 0x13, 0x1d, 0xe3, 0x64, 0x7f, 0x08, 0x1d, 0x49, 0xad, 0x72, 0x86, 0xbb, 0x64, + 0xc8, 0x0a, 0x52, 0xf3, 0xa8, 0x6f, 0x17, 0xc8, 0xfc, 0x99, 0xc4, 0x5a, 0x8f, 0x90, 0x17, 0xb1, + 0xd5, 0x2d, 0x01, 0xc5, 0xbe, 0x6e, 0xc2, 0xb2, 0x21, 0x48, 0x78, 0xea, 0xf3, 0x58, 0x57, 0xc6, + 0x56, 0x9c, 0x9e, 0x86, 0x1e, 0x68, 0x20, 0xbb, 0x9f, 0x93, 0x59, 0x93, 0xb1, 0x98, 0xa6, 0x31, + 0x5d, 0x8c, 0xcd, 0x18, 0x6e, 0xdb, 0x4f, 0xa1, 0x89, 0xb4, 0xa0, 0x86, 0xef, 0xeb, 0xbf, 0xc3, + 0x3a, 0xd0, 0x34, 0xa3, 0xf6, 0x2b, 0xac, 0x07, 0x6d, 0xba, 0xce, 0x42, 0xb8, 0xa5, 0xe1, 0x1f, + 0xaf, 0x42, 0x67, 0x2f, 0x96, 0x2a, 0xcd, 0xb4, 0x08, 0x17, 0x97, 0x36, 0xea, 0x74, 0x69, 0xc3, + 0x54, 0x50, 0xea, 0xcf, 0xa0, 0x0a, 0xca, 0x4f, 0xa1, 0x69, 0xee, 0x07, 0x99, 0xc0, 0xf7, 0xb9, + 0x97, 0x8b, 0x2c, 0x0d, 0xdb, 0x82, 0x56, 0x60, 0x2e, 0x2e, 0x99, 0x34, 0x7e, 0xe9, 0x36, 0x91, + 0xbd, 0xd2, 0xe4, 0xe4, 0x34, 0xec, 0x03, 0xa8, 0x7a, 0xe3, 0xb1, 0x39, 0xf5, 0xae, 0x14, 0xa4, + 0xe4, 0xc4, 0x38, 0x88, 0x63, 0x77, 0xa1, 0x4d, 0xea, 0x93, 0x2a, 0x59, 0x1a, 0xf3, 0x63, 0xda, + 0x32, 0x19, 0xad, 0x51, 0x29, 0x66, 0x7e, 0x17, 0xda, 0x91, 0x10, 0x89, 0xee, 0xd0, 0x9c, 0xef, 0x60, 0x8b, 0x1b, 0x9c, 0x56, 0x64, 0xcb, 0x1c, 0x6e, 0x41, 0x03, 0xdd, 0x63, 0x91, 0x18, 0xb7, - 0xb2, 0x34, 0x0f, 0x4a, 0xf2, 0x3b, 0x75, 0x89, 0x3f, 0x6c, 0x07, 0x40, 0xcb, 0x3f, 0x8d, 0xdc, - 0x5e, 0x64, 0x47, 0x9e, 0xcf, 0xc3, 0x4d, 0x6a, 0x53, 0x7b, 0x0f, 0x60, 0xa0, 0x73, 0x37, 0xa5, - 0x9e, 0x60, 0xeb, 0xff, 0x6c, 0xcf, 0xf9, 0x74, 0xa0, 0xd3, 0x4f, 0xe7, 0xd3, 0x83, 0x1f, 0x43, - 0x33, 0xd1, 0xc9, 0x0b, 0xd2, 0x30, 0x9d, 0x9d, 0xb5, 0xa2, 0xab, 0xc9, 0x6a, 0x38, 0x96, 0x82, - 0xfd, 0x01, 0xf4, 0x75, 0x9d, 0xda, 0xd8, 0x44, 0xf1, 0x29, 0xf2, 0x35, 0x77, 0xb5, 0x64, 0x2e, - 0xc8, 0xef, 0xf4, 0xd4, 0x5c, 0xcc, 0xff, 0xe7, 0xd0, 0x2b, 0x4a, 0xfd, 0x7d, 0x2f, 0x26, 0xbd, - 0x43, 0xd1, 0x74, 0xdb, 0xbd, 0x7c, 0x5a, 0x71, 0xba, 0xbc, 0x7c, 0x76, 0xd9, 0x82, 0x86, 0xa9, - 0x9d, 0x1c, 0x50, 0xaf, 0xd2, 0xcd, 0x53, 0x5d, 0x2d, 0xe5, 0x18, 0x3c, 0xf2, 0xb2, 0x28, 0x0b, - 0x23, 0xc7, 0x6a, 0x8e, 0x97, 0x79, 0x4d, 0x98, 0xd3, 0xce, 0xcb, 0xc1, 0xd8, 0xa3, 0xf9, 0x32, - 0x35, 0x5d, 0x8e, 0xb5, 0x4e, 0x5d, 0xaf, 0x5c, 0xd0, 0x55, 0x57, 0x65, 0x39, 0xab, 0xc9, 0x42, - 0xb5, 0xdb, 0x1d, 0x68, 0x89, 0x34, 0xa0, 0xca, 0x5a, 0xca, 0xf5, 0x12, 0x3f, 0xa9, 0x3a, 0x4f, - 0x5f, 0x8a, 0x22, 0xe5, 0xd1, 0x14, 0xba, 0x81, 0x8e, 0x45, 0x92, 0x0a, 0xf2, 0x02, 0x49, 0xc5, - 0x5d, 0x5a, 0x76, 0x2c, 0x0c, 0x9e, 0x14, 0xdc, 0x87, 0xd0, 0xb4, 0x15, 0xa1, 0x9b, 0x4b, 0x94, - 0x16, 0xc5, 0x3e, 0x85, 0xd5, 0x79, 0x85, 0x26, 0x87, 0x97, 0x97, 0xa8, 0xfb, 0x73, 0xfa, 0x0b, - 0xad, 0x71, 0x3d, 0x0a, 0x67, 0xa1, 0x1a, 0x0e, 0x97, 0x0e, 0x3f, 0x1a, 0x81, 0xe7, 0x23, 0x93, - 0x22, 0xb8, 0xb2, 0x7c, 0x3e, 0x32, 0x69, 0x84, 0x21, 0x34, 0x43, 0xf9, 0x38, 0x4c, 0xa5, 0x1a, - 0x5e, 0xb5, 0xd6, 0x91, 0x9a, 0x6c, 0x13, 0x1a, 0xa1, 0x44, 0x33, 0x31, 0xbc, 0x66, 0xaf, 0xd1, - 0x91, 0xd1, 0xb8, 0x0d, 0x0d, 0x53, 0x2d, 0x7b, 0x63, 0x69, 0x47, 0x9b, 0x9a, 0x74, 0xc7, 0x50, - 0xb0, 0x1f, 0x43, 0x93, 0x4a, 0x25, 0x45, 0x32, 0x7c, 0x7f, 0x51, 0x02, 0x74, 0xbd, 0xa2, 0xd3, - 0x88, 0x74, 0xdd, 0xe2, 0xc7, 0xd0, 0xb4, 0x4e, 0xca, 0x68, 0x51, 0xaa, 0x8d, 0xb3, 0xe2, 0x58, - 0x0a, 0x76, 0x13, 0xea, 0x33, 0xd4, 0x63, 0xc3, 0x0f, 0x16, 0x77, 0xa8, 0x56, 0x6f, 0x1a, 0xcb, - 0xfe, 0x1e, 0x5c, 0x2d, 0x17, 0x1b, 0xda, 0x4a, 0x44, 0x13, 0x00, 0xbc, 0x49, 0x7d, 0xdf, 0xbf, - 0x40, 0x54, 0xe6, 0x6b, 0x16, 0x9d, 0xcb, 0xc9, 0x1b, 0x8a, 0x19, 0x3f, 0xcf, 0xd5, 0x3d, 0xee, - 0xae, 0xe1, 0x2d, 0x5b, 0xc6, 0xb8, 0x6c, 0x30, 0xac, 0x11, 0x20, 0x3b, 0xf3, 0x05, 0x74, 0xc7, - 0xd9, 0xeb, 0xd7, 0xe7, 0x36, 0x78, 0xfd, 0x11, 0xf5, 0x2b, 0x1d, 0xc1, 0x4b, 0xf5, 0x8d, 0x4e, - 0x67, 0x5c, 0x2a, 0x76, 0xbc, 0x0c, 0x4d, 0x3f, 0x76, 0xbd, 0x20, 0x48, 0x87, 0x5b, 0xba, 0xbe, - 0xd1, 0x8f, 0x77, 0x83, 0x80, 0x6e, 0x0d, 0x8b, 0x84, 0xd3, 0x15, 0x3d, 0x37, 0x0c, 0x86, 0x3f, - 0xd6, 0x86, 0xc7, 0x82, 0xf6, 0x03, 0xba, 0x56, 0x6c, 0xcf, 0xad, 0x61, 0x30, 0xbc, 0x6d, 0xae, - 0x15, 0x1b, 0xd0, 0x7e, 0x80, 0x8e, 0x27, 0x3a, 0xf9, 0x16, 0x32, 0xfc, 0x58, 0x27, 0x04, 0x66, - 0xde, 0xd9, 0xa1, 0x01, 0xe1, 0x26, 0xd5, 0xe9, 0x35, 0x52, 0x5b, 0x77, 0x16, 0x37, 0x69, 0x9e, - 0xa6, 0x74, 0xda, 0x61, 0x9e, 0xb1, 0xa4, 0x8d, 0x4d, 0xaa, 0xc8, 0x8d, 0x76, 0x86, 0x9f, 0x2c, - 0x6f, 0x6c, 0x93, 0x85, 0xc5, 0x8d, 0x6d, 0x13, 0xb2, 0x3b, 0x00, 0x5a, 0x67, 0x91, 0xc2, 0xd9, - 0x5e, 0xec, 0x93, 0x9f, 0x06, 0x1c, 0x7d, 0x19, 0x80, 0x54, 0xcd, 0x0e, 0x00, 0x85, 0xdf, 0x75, - 0x9f, 0xbb, 0x8b, 0x7d, 0x72, 0xef, 0xde, 0x69, 0xbf, 0xcc, 0x1d, 0xfd, 0xbb, 0xd0, 0xce, 0xd0, - 0x8f, 0x47, 0x4f, 0x7a, 0x78, 0x6f, 0x51, 0x98, 0xad, 0x8b, 0xef, 0xb4, 0x32, 0xf3, 0x84, 0x2f, - 0x21, 0xdb, 0x43, 0x6e, 0xc8, 0xf0, 0xfe, 0xe2, 0x4b, 0xf2, 0x73, 0x80, 0x43, 0x26, 0x4a, 0x1f, - 0x09, 0x3e, 0x87, 0x8e, 0x66, 0x9a, 0xee, 0xb4, 0xb3, 0x28, 0x23, 0x85, 0x5f, 0xe3, 0x68, 0xee, - 0xea, 0x6e, 0x37, 0xa1, 0xee, 0x25, 0x49, 0x74, 0x3e, 0xfc, 0x74, 0x51, 0xc2, 0x77, 0x11, 0xec, - 0x68, 0x2c, 0x8a, 0xd2, 0x2c, 0x8b, 0x54, 0x68, 0xeb, 0xf7, 0x3f, 0x5b, 0x14, 0xa5, 0xd2, 0x85, - 0x28, 0xa7, 0x33, 0x2b, 0xdd, 0x8e, 0xba, 0x03, 0xad, 0x44, 0x48, 0xe5, 0x06, 0xb3, 0x68, 0xf8, - 0xf9, 0x92, 0x19, 0xd1, 0x55, 0xe8, 0x4e, 0x33, 0x31, 0x65, 0xfc, 0x73, 0xd7, 0xf4, 0x7e, 0x32, - 0x7f, 0x4d, 0x8f, 0xed, 0x40, 0x77, 0x26, 0xe2, 0x89, 0x08, 0x4e, 0x34, 0xf7, 0x7f, 0x5a, 0xae, - 0x71, 0x3e, 0x40, 0x0c, 0x71, 0xbe, 0x63, 0x88, 0xb0, 0xa1, 0xcf, 0x76, 0xbf, 0xae, 0xb5, 0xd6, - 0x06, 0xec, 0xd7, 0xb5, 0xd6, 0x87, 0x83, 0x9b, 0x4e, 0x47, 0xd2, 0xad, 0x79, 0x1a, 0x62, 0xf4, - 0x39, 0x74, 0x77, 0xe9, 0x1f, 0x03, 0x42, 0x49, 0x7a, 0xf4, 0x26, 0xd4, 0xf2, 0xfc, 0x7c, 0xae, - 0xa0, 0x89, 0xe2, 0x35, 0xdf, 0x8f, 0xc7, 0xc2, 0x21, 0xf4, 0xe8, 0x5f, 0xd7, 0xa0, 0x71, 0x24, - 0xb2, 0xd4, 0xe7, 0xdf, 0x7e, 0x97, 0xe4, 0x47, 0x56, 0xca, 0xe2, 0xa2, 0x5c, 0x58, 0x0b, 0x14, - 0xa1, 0x17, 0xeb, 0x13, 0xdb, 0x45, 0xea, 0x7f, 0x03, 0xea, 0xfa, 0x68, 0xa8, 0x23, 0xca, 0xba, - 0x41, 0x3b, 0x2c, 0x93, 0x53, 0xfa, 0x5b, 0x00, 0x93, 0x23, 0xa9, 0x39, 0x60, 0x41, 0xfb, 0x01, + 0xb2, 0x34, 0x0f, 0x4a, 0xf2, 0x3b, 0x75, 0x89, 0x3f, 0x6c, 0x1b, 0x40, 0xcb, 0x3f, 0x8d, 0xdc, + 0x9e, 0x67, 0x47, 0x9e, 0xcf, 0xc3, 0x4d, 0x6a, 0x53, 0x7b, 0x0f, 0xa0, 0xaf, 0x73, 0x37, 0xa5, + 0x9e, 0x60, 0xeb, 0xff, 0x6c, 0xcf, 0xd9, 0x74, 0xa0, 0xb3, 0x9c, 0xce, 0xa6, 0x07, 0x3f, 0x81, + 0x66, 0xa2, 0x93, 0x17, 0xa4, 0x61, 0x3a, 0xdb, 0xab, 0x45, 0x57, 0x93, 0xd5, 0x70, 0x2c, 0x05, + 0xfb, 0x03, 0x58, 0xd6, 0x75, 0x6a, 0x23, 0x13, 0xc5, 0xa7, 0xc8, 0xd7, 0xcc, 0xd5, 0x92, 0x99, + 0x20, 0xbf, 0xd3, 0x53, 0x33, 0x31, 0xff, 0x5f, 0x40, 0xaf, 0x28, 0xf5, 0xf7, 0xbd, 0x98, 0xf4, + 0x0e, 0x45, 0xd3, 0x6d, 0xf7, 0xf2, 0x69, 0xc5, 0xe9, 0xf2, 0xf2, 0xd9, 0x65, 0x13, 0x1a, 0xa6, + 0x76, 0xb2, 0x4f, 0xbd, 0x4a, 0x37, 0x4f, 0x75, 0xb5, 0x94, 0x63, 0xf0, 0xc8, 0xcb, 0xa2, 0x2c, + 0x8c, 0x1c, 0xab, 0x19, 0x5e, 0xe6, 0x35, 0x61, 0x4e, 0x3b, 0x2f, 0x07, 0x63, 0x8f, 0x66, 0xcb, + 0xd4, 0x74, 0x39, 0xd6, 0x1a, 0x75, 0xbd, 0x72, 0x4e, 0x57, 0x5d, 0x95, 0xe5, 0xac, 0x24, 0x73, + 0xd5, 0x6e, 0x77, 0xa0, 0x25, 0xd2, 0x80, 0x2a, 0x6b, 0x29, 0xd7, 0x4b, 0xfc, 0xa4, 0xea, 0x3c, + 0x7d, 0x29, 0x8a, 0x94, 0x47, 0x53, 0xe8, 0x06, 0x3a, 0x16, 0x49, 0x2a, 0xc8, 0x0b, 0x24, 0x15, + 0x77, 0x69, 0xd1, 0xb1, 0x30, 0x78, 0x52, 0x70, 0x1f, 0x41, 0xd3, 0x56, 0x84, 0x6e, 0x2c, 0x50, + 0x5a, 0x14, 0xfb, 0x0c, 0x56, 0x66, 0x15, 0x9a, 0x1c, 0x5c, 0x5e, 0xa0, 0x5e, 0x9e, 0xd1, 0x5f, + 0x68, 0x8d, 0xeb, 0x51, 0x38, 0x0d, 0xd5, 0x60, 0xb0, 0x70, 0xf8, 0xd1, 0x08, 0x3c, 0x1f, 0x99, + 0x14, 0xc1, 0x95, 0xc5, 0xf3, 0x91, 0x49, 0x23, 0x0c, 0xa0, 0x19, 0xca, 0xc7, 0x61, 0x2a, 0xd5, + 0xe0, 0xaa, 0xb5, 0x8e, 0xd4, 0x64, 0x1b, 0xd0, 0x08, 0x25, 0x9a, 0x89, 0xc1, 0x35, 0x7b, 0x8d, + 0x8e, 0x8c, 0xc6, 0x6d, 0x68, 0x98, 0x6a, 0xd9, 0x1b, 0x0b, 0x3b, 0xda, 0xd4, 0xa4, 0x3b, 0x86, + 0x82, 0xfd, 0x04, 0x9a, 0x54, 0x2a, 0x29, 0x92, 0xc1, 0x07, 0xf3, 0x12, 0xa0, 0xeb, 0x15, 0x9d, + 0x46, 0xa4, 0xeb, 0x16, 0x3f, 0x81, 0xa6, 0x75, 0x52, 0x86, 0xf3, 0x52, 0x6d, 0x9c, 0x15, 0xc7, + 0x52, 0xb0, 0x9b, 0x50, 0x9f, 0xa2, 0x1e, 0x1b, 0x7c, 0x38, 0xbf, 0x43, 0xb5, 0x7a, 0xd3, 0x58, + 0xf6, 0xf7, 0xe0, 0x6a, 0xb9, 0xd8, 0xd0, 0x56, 0x22, 0x9a, 0x00, 0xe0, 0x4d, 0xea, 0xfb, 0xc1, + 0x39, 0xa2, 0x32, 0x5b, 0xb3, 0xe8, 0x5c, 0x4e, 0x2e, 0x28, 0x66, 0xfc, 0x22, 0x57, 0xf7, 0xb8, + 0xbb, 0x06, 0xb7, 0x6c, 0x19, 0xe3, 0xa2, 0xc1, 0xb0, 0x46, 0x80, 0xec, 0xcc, 0x97, 0xd0, 0x1d, + 0x65, 0xaf, 0x5f, 0x9f, 0xd9, 0xe0, 0xf5, 0xc7, 0xd4, 0xaf, 0x74, 0x04, 0x2f, 0xd5, 0x37, 0x3a, + 0x9d, 0x51, 0xa9, 0xd8, 0xf1, 0x32, 0x34, 0xfd, 0xd8, 0xf5, 0x82, 0x20, 0x1d, 0x6c, 0xea, 0xfa, + 0x46, 0x3f, 0xde, 0x09, 0x02, 0xba, 0x35, 0x2c, 0x12, 0x4e, 0x57, 0xf4, 0xdc, 0x30, 0x18, 0xfc, + 0x44, 0x1b, 0x1e, 0x0b, 0xda, 0x0b, 0xe8, 0x5a, 0xb1, 0x3d, 0xb7, 0x86, 0xc1, 0xe0, 0xb6, 0xb9, + 0x56, 0x6c, 0x40, 0x7b, 0x01, 0x3a, 0x9e, 0xe8, 0xe4, 0x5b, 0xc8, 0xe0, 0x13, 0x9d, 0x10, 0x98, + 0x7a, 0xa7, 0x07, 0x06, 0x84, 0x9b, 0x54, 0xa7, 0xd7, 0x48, 0x6d, 0xdd, 0x99, 0xdf, 0xa4, 0x79, + 0x9a, 0xd2, 0x69, 0x87, 0x79, 0xc6, 0x92, 0x36, 0x36, 0xa9, 0x22, 0x37, 0xda, 0x1e, 0x7c, 0xba, + 0xb8, 0xb1, 0x4d, 0x16, 0x16, 0x37, 0xb6, 0x4d, 0xc8, 0x6e, 0x03, 0x68, 0x9d, 0x45, 0x0a, 0x67, + 0x6b, 0xbe, 0x4f, 0x7e, 0x1a, 0x70, 0xf4, 0x65, 0x00, 0x52, 0x35, 0xdb, 0x00, 0x14, 0x7e, 0xd7, + 0x7d, 0xee, 0xce, 0xf7, 0xc9, 0xbd, 0x7b, 0xa7, 0xfd, 0x32, 0x77, 0xf4, 0xef, 0x42, 0x3b, 0x43, + 0x3f, 0x1e, 0x3d, 0xe9, 0xc1, 0xbd, 0x79, 0x61, 0xb6, 0x2e, 0xbe, 0xd3, 0xca, 0xcc, 0x13, 0xbe, + 0x84, 0x6c, 0x0f, 0xb9, 0x21, 0x83, 0xfb, 0xf3, 0x2f, 0xc9, 0xcf, 0x01, 0x0e, 0x99, 0x28, 0x7d, + 0x24, 0xf8, 0x02, 0x3a, 0x9a, 0x69, 0xba, 0xd3, 0xf6, 0xbc, 0x8c, 0x14, 0x7e, 0x8d, 0xa3, 0xb9, + 0xab, 0xbb, 0xdd, 0x84, 0xba, 0x97, 0x24, 0xd1, 0xd9, 0xe0, 0xb3, 0x79, 0x09, 0xdf, 0x41, 0xb0, + 0xa3, 0xb1, 0x28, 0x4a, 0xd3, 0x2c, 0x52, 0xa1, 0xad, 0xdf, 0xff, 0x7c, 0x5e, 0x94, 0x4a, 0x17, + 0xa2, 0x9c, 0xce, 0xb4, 0x74, 0x3b, 0xea, 0x0e, 0xb4, 0x12, 0x21, 0x95, 0x1b, 0x4c, 0xa3, 0xc1, + 0x17, 0x0b, 0x66, 0x44, 0x57, 0xa1, 0x3b, 0xcd, 0xc4, 0x94, 0xf1, 0xcf, 0x5c, 0xd3, 0xfb, 0xe9, + 0xec, 0x35, 0x3d, 0xb6, 0x0d, 0xdd, 0xa9, 0x88, 0xc7, 0x22, 0x38, 0xd6, 0xdc, 0xff, 0x59, 0xb9, + 0xc6, 0x79, 0x1f, 0x31, 0xc4, 0xf9, 0x8e, 0x21, 0xc2, 0x86, 0x3e, 0xdb, 0xfd, 0xba, 0xd6, 0x5a, + 0xed, 0xb3, 0x5f, 0xd7, 0x5a, 0x1f, 0xf5, 0x6f, 0x3a, 0x1d, 0x49, 0xb7, 0xe6, 0x69, 0x88, 0xe1, + 0x17, 0xd0, 0xdd, 0xa1, 0x7f, 0x0c, 0x08, 0x25, 0xe9, 0xd1, 0x9b, 0x50, 0xcb, 0xf3, 0xf3, 0xb9, + 0x82, 0x26, 0x8a, 0xd7, 0x7c, 0x2f, 0x1e, 0x09, 0x87, 0xd0, 0xc3, 0x7f, 0x53, 0x83, 0xc6, 0xa1, + 0xc8, 0x52, 0x9f, 0xbf, 0xfd, 0x2e, 0xc9, 0x7b, 0x56, 0xca, 0xe2, 0xa2, 0x5c, 0x58, 0x0b, 0x14, + 0xa1, 0xe7, 0xeb, 0x13, 0xdb, 0x45, 0xea, 0x7f, 0x1d, 0xea, 0xfa, 0x68, 0xa8, 0x23, 0xca, 0xba, + 0x41, 0x3b, 0x2c, 0x93, 0x13, 0xfa, 0x5b, 0x00, 0x93, 0x23, 0xa9, 0x39, 0x60, 0x41, 0x7b, 0x01, 0x85, 0x8e, 0x2c, 0x01, 0x6d, 0xe1, 0x86, 0x09, 0x0d, 0x1b, 0x20, 0x6d, 0x64, 0x5b, 0x56, 0xd0, - 0x7c, 0x43, 0x59, 0xc1, 0x7b, 0x50, 0x8b, 0x6d, 0x25, 0x7b, 0x8e, 0xa7, 0x3b, 0xd9, 0x04, 0x67, - 0xb7, 0x21, 0xbf, 0x00, 0x63, 0x5c, 0x92, 0x37, 0x5f, 0x90, 0xd9, 0x81, 0x76, 0xfe, 0x1f, 0x13, - 0xc6, 0x0b, 0xd9, 0xd8, 0x2e, 0xfe, 0x75, 0xe2, 0xd8, 0x3e, 0x39, 0x05, 0xd9, 0xdb, 0x93, 0xe3, - 0x9d, 0x1f, 0x94, 0x1c, 0x37, 0x47, 0x34, 0x5f, 0xc4, 0x52, 0x99, 0xe0, 0x58, 0x33, 0x94, 0x7b, - 0xd8, 0x64, 0xbf, 0x07, 0xbd, 0x94, 0xfb, 0x2f, 0xdd, 0x99, 0x9c, 0xe8, 0x57, 0xf4, 0xca, 0x97, - 0xf0, 0x66, 0x72, 0xf2, 0x25, 0x25, 0xee, 0xcd, 0x89, 0xa9, 0x83, 0xb4, 0x07, 0x72, 0x42, 0xa3, - 0x7e, 0x0c, 0x6b, 0x33, 0x3e, 0x3b, 0xe1, 0xa9, 0x9c, 0x86, 0x89, 0x55, 0xb5, 0x7d, 0x2a, 0x30, - 0x18, 0x14, 0x08, 0x3d, 0x97, 0xd1, 0x3f, 0xaa, 0x40, 0x0b, 0xb9, 0x88, 0xb2, 0xc4, 0x18, 0xd4, - 0x66, 0x7e, 0x92, 0x19, 0x47, 0x98, 0x9e, 0xcd, 0xff, 0x56, 0x68, 0x29, 0x31, 0xff, 0x5b, 0x41, - 0x6b, 0xa8, 0x53, 0x3e, 0xf4, 0xac, 0xef, 0x81, 0x9f, 0x53, 0x54, 0x50, 0x4b, 0x86, 0x6d, 0xb2, - 0x4b, 0xd0, 0xf0, 0x63, 0x3a, 0x0d, 0xeb, 0xd4, 0x59, 0xdd, 0x8f, 0xf1, 0x14, 0xac, 0xc1, 0x45, - 0x35, 0x74, 0xdd, 0x8f, 0xf7, 0x83, 0xb3, 0xd1, 0xbf, 0xaf, 0xc0, 0xda, 0x61, 0x2a, 0x7c, 0x2e, - 0xe5, 0x13, 0x34, 0xe4, 0x94, 0xa6, 0xc0, 0x37, 0x52, 0x54, 0x57, 0x67, 0x04, 0xe8, 0x19, 0x65, - 0x58, 0x87, 0x2a, 0xf2, 0xe3, 0x46, 0xd5, 0x69, 0x13, 0x84, 0x4e, 0x1b, 0x39, 0xba, 0x94, 0xfe, - 0xd6, 0x68, 0x8a, 0x07, 0xdf, 0x84, 0x7e, 0x91, 0x58, 0x29, 0x65, 0xea, 0x8b, 0x6b, 0xa6, 0x34, - 0xca, 0x75, 0xe8, 0x98, 0x7a, 0x0a, 0x1a, 0x46, 0x87, 0xf8, 0x41, 0x83, 0x8e, 0xcc, 0x2c, 0xb4, - 0x72, 0x20, 0xbc, 0x0e, 0xea, 0x6b, 0x75, 0x81, 0xe8, 0xd1, 0xdf, 0x87, 0xc1, 0x61, 0xca, 0x13, - 0x2f, 0xe5, 0x54, 0x5f, 0x41, 0x2c, 0xde, 0x84, 0x46, 0xc4, 0xe3, 0x89, 0x49, 0xe9, 0x57, 0x1d, - 0xd3, 0xca, 0xff, 0x72, 0x64, 0xa5, 0xf4, 0x97, 0x23, 0xc8, 0xea, 0x94, 0x7b, 0xe6, 0x9f, 0x49, - 0xe8, 0x19, 0xb7, 0x20, 0x1e, 0x19, 0xf5, 0xb9, 0xa8, 0xe5, 0xe8, 0x86, 0xb9, 0xcb, 0x76, 0x12, - 0xc6, 0x54, 0x9b, 0x46, 0x77, 0xd9, 0x1e, 0x84, 0xf1, 0xe8, 0x3f, 0x54, 0xa1, 0x63, 0xf8, 0x49, - 0x2f, 0xd7, 0x6b, 0x59, 0xc9, 0xd7, 0x72, 0x00, 0x55, 0xf9, 0x22, 0x32, 0x8b, 0x8b, 0x8f, 0xec, - 0x53, 0xa8, 0x46, 0xe1, 0xcc, 0x1c, 0x71, 0xae, 0xcd, 0x99, 0xab, 0xf9, 0x55, 0x31, 0x82, 0x87, - 0xd4, 0xa8, 0x23, 0xe9, 0x1e, 0x2c, 0x8a, 0xb8, 0xe1, 0x24, 0x9a, 0x8e, 0x33, 0xdc, 0x47, 0xc8, - 0x23, 0xcf, 0xa7, 0x32, 0x03, 0xab, 0x1c, 0x7a, 0x4e, 0xdb, 0x40, 0xf6, 0x03, 0xf6, 0x19, 0xb4, - 0xf2, 0x40, 0xa5, 0x3d, 0xd4, 0xa8, 0xb3, 0x78, 0x7b, 0xef, 0xe9, 0xf1, 0x59, 0x6c, 0x23, 0x91, - 0xe6, 0x65, 0x39, 0x25, 0xfb, 0x03, 0xe8, 0x4a, 0x2e, 0xa5, 0xbe, 0xa8, 0x38, 0x16, 0x46, 0x69, - 0x5c, 0x2a, 0x9f, 0x57, 0x08, 0x8b, 0x5f, 0x6d, 0xb7, 0x88, 0x2c, 0x40, 0xec, 0x4b, 0xe8, 0xdb, - 0xfe, 0x91, 0x98, 0x4c, 0xf2, 0x40, 0xfa, 0xb5, 0xa5, 0x11, 0x9e, 0x10, 0xba, 0x34, 0x4e, 0x4f, - 0x96, 0x11, 0xec, 0x57, 0xd0, 0x4f, 0xf4, 0x1a, 0xbb, 0xa6, 0x94, 0x47, 0x2b, 0x9f, 0xab, 0x73, - 0xde, 0xd5, 0x9c, 0x0c, 0x14, 0x57, 0x89, 0x0a, 0xb8, 0x5c, 0xbe, 0xb4, 0xab, 0x33, 0x2b, 0x73, - 0x97, 0x76, 0x47, 0xff, 0xbb, 0x0a, 0x9d, 0xd2, 0xa7, 0xd1, 0x9f, 0xc8, 0x48, 0x9e, 0xda, 0x22, - 0x1a, 0x7c, 0x46, 0xd8, 0x54, 0x48, 0x5b, 0x13, 0x42, 0xcf, 0x08, 0x4b, 0x45, 0x9e, 0x1b, 0xa7, - 0x67, 0x7c, 0xa1, 0x39, 0xad, 0xea, 0x6b, 0xd5, 0xb4, 0x72, 0x35, 0xa7, 0x5b, 0x00, 0xf7, 0x03, - 0xfa, 0xb7, 0x19, 0x4f, 0x79, 0x27, 0x9e, 0xb4, 0xd5, 0x48, 0x79, 0x1b, 0x77, 0xfd, 0x4b, 0x9e, - 0xe2, 0x5c, 0x6c, 0x2e, 0xd0, 0x34, 0x51, 0x20, 0x48, 0x51, 0xbe, 0x16, 0xb1, 0xce, 0x03, 0x76, - 0x9d, 0x16, 0x02, 0xbe, 0x16, 0x31, 0x75, 0x33, 0xcb, 0x6f, 0xf2, 0xd9, 0xb6, 0x89, 0xea, 0xf0, - 0x45, 0xc6, 0xd1, 0x4d, 0x0d, 0xe8, 0x3e, 0x43, 0xdb, 0x69, 0x52, 0x5b, 0xa7, 0xd8, 0xc9, 0x9f, - 0x7e, 0xe5, 0x85, 0x8a, 0xe4, 0x4c, 0x64, 0xca, 0x70, 0x68, 0x15, 0x11, 0x5f, 0x79, 0xa1, 0x3a, - 0xd6, 0x60, 0x76, 0xdf, 0x5c, 0x53, 0x2a, 0xd3, 0xba, 0x78, 0x18, 0xd0, 0x51, 0x30, 0xb6, 0x40, - 0x7f, 0xc4, 0xe9, 0xff, 0x56, 0x66, 0x9e, 0x4a, 0xc3, 0x33, 0x11, 0xa3, 0x59, 0x54, 0xe1, 0x4b, - 0x5e, 0xfc, 0xc3, 0x4d, 0xcb, 0x59, 0xcf, 0x91, 0x4f, 0x09, 0x47, 0x89, 0x93, 0xe7, 0xb0, 0xc5, - 0xcf, 0x92, 0x28, 0xf4, 0xc3, 0x85, 0x0b, 0x82, 0xae, 0xef, 0x49, 0xe5, 0xa6, 0x5c, 0x65, 0x69, - 0x2c, 0x29, 0xc0, 0x63, 0xaa, 0x34, 0x3e, 0xb0, 0xf4, 0xe5, 0x4b, 0x83, 0x7b, 0x9e, 0x54, 0x8e, - 0xa6, 0x7d, 0x9a, 0x45, 0x11, 0x32, 0x21, 0x4f, 0xd8, 0xe8, 0x5a, 0x8d, 0xa6, 0xd4, 0xa9, 0x9a, - 0xd1, 0x7f, 0xad, 0xc0, 0xda, 0x92, 0x58, 0xa2, 0x6b, 0x8c, 0x22, 0x69, 0xf3, 0xcb, 0x5d, 0xa7, - 0x81, 0xcd, 0xfd, 0x80, 0x10, 0x6a, 0xa6, 0x6c, 0x66, 0x19, 0x11, 0x6a, 0x86, 0x7b, 0xee, 0x12, - 0x34, 0xd4, 0x19, 0x2d, 0xb9, 0xd6, 0x2c, 0x75, 0x75, 0x86, 0x6b, 0xbd, 0x0b, 0xed, 0x48, 0x4c, - 0xdc, 0x88, 0xbf, 0xe4, 0xfa, 0xd6, 0x76, 0x7f, 0xe7, 0xc3, 0xb7, 0xec, 0x87, 0xed, 0x27, 0x62, - 0xf2, 0x04, 0x69, 0x9d, 0x56, 0x64, 0x9e, 0x46, 0xbf, 0x86, 0x96, 0x85, 0xb2, 0x36, 0xd4, 0x1f, - 0xf2, 0x93, 0x6c, 0x32, 0x78, 0x87, 0xb5, 0xa0, 0x86, 0x3d, 0x06, 0x15, 0x7c, 0xfa, 0xca, 0x4b, - 0xe3, 0xc1, 0x0a, 0xa2, 0x1f, 0xa5, 0xa9, 0x48, 0x07, 0x55, 0x7c, 0x3c, 0xf4, 0xe2, 0xd0, 0x1f, - 0xd4, 0xf0, 0xf1, 0xb1, 0xa7, 0xbc, 0x68, 0x50, 0x1f, 0xfd, 0xb6, 0x0e, 0xad, 0x43, 0xf3, 0x76, - 0xf6, 0x10, 0x7a, 0xf9, 0x9f, 0x03, 0x5d, 0x1c, 0xa3, 0x3a, 0x5c, 0x7c, 0xa0, 0x18, 0x55, 0x37, - 0x29, 0xb5, 0x16, 0xff, 0x62, 0x68, 0x65, 0xe9, 0x2f, 0x86, 0xde, 0x85, 0xea, 0x8b, 0xf4, 0x7c, - 0xbe, 0xb4, 0xf1, 0x30, 0xf2, 0x62, 0x07, 0xc1, 0xec, 0x3e, 0x74, 0x28, 0x7b, 0x24, 0xc9, 0xd1, - 0x32, 0x71, 0x9d, 0xf2, 0x1f, 0x4f, 0x11, 0xdc, 0x01, 0x24, 0x32, 0xce, 0xd8, 0x36, 0xb4, 0xfc, - 0x69, 0x18, 0x05, 0x29, 0x8f, 0x4d, 0xd9, 0x30, 0x5b, 0x9e, 0xb2, 0x93, 0xd3, 0xb0, 0x3f, 0x84, - 0x41, 0x58, 0xc4, 0xa5, 0x8a, 0x64, 0xe1, 0x9c, 0x72, 0x2b, 0x45, 0xae, 0x9c, 0xd5, 0x12, 0x39, - 0x59, 0xff, 0xe2, 0x5a, 0x73, 0xb3, 0x7c, 0xad, 0x59, 0xff, 0x39, 0x0c, 0x99, 0xe8, 0x56, 0x7e, - 0xaa, 0x45, 0x0b, 0x7d, 0xcb, 0xf8, 0x55, 0xed, 0xc5, 0x63, 0x80, 0xf5, 0x0a, 0x8c, 0x7f, 0xf5, - 0x21, 0xf4, 0xd1, 0x5f, 0x73, 0xb5, 0x9b, 0x87, 0x4a, 0x17, 0xcc, 0xbf, 0x30, 0x64, 0x72, 0xfa, - 0x10, 0x1d, 0x3d, 0x14, 0xc6, 0x9b, 0xd0, 0xb7, 0xdf, 0x62, 0x6a, 0xcb, 0x3a, 0x26, 0x99, 0x68, - 0xa0, 0xba, 0xba, 0x6c, 0x1b, 0xd6, 0xfd, 0xa9, 0x17, 0xc7, 0x3c, 0x72, 0x4f, 0xb2, 0xf1, 0xd8, - 0x5a, 0xd8, 0x2e, 0x85, 0x4d, 0xd7, 0x0c, 0xea, 0x01, 0x61, 0xc8, 0xd0, 0x8e, 0xa0, 0x17, 0x87, - 0x91, 0xce, 0x0d, 0x90, 0x37, 0xd1, 0x23, 0xca, 0x4e, 0x1c, 0x46, 0x94, 0x1c, 0x40, 0x9f, 0xe2, - 0x97, 0x30, 0xc8, 0xb2, 0x30, 0x90, 0xae, 0x12, 0xf6, 0x7f, 0x75, 0x4c, 0x84, 0xb9, 0x14, 0xb3, - 0x79, 0x9e, 0x85, 0xc1, 0xb1, 0x30, 0xff, 0xac, 0xd3, 0x23, 0x7a, 0xdb, 0x1c, 0xfd, 0x12, 0xba, - 0x65, 0xd9, 0x41, 0x59, 0xa4, 0x43, 0xf5, 0xe0, 0x1d, 0x06, 0xd0, 0x78, 0x2a, 0xd2, 0x99, 0x17, - 0x0d, 0x2a, 0xf8, 0xac, 0x2f, 0xfb, 0x0f, 0x56, 0x58, 0x17, 0x5a, 0xf6, 0x90, 0x38, 0xa8, 0x9a, - 0xb4, 0xcd, 0xcf, 0xa1, 0x65, 0xff, 0x2e, 0x88, 0xfe, 0x6a, 0x45, 0x04, 0x5c, 0x7b, 0xbd, 0xa6, - 0x86, 0x0f, 0x01, 0xe4, 0xf1, 0xda, 0x7f, 0x00, 0x5b, 0x29, 0xfe, 0x01, 0x6c, 0xf4, 0x47, 0xd0, - 0x2d, 0x4f, 0xd1, 0x06, 0x22, 0x2b, 0x45, 0x20, 0xf2, 0x82, 0x5e, 0x94, 0x5b, 0x4e, 0xc5, 0xcc, - 0x2d, 0x39, 0x66, 0x2d, 0x04, 0xe0, 0x6b, 0x46, 0xff, 0xb0, 0x02, 0x75, 0x3a, 0x3a, 0x91, 0x29, - 0xc6, 0x87, 0x62, 0x07, 0xd5, 0x9d, 0x36, 0x41, 0xfe, 0x1f, 0xee, 0x2c, 0xe5, 0x89, 0xa9, 0xda, - 0x5b, 0x13, 0x53, 0xb7, 0xff, 0xac, 0x02, 0x0d, 0xfd, 0xaf, 0x6b, 0x6c, 0x0d, 0x7a, 0xcf, 0xe3, - 0xd3, 0x58, 0xbc, 0x8a, 0x35, 0x60, 0xf0, 0x0e, 0x5b, 0x87, 0x55, 0xcb, 0x7b, 0xf3, 0xf7, 0x6e, - 0x83, 0x0a, 0x1b, 0x40, 0x97, 0x56, 0xd7, 0x42, 0x56, 0xd8, 0xbb, 0x30, 0x34, 0xd6, 0xf4, 0x21, - 0x2a, 0x63, 0xa1, 0xc2, 0xf1, 0xb9, 0xc5, 0x56, 0xd9, 0x2a, 0x74, 0x8e, 0x94, 0x48, 0x8e, 0x78, - 0x1c, 0x84, 0xf1, 0x64, 0x50, 0x63, 0x43, 0xd8, 0xb0, 0xa3, 0xea, 0x7f, 0x26, 0x7b, 0x1c, 0xc6, - 0xa1, 0x9c, 0x0e, 0xea, 0xec, 0x1a, 0x5c, 0xbe, 0x08, 0xb3, 0xeb, 0x9f, 0x0e, 0x1a, 0xb7, 0x3f, - 0x03, 0xb6, 0xfc, 0x47, 0x66, 0x38, 0xfa, 0x13, 0x3e, 0xf1, 0xfc, 0xf3, 0xbd, 0x48, 0x48, 0x14, - 0x8a, 0x1e, 0xb4, 0x8b, 0x5e, 0x95, 0xdb, 0x8f, 0xa1, 0xa1, 0xff, 0x79, 0xae, 0xf4, 0x7d, 0x1a, - 0x30, 0x78, 0x07, 0x3b, 0xa3, 0xc9, 0x09, 0xe3, 0xc9, 0x53, 0x7e, 0xa6, 0xb4, 0x22, 0x7c, 0xe2, - 0x49, 0x35, 0x58, 0x61, 0x7d, 0x00, 0xf3, 0x09, 0x8f, 0xe2, 0x60, 0x50, 0x7d, 0xb0, 0xf7, 0xe7, - 0xbf, 0x7b, 0xaf, 0xf2, 0x17, 0xbf, 0x7b, 0xaf, 0xf2, 0x5f, 0x7e, 0xf7, 0xde, 0x3b, 0x7f, 0xfa, - 0x97, 0xef, 0x55, 0xbe, 0xbe, 0x5f, 0xfa, 0x5f, 0x3d, 0x63, 0x89, 0xa8, 0x16, 0xe0, 0x6e, 0x6e, - 0x96, 0xee, 0x26, 0xa7, 0x93, 0xbb, 0xc9, 0xc9, 0x5d, 0x2b, 0xe7, 0x27, 0x0d, 0xfa, 0xbb, 0xbc, - 0x4f, 0xff, 0x6f, 0x00, 0x00, 0x00, 0xff, 0xff, 0x1b, 0xfe, 0xe5, 0x5d, 0xad, 0x4f, 0x00, 0x00, + 0xbc, 0xa0, 0xac, 0xe0, 0x7d, 0xa8, 0xc5, 0xb6, 0x92, 0x3d, 0xc7, 0xd3, 0x9d, 0x6c, 0x82, 0xb3, + 0xdb, 0x90, 0x5f, 0x80, 0x31, 0x2e, 0xc9, 0xc5, 0x17, 0x64, 0xb6, 0xa1, 0x9d, 0xff, 0xc7, 0x84, + 0xf1, 0x42, 0xd6, 0xb7, 0x8a, 0x7f, 0x9d, 0x38, 0xb2, 0x4f, 0x4e, 0x41, 0xf6, 0xe6, 0xe4, 0x78, + 0xe7, 0x47, 0x25, 0xc7, 0xcd, 0x11, 0xcd, 0x17, 0xb1, 0x54, 0x26, 0x38, 0xd6, 0x0c, 0xe5, 0x2e, + 0x36, 0xd9, 0xef, 0x41, 0x2f, 0xe5, 0xfe, 0x4b, 0x77, 0x2a, 0xc7, 0xfa, 0x15, 0xbd, 0xf2, 0x25, + 0xbc, 0xa9, 0x1c, 0x7f, 0x4d, 0x89, 0x7b, 0x73, 0x62, 0xea, 0x20, 0xed, 0xbe, 0x1c, 0xd3, 0xa8, + 0x9f, 0xc0, 0xea, 0x94, 0x4f, 0x8f, 0x79, 0x2a, 0x27, 0x61, 0x62, 0x55, 0xed, 0x32, 0x15, 0x18, + 0xf4, 0x0b, 0x84, 0x9e, 0xcb, 0xf0, 0x1f, 0x55, 0xa0, 0x85, 0x5c, 0x44, 0x59, 0x62, 0x0c, 0x6a, + 0x53, 0x3f, 0xc9, 0x8c, 0x23, 0x4c, 0xcf, 0xe6, 0x7f, 0x2b, 0xb4, 0x94, 0x98, 0xff, 0xad, 0xa0, + 0x35, 0xd4, 0x29, 0x1f, 0x7a, 0xd6, 0xf7, 0xc0, 0xcf, 0x28, 0x2a, 0xa8, 0x25, 0xc3, 0x36, 0xd9, + 0x25, 0x68, 0xf8, 0x31, 0x9d, 0x86, 0x75, 0xea, 0xac, 0xee, 0xc7, 0x78, 0x0a, 0xd6, 0xe0, 0xa2, + 0x1a, 0xba, 0xee, 0xc7, 0x7b, 0xc1, 0xe9, 0xf0, 0x3f, 0x54, 0x60, 0xf5, 0x20, 0x15, 0x3e, 0x97, + 0xf2, 0x09, 0x1a, 0x72, 0x4a, 0x53, 0xe0, 0x1b, 0x29, 0xaa, 0xab, 0x33, 0x02, 0xf4, 0x8c, 0x32, + 0xac, 0x43, 0x15, 0xf9, 0x71, 0xa3, 0xea, 0xb4, 0x09, 0x42, 0xa7, 0x8d, 0x1c, 0x5d, 0x4a, 0x7f, + 0x6b, 0x34, 0xc5, 0x83, 0x6f, 0xc2, 0x72, 0x91, 0x58, 0x29, 0x65, 0xea, 0x8b, 0x6b, 0xa6, 0x34, + 0xca, 0x75, 0xe8, 0x98, 0x7a, 0x0a, 0x1a, 0x46, 0x87, 0xf8, 0x41, 0x83, 0x0e, 0xcd, 0x2c, 0xb4, + 0x72, 0x20, 0xbc, 0x0e, 0xea, 0x6b, 0x75, 0x81, 0xe8, 0xe1, 0xdf, 0x87, 0xfe, 0x41, 0xca, 0x13, + 0x2f, 0xe5, 0x54, 0x5f, 0x41, 0x2c, 0xde, 0x80, 0x46, 0xc4, 0xe3, 0xb1, 0x49, 0xe9, 0x57, 0x1d, + 0xd3, 0xca, 0xff, 0x72, 0x64, 0xa9, 0xf4, 0x97, 0x23, 0xc8, 0xea, 0x94, 0x7b, 0xe6, 0x9f, 0x49, + 0xe8, 0x19, 0xb7, 0x20, 0x1e, 0x19, 0xf5, 0xb9, 0xa8, 0xe5, 0xe8, 0x86, 0xb9, 0xcb, 0x76, 0x1c, + 0xc6, 0x54, 0x9b, 0x46, 0x77, 0xd9, 0x1e, 0x84, 0xf1, 0xf0, 0x5f, 0xd6, 0xa1, 0x63, 0xf8, 0x49, + 0x2f, 0xd7, 0x6b, 0x59, 0xc9, 0xd7, 0xb2, 0x0f, 0x55, 0xf9, 0x22, 0x32, 0x8b, 0x8b, 0x8f, 0xec, + 0x33, 0xa8, 0x46, 0xe1, 0xd4, 0x1c, 0x71, 0xae, 0xcd, 0x98, 0xab, 0xd9, 0x55, 0x31, 0x82, 0x87, + 0xd4, 0xa8, 0x23, 0xe9, 0x1e, 0x2c, 0x8a, 0xb8, 0xe1, 0x24, 0x9a, 0x8e, 0x53, 0xdc, 0x47, 0xc8, + 0x23, 0xcf, 0xa7, 0x32, 0x03, 0xab, 0x1c, 0x7a, 0x4e, 0xdb, 0x40, 0xf6, 0x02, 0xf6, 0x39, 0xb4, + 0xf2, 0x40, 0xa5, 0x3d, 0xd4, 0xa8, 0xd3, 0x78, 0x6b, 0xf7, 0xe9, 0xd1, 0x69, 0x6c, 0x23, 0x91, + 0xe6, 0x65, 0x39, 0x25, 0xfb, 0x03, 0xe8, 0x4a, 0x2e, 0xa5, 0xbe, 0xa8, 0x38, 0x12, 0x46, 0x69, + 0x5c, 0x2a, 0x9f, 0x57, 0x08, 0x8b, 0x5f, 0x6d, 0xb7, 0x88, 0x2c, 0x40, 0xec, 0x6b, 0x58, 0xb6, + 0xfd, 0x23, 0x31, 0x1e, 0xe7, 0x81, 0xf4, 0x6b, 0x0b, 0x23, 0x3c, 0x21, 0x74, 0x69, 0x9c, 0x9e, + 0x2c, 0x23, 0xd8, 0x57, 0xb0, 0x9c, 0xe8, 0x35, 0x76, 0x4d, 0x29, 0x8f, 0x56, 0x3e, 0x57, 0x67, + 0xbc, 0xab, 0x19, 0x19, 0x28, 0xae, 0x12, 0x15, 0x70, 0xb9, 0x78, 0x69, 0x57, 0x67, 0x56, 0x66, + 0x2f, 0xed, 0x72, 0xd8, 0x30, 0x7f, 0x56, 0x31, 0x4a, 0x3d, 0xba, 0x7c, 0xa8, 0x2b, 0x38, 0x6c, + 0xcd, 0xdd, 0xdd, 0x85, 0x15, 0xc3, 0x17, 0x6e, 0xe9, 0xcb, 0xa9, 0x8f, 0x4d, 0x17, 0x2a, 0xf0, + 0x30, 0xc5, 0x16, 0xeb, 0xe9, 0x39, 0x28, 0xb6, 0x05, 0x6b, 0xe6, 0x35, 0xfc, 0x94, 0xfb, 0x99, + 0xb9, 0x77, 0x4d, 0x2a, 0xaa, 0xeb, 0xac, 0x6a, 0xd4, 0x23, 0x8b, 0xd9, 0x0b, 0xae, 0x7e, 0x05, + 0x57, 0x2e, 0x7c, 0xc5, 0xdb, 0xca, 0x07, 0x7a, 0xe5, 0x5b, 0x91, 0xff, 0xbb, 0x0a, 0x9d, 0xd2, + 0xd2, 0xd1, 0x9f, 0xe4, 0x48, 0x9e, 0xda, 0x22, 0x21, 0x7c, 0x46, 0xd8, 0x44, 0x48, 0x5b, 0xf3, + 0x42, 0xcf, 0x08, 0x4b, 0x45, 0x9e, 0xfb, 0xa7, 0x67, 0x64, 0xa8, 0x39, 0x8d, 0x9b, 0xe9, 0xd7, + 0xf4, 0x6d, 0xde, 0x02, 0xb8, 0x17, 0xd0, 0xbf, 0xe9, 0x78, 0xca, 0x3b, 0xf6, 0xa4, 0xad, 0xb6, + 0xca, 0xdb, 0xa8, 0xd5, 0x5e, 0xf2, 0x14, 0xe7, 0x62, 0x73, 0x9d, 0xa6, 0x89, 0x02, 0x4f, 0x86, + 0xe0, 0xb5, 0x88, 0x75, 0x9e, 0xb3, 0xeb, 0xb4, 0x10, 0xf0, 0xad, 0x88, 0xa9, 0x9b, 0x11, 0x6f, + 0x93, 0xaf, 0xb7, 0x4d, 0x54, 0xf7, 0x2f, 0x32, 0x8e, 0x6e, 0x78, 0x40, 0xf7, 0x35, 0xda, 0x4e, + 0x93, 0xda, 0xba, 0x84, 0x80, 0xce, 0x0b, 0xaf, 0xbc, 0x50, 0xd1, 0x3e, 0x12, 0x99, 0x32, 0x12, + 0xb0, 0x82, 0x88, 0x6f, 0xbc, 0x50, 0x1d, 0x69, 0x30, 0xbb, 0x6f, 0xae, 0x61, 0x95, 0x69, 0x5d, + 0x3c, 0xec, 0xe8, 0x28, 0x1f, 0x9b, 0xa3, 0x3f, 0xe4, 0xf4, 0x7f, 0x32, 0x53, 0x4f, 0xa5, 0xe1, + 0xa9, 0x88, 0xd1, 0xec, 0xab, 0xf0, 0x25, 0x2f, 0xfe, 0xc1, 0xa7, 0xe5, 0xac, 0xe5, 0xc8, 0xa7, + 0x84, 0xa3, 0xc4, 0xd0, 0x73, 0xd8, 0xe4, 0xa7, 0x49, 0x14, 0xfa, 0xe1, 0xdc, 0x05, 0x48, 0xd7, + 0xf7, 0xa4, 0x72, 0x53, 0xae, 0xb2, 0x34, 0x96, 0x14, 0xc0, 0x32, 0x55, 0x28, 0x1f, 0x5a, 0xfa, + 0xf2, 0xa5, 0xc8, 0x5d, 0x4f, 0x2a, 0x47, 0xd3, 0x3e, 0xcd, 0xa2, 0x08, 0x99, 0x90, 0x27, 0xa4, + 0x74, 0x2d, 0x4a, 0x53, 0xea, 0x54, 0xd4, 0xf0, 0xbf, 0x56, 0x60, 0x75, 0x61, 0xdb, 0xa1, 0xeb, + 0x8f, 0x5b, 0xce, 0xe6, 0xcf, 0xbb, 0x4e, 0x03, 0x9b, 0x7b, 0x01, 0x21, 0xd4, 0x54, 0xd9, 0xcc, + 0x39, 0x22, 0xd4, 0x14, 0x75, 0xca, 0x25, 0x68, 0xa8, 0x53, 0x5a, 0x72, 0xad, 0x39, 0xeb, 0xea, + 0x14, 0xd7, 0x7a, 0x07, 0xda, 0x91, 0x18, 0xbb, 0x11, 0x7f, 0xc9, 0xf5, 0xad, 0xf4, 0xe5, 0xed, + 0x8f, 0xde, 0xb0, 0xdf, 0xb7, 0x9e, 0x88, 0xf1, 0x13, 0xa4, 0x75, 0x5a, 0x91, 0x79, 0x1a, 0xfe, + 0x1a, 0x5a, 0x16, 0xca, 0xda, 0x50, 0x7f, 0xc8, 0x8f, 0xb3, 0x71, 0xff, 0x1d, 0xd6, 0x82, 0x1a, + 0xf6, 0xe8, 0x57, 0xf0, 0xe9, 0x1b, 0x2f, 0x8d, 0xfb, 0x4b, 0x88, 0x7e, 0x94, 0xa6, 0x22, 0xed, + 0x57, 0xf1, 0xf1, 0xc0, 0x8b, 0x43, 0xbf, 0x5f, 0xc3, 0xc7, 0xc7, 0x9e, 0xf2, 0xa2, 0x7e, 0x7d, + 0xf8, 0xdb, 0x3a, 0xb4, 0x0e, 0xcc, 0xdb, 0xd9, 0x43, 0xe8, 0xe5, 0x7f, 0x7e, 0x74, 0x7e, 0x0c, + 0xee, 0x60, 0xfe, 0x81, 0x62, 0x70, 0xdd, 0xa4, 0xd4, 0x9a, 0xff, 0x0b, 0xa5, 0xa5, 0x85, 0xbf, + 0x50, 0x7a, 0x17, 0xaa, 0x2f, 0xd2, 0xb3, 0xd9, 0xd2, 0xcd, 0x83, 0xc8, 0x8b, 0x1d, 0x04, 0xb3, + 0xfb, 0xd0, 0xa1, 0xec, 0x98, 0x24, 0x47, 0xd2, 0xc4, 0xad, 0xca, 0x7f, 0xac, 0x45, 0x70, 0x07, + 0x90, 0xc8, 0x38, 0x9b, 0x5b, 0xd0, 0xf2, 0x27, 0x61, 0x14, 0xa4, 0x3c, 0x36, 0x65, 0xd1, 0x6c, + 0x71, 0xca, 0x4e, 0x4e, 0xc3, 0xfe, 0x10, 0xfa, 0x61, 0x11, 0x77, 0x2b, 0x92, 0xa1, 0x33, 0xca, + 0xbb, 0x14, 0x99, 0x73, 0x56, 0x4a, 0xe4, 0xe4, 0xdd, 0x14, 0xd7, 0xb6, 0x9b, 0xe5, 0x6b, 0xdb, + 0xfa, 0xcf, 0x6f, 0xc8, 0x05, 0x69, 0xe5, 0xa7, 0x76, 0xf4, 0x40, 0x6e, 0x19, 0xbf, 0xb1, 0x3d, + 0x7f, 0xcc, 0xb1, 0x5e, 0x8f, 0xf1, 0x1f, 0x3f, 0x82, 0x65, 0xf4, 0x47, 0x5d, 0xed, 0xc6, 0xa2, + 0x51, 0x01, 0xf3, 0x2f, 0x13, 0x99, 0x9c, 0x3c, 0x44, 0x47, 0x16, 0x85, 0xf1, 0x26, 0x2c, 0xdb, + 0x6f, 0x31, 0xb5, 0x73, 0x1d, 0x93, 0x2c, 0x35, 0x50, 0x5d, 0x3d, 0xb7, 0x05, 0x6b, 0xfe, 0xc4, + 0x8b, 0x63, 0x1e, 0xb9, 0xc7, 0xd9, 0x68, 0x64, 0x3d, 0x88, 0x2e, 0x85, 0x85, 0x57, 0x0d, 0xea, + 0x01, 0x61, 0xc8, 0x91, 0x18, 0x42, 0x2f, 0x0e, 0x23, 0x9d, 0xfb, 0x20, 0x6f, 0xa9, 0x47, 0x94, + 0x9d, 0x38, 0x8c, 0x28, 0xf9, 0x81, 0x3e, 0xd3, 0xaf, 0xa0, 0x9f, 0x65, 0x61, 0x20, 0x5d, 0x25, + 0xec, 0xff, 0x06, 0x99, 0x08, 0x7a, 0x29, 0x26, 0xf5, 0x3c, 0x0b, 0x83, 0x23, 0x61, 0xfe, 0x39, + 0xa8, 0x47, 0xf4, 0xb6, 0x39, 0xfc, 0x15, 0x74, 0xcb, 0xb2, 0x83, 0xb2, 0x48, 0x41, 0x83, 0xfe, + 0x3b, 0x0c, 0xa0, 0xf1, 0x54, 0xa4, 0x53, 0x2f, 0xea, 0x57, 0xf0, 0x59, 0x2b, 0xf3, 0xfe, 0x12, + 0xeb, 0x42, 0xcb, 0x1e, 0x82, 0xfb, 0x55, 0x93, 0x96, 0xfa, 0x05, 0xb4, 0xec, 0xdf, 0x21, 0xd1, + 0x5f, 0xc9, 0x88, 0x80, 0x6b, 0xaf, 0xde, 0xd4, 0x28, 0x22, 0x80, 0x3c, 0x7a, 0xfb, 0x0f, 0x67, + 0x4b, 0xc5, 0x3f, 0x9c, 0x0d, 0xff, 0x08, 0xba, 0xe5, 0x29, 0xda, 0x40, 0x6b, 0xa5, 0x08, 0xb4, + 0x9e, 0xd3, 0x8b, 0x72, 0xe7, 0xa9, 0x98, 0xba, 0x25, 0xc7, 0xb3, 0x85, 0x00, 0x7c, 0xcd, 0xf0, + 0x1f, 0x56, 0xa0, 0x4e, 0x47, 0x43, 0x72, 0x35, 0xf0, 0xa1, 0xd8, 0x41, 0x75, 0xa7, 0x4d, 0x90, + 0xff, 0x87, 0x3b, 0x59, 0x79, 0xe2, 0xad, 0xf6, 0xc6, 0xc4, 0xdb, 0xed, 0x3f, 0xab, 0x40, 0x43, + 0xff, 0xab, 0x1c, 0x5b, 0x85, 0xde, 0xf3, 0xf8, 0x24, 0x16, 0xaf, 0x62, 0x0d, 0xe8, 0xbf, 0xc3, + 0xd6, 0x60, 0xc5, 0xf2, 0xde, 0xfc, 0x7d, 0x5d, 0xbf, 0xc2, 0xfa, 0xd0, 0xa5, 0xd5, 0xb5, 0x90, + 0x25, 0xf6, 0x2e, 0x0c, 0x8c, 0xb7, 0xf0, 0x10, 0x95, 0xb1, 0x50, 0xe1, 0xe8, 0xcc, 0x62, 0xab, + 0x6c, 0x05, 0x3a, 0x87, 0x4a, 0x24, 0x87, 0x3c, 0x0e, 0xc2, 0x78, 0xdc, 0xaf, 0xb1, 0x01, 0xac, + 0xdb, 0x51, 0xf5, 0x3f, 0xaf, 0x3d, 0x0e, 0xe3, 0x50, 0x4e, 0xfa, 0x75, 0x76, 0x0d, 0x2e, 0x9f, + 0x87, 0xd9, 0xf1, 0x4f, 0xfa, 0x8d, 0xdb, 0x9f, 0x03, 0x5b, 0xfc, 0xa3, 0x36, 0x1c, 0xfd, 0x09, + 0x1f, 0x7b, 0xfe, 0xd9, 0x6e, 0x24, 0x24, 0x0a, 0x45, 0x0f, 0xda, 0x45, 0xaf, 0xca, 0xed, 0xc7, + 0xd0, 0xd0, 0xff, 0xac, 0x57, 0xfa, 0x3e, 0x0d, 0xe8, 0xbf, 0x83, 0x9d, 0xd1, 0xe4, 0x84, 0xf1, + 0xf8, 0x29, 0x3f, 0x55, 0x5a, 0x11, 0x3e, 0xf1, 0xa4, 0xea, 0x2f, 0xb1, 0x65, 0x00, 0xf3, 0x09, + 0x8f, 0xe2, 0xa0, 0x5f, 0x7d, 0xb0, 0xfb, 0xe7, 0xbf, 0x7b, 0xbf, 0xf2, 0x17, 0xbf, 0x7b, 0xbf, + 0xf2, 0x5f, 0x7e, 0xf7, 0xfe, 0x3b, 0x7f, 0xfa, 0x97, 0xef, 0x57, 0xbe, 0xbd, 0x5f, 0xfa, 0xdf, + 0x40, 0x63, 0x89, 0xa8, 0xd6, 0xe1, 0x6e, 0x6e, 0x96, 0xee, 0x26, 0x27, 0xe3, 0xbb, 0xc9, 0xf1, + 0x5d, 0x2b, 0xe7, 0xc7, 0x0d, 0xfa, 0x3b, 0xc0, 0xcf, 0xfe, 0x6f, 0x00, 0x00, 0x00, 0xff, 0xff, + 0x8e, 0x30, 0xc6, 0xcd, 0x8d, 0x50, 0x00, 0x00, } func (m *Message) Marshal() (dAtA []byte, err error) { @@ -11978,6 +12004,30 @@ func (m *ProcessInfo) MarshalToSizedBuffer(dAtA []byte) (int, error) { i -= len(m.XXX_unrecognized) copy(dAtA[i:], m.XXX_unrecognized) } + if len(m.RemoteExecutionId) > 0 { + i -= len(m.RemoteExecutionId) + copy(dAtA[i:], m.RemoteExecutionId) + i = encodeVarintPipeline(dAtA, i, uint64(len(m.RemoteExecutionId))) + i-- + dAtA[i] = 0x62 + } + if len(m.RemoteFragmentCounts) > 0 { + for k := range m.RemoteFragmentCounts { + v := m.RemoteFragmentCounts[k] + baseI := i + i = encodeVarintPipeline(dAtA, i, uint64(v)) + i-- + dAtA[i] = 0x10 + i -= len(k) + copy(dAtA[i:], k) + i = encodeVarintPipeline(dAtA, i, uint64(len(k))) + i-- + dAtA[i] = 0xa + i = encodeVarintPipeline(dAtA, i, uint64(baseI-i)) + i-- + dAtA[i] = 0x5a + } + } if m.AffectedRows != 0 { i = encodeVarintPipeline(dAtA, i, uint64(m.AffectedRows)) i-- @@ -14884,6 +14934,18 @@ func (m *ProcessInfo) ProtoSize() (n int) { if m.AffectedRows != 0 { n += 1 + sovPipeline(uint64(m.AffectedRows)) } + if len(m.RemoteFragmentCounts) > 0 { + for k, v := range m.RemoteFragmentCounts { + _ = k + _ = v + mapEntrySize := 1 + len(k) + sovPipeline(uint64(len(k))) + 1 + sovPipeline(uint64(v)) + n += mapEntrySize + 1 + sovPipeline(uint64(mapEntrySize)) + } + } + l = len(m.RemoteExecutionId) + if l > 0 { + n += 1 + l + sovPipeline(uint64(l)) + } if m.XXX_unrecognized != nil { n += len(m.XXX_unrecognized) } @@ -30669,6 +30731,153 @@ func (m *ProcessInfo) Unmarshal(dAtA []byte) error { break } } + case 11: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field RemoteFragmentCounts", wireType) + } + var msglen int + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + msglen |= int(b&0x7F) << shift + if b < 0x80 { + break + } + } + if msglen < 0 { + return ErrInvalidLengthPipeline + } + postIndex := iNdEx + msglen + if postIndex < 0 { + return ErrInvalidLengthPipeline + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + if m.RemoteFragmentCounts == nil { + m.RemoteFragmentCounts = make(map[string]uint32) + } + var mapkey string + var mapvalue uint32 + for iNdEx < postIndex { + entryPreIndex := iNdEx + var wire uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + wire |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + fieldNum := int32(wire >> 3) + if fieldNum == 1 { + var stringLenmapkey uint64 + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + stringLenmapkey |= uint64(b&0x7F) << shift + if b < 0x80 { + break + } + } + intStringLenmapkey := int(stringLenmapkey) + if intStringLenmapkey < 0 { + return ErrInvalidLengthPipeline + } + postStringIndexmapkey := iNdEx + intStringLenmapkey + if postStringIndexmapkey < 0 { + return ErrInvalidLengthPipeline + } + if postStringIndexmapkey > l { + return io.ErrUnexpectedEOF + } + mapkey = string(dAtA[iNdEx:postStringIndexmapkey]) + iNdEx = postStringIndexmapkey + } else if fieldNum == 2 { + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + mapvalue |= uint32(b&0x7F) << shift + if b < 0x80 { + break + } + } + } else { + iNdEx = entryPreIndex + skippy, err := skipPipeline(dAtA[iNdEx:]) + if err != nil { + return err + } + if (skippy < 0) || (iNdEx+skippy) < 0 { + return ErrInvalidLengthPipeline + } + if (iNdEx + skippy) > postIndex { + return io.ErrUnexpectedEOF + } + iNdEx += skippy + } + } + m.RemoteFragmentCounts[mapkey] = mapvalue + iNdEx = postIndex + case 12: + if wireType != 2 { + return fmt.Errorf("proto: wrong wireType = %d for field RemoteExecutionId", wireType) + } + var byteLen int + for shift := uint(0); ; shift += 7 { + if shift >= 64 { + return ErrIntOverflowPipeline + } + if iNdEx >= l { + return io.ErrUnexpectedEOF + } + b := dAtA[iNdEx] + iNdEx++ + byteLen |= int(b&0x7F) << shift + if b < 0x80 { + break + } + } + if byteLen < 0 { + return ErrInvalidLengthPipeline + } + postIndex := iNdEx + byteLen + if postIndex < 0 { + return ErrInvalidLengthPipeline + } + if postIndex > l { + return io.ErrUnexpectedEOF + } + m.RemoteExecutionId = append(m.RemoteExecutionId[:0], dAtA[iNdEx:postIndex]...) + if m.RemoteExecutionId == nil { + m.RemoteExecutionId = []byte{} + } + iNdEx = postIndex default: iNdEx = preIndex skippy, err := skipPipeline(dAtA[iNdEx:]) diff --git a/pkg/sql/colexec/aggexec/maxby.go b/pkg/sql/colexec/aggexec/maxby.go index 203da3e60b83d..83a24510e0261 100644 --- a/pkg/sql/colexec/aggexec/maxby.go +++ b/pkg/sql/colexec/aggexec/maxby.go @@ -332,7 +332,15 @@ func compactMaxByStateVector(vec *vector.Vector, usage *maxByVarlenaUsage, mp *m usage.staleBytes <= usage.liveBytes+maxByVarlenaCompactionSlack { return nil } - compact, err := vec.CloneToFlatCompact(mp) + var ( + compact *vector.Vector + err error + ) + if selection := vec.AllocationAccountSelection(); selection != nil { + compact, err = vec.CloneToFlatCompactWithAllocation(mp, selection) + } else { + compact, err = vec.CloneToFlatCompact(mp) + } if err != nil { return err } diff --git a/pkg/sql/colexec/aggexec/maxby_test.go b/pkg/sql/colexec/aggexec/maxby_test.go index a3d8adb92e30e..9ec5574c2b537 100644 --- a/pkg/sql/colexec/aggexec/maxby_test.go +++ b/pkg/sql/colexec/aggexec/maxby_test.go @@ -80,6 +80,54 @@ func TestMaxByCompactsReplacedVarlenaState(t *testing.T) { require.Less(t, state.Allocated(), 2<<20, "winner state must be bounded by live groups, not by replaced input rows") } +func TestCompactMaxByStateVectorPreservesAllocationOwner(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(8 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + mpool.AllocationSite(2), + mpool.AllocationSite(3), + mpool.AllocationSite(4), + ) + require.NoError(t, err) + vec := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) + require.NoError(t, vec.SetAllocationAccount(selection)) + defer func() { + vec.Free(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) + }() + + value := []byte(strings.Repeat("x", 4096)) + require.NoError(t, vector.AppendBytes(vec, value, false, mp)) + for i := 0; i < 400; i++ { + value[0] = byte(i) + require.NoError(t, vec.SetRawBytesAt(0, value, mp)) + } + before := account.Snapshot() + require.Greater(t, before.Used, uint64(maxByVarlenaCompactionSlack)) + + usage := &maxByVarlenaUsage{ + liveBytes: maxByAreaBytes(vec.GetRawBytesAt(0)), + staleBytes: len(vec.GetArea()) - maxByAreaBytes(vec.GetRawBytesAt(0)), + } + require.NoError(t, compactMaxByStateVector(vec, usage, mp)) + require.Same(t, selection, vec.AllocationAccountSelection()) + require.Equal(t, value, vec.GetBytesAt(0)) + after := account.Snapshot() + require.Less(t, after.Used, before.Used) + require.GreaterOrEqual(t, after.Peak, before.Used) +} + func TestMaxByTracksManyGroupVarlenaUsageIncrementally(t *testing.T) { mp := mpool.MustNewZero() params := []types.Type{types.T_varchar.ToType(), types.T_int64.ToType(), types.T_int64.ToType()} diff --git a/pkg/sql/colexec/allocation_state.go b/pkg/sql/colexec/allocation_state.go new file mode 100644 index 0000000000000..f1a2a9c7a78fb --- /dev/null +++ b/pkg/sql/colexec/allocation_state.go @@ -0,0 +1,65 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package colexec + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +// NewAccountedBitmap creates a bitmap whose complete backing capacity belongs +// to the statement allocation account. The caller owns the returned bitmap +// until FreeAccountedBitmap or an explicit ownership transfer. +func NewAccountedBitmap( + rows int64, + mp *mpool.MPool, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) (*bitmap.Bitmap, error) { + if rows < 0 || rows > math.MaxInt64-63 || mp == nil || account == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + words := (rows + 63) / 64 + if words > int64(math.MaxInt) { + return nil, mpool.ErrAllocationAccountInvalid + } + storage, err := mpool.MakeSliceAccounted[uint64]( + int(words), + mp, + account, + owner, + site, + ) + if err != nil { + return nil, err + } + value := new(bitmap.Bitmap) + value.InstallExternalStorage(storage) + value.InitWithSize(rows) + return value, nil +} + +func FreeAccountedBitmap(value *bitmap.Bitmap, mp *mpool.MPool) { + if value == nil || !value.HasExternalStorage() { + return + } + storage := value.ReleaseExternalStorage() + if cap(storage) > 0 { + mpool.FreeSlice(mp, storage) + } +} diff --git a/pkg/sql/colexec/connector/types.go b/pkg/sql/colexec/connector/types.go index d7e02fb5af4ba..5169fafc3dec5 100644 --- a/pkg/sql/colexec/connector/types.go +++ b/pkg/sql/colexec/connector/types.go @@ -17,6 +17,7 @@ package connector import ( "context" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/pSpool" @@ -30,8 +31,9 @@ var _ vm.Operator = new(Connector) type Connector struct { ctr container - Reg *process.WaitRegister - cleanupSpool *pSpool.PipelineSpool + Reg *process.WaitRegister + cleanupSpool *pSpool.PipelineSpool + allocationAccount *mpool.AllocationAccount vm.OperatorBase } @@ -73,6 +75,45 @@ func (connector *Connector) WithReg(reg *process.WaitRegister) *Connector { return connector } +func (connector *Connector) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if connector.allocationAccount != nil && connector.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + connector.allocationAccount = account + return nil +} + +// ActivatesAllocationAccountLifecycle reports that Connector only participates +// in an account already required by an allocation-producing operator. +func (connector *Connector) ActivatesAllocationAccountLifecycle() bool { + return false +} + +func (connector *Connector) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if connector.allocationAccount == nil { + return nil + } + if connector.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if connector.ctr.sp != nil { + return mpool.ErrAllocationAccountInvariant + } + if connector.cleanupSpool != nil { + connector.cleanupSpool.FinalizeAfterConsumersQuiesced() + connector.cleanupSpool = nil + } + connector.allocationAccount = nil + return nil +} + func (connector *Connector) Release() { if connector != nil { reuse.Free[Connector](connector, nil) @@ -100,7 +141,11 @@ func (connector *Connector) Reset(proc *process.Process, pipelineFailed bool, er abortErr = fallbackErr } sp.Abort(abortErr) - connector.cleanupSpool = nil + if connector.allocationAccount != nil { + connector.cleanupSpool = sp + } else { + connector.cleanupSpool = nil + } } connector.ctr.sp = nil } else if terminalSignal.EventType == process.EventEnd && !terminalDelivered { @@ -146,6 +191,10 @@ func (connector *Connector) CleanupDeferredSpool() { if connector.cleanupSpool == nil { return } + if connector.allocationAccount != nil { + connector.cleanupSpool.ReleaseReusableCacheAfterProducerQuiesced() + return + } connector.cleanupSpool.ForceCleanupAfterTerminalSignal() connector.cleanupSpool = nil } diff --git a/pkg/sql/colexec/connector/types_test.go b/pkg/sql/colexec/connector/types_test.go index 34dd9b182ad94..824d38ccd8e95 100644 --- a/pkg/sql/colexec/connector/types_test.go +++ b/pkg/sql/colexec/connector/types_test.go @@ -378,6 +378,130 @@ func TestConnectorResetEndPreservesQueuedSpoolBatchUntilDeferredCleanup(t *testi require.Equal(t, int64(0), mp.CurrNB()) } +func TestConnectorAllocationClearFinalizesAbortedSpool(t *testing.T) { + testConnectorAllocationClearFinalizesSpool(t, true) +} + +func TestConnectorAccountedDeferredCleanupReleasesReusableCache(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newConnectorSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + sp.ReleaseCurrent(0) + require.Positive(t, mp.CurrNB()) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + conn := &Connector{cleanupSpool: sp} + require.NoError(t, conn.SetAllocationAccount(account)) + conn.CleanupDeferredSpool() + require.Same(t, sp, conn.cleanupSpool) + require.Zero(t, mp.CurrNB()) + require.NoError(t, conn.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestConnectorAllocationAccountContract(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 1) + require.NoError(t, err) + first, err := registry.Open(1) + require.NoError(t, err) + second, err := registry.Open(1) + require.NoError(t, err) + conn := &Connector{} + require.False(t, conn.ActivatesAllocationAccountLifecycle()) + require.ErrorIs(t, conn.SetAllocationAccount(nil), mpool.ErrAllocationAccountInvalid) + require.NoError(t, conn.SetAllocationAccount(first)) + require.ErrorIs(t, conn.SetAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + require.ErrorIs(t, conn.ClearAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + conn.ctr.sp = &pSpool.PipelineSpool{} + require.ErrorIs(t, conn.ClearAllocationAccount(first), mpool.ErrAllocationAccountInvariant) + conn.ctr.sp = nil + require.NoError(t, conn.ClearAllocationAccount(first)) + require.NoError(t, conn.ClearAllocationAccount(first)) + _, _, err = registry.CompleteTerminal(first) + require.NoError(t, err) + _, _, err = registry.CompleteTerminal(second) + require.NoError(t, err) +} + +func TestConnectorAllocationClearFinalizesTerminalSpoolPending(t *testing.T) { + testConnectorAllocationClearFinalizesSpool(t, false) +} + +func testConnectorAllocationClearFinalizesSpool(t *testing.T, abort bool) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 102, + 103, + 104, + 105, + ) + require.NoError(t, err) + src := batch.NewOffHeapWithSize(1) + require.NoError(t, src.SetAllocationAccount(selection)) + src.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, vector.AppendFixed(src.Vecs[0], int64(1), false, mp)) + src.SetRowCount(1) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + conn := &Connector{} + require.NoError(t, conn.SetAllocationAccount(account)) + if abort { + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + conn.ctr.sp = sp + conn.Reg = process.NewPipelineEdge(1, 0) + conn.Reset(nil, true, moerr.NewInternalErrorNoCtx("pipeline failed")) + require.Same(t, sp, conn.cleanupSpool) + sp.ReleaseCurrent(0) + } else { + conn.cleanupSpool = sp + sp.ForceCleanupAfterTerminalSignal() + } + conn.CleanupDeferredSpool() + require.Same(t, sp, conn.cleanupSpool) + require.NoError(t, conn.ClearAllocationAccount(account)) + require.Nil(t, conn.cleanupSpool) + src.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + func newConnectorSpoolTestBatch(t *testing.T, mp *mpool.MPool, rows int) *batch.Batch { t.Helper() src := batch.NewWithSize(1) diff --git a/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..b0c6f5132c28e --- /dev/null +++ b/pkg/sql/colexec/dedupjoin/allocation_test_helpers_test.go @@ -0,0 +1,125 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package dedupjoin + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +type testAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} + +func TestDedupJoinResultAndFinalizeBatchesUseAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &DedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + RightTypes: []types.Type{types.T_int64.ToType()}, + } + account := installTestAllocation(t, arg) + require.NoError(t, arg.resetRBat()) + require.Same(t, arg.resultAllocation, arg.ctr.rbat.Vecs[0].AllocationAccountSelection()) + + arg.ctr.matched = &bitmap.Bitmap{} + first := batch.NewWithSize(0) + first.SetRowCount(colexec.DefaultBatchSize) + second := batch.NewWithSize(0) + second.SetRowCount(1) + arg.ctr.batches = []*batch.Batch{first, second} + require.NoError(t, arg.ctr.finalize(arg, proc)) + require.Len(t, arg.ctr.buf, 2) + for _, result := range arg.ctr.buf { + require.Same(t, arg.resultAllocation, result.Vecs[0].AllocationAccountSelection()) + } + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.rbat) + require.Empty(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestDedupJoinResultAndFinalizeBatchesHonorAllocationCapacity(t *testing.T) { + newAccount := func(t *testing.T) *mpool.AllocationAccount { + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + return account + } + + t.Run("probe result", func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + account := newAccount(t) + arg := &DedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetRBat()) + err := vector.AppendFixed(arg.ctr.rbat.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) + }) + + t.Run("multi batch finalize", func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + account := newAccount(t) + arg := &DedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + arg.ctr.matched = &bitmap.Bitmap{} + first := batch.NewWithSize(0) + first.SetRowCount(colexec.DefaultBatchSize) + second := batch.NewWithSize(0) + second.SetRowCount(1) + arg.ctr.batches = []*batch.Batch{first, second} + err := arg.ctr.finalize(arg, proc) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) + }) +} diff --git a/pkg/sql/colexec/dedupjoin/expression_memory_test.go b/pkg/sql/colexec/dedupjoin/expression_memory_test.go deleted file mode 100644 index 4265a7313bb4b..0000000000000 --- a/pkg/sql/colexec/dedupjoin/expression_memory_test.go +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed 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. - -package dedupjoin - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestDedupJoinResetReleasesProbeExpressionLease(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I32Val{I32Val: 1}, - }}, - } - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - arg := &DedupJoin{} - arg.ctr.evecs = []evalVector{{executor: executors[0]}} - arg.ctr.vecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionLease = lease - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, generation.Used()) - - arg.Reset(proc, false, nil) - require.Zero(t, generation.Used()) - require.Nil(t, arg.ctr.evecs) - require.Nil(t, arg.ctr.vecs) - require.Nil(t, arg.ctr.probeExpressionLease) -} diff --git a/pkg/sql/colexec/dedupjoin/join.go b/pkg/sql/colexec/dedupjoin/join.go index c7a60ba00ef3e..d6dbebb91babf 100644 --- a/pkg/sql/colexec/dedupjoin/join.go +++ b/pkg/sql/colexec/dedupjoin/join.go @@ -19,9 +19,9 @@ import ( "strings" "github.com/matrixorigin/matrixone/pkg/catalog" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -138,6 +138,9 @@ func (dedupJoin *DedupJoin) OpType() vm.OpType { return vm.DedupJoin } func (dedupJoin *DedupJoin) Prepare(proc *process.Process) (err error) { + if dedupJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if dedupJoin.OpAnalyzer == nil { dedupJoin.OpAnalyzer = process.NewAnalyzer(dedupJoin.GetIdx(), dedupJoin.IsFirst, dedupJoin.IsLast, "dedup join") } else { @@ -148,13 +151,21 @@ func (dedupJoin *DedupJoin) Prepare(proc *process.Process) (err error) { newUpdateExecs := len(dedupJoin.ctr.exprExecs) == 0 && len(dedupJoin.UpdateColExprList) > 0 var evalExecs, updateExecs []colexec.ExpressionExecutor if newEvalVectors { - evalExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, dedupJoin.Conditions[0]) + evalExecs, err = hashbuild.NewExpressionExecutors( + proc, + dedupJoin.Conditions[0], + dedupJoin.allocationAccount, + ) if err != nil { return err } } if newUpdateExecs { - updateExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, dedupJoin.UpdateColExprList) + updateExecs, err = hashbuild.NewExpressionExecutors( + proc, + dedupJoin.UpdateColExprList, + dedupJoin.allocationAccount, + ) if err != nil { for _, exec := range evalExecs { exec.Free() @@ -249,23 +260,30 @@ func (dedupJoin *DedupJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.cleanBuf(proc) // Clear previous bucket state before advancing. ctr.cleanBucketState(proc) + var allocationErr error ok, bktErr := ctr.spillEngine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, res spillutil.BucketResult) { if res == spillutil.BucketReady { ctr.mp = jm ctr.batches = jm.GetBatches() ctr.batchRowCount = jm.GetRowCount() - ctr.matched = &bitmap.Bitmap{} - if dedupJoin.OnDuplicateAction != plan.Node_UPDATE { - ctr.matched.InitWithSize(ctr.batchRowCount) - } else { - ctr.matched.InitWithSize(int64(jm.GetGroupCount())) + rows := ctr.batchRowCount + if dedupJoin.OnDuplicateAction == plan.Node_UPDATE { + rows = int64(jm.GetGroupCount()) } + ctr.matched, allocationErr = colexec.NewAccountedBitmap( + rows, proc.Mp(), dedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteMatched, + ) } }) if bktErr != nil { return result, hashbuild.TerminalBudgetError(proc.Ctx, bktErr) } + if allocationErr != nil { + return result, allocationErr + } if ok && ctr.mp != nil { // BucketReady: init capture buffers for REPLACE spill path. if ctr.batchRowCount > 0 && len(dedupJoin.OldColCapturePlaceholderIdxList) > 0 { @@ -312,22 +330,13 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce if takeErr != nil { return takeErr } - probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) - for i := range ctr.evecs { - probeExecutors[i] = ctr.evecs[i].executor - } - probeExpressionLease, leaseErr := hashbuild.NewExpressionMemoryLease( - budget, dedupJoin.Conditions[0], probeExecutors, false) - if leaseErr != nil { + if dedupJoin.allocationAccount == nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - return leaseErr + return mpool.ErrAllocationAccountInvalid } - ctr.probeExpressionLease = probeExpressionLease - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ BuildKeyExprs: dedupJoin.Conditions[1], ProbeKeyExprs: dedupJoin.Conditions[0], SpillThreshold: ctr.spillThreshold, @@ -343,13 +352,16 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce DedupDeleteMarkerColIdx: dedupJoin.DedupDeleteMarkerColIdx, DedupDeleteKeepColIdxList: dedupJoin.DedupDeleteKeepColIdxList, Budget: budget, - ProbeExpressionLease: probeExpressionLease, - }) - if len(payload.Files) > 0 { - engine.InitFromSpilledFiles(payload.Files) - } else { - engine.InitFromSpilledMap(payload.LegacyFds) + }, dedupJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) + if engineErr != nil { + _ = payload.Close() + ctr.mp.Free() + ctr.mp = nil + ctr.cleanEvalVectors() + return engineErr } + engine.InitFromSpilledFiles(payload.Files) + ctr.spillEngine = engine if err := engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { input, err := vm.ChildrenCall(dedupJoin.GetChildren(0), proc, analyzer) @@ -366,10 +378,10 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce ctr.mp.Free() ctr.mp = nil engine.Cleanup(proc) + ctr.spillEngine = nil return err } ctr.mp.Free() - ctr.spillEngine = engine ctr.mp = nil return } @@ -380,11 +392,19 @@ func (dedupJoin *DedupJoin) build(analyzer process.Analyzer, proc *process.Proce ctr.batches = ctr.mp.GetBatches() ctr.batchRowCount = ctr.mp.GetRowCount() if ctr.batchRowCount > 0 { - ctr.matched = &bitmap.Bitmap{} - if dedupJoin.OnDuplicateAction != plan.Node_UPDATE { - ctr.matched.InitWithSize(ctr.batchRowCount) - } else { - ctr.matched.InitWithSize(int64(ctr.mp.GetGroupCount())) + rows := ctr.batchRowCount + if dedupJoin.OnDuplicateAction == plan.Node_UPDATE { + rows = int64(ctr.mp.GetGroupCount()) + } + ctr.matched, err = colexec.NewAccountedBitmap( + rows, + proc.Mp(), + dedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteMatched, + ) + if err != nil { + return err } } if ctr.batchRowCount > 0 && len(dedupJoin.OldColCapturePlaceholderIdxList) > 0 { @@ -411,7 +431,13 @@ func (ctr *container) initCaptureBuffers(ap *DedupJoin, proc *process.Process) e ctr.capturedVecs = make([]*vector.Vector, n) for i, probePos := range ap.OldColCaptureProbeIdxList { typ := ap.LeftTypes[probePos] - vec := vector.NewOffHeapVecWithType(typ) + vec, err := vector.NewOffHeapVecWithTypeAndAllocation( + typ, + ap.stateAllocation, + ) + if err != nil { + return err + } if err := vector.AppendMultiFixed(vec, 0, true, int(ctr.batchRowCount), proc.Mp()); err != nil { vec.Free(proc.Mp()) ctr.capturedVecs[i] = nil @@ -419,8 +445,18 @@ func (ctr *container) initCaptureBuffers(ap *DedupJoin, proc *process.Process) e } ctr.capturedVecs[i] = vec } - ctr.captured = &bitmap.Bitmap{} - ctr.captured.InitWithSize(ctr.batchRowCount) + var err error + ctr.captured, err = colexec.NewAccountedBitmap( + ctr.batchRowCount, + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteCaptured, + ) + if err != nil { + ctr.cleanCaptured(proc) + return err + } ctr.captureResultIdx = make([]int32, len(ap.Result)) for j := range ctr.captureResultIdx { ctr.captureResultIdx[j] = -1 @@ -439,6 +475,7 @@ func (ctr *container) initCaptureBuffers(ap *DedupJoin, proc *process.Process) e return nil } func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { + var err error if ap.needsFinalizeMerge() { if !ap.IsMerger { if ap.Mailbox == nil { @@ -468,6 +505,7 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { // point Reset/Free still owns and releases these vectors. ctr.captured = nil ctr.capturedVecs = nil + ctr.matched = nil // Publication, not acknowledgement, is the worker's single status // for this round. Mark it before waiting so concurrent cancellation // cannot make Reset enqueue a duplicate abort status. @@ -570,12 +608,18 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { cIdx := ctr.captureResultIdx[j] if ctr.captured != nil && ctr.captured.Count() > 0 { typ := ap.RightTypes[rp.Pos] - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(typ) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(typ) + if err != nil { + return err + } if err := ap.ctr.buf[i].Vecs[j].UnionBatch(ctr.capturedVecs[cIdx], capOffset, batSize, nil, proc.Mp()); err != nil { return err } } else { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[i].Vecs[j], 0, true, batSize, proc.Mp()); err != nil { return err } @@ -592,13 +636,19 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { // its own valid vector; ownership transfer here would // leave later references reading a nil vector. typ := ap.RightTypes[rp.Pos] - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(typ) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(typ) + if err != nil { + return err + } if err := vector.GetUnionAllFunction(typ, proc.Mp())(ap.ctr.buf[i].Vecs[j], bat.Vecs[rp.Pos]); err != nil { return err } } } else { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[i].Vecs[j], 0, true, batSize, proc.Mp()); err != nil { return err } @@ -613,7 +663,18 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { return nil } ctr.matched.Negate() - sels := make([]int32, 0, count) + sels, err := mpool.MakeSliceAccounted[int32]( + count, + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteFinalizeSelections, + ) + if err != nil { + return err + } + sels = sels[:0] + defer mpool.FreeSlice(proc.Mp(), sels) itr := ctr.matched.Iterator() for itr.HasNext() { r := itr.Next() @@ -631,12 +692,18 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { ap.ctr.buf[i] = batch.NewOffHeapWithSize(len(ap.Result)) for j, rp := range ap.Result { if rp.Rel == 1 { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } if err := unionSelsByBatch(ap.ctr.buf[i].Vecs[j], ctr.batches, rp.Pos, newSels, proc); err != nil { return err } } else { - ap.ctr.buf[i].Vecs[j] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[i].Vecs[j], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[i].Vecs[j], 0, true, len(newSels), proc.Mp()); err != nil { return err } @@ -662,7 +729,10 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { ap.ctr.buf[batIdx] = batch.NewOffHeapWithSize(len(ap.Result)) for i, rp := range ap.Result { if rp.Rel == 1 { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } for _, sel := range sels[fillCnt : fillCnt+batSize] { idx1, idx2 := sel/colexec.DefaultBatchSize, sel%colexec.DefaultBatchSize if err := ap.ctr.buf[batIdx].Vecs[i].UnionOne(ctr.batches[idx1].Vecs[rp.Pos], int64(idx2), proc.Mp()); err != nil { @@ -670,7 +740,10 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { } } } else { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } if err := vector.AppendMultiFixed(ap.ctr.buf[batIdx].Vecs[i], 0, true, batSize, proc.Mp()); err != nil { return err } @@ -694,9 +767,15 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error { ap.ctr.buf[batIdx] = batch.NewOffHeapWithSize(len(ap.Result)) for i, rp := range ap.Result { if rp.Rel == 1 { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.RightTypes[rp.Pos]) + if err != nil { + return err + } } else { - ap.ctr.buf[batIdx].Vecs[i] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) + ap.ctr.buf[batIdx].Vecs[i], err = ap.newResultVector(ap.LeftTypes[rp.Pos]) + if err != nil { + return err + } } } } @@ -791,8 +870,10 @@ func (ctr *container) withRestoredJoinBat1Vectors(updateCols []int32, fn func() } func (ctr *container) probe(bat *batch.Batch, ap *DedupJoin, proc *process.Process, analyzer process.Analyzer, result *vm.CallResult) error { - ap.resetRBat() - err := ctr.evalJoinConditionBudgeted(bat, proc) + if err := ap.resetRBat(); err != nil { + return err + } + err := ctr.evalJoinCondition(bat, proc) if err != nil { return err } @@ -816,7 +897,10 @@ func (ctr *container) probe(bat *batch.Batch, ap *DedupJoin, proc *process.Proce if n > hashmap.UnitLimit { n = hashmap.UnitLimit } - vals, zvals := itr.Find(i, n, ctr.vecs) + vals, zvals, err := itr.Find(i, n, ctr.vecs) + if err != nil { + return err + } for k := 0; k < n; k++ { if zvals[k] == 0 || vals[k] == 0 { continue @@ -973,16 +1057,6 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) return nil } -func (ctr *container) evalJoinConditionBudgeted(bat *batch.Batch, proc *process.Process) error { - if ctr.probeExpressionLease == nil { - return ctr.evalJoinCondition(bat, proc) - } - return ctr.probeExpressionLease.Eval(proc, []*batch.Batch{bat}, bat.RowCount(), func(i int, vec *vector.Vector) error { - ctr.vecs[i] = vec - ctr.evecs[i].vec = vec - return nil - }) -} func unionSelsByBatch(dst *vector.Vector, batches []*batch.Batch, colPos int32, sels []int32, proc *process.Process) error { if len(sels) <= 16 { for _, sel := range sels { @@ -1016,7 +1090,11 @@ func unionSelsByBatch(dst *vector.Vector, batches []*batch.Batch, colPos int32, } return nil } -func (dedupJoin *DedupJoin) resetRBat() { +func (dedupJoin *DedupJoin) newResultVector(typ types.Type) (*vector.Vector, error) { + return vector.NewOffHeapVecWithTypeAndAllocation(typ, dedupJoin.resultAllocation) +} + +func (dedupJoin *DedupJoin) resetRBat() error { ctr := &dedupJoin.ctr if ctr.rbat != nil { ctr.rbat.CleanOnlyData() @@ -1029,5 +1107,11 @@ func (dedupJoin *DedupJoin) resetRBat() { ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(dedupJoin.RightTypes[rp.Pos]) } } + if err := ctr.rbat.SetAllocationAccount(dedupJoin.resultAllocation); err != nil { + ctr.rbat.Clean(nil) + ctr.rbat = nil + return err + } } + return nil } diff --git a/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go b/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go index d14a20d3f0afc..8cdd73e366e8b 100644 --- a/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go +++ b/pkg/sql/colexec/dedupjoin/join_finalize_optimize_test.go @@ -40,6 +40,7 @@ func runFinalizeFixture( buildBat, probeBat *batch.Batch, ) []*batch.Batch { t.Helper() + installTestAllocation(t, dedupArg, buildArg) buildArg.Children = nil buildArg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBat})) diff --git a/pkg/sql/colexec/dedupjoin/join_test.go b/pkg/sql/colexec/dedupjoin/join_test.go index 24a230c1cdb17..9d9ea4992968b 100644 --- a/pkg/sql/colexec/dedupjoin/join_test.go +++ b/pkg/sql/colexec/dedupjoin/join_test.go @@ -18,6 +18,7 @@ import ( "bytes" "context" "errors" + "io" "os" "sync" "testing" @@ -58,6 +59,30 @@ type joinTestCase struct { barg *hashbuild.HashBuild } +func newDedupTestSpillEngine( + t *testing.T, + cfg spillutil.SpillEngineConfig, +) *spillutil.SpillEngine { + t.Helper() + if cfg.Budget == nil { + budget := process.MustNewHashBuildBudget(1<<60, 1<<60) + var err error + cfg.Budget, err = budget.OpenGeneration(1) + require.NoError(t, err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, cfg.Budget) + require.NoError(t, err) + engine, err := spillutil.NewSpillEngine( + cfg, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + return engine +} + func TestDedupFinalizeCleansConsumedBuffer(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) baseline := proc.Mp().CurrNB() @@ -71,7 +96,10 @@ func TestDedupFinalizeCleansConsumedBuffer(t *testing.T) { arg.ctr.state = Finalize arg.ctr.buf = []*batch.Batch{bat} arg.ctr.lastPos = 1 - arg.ctr.spillEngine = spillutil.NewSpillEngine(spillutil.SpillEngineConfig{}) + arg.ctr.spillEngine = newDedupTestSpillEngine( + t, + spillutil.SpillEngineConfig{}, + ) res, err := arg.Call(proc) require.NoError(t, err) @@ -112,14 +140,32 @@ func writeDedupSpillBatch(t *testing.T, proc *process.Process, name string, valu require.NoError(t, err) fd, err := spillfs.CreateAndRemoveFile(proc.Ctx, name) require.NoError(t, err) - w := spillutil.BucketWriter{Name: name, Fd: fd} bat := batch.NewWithSize(1) bat.Vecs[0] = testutil.MakeInt32Vector([]int32{value}, nil, proc.Mp()) bat.SetRowCount(1) - var buf bytes.Buffer - require.NoError(t, spillutil.FlushBucketBatch(proc, bat, &w, &buf, nil)) + var payload bytes.Buffer + require.NoError(t, bat.MarshalBinaryWithGroupingTo(&payload)) + rows, size, magic := int64(1), int64(payload.Len()), uint64(spillutil.SpillMagic) + for _, part := range [][]byte{ + types.EncodeInt64(&rows), + types.EncodeInt64(&size), + payload.Bytes(), + types.EncodeUint64(&magic), + } { + _, err = fd.Write(part) + require.NoError(t, err) + } + _, err = fd.Seek(0, io.SeekStart) + require.NoError(t, err) bat.Clean(proc.Mp()) - return w.HandOffFd() + return fd +} + +func newDedupSpillFile(t *testing.T, fd *os.File, rows int64) *message.SpillFile { + t.Helper() + info, err := fd.Stat() + require.NoError(t, err) + return message.NewSpillFile(fd, rows, uint64(info.Size()), nil) } func TestDedupSpillAdvancesAfterOutput(t *testing.T) { @@ -127,15 +173,23 @@ func TestDedupSpillAdvancesAfterOutput(t *testing.T) { baseline := proc.Mp().CurrNB() typ := types.T_int32.ToType() conditions := [][]*plan.Expr{{newExpr(0, typ)}, {newExpr(0, typ)}} - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine := newDedupTestSpillEngine(t, spillutil.SpillEngineConfig{ BuildKeyExprs: conditions[1], NeedBatches: true, NeedsBuildForEmptyProbe: true, IsDedup: true, }) - engine.InitFromSpilledMap([]*os.File{ - writeDedupSpillBatch(t, proc, "dedup_bucket_1", 1), - writeDedupSpillBatch(t, proc, "dedup_bucket_2", 2), + engine.InitFromSpilledFiles([]*message.SpillFile{ + newDedupSpillFile( + t, + writeDedupSpillBatch(t, proc, "dedup_bucket_1", 1), + 1, + ), + newDedupSpillFile( + t, + writeDedupSpillBatch(t, proc, "dedup_bucket_2", 2), + 1, + ), }) arg := &DedupJoin{ @@ -144,6 +198,7 @@ func TestDedupSpillAdvancesAfterOutput(t *testing.T) { Result: []colexec.ResultPos{{Rel: 1, Pos: 0}}, OnDuplicateAction: plan.Node_FAIL, } + installTestAllocation(t, arg) require.NoError(t, arg.Prepare(proc)) arg.ctr.state = Finalize arg.ctr.spillEngine = engine @@ -314,6 +369,7 @@ func TestDedupPrepareFailureCanRetry(t *testing.T) { Conditions: [][]*plan.Expr{{valid}, {valid}}, UpdateColExprList: []*plan.Expr{valid, invalid}, } + installTestAllocation(t, arg) require.Error(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.vecs) @@ -523,7 +579,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] // }, //}) tag++ - return joinTestCase{ + tc := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -555,6 +611,8 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] JoinMapRefCnt: 1, }, } + installTestAllocation(t, tc.arg, tc.barg) + return tc } func resetChildren(arg *DedupJoin, m *mpool.MPool) { @@ -653,6 +711,7 @@ func TestDedupJoinCapture(t *testing.T) { JoinMapTag: curTag, JoinMapRefCnt: 1, } + installTestAllocation(t, dedupArg, buildArg) // Set up children buildOp := colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBat}) @@ -755,6 +814,7 @@ func TestDedupJoinCapturePartialMatch(t *testing.T) { JoinMapTag: curTag, JoinMapRefCnt: 1, } + installTestAllocation(t, dedupArg, buildArg) buildOp := colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBat}) buildArg.Children = nil @@ -849,6 +909,7 @@ func TestDedupJoinCaptureReset(t *testing.T) { JoinMapTag: curTag, JoinMapRefCnt: 1, } + installTestAllocation(t, dedupArg, buildArg) // --- First run --- buildBat1 := makeInt32Batch(proc.Mp(), [][]int32{{10, 20}, {0, 0}}, [][]uint64{nil, {0, 1}}) @@ -1632,6 +1693,7 @@ func TestDedupFinalizeParallelMergePreservesDataAcrossReset(t *testing.T) { IsMerger: false, Mailbox: mailbox, } + installTestAllocation(t, arg, workerArg) cleaned := false t.Cleanup(func() { if !cleaned { diff --git a/pkg/sql/colexec/dedupjoin/key_contract_test.go b/pkg/sql/colexec/dedupjoin/key_contract_test.go index 018ee5f4e6f47..199f54edfbfff 100644 --- a/pkg/sql/colexec/dedupjoin/key_contract_test.go +++ b/pkg/sql/colexec/dedupjoin/key_contract_test.go @@ -189,6 +189,7 @@ func runDedupJoinDoubleSignedZeroContract( if mode.shuffle { buildArg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: joinMapTag + 8000} } + installTestAllocation(t, dedupArg, buildArg) buildArg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBatch})) spillBefore := promtestutil.ToFloat64( diff --git a/pkg/sql/colexec/dedupjoin/terminal_budget_test.go b/pkg/sql/colexec/dedupjoin/terminal_budget_test.go index 5a92ea6e3fe93..ec706bf726166 100644 --- a/pkg/sql/colexec/dedupjoin/terminal_budget_test.go +++ b/pkg/sql/colexec/dedupjoin/terminal_budget_test.go @@ -42,7 +42,7 @@ func TestDedupJoinCallConvertsTerminalBudgetAdmission(t *testing.T) { admission := &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, + Component: process.HashBuildBudgetComponentMemory, Requested: 2, Used: 1, Cap: 1, diff --git a/pkg/sql/colexec/dedupjoin/types.go b/pkg/sql/colexec/dedupjoin/types.go index e307bfbef3f02..7831404cabe6b 100644 --- a/pkg/sql/colexec/dedupjoin/types.go +++ b/pkg/sql/colexec/dedupjoin/types.go @@ -21,6 +21,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -43,6 +44,23 @@ const ( End ) +const ( + dedupJoinAllocationSiteMatched mpool.AllocationSite = iota + 82 + dedupJoinAllocationSiteCaptured + dedupJoinAllocationSiteCaptureData + dedupJoinAllocationSiteCaptureArea + dedupJoinAllocationSiteCaptureNulls + dedupJoinAllocationSiteCaptureGrouping + dedupJoinAllocationSiteFinalizeSelections +) + +const ( + dedupJoinAllocationSiteResultData mpool.AllocationSite = iota + 110 + dedupJoinAllocationSiteResultArea + dedupJoinAllocationSiteResultNulls + dedupJoinAllocationSiteResultGrouping +) + // WorkerJoinMsg carries per-worker state from non-merger workers to the // merger worker at finalize time. Regular DEDUP JOIN only populates matched; // the REPLACE INTO merged main-table scan path (OldColCapture) additionally @@ -199,6 +217,8 @@ func freeCapturedVecs(vecs []*vector.Vector, proc *process.Process) { func freeWorkerJoinMsg(msg *WorkerJoinMsg, proc *process.Process) { if msg != nil { + colexec.FreeAccountedBitmap(msg.matched, proc.Mp()) + colexec.FreeAccountedBitmap(msg.captured, proc.Mp()) freeCapturedVecs(msg.capturedVecs, proc) } } @@ -257,10 +277,6 @@ type container struct { // Spill support for large build sides. spillEngine *spillutil.SpillEngine spillThreshold int64 - // Non-nil only for spilled joins, where probe expressions are part of the - // shared HashBuild/spill working set. Resident probe expressions remain - // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease } type DedupJoin struct { @@ -298,10 +314,76 @@ type DedupJoin struct { // main-table scan path; empty for regular INSERT/UPDATE. OldColCapturePlaceholderIdxList []int32 OldColCaptureProbeIdxList []int32 + allocationAccount *mpool.AllocationAccount + stateAllocation *vector.AllocationAccountSelection + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } +func (dedupJoin *DedupJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if dedupJoin.allocationAccount != nil && + dedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if dedupJoin.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteCaptureData, + dedupJoinAllocationSiteCaptureArea, + dedupJoinAllocationSiteCaptureNulls, + dedupJoinAllocationSiteCaptureGrouping, + ) + if err != nil { + return err + } + resultSelection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + dedupJoinAllocationSiteResultData, + dedupJoinAllocationSiteResultArea, + dedupJoinAllocationSiteResultNulls, + dedupJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } + dedupJoin.allocationAccount = account + dedupJoin.stateAllocation = selection + dedupJoin.resultAllocation = resultSelection + return nil +} + +func (dedupJoin *DedupJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if dedupJoin.allocationAccount == nil { + return nil + } + if dedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if dedupJoin.ctr.mp != nil || dedupJoin.ctr.spillEngine != nil || + len(dedupJoin.ctr.evecs) != 0 || len(dedupJoin.ctr.exprExecs) != 0 || + dedupJoin.ctr.matched != nil || dedupJoin.ctr.captured != nil || + len(dedupJoin.ctr.capturedVecs) != 0 || dedupJoin.ctr.rbat != nil || + len(dedupJoin.ctr.buf) != 0 { + return mpool.ErrAllocationAccountInvariant + } + dedupJoin.allocationAccount = nil + dedupJoin.stateAllocation = nil + dedupJoin.resultAllocation = nil + return nil +} + func (dedupJoin *DedupJoin) GetOperatorBase() *vm.OperatorBase { return &dedupJoin.OperatorBase } @@ -370,19 +452,14 @@ func (dedupJoin *DedupJoin) Reset(proc *process.Process, pipelineFailed bool, er } ctr.maxAllocSize = 0 - ctr.cleanBuf(proc) + ctr.cleanResultBatches(proc) ctr.cleanBucketState(proc) - ctr.resetExprExecutor() + ctr.cleanExprExecutor() if ctr.spillEngine != nil { ctr.spillEngine.Cleanup(proc) ctr.spillEngine = nil } - if ctr.probeExpressionLease != nil { - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - } else { - ctr.resetEvalVectors() - } + ctr.cleanEvalVectors() ctr.roundStatusPublished = false ctr.state = Build ctr.lastPos = 0 @@ -395,7 +472,7 @@ func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err // reopened prepared-pipeline generation stopped from Free. dedupJoin.Mailbox.drain(proc) } - ctr.cleanBuf(proc) + ctr.cleanResultBatches(proc) ctr.cleanBucketState(proc) ctr.cleanBatch(proc) ctr.cleanExprExecutor() @@ -404,24 +481,19 @@ func (dedupJoin *DedupJoin) Free(proc *process.Process, pipelineFailed bool, err ctr.spillEngine = nil } ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() } func (dedupJoin *DedupJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { return input, nil } -func (ctr *container) resetExprExecutor() { - for i := range ctr.exprExecs { - ctr.exprExecs[i].ResetForNextQuery() - } -} - func (ctr *container) cleanExprExecutor() { for i := range ctr.exprExecs { - ctr.exprExecs[i].Free() - ctr.exprExecs[i] = nil + if ctr.exprExecs[i] != nil { + ctr.exprExecs[i].Free() + } } + ctr.exprExecs = nil } func (ctr *container) cleanBuf(proc *process.Process) { @@ -433,6 +505,14 @@ func (ctr *container) cleanBuf(proc *process.Process) { ctr.buf = nil } +func (ctr *container) cleanResultBatches(proc *process.Process) { + ctr.cleanBuf(proc) + if ctr.rbat != nil { + ctr.rbat.Clean(proc.GetMPool()) + ctr.rbat = nil + } +} + func (ctr *container) cleanCaptured(proc *process.Process) { for _, v := range ctr.capturedVecs { if v != nil { @@ -440,6 +520,7 @@ func (ctr *container) cleanCaptured(proc *process.Process) { } } ctr.capturedVecs = nil + colexec.FreeAccountedBitmap(ctr.captured, proc.Mp()) ctr.captured = nil ctr.captureResultIdx = nil } @@ -471,10 +552,12 @@ func (ctr *container) cleanBucketState(proc *process.Process) { ctr.cleanHashMap() ctr.batches = nil ctr.batchRowCount = 0 + colexec.FreeAccountedBitmap(ctr.matched, proc.Mp()) ctr.matched = nil } func (ctr *container) cleanHashMap() { + hashmap.IteratorClearOwner(ctr.cachedItr) ctr.cachedItr = nil if ctr.mp != nil { ctr.mp.Free() @@ -492,18 +575,3 @@ func (ctr *container) cleanEvalVectors() { ctr.evecs = nil ctr.vecs = nil } - -func (ctr *container) resetEvalVectors() { - for i := range ctr.evecs { - if ctr.evecs[i].executor != nil { - ctr.evecs[i].executor.ResetForNextQuery() - } - } -} - -func (ctr *container) releaseProbeExpressionLease() { - if ctr.probeExpressionLease != nil { - ctr.probeExpressionLease.Release() - ctr.probeExpressionLease = nil - } -} diff --git a/pkg/sql/colexec/dispatch/dispatch_test.go b/pkg/sql/colexec/dispatch/dispatch_test.go index 1b39a02f1ec73..ab3b3ea2068bc 100644 --- a/pkg/sql/colexec/dispatch/dispatch_test.go +++ b/pkg/sql/colexec/dispatch/dispatch_test.go @@ -897,6 +897,131 @@ func TestDispatchResetEndPreservesQueuedBroadcastBatchUntilDeferredCleanup(t *te require.Equal(t, int64(0), mp.CurrNB()) } +func TestDispatchAllocationClearFinalizesTerminalSpoolPending(t *testing.T) { + testDispatchAllocationClearFinalizesSpool(t, false) +} + +func TestDispatchAccountedDeferredCleanupReleasesReusableCache(t *testing.T) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + srcMP := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(srcMP) + }) + src := newDispatchSpoolTestBatch(t, srcMP, 1024) + t.Cleanup(func() { + src.Clean(srcMP) + }) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + sp.ReleaseCurrent(0) + require.Positive(t, mp.CurrNB()) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + d := &Dispatch{cleanupSpool: sp} + require.NoError(t, d.SetAllocationAccount(account)) + d.CleanupDeferredSpool() + require.Same(t, sp, d.cleanupSpool) + require.Zero(t, mp.CurrNB()) + require.NoError(t, d.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestDispatchAllocationAccountContract(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 1) + require.NoError(t, err) + first, err := registry.Open(1) + require.NoError(t, err) + second, err := registry.Open(1) + require.NoError(t, err) + d := &Dispatch{} + require.False(t, d.ActivatesAllocationAccountLifecycle()) + require.ErrorIs(t, d.SetAllocationAccount(nil), mpool.ErrAllocationAccountInvalid) + require.NoError(t, d.SetAllocationAccount(first)) + require.ErrorIs(t, d.SetAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + require.ErrorIs(t, d.ClearAllocationAccount(second), mpool.ErrAllocationAccountMismatch) + d.ctr = &container{sp: &pSpool.PipelineSpool{}} + require.ErrorIs(t, d.ClearAllocationAccount(first), mpool.ErrAllocationAccountInvariant) + d.ctr = nil + require.NoError(t, d.ClearAllocationAccount(first)) + require.NoError(t, d.ClearAllocationAccount(first)) + _, _, err = registry.CompleteTerminal(first) + require.NoError(t, err) + _, _, err = registry.CompleteTerminal(second) + require.NoError(t, err) +} + +func TestDispatchAllocationClearFinalizesAbortedSpool(t *testing.T) { + testDispatchAllocationClearFinalizesSpool(t, true) +} + +func testDispatchAllocationClearFinalizesSpool(t *testing.T, abort bool) { + mp := mpool.MustNewZeroNoFixed() + t.Cleanup(func() { + mpool.DeleteMPool(mp) + }) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 102, + 103, + 104, + 105, + ) + require.NoError(t, err) + src := batch.NewOffHeapWithSize(1) + require.NoError(t, src.SetAllocationAccount(selection)) + src.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, vector.AppendFixed(src.Vecs[0], int64(1), false, mp)) + src.SetRowCount(1) + + sp := pSpool.InitMyPipelineSpool(mp, 1) + done, err := sp.SendBatch(context.Background(), 0, src, nil) + require.NoError(t, err) + require.False(t, done) + + d := &Dispatch{} + require.NoError(t, d.SetAllocationAccount(account)) + if abort { + got, info := sp.ReceiveBatch(0) + require.NoError(t, info) + require.NotNil(t, got) + d.ctr = &container{sp: sp} + d.Reset(nil, true, moerr.NewInternalErrorNoCtx("pipeline failed")) + require.Same(t, sp, d.cleanupSpool) + sp.ReleaseCurrent(0) + } else { + d.cleanupSpool = sp + sp.ForceCleanupAfterTerminalSignal() + } + d.CleanupDeferredSpool() + require.Same(t, sp, d.cleanupSpool) + require.NoError(t, d.ClearAllocationAccount(account)) + require.Nil(t, d.cleanupSpool) + + src.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + _, err = registry.Finalize(account) + require.NoError(t, err) +} + // TestReceiverDone_OldBehavior tests the old behavior (kept for backward compatibility verification) func TestReceiverDone_OldBehavior(t *testing.T) { proc := testutil.NewProcess(t) diff --git a/pkg/sql/colexec/dispatch/types.go b/pkg/sql/colexec/dispatch/types.go index 3b73c149ea50c..f9e8b76b3312d 100644 --- a/pkg/sql/colexec/dispatch/types.go +++ b/pkg/sql/colexec/dispatch/types.go @@ -76,8 +76,9 @@ type container struct { } type Dispatch struct { - ctr *container - cleanupSpool *pSpool.PipelineSpool + ctr *container + cleanupSpool *pSpool.PipelineSpool + allocationAccount *mpool.AllocationAccount // MaterializedSource is used by a multi-reference CTE whose consumers can // have execution dependencies on one another. It is local-only and bypasses @@ -109,6 +110,45 @@ func (dispatch *Dispatch) GetOperatorBase() *vm.OperatorBase { return &dispatch.OperatorBase } +func (dispatch *Dispatch) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if dispatch.allocationAccount != nil && dispatch.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + dispatch.allocationAccount = account + return nil +} + +// ActivatesAllocationAccountLifecycle reports that Dispatch only participates +// in an account already required by an allocation-producing operator. +func (dispatch *Dispatch) ActivatesAllocationAccountLifecycle() bool { + return false +} + +func (dispatch *Dispatch) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if dispatch.allocationAccount == nil { + return nil + } + if dispatch.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if dispatch.ctr != nil && dispatch.ctr.sp != nil { + return mpool.ErrAllocationAccountInvariant + } + if dispatch.cleanupSpool != nil { + dispatch.cleanupSpool.FinalizeAfterConsumersQuiesced() + dispatch.cleanupSpool = nil + } + dispatch.allocationAccount = nil + return nil +} + func init() { reuse.CreatePool[Dispatch]( func() *Dispatch { @@ -286,7 +326,11 @@ func (dispatch *Dispatch) Reset(proc *process.Process, pipelineFailed bool, err abortErr = fallbackErr } sp.Abort(abortErr) - dispatch.cleanupSpool = nil + if dispatch.allocationAccount != nil { + dispatch.cleanupSpool = sp + } else { + dispatch.cleanupSpool = nil + } } dispatch.ctr.sp = nil } else { @@ -308,6 +352,10 @@ func (dispatch *Dispatch) CleanupDeferredSpool() { if dispatch.cleanupSpool == nil { return } + if dispatch.allocationAccount != nil { + dispatch.cleanupSpool.ReleaseReusableCacheAfterProducerQuiesced() + return + } dispatch.cleanupSpool.ForceCleanupAfterTerminalSignal() dispatch.cleanupSpool = nil } diff --git a/pkg/sql/colexec/evalExpression.go b/pkg/sql/colexec/evalExpression.go index a2c98ba315d91..217d6f7587006 100644 --- a/pkg/sql/colexec/evalExpression.go +++ b/pkg/sql/colexec/evalExpression.go @@ -92,9 +92,20 @@ type ExpressionExecutor interface { } func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs []*plan.Expr) (executors []ExpressionExecutor, err error) { + return NewExpressionExecutorsFromPlanExpressionsWithAllocation(proc, planExprs, nil) +} + +// NewExpressionExecutorsFromPlanExpressionsWithAllocation builds a complete +// expression tree whose owned MPool vectors use one immutable allocation +// selection. Borrowed input vectors remain owned and charged by their source. +func NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc *process.Process, + planExprs []*plan.Expr, + selection *vector.AllocationAccountSelection, +) (executors []ExpressionExecutor, err error) { executors = make([]ExpressionExecutor, len(planExprs)) for i := range executors { - executors[i], err = NewExpressionExecutor(proc, planExprs[i]) + executors[i], err = NewExpressionExecutorWithAllocation(proc, planExprs[i], selection) if err != nil { for j := 0; j < i; j++ { executors[j].Free() @@ -106,10 +117,18 @@ func NewExpressionExecutorsFromPlanExpressions(proc *process.Process, planExprs } func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (ExpressionExecutor, error) { + return NewExpressionExecutorWithAllocation(proc, planExpr, nil) +} + +func NewExpressionExecutorWithAllocation( + proc *process.Process, + planExpr *plan.Expr, + selection *vector.AllocationAccountSelection, +) (ExpressionExecutor, error) { switch t := planExpr.Expr.(type) { case *plan.Expr_Lit: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec, err := generateConstExpressionExecutor(proc, typ, t.Lit) + vec, err := generateConstExpressionExecutor(proc, typ, t.Lit, selection) if err != nil { return nil, err } @@ -117,17 +136,21 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi case *plan.Expr_T: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - vec := vector.NewConstNull(typ, 1, proc.Mp()) + vec, err := newExpressionConstNull(typ, 1, selection) + if err != nil { + return nil, err + } return NewFixedVectorExpressionExecutor(proc.Mp(), false, vec), nil case *plan.Expr_Col: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ce := NewColumnExpressionExecutor() *ce = ColumnExpressionExecutor{ - mp: proc.Mp(), - relIndex: int(t.Col.RelPos), - colIndex: int(t.Col.ColPos), - typ: typ, + mp: proc.Mp(), + relIndex: int(t.Col.RelPos), + colIndex: int(t.Col.ColPos), + typ: typ, + allocation: selection, } // [issue#19574] // if < 0, it's special for agg or others. @@ -138,17 +161,20 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi case *plan.Expr_P: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - return NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ), nil + executor := NewParamExpressionExecutor(proc.Mp(), int(t.P.Pos), typ) + executor.allocation = selection + return executor, nil case *plan.Expr_V: typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) ve := NewVarExpressionExecutor() *ve = VarExpressionExecutor{ - mp: proc.Mp(), - name: t.V.Name, - system: t.V.System, - global: t.V.Global, - typ: typ, + mp: proc.Mp(), + name: t.V.Name, + system: t.V.System, + global: t.V.Global, + typ: typ, + allocation: selection, } return ve, nil @@ -156,6 +182,7 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi vec := vector.NewVec(types.T_any.ToType()) err := vec.UnmarshalBinary(t.Vec.Data) if err != nil { + vec.Free(proc.Mp()) return nil, err } return NewFixedVectorExpressionExecutor(proc.Mp(), true, vec), nil @@ -164,9 +191,12 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi executor := NewListExpressionExecutor() resultVecTyp := t.List.List[0].GetTyp() typ := types.New(types.T(resultVecTyp.Id), resultVecTyp.Width, resultVecTyp.Scale) - executor.Init(proc, typ, len(t.List.List)) + if err := executor.init(proc, typ, len(t.List.List), selection); err != nil { + executor.Free() + return nil, err + } for i := range executor.parameterExecutor { - subExecutor, paramErr := NewExpressionExecutor(proc, t.List.List[i]) + subExecutor, paramErr := NewExpressionExecutorWithAllocation(proc, t.List.List[i], selection) if paramErr != nil { executor.Free() return nil, paramErr @@ -196,13 +226,13 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi } typ := types.New(types.T(planExpr.Typ.Id), planExpr.Typ.Width, planExpr.Typ.Scale) - if err = executor.Init(proc, len(t.F.Args), typ); err != nil { + if err = executor.init(proc, len(t.F.Args), typ, selection); err != nil { executor.Free() return nil, err } for i := range executor.parameterExecutor { - subExecutor, paramErr := NewExpressionExecutor(proc, t.F.Args[i]) + subExecutor, paramErr := NewExpressionExecutorWithAllocation(proc, t.F.Args[i], selection) if paramErr != nil { executor.Free() return nil, paramErr @@ -215,6 +245,66 @@ func NewExpressionExecutor(proc *process.Process, planExpr *plan.Expr) (Expressi return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("unsupported expression executor for %v now", planExpr)) } +func newExpressionOffHeapVector( + typ types.Type, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewOffHeapVecWithType(typ), nil + } + return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) +} + +func newExpressionConstNull( + typ types.Type, + length int, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstNull(typ, length, nil), nil + } + return vector.NewConstNullWithAllocation(typ, length, selection) +} + +func newExpressionConstFixed[T any]( + typ types.Type, + value T, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstFixed(typ, value, length, mp) + } + return vector.NewConstFixedWithAllocation(typ, value, length, mp, selection) +} + +func newExpressionConstBytes( + typ types.Type, + value []byte, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstBytes(typ, value, length, mp) + } + return vector.NewConstBytesWithAllocation(typ, value, length, mp, selection) +} + +func newExpressionConstArray[T types.ArrayElement]( + typ types.Type, + value []T, + length int, + mp *mpool.MPool, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return vector.NewConstArray(typ, value, length, mp) + } + return vector.NewConstArrayWithAllocation(typ, value, length, mp, selection) +} + // FixedVectorExpressionExecutor // the content of its vector is fixed. // e.g. @@ -230,7 +320,8 @@ type FixedVectorExpressionExecutor struct { } type FunctionExpressionExecutor struct { - m *mpool.MPool + m *mpool.MPool + allocation *vector.AllocationAccountSelection // resultType is the declared function return type. Some built-ins refine // result metadata (for example temporal scale or decimal width/scale) at // runtime, so reusable result vectors must start each evaluation from this @@ -262,9 +353,10 @@ type FunctionExpressionExecutor struct { } type ColumnExpressionExecutor struct { - mp *mpool.MPool - relIndex int - colIndex int + mp *mpool.MPool + allocation *vector.AllocationAccountSelection + relIndex int + colIndex int // result type. typ types.Type @@ -283,8 +375,9 @@ func (expr *ColumnExpressionExecutor) GetColIndex() int { } type ParamExpressionExecutor struct { - mp *mpool.MPool - null *vector.Vector + mp *mpool.MPool + allocation *vector.AllocationAccountSelection + null *vector.Vector // maskedNull is separate from null/vec because it is not a resolved // parameter value and must never participate in the folded-value cache. maskedNull *vector.Vector @@ -298,7 +391,11 @@ type ParamExpressionExecutor struct { func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var err error + expr.maskedNull, err = newExpressionConstNull(expr.typ, 1, expr.allocation) + if err != nil { + return nil, err + } } return expr.maskedNull, nil } @@ -319,13 +416,18 @@ func (expr *ParamExpressionExecutor) Eval(proc *process.Process, batches []*batc if val == nil { if expr.null == nil { - expr.null = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + expr.null, err = newExpressionConstNull(expr.typ, 1, expr.allocation) + if err != nil { + return nil, err + } } return expr.null, nil } if expr.vec == nil { - expr.vec, err = vector.NewConstBytes(expr.typ, val, 1, proc.Mp()) + expr.vec, err = newExpressionConstBytes( + expr.typ, val, 1, proc.Mp(), expr.allocation, + ) } else { err = vector.SetConstBytes(expr.vec, val, 1, proc.GetMPool()) } @@ -373,8 +475,9 @@ func (expr *ParamExpressionExecutor) IsColumnExpr() bool { } type VarExpressionExecutor struct { - mp *mpool.MPool - null *vector.Vector + mp *mpool.MPool + allocation *vector.AllocationAccountSelection + null *vector.Vector // maskedNull lets a skipped variable avoid the resolver without changing // the value cache used by a later selected evaluation. maskedNull *vector.Vector @@ -389,7 +492,11 @@ type VarExpressionExecutor struct { func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if noRowsSelected(selectList, expressionRowCount(batches)) { if expr.maskedNull == nil { - expr.maskedNull = vector.NewConstNull(expr.typ, 1, proc.GetMPool()) + var err error + expr.maskedNull, err = newExpressionConstNull(expr.typ, 1, expr.allocation) + if err != nil { + return nil, err + } } return expr.maskedNull, nil } @@ -411,7 +518,9 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. if val == nil { if expr.null == nil { - expr.null, err = util.GenVectorByVarValue(proc, expr.typ, nil) + expr.null, err = util.GenVectorByVarValueWithAllocation( + proc, expr.typ, nil, expr.allocation, + ) } if err == nil { expr.null.SetIsBin(isBin) @@ -420,7 +529,9 @@ func (expr *VarExpressionExecutor) Eval(proc *process.Process, batches []*batch. } if expr.vec == nil { - expr.vec, err = util.GenVectorByVarValue(proc, expr.typ, val) + expr.vec, err = util.GenVectorByVarValueWithAllocation( + proc, expr.typ, val, expr.allocation, + ) } else { switch v := val.(type) { case []byte: @@ -474,7 +585,8 @@ func (expr *VarExpressionExecutor) IsColumnExpr() bool { } type ListExpressionExecutor struct { - mp *mpool.MPool + mp *mpool.MPool + allocation *vector.AllocationAccountSelection typ types.Type resultVector *vector.Vector @@ -484,11 +596,17 @@ type ListExpressionExecutor struct { func (expr *ListExpressionExecutor) Eval(proc *process.Process, batches []*batch.Batch, selectList []bool) (*vector.Vector, error) { if expr.resultVector == nil { - expr.resultVector = vector.NewOffHeapVecWithType(expr.typ) + var err error + expr.resultVector, err = newExpressionOffHeapVector(expr.typ, expr.allocation) + if err != nil { + return nil, err + } } else { expr.resultVector.CleanOnlyData() } - expr.resultVector.PreExtend(len(expr.parameterExecutor), proc.Mp()) + if err := expr.resultVector.PreExtend(len(expr.parameterExecutor), proc.Mp()); err != nil { + return nil, err + } for i := range expr.parameterExecutor { vec, err := expr.parameterExecutor[i].Eval(proc, batches, selectList) if err != nil { @@ -517,7 +635,9 @@ func (expr *ListExpressionExecutor) Free() { return } for _, e := range expr.parameterExecutor { - e.Free() + if e != nil { + e.Free() + } } if expr.resultVector != nil { expr.resultVector.Free(expr.mp) @@ -531,12 +651,26 @@ func (expr *ListExpressionExecutor) IsColumnExpr() bool { } func (expr *ListExpressionExecutor) Init(proc *process.Process, typ types.Type, parameterNum int) { + if err := expr.init(proc, typ, parameterNum, nil); err != nil { + panic(err) + } +} + +func (expr *ListExpressionExecutor) init( + proc *process.Process, + typ types.Type, + parameterNum int, + selection *vector.AllocationAccountSelection, +) error { m := proc.Mp() expr.typ = typ expr.mp = m + expr.allocation = selection expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - expr.resultVector = vector.NewOffHeapVecWithType(typ) + var err error + expr.resultVector, err = newExpressionOffHeapVector(typ, selection) + return err } func (expr *ListExpressionExecutor) SetParameter(index int, executor ExpressionExecutor) { @@ -553,14 +687,24 @@ func (expr *FunctionExpressionExecutor) Init( proc *process.Process, parameterNum int, retType types.Type) (err error) { + return expr.init(proc, parameterNum, retType, nil) +} + +func (expr *FunctionExpressionExecutor) init( + proc *process.Process, + parameterNum int, + retType types.Type, + selection *vector.AllocationAccountSelection, +) (err error) { m := proc.Mp() expr.m = m + expr.allocation = selection expr.resultType = retType expr.parameterResults = make([]*vector.Vector, parameterNum) expr.parameterExecutor = make([]ExpressionExecutor, parameterNum) - expr.resultVector = vector.NewFunctionResultWrapper(retType, m) + expr.resultVector, err = vector.NewFunctionResultWrapperWithAllocation(retType, m, selection) return err } @@ -624,14 +768,17 @@ func (expr *FunctionExpressionExecutor) EvalIff(proc *process.Process, batches [ return err } } else { - expr.parameterResults[1] = expr.iffNullResult(0, rowCount) + expr.parameterResults[1], err = expr.iffNullResult(0, rowCount) + if err != nil { + return err + } } if hasSelectedRows(falseBranch) { expr.parameterResults[2], err = expr.parameterExecutor[2].Eval(proc, batches, falseBranch) return err } - expr.parameterResults[2] = expr.iffNullResult(1, rowCount) - return nil + expr.parameterResults[2], err = expr.iffNullResult(1, rowCount) + return err } func hasSelectedRows(selectList []bool) bool { @@ -643,19 +790,24 @@ func hasSelectedRows(selectList []bool) bool { return false } -func (expr *FunctionExpressionExecutor) iffNullResult(index, length int) *vector.Vector { +func (expr *FunctionExpressionExecutor) iffNullResult(index, length int) (*vector.Vector, error) { typ := expr.resultType result := expr.iffNullResults[index] if result == nil || *result.GetType() != typ { if result != nil { result.Free(expr.m) + expr.iffNullResults[index] = nil + } + var err error + result, err = newExpressionConstNull(typ, length, expr.allocation) + if err != nil { + return nil, err } - result = vector.NewConstNull(typ, length, expr.m) expr.iffNullResults[index] = result } else { result.SetLength(length) } - return result + return result, nil } func (expr *FunctionExpressionExecutor) EvalCase(proc *process.Process, batches []*batch.Batch, selectList []bool) (err error) { @@ -786,7 +938,13 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( if rowAligned && !parameter.IsConst() { selected := expr.selectedParameterVectors[i] if selected == nil { - selected = vector.NewOffHeapVecWithType(*parameter.GetType()) + var err error + selected, err = newExpressionOffHeapVector( + *parameter.GetType(), expr.allocation, + ) + if err != nil { + return nil, err + } expr.selectedParameterVectors[i] = selected } else { selected.Reset(*parameter.GetType()) @@ -806,7 +964,13 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( return nil, err } if expr.selectedResult == nil { - expr.selectedResult = vector.NewFunctionResultWrapper(expr.resultType, expr.m) + var err error + expr.selectedResult, err = vector.NewFunctionResultWrapperWithAllocation( + expr.resultType, expr.m, expr.allocation, + ) + if err != nil { + return nil, err + } } expr.resetResultType(expr.selectedResult) if err := expr.selectedResult.PreExtendAndReset(selectedCount); err != nil { @@ -826,7 +990,13 @@ func (expr *FunctionExpressionExecutor) evalSelectedRows( result.SetIsBin(runtimeIsBin) result.ResetWithSameType() if expr.selectedNullResult == nil { - expr.selectedNullResult = vector.NewConstNull(runtimeType, 1, expr.m) + var err error + expr.selectedNullResult, err = newExpressionConstNull( + runtimeType, 1, expr.allocation, + ) + if err != nil { + return nil, err + } } else { expr.selectedNullResult.SetType(runtimeType) expr.selectedNullResult.SetLength(1) @@ -1008,19 +1178,27 @@ func (expr *ColumnExpressionExecutor) Eval(_ *process.Process, batches []*batch. vec := batches[relIndex].Vecs[expr.colIndex] if vec.IsConstNull() { - vec = expr.getConstNullVec(expr.typ, vec.Length()) + var err error + vec, err = expr.getConstNullVec(expr.typ, vec.Length()) + if err != nil { + return nil, err + } } return vec, nil } -func (expr *ColumnExpressionExecutor) getConstNullVec(typ types.Type, length int) *vector.Vector { +func (expr *ColumnExpressionExecutor) getConstNullVec(typ types.Type, length int) (*vector.Vector, error) { if expr.nullVecCache != nil { expr.nullVecCache.SetType(typ) expr.nullVecCache.SetLength(length) } else { - expr.nullVecCache = vector.NewConstNull(typ, length, expr.mp) + var err error + expr.nullVecCache, err = newExpressionConstNull(typ, length, expr.allocation) + if err != nil { + return nil, err + } } - return expr.nullVecCache + return expr.nullVecCache, nil } func (expr *ColumnExpressionExecutor) EvalWithoutResultReusing(proc *process.Process, batches []*batch.Batch, _ []bool) (*vector.Vector, error) { @@ -1077,109 +1255,116 @@ func (expr *FixedVectorExpressionExecutor) IsColumnExpr() bool { return false } -func generateConstExpressionExecutor(proc *process.Process, typ types.Type, con *plan.Literal) (vec *vector.Vector, err error) { +func generateConstExpressionExecutor( + proc *process.Process, + typ types.Type, + con *plan.Literal, + selection *vector.AllocationAccountSelection, +) (vec *vector.Vector, err error) { if con.GetIsnull() { - vec = vector.NewConstNull(typ, 1, proc.Mp()) + vec, err = newExpressionConstNull(typ, 1, selection) } else { switch val := con.GetValue().(type) { case *plan.Literal_Bval: - vec, err = vector.NewConstFixed(constBType, val.Bval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constBType, val.Bval, 1, proc.Mp(), selection) case *plan.Literal_I8Val: - vec, err = vector.NewConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI8Type, int8(val.I8Val), 1, proc.Mp(), selection) case *plan.Literal_I16Val: - vec, err = vector.NewConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI16Type, int16(val.I16Val), 1, proc.Mp(), selection) case *plan.Literal_I32Val: - vec, err = vector.NewConstFixed(constI32Type, val.I32Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI32Type, val.I32Val, 1, proc.Mp(), selection) case *plan.Literal_I64Val: - vec, err = vector.NewConstFixed(constI64Type, val.I64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constI64Type, val.I64Val, 1, proc.Mp(), selection) case *plan.Literal_U8Val: - vec, err = vector.NewConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU8Type, uint8(val.U8Val), 1, proc.Mp(), selection) case *plan.Literal_U16Val: - vec, err = vector.NewConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU16Type, uint16(val.U16Val), 1, proc.Mp(), selection) case *plan.Literal_U32Val: - vec, err = vector.NewConstFixed(constU32Type, val.U32Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU32Type, val.U32Val, 1, proc.Mp(), selection) case *plan.Literal_U64Val: if typ.Oid == types.T_bit { - vec, err = vector.NewConstFixed(typ, val.U64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, val.U64Val, 1, proc.Mp(), selection) } else { - vec, err = vector.NewConstFixed(constU64Type, val.U64Val, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constU64Type, val.U64Val, 1, proc.Mp(), selection) } case *plan.Literal_Fval: - vec, err = vector.NewConstFixed(constFType, val.Fval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constFType, val.Fval, 1, proc.Mp(), selection) case *plan.Literal_Dval: - vec, err = vector.NewConstFixed(constDType, val.Dval, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constDType, val.Dval, 1, proc.Mp(), selection) case *plan.Literal_Dateval: - vec, err = vector.NewConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constDateType, types.Date(val.Dateval), 1, proc.Mp(), selection) case *plan.Literal_Timeval: - vec, err = vector.NewConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, types.Time(val.Timeval), 1, proc.Mp(), selection) case *plan.Literal_Datetimeval: - vec, err = vector.NewConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, types.Datetime(val.Datetimeval), 1, proc.Mp(), selection) case *plan.Literal_Decimal64Val: cd64 := val.Decimal64Val d64 := types.Decimal64(cd64.A) - vec, err = vector.NewConstFixed(typ, d64, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, d64, 1, proc.Mp(), selection) case *plan.Literal_Decimal128Val: cd128 := val.Decimal128Val d128 := types.Decimal128{B0_63: uint64(cd128.A), B64_127: uint64(cd128.B)} - vec, err = vector.NewConstFixed(typ, d128, 1, proc.Mp()) + vec, err = newExpressionConstFixed(typ, d128, 1, proc.Mp(), selection) case *plan.Literal_Timestampval: scale := typ.Scale if scale < 0 || scale > 6 { return nil, moerr.NewErrTooBigPrecision(proc.Ctx, int64(scale), "TIMESTAMP", 6) } - vec, err = vector.NewConstFixed(constTimestampTypes[scale], types.Timestamp(val.Timestampval), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constTimestampTypes[scale], types.Timestamp(val.Timestampval), 1, proc.Mp(), selection) case *plan.Literal_Sval: sval := val.Sval // Distinguish binary with non-binary string. if typ.Oid == types.T_binary || typ.Oid == types.T_varbinary || typ.Oid == types.T_blob { - vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) } else if typ.Oid == types.T_geometry { - vec, err = vector.NewConstBytes(typ, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(typ, []byte(sval), 1, proc.Mp(), selection) } else if typ.Oid == types.T_array_float32 { array, err1 := types.StringToArray[float32](sval) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) } else if typ.Oid == types.T_array_float64 { array, err1 := types.StringToArray[float64](sval) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstArray(typ, array, 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, array, 1, proc.Mp(), selection) } else if typ.Oid == types.T_datalink { _, _, err1 := datalink.ParseDatalink(sval, proc) if err1 != nil { return nil, err1 } - vec, err = vector.NewConstBytes(constBinType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constBinType, []byte(sval), 1, proc.Mp(), selection) } else { - vec, err = vector.NewConstBytes(constSType, []byte(sval), 1, proc.Mp()) + vec, err = newExpressionConstBytes(constSType, []byte(sval), 1, proc.Mp(), selection) } case *plan.Literal_Defaultval: defaultVal := val.Defaultval - vec, err = vector.NewConstFixed(constBType, defaultVal, 1, proc.Mp()) + vec, err = newExpressionConstFixed(constBType, defaultVal, 1, proc.Mp(), selection) case *plan.Literal_EnumVal: - vec, err = vector.NewConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp()) + vec, err = newExpressionConstFixed(constEnumType, types.Enum(val.EnumVal), 1, proc.Mp(), selection) case *plan.Literal_VecVal: switch typ.Oid { case types.T_array_float32: - vec, err = vector.NewConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[float32]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_float64: - vec, err = vector.NewConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[float64]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_bf16: - vec, err = vector.NewConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[types.BF16]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_float16: - vec, err = vector.NewConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[types.Float16]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_int8: - vec, err = vector.NewConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[int8]([]byte(val.VecVal)), 1, proc.Mp(), selection) case types.T_array_uint8: - vec, err = vector.NewConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp()) + vec, err = newExpressionConstArray(typ, types.BytesToArray[uint8]([]byte(val.VecVal)), 1, proc.Mp(), selection) } default: return nil, moerr.NewNYI(proc.Ctx, fmt.Sprintf("const expression %v", con.GetValue())) } - vec.SetIsBin(con.IsBin) + if err == nil { + vec.SetIsBin(con.IsBin) + } } return vec, err } diff --git a/pkg/sql/colexec/evalExpressionAllocation_test.go b/pkg/sql/colexec/evalExpressionAllocation_test.go new file mode 100644 index 0000000000000..32dc2e9a299f0 --- /dev/null +++ b/pkg/sql/colexec/evalExpressionAllocation_test.go @@ -0,0 +1,106 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package colexec + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +func TestAccountedExpressionTreeCoversNestedAndSelectedResults(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + typ := types.T_varchar.ToType() + column := &plan.Expr{ + Typ: plan.Type{Id: int32(typ.Oid), Width: typ.Width}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + } + bindFunction := func(name string, args ...*plan.Expr) *plan.Expr { + argTypes := make([]types.Type, len(args)) + for i := range args { + argTypes[i] = types.New( + types.T(args[i].Typ.Id), args[i].Typ.Width, args[i].Typ.Scale, + ) + } + fn, bindErr := function.GetFunctionByName(proc.Ctx, name, argTypes) + require.NoError(t, bindErr) + retType := fn.GetReturnType() + return &plan.Expr{ + Typ: plan.Type{ + Id: int32(retType.Oid), Width: retType.Width, Scale: retType.Scale, + }, + Expr: &plan.Expr_F{F: &plan.Function{ + Func: &plan.ObjectRef{Obj: fn.GetEncodedOverloadID(), ObjName: name}, + Args: args, + }}, + } + } + literal := &plan.Expr{ + Typ: plan.Type{Id: int32(typ.Oid), Width: typ.Width}, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Value: &plan.Literal_Sval{Sval: "-"}, + }}, + } + expression := bindFunction("concat", bindFunction("lower", column), column, literal) + executor, err := NewExpressionExecutorWithAllocation(proc, expression, selection) + require.NoError(t, err) + root := executor.(*FunctionExpressionExecutor) + nested := root.parameterExecutor[0].(*FunctionExpressionExecutor) + + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector( + []string{"AA", "BB", "CC", "DD"}, nil, proc.Mp(), + ) + input.SetRowCount(4) + defer input.Clean(proc.Mp()) + result, err := executor.Eval( + proc, []*batch.Batch{input}, []bool{true, false, true, false}, + ) + require.NoError(t, err) + require.Equal(t, []string{"aaAA-", "", "ccCC-", ""}, vector.InefficientMustStrCol(result)) + + assertFunctionStorage := func(function *FunctionExpressionExecutor) { + t.Helper() + require.Same(t, selection, function.resultVector.GetResultVector().AllocationAccountSelection()) + require.Same(t, selection, function.selectedResult.GetResultVector().AllocationAccountSelection()) + for _, selected := range function.selectedParameterVectors { + if selected != nil { + require.Same(t, selection, selected.AllocationAccountSelection()) + } + } + } + assertFunctionStorage(root) + assertFunctionStorage(nested) + fixed := root.parameterExecutor[2].(*FixedVectorExpressionExecutor) + require.Same(t, selection, fixed.resultVector.AllocationAccountSelection()) + require.Positive(t, account.Snapshot().Used) + + executor.Free() + require.Zero(t, account.Snapshot().Used) +} diff --git a/pkg/sql/colexec/evalExpressionMemory_test.go b/pkg/sql/colexec/evalExpressionMemory_test.go deleted file mode 100644 index 300fa5439157a..0000000000000 --- a/pkg/sql/colexec/evalExpressionMemory_test.go +++ /dev/null @@ -1,156 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed 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. - -package colexec - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestExpressionExecutorRetainedBytesExcludesBorrowedPlanVector(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - source := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - data, err := source.MarshalBinary() - require.NoError(t, err) - source.Free(proc.Mp()) - require.Zero(t, proc.Mp().CurrNB()) - - executor, err := NewExpressionExecutor(proc, &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Vec{Vec: &plan.LiteralVec{ - Len: 3, - Data: data, - }}, - }) - require.NoError(t, err) - fixed, ok := executor.(*FixedVectorExpressionExecutor) - require.True(t, ok) - require.True(t, fixed.resultVector.NeedDup()) - require.Positive(t, fixed.resultVector.Allocated()) - - retained, known := ExpressionExecutorRetainedBytes(executor) - require.True(t, known) - require.Zero(t, retained) - executor.Free() - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestExpressionExecutorRetainedBytesCoversOwnedExecutorTree(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - newVec := func(values ...int32) *vector.Vector { - return testutil.MakeInt32Vector(values, nil, proc.Mp()) - } - - fixedVec := newVec(1) - listResult := newVec(1, 2) - paramNull := newVec(1) - paramMaskedNull := newVec(2) - paramVec := newVec(3) - varNull := newVec(4) - varMaskedNull := newVec(5) - varVec := newVec(6) - selectedParameter := newVec(7, 8) - selectedNull := vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()) - iffNull0 := vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()) - iffNull1 := vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()) - ownedVectors := []*vector.Vector{ - fixedVec, - listResult, - paramNull, - paramMaskedNull, - paramVec, - varNull, - varMaskedNull, - varVec, - selectedParameter, - selectedNull, - iffNull0, - iffNull1, - } - - result := vector.NewFunctionResultWrapper(types.T_int32.ToType(), proc.Mp()) - require.NoError(t, result.PreExtendAndReset(4)) - selectedResult := vector.NewFunctionResultWrapper(types.T_int32.ToType(), proc.Mp()) - require.NoError(t, selectedResult.PreExtendAndReset(2)) - - fixed := &FixedVectorExpressionExecutor{resultVector: fixedVec} - list := &ListExpressionExecutor{ - resultVector: listResult, - parameterExecutor: []ExpressionExecutor{fixed}, - } - param := &ParamExpressionExecutor{ - null: paramNull, - maskedNull: paramMaskedNull, - vec: paramVec, - } - variable := &VarExpressionExecutor{ - null: varNull, - maskedNull: varMaskedNull, - vec: varVec, - } - function := &FunctionExpressionExecutor{ - resultVector: result, - selectedResult: selectedResult, - selectedNullResult: selectedNull, - selectedParameterVectors: []*vector.Vector{selectedParameter}, - parameterExecutor: []ExpressionExecutor{param}, - iffNullResults: [2]*vector.Vector{iffNull0, iffNull1}, - } - column := &ColumnExpressionExecutor{ - nullVecCache: vector.NewConstNull(types.T_int32.ToType(), 2, proc.Mp()), - } - - executors := []ExpressionExecutor{list, variable, function, column, nil} - retained, known := ExpressionExecutorsRetainedBytes(executors) - require.True(t, known) - require.Positive(t, retained) - for _, executor := range executors { - _, known = ExpressionExecutorRetainedBytes(executor) - require.True(t, known) - } - _, known = ExpressionExecutorRetainedBytes(retainedBytesUnknownExecutor{}) - require.False(t, known) - - result.Free() - selectedResult.Free() - column.nullVecCache.Free(proc.Mp()) - for _, vec := range ownedVectors { - vec.Free(proc.Mp()) - } - require.Zero(t, proc.Mp().CurrNB()) -} - -type retainedBytesUnknownExecutor struct{} - -func (retainedBytesUnknownExecutor) Eval(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (retainedBytesUnknownExecutor) EvalWithoutResultReusing(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (retainedBytesUnknownExecutor) ResetForNextQuery() {} -func (retainedBytesUnknownExecutor) Free() {} -func (retainedBytesUnknownExecutor) IsColumnExpr() bool { return false } -func (retainedBytesUnknownExecutor) TypeName() string { return "unknown" } diff --git a/pkg/sql/colexec/fuzzyfilter/filter.go b/pkg/sql/colexec/fuzzyfilter/filter.go index f51b12c88d64e..1f6ff802d1252 100644 --- a/pkg/sql/colexec/fuzzyfilter/filter.go +++ b/pkg/sql/colexec/fuzzyfilter/filter.go @@ -19,6 +19,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/bloomfilter" "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -148,9 +149,19 @@ func (fuzzyFilter *FuzzyFilter) Prepare(proc *process.Process) (err error) { ) != keycodec.ExactRuntimeFilterUnsupported } if ctr.runtimeFilterUsable { + if fuzzyFilter.allocationAccount == nil || + fuzzyFilter.runtimeFilterAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } if ctr.pass2RuntimeFilter == nil { - ctr.pass2RuntimeFilter = vector.NewOffHeapVecWithType( - plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp)) + ctr.pass2RuntimeFilter, err = + vector.NewOffHeapVecWithTypeAndAllocation( + plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp), + fuzzyFilter.runtimeFilterAllocation, + ) + if err != nil { + return err + } } } else if ctr.pass2RuntimeFilter != nil { // FuzzyFilter must still execute its uniqueness check, but an @@ -391,7 +402,7 @@ func (fuzzyFilter *FuzzyFilter) handleRuntimeFilter(proc *process.Process) error } if encoding == keycodec.ExactRuntimeFilterFloatZeroClosed { if err := runtimefilter.CloseFloatSignedZero( - ctr.pass2RuntimeFilter, proc.Mp(), nil); err != nil { + ctr.pass2RuntimeFilter, proc.Mp()); err != nil { if fuzzyFilter.fallbackRuntimeFilter(proc, err) { return nil } @@ -410,16 +421,13 @@ func (fuzzyFilter *FuzzyFilter) handleRuntimeFilter(proc *process.Process) error // Reset bitmap before sort to avoid corruption. ctr.pass2RuntimeFilter.GetNulls().Reset() ctr.pass2RuntimeFilter.InplaceSort() - budget, err := proc.GetHashBuildBudget() - if err != nil { - if fuzzyFilter.fallbackRuntimeFilter(proc, err) { - return nil - } - fuzzyFilter.abandonRuntimeFilter(proc) - return err - } data, release, err := runtimefilter.MarshalExactFilterVector( - ctr.pass2RuntimeFilter, budget) + ctr.pass2RuntimeFilter, + proc.Mp(), + fuzzyFilter.allocationAccount, + fuzzyFilterAllocationOwner, + fuzzyFilterAllocationSiteRuntimeFilterPayload, + ) if err != nil { if fuzzyFilter.fallbackRuntimeFilter(proc, err) { return nil @@ -536,11 +544,6 @@ func (fuzzyFilter *FuzzyFilter) generate() error { ctr := &fuzzyFilter.ctr rbat := batch.NewWithSize(1) rbat.SetVector(0, vector.NewVec(plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp))) - // Runtime-filter retention is optional and can grow to the configured IN - // cardinality. Keep it off-heap so the process pool can reject growth - // recoverably; appendPassToRuntimeFilter then abandons it and sends PASS. - ctr.pass2RuntimeFilter = vector.NewOffHeapVecWithType( - plan.MakeTypeByPlan2Type(fuzzyFilter.PkTyp)) ctr.rbat = rbat return nil } diff --git a/pkg/sql/colexec/fuzzyfilter/filter_test.go b/pkg/sql/colexec/fuzzyfilter/filter_test.go index 57a42546aa39d..631e485f5f5ff 100644 --- a/pkg/sql/colexec/fuzzyfilter/filter_test.go +++ b/pkg/sql/colexec/fuzzyfilter/filter_test.go @@ -160,7 +160,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(101, probeType, payloadType) arg, proc := newRuntimeFilterTest(t, spec, payloadType) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.False(t, arg.ctr.runtimeFilterUsable) require.Nil(t, arg.ctr.pass2RuntimeFilter) @@ -187,7 +187,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(102, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.True(t, arg.ctr.runtimeFilterUsable) require.NotNil(t, arg.ctr.pass2RuntimeFilter) require.Zero(t, arg.ctr.pass2RuntimeFilter.Length()) @@ -206,7 +206,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(103, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.True(t, arg.ctr.runtimeFilterUsable) payload := testutil.MakeInt64Vector([]int64{7}, nil, proc.Mp()) require.NoError(t, arg.appendPassToRuntimeFilter(payload, proc)) @@ -225,7 +225,7 @@ func TestRuntimeFilterContract(t *testing.T) { spec := newRuntimeFilterSpec(104, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) payload := testutil.MakeInt32Vector([]int32{7, 3}, nil, proc.Mp()) require.NoError(t, arg.appendPassToRuntimeFilter(payload, proc)) @@ -275,7 +275,7 @@ func TestFuzzyRuntimeFilterCopyFailureFailsOpen(t *testing.T) { arg := newArgument(typ) arg.N = 1 arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) sourceMP := mpool.MustNewZero() payload := vector.NewVec(typ) @@ -338,7 +338,7 @@ func TestFuzzyRuntimeFilterClosureFailureFailsOpen(t *testing.T) { arg := newArgument(typ) arg.N = 1 arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.NoError(t, vector.AppendFixed( arg.ctr.pass2RuntimeFilter, float64(0), false, limited)) for arg.ctr.pass2RuntimeFilter.Length() < @@ -381,17 +381,26 @@ func TestFuzzyRuntimeFilterBudgetErrorPolicy(t *testing.T) { typ := types.T_int32.ToType() spec := newRuntimeFilterSpec(109, typ, typ) arg, proc := newRuntimeFilterTest(t, spec, typ) - require.NoError(t, arg.Prepare(proc)) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController( + 2*generation.Cap(), generation) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + prepareFuzzyFilter(t, arg, proc) require.NoError(t, vector.AppendFixed( arg.ctr.pass2RuntimeFilter, int32(1), false, proc.Mp())) - generation, err := proc.GetHashBuildBudget() - require.NoError(t, err) - var held *process.HashBuildReservation + var filler []byte if test.closed { generation.Close() } else { - held, err = generation.Reserve(generation.Cap()) + remaining := generation.Cap() - generation.Used() + filler, err = proc.Mp().AllocAccounted( + int(remaining), account, 63, 255) require.NoError(t, err) } @@ -410,7 +419,8 @@ func TestFuzzyRuntimeFilterBudgetErrorPolicy(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) require.Equal(t, int64(1), stats["FuzzyFilterRuntimeFilterBudgetFallbacks"]) - require.True(t, held.Release()) + proc.Mp().Free(filler) + generation.Close() } require.False(t, arg.ctr.runtimeFilterUsable) require.Nil(t, arg.ctr.pass2RuntimeFilter) @@ -447,7 +457,7 @@ func TestFuzzyCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { require.NoError(t, buildChild.Prepare(proc)) require.NoError(t, probeChild.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) _, err := vm.Exec(arg, proc) require.ErrorIs(t, err, buildErr) @@ -473,7 +483,7 @@ func TestFuzzyCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) proc.GetMessageBoard().Reset() - require.NoError(t, arg.Prepare(proc)) + prepareFuzzyFilter(t, arg, proc) require.False(t, arg.ctr.runtimeFilterDone, "Prepare must open the terminal gate for the next generation") arg.finalizeBuildFailure(proc) @@ -581,6 +591,22 @@ func newRuntimeFilterTest( return arg, proc } +func prepareFuzzyFilter( + t *testing.T, + arg *FuzzyFilter, + proc *process.Process, +) { + t.Helper() + if arg.allocationAccount == nil { + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + } + require.NoError(t, arg.Prepare(proc)) +} + func receiveRuntimeFilter( t *testing.T, proc *process.Process, diff --git a/pkg/sql/colexec/fuzzyfilter/types.go b/pkg/sql/colexec/fuzzyfilter/types.go index 40f18329461ad..48627bcfb8908 100644 --- a/pkg/sql/colexec/fuzzyfilter/types.go +++ b/pkg/sql/colexec/fuzzyfilter/types.go @@ -17,6 +17,7 @@ package fuzzyfilter import ( "github.com/matrixorigin/matrixone/pkg/common/bloomfilter" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -37,6 +38,16 @@ const ( End ) +const fuzzyFilterAllocationOwner mpool.AllocationOwner = 1 + +const ( + fuzzyFilterAllocationSiteRuntimeFilterData mpool.AllocationSite = iota + 1 + fuzzyFilterAllocationSiteRuntimeFilterArea + fuzzyFilterAllocationSiteRuntimeFilterNulls + fuzzyFilterAllocationSiteRuntimeFilterGrouping + fuzzyFilterAllocationSiteRuntimeFilterPayload +) + type container struct { state int @@ -55,7 +66,9 @@ type container struct { } type FuzzyFilter struct { - ctr container + ctr container + allocationAccount *mpool.AllocationAccount + runtimeFilterAllocation *vector.AllocationAccountSelection // Estimates of the number of data items obtained from statistical information N float64 @@ -69,6 +82,51 @@ type FuzzyFilter struct { vm.OperatorBase } +func (fuzzyFilter *FuzzyFilter) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil { + return mpool.ErrAllocationAccountInvalid + } + if fuzzyFilter.allocationAccount != nil { + if fuzzyFilter.allocationAccount == account { + return nil + } + return mpool.ErrAllocationAccountMismatch + } + selection, err := vector.NewAllocationAccountSelection( + account, + fuzzyFilterAllocationOwner, + fuzzyFilterAllocationSiteRuntimeFilterData, + fuzzyFilterAllocationSiteRuntimeFilterArea, + fuzzyFilterAllocationSiteRuntimeFilterNulls, + fuzzyFilterAllocationSiteRuntimeFilterGrouping, + ) + if err != nil { + return err + } + fuzzyFilter.allocationAccount = account + fuzzyFilter.runtimeFilterAllocation = selection + return nil +} + +func (fuzzyFilter *FuzzyFilter) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if fuzzyFilter.allocationAccount == nil { + return nil + } + if fuzzyFilter.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if fuzzyFilter.ctr.pass2RuntimeFilter != nil { + return mpool.ErrAllocationAccountInvariant + } + fuzzyFilter.allocationAccount = nil + fuzzyFilter.runtimeFilterAllocation = nil + return nil +} + func (fuzzyFilter *FuzzyFilter) GetOperatorBase() *vm.OperatorBase { return &fuzzyFilter.OperatorBase } @@ -127,7 +185,8 @@ func (fuzzyFilter *FuzzyFilter) Reset(proc *process.Process, pipelineFailed bool ctr.runtimeFilterUsable = false ctr.collisionCnt = 0 if ctr.pass2RuntimeFilter != nil { - ctr.pass2RuntimeFilter.CleanOnlyData() + ctr.pass2RuntimeFilter.Free(proc.Mp()) + ctr.pass2RuntimeFilter = nil } if ctr.rbat != nil { ctr.rbat.CleanOnlyData() diff --git a/pkg/sql/colexec/group/helper.go b/pkg/sql/colexec/group/helper.go index 4e5f408a927aa..d6b2d81b40759 100644 --- a/pkg/sql/colexec/group/helper.go +++ b/pkg/sql/colexec/group/helper.go @@ -391,7 +391,9 @@ func (ctr *container) spillDataToDisk(proc *process.Process, opAnalyzer process. } } gbBatch.SetRowCount(int(cnt)) - gbBatch.MarshalBinaryWithBuffer(buf, false) + if _, err := gbBatch.MarshalBinaryWithBuffer(buf, false); err != nil { + return 0, 0, err + } // write marker var magic uint64 = 0x12345678DEADBEEF diff --git a/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go b/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..c06b9c0da4696 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/allocation_test_helpers_test.go @@ -0,0 +1,86 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func newTestAllocationAccount(t testing.TB) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + return account +} + +func installTestHashBuildBudget( + t testing.TB, + op *HashBuild, + generation *process.HashBuildBudgetGeneration, +) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, generation) + require.NoError(t, err) + replaceTestHashBuildAllocation(t, op, account) + op.ctr.hashmapBuilder.setBudget(generation) +} + +func installTestProcessHashBuildBudget( + t testing.TB, + op *HashBuild, + proc *process.Process, +) *process.HashBuildBudgetGeneration { + t.Helper() + generation, err := proc.GetHashBuildBudget() + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, generation) + require.NoError(t, err) + replaceTestHashBuildAllocation(t, op, account) + return generation +} + +func newTestHashmapBuilder(t testing.TB) *HashmapBuilder { + t.Helper() + builder := &HashmapBuilder{} + require.NoError(t, builder.SetAllocationAccount(newTestAllocationAccount(t))) + return builder +} + +func installTestHashBuildAllocation(t testing.TB, op *HashBuild) { + t.Helper() + require.NoError(t, op.SetAllocationAccount(newTestAllocationAccount(t))) +} + +func replaceTestHashBuildAllocation( + t testing.TB, + op *HashBuild, + account *mpool.AllocationAccount, +) { + t.Helper() + if current := op.ctr.hashmapBuilder.mapAllocationAccount; current != nil { + require.NoError(t, op.ClearAllocationAccount(current)) + } + require.NoError(t, op.SetAllocationAccount(account)) +} diff --git a/pkg/sql/colexec/hashbuild/budget.go b/pkg/sql/colexec/hashbuild/budget.go index ab38dc0d2fc99..31d9fd1ebc05e 100644 --- a/pkg/sql/colexec/hashbuild/budget.go +++ b/pkg/sql/colexec/hashbuild/budget.go @@ -15,766 +15,123 @@ package hashbuild import ( - "math" - "sync" - "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/hashtable" "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/runtimefilter" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" ) -type hashMapResizeReservation struct { - owner *hashMapReservationOwner - token *process.HashBuildReservation -} - -func (r *hashMapResizeReservation) Commit(plan hashtable.ResizePlan) { - r.owner.commit(r.token, plan.ReuseCurrentBlocks) - r.token = nil -} - -func (r *hashMapResizeReservation) Rollback() { - if r.token != nil { - r.token.Release() - r.token = nil - } -} - -// hashMapReservationOwner follows the physical hash table across producer to -// JoinMap ownership transfer. Full-table replacement swaps the retained token; -// segmented growth keeps the existing tokens and adds one for the appended -// blocks. Resize callbacks retain this owner so consumer growth never stores -// reservations back into a reused producer. -type hashMapReservationOwner struct { - mu sync.Mutex - tokens []*process.HashBuildReservation -} - -func (o *hashMapReservationOwner) commit(token *process.HashBuildReservation, reuseCurrent bool) { - o.mu.Lock() - if reuseCurrent { - o.tokens = append(o.tokens, token) - o.mu.Unlock() - return - } - old := o.tokens - o.tokens = []*process.HashBuildReservation{token} - o.mu.Unlock() - for _, reservation := range old { - reservation.Release() - } -} - -func (o *hashMapReservationOwner) release() { - if o == nil { - return - } - o.mu.Lock() - tokens := o.tokens - o.tokens = nil - o.mu.Unlock() - for _, token := range tokens { - token.Release() - } -} - -func (hb *HashmapBuilder) setBudget(budget *process.HashBuildBudgetGeneration) { +// setBudget retains the statement generation for non-memory resource ledgers +// such as spill files and disk. Physical memory admission is exclusively +// driven by the allocation account installed through SetAllocationAccount. +func (hb *HashmapBuilder) setBudget( + budget *process.HashBuildBudgetGeneration, +) { hb.budget = budget } // SetBudget is the exported boundary used by spill and integration tests. -func (hb *HashmapBuilder) SetBudget(budget *process.HashBuildBudgetGeneration) { hb.setBudget(budget) } - -func (hb *HashmapBuilder) reserveInitialMap(size int64) error { - if hb.budget == nil || size <= 0 { - return nil - } - reservation, err := hb.budget.Reserve(uint64(size)) - if err != nil { - return err - } - hb.mapReservation = &hashMapReservationOwner{tokens: []*process.HashBuildReservation{reservation}} - return nil -} - -func resizeAdmission(budget *process.HashBuildBudgetGeneration, owner *hashMapReservationOwner, plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - if budget == nil || plan.AdditionalBytes == 0 { - return nil, nil - } - token, err := budget.Reserve(plan.AdditionalBytes) - if err != nil { - return nil, err - } - return &hashMapResizeReservation{owner: owner, token: token}, nil +func (hb *HashmapBuilder) SetBudget( + budget *process.HashBuildBudgetGeneration, +) { + hb.setBudget(budget) } -// NewBudgetedEmptyJoinMap creates an initially empty JoinMap whose complete -// physical hash-table lifetime is charged to budget. The initial allocation is -// admitted before touching the mpool, every later resize uses the same -// generation, and JoinMap.Free releases all retained reservations. -// -// This is used by consumers that must grow a hash table from probe-side keys -// (for example RightDedupJoin after an empty build partition). Such maps cannot -// use the regular HashmapBuilder ownership transfer because there is no build -// batch to publish. -func NewBudgetedEmptyJoinMap( +// NewAccountedEmptyJoinMap creates a consumer-grown map under the statement +// allocation generation. The map and its string-key iterator scratch carry +// the same immutable provenance as producer-built maps. +func NewAccountedEmptyJoinMap( keyWidth int, - budget *process.HashBuildBudgetGeneration, + account *mpool.AllocationAccount, mp *mpool.MPool, ) (*message.JoinMap, error) { - if budget == nil || mp == nil { - return nil, process.ErrHashBuildBudgetInvalid - } - - initialBytes := hashtable.Int64HashMapInitialAllocationBytes() - if keyWidth > 8 { - initialBytes = hashtable.StringHashMapInitialAllocationBytes() - } - initial, err := budget.Reserve(initialBytes) + if account == nil || mp == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + selection, err := hashtable.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashCell, + HashBuildAllocationSiteHashDescriptor, + ) if err != nil { return nil, err } - owner := &hashMapReservationOwner{ - tokens: []*process.HashBuildReservation{initial}, - } - var ( intHashMap *hashmap.IntHashMap strHashMap *hashmap.StrHashMap ) if keyWidth <= 8 { - intHashMap, err = hashmap.NewIntHashMap(false, mp) - if err == nil { - intHashMap.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - } + intHashMap, err = hashmap.NewIntHashMapWithAllocation( + false, + mp, + selection, + ) } else { - strHashMap, err = hashmap.NewStrHashMap(false, mp) - if err == nil { - strHashMap.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - } + iteratorAllocation, allocationErr := hashmap.NewIteratorAllocation( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashIterator, + ) + if allocationErr != nil { + return nil, allocationErr + } + strHashMap, err = hashmap.NewStrHashMapWithAllocations( + false, + mp, + selection, + iteratorAllocation, + ) } if err != nil { - owner.release() return nil, err } - jm := message.NewJoinMap(message.GroupSels{}, intHashMap, strHashMap, nil, nil, mp) - jm.SetMemoryRelease(owner.release) + jm := message.NewJoinMap( + message.GroupSels{}, + intHashMap, + strHashMap, + nil, + nil, + mp, + ) jm.IncRef(1) return jm, nil } -func (hb *HashmapBuilder) attachIntHashMapAdmission(m *hashmap.IntHashMap) error { - owner := hb.mapReservation - budget := hb.budget - m.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - return nil -} - -func (hb *HashmapBuilder) attachStrHashMapAdmission(m *hashmap.StrHashMap) error { - owner := hb.mapReservation - budget := hb.budget - m.SetResizeAdmission(func(plan hashtable.ResizePlan) (hashtable.ResizeReservation, error) { - return resizeAdmission(budget, owner, plan) - }) - return nil -} - -func batchesAllocated(batches []*batch.Batch) uint64 { - var total uint64 - for _, bat := range batches { - if bat != nil { - total += uint64(bat.Allocated()) - } - } - return total -} - -type batchCopyAllocationSnapshot struct { - length int - tail *batch.Batch - tailAllocated uint64 -} - -func snapshotBatchCopyAllocation(batches []*batch.Batch) (batchCopyAllocationSnapshot, error) { - snapshot := batchCopyAllocationSnapshot{length: len(batches)} - if snapshot.length == 0 { - return snapshot, nil - } - snapshot.tail = batches[snapshot.length-1] - if snapshot.tail == nil { - return batchCopyAllocationSnapshot{}, process.ErrHashBuildBudgetInvalid - } - allocated := snapshot.tail.Allocated() - if allocated < 0 { - return batchCopyAllocationSnapshot{}, process.ErrHashBuildBudgetInvalid - } - snapshot.tailAllocated = uint64(allocated) - return snapshot, nil -} - -// batchCopyAllocatedDelta relies on CopyIntoBatches' append-only contract: it -// may grow the old partial tail and append destination batches. A full-size -// source can swap one new batch with that partial tail, so inspect the old tail -// plus the appended suffix by identity instead of rescanning every retained -// batch. Across a build this keeps retained-copy accounting linear in the -// number of destination batches rather than quadratic. -func batchCopyAllocatedDelta( - batches []*batch.Batch, - snapshot batchCopyAllocationSnapshot, -) (uint64, error) { - if snapshot.length < 0 || len(batches) < snapshot.length { - return 0, process.ErrHashBuildBudgetInvalid - } - start := 0 - seenTail := snapshot.length == 0 - if snapshot.length > 0 { - if snapshot.tail == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - start = snapshot.length - 1 - } - var delta uint64 - for i := start; i < len(batches); i++ { - bat := batches[i] - if bat == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - allocated := bat.Allocated() - if allocated < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - value := uint64(allocated) - if bat == snapshot.tail { - if seenTail || value < snapshot.tailAllocated { - return 0, process.ErrHashBuildBudgetInvalid - } - seenTail = true - value -= snapshot.tailAllocated - } - if delta > math.MaxUint64-value { - return 0, process.ErrHashBuildBudgetInvalid - } - delta += value - } - if !seenTail { - return 0, process.ErrHashBuildBudgetInvalid - } - return delta, nil -} - -func (hb *HashmapBuilder) copyBuildBatch(src *batch.Batch, proc *process.Process) error { - if hb.budget == nil { - return hb.Batches.CopyIntoBatches(src, proc) - } - projection, err := hb.projectedBatchCopy(src) - if err != nil { - return err - } - return hb.copyBuildBatchProjected(src, proc, projection) -} - -func (hb *HashmapBuilder) copyBuildBatchProjected( +func (hb *HashmapBuilder) copyBuildBatch( src *batch.Batch, proc *process.Process, - projection batchCopyProjection, ) error { - if hb.budget == nil { - if err := hb.Batches.CopyIntoBatches(src, proc); err != nil { - // CopyIntoBatches destroys every retained destination on failure. - // Keep the derived tail state transactional with that owner cleanup. - hb.retainedSpillTailSelected = 0 - return err - } - hb.retainedSpillTailSelected = projection.nextTailSelected - return nil - } - reservation, err := hb.budget.Reserve(projection.admissionBytes) - if err != nil { - return err - } - snapshot, err := snapshotBatchCopyAllocation(hb.Batches.Buf) - if err != nil { - reservation.Release() - return err - } - if err = hb.Batches.CopyIntoBatches(src, proc); err != nil { - reservation.Release() - hb.releaseBatchReservations() - hb.retainedSpillTailSelected = 0 - return err - } - actual, err := batchCopyAllocatedDelta(hb.Batches.Buf, snapshot) - if err != nil { - hb.Batches.Clean(proc.Mp()) - hb.retainedSpillTailSelected = 0 - reservation.Release() - hb.releaseBatchReservations() - return err - } - metadata, ok := retainedMetadataAllowance(src) - if !ok || actual > math.MaxUint64-metadata { - hb.Batches.Clean(proc.Mp()) - hb.retainedSpillTailSelected = 0 - reservation.Release() - hb.releaseBatchReservations() - return process.ErrHashBuildBudgetInvalid - } - actual += metadata - if actual > projection.admissionBytes { - // This indicates an incomplete pre-allocation bound. Fail closed after - // cleaning; never legitimize the excess with post-allocation admission. - hb.Batches.Clean(proc.Mp()) - hb.retainedSpillTailSelected = 0 - reservation.Release() - hb.releaseBatchReservations() - return process.ErrHashBuildBudgetInvalid - } - if _, err = reservation.ReconcileDown(actual); err != nil { - hb.Batches.Clean(proc.Mp()) - hb.retainedSpillTailSelected = 0 - reservation.Release() - hb.releaseBatchReservations() - return err - } - hb.batchReservations = append(hb.batchReservations, reservation) - hb.retainedSpillTailSelected = projection.nextTailSelected - return nil -} - -// CopyBuildBatch is an exported compatibility wrapper. -func (hb *HashmapBuilder) CopyBuildBatch(src *batch.Batch, proc *process.Process) error { - return hb.copyBuildBatch(src, proc) -} - -func retainedMetadataAllowance(src *batch.Batch) (uint64, bool) { - if src == nil { - return 0, false - } - rows := uint64(src.RowCount()) - columns := uint64(len(src.Vecs)) - if columns > (math.MaxUint64-16)/8 { - return 0, false - } - perRow := uint64(16) + columns*8 - if rows > 0 && perRow > math.MaxUint64/rows { - return 0, false - } - return rows * perRow, true -} - -func projectedSpillSelectedAppendBytes( - src *vector.Vector, - rows int, - payloadBytes uint64, -) (uint64, error) { - if src == nil || rows < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - typeSize := src.GetType().TypeSize() - if typeSize < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - selected, err := spillCheckedMul(uint64(rows), uint64(typeSize)) - if err != nil { - return 0, err - } - return spillCheckedAdd(selected, payloadBytes) -} - -// projectedPartialTailReplacement follows UnionBatch's allocation order. The -// existing tail reservation covers old capacities. At each grow, admission -// needs the complete replacement capacity plus deltas retained by earlier -// grows. appendedSelected is the logical spill materialization contributed by -// this source range, computed from the same varlen scan used for allocation. -func projectedPartialTailReplacement( - tail, src *batch.Batch, - appendRows int, -) (peak, retained, appendedSelected uint64, err error) { - if tail == nil || src == nil || appendRows < 0 || len(tail.Vecs) != len(src.Vecs) { - return 0, 0, 0, process.ErrHashBuildBudgetInvalid + if hb.batchAllocation == nil { + return mpool.ErrAllocationAccountInvalid } - for i, srcVec := range src.Vecs { - dstVec := tail.Vecs[i] - if dstVec == nil || srcVec == nil || dstVec.Length() > math.MaxInt-appendRows { - return 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - - typeSize := srcVec.GetType().TypeSize() - requiredRows := dstVec.Length() + appendRows - if typeSize < 0 || (typeSize > 0 && requiredRows > math.MaxInt/typeSize) { - return 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - oldDataCap := cap(dstVec.GetData()) - if requiredData := requiredRows * typeSize; requiredData > oldDataCap { - newCap, ok := mpool.GrowCapacity(int64(oldDataCap), int64(requiredData)) - if !ok || retained > math.MaxUint64-uint64(newCap) { - return 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - if candidate := retained + uint64(newCap); candidate > peak { - peak = candidate - } - retained += uint64(newCap) - uint64(oldDataCap) - } - - areaBytes, selectedPayload, areaErr := - unionBatchAreaProjection(srcVec, 0, appendRows) - if areaErr != nil || areaBytes > math.MaxInt-len(dstVec.GetArea()) { - return 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - selected, selectedErr := projectedSpillSelectedAppendBytes( - srcVec, appendRows, selectedPayload) - if selectedErr != nil { - return 0, 0, 0, selectedErr - } - appendedSelected, selectedErr = spillCheckedAdd(appendedSelected, selected) - if selectedErr != nil { - return 0, 0, 0, selectedErr - } - oldAreaCap := cap(dstVec.GetArea()) - requiredArea := len(dstVec.GetArea()) + areaBytes - if requiredArea > oldAreaCap { - newCap, ok := mpool.GrowCapacity(int64(oldAreaCap), int64(requiredArea)) - if !ok || retained > math.MaxUint64-uint64(newCap) { - return 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - if candidate := retained + uint64(newCap); candidate > peak { - peak = candidate - } - retained += uint64(newCap) - uint64(oldAreaCap) - } - } - return peak, retained, appendedSelected, nil -} - -func projectedPartialTailReplacementBytes( - tail, src *batch.Batch, - appendRows int, -) (peak, retained uint64, err error) { - peak, retained, _, err = projectedPartialTailReplacement(tail, src, appendRows) - return peak, retained, err -} - -type batchCopyProjection struct { - admissionBytes uint64 - maxRetainedSelected uint64 - maxRetainedRows int - nextTailSelected uint64 - columns int + return hb.Batches.CopyIntoBatchesWithAllocation( + src, + proc, + hb.batchAllocation, + ) } -// projectedNewDestinationAllocation follows CopyIntoBatches for destinations -// that start empty. In addition to the total pre-allocation charge, it records -// the largest individual destination. HashBuild reuses that already-required -// projection to prove the future spill of every retained destination without -// adding another varlen row scan to the non-spill path. -func projectedNewDestinationAllocation( +// CopyBuildBatch is the exported boundary used by spill and integration tests. +func (hb *HashmapBuilder) CopyBuildBatch( src *batch.Batch, - start, rows int, -) (total uint64, maxRows int, maxSelected, lastSelected uint64, err error) { - if src == nil || start < 0 || rows < 0 || start > src.RowCount() || rows > src.RowCount()-start { - return 0, 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - end := start + rows - add := func(target *uint64, value uint64) error { - if *target > math.MaxUint64-value { - return process.ErrHashBuildBudgetInvalid - } - *target += value - return nil - } - for offset := start; offset < end; { - segmentRows := end - offset - if segmentRows > colexec.DefaultBatchSize { - segmentRows = colexec.DefaultBatchSize - } - var segmentAllocated uint64 - var segmentSelected uint64 - for _, vec := range src.Vecs { - if vec == nil { - return 0, 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - typeSize := vec.GetType().TypeSize() - if typeSize < 0 || (typeSize > 0 && segmentRows > math.MaxInt/typeSize) { - return 0, 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - dataCap, ok := mpool.GrowCapacity(0, int64(segmentRows*typeSize)) - if !ok || dataCap < 0 { - return 0, 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - if err = add(&segmentAllocated, uint64(dataCap)); err != nil { - return 0, 0, 0, 0, err - } - areaBytes := 0 - var selectedPayload uint64 - if vec.GetType().IsVarlen() { - var areaErr error - areaBytes, selectedPayload, areaErr = - unionBatchAreaProjection(vec, offset, segmentRows) - if areaErr != nil { - return 0, 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - areaCap, ok := mpool.GrowCapacity(0, int64(areaBytes)) - if !ok || areaCap < 0 { - return 0, 0, 0, 0, process.ErrHashBuildBudgetInvalid - } - if err = add(&segmentAllocated, uint64(areaCap)); err != nil { - return 0, 0, 0, 0, err - } - } - selected, selectedErr := projectedSpillSelectedAppendBytes( - vec, segmentRows, selectedPayload) - if selectedErr != nil { - return 0, 0, 0, 0, selectedErr - } - if err = add(&segmentSelected, selected); err != nil { - return 0, 0, 0, 0, err - } - } - if err = add(&total, segmentAllocated); err != nil { - return 0, 0, 0, 0, err - } - if segmentSelected > maxSelected { - maxSelected = segmentSelected - } - lastSelected = segmentSelected - if segmentRows > maxRows { - maxRows = segmentRows - } - offset += segmentRows - } - return total, maxRows, maxSelected, lastSelected, nil -} - -func projectedNewDestinationBytes(src *batch.Batch, start, rows int) (uint64, error) { - total, _, _, _, err := projectedNewDestinationAllocation(src, start, rows) - return total, err -} - -func (hb *HashmapBuilder) projectedBatchCopy(src *batch.Batch) (batchCopyProjection, error) { - if src == nil || src.RowCount() < 0 { - return batchCopyProjection{}, process.ErrHashBuildBudgetInvalid - } - projection := batchCopyProjection{columns: len(src.Vecs)} - rows := uint64(src.RowCount()) - last := len(hb.Batches.Buf) - 1 - hadPartialTail := last >= 0 && hb.Batches.Buf[last] != nil && - hb.Batches.Buf[last].RowCount() != colexec.DefaultBatchSize - hasPartialTail := rows != uint64(colexec.DefaultBatchSize) && hadPartialTail - appendRows := 0 - if hasPartialTail { - // CopyIntoBatches appends into the partial tail. Derive each replacement - // from the destination's old capacity and the actual old+append target. - // A flat 1.25x multiplier is not a bound: GrowCapacity can take repeated - // 1.25x steps before reaching the required size. - tail := hb.Batches.Buf[last] - if tail.RowCount() < 0 || tail.RowCount() >= colexec.DefaultBatchSize { - return batchCopyProjection{}, process.ErrHashBuildBudgetInvalid - } - appendRows = colexec.DefaultBatchSize - tail.RowCount() - if appendRows > src.RowCount() { - appendRows = src.RowCount() - } - replacementPeak, retainedDelta, appendedSelected, err := - projectedPartialTailReplacement(tail, src, appendRows) - if err != nil { - return batchCopyProjection{}, err - } - projected := replacementPeak - combinedSelected, err := spillCheckedAdd( - hb.retainedSpillTailSelected, appendedSelected) - if err != nil { - return batchCopyProjection{}, err - } - projection.maxRetainedRows = tail.RowCount() + appendRows - projection.maxRetainedSelected = combinedSelected - if projection.maxRetainedRows < colexec.DefaultBatchSize { - projection.nextTailSelected = combinedSelected - } - if appendRows < src.RowCount() { - // After the tail grow finishes, its retained delta stays live while - // CopyIntoBatches materializes the remaining source rows. - remaining, maxRows, maxSelected, lastSelected, err := projectedNewDestinationAllocation( - src, appendRows, src.RowCount()-appendRows, - ) - if err != nil { - return batchCopyProjection{}, err - } - if retainedDelta > math.MaxUint64-remaining { - return batchCopyProjection{}, process.ErrHashBuildBudgetInvalid - } - if retained := retainedDelta + remaining; retained > projected { - projected = retained - } - if maxSelected > projection.maxRetainedSelected { - projection.maxRetainedSelected = maxSelected - } - if maxRows > projection.maxRetainedRows { - projection.maxRetainedRows = maxRows - } - remainingRows := src.RowCount() - appendRows - if remainingRows%colexec.DefaultBatchSize != 0 { - projection.nextTailSelected = lastSelected - } else { - projection.nextTailSelected = 0 - } - } - projection.admissionBytes, err = projectedBatchCopyWithMetadata(src, projected) - return projection, err - } - projected, maxRows, maxSelected, lastSelected, err := projectedNewDestinationAllocation( - src, 0, src.RowCount()) - if err != nil { - return batchCopyProjection{}, err - } - projection.admissionBytes, err = projectedBatchCopyWithMetadata(src, projected) - projection.maxRetainedRows = maxRows - projection.maxRetainedSelected = maxSelected - if src.RowCount() == colexec.DefaultBatchSize && hadPartialTail { - // CopyIntoBatches swaps the new full batch before the old partial tail; - // the cached tail itself is unchanged. - projection.nextTailSelected = hb.retainedSpillTailSelected - } else if src.RowCount()%colexec.DefaultBatchSize != 0 { - projection.nextTailSelected = lastSelected - } - return projection, err -} - -func (hb *HashmapBuilder) projectedBatchCopyBytes(src *batch.Batch) (uint64, error) { - projection, err := hb.projectedBatchCopy(src) - return projection.admissionBytes, err -} - -func projectedBatchCopyWithMetadata(src *batch.Batch, projected uint64) (uint64, error) { - // Vector null bitmaps and batch/vector slice metadata live on the Go heap - // and are therefore not included in Batch.Allocated. Charge a deliberately - // conservative per-row allowance that also scales with the column count. - // The source remains caller-owned, any retained tail already has its own - // reservation, and CopyIntoBatches reconciles this reservation to the actual - // retained delta below. - metadata, ok := retainedMetadataAllowance(src) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - const batchAllocationSlack = uint64(64 << 10) - if projected > math.MaxUint64-metadata || - projected+metadata > math.MaxUint64-batchAllocationSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return projected + metadata + batchAllocationSlack, nil + proc *process.Process, +) error { + return hb.copyBuildBatch(src, proc) } func (hb *HashmapBuilder) cleanBatches(proc *process.Process) { hb.Batches.Clean(proc.Mp()) - hb.retainedSpillTailSelected = 0 - hb.releaseBatchReservations() -} - -func (hb *HashmapBuilder) buildAuxBytes( - needUniqueVec bool, -) (uint64, error) { - uniqueBytes, err := hb.uniqueJoinKeyBytes() - if err != nil { - return 0, err - } - return hb.buildAuxBytesWithUniqueProjection( - needUniqueVec, uniqueBytes) -} - -func (hb *HashmapBuilder) uniqueJoinKeyBytes() (uint64, error) { - var total uint64 - for _, vec := range hb.UniqueJoinKeys { - if vec == nil { - continue - } - allocated := vec.Allocated() - if allocated < 0 || - total > math.MaxUint64-uint64(allocated) { - return 0, process.ErrHashBuildBudgetInvalid - } - total += uint64(allocated) - } - return total, nil -} - -func (hb *HashmapBuilder) buildAuxBytesWithUniqueProjection( - needUniqueVec bool, - uniqueBytes uint64, -) (uint64, error) { - // Covers mandatory hashmap/sels scratch plus the selected runtime-filter - // key vectors' actual persistent capacities. Before their first append, a - // bounded source-relative estimate admits the optional owner; every grow is - // then preflighted against its exact mpool capacity. Retained - // build batches are already charged by batchReservations, expression results - // have their own reservations, and runtime-filter serialization is admitted - // separately. Charging multiple whole-batch copies here double-counts those - // owners and can reject a build before any auxiliary allocation occurs. - bytes := batchesAllocated(hb.Batches.Buf) - if needUniqueVec { - growthSlack := bytes / 4 - if bytes%4 != 0 { - growthSlack++ - } - if uniqueBytes > growthSlack { - growthSlack = uniqueBytes - } - if bytes > math.MaxUint64-growthSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - bytes += growthSlack - } - rowCount := hb.InputBatchRowCount - if hb.hashMapRowCountSet { - rowCount = hb.hashMapRowCount - } - rows := uint64(rowCount) - const iteratorScratch = uint64(640 << 10) - if rows > math.MaxUint64/64 || bytes > math.MaxUint64-rows*64 || bytes+rows*64 > math.MaxUint64-iteratorScratch { - return 0, process.ErrHashBuildBudgetInvalid - } - bytes += rows*64 + iteratorScratch - return bytes, nil -} - -func (hb *HashmapBuilder) reserveBuildAux(needUniqueVec bool) error { - if hb.budget == nil { - return nil - } - if hb.auxReservation != nil { - // BuildHashmap can be retried on the same retained batches with a - // different optional-runtime-filter decision. Reconcile the existing - // owner instead of silently retaining the previous projection (or, - // worse, collecting optional keys under a mandatory-only charge). - return hb.resizeBuildAuxReservation(needUniqueVec) - } - bytes, err := hb.buildAuxBytes(needUniqueVec) - if err != nil { - return err - } - token, err := hb.budget.Reserve(bytes) - if err != nil { - return err - } - hb.auxReservation = token - return nil } -// abandonOptionalRuntimeFilterKeys removes only the exact-filter owner from an -// in-progress mandatory map build. No map or input batch is replayed. The -// persistent auxiliary reservation is reconciled to the same projection used -// by a build which never requested UniqueJoinKeys. +// abandonOptionalRuntimeFilterKeys removes only the optional exact-filter +// owner from an in-progress mandatory map build. Physical vector frees are the +// single source of truth for releasing the account charge. func (hb *HashmapBuilder) abandonOptionalRuntimeFilterKeys( proc *process.Process, ) error { @@ -786,9 +143,8 @@ func (hb *HashmapBuilder) abandonOptionalRuntimeFilterKeys( } // fallbackOptionalRuntimeFilterCollection converts only a proven optional -// cause into in-place key abandonment. Fatal causes are returned unchanged, -// leaving the fallback bit untouched and builder ownership with terminal -// cleanup. +// allocation failure into in-place key abandonment. Fatal causes keep builder +// ownership with terminal cleanup. func (hb *HashmapBuilder) fallbackOptionalRuntimeFilterCollection( proc *process.Process, cause error, @@ -797,19 +153,18 @@ func (hb *HashmapBuilder) fallbackOptionalRuntimeFilterCollection( runtimefilter.OptionalFallbackNone { return cause } - if err := hb.abandonOptionalRuntimeFilterKeys(proc); err != nil { - return err - } - return nil + return hb.abandonOptionalRuntimeFilterKeys(proc) } -// releaseOptionalRuntimeFilterKeys drops terminal producer-only state without -// marking collection fallback. The JoinMap retains only the mandatory -// auxiliary projection, so its transferred budget owner must be reconciled -// before publication. +// releaseOptionalRuntimeFilterKeys drops terminal producer-only vectors. The +// backing MPool allocations release their allocation-account charges exactly +// once when each vector is freed. func (hb *HashmapBuilder) releaseOptionalRuntimeFilterKeys( proc *process.Process, ) error { + if proc == nil { + return process.ErrHashBuildBudgetInvalid + } for i := range hb.UniqueJoinKeys { if hb.UniqueJoinKeys[i] != nil { hb.UniqueJoinKeys[i].Free(proc.Mp()) @@ -817,343 +172,32 @@ func (hb *HashmapBuilder) releaseOptionalRuntimeFilterKeys( } hb.UniqueJoinKeys = nil hb.uniqueSels = nil - if hb.auxReservation == nil { - return nil - } - required, err := hb.buildAuxBytes(false) - if err != nil { - return err - } - if required > hb.auxReservation.Size() { - return process.ErrHashBuildBudgetInvalid - } - _, err = hb.auxReservation.ReconcileDown(required) - return err -} - -func (hb *HashmapBuilder) resizeBuildAuxReservation( - needUniqueVec bool, -) error { - if hb.budget == nil { - return nil - } - if hb.auxReservation == nil { - return process.ErrHashBuildBudgetInvalid - } - target, err := hb.buildAuxBytes(needUniqueVec) - if err != nil { - return err - } - current := hb.auxReservation.Size() - switch { - case current < target: - return hb.auxReservation.Grow(target - current) - case current > target: - _, err = hb.auxReservation.ReconcileDown(target) - return err - default: - return nil - } + return nil } -// prepareCanonicalRuntimeFilterCollection first resizes the mandatory -// auxiliary owner for a Dedup input after its in-place canonical rewrite. It -// then attempts the optional UniqueJoinKeys delta. Failure of only that delta -// disables the runtime filter without failing the canonical map build. +// prepareCanonicalRuntimeFilterCollection restarts optional key collection +// after a destructive Dedup rewrite. The rewritten mandatory map is already +// charged at its physical allocation sites; optional growth may fail open at +// the vector allocation boundary during the rebuild. func (hb *HashmapBuilder) prepareCanonicalRuntimeFilterCollection( requested bool, ) (bool, error) { - if err := hb.resizeBuildAuxReservation(false); err != nil { - return false, err - } if !requested { return false, nil } - if err := hb.resizeBuildAuxReservation(true); err != nil { - if runtimefilter.ClassifyOptionalFallback(err) != - runtimefilter.OptionalFallbackBudgetAdmission { - return false, err - } - hb.runtimeFilterCollectionFallback = true - return false, nil - } hb.runtimeFilterCollectionFallback = false return true, nil } -func uniqueAppendAreaBytes(src *vector.Vector, start, rows int, sels []int64) (int, error) { - if src == nil || !src.GetType().IsVarlen() { - return 0, nil - } - if start < 0 || rows < 0 || (sels == nil && (start > src.Length() || rows > src.Length()-start)) { - return 0, process.ErrHashBuildBudgetInvalid - } - values, _ := vector.MustVarlenaRawData(src) - areaBytes := 0 - for i := 0; i < rows; i++ { - idx := start + i - if sels != nil { - if i >= len(sels) || sels[i] < 0 || sels[i] >= int64(src.Length()) { - return 0, process.ErrHashBuildBudgetInvalid - } - idx = int(sels[i]) - } - if src.IsConst() { - idx = 0 - } - if idx < 0 || idx >= len(values) || - (!src.GetNulls().EmptyByFlag() && src.GetNulls().Contains(uint64(idx))) || - values[idx].IsSmall() { - continue - } - _, valueLen := values[idx].OffsetLen() - valueBytes := int(valueLen) - if areaBytes > math.MaxInt-valueBytes { - return 0, process.ErrHashBuildBudgetInvalid - } - areaBytes += valueBytes - } - return areaBytes, nil -} - -// logicalAppendAreaBytes sums the payload that UnionInt32 will materialize for -// a contiguous non-const source range. It is intentionally separate from the -// general selected-row helper above: retained-copy projection runs once per -// ingress batch, so its common no-area and no-null paths avoid per-row class, -// bitmap, and selection checks. -func logicalAppendAreaBytes(src *vector.Vector, start, rows int) (uint64, error) { - if src == nil || !src.GetType().IsVarlen() || start < 0 || rows < 0 || - start > src.Length() || rows > src.Length()-start || src.IsConst() { - return 0, process.ErrHashBuildBudgetInvalid - } - if rows == 0 || len(src.GetArea()) == 0 { - return 0, nil - } - values, _ := vector.MustVarlenaRawData(src) - end := start + rows - if end > len(values) { - return 0, process.ErrHashBuildBudgetInvalid - } - var payload uint64 - if src.GetNulls().EmptyByFlag() { - for i := start; i < end; i++ { - if values[i].IsSmall() { - continue - } - _, length := values[i].OffsetLen() - if payload > math.MaxUint64-uint64(length) { - return 0, process.ErrHashBuildBudgetInvalid - } - payload += uint64(length) - } - return payload, nil - } - for i := start; i < end; i++ { - if src.GetNulls().Contains(uint64(i)) { - continue - } - if values[i].IsSmall() { - continue - } - _, length := values[i].OffsetLen() - if payload > math.MaxUint64-uint64(length) { - return 0, process.ErrHashBuildBudgetInvalid - } - payload += uint64(length) - } - return payload, nil -} - -// unionBatchAreaProjection mirrors Vector.UnionBatch's flags=nil varlen paths -// and keeps physical retention separate from logical spill materialization. -// A whole-vector copy retains the complete source area, including stale bytes, -// while shared varlena descriptors can make a later UnionInt32 copy the same -// payload once per logical row. One live-descriptor scan therefore supplies the -// exact selected payload without treating vector class as an ownership proxy. -func unionBatchAreaProjection( - src *vector.Vector, - start, rows int, -) (physicalBytes int, selectedPayload uint64, err error) { - if src == nil || !src.GetType().IsVarlen() { - return 0, 0, nil - } - if start < 0 || rows < 0 || - start > src.Length() || rows > src.Length()-start { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - if rows == 0 { - return 0, 0, nil - } - if len(src.GetArea()) == 0 { - return 0, 0, nil - } - if src.IsConst() { - // UnionBatch materializes the constant payload once and broadcasts its - // varlena header to the appended logical rows. UnionInt32 later copies - // that referenced payload once for every selected row. - physicalBytes, err = uniqueAppendAreaBytes(src, 0, 1, nil) - if err != nil { - return 0, 0, err - } - selectedPayload, err = spillCheckedMul( - uint64(physicalBytes), uint64(rows)) - return physicalBytes, selectedPayload, err - } - if start == 0 && rows == src.Length() { - if src.VarlenaAreaIsDisjoint() { - // The retained full-vector copy preserves the complete physical area. - // Disjoint live descriptors prove their logical payload cannot exceed - // that area even when it contains dead bytes. This avoids an - // O(columns*rows) pre-spill scan for ordinary append-built vectors. - return len(src.GetArea()), uint64(len(src.GetArea())), nil - } - livePayload, err := logicalAppendAreaBytes(src, start, rows) - if err != nil { - return 0, 0, err - } - return len(src.GetArea()), livePayload, nil - } - livePayload, err := logicalAppendAreaBytes(src, start, rows) - if err != nil { - return 0, 0, err - } - if livePayload > uint64(math.MaxInt) { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - return int(livePayload), livePayload, nil -} - -// unionBatchAreaBytes is kept as the allocation-only projection used by -// focused CopyIntoBatches admission tests. -func unionBatchAreaBytes(src *vector.Vector, start, rows int) (int, error) { - physical, _, err := unionBatchAreaProjection(src, start, rows) - return physical, err -} - -func (hb *HashmapBuilder) reserveUniqueAppendOverlap(dst *vector.Vector, rows, areaBytes int) (*process.HashBuildReservation, error) { - if hb.budget == nil { - return nil, nil - } - if dst == nil || rows < 0 || areaBytes < 0 { - return nil, process.ErrHashBuildBudgetInvalid - } - typeSize := dst.GetType().TypeSize() - if typeSize < 0 || dst.Length() > math.MaxInt-rows || - (typeSize > 0 && dst.Length()+rows > math.MaxInt/typeSize) { - return nil, process.ErrHashBuildBudgetInvalid - } - requiredData := (dst.Length() + rows) * typeSize - dataCapacity, ok := mpool.GrowCapacity( - int64(cap(dst.GetData())), int64(requiredData)) - if !ok || dataCapacity < 0 { - return nil, process.ErrHashBuildBudgetInvalid - } - var overlap uint64 - if requiredData > cap(dst.GetData()) { - overlap = uint64(cap(dst.GetData())) - } - if len(dst.GetArea()) > math.MaxInt-areaBytes { - return nil, process.ErrHashBuildBudgetInvalid - } - requiredArea := len(dst.GetArea()) + areaBytes - areaCapacity, ok := mpool.GrowCapacity( - int64(cap(dst.GetArea())), int64(requiredArea)) - if !ok || areaCapacity < 0 { - return nil, process.ErrHashBuildBudgetInvalid - } - if requiredArea > cap(dst.GetArea()) { - if overlap > math.MaxUint64-uint64(cap(dst.GetArea())) { - return nil, process.ErrHashBuildBudgetInvalid - } - overlap += uint64(cap(dst.GetArea())) - } - if requiredData <= cap(dst.GetData()) && - requiredArea <= cap(dst.GetArea()) { - // The persistent capacities were admitted by their preceding grow. - // Avoid rescanning every retained batch for each UnitLimit append which - // stays within those capacities; only allocator growth changes either - // the retained owner or the temporary replacement overlap. - return nil, nil - } - - currentUnique, err := hb.uniqueJoinKeyBytes() - if err != nil { - return nil, err - } - oldCapacity := uint64(cap(dst.GetData())) + - uint64(cap(dst.GetArea())) - newCapacity := uint64(dataCapacity) + uint64(areaCapacity) - if currentUnique < oldCapacity || - currentUnique-oldCapacity > math.MaxUint64-newCapacity { - return nil, process.ErrHashBuildBudgetInvalid - } - projectedUnique := currentUnique - oldCapacity + newCapacity - target, err := hb.buildAuxBytesWithUniqueProjection( - true, projectedUnique) - if err != nil { - return nil, err - } - if hb.auxReservation == nil { - return nil, process.ErrHashBuildBudgetInvalid - } - if current := hb.auxReservation.Size(); current < target { - if err = hb.auxReservation.Grow(target - current); err != nil { - return nil, err - } - } - if overlap == 0 { - return nil, nil - } - return hb.budget.Reserve(overlap) -} - -func (hb *HashmapBuilder) marshalRuntimeFilterVector(vec *vector.Vector) ([]byte, func(), error) { - return runtimefilter.MarshalExactFilterVector(vec, hb.budget) -} - -func (hb *HashmapBuilder) releaseBatchReservations() { - for _, reservation := range hb.batchReservations { - reservation.Release() - } - hb.batchReservations = nil -} - -func (hb *HashmapBuilder) releaseReservations() { - hb.releaseMapReservation() - hb.releaseBatchReservations() - if hb.auxReservation != nil { - hb.auxReservation.Release() - hb.auxReservation = nil - } -} - -func (hb *HashmapBuilder) releaseMapReservation() { - if hb.mapReservation != nil { - hb.mapReservation.release() - hb.mapReservation = nil - } -} - -func (hb *HashmapBuilder) detachReservations() func() { - mapOwner := hb.mapReservation - hb.mapReservation = nil - reservations := make([]*process.HashBuildReservation, 0, 1+len(hb.batchReservations)) - for _, reservation := range hb.batchReservations { - if token := reservation.Transfer(); token != nil { - reservations = append(reservations, token) - } - } - hb.batchReservations = nil - if hb.auxReservation != nil { - if token := hb.auxReservation.Transfer(); token != nil { - reservations = append(reservations, token) - } - hb.auxReservation = nil - } - return func() { - mapOwner.release() - for _, reservation := range reservations { - reservation.Release() - } - } +func (hb *HashmapBuilder) marshalRuntimeFilterVector( + vec *vector.Vector, + mp *mpool.MPool, +) ([]byte, func(), error) { + return runtimefilter.MarshalExactFilterVector( + vec, + mp, + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteRuntimeFilterPayload, + ) } diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go index f68b0a69e9974..065b5a04afbb9 100644 --- a/pkg/sql/colexec/hashbuild/build.go +++ b/pkg/sql/colexec/hashbuild/build.go @@ -77,12 +77,20 @@ func (hashBuild *HashBuild) Prepare(proc *process.Process) (err error) { return TerminalBudgetError(proc.Ctx, err) } hashBuild.ctr.hashmapBuilder.setBudget(budget) + if hashBuild.ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if hashBuild.IsShuffle && hashBuild.RuntimeFilterSpec == nil { return moerr.NewInternalError(proc.Ctx, "shuffle hash build must have runtime filter") } if !hashBuild.NeedHashMap { return nil } + if hashBuild.IsShuffle { + if err = hashBuild.installRecoveryCapacity(budget); err != nil { + return TerminalBudgetError(proc.Ctx, err) + } + } hashBuild.ctr.hashmapBuilder.IsDedup = hashBuild.IsDedup hashBuild.ctr.hashmapBuilder.DedupBuildKeepLast = hashBuild.DedupBuildKeepLast @@ -126,71 +134,11 @@ func (hashBuild *HashBuild) Call(proc *process.Process) (vm.CallResult, error) { ctr.state = SendJoinMap case SendJoinMap: - ctr.terminalMu.Lock() - if hashBuild.JoinMapTag <= 0 { - ctr.terminalMu.Unlock() - err := moerr.NewInternalError(proc.Ctx, "wrong joinmap message tag!") - hashBuild.finalizeBuildFailure(proc, err) - return result, err - } - if atomic.LoadUint32(&ctr.terminalPublished) != 0 { - ctr.terminalMu.Unlock() - return result, moerr.NewQueryInterrupted(proc.Ctx) - } - - var jm *message.JoinMap - spillMode := len(ctr.spilledFds) > 0 - var spillPayloadErr error - - if ctr.hashmapBuilder.InputBatchRowCount > 0 { - if spillMode { - // In spill mode: send empty JoinMap with spill fds, no batches - jm = message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) - } else { - // Normal mode: send hashmap and batches - jm = ctr.hashmapBuilder.GetJoinMap(proc.Mp()) - jm.SetPushedRuntimeFilterIn(ctr.runtimeFilterIn) - } - jm.SetRowCount(int64(ctr.hashmapBuilder.InputBatchRowCount)) - jm.SetHasNullKey(ctr.hashmapBuilder.HasNullKey) - jm.IncRef(hashBuild.JoinMapRefCnt) - if spillMode { - payload := message.SpillBuildPayload{LegacyFds: ctr.spilledFds} - if ctr.spillBundle != nil { - payload = message.SpillBuildPayload{ - Files: ctr.spillBundle.accountedFiles(), - BudgetRef: ctr.hashmapBuilder.budget, - } - } - spillPayloadErr = jm.SetSpillBuildPayload(payload) - if spillPayloadErr == nil { - ctr.spilledFds = nil // ownership transferred - ctr.spillBundle = nil - } - } - } - - if spillPayloadErr != nil { - jm.FreeMemory() - ctr.terminalMu.Unlock() - err := moerr.NewInternalError(proc.Ctx, spillPayloadErr.Error()) + if err := hashBuild.sendJoinMap(proc); err != nil { hashBuild.finalizeBuildFailure(proc, err) return result, err } - - if !hashBuild.publishJoinMap(proc, jm) { - // Reset/Free may have won the terminal gate concurrently during - // cancellation. Keep the producer side successful only if this - // publication won; consumers must never see two terminal values. - if jm != nil { - jm.FreeMemory() - } - ctr.terminalMu.Unlock() - return result, moerr.NewQueryInterrupted(proc.Ctx) - } - ctr.state = SendSucceed - ctr.terminalMu.Unlock() case SendSucceed: result.Batch = nil @@ -200,6 +148,79 @@ func (hashBuild *HashBuild) Call(proc *process.Process) (vm.CallResult, error) { } } +// sendJoinMap serializes terminal publication with Reset and Free. The defer +// is part of the lifecycle contract: allocation, spill-payload, and message +// hooks may panic, and cleanup must never deadlock trying to reacquire this +// mutex while recovering the active statement. +func (hashBuild *HashBuild) sendJoinMap(proc *process.Process) error { + ctr := &hashBuild.ctr + ctr.terminalMu.Lock() + defer ctr.terminalMu.Unlock() + + if hashBuild.JoinMapTag <= 0 { + return moerr.NewInternalError(proc.Ctx, "wrong joinmap message tag!") + } + if hashBuild.JoinMapRefCnt <= 0 { + return moerr.NewInternalErrorf( + proc.Ctx, + "invalid join map reference count: %d", + hashBuild.JoinMapRefCnt, + ) + } + if atomic.LoadUint32(&ctr.terminalPublished) != 0 { + return moerr.NewQueryInterrupted(proc.Ctx) + } + + var jm *message.JoinMap + joinMapOwned := false + defer func() { + if joinMapOwned && jm != nil { + jm.FreeMemory() + } + }() + spillMode := len(ctr.spilledFds) > 0 + if ctr.hashmapBuilder.InputBatchRowCount > 0 { + if spillMode { + jm = message.NewJoinMap( + message.GroupSels{}, nil, nil, nil, nil, proc.Mp(), + ) + } else { + jm = ctr.hashmapBuilder.GetJoinMap(proc.Mp()) + if jm == nil { + return process.ErrHashBuildBudgetInvalid + } + joinMapOwned = true + jm.SetPushedRuntimeFilterIn(ctr.runtimeFilterIn) + } + if spillMode { + joinMapOwned = true + } + jm.SetRowCount(int64(ctr.hashmapBuilder.InputBatchRowCount)) + jm.SetHasNullKey(ctr.hashmapBuilder.HasNullKey) + jm.IncRef(hashBuild.JoinMapRefCnt) + if spillMode { + if ctr.spillBundle == nil || ctr.hashmapBuilder.budget == nil { + return process.ErrHashBuildBudgetInvalid + } + payload := message.SpillBuildPayload{ + Files: ctr.spillBundle.accountedFiles(), + BudgetRef: ctr.hashmapBuilder.budget, + } + if err := jm.SetSpillBuildPayload(payload); err != nil { + return moerr.NewInternalError(proc.Ctx, err.Error()) + } + ctr.spilledFds = nil + ctr.spillBundle = nil + } + } + + if !hashBuild.publishJoinMap(proc, jm) { + return moerr.NewQueryInterrupted(proc.Ctx) + } + joinMapOwned = false + return nil +} + // finalizeBuildFailure publishes every producer-side dependency before Call // returns. Consumers may already be blocked in ReceiveJoinMap/RuntimeFilter; // deferring publication until Reset could deadlock a pipeline scheduler that @@ -214,59 +235,25 @@ func (hashBuild *HashBuild) finalizeBuildFailure(proc *process.Process, err erro } } -func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyzer) error { +func (hashBuild *HashBuild) build( + proc *process.Process, + analyzer process.Analyzer, +) (retErr error) { ctr := &hashBuild.ctr + ctr.spillConditions = hashBuild.Conditions spillMode := false var spillFiles []*os.File bundleTransferred := false - ensureExpressionRecovery := func(rows int) error { - lease := ctr.hashmapBuilder.expressionLease - if lease == nil || lease.Len() != len(ctr.hashmapBuilder.executors) || - len(ctr.hashmapBuilder.executors) != len(hashBuild.Conditions) { - return process.ErrHashBuildBudgetInvalid - } - return lease.EnsureRunRecovery(proc, rows) - } - ensureDirectRecovery := func(bat *batch.Batch) error { - if bat == nil { + ensureRecovery := func(projection recoveryBatchProjection) error { + if !hashBuild.IsShuffle { return nil } - if err := ensureExpressionRecovery(bat.RowCount()); err != nil { - return err - } - return ctr.ensureDirectSpillRecovery(bat, analyzer) - } - ensureRetainedRecovery := func(projection batchCopyProjection) error { - if err := ensureExpressionRecovery(projection.maxRetainedRows); err != nil { - return err - } - return ctr.ensureRetainedSpillRecovery(projection, analyzer) - } - ensureDirectRecoveryWithReclaim := func(bat *batch.Batch) error { - err := ensureDirectRecovery(bat) - if !spillMode || !errors.Is(err, process.ErrHashBuildBudgetAdmission) { - return err - } - reclaimed, reclaimErr := ctr.reclaimOptionalSpillCoalesce( - proc, spillFiles, analyzer) - if reclaimErr != nil { - return reclaimErr - } - if !reclaimed { - // The rejection came from mandatory owners or sibling pressure. A - // second identical admission cannot make progress. - return err - } - return ensureDirectRecovery(bat) - } - spillBatch := func(bat *batch.Batch, sourceAlreadyCharged bool) error { - return ctr.spillBatchBounded( + return ctr.ensureRecoveryCapacity( proc, - bat, - spillFiles, + hashBuild.Conditions, + projection, analyzer, - sourceAlreadyCharged, ) } @@ -281,13 +268,20 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz ctr.spillBundle.release() ctr.spillBundle = nil } + ctr.freeSpillExprExecs() + ctr.spillConditions = nil // Build-key executors are producer scratch. No consumer reads them after - // build() returns, so release their retained vectors and expression lease - // here instead of holding both until pipeline Reset. + // build() returns, so release their retained vectors here instead of + // holding them until pipeline Reset. ctr.hashmapBuilder.FreeTemporaryVectors(proc) ctr.hashmapBuilder.FreeExecutors() ctr.dropSpillScratchBuffers() - ctr.releaseSpillScratchReservation() + if err := hashBuild.releaseRecoveryCapacity( + ctr.hashmapBuilder.mapAllocationAccount, + true, + ); err != nil { + retErr = errors.Join(retErr, err) + } }() startSpill := func() error { @@ -308,17 +302,6 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz hashBuild.JoinMapRefCnt, ) } - for _, condition := range hashBuild.Conditions { - if condition == nil { - return process.ErrHashBuildBudgetInvalid - } - } - execs := ctr.hashmapBuilder.executors - expressionLease := ctr.hashmapBuilder.expressionLease - if expressionLease == nil || expressionLease.Len() != len(execs) || - len(execs) != len(hashBuild.Conditions) { - return process.ErrHashBuildBudgetInvalid - } if spillFiles == nil { spillFiles = make([]*os.File, spillNumBuckets) } @@ -327,6 +310,14 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // Drain retained copies oldest-first. Each successful partition is // followed immediately by reservation and mpool release, so the source // batch and one partition scratch are the only simultaneous peaks. + execs := ctr.spillExprExecs + if len(ctr.hashmapBuilder.Batches.Buf) > 0 { + var err error + execs, err = ctr.initSpillExprExecs(proc, hashBuild.Conditions) + if err != nil { + return err + } + } for len(ctr.hashmapBuilder.Batches.Buf) > 0 { if err := checkHashBuildCanceled(proc); err != nil { return err @@ -338,13 +329,23 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz } continue } - if err := spillBatch(bat, true); err != nil { + if err := ctr.spillBatchWithPressure(proc, bat, spillFiles, execs, analyzer, true); err != nil { return err } if err := ctr.hashmapBuilder.CleanCopiedBatchAt(0, proc); err != nil { return err } } + // No retained state remains after the drain. Drop every mandatory + // recovery-class borrower before returning the conservative floor, then + // let direct sources use ordinary allocation-led admission. + ctr.dropMandatorySpillRecoveryScratch() + if err := hashBuild.releaseRecoveryCapacity( + ctr.hashmapBuilder.mapAllocationAccount, + true, + ); err != nil { + return err + } v2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1").Inc() return nil } @@ -352,12 +353,12 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz if err := startSpill(); err != nil { return err } - if hashBuild.IsShuffle { - if err := ensureDirectRecoveryWithReclaim(bat); err != nil { - return err - } - } - return spillBatch(bat, false) + // Recovery headroom proves that already-retained batches can be drained; + // an upstream-owned direct source cannot strand retained state. Admit its + // scratch at the physical allocation sites instead of applying the + // conservative retained-batch projection as a query-fatal gate. + return ctr.spillBatchWithPressure( + proc, bat, spillFiles, ctr.spillExprExecs, analyzer, false) } for { @@ -406,29 +407,21 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // Store original batch retainedMemBefore := ctr.hashmapBuilder.Batches.MemSize + var projection recoveryBatchProjection if hashBuild.IsShuffle { - // The retained destination may differ materially from ingress: const - // vectors become ordinary vectors and partial tails grow. Reuse the one - // unavoidable copy allocation projection to reserve the largest - // destination's future spill peak before any source row is retained. - var projection batchCopyProjection - projection, err = ctr.hashmapBuilder.projectedBatchCopy(result.Batch) - if err == nil { - err = ensureRetainedRecovery(projection) - } + projection, err = ctr.hashmapBuilder.projectRetainedRecovery(result.Batch) if err == nil { - err = ctr.hashmapBuilder.copyBuildBatchProjected(result.Batch, proc, projection) + err = ensureRecovery(projection) } - } else { + } + if err == nil { err = ctr.hashmapBuilder.copyBuildBatch(result.Batch, proc) } if err != nil { - if hashBuild.IsShuffle && errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if hashBuild.IsShuffle && IsRetryableMemoryCapacity(err) { // The source batch is still owned by the upstream operator. Do - // not retry CopyIntoBatches (or increment row count again). Every - // direct transition drains older retained copies under their existing - // guarantee, then gives mandatory recovery priority over optional - // write coalescing before writing this upstream-owned batch. + // not retry CopyIntoBatches (or increment row count again); enter + // spill recovery and write this batch directly. if err := spillDirect(result.Batch); err != nil { return err } @@ -436,6 +429,10 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz } return err } + if hashBuild.IsShuffle { + ctr.hashmapBuilder.retainedSpillTailSelected = + projection.nextTailSelected + } // Representation expansion (including const sources) or completion of a // partial retained batch can increase MemSize by more than the source's @@ -515,7 +512,7 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz "HashBuildRuntimeFilterCollectionFallbacks", 1) } if err != nil { - if !hashBuild.IsShuffle || !errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if !hashBuild.IsShuffle || !IsRetryableMemoryCapacity(err) { return err } if !rebuildSafe { @@ -526,9 +523,9 @@ func (hashBuild *HashBuild) build(proc *process.Process, analyzer process.Analyz // publishing a semantically incomplete spill payload. return err } - // Preserve the copied batches and discard only partial map state. Every - // retained destination already owns a recovery high-water lease, so a - // hard map-budget rejection cannot strand the build in memory. + // Preserve the copied batches and discard only partial map state. + // Scratch is admitted lazily while draining; failure remains a + // controlled resource error rather than an allocation past the cap. ctr.hashmapBuilder.FreeHashMapOnly(proc) if err := startSpill(); err != nil { return err @@ -787,6 +784,14 @@ func (hashBuild *HashBuild) handleRuntimeFilter( } ctr.hashmapBuilder.uniqueKeySlots = nil }() + // A spilled build has no resident unique-key vector. Treating that absence + // as an empty build would publish DROP and incorrectly discard every probe + // row, so spill always disables this optional optimization. + if len(ctr.spilledFds) > 0 { + runtimeFilter.Typ = message.RuntimeFilter_PASS + hashBuild.sendRuntimeFilter(runtimeFilter, spec, proc) + return nil + } // send the unique join keys (doc_id membership pushdown) when requested if spec.UseMembershipFilter { @@ -821,13 +826,20 @@ func (hashBuild *HashBuild) handleRuntimeFilter( return nil } rowCount := keyVec.Length() + if keyVec.GetGrouping().GetBitmap().CountRange( + 0, uint64(keyVec.Length()), + ) > 0 { + runtimeFilter.Typ = message.RuntimeFilter_PASS + hashBuild.sendRuntimeFilter(runtimeFilter, spec, proc) + return nil + } // Always send the unique join keys; the consumer (ivfflat / fulltext // search) decides whether to use them as an exact pk IN filter or to // build a membership filter, based on its own threshold. runtimeFilter.Typ = message.RuntimeFilter_UNIQUEJOINKEYS - data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector(keyVec) + data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector(keyVec, proc.Mp()) if err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -904,15 +916,6 @@ func (hashBuild *HashBuild) handleRuntimeFilter( if err := runtimefilter.CloseFloatSignedZero( keyVec, proc.Mp(), - func() (func(), error) { - overlap, err := ctr.hashmapBuilder.reserveUniqueAppendOverlap(keyVec, 1, 0) - if err != nil || overlap == nil { - return nil, err - } - return func() { - overlap.Release() - }, nil - }, ); err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -928,7 +931,10 @@ func (hashBuild *HashBuild) handleRuntimeFilter( } keyVec.GetNulls().Reset() keyVec.InplaceSort() - data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector(keyVec) + data, release, err := ctr.hashmapBuilder.marshalRuntimeFilterVector( + keyVec, + proc.Mp(), + ) if err != nil { if hashBuild.fallbackOptionalRuntimeFilter(err, &runtimeFilter, spec, proc) { return nil @@ -1023,12 +1029,11 @@ func (hashBuild *HashBuild) handleSerializedRuntimeFilter( } // materializeSerializedRuntimeFilter evaluates one proven serial/serial_full -// contract under the same query-wide HashBuild budget as the map and unique -// component vectors. It reuses the production component encoders, but +// contract under the same physical allocation account as the map and unique +// component vectors. It reuses the production component encoders and // precomputes a tight output-area bound from the actual unique values. The -// generic expression estimator must not be used here: a serial result is typed -// VARCHAR(max), which would reserve 64 KiB per tiny integer tuple and turn a -// useful index filter into PASS. +// account observes the actual vector growth rather than an estimated duplicate +// reservation. func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( proc *process.Process, spec *plan.RuntimeFilterSpec, @@ -1063,27 +1068,21 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( if err != nil { return nil, nil, 0, false, err } - peak, err := serializedRuntimeFilterAllocationPeak( - rowCount, areaBound, maxRowBound) - if err != nil { - return nil, nil, 0, false, err - } - - var reservation *process.HashBuildReservation - if budget := hashBuild.ctr.hashmapBuilder.budget; budget != nil { - reservation, err = budget.Reserve(peak) - if err != nil { - return nil, nil, 0, false, err - } - defer reservation.Release() - } payloadType, ok := planExprType( runtimefilter.BuildKeyExpr(spec)) - if !ok || areaBound > uint64(math.MaxInt) { + if !ok || areaBound > uint64(math.MaxInt) || + maxRowBound > uint64(math.MaxInt) || + hashBuild.ctr.hashmapBuilder.uniqueKeyAllocation == nil { return nil, nil, 0, false, nil } - payload := vector.NewOffHeapVecWithType(payloadType) + payload, err := vector.NewOffHeapVecWithTypeAndAllocation( + payloadType, + hashBuild.ctr.hashmapBuilder.uniqueKeyAllocation, + ) + if err != nil { + return nil, nil, 0, false, err + } defer payload.Free(proc.Mp()) if err = payload.PreExtendWithArea( rowCount, int(areaBound), proc.Mp(), @@ -1096,8 +1095,23 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( if packerSize == 0 { packerSize = 1 } - packer := types.NewPackerWithSize(packerSize) - defer packer.Close() + scratch, err := mpool.NewAccountedBuffer( + proc.Mp(), + hashBuild.ctr.hashmapBuilder.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteRuntimeFilterScratch, + ) + if err != nil { + return nil, nil, 0, false, err + } + defer scratch.Free() + if err = scratch.Resize(int(packerSize)); err != nil { + if mpool.IsRetryableAllocationCapacity(err) { + err = runtimefilter.MarkOptionalAllocationError(err) + } + return nil, nil, 0, false, err + } + packer := types.NewPackerWithFixedBuffer(scratch.Bytes()) for row := 0; row < rowCount; row++ { if row&8191 == 0 { @@ -1119,6 +1133,12 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( } encoders[i](component, row, packer) } + if err = packer.Err(); err != nil { + return nil, nil, 0, false, errors.Join( + mpool.ErrAllocationAccountInvariant, + err, + ) + } if rowIsNull { // serial is NULL if any component is NULL. NULL build keys never // match SQL equality, so omit them rather than turning a reset null @@ -1149,7 +1169,10 @@ func (hashBuild *HashBuild) materializeSerializedRuntimeFilter( } payload.InplaceSort() data, release, err = - hashBuild.ctr.hashmapBuilder.marshalRuntimeFilterVector(payload) + hashBuild.ctr.hashmapBuilder.marshalRuntimeFilterVector( + payload, + proc.Mp(), + ) if err != nil { if release != nil { release() @@ -1211,53 +1234,6 @@ func serializedRuntimeFilterBounds( return areaBytes, maxRowBytes, nil } -func serializedRuntimeFilterAllocationPeak( - rowCount int, - areaBytes uint64, - maxRowBytes uint64, -) (uint64, error) { - if rowCount < 0 || - uint64(rowCount) > math.MaxUint64/types.VarlenaSize || - areaBytes > math.MaxInt64 { - return 0, process.ErrHashBuildBudgetInvalid - } - packerRequest := maxRowBytes - if packerRequest == 0 { - packerRequest = 1 - } - packerCapacity, ok := types.PackerAllocationSize(packerRequest) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - dataBytes := uint64(rowCount) * types.VarlenaSize - if dataBytes > math.MaxInt64 { - return 0, process.ErrHashBuildBudgetInvalid - } - dataCapacity, ok := mpool.GrowCapacity(0, int64(dataBytes)) - if !ok || dataCapacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - areaCapacity, ok := mpool.GrowCapacity(0, int64(areaBytes)) - if !ok || areaCapacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - // The output vector is pre-extended, so it has no allocate-copy-free - // growth overlap. Account the packer's actual size class rather than its - // requested slice: rounding can approach another full request. - peak := uint64(dataCapacity) - for _, part := range []uint64{ - uint64(areaCapacity), - packerCapacity, - (uint64(rowCount) + 7) / 8, - } { - if peak > math.MaxUint64-part { - return 0, process.ErrHashBuildBudgetInvalid - } - peak += part - } - return peak, nil -} - // Runtime filters are optional probe-side optimizations. Fail open only for a // query/CN admission rejection or an allocation error marked at an exact // optional payload/vector boundary. Cancellation, contract violations, and @@ -1287,6 +1263,17 @@ func (hashBuild *HashBuild) fallbackOptionalRuntimeFilter( } else { stats.AddExtraStat( "HashBuildRuntimeFilterAllocationFallbacks", 1) + if account := hashBuild.ctr.hashmapBuilder.mapAllocationAccount; account != nil { + snapshot := account.Snapshot() + stats.SetMaxExtraStat( + "HashBuildRuntimeFilterBudgetFallbackUsedBytes", + hashBuildStatInt64(snapshot.Used), + ) + stats.SetMaxExtraStat( + "HashBuildRuntimeFilterBudgetFallbackCapBytes", + hashBuildStatInt64(snapshot.Limit), + ) + } } } *runtimeFilter = message.RuntimeFilterMessage{ diff --git a/pkg/sql/colexec/hashbuild/build_test.go b/pkg/sql/colexec/hashbuild/build_test.go index 724fc21e93770..a8ab28ae0148f 100644 --- a/pkg/sql/colexec/hashbuild/build_test.go +++ b/pkg/sql/colexec/hashbuild/build_test.go @@ -20,6 +20,7 @@ import ( "errors" "fmt" "math" + "os" "strings" "sync" "testing" @@ -50,46 +51,6 @@ const ( BenchmarkRows = 100000 // default rows for benchmark ) -// gatedHashBuildInput exposes a deterministic boundary after HashBuild has -// retained the first batch and before it can pull the second one. Tests use -// that boundary to model other workers charging the same query generation -// without scheduler sleeps or polling. -type gatedHashBuildInput struct { - *colexec.MockOperator - batches []*batch.Batch - current int - firstBatchRetained chan struct{} - continueInput chan struct{} -} - -func newGatedHashBuildInput(batches ...*batch.Batch) *gatedHashBuildInput { - return &gatedHashBuildInput{ - MockOperator: colexec.NewMockOperator(), - batches: batches, - firstBatchRetained: make(chan struct{}), - continueInput: make(chan struct{}), - } -} - -func (op *gatedHashBuildInput) Call(proc *process.Process) (vm.CallResult, error) { - result := vm.NewCallResult() - if op.current >= len(op.batches) { - result.Status = vm.ExecStop - return result, nil - } - if op.current == 1 { - close(op.firstBatchRetained) - select { - case <-op.continueInput: - case <-proc.Ctx.Done(): - return result, context.Cause(proc.Ctx) - } - } - result.Batch = op.batches[op.current] - op.current++ - return result, nil -} - func runtimeFilterPlanType(typ types.Type) *plan.Type { return &plan.Type{ Id: int32(typ.Oid), @@ -207,6 +168,7 @@ func TestHashBuildRepeatedResetFinalizesRuntimeFilterOnce(t *testing.T) { func TestBroadcastBudgetFailureUnblocksAllConsumers(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + installTestProcessHashBuildBudget(t, tc.arg, tc.proc) tc.arg.SetChildren([]vm.Operator{tc.marg}) require.NoError(t, tc.marg.Prepare(tc.proc)) require.NoError(t, tc.arg.Prepare(tc.proc)) @@ -271,6 +233,7 @@ func TestHashBuildPrepareConvertsTerminalBudgetAdmission(t *testing.T) { Idx: 0, }}, } + installTestProcessHashBuildBudget(t, arg, proc) var prepareErr error t.Cleanup(func() { arg.Free(proc, true, prepareErr) @@ -293,6 +256,7 @@ func TestHashBuildPrepareConvertsTerminalBudgetAdmission(t *testing.T) { func TestHashBuildWithoutMapStillBudgetsRetainedBatches(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, nil) + installTestProcessHashBuildBudget(t, tc.arg, tc.proc) tc.arg.NeedHashMap = false tc.arg.NeedBatches = true tc.arg.SetChildren([]vm.Operator{tc.marg}) @@ -321,23 +285,6 @@ func TestShuffleWithoutMapRejectsMissingRuntimeFilter(t *testing.T) { tc.arg.Free(tc.proc, true, nil) } -func TestHashBuildFreeWithoutResetReleasesOwnedMemory(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - require.NoError(t, tc.arg.Prepare(tc.proc)) - budget, err := tc.proc.GetHashBuildBudget() - require.NoError(t, err) - input := newBatch(tc.types, tc.proc, 100) - require.NoError(t, tc.arg.ctr.hashmapBuilder.copyBuildBatch(input, tc.proc)) - tc.arg.ctr.hashmapBuilder.InputBatchRowCount = input.RowCount() - input.Clean(tc.proc.Mp()) - require.NoError(t, tc.arg.ctr.hashmapBuilder.BuildHashmap(false, false, false, tc.proc)) - require.Greater(t, budget.Used(), uint64(0)) - - buildErr := errors.New("injected build failure") - tc.arg.Free(tc.proc, true, buildErr) - require.Zero(t, budget.Used()) -} - func BenchmarkBuild(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []buildTestCase{ @@ -390,25 +337,27 @@ func newTestCase(t testing.TB, flgs []bool, ts []types.Type, cs []*plan.Expr) bu proc.Reg.MergeReceivers[0] = &process.WaitRegister{ Ch2: make(chan process.PipelineSignal, 10), } + arg := &HashBuild{ + JoinMapTag: 1, + JoinMapRefCnt: 1, + Conditions: cs, + NeedHashMap: true, + OperatorBase: vm.OperatorBase{ + OperatorInfo: vm.OperatorInfo{ + Idx: 0, + IsFirst: false, + IsLast: false, + }, + }, + } + installTestHashBuildAllocation(t, arg) return buildTestCase{ types: ts, flgs: flgs, proc: proc, cancel: cancel, - arg: &HashBuild{ - JoinMapTag: 1, - JoinMapRefCnt: 1, - Conditions: cs, - NeedHashMap: true, - OperatorBase: vm.OperatorBase{ - OperatorInfo: vm.OperatorInfo{ - Idx: 0, - IsFirst: false, - IsLast: false, - }, - }, - }, - marg: &merge.Merge{}, + arg: arg, + marg: &merge.Merge{}, } } @@ -419,6 +368,7 @@ func TestHashBuildPrepareDropsPriorGenerationSpillFileService(t *testing.T) { require.NoError(t, err) arg := &HashBuild{NeedHashMap: false} + installTestHashBuildAllocation(t, arg) arg.ctr.spillFS = prior require.NoError(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.spillFS, "a reused operator must not retain the prior Process service") @@ -559,6 +509,7 @@ func TestHashBuildWithRuntimeFilter(t *testing.T) { }, }, } + installTestHashBuildAllocation(t, arg) err := arg.Prepare(proc) require.NoError(t, err) @@ -670,14 +621,13 @@ func TestHashBuildOptionalRuntimeFilterCollectionFallsBackToJoinMap( tc.arg.RuntimeFilterSpec = rawRuntimeFilterSpec( tc.arg.JoinMapTag+500, 100, typ) tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) providerCalls := 0 forcedCollectionReject := false @@ -762,14 +712,13 @@ func TestHashBuildClosedMapBudgetDoesNotRecordCollectionFallback( tc.arg.RuntimeFilterSpec = rawRuntimeFilterSpec( tc.arg.JoinMapTag+501, 100, typ) tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) providerCalls := 0 forcedClosed := false @@ -814,9 +763,6 @@ func TestHashBuildRuntimeFilterFallbackStatsTriggerDiagnostics(t *testing.T) { "HashBuildRuntimeFilterCollectionFallbacks", "HashBuildRuntimeFilterBudgetFallbacks", "HashBuildRuntimeFilterAllocationFallbacks", - "HashBuildSpillScratchReserveRejects", - "HashBuildSpillScratchGrowRejects", - "HashBuildSpillScratchGrowCount", } { t.Run(stat, func(t *testing.T) { require.True(t, hasHashBuildDiagnosticStats( @@ -836,13 +782,12 @@ func TestHashmapBuilderUniqueGrowthFailureAbandonsOptionalKeysInPlace( []types.Type{typ}, []*plan.Expr{newExpr(0, typ)}, ) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.arg.Prepare(tc.proc)) const uniqueGrowthRows = hashmap.UnitLimit * 2 input := newBatch( @@ -1042,14 +987,13 @@ func TestShuffleDedupAdmissionAfterRewriteDoesNotSpillPartialInput( Tag: tc.arg.JoinMapTag + 700, } tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - const capBytes = uint64(64 << 20) aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) generation, err := aggregate.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) // Ingress happens before BuildHashmap initializes this phase. Mark the // retained source safe so the provider rejects only after Dedup crosses // its explicit in-place rewrite boundary. @@ -1258,6 +1202,10 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. RuntimeFilterSpec: spec, } arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, arg, generation) keyVec := vector.NewOffHeapVecWithType(typ) require.NoError(t, keyVec.PreExtend(256, mp)) @@ -1265,13 +1213,6 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. arg.ctr.hashmapBuilder.InputBatchRowCount = keyVec.Length() arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{keyVec} - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - arg.ctr.hashmapBuilder.setBudget(generation) - require.NoError(t, arg.ctr.hashmapBuilder.reserveBuildAux(true)) - usedWithUniqueKeys := generation.Used() - var filler []byte defer func() { if filler != nil { @@ -1293,7 +1234,7 @@ func TestHashBuildFloatRuntimeFilterAllocationFailureFallsBackToPass(t *testing. require.False(t, arg.ctr.runtimeFilterIn) require.Nil(t, arg.ctr.hashmapBuilder.UniqueJoinKeys) require.Zero(t, generation.RejectCount()) - require.Less(t, generation.Used(), usedWithUniqueKeys) + require.Zero(t, generation.Used()) extra := arg.OpAnalyzer.GetOpStats().ExtraStats require.Equal(t, int64(1), extra["HashBuildRuntimeFilterAllocationFallbacks"]) @@ -1570,15 +1511,14 @@ func TestRuntimeFilterExplicitDecimalContractProducesIn(t *testing.T) { spec := rawRuntimeFilterSpec(105, 100, decimalType) tc.arg.RuntimeFilterSpec = spec tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, tc.arg, generation) payload := vector.NewVec(decimalType) require.NoError(t, vector.AppendFixed( payload, types.Decimal64(1000), false, tc.proc.Mp())) tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{payload} - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) require.True(t, tc.arg.ctr.runtimeFilterIn) @@ -1618,15 +1558,14 @@ func TestDirectRuntimeFilterUsesDeclaredHashSlot(t *testing.T) { spec.BuildExpr = newExpr(1, typ) tc.arg.RuntimeFilterSpec = spec tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 2 + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, tc.arg, generation) tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ testutil.MakeInt32Vector([]int32{901, 902}, nil, tc.proc.Mp()), testutil.MakeInt32Vector([]int32{11, 12}, nil, tc.proc.Mp()), } - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) receiver := message.NewMessageReceiver( []int32{spec.Tag}, @@ -1704,6 +1643,10 @@ func TestHashBuildSerializedRuntimeFilterAllocationFailureFallsBackToPass(t *tes RuntimeFilterSpec: spec, } arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, arg, generation) arg.ctr.hashmapBuilder.InputBatchRowCount = 1 arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ testutil.MakeInt32Vector([]int32{1}, nil, mp), @@ -1722,11 +1665,6 @@ func TestHashBuildSerializedRuntimeFilterAllocationFailureFallsBackToPass(t *tes } }) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - arg.ctr.hashmapBuilder.setBudget(generation) - var filler []byte defer func() { if filler != nil { @@ -1788,6 +1726,10 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { []types.Type{componentType, componentType}, true) tc.arg.RuntimeFilterSpec = spec tc.arg.ctr.hashmapBuilder.InputBatchRowCount = rowCount + budget := process.MustNewHashBuildBudget(512<<10, 512<<10) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + installTestHashBuildBudget(t, tc.arg, generation) first := make([]int32, rowCount) second := make([]int32, rowCount) @@ -1816,11 +1758,6 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { // The generic VARCHAR(max) estimator would request roughly 64 MiB for // these tiny tuples. The tuple-specific bound must fit comfortably here. - budget := process.MustNewHashBuildBudget(512<<10, 512<<10) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) require.True(t, tc.arg.ctr.runtimeFilterIn) @@ -1868,74 +1805,45 @@ func TestSerializedRuntimeFilterUsesTightBudgetAndProducesIn(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestSerializedRuntimeFilterBudgetAccountsPackerSizeClass(t *testing.T) { +func TestSerializedRuntimeFilterScratchUsesPhysicalAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) componentType := types.T_varchar.ToType() - tc := newTestCase( - t, - []bool{false}, - []types.Type{componentType}, - []*plan.Expr{newExpr(0, componentType)}, - ) spec := makeSerializedRuntimeFilterSpec( - t, tc.proc, 107, 2, []types.Type{componentType}, false) - tc.arg.RuntimeFilterSpec = spec - tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 - value := strings.Repeat("x", 128<<10) - keys := []*vector.Vector{ - testutil.MakeVarcharVector([]string{value}, nil, tc.proc.Mp()), - } - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = keys - - service := tc.proc.GetService() - rt := moruntime.ServiceRuntime(service) - original, hadOriginal := rt.GetGlobalVariables(moruntime.MOProtocolVersion) - rt.SetGlobalVariables( - moruntime.MOProtocolVersion, defines.MORPCVersion8) - t.Cleanup(func() { - if hadOriginal { - rt.SetGlobalVariables(moruntime.MOProtocolVersion, original) - } else { - rt.SetGlobalVariables( - moruntime.MOProtocolVersion, defines.MORPCLatestVersion) - } - }) - - areaBytes, maxRowBytes, err := - serializedRuntimeFilterBounds( - tc.proc, keys, []int{0}, 1, false) - require.NoError(t, err) - require.Greater(t, maxRowBytes, uint64(128<<10)) - packerBytes, ok := types.PackerAllocationSize(maxRowBytes) - require.True(t, ok) - require.Equal(t, uint64(256<<10), packerBytes) - peak, err := serializedRuntimeFilterAllocationPeak( - 1, areaBytes, maxRowBytes) - require.NoError(t, err) - - // One byte below the true peak must fail open before constructing the - // packer. Counting only its requested slice would incorrectly admit it. - budget := process.MustNewHashBuildBudget(peak-1, peak-1) + t, proc, 107, 2, []types.Type{componentType}, false) + arg := &HashBuild{RuntimeFilterSpec: spec} + budget := process.MustNewHashBuildBudget(450<<10, 450<<10) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, arg, generation) - require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) - receiver := message.NewMessageReceiver( - []int32{spec.Tag}, - message.AddrBroadCastOnCurrentCN(), - tc.proc.GetMessageBoard(), - ) - msgs, _, err := receiver.ReceiveMessage(false, tc.proc.Ctx) - require.NoError(t, err) - require.Len(t, msgs, 1) - require.Equal(t, int32(message.RuntimeFilter_PASS), - msgs[0].(message.RuntimeFilterMessage).Typ) + value := strings.Repeat("x", 300<<10) + arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()), + } + arg.ctr.hashmapBuilder.InputBatchRowCount = 1 + + data, release, rows, usable, err := + arg.materializeSerializedRuntimeFilter( + proc, + spec, + []types.Type{componentType}, + 1, + ) + require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) + require.Equal(t, runtimefilter.OptionalFallbackBudgetAdmission, + runtimefilter.ClassifyOptionalFallback(err)) + require.Nil(t, data) + require.Nil(t, release) + require.Zero(t, rows) + require.False(t, usable) + require.Equal(t, uint64(1), generation.RejectCount()) require.Zero(t, generation.Used()) - require.Zero(t, generation.Peak()) + require.NotZero(t, generation.Peak()) + arg.ctr.hashmapBuilder.Free(proc) generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) } func TestSerializedRuntimeFilterBoundsObserveCancellation(t *testing.T) { @@ -2106,14 +2014,13 @@ func TestRuntimeFilterMarshalBudgetAdmissionFallsBackToPass(t *testing.T) { tc.arg.RuntimeFilterSpec = spec tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ - testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()), - } - budget := process.MustNewHashBuildBudget(1, 1) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()), + } require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) require.True(t, tc.arg.ctr.runtimeFilterDone) @@ -2152,20 +2059,16 @@ func TestRuntimeFilterMarshalUsesSinglePayloadBudget(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()) - wireBytes := uint64(1+len(types.EncodeType(vec.GetType()))+4*4+1) + - uint64(len(vec.GetData())+len(vec.GetArea())) - projected := wireBytes + 64<<10 - - budget := process.MustNewHashBuildBudget(projected, projected) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) - data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec) + data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) require.NotEmpty(t, data) - require.Equal(t, projected, generation.Peak()) - require.LessOrEqual(t, generation.Used(), projected) + require.Equal(t, uint64(cap(data)), generation.Peak()) + require.Equal(t, uint64(cap(data)), generation.Used()) require.NotNil(t, release) release() require.Zero(t, generation.Used()) @@ -2184,20 +2087,16 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { values[i] = strings.Repeat("x", 1024+i) } vec := testutil.MakeVarcharVector(values, nil, tc.proc.Mp()) - wireBytes := uint64(1+len(types.EncodeType(vec.GetType()))+4*4+1) + - uint64(len(vec.GetData())+len(vec.GetArea())) - projected := wireBytes + 64<<10 - - budget := process.MustNewHashBuildBudget(projected, projected) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) + installTestHashBuildBudget(t, tc.arg, generation) - data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec) + data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector(vec, tc.proc.Mp()) require.NoError(t, err) require.NotEmpty(t, data) - require.Equal(t, projected, generation.Peak()) - require.LessOrEqual(t, generation.Used(), projected) + require.Equal(t, uint64(cap(data)), generation.Peak()) + require.Equal(t, uint64(cap(data)), generation.Used()) release() require.Zero(t, generation.Used()) @@ -2207,76 +2106,221 @@ func TestRuntimeFilterMarshalSinglePayloadCoversVarlenaPeak(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestRuntimeFilterMarshalClosedBudgetRemainsFatal(t *testing.T) { +func TestRuntimeFilterMarshalAccountedPayloadMessageLifecycle(t *testing.T) { + tc := newTestCase(t, []bool{true}, []types.Type{types.T_varchar.ToType()}, + []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) + vec := testutil.MakeVarcharVector( + []string{strings.Repeat("x", 4<<10), strings.Repeat("y", 8<<10)}, + nil, + tc.proc.Mp(), + ) + + const limit = uint64(1 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(2*limit, generation) + require.NoError(t, err) + tc.arg.NeedHashMap = true + replaceTestHashBuildAllocation(t, tc.arg, account) + tc.arg.ctr.hashmapBuilder.setBudget(generation) + + data, release, err := tc.arg.ctr.hashmapBuilder.marshalRuntimeFilterVector( + vec, + tc.proc.Mp(), + ) + require.NoError(t, err) + require.NotEmpty(t, data) + require.NotNil(t, release) + snapshot := account.Snapshot() + require.Positive(t, snapshot.Used) + require.Equal(t, snapshot.Used, generation.Used()) + + spec := &plan.RuntimeFilterSpec{Tag: 103} + runtimeFilter := message.RuntimeFilterMessage{ + Tag: spec.Tag, + Typ: message.RuntimeFilter_IN, + Card: 2, + Data: data, + } + runtimeFilter.SetMemoryRelease(release) + message.SendRuntimeFilter(runtimeFilter, spec, tc.proc.GetMessageBoard()) + require.True(t, tc.proc.GetMessageBoard().CloseAndDrain()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + vec.Free(tc.proc.Mp()) + + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + generation.Close() + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + +func TestRuntimeFilterMarshalAccountedOneByteShortFallsBackToPass(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - spec := rawRuntimeFilterSpec(102, 100, types.T_int32.ToType()) - tc.arg.RuntimeFilterSpec = spec - tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") - tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 - tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ - testutil.MakeInt32Vector([]int32{1}, nil, tc.proc.Mp()), - } + vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()) + size, err := vec.MarshalBinarySize() + require.NoError(t, err) + capacity, ok := mpool.GrowCapacity(0, int64(size)) + require.True(t, ok) + require.Positive(t, capacity) + vec.Free(tc.proc.Mp()) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + limit := uint64(capacity - 1) + budget := process.MustNewHashBuildBudget(limit, limit) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - generation.Close() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(2*limit, generation) + require.NoError(t, err) + tc.arg.NeedHashMap = true + replaceTestHashBuildAllocation(t, tc.arg, account) tc.arg.ctr.hashmapBuilder.setBudget(generation) + tc.arg.RuntimeFilterSpec = rawRuntimeFilterSpec( + 104, 100, types.T_int32.ToType()) + tc.arg.OpAnalyzer = process.NewAnalyzer(0, false, false, "hash build") + tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 4 + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()), + } - err = tc.arg.handleRuntimeFilter(tc.proc) - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.Nil(t, tc.arg.ctr.hashmapBuilder.UniqueJoinKeys) - require.False(t, tc.arg.ctr.runtimeFilterDone) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildRuntimeFilterBudgetFallbacks"]) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildRuntimeFilterAllocationFallbacks"]) + require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Equal(t, int64(1), + tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildRuntimeFilterBudgetFallbacks"]) receiver := message.NewMessageReceiver( - []int32{spec.Tag}, message.AddrBroadCastOnCurrentCN(), tc.proc.GetMessageBoard()) - msgs, done, receiveErr := receiver.ReceiveMessage(false, tc.proc.Ctx) - require.NoError(t, receiveErr) + []int32{tc.arg.RuntimeFilterSpec.Tag}, + message.AddrBroadCastOnCurrentCN(), + tc.proc.GetMessageBoard(), + ) + msgs, done, err := receiver.ReceiveMessage(false, tc.proc.Ctx) + require.NoError(t, err) require.False(t, done) - require.Empty(t, msgs) + require.Len(t, msgs, 1) + runtimeFilter, ok := msgs[0].(message.RuntimeFilterMessage) + require.True(t, ok) + require.Equal(t, int32(message.RuntimeFilter_PASS), runtimeFilter.Typ) + require.Empty(t, runtimeFilter.Data) + require.True(t, tc.proc.GetMessageBoard().CloseAndDrain()) + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + generation.Close() tc.proc.Free() require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestHashBuildMultipleTypes(t *testing.T) { - tests := []struct { - name string - typ types.Type - }{ - {"int16", types.T_int16.ToType()}, - {"int32", types.T_int32.ToType()}, - {"int64", types.T_int64.ToType()}, - {"uint8", types.T_uint8.ToType()}, - {"varchar", types.T_varchar.ToType()}, +func TestRuntimeFilterWithGroupingKeyFallsBackToPass(t *testing.T) { + tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: 105, + UpperLimit: 100, + UseMembershipFilter: true, } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{tt.typ}, []*plan.Expr{newExpr(0, tt.typ)}) - err := tc.marg.Prepare(tc.proc) - require.NoError(t, err) - err = tc.arg.Prepare(tc.proc) - require.NoError(t, err) - tc.arg.SetChildren([]vm.Operator{tc.marg}) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(newBatch(tc.types, tc.proc, Rows), nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(batch.EmptyBatch, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - ok, err := vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - require.Equal(t, vm.ExecStop, ok.Status) - tc.arg.Free(tc.proc, false, nil) - tc.proc.Free() - }) + tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 + tc.arg.ctr.hashmapBuilder.UniqueJoinKeys = []*vector.Vector{ + vector.NewRollupConst(types.T_int32.ToType(), 1, tc.proc.Mp()), } -} -func TestHashBuildNullable(t *testing.T) { - tc := newTestCase(t, []bool{true}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - err := tc.marg.Prepare(tc.proc) + require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) + receiver := message.NewMessageReceiver( + []int32{tc.arg.RuntimeFilterSpec.Tag}, + message.AddrBroadCastOnCurrentCN(), + tc.proc.GetMessageBoard(), + ) + msgs, done, err := receiver.ReceiveMessage(false, tc.proc.Ctx) + require.NoError(t, err) + require.False(t, done) + require.Len(t, msgs, 1) + runtimeFilter, ok := msgs[0].(message.RuntimeFilterMessage) + require.True(t, ok) + require.Equal(t, int32(message.RuntimeFilter_PASS), runtimeFilter.Typ) + require.Empty(t, runtimeFilter.Data) + require.Nil(t, tc.arg.ctr.hashmapBuilder.UniqueJoinKeys) + + tc.arg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + +func TestSpilledBuildRuntimeFilterPassesInsteadOfDropping(t *testing.T) { + tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: 106, + UpperLimit: 100, + Expr: newExpr(0, types.T_int32.ToType()), + } + tc.arg.ctr.hashmapBuilder.InputBatchRowCount = 1 + file, err := os.CreateTemp(t.TempDir(), "hashbuild-spilled-runtime-filter") + require.NoError(t, err) + tc.arg.ctr.spilledFds = []*os.File{file} + + require.NoError(t, tc.arg.handleRuntimeFilter(tc.proc)) + receiver := message.NewMessageReceiver( + []int32{tc.arg.RuntimeFilterSpec.Tag}, + message.AddrBroadCastOnCurrentCN(), + tc.proc.GetMessageBoard(), + ) + msgs, done, err := receiver.ReceiveMessage(false, tc.proc.Ctx) + require.NoError(t, err) + require.False(t, done) + require.Len(t, msgs, 1) + runtimeFilter, ok := msgs[0].(message.RuntimeFilterMessage) + require.True(t, ok) + require.Equal(t, int32(message.RuntimeFilter_PASS), runtimeFilter.Typ) + require.Empty(t, runtimeFilter.Data) + + tc.arg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + +func TestHashBuildMultipleTypes(t *testing.T) { + tests := []struct { + name string + typ types.Type + }{ + {"int16", types.T_int16.ToType()}, + {"int32", types.T_int32.ToType()}, + {"int64", types.T_int64.ToType()}, + {"uint8", types.T_uint8.ToType()}, + {"varchar", types.T_varchar.ToType()}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tc := newTestCase(t, []bool{false}, []types.Type{tt.typ}, []*plan.Expr{newExpr(0, tt.typ)}) + err := tc.marg.Prepare(tc.proc) + require.NoError(t, err) + err = tc.arg.Prepare(tc.proc) + require.NoError(t, err) + tc.arg.SetChildren([]vm.Operator{tc.marg}) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(newBatch(tc.types, tc.proc, Rows), nil, tc.proc.Mp()) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(batch.EmptyBatch, nil, tc.proc.Mp()) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) + ok, err := vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + require.Equal(t, vm.ExecStop, ok.Status) + tc.arg.Free(tc.proc, false, nil) + tc.proc.Free() + }) + } +} + +func TestHashBuildNullable(t *testing.T) { + tc := newTestCase(t, []bool{true}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) + err := tc.marg.Prepare(tc.proc) require.NoError(t, err) err = tc.arg.Prepare(tc.proc) require.NoError(t, err) @@ -2353,6 +2397,7 @@ func TestHashBuildRuntimeFilterWithNulls(t *testing.T) { }, }, } + installTestHashBuildAllocation(t, arg) err := arg.Prepare(proc) require.NoError(t, err) @@ -2405,6 +2450,7 @@ func TestHashBuildRuntimeFilterWithNullsHashOnPK(t *testing.T) { }, }, } + installTestHashBuildAllocation(t, arg) err := arg.Prepare(proc) require.NoError(t, err) @@ -2429,55 +2475,6 @@ func TestHashBuildRuntimeFilterWithNullsHashOnPK(t *testing.T) { proc.Free() } -func TestHashBuildIsShuffle(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - budget, budgetErr := tc.proc.GetHashBuildBudget() - require.NoError(t, budgetErr) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.TrackNullKeys = true - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: 2} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - for cycle := 0; cycle < 2; cycle++ { - if cycle > 0 { - tc.marg.Reset(tc.proc, false, nil) - tc.proc.GetMessageBoard().Reset() - } - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - build := batch.NewWithSize(1) - var buildNulls []uint64 - if cycle == 0 { - buildNulls = []uint64{1} - } - build.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 0, 2}, buildNulls, tc.proc.Mp()) - build.SetRowCount(3) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(batch.EmptyBatch, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err := vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - result, receiveErr := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, receiveErr) - require.True(t, result.IsSuccess(), "cycle %d must publish a spilled JoinMap", cycle) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.Len(t, spillPayload.Files, spillNumBuckets) - require.Same(t, budget, spillPayload.BudgetRef) - require.NoError(t, spillPayload.Close()) - require.Zero(t, budget.Used()) - require.Zero(t, budget.SpillDiskUsed()) - require.Zero(t, budget.SpillFDUsed()) - tc.arg.Reset(tc.proc, false, nil) - } - tc.arg.Free(tc.proc, false, nil) - tc.proc.Free() -} - func TestBroadcastHashBuildParallelConsumersStayResident(t *testing.T) { tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) tc.arg.IsShuffle = false @@ -2559,905 +2556,224 @@ func TestHashBuildRejectsSharedSpillPayload(t *testing.T) { tc.proc.Free() } -func TestShuffleHashBuildRecoveryLeaseDoesNotForceSpill(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_varchar.ToType()}, []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3500} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(8 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - payload := make([]byte, 1<<20) - for i := range payload { - payload[i] = 'x' - } - build := batch.NewWithSize(1) - build.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), payload, 1, tc.proc.Mp()) - require.NoError(t, err) - build.SetRowCount(1) - - directNeed, err := spillBudgetBytes(build) - require.NoError(t, err) - require.Less(t, directNeed, capBytes) +func TestHashBuildRejectsNonPositiveJoinMapRefCountBeforeTransfer(t *testing.T) { + for _, refCount := range []int32{0, -1} { + t.Run(fmt.Sprintf("ref-%d", refCount), func(t *testing.T) { + tc := newTestCase( + t, + []bool{false}, + []types.Type{types.T_int32.ToType()}, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + ) + tc.arg.JoinMapRefCnt = refCount + tc.arg.SpillThreshold = math.MaxInt64 + tc.arg.SetChildren([]vm.Operator{tc.marg}) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) + account := tc.arg.ctr.hashmapBuilder.mapAllocationAccount - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err = vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) + build := batch.NewWithSize(1) + build.Vecs[0] = testutil.MakeInt32Vector( + []int32{1, 2, 3}, nil, tc.proc.Mp(), + ) + build.SetRowCount(3) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.False(t, jm.IsSpilled(), - "a recovery lease changes admission ownership, not the local spill policy") - require.Equal(t, int64(1), jm.GetRowCount()) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) - require.Positive(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillRecoveryReservedBytes"]) - jm.Free() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) + _, buildErr := vm.Exec(tc.arg, tc.proc) + require.ErrorContains(t, buildErr, "invalid join map reference count") + result, err := message.ReceiveJoinMapResult( + tc.arg.JoinMapTag, + false, + 0, + tc.proc.GetMessageBoard(), + tc.proc.Ctx, + ) + require.NoError(t, err) + require.True(t, result.IsBuildError()) + require.Nil(t, result.JoinMap()) - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) + tc.arg.Reset(tc.proc, true, buildErr) + tc.marg.Reset(tc.proc, true, buildErr) + require.Zero(t, account.Snapshot().Used) + tc.arg.Free(tc.proc, true, buildErr) + tc.proc.Free() + }) + } } -func TestShuffleHashBuildSpillsBeforeRetainingThresholdCrossingBatch(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_varchar.ToType()}, []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) +func TestShuffleHashBuildAccountedSpillLifecycle(t *testing.T) { + tc := newTestCase( + t, + []bool{false}, + []types.Type{types.T_int64.ToType()}, + []*plan.Expr{newExpr(0, types.T_int64.ToType())}, + ) tc.arg.IsShuffle = true tc.arg.ShuffleIdx = 0 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3502} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - makeBuildBatch := func() *batch.Batch { - values := make([]string, colexec.DefaultBatchSize) - for i := range values { - values[i] = strings.Repeat("x", 256) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, tc.proc.Mp()) - bat.SetRowCount(len(values)) - return bat + tc.arg.SpillThreshold = 1 + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: tc.arg.JoinMapTag + 4_500, } - first := makeBuildBatch() - second := makeBuildBatch() - inputSize := int64(first.Size()) - tc.arg.SpillThreshold = inputSize + 1 - tc.arg.ctr.setSpillThreshold(tc.arg.SpillThreshold) - - projection, err := tc.arg.ctr.hashmapBuilder.projectedBatchCopy(first) - require.NoError(t, err) - directRecovery, err := spillBudgetBytes(first) - require.NoError(t, err) - directRecovery, err = spillRecoveryReservationBytes(directRecovery) - require.NoError(t, err) - retainedRecovery, err := spillRetainedRecoveryBudgetBytes(projection) + tc.arg.SetChildren([]vm.Operator{tc.marg}) + const limit = uint64(8 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) require.NoError(t, err) - retainedRecovery, err = spillRecoveryReservationBytes(retainedRecovery) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) require.NoError(t, err) - - // Admit the first recovery lease and its retained-copy allocation together. - // The second batch crosses the local threshold and must be spilled directly, - // without a second retained copy or any failed shared-budget admission. - coalesceSlack := uint64(spillNumBuckets * spillWriteCoalesceSize) - capBytes := max(directRecovery, retainedRecovery) + - projection.admissionBytes + coalesceSlack - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) + account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) + replaceTestHashBuildAllocation(t, tc.arg, account) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) tc.arg.ctr.hashmapBuilder.setBudget(generation) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(first, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(second, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.NoError(t, buildErr) - + build := newBatch(tc.types, tc.proc, colexec.DefaultBatchSize) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + build, + nil, + tc.proc.Mp(), + ) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + nil, + nil, + tc.proc.Mp(), + ) + _, err = vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) + tc.arg.JoinMapTag, + true, + tc.arg.ShuffleIdx, + tc.proc.GetMessageBoard(), + tc.proc.Ctx, + ) require.NoError(t, err) require.True(t, result.IsSuccess()) jm := result.JoinMap() require.NotNil(t, jm) require.True(t, jm.IsSpilled()) - require.Equal(t, int64(2*colexec.DefaultBatchSize), jm.GetRowCount()) + require.Equal(t, int64(colexec.DefaultBatchSize), jm.GetRowCount()) payload, err := jm.TakeSpillBuildPayload() require.NoError(t, err) require.NoError(t, payload.Close()) - - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"]) - require.Zero(t, extra["HashBuildSpillScratchReserveRejects"]) - require.Zero(t, extra["QueryHashBudgetRejects"], - "pre-copy thresholding must not consume recovery headroom first") - require.Empty(t, tc.arg.ctr.hashmapBuilder.Batches.Buf) + require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) require.Zero(t, generation.SpillDiskUsed()) require.Zero(t, generation.SpillFDUsed()) tc.arg.Reset(tc.proc, false, nil) tc.marg.Reset(tc.proc, false, nil) + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) tc.arg.Free(tc.proc, false, nil) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - generation.Close() tc.proc.Free() require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestShuffleHashBuildPreservesRecoveryHeadroomAcrossSharedBudgetPressure(t *testing.T) { - typ := types.T_int32.ToType() - tc := newTestCase(t, []bool{false}, []types.Type{typ}, []*plan.Expr{newExpr(0, typ)}) +func TestShuffleHashBuildDirectSpillUsesActualAllocation(t *testing.T) { + tc := newTestCase( + t, + []bool{false, false}, + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + nil, + ) + tc.arg.Conditions = []*plan.Expr{makeIssue26454ConcatKey(t, tc.proc)} tc.arg.IsShuffle = true tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 2 * colexec.DefaultBatchSize - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3503} - - makeBuildBatch := func(base int32) *batch.Batch { - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = base + int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, tc.proc.Mp()) - bat.SetRowCount(len(values)) - return bat + tc.arg.SpillThreshold = 1 + tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: tc.arg.JoinMapTag + 4_501, } - first := makeBuildBatch(0) - second := makeBuildBatch(colexec.DefaultBatchSize) - input := newGatedHashBuildInput(first, second) - tc.arg.SetChildren([]vm.Operator{input}) - require.NoError(t, input.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) + tc.arg.SetChildren([]vm.Operator{tc.marg}) - directNeed, err := spillBudgetBytes(first) - require.NoError(t, err) - retainedNeed, err := spillScratchBudgetBytes(first, true) + const ( + limit = uint64(8 << 20) + rows = 128 + ) + estimated, err := expressionRecoveryBytes( + tc.proc, + tc.arg.Conditions, + rows, + false, + ) require.NoError(t, err) - recoveryNeed := max(directNeed, retainedNeed) - require.Positive(t, recoveryNeed) + require.Greater(t, estimated, limit, + "fixture must exceed the conservative retained-recovery projection") - const capBytes = uint64(64 << 20) - require.Less(t, recoveryNeed, capBytes/2) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) + budget := process.MustNewHashBuildBudget(limit, limit) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - type execOutcome struct { - result vm.CallResult - err error - } - execDone := make(chan execOutcome, 1) - go func() { - result, execErr := vm.Exec(tc.arg, tc.proc) - execDone <- execOutcome{result: result, err: execErr} - }() - - select { - case <-input.firstBatchRetained: - case outcome := <-execDone: - t.Fatalf("HashBuild stopped before retaining the first batch: result=%+v err=%v", outcome.result, outcome.err) - } - continued := false - defer func() { - if !continued { - close(input.continueInput) - } - }() - - usedAfterFirst := generation.Used() - require.Less(t, usedAfterFirst+recoveryNeed, capBytes) - // Model sibling HashBuild workers consuming every byte except one less than - // this worker needs to recover. A valid recovery lease is already charged - // and remains usable; a lazy design has no way to start spilling here. - blockerBytes := capBytes - usedAfterFirst - (retainedNeed - 1) - blocker, err := generation.Reserve(blockerBytes) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + replaceTestHashBuildAllocation(t, tc.arg, account) + require.NoError(t, tc.marg.Prepare(tc.proc)) + require.NoError(t, tc.arg.Prepare(tc.proc)) - close(input.continueInput) - continued = true - outcome := <-execDone - require.NoError(t, outcome.err) - require.Equal(t, vm.ExecStop, outcome.result.Status) + build := newBatch(tc.types, tc.proc, rows) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + build, + nil, + tc.proc.Mp(), + ) + tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly( + nil, + nil, + tc.proc.Mp(), + ) + _, err = vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + require.Zero(t, generation.RejectCount(), + "an unretained direct source must not hit the recovery projection gate") + require.Less(t, generation.Peak(), limit) result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) + tc.arg.JoinMapTag, + true, + tc.arg.ShuffleIdx, + tc.proc.GetMessageBoard(), + tc.proc.Ctx, + ) require.NoError(t, err) require.True(t, result.IsSuccess()) jm := result.JoinMap() require.NotNil(t, jm) require.True(t, jm.IsSpilled()) - require.Equal(t, int64(2*colexec.DefaultBatchSize), jm.GetRowCount()) + require.Equal(t, int64(rows), jm.GetRowCount()) payload, err := jm.TakeSpillBuildPayload() require.NoError(t, err) require.NoError(t, payload.Close()) - - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"]) - require.Zero(t, extra["HashBuildSpillScratchReserveRejects"]) - require.Zero(t, extra["HashBuildSpillScratchGrowRejects"]) - - require.True(t, blocker.Release()) + require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) + tc.arg.Reset(tc.proc, false, nil) + tc.marg.Reset(tc.proc, false, nil) + require.NoError(t, tc.arg.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) tc.arg.Free(tc.proc, false, nil) - input.Reset(tc.proc, false, nil) - input.Free(tc.proc, false, nil) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - generation.Close() tc.proc.Free() require.Zero(t, tc.proc.Mp().CurrNB()) } -func TestShuffleHashBuildSerialFullRecoverySurvivesSharedBudgetPressure(t *testing.T) { - bindProc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - left := newExpr(0, types.T_int32.ToType()) - right := newExpr(1, types.T_int32.ToType()) - serialFull, err := plan2.BindFuncExprImplByPlanExpr( - bindProc.Ctx, - "serial_full", - []*plan.Expr{left, right}, - ) +func TestObserveHashBuildBudgetUsesGenerationSnapshot(t *testing.T) { + budget := process.MustNewHashBuildBudget(1024, 1024) + generation, err := budget.OpenGeneration(1) require.NoError(t, err) - bindProc.Free() - - typ := types.T_int32.ToType() - tc := newTestCase( - t, - []bool{false, false}, - []types.Type{typ, typ}, - []*plan.Expr{serialFull}, - ) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 2 * colexec.DefaultBatchSize - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3505} - - makeBuildBatch := func(base int32) *batch.Batch { - leftValues := make([]int32, colexec.DefaultBatchSize) - rightValues := make([]int32, colexec.DefaultBatchSize) - for i := range leftValues { - leftValues[i] = base + int32(i) - rightValues[i] = int32(i%97 + 1) - } - bat := batch.NewWithSize(2) - bat.Vecs[0] = testutil.MakeInt32Vector(leftValues, nil, tc.proc.Mp()) - bat.Vecs[1] = testutil.MakeInt32Vector(rightValues, nil, tc.proc.Mp()) - bat.SetRowCount(len(leftValues)) - return bat - } - first := makeBuildBatch(0) - second := makeBuildBatch(colexec.DefaultBatchSize) - input := newGatedHashBuildInput(first, second) - tc.arg.SetChildren([]vm.Operator{input}) - - // Initialize the Process-owned generation before Prepare so both the build - // expression lease and every later recovery owner charge the same ledger. - const capBytes = uint64(64 << 20) - tc.proc.Base.Lim.Size = int64(capBytes) - generation, err := tc.proc.GetHashBuildBudget() + registry, err := mpool.NewAllocationAccountRegistry(1, 4) require.NoError(t, err) - require.NoError(t, input.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - type execOutcome struct { - result vm.CallResult - err error - } - execDone := make(chan execOutcome, 1) - go func() { - result, execErr := vm.Exec(tc.arg, tc.proc) - execDone <- execOutcome{result: result, err: execErr} - }() - - select { - case <-input.firstBatchRetained: - case outcome := <-execDone: - t.Fatalf("HashBuild stopped before retaining the first batch: result=%+v err=%v", outcome.result, outcome.err) - } - continued := false - defer func() { - if !continued { - close(input.continueInput) - } - }() - - usedAfterFirst := generation.Used() - require.Less(t, usedAfterFirst, capBytes) - blocker, err := generation.Reserve(capBytes - usedAfterFirst) - require.NoError(t, err) - require.Equal(t, capBytes, generation.Used()) - - close(input.continueInput) - continued = true - outcome := <-execDone - require.NoError(t, outcome.err) - require.Equal(t, vm.ExecStop, outcome.result.Status) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(2*colexec.DefaultBatchSize), jm.GetRowCount()) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, payload.Close()) - require.Equal(t, int64(1), - tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) - - require.True(t, blocker.Release()) - require.Zero(t, generation.Used()) - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - input.Reset(tc.proc, false, nil) - input.Free(tc.proc, false, nil) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildRecoveryLeaseSharedGenerationFanout(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - defer bat.Clean(proc.Mp()) - - var projectionBuilder HashmapBuilder - projection, err := projectionBuilder.projectedBatchCopy(bat) - require.NoError(t, err) - directNeed, err := spillBudgetBytes(bat) - require.NoError(t, err) - directNeed, err = spillRecoveryReservationBytes(directNeed) - require.NoError(t, err) - retainedNeed, err := spillRetainedRecoveryBudgetBytes(projection) - require.NoError(t, err) - retainedNeed, err = spillRecoveryReservationBytes(retainedNeed) - require.NoError(t, err) - perWorker := max(directNeed, retainedNeed) - - for _, workers := range []int{1, 8, 64} { - t.Run(fmt.Sprintf("workers-%d", workers), func(t *testing.T) { - capBytes, err := spillCheckedMul(perWorker, uint64(workers)) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - containers := make([]container, workers) - analyzers := make([]process.Analyzer, workers) - start := make(chan struct{}) - errs := make(chan error, workers) - var wg sync.WaitGroup - for i := range containers { - containers[i].hashmapBuilder.setBudget(generation) - analyzers[i] = process.NewAnalyzer(i, false, false, "recovery fanout") - wg.Add(1) - go func(i int) { - defer wg.Done() - <-start - if reserveErr := containers[i].ensureDirectSpillRecovery(bat, analyzers[i]); reserveErr != nil { - errs <- reserveErr - return - } - errs <- containers[i].ensureRetainedSpillRecovery(projection, analyzers[i]) - }(i) - } - close(start) - wg.Wait() - close(errs) - for reserveErr := range errs { - require.NoError(t, reserveErr) - } - require.Equal(t, capBytes, generation.Used()) - for i := range containers { - require.Equal(t, perWorker, containers[i].spillScratchBase) - require.NotNil(t, containers[i].spillScratchReservation) - } - - wg = sync.WaitGroup{} - for i := range containers { - wg.Add(1) - go func(i int) { - defer wg.Done() - containers[i].releaseSpillScratchReservation() - containers[i].releaseSpillScratchReservation() - }(i) - } - wg.Wait() - require.Zero(t, generation.Used()) - generation.Close() - }) - } -} - -func TestShuffleHashBuildCancellationReleasesRecoveryLease(t *testing.T) { - typ := types.T_int32.ToType() - tc := newTestCase(t, []bool{false}, []types.Type{typ}, []*plan.Expr{newExpr(0, typ)}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = math.MaxInt64 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3504} - - ctx, cancel := context.WithCancelCause(tc.proc.Ctx) - process.ReplacePipelineCtx(tc.proc, ctx, cancel) - makeBuildBatch := func(base int32) *batch.Batch { - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = base + int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, tc.proc.Mp()) - bat.SetRowCount(len(values)) - return bat - } - first := makeBuildBatch(0) - second := makeBuildBatch(colexec.DefaultBatchSize) - input := newGatedHashBuildInput(first, second) - tc.arg.SetChildren([]vm.Operator{input}) - require.NoError(t, input.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - type execOutcome struct { - result vm.CallResult - err error - } - execDone := make(chan execOutcome, 1) - go func() { - result, execErr := vm.Exec(tc.arg, tc.proc) - execDone <- execOutcome{result: result, err: execErr} - }() - select { - case <-input.firstBatchRetained: - case outcome := <-execDone: - t.Fatalf("HashBuild stopped before the cancellation boundary: result=%+v err=%v", outcome.result, outcome.err) - } - require.NotNil(t, tc.arg.ctr.spillScratchReservation) - require.Positive(t, generation.Used()) - - cancel(context.Canceled) - outcome := <-execDone - require.ErrorIs(t, outcome.err, context.Canceled) - require.Nil(t, tc.arg.ctr.spillScratchReservation, - "build terminal cleanup owns the recovery lease") - require.Positive(t, generation.Used(), - "the copied batch remains owned until pipeline Reset") - - terminal, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), context.Background()) - require.NoError(t, err) - require.True(t, terminal.IsBuildError()) - - tc.arg.Reset(tc.proc, true, outcome.err) - tc.arg.Reset(tc.proc, true, outcome.err) - input.Reset(tc.proc, true, outcome.err) - require.Zero(t, generation.Used()) - tc.arg.Free(tc.proc, true, outcome.err) - tc.arg.Free(tc.proc, true, outcome.err) - input.Free(tc.proc, true, outcome.err) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildRecoveryAdmissionFailsBeforeRetain(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_varchar.ToType()}, []*plan.Expr{newExpr(0, types.T_varchar.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ - Tag: tc.arg.JoinMapTag + 3501, - } - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - payload := bytes.Repeat([]byte{'x'}, 1<<20) - build := batch.NewWithSize(1) - var err error - build.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), payload, 1, tc.proc.Mp()) - require.NoError(t, err) - build.SetRowCount(1) - - // Leave enough budget for the retained copy in isolation, but not for the - // direct recovery path. Retention is forbidden: failing before the copy is - // the only state that neither exceeds the cap nor strands owned build rows. - copyPeak, err := tc.arg.ctr.hashmapBuilder.projectedBatchCopyBytes(build) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(copyPeak, copyPeak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.Error(t, buildErr) - require.True(t, moerr.IsMoErrCode(buildErr, moerr.ErrOOM), - "terminal spill admission must use the resource-exhausted wire code") - require.Contains(t, buildErr.Error(), "hash build memory budget exceeded") - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"], - "the terminal direct-spill transition is counted even when admission rejects before I/O") - require.Equal(t, int64(1), extra["HashBuildSpillRecoveryReserveRejects"]) - require.Zero(t, extra["HashBuildSpillScratchReserveRejects"]) - require.Equal(t, int64(1), extra["QueryHashBudgetRejects"]) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsBuildError()) - require.Equal(t, buildErr.Error(), result.BuildError().Error()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - build.Clean(tc.proc.Mp()) - require.Zero(t, generation.Used()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildSpillModeReclaimsOptionalCacheForRecoveryGrowth(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3510} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - first := newBatch(tc.types, tc.proc, 1) - second := newBatch(tc.types, tc.proc, 2048) - firstNeed, err := spillBudgetBytes(first) - require.NoError(t, err) - firstNeed, err = spillRecoveryReservationBytes(firstNeed) - require.NoError(t, err) - secondNeed, err := spillBudgetBytes(second) - require.NoError(t, err) - secondNeed, err = spillRecoveryReservationBytes(secondNeed) - require.NoError(t, err) - require.Equal(t, firstNeed+uint64(spillWriteCoalesceSize), secondNeed, - "fixture needs one optional-cache quantum between recovery high waters") - - budget := process.MustNewHashBuildBudget(secondNeed, secondNeed) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(first, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(second, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - result, buildErr := vm.Exec(tc.arg, tc.proc) - require.NoError(t, buildErr) - require.Equal(t, vm.ExecStop, result.Status) - - joinResult, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, joinResult.IsSuccess()) - jm := joinResult.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(2049), jm.GetRowCount()) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, payload.Close()) - - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"]) - require.Equal(t, int64(1), extra["HashBuildSpillRecoveryGrowRejects"]) - require.Equal(t, int64(1), extra["HashBuildSpillRecoveryGrowCount"]) - require.Equal(t, int64(secondNeed-firstNeed), - extra["HashBuildSpillRecoveryGrowBytes"]) - require.Equal(t, int64(1), extra["HashBuildCoalesceGrowRejects"], - "one failed optional probe switches the rest of this batch to write-through") - require.Equal(t, int64(2), extra["QueryHashBudgetRejects"]) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildSpillModeGrowthRejectReleasesRecoveryLease(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3503} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - first := newBatch(tc.types, tc.proc, 1) - second := newBatch(tc.types, tc.proc, 8*colexec.DefaultBatchSize) - firstNeed, err := spillBudgetBytes(first) - require.NoError(t, err) - firstNeed, err = spillRecoveryReservationBytes(firstNeed) - require.NoError(t, err) - secondNeed, err := spillBudgetBytes(second) - require.NoError(t, err) - secondNeed, err = spillRecoveryReservationBytes(secondNeed) - require.NoError(t, err) - capBytes := firstNeed + uint64(spillWriteCoalesceSize) - require.Greater(t, secondNeed, capBytes) - - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(first, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(second, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.True(t, moerr.IsMoErrCode(buildErr, moerr.ErrOOM)) - require.Contains(t, buildErr.Error(), "hash build memory budget exceeded") - - extra := tc.arg.OpAnalyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillStarts"]) - require.Equal(t, int64(2), extra["HashBuildSpillRecoveryGrowRejects"], - "reclaim retries once, then preserves a genuine over-cap failure") - require.Zero(t, extra["HashBuildCoalesceGrowRejects"]) - require.Equal(t, int64(2), extra["QueryHashBudgetRejects"]) - require.Nil(t, tc.arg.ctr.spillScratchReservation) - require.Zero(t, tc.arg.ctr.spillScratchBase) - require.Zero(t, generation.Used()) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsBuildError()) - require.Equal(t, buildErr.Error(), result.BuildError().Error()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - first.Clean(tc.proc.Mp()) - second.Clean(tc.proc.Mp()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildSpillsExpressionKey(t *testing.T) { - bindProc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - col := newExpr(0, types.T_int32.ToType()) - modulo, err := plan2.BindFuncExprImplByPlanExpr( - bindProc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{modulo}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 4000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - build := batch.NewWithSize(1) - build.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, tc.proc.Mp()) - build.SetRowCount(4) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(build, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - - _, err = vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(4), jm.GetRowCount()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, spillPayload.Close()) - budget, err := tc.proc.GetHashBuildBudget() - require.NoError(t, err) - require.Zero(t, budget.Used()) - require.Zero(t, budget.SpillDiskUsed()) - require.Zero(t, budget.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.proc.Free() - bindProc.Free() -} - -func TestShuffleHashBuildResizeRejectReleasesPartialMapAndSpills(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := aggregate.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - // With one full ingress batch, the first four admissions are emergency - // scratch, retained copy, build auxiliary memory, and the initial map. - // Reject exactly the fifth admission: the first resize after insertion. - providerCalls := 0 - forcedResizeReject := false - aggregate.SetAggregateCapProvider(func() (uint64, error) { - providerCalls++ - if providerCalls == 5 { - forcedResizeReject = true - return generation.Used(), nil - } - return capBytes, nil - }) - - bat := newBatch(tc.types, tc.proc, 8192) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(bat, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err = vm.Exec(tc.arg, tc.proc) + account, err := registry.OpenWithController(2048, generation) require.NoError(t, err) - require.True(t, forcedResizeReject) - require.Zero(t, generation.Used(), "partial map, retained batches, and emergency scratch must be released before publication") - - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.NotNil(t, jm) - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(8192), jm.GetRowCount()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, spillPayload.Close()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildClosedMapBudgetDoesNotSpill(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 3001} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := aggregate.OpenGeneration(1) + mp := mpool.MustNewZero() + allocation, err := mp.AllocAccounted(128, account, HashBuildAllocationOwner, HashBuildAllocationSiteHashCell) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - // With one full ingress batch, the fourth admission is the initial hashmap. - // A lifecycle failure there is fatal and must never enter spill recovery. - providerCalls := 0 - forcedClosed := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - } - aggregate.SetAggregateCapProvider(func() (uint64, error) { - providerCalls++ - if providerCalls == 4 { - return 0, forcedClosed - } - return capBytes, nil - }) - - bat := newBatch(tc.types, tc.proc, 8192) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(bat, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.Same(t, forcedClosed, buildErr) - require.ErrorIs(t, buildErr, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, buildErr, process.ErrHashBuildBudgetAdmission) - require.Equal(t, 4, providerCalls) - require.Zero(t, tc.arg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) - require.Empty(t, tc.arg.ctr.spilledFds) - require.Nil(t, tc.arg.ctr.spillBundle) - - result, err := message.ReceiveJoinMapResult( - tc.arg.JoinMapTag, - true, - tc.arg.ShuffleIdx, - tc.proc.GetMessageBoard(), - tc.proc.Ctx, - ) - require.NoError(t, err) - require.True(t, result.IsBuildError()) - require.False(t, result.IsSuccess()) - require.Nil(t, result.JoinMap()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - generation.Close() - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestObserveHashBuildBudgetUsesGenerationSnapshot(t *testing.T) { - budget := process.MustNewHashBuildBudget(1024, 1024) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reservation, err := generation.Reserve(128) - require.NoError(t, err) - _, err = generation.Reserve(1024) + _, err = mp.AllocAccounted(1024, account, HashBuildAllocationOwner, HashBuildAllocationSiteHashCell) require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - reservation.Release() + mp.Free(allocation) analyzer := process.NewAnalyzer(0, false, false, "hash build") observeHashBuildBudget(analyzer, generation) @@ -3471,115 +2787,7 @@ func TestObserveHashBuildBudgetUsesGenerationSnapshot(t *testing.T) { observeHashBuildBudget(analyzer, generation) require.Equal(t, int64(1), extra["QueryHashBudgetRejects"]) generation.Close() -} - -func TestShuffleHashBuildSpillFailureReleasesEmergencyResources(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 4000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - runtimeFilterReceiver := message.NewMessageReceiver( - []int32{tc.arg.RuntimeFilterSpec.Tag}, - message.AddrBroadCastOnCurrentCN(), - tc.proc.GetMessageBoard(), - ) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - aggregate := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := aggregate.OpenGenerationWithSpillCaps(1, 64<<20, 1, 32) + account.Seal() + _, err = registry.Finalize(account) require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - - bat := newBatch(tc.types, tc.proc, 8192) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(bat, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, buildErr := vm.Exec(tc.arg, tc.proc) - require.True(t, moerr.IsMoErrCode(buildErr, moerr.ErrOOM)) - require.Contains(t, buildErr.Error(), "hash build spill disk budget exceeded") - require.Contains(t, buildErr.Error(), "requested=") - require.Contains(t, buildErr.Error(), "processLimitationSpillSize") - require.NotErrorIs(t, buildErr, process.ErrHashBuildBudgetAdmission) - require.Nil(t, tc.arg.ctr.spillScratchReservation) - require.Zero(t, cap(tc.arg.ctr.spillHashValues)) - require.Zero(t, cap(tc.arg.ctr.spillSelection)) - require.Zero(t, cap(tc.arg.ctr.spillKeyVecs)) - require.Zero(t, tc.arg.ctr.spillWriteBuf.Cap()) - - tc.arg.Reset(tc.proc, true, buildErr) - tc.arg.Reset(tc.proc, true, buildErr) - runtimeFilters, done, receiveErr := runtimeFilterReceiver.ReceiveMessage( - false, tc.proc.Ctx) - require.NoError(t, receiveErr) - require.False(t, done) - require.Len(t, runtimeFilters, 1, - "repeated Reset must publish one terminal value per generation") - tc.arg.Free(tc.proc, true, buildErr) - tc.arg.Free(tc.proc, true, buildErr) - tc.marg.Reset(tc.proc, true, buildErr) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) -} - -func TestShuffleHashBuildDrainsRetainedBatchBeforeGrowingScratch(t *testing.T) { - tc := newTestCase(t, []bool{false}, []types.Type{types.T_int32.ToType()}, []*plan.Expr{newExpr(0, types.T_int32.ToType())}) - tc.arg.IsShuffle = true - tc.arg.ShuffleIdx = 0 - tc.arg.SpillThreshold = 1 << 30 - tc.arg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: tc.arg.JoinMapTag + 5000} - tc.arg.SetChildren([]vm.Operator{tc.marg}) - require.NoError(t, tc.marg.Prepare(tc.proc)) - require.NoError(t, tc.arg.Prepare(tc.proc)) - - const capBytes = uint64(64 << 20) - aggregate := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := aggregate.OpenGeneration(1) - require.NoError(t, err) - tc.arg.ctr.hashmapBuilder.setBudget(generation) - providerCalls := 0 - forcedGrowReject := false - aggregate.SetAggregateCapProvider(func() (uint64, error) { - providerCalls++ - // First ingress reserves direct scratch, grows it for future retained - // drain, then reserves its copy. Reject the fourth admission: growing - // direct scratch for the larger second ingress while that copy is live. - if providerCalls == 4 { - forcedGrowReject = true - return generation.Used(), nil - } - return capBytes, nil - }) - - first := newBatch(tc.types, tc.proc, 8192) - second := newBatch(tc.types, tc.proc, 65536) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(first, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(second, nil, tc.proc.Mp()) - tc.proc.Reg.MergeReceivers[0].Ch2 <- process.NewPipelineSignalToDirectly(nil, nil, tc.proc.Mp()) - _, err = vm.Exec(tc.arg, tc.proc) - require.NoError(t, err) - require.True(t, forcedGrowReject) - require.Zero(t, generation.Used()) - - result, err := message.ReceiveJoinMapResult(tc.arg.JoinMapTag, true, tc.arg.ShuffleIdx, tc.proc.GetMessageBoard(), tc.proc.Ctx) - require.NoError(t, err) - require.True(t, result.IsSuccess()) - jm := result.JoinMap() - require.True(t, jm.IsSpilled()) - require.Equal(t, int64(73728), jm.GetRowCount()) - spillPayload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.NoError(t, spillPayload.Close()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - - tc.arg.Reset(tc.proc, false, nil) - tc.arg.Free(tc.proc, false, nil) - tc.marg.Reset(tc.proc, false, nil) - tc.proc.Free() - require.Zero(t, tc.proc.Mp().CurrNB()) } diff --git a/pkg/sql/colexec/hashbuild/dedup_memory.go b/pkg/sql/colexec/hashbuild/dedup_memory.go new file mode 100644 index 0000000000000..7d285d31793ee --- /dev/null +++ b/pkg/sql/colexec/hashbuild/dedup_memory.go @@ -0,0 +1,97 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +func makeDedupSlice[T any]( + hb *HashmapBuilder, + n int, + mp *mpool.MPool, + site mpool.AllocationSite, +) ([]T, error) { + if n < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + if hb.mapAllocationAccount == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + return mpool.MakeSliceAccounted[T]( + n, + mp, + hb.mapAllocationAccount, + HashBuildAllocationOwner, + site, + ) +} + +func freeDedupSlice[T any](hb *HashmapBuilder, values []T, mp *mpool.MPool) { + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} + +func (hb *HashmapBuilder) newDedupBitmap( + rows int, + mp *mpool.MPool, + site mpool.AllocationSite, +) (*bitmap.Bitmap, error) { + if rows < 0 || rows > math.MaxInt-63 { + return nil, mpool.ErrAllocationAccountInvalid + } + bm := &bitmap.Bitmap{} + if hb.mapAllocationAccount == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + words := (rows + 63) / 64 + storage, err := mpool.MakeSliceAccounted[uint64]( + words, + mp, + hb.mapAllocationAccount, + HashBuildAllocationOwner, + site, + ) + if err != nil { + return nil, err + } + bm.InstallExternalStorage(storage) + bm.InitWithSize(int64(rows)) + return bm, nil +} + +func releaseDedupBitmap(bm *bitmap.Bitmap, mp *mpool.MPool) { + if bm == nil || !bm.HasExternalStorage() { + return + } + storage := bm.ReleaseExternalStorage() + if cap(storage) > 0 { + mpool.FreeSlice(mp, storage) + } +} + +func (hb *HashmapBuilder) freeIgnoreRows(mp *mpool.MPool) { + releaseDedupBitmap(hb.IgnoreRows, mp) + hb.IgnoreRows = nil +} + +func (hb *HashmapBuilder) freeDelRows(mp *mpool.MPool) { + releaseDedupBitmap(hb.DelRows, mp) + hb.DelRows = nil +} diff --git a/pkg/sql/colexec/hashbuild/errors.go b/pkg/sql/colexec/hashbuild/errors.go index 2a1b2bf46b6aa..36957706ebc3b 100644 --- a/pkg/sql/colexec/hashbuild/errors.go +++ b/pkg/sql/colexec/hashbuild/errors.go @@ -20,6 +20,7 @@ import ( "strings" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -27,7 +28,7 @@ import ( // operator's public Call boundary. Spill and other recovery paths must keep the // typed admission error until they have exhausted every recovery option. func TerminalBudgetError(ctx context.Context, err error) error { - if err == nil || !errors.Is(err, process.ErrHashBuildBudgetAdmission) { + if err == nil { return err } // A joined lifecycle/accounting failure is not a capacity rejection. Keep @@ -40,6 +41,24 @@ func TerminalBudgetError(ctx context.Context, err error) error { var budgetErr *process.HashBuildBudgetError if !errors.As(err, &budgetErr) || budgetErr.Kind != process.HashBuildBudgetErrorAdmission { + switch { + case mpool.AllocationFailureReasonOf(err) == + mpool.AllocationFailureCapacity && + !mpool.IsMPoolCapacityFailure(err): + return moerr.NewResourceExhaustedf(ctx, + "hash build memory budget exceeded; reduce join build width or query concurrency, increase processLimitationSize, or lower join_spill_mem for an eligible shuffle join") + case errors.Is(err, process.ErrHashBuildBudgetAdmission): + return moerr.NewResourceExhaustedf(ctx, + "hash build resource budget exceeded; inspect hash-build budget metrics and resource limits") + default: + return err + } + } + if budgetErr.Component == 0 { + reason := terminalBudgetReason(budgetErr.Message) + if reason != "" { + return moerr.NewResourceExhaustedf(ctx, "%s", reason) + } return moerr.NewResourceExhaustedf( ctx, "hash build resource budget exceeded; inspect hash-build budget metrics and resource limits", @@ -48,14 +67,14 @@ func TerminalBudgetError(ctx context.Context, err error) error { reason := terminalBudgetReason(budgetErr.Message) var resource, action string - switch budgetErr.Resource { - case process.HashBuildBudgetResourceMemory: + switch budgetErr.Component { + case process.HashBuildBudgetComponentMemory: resource = "memory" action = "reduce join build width or query concurrency, increase processLimitationSize, or lower join_spill_mem for an eligible shuffle join; automatic spill can still exhaust recovery headroom for wide or skewed partitions" - case process.HashBuildBudgetResourceSpillDisk: + case process.HashBuildBudgetComponentSpillDisk: resource = "spill disk" action = "free spill storage or increase processLimitationSpillSize" - case process.HashBuildBudgetResourceSpillFD: + case process.HashBuildBudgetComponentSpillFD: resource = "spill file descriptor" action = "reduce concurrent spill work or raise the CN open-file limit" default: diff --git a/pkg/sql/colexec/hashbuild/errors_test.go b/pkg/sql/colexec/hashbuild/errors_test.go index b6d890a976cbb..d7243eee729f7 100644 --- a/pkg/sql/colexec/hashbuild/errors_test.go +++ b/pkg/sql/colexec/hashbuild/errors_test.go @@ -20,6 +20,7 @@ import ( "testing" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) @@ -32,18 +33,18 @@ func TestTerminalBudgetError(t *testing.T) { }) for _, tc := range []struct { - name string - resource process.HashBuildBudgetResource - want []string + name string + component process.HashBuildBudgetComponent + want []string }{ - {"memory", process.HashBuildBudgetResourceMemory, []string{"memory", "requested=3", "used=5", "limit=7", "build width", "processLimitationSize", "join_spill_mem", "recovery headroom"}}, - {"spill disk", process.HashBuildBudgetResourceSpillDisk, []string{"spill disk", "requested=3", "used=5", "limit=7", "processLimitationSpillSize"}}, - {"spill fd", process.HashBuildBudgetResourceSpillFD, []string{"spill file descriptor", "requested=3", "used=5", "limit=7", "open-file limit"}}, + {"memory", process.HashBuildBudgetComponentMemory, []string{"memory", "requested=3", "used=5", "limit=7", "build width", "processLimitationSize", "join_spill_mem", "recovery headroom"}}, + {"spill disk", process.HashBuildBudgetComponentSpillDisk, []string{"spill disk", "requested=3", "used=5", "limit=7", "processLimitationSpillSize"}}, + {"spill fd", process.HashBuildBudgetComponentSpillFD, []string{"spill file descriptor", "requested=3", "used=5", "limit=7", "open-file limit"}}, } { t.Run(tc.name, func(t *testing.T) { err := TerminalBudgetError(context.Background(), &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: tc.resource, + Component: tc.component, Requested: 3, Used: 5, Cap: 7, @@ -72,7 +73,7 @@ func TestTerminalBudgetError(t *testing.T) { t.Run("resource admission keeps spill depth context", func(t *testing.T) { err := TerminalBudgetError(context.Background(), &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, + Component: process.HashBuildBudgetComponentMemory, Requested: 3, Used: 5, Cap: 7, @@ -91,6 +92,29 @@ func TestTerminalBudgetError(t *testing.T) { require.NotContains(t, err.Error(), process.ErrHashBuildBudgetAdmission.Error()) }) + t.Run("physical capacity is terminal resource exhaustion", func(t *testing.T) { + err := TerminalBudgetError( + context.Background(), mpool.ErrAllocationAccountCapacity) + require.True(t, moerr.IsMoErrCode(err, moerr.ErrOOM)) + require.Contains(t, err.Error(), "hash build memory budget exceeded") + require.Contains(t, err.Error(), "processLimitationSize") + }) + + t.Run("mpool capacity preserves allocator error", func(t *testing.T) { + capacity := moerr.NewMPoolCapacityNoCtxf("mpool out of space") + require.Same(t, capacity, + TerminalBudgetError(context.Background(), capacity)) + }) + + t.Run("physical lifecycle failure stays fatal", func(t *testing.T) { + joined := errors.Join( + mpool.ErrAllocationAccountCapacity, + mpool.ErrAllocationAccountSealed, + ) + require.Same(t, joined, + TerminalBudgetError(context.Background(), joined)) + }) + for _, lifecycle := range []error{ process.ErrHashBuildBudgetClosed, process.ErrHashBuildBudgetInvalid, diff --git a/pkg/sql/colexec/hashbuild/expression_memory.go b/pkg/sql/colexec/hashbuild/expression_memory.go index 653f8096f93a5..3f2b776d5b34c 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory.go +++ b/pkg/sql/colexec/hashbuild/expression_memory.go @@ -15,753 +15,63 @@ package hashbuild import ( - "math" - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/process" ) -type expressionMemoryLeaseSlot struct { - tokens []*process.HashBuildReservation - admittedPeak uint64 - mayReplaceWithinBound bool - recoveryPeak uint64 - recoveryMayReplace bool - recoveryCandidate uint64 - recoveryCandidateReplace bool -} - -// ExpressionMemoryLease couples one stable expression-executor set to the -// HashBuild reservation that covers its retained vector capacity. It is not -// safe for concurrent use. -// -// Run retains the largest admitted expression bound independently for every -// root executor. Re-evaluating a root for an equal or smaller batch therefore -// reuses its existing high-water reservation. Growth admits only the part of -// that root's allocate-copy-free peak not covered by the old reservation, then -// reconciles the transient charge into the new high water. Unrelated roots are -// never double-charged. -// -// The owner must Free every executor and duplicate vector covered by the lease -// before calling Release. A lease cannot cross statement budget generations. -type ExpressionMemoryLease struct { - budget *process.HashBuildBudgetGeneration - exprs []*plan.Expr - executors []colexec.ExpressionExecutor - duplicate bool - slots []expressionMemoryLeaseSlot - recoveryReservation *process.HashBuildReservation - recoveryRows int - recoveryReady bool - recoveryReconcile bool - released bool -} +const ( + hashBuildAllocationSiteExpressionData mpool.AllocationSite = iota + 98 + hashBuildAllocationSiteExpressionArea + hashBuildAllocationSiteExpressionNulls + hashBuildAllocationSiteExpressionGrouping +) -// NewBudgetedExpressionExecutors admits the mpool-backed capacity owned by -// constant children before constructing them. The returned lease adopts those -// reservations, so construction and later evaluation have one continuous -// budget lifetime. -func NewBudgetedExpressionExecutors( +// NewExpressionExecutors constructs expression trees used by HashBuild and +// join operators. Every MPool vector owned by the tree, including nested +// function results and reusable selection buffers, shares the query account. +func NewExpressionExecutors( proc *process.Process, - budget *process.HashBuildBudgetGeneration, exprs []*plan.Expr, - duplicate bool, -) ([]colexec.ExpressionExecutor, *ExpressionMemoryLease, error) { - if budget == nil { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) - if err != nil { - return nil, nil, err - } - lease, err := NewExpressionMemoryLease(nil, exprs, executors, duplicate) - if err != nil { - for _, executor := range executors { - executor.Free() - } - return nil, nil, err - } - return executors, lease, err - } - - executors := make([]colexec.ExpressionExecutor, len(exprs)) - lease := &ExpressionMemoryLease{ - budget: budget, - exprs: exprs, - executors: executors, - duplicate: duplicate, - slots: make([]expressionMemoryLeaseSlot, len(exprs)), - } - cleanup := func() { - for _, executor := range executors { - if executor != nil { - executor.Free() - } - } - lease.Release() - } - - for i, expr := range exprs { - initial, err := expressionInitialOwnedBytes(expr) - if err != nil { - cleanup() - return nil, nil, err - } - if initial > 0 { - token, err := budget.Reserve(initial) - if err != nil { - cleanup() - return nil, nil, err - } - lease.slots[i].tokens = append(lease.slots[i].tokens, token) - lease.slots[i].admittedPeak = initial - } - - executor, err := colexec.NewExpressionExecutor(proc, expr) - if err != nil { - cleanup() - return nil, nil, err - } - executors[i] = executor - - retained, ok := colexec.ExpressionExecutorRetainedBytes(executor) - if !ok || retained > initial { - cleanup() - return nil, nil, process.ErrHashBuildBudgetInvalid - } - slot := &lease.slots[i] - slot.mayReplaceWithinBound = expressionExecutorMayGrowWithinBound(expr) - if len(slot.tokens) > 0 { - token := slot.tokens[0] - if retained == 0 { - token.Release() - slot.tokens = nil - } else if _, err = token.ReconcileDown(retained); err != nil { - cleanup() - return nil, nil, err - } - } - slot.admittedPeak = retained - } - return executors, lease, nil + account *mpool.AllocationAccount, +) ([]colexec.ExpressionExecutor, error) { + return newExpressionExecutorsWithCapacityClass( + proc, + exprs, + account, + mpool.AllocationCapacityClassDefault, + ) } -func NewExpressionMemoryLease( - budget *process.HashBuildBudgetGeneration, +func newExpressionExecutorsWithCapacityClass( + proc *process.Process, exprs []*plan.Expr, - executors []colexec.ExpressionExecutor, - duplicate bool, -) (*ExpressionMemoryLease, error) { - if len(exprs) != len(executors) { + account *mpool.AllocationAccount, + capacityClass mpool.AllocationCapacityClass, +) ([]colexec.ExpressionExecutor, error) { + if len(exprs) == 0 { return nil, process.ErrHashBuildBudgetInvalid } - lease := &ExpressionMemoryLease{ - budget: budget, - exprs: exprs, - executors: executors, - duplicate: duplicate, - slots: make([]expressionMemoryLeaseSlot, len(executors)), - } - if budget == nil { - return lease, nil - } - - for i, executor := range executors { - retained, ok := colexec.ExpressionExecutorRetainedBytes(executor) - if !ok { - lease.Release() - return nil, process.ErrHashBuildBudgetInvalid - } - lease.slots[i].admittedPeak = retained - lease.slots[i].mayReplaceWithinBound = expressionExecutorMayGrowWithinBound(exprs[i]) - if retained == 0 { - continue - } - token, err := budget.Reserve(retained) - if err != nil { - lease.Release() - return nil, err - } - lease.slots[i].tokens = append(lease.slots[i].tokens, token) - } - return lease, nil -} - -func expressionInitialOwnedBytes(expr *plan.Expr) (uint64, error) { - if expr == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - switch typed := expr.Expr.(type) { - case *plan.Expr_Lit: - if typed.Lit == nil || typed.Lit.GetIsnull() { - return 0, nil - } - return literalInitialOwnedBytes(types.T(expr.Typ.Id), typed.Lit) - case *plan.Expr_List: - return expressionListInitialOwnedBytes(typed.List.GetList()) - case *plan.Expr_F: - children, err := expressionListInitialOwnedBytes(typed.F.GetArgs()) - if err != nil { - return 0, err - } - own := expressionFunctionInitialOwnedBytes(typed.F) - if children > math.MaxUint64-own { - return 0, process.ErrHashBuildBudgetInvalid - } - return children + own, nil - default: - return 0, nil - } -} - -func expressionFunctionInitialOwnedBytes(fn *plan.Function) uint64 { - if fn == nil || fn.Func == nil { - return 0 - } - fid, _ := function.DecodeOverloadID(fn.Func.Obj) - if fid == function.SERIAL || fid == function.SERIAL_FULL { - return types.DefaultPackerCapacity() - } - return 0 -} - -func expressionListInitialOwnedBytes(exprs []*plan.Expr) (uint64, error) { - var total uint64 for _, expr := range exprs { - size, err := expressionInitialOwnedBytes(expr) - if err != nil || total > math.MaxUint64-size { - return 0, process.ErrHashBuildBudgetInvalid - } - total += size - } - return total, nil -} - -func expressionExecutorMayGrowWithinBound(expr *plan.Expr) bool { - return expressionExecutorMayGrowWithinSelection(expr, false) -} - -// expressionExecutorMayGrowWithinSelection reports whether an executor tree -// can replace retained mpool capacity without increasing its input row bound. -// Varlena values can change size at the same row count. A function evaluated -// through a partial flow-control mask can likewise grow its cached compacted -// parameters and selected result when only the mask distribution changes. -func expressionExecutorMayGrowWithinSelection( - expr *plan.Expr, - mayReceivePartialSelection bool, -) bool { - if expr == nil { - return true - } - switch typed := expr.Expr.(type) { - case *plan.Expr_Col, *plan.Expr_Lit, *plan.Expr_T, *plan.Expr_Vec: - return false - case *plan.Expr_F: - if typed.F == nil || types.T(expr.Typ.Id).FixedLength() < 0 || - mayReceivePartialSelection { - return true - } - fid := int32(-1) - if typed.F.Func != nil { - fid, _ = function.DecodeOverloadID(typed.F.Func.Obj) - } - for i, arg := range typed.F.GetArgs() { - if expressionExecutorMayGrowWithinSelection( - arg, - expressionChildMayReceivePartialSelection( - fid, i, mayReceivePartialSelection), - ) { - return true - } - } - return false - case *plan.Expr_List: - if types.T(expr.Typ.Id).FixedLength() < 0 { - return true - } - for _, item := range typed.List.GetList() { - if expressionExecutorMayGrowWithinSelection( - item, mayReceivePartialSelection) { - return true - } - } - return false - case *plan.Expr_P, *plan.Expr_V: - return types.T(expr.Typ.Id).FixedLength() < 0 - default: - return true - } -} - -func expressionChildMayReceivePartialSelection( - fid int32, - argument int, - parentMayReceivePartialSelection bool, -) bool { - switch fid { - case function.IFF, function.CASE, function.COALESCE: - // The first argument inherits the caller's mask. Every later argument - // can receive a mask narrowed by an earlier condition/value. - return parentMayReceivePartialSelection || argument > 0 - default: - return parentMayReceivePartialSelection - } -} - -func literalInitialOwnedBytes(oid types.T, literal *plan.Literal) (uint64, error) { - var dataBytes uint64 - var payloadBytes uint64 - switch value := literal.GetValue().(type) { - case *plan.Literal_Bval, *plan.Literal_I8Val, *plan.Literal_U8Val, *plan.Literal_Defaultval: - dataBytes = 1 - case *plan.Literal_I16Val, *plan.Literal_U16Val, *plan.Literal_EnumVal: - dataBytes = 2 - case *plan.Literal_I32Val, *plan.Literal_U32Val, *plan.Literal_Fval, *plan.Literal_Dateval: - dataBytes = 4 - case *plan.Literal_I64Val, *plan.Literal_U64Val, *plan.Literal_Dval, - *plan.Literal_Timeval, *plan.Literal_Datetimeval, - *plan.Literal_Decimal64Val, *plan.Literal_Timestampval: - dataBytes = 8 - case *plan.Literal_Decimal128Val: - dataBytes = 16 - case *plan.Literal_Sval: - dataBytes = types.VarlenaSize - switch oid { - case types.T_array_float32: - if uint64(len(value.Sval)) > math.MaxUint64/4 { - return 0, process.ErrHashBuildBudgetInvalid - } - // The textual representation has at least one byte per element. - // Reserve a parsing-allocation-free upper bound; construction later - // reconciles it to the actual binary payload. - payloadBytes = uint64(len(value.Sval)) * 4 - case types.T_array_float64: - if uint64(len(value.Sval)) > math.MaxUint64/8 { - return 0, process.ErrHashBuildBudgetInvalid - } - payloadBytes = uint64(len(value.Sval)) * 8 - default: - payloadBytes = uint64(len(value.Sval)) + if expr == nil { + return nil, process.ErrHashBuildBudgetInvalid } - case *plan.Literal_VecVal: - dataBytes = types.VarlenaSize - payloadBytes = uint64(len(value.VecVal)) - default: - // Unsupported literal kinds are rejected by the expression factory - // before they can own an mpool-backed vector. - return 0, nil } - - dataCapacity, err := initialAllocationCapacity(dataBytes) + selection, err := vector.NewAllocationAccountSelectionWithCapacityClass( + account, + HashBuildAllocationOwner, + hashBuildAllocationSiteExpressionData, + hashBuildAllocationSiteExpressionArea, + hashBuildAllocationSiteExpressionNulls, + hashBuildAllocationSiteExpressionGrouping, + capacityClass, + ) if err != nil { - return 0, err - } - if payloadBytes <= types.VarlenaInlineSize { - return dataCapacity, nil - } - areaCapacity, err := initialAllocationCapacity(payloadBytes) - if err != nil || dataCapacity > math.MaxUint64-areaCapacity { - return 0, process.ErrHashBuildBudgetInvalid - } - return dataCapacity + areaCapacity, nil -} - -func initialAllocationCapacity(required uint64) (uint64, error) { - if required == 0 { - return 0, nil - } - if required > math.MaxInt64 { - return 0, process.ErrHashBuildBudgetInvalid - } - capacity, ok := mpool.GrowCapacity(0, int64(required)) - if !ok || capacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - return uint64(capacity), nil -} - -// EnsureRunRecovery admits the complete expression-executor peak needed to -// evaluate batches up to rows without another budget reservation. HashBuild -// calls it before retaining a spillable batch, so spill can always evaluate a -// key and release the first retained batch even when sibling workers consume -// all remaining query headroom. -// -// Root executors run sequentially but retain their result capacities. The -// first-run bound therefore replaces each root's old admitted capacity with -// its target peak in evaluation order. Executors whose capacity can still -// vary inside the same row bound (varlena values or flow-control selected -// scratch) may always overlap one old result with one replacement. A stable -// fixed-width root needs the same protection only until it has successfully -// reached this row high water. One aggregate standby token covers the largest -// such overlap instead of pessimistically reserving two copies for every root. -func (l *ExpressionMemoryLease) EnsureRunRecovery( - proc *process.Process, - rows int, -) error { - if l == nil { - return process.ErrHashBuildBudgetInvalid - } - if l.released { - return process.ErrHashBuildReservationInactive - } - if rows < 0 || len(l.exprs) != len(l.executors) || len(l.slots) != len(l.executors) { - return process.ErrHashBuildBudgetInvalid - } - if l.budget == nil || (l.recoveryReady && rows <= l.recoveryRows) { - return nil - } - - var running uint64 - for i := range l.slots { - retained, ok := colexec.ExpressionExecutorRetainedBytes(l.executors[i]) - if !ok || retained > l.slots[i].admittedPeak { - return process.ErrHashBuildBudgetInvalid - } - if running > math.MaxUint64-l.slots[i].admittedPeak { - return process.ErrHashBuildBudgetInvalid - } - running += l.slots[i].admittedPeak - } - - target := running - var replacementOverlap uint64 - for i, expr := range l.exprs { - peak, err := expressionVectorPeak(proc, expr, rows, l.duplicate) - if err != nil { - return err - } - slot := &l.slots[i] - slot.recoveryCandidate = peak - slot.recoveryCandidateReplace = slot.mayReplaceWithinBound - - allocation := peak - nextRetained := peak - if slot.admittedPeak > nextRetained { - // Executors retain reusable high-water capacity; evaluating a - // smaller variable-width value does not prove that old capacity was - // returned to mpool. - nextRetained = slot.admittedPeak - } - if !slot.mayReplaceWithinBound && peak <= slot.admittedPeak { - // A fixed-width executor whose high water already covers this row - // bound reuses its vectors and retains the existing capacity. - allocation = 0 - nextRetained = slot.admittedPeak - } else if !slot.mayReplaceWithinBound { - // Until this exact high water has executed successfully, a smaller - // fixed-width result can still overlap its later replacement. - slot.recoveryCandidateReplace = true - } - if running > math.MaxUint64-allocation { - return process.ErrHashBuildBudgetInvalid - } - if candidate := running + allocation; candidate > target { - target = candidate - } - if running < slot.admittedPeak || - running-slot.admittedPeak > math.MaxUint64-nextRetained { - return process.ErrHashBuildBudgetInvalid - } - running = running - slot.admittedPeak + nextRetained - if slot.recoveryCandidateReplace && peak > replacementOverlap { - replacementOverlap = peak - } - } - if running > math.MaxUint64-replacementOverlap { - return process.ErrHashBuildBudgetInvalid - } - if repeated := running + replacementOverlap; repeated > target { - target = repeated - } - - reserved := l.Reserved() - if target > reserved { - growth := target - reserved - if l.recoveryReservation == nil { - reservation, err := l.budget.Reserve(growth) - if err != nil { - return err - } - l.recoveryReservation = reservation - } else if err := l.recoveryReservation.Grow(growth); err != nil { - return err - } - } - - for i := range l.slots { - l.slots[i].recoveryPeak = l.slots[i].recoveryCandidate - l.slots[i].recoveryMayReplace = l.slots[i].recoveryCandidateReplace - } - l.recoveryRows = rows - l.recoveryReady = true - return nil -} - -func (l *ExpressionMemoryLease) reconcileRecoveryAfterRun() error { - if l == nil || l.recoveryReservation == nil { - return nil - } - var steady uint64 - var replacementOverlap uint64 - for i := range l.slots { - slot := &l.slots[i] - retained := slot.admittedPeak - if slot.recoveryPeak > retained { - retained = slot.recoveryPeak - } - if steady > math.MaxUint64-retained { - return process.ErrHashBuildBudgetInvalid - } - steady += retained - if slot.recoveryMayReplace && slot.recoveryPeak > replacementOverlap { - replacementOverlap = slot.recoveryPeak - } - } - if steady > math.MaxUint64-replacementOverlap { - return process.ErrHashBuildBudgetInvalid - } - steady += replacementOverlap - reserved := l.Reserved() - if reserved <= steady { - return nil - } - release := reserved - steady - recoverySize := l.recoveryReservation.Size() - if release > recoverySize { - release = recoverySize - } - _, err := l.recoveryReservation.ReconcileDown(recoverySize - release) - return err -} - -// Run admits and evaluates each root in index order. Growth keeps the root's -// old retained charge live, reserves only the uncovered allocate-copy-free -// overlap, and reconciles that overlap into the new high-water charge. -func (l *ExpressionMemoryLease) Run( - proc *process.Process, - rows int, - fn func(index int) error, -) (err error) { - if fn == nil { - return process.ErrHashBuildBudgetInvalid - } - if l == nil { - return process.ErrHashBuildBudgetInvalid - } - if l.released { - return process.ErrHashBuildReservationInactive - } - if rows < 0 { - return process.ErrHashBuildBudgetInvalid - } - if l.budget != nil && l.recoveryReady && rows <= l.recoveryRows { - for i, expr := range l.exprs { - peak := l.slots[i].recoveryPeak - if rows < l.recoveryRows { - var peakErr error - peak, peakErr = expressionVectorPeak(proc, expr, rows, l.duplicate) - if peakErr != nil { - return peakErr - } - } - evalErr := fn(i) - if l.released { - return evalErr - } - if peak > l.slots[i].admittedPeak { - l.slots[i].admittedPeak = peak - } - if evalErr == nil && !l.slots[i].mayReplaceWithinBound && rows == l.recoveryRows && - l.slots[i].recoveryMayReplace { - l.slots[i].recoveryMayReplace = false - l.recoveryReconcile = true - } - if evalErr != nil { - return evalErr - } - } - if l.recoveryReconcile { - // Return only replacement headroom whose fixed-width root has reached - // the prepared row high water. Keep every retained capacity, every root - // still growing toward that high water, and variable replacement owned. - if err := l.reconcileRecoveryAfterRun(); err != nil { - return err - } - l.recoveryReconcile = false - } - return nil - } - if l.recoveryReady && rows > l.recoveryRows { - // A larger unprepared evaluation can change the retained-capacity - // state. Keep the standby token owned, but require the next recovery - // preflight to recompute its guarantee from that new state. - l.recoveryReady = false - } - - for i, expr := range l.exprs { - if l.budget == nil { - if err := fn(i); err != nil { - return err - } - continue - } - - peak, peakErr := expressionVectorPeak(proc, expr, rows, l.duplicate) - if peakErr != nil { - return peakErr - } - slot := &l.slots[i] - if peak <= slot.admittedPeak && !slot.mayReplaceWithinBound { - if err := fn(i); err != nil { - return err - } - continue - } - - retained, ok := colexec.ExpressionExecutorRetainedBytes(l.executors[i]) - if !ok || retained > slot.admittedPeak { - return process.ErrHashBuildBudgetInvalid - } - var transient uint64 - if retained > math.MaxUint64-peak { - return process.ErrHashBuildBudgetInvalid - } - physicalPeak := retained + peak - if physicalPeak > slot.admittedPeak { - transient = physicalPeak - slot.admittedPeak - } - if transient == 0 { - if err := fn(i); err != nil { - return err - } - continue - } - - candidate, reserveErr := l.budget.Reserve(transient) - if reserveErr != nil { - return reserveErr - } - evalErr := fn(i) - if l.released { - candidate.Release() - return evalErr - } - if peak > slot.admittedPeak { - growth := peak - slot.admittedPeak - if _, reconcileErr := candidate.ReconcileDown(growth); reconcileErr != nil { - slot.tokens = append(slot.tokens, candidate) - slot.admittedPeak += transient - if evalErr != nil { - return evalErr - } - return reconcileErr - } - slot.tokens = append(slot.tokens, candidate) - slot.admittedPeak = peak - } else { - candidate.Release() - } - if evalErr != nil { - return evalErr - } - } - return nil -} - -// Eval evaluates the executors owned by the lease in the same per-root order -// used for admission. consume receives each successfully evaluated vector -// before the next root is admitted. -func (l *ExpressionMemoryLease) Eval( - proc *process.Process, - bats []*batch.Batch, - rows int, - consume func(index int, vec *vector.Vector) error, -) error { - if l == nil || consume == nil { - return process.ErrHashBuildBudgetInvalid - } - return l.Run(proc, rows, func(index int) error { - vec, err := l.executors[index].Eval(proc, bats, nil) - if err != nil { - return err - } - return consume(index, vec) - }) -} - -func (l *ExpressionMemoryLease) Reserved() uint64 { - if l == nil || l.released { - return 0 - } - var total uint64 - for i := range l.slots { - for _, token := range l.slots[i].tokens { - size := token.Size() - if total > math.MaxUint64-size { - return math.MaxUint64 - } - total += size - } - } - if l.recoveryReservation != nil { - size := l.recoveryReservation.Size() - if total > math.MaxUint64-size { - return math.MaxUint64 - } - total += size - } - return total -} - -func (l *ExpressionMemoryLease) Len() int { - if l == nil || l.released { - return 0 - } - return len(l.executors) -} - -// Retained returns the current mpool-backed capacity physically owned by the -// executor set. Reserved may be larger: the documented delta is the retained -// high-water admission bound kept available for safe executor reuse. -func (l *ExpressionMemoryLease) Retained() (uint64, bool) { - if l == nil { - return 0, true - } - if l.released { - return 0, false - } - return colexec.ExpressionExecutorsRetainedBytes(l.executors) -} - -func (l *ExpressionMemoryLease) Release() { - if l == nil || l.released { - return - } - l.released = true - for i := range l.slots { - for _, token := range l.slots[i].tokens { - token.Release() - } - l.slots[i].tokens = nil - l.slots[i].admittedPeak = 0 - l.slots[i].recoveryPeak = 0 - l.slots[i].recoveryMayReplace = false - l.slots[i].recoveryCandidate = 0 - l.slots[i].recoveryCandidateReplace = false - } - if l.recoveryReservation != nil { - l.recoveryReservation.Release() - l.recoveryReservation = nil + return nil, err } - l.recoveryRows = 0 - l.recoveryReady = false - l.recoveryReconcile = false - l.slots = nil - l.exprs = nil - l.executors = nil - l.budget = nil + return colexec.NewExpressionExecutorsFromPlanExpressionsWithAllocation( + proc, exprs, selection, + ) } diff --git a/pkg/sql/colexec/hashbuild/expression_memory_test.go b/pkg/sql/colexec/hashbuild/expression_memory_test.go index d5213bc5485a9..4352f439d76eb 100644 --- a/pkg/sql/colexec/hashbuild/expression_memory_test.go +++ b/pkg/sql/colexec/hashbuild/expression_memory_test.go @@ -4,7 +4,7 @@ // 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 +// 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, @@ -15,9 +15,6 @@ package hashbuild import ( - "context" - "errors" - "math" "strings" "testing" @@ -26,1626 +23,55 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) -func makeExpressionLeaseTestExpr(t *testing.T, proc *process.Process) *plan.Expr { - t.Helper() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - return expr -} - -func makeSerialExpressionLeaseTestExpr( - t *testing.T, - proc *process.Process, - name string, - argTypes ...types.Type, -) *plan.Expr { - t.Helper() - args := make([]*plan.Expr, len(argTypes)) - for i, typ := range argTypes { - args[i] = &plan.Expr{ - Typ: plan.Type{ - Id: int32(typ.Oid), - Width: typ.Width, - Scale: typ.Scale, - }, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: int32(i)}}, - } - } - expr, err := plan2.BindFuncExprImplByPlanExpr(proc.Ctx, name, args) +func TestRecoveryProjectionPreservesLogicalVarlenaMultiplicity(t *testing.T) { + mp := mpool.MustNewZero() + const rows = 4 + payload := strings.Repeat("x", 64) + constant, err := vector.NewConstBytes( + types.T_varchar.ToType(), []byte(payload), rows, mp) require.NoError(t, err) - return expr -} - -func makeSerialExpressionLeaseTestBatch( - t *testing.T, - proc *process.Process, - rows int, - typ types.Type, - value string, -) *batch.Batch { - return makeSerialExpressionLeaseTestBatchWithColumns(t, proc, rows, typ, value, 2) -} - -func makeSerialExpressionLeaseTestBatchWithColumns( - t *testing.T, - proc *process.Process, - rows int, - typ types.Type, - value string, - columns int, -) *batch.Batch { - t.Helper() - values := make([]string, rows) - for i := range values { - values[i] = value - } - bat := batch.NewWithSize(columns) - for i := range bat.Vecs { - bat.Vecs[i] = testutil.NewStringVector(rows, typ, proc.Mp(), false, nil, values) - require.NotNil(t, bat.Vecs[i]) - } - bat.SetRowCount(rows) - return bat -} - -func makeExpressionLeaseTestBatch(proc *process.Process, rows int) *batch.Batch { - values := make([]int32, rows) - for i := range values { - values[i] = int32(i) - } + defer constant.Free(mp) bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) + bat.Vecs[0] = constant bat.SetRowCount(rows) - return bat -} - -func makeMaxArrayLeaseTestVector[T types.ArrayElement]( - t *testing.T, - proc *process.Process, - oid types.T, -) *vector.Vector { - t.Helper() - typ := types.New(oid, types.MaxArrayDimension, 0) - vec := vector.NewVec(typ) - values := make([]T, types.MaxArrayDimension) - require.NoError(t, vector.AppendArrayList( - vec, - [][]T{values, values}, - nil, - proc.Mp(), - )) - return vec -} - -func evalExpressionLeaseTestExecutors( - proc *process.Process, - executors []colexec.ExpressionExecutor, - bat *batch.Batch, -) error { - for _, executor := range executors { - if _, err := executor.Eval(proc, []*batch.Batch{bat}, nil); err != nil { - return err - } - } - return nil -} -func freeExpressionLeaseTestExecutors(executors []colexec.ExpressionExecutor) { - for _, executor := range executors { - executor.Free() - } -} - -func TestExpressionChildMayReceivePartialSelection(t *testing.T) { - for _, test := range []struct { - name string - fid int32 - argument int - parent bool - want bool - }{ - {name: "ordinary root", fid: -1, argument: 1}, - {name: "ordinary nested", fid: -1, argument: 0, parent: true, want: true}, - {name: "iff condition", fid: function.IFF, argument: 0}, - {name: "iff branch", fid: function.IFF, argument: 1, want: true}, - {name: "case first condition", fid: function.CASE, argument: 0}, - {name: "case later condition", fid: function.CASE, argument: 2, want: true}, - {name: "coalesce first value", fid: function.COALESCE, argument: 0}, - {name: "coalesce later value", fid: function.COALESCE, argument: 1, want: true}, - {name: "nested first argument", fid: function.COALESCE, argument: 0, parent: true, want: true}, - } { - t.Run(test.name, func(t *testing.T) { - require.Equal(t, test.want, expressionChildMayReceivePartialSelection( - test.fid, test.argument, test.parent)) - }) - } -} - -func TestExpressionMemoryLeaseReusesRetainedHighWater(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - - initialRetained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Positive(t, initialRetained) - largePeak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - budgetCap := initialRetained + largePeak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) + selected, err := projectedSelectedRange(bat, 0, rows) require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - require.Equal(t, initialRetained, generation.Used()) - - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer large.Clean(proc.Mp()) - require.NoError(t, lease.Run(proc, large.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, large) - })) - require.Equal(t, largePeak, generation.Used()) - reservesAfterLarge := generation.ReserveCount() - - retainedAfterLarge, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Greater(t, retainedAfterLarge, initialRetained) - leaseRetained, ok := lease.Retained() - require.True(t, ok) - require.Equal(t, retainedAfterLarge, leaseRetained) - require.GreaterOrEqual(t, lease.Reserved(), leaseRetained) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.Run(proc, small.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, small) - })) - require.Equal(t, reservesAfterLarge, generation.ReserveCount()) - require.Equal(t, largePeak, generation.Used()) - - for _, executor := range executors { - executor.ResetForNextQuery() - } - retainedAfterReset, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Equal(t, retainedAfterLarge, retainedAfterReset) - require.NoError(t, lease.Run(proc, large.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, large) - })) - require.Equal(t, reservesAfterLarge, generation.ReserveCount()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) + require.Equal(t, uint64(rows*(types.VarlenaSize+len(payload))), selected) } -func TestExpressionMemoryLeasePreAdmitsRecoveryRun(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - const budgetCap = uint64(8 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - - input := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer input.Clean(proc.Mp()) - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.EnsureRunRecovery(proc, input.RowCount())) - require.Positive(t, lease.Reserved()) - require.Equal(t, lease.Reserved(), generation.Used()) - preparedReserved := lease.Reserved() - - blocker, err := generation.Reserve(budgetCap - generation.Used()) - require.NoError(t, err) - require.Equal(t, budgetCap, generation.Used()) - reservesAtFullBudget := generation.ReserveCount() - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{small}, - small.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Equal(t, budgetCap, generation.Used(), - "a smaller first run must retain the future fixed-width growth overlap") - reconcilesBeforeHighWater := generation.ReconcileCount() - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - input.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Less(t, lease.Reserved(), preparedReserved, - "the fixed-width standby must be returned after reaching its high water") - require.Greater(t, generation.ReconcileCount(), reconcilesBeforeHighWater) - reconcilesAfterHighWater := generation.ReconcileCount() - for _, steady := range []*batch.Batch{input, small, input, small} { - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{steady}, - steady.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - require.Equal(t, reservesAtFullBudget, generation.ReserveCount(), - "a prepared recovery run must not reserve after sibling pressure arrives") - require.Equal(t, reconcilesAfterHighWater, generation.ReconcileCount(), - "steady fixed-width evaluation must not reconcile per batch") - - require.True(t, blocker.Release()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRecoveryGrowFailureKeepsPriorGuarantee(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - const budgetCap = uint64(8 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.EnsureRunRecovery(proc, small.RowCount())) - oldReserved := lease.Reserved() - blocker, err := generation.Reserve(budgetCap - generation.Used()) - require.NoError(t, err) - - err = lease.EnsureRunRecovery(proc, colexec.DefaultBatchSize) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, oldReserved, lease.Reserved()) - require.True(t, lease.recoveryReady) - require.Equal(t, small.RowCount(), lease.recoveryRows) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{small}, - small.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - ), "the prior retained-batch bound must remain drainable") - - require.True(t, blocker.Release()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeasePreAdmitsVariableReplacement(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - peak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - require.Less(t, peak, uint64(32<<20)) - - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - require.NoError(t, lease.EnsureRunRecovery(proc, 1)) - require.Equal(t, 2*peak, lease.Reserved(), - "one variable root needs its retained result plus one replacement") - - makeInput := func(value string) *batch.Batch { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()) - bat.SetRowCount(1) - return bat - } - short := makeInput("a") - long := makeInput(strings.Repeat("b", 64<<10)) - defer short.Clean(proc.Mp()) - defer long.Clean(proc.Mp()) - - blocker, err := generation.Reserve(budgetCap - generation.Used()) - require.NoError(t, err) - reservesAtFullBudget := generation.ReserveCount() - for _, input := range []*batch.Batch{short, long} { - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - 1, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - require.Equal(t, reservesAtFullBudget, generation.ReserveCount(), - "variable-size reuse must consume the pre-admitted replacement overlap") - - require.True(t, blocker.Release()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionSerialResultPeakUsesEncodingContract(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - const rows = colexec.DefaultBatchSize - - for _, name := range []string{"serial", "serial_full"} { - t.Run(name, func(t *testing.T) { - expr := makeSerialExpressionLeaseTestExpr( - t, proc, name, types.T_int32.ToType(), types.T_int32.ToType()) - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - resultPeak, err := expressionVarlenaWidthPeak(12, rows) - require.NoError(t, err) - packerPeak, ok := types.PackerCapacityUpperBound(12, 6) - require.True(t, ok) - require.Equal(t, resultPeak+packerPeak, peak) - - generic, err := expressionTypePeak(expr.Typ, rows) - require.NoError(t, err) - require.Less(t, peak, generic) - withReplacement, err := expressionVectorPeak(proc, expr, rows, true) - require.NoError(t, err) - require.Equal(t, peak+resultPeak, withReplacement, - "duplicating the result must not duplicate the sole Packer owner") - require.Less(t, 2*peak, uint64(1<<30), - "a normal serial key must fit an otherwise empty 1 GiB budget") - }) - } -} - -func TestExpressionSerialPackerRecoveryBoundary(t *testing.T) { - const ( - rows = 1 - columns = 4 - width = 65535 - ) - typ := types.New(types.T_varbinary, width, 0) - value := string(make([]byte, width)) - - for _, name := range []string{"serial", "serial_full"} { - for _, delta := range []int64{-1, 0, 1} { - t.Run(name+"/"+map[int64]string{-1: "N-1", 0: "N", 1: "N+1"}[delta], func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - argTypes := make([]types.Type, columns) - for i := range argTypes { - argTypes[i] = typ - } - expr := makeSerialExpressionLeaseTestExpr(t, proc, name, argTypes...) - component, supported := function.SerialEncodedTypeSizeBound(typ) - require.True(t, supported) - payload := uint64(columns) * component - resultPeak, err := expressionVarlenaWidthPeak(payload, rows) - require.NoError(t, err) - packerPeak, ok := types.PackerCapacityUpperBound(payload, component) - require.True(t, ok) - require.Equal(t, uint64(1<<20), packerPeak, - "the bound must follow the real class allocator and append contract") - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - require.Equal(t, resultPeak+packerPeak, peak) - required := 2 * peak - capBytes := uint64(int64(required) + delta) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - - err = lease.EnsureRunRecovery(proc, rows) - if delta < 0 { - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - generation.Close() - proc.Free() - return - } - require.NoError(t, err) - require.Equal(t, required, lease.Reserved()) - blocker, err := generation.Reserve(capBytes - generation.Used()) - require.NoError(t, err) - reservesAtSaturation := generation.ReserveCount() - input := makeSerialExpressionLeaseTestBatchWithColumns( - t, proc, rows, typ, value, columns) - for range 2 { - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - require.Equal(t, reservesAtSaturation, generation.ReserveCount()) - retained, ok := lease.Retained() - require.True(t, ok) - require.Greater(t, retained, resultPeak, - "the physical oracle must include the real retained Packer") - require.LessOrEqual(t, retained, peak) - - input.Clean(proc.Mp()) - require.True(t, blocker.Release()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - generation.Close() - proc.Free() - require.Zero(t, proc.Mp().CurrNB()) - }) - } - } -} - -func TestExpressionSerialPhysicalCapacityBoundAndReuse(t *testing.T) { - const rows = colexec.DefaultBatchSize - for _, name := range []string{"serial", "serial_full"} { - for _, oid := range []types.T{types.T_binary, types.T_varbinary} { - t.Run(name+"/"+oid.String(), func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - typ := types.New(oid, 128, 0) - expr := makeSerialExpressionLeaseTestExpr(t, proc, name, typ, typ) - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - require.LessOrEqual(t, peak, uint64(math.MaxUint64/2)) - budgetCap := 2 * peak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - defer func() { - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - generation.Close() - proc.Free() - require.Zero(t, proc.Mp().CurrNB()) - }() - - input := makeSerialExpressionLeaseTestBatch( - t, proc, rows, typ, string(make([]byte, 128))) - defer input.Clean(proc.Mp()) - eval := func() error { - return lease.Eval( - proc, - []*batch.Batch{input}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - ) - } - - require.NoError(t, eval()) - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, peak, - "the pre-admitted peak must cover the executor's physical capacities") - require.NoError(t, eval(), - "a second batch within the same row/type bound must not fail the ownership audit") - }) - } - } -} - -func TestExpressionSerialPhysicalCapacityRecoveryBoundary(t *testing.T) { - const rows = colexec.DefaultBatchSize - for _, name := range []string{"serial", "serial_full"} { - for _, oid := range []types.T{types.T_binary, types.T_varbinary} { - for _, delta := range []int64{-1, 0, 1} { - t.Run(name+"/"+oid.String()+"/"+ - map[int64]string{-1: "N-1", 0: "N", 1: "N+1"}[delta], func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - typ := types.New(oid, 128, 0) - expr := makeSerialExpressionLeaseTestExpr(t, proc, name, typ, typ) - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - require.LessOrEqual(t, peak, uint64(math.MaxUint64/2)) - required := 2 * peak - require.Positive(t, required) - capBytes := required - if delta < 0 { - capBytes-- - } else { - capBytes += uint64(delta) - } - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - defer func() { - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - generation.Close() - proc.Free() - require.Zero(t, proc.Mp().CurrNB()) - }() - - err = lease.EnsureRunRecovery(proc, rows) - if delta < 0 { - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, types.DefaultPackerCapacity(), generation.Used(), - "failed recovery growth must keep the live constructor allocation owned") - return - } - require.NoError(t, err) - require.Equal(t, required, lease.Reserved()) - blocker, err := generation.Reserve(capBytes - generation.Used()) - require.NoError(t, err) - defer blocker.Release() - reservesAtSaturation := generation.ReserveCount() - - input := makeSerialExpressionLeaseTestBatch( - t, proc, rows, typ, string(make([]byte, 128))) - defer input.Clean(proc.Mp()) - for range 2 { - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - require.Equal(t, reservesAtSaturation, generation.ReserveCount(), - "prepared first evaluation and reuse must not reserve under sibling pressure") - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, peak) - }) - } - } - } -} - -func TestExpressionSerialRecoveryAdmissionAndReplacementMatrix(t *testing.T) { - const rows = colexec.DefaultBatchSize - for _, name := range []string{"serial", "serial_full"} { - t.Run(name, func(t *testing.T) { - for _, delta := range []int64{-1, 0, 1} { - t.Run(map[int64]string{-1: "N-1", 0: "N", 1: "N+1"}[delta], func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - expr := makeSerialExpressionLeaseTestExpr( - t, - proc, - name, - types.New(types.T_varchar, 32, 0), - types.New(types.T_varchar, 32, 0), - ) - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - required := 2 * peak - capBytes := uint64(int64(required) + delta) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - - err = lease.EnsureRunRecovery(proc, rows) - if delta < 0 { - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, types.DefaultPackerCapacity(), generation.Used(), - "failed recovery growth must keep the live constructor allocation owned") - } else { - require.NoError(t, err) - require.Equal(t, required, lease.Reserved()) - blocker, reserveErr := generation.Reserve(capBytes - generation.Used()) - require.NoError(t, reserveErr) - reservesAtSaturation := generation.ReserveCount() - - short := makeSerialExpressionLeaseTestBatch( - t, proc, rows, types.New(types.T_varchar, 32, 0), "a") - wide := makeSerialExpressionLeaseTestBatch( - t, proc, rows, types.New(types.T_varchar, 32, 0), strings.Repeat("x", 32)) - for _, input := range []*batch.Batch{short, wide, short} { - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - require.Equal(t, reservesAtSaturation, generation.ReserveCount(), - "first run, wider replacement, and reuse must consume only pre-admitted ownership") - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved()) - short.Clean(proc.Mp()) - wide.Clean(proc.Mp()) - require.True(t, blocker.Release()) - } - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - generation.Close() - proc.Free() - require.Zero(t, proc.Mp().CurrNB()) - }) - } - }) - } -} - -func TestExpressionSerialRecoverySharedGenerationFanout(t *testing.T) { - const ( - rows = colexec.DefaultBatchSize - workers = 8 - ) - for _, name := range []string{"serial", "serial_full"} { - t.Run(name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - expr := makeSerialExpressionLeaseTestExpr( - t, proc, name, types.T_int32.ToType(), types.T_int32.ToType()) - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - perWorker := 2 * peak - budgetCap, err := spillCheckedMul(perWorker, workers) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - executorSets := make([][]colexec.ExpressionExecutor, workers) - leases := make([]*ExpressionMemoryLease, workers) - for i := range workers { - executorSets[i], leases[i], err = NewBudgetedExpressionExecutors( - proc, generation, []*plan.Expr{expr}, false) - require.NoError(t, err) - require.NoError(t, leases[i].EnsureRunRecovery(proc, rows)) - require.Equal(t, perWorker, leases[i].Reserved()) - } - require.Equal(t, budgetCap, generation.Used(), - "every worker must own its complete replacement peak") - - values := make([]int32, rows) - for i := range values { - values[i] = int32(i) - } - input := batch.NewWithSize(2) - input.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - input.Vecs[1] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - input.SetRowCount(rows) - reservesAtSaturation := generation.ReserveCount() - for i := range workers { - require.NoError(t, leases[i].Eval( - proc, - []*batch.Batch{input}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - require.Equal(t, reservesAtSaturation, generation.ReserveCount(), - "fanout evaluation must not perform a late shared-budget admission") - - input.Clean(proc.Mp()) - for i := range workers { - freeExpressionLeaseTestExecutors(executorSets[i]) - leases[i].Release() - } - require.Zero(t, generation.Used()) - generation.Close() - proc.Free() - require.Zero(t, proc.Mp().CurrNB()) - }) - } -} - -func TestExpressionMemoryLeaseRecoveryUsesSequentialRootPeak(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - peak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - - budget := process.MustNewHashBuildBudget(4*peak, 4*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - exprs := []*plan.Expr{expr, expr} - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, exprs, false) - require.NoError(t, err) - require.Zero(t, lease.Reserved()) - require.NoError(t, lease.EnsureRunRecovery(proc, 1)) - require.Equal(t, 3*peak, lease.Reserved(), - "sequential roots share one replacement standby instead of each reserving two copies") - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRecoveryKeepsPriorRootHighWater(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - largePeak, err := expressionVectorPeak(proc, expr, 2, false) - require.NoError(t, err) - require.Greater(t, largePeak, smallPeak) - - const budgetCap = uint64(128 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - exprs := []*plan.Expr{expr, expr} - executors, lease, err := NewBudgetedExpressionExecutors( - proc, generation, exprs, false) - require.NoError(t, err) - - large := batch.NewWithSize(1) - large.Vecs[0] = testutil.MakeVarcharVector([]string{"a", "b"}, nil, proc.Mp()) - large.SetRowCount(2) - defer large.Clean(proc.Mp()) - stopAfterFirstRoot := errors.New("stop after first root") - err = lease.Run(proc, large.RowCount(), func(index int) error { - _, evalErr := executors[index].Eval( - proc, []*batch.Batch{large}, nil) - if evalErr != nil { - return evalErr - } - return stopAfterFirstRoot - }) - require.ErrorIs(t, err, stopAfterFirstRoot) - require.Equal(t, largePeak, lease.Reserved()) - - require.NoError(t, lease.EnsureRunRecovery(proc, 1)) - require.Equal(t, largePeak+2*smallPeak, lease.Reserved(), - "a smaller re-preflight must retain the prior root plus the next root and its replacement") - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionTypePeakUsesArrayElementWidth(t *testing.T) { - for _, tc := range []struct { - oid types.T - elementWidth uint64 - }{ - {oid: types.T_array_float64, elementWidth: 8}, - {oid: types.T_array_float32, elementWidth: 4}, - {oid: types.T_array_bf16, elementWidth: 2}, - {oid: types.T_array_float16, elementWidth: 2}, - {oid: types.T_array_int8, elementWidth: 1}, - {oid: types.T_array_uint8, elementWidth: 1}, - } { - t.Run(tc.oid.String(), func(t *testing.T) { - peak, err := expressionTypePeak(plan.Type{ - Id: int32(tc.oid), - Width: types.MaxArrayDimension, - }, 1) - require.NoError(t, err) - want, err := expressionVarlenaWidthPeak( - uint64(types.MaxArrayDimension)*tc.elementWidth, 1) - require.NoError(t, err) - require.Equal(t, want, peak) - }) - } -} - -func TestExpressionMemoryLeaseCoversMaxNarrowArrayPayload(t *testing.T) { - for _, tc := range []struct { - oid types.T - makeInput func(*testing.T, *process.Process) *vector.Vector - }{ - { - oid: types.T_array_bf16, - makeInput: func(t *testing.T, proc *process.Process) *vector.Vector { - return makeMaxArrayLeaseTestVector[types.BF16]( - t, proc, types.T_array_bf16) - }, - }, - { - oid: types.T_array_float16, - makeInput: func(t *testing.T, proc *process.Process) *vector.Vector { - return makeMaxArrayLeaseTestVector[types.Float16]( - t, proc, types.T_array_float16) - }, - }, - } { - t.Run(tc.oid.String(), func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - arrayType := plan.Type{ - Id: int32(tc.oid), - Width: types.MaxArrayDimension, - } - condition := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_bool)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - left := &plan.Expr{ - Typ: arrayType, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 1}}, - } - right := &plan.Expr{ - Typ: arrayType, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 2}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "iff", - []*plan.Expr{condition, left, right}, - ) - require.NoError(t, err) - - peak, err := expressionVectorPeak(proc, expr, 2, false) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(2*peak, 2*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - require.NoError(t, err) - - input := batch.NewWithSize(3) - input.Vecs[0] = testutil.MakeBoolVector( - []bool{true, false}, nil, proc.Mp()) - input.Vecs[1] = tc.makeInput(t, proc) - input.Vecs[2] = tc.makeInput(t, proc) - input.SetRowCount(2) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{input}, - input.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved(), - "retained max-width array payload must remain within admission") - - input.Clean(proc.Mp()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - require.Zero(t, proc.Mp().CurrNB()) - }) - } -} - -func TestExpressionMemoryLeaseGrowthRequiresReplacementPeak(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - largePeak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - - t.Run("reject before evaluation", func(t *testing.T) { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.LessOrEqual(t, retained, smallPeak) - budgetCap := retained + largePeak - 1 - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - require.NoError(t, lease.Run(proc, 1, func(_ int) error { return nil })) - require.Equal(t, smallPeak, generation.Used()) - evaluated := false - err = lease.Run(proc, colexec.DefaultBatchSize, func(_ int) error { - evaluated = true - return nil - }) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.False(t, evaluated) - require.Equal(t, smallPeak, generation.Used()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - }) - - t.Run("commit exact replacement peak", func(t *testing.T) { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.LessOrEqual(t, retained, smallPeak) - budgetCap := retained + largePeak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(2) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - require.NoError(t, lease.Run(proc, 1, func(_ int) error { return nil })) - require.NoError(t, lease.Run(proc, colexec.DefaultBatchSize, func(_ int) error { return nil })) - require.Equal(t, largePeak, generation.Used()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) - }) -} - -func TestExpressionMemoryLeaseGrowsRootsIndependently(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - exprs := []*plan.Expr{expr, expr} - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) - require.NoError(t, err) - - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - largePeak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - require.Greater(t, largePeak, smallPeak) - for _, executor := range executors { - retained, ok := colexec.ExpressionExecutorRetainedBytes(executor) - require.True(t, ok) - require.LessOrEqual(t, retained, smallPeak) - } - - // Sequential root growth peaks at old(root 2) + new(root 1) + - // new(root 2). An aggregate replacement would incorrectly also charge - // old(root 1) and reject this exact-cap admission. - budgetCap := smallPeak + 2*largePeak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, exprs, executors, false) - require.NoError(t, err) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{small}, - small.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Equal(t, 2*smallPeak, generation.Used()) - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer large.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{large}, - large.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Equal(t, 2*largePeak, generation.Used()) - require.LessOrEqual(t, generation.Peak(), budgetCap) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseCoversVariableWidthReuseOverlap(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr(proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - peak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - - budget := process.MustNewHashBuildBudget(2*peak, 2*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, - []*plan.Expr{expr}, - executors, - false, - ) - require.NoError(t, err) - - eval := func(value string) { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector([]string{value}, nil, proc.Mp()) - bat.SetRowCount(1) - defer bat.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{bat}, - 1, - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - eval("a") - require.Equal(t, peak, generation.Used()) - retained, ok := lease.Retained() - require.True(t, ok) - require.Positive(t, retained) - reserves := generation.ReserveCount() - - eval(strings.Repeat("b", 64<<10)) - require.Equal(t, reserves+1, generation.ReserveCount(), - "same-row variable-width growth needs a transient overlap reservation") - require.Equal(t, peak, generation.Used(), - "transient overlap must not inflate the retained high-water charge") - require.GreaterOrEqual(t, generation.Peak(), peak+retained) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseCoversFlowControlSelectedScratch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - condition := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_bool)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - makeLower := func(colPos int32) *plan.Expr { - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr(proc.Ctx, "lower", []*plan.Expr{col}) - require.NoError(t, err) - return expr - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "iff", - []*plan.Expr{condition, makeLower(1), makeLower(2)}, - ) - require.NoError(t, err) - peak, err := expressionVectorPeak(proc, expr, 4, false) - require.NoError(t, err) - rootOutput, err := expressionTypePeak(expr.Typ, 4) - require.NoError(t, err) - expectedPeak := rootOutput - for _, branch := range expr.GetF().Args[1:] { - branchOutput, branchErr := expressionTypePeak(branch.Typ, 4) - require.NoError(t, branchErr) - selectedParameter, parameterErr := expressionTypePeak(branch.GetF().Args[0].Typ, 4) - require.NoError(t, parameterErr) - expectedPeak += branchOutput + branchOutput + selectedParameter - } - require.Equal(t, expectedPeak, peak, - "flow-control branches need ordinary output, selected result, and selected parameter capacity") - budget := process.MustNewHashBuildBudget(2*peak, 2*peak) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - require.NoError(t, err) - - eval := func(width int) { - bat := batch.NewWithSize(3) - bat.Vecs[0] = testutil.MakeBoolVector( - []bool{true, false, true, false}, - nil, - proc.Mp(), - ) - left := []string{ - strings.Repeat("A", width), - strings.Repeat("B", width), - strings.Repeat("C", width), - strings.Repeat("D", width), - } - right := []string{ - strings.Repeat("E", width), - strings.Repeat("F", width), - strings.Repeat("G", width), - strings.Repeat("H", width), - } - bat.Vecs[1] = testutil.MakeVarcharVector(left, nil, proc.Mp()) - bat.Vecs[2] = testutil.MakeVarcharVector(right, nil, proc.Mp()) - bat.SetRowCount(4) - defer bat.Clean(proc.Mp()) - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{bat}, - bat.RowCount(), - func(_ int, _ *vector.Vector) error { return nil }, - )) - } - eval(8) - retained, ok := lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved(), - "selected result and parameter scratch must be covered by the admitted peak") - eval(4 << 10) - retained, ok = lease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, lease.Reserved()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseKeepsFixedFlowControlReplacementStandby(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - condition := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_bool)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - makePlusZero := func(colPos int32) *plan.Expr { - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "+", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(0)}, - ) - require.NoError(t, err) - return expr - } - expr, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "iff", - []*plan.Expr{condition, makePlusZero(1), makePlusZero(2)}, - ) - require.NoError(t, err) - - const rows = colexec.DefaultBatchSize - peak, err := expressionVectorPeak(proc, expr, rows, false) - require.NoError(t, err) - budgetCap := 4 * peak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - require.NoError(t, err) - require.NoError(t, lease.EnsureRunRecovery(proc, rows)) - prepared := lease.Reserved() - require.Greater(t, prepared, peak) - reconcilesAfterPrepare := generation.ReconcileCount() - - makeInput := func(trueRows int) *batch.Batch { - mask := make([]bool, rows) - left := make([]int32, rows) - right := make([]int32, rows) - for row := range rows { - mask[row] = row < trueRows - left[row] = int32(row) - right[row] = int32(rows - row) - } - bat := batch.NewWithSize(3) - bat.Vecs[0] = testutil.MakeBoolVector(mask, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeInt32Vector(left, nil, proc.Mp()) - bat.Vecs[2] = testutil.MakeInt32Vector(right, nil, proc.Mp()) - bat.SetRowCount(rows) - return bat - } - mostlyFalse := makeInput(1) - mostlyTrue := makeInput(rows - 1) - defer mostlyFalse.Clean(proc.Mp()) - defer mostlyTrue.Clean(proc.Mp()) - - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{mostlyFalse}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - )) - require.Equal(t, prepared, lease.Reserved(), - "one mask distribution cannot settle another distribution's scratch peak") - require.Equal(t, lease.Reserved(), generation.Used()) - require.Equal(t, reconcilesAfterPrepare, generation.ReconcileCount(), - "a mask-dependent fixed tree must keep its replacement standby") - retainedBefore, ok := lease.Retained() - require.True(t, ok) - - blocker, err := generation.Reserve(budgetCap - generation.Used()) - require.NoError(t, err) - reservesAtFullBudget := generation.ReserveCount() - require.NoError(t, lease.Eval( - proc, - []*batch.Batch{mostlyTrue}, - rows, - func(_ int, _ *vector.Vector) error { return nil }, - )) - retainedAfter, ok := lease.Retained() - require.True(t, ok) - require.Greater(t, retainedAfter, retainedBefore, - "changing only the mask must exercise fixed-width selected-scratch growth") - require.Equal(t, reservesAtFullBudget, generation.ReserveCount()) - require.Equal(t, reconcilesAfterPrepare, generation.ReconcileCount()) - require.LessOrEqual(t, retainedAfter, lease.Reserved()) - - require.True(t, blocker.Release()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestBudgetedExpressionConstructionAdmitsBeforeMpoolAllocation(t *testing.T) { +func TestRecoveryProjectionUsesIncrementalPartialTail(t *testing.T) { mp := mpool.MustNewZero() - proc := testutil.NewProcessWithMPool(t, "", mp) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_Sval{Sval: strings.Repeat("x", 64<<10)}, - }}, - } - initial, err := expressionInitialOwnedBytes(expr) - require.NoError(t, err) - require.Positive(t, initial) - - budget := process.MustNewHashBuildBudget(initial-1, initial-1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - epoch := mp.StartResourcePeakEpoch() - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - peak, ok := mp.EndResourcePeakEpoch(epoch) - require.True(t, ok) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, executors) - require.Nil(t, lease) - require.Zero(t, peak, "budget rejection must happen before constructing the literal vector") - require.Zero(t, generation.Used()) - require.Zero(t, mp.CurrNB()) -} - -func TestBudgetedSerialConstructionAdmitsPackerBeforeAllocation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeSerialExpressionLeaseTestExpr( - t, proc, "serial_full", types.T_int32.ToType(), types.T_int32.ToType()) - initial, err := expressionInitialOwnedBytes(expr) - require.NoError(t, err) - require.Equal(t, types.DefaultPackerCapacity(), initial) - - budget := process.MustNewHashBuildBudget(initial-1, initial-1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - executors, lease, err := NewBudgetedExpressionExecutors( - proc, - generation, - []*plan.Expr{expr}, - false, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, executors) - require.Nil(t, lease) - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRetainsFailedEvaluationBound(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - peak, err := expressionVectorPeak(proc, expr, 32, false) - require.NoError(t, err) - budgetCap := retained + peak - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease(generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - wantErr := errors.New("expression evaluation failed") - require.ErrorIs(t, lease.Run(proc, 32, func(_ int) error { - return wantErr - }), wantErr) - require.Equal(t, peak, generation.Used(), - "a failed evaluator may retain partially grown buffers") - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseDoesNotShrinkAdoptedCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize*8) - defer large.Clean(proc.Mp()) - require.NoError(t, evalExpressionLeaseTestExecutors(proc, executors, large)) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(executors) - require.True(t, ok) - require.Positive(t, retained) - - smallPeak, err := expressionVectorPeak(proc, expr, 1, false) - require.NoError(t, err) - require.Less(t, smallPeak, retained) - budget := process.MustNewHashBuildBudget(retained, retained) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - reservesAfterAdoption := generation.ReserveCount() - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, lease.Run(proc, small.RowCount(), func(_ int) error { - return evalExpressionLeaseTestExecutors(proc, executors, small) - })) - require.Equal(t, reservesAfterAdoption, generation.ReserveCount()) - require.Equal(t, retained, generation.Used()) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRejectsUnknownExecutorOwnership(t *testing.T) { - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - _, err = NewExpressionMemoryLease( - generation, - []*plan.Expr{{Typ: plan.Type{Id: int32(types.T_int32)}}}, - []colexec.ExpressionExecutor{unknownExpressionLeaseExecutor{}}, - false, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Zero(t, generation.Used()) -} - -func TestExpressionMemoryLeaseRejectsInvalidCalls(t *testing.T) { - _, err := NewExpressionMemoryLease( - nil, - []*plan.Expr{{Typ: plan.Type{Id: int32(types.T_int32)}}}, - nil, - false, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - var nilLease *ExpressionMemoryLease - require.ErrorIs(t, - nilLease.Run(nil, 0, func(int) error { return nil }), - process.ErrHashBuildBudgetInvalid, - ) - require.ErrorIs(t, - nilLease.Run(nil, 0, nil), - process.ErrHashBuildBudgetInvalid, - ) - require.ErrorIs(t, - nilLease.Eval(nil, nil, 0, nil), - process.ErrHashBuildBudgetInvalid, - ) - require.Zero(t, nilLease.Reserved()) - require.Zero(t, nilLease.Len()) - retained, ok := nilLease.Retained() - require.True(t, ok) - require.Zero(t, retained) - nilLease.Release() + tail := batch.NewWithSize(1) + tail.Vecs[0] = testutil.MakeInt64Vector([]int64{1, 2, 3}, nil, mp) + tail.SetRowCount(3) + defer tail.Clean(mp) + source := batch.NewWithSize(1) + source.Vecs[0] = testutil.MakeInt64Vector([]int64{4, 5, 6, 7}, nil, mp) + source.SetRowCount(4) + defer source.Clean(mp) - emptyLease := &ExpressionMemoryLease{} - require.ErrorIs(t, - emptyLease.Run(nil, -1, func(int) error { return nil }), - process.ErrHashBuildBudgetInvalid, - ) -} - -func TestExpressionMemoryAccountingHelperBoundaries(t *testing.T) { - size, err := expressionInitialOwnedBytes(nil) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Zero(t, size) - - size, err = literalInitialOwnedBytes(types.T_int32, &plan.Literal{}) - require.NoError(t, err) - require.Zero(t, size) - - require.True(t, expressionExecutorMayGrowWithinBound(nil)) - require.True(t, expressionExecutorMayGrowWithinBound(&plan.Expr{})) - - size, err = initialAllocationCapacity(0) - require.NoError(t, err) - require.Zero(t, size) -} - -func TestExpressionMemoryLeaseEvalPropagatesExecutorError(t *testing.T) { - expected := errors.New("expression evaluation failed") - executor := failingExpressionLeaseExecutor{err: expected} - lease, err := NewExpressionMemoryLease( - nil, - []*plan.Expr{{Typ: plan.Type{Id: int32(types.T_int32)}}}, - []colexec.ExpressionExecutor{executor}, - false, - ) + builder := HashmapBuilder{} + builder.Batches.Buf = []*batch.Batch{tail} + builder.retainedSpillTailSelected = 3 * 8 + projection, err := builder.projectRetainedRecovery(source) require.NoError(t, err) - defer lease.Release() - - consumed := false - err = lease.Eval(nil, nil, 0, func(int, *vector.Vector) error { - consumed = true - return nil - }) - require.ErrorIs(t, err, expected) - require.False(t, consumed) + require.Equal(t, 7, projection.maxRows) + require.Equal(t, uint64(7*8), projection.maxSelected) + require.Equal(t, uint64(7*8), projection.nextTailSelected) } -func TestExpressionMemoryLeaseReleaseIsTerminal(t *testing.T) { +func TestExpressionRecoveryIncludesReplacementOverlap(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - freeExpressionLeaseTestExecutors(executors) - lease.Release() - lease.Release() - require.Zero(t, lease.Reserved()) - _, ok := lease.Retained() - require.False(t, ok) - require.Zero(t, generation.Used()) - - called := false - err = lease.Run(proc, 1, func(_ int) error { - called = true - return nil - }) - require.ErrorIs(t, err, process.ErrHashBuildReservationInactive) - require.False(t, called) - require.Zero(t, generation.Used()) - - executors, err = colexec.NewExpressionExecutorsFromPlanExpressions( - proc, []*plan.Expr{expr}) + expr := makeIssue26454ConcatKey(t, proc) + peak, err := expressionVectorPeak(proc, expr, 1024, false) require.NoError(t, err) - lease, err = NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) + recovery, err := expressionRecoveryBytes(proc, []*plan.Expr{expr}, 1024, false) require.NoError(t, err) - require.NoError(t, lease.Run(proc, 1, func(_ int) error { - freeExpressionLeaseTestExecutors(executors) - lease.Release() - return nil - })) - require.Zero(t, generation.Used(), - "release during evaluation must also discard the pending replacement") -} - -func TestExpressionMemoryLeaseCancellationAndRepeatedGenerations(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := makeExpressionLeaseTestExpr(t, proc) - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - - runGeneration := func(id uint64, callbackErr error) { - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions( - proc, []*plan.Expr{expr}) - require.NoError(t, err) - generation, err := budget.OpenGeneration(id) - require.NoError(t, err) - lease, err := NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - input := makeExpressionLeaseTestBatch(proc, 128) - - err = lease.Run(proc, input.RowCount(), func(_ int) error { - if evalErr := evalExpressionLeaseTestExecutors( - proc, executors, input); evalErr != nil { - return evalErr - } - return callbackErr - }) - if callbackErr != nil { - require.ErrorIs(t, err, callbackErr) - } else { - require.NoError(t, err) - } - require.Positive(t, generation.Used()) - - input.Clean(proc.Mp()) - freeExpressionLeaseTestExecutors(executors) - lease.Release() - generation.Close() - require.Zero(t, generation.Used()) - require.Zero(t, proc.Mp().CurrNB()) - } - - runGeneration(1, context.Canceled) - runGeneration(2, nil) -} - -type unknownExpressionLeaseExecutor struct{} - -func (unknownExpressionLeaseExecutor) Eval(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (unknownExpressionLeaseExecutor) EvalWithoutResultReusing(*process.Process, []*batch.Batch, []bool) (*vector.Vector, error) { - return nil, nil -} -func (unknownExpressionLeaseExecutor) ResetForNextQuery() {} -func (unknownExpressionLeaseExecutor) Free() {} -func (unknownExpressionLeaseExecutor) IsColumnExpr() bool { return false } -func (unknownExpressionLeaseExecutor) TypeName() string { return "unknown" } - -type failingExpressionLeaseExecutor struct { - unknownExpressionLeaseExecutor - err error -} - -func (f failingExpressionLeaseExecutor) Eval( - *process.Process, - []*batch.Batch, - []bool, -) (*vector.Vector, error) { - return nil, f.err + require.Equal(t, 2*peak, recovery) } diff --git a/pkg/sql/colexec/hashbuild/expression_recovery.go b/pkg/sql/colexec/hashbuild/expression_recovery.go new file mode 100644 index 0000000000000..d17aa1410f925 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/expression_recovery.go @@ -0,0 +1,387 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/plan/function" + "github.com/matrixorigin/matrixone/pkg/vm/process" +) + +// expressionRecoveryBytes bounds all retained expression roots plus one +// allocate-copy-free replacement. Roots run sequentially, so one largest-root +// overlap is sufficient even when several roots retain reusable vectors. +func expressionRecoveryBytes( + proc *process.Process, + exprs []*plan.Expr, + rows int, + duplicate bool, +) (uint64, error) { + if rows < 0 || len(exprs) == 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + var total, replacement uint64 + for _, expr := range exprs { + peak, err := expressionVectorPeak(proc, expr, rows, duplicate) + if err != nil { + return 0, err + } + total, err = recoveryCheckedAdd(total, peak) + if err != nil { + return 0, err + } + if peak > replacement { + replacement = peak + } + } + return recoveryCheckedAdd(total, replacement) +} + +// expressionVectorPeak is an execution-before-allocation upper bound for one +// expression executor tree based on the declared SQL result types. +func expressionVectorPeak( + proc *process.Process, + expr *plan.Expr, + rows int, + duplicate bool, +) (uint64, error) { + if expr == nil || rows < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + total, root, err := expressionTreePeakWithSelection( + proc, expr, uint64(rows), false) + if err != nil { + return 0, err + } + if duplicate { + return recoveryCheckedAdd(total, root) + } + return total, nil +} + +func expressionTreePeakWithSelection( + proc *process.Process, + expr *plan.Expr, + rows uint64, + mayReceivePartialSelection bool, +) (total uint64, output uint64, err error) { + if expr == nil { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + switch node := expr.Expr.(type) { + case *plan.Expr_Col: + return 0, 0, nil + case *plan.Expr_F: + if node.F == nil { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + fid := int32(-1) + if node.F.Func != nil { + fid, _ = function.DecodeOverloadID(node.F.Func.Obj) + } + for index, arg := range node.F.Args { + child, _, childErr := expressionTreePeakWithSelection( + proc, + arg, + rows, + expressionChildMayReceivePartialSelection( + fid, index, mayReceivePartialSelection), + ) + if childErr != nil { + return 0, 0, childErr + } + total, err = recoveryCheckedAdd(total, child) + if err != nil { + return 0, 0, err + } + } + case *plan.Expr_P: + if node.P == nil || proc == nil || proc.GetPrepareParams() == nil { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + paramPeak, paramErr := expressionParamPeak(proc, node.P.Pos) + if paramErr != nil { + return 0, 0, paramErr + } + typePeak, typeErr := expressionTypePeak(expr.Typ, 1) + if typeErr != nil { + return 0, 0, typeErr + } + output = max(paramPeak, typePeak) + return output, output, nil + case *plan.Expr_Lit, *plan.Expr_V, *plan.Expr_Raw, *plan.Expr_Vec, + *plan.Expr_Fold, *plan.Expr_T: + // Leaf executors may materialize their declared output below. + default: + return 0, 0, process.ErrHashBuildBudgetInvalid + } + + output, err = expressionResultPeak(expr, rows) + if err != nil { + return 0, 0, err + } + total, err = recoveryCheckedAdd(total, output) + if err != nil { + return 0, 0, err + } + private, err := expressionFunctionPrivatePeak(expr) + if err != nil { + return 0, 0, err + } + total, err = recoveryCheckedAdd(total, private) + if err != nil { + return 0, 0, err + } + + if _, isFunction := expr.Expr.(*plan.Expr_F); mayReceivePartialSelection && isFunction { + total, err = recoveryCheckedAdd(total, output) + if err != nil { + return 0, 0, err + } + for _, arg := range nodeFunctionArgs(expr) { + switch arg.Expr.(type) { + case *plan.Expr_Col, *plan.Expr_F: + selected, selectedErr := expressionTypePeak(arg.Typ, rows) + if selectedErr != nil { + return 0, 0, selectedErr + } + total, err = recoveryCheckedAdd(total, selected) + if err != nil { + return 0, 0, err + } + } + } + } + return total, output, nil +} + +func expressionResultPeak(expr *plan.Expr, rows uint64) (uint64, error) { + fn, ok := expr.Expr.(*plan.Expr_F) + if !ok || fn.F == nil || fn.F.Func == nil { + return expressionTypePeak(expr.Typ, rows) + } + fid, _ := function.DecodeOverloadID(fn.F.Func.Obj) + if fid != function.SERIAL && fid != function.SERIAL_FULL { + return expressionTypePeak(expr.Typ, rows) + } + payload, _, supported, err := serialExpressionPackerBounds(fn.F) + if err != nil { + return 0, err + } + if !supported { + return expressionTypePeak(expr.Typ, rows) + } + return expressionVarlenaWidthPeak(payload, rows) +} + +func expressionFunctionPrivatePeak(expr *plan.Expr) (uint64, error) { + fn, ok := expr.Expr.(*plan.Expr_F) + if !ok || fn.F == nil || fn.F.Func == nil { + return 0, nil + } + fid, _ := function.DecodeOverloadID(fn.F.Func.Obj) + if fid != function.SERIAL && fid != function.SERIAL_FULL { + return 0, nil + } + payload, maxAppend, supported, err := serialExpressionPackerBounds(fn.F) + if err != nil { + return 0, err + } + if !supported { + return types.DefaultPackerCapacity(), nil + } + capacity, ok := types.PackerCapacityUpperBound(payload, maxAppend) + if !ok { + return 0, process.ErrHashBuildBudgetInvalid + } + return capacity, nil +} + +func serialExpressionPackerBounds(fn *plan.Function) ( + payload uint64, + maxAppend uint64, + supported bool, + err error, +) { + if fn == nil { + return 0, 0, false, process.ErrHashBuildBudgetInvalid + } + for _, arg := range fn.Args { + if arg == nil { + return 0, 0, false, process.ErrHashBuildBudgetInvalid + } + component, ok := function.SerialEncodedTypeSizeBound(types.New( + types.T(arg.Typ.Id), arg.Typ.Width, arg.Typ.Scale, + )) + if !ok { + return 0, 0, false, nil + } + payload, err = recoveryCheckedAdd(payload, component) + if err != nil { + return 0, 0, false, err + } + maxAppend = max(maxAppend, component) + } + return payload, maxAppend, true, nil +} + +func nodeFunctionArgs(expr *plan.Expr) []*plan.Expr { + if node, ok := expr.Expr.(*plan.Expr_F); ok && node.F != nil { + return node.F.Args + } + return nil +} + +func expressionChildMayReceivePartialSelection( + fid int32, + argument int, + parentMayReceivePartialSelection bool, +) bool { + switch fid { + case function.IFF, function.CASE, function.COALESCE: + return parentMayReceivePartialSelection || argument > 0 + default: + return parentMayReceivePartialSelection + } +} + +func expressionParamPeak(proc *process.Process, pos int32) (uint64, error) { + value, err := proc.GetPrepareParamsAt(int(pos)) + if err != nil { + return 0, err + } + if value == nil { + return 0, nil + } + header, ok := mpool.GrowCapacity(0, int64(types.VarlenaSize)) + if !ok || header < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + peak := uint64(header) + if len(value) <= types.VarlenaInlineSize { + return peak, nil + } + area, ok := mpool.GrowCapacity(0, int64(len(value))) + if !ok || area < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + return recoveryCheckedAdd(peak, uint64(area)) +} + +func expressionTypePeak(typ plan.Type, rows uint64) (uint64, error) { + oid := types.T(typ.Id) + width := int64(oid.FixedLength()) + if width >= 0 { + return expressionFixedWidthPeak(uint64(max(width, 1)), rows) + } + width = int64(typ.Width) + hardMax := int64(types.MaxVarcharLen) + if oid.IsArrayRelate() { + elementWidth := int64(oid.ToType().GetArrayElementSize()) + width *= elementWidth + hardMax = int64(types.MaxArrayDimension) * elementWidth + } else { + switch oid { + case types.T_blob, types.T_text, types.T_json, types.T_datalink, + types.T_geometry, types.T_geometry32: + hardMax = int64(types.MaxBlobLen) + } + } + if width > hardMax { + hardMax = width + } + return expressionVarlenaWidthPeak(uint64(max(hardMax, 1)), rows) +} + +const ( + expressionPerRowAllowance = uint64(32) + expressionAllocationSlack = uint64(64 << 10) +) + +func expressionAllocationCapacityUpperBound(required uint64) (uint64, error) { + if required == 0 { + return 0, nil + } + if mpool.CapLimit <= 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + limit := uint64(mpool.CapLimit) + if required >= limit { + return limit, nil + } + capacity, ok := mpool.GrowCapacity(int64(required-1), int64(required)) + if !ok || capacity < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + return uint64(capacity), nil +} + +func expressionFixedWidthPeak(width, rows uint64) (uint64, error) { + data, err := recoveryCheckedMul(max(width, 1), rows) + if err != nil { + return 0, err + } + data, err = expressionAllocationCapacityUpperBound(data) + if err != nil { + return 0, err + } + allowance, err := recoveryCheckedMul(rows, expressionPerRowAllowance) + if err != nil { + return 0, err + } + total, err := recoveryCheckedAdd(data, allowance) + if err != nil { + return 0, err + } + return recoveryCheckedAdd(total, expressionAllocationSlack) +} + +func expressionVarlenaWidthPeak(width, rows uint64) (uint64, error) { + descriptors, err := recoveryCheckedMul(rows, uint64(types.VarlenaSize)) + if err != nil { + return 0, err + } + descriptors, err = expressionAllocationCapacityUpperBound(descriptors) + if err != nil { + return 0, err + } + var area uint64 + if width > uint64(types.VarlenaInlineSize) { + area, err = recoveryCheckedMul(width, rows) + if err != nil { + return 0, err + } + area, err = expressionAllocationCapacityUpperBound(area) + if err != nil { + return 0, err + } + } + metadata, err := recoveryCheckedMul( + rows, expressionPerRowAllowance-uint64(types.VarlenaSize)) + if err != nil { + return 0, err + } + total, err := recoveryCheckedAdd(descriptors, area) + if err != nil { + return 0, err + } + total, err = recoveryCheckedAdd(total, metadata) + if err != nil { + return 0, err + } + return recoveryCheckedAdd(total, expressionAllocationSlack) +} diff --git a/pkg/sql/colexec/hashbuild/expression_test_helpers_test.go b/pkg/sql/colexec/hashbuild/expression_test_helpers_test.go new file mode 100644 index 0000000000000..88bec2b987dd3 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/expression_test_helpers_test.go @@ -0,0 +1,94 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func makeIssue26454ConcatKey(t testing.TB, proc *process.Process) *plan.Expr { + t.Helper() + cast := func(colPos int32) *plan.Expr { + col := &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_int32)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: colPos}}, + } + targetType := plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "cast", + []*plan.Expr{ + col, + { + Typ: targetType, + Expr: &plan.Expr_T{T: &plan.TargetType{}}, + }, + }, + ) + require.NoError(t, err) + return expr + } + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "concat", + []*plan.Expr{ + cast(0), + plan2.MakePlan2StringConstExprWithType("-"), + cast(1), + }, + ) + require.NoError(t, err) + return expr +} + +func makeIssue26454CaseKey(t testing.TB, proc *process.Process) *plan.Expr { + t.Helper() + column := &plan.Expr{ + Typ: plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + }, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + } + condition, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "=", + []*plan.Expr{ + column, + plan2.MakePlan2StringConstExprWithType("ATM_CON"), + }, + ) + require.NoError(t, err) + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "case", + []*plan.Expr{ + condition, + plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_HEADERS"), + plan2.MakePlan2StringConstExprWithType("CON_CONTRACT_DOC"), + }, + ) + require.NoError(t, err) + return expr +} diff --git a/pkg/sql/colexec/hashbuild/hashmap.go b/pkg/sql/colexec/hashbuild/hashmap.go index dc2c14f750562..58f524b56b978 100644 --- a/pkg/sql/colexec/hashbuild/hashmap.go +++ b/pkg/sql/colexec/hashbuild/hashmap.go @@ -31,7 +31,6 @@ import ( "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/runtimefilter" - "github.com/matrixorigin/matrixone/pkg/sql/plan/function" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -74,18 +73,10 @@ type HashmapBuilder struct { dedupDeleteKeepColIdxList []int32 DelRows *bitmap.Bitmap budget *process.HashBuildBudgetGeneration - mapReservation *hashMapReservationOwner - batchReservations []*process.HashBuildReservation - // retainedSpillTailSelected is the logical materialized size of the one - // partial CopyIntoBatches tail. Varlena descriptors may share one physical - // payload while a later spill selection repeats it per logical row; tracking - // the logical sum incrementally avoids both an unsafe allocation proxy and - // repeatedly rescanning the growing tail. - retainedSpillTailSelected uint64 - auxReservation *process.HashBuildReservation keyExprs []*plan.Expr - expressionLease *ExpressionMemoryLease - + // retainedSpillTailSelected is the logical spill materialization of the + // one partial CopyIntoBatches tail. It avoids rescanning that growing tail. + retainedSpillTailSelected uint64 // Exact runtime-filter keys are an optional owner inside the mandatory // JoinMap build. The fallback bit is observed by HashBuild for diagnostics. // @@ -96,6 +87,12 @@ type HashmapBuilder struct { // must not be retried or re-spilled. runtimeFilterCollectionFallback bool retainedBatchRecoverySafe bool + mapAllocationAccount *mpool.AllocationAccount + mapAllocation *hashtable.AllocationAccountSelection + iteratorAllocation *hashmap.IteratorAllocation + batchAllocation *vector.AllocationAccountSelection + uniqueKeyAllocation *vector.AllocationAccountSelection + recoveryCapacityClass mpool.AllocationCapacityClass } func (hb *HashmapBuilder) GetSize() int64 { @@ -126,7 +123,8 @@ func (hb *HashmapBuilder) GetJoinMap(mp *mpool.MPool) *message.JoinMap { } sels := hb.Sels hb.Sels = message.GroupSels{} - jm := message.NewJoinMap(sels, hb.IntHashMap, hb.StrHashMap, hb.DelRows, hb.Batches.Buf, mp) + jmDelRows := hb.DelRows + jm := message.NewJoinMap(sels, hb.IntHashMap, hb.StrHashMap, jmDelRows, hb.Batches.Buf, mp) jm.SetHasNullKey(hb.HasNullKey) hb.IntHashMap = nil hb.StrHashMap = nil @@ -134,15 +132,12 @@ func (hb *HashmapBuilder) GetJoinMap(mp *mpool.MPool) *message.JoinMap { hb.Batches.Reset() hb.retainedSpillTailSelected = 0 // Iterators are producer scratch and are not part of JoinMap ownership. - // Drop budgeted cached backing before transferring the encompassing aux - // reservation to a consumer that may free it immediately after publication. hb.detachAndPruneCachedIterators() - hb.IgnoreRows = nil + hb.freeIgnoreRows(mp) hb.uniqueSels = nil hb.curVecs = nil - release := hb.detachReservations() jm.SetMemoryRelease(func() { - release() + releaseDedupBitmap(jmDelRows, mp) }) return jm } @@ -164,7 +159,20 @@ func (hb *HashmapBuilder) observeNullKeys(keyVecs []*vector.Vector) { return } for _, vec := range keyVecs { - if vec.HasNull() { + if vec == nil { + continue + } + rows := uint64(vec.Length()) + if vec.IsConstNull() { + if vec.GetGrouping().GetBitmap().CountRange(0, rows) < vec.Length() { + hb.HasNullKey = true + return + } + continue + } + if vec.GetNulls().GetBitmap().AnySetNotIn( + vec.GetGrouping().GetBitmap(), 0, rows, + ) { hb.HasNullKey = true return } @@ -190,11 +198,11 @@ func (hb *HashmapBuilder) Prepare( } keyWidth += width } - executors, expressionLease, err := NewBudgetedExpressionExecutors( + executors, err := newExpressionExecutorsWithCapacityClass( proc, - hb.budget, keyCols, - needDupVec, + hb.mapAllocationAccount, + hb.recoveryCapacityClass, ) if err != nil { return err @@ -202,7 +210,6 @@ func (hb *HashmapBuilder) Prepare( hb.needDupVec = needDupVec hb.executors = executors hb.keyExprs = keyCols - hb.expressionLease = expressionLease hb.keyWidth = keyWidth hb.InputBatchRowCount = 0 hb.hashMapRowCount = 0 @@ -237,8 +244,8 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.retainedSpillTailSelected = 0 hb.IntHashMap = nil hb.StrHashMap = nil - hb.IgnoreRows = nil - hb.DelRows = nil + hb.freeIgnoreRows(proc.Mp()) + hb.freeDelRows(proc.Mp()) for i := range hb.UniqueJoinKeys { if hb.UniqueJoinKeys[i] != nil { hb.UniqueJoinKeys[i].Free(proc.Mp()) @@ -247,8 +254,8 @@ func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) hb.UniqueJoinKeys = nil hb.uniqueKeySlots = nil // Function executors retain result-vector capacity across ResetForNextQuery. - // Free them before releasing expression reservations; Prepare recreates the - // executor set for the next generation. + // Destroy them here; immutable allocation selections remain installed until + // the statement lifecycle calls ClearAllocationAccount. hb.FreeExecutors() } @@ -258,6 +265,8 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { hb.cachedStrIterator = nil hb.FreeHashMapAndBatches(proc) hb.FreeTemporaryVectors(proc) + hb.freeIgnoreRows(proc.Mp()) + hb.freeDelRows(proc.Mp()) hb.needDupVec = false hb.HasNullKey = false hb.Batches.Reset() @@ -272,6 +281,8 @@ func (hb *HashmapBuilder) Free(proc *process.Process) { } hb.UniqueJoinKeys = nil hb.uniqueKeySlots = nil + hb.runtimeFilterCollectionFallback = false + hb.retainedBatchRecoverySafe = false } func (hb *HashmapBuilder) FreeExecutors() { @@ -282,7 +293,6 @@ func (hb *HashmapBuilder) FreeExecutors() { } hb.executors = nil hb.keyExprs = nil - hb.releaseExpressionLease() } func (hb *HashmapBuilder) FreeTemporaryVectors(proc *process.Process) { @@ -308,7 +318,8 @@ func (hb *HashmapBuilder) FreeHashMapAndBatches(proc *process.Process) { hb.Sels.Free(proc.Mp()) hb.Batches.Clean(proc.Mp()) hb.retainedSpillTailSelected = 0 - hb.releaseReservations() + hb.freeIgnoreRows(proc.Mp()) + hb.freeDelRows(proc.Mp()) } // evalBatch evaluates join key expressions for one batch, storing results in hb.curVecs. @@ -341,13 +352,9 @@ func (hb *HashmapBuilder) evalBatch(batchIdx int, proc *process.Process) error { return nil } var err error - if hb.expressionLease != nil { - err = hb.expressionLease.Run(proc, bat.RowCount(), evalOne) - } else { - for idx := range hb.executors { - if err = evalOne(idx); err != nil { - break - } + for idx := range hb.executors { + if err = evalOne(idx); err != nil { + break } } if err != nil { @@ -359,402 +366,42 @@ func (hb *HashmapBuilder) evalBatch(batchIdx int, proc *process.Process) error { func (hb *HashmapBuilder) abortExpressionEval(proc *process.Process) { // Eval may allocate cached child/result vectors before returning an error. - // Destroy the complete executor tree before releasing its retained lease. + // Destroy the complete executor tree so every exact allocation is released. hb.FreeTemporaryVectors(proc) hb.FreeExecutors() } -// expressionVectorPeak is an execution-before-allocation upper bound based on -// the SQL result type. Varlena widths use the declared maximum (or the engine -// maximum when absent), so input-dependent expanding functions are rejected -// by admission before Eval instead of allocating first. -func expressionVectorPeak(proc *process.Process, expr *plan.Expr, rows int, duplicate bool) (uint64, error) { - if expr == nil || rows < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - total, root, err := expressionTreePeak(proc, expr, uint64(rows)) - if err != nil { - return 0, err - } - if duplicate { - if total > math.MaxUint64-root { - return 0, process.ErrHashBuildBudgetInvalid - } - total += root - } - return total, nil -} - -// ExpressionVectorPeak exposes the same execution-before-allocation bound used -// by HashmapBuilder to spill/re-spill callers. Expression evaluators cache -// intermediate and result vectors, so callers must keep the returned amount -// reserved until the corresponding executor tree is freed or evaluated again -// under a replacement reservation. -func ExpressionVectorPeak(proc *process.Process, expr *plan.Expr, rows int, duplicate bool) (uint64, error) { - return expressionVectorPeak(proc, expr, rows, duplicate) -} - -func expressionTreePeak(proc *process.Process, expr *plan.Expr, rows uint64) (total uint64, output uint64, err error) { - return expressionTreePeakWithSelection(proc, expr, rows, false) -} - -func expressionTreePeakWithSelection( - proc *process.Process, - expr *plan.Expr, - rows uint64, - mayReceivePartialSelection bool, -) (total uint64, output uint64, err error) { - if expr == nil { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - switch node := expr.Expr.(type) { - case *plan.Expr_Col: - return 0, 0, nil - case *plan.Expr_F: - if node.F == nil { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - var fid int32 = -1 - if node.F.Func != nil { - fid, _ = function.DecodeOverloadID(node.F.Func.Obj) - } - for i, arg := range node.F.Args { - child, _, childErr := expressionTreePeakWithSelection( - proc, - arg, - rows, - expressionChildMayReceivePartialSelection( - fid, i, mayReceivePartialSelection), - ) - if childErr != nil || total > math.MaxUint64-child { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += child - } - case *plan.Expr_P: - if node.P == nil || proc == nil || proc.GetPrepareParams() == nil { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - paramPeak, paramErr := expressionParamPeak(proc, node.P.Pos) - if paramErr != nil { - return 0, 0, paramErr - } - typePeak, typeErr := expressionTypePeak(expr.Typ, 1) - if typeErr != nil { - return 0, 0, typeErr - } - if paramPeak > typePeak { - output = paramPeak - } else { - output = typePeak - } - return output, output, nil - case *plan.Expr_Lit, *plan.Expr_V, *plan.Expr_Raw, *plan.Expr_Vec, *plan.Expr_Fold, *plan.Expr_T: - // These executors may materialize a vector but have no child expression - // tree. Expr_T is the target-type argument used by CAST/bit_cast and is - // evaluated as a fixed vector. Charge their declared output below. - default: - // Window, subquery, correlated, list and max nodes do not - // expose a bounded vector-evaluator tree here. - return 0, 0, process.ErrHashBuildBudgetInvalid - } - output, err = expressionResultPeak(expr, rows) - if err != nil || total > math.MaxUint64-output { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += output - private, privateErr := expressionFunctionPrivatePeak(expr) - if privateErr != nil || total > math.MaxUint64-private { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += private - - if _, isFunction := expr.Expr.(*plan.Expr_F); mayReceivePartialSelection && isFunction { - // A partially selected function retains both its ordinary full-row - // result and a selected-result scratch vector. Row-aligned column and - // non-folded function parameters are also copied into retained selected - // parameter vectors before the function executes. - if total > math.MaxUint64-output { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += output - for _, arg := range nodeFunctionArgs(expr) { - switch arg.Expr.(type) { - case *plan.Expr_Col, *plan.Expr_F: - selectedParameter, selectedErr := expressionTypePeak(arg.Typ, rows) - if selectedErr != nil || total > math.MaxUint64-selectedParameter { - return 0, 0, process.ErrHashBuildBudgetInvalid - } - total += selectedParameter - } - } - } - return total, output, nil -} - -// expressionResultPeak keeps the generic SQL-type bound for ordinary -// functions, but lets functions with a stronger allocation contract provide a -// tighter result-vector bound. serial and serial_full are the first such -// functions: their encoded result is the sum of the component encodings, not -// an arbitrary VARCHAR(max) value. Their retained Packer is charged separately -// by expressionFunctionPrivatePeak so duplicate-result ownership does not -// duplicate the sole function operator. -func expressionResultPeak(expr *plan.Expr, rows uint64) (uint64, error) { - if expr == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - fn, ok := expr.Expr.(*plan.Expr_F) - if !ok || fn.F == nil || fn.F.Func == nil { - return expressionTypePeak(expr.Typ, rows) - } - fid, _ := function.DecodeOverloadID(fn.F.Func.Obj) - if fid != function.SERIAL && fid != function.SERIAL_FULL { - return expressionTypePeak(expr.Typ, rows) - } - - payloadPerRow, _, supported, err := serialExpressionPackerBounds(fn.F) - if err != nil { - return 0, err - } - if !supported { - // Keep the pre-existing representation-independent bound if the - // encoder contract does not recognize a planner type. - return expressionTypePeak(expr.Typ, rows) - } - return expressionVarlenaWidthPeak(payloadPerRow, rows) -} - -func expressionFunctionPrivatePeak(expr *plan.Expr) (uint64, error) { - fn, ok := expr.Expr.(*plan.Expr_F) - if !ok || fn.F == nil || fn.F.Func == nil { - return 0, nil - } - fid, _ := function.DecodeOverloadID(fn.F.Func.Obj) - if fid != function.SERIAL && fid != function.SERIAL_FULL { - return 0, nil - } - payload, maxAppend, supported, err := serialExpressionPackerBounds(fn.F) - if err != nil { - return 0, err - } - if !supported { - // getPackFun resolves every component before encoding the first row. - // An unsupported component can therefore retain only the constructor's - // initial Packer allocation before Eval fails. - return types.DefaultPackerCapacity(), nil - } - capacity, ok := types.PackerCapacityUpperBound(payload, maxAppend) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - return capacity, nil -} - -// serialExpressionPackerBounds returns both the maximum encoded row length and -// the largest single append issued by its component encoders. A component -// cannot append more in one call than its complete encoded-size bound, so the -// maximum component bound is also a representation-independent append bound. -func serialExpressionPackerBounds(fn *plan.Function) ( - payload uint64, - maxAppend uint64, - supported bool, - err error, -) { - if fn == nil { - return 0, 0, false, process.ErrHashBuildBudgetInvalid - } - for _, arg := range fn.Args { - if arg == nil { - return 0, 0, false, process.ErrHashBuildBudgetInvalid +// hasGroupingKey reports whether a direct build key column contains any +// GROUPING sentinel in any retained batch. Such maps use the string encoder's +// explicit key domains; an IntHashMap cannot represent a sentinel outside the +// complete uint64 raw-value domain without collisions. +func (hb *HashmapBuilder) hasGroupingKey() bool { + for _, executor := range hb.executors { + if !executor.IsColumnExpr() { + continue } - component, ok := function.SerialEncodedTypeSizeBound(types.New( - types.T(arg.Typ.Id), arg.Typ.Width, arg.Typ.Scale, - )) + column, ok := executor.(*colexec.ColumnExpressionExecutor) if !ok { - return 0, 0, false, nil - } - if payload > math.MaxUint64-component { - return 0, 0, false, process.ErrHashBuildBudgetInvalid - } - payload += component - if component > maxAppend { - maxAppend = component - } - } - return payload, maxAppend, true, nil -} - -func nodeFunctionArgs(expr *plan.Expr) []*plan.Expr { - if node, ok := expr.Expr.(*plan.Expr_F); ok && node.F != nil { - return node.F.Args - } - return nil -} - -// expressionParamPeak returns an upper bound for the allocations made by a -// non-null ParamExpressionExecutor. Params are materialized as one-element -// const vectors, whose data is one varlena header and whose area is allocated -// only for payloads that do not fit in that header. -func expressionParamPeak(proc *process.Process, pos int32) (uint64, error) { - val, err := proc.GetPrepareParamsAt(int(pos)) - if err != nil { - return 0, err - } - if val == nil { - return 0, nil - } - - headerCap, ok := mpool.GrowCapacity(0, int64(types.VarlenaSize)) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - peak := uint64(headerCap) - if len(val) <= types.VarlenaInlineSize { - return peak, nil - } - - areaCap, ok := mpool.GrowCapacity(0, int64(len(val))) - if !ok || areaCap < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - if uint64(areaCap) > math.MaxUint64-peak { - return 0, process.ErrHashBuildBudgetInvalid - } - return peak + uint64(areaCap), nil -} - -func expressionTypePeak(typ plan.Type, rows uint64) (uint64, error) { - oid := types.T(typ.Id) - width := int64(oid.FixedLength()) - if width >= 0 { - if width < 1 { - width = 1 - } - return expressionFixedWidthPeak(uint64(width), rows) - } - - width = int64(typ.Width) - hardMax := int64(types.MaxVarcharLen) - if oid.IsArrayRelate() { - elementWidth := int64(oid.ToType().GetArrayElementSize()) - width *= elementWidth - hardMax = int64(types.MaxArrayDimension) * elementWidth - } else { - switch oid { - case types.T_blob, types.T_text, types.T_json, types.T_datalink, - types.T_geometry, types.T_geometry32: - hardMax = int64(types.MaxBlobLen) - } - } - if width > hardMax { - // Never clamp a declared bound downward. Array width is declared - // in elements, while every other varlena width is in bytes. - hardMax = width - } - width = hardMax - if width < 1 { - width = 1 - } - return expressionVarlenaWidthPeak(uint64(width), rows) -} - -const ( - expressionPerRowAllowance = uint64(32) - expressionAllocationSlack = uint64(64 << 10) -) - -// expressionAllocationCapacityUpperBound bounds the capacity retained after -// any sequence of GrowCapacity calls whose logical requirement never exceeds -// required. The last growth either allocates required directly or starts from -// a capacity below required. GrowCapacity's single-step growth is monotonic in -// that starting capacity, so required-1 covers every incremental append -// history without replaying an O(rows) growth sequence during HashBuild. -func expressionAllocationCapacityUpperBound(required uint64) (uint64, error) { - if required == 0 { - return 0, nil - } - if mpool.CapLimit <= 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - limit := uint64(mpool.CapLimit) - if required >= limit { - // A successful mpool allocation cannot retain more than CapLimit. - // Values which exceed the allocator's own limit still fail in Eval. - return limit, nil - } - capacity, ok := mpool.GrowCapacity(int64(required-1), int64(required)) - if !ok || capacity < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - return uint64(capacity), nil -} - -func expressionFixedWidthPeak(width, rows uint64) (uint64, error) { - if width == 0 { - width = 1 - } - if rows != 0 && width > math.MaxUint64/rows { - return 0, process.ErrHashBuildBudgetInvalid - } - dataCapacity, err := expressionAllocationCapacityUpperBound(width * rows) - if err != nil || rows > math.MaxUint64/expressionPerRowAllowance { - return 0, process.ErrHashBuildBudgetInvalid - } - allowance := rows * expressionPerRowAllowance - if dataCapacity > math.MaxUint64-allowance || - dataCapacity+allowance > math.MaxUint64-expressionAllocationSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return dataCapacity + allowance + expressionAllocationSlack, nil -} - -func expressionVarlenaWidthPeak(width, rows uint64) (uint64, error) { - if width == 0 { - width = 1 - } - descriptorWidth := uint64(types.VarlenaSize) - if descriptorWidth > expressionPerRowAllowance || - (rows != 0 && descriptorWidth > math.MaxUint64/rows) { - return 0, process.ErrHashBuildBudgetInvalid - } - dataCapacity, err := expressionAllocationCapacityUpperBound( - rows * descriptorWidth) - if err != nil { - return 0, err - } - - var areaCapacity uint64 - if width > uint64(types.VarlenaInlineSize) { - if rows != 0 && width > math.MaxUint64/rows { - return 0, process.ErrHashBuildBudgetInvalid - } - areaCapacity, err = expressionAllocationCapacityUpperBound(width * rows) - if err != nil { - return 0, err + continue + } + // HashBuild evaluates every key against one retained build batch. + // Some join planners preserve the original build relation index in a + // direct column expression, but ColumnExpressionExecutor deliberately + // resolves that expression against the only input batch. Inspect the + // same physical column here; filtering on RelPos would miss GROUPING + // sentinels for DedupJoin and RightDedupJoin. + position := column.GetColIndex() + for _, bat := range hb.Batches.Buf { + if bat != nil && position >= 0 && position < len(bat.Vecs) && + bat.Vecs[position] != nil && + bat.Vecs[position].GetGrouping().GetBitmap().CountRange( + 0, uint64(bat.Vecs[position].Length()), + ) > 0 { + return true + } } } - - // The historical per-row allowance included the varlena descriptor. Its - // physical capacity is now charged above, leaving the non-mpool metadata - // allowance unchanged instead of double-counting the descriptor. - metadataPerRow := expressionPerRowAllowance - descriptorWidth - if rows != 0 && metadataPerRow > math.MaxUint64/rows { - return 0, process.ErrHashBuildBudgetInvalid - } - metadata := rows * metadataPerRow - if dataCapacity > math.MaxUint64-areaCapacity || - dataCapacity+areaCapacity > math.MaxUint64-metadata || - dataCapacity+areaCapacity+metadata > math.MaxUint64-expressionAllocationSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return dataCapacity + areaCapacity + metadata + expressionAllocationSlack, nil -} - -func (hb *HashmapBuilder) releaseExpressionLease() { - if hb.expressionLease != nil { - hb.expressionLease.Release() - hb.expressionLease = nil - } + return false } func (hb *HashmapBuilder) BuildHashmap(hashOnPK bool, needAllocateSels bool, needUniqueVec bool, proc *process.Process) (retErr error) { @@ -803,28 +450,11 @@ func (hb *HashmapBuilder) buildHashmap( if hb.InputBatchRowCount == 0 { return nil } - if err := hb.reserveBuildAux(needUniqueVec); err != nil { - if !needUniqueVec { - return err - } - if runtimefilter.ClassifyOptionalFallback(err) != - runtimefilter.OptionalFallbackBudgetAdmission { - return err - } - // The extra auxiliary charge exists only for optional exact-filter key - // retention. Retry the admission in place without that owner before - // allocating or mutating the mandatory map. - needUniqueVec = false - if err = hb.reserveBuildAux(false); err != nil { - return err - } - // Linearize the fallback only after mandatory admission succeeds. A - // failed retry is a fatal build, not a successful optional downgrade. - hb.runtimeFilterCollectionFallback = true - } dedupBuildKeepLast = dedupBuildKeepLast && hb.IsDedup && hb.OnDuplicateAction == plan.Node_FAIL defer func() { if retErr != nil { + hashmap.IteratorClearOwner(hb.cachedIntIterator) + hashmap.IteratorClearOwner(hb.cachedStrIterator) hb.cachedIntIterator = nil hb.cachedStrIterator = nil } @@ -840,17 +470,23 @@ func (hb *HashmapBuilder) buildHashmap( var err error var itr hashmap.Iterator - if hb.keyWidth <= 8 { - if err = hb.reserveInitialMap(int64(hashtable.Int64HashMapInitialAllocationBytes())); err != nil { - return err - } - if hb.IntHashMap, err = hashmap.NewIntHashMap(false, proc.Mp()); err != nil { - hb.releaseMapReservation() - return err - } - if err = hb.attachIntHashMapAdmission(hb.IntHashMap); err != nil { - hb.IntHashMap.Free() - hb.IntHashMap = nil + hasGroupingKey := hb.hasGroupingKey() + useIntHashMap := hb.keyWidth <= 8 && !hasGroupingKey + if hb.mapAllocation == nil || hb.mapAllocationAccount == nil || + hb.iteratorAllocation == nil || hb.batchAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } + if useIntHashMap { + hb.IntHashMap, err = hashmap.NewIntHashMapWithAllocation( + false, + proc.Mp(), + hb.mapAllocation, + ) + if err != nil { + if hb.IntHashMap != nil { + hb.IntHashMap.Free() + hb.IntHashMap = nil + } return err } if hb.cachedIntIterator != nil { @@ -861,17 +497,23 @@ func (hb *HashmapBuilder) buildHashmap( hb.cachedIntIterator = itr } } else { - if err = hb.reserveInitialMap(int64(hashtable.StringHashMapInitialAllocationBytes())); err != nil { - return err - } - if hb.StrHashMap, err = hashmap.NewStrHashMap(false, proc.Mp()); err != nil { - hb.releaseMapReservation() + hb.StrHashMap, err = hashmap.NewStrHashMapWithAllocations( + false, + proc.Mp(), + hb.mapAllocation, + hb.iteratorAllocation, + ) + if err != nil { + if hb.StrHashMap != nil { + hb.StrHashMap.Free() + hb.StrHashMap = nil + } return err } - if err = hb.attachStrHashMapAdmission(hb.StrHashMap); err != nil { - hb.StrHashMap.Free() - hb.StrHashMap = nil - return err + if hasGroupingKey { + if err = hb.StrHashMap.SetGroupingAware(); err != nil { + return err + } } if hb.cachedStrIterator != nil { hashmap.IteratorChangeOwner(hb.cachedStrIterator, hb.StrHashMap) @@ -884,7 +526,7 @@ func (hb *HashmapBuilder) buildHashmap( if hashOnPK || hb.IsDedup { // if hash on primary key, prealloc hashmap size to the count of batch - if hb.keyWidth <= 8 { + if useIntHashMap { err = hb.IntHashMap.PreAlloc(uint64(hb.InputBatchRowCount)) if err != nil { return err @@ -898,14 +540,37 @@ func (hb *HashmapBuilder) buildHashmap( } if needAllocateSels { - if err := hb.Sels.Init(hb.InputBatchRowCount, proc.Mp()); err != nil { + err = hb.Sels.InitWithAllocation( + hb.InputBatchRowCount, + proc.Mp(), + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteGroupSels, + ) + if err != nil { return err } } if hb.IsDedup && (hb.OnDuplicateAction == plan.Node_IGNORE || dedupBuildKeepLast) { - hb.IgnoreRows = &bitmap.Bitmap{} - hb.IgnoreRows.InitWithSize(int64(hb.InputBatchRowCount)) + hb.IgnoreRows, err = hb.newDedupBitmap( + hb.InputBatchRowCount, + proc.Mp(), + HashBuildAllocationSiteDedupIgnoreBitmap, + ) + if err != nil { + return err + } + } + if hb.delColIdx != -1 && hb.DelRows == nil { + hb.DelRows, err = hb.newDedupBitmap( + hb.InputBatchRowCount, + proc.Mp(), + HashBuildAllocationSiteDedupDeleteBitmap, + ) + if err != nil { + return err + } } var ( @@ -920,21 +585,55 @@ func (hb *HashmapBuilder) buildHashmap( ignoreCandidateOwnsKey []bool ignoreCandidateOldKey []*vector.Vector ) + cleanupDedupScratch := func() { + freeDedupSlice(hb, lastRows, proc.Mp()) + lastRows = nil + freeDedupSlice(hb, ignoreSurvivorRows, proc.Mp()) + ignoreSurvivorRows = nil + freeDedupSlice(hb, ignoreSurvivorOwnsKey, proc.Mp()) + ignoreSurvivorOwnsKey = nil + } + defer cleanupDedupScratch() if dedupBuildKeepLast { - lastRows = make([]int64, hb.InputBatchRowCount+1) + lastRows, err = makeDedupSlice[int64]( + hb, + hb.InputBatchRowCount+1, + proc.Mp(), + HashBuildAllocationSiteDedupLastRows, + ) + if err != nil { + return err + } for i := range lastRows { lastRows[i] = -1 } } if hb.IsDedup && hb.OnDuplicateAction == plan.Node_IGNORE && hb.delColIdx >= 0 { - ignoreSurvivorRows = make([]int64, hb.InputBatchRowCount+1) - ignoreSurvivorOwnsKey = make([]bool, hb.InputBatchRowCount+1) + ignoreSurvivorRows, err = makeDedupSlice[int64]( + hb, + hb.InputBatchRowCount+1, + proc.Mp(), + HashBuildAllocationSiteDedupSurvivorRows, + ) + if err != nil { + return err + } + ignoreSurvivorOwnsKey, err = makeDedupSlice[bool]( + hb, + hb.InputBatchRowCount+1, + proc.Mp(), + HashBuildAllocationSiteDedupSurvivorOwnsKey, + ) + if err != nil { + return err + } ignoreBuildGroups = make([]uint64, hashmap.UnitLimit) ignoreBuildZvals = make([]int64, hashmap.UnitLimit) ignoreCandidateOwnsKey = make([]bool, hashmap.UnitLimit) ignoreCandidateOldKey = make([]*vector.Vector, 1) } +buildUnits: for i := 0; i < hb.InputBatchRowCount; i += hashmap.UnitLimit { if i%(hashmap.UnitLimit*32) == 0 { if err := checkHashBuildCanceled(proc); err != nil { @@ -950,7 +649,7 @@ func (hb *HashmapBuilder) buildHashmap( // if not hash on primary key, estimate the hashmap size after 8192 rows //preAlloc to improve performance and reduce memory reAlloc if !hashOnPK && !hb.IsDedup && hb.InputBatchRowCount > hashmap.HashMapSizeThreshHold && i == hashmap.HashMapSizeEstimate { - if hb.keyWidth <= 8 { + if useIntHashMap { groupCount := hb.IntHashMap.GroupCount() rate := float64(groupCount) / float64(i) hashmapCount := uint64(float64(hb.InputBatchRowCount) * rate) @@ -995,7 +694,14 @@ func (hb *HashmapBuilder) buildHashmap( zvals = ignoreBuildZvals[:n] clear(ignoreCandidateOwnsKey[:n]) ignoreCandidateOldKey[0] = hb.Batches.Buf[vecIdx1].Vecs[hb.delColIdx] - oldVals, oldZvals := itr.Find(vecIdx2, n, ignoreCandidateOldKey) + oldVals, oldZvals, findErr := itr.Find( + vecIdx2, + n, + ignoreCandidateOldKey, + ) + if findErr != nil { + return findErr + } for k := 0; k < n; k++ { ignoreCandidateOwnsKey[k] = zvals[k] != 0 && oldZvals[k] != 0 && vals[k] != 0 && oldVals[k] == vals[k] } @@ -1080,11 +786,31 @@ func (hb *HashmapBuilder) buildHashmap( if needUniqueVec { if len(hb.UniqueJoinKeys) == 0 { + if hb.uniqueKeyAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } hb.UniqueJoinKeys = make([]*vector.Vector, len(hb.executors)) for j, vec := range hb.curVecs { - if hb.collectUniqueKeySlot(j) { - hb.UniqueJoinKeys[j] = - vector.NewOffHeapVecWithType(*vec.GetType()) + if !hb.collectUniqueKeySlot(j) { + continue + } + hb.UniqueJoinKeys[j], err = vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), + hb.uniqueKeyAllocation, + ) + if err != nil { + cause := err + if mpool.IsRetryableAllocationCapacity(err) { + cause = runtimefilter.MarkOptionalAllocationError(err) + } + if fatalErr := hb.fallbackOptionalRuntimeFilterCollection( + proc, + cause, + ); fatalErr != nil { + return fatalErr + } + needUniqueVec = false + continue buildUnits } } } @@ -1094,32 +820,12 @@ func (hb *HashmapBuilder) buildHashmap( if !hb.collectUniqueKeySlot(j) { continue } - areaBytes, reserveErr := - unionBatchAreaBytes(vec, vecIdx2, n) - if reserveErr != nil { - // Range and overflow failures contradict the collection - // oracle; they are never optional allocation failures. - return reserveErr - } - overlap, reserveErr := hb.reserveUniqueAppendOverlap(hb.UniqueJoinKeys[j], n, areaBytes) - if reserveErr != nil { - if fatalErr := - hb.fallbackOptionalRuntimeFilterCollection( - proc, reserveErr); fatalErr != nil { - return fatalErr - } - needUniqueVec = false - break - } err = hb.UniqueJoinKeys[j].UnionBatch(vec, int64(vecIdx2), n, nil, proc.Mp()) - if overlap != nil { - overlap.Release() - } if err != nil { - // With the range and capacity oracle above satisfied, - // UnionBatch error returns are only mpool growth failures. - allocationErr := - runtimefilter.MarkOptionalAllocationError(err) + allocationErr := err + if mpool.IsRetryableAllocationCapacity(err) { + allocationErr = runtimefilter.MarkOptionalAllocationError(err) + } if fatalErr := hb.fallbackOptionalRuntimeFilterCollection( proc, allocationErr); fatalErr != nil { @@ -1146,31 +852,12 @@ func (hb *HashmapBuilder) buildHashmap( if !hb.collectUniqueKeySlot(j) { continue } - areaBytes, reserveErr := uniqueAppendAreaBytes(vec, 0, len(newSels), newSels) - if reserveErr != nil { - // Selector/range/overflow failures are collection - // contract errors and remain fatal. - return reserveErr - } - overlap, reserveErr := hb.reserveUniqueAppendOverlap(hb.UniqueJoinKeys[j], len(newSels), areaBytes) - if reserveErr != nil { - if fatalErr := - hb.fallbackOptionalRuntimeFilterCollection( - proc, reserveErr); fatalErr != nil { - return fatalErr - } - needUniqueVec = false - break - } err = hb.UniqueJoinKeys[j].Union(vec, newSels, proc.Mp()) - if overlap != nil { - overlap.Release() - } if err != nil { - // With generated selectors and the capacity oracle above - // satisfied, Union error returns are mpool growth failures. - allocationErr := - runtimefilter.MarkOptionalAllocationError(err) + allocationErr := err + if mpool.IsRetryableAllocationCapacity(err) { + allocationErr = runtimefilter.MarkOptionalAllocationError(err) + } if fatalErr := hb.fallbackOptionalRuntimeFilterCollection( proc, allocationErr); fatalErr != nil { @@ -1204,6 +891,7 @@ func (hb *HashmapBuilder) buildHashmap( hb.InputBatchRowCount = totalRowCount } hb.hashMapRowCount = hb.InputBatchRowCount + cleanupDedupScratch() hb.resetHashStateForRebuild(proc) needUniqueVec, err = hb.prepareCanonicalRuntimeFilterCollection( runtimeFilterRequested) @@ -1232,7 +920,8 @@ func (hb *HashmapBuilder) buildHashmap( } hb.InputBatchRowCount = hb.Batches.RowCount() hb.hashMapRowCount = hb.InputBatchRowCount - hb.DelRows = nil + cleanupDedupScratch() + hb.freeDelRows(proc.Mp()) hb.resetHashStateForRebuild(proc) needUniqueVec, err = hb.prepareCanonicalRuntimeFilterCollection( runtimeFilterRequested) @@ -1244,8 +933,18 @@ func (hb *HashmapBuilder) buildHashmap( if hb.delColIdx != -1 { if hb.DelRows == nil { - hb.DelRows = &bitmap.Bitmap{} - hb.DelRows.InitWithSize(int64(max(cardinality, uint64(hb.Batches.RowCount())))) + delRows := max(cardinality, uint64(hb.Batches.RowCount())) + if delRows > uint64(math.MaxInt) { + return process.ErrHashBuildBudgetInvalid + } + hb.DelRows, err = hb.newDedupBitmap( + int(delRows), + proc.Mp(), + HashBuildAllocationSiteDedupDeleteBitmap, + ) + if err != nil { + return err + } } // Scan every build row, including the delete-only rows appended by @@ -1282,7 +981,14 @@ func (hb *HashmapBuilder) buildHashmap( if err = hb.evalBatch(vecIdx1, proc); err != nil { return err } - newVals, newZvals := itr.Find(vecIdx2, n, hb.curVecs) + newVals, newZvals, findErr := itr.Find( + vecIdx2, + n, + hb.curVecs, + ) + if findErr != nil { + return findErr + } for k := 0; k < n; k++ { buildGroups[k] = 0 if newZvals[k] != 0 { @@ -1291,7 +997,10 @@ func (hb *HashmapBuilder) buildHashmap( } } tmpVecs[0] = hb.Batches.Buf[vecIdx1].Vecs[hb.delColIdx] - vals, zvals := itr.Find(vecIdx2, n, tmpVecs) + vals, zvals, findErr := itr.Find(vecIdx2, n, tmpVecs) + if findErr != nil { + return findErr + } for k, v := range vals[:n] { if zvals[k] == 0 || v == 0 { @@ -1334,7 +1043,6 @@ func (hb *HashmapBuilder) resetHashStateForRebuild(proc *process.Process) { hb.StrHashMap.Free() hb.StrHashMap = nil } - hb.releaseMapReservation() hb.Sels.Free(proc.Mp()) for i := range hb.UniqueJoinKeys { if hb.UniqueJoinKeys[i] != nil { @@ -1348,20 +1056,15 @@ func (hb *HashmapBuilder) resetHashStateForRebuild(proc *process.Process) { hb.executors[i].ResetForNextQuery() } } - hb.IgnoreRows = nil + hb.freeIgnoreRows(proc.Mp()) } // FreeHashMapOnly discards a partial hash build while preserving the copied -// build batches and their reservations. It is the supported transition from a -// failed BuildHashmap attempt to either a less memory-intensive rebuild or -// bounded spill recovery. +// build batches for bounded spill recovery. It is the only supported +// transition from a failed BuildHashmap attempt to re-spill. func (hb *HashmapBuilder) FreeHashMapOnly(proc *process.Process) { hb.resetHashStateForRebuild(proc) - hb.DelRows = nil - if hb.auxReservation != nil { - hb.auxReservation.Release() - hb.auxReservation = nil - } + hb.freeDelRows(proc.Mp()) } func (hb *HashmapBuilder) keepDiscardedRowsForDelete(proc *process.Process) error { @@ -1369,11 +1072,19 @@ func (hb *HashmapBuilder) keepDiscardedRowsForDelete(proc *process.Process) erro return hb.Batches.Shrink(hb.IgnoreRows, proc) } - activeRows := hb.IgnoreRows.Clone() - activeRows.Negate() - activeCount := activeRows.Count() + activeCount := int(hb.IgnoreRows.Len()) - hb.IgnoreRows.Count() - discardedWithDeletes := make([]int32, 0, hb.IgnoreRows.Count()) + discardedStorage, err := makeDedupSlice[int32]( + hb, + hb.IgnoreRows.Count(), + proc.Mp(), + HashBuildAllocationSiteDedupDiscardedRows, + ) + if err != nil { + return err + } + defer freeDedupSlice(hb, discardedStorage, proc.Mp()) + discardedWithDeletes := discardedStorage[:0] itr := hb.IgnoreRows.Iterator() for itr.HasNext() { row := itr.Next() @@ -1416,8 +1127,19 @@ func (hb *HashmapBuilder) keepDiscardedRowsForDelete(proc *process.Process) erro return err } - hb.DelRows = &bitmap.Bitmap{} - hb.DelRows.InitWithSize(int64(activeCount + len(discardedWithDeletes))) + newRows := activeCount + len(discardedWithDeletes) + if hb.DelRows == nil { + hb.DelRows, err = hb.newDedupBitmap( + newRows, + proc.Mp(), + HashBuildAllocationSiteDedupDeleteBitmap, + ) + if err != nil { + return err + } + } else { + hb.DelRows.InitWithSize(int64(newRows)) + } for i := range discardedWithDeletes { hb.DelRows.Add(uint64(activeCount + i)) } @@ -1431,9 +1153,29 @@ func (hb *HashmapBuilder) makeDeleteOnlyBatch(rows []int32, proc *process.Proces } bat := batch.NewOffHeapWithSize(len(hb.Batches.Buf[0].Vecs)) + if hb.mapAllocationAccount == nil { + bat.Clean(proc.Mp()) + return nil, mpool.ErrAllocationAccountInvalid + } + selection, err := vector.NewAllocationAccountSelection( + hb.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildAllocationSiteDedupDeleteOnlyData, + HashBuildAllocationSiteDedupDeleteOnlyArea, + HashBuildAllocationSiteDedupDeleteOnlyNulls, + HashBuildAllocationSiteDedupDeleteOnlyGrouping, + ) + if err != nil { + bat.Clean(proc.Mp()) + return nil, err + } + if err = bat.SetAllocationAccount(selection); err != nil { + bat.Clean(proc.Mp()) + return nil, err + } bat.Attrs = hb.Batches.Buf[0].Attrs for colIdx, vec := range hb.Batches.Buf[0].Vecs { - bat.Vecs[colIdx] = vector.NewOffHeapVecWithType(*vec.GetType()) + bat.SetVector(int32(colIdx), vector.NewOffHeapVecWithType(*vec.GetType())) } cleanOnErr := true @@ -1491,16 +1233,10 @@ func (hb *HashmapBuilder) detachAndPruneCachedIterators() { } if hb.cachedStrIterator != nil { if hashmap.StrIteratorCapacity(hb.cachedStrIterator) > hashmap.MaxStrIteratorCapacity { + hashmap.IteratorClearOwner(hb.cachedStrIterator) hb.cachedStrIterator = nil return } hashmap.IteratorClearOwner(hb.cachedStrIterator) } - if hb.budget != nil { - // Budgeted builds charge iterator scratch only for the execution that - // allocated it. Do not retain Go backing arrays in the pooled operator - // after that reservation is released or transferred. - hb.cachedIntIterator = nil - hb.cachedStrIterator = nil - } } diff --git a/pkg/sql/colexec/hashbuild/hashmap_test.go b/pkg/sql/colexec/hashbuild/hashmap_test.go index faa9fc9e30dd0..37d6bc517f148 100644 --- a/pkg/sql/colexec/hashbuild/hashmap_test.go +++ b/pkg/sql/colexec/hashbuild/hashmap_test.go @@ -15,9 +15,7 @@ package hashbuild import ( - "context" - "errors" - "math" + "fmt" "reflect" "strconv" "strings" @@ -32,1283 +30,666 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/runtimefilter" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) -func TestBuildHashMap(t *testing.T) { - var hb HashmapBuilder - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - err := hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc) - require.NoError(t, err) - - inputBatch := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, int(100000), proc.Mp()) - err = hb.Batches.CopyIntoBatches(inputBatch, proc) - hb.InputBatchRowCount = inputBatch.RowCount() - inputBatch.Clean(proc.Mp()) - require.NoError(t, err) - - err = hb.BuildHashmap(false, true, true, proc) - require.NoError(t, err) - require.Less(t, int64(0), hb.GetSize()) - require.Less(t, uint64(0), hb.GetGroupCount()) - hb.Reset(proc, true) - hb.Free(proc) - require.Equal(t, int64(0), proc.Mp().CurrNB()) -} - -func TestBuildHashmapOptionalAuxClosedBudgetRemainsFatal(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - generation.Close() - - hb := HashmapBuilder{InputBatchRowCount: 1} - hb.setBudget(generation) - err = hb.BuildHashmap(false, false, true, proc) - require.Error(t, err) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, err, &budgetErr) - require.Equal(t, process.HashBuildBudgetErrorClosed, budgetErr.Kind) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback) - require.Nil(t, hb.UniqueJoinKeys) - require.Zero(t, generation.Used()) - hb.Free(proc) -} - -func TestBuildHashmapMandatoryAuxRetryFailureDoesNotRecordFallback(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - budget := process.MustNewHashBuildBudget(1, 1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - hb := HashmapBuilder{InputBatchRowCount: 1} - hb.setBudget(generation) - err = hb.BuildHashmap(false, false, true, proc) - require.Error(t, err) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, err, &budgetErr) - require.Equal(t, process.HashBuildBudgetErrorAdmission, budgetErr.Kind) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback, - "fatal mandatory retry must not be counted as an optional fallback") - require.Equal(t, uint64(2), generation.RejectCount()) - require.Zero(t, generation.Used()) - hb.Free(proc) -} - -func TestPrepareCanonicalRuntimeFilterCollectionClosedBudgetRemainsFatal( - t *testing.T, -) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 16, proc.Mp()) - defer input.Clean(proc.Mp()) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - hb := HashmapBuilder{ - Batches: colexec.Batches{Buf: []*batch.Batch{input}}, - InputBatchRowCount: input.RowCount(), - } - hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(false)) - generation.Close() - - collect, err := hb.prepareCanonicalRuntimeFilterCollection(true) - require.Error(t, err) - require.False(t, collect) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, err, &budgetErr) - require.Equal(t, process.HashBuildBudgetErrorClosed, budgetErr.Kind) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback) - hb.Batches.Buf = nil - hb.releaseReservations() - require.Zero(t, generation.Used()) -} - -func TestOptionalRuntimeFilterCollectionCleanupFailureRemainsFatal( - t *testing.T, -) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - key := testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - hb := HashmapBuilder{ - InputBatchRowCount: 1, - UniqueJoinKeys: []*vector.Vector{key}, +func TestObserveNullKeysUsesColumnLevelGroupingSentinel(t *testing.T) { + mp := mpool.MustNewZero() + newVec := func(nullRows, groupingRows []uint64) *vector.Vector { + vec := vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList( + vec, + []int32{1, 2}, + nil, + mp, + )) + for _, row := range nullRows { + vec.GetNulls().Add(row) + } + for _, row := range groupingRows { + vec.GetGrouping().Add(row) + } + return vec } - hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(true)) - require.True(t, hb.auxReservation.Release()) - err = hb.fallbackOptionalRuntimeFilterCollection( - proc, - runtimefilter.MarkOptionalAllocationError( - errors.New("mpool allocation failed")), - ) - require.ErrorIs(t, err, process.ErrHashBuildReservationInactive) - fallback, _ := hb.runtimeFilterFallbackState() - require.False(t, fallback) - require.Nil(t, hb.UniqueJoinKeys) - require.Zero(t, generation.Used()) - hb.releaseReservations() -} - -func TestBuildHashmapUniqueUnionAllocationFailureFallsBack(t *testing.T) { - for _, test := range []struct { - name string - hashOnPK bool + tests := []struct { + name string + nullRows []uint64 + groupingRows []uint64 + want bool }{ - {name: "union"}, - {name: "union-batch", hashOnPK: true}, - } { + { + name: "grouping sentinel masks null in same row", + nullRows: []uint64{0}, + groupingRows: []uint64{0}, + want: false, + }, + { + name: "full grouping is sentinel", + nullRows: []uint64{0, 1}, + groupingRows: []uint64{0, 1}, + want: false, + }, + { + name: "partial grouping without null", + groupingRows: []uint64{0}, + want: false, + }, + { + name: "null outside partial grouping is retained", + nullRows: []uint64{1}, + groupingRows: []uint64{0}, + want: true, + }, + } + for _, test := range tests { t.Run(test.name, func(t *testing.T) { - testBuildHashmapUniqueUnionAllocationFailureFallsBack( - t, test.hashOnPK) + vec := newVec(test.nullRows, test.groupingRows) + defer vec.Free(mp) + builder := HashmapBuilder{TrackNullKeys: true} + builder.observeNullKeys([]*vector.Vector{vec}) + require.Equal(t, test.want, builder.HasNullKey) }) } } -func testBuildHashmapUniqueUnionAllocationFailureFallsBack( - t *testing.T, - hashOnPK bool, -) { - mp, err := mpool.NewMPool(t.Name(), 8<<20, mpool.NoFixed) - require.NoError(t, err) - proc := testutil.NewProcessWithMPool(t, "", mp) - - var hb HashmapBuilder - require.NoError(t, hb.Prepare( - []*plan.Expr{newExpr(0, types.T_int32.ToType())}, - -1, -1, nil, proc)) - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 16, mp) - hb.Batches.Buf = []*batch.Batch{input} - hb.InputBatchRowCount = input.RowCount() - - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb.setBudget(generation) - - var filler []byte - defer func() { - hb.Free(proc) - require.Zero(t, generation.Used()) - if filler != nil { - mp.Free(filler) - } - generation.Close() - proc.Free() - require.Zero(t, mp.CurrNB()) - }() - - // Calibrate the deterministic mandatory map footprint, then leave exactly - // that much headroom. The second build can recreate its required map, while - // the first optional-key Union allocation must fail at the mpool boundary. - retainedBytes := mp.CurrNB() - require.NoError(t, hb.BuildHashmap(hashOnPK, false, false, proc)) - mapBytes := mp.CurrNB() - retainedBytes - require.Greater(t, mapBytes, int64(0)) - hb.FreeHashMapOnly(proc) - require.Equal(t, retainedBytes, mp.CurrNB()) - - fillerBytes := mp.Cap() - mp.CurrNB() - mapBytes - require.Greater(t, fillerBytes, int64(0)) - filler, err = mp.Alloc(int(fillerBytes), true) - require.NoError(t, err) - - require.NoError(t, hb.BuildHashmap(hashOnPK, false, true, proc)) - fallback, rebuildSafe := hb.runtimeFilterFallbackState() - require.True(t, fallback) - require.True(t, rebuildSafe) - require.Nil(t, hb.UniqueJoinKeys) - require.Greater(t, hb.GetGroupCount(), uint64(0)) - require.Zero(t, generation.RejectCount(), - "mpool failure must not be misclassified as budget admission") -} - -func TestBuildHashMapBudgetRejectsResizeAndReleasesOnReset(t *testing.T) { - const budgetCap = uint64(1 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) - - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 10_000, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - - err = hb.BuildHashmap(false, false, false, proc) - require.Error(t, err) - require.True(t, errors.Is(err, process.ErrHashBuildBudgetAdmission)) - require.Greater(t, generation.Used(), uint64(0)) - - hb.Reset(proc, true) - require.Zero(t, generation.Used()) -} - -func TestBuildHashMapCancellationReleasesRetainedBudgetOnReset(t *testing.T) { - const budgetCap = uint64(16 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) - - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 10_000, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - require.Positive(t, generation.Used(), "retained build input must own budget before cancellation") - - proc.Cancel(context.Canceled) - err = hb.BuildHashmap(false, false, false, proc) - require.ErrorIs(t, err, context.Canceled) - require.Nil(t, hb.IntHashMap) - require.Nil(t, hb.StrHashMap) - - hb.Reset(proc, true) - hb.Free(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestPublishedJoinMapResizeKeepsReservationWithConsumer(t *testing.T) { - const budgetCap = uint64(16 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 100, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - require.NoError(t, hb.BuildHashmap(false, false, false, proc)) - - jm := hb.GetJoinMap(proc.Mp()) - require.NotNil(t, jm) - jm.IncRef(1) - usedBeforeResize := generation.Used() - secondGeneration, err := budget.OpenGeneration(2) - require.NoError(t, err) - hb.setBudget(secondGeneration) - require.NoError(t, jm.PreAlloc(100_000)) - require.Greater(t, generation.Used(), usedBeforeResize) - require.Zero(t, secondGeneration.Used(), "published map must retain its original generation") - jm.Free() - require.Zero(t, generation.Used()) - - hb.Reset(proc, false) -} - -func TestHashMapReservationOwnerRetainsSegmentedGrowthTokens(t *testing.T) { - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - initial, err := generation.Reserve(100) - require.NoError(t, err) - owner := &hashMapReservationOwner{tokens: []*process.HashBuildReservation{initial}} - - incremental, err := generation.Reserve(50) - require.NoError(t, err) - (&hashMapResizeReservation{owner: owner, token: incremental}).Commit( - hashtable.ResizePlan{ReuseCurrentBlocks: true}, - ) - require.Equal(t, uint64(150), generation.Used()) - - replacement, err := generation.Reserve(200) - require.NoError(t, err) - (&hashMapResizeReservation{owner: owner, token: replacement}).Commit(hashtable.ResizePlan{}) - require.Equal(t, uint64(200), generation.Used()) - - owner.release() - require.Zero(t, generation.Used()) -} - -func TestBudgetedEmptyJoinMapRejectsUnadmittedAllocationAndResize(t *testing.T) { - for _, tc := range []struct { - name string - keyWidth int - initialBytes uint64 +func TestBuildHashmapPreservesRowwiseGroupingAcrossCopiedBatchMerge(t *testing.T) { + for _, test := range []struct { + name string + groupFirst bool + columns int }{ - {name: "int", keyWidth: 4, initialBytes: hashtable.Int64HashMapInitialAllocationBytes()}, - {name: "string", keyWidth: 128, initialBytes: hashtable.StringHashMapInitialAllocationBytes()}, - } { - t.Run(tc.name, func(t *testing.T) { - mp := mpool.MustNewZero() - - tooSmall := process.MustNewHashBuildBudget(tc.initialBytes-1, tc.initialBytes-1) - tooSmallGeneration, err := tooSmall.OpenGeneration(1) - require.NoError(t, err) - jm, err := NewBudgetedEmptyJoinMap(tc.keyWidth, tooSmallGeneration, mp) - require.Nil(t, jm) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Zero(t, tooSmallGeneration.Used()) - require.Zero(t, mp.CurrNB()) - - exact := process.MustNewHashBuildBudget(tc.initialBytes, tc.initialBytes) - generation, err := exact.OpenGeneration(2) - require.NoError(t, err) - jm, err = NewBudgetedEmptyJoinMap(tc.keyWidth, generation, mp) - require.NoError(t, err) - require.Equal(t, tc.initialBytes, generation.Used()) - require.Equal(t, int64(tc.initialBytes), mp.CurrNB()) - - err = jm.PreAlloc(10_000) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, tc.initialBytes, generation.Used(), - "rejected growth must roll back its temporary reservation") - - jm.Free() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, mp.CurrNB()) - }) - } -} - -func TestCopyBuildBatchBudgetsSmallIngressAfterFullBatches(t *testing.T) { - const budgetCap = uint64(32 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for _, rows := range []int{colexec.DefaultBatchSize, colexec.DefaultBatchSize} { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 100, proc.Mp()) - projected, err := hb.projectedBatchCopyBytes(input) - require.NoError(t, err) - destination, err := projectedNewDestinationBytes(input, 0, input.RowCount()) - require.NoError(t, err) - metadata, ok := retainedMetadataAllowance(input) - require.True(t, ok) - require.Equal(t, destination+metadata+uint64(64<<10), projected, - "a small ingress must not be projected as a full 8192-row allocation") - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - - require.Len(t, hb.Batches.Buf, 3) - require.Equal(t, 100, hb.Batches.Buf[2].RowCount()) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchBudgetsPartialTailGrowth(t *testing.T) { - const budgetCap = uint64(32 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for range 1000 { - // Deep spill partitions contain many tiny records. They coalesce into - // one physical batch whose vector capacity grows geometrically. - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 7, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 1) - require.Equal(t, 7000, hb.Batches.Buf[0].RowCount()) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchBudgetsWideVarcharPartialTailReplacement(t *testing.T) { - const budgetCap = uint64(128 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]string, colexec.DefaultBatchSize/2) - for i := range values { - values[i] = strings.Repeat("x", 1024) - } - input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - input.SetRowCount(len(values)) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - hb.setBudget(generation) - defer hb.FreeHashMapAndBatches(proc) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Len(t, hb.Batches.Buf, 1) - require.Equal(t, colexec.DefaultBatchSize, hb.Batches.Buf[0].RowCount()) -} - -func TestCopyBuildBatchProjectedWithoutBudgetTracksPartialTail(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 3, proc.Mp()) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - projection, err := hb.projectedBatchCopy(input) - require.NoError(t, err) - require.Positive(t, projection.nextTailSelected) - require.NoError(t, hb.copyBuildBatchProjected(input, proc, projection)) - defer hb.cleanBatches(proc) - - require.Len(t, hb.Batches.Buf, 1) - require.Equal(t, input.RowCount(), hb.Batches.Buf[0].RowCount()) - require.Equal(t, projection.nextTailSelected, hb.retainedSpillTailSelected) - require.Empty(t, hb.batchReservations) -} - -func TestCopyBuildBatchProjectedWithoutBudgetFailureClearsTailProjection(t *testing.T) { - mp, err := mpool.NewMPool(t.Name(), 1<<20, mpool.NoFixed) - require.NoError(t, err) - proc := testutil.NewProcessWithMPool(t, "", mp) - - first := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 1, proc.Mp()) - second := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, - colexec.DefaultBatchSize-1, proc.Mp()) - var filler []byte - defer func() { - if filler != nil { - mp.Free(filler) - } - first.Clean(proc.Mp()) - second.Clean(proc.Mp()) - proc.Free() - require.Zero(t, mp.CurrNB()) - }() - - var hb HashmapBuilder - projection, err := hb.projectedBatchCopy(first) - require.NoError(t, err) - require.NoError(t, hb.copyBuildBatchProjected(first, proc, projection)) - require.Positive(t, hb.retainedSpillTailSelected) - - projection, err = hb.projectedBatchCopy(second) - require.NoError(t, err) - filler, err = mp.Alloc(int(mp.Cap()-mp.CurrNB()), true) - require.NoError(t, err) - require.Equal(t, mp.Cap(), mp.CurrNB()) - - err = hb.copyBuildBatchProjected(second, proc, projection) - require.Error(t, err) - require.Empty(t, hb.Batches.Buf) - require.Zero(t, hb.retainedSpillTailSelected) - require.Empty(t, hb.batchReservations) -} - -func TestProjectedPartialTailReplacementMatchesUnionBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - tail := batch.NewWithSize(2) - tail.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - tail.Vecs[1] = testutil.MakeVarcharVector([]string{strings.Repeat("a", 1024)}, nil, proc.Mp()) - tail.SetRowCount(1) - defer tail.Clean(proc.Mp()) - - src := batch.NewWithSize(2) - src.Vecs[0] = testutil.MakeInt32Vector([]int32{2, 3}, nil, proc.Mp()) - constVec, err := vector.NewConstBytes( - types.T_varchar.ToType(), - []byte(strings.Repeat("b", 1024)), - 2, - proc.Mp(), - ) - require.NoError(t, err) - src.Vecs[1] = constVec - src.SetRowCount(2) - defer src.Clean(proc.Mp()) - - peak, retained, err := projectedPartialTailReplacementBytes(tail, src, src.RowCount()) - require.NoError(t, err) - require.GreaterOrEqual(t, peak, retained) - before := tail.Allocated() - for i := range tail.Vecs { - require.NoError(t, tail.Vecs[i].UnionBatch(src.Vecs[i], 0, src.RowCount(), nil, proc.Mp())) - } - tail.AddRowCount(src.RowCount()) - require.Equal(t, uint64(tail.Allocated()-before), retained) - - inline := batch.NewWithSize(1) - inline.Vecs[0] = testutil.MakeVarcharVector([]string{"small"}, nil, proc.Mp()) - inline.SetRowCount(1) - defer inline.Clean(proc.Mp()) - preallocated, err := proc.NewBatchFromSrc(inline, colexec.DefaultBatchSize) - require.NoError(t, err) - defer preallocated.Clean(proc.Mp()) - require.NoError(t, preallocated.Vecs[0].UnionBatch(inline.Vecs[0], 0, 1, nil, proc.Mp())) - preallocated.AddRowCount(1) - peak, retained, err = projectedPartialTailReplacementBytes(preallocated, inline, 1) - require.NoError(t, err) - require.Zero(t, peak) - require.Zero(t, retained) -} - -func TestCopyBuildBatchBudgetsPartialTailWithRemainder(t *testing.T) { - const budgetCap = uint64(16 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - var hb HashmapBuilder - hb.setBudget(generation) - defer hb.FreeHashMapAndBatches(proc) - for _, rows := range []int{ - colexec.DefaultBatchSize, - colexec.DefaultBatchSize, - colexec.DefaultBatchSize - 1, - 2, + {name: "grouping then ordinary", groupFirst: true, columns: 1}, + {name: "ordinary then grouping", groupFirst: false, columns: 1}, + {name: "multi-column grouping pattern", groupFirst: true, columns: 2}, } { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 4) - require.Equal(t, 1, hb.Batches.Buf[3].RowCount()) -} - -func TestBatchCopyAllocatedDeltaTracksOnlyChangedSuffix(t *testing.T) { - tests := []struct { - name string - retainedRows int - ingressRows int - }{ - {name: "empty/multiple-destinations", ingressRows: 2*colexec.DefaultBatchSize + 7}, - {name: "full-tail/append", retainedRows: colexec.DefaultBatchSize, ingressRows: colexec.DefaultBatchSize}, - {name: "partial-tail/grow-and-append", retainedRows: 7000, ingressRows: 2000}, - {name: "partial-tail/full-ingress-swap", retainedRows: 7, ingressRows: colexec.DefaultBatchSize}, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { + t.Run(test.name, func(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() + builder := newTestHashmapBuilder(t) + defer builder.Free(proc) - var batches colexec.Batches - defer batches.Clean(proc.Mp()) - if tc.retainedRows > 0 { - retained := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, tc.retainedRows, proc.Mp()) - defer retained.Clean(proc.Mp()) - require.NoError(t, batches.CopyIntoBatches(retained, proc)) + exprs := make([]*plan.Expr, test.columns) + for column := range exprs { + exprs[column] = newExpr(int32(column), types.T_int32.ToType()) + } + require.NoError(t, builder.Prepare(exprs, -1, -1, nil, proc)) + + appendInput := func(grouping bool, groupingColumn int) { + input := batch.NewWithSize(test.columns) + for column := 0; column < test.columns; column++ { + if grouping && column == groupingColumn { + input.Vecs[column] = vector.NewRollupConst( + types.T_int32.ToType(), 1, proc.Mp(), + ) + } else { + input.Vecs[column] = vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixed( + input.Vecs[column], int32(column), false, proc.Mp(), + )) + } + } + input.SetRowCount(1) + require.NoError(t, builder.CopyBuildBatch(input, proc)) + builder.InputBatchRowCount++ + input.Clean(proc.Mp()) + } + + appendInput(test.groupFirst, 0) + if test.columns == 2 { + // A different grouping column in the later row proves that the + // per-column bit pattern survives tail coalescing. + appendInput(true, 1) + } else { + appendInput(!test.groupFirst, 0) } - before := batchesAllocated(batches.Buf) - snapshot, err := snapshotBatchCopyAllocation(batches.Buf) - require.NoError(t, err) - ingress := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, tc.ingressRows, proc.Mp()) - defer ingress.Clean(proc.Mp()) - require.NoError(t, batches.CopyIntoBatches(ingress, proc)) + require.Len(t, builder.Batches.Buf, 1) + require.Equal(t, 2, builder.Batches.Buf[0].RowCount()) + require.False(t, builder.Batches.Buf[0].Vecs[0].IsGrouping()) + require.True(t, builder.Batches.Buf[0].Vecs[0].HasGrouping()) - after := batchesAllocated(batches.Buf) - require.GreaterOrEqual(t, after, before) - delta, err := batchCopyAllocatedDelta(batches.Buf, snapshot) - require.NoError(t, err) - require.Equal(t, after-before, delta) + require.NoError(t, builder.BuildHashmap(false, false, false, proc)) + require.Nil(t, builder.IntHashMap) + require.NotNil(t, builder.StrHashMap) + require.Equal(t, uint64(2), builder.StrHashMap.GroupCount()) }) } } -func TestBatchCopyAllocatedDeltaRejectsLostTail(t *testing.T) { +func TestBuildHashmapDetectsGroupingForOriginalBuildRelation(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - retained := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 7, proc.Mp()) - defer retained.Clean(proc.Mp()) - snapshot, err := snapshotBatchCopyAllocation([]*batch.Batch{retained}) - require.NoError(t, err) - - replacement := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, true, 7, proc.Mp()) - defer replacement.Clean(proc.Mp()) - _, err = batchCopyAllocatedDelta([]*batch.Batch{replacement}, snapshot) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestProjectedPartialTailReplacementRejectsInvalidInputs(t *testing.T) { - _, _, err := projectedPartialTailReplacementBytes(nil, nil, -1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = (&HashmapBuilder{}).projectedBatchCopyBytes(nil) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) + builder := newTestHashmapBuilder(t) + defer builder.Free(proc) - tail := batch.NewOffHeapWithSize(1) - src := batch.NewOffHeapWithSize(1) - _, _, err = projectedPartialTailReplacementBytes(tail, src, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) + expr := newExpr(0, types.T_int32.ToType()) + expr.GetCol().RelPos = 1 + require.NoError(t, builder.Prepare([]*plan.Expr{expr}, -1, -1, nil, proc)) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - src.Vecs[0] = testutil.MakeVarcharVector([]string{strings.Repeat("x", 32)}, nil, proc.Mp()) - src.SetRowCount(1) - defer src.Clean(proc.Mp()) - tail.Vecs[0] = vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - defer tail.Clean(proc.Mp()) - _, _, err = projectedPartialTailReplacementBytes(tail, src, 2) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - invalidTail := batch.NewOffHeapWithSize(1) - invalidTail.SetRowCount(-1) - hb := HashmapBuilder{} - hb.Batches.Buf = []*batch.Batch{invalidTail} - _, err = hb.projectedBatchCopyBytes(src) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - mismatchedTail := batch.NewOffHeapWithSize(0) - hb.Batches.Buf = []*batch.Batch{mismatchedTail} - _, err = hb.projectedBatchCopyBytes(src) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestCopyBuildBatchUsesProjectedDestinationCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - values := make([]string, 4096) - for i := range values { - values[i] = strings.Repeat("x", 1024) - } input := batch.NewWithSize(1) - input.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - input.SetRowCount(len(values)) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - projected, err := hb.projectedBatchCopyBytes(input) - require.NoError(t, err) - destination, err := projectedNewDestinationBytes(input, 0, input.RowCount()) - require.NoError(t, err) - metadata, ok := retainedMetadataAllowance(input) - require.True(t, ok) - const wantSlack = uint64(64 << 10) - require.Equal(t, destination+metadata+wantSlack, projected) + input.Vecs[0] = vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList( + input.Vecs[0], []int32{0, 0}, nil, proc.Mp(), + )) + input.Vecs[0].GetGrouping().Add(1) + input.SetRowCount(2) + require.NoError(t, builder.CopyBuildBatch(input, proc)) + builder.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) - budget := process.MustNewHashBuildBudget(projected, projected) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Equal(t, projected, generation.Peak()) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) + require.NoError(t, builder.BuildHashmap(false, false, false, proc)) + require.Nil(t, builder.IntHashMap) + require.NotNil(t, builder.StrHashMap) + require.Equal(t, uint64(2), builder.StrHashMap.GroupCount()) } -func TestCopyBuildBatchSplitsLargeIngressWithinProjection(t *testing.T) { +func TestBuildHashMap(t *testing.T) { + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 50_000 - input := testutil.NewBatch([]types.Type{types.T_uuid.ToType()}, true, rows, proc.Mp()) - part, err := vector.NewConstFixed[int32](types.T_int32.ToType(), 1, rows, proc.Mp()) - require.NoError(t, err) - input.Vecs = append(input.Vecs, part) - defer input.Clean(proc.Mp()) - - budget := process.MustNewHashBuildBudget(1<<30, 1<<30) - generation, err := budget.OpenGeneration(1) + err := hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc) require.NoError(t, err) - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Equal(t, 2*rows, hb.Batches.RowCount()) - require.Len(t, hb.Batches.Buf, 13) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchSplitsLargeConstVarcharIngressWithinProjection(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 50_000 - input := batch.NewWithSize(1) - value := make([]byte, 1<<20) - vec, err := vector.NewConstBytes(types.T_varchar.ToType(), value, rows, proc.Mp()) + inputBatch := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, int(100000), proc.Mp()) + err = hb.CopyBuildBatch(inputBatch, proc) + hb.InputBatchRowCount = inputBatch.RowCount() + inputBatch.Clean(proc.Mp()) require.NoError(t, err) - input.Vecs[0] = vec - input.SetRowCount(rows) - defer input.Clean(proc.Mp()) - budget := process.MustNewHashBuildBudget(1<<30, 1<<30) - generation, err := budget.OpenGeneration(1) + err = hb.BuildHashmap(false, true, true, proc) require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.Equal(t, 2*rows, hb.Batches.RowCount()) - require.Len(t, hb.Batches.Buf, 13) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) + require.Less(t, int64(0), hb.GetSize()) + require.Less(t, uint64(0), hb.GetGroupCount()) + hb.Reset(proc, true) + hb.Free(proc) + require.Equal(t, int64(0), proc.Mp().CurrNB()) } -func TestCopyBuildBatchManyExactSegmentsAvoidsFalseAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 64 * colexec.DefaultBatchSize - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - defer input.Clean(proc.Mp()) - - destination, err := projectedNewDestinationBytes(input, 0, input.RowCount()) +func TestHashmapBuilderPhysicalAllocationsChargeOnce(t *testing.T) { + const budgetCap = uint64(16 << 20) + budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) require.NoError(t, err) - metadata, ok := retainedMetadataAllowance(input) - require.True(t, ok) - budgetCap := 2*(destination+metadata) + uint64(512<<10) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.NoError(t, hb.copyBuildBatch(input, proc)) - require.LessOrEqual(t, generation.Used(), budgetCap) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchSharedVarlenaRejectsBeforeAllocation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const rows = 50_000 - value := make([]byte, 1<<10) - constVec, err := vector.NewConstBytes(types.T_varchar.ToType(), value, rows, proc.Mp()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) require.NoError(t, err) - defer constVec.Free(proc.Mp()) - - flat := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - require.NoError(t, flat.UnionBatch(constVec, 0, rows, nil, proc.Mp())) - require.False(t, flat.IsConst()) - require.Equal(t, len(value), len(flat.GetArea())) - input := batch.NewWithSize(1) - input.Vecs[0] = flat - input.SetRowCount(rows) - defer input.Clean(proc.Mp()) - - const budgetCap = uint64(10 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) + account, err := registry.OpenWithController(budgetCap, generation) require.NoError(t, err) - var hb HashmapBuilder + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder hb.setBudget(generation) - projected, err := hb.projectedBatchCopyBytes(input) - require.NoError(t, err) - require.GreaterOrEqual(t, projected, uint64(rows*len(value))) - err = hb.copyBuildBatch(input, proc) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Empty(t, hb.Batches.Buf) - require.Zero(t, generation.Used()) -} - -func TestCopyBuildBatchWholeSharedVarlenaAvoidsFalseAdmission(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - const rows = colexec.DefaultBatchSize - value := make([]byte, 1<<10) - constVec, err := vector.NewConstBytes(types.T_varchar.ToType(), value, rows, proc.Mp()) - require.NoError(t, err) - defer constVec.Free(proc.Mp()) - - flat := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - require.NoError(t, flat.UnionBatch(constVec, 0, rows, nil, proc.Mp())) - require.False(t, flat.IsConst()) - input := batch.NewWithSize(1) - input.Vecs[0] = flat - input.SetRowCount(rows) - defer input.Clean(proc.Mp()) - - const budgetCap = uint64(2 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) + require.NoError(t, hb.Prepare( + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + -1, + -1, + nil, + proc, + )) + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType()}, + true, + 10_000, + proc.Mp(), + ) require.NoError(t, hb.copyBuildBatch(input, proc)) - require.LessOrEqual(t, generation.Used(), budgetCap) - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) -} - -func TestReserveBuildAuxChargesOneRetainedCopy(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, colexec.DefaultBatchSize, proc.Mp()) - defer input.Clean(proc.Mp()) - - var hb HashmapBuilder - hb.Batches.Buf = []*batch.Batch{input} - hb.InputBatchRowCount = input.RowCount() - retained := batchesAllocated(hb.Batches.Buf) - const iteratorScratch = uint64(640 << 10) - want := retained + (retained+3)/4 + uint64(input.RowCount())*64 + iteratorScratch - - budget := process.MustNewHashBuildBudget(want, want) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb.setBudget(generation) - require.NoError(t, hb.reserveBuildAux(true)) - require.Equal(t, want, generation.Used()) - hb.releaseReservations() - require.Zero(t, generation.Used()) - // The batch belongs to the test rather than batchReservations. - hb.Batches.Buf = nil -} - -func TestReserveUniqueAppendOverlapChargesReplacedCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]string, 4_096) - for i := range values { - values[i] = strings.Repeat("x", 1_024) + require.NotEmpty(t, hb.Batches.Buf) + for _, copied := range hb.Batches.Buf { + require.Same(t, hb.batchAllocation, copied.AllocationAccountSelection()) } - dst := testutil.MakeVarcharVector(values, nil, proc.Mp()) - defer dst.Free(proc.Mp()) - extraArea := cap(dst.GetArea()) - len(dst.GetArea()) + 1 - src := testutil.MakeVarcharVector([]string{strings.Repeat("y", extraArea)}, nil, proc.Mp()) - defer src.Free(proc.Mp()) - - want := uint64(cap(dst.GetData()) + cap(dst.GetArea())) - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb := HashmapBuilder{ - budget: generation, - UniqueJoinKeys: []*vector.Vector{dst}, - } - require.NoError(t, hb.reserveBuildAux(true)) - areaBytes, err := uniqueAppendAreaBytes(src, 0, 1, nil) - require.NoError(t, err) - token, err := hb.reserveUniqueAppendOverlap(dst, 1, areaBytes) - require.NoError(t, err) - require.NotNil(t, token) - require.Equal(t, want, token.Size()) - token.Release() - require.Equal(t, hb.auxReservation.Size(), generation.Used()) - hb.releaseReservations() - require.Zero(t, generation.Used()) - - largeValue := strings.Repeat("z", 100) - selected := testutil.MakeVarcharVector([]string{"a", largeValue}, nil, proc.Mp()) - defer selected.Free(proc.Mp()) - selectedArea, err := uniqueAppendAreaBytes(selected, 0, 1, []int64{0}) - require.NoError(t, err) - require.Zero(t, selectedArea) - selectedArea, err = uniqueAppendAreaBytes(selected, 0, 1, []int64{1}) - require.NoError(t, err) - require.Equal(t, len(largeValue), selectedArea) - _, err = uniqueAppendAreaBytes(selected, -1, 1, nil) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = uniqueAppendAreaBytes(selected, 0, 1, []int64{int64(selected.Length())}) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = uniqueAppendAreaBytes(selected, 0, 2, []int64{0}) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - fixed := testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - defer fixed.Free(proc.Mp()) - selectedArea, err = uniqueAppendAreaBytes(fixed, math.MaxInt, math.MaxInt, nil) - require.NoError(t, err) - require.Zero(t, selectedArea) - - constValue, err := vector.NewConstBytes(types.T_varchar.ToType(), []byte(largeValue), 2, proc.Mp()) - require.NoError(t, err) - defer constValue.Free(proc.Mp()) - selectedArea, err = uniqueAppendAreaBytes(constValue, 0, 2, nil) - require.NoError(t, err) - require.Equal(t, 2*len(largeValue), selectedArea) - - noBudget := HashmapBuilder{} - token, err = noBudget.reserveUniqueAppendOverlap(dst, 1, 1) - require.NoError(t, err) - require.Nil(t, token) - token, err = hb.reserveUniqueAppendOverlap(nil, 1, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - token, err = hb.reserveUniqueAppendOverlap(dst, -1, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - token, err = hb.reserveUniqueAppendOverlap(dst, 1, -1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestUniqueAppendBudgetIncludesDeadAreaCopiedByUnionBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - src := testutil.MakeVarcharVector( - []string{"inline", strings.Repeat("d", 128<<10)}, nil, proc.Mp()) - defer src.Free(proc.Mp()) - // SetLength leaves the second value's area allocation behind. The sole live - // row is inline, but UnionBatch's whole-vector fast path copies all of area. - src.SetLength(1) - liveArea, err := uniqueAppendAreaBytes(src, 0, 1, nil) - require.NoError(t, err) - require.Zero(t, liveArea) - unionArea, err := unionBatchAreaBytes(src, 0, 1) - require.NoError(t, err) - require.Equal(t, len(src.GetArea()), unionArea) - require.Greater(t, unionArea, 0) - - dst := vector.NewOffHeapVecWithType(types.T_varchar.ToType()) - defer dst.Free(proc.Mp()) - const mandatoryAux = uint64(640 << 10) - budget := process.MustNewHashBuildBudget(mandatoryAux, mandatoryAux) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - hb := HashmapBuilder{ - budget: generation, - UniqueJoinKeys: []*vector.Vector{dst}, - } - require.NoError(t, hb.reserveBuildAux(true)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) - _, err = hb.reserveUniqueAppendOverlap(dst, 1, unionArea) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Zero(t, dst.Length()) - require.Zero(t, dst.Allocated(), - "admission must fail before UnionBatch allocates copied dead area") + require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + require.Positive(t, account.Snapshot().Used) + require.Equal(t, account.Snapshot().Used, generation.Used()) - hb.releaseReservations() + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(2) + hb.Reset(proc, false) + beforeResize := account.Snapshot().Used + require.NoError(t, jm.PreAlloc(100_000)) + require.Greater(t, account.Snapshot().Used, beforeResize) + beforeFirstConsumer := account.Snapshot().Used + jm.Free() + require.Equal(t, beforeFirstConsumer, account.Snapshot().Used) + jm.Free() + require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) - generation.Close() + + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) } -func TestUnionBatchAreaProjectionBoundaries(t *testing.T) { +func TestHashmapBuilderAccountedBatchCopyOneByteShortRollsBack(t *testing.T) { + const budgetCap = uint64(64 << 20) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_varchar.ToType()}, + true, + 10_000, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) - inline := testutil.MakeVarcharVector([]string{"inline"}, nil, proc.Mp()) - defer inline.Free(proc.Mp()) - require.Empty(t, inline.GetArea()) + measure := func(limit uint64, metadataSlots uint64) ( + mpool.AllocationAccountSnapshot, + uint64, + error, + ) { + budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + + copyErr := hb.copyBuildBatch(input, proc) + snapshot := account.Snapshot() + metadataPeak := registry.PeakAllocationMetadata() + if copyErr == nil { + hb.cleanBatches(proc) + } + require.Empty(t, hb.Batches.Buf) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + terminal, first, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + return snapshot, metadataPeak, copyErr + } - physical, selected, err := unionBatchAreaProjection(inline, 0, 0) + probe, metadataPeak, err := measure(budgetCap, 128) require.NoError(t, err) - require.Zero(t, physical) - require.Zero(t, selected) + require.Positive(t, probe.Peak) + require.Positive(t, metadataPeak) + rejected, _, err := measure(probe.Peak-1, 128) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, rejected.Used) + rejected, _, err = measure(budgetCap, metadataPeak-1) + require.ErrorIs(t, err, mpool.ErrAllocationMetadataSlots) + require.Zero(t, rejected.Used) +} - physical, selected, err = unionBatchAreaProjection(inline, 0, 1) +func TestAccountedJoinMapTransfersBatchesAndGroupSelsToLastConsumer(t *testing.T) { + const budgetCap = uint64(16 << 20) + budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) require.NoError(t, err) - require.Zero(t, physical) - require.Zero(t, selected) + account, err := registry.OpenWithController(budgetCap, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + require.NoError(t, hb.Prepare( + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + -1, + -1, + nil, + proc, + )) + input := makeIntKeyValueBatch( + proc, + []int32{1, 1, 2, 2}, + []int32{10, 20, 30, 40}, + ) + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + require.NoError(t, hb.BuildHashmap(false, true, false, proc)) + require.Positive(t, hb.Sels.Size()) - _, _, err = unionBatchAreaProjection(inline, -1, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(2) + hb.Reset(proc, false) + require.Equal(t, []int32{0, 1}, jm.GetSels(0)) + require.Equal(t, []int32{2, 3}, jm.GetSels(1)) + live := account.Snapshot().Used + require.Positive(t, live) + jm.Free() + require.Equal(t, live, account.Snapshot().Used) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) - selected, err = logicalAppendAreaBytes(inline, 0, 0) + terminal, first, err := registry.CompleteTerminal(account) require.NoError(t, err) - require.Zero(t, selected) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) } -func TestUnionBatchAreaProjectionWideFlatAndSharedRepresentations(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - typ := types.New(types.T_varchar, 128, 0) - const rows = 128 - payload := make([]byte, 128) +func TestAccountedEmptyJoinMapUsesPhysicalAllocationAsSoleCharge(t *testing.T) { + for _, tc := range []struct { + name string + keyWidth int + initialBytes uint64 + }{ + {name: "int", keyWidth: 4, initialBytes: hashtable.Int64HashMapInitialAllocationBytes()}, + {name: "string", keyWidth: 128, initialBytes: hashtable.StringHashMapInitialAllocationBytes()}, + } { + t.Run(tc.name, func(t *testing.T) { + const capBytes = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + mp := mpool.MustNewZero() - flat := vector.NewVec(typ) - for range rows { - require.NoError(t, vector.AppendBytes(flat, payload, false, proc.Mp())) + jm, err := NewAccountedEmptyJoinMap(tc.keyWidth, account, mp) + require.NoError(t, err) + descriptorBytes := hashtable.HashMapBlockDescriptorBytes() + expectedInitial := tc.initialBytes + descriptorBytes + require.Equal(t, expectedInitial, account.Snapshot().Used) + require.Equal(t, expectedInitial, generation.Used()) + + require.NoError(t, jm.PreAlloc(10_000)) + require.Equal(t, account.Snapshot().Used, generation.Used()) + require.Equal(t, account.Snapshot().Used, generation.Used()) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) } - defer flat.Free(proc.Mp()) - require.True(t, flat.VarlenaAreaIsDisjoint()) - physical, selected, err := unionBatchAreaProjection(flat, 0, rows) - require.NoError(t, err) - require.Equal(t, rows*len(payload), physical) - require.Equal(t, uint64(physical), selected) +} - constant, err := vector.NewConstBytes(typ, payload, rows, proc.Mp()) +func TestAccountedEmptyJoinMapInitialFailureRollsBackController(t *testing.T) { + initial := hashtable.Int64HashMapInitialAllocationBytes() + + hashtable.HashMapBlockDescriptorBytes() + budget := process.MustNewHashBuildBudget(initial, initial) + generation, err := budget.OpenGeneration(1) require.NoError(t, err) - defer constant.Free(proc.Mp()) - shared := vector.NewVec(typ) - require.NoError(t, shared.UnionBatch(constant, 0, rows, nil, proc.Mp())) - defer shared.Free(proc.Mp()) - require.False(t, shared.IsConst()) - require.False(t, shared.VarlenaAreaIsDisjoint()) - require.Equal(t, len(payload), len(shared.GetArea())) - - physical, selected, err = unionBatchAreaProjection(shared, 0, rows) + registry, err := mpool.NewAllocationAccountRegistry(1, 2) require.NoError(t, err) - require.Equal(t, len(payload), physical) - require.Equal(t, uint64(rows*len(payload)), selected, - "a shared flat representation must retain the descriptor-scan fallback") + account, err := registry.OpenWithController(initial-1, generation) + require.NoError(t, err) + mp := mpool.MustNewZero() - physical, selected, err = unionBatchAreaProjection(flat, 1, rows-1) + jm, err := NewAccountedEmptyJoinMap(4, account, mp) + require.Nil(t, jm) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Zero(t, registry.LiveAllocationMetadata()) + require.Zero(t, mp.CurrNB()) + _, _, err = registry.CompleteTerminal(account) require.NoError(t, err) - require.Equal(t, (rows-1)*len(payload), physical) - require.Equal(t, uint64(physical), selected, - "a partial range must retain its exact-copy path") } -func TestCleanCopiedBatchReleasesCoalescedIngressReservations(t *testing.T) { - const budgetCap = uint64(4 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) +func TestAccountedJoinMapLateFreeKeepsOriginalGeneration(t *testing.T) { + const capBytes = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + firstGeneration, err := budget.OpenGeneration(1) require.NoError(t, err) - generation, err := budget.OpenGeneration(1) + secondGeneration, err := budget.OpenGeneration(2) require.NoError(t, err) - - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for range 2 { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, colexec.DefaultBatchSize/2, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 1, "small ingress batches should coalesce") - require.Len(t, hb.batchReservations, 2, "reservations follow ingress, not physical batches") - require.Greater(t, generation.Used(), uint64(0)) - require.NoError(t, hb.CleanCopiedBatchAt(0, proc)) - require.Empty(t, hb.Batches.Buf) - require.Empty(t, hb.batchReservations) - require.Zero(t, generation.Used()) -} - -func TestDrainCopiedBatchesReleasesBeforeSubsequentAdmission(t *testing.T) { - const budgetCap = uint64(4 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) + registry, err := mpool.NewAllocationAccountRegistry(2, 16) require.NoError(t, err) - generation, err := budget.OpenGeneration(1) + firstAccount, err := registry.OpenWithController(capBytes, firstGeneration) require.NoError(t, err) + mp := mpool.MustNewZero() + jm, err := NewAccountedEmptyJoinMap(4, firstAccount, mp) + require.NoError(t, err) + firstUsed := firstGeneration.Used() + require.Positive(t, firstUsed) - var hb HashmapBuilder - hb.setBudget(generation) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - for range 2 { - input := testutil.NewBatch( - []types.Type{types.T_int32.ToType()}, - true, - colexec.DefaultBatchSize/2, - proc.Mp(), - ) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) - } - require.Len(t, hb.Batches.Buf, 1, "small ingress batches should coalesce") - require.Len(t, hb.batchReservations, 2, "reservations follow ingress, not physical batches") - - visits := 0 - require.NoError(t, hb.DrainCopiedBatches(proc, func(bat *batch.Batch) error { - visits++ - require.NotNil(t, bat) - require.Positive(t, generation.Used(), "physical batch must remain charged while visited") - return nil - })) - require.Equal(t, 1, visits) - require.Empty(t, hb.Batches.Buf) - require.Empty(t, hb.batchReservations) - require.Zero(t, generation.Used(), "the final physical batch must release every coalesced ingress charge") - - // Model the expression/scatter/read reservation that follows a re-spill - // drain. It can consume the complete cap only after stale batch ownership - // has been removed from the ledger. - next, err := generation.Reserve(budgetCap) + secondAccount, err := registry.OpenWithController(capBytes, secondGeneration) require.NoError(t, err) - require.True(t, next.Release()) - require.Zero(t, generation.Used()) -} + require.Zero(t, secondGeneration.Used()) + jm.Free() + require.Zero(t, firstGeneration.Used()) + require.Zero(t, firstAccount.Snapshot().Used) + require.Zero(t, secondGeneration.Used()) -func TestDrainCopiedBatchesVisitFailureRetainsOwnership(t *testing.T) { - const budgetCap = uint64(4 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) + _, _, err = registry.CompleteTerminal(firstAccount) require.NoError(t, err) - generation, err := budget.OpenGeneration(1) + _, _, err = registry.CompleteTerminal(secondAccount) require.NoError(t, err) +} - var hb HashmapBuilder - hb.setBudget(generation) +func TestAccountedRuntimeFilterUniqueKeysDegradeWithoutFailingHashBuild(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - for _, rows := range []int{ - colexec.DefaultBatchSize / 2, - colexec.DefaultBatchSize / 2, - 1024, - } { - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, rows, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - input.Clean(proc.Mp()) + values := make([]string, 1_024) + for i := range values { + values[i] = strconv.Itoa(i) + strings.Repeat(string(rune('a'+i%26)), 4<<10) } - require.Len(t, hb.Batches.Buf, 2) - require.Len(t, hb.batchReservations, 3) - - wantErr := errors.New("visit failed") - visits := 0 - require.ErrorIs(t, hb.DrainCopiedBatches(proc, func(*batch.Batch) error { - visits++ - if visits == 1 { - return nil + input := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeVarcharVector(values, nil, proc.Mp()), + }, nil) + defer input.Clean(proc.Mp()) + exprs := []*plan.Expr{newExpr(0, types.T_varchar.ToType())} + + run := func(limit uint64, needUnique bool) (mpool.AllocationAccountSnapshot, bool) { + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + builder := &HashmapBuilder{} + builder.SetBudget(generation) + require.NoError(t, builder.SetAllocationAccount(account)) + require.NoError(t, builder.Prepare(exprs, -1, -1, nil, proc)) + require.NoError(t, builder.CopyBuildBatch(input, proc)) + builder.InputBatchRowCount = input.RowCount() + require.NoError(t, builder.BuildHashmap(false, false, needUnique, proc)) + snapshot := account.Snapshot() + fallback, _ := builder.runtimeFilterFallbackState() + if needUnique { + require.True(t, fallback) + require.Empty(t, builder.UniqueJoinKeys) + require.NotNil(t, builder.StrHashMap) + require.Equal(t, uint64(input.RowCount()), builder.StrHashMap.GroupCount()) } - return wantErr - }), wantErr) - require.Equal(t, 2, visits) - require.Len(t, hb.Batches.Buf, 1, "the failed current batch remains owned after prior batches drain") - require.Len(t, hb.batchReservations, 3, - "coalesced ingress reservations stay conservative until terminal cleanup") - require.Positive(t, generation.Used()) - - hb.FreeHashMapAndBatches(proc) - require.Zero(t, generation.Used()) + builder.Free(proc) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + return snapshot, fallback + } + + baseline, fallback := run(64<<20, false) + require.False(t, fallback) + require.Positive(t, baseline.Peak) + // The exact baseline peak is sufficient for the required hash build, but + // not for a second, optional copy of the 4 MiB runtime-filter key payload. + constrained, fallback := run(baseline.Peak, true) + require.True(t, fallback) + require.LessOrEqual(t, constrained.Peak, baseline.Peak) } -func TestSpillExpressionHashKeyUsesBoundedAdmission(t *testing.T) { - var ctr container +func TestSpillExpressionStorageUsesRetainedAccount(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) + budget := process.MustNewHashBuildBudget(16<<20, 16<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - ctr.hashmapBuilder.setBudget(generation) - expr := makeExpressionLeaseTestExpr(t, proc) - err = initSpillExprExecsForTest(&ctr, proc, []*plan.Expr{expr}) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) require.NoError(t, err) - require.NoError(t, ctr.hashmapBuilder.expressionLease.Run(proc, 8192, func(_ int) error { return nil })) - require.Positive(t, ctr.hashmapBuilder.expressionLease.Reserved()) - require.Equal(t, ctr.hashmapBuilder.expressionLease.Reserved(), generation.Used()) - ctr.hashmapBuilder.FreeExecutors() - require.Zero(t, generation.Used()) -} - -func TestExpressionHashKeyReservesDeclaredPeakBeforeEval(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(96<<10, 96<<10) + account, err := registry.OpenWithController(16<<20, generation) require.NoError(t, err) - generation, err := budget.OpenGeneration(1) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + ctr := &op.ctr + ctr.hashmapBuilder.setBudget(generation) + expr := makeIssue26454ConcatKey(t, proc) + executors, err := ctr.initSpillExprExecs(proc, []*plan.Expr{expr}) + require.NoError(t, err) + constructorUsed := account.Snapshot().Used + input := batch.NewWithSize(2) + input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) + input.Vecs[1] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) + input.SetRowCount(2) + defer input.Clean(proc.Mp()) + result, err := executors[0].Eval(proc, []*batch.Batch{input}, nil) require.NoError(t, err) + require.Equal(t, []string{"1-3", "2-4"}, vector.InefficientMustStrCol(result)) + require.Greater(t, account.Snapshot().Used, constructorUsed) + require.Equal(t, account.Snapshot().Used, generation.Used()) - var hb HashmapBuilder - hb.setBudget(generation) - require.NoError(t, hb.Prepare([]*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_F{F: &plan.Function{}}, - }}, -1, -1, nil, proc)) - input := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 1, proc.Mp()) - require.NoError(t, hb.copyBuildBatch(input, proc)) - hb.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - err = hb.BuildHashmap(false, true, false, proc) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - hb.Free(proc) + ctr.freeSpillExprExecs() + require.Zero(t, account.Snapshot().Used) require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) } -func TestExpressionHashKeyAcceptsCastTargetType(t *testing.T) { - proc := testutil.NewProcess(t) +func TestSpillExpressionStorageHonorsAccountCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_F{F: &plan.Function{Args: []*plan.Expr{ - newExpr(0, types.T_int64.ToType()), - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_T{T: &plan.TargetType{}}, - }, - }}}, + run := func(limit uint64) (uint64, error) { + budget := process.MustNewHashBuildBudget(16<<20, 16<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + op.ctr.hashmapBuilder.setBudget(generation) + executors, evalErr := op.ctr.initSpillExprExecs( + proc, + []*plan.Expr{makeIssue26454ConcatKey(t, proc)}, + ) + if evalErr == nil { + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + true, + 10_000, + proc.Mp(), + ) + _, evalErr = executors[0].Eval(proc, []*batch.Batch{input}, nil) + input.Clean(proc.Mp()) + } + peak := account.Snapshot().Peak + op.ctr.freeSpillExprExecs() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + _, _, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + return peak, evalErr } - peak, err := expressionVectorPeak(proc, expr, 1024, false) + peak, err := run(16 << 20) require.NoError(t, err) - outputPeak, err := expressionFixedWidthPeak(uint64(types.T_int32.TypeLen()), 1024) - require.NoError(t, err) - require.Equal(t, 2*outputPeak, peak, - "charge the target-type and cast result vectors") + require.Greater(t, peak, uint64(1)) + _, err = run(peak - 1) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) } -func TestPreparedParamExpressionPeakUsesConstCardinality(t *testing.T) { - proc := testutil.NewProcess(t) +func TestIssue26454ExpressionKeyBuildUsesActualCapacity(t *testing.T) { + const capBytes = uint64(16 << 20) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - proc.SetPrepareParams(params) - require.NoError(t, vector.AppendBytes(params, []byte("prepared"), false, proc.Mp())) - require.NoError(t, vector.AppendBytes(params, nil, true, proc.Mp())) - - paramExpr := func(pos int32) *plan.Expr { - return &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_text), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: pos}}, - } - } + for _, tc := range []struct { + name string + expr *plan.Expr + input func() *batch.Batch + }{ + { + name: "concat cast key", + expr: makeIssue26454ConcatKey(t, proc), + input: func() *batch.Batch { + return testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_int32.ToType()}, + true, + 10_000, + proc.Mp(), + ) + }, + }, + { + name: "case equality key", + expr: makeIssue26454CaseKey(t, proc), + input: func() *batch.Batch { + values := make([]string, 10_000) + for i := range values { + if i%2 == 0 { + values[i] = "ATM_CON" + } else { + values[i] = "OTHER" + } + } + bat := batch.NewWithSize(1) + bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) + bat.SetRowCount(len(values)) + return bat + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + require.NoError(t, hb.Prepare([]*plan.Expr{tc.expr}, -1, -1, nil, proc)) + input := tc.input() + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + require.LessOrEqual(t, generation.Used(), capBytes) - peakOne, err := expressionVectorPeak(proc, paramExpr(0), 1, false) - require.NoError(t, err) - peakBatch, err := expressionVectorPeak(proc, paramExpr(0), colexec.DefaultBatchSize, false) - require.NoError(t, err) - require.Equal(t, peakOne, peakBatch, "const parameter admission must not scale with input rows") - peakNull, err := expressionVectorPeak(proc, paramExpr(1), colexec.DefaultBatchSize, false) - require.NoError(t, err) - require.Equal(t, peakOne, peakNull, "null parameter keeps the declared one-row type bound") + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(1) + hb.Reset(proc, false) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) + } } func TestPreparedParamExpressionExecutorRemainsConst(t *testing.T) { @@ -1346,93 +727,14 @@ func TestPreparedParamExpressionExecutorRemainsConst(t *testing.T) { } } -func TestPreparedParamExpressionPeakNestedFunctionCardinality(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - require.NoError(t, vector.AppendBytes(params, []byte("prepared"), false, proc.Mp())) - proc.SetPrepareParams(params) - - param := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_text), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: 0}}, - } - cast := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int64)}, - Expr: &plan.Expr_F{F: &plan.Function{Args: []*plan.Expr{ - param, - {Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_T{T: &plan.TargetType{}}}, - }}}, - } - modulo := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int64)}, - Expr: &plan.Expr_F{F: &plan.Function{Args: []*plan.Expr{ - {Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{RelPos: 0, ColPos: 0}}}, - cast, - }}}, - } - - paramTotal, paramOutput, err := expressionTreePeak(proc, param, colexec.DefaultBatchSize) - require.NoError(t, err) - paramOne, _, err := expressionTreePeak(proc, param, 1) - require.NoError(t, err) - require.Equal(t, paramOne, paramTotal) - _, rootOutput, err := expressionTreePeak(proc, modulo, colexec.DefaultBatchSize) - require.NoError(t, err) - rootTypePeak, err := expressionTypePeak(modulo.Typ, colexec.DefaultBatchSize) - require.NoError(t, err) - require.Equal(t, rootTypePeak, rootOutput, "function output remains sized for input rows") - require.Greater(t, paramOutput, uint64(0)) -} - -func TestPreparedParamExpressionPeakRejectsInvalidPosition(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - proc.SetPrepareParams(params) - require.NoError(t, vector.AppendBytes(params, []byte("prepared"), false, proc.Mp())) - - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_text)}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: -1}}, - } - _, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.Error(t, err) -} - -func TestPreparedParamExpressionPeakAccountsLargePayload(t *testing.T) { - proc := testutil.NewProcess(t) - defer proc.Free() - params := vector.NewVec(types.T_text.ToType()) - defer params.Free(proc.Mp()) - payload := make([]byte, types.MaxBlobLen+1) - require.NoError(t, vector.AppendBytes(params, payload, false, proc.Mp())) - proc.SetPrepareParams(params) - - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_varchar), Width: types.MaxVarcharLen}, - Expr: &plan.Expr_P{P: &plan.ParamRef{Pos: 0}}, - } - peak, err := expressionVectorPeak(proc, expr, colexec.DefaultBatchSize, false) - require.NoError(t, err) - header, ok := mpool.GrowCapacity(0, int64(types.VarlenaSize)) - require.True(t, ok) - area, ok := mpool.GrowCapacity(0, int64(len(payload))) - require.True(t, ok) - require.GreaterOrEqual(t, peak, uint64(header)+uint64(area)) - require.Greater(t, peak, uint64(types.MaxBlobLen)) -} - func TestGetJoinMapTransfersGroupSels(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) input := makeIntKeyValueBatch(proc, []int32{1, 1}, []int32{10, 20}) - require.NoError(t, hb.Batches.CopyIntoBatches(input, proc)) + require.NoError(t, hb.CopyBuildBatch(input, proc)) hb.InputBatchRowCount = input.RowCount() input.Clean(proc.Mp()) @@ -1456,7 +758,7 @@ func TestGetJoinMapTransfersGroupSels(t *testing.T) { } func TestDedupUpdateBuildGroupsNullKeysSeparately(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_UPDATE @@ -1480,7 +782,7 @@ func TestDedupUpdateBuildGroupsNullKeysSeparately(t *testing.T) { bat := batch.New([]string{"id"}) bat.SetVector(0, keyVec) bat.SetRowCount(rows) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1500,7 +802,7 @@ func TestDedupUpdateBuildGroupsNullKeysSeparately(t *testing.T) { func TestHashMapAllocAndFree(t *testing.T) { mp := mpool.MustNewZero() - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) var err error hb.IntHashMap, err = hashmap.NewIntHashMap(false, mp) require.NoError(t, err) @@ -1525,7 +827,7 @@ func TestHashMapAllocAndFree(t *testing.T) { } func TestIteratorReuseAcrossBuilds(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -1533,7 +835,7 @@ func TestIteratorReuseAcrossBuilds(t *testing.T) { defer b.Clean(proc.Mp()) hb.InputBatchRowCount = b.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(b, proc)) + require.NoError(t, hb.CopyBuildBatch(b, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) itr1 := hb.cachedIntIterator @@ -1544,14 +846,15 @@ func TestIteratorReuseAcrossBuilds(t *testing.T) { require.Same(t, itr1, hb.cachedIntIterator) // Next build should reuse the same iterator instance. + require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = b.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(b, proc)) + require.NoError(t, hb.CopyBuildBatch(b, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.Same(t, itr1, hb.cachedIntIterator) } func TestStrIteratorCapacityPrune(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) @@ -1563,7 +866,7 @@ func TestStrIteratorCapacityPrune(t *testing.T) { bat.SetVector(0, vec) bat.SetRowCount(1) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) require.Greater(t, hashmap.StrIteratorCapacity(hb.cachedStrIterator), hashmap.MaxStrIteratorCapacity) @@ -1573,7 +876,7 @@ func TestStrIteratorCapacityPrune(t *testing.T) { } func TestStrIteratorBelowThresholdIsKept(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) @@ -1586,7 +889,7 @@ func TestStrIteratorBelowThresholdIsKept(t *testing.T) { bat.SetVector(0, vec) bat.SetRowCount(vec.Length()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) require.Less(t, hashmap.StrIteratorCapacity(hb.cachedStrIterator), hashmap.MaxStrIteratorCapacity) @@ -1596,7 +899,7 @@ func TestStrIteratorBelowThresholdIsKept(t *testing.T) { } func TestResetWithHashTableSentKeepsCache(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -1604,7 +907,7 @@ func TestResetWithHashTableSentKeepsCache(t *testing.T) { b := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 8, proc.Mp()) defer b.Clean(proc.Mp()) hb.InputBatchRowCount = b.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(b, proc)) + require.NoError(t, hb.CopyBuildBatch(b, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -1614,7 +917,7 @@ func TestResetWithHashTableSentKeepsCache(t *testing.T) { } func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // First int build @@ -1622,7 +925,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { bInt := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 4, proc.Mp()) defer bInt.Clean(proc.Mp()) hb.InputBatchRowCount = bInt.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bInt, proc)) + require.NoError(t, hb.CopyBuildBatch(bInt, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) require.Nil(t, hb.cachedStrIterator) @@ -1639,7 +942,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { bat.SetVector(0, vec) bat.SetRowCount(1) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) @@ -1653,7 +956,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { bInt2 := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 4, proc.Mp()) defer bInt2.Clean(proc.Mp()) hb.InputBatchRowCount = bInt2.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bInt2, proc)) + require.NoError(t, hb.CopyBuildBatch(bInt2, proc)) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -1663,7 +966,7 @@ func TestAlternateIntStrBuildsReuseIndependently(t *testing.T) { } func TestBuildHashmapWithZeroInputKeepsCachesUntouched(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -1677,7 +980,7 @@ func TestBuildHashmapWithZeroInputKeepsCachesUntouched(t *testing.T) { } func TestDedupBuildDuplicateKeyStillFailsByDefault(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_FAIL @@ -1691,7 +994,7 @@ func TestDedupBuildDuplicateKeyStillFailsByDefault(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{1, 1}, []int32{10, 20}) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1701,7 +1004,7 @@ func TestDedupBuildDuplicateKeyStillFailsByDefault(t *testing.T) { } func TestDedupBuildKeepLastForReplace(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.DedupBuildKeepLast = true @@ -1716,7 +1019,7 @@ func TestDedupBuildKeepLastForReplace(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{1, 1, 2}, []int32{10, 20, 30}) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1734,7 +1037,7 @@ func TestDedupBuildKeepLastForReplace(t *testing.T) { } func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.DedupBuildKeepLast = true @@ -1755,7 +1058,7 @@ func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { []int32{100, 0, 0}, []uint64{1, 2}, ) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1775,6 +1078,113 @@ func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { require.Equal(t, int32(100), markers[2]) } +func TestAccountedDedupScratchAndDeleteBitmapFollowJoinMapLifetime(t *testing.T) { + const capBytes = uint64(64 << 20) + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) + require.NoError(t, err) + account, err := registry.OpenWithController(capBytes, generation) + require.NoError(t, err) + + var op HashBuild + op.NeedHashMap = true + require.NoError(t, op.SetAllocationAccount(account)) + hb := &op.ctr.hashmapBuilder + hb.setBudget(generation) + hb.IsDedup = true + hb.DedupBuildKeepLast = true + hb.OnDuplicateAction = plan.Node_FAIL + hb.DedupColName = "id" + hb.DedupColTypes = []plan.Type{newExpr(0, types.T_int32.ToType()).Typ} + require.NoError(t, hb.Prepare( + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, + -1, + 2, + []int32{2}, + proc, + )) + input := makeIntKeyValueBatchWithMarker( + proc, + []int32{1, 1, 2}, + []int32{10, 20, 30}, + []int32{100, 0, 0}, + []uint64{1, 2}, + ) + require.NoError(t, hb.copyBuildBatch(input, proc)) + hb.InputBatchRowCount = input.RowCount() + input.Clean(proc.Mp()) + + require.NoError(t, hb.BuildHashmap(false, false, false, proc)) + require.NotNil(t, hb.DelRows) + require.True(t, hb.DelRows.HasExternalStorage()) + require.True(t, hb.DelRows.Contains(2)) + require.Equal(t, account.Snapshot().Used, generation.Used()) + require.Positive(t, account.Snapshot().Used) + + jm := hb.GetJoinMap(proc.Mp()) + require.NotNil(t, jm) + jm.IncRef(1) + hb.Reset(proc, false) + // DelRows remains physically owned by the consumer together with the map. + require.Positive(t, account.Snapshot().Used) + require.True(t, jm.IsDeleted(2)) + jm.Free() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.NoError(t, op.ClearAllocationAccount(account)) + terminal, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) +} + +func TestAccountedDedupBitmapExactBoundaryRollsBack(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + for _, tc := range []struct { + name string + cap uint64 + wantErr bool + }{ + {name: "one byte short", cap: 7, wantErr: true}, + {name: "exact", cap: 8}, + } { + t.Run(tc.name, func(t *testing.T) { + budget := process.MustNewHashBuildBudget(tc.cap, tc.cap) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.OpenWithController(tc.cap, generation) + require.NoError(t, err) + hb := &HashmapBuilder{mapAllocationAccount: account} + bm, err := hb.newDedupBitmap( + 64, + proc.Mp(), + HashBuildAllocationSiteDedupIgnoreBitmap, + ) + if tc.wantErr { + require.Error(t, err) + require.True(t, IsRetryableMemoryCapacity(err)) + require.Nil(t, bm) + require.Zero(t, account.Snapshot().Used) + } else { + require.NoError(t, err) + require.Equal(t, uint64(8), account.Snapshot().Used) + releaseDedupBitmap(bm, proc.Mp()) + require.Zero(t, account.Snapshot().Used) + } + terminal, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + }) + } +} + // TestDedupBuildKeepLastMarksConflictBucketForDiscardedFanout reproduces the // REPLACE multi-UK fan-out case (issue #24428) at the hashbuild layer: one new // row (same new PK) fans out to several build rows that carry DIFFERENT old @@ -1783,7 +1193,7 @@ func TestDedupBuildKeepLastPreservesDeleteOnlyRows(t *testing.T) { // row's old PK equals the surviving row's new key, otherwise the dedup-join // probe side raises a false DuplicateEntry for the existing row REPLACE removes. func TestDedupBuildKeepLastMarksConflictBucketForDiscardedFanout(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.DedupBuildKeepLast = true @@ -1808,7 +1218,7 @@ func TestDedupBuildKeepLastMarksConflictBucketForDiscardedFanout(t *testing.T) { []int32{100, 200, 300}, nil, ) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1850,7 +1260,7 @@ func TestDedupBuildIgnoreOnlyMarksCandidateOwnOldKey(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_IGNORE @@ -1862,7 +1272,7 @@ func TestDedupBuildIgnoreOnlyMarksCandidateOwnOldKey(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, 1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, tt.newKeys, tt.oldKeys) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1876,7 +1286,7 @@ func TestDedupBuildIgnoreOnlyMarksCandidateOwnOldKey(t *testing.T) { func TestDedupBuildIgnorePrefersOriginalKeyOwner(t *testing.T) { for _, oldKeys := range [][]int32{{1, 2}, {2, 1}} { t.Run(strings.Join([]string{strconv.Itoa(int(oldKeys[0])), strconv.Itoa(int(oldKeys[1]))}, "_"), func(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_IGNORE @@ -1888,7 +1298,7 @@ func TestDedupBuildIgnorePrefersOriginalKeyOwner(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, 1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{2, 2}, oldKeys) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1903,7 +1313,7 @@ func TestDedupBuildIgnorePrefersOriginalKeyOwner(t *testing.T) { } func TestDedupBuildIgnoreRebuildsAfterOwnerReplacement(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) hb.IsDedup = true hb.OnDuplicateAction = plan.Node_IGNORE @@ -1915,7 +1325,7 @@ func TestDedupBuildIgnoreRebuildsAfterOwnerReplacement(t *testing.T) { require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, 1, -1, nil, proc)) bat := makeIntKeyValueBatch(proc, []int32{2, 1, 2}, []int32{1, 3, 2}) - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) hb.InputBatchRowCount = bat.RowCount() bat.Clean(proc.Mp()) @@ -1931,7 +1341,7 @@ func TestDedupBuildIgnoreRebuildsAfterOwnerReplacement(t *testing.T) { } func TestBuildHashmapErrorDoesNotLeakIterators(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) mp := mpool.MustNewZero() proc := testutil.NewProcessWithMPool(t, "", mp) @@ -1956,13 +1366,13 @@ func TestBuildHashmapErrorDoesNotLeakIterators(t *testing.T) { intBat := batch.New([]string{"col"}) intBat.SetVector(0, intVec) intBat.SetRowCount(1) - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) } func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -1971,7 +1381,7 @@ func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { // First build: should allocate uniqueSels hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, true, proc)) require.NotNil(t, hb.uniqueSels) require.Greater(t, cap(hb.uniqueSels), 0) @@ -1986,7 +1396,7 @@ func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { hb.InputBatchRowCount = bat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, true, proc)) require.NotNil(t, hb.uniqueSels) require.Greater(t, len(hb.uniqueSels), 0) @@ -1994,7 +1404,7 @@ func TestBuildHashmapReuseUniqueSelsBuffer(t *testing.T) { } func TestBuildHashmapDoesNotCreateUniqueSelsWhenNotNeeded(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) @@ -2002,13 +1412,13 @@ func TestBuildHashmapDoesNotCreateUniqueSelsWhenNotNeeded(t *testing.T) { defer bat.Clean(proc.Mp()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.Nil(t, hb.uniqueSels, "should not allocate uniqueSels when needUniqueVec is false") } func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // Build once to create cached str iterator. @@ -2016,7 +1426,7 @@ func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { bat := makeStrBatch(t, 4, proc) defer bat.Clean(proc.Mp()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) @@ -2032,7 +1442,7 @@ func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { hb.InputBatchRowCount = bat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) rv := reflect.ValueOf(hb.cachedStrIterator).Elem() @@ -2042,7 +1452,7 @@ func TestCachedStrIteratorOwnerClearedBeforeReuse(t *testing.T) { } func TestSwitchKeyTypeCreatesCorrectIterator(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // Build int first. @@ -2050,7 +1460,7 @@ func TestSwitchKeyTypeCreatesCorrectIterator(t *testing.T) { intBat := makeIntBatch(t, 2, proc) defer intBat.Clean(proc.Mp()) hb.InputBatchRowCount = intBat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2063,13 +1473,13 @@ func TestSwitchKeyTypeCreatesCorrectIterator(t *testing.T) { hb.InputBatchRowCount = strBat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(strBat, proc)) + require.NoError(t, hb.CopyBuildBatch(strBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) } func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // Build once to create cached int iterator and bind to map A. @@ -2077,7 +1487,7 @@ func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { bat := makeIntBatch(t, 4, proc) defer bat.Clean(proc.Mp()) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2093,7 +1503,7 @@ func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { hb.InputBatchRowCount = bat.RowCount() hb.Batches.Reset() hb.Batches.Buf = nil - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) // Owner should now be non-nil and point to the new map (i.e., not staleMap). @@ -2104,7 +1514,7 @@ func TestCachedIteratorOwnerClearedBeforeReuse(t *testing.T) { } func TestFreeThenBuildRepopulatesCache(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) // First build to populate cache. @@ -2114,7 +1524,7 @@ func TestFreeThenBuildRepopulatesCache(t *testing.T) { intBat.SetVector(0, intVec) intBat.SetRowCount(2) hb.InputBatchRowCount = intBat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) @@ -2126,7 +1536,7 @@ func TestFreeThenBuildRepopulatesCache(t *testing.T) { // Build again after Free should succeed and repopulate cache. require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = intBat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(intBat, proc)) + require.NoError(t, hb.CopyBuildBatch(intBat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedIntIterator) } @@ -2148,7 +1558,7 @@ func (f failingExecutor) ResetForNextQuery() {} // Benchmarks: cached vs new iterator paths for int/str. func BenchmarkBuildHashmapCachedInt(b *testing.B) { proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) data := makeIntBatch(b, 1024, proc) defer data.Clean(proc.Mp()) @@ -2156,7 +1566,7 @@ func BenchmarkBuildHashmapCachedInt(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Reset(proc, true) } @@ -2164,7 +1574,7 @@ func BenchmarkBuildHashmapCachedInt(b *testing.B) { func BenchmarkBuildHashmapCachedStr(b *testing.B) { proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) data := makeStrBatch(b, 1024, proc) defer data.Clean(proc.Mp()) @@ -2172,7 +1582,7 @@ func BenchmarkBuildHashmapCachedStr(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Reset(proc, true) } @@ -2236,10 +1646,10 @@ func BenchmarkBuildHashmapColdInt(b *testing.B) { defer data.Clean(proc.Mp()) b.ResetTimer() for i := 0; i < b.N; i++ { - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_int32.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Free(proc) } @@ -2251,17 +1661,172 @@ func BenchmarkBuildHashmapColdStr(b *testing.B) { defer data.Clean(proc.Mp()) b.ResetTimer() for i := 0; i < b.N; i++ { - hb := &HashmapBuilder{} + hb := newTestHashmapBuilder(b) require.NoError(b, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) hb.InputBatchRowCount = data.RowCount() - require.NoError(b, hb.Batches.CopyIntoBatches(data, proc)) + require.NoError(b, hb.CopyBuildBatch(data, proc)) require.NoError(b, hb.BuildHashmap(false, false, false, proc)) hb.Free(proc) } } +func BenchmarkCopyBuildBatchAccounting(b *testing.B) { + const capBytes = uint64(256 << 20) + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType(), types.T_varchar.ToType()}, + true, + colexec.DefaultBatchSize, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err := budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + if err != nil { + b.Fatal(err) + } + account, err := registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + hb := &HashmapBuilder{} + if err = hb.SetAllocationAccount(account); err != nil { + b.Fatal(err) + } + hb.setBudget(generation) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err = hb.copyBuildBatch(input, proc); err != nil { + b.Fatal(err) + } + hb.cleanBatches(proc) + } + b.StopTimer() + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + if err = hb.ClearAllocationAccount(account); err != nil { + b.Fatal(err) + } + if _, _, err = registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } +} + +// BenchmarkResidentHashBuildAccounting compares a local physical account with +// the production shared budget controller across the complete resident owner +// closure: copied batches, key expression, hash cells/descriptors, and terminal +// release all run on every iteration. The builder intentionally has no +// unaccounted mode. +// The 32-row case models high-frequency TP statements; 8,192 rows exercises a +// full physical batch without entering spill. +func BenchmarkResidentHashBuildAccounting(b *testing.B) { + const capBytes = uint64(512 << 20) + for _, rows := range []int{32, colexec.DefaultBatchSize} { + for _, stringKey := range []bool{false, true} { + kind := "int" + if stringKey { + kind = "varchar" + } + for _, controlled := range []bool{false, true} { + mode := "local-account" + if controlled { + mode = "budget-controlled" + } + b.Run(fmt.Sprintf("%s/rows-%d/%s", kind, rows, mode), func(b *testing.B) { + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + var input *batch.Batch + var keyType types.Type + if stringKey { + input = makeStrBatch(b, rows, proc) + keyType = types.T_varchar.ToType() + } else { + input = makeIntBatch(b, rows, proc) + keyType = types.T_int32.ToType() + } + defer input.Clean(proc.Mp()) + + var generation *process.HashBuildBudgetGeneration + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + if err != nil { + b.Fatal(err) + } + var account *mpool.AllocationAccount + if controlled { + budget := process.MustNewHashBuildBudget(capBytes, capBytes) + generation, err = budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + account, err = registry.OpenWithController(capBytes, generation) + if err != nil { + b.Fatal(err) + } + } else { + account, err = registry.Open(capBytes) + if err != nil { + b.Fatal(err) + } + } + + b.ReportAllocs() + b.SetBytes(int64(input.Size())) + b.ResetTimer() + for range b.N { + hb := &HashmapBuilder{} + if controlled { + hb.SetBudget(generation) + } + if err := hb.SetAllocationAccount(account); err != nil { + b.Fatal(err) + } + if err := hb.Prepare( + []*plan.Expr{newExpr(0, keyType)}, + -1, + -1, + nil, + proc, + ); err != nil { + b.Fatal(err) + } + hb.InputBatchRowCount = input.RowCount() + if err := hb.CopyBuildBatch(input, proc); err != nil { + b.Fatal(err) + } + if err := hb.BuildHashmap(false, false, false, proc); err != nil { + b.Fatal(err) + } + hb.Free(proc) + } + b.StopTimer() + if account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", account.Snapshot().Used) + } + if _, _, err := registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + if controlled { + if generation.Used() != 0 { + b.Fatalf("generation used = %d", generation.Used()) + } + generation.Close() + } + }) + } + } + } +} + func TestExtractRestoreCachedIterators(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) mp := mpool.MustNewZero() intMap, err := hashmap.NewIntHashMap(false, mp) @@ -2288,7 +1853,7 @@ func TestExtractRestoreCachedIterators(t *testing.T) { } func TestStrIteratorLargeStringTriggersPrune(t *testing.T) { - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) require.NoError(t, hb.Prepare([]*plan.Expr{newExpr(0, types.T_varchar.ToType())}, -1, -1, nil, proc)) @@ -2301,7 +1866,7 @@ func TestStrIteratorLargeStringTriggersPrune(t *testing.T) { bat.SetRowCount(1) hb.InputBatchRowCount = bat.RowCount() - require.NoError(t, hb.Batches.CopyIntoBatches(bat, proc)) + require.NoError(t, hb.CopyBuildBatch(bat, proc)) require.NoError(t, hb.BuildHashmap(false, false, false, proc)) require.NotNil(t, hb.cachedStrIterator) @@ -2316,7 +1881,7 @@ func TestStrIteratorLargeStringTriggersPrune(t *testing.T) { // curVecs or UniqueJoinKeys contained nil pointers. func TestResetWithNilPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Test case 1: curVecs with nil pointers and needDupVec = true hb.needDupVec = true @@ -2340,7 +1905,7 @@ func TestResetWithNilPointers(t *testing.T) { // TestResetWithMixedNilAndValidPointers tests Reset() with a mix of nil and valid vectors func TestResetWithMixedNilAndValidPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Create some valid vectors vec1 := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) @@ -2363,7 +1928,7 @@ func TestResetWithMixedNilAndValidPointers(t *testing.T) { // TestFreeWithNilPointers tests that Free() handles nil pointers gracefully func TestFreeWithNilPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Test case: UniqueJoinKeys with nil pointers hb.UniqueJoinKeys = make([]*vector.Vector, 3) @@ -2380,7 +1945,7 @@ func TestFreeWithNilPointers(t *testing.T) { // TestFreeWithMixedNilAndValidPointers tests Free() with a mix of nil and valid vectors func TestFreeWithMixedNilAndValidPointers(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - var hb HashmapBuilder + hb := newTestHashmapBuilder(t) // Create some valid vectors vec1 := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) diff --git a/pkg/sql/colexec/hashbuild/nonspill_benchmark_test.go b/pkg/sql/colexec/hashbuild/nonspill_benchmark_test.go index 68c36e4a3cea3..fd2be704c0bc4 100644 --- a/pkg/sql/colexec/hashbuild/nonspill_benchmark_test.go +++ b/pkg/sql/colexec/hashbuild/nonspill_benchmark_test.go @@ -163,6 +163,8 @@ func benchmarkHashBuildNonSpillE2E( budget, err := proc.GetHashBuildBudget() require.NoError(b, err) require.NotNil(b, budget) + registry, err := budget.AllocationAccountRegistry() + require.NoError(b, err) keyExpr := nonSpillBenchmarkColumnExpr(0, keyType) if computed { keyExpr, err = plan2.BindFuncExprImplByPlanExpr( @@ -185,6 +187,12 @@ func benchmarkHashBuildNonSpillE2E( for i := 0; i < b.N; i++ { child := colexec.NewMockOperator().WithBatchs(inputs) arg := newNonSpillBenchmarkHashBuild(keyExpr, child) + account, openErr := registry.OpenWithController( + budget.Snapshot().Cap, + budget, + ) + require.NoError(b, openErr) + require.NoError(b, arg.SetAllocationAccount(account)) require.NoError(b, child.Prepare(proc)) require.NoError(b, arg.Prepare(proc)) before := budget.Snapshot() @@ -221,6 +229,9 @@ func benchmarkHashBuildNonSpillE2E( joinMap.Free() arg.Reset(proc, false, nil) arg.Free(proc, false, nil) + require.NoError(b, arg.ClearAllocationAccount(account)) + _, _, terminalErr := registry.CompleteTerminal(account) + require.NoError(b, terminalErr) child.Release() require.Zero(b, budget.Used()) proc.GetMessageBoard().Reset() diff --git a/pkg/sql/colexec/hashbuild/pressure.go b/pkg/sql/colexec/hashbuild/pressure.go new file mode 100644 index 0000000000000..851f7f05b9592 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/pressure.go @@ -0,0 +1,213 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "errors" + "fmt" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/process" +) + +// MemoryPressureReason is the single classification used by HashBuild and all +// spilled join consumers. Only Capacity may enter reclaim/spill/reduce/degrade +// control flow. Sealed and invariant failures are lifecycle bugs and must +// remain terminal. HashBuildBudgetError now exposes disjoint lifecycle and +// capacity identities; this classifier is the one control-flow boundary for +// physical-allocation and spill-resource failures. +type MemoryPressureReason uint8 + +const ( + MemoryPressureNone MemoryPressureReason = iota + MemoryPressureCapacity + MemoryPressureSealed + MemoryPressureMismatch + MemoryPressureAllocatorLimit + MemoryPressureInvariant + MemoryPressureInvalid + MemoryPressureMinimumUnit + MemoryPressureSpillDiskLimit + MemoryPressureSpillFDLimit +) + +func MemoryPressureReasonOf(err error) MemoryPressureReason { + if err == nil { + return MemoryPressureNone + } + var minimum *MinimumAllocationPressureError + if errors.As(err, &minimum) { + return MemoryPressureMinimumUnit + } + + // Physical-account ownership and lifecycle failures dominate any joined + // logical budget error. Otherwise a sealed or invariant account could be + // mistaken for retryable memory pressure. + switch mpool.AllocationFailureReasonOf(err) { + case mpool.AllocationFailureCapacity: + return MemoryPressureCapacity + case mpool.AllocationFailureSealed, + mpool.AllocationFailureSuspended: + return MemoryPressureSealed + case mpool.AllocationFailureMismatch: + return MemoryPressureMismatch + case mpool.AllocationFailureAllocatorLimit: + return MemoryPressureAllocatorLimit + case mpool.AllocationFailureInvariant: + return MemoryPressureInvariant + } + + var budgetErr *process.HashBuildBudgetError + if errors.As(err, &budgetErr) { + switch budgetErr.Kind { + case process.HashBuildBudgetErrorAdmission: + switch budgetErr.Component { + case process.HashBuildBudgetComponentMemory: + return MemoryPressureCapacity + case process.HashBuildBudgetComponentSpillDisk: + return MemoryPressureSpillDiskLimit + case process.HashBuildBudgetComponentSpillFD: + return MemoryPressureSpillFDLimit + default: + return MemoryPressureInvalid + } + case process.HashBuildBudgetErrorClosed: + return MemoryPressureSealed + case process.HashBuildBudgetErrorInvalid, + process.HashBuildBudgetErrorCeilingMissing: + return MemoryPressureInvalid + default: + return MemoryPressureInvalid + } + } + + // Resource-ledger helpers may still return a bare lifecycle sentinel. + if errors.Is(err, process.ErrHashBuildBudgetClosed) { + return MemoryPressureSealed + } + if errors.Is(err, process.ErrHashBuildBudgetInvalid) || + errors.Is(err, process.ErrHashBuildCeilingMissing) { + return MemoryPressureInvalid + } + return MemoryPressureNone +} + +func IsRetryableMemoryCapacity(err error) bool { + return MemoryPressureReasonOf(err) == MemoryPressureCapacity +} + +// MinimumAllocationPressureError means the operation has already reclaimed +// optional storage and reduced itself to one indivisible input unit. It does +// not unwrap the last capacity error: callers must not mistake the terminal +// boundary for another retryable admission failure. +type MinimumAllocationPressureError struct { + Owner string + Site string + Response string + Used uint64 + Limit uint64 +} + +func (e *MinimumAllocationPressureError) Error() string { + if e == nil { + return "minimum allocation cannot be admitted" + } + return fmt.Sprintf( + "minimum allocation cannot be admitted: owner=%s site=%s response=%s used=%d limit=%d", + e.Owner, + e.Site, + e.Response, + e.Used, + e.Limit, + ) +} + +func NewMinimumAllocationPressureError( + owner string, + site string, + account *mpool.AllocationAccount, +) error { + err := &MinimumAllocationPressureError{ + Owner: owner, + Site: site, + Response: "reclaim/reduce/degrade exhausted", + } + if account != nil { + snapshot := account.Snapshot() + err.Used = snapshot.Used + err.Limit = snapshot.Limit + } + return err +} + +// PressureProgress is the monotonic proof required before retrying one +// logical operation. A retry is legal only after memory was reclaimed, spill +// state advanced, the input unit shrank, or optional work was disabled. +type PressureProgress struct { + Used uint64 + SpillEpoch uint64 + InputUnits int + OptionalDisabled bool +} + +type PressureRetryGuard struct { + previous PressureProgress + attempts int + limit int +} + +func NewPressureRetryGuard(initial PressureProgress, limit int) *PressureRetryGuard { + if limit <= 0 { + limit = 64 + } + return &PressureRetryGuard{previous: initial, limit: limit} +} + +func (g *PressureRetryGuard) Advance(next PressureProgress) error { + if g == nil || next.InputUnits < 0 { + return process.ErrHashBuildBudgetInvalid + } + if g.attempts >= g.limit { + return errors.Join( + process.ErrHashBuildBudgetInvalid, + moerr.NewInternalErrorNoCtx( + "memory-pressure retry limit exceeded", + ), + ) + } + progress := next.Used < g.previous.Used || + next.SpillEpoch > g.previous.SpillEpoch || + (g.previous.InputUnits > 0 && next.InputUnits < g.previous.InputUnits) || + (!g.previous.OptionalDisabled && next.OptionalDisabled) + if !progress { + return errors.Join( + process.ErrHashBuildBudgetInvalid, + moerr.NewInternalErrorNoCtx( + "memory-pressure retry made no progress", + ), + ) + } + g.previous = next + g.attempts++ + return nil +} + +func (g *PressureRetryGuard) Attempts() int { + if g == nil { + return 0 + } + return g.attempts +} diff --git a/pkg/sql/colexec/hashbuild/pressure_test.go b/pkg/sql/colexec/hashbuild/pressure_test.go new file mode 100644 index 0000000000000..799a3a08ffc22 --- /dev/null +++ b/pkg/sql/colexec/hashbuild/pressure_test.go @@ -0,0 +1,88 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "errors" + "fmt" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func TestMemoryPressureReasonSeparatesCapacityFromLifecycle(t *testing.T) { + tests := []struct { + err error + reason MemoryPressureReason + }{ + {nil, MemoryPressureNone}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission}, MemoryPressureInvalid}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentMemory}, MemoryPressureCapacity}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillDisk}, MemoryPressureSpillDiskLimit}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorAdmission, Component: process.HashBuildBudgetComponentSpillFD}, MemoryPressureSpillFDLimit}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorClosed}, MemoryPressureSealed}, + {&process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorInvalid}, MemoryPressureInvalid}, + {fmt.Errorf("wrapped: %w", process.ErrHashBuildBudgetAdmission), MemoryPressureNone}, + {mpool.ErrAllocationAccountCapacity, MemoryPressureCapacity}, + {moerr.NewMPoolCapacityNoCtxf("test"), MemoryPressureCapacity}, + {mpool.ErrAllocationMetadataSlots, MemoryPressureCapacity}, + {mpool.ErrAllocationAccountSealed, MemoryPressureSealed}, + {mpool.ErrAllocationAccountMismatch, MemoryPressureMismatch}, + {mpool.ErrAllocationAllocatorLimit, MemoryPressureAllocatorLimit}, + {mpool.ErrAllocationAccountInvariant, MemoryPressureInvariant}, + {NewMinimumAllocationPressureError("hashbuild", "spill", nil), MemoryPressureMinimumUnit}, + } + for _, test := range tests { + require.Equal(t, test.reason, MemoryPressureReasonOf(test.err)) + require.Equal(t, test.reason == MemoryPressureCapacity, IsRetryableMemoryCapacity(test.err)) + } + + memoryAdmission := &process.HashBuildBudgetError{ + Kind: process.HashBuildBudgetErrorAdmission, + Component: process.HashBuildBudgetComponentMemory, + } + for _, test := range []struct { + err error + reason MemoryPressureReason + }{ + {errors.Join(memoryAdmission, mpool.ErrAllocationAccountInvariant), MemoryPressureInvariant}, + {errors.Join(memoryAdmission, mpool.ErrAllocationAccountInvalid), MemoryPressureInvariant}, + {errors.Join(memoryAdmission, mpool.ErrAllocationAccountSealed), MemoryPressureSealed}, + } { + require.Equal(t, test.reason, MemoryPressureReasonOf(test.err)) + require.False(t, IsRetryableMemoryCapacity(test.err)) + } +} + +func TestPressureRetryGuardRequiresMonotonicProgress(t *testing.T) { + initial := PressureProgress{Used: 100, SpillEpoch: 1, InputUnits: 16} + for _, next := range []PressureProgress{ + {Used: 99, SpillEpoch: 1, InputUnits: 16}, + {Used: 99, SpillEpoch: 2, InputUnits: 16}, + {Used: 99, SpillEpoch: 2, InputUnits: 8}, + {Used: 99, SpillEpoch: 2, InputUnits: 8, OptionalDisabled: true}, + } { + guard := NewPressureRetryGuard(initial, 1) + require.NoError(t, guard.Advance(next)) + require.Equal(t, 1, guard.Attempts()) + require.Error(t, guard.Advance(next), "the retry limit remains fail-closed") + } + guard := NewPressureRetryGuard(initial, 4) + require.Error(t, guard.Advance(initial)) + require.Zero(t, guard.Attempts()) +} diff --git a/pkg/sql/colexec/hashbuild/recovery_projection.go b/pkg/sql/colexec/hashbuild/recovery_projection.go new file mode 100644 index 0000000000000..5b3ae5b171f7e --- /dev/null +++ b/pkg/sql/colexec/hashbuild/recovery_projection.go @@ -0,0 +1,321 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashbuild + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/vm/process" +) + +const recoveryCapacityQuantum = uint64(64 << 10) + +func recoveryCheckedAdd(left, right uint64) (uint64, error) { + if left > math.MaxUint64-right { + return 0, process.ErrHashBuildBudgetInvalid + } + return left + right, nil +} + +func recoveryCheckedMul(left, right uint64) (uint64, error) { + if left != 0 && right > math.MaxUint64/left { + return 0, process.ErrHashBuildBudgetInvalid + } + return left * right, nil +} + +func roundRecoveryCapacity(size uint64) (uint64, error) { + if size == 0 { + return 0, nil + } + if size > math.MaxUint64-(recoveryCapacityQuantum-1) { + return 0, process.ErrHashBuildBudgetInvalid + } + return (size + recoveryCapacityQuantum - 1) &^ (recoveryCapacityQuantum - 1), nil +} + +// recoveryBatchProjection is the maximum logical batch materialization that +// one CopyIntoBatches call can add, plus the new logical size of its partial +// tail. Physical copy allocation is admitted independently by AllocationAccount. +type recoveryBatchProjection struct { + maxRows int + maxSelected uint64 + nextTailSelected uint64 + columns int +} + +func (ctr *container) ensureRecoveryCapacity( + proc *process.Process, + exprs []*plan.Expr, + projection recoveryBatchProjection, + analyzer process.Analyzer, +) error { + if ctr.recoveryCapacity == nil || projection.maxRows <= 0 { + return process.ErrHashBuildBudgetInvalid + } + expression := ctr.expressionRecoveryPeak + if projection.maxRows > ctr.expressionRecoveryRows { + var err error + expression, err = expressionRecoveryBytes( + proc, + exprs, + projection.maxRows, + ctr.hashmapBuilder.needDupVec, + ) + if err != nil { + return err + } + } + spill, err := spillRecoveryPeak(projection) + if err != nil { + return err + } + expression = max(expression, ctr.expressionRecoveryPeak) + spill = max(spill, ctr.spillRecoveryPeak) + target, err := recoveryCheckedAdd(expression, spill) + if err != nil { + return err + } + target, err = roundRecoveryCapacity(target) + if err != nil { + return err + } + if err = ctr.recoveryCapacity.EnsureCapacity(target); err != nil { + if analyzer != nil { + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillRecoveryReserveRejects", 1) + } + return err + } + ctr.expressionRecoveryPeak = expression + ctr.expressionRecoveryRows = max( + ctr.expressionRecoveryRows, projection.maxRows) + ctr.spillRecoveryPeak = spill + if analyzer != nil { + analyzer.GetOpStats().SetMaxExtraStat( + "HashBuildSpillRecoveryReservedBytes", hashBuildStatInt64(target)) + } + return nil +} + +func (hb *HashmapBuilder) projectRetainedRecovery( + src *batch.Batch, +) (recoveryBatchProjection, error) { + if src == nil || src.RowCount() <= 0 { + return recoveryBatchProjection{}, process.ErrHashBuildBudgetInvalid + } + projection := recoveryBatchProjection{columns: len(src.Vecs)} + last := len(hb.Batches.Buf) - 1 + hasTail := last >= 0 && hb.Batches.Buf[last] != nil && + hb.Batches.Buf[last].RowCount() != colexec.DefaultBatchSize + + // A full source is copied as its own destination and swapped ahead of an + // existing partial tail, so that tail remains unchanged. + if src.RowCount() == colexec.DefaultBatchSize { + selected, err := projectedSelectedRange(src, 0, src.RowCount()) + if err != nil { + return recoveryBatchProjection{}, err + } + projection.maxRows = src.RowCount() + projection.maxSelected = selected + if hasTail { + projection.nextTailSelected = hb.retainedSpillTailSelected + } + return projection, nil + } + + offset := 0 + if hasTail { + tailRows := hb.Batches.Buf[last].RowCount() + appendRows := min(colexec.DefaultBatchSize-tailRows, src.RowCount()) + appended, err := projectedSelectedRange(src, 0, appendRows) + if err != nil { + return recoveryBatchProjection{}, err + } + combined, err := recoveryCheckedAdd(hb.retainedSpillTailSelected, appended) + if err != nil { + return recoveryBatchProjection{}, err + } + projection.maxRows = tailRows + appendRows + projection.maxSelected = combined + if projection.maxRows < colexec.DefaultBatchSize { + projection.nextTailSelected = combined + } + offset = appendRows + } + + for offset < src.RowCount() { + rows := min(colexec.DefaultBatchSize, src.RowCount()-offset) + selected, err := projectedSelectedRange(src, offset, rows) + if err != nil { + return recoveryBatchProjection{}, err + } + if rows > projection.maxRows { + projection.maxRows = rows + } + if selected > projection.maxSelected { + projection.maxSelected = selected + } + if rows < colexec.DefaultBatchSize { + projection.nextTailSelected = selected + } else if offset+rows == src.RowCount() { + projection.nextTailSelected = 0 + } + offset += rows + } + return projection, nil +} + +func projectedSelectedRange(src *batch.Batch, start, rows int) (uint64, error) { + if src == nil || start < 0 || rows < 0 || start > src.RowCount() || + rows > src.RowCount()-start { + return 0, process.ErrHashBuildBudgetInvalid + } + var total uint64 + for _, vec := range src.Vecs { + if vec == nil { + return 0, process.ErrHashBuildBudgetInvalid + } + typeSize := vec.GetType().TypeSize() + if typeSize < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + descriptors, err := recoveryCheckedMul(uint64(rows), uint64(typeSize)) + if err != nil { + return 0, err + } + total, err = recoveryCheckedAdd(total, descriptors) + if err != nil { + return 0, err + } + _, payload, err := unionBatchAreaProjection(vec, start, rows) + if err != nil { + return 0, err + } + total, err = recoveryCheckedAdd(total, payload) + if err != nil { + return 0, err + } + } + return total, nil +} + +// projectedUnionArea mirrors the logical payload that UnionBatch retains and +// UnionInt32 later materializes. The physical result is retained for focused +// allocation tests and documents the const/full-vector distinction. +func unionBatchAreaProjection( + src *vector.Vector, + start, rows int, +) (physical int, selected uint64, err error) { + if src == nil || !src.GetType().IsVarlen() { + return 0, 0, nil + } + if start < 0 || rows < 0 || start > src.Length() || rows > src.Length()-start { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + if rows == 0 || len(src.GetArea()) == 0 { + return 0, 0, nil + } + if src.IsConst() { + payload, err := selectedVarlenaPayload(src, 0, 1) + if err != nil || payload > math.MaxInt { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + selected, err = recoveryCheckedMul(payload, uint64(rows)) + return int(payload), selected, err + } + if start == 0 && rows == src.Length() && src.VarlenaAreaIsDisjoint() { + return len(src.GetArea()), uint64(len(src.GetArea())), nil + } + payload, err := selectedVarlenaPayload(src, start, rows) + if err != nil || payload > math.MaxInt { + return 0, 0, process.ErrHashBuildBudgetInvalid + } + if start == 0 && rows == src.Length() { + return len(src.GetArea()), payload, nil + } + return int(payload), payload, nil +} + +func selectedVarlenaPayload(src *vector.Vector, start, rows int) (uint64, error) { + values, _ := vector.MustVarlenaRawData(src) + end := start + rows + if start < 0 || rows < 0 || end > len(values) { + return 0, process.ErrHashBuildBudgetInvalid + } + var payload uint64 + for row := start; row < end; row++ { + index := row + if src.IsConst() { + index = 0 + } + if src.GetNulls().Contains(uint64(index)) || values[index].IsSmall() { + continue + } + _, length := values[index].OffsetLen() + var err error + payload, err = recoveryCheckedAdd(payload, uint64(length)) + if err != nil { + return 0, err + } + } + return payload, nil +} + +func spillRecoveryPeak(projection recoveryBatchProjection) (uint64, error) { + if projection.maxRows <= 0 || projection.columns < 0 { + return 0, process.ErrHashBuildBudgetInvalid + } + rowScratch, err := recoveryCheckedMul(uint64(projection.maxRows), 12) + if err != nil { + return 0, err + } + // Null/grouping bitmaps, vector growth rounding, and record framing have a + // bounded per-column allowance. The selected batch and marshal buffer are + // simultaneously live, hence two copies of the logical materialization. + columnSlack, err := recoveryCheckedMul(uint64(projection.columns), 16<<10) + if err != nil { + return 0, err + } + slack, err := recoveryCheckedAdd(64<<10, columnSlack) + if err != nil { + return 0, err + } + selected, err := recoveryCheckedAdd(projection.maxSelected, slack) + if err != nil { + return 0, err + } + total, err := recoveryCheckedAdd(rowScratch, selected) + if err != nil { + return 0, err + } + columnMarshalSlack, err := recoveryCheckedMul(uint64(projection.columns), 128) + if err != nil { + return 0, err + } + marshalSlack, err := recoveryCheckedAdd(64<<10, columnMarshalSlack) + if err != nil { + return 0, err + } + marshal, err := recoveryCheckedAdd(selected, marshalSlack) + if err != nil { + return 0, err + } + return recoveryCheckedAdd(total, marshal) +} diff --git a/pkg/sql/colexec/hashbuild/spill.go b/pkg/sql/colexec/hashbuild/spill.go index 88a1d4a980ba5..73baf2db3e290 100644 --- a/pkg/sql/colexec/hashbuild/spill.go +++ b/pkg/sql/colexec/hashbuild/spill.go @@ -15,8 +15,6 @@ package hashbuild import ( - "bytes" - "errors" "fmt" "io" "math" @@ -24,10 +22,11 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/fileservice" + "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -35,464 +34,31 @@ import ( const ( spillNumBuckets = 32 spillMagic = 0x12345678DEADBEEF + spillBufferSize = 8192 // Buffer 8192 rows before flushing // Serialized records are accumulated per bucket across source batches. // Allocation is admitted lazily against the lifecycle scratch lease and // falls back to direct writes when the hard budget has no headroom. spillWriteCoalesceSize = 64 << 10 ) -func spillCheckedAdd(total, value uint64) (uint64, error) { - if total > math.MaxUint64-value { - return 0, process.ErrHashBuildBudgetInvalid - } - return total + value, nil -} - -func spillCheckedMul(left, right uint64) (uint64, error) { - if left != 0 && right > math.MaxUint64/left { - return 0, process.ErrHashBuildBudgetInvalid - } - return left * right, nil -} - -func spillCapacityReplacementOverlap(rows, keys, hashCap, rowIDCap, keyCap int) (uint64, error) { - var overlap uint64 - add := func(required, current int, width uint64) error { - if required < 0 || current < 0 { - return process.ErrHashBuildBudgetInvalid - } - if required <= current { - return nil - } - old, err := spillCheckedMul(uint64(current), width) - if err != nil { - return err - } - overlap, err = spillCheckedAdd(overlap, old) - return err - } - if err := add(keys, keyCap, 8); err != nil { - return 0, err - } - if err := add(rows, hashCap, 8); err != nil { - return 0, err - } - if err := add(rows, rowIDCap, 4); err != nil { - return 0, err - } - return overlap, nil -} - -// spillMaterializedBytes models the batch that spillBatchBounded creates with -// UnionInt32. It follows vector materialization semantics instead of retained -// capacity or stale logical length: fixed-width descriptors are per output -// row, null payload is skipped, and const varlen payload is copied once. -func spillMaterializedBytes(bat *batch.Batch) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, nil - } - liveRows := uint64(bat.RowCount()) - var materialized uint64 - for _, vec := range bat.Vecs { - if vec == nil { - return 0, process.ErrHashBuildBudgetInvalid - } - typeSize := vec.GetType().TypeSize() - if typeSize < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - descriptors, err := spillCheckedMul(liveRows, uint64(typeSize)) - if err != nil { - return 0, err - } - if materialized, err = spillCheckedAdd(materialized, descriptors); err != nil { - return 0, err - } - if !vec.GetType().IsVarlen() || vec.IsConstNull() { - continue - } - - values, _ := vector.MustVarlenaRawData(vec) - valueRows := liveRows - if vec.IsConst() { - valueRows = 1 - } - if valueRows == 0 || valueRows > uint64(len(values)) { - return 0, process.ErrHashBuildBudgetInvalid - } - var livePayload uint64 - hasNull := !vec.GetNulls().EmptyByFlag() - for row := uint64(0); row < valueRows; row++ { - if hasNull && vec.GetNulls().Contains(row) { - continue - } - if values[row].IsSmall() { - continue - } - _, length := values[row].OffsetLen() - if livePayload, err = spillCheckedAdd(livePayload, uint64(length)); err != nil { - return 0, err - } - } - - if materialized, err = spillCheckedAdd(materialized, livePayload); err != nil { - return 0, err - } - } - return materialized, nil -} - -func spillMarshalSlack(columns uint64) (uint64, error) { - const ( - fixedSlack = uint64(64 << 10) - perColumnSlack = uint64(128) - ) - if columns > (math.MaxUint64-fixedSlack)/perColumnSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return fixedSlack + columns*perColumnSlack, nil -} - -func spillMaterializationSlack(columns uint64) (uint64, error) { - const ( - fixedSlack = uint64(64 << 10) - perColumnSlack = uint64(16 << 10) - ) - if columns > (math.MaxUint64-fixedSlack)/perColumnSlack { - return 0, process.ErrHashBuildBudgetInvalid - } - return fixedSlack + columns*perColumnSlack, nil -} - -// spillPeakBudgetFor accounts each simultaneously live component explicitly. -// inputBytes is zero for a retained batch whose source reservation is already -// owned by HashBuild. -func spillPeakBudgetFor(rows, inputBytes, selectedBytes, columns uint64) (uint64, error) { - rowScratch, err := spillCheckedMul(rows, 12) // hashes + one row-id array - if err != nil { - return 0, err - } - total, err := spillCheckedAdd(rowScratch, inputBytes) - if err != nil { - return 0, err - } - if total, err = spillCheckedAdd(total, selectedBytes); err != nil { - return 0, err - } - // MarshalBinary creates one serialized payload. The selected estimate - // already includes its fixed-width data and varlen area, so charge that - // payload once plus bounded framing/allocation slack. - marshalSlack, err := spillMarshalSlack(columns) - if err != nil { - return 0, err - } - marshalBytes, err := spillCheckedAdd(selectedBytes, marshalSlack) - if err != nil { - return 0, err - } - if total, err = spillCheckedAdd(total, marshalBytes); err != nil { - return 0, err - } - if total > uint64(^uint(0)>>1) { - return 0, process.ErrHashBuildBudgetInvalid - } - return total, nil -} - -// spillBudgetBytes admits only the actual direct-spill path for the current -// input. It never projects a hypothetical retained batch. -func spillBudgetBytes(bat *batch.Batch) (uint64, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, nil - } - rows := uint64(bat.RowCount()) - selected, err := spillMaterializedBytes(bat) - if err != nil { - return 0, err - } - materializationSlack, err := spillMaterializationSlack(uint64(len(bat.Vecs))) - if err != nil { - return 0, err - } - if selected, err = spillCheckedAdd(selected, materializationSlack); err != nil { - return 0, err - } - return spillPeakBudgetFor(rows, uint64(bat.Allocated()), selected, uint64(len(bat.Vecs))) -} - -// spillScratchBudgetBytes returns the incremental spill charge. A copied -// build batch remains covered by HashmapBuilder.batchReservations while it is -// drained, so charging its source footprint again would double count it. An -// upstream batch has no HashBuild-owned reservation and keeps the full charge. -func spillScratchBudgetBytes(bat *batch.Batch, sourceAlreadyCharged bool) (uint64, error) { - need, err := spillBudgetBytes(bat) - if err != nil || !sourceAlreadyCharged || bat == nil || bat.RowCount() <= 0 { - return need, err - } - // copyBuildBatch reconciles its retained reservation against Allocated - // (plus metadata), so only that proven charge may be subtracted here. - source := uint64(bat.Allocated()) - if source > need { - return 0, process.ErrHashBuildBudgetInvalid - } - return need - source, nil -} - -const spillRecoveryReservationQuantum = uint64(64 << 10) - -// spillRecoveryReservationBytes rounds a recovery high-water mark to a small, -// fixed allocation quantum. The lease is per HashBuild execution, not per -// batch: rounding avoids rescanning near-identical varlen batches merely -// because their payload differs by a few bytes, while keeping the bounded -// over-reservation independent of row count, fanout, and query shape. -func spillRecoveryReservationBytes(need uint64) (uint64, error) { - if need == 0 { - return 0, nil - } - if need > math.MaxUint64-(spillRecoveryReservationQuantum-1) { - return 0, process.ErrHashBuildBudgetInvalid - } - return (need + spillRecoveryReservationQuantum - 1) &^ (spillRecoveryReservationQuantum - 1), nil -} - -// spillDirectRecoveryBudgetUpper is a row-scan-free upper bound used to decide -// whether the current recovery lease can already spill an upstream batch. -// Fixed-width and const batches are exact apart from bounded slack. A regular -// varlen vector may have every descriptor reference the same physical area, so -// rows*area is the smallest representation-independent bound available without -// inspecting descriptors. Such batches are scanned exactly only when this -// conservative bound crosses the retained high-water mark; this path is cold -// for normal resident builds. -func spillDirectRecoveryBudgetUpper(bat *batch.Batch) (uint64, bool, error) { - if bat == nil || bat.RowCount() <= 0 { - return 0, false, nil - } - rows := uint64(bat.RowCount()) - var selected uint64 - hasVarlen := false - for _, vec := range bat.Vecs { - if vec == nil { - return 0, false, process.ErrHashBuildBudgetInvalid - } - typeSize := vec.GetType().TypeSize() - if typeSize < 0 { - return 0, false, process.ErrHashBuildBudgetInvalid - } - descriptors, err := spillCheckedMul(rows, uint64(typeSize)) - if err != nil { - return 0, false, err - } - selected, err = spillCheckedAdd(selected, descriptors) - if err != nil { - return 0, false, err - } - if vec.GetType().IsVarlen() && !vec.IsConstNull() { - hasVarlen = true - payloadUpper := uint64(len(vec.GetArea())) - if !vec.IsConst() { - payloadUpper, err = spillCheckedMul(payloadUpper, rows) - if err != nil { - return 0, false, err - } - } - selected, err = spillCheckedAdd(selected, payloadUpper) - if err != nil { - return 0, false, err - } - } - } - materializationSlack, err := spillMaterializationSlack(uint64(len(bat.Vecs))) - if err != nil { - return 0, false, err +func (ctr *container) dropSpillScratchBuffers() { + if cap(ctr.spillHashValues) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillHashValues) } - selected, err = spillCheckedAdd(selected, materializationSlack) - if err != nil { - return 0, false, err + if cap(ctr.spillBucketRowIds) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillBucketRowIds) } - allocated := bat.Allocated() - if allocated < 0 { - return 0, false, process.ErrHashBuildBudgetInvalid + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil } - need, err := spillPeakBudgetFor(rows, uint64(allocated), selected, uint64(len(bat.Vecs))) - return need, hasVarlen, err -} - -// spillRetainedRecoveryBudgetBytes turns the allocation projection already -// required by CopyIntoBatches into a future-drain proof. The projection tracks -// logical spill materialization rather than physical retained allocation: -// ordinary non-const descriptors can also share one retained payload while a -// later spill selection copies it once per row. Source memory itself is covered -// separately by the retained-batch reservation. -func spillRetainedRecoveryBudgetBytes(projection batchCopyProjection) (uint64, error) { - if projection.maxRetainedRows <= 0 || projection.columns < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - materializationSlack, err := spillMaterializationSlack(uint64(projection.columns)) - if err != nil { - return 0, err - } - selected, err := spillCheckedAdd( - projection.maxRetainedSelected, - materializationSlack, - ) - if err != nil { - return 0, err - } - return spillPeakBudgetFor( - uint64(projection.maxRetainedRows), - 0, - selected, - uint64(projection.columns), - ) -} - -func (ctr *container) ensureSpillRecoveryReservationBytes( - need uint64, - analyzer process.Analyzer, -) error { - if ctr.hashmapBuilder.budget == nil || need == 0 || need <= ctr.spillScratchBase { - return nil - } - var err error - if ctr.spillScratchReservation == nil { - ctr.spillScratchReservation, err = ctr.hashmapBuilder.budget.Reserve(need) - if err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildSpillRecoveryReserveRejects", 1) - return err + for i := range ctr.spillAccountedBuckets { + if ctr.spillAccountedBuckets[i] != nil { + ctr.spillAccountedBuckets[i].Free() + ctr.spillAccountedBuckets[i] = nil } - ctr.spillScratchBase = need - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillRecoveryReservedBytes", hashBuildStatInt64(need)) - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", hashBuildStatInt64(need)) - return nil - } - - grow := need - ctr.spillScratchBase - if err = ctr.spillScratchReservation.Grow(grow); err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildSpillRecoveryGrowRejects", 1) - return err } - ctr.spillScratchBase = need - analyzer.GetOpStats().AddExtraStat("HashBuildSpillRecoveryGrowCount", 1) - analyzer.GetOpStats().AddExtraStat( - "HashBuildSpillRecoveryGrowBytes", hashBuildStatInt64(grow)) - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillRecoveryReservedBytes", hashBuildStatInt64(need)) - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", hashBuildStatInt64(need)) - return nil -} - -func (ctr *container) ensureDirectSpillRecovery( - bat *batch.Batch, - analyzer process.Analyzer, -) error { - upper, hasVarlen, err := spillDirectRecoveryBudgetUpper(bat) - if err != nil { - return err - } - if upper <= ctr.spillScratchBase { - return nil - } - need := upper - if hasVarlen { - // The cheap upper includes dead/null varlena area. Pay the exact live-row - // scan only when a larger lease may be required, avoiding both false - // admission failure and a scan on steady-state batches. - need, err = spillBudgetBytes(bat) - if err != nil { - return err - } - if need <= ctr.spillScratchBase { - return nil - } - } - need, err = spillRecoveryReservationBytes(need) - if err != nil { - return err - } - return ctr.ensureSpillRecoveryReservationBytes(need, analyzer) -} - -func (ctr *container) ensureRetainedSpillRecovery( - projection batchCopyProjection, - analyzer process.Analyzer, -) error { - need, err := spillRetainedRecoveryBudgetBytes(projection) - if err != nil { - return err - } - need, err = spillRecoveryReservationBytes(need) - if err != nil { - return err - } - return ctr.ensureSpillRecoveryReservationBytes(need, analyzer) -} - -func (ctr *container) growSpillScratchTransient( - required uint64, - analyzer process.Analyzer, -) (uint64, bool, error) { - if ctr.hashmapBuilder.budget == nil || ctr.spillScratchReservation == nil || - required <= ctr.spillScratchBase { - return 0, false, nil - } - oldSize := ctr.spillScratchReservation.Size() - if err := ctr.spillScratchReservation.Grow(required - ctr.spillScratchBase); err != nil { - return 0, false, err - } - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) - return oldSize, true, nil -} - -// growSpillScratchTransientWithReclaim gives a mandatory replacement overlap -// one deterministic retry after returning optional coalesce ownership. Callers -// invoke it before mutating the allocation being replaced; flushing previously -// buffered records is therefore safe and the current allocation is never -// replayed. -func (ctr *container) growSpillScratchTransientWithReclaim( - proc *process.Process, - files []*os.File, - required uint64, - analyzer process.Analyzer, -) (uint64, bool, error) { - oldSize, grew, err := ctr.growSpillScratchTransient(required, analyzer) - if !errors.Is(err, process.ErrHashBuildBudgetAdmission) { - return oldSize, grew, err - } - reclaimed, reclaimErr := ctr.reclaimOptionalSpillCoalesce(proc, files, analyzer) - if reclaimErr != nil { - return 0, false, reclaimErr - } - if !reclaimed { - return 0, false, err - } - return ctr.growSpillScratchTransient(required, analyzer) -} - -func (ctr *container) restoreSpillScratchTransient(oldSize uint64, grew bool) error { - if !grew { - return nil - } - _, err := ctr.spillScratchReservation.ReconcileDown(oldSize) - return err -} - -func (ctr *container) releaseSpillScratchReservation() { - if ctr.spillScratchReservation != nil { - ctr.spillScratchReservation.Release() - ctr.spillScratchReservation = nil - } - ctr.spillScratchBase = 0 -} - -func (ctr *container) dropSpillScratchBuffers() { - for bucket := range ctr.spillBucketWriteBufs { - ctr.spillBucketWriteBufs[bucket] = bytes.Buffer{} + for bucket := range ctr.spillBucketWriteRows { ctr.spillBucketWriteRows[bucket] = 0 } ctr.spillHashValues = nil @@ -503,99 +69,91 @@ func (ctr *container) dropSpillScratchBuffers() { for i := range ctr.spillBucketOffsets { ctr.spillBucketOffsets[i] = 0 } - ctr.spillSelection = nil ctr.spillKeyVecs = nil - ctr.spillWriteBuf = bytes.Buffer{} -} - -// reclaimOptionalSpillCoalesce gives mandatory recovery reservations priority -// over the write-coalescing cache. It is called only at a quiescent mandatory -// allocation boundary: between spill batches or before the current bucket's -// replacement allocation. Transient scratch has been restored there, so every -// byte above spillScratchBase is owned by optional per-bucket caches. Flush -// their pending records, drop their backing, and return that charge to the -// recovery floor before retrying mandatory admission once. -func (ctr *container) reclaimOptionalSpillCoalesce( - proc *process.Process, - files []*os.File, - analyzer process.Analyzer, -) (bool, error) { - if ctr.spillScratchReservation == nil { - if ctr.spillScratchBase != 0 { - return false, process.ErrHashBuildBudgetInvalid - } - return false, nil - } - current := ctr.spillScratchReservation.Size() - if current < ctr.spillScratchBase { - return false, process.ErrHashBuildBudgetInvalid - } - if current == ctr.spillScratchBase { - return false, nil - } - if err := ctr.flushSpillBuffers(proc, files, analyzer); err != nil { - return false, err - } - for bucket := range ctr.spillBucketWriteBufs { - ctr.spillBucketWriteBufs[bucket] = bytes.Buffer{} - ctr.spillBucketWriteRows[bucket] = 0 - } - _, err := ctr.spillScratchReservation.ReconcileDown(ctr.spillScratchBase) - return err == nil, err -} - -func spillMarshalGrowBytes(bat *batch.Batch) (uint64, error) { - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size + ctr.spillAllocationMP = nil + ctr.spillCoalesceDisabled = false +} + +func growHashBuildSpillSlice[T any]( + values []T, + length int, + mp *mpool.MPool, + account *mpool.AllocationAccount, + site mpool.AllocationSite, + capacityClass mpool.AllocationCapacityClass, +) ([]T, error) { + if length < 0 || account == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + if length <= cap(values) { + return values[:length], nil + } + capacity := cap(values) + if capacity == 0 { + capacity = 1 + } + for capacity < length { + if capacity > math.MaxInt/2 { + capacity = length + break + } + capacity *= 2 + } + next, err := mpool.MakeSliceAccountedWithCapacityClass[T]( + capacity, + mp, + account, + HashBuildAllocationOwner, + site, + capacityClass, + ) + if err != nil { + return nil, err } - columns := uint64(len(bat.Vecs)) - if columns > (math.MaxUint64-24)/128 { - return 0, process.ErrHashBuildBudgetInvalid + copy(next, values) + if cap(values) > 0 { + mpool.FreeSlice(mp, values) } - return spillCheckedAdd(base, columns*128+24) + return next[:length], nil } -func marshalSpillRecord(bat *batch.Batch, buf *bytes.Buffer) (int64, error) { - if bat == nil || bat.RowCount() == 0 { +func marshalSpillRecordAccounted( + bat *batch.Batch, + buf *mpool.AccountedBuffer, +) (int64, error) { + if bat == nil || bat.RowCount() == 0 || buf == nil { return 0, nil } - cnt := int64(bat.RowCount()) buf.Reset() - grow, err := spillMarshalGrowBytes(bat) - if err != nil { - return 0, err - } - if grow > uint64(math.MaxInt) { + batchSize, err := bat.MarshalBinaryWithGroupingSize() + if err != nil || batchSize > math.MaxInt-24 { + if err != nil { + return 0, err + } return 0, process.ErrHashBuildBudgetInvalid } - if uint64(buf.Cap()) < grow { - // Drop a smaller retained buffer before allocating the final capacity; - // otherwise bytes.Buffer's geometric growth recreates the multiplier - // that admission intentionally removed. - *buf = *bytes.NewBuffer(make([]byte, 0, int(grow))) + if err := buf.EnsureCapacity(batchSize + 24); err != nil { + return 0, err + } + if err := buf.WriteInt64(cnt); err != nil { + return 0, err } - buf.Write(types.EncodeInt64(&cnt)) - // Reserve space for batchSize (filled in after marshalling) batchSizePos := buf.Len() - var zero int64 - buf.Write(types.EncodeInt64(&zero)) - - // Write batch data directly to spillWriteBuf. The bounded partition path - // reserves this buffer's conservative upper bound before entering here. - batchStartPos := buf.Len() - if _, err := bat.MarshalBinaryWithBuffer(buf, false); err != nil { + if err := buf.WriteInt64(0); err != nil { + return 0, err + } + batchStart := buf.Len() + if err := bat.MarshalBinaryWithGroupingTo(buf); err != nil { + return 0, err + } + serializedSize := int64(buf.Len() - batchStart) + if err := buf.SetInt64(batchSizePos, serializedSize); err != nil { + return 0, err + } + if err := buf.WriteUint64(uint64(spillMagic)); err != nil { return 0, err } - batchSize := int64(buf.Len() - batchStartPos) - - // Write batchSize at reserved position - batchSizeBytes := types.EncodeInt64(&batchSize) - copy(buf.Bytes()[batchSizePos:batchSizePos+len(batchSizeBytes)], batchSizeBytes) - - magic := uint64(spillMagic) - buf.Write(types.EncodeUint64(&magic)) return cnt, nil } @@ -617,15 +175,12 @@ func (ctr *container) writeSpillPayload( return err } - var err error - if ctr.hashmapBuilder.budget != nil { - if ctr.spillBundle == nil { - return process.ErrHashBuildBudgetInvalid - } - _, _, err = ctr.spillBundle.growDisk(file, ctr.hashmapBuilder.budget, uint64(len(payload))) - if err != nil { - return err - } + if ctr.hashmapBuilder.budget == nil || ctr.spillBundle == nil { + return process.ErrHashBuildBudgetInvalid + } + _, _, err := ctr.spillBundle.growDisk(file, ctr.hashmapBuilder.budget, uint64(len(payload))) + if err != nil { + return err } if err := checkHashBuildCanceled(proc); err != nil { return err @@ -637,12 +192,10 @@ func (ctr *container) writeSpillPayload( if written != len(payload) { return io.ErrShortWrite } - if ctr.hashmapBuilder.budget != nil { - // The exact payload length was admitted. Record logical ownership only - // after the full write; partial writes retain the conservative charge - // until the enclosing bundle closes the file. - ctr.spillBundle.recordDiskWrite(file, rows, uint64(written)) - } + // The exact payload length was admitted. Record logical ownership only + // after the full write; partial writes retain the conservative charge + // until the enclosing bundle closes the file. + ctr.spillBundle.recordDiskWrite(file, rows, uint64(written)) if analyzer != nil { analyzer.Spill(int64(written)) analyzer.SpillRows(rows) @@ -679,12 +232,12 @@ func (ctr *container) ensureSpillFile(proc *process.Process, files []*os.File, b return nil, err } name := fmt.Sprintf("join_%s_%d_build", ctr.spillUUID, bucket) - var fdToken *process.HashBuildSpillFDReservation - if ctr.hashmapBuilder.budget != nil { - fdToken, err = ctr.hashmapBuilder.budget.ReserveSpillFD(1) - if err != nil { - return nil, err - } + if ctr.hashmapBuilder.budget == nil { + return nil, process.ErrHashBuildBudgetInvalid + } + fdToken, err := ctr.hashmapBuilder.budget.ReserveSpillFD(1) + if err != nil { + return nil, err } f, err := spillfs.CreateAndRemoveFile(proc.Ctx, name) if err != nil { @@ -694,12 +247,10 @@ func (ctr *container) ensureSpillFile(proc *process.Process, files []*os.File, b return nil, err } files[bucket] = f - if fdToken != nil { - if ctr.spillBundle == nil { - ctr.spillBundle = &spillFileBundle{} - } - ctr.spillBundle.addFD(f, bucket, fdToken) + if ctr.spillBundle == nil { + ctr.spillBundle = &spillFileBundle{} } + ctr.spillBundle.addFD(f, bucket, fdToken) return f, nil } @@ -709,97 +260,64 @@ func (ctr *container) ensureSpillFile(proc *process.Process, files []*os.File, b // bucket. One selected batch is reused as each bucket is materialized and // marshaled before advancing; serialized records are coalesced until the // bounded buffers or final handoff flush. -func (ctr *container) spillBatchBounded( - proc *process.Process, - bat *batch.Batch, - files []*os.File, - analyzer process.Analyzer, - sourceAlreadyCharged bool, -) error { +func (ctr *container) spillBatchBounded(proc *process.Process, bat *batch.Batch, files []*os.File, executors []colexec.ExpressionExecutor, analyzer process.Analyzer, sourceAlreadyCharged bool) error { if bat == nil || bat.RowCount() == 0 { return nil } if err := checkHashBuildCanceled(proc); err != nil { return err } - expressionLease := ctr.hashmapBuilder.expressionLease - if expressionLease == nil { - return process.ErrHashBuildBudgetInvalid - } - rows := bat.RowCount() - keyCount := expressionLease.Len() - if keyCount == 0 { - return process.ErrHashBuildBudgetInvalid + if ctr.spillBatchAllocation == nil || + ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid } - need, err := spillScratchBudgetBytes(bat, sourceAlreadyCharged) - if err != nil { - return err + if ctr.spillAllocationMP != nil && ctr.spillAllocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid } - // Scratch belongs to the execution generation, not to one batch. Establish - // and grow the lease lazily before allocating spill buffers. Keep it live - // while capacities are retained and release it from Reset/Free/build cleanup - // exactly once. - if ctr.hashmapBuilder.budget != nil { - if ctr.spillScratchReservation == nil { - ctr.spillScratchReservation, err = ctr.hashmapBuilder.budget.Reserve(need) - if err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchReserveRejects", 1) - return err - } - ctr.spillScratchBase = need - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) - } else if need > ctr.spillScratchBase { - grow := need - ctr.spillScratchBase - if err := ctr.spillScratchReservation.Grow(grow); err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchGrowRejects", 1) - return err - } - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchGrowCount", 1) - analyzer.GetOpStats().AddExtraStat("HashBuildSpillScratchGrowBytes", hashBuildStatInt64(grow)) - ctr.spillScratchBase = need - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) + ctr.spillAllocationMP = proc.Mp() + if !sourceAlreadyCharged { + externalBytes := bat.Allocated() + if size := bat.Size(); size > externalBytes { + externalBytes = size } + analyzer.GetOpStats().SetMaxExtraStat( + "HashBuildSpillBorrowedSourceBytes", + int64(externalBytes), + ) + } + + rows := bat.RowCount() + if !keycodec.ValidVectors(bat.Vecs, rows) { + return process.ErrHashBuildBudgetInvalid } + var err error - replacementOverlap, err := spillCapacityReplacementOverlap( + if cap(ctr.spillKeyVecs) < len(executors) { + ctr.spillKeyVecs = make([]*vector.Vector, len(executors)) + } + ctr.spillHashValues, err = growHashBuildSpillSlice( + ctr.spillHashValues, rows, - keyCount, - cap(ctr.spillHashValues), - cap(ctr.spillBucketRowIds), - cap(ctr.spillKeyVecs), + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildSpillAllocationSiteHashValues, + ctr.recoveryCapacityClass, ) if err != nil { return err } - replacementPeak, err := spillCheckedAdd(need, replacementOverlap) - if err != nil { - return err - } - oldScratchSize, grewScratch, err := ctr.growSpillScratchTransientWithReclaim( - proc, files, replacementPeak, analyzer) + ctr.spillBucketRowIds, err = growHashBuildSpillSlice( + ctr.spillBucketRowIds, + rows, + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildSpillAllocationSiteRowIDs, + ctr.recoveryCapacityClass, + ) if err != nil { return err } - - if cap(ctr.spillKeyVecs) < keyCount { - ctr.spillKeyVecs = make([]*vector.Vector, keyCount) - } - if cap(ctr.spillHashValues) < rows { - ctr.spillHashValues = make([]uint64, rows) - } - if cap(ctr.spillBucketRowIds) < rows { - ctr.spillBucketRowIds = make([]int32, rows) - } - if err := ctr.restoreSpillScratchTransient(oldScratchSize, grewScratch); err != nil { - return err - } - keyVecs := ctr.spillKeyVecs[:keyCount] + keyVecs := ctr.spillKeyVecs[:len(executors)] var selected *batch.Batch defer func() { if selected != nil { @@ -809,21 +327,27 @@ func (ctr *container) spillBatchBounded( ctr.spillKeyVecs[i] = nil } }() - err = expressionLease.Eval( - proc, - []*batch.Batch{bat}, - bat.RowCount(), - func(i int, vec *vector.Vector) error { + evalOne := func(i int) error { + vec, evalErr := executors[i].Eval(proc, []*batch.Batch{bat}, nil) + if evalErr == nil { keyVecs[i] = vec - return nil - }, - ) + } + return evalErr + } + for i := range executors { + if err = evalOne(i); err != nil { + break + } + } if err != nil { - // Eval may leave child/result vectors cached. The caller that owns the - // executor set keeps its lease charged until it destroys the complete - // tree; this function must not guess or duplicate that ownership. + // Eval may leave child/result allocations cached. Destroy the tree so a + // pressure retry starts from the exact post-rollback account state. + ctr.freeSpillExprExecs() return err } + if !keycodec.ValidVectors(keyVecs, rows) { + return process.ErrHashBuildBudgetInvalid + } if err := checkHashBuildCanceled(proc); err != nil { return err } @@ -835,10 +359,6 @@ func (ctr *container) spillBatchBounded( if err := checkHashBuildCanceled(proc); err != nil { return err } - // Keep the legacy spillSelection field as an alias for callers/tests that - // inspect it. It intentionally points at the same backing array: no second - // row-id allocation is made. - ctr.spillSelection = ctr.spillBucketRowIds counts := ctr.spillBucketCounts[:] for i := range counts { counts[i] = 0 @@ -862,11 +382,6 @@ func (ctr *container) spillBatchBounded( writePos[bucket] = pos + 1 } - // Coalescing is optional. Once one bucket cannot admit a new cache in this - // batch, later buckets without an already-owned cache write through instead - // of repeating the same fanout-sized budget rejection. The next ingress - // batch probes again, so released sibling headroom is still discoverable. - allowNewCoalesce := true for bucket := 0; bucket < spillNumBuckets; bucket++ { if err := checkHashBuildCanceled(proc); err != nil { return err @@ -877,43 +392,313 @@ func (ctr *container) spillBatchBounded( } if selected == nil { selected = batch.NewOffHeapWithSize(len(bat.Vecs)) - selected.Attrs = bat.Attrs + if err := selected.SetAllocationAccount( + ctr.spillBatchAllocation, + ); err != nil { + return err + } for i, vec := range bat.Vecs { if vec == nil { return process.ErrHashBuildBudgetInvalid } - selected.Vecs[i] = vector.NewOffHeapVecWithType(*vec.GetType()) + selected.Vecs[i], err = + vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), + ctr.spillBatchAllocation, + ) + if err != nil { + return err + } } } - selected.CleanOnlyData() - sels := ctr.spillBucketRowIds[start:end] - n := int(end - start) - var spillErr error - for i, vec := range bat.Vecs { - if vec == nil { - spillErr = process.ErrHashBuildBudgetInvalid - break + cursor := start + for cursor < end { + attemptEnd := end + reclaimedMinimum := false + for { + selected.CleanOnlyData() + sels := ctr.spillBucketRowIds[cursor:attemptEnd] + n := int(attemptEnd - cursor) + var spillErr error + for i, vec := range bat.Vecs { + if vec == nil { + spillErr = process.ErrHashBuildBudgetInvalid + break + } + if spillErr = selected.Vecs[i].PreExtend(n, proc.Mp()); spillErr != nil { + break + } + if spillErr = selected.Vecs[i].UnionInt32(vec, sels, proc.Mp()); spillErr != nil { + break + } + } + if spillErr == nil { + selected.SetRowCount(n) + var file *os.File + file, spillErr = ctr.ensureSpillFile(proc, files, int(bucket)) + if spillErr == nil { + spillErr = ctr.appendSpillRecord( + proc, + file, + int(bucket), + selected, + analyzer, + ) + } + } + selected.CleanOnlyData() + if spillErr == nil { + cursor = attemptEnd + break + } + if !IsRetryableMemoryCapacity(spillErr) { + return spillErr + } + if err := checkHashBuildCanceled(proc); err != nil { + return err + } + if n > 1 { + attemptEnd = cursor + int32((n+1)/2) + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillBatchReductions", + 1, + ) + continue + } + if !reclaimedMinimum { + before := ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used + if err := ctr.reclaimOptionalSpillBuffers( + proc, + files, + analyzer, + ); err != nil { + return err + } + reclaimedMinimum = true + after := ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used + if after >= before { + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillMinimumUnitErrors", + 1, + ) + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-selected-or-codec", + ctr.hashmapBuilder.mapAllocationAccount, + ) + } + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillOptionalReclaims", + 1, + ) + continue + } + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillMinimumUnitErrors", + 1, + ) + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-selected-or-codec", + ctr.hashmapBuilder.mapAllocationAccount, + ) + } + } + } + return nil +} + +// reclaimOptionalSpillBuffers publishes already completed coalesced records, +// then drops codec/coalesce capacity. The current selected record has not been +// published when this is called, so retrying that one record is idempotent. +func (ctr *container) reclaimOptionalSpillBuffers( + proc *process.Process, + files []*os.File, + analyzer process.Analyzer, +) error { + for bucket, buffer := range ctr.spillAccountedBuckets { + if buffer == nil { + continue + } + if buffer.Len() > 0 { + if bucket >= len(files) || files[bucket] == nil { + return process.ErrHashBuildBudgetInvalid + } + if err := ctr.flushPendingSpillBucket( + proc, + files[bucket], + bucket, + analyzer, + ); err != nil { + return err + } + } + buffer.Free() + ctr.spillAccountedBuckets[bucket] = nil + } + ctr.spillCoalesceDisabled = true + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil + } + return nil +} + +func (ctr *container) releaseSpillComputeScratch() { + if ctr.spillBatchAllocation == nil || ctr.spillAllocationMP == nil { + return + } + if cap(ctr.spillHashValues) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillHashValues) + } + if cap(ctr.spillBucketRowIds) > 0 { + mpool.FreeSlice(ctr.spillAllocationMP, ctr.spillBucketRowIds) + } + ctr.spillHashValues = nil + ctr.spillBucketRowIds = nil +} + +// dropMandatorySpillRecoveryScratch releases only allocations that borrow the +// retained-state recovery floor. Optional coalescing buffers use ordinary +// admission and may keep their already-produced records across the transition. +func (ctr *container) dropMandatorySpillRecoveryScratch() { + ctr.freeSpillExprExecs() + ctr.releaseSpillComputeScratch() + if ctr.spillAccountedWrite != nil { + ctr.spillAccountedWrite.Free() + ctr.spillAccountedWrite = nil + } +} + +// spillBatchWithPressure retries only the unpublished prefix of an exact +// spill operation. Hash/expression capacity failures happen before any bucket +// write; selected/codec failures are handled transactionally inside +// spillBatchBounded. Each retry halves the input or reclaims memory, and a +// one-row failure becomes a controlled minimum-unit error. +func (ctr *container) spillBatchWithPressure( + proc *process.Process, + bat *batch.Batch, + files []*os.File, + executors []colexec.ExpressionExecutor, + analyzer process.Analyzer, + sourceAlreadyCharged bool, +) error { + if ctr.spillBatchAllocation == nil || bat == nil || bat.RowCount() == 0 { + return ctr.spillBatchBounded( + proc, + bat, + files, + executors, + analyzer, + sourceAlreadyCharged, + ) + } + rows := bat.RowCount() + chunk := rows + minimumRetried := false + guard := NewPressureRetryGuard(PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: chunk, + OptionalDisabled: ctr.spillCoalesceDisabled, + }, 64) + for start := 0; start < rows; { + if len(ctr.hashmapBuilder.executors) == 0 { + var err error + executors, err = ctr.initSpillExprExecs( + proc, ctr.spillConditions) + if err != nil { + return err } - if spillErr = selected.Vecs[i].PreExtend(n, proc.Mp()); spillErr != nil { - break + } + end := rows + if chunk < rows-start { + end = start + chunk + } + current := bat + if start != 0 || end != rows { + var err error + current, err = bat.WindowWithAllocation( + start, end, proc.Mp(), ctr.spillBatchAllocation, + ) + if err != nil { + return err } - if spillErr = selected.Vecs[i].UnionInt32(vec, sels[:n], proc.Mp()); spillErr != nil { - break + } + err := ctr.spillBatchBounded( + proc, + current, + files, + executors, + analyzer, + sourceAlreadyCharged, + ) + if current != bat { + current.Clean(proc.Mp()) + } + if err == nil { + start = end + minimumRetried = false + nextUnits := chunk + if remaining := rows - start; remaining < nextUnits { + nextUnits = remaining } + guard = NewPressureRetryGuard(PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: nextUnits, + OptionalDisabled: ctr.spillCoalesceDisabled, + }, 64) + continue + } + if !IsRetryableMemoryCapacity(err) { + return err } - if spillErr == nil { - selected.SetRowCount(n) - var file *os.File - file, spillErr = ctr.ensureSpillFile(proc, files, int(bucket)) - if spillErr == nil { - spillErr = ctr.appendSpillRecord( - proc, files, file, int(bucket), selected, need, analyzer, &allowNewCoalesce) + if cancelErr := checkHashBuildCanceled(proc); cancelErr != nil { + return cancelErr + } + ctr.releaseSpillComputeScratch() + attempted := end - start + if attempted <= 1 { + if !minimumRetried { + if reclaimErr := ctr.reclaimOptionalSpillBuffers( + proc, + files, + analyzer, + ); reclaimErr != nil { + return reclaimErr + } + next := PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: attempted, + OptionalDisabled: ctr.spillCoalesceDisabled, + } + if guard.Advance(next) != nil { + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-hash-or-expression", + ctr.hashmapBuilder.mapAllocationAccount, + ) + } + minimumRetried = true + analyzer.GetOpStats().AddExtraStat( + "HashBuildSpillMinimumRetries", + 1, + ) + continue } + return NewMinimumAllocationPressureError( + "hashbuild", + "spill-hash-or-expression", + ctr.hashmapBuilder.mapAllocationAccount, + ) } - selected.CleanOnlyData() - if spillErr != nil { - return spillErr + chunk = (attempted + 1) / 2 + if err := guard.Advance(PressureProgress{ + Used: ctr.hashmapBuilder.mapAllocationAccount.Snapshot().Used, + InputUnits: chunk, + }); err != nil { + return err } + analyzer.GetOpStats().AddExtraStat("HashBuildSpillInputReductions", 1) } return nil } @@ -924,95 +709,95 @@ func (ctr *container) spillBatchBounded( // temporary copy can be retained. func (ctr *container) appendSpillRecord( proc *process.Process, - files []*os.File, file *os.File, bucket int, bat *batch.Batch, - scratchNeed uint64, analyzer process.Analyzer, - allowNewCoalesce *bool, ) error { if bucket < 0 || bucket >= spillNumBuckets { return process.ErrHashBuildBudgetInvalid } - grow, err := spillMarshalGrowBytes(bat) - if err != nil { - return err - } - var oldScratchSize uint64 - var grewScratch bool - if old := uint64(ctr.spillWriteBuf.Cap()); ctr.hashmapBuilder.budget != nil && old > 0 && old < grow { - peak, addErr := spillCheckedAdd(scratchNeed, old) - if addErr != nil { - return addErr - } - oldScratchSize, grewScratch, err = ctr.growSpillScratchTransientWithReclaim( - proc, files, peak, analyzer) + if ctr.spillAllocationMP != proc.Mp() || + ctr.hashmapBuilder.mapAllocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } + if ctr.spillAccountedWrite == nil { + var err error + ctr.spillAccountedWrite, err = mpool.NewAccountedBufferWithCapacityClass( + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteMarshalBuffer, + ctr.recoveryCapacityClass, + ) if err != nil { return err } } - cnt, err := marshalSpillRecord(bat, &ctr.spillWriteBuf) - if restoreErr := ctr.restoreSpillScratchTransient(oldScratchSize, grewScratch); restoreErr != nil { - return restoreErr - } + cnt, err := marshalSpillRecordAccounted(bat, ctr.spillAccountedWrite) if err != nil { return err } - payload := ctr.spillWriteBuf.Bytes() - buf := &ctr.spillBucketWriteBufs[bucket] - if buf.Len() > 0 && buf.Len()+len(payload) > spillWriteCoalesceSize { - if err := ctr.flushPendingSpillBucket(proc, file, bucket, analyzer); err != nil { + payload := ctr.spillAccountedWrite.Bytes() + if ctr.spillCoalesceDisabled { + return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) + } + buffer := ctr.spillAccountedBuckets[bucket] + if buffer != nil && buffer.Len() > 0 && + buffer.Len()+len(payload) > spillWriteCoalesceSize { + if err := ctr.flushPendingSpillBucket( + proc, + file, + bucket, + analyzer, + ); err != nil { return err } } if len(payload) > spillWriteCoalesceSize { return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) } - if buf.Len() == 0 { - if buf.Cap() < spillWriteCoalesceSize && - allowNewCoalesce != nil && !*allowNewCoalesce { - return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) + if buffer == nil { + buffer, err = mpool.NewAccountedBuffer( + proc.Mp(), + ctr.hashmapBuilder.mapAllocationAccount, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteCoalesceBuffer, + ) + if err != nil { + return err } - if !ctr.ensureSpillCoalesceCapacity(buf, analyzer) { - if allowNewCoalesce != nil { - *allowNewCoalesce = false + ctr.spillAccountedBuckets[bucket] = buffer + } + if buffer.Len() == 0 && buffer.Cap() < spillWriteCoalesceSize { + if err := buffer.EnsureCapacity(spillWriteCoalesceSize); err != nil { + if mpool.IsRetryableAllocationCapacity(err) { + return ctr.writeSpillPayload( + proc, + file, + payload, + cnt, + analyzer, + ) } - return ctr.writeSpillPayload(proc, file, payload, cnt, analyzer) - } - if buf.Cap() < spillWriteCoalesceSize { - *buf = *bytes.NewBuffer(make([]byte, 0, spillWriteCoalesceSize)) + return err } } - _, _ = buf.Write(payload) + if _, err := buffer.Write(payload); err != nil { + return err + } ctr.spillBucketWriteRows[bucket] += cnt - if buf.Len() >= spillWriteCoalesceSize { - return ctr.flushPendingSpillBucket(proc, file, bucket, analyzer) + if buffer.Len() >= spillWriteCoalesceSize { + return ctr.flushPendingSpillBucket( + proc, + file, + bucket, + analyzer, + ) } return nil } -func (ctr *container) ensureSpillCoalesceCapacity(buf *bytes.Buffer, analyzer process.Analyzer) bool { - if buf == nil || buf.Cap() >= spillWriteCoalesceSize { - return true - } - if ctr.hashmapBuilder.budget == nil || ctr.spillScratchReservation == nil { - return ctr.hashmapBuilder.budget == nil - } - additional := uint64(spillWriteCoalesceSize - buf.Cap()) - if err := ctr.spillScratchReservation.Grow(additional); err != nil { - analyzer.GetOpStats().AddExtraStat("HashBuildCoalesceGrowRejects", 1) - return false - } - analyzer.GetOpStats().AddExtraStat("HashBuildCoalesceGrowCount", 1) - analyzer.GetOpStats().AddExtraStat("HashBuildCoalesceGrowBytes", hashBuildStatInt64(additional)) - analyzer.GetOpStats().SetMaxExtraStat( - "HashBuildSpillScratchPeakBytes", - hashBuildStatInt64(ctr.spillScratchReservation.Size()), - ) - return true -} - func (ctr *container) flushPendingSpillBucket( proc *process.Process, file *os.File, @@ -1022,16 +807,15 @@ func (ctr *container) flushPendingSpillBucket( if bucket < 0 || bucket >= spillNumBuckets { return process.ErrHashBuildBudgetInvalid } - buf := &ctr.spillBucketWriteBufs[bucket] - if buf.Len() == 0 { + rows := ctr.spillBucketWriteRows[bucket] + buffer := ctr.spillAccountedBuckets[bucket] + if buffer == nil || buffer.Len() == 0 { return nil } - rows := ctr.spillBucketWriteRows[bucket] - payload := buf.Bytes() - err := ctr.writeSpillPayload(proc, file, payload, rows, analyzer) + err := ctr.writeSpillPayload(proc, file, buffer.Bytes(), rows, analyzer) // Clear even on a failed/partial write. A caller's enclosing failure path // owns cleanup, and retrying the same bytes could duplicate records. - buf.Reset() + buffer.Reset() ctr.spillBucketWriteRows[bucket] = 0 return err } @@ -1043,17 +827,25 @@ func (ctr *container) flushPendingSpillBucket( func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, analyzer process.Analyzer) error { var firstErr error for bucket := 0; bucket < spillNumBuckets; bucket++ { - if ctr.spillBucketWriteBufs[bucket].Len() == 0 { + pending := 0 + if ctr.spillAccountedBuckets[bucket] != nil { + pending = ctr.spillAccountedBuckets[bucket].Len() + } + if pending == 0 { continue } if firstErr != nil { - ctr.spillBucketWriteBufs[bucket].Reset() + if ctr.spillAccountedBuckets[bucket] != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 continue } if err := checkHashBuildCanceled(proc); err != nil { firstErr = err - ctr.spillBucketWriteBufs[bucket].Reset() + if ctr.spillAccountedBuckets[bucket] != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 continue } @@ -1063,7 +855,9 @@ func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, } if file == nil { firstErr = process.ErrHashBuildBudgetInvalid - ctr.spillBucketWriteBufs[bucket].Reset() + if ctr.spillAccountedBuckets[bucket] != nil { + ctr.spillAccountedBuckets[bucket].Reset() + } ctr.spillBucketWriteRows[bucket] = 0 continue } @@ -1074,6 +868,40 @@ func (ctr *container) flushSpillBuffers(proc *process.Process, files []*os.File, return firstErr } +// initSpillExprExecs reuses the HashmapBuilder key executors. Spill and normal +// build are mutually exclusive after the transition, so a second executor +// tree would only duplicate retained capacity and its lifecycle. +func (ctr *container) initSpillExprExecs(proc *process.Process, conditions []*plan.Expr) ([]colexec.ExpressionExecutor, error) { + for _, condition := range conditions { + if condition == nil { + return nil, &process.HashBuildBudgetError{Kind: process.HashBuildBudgetErrorInvalid, Message: "nil shuffle spill key"} + } + } + ctr.spillConditions = conditions + if len(ctr.hashmapBuilder.executors) != len(conditions) { + ctr.hashmapBuilder.FreeExecutors() + execs, err := newExpressionExecutorsWithCapacityClass( + proc, + conditions, + ctr.hashmapBuilder.mapAllocationAccount, + ctr.recoveryCapacityClass, + ) + if err != nil { + return nil, err + } + ctr.hashmapBuilder.executors = execs + ctr.hashmapBuilder.keyExprs = conditions + } + ctr.spillExprExecs = ctr.hashmapBuilder.executors + return ctr.spillExprExecs, nil +} + +// freeSpillExprExecs clears the spill alias and releases its builder-owned tree. +func (ctr *container) freeSpillExprExecs() { + ctr.spillExprExecs = nil + ctr.hashmapBuilder.FreeExecutors() +} + func (ctr *container) memUsed() int64 { sz := ctr.hashmapBuilder.GetSize() + ctr.hashmapBuilder.Batches.MemSize batches := ctr.hashmapBuilder.Batches.Buf diff --git a/pkg/sql/colexec/hashbuild/spill_test.go b/pkg/sql/colexec/hashbuild/spill_test.go index 7e07f2c4b40ac..407587b8e77d4 100644 --- a/pkg/sql/colexec/hashbuild/spill_test.go +++ b/pkg/sql/colexec/hashbuild/spill_test.go @@ -16,10 +16,9 @@ package hashbuild import ( "bufio" - "bytes" "context" + "errors" "io" - "math" "os" "strings" "testing" @@ -35,2143 +34,248 @@ import ( "github.com/stretchr/testify/require" ) -func initSpillExprExecsForTest( - ctr *container, - proc *process.Process, - conditions []*plan.Expr, -) error { - for _, condition := range conditions { - if condition == nil { - return process.ErrHashBuildBudgetInvalid - } - } - ctr.hashmapBuilder.FreeExecutors() - if err := ctr.hashmapBuilder.Prepare( - conditions, -1, -1, nil, proc, - ); err != nil { - return err - } - return nil -} - -func TestComputeXXHashBuild(t *testing.T) { - mp := mpool.MustNewZero() - - t.Run("empty", func(t *testing.T) { - computeXXHash(nil, nil) - }) - - t.Run("single_column", func(t *testing.T) { - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, mp) - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec}, hashValues) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, hashValues[0], hashValues[1]) - }) - - t.Run("multiple_columns", func(t *testing.T) { - vec1 := testutil.MakeInt32Vector([]int32{1, 2}, nil, mp) - vec2 := testutil.MakeVarcharVector([]string{"a", "b"}, nil, mp) - hashValues := make([]uint64, 2) - computeXXHash([]*vector.Vector{vec1, vec2}, hashValues) - require.NotEqual(t, hashValues[0], hashValues[1]) - }) - - t.Run("const_vector", func(t *testing.T) { - vec := testutil.MakeInt32Vector([]int32{5}, nil, mp) - vec.SetClass(vector.CONSTANT) - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec}, hashValues) - require.Equal(t, hashValues[0], hashValues[1]) - }) -} - -func TestMarshalSpillRecordBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - t.Run("empty batch", func(t *testing.T) { - var buf bytes.Buffer - cnt, err := marshalSpillRecord(nil, &buf) - require.NoError(t, err) - require.Zero(t, cnt) - require.Zero(t, buf.Len()) - }) - - t.Run("framed batch", func(t *testing.T) { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.SetRowCount(3) - defer bat.Clean(proc.Mp()) - - var buf bytes.Buffer - cnt, err := marshalSpillRecord(bat, &buf) - require.NoError(t, err) - require.Equal(t, int64(3), cnt) - require.GreaterOrEqual(t, buf.Len(), 24) - require.Equal(t, cnt, types.DecodeInt64(buf.Bytes()[:8])) - payloadBytes := types.DecodeInt64(buf.Bytes()[8:16]) - require.Positive(t, payloadBytes) - require.Equal(t, int64(buf.Len()), 16+payloadBytes+8) - require.Equal(t, uint64(spillMagic), - types.DecodeUint64(buf.Bytes()[buf.Len()-8:])) - }) -} - -func TestShouldSpillBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - t.Run("not_shuffle", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: false, - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1) - bat := batch.NewWithSize(0) - bat.SetRowCount(1) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - require.False(t, hb.shouldSpillBatches()) - }) - - t.Run("no_hashmap", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: true, - } - hb.ctr.setSpillThreshold(1) - bat := batch.NewWithSize(0) - bat.SetRowCount(1) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - require.False(t, hb.shouldSpillBatches()) - }) - - t.Run("below_threshold", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 1024 * 1024, // 1MB - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1024 * 1024) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{ - {Vecs: []*vector.Vector{testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp())}}, - } - require.False(t, hb.shouldSpillBatches()) - }) - - t.Run("above_threshold", func(t *testing.T) { - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 1, // 1 byte - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1) - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5}, nil, proc.Mp()) - bat.SetRowCount(5) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - hb.ctr.hashmapBuilder.InputBatchRowCount = bat.RowCount() - require.True(t, hb.shouldSpillBatches()) - }) -} - -func TestShouldSpillBeforeRetain(t *testing.T) { - t.Run("byte threshold predicts crossing batch", func(t *testing.T) { - hb := &HashBuild{IsShuffle: true, NeedHashMap: true} - hb.ctr.setSpillThreshold(100_001) - hb.ctr.hashmapBuilder.Batches.MemSize = 60_000 - hb.ctr.hashmapBuilder.InputBatchRowCount = 2 - - require.False(t, hb.shouldSpillBeforeRetain(40_001), - "the byte convention spills only after the threshold") - require.True(t, hb.shouldSpillBeforeRetain(40_002), - "the crossing batch must be routed directly before it consumes headroom") - }) - - t.Run("row threshold already includes ingress batch", func(t *testing.T) { - hb := &HashBuild{IsShuffle: true, NeedHashMap: true} - hb.ctr.setSpillThreshold(10) - hb.ctr.hashmapBuilder.InputBatchRowCount = 9 - require.False(t, hb.shouldSpillBeforeRetain(1)) - hb.ctr.hashmapBuilder.InputBatchRowCount = 10 - require.True(t, hb.shouldSpillBeforeRetain(1)) - }) - - t.Run("ineligible topology stays resident", func(t *testing.T) { - hb := &HashBuild{IsShuffle: false, NeedHashMap: true} - hb.ctr.setSpillThreshold(1) - hb.ctr.hashmapBuilder.InputBatchRowCount = 1 - require.False(t, hb.shouldSpillBeforeRetain(math.MaxInt64)) - }) - - t.Run("size overflow fails toward spill", func(t *testing.T) { - hb := &HashBuild{IsShuffle: true, NeedHashMap: true} - hb.ctr.setSpillThreshold(100_001) - hb.ctr.hashmapBuilder.Batches.MemSize = math.MaxInt64 - 1 - require.True(t, hb.shouldSpillBeforeRetain(2)) - }) -} - -func TestMemUsedIncludesPartialTailAfterFullBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - full := batch.NewWithSize(0) - full.SetRowCount(colexec.DefaultBatchSize) - partial := batch.NewWithSize(1) - partial.Vecs[0] = testutil.MakeVarcharVector([]string{"partial-tail"}, nil, proc.Mp()) - partial.SetRowCount(1) - defer partial.Clean(proc.Mp()) - - ctr := container{} - ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{full, partial} - ctr.hashmapBuilder.Batches.MemSize = 60_000 - require.Equal(t, int64(60_000+partial.Size()), ctr.memUsed()) -} - -func TestHashDistributionBuild(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, - 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30}, nil, mp) - - hashValues := make([]uint64, 30) - computeXXHash([]*vector.Vector{vec}, hashValues) - - bucketCounts := make([]int, spillNumBuckets) - for _, hash := range hashValues { - bucketId := hash & (spillNumBuckets - 1) - bucketCounts[bucketId]++ - } - - // At least some buckets should have values - nonEmptyBuckets := 0 - for _, count := range bucketCounts { - if count > 0 { - nonEmptyBuckets++ - } - } - require.Greater(t, nonEmptyBuckets, 1) -} - -func TestMultipleDataTypesBuild(t *testing.T) { - mp := mpool.MustNewZero() - - tests := []struct { - name string - vec *vector.Vector - }{ - {"int8", testutil.MakeInt8Vector([]int8{1, 2, 3}, nil, mp)}, - {"int16", testutil.MakeInt16Vector([]int16{100, 200, 300}, nil, mp)}, - {"int64", testutil.MakeInt64Vector([]int64{1000, 2000, 3000}, nil, mp)}, - {"uint32", testutil.MakeUint32Vector([]uint32{10, 20, 30}, nil, mp)}, - {"float32", testutil.MakeFloat32Vector([]float32{1.1, 2.2, 3.3}, nil, mp)}, - {"float64", testutil.MakeFloat64Vector([]float64{10.1, 20.2, 30.3}, nil, mp)}, - {"varchar", testutil.MakeVarcharVector([]string{"abc", "def", "ghi"}, nil, mp)}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{tt.vec}, hashValues) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, hashValues[0], hashValues[1]) - }) - } -} - -func TestNullValuesBuild(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, []uint64{1}, mp) - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec}, hashValues) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, uint64(0), hashValues[2]) -} - -func TestFileWriteErrorBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - analyzer := process.NewAnalyzer(0, false, false, "test") - spillfs, _ := proc.GetSpillFileService() - file, _ := spillfs.CreateFile(context.Background(), "test_error_build") - file.Close() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - defer bat.Clean(proc.Mp()) - - ctr := &container{spillUUID: t.Name()} - var buf bytes.Buffer - cnt, err := marshalSpillRecord(bat, &buf) - require.NoError(t, err) - err = ctr.writeSpillPayload(proc, file, buf.Bytes(), cnt, analyzer) - require.Error(t, err) - - spillfs.RemoveFile(context.Background(), "test_error_build") -} - -func TestWriteSpillPayloadCancellationStopsBeforePhysicalWrite(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - file, err := spillfs.CreateFile(context.Background(), t.Name()) - require.NoError(t, err) - defer func() { - require.NoError(t, file.Close()) - require.NoError(t, spillfs.RemoveFile(context.Background(), t.Name())) - }() - - proc.Cancel(context.Canceled) - analyzer := process.NewAnalyzer(0, false, false, "test") - err = (&container{}).writeSpillPayload(proc, file, []byte("stale spill payload"), 1, analyzer) - require.ErrorIs(t, err, context.Canceled) - - info, err := file.Stat() - require.NoError(t, err) - require.Zero(t, info.Size(), "canceled spill must not start physical I/O") - require.Zero(t, analyzer.GetOpStats().SpillSize) - require.Zero(t, analyzer.GetOpStats().SpillRows) -} - -func TestSpillBatchPartitioning(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - // Create batch with known values - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5, 6, 7, 8}, nil, proc.Mp()) - bat.SetRowCount(8) - defer bat.Clean(proc.Mp()) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - err = ctr.spillBatchBounded(proc, bat, files, analyzer, false) - require.NoError(t, err) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - - hashes := make([]uint64, bat.RowCount()) - computeXXHash(bat.Vecs, hashes) - expectedBuckets := make(map[int]struct{}) - for _, hash := range hashes { - expectedBuckets[int(hash&(spillNumBuckets-1))] = struct{}{} - } - require.Greater(t, len(expectedBuckets), 1) - for bucket, file := range files { - _, expected := expectedBuckets[bucket] - require.Equalf(t, expected, file != nil, "bucket %d", bucket) - } - require.Equal(t, int64(bat.RowCount()), analyzer.GetOpStats().SpillRows) -} - -func TestSpillBatchEmptyInput(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{}, nil, proc.Mp()) - bat.SetRowCount(0) - defer bat.Clean(proc.Mp()) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - err = ctr.spillBatchBounded(proc, bat, files, analyzer, false) - require.NoError(t, err) - for _, file := range files { - require.Nil(t, file) - } - require.Zero(t, analyzer.GetOpStats().SpillRows) -} - -func TestSpillBatchLargeInputPreservesRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - size := 2 * colexec.DefaultBatchSize - values := make([]int32, size) - for i := range values { - values[i] = int32(i) - } - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(size) - defer bat.Clean(proc.Mp()) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - err = ctr.spillBatchBounded(proc, bat, files, analyzer, false) - require.NoError(t, err) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - require.Equal(t, int64(bat.RowCount()), analyzer.GetOpStats().SpillRows) -} - -func TestSpillBatchNullKeyPreservesRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, []uint64{1}, proc.Mp()) // null at index 1 - bat.SetRowCount(4) - defer bat.Clean(proc.Mp()) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - err = ctr.spillBatchBounded(proc, bat, files, analyzer, false) - require.NoError(t, err) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - require.Equal(t, int64(bat.RowCount()), analyzer.GetOpStats().SpillRows) -} - -func TestSpillBatchMultiColumnPreservesRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - bat := batch.NewWithSize(2) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector([]string{"a", "b", "c"}, nil, proc.Mp()) - bat.SetRowCount(3) - defer bat.Clean(proc.Mp()) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - { - Typ: plan.Type{Id: int32(types.T_varchar)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 1}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - err = ctr.spillBatchBounded(proc, bat, files, analyzer, false) - require.NoError(t, err) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - require.Equal(t, int64(bat.RowCount()), analyzer.GetOpStats().SpillRows) -} - -func TestShouldSpillBatchesRowThreshold(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 10, // Small row threshold - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(10) - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.SetRowCount(3) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - hb.ctr.hashmapBuilder.InputBatchRowCount = bat.RowCount() - - require.False(t, hb.shouldSpillBatches()) - - // Add more batches to exceed threshold - for i := 0; i < 10; i++ { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{int32(i)}, nil, proc.Mp()) - bat.SetRowCount(1) - hb.ctr.hashmapBuilder.Batches.Buf = append(hb.ctr.hashmapBuilder.Batches.Buf, bat) - hb.ctr.hashmapBuilder.InputBatchRowCount += bat.RowCount() - } - - require.True(t, hb.shouldSpillBatches()) -} - -func TestShouldSpillBatchesMemThreshold(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - hb := &HashBuild{ - IsShuffle: true, - SpillThreshold: 1024 * 1024, // 1MB - NeedHashMap: true, - } - hb.ctr.setSpillThreshold(1024 * 1024) - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - hb.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} - - require.False(t, hb.shouldSpillBatches()) -} - -func TestHashWithConstVector(t *testing.T) { - mp := mpool.MustNewZero() - - vec := testutil.MakeInt32Vector([]int32{42}, nil, mp) - vec.SetClass(vector.CONSTANT) - - hashValues := make([]uint64, 10) - computeXXHash([]*vector.Vector{vec}, hashValues) - - // All values should be the same for const vector - for i := 1; i < len(hashValues); i++ { - require.Equal(t, hashValues[0], hashValues[i]) - } -} - -func TestHashMultiColumnCombinations(t *testing.T) { - mp := mpool.MustNewZero() - - vec1 := testutil.MakeInt32Vector([]int32{1, 1, 2}, nil, mp) - vec2 := testutil.MakeVarcharVector([]string{"a", "b", "a"}, nil, mp) - - hashValues := make([]uint64, 3) - computeXXHash([]*vector.Vector{vec1, vec2}, hashValues) - - // Different combinations should produce different hashes - require.NotEqual(t, hashValues[0], hashValues[1]) - require.NotEqual(t, hashValues[0], hashValues[2]) -} - -func TestSpillBatchSingleRowUsesOneBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - // Single value should go to one bucket - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - defer bat.Clean(proc.Mp()) - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - err = ctr.spillBatchBounded(proc, bat, files, analyzer, false) - require.NoError(t, err) - - nonEmptyBuckets := 0 - for _, file := range files { - if file != nil { - nonEmptyBuckets++ - } - } - require.Equal(t, 1, nonEmptyBuckets) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - require.Equal(t, int64(1), analyzer.GetOpStats().SpillRows) -} - -func TestSpillScratchSlicesReuseAcrossEqualBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - file.Close() - } - } - }() - - conditions := []*plan.Expr{ - { - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{ - Col: &plan.ColRef{ColPos: 0}, - }, - }, - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - ctr := &container{spillUUID: t.Name()} - - err := initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - - // First batch - bat1 := batch.NewWithSize(1) - bat1.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat1.SetRowCount(2) - defer bat1.Clean(proc.Mp()) - - err = ctr.spillBatchBounded(proc, bat1, files, analyzer, false) - require.NoError(t, err) - require.Len(t, ctr.spillHashValues, bat1.RowCount()) - require.Len(t, ctr.spillBucketRowIds, bat1.RowCount()) - firstHashStorage := &ctr.spillHashValues[0] - firstRowIDStorage := &ctr.spillBucketRowIds[0] - - // Second batch - buffers should be reused - bat2 := batch.NewWithSize(1) - bat2.Vecs[0] = testutil.MakeInt32Vector([]int32{3, 4}, nil, proc.Mp()) - bat2.SetRowCount(2) - defer bat2.Clean(proc.Mp()) - - err = ctr.spillBatchBounded(proc, bat2, files, analyzer, false) - require.NoError(t, err) - require.Same(t, firstHashStorage, &ctr.spillHashValues[0]) - require.Same(t, firstRowIDStorage, &ctr.spillBucketRowIds[0]) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - require.Equal(t, int64(4), analyzer.GetOpStats().SpillRows) -} - -func TestSpillExpressionLeaseRetainsLargeBatchHighWater(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(256<<20, 256<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - _ = file.Close() - } - } - }() - expr := makeExpressionLeaseTestExpr(t, proc) - ctr := &container{spillUUID: t.Name()} - ctr.hashmapBuilder.setBudget(generation) - err = initSpillExprExecsForTest(ctr, proc, []*plan.Expr{expr}) - require.NoError(t, err) - require.NotNil(t, ctr.hashmapBuilder.expressionLease) - defer ctr.hashmapBuilder.FreeExecutors() - defer ctr.dropSpillScratchBuffers() - defer ctr.releaseSpillScratchReservation() - - analyzer := process.NewAnalyzer(0, false, false, "test") - large := makeExpressionLeaseTestBatch(proc, colexec.DefaultBatchSize) - defer large.Clean(proc.Mp()) - require.NoError(t, ctr.spillBatchBounded(proc, large, files, analyzer, false)) - largeReserved := ctr.hashmapBuilder.expressionLease.Reserved() - require.Positive(t, largeReserved) - - small := makeExpressionLeaseTestBatch(proc, 1) - defer small.Clean(proc.Mp()) - require.NoError(t, ctr.spillBatchBounded(proc, small, files, analyzer, false)) - require.Equal(t, largeReserved, ctr.hashmapBuilder.expressionLease.Reserved(), - "a small spill batch must not release retained executor headroom") - retained, ok := ctr.hashmapBuilder.expressionLease.Retained() - require.True(t, ok) - require.LessOrEqual(t, retained, ctr.hashmapBuilder.expressionLease.Reserved()) -} - -func TestSpillWriteCoalescesAcrossBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(8 << 20) - require.NoError(t, err) - defer generation.Close() - files := make([]*os.File, spillNumBuckets) - conditions := []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} - ctr := &container{spillUUID: t.Name()} - ctr.hashmapBuilder.setBudget(generation) - cleanup := func() { - for i, file := range files { - if file != nil { - _ = file.Close() - files[i] = nil - } - } - if ctr.spillBundle != nil { - ctr.spillBundle.release() - ctr.spillBundle = nil - } - ctr.hashmapBuilder.FreeExecutors() - ctr.dropSpillScratchBuffers() - ctr.releaseSpillScratchReservation() - } - defer cleanup() - err = initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - analyzer := process.NewAnalyzer(0, false, false, "test") - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 1, 1}, nil, proc.Mp()) - bat.SetRowCount(3) - defer bat.Clean(proc.Mp()) - for i := 0; i < 2; i++ { - require.NoError(t, ctr.spillBatchBounded(proc, bat, files, analyzer, false)) - } - var pending int - for i := range ctr.spillBucketWriteBufs { - pending += ctr.spillBucketWriteBufs[i].Len() - } - require.Positive(t, pending) - var file *os.File - for _, f := range files { - if f != nil { - file = f - break - } - } - require.NotNil(t, file) - stat, err := file.Stat() - require.NoError(t, err) - require.Zero(t, stat.Size(), "records stay pending until the handoff flush") - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) - stat, err = file.Stat() - require.NoError(t, err) - require.Positive(t, stat.Size()) - for i := range ctr.spillBucketWriteBufs { - require.Zero(t, ctr.spillBucketWriteBufs[i].Len()) - } - _, err = file.Seek(0, io.SeekStart) - require.NoError(t, err) - reader := bufio.NewReader(file) - var totalRows int64 - for { - var header [16]byte - _, err = io.ReadFull(reader, header[:]) - if err == io.EOF { - break - } - require.NoError(t, err) - cnt := types.DecodeInt64(header[:8]) - payload := types.DecodeInt64(header[8:]) - require.GreaterOrEqual(t, cnt, int64(0)) - require.GreaterOrEqual(t, payload, int64(0)) - _, err = io.CopyN(io.Discard, reader, payload) - require.NoError(t, err) - var magic [8]byte - _, err = io.ReadFull(reader, magic[:]) - require.NoError(t, err) - require.Equal(t, uint64(spillMagic), types.DecodeUint64(magic[:])) - totalRows += cnt - } - require.Equal(t, int64(6), totalRows) - scratchPeak := analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchPeakBytes"] - require.GreaterOrEqual(t, scratchPeak, hashBuildStatInt64(ctr.spillScratchReservation.Size())) - require.Greater(t, scratchPeak, hashBuildStatInt64(ctr.spillScratchBase), - "scratch peak must include retained coalesce buffers above the base lease") - cleanup() - require.Zero(t, generation.Used()) -} - -func TestReclaimOptionalSpillCoalesceReleasesRecoveryHeadroom(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - files := make([]*os.File, spillNumBuckets) - ctr := &container{spillUUID: t.Name()} - ctr.hashmapBuilder.setBudget(generation) - cleanup := func() { - for i, file := range files { - if file != nil { - _ = file.Close() - files[i] = nil - } - } - if ctr.spillBundle != nil { - ctr.spillBundle.release() - ctr.spillBundle = nil - } - ctr.hashmapBuilder.FreeExecutors() - ctr.dropSpillScratchBuffers() - ctr.releaseSpillScratchReservation() - } - defer cleanup() - conditions := []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} - err = initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - analyzer := process.NewAnalyzer(0, false, false, "recovery priority") - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 1, 1}, nil, proc.Mp()) - bat.SetRowCount(3) - defer bat.Clean(proc.Mp()) - require.NoError(t, ctr.spillBatchBounded(proc, bat, files, analyzer, false)) - require.NotNil(t, ctr.spillScratchReservation) - base := ctr.spillScratchBase - beforeSize := ctr.spillScratchReservation.Size() - beforeUsed := generation.Used() - require.Greater(t, beforeSize, base) - - reclaimed, err := ctr.reclaimOptionalSpillCoalesce(proc, files, analyzer) - require.NoError(t, err) - require.True(t, reclaimed) - require.Equal(t, base, ctr.spillScratchReservation.Size()) - require.Equal(t, beforeUsed-(beforeSize-base), generation.Used()) - for bucket := range ctr.spillBucketWriteBufs { - require.Zero(t, ctr.spillBucketWriteBufs[bucket].Len()) - require.Zero(t, ctr.spillBucketWriteBufs[bucket].Cap()) - require.Zero(t, ctr.spillBucketWriteRows[bucket]) - } - // The transition is idempotent and must not shrink the mandatory floor. - reclaimed, err = ctr.reclaimOptionalSpillCoalesce(proc, files, analyzer) - require.NoError(t, err) - require.False(t, reclaimed) - require.Equal(t, base, ctr.spillScratchReservation.Size()) - - cleanup() - require.Zero(t, generation.Used()) -} - -func TestSpillScratchBudgetDoesNotDoubleChargeRetainedSource(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = int32(i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - defer bat.Clean(proc.Mp()) - - fullNeed, err := spillScratchBudgetBytes(bat, false) - require.NoError(t, err) - retainedNeed, err := spillScratchBudgetBytes(bat, true) - require.NoError(t, err) - source := uint64(bat.Allocated()) - require.Equal(t, source, fullNeed-retainedNeed) - require.Positive(t, retainedNeed) - - // The retained source has its own batch reservation. Only the incremental - // scratch must be admitted when the real spill path starts. - proofCap := source + fullNeed - 1 - proofBudget := process.MustNewHashBuildBudget(proofCap, proofCap) - proofGeneration, err := proofBudget.OpenGeneration(1) - require.NoError(t, err) - proofSource, err := proofGeneration.Reserve(source) - require.NoError(t, err) - _, err = proofGeneration.Reserve(fullNeed) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - proofScratch, err := proofGeneration.Reserve(retainedNeed) - require.NoError(t, err) - proofScratch.Release() - proofSource.Release() - require.Zero(t, proofGeneration.Used()) -} - -func TestSpillScratchLazyGrowSucceeds(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - values := make([]int32, colexec.DefaultBatchSize) - for i := range values { - values[i] = int32(i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - defer bat.Clean(proc.Mp()) - - need, err := spillScratchBudgetBytes(bat, true) - require.NoError(t, err) - require.Greater(t, need, uint64(1)) - source := uint64(bat.Allocated()) - const slack = uint64(2 << 20) - capBytes := source + need + slack - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - retainedToken, err := generation.Reserve(source) - require.NoError(t, err) - defer retainedToken.Release() - scratchToken, err := generation.Reserve(need - 1) - require.NoError(t, err) - - files := make([]*os.File, spillNumBuckets) - defer func() { - for _, file := range files { - if file != nil { - _ = file.Close() - } - } - }() - conditions := []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} - ctr := &container{ - spillUUID: t.Name(), - spillScratchReservation: scratchToken, - spillScratchBase: need - 1, - } - ctr.hashmapBuilder.setBudget(generation) - err = initSpillExprExecsForTest(ctr, proc, conditions) - require.NoError(t, err) - defer ctr.hashmapBuilder.FreeExecutors() - defer ctr.dropSpillScratchBuffers() - defer ctr.releaseSpillScratchReservation() - - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, ctr.spillBatchBounded(proc, bat, files, analyzer, true)) - require.Equal(t, need, ctr.spillScratchBase) - require.GreaterOrEqual(t, scratchToken.Size(), need) - require.Equal(t, int64(1), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchGrowCount"]) - require.Equal(t, int64(1), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchGrowBytes"]) - require.NoError(t, ctr.flushSpillBuffers(proc, files, analyzer)) -} - -func TestSpillScratchLazyGrowRejectPreservesRetainedSource(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, proc.Mp()) - bat.SetRowCount(4) - defer bat.Clean(proc.Mp()) - originalAllocated := bat.Allocated() - - need, err := spillScratchBudgetBytes(bat, true) - require.NoError(t, err) - require.Greater(t, need, uint64(1)) - source := uint64(bat.Allocated()) - capBytes := source + need - 1 - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - sourceToken, err := generation.Reserve(source) - require.NoError(t, err) - scratchToken, err := generation.Reserve(need - 1) - require.NoError(t, err) - - files := make([]*os.File, spillNumBuckets) - ctr := &container{ - spillUUID: t.Name(), - spillScratchReservation: scratchToken, - spillScratchBase: need - 1, - } - ctr.hashmapBuilder.setBudget(generation) - err = initSpillExprExecsForTest(ctr, proc, []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }}) - require.NoError(t, err) - analyzer := process.NewAnalyzer(0, false, false, "lazy spill reject") - err = ctr.spillBatchBounded(proc, bat, files, analyzer, true) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, int64(1), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchGrowRejects"]) - require.Equal(t, need-1, scratchToken.Size()) - require.Equal(t, capBytes, generation.Used()) - require.Equal(t, 4, bat.RowCount()) - require.Equal(t, originalAllocated, bat.Allocated()) - require.Nil(t, ctr.spillHashValues) - require.Nil(t, ctr.spillBucketRowIds) - for _, file := range files { - require.Nil(t, file) - } - - ctr.hashmapBuilder.FreeExecutors() - ctr.releaseSpillScratchReservation() - sourceToken.Release() - require.Zero(t, generation.Used()) -} - -func TestFlushSpillBuffersCancellationDiscardsPendingWrites(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - - ctr := &container{} - for _, bucket := range []int{0, spillNumBuckets - 1} { - _, err := ctr.spillBucketWriteBufs[bucket].Write([]byte("pending")) - require.NoError(t, err) - ctr.spillBucketWriteRows[bucket] = 1 - } - proc.Cancel(context.Canceled) - - err := ctr.flushSpillBuffers(proc, nil, process.NewAnalyzer(0, false, false, "test")) - require.ErrorIs(t, err, context.Canceled) - for bucket := 0; bucket < spillNumBuckets; bucket++ { - require.Zero(t, ctr.spillBucketWriteBufs[bucket].Len()) - require.Zero(t, ctr.spillBucketWriteRows[bucket]) - } +type spillTestHarness struct { + op *HashBuild + proc *process.Process + generation *process.HashBuildBudgetGeneration + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + files []*os.File } -func TestSpillMaterializedBytesDoesNotScaleShuffledConstVector(t *testing.T) { +func newSpillTestHarness(t *testing.T, limit uint64) *spillTestHarness { + t.Helper() proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const sourceRows = 32 * 1024 - bat := batch.NewWithSize(2) - values := make([]int32, sourceRows) - for i := range values { - values[i] = int32(i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - var err error - bat.Vecs[1], err = vector.NewConstBytes( - types.T_varchar.ToType(), - []byte("test create big fulltext index"), - sourceRows, - proc.Mp(), - ) - require.NoError(t, err) - bat.SetRowCount(sourceRows) - defer bat.Clean(proc.Mp()) - - // Batch.Shuffle intentionally leaves a const vector untouched while it - // changes the batch cardinality. This is the shape produced by the failed - // generate_series + const-varchar BVT query. - require.NoError(t, bat.Shuffle([]int64{0}, proc.Mp())) - require.Equal(t, 1, bat.RowCount()) - require.Equal(t, sourceRows, bat.Vecs[1].Length()) - - legacySource := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > legacySource { - legacySource = size - } - legacyScaled := legacySource * uint64(colexec.DefaultBatchSize) - legacyMetadata, ok := retainedMetadataAllowance(bat) - require.True(t, ok) - legacyScaled += legacyMetadata * uint64(colexec.DefaultBatchSize) - legacyNeed, err := spillPeakBudgetFor(uint64(colexec.DefaultBatchSize), 0, legacyScaled, uint64(len(bat.Vecs))) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) require.NoError(t, err) - require.Greater(t, legacyNeed, uint64(10<<30), - "the old logical-size extrapolation must reproduce the false 10 GiB rejection") - - materialized, err := spillMaterializedBytes(bat) + registry, err := mpool.NewAllocationAccountRegistry(1, 256) require.NoError(t, err) - wantMaterialized := uint64(types.T_int32.ToType().TypeSize()) + - uint64(types.T_varchar.ToType().TypeSize()) + - uint64(len("test create big fulltext index")) - require.Equal(t, wantMaterialized, materialized, - "lazy admission must use the batch's live rows and const payload") - - directNeed, err := spillBudgetBytes(bat) + account, err := registry.OpenWithController(limit, generation) require.NoError(t, err) - require.Less(t, directNeed, uint64(16<<20), - "lazy scratch admission must not scale stale logical length") -} - -func TestSpillMaterializedBytesDoesNotScaleRetainedVectorCapacity(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - const sourceRows = 32 * 1024 - bat := batch.NewWithSize(2) - values := make([]int32, sourceRows) - strings := make([]string, sourceRows) - for i := range values { - values[i] = int32(i) - strings[i] = "test create big fulltext index" + op := &HashBuild{NeedHashMap: true} + require.NoError(t, op.SetAllocationAccount(account)) + op.ctr.hashmapBuilder.setBudget(generation) + op.ctr.spillUUID = t.Name() + return &spillTestHarness{ + op: op, + proc: proc, + generation: generation, + registry: registry, + account: account, + files: make([]*os.File, spillNumBuckets), } - bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector(strings, nil, proc.Mp()) - bat.SetRowCount(sourceRows) - defer bat.Clean(proc.Mp()) - - // Reused shuffle and table-function batches keep their allocation while - // publishing a tiny final batch. Only the first row is live, but Allocated - // still describes the original 32K-row capacity. - bat.Vecs[0].SetLength(1) - bat.Vecs[1].SetLength(1) - bat.SetRowCount(1) - require.Greater(t, bat.Allocated(), 1<<20) - - materialized, err := spillMaterializedBytes(bat) - require.NoError(t, err) - wantMaterialized := uint64(types.T_int32.ToType().TypeSize()) + - uint64(types.T_varchar.ToType().TypeSize()) + - uint64(len(strings[0])) - require.Equal(t, wantMaterialized, materialized) - - directNeed, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Less(t, directNeed, uint64(16<<20), - "source capacity is charged once, never extrapolated per live row") -} - -func TestSpillProjectedSourceSkipsStaleNullVarlenaPayload(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector([]string{"x"}, []uint64{0}, proc.Mp()) - bat.SetRowCount(1) - defer bat.Clean(proc.Mp()) - - // A null append does not overwrite a reused varlen slot. Plant the stale - // non-inline header that such a slot can retain; UnionInt32 skips the null - // value, so this dead payload must not be projected into the spill batch. - values, _ := vector.MustVarlenaRawData(bat.Vecs[0]) - const staleLen = uint32(1 << 20) - values[0].SetOffsetLen(0, staleLen) - - source, err := spillMaterializedBytes(bat) - require.NoError(t, err) - require.Equal(t, uint64(bat.Vecs[0].GetType().TypeSize()), source) - - need, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Less(t, need, uint64(16<<20)) -} - -func TestSpillMaterializedBytesBoundaryInputs(t *testing.T) { - source, err := spillMaterializedBytes(nil) - require.NoError(t, err) - require.Zero(t, source) - - invalid := batch.NewWithSize(1) - invalid.SetRowCount(1) - _, err = spillMaterializedBytes(invalid) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillBudgetBytes(invalid) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - short := batch.NewWithSize(1) - short.Vecs[0] = testutil.MakeVarcharVector([]string{"x"}, nil, proc.Mp()) - short.SetRowCount(2) - defer short.Clean(proc.Mp()) - _, err = spillMaterializedBytes(short) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - constNull := batch.NewWithSize(1) - constNull.Vecs[0] = vector.NewConstNull(types.T_varchar.ToType(), 1, proc.Mp()) - constNull.SetRowCount(1) - defer constNull.Clean(proc.Mp()) - - source, err = spillMaterializedBytes(constNull) - require.NoError(t, err) - require.Equal(t, uint64(types.T_varchar.ToType().TypeSize()), source) -} - -func TestSpillBudgetArithmeticFailsClosed(t *testing.T) { - value, err := spillCheckedAdd(math.MaxUint64-1, 1) - require.NoError(t, err) - require.Equal(t, uint64(math.MaxUint64), value) - _, err = spillCheckedAdd(math.MaxUint64, 1) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - value, err = spillCheckedMul(math.MaxUint64, 1) - require.NoError(t, err) - require.Equal(t, uint64(math.MaxUint64), value) - _, err = spillCheckedMul(math.MaxUint64, 2) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - _, err = spillPeakBudgetFor(math.MaxUint64, 0, 0, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillPeakBudgetFor(0, math.MaxUint64, 1, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillPeakBudgetFor(0, 0, math.MaxUint64, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestSpillCapacityReplacementOverlapChargesOldArrays(t *testing.T) { - got, err := spillCapacityReplacementOverlap(16, 4, 8, 8, 2) - require.NoError(t, err) - require.Equal(t, uint64(8*8+8*4+2*8), got) - - got, err = spillCapacityReplacementOverlap(8, 2, 8, 8, 2) - require.NoError(t, err) - require.Zero(t, got) - - _, err = spillCapacityReplacementOverlap(-1, 0, 0, 0, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, err = spillCapacityReplacementOverlap(math.MaxInt, 0, math.MaxInt-1, 0, 0) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) } -func TestSpillReplacementPeakReusesHighWaterLease(t *testing.T) { - budget := process.MustNewHashBuildBudget(120, 120) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - token, err := generation.Reserve(100) - require.NoError(t, err) - ctr := container{ - hashmapBuilder: HashmapBuilder{budget: generation}, - spillScratchReservation: token, - spillScratchBase: 100, +func (h *spillTestHarness) close(t *testing.T) { + t.Helper() + h.op.ctr.dropSpillScratchBuffers() + h.op.ctr.freeSpillExprExecs() + for _, file := range h.files { + if file != nil { + require.NoError(t, file.Close()) + } } - analyzer := process.NewAnalyzer(0, false, false, "replacement peak") - - oldSize, grew, err := ctr.growSpillScratchTransient(90, analyzer) - require.NoError(t, err) - require.False(t, grew) - require.Zero(t, oldSize) - require.Equal(t, uint64(100), generation.Used()) - - oldSize, grew, err = ctr.growSpillScratchTransient(110, analyzer) + if h.op.ctr.spillBundle != nil { + h.op.ctr.spillBundle.release() + h.op.ctr.spillBundle = nil + } + require.Zero(t, h.account.Snapshot().Used) + require.Zero(t, h.generation.Used()) + require.NoError(t, h.op.ClearAllocationAccount(h.account)) + terminal, first, err := h.registry.CompleteTerminal(h.account) require.NoError(t, err) - require.True(t, grew) - require.Equal(t, uint64(100), oldSize) - require.Equal(t, uint64(110), generation.Used()) - require.Equal(t, int64(110), analyzer.GetOpStats().ExtraStats["HashBuildSpillScratchPeakBytes"]) - require.NoError(t, ctr.restoreSpillScratchTransient(oldSize, grew)) - require.Equal(t, uint64(100), generation.Used()) - require.NoError(t, ctr.restoreSpillScratchTransient(0, false)) - - _, grew, err = ctr.growSpillScratchTransient(121, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.False(t, grew) - require.Equal(t, uint64(100), generation.Used()) - - token.Release() - require.Zero(t, generation.Used()) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + h.proc.Free() } -func TestSpillReplacementPeakReclaimsOptionalBeforeRetry(t *testing.T) { - const ( - base = uint64(spillWriteCoalesceSize) - capBytes = 2 * base - ) - for _, test := range []struct { - name string - required uint64 - wantRejects uint64 - wantError bool - }{ - {name: "reclaimed overlap fits", required: capBytes, wantRejects: 1}, - {name: "mandatory overlap exceeds cap", required: capBytes + 1, wantRejects: 2, wantError: true}, - } { - t.Run(test.name, func(t *testing.T) { - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) +func spillFileRows(t *testing.T, files []*os.File) int64 { + t.Helper() + var total int64 + for _, file := range files { + if file == nil { + continue + } + _, err := file.Seek(0, io.SeekStart) + require.NoError(t, err) + reader := bufio.NewReader(file) + for { + var header [16]byte + _, err = io.ReadFull(reader, header[:]) + if err == io.EOF { + break + } require.NoError(t, err) - defer generation.Close() - token, err := generation.Reserve(capBytes) + rows := types.DecodeInt64(header[:8]) + payload := types.DecodeInt64(header[8:]) + require.GreaterOrEqual(t, rows, int64(0)) + require.GreaterOrEqual(t, payload, int64(0)) + _, err = io.CopyN(io.Discard, reader, payload) require.NoError(t, err) - - ctr := container{ - hashmapBuilder: HashmapBuilder{budget: generation}, - spillScratchReservation: token, - spillScratchBase: base, - } - ctr.spillBucketWriteBufs[0] = *bytes.NewBuffer( - make([]byte, 0, spillWriteCoalesceSize)) - analyzer := process.NewAnalyzer(0, false, false, "replacement reclaim") - - oldSize, grew, err := ctr.growSpillScratchTransientWithReclaim( - nil, nil, test.required, analyzer) - if test.wantError { - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.False(t, grew) - require.Zero(t, oldSize) - require.Equal(t, base, token.Size()) - } else { - require.NoError(t, err) - require.True(t, grew) - require.Equal(t, base, oldSize) - require.Equal(t, test.required, token.Size()) - require.NoError(t, ctr.restoreSpillScratchTransient(oldSize, grew)) - require.Equal(t, base, token.Size()) - } - require.Equal(t, test.wantRejects, generation.RejectCount()) - require.Zero(t, ctr.spillBucketWriteBufs[0].Cap()) - token.Release() - require.Zero(t, generation.Used()) - }) + var magic [8]byte + _, err = io.ReadFull(reader, magic[:]) + require.NoError(t, err) + require.Equal(t, uint64(spillMagic), types.DecodeUint64(magic[:])) + total += rows + } } + return total } -func TestSpillPeakChargesSerializedPayloadOnce(t *testing.T) { - const ( - rows = uint64(8192) - inputBytes = uint64(3 << 20) - selectedBytes = uint64(5 << 20) - ) - got, err := spillPeakBudgetFor(rows, inputBytes, selectedBytes, 0) - require.NoError(t, err) - want := rows*12 + inputBytes + selectedBytes + selectedBytes + 64*1024 - require.Equal(t, want, got) +func TestComputeXXHashBuild(t *testing.T) { + mp := mpool.MustNewZero() + first := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, mp) + second := testutil.MakeVarcharVector([]string{"a", "b", "c"}, nil, mp) + defer first.Free(mp) + defer second.Free(mp) + hashes := make([]uint64, 3) + computeXXHash([]*vector.Vector{first, second}, hashes) + require.NotEqual(t, hashes[0], hashes[1]) + + constant := testutil.MakeInt32Vector([]int32{5}, nil, mp) + defer constant.Free(mp) + constant.SetClass(vector.CONSTANT) + computeXXHash([]*vector.Vector{constant}, hashes) + require.Equal(t, hashes[0], hashes[1]) + require.Equal(t, hashes[1], hashes[2]) } -func TestMarshalSpillRecordPreallocatesSinglePayload(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - bat := batch.NewWithSize(1) - var err error - bat.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 4<<20), 64, proc.Mp(), +func TestShouldSpillBatches(t *testing.T) { + bat := batch.NewWithSize(0) + bat.SetRowCount(2) + op := &HashBuild{IsShuffle: true, NeedHashMap: true} + op.ctr.setSpillThreshold(1) + op.ctr.hashmapBuilder.Batches.Buf = []*batch.Batch{bat} + op.ctr.hashmapBuilder.InputBatchRowCount = bat.RowCount() + require.True(t, op.shouldSpillBatches()) + op.IsShuffle = false + require.False(t, op.shouldSpillBatches()) + op.IsShuffle = true + op.NeedHashMap = false + require.False(t, op.shouldSpillBatches()) +} + +func TestAccountedSpillAdaptsAndPreservesRows(t *testing.T) { + h := newSpillTestHarness(t, 80<<10) + defer h.close(t) + values := make([]int64, colexec.DefaultBatchSize) + for i := range values { + values[i] = int64(i) + } + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeInt64Vector(values, nil, h.proc.Mp()) + input.SetRowCount(len(values)) + defer input.Clean(h.proc.Mp()) + executors, err := h.op.ctr.initSpillExprExecs( + h.proc, + []*plan.Expr{newExpr(0, types.T_int64.ToType())}, ) require.NoError(t, err) - bat.SetRowCount(64) - defer bat.Clean(proc.Mp()) - - buf := bytes.NewBuffer(make([]byte, 0, 1<<20)) - _, err = marshalSpillRecord(bat, buf) - require.NoError(t, err) - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size - } - require.Equal(t, base+128+24, uint64(buf.Cap())) - - small := batch.NewWithSize(1) - small.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 1024), 1, proc.Mp(), + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, h.op.ctr.spillBatchWithPressure( + h.proc, input, h.files, executors, analyzer, false, + )) + require.Positive(t, analyzer.GetOpStats().ExtraStats["HashBuildSpillInputReductions"]) + require.NoError(t, h.op.ctr.flushSpillBuffers(h.proc, h.files, analyzer)) + require.Equal(t, int64(len(values)), spillFileRows(t, h.files)) +} + +func TestAccountedSpillCoalescesWithoutDuplicateOwnership(t *testing.T) { + h := newSpillTestHarness(t, 8<<20) + defer h.close(t) + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 1, 1}, nil, h.proc.Mp()) + input.SetRowCount(3) + defer input.Clean(h.proc.Mp()) + executors, err := h.op.ctr.initSpillExprExecs( + h.proc, + []*plan.Expr{newExpr(0, types.T_int32.ToType())}, ) require.NoError(t, err) - small.SetRowCount(1) - defer small.Clean(proc.Mp()) - _, err = marshalSpillRecord(small, buf) - require.NoError(t, err, "a retained large serialization buffer must be reusable for a smaller batch") -} - -func TestSpillLazyReservationBoundaryInputs(t *testing.T) { - budget := process.MustNewHashBuildBudget(1, 1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "spill reservation boundary") - - require.NoError(t, ctr.spillBatchBounded(nil, nil, nil, analyzer, false)) - require.Nil(t, ctr.spillScratchReservation) - - invalid := batch.NewWithSize(1) - invalid.SetRowCount(1) - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - err = ctr.spillBatchBounded( - proc, invalid, make([]*os.File, spillNumBuckets), analyzer, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Nil(t, ctr.spillScratchReservation) - require.Zero(t, generation.Used()) + analyzer := process.NewAnalyzer(0, false, false, "test") + for range 2 { + require.NoError(t, h.op.ctr.spillBatchWithPressure( + h.proc, input, h.files, executors, analyzer, false, + )) + } + var pending int + for _, buffer := range h.op.ctr.spillAccountedBuckets { + if buffer != nil { + pending += buffer.Len() + } + } + require.Positive(t, pending) + require.NoError(t, h.op.ctr.flushSpillBuffers(h.proc, h.files, analyzer)) + require.Equal(t, int64(6), spillFileRows(t, h.files)) + require.Equal(t, h.account.Snapshot().Used, h.generation.Used()) } -func TestSpillBatchRejectsMissingExpressionOwnerBeforeScratchAdmission(t *testing.T) { +func TestSpillWithoutAllocationAccountFailsClosed(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) + bat := testutil.NewBatch([]types.Type{types.T_int32.ToType()}, true, 1, proc.Mp()) defer bat.Clean(proc.Mp()) - - need, err := spillBudgetBytes(bat) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(need, need) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - - err = ctr.spillBatchBounded( + err := (&container{}).spillBatchBounded( proc, bat, make([]*os.File, spillNumBuckets), - process.NewAnalyzer(0, false, false, "missing expression owner"), + nil, + process.NewAnalyzer(0, false, false, "test"), false, ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Nil(t, ctr.spillScratchReservation) - require.Zero(t, generation.Used()) - require.Zero(t, generation.ReserveCount()) - require.Zero(t, generation.RejectCount()) -} - -func TestSpillMaterializedBytesFollowsConstUnionSemantics(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - payload := make([]byte, 1<<20) - for i := range payload { - payload[i] = 'x' - } - const rows = 64 - source := batch.NewWithSize(1) - var err error - source.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), payload, rows, proc.Mp()) - require.NoError(t, err) - source.SetRowCount(rows) - defer source.Clean(proc.Mp()) - - directBytes, err := spillMaterializedBytes(source) - require.NoError(t, err) - require.Equal(t, - uint64(rows*types.T_varchar.ToType().TypeSize()+len(payload)), - directBytes, - "direct UnionInt32 copies one payload and broadcasts its descriptor") - - selected := batch.NewWithSize(1) - selected.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - defer selected.Clean(proc.Mp()) - sels := make([]int32, rows) - for i := range sels { - sels[i] = int32(i) - } - require.NoError(t, selected.Vecs[0].PreExtend(rows, proc.Mp())) - require.NoError(t, selected.Vecs[0].UnionInt32(source.Vecs[0], sels, proc.Mp())) - selected.SetRowCount(rows) - require.GreaterOrEqual(t, directBytes, uint64(selected.Allocated())) - - var retained colexec.Batches - defer retained.Clean(proc.Mp()) - require.NoError(t, retained.CopyIntoBatches(source, proc)) - require.Len(t, retained.Buf, 1) - require.False(t, retained.Buf[0].Vecs[0].IsConst()) - retainedBytes, err := spillMaterializedBytes(retained.Buf[0]) - require.NoError(t, err) - require.Equal(t, - uint64(rows)*(uint64(types.T_varchar.ToType().TypeSize())+uint64(len(payload))), - retainedBytes, - "the actual retained batch is non-const, so selection copies each value") - require.Greater(t, retainedBytes, directBytes*32) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvalid) } -func TestSpillMaterializedEstimateCoversRetainedConstCopy(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - payload := make([]byte, 1<<10) - for i := range payload { - payload[i] = 'x' - } - const ( - inputRows = 4 - totalRows = inputRows * 2 +func TestSpillMinimumUnitPressureIsControlled(t *testing.T) { + h := newSpillTestHarness(t, 1<<10) + defer h.close(t) + input := batch.NewWithSize(1) + input.Vecs[0] = testutil.MakeVarcharVector( + []string{strings.Repeat("x", 64<<10)}, nil, h.proc.Mp(), + ) + input.SetRowCount(1) + defer input.Clean(h.proc.Mp()) + executors, err := h.op.ctr.initSpillExprExecs( + h.proc, + []*plan.Expr{newExpr(0, types.T_varchar.ToType())}, ) - source := batch.NewWithSize(1) - var err error - source.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), payload, inputRows, proc.Mp()) - require.NoError(t, err) - source.SetRowCount(inputRows) - defer source.Clean(proc.Mp()) - - var retained colexec.Batches - defer retained.Clean(proc.Mp()) - require.NoError(t, retained.CopyIntoBatches(source, proc)) - require.NoError(t, retained.CopyIntoBatches(source, proc)) - require.Len(t, retained.Buf, 1) - require.Equal(t, totalRows, retained.Buf[0].RowCount()) - require.False(t, retained.Buf[0].Vecs[0].IsConst(), - "CopyIntoBatches materializes const ingress as retained row values") - - estimated, err := spillMaterializedBytes(retained.Buf[0]) - require.NoError(t, err) - selected := batch.NewWithSize(1) - selected.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - defer selected.Clean(proc.Mp()) - sels := make([]int32, totalRows) - for i := range sels { - sels[i] = int32(i) - } - require.NoError(t, selected.Vecs[0].PreExtend(totalRows, proc.Mp())) - require.NoError(t, selected.Vecs[0].UnionInt32(retained.Buf[0].Vecs[0], sels, proc.Mp())) - selected.SetRowCount(totalRows) - require.GreaterOrEqual(t, estimated, uint64(selected.Allocated())) -} - -func TestSpillMaterializedEstimateFollowsFullBatchCloneToSemantics(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - makeConstBatch := func(payloadBytes int) *batch.Batch { - payload := make([]byte, payloadBytes) - for i := range payload { - payload[i] = 'x' - } - source := batch.NewWithSize(1) - var err error - source.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), - payload, - colexec.DefaultBatchSize, - proc.Mp(), - ) - require.NoError(t, err) - source.SetRowCount(colexec.DefaultBatchSize) - return source - } - - large := makeConstBatch(1 << 20) - defer large.Clean(proc.Mp()) - directNeed, err := spillBudgetBytes(large) - require.NoError(t, err) - require.Less(t, directNeed, uint64(16<<20)) - - var retainedLarge colexec.Batches - defer retainedLarge.Clean(proc.Mp()) - require.NoError(t, retainedLarge.CopyIntoBatches(large, proc)) - require.Len(t, retainedLarge.Buf, 1) - require.False(t, retainedLarge.Buf[0].Vecs[0].IsConst(), - "Batch.Dup delegates to Batch.CloneTo/UnionBatch and does not call Vector.Dup") - require.Equal(t, 1<<20, len(retainedLarge.Buf[0].Vecs[0].GetArea())) - retainedNeed, err := spillMaterializedBytes(retainedLarge.Buf[0]) - require.NoError(t, err) - wantRetained := uint64(colexec.DefaultBatchSize) * - (uint64(1<<20) + uint64(types.T_varchar.ToType().TypeSize())) - require.Equal(t, wantRetained, retainedNeed, - "the actual non-const retained batch materializes one MiB plus one descriptor per row") - - // Materialize a smaller exact-full-batch payload end-to-end without - // allocating the MiB case's 8 GiB selected area. - small := makeConstBatch(4 << 10) - defer small.Clean(proc.Mp()) - var retainedSmall colexec.Batches - defer retainedSmall.Clean(proc.Mp()) - require.NoError(t, retainedSmall.CopyIntoBatches(small, proc)) - require.False(t, retainedSmall.Buf[0].Vecs[0].IsConst()) - - estimated, err := spillMaterializedBytes(retainedSmall.Buf[0]) - require.NoError(t, err) - selected := batch.NewWithSize(1) - selected.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - defer selected.Clean(proc.Mp()) - sels := make([]int32, colexec.DefaultBatchSize) - for i := range sels { - sels[i] = int32(i) - } - require.NoError(t, selected.Vecs[0].PreExtend(colexec.DefaultBatchSize, proc.Mp())) - require.NoError(t, selected.Vecs[0].UnionInt32(retainedSmall.Buf[0].Vecs[0], sels, proc.Mp())) - selected.SetRowCount(colexec.DefaultBatchSize) - require.Equal(t, colexec.DefaultBatchSize*(4<<10), len(selected.Vecs[0].GetArea())) - require.GreaterOrEqual(t, estimated, uint64(selected.Allocated())) -} - -func TestRetainedSpillRecoveryProjectionMatrix(t *testing.T) { - type ingressSpec struct { - kind string - rows int - payloadBytes int - } - tests := []struct { - name string - ingresses []ingressSpec - }{ - { - name: "fixed-boundaries-and-full-tail-swap", - ingresses: []ingressSpec{ - {kind: "fixed", rows: colexec.DefaultBatchSize - 1}, - {kind: "fixed", rows: 1}, - {kind: "fixed", rows: colexec.DefaultBatchSize}, - {kind: "fixed", rows: colexec.DefaultBatchSize + 1}, - }, - }, - { - name: "const-varlen-incremental-tail", - ingresses: []ingressSpec{ - {kind: "const", rows: 1, payloadBytes: 1024}, - {kind: "const", rows: 31, payloadBytes: 2048}, - {kind: "const", rows: colexec.DefaultBatchSize - 32, payloadBytes: 512}, - }, - }, - { - name: "const-varlen-exact-full-batch", - ingresses: []ingressSpec{ - {kind: "const", rows: colexec.DefaultBatchSize, payloadBytes: 4096}, - }, - }, - { - name: "shared-varlen-exact-full-batch", - ingresses: []ingressSpec{ - {kind: "shared", rows: colexec.DefaultBatchSize, payloadBytes: 1024}, - }, - }, - { - name: "nullable-varlen-multiple-destinations", - ingresses: []ingressSpec{ - {kind: "nullable", rows: 2*colexec.DefaultBatchSize + 17, payloadBytes: 64}, - }, - }, - { - name: "partial-tail-with-exact-full-remainder", - ingresses: []ingressSpec{ - {kind: "fixed", rows: colexec.DefaultBatchSize - 1}, - {kind: "fixed", rows: colexec.DefaultBatchSize + 1}, - }, - }, - { - name: "partial-tail-with-larger-varlen-remainder", - ingresses: []ingressSpec{ - {kind: "const", rows: colexec.DefaultBatchSize - 1, payloadBytes: 1}, - {kind: "const", rows: colexec.DefaultBatchSize + 2, payloadBytes: 64}, - }, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(4<<30, 4<<30) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - var hb HashmapBuilder - hb.setBudget(generation) - defer hb.cleanBatches(proc) - - var recoveryHighWater uint64 - for ingressIndex, spec := range test.ingresses { - source := batch.NewWithSize(1) - switch spec.kind { - case "fixed": - values := make([]int32, spec.rows) - for i := range values { - values[i] = int32(ingressIndex*100_000 + i) - } - source.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) - case "const": - payload := bytes.Repeat([]byte{'x'}, spec.payloadBytes) - source.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), payload, spec.rows, proc.Mp()) - require.NoError(t, err) - case "shared": - // UnionBatch deliberately flattens a constant while retaining one - // physical payload shared by every descriptor. This is a regular - // non-const vector, so recovery projection must follow logical row - // references rather than use class as a proxy for area ownership. - payload := bytes.Repeat([]byte{'x'}, spec.payloadBytes) - constant, constErr := vector.NewConstBytes( - types.T_varchar.ToType(), payload, spec.rows, proc.Mp()) - require.NoError(t, constErr) - source.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - require.NoError(t, source.Vecs[0].UnionBatch( - constant, 0, spec.rows, nil, proc.Mp())) - constant.Free(proc.Mp()) - require.False(t, source.Vecs[0].IsConst()) - require.Equal(t, spec.payloadBytes, len(source.Vecs[0].GetArea())) - case "nullable": - values := make([]string, spec.rows) - nulls := make([]uint64, 0, spec.rows/5+1) - for i := range values { - values[i] = strings.Repeat(string(rune('a'+i%17)), spec.payloadBytes) - if i%5 == 0 { - nulls = append(nulls, uint64(i)) - } - } - source.Vecs[0] = testutil.MakeVarcharVector(values, nulls, proc.Mp()) - default: - t.Fatalf("unknown ingress kind %q", spec.kind) - } - source.SetRowCount(spec.rows) - - projection, projectionErr := hb.projectedBatchCopy(source) - require.NoError(t, projectionErr) - projectedNeed, projectionErr := spillRetainedRecoveryBudgetBytes(projection) - require.NoError(t, projectionErr) - projectedNeed, projectionErr = spillRecoveryReservationBytes(projectedNeed) - require.NoError(t, projectionErr) - recoveryHighWater = max(recoveryHighWater, projectedNeed) - - require.NoError(t, hb.copyBuildBatchProjected(source, proc, projection)) - require.Equal(t, projection.nextTailSelected, hb.retainedSpillTailSelected) - source.Clean(proc.Mp()) - - for retainedIndex, retained := range hb.Batches.Buf { - actualNeed, actualErr := spillScratchBudgetBytes(retained, true) - require.NoError(t, actualErr) - require.LessOrEqualf(t, actualNeed, recoveryHighWater, - "ingress=%d retained=%d rows=%d", ingressIndex, retainedIndex, retained.RowCount()) - } - } - - if test.name == "const-varlen-exact-full-batch" { - wantSelected := uint64(colexec.DefaultBatchSize) * - (uint64(types.T_varchar.ToType().TypeSize()) + 4096) - actualSelected, err := spillMaterializedBytes(hb.Batches.Buf[0]) - require.NoError(t, err) - require.Equal(t, wantSelected, actualSelected) - } - if test.name == "partial-tail-with-exact-full-remainder" { - require.Zero(t, hb.retainedSpillTailSelected) - require.Equal(t, colexec.DefaultBatchSize, - hb.Batches.Buf[len(hb.Batches.Buf)-1].RowCount()) - } - if test.name == "partial-tail-with-larger-varlen-remainder" { - tail := hb.Batches.Buf[len(hb.Batches.Buf)-1] - require.Equal(t, 1, tail.RowCount()) - actualSelected, err := spillMaterializedBytes(tail) - require.NoError(t, err) - require.Equal(t, actualSelected, hb.retainedSpillTailSelected) - } - - hb.cleanBatches(proc) - require.Zero(t, hb.retainedSpillTailSelected) - require.Zero(t, generation.Used()) - require.Zero(t, proc.Mp().CurrNB()) - }) - } -} - -func TestSpillRecoveryReservationRoundingBoundaries(t *testing.T) { - tests := []struct { - need uint64 - want uint64 - }{ - {need: 0, want: 0}, - {need: 1, want: spillRecoveryReservationQuantum}, - {need: spillRecoveryReservationQuantum - 1, want: spillRecoveryReservationQuantum}, - {need: spillRecoveryReservationQuantum, want: spillRecoveryReservationQuantum}, - {need: spillRecoveryReservationQuantum + 1, want: 2 * spillRecoveryReservationQuantum}, - } - for _, test := range tests { - got, err := spillRecoveryReservationBytes(test.need) - require.NoError(t, err) - require.Equal(t, test.want, got) - } - _, err := spillRecoveryReservationBytes(math.MaxUint64) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) -} - -func TestSpillRecoveryReservationLifecycle(t *testing.T) { - t.Run("reserve-reuse-grow-release", func(t *testing.T) { - const capBytes = 2 * spillRecoveryReservationQuantum - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "recovery lifecycle") - - require.NoError(t, ctr.ensureSpillRecoveryReservationBytes( - spillRecoveryReservationQuantum, analyzer)) - token := ctr.spillScratchReservation - require.NotNil(t, token) - require.Equal(t, spillRecoveryReservationQuantum, token.Size()) - require.Equal(t, spillRecoveryReservationQuantum, generation.Used()) - - reserveCount := generation.ReserveCount() - require.NoError(t, ctr.ensureSpillRecoveryReservationBytes(1, analyzer)) - require.Same(t, token, ctr.spillScratchReservation) - require.Equal(t, reserveCount, generation.ReserveCount()) - - require.NoError(t, ctr.ensureSpillRecoveryReservationBytes(capBytes, analyzer)) - require.Same(t, token, ctr.spillScratchReservation) - require.Equal(t, capBytes, ctr.spillScratchBase) - require.Equal(t, capBytes, token.Size()) - require.Equal(t, capBytes, generation.Used()) - require.Equal(t, reserveCount+1, generation.ReserveCount()) - - extra := analyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillRecoveryGrowCount"]) - require.Equal(t, int64(spillRecoveryReservationQuantum), - extra["HashBuildSpillRecoveryGrowBytes"]) - require.Equal(t, int64(capBytes), extra["HashBuildSpillRecoveryReservedBytes"]) - - ctr.releaseSpillScratchReservation() - require.Nil(t, ctr.spillScratchReservation) - require.Zero(t, ctr.spillScratchBase) - require.Zero(t, generation.Used()) - }) - - t.Run("grow-rejection-preserves-old-lease", func(t *testing.T) { - const capBytes = spillRecoveryReservationQuantum - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "recovery grow rejection") - require.NoError(t, ctr.ensureSpillRecoveryReservationBytes(capBytes, analyzer)) - token := ctr.spillScratchReservation - - err = ctr.ensureSpillRecoveryReservationBytes(2*capBytes, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Same(t, token, ctr.spillScratchReservation) - require.Equal(t, capBytes, ctr.spillScratchBase) - require.Equal(t, capBytes, token.Size()) - require.Equal(t, capBytes, generation.Used()) - require.Equal(t, uint64(1), generation.ReserveCount()) - require.Equal(t, uint64(1), generation.RejectCount()) - - extra := analyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["HashBuildSpillRecoveryGrowRejects"]) - require.Zero(t, extra["HashBuildSpillRecoveryGrowCount"]) - require.NoError(t, ctr.ensureSpillRecoveryReservationBytes(capBytes, analyzer), - "a failed grow must leave the prior recovery lease reusable") - - ctr.releaseSpillScratchReservation() - require.Zero(t, generation.Used()) - }) -} - -func TestSpillRecoveryRejectsInvalidProofsWithoutChargingBudget(t *testing.T) { - const capBytes = spillRecoveryReservationQuantum - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "invalid recovery proof") - - upper, hasVarlen, err := spillDirectRecoveryBudgetUpper(nil) - require.NoError(t, err) - require.Zero(t, upper) - require.False(t, hasVarlen) - require.NoError(t, ctr.ensureDirectSpillRecovery(nil, analyzer)) - - malformed := batch.NewOffHeapWithSize(1) - malformed.SetRowCount(1) - err = ctr.ensureDirectSpillRecovery(malformed, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Nil(t, ctr.spillScratchReservation) - require.Zero(t, generation.Used()) - - for _, test := range []struct { - name string - projection batchCopyProjection - }{ - {name: "empty-destination", projection: batchCopyProjection{}}, - { - name: "selected-size-overflow", - projection: batchCopyProjection{ - maxRetainedRows: 1, - maxRetainedSelected: math.MaxUint64, - columns: 1, - }, - }, - } { - t.Run(test.name, func(t *testing.T) { - err := ctr.ensureRetainedSpillRecovery(test.projection, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.Nil(t, ctr.spillScratchReservation) - require.Zero(t, generation.Used()) - }) - } -} - -func TestSpillDirectRecoveryUpperBoundsExactMatrix(t *testing.T) { - tests := []struct { - name string - makeBatch func(*testing.T, *process.Process) *batch.Batch - hasVarlen bool - exactUpper bool - }{ - { - name: "fixed", - makeBatch: func(_ *testing.T, proc *process.Process) *batch.Batch { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, proc.Mp()) - bat.SetRowCount(4) - return bat - }, - exactUpper: true, - }, - { - name: "const-varlen", - makeBatch: func(t *testing.T, proc *process.Process) *batch.Batch { - bat := batch.NewWithSize(1) - var err error - bat.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), bytes.Repeat([]byte{'x'}, 4096), - colexec.DefaultBatchSize, proc.Mp()) - require.NoError(t, err) - bat.SetRowCount(colexec.DefaultBatchSize) - return bat - }, - hasVarlen: true, - exactUpper: true, - }, - { - name: "shared-varlen", - makeBatch: func(t *testing.T, proc *process.Process) *batch.Batch { - constant, err := vector.NewConstBytes( - types.T_varchar.ToType(), bytes.Repeat([]byte{'x'}, 4096), - colexec.DefaultBatchSize, proc.Mp()) - require.NoError(t, err) - defer constant.Free(proc.Mp()) - bat := batch.NewWithSize(1) - bat.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) - require.NoError(t, bat.Vecs[0].UnionBatch( - constant, 0, colexec.DefaultBatchSize, nil, proc.Mp())) - require.False(t, bat.Vecs[0].IsConst()) - require.Equal(t, 4096, len(bat.Vecs[0].GetArea())) - bat.SetRowCount(colexec.DefaultBatchSize) - return bat - }, - hasVarlen: true, - exactUpper: true, - }, - { - name: "nullable-varlen", - makeBatch: func(_ *testing.T, proc *process.Process) *batch.Batch { - const rows = 1024 - values := make([]string, rows) - nulls := make([]uint64, 0, rows/2) - for i := range values { - values[i] = strings.Repeat("x", 128) - if i%2 == 0 { - nulls = append(nulls, uint64(i)) - } - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nulls, proc.Mp()) - bat.SetRowCount(rows) - return bat - }, - hasVarlen: true, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - bat := test.makeBatch(t, proc) - defer bat.Clean(proc.Mp()) - - upper, hasVarlen, err := spillDirectRecoveryBudgetUpper(bat) - require.NoError(t, err) - exact, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Equal(t, test.hasVarlen, hasVarlen) - require.GreaterOrEqual(t, upper, exact) - if test.exactUpper { - require.Equal(t, exact, upper) - } - - capBytes, err := spillRecoveryReservationBytes(exact) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(capBytes, capBytes) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - analyzer := process.NewAnalyzer(0, false, false, "direct recovery upper") - require.NoError(t, ctr.ensureDirectSpillRecovery(bat, analyzer)) - reserved := ctr.spillScratchBase - reserveCount := generation.ReserveCount() - require.NoError(t, ctr.ensureDirectSpillRecovery(bat, analyzer)) - require.Equal(t, reserved, ctr.spillScratchBase) - require.Equal(t, reserveCount, generation.ReserveCount()) - ctr.releaseSpillScratchReservation() - require.Zero(t, generation.Used()) - }) - } -} - -func TestSpillBatchLazyReservationFailsClosed(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, nil, proc.Mp()) - bat.SetRowCount(4) - defer bat.Clean(proc.Mp()) - - need, err := spillBudgetBytes(bat) - require.NoError(t, err) - require.Positive(t, need) - budget := process.MustNewHashBuildBudget(need-1, need-1) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - ctr := &container{} - ctr.hashmapBuilder.setBudget(generation) - err = initSpillExprExecsForTest(ctr, proc, []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }}) require.NoError(t, err) - files := make([]*os.File, spillNumBuckets) - err = ctr.spillBatchBounded( - proc, bat, files, - process.NewAnalyzer(0, false, false, "direct spill reject"), false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, ctr.spillScratchReservation) - require.Nil(t, ctr.spillHashValues) - require.Nil(t, ctr.spillBucketRowIds) - require.Equal(t, 4, bat.RowCount()) - for _, file := range files { - require.Nil(t, file) - } - ctr.hashmapBuilder.FreeExecutors() - require.Zero(t, generation.Used()) + err = h.op.ctr.spillBatchWithPressure( + h.proc, + input, + h.files, + executors, + process.NewAnalyzer(0, false, false, "test"), + false, + ) + var minimum *MinimumAllocationPressureError + require.True(t, errors.As(err, &minimum), "unexpected error: %v", err) } -func TestEnsureSpillFile(t *testing.T) { +func TestWriteSpillPayloadCancellationStopsBeforeIO(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - ctr := &container{spillUUID: "test_ensure"} - files := make([]*os.File, spillNumBuckets) - - // First call creates a file. - f, err := ctr.ensureSpillFile(proc, files, 3) - require.NoError(t, err) - require.NotNil(t, f) - require.Equal(t, f, files[3]) - defer f.Close() - - // Second call returns cached file. - f2, err := ctr.ensureSpillFile(proc, files, 3) + ctx, cancel := context.WithCancelCause(proc.Ctx) + process.ReplacePipelineCtx(proc, ctx, cancel) + spillfs, err := proc.GetSpillFileService() require.NoError(t, err) - require.Same(t, f, f2, "should return the same file object") - - // Different bucket creates a different file. - f3, err := ctr.ensureSpillFile(proc, files, 7) + file, err := spillfs.CreateFile(context.Background(), t.Name()) require.NoError(t, err) - require.NotNil(t, f3) - require.NotEqual(t, f.Fd(), f3.Fd()) - defer f3.Close() - - // Untouched buckets remain nil. - require.Nil(t, files[0]) - require.Nil(t, files[1]) -} - -func TestCleanupSpillFiles(t *testing.T) { - // Create temp files to simulate spill fds. - var fds []*os.File - for i := 0; i < 3; i++ { - f, err := os.CreateTemp("", "test_cleanup_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - fds = append(fds, f) - } - // Include a nil entry. - fds = append(fds, nil) - - hb := &HashBuild{ctr: container{spilledFds: fds}} - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - hb.cleanupSpillFiles(proc) - require.Nil(t, hb.ctr.spilledFds) - - // Verify all files are closed (writing should fail). - for _, f := range fds[:3] { - _, err := f.Write([]byte("x")) - require.Error(t, err, "file should be closed") - } + defer func() { + require.NoError(t, file.Close()) + require.NoError(t, spillfs.RemoveFile(context.Background(), t.Name())) + }() + proc.Cancel(context.Canceled) + err = (&container{}).writeSpillPayload( + proc, + file, + []byte("stale"), + 1, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorIs(t, err, context.Canceled) + info, statErr := file.Stat() + require.NoError(t, statErr) + require.Zero(t, info.Size()) } diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go index 0a3bd6a90fe8f..8e19ad938e113 100644 --- a/pkg/sql/colexec/hashbuild/types.go +++ b/pkg/sql/colexec/hashbuild/types.go @@ -15,14 +15,16 @@ package hashbuild import ( - "bytes" "os" "sync" "sync/atomic" + "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/hashtable" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/logutil" @@ -46,6 +48,53 @@ const ( SendSucceed ) +const HashBuildAllocationOwner mpool.AllocationOwner = 1 + +const ( + HashBuildSpillAllocationSiteSelectedData mpool.AllocationSite = iota + 64 + HashBuildSpillAllocationSiteSelectedArea + HashBuildSpillAllocationSiteSelectedNulls + HashBuildSpillAllocationSiteSelectedGrouping + HashBuildSpillAllocationSiteHashValues + HashBuildSpillAllocationSiteRowIDs + HashBuildSpillAllocationSiteMarshalBuffer + HashBuildSpillAllocationSiteCoalesceBuffer +) + +const ( + HashBuildAllocationSiteHashCell mpool.AllocationSite = iota + 24 + HashBuildAllocationSiteHashDescriptor + HashBuildAllocationSiteBatchData + HashBuildAllocationSiteBatchArea + HashBuildAllocationSiteBatchNulls + HashBuildAllocationSiteBatchGrouping + HashBuildAllocationSiteGroupSels + HashBuildAllocationSiteHashIterator +) + +// Runtime-filter keys and their published wire payload have lifetimes that +// differ from copied build batches: keys die after publication while the +// payload lives on the message board until every receiver destroys it. Keep +// their sites distinct from both the builder and SpillEngine ranges. +const ( + HashBuildAllocationSiteUniqueKeyData mpool.AllocationSite = iota + 44 + HashBuildAllocationSiteUniqueKeyArea + HashBuildAllocationSiteUniqueKeyNulls + HashBuildAllocationSiteUniqueKeyGrouping + HashBuildAllocationSiteRuntimeFilterPayload + HashBuildAllocationSiteRuntimeFilterScratch + HashBuildAllocationSiteDedupIgnoreBitmap + HashBuildAllocationSiteDedupDeleteBitmap + HashBuildAllocationSiteDedupLastRows + HashBuildAllocationSiteDedupSurvivorRows + HashBuildAllocationSiteDedupSurvivorOwnsKey + HashBuildAllocationSiteDedupDiscardedRows + HashBuildAllocationSiteDedupDeleteOnlyData + HashBuildAllocationSiteDedupDeleteOnlyArea + HashBuildAllocationSiteDedupDeleteOnlyNulls + HashBuildAllocationSiteDedupDeleteOnlyGrouping +) + type container struct { state int runtimeFilterIn bool @@ -72,26 +121,24 @@ type container struct { // input batch. spillBucketOffsets identifies each bucket's sub-slice; // keeping one array avoids the 32 independent append/growth paths used by // the old scatter implementation. - spillBucketRowIds []int32 - spillBucketCounts [spillNumBuckets]int32 - spillBucketOffsets [spillNumBuckets + 1]int32 - spillSelection []int32 - spillWriteBuf bytes.Buffer - // spillBucketWriteBufs coalesce serialized records across source batches. - // Each buffer is bounded by spillWriteCoalesceSize (plus bytes.Buffer's - // bounded growth slack), so fanout does not imply fanout-sized vectors. - spillBucketWriteBufs [spillNumBuckets]bytes.Buffer - spillBucketWriteRows [spillNumBuckets]int64 - spillKeyVecs []*vector.Vector - // spillScratchReservation is the query/CN-charged recovery lease for this - // execution. Shuffle HashBuild establishes it before retaining spillable - // data; spill then reuses the same lease for its bounded scratch buffers. - // Direct spill callers may still establish it lazily. Reset, Free, and the - // build terminal cleanup all release it idempotently. - spillScratchReservation *process.HashBuildReservation - // spillScratchBase is the retained scratch floor. Coalesce-buffer growth is - // charged on top and must never be mistaken for this floor. - spillScratchBase uint64 + spillBucketRowIds []int32 + spillBucketCounts [spillNumBuckets]int32 + spillBucketOffsets [spillNumBuckets + 1]int32 + spillBucketWriteRows [spillNumBuckets]int64 + spillKeyVecs []*vector.Vector + spillBatchAllocation *vector.AllocationAccountSelection + spillAllocationMP *mpool.MPool + spillAccountedWrite *mpool.AccountedBuffer + spillAccountedBuckets [spillNumBuckets]*mpool.AccountedBuffer + spillCoalesceDisabled bool + recoveryCapacity *process.HashBuildRecoveryCapacity + recoveryCapacityClass mpool.AllocationCapacityClass + expressionRecoveryPeak uint64 + expressionRecoveryRows int + spillRecoveryPeak uint64 + // cached expression executors for spill (reused across batches) + spillExprExecs []colexec.ExpressionExecutor + spillConditions []*plan.Expr } // spillFileBundle is deliberately owned by hashbuild. Build converts each @@ -167,7 +214,7 @@ func (b *spillFileBundle) growDisk(file *os.File, budget *process.HashBuildBudge b.mu.Lock() defer b.mu.Unlock() if b.released { - return 0, false, process.ErrHashBuildReservationInactive + return 0, false, process.ErrHashBuildSpillReservationInactive } if b.entries == nil { b.entries = make(map[*os.File]*spillFileEntry) @@ -265,6 +312,246 @@ func (hashBuild *HashBuild) GetOperatorBase() *vm.OperatorBase { return &hashBuild.OperatorBase } +// SetAllocationAccount selects immutable provenance for the hash-table owner +// before Prepare. Compile invokes it once for each execution attempt; Reset +// clears the selection only after producer or JoinMap ownership has moved on. +func (hashBuild *HashBuild) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + selection, err := vector.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteSelectedData, + HashBuildSpillAllocationSiteSelectedArea, + HashBuildSpillAllocationSiteSelectedNulls, + HashBuildSpillAllocationSiteSelectedGrouping, + ) + if err != nil { + return err + } + if err := hashBuild.ctr.hashmapBuilder.SetAllocationAccount(account); err != nil { + return err + } + hashBuild.ctr.spillBatchAllocation = selection + return nil +} + +func (hashBuild *HashBuild) installRecoveryCapacity( + budget *process.HashBuildBudgetGeneration, +) error { + ctr := &hashBuild.ctr + account := ctr.hashmapBuilder.mapAllocationAccount + if account == nil { + return mpool.ErrAllocationAccountInvariant + } + if ctr.recoveryCapacity != nil { + if ctr.recoveryCapacityClass == mpool.AllocationCapacityClassDefault || + ctr.hashmapBuilder.recoveryCapacityClass != ctr.recoveryCapacityClass { + return mpool.ErrAllocationAccountInvariant + } + return nil + } + if ctr.recoveryCapacityClass != mpool.AllocationCapacityClassDefault { + return mpool.ErrAllocationAccountInvariant + } + capacity, err := process.NewHashBuildRecoveryCapacity(budget) + if err != nil { + return err + } + class, err := account.RegisterCapacityController(capacity) + if err != nil { + _ = capacity.Close() + return err + } + selection, err := vector.NewAllocationAccountSelectionWithCapacityClass( + account, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteSelectedData, + HashBuildSpillAllocationSiteSelectedArea, + HashBuildSpillAllocationSiteSelectedNulls, + HashBuildSpillAllocationSiteSelectedGrouping, + class, + ) + if err != nil { + _ = account.UnregisterCapacityController(class, capacity) + _ = capacity.Close() + return err + } + ctr.recoveryCapacity = capacity + ctr.recoveryCapacityClass = class + ctr.spillBatchAllocation = selection + ctr.hashmapBuilder.recoveryCapacityClass = class + return nil +} + +// releaseRecoveryCapacity returns recovery headroom after retained spill state +// has been drained or build reaches a terminal result. restoreDefault keeps +// later direct/test/reuse allocations on the statement's ordinary controller; +// statement teardown passes false and drops the selection immediately afterward. +func (hashBuild *HashBuild) releaseRecoveryCapacity( + account *mpool.AllocationAccount, + restoreDefault bool, +) error { + ctr := &hashBuild.ctr + if ctr.recoveryCapacity == nil { + return nil + } + if account == nil || ctr.recoveryCapacityClass == + mpool.AllocationCapacityClassDefault { + return mpool.ErrAllocationAccountInvariant + } + capacity := ctr.recoveryCapacity + class := ctr.recoveryCapacityClass + if err := capacity.Close(); err != nil { + return err + } + if err := account.UnregisterCapacityController(class, capacity); err != nil { + return err + } + ctr.recoveryCapacity = nil + ctr.recoveryCapacityClass = mpool.AllocationCapacityClassDefault + ctr.expressionRecoveryPeak = 0 + ctr.expressionRecoveryRows = 0 + ctr.spillRecoveryPeak = 0 + ctr.hashmapBuilder.recoveryCapacityClass = + mpool.AllocationCapacityClassDefault + if !restoreDefault { + ctr.spillBatchAllocation = nil + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildSpillAllocationSiteSelectedData, + HashBuildSpillAllocationSiteSelectedArea, + HashBuildSpillAllocationSiteSelectedNulls, + HashBuildSpillAllocationSiteSelectedGrouping, + ) + if err != nil { + ctr.spillBatchAllocation = nil + return err + } + ctr.spillBatchAllocation = selection + return nil +} + +// SetAllocationAccount installs the physical allocation provenance shared by +// the producer HashBuild and SpillEngine rebuild builders. A builder is always +// single-generation and clears the selection only after all owned resources +// have either been freed or transferred to a JoinMap. +func (hb *HashmapBuilder) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := hb + if builder.mapAllocationAccount != nil { + if builder.mapAllocationAccount == account { + return nil + } + return mpool.ErrAllocationAccountMismatch + } + selection, err := hashtable.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashCell, + HashBuildAllocationSiteHashDescriptor, + ) + if err != nil { + return err + } + iteratorAllocation, err := hashmap.NewIteratorAllocation( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteHashIterator, + ) + if err != nil { + return err + } + batchSelection, err := vector.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteBatchData, + HashBuildAllocationSiteBatchArea, + HashBuildAllocationSiteBatchNulls, + HashBuildAllocationSiteBatchGrouping, + ) + if err != nil { + return err + } + uniqueKeySelection, err := vector.NewAllocationAccountSelection( + account, + HashBuildAllocationOwner, + HashBuildAllocationSiteUniqueKeyData, + HashBuildAllocationSiteUniqueKeyArea, + HashBuildAllocationSiteUniqueKeyNulls, + HashBuildAllocationSiteUniqueKeyGrouping, + ) + if err != nil { + return err + } + builder.mapAllocationAccount = account + builder.mapAllocation = selection + builder.iteratorAllocation = iteratorAllocation + builder.batchAllocation = batchSelection + builder.uniqueKeyAllocation = uniqueKeySelection + return nil +} + +func (hashBuild *HashBuild) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := &hashBuild.ctr.hashmapBuilder + if len(hashBuild.ctr.spillExprExecs) != 0 { + return mpool.ErrAllocationAccountInvariant + } + if hashBuild.ctr.spillAllocationMP != nil || + hashBuild.ctr.spillAccountedWrite != nil { + return mpool.ErrAllocationAccountInvariant + } + for _, buffer := range hashBuild.ctr.spillAccountedBuckets { + if buffer != nil { + return mpool.ErrAllocationAccountInvariant + } + } + if hashBuild.ctr.recoveryCapacity != nil { + if err := hashBuild.releaseRecoveryCapacity(account, false); err != nil { + return err + } + } + if err := builder.ClearAllocationAccount(account); err != nil { + return err + } + hashBuild.ctr.spillBatchAllocation = nil + return nil +} + +// ClearAllocationAccount verifies that no builder-owned object can allocate +// through the generation before dropping its selections. +func (hb *HashmapBuilder) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + builder := hb + if builder.mapAllocationAccount == nil { + return nil + } + if builder.mapAllocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if builder.IntHashMap != nil || builder.StrHashMap != nil || + len(builder.Batches.Buf) != 0 || builder.Sels.Size() != 0 || + len(builder.executors) != 0 || len(builder.curVecs) != 0 || + len(builder.UniqueJoinKeys) != 0 || + builder.IgnoreRows != nil || builder.DelRows != nil { + return mpool.ErrAllocationAccountInvariant + } + builder.mapAllocationAccount = nil + builder.mapAllocation = nil + builder.iteratorAllocation = nil + builder.batchAllocation = nil + builder.uniqueKeyAllocation = nil + builder.recoveryCapacityClass = mpool.AllocationCapacityClassDefault + return nil +} + func init() { reuse.CreatePool[HashBuild]( func() *HashBuild { @@ -315,14 +602,13 @@ func (hashBuild *HashBuild) Reset(proc *process.Process, pipelineFailed bool, er hashBuild.publishBuildError(proc, err) } else { // Preserve the established nil JoinMap convention for a true empty - // build and for legacy cleanup paths that completed without a map. + // build and for cleanup paths that completed without a map. hashBuild.publishJoinMap(proc, nil) } } hashBuild.ctr.hashmapBuilder.Reset(proc, !mapSucceed) hashBuild.ctr.dropSpillScratchBuffers() - hashBuild.ctr.releaseSpillScratchReservation() // Only clean up build files when the join map was NOT successfully sent. // When mapSucceed=true, hashjoin owns the files and deletes them after reading. if !mapSucceed { @@ -363,8 +649,8 @@ func (hashBuild *HashBuild) Free(proc *process.Process, pipelineFailed bool, err hashBuild.cleanupSpillFiles(proc) hashBuild.ctr.spillFS = nil hashBuild.ctr.hashmapBuilder.Free(proc) + hashBuild.ctr.freeSpillExprExecs() hashBuild.ctr.dropSpillScratchBuffers() - hashBuild.ctr.releaseSpillScratchReservation() } func (hashBuild *HashBuild) logDiagnostics(proc *process.Process, pipelineFailed bool, err error) { @@ -396,25 +682,23 @@ func hasHashBuildDiagnosticStats(extra map[string]int64) bool { extra["HashBuildRuntimeFilterCollectionFallbacks"] != 0 || extra["HashBuildRuntimeFilterBudgetFallbacks"] != 0 || extra["HashBuildRuntimeFilterAllocationFallbacks"] != 0 || - extra["HashBuildSpillRecoveryReserveRejects"] != 0 || - extra["HashBuildSpillRecoveryGrowRejects"] != 0 || - extra["HashBuildSpillRecoveryGrowCount"] != 0 || - extra["HashBuildSpillScratchReserveRejects"] != 0 || - extra["HashBuildSpillScratchGrowRejects"] != 0 || - extra["HashBuildSpillScratchGrowCount"] != 0 + extra["HashBuildSpillRecoveryReserveRejects"] != 0 } func (hashBuild *HashBuild) publishJoinMap(proc *process.Process, jm *message.JoinMap) bool { if !atomic.CompareAndSwapUint32(&hashBuild.ctr.terminalPublished, 0, 1) { return false } - message.SendJoinMapResult( + if !message.SendJoinMapResult( message.NewJoinMapResult(jm), hashBuild.JoinMapTag, hashBuild.IsShuffle, hashBuild.ShuffleIdx, proc.GetMessageBoard(), - ) + ) { + atomic.StoreUint32(&hashBuild.ctr.terminalPublished, 0) + return false + } return true } @@ -422,13 +706,16 @@ func (hashBuild *HashBuild) publishBuildError(proc *process.Process, err error) if !atomic.CompareAndSwapUint32(&hashBuild.ctr.terminalPublished, 0, 1) { return false } - message.FinalizeJoinMapBuildError( + if !message.FinalizeJoinMapBuildError( proc.GetMessageBoard(), hashBuild.JoinMapTag, hashBuild.IsShuffle, hashBuild.ShuffleIdx, err, - ) + ) { + atomic.StoreUint32(&hashBuild.ctr.terminalPublished, 0) + return false + } return true } @@ -448,9 +735,8 @@ func (hashBuild *HashBuild) cleanupSpillFiles(proc *process.Process) { } } -// CleanCopiedBatchAt is the lifecycle hook used by bounded initial spill. -// HashBuild keeps this wrapper on the operator side so batch reservation -// ownership remains private to the hashbuild package. +// CleanCopiedBatchAt releases one retained build batch after it has been +// durably transferred to spill storage. func (hb *HashmapBuilder) CleanCopiedBatchAt(idx int, proc *process.Process) error { if idx < 0 || idx >= len(hb.Batches.Buf) { return process.ErrHashBuildBudgetInvalid @@ -466,23 +752,15 @@ func (hb *HashmapBuilder) CleanCopiedBatchAt(idx int, proc *process.Process) err hb.Batches.MemSize += int64(bat.Size()) } } - // CopyIntoBatches can coalesce several ingress batches into one physical - // batch (and can reorder a full batch around a partial tail), so an ingress - // reservation cannot be matched safely to Batches.Buf[idx]. Keep the - // conservative charges until the last physical batch has been dropped. if len(hb.Batches.Buf) == 0 { hb.retainedSpillTailSelected = 0 - hb.releaseBatchReservations() } return nil } // DrainCopiedBatches visits and then releases every retained physical build // batch. A failed visit leaves the current and remaining batches owned by the -// builder so its normal cleanup path can release them. CopyIntoBatches can -// coalesce several ingress batches into one physical batch, so the associated -// reservations are released together only after the final physical batch is -// destroyed. +// builder so its normal cleanup path can release them. func (hb *HashmapBuilder) DrainCopiedBatches( proc *process.Process, visit func(*batch.Batch) error, @@ -511,7 +789,6 @@ func (hb *HashmapBuilder) DrainCopiedBatches( hb.Batches.Buf = nil hb.Batches.MemSize = 0 hb.retainedSpillTailSelected = 0 - hb.releaseBatchReservations() return nil } diff --git a/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..5c00c5dca350d --- /dev/null +++ b/pkg/sql/colexec/hashjoin/allocation_test_helpers_test.go @@ -0,0 +1,84 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashjoin + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +type testAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} + +func TestHashJoinResultBatchUsesAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &HashJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + account := installTestAllocation(t, arg) + require.NoError(t, arg.resetResultBat()) + require.Same(t, arg.resultAllocation, arg.ctr.resBat.Vecs[0].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp())) + used := account.Snapshot().Used + require.Positive(t, used) + require.NoError(t, arg.resetResultBat()) + require.Equal(t, used, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.resBat) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestHashJoinResultBatchHonorsAllocationCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + arg := &HashJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetResultBat()) + err = vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) +} diff --git a/pkg/sql/colexec/hashjoin/bitmap_mailbox.go b/pkg/sql/colexec/hashjoin/bitmap_mailbox.go new file mode 100644 index 0000000000000..519207b70754f --- /dev/null +++ b/pkg/sql/colexec/hashjoin/bitmap_mailbox.go @@ -0,0 +1,98 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashjoin + +import ( + "context" + "sync" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" +) + +// BitmapMailbox is the single ownership boundary for the parallel right-join +// bitmap exchange. A successful Send transfers ownership to the mailbox. Once +// sealed, late senders retain ownership and normal operator cleanup frees it. +type BitmapMailbox struct { + mu sync.Mutex + sealed bool + ch chan *bitmap.Bitmap +} + +func NewBitmapMailbox(workers int) *BitmapMailbox { + if workers < 1 { + workers = 1 + } + return &BitmapMailbox{ch: make(chan *bitmap.Bitmap, workers)} +} + +func (m *BitmapMailbox) Send(value *bitmap.Bitmap) bool { + if m == nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.sealed { + return false + } + // Every non-merger publishes at most once and the mailbox capacity equals + // the worker count, so publication cannot block while holding mu. + m.ch <- value + return true +} + +func (m *BitmapMailbox) Receive( + ctx context.Context, +) (*bitmap.Bitmap, bool) { + if m == nil || ctx == nil { + return nil, false + } + select { + case <-ctx.Done(): + return nil, false + case value := <-m.ch: + return value, true + } +} + +// SealAndDrain makes cancellation order-independent. It owns and frees every +// value already transferred into the mailbox; concurrent or later Send calls +// fail and leave ownership with their sender. +func (m *BitmapMailbox) SealAndDrain(mp *mpool.MPool) { + if m == nil { + return + } + m.mu.Lock() + m.sealed = true + for { + select { + case value := <-m.ch: + colexec.FreeAccountedBitmap(value, mp) + default: + m.mu.Unlock() + return + } + } +} + +func (m *BitmapMailbox) Terminal() bool { + if m == nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + return m.sealed && len(m.ch) == 0 +} diff --git a/pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go b/pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go new file mode 100644 index 0000000000000..763bf52d0f1b9 --- /dev/null +++ b/pkg/sql/colexec/hashjoin/bitmap_mailbox_test.go @@ -0,0 +1,105 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package hashjoin + +import ( + "sync" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/bitmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/stretchr/testify/require" +) + +func newMailboxAccountedBitmap( + t *testing.T, + mp *mpool.MPool, + account *mpool.AllocationAccount, +) *bitmap.Bitmap { + t.Helper() + value, err := colexec.NewAccountedBitmap( + 1024, + mp, + account, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteMatchedRows, + ) + require.NoError(t, err) + return value +} + +func TestBitmapMailboxSealOwnsQueuedAndRejectsLateTransfer(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + mp := mpool.MustNewZero() + mailbox := NewBitmapMailbox(2) + + queued := newMailboxAccountedBitmap(t, mp, account) + require.True(t, mailbox.Send(queued)) + mailbox.SealAndDrain(mp) + require.Zero(t, account.Snapshot().Used) + require.Empty(t, mailbox.ch) + + late := newMailboxAccountedBitmap(t, mp, account) + require.False(t, mailbox.Send(late)) + require.NotZero(t, account.Snapshot().Used) + colexec.FreeAccountedBitmap(late, mp) + require.Zero(t, account.Snapshot().Used) + + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) +} + +func TestBitmapMailboxConcurrentSealPreservesSingleOwner(t *testing.T) { + const workers = 16 + registry, err := mpool.NewAllocationAccountRegistry(1, workers+1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + mp := mpool.MustNewZero() + mailbox := NewBitmapMailbox(workers) + values := make([]*bitmap.Bitmap, workers) + for i := range values { + values[i] = newMailboxAccountedBitmap(t, mp, account) + } + + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(workers) + for _, value := range values { + go func(value *bitmap.Bitmap) { + defer wg.Done() + <-start + if !mailbox.Send(value) { + colexec.FreeAccountedBitmap(value, mp) + } + }(value) + } + close(start) + mailbox.SealAndDrain(mp) + wg.Wait() + mailbox.SealAndDrain(mp) + + require.Zero(t, account.Snapshot().Used) + require.Empty(t, mailbox.ch) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/sql/colexec/hashjoin/expression_memory_test.go b/pkg/sql/colexec/hashjoin/expression_memory_test.go deleted file mode 100644 index 747d35fa2341a..0000000000000 --- a/pkg/sql/colexec/hashjoin/expression_memory_test.go +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed 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. - -package hashjoin - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestHashJoinResetReleasesProbeExpressionLease(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I32Val{I32Val: 1}, - }}, - } - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - arg := &HashJoin{} - arg.ctr.eqCondExecs = executors - arg.ctr.eqCondVecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionLease = lease - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, generation.Used()) - - arg.Reset(proc, false, nil) - require.Zero(t, generation.Used()) - require.Nil(t, arg.ctr.eqCondExecs) - require.Nil(t, arg.ctr.eqCondVecs) - require.Nil(t, arg.ctr.probeExpressionLease) -} diff --git a/pkg/sql/colexec/hashjoin/join.go b/pkg/sql/colexec/hashjoin/join.go index fc932967c3060..84d6bba355821 100644 --- a/pkg/sql/colexec/hashjoin/join.go +++ b/pkg/sql/colexec/hashjoin/join.go @@ -17,9 +17,9 @@ package hashjoin import ( "bytes" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -76,6 +76,9 @@ func (hashJoin *HashJoin) Prepare(proc *process.Process) (err error) { return err } } + if hashJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if hashJoin.OpAnalyzer == nil { hashJoin.OpAnalyzer = process.NewAnalyzer(hashJoin.GetIdx(), hashJoin.IsFirst, hashJoin.IsLast, opName) @@ -91,20 +94,30 @@ func (hashJoin *HashJoin) Prepare(proc *process.Process) (err error) { } if len(ctr.eqCondVecs) == 0 { - eqCondExecs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, hashJoin.EqConds[0]) + eqCondExecs, err := hashbuild.NewExpressionExecutors( + proc, + hashJoin.EqConds[0], + hashJoin.allocationAccount, + ) if err != nil { return err } var nonEqCondExec colexec.ExpressionExecutor if hashJoin.NonEqCond != nil { - nonEqCondExec, err = colexec.NewExpressionExecutor(proc, hashJoin.NonEqCond) + var nonEqExecs []colexec.ExpressionExecutor + nonEqExecs, err = hashbuild.NewExpressionExecutors( + proc, + []*plan.Expr{hashJoin.NonEqCond}, + hashJoin.allocationAccount, + ) if err != nil { for _, exec := range eqCondExecs { exec.Free() } return err } + nonEqCondExec = nonEqExecs[0] } ctr.eqCondVecs = make([]*vector.Vector, len(hashJoin.EqConds[0])) @@ -202,7 +215,9 @@ func (hashJoin *HashJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.lastIdx = 0 } - hashJoin.resetResultBat() + if err = hashJoin.resetResultBat(); err != nil { + return result, err + } for i, rp := range hashJoin.ResultCols { if rp.Rel == 0 { ctr.resBat.Vecs[i].SetSorted(ctr.leftBat.Vecs[rp.Pos].GetSorted()) @@ -282,7 +297,7 @@ func (hashJoin *HashJoin) Call(proc *process.Process) (vm.CallResult, error) { // For spilled join, clean up current bucket and move to next if (ctr.spillEngine != nil) && (ctr.spillEngine.HasMoreBuckets() || ctr.spillEngine.IsProbing()) { - ctr.rightRowsMatched = nil + ctr.freeRightRowsMatched(proc) ctr.cleanHashMap() ctr.state = Probe } @@ -337,18 +352,13 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process if takeErr != nil { return takeErr } - probeExpressionLease, leaseErr := hashbuild.NewExpressionMemoryLease( - budget, hashJoin.EqConds[0], ctr.eqCondExecs, false) - if leaseErr != nil { + if hashJoin.allocationAccount == nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil - ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() - return leaseErr + return mpool.ErrAllocationAccountInvalid } - ctr.probeExpressionLease = probeExpressionLease - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ BuildKeyExprs: hashJoin.EqConds[1], ProbeKeyExprs: hashJoin.EqConds[0], SpillThreshold: ctr.spillThreshold, @@ -358,13 +368,16 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process NeedAllocateSels: !hashJoin.HashOnPK, NeedBatches: hashJoin.NeedBuildBatches(), Budget: budget, - ProbeExpressionLease: probeExpressionLease, - }) - if len(payload.Files) > 0 { - engine.InitFromSpilledFiles(payload.Files) - } else { - engine.InitFromSpilledMap(payload.LegacyFds) + }, hashJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) + if engineErr != nil { + _ = payload.Close() + ctr.mp.Free() + ctr.mp = nil + ctr.cleanEqCondExecutors() + return engineErr } + engine.InitFromSpilledFiles(payload.Files) + ctr.spillEngine = engine if err := engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { input, err := vm.ChildrenCall(hashJoin.GetChildren(0), proc, analyzer) @@ -381,10 +394,10 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process ctr.mp.Free() ctr.mp = nil engine.Cleanup(proc) + ctr.spillEngine = nil return err } ctr.mp.Free() - ctr.spillEngine = engine ctr.mp = nil return nil } @@ -399,8 +412,16 @@ func (hashJoin *HashJoin) build(analyzer process.Analyzer, proc *process.Process if hashJoin.EmitUnmatchedBuild() { if ctr.rightRowCnt > 0 { - ctr.rightRowsMatched = &bitmap.Bitmap{} - ctr.rightRowsMatched.InitWithSize(ctr.rightRowCnt) + ctr.rightRowsMatched, err = colexec.NewAccountedBitmap( + ctr.rightRowCnt, + proc.Mp(), + hashJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteMatchedRows, + ) + if err != nil { + return err + } } } @@ -441,6 +462,7 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p // Load next bucket via engine convenience method. if ctr.mp == nil { + var allocationErr error ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, res spillutil.BucketResult) { if res == spillutil.BucketReady { @@ -449,8 +471,14 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p ctr.rightRowCnt = jm.GetRowCount() ctr.probeHashOnPK = hashJoin.HashOnPK || ctr.mp.HashOnUnique() if hashJoin.EmitUnmatchedBuild() && ctr.rightRowCnt > 0 { - ctr.rightRowsMatched = &bitmap.Bitmap{} - ctr.rightRowsMatched.InitWithSize(ctr.rightRowCnt) + ctr.rightRowsMatched, allocationErr = + colexec.NewAccountedBitmap( + ctr.rightRowCnt, + proc.Mp(), + hashJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteMatchedRows, + ) ctr.rightMatchedIter = nil } } @@ -458,9 +486,13 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p if err != nil { return result, err } + if allocationErr != nil { + return result, allocationErr + } if !ok { return result, nil } + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil ctr.probeState = psNextBatch ctr.lastIdx = 0 @@ -471,7 +503,7 @@ func (hashJoin *HashJoin) getSpilledInputBatch(proc *process.Process, analyzer p } func (ctr *container) probe(hashJoin *HashJoin, proc *process.Process, result *vm.CallResult) error { - err := ctr.evalJoinConditionBudgeted(ctr.leftBat, proc) + err := ctr.evalJoinCondition(ctr.leftBat, proc) if err != nil { return err } @@ -496,7 +528,15 @@ func (ctr *container) probe(hashJoin *HashJoin, proc *process.Process, result *v case psNextBatch: if ctr.lastIdx < leftRowCnt { hashBatch := min(leftRowCnt-ctr.lastIdx, hashmap.UnitLimit) - ctr.vs, ctr.zvs = ctr.itr.Find(ctr.lastIdx, hashBatch, ctr.eqCondVecs) + var err error + ctr.vs, ctr.zvs, err = ctr.itr.Find( + ctr.lastIdx, + hashBatch, + ctr.eqCondVecs, + ) + if err != nil { + return err + } ctr.vsIdx = 0 ctr.probeState = psBatchRow } else { @@ -787,7 +827,7 @@ func (ctr *container) appendMarkForEmptyBuildBucket(marker *vector.Vector, proc return vector.SetConstNull(marker, rowCnt, proc.Mp()) } - if err := ctr.evalJoinConditionBudgeted(ctr.leftBat, proc); err != nil { + if err := ctr.evalJoinCondition(ctr.leftBat, proc); err != nil { return err } if err := vector.AppendMultiFixed(marker, false, false, rowCnt, proc.Mp()); err != nil { @@ -795,12 +835,18 @@ func (ctr *container) appendMarkForEmptyBuildBucket(marker *vector.Vector, proc } for _, vec := range ctr.eqCondVecs { if vec.IsConstNull() { + if err := marker.PreExtendNulls(rowCnt, proc.Mp()); err != nil { + return err + } marker.GetNulls().AddRange(0, uint64(rowCnt)) return nil } if !vec.GetNulls().Any() { continue } + if err := marker.PreExtendNulls(rowCnt, proc.Mp()); err != nil { + return err + } nulls.Or(marker.GetNulls(), vec.GetNulls(), marker.GetNulls()) } return nil @@ -815,34 +861,34 @@ func (ctr *container) syncBitmap(hashJoin *HashJoin, proc *process.Process) erro if hashJoin.NumCPU > 1 { if !hashJoin.IsMerger { - hashJoin.Channel <- ctr.rightRowsMatched + if hashJoin.Mailbox.Send(ctr.rightRowsMatched) { + ctr.rightRowsMatched = nil + } return nil } else { matchedCnt := ctr.rightRowsMatched.Count() for cnt := 1; cnt < int(hashJoin.NumCPU); cnt++ { - v := colexec.ReceiveBitmapFromChannel(proc.Ctx, hashJoin.Channel) - if v == nil { + v, received := hashJoin.Mailbox.Receive(proc.Ctx) + if !received || v == nil { // A worker was torn down before syncing (its Reset sends - // nil) or the context was canceled. The merge is aborted, - // but keep draining this generation's remaining messages - // so no stale bitmap is left behind in the shared - // channel, then bail out without initializing the + // nil) or the context was canceled. Sealing transfers all + // already-published values to cleanup and makes late + // publishers retain their own value. Bail out without initializing the // iterator — Call routes to End and nothing is finalized. - for cnt++; cnt < int(hashJoin.NumCPU); cnt++ { - colexec.ReceiveBitmapFromChannel(proc.Ctx, hashJoin.Channel) - } + hashJoin.Mailbox.SealAndDrain(proc.Mp()) return nil } matchedCnt += v.Count() ctr.rightRowsMatched.Or(v) + colexec.FreeAccountedBitmap(v, proc.Mp()) } if ctr.probeSingle && matchedCnt > ctr.rightRowsMatched.Count() { return moerr.NewErrSubqueryNo1Row(proc.Ctx) } - close(hashJoin.Channel) + hashJoin.Mailbox.SealAndDrain(proc.Mp()) } } @@ -856,7 +902,9 @@ func (ctr *container) syncBitmap(hashJoin *HashJoin, proc *process.Process) erro } func (ctr *container) finalize(hashJoin *HashJoin, proc *process.Process, result *vm.CallResult) error { - hashJoin.resetResultBat() + if err := hashJoin.resetResultBat(); err != nil { + return err + } rowCnt := 0 for ; rowCnt < colexec.DefaultBatchSize && ctr.rightMatchedIter.HasNext(); rowCnt++ { @@ -984,17 +1032,7 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) return nil } -func (ctr *container) evalJoinConditionBudgeted(bat *batch.Batch, proc *process.Process) error { - if ctr.probeExpressionLease == nil { - return ctr.evalJoinCondition(bat, proc) - } - return ctr.probeExpressionLease.Eval(proc, []*batch.Batch{bat}, bat.RowCount(), func(i int, vec *vector.Vector) error { - ctr.eqCondVecs[i] = vec - return nil - }) -} - -func (hashJoin *HashJoin) resetResultBat() { +func (hashJoin *HashJoin) resetResultBat() error { ctr := &hashJoin.ctr if ctr.resBat != nil { ctr.resBat.CleanOnlyData() @@ -1015,5 +1053,11 @@ func (hashJoin *HashJoin) resetResultBat() { ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(types.T_bool.ToType()) } } + if err := ctr.resBat.SetAllocationAccount(hashJoin.resultAllocation); err != nil { + ctr.resBat.Clean(nil) + ctr.resBat = nil + return err + } } + return nil } diff --git a/pkg/sql/colexec/hashjoin/join_test.go b/pkg/sql/colexec/hashjoin/join_test.go index e449f0f4a6ff3..9d87669172458 100644 --- a/pkg/sql/colexec/hashjoin/join_test.go +++ b/pkg/sql/colexec/hashjoin/join_test.go @@ -61,6 +61,7 @@ func TestHashJoinPrepareFailureCanRetry(t *testing.T) { EqConds: [][]*plan.Expr{{valid}, {valid}}, NonEqCond: invalid, } + installTestAllocation(t, arg) require.Error(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.eqCondVecs) @@ -786,9 +787,9 @@ func TestHashJoinSingleRejectsDuplicateMatchesAcrossWorkers(t *testing.T) { JoinType: plan.Node_SINGLE, NumCPU: 2, IsMerger: true, - Channel: make(chan *bitmap.Bitmap, 1), + Mailbox: NewBitmapMailbox(2), } - hashJoin.Channel <- remoteMatches + require.True(t, hashJoin.Mailbox.Send(remoteMatches)) ctr := container{rightRowsMatched: localMatches, probeSingle: true} err := ctr.syncBitmap(hashJoin, proc) @@ -820,14 +821,14 @@ func TestHashJoinMergerSyncBitmapAborted(t *testing.T) { IsRightJoin: true, NumCPU: 3, IsMerger: true, - Channel: make(chan *bitmap.Bitmap, 3), + Mailbox: NewBitmapMailbox(3), ResultCols: []colexec.ResultPos{colexec.NewResultPos(1, 0)}, RightTypes: []types.Type{types.T_int32.ToType()}, } // Worker A was torn down before syncing (its Reset sends nil); worker B // synced normally and its bitmap lands after the abort marker. - hashJoin.Channel <- nil - hashJoin.Channel <- staleMatches + require.True(t, hashJoin.Mailbox.Send(nil)) + require.True(t, hashJoin.Mailbox.Send(staleMatches)) hashJoin.ctr.state = SyncBitmap hashJoin.ctr.rightRowsMatched = matched hashJoin.ctr.rightBats = []*batch.Batch{rightBat} @@ -836,13 +837,13 @@ func TestHashJoinMergerSyncBitmapAborted(t *testing.T) { require.NoError(t, err) require.Nil(t, result.Batch) require.Equal(t, vm.ExecStop, result.Status) - // Worker B's bitmap must not be left behind in the shared channel. - require.Empty(t, hashJoin.Channel) + // Worker B's bitmap must not be left behind in the shared mailbox. + require.Empty(t, hashJoin.Mailbox.ch) // The merger already synced this generation, so Reset must not push the // nil abort marker either. hashJoin.Reset(proc, false, nil) - require.Empty(t, hashJoin.Channel) + require.Empty(t, hashJoin.Mailbox.ch) // Next generation over the same operator and channel: a clean sync must // only observe this generation's bitmaps. @@ -856,8 +857,9 @@ func TestHashJoinMergerSyncBitmapAborted(t *testing.T) { workerMatches2.InitWithSize(4) workerMatches2.Add(2) - hashJoin.Channel <- workerMatches1 - hashJoin.Channel <- workerMatches2 + hashJoin.Mailbox = NewBitmapMailbox(3) + require.True(t, hashJoin.Mailbox.Send(workerMatches1)) + require.True(t, hashJoin.Mailbox.Send(workerMatches2)) hashJoin.ctr.state = SyncBitmap hashJoin.ctr.rightRowsMatched = matched2 hashJoin.ctr.rightBats = []*batch.Batch{rightBat} @@ -896,11 +898,11 @@ func TestHashJoinMergerFinalizeEmitsUnmatchedBuildRows(t *testing.T) { IsRightJoin: true, NumCPU: 2, IsMerger: true, - Channel: make(chan *bitmap.Bitmap, 2), + Mailbox: NewBitmapMailbox(2), ResultCols: []colexec.ResultPos{colexec.NewResultPos(1, 0)}, RightTypes: []types.Type{types.T_int32.ToType()}, } - hashJoin.Channel <- remoteMatches + require.True(t, hashJoin.Mailbox.Send(remoteMatches)) hashJoin.ctr.state = SyncBitmap hashJoin.ctr.rightRowsMatched = matched hashJoin.ctr.rightBats = []*batch.Batch{rightBat} @@ -1049,7 +1051,7 @@ func newTestCaseWithMPool( resultBatch.Vecs[i] = bat.Vecs[rp[i].Pos] } tag++ - return joinTestCase{ + tc := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -1088,6 +1090,8 @@ func newTestCaseWithMPool( }, resultBatch: resultBatch, } + installTestAllocation(t, tc.arg, tc.barg) + return tc } func resetChildren(arg *HashJoin, m *mpool.MPool) { diff --git a/pkg/sql/colexec/hashjoin/key_contract_test.go b/pkg/sql/colexec/hashjoin/key_contract_test.go index 9f6bde1f22f51..ce84365955879 100644 --- a/pkg/sql/colexec/hashjoin/key_contract_test.go +++ b/pkg/sql/colexec/hashjoin/key_contract_test.go @@ -351,6 +351,7 @@ func runHashJoinKeyContract( if mode.shuffle { buildArg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: joinMapTag + 7000} } + installTestAllocation(t, arg, buildArg) var build, probe *batch.Batch defer func() { arg.Free(proc, false, nil) diff --git a/pkg/sql/colexec/hashjoin/mark_spill_test.go b/pkg/sql/colexec/hashjoin/mark_spill_test.go index 6549f6b56eb77..72927a3642172 100644 --- a/pkg/sql/colexec/hashjoin/mark_spill_test.go +++ b/pkg/sql/colexec/hashjoin/mark_spill_test.go @@ -15,7 +15,6 @@ package hashjoin import ( - "os" "testing" "github.com/matrixorigin/matrixone/pkg/container/batch" @@ -191,7 +190,7 @@ func TestHashMarkJoinEmptySpillBucketTruthTable(t *testing.T) { tc.arg.ctr.leftBat = probe tc.arg.ctr.globalBuildRowCnt = tt.globalBuildRowCnt tc.arg.ctr.buildHasNullKey = tt.buildHasNullKey - tc.arg.resetResultBat() + require.NoError(t, tc.arg.resetResultBat()) var result vm.CallResult require.NoError(t, tc.arg.ctr.emptyProbe(tc.arg, tc.proc, &result)) @@ -220,6 +219,7 @@ func TestHashMarkJoinEmptySpillBucketTruthTable(t *testing.T) { // spill scatter phase, and an empty global build makes even NULL probes FALSE. func TestHashMarkJoinSpilledEmptyBuild(t *testing.T) { tc := newMarkSpillTestCase(t) + generation, registry, account := installHashJoinTestAllocation(t, tc.arg) probe := batch.NewWithSize(1) probe.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 0}, []uint64{1}, tc.proc.Mp()) probe.SetRowCount(2) @@ -229,14 +229,14 @@ func TestHashMarkJoinSpilledEmptyBuild(t *testing.T) { jm.SetRowCount(0) jm.IncRef(1) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ - LegacyFds: make([]*os.File, spillutil.SpillNumBuckets), + Files: make([]*message.SpillFile, spillutil.SpillNumBuckets), + BudgetRef: generation, })) message.SendMessage(message.JoinMapMsg{ - JoinMapPtr: jm, + Result: message.NewJoinMapResult(jm), IsShuffle: true, ShuffleIdx: tc.arg.ShuffleIdx, Tag: tc.arg.JoinMapTag, - Spilled: true, }, tc.proc.GetMessageBoard()) require.NoError(t, tc.arg.Prepare(tc.proc)) @@ -245,4 +245,7 @@ func TestHashMarkJoinSpilledEmptyBuild(t *testing.T) { 1: {value: false}, }, collectMarkResults(t, &tc)) finishMarkSpillTest(t, &tc) + require.Zero(t, account.Snapshot().Used) + _, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) } diff --git a/pkg/sql/colexec/hashjoin/spill_diskv2_test.go b/pkg/sql/colexec/hashjoin/spill_diskv2_test.go index 64d28bb960aad..b9ce3e4eb8fa1 100644 --- a/pkg/sql/colexec/hashjoin/spill_diskv2_test.go +++ b/pkg/sql/colexec/hashjoin/spill_diskv2_test.go @@ -17,15 +17,17 @@ package hashjoin import ( "bytes" "context" - "os" + "io" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/defines" "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) @@ -74,23 +76,36 @@ func TestHashJoinSpillDiskV2(t *testing.T) { buildBat.Vecs[0] = testutil.MakeInt32Vector(buildVals, nil, proc.Mp()) buildBat.SetRowCount(100) - // spill the build bucket via FlushBucketBatch + // Write one production-format spill record to the DISK-V2 file. buildFile, err := spillfs.CreateAndRemoveFile(context.Background(), "diskv2_build") require.NoError(t, err) - var buf bytes.Buffer - bw := spillutil.BucketWriter{Name: "diskv2_build", Fd: buildFile} - err = spillutil.FlushBucketBatch(proc, buildBat, &bw, &buf, nil) + var payload bytes.Buffer + err = buildBat.MarshalBinaryWithGroupingTo(&payload) + require.NoError(t, err) + rows, size, magic := int64(buildBat.RowCount()), int64(payload.Len()), uint64(spillutil.SpillMagic) + for _, part := range [][]byte{ + types.EncodeInt64(&rows), + types.EncodeInt64(&size), + payload.Bytes(), + types.EncodeUint64(&magic), + } { + _, err = buildFile.Write(part) + require.NoError(t, err) + } + _, err = buildFile.Seek(0, io.SeekStart) + require.NoError(t, err) + info, err := buildFile.Stat() require.NoError(t, err) - buildFd := bw.HandOffFd() - require.NotNil(t, buildFd) // rebuild via SpillEngine - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine := newAccountedTestSpillEngine(t, spillutil.SpillEngineConfig{ BuildKeyExprs: makeKeyExpr(), NeedsBuildForEmptyProbe: true, NeedBatches: true, }) - engine.InitFromSpilledMap([]*os.File{buildFd}) + engine.InitFromSpilledFiles([]*message.SpillFile{ + message.NewSpillFile(buildFile, 100, uint64(info.Size()), nil), + }) analyzer := process.NewAnalyzer(0, false, false, "test") jm, res, err := engine.RebuildHashmap(proc, analyzer) diff --git a/pkg/sql/colexec/hashjoin/spill_integration_test.go b/pkg/sql/colexec/hashjoin/spill_integration_test.go index e3261f099ba36..22ac081d58765 100644 --- a/pkg/sql/colexec/hashjoin/spill_integration_test.go +++ b/pkg/sql/colexec/hashjoin/spill_integration_test.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/sql/colexec/spillutil" "github.com/matrixorigin/matrixone/pkg/testutil" metricv2 "github.com/matrixorigin/matrixone/pkg/util/metric/v2" @@ -39,13 +40,56 @@ func makeKeyExpr() []*plan.Expr { }} } +func newAccountedTestSpillEngine( + t *testing.T, + cfg spillutil.SpillEngineConfig, +) *spillutil.SpillEngine { + t.Helper() + if cfg.Budget == nil { + budget := process.MustNewHashBuildBudget(1<<60, 1<<60) + var err error + cfg.Budget, err = budget.OpenGeneration(1) + require.NoError(t, err) + } + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, cfg.Budget) + require.NoError(t, err) + engine, err := spillutil.NewSpillEngine( + cfg, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + return engine +} + +func installHashJoinTestAllocation( + t *testing.T, + join *HashJoin, +) (*process.HashBuildBudgetGeneration, *mpool.AllocationAccountRegistry, *mpool.AllocationAccount) { + t.Helper() + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(64<<20, generation) + require.NoError(t, err) + if join.allocationAccount != nil { + require.NoError(t, join.ClearAllocationAccount(join.allocationAccount)) + } + require.NoError(t, join.SetAllocationAccount(account)) + return generation, registry, account +} + // TestGetSpilledInputBatchNoBuckets verifies that getSpilledInputBatch // returns nil when the engine has no buckets. func TestGetSpilledInputBatchNoBuckets(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{}) + engine := newAccountedTestSpillEngine(t, spillutil.SpillEngineConfig{}) hashJoin := &HashJoin{ctr: container{spillEngine: engine}} result, err := hashJoin.getSpilledInputBatch(proc, process.NewAnalyzer(0, false, false, "test")) require.NoError(t, err) @@ -168,6 +212,152 @@ func TestShuffleJoinFiniteBudgetInitialSpillAndReSpill(t *testing.T) { require.Zero(t, tc.proc.Mp().CurrNB()) } +func TestShuffleJoinSpillUsesCanonicalGroupingPartitionKey(t *testing.T) { + for _, test := range []struct { + name string + typ types.Type + rows int + spillThreshold int64 + wantRespill bool + probe func(*process.Process) *vector.Vector + build func(*process.Process) *vector.Vector + }{ + { + name: "varchar", + typ: types.T_varchar.ToType(), + rows: 1, + spillThreshold: 1, + probe: func(proc *process.Process) *vector.Vector { + return testutil.MakeVarcharVector([]string{"probe"}, nil, proc.Mp()) + }, + build: func(proc *process.Process) *vector.Vector { + return testutil.MakeVarcharVector([]string{"build"}, nil, proc.Mp()) + }, + }, + { + name: "int32", + typ: types.T_int32.ToType(), + rows: 1, + spillThreshold: 1, + probe: func(proc *process.Process) *vector.Vector { + return testutil.MakeInt32Vector([]int32{222}, nil, proc.Mp()) + }, + build: func(proc *process.Process) *vector.Vector { + return testutil.MakeInt32Vector([]int32{111}, nil, proc.Mp()) + }, + }, + { + name: "scaled-float32", + rows: 8192, + spillThreshold: 50, + wantRespill: true, + typ: func() types.Type { + typ := types.T_float32.ToType() + typ.Scale = 2 + return typ + }(), + probe: func(proc *process.Process) *vector.Vector { + values := make([]float32, 8192) + values[0] = 1.234 + for i := 1; i < len(values); i++ { + values[i] = float32(i)/10 + 0.001 + } + vec := testutil.MakeFloat32Vector(values, nil, proc.Mp()) + vec.GetType().Scale = 2 + return vec + }, + build: func(proc *process.Process) *vector.Vector { + values := make([]float32, 8192) + values[0] = 9.876 + for i := 1; i < len(values); i++ { + values[i] = float32(i) / 10 + } + vec := testutil.MakeFloat32Vector(values, nil, proc.Mp()) + vec.GetType().Scale = 2 + return vec + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + keyExpr := []*plan.Expr{{ + Typ: plan.Type{ + Id: int32(test.typ.Oid), + Width: test.typ.Width, + Scale: test.typ.Scale, + }, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ + ColPos: 0, + }}, + }} + tc := newTestCase( + t, + []bool{false}, + []types.Type{test.typ}, + []colexec.ResultPos{colexec.NewResultPos(0, 0)}, + [][]*plan.Expr{keyExpr, keyExpr}, + ) + tc.arg.NonEqCond = nil + tc.arg.IsShuffle = true + tc.arg.ShuffleIdx = 0 + tc.arg.SpillThreshold = test.spillThreshold + tc.barg.IsShuffle = true + tc.barg.ShuffleIdx = 0 + tc.barg.SpillThreshold = test.spillThreshold + tc.barg.NeedBatches = false + tc.barg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{ + Tag: tc.arg.JoinMapTag + 1_500, + } + + probe := batch.NewWithSize(1) + probe.Vecs[0] = test.probe(tc.proc) + probe.SetRowCount(test.rows) + build := batch.NewWithSize(1) + build.Vecs[0] = test.build(tc.proc) + build.SetRowCount(test.rows) + probe.Vecs[0].GetGrouping().Add(0) + build.Vecs[0].GetGrouping().Add(0) + resetChildrenWithBatch(tc.arg, probe) + resetHashBuildChildrenWithBatch(tc.barg, build) + + spillBefore := promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")) + respillBefore := promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("respill", "2")) + + require.NoError(t, tc.arg.Prepare(tc.proc)) + require.NoError(t, tc.barg.Prepare(tc.proc)) + _, err := vm.Exec(tc.barg, tc.proc) + require.NoError(t, err) + + rows := 0 + for { + result, err := vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + if result.Batch != nil { + rows += result.Batch.RowCount() + } + if result.Status == vm.ExecStop { + break + } + } + require.Equal(t, test.rows, rows) + require.Greater(t, promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")), spillBefore) + if test.wantRespill { + require.Greater(t, promtestutil.ToFloat64( + metricv2.HashBuildSpillDepthCounter.WithLabelValues("respill", "2")), respillBefore) + } + + tc.arg.Reset(tc.proc, false, nil) + tc.barg.Reset(tc.proc, false, nil) + tc.arg.Free(tc.proc, false, nil) + tc.barg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) + }) + } +} + func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { tc := newTestCase( t, @@ -179,10 +369,10 @@ func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { // This cap admits one bounded scatter pass and per-bucket rebuild, but not // the complete 8K-row retained build/map. The very high soft threshold // proves that spill is entered from hard admission rejection, not policy. - // Lazy spill scratch no longer consumes resident headroom before spill. - // Keep this cap below the resident hashmap peak while leaving enough room - // for the bounded scatter pass after the rejected map is released. - tc.proc.Base.Lim.Size = 1536 << 10 + // Leave enough capacity for the mandatory bounded recovery pass itself. + // The retained 8K-row build/map still cannot fit, so ordinary allocation + // admission is what drives the operator into spill mode. + tc.proc.Base.Lim.Size = 512 << 10 tc.proc.Base.Lim.SpillSize = 64 << 20 const rows = 8192 @@ -207,8 +397,18 @@ func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { buildInput := colexec.NewMockOperator().WithBatchs([]*batch.Batch{build1, build2}) tc.barg.Children = nil tc.barg.AppendChild(buildInput) + oldAccount := tc.arg.allocationAccount + require.NoError(t, tc.arg.ClearAllocationAccount(oldAccount)) + require.NoError(t, tc.barg.ClearAllocationAccount(oldAccount)) + budget, err := tc.proc.GetHashBuildBudget() + require.NoError(t, err) + registry, err := budget.AllocationAccountRegistry() + require.NoError(t, err) + account, err := registry.OpenWithController(budget.Snapshot().Cap, budget) + require.NoError(t, err) + require.NoError(t, tc.arg.SetAllocationAccount(account)) + require.NoError(t, tc.barg.SetAllocationAccount(account)) - rejectBefore := promtestutil.ToFloat64(metricv2.HashBuildBudgetEventCounter.WithLabelValues("memory", "reject", "query")) spillBefore := promtestutil.ToFloat64(metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")) require.NoError(t, tc.arg.Prepare(tc.proc)) require.NoError(t, tc.barg.Prepare(tc.proc)) @@ -228,13 +428,12 @@ func TestShuffleJoinHardBudgetRejectTransitionsToSpill(t *testing.T) { } } require.ElementsMatch(t, values, resultValues) - require.Greater(t, promtestutil.ToFloat64(metricv2.HashBuildBudgetEventCounter.WithLabelValues("memory", "reject", "query")), rejectBefore) + require.Positive(t, tc.barg.OpAnalyzer.GetOpStats().ExtraStats["HashBuildSpillStarts"]) + require.LessOrEqual(t, account.Snapshot().Peak, account.Snapshot().Limit) require.Greater(t, promtestutil.ToFloat64(metricv2.HashBuildSpillDepthCounter.WithLabelValues("spill", "1")), spillBefore) tc.arg.Free(tc.proc, false, nil) tc.barg.Free(tc.proc, false, nil) - budget, err := tc.proc.GetHashBuildBudget() - require.NoError(t, err) require.Zero(t, budget.Used()) require.Zero(t, budget.SpillDiskUsed()) require.Zero(t, budget.SpillFDUsed()) diff --git a/pkg/sql/colexec/hashjoin/terminal_budget_test.go b/pkg/sql/colexec/hashjoin/terminal_budget_test.go index 3ee3c2b8441c4..8763dd9348ca5 100644 --- a/pkg/sql/colexec/hashjoin/terminal_budget_test.go +++ b/pkg/sql/colexec/hashjoin/terminal_budget_test.go @@ -42,7 +42,7 @@ func TestHashJoinCallConvertsTerminalBudgetAdmission(t *testing.T) { admission := &process.HashBuildBudgetError{ Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, + Component: process.HashBuildBudgetComponentMemory, Requested: 2, Used: 1, Cap: 1, diff --git a/pkg/sql/colexec/hashjoin/types.go b/pkg/sql/colexec/hashjoin/types.go index c0330797c94c7..77adc90291530 100644 --- a/pkg/sql/colexec/hashjoin/types.go +++ b/pkg/sql/colexec/hashjoin/types.go @@ -17,6 +17,7 @@ package hashjoin import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -48,6 +49,15 @@ const ( psBatchRow ) +const hashJoinAllocationSiteMatchedRows mpool.AllocationSite = 80 + +const ( + hashJoinAllocationSiteResultData mpool.AllocationSite = iota + 102 + hashJoinAllocationSiteResultArea + hashJoinAllocationSiteResultNulls + hashJoinAllocationSiteResultGrouping +) + type container struct { state int itr hashmap.Iterator @@ -104,13 +114,9 @@ type container struct { maxAllocSize int64 // spill support - spillEngine *spillutil.SpillEngine - spillThreshold int64 - // Non-nil only for spilled joins, where probe expressions are part of the - // shared HashBuild/spill working set. Resident probe expressions remain - // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - probeBucketActive bool // true while reading probe batches from a bucket + spillEngine *spillutil.SpillEngine + spillThreshold int64 + probeBucketActive bool // true while reading probe batches from a bucket } type HashJoin struct { @@ -125,7 +131,7 @@ type HashJoin struct { NonEqCond *plan.Expr EqConds [][]*plan.Expr - Channel chan *bitmap.Bitmap + Mailbox *BitmapMailbox NumCPU uint64 HashOnPK bool @@ -138,10 +144,65 @@ type HashJoin struct { RuntimeFilterSpecs []*plan.RuntimeFilterSpec JoinMapTag int32 SpillThreshold int64 + allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } +func (hashJoin *HashJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if hashJoin.allocationAccount != nil && + hashJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if hashJoin.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + hashJoinAllocationSiteResultData, + hashJoinAllocationSiteResultArea, + hashJoinAllocationSiteResultNulls, + hashJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } + hashJoin.allocationAccount = account + hashJoin.resultAllocation = selection + return nil +} + +func (hashJoin *HashJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if hashJoin.allocationAccount == nil { + return nil + } + if hashJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if hashJoin.ctr.mp != nil || hashJoin.ctr.spillEngine != nil || + len(hashJoin.ctr.eqCondExecs) != 0 || + hashJoin.ctr.nonEqCondExec != nil || + hashJoin.ctr.rightRowsMatched != nil || + hashJoin.ctr.resBat != nil { + return mpool.ErrAllocationAccountInvariant + } + if hashJoin.NumCPU > 1 && !hashJoin.Mailbox.Terminal() { + return mpool.ErrAllocationAccountInvariant + } + hashJoin.allocationAccount = nil + hashJoin.resultAllocation = nil + return nil +} + func (hashJoin *HashJoin) GetOperatorBase() *vm.OperatorBase { return &hashJoin.OperatorBase } @@ -191,22 +252,23 @@ func (hashJoin *HashJoin) ExecProjection(proc *process.Process, input *batch.Bat func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err error) { ctr := &hashJoin.ctr + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil if !ctr.bitmapSynced && hashJoin.NumCPU > 1 && !hashJoin.IsMerger { - hashJoin.Channel <- nil + hashJoin.Mailbox.Send(nil) } - // SpillEngine borrows the probe executor lease. End that borrow before the - // join frees the executors and releases their reservation. - ctr.cleanBucketBatches(proc) - if ctr.probeExpressionLease != nil { - ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() - } else { - ctr.resetEqCondExecutors() + if hashJoin.NumCPU > 1 && hashJoin.IsMerger { + hashJoin.Mailbox.SealAndDrain(proc.Mp()) } + ctr.cleanBucketBatches(proc) + ctr.cleanEqCondExecutors() ctr.cleanHashMap() - ctr.resetNonEqCondExecutor() - ctr.rightRowsMatched = nil + ctr.cleanNonEqCondExecutor() + if ctr.resBat != nil { + ctr.resBat.Clean(proc.GetMPool()) + ctr.resBat = nil + } + ctr.freeRightRowsMatched(proc) ctr.rightMatchedIter = nil ctr.skipProbe = false ctr.bitmapSynced = false @@ -216,7 +278,6 @@ func (hashJoin *HashJoin) Reset(proc *process.Process, pipelineFailed bool, err ctr.state = Build ctr.probeState = psNextBatch ctr.lastIdx = 0 - if hashJoin.OpAnalyzer != nil { hashJoin.OpAnalyzer.Alloc(ctr.maxAllocSize) } @@ -229,17 +290,10 @@ func (hashJoin *HashJoin) Free(proc *process.Process, pipelineFailed bool, err e ctr.cleanBatch(proc) ctr.cleanBucketBatches(proc) ctr.cleanEqCondExecutors() - ctr.releaseProbeExpressionLease() ctr.cleanHashMap() ctr.cleanNonEqCondExecutor() } -func (ctr *container) resetNonEqCondExecutor() { - if ctr.nonEqCondExec != nil { - ctr.nonEqCondExec.ResetForNextQuery() - } -} - func (ctr *container) cleanNonEqCondExecutor() { if ctr.nonEqCondExec != nil { ctr.nonEqCondExec.Free() @@ -259,9 +313,12 @@ func (ctr *container) cleanBatch(proc *process.Process) { ctr.joinBats[i] = nil } } - if ctr.rightRowsMatched != nil { - ctr.rightRowsMatched = nil - } + ctr.freeRightRowsMatched(proc) +} + +func (ctr *container) freeRightRowsMatched(proc *process.Process) { + colexec.FreeAccountedBitmap(ctr.rightRowsMatched, proc.Mp()) + ctr.rightRowsMatched = nil } func (ctr *container) cleanBucketBatches(proc *process.Process) { @@ -273,6 +330,8 @@ func (ctr *container) cleanBucketBatches(proc *process.Process) { } func (ctr *container) cleanHashMap() { + hashmap.IteratorClearOwner(ctr.itr) + ctr.itr = nil if ctr.mp != nil { ctr.mp.Free() ctr.mp = nil @@ -289,21 +348,6 @@ func (ctr *container) cleanEqCondExecutors() { ctr.eqCondVecs = nil } -func (ctr *container) resetEqCondExecutors() { - for i := range ctr.eqCondExecs { - if ctr.eqCondExecs[i] != nil { - ctr.eqCondExecs[i].ResetForNextQuery() - } - } -} - -func (ctr *container) releaseProbeExpressionLease() { - if ctr.probeExpressionLease != nil { - ctr.probeExpressionLease.Release() - ctr.probeExpressionLease = nil - } -} - func (hashJoin *HashJoin) IsInner() bool { return hashJoin.JoinType == plan.Node_INNER } diff --git a/pkg/sql/colexec/indexbuild/build.go b/pkg/sql/colexec/indexbuild/build.go index 948f9e49e8221..949373430e169 100644 --- a/pkg/sql/colexec/indexbuild/build.go +++ b/pkg/sql/colexec/indexbuild/build.go @@ -18,6 +18,7 @@ import ( "bytes" "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -66,9 +67,38 @@ func (indexBuild *IndexBuild) Prepare(proc *process.Process) (err error) { ctr.runtimeFilterUsable = runtimefilter.ExactKeyEncoding( spec, declaredType) != keycodec.ExactRuntimeFilterUnsupported } + if ctr.runtimeFilterUsable && + (indexBuild.allocationAccount == nil || + indexBuild.runtimeFilterAllocation == nil) { + return mpool.ErrAllocationAccountInvalid + } return nil } +func (indexBuild *IndexBuild) newRuntimeFilterBatch( + typ types.Type, +) (*batch.Batch, error) { + if indexBuild.runtimeFilterAllocation == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + vec, err := vector.NewOffHeapVecWithTypeAndAllocation( + typ, + indexBuild.runtimeFilterAllocation, + ) + if err != nil { + return nil, err + } + buf := batch.NewOffHeapWithSize(1) + if err = buf.SetAllocationAccount( + indexBuild.runtimeFilterAllocation, + ); err != nil { + vec.Free(nil) + return nil, err + } + buf.SetVector(0, vec) + return buf, nil +} + func (indexBuild *IndexBuild) Call(proc *process.Process) (vm.CallResult, error) { analyzer := indexBuild.OpAnalyzer @@ -195,9 +225,11 @@ func (ctr *container) collectBuildBatches(indexBuild *IndexBuild, proc *process. // expanded into an equally large retained vector. if !inputVec.IsConstNull() && inputVec.Length() > 0 { if ctr.buf == nil { - ctr.buf = batch.NewOffHeapWithSize(1) - ctr.buf.Vecs[0] = vector.NewOffHeapVecWithType( + ctr.buf, err = indexBuild.newRuntimeFilterBatch( *inputVec.GetType()) + if err != nil { + return err + } } if err = ctr.buf.UnionOne(result.Batch, 0, proc.Mp()); err != nil { err = runtimefilter.MarkOptionalAllocationError(err) @@ -214,9 +246,11 @@ func (ctr *container) collectBuildBatches(indexBuild *IndexBuild, proc *process. // cardinality-bounded copy. Off-heap growth is tracked by the // process pool and can fail open to PASS instead of ending in // an unrecoverable Go-heap OOM. - ctr.buf = batch.NewOffHeapWithSize(1) - ctr.buf.Vecs[0] = vector.NewOffHeapVecWithType( + ctr.buf, err = indexBuild.newRuntimeFilterBatch( *inputVec.GetType()) + if err != nil { + return err + } } ctr.buf, err = ctr.buf.AppendWithCopy(proc.Ctx, proc.Mp(), result.Batch) if err != nil { @@ -286,7 +320,11 @@ func (ctr *container) handleRuntimeFilter(ap *IndexBuild, proc *process.Process) // Batch.Dup preserves a first-batch constant vector. Materialize only // its one distinct value: expanding every repeated row would waste // memory and AppendFixed cannot add signed-zero closure to a const vec. - flat := vector.NewOffHeapVecWithType(*vec.GetType()) + flat, err := vector.NewOffHeapVecWithTypeAndAllocation( + *vec.GetType(), ap.runtimeFilterAllocation) + if err != nil { + return err + } if !vec.IsConstNull() && vec.Length() > 0 { if err := flat.UnionOne(vec, 0, proc.Mp()); err != nil { flat.Free(proc.Mp()) @@ -316,7 +354,7 @@ func (ctr *container) handleRuntimeFilter(ap *IndexBuild, proc *process.Process) return nil } if encoding == keycodec.ExactRuntimeFilterFloatZeroClosed { - if err := runtimefilter.CloseFloatSignedZero(vec, proc.Mp(), nil); err != nil { + if err := runtimefilter.CloseFloatSignedZero(vec, proc.Mp()); err != nil { if ctr.fallbackRuntimeFilter(ap, proc, err) { return nil } @@ -335,15 +373,13 @@ func (ctr *container) handleRuntimeFilter(ap *IndexBuild, proc *process.Process) // NULLs are irrelevant for IN-filter: clear bitmap before sort. vec.GetNulls().Reset() vec.InplaceSort() - budget, err := proc.GetHashBuildBudget() - if err != nil { - if ctr.fallbackRuntimeFilter(ap, proc, err) { - return nil - } - ctr.abandonRuntimeFilter(proc) - return err - } - data, release, err := runtimefilter.MarshalExactFilterVector(vec, budget) + data, release, err := runtimefilter.MarshalExactFilterVector( + vec, + proc.Mp(), + ap.allocationAccount, + indexBuildAllocationOwner, + indexBuildAllocationSiteRuntimeFilterPayload, + ) if err != nil { if ctr.fallbackRuntimeFilter(ap, proc, err) { return nil diff --git a/pkg/sql/colexec/indexbuild/build_test.go b/pkg/sql/colexec/indexbuild/build_test.go index 122fb648ca0e5..4da04788d46b0 100644 --- a/pkg/sql/colexec/indexbuild/build_test.go +++ b/pkg/sql/colexec/indexbuild/build_test.go @@ -67,6 +67,22 @@ func indexBuildTestProcess(t *testing.T) *process.Process { return proc } +func prepareIndexBuild( + t *testing.T, + arg *IndexBuild, + proc *process.Process, +) { + t.Helper() + if arg.allocationAccount == nil { + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + } + require.NoError(t, arg.Prepare(proc)) +} + func indexBuildBatch(vec *vector.Vector, rows int) *batch.Batch { bat := batch.NewWithSize(1) bat.Vecs[0] = vec @@ -116,7 +132,7 @@ func executeIndexBuild( arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -149,7 +165,7 @@ func TestIndexBuildExactRuntimeFilterContract(t *testing.T) { child := colexec.NewMockOperator() arg := NewArgument() arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) arg.ctr.buf = indexBuildBatch(nil, 1) require.NotPanics(t, func() { require.NoError(t, arg.ctr.handleRuntimeFilter(arg, proc)) @@ -203,7 +219,7 @@ func TestIndexBuildExactRuntimeFilterContract(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -225,7 +241,7 @@ func TestIndexBuildExactRuntimeFilterContract(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -294,7 +310,7 @@ func TestIndexBuildFloatRuntimeFilterClosesConstSignedZero(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) result, err := vm.Exec(arg, proc) require.NoError(t, err) require.Equal(t, vm.ExecStop, result.Status) @@ -353,7 +369,7 @@ func TestIndexBuildRuntimeFilterCopyFailureFailsOpen(t *testing.T) { arg.RuntimeFilterSpec = spec arg.AppendChild(child) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) filler, err := limited.Alloc( int(limited.Cap()-limited.CurrNB()), true) require.NoError(t, err) @@ -413,7 +429,7 @@ func TestIndexBuildRuntimeFilterClosureFailureFailsOpen(t *testing.T) { proc.SetMessageBoard(message.NewMessageBoard()) arg := NewArgument() arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) arg.ctr.buf = batch.NewOffHeapWithSize(1) arg.ctr.buf.Vecs[0] = vector.NewOffHeapVecWithType(typ) @@ -460,19 +476,28 @@ func TestIndexBuildRuntimeFilterBudgetErrorPolicy(t *testing.T) { spec := indexBuildRawSpec(111, 16, typ) arg := NewArgument() arg.RuntimeFilterSpec = spec - require.NoError(t, arg.Prepare(proc)) + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController( + 2*generation.Cap(), generation) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + prepareIndexBuild(t, arg, proc) arg.ctr.buf = indexBuildBatch( testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()), 3, ) - generation, err := proc.GetHashBuildBudget() - require.NoError(t, err) - var held *process.HashBuildReservation + var filler []byte if test.closed { generation.Close() } else { - held, err = generation.Reserve(generation.Cap()) + remaining := generation.Cap() - generation.Used() + filler, err = proc.Mp().AllocAccounted( + int(remaining), account, 63, 255) require.NoError(t, err) } @@ -490,7 +515,8 @@ func TestIndexBuildRuntimeFilterBudgetErrorPolicy(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) require.Equal(t, int64(1), stats["IndexBuildRuntimeFilterBudgetFallbacks"]) - require.True(t, held.Release()) + proc.Mp().Free(filler) + generation.Close() } require.False(t, arg.ctr.runtimeFilterUsable) require.Nil(t, arg.ctr.buf) @@ -545,7 +571,7 @@ func TestIndexBuildCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { ) require.NoError(t, child.Prepare(proc)) - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) _, err := vm.Exec(arg, proc) require.ErrorIs(t, err, buildErr) @@ -571,7 +597,7 @@ func TestIndexBuildCallErrorUnblocksRuntimeFilterBeforeReset(t *testing.T) { require.True(t, arg.ctr.runtimeFilterDone) proc.GetMessageBoard().Reset() - require.NoError(t, arg.Prepare(proc)) + prepareIndexBuild(t, arg, proc) require.False(t, arg.ctr.runtimeFilterDone, "Prepare must open the terminal gate for the next generation") arg.finalizeBuildFailure(proc) diff --git a/pkg/sql/colexec/indexbuild/types.go b/pkg/sql/colexec/indexbuild/types.go index 2995f54323d5c..7eae66d34a174 100644 --- a/pkg/sql/colexec/indexbuild/types.go +++ b/pkg/sql/colexec/indexbuild/types.go @@ -15,8 +15,10 @@ package indexbuild import ( + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -31,6 +33,16 @@ const ( End ) +const indexBuildAllocationOwner mpool.AllocationOwner = 1 + +const ( + indexBuildAllocationSiteRuntimeFilterData mpool.AllocationSite = iota + 1 + indexBuildAllocationSiteRuntimeFilterArea + indexBuildAllocationSiteRuntimeFilterNulls + indexBuildAllocationSiteRuntimeFilterGrouping + indexBuildAllocationSiteRuntimeFilterPayload +) + type container struct { state int buf *batch.Batch @@ -39,11 +51,58 @@ type container struct { } type IndexBuild struct { - ctr container - RuntimeFilterSpec *plan.RuntimeFilterSpec + ctr container + RuntimeFilterSpec *plan.RuntimeFilterSpec + allocationAccount *mpool.AllocationAccount + runtimeFilterAllocation *vector.AllocationAccountSelection vm.OperatorBase } +func (indexBuild *IndexBuild) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil { + return mpool.ErrAllocationAccountInvalid + } + if indexBuild.allocationAccount != nil { + if indexBuild.allocationAccount == account { + return nil + } + return mpool.ErrAllocationAccountMismatch + } + selection, err := vector.NewAllocationAccountSelection( + account, + indexBuildAllocationOwner, + indexBuildAllocationSiteRuntimeFilterData, + indexBuildAllocationSiteRuntimeFilterArea, + indexBuildAllocationSiteRuntimeFilterNulls, + indexBuildAllocationSiteRuntimeFilterGrouping, + ) + if err != nil { + return err + } + indexBuild.allocationAccount = account + indexBuild.runtimeFilterAllocation = selection + return nil +} + +func (indexBuild *IndexBuild) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if indexBuild.allocationAccount == nil { + return nil + } + if indexBuild.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if indexBuild.ctr.buf != nil { + return mpool.ErrAllocationAccountInvariant + } + indexBuild.allocationAccount = nil + indexBuild.runtimeFilterAllocation = nil + return nil +} + func (indexBuild *IndexBuild) GetOperatorBase() *vm.OperatorBase { return &indexBuild.OperatorBase } @@ -92,7 +151,8 @@ func (indexBuild *IndexBuild) Reset(proc *process.Process, pipelineFailed bool, indexBuild.ctr.state = ReceiveBatch indexBuild.ctr.runtimeFilterUsable = false if indexBuild.ctr.buf != nil { - indexBuild.ctr.buf.CleanOnlyData() + indexBuild.ctr.buf.Clean(proc.Mp()) + indexBuild.ctr.buf = nil } } diff --git a/pkg/sql/colexec/intersect/intersect.go b/pkg/sql/colexec/intersect/intersect.go index 74ca37fdd4b0e..4767df195d840 100644 --- a/pkg/sql/colexec/intersect/intersect.go +++ b/pkg/sql/colexec/intersect/intersect.go @@ -172,7 +172,10 @@ func (intersect *Intersect) probeHashTable(proc *process.Process, analyzer proce copy(needInsert, resetsNeedInsert) insertcnt := 0 - vs, zs := itr.Find(i, n, input.Batch.Vecs) + vs, zs, err := itr.Find(i, n, input.Batch.Vecs) + if err != nil { + return false, err + } for j, v := range vs { diff --git a/pkg/sql/colexec/intersectall/intersectall.go b/pkg/sql/colexec/intersectall/intersectall.go index 2d283d4412e64..5c78cfcb7d808 100644 --- a/pkg/sql/colexec/intersectall/intersectall.go +++ b/pkg/sql/colexec/intersectall/intersectall.go @@ -194,7 +194,10 @@ func (intersectAll *IntersectAll) probe(proc *process.Process, analyzer process. copy(ctr.inserted[:n], ctr.resetInserted[:n]) cnt = 0 - vs, _ := itr.Find(i, n, input.Batch.Vecs) + vs, _, err := itr.Find(i, n, input.Batch.Vecs) + if err != nil { + return false, err + } for j, v := range vs { // not found diff --git a/pkg/sql/colexec/join_util.go b/pkg/sql/colexec/join_util.go index fcd63a9099f90..e61fdd3346196 100644 --- a/pkg/sql/colexec/join_util.go +++ b/pkg/sql/colexec/join_util.go @@ -60,22 +60,106 @@ func (bs *Batches) Reset() { // copy from input batch into batches // the batches structure hold data in fix size 8192 rows, and continue to append from next batch -// if error return , the batches will clean itself func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err error) { - defer func() { - if err != nil { - bs.Clean(proc.Mp()) + return bs.CopyIntoBatchesWithAllocation(src, proc, nil) +} + +// CopyIntoBatchesWithAllocation selects provenance for every retained vector +// destination. The Go descriptors are bounded by one Batch per 8,192 rows and +// one Vector pointer per input column; physical data, area, null, and grouping +// buffers are allocation-accounted and remain owned by the copied batches. +// +// The append is transactional. In particular, an allocation rejection while +// copying a later input must not destroy batches retained from earlier inputs: +// HashBuild needs those batches intact to recover by spilling them. Existing +// vectors are restored to logical append checkpoints on failure; successful +// capacity growth remains owned and reusable. This avoids repeatedly copying +// a partial 8,192-row tail as small input batches arrive. +func (bs *Batches) CopyIntoBatchesWithAllocation( + src *batch.Batch, + proc *process.Process, + selection *vector.AllocationAccountSelection, +) (err error) { + if len(bs.Buf) > 0 && + !vector.AllocationAccountSelectionsEqual( + bs.Buf[len(bs.Buf)-1].AllocationAccountSelection(), + selection, + ) { + return mpool.ErrAllocationAccountMismatch + } + + originalLen := len(bs.Buf) + originalMemSize := bs.MemSize + originalNil := bs.Buf == nil + var originalTail *batch.Batch + originalTailRows := 0 + var localTailCheckpoints [16]vector.AppendCheckpoint + var tailCheckpoints []vector.AppendCheckpoint + if originalLen > 0 && bs.Buf[originalLen-1].RowCount() != DefaultBatchSize { + originalTail = bs.Buf[originalLen-1] + originalTailRows = originalTail.RowCount() + if len(originalTail.Vecs) > len(localTailCheckpoints) { + tailCheckpoints = make([]vector.AppendCheckpoint, len(originalTail.Vecs)) + } else { + tailCheckpoints = localTailCheckpoints[:len(originalTail.Vecs)] } - }() + for i := range originalTail.Vecs { + tailCheckpoints[i] = originalTail.Vecs[i].MakeAppendCheckpoint() + } + } + if err = bs.copyIntoBatches(src, proc, selection); err == nil { + return nil + } + for i := originalLen; i < len(bs.Buf); i++ { + bs.Buf[i].Clean(proc.Mp()) + } + bs.Buf = bs.Buf[:originalLen] + if originalNil { + bs.Buf = nil + } + if originalTail != nil { + for i := range originalTail.Vecs { + originalTail.Vecs[i].RollbackAppend( + tailCheckpoints[i], + DefaultBatchSize-originalTailRows, + ) + } + originalTail.SetRowCount(originalTailRows) + } + bs.MemSize = originalMemSize + return err +} +func (bs *Batches) copyIntoBatches( + src *batch.Batch, + proc *process.Process, + selection *vector.AllocationAccountSelection, +) (err error) { if bs.Buf == nil { bs.Buf = make([]*batch.Batch, 0, 16) } + if len(bs.Buf) > 0 && + !vector.AllocationAccountSelectionsEqual( + bs.Buf[len(bs.Buf)-1].AllocationAccountSelection(), + selection, + ) { + return mpool.ErrAllocationAccountMismatch + } var tmp *batch.Batch if src.RowCount() == DefaultBatchSize { - tmp, err = src.Dup(proc.Mp()) + if selection == nil { + tmp, err = src.Dup(proc.Mp()) + } else { + tmp, err = proc.NewBatchFromSrcWithAllocation(src, 0, selection) + if err == nil { + err = src.CloneTo(tmp, proc.Mp()) + } + } if err != nil { + if tmp != nil { + tmp.Clean(proc.Mp()) + } return err } bs.MemSize += int64(tmp.Size()) @@ -96,12 +180,22 @@ func (bs *Batches) CopyIntoBatches(src *batch.Batch, proc *process.Process) (err lenBuf := len(bs.Buf) if lenBuf > 0 && bs.Buf[lenBuf-1].RowCount() != DefaultBatchSize { tmp = bs.Buf[lenBuf-1] + if !vector.AllocationAccountSelectionsEqual( + tmp.AllocationAccountSelection(), + selection, + ) { + return mpool.ErrAllocationAccountMismatch + } } else { preAllocSize := length - offset if preAllocSize > DefaultBatchSize { preAllocSize = DefaultBatchSize } - tmp, err = proc.NewBatchFromSrc(src, preAllocSize) + tmp, err = proc.NewBatchFromSrcWithAllocation( + src, + preAllocSize, + selection, + ) if err != nil { return err } @@ -123,45 +217,63 @@ func (bs *Batches) Shrink(ignoreRow *bitmap.Bitmap, proc *process.Process) error if ignoreRow.Count() == 0 { return nil } + if len(bs.Buf) == 0 || bs.Buf[0] == nil { + return mpool.ErrAllocationAccountInvalid + } ignoreRow.Negate() - count := int64(ignoreRow.Count()) - sels := make([]int32, 0, count) - itr := ignoreRow.Iterator() - for itr.HasNext() { - r := itr.Next() - sels = append(sels, int32(r)) + // Build the replacement privately and stream the active row IDs directly + // from the bitmap. The old implementation materialized one Go int32 per + // row and silently dropped the copied-batch allocation provenance. + count := ignoreRow.Count() + n := (count + DefaultBatchSize - 1) / DefaultBatchSize + if n == 0 { + n = 1 } - - n := (len(sels)-1)/DefaultBatchSize + 1 + selection := bs.Buf[0].AllocationAccountSelection() newBuf := make([]*batch.Batch, n) + cleanup := true + defer func() { + if cleanup { + for _, bat := range newBuf { + if bat != nil { + bat.Clean(proc.Mp()) + } + } + // Preserve the caller's ignore-row checkpoint on failure. + ignoreRow.Negate() + } + }() for i := range newBuf { - newBuf[i] = batch.NewOffHeapWithSize(len(bs.Buf[i].Vecs)) - for j, vec := range bs.Buf[0].Vecs { - newBuf[i].Vecs[j] = vector.NewOffHeapVecWithType(*vec.GetType()) + newBuf[i] = batch.NewOffHeapWithSize(len(bs.Buf[0].Vecs)) + if err := newBuf[i].SetAllocationAccount(selection); err != nil { + return err } - var newsels []int32 - if (i+1)*DefaultBatchSize <= len(sels) { - newsels = sels[i*DefaultBatchSize : (i+1)*DefaultBatchSize] - } else { - newsels = sels[i*DefaultBatchSize:] + for j, vec := range bs.Buf[0].Vecs { + newBuf[i].SetVector(int32(j), vector.NewOffHeapVecWithType(*vec.GetType())) } - for _, sel := range newsels { - idx1, idx2 := sel/DefaultBatchSize, sel%DefaultBatchSize - for j, vec := range bs.Buf[idx1].Vecs { - if err := newBuf[i].Vecs[j].UnionOne(vec, int64(idx2), proc.Mp()); err != nil { - for k := 0; k <= i; k++ { - newBuf[k].Clean(proc.Mp()) - } - return err - } + } + itr := ignoreRow.Iterator() + outRow := 0 + for itr.HasNext() { + sel := int(itr.Next()) + srcBatch, srcRow := sel/DefaultBatchSize, sel%DefaultBatchSize + dstBatch := outRow / DefaultBatchSize + for j, vec := range bs.Buf[srcBatch].Vecs { + if err := newBuf[dstBatch].Vecs[j].UnionOne(vec, int64(srcRow), proc.Mp()); err != nil { + return err } } - newBuf[i].SetRowCount(len(newsels)) + newBuf[dstBatch].AddRowCount(1) + outRow++ } bs.Clean(proc.Mp()) bs.Buf = newBuf + for _, bat := range newBuf { + bs.MemSize += int64(bat.Size()) + } + cleanup = false return nil } @@ -181,8 +293,21 @@ func appendToFixedSizeFromOffset(dst *batch.Batch, src *batch.Batch, offset int, if length+offset > src.RowCount() { length = src.RowCount() - offset } + var localCheckpoints [16]vector.AppendCheckpoint + var checkpoints []vector.AppendCheckpoint + if len(dst.Vecs) > len(localCheckpoints) { + checkpoints = make([]vector.AppendCheckpoint, len(dst.Vecs)) + } else { + checkpoints = localCheckpoints[:len(dst.Vecs)] + } + for i := range dst.Vecs { + checkpoints[i] = dst.Vecs[i].MakeAppendCheckpoint() + } for i := range dst.Vecs { if err = dst.Vecs[i].UnionBatch(src.Vecs[i], int64(offset), length, nil, proc.Mp()); err != nil { + for j := 0; j <= i; j++ { + dst.Vecs[j].RollbackAppend(checkpoints[j], length) + } return 0, err } dst.Vecs[i].SetSorted(false) diff --git a/pkg/sql/colexec/join_util_test.go b/pkg/sql/colexec/join_util_test.go index 8d4dadecd30af..a872de2d164a8 100644 --- a/pkg/sql/colexec/join_util_test.go +++ b/pkg/sql/colexec/join_util_test.go @@ -15,6 +15,7 @@ package colexec import ( + "sync/atomic" "testing" "github.com/stretchr/testify/require" @@ -22,9 +23,69 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/testutil" ) +type testAppendCapacityController struct { + limit atomic.Uint64 + used atomic.Uint64 +} + +func (c *testAppendCapacityController) AcquireAllocationCapacity(size uint64) error { + for { + used := c.used.Load() + limit := c.limit.Load() + if size > limit || used > limit-size { + return mpool.ErrAllocationAccountCapacity + } + if c.used.CompareAndSwap(used, used+size) { + return nil + } + } +} + +func (c *testAppendCapacityController) ReleaseAllocationCapacity(size uint64) { + for { + used := c.used.Load() + if size > used { + panic("test allocation capacity underflow") + } + if c.used.CompareAndSwap(used, used-size) { + return + } + } +} + +func BenchmarkCopyIntoBatchesPartialTail(b *testing.B) { + const rowsPerInput = 128 + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{ + types.T_int64.ToType(), + types.T_int64.ToType(), + types.T_int64.ToType(), + types.T_int64.ToType(), + }, + true, + rowsPerInput, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + b.ResetTimer() + for range b.N { + var batches Batches + for rows := 0; rows < DefaultBatchSize; rows += rowsPerInput { + if err := batches.CopyIntoBatches(input, proc); err != nil { + b.Fatal(err) + } + } + batches.Clean(proc.Mp()) + } +} + func TestBatches(t *testing.T) { var batches Batches proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) @@ -76,3 +137,183 @@ func TestBatches(t *testing.T) { batches.Clean(proc.Mp()) require.Equal(t, int64(0), proc.Mp().CurrNB()) } + +func TestBatchesShrinkPreservesAllocationAndRollback(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + input := testutil.NewBatch( + []types.Type{types.T_int32.ToType()}, + true, + DefaultBatchSize, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + measure := func(limit uint64, shrink bool) (uint64, error) { + registry, err := mpool.NewAllocationAccountRegistry(1, 32) + require.NoError(t, err) + account, err := registry.Open(limit) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, 1, 1, 2, 3, 4) + require.NoError(t, err) + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(input, proc, selection)) + before := account.Snapshot().Used + require.Positive(t, before) + var shrinkErr error + if shrink { + ignore := &bitmap.Bitmap{} + ignore.InitWithSize(DefaultBatchSize) + ignore.Add(0) + shrinkErr = batches.Shrink(ignore, proc) + if shrinkErr == nil { + require.Equal(t, DefaultBatchSize-1, batches.RowCount()) + for _, bat := range batches.Buf { + require.Same(t, selection, bat.AllocationAccountSelection()) + } + } else { + require.Equal(t, 1, ignore.Count(), "failed shrink restores ignore-row checkpoint") + require.Equal(t, DefaultBatchSize, batches.RowCount()) + require.Equal(t, before, account.Snapshot().Used) + } + } + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + return before, shrinkErr + } + + used, err := measure(1<<20, false) + require.NoError(t, err) + _, err = measure(used, true) + require.Error(t, err) + require.True(t, mpool.IsRetryableAllocationCapacity(err)) + _, err = measure(1<<20, true) + require.NoError(t, err) +} + +func TestCopyIntoBatchesAcceptsEquivalentAllocationSelection(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + first, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + second, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + require.NotSame(t, first, second) + input := testutil.NewBatch( + []types.Type{types.T_int64.ToType()}, + true, + 32, + proc.Mp(), + ) + defer input.Clean(proc.Mp()) + + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(input, proc, first)) + require.NoError(t, batches.CopyIntoBatchesWithAllocation(input, proc, second)) + require.Len(t, batches.Buf, 1) + require.Equal(t, 64, batches.RowCount()) + require.Same(t, first, batches.Buf[0].AllocationAccountSelection()) + + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestCopyIntoBatchesAllocationFailureRollsBackPartialTail(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + controller := &testAppendCapacityController{} + controller.limit.Store(1 << 60) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, controller) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + typesInTail := []types.Type{types.T_int64.ToType(), types.T_int64.ToType()} + initial := testutil.NewBatch(typesInTail, true, 1, proc.Mp()) + defer initial.Clean(proc.Mp()) + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(initial, proc, selection)) + require.Len(t, batches.Buf, 1) + require.Equal(t, 1, batches.RowCount()) + firstBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])...) + secondBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])...) + + const appendRows = 128 + oldCapacity := cap(batches.Buf[0].Vecs[0].GetData()) + requiredBytes := (batches.RowCount() + appendRows) * types.T_int64.ToType().TypeSize() + newCapacity, ok := mpool.GrowCapacity(int64(oldCapacity), int64(requiredBytes)) + require.True(t, ok) + require.Greater(t, newCapacity, int64(oldCapacity)) + controller.limit.Store(controller.used.Load() + uint64(newCapacity)) + + more := testutil.NewBatch(typesInTail, true, appendRows, proc.Mp()) + defer more.Clean(proc.Mp()) + err = batches.CopyIntoBatchesWithAllocation(more, proc, selection) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Len(t, batches.Buf, 1) + require.Equal(t, 1, batches.RowCount()) + require.Equal(t, firstBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])) + require.Equal(t, secondBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])) + + controller.limit.Store(1 << 60) + require.NoError(t, batches.CopyIntoBatchesWithAllocation(more, proc, selection)) + require.Equal(t, 1+appendRows, batches.RowCount()) + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, controller.used.Load()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestCopyIntoBatchesFailureRollsBackEarlierTailChunk(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + controller := &testAppendCapacityController{} + controller.limit.Store(1 << 60) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<60, controller) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + typesInTail := []types.Type{types.T_int64.ToType(), types.T_int64.ToType()} + initial := testutil.NewBatch(typesInTail, true, DefaultBatchSize-2, proc.Mp()) + defer initial.Clean(proc.Mp()) + var batches Batches + require.NoError(t, batches.CopyIntoBatchesWithAllocation(initial, proc, selection)) + firstBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])...) + secondBefore := append([]int64(nil), vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])...) + + // Filling the last two rows needs no growth. Reject creation of the next + // batch, after that first chunk has already succeeded. + controller.limit.Store(controller.used.Load()) + more := testutil.NewBatch(typesInTail, true, 128, proc.Mp()) + defer more.Clean(proc.Mp()) + err = batches.CopyIntoBatchesWithAllocation(more, proc, selection) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Len(t, batches.Buf, 1) + require.Equal(t, DefaultBatchSize-2, batches.RowCount()) + require.Equal(t, firstBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[0])) + require.Equal(t, secondBefore, vector.MustFixedColNoTypeCheck[int64](batches.Buf[0].Vecs[1])) + + controller.limit.Store(1 << 60) + require.NoError(t, batches.CopyIntoBatchesWithAllocation(more, proc, selection)) + require.Equal(t, DefaultBatchSize-2+128, batches.RowCount()) + batches.Clean(proc.Mp()) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, controller.used.Load()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} diff --git a/pkg/sql/colexec/limit/limit_test.go b/pkg/sql/colexec/limit/limit_test.go index 25cb7316e6355..36d1e439ddfe4 100644 --- a/pkg/sql/colexec/limit/limit_test.go +++ b/pkg/sql/colexec/limit/limit_test.go @@ -22,6 +22,8 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -184,6 +186,48 @@ func TestLimitDoesNotMutateInputBatch(t *testing.T) { require.Zero(t, proc.Mp().CurrNB()) } +func TestLimitResetReleasesCopiedAllocationAccountData(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(1) + input.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 32 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + } + input.SetRowCount(32) + + arg := NewArgument().WithLimit(plan2.MakePlan2Uint64ConstExprWithType(1)) + child := colexec.NewMockOperator().WithBatchs([]*batch.Batch{input}) + arg.AppendChild(child) + require.NoError(t, arg.Prepare(proc)) + result, err := arg.Call(proc) + require.NoError(t, err) + require.Equal(t, 1, result.Batch.RowCount()) + + // Pipeline cleanup resets children before parents. Simulate HashJoin + // releasing its result batch, then verify Limit releases its accounted copy. + input.Clean(proc.Mp()) + require.Positive(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + arg.Free(proc, false, nil) + child.Free(proc, false, nil) + arg.Release() + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func BenchmarkLimit(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []limitTestCase{ diff --git a/pkg/sql/colexec/limit/types.go b/pkg/sql/colexec/limit/types.go index 04e143cf327c7..a6b8f75441783 100644 --- a/pkg/sql/colexec/limit/types.go +++ b/pkg/sql/colexec/limit/types.go @@ -78,6 +78,12 @@ func (limit *Limit) Reset(proc *process.Process, pipelineFailed bool, err error) if limit.ctr.limitExecutor != nil { limit.ctr.limitExecutor.ResetForNextQuery() } + if limit.ctr.buf.HasAllocationAccount() { + // Prepared operators may reuse ordinary buffers across executions, but an + // accounted buffer belongs to exactly one execution generation. + limit.ctr.buf.Clean(proc.Mp()) + limit.ctr.buf = nil + } limit.ctr.seen = 0 } diff --git a/pkg/sql/colexec/loopjoin/join.go b/pkg/sql/colexec/loopjoin/join.go index e131626e5cc4a..b1e7a0fece2b7 100644 --- a/pkg/sql/colexec/loopjoin/join.go +++ b/pkg/sql/colexec/loopjoin/join.go @@ -17,13 +17,14 @@ package loopjoin import ( "bytes" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/util/resource" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" @@ -58,6 +59,9 @@ func (loopJoin *LoopJoin) OpType() vm.OpType { func (loopJoin *LoopJoin) Prepare(proc *process.Process) error { var err error + if loopJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if loopJoin.OpAnalyzer == nil { loopJoin.OpAnalyzer = process.NewAnalyzer(loopJoin.GetIdx(), loopJoin.IsFirst, loopJoin.IsLast, opName) } else { @@ -65,10 +69,16 @@ func (loopJoin *LoopJoin) Prepare(proc *process.Process) error { } if loopJoin.NonEqCond != nil && loopJoin.ctr.expr == nil { - loopJoin.ctr.expr, err = colexec.NewExpressionExecutor(proc, loopJoin.NonEqCond) + var execs []colexec.ExpressionExecutor + execs, err = hashbuild.NewExpressionExecutors( + proc, + []*plan.Expr{loopJoin.NonEqCond}, + loopJoin.allocationAccount, + ) if err != nil { return err } + loopJoin.ctr.expr = execs[0] } return err } @@ -90,7 +100,9 @@ func (loopJoin *LoopJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.state = End } else { if loopJoin.JoinType == plan.Node_OUTER && ctr.mp != nil { - ctr.initRightMatchedBitmap() + if err = ctr.initRightMatchedBitmap(loopJoin, proc); err != nil { + return result, err + } } ctr.state = Probe } @@ -119,7 +131,9 @@ func (loopJoin *LoopJoin) Call(proc *process.Process) (vm.CallResult, error) { ctr.batIdx = 0 } - loopJoin.resetResultBat() + if err = loopJoin.resetResultBat(); err != nil { + return result, err + } for i, rp := range loopJoin.ResultCols { if rp.Rel == 0 { ctr.resBat.Vecs[i].SetSorted(ctr.inBat.Vecs[rp.Pos].GetSorted()) @@ -424,7 +438,7 @@ func (ctr *container) probe(ap *LoopJoin, proc *process.Process, result *vm.Call return nil } -func (loopJoin *LoopJoin) resetResultBat() { +func (loopJoin *LoopJoin) resetResultBat() error { ctr := &loopJoin.ctr if ctr.resBat != nil { ctr.resBat.CleanOnlyData() @@ -433,58 +447,85 @@ func (loopJoin *LoopJoin) resetResultBat() { ctr.resBat.Vecs[i].SetLength(0) } } else { - ctr.resBat = batch.NewWithSize(len(loopJoin.ResultCols)) + ctr.resBat = batch.NewOffHeapWithSize(len(loopJoin.ResultCols)) for i, rp := range loopJoin.ResultCols { switch rp.Rel { case 0: - ctr.resBat.Vecs[i] = vector.NewVec(*ctr.inBat.Vecs[rp.Pos].GetType()) + var leftType types.Type + if ctr.inBat != nil && int(rp.Pos) < len(ctr.inBat.Vecs) { + leftType = *ctr.inBat.Vecs[rp.Pos].GetType() + } else if int(rp.Pos) < len(loopJoin.LeftTypes) { + leftType = loopJoin.LeftTypes[rp.Pos] + } else { + ctr.resBat.Clean(nil) + ctr.resBat = nil + return process.ErrHashBuildBudgetInvalid + } + ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(leftType) case 1: - ctr.resBat.Vecs[i] = vector.NewVec(loopJoin.RightTypes[rp.Pos]) + ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(loopJoin.RightTypes[rp.Pos]) case -1: - ctr.resBat.Vecs[i] = vector.NewVec(types.T_bool.ToType()) + ctr.resBat.Vecs[i] = vector.NewOffHeapVecWithType(types.T_bool.ToType()) } } + if err := ctr.resBat.SetAllocationAccount(loopJoin.resultAllocation); err != nil { + ctr.resBat.Clean(nil) + ctr.resBat = nil + return err + } } + return nil } // initRightMatchedBitmap allocates the per-build-row matched bitmap. -func (ctr *container) initRightMatchedBitmap() { +func (ctr *container) initRightMatchedBitmap( + ap *LoopJoin, + proc *process.Process, +) error { bats := ctr.mp.GetBatches() - ctr.rightBatchOffset = make([]uint64, len(bats)) + var err error + ctr.rightBatchOffset, err = mpool.MakeSliceAccounted[uint64]( + len(bats), + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + loopJoinAllocationSiteBatchOffsets, + ) + if err != nil { + return err + } var total uint64 for i, b := range bats { ctr.rightBatchOffset[i] = total total += uint64(b.RowCount()) } - ctr.rightRowsMatched = &bitmap.Bitmap{} - ctr.rightRowsMatched.InitWithSize(int64(total)) + if total > uint64(^uint64(0)>>1) { + ctr.cleanRightMatchState(proc) + return mpool.ErrAllocationAccountInvalid + } + ctr.rightRowsMatched, err = colexec.NewAccountedBitmap( + int64(total), + proc.Mp(), + ap.allocationAccount, + hashbuild.HashBuildAllocationOwner, + loopJoinAllocationSiteMatched, + ) + if err != nil { + ctr.cleanRightMatchState(proc) + return err + } + return nil } // finalize emits one batch worth of unmatched build rows with NULL probe // columns. Iterator is monotonic, so rightMatchedBat only advances. func (ctr *container) finalize(ap *LoopJoin, proc *process.Process, result *vm.CallResult) error { bats := ctr.mp.GetBatches() - if ctr.resBat == nil { - ctr.resBat = batch.NewWithSize(len(ap.ResultCols)) - for i, rp := range ap.ResultCols { - switch rp.Rel { - case 0: - ctr.resBat.Vecs[i] = vector.NewVec(ap.LeftTypes[rp.Pos]) - case 1: - ctr.resBat.Vecs[i] = vector.NewVec(ap.RightTypes[rp.Pos]) - default: - ctr.resBat.Vecs[i] = vector.NewVec(types.T_bool.ToType()) - } - } - } else { - ctr.resBat.CleanOnlyData() - for i := range ctr.resBat.Vecs { - ctr.resBat.Vecs[i].SetClass(vector.FLAT) - ctr.resBat.Vecs[i].SetLength(0) - } + if err := ap.resetResultBat(); err != nil { + return err } rowCnt := 0 diff --git a/pkg/sql/colexec/loopjoin/join_test.go b/pkg/sql/colexec/loopjoin/join_test.go index 78146989376a4..fece6216d6c01 100644 --- a/pkg/sql/colexec/loopjoin/join_test.go +++ b/pkg/sql/colexec/loopjoin/join_test.go @@ -52,6 +52,67 @@ type joinTestCase struct { resultBatch *batch.Batch } +type loopJoinTestAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installLoopJoinTestAllocation( + t testing.TB, + owners ...loopJoinTestAllocationOwner, +) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} + +func TestLoopJoinResultBatchUsesAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &LoopJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + account := installLoopJoinTestAllocation(t, arg) + require.NoError(t, arg.resetResultBat()) + require.Same(t, arg.resultAllocation, arg.ctr.resBat.Vecs[0].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp())) + used := account.Snapshot().Used + require.Positive(t, used) + require.NoError(t, arg.resetResultBat()) + require.Equal(t, used, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.resBat) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestLoopJoinResultBatchHonorsAllocationCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + arg := &LoopJoin{ + ResultCols: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetResultBat()) + err = vector.AppendFixed(arg.ctr.resBat.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + var ( tag int32 ) @@ -69,6 +130,36 @@ func TestString(t *testing.T) { } } +func TestResetRebuildsExpressionForNextAllocationGeneration(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + join := &LoopJoin{ + NonEqCond: &plan.Expr{ + Typ: plan.Type{Id: int32(types.T_bool)}, + Expr: &plan.Expr_Lit{Lit: &plan.Literal{ + Value: &plan.Literal_Bval{Bval: true}, + }}, + }, + } + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + + for range 2 { + account, openErr := registry.Open(1 << 20) + require.NoError(t, openErr) + require.NoError(t, join.SetAllocationAccount(account)) + require.NoError(t, join.Prepare(proc)) + require.NotNil(t, join.ctr.expr) + + join.Reset(proc, false, nil) + require.Nil(t, join.ctr.expr) + require.NoError(t, join.ClearAllocationAccount(account)) + terminal, _, terminalErr := registry.CompleteTerminal(account) + require.NoError(t, terminalErr) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.State) + } +} + func TestJoin(t *testing.T) { for _, tc := range makeTestCases(t) { @@ -296,6 +387,7 @@ func TestLoopJoinFinalizeResetsAfterPreviousEmptyProbe(t *testing.T) { }, }, } + installLoopJoinTestAllocation(t, join, build) resetChildrenWithBatch(join, makeInt32LoopJoinBatch(proc.Mp(), []int32{7})) resetHashBuildChildrenWithBatch(build, batch.EmptyBatch) @@ -398,6 +490,7 @@ func TestMarkJoinEmitsOneRowPerProbeRowAcrossBuildBatches(t *testing.T) { }, }, } + installLoopJoinTestAllocation(t, join, build) build.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{ makeInt32LoopJoinBatch(proc.Mp(), []int32{1}), makeInt32LoopJoinBatch(proc.Mp(), []int32{1}), @@ -487,6 +580,7 @@ func TestMarkJoinResumesAfterDefaultBatchSize(t *testing.T) { }, }, } + installLoopJoinTestAllocation(t, join, build) build.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{ makeInt32LoopJoinBatch(proc.Mp(), []int32{-1}), })) @@ -602,7 +696,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result resultBatch.Vecs[i] = bat.Vecs[rp[i].Pos] } tag++ - return joinTestCase{ + testCase := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -634,6 +728,8 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result }, resultBatch: resultBatch, } + installLoopJoinTestAllocation(t, testCase.arg, testCase.barg) + return testCase } func resetChildren(arg *LoopJoin, m *mpool.MPool) { diff --git a/pkg/sql/colexec/loopjoin/types.go b/pkg/sql/colexec/loopjoin/types.go index d18eba6acd383..791e5915272b0 100644 --- a/pkg/sql/colexec/loopjoin/types.go +++ b/pkg/sql/colexec/loopjoin/types.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -37,6 +38,18 @@ const ( End ) +const ( + loopJoinAllocationSiteMatched mpool.AllocationSite = iota + 92 + loopJoinAllocationSiteBatchOffsets +) + +const ( + loopJoinAllocationSiteResultData mpool.AllocationSite = iota + 106 + loopJoinAllocationSiteResultArea + loopJoinAllocationSiteResultNulls + loopJoinAllocationSiteResultGrouping +) + type container struct { state int probeIdx int @@ -60,18 +73,68 @@ type container struct { } type LoopJoin struct { - ctr container - LeftTypes []types.Type - RightTypes []types.Type - NonEqCond *plan.Expr - ResultCols []colexec.ResultPos - JoinMapTag int32 - JoinType plan.Node_JoinType - MarkPos int + ctr container + LeftTypes []types.Type + RightTypes []types.Type + NonEqCond *plan.Expr + ResultCols []colexec.ResultPos + JoinMapTag int32 + JoinType plan.Node_JoinType + MarkPos int + allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } +func (loopJoin *LoopJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if loopJoin.allocationAccount != nil && + loopJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if loopJoin.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + loopJoinAllocationSiteResultData, + loopJoinAllocationSiteResultArea, + loopJoinAllocationSiteResultNulls, + loopJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } + loopJoin.allocationAccount = account + loopJoin.resultAllocation = selection + return nil +} + +func (loopJoin *LoopJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if loopJoin.allocationAccount == nil { + return nil + } + if loopJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + ctr := &loopJoin.ctr + if ctr.mp != nil || ctr.expr != nil || ctr.rightRowsMatched != nil || + len(ctr.rightBatchOffset) != 0 || ctr.resBat != nil { + return mpool.ErrAllocationAccountInvariant + } + loopJoin.allocationAccount = nil + loopJoin.resultAllocation = nil + return nil +} + func (loopJoin *LoopJoin) GetOperatorBase() *vm.OperatorBase { return &loopJoin.OperatorBase } @@ -106,12 +169,18 @@ func (loopJoin *LoopJoin) Release() { func (loopJoin *LoopJoin) Reset(proc *process.Process, pipelineFailed bool, err error) { ctr := &loopJoin.ctr - ctr.resetNonEqCondExecutor() + // The executor owns allocations from this execution generation. Prepared + // statements must rebuild it after the next account is installed instead + // of carrying generation-bound storage across Reset. + ctr.cleanNonEqCondExecutor() ctr.cleanHashMap() + if ctr.resBat != nil { + ctr.resBat.Clean(proc.GetMPool()) + ctr.resBat = nil + } ctr.state = Build ctr.inBat = nil - ctr.rightRowsMatched = nil - ctr.rightBatchOffset = nil + ctr.cleanRightMatchState(proc) ctr.rightMatchedIter = nil ctr.rightMatchedBat = 0 } @@ -121,7 +190,19 @@ func (loopJoin *LoopJoin) Free(proc *process.Process, pipelineFailed bool, err e ctr.cleanBatch(proc.Mp()) ctr.cleanNonEqCondExecutor() + ctr.cleanRightMatchState(proc) + +} +func (ctr *container) cleanRightMatchState(proc *process.Process) { + colexec.FreeAccountedBitmap(ctr.rightRowsMatched, proc.Mp()) + ctr.rightRowsMatched = nil + if cap(ctr.rightBatchOffset) > 0 { + mpool.FreeSlice(proc.Mp(), ctr.rightBatchOffset) + } + ctr.rightBatchOffset = nil + ctr.rightMatchedIter = nil + ctr.rightMatchedBat = 0 } func (loopJoin *LoopJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { @@ -139,12 +220,6 @@ func (ctr *container) cleanBatch(mp *mpool.MPool) { } } -func (ctr *container) resetNonEqCondExecutor() { - if ctr.expr != nil { - ctr.expr.ResetForNextQuery() - } -} - func (ctr *container) cleanNonEqCondExecutor() { if ctr.expr != nil { ctr.expr.Free() diff --git a/pkg/sql/colexec/mergeorder/order.go b/pkg/sql/colexec/mergeorder/order.go index 9904728afc619..47a4653bc8c3c 100644 --- a/pkg/sql/colexec/mergeorder/order.go +++ b/pkg/sql/colexec/mergeorder/order.go @@ -278,12 +278,8 @@ func (ctr *container) removeInMemoryBatch(proc *process.Process, index int) erro if ctr.inMemoryHeap != nil { heap.Remove(ctr.inMemoryHeap, ctr.inMemoryHeapPos[index]) } - for i := range cols { - if batchContainsVector(bat, cols[i]) { - continue - } - cols[i].Free(proc.GetMPool()) - } + freeOrderColumns(proc.GetMPool(), bat, cols) + bat.Clean(proc.GetMPool()) ctr.batchList[index] = nil ctr.orderCols[index] = nil ctr.indexList[index] = -1 diff --git a/pkg/sql/colexec/mergeorder/order_test.go b/pkg/sql/colexec/mergeorder/order_test.go index 635b364a85f50..50549da223a1f 100644 --- a/pkg/sql/colexec/mergeorder/order_test.go +++ b/pkg/sql/colexec/mergeorder/order_test.go @@ -1080,6 +1080,74 @@ func TestMergeOrderResetAndOpType(t *testing.T) { arg.Free(proc, false, nil) } +func TestMergeOrderResetReleasesAccountedResult(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + result := batch.NewOffHeapWithSize(1) + result.Vecs[0] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, result.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(result.Vecs[0], int64(1), false, proc.Mp())) + result.SetRowCount(1) + require.Positive(t, account.Snapshot().Used) + + arg := &MergeOrder{} + arg.ctr.buf = result + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + arg.Free(proc, false, nil) + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + +func TestRemoveInMemoryBatchReleasesAccountedBatch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(2) + for i := range input.Vecs { + input.Vecs[i] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + } + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 17 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + require.NoError(t, vector.AppendFixed(input.Vecs[1], int64(i), false, proc.Mp())) + } + input.SetRowCount(17) + require.Positive(t, account.Snapshot().Used) + + ctr := container{ + batchList: []*batch.Batch{input}, + orderCols: [][]*vector.Vector{{input.Vecs[0]}}, + indexList: []int64{17}, + spillMemUsage: int64(input.Size()), + } + require.NoError(t, ctr.removeInMemoryBatch(proc, 0)) + require.Nil(t, ctr.batchList[0]) + require.Nil(t, ctr.orderCols[0]) + require.Zero(t, ctr.spillMemUsage) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func TestSpillHelperBranches(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer func() { diff --git a/pkg/sql/colexec/mergeorder/types.go b/pkg/sql/colexec/mergeorder/types.go index e9e8f6caac2e1..3f944e95a4790 100644 --- a/pkg/sql/colexec/mergeorder/types.go +++ b/pkg/sql/colexec/mergeorder/types.go @@ -174,7 +174,14 @@ func (mergeOrder *MergeOrder) Reset(proc *process.Process, pipelineFailed bool, } } if ctr.buf != nil { - ctr.buf.CleanOnlyData() + if ctr.buf.HasAllocationAccount() { + // The final merge batch may directly own a child execution's result. + // Accounted storage cannot survive that execution's Reset boundary. + ctr.buf.Clean(proc.Mp()) + ctr.buf = nil + } else { + ctr.buf.CleanOnlyData() + } } } @@ -208,15 +215,13 @@ func (mergeOrder *MergeOrder) cleanBatchAndCol(proc *process.Process) { mp := proc.Mp() ctr := &mergeOrder.ctr for i := range ctr.batchList { + if ctr.batchList[i] != nil && i < len(ctr.orderCols) && ctr.orderCols[i] != nil { + freeOrderColumns(mp, ctr.batchList[i], ctr.orderCols[i]) + } if ctr.batchList[i] != nil { ctr.batchList[i].Clean(mp) } } - for i := range ctr.orderCols { - if ctr.orderCols[i] != nil { - freeOrderColumns(mp, ctr.batchList[i], ctr.orderCols[i]) - } - } } func (ctr *container) cleanupSpill(proc *process.Process) { diff --git a/pkg/sql/colexec/multi_update/s3writer_delegate.go b/pkg/sql/colexec/multi_update/s3writer_delegate.go index 962e43fb001fe..91d6a534776fa 100644 --- a/pkg/sql/colexec/multi_update/s3writer_delegate.go +++ b/pkg/sql/colexec/multi_update/s3writer_delegate.go @@ -39,6 +39,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/deletion" plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + sqlutil "github.com/matrixorigin/matrixone/pkg/sql/util" "github.com/matrixorigin/matrixone/pkg/vm/engine/disttae" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/options" @@ -286,7 +287,9 @@ func (writer *s3WriterDelegate) append( if tableType == UpdateMainTable { if mainTableNullPkFilter { var checked *batch.Batch - if checked, err = projBat.Clone(mp, false); err != nil { + // This validation copy leaves any allocation-accounted join owner; + // it is short-lived and never published back into that owner. + if checked, err = sqlutil.CopyBatch(projBat, proc); err != nil { return } nulls := checked.Vecs[mainTablePkProjectIdx].GetNulls().GetBitmap().Clone() @@ -321,7 +324,9 @@ func (writer *s3WriterDelegate) append( // Clone because SelectColumns shares vectors, and ShrinkByMask // modifies in-place. var filtered *batch.Batch - if filtered, err = projBat.Clone(mp, false); err != nil { + // The sinker owns the filtered copy independently of the input + // pipeline, so cross the allocation ownership boundary explicitly. + if filtered, err = sqlutil.CopyBatch(projBat, proc); err != nil { return } nullIdx := writer.sortIndexes[i] diff --git a/pkg/sql/colexec/offset/offset_test.go b/pkg/sql/colexec/offset/offset_test.go index 4c13d023f9d7a..2df732ac35c92 100644 --- a/pkg/sql/colexec/offset/offset_test.go +++ b/pkg/sql/colexec/offset/offset_test.go @@ -23,6 +23,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -140,6 +141,46 @@ func TestOffset(t *testing.T) { } } +func TestOffsetResetReleasesCopiedAllocationAccountData(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(1) + input.SetVector(0, vector.NewOffHeapVecWithType(types.T_int64.ToType())) + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 32 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + } + input.SetRowCount(32) + + arg := NewArgument().WithOffset(plan2.MakePlan2Uint64ConstExprWithType(1)) + child := colexec.NewMockOperator().WithBatchs([]*batch.Batch{input}) + arg.AppendChild(child) + require.NoError(t, arg.Prepare(proc)) + result, err := arg.Call(proc) + require.NoError(t, err) + require.Equal(t, 31, result.Batch.RowCount()) + + input.Clean(proc.Mp()) + require.Positive(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.buf) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + arg.Free(proc, false, nil) + child.Free(proc, false, nil) + arg.Release() + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func BenchmarkOffset(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []offsetTestCase{ diff --git a/pkg/sql/colexec/offset/types.go b/pkg/sql/colexec/offset/types.go index c6f3ef638aafc..bd52640ce925f 100644 --- a/pkg/sql/colexec/offset/types.go +++ b/pkg/sql/colexec/offset/types.go @@ -79,7 +79,14 @@ func (offset *Offset) Reset(proc *process.Process, pipelineFailed bool, err erro offset.ctr.offsetExecutor.ResetForNextQuery() } if offset.ctr.buf != nil { - offset.ctr.buf.CleanOnlyData() + if offset.ctr.buf.HasAllocationAccount() { + // Do not carry an execution-scoped allocation selection into the next + // prepared execution. + offset.ctr.buf.Clean(proc.Mp()) + offset.ctr.buf = nil + } else { + offset.ctr.buf.CleanOnlyData() + } } offset.ctr.seen = 0 } diff --git a/pkg/sql/colexec/order/order_test.go b/pkg/sql/colexec/order/order_test.go index 66bdc0fbb8fad..4fe2e60929321 100644 --- a/pkg/sql/colexec/order/order_test.go +++ b/pkg/sql/colexec/order/order_test.go @@ -21,6 +21,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -94,6 +95,43 @@ func TestOrder(t *testing.T) { } } +func TestOrderResetReleasesPartiallyAccumulatedAccountedBatch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + input := batch.NewOffHeapWithSize(2) + for i := range input.Vecs { + input.Vecs[i] = vector.NewOffHeapVecWithType(types.T_int64.ToType()) + } + require.NoError(t, input.SetAllocationAccount(selection)) + for i := range 64 { + require.NoError(t, vector.AppendFixed(input.Vecs[0], int64(i), false, proc.Mp())) + require.NoError(t, vector.AppendFixed(input.Vecs[1], int64(i), false, proc.Mp())) + } + input.SetRowCount(64) + + arg := &Order{} + _, err = arg.ctr.appendBatch(proc, input) + require.NoError(t, err) + input.Clean(proc.Mp()) + require.Positive(t, account.Snapshot().Used) + + arg.Reset(proc, true, nil) + require.Nil(t, arg.ctr.batWaitForSort) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + + arg.Free(proc, true, nil) + proc.Free() + require.Zero(t, proc.Mp().CurrNB()) +} + func BenchmarkOrder(b *testing.B) { for i := 0; i < b.N; i++ { tcs := []orderTestCase{ diff --git a/pkg/sql/colexec/order/types.go b/pkg/sql/colexec/order/types.go index 6055dab28a94c..dbab07cd34823 100644 --- a/pkg/sql/colexec/order/types.go +++ b/pkg/sql/colexec/order/types.go @@ -84,7 +84,11 @@ type container struct { func (order *Order) Reset(proc *process.Process, pipelineFailed bool, err error) { ctr := &order.ctr if ctr.batWaitForSort != nil { - if ctr.batWaitForSort.RowCount() > colexec.DefaultBatchSize { + if ctr.batWaitForSort.HasAllocationAccount() || + ctr.batWaitForSort.RowCount() > colexec.DefaultBatchSize { + // A partially accumulated sort batch can survive when an upstream + // pipeline is stopped before sortAndSend transfers it to rbat. Its + // allocation account belongs to the completed execution generation. ctr.batWaitForSort.Clean(proc.Mp()) ctr.batWaitForSort = nil } else { diff --git a/pkg/sql/colexec/product/product.go b/pkg/sql/colexec/product/product.go index 4826a46fd005d..e6edd1bc8893b 100644 --- a/pkg/sql/colexec/product/product.go +++ b/pkg/sql/colexec/product/product.go @@ -17,6 +17,7 @@ package product import ( "bytes" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" @@ -38,6 +39,9 @@ func (product *Product) OpType() vm.OpType { } func (product *Product) Prepare(proc *process.Process) error { + if product.allocationAccount == nil || product.resultAllocation == nil { + return mpool.ErrAllocationAccountInvalid + } if product.OpAnalyzer == nil { product.OpAnalyzer = process.NewAnalyzer(product.GetIdx(), product.IsFirst, product.IsLast, "cross join") } else { @@ -94,20 +98,30 @@ func (product *Product) Call(proc *process.Process) (vm.CallResult, error) { continue } } - if ctr.bat == nil { + if ctr.mp == nil { ctr.inBat = nil continue } if ctr.rbat == nil { + buildBat := ctr.firstBuildBatch() + if buildBat == nil { + ctr.inBat = nil + continue + } ctr.rbat = batch.NewOffHeapWithSize(len(product.Result)) for i, rp := range product.Result { if rp.Rel == 0 { ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(*ctr.inBat.Vecs[rp.Pos].GetType()) } else { - ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(*ctr.bat.Vecs[rp.Pos].GetType()) + ctr.rbat.Vecs[i] = vector.NewOffHeapVecWithType(*buildBat.Vecs[rp.Pos].GetType()) } } + if err := ctr.rbat.SetAllocationAccount(product.resultAllocation); err != nil { + ctr.rbat.Clean(proc.Mp()) + ctr.rbat = nil + return result, err + } } else { ctr.rbat.CleanOnlyData() } @@ -137,48 +151,63 @@ func (product *Product) build(proc *process.Process, analyzer process.Analyzer) if mp == nil { return nil } - batches := mp.GetBatches() - - //maybe optimize this in the future - for i := range batches { - ctr.bat, err = ctr.bat.AppendWithCopy(proc.Ctx, proc.Mp(), batches[i]) - if err != nil { - return err - } - } - mp.Free() + ctr.mp = mp return nil } func (ctr *container) probe(ap *Product, proc *process.Process, result *vm.CallResult) error { count := ctr.inBat.RowCount() - count2 := ctr.bat.RowCount() - var i, j int - for j = ctr.probeIdx; j < count2; j++ { - for i = 0; i < count; i++ { - for k, rp := range ap.Result { - if rp.Rel == 0 { - if err := ctr.rbat.Vecs[k].UnionOne(ctr.inBat.Vecs[rp.Pos], int64(i), proc.Mp()); err != nil { - return err - } - } else { - if err := ctr.rbat.Vecs[k].UnionOne(ctr.bat.Vecs[rp.Pos], int64(j), proc.Mp()); err != nil { - return err + batches := ctr.mp.GetBatches() + for ctr.buildBatIdx < len(batches) { + buildBat := batches[ctr.buildBatIdx] + if buildBat == nil || buildBat.RowCount() == 0 { + ctr.buildBatIdx++ + ctr.buildRowIdx = 0 + continue + } + for row := ctr.buildRowIdx; row < buildBat.RowCount(); row++ { + for probeRow := 0; probeRow < count; probeRow++ { + for k, rp := range ap.Result { + if rp.Rel == 0 { + if err := ctr.rbat.Vecs[k].UnionOne(ctr.inBat.Vecs[rp.Pos], int64(probeRow), proc.Mp()); err != nil { + return err + } + } else { + if err := ctr.rbat.Vecs[k].UnionOne(buildBat.Vecs[rp.Pos], int64(row), proc.Mp()); err != nil { + return err + } } } } + ctr.rbat.AddRowCount(count) + ctr.buildRowIdx = row + 1 + if ctr.rbat.RowCount() >= colexec.DefaultBatchSize { + if ctr.buildRowIdx == buildBat.RowCount() { + ctr.buildBatIdx++ + ctr.buildRowIdx = 0 + } + result.Batch = ctr.rbat + return nil + } } - if ctr.rbat.Vecs[0].Length() >= colexec.DefaultBatchSize { - result.Batch = ctr.rbat - ctr.rbat.SetRowCount(ctr.rbat.Vecs[0].Length()) - ctr.probeIdx = j + 1 - return nil - } + ctr.buildBatIdx++ + ctr.buildRowIdx = 0 } - // ctr.rbat.AddRowCount(count * count2) - ctr.probeIdx = 0 - ctr.rbat.SetRowCount(ctr.rbat.Vecs[0].Length()) + ctr.buildBatIdx = 0 + ctr.buildRowIdx = 0 result.Batch = ctr.rbat ctr.inBat = nil return nil } + +func (ctr *container) firstBuildBatch() *batch.Batch { + if ctr.mp == nil { + return nil + } + for _, bat := range ctr.mp.GetBatches() { + if bat != nil && bat.RowCount() > 0 { + return bat + } + } + return nil +} diff --git a/pkg/sql/colexec/product/product_test.go b/pkg/sql/colexec/product/product_test.go index 77aedf8011c95..9f3c13a81f933 100644 --- a/pkg/sql/colexec/product/product_test.go +++ b/pkg/sql/colexec/product/product_test.go @@ -73,6 +73,12 @@ func TestPrepare(t *testing.T) { } } +func TestPrepareRequiresAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + require.ErrorIs(t, (&Product{}).Prepare(proc), mpool.ErrAllocationAccountInvalid) +} + func TestProduct(t *testing.T) { for _, tc := range makeTestCases(t) { @@ -98,6 +104,7 @@ func TestProduct(t *testing.T) { tc.arg.Reset(tc.proc, false, nil) tc.barg.Reset(tc.proc, false, nil) + require.Zero(t, tc.arg.allocationAccount.Snapshot().Used) resetChildren(tc.arg, tc.proc.Mp()) resetHashBuildChildren(tc.barg, tc.proc.Mp()) @@ -122,6 +129,7 @@ func TestProduct(t *testing.T) { tc.arg.Reset(tc.proc, false, nil) tc.barg.Reset(tc.proc, false, nil) + require.Zero(t, tc.arg.allocationAccount.Snapshot().Used) tc.arg.Free(tc.proc, false, nil) tc.barg.Free(tc.proc, false, nil) @@ -130,6 +138,51 @@ func TestProduct(t *testing.T) { } } +func TestProductConsumesMultipleBuildBatchesWithoutCopy(t *testing.T) { + tc := newTestCase( + t, + []bool{false}, + []types.Type{types.T_int32.ToType()}, + []colexec.ResultPos{ + colexec.NewResultPos(0, 0), + colexec.NewResultPos(1, 0), + }, + ) + probe := colexec.MakeMockBatchs(tc.proc.Mp()) + build1 := colexec.MakeMockBatchs(tc.proc.Mp()) + build2 := colexec.MakeMockBatchs(tc.proc.Mp()) + tc.arg.Children = nil + tc.arg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{probe})) + tc.barg.Children = nil + tc.barg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{build1, build2})) + + require.NoError(t, tc.arg.Prepare(tc.proc)) + require.NoError(t, tc.barg.Prepare(tc.proc)) + _, err := vm.Exec(tc.barg, tc.proc) + require.NoError(t, err) + wantRows := probe.RowCount() * (build1.RowCount() + build2.RowCount()) + rows := 0 + for { + result, err := vm.Exec(tc.arg, tc.proc) + require.NoError(t, err) + if result.Batch != nil { + rows += result.Batch.RowCount() + } + if result.Status == vm.ExecStop { + break + } + } + require.Equal(t, wantRows, rows) + + tc.arg.Reset(tc.proc, false, nil) + tc.barg.Reset(tc.proc, false, nil) + require.Zero(t, tc.arg.allocationAccount.Snapshot().Used) + tc.arg.Free(tc.proc, false, nil) + tc.barg.Free(tc.proc, false, nil) + tc.proc.Free() + require.Zero(t, tc.proc.Mp().CurrNB()) +} + /* func BenchmarkProduct(b *testing.B) { for i := 0; i < b.N; i++ { @@ -170,7 +223,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result resultBatch.Vecs[i] = vector.NewVec(*bat.Vecs[rp[i].Pos].GetType()) } tag++ - return productTestCase{ + tc := productTestCase{ types: ts, flgs: flgs, proc: proc, @@ -200,6 +253,13 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []colexec.Result }, resultBatch: resultBatch, } + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, tc.arg.SetAllocationAccount(account)) + require.NoError(t, tc.barg.SetAllocationAccount(account)) + return tc } func resetChildren(arg *Product, m *mpool.MPool) { bat := colexec.MakeMockBatchs(m) diff --git a/pkg/sql/colexec/product/types.go b/pkg/sql/colexec/product/types.go index 733b388af223b..e548aba211337 100644 --- a/pkg/sql/colexec/product/types.go +++ b/pkg/sql/colexec/product/types.go @@ -18,8 +18,11 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" ) @@ -32,19 +35,30 @@ const ( ) type container struct { - state int - probeIdx int - bat *batch.Batch - rbat *batch.Batch - inBat *batch.Batch + state int + buildBatIdx int + buildRowIdx int + rbat *batch.Batch + inBat *batch.Batch + mp *message.JoinMap } +const ( + productAllocationSiteResultData mpool.AllocationSite = iota + 94 + productAllocationSiteResultArea + productAllocationSiteResultNulls + productAllocationSiteResultGrouping +) + type Product struct { ctr container Result []colexec.ResultPos IsShuffle bool JoinMapTag int32 + allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection + vm.OperatorBase } @@ -52,6 +66,52 @@ func (product *Product) GetOperatorBase() *vm.OperatorBase { return &product.OperatorBase } +func (product *Product) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if product.allocationAccount != nil && + product.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if product.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + productAllocationSiteResultData, + productAllocationSiteResultArea, + productAllocationSiteResultNulls, + productAllocationSiteResultGrouping, + ) + if err != nil { + return err + } + product.allocationAccount = account + product.resultAllocation = selection + return nil +} + +func (product *Product) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if product.allocationAccount == nil { + return nil + } + if product.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if product.ctr.mp != nil || product.ctr.rbat != nil { + return mpool.ErrAllocationAccountInvariant + } + product.allocationAccount = nil + product.resultAllocation = nil + return nil +} + func init() { reuse.CreatePool[Product]( func() *Product { @@ -80,15 +140,8 @@ func (product *Product) Release() { } func (product *Product) Reset(proc *process.Process, pipelineFailed bool, err error) { - if product.ctr.bat != nil { - product.ctr.bat.CleanOnlyData() - } - if product.ctr.rbat != nil { - product.ctr.rbat.CleanOnlyData() - } - product.ctr.inBat = nil + product.ctr.cleanBatch(proc.Mp()) product.ctr.state = Build - product.ctr.probeIdx = 0 } func (product *Product) Free(proc *process.Process, pipelineFailed bool, err error) { @@ -100,13 +153,15 @@ func (product *Product) ExecProjection(proc *process.Process, input *batch.Batch } func (ctr *container) cleanBatch(mp *mpool.MPool) { - if ctr.bat != nil { - ctr.bat.Clean(mp) - ctr.bat = nil - } if ctr.rbat != nil { ctr.rbat.Clean(mp) ctr.rbat = nil } + if ctr.mp != nil { + ctr.mp.Free() + ctr.mp = nil + } ctr.inBat = nil + ctr.buildBatIdx = 0 + ctr.buildRowIdx = 0 } diff --git a/pkg/sql/colexec/productl2/joinmap_account_test.go b/pkg/sql/colexec/productl2/joinmap_account_test.go new file mode 100644 index 0000000000000..478f732ab9482 --- /dev/null +++ b/pkg/sql/colexec/productl2/joinmap_account_test.go @@ -0,0 +1,97 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package productl2 + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vectorindex/metric" + "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/stretchr/testify/require" +) + +func TestProductL2ReleasesProducerAccountedJoinMap(t *testing.T) { + mp := mpool.MustNewZero() + proc := testutil.NewProcessWithMPool(t, "", mp) + proc.SetMessageBoard(message.NewMessageBoard()) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<12) + require.NoError(t, err) + account, err := registry.Open(1 << 30) + require.NoError(t, err) + + arrayType := types.T_array_float32.ToType() + arrayType.Width = 2 + build := batch.NewWithSize(2) + build.Vecs[0] = vector.NewVec(arrayType) + require.NoError(t, vector.AppendArrayList( + build.Vecs[0], [][]float32{{0, 0}, {10, 10}}, nil, mp, + )) + build.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + build.Vecs[1], []int64{10, 20}, nil, mp, + )) + build.SetRowCount(2) + + probe := batch.NewWithSize(1) + probe.Vecs[0] = vector.NewVec(arrayType) + require.NoError(t, vector.AppendArrayList( + probe.Vecs[0], [][]float32{{1, 1}}, nil, mp, + )) + probe.SetRowCount(1) + + const tag = int32(7001) + producer := &hashbuild.HashBuild{ + NeedBatches: true, + JoinMapTag: tag, + JoinMapRefCnt: 1, + } + producer.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{build})) + consumer := &Productl2{ + Result: []colexec.ResultPos{colexec.NewResultPos(1, 1)}, + OnExpr: onExprWithProbeCol(0), + JoinMapTag: tag, + VectorOpType: metric.OpType_L2Distance, + } + consumer.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{probe})) + require.NoError(t, producer.SetAllocationAccount(account)) + require.NoError(t, producer.Prepare(proc)) + require.NoError(t, consumer.Prepare(proc)) + + _, err = vm.Exec(producer, proc) + require.NoError(t, err) + result, err := vm.Exec(consumer, proc) + require.NoError(t, err) + require.NotNil(t, result.Batch) + require.Equal(t, []int64{10}, vector.MustFixedColNoTypeCheck[int64](result.Batch.Vecs[0])) + + consumer.Reset(proc, false, nil) + producer.Reset(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, producer.ClearAllocationAccount(account)) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + consumer.Free(proc, false, nil) + producer.Free(proc, false, nil) + proc.Free() + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/sql/colexec/productl2/product_l2.go b/pkg/sql/colexec/productl2/product_l2.go index 468d0bb76cb2b..e94a0a4648d46 100644 --- a/pkg/sql/colexec/productl2/product_l2.go +++ b/pkg/sql/colexec/productl2/product_l2.go @@ -186,8 +186,16 @@ func (productl2 *Productl2) build(proc *process.Process, analyzer process.Analyz return nil } batches := mp.GetBatches() - //maybe optimize this in the future + // ProductL2 index/scratch is outside the first HashBuild accounting domain. + // Create an explicit unaccounted destination instead of letting a nil + // AppendWithCopy clone inherit the producer's allocation selection. for i := range batches { + if ctr.bat == nil { + ctr.bat = batch.NewOffHeapWithSize(len(batches[i].Vecs)) + for j, source := range batches[i].Vecs { + ctr.bat.Vecs[j] = vector.NewOffHeapVecWithType(*source.GetType()) + } + } ctr.bat, err = ctr.bat.AppendWithCopy(proc.Ctx, proc.Mp(), batches[i]) if err != nil { return err diff --git a/pkg/sql/colexec/receiver_operator.go b/pkg/sql/colexec/receiver_operator.go deleted file mode 100644 index 4f9cab8451f2b..0000000000000 --- a/pkg/sql/colexec/receiver_operator.go +++ /dev/null @@ -1,32 +0,0 @@ -// Copyright 2023 Matrix Origin -// -// Licensed 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. - -package colexec - -import ( - "context" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" -) - -func ReceiveBitmapFromChannel(usr context.Context, ch chan *bitmap.Bitmap) *bitmap.Bitmap { - select { - case <-usr.Done(): - return nil - case bm, ok := <-ch: - if !ok { - return nil - } - return bm - } -} diff --git a/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go b/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go new file mode 100644 index 0000000000000..08768aaa6b6f6 --- /dev/null +++ b/pkg/sql/colexec/rightdedupjoin/allocation_test_helpers_test.go @@ -0,0 +1,84 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package rightdedupjoin + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/stretchr/testify/require" +) + +type testAllocationOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error +} + +func installTestAllocation(t testing.TB, owners ...testAllocationOwner) *mpool.AllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + for _, owner := range owners { + require.NoError(t, owner.SetAllocationAccount(account)) + } + return account +} + +func TestRightDedupJoinResultBatchUsesAllocationAccount(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + arg := &RightDedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + account := installTestAllocation(t, arg) + require.NoError(t, arg.resetResultBatch()) + require.Same(t, arg.resultAllocation, arg.ctr.resultBatch.Vecs[0].AllocationAccountSelection()) + require.NoError(t, vector.AppendFixed(arg.ctr.resultBatch.Vecs[0], int64(1), false, proc.Mp())) + used := account.Snapshot().Used + require.Positive(t, used) + require.NoError(t, arg.resetResultBatch()) + require.Equal(t, used, account.Snapshot().Used) + + arg.Reset(proc, false, nil) + require.Nil(t, arg.ctr.resultBatch) + require.Zero(t, account.Snapshot().Used) + require.NoError(t, arg.ClearAllocationAccount(account)) +} + +func TestRightDedupJoinResultBatchHonorsAllocationCapacity(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1) + require.NoError(t, err) + arg := &RightDedupJoin{ + Result: []colexec.ResultPos{{Rel: 0, Pos: 0}}, + LeftTypes: []types.Type{types.T_int64.ToType()}, + } + require.NoError(t, arg.SetAllocationAccount(account)) + require.NoError(t, arg.resetResultBatch()) + err = vector.AppendFixed(arg.ctr.resultBatch.Vecs[0], int64(1), false, proc.Mp()) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Zero(t, account.Snapshot().Used) + arg.Reset(proc, false, nil) + require.NoError(t, arg.ClearAllocationAccount(account)) +} diff --git a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go b/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go deleted file mode 100644 index 971f251b0991f..0000000000000 --- a/pkg/sql/colexec/rightdedupjoin/expression_memory_test.go +++ /dev/null @@ -1,64 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// Licensed 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. - -package rightdedupjoin - -import ( - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" - "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - "github.com/matrixorigin/matrixone/pkg/testutil" - "github.com/matrixorigin/matrixone/pkg/vm/process" - "github.com/stretchr/testify/require" -) - -func TestRightDedupJoinResetReleasesProbeExpressionLease(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - expr := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Lit{Lit: &plan.Literal{ - Value: &plan.Literal_I32Val{I32Val: 1}, - }}, - } - executors, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{expr}) - require.NoError(t, err) - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - lease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{expr}, executors, false) - require.NoError(t, err) - - arg := &RightDedupJoin{} - arg.ctr.evecs = []evalVector{{executor: executors[0]}} - arg.ctr.vecs = make([]*vector.Vector, len(executors)) - arg.ctr.probeExpressionLease = lease - input := batch.NewWithSize(0) - input.SetRowCount(4) - require.NoError(t, arg.ctr.evalJoinConditionBudgeted(input, proc)) - require.Positive(t, generation.Used()) - - arg.Reset(proc, false, nil) - require.Zero(t, generation.Used()) - require.Nil(t, arg.ctr.evecs) - require.Nil(t, arg.ctr.vecs) - require.Nil(t, arg.ctr.probeExpressionLease) -} diff --git a/pkg/sql/colexec/rightdedupjoin/join.go b/pkg/sql/colexec/rightdedupjoin/join.go index 1289d9c6dee4b..27b9f3225cc61 100644 --- a/pkg/sql/colexec/rightdedupjoin/join.go +++ b/pkg/sql/colexec/rightdedupjoin/join.go @@ -19,9 +19,9 @@ import ( "strings" "github.com/matrixorigin/matrixone/pkg/catalog" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -47,6 +47,9 @@ func (rightDedupJoin *RightDedupJoin) OpType() vm.OpType { } func (rightDedupJoin *RightDedupJoin) Prepare(proc *process.Process) (err error) { + if rightDedupJoin.allocationAccount == nil { + return mpool.ErrAllocationAccountInvalid + } if rightDedupJoin.OpAnalyzer == nil { rightDedupJoin.OpAnalyzer = process.NewAnalyzer(rightDedupJoin.GetIdx(), rightDedupJoin.IsFirst, rightDedupJoin.IsLast, "dedup join") } else { @@ -58,13 +61,21 @@ func (rightDedupJoin *RightDedupJoin) Prepare(proc *process.Process) (err error) newUpdateExecs := len(rightDedupJoin.ctr.exprExecs) == 0 && len(rightDedupJoin.UpdateColExprList) > 0 var evalExecs, updateExecs []colexec.ExpressionExecutor if newEvalVectors { - evalExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, rightDedupJoin.Conditions[0]) + evalExecs, err = hashbuild.NewExpressionExecutors( + proc, + rightDedupJoin.Conditions[0], + rightDedupJoin.allocationAccount, + ) if err != nil { return err } } if newUpdateExecs { - updateExecs, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, rightDedupJoin.UpdateColExprList) + updateExecs, err = hashbuild.NewExpressionExecutors( + proc, + rightDedupJoin.UpdateColExprList, + rightDedupJoin.allocationAccount, + ) if err != nil { for _, exec := range evalExecs { exec.Free() @@ -147,7 +158,7 @@ func (rightDedupJoin *RightDedupJoin) Call(proc *process.Process) (vm.CallResult if ctr.spillEngine != nil { // Clear previous bucket state before advancing. ctr.cleanHashMap() - ctr.matched = nil + ctr.cleanBitmap(proc) ctr.groupCount = 0 ctr.buildGroupCount = 0 var initErr error @@ -159,8 +170,13 @@ func (rightDedupJoin *RightDedupJoin) Call(proc *process.Process) (vm.CallResult ctr.groupCount = jm.GetGroupCount() ctr.buildGroupCount = ctr.groupCount if !proc.GetTxnOperator().Txn().IsPessimistic() && ctr.buildGroupCount > 0 { - ctr.matched = &bitmap.Bitmap{} - ctr.matched.InitWithSize(int64(ctr.buildGroupCount)) + ctr.matched, initErr = colexec.NewAccountedBitmap( + int64(ctr.buildGroupCount), + proc.Mp(), + rightDedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + rightDedupJoinAllocationSiteMatched, + ) } case spillutil.BucketEmptyBuild: ctr.mp, initErr = rightDedupJoin.newEmptyJoinMap(proc) @@ -212,35 +228,29 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro if takeErr != nil { return takeErr } - probeExecutors := make([]colexec.ExpressionExecutor, len(ctr.evecs)) - for i := range ctr.evecs { - probeExecutors[i] = ctr.evecs[i].executor - } - probeExpressionLease, leaseErr := hashbuild.NewExpressionMemoryLease( - budget, rightDedupJoin.Conditions[0], probeExecutors, false) - if leaseErr != nil { + if rightDedupJoin.allocationAccount == nil { _ = payload.Close() ctr.mp.Free() ctr.mp = nil - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - return leaseErr + return mpool.ErrAllocationAccountInvalid } - ctr.probeExpressionLease = probeExpressionLease - engine := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ + engine, engineErr := spillutil.NewSpillEngine(spillutil.SpillEngineConfig{ BuildKeyExprs: rightDedupJoin.Conditions[1], ProbeKeyExprs: rightDedupJoin.Conditions[0], SpillThreshold: ctr.spillThreshold, NeedsProbeForEmptyBuild: true, MergeProbeBatches: true, Budget: budget, - ProbeExpressionLease: probeExpressionLease, - }) - if len(payload.Files) > 0 { - engine.InitFromSpilledFiles(payload.Files) - } else { - engine.InitFromSpilledMap(payload.LegacyFds) + }, rightDedupJoin.allocationAccount, hashbuild.HashBuildAllocationOwner) + if engineErr != nil { + _ = payload.Close() + ctr.mp.Free() + ctr.mp = nil + ctr.cleanEvalVectors() + return engineErr } + engine.InitFromSpilledFiles(payload.Files) + ctr.spillEngine = engine if err := engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { input, err := vm.ChildrenCall(rightDedupJoin.GetChildren(0), proc, analyzer) @@ -257,18 +267,27 @@ func (rightDedupJoin *RightDedupJoin) build(analyzer process.Analyzer, proc *pro ctr.mp.Free() ctr.mp = nil engine.Cleanup(proc) + ctr.spillEngine = nil return err } ctr.mp.Free() - ctr.spillEngine = engine ctr.mp = nil return } ctr.groupCount = ctr.mp.GetGroupCount() ctr.buildGroupCount = ctr.groupCount - if !proc.GetTxnOperator().Txn().IsPessimistic() { - ctr.matched = &bitmap.Bitmap{} - ctr.matched.InitWithSize(int64(ctr.buildGroupCount)) + if !proc.GetTxnOperator().Txn().IsPessimistic() && + ctr.buildGroupCount > 0 { + ctr.matched, err = colexec.NewAccountedBitmap( + int64(ctr.buildGroupCount), + proc.Mp(), + rightDedupJoin.allocationAccount, + hashbuild.HashBuildAllocationOwner, + rightDedupJoinAllocationSiteMatched, + ) + if err != nil { + return err + } } } @@ -286,15 +305,18 @@ func (rightDedupJoin *RightDedupJoin) newEmptyJoinMap(proc *process.Process) (*m keyWidth += width } - budget, err := proc.GetHashBuildBudget() - if err != nil { - return nil, err + if rightDedupJoin.allocationAccount == nil { + return nil, mpool.ErrAllocationAccountInvalid } - return hashbuild.NewBudgetedEmptyJoinMap(keyWidth, budget, proc.Mp()) + return hashbuild.NewAccountedEmptyJoinMap( + keyWidth, + rightDedupJoin.allocationAccount, + proc.Mp(), + ) } func (ctr *container) probe(bat *batch.Batch, ap *RightDedupJoin, proc *process.Process, analyzer process.Analyzer, result *vm.CallResult) error { - err := ctr.evalJoinConditionBudgeted(bat, proc) + err := ctr.evalJoinCondition(bat, proc) if err != nil { return err } @@ -379,16 +401,8 @@ func (ctr *container) probe(bat *batch.Batch, ap *RightDedupJoin, proc *process. } } - ctr.resetResultBatch() - if ctr.resultBatch == nil { - ctr.resultBatch = batch.NewOffHeapWithSize(len(ap.Result)) - for i, rp := range ap.Result { - if rp.Rel == 0 { - ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType(ap.LeftTypes[rp.Pos]) - } else { - ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType(ap.RightTypes[rp.Pos]) - } - } + if err := ap.resetResultBatch(); err != nil { + return err } for i, rp := range ap.Result { @@ -419,14 +433,3 @@ func (ctr *container) evalJoinCondition(bat *batch.Batch, proc *process.Process) } return nil } - -func (ctr *container) evalJoinConditionBudgeted(bat *batch.Batch, proc *process.Process) error { - if ctr.probeExpressionLease == nil { - return ctr.evalJoinCondition(bat, proc) - } - return ctr.probeExpressionLease.Eval(proc, []*batch.Batch{bat}, bat.RowCount(), func(i int, vec *vector.Vector) error { - ctr.vecs[i] = vec - ctr.evecs[i].vec = vec - return nil - }) -} diff --git a/pkg/sql/colexec/rightdedupjoin/join_test.go b/pkg/sql/colexec/rightdedupjoin/join_test.go index 904c953db2718..19702a8f41617 100644 --- a/pkg/sql/colexec/rightdedupjoin/join_test.go +++ b/pkg/sql/colexec/rightdedupjoin/join_test.go @@ -18,7 +18,6 @@ import ( "bytes" "context" "fmt" - "os" "testing" "github.com/golang/mock/gomock" @@ -104,6 +103,7 @@ func runRightDedupCase(t *testing.T, buildVals, probeVals []int32, pessimistic, JoinMapTag: curTag, } arg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{probeBat})) + installTestAllocation(t, arg, buildArg) require.NoError(t, buildArg.Prepare(proc)) require.NoError(t, arg.Prepare(proc)) @@ -154,6 +154,13 @@ func TestRightDedupDuplicateTracking(t *testing.T) { func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBatches bool) { proc, ctrl := newRightDedupTestProcess(t, pessimistic) defer ctrl.Finish() + budget := process.MustNewHashBuildBudget(64<<20, 64<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + account, err := registry.OpenWithController(64<<20, generation) + require.NoError(t, err) typ := types.T_int32.ToType() tag++ curTag := tag @@ -171,6 +178,7 @@ func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBa JoinMapTag: curTag, SpillThreshold: 1, } + require.NoError(t, arg.SetAllocationAccount(account)) probeValues := [][]int32{{1}, {2}} if duplicateAcrossBatches { @@ -188,14 +196,14 @@ func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBa jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) jm.IncRef(1) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ - LegacyFds: make([]*os.File, spillutil.SpillNumBuckets), + Files: make([]*message.SpillFile, spillutil.SpillNumBuckets), + BudgetRef: generation, })) message.SendMessage(message.JoinMapMsg{ - JoinMapPtr: jm, + Result: message.NewJoinMapResult(jm), IsShuffle: true, ShuffleIdx: 0, Tag: curTag, - Spilled: true, }, proc.GetMessageBoard()) require.NoError(t, arg.Prepare(proc)) @@ -220,6 +228,9 @@ func runRightDedupSpilledEmptyBuild(t *testing.T, pessimistic, duplicateAcrossBa } arg.Free(proc, false, nil) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) proc.Free() require.Equal(t, int64(0), proc.Mp().CurrNB()) } @@ -244,6 +255,7 @@ func TestRightDedupResetAndPrepareRetry(t *testing.T) { Conditions: [][]*plan.Expr{{valid}, {valid}}, UpdateColExprList: []*plan.Expr{valid, invalid}, } + installTestAllocation(t, arg) require.Error(t, arg.Prepare(proc)) require.Nil(t, arg.ctr.vecs) @@ -268,6 +280,7 @@ func TestRightDedupEmptyMapUsesEvaluatedKeyType(t *testing.T) { LeftTypes: []types.Type{types.T_int32.ToType()}, Conditions: [][]*plan.Expr{{newExpr(0, varcharTyp)}, {newExpr(0, varcharTyp)}}, } + installTestAllocation(t, arg) jm, err := arg.newEmptyJoinMap(proc) require.NoError(t, err) require.NoError(t, jm.PreAlloc(2)) @@ -311,6 +324,11 @@ func TestRightDedupEmptyBuildProbeMapHonorsHashBuildBudget(t *testing.T) { DedupColTypes: []plan.Type{{Id: int32(types.T_int32)}}, JoinMapTag: tag, } + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(initialBytes, budget) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) arg.AppendChild(probeSource) var callErr error t.Cleanup(func() { @@ -338,8 +356,8 @@ func TestRightDedupEmptyBuildProbeMapHonorsHashBuildBudget(t *testing.T) { require.NotContains(t, callErr.Error(), "convert go error") require.NotContains(t, callErr.Error(), process.ErrHashBuildBudgetAdmission.Error()) require.Contains(t, callErr.Error(), "hash build memory budget exceeded") - require.Equal(t, initialBytes, budget.Used(), - "the admitted initial table remains owned until operator cleanup") + require.Zero(t, budget.Used(), + "failed probe-map construction must roll back its physical allocation") } var ( @@ -537,7 +555,7 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] // }, //}) tag++ - return joinTestCase{ + tc := joinTestCase{ types: ts, flgs: flgs, proc: proc, @@ -569,6 +587,8 @@ func newTestCase(t *testing.T, flgs []bool, ts []types.Type, rp []int32, cs [][] JoinMapRefCnt: 1, }, } + installTestAllocation(t, tc.arg, tc.barg) + return tc } func resetChildren(arg *RightDedupJoin, m *mpool.MPool) { diff --git a/pkg/sql/colexec/rightdedupjoin/key_contract_test.go b/pkg/sql/colexec/rightdedupjoin/key_contract_test.go index ecdc21fe2ed84..3ca8e1531e632 100644 --- a/pkg/sql/colexec/rightdedupjoin/key_contract_test.go +++ b/pkg/sql/colexec/rightdedupjoin/key_contract_test.go @@ -160,6 +160,7 @@ func runRightDedupJoinDoubleSignedZeroContract( if mode.shuffle { buildArg.RuntimeFilterSpec = &plan.RuntimeFilterSpec{Tag: joinMapTag + 9000} } + installTestAllocation(t, rightDedupArg, buildArg) buildArg.AppendChild(colexec.NewMockOperator().WithBatchs([]*batch.Batch{buildBatch})) spillBefore := promtestutil.ToFloat64( diff --git a/pkg/sql/colexec/rightdedupjoin/types.go b/pkg/sql/colexec/rightdedupjoin/types.go index e5f7d820657a7..54bbd08894387 100644 --- a/pkg/sql/colexec/rightdedupjoin/types.go +++ b/pkg/sql/colexec/rightdedupjoin/types.go @@ -17,6 +17,7 @@ package rightdedupjoin import ( "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/hashmap" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/common/reuse" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -39,6 +40,15 @@ const ( End ) +const rightDedupJoinAllocationSiteMatched mpool.AllocationSite = 90 + +const ( + rightDedupJoinAllocationSiteResultData mpool.AllocationSite = iota + 114 + rightDedupJoinAllocationSiteResultArea + rightDedupJoinAllocationSiteResultNulls + rightDedupJoinAllocationSiteResultGrouping +) + type evalVector struct { executor colexec.ExpressionExecutor vec *vector.Vector @@ -64,11 +74,7 @@ type container struct { spillEngine *spillutil.SpillEngine spillThreshold int64 - // Non-nil only for spilled joins, where probe expressions are part of the - // shared HashBuild/spill working set. Resident probe expressions remain - // under normal process/mpool accounting; this is not a general query budget. - probeExpressionLease *hashbuild.ExpressionMemoryLease - resultBatch *batch.Batch + resultBatch *batch.Batch } type RightDedupJoin struct { @@ -90,10 +96,63 @@ type RightDedupJoin struct { DelColIdx int32 UpdateColIdxList []int32 UpdateColExprList []*plan.Expr + allocationAccount *mpool.AllocationAccount + resultAllocation *vector.AllocationAccountSelection vm.OperatorBase } +func (rightDedupJoin *RightDedupJoin) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if account == nil || account.Handle() == 0 { + return mpool.ErrAllocationAccountInvalid + } + if rightDedupJoin.allocationAccount != nil && + rightDedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if rightDedupJoin.allocationAccount == account { + return nil + } + selection, err := vector.NewAllocationAccountSelection( + account, + hashbuild.HashBuildAllocationOwner, + rightDedupJoinAllocationSiteResultData, + rightDedupJoinAllocationSiteResultArea, + rightDedupJoinAllocationSiteResultNulls, + rightDedupJoinAllocationSiteResultGrouping, + ) + if err != nil { + return err + } + rightDedupJoin.allocationAccount = account + rightDedupJoin.resultAllocation = selection + return nil +} + +func (rightDedupJoin *RightDedupJoin) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if rightDedupJoin.allocationAccount == nil { + return nil + } + if rightDedupJoin.allocationAccount != account { + return mpool.ErrAllocationAccountMismatch + } + if rightDedupJoin.ctr.mp != nil || + rightDedupJoin.ctr.spillEngine != nil || + len(rightDedupJoin.ctr.evecs) != 0 || + len(rightDedupJoin.ctr.exprExecs) != 0 || + rightDedupJoin.ctr.matched != nil || + rightDedupJoin.ctr.resultBatch != nil { + return mpool.ErrAllocationAccountInvariant + } + rightDedupJoin.allocationAccount = nil + rightDedupJoin.resultAllocation = nil + return nil +} + func (rightDedupJoin *RightDedupJoin) GetOperatorBase() *vm.OperatorBase { return &rightDedupJoin.OperatorBase } @@ -131,30 +190,26 @@ func (rightDedupJoin *RightDedupJoin) Reset(proc *process.Process, pipelineFaile rightDedupJoin.OpAnalyzer.Alloc(ctr.maxAllocSize) } ctr.maxAllocSize = 0 + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil ctr.groupCount = 0 ctr.buildGroupCount = 0 - ctr.cleanBitmap() + ctr.cleanBitmap(proc) ctr.cleanHashMap() - ctr.resetResultBatch() - ctr.resetExprExecutor() + ctr.cleanResultBatch(proc) + ctr.cleanExprExecutor() if ctr.spillEngine != nil { ctr.spillEngine.Cleanup(proc) ctr.spillEngine = nil } - if ctr.probeExpressionLease != nil { - ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() - } else { - ctr.resetEvalVectors() - } + ctr.cleanEvalVectors() ctr.state = Build } func (rightDedupJoin *RightDedupJoin) Free(proc *process.Process, pipelineFailed bool, err error) { ctr := &rightDedupJoin.ctr - ctr.cleanBitmap() + ctr.cleanBitmap(proc) ctr.cleanHashMap() ctr.cleanResultBatch(proc) ctr.cleanExprExecutor() @@ -163,27 +218,23 @@ func (rightDedupJoin *RightDedupJoin) Free(proc *process.Process, pipelineFailed ctr.spillEngine = nil } ctr.cleanEvalVectors() - ctr.releaseProbeExpressionLease() } func (rightDedupJoin *RightDedupJoin) ExecProjection(proc *process.Process, input *batch.Batch) (*batch.Batch, error) { return input, nil } -func (ctr *container) resetExprExecutor() { - for i := range ctr.exprExecs { - ctr.exprExecs[i].ResetForNextQuery() - } -} - func (ctr *container) cleanExprExecutor() { for i := range ctr.exprExecs { - ctr.exprExecs[i].Free() - ctr.exprExecs[i] = nil + if ctr.exprExecs[i] != nil { + ctr.exprExecs[i].Free() + } } + ctr.exprExecs = nil } func (ctr *container) cleanHashMap() { + hashmap.IteratorClearOwner(ctr.itr) ctr.itr = nil if ctr.mp != nil { ctr.mp.Free() @@ -191,7 +242,8 @@ func (ctr *container) cleanHashMap() { } } -func (ctr *container) cleanBitmap() { +func (ctr *container) cleanBitmap(proc *process.Process) { + colexec.FreeAccountedBitmap(ctr.matched, proc.Mp()) ctr.matched = nil } @@ -206,6 +258,32 @@ func (ctr *container) resetResultBatch() { } } +func (rightDedupJoin *RightDedupJoin) resetResultBatch() error { + ctr := &rightDedupJoin.ctr + ctr.resetResultBatch() + if ctr.resultBatch != nil { + return nil + } + ctr.resultBatch = batch.NewOffHeapWithSize(len(rightDedupJoin.Result)) + for i, rp := range rightDedupJoin.Result { + if rp.Rel == 0 { + ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType( + rightDedupJoin.LeftTypes[rp.Pos], + ) + } else { + ctr.resultBatch.Vecs[i] = vector.NewOffHeapVecWithType( + rightDedupJoin.RightTypes[rp.Pos], + ) + } + } + if err := ctr.resultBatch.SetAllocationAccount(rightDedupJoin.resultAllocation); err != nil { + ctr.resultBatch.Clean(nil) + ctr.resultBatch = nil + return err + } + return nil +} + func (ctr *container) cleanResultBatch(proc *process.Process) { if ctr.resultBatch != nil { ctr.resultBatch.Clean(proc.Mp()) @@ -223,18 +301,3 @@ func (ctr *container) cleanEvalVectors() { ctr.evecs = nil ctr.vecs = nil } - -func (ctr *container) resetEvalVectors() { - for i := range ctr.evecs { - if ctr.evecs[i].executor != nil { - ctr.evecs[i].executor.ResetForNextQuery() - } - } -} - -func (ctr *container) releaseProbeExpressionLease() { - if ctr.probeExpressionLease != nil { - ctr.probeExpressionLease.Release() - ctr.probeExpressionLease = nil - } -} diff --git a/pkg/sql/colexec/runtimefilter/contract.go b/pkg/sql/colexec/runtimefilter/contract.go index fbfd8e2423dc6..453c09dd7129f 100644 --- a/pkg/sql/colexec/runtimefilter/contract.go +++ b/pkg/sql/colexec/runtimefilter/contract.go @@ -18,7 +18,6 @@ package runtimefilter import ( - "bytes" "context" "errors" "math" @@ -47,7 +46,8 @@ func (e *optionalAllocationError) Unwrap() error { return e.cause } // MarkOptionalAllocationError preserves the allocation error while giving a // runtime-filter producer a narrow fail-open classification. func MarkOptionalAllocationError(err error) error { - if err == nil || IsOptionalAllocationError(err) { + if err == nil || IsOptionalAllocationError(err) || + !mpool.IsRetryableAllocationCapacity(err) { return err } return &optionalAllocationError{cause: err} @@ -79,21 +79,25 @@ func ClassifyOptionalFallback(err error) OptionalFallbackKind { errors.Is(err, context.DeadlineExceeded) { return OptionalFallbackNone } + allocationReason := mpool.AllocationFailureReasonOf(err) + if allocationReason != mpool.AllocationFailureNone && + allocationReason != mpool.AllocationFailureCapacity { + return OptionalFallbackNone + } // Reject every known fatal branch before accepting an admission branch. // This also keeps errors.Join(admission, fatal) fatal regardless of the // traversal order chosen by errors.As below. if errors.Is(err, process.ErrHashBuildBudgetClosed) || errors.Is(err, process.ErrHashBuildBudgetInvalid) || - errors.Is(err, process.ErrHashBuildCeilingMissing) || - errors.Is(err, process.ErrHashBuildReservationInactive) || - errors.Is(err, process.ErrHashBuildReservationUpward) { + errors.Is(err, process.ErrHashBuildCeilingMissing) { return OptionalFallbackNone } var budgetErr *process.HashBuildBudgetError if errors.As(err, &budgetErr) { if budgetErr != nil && - budgetErr.Kind == process.HashBuildBudgetErrorAdmission { + budgetErr.Kind == process.HashBuildBudgetErrorAdmission && + budgetErr.Component == process.HashBuildBudgetComponentMemory { return OptionalFallbackBudgetAdmission } return OptionalFallbackNone @@ -105,7 +109,8 @@ func ClassifyOptionalFallback(err error) OptionalFallbackKind { return OptionalFallbackNone } - if IsOptionalAllocationError(err) { + if IsOptionalAllocationError(err) && + mpool.IsRetryableAllocationCapacity(err) { return OptionalFallbackAllocation } return OptionalFallbackNone @@ -292,12 +297,10 @@ func planType(typ plan.Type) types.Type { } // CloseFloatSignedZero appends the complementary representation when an exact -// float payload contains only one of +0 and -0. beforeAppend lets budgeted -// producers reserve the vector-growth overlap before the allocation occurs. +// float payload contains only one of +0 and -0. func CloseFloatSignedZero( vec *vector.Vector, mp *mpool.MPool, - beforeAppend func() (release func(), err error), ) error { if vec == nil || mp == nil { return moerr.NewInternalErrorNoCtx("invalid float runtime-filter vector") @@ -347,18 +350,6 @@ func CloseFloatSignedZero( return nil } - var release func() - var err error - if beforeAppend != nil { - release, err = beforeAppend() - if err != nil { - return err - } - } - if release != nil { - defer release() - } - if vec.GetType().Oid == types.T_float32 { value := float32(0) if hasPositiveZero { @@ -375,84 +366,41 @@ func CloseFloatSignedZero( vector.AppendFixed(vec, value, false, mp)) } -// MarshalExactFilterVector serializes an exact-filter vector under the -// statement/CN hash-build budget. Runtime-filter payloads live on the Go heap, -// outside mpool accounting, so every producer must retain this reservation -// until the MessageBoard destroys the message. -// -// Exact IN payloads have already discarded NULL. Requiring an empty null -// bitmap makes the wire size exact before allocation and avoids a second, -// unbudgeted roaring-bitmap serialization. +// MarshalExactFilterVector serializes an exact-filter vector into physical +// MPool storage owned by the statement allocation account. The returned +// release closure transfers that storage lifetime to the MessageBoard. func MarshalExactFilterVector( vec *vector.Vector, - budget *process.HashBuildBudgetGeneration, + mp *mpool.MPool, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, ) ([]byte, func(), error) { - if vec == nil || budget == nil || vec.GetNulls().Any() { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - - length := vec.Length() - typeSize := vec.GetType().TypeSize() - if length < 0 || uint64(length) > math.MaxUint32 || typeSize < 0 { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - dataBytes := uint64(typeSize) - if !vec.IsConst() { - if typeSize > 0 && uint64(length) > math.MaxUint64/uint64(typeSize) { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - dataBytes *= uint64(length) - } else if vec.IsConstNull() { - dataBytes = 0 - } - areaBytes := uint64(len(vec.GetArea())) - if dataBytes > math.MaxUint32 || areaBytes > math.MaxUint32 || - dataBytes > uint64(len(vec.GetData())) { + if vec == nil || mp == nil || account == nil || vec.GetNulls().Any() { return nil, nil, process.ErrHashBuildBudgetInvalid } - - // class + encoded type + length/data/area/null lengths + sorted flag. - headerBytes := uint64(1 + len(types.EncodeType(vec.GetType())) + 4*4 + 1) - if dataBytes > math.MaxUint64-headerBytes || - areaBytes > math.MaxUint64-headerBytes-dataBytes { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - wireBytes := headerBytes + dataBytes + areaBytes - if wireBytes > uint64(math.MaxInt) { - return nil, nil, process.ErrHashBuildBudgetInvalid - } - - // bytes.Buffer's visible capacity may be rounded above Grow's request. - // Reserve a bounded allocator overlap, verify it after allocation, then - // reconcile to the capacity retained by the message. - const allocationSlack = uint64(64 << 10) - if wireBytes > math.MaxUint64-allocationSlack { - return nil, nil, process.ErrHashBuildBudgetInvalid + plan, err := vec.PrepareMarshalBinary() + if err != nil { + return nil, nil, err } - projected := wireBytes + allocationSlack - token, err := budget.Reserve(projected) + buf, err := mpool.NewAccountedBuffer(mp, account, owner, site) if err != nil { return nil, nil, err } - - var buf bytes.Buffer - buf.Grow(int(wireBytes)) - if uint64(buf.Cap()) > projected { - token.Release() - return nil, nil, process.ErrHashBuildBudgetInvalid + if err = buf.EnsureCapacity(plan.Size()); err != nil { + buf.Free() + if mpool.IsRetryableAllocationCapacity(err) { + err = MarkOptionalAllocationError(err) + } + return nil, nil, err } - if err = vec.MarshalBinaryWithBuffer(&buf); err != nil { - token.Release() + if err = plan.MarshalTo(buf); err != nil { + buf.Free() return nil, nil, err } - data := buf.Bytes() - if uint64(len(data)) != wireBytes || uint64(cap(data)) > projected { - token.Release() + if buf.Len() != plan.Size() { + buf.Free() return nil, nil, process.ErrHashBuildBudgetInvalid } - if _, err = token.ReconcileDown(uint64(cap(data))); err != nil { - token.Release() - return nil, nil, err - } - return data, func() { token.Release() }, nil + return buf.Bytes(), buf.Free, nil } diff --git a/pkg/sql/colexec/runtimefilter/contract_test.go b/pkg/sql/colexec/runtimefilter/contract_test.go index b0dc73008731a..9a1f9bff1329e 100644 --- a/pkg/sql/colexec/runtimefilter/contract_test.go +++ b/pkg/sql/colexec/runtimefilter/contract_test.go @@ -20,6 +20,7 @@ import ( "testing" "github.com/matrixorigin/matrixone/pkg/common/hashmap/keycodec" + "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -30,13 +31,22 @@ import ( ) func TestClassifyOptionalFallbackFatalFirst(t *testing.T) { - mpoolErr := errors.New("mpool allocation failed") + mpoolErr := mpool.ErrAllocationAccountCapacity providerErr := errors.New("budget provider failed") var nilBudgetErr *process.HashBuildBudgetError - budgetErr := func(kind process.HashBuildBudgetErrorKind) error { - return &process.HashBuildBudgetError{Kind: kind} + budgetErr := func( + kind process.HashBuildBudgetErrorKind, + component process.HashBuildBudgetComponent, + ) error { + return &process.HashBuildBudgetError{Kind: kind, Component: component} } marked := MarkOptionalAllocationError + memoryAdmission := func() error { + return budgetErr( + process.HashBuildBudgetErrorAdmission, + process.HashBuildBudgetComponentMemory, + ) + } tests := []struct { name string @@ -45,26 +55,34 @@ func TestClassifyOptionalFallbackFatalFirst(t *testing.T) { }{ {name: "nil", want: OptionalFallbackNone}, {name: "typed nil budget error", err: nilBudgetErr, want: OptionalFallbackNone}, - {name: "typed admission", err: budgetErr(process.HashBuildBudgetErrorAdmission), want: OptionalFallbackBudgetAdmission}, - {name: "marked typed admission", err: marked(budgetErr(process.HashBuildBudgetErrorAdmission)), want: OptionalFallbackBudgetAdmission}, - {name: "joined admission and closed", err: errors.Join(budgetErr(process.HashBuildBudgetErrorAdmission), budgetErr(process.HashBuildBudgetErrorClosed)), want: OptionalFallbackNone}, + {name: "typed memory admission", err: memoryAdmission(), want: OptionalFallbackBudgetAdmission}, + {name: "marked typed memory admission", err: marked(memoryAdmission()), want: OptionalFallbackBudgetAdmission}, + {name: "typed spill disk admission", err: budgetErr(process.HashBuildBudgetErrorAdmission, process.HashBuildBudgetComponentSpillDisk), want: OptionalFallbackNone}, + {name: "typed spill fd admission", err: budgetErr(process.HashBuildBudgetErrorAdmission, process.HashBuildBudgetComponentSpillFD), want: OptionalFallbackNone}, + {name: "typed admission without component", err: budgetErr(process.HashBuildBudgetErrorAdmission, 0), want: OptionalFallbackNone}, + {name: "joined admission and closed", err: errors.Join(memoryAdmission(), budgetErr(process.HashBuildBudgetErrorClosed, 0)), want: OptionalFallbackNone}, {name: "marked mpool allocation", err: marked(mpoolErr), want: OptionalFallbackAllocation}, {name: "plain mpool error", err: mpoolErr, want: OptionalFallbackNone}, + {name: "marked sealed", err: marked(mpool.ErrAllocationAccountSealed), want: OptionalFallbackNone}, + {name: "marked mismatch", err: marked(mpool.ErrAllocationAccountMismatch), want: OptionalFallbackNone}, + {name: "marked invariant", err: marked(mpool.ErrAllocationAccountInvariant), want: OptionalFallbackNone}, + {name: "marked allocator limit", err: marked(mpool.ErrAllocationAllocatorLimit), want: OptionalFallbackNone}, + {name: "marked suspended", err: marked(mpool.ErrAllocationAdmissionSuspended), want: OptionalFallbackNone}, + {name: "marked joined capacity and invariant", err: marked(errors.Join(mpool.ErrAllocationAccountCapacity, mpool.ErrAllocationAccountInvariant)), want: OptionalFallbackNone}, + {name: "marked joined mpool capacity and invalid", err: marked(errors.Join(moerr.NewMPoolCapacityNoCtxf("test"), mpool.ErrAllocationAccountInvalid)), want: OptionalFallbackNone}, {name: "plain provider error", err: providerErr, want: OptionalFallbackNone}, {name: "raw admission sentinel", err: process.ErrHashBuildBudgetAdmission, want: OptionalFallbackNone}, - {name: "typed closed", err: budgetErr(process.HashBuildBudgetErrorClosed), want: OptionalFallbackNone}, - {name: "marked typed closed", err: marked(budgetErr(process.HashBuildBudgetErrorClosed)), want: OptionalFallbackNone}, - {name: "typed invalid", err: budgetErr(process.HashBuildBudgetErrorInvalid), want: OptionalFallbackNone}, - {name: "marked typed invalid", err: marked(budgetErr(process.HashBuildBudgetErrorInvalid)), want: OptionalFallbackNone}, - {name: "typed ceiling missing", err: budgetErr(process.HashBuildBudgetErrorCeilingMissing), want: OptionalFallbackNone}, - {name: "marked typed ceiling missing", err: marked(budgetErr(process.HashBuildBudgetErrorCeilingMissing)), want: OptionalFallbackNone}, + {name: "typed closed", err: budgetErr(process.HashBuildBudgetErrorClosed, 0), want: OptionalFallbackNone}, + {name: "marked typed closed", err: marked(budgetErr(process.HashBuildBudgetErrorClosed, 0)), want: OptionalFallbackNone}, + {name: "typed invalid", err: budgetErr(process.HashBuildBudgetErrorInvalid, 0), want: OptionalFallbackNone}, + {name: "marked typed invalid", err: marked(budgetErr(process.HashBuildBudgetErrorInvalid, 0)), want: OptionalFallbackNone}, + {name: "typed ceiling missing", err: budgetErr(process.HashBuildBudgetErrorCeilingMissing, 0), want: OptionalFallbackNone}, + {name: "marked typed ceiling missing", err: marked(budgetErr(process.HashBuildBudgetErrorCeilingMissing, 0)), want: OptionalFallbackNone}, {name: "marked canceled", err: marked(context.Canceled), want: OptionalFallbackNone}, {name: "marked deadline", err: marked(context.DeadlineExceeded), want: OptionalFallbackNone}, {name: "marked raw closed", err: marked(process.ErrHashBuildBudgetClosed), want: OptionalFallbackNone}, {name: "marked raw invalid", err: marked(process.ErrHashBuildBudgetInvalid), want: OptionalFallbackNone}, {name: "marked raw ceiling", err: marked(process.ErrHashBuildCeilingMissing), want: OptionalFallbackNone}, - {name: "marked inactive reservation", err: marked(process.ErrHashBuildReservationInactive), want: OptionalFallbackNone}, - {name: "marked upward reconciliation", err: marked(process.ErrHashBuildReservationUpward), want: OptionalFallbackNone}, } for _, test := range tests { @@ -274,10 +292,14 @@ func TestMarshalExactFilterVectorUsesWireSizedBudget(t *testing.T) { aggregate := process.MustNewHashBuildBudget(1<<20, 1<<20) budget, err := aggregate.OpenGeneration(1) require.NoError(t, err) - data, release, err := MarshalExactFilterVector(vec, budget) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<20, budget) + require.NoError(t, err) + data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.NoError(t, err) require.Len(t, data, 34+vec.Length()) - // The retained charge is the actual bytes.Buffer capacity, not a + // The retained charge is the actual caller-owned buffer capacity, not a // row-count-derived metadata estimate. require.LessOrEqual(t, budget.Used(), uint64(2*len(data))) require.NotZero(t, budget.Used()) @@ -297,7 +319,11 @@ func TestMarshalExactFilterVectorAdmissionFailsBeforeAllocation(t *testing.T) { aggregate := process.MustNewHashBuildBudget(1, 1) budget, err := aggregate.OpenGeneration(1) require.NoError(t, err) - data, release, err := MarshalExactFilterVector(vec, budget) + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.OpenWithController(1<<20, budget) + require.NoError(t, err) + data, release, err := MarshalExactFilterVector(vec, mp, account, 1, 1) require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) require.Nil(t, data) require.Nil(t, release) diff --git a/pkg/sql/colexec/sample/sample.go b/pkg/sql/colexec/sample/sample.go index a46ef5294c69c..6bb509b8eac65 100644 --- a/pkg/sql/colexec/sample/sample.go +++ b/pkg/sql/colexec/sample/sample.go @@ -214,42 +214,204 @@ func (ctr *container) evaluateSampleAndGroupByColumns(proc *process.Process, bat } func (ctr *container) hashAndSample(bat *batch.Batch, proc *process.Process) (err error) { - var iterator hashmap.Iterator - var groupList []uint64 count := bat.RowCount() + if !hasGroupingRows(ctr.groupVectors) { + return ctr.hashNormalRows(bat, proc, 0, count) + } + if err = ctr.enableGroupingDomain(proc); err != nil { + return err + } + groupingIterator := ctr.groupingHashMap.NewIterator() + var normalIterator hashmap.Iterator + + for offset := 0; offset < count; { + grouping := rowHasGrouping(ctr.groupVectors, offset) + end := offset + 1 + for end < count && rowHasGrouping(ctr.groupVectors, end) == grouping { + end++ + } + if grouping { + err = ctr.hashRows( + bat, + groupingIterator, + &ctr.groupingGroupIDs, + offset, + end-offset, + true, + ) + } else { + if normalIterator == nil { + normalIterator, err = ctr.normalIterator(proc) + if err != nil { + return err + } + } + err = ctr.hashRows( + bat, + normalIterator, + &ctr.normalGroupIDs, + offset, + end-offset, + true, + ) + } + if err != nil { + return err + } + offset = end + } + return nil +} + +func hasGroupingRows(vecs []*vector.Vector) bool { + for _, vec := range vecs { + if vec != nil && vec.HasGrouping() { + return true + } + } + return false +} + +func rowHasGrouping(vecs []*vector.Vector, row int) bool { + for _, vec := range vecs { + if vec != nil && vec.GetGrouping().Contains(uint64(row)) { + return true + } + } + return false +} + +func (ctr *container) normalIterator(proc *process.Process) (hashmap.Iterator, error) { + var err error if ctr.useIntHashMap { if ctr.intHashMap == nil { ctr.intHashMap, err = hashmap.NewIntHashMap(ctr.groupVectorsNullable, proc.Mp()) if err != nil { - return err + return nil, err } } - iterator = ctr.intHashMap.NewIterator() + return ctr.intHashMap.NewIterator(), nil } else { if ctr.strHashMap == nil { ctr.strHashMap, err = hashmap.NewStrHashMap(ctr.groupVectorsNullable, proc.Mp()) if err != nil { - return err + return nil, err } } - iterator = ctr.strHashMap.NewIterator() + return ctr.strHashMap.NewIterator(), nil } +} - for i := 0; i < count; i += hashmap.UnitLimit { - n := count - i +func (ctr *container) normalGroupCount() uint64 { + if ctr.useIntHashMap && ctr.intHashMap != nil { + return ctr.intHashMap.GroupCount() + } + if !ctr.useIntHashMap && ctr.strHashMap != nil { + return ctr.strHashMap.GroupCount() + } + return 0 +} + +func (ctr *container) enableGroupingDomain(proc *process.Process) error { + if ctr.groupingHashMap != nil { + return nil + } + groupingMap, err := hashmap.NewStrHashMap( + ctr.groupVectorsNullable, + proc.Mp(), + ) + if err != nil { + return err + } + if err = groupingMap.SetGroupingAware(); err != nil { + groupingMap.Free() + return err + } + + normalGroups := ctr.normalGroupCount() + ctr.normalGroupIDs = make([]uint64, normalGroups+1) + for i := uint64(1); i <= normalGroups; i++ { + ctr.normalGroupIDs[i] = i + } + ctr.groupingGroupIDs = []uint64{0} + ctr.nextGroupID = normalGroups + ctr.groupingHashMap = groupingMap + return nil +} + +func (ctr *container) globalGroupIDs( + local []uint64, + translation *[]uint64, +) []uint64 { + ids := *translation + for i, localID := range local { + if localID == 0 { + continue + } + for uint64(len(ids)) <= localID { + ctr.nextGroupID++ + ids = append(ids, ctr.nextGroupID) + } + local[i] = ids[localID] + } + *translation = ids + return local +} + +func (ctr *container) hashNormalRows( + bat *batch.Batch, + proc *process.Process, + offset int, + count int, +) error { + iterator, err := ctr.normalIterator(proc) + if err != nil { + return err + } + return ctr.hashRows( + bat, + iterator, + &ctr.normalGroupIDs, + offset, + count, + ctr.groupingHashMap != nil, + ) +} + +func (ctr *container) hashRows( + bat *batch.Batch, + iterator hashmap.Iterator, + translation *[]uint64, + offset int, + count int, + translate bool, +) error { + end := offset + count + for offset < end { + n := end - offset if n > hashmap.UnitLimit { n = hashmap.UnitLimit } - groupList, _, err = iterator.Insert(i, n, ctr.groupVectors) + groupList, _, err := iterator.Insert(offset, n, ctr.groupVectors) if err != nil { return err } - err = ctr.samplePool.BatchSample(i, n, groupList, ctr.sampleVectors, ctr.groupVectors, bat) - if err != nil { + if translate { + groupList = ctr.globalGroupIDs(groupList[:n], translation) + } + if err = ctr.samplePool.BatchSample( + offset, + n, + groupList, + ctr.sampleVectors, + ctr.groupVectors, + bat, + ); err != nil { return err } + offset += n } - return + return nil } diff --git a/pkg/sql/colexec/sample/sample_test.go b/pkg/sql/colexec/sample/sample_test.go index b51a96929c6f1..b589b593ce5c3 100644 --- a/pkg/sql/colexec/sample/sample_test.go +++ b/pkg/sql/colexec/sample/sample_test.go @@ -20,6 +20,7 @@ import ( "runtime/debug" "testing" + "github.com/matrixorigin/matrixone/pkg/common/hashmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -263,6 +264,340 @@ func TestSamplePool(t *testing.T) { require.Equal(t, int64(0), proc.Mp().CurrNB()) } +func TestSampleSeparatesGroupingKeyDomain(t *testing.T) { + for _, test := range []struct { + name string + pool func(*process.Process) *sPool + }{ + { + name: "row", + pool: func(proc *process.Process) *sPool { + return newSamplePoolByRows(proc, 1, 1, false) + }, + }, + { + name: "percent", + pool: func(proc *process.Process) *sPool { + return newSamplePoolByPercent(proc, 100, 1) + }, + }, + { + name: "merge", + pool: func(proc *process.Process) *sPool { + return newSamplePoolByRowsForMerge(proc, 1, 1, false) + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: test.pool(proc), + } + defer freeSampleHashContainer(ctr) + + ordinary := makeSampleGroupingBatch( + t, + proc, + []int64{1, 2}, + []int64{10, 20}, + false, + test.name == "merge", + ) + ctr.groupVectors = ordinary.Vecs[:1] + ctr.sampleVectors = ordinary.Vecs[1:2] + require.NoError(t, ctr.hashAndSample(ordinary, proc)) + ordinary.Clean(proc.Mp()) + + rollup := makeSampleGroupingBatch( + t, + proc, + []int64{0}, + []int64{30}, + true, + test.name == "merge", + ) + ctr.groupVectors = rollup.Vecs[:1] + ctr.sampleVectors = rollup.Vecs[1:2] + require.NoError(t, ctr.hashAndSample(rollup, proc)) + rollup.Clean(proc.Mp()) + + result, err := ctr.samplePool.Result(true) + require.NoError(t, err) + require.Equal(t, 3, result.RowCount()) + require.Equal(t, 1, result.Vecs[0].GetGrouping().Count()) + result.Clean(proc.Mp()) + }) + } +} + +func TestSampleSeparatesSQLNullFromGrouping(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + groupVectorsNullable: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + } + defer freeSampleHashContainer(ctr) + + nullBatch := batch.NewWithSize(2) + nullBatch.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed( + nullBatch.Vecs[0], + int64(0), + true, + proc.Mp(), + )) + nullBatch.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed( + nullBatch.Vecs[1], + int64(10), + false, + proc.Mp(), + )) + nullBatch.SetRowCount(1) + ctr.groupVectors = nullBatch.Vecs[:1] + ctr.sampleVectors = nullBatch.Vecs[1:] + require.NoError(t, ctr.hashAndSample(nullBatch, proc)) + nullBatch.Clean(proc.Mp()) + + rollup := makeSampleGroupingBatch( + t, + proc, + []int64{0}, + []int64{20}, + true, + false, + ) + ctr.groupVectors = rollup.Vecs[:1] + ctr.sampleVectors = rollup.Vecs[1:] + require.NoError(t, ctr.hashAndSample(rollup, proc)) + rollup.Clean(proc.Mp()) + + result, err := ctr.samplePool.Result(true) + require.NoError(t, err) + require.Equal(t, 2, result.RowCount()) + require.Equal(t, 2, result.Vecs[0].GetNulls().Count()) + require.Equal(t, 1, result.Vecs[0].GetGrouping().Count()) + result.Clean(proc.Mp()) +} + +func TestSampleSeparatesPartialRollupKeys(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + } + defer freeSampleHashContainer(ctr) + + input := batch.NewWithSize(3) + for column := 0; column < 2; column++ { + input.Vecs[column] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + input.Vecs[column], + []int64{0, 0}, + nil, + proc.Mp(), + )) + } + input.Vecs[0].GetGrouping().Add(0) + input.Vecs[1].GetGrouping().Add(1) + input.Vecs[2] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList( + input.Vecs[2], + []int64{10, 20}, + nil, + proc.Mp(), + )) + input.SetRowCount(2) + ctr.groupVectors = input.Vecs[:2] + ctr.sampleVectors = input.Vecs[2:] + require.NoError(t, ctr.hashAndSample(input, proc)) + input.Clean(proc.Mp()) + + result, err := ctr.samplePool.Result(true) + require.NoError(t, err) + require.Equal(t, 2, result.RowCount()) + require.Equal(t, 1, result.Vecs[0].GetGrouping().Count()) + require.Equal(t, 1, result.Vecs[1].GetGrouping().Count()) + result.Clean(proc.Mp()) +} + +func TestSampleAlternatingGroupingReusesIterators(t *testing.T) { + proc := testutil.NewProcess(t) + defer proc.Free() + input := makeSampleGroupingBatch( + t, + proc, + makeSequence(hashmap.UnitLimit), + makeSequence(hashmap.UnitLimit), + false, + false, + ) + defer input.Clean(proc.Mp()) + for row := 0; row < input.RowCount(); row += 2 { + input.Vecs[0].GetGrouping().Add(uint64(row)) + } + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + groupVectors: input.Vecs[:1], + sampleVectors: input.Vecs[1:], + } + defer freeSampleHashContainer(ctr) + require.NoError(t, ctr.hashAndSample(input, proc)) + + var runErr error + allocations := testing.AllocsPerRun(20, func() { + runErr = ctr.hashAndSample(input, proc) + }) + require.NoError(t, runErr) + require.Less(t, allocations, float64(32)) +} + +func BenchmarkSampleGroupedHashFastPath(b *testing.B) { + proc := testutil.NewProcess(b) + defer proc.Free() + input := makeSampleGroupingBatch( + b, + proc, + makeSequence(256), + makeSequence(256), + false, + false, + ) + defer input.Clean(proc.Mp()) + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + groupVectors: input.Vecs[:1], + sampleVectors: input.Vecs[1:], + } + defer freeSampleHashContainer(ctr) + require.NoError(b, ctr.hashAndSample(input, proc)) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err := ctr.hashAndSample(input, proc); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkSampleAlternatingGrouping(b *testing.B) { + proc := testutil.NewProcess(b) + defer proc.Free() + input := makeSampleGroupingBatch( + b, + proc, + makeSequence(hashmap.UnitLimit), + makeSequence(hashmap.UnitLimit), + false, + false, + ) + defer input.Clean(proc.Mp()) + for row := 0; row < input.RowCount(); row += 2 { + input.Vecs[0].GetGrouping().Add(uint64(row)) + } + ctr := &container{ + isGroupBy: true, + useIntHashMap: true, + samplePool: newSamplePoolByRows(proc, 1, 1, false), + groupVectors: input.Vecs[:1], + sampleVectors: input.Vecs[1:], + } + defer freeSampleHashContainer(ctr) + require.NoError(b, ctr.hashAndSample(input, proc)) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if err := ctr.hashAndSample(input, proc); err != nil { + b.Fatal(err) + } + } +} + +func makeSequence(count int) []int64 { + values := make([]int64, count) + for i := range values { + values[i] = int64(i) + } + return values +} + +func makeSampleGroupingBatch( + tb testing.TB, + proc *process.Process, + groups []int64, + samples []int64, + rollup bool, + merge bool, +) *batch.Batch { + tb.Helper() + columns := 2 + if merge { + columns++ + } + bat := batch.NewWithSize(columns) + if rollup { + bat.Vecs[0] = vector.NewRollupConst( + types.T_int64.ToType(), + len(groups), + proc.Mp(), + ) + } else { + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(tb, vector.AppendFixedList( + bat.Vecs[0], + groups, + nil, + proc.Mp(), + )) + } + bat.Vecs[1] = vector.NewVec(types.T_int64.ToType()) + require.NoError(tb, vector.AppendFixedList( + bat.Vecs[1], + samples, + nil, + proc.Mp(), + )) + if merge { + var err error + bat.Vecs[2], err = vector.NewConstFixed( + types.T_int64.ToType(), + int64(len(groups)), + len(groups), + proc.Mp(), + ) + require.NoError(tb, err) + } + bat.SetRowCount(len(groups)) + return bat +} + +func freeSampleHashContainer(ctr *container) { + if ctr.intHashMap != nil { + ctr.intHashMap.Free() + } + if ctr.strHashMap != nil { + ctr.strHashMap.Free() + } + if ctr.groupingHashMap != nil { + ctr.groupingHashMap.Free() + } + ctr.samplePool.Free() +} + func genSampleBatch(proc *process.Process, rows [][]int64) (*batch.Batch, error) { b := batch.NewWithSize(len(rows[0])) diff --git a/pkg/sql/colexec/sample/types.go b/pkg/sql/colexec/sample/types.go index e10009533afee..95c3b4c0944bd 100644 --- a/pkg/sql/colexec/sample/types.go +++ b/pkg/sql/colexec/sample/types.go @@ -83,8 +83,16 @@ type container struct { buf *batch.Batch // hash map related. - intHashMap *hashmap.IntHashMap - strHashMap *hashmap.StrHashMap + intHashMap *hashmap.IntHashMap + strHashMap *hashmap.StrHashMap + groupingHashMap *hashmap.StrHashMap + + // A grouping-aware key has a domain that cannot be represented by the + // normal IntHashMap and cannot be installed into a populated StrHashMap. + // These tables translate each map's local IDs into one sample-pool domain. + normalGroupIDs []uint64 + groupingGroupIDs []uint64 + nextGroupID uint64 } func init() { @@ -209,6 +217,9 @@ func (sample *Sample) Free(proc *process.Process, pipelineFailed bool, err error if sample.ctr.strHashMap != nil { sample.ctr.strHashMap.Free() } + if sample.ctr.groupingHashMap != nil { + sample.ctr.groupingHashMap.Free() + } for _, executor := range sample.ctr.sampleExecutors { if executor != nil { executor.Free() diff --git a/pkg/sql/colexec/shuffle/shufflepool.go b/pkg/sql/colexec/shuffle/shufflepool.go index d038e2db6abd8..9fbec8594b4be 100644 --- a/pkg/sql/colexec/shuffle/shufflepool.go +++ b/pkg/sql/colexec/shuffle/shufflepool.go @@ -618,7 +618,7 @@ func (sp *ShufflePool) tryWrite( sp.batchLocks[bucket].Unlock() break } - readyDelta := sp.batchSets[bucket].ReadyDelta(len(chunk)) + readyDelta := sp.batchSets[bucket].ReadyDeltaFor(srcBatch, len(chunk)) wait, ok := sp.reserveReady(int32(bucket), readyDelta) if !ok { sp.batchLocks[bucket].Unlock() diff --git a/pkg/sql/colexec/shuffle/shufflepool_test.go b/pkg/sql/colexec/shuffle/shufflepool_test.go index ab8ea47204816..a8d2f0acb1200 100644 --- a/pkg/sql/colexec/shuffle/shufflepool_test.go +++ b/pkg/sql/colexec/shuffle/shufflepool_test.go @@ -182,6 +182,57 @@ func TestShufflePoolBoundsReadyBatchesAndResumes(t *testing.T) { require.Equal(t, int64(0), proc.Mp().CurrNB()) } +func TestShufflePoolReservesReadyCreditForProvenanceChange(t *testing.T) { + mp := mpool.MustNewZero() + proc := testutil.NewProcessWithMPool(t, "", mp) + defer proc.Free() + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + sp := NewShufflePool(1, 1, true) + unaccounted := testutil.NewBatch([]types.Type{types.T_int64.ToType()}, false, 2, mp) + done, err := writeBatchToBucketForTest(sp, unaccounted, proc, 0) + require.NoError(t, err) + require.True(t, done) + require.Zero(t, sp.readyCount) + + accounted := batch.NewWithSchema( + true, + nil, + []types.Type{types.T_int64.ToType()}, + ) + require.NoError(t, accounted.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(accounted.Vecs[0], int64(7), false, mp)) + accounted.SetRowCount(1) + done, err = writeBatchToBucketForTest(sp, accounted, proc, 0) + require.NoError(t, err) + require.True(t, done) + require.Equal(t, 1, sp.readyCount) + + ready := sp.getAnyFullBatch() + require.NotNil(t, ready) + require.Equal(t, 2, ready.RowCount()) + sp.discardBatch(ready, mp) + require.Zero(t, sp.readyCount) + tail := sp.getAnyLastBatch() + require.NotNil(t, tail) + require.Equal(t, 1, tail.RowCount()) + require.Same(t, selection, tail.AllocationAccountSelection()) + sp.discardBatch(tail, mp) + + unaccounted.Clean(mp) + accounted.Clean(mp) + sp.abort(mp) + require.Zero(t, account.Seal().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal(t, int64(0), mp.CurrNB()) +} + func TestShufflePoolFixedBucketsHaveIndependentBackpressure(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() diff --git a/pkg/sql/colexec/spillutil/allocation_account.go b/pkg/sql/colexec/spillutil/allocation_account.go new file mode 100644 index 0000000000000..d860a197ca62a --- /dev/null +++ b/pkg/sql/colexec/spillutil/allocation_account.go @@ -0,0 +1,187 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package spillutil + +import ( + "math" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" +) + +// Spill allocation sites occupy a dedicated range within the HashBuild owner. +const ( + SpillAllocationSiteDecodedData mpool.AllocationSite = iota + 32 + SpillAllocationSiteDecodedArea + SpillAllocationSiteSelectedData + SpillAllocationSiteSelectedArea + SpillAllocationSiteHashValues + SpillAllocationSiteRowIDs + SpillAllocationSiteMarshalBuffer + SpillAllocationSiteCoalesceBuffer + SpillAllocationSiteDecodedNulls + SpillAllocationSiteDecodedGrouping + SpillAllocationSiteSelectedNulls + SpillAllocationSiteSelectedGrouping +) + +// SpillAllocationAccount is the allocation provenance for one spill +// engine. +type SpillAllocationAccount struct { + account *mpool.AllocationAccount + owner mpool.AllocationOwner + + decoded *vector.AllocationAccountSelection + selected *vector.AllocationAccountSelection +} + +func NewSpillAllocationAccount( + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, +) (*SpillAllocationAccount, error) { + decoded, err := vector.NewAllocationAccountSelection( + account, + owner, + SpillAllocationSiteDecodedData, + SpillAllocationSiteDecodedArea, + SpillAllocationSiteDecodedNulls, + SpillAllocationSiteDecodedGrouping, + ) + if err != nil { + return nil, err + } + selected, err := vector.NewAllocationAccountSelection( + account, + owner, + SpillAllocationSiteSelectedData, + SpillAllocationSiteSelectedArea, + SpillAllocationSiteSelectedNulls, + SpillAllocationSiteSelectedGrouping, + ) + if err != nil { + return nil, err + } + return &SpillAllocationAccount{ + account: account, + owner: owner, + decoded: decoded, + selected: selected, + }, nil +} + +func (a *SpillAllocationAccount) validate() error { + if a == nil || a.account == nil || a.account.Handle() == 0 || + a.owner < mpool.AllocationOwnerMin || + a.owner > mpool.AllocationOwnerMax || + a.decoded == nil || a.selected == nil { + return mpool.ErrAllocationAccountInvalid + } + return nil +} + +func newSpillBatch( + size int, + selection *vector.AllocationAccountSelection, +) (*batch.Batch, error) { + if selection == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + bat := batch.NewOffHeapWithSize(size) + if err := bat.SetAllocationAccount(selection); err != nil { + bat.Clean(nil) + return nil, err + } + return bat, nil +} + +func newSpillVector( + typ types.Type, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { + if selection == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + return vector.NewOffHeapVecWithTypeAndAllocation(typ, selection) +} + +func growSpillSlice[T any]( + values []T, + length int, + mp *mpool.MPool, + allocation *SpillAllocationAccount, + site mpool.AllocationSite, +) ([]T, error) { + if length < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } + if length <= cap(values) { + return values[:length], nil + } + if err := allocation.validate(); err != nil { + return nil, err + } + newCapacity := cap(values) + if newCapacity == 0 { + newCapacity = 1 + } + for newCapacity < length { + if newCapacity > math.MaxInt/2 { + newCapacity = length + break + } + newCapacity *= 2 + } + next, err := mpool.MakeSliceAccounted[T]( + newCapacity, + mp, + allocation.account, + allocation.owner, + site, + ) + if err != nil { + return nil, err + } + copy(next, values) + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } + return next[:length], nil +} + +func freeSpillSlice[T any]( + values []T, + mp *mpool.MPool, +) { + if cap(values) > 0 { + mpool.FreeSlice(mp, values) + } +} + +func (a *SpillAllocationAccount) newBuffer( + mp *mpool.MPool, + site mpool.AllocationSite, +) (*mpool.AccountedBuffer, error) { + if err := a.validate(); err != nil { + return nil, err + } + return mpool.NewAccountedBuffer( + mp, + a.account, + a.owner, + site, + ) +} diff --git a/pkg/sql/colexec/spillutil/allocation_account_test.go b/pkg/sql/colexec/spillutil/allocation_account_test.go new file mode 100644 index 0000000000000..249daac7e970b --- /dev/null +++ b/pkg/sql/colexec/spillutil/allocation_account_test.go @@ -0,0 +1,833 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package spillutil + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +type testSpillAllocationAccount struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + allocation *SpillAllocationAccount + generation *process.HashBuildBudgetGeneration +} + +func makeTestCastKeyExpr( + t testing.TB, + proc *process.Process, +) []*plan.Expr { + t.Helper() + expr, err := plan2.BindFuncExprImplByPlanExpr( + proc.Ctx, + "cast", + []*plan.Expr{ + { + Typ: plan.Type{Id: int32(types.T_int32)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + }, + { + Typ: plan.Type{ + Id: int32(types.T_varchar), + Width: types.MaxVarcharLen, + }, + Expr: &plan.Expr_T{T: &plan.TargetType{}}, + }, + }, + ) + require.NoError(t, err) + return []*plan.Expr{expr} +} + +func TestNewSpillEngineRequiresBudgetGeneration(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + _, err = NewSpillEngine( + SpillEngineConfig{}, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) +} + +func newTestSpillAllocationAccount( + t testing.TB, + limit uint64, + metadataSlots uint64, +) testSpillAllocationAccount { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, metadataSlots) + require.NoError(t, err) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + allocation, err := NewSpillAllocationAccount(account, 2) + require.NoError(t, err) + return testSpillAllocationAccount{ + registry: registry, + account: account, + allocation: allocation, + generation: generation, + } +} + +func finalizeTestSpillAllocationAccount( + t testing.TB, + state testSpillAllocationAccount, +) { + t.Helper() + snapshot := state.account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + _, err := state.registry.Finalize(state.account) + require.NoError(t, err) +} + +func writeSpillAllocationTestFile( + t testing.TB, + bat *batch.Batch, + truncate int, +) *os.File { + t.Helper() + payload := marshalTestSpillRecord(bat) + if truncate > 0 { + payload = payload[:len(payload)-truncate] + } + path := filepath.Join(t.TempDir(), "spill.bin") + require.NoError(t, os.WriteFile(path, payload, 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + return file +} + +func writeSpillAllocationTestRecords( + t testing.TB, + batches ...*batch.Batch, +) *os.File { + t.Helper() + payload := make([]byte, 0) + for _, bat := range batches { + payload = append(payload, marshalTestSpillRecord(bat)...) + } + path := filepath.Join(t.TempDir(), "spill-records.bin") + require.NoError(t, os.WriteFile(path, payload, 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + return file +} + +func TestSpillAllocationAccountDecodedBatchLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-decoded"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 8<<20, 16) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 4, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4}, + ), + testutil.NewVector( + 4, + types.T_varchar.ToType(), + proc.Mp(), + false, + []string{"a", "payload-longer-than-inline", "c", "d"}, + ), + }, nil) + defer source.Clean(proc.Mp()) + + reader := BucketReader{ + fd: writeSpillAllocationTestFile(t, source, 0), + allocation: state.allocation, + } + reuse := batch.NewOffHeapWithSize(0) + decoded, err := reader.ReadBatch(proc, reuse) + require.NoError(t, err) + require.Same(t, reuse, decoded) + require.NotNil(t, decoded.AllocationAccountSelection()) + for _, vec := range decoded.Vecs { + require.NotNil(t, vec.AllocationAccountSelection()) + } + require.Positive(t, state.account.Snapshot().Used) + + reader.Close() + reuse.Clean(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + + reader = BucketReader{ + fd: writeSpillAllocationTestRecords(t, source, source), + mergeRecords: true, + allocation: state.allocation, + } + reuse = batch.NewOffHeapWithSize(0) + decoded, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + require.Equal(t, 2*source.RowCount(), decoded.RowCount()) + require.Positive(t, state.account.Snapshot().Used) + reader.Close() + reuse.Clean(proc.Mp()) + require.Zero(t, state.account.Snapshot().Used) + + reader = BucketReader{ + fd: writeSpillAllocationTestFile(t, source, 1), + allocation: state.allocation, + } + reuse = batch.NewOffHeapWithSize(0) + _, err = reader.ReadBatch(proc, reuse) + require.Error(t, err) + require.Zero(t, state.account.Snapshot().Used) + require.Zero(t, state.registry.LiveAllocationMetadata()) + reader.Close() + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountDecodedReuseRetriesFromCleanRecord(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-decoded-retry"), + ) + defer proc.Free() + makeSource := func(width int) *batch.Batch { + values := make([]string, 1_024) + for i := range values { + values[i] = strings.Repeat("x", width) + } + return testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + len(values), + types.T_varchar.ToType(), + proc.Mp(), + false, + values, + ), + }, nil) + } + first := makeSource(512) + second := makeSource(2_048) + defer first.Clean(proc.Mp()) + defer second.Clean(proc.Mp()) + + measure := newTestSpillAllocationAccount(t, 64<<20, 128) + reader := BucketReader{ + fd: writeSpillAllocationTestRecords(t, first, second), + allocation: measure.allocation, + } + reuse := batch.NewOffHeapWithSize(0) + _, err := reader.ReadBatch(proc, reuse) + require.NoError(t, err) + firstUsed := measure.account.Snapshot().Used + require.Positive(t, firstUsed) + reuse.Clean(proc.Mp()) + require.NoError(t, reuse.SetAllocationAccount(measure.allocation.decoded)) + _, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + secondUsed := measure.account.Snapshot().Used + require.Positive(t, secondUsed) + reuse.Clean(proc.Mp()) + reader.Close() + finalizeTestSpillAllocationAccount(t, measure) + + limit := max(firstUsed, secondUsed) + 128<<10 + require.Less(t, limit, firstUsed+secondUsed) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 128) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + allocation, err := NewSpillAllocationAccount(account, 2) + require.NoError(t, err) + reader = BucketReader{ + fd: writeSpillAllocationTestRecords(t, first, second), + allocation: allocation, + } + reuse = batch.NewOffHeapWithSize(0) + _, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + rejects := generation.Snapshot().RejectCount + _, err = reader.ReadBatch(proc, reuse) + require.NoError(t, err) + require.Equal(t, rejects, generation.Snapshot().RejectCount, + "the local account rejects the overlap before the shared controller") + require.Equal(t, uint64(1), reader.cleanRetries, + "replacement overlap must exercise the clean-record retry") + require.Equal(t, + account.Snapshot().Used, + generation.Used(), + "decoded payloads are charged only by their physical allocations", + ) + reuse.Clean(proc.Mp()) + reader.Close() + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestSpillAllocationAccountScatterScratchLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 1<<20, 64) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + require.NoError(t, err) + + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 8, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := engine.makeBucketWriters("spill_allocation_scatter") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, engine.scatterBatchBounded( + proc, + source, + source.Vecs, + writers, + 0, + false, + analyzer, + )) + require.Len(t, engine.scatterHashValues, source.RowCount()) + require.Len(t, engine.scatterBucketRowIds, source.RowCount()) + snapshot := state.account.Snapshot() + require.Greater( + t, + snapshot.Used, + uint64(source.RowCount()*(8+4)), + ) + require.Greater(t, snapshot.Peak, snapshot.Used) + require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) + var writtenRows int64 + for i := range writers { + writtenRows += writers[i].Rows + } + require.Equal(t, int64(source.RowCount()), writtenRows) + require.Equal(t, snapshot.Used, state.account.Snapshot().Used) + + engine.releaseScatterScratch() + require.Zero(t, state.account.Snapshot().Used) + engine.Cleanup(proc) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountScatterDoesNotReadmitBorrowedSource(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter-source"), + ) + defer proc.Free() + const limit = uint64(8 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 64) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + allocation, err := NewSpillAllocationAccount(account, 2) + require.NoError(t, err) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: generation}, + allocation, + ) + require.NoError(t, err) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 8, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := engine.makeBucketWriters("spill_allocation_scatter_source") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + require.NoError(t, engine.scatterBatchBounded( + proc, + source, + source.Vecs, + writers, + 0, + false, + process.NewAnalyzer(0, false, false, "test"), + )) + snapshot := generation.Snapshot() + require.Equal(t, account.Snapshot().Used, snapshot.Used, + "borrowed input is already live; only new private spill bytes are admitted") + require.Greater(t, snapshot.PeakUsed, snapshot.Used) + + engine.releaseScatterScratch() + engine.Cleanup(proc) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestSpillAllocationAccountMarshalBufferLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-marshal"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 1<<20, 8) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 4, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4}, + ), + }, nil) + defer source.Clean(proc.Mp()) + + accounted, err := state.allocation.newBuffer( + proc.Mp(), + SpillAllocationSiteMarshalBuffer, + ) + require.NoError(t, err) + require.NoError(t, marshalSpillRecordTo(source, accounted)) + require.Equal(t, marshalTestSpillRecord(source), accounted.Bytes()) + used := state.account.Snapshot().Used + require.Positive(t, used) + + accounted.Reset() + require.Zero(t, accounted.Len()) + require.Equal(t, used, state.account.Snapshot().Used) + require.NoError(t, marshalSpillRecordTo(source, accounted)) + require.Equal(t, used, state.account.Snapshot().Used) + + accounted.Free() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountCoalesceAdmissionFallback(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-coalesce-fallback"), + ) + defer proc.Free() + // The record buffer consumes the only metadata slot. Coalescing is an + // optional optimization, so its admission failure must fall back to one + // direct write instead of failing the scatter. + state := newTestSpillAllocationAccount(t, 1<<20, 1) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + require.NoError(t, err) + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + 2, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := engine.makeBucketWriters("spill_allocation_coalesce_fallback") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + + require.NoError(t, engine.appendScatterRecord( + proc, + source, + &writers[0], + 0, + process.NewAnalyzer(0, false, false, "test"), + )) + require.Equal(t, int64(source.RowCount()), writers[0].Rows) + require.Nil(t, engine.scatterAccountedWriteBuffers[0].Bytes()) + require.Equal(t, uint64(1), state.registry.PeakAllocationMetadata()) + + engine.releaseScatterScratch() + engine.Cleanup(proc) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountScatterFailureCleanup(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter-failure"), + ) + defer proc.Free() + const rows = 8 + state := newTestSpillAllocationAccount(t, rows*(8+4), 8) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + require.NoError(t, err) + + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.NewVector( + rows, + types.T_int64.ToType(), + proc.Mp(), + false, + []int64{1, 2, 3, 4, 5, 6, 7, 8}, + ), + }, nil) + defer source.Clean(proc.Mp()) + writers := engine.makeBucketWriters("spill_allocation_scatter_failure") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + err = engine.scatterBatchBounded( + proc, + source, + source.Vecs, + writers, + 0, + false, + process.NewAnalyzer(0, false, false, "test"), + ) + require.Equal( + t, + hashbuild.MemoryPressureMinimumUnit, + hashbuild.MemoryPressureReasonOf(err), + ) + require.Equal(t, uint64(rows*(8+4)), state.account.Snapshot().Used) + + engine.releaseScatterScratch() + require.Zero(t, state.account.Snapshot().Used) + engine.Cleanup(proc) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountScatterReducesUnpublishedInput(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-scatter-reduce"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 80<<10, 128) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + require.NoError(t, err) + values := make([]int64, 8_192) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + writers := engine.makeBucketWriters("spill_allocation_scatter_reduce") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, engine.scatterBatchWithPressure( + proc, + source, + source.Vecs, + writers, + 0, + false, + analyzer, + )) + require.Positive(t, + analyzer.GetOpStats().ExtraStats["JoinSpillInputReductions"]) + require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) + var rows int64 + for i := range writers { + rows += writers[i].Rows + } + require.Equal(t, int64(len(values)), rows) + + engine.releaseScatterScratch() + engine.Cleanup(proc) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) +} + +func TestSpillAllocationAccountExpressionPressureReducesBeforePublication(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-expression-reduce"), + ) + defer proc.Free() + state := newTestSpillAllocationAccount(t, 2<<20, 4_096) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + require.NoError(t, err) + values := make([]int64, 257) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + writers := engine.makeBucketWriters("spill_allocation_expression_reduce") + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "test") + require.NoError(t, engine.scatterEvaluatedBatchWithPressure( + proc, + source, + writers, + 0, + false, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + if current.RowCount() > 32 { + return nil, mpool.ErrAllocationAccountCapacity + } + return current.Vecs, nil + }, + )) + require.Positive(t, + analyzer.GetOpStats().ExtraStats["JoinSpillExpressionInputReductions"]) + require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) + var rows int64 + for i := range writers { + rows += writers[i].Rows + } + require.Equal(t, int64(len(values)), rows, + "evaluation retries must not duplicate or omit published rows") + + engine.releaseScatterScratch() + engine.Cleanup(proc) + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) +} + +// BenchmarkSpillScatterAccounting measures the steady streaming closure with +// the same selected-vector, hash/row-ID, marshal, and coalesce owners used by +// initial and recursive join spill. /dev/null keeps the benchmark bounded and +// retains the write syscall without turning repeated measurements into a disk +// capacity test. +func BenchmarkSpillScatterAccounting(b *testing.B) { + proc := testutil.NewProcessWithMPool(b, "", mpool.MustNewZero()) + defer proc.Free() + values := make([]int64, 4_096) + for i := range values { + values[i] = int64(i) + } + source := testutil.NewBatchWithVectors([]*vector.Vector{ + testutil.MakeInt64Vector(values, nil, proc.Mp()), + }, nil) + defer source.Clean(proc.Mp()) + + state := newTestSpillAllocationAccount(b, 64<<20, 4_096) + engine, err := newSpillEngine( + SpillEngineConfig{Budget: state.generation}, + state.allocation, + ) + if err != nil { + b.Fatal(err) + } + writers := engine.makeBucketWriters("benchmark-discard") + for i := range writers { + writers[i].Fd, err = os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + b.Fatal(err) + } + writers[i].diskReservation, err = state.generation.ReserveSpillDisk(0) + if err != nil { + b.Fatal(err) + } + } + defer func() { + for i := range writers { + writers[i].Close() + } + }() + analyzer := process.NewAnalyzer(0, false, false, "benchmark") + + b.ReportAllocs() + b.SetBytes(int64(source.Size())) + b.ResetTimer() + for range b.N { + if err = engine.scatterBatchWithPressure( + proc, + source, + source.Vecs, + writers, + 0, + false, + analyzer, + ); err != nil { + b.Fatal(err) + } + if err = engine.flushScatterBuffers(proc, writers, analyzer); err != nil { + b.Fatal(err) + } + for i := range writers { + if _, err = writers[i].diskReservation.ReconcileDown(0); err != nil { + b.Fatal(err) + } + writers[i].Rows = 0 + writers[i].Bytes = 0 + } + } + b.StopTimer() + engine.Cleanup(proc) + if state.account.Snapshot().Used != 0 { + b.Fatalf("account used = %d", state.account.Snapshot().Used) + } + finalizeTestSpillAllocationAccount(b, state) +} + +func TestSpillAllocationAccountRebuildAndRecursiveSpillLifecycle(t *testing.T) { + proc := testutil.NewProcessWithMPool( + t, + "", + mpool.MustNew("spill-allocation-rebuild"), + ) + defer proc.Free() + const limit = uint64(64 << 20) + budget := process.MustNewHashBuildBudget(limit, limit) + generation, err := budget.OpenGenerationWithSpillCaps( + 1, + limit, + 1<<30, + 4_096, + ) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.OpenWithController(limit, generation) + require.NoError(t, err) + engine, err := NewSpillEngine( + SpillEngineConfig{ + BuildKeyExprs: makeTestCastKeyExpr(t, proc), + Budget: generation, + SpillThreshold: 100, + NeedsBuildForEmptyProbe: true, + }, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + + values := make([]int32, 5_000) + for i := range values { + values[i] = int32(i) + } + source := makeInt32Batch(proc, values) + fd := writeBuildFile(proc, "accounted_recursive_build", source) + source.Clean(proc.Mp()) + initTestSpillFiles(engine, []*os.File{fd}, int64(len(values))) + analyzer := process.NewAnalyzer(0, false, false, "test") + + respills := 0 + ready := 0 + expressionStorageObserved := false + for steps := 0; engine.HasMoreBuckets(); steps++ { + require.Less(t, steps, 4_096, "recursive spill queue made no progress") + jm, result, rebuildErr := engine.RebuildHashmap(proc, analyzer) + require.NoError(t, rebuildErr) + switch result { + case BucketReSpilled: + respills++ + expressionStorageObserved = expressionStorageObserved || + (len(engine.keyExecs) == 1 && account.Snapshot().Used > 0) + case BucketReady: + ready++ + require.NotNil(t, jm) + jm.Free() + case BucketSkip, BucketEmptyBuild: + require.Nil(t, jm) + default: + require.NotEqual(t, BucketQueueEmpty, result) + } + } + require.Positive(t, respills) + require.Positive(t, ready) + require.True(t, expressionStorageObserved) + require.Positive(t, account.Snapshot().Peak) + + engine.Cleanup(proc) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + require.Zero(t, generation.SpillDiskUsed()) + require.Zero(t, generation.SpillFDUsed()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} diff --git a/pkg/sql/colexec/spillutil/exact_test_helpers_test.go b/pkg/sql/colexec/spillutil/exact_test_helpers_test.go new file mode 100644 index 0000000000000..9a875dd28bdea --- /dev/null +++ b/pkg/sql/colexec/spillutil/exact_test_helpers_test.go @@ -0,0 +1,91 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package spillutil + +import ( + "bytes" + "os" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func newExactTestSpillEngine( + t testing.TB, + cfg SpillEngineConfig, +) *SpillEngine { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) + require.NoError(t, err) + if cfg.Budget == nil { + budget := process.MustNewHashBuildBudget(1<<60, 1<<60) + cfg.Budget, err = budget.OpenGeneration(1) + require.NoError(t, err) + } + account, err := registry.OpenWithController(1<<60, cfg.Budget) + require.NoError(t, err) + engine, err := NewSpillEngine( + cfg, + account, + hashbuild.HashBuildAllocationOwner, + ) + require.NoError(t, err) + return engine +} + +func initTestSpillFiles(engine *SpillEngine, fds []*os.File, rows ...int64) { + if len(rows) != len(fds) { + panic("spill test file/row metadata mismatch") + } + files := make([]*message.SpillFile, len(fds)) + for i, fd := range fds { + if fd != nil { + files[i] = newTestSpillFile(fd, rows[i]) + } + } + engine.InitFromSpilledFiles(files) +} + +func newTestSpillFile(fd *os.File, rows int64) *message.SpillFile { + info, err := fd.Stat() + if err != nil { + panic(err) + } + return message.NewSpillFile(fd, rows, uint64(info.Size()), nil) +} + +type testSpillRecordBuffer struct { + bytes.Buffer +} + +func (b *testSpillRecordBuffer) EnsureCapacity(required int) error { + if b.Cap() < required { + b.Buffer = *bytes.NewBuffer(make([]byte, 0, required)) + } + return nil +} + +func marshalTestSpillRecord(bat *batch.Batch) []byte { + var encoded testSpillRecordBuffer + if err := marshalSpillRecordTo(bat, &encoded); err != nil { + panic(err) + } + return bytes.Clone(encoded.Bytes()) +} diff --git a/pkg/sql/colexec/spillutil/join_spill.go b/pkg/sql/colexec/spillutil/join_spill.go index 70f0aa7f1e468..ef881bea66daa 100644 --- a/pkg/sql/colexec/spillutil/join_spill.go +++ b/pkg/sql/colexec/spillutil/join_spill.go @@ -16,8 +16,7 @@ package spillutil import ( - "bufio" - "bytes" + "encoding/binary" "errors" "fmt" "io" @@ -53,7 +52,6 @@ const ( // spill records without retaining the pre-admission unmarshal estimate // for a large record until the reader closes. The additive bound makes the // long-lived charge independent of the largest serialized payload seen. - decodedBatchLeaseSlack = 1 << 20 ) // SpillBucket holds file descriptors for one spilled bucket. @@ -77,621 +75,245 @@ func checkSpillCanceled(proc *process.Process) error { } } -// BucketReader reads serialized batch records from an fd. +// BucketReader decodes one move-only spill file. A pending header replaces +// bufio.Peek, so all data-scaled decode storage is owned by accounted vectors +// instead of an untracked Go-heap buffer. type BucketReader struct { - fd *os.File - reader *bufio.Reader - buf [16]byte - budget *process.HashBuildBudgetGeneration - reservation *process.HashBuildReservation - batchToken *process.HashBuildReservation - batchCharge uint64 - spillFile *message.SpillFile - mergeRecords bool + fd *os.File + header [16]byte + headerPending bool + spillFile *message.SpillFile + mergeRecords bool + allocation *SpillAllocationAccount + cleanRetries uint64 + schema []types.Type } -func (r *BucketReader) ReadBatch(proc *process.Process, reuseBat *batch.Batch) (*batch.Batch, error) { +func (r *BucketReader) ReadBatch( + proc *process.Process, + reuseBat *batch.Batch, +) (*batch.Batch, error) { if err := checkSpillCanceled(proc); err != nil { return nil, err } if r.fd == nil { return nil, io.EOF } - if r.reader == nil { - r.reader = bufio.NewReaderSize(r.fd, 4*1024*1024) + if reuseBat == nil { + return nil, moerr.NewInvalidInput( + proc.Ctx, + "spill batch reader requires a reuse batch", + ) + } + if r.allocation == nil { + return nil, mpool.ErrAllocationAccountInvalid } - _, token, charge, err := r.readBatchRecord(proc, reuseBat, r.batchToken, r.batchCharge, true) - if err != nil { - r.releaseReadBatch(proc, reuseBat, token) + if err := reuseBat.SetAllocationAccount(r.allocation.decoded); err != nil { + return nil, err + } + if _, err := r.readBatchRecord(proc, reuseBat); err != nil { + reuseBat.Clean(proc.Mp()) return nil, err } - oldToken := r.batchToken - r.batchToken = token - r.batchCharge = charge - if oldToken != nil && oldToken != token { - oldToken.Release() + if err := r.validateSchema(proc, reuseBat); err != nil { + reuseBat.Clean(proc.Mp()) + return nil, err } if !r.mergeRecords { if err := checkSpillCanceled(proc); err != nil { - r.releaseReadBatch(proc, reuseBat, nil) + reuseBat.Clean(proc.Mp()) return nil, err } return reuseBat, nil } - // Merge adjacent records up to the bounded historical batch payload. - // This preserves dedup/outer-join behaviour across small source batches - // without retaining one selected batch per bucket during scatter. + for reuseBat.RowCount() < colexec.DefaultBatchSize { - if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) + nextRows, err := r.peekRecordRows(proc) + if err == io.EOF { + break } - header, err := r.reader.Peek(16) if err != nil { - if err == io.EOF && len(header) == 0 { - break - } - if err == io.EOF { - err = io.ErrUnexpectedEOF - } - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) - } - if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) + return nil, r.mergeReadError(proc, reuseBat, nil, err) } - nextRows := types.DecodeInt64(header[:8]) - nextBatchSize := types.DecodeInt64(header[8:16]) - if nextRows < 0 || nextBatchSize < 0 { - return nil, r.mergeReadError( - proc, - reuseBat, - nil, - nil, - moerr.NewInternalError(proc.Ctx, "negative spill batch header"), - ) - } - // A source record is an indivisible ownership and budget unit. Leave it - // for the next ReadBatch rather than consuming it and growing this batch - // beyond the advertised merge bound. if nextRows > int64(colexec.DefaultBatchSize-reuseBat.RowCount()) { break } - next := batch.NewOffHeapWithSize(0) - _, nextToken, _, err := r.readBatchRecord(proc, next, nil, 0, false) + next, err := newSpillBatch(0, r.allocation.decoded) if err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) + return nil, r.mergeReadError(proc, reuseBat, nil, err) } - if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) - } - var mergeToken *process.HashBuildReservation - if r.budget != nil { - // Keep the current destination (O) and the source record (N) live - // while admitting the final destination (D). UnionBatch may retain - // rounded capacities larger than O+N, so reserving O+N here is not a - // safe admission bound. - predicted, ok := predictMergedRetainedBytes(reuseBat, next) - if !ok { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, process.ErrHashBuildBudgetInvalid) - } - mergeToken, err = r.budget.Reserve(predicted) - if err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err) - } + if _, err := r.readBatchRecord(proc, next); err != nil { + return nil, r.mergeReadError(proc, reuseBat, next, err) } - if len(reuseBat.Vecs) != len(next.Vecs) { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, process.ErrHashBuildBudgetInvalid, mergeToken) + if err := r.validateSchema(proc, next); err != nil { + return nil, r.mergeReadError(proc, reuseBat, next, err) } for i := range next.Vecs { - if err := reuseBat.Vecs[i].UnionBatch(next.Vecs[i], 0, next.RowCount(), nil, proc.Mp()); err != nil { - return nil, r.mergeReadError(proc, reuseBat, next, nextToken, err, mergeToken) + if err := reuseBat.Vecs[i].UnionBatch( + next.Vecs[i], + 0, + next.RowCount(), + nil, + proc.Mp(), + ); err != nil { + return nil, r.mergeReadError(proc, reuseBat, next, err) } } - reuseBat.SetRowCount(reuseBat.RowCount() + next.RowCount()) + reuseBat.AddRowCount(next.RowCount()) next.Clean(proc.Mp()) - if mergeToken != nil { - actual, ok := batchRetainedBytes(reuseBat) - if !ok { - return nil, r.mergeReadError(proc, reuseBat, nil, nextToken, process.ErrHashBuildBudgetInvalid, mergeToken) - } - if err := reconcileReadReservation(mergeToken, actual); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nextToken, err, mergeToken) - } - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 - } - if nextToken != nil { - nextToken.Release() - } - r.batchToken = mergeToken - r.batchCharge = actual - } } if err := checkSpillCanceled(proc); err != nil { - return nil, r.mergeReadError(proc, reuseBat, nil, nil, err) + return nil, r.mergeReadError(proc, reuseBat, nil, err) } return reuseBat, nil } -// mergeReadError unwinds all ownership acquired while appending a source -// record. The destination may have been partially mutated by UnionBatch, so it -// is cleaned as well. Reservations are exactly-once tokens; releasing an -// already released token is harmless and keeps every error path symmetric. -func (r *BucketReader) mergeReadError(proc *process.Process, dst, src *batch.Batch, srcToken *process.HashBuildReservation, err error, extra ...*process.HashBuildReservation) error { - if src != nil { - src.Clean(proc.Mp()) - } - if srcToken != nil { - srcToken.Release() +func (r *BucketReader) validateSchema( + proc *process.Process, + bat *batch.Batch, +) error { + if bat == nil { + return process.ErrHashBuildBudgetInvalid } - for _, token := range extra { - if token != nil { - token.Release() + if r.schema == nil { + r.schema = make([]types.Type, len(bat.Vecs)) + for i, vec := range bat.Vecs { + if vec == nil { + return moerr.NewInternalError(proc.Ctx, "nil vector in spill batch") + } + r.schema[i] = *vec.GetType() } - } - if dst != nil { - dst.Clean(proc.Mp()) - } - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 - } - return err -} - -func addUint64(a, b uint64) (uint64, bool) { - if a > math.MaxUint64-b { - return 0, false - } - return a + b, true -} - -func batchRetainedBytes(bat *batch.Batch) (uint64, bool) { - if bat == nil || bat.RowCount() < 0 { - return 0, false - } - actual := uint64(bat.Allocated()) - metadata, ok := batchRetainedMetadataBytes(uint64(bat.RowCount()), uint64(len(bat.Vecs))) - if !ok { - return 0, false - } - return addUint64(actual, metadata) -} - -func batchRetainedMetadataBytes(rows, cols uint64) (uint64, bool) { - if cols > (math.MaxUint64-16)/8 { - return 0, false - } - metadata := uint64(16) + cols*8 - if rows > 0 && metadata > math.MaxUint64/rows { - return 0, false - } - return rows * metadata, true -} - -// reconcileReadReservation shrinks a conservative read reservation to the -// retained batch size. ReconcileDown already validates that actual does not -// exceed the reservation, so callers do not need a separate Size call (and a -// second acquisition of the shared hash-build budget mutex). -func reconcileReadReservation(token *process.HashBuildReservation, actual uint64) error { - if token == nil { return nil } - if _, err := token.ReconcileDown(actual); err != nil { - if errors.Is(err, process.ErrHashBuildReservationUpward) { - return process.ErrHashBuildBudgetInvalid - } - return err - } - return nil -} - -// predictMergedRetainedBytes computes the retained upper bound after the exact -// full-record UnionBatch append used by ReadBatch. It mirrors Vector.extend's -// data-cap growth and UnionBatch's varlen fast path (which appends a complete -// non-const source area in one operation). No destination mutation is performed. -func predictMergedRetainedBytes(dst, src *batch.Batch) (uint64, bool) { - if dst == nil || src == nil || dst.RowCount() < 0 || src.RowCount() < 0 || len(dst.Vecs) != len(src.Vecs) { - return 0, false - } - oldRows, ok := intToUint64(dst.RowCount()) - if !ok { - return 0, false - } - srcRows, ok := intToUint64(src.RowCount()) - if !ok { - return 0, false + if len(bat.Vecs) != len(r.schema) { + return moerr.NewInternalError(proc.Ctx, "spill batch schema changed") } - mergedRows, ok := addUint64(oldRows, srcRows) - if !ok || mergedRows > uint64(maxIntValue()) { - return 0, false - } - - var allocated uint64 - for i := range dst.Vecs { - dv, sv := dst.Vecs[i], src.Vecs[i] - if dv == nil || sv == nil || *dv.GetType() != *sv.GetType() || dv.Length() != dst.RowCount() || sv.Length() != src.RowCount() { - return 0, false - } - typeSize := dv.GetType().TypeSize() - if typeSize < 0 { - return 0, false - } - dataRequired, ok := mulUint64(mergedRows, uint64(typeSize)) - if !ok || dataRequired > uint64(math.MaxInt64) { - return 0, false + for i, vec := range bat.Vecs { + if vec == nil || !r.schema[i].Eq(*vec.GetType()) { + return moerr.NewInternalError(proc.Ctx, "spill batch schema changed") } - dataCap, ok := predictedCapacity(cap(dv.GetData()), dataRequired) - if !ok { - return 0, false - } - if allocated, ok = addUint64(allocated, dataCap); !ok { - return 0, false - } - - if !dv.GetType().IsVarlen() { - continue - } - areaAdd, ok := mergedVarlenAreaAdd(sv, srcRows) - if !ok { - return 0, false - } - areaRequired, ok := addUint64(uint64(len(dv.GetArea())), areaAdd) - if !ok || areaRequired > uint64(math.MaxInt64) { - return 0, false - } - areaCap, ok := predictedCapacity(cap(dv.GetArea()), areaRequired) - if !ok { - return 0, false - } - if allocated, ok = addUint64(allocated, areaCap); !ok { - return 0, false - } - } - - cols := uint64(len(dst.Vecs)) - if cols > (math.MaxUint64-16)/8 { - return 0, false - } - metadata := 16 + cols*8 - rowMetadata, ok := mulUint64(mergedRows, metadata) - if !ok { - return 0, false } - return addUint64(allocated, rowMetadata) + return nil } -func mergedVarlenAreaAdd(src *vector.Vector, rows uint64) (uint64, bool) { - if src == nil || !src.GetType().IsVarlen() { - return 0, false - } - if rows == 0 { - return 0, true +func (r *BucketReader) peekRecordRows(proc *process.Process) (int64, error) { + if err := checkSpillCanceled(proc); err != nil { + return 0, err } - if src.IsConst() { - if src.IsConstNull() { - return 0, true - } - // UnionBatch materializes one const value and broadcasts its header. An - // inline value needs no area; a non-inline value appends exactly once. - if len(src.GetData()) < src.GetType().TypeSize() { - return 0, false - } - values := vector.MustFixedColNoTypeCheck[types.Varlena](src) - if len(values) != 1 { - return 0, false - } - value := &values[0] - if value.IsSmall() { - return 0, true - } - off, length := value.OffsetLen() - end, ok := addUint64(uint64(off), uint64(length)) - if !ok || end > uint64(len(src.GetArea())) { - return 0, false + if !r.headerPending { + if _, err := io.ReadFull(r.fd, r.header[:]); err != nil { + return 0, err } - return uint64(length), true - } - - // The full-record fast path copies the complete source area once, including - // stale bytes. Header validation remains UnionBatch's responsibility; avoid - // adding another per-row scan on the spill rebuild hot path. - return uint64(len(src.GetArea())), true -} - -func predictedCapacity(oldCap int, required uint64) (uint64, bool) { - if oldCap < 0 || uint64(oldCap) > uint64(math.MaxInt64) || required > uint64(math.MaxInt64) { - return 0, false - } - if required <= uint64(oldCap) { - return uint64(oldCap), true - } - cap, ok := mpool.GrowCapacity(int64(oldCap), int64(required)) - if !ok || cap < 0 { - return 0, false - } - return uint64(cap), true -} - -func intToUint64(v int) (uint64, bool) { - if v < 0 { - return 0, false - } - return uint64(v), true -} - -func mulUint64(a, b uint64) (uint64, bool) { - if a != 0 && b > math.MaxUint64/a { - return 0, false - } - return a * b, true -} - -func batchPayloadWithAllocationSlack(payload, columns uint64) (uint64, bool) { - const perVectorAllocationSlack = uint64(64 << 10) - if columns >= math.MaxUint64/perVectorAllocationSlack { - return 0, false - } - allocationSlack := (columns + 1) * perVectorAllocationSlack - if payload > math.MaxUint64-allocationSlack { - return 0, false - } - return payload + allocationSlack, true -} - -func decodedBatchProjectedBytes(payload uint64, rows int64, columns int32) (uint64, bool) { - if rows < 0 || columns < 0 { - return 0, false - } - projected, ok := batchPayloadWithAllocationSlack(payload, uint64(columns)) - if !ok { - return 0, false - } - metadata, ok := batchRetainedMetadataBytes(uint64(rows), uint64(columns)) - if !ok { - return 0, false - } - return addUint64(projected, metadata) -} - -func decodedBatchReusePeakBytes(retained, projected, payload uint64) (uint64, bool) { - // For large buffers mpool.Grow follows Go's 1.25x growth policy. The old - // allocation remains live until the replacement is allocated and copied. - // Small-buffer doubling is bounded by the per-vector slack already included - // in projected. - growthSlack := payload / 4 - if payload%4 != 0 { - growthSlack++ - } - newAllocation, ok := addUint64(projected, growthSlack) - if !ok { - return 0, false - } - return addUint64(retained, newAllocation) -} - -func maxIntValue() int { - return int(^uint(0) >> 1) -} - -func marshalSpillRecordGrowBytes(bat *batch.Batch) (uint64, bool) { - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size + r.headerPending = true } - columns := uint64(len(bat.Vecs)) - if columns > (math.MaxUint64-24)/128 { - return 0, false - } - return addUint64(base, columns*128+24) -} - -func (r *BucketReader) releaseReadBatch(proc *process.Process, bat *batch.Batch, token *process.HashBuildReservation) { - if bat != nil { - bat.Clean(proc.Mp()) - } - if token != nil { - token.Release() - } - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 + rows := types.DecodeInt64(r.header[:8]) + batchSize := types.DecodeInt64(r.header[8:]) + if rows < 0 || batchSize < 0 { + return 0, moerr.NewInternalError( + proc.Ctx, + "negative spill batch header", + ) } + return rows, nil } func (r *BucketReader) readBatchRecord( proc *process.Process, reuseBat *batch.Batch, - token *process.HashBuildReservation, - charge uint64, - retainLease bool, -) (*batch.Batch, *process.HashBuildReservation, uint64, error) { +) (*batch.Batch, error) { if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err + return nil, err } - if _, err := io.ReadFull(r.reader, r.buf[:]); err != nil { - if err == io.EOF { - return nil, token, charge, io.EOF + if !r.headerPending { + if _, err := io.ReadFull(r.fd, r.header[:]); err != nil { + return nil, err } - return nil, token, charge, err } - cnt := types.DecodeInt64(r.buf[:8]) - batchSize := types.DecodeInt64(r.buf[8:16]) + r.headerPending = false + cnt := types.DecodeInt64(r.header[:8]) + batchSize := types.DecodeInt64(r.header[8:]) if cnt < 0 || batchSize < 0 { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "negative spill batch header") + return nil, moerr.NewInternalError( + proc.Ctx, + "negative spill batch header", + ) } - if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err + payloadOffset, err := r.fd.Seek(0, io.SeekCurrent) + if err != nil { + return nil, err } - if r.budget != nil { - payload := uint64(batchSize) - if payload > uint64(maxIntValue())-(64<<10) { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - // The batch payload starts with row count and vector count. Peek only the - // fixed header so allocator rounding can be bounded per decoded vector - // before UnmarshalFromReader performs any allocation. - header, err := r.reader.Peek(12) - if err != nil { - return nil, token, charge, err - } - rows := types.DecodeInt64(header[:8]) - columns := types.DecodeInt32(header[8:12]) - if rows != cnt { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "row count mismatch") - } - projected, ok := decodedBatchProjectedBytes(payload, rows, columns) - if !ok { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - // The serialized payload already includes every vector's data, area, null - // bitmap, and headers. Reserve one decoded payload plus bounded allocator - // slack, then reconcile to the retained capacities reported by the batch. - // Multiplying the complete payload rejects large spill records before - // UnmarshalFromReader can establish their actual retained footprint. - if token == nil { - // A caller-provided reuse batch has no budget ownership on the first - // read. Drop it before admitting the decoded payload. - reuseBat.Clean(proc.Mp()) - var err error - token, err = r.budget.Reserve(projected) - if err != nil { - return nil, nil, 0, err - } - charge = projected - } else { - // Reusing vectors can briefly keep their old allocation alive while - // mpool.Grow allocates the replacement. Admit one complete decoded - // payload above the retained lease before unmarshal. If that transient - // peak does not fit, release the old batch and decode from a clean - // batch so a valid single payload is not rejected. - retained, retainedOK := batchRetainedBytes(reuseBat) - peak, peakOK := decodedBatchReusePeakBytes(retained, projected, payload) - var growErr error - if retainedOK && peakOK && peak > charge { - growErr = token.Grow(peak - charge) - } - if growErr != nil && - !errors.Is(growErr, process.ErrHashBuildBudgetAdmission) { - return nil, token, charge, growErr - } - if !retainedOK || !peakOK || growErr != nil { - reuseBat.Clean(proc.Mp()) - token.Release() - token = nil - var err error - token, err = r.budget.Reserve(projected) - if err != nil { - return nil, nil, 0, err - } - charge = projected - } else if peak > charge { - charge = peak - } + decode := func() (io.LimitedReader, error) { + reuseBat.CleanOnlyData() + if err := checkSpillCanceled(proc); err != nil { + return io.LimitedReader{}, err } + limited := io.LimitedReader{R: r.fd, N: batchSize} + return limited, reuseBat.UnmarshalFromReaderWithGrouping(&limited, proc.Mp()) } - - reuseBat.CleanOnlyData() - if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err - } - - limitReader := io.LimitedReader{R: r.reader, N: batchSize} - if err := reuseBat.UnmarshalFromReader(&limitReader, proc.Mp()); err != nil { - return nil, token, charge, err + limited, decodeErr := decode() + if decodeErr != nil && mpool.IsRetryableAllocationCapacity(decodeErr) { + // Reuse owns the old capacity while a replacement is allocated. Release + // it, rewind the unpublished record, and retry against actual allocation. + reuseBat.Clean(proc.Mp()) + if err := reuseBat.SetAllocationAccount(r.allocation.decoded); err != nil { + return nil, err + } + if _, err := r.fd.Seek(payloadOffset, io.SeekStart); err != nil { + return nil, err + } + r.cleanRetries++ + limited, decodeErr = decode() } - if err := checkSpillCanceled(proc); err != nil { - return nil, token, charge, err + if decodeErr != nil { + return nil, decodeErr } - - // Verify the batch unmarshal consumed exactly batchSize bytes. - if limitReader.N > 0 { - return nil, token, charge, moerr.NewInternalErrorf(proc.Ctx, "batch unmarshal did not consume all bytes: %d remaining", limitReader.N) + if limited.N != 0 { + return nil, moerr.NewInternalErrorf( + proc.Ctx, + "batch unmarshal did not consume all bytes: %d remaining", + limited.N, + ) } - - // Read magic (8 bytes) - if _, err := io.ReadFull(r.reader, r.buf[:8]); err != nil { - return nil, token, charge, err + if _, err := io.ReadFull(r.fd, r.header[:8]); err != nil { + return nil, err } - if types.DecodeUint64(r.buf[:8]) != SpillMagic { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "corrupted spill file") + if types.DecodeUint64(r.header[:8]) != SpillMagic { + return nil, moerr.NewInternalError(proc.Ctx, "corrupted spill file") } - if reuseBat.RowCount() != int(cnt) { - return nil, token, charge, moerr.NewInternalError(proc.Ctx, "row count mismatch") - } - if token != nil { - actual, ok := batchRetainedBytes(reuseBat) - if !ok || actual > charge { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - target := actual - if retainLease { - withSlack, ok := addUint64(actual, decodedBatchLeaseSlack) - if !ok { - return nil, token, charge, process.ErrHashBuildBudgetInvalid - } - if withSlack < charge { - target = withSlack - } else { - target = charge - } - } - if target < charge { - if err := reconcileReadReservation(token, target); err != nil { - return nil, token, charge, err - } - charge = target - } + return nil, moerr.NewInternalError(proc.Ctx, "row count mismatch") } - return reuseBat, token, charge, nil + return reuseBat, checkSpillCanceled(proc) } -func (r *BucketReader) ResetForFd(fd *os.File) { - r.closeCurrentFile() - if fd == nil { - return +func (r *BucketReader) mergeReadError( + proc *process.Process, + dst *batch.Batch, + src *batch.Batch, + err error, +) error { + if src != nil { + src.Clean(proc.Mp()) } - r.fd = fd - if r.reader == nil { - r.reader = bufio.NewReaderSize(fd, 4*1024*1024) - } else { - r.reader.Reset(fd) + if dst != nil { + dst.Clean(proc.Mp()) } + return err } -func (r *BucketReader) ResetForSpillFile(file *message.SpillFile) { +func (r *BucketReader) ResetForSpillFile(file *message.SpillFile) error { r.closeCurrentFile() if file == nil { - return - } - r.spillFile = file - r.fd = file.File() - if r.reader == nil { - r.reader = bufio.NewReaderSize(r.fd, 4*1024*1024) - } else { - r.reader.Reset(r.fd) - } -} - -// EnsureBuffer admits the reader's fixed backing allocation before creating -// the bufio.Reader. Rebuild and probe readers are shared one at a time, so a -// single reservation bounds their peak rather than charging one per bucket. -func (r *BucketReader) EnsureBuffer(budget *process.HashBuildBudgetGeneration) error { - r.budget = budget - if r.reader != nil || budget == nil { return nil } - const size = uint64(64 << 10) - token, err := budget.Reserve(size) - if err != nil { + if err := file.Validate(); err != nil { + _ = file.Close() return err } - r.reservation = token - r.reader = bufio.NewReaderSize(nil, int(size)) + r.spillFile = file + r.fd = file.File() + r.headerPending = false + r.schema = nil return nil } @@ -703,25 +325,15 @@ func (r *BucketReader) closeCurrentFile() { r.fd = nil } if r.fd != nil { - r.fd.Close() + _ = r.fd.Close() r.fd = nil } + r.headerPending = false + r.schema = nil } func (r *BucketReader) Close() { r.closeCurrentFile() - if r.batchToken != nil { - r.batchToken.Release() - r.batchToken = nil - r.batchCharge = 0 - } - if r.reservation != nil { - r.reservation.Release() - r.reservation = nil - } - // A retained bufio.Reader capacity must remain charged. Drop it when the - // reservation is released; ResetForFd/ResetForSpillFile keep both alive. - r.reader = nil } // BucketWriter writes serialized batch records to an fd. @@ -781,23 +393,6 @@ func (w *BucketWriter) Close() { } } -func (w *BucketWriter) HandOffFd() *os.File { - if w.Fd == nil { - return nil - } - // A raw descriptor cannot carry accounting ownership. Budgeted writers - // must use handOffSpillFile; retain ownership here so Close can unwind it. - if w.fdReservation != nil || w.diskReservation != nil { - return nil - } - if _, err := w.Fd.Seek(0, io.SeekStart); err != nil { - return nil - } - fd := w.Fd - w.Fd = nil - return fd -} - func (w *BucketWriter) handOffSpillFile() (*message.SpillFile, error) { if w.Fd == nil { return nil, nil @@ -832,50 +427,84 @@ func MakeBucketWriters(prefix string) []BucketWriter { return writers } -// FlushBucketBatch writes bat to w, creating the spill file on first write. -// If analyzer is non-nil, spill bytes/rows are tracked. -func FlushBucketBatch(proc *process.Process, bat *batch.Batch, w *BucketWriter, bucketBuf *bytes.Buffer, analyzer process.Analyzer) error { - if bat == nil || bat.RowCount() == 0 { - return nil - } - // Serialize before creating the file. This admits marshal scratch and the - // exact disk extent before CreateAndRemoveFile/write, so a rejected write - // leaves both the writer and source batch intact. - cnt := int64(bat.RowCount()) - if err := marshalSpillRecord(bat, bucketBuf); err != nil { - return err - } - return writeBucketPayload(proc, bucketBuf.Bytes(), cnt, w, analyzer) +type spillRecordBuffer interface { + io.Writer + Bytes() []byte + EnsureCapacity(int) error + Len() int + Reset() } -func marshalSpillRecord(bat *batch.Batch, buf *bytes.Buffer) error { +func marshalSpillRecordTo( + bat *batch.Batch, + buf spillRecordBuffer, +) error { if bat == nil || bat.RowCount() == 0 { return nil } cnt := int64(bat.RowCount()) buf.Reset() - grow, ok := marshalSpillRecordGrowBytes(bat) - if !ok || grow > uint64(maxIntValue()) { + batchSize, err := bat.MarshalBinaryWithGroupingSize() + if err != nil || batchSize > math.MaxInt-24 { + if err != nil { + return err + } return process.ErrHashBuildBudgetInvalid } - if uint64(buf.Cap()) < grow { - // Allocate the final serialization capacity in one step. Retaining a - // smaller bytes.Buffer while it grows geometrically would invalidate the - // single-payload admission estimate. - *buf = *bytes.NewBuffer(make([]byte, 0, int(grow))) + if err := buf.EnsureCapacity(batchSize + 24); err != nil { + return err + } + if err := writeSpillRecordInt64(buf, cnt); err != nil { + return err } - buf.Write(types.EncodeInt64(&cnt)) batchSizePos := buf.Len() - var zero int64 - buf.Write(types.EncodeInt64(&zero)) + if err := writeSpillRecordInt64(buf, 0); err != nil { + return err + } batchStart := buf.Len() - if _, err := bat.MarshalBinaryWithBuffer(buf, false); err != nil { + if err := bat.MarshalBinaryWithGroupingTo(buf); err != nil { + return err + } + serializedSize := int64(buf.Len() - batchStart) + if setter, ok := buf.(interface{ SetInt64(int, int64) error }); ok { + if err := setter.SetInt64(batchSizePos, serializedSize); err != nil { + return err + } + } else { + binary.NativeEndian.PutUint64( + buf.Bytes()[batchSizePos:batchSizePos+8], + uint64(serializedSize), + ) + } + return writeSpillRecordUint64(buf, uint64(SpillMagic)) +} + +func writeSpillRecordInt64(w io.Writer, value int64) error { + if typed, ok := w.(interface{ WriteInt64(int64) error }); ok { + return typed.WriteInt64(value) + } + var data [8]byte + binary.NativeEndian.PutUint64(data[:], uint64(value)) + return writeSpillRecordBytes(w, data[:]) +} + +func writeSpillRecordUint64(w io.Writer, value uint64) error { + if typed, ok := w.(interface{ WriteUint64(uint64) error }); ok { + return typed.WriteUint64(value) + } + var data [8]byte + binary.NativeEndian.PutUint64(data[:], value) + return writeSpillRecordBytes(w, data[:]) +} + +func writeSpillRecordBytes(w io.Writer, value []byte) error { + written, err := w.Write(value) + if err != nil { return err } - batchSize := int64(buf.Len() - batchStart) - copy(buf.Bytes()[batchSizePos:batchSizePos+8], types.EncodeInt64(&batchSize)) - magic := uint64(SpillMagic) - buf.Write(types.EncodeUint64(&magic)) + if written != len(value) { + return io.ErrShortWrite + } return nil } @@ -885,6 +514,9 @@ func writeBucketPayload(proc *process.Process, payload []byte, rows int64, w *Bu if w == nil || len(payload) == 0 { return process.ErrHashBuildBudgetInvalid } + if w.Budget == nil { + return process.ErrHashBuildBudgetInvalid + } if err := checkSpillCanceled(proc); err != nil { return err } @@ -901,32 +533,26 @@ func writeBucketPayload(proc *process.Process, payload []byte, rows int64, w *Bu } _, _ = w.diskReservation.ReconcileDown(oldDiskSize) } - if w.Budget != nil { - // Keep one growable disk token per file so bookkeeping remains bounded - // even when the input arrives as millions of tiny batches. - if w.diskReservation == nil { - diskToken, err := w.Budget.ReserveSpillDisk(uint64(len(payload))) - if err != nil { - return err - } - w.diskReservation = diskToken - newDiskToken = true - } else { - oldDiskSize = w.diskReservation.Size() - if err := w.diskReservation.Grow(uint64(len(payload))); err != nil { - return err - } + // Keep one growable disk token per file so bookkeeping remains bounded + // even when the input arrives as millions of tiny batches. + if w.diskReservation == nil { + diskToken, err := w.Budget.ReserveSpillDisk(uint64(len(payload))) + if err != nil { + return err + } + w.diskReservation = diskToken + newDiskToken = true + } else { + oldDiskSize = w.diskReservation.Size() + if err := w.diskReservation.Grow(uint64(len(payload))); err != nil { + return err } } if !w.Created() { - var fdToken *process.HashBuildSpillFDReservation - var err error - if w.Budget != nil { - fdToken, err = w.Budget.ReserveSpillFD(1) - if err != nil { - rollbackDisk() - return err - } + fdToken, err := w.Budget.ReserveSpillFD(1) + if err != nil { + rollbackDisk() + return err } fs, err := w.getSpillFileService(proc) if err != nil { @@ -966,11 +592,6 @@ func writeBucketPayload(proc *process.Process, payload []byte, rows int64, w *Bu return nil } -// hashCombine merges a new hash value into a running hash state (Boost-style). -func hashCombine(h, val uint64) uint64 { - return keycodec.HashCombine(h, val) -} - // ComputeXXHash evaluates key vectors and computes XXHash64 values using // column-at-a-time processing for better cache locality. seed initialises every // hash slot so different spill depths produce different bucket distributions. @@ -982,7 +603,10 @@ func ComputeXXHash(keyVecs []*vector.Vector, hashValues []uint64, seed uint64) { // id array in two linear passes. This replaces the historical bucket-by-bucket // scan of hashValues (which revisited every row once for each bucket). func classifyRows(hashValues []uint64, bucketCount int, shift uint64, rowIDs []int32, counts []int32, offsets []int32) error { - if bucketCount <= 0 || bucketCount&(bucketCount-1) != 0 || shift >= 64 || len(rowIDs) < len(hashValues) || len(counts) < bucketCount || len(offsets) < bucketCount+1 { + if bucketCount <= 0 || bucketCount > SpillNumBuckets || + bucketCount&(bucketCount-1) != 0 || shift >= 64 || + len(rowIDs) < len(hashValues) || len(counts) < bucketCount || + len(offsets) < bucketCount+1 { return process.ErrHashBuildBudgetInvalid } for i := 0; i < bucketCount; i++ { @@ -997,363 +621,264 @@ func classifyRows(hashValues []uint64, bucketCount int, shift uint64, rowIDs []i offsets[i+1] = offsets[i] + counts[i] } var writePos [SpillNumBuckets]int32 - if bucketCount <= len(writePos) { - copy(writePos[:bucketCount], offsets[:bucketCount]) - for row, hash := range hashValues { - bucket := int((hash >> shift) & mask) - pos := writePos[bucket] - rowIDs[pos] = int32(row) - writePos[bucket] = pos + 1 - } - return nil - } - // SpillNumBuckets is the production fanout. Keep the helper correct for - // callers using another power-of-two fanout without allocating a second - // row-id structure. - positions := make([]int32, bucketCount) - copy(positions, offsets[:bucketCount]) + copy(writePos[:bucketCount], offsets[:bucketCount]) for row, hash := range hashValues { bucket := int((hash >> shift) & mask) - pos := positions[bucket] + pos := writePos[bucket] rowIDs[pos] = int32(row) - positions[bucket] = pos + 1 + writePos[bucket] = pos + 1 } return nil } -// scatterImpl is the internal implementation that accepts reusable buffers. -func scatterImpl( +// scatterBatchBounded writes one bucket at a time. It keeps one selected batch +// and one row-id slice, flushing them before advancing to the next bucket. +func (e *SpillEngine) scatterBatchBounded( proc *process.Process, bat *batch.Batch, keyVecs []*vector.Vector, writers []BucketWriter, - buffers []*batch.Batch, - seed uint64, - bucketBuf *bytes.Buffer, + partitionLevel uint64, + sourceAlreadyCharged bool, analyzer process.Analyzer, - reuseHashValues *[]uint64, - reuseBucketRowIds *[][]int32, -) error { - rowCount := bat.RowCount() - if rowCount == 0 { +) (retErr error) { + if bat == nil || bat.RowCount() == 0 { return nil } - - var hashValues []uint64 - if reuseHashValues != nil && cap(*reuseHashValues) >= rowCount { - hashValues = (*reuseHashValues)[:rowCount] - } else { - hashValues = make([]uint64, rowCount) - if reuseHashValues != nil { - *reuseHashValues = hashValues - } + if err := checkSpillCanceled(proc); err != nil { + return err } - ComputeXXHash(keyVecs, hashValues, seed) - - if len(writers) == 0 || len(writers)&(len(writers)-1) != 0 { + if len(writers) == 0 || len(writers) > SpillNumBuckets { return process.ErrHashBuildBudgetInvalid } - // Build one contiguous row-id array, then expose each bucket as a slice of - // that array for compatibility with the buffered path. - var bucketRowIds [][]int32 - if reuseBucketRowIds != nil { - bucketRowIds = *reuseBucketRowIds - if cap(bucketRowIds) < len(writers) { - bucketRowIds = make([][]int32, len(writers)) - *reuseBucketRowIds = bucketRowIds - } else { - bucketRowIds = bucketRowIds[:len(writers)] - } - } else { - bucketRowIds = make([][]int32, len(writers)) - } - var rowIDs []int32 - if len(bucketRowIds) > 0 && cap(bucketRowIds[0]) >= rowCount { - rowIDs = bucketRowIds[0][:rowCount] - } else { - rowIDs = make([]int32, rowCount) + if e.allocation == nil { + return mpool.ErrAllocationAccountInvalid } - var countsFixed [SpillNumBuckets]int32 - var offsetsFixed [SpillNumBuckets + 1]int32 - counts := countsFixed[:len(writers)] - offsets := offsetsFixed[:len(writers)+1] - if len(writers) > SpillNumBuckets { - counts = make([]int32, len(writers)) - offsets = make([]int32, len(writers)+1) - } - if err := classifyRows(hashValues, len(writers), 0, rowIDs, counts, offsets); err != nil { - return err - } - for i := range bucketRowIds { - bucketRowIds[i] = rowIDs[offsets[i]:offsets[i+1]] + rows := bat.RowCount() + if !keycodec.ValidVectors(bat.Vecs, rows) || + !keycodec.ValidVectors(keyVecs, rows) { + return process.ErrHashBuildBudgetInvalid } - - // Only iterate non-empty buckets. - for bucketId, sels := range bucketRowIds { - if len(sels) == 0 { - continue - } - if writers[bucketId].Name == "" { - continue // disabled bucket — discard rows - } - buf := buffers[bucketId] - if buf == nil { - buf = batch.NewOffHeapWithSize(len(bat.Vecs)) - for j, vec := range bat.Vecs { - buf.Vecs[j] = vector.NewOffHeapVecWithType(*vec.GetType()) - buf.Vecs[j].PreExtend(8192, proc.Mp()) - } - buffers[bucketId] = buf - } - for j, vec := range bat.Vecs { - if err := buf.Vecs[j].UnionInt32(vec, sels, proc.Mp()); err != nil { - return err - } + var selected *batch.Batch + defer func() { + if selected != nil { + selected.Clean(proc.Mp()) + selected = nil } - buf.SetRowCount(buf.RowCount() + len(sels)) - if buf.RowCount() >= 8192 { - if err := FlushBucketBatch(proc, buf, &writers[bucketId], bucketBuf, analyzer); err != nil { - return err - } - buf.CleanOnlyData() + if retErr != nil && !hashbuild.IsRetryableMemoryCapacity(retErr) { + e.discardScatterBuffers() } - } - - return nil -} - -// scatterBatch scatters bat using the engine's reusable hash/row-id buffers. -func (e *SpillEngine) scatterBatch( - proc *process.Process, - bat *batch.Batch, - keyVecs []*vector.Vector, - writers []BucketWriter, - buffers []*batch.Batch, - partitionLevel uint64, - sourceAlreadyCharged bool, - analyzer process.Analyzer, -) error { - return e.scatterBatchBounded(proc, bat, keyVecs, writers, partitionLevel, sourceAlreadyCharged, analyzer) -} - -func scatterTransientBudgetBytes(bat *batch.Batch, sourceAlreadyCharged bool) (uint64, error) { - if bat == nil || bat.RowCount() < 0 { - return 0, process.ErrHashBuildBudgetInvalid - } - allocated := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > allocated { - allocated = size - } - return scatterTransientBudgetFor(allocated, uint64(len(bat.Vecs)), sourceAlreadyCharged) -} - -func scatterTransientBudgetFor(allocated, columns uint64, sourceAlreadyCharged bool) (uint64, error) { - oneMaterializedBatch, ok := batchPayloadWithAllocationSlack(allocated, columns) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } - // The selected batch and serialized payload are distinct live objects, so - // each receives one source-sized estimate plus bounded allocator/framing - // slack. The row/hash arrays are accounted separately from their capacities. - need, ok := addUint64(oneMaterializedBatch, oneMaterializedBatch) - if !ok { - return 0, process.ErrHashBuildBudgetInvalid - } + }() if !sourceAlreadyCharged { - if need, ok = addUint64(need, allocated); !ok { - return 0, process.ErrHashBuildBudgetInvalid + // The child batch is borrowed and already physically live. Rejecting a + // new logical token cannot reclaim it, so observe it while exact-account + // admission governs every new scatter allocation. + externalBytes := bat.Allocated() + if size := bat.Size(); size > externalBytes { + externalBytes = size } - } - return need, nil -} - -// reserveRebuildScatterScratch protects the one-batch repartition workspace -// before the rebuild retains another decoded batch. The reservation is only an -// accounting lease: no scatter buffers are allocated until re-spill actually -// starts. Keeping this floor lets a copy admission fail early enough that the -// already-retained batches can still be repartitioned under the same hard cap. -func (e *SpillEngine) reserveRebuildScatterScratch( - builder *hashbuild.HashmapBuilder, - bat *batch.Batch, - analyzer process.Analyzer, -) error { - if e.cfg.Budget == nil { - return nil - } - if builder == nil || bat == nil || bat.RowCount() < 0 { - return process.ErrHashBuildBudgetInvalid + analyzer.GetOpStats().SetMaxExtraStat( + "JoinSpillBorrowedSourceBytes", + int64(externalBytes), + ) } - allocated := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > allocated { - allocated = size + if e.allocationMP != nil && e.allocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid } - rows := bat.RowCount() - columns := len(bat.Vecs) - - // CopyIntoBatches may complete a partial physical tail with this record. - // Bound that resulting batch before the copy; reserving only for either - // input independently is not enough when two small records coalesce. - batches := builder.Batches.Buf - if bat.RowCount() != colexec.DefaultBatchSize && len(batches) > 0 { - tail := batches[len(batches)-1] - if tail == nil { - return process.ErrHashBuildBudgetInvalid - } - if tail.RowCount() != colexec.DefaultBatchSize { - merged, ok := predictMergedRetainedBytes(tail, bat) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if merged > allocated { - allocated = merged - } - if tail.RowCount() > math.MaxInt-bat.RowCount() { - return process.ErrHashBuildBudgetInvalid - } - rows += tail.RowCount() - if len(tail.Vecs) > columns { - columns = len(tail.Vecs) - } - } + e.allocationMP = proc.Mp() + var err error + e.scatterHashValues, err = growSpillSlice( + e.scatterHashValues, + rows, + proc.Mp(), + e.allocation, + SpillAllocationSiteHashValues, + ) + if err != nil { + return err } - - retained, ok := e.scatterRetainedBytes() - if !ok { - return process.ErrHashBuildBudgetInvalid + hashValues := e.scatterHashValues[:rows] + // Re-spill must consume fresh hash bits. Merely changing the initial seed + // leaves the low five bits correlated with the parent partition and can put + // every parent row into one child, making repartition unable to progress. + // Level zero uses bits 0..4, level one bits 5..9, and so on. + ComputeXXHash(keyVecs, hashValues, 0) + if err := checkSpillCanceled(proc); err != nil { + return err } - growth, ok := e.scatterCapacityGrowthBytes(rows, len(e.cfg.BuildKeyExprs)) - if !ok { + shift := partitionLevel * 5 + if shift >= 64 { return process.ErrHashBuildBudgetInvalid } - transient, err := scatterTransientBudgetFor(allocated, uint64(columns), true) + e.scatterBucketRowIds, err = growSpillSlice( + e.scatterBucketRowIds, + rows, + proc.Mp(), + e.allocation, + SpillAllocationSiteRowIDs, + ) if err != nil { return err } - need, ok := addUint64(retained, growth) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if need, ok = addUint64(need, transient); !ok { - return process.ErrHashBuildBudgetInvalid + if cap(e.keyVecs) < len(keyVecs) { + e.keyVecs = make([]*vector.Vector, len(keyVecs)) } - - return e.reserveRebuildScratchFloor(need, analyzer) -} - -func (e *SpillEngine) reserveRebuildScratchFloor(need uint64, analyzer process.Analyzer) error { - if need == 0 || e.cfg.Budget == nil { - return nil + if err := classifyRows(hashValues, len(writers), shift, e.scatterBucketRowIds, e.scatterBucketCounts[:], e.scatterBucketOffsets[:]); err != nil { + return err } - var err error - if e.scatterScratchReservation == nil { - e.scatterScratchReservation, err = e.cfg.Budget.Reserve(need) - if err != nil { - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchReserveRejects", 1) - } + for bucketID := range writers { + if err := checkSpillCanceled(proc); err != nil { return err } - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchReserveCount", 1) + start, end := e.scatterBucketOffsets[bucketID], e.scatterBucketOffsets[bucketID+1] + if start == end || writers[bucketID].Name == "" { + continue } - } else if current := e.scatterScratchReservation.Size(); need > current { - grow := need - current - if err = e.scatterScratchReservation.Grow(grow); err != nil { - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchGrowRejects", 1) + if selected == nil { + selected, err = newSpillBatch( + len(bat.Vecs), + e.allocation.selected, + ) + if err != nil { + return err + } + for j, vec := range bat.Vecs { + selected.Vecs[j], err = newSpillVector( + *vec.GetType(), + e.allocation.selected, + ) + if err != nil { + return err + } } - return err } - if analyzer != nil { - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchGrowCount", 1) - analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildScratchGrowBytes", spillStatInt64(grow)) + cursor := start + for cursor < end { + attemptEnd := end + reclaimedMinimum := false + for { + selected.CleanOnlyData() + sels := e.scatterBucketRowIds[cursor:attemptEnd] + var scatterErr error + for j, vec := range bat.Vecs { + if scatterErr = selected.Vecs[j].UnionInt32( + vec, + sels, + proc.Mp(), + ); scatterErr != nil { + break + } + } + if scatterErr == nil { + selected.SetRowCount(len(sels)) + scatterErr = e.appendScatterRecord( + proc, + selected, + &writers[bucketID], + bucketID, + analyzer, + ) + } + selected.CleanOnlyData() + if scatterErr == nil { + cursor = attemptEnd + break + } + if !hashbuild.IsRetryableMemoryCapacity(scatterErr) { + return scatterErr + } + if err := checkSpillCanceled(proc); err != nil { + return err + } + n := int(attemptEnd - cursor) + if n > 1 { + attemptEnd = cursor + int32((n+1)/2) + analyzer.GetOpStats().AddExtraStat( + "JoinSpillBatchReductions", + 1, + ) + continue + } + if !reclaimedMinimum { + before := e.allocation.account.Snapshot().Used + if err := e.reclaimOptionalScatterBuffers( + proc, + writers, + analyzer, + ); err != nil { + return err + } + reclaimedMinimum = true + after := e.allocation.account.Snapshot().Used + if after >= before { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-selected-or-codec", + e.allocation.account, + ) + } + analyzer.GetOpStats().AddExtraStat( + "JoinSpillOptionalReclaims", + 1, + ) + continue + } + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-selected-or-codec", + e.allocation.account, + ) + } } } - if need > e.scatterScratchFloor { - e.scatterScratchFloor = need - } - if analyzer != nil { - analyzer.GetOpStats().SetMaxExtraStat( - "JoinSpillRebuildScratchFloorBytes", - spillStatInt64(e.scatterScratchReservation.Size()), - ) - } return nil } -func spillStatInt64(v uint64) int64 { - if v > math.MaxInt64 { - return math.MaxInt64 - } - return int64(v) -} - -func (e *SpillEngine) scatterRetainedBytes() (uint64, bool) { - actual := uint64(0) - add := func(v uint64) bool { - var ok bool - actual, ok = addUint64(actual, v) - return ok - } - mul := func(v, n uint64) (uint64, bool) { - if n != 0 && v > math.MaxUint64/n { - return 0, false +func (e *SpillEngine) reclaimOptionalScatterBuffers( + proc *process.Process, + writers []BucketWriter, + analyzer process.Analyzer, +) error { + for bucket, buffer := range e.scatterAccountedWriteBuffers { + if buffer == nil { + continue } - return v * n, true - } - hashBytes, hashOK := mul(uint64(cap(e.scatterHashValues)), 8) - rowIDBytes, rowIDOK := mul(uint64(cap(e.scatterBucketRowIds)), 4) - keyBytes, keyOK := mul(uint64(cap(e.keyVecs)), 8) - countBytes, countOK := mul(uint64(len(e.scatterBucketCounts)), 4) - offsetBytes, offsetOK := mul(uint64(len(e.scatterBucketOffsets)), 4) - if !hashOK || !rowIDOK || !keyOK || !countOK || !offsetOK || - !add(hashBytes) || !add(rowIDBytes) || !add(keyBytes) || !add(countBytes) || !add(offsetBytes) || - !add(uint64(e.scatterWriteBuf.Cap())) { - return 0, false - } - for i := range e.scatterWriteBuffers { - if !add(uint64(e.scatterWriteBuffers[i].Cap())) { - return 0, false + if buffer.Len() > 0 { + if bucket >= len(writers) { + return process.ErrHashBuildBudgetInvalid + } + if err := e.flushPendingScatterBucket( + proc, + &writers[bucket], + bucket, + analyzer, + ); err != nil { + return err + } } + buffer.Free() + e.scatterAccountedWriteBuffers[bucket] = nil + } + e.scatterCoalesceDisabled = true + if e.scatterAccountedWriteBuf != nil { + e.scatterAccountedWriteBuf.Free() + e.scatterAccountedWriteBuf = nil } - return actual, true + return nil } -func (e *SpillEngine) scatterCapacityGrowthBytes(rows, keys int) (uint64, bool) { - if rows < 0 || keys < 0 { - return 0, false - } - var growth uint64 - addGrowth := func(required, current uint64) bool { - if required <= current { - return true - } - var ok bool - // make allocates the complete replacement before assignment drops the - // old slice. retained already includes current, so admit all of required. - growth, ok = addUint64(growth, required) - return ok - } - rowCount := uint64(rows) - keyCount := uint64(keys) - if rowCount > math.MaxUint64/8 || keyCount > math.MaxUint64/8 { - return 0, false - } - if !addGrowth(rowCount*8, uint64(cap(e.scatterHashValues))*8) || - !addGrowth(rowCount*4, uint64(cap(e.scatterBucketRowIds))*4) || - !addGrowth(keyCount*8, uint64(cap(e.keyVecs))*8) { - return 0, false +func (e *SpillEngine) releaseScatterComputeScratch() { + if e.allocationMP == nil { + return } - return growth, true + freeSpillSlice(e.scatterHashValues, e.allocationMP) + freeSpillSlice(e.scatterBucketRowIds, e.allocationMP) + e.scatterHashValues = nil + e.scatterBucketRowIds = nil } -// scatterBatchBounded writes one bucket at a time. The historical path kept -// SpillNumBuckets selected batches alive for the full input stream; that made -// a repartition pass itself exceed the hash-build budget. This implementation -// keeps one selected batch and one row-id slice, flushing it before advancing -// to the next bucket. -func (e *SpillEngine) scatterBatchBounded( +func (e *SpillEngine) scatterBatchWithPressure( proc *process.Process, bat *batch.Batch, keyVecs []*vector.Vector, @@ -1361,118 +886,276 @@ func (e *SpillEngine) scatterBatchBounded( partitionLevel uint64, sourceAlreadyCharged bool, analyzer process.Analyzer, -) (retErr error) { +) error { if bat == nil || bat.RowCount() == 0 { - return nil - } - if err := checkSpillCanceled(proc); err != nil { - return err - } - if len(writers) == 0 || len(writers) > SpillNumBuckets { - return process.ErrHashBuildBudgetInvalid + return e.scatterBatchBounded( + proc, + bat, + keyVecs, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) } rows := bat.RowCount() - var selected *batch.Batch - defer func() { - if selected != nil { - selected.Clean(proc.Mp()) - selected = nil - } - reconcileErr := e.reconcileScatterScratch() - if reconcileErr != nil && retErr == nil { - retErr = reconcileErr + chunk := rows + minimumRetried := false + guard := hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: rows, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + for start := 0; start < rows; { + end := rows + if chunk < rows-start { + end = start + chunk + } + current := bat + currentKeys := keyVecs + if start != 0 || end != rows { + var err error + current, err = bat.WindowWithAllocation( + start, end, proc.Mp(), e.allocation.selected, + ) + if err != nil { + return err + } + currentKeys = make([]*vector.Vector, len(keyVecs)) + for i, key := range keyVecs { + currentKeys[i], err = key.WindowWithAllocation( + start, end, proc.Mp(), e.allocation.selected, + ) + if err != nil { + for j := 0; j < i; j++ { + currentKeys[j].Free(proc.Mp()) + } + current.Clean(proc.Mp()) + return err + } + } } - if retErr != nil { - e.discardScatterBuffers() + err := e.scatterBatchBounded( + proc, + current, + currentKeys, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) + if current != bat { + for _, key := range currentKeys { + key.Free(proc.Mp()) + } + current.Clean(proc.Mp()) + } + if err == nil { + start = end + minimumRetried = false + nextUnits := chunk + if remaining := rows - start; remaining < nextUnits { + nextUnits = remaining + } + guard = hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: nextUnits, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + continue } - }() - if e.cfg.Budget != nil { - // Start with retained capacities already owned by this token, add only - // row/hash capacity growth, then add each per-batch transient once. - retained, ok := e.scatterRetainedBytes() - if !ok { - return process.ErrHashBuildBudgetInvalid - } - growth, ok := e.scatterCapacityGrowthBytes(rows, len(keyVecs)) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - transient, err := scatterTransientBudgetBytes(bat, sourceAlreadyCharged) - if err != nil { + if !hashbuild.IsRetryableMemoryCapacity(err) { return err } - need, ok := addUint64(retained, growth) - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if need, ok = addUint64(need, transient); !ok { - return process.ErrHashBuildBudgetInvalid - } - if e.scatterScratchReservation == nil { - e.scatterScratchReservation, err = e.cfg.Budget.Reserve(need) - } else if current := e.scatterScratchReservation.Size(); need > current { - // Grow the retained scratch token to the complete batch peak. Its - // current hash/row-id/coalesce capacities are components of need, - // not an additional allocation to charge a second time. - err = e.scatterScratchReservation.Grow(need - current) + if cancelErr := checkSpillCanceled(proc); cancelErr != nil { + return cancelErr + } + e.releaseScatterComputeScratch() + attempted := end - start + if attempted <= 1 { + if !minimumRetried { + if reclaimErr := e.reclaimOptionalScatterBuffers( + proc, + writers, + analyzer, + ); reclaimErr != nil { + return reclaimErr + } + next := hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: attempted, + OptionalDisabled: e.scatterCoalesceDisabled, + } + if guard.Advance(next) != nil { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-hash", + e.allocation.account, + ) + } + minimumRetried = true + analyzer.GetOpStats().AddExtraStat( + "JoinSpillMinimumRetries", + 1, + ) + continue + } + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-hash", + e.allocation.account, + ) } - if err != nil { + chunk = (attempted + 1) / 2 + if err := guard.Advance(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: chunk, + }); err != nil { return err } + analyzer.GetOpStats().AddExtraStat("JoinSpillInputReductions", 1) } - - if cap(e.scatterHashValues) < rows { - e.scatterHashValues = make([]uint64, rows) - } - hashValues := e.scatterHashValues[:rows] - // Re-spill must consume fresh hash bits. Merely changing the initial seed - // leaves the low five bits correlated with the parent partition and can put - // every parent row into one child, making repartition unable to progress. - // Level zero uses bits 0..4, level one bits 5..9, and so on. - ComputeXXHash(keyVecs, hashValues, 0) - if err := checkSpillCanceled(proc); err != nil { - return err - } - shift := partitionLevel * 5 - if shift >= 64 { - return process.ErrHashBuildBudgetInvalid - } - if cap(e.scatterBucketRowIds) < rows { - e.scatterBucketRowIds = make([]int32, rows) - } - if cap(e.keyVecs) < len(keyVecs) { - e.keyVecs = make([]*vector.Vector, len(keyVecs)) + return nil +} + +// scatterEvaluatedBatchWithPressure extends the same unpublished-input +// checkpoint across key evaluation and scatter. Exact expression executors may +// retain successfully admitted capacities after a later child/result growth +// fails; evaluating a smaller immutable window can then reuse those capacities +// without replaying any bucket record. scatterBatchWithPressure owns the +// transactional boundary after evaluation, so a capacity error returned here +// has not published the current window. +func (e *SpillEngine) scatterEvaluatedBatchWithPressure( + proc *process.Process, + bat *batch.Batch, + writers []BucketWriter, + partitionLevel uint64, + sourceAlreadyCharged bool, + analyzer process.Analyzer, + eval func(*batch.Batch) ([]*vector.Vector, error), +) error { + if bat == nil || bat.RowCount() == 0 { + return nil } - if err := classifyRows(hashValues, len(writers), shift, e.scatterBucketRowIds, e.scatterBucketCounts[:], e.scatterBucketOffsets[:]); err != nil { - return err + if eval == nil { + return process.ErrHashBuildBudgetInvalid } - for bucketID := range writers { - if err := checkSpillCanceled(proc); err != nil { - return err + + rows := bat.RowCount() + chunk := rows + minimumRetried := false + guard := hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: chunk, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + for start := 0; start < rows; { + end := rows + if chunk < rows-start { + end = start + chunk + } + current := bat + if start != 0 || end != rows { + var err error + current, err = bat.WindowWithAllocation( + start, end, proc.Mp(), e.allocation.selected, + ) + if err != nil { + return err + } } - start, end := e.scatterBucketOffsets[bucketID], e.scatterBucketOffsets[bucketID+1] - if start == end || writers[bucketID].Name == "" { - continue + + var err error + if err = checkSpillCanceled(proc); err == nil { + var keyVecs []*vector.Vector + keyVecs, err = eval(current) + if err == nil { + err = checkSpillCanceled(proc) + } + if err == nil { + err = e.scatterBatchWithPressure( + proc, + current, + keyVecs, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + ) + } } - sels := e.scatterBucketRowIds[start:end] - if selected == nil { - selected = batch.NewOffHeapWithSize(len(bat.Vecs)) - for j, vec := range bat.Vecs { - selected.Vecs[j] = vector.NewOffHeapVecWithType(*vec.GetType()) + if current != bat { + current.Clean(proc.Mp()) + } + if err == nil { + start = end + minimumRetried = false + nextUnits := chunk + if remaining := rows - start; remaining < nextUnits { + nextUnits = remaining } + guard = hashbuild.NewPressureRetryGuard(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: nextUnits, + OptionalDisabled: e.scatterCoalesceDisabled, + }, 64) + continue } - selected.CleanOnlyData() - for j, vec := range bat.Vecs { - if err := selected.Vecs[j].UnionInt32(vec, sels, proc.Mp()); err != nil { - selected.CleanOnlyData() + if !hashbuild.IsRetryableMemoryCapacity(err) { + return err + } + if cancelErr := checkSpillCanceled(proc); cancelErr != nil { + return cancelErr + } + + e.releaseScatterComputeScratch() + attempted := end - start + if attempted > 1 { + chunk = (attempted + 1) / 2 + if err := guard.Advance(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: chunk, + OptionalDisabled: e.scatterCoalesceDisabled, + }); err != nil { return err } + analyzer.GetOpStats().AddExtraStat( + "JoinSpillExpressionInputReductions", + 1, + ) + continue } - selected.SetRowCount(len(sels)) - if err := e.appendScatterRecord(proc, selected, &writers[bucketID], bucketID, analyzer); err != nil { - selected.CleanOnlyData() - return err + + if minimumRetried { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-expression", + e.allocation.account, + ) } + if reclaimErr := e.reclaimOptionalScatterBuffers( + proc, + writers, + analyzer, + ); reclaimErr != nil { + return reclaimErr + } + if err := guard.Advance(hashbuild.PressureProgress{ + Used: e.allocation.account.Snapshot().Used, + InputUnits: attempted, + OptionalDisabled: e.scatterCoalesceDisabled, + }); err != nil { + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + "scatter-expression", + e.allocation.account, + ) + } + minimumRetried = true + analyzer.GetOpStats().AddExtraStat( + "JoinSpillExpressionMinimumRetries", + 1, + ) } return nil } @@ -1481,59 +1164,114 @@ func (e *SpillEngine) appendScatterRecord(proc *process.Process, bat *batch.Batc if bucket < 0 || bucket >= SpillNumBuckets || writer == nil { return process.ErrHashBuildBudgetInvalid } - cnt := int64(bat.RowCount()) - if err := marshalSpillRecord(bat, &e.scatterWriteBuf); err != nil { + return e.appendAccountedScatterRecord( + proc, + bat, + writer, + bucket, + int64(bat.RowCount()), + analyzer, + ) +} + +func (e *SpillEngine) appendAccountedScatterRecord( + proc *process.Process, + bat *batch.Batch, + writer *BucketWriter, + bucket int, + rows int64, + analyzer process.Analyzer, +) error { + if e.allocationMP != nil && e.allocationMP != proc.Mp() { + return mpool.ErrAllocationAccountInvalid + } + e.allocationMP = proc.Mp() + if e.scatterAccountedWriteBuf == nil { + var err error + e.scatterAccountedWriteBuf, err = e.allocation.newBuffer( + proc.Mp(), + SpillAllocationSiteMarshalBuffer, + ) + if err != nil { + return err + } + } + if err := marshalSpillRecordTo( + bat, + e.scatterAccountedWriteBuf, + ); err != nil { return err } - payload := e.scatterWriteBuf.Bytes() - buf := &e.scatterWriteBuffers[bucket] - if buf.Len() > 0 && buf.Len()+len(payload) > spillWriteCoalesceSize { - if err := e.flushPendingScatterBucket(proc, writer, bucket, analyzer); err != nil { + payload := e.scatterAccountedWriteBuf.Bytes() + if e.scatterCoalesceDisabled { + return writeBucketPayload(proc, payload, rows, writer, analyzer) + } + buf := e.scatterAccountedWriteBuffers[bucket] + if buf != nil && buf.Len() > 0 && + buf.Len()+len(payload) > spillWriteCoalesceSize { + if err := e.flushPendingScatterBucket( + proc, + writer, + bucket, + analyzer, + ); err != nil { return err } } if len(payload) > spillWriteCoalesceSize { - return writeBucketPayload(proc, payload, cnt, writer, analyzer) + return writeBucketPayload(proc, payload, rows, writer, analyzer) } - if buf.Len() == 0 { - if !e.ensureScatterCoalesceCapacity(buf) { - return writeBucketPayload(proc, payload, cnt, writer, analyzer) + if buf == nil { + var err error + buf, err = e.allocation.newBuffer( + proc.Mp(), + SpillAllocationSiteCoalesceBuffer, + ) + if err != nil { + return err } - if buf.Cap() < spillWriteCoalesceSize { - *buf = *bytes.NewBuffer(make([]byte, 0, spillWriteCoalesceSize)) + e.scatterAccountedWriteBuffers[bucket] = buf + } + if buf.Len() == 0 && buf.Cap() < spillWriteCoalesceSize { + if err := buf.EnsureCapacity(spillWriteCoalesceSize); err != nil { + if errors.Is(err, mpool.ErrAllocationAccountCapacity) || + errors.Is(err, mpool.ErrAllocationMetadataSlots) { + return writeBucketPayload( + proc, + payload, + rows, + writer, + analyzer, + ) + } + return err } } - _, _ = buf.Write(payload) - e.scatterWriteRows[bucket] += cnt + if _, err := buf.Write(payload); err != nil { + return err + } + e.scatterWriteRows[bucket] += rows if buf.Len() >= spillWriteCoalesceSize { return e.flushPendingScatterBucket(proc, writer, bucket, analyzer) } return nil } -func (e *SpillEngine) ensureScatterCoalesceCapacity(buf *bytes.Buffer) bool { - if buf == nil || buf.Cap() >= spillWriteCoalesceSize { - return true - } - if e.cfg.Budget == nil || e.scatterScratchReservation == nil { - return e.cfg.Budget == nil - } - additional := uint64(spillWriteCoalesceSize - buf.Cap()) - if err := e.scatterScratchReservation.Grow(additional); err != nil { - return false - } - return true -} - func (e *SpillEngine) flushPendingScatterBucket(proc *process.Process, writer *BucketWriter, bucket int, analyzer process.Analyzer) error { if bucket < 0 || bucket >= SpillNumBuckets || writer == nil { return process.ErrHashBuildBudgetInvalid } - buf := &e.scatterWriteBuffers[bucket] - if buf.Len() == 0 { + buf := e.scatterAccountedWriteBuffers[bucket] + if buf == nil || buf.Len() == 0 { return nil } - err := writeBucketPayload(proc, buf.Bytes(), e.scatterWriteRows[bucket], writer, analyzer) + err := writeBucketPayload( + proc, + buf.Bytes(), + e.scatterWriteRows[bucket], + writer, + analyzer, + ) buf.Reset() e.scatterWriteRows[bucket] = 0 return err @@ -1545,7 +1283,11 @@ func (e *SpillEngine) flushPendingScatterBucket(proc *process.Process, writer *B func (e *SpillEngine) flushScatterBuffers(proc *process.Process, writers []BucketWriter, analyzer process.Analyzer) error { var firstErr error for bucket := 0; bucket < SpillNumBuckets; bucket++ { - if e.scatterWriteBuffers[bucket].Len() == 0 { + pending := 0 + if e.scatterAccountedWriteBuffers[bucket] != nil { + pending = e.scatterAccountedWriteBuffers[bucket].Len() + } + if pending == 0 { continue } var writer *BucketWriter @@ -1560,8 +1302,10 @@ func (e *SpillEngine) flushScatterBuffers(proc *process.Process, writers []Bucke } func (e *SpillEngine) discardScatterBuffers() { - for bucket := range e.scatterWriteBuffers { - e.scatterWriteBuffers[bucket].Reset() + for bucket := range e.scatterAccountedWriteBuffers { + if e.scatterAccountedWriteBuffers[bucket] != nil { + e.scatterAccountedWriteBuffers[bucket].Reset() + } e.scatterWriteRows[bucket] = 0 } } @@ -1571,51 +1315,36 @@ func (e *SpillEngine) discardScatterBuffers() { // charged while the next child hashmap is rebuilt. Cleanup calls this method // as an idempotent fallback for cancellation paths. func (e *SpillEngine) releaseScatterScratch() { + freeSpillSlice( + e.scatterHashValues, + e.allocationMP, + ) + freeSpillSlice( + e.scatterBucketRowIds, + e.allocationMP, + ) e.scatterHashValues = nil e.scatterBucketRowIds = nil e.keyVecs = nil - e.scatterWriteBuf = bytes.Buffer{} + if e.scatterAccountedWriteBuf != nil { + e.scatterAccountedWriteBuf.Free() + e.scatterAccountedWriteBuf = nil + } for i := range e.scatterBucketCounts { e.scatterBucketCounts[i] = 0 } for i := range e.scatterBucketOffsets { e.scatterBucketOffsets[i] = 0 } - for i := range e.scatterWriteBuffers { - e.scatterWriteBuffers[i] = bytes.Buffer{} - e.scatterWriteRows[i] = 0 - } - if e.scatterScratchReservation != nil { - e.scatterScratchReservation.Release() - e.scatterScratchReservation = nil - } - e.scatterScratchFloor = 0 -} - -// reconcileScatterScratch leaves only the capacities retained by the engine -// charged after a batch completes. The source batch, selected vectors, and -// marshal buffer are transient and must not pin budget across the queue. -func (e *SpillEngine) reconcileScatterScratch() error { - if e.scatterScratchReservation == nil { - return nil - } - actual, ok := e.scatterRetainedBytes() - if !ok { - return process.ErrHashBuildBudgetInvalid - } - if actual < e.scatterScratchFloor { - actual = e.scatterScratchFloor - } - reserved := e.scatterScratchReservation.Size() - if actual > reserved { - return process.ErrHashBuildBudgetInvalid - } - if actual < reserved { - if _, err := e.scatterScratchReservation.ReconcileDown(actual); err != nil { - return err + for i := range e.scatterAccountedWriteBuffers { + if e.scatterAccountedWriteBuffers[i] != nil { + e.scatterAccountedWriteBuffers[i].Free() + e.scatterAccountedWriteBuffers[i] = nil } + e.scatterWriteRows[i] = 0 } - return nil + e.allocationMP = nil + e.scatterCoalesceDisabled = false } // ReusableBufferPool maintains a persistent pool of spill buffers, preserving @@ -1676,12 +1405,8 @@ type SpillEngineConfig struct { // Budget is the statement generation shared with HashBuild. Rebuild and // re-spill must charge this exact generation; creating a fresh generation // would bypass aggregate admission and make ownership impossible to audit. - Budget *process.HashBuildBudgetGeneration - // ProbeExpressionLease is owned by the consuming join operator and borrowed - // by SpillEngine while it scatters or re-scatters probe batches. The join - // must free its probe executors before releasing this lease. - ProbeExpressionLease *hashbuild.ExpressionMemoryLease - MaxQueue int + Budget *process.HashBuildBudgetGeneration + MaxQueue int } // BucketResult encodes the outcome of a RebuildHashmap call. @@ -1697,60 +1422,90 @@ const ( // SpillEngine owns the spill bucket queue and drives the probe-batch loop. type SpillEngine struct { - cfg SpillEngineConfig - buckets []SpillBucket - spillFS spillFileServiceCache + cfg SpillEngineConfig + buckets []SpillBucket + spillFS spillFileServiceCache + allocation *SpillAllocationAccount + allocationMP *mpool.MPool // Current bucket state - buildReader BucketReader - probeReader BucketReader - buildReadBatch *batch.Batch - probeReadBatch *batch.Batch + buildReader BucketReader + probeReader BucketReader + buildReadBatch *batch.Batch + probeReadBatch *batch.Batch + probeExpected int64 + probeDecoded int64 + probeExpectedSet bool // Reusable scatter state buildPool ReusableBufferPool probePool ReusableBufferPool // Cached key executors for re-spill - keyExecs []colexec.ExpressionExecutor - keyVecs []*vector.Vector - buildExprLease *hashbuild.ExpressionMemoryLease + keyExecs []colexec.ExpressionExecutor + keyVecs []*vector.Vector // Reusable scatter buffers to avoid per-batch allocations. - scatterHashValues []uint64 - scatterBucketRowIds []int32 - scatterBucketCounts [SpillNumBuckets]int32 - scatterBucketOffsets [SpillNumBuckets + 1]int32 - scatterWriteBuf bytes.Buffer - scatterWriteBuffers [SpillNumBuckets]bytes.Buffer - scatterWriteRows [SpillNumBuckets]int64 - // The lease follows reusable scratch capacities within one rebuild/scatter - // phase. releaseScatterScratch drops both the backing arrays and this token; - // Cleanup is the idempotent terminal fallback. - scatterScratchReservation *process.HashBuildReservation - // scatterScratchFloor is pre-admitted only while rebuilding an already - // spilled bucket. It keeps one bounded repartition workspace available if the - // next retained-copy admission or threshold decision requires re-spill. It - // is a conservative bound, not a measurement of later physical allocations. - scatterScratchFloor uint64 - - // probeKeyEval is the unbudgeted fallback for probe re-scatter. Production - // spilled joins evaluate the probe executors owned by ProbeExpressionLease. + scatterHashValues []uint64 + scatterBucketRowIds []int32 + scatterBucketCounts [SpillNumBuckets]int32 + scatterBucketOffsets [SpillNumBuckets + 1]int32 + scatterAccountedWriteBuf *mpool.AccountedBuffer + scatterAccountedWriteBuffers [SpillNumBuckets]*mpool.AccountedBuffer + scatterCoalesceDisabled bool + scatterWriteRows [SpillNumBuckets]int64 + // probeKeyEval evaluates the consuming join's allocation-accounted probe + // executors during re-scatter. probeKeyEval func(*batch.Batch) ([]*vector.Vector, error) } -// NewSpillEngine creates an engine from configuration. Call InitFromSpilledMap next. -func NewSpillEngine(cfg SpillEngineConfig) *SpillEngine { +// NewSpillEngine binds every spill allocation to one execution generation. +// A spill engine cannot exist outside that account. +func NewSpillEngine( + cfg SpillEngineConfig, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, +) (*SpillEngine, error) { + if account == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + allocation, err := NewSpillAllocationAccount(account, owner) + if err != nil { + return nil, err + } + return newSpillEngine(cfg, allocation) +} + +func newSpillEngine( + cfg SpillEngineConfig, + allocation *SpillAllocationAccount, +) (*SpillEngine, error) { + if err := allocation.validate(); err != nil { + return nil, err + } + if cfg.Budget == nil { + return nil, process.ErrHashBuildBudgetInvalid + } + if cfg.Budget.Closed() { + return nil, process.ErrHashBuildBudgetClosed + } if cfg.MaxQueue <= 0 { cfg.MaxQueue = SpillNumBuckets * SpillNumBuckets } - return &SpillEngine{cfg: cfg} + engine := &SpillEngine{ + cfg: cfg, + allocation: allocation, + } + engine.buildReader.allocation = allocation + engine.probeReader.allocation = allocation + return engine, nil } func (e *SpillEngine) makeBucketWriters(prefix string) []BucketWriter { writers := MakeBucketWriters(prefix) for i := range writers { writers[i].spillFS = &e.spillFS + writers[i].Budget = e.cfg.Budget } return writers } @@ -1762,51 +1517,62 @@ func TakeSpillBuildPayload( proc *process.Process, jm *message.JoinMap, ) (message.SpillBuildPayload, *process.HashBuildBudgetGeneration, error) { + if jm == nil { + return message.SpillBuildPayload{}, nil, moerr.NewInternalError( + proc.Ctx, + message.ErrSpillBuildPayloadEmpty.Error(), + ) + } + expectedRows := jm.GetRowCount() payload, err := jm.TakeSpillBuildPayload() if err != nil { return message.SpillBuildPayload{}, nil, moerr.NewInternalError(proc.Ctx, err.Error()) } - - var budget *process.HashBuildBudgetGeneration - if len(payload.Files) > 0 { - var ok bool - budget, ok = payload.BudgetRef.(*process.HashBuildBudgetGeneration) - if !ok || budget == nil { - _ = payload.Close() - return message.SpillBuildPayload{}, nil, moerr.NewInternalError( - proc.Ctx, - "spilled join map is missing its producer budget generation", - ) + actualRows := int64(0) + validRows := expectedRows >= 0 + for _, file := range payload.Files { + if file == nil { + continue } - } else { - budget, err = proc.GetHashBuildBudget() - if err != nil { + rows := file.Rows() + if rows <= 0 || actualRows > math.MaxInt64-rows { + validRows = false + break + } + actualRows += rows + } + if !validRows || actualRows != expectedRows { + _ = payload.Close() + return message.SpillBuildPayload{}, nil, spillRowCountMismatch( + proc, + "build payload", + expectedRows, + actualRows, + ) + } + for _, file := range payload.Files { + if file == nil { + continue + } + if err := file.Validate(); err != nil { _ = payload.Close() return message.SpillBuildPayload{}, nil, err } } - return payload, budget, nil -} -// InitFromSpilledMap creates SpillBucket entries from build FDs. -// Empty (nil) FDs become placeholder buckets for outer-join semantics. -func (e *SpillEngine) InitFromSpilledMap(buildFds []*os.File) { - e.buckets = make([]SpillBucket, 0, len(buildFds)) - for _, fd := range buildFds { - var file *message.SpillFile - if fd != nil { - file = message.NewSpillFile(fd, 0, 0, nil) - } - e.buckets = append(e.buckets, SpillBucket{ - BuildFd: file, - Depth: 1, - }) + budget, ok := payload.BudgetRef.(*process.HashBuildBudgetGeneration) + if !ok || budget == nil { + _ = payload.Close() + return message.SpillBuildPayload{}, nil, moerr.NewInternalError( + proc.Ctx, + "spilled join map is missing its producer budget generation", + ) } + return payload, budget, nil } -// InitFromSpilledFiles is the ownership-preserving counterpart of the legacy -// descriptor initializer. Each SpillFile remains the sole owner of its fd and -// reservations while it moves through the bucket queue. +// InitFromSpilledFiles transfers the sole ownership of each build spill file +// and its resource reservations into the bucket queue. func (e *SpillEngine) InitFromSpilledFiles(files []*message.SpillFile) { e.buckets = make([]SpillBucket, 0, len(files)) for _, file := range files { @@ -1837,11 +1603,8 @@ func (e *SpillEngine) ScatterProbeTable( // The build payload defines the partition fanout. Using the production // maximum unconditionally would hash probe rows into writers that have no // corresponding build bucket; those files are never handed off and their - // rows would be silently discarded for legacy or reduced-fanout payloads. + // rows would be silently discarded for reduced-fanout payloads. writers := e.makeBucketWriters("probe")[:bucketCount] - for i := range writers { - writers[i].Budget = e.cfg.Budget - } // Disable writers for empty-build buckets unless outer join requires probe output. if !e.cfg.NeedsProbeForEmptyBuild { @@ -1883,14 +1646,17 @@ func (e *SpillEngine) ScatterProbeTable( if bat.IsEmpty() { continue } - keyVecs, err := e.evalProbeKeys(proc, bat, evalKeysFn) - if err != nil { - return err - } - if err := checkSpillCanceled(proc); err != nil { - return err - } - if err := e.scatterBatch(proc, bat, keyVecs, writers, nil, 0, false, analyzer); err != nil { + if err := e.scatterEvaluatedBatchWithPressure( + proc, + bat, + writers, + 0, + false, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + return e.evalProbeKeys(proc, current, evalKeysFn) + }, + ); err != nil { return err } } @@ -1935,11 +1701,23 @@ func (e *SpillEngine) NextProbeBatch(proc *process.Process) (*batch.Batch, error return nil, nil } if e.probeReadBatch == nil { - e.probeReadBatch = batch.NewOffHeapWithSize(0) + var err error + e.probeReadBatch, err = newSpillBatch(0, e.allocation.decoded) + if err != nil { + return nil, err + } } e.probeReader.mergeRecords = e.cfg.MergeProbeBatches || e.cfg.IsDedup bat, err := e.probeReader.ReadBatch(proc, e.probeReadBatch) if err == io.EOF { + if e.probeExpectedSet && e.probeDecoded != e.probeExpected { + return nil, spillRowCountMismatch( + proc, + "probe", + e.probeExpected, + e.probeDecoded, + ) + } return nil, nil } if err != nil { @@ -1948,30 +1726,63 @@ func (e *SpillEngine) NextProbeBatch(proc *process.Process) (*batch.Batch, error // Cancellation can race the reader's final record-boundary check. Do not // hand a freshly decoded batch to the join probe loop after that point. if err := checkSpillCanceled(proc); err != nil { - e.probeReader.releaseReadBatch(proc, e.probeReadBatch, nil) + e.probeReadBatch.Clean(proc.Mp()) return nil, err } + rows := int64(bat.RowCount()) + if rows < 0 || e.probeDecoded > math.MaxInt64-rows { + e.probeReadBatch.Clean(proc.Mp()) + return nil, spillRowCountMismatch( + proc, + "probe", + e.probeExpected, + math.MaxInt64, + ) + } + decoded := e.probeDecoded + rows + if e.probeExpectedSet && decoded > e.probeExpected { + e.probeReadBatch.Clean(proc.Mp()) + return nil, spillRowCountMismatch( + proc, + "probe", + e.probeExpected, + decoded, + ) + } + e.probeDecoded = decoded return bat, nil } -// builderMemSize computes total memory used by a HashmapBuilder during the rebuild -// loop. MemSize covers completed fixed-size batches; include the one permitted -// partial tail as well. The full scan is only a fallback for directly assembled -// state where MemSize has not been maintained. -func builderMemSize(builder *hashbuild.HashmapBuilder) int64 { - sz := builder.GetSize() + builder.Batches.MemSize - batches := builder.Batches.Buf - if builder.Batches.MemSize == 0 { - for _, b := range builder.Batches.Buf { - sz += int64(b.Size()) - } - } else if len(batches) > 0 { - tail := batches[len(batches)-1] - if tail != nil && tail.RowCount() != colexec.DefaultBatchSize { - sz += int64(tail.Size()) - } +func (e *SpillEngine) startProbe(file *message.SpillFile, expected int64) error { + if err := e.probeReader.ResetForSpillFile(file); err != nil { + return err } - return sz + e.probeExpected = expected + e.probeDecoded = 0 + e.probeExpectedSet = true + return nil +} + +func spillRowCountMismatch( + proc *process.Process, + side string, + expected int64, + actual int64, +) error { + return moerr.NewInternalErrorf( + proc.Ctx, + "corrupted spill %s row count: expected=%d actual=%d", + side, + expected, + actual, + ) +} + +// builderMemSize computes total memory used by a HashmapBuilder during the +// rebuild loop. GetSize covers hashmap structures and Batches.MemSize covers +// the raw accumulated batches maintained by the builder's copy API. +func builderMemSize(builder *hashbuild.HashmapBuilder) int64 { + return builder.GetSize() + builder.Batches.MemSize } func shouldReSpillBeforeRetain( @@ -2008,6 +1819,14 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketQueueEmpty, nil } bucket := e.buckets[0] + if (bucket.BuildFd == nil && bucket.BuildRows != 0) || + (bucket.ProbeFd == nil && bucket.ProbeRows != 0) || + bucket.BuildRows < 0 || bucket.ProbeRows < 0 { + return nil, BucketSkip, moerr.NewInternalError( + proc.Ctx, + "corrupted spill bucket file/row metadata", + ) + } // A build-only bucket cannot contribute to joins that never emit unmatched // build rows. Close and pop it before allocating a reader, copying batches, @@ -2026,11 +1845,10 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana e.buckets[0].ProbeFd = nil // transferred to reader below; prevent Cleanup double-close e.buckets = e.buckets[1:] if e.cfg.NeedsProbeForEmptyBuild && bucket.ProbeFd != nil { - if err := e.probeReader.EnsureBuffer(e.cfg.Budget); err != nil { - bucket.ProbeFd.Close() + if err := e.startProbe(bucket.ProbeFd, bucket.ProbeRows); err != nil { + bucket.ProbeFd = nil return nil, BucketSkip, err } - e.probeReader.ResetForSpillFile(bucket.ProbeFd) bucket.ProbeFd = nil return nil, BucketEmptyBuild, nil } @@ -2041,7 +1859,15 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana } builder := &hashbuild.HashmapBuilder{} + // The rebuild builder is stack-owned until GetJoinMap detaches its durable + // state. Keep a panic-safe terminal guard: the outer pipeline recover cannot + // otherwise reach a local builder abandoned by expression, hash, or spill + // hooks. + defer builder.Free(proc) builder.SetBudget(e.cfg.Budget) + if err := builder.SetAllocationAccount(e.allocation.account); err != nil { + return nil, BucketSkip, err + } builder.IsDedup = e.cfg.IsDedup builder.OnDuplicateAction = e.cfg.OnDuplicateAction builder.DedupBuildKeepLast = e.cfg.DedupBuildKeepLast @@ -2052,17 +1878,18 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketSkip, err } - if err := e.buildReader.EnsureBuffer(e.cfg.Budget); err != nil { - builder.Free(proc) - bucket.BuildFd.Close() - bucket.BuildFd = nil + if err := e.buildReader.ResetForSpillFile(bucket.BuildFd); err != nil { return nil, BucketSkip, err } - e.buildReader.ResetForSpillFile(bucket.BuildFd) e.buckets[0].BuildFd = nil // prevent Cleanup double-close on error defer e.buildReader.closeCurrentFile() if e.buildReadBatch == nil { - e.buildReadBatch = batch.NewOffHeapWithSize(0) + readBatch, err := newSpillBatch(0, e.allocation.decoded) + if err != nil { + builder.Free(proc) + return nil, BucketSkip, err + } + e.buildReadBatch = readBatch } // A rebuild may pre-admit one scatter workspace so a retained-copy reject // can still repartition the batches already owned by the builder. Release it @@ -2102,17 +1929,6 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketSkip, err } if bucket.Depth < SpillMaxPass { - if err := e.reserveRebuildScatterScratch(builder, bat, analyzer); err != nil { - // Scratch is contingency headroom, not a prerequisite for a - // bucket that may still rebuild within the cap. Admission misses - // are observable but best-effort; lifecycle/accounting failures - // remain terminal and are returned unchanged. - if !isBudgetAdmission(err) { - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) - return nil, BucketSkip, err - } - } if shouldReSpillBeforeRetain(builder, bat, e.cfg.SpillThreshold) { if analyzer != nil { analyzer.GetOpStats().AddExtraStat("JoinSpillRebuildPreCopyReSpillAttempts", 1) @@ -2132,7 +1948,7 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana builder.FreeHashMapAndBatches(proc) builder.Free(proc) if isBudgetAdmission(err) { - return nil, BucketSkip, noProgressError(bucket.Depth, err) + return nil, BucketSkip, noProgressError(proc, bucket.Depth) } return nil, BucketSkip, err } @@ -2154,6 +1970,16 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana return nil, BucketReSpilled, nil } } + if int64(builder.InputBatchRowCount) != bucket.BuildRows { + builder.FreeHashMapAndBatches(proc) + builder.Free(proc) + return nil, BucketSkip, spillRowCountMismatch( + proc, + "build", + bucket.BuildRows, + int64(builder.InputBatchRowCount), + ) + } if err := checkSpillCanceled(proc); err != nil { builder.FreeHashMapAndBatches(proc) @@ -2184,7 +2010,7 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana builder.FreeHashMapAndBatches(proc) builder.Free(proc) if isBudgetAdmission(err) { - return nil, BucketSkip, noProgressError(bucket.Depth, err) + return nil, BucketSkip, noProgressError(proc, bucket.Depth) } return nil, BucketSkip, err } @@ -2197,7 +2023,14 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana } } - jm := builder.GetJoinMap(proc.Mp()) + var jm *message.JoinMap + joinMapTransferred := false + defer func() { + if jm != nil && !joinMapTransferred { + jm.FreeMemory() + } + }() + jm = builder.GetJoinMap(proc.Mp()) if jm == nil { // GetJoinMap transfers nothing when the decoded build contains no // rows. Release executors and every residual builder allocation before @@ -2206,11 +2039,10 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana e.buckets[0].ProbeFd = nil // transferred to reader below; prevent Cleanup double-close e.buckets = e.buckets[1:] if e.cfg.NeedsProbeForEmptyBuild && bucket.ProbeFd != nil { - if err := e.probeReader.EnsureBuffer(e.cfg.Budget); err != nil { - bucket.ProbeFd.Close() + if err := e.startProbe(bucket.ProbeFd, bucket.ProbeRows); err != nil { + bucket.ProbeFd = nil return nil, BucketSkip, err } - e.probeReader.ResetForSpillFile(bucket.ProbeFd) bucket.ProbeFd = nil return nil, BucketEmptyBuild, nil } @@ -2227,14 +2059,13 @@ func (e *SpillEngine) RebuildHashmap(proc *process.Process, analyzer process.Ana // Pop the head bucket and open probe reader. e.buckets = e.buckets[1:] if bucket.ProbeFd != nil { - if err := e.probeReader.EnsureBuffer(e.cfg.Budget); err != nil { - bucket.ProbeFd.Close() - jm.Free() + if err := e.startProbe(bucket.ProbeFd, bucket.ProbeRows); err != nil { + bucket.ProbeFd = nil return nil, BucketSkip, err } - e.probeReader.ResetForSpillFile(bucket.ProbeFd) bucket.ProbeFd = nil } + joinMapTransferred = true return jm, BucketReady, nil } @@ -2247,13 +2078,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal // executors, so the two equivalent retained working sets never overlap. builder.FreeExecutors() buildWriters := e.makeBucketWriters("build_sub") - for i := range buildWriters { - buildWriters[i].Budget = e.cfg.Budget - } probeWriters := e.makeBucketWriters("probe_sub") - for i := range probeWriters { - probeWriters[i].Budget = e.cfg.Budget - } partitionLevel := uint64(bucket.Depth) probeFdConsumed := false @@ -2287,31 +2112,30 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal // Cache key executors. if len(e.keyExecs) != len(e.cfg.BuildKeyExprs) { - execs, lease, err := hashbuild.NewBudgetedExpressionExecutors( + execs, err := hashbuild.NewExpressionExecutors( proc, - e.cfg.Budget, e.cfg.BuildKeyExprs, - false, + e.allocation.account, ) + if err != nil { + for _, exec := range execs { + exec.Free() + } + } if err != nil { return nil, err } e.freeKeyExecs() e.keyExecs = execs - e.buildExprLease = lease } // evalAndScatter builds key vectors using the given executors and scatters. evalAndScatter := func( bat *batch.Batch, writers []BucketWriter, - buffers []*batch.Batch, execs []colexec.ExpressionExecutor, sourceAlreadyCharged bool, ) error { - if err := checkSpillCanceled(proc); err != nil { - return err - } if cap(e.keyVecs) < len(execs) { e.keyVecs = make([]*vector.Vector, len(execs)) } @@ -2321,24 +2145,34 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal keyVecs[i] = nil } }() - err := e.buildExprLease.Run(proc, bat.RowCount(), func(i int) error { - vec, evalErr := execs[i].Eval(proc, []*batch.Batch{bat}, nil) - if evalErr != nil { - return evalErr - } - keyVecs[i] = vec - return nil - }) - if err != nil { - // Eval may leave newly allocated child/result vectors cached. - // Destroy the owned executor tree before releasing its lease. - e.freeKeyExecs() - return err - } - if err := checkSpillCanceled(proc); err != nil { - return err - } - return e.scatterBatch(proc, bat, keyVecs, writers, nil, partitionLevel, sourceAlreadyCharged, analyzer) + return e.scatterEvaluatedBatchWithPressure( + proc, + bat, + writers, + partitionLevel, + sourceAlreadyCharged, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + for i := range keyVecs { + keyVecs[i] = nil + } + for i := range execs { + vec, err := execs[i].Eval(proc, []*batch.Batch{current}, nil) + if err != nil { + // Exact capacity pressure keeps the executor tree as the + // rollback checkpoint: admitted child/result capacities may + // make a smaller immutable window fit. Every other failure is + // terminal and can destroy the private tree immediately. + if !hashbuild.IsRetryableMemoryCapacity(err) { + e.freeKeyExecs() + } + return nil, err + } + keyVecs[i] = vec + } + return keyVecs, nil + }, + ) } var buildRows int64 @@ -2346,7 +2180,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if b != nil { buildRows += int64(b.RowCount()) } - if err := evalAndScatter(b, buildWriters, nil, e.keyExecs, true); err != nil { + if err := evalAndScatter(b, buildWriters, e.keyExecs, true); err != nil { return err } return nil @@ -2356,7 +2190,7 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if pending != nil && pending.RowCount() > 0 { // pending is the current BucketReader batch whose copy admission failed; // the reader keeps its batch token live until the next ReadBatch. - if err := evalAndScatter(pending, buildWriters, nil, e.keyExecs, true); err != nil { + if err := evalAndScatter(pending, buildWriters, e.keyExecs, true); err != nil { return nil, err } } @@ -2376,24 +2210,35 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal return nil, err } buildRows += int64(bat.RowCount()) - if err := evalAndScatter(bat, buildWriters, nil, e.keyExecs, true); err != nil { + if err := evalAndScatter(bat, buildWriters, e.keyExecs, true); err != nil { return nil, err } } + if buildRows != bucket.BuildRows { + return nil, spillRowCountMismatch( + proc, + "build", + bucket.BuildRows, + buildRows, + ) + } if err := e.flushScatterBuffers(proc, buildWriters, analyzer); err != nil { return nil, err } if e.probeReadBatch == nil { - e.probeReadBatch = batch.NewOffHeapWithSize(0) + readBatch, err := newSpillBatch(0, e.allocation.decoded) + if err != nil { + return nil, err + } + e.probeReadBatch = readBatch } - // Scatter probe file. Reuse reader's 4 MiB buffer from the build pass. + var probeRows int64 if bucket.ProbeFd != nil { - if err := reader.EnsureBuffer(e.cfg.Budget); err != nil { + if err := reader.ResetForSpillFile(bucket.ProbeFd); err != nil { return nil, err } - reader.ResetForSpillFile(bucket.ProbeFd) probeFdConsumed = true // Disable probe writers for empty sub-build buckets (unless outer join). if !e.cfg.NeedsProbeForEmptyBuild { @@ -2414,10 +2259,19 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if err != nil { return nil, err } - if err := scatterProbe(proc, e, bat, probeWriters, nil, partitionLevel, analyzer); err != nil { + probeRows += int64(bat.RowCount()) + if err := scatterProbe(proc, e, bat, probeWriters, partitionLevel, analyzer); err != nil { return nil, err } } + if probeRows != bucket.ProbeRows { + return nil, spillRowCountMismatch( + proc, + "probe", + bucket.ProbeRows, + probeRows, + ) + } if err := e.flushScatterBuffers(proc, probeWriters, analyzer); err != nil { return nil, err } @@ -2441,8 +2295,8 @@ func (e *SpillEngine) reSpillBucket(proc *process.Process, analyzer process.Anal if enqueue { if len(e.buckets)-1+len(subBuckets)+1 > e.cfg.MaxQueue { return nil, &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Message: fmt.Sprintf("join spill queue limit exceeded (limit=%d); reduce join-key skew or increase processLimitationSize", e.cfg.MaxQueue), + Kind: process.HashBuildBudgetErrorInvalid, + Message: fmt.Sprintf("spill queue limit exceeded: limit=%d", e.cfg.MaxQueue), } } buildFile, err := buildWriters[i].handOffSpillFile() @@ -2574,51 +2428,30 @@ func (e *SpillEngine) AdvanceToNextBucket( return true, nil } -// scatterProbe evaluates probe-side keys (EqConds[0]) for probe re-scatter. -// It uses the borrowed probe lease, not build-side keyExecs; probeKeyEval is -// retained only as the unbudgeted fallback. -func scatterProbe(proc *process.Process, e *SpillEngine, bat *batch.Batch, writers []BucketWriter, buffers []*batch.Batch, seed uint64, analyzer process.Analyzer) error { - keyVecs, err := e.evalProbeKeys(proc, bat, e.probeKeyEval) - if err != nil { - return err - } - return e.scatterBatch(proc, bat, keyVecs, writers, buffers, seed, true, analyzer) +// scatterProbe evaluates the consuming join's probe-side keys for re-scatter. +func scatterProbe(proc *process.Process, e *SpillEngine, bat *batch.Batch, writers []BucketWriter, seed uint64, analyzer process.Analyzer) error { + return e.scatterEvaluatedBatchWithPressure( + proc, + bat, + writers, + seed, + true, + analyzer, + func(current *batch.Batch) ([]*vector.Vector, error) { + return e.evalProbeKeys(proc, current, e.probeKeyEval) + }, + ) } func (e *SpillEngine) evalProbeKeys( proc *process.Process, bat *batch.Batch, - fallback func(*batch.Batch) ([]*vector.Vector, error), + eval func(*batch.Batch) ([]*vector.Vector, error), ) ([]*vector.Vector, error) { - if e.cfg.ProbeExpressionLease == nil { - if fallback == nil { - return nil, process.ErrHashBuildBudgetInvalid - } - return fallback(bat) - } - if e.cfg.ProbeExpressionLease.Len() != len(e.cfg.ProbeKeyExprs) { + if eval == nil { return nil, process.ErrHashBuildBudgetInvalid } - if cap(e.keyVecs) < len(e.cfg.ProbeKeyExprs) { - e.keyVecs = make([]*vector.Vector, len(e.cfg.ProbeKeyExprs)) - } - keyVecs := e.keyVecs[:len(e.cfg.ProbeKeyExprs)] - err := e.cfg.ProbeExpressionLease.Eval( - proc, - []*batch.Batch{bat}, - bat.RowCount(), - func(index int, vec *vector.Vector) error { - keyVecs[index] = vec - return nil - }, - ) - if err != nil { - for i := range keyVecs { - keyVecs[i] = nil - } - return nil, err - } - return keyVecs, nil + return eval(bat) } func (e *SpillEngine) freeKeyExecs() { @@ -2628,33 +2461,19 @@ func (e *SpillEngine) freeKeyExecs() { } } e.keyExecs = nil - if e.buildExprLease != nil { - e.buildExprLease.Release() - e.buildExprLease = nil - } } func isBudgetAdmission(err error) bool { - return err != nil && - errors.Is(err, process.ErrHashBuildBudgetAdmission) + return hashbuild.IsRetryableMemoryCapacity(err) } -func noProgressError(depth int, cause error) error { - budgetErr := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Message: fmt.Sprintf("join spill cannot make progress at depth %d; reduce join-key skew or increase processLimitationSize", depth), - } - if cause != nil { - var budgetCause *process.HashBuildBudgetError - if errors.As(cause, &budgetCause) && budgetCause.Kind == process.HashBuildBudgetErrorAdmission { - budgetErr.Resource = budgetCause.Resource - budgetErr.Requested = budgetCause.Requested - budgetErr.Used = budgetCause.Used - budgetErr.Cap = budgetCause.Cap - budgetErr.Message = fmt.Sprintf("join spill cannot make progress at depth %d", depth) - } - } - return budgetErr +func noProgressError(proc *process.Process, depth int) error { + _ = proc + return hashbuild.NewMinimumAllocationPressureError( + "join-spill", + fmt.Sprintf("partition-depth-%d", depth), + nil, + ) } // Cleanup releases all engine resources. diff --git a/pkg/sql/colexec/spillutil/join_spill_test.go b/pkg/sql/colexec/spillutil/join_spill_test.go index 3ca3f369d8d18..b3bd33ef23e90 100644 --- a/pkg/sql/colexec/spillutil/join_spill_test.go +++ b/pkg/sql/colexec/spillutil/join_spill_test.go @@ -15,4478 +15,577 @@ package spillutil import ( - "bufio" "bytes" "context" - "errors" "fmt" "io" - "math" "os" - "runtime" - "strings" "testing" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" - "github.com/matrixorigin/matrixone/pkg/defines" - "github.com/matrixorigin/matrixone/pkg/fileservice" "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/matrixorigin/matrixone/pkg/sql/colexec" "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" - plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" "github.com/matrixorigin/matrixone/pkg/testutil" "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "github.com/stretchr/testify/require" ) -type boundaryCancelReader struct { - reader *bytes.Reader - boundary int64 - read int64 - cancel func() - canceled bool +func makeTestKeyExpr() []*plan.Expr { + return []*plan.Expr{{ + Typ: plan.Type{Id: int32(types.T_int32), Width: 32}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, + }} +} + +func makeTestEvalKeysFn() func(*batch.Batch) ([]*vector.Vector, error) { + return func(bat *batch.Batch) ([]*vector.Vector, error) { + return bat.Vecs[:1], nil + } +} + +func makeInt32Batch(proc *process.Process, values []int32) *batch.Batch { + bat := batch.NewWithSize(1) + bat.Vecs[0] = testutil.MakeInt32Vector(values, nil, proc.Mp()) + bat.SetRowCount(len(values)) + return bat } -func (r *boundaryCancelReader) Read(p []byte) (int, error) { - if !r.canceled && r.read >= r.boundary { - r.canceled = true - r.cancel() +func writeBuildFile( + proc *process.Process, + name string, + bat *batch.Batch, +) *os.File { + return writeBuildRecords(proc, name, bat) +} + +func writeBuildRecords( + proc *process.Process, + name string, + batches ...*batch.Batch, +) *os.File { + spillfs, err := proc.GetSpillFileService() + if err != nil { + panic(err) } - if !r.canceled { - remaining := r.boundary - r.read - if int64(len(p)) > remaining { - p = p[:remaining] + file, err := spillfs.CreateAndRemoveFile(context.Background(), name) + if err != nil { + panic(err) + } + for _, bat := range batches { + if _, err := file.Write(marshalTestSpillRecord(bat)); err != nil { + panic(err) } } - n, err := r.reader.Read(p) - r.read += int64(n) - return n, err + if _, err := file.Seek(0, io.SeekStart); err != nil { + panic(err) + } + return file +} + +func makeCorruptBatchFile(t *testing.T) *os.File { + t.Helper() + file, err := os.CreateTemp(t.TempDir(), "corrupt-spill") + require.NoError(t, err) + rowCount, batchSize := int64(1), int64(1) + var encoded bytes.Buffer + encoded.Write(types.EncodeInt64(&rowCount)) + encoded.Write(types.EncodeInt64(&batchSize)) + encoded.WriteByte(0xff) + _, err = file.Write(encoded.Bytes()) + require.NoError(t, err) + _, err = file.Seek(0, io.SeekStart) + require.NoError(t, err) + return file } func TestTakeSpillBuildPayloadRejectsWrongBudgetRef(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - _, _, err := TakeSpillBuildPayload(proc, nil) require.ErrorContains(t, err, message.ErrSpillBuildPayloadEmpty.Error()) fd, err := os.CreateTemp(t.TempDir(), "wrong-budget-ref") require.NoError(t, err) - t.Cleanup(func() { _ = fd.Close() }) + _, err = fd.Write([]byte{1}) + require.NoError(t, err) + _, err = fd.Seek(0, io.SeekStart) + require.NoError(t, err) releases := 0 file := message.NewSpillFile(fd, 1, 1, func() { releases++ }) - jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) + jm := message.NewJoinMap( + message.GroupSels{}, nil, nil, nil, nil, proc.Mp(), + ) + jm.SetRowCount(1) jm.IncRef(1) - jmFreed := false - t.Cleanup(func() { - if !jmFreed { - jm.Free() - } - }) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ Files: []*message.SpillFile{file}, BudgetRef: struct{}{}, })) - _, _, err = TakeSpillBuildPayload(proc, jm) require.ErrorContains(t, err, "missing its producer budget generation") require.Equal(t, 1, releases) - _, err = fd.Stat() - require.Error(t, err) - jm.Free() - jmFreed = true - require.Equal(t, 1, releases) } -func TestTakeSpillBuildPayloadLegacyResolvesConsumerBudget(t *testing.T) { +func TestTakeSpillBuildPayloadRejectsGlobalRowMismatch(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - fd, err := os.CreateTemp(t.TempDir(), "legacy-build-payload") + budget := process.MustNewHashBuildBudget(1<<20, 1<<20) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + fd, err := os.CreateTemp(t.TempDir(), "payload-row-mismatch") require.NoError(t, err) - t.Cleanup(func() { _ = fd.Close() }) - jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, nil, proc.Mp()) + releases := 0 + jm := message.NewJoinMap( + message.GroupSels{}, nil, nil, nil, nil, proc.Mp(), + ) + jm.SetRowCount(2) jm.IncRef(1) - jmFreed := false - t.Cleanup(func() { - if !jmFreed { - jm.Free() - } - }) require.NoError(t, jm.SetSpillBuildPayload(message.SpillBuildPayload{ - LegacyFds: []*os.File{fd}, + Files: []*message.SpillFile{ + message.NewSpillFile(fd, 1, 0, func() { releases++ }), + }, + BudgetRef: generation, })) - - wantBudget, err := proc.GetHashBuildBudget() - require.NoError(t, err) - payload, budget, err := TakeSpillBuildPayload(proc, jm) - require.NoError(t, err) - require.Same(t, fd, payload.LegacyFds[0]) - require.Same(t, wantBudget, budget) - t.Cleanup(func() { _ = payload.Close() }) - require.NoError(t, payload.Close()) - _, err = fd.Stat() - require.Error(t, err) - + _, _, err = TakeSpillBuildPayload(proc, jm) + require.ErrorContains(t, err, "row count") + require.Equal(t, 1, releases) jm.Free() - jmFreed = true -} - -func TestComputeXXHash(t *testing.T) { - mp := mpool.MustNewZero() - ComputeXXHash(nil, nil, 0) - - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, mp) - defer vec.Free(mp) - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{vec}, hashValues, 0) - require.NotEqual(t, uint64(0), hashValues[0]) - - constVec, err := vector.NewConstFixed(types.T_int32.ToType(), int32(7), 3, mp) - require.NoError(t, err) - defer constVec.Free(mp) - ComputeXXHash([]*vector.Vector{constVec}, hashValues, 1) - require.Equal(t, hashValues[0], hashValues[2]) - - shortVec := testutil.MakeInt32Vector([]int32{9}, nil, mp) - defer shortVec.Free(mp) - ComputeXXHash([]*vector.Vector{shortVec}, hashValues, 2) - - constNull := vector.NewConstNull(types.T_int32.ToType(), 3, mp) - defer constNull.Free(mp) - nullHashes := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{constNull}, nullHashes, 7) - require.Equal(t, hashCombine(uint64(7), uint64(0)), nullHashes[0]) - require.Equal(t, nullHashes[0], nullHashes[2]) } func TestClassifyRowsConservesRows(t *testing.T) { hashes := make([]uint64, 257) for i := range hashes { - // Include skew, an empty bucket, and rows that differ only at the - // re-spill bit offset. hashes[i] = uint64(i%7) | (uint64(i&3) << 5) } rowIDs := make([]int32, len(hashes)) counts := make([]int32, SpillNumBuckets) offsets := make([]int32, SpillNumBuckets+1) - require.NoError(t, classifyRows(hashes, SpillNumBuckets, 0, rowIDs, counts, offsets)) - require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) - seen := make([]bool, len(hashes)) - for bucket := 0; bucket < SpillNumBuckets; bucket++ { - for _, rowID := range rowIDs[offsets[bucket]:offsets[bucket+1]] { + for _, shift := range []uint64{0, 5} { + require.NoError(t, classifyRows( + hashes, + SpillNumBuckets, + shift, + rowIDs, + counts, + offsets, + )) + require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) + seen := make([]bool, len(hashes)) + for _, rowID := range rowIDs { row := int(rowID) require.GreaterOrEqual(t, row, 0) require.Less(t, row, len(hashes)) require.False(t, seen[row]) seen[row] = true - require.Equal(t, bucket, int(hashes[row]&(SpillNumBuckets-1))) - } - } - for _, ok := range seen { - require.True(t, ok) - } - - // Re-spill consumes the next five hash bits without changing the row - // conservation invariant. - require.NoError(t, classifyRows(hashes, SpillNumBuckets, 5, rowIDs, counts, offsets)) - require.Equal(t, int32(len(hashes)), offsets[SpillNumBuckets]) -} - -func legacyClassifyRows(hashes []uint64, rowIDs []int32) { - pos := 0 - for bucket := uint64(0); bucket < SpillNumBuckets; bucket++ { - for row, hash := range hashes { - if hash&(SpillNumBuckets-1) == bucket { - rowIDs[pos] = int32(row) - pos++ - } } } + require.ErrorIs(t, classifyRows( + []uint64{0}, + SpillNumBuckets*2, + 0, + make([]int32, 1), + make([]int32, SpillNumBuckets*2), + make([]int32, SpillNumBuckets*2+1), + ), process.ErrHashBuildBudgetInvalid) } -func BenchmarkClassifyRows(b *testing.B) { - hashes := make([]uint64, 8192) - for i := range hashes { - hashes[i] = uint64(i*2654435761) ^ uint64(i>>3) - } - rowIDs := make([]int32, len(hashes)) - counts := make([]int32, SpillNumBuckets) - offsets := make([]int32, SpillNumBuckets+1) - b.Run("counts_prefix_rowids", func(b *testing.B) { - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - if err := classifyRows(hashes, SpillNumBuckets, 0, rowIDs, counts, offsets); err != nil { - b.Fatal(err) - } - } - }) - b.Run("legacy_bucket_scan", func(b *testing.B) { - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - legacyClassifyRows(hashes, rowIDs) - } - }) -} - -func TestBucketWriterAccountedHandOffSeekFailureRetainsOwnership(t *testing.T) { +func TestAccountedBucketReaderRoundTripAndCorruption(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) + budget := process.MustNewHashBuildBudget(8<<20, 8<<20) generation, err := budget.OpenGeneration(1) require.NoError(t, err) - fdToken, err := generation.ReserveSpillFD(1) - require.NoError(t, err) - fd, err := os.CreateTemp(t.TempDir(), "closed-spill") + source := makeInt32Batch(proc, []int32{1, 2, 3}) + defer source.Clean(proc.Mp()) + registry, err := mpool.NewAllocationAccountRegistry(1, 1<<20) require.NoError(t, err) - require.NoError(t, fd.Close()) - w := BucketWriter{Fd: fd, fdReservation: fdToken} - file, err := w.handOffSpillFile() - require.Error(t, err) - require.Nil(t, file) - require.Same(t, fd, w.Fd, "failed rewind must retain file ownership") - require.Equal(t, uint64(1), generation.SpillFDUsed()) - w.Close() - require.Zero(t, generation.SpillFDUsed()) -} - -func TestFlushBucketBatchAndReadRoundtrip(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() + account, err := registry.OpenWithController(8<<20, generation) require.NoError(t, err) - - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_rt") + allocation, err := NewSpillAllocationAccount( + account, + hashbuild.HashBuildAllocationOwner, + ) require.NoError(t, err) - defer f.Close() - - var buf bytes.Buffer - w := BucketWriter{Name: "test_rt", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{10, 20, 30}, nil, proc.Mp()) - bat.SetRowCount(3) - require.NoError(t, FlushBucketBatch(proc, nil, &w, &buf, nil)) - err = FlushBucketBatch(proc, bat, &w, &buf, process.NewAnalyzer(0, false, false, "test")) + reader := &BucketReader{ + fd: writeBuildFile(proc, t.Name(), source), + allocation: allocation, + } + decoded, err := newSpillBatch(0, reader.allocation.decoded) require.NoError(t, err) - - fd := w.HandOffFd() - reader := BucketReader{fd: fd} - reuseBat := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuseBat) + got, err := reader.ReadBatch(proc, decoded) require.NoError(t, err) require.Equal(t, 3, got.RowCount()) + require.Equal(t, []int32{1, 2, 3}, vector.MustFixedColNoTypeCheck[int32](got.Vecs[0])) + got.Clean(proc.Mp()) + _, err = reader.ReadBatch(proc, decoded) + require.ErrorIs(t, err, io.EOF) reader.Close() -} - -func TestBucketReaderAccountedLifecycle(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - bat := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - defer bat.Clean(proc.Mp()) - var buf bytes.Buffer - w := BucketWriter{Name: "test_accounted_reader", Budget: generation} - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - file, err := w.handOffSpillFile() - require.NoError(t, err) - require.NotNil(t, file) - require.Positive(t, generation.SpillDiskUsed()) - require.Equal(t, uint64(1), generation.SpillFDUsed()) - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForSpillFile(file) - reuseBat := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuseBat) + require.Zero(t, account.Snapshot().Used) + _, _, err = registry.CompleteTerminal(account) require.NoError(t, err) - require.Equal(t, 4, got.RowCount()) - require.Positive(t, generation.Used()) - - reuseBat.Clean(proc.Mp()) - reader.Close() require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} -func TestReconcileReadReservation(t *testing.T) { - budget := process.MustNewHashBuildBudget(1<<20, 1<<20) - generation, err := budget.OpenGeneration(1) + corruptState := newTestSpillAllocationAccount(t, 8<<20, 16) + corrupt := &BucketReader{ + fd: makeCorruptBatchFile(t), + allocation: corruptState.allocation, + } + bad, err := newSpillBatch(0, corruptState.allocation.decoded) require.NoError(t, err) - - t.Run("shrink", func(t *testing.T) { - token, err := generation.Reserve(1024) - require.NoError(t, err) - require.NoError(t, reconcileReadReservation(token, 256)) - require.Equal(t, uint64(256), generation.Used()) - require.True(t, token.Release()) - require.Zero(t, generation.Used()) - }) - - t.Run("underestimated-retained-bytes", func(t *testing.T) { - token, err := generation.Reserve(256) - require.NoError(t, err) - require.ErrorIs(t, reconcileReadReservation(token, 257), process.ErrHashBuildBudgetInvalid) - // Failed upward reconciliation keeps the original token live so both - // reader cleanup paths can release the complete reservation exactly once. - require.Equal(t, uint64(256), generation.Used()) - require.True(t, token.Release()) - require.Zero(t, generation.Used()) - }) + _, err = corrupt.ReadBatch(proc, bad) + require.Error(t, err) + corrupt.Close() + finalizeTestSpillAllocationAccount(t, corruptState) } -func TestPredictMergedRetainedBytesMatchesUnionBatch(t *testing.T) { +func TestBucketReaderRejectsSchemaChangeBeforeMerge(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - tests := []struct { - name string - dst func() (*batch.Batch, *batch.Batch) - }{ - { - name: "fixed-and-varlen-multi-column", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(2) - dst.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - dst.Vecs[1] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(2) - src.Vecs[0] = testutil.MakeInt32Vector([]int32{3, 4, 5}, nil, proc.Mp()) - src.Vecs[1] = testutil.MakeVarcharVector([]string{"right", "hand", "rows"}, nil, proc.Mp()) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-fixed", - dst: func() (*batch.Batch, *batch.Batch) { - dst := makeInt32Batch(proc, []int32{1, 2}) - src := batch.NewWithSize(1) - var err error - src.Vecs[0], err = vector.NewConstFixed(types.T_int32.ToType(), int32(9), 3, proc.Mp()) - require.NoError(t, err) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-inline-varlen", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(1) - dst.Vecs[0] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(1) - var err error - src.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("inline"), 3, proc.Mp()) - require.NoError(t, err) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-non-inline", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(1) - dst.Vecs[0] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(1) - var err error - src.Vecs[0], err = vector.NewConstBytes(types.T_varchar.ToType(), []byte("a sufficiently long constant value"), 3, proc.Mp()) - require.NoError(t, err) - src.SetRowCount(3) - return dst, src - }, - }, - { - name: "const-null", - dst: func() (*batch.Batch, *batch.Batch) { - dst := batch.NewWithSize(1) - dst.Vecs[0] = testutil.MakeVarcharVector([]string{"left", "side"}, nil, proc.Mp()) - dst.SetRowCount(2) - src := batch.NewWithSize(1) - src.Vecs[0] = vector.NewConstNull(types.T_varchar.ToType(), 3, proc.Mp()) - src.SetRowCount(3) - return dst, src - }, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - dst, src := tc.dst() - defer dst.Clean(proc.Mp()) - defer src.Clean(proc.Mp()) - predicted, ok := predictMergedRetainedBytes(dst, src) - require.True(t, ok) - require.NoError(t, dst.UnionWindow(src, 0, src.RowCount(), proc.Mp())) - actual, ok := batchRetainedBytes(dst) - require.True(t, ok) - require.LessOrEqual(t, actual, predicted) - require.Equal(t, src.RowCount()+2, dst.RowCount()) - }) + textBatch := batch.NewWithSize(1) + textBatch.Vecs[0] = vector.NewVec(types.T_varchar.ToType()) + require.NoError(t, vector.AppendBytes( + textBatch.Vecs[0], []byte("x"), false, proc.Mp(), + )) + textBatch.SetRowCount(1) + defer textBatch.Clean(proc.Mp()) + intBatch := makeInt32Batch(proc, []int32{1}) + defer intBatch.Clean(proc.Mp()) + + state := newTestSpillAllocationAccount(t, 8<<20, 16) + reader := &BucketReader{ + fd: writeBuildRecords( + proc, + t.Name(), + textBatch, + intBatch, + ), + mergeRecords: true, + allocation: state.allocation, } + reuse, err := newSpillBatch(0, state.allocation.decoded) + require.NoError(t, err) + _, err = reader.ReadBatch(proc, reuse) + require.ErrorContains(t, err, "spill batch schema changed") + reader.Close() + require.Zero(t, state.account.Snapshot().Used) + finalizeTestSpillAllocationAccount(t, state) } -func TestPredictMergedRetainedBytesAdmissionBudget(t *testing.T) { +func TestRebuildHashmapBasic(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - dstVals := make([]int32, 100) - srcVals := make([]int32, 100) - for i := range dstVals { - dstVals[i] = int32(i) - srcVals[i] = int32(i + len(dstVals)) - } - dst := makeInt32Batch(proc, dstVals) - src := makeInt32Batch(proc, srcVals) - defer dst.Clean(proc.Mp()) - defer src.Clean(proc.Mp()) - old, ok := batchRetainedBytes(dst) - require.True(t, ok) - next, ok := batchRetainedBytes(src) - require.True(t, ok) - predicted, ok := predictMergedRetainedBytes(dst, src) - require.True(t, ok) - require.Greater(t, predicted, old+next, "rounded destination growth must be admitted independently") - - reserveAll := func(t *testing.T, cap uint64) { - budget := process.MustNewHashBuildBudget(cap, cap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - o, err := generation.Reserve(old) - require.NoError(t, err) - n, err := generation.Reserve(next) - require.NoError(t, err) - d, err := generation.Reserve(predicted) - require.NoError(t, err) - o.Release() - n.Release() - d.Release() - require.Zero(t, generation.Used()) + values := make([]int32, 100) + for i := range values { + values[i] = int32(i) } - reserveAll(t, old+next+predicted) - - budget := process.MustNewHashBuildBudget(old+next+predicted-1, old+next+predicted-1) - generation, err := budget.OpenGeneration(2) - require.NoError(t, err) - o, err := generation.Reserve(old) + build := makeInt32Batch(proc, values) + defer build.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedsBuildForEmptyProbe: true, + }) + initTestSpillFiles(engine, []*os.File{ + writeBuildFile(proc, t.Name(), build), + }, int64(len(values))) + jm, result, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) require.NoError(t, err) - n, err := generation.Reserve(next) + require.Equal(t, BucketReady, result) + require.Equal(t, int64(100), jm.GetRowCount()) + jm.Free() + _, result, err = engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) require.NoError(t, err) - _, err = generation.Reserve(predicted) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - // Admission happens before UnionBatch, so the destination remains intact. - require.Equal(t, len(dstVals), dst.RowCount()) - require.Equal(t, dstVals[0], vector.GetFixedAtNoTypeCheck[int32](dst.Vecs[0], 0)) - o.Release() - n.Release() - require.Zero(t, generation.Used()) + require.Equal(t, BucketQueueEmpty, result) + engine.Cleanup(proc) } -func TestBucketWriterAggregatesDiskAccountingPerFile(t *testing.T) { +func TestReSpillConservesBuildAndProbeRows(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) + values := make([]int32, 5_000) + for i := range values { + values[i] = int32(i) + } + build := makeInt32Batch(proc, values) + probe := makeInt32Batch(proc, values) + defer build.Clean(proc.Mp()) + defer probe.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + SpillThreshold: 500, + NeedsBuildForEmptyProbe: true, + NeedsProbeForEmptyBuild: true, + }) + initTestSpillFiles(engine, []*os.File{ + writeBuildFile(proc, t.Name()+"-build", build), + }, int64(len(values))) + engine.buckets[0].ProbeFd = newTestSpillFile( + writeBuildFile(proc, t.Name()+"-probe", probe), + int64(len(values)), + ) + engine.buckets[0].ProbeRows = int64(len(values)) + engine.probeKeyEval = makeTestEvalKeysFn() + jm, result, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) require.NoError(t, err) - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - defer bat.Clean(proc.Mp()) - - w := BucketWriter{Name: "aggregate_disk_token", Budget: generation} - var buf bytes.Buffer - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - first := w.diskReservation - require.NotNil(t, first) - firstSize := first.Size() - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - require.Same(t, first, w.diskReservation) - require.Greater(t, w.diskReservation.Size(), firstSize) - require.Equal(t, w.diskReservation.Size(), generation.SpillDiskUsed()) - w.Close() - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) + require.Nil(t, jm) + require.Equal(t, BucketReSpilled, result) + var buildRows, probeRows int64 + for _, bucket := range engine.buckets { + buildRows += bucket.BuildRows + probeRows += bucket.ProbeRows + } + require.Equal(t, int64(len(values)), buildRows) + require.Equal(t, int64(len(values)), probeRows) + engine.Cleanup(proc) } -func TestBucketReaderEOF(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - var buf bytes.Buffer - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_eof") - require.NoError(t, err) - w := BucketWriter{Name: "test_eof", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - FlushBucketBatch(proc, bat, &w, &buf, nil) - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - _, err = reader.ReadBatch(proc, reuseBat) - require.Equal(t, io.EOF, err) - reader.Close() +func TestSpillRejectsCompleteRecordTruncation(t *testing.T) { + for _, test := range []struct { + name string + threshold int64 + }{ + {name: "rebuild", threshold: 1 << 30}, + {name: "re-spill", threshold: 1}, + } { + for _, metadataRows := range []int64{0, 6} { + t.Run(fmt.Sprintf("%s/metadata-%d", test.name, metadataRows), func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + first := makeInt32Batch(proc, []int32{1, 2, 3}) + defer first.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + SpillThreshold: test.threshold, + NeedsBuildForEmptyProbe: true, + }) + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile( + writeBuildFile(proc, fmt.Sprintf("truncate-%s-%d", test.name, metadataRows), first), + metadataRows, + ), + }) + _, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorContains(t, err, "row count") + engine.Cleanup(proc) + }) + } + } } -func TestBucketReaderCorruptedMagic(t *testing.T) { +func TestSpillRejectsPhysicalTruncationBeforeFirstRecord(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) + baseline := proc.Mp().CurrNB() + first := makeInt32Batch(proc, []int32{1}) + second := makeInt32Batch(proc, []int32{2}) + defer first.Clean(proc.Mp()) + defer second.Clean(proc.Mp()) - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_corrupt") + file := writeBuildRecords(proc, t.Name(), first, second) + info, err := file.Stat() require.NoError(t, err) - - // Write a valid batch via FlushBucketBatch, then corrupt the magic. - var buf bytes.Buffer - w := BucketWriter{Name: "test_corrupt", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) + require.NoError(t, file.Truncate(int64(len(marshalTestSpillRecord(first))))) + _, err = file.Seek(0, io.SeekStart) require.NoError(t, err) - // Overwrite last 8 bytes (magic) with zeros. - f.Seek(-8, io.SeekEnd) - var zeroMagic uint64 - f.Write(types.EncodeUint64(&zeroMagic)) - f.Seek(0, io.SeekStart) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedsBuildForEmptyProbe: true, + }) + engine.InitFromSpilledFiles([]*message.SpillFile{ + message.NewSpillFile(file, 2, uint64(info.Size()), nil), + }) + jm, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.Nil(t, jm) + require.ErrorContains(t, err, "corrupted spill file size") + engine.Cleanup(proc) + require.Zero(t, engine.allocation.account.Snapshot().Used) + require.Equal(t, baseline, proc.Mp().CurrNB()) +} - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(f) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.Error(t, err) - require.Contains(t, err.Error(), "corrupted") - require.Equal(t, uint64(64<<10), generation.Used(), "failed read must release its decoded-batch lease") - reader.Close() - require.Zero(t, generation.Used()) - f.Close() +func TestProbeRejectsCompleteRecordTruncation(t *testing.T) { + for _, test := range []struct { + name string + values []int32 + metadataRows int64 + rebuildError bool + firstError bool + }{ + {name: "zero metadata", values: []int32{1}, metadataRows: 0, rebuildError: true}, + {name: "row excess", values: []int32{1, 2}, metadataRows: 1, firstError: true}, + {name: "complete record truncation", values: []int32{1}, metadataRows: 2}, + } { + t.Run(test.name, func(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + defer proc.Free() + build := makeInt32Batch(proc, []int32{1}) + probe := makeInt32Batch(proc, test.values) + defer build.Clean(proc.Mp()) + defer probe.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedsBuildForEmptyProbe: true, + }) + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile(writeBuildFile(proc, "probe-build", build), 1), + }) + engine.buckets[0].ProbeFd = newTestSpillFile( + writeBuildFile(proc, "probe-data", probe), + test.metadataRows, + ) + engine.buckets[0].ProbeRows = test.metadataRows + jm, result, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + if test.rebuildError { + require.Nil(t, jm) + require.ErrorContains(t, err, "row count") + engine.Cleanup(proc) + return + } + require.NoError(t, err) + require.Equal(t, BucketReady, result) + jm.Free() + got, err := engine.NextProbeBatch(proc) + if test.firstError { + require.Nil(t, got) + require.ErrorContains(t, err, "row count") + } else { + require.NoError(t, err) + require.Equal(t, 1, got.RowCount()) + _, err = engine.NextProbeBatch(proc) + require.ErrorContains(t, err, "row count") + } + engine.Cleanup(proc) + }) + } } -func TestBucketReaderTruncatedMagic(t *testing.T) { +func TestReSpillRejectsProbeRowMetadata(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1}) - fd := writeBuildFile(proc, "test_truncated_magic", bat) - bat.Clean(proc.Mp()) - info, err := fd.Stat() - require.NoError(t, err) - require.NoError(t, fd.Truncate(info.Size()-4)) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, io.ErrUnexpectedEOF) - reader.Close() -} - -func TestBucketWriterHandOffFd(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - spillfs, _ := proc.GetSpillFileService() - f, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_handoff") - w := BucketWriter{Fd: f} - fd := w.HandOffFd() - require.NotNil(t, fd) - require.Nil(t, w.Fd) - require.False(t, w.Created()) - fd.Close() - - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(2) - require.NoError(t, err) - accounted := BucketWriter{Name: "test_accounted_raw_handoff", Budget: generation} - bat := makeInt32Batch(proc, []int32{1}) - var buf bytes.Buffer - require.NoError(t, FlushBucketBatch(proc, bat, &accounted, &buf, nil)) - bat.Clean(proc.Mp()) - require.Nil(t, accounted.HandOffFd(), "raw handoff must not orphan accounting tokens") - accounted.Close() - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} - -func TestMakeBucketWriters(t *testing.T) { - writers := MakeBucketWriters("test") - require.Equal(t, SpillNumBuckets, len(writers)) - for i := range writers { - require.NotEmpty(t, writers[i].Name) - require.Nil(t, writers[i].Fd) - } -} - -type countingMutableFileService struct { - fileservice.MutableFileService - ensureCalls int - closeCalls int -} - -func (s *countingMutableFileService) EnsureDir(ctx context.Context, path string) error { - s.ensureCalls++ - return s.MutableFileService.EnsureDir(ctx, path) -} - -func (s *countingMutableFileService) Close(ctx context.Context) { - s.closeCalls++ - s.MutableFileService.Close(ctx) -} - -func TestSpillEngineSharesFileServiceAcrossWriters(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - local, err := fileservice.Get[fileservice.MutableFileService]( - proc.GetFileService(), - defines.LocalFileServiceName, - ) - require.NoError(t, err) - countingLocal := &countingMutableFileService{MutableFileService: local} - services, err := fileservice.NewFileServices("", countingLocal) - require.NoError(t, err) - proc.SetFileService(services) - - engine := NewSpillEngine(SpillEngineConfig{}) - first := engine.makeBucketWriters("cached_first") - second := engine.makeBucketWriters("cached_second") - require.Same(t, first[0].spillFS, second[0].spillFS) - require.Same(t, &engine.spillFS, first[0].spillFS) - require.Zero(t, countingLocal.ensureCalls, "resolution remains lazy") - - require.NoError(t, writeBucketPayload(proc, []byte("first"), 1, &first[0], nil)) - require.Equal(t, 1, countingLocal.ensureCalls) - cached := engine.spillFS.fs - require.NotNil(t, cached) - - require.NoError(t, writeBucketPayload(proc, []byte("second"), 1, &second[1], nil)) - require.Equal(t, 1, countingLocal.ensureCalls, "all engine writers reuse one resolved service") - require.Equal(t, cached, engine.spillFS.fs) - - first[0].Close() - second[1].Close() - - // The service is borrowed. Cleanup releases engine-owned files and memory, - // but must neither close nor invalidate the process-owned service. - borrowed := &countingMutableFileService{MutableFileService: cached} - engine.spillFS.fs = borrowed - engine.Cleanup(proc) - require.Zero(t, borrowed.closeCalls) - file, err := borrowed.CreateAndRemoveFile(proc.Ctx, "after_engine_cleanup") - require.NoError(t, err) - require.NoError(t, file.Close()) - - // A writer constructed outside SpillEngine has no shared cache and keeps - // the historical process lookup fallback. - direct := BucketWriter{Name: "direct_writer_fallback"} - require.NoError(t, writeBucketPayload(proc, []byte("direct"), 1, &direct, nil)) - require.Equal(t, 2, countingLocal.ensureCalls) - direct.Close() -} - -func TestScatterProbeTableRejectsRecursiveMarker(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - marker := batch.NewWithSize(0) - marker.SetRowCount(1) - marker.SetLast() - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap([]*os.File{nil}) - called := false - err := engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - if called { - return nil, nil - } - called = true - return marker, nil - }, - nil, - func(*batch.Batch) ([]*vector.Vector, error) { - t.Fatal("recursive marker must not be evaluated as data") - return nil, nil - }, - ) - require.Error(t, err) - require.Contains(t, err.Error(), "recursive input") - engine.Cleanup(proc) -} - -func TestScatterProbeTableErrors(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - wantErr := errors.New("scatter probe failure") - - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap([]*os.File{nil}) - err := engine.ScatterProbeTable(proc, - func() (*batch.Batch, error) { return nil, wantErr }, nil, - func(*batch.Batch) ([]*vector.Vector, error) { return nil, nil }) - require.ErrorIs(t, err, wantErr) - engine.Cleanup(proc) - - bat := makeInt32Batch(proc, []int32{1}) - engine = NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap([]*os.File{nil}) - err = engine.ScatterProbeTable(proc, - func() (*batch.Batch, error) { return bat, nil }, nil, - func(*batch.Batch) ([]*vector.Vector, error) { return nil, wantErr }) - require.ErrorIs(t, err, wantErr) - engine.Cleanup(proc) - bat.Clean(proc.Mp()) -} - -func TestReusableBufferPool(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - pool := ReusableBufferPool{} - bufs := pool.Acquire(SpillNumBuckets) - require.Equal(t, SpillNumBuckets, len(bufs)) - for i := range bufs { - require.Nil(t, bufs[i]) - } - pool.Release(proc) -} - -func TestBucketReaderCancellationStopsBeforeMergingNextRecord(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - - first := makeInt32Batch(proc, []int32{1, 2}) - second := makeInt32Batch(proc, []int32{3, 4}) - var encoded bytes.Buffer - require.NoError(t, marshalSpillRecord(first, &encoded)) - firstRecord := bytes.Clone(encoded.Bytes()) - require.NoError(t, marshalSpillRecord(second, &encoded)) - stream := append(firstRecord, encoded.Bytes()...) - first.Clean(proc.Mp()) - second.Clean(proc.Mp()) - - source := &boundaryCancelReader{ - reader: bytes.NewReader(stream), - boundary: int64(len(firstRecord)), - cancel: func() { proc.Cancel(context.Canceled) }, - } - fd, err := os.CreateTemp(t.TempDir(), "bucket-reader-cancel") - require.NoError(t, err) - reader := BucketReader{ - fd: fd, - reader: bufio.NewReaderSize(source, 16), - mergeRecords: true, - } - reuseBat := batch.NewOffHeapWithSize(0) - - got, err := reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, context.Canceled) - require.Nil(t, got) - require.True(t, source.canceled) - require.Equal(t, int64(len(firstRecord)+16), source.read, - "reader may inspect the next header but must not decode its payload after cancellation") - require.Zero(t, reuseBat.RowCount()) - - reader.Close() - reuseBat.Clean(proc.Mp()) - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestBucketReaderEmptyFile(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - r := BucketReader{} - r.ResetForFd(nil) - reuseBat := batch.NewOffHeapWithSize(0) - _, err := r.ReadBatch(proc, reuseBat) - require.Equal(t, io.EOF, err) -} - -func TestLazySpillFileCreation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_lazy") - for i := range writers { - require.Nil(t, writers[i].Fd, "all writers should start with nil Fd") - } - - // Write a batch that will populate some buckets - var buf bytes.Buffer - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5}, nil, proc.Mp()) - bat.SetRowCount(5) - buffers := make([]*batch.Batch, len(writers)) - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Flush remaining buffers — files are created lazily on first write - for i, b := range buffers { - if b != nil && b.RowCount() > 0 { - err := FlushBucketBatch(proc, b, &writers[i], &buf, nil) - require.NoError(t, err) - require.True(t, writers[i].Created(), "writer should have created file on first flush") - } - } - - // Clean up - for i := range writers { - writers[i].Close() - } -} - -func TestReaderRowCountMismatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_mismatch") - require.NoError(t, err) - - // Write valid batch then corrupt the row count in the header. - var buf bytes.Buffer - w := BucketWriter{Name: "test", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - - // Overwrite the count (first 8 bytes) with a wrong value. - f.Seek(0, io.SeekStart) - wrongCnt := int64(999) - f.Write(types.EncodeInt64(&wrongCnt)) - f.Seek(0, io.SeekStart) - - reader := BucketReader{} - reader.ResetForFd(f) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.Error(t, err) - require.Contains(t, err.Error(), "mismatch") - reader.Close() - f.Close() -} - -func TestScatterBatchDistribution(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - require.NoError(t, scatterImpl(proc, batch.NewWithSize(0), nil, nil, nil, 0, nil, nil, nil, nil)) - - writers := MakeBucketWriters("test_dist") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - nRows := SpillNumBuckets * 100 - vals := make([]int32, nRows) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(nRows) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Most buckets should have data with enough rows. - nonEmpty := 0 - for _, b := range buffers { - if b != nil && b.RowCount() > 0 { - nonEmpty++ - } - } - require.Greater(t, nonEmpty, SpillNumBuckets/2, "at least half the buckets should have data") - - // Total rows should be preserved. - totalRows := 0 - for i, b := range buffers { - if b != nil { - FlushBucketBatch(proc, b, &writers[i], &buf, nil) - totalRows += b.RowCount() - } - } - require.Equal(t, nRows, totalRows) -} - -func TestBucketReaderPartialRead(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_partial") - require.NoError(t, err) - - // Write incomplete data (only a count, no batch body). - cnt := int64(5) - f.Write(types.EncodeInt64(&cnt)) - f.Seek(0, io.SeekStart) - - reader := BucketReader{} - reader.ResetForFd(f) - reuseBat := batch.NewOffHeapWithSize(0) - _, err = reader.ReadBatch(proc, reuseBat) - require.Error(t, err) - reader.Close() - f.Close() -} - -func TestBucketReaderDoubleClose(t *testing.T) { - r := BucketReader{} - r.Close() - r.Close() // should not panic -} - -func TestComputeXXHashWithNulls(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3}, []uint64{1}, mp) // null at index 1 - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{vec}, hashValues, 0) - require.NotEqual(t, uint64(0), hashValues[0]) - require.NotEqual(t, uint64(0), hashValues[2]) -} - -func TestComputeXXHashMultipleColumns(t *testing.T) { - mp := mpool.MustNewZero() - vec1 := testutil.MakeInt32Vector([]int32{1, 1, 2}, nil, mp) - vec2 := testutil.MakeVarcharVector([]string{"a", "b", "a"}, nil, mp) - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{vec1, vec2}, hashValues, 0) - // Same (col1, col2) pairs should hash differently. - require.NotEqual(t, hashValues[0], hashValues[1]) - require.NotEqual(t, hashValues[0], hashValues[2]) -} - -func TestHandOffFdSeeksToStart(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_seek") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test_seek", Fd: f} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{10, 20, 30}, nil, proc.Mp()) - bat.SetRowCount(3) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - - // Position should be past data. - pos, _ := w.Fd.Seek(0, io.SeekCurrent) - require.Greater(t, pos, int64(0)) - - fd := w.HandOffFd() - require.NotNil(t, fd) - pos, _ = fd.Seek(0, io.SeekCurrent) - require.Equal(t, int64(0), pos, "HandOffFd must seek to start") - fd.Close() -} - -func TestComputeXXHashMultipleTypes(t *testing.T) { - mp := mpool.MustNewZero() - tests := []struct { - name string - vec *vector.Vector - }{ - {"int8", testutil.MakeInt8Vector([]int8{1, 2, 3}, nil, mp)}, - {"int16", testutil.MakeInt16Vector([]int16{100, 200, 300}, nil, mp)}, - {"int64", testutil.MakeInt64Vector([]int64{1000, 2000, 3000}, nil, mp)}, - {"uint32", testutil.MakeUint32Vector([]uint32{10, 20, 30}, nil, mp)}, - {"float32", testutil.MakeFloat32Vector([]float32{1.1, 2.2, 3.3}, nil, mp)}, - {"float64", testutil.MakeFloat64Vector([]float64{10.1, 20.2, 30.3}, nil, mp)}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - hashValues := make([]uint64, 3) - ComputeXXHash([]*vector.Vector{tt.vec}, hashValues, 0) - require.NotEqual(t, uint64(0), hashValues[0]) - }) - } -} - -func TestScatterBatchLargeData(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_large") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - // Large enough to trigger internal flush (>8192 rows). - size := 10000 - vals := make([]int32, size) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(size) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Verify total rows preserved. - totalRows := 0 - for i, b := range buffers { - if b != nil { - FlushBucketBatch(proc, b, &writers[i], &buf, nil) - totalRows += b.RowCount() - } - } - require.Equal(t, size, totalRows) -} - -func TestResetForFdReusesReader(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - var buf bytes.Buffer - f1, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_reuse_1") - w1 := BucketWriter{Name: "test1", Fd: f1} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bat.SetRowCount(3) - FlushBucketBatch(proc, bat, &w1, &buf, nil) - fd1 := w1.HandOffFd() - - f2, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_reuse_2") - w2 := BucketWriter{Name: "test2", Fd: f2} - FlushBucketBatch(proc, bat, &w2, &buf, nil) - fd2 := w2.HandOffFd() - - r := BucketReader{} - r.ResetForFd(fd1) - require.NotNil(t, r.fd) - - // Second ResetForFd reuses internal state. - r.ResetForFd(fd2) - require.NotNil(t, r.fd) - - r.Close() - fd1.Close() -} - -func TestFlushBucketBatchMultipleCalls(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_multi_flush") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test_multi_flush", Fd: f} - totalRows := 0 - for i := 0; i < 10; i++ { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{int32(i)}, nil, proc.Mp()) - bat.SetRowCount(1) - err := FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - totalRows++ - } - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - readRows := 0 - for { - got, err := reader.ReadBatch(proc, reuseBat) - if err == io.EOF { - break - } - require.NoError(t, err) - readRows += got.RowCount() - } - require.Equal(t, totalRows, readRows) - reader.Close() -} - -func TestHashDistribution(t *testing.T) { - mp := mpool.MustNewZero() - vec := testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil, mp) - hashValues := make([]uint64, 20) - ComputeXXHash([]*vector.Vector{vec}, hashValues, 0) - - bucketCounts := make([]int, SpillNumBuckets) - for _, h := range hashValues { - bucketCounts[h&(SpillNumBuckets-1)]++ - } - nonEmpty := 0 - for _, c := range bucketCounts { - if c > 0 { - nonEmpty++ - } - } - require.Greater(t, nonEmpty, 1, "hashes must distribute across multiple buckets") -} - -func TestSpillFileCleanup(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - - file, err := spillfs.CreateFile(context.Background(), "test_cleanup") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test_cleanup", Fd: file} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - file.Close() - - // File should still exist (it was CreateFile, not CreateAndRemoveFile). - f2, err := spillfs.OpenFile(context.Background(), "test_cleanup") - require.NoError(t, err) - f2.Close() - - spillfs.RemoveFile(context.Background(), "test_cleanup") -} - -func TestFileWriteError(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - file, err := spillfs.CreateFile(context.Background(), "test_error") - require.NoError(t, err) - file.Close() // close before write - - var buf bytes.Buffer - w := BucketWriter{Name: "test_error", Fd: file} - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1}, nil, proc.Mp()) - bat.SetRowCount(1) - err = FlushBucketBatch(proc, bat, &w, &buf, nil) - require.Error(t, err) - - spillfs.RemoveFile(context.Background(), "test_error") -} - -func TestScatterBatchWithNulls(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_null") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4}, []uint64{1}, proc.Mp()) // null at index 1 - bat.SetRowCount(4) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - // Should not panic with nulls. -} - -func TestReaderBatchReuse(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - var buf bytes.Buffer - f, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_reuse_read") - w := BucketWriter{Name: "test", Fd: f} - - // Write batches with different sizes. - for _, size := range []int{5, 2} { - vals := make([]int32, size) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(size) - FlushBucketBatch(proc, bat, &w, &buf, nil) - } - - fd := w.HandOffFd() - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - before := generation.Snapshot() - - // Read with the same reuseBat. The second record admits the one-time - // old-plus-new transition and then keeps the bounded high-water lease. - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 5, got.RowCount()) - afterFirst := generation.Snapshot() - require.Equal(t, before.ReserveCount+1, afterFirst.ReserveCount) - require.Equal(t, before.ReconcileCount, afterFirst.ReconcileCount) - - got, err = reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - afterSecond := generation.Snapshot() - require.Equal(t, afterFirst.ReserveCount+1, afterSecond.ReserveCount) - require.Equal(t, afterFirst.ReconcileCount, afterSecond.ReconcileCount) - require.Equal(t, afterFirst.ReleaseCount, afterSecond.ReleaseCount) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestReaderBatchLeaseGrowsForLargerRecord(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_lease_grow") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_lease_grow", Fd: f} - for _, size := range []int{2, 1_000} { - vals := make([]int32, size) - bat := makeInt32Batch(proc, vals) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(w.HandOffFd()) - reuseBat := batch.NewOffHeapWithSize(0) - before := generation.Snapshot() - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - afterFirst := generation.Snapshot() - require.Equal(t, before.ReserveCount+1, afterFirst.ReserveCount) - - got, err = reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 1_000, got.RowCount()) - afterSecond := generation.Snapshot() - require.Equal(t, afterFirst.ReserveCount+1, afterSecond.ReserveCount, "larger record should grow the existing lease once") - require.Greater(t, afterSecond.Used, afterFirst.Used) - require.Equal(t, afterFirst.ReconcileCount, afterSecond.ReconcileCount) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestDecodedBatchReusePeakCoversMpoolGrowth(t *testing.T) { - const oldCapacity = int64(4 << 20) - required := oldCapacity + 1 - newCapacity, ok := mpool.GrowCapacity(oldCapacity, required) - require.True(t, ok) - projected, ok := decodedBatchProjectedBytes(uint64(required), 1, 1) - require.True(t, ok) - peak, ok := decodedBatchReusePeakBytes(uint64(oldCapacity), projected, uint64(required)) - require.True(t, ok) - require.GreaterOrEqual(t, peak, uint64(oldCapacity+newCapacity)) -} - -func TestReaderBatchReuseFallsBackBeforeTransientGrowth(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_reuse_transient") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_reuse_transient", Fd: f} - for _, width := range []int{1_024, 1_025} { - values := make([]string, 4_096) - for i := range values { - values[i] = strings.Repeat("x", width) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - - fd := w.HandOffFd() - var projected [2]uint64 - for i := range projected { - var header [16]byte - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - rows := types.DecodeInt64(header[:8]) - payload := types.DecodeInt64(header[8:]) - require.Positive(t, rows) - require.Positive(t, payload) - var ok bool - projected[i], ok = decodedBatchProjectedBytes(uint64(payload), rows, 1) - require.True(t, ok) - _, err = fd.Seek(payload+8, io.SeekCurrent) - require.NoError(t, err) - } - require.Greater(t, projected[1], projected[0]) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - // The cap fits the retained record plus one logical new payload, but not the - // allocator's 1.25x replacement capacity. The reader must release the old - // lease and decode the second record fresh. - cap := uint64(64<<10) + projected[0] + projected[1] - budget := process.MustNewHashBuildBudget(cap, cap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - baseline := uint64(proc.Mp().CurrNB()) - epoch := proc.Mp().StartResourcePeakEpoch() - require.NotNil(t, epoch) - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 4_096, got.RowCount()) - got, err = reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 4_096, got.RowCount()) - - peak, ok := proc.Mp().EndResourcePeakEpoch(epoch) - require.True(t, ok) - require.LessOrEqual(t, peak, baseline+projected[1]) - require.LessOrEqual(t, generation.Peak(), cap) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestReaderBatchLeaseUsesSinglePayloadEstimate(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget := process.MustNewHashBuildBudget(64<<20, 64<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_lease_trim") - require.NoError(t, err) - values := make([]string, 4_096) - for i := range values { - values[i] = strings.Repeat("x", 1_024) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(values, nil, proc.Mp()) - bat.SetRowCount(len(values)) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_lease_trim", Fd: f} - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - fd := w.HandOffFd() - var header [16]byte - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - payload := types.DecodeInt64(header[8:]) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - before := generation.Snapshot() - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - actual, ok := batchRetainedBytes(got) - require.True(t, ok) - after := generation.Snapshot() - require.Equal(t, before.ReconcileCount, after.ReconcileCount) - require.GreaterOrEqual(t, after.Used, before.Used+actual) - projected, ok := decodedBatchProjectedBytes(uint64(payload), int64(got.RowCount()), int32(len(got.Vecs))) - require.True(t, ok) - require.Equal(t, before.Used+projected, after.Used) - - reuseBat.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestMarshalSpillRecordPreallocatesSinglePayload(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - bat := batch.NewWithSize(1) - var err error - bat.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 4<<20), 64, proc.Mp(), - ) - require.NoError(t, err) - bat.SetRowCount(64) - defer bat.Clean(proc.Mp()) - - buf := bytes.NewBuffer(make([]byte, 0, 1<<20)) - require.NoError(t, marshalSpillRecord(bat, buf)) - base := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > base { - base = size - } - require.Equal(t, base+128+24, uint64(buf.Cap())) - - small := batch.NewWithSize(1) - small.Vecs[0], err = vector.NewConstBytes( - types.T_varchar.ToType(), make([]byte, 1024), 1, proc.Mp(), - ) - require.NoError(t, err) - small.SetRowCount(1) - defer small.Clean(proc.Mp()) - require.NoError(t, marshalSpillRecord(small, buf), - "a retained large serialization buffer must be reusable for a smaller batch") -} - -func TestReaderBatchLeaseGrowRejectionReleasesToken(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_lease_reject") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_lease_reject", Fd: f} - for _, size := range []int{2, 1_000} { - bat := makeInt32Batch(proc, make([]int32, size)) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - fd := w.HandOffFd() - - var header [16]byte - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - firstPayload := types.DecodeInt64(header[8:]) - _, err = fd.Seek(firstPayload+8, io.SeekCurrent) - require.NoError(t, err) - _, err = io.ReadFull(fd, header[:]) - require.NoError(t, err) - secondPayload := types.DecodeInt64(header[8:]) - require.Greater(t, secondPayload, firstPayload) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - firstProjected, ok := decodedBatchProjectedBytes(uint64(firstPayload), 2, 1) - require.True(t, ok) - secondProjected, ok := decodedBatchProjectedBytes(uint64(secondPayload), 1_000, 1) - require.True(t, ok) - cap := uint64(64<<10) + firstProjected + (secondProjected-firstProjected)/2 - budget := process.MustNewHashBuildBudget(cap, cap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - _, err = reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, uint64(64<<10), generation.Used(), "grow rejection must release the existing decoded-batch lease") - - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestReaderBatchClosedLeaseDoesNotRetryAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_read_closed_lease") - require.NoError(t, err) - var buf bytes.Buffer - w := BucketWriter{Name: "test_read_closed_lease", Fd: f} - for _, size := range []int{2, 1_000} { - bat := makeInt32Batch(proc, make([]int32, size)) - require.NoError(t, FlushBucketBatch(proc, bat, &w, &buf, nil)) - bat.Clean(proc.Mp()) - } - - budget := process.MustNewHashBuildBudget(8<<20, 8<<20) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - reader := BucketReader{} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(w.HandOffFd()) - reuseBat := batch.NewOffHeapWithSize(0) - - got, err := reader.ReadBatch(proc, reuseBat) - require.NoError(t, err) - require.Equal(t, 2, got.RowCount()) - before := generation.Snapshot() - generation.Close() - - _, err = reader.ReadBatch(proc, reuseBat) - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - after := generation.Snapshot() - require.Equal(t, before.RejectCount+1, after.RejectCount, - "a closed lease must not fall through to a second Reserve attempt") - require.Nil(t, reader.batchToken) - require.Zero(t, reader.batchCharge) - require.Zero(t, reuseBat.RowCount()) - - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestScatterWithMultiColumn(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_multi_col") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - bat := batch.NewWithSize(2) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3, 4, 5}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeVarcharVector([]string{"a", "b", "c", "d", "e"}, nil, proc.Mp()) - bat.SetRowCount(5) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // All 5 rows must be distributed across buffers. - totalRows := 0 - hasTwoCols := false - for _, b := range buffers { - if b != nil && b.RowCount() > 0 { - totalRows += b.RowCount() - if len(b.Vecs) == 2 { - hasTwoCols = true - } - } - } - require.Equal(t, 5, totalRows, "all rows must be accounted for in buffers") - require.True(t, hasTwoCols, "buffer batches must preserve column count") -} - -func TestScatterLargeVarchar(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_large_varchar") - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - - size := 100 - vals := make([]string, size) - for i := range vals { - vals[i] = fmt.Sprintf("large_string_value_%d_with_padding", i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeVarcharVector(vals, nil, proc.Mp()) - bat.SetRowCount(size) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - totalRows := 0 - for i, b := range buffers { - if b != nil && b.RowCount() > 0 { - FlushBucketBatch(proc, b, &writers[i], &buf, nil) - totalRows += b.RowCount() - } - } - require.Equal(t, size, totalRows) -} - -func TestBucketBufferReuse(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_buf_reuse") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test", Fd: f} - - // Reuse same writer across multiple batches. - for range 2 { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2}, nil, proc.Mp()) - bat.SetRowCount(2) - err := FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - } - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - totalRows := 0 - for { - got, err := reader.ReadBatch(proc, reuseBat) - if err == io.EOF { - break - } - require.NoError(t, err) - totalRows += got.RowCount() - } - require.Equal(t, 4, totalRows) // 2 batches × 2 rows - reader.Close() -} - -func TestReusableBufferPoolWithData(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - pool := ReusableBufferPool{} - bufs := pool.Acquire(3) - - // Populate buffers with data. - for i := range bufs { - bufs[i] = batch.NewWithSize(1) - bufs[i].Vecs[0] = testutil.MakeInt32Vector([]int32{1, 2, 3}, nil, proc.Mp()) - bufs[i].SetRowCount(3) - } - - // Release should clean everything. - pool.Release(proc) -} - -func TestSpillFileFormatMultipleBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - f, err := spillfs.CreateAndRemoveFile(context.Background(), "test_format") - require.NoError(t, err) - - var buf bytes.Buffer - w := BucketWriter{Name: "test", Fd: f} - for i := 0; i < 3; i++ { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector([]int32{int32(i * 10), int32(i*10 + 1)}, nil, proc.Mp()) - bat.SetRowCount(2) - err := FlushBucketBatch(proc, bat, &w, &buf, nil) - require.NoError(t, err) - } - - fd := w.HandOffFd() - reader := BucketReader{} - reader.ResetForFd(fd) - reuseBat := batch.NewOffHeapWithSize(0) - batchCount := 0 - totalRows := 0 - for { - got, err := reader.ReadBatch(proc, reuseBat) - if err == io.EOF { - break - } - require.NoError(t, err) - batchCount++ - totalRows += got.RowCount() - } - require.Equal(t, 3, batchCount) - require.Equal(t, 6, totalRows) - reader.Close() -} - -func TestBucketReaderMergesAdjacentAccountedRecords(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var buf bytes.Buffer - writer := BucketWriter{Name: "merge_records", Budget: generation} - for _, values := range [][]int32{{1, 2}, {3, 4, 5}} { - bat := makeInt32Batch(proc, values) - require.NoError(t, FlushBucketBatch(proc, bat, &writer, &buf, nil)) - bat.Clean(proc.Mp()) - } - file, err := writer.handOffSpillFile() - require.NoError(t, err) - - reader := BucketReader{mergeRecords: true} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForSpillFile(file) - reuse := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 5, got.RowCount()) - require.Positive(t, generation.Used()) - _, err = reader.ReadBatch(proc, reuse) - require.ErrorIs(t, err, io.EOF) - reuse.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} - -func TestBucketReaderMergeRejectsTruncatedTrailingHeader(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - var payload bytes.Buffer - require.NoError(t, marshalSpillRecord(bat, &payload)) - bat.Clean(proc.Mp()) - - fd, err := os.CreateTemp(t.TempDir(), "truncated-spill") - require.NoError(t, err) - t.Cleanup(func() { _ = fd.Close() }) - _, err = fd.Write(payload.Bytes()) - require.NoError(t, err) - // A clean file boundary has zero bytes left. Any non-empty fragment of the - // next 16-byte frame header is corruption and must not be accepted as EOF. - _, err = fd.Write(types.EncodeInt64(new(int64))) - require.NoError(t, err) - _, err = fd.Seek(0, io.SeekStart) - require.NoError(t, err) - - reader := BucketReader{mergeRecords: true} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForFd(fd) - reuse := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuse) - require.Nil(t, got) - require.ErrorIs(t, err, io.ErrUnexpectedEOF) - require.Zero(t, reuse.RowCount()) - require.Nil(t, reader.batchToken) - require.Zero(t, reader.batchCharge) - - reuse.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) -} - -func TestBucketReaderMergeRecordsRespectsBatchBoundary(t *testing.T) { - tests := []struct { - name string - recordRows []int - wantRows []int - }{ - { - name: "two medium records stay separate", - recordRows: []int{5000, 5000}, - wantRows: []int{5000, 5000}, - }, - { - name: "records exactly fill the boundary", - recordRows: []int{8191, 1}, - wantRows: []int{8192}, - }, - { - name: "record crossing the boundary stays separate", - recordRows: []int{8191, 2}, - wantRows: []int{8191, 2}, - }, - { - name: "one oversized source record remains indivisible", - recordRows: []int{9000}, - wantRows: []int{9000}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - var buf bytes.Buffer - writer := BucketWriter{Name: "merge_boundary", Budget: generation} - for record, rows := range tt.recordRows { - values := make([]int32, rows) - for row := range values { - values[row] = int32(record*10000 + row) - } - bat := makeInt32Batch(proc, values) - require.NoError(t, FlushBucketBatch(proc, bat, &writer, &buf, nil)) - bat.Clean(proc.Mp()) - } - file, err := writer.handOffSpillFile() - require.NoError(t, err) - - reader := BucketReader{mergeRecords: true} - require.NoError(t, reader.EnsureBuffer(generation)) - reader.ResetForSpillFile(file) - reuse := batch.NewOffHeapWithSize(0) - var gotRows []int - for { - got, err := reader.ReadBatch(proc, reuse) - if err == io.EOF { - break - } - require.NoError(t, err) - gotRows = append(gotRows, got.RowCount()) - } - require.Equal(t, tt.wantRows, gotRows) - - reuse.Clean(proc.Mp()) - reader.Close() - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - }) - } -} - -func TestBucketReaderMergeErrorReleasesAllOwnership(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - readerToken, err := generation.Reserve(1) - require.NoError(t, err) - sourceToken, err := generation.Reserve(1) - require.NoError(t, err) - extraToken, err := generation.Reserve(1) - require.NoError(t, err) - reader := BucketReader{batchToken: readerToken, batchCharge: 1} - dst := makeInt32Batch(proc, []int32{1}) - src := makeInt32Batch(proc, []int32{2}) - want := errors.New("merge failed") - require.ErrorIs(t, reader.mergeReadError(proc, dst, src, sourceToken, want, extraToken, nil), want) - require.Nil(t, reader.batchToken) - require.Zero(t, reader.batchCharge) - require.Zero(t, generation.Used()) - require.True(t, readerToken.Released()) - require.True(t, sourceToken.Released()) - require.True(t, extraToken.Released()) - - require.ErrorIs(t, reader.mergeReadError(proc, nil, nil, nil, want, nil), want) -} - -func TestSpillEngineInitFromOwnedFilesAndErrorClassification(t *testing.T) { - first, err := os.CreateTemp(t.TempDir(), "owned-build") - require.NoError(t, err) - owned := message.NewSpillFile(first, 7, 11, nil) - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledFiles([]*message.SpillFile{owned, nil}) - require.Len(t, engine.buckets, 2) - require.Same(t, owned, engine.buckets[0].BuildFd) - require.Equal(t, int64(7), engine.buckets[0].BuildRows) - require.Equal(t, 1, engine.buckets[0].Depth) - require.Nil(t, engine.buckets[1].BuildFd) - require.Zero(t, engine.buckets[1].BuildRows) - - require.False(t, isBudgetAdmission(nil)) - require.False(t, isBudgetAdmission(io.EOF)) - require.True(t, isBudgetAdmission(process.ErrHashBuildBudgetAdmission)) - require.False(t, isBudgetAdmission(process.ErrHashBuildBudgetClosed)) - require.True(t, isBudgetAdmission(&process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - })) - require.False(t, isBudgetAdmission(&process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - })) - logicalNoProgress := noProgressError(3, nil) - require.ErrorIs(t, logicalNoProgress, process.ErrHashBuildBudgetAdmission) - require.Contains(t, logicalNoProgress.Error(), "depth 3") - require.Contains(t, logicalNoProgress.Error(), "reduce join-key skew") - require.NotContains(t, logicalNoProgress.Error(), process.ErrHashBuildBudgetAdmission.Error()) - - memoryAdmission := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorAdmission, - Resource: process.HashBuildBudgetResourceMemory, - Requested: 11, - Used: 13, - Cap: 17, - } - budgetNoProgress := noProgressError(4, memoryAdmission) - var budgetErr *process.HashBuildBudgetError - require.ErrorAs(t, budgetNoProgress, &budgetErr) - require.Equal(t, process.HashBuildBudgetResourceMemory, budgetErr.Resource) - require.Equal(t, uint64(11), budgetErr.Requested) - require.Equal(t, uint64(13), budgetErr.Used) - require.Equal(t, uint64(17), budgetErr.Cap) - require.Equal(t, "join spill cannot make progress at depth 4", budgetErr.Message) - require.NoError(t, owned.Close()) -} - -func TestSpillSizeHelpersRejectInvalidAndOverflowInputs(t *testing.T) { - require.ErrorIs(t, writeBucketPayload(nil, nil, 0, nil, nil), process.ErrHashBuildBudgetInvalid) - require.NoError(t, marshalSpillRecord(nil, &bytes.Buffer{})) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - var header bytes.Buffer - negativeRows := int64(-1) - zero := int64(0) - header.Write(types.EncodeInt64(&negativeRows)) - header.Write(types.EncodeInt64(&zero)) - reader := BucketReader{reader: bufio.NewReader(&header)} - _, _, _, err := reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.Error(t, err) - - makeHeader := func(batchSize int64) *bufio.Reader { - var data bytes.Buffer - rows := int64(0) - data.Write(types.EncodeInt64(&rows)) - data.Write(types.EncodeInt64(&batchSize)) - data.Write(make([]byte, 12)) - return bufio.NewReader(&data) - } - budgetForHeader, err := process.NewHashBuildBudget(1, 1) - require.NoError(t, err) - headerGeneration, err := budgetForHeader.OpenGeneration(1) - require.NoError(t, err) - reader = BucketReader{reader: makeHeader(math.MaxInt64), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - - var truncatedPayload bytes.Buffer - truncatedPayload.Write(types.EncodeInt64(&zero)) - one := int64(1) - truncatedPayload.Write(types.EncodeInt64(&one)) - truncatedPayload.Write(make([]byte, 4)) - reader = BucketReader{reader: bufio.NewReader(&truncatedPayload), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.Error(t, err) - - var mismatchedRows bytes.Buffer - twelve := int64(12) - mismatchedRows.Write(types.EncodeInt64(&one)) - mismatchedRows.Write(types.EncodeInt64(&twelve)) - mismatchedRows.Write(types.EncodeInt64(&zero)) - mismatchedRows.Write(make([]byte, 4)) - reader = BucketReader{reader: bufio.NewReader(&mismatchedRows), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.Error(t, err) - require.Contains(t, err.Error(), "row count mismatch") - - reader = BucketReader{reader: makeHeader(1), budget: headerGeneration} - _, _, _, err = reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), nil, 0, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - headerToken, err := headerGeneration.Reserve(1) - require.NoError(t, err) - reader = BucketReader{reader: makeHeader(1), budget: headerGeneration} - _, returnedToken, _, err := reader.readBatchRecord(proc, batch.NewOffHeapWithSize(0), headerToken, 1, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, returnedToken) - require.True(t, headerToken.Released()) - - _, ok := addUint64(math.MaxUint64, 1) - require.False(t, ok) - _, ok = mulUint64(math.MaxUint64, 2) - require.False(t, ok) - _, ok = decodedBatchReusePeakBytes(math.MaxUint64, 1, 1) - require.False(t, ok) - _, ok = decodedBatchReusePeakBytes(0, math.MaxUint64, 1) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(0, -1, 0) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(0, 0, -1) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(math.MaxUint64, 0, 0) - require.False(t, ok) - _, ok = decodedBatchProjectedBytes(0, math.MaxInt64, math.MaxInt32) - require.False(t, ok) - _, ok = batchRetainedMetadataBytes(1, math.MaxUint64) - require.False(t, ok) - _, ok = batchRetainedMetadataBytes(math.MaxUint64, 1) - require.False(t, ok) - _, ok = batchPayloadWithAllocationSlack(0, math.MaxUint64) - require.False(t, ok) - _, ok = batchPayloadWithAllocationSlack(math.MaxUint64, 0) - require.False(t, ok) - _, ok = intToUint64(-1) - require.False(t, ok) - _, ok = predictedCapacity(-1, 1) - require.False(t, ok) - _, ok = predictedCapacity(1, math.MaxUint64) - require.False(t, ok) - - _, ok = batchRetainedBytes(nil) - require.False(t, ok) - invalidRows := batch.NewOffHeapWithSize(0) - invalidRows.SetRowCount(-1) - _, ok = batchRetainedBytes(invalidRows) - require.False(t, ok) - _, err = scatterTransientBudgetBytes(nil, false) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - _, ok = (&SpillEngine{}).scatterCapacityGrowthBytes(-1, 0) - require.False(t, ok) - _, ok = (&SpillEngine{}).scatterCapacityGrowthBytes(math.MaxInt, 0) - require.False(t, ok) - - valid := batch.NewOffHeapWithSize(0) - valid.SetRowCount(0) - _, ok = predictMergedRetainedBytes(nil, valid) - require.False(t, ok) - invalidRows = batch.NewOffHeapWithSize(0) - invalidRows.SetRowCount(-1) - _, ok = predictMergedRetainedBytes(invalidRows, valid) - require.False(t, ok) - mismatched := batch.NewOffHeapWithSize(1) - mismatched.SetRowCount(0) - _, ok = predictMergedRetainedBytes(valid, mismatched) - require.False(t, ok) - nilVectorDst := batch.NewOffHeapWithSize(1) - nilVectorDst.SetRowCount(0) - nilVectorSrc := batch.NewOffHeapWithSize(1) - nilVectorSrc.SetRowCount(0) - _, ok = predictMergedRetainedBytes(nilVectorDst, nilVectorSrc) - require.False(t, ok) - hugeRowsDst := batch.NewOffHeapWithSize(0) - hugeRowsDst.SetRowCount(maxIntValue()) - hugeRowsSrc := batch.NewOffHeapWithSize(0) - hugeRowsSrc.SetRowCount(1) - _, ok = predictMergedRetainedBytes(hugeRowsDst, hugeRowsSrc) - require.False(t, ok) - - mp := mpool.MustNewZero() - fixed := testutil.MakeInt32Vector([]int32{1}, nil, mp) - defer fixed.Free(mp) - _, ok = mergedVarlenAreaAdd(nil, 1) - require.False(t, ok) - _, ok = mergedVarlenAreaAdd(fixed, 1) - require.False(t, ok) - constNull := vector.NewConstNull(types.T_varchar.ToType(), 1, mp) - defer constNull.Free(mp) - bytes, ok := mergedVarlenAreaAdd(constNull, 0) - require.True(t, ok) - require.Zero(t, bytes) - bytes, ok = mergedVarlenAreaAdd(constNull, 1) - require.True(t, ok) - require.Zero(t, bytes) - - require.NoError(t, reconcileReadReservation(nil, 0)) - budget, err := process.NewHashBuildBudget(10, 10) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - token, err := generation.Reserve(1) - require.NoError(t, err) - require.ErrorIs(t, reconcileReadReservation(token, 2), process.ErrHashBuildBudgetInvalid) - require.True(t, token.Release()) - require.ErrorIs(t, reconcileReadReservation(token, 0), process.ErrHashBuildReservationInactive) -} - -func TestScatterSkipsDisabledWriters(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - writers := MakeBucketWriters("test_skip") - for i := 0; i < len(writers); i += 2 { - writers[i].Name = "" // disable even buckets - } - - buffers := make([]*batch.Batch, len(writers)) - var buf bytes.Buffer - nRows := SpillNumBuckets * 100 - vals := make([]int32, nRows) - for i := range vals { - vals[i] = int32(i) - } - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(nRows) - - err := scatterImpl(proc, bat, bat.Vecs[:1], writers, buffers, 0, &buf, nil, nil, nil) - require.NoError(t, err) - - // Disabled buckets must have no buffer and no file created. - for i := 0; i < len(writers); i += 2 { - require.Nil(t, buffers[i], "disabled bucket %d must have no buffer", i) - require.False(t, writers[i].Created(), "disabled bucket %d must not have file", i) - } - - // Enabled buckets should have received data. - hashValues := make([]uint64, nRows) - ComputeXXHash(bat.Vecs[:1], hashValues, 0) - var expectedOddRows int - for _, h := range hashValues { - if h&uint64(SpillNumBuckets-1)&1 == 1 { - expectedOddRows++ - } - } - require.Greater(t, expectedOddRows, 0) - - var oddRows int - for i := 1; i < len(writers); i += 2 { - if buffers[i] != nil { - oddRows += buffers[i].RowCount() - } - } - require.Equal(t, expectedOddRows, oddRows, "all odd-bucket rows must be in buffers") -} - -// --- SpillEngine tests --- - -func makeTestKeyExpr() []*plan.Expr { - return []*plan.Expr{{ - Typ: plan.Type{Id: int32(types.T_int32), Width: 32}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - }} -} - -func makeTestEvalKeysFn() func(*batch.Batch) ([]*vector.Vector, error) { - return func(bat *batch.Batch) ([]*vector.Vector, error) { - return bat.Vecs[:1], nil - } -} - -func TestScatterProbeAdmitsExpressionBeforeEvaluation(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - modulo, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - execs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{modulo}) - require.NoError(t, err) - budget, err := process.NewHashBuildBudget(8<<20, 8<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - probeLease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{modulo}, execs, false) - require.NoError(t, err) - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: []*plan.Expr{modulo}, - Budget: generation, - ProbeExpressionLease: probeLease, - }) - engine.InitFromSpilledMap(make([]*os.File, SpillNumBuckets)) - input := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - defer input.Clean(proc.Mp()) - childrenCalls := 0 - fallbackCalled := false - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalls++ - if childrenCalls == 1 { - return input, nil - } - return nil, nil - }, - process.NewAnalyzer(0, false, false, "test"), - func(*batch.Batch) ([]*vector.Vector, error) { - fallbackCalled = true - return nil, errors.New("budgeted probe must evaluate its leased executors") - }, - ) - require.NoError(t, err) - require.Equal(t, 2, childrenCalls) - require.False(t, fallbackCalled) - require.Positive(t, probeLease.Reserved()) - require.Positive(t, generation.Used()) - engine.Cleanup(proc) - require.Positive(t, generation.Used(), "SpillEngine only borrows the probe lease") - for _, exec := range execs { - exec.Free() - } - probeLease.Release() - require.Zero(t, generation.Used()) -} - -func TestScatterProbeExpressionAdmissionRejectsBeforeEval(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - modulo, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - execs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, []*plan.Expr{modulo}) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(execs) - require.True(t, ok) - peak, err := hashbuild.ExpressionVectorPeak(proc, modulo, 4, false) - require.NoError(t, err) - budgetCap := retained + peak - 1 - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - probeLease, err := hashbuild.NewExpressionMemoryLease( - generation, []*plan.Expr{modulo}, execs, false) - require.NoError(t, err) - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: []*plan.Expr{modulo}, - Budget: generation, - ProbeExpressionLease: probeLease, - }) - engine.InitFromSpilledMap(make([]*os.File, SpillNumBuckets)) - input := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - defer input.Clean(proc.Mp()) - childrenCalls := 0 - evalCalled := false - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalls++ - if childrenCalls == 1 { - return input, nil - } - return nil, nil - }, - process.NewAnalyzer(0, false, false, "test"), - func(*batch.Batch) ([]*vector.Vector, error) { - evalCalled = true - return nil, nil - }, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, 1, childrenCalls) - require.False(t, evalCalled) - engine.Cleanup(proc) - for _, exec := range execs { - exec.Free() - } - probeLease.Release() - require.Zero(t, generation.Used()) -} - -func makeInt32Batch(proc *process.Process, vals []int32) *batch.Batch { - bat := batch.NewWithSize(1) - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.SetRowCount(len(vals)) - return bat -} - -func makeInt32PayloadBatch(t *testing.T, proc *process.Process, start, rows, payloadBytes int) *batch.Batch { - t.Helper() - bat := batch.NewWithSize(2) - vals := make([]int32, rows) - for i := range vals { - vals[i] = int32(start + i) - } - bat.Vecs[0] = testutil.MakeInt32Vector(vals, nil, proc.Mp()) - bat.Vecs[1] = vector.NewVec(types.T_varchar.ToType()) - payload := bytes.Repeat([]byte{'x'}, payloadBytes) - for i := 0; i < rows; i++ { - require.NoError(t, vector.AppendBytes(bat.Vecs[1], payload, false, proc.Mp())) - } - bat.SetRowCount(rows) - return bat -} - -func makeDedupKeepLastSpillBatch(proc *process.Process) *batch.Batch { - bat := batch.NewWithSize(3) - bat.Vecs[0] = testutil.MakeInt32Vector( - []int32{1, 1, 2}, nil, proc.Mp()) - bat.Vecs[1] = testutil.MakeInt32Vector( - []int32{10, 20, 30}, nil, proc.Mp()) - bat.Vecs[2] = testutil.MakeInt32Vector( - []int32{100, 0, 0}, []uint64{1, 2}, proc.Mp()) - bat.SetRowCount(3) - return bat -} - -func runtimeStackHasFunctionSuffix(suffix string) bool { - var callers [32]uintptr - n := runtime.Callers(2, callers[:]) - frames := runtime.CallersFrames(callers[:n]) - for { - frame, more := frames.Next() - if strings.HasSuffix(frame.Function, suffix) { - return true - } - if !more { - return false - } - } -} - -func writeBuildFile(proc *process.Process, name string, bat *batch.Batch) *os.File { - return writeBuildRecords(proc, name, bat) -} - -func writeBuildRecords(proc *process.Process, name string, batches ...*batch.Batch) *os.File { - spillfs, _ := proc.GetSpillFileService() - f, _ := spillfs.CreateAndRemoveFile(context.Background(), name) - var buf bytes.Buffer - w := BucketWriter{Name: name, Fd: f} - for _, bat := range batches { - FlushBucketBatch(proc, bat, &w, &buf, nil) - } - return w.HandOffFd() -} - -func makeCorruptBatchFile(t *testing.T) *os.File { - f, err := os.CreateTemp(t.TempDir(), "corrupt-spill") - require.NoError(t, err) - rowCount, batchSize := int64(1), int64(1) - var buf bytes.Buffer - buf.Write(types.EncodeInt64(&rowCount)) - buf.Write(types.EncodeInt64(&batchSize)) - buf.WriteByte(0xff) - _, err = f.Write(buf.Bytes()) - require.NoError(t, err) - _, err = f.Seek(0, io.SeekStart) - require.NoError(t, err) - return f -} - -func TestInitFromSpilledMapMixed(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd1 := writeBuildFile(proc, "test_mixed_1", bat) - fd2 := writeBuildFile(proc, "test_mixed_2", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd1, nil, fd2}) - require.Equal(t, 3, len(engine.buckets)) - require.NotNil(t, engine.buckets[0].BuildFd) - require.Nil(t, engine.buckets[1].BuildFd) - require.NotNil(t, engine.buckets[2].BuildFd) - engine.Cleanup(proc) -} - -func TestRebuildHashmapBasic(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 100) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_rebuild", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.NotNil(t, jm) - require.Equal(t, int64(100), jm.GetRowCount()) - - jm2, res2, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketQueueEmpty, res2) - require.Nil(t, jm2) - - jm.Free() - engine.Cleanup(proc) -} - -func TestRebuildHashmapCancellationKeepsFileOwnedUntilCleanup(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - build := makeInt32Batch(proc, []int32{1, 2, 3, 4}) - var serialized bytes.Buffer - writer := BucketWriter{Name: "rebuild_cancel", Budget: generation} - defer writer.Close() - require.NoError(t, FlushBucketBatch(proc, build, &writer, &serialized, nil)) - build.Clean(proc.Mp()) - file, err := writer.handOffSpillFile() - require.NoError(t, err) - require.Positive(t, generation.SpillDiskUsed()) - require.Positive(t, generation.SpillFDUsed()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - Budget: generation, - }) - engine.InitFromSpilledFiles([]*message.SpillFile{file}) - defer engine.Cleanup(proc) - - proc.Cancel(context.Canceled) - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.ErrorIs(t, err, context.Canceled) - require.Nil(t, jm) - require.Equal(t, BucketSkip, res) - require.True(t, engine.HasMoreBuckets(), "cancellation must leave the queued file with the engine cleanup owner") - require.Positive(t, generation.SpillDiskUsed()) - require.Positive(t, generation.SpillFDUsed()) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestRebuildHashmapRespectsNeedFlags(t *testing.T) { - tests := []struct { - name string - needAllocateSels bool - needBatches bool - }{ - {name: "sels only", needAllocateSels: true}, - {name: "batches only", needBatches: true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 1}) - fd := writeBuildFile(proc, "test_rebuild_flags", bat) - bat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedAllocateSels: tt.needAllocateSels, - NeedBatches: tt.needBatches, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.NotNil(t, jm) - - if tt.needAllocateSels { - require.Equal(t, []int32{0, 1}, jm.GetSels(0)) - } else { - require.Nil(t, jm.GetSels(0)) - } - if tt.needBatches { - require.Len(t, jm.GetBatches(), 1) - require.Equal(t, 2, jm.GetBatches()[0].RowCount()) - } else { - require.Empty(t, jm.GetBatches()) - } - - jm.Free() - engine.Cleanup(proc) - require.Equal(t, int64(0), proc.Mp().CurrNB()) - }) - } -} - -func TestRebuildHashmapWithoutBatchesDropsBatchBudgetBeforeProbe(t *testing.T) { - run := func(t *testing.T, needBatches bool) uint64 { - t.Helper() - const budgetCap = uint64(64 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - values := make([]int32, colexec.DefaultBatchSize/2) - for i := range values { - values[i] = int32(i) - } - buildBat := makeInt32Batch(proc, values) - buildFd := writeBuildFile(proc, "test_rebuild_batch_budget", buildBat) - buildBat.Clean(proc.Mp()) - probeBat := makeInt32Batch(proc, []int32{1}) - probeFd := writeBuildFile(proc, "test_rebuild_batch_budget_probe", probeBat) - probeBat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedBatches: needBatches, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 1, 0, nil) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.NotNil(t, jm) - if needBatches { - require.NotEmpty(t, jm.GetBatches()) - } else { - require.Empty(t, jm.GetBatches()) - } - used := generation.Used() - require.Positive(t, used, "hash map and reader ownership remain live during probe") - - jm.Free() - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - return used - } - - withoutBatches := run(t, false) - withBatches := run(t, true) - require.Greater(t, withBatches, withoutBatches, - "NeedBatches=false must not transfer destroyed batch reservations into the JoinMap") -} - -func TestRebuildHashmapEmptyBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{nil}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketSkip, res) - require.Nil(t, jm) - require.False(t, engine.HasMoreBuckets()) - engine.Cleanup(proc) -} - -func TestRebuildHashmapEmptyBuildOuterJoin(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - var buf bytes.Buffer - spillfs, _ := proc.GetSpillFileService() - f, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_outer_probe") - bat := makeInt32Batch(proc, []int32{1, 2}) - w := BucketWriter{Name: "test_outer_probe", Fd: f} - FlushBucketBatch(proc, bat, &w, &buf, nil) - probeFd := w.HandOffFd() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - }) - engine.InitFromSpilledMap([]*os.File{nil}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketEmptyBuild, res) - require.Nil(t, jm) - require.True(t, engine.IsProbing()) - require.False(t, engine.HasMoreBuckets()) - - got, err := engine.NextProbeBatch(proc) - require.NoError(t, err) - require.NotNil(t, got) - require.Equal(t, 2, got.RowCount()) - - engine.Cleanup(proc) -} - -func TestRebuildHashmapEmptyFile(t *testing.T) { - for _, keepProbe := range []bool{false, true} { - t.Run(fmt.Sprintf("keep_probe_%t", keepProbe), func(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - baseline := proc.Mp().CurrNB() - budget, err := process.NewHashBuildBudget(16<<20, 16<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - spillfs, err := proc.GetSpillFileService() - require.NoError(t, err) - buildFd, err := spillfs.CreateAndRemoveFile(proc.Ctx, "test_empty_build_file") - require.NoError(t, err) - probeBat := makeInt32Batch(proc, []int32{1}) - probeFd := writeBuildFile(proc, "test_empty_build_probe", probeBat) - probeBat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - // A literal executor owns an mpool vector as soon as Prepare - // succeeds. This makes the empty-build branch's builder.Free - // observable instead of relying on a zero-allocation column - // executor. - BuildKeyExprs: []*plan.Expr{ - plan2.MakePlan2Int32ConstExprWithType(1), - }, - NeedsProbeForEmptyBuild: keepProbe, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Nil(t, jm) - if keepProbe { - require.Equal(t, BucketEmptyBuild, res) - require.True(t, engine.IsProbing()) - } else { - require.Equal(t, BucketSkip, res) - } - engine.Cleanup(proc) - require.Equal(t, baseline, proc.Mp().CurrNB()) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - }) - } -} - -func TestScatterProbeTable(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{10, 20, 30}) - fd := writeBuildFile(proc, "test_sp_build", bat) - bat.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - // Use many rows to ensure distribution across buckets. - vals := make([]int32, 1000) - for i := range vals { - vals[i] = int32(i) - } - batches := []*batch.Batch{makeInt32Batch(proc, vals)} - defer batches[0].Clean(proc.Mp()) - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil - } - b := batches[idx] - idx++ - return b, nil - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - require.NotNil(t, engine.probeKeyEval) - require.Len(t, engine.buckets, 1) - require.NotNil(t, engine.buckets[0].ProbeFd) - require.Equal(t, int64(len(vals)), engine.buckets[0].ProbeRows, - "probe partitioning must conserve every row at the build payload's fanout") - - engine.Cleanup(proc) -} - -func TestScatterProbeTableRejectsInvalidBuildFanoutBeforeInput(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - for _, bucketCount := range []int{0, 3, SpillNumBuckets + 1} { - t.Run(fmt.Sprintf("buckets_%d", bucketCount), func(t *testing.T) { - engine := NewSpillEngine(SpillEngineConfig{}) - engine.InitFromSpilledMap(make([]*os.File, bucketCount)) - inputCalled := false - err := engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - inputCalled = true - return nil, nil - }, - process.NewAnalyzer(0, false, false, "test"), - makeTestEvalKeysFn(), - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetInvalid) - require.False(t, inputCalled, "invalid fanout must fail before consuming probe input") - engine.Cleanup(proc) - }) - } -} - -func TestScatterProbeCancellationReleasesPhysicalAndMemoryBudget(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - budget, err := process.NewHashBuildBudget(64<<20, 64<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - // Identical keys force one large selected payload through a real spill - // writer before the second upstream call cancels the pipeline. - const rows = 8192 - keys := make([]int32, rows) - payload := make([]string, rows) - for i := range keys { - keys[i] = 7 - payload[i] = strings.Repeat("x", 128) - } - input := batch.NewWithSize(2) - input.Vecs[0] = testutil.MakeInt32Vector(keys, nil, proc.Mp()) - input.Vecs[1] = testutil.MakeVarcharVector(payload, nil, proc.Mp()) - input.SetRowCount(rows) - defer func() { - if input != nil { - input.Clean(proc.Mp()) - } - }() - - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - Budget: generation, - }) - engine.InitFromSpilledMap(make([]*os.File, SpillNumBuckets)) - analyzer := process.NewAnalyzer(0, false, false, "test") - childrenCalls := 0 - var peakDisk, peakFD uint64 - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalls++ - if childrenCalls == 1 { - return input, nil - } - peakDisk = generation.SpillDiskUsed() - peakFD = generation.SpillFDUsed() - proc.Cancel(context.Canceled) - return input, nil - }, - analyzer, - makeTestEvalKeysFn(), - ) - require.ErrorIs(t, err, context.Canceled) - require.Equal(t, 2, childrenCalls) - require.Positive(t, peakDisk, "first batch must reach a physical spill file") - require.Positive(t, peakFD, "first batch must own an admitted spill descriptor") - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - input.Clean(proc.Mp()) - input = nil - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestSpillEntryPointsRejectPreCanceledProcessWithoutOwnershipTransfer(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - t.Cleanup(proc.Free) - ctx, cancel := context.WithCancelCause(proc.Ctx) - process.ReplacePipelineCtx(proc, ctx, cancel) - proc.Cancel(context.Canceled) - - reuseBat := batch.NewOffHeapWithSize(0) - reuseCleaned := false - cleanReuse := func() { - if !reuseCleaned { - reuseBat.Clean(proc.Mp()) - reuseCleaned = true - } - } - t.Cleanup(cleanReuse) - reader := BucketReader{} - got, err := reader.ReadBatch(proc, reuseBat) - require.Nil(t, got) - require.ErrorIs(t, err, context.Canceled) - require.Zero(t, reuseBat.RowCount()) - - writer := BucketWriter{Name: "must_not_be_created"} - t.Cleanup(writer.Close) - err = writeBucketPayload( - proc, - []byte{1}, - 1, - &writer, - process.NewAnalyzer(0, false, false, "test"), - ) - require.ErrorIs(t, err, context.Canceled) - require.False(t, writer.Created()) - require.Nil(t, writer.diskReservation) - require.Nil(t, writer.fdReservation) - - input := makeInt32Batch(proc, []int32{1}) - inputCleaned := false - cleanInput := func() { - if !inputCleaned { - input.Clean(proc.Mp()) - inputCleaned = true - } - } - t.Cleanup(cleanInput) - engine := NewSpillEngine(SpillEngineConfig{}) - engineCleaned := false - cleanEngine := func() { - if !engineCleaned { - engine.Cleanup(proc) - engineCleaned = true - } - } - t.Cleanup(cleanEngine) - scatterWriters := []BucketWriter{{Name: "must_not_be_created"}} - t.Cleanup(scatterWriters[0].Close) - err = engine.scatterBatch( - proc, - input, - []*vector.Vector{input.Vecs[0]}, - scatterWriters, - nil, - 0, - false, - process.NewAnalyzer(0, false, false, "test"), - ) - require.ErrorIs(t, err, context.Canceled) - require.False(t, scatterWriters[0].Created()) - require.Nil(t, scatterWriters[0].diskReservation) - require.Nil(t, scatterWriters[0].fdReservation) - - engine.InitFromSpilledMap([]*os.File{nil}) - childrenCalled := false - err = engine.ScatterProbeTable( - proc, - func() (*batch.Batch, error) { - childrenCalled = true - return input, nil - }, - process.NewAnalyzer(0, false, false, "test"), - makeTestEvalKeysFn(), - ) - require.ErrorIs(t, err, context.Canceled) - require.False(t, childrenCalled) - - probeFd, err := os.CreateTemp(t.TempDir(), "pre-canceled-probe") - require.NoError(t, err) - t.Cleanup(func() { _ = probeFd.Close() }) - probeReleases := 0 - probeFile := message.NewSpillFile(probeFd, 1, 1, func() { probeReleases++ }) - engine.probeReader.ResetForSpillFile(probeFile) - got, err = engine.NextProbeBatch(proc) - require.Nil(t, got) - require.ErrorIs(t, err, context.Canceled) - require.Zero(t, probeReleases) - _, err = probeFd.Stat() - require.NoError(t, err) - - reSpillFd, err := os.CreateTemp(t.TempDir(), "pre-canceled-respill") - require.NoError(t, err) - t.Cleanup(func() { _ = reSpillFd.Close() }) - reSpillReleases := 0 - reSpillFile := message.NewSpillFile(reSpillFd, 1, 1, func() { reSpillReleases++ }) - t.Cleanup(func() { _ = reSpillFile.Close() }) - subBuckets, err := engine.reSpillBucket( - proc, - process.NewAnalyzer(0, false, false, "test"), - SpillBucket{BuildFd: reSpillFile}, - nil, - nil, - nil, - ) - require.Nil(t, subBuckets) - require.ErrorIs(t, err, context.Canceled) - require.Zero(t, reSpillReleases) - _, err = reSpillFd.Stat() - require.NoError(t, err) - require.NoError(t, reSpillFile.Close()) - require.Equal(t, 1, reSpillReleases) - _, err = reSpillFd.Stat() - require.Error(t, err) - - cleanEngine() - require.Equal(t, 1, probeReleases) - _, err = probeFd.Stat() - require.Error(t, err) - cleanInput() - cleanReuse() - require.Zero(t, proc.Mp().CurrNB()) -} - -func TestScatterProbeTableSkipEmptyBuild(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd1 := writeBuildFile(proc, "test_skip_build_1", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd1, nil}) - - // Use many rows to ensure distribution. - vals := make([]int32, 1000) - for i := range vals { - vals[i] = int32(i) - } - batches := []*batch.Batch{makeInt32Batch(proc, vals)} - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil - } - b := batches[idx] - idx++ - return b, nil - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - - // Bucket 1 (nil build, not outer join) should have no probe data. - require.Nil(t, engine.buckets[1].ProbeFd) - - engine.Cleanup(proc) -} - -func TestScatterProbeTableWithEmptyBatches(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2}) - fd := writeBuildFile(proc, "test_empty_bat", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - batches := []*batch.Batch{ - batch.NewWithSize(0), - makeInt32Batch(proc, []int32{5, 6}), - } - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil - } - b := batches[idx] - idx++ - return b, nil - } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - - engine.Cleanup(proc) -} - -func TestNextProbeBatch(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - // Write build file. - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd := writeBuildFile(proc, "test_npb_build", bat) - - // Write probe file manually (ensures bucket 0 has probe data). - probeFd := writeBuildFile(proc, "test_npb_probe", makeInt32Batch(proc, []int32{5, 6, 7})) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReady, res) - require.True(t, engine.IsProbing()) - - got, err := engine.NextProbeBatch(proc) - require.NoError(t, err) - require.NotNil(t, got) - require.Equal(t, 3, got.RowCount()) - - got2, err2 := engine.NextProbeBatch(proc) - require.NoError(t, err2) - require.Nil(t, got2) - - engine.FinishBucket() - require.False(t, engine.IsProbing()) - got3, err3 := engine.NextProbeBatch(proc) - require.NoError(t, err3) - require.Nil(t, got3) - - jm.Free() - engine.Cleanup(proc) -} - -func TestCorruptSpillBatchErrors(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{}) - engine.probeReader.ResetForFd(makeCorruptBatchFile(t)) - _, err := engine.NextProbeBatch(proc) - require.Error(t, err) - engine.Cleanup(proc) - - engine = NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - }) - engine.InitFromSpilledMap([]*os.File{makeCorruptBatchFile(t)}) - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.Error(t, err) - require.Nil(t, jm) - require.Equal(t, BucketSkip, res) - engine.Cleanup(proc) -} - -func TestRebuildSkipsBuildOnlyBucketWhenJoinCannotUseIt(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - fd := makeCorruptBatchFile(t) - releases := 0 - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledFiles([]*message.SpillFile{ - message.NewSpillFile(fd, 1, 17, func() { releases++ }), - }) - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err, "an irrelevant build-only file must not be read") - require.Nil(t, jm) - require.Equal(t, BucketSkip, res) - require.False(t, engine.HasMoreBuckets()) - require.Equal(t, 1, releases) - _, err = fd.Stat() - require.Error(t, err) - - engine.Cleanup(proc) - require.Equal(t, 1, releases) -} - -func TestAdvanceToNextBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 50) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_advance_build", bat) - probeFd := writeBuildFile(proc, "test_advance_probe", makeInt32Batch(proc, []int32{5, 6})) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - var capturedJM *message.JoinMap - var capturedRes BucketResult - analyzer := process.NewAnalyzer(0, false, false, "test") - - ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, res BucketResult) { - capturedJM = jm - capturedRes = res - }) - require.NoError(t, err) - require.True(t, ok) - require.NotNil(t, capturedJM) - require.Equal(t, BucketReady, capturedRes) - require.True(t, engine.IsProbing()) - require.False(t, engine.HasMoreBuckets()) - - capturedJM.Free() - engine.FinishBucket() - ok, err = engine.AdvanceToNextBucket(proc, analyzer, func(*message.JoinMap, BucketResult) { - t.Fatal("queue exhaustion must not invoke callback") - }) - require.NoError(t, err) - require.False(t, ok) - engine.Cleanup(proc) -} - -func TestReSpillBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_respill_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 100, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReSpilled, res) - require.Nil(t, jm) - - // Drain all remaining buckets. - for engine.HasMoreBuckets() { - jm2, _, err2 := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err2) - if jm2 != nil { - jm2.Free() - } - } - - engine.Cleanup(proc) -} - -func TestRebuildHashmapKeepsScratchHeadroomForCopyAdmissionReSpill(t *testing.T) { - const ( - budgetCap = uint64(12 << 20) - recordRows = colexec.DefaultBatchSize - recordCount = 3 - payloadSize = 256 - ) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - records := make([]*batch.Batch, recordCount) - for i := range records { - records[i] = makeInt32PayloadBatch(t, proc, i*recordRows, recordRows, payloadSize) - } - buildFd := writeBuildRecords(proc, "rebuild_scratch_headroom", records...) - for i := range records { - records[i].Clean(proc.Mp()) - } - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - analyzer := process.NewAnalyzer(0, false, false, "test") - - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Nil(t, jm) - require.Equal(t, BucketReSpilled, result) - var childRows int64 - for _, child := range engine.buckets { - childRows += child.BuildRows - } - require.Equal(t, int64(recordRows*recordCount), childRows) - extra := analyzer.GetOpStats().ExtraStats - require.Positive(t, extra["JoinSpillRebuildScratchFloorBytes"]) - require.Equal(t, int64(1), extra["JoinSpillRebuildCopyAdmissionReSpillAttempts"]) - require.Zero(t, extra["JoinSpillRebuildPreCopyReSpillAttempts"]) - - // The rejected record is reader-owned pending state: retained, pending, and - // unread rows must form an exact partition of the original stream. - seen := make([]uint8, recordRows*recordCount) - reuse := batch.NewOffHeapWithSize(0) - reader := BucketReader{} - for i := range engine.buckets { - file := engine.buckets[i].BuildFd - engine.buckets[i].BuildFd = nil - reader.ResetForSpillFile(file) - for { - bat, readErr := reader.ReadBatch(proc, reuse) - if readErr == io.EOF { - break - } - require.NoError(t, readErr) - for _, key := range vector.MustFixedColNoTypeCheck[int32](bat.Vecs[0]) { - require.GreaterOrEqual(t, key, int32(0)) - require.Less(t, key, int32(len(seen))) - seen[key]++ - } - } - reader.closeCurrentFile() - } - reader.Close() - reuse.Clean(proc.Mp()) - for key, count := range seen { - require.Equalf(t, uint8(1), count, "key %d must be emitted exactly once", key) - } - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildScratchAdmissionIsBestEffortForResidentBucket(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeInt32PayloadBatch(t, proc, 0, 1024, 1024) - buildFd := writeBuildFile(proc, "rebuild_scratch_best_effort", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedBatches: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - floorRejected := false - budget.SetAggregateCapProvider(func() (uint64, error) { - if !floorRejected && runtimeStackHasFunctionSuffix( - "spillutil.(*SpillEngine).reserveRebuildScatterScratch", - ) { - floorRejected = true - return generation.Used(), nil - } - return budgetCap, nil - }) - analyzer := process.NewAnalyzer(0, false, false, "test") - - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.True(t, floorRejected) - require.Equal(t, BucketReady, result) - require.NotNil(t, jm) - require.Equal(t, int64(1024), jm.GetRowCount()) - require.Equal(t, int64(1), - analyzer.GetOpStats().ExtraStats["JoinSpillRebuildScratchReserveRejects"]) - jm.Free() - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildScratchLifecycleFailureIsNotRecoveredAsAdmission(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeInt32Batch(proc, []int32{1, 2, 3}) - buildFd := writeBuildFile(proc, "rebuild_scratch_closed", build) - build.Clean(proc.Mp()) - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - closedErr := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - Message: "forced closed rebuild scratch budget", - } - budget.SetAggregateCapProvider(func() (uint64, error) { - if runtimeStackHasFunctionSuffix( - "spillutil.(*SpillEngine).reserveRebuildScatterScratch", - ) { - return 0, closedErr - } - return budgetCap, nil - }) - analyzer := process.NewAnalyzer(0, false, false, "test") - - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.Same(t, closedErr, err) - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, jm) - require.Equal(t, BucketSkip, result) - require.Zero(t, analyzer.GetOpStats().ExtraStats["JoinSpillRebuildCopyAdmissionReSpillAttempts"]) - require.Zero(t, analyzer.GetOpStats().ExtraStats["JoinSpillRebuildPreCopyReSpillAttempts"]) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildHashmapRejectsReSpillAfterDedupRewrite(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeDedupKeepLastSpillBatch(proc) - buildFd := writeBuildFile(proc, "dedup_unsafe_respill", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedBatches: true, - IsDedup: true, - OnDuplicateAction: plan.Node_FAIL, - DedupBuildKeepLast: true, - DedupColName: "id", - DedupColTypes: []plan.Type{{Id: int32(types.T_int32)}}, - DelColIdx: -1, - DedupDeleteMarkerColIdx: 2, - DedupDeleteKeepColIdxList: []int32{2}, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - - forcedUnsafeReject := false - budget.SetAggregateCapProvider(func() (uint64, error) { - // keepDiscardedRowsForDelete has already compacted the retained input - // when it asks copyBuildBatch to admit the delete-only rows. Reject that - // exact transition without depending on a fragile global call ordinal. - if runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).keepDiscardedRowsForDelete", - ) { - forcedUnsafeReject = true - return max(uint64(1), generation.Used()), nil - } - return budgetCap, nil - }) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.True(t, forcedUnsafeReject) - require.Nil(t, jm) - require.Equal(t, BucketSkip, result) - require.Len(t, engine.buckets, 1, - "unsafe recovery must not replace the parent with child buckets") - require.Equal(t, 1, engine.buckets[0].Depth) - require.Nil(t, engine.buckets[0].BuildFd, - "the consumed parent file must not be republished as a JoinMap or child") - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildHashmapReSpillsAdmissionBeforeDedupRewrite(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeDedupKeepLastSpillBatch(proc) - buildFd := writeBuildFile(proc, "dedup_safe_respill", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - NeedBatches: true, - IsDedup: true, - OnDuplicateAction: plan.Node_FAIL, - DedupBuildKeepLast: true, - DedupColName: "id", - DedupColTypes: []plan.Type{{Id: int32(types.T_int32)}}, - DelColIdx: -1, - DedupDeleteMarkerColIdx: 2, - DedupDeleteKeepColIdxList: []int32{2}, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - - forcedSafeReject := false - budget.SetAggregateCapProvider(func() (uint64, error) { - // The first budget request made from buildHashmap is reserveBuildAux, - // before any Dedup batch rewrite. Reject once; re-spill itself does not - // call buildHashmap and therefore retains the normal cap. - if !forcedSafeReject && - runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).buildHashmap", - ) && - !runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).keepDiscardedRowsForDelete", - ) { - forcedSafeReject = true - return max(uint64(1), generation.Used()), nil - } - return budgetCap, nil - }) - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, result, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.True(t, forcedSafeReject) - require.Nil(t, jm) - require.Equal(t, BucketReSpilled, result) - require.NotEmpty(t, engine.buckets) - var childRows int64 - for _, child := range engine.buckets { - require.Equal(t, 2, child.Depth) - childRows += child.BuildRows - } - require.Equal(t, int64(3), childRows, - "safe recovery must conserve the original retained rows") - require.Equal(t, int64(1), - analyzer.GetOpStats().ExtraStats["JoinSpillRebuildMapAdmissionReSpillAttempts"]) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestRebuildHashmapClosedBudgetDoesNotReSpill(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - build := makeInt32Batch(proc, []int32{1, 2, 3}) - buildFd := writeBuildFile(proc, "closed_budget_no_respill", build) - build.Clean(proc.Mp()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1 << 30, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - - closedErr := &process.HashBuildBudgetError{ - Kind: process.HashBuildBudgetErrorClosed, - Message: "forced closed hash-build budget", - } - forcedClosed := false - budget.SetAggregateCapProvider(func() (uint64, error) { - if !forcedClosed && runtimeStackHasFunctionSuffix( - "hashbuild.(*HashmapBuilder).buildHashmap", - ) { - forcedClosed = true - return 0, closedErr - } - return budgetCap, nil - }) - - jm, result, err := engine.RebuildHashmap( - proc, process.NewAnalyzer(0, false, false, "test")) - require.True(t, forcedClosed) - require.Same(t, closedErr, err, - "a lifecycle failure must be returned unchanged") - require.ErrorIs(t, err, process.ErrHashBuildBudgetClosed) - require.NotErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Nil(t, jm) - require.Equal(t, BucketSkip, result) - require.Len(t, engine.buckets, 1, - "a lifecycle failure must not replace the parent with child buckets") - require.Equal(t, 1, engine.buckets[0].Depth) - require.Nil(t, engine.buckets[0].BuildFd) - - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) - require.Zero(t, proc.Mp().CurrNB()) - generation.Close() - proc.Free() -} - -func TestReSpillReleasesBuilderExecutorsBeforeReplacementAdmission(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - col := &plan.Expr{ - Typ: plan.Type{Id: int32(types.T_int32)}, - Expr: &plan.Expr_Col{Col: &plan.ColRef{ColPos: 0}}, - } - modulo, err := plan2.BindFuncExprImplByPlanExpr( - proc.Ctx, - "%", - []*plan.Expr{col, plan2.MakePlan2Int32ConstExprWithType(2)}, - ) - require.NoError(t, err) - exprs := []*plan.Expr{modulo} - - probeExecs, err := colexec.NewExpressionExecutorsFromPlanExpressions(proc, exprs) - require.NoError(t, err) - retained, ok := colexec.ExpressionExecutorsRetainedBytes(probeExecs) - require.True(t, ok) - require.Positive(t, retained) - for _, executor := range probeExecs { - executor.Free() - } - - // The cap intentionally fits exactly one executor set. reSpillBucket must - // release the failed builder's equivalent set before constructing its own. - budget := process.MustNewHashBuildBudget(retained, retained) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - builder := &hashbuild.HashmapBuilder{} - builder.SetBudget(generation) - require.NoError(t, builder.Prepare(exprs, -1, -1, nil, proc)) - require.Equal(t, retained, generation.Used()) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: exprs, - Budget: generation, - }) - subBuckets, err := engine.reSpillBucket( - proc, - process.NewAnalyzer(0, false, false, "test"), - SpillBucket{}, - builder, - &BucketReader{}, - nil, - ) - require.NoError(t, err) - require.Empty(t, subBuckets) - require.Equal(t, retained, generation.Used()) - require.NotNil(t, engine.buildExprLease) - - builder.Free(proc) - engine.Cleanup(proc) - require.Zero(t, generation.Used()) -} - -func TestReSpillBucketReleasesDrainedBatchBudget(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget, err := process.NewHashBuildBudget(budgetCap, budgetCap) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - builder := &hashbuild.HashmapBuilder{} - builder.SetBudget(generation) - require.NoError(t, builder.Prepare(makeTestKeyExpr(), -1, -1, nil, proc)) - defer builder.Free(proc) - - values := make([]int32, colexec.DefaultBatchSize/2) - for i := range values { - values[i] = int32(i) - } - input := makeInt32Batch(proc, values) - require.NoError(t, builder.CopyBuildBatch(input, proc)) - builder.InputBatchRowCount = input.RowCount() - input.Clean(proc.Mp()) - batchCharge := generation.Used() - require.Positive(t, batchCharge) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - Budget: generation, - }) - subBuckets, err := engine.reSpillBucket( - proc, - process.NewAnalyzer(0, false, false, "test"), - SpillBucket{Depth: 1, BuildRows: int64(len(values))}, - builder, - &BucketReader{}, - nil, - ) - require.NoError(t, err) - for i := range subBuckets { - if subBuckets[i].BuildFd != nil { - require.NoError(t, subBuckets[i].BuildFd.Close()) - } - if subBuckets[i].ProbeFd != nil { - require.NoError(t, subBuckets[i].ProbeFd.Close()) - } - } - engine.Cleanup(proc) - - require.Empty(t, builder.Batches.Buf) - require.Less(t, generation.Used(), batchCharge, - "re-spill must not retain the destroyed build-batch reservation") - require.Zero(t, generation.Used()) - - // Model the scratch/read admission that follows the drain. The full cap is - // available only when re-spill released the stale batch ownership itself, - // instead of relying on its caller to free the builder later. - next, err := generation.Reserve(budgetCap) - require.NoError(t, err) - require.True(t, next.Release()) -} - -func TestReSpillConservesBuildAndProbeRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) - } - build := makeInt32Batch(proc, vals) - probe := makeInt32Batch(proc, vals) - buildFd := writeBuildFile(proc, "test_conserve_build", build) - probeFd := writeBuildFile(proc, "test_conserve_probe", probe) - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - SpillThreshold: 500, - NeedsBuildForEmptyProbe: true, - NeedsProbeForEmptyBuild: true, - }) - engine.InitFromSpilledMap([]*os.File{buildFd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, int64(len(vals)), 0, nil) - engine.probeKeyEval = makeTestEvalKeysFn() - - jm, res, err := engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - require.Equal(t, BucketReSpilled, res) - require.Nil(t, jm) - var buildRows, probeRows, largest int64 - for _, child := range engine.buckets { - buildRows += child.BuildRows - probeRows += child.ProbeRows - if child.BuildRows > largest { - largest = child.BuildRows - } - } - require.Equal(t, int64(len(vals)), buildRows) - require.Equal(t, int64(len(vals)), probeRows) - require.Less(t, largest, int64(len(vals))) - - for engine.HasMoreBuckets() { - jm, _, err = engine.RebuildHashmap(proc, process.NewAnalyzer(0, false, false, "test")) - require.NoError(t, err) - if jm != nil { - jm.Free() - } - } - engine.Cleanup(proc) -} - -func TestReSpillDepthLimit(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(64<<10, 64<<10) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - bat := makeInt32Batch(proc, []int32{1, 2, 3, 4, 5}) - fd := writeBuildFile(proc, "test_depth_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsBuildForEmptyProbe: true, - SpillThreshold: 1, - Budget: generation, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].Depth = SpillMaxPass - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.Error(t, err, "depth limit must not force an over-budget hashmap build") - require.Equal(t, BucketSkip, res) - require.Nil(t, jm) - engine.Cleanup(proc) -} - -func TestReSpillWithProbe(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) - } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_rsp_probe_build", bat) - - // Manually write probe file. - probeFd := writeBuildFile(proc, "test_rsp_probe", makeInt32Batch(proc, []int32{100, 200, 300})) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - SpillThreshold: 100, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - engine.buckets[0].ProbeFd = message.NewSpillFile(probeFd, 0, 0, nil) - - // Set probeKeyEval so scatterProbe works during re-spill. - engine.probeKeyEval = makeTestEvalKeysFn() - - analyzer := process.NewAnalyzer(0, false, false, "test") - jm, res, err := engine.RebuildHashmap(proc, analyzer) - require.NoError(t, err) - require.Equal(t, BucketReSpilled, res) - require.Nil(t, jm) - - for engine.HasMoreBuckets() { - jm2, _, err := engine.RebuildHashmap(proc, analyzer) - if jm2 != nil { - jm2.Free() - } - if err != nil { - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - break - } - } - - engine.Cleanup(proc) -} - -func TestAdvanceToNextBucketReSpilled(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - vals := make([]int32, 5000) - for i := range vals { - vals[i] = int32(i) + values := make([]int32, 100) + for i := range values { + values[i] = int32(i) } - bat := makeInt32Batch(proc, vals) - fd := writeBuildFile(proc, "test_adv_re_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ + build := makeInt32Batch(proc, values) + probe := makeInt32Batch(proc, []int32{1}) + defer build.Clean(proc.Mp()) + defer probe.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ BuildKeyExprs: makeTestKeyExpr(), + SpillThreshold: 1, NeedsBuildForEmptyProbe: true, - SpillThreshold: 100, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - analyzer := process.NewAnalyzer(0, false, false, "test") - - callbackCalled := false - ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, _ BucketResult) { - callbackCalled = true - if jm != nil { - jm.Free() - } - }) - require.NoError(t, err) - require.True(t, ok) - require.False(t, callbackCalled, "re-spill is consumed before the callback") - - for engine.HasMoreBuckets() { - ok, err := engine.AdvanceToNextBucket(proc, analyzer, func(jm *message.JoinMap, _ BucketResult) { - if jm != nil { - jm.Free() - } - }) - require.NoError(t, err) - if ok { - engine.FinishBucket() - } - } - // Test passes if the loop terminates without errors. - engine.Cleanup(proc) -} - -func TestBuilderMemSize(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - builder := &hashbuild.HashmapBuilder{} - err := builder.Prepare(makeTestKeyExpr(), -1, -1, nil, proc) - require.NoError(t, err) - - sz := builderMemSize(builder) - require.Equal(t, int64(0), sz) - - bat := makeInt32Batch(proc, []int32{1, 2, 3, 4, 5}) - err = builder.Batches.CopyIntoBatches(bat, proc) - require.NoError(t, err) - builder.InputBatchRowCount += bat.RowCount() - - sz2 := builderMemSize(builder) - require.Greater(t, sz2, int64(0), "size should grow after adding batches") - - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) -} - -func TestBuilderMemSizeIncludesCompletedBatchesAndPartialTail(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - builder := &hashbuild.HashmapBuilder{} - fullValues := make([]int32, colexec.DefaultBatchSize) - full := makeInt32Batch(proc, fullValues) - partial := makeInt32Batch(proc, []int32{1, 2, 3}) - require.NoError(t, builder.Batches.CopyIntoBatches(full, proc)) - require.NoError(t, builder.Batches.CopyIntoBatches(partial, proc)) - require.Len(t, builder.Batches.Buf, 2) - require.Equal(t, colexec.DefaultBatchSize, builder.Batches.Buf[0].RowCount()) - require.Equal(t, 3, builder.Batches.Buf[1].RowCount()) - - want := builder.Batches.MemSize + int64(builder.Batches.Buf[1].Size()) - require.Equal(t, want, builderMemSize(builder)) - - full.Clean(proc.Mp()) - partial.Clean(proc.Mp()) - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) -} - -func TestShouldReSpillBeforeRetainUsesPredictedBytesAndRows(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - builder := &hashbuild.HashmapBuilder{} - retained := makeInt32Batch(proc, []int32{1, 2, 3}) - next := makeInt32Batch(proc, []int32{4, 5}) - require.NoError(t, builder.Batches.CopyIntoBatches(retained, proc)) - builder.InputBatchRowCount = retained.RowCount() - builder.Batches.MemSize = 200_000 - - predictedBytes := builderMemSize(builder) + int64(next.Size()) - require.False(t, shouldReSpillBeforeRetain(builder, next, predictedBytes)) - require.True(t, shouldReSpillBeforeRetain(builder, next, predictedBytes-1)) - require.False(t, shouldReSpillBeforeRetain(builder, next, 6)) - require.True(t, shouldReSpillBeforeRetain(builder, next, 5)) - require.False(t, shouldReSpillBeforeRetain(builder, next, 0)) - - retained.Clean(proc.Mp()) - next.Clean(proc.Mp()) - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) -} - -func TestRebuildScratchFloorBoundsCoalescedPhysicalBatch(t *testing.T) { - const budgetCap = uint64(64 << 20) - budget := process.MustNewHashBuildBudget(budgetCap, budgetCap) - generation, err := budget.OpenGeneration(1) - require.NoError(t, err) - defer generation.Close() - - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - builder := &hashbuild.HashmapBuilder{} - builder.SetBudget(generation) - require.NoError(t, builder.Prepare(makeTestKeyExpr(), -1, -1, nil, proc)) - firstRows := make([]int32, colexec.DefaultBatchSize/2) - first := makeInt32Batch(proc, firstRows) - second := makeInt32Batch(proc, firstRows) - require.NoError(t, builder.CopyBuildBatch(first, proc)) - builder.InputBatchRowCount = first.RowCount() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - Budget: generation, - }) - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.reserveRebuildScatterScratch(builder, second, analyzer)) - require.Positive(t, engine.scatterScratchFloor) - require.Equal(t, engine.scatterScratchFloor, engine.scatterScratchReservation.Size()) - require.NoError(t, builder.CopyBuildBatch(second, proc)) - builder.InputBatchRowCount += second.RowCount() - require.Len(t, builder.Batches.Buf, 1, "two half records must coalesce") - physical := builder.Batches.Buf[0] - transient, err := scatterTransientBudgetBytes(physical, true) - require.NoError(t, err) - growth, ok := engine.scatterCapacityGrowthBytes(physical.RowCount(), 1) - require.True(t, ok) - require.GreaterOrEqual(t, engine.scatterScratchFloor, transient+growth) - - require.NoError(t, engine.reconcileScatterScratch()) - require.Equal(t, engine.scatterScratchFloor, engine.scatterScratchReservation.Size(), - "repartition headroom must survive per-batch reconciliation") - extra := analyzer.GetOpStats().ExtraStats - require.Equal(t, int64(1), extra["JoinSpillRebuildScratchReserveCount"]) - require.Equal(t, spillStatInt64(engine.scatterScratchFloor), extra["JoinSpillRebuildScratchFloorBytes"]) - - first.Clean(proc.Mp()) - second.Clean(proc.Mp()) - builder.FreeHashMapAndBatches(proc) - builder.Free(proc) - engine.Cleanup(proc) - require.Zero(t, generation.Used()) -} - -func TestFinishBucket(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), + NeedsProbeForEmptyBuild: true, }) - engine.InitFromSpilledMap([]*os.File{nil}) - - fd, err := os.CreateTemp(t.TempDir(), "probe") - require.NoError(t, err) - engine.probeReader.ResetForFd(fd) - require.True(t, engine.IsProbing()) - - engine.FinishBucket() - require.False(t, engine.IsProbing()) - - engine.FinishBucket() - require.False(t, engine.IsProbing()) - - engine.Cleanup(proc) -} - -func TestCleanupSpillEngine(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd1 := writeBuildFile(proc, "test_cl_build", bat) - - var buf bytes.Buffer - spillfs, _ := proc.GetSpillFileService() - probeFile, _ := spillfs.CreateAndRemoveFile(context.Background(), "test_cl_probe") - pw := BucketWriter{Name: "test_cl_probe", Fd: probeFile} - FlushBucketBatch(proc, bat, &pw, &buf, nil) - fd2 := pw.HandOffFd() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile(writeBuildFile(proc, t.Name()+"-build", build), int64(len(values))), }) - engine.InitFromSpilledMap([]*os.File{fd1}) - engine.probeReader.ResetForFd(fd2) - - engine.buildReadBatch = batch.NewOffHeapWithSize(0) - engine.probeReadBatch = batch.NewOffHeapWithSize(0) - - engine.keyExecs = make([]colexec.ExpressionExecutor, 1) - exec, _ := colexec.NewExpressionExecutor(proc, makeTestKeyExpr()[0]) - engine.keyExecs[0] = exec - + engine.buckets[0].ProbeFd = newTestSpillFile( + writeBuildFile(proc, t.Name()+"-probe", probe), + 0, + ) + engine.probeKeyEval = makeTestEvalKeysFn() + _, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorContains(t, err, "row count") engine.Cleanup(proc) - - require.False(t, engine.IsProbing()) - require.Nil(t, engine.buckets) - require.Nil(t, engine.buildReadBatch) - require.Nil(t, engine.probeReadBatch) - require.Nil(t, engine.keyExecs) - - b := make([]byte, 1) - _, err := fd1.Read(b) - require.Error(t, err, "fd1 should be closed") - _, err = fd2.Read(b) - require.Error(t, err, "fd2 should be closed") -} - -func TestScatterProbeFunctionUsesStoredEval(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - }) - - evalCalled := false - engine.probeKeyEval = func(bat *batch.Batch) ([]*vector.Vector, error) { - evalCalled = true - return []*vector.Vector{bat.Vecs[0]}, nil - } - - writers := MakeBucketWriters("test_scatter_func") - buffers := make([]*batch.Batch, len(writers)) - bat := makeInt32Batch(proc, []int32{5, 15, 25}) - - err := scatterProbe(proc, engine, bat, writers, buffers, 1, nil) - require.NoError(t, err) - require.True(t, evalCalled, "probeKeyEval must be used for scatterProbe") - - wantErr := errors.New("probe key evaluation failed") - engine.probeKeyEval = func(*batch.Batch) ([]*vector.Vector, error) { return nil, wantErr } - require.ErrorIs(t, scatterProbe(proc, engine, bat, writers, buffers, 1, nil), wantErr) - - for i := range writers { - writers[i].Close() - } } -func TestScatterPeakDoesNotDoubleChargeReservedSource(t *testing.T) { +func TestRebuildRejectsRowsWithoutFile(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - values := make([]int32, 8192) - for i := range values { - values[i] = int32(i) - } - bat := makeInt32Batch(proc, values) - defer bat.Clean(proc.Mp()) - - charged, err := scatterTransientBudgetBytes(bat, true) - require.NoError(t, err) - uncharged, err := scatterTransientBudgetBytes(bat, false) - require.NoError(t, err) - source := uint64(bat.Allocated()) - if size := uint64(bat.Size()); size > source { - source = size - } - require.Equal(t, source, uncharged-charged) - - emptyEngine := NewSpillEngine(SpillEngineConfig{}) - retained, ok := emptyEngine.scatterRetainedBytes() - require.True(t, ok) - growth, ok := emptyEngine.scatterCapacityGrowthBytes(bat.RowCount(), 1) - require.True(t, ok) - marshalOverlap, ok := marshalSpillRecordGrowBytes(bat) - require.True(t, ok) - capacity := source + retained + growth + charged + marshalOverlap - budget, err := process.NewHashBuildBudget(capacity, capacity) - require.NoError(t, err) - generation, err := budget.OpenGeneration(capacity) - require.NoError(t, err) - defer generation.Close() - sourceReservation, err := generation.Reserve(source) - require.NoError(t, err) - defer sourceReservation.Release() - - engine := NewSpillEngine(SpillEngineConfig{Budget: generation}) - writers := MakeBucketWriters("test_scatter_charged_source") - defer func() { - for i := range writers { - writers[i].Close() - } + for _, bucket := range []SpillBucket{ + {BuildRows: 1}, + {ProbeRows: 1}, + } { + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + }) + engine.buckets = []SpillBucket{bucket} + _, _, err := engine.RebuildHashmap( + proc, + process.NewAnalyzer(0, false, false, "test"), + ) + require.ErrorContains(t, err, "file/row metadata") engine.Cleanup(proc) - }() - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded( - proc, bat, []*vector.Vector{bat.Vecs[0]}, writers, 0, true, analyzer, - )) -} - -func TestScatterCapacityGrowthChargesCompleteReplacement(t *testing.T) { - engine := NewSpillEngine(SpillEngineConfig{}) - engine.scatterHashValues = make([]uint64, 8) - engine.scatterBucketRowIds = make([]int32, 8) - engine.keyVecs = make([]*vector.Vector, 1) - growth, ok := engine.scatterCapacityGrowthBytes(9, 2) - require.True(t, ok) - require.Equal(t, uint64(9*8+9*4+2*8), growth) -} - -func TestScatterScratchLifecycle(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(64<<20, 64<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(64 << 20) - require.NoError(t, err) - defer generation.Close() - - engine := NewSpillEngine(SpillEngineConfig{Budget: generation}) - writers := MakeBucketWriters("test_scatter_scratch") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - values := make([]int32, 8192) - for i := range values { - values[i] = int32(i) - } - bat := makeInt32Batch(proc, values) - defer bat.Clean(proc.Mp()) - analyzer := process.NewAnalyzer(0, false, false, "test") - keys := []*vector.Vector{bat.Vecs[0]} - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.NotNil(t, engine.scatterScratchReservation) - firstHashCap := cap(engine.scatterHashValues) - firstRowIDCap := cap(engine.scatterBucketRowIds) - require.Equal(t, len(values), firstHashCap) - require.Equal(t, len(values), firstRowIDCap) - firstHash := &engine.scatterHashValues[0] - firstRowID := &engine.scatterBucketRowIds[0] - retained := generation.Used() - firstReserveCount := generation.ReserveCount() - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.Greater(t, generation.ReserveCount(), firstReserveCount, "each batch peak must be admitted above retained scratch") - require.Equal(t, retained, generation.Used(), "batch peak must reconcile to retained scratch") - require.Equal(t, firstHashCap, cap(engine.scatterHashValues)) - require.Equal(t, firstRowIDCap, cap(engine.scatterBucketRowIds)) - require.Equal(t, firstHash, &engine.scatterHashValues[0]) - require.Equal(t, firstRowID, &engine.scatterBucketRowIds[0]) - - for i := range writers { - writers[i].Close() } - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Cap()) - } - // Cleanup is an idempotent terminal release point. - engine.Cleanup(proc) - require.Zero(t, generation.Used()) } -func TestScatterScratchRejectsPeakAboveRetainedBudget(t *testing.T) { +func TestReSpillOmitsUnusedBatchMetadata(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - budget, err := process.NewHashBuildBudget(1<<20, 1<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(1 << 20) - require.NoError(t, err) - defer generation.Close() - engine := NewSpillEngine(SpillEngineConfig{Budget: generation}) - writers := MakeBucketWriters("test_scatter_peak_reject") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - values := make([]int32, 8192) + values := make([]int32, 100) for i := range values { values[i] = int32(i) } - bat := makeInt32Batch(proc, values) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - rowCap := cap(engine.scatterBucketRowIds) - used := generation.Used() - require.Positive(t, used) - largerValues := make([]int32, len(values)*2) - for i := range largerValues { - largerValues[i] = int32(i) - } - larger := makeInt32Batch(proc, largerValues) - defer larger.Clean(proc.Mp()) - err = engine.scatterBatchBounded( - proc, larger, []*vector.Vector{larger.Vecs[0]}, writers, 0, false, analyzer, - ) - require.ErrorIs(t, err, process.ErrHashBuildBudgetAdmission) - require.Equal(t, rowCap, cap(engine.scatterBucketRowIds), "rejection must precede new scratch allocation") - engine.Cleanup(proc) - require.Zero(t, generation.Used()) -} - -func TestScatterPhaseReleasesScratchKeepsSpillOwnership(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - budget, err := process.NewHashBuildBudget(64<<20, 64<<20) - require.NoError(t, err) - generation, err := budget.OpenGeneration(64 << 20) - require.NoError(t, err) - defer generation.Close() - buildBat := makeInt32Batch(proc, []int32{1}) - buildFd := writeBuildFile(proc, "test_phase_build", buildBat) - defer buildBat.Clean(proc.Mp()) - engine := NewSpillEngine(SpillEngineConfig{ - ProbeKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - Budget: generation, + build := makeInt32Batch(proc, values) + build.Attrs = []string{"key"} + defer build.Clean(proc.Mp()) + engine := newExactTestSpillEngine(t, SpillEngineConfig{ + BuildKeyExprs: makeTestKeyExpr(), + NeedBatches: true, + SpillThreshold: 1, + NeedsBuildForEmptyProbe: true, }) - buildFds := make([]*os.File, SpillNumBuckets) - buildFds[0] = buildFd - engine.InitFromSpilledMap(buildFds) - probeBat := makeInt32Batch(proc, []int32{2, 2, 2, 2}) - defer probeBat.Clean(proc.Mp()) - childrenDone := false - analyzer := process.NewAnalyzer(0, false, false, "test") - err = engine.ScatterProbeTable(proc, func() (*batch.Batch, error) { - if childrenDone { - return nil, nil - } - childrenDone = true - return probeBat, nil - }, analyzer, func(bat *batch.Batch) ([]*vector.Vector, error) { - return []*vector.Vector{bat.Vecs[0]}, nil + engine.InitFromSpilledFiles([]*message.SpillFile{ + newTestSpillFile(writeBuildFile(proc, t.Name(), build), int64(len(values))), }) - require.NoError(t, err) - require.Zero(t, generation.Used(), "scatter memory scratch must end with the phase") - require.Nil(t, engine.scatterScratchReservation) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Cap()) - } - require.Positive(t, generation.SpillDiskUsed(), "handed-off probe file keeps disk accounting") - require.Positive(t, generation.SpillFDUsed(), "handed-off probe file keeps FD accounting") - engine.Cleanup(proc) - require.Zero(t, generation.Used()) - require.Zero(t, generation.SpillDiskUsed()) - require.Zero(t, generation.SpillFDUsed()) -} - -func TestScatterCoalescesAcrossBatchesUntilFlush(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - engine := NewSpillEngine(SpillEngineConfig{}) - writers := MakeBucketWriters("test_scatter_coalesce") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - bat := makeInt32Batch(proc, []int32{1, 1, 1}) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - var pending int - for i := range engine.scatterWriteBuffers { - pending += engine.scatterWriteBuffers[i].Len() - } - require.Positive(t, pending) - for i := range writers { - require.Zero(t, writers[i].Rows) - } - require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Len()) - } - var rows int64 - for i := range writers { - rows += writers[i].Rows - } - require.Equal(t, int64(6), rows) -} - -func TestScatterCoalescedRecordRoundTrip(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - engine := NewSpillEngine(SpillEngineConfig{}) - writers := MakeBucketWriters("test_scatter_coalesce_roundtrip") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - bat := makeInt32Batch(proc, []int32{7, 7, 7}) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} + engine.buckets[0].Depth = SpillMaxPass - 1 analyzer := process.NewAnalyzer(0, false, false, "test") - for i := 0; i < 3; i++ { - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - } - require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) - var target *BucketWriter - for i := range writers { - if writers[i].Rows > 0 { - target = &writers[i] - break - } - } - require.NotNil(t, target) - require.Equal(t, int64(9), target.Rows) - _, err := target.Fd.Seek(0, io.SeekStart) - require.NoError(t, err) - reader := BucketReader{fd: target.Fd} - reuse := batch.NewOffHeapWithSize(0) - got, err := reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 3, got.RowCount()) - got, err = reader.ReadBatch(proc, reuse) + _, result, err := engine.RebuildHashmap(proc, analyzer) require.NoError(t, err) - require.Equal(t, 3, got.RowCount()) - got, err = reader.ReadBatch(proc, reuse) - require.NoError(t, err) - require.Equal(t, 3, got.RowCount()) - _, err = reader.ReadBatch(proc, reuse) - require.ErrorIs(t, err, io.EOF) - reader.Close() -} - -func TestScatterCoalesceFlushErrorClearsPending(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - engine := NewSpillEngine(SpillEngineConfig{}) - writers := MakeBucketWriters("test_scatter_coalesce_error") - defer func() { - for i := range writers { - writers[i].Close() - } - }() - bat := makeInt32Batch(proc, []int32{11, 11, 11}) - defer bat.Clean(proc.Mp()) - keys := []*vector.Vector{bat.Vecs[0]} - analyzer := process.NewAnalyzer(0, false, false, "test") - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.NoError(t, engine.flushScatterBuffers(proc, writers, analyzer)) - var target *BucketWriter - for i := range writers { - if writers[i].Rows > 0 { - target = &writers[i] - break - } - } - require.NotNil(t, target) - require.NoError(t, engine.scatterBatchBounded(proc, bat, keys, writers, 0, false, analyzer)) - require.Positive(t, engine.scatterWriteBuffers[targetIndex(writers, target)].Len()) - require.NoError(t, target.Fd.Close()) - require.Error(t, engine.flushScatterBuffers(proc, writers, analyzer)) - for i := range engine.scatterWriteBuffers { - require.Zero(t, engine.scatterWriteBuffers[i].Len()) - } -} + require.Equal(t, BucketReSpilled, result) -func targetIndex(writers []BucketWriter, target *BucketWriter) int { - for i := range writers { - if &writers[i] == target { - return i + found := false + for len(engine.buckets) > 0 { + jm, next, err := engine.RebuildHashmap(proc, analyzer) + require.NoError(t, err) + if next != BucketReady { + continue } - } - return -1 -} - -func TestScatterProbeTableOuterJoinKeepsProbe(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{10, 20, 30}) - fd1 := writeBuildFile(proc, "test_outer_build", bat) - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: makeTestKeyExpr(), - NeedsProbeForEmptyBuild: true, - }) - engine.InitFromSpilledMap([]*os.File{fd1, nil}) - - vals := make([]int32, 1000) - for i := range vals { - vals[i] = int32(i) - } - batches := []*batch.Batch{makeInt32Batch(proc, vals)} - idx := 0 - children := func() (*batch.Batch, error) { - if idx >= len(batches) { - return nil, nil + found = true + for _, bat := range jm.GetBatches() { + require.Empty(t, bat.Attrs) } - b := batches[idx] - idx++ - return b, nil + jm.Free() } - - analyzer := process.NewAnalyzer(0, false, false, "test") - err := engine.ScatterProbeTable(proc, children, analyzer, makeTestEvalKeysFn()) - require.NoError(t, err) - - require.NotNil(t, engine.buckets[0].ProbeFd) - require.NotNil(t, engine.buckets[1].ProbeFd, "outer join must keep probe for empty build") - + require.True(t, found) engine.Cleanup(proc) } func TestCleanupDoubleSafe(t *testing.T) { proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) defer proc.Free() - - engine := NewSpillEngine(SpillEngineConfig{ + engine := newExactTestSpillEngine(t, SpillEngineConfig{ BuildKeyExprs: makeTestKeyExpr(), }) - engine.InitFromSpilledMap([]*os.File{nil, nil, nil}) - + initTestSpillFiles(engine, []*os.File{nil, nil, nil}, 0, 0, 0) engine.Cleanup(proc) engine.Cleanup(proc) - require.Nil(t, engine.buckets) } - -// TestRebuildHashmapPrepareError covers the builder.Free(proc) path -// when HashmapBuilder.Prepare fails (e.g., with an invalid key expression). -func TestRebuildHashmapPrepareError(t *testing.T) { - proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) - defer proc.Free() - - bat := makeInt32Batch(proc, []int32{1, 2, 3}) - fd := writeBuildFile(proc, "test_prep_err", bat) - - // Use an expression that will fail in Prepare (nil Expr field). - badExpr := []*plan.Expr{{}} - - engine := NewSpillEngine(SpillEngineConfig{ - BuildKeyExprs: badExpr, - NeedsBuildForEmptyProbe: true, - }) - engine.InitFromSpilledMap([]*os.File{fd}) - - callbackCalled := false - ok, err := engine.AdvanceToNextBucket(proc, process.NewAnalyzer(0, false, false, "test"), - func(*message.JoinMap, BucketResult) { callbackCalled = true }) - require.Error(t, err) - require.False(t, ok) - require.False(t, callbackCalled) - - engine.Cleanup(proc) -} diff --git a/pkg/sql/compile/allocation_account_lifecycle.go b/pkg/sql/compile/allocation_account_lifecycle.go new file mode 100644 index 0000000000000..02986242bc587 --- /dev/null +++ b/pkg/sql/compile/allocation_account_lifecycle.go @@ -0,0 +1,444 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package compile + +import ( + "errors" + "reflect" + "sync" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/matrixorigin/matrixone/pkg/vm/process" +) + +func allocationLifecycleCall(call func() error) (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = errors.Join( + err, + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtxf( + "allocation lifecycle panic: %v", recovered, + ), + ) + } + }() + return call() +} + +// joinAllocationLifecycleErrors keeps a lone failure's concrete type intact +// and avoids rejoining the same terminal failure. errors.Join wraps even one +// non-nil error, which would turn a statement *moerr.Error into a generic Go +// error before it crosses the pipeline wire. +func joinAllocationLifecycleErrors(primary, secondary error) error { + if primary == nil { + return secondary + } + if secondary == nil { + return primary + } + if reflect.TypeOf(primary).Comparable() && primary == secondary { + return primary + } + return errors.Join(primary, secondary) +} + +type executionAllocationAccountOwner interface { + SetAllocationAccount(*mpool.AllocationAccount) error + ClearAllocationAccount(*mpool.AllocationAccount) error +} + +type executionAllocationAccountActivationPolicy interface { + ActivatesAllocationAccountLifecycle() bool +} + +func hasAllocationAccountActivator( + owners []executionAllocationAccountOwner, +) bool { + for _, owner := range owners { + policy, ok := owner.(executionAllocationAccountActivationPolicy) + if !ok || policy.ActivatesAllocationAccountLifecycle() { + return true + } + } + return false +} + +// statementAllocationAttempt owns one local execution generation. The +// MessageBoard pointer is captured at open so prepared/retry Reset cannot make +// terminal cleanup drain a newer board. +type statementAllocationAttempt struct { + registry *mpool.AllocationAccountRegistry + account *mpool.AllocationAccount + board *message.MessageBoard + exporter func(mpool.AllocationAccountTerminalSnapshot) + + ownersMu sync.Mutex + owners []executionAllocationAccountOwner + ownerSet map[executionAllocationAccountOwner]struct{} + closing bool + + prepareOnce sync.Once + completeOnce sync.Once + snapshot mpool.AllocationAccountTerminalSnapshot + prepareErr error + completeErr error +} + +func (c *Compile) beginAllocationAccountAttempt() ( + *statementAllocationAttempt, + error, +) { + if c == nil || c.allocationAccountRegistry == nil { + return nil, nil + } + if c.proc == nil || c.MessageBoard == nil || c.allocationAttempt != nil || + c.allocationTerminalExporter == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + owners := c.allocationAccountOwners + var err error + if owners == nil { + owners, err = collectAllocationAccountOwners(c.scopes) + if err != nil { + return nil, err + } + } + c.allocationAccountOwners = nil + if c.allocationControllerProvider == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + controller, err := c.allocationControllerProvider() + if err != nil { + return nil, err + } + if controller == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + account, err := c.allocationAccountRegistry.OpenWithController( + c.allocationAccountLimit, + controller, + ) + if err != nil { + return nil, err + } + owners, err = configureAllocationAccountOwners(owners, account) + if err != nil { + var snapshot mpool.AllocationAccountTerminalSnapshot + var first bool + finalizeErr := allocationLifecycleCall(func() error { + var terminalErr error + snapshot, first, terminalErr = c.allocationAccountRegistry. + CompleteTerminalWithError(account, err) + return terminalErr + }) + if first { + finalizeErr = joinAllocationLifecycleErrors( + finalizeErr, + allocationLifecycleCall(func() error { + c.allocationTerminalExporter(snapshot) + return nil + }), + ) + } + if finalizeErr != nil { + return nil, joinAllocationLifecycleErrors(err, finalizeErr) + } + return nil, err + } + attempt := &statementAllocationAttempt{ + registry: c.allocationAccountRegistry, + account: account, + board: c.MessageBoard, + exporter: c.allocationTerminalExporter, + owners: owners, + ownerSet: make(map[executionAllocationAccountOwner]struct{}, len(owners)), + } + for _, owner := range owners { + attempt.ownerSet[owner] = struct{}{} + } + c.allocationAttempt = attempt + return attempt, nil +} + +// attachRuntimeOwners binds operators cloned after runOnce starts to the same +// attempt. Parallel scan/load workers are execution-local and do not exist +// when the template scopes are collected. +func (a *statementAllocationAttempt) attachRuntimeOwners(scopes []*Scope) error { + if a == nil || a.account == nil { + return mpool.ErrAllocationAccountInvariant + } + owners, err := collectAllocationAccountOwners(scopes) + if err != nil || len(owners) == 0 { + return err + } + + a.ownersMu.Lock() + defer a.ownersMu.Unlock() + if a.closing { + return mpool.ErrAllocationAccountInvariant + } + newOwners := make([]executionAllocationAccountOwner, 0, len(owners)) + for _, owner := range owners { + if _, exists := a.ownerSet[owner]; !exists { + newOwners = append(newOwners, owner) + } + } + configured, err := configureAllocationAccountOwners(newOwners, a.account) + if err != nil { + return err + } + for _, owner := range configured { + a.ownerSet[owner] = struct{}{} + } + a.owners = append(a.owners, configured...) + return nil +} + +func (c *Compile) attachRuntimeAllocationOwners(scopes []*Scope) error { + if c == nil { + return mpool.ErrAllocationAccountInvariant + } + if c.allocationAttempt == nil { + owners, err := collectAllocationAccountOwners(scopes) + if err != nil { + return err + } + if !hasAllocationAccountActivator(owners) { + return nil + } + return mpool.ErrAllocationAccountInvariant + } + return c.allocationAttempt.attachRuntimeOwners(scopes) +} + +func configureAllocationAccountOwners( + owners []executionAllocationAccountOwner, + account *mpool.AllocationAccount, +) ([]executionAllocationAccountOwner, error) { + configured := make([]executionAllocationAccountOwner, 0, len(owners)) + rollback := func(cause error) error { + for i := len(configured) - 1; i >= 0; i-- { + cause = joinAllocationLifecycleErrors( + cause, + allocationLifecycleCall(func() error { + return configured[i].ClearAllocationAccount(account) + }), + ) + } + return cause + } + for _, owner := range owners { + if err := allocationLifecycleCall(func() error { + return owner.SetAllocationAccount(account) + }); err != nil { + return nil, rollback(err) + } + configured = append(configured, owner) + } + return configured, nil +} + +func collectAllocationAccountOwners( + scopes []*Scope, +) ([]executionAllocationAccountOwner, error) { + owners := make([]executionAllocationAccountOwner, 0) + seen := make(map[executionAllocationAccountOwner]struct{}) + var inspect func(*Scope) error + inspect = func(scope *Scope) error { + if scope == nil { + return nil + } + if err := vm.HandleAllOp(scope.RootOp, func(_ vm.Operator, op vm.Operator) error { + if owner, ok := op.(executionAllocationAccountOwner); ok { + if _, exists := seen[owner]; !exists { + seen[owner] = struct{}{} + owners = append(owners, owner) + } + } + return nil + }); err != nil { + return err + } + for _, preScope := range scope.PreScopes { + if err := inspect(preScope); err != nil { + return err + } + } + return nil + } + for _, scope := range scopes { + if err := inspect(scope); err != nil { + return nil, err + } + } + return owners, nil +} + +// ensureAllocationAccountLifecycle installs one account whenever the physical +// plan contains a HashBuild/join allocation owner. Implementing the owner +// contract is the boundary: there is no per-owner activation switch. +func (c *Compile) ensureAllocationAccountLifecycle( + exporter func(mpool.AllocationAccountTerminalSnapshot), +) error { + if c == nil { + return nil + } + owners, err := collectAllocationAccountOwners(c.scopes) + if err != nil { + return err + } + c.allocationAccountOwners = owners + if !hasAllocationAccountActivator(owners) { + c.allocationAccountOwners = nil + if c.allocationControllerProvider != nil { + c.allocationAccountRegistry = nil + c.allocationAccountLimit = 0 + c.allocationControllerProvider = nil + c.allocationTerminalExporter = nil + } + return nil + } + if exporter == nil { + return mpool.ErrAllocationAccountInvariant + } + if c.proc == nil { + return mpool.ErrAllocationAccountInvariant + } + budget, err := c.proc.GetHashBuildBudget() + if err != nil { + return err + } + registry, err := budget.AllocationAccountRegistry() + if err != nil { + return err + } + limit := budget.Snapshot().Cap + if limit == 0 { + return mpool.ErrAllocationAccountInvariant + } + c.allocationAccountRegistry = registry + c.allocationAccountLimit = limit + c.allocationControllerProvider = func() (mpool.AllocationCapacityController, error) { + if budget.Closed() { + return nil, process.ErrHashBuildBudgetClosed + } + return budget, nil + } + c.allocationTerminalExporter = exporter + return nil +} + +func (a *statementAllocationAttempt) finish() ( + mpool.AllocationAccountTerminalSnapshot, + error, +) { + if a == nil { + return mpool.AllocationAccountTerminalSnapshot{}, nil + } + a.prepareTerminal(true) + return a.completeTerminal() +} + +// prepareTerminal closes the operator-owned part of an attempt after every +// scope producer has quiesced. A coordinator-owned board can be closed here. +// Remote fragments share one board on a CN, so their statement group closes it +// only after every expected fragment has reached this boundary. +func (a *statementAllocationAttempt) prepareTerminal(closeBoard bool) error { + if a == nil { + return nil + } + a.prepareOnce.Do(func() { + if closeBoard { + a.prepareErr = joinAllocationLifecycleErrors( + a.prepareErr, + allocationLifecycleCall(func() error { + a.board.CloseAndDrain() + return nil + }), + ) + } + a.ownersMu.Lock() + a.closing = true + owners := a.owners + a.owners = nil + a.ownerSet = nil + a.ownersMu.Unlock() + for i := len(owners) - 1; i >= 0; i-- { + a.prepareErr = joinAllocationLifecycleErrors( + a.prepareErr, + allocationLifecycleCall(func() error { + return owners[i].ClearAllocationAccount(a.account) + }), + ) + } + }) + return a.prepareErr + +} + +func (a *statementAllocationAttempt) completeTerminal() ( + mpool.AllocationAccountTerminalSnapshot, + error, +) { + if a == nil { + return mpool.AllocationAccountTerminalSnapshot{}, nil + } + a.completeOnce.Do(func() { + prepareErr := a.prepareTerminal(false) + var first bool + a.completeErr = allocationLifecycleCall(func() error { + var terminalErr error + a.snapshot, first, terminalErr = a.registry.CompleteTerminalWithError( + a.account, + prepareErr, + ) + return terminalErr + }) + if first && a.exporter != nil { + a.completeErr = joinAllocationLifecycleErrors( + a.completeErr, + allocationLifecycleCall(func() error { + a.exporter(a.snapshot) + return nil + }), + ) + } + }) + return a.snapshot, a.completeErr +} + +func (c *Compile) finishAllocationAccountAttempt() error { + if c == nil || c.allocationAttempt == nil { + return nil + } + attempt := c.allocationAttempt + c.allocationAttempt = nil + _, err := attempt.finish() + return err +} + +func (c *Compile) copyAllocationAccountLifecycleTo(dst *Compile) { + if c == nil || dst == nil { + return + } + dst.allocationAccountRegistry = c.allocationAccountRegistry + dst.allocationAccountLimit = c.allocationAccountLimit + dst.allocationTerminalExporter = c.allocationTerminalExporter + dst.allocationControllerProvider = c.allocationControllerProvider +} diff --git a/pkg/sql/compile/allocation_account_lifecycle_test.go b/pkg/sql/compile/allocation_account_lifecycle_test.go new file mode 100644 index 0000000000000..3ffb90ea636a0 --- /dev/null +++ b/pkg/sql/compile/allocation_account_lifecycle_test.go @@ -0,0 +1,699 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package compile + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/pb/txn" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/connector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/dispatch" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/hashbuild" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/product" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm" + "github.com/matrixorigin/matrixone/pkg/vm/engine" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/matrixorigin/matrixone/pkg/vm/process" + "github.com/stretchr/testify/require" +) + +func TestJoinAllocationLifecycleErrorsPreservesSingle(t *testing.T) { + primary := moerr.NewDuplicateEntryNoCtx("duplicate", "primary") + secondary := errors.New("cleanup failed") + + require.Same(t, primary, joinAllocationLifecycleErrors(primary, nil)) + require.Same(t, secondary, joinAllocationLifecycleErrors(nil, secondary)) + require.Same(t, primary, joinAllocationLifecycleErrors(primary, primary)) + + joined := joinAllocationLifecycleErrors(primary, secondary) + require.ErrorIs(t, joined, primary) + require.ErrorIs(t, joined, secondary) +} + +type allocationLifecycleErrorOperator struct { + *colexec.MockOperator + err error + account *mpool.AllocationAccount +} + +type allocationLifecycleOwnerOperator struct { + *colexec.MockOperator + account *mpool.AllocationAccount + failSet bool + failClear bool + panicClear bool + clears int + released bool + releaseSawLiveAccount bool +} + +func (op *allocationLifecycleOwnerOperator) Release() { + op.released = true + op.releaseSawLiveAccount = op.account != nil +} + +func (op *allocationLifecycleOwnerOperator) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.failSet { + return mpool.ErrAllocationAccountMismatch + } + if op.account != nil && op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = account + return nil +} + +func (op *allocationLifecycleOwnerOperator) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.panicClear { + panic("test allocation owner clear panic") + } + if op.account == nil { + return nil + } + if op.failClear { + return mpool.ErrAllocationAccountInvariant + } + if op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = nil + op.clears++ + return nil +} + +func TestStatementAllocationAttemptOwnerTeardownPanicIsTerminalFailure(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + panicClear: true, + } + c.scopes = []*Scope{{RootOp: owner}} + + _, err = c.beginAllocationAccountAttempt() + require.NoError(t, err) + err = c.finishAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + require.Zero(t, registry.LiveAllocationMetadata()) +} + +func (op *allocationLifecycleErrorOperator) Call( + *process.Process, +) (vm.CallResult, error) { + return vm.CancelResult, op.err +} + +func (op *allocationLifecycleErrorOperator) SetAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.account != nil && op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = account + return nil +} + +func (op *allocationLifecycleErrorOperator) ClearAllocationAccount( + account *mpool.AllocationAccount, +) error { + if op.account != account { + return mpool.ErrAllocationAccountMismatch + } + op.account = nil + return nil +} + +type allocationLifecycleTestController struct{} + +func (*allocationLifecycleTestController) AcquireAllocationCapacity(uint64) error { + return nil +} + +func (*allocationLifecycleTestController) ReleaseAllocationCapacity(uint64) {} + +func newRunLifecycleCompile( + t *testing.T, +) (*Compile, *mpool.AllocationAccountRegistry) { + t.Helper() + proc := testutil.NewProcess(t) + ctrl := gomock.NewController(t) + txnClient, txnOperator := newTestTxnClientAndOpWithIsolation( + ctrl, + txn.TxnIsolation_RC, + ) + proc.Base.TxnClient = txnClient + proc.Base.TxnOperator = txnOperator + proc.ReplaceTopCtx(context.Background()) + c := NewCompile( + "local", + "", + "select 1", + "", + "", + nil, + proc, + nil, + false, + nil, + time.Now(), + ) + c.pn = &plan.Plan{Plan: &plan.Plan_Query{Query: &plan.Query{}}} + c.anal = newAnalyzeModule() + budget, err := proc.GetHashBuildBudget() + require.NoError(t, err) + registry, err := budget.AllocationAccountRegistry() + require.NoError(t, err) + return c, registry +} + +func newTestAllocationLifecycleCompile( + t *testing.T, + registry *mpool.AllocationAccountRegistry, + exporter func(mpool.AllocationAccountTerminalSnapshot), +) *Compile { + t.Helper() + return &Compile{ + proc: testutil.NewProcess(t), + MessageBoard: message.NewMessageBoard(), + allocationAccountRegistry: registry, + allocationAccountLimit: 1 << 20, + allocationControllerProvider: func() (mpool.AllocationCapacityController, error) { + return &allocationLifecycleTestController{}, nil + }, + allocationTerminalExporter: exporter, + } +} + +func TestStatementAllocationAttemptZeroTerminalExportsOnce(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotNil(t, attempt) + buffer, err := c.proc.Mp().AllocAccounted( + 64, + attempt.account, + 1, + 1, + ) + require.NoError(t, err) + c.proc.Mp().Free(buffer) + + require.NoError(t, c.finishAllocationAccountAttempt()) + require.Len(t, exported, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, exported[0].State) + require.Zero(t, exported[0].Used) + require.Equal(t, uint64(64), exported[0].Peak) + _, ok := registry.Resolve(exported[0].Handle) + require.False(t, ok) + + repeated, err := attempt.finish() + require.NoError(t, err) + require.Equal(t, exported[0], repeated) + require.Len(t, exported, 1) + require.NotSame(t, c.MessageBoard, c.MessageBoard.Reset()) +} + +func TestStatementAllocationAttemptLateFreeDrainsTombstone(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + buffer, err := c.proc.Mp().AllocAccounted( + 64, + attempt.account, + 1, + 1, + ) + require.NoError(t, err) + + err = c.finishAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + require.Equal(t, uint64(cap(buffer)), exported[0].Used) + require.True(t, registry.AdmissionSuspended()) + _, err = registry.Open(1) + require.ErrorIs(t, err, mpool.ErrAllocationAdmissionSuspended) + + // The physical allocation retains its original account after the producer + // process has detached the generation. Its normal Free drains the + // tombstone; no synthetic release is needed. + c.proc.Mp().Free(buffer) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(exported[0].Handle) + require.False(t, ok) + + c.MessageBoard = c.MessageBoard.Reset() + next, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotEqual(t, attempt.account.Handle(), next.account.Handle()) + require.NoError(t, c.finishAllocationAccountAttempt()) + require.Len(t, exported, 2) +} + +func TestStatementAllocationAttemptConcurrentTerminalIsOneShot(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exports atomic.Int32 + c := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + exports.Add(1) + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + + const contenders = 128 + start := make(chan struct{}) + errs := make(chan error, contenders) + var wait sync.WaitGroup + wait.Add(contenders) + for range contenders { + go func() { + defer wait.Done() + <-start + _, finishErr := attempt.finish() + errs <- finishErr + }() + } + close(start) + wait.Wait() + close(errs) + for finishErr := range errs { + require.NoError(t, finishErr) + } + require.Equal(t, int32(1), exports.Load()) + c.allocationAttempt = nil +} + +func TestStatementAllocationAttemptRejectsOverlappingOpen(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + c := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + }) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + + _, err = c.beginAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.NoError(t, c.finishAllocationAccountAttempt()) + _, finishErr := attempt.finish() + require.NoError(t, finishErr) + + next, err := registry.Open(1) + require.NoError(t, err, "rejected overlap must not consume a slot") + _, _, err = registry.CompleteTerminal(next) + require.NoError(t, err) +} + +func TestStatementAllocationAttemptRequiresTerminalExporter(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + c := newTestAllocationLifecycleCompile(t, registry, nil) + + _, err = c.beginAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + account, err := registry.Open(1) + require.NoError(t, err, "failed begin must not consume a generation slot") + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestStatementAllocationAttemptOwnerConfigurationRollsBack(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + configured := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + rejected := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + failSet: true, + } + c.scopes = []*Scope{ + {RootOp: configured}, + {RootOp: rejected}, + } + + _, err = c.beginAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountMismatch) + require.Nil(t, configured.account) + require.Equal(t, 1, configured.clears) + require.Nil(t, c.allocationAttempt) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + + account, err := registry.Open(1) + require.NoError(t, err, "failed owner configuration leaked its registry slot") + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestAllocationAccountConfiguresEveryStatementOwner(t *testing.T) { + first := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + second := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + scopes := []*Scope{{RootOp: first}, {RootOp: second}} + owners, err := collectAllocationAccountOwners(scopes) + require.NoError(t, err) + require.Len(t, owners, 2) + + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + configured, err := configureAllocationAccountOwners(owners, account) + require.NoError(t, err) + require.Len(t, configured, 2) + require.Same(t, account, first.account) + require.Same(t, account, second.account) + for i := len(configured) - 1; i >= 0; i-- { + require.NoError(t, configured[i].ClearAllocationAccount(account)) + } + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestAllocationAccountTransportParticipantsDoNotActivateLifecycle(t *testing.T) { + connectorOp := connector.NewArgument() + dispatchOp := dispatch.NewArgument() + t.Cleanup(connectorOp.Release) + t.Cleanup(dispatchOp.Release) + c := &Compile{ + proc: testutil.NewProcess(t), + MessageBoard: message.NewMessageBoard(), + scopes: []*Scope{ + {RootOp: connectorOp}, + {RootOp: dispatchOp}, + }, + } + require.NoError(t, c.ensureAllocationAccountLifecycle(func( + mpool.AllocationAccountTerminalSnapshot, + ) { + })) + require.Nil(t, c.allocationAccountOwners) + require.Nil(t, c.allocationControllerProvider) + require.Nil(t, c.allocationAccountRegistry) + require.NoError(t, c.attachRuntimeAllocationOwners(c.scopes)) + + transportOwners := []executionAllocationAccountOwner{ + connectorOp, + dispatchOp, + } + require.False(t, hasAllocationAccountActivator(transportOwners)) + + active := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + require.True(t, hasAllocationAccountActivator(append(transportOwners, active))) + require.ErrorIs(t, c.attachRuntimeAllocationOwners([]*Scope{{ + RootOp: active, + }}), mpool.ErrAllocationAccountInvariant) +} + +func TestAllocationAccountCollectsProductConsumerAndHashBuild(t *testing.T) { + consumer := product.NewArgument() + producer := hashbuild.NewArgument() + scopes := []*Scope{{ + RootOp: consumer, + PreScopes: []*Scope{ + {RootOp: producer}, + }, + }} + owners, err := collectAllocationAccountOwners(scopes) + require.NoError(t, err) + require.Len(t, owners, 2) + + registry, err := mpool.NewAllocationAccountRegistry(1, 32) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + configured, err := configureAllocationAccountOwners(owners, account) + require.NoError(t, err) + require.Len(t, configured, 2) + for i := len(configured) - 1; i >= 0; i-- { + require.NoError(t, configured[i].ClearAllocationAccount(account)) + } + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) + consumer.Release() + producer.Release() +} + +func TestParallelRuntimeClonesJoinAllocationAttempt(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 16) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + template := hashbuild.NewArgument() + template.NeedHashMap = false + source := &Scope{ + RootOp: template, + Proc: c.proc, + NodeInfo: engine.Node{ + Mcpu: 2, + }, + } + c.scopes = []*Scope{source} + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotNil(t, attempt) + + parallel, workers := newParallelScope(source) + require.Len(t, workers, 2) + for _, worker := range workers { + hb := worker.RootOp.(*hashbuild.HashBuild) + require.ErrorIs(t, hb.Prepare(worker.Proc), mpool.ErrAllocationAccountInvalid) + } + require.NoError(t, c.attachRuntimeAllocationOwners(workers)) + for _, worker := range workers { + hb := worker.RootOp.(*hashbuild.HashBuild) + require.NoError(t, hb.Prepare(worker.Proc)) + } + + require.NoError(t, c.finishAllocationAccountAttempt()) + require.Len(t, exported, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, exported[0].State) + require.Zero(t, exported[0].Used) + parallel.release() + template.Release() +} + +func TestStatementAllocationAttemptOwnerTeardownFailureExportsFailure(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + failClear: true, + } + c.scopes = []*Scope{{RootOp: owner}} + + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + err = c.finishAllocationAccountAttempt() + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Len(t, exported, 1) + require.Equal( + t, + mpool.AllocationAccountTerminalInvariantFailure, + exported[0].State, + ) + require.Zero(t, exported[0].Used) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(attempt.account.Handle()) + require.False(t, ok) +} + +func TestCompileClearFinalizesAllocationOwnerBeforeRelease(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + var exported []mpool.AllocationAccountTerminalSnapshot + c := newTestAllocationLifecycleCompile(t, registry, func( + snapshot mpool.AllocationAccountTerminalSnapshot, + ) { + exported = append(exported, snapshot) + }) + c.affectRows = &atomic.Uint64{} + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + c.scopes = []*Scope{{RootOp: owner}} + + _, err = c.beginAllocationAccountAttempt() + require.NoError(t, err) + c.clear() + + require.True(t, owner.released) + require.False(t, owner.releaseSawLiveAccount) + require.Nil(t, owner.account) + require.Equal(t, 1, owner.clears) + require.Len(t, exported, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, exported[0].State) + require.Zero(t, exported[0].Used) +} + +func TestCompileAutomaticallyActivatesCompleteHashTableOwner(t *testing.T) { + proc := testutil.NewProcess(t) + c := &Compile{ + proc: proc, + MessageBoard: message.NewMessageBoard(), + } + owner := hashbuild.NewArgument() + owner.NeedHashMap = true + owner.Conditions = []*plan.Expr{{ + Typ: plan.Type{Id: int32(types.T_int64)}, + Expr: &plan.Expr_Col{Col: &plan.ColRef{}}, + }} + c.scopes = []*Scope{{RootOp: owner}} + + require.NoError(t, c.ensureAllocationAccountLifecycle(func( + mpool.AllocationAccountTerminalSnapshot, + ) { + })) + require.NotNil(t, c.allocationControllerProvider) + require.NotNil(t, c.allocationAccountRegistry) + attempt, err := c.beginAllocationAccountAttempt() + require.NoError(t, err) + require.NotNil(t, attempt) + require.NoError(t, owner.ClearAllocationAccount(attempt.account)) + require.NoError(t, c.finishAllocationAccountAttempt()) + _, ok := c.allocationAccountRegistry.Resolve(attempt.account.Handle()) + require.False(t, ok) + owner.Release() +} + +func TestCompileRunFinalizesAllocationAttemptOnCancellation(t *testing.T) { + c, registry := newRunLifecycleCompile(t) + // The canceled outer context is observed after runOnce, after the + // allocation generation has opened. + canceled, cancel := context.WithCancel(context.Background()) + cancel() + c.proc.ReplaceTopCtx(canceled) + scope := newScope(magicType(255)) + owner := &allocationLifecycleOwnerOperator{MockOperator: colexec.NewMockOperator()} + scope.RootOp = owner + c.scopes = []*Scope{scope} + + _, err := c.Run(0) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, owner.account) + require.Zero(t, registry.LiveAllocationMetadata()) + c.Release() +} + +func TestCompileRunFinalizesAllocationAttemptOnExecutionError(t *testing.T) { + c, registry := newRunLifecycleCompile(t) + executionErr := moerr.NewInternalErrorNoCtx("allocation lifecycle test") + scope := newScope(Normal) + scope.Proc = c.proc.NewNoContextChildProc(0) + owner := &allocationLifecycleErrorOperator{ + MockOperator: colexec.NewMockOperator(), + err: executionErr, + } + scope.RootOp = owner + c.scopes = []*Scope{scope} + + _, err := c.Run(0) + require.ErrorIs(t, err, executionErr) + require.Nil(t, owner.account) + require.Zero(t, registry.LiveAllocationMetadata()) + c.Release() +} + +func TestCompileRunFinalizesAllocationAttemptOnPanic(t *testing.T) { + c, registry := newRunLifecycleCompile(t) + scope := newScope(magicType(255)) + owner := &allocationLifecycleOwnerOperator{MockOperator: colexec.NewMockOperator()} + scope.RootOp = owner + c.scopes = []*Scope{scope} + // Force the panic after beginAllocationAccountAttempt and before runOnce. + c.lockMeta = nil + + require.Panics(t, func() { + _, _ = c.Run(0) + }) + require.Nil(t, owner.account) + require.Zero(t, registry.LiveAllocationMetadata()) + c.Release() +} diff --git a/pkg/sql/compile/analyze_module.go b/pkg/sql/compile/analyze_module.go index 4009c26681899..cff5353e893d9 100644 --- a/pkg/sql/compile/analyze_module.go +++ b/pkg/sql/compile/analyze_module.go @@ -17,6 +17,7 @@ package compile import ( "bytes" "fmt" + "slices" "strings" "sync" @@ -44,15 +45,18 @@ type AnalyzeModule struct { phyPlan *models.PhyPlan remotePhyPlans []models.PhyPlan // Added read-write lock - mu sync.RWMutex - retryTimes int - explainPhyBuffer *bytes.Buffer - remoteUsage resource.Usage - remoteMemory resource.MemoryTotals - remoteQuality resource.QualityFlags - remoteMissingFragments uint64 - remoteMissingMemoryDomains uint64 - remoteReports uint64 + mu sync.RWMutex + retryTimes int + explainPhyBuffer *bytes.Buffer + remoteUsage resource.Usage + remoteMemory resource.MemoryTotals + remoteAllocation resource.AllocationAccountTotals + remoteQuality resource.QualityFlags + remoteMissingFragments uint64 + remoteMissingMemoryDomains uint64 + remoteReports uint64 + remotePendingAllocationGroups map[string]uint64 + remoteCompletedAllocationGroups map[string]struct{} } // remoteResourceSnapshot is a by-value view of the terminal resource facts @@ -60,12 +64,15 @@ type AnalyzeModule struct { // separately from the nested missing counts so each hop can account for the // remote scopes it expected to hear from exactly once. type remoteResourceSnapshot struct { - Usage resource.Usage - Memory resource.MemoryTotals - Quality resource.QualityFlags - MissingFragmentCount uint64 - MissingMemoryDomainCount uint64 - DirectReportCount uint64 + Usage resource.Usage + Memory resource.MemoryTotals + Allocation resource.AllocationAccountTotals + Quality resource.QualityFlags + MissingFragmentCount uint64 + MissingMemoryDomainCount uint64 + DirectReportCount uint64 + PendingAllocationGroups []remoteAllocationGroupPending + CompletedAllocationGroups []string } // Reset When Compile reused, reset AnalyzeModule to prevent resource accumulation @@ -82,10 +89,13 @@ func (anal *AnalyzeModule) Reset(isPrepare bool, isTpQuery bool) { anal.retryTimes = 0 anal.remoteUsage = resource.Usage{} anal.remoteMemory = resource.MemoryTotals{} + anal.remoteAllocation = resource.AllocationAccountTotals{} anal.remoteQuality = 0 anal.remoteMissingFragments = 0 anal.remoteMissingMemoryDomains = 0 anal.remoteReports = 0 + anal.remotePendingAllocationGroups = nil + anal.remoteCompletedAllocationGroups = nil if anal.qry != nil { for _, node := range anal.qry.Nodes { if node.AnalyzeInfo == nil { @@ -102,8 +112,11 @@ func (anal *AnalyzeModule) Reset(isPrepare bool, isTpQuery bool) { func (anal *AnalyzeModule) appendRemoteResource( delta resource.Delta, memory resource.MemoryTotals, + allocation resource.AllocationAccountTotals, missingFragments uint64, missingMemoryDomains uint64, + pendingAllocationGroups []remoteAllocationGroupPending, + completedAllocationGroups []string, ) { if anal == nil { return @@ -113,14 +126,77 @@ func (anal *AnalyzeModule) appendRemoteResource( quality := delta.Quality quality |= resource.MergeUsage(&anal.remoteUsage, delta.Usage) quality |= resource.MergeMemoryTotals(&anal.remoteMemory, memory) + quality |= resource.MergeAllocationAccountTotals( + &anal.remoteAllocation, + allocation, + ) anal.remoteMissingFragments, quality = addCheckedRemoteCounter( anal.remoteMissingFragments, missingFragments, quality) anal.remoteMissingMemoryDomains, quality = addCheckedRemoteCounter( anal.remoteMissingMemoryDomains, missingMemoryDomains, quality) anal.remoteReports, quality = addCheckedRemoteCounter(anal.remoteReports, 1, quality) + if len(completedAllocationGroups) > 0 && anal.remoteCompletedAllocationGroups == nil { + anal.remoteCompletedAllocationGroups = make(map[string]struct{}) + } + for _, key := range completedAllocationGroups { + if key == "" { + quality |= resource.QualityInvariantFailure + continue + } + anal.remoteCompletedAllocationGroups[key] = struct{}{} + delete(anal.remotePendingAllocationGroups, key) + } + if len(pendingAllocationGroups) > 0 && anal.remotePendingAllocationGroups == nil { + anal.remotePendingAllocationGroups = make(map[string]uint64) + } + for _, signal := range pendingAllocationGroups { + if signal.Key == "" || signal.Count == 0 { + quality |= resource.QualityInvariantFailure + continue + } + if _, completed := anal.remoteCompletedAllocationGroups[signal.Key]; !completed { + anal.remotePendingAllocationGroups[signal.Key], quality = + addCheckedRemoteCounter( + anal.remotePendingAllocationGroups[signal.Key], + signal.Count, + quality, + ) + } + } anal.remoteQuality |= quality } +func sortedRemoteAllocationGroupKeys(values map[string]struct{}) []string { + if len(values) == 0 { + return nil + } + keys := make([]string, 0, len(values)) + for key := range values { + keys = append(keys, key) + } + slices.Sort(keys) + return keys +} + +func sortedRemoteAllocationGroupPending( + values map[string]uint64, +) []remoteAllocationGroupPending { + if len(values) == 0 { + return nil + } + signals := make([]remoteAllocationGroupPending, 0, len(values)) + for key, count := range values { + signals = append(signals, remoteAllocationGroupPending{ + Key: key, + Count: count, + }) + } + slices.SortFunc(signals, func(a, b remoteAllocationGroupPending) int { + return strings.Compare(a.Key, b.Key) + }) + return signals +} + func (anal *AnalyzeModule) remoteResourceSummary() remoteResourceSnapshot { if anal == nil { return remoteResourceSnapshot{} @@ -130,10 +206,17 @@ func (anal *AnalyzeModule) remoteResourceSummary() remoteResourceSnapshot { return remoteResourceSnapshot{ Usage: anal.remoteUsage, Memory: anal.remoteMemory, + Allocation: anal.remoteAllocation, Quality: anal.remoteQuality, MissingFragmentCount: anal.remoteMissingFragments, MissingMemoryDomainCount: anal.remoteMissingMemoryDomains, DirectReportCount: anal.remoteReports, + PendingAllocationGroups: sortedRemoteAllocationGroupPending( + anal.remotePendingAllocationGroups, + ), + CompletedAllocationGroups: sortedRemoteAllocationGroupKeys( + anal.remoteCompletedAllocationGroups, + ), } } diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go index 7cdfd302f7435..0ddd40cbee114 100644 --- a/pkg/sql/compile/compile.go +++ b/pkg/sql/compile/compile.go @@ -28,6 +28,7 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/parquet-go/parquet-go" "github.com/matrixorigin/matrixone/pkg/catalog" @@ -240,6 +241,8 @@ func (c *Compile) Reset(proc *process.Process, startAt time.Time, fill func(*bat c.MessageBoard = c.MessageBoard.Reset() proc.SetMessageBoard(c.MessageBoard) + c.remoteFragmentCounts = nil + c.remoteExecutionID = uuid.Nil c.counterSet.Reset() for _, f := range c.fuzzys { @@ -277,6 +280,12 @@ func (c *Compile) clear() { if c.anal != nil { c.anal.release() } + // The attempt owns references to allocation-aware operators. Finalize it + // before Scope.release returns those operators to reuse pools; otherwise a + // defensive cleanup path could clear an already-reset or reused owner. + if err := c.finishAllocationAccountAttempt(); err != nil { + logutil.Errorf("allocation account terminal cleanup failed: %v", err) + } for i := range c.scopes { c.scopes[i].release() } @@ -317,6 +326,14 @@ func (c *Compile) clear() { c.needLockMeta = false c.isInternal = false c.resourceAttemptOwnerEligible = false + c.allocationAccountRegistry = nil + c.allocationAccountLimit = 0 + c.allocationControllerProvider = nil + c.allocationTerminalExporter = nil + c.allocationAccountOwners = nil + c.allocationAttempt = nil + c.remoteFragmentCounts = nil + c.remoteExecutionID = uuid.Nil c.isPrepare = false c.hasMergeOp = false c.needBlock = false @@ -616,11 +633,7 @@ func (c *Compile) prePipelineInitializer() (err error) { func newMaterializedSpillBudget(proc *process.Process) materialized.SpillBudget { return materialized.SpillBudget{ ReserveMemory: func(size uint64) (materialized.Reservation, error) { - budget, err := proc.GetHashBuildBudget() - if err != nil { - return nil, err - } - return budget.Reserve(size) + return proc.GetCTEMemoryBudget().Reserve(proc.Ctx, size) }, ReserveDisk: func(size uint64) (materialized.GrowingReservation, error) { budget, err := proc.GetHashBuildBudget() diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index f8b44fe215a9e..59c82c6939d9b 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -22,7 +22,9 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" commonutil "github.com/matrixorigin/matrixone/pkg/common/util" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/defines" @@ -258,8 +260,22 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { attemptAnal := runC.anal var coordinatorPhaseStart time.Time var coordinatorPhaseBase time.Duration + var allocationAttempt *statementAllocationAttempt + finishAllocationAttempt := func() error { + if allocationAttempt == nil { + return nil + } + attempt := allocationAttempt + allocationAttempt = nil + if runC != nil && runC.allocationAttempt == attempt { + runC.allocationAttempt = nil + } + _, finishErr := attempt.finish() + return finishErr + } defer func() { if recovered := recover(); recovered != nil { + _ = finishAllocationAttempt() if attemptOpen { if !coordinatorPhaseStart.IsZero() { attemptPreRunWall = coordinatorPhaseBase + time.Since(coordinatorPhaseStart) @@ -288,7 +304,30 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // Before compile.runOnce, Reset the 'StatsInfo' execution related resources in context // running. - if err = runC.prePipelineInitializer(); err == nil { + if runC.remoteFragmentCounts == nil { + runC.remoteFragmentCounts = collectRemoteFragmentCounts(runC.scopes, runC.addr) + } + // A retry is a new physical execution generation. Reusing the previous + // ID could attach late RPCs from the failed generation to the new + // generation's shared board and terminal-account group. + if len(runC.remoteFragmentCounts) > 0 { + runC.remoteExecutionID = newRemoteExecutionID() + } else { + runC.remoteExecutionID = uuid.Nil + } + exporter := func(snapshot mpool.AllocationAccountTerminalSnapshot) { + if resourceRecorder != nil { + resourceRecorder.recordAllocationAccountTerminal(snapshot) + } + } + err = runC.ensureAllocationAccountLifecycle(exporter) + if err == nil { + allocationAttempt, err = runC.beginAllocationAccountAttempt() + } + if err == nil { + err = runC.prePipelineInitializer() + } + if err == nil { preRunWall = carriedPreRunWall + time.Since(preRunOnceStart) attemptPreRunWall = preRunWall runC.MessageBoard.BeforeRunonce() @@ -315,6 +354,15 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { attemptPreRunWall = preRunWall coordinatorPhaseStart = time.Time{} coordinatorPhaseBase = 0 + if terminalErr := finishAllocationAttempt(); terminalErr != nil { + err = joinAllocationLifecycleErrors(err, terminalErr) + resourceRecorder.finishAttempt( + uint64(retryTimes), attemptStart, preRunWall, attemptRemoteWait, stats, + attemptScopes, attemptAnal, c.addr, false, + ) + attemptOpen = false + return nil, err + } c.fatalLog(retryTimes, err) if !c.canRetry(err) { @@ -450,6 +498,15 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { // outcome. The panic defer above remains the single terminal owner until // this call returns. c.AnalyzeExecPlan(runC, queryResult, stats, isExplainPhyPlan, option) + if terminalErr := finishAllocationAttempt(); terminalErr != nil { + err = joinAllocationLifecycleErrors(err, terminalErr) + resourceRecorder.finishAttempt( + uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, + attemptScopes, attemptAnal, c.addr, false, + ) + attemptOpen = false + return nil, err + } resourceRecorder.finishAttempt( uint64(retryTimes), attemptStart, attemptPreRunWall, attemptRemoteWait, stats, @@ -732,6 +789,7 @@ func (c *Compile) buildRetryCompile(defChanged bool) (*Compile, error) { var e error runC := NewCompile(c.addr, c.db, c.sql, c.tenant, c.uid, c.e, c.proc, c.stmt, c.isInternal, c.cnLabel, c.startAt) + c.copyAllocationAccountLifecycleTo(runC) runC.SetQuerySchedulingIntent(c.querySchedulingIntent) runC.SetSchedulingTraceRecorder(c.schedulingTrace) runC.SetOriginSQL(c.originSQL) diff --git a/pkg/sql/compile/compile_test.go b/pkg/sql/compile/compile_test.go index 2e06818e60645..9d5838ff3f186 100644 --- a/pkg/sql/compile/compile_test.go +++ b/pkg/sql/compile/compile_test.go @@ -136,7 +136,6 @@ func TestCompileRunPreservesBinaryPrepareParamAcrossRetries(t *testing.T) { require.Zero(t, params.Length()) require.Nil(t, params.GetData()) require.Nil(t, params.GetArea()) - c.Release() proc.Free() proc.GetSessionInfo().Buf.Free() diff --git a/pkg/sql/compile/operator.go b/pkg/sql/compile/operator.go index a525df21e6151..1a0791c4aa31a 100644 --- a/pkg/sql/compile/operator.go +++ b/pkg/sql/compile/operator.go @@ -27,7 +27,6 @@ import ( "github.com/gogo/protobuf/proto" "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/catalog" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/sqlquote" "github.com/matrixorigin/matrixone/pkg/container/batch" @@ -130,14 +129,14 @@ func mergeReceiverChannelBufferSize(s *Scope) int { type operatorDupContext struct { shufflePools map[*shuffle.Shuffle]*shuffle.ShufflePool - hashJoinChannels map[*hashjoin.HashJoin]chan *bitmap.Bitmap + hashJoinMailboxes map[*hashjoin.HashJoin]*hashjoin.BitmapMailbox dedupJoinMailboxes map[*dedupjoin.DedupJoin]*dedupjoin.WorkerJoinMailbox } func newOperatorDupContext() *operatorDupContext { return &operatorDupContext{ shufflePools: make(map[*shuffle.Shuffle]*shuffle.ShufflePool), - hashJoinChannels: make(map[*hashjoin.HashJoin]chan *bitmap.Bitmap), + hashJoinMailboxes: make(map[*hashjoin.HashJoin]*hashjoin.BitmapMailbox), dedupJoinMailboxes: make(map[*dedupjoin.DedupJoin]*dedupjoin.WorkerJoinMailbox), } } @@ -230,12 +229,12 @@ func dupOperatorWithContext(sourceOp vm.Operator, index int, maxParallel int, du op.CanSkipProbe = t.CanSkipProbe op.IsShuffle = t.IsShuffle if !t.IsShuffle { - channel := dupCtx.hashJoinChannels[t] - if channel == nil { - channel = make(chan *bitmap.Bitmap, maxParallel) - dupCtx.hashJoinChannels[t] = channel + mailbox := dupCtx.hashJoinMailboxes[t] + if mailbox == nil { + mailbox = hashjoin.NewBitmapMailbox(maxParallel) + dupCtx.hashJoinMailboxes[t] = mailbox } - op.Channel = channel + op.Mailbox = mailbox op.NumCPU = uint64(maxParallel) op.IsMerger = (index == 0) } diff --git a/pkg/sql/compile/operator_test.go b/pkg/sql/compile/operator_test.go index e2e19ef55291d..ba15db4ea7119 100644 --- a/pkg/sql/compile/operator_test.go +++ b/pkg/sql/compile/operator_test.go @@ -19,7 +19,6 @@ import ( "math" "testing" - "github.com/matrixorigin/matrixone/pkg/common/bitmap" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" @@ -599,21 +598,21 @@ func TestDupOperatorDedupJoinSharesMailboxOnlyWithinGeneration(t *testing.T) { require.NotSame(t, dup1.Mailbox, nextGeneration.Mailbox) } -func TestDupOperatorHashJoinSharesChannelOnlyWithinGeneration(t *testing.T) { +func TestDupOperatorHashJoinSharesMailboxOnlyWithinGeneration(t *testing.T) { op := hashjoin.NewArgument() - staleChannel := make(chan *bitmap.Bitmap, 2) - close(staleChannel) - op.Channel = staleChannel + staleMailbox := hashjoin.NewBitmapMailbox(2) + staleMailbox.SealAndDrain(mpool.MustNewZero()) + op.Mailbox = staleMailbox dupCtx := newOperatorDupContext() dup1 := dupOperatorWithContext(op, 0, 2, dupCtx).(*hashjoin.HashJoin) dup2 := dupOperatorWithContext(op, 1, 2, dupCtx).(*hashjoin.HashJoin) - require.Equal(t, staleChannel, op.Channel, "duplicating must not mutate the reusable template") - require.NotEqual(t, staleChannel, dup1.Channel, "a stale closed template channel must not enter a new execution") - require.Equal(t, dup1.Channel, dup2.Channel) + require.Same(t, staleMailbox, op.Mailbox, "duplicating must not mutate the reusable template") + require.NotSame(t, staleMailbox, dup1.Mailbox, "a stale template mailbox must not enter a new execution") + require.Same(t, dup1.Mailbox, dup2.Mailbox) nextGeneration := dupOperatorWithContext(op, 0, 2, newOperatorDupContext()).(*hashjoin.HashJoin) - require.NotEqual(t, dup1.Channel, nextGeneration.Channel) + require.NotSame(t, dup1.Mailbox, nextGeneration.Mailbox) } func TestDupOperatorAssignsSharedShuffleConsumerIndex(t *testing.T) { diff --git a/pkg/sql/compile/remote_allocation_statement_group.go b/pkg/sql/compile/remote_allocation_statement_group.go new file mode 100644 index 0000000000000..46a2e321d9765 --- /dev/null +++ b/pkg/sql/compile/remote_allocation_statement_group.go @@ -0,0 +1,550 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package compile + +import ( + "errors" + "sync" + "time" + + "github.com/google/uuid" + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/logutil" + "github.com/matrixorigin/matrixone/pkg/util/resource" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "go.uber.org/zap" +) + +func newRemoteExecutionID() uuid.UUID { + return uuid.New() +} + +func remoteMessageBoardID( + statementID uuid.UUID, + remoteExecutionID uuid.UUID, +) uuid.UUID { + if remoteExecutionID != uuid.Nil { + return remoteExecutionID + } + return statementID +} + +func remoteAllocationStatementGroupKey( + remoteExecutionID uuid.UUID, + address string, +) string { + if remoteExecutionID == uuid.Nil || address == "" { + return "" + } + return remoteExecutionID.String() + "@" + address +} + +// A missing planned RPC means the coordinator failed before dispatch +// completed. Bound the orphan lifetime by the same five-minute interval used +// by blocking MessageBoard receives. The timer is canceled as soon as every +// expected fragment has registered, so it never limits a fully dispatched +// statement's execution time. +var remoteAllocationStatementRegistrationTimeout = 5 * time.Minute + +// A late RPC can carry a new MessageBoard, so keep a bounded record of an +// incomplete execution after its active group has been released. A sender +// without a caller deadline can remain in flight for MaxRpcTime; expiring the +// key earlier would let that same physical generation reopen. The record +// contains no statement resources. +const remoteAllocationStatementTombstoneTimeout = MaxRpcTime + +// Reserve tombstone capacity when a remote execution is first admitted. An +// active group can become a tombstone after a partial dispatch, so bounding +// only the tombstone map would admit more generations than can later be +// retained safely. Existing fragments do not consume another reservation. +const remoteAllocationStatementGenerationLimit = 4096 + +// collectRemoteFragmentCounts computes the number of pipeline RPCs that the +// complete physical scope graph will send to each CN. The execution address +// changes when traversal crosses a Remote scope: nested scopes targeting that +// same address execute inside the received pipeline and do not create another +// RPC, while a different target does. +func collectRemoteFragmentCounts( + scopes []*Scope, + rootAddress string, +) map[string]uint32 { + counts := make(map[string]uint32) + var visit func(*Scope, string) + visit = func(scope *Scope, executionAddress string) { + if scope == nil { + return + } + if scope.Magic == Remote && !scope.ipAddrMatch(executionAddress) { + target := scope.NodeInfo.Addr + counts[target]++ + executionAddress = target + } + for _, pre := range scope.PreScopes { + visit(pre, executionAddress) + } + } + for _, scope := range scopes { + visit(scope, rootAddress) + } + return counts +} + +func validateRemoteAllocationTopologyCapability( + scopes []*Scope, + remoteFragmentCounts map[string]uint32, +) error { + if len(remoteFragmentCounts) > 0 { + return nil + } + owners, err := collectAllocationAccountOwners(scopes) + if err != nil { + return err + } + if !hasAllocationAccountActivator(owners) { + return nil + } + return moerr.NewNotSupportedNoCtx( + "remote allocation-accounted execution requires fragment topology metadata", + ) +} + +var remoteAllocationStatementGroups = struct { + sync.Mutex + byBoard map[*message.MessageBoard]*remoteAllocationStatementGroup + byKey map[string]*remoteAllocationStatementGroup + tombstones map[string]*remoteAllocationStatementTombstone +}{ + byBoard: make(map[*message.MessageBoard]*remoteAllocationStatementGroup), + byKey: make(map[string]*remoteAllocationStatementGroup), + tombstones: make(map[string]*remoteAllocationStatementTombstone), +} + +type remoteAllocationStatementTombstone struct { + timer *time.Timer +} + +// remoteAllocationStatementGroup is the terminal owner for all pipeline RPCs +// of one statement that execute on one CN. Those RPCs share a MessageBoard, so +// no individual fragment may close it or validate transferred allocations +// while a sibling can still consume them. +type remoteAllocationStatementGroup struct { + key string + board *message.MessageBoard + expected uint32 + registered uint32 + finished uint32 + attempts []*statementAllocationAttempt + pools []*mpool.MPool + participants []*remoteAllocationStatementParticipant + timer *time.Timer + expired bool + finalized bool + err error +} + +type remoteAllocationStatementParticipant struct { + group *remoteAllocationStatementGroup + cancel func(error) + finished bool + + stageOnce sync.Once + finishOnce sync.Once + terminal remoteAllocationStatementTerminal + err error +} + +type remoteAllocationStatementTerminal struct { + allocation []mpool.AllocationAccountTerminalSnapshot + memory resource.MemoryTotals + quality resource.QualityFlags + complete bool +} + +func acquireRemoteAllocationStatementParticipant( + key string, + board *message.MessageBoard, + expected uint32, + cancel func(error), +) (*remoteAllocationStatementParticipant, error) { + if key == "" || board == nil { + return nil, mpool.ErrAllocationAccountInvariant + } + if expected == 0 { + return nil, mpool.ErrAllocationAccountInvariant + } + + remoteAllocationStatementGroups.Lock() + defer remoteAllocationStatementGroups.Unlock() + if remoteAllocationStatementGroups.tombstones[key] != nil { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx("remote allocation statement group already aborted"), + ) + } + group := remoteAllocationStatementGroups.byBoard[board] + if group == nil { + if remoteAllocationStatementGroups.byKey[key] != nil { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx("remote allocation statement group key already registered"), + ) + } + if len(remoteAllocationStatementGroups.byKey)+ + len(remoteAllocationStatementGroups.tombstones) >= + remoteAllocationStatementGenerationLimit { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx( + "remote allocation statement generation capacity reached", + ), + ) + } + group = &remoteAllocationStatementGroup{ + key: key, + board: board, + expected: expected, + } + remoteAllocationStatementGroups.byBoard[board] = group + remoteAllocationStatementGroups.byKey[key] = group + } + if group.key != key || group.expected != expected || group.finalized || + group.expired || group.registered >= group.expected { + return nil, errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx("invalid remote allocation statement group registration"), + ) + } + participant := &remoteAllocationStatementParticipant{ + group: group, + cancel: cancel, + } + group.registered++ + group.participants = append(group.participants, participant) + if group.registered == group.expected { + if group.timer != nil { + group.timer.Stop() + group.timer = nil + } + } else if group.timer == nil { + group.timer = time.AfterFunc( + remoteAllocationStatementRegistrationTimeout, + func() { expireRemoteAllocationStatementGroup(group) }, + ) + } + return participant, nil +} + +// stage clears fragment-local operators while they are still reachable and +// retains the fragment MPool for a statement-boundary snapshot. Account and +// allocator completion are deferred until every expected fragment has +// finished. +func (p *remoteAllocationStatementParticipant) stage( + attempt *statementAllocationAttempt, + pool *mpool.MPool, +) { + if p == nil || p.group == nil { + return + } + p.stageOnce.Do(func() { + if attempt != nil { + attempt.exporter = nil + } + + remoteAllocationStatementGroups.Lock() + if p.group.finalized { + p.err = mpool.ErrAllocationAccountInvariant + remoteAllocationStatementGroups.Unlock() + return + } + if attempt != nil { + p.group.attempts = append(p.group.attempts, attempt) + } + if pool != nil { + p.group.pools = append(p.group.pools, pool) + } + remoteAllocationStatementGroups.Unlock() + + // Transfer terminal ownership before clearing operators. If a cleanup + // hook panics, the handler defer can still finish this participant and + // the group retains the account and allocator domain. + if attempt != nil { + _ = attempt.prepareTerminal(false) + } + }) +} + +// finish marks one remote fragment quiescent without imposing a cross-RPC +// response barrier. Nested remote paths can re-enter the same CN, so waiting +// for sibling completion here would create a B -> C -> B dependency cycle. +// The fragment that completes the group publishes the aggregate exactly once; +// earlier responses carry no duplicate terminal totals. +func (p *remoteAllocationStatementParticipant) finish(cause error) ( + remoteAllocationStatementTerminal, + error, +) { + if p == nil || p.group == nil { + return remoteAllocationStatementTerminal{}, nil + } + p.stage(nil, nil) + p.finishOnce.Do(func() { + remoteAllocationStatementGroups.Lock() + group := p.group + if group.finalized || + (!group.expired && group.finished >= group.expected) || + (group.expired && group.finished >= group.registered) { + p.err = joinAllocationLifecycleErrors( + p.err, + mpool.ErrAllocationAccountInvariant, + ) + remoteAllocationStatementGroups.Unlock() + return + } + abort := cause != nil && !group.expired + if cause != nil { + group.err = joinAllocationLifecycleErrors(group.err, cause) + group.expired = true + if group.timer != nil { + group.timer.Stop() + group.timer = nil + } + } + group.finished++ + p.finished = true + var cancels []func(error) + if abort { + cancels = activeRemoteAllocationStatementCancelsLocked(group) + } + if !group.expired && group.registered < group.expected && group.timer == nil { + group.timer = time.AfterFunc( + remoteAllocationStatementRegistrationTimeout, + func() { expireRemoteAllocationStatementGroup(group) }, + ) + } + complete := group.expired && group.finished == group.registered || + group.registered == group.expected && group.finished == group.expected + if !complete { + remoteAllocationStatementGroups.Unlock() + if abort { + abortErr := allocationLifecycleCall(func() error { + group.board.Close() + return nil + }) + abortErr = joinAllocationLifecycleErrors( + abortErr, + cancelRemoteAllocationStatementParticipants(cancels, cause), + ) + if abortErr != nil { + remoteAllocationStatementGroups.Lock() + group.err = joinAllocationLifecycleErrors(group.err, abortErr) + remoteAllocationStatementGroups.Unlock() + p.err = joinAllocationLifecycleErrors(p.err, abortErr) + } + } + } else { + attempts, pools := takeRemoteAllocationStatementGroupLocked(group) + terminalErr := group.err + remoteAllocationStatementGroups.Unlock() + defer releaseRemoteAllocationStatementGroup(group) + p.terminal, terminalErr = completeRemoteAllocationStatementGroup( + group, + attempts, + pools, + terminalErr, + ) + p.err = joinAllocationLifecycleErrors(p.err, terminalErr) + } + }) + return p.terminal, p.err +} + +func releaseRemoteAllocationStatementGroup(group *remoteAllocationStatementGroup) { + remoteAllocationStatementGroups.Lock() + if remoteAllocationStatementGroups.byBoard[group.board] == group { + delete(remoteAllocationStatementGroups.byBoard, group.board) + } + if remoteAllocationStatementGroups.byKey[group.key] == group { + delete(remoteAllocationStatementGroups.byKey, group.key) + } + remoteAllocationStatementGroups.Unlock() +} + +func cancelRemoteAllocationStatementParticipants( + cancels []func(error), + cause error, +) error { + var err error + for _, cancel := range cancels { + if cancel != nil { + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { + cancel(cause) + return nil + })) + } + } + return err +} + +func activeRemoteAllocationStatementCancelsLocked( + group *remoteAllocationStatementGroup, +) []func(error) { + cancels := make([]func(error), 0, len(group.participants)) + for _, participant := range group.participants { + if participant != nil && !participant.finished && participant.cancel != nil { + cancels = append(cancels, participant.cancel) + } + } + return cancels +} + +func takeRemoteAllocationStatementGroupLocked( + group *remoteAllocationStatementGroup, +) ([]*statementAllocationAttempt, []*mpool.MPool) { + group.finalized = true + if group.expired && group.registered < group.expected { + installRemoteAllocationStatementTombstoneLocked(group.key) + } + if group.timer != nil { + group.timer.Stop() + group.timer = nil + } + attempts := group.attempts + group.attempts = nil + pools := group.pools + group.pools = nil + return attempts, pools +} + +func installRemoteAllocationStatementTombstoneLocked(key string) { + if remoteAllocationStatementGroups.tombstones[key] != nil { + return + } + tombstone := &remoteAllocationStatementTombstone{} + remoteAllocationStatementGroups.tombstones[key] = tombstone + tombstone.timer = time.AfterFunc( + remoteAllocationStatementTombstoneTimeout, + func() { + remoteAllocationStatementGroups.Lock() + if remoteAllocationStatementGroups.tombstones[key] == tombstone { + delete(remoteAllocationStatementGroups.tombstones, key) + } + remoteAllocationStatementGroups.Unlock() + }, + ) +} + +func completeRemoteAllocationStatementGroup( + group *remoteAllocationStatementGroup, + attempts []*statementAllocationAttempt, + pools []*mpool.MPool, + terminalErr error, +) ( + remoteAllocationStatementTerminal, + error, +) { + terminalErr = joinAllocationLifecycleErrors( + terminalErr, + allocationLifecycleCall(func() error { + group.board.CloseAndDrain() + return nil + }), + ) + terminal := remoteAllocationStatementTerminal{ + complete: true, + allocation: make( + []mpool.AllocationAccountTerminalSnapshot, + 0, + len(attempts), + ), + } + for _, attempt := range attempts { + snapshot, err := attempt.completeTerminal() + terminal.allocation = append(terminal.allocation, snapshot) + terminalErr = joinAllocationLifecycleErrors(terminalErr, err) + } + for _, pool := range pools { + terminalErr = joinAllocationLifecycleErrors( + terminalErr, + allocationLifecycleCall(func() error { + domain, quality := pool.ResourceSnapshot() + terminal.quality |= quality | + resource.MergeMemoryDomain(&terminal.memory, domain) + return nil + }), + ) + } + return terminal, terminalErr +} + +func expireRemoteAllocationStatementGroup( + group *remoteAllocationStatementGroup, +) { + remoteAllocationStatementGroups.Lock() + if group.finalized || group.registered == group.expected { + remoteAllocationStatementGroups.Unlock() + return + } + group.expired = true + timeoutErr := moerr.NewInternalErrorNoCtx( + "remote allocation statement group registration timed out", + ) + group.err = joinAllocationLifecycleErrors(group.err, timeoutErr) + expected, registered, finished := group.expected, group.registered, group.finished + cancels := activeRemoteAllocationStatementCancelsLocked(group) + var attempts []*statementAllocationAttempt + var pools []*mpool.MPool + complete := group.finished == group.registered + if complete { + attempts, pools = takeRemoteAllocationStatementGroupLocked(group) + } + terminalErr := group.err + remoteAllocationStatementGroups.Unlock() + + if complete { + defer releaseRemoteAllocationStatementGroup(group) + _, terminalErr = completeRemoteAllocationStatementGroup( + group, + attempts, + pools, + terminalErr, + ) + } else { + terminalErr = joinAllocationLifecycleErrors( + terminalErr, + allocationLifecycleCall(func() error { + group.board.Close() + return nil + }), + ) + } + cancelErr := cancelRemoteAllocationStatementParticipants(cancels, timeoutErr) + terminalErr = joinAllocationLifecycleErrors(terminalErr, cancelErr) + if !complete && cancelErr != nil { + remoteAllocationStatementGroups.Lock() + if !group.finalized { + group.err = joinAllocationLifecycleErrors(group.err, cancelErr) + } + remoteAllocationStatementGroups.Unlock() + } + fields := []zap.Field{ + zap.Uint32("expected-fragments", expected), + zap.Uint32("registered-fragments", registered), + zap.Uint32("finished-fragments", finished), + } + if terminalErr != nil { + fields = append(fields, zap.Error(terminalErr)) + } + logutil.Warn("expired incomplete remote allocation statement group", fields...) +} diff --git a/pkg/sql/compile/remote_allocation_statement_group_test.go b/pkg/sql/compile/remote_allocation_statement_group_test.go new file mode 100644 index 0000000000000..f45d22c4af06e --- /dev/null +++ b/pkg/sql/compile/remote_allocation_statement_group_test.go @@ -0,0 +1,667 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package compile + +import ( + "context" + "errors" + "fmt" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/sql/colexec" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/connector" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/dispatch" + "github.com/matrixorigin/matrixone/pkg/vm/engine" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/stretchr/testify/require" +) + +type remoteAllocationAccountedMessage struct { + mp *mpool.MPool + buffer []byte + destroyed *atomic.Int32 +} + +func remoteAllocationStatementGroupRegistered(board *message.MessageBoard) bool { + remoteAllocationStatementGroups.Lock() + defer remoteAllocationStatementGroups.Unlock() + _, registered := remoteAllocationStatementGroups.byBoard[board] + return registered +} + +func acquireRemoteAllocationStatementTestParticipant( + t *testing.T, + board *message.MessageBoard, + expected uint32, + cancel func(error), +) (*remoteAllocationStatementParticipant, error) { + t.Helper() + key := fmt.Sprintf("test@%p", board) + t.Cleanup(func() { clearRemoteAllocationStatementTestTombstone(key) }) + return acquireRemoteAllocationStatementParticipant(key, board, expected, cancel) +} + +func clearRemoteAllocationStatementTestTombstone(key string) { + remoteAllocationStatementGroups.Lock() + if tombstone := remoteAllocationStatementGroups.tombstones[key]; tombstone != nil { + tombstone.timer.Stop() + delete(remoteAllocationStatementGroups.tombstones, key) + } + remoteAllocationStatementGroups.Unlock() +} + +func (m *remoteAllocationAccountedMessage) Serialize() []byte { return nil } + +func (m *remoteAllocationAccountedMessage) Deserialize([]byte) message.Message { + return m +} + +func (m *remoteAllocationAccountedMessage) NeedBlock() bool { return true } + +func (m *remoteAllocationAccountedMessage) GetMsgTag() int32 { return 1 } + +func (m *remoteAllocationAccountedMessage) GetReceiverAddr() message.MessageAddress { + return message.AddrBroadCastOnCurrentCN() +} + +func (m *remoteAllocationAccountedMessage) DebugString() string { + return "remote allocation-accounted message" +} + +func (m *remoteAllocationAccountedMessage) Destroy() { + if m.buffer != nil { + m.mp.Free(m.buffer) + m.buffer = nil + m.destroyed.Add(1) + } +} + +func TestCollectRemoteFragmentCountsCarriesExecutionAddress(t *testing.T) { + remoteC := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-c:6001"}, + } + sameB := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-b:6001"}, + PreScopes: []*Scope{remoteC}, + } + firstB := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-b:6001"}, + PreScopes: []*Scope{sameB}, + } + secondB := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-b:6001"}, + } + local := &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "cn-a:6001"}, + } + + require.Equal(t, map[string]uint32{ + "cn-b:6001": 2, + "cn-c:6001": 1, + }, collectRemoteFragmentCounts( + []*Scope{firstB, nil, secondB, local}, + "cn-a:6001", + )) +} + +func TestRemoteExecutionIDSeparatesRetryMessageBoards(t *testing.T) { + statementID := newRemoteExecutionID() + firstAttempt := newRemoteExecutionID() + secondAttempt := newRemoteExecutionID() + require.NotEqual(t, firstAttempt, secondAttempt) + require.Equal(t, statementID, remoteMessageBoardID(statementID, [16]byte{})) + require.Equal(t, firstAttempt, remoteMessageBoardID(statementID, firstAttempt)) + require.Equal(t, secondAttempt, remoteMessageBoardID(statementID, secondAttempt)) + require.NotEqual(t, + remoteAllocationStatementGroupKey(firstAttempt, "cn-a:6001"), + remoteAllocationStatementGroupKey(secondAttempt, "cn-a:6001"), + ) +} + +func TestRemoteAllocationTopologyCapabilityIsRequiredForOwners(t *testing.T) { + owner := &allocationLifecycleOwnerOperator{ + MockOperator: colexec.NewMockOperator(), + } + scopes := []*Scope{{RootOp: owner}} + require.Error(t, validateRemoteAllocationTopologyCapability(scopes, nil)) + require.NoError(t, validateRemoteAllocationTopologyCapability( + scopes, + map[string]uint32{"cn-a:6001": 1}, + )) + + connectorOp := connector.NewArgument() + dispatchOp := dispatch.NewArgument() + t.Cleanup(connectorOp.Release) + t.Cleanup(dispatchOp.Release) + require.NoError(t, validateRemoteAllocationTopologyCapability([]*Scope{ + {RootOp: connectorOp}, + {RootOp: dispatchOp}, + }, nil)) +} + +func TestRemoteAllocationStatementGroupDefersSharedBoardTerminal(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(2, 4) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + + buffer, err := producer.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, err) + var destroyed atomic.Int32 + message.SendMessage(&remoteAllocationAccountedMessage{ + mp: producer.proc.Mp(), + buffer: buffer, + destroyed: &destroyed, + }, board) + + first, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) + require.NoError(t, err) + first.stage(attempt, producer.proc.Mp()) + terminal, err := first.finish(nil) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + require.False(t, terminal.complete) + require.Zero(t, destroyed.Load()) + require.Equal(t, uint64(cap(buffer)), attempt.account.Snapshot().Used) + require.NotContains(t, board.DebugString(), "closed") + require.False(t, registry.AdmissionSuspended()) + + second, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) + require.NoError(t, err) + // The second fragment has no allocation owner. It still participates in + // the statement boundary and, as the last fragment, drains the producer's + // queued ownership before completing that producer's account. + terminal, err = second.finish(nil) + require.NoError(t, err) + require.True(t, terminal.complete) + require.Len(t, terminal.allocation, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) + require.Zero(t, terminal.allocation[0].Used) + require.Equal(t, uint64(cap(buffer)), terminal.allocation[0].Peak) + require.Equal(t, uint64(cap(buffer)), terminal.memory.AllocatedBytes) + require.Equal(t, uint64(cap(buffer)), terminal.memory.FreedBytes) + require.Zero(t, terminal.memory.LiveBytesAtSeal) + require.Zero(t, terminal.quality) + require.Equal(t, int32(1), destroyed.Load()) + require.Contains(t, board.DebugString(), "closed") + require.False(t, registry.AdmissionSuspended()) + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupPreservesStatementError(t *testing.T) { + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, + 1, + nil, + ) + require.NoError(t, err) + + primary := moerr.NewDuplicateEntryNoCtx("duplicate", "primary") + terminal, err := participant.finish(primary) + require.True(t, terminal.complete) + require.Same(t, primary, err) + require.Same(t, primary, moerr.ConvertGoError(context.Background(), err)) +} + +func TestRemoteAllocationStatementGroupRejectsTopologyMismatch(t *testing.T) { + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) + require.NoError(t, err) + _, err = acquireRemoteAllocationStatementTestParticipant(t, board, 3, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + + second, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) + require.NoError(t, err) + _, err = participant.finish(nil) + require.NoError(t, err) + _, err = second.finish(nil) + require.NoError(t, err) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupExpiresMissingFragment(t *testing.T) { + previousTimeout := remoteAllocationStatementRegistrationTimeout + remoteAllocationStatementRegistrationTimeout = 10 * time.Millisecond + t.Cleanup(func() { + remoteAllocationStatementRegistrationTimeout = previousTimeout + }) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("expired remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + buffer, err := producer.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, err) + var destroyed atomic.Int32 + message.SendMessage(&remoteAllocationAccountedMessage{ + mp: producer.proc.Mp(), + buffer: buffer, + destroyed: &destroyed, + }, board) + + participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) + require.NoError(t, err) + participant.stage(attempt, producer.proc.Mp()) + terminal, err := participant.finish(nil) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + require.False(t, terminal.complete) + require.Eventually(t, func() bool { + return destroyed.Load() == 1 && + registry.LiveAllocationMetadata() == 0 && + !remoteAllocationStatementGroupRegistered(board) + }, time.Second, time.Millisecond) + require.Contains(t, board.DebugString(), "closed") + require.False(t, registry.AdmissionSuspended()) +} + +func TestRemoteAllocationStatementRegistrationTimerStartsBeforeFinish(t *testing.T) { + previousTimeout := remoteAllocationStatementRegistrationTimeout + remoteAllocationStatementRegistrationTimeout = 10 * time.Millisecond + t.Cleanup(func() { + remoteAllocationStatementRegistrationTimeout = previousTimeout + }) + + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("expired remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + buffer, err := producer.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, err) + var destroyed atomic.Int32 + message.SendMessage(&remoteAllocationAccountedMessage{ + mp: producer.proc.Mp(), + buffer: buffer, + destroyed: &destroyed, + }, board) + + canceled := make(chan error, 1) + participant, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, + 2, + func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + + select { + case cause := <-canceled: + require.Error(t, cause) + case <-time.After(time.Second): + t.Fatal("registration timeout did not cancel the active fragment") + } + require.Eventually(t, func() bool { + return strings.Contains(board.DebugString(), "closed") + }, time.Second, time.Millisecond) + // Closing wakes the active fragment but cannot destroy ownership that the + // fragment may still be consuming. The last finish performs the drain. + require.Zero(t, destroyed.Load()) + require.Equal(t, uint64(cap(buffer)), attempt.account.Snapshot().Used) + + participant.stage(attempt, producer.proc.Mp()) + terminal, err := participant.finish(errors.New("active fragment observed cancellation")) + require.Error(t, err) + require.True(t, terminal.complete) + require.Len(t, terminal.allocation, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) + require.Zero(t, terminal.allocation[0].Used) + require.Equal(t, int32(1), destroyed.Load()) + require.Zero(t, terminal.memory.LiveBytesAtSeal) + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupFailureAbortsMissingFragment(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + board := message.NewMessageBoard() + producer := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("remote statement group must own terminal export") + }) + producer.MessageBoard = board + attempt, err := producer.beginAllocationAccountAttempt() + require.NoError(t, err) + participant, err := acquireRemoteAllocationStatementTestParticipant(t, board, 2, nil) + require.NoError(t, err) + participant.stage(attempt, producer.proc.Mp()) + + failure := errors.New("remote fragment failed") + terminal, err := participant.finish(failure) + require.ErrorIs(t, err, failure) + require.True(t, terminal.complete) + require.Len(t, terminal.allocation, 1) + require.Equal(t, mpool.AllocationAccountTerminalValid, terminal.allocation[0].State) + require.Contains(t, board.DebugString(), "closed") + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupRejectsLateFragmentAfterAbort(t *testing.T) { + key := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + t.Cleanup(func() { clearRemoteAllocationStatementTestTombstone(key) }) + + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementParticipant( + key, + board, + 2, + nil, + ) + require.NoError(t, err) + terminal, err := participant.finish(errors.New("first fragment aborted")) + require.Error(t, err) + require.True(t, terminal.complete) + require.False(t, remoteAllocationStatementGroupRegistered(board)) + + lateBoard := message.NewMessageBoard() + _, err = acquireRemoteAllocationStatementParticipant(key, lateBoard, 2, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + lateBoard.CloseAndDrain() + + retryKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + retryBoard := message.NewMessageBoard() + retry, err := acquireRemoteAllocationStatementParticipant( + retryKey, + retryBoard, + 1, + nil, + ) + require.NoError(t, err) + terminal, err = retry.finish(nil) + require.NoError(t, err) + require.True(t, terminal.complete) +} + +func TestRemoteAllocationStatementGroupExpiryRejectsLateFragment(t *testing.T) { + require.Equal(t, MaxRpcTime, remoteAllocationStatementTombstoneTimeout) + key := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + t.Cleanup(func() { clearRemoteAllocationStatementTestTombstone(key) }) + + board := message.NewMessageBoard() + participant, err := acquireRemoteAllocationStatementParticipant( + key, + board, + 2, + nil, + ) + require.NoError(t, err) + expireRemoteAllocationStatementGroup(participant.group) + terminal, err := participant.finish(nil) + require.Error(t, err) + require.True(t, terminal.complete) + require.False(t, remoteAllocationStatementGroupRegistered(board)) + + lateBoard := message.NewMessageBoard() + _, err = acquireRemoteAllocationStatementParticipant(key, lateBoard, 2, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + lateBoard.CloseAndDrain() + + retryKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + retryBoard := message.NewMessageBoard() + retry, err := acquireRemoteAllocationStatementParticipant( + retryKey, + retryBoard, + 1, + nil, + ) + require.NoError(t, err) + terminal, err = retry.finish(nil) + require.NoError(t, err) + require.True(t, terminal.complete) +} + +func TestRemoteAllocationStatementGenerationCapacity(t *testing.T) { + activeKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + activeBoard := message.NewMessageBoard() + first, err := acquireRemoteAllocationStatementParticipant( + activeKey, + activeBoard, + 3, + nil, + ) + require.NoError(t, err) + t.Cleanup(func() { + remoteAllocationStatementGroups.Lock() + if group := remoteAllocationStatementGroups.byKey[activeKey]; group != nil { + if group.timer != nil { + group.timer.Stop() + } + delete(remoteAllocationStatementGroups.byBoard, group.board) + delete(remoteAllocationStatementGroups.byKey, activeKey) + } + if tombstone := remoteAllocationStatementGroups.tombstones[activeKey]; tombstone != nil { + if tombstone.timer != nil { + tombstone.timer.Stop() + } + delete(remoteAllocationStatementGroups.tombstones, activeKey) + } + remoteAllocationStatementGroups.Unlock() + activeBoard.CloseAndDrain() + }) + + remoteAllocationStatementGroups.Lock() + for i := 1; i < remoteAllocationStatementGenerationLimit; i++ { + key := fmt.Sprintf("capacity-tombstone-%d", i) + remoteAllocationStatementGroups.tombstones[key] = + &remoteAllocationStatementTombstone{} + } + remoteAllocationStatementGroups.Unlock() + t.Cleanup(func() { + remoteAllocationStatementGroups.Lock() + for i := 1; i < remoteAllocationStatementGenerationLimit; i++ { + delete( + remoteAllocationStatementGroups.tombstones, + fmt.Sprintf("capacity-tombstone-%d", i), + ) + } + remoteAllocationStatementGroups.Unlock() + }) + + // The reservation belongs to the generation, not each fragment. + second, err := acquireRemoteAllocationStatementParticipant( + activeKey, + activeBoard, + 3, + nil, + ) + require.NoError(t, err) + + newKey := remoteAllocationStatementGroupKey(newRemoteExecutionID(), "cn-a:6001") + newBoard := message.NewMessageBoard() + _, err = acquireRemoteAllocationStatementParticipant(newKey, newBoard, 1, nil) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + newBoard.CloseAndDrain() + + terminal, err := first.finish(errors.New("partial dispatch failed")) + require.NoError(t, err) + require.False(t, terminal.complete) + terminal, err = second.finish(errors.New("sibling canceled")) + require.Error(t, err) + require.True(t, terminal.complete) + + remoteAllocationStatementGroups.Lock() + tombstoneCount := len(remoteAllocationStatementGroups.tombstones) + remoteAllocationStatementGroups.Unlock() + require.Equal(t, remoteAllocationStatementGenerationLimit, tombstoneCount) +} + +func TestRemoteAllocationStatementGroupFailureCancelsActiveSibling(t *testing.T) { + board := message.NewMessageBoard() + canceled := make(chan error, 2) + first, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, + 3, + func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + second, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, + 3, + func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + + failure := errors.New("first remote fragment failed") + terminal, err := first.finish(failure) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + require.False(t, terminal.complete) + // The failing participant is already quiescent; only its active sibling + // needs cancellation. + require.ErrorIs(t, <-canceled, failure) + select { + case unexpected := <-canceled: + t.Fatalf("finished participant was canceled: %v", unexpected) + default: + } + require.Contains(t, board.DebugString(), "closed") + + terminal, err = second.finish(errors.New("active sibling canceled")) + require.Error(t, err) + require.Empty(t, terminal.allocation) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} + +func TestRemoteAllocationStatementGroupExpirationWaitsForActiveFragment(t *testing.T) { + previousTimeout := remoteAllocationStatementRegistrationTimeout + remoteAllocationStatementRegistrationTimeout = 10 * time.Millisecond + t.Cleanup(func() { + remoteAllocationStatementRegistrationTimeout = previousTimeout + }) + + registry, err := mpool.NewAllocationAccountRegistry(2, 2) + require.NoError(t, err) + board := message.NewMessageBoard() + newAttempt := func() (*Compile, *statementAllocationAttempt, []byte) { + c := newTestAllocationLifecycleCompile(t, registry, func( + mpool.AllocationAccountTerminalSnapshot, + ) { + t.Fatal("remote statement group must own terminal export") + }) + c.MessageBoard = board + attempt, openErr := c.beginAllocationAccountAttempt() + require.NoError(t, openErr) + buffer, allocErr := c.proc.Mp().AllocAccounted( + 64, + attempt.account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + ) + require.NoError(t, allocErr) + return c, attempt, buffer + } + + firstCompile, firstAttempt, firstBuffer := newAttempt() + secondCompile, secondAttempt, secondBuffer := newAttempt() + canceled := make(chan error, 2) + first, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, 3, func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + second, err := acquireRemoteAllocationStatementTestParticipant( + t, + board, 3, func(cause error) { canceled <- cause }, + ) + require.NoError(t, err) + firstCompile.proc.Mp().Free(firstBuffer) + first.stage(firstAttempt, firstCompile.proc.Mp()) + terminal, err := first.finish(nil) + require.NoError(t, err) + require.Empty(t, terminal.allocation) + + require.Eventually(t, func() bool { + return strings.Contains(board.DebugString(), "closed") + }, time.Second, time.Millisecond) + require.Error(t, <-canceled) + select { + case unexpected := <-canceled: + t.Fatalf("finished participant was canceled: %v", unexpected) + default: + } + // Expiration closes the shared transport but cannot terminally inspect a + // registered fragment that is still executing. + _, live := registry.Resolve(secondAttempt.account.Handle()) + require.True(t, live) + require.Equal(t, uint64(cap(secondBuffer)), secondAttempt.account.Snapshot().Used) + require.False(t, secondAttempt.account.Snapshot().Sealed) + + secondCompile.proc.Mp().Free(secondBuffer) + second.stage(secondAttempt, secondCompile.proc.Mp()) + terminal, err = second.finish(errors.New("active fragment observed cancellation")) + require.Error(t, err) + require.Len(t, terminal.allocation, 2) + for _, snapshot := range terminal.allocation { + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + require.Zero(t, snapshot.Used) + } + require.Zero(t, terminal.memory.LiveBytesAtSeal) + require.Zero(t, terminal.quality) + require.Zero(t, registry.LiveAllocationMetadata()) + require.False(t, remoteAllocationStatementGroupRegistered(board)) +} diff --git a/pkg/sql/compile/remoterun.go b/pkg/sql/compile/remoterun.go index 525a9182bf07e..2863c8794f6b9 100644 --- a/pkg/sql/compile/remoterun.go +++ b/pkg/sql/compile/remoterun.go @@ -17,6 +17,7 @@ package compile import ( "context" "fmt" + "maps" "time" "unsafe" @@ -170,11 +171,17 @@ func decodeScope(data []byte, proc *process.Process, isRemote bool, eng engine.E func encodeProcessInfo( proc *process.Process, sql string, + remoteFragmentCounts map[string]uint32, + remoteExecutionID uuid.UUID, ) ([]byte, error) { v, err := proc.BuildProcessInfo(sql) if err != nil { return nil, err } + v.RemoteFragmentCounts = maps.Clone(remoteFragmentCounts) + if remoteExecutionID != uuid.Nil { + v.RemoteExecutionId = append([]byte(nil), remoteExecutionID[:]...) + } return v.Marshal() } diff --git a/pkg/sql/compile/remoterunClient.go b/pkg/sql/compile/remoterunClient.go index df39a870bac69..a7045658d91fc 100644 --- a/pkg/sql/compile/remoterunClient.go +++ b/pkg/sql/compile/remoterunClient.go @@ -23,6 +23,7 @@ import ( "sync" "time" + "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/cnservice/cnclient" "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/morpc" @@ -69,7 +70,13 @@ func (s *Scope) remoteRun(c *Compile) (sender *messageSenderOnClient, err error) // encode structures which need to send. var scopeEncodeData, processEncodeData []byte var withoutOutput, folded bool - scopeEncodeData, withoutOutput, processEncodeData, folded, err = prepareRemoteRunSendingData(c.sql, s, c.proc) + scopeEncodeData, withoutOutput, processEncodeData, folded, err = prepareRemoteRunSendingData( + c.sql, + s, + c.proc, + c.remoteFragmentCounts, + c.remoteExecutionID, + ) if err != nil { return nil, err } @@ -210,7 +217,13 @@ func findPipelineExternalLocalReceiver(s *Scope) vm.Operator { return nil } -func prepareRemoteRunSendingData(sqlStr string, s *Scope, proc *process.Process) (scopeData []byte, withoutOutput bool, processData []byte, folded bool, err error) { +func prepareRemoteRunSendingData( + sqlStr string, + s *Scope, + proc *process.Process, + remoteFragmentCounts map[string]uint32, + remoteExecutionID uuid.UUID, +) (scopeData []byte, withoutOutput bool, processData []byte, folded bool, err error) { encodedScope, withoutOutput := getScopeForRemoteRunEncoding(s) encodedScope, folded, err = foldVarExprsInRemoteRunScope(encodedScope, proc) if err != nil { @@ -223,7 +236,12 @@ func prepareRemoteRunSendingData(sqlStr string, s *Scope, proc *process.Process) } // Encode the Process related information. - if processData, err = encodeProcessInfo(s.Proc, sqlStr); err != nil { + if processData, err = encodeProcessInfo( + s.Proc, + sqlStr, + remoteFragmentCounts, + remoteExecutionID, + ); err != nil { return nil, false, nil, false, err } @@ -835,8 +853,11 @@ func (sender *messageSenderOnClient) dealRemoteTerminal(data []byte) error { sender.anal.appendRemoteResource( envelope.Delta, envelope.Memory, + envelope.Allocation, envelope.MissingFragmentCount, envelope.MissingMemoryDomainCount, + envelope.PendingAllocationGroups, + envelope.CompletedAllocationGroups, ) } sender.terminalSeen = true diff --git a/pkg/sql/compile/remoterunServer.go b/pkg/sql/compile/remoterunServer.go index 297a4657d40e1..680d31f7050b9 100644 --- a/pkg/sql/compile/remoterunServer.go +++ b/pkg/sql/compile/remoterunServer.go @@ -19,6 +19,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" "math" "sync" "time" @@ -51,6 +52,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/vm" "github.com/matrixorigin/matrixone/pkg/vm/engine" "github.com/matrixorigin/matrixone/pkg/vm/engine/disttae" + "github.com/matrixorigin/matrixone/pkg/vm/message" "github.com/matrixorigin/matrixone/pkg/vm/process" "go.uber.org/zap" ) @@ -188,7 +190,7 @@ func (receiver *messageReceiverOnServer) waitUntilDisconnectedOrCancelled() { } } -func handlePipelineMessage(receiver *messageReceiverOnServer) error { +func handlePipelineMessage(receiver *messageReceiverOnServer) (err error) { switch receiver.messageTyp { case pipeline.Method_PrepareDoneNotifyMessage: @@ -251,19 +253,108 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) error { if errBuildCompile != nil { return errBuildCompile } + var allocationAttempt *statementAllocationAttempt + var allocationParticipant *remoteAllocationStatementParticipant + var allocationGroupKey string + var localAllocation resource.AllocationAccountTotals + var localAllocationQuality resource.QualityFlags var runErr error + sharedMessageBoard := runCompile.MessageBoard + memoryPool := runCompile.proc.Mp() + statementGroupEnabled := len(runCompile.remoteFragmentCounts) > 0 + participantFinished := false + // This outer defer is the last-resort owner for a participant if the + // normal terminal path itself panics. Keep the Compile alive until after + // this guard runs: Release can make the pooled object reachable by a new + // RPC while terminal cleanup still needs its stable attempt references. defer func() { + recovered := recover() + defer runCompile.Release() + if allocationParticipant != nil && !participantFinished { + if allocationAttempt != nil { + runCompile.allocationAttempt = nil + } + allocationParticipant.stage(allocationAttempt, memoryPool) + cause := err + if recovered != nil { + cause = joinAllocationLifecycleErrors( + cause, + moerr.ConvertPanicError(receiver.messageCtx, recovered), + ) + } + _, terminalErr := allocationParticipant.finish(cause) + err = joinAllocationLifecycleErrors(err, terminalErr) + participantFinished = true + } + if recovered != nil { + panic(recovered) + } + }() + defer func() { + if recovered := recover(); recovered != nil { + err = joinAllocationLifecycleErrors( + err, + moerr.ConvertPanicError(receiver.messageCtx, recovered), + ) + } // Capture operator and descendant facts before cleanup. The MPool // snapshot intentionally follows Compile.clear so temporary execution // allocations are released before LiveBytesAtSeal is measured. The // descendant snapshot is already reduced under AnalyzeModule's mutex; // sender quiescence remains the lifecycle contract for this boundary. - localDelta := collectScopeResourceDelta(runCompile.scopes, receiver.cnInformation.cnAddr) - descendant := runCompile.anal.remoteResourceSummary() - expectedDirect := countExpectedRemoteScopes(runCompile.scopes, receiver.cnInformation.cnAddr) - memoryPool := runCompile.proc.Mp() - runCompile.clear() - localMemory, localMemoryQuality := memoryPool.ResourceSnapshot() + var localDelta resource.Delta + var descendant remoteResourceSnapshot + var expectedDirect uint64 + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { + localDelta = collectScopeResourceDelta( + runCompile.scopes, + receiver.cnInformation.cnAddr, + ) + descendant = runCompile.anal.remoteResourceSummary() + expectedDirect = countExpectedRemoteScopes( + runCompile.scopes, + receiver.cnInformation.cnAddr, + ) + return nil + })) + var terminal remoteAllocationStatementTerminal + if allocationAttempt != nil { + runCompile.allocationAttempt = nil + } + allocationParticipant.stage(allocationAttempt, memoryPool) + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { + if statementGroupEnabled { + // The remote statement group, rather than any one fragment Compile, + // owns the shared multi-CN board. Detach it before clear resets the + // fragment-local Compile state. A fragment rejected before it joins a + // group is terminal for the statement, so it closes its unowned board. + if allocationParticipant == nil { + sharedMessageBoard.CloseAndDrain() + } + runCompile.MessageBoard = message.NewMessageBoard() + runCompile.proc.SetMessageBoard(runCompile.MessageBoard) + } + runCompile.clear() + return nil + })) + terminal, terminalErr := allocationParticipant.finish(err) + participantFinished = true + err = joinAllocationLifecycleErrors(err, terminalErr) + for _, snapshot := range terminal.allocation { + localAllocationQuality |= localAllocation.AddGeneration( + snapshot.Peak, + snapshot.Used, + snapshot.State == mpool.AllocationAccountTerminalValid, + ) + } + var localMemory resource.MemoryDomainSummary + var localMemoryQuality resource.QualityFlags + if !statementGroupEnabled { + err = joinAllocationLifecycleErrors(err, allocationLifecycleCall(func() error { + localMemory, localMemoryQuality = memoryPool.ResourceSnapshot() + return nil + })) + } aggregate := composeRemoteResourceAggregate( localDelta, localMemory, @@ -271,13 +362,60 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) error { descendant, expectedDirect, ) + aggregate.Delta.Quality |= localAllocationQuality | + resource.MergeAllocationAccountTotals( + &aggregate.Allocation, + localAllocation, + ) + aggregate.Delta.Quality |= terminal.quality | + resource.MergeMemoryTotals(&aggregate.Memory, terminal.memory) + if allocationParticipant != nil { + addRemoteAllocationGroupSignal( + &aggregate, + allocationGroupKey, + terminal.complete, + ) + } receiver.resourceDelta = aggregate.Delta receiver.resourceMemory = aggregate.Memory + receiver.resourceAllocation = aggregate.Allocation receiver.resourceMissingFragments = aggregate.MissingFragmentCount receiver.resourceMissingMemoryDomains = aggregate.MissingMemoryDomainCount - - runCompile.Release() + receiver.resourcePendingAllocationGroups = aggregate.PendingAllocationGroups + receiver.resourceCompletedAllocationGroups = aggregate.CompletedAllocationGroups }() + if statementGroupEnabled { + expectedFragments, ok := runCompile.remoteFragmentCounts[runCompile.addr] + if !ok || expectedFragments == 0 { + return errors.Join( + mpool.ErrAllocationAccountInvariant, + moerr.NewInternalErrorNoCtx( + "remote fragment topology has no local CN entry", + ), + ) + } + // Capture the function value for this execution generation. Both the + // Compile and Process fields are reset after an early sibling returns; + // looking up proc.Cancel later could cancel an unrelated generation. + remoteCancel := runCompile.proc.Cancel + allocationGroupKey = remoteAllocationStatementGroupKey( + runCompile.remoteExecutionID, + runCompile.addr, + ) + allocationParticipant, runErr = acquireRemoteAllocationStatementParticipant( + allocationGroupKey, + runCompile.MessageBoard, + expectedFragments, + func(cause error) { + if remoteCancel != nil { + remoteCancel(cause) + } + }, + ) + if runErr != nil { + return runErr + } + } // decode and running the pipeline. s, runErr := decodeScope(receiver.scopeData, runCompile.proc, true, runCompile.e) @@ -301,6 +439,23 @@ func handlePipelineMessage(receiver *messageReceiverOnServer) error { } runCompile.scopes = []*Scope{s} + if runErr = validateRemoteAllocationTopologyCapability( + runCompile.scopes, + runCompile.remoteFragmentCounts, + ); runErr != nil { + return runErr + } + if statementGroupEnabled { + runErr = runCompile.ensureAllocationAccountLifecycle( + func(mpool.AllocationAccountTerminalSnapshot) {}, + ) + if runErr == nil { + allocationAttempt, runErr = runCompile.beginAllocationAccountAttempt() + } + } + if runErr != nil { + return runErr + } runCompile.InitPipelineContextToExecuteQuery() normalizeRemoteDispatchReceiverAddresses(s, runCompile.addr) @@ -525,9 +680,11 @@ type processHelper struct { txnClient client.TxnClient sessionInfo process.SessionInfo //analysisNodeList []int32 - StmtId uuid.UUID - prepareParams pipeline.PrepareParamInfo - affectedRows int64 + StmtId uuid.UUID + prepareParams pipeline.PrepareParamInfo + affectedRows int64 + remoteFragmentCounts map[string]uint32 + remoteExecutionID uuid.UUID } // messageReceiverOnServer supported a series methods to write back results. @@ -556,11 +713,14 @@ type messageReceiverOnServer struct { colexecServer *colexec.Server // result. - phyPlan *models.PhyPlan - resourceDelta resource.Delta - resourceMemory resource.MemoryTotals - resourceMissingFragments uint64 - resourceMissingMemoryDomains uint64 + phyPlan *models.PhyPlan + resourceDelta resource.Delta + resourceMemory resource.MemoryTotals + resourceAllocation resource.AllocationAccountTotals + resourceMissingFragments uint64 + resourceMissingMemoryDomains uint64 + resourcePendingAllocationGroups []remoteAllocationGroupPending + resourceCompletedAllocationGroups []string } func newMessageReceiverOnServer( @@ -715,10 +875,16 @@ func (receiver *messageReceiverOnServer) newCompile() (*Compile, error) { c := allocateNewCompile(proc) c.execType = plan2.ExecTypeAP_MULTICN c.e = cnInfo.storeEngine - c.MessageBoard = c.MessageBoard.SetMultiCN(c.GetMessageCenter(), c.proc.GetStmtProfile().GetStmtId()) + messageBoardID := remoteMessageBoardID( + c.proc.GetStmtProfile().GetStmtId(), + pHelper.remoteExecutionID, + ) + c.MessageBoard = c.MessageBoard.SetMultiCN(c.GetMessageCenter(), messageBoardID) c.proc.SetMessageBoard(c.MessageBoard) c.anal = newAnalyzeModule() c.addr = receiver.cnInformation.cnAddr + c.remoteFragmentCounts = maps.Clone(pHelper.remoteFragmentCounts) + c.remoteExecutionID = pHelper.remoteExecutionID // a method to send back. c.execType = plan2.ExecTypeAP_MULTICN @@ -858,11 +1024,14 @@ func (receiver *messageReceiverOnServer) sendEndMessage() error { func (receiver *messageReceiverOnServer) setTerminalAnalysis(message *pipeline.Message) error { envelope := remoteTerminalEnvelope{ - TerminalResourceVersion: remoteTerminalResourceVersion, - Delta: receiver.resourceDelta, - Memory: receiver.resourceMemory, - MissingFragmentCount: receiver.resourceMissingFragments, - MissingMemoryDomainCount: receiver.resourceMissingMemoryDomains, + TerminalResourceVersion: remoteTerminalResourceVersion, + Delta: receiver.resourceDelta, + Memory: receiver.resourceMemory, + Allocation: receiver.resourceAllocation, + MissingFragmentCount: receiver.resourceMissingFragments, + MissingMemoryDomainCount: receiver.resourceMissingMemoryDomains, + PendingAllocationGroups: receiver.resourcePendingAllocationGroups, + CompletedAllocationGroups: receiver.resourceCompletedAllocationGroups, } if receiver.phyPlan != nil { envelope.PhyPlan = *receiver.phyPlan @@ -883,12 +1052,24 @@ func generateProcessHelper(ctx context.Context, data []byte, cli client.TxnClien } result := processHelper{ - id: procInfo.Id, - lim: process.ConvertToProcessLimitation(procInfo.Lim), - unixTime: procInfo.UnixTime, - accountId: procInfo.AccountId, - txnClient: cli, - affectedRows: procInfo.AffectedRows, + id: procInfo.Id, + lim: process.ConvertToProcessLimitation(procInfo.Lim), + unixTime: procInfo.UnixTime, + accountId: procInfo.AccountId, + txnClient: cli, + affectedRows: procInfo.AffectedRows, + remoteFragmentCounts: maps.Clone(procInfo.RemoteFragmentCounts), + } + if len(procInfo.RemoteExecutionId) > 0 { + result.remoteExecutionID, err = uuid.FromBytes(procInfo.RemoteExecutionId) + if err != nil { + return processHelper{}, err + } + } + if (len(result.remoteFragmentCounts) == 0) != (result.remoteExecutionID == uuid.Nil) { + return processHelper{}, moerr.NewInternalErrorNoCtx( + "incomplete remote allocation lifecycle metadata", + ) } result.txnOperator, err = cli.NewWithSnapshot(ctx, procInfo.Snapshot) if err != nil { diff --git a/pkg/sql/compile/remoterun_test.go b/pkg/sql/compile/remoterun_test.go index 37b083730fcc8..5fbf1ed346da9 100644 --- a/pkg/sql/compile/remoterun_test.go +++ b/pkg/sql/compile/remoterun_test.go @@ -145,8 +145,41 @@ func Test_EncodeProcessInfo(t *testing.T) { SqlHelper: nil, } - _, err := encodeProcessInfo(proc, "") + remoteExecutionID := uuid.New() + data, err := encodeProcessInfo(proc, "", map[string]uint32{ + "cn-a:6001": 2, + "cn-b:6001": 1, + }, remoteExecutionID) require.Nil(t, err) + restored := new(pipeline.ProcessInfo) + require.NoError(t, restored.Unmarshal(data)) + require.Equal(t, map[string]uint32{ + "cn-a:6001": 2, + "cn-b:6001": 1, + }, restored.RemoteFragmentCounts) + restoredExecutionID, err := uuid.FromBytes(restored.RemoteExecutionId) + require.NoError(t, err) + require.Equal(t, remoteExecutionID, restoredExecutionID) +} + +func TestGenerateProcessHelperRejectsIncompleteRemoteLifecycleMetadata(t *testing.T) { + tests := []pipeline.ProcessInfo{ + {RemoteFragmentCounts: map[string]uint32{"cn-a:6001": 1}}, + {RemoteExecutionId: func() []byte { + id := uuid.New() + return id[:] + }()}, + { + RemoteFragmentCounts: map[string]uint32{"cn-a:6001": 1}, + RemoteExecutionId: []byte{1}, + }, + } + for i := range tests { + data, err := tests[i].Marshal() + require.NoError(t, err) + _, err = generateProcessHelper(context.Background(), data, nil) + require.Error(t, err) + } } func Test_refactorScope(t *testing.T) { @@ -2758,7 +2791,7 @@ func Test_prepareRemoteRunSendingData(t *testing.T) { Proc: proc, RootOp: connector.NewArgument(), } - _, withoutOut, _, _, err := prepareRemoteRunSendingData("", s1, proc) + _, withoutOut, _, _, err := prepareRemoteRunSendingData("", s1, proc, nil, uuid.Nil) require.NoError(t, err) require.False(t, withoutOut) require.NotNil(t, s1.RootOp) @@ -2772,7 +2805,7 @@ func Test_prepareRemoteRunSendingData(t *testing.T) { } s2.RootOp.AppendChild(value_scan.NewArgument()) originChild := s2.RootOp.GetOperatorBase().GetChildren(0) - _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s2, proc) + _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s2, proc, nil, uuid.Nil) require.NoError(t, err) require.False(t, withoutOut) require.Equal(t, 1, s2.RootOp.GetOperatorBase().NumChildren()) @@ -2785,7 +2818,7 @@ func Test_prepareRemoteRunSendingData(t *testing.T) { RootOp: value_scan.NewArgument(), } s3.RootOp.AppendChild(value_scan.NewArgument()) - _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s3, proc) + _, withoutOut, _, _, err = prepareRemoteRunSendingData("", s3, proc, nil, uuid.Nil) require.NoError(t, err) require.True(t, withoutOut) } @@ -2813,7 +2846,7 @@ func TestPrepareRemoteRunSendingDataKeepsConnectorChildTableFunctionParams(t *te RootOp: conn, } - scopeData, withoutOut, _, _, err := prepareRemoteRunSendingData("", s, proc) + scopeData, withoutOut, _, _, err := prepareRemoteRunSendingData("", s, proc, nil, uuid.Nil) require.NoError(t, err) require.False(t, withoutOut) diff --git a/pkg/sql/compile/resource_accounting.go b/pkg/sql/compile/resource_accounting.go index e0cf9d1fa61dc..b921571570135 100644 --- a/pkg/sql/compile/resource_accounting.go +++ b/pkg/sql/compile/resource_accounting.go @@ -17,9 +17,12 @@ package compile import ( "context" "math" + "slices" + "strings" "sync/atomic" "time" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/sql/models" "github.com/matrixorigin/matrixone/pkg/util/resource" "github.com/matrixorigin/matrixone/pkg/util/trace/impl/motrace/statistic" @@ -27,14 +30,16 @@ import ( ) type executionResourceRecorder struct { - root *resource.Root - stats *statistic.StatsInfo - execution resource.ExecutionSummary - published bool - ownsAttempts bool + root *resource.Root + stats *statistic.StatsInfo + execution resource.ExecutionSummary + published bool + ownsAttempts bool + pendingAllocation resource.AllocationAccountTotals + pendingAllocationQuality resource.QualityFlags } -const remoteTerminalResourceVersion = 1 +const remoteTerminalResourceVersion = 3 // remoteTerminalEnvelope keeps PhyPlan fields at the top level so clients from // before resource accounting can still decode the terminal plan during a @@ -42,20 +47,44 @@ const remoteTerminalResourceVersion = 1 // appended resource facts from a legacy bare PhyPlan payload. type remoteTerminalEnvelope struct { models.PhyPlan - TerminalResourceVersion uint32 `json:"terminal_resource_version,omitempty"` - Delta resource.Delta `json:"resource_delta"` - Memory resource.MemoryTotals `json:"memory"` - MissingFragmentCount uint64 `json:"missing_fragment_count,omitempty"` - MissingMemoryDomainCount uint64 `json:"missing_memory_domain_count,omitempty"` + TerminalResourceVersion uint32 `json:"terminal_resource_version,omitempty"` + Delta resource.Delta `json:"resource_delta"` + Memory resource.MemoryTotals `json:"memory"` + Allocation resource.AllocationAccountTotals `json:"allocation_account"` + MissingFragmentCount uint64 `json:"missing_fragment_count,omitempty"` + MissingMemoryDomainCount uint64 `json:"missing_memory_domain_count,omitempty"` + PendingAllocationGroups []remoteAllocationGroupPending `json:"pending_allocation_groups,omitempty"` + CompletedAllocationGroups []string `json:"completed_allocation_groups,omitempty"` +} + +type remoteAllocationGroupPending struct { + Key string `json:"key"` + Count uint64 `json:"count"` } // remoteResourceAggregate is the already-reduced terminal output sent by one // remote hop. Delta contains local plus descendant usage and quality. type remoteResourceAggregate struct { - Delta resource.Delta - Memory resource.MemoryTotals - MissingFragmentCount uint64 - MissingMemoryDomainCount uint64 + Delta resource.Delta + Memory resource.MemoryTotals + Allocation resource.AllocationAccountTotals + MissingFragmentCount uint64 + MissingMemoryDomainCount uint64 + PendingAllocationGroups []remoteAllocationGroupPending + CompletedAllocationGroups []string +} + +func (r *executionResourceRecorder) recordAllocationAccountTerminal( + snapshot mpool.AllocationAccountTerminalSnapshot, +) { + if r == nil { + return + } + r.pendingAllocationQuality |= r.pendingAllocation.AddGeneration( + snapshot.Peak, + snapshot.Used, + snapshot.State == mpool.AllocationAccountTerminalValid, + ) } func newExecutionResourceRecorder( @@ -117,6 +146,24 @@ func (r *executionResourceRecorder) finishAttempt( remote, countExpectedRemoteScopes(scopes, localAddress), ) + var pending uint64 + for _, signal := range remoteAggregate.PendingAllocationGroups { + pending, remoteAggregate.Delta.Quality = addCheckedRemoteCounter( + pending, + signal.Count, + remoteAggregate.Delta.Quality, + ) + } + if pending > 0 { + remoteAggregate.MissingMemoryDomainCount, remoteAggregate.Delta.Quality = + addCheckedRemoteCounter( + remoteAggregate.MissingMemoryDomainCount, + pending, + remoteAggregate.Delta.Quality, + ) + remoteAggregate.Delta.Quality |= resource.QualityPartial | + resource.QualityMissingMemoryDomain + } delta = remoteAggregate.Delta var coordinator resource.LocalRecorder @@ -170,6 +217,17 @@ func (r *executionResourceRecorder) finishAttempt( summary := resource.AttemptSummary{WallNS: wallNS} summary.Quality |= delta.Quality | resource.MergeUsage(&summary.Usage, delta.Usage) summary.Quality |= resource.MergeMemoryTotals(&summary.Memory, remoteAggregate.Memory) + summary.Quality |= resource.MergeAllocationAccountTotals( + &summary.Allocation, + remoteAggregate.Allocation, + ) + summary.Quality |= r.pendingAllocationQuality | + resource.MergeAllocationAccountTotals( + &summary.Allocation, + r.pendingAllocation, + ) + r.pendingAllocation = resource.AllocationAccountTotals{} + r.pendingAllocationQuality = 0 var quality resource.QualityFlags summary.MissingFragmentCount, quality = addCheckedRemoteCounter( summary.MissingFragmentCount, remoteAggregate.MissingFragmentCount, summary.Quality) @@ -221,6 +279,80 @@ func addCheckedRemoteCounter(value, add uint64, quality resource.QualityFlags) ( return value + add, quality } +func reduceRemoteAllocationGroupSignals( + pending []remoteAllocationGroupPending, + completed []string, +) ([]remoteAllocationGroupPending, []string, resource.QualityFlags) { + var quality resource.QualityFlags + completedSet := make(map[string]struct{}, len(completed)) + for _, key := range completed { + if key == "" { + quality |= resource.QualityInvariantFailure + continue + } + completedSet[key] = struct{}{} + } + pendingCounts := make(map[string]uint64, len(pending)) + for _, signal := range pending { + if signal.Key == "" || signal.Count == 0 { + quality |= resource.QualityInvariantFailure + continue + } + if _, resolved := completedSet[signal.Key]; !resolved { + pendingCounts[signal.Key], quality = addCheckedRemoteCounter( + pendingCounts[signal.Key], + signal.Count, + quality, + ) + } + } + pending = pending[:0] + for key, count := range pendingCounts { + pending = append(pending, remoteAllocationGroupPending{ + Key: key, + Count: count, + }) + } + completed = completed[:0] + for key := range completedSet { + completed = append(completed, key) + } + slices.SortFunc(pending, func(a, b remoteAllocationGroupPending) int { + return strings.Compare(a.Key, b.Key) + }) + slices.Sort(completed) + return pending, completed, quality +} + +func addRemoteAllocationGroupSignal( + aggregate *remoteResourceAggregate, + key string, + completed bool, +) { + if aggregate == nil { + return + } + if completed { + aggregate.CompletedAllocationGroups = append( + aggregate.CompletedAllocationGroups, + key, + ) + } else { + aggregate.PendingAllocationGroups = append( + aggregate.PendingAllocationGroups, + remoteAllocationGroupPending{Key: key, Count: 1}, + ) + } + var quality resource.QualityFlags + aggregate.PendingAllocationGroups, + aggregate.CompletedAllocationGroups, + quality = reduceRemoteAllocationGroupSignals( + aggregate.PendingAllocationGroups, + aggregate.CompletedAllocationGroups, + ) + aggregate.Delta.Quality |= quality +} + // composeRemoteResourceAggregate composes one hop's captured local resource // facts with an already-reduced descendant aggregate. It is pure so every // remote hop and the coordinator use the same merge algebra. @@ -236,14 +368,33 @@ func composeRemoteResourceAggregate( result.Delta.Quality |= resource.MergeUsage(&result.Delta.Usage, descendant.Usage) result.Delta.Quality |= resource.MergeMemoryDomain(&result.Memory, localMemory) result.Delta.Quality |= resource.MergeMemoryTotals(&result.Memory, descendant.Memory) + result.Delta.Quality |= resource.MergeAllocationAccountTotals( + &result.Allocation, + descendant.Allocation, + ) result.MissingFragmentCount = descendant.MissingFragmentCount result.MissingMemoryDomainCount = descendant.MissingMemoryDomainCount + result.PendingAllocationGroups = append( + []remoteAllocationGroupPending(nil), + descendant.PendingAllocationGroups..., + ) + result.CompletedAllocationGroups = append( + []string(nil), descendant.CompletedAllocationGroups..., + ) if descendant.MissingFragmentCount > 0 { result.Delta.Quality |= resource.QualityPartial | resource.QualityMissingFragment } if descendant.MissingMemoryDomainCount > 0 { result.Delta.Quality |= resource.QualityPartial | resource.QualityMissingMemoryDomain } + var groupQuality resource.QualityFlags + result.PendingAllocationGroups, + result.CompletedAllocationGroups, + groupQuality = reduceRemoteAllocationGroupSignals( + result.PendingAllocationGroups, + result.CompletedAllocationGroups, + ) + result.Delta.Quality |= groupQuality if descendant.DirectReportCount < expectedDirect { directMissing := expectedDirect - descendant.DirectReportCount result.MissingFragmentCount, result.Delta.Quality = addCheckedRemoteCounter( diff --git a/pkg/sql/compile/resource_accounting_test.go b/pkg/sql/compile/resource_accounting_test.go index 216afa990782d..ff60d30fa8dee 100644 --- a/pkg/sql/compile/resource_accounting_test.go +++ b/pkg/sql/compile/resource_accounting_test.go @@ -22,6 +22,7 @@ import ( "testing" "time" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/sql/colexec/value_scan" "github.com/matrixorigin/matrixone/pkg/sql/models" "github.com/matrixorigin/matrixone/pkg/testutil" @@ -51,8 +52,11 @@ func TestExecutionResourceRecorder(t *testing.T) { MaxDomainPeakLiveBytes: 80, SumDomainPeakLiveBytesBound: 80, }, + resource.AllocationAccountTotals{}, 0, 0, + nil, + nil, ) recorder.finishAttempt( 0, @@ -75,6 +79,42 @@ func TestExecutionResourceRecorder(t *testing.T) { require.Zero(t, summary.Quality&resource.QualityMissingFragment) } +func TestExecutionResourceRecorderPublishesAllocationTerminal(t *testing.T) { + root := resource.NewRoot(resource.ConnExternal) + recorder := newExecutionResourceRecorder( + resource.ContextWithRoot(context.Background(), root), + true, + ) + require.NotNil(t, recorder) + recorder.recordAllocationAccountTerminal( + mpool.AllocationAccountTerminalSnapshot{ + AllocationAccountSnapshot: mpool.AllocationAccountSnapshot{ + Peak: 64, + }, + State: mpool.AllocationAccountTerminalValid, + }, + ) + recorder.finishAttempt( + 0, + time.Now(), + 0, + 0, + nil, + nil, + nil, + "", + false, + ) + recorder.publish() + + summary := root.PreResponseSummary() + require.Equal(t, uint64(1), summary.Allocation.GenerationCount) + require.Equal(t, uint64(1), summary.Allocation.ValidGenerationCount) + require.Equal(t, uint64(64), summary.Allocation.MaxGenerationPeak) + require.Zero(t, summary.Allocation.LiveBytesAtTerminal) + require.Zero(t, summary.Quality&resource.QualityInvariantFailure) +} + func TestExplainPhyBufferUsesPublishedCurrentAttempt(t *testing.T) { root := resource.NewRoot(resource.ConnExternal) require.True(t, root.MergeExecution(resource.ExecutionSummary{ @@ -258,6 +298,17 @@ func TestRemoteTerminalEnvelope(t *testing.T) { MaxDomainPeakLiveBytes: 15, SumDomainPeakLiveBytesBound: 15, }, + Allocation: resource.AllocationAccountTotals{ + GenerationCount: 1, + ValidGenerationCount: 1, + MaxGenerationPeak: 17, + SumGenerationPeak: 17, + }, + PendingAllocationGroups: []remoteAllocationGroupPending{{ + Key: "pending@cn", + Count: 2, + }}, + CompletedAllocationGroups: []string{"completed@cn"}, } data, err := json.Marshal(envelope) require.NoError(t, err) @@ -267,8 +318,15 @@ func TestRemoteTerminalEnvelope(t *testing.T) { summary := anal.remoteResourceSummary() require.Equal(t, uint64(1), summary.DirectReportCount) require.Equal(t, uint64(11), summary.Usage.ExclusiveActiveNS) + require.Equal(t, []remoteAllocationGroupPending{{ + Key: "pending@cn", + Count: 2, + }}, summary.PendingAllocationGroups) + require.Equal(t, []string{"completed@cn"}, summary.CompletedAllocationGroups) require.Equal(t, uint64(12), summary.Usage.S3ReadBytes) require.Equal(t, uint64(15), summary.Memory.MaxDomainPeakLiveBytes) + require.Equal(t, uint64(1), summary.Allocation.GenerationCount) + require.Equal(t, uint64(17), summary.Allocation.MaxGenerationPeak) require.NotZero(t, summary.Quality&resource.QualityPartial) require.Len(t, anal.remotePhyPlans, 1) require.Equal(t, "Merge", anal.remotePhyPlans[0].LocalScope[0].Magic) @@ -444,7 +502,15 @@ func TestRemoteResourceCounterSaturates(t *testing.T) { remoteMissingMemoryDomains: math.MaxUint64, remoteReports: math.MaxUint64, } - anal.appendRemoteResource(resource.Delta{}, resource.MemoryTotals{}, 1, 1) + anal.appendRemoteResource( + resource.Delta{}, + resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, + 1, + 1, + nil, + nil, + ) snapshot := anal.remoteResourceSummary() require.Equal(t, uint64(math.MaxUint64), snapshot.MissingFragmentCount) require.Equal(t, uint64(math.MaxUint64), snapshot.MissingMemoryDomainCount) @@ -466,12 +532,123 @@ func TestAnalyzeModuleResetClearsRemoteResourceAggregate(t *testing.T) { anal := &AnalyzeModule{} anal.appendRemoteResource( resource.Delta{Usage: resource.Usage{S3ReadBytes: 11}, Quality: resource.QualityPartial}, - resource.MemoryTotals{MaxDomainPeakLiveBytes: 8}, 2, 3) + resource.MemoryTotals{MaxDomainPeakLiveBytes: 8}, + resource.AllocationAccountTotals{}, + 2, + 3, + []remoteAllocationGroupPending{{Key: "pending", Count: 1}}, + []string{"completed"}, + ) anal.Reset(false, false) snapshot := anal.remoteResourceSummary() require.Equal(t, remoteResourceSnapshot{}, snapshot) } +func TestAnalyzeModuleResolvesRemoteAllocationGroupsInEitherOrder(t *testing.T) { + for _, tc := range []struct { + name string + completeFirst bool + }{ + {name: "pending-before-complete"}, + {name: "complete-before-pending", completeFirst: true}, + } { + t.Run(tc.name, func(t *testing.T) { + anal := &AnalyzeModule{} + pending := []remoteAllocationGroupPending{{ + Key: "execution@cn", + Count: 1, + }} + var firstPending, secondPending []remoteAllocationGroupPending + var firstCompleted, secondCompleted []string + if tc.completeFirst { + firstCompleted = []string{"execution@cn"} + secondPending = pending + } else { + firstPending = pending + secondCompleted = []string{"execution@cn"} + } + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + firstPending, firstCompleted, + ) + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + secondPending, secondCompleted, + ) + + snapshot := anal.remoteResourceSummary() + require.Empty(t, snapshot.PendingAllocationGroups) + require.Equal(t, []string{"execution@cn"}, snapshot.CompletedAllocationGroups) + require.Zero(t, snapshot.Quality&resource.QualityInvariantFailure) + }) + } +} + +func TestExecutionResourceRecorderMarksUnresolvedAllocationGroupPartial(t *testing.T) { + root := resource.NewRoot(resource.ConnExternal) + recorder := newExecutionResourceRecorder( + resource.ContextWithRoot(context.Background(), root), + true, + ) + require.NotNil(t, recorder) + anal := &AnalyzeModule{} + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + []remoteAllocationGroupPending{{Key: "execution@cn", Count: 1}}, nil, + ) + + recorder.finishAttempt( + 0, time.Now(), 0, 0, nil, nil, anal, "local:6001", false, + ) + recorder.publish() + + summary := root.PreResponseSummary() + require.Equal(t, uint64(1), summary.MissingMemoryDomainCount) + require.NotZero(t, summary.Quality&resource.QualityPartial) + require.NotZero(t, summary.Quality&resource.QualityMissingMemoryDomain) + require.Zero(t, summary.Quality&resource.QualityMissingFragment) +} + +func TestExecutionResourceRecorderPreservesPendingGroupCardinality(t *testing.T) { + root := resource.NewRoot(resource.ConnExternal) + recorder := newExecutionResourceRecorder( + resource.ContextWithRoot(context.Background(), root), + true, + ) + require.NotNil(t, recorder) + anal := &AnalyzeModule{} + for range 3 { + anal.appendRemoteResource( + resource.Delta{}, resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, 0, 0, + []remoteAllocationGroupPending{{Key: "execution@cn", Count: 1}}, + nil, + ) + } + scopes := make([]*Scope, 4) + for i := range scopes { + scopes[i] = &Scope{ + Magic: Remote, + NodeInfo: engine.Node{Addr: "remote:6001"}, + } + } + + recorder.finishAttempt( + 0, time.Now(), 0, 0, nil, scopes, anal, "local:6001", false, + ) + recorder.publish() + + summary := root.PreResponseSummary() + require.Equal(t, uint64(1), summary.MissingFragmentCount) + require.Equal(t, uint64(4), summary.MissingMemoryDomainCount) + require.NotZero(t, summary.Quality&resource.QualityPartial) + require.NotZero(t, summary.Quality&resource.QualityMissingFragment) + require.NotZero(t, summary.Quality&resource.QualityMissingMemoryDomain) +} + func TestAnalyzeModuleRemoteResourceConcurrentAccess(t *testing.T) { anal := &AnalyzeModule{} var wg sync.WaitGroup @@ -480,7 +657,15 @@ func TestAnalyzeModuleRemoteResourceConcurrentAccess(t *testing.T) { go func() { defer wg.Done() for j := 0; j < 100; j++ { - anal.appendRemoteResource(resource.Delta{Usage: resource.Usage{SpillBytes: 1}}, resource.MemoryTotals{}, 1, 1) + anal.appendRemoteResource( + resource.Delta{Usage: resource.Usage{SpillBytes: 1}}, + resource.MemoryTotals{}, + resource.AllocationAccountTotals{}, + 1, + 1, + []remoteAllocationGroupPending{{Key: "pending", Count: 1}}, + []string{"completed"}, + ) _ = anal.remoteResourceSummary() } }() diff --git a/pkg/sql/compile/scope.go b/pkg/sql/compile/scope.go index 37d329b63c001..3b11a1d798dcf 100644 --- a/pkg/sql/compile/scope.go +++ b/pkg/sql/compile/scope.go @@ -751,6 +751,10 @@ func buildLoadParallelRun(s *Scope, c *Compile) (*Scope, error) { return nil, err } } + if err := c.attachRuntimeAllocationOwners(ss); err != nil { + s.discardParallelGeneration(ms) + return nil, err + } return ms, nil } @@ -802,6 +806,10 @@ func buildScanParallelRun(s *Scope, c *Compile) (*Scope, error) { RecvMsgList: recvMsgList, } } + if err := c.attachRuntimeAllocationOwners(ss); err != nil { + s.discardParallelGeneration(ms) + return nil, err + } return ms, nil } diff --git a/pkg/sql/compile/sql_executor.go b/pkg/sql/compile/sql_executor.go index ac962d14f1c7d..605cb7d6b05e5 100644 --- a/pkg/sql/compile/sql_executor.go +++ b/pkg/sql/compile/sql_executor.go @@ -534,26 +534,28 @@ func (exec *txnExecutor) Exec( // the bat is valid only in current method. So we need copy data. // FIXME: add a custom streaming apply handler to consume readed data. Now // our current internal sql will never read too much data. - rows, err := bat.Clone(exec.s.mp, streaming) + // Internal executor results outlive Compile.Run. Keep the + // existing physical clone mode, but end statement allocation + // ownership before returning or publishing the batch. + rows, err := cloneInternalExecutorResultBatch( + bat, + exec.s.mp, + streaming, + ) if err != nil { return err } if streaming { stream_result := executor.NewResult(exec.s.mp) - for len(stream_chan) == cap(stream_chan) { - select { - case <-proc.Ctx.Done(): - err_chan <- moerr.NewInternalError(proc.Ctx, "context cancelled") - return moerr.NewInternalError(proc.Ctx, "context cancelled") - case <-exec.ctx.Done(): - err_chan <- exec.ctx.Err() - return exec.ctx.Err() - default: - time.Sleep(1 * time.Millisecond) - } - } stream_result.Batches = []*batch.Batch{rows} - stream_chan <- stream_result + if err := publishInternalExecutorStreamResult( + proc.Ctx, + exec.ctx, + stream_chan, + stream_result, + ); err != nil { + return err + } } else { batches = append(batches, rows) } @@ -599,6 +601,39 @@ func (exec *txnExecutor) Exec( return result, nil } +func cloneInternalExecutorResultBatch( + bat *batch.Batch, + mp *mpool.MPool, + streaming bool, +) (*batch.Batch, error) { + return bat.CloneWithoutAllocationAccount(mp, streaming) +} + +// publishInternalExecutorStreamResult transfers result ownership only after a +// successful send. Cancellation keeps ownership here and closes the result. +func publishInternalExecutorStreamResult( + procCtx context.Context, + execCtx context.Context, + streamChan chan executor.Result, + result executor.Result, +) error { + select { + case streamChan <- result: + return nil + default: + } + select { + case streamChan <- result: + return nil + case <-procCtx.Done(): + result.Close() + return moerr.NewInternalError(procCtx, "context cancelled") + case <-execCtx.Done(): + result.Close() + return execCtx.Err() + } +} + func (exec *txnExecutor) LockTable(table string) error { txnOp := exec.opts.Txn() ctx := exec.ctx diff --git a/pkg/sql/compile/sql_executor_allocation_test.go b/pkg/sql/compile/sql_executor_allocation_test.go new file mode 100644 index 0000000000000..7382860e6292f --- /dev/null +++ b/pkg/sql/compile/sql_executor_allocation_test.go @@ -0,0 +1,138 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package compile + +import ( + "context" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/util/executor" + "github.com/stretchr/testify/require" +) + +func TestInternalExecutorResultEndsStatementAllocationOwnership(t *testing.T) { + for _, streaming := range []bool{false, true} { + t.Run(map[bool]string{false: "retained", true: "streaming"}[streaming], func(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + + source := batch.NewWithSchema( + true, + []string{"id", "value"}, + []types.Type{types.T_int64.ToType(), types.T_varchar.ToType()}, + ) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + require.NoError(t, vector.AppendBytes(source.Vecs[1], []byte("result"), false, mp)) + source.SetRowCount(1) + + cloned, err := cloneInternalExecutorResultBatch(source, mp, streaming) + require.NoError(t, err) + require.Nil(t, cloned.AllocationAccountSelection()) + for _, vec := range cloned.Vecs { + require.Nil(t, vec.AllocationAccountSelection()) + } + + var result executor.Result + if streaming { + results := make(chan executor.Result, 1) + published := executor.NewResult(mp) + published.Batches = []*batch.Batch{cloned} + results <- published + result = <-results + } else { + result = executor.NewResult(mp) + result.Batches = []*batch.Batch{cloned} + } + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + + require.Equal( + t, + int64(42), + vector.GetFixedAtNoTypeCheck[int64](result.Batches[0].Vecs[0], 0), + ) + require.Equal(t, []byte("result"), result.Batches[0].Vecs[1].GetBytesAt(0)) + result.Close() + require.Zero(t, mp.CurrNB()) + }) + } +} + +func TestInternalExecutorStreamCancellationCleansUnpublishedResult(t *testing.T) { + mp := mpool.MustNewZero() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection(account, 1, 1, 2, 3, 4) + require.NoError(t, err) + + source := batch.NewWithSchema( + true, + []string{"value"}, + []types.Type{types.T_int64.ToType()}, + ) + require.NoError(t, source.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(source.Vecs[0], int64(42), false, mp)) + source.SetRowCount(1) + + cloned, err := cloneInternalExecutorResultBatch(source, mp, true) + require.NoError(t, err) + result := executor.NewResult(mp) + result.Batches = []*batch.Batch{cloned} + results := make(chan executor.Result, 1) + results <- executor.NewResult(mp) + procCtx, cancel := context.WithCancel(context.Background()) + cancel() + + err = publishInternalExecutorStreamResult( + procCtx, + context.Background(), + results, + result, + ) + require.Error(t, err) + require.Len(t, results, 1) + + source.Clean(mp) + snapshot := account.Seal() + require.Zero(t, snapshot.Used) + require.Zero(t, registry.LiveAllocationMetadata()) + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Zero(t, mp.CurrNB()) +} diff --git a/pkg/sql/compile/types.go b/pkg/sql/compile/types.go index 4b7af939ae88b..c4365a7b82908 100644 --- a/pkg/sql/compile/types.go +++ b/pkg/sql/compile/types.go @@ -21,6 +21,7 @@ import ( "github.com/google/uuid" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" icebergapi "github.com/matrixorigin/matrixone/pkg/iceberg/api" @@ -343,6 +344,14 @@ type Compile struct { // resourceAttemptOwnerEligible is set only for the top-level statement // Compile. The statement root still arbitrates the single actual owner. resourceAttemptOwnerEligible bool + allocationAccountRegistry *mpool.AllocationAccountRegistry + allocationAccountLimit uint64 + allocationControllerProvider func() (mpool.AllocationCapacityController, error) + allocationTerminalExporter func(mpool.AllocationAccountTerminalSnapshot) + allocationAccountOwners []executionAllocationAccountOwner + allocationAttempt *statementAllocationAttempt + remoteFragmentCounts map[string]uint32 + remoteExecutionID uuid.UUID hasMergeOp bool // ncpu set as system.GoRoutines() while NewCompile, instead of global static value. diff --git a/pkg/sql/plan/apply_indices_test.go b/pkg/sql/plan/apply_indices_test.go index 1cc81da1e0b21..4d5b4775caf82 100644 --- a/pkg/sql/plan/apply_indices_test.go +++ b/pkg/sql/plan/apply_indices_test.go @@ -752,6 +752,7 @@ func TestUniqueIndexRuntimeFilterUsesSelectedHashSlot(t *testing.T) { } arg.RuntimeFilterSpec = spec arg.AppendChild(child) + registry, account := installIndexPlanHashBuildAllocation(t, arg) require.NoError(t, child.Prepare(proc)) require.NoError(t, arg.Prepare(proc)) result, err := vm.Exec(arg, proc) @@ -778,10 +779,11 @@ func TestUniqueIndexRuntimeFilterUsesSelectedHashSlot(t *testing.T) { payload.Free(proc.Mp()) runtimeFilter.Destroy() arg.Free(proc, false, nil) + proc.GetMessageBoard().Reset() + finishIndexPlanHashBuildAllocation(t, registry, account, arg) child.Free(proc, false, nil) arg.Release() child.Release() - proc.GetMessageBoard().Reset() proc.Free() require.Zero(t, proc.Mp().CurrNB()) } @@ -858,6 +860,7 @@ func TestIndexJoinGeneratedSerializedRuntimeFilterExecutesEndToEnd(t *testing.T) planpb.Type{Id: int32(types.T_int32)}, 0, 0)} arg.RuntimeFilterSpec = spec arg.AppendChild(child) + registry, account := installIndexPlanHashBuildAllocation(t, arg) require.NoError(t, child.Prepare(proc)) require.NoError(t, arg.Prepare(proc)) result, err := vm.Exec(arg, proc) @@ -916,14 +919,42 @@ func TestIndexJoinGeneratedSerializedRuntimeFilterExecutesEndToEnd(t *testing.T) payload.Free(proc.Mp()) runtimeFilter.Destroy() arg.Free(proc, false, nil) + proc.GetMessageBoard().Reset() + finishIndexPlanHashBuildAllocation(t, registry, account, arg) child.Free(proc, false, nil) arg.Release() child.Release() - proc.GetMessageBoard().Reset() proc.Free() require.Zero(t, proc.Mp().CurrNB()) } +func installIndexPlanHashBuildAllocation( + t testing.TB, + arg *hashbuild.HashBuild, +) (*mpool.AllocationAccountRegistry, *mpool.AllocationAccount) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 4_096) + require.NoError(t, err) + account, err := registry.Open(1 << 60) + require.NoError(t, err) + require.NoError(t, arg.SetAllocationAccount(account)) + return registry, account +} + +func finishIndexPlanHashBuildAllocation( + t testing.TB, + registry *mpool.AllocationAccountRegistry, + account *mpool.AllocationAccount, + arg *hashbuild.HashBuild, +) { + t.Helper() + require.NoError(t, arg.ClearAllocationAccount(account)) + snapshot, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Zero(t, snapshot.Used) +} + func TestForceIndexForJoinBuildsRightAccessWithoutReorder(t *testing.T) { builder, joinID, leftScanID, _ := makeIndexHintJoinBuilder(t) joinNode := builder.qry.Nodes[joinID] diff --git a/pkg/sql/util/copy_batch_test.go b/pkg/sql/util/copy_batch_test.go index 643a56073c103..659ff1db4f210 100644 --- a/pkg/sql/util/copy_batch_test.go +++ b/pkg/sql/util/copy_batch_test.go @@ -91,6 +91,49 @@ func TestCopyBatchCompactsAndFlattens(t *testing.T) { require.Equal(t, int64(0), mp.CurrNB()) } +func TestCopyBatchCrossesAllocationOwnershipBoundary(t *testing.T) { + proc := testutil.NewProcess(t) + mp := proc.Mp() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + 1, + 1, + 2, + 3, + 4, + ) + require.NoError(t, err) + + src := batch.NewOffHeapWithSize(1) + vec := vector.NewOffHeapVecWithType(types.T_int64.ToType()) + require.NoError(t, vec.SetAllocationAccount(selection)) + require.NoError(t, vector.AppendFixed(vec, int64(42), false, mp)) + src.SetVector(0, vec) + src.SetRowCount(1) + sourceUsed := account.Snapshot().Used + require.Positive(t, sourceUsed) + + got, err := CopyBatch(src, proc) + require.NoError(t, err) + require.Nil(t, got.Vecs[0].AllocationAccountSelection()) + require.Equal(t, int64(42), vector.GetFixedAtNoTypeCheck[int64](got.Vecs[0], 0)) + require.Equal(t, sourceUsed, account.Snapshot().Used) + + src.Clean(mp) + require.Zero(t, account.Snapshot().Used) + require.Equal(t, int64(42), vector.GetFixedAtNoTypeCheck[int64](got.Vecs[0], 0)) + got.Clean(mp) + + account.Seal() + _, err = registry.Finalize(account) + require.NoError(t, err) + require.Equal(t, int64(0), mp.CurrNB()) +} + func BenchmarkCopyBatchCompact(b *testing.B) { proc := testutil.NewProcess(b) mp := proc.Mp() diff --git a/pkg/sql/util/eval_expr_util.go b/pkg/sql/util/eval_expr_util.go index da58ac7b71d38..56bf9d641aee3 100644 --- a/pkg/sql/util/eval_expr_util.go +++ b/pkg/sql/util/eval_expr_util.go @@ -111,13 +111,28 @@ func DecodeBinaryString(s string) ([]byte, error) { } func GenVectorByVarValue(proc *process.Process, typ types.Type, val any) (*vector.Vector, error) { + return GenVectorByVarValueWithAllocation(proc, typ, val, nil) +} + +func GenVectorByVarValueWithAllocation( + proc *process.Process, + typ types.Type, + val any, + selection *vector.AllocationAccountSelection, +) (*vector.Vector, error) { if val == nil { - vec := vector.NewConstNull(typ, 1, proc.Mp()) - return vec, nil - } else { - strVal := getVal(val) + if selection == nil { + return vector.NewConstNull(typ, 1, proc.Mp()), nil + } + return vector.NewConstNullWithAllocation(typ, 1, selection) + } + strVal := getVal(val) + if selection == nil { return vector.NewConstBytes(typ, []byte(strVal), 1, proc.Mp()) } + return vector.NewConstBytesWithAllocation( + typ, []byte(strVal), 1, proc.Mp(), selection, + ) } func AppendAnyToStringVector(proc *process.Process, val any, vec *vector.Vector) error { diff --git a/pkg/sql/util/util.go b/pkg/sql/util/util.go index f329edce53af4..00930be1ea470 100644 --- a/pkg/sql/util/util.go +++ b/pkg/sql/util/util.go @@ -21,6 +21,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/catalog" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/defines" "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" "github.com/matrixorigin/matrixone/pkg/vm/process" @@ -44,7 +45,18 @@ func CopyBatch(bat *batch.Batch, proc *process.Process) (*batch.Batch, error) { rbat := batch.NewWithSize(len(bat.Vecs)) rbat.Attrs = append(rbat.Attrs, bat.Attrs...) for i, srcVec := range bat.Vecs { - vec, err := srcVec.CloneToFlatCompact(proc.Mp()) + var ( + vec *vector.Vector + err error + ) + if srcVec.AllocationAccountSelection() != nil { + // CopyBatch is an ownership boundary: the source keeps its physical + // account until Free, while the independent destination belongs to + // the generic downstream batch owner. + vec, err = srcVec.CloneToFlatCompactWithAllocation(proc.Mp(), nil) + } else { + vec, err = srcVec.CloneToFlatCompact(proc.Mp()) + } if err != nil { rbat.Clean(proc.Mp()) return nil, err diff --git a/pkg/util/resource/summary.go b/pkg/util/resource/summary.go index 3eef2ffd9613a..74389d8619690 100644 --- a/pkg/util/resource/summary.go +++ b/pkg/util/resource/summary.go @@ -57,6 +57,106 @@ type MemoryTotals struct { CrossPoolFreeCount uint64 } +// AllocationAccountTotals is the fixed-size terminal observation of accounted +// allocation generations. It is diagnostic only: these bytes are a subset of +// allocator memory and are never added to MemoryTotals or fed back into +// admission. +type AllocationAccountTotals struct { + GenerationCount uint64 + ValidGenerationCount uint64 + FailedGenerationCount uint64 + MaxGenerationPeak uint64 + SumGenerationPeak uint64 + LiveBytesAtTerminal uint64 +} + +func (t *AllocationAccountTotals) AddGeneration( + peak uint64, + liveAtTerminal uint64, + valid bool, +) QualityFlags { + var quality QualityFlags + t.GenerationCount, quality = addChecked(t.GenerationCount, 1, quality) + if valid { + t.ValidGenerationCount, quality = addChecked( + t.ValidGenerationCount, + 1, + quality, + ) + } else { + t.FailedGenerationCount, quality = addChecked( + t.FailedGenerationCount, + 1, + quality|QualityInvariantFailure, + ) + } + if peak > t.MaxGenerationPeak { + t.MaxGenerationPeak = peak + } + t.SumGenerationPeak, quality = addChecked( + t.SumGenerationPeak, + peak, + quality, + ) + t.LiveBytesAtTerminal, quality = addChecked( + t.LiveBytesAtTerminal, + liveAtTerminal, + quality, + ) + if liveAtTerminal != 0 { + quality |= QualityNonZeroLiveAtSeal | QualityInvariantFailure + } + return quality +} + +func MergeAllocationAccountTotals( + dst *AllocationAccountTotals, + delta AllocationAccountTotals, +) QualityFlags { + var quality QualityFlags + dst.GenerationCount, quality = addChecked( + dst.GenerationCount, + delta.GenerationCount, + quality, + ) + dst.ValidGenerationCount, quality = addChecked( + dst.ValidGenerationCount, + delta.ValidGenerationCount, + quality, + ) + dst.FailedGenerationCount, quality = addChecked( + dst.FailedGenerationCount, + delta.FailedGenerationCount, + quality, + ) + dst.SumGenerationPeak, quality = addChecked( + dst.SumGenerationPeak, + delta.SumGenerationPeak, + quality, + ) + dst.LiveBytesAtTerminal, quality = addChecked( + dst.LiveBytesAtTerminal, + delta.LiveBytesAtTerminal, + quality, + ) + if delta.MaxGenerationPeak > dst.MaxGenerationPeak { + dst.MaxGenerationPeak = delta.MaxGenerationPeak + } + if delta.FailedGenerationCount != 0 || delta.LiveBytesAtTerminal != 0 { + quality |= QualityInvariantFailure + } + if delta.ValidGenerationCount > delta.GenerationCount || + delta.FailedGenerationCount > + delta.GenerationCount-delta.ValidGenerationCount || + delta.MaxGenerationPeak > delta.SumGenerationPeak { + quality |= QualityInvariantFailure + } + if delta.LiveBytesAtTerminal != 0 { + quality |= QualityNonZeroLiveAtSeal + } + return quality +} + // MergeMemoryDomain merges one physical domain exactly once. func MergeMemoryDomain(dst *MemoryTotals, domain MemoryDomainSummary) QualityFlags { flags := domain.Validate() @@ -74,8 +174,9 @@ func MergeMemoryDomain(dst *MemoryTotals, domain MemoryDomainSummary) QualityFla // AttemptSummary is the immutable result of one compile/run generation. type AttemptSummary struct { - Usage Usage - Memory MemoryTotals + Usage Usage + Memory MemoryTotals + Allocation AllocationAccountTotals WallNS uint64 MissingFragmentCount uint64 @@ -85,8 +186,9 @@ type AttemptSummary struct { // ExecutionSummary is fixed-size in retry count. type ExecutionSummary struct { - Usage Usage - Memory MemoryTotals + Usage Usage + Memory MemoryTotals + Allocation AllocationAccountTotals AttemptCount uint64 RetryWallNS uint64 @@ -100,6 +202,7 @@ type ExecutionSummary struct { func (s *ExecutionSummary) AddAttempt(attempt AttemptSummary, retried bool) { s.Quality |= attempt.Quality | MergeUsage(&s.Usage, attempt.Usage) s.Quality |= MergeMemoryTotals(&s.Memory, attempt.Memory) + s.Quality |= MergeAllocationAccountTotals(&s.Allocation, attempt.Allocation) s.AttemptCount, s.Quality = addChecked(s.AttemptCount, 1, s.Quality) if retried { s.RetryWallNS, s.Quality = addChecked(s.RetryWallNS, attempt.WallNS, s.Quality) @@ -123,8 +226,9 @@ const ( // algebra. Serialization and plan diagnostics consume this value but never add // resources independently. type StatementResourceSummary struct { - Usage Usage - Memory MemoryTotals + Usage Usage + Memory MemoryTotals + Allocation AllocationAccountTotals StatementWallNS uint64 AttemptCount uint64 @@ -140,6 +244,7 @@ type StatementResourceSummary struct { func (s *StatementResourceSummary) MergeExecution(execution ExecutionSummary) { s.Quality |= execution.Quality | MergeUsage(&s.Usage, execution.Usage) s.Quality |= MergeMemoryTotals(&s.Memory, execution.Memory) + s.Quality |= MergeAllocationAccountTotals(&s.Allocation, execution.Allocation) s.AttemptCount, s.Quality = addChecked(s.AttemptCount, execution.AttemptCount, s.Quality) s.RetryWallNS, s.Quality = addChecked(s.RetryWallNS, execution.RetryWallNS, s.Quality) s.MissingFragmentCount, s.Quality = addChecked( @@ -154,6 +259,7 @@ func (s *StatementResourceSummary) MergeExecution(execution ExecutionSummary) { func (s *StatementResourceSummary) Merge(other StatementResourceSummary) { s.Quality |= other.Quality | QualityAggregated | MergeUsage(&s.Usage, other.Usage) s.Quality |= MergeMemoryTotals(&s.Memory, other.Memory) + s.Quality |= MergeAllocationAccountTotals(&s.Allocation, other.Allocation) s.StatementWallNS, s.Quality = addChecked(s.StatementWallNS, other.StatementWallNS, s.Quality) s.AttemptCount, s.Quality = addChecked(s.AttemptCount, other.AttemptCount, s.Quality) s.RetryWallNS, s.Quality = addChecked(s.RetryWallNS, other.RetryWallNS, s.Quality) diff --git a/pkg/util/resource/usage_test.go b/pkg/util/resource/usage_test.go index 3968db4711d0d..1bf6d14ecdbe3 100644 --- a/pkg/util/resource/usage_test.go +++ b/pkg/util/resource/usage_test.go @@ -85,6 +85,30 @@ func TestMergeUsageOverflowSaturates(t *testing.T) { } } +func TestAllocationAccountTotalsMergeAndFailureQuality(t *testing.T) { + var valid AllocationAccountTotals + if quality := valid.AddGeneration(10, 0, true); quality != 0 { + t.Fatalf("valid generation quality = %v", quality) + } + var failed AllocationAccountTotals + quality := failed.AddGeneration(20, 5, false) + if quality&QualityInvariantFailure == 0 || + quality&QualityNonZeroLiveAtSeal == 0 { + t.Fatalf("failure quality = %v", quality) + } + + quality = MergeAllocationAccountTotals(&valid, failed) + if valid.GenerationCount != 2 || valid.ValidGenerationCount != 1 || + valid.FailedGenerationCount != 1 || valid.MaxGenerationPeak != 20 || + valid.SumGenerationPeak != 30 || valid.LiveBytesAtTerminal != 5 { + t.Fatalf("merged allocation totals = %+v", valid) + } + if quality&QualityInvariantFailure == 0 || + quality&QualityNonZeroLiveAtSeal == 0 { + t.Fatalf("merge quality = %v", quality) + } +} + func TestLocalRecorder(t *testing.T) { var recorder LocalRecorder recorder.AddActiveInterval(100, 10, 20) diff --git a/pkg/vm/engine/disttae/txn_table.go b/pkg/vm/engine/disttae/txn_table.go index c934427bbd039..a1d280ba1e860 100644 --- a/pkg/vm/engine/disttae/txn_table.go +++ b/pkg/vm/engine/disttae/txn_table.go @@ -3599,6 +3599,12 @@ func (tbl *txnTable) GetExtraInfo() *api.SchemaExtra { // If v has no NULLs it returns Dup(v). The caller must Free the result. func dupVectorWithoutNulls(v *vector.Vector, mp *mpool.MPool) (*vector.Vector, error) { if !v.HasNull() { + if v.AllocationAccountSelection() != nil { + // PK validation borrows caller-owned data. Its locally sorted copy is a + // short-lived transaction-engine owner, not a continuation of the + // statement owner, so make that ownership exit explicit. + return v.DupOffHeapWithAllocation(mp, nil) + } return v.Dup(mp) } filtered := vector.NewVec(*v.GetType()) diff --git a/pkg/vm/engine/disttae/txn_test.go b/pkg/vm/engine/disttae/txn_test.go index bf4ac5e1c564d..aae5464e30d2a 100644 --- a/pkg/vm/engine/disttae/txn_test.go +++ b/pkg/vm/engine/disttae/txn_test.go @@ -26,6 +26,7 @@ import ( "github.com/golang/mock/gomock" "github.com/matrixorigin/matrixone/pkg/catalog" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" @@ -1560,6 +1561,59 @@ func TestDupVectorWithoutNulls(t *testing.T) { }) } +func TestDupVectorWithoutNullsLeavesSealedStatementOwner(t *testing.T) { + proc := testutil.NewProc(t) + mp := proc.Mp() + + for _, tc := range []struct { + name string + withNulls bool + }{ + {name: "no nulls"}, + {name: "with nulls", withNulls: true}, + } { + t.Run(tc.name, func(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + selection, err := vector.NewAllocationAccountSelection( + account, + mpool.AllocationOwner(1), + mpool.AllocationSite(1), + mpool.AllocationSite(2), + mpool.AllocationSite(3), + mpool.AllocationSite(4), + ) + require.NoError(t, err) + source, err := vector.NewOffHeapVecWithTypeAndAllocation( + types.T_int64.ToType(), + selection, + ) + require.NoError(t, err) + require.NoError(t, vector.AppendFixed(source, int64(1), false, mp)) + if tc.withNulls { + require.NoError(t, vector.AppendFixed(source, int64(0), true, mp)) + } + require.NoError(t, vector.AppendFixed(source, int64(2), false, mp)) + + used := account.Seal().Used + require.NotZero(t, used) + out, err := dupVectorWithoutNulls(source, mp) + require.NoError(t, err) + require.Nil(t, out.AllocationAccountSelection()) + require.Equal(t, used, account.Snapshot().Used) + require.Equal(t, []int64{1, 2}, vector.MustFixedColWithTypeCheck[int64](out)) + + out.Free(mp) + source.Free(mp) + require.Zero(t, account.Snapshot().Used) + _, err = registry.Finalize(account) + require.NoError(t, err) + }) + } +} + func newInt64BatchForTest( t *testing.T, proc *process.Process, diff --git a/pkg/vm/engine/tae/txn/txnimpl/anode.go b/pkg/vm/engine/tae/txn/txnimpl/anode.go index b41477515ab64..05990d6aa7ef8 100644 --- a/pkg/vm/engine/tae/txn/txnimpl/anode.go +++ b/pkg/vm/engine/tae/txn/txnimpl/anode.go @@ -19,6 +19,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/nulls" + cnvector "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/objectio" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/catalog" "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/common" @@ -118,6 +119,7 @@ func (n *anode) Append(data *containers.Batch, offset uint32) (an uint32, err er from := uint32(n.data.Length()) an = n.PrepareAppend(data, offset) + fakePKProvided := false for _, attr := range data.Attrs { if attr == catalog.PhyAddrColumnName { continue @@ -128,12 +130,27 @@ func (n *anode) Append(data *containers.Batch, offset uint32) (an uint32, err er // } // } def := schema.ColDefs[schema.GetColIdx(attr)] + fakePKProvided = fakePKProvided || def.FakePK destVec := n.data.Vecs[def.Idx] // logutil.Infof("destVec: %s, %d, %d", destVec.String(), cnt, data.Length()) destVec.ExtendWithOffset(data.Vecs[def.Idx], int(offset), int(an)) } + if schema.HasFakePK() && !fakePKProvided { + fakePK := n.data.Vecs[schema.GetPrimaryKey().Idx] + if err = cnvector.AppendMultiFixed( + fakePK.GetDownstreamVector(), + uint64(0), + false, + int(an), + fakePK.GetAllocator(), + ); err != nil { + return + } + } + if err = n.FillPhyAddrColumn(from, an); err != nil { + return + } n.rows = uint32(n.data.Length()) - err = n.FillPhyAddrColumn(from, an) return } diff --git a/pkg/vm/engine/tae/txn/txnimpl/table.go b/pkg/vm/engine/tae/txn/txnimpl/table.go index d334cd820d779..1bfec495efafb 100644 --- a/pkg/vm/engine/tae/txn/txnimpl/table.go +++ b/pkg/vm/engine/tae/txn/txnimpl/table.go @@ -558,6 +558,7 @@ func (tbl *txnTable) TransferDeletes( } deletes.Vecs[i].CompactByBitmap(&transferd) } + tbl.tombstoneTable.tableSpace.node.rows = uint32(deletes.Length()) return } diff --git a/pkg/vm/engine/test/workspace_test.go b/pkg/vm/engine/test/workspace_test.go index 5946fd9176e7f..ec4c0f96cf157 100644 --- a/pkg/vm/engine/test/workspace_test.go +++ b/pkg/vm/engine/test/workspace_test.go @@ -1079,10 +1079,15 @@ func Test_MultiTxnInsertDelete(t *testing.T) { { require.NoError(t, testutil.WriteToRelation(ctx, txn, relation, bat2, false, true)) + localPKs := vector.MustFixedColWithTypeCheck[int64](bat2.Vecs[primaryKeyIdx]) + localPKOffset := 0 txn.GetWorkspace().(*disttae.Transaction).ForEachTableWrites( relation.GetDBID(ctx), relation.GetTableID(ctx), 1, func(entry disttae.Entry) { waitedDeletes := vector.MustFixedColWithTypeCheck[types.Rowid](entry.Bat().GetVector(0)) require.NoError(t, vector.AppendFixedList[types.Rowid](tombstoneBat.Vecs[0], waitedDeletes, nil, mp)) + require.NoError(t, vector.AppendFixedList[int64]( + tombstoneBat.Vecs[1], localPKs[localPKOffset:localPKOffset+len(waitedDeletes)], nil, mp)) + localPKOffset += len(waitedDeletes) tombstoneBat.SetRowCount(tombstoneBat.RowCount() + len(waitedDeletes)) }) @@ -1882,10 +1887,15 @@ func Test_MultiTxnRollbackStatement(t *testing.T) { require.NoError(t, txn.GetWorkspace().RollbackLastStatement(ctx)) require.NoError(t, txn.GetWorkspace().IncrStatementID(ctx, false)) + localPKs := vector.MustFixedColWithTypeCheck[int64](bat2.Vecs[primaryKeyIdx]) + localPKOffset := 0 txn.GetWorkspace().(*disttae.Transaction).ForEachTableWrites( relation.GetDBID(ctx), relation.GetTableID(ctx), 1, func(entry disttae.Entry) { waitedDeletes := vector.MustFixedColWithTypeCheck[types.Rowid](entry.Bat().GetVector(0)) require.NoError(t, vector.AppendFixedList[types.Rowid](tombstoneBat.Vecs[0], waitedDeletes, nil, mp)) + require.NoError(t, vector.AppendFixedList[int64]( + tombstoneBat.Vecs[1], localPKs[localPKOffset:localPKOffset+len(waitedDeletes)], nil, mp)) + localPKOffset += len(waitedDeletes) tombstoneBat.SetRowCount(tombstoneBat.RowCount() + len(waitedDeletes)) }) diff --git a/pkg/vm/message/group_sels_test.go b/pkg/vm/message/group_sels_test.go index d6b3aa5132711..bf5fab352ddac 100644 --- a/pkg/vm/message/group_sels_test.go +++ b/pkg/vm/message/group_sels_test.go @@ -25,6 +25,25 @@ func testMp() *mpool.MPool { return mpool.MustNewZero() } +func initTestGroupSels( + t *testing.T, + sels *GroupSels, + n int, + mp *mpool.MPool, +) { + t.Helper() + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + require.NoError(t, sels.InitWithAllocation(n, mp, account, 1, 1)) + t.Cleanup(func() { + sels.Free(mp) + _, _, err := registry.CompleteTerminal(account) + require.NoError(t, err) + }) +} + func TestGroupSels_NilBeforeInit(t *testing.T) { var js GroupSels require.NoError(t, js.Finalize(3, 3, testMp())) @@ -35,7 +54,7 @@ func TestGroupSels_NilBeforeInit(t *testing.T) { func TestGroupSels_AllUnique(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(3, mp)) + initTestGroupSels(t, &js, 3, mp) js.Insert(0, 0) js.Insert(1, 1) js.Insert(2, 2) @@ -47,7 +66,7 @@ func TestGroupSels_AllUnique(t *testing.T) { func TestGroupSels_Normal0Based(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(4, mp)) + initTestGroupSels(t, &js, 4, mp) js.Insert(0, 10) js.Insert(1, 11) js.Insert(0, 12) @@ -63,7 +82,7 @@ func TestGroupSels_Normal0Based(t *testing.T) { func TestGroupSels_Dedup1Based(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(3, mp)) + initTestGroupSels(t, &js, 3, mp) js.Insert(1, 0) js.Insert(2, 1) js.Insert(1, 2) @@ -78,7 +97,7 @@ func TestGroupSels_Dedup1Based(t *testing.T) { func TestGroupSels_Free(t *testing.T) { mp := testMp() var js GroupSels - require.NoError(t, js.Init(2, mp)) + initTestGroupSels(t, &js, 2, mp) js.Insert(0, 5) js.Free(mp) require.NoError(t, js.Finalize(1, 1, mp)) @@ -89,7 +108,7 @@ func TestGroupSels_AllNulls(t *testing.T) { // all rows are null — Init called but Insert never called mp := testMp() var js GroupSels - require.NoError(t, js.Init(3, mp)) + initTestGroupSels(t, &js, 3, mp) // no Insert calls require.NoError(t, js.Finalize(0, 3, mp)) require.Nil(t, js.offsets) @@ -100,7 +119,7 @@ func TestGroupSels_NullsSkipped(t *testing.T) { // must NOT trigger all-unique path mp := testMp() var js GroupSels - require.NoError(t, js.Init(4, mp)) + initTestGroupSels(t, &js, 4, mp) js.Insert(0, 0) // row 0 → group 0 js.Insert(1, 1) // row 1 → group 1 // row 2 is null, skipped @@ -112,3 +131,66 @@ func TestGroupSels_NullsSkipped(t *testing.T) { require.ElementsMatch(t, []int32{3}, js.Get(2)) js.Free(mp) } + +func TestGroupSelsAllocationAccountLifecycleAndRollback(t *testing.T) { + const ( + owner mpool.AllocationOwner = 1 + site mpool.AllocationSite = 30 + ) + for _, tc := range []struct { + name string + limit uint64 + metadataSlots uint64 + wantErr error + }{ + {name: "exact", limit: 64, metadataSlots: 3}, + { + name: "one byte short", + limit: 63, + metadataSlots: 3, + wantErr: mpool.ErrAllocationAccountCapacity, + }, + { + name: "one metadata slot short", + limit: 64, + metadataSlots: 2, + wantErr: mpool.ErrAllocationMetadataSlots, + }, + } { + t.Run(tc.name, func(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, tc.metadataSlots) + require.NoError(t, err) + account, err := registry.Open(tc.limit) + require.NoError(t, err) + mp := testMp() + var sels GroupSels + require.NoError(t, sels.InitWithAllocation(4, mp, account, owner, site)) + sels.Insert(0, 0) + sels.Insert(0, 1) + sels.Insert(1, 2) + sels.Insert(1, 3) + + err = sels.Finalize(2, 4, mp) + if tc.wantErr == nil { + require.NoError(t, err) + require.Equal(t, uint64(64), account.Snapshot().Peak) + require.Equal(t, uint64(32), account.Snapshot().Used) + require.ElementsMatch(t, []int32{0, 1}, sels.Get(0)) + require.ElementsMatch(t, []int32{2, 3}, sels.Get(1)) + } else { + require.ErrorIs(t, err, tc.wantErr) + require.Nil(t, sels.offsets) + require.Nil(t, sels.vals) + require.NotNil(t, sels.tmp) + require.Equal(t, uint64(32), account.Snapshot().Used) + } + sels.Free(mp) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, registry.LiveAllocationMetadata()) + snapshot, first, err := registry.CompleteTerminal(account) + require.NoError(t, err) + require.True(t, first) + require.Equal(t, mpool.AllocationAccountTerminalValid, snapshot.State) + }) + } +} diff --git a/pkg/vm/message/joinMapDependency_test.go b/pkg/vm/message/joinMapDependency_test.go index 77b3965369220..50513eca3cc4c 100644 --- a/pkg/vm/message/joinMapDependency_test.go +++ b/pkg/vm/message/joinMapDependency_test.go @@ -48,6 +48,19 @@ func TestJoinMapResultDistinguishesSuccessEmptyAndBuildError(t *testing.T) { require.Equal(t, baseErr.ErrorCode(), got.ErrorCode()) } +func TestSendJoinMapResultRetainsOwnershipWhenBoardUnavailable(t *testing.T) { + jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, nil) + require.False(t, SendJoinMapResult( + NewJoinMapResult(jm), + 1, + false, + 0, + nil, + )) + require.True(t, jm.IsValid()) + jm.FreeMemory() +} + func TestRuntimeFilterMemoryReleaseIsSharedAcrossMessageCopies(t *testing.T) { var releases atomic.Int32 msg := RuntimeFilterMessage{Data: make([]byte, 128)} diff --git a/pkg/vm/message/joinMapMsg.go b/pkg/vm/message/joinMapMsg.go index 3e0e1ed1fb2cb..d7385949e22a9 100644 --- a/pkg/vm/message/joinMapMsg.go +++ b/pkg/vm/message/joinMapMsg.go @@ -17,6 +17,7 @@ package message import ( "bytes" "context" + "math" "os" "strconv" "sync" @@ -39,22 +40,69 @@ type GroupSels struct { // tmp holds (groupID, rowID) pairs during the build phase, before Finalize. tmp []int32 + + account *mpool.AllocationAccount + owner mpool.AllocationOwner + site mpool.AllocationSite } func freeSlice[T any](mp *mpool.MPool, s []T) { mpool.FreeSlice(mp, s[:cap(s)]) } -func (sels *GroupSels) Init(n int, mp *mpool.MPool) error { +// InitWithAllocation makes the complete temporary/final row-index owner use +// one immutable allocation generation. GroupSels is copied into JoinMap at +// publication, so its physical slices retain this provenance until the last +// consumer frees the map. +func (sels *GroupSels) InitWithAllocation( + n int, + mp *mpool.MPool, + account *mpool.AllocationAccount, + owner mpool.AllocationOwner, + site mpool.AllocationSite, +) error { + if n < 0 || n > math.MaxInt/2 { + return mpool.ErrAllocationAccountInvalid + } + if sels.tmp != nil || sels.vals != nil || sels.offsets != nil { + return mpool.ErrAllocationAccountInvariant + } + if account == nil || account.Handle() == 0 || + owner < mpool.AllocationOwnerMin || owner > mpool.AllocationOwnerMax || + site < mpool.AllocationSiteMin { + return mpool.ErrAllocationAccountInvalid + } var err error - sels.tmp, err = mpool.MakeSlice[int32](n*2, mp, false) + sels.tmp, err = mpool.MakeSliceAccounted[int32]( + n*2, + mp, + account, + owner, + site, + ) if err != nil { return err } + sels.account = account + sels.owner = owner + sels.site = site sels.tmp = sels.tmp[:0] return nil } +func (sels *GroupSels) makeSlice(n int, mp *mpool.MPool) ([]int32, error) { + if sels.account == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + return mpool.MakeSliceAccounted[int32]( + n, + mp, + sels.account, + sels.owner, + sels.site, + ) +} + func (sels *GroupSels) Free(mp *mpool.MPool) { if mp != nil { freeSlice(mp, sels.vals) @@ -64,6 +112,9 @@ func (sels *GroupSels) Free(mp *mpool.MPool) { sels.vals = nil sels.offsets = nil sels.tmp = nil + sels.account = nil + sels.owner = 0 + sels.site = 0 } func (sels *GroupSels) Size() int64 { @@ -107,8 +158,7 @@ func (sels *GroupSels) Finalize(groupCount int, inputRowCount int, mp *mpool.MPo } } // groupCount+2: +1 for sentinel, +1 for 1-based callers (dedup UPDATE uses keys 1..groupCount) - var err error - sels.offsets, err = mpool.MakeSlice[int32](groupCount+2, mp, false) + offsets, err := sels.makeSlice(groupCount+2, mp) if err != nil { return err } @@ -116,26 +166,29 @@ func (sels *GroupSels) Finalize(groupCount int, inputRowCount int, mp *mpool.MPo // count occurrences per group for i := 0; i < len(sels.tmp); i += 2 { k := sels.tmp[i] - sels.offsets[k+1]++ + offsets[k+1]++ } // prefix sum - for i := int32(1); i < int32(len(sels.offsets)); i++ { - sels.offsets[i] += sels.offsets[i-1] + for i := int32(1); i < int32(len(offsets)); i++ { + offsets[i] += offsets[i-1] } // scatter vals using offsets as write cursors, then recover - sels.vals, err = mpool.MakeSlice[int32](n, mp, false) + vals, err := sels.makeSlice(n, mp) if err != nil { + freeSlice(mp, offsets) return err } for i := 0; i < len(sels.tmp); i += 2 { k := sels.tmp[i] v := sels.tmp[i+1] - sels.vals[sels.offsets[k]] = v - sels.offsets[k]++ + vals[offsets[k]] = v + offsets[k]++ } // recover offsets: shift right by one - copy(sels.offsets[1:], sels.offsets[:len(sels.offsets)-1]) - sels.offsets[0] = 0 + copy(offsets[1:], offsets[:len(offsets)-1]) + offsets[0] = 0 + sels.vals = vals + sels.offsets = offsets freeSlice(mp, sels.tmp) sels.tmp = nil return nil @@ -165,9 +218,9 @@ type JoinMap struct { memoryReleaseOnce sync.Once // A resident JoinMap may be broadcast to multiple consumers, but a spill - // payload is move-only. Keep the complete payload behind one lock so files, - // legacy descriptors, and the producer budget generation cannot be claimed - // by different consumers. + // payload is move-only. Keep the complete payload behind one lock so files + // and the producer budget generation cannot be claimed by different + // consumers. spillMu sync.Mutex spilled atomic.Bool spillPayload SpillBuildPayload @@ -177,9 +230,7 @@ type JoinMap struct { var ( ErrSpillBuildPayloadEmpty = moerr.NewInternalErrorNoCtx("spill build payload is empty") - ErrSpillBuildPayloadMixed = moerr.NewInternalErrorNoCtx("spill build payload mixes accounted files and legacy descriptors") ErrSpillBuildBudgetRef = moerr.NewInternalErrorNoCtx("accounted spill build payload is missing its budget reference") - ErrSpillBuildLegacyBudget = moerr.NewInternalErrorNoCtx("legacy spill build payload must not carry a budget reference") ErrSpillBuildPayloadSet = moerr.NewInternalErrorNoCtx("spill build payload is already set") ErrSpillBuildPayloadTaken = moerr.NewInternalErrorNoCtx("spill build payload is already taken") ErrSpillBuildShared = moerr.NewInternalErrorNoCtx("spill build payload requires exactly one consumer") @@ -191,7 +242,6 @@ var ( // transferring the files to SpillEngine. type SpillBuildPayload struct { Files []*SpillFile - LegacyFds []*os.File BudgetRef any } @@ -212,15 +262,6 @@ func (p *SpillBuildPayload) Close() error { } } p.Files = nil - for i, fd := range p.LegacyFds { - if fd != nil { - if err := fd.Close(); err != nil && firstErr == nil { - firstErr = err - } - p.LegacyFds[i] = nil - } - } - p.LegacyFds = nil p.BudgetRef = nil return firstErr } @@ -268,6 +309,40 @@ func (f *SpillFile) Bytes() uint64 { return f.bytes } +// Validate proves that the move-only descriptor still names the complete +// physical file recorded by its producer. Writers publish Rows and Bytes only +// after complete-record writes, so a non-empty spill file must have positive +// metadata and an exact physical size before any record is decoded. +func (f *SpillFile) Validate() error { + if f == nil { + return moerr.NewInternalErrorNoCtx("nil spill file") + } + f.mu.Lock() + defer f.mu.Unlock() + if f.fd == nil { + return moerr.NewInternalErrorNoCtx("invalid spill file metadata") + } + if f.rows <= 0 { + return moerr.NewInternalErrorNoCtx("corrupted spill file row count metadata") + } + if f.bytes == 0 { + return moerr.NewInternalErrorNoCtx("corrupted spill file size metadata") + } + info, err := f.fd.Stat() + if err != nil { + return err + } + if info.Size() < 0 || uint64(info.Size()) != f.bytes { + return moerr.NewInternalErrorf( + context.Background(), + "corrupted spill file size: expected=%d actual=%d", + f.bytes, + info.Size(), + ) + } + return nil +} + func (f *SpillFile) Close() error { if f == nil { return nil @@ -406,18 +481,12 @@ func (jm *JoinMap) SetMemoryRelease(release func()) { // // On error, ownership remains with the caller. func (jm *JoinMap) SetSpillBuildPayload(payload SpillBuildPayload) error { - if jm == nil || (len(payload.Files) == 0 && len(payload.LegacyFds) == 0) { + if jm == nil || len(payload.Files) == 0 { return ErrSpillBuildPayloadEmpty } - if len(payload.Files) > 0 && len(payload.LegacyFds) > 0 { - return ErrSpillBuildPayloadMixed - } - if len(payload.Files) > 0 && payload.BudgetRef == nil { + if payload.BudgetRef == nil { return ErrSpillBuildBudgetRef } - if len(payload.LegacyFds) > 0 && payload.BudgetRef != nil { - return ErrSpillBuildLegacyBudget - } jm.spillMu.Lock() defer jm.spillMu.Unlock() if jm.spillPayloadTaken { @@ -436,8 +505,8 @@ func (jm *JoinMap) SetSpillBuildPayload(payload SpillBuildPayload) error { return nil } -// TakeSpillBuildPayload atomically transfers files, legacy descriptors, and -// the producer budget generation to the sole spill consumer. +// TakeSpillBuildPayload atomically transfers files and the producer budget +// generation to the sole spill consumer. func (jm *JoinMap) TakeSpillBuildPayload() (SpillBuildPayload, error) { if jm == nil { return SpillBuildPayload{}, ErrSpillBuildPayloadEmpty @@ -515,25 +584,10 @@ func (jm *JoinMap) PreAlloc(n uint64) error { } type JoinMapMsg struct { - JoinMapPtr *JoinMap IsShuffle bool ShuffleIdx int32 Tag int32 - Spilled bool - // Result is the terminal dependency state. The zero value is retained for - // source compatibility with older direct JoinMapMsg literals; those are - // interpreted as an explicit successful result (including nil for an empty - // build) by terminalResult. - Result JoinMapResult -} - -func (t JoinMapMsg) terminalResult() JoinMapResult { - if t.Result.Finalized() { - return t.Result - } - // Legacy messages predate the explicit result field. A nil JoinMap in a - // legacy message is the established empty-build success convention. - return NewJoinMapResult(t.JoinMapPtr) + Result JoinMapResult } func (t JoinMapMsg) Serialize() []byte { @@ -549,10 +603,7 @@ func (t JoinMapMsg) NeedBlock() bool { } func (t JoinMapMsg) Destroy() { - jm := t.JoinMapPtr - if jm == nil && t.Result.IsSuccess() { - jm = t.Result.JoinMap() - } + jm := t.Result.JoinMap() if jm != nil { jm.FreeMemory() } @@ -568,13 +619,13 @@ func (t JoinMapMsg) DebugString() string { if t.IsShuffle { buf.WriteString("shuffle index " + strconv.Itoa(int(t.ShuffleIdx)) + "\n") } - if t.JoinMapPtr != nil { - buf.WriteString("joinmap rowcnt " + strconv.Itoa(int(t.JoinMapPtr.rowCnt)) + "\n") - buf.WriteString("joinmap refcnt " + strconv.Itoa(int(t.JoinMapPtr.GetRefCount())) + "\n") + if jm := t.Result.JoinMap(); jm != nil { + buf.WriteString("joinmap rowcnt " + strconv.Itoa(int(jm.rowCnt)) + "\n") + buf.WriteString("joinmap refcnt " + strconv.Itoa(int(jm.GetRefCount())) + "\n") } else if t.Result.IsBuildError() { buf.WriteString("joinmap build error " + t.Result.BuildError().Error() + "\n") } else { - buf.WriteString("joinmapPtr is nil \n") + buf.WriteString("joinmap is nil \n") } return buf.String() } @@ -620,7 +671,7 @@ func ReceiveJoinMapResult(tag int32, isShuffle bool, shuffleIdx int32, mb *Messa continue } } - result := msg.terminalResult() + result := msg.Result if !result.Finalized() { // A malformed/zero result must not be interpreted as empty. Keep // waiting for the producer's terminal publication. @@ -642,36 +693,24 @@ func ReceiveJoinMapResult(tag int32, isShuffle bool, shuffleIdx int32, mb *Messa } // SendJoinMapResult publishes one terminal dependency value without waiting -// for any consumer acknowledgement. The caller owns exactly-once admission -// (typically an atomic generation gate in HashBuild); this function only -// performs the non-blocking MessageBoard publication. -func SendJoinMapResult(result JoinMapResult, tag int32, isShuffle bool, shuffleIdx int32, mb *MessageBoard) { - if !result.Finalized() { - return +// for consumer acknowledgement. True means the MessageBoard accepted the +// value's ownership (or was already closed and destroyed it); false leaves +// ownership with the caller. +func SendJoinMapResult(result JoinMapResult, tag int32, isShuffle bool, shuffleIdx int32, mb *MessageBoard) bool { + if !result.Finalized() || mb == nil || mb.rwMutex == nil { + return false } msg := JoinMapMsg{ - JoinMapPtr: result.JoinMap(), IsShuffle: isShuffle, ShuffleIdx: shuffleIdx, Tag: tag, Result: result, } - if jm := result.JoinMap(); jm != nil { - msg.Spilled = jm.IsSpilled() - } SendMessage(msg, mb) + return true } // FinalizeJoinMapBuildError publishes a typed BuildError terminal value. -// It is kept separate from FinalizeJoinMapMessage so legacy nil-map empty -// build compatibility cannot accidentally turn an admission failure into a -// successful empty dependency. -func FinalizeJoinMapBuildError(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, err error) { - SendJoinMapResult(NewJoinMapBuildErrorResult(err), tag, isShuffle, shuffleIdx, mb) -} - -func FinalizeJoinMapMessage(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, sendMapSucceed bool) { - if !sendMapSucceed { - SendJoinMapResult(NewJoinMapResult(nil), tag, isShuffle, shuffleIdx, mb) - } +func FinalizeJoinMapBuildError(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, err error) bool { + return SendJoinMapResult(NewJoinMapBuildErrorResult(err), tag, isShuffle, shuffleIdx, mb) } diff --git a/pkg/vm/message/message.go b/pkg/vm/message/message.go index f8590dd81f33b..44eba0cda7299 100644 --- a/pkg/vm/message/message.go +++ b/pkg/vm/message/message.go @@ -75,6 +75,8 @@ type MessageCenter struct { type MessageBoard struct { reset bool // for debug purpose + closed bool + drained bool multiCN bool stmtId uuid.UUID messageCenter *MessageCenter @@ -98,10 +100,18 @@ func (m *MessageBoard) finalize() { } func (m *MessageBoard) DebugString() string { + if m == nil || m.rwMutex == nil { + return "messageBoard is nil\n" + } + m.rwMutex.RLock() + defer m.rwMutex.RUnlock() buf := bytes.NewBuffer(make([]byte, 0, 400)) if m.reset { buf.WriteString("messageBoard has been reseted!\n") } + if m.closed { + buf.WriteString("messageBoard has been closed!\n") + } if m.multiCN { buf.WriteString("messageBoard on MultiCN\n") } else { @@ -135,26 +145,99 @@ func (m *MessageBoard) BeforeRunonce() { // call this before runonce m.rwMutex.Lock() defer m.rwMutex.Unlock() - m.reset = false + if !m.closed { + m.reset = false + } } func (m *MessageBoard) Reset() *MessageBoard { - if m.multiCN { - m.messageCenter.RwMutex.Lock() - delete(m.messageCenter.StmtIDToBoard, m.stmtId) - m.messageCenter.RwMutex.Unlock() + m.rwMutex.RLock() + multiCN := m.multiCN + center := m.messageCenter + stmtID := m.stmtId + m.rwMutex.RUnlock() + if multiCN { + center.RwMutex.Lock() + delete(center.StmtIDToBoard, stmtID) + center.RwMutex.Unlock() // other pipeline could still access thie messageBoard // so reset current message board to a new one return NewMessageBoard() } m.rwMutex.Lock() defer m.rwMutex.Unlock() + if m.closed { + return NewMessageBoard() + } m.cleanupQueuedMessagesLocked() m.multiCN = false m.reset = true return m } +// Close prevents publication and wakes blocked receivers without destroying +// queued ownership. It is safe while consumers are still unwinding; the +// terminal owner must call CloseAndDrain after every producer and consumer is +// quiescent. True identifies the call that first closed the board. +func (m *MessageBoard) Close() bool { + return m.close(false) +} + +// CloseAndDrain is the terminal MessageBoard boundary for one execution +// attempt. Callers invoke it only after all scope and remote-notifier producers +// and consumers are quiescent. It removes a multi-CN registration, destroys +// every queued ownership-bearing message, and prevents a late producer from +// republishing into the closed generation. The operation is idempotent; true +// identifies the call that performed the drain, including a drain after an +// earlier Close. +func (m *MessageBoard) CloseAndDrain() bool { + return m.close(true) +} + +func (m *MessageBoard) close(drain bool) bool { + if m == nil || m.rwMutex == nil { + return false + } + m.rwMutex.RLock() + multiCN := m.multiCN + center := m.messageCenter + stmtID := m.stmtId + m.rwMutex.RUnlock() + if multiCN && center != nil { + center.RwMutex.Lock() + if center.StmtIDToBoard[stmtID] == m { + delete(center.StmtIDToBoard, stmtID) + } + center.RwMutex.Unlock() + } + + m.rwMutex.Lock() + defer m.rwMutex.Unlock() + firstClose := !m.closed + if firstClose { + m.closed = true + m.reset = true + for _, waiter := range m.waiters { + if waiter == nil { + continue + } + select { + case waiter <- true: + default: + } + } + } + if !drain { + return firstClose + } + if m.drained { + return false + } + m.drained = true + m.cleanupQueuedMessagesLocked() + return true +} + func (m *MessageBoard) cleanupQueuedMessages() { if m == nil || m.rwMutex == nil { return @@ -198,6 +281,11 @@ func NewMessageReceiver(tags []int32, addr MessageAddress, mb *MessageBoard) *Me func SendMessage(m Message, mb *MessageBoard) { if m.GetReceiverAddr().CnAddr == CURRENTCN { // message for current CN mb.rwMutex.Lock() + if mb.closed { + mb.rwMutex.Unlock() + m.Destroy() + return + } mb.messages = append(mb.messages, &m) if m.NeedBlock() { // broadcast for block message @@ -214,7 +302,7 @@ func SendMessage(m Message, mb *MessageBoard) { } } -func (mr *MessageReceiver) receiveMessageNonBlock() []Message { +func (mr *MessageReceiver) receiveMessageNonBlock() ([]Message, bool) { mr.mb.rwMutex.RLock() defer mr.mb.rwMutex.RUnlock() var result []Message @@ -235,25 +323,44 @@ func (mr *MessageReceiver) receiveMessageNonBlock() []Message { } } } - return result + return result, mr.mb.closed } func (mr *MessageReceiver) ReceiveMessage(needBlock bool, ctx context.Context) ([]Message, bool, error) { - var result = mr.receiveMessageNonBlock() + result, closed := mr.receiveMessageNonBlock() if !needBlock || len(result) > 0 { return result, false, nil } + if closed { + return result, false, moerr.NewInternalErrorNoCtx( + "message board is closed", + ) + } if mr.waiter == nil { mr.waiter = make(chan bool, 1) mr.mb.rwMutex.Lock() - mr.mb.waiters = append(mr.mb.waiters, mr.waiter) + if mr.mb.closed { + closed = true + } else { + mr.mb.waiters = append(mr.mb.waiters, mr.waiter) + } mr.mb.rwMutex.Unlock() + if closed { + return result, false, moerr.NewInternalErrorNoCtx( + "message board is closed", + ) + } } for { - result = mr.receiveMessageNonBlock() + result, closed = mr.receiveMessageNonBlock() if len(result) > 0 { break } + if closed { + return result, false, moerr.NewInternalErrorNoCtx( + "message board is closed", + ) + } timeout := messageTimeout if mr.debug { timeout = 1 * time.Second diff --git a/pkg/vm/message/message_test.go b/pkg/vm/message/message_test.go index 2a5c2adfcd450..07a87eb520dfd 100644 --- a/pkg/vm/message/message_test.go +++ b/pkg/vm/message/message_test.go @@ -15,6 +15,7 @@ package message import ( + "context" "os" "runtime" "runtime/debug" @@ -35,6 +36,32 @@ type testMessage struct { destroyed *atomic.Int32 } +type accountedTestMessage struct { + mp *mpool.MPool + buffer []byte +} + +func (m *accountedTestMessage) Serialize() []byte { return nil } + +func (m *accountedTestMessage) Deserialize([]byte) Message { return m } + +func (m *accountedTestMessage) NeedBlock() bool { return true } + +func (m *accountedTestMessage) GetMsgTag() int32 { return 1 } + +func (m *accountedTestMessage) GetReceiverAddr() MessageAddress { + return AddrBroadCastOnCurrentCN() +} + +func (m *accountedTestMessage) DebugString() string { return "accounted test message" } + +func (m *accountedTestMessage) Destroy() { + if m.buffer != nil { + m.mp.Free(m.buffer) + m.buffer = nil + } +} + func (m testMessage) Serialize() []byte { return nil } @@ -75,7 +102,7 @@ func TestJoinMapMsgDestroyReleasesJoinMapMemory(t *testing.T) { shm: shm, } - JoinMapMsg{JoinMapPtr: jm, Tag: 1}.Destroy() + JoinMapMsg{Result: NewJoinMapResult(jm), Tag: 1}.Destroy() require.Nil(t, jm.shm) require.False(t, jm.valid) @@ -97,6 +124,102 @@ func TestMessageBoardResetDestroysQueuedMessages(t *testing.T) { require.Empty(t, mb.waiters) } +func TestMessageBoardCloseAndDrainRejectsLateMessages(t *testing.T) { + var destroyed atomic.Int32 + mb := NewMessageBoard() + SendMessage(testMessage{tag: 1, destroyed: &destroyed}, mb) + + receiver := NewMessageReceiver( + []int32{2}, + AddrBroadCastOnCurrentCN(), + mb, + ) + waiting := make(chan error, 1) + go func() { + _, _, err := receiver.ReceiveMessage(true, context.Background()) + waiting <- err + }() + + require.True(t, mb.CloseAndDrain()) + require.False(t, mb.CloseAndDrain()) + require.ErrorContains(t, <-waiting, "message board is closed") + require.Equal(t, int32(1), destroyed.Load()) + require.Empty(t, mb.messages) + require.Empty(t, mb.waiters) + + SendMessage(testMessage{tag: 2, destroyed: &destroyed}, mb) + require.Equal(t, int32(2), destroyed.Load()) + require.NotSame(t, mb, mb.Reset()) +} + +func TestMessageBoardCloseAndDrainRemovesMultiCNRegistration(t *testing.T) { + center := &MessageCenter{ + StmtIDToBoard: make(map[uuid.UUID]*MessageBoard), + RwMutex: &sync.Mutex{}, + } + stmtID := uuid.New() + mb := NewMessageBoard().SetMultiCN(center, stmtID) + require.Same(t, mb, center.StmtIDToBoard[stmtID]) + + require.True(t, mb.CloseAndDrain()) + _, ok := center.StmtIDToBoard[stmtID] + require.False(t, ok) +} + +func TestMessageBoardCloseDefersQueuedOwnershipDrain(t *testing.T) { + mb := NewMessageBoard() + var destroyed atomic.Int32 + SendMessage(testMessage{tag: 1, destroyed: &destroyed}, mb) + + receiver := NewMessageReceiver( + []int32{2}, + AddrBroadCastOnCurrentCN(), + mb, + ) + waiting := make(chan error, 1) + go func() { + _, _, err := receiver.ReceiveMessage(true, context.Background()) + waiting <- err + }() + + require.True(t, mb.Close()) + require.False(t, mb.Close()) + require.ErrorContains(t, <-waiting, "message board is closed") + require.Zero(t, destroyed.Load()) + require.Len(t, mb.messages, 1) + + require.True(t, mb.CloseAndDrain()) + require.False(t, mb.CloseAndDrain()) + require.Equal(t, int32(1), destroyed.Load()) + require.Empty(t, mb.messages) +} + +func TestClosedMessageBoardLatePayloadDrainsOriginalGeneration(t *testing.T) { + registry, err := mpool.NewAllocationAccountRegistry(1, 1) + require.NoError(t, err) + account, err := registry.Open(1 << 20) + require.NoError(t, err) + mp := mpool.MustNewZero() + buffer, err := mp.AllocAccounted(64, account, 1, 1) + require.NoError(t, err) + + mb := NewMessageBoard() + require.True(t, mb.CloseAndDrain()) + terminal, first, err := registry.CompleteTerminal(account) + require.True(t, first) + require.ErrorIs(t, err, mpool.ErrAllocationAccountInvariant) + require.Equal(t, uint64(cap(buffer)), terminal.Used) + require.True(t, registry.AdmissionSuspended()) + + // A producer that already owns a payload cannot republish it after the + // attempt boundary. Destroy performs the physical Free against the account + // captured by the allocation, even though that generation is now sealed. + SendMessage(&accountedTestMessage{mp: mp, buffer: buffer}, mb) + require.False(t, registry.AdmissionSuspended()) + _, ok := registry.Resolve(account.Handle()) + require.False(t, ok) +} + func TestMessageBoardFinalizerDestroysQueuedMessages(t *testing.T) { var destroyed atomic.Int32 @@ -120,164 +243,6 @@ func TestMessageBoardFinalizerDestroysQueuedMessages(t *testing.T) { }, 5*time.Second, 20*time.Millisecond) } -func TestLegacySpillBuildPayload(t *testing.T) { - t.Run("transfers_ownership", func(t *testing.T) { - f1, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f1.Name()) - f2, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f2.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f1, f2}, - })) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - - require.Len(t, payload.LegacyFds, 2) - require.Same(t, f1, payload.LegacyFds[0]) - require.Same(t, f2, payload.LegacyFds[1]) - - require.NoError(t, payload.Close()) - }) - - t.Run("second_call_returns_explicit_error", func(t *testing.T) { - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f}, - })) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - - _, err = jm.TakeSpillBuildPayload() - require.ErrorIs(t, err, ErrSpillBuildPayloadTaken) - require.NoError(t, payload.Close()) - }) - - t.Run("empty_payload_is_rejected", func(t *testing.T) { - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.ErrorIs(t, jm.SetSpillBuildPayload(SpillBuildPayload{}), ErrSpillBuildPayloadEmpty) - }) - - t.Run("mixed_payload_is_rejected", func(t *testing.T) { - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.ErrorIs(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - Files: []*SpillFile{nil}, - LegacyFds: []*os.File{nil}, - BudgetRef: struct{}{}, - }), ErrSpillBuildPayloadMixed) - }) - - t.Run("legacy_payload_with_budget_is_rejected", func(t *testing.T) { - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - require.ErrorIs(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{nil}, - BudgetRef: struct{}{}, - }), ErrSpillBuildLegacyBudget) - }) - - t.Run("free_before_set_rejects_late_publication", func(t *testing.T) { - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mpool.MustNewZero()) - jm.IncRef(1) - jm.FreeMemory() - payload := SpillBuildPayload{LegacyFds: []*os.File{f}} - require.ErrorIs(t, jm.SetSpillBuildPayload(payload), ErrSpillBuildPayloadTaken) - _, err = f.Stat() - require.NoError(t, err, "rejected late payload remains caller-owned") - require.NoError(t, payload.Close()) - }) -} - -func TestFreeMemoryClosesSpillFds(t *testing.T) { - t.Run("closes_all_fds", func(t *testing.T) { - mp := mpool.MustNewZero() - f1, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f1.Name()) - f2, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f2.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mp) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f1, f2}, - })) - jm.FreeMemory() - - require.False(t, jm.valid) - - // Verify fds are closed - _, err = f1.Stat() - require.Error(t, err) - _, err = f2.Stat() - require.Error(t, err) - }) - - t.Run("handles_nil_in_fd_slice", func(t *testing.T) { - mp := mpool.MustNewZero() - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mp) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f, nil}, - })) - jm.FreeMemory() // must not panic on nil entry - }) - - t.Run("take_then_free_does_not_double_close", func(t *testing.T) { - mp := mpool.MustNewZero() - f, err := os.CreateTemp("", "test_fd_*") - require.NoError(t, err) - defer os.Remove(f.Name()) - - jm := NewJoinMap(GroupSels{}, nil, nil, nil, nil, mp) - jm.IncRef(1) - require.NoError(t, jm.SetSpillBuildPayload(SpillBuildPayload{ - LegacyFds: []*os.File{f}, - })) - payload, err := jm.TakeSpillBuildPayload() - require.NoError(t, err) - require.Len(t, payload.LegacyFds, 1) - - // FreeMemory after TakeSpillBuildPayload should not close the fds. - jm.FreeMemory() - - // fd is still open (caller owns it) - _, err = f.Stat() - require.NoError(t, err) - require.NoError(t, payload.Close()) - }) - - t.Run("double_free_safe", func(t *testing.T) { - mp := mpool.MustNewZero() - jm := &JoinMap{ - valid: true, - mpool: mp, - } - jm.FreeMemory() - jm.FreeMemory() // must not panic - }) -} - func TestAccountedSpillFileOwnership(t *testing.T) { var releases atomic.Int32 newFile := func() (*SpillFile, string) { diff --git a/pkg/vm/process/cte_memory_budget.go b/pkg/vm/process/cte_memory_budget.go index eb3de96d2cf3e..af68a441cd81f 100644 --- a/pkg/vm/process/cte_memory_budget.go +++ b/pkg/vm/process/cte_memory_budget.go @@ -113,15 +113,15 @@ func (r *CTEMemoryReservation) Resize(ctx context.Context, bytes uint64) error { return nil } -func (r *CTEMemoryReservation) Release() { +func (r *CTEMemoryReservation) Release() bool { if r == nil || r.budget == nil { - return + return false } b := r.budget b.mu.Lock() defer b.mu.Unlock() if !r.active { - return + return false } if !b.closed { if r.bytes <= b.used { @@ -132,6 +132,7 @@ func (r *CTEMemoryReservation) Release() { } r.bytes = 0 r.active = false + return true } func (r *CTEMemoryReservation) Bytes() uint64 { diff --git a/pkg/vm/process/hashbuild_budget.go b/pkg/vm/process/hashbuild_budget.go index dd8dba0548975..a5f06f216cdb2 100644 --- a/pkg/vm/process/hashbuild_budget.go +++ b/pkg/vm/process/hashbuild_budget.go @@ -15,6 +15,7 @@ package process import ( + "errors" "fmt" "math" "sync" @@ -30,6 +31,20 @@ import ( const hashBuildMinimumReserve = uint64(4 << 30) +const ( + hashBuildAllocationGenerationSlots = uint32(131_072) + // Account metadata lives outside the MPool payload cap. Bound its worst-case + // Go-heap footprint to half of hashBuildMinimumReserve: 128 bytes covers one + // sparse pointer/lease-map entry, and 16,777,216 live entries consume at most + // 2 GiB by construction. Small aggregate caps use the tighter byte + // conservation bound because every physical allocation owns at least one + // byte. Slot exhaustion is real metadata capacity pressure, not an + // estimator rejection. + hashBuildAllocationMetadataBytesPerSlot = uint64(128) + hashBuildAllocationMetadataHeadroom = hashBuildMinimumReserve / 2 + hashBuildAllocationMetadataMaxSlots = hashBuildAllocationMetadataHeadroom / hashBuildAllocationMetadataBytesPerSlot +) + const ( // Keep a process-wide reserve for listeners, RPC connections, object // storage, logs, and other descriptors that are not represented by the @@ -111,18 +126,12 @@ func observeHashBuildBudget(component, event, scope string, bytes uint64) { // process rather than the SQL layer so that operators and remote execution // code can make an admission decision without importing frontend packages. var ( - ErrHashBuildBudgetAdmission = moerr.NewInternalErrorNoCtx("hash build budget admission rejected") - // ErrHashBuildBudgetRejected is kept as a more discoverable spelling of the - // admission sentinel. It is the same value, so errors.Is works with either. - ErrHashBuildBudgetRejected = ErrHashBuildBudgetAdmission - ErrHashBuildBudgetClosed = moerr.NewInternalErrorNoCtx("hash build budget is closed") - ErrHashBuildBudgetInvalid = moerr.NewInternalErrorNoCtx("invalid hash build budget") - ErrHashBuildCeilingMissing = moerr.NewInternalErrorNoCtx("hash build budget ceiling unavailable") - ErrHashBuildBudgetUnavailable = ErrHashBuildCeilingMissing - ErrHashBuildReservationInactive = moerr.NewInternalErrorNoCtx("hash build reservation is inactive") - ErrHashBuildReservationUpward = moerr.NewInternalErrorNoCtx("hash build reservation reconciliation would increase charge") - ErrHashBuildReservationReconcileUpward = ErrHashBuildReservationUpward - ErrHashBuildReservationClosed = ErrHashBuildReservationInactive + ErrHashBuildBudgetAdmission = moerr.NewInternalErrorNoCtx("hash build budget admission rejected") + ErrHashBuildBudgetClosed = moerr.NewInternalErrorNoCtx("hash build budget is closed") + ErrHashBuildBudgetInvalid = moerr.NewInternalErrorNoCtx("invalid hash build budget") + ErrHashBuildCeilingMissing = moerr.NewInternalErrorNoCtx("hash build budget ceiling unavailable") + ErrHashBuildSpillReservationInactive = moerr.NewInternalErrorNoCtx("hash build spill reservation is inactive") + ErrHashBuildSpillReservationUpward = moerr.NewInternalErrorNoCtx("hash build spill reservation reconciliation would increase charge") ) // HashBuildBudgetErrorKind identifies the class of a budget error. @@ -135,16 +144,17 @@ const ( HashBuildBudgetErrorCeilingMissing ) -// HashBuildBudgetResource identifies the finite resource whose admission -// failed. It is separate from Kind: every resource can reject capacity, while -// lifecycle and accounting failures remain resource-independent. -type HashBuildBudgetResource uint8 +// HashBuildBudgetComponent identifies the independently bounded resource that +// rejected an admission. Zero is invalid: every admission error must name its +// physical resource. A spill-disk or spill-FD rejection must never enter the +// memory reclaim/reduce loop because reducing an in-memory batch cannot create +// either resource and may replay already-published spill records. +type HashBuildBudgetComponent uint8 const ( - HashBuildBudgetResourceUnknown HashBuildBudgetResource = iota - HashBuildBudgetResourceMemory - HashBuildBudgetResourceSpillDisk - HashBuildBudgetResourceSpillFD + HashBuildBudgetComponentMemory HashBuildBudgetComponent = iota + 1 + HashBuildBudgetComponentSpillDisk + HashBuildBudgetComponentSpillFD ) // HashBuildBudgetError carries bounded, observational details for an @@ -153,7 +163,7 @@ const ( // to inspect; they are never produced by overflowing arithmetic. type HashBuildBudgetError struct { Kind HashBuildBudgetErrorKind - Resource HashBuildBudgetResource + Component HashBuildBudgetComponent Requested uint64 Used uint64 Cap uint64 @@ -206,7 +216,7 @@ func (e *HashBuildBudgetError) Is(target error) bool { if e == nil { return false } - if target == ErrHashBuildBudgetAdmission || target == ErrHashBuildBudgetRejected { + if target == ErrHashBuildBudgetAdmission { return e.Kind == HashBuildBudgetErrorAdmission } switch e.Kind { @@ -257,6 +267,10 @@ type HashBuildBudget struct { spillFDConfiguredCap uint64 spillFDCap uint64 spillFDUsed uint64 + + allocationRegistryOnce sync.Once + allocationRegistry *mpool.AllocationAccountRegistry + allocationRegistryErr error } // NewHashBuildBudget creates a local-CN budget. Both caps are finite and @@ -407,10 +421,10 @@ func (b *HashBuildBudget) SetSpillCaps(diskBytes, fds uint64) error { } effectiveFDCap := clampSpillFDCap(fds, processLimit, limitKnown) if b.spillDiskUsed > diskBytes { - return newAdmissionError(HashBuildBudgetResourceSpillDisk, 0, b.spillDiskUsed, diskBytes) + return newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, 0, b.spillDiskUsed, diskBytes) } if b.spillFDUsed > effectiveFDCap { - return newAdmissionError(HashBuildBudgetResourceSpillFD, 0, b.spillFDUsed, effectiveFDCap) + return newComponentAdmissionError(HashBuildBudgetComponentSpillFD, 0, b.spillFDUsed, effectiveFDCap) } b.spillDiskCap = diskBytes b.spillFDConfiguredCap = fds @@ -454,17 +468,6 @@ func MustNewHashBuildBudget(aggregateCap, queryCap uint64) *HashBuildBudget { return b } -func NewHashBuildBudgetWithSpillCaps(aggregateCap, queryCap, spillDiskCap, spillFDCap uint64) (*HashBuildBudget, error) { - b, err := NewHashBuildBudget(aggregateCap, queryCap) - if err != nil { - return nil, err - } - if err = b.SetSpillCaps(spillDiskCap, spillFDCap); err != nil { - return nil, err - } - return b, nil -} - // AggregateCap returns the configured local-CN cap. func (b *HashBuildBudget) AggregateCap() uint64 { if b == nil { @@ -475,10 +478,6 @@ func (b *HashBuildBudget) AggregateCap() uint64 { return b.aggregateCap } -// CNHashCap is an alias useful to callers that describe the aggregate as the -// CN hash cap. -func (b *HashBuildBudget) CNHashCap() uint64 { return b.AggregateCap() } - // QueryCap returns the per-generation, per-target-CN cap. func (b *HashBuildBudget) QueryCap() uint64 { if b == nil { @@ -499,15 +498,6 @@ func (b *HashBuildBudget) AggregateUsed() uint64 { return b.aggregateUsed } -// CNHashUsed is an alias for AggregateUsed. -func (b *HashBuildBudget) CNHashUsed() uint64 { return b.AggregateUsed() } - -// Current is a concise alias for AggregateUsed. -func (b *HashBuildBudget) Current() uint64 { return b.AggregateUsed() } - -// Capacity is a concise alias for AggregateCap. -func (b *HashBuildBudget) Capacity() uint64 { return b.AggregateCap() } - // Closed reports whether no new generation or reservation may be opened. func (b *HashBuildBudget) Closed() bool { if b == nil { @@ -835,6 +825,8 @@ type HashBuildBudgetGeneration struct { peakUsed uint64 } +var _ mpool.AllocationCapacityController = (*HashBuildBudgetGeneration)(nil) + // HashBuildBudgetGenerationSnapshot is an immutable fixed-cardinality view. type HashBuildBudgetGenerationSnapshot struct { ID, Cap, Used, PeakUsed uint64 @@ -844,14 +836,6 @@ type HashBuildBudgetGenerationSnapshot struct { Closed bool } -// HashBuildGeneration is a shorter spelling retained for call sites. -type HashBuildGeneration = HashBuildBudgetGeneration - -// HashBuildQueryBudget makes the per-generation/per-target-CN scope explicit -// at call sites. It is an alias, so tokens and methods retain one ownership -// implementation. -type HashBuildQueryBudget = HashBuildBudgetGeneration - // OpenGeneration opens a per-statement execution generation. The budget's // query cap is copied by reference (and remains immutable), while used bytes // belong solely to the returned generation. @@ -962,10 +946,7 @@ func (b *HashBuildBudget) openProcessGeneration( } } - memoryCap := requestedMemoryCap - if memoryCap > b.aggregateCap { - memoryCap = b.aggregateCap - } + memoryCap := min(requestedMemoryCap, b.aggregateCap) if memoryCap == 0 { return nil, &HashBuildBudgetError{ Kind: HashBuildBudgetErrorInvalid, @@ -978,17 +959,12 @@ func (b *HashBuildBudget) openProcessGeneration( if spillDiskCap == 0 { spillDiskCap = defaultSpillCap(memoryCap) } - if spillDiskCap > b.spillDiskCap { - spillDiskCap = b.spillDiskCap - } - configuredFDCap := configuredSpillFDCap(memoryCap) - if configuredFDCap > b.spillFDConfiguredCap { - configuredFDCap = b.spillFDConfiguredCap - } - effectiveFDCap := configuredFDCap - if effectiveFDCap > b.spillFDCap { - effectiveFDCap = b.spillFDCap - } + spillDiskCap = min(spillDiskCap, b.spillDiskCap) + configuredFDCap := min( + configuredSpillFDCap(memoryCap), + b.spillFDConfiguredCap, + ) + effectiveFDCap := min(configuredFDCap, b.spillFDCap) return &HashBuildBudgetGeneration{ budget: b, @@ -1000,25 +976,6 @@ func (b *HashBuildBudget) openProcessGeneration( }, nil } -// OpenGenerationWithLimits is a compatibility spelling for explicit spill caps. -func (b *HashBuildBudget) OpenGenerationWithLimits(id, memoryCap, spillDiskCap, spillFDCap uint64) (*HashBuildBudgetGeneration, error) { - return b.OpenGenerationWithSpillCaps(id, memoryCap, spillDiskCap, spillFDCap) -} - -func (b *HashBuildBudget) OpenGenerationWithCapAndSpill(id, memoryCap, spillDiskCap, spillFDCap uint64) (*HashBuildBudgetGeneration, error) { - return b.OpenGenerationWithSpillCaps(id, memoryCap, spillDiskCap, spillFDCap) -} - -// NewGeneration is an alias for OpenGeneration. -func (b *HashBuildBudget) NewGeneration(id uint64) (*HashBuildBudgetGeneration, error) { - return b.OpenGeneration(id) -} - -// OpenQueryBudget is the explicit per-query-CN spelling of OpenGeneration. -func (b *HashBuildBudget) OpenQueryBudget(id uint64) (*HashBuildQueryBudget, error) { - return b.OpenGeneration(id) -} - // ID returns the execution generation identity. func (g *HashBuildBudgetGeneration) ID() uint64 { if g == nil { @@ -1035,12 +992,6 @@ func (g *HashBuildBudgetGeneration) Cap() uint64 { return g.cap } -// QueryCap returns this generation's query-CN cap. -func (g *HashBuildBudgetGeneration) QueryCap() uint64 { return g.Cap() } - -// Capacity is a concise alias for Cap. -func (g *HashBuildBudgetGeneration) Capacity() uint64 { return g.Cap() } - // Used reports bytes reserved by this generation. func (g *HashBuildBudgetGeneration) Used() uint64 { if g == nil || g.budget == nil { @@ -1051,6 +1002,19 @@ func (g *HashBuildBudgetGeneration) Used() uint64 { return g.used } +// Peak reports the maximum physically owned bytes observed by this +// generation. It is observational only; admission and release remain owned by +// AllocationAccount-backed MPool allocations. +func (g *HashBuildBudgetGeneration) Peak() uint64 { + return g.Snapshot().PeakUsed +} + +// RejectCount reports physical allocation-capacity rejections observed by +// this generation. +func (g *HashBuildBudgetGeneration) RejectCount() uint64 { + return g.Snapshot().RejectCount +} + func (g *HashBuildBudgetGeneration) SpillDiskCap() uint64 { if g == nil || g.budget == nil { return 0 @@ -1098,17 +1062,6 @@ func (g *HashBuildBudgetGeneration) Snapshot() HashBuildBudgetGenerationSnapshot } } -// Stats is an alias retained for observability call sites. -func (g *HashBuildBudgetGeneration) Stats() HashBuildBudgetGenerationSnapshot { return g.Snapshot() } -func (g *HashBuildBudgetGeneration) Peak() uint64 { return g.Snapshot().PeakUsed } -func (g *HashBuildBudgetGeneration) ReserveCount() uint64 { return g.Snapshot().ReserveCount } -func (g *HashBuildBudgetGeneration) RejectCount() uint64 { return g.Snapshot().RejectCount } -func (g *HashBuildBudgetGeneration) ReconcileCount() uint64 { return g.Snapshot().ReconcileCount } -func (g *HashBuildBudgetGeneration) ReleaseCount() uint64 { return g.Snapshot().ReleaseCount } - -// Current is a concise alias for Used. -func (g *HashBuildBudgetGeneration) Current() uint64 { return g.Used() } - // Closed reports whether this generation rejects new reservations. func (g *HashBuildBudgetGeneration) Closed() bool { if g == nil || g.budget == nil { @@ -1119,6 +1072,38 @@ func (g *HashBuildBudgetGeneration) Closed() bool { return g.closed || g.budget.closed } +// AllocationAccountRegistry returns the bounded CN-local registry shared by +// every HashBuild generation under this aggregate budget. The slot +// bound follows a conservation fact rather than a per-operator multiplier: +// every live allocation owns at least one byte and all accounts share the +// aggregate byte cap, so live metadata cannot exceed aggregate capacity. A +// second fixed bound reserves at most 2 GiB of the existing 4 GiB CN +// headroom at 128 bytes per metadata entry. The registry stores only the +// resulting scalar limit; it does not preallocate one object per slot. +func (g *HashBuildBudgetGeneration) AllocationAccountRegistry() ( + *mpool.AllocationAccountRegistry, + error, +) { + if g == nil || g.budget == nil { + return nil, ErrHashBuildBudgetInvalid + } + b := g.budget + b.allocationRegistryOnce.Do(func() { + capBytes := b.AggregateCap() + if capBytes == 0 { + b.allocationRegistryErr = ErrHashBuildBudgetInvalid + return + } + allocationSlots := min(capBytes, hashBuildAllocationMetadataMaxSlots) + b.allocationRegistry, b.allocationRegistryErr = + mpool.NewAllocationAccountRegistry( + hashBuildAllocationGenerationSlots, + allocationSlots, + ) + }) + return b.allocationRegistry, b.allocationRegistryErr +} + // Close rejects future reservations for this generation while allowing all // currently live tokens to release. It is idempotent. func (g *HashBuildBudgetGeneration) Close() { @@ -1130,12 +1115,56 @@ func (g *HashBuildBudgetGeneration) Close() { g.budget.mu.Unlock() } -// Reserve performs the required two-level sequence: charge CN aggregate, -// then charge query-CN. If query-CN rejects, aggregate is rolled back before -// returning, so callers never observe a partial reservation. -func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, error) { +// AcquireAllocationCapacity applies the HashBuild query/CN policy to a physical +// MPool allocation. It creates no independently releasable reservation token: +// the physical allocation lease is the sole release owner. +func (g *HashBuildBudgetGeneration) AcquireAllocationCapacity(size uint64) error { + if size == 0 { + return nil + } + err := g.acquireMemory(size) + if err == nil { + return nil + } + switch { + case errors.Is(err, ErrHashBuildBudgetClosed): + return errors.Join(mpool.ErrAllocationAccountSealed, err) + case errors.Is(err, ErrHashBuildBudgetAdmission): + return errors.Join(mpool.ErrAllocationAccountCapacity, err) + default: + return errors.Join(mpool.ErrAllocationAccountInvariant, err) + } +} + +// ReleaseAllocationCapacity is called only by physical MPool Free through the +// allocation account. A mismatch is an ownership invariant failure. +func (g *HashBuildBudgetGeneration) ReleaseAllocationCapacity(size uint64) { + if size == 0 { + return + } + if g == nil || g.budget == nil { + panic("nil hash build allocation capacity controller") + } + b := g.budget + b.mu.Lock() + if g.used < size || b.aggregateUsed < size { + b.mu.Unlock() + panic("hash build allocation capacity release underflow") + } + g.used -= size + b.aggregateUsed -= size + g.releaseCount++ + b.mu.Unlock() + observeHashBuildBudget("memory", "release", "query", size) + observeHashBuildBudget("memory", "release", "cn", size) +} + +// acquireMemory admits one physical MPool allocation. The allocation account +// is the only owner of the charge and releases it from MPool.Free; there is no +// parallel estimate/reservation token. +func (g *HashBuildBudgetGeneration) acquireMemory(size uint64) error { if g == nil || g.budget == nil { - return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Message: "nil hash build generation"} + return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Message: "nil hash build generation"} } b := g.budget // A closed budget/generation has a deterministic lifecycle result and does @@ -1146,7 +1175,7 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, observeHashBuildBudget("memory", "reject", "query", size) err := &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: size, Used: g.used, Cap: g.cap} b.mu.Unlock() - return nil, err + return err } // The common cached-cap path decides whether a refresh is needed and updates @@ -1156,16 +1185,19 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, if cached { if err != nil { b.mu.Unlock() - return nil, err + return err } - token, firstErr, aggregateRejected := g.reserveLocked(size, false) + firstErr, aggregateRejected := g.acquireMemoryLocked( + size, + false, + ) b.mu.Unlock() - if token != nil { + if firstErr == nil && !aggregateRejected { observeHashBuildBudget("memory", "reserve", "query", size) observeHashBuildBudget("memory", "reserve", "cn", size) } if !aggregateRejected { - return token, firstErr + return firstErr } } else { b.mu.Unlock() @@ -1175,17 +1207,20 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, var refreshed bool epoch, hasProvider, refreshed, err = b.refreshAggregateCap(false, 0) if err != nil { - return nil, err + return err } b.mu.Lock() - token, firstErr, aggregateRejected := g.reserveLocked(size, false) + firstErr, aggregateRejected := g.acquireMemoryLocked( + size, + false, + ) b.mu.Unlock() - if token != nil { + if firstErr == nil && !aggregateRejected { observeHashBuildBudget("memory", "reserve", "query", size) observeHashBuildBudget("memory", "reserve", "cn", size) } if !aggregateRejected { - return token, firstErr + return firstErr } if refreshed { hasProvider = false @@ -1197,32 +1232,39 @@ func (g *HashBuildBudgetGeneration) Reserve(size uint64) (*HashBuildReservation, // The epoch check turns concurrent retries into a single-flight operation. if hasProvider { if _, _, _, err = b.refreshAggregateCap(true, epoch); err != nil { - return nil, err + return err } } b.mu.Lock() - token, err, aggregateRejected := g.reserveLocked(size, true) + err, aggregateRejected := g.acquireMemoryLocked( + size, + true, + ) b.mu.Unlock() - if token != nil { + if err == nil && !aggregateRejected { observeHashBuildBudget("memory", "reserve", "query", size) observeHashBuildBudget("memory", "reserve", "cn", size) } if aggregateRejected { - return nil, err + return err } - return token, err + return err } -// reserveLocked attempts one memory reservation. b.mu must be held. The bool -// result identifies an aggregate-cap failure so Reserve can trigger a forced +// acquireMemoryLocked attempts one physical-memory admission. b.mu must be +// held. The bool result identifies an aggregate-cap failure so the caller can +// trigger a forced // live-ceiling refresh without counting a transient failure as a rejection. -func (g *HashBuildBudgetGeneration) reserveLocked(size uint64, recordAggregateReject bool) (*HashBuildReservation, error, bool) { +func (g *HashBuildBudgetGeneration) acquireMemoryLocked( + size uint64, + recordAggregateReject bool, +) (error, bool) { b := g.budget if b.closed || g.closed { g.rejectCount++ observeHashBuildBudget("memory", "reject", "query", size) - return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: size, Used: g.used, Cap: g.cap}, false + return &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: size, Used: g.used, Cap: g.cap}, false } // Check by subtraction rather than used+size: this is safe for // math.MaxUint64 and rejects every overflow-sized request. @@ -1231,7 +1273,7 @@ func (g *HashBuildBudgetGeneration) reserveLocked(size uint64, recordAggregateRe g.rejectCount++ observeHashBuildBudget("memory", "reject", "cn", size) } - return nil, newAdmissionError(HashBuildBudgetResourceMemory, size, b.aggregateUsed, b.aggregateCap), true + return newAdmissionError(size, b.aggregateUsed, b.aggregateCap), true } b.aggregateUsed += size if g.used > g.cap || size > g.cap-g.used { @@ -1239,150 +1281,34 @@ func (g *HashBuildBudgetGeneration) reserveLocked(size uint64, recordAggregateRe b.aggregateUsed -= size g.rejectCount++ observeHashBuildBudget("memory", "reject", "query", size) - return nil, newAdmissionError(HashBuildBudgetResourceMemory, size, g.used, g.cap), false + return newAdmissionError(size, g.used, g.cap), false } g.used += size g.reserveCount++ if g.used > g.peakUsed { g.peakUsed = g.used } - return &HashBuildReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil, false -} - -// TryReserve is a boolean convenience for admission-only call sites. -func (g *HashBuildBudgetGeneration) TryReserve(size uint64) bool { - t, err := g.Reserve(size) - if err != nil { - return false - } - // A TryReserve caller has no token to retain; immediately release it. Use - // Release rather than manually decrementing to preserve exactly-once state. - t.Release() - return true -} - -// Grow increases a live memory reservation atomically. It is used for the -// Shuffle emergency spill-scratch lease so retained copies cannot consume the -// memory required to recover from a later admission rejection. -func (r *HashBuildReservation) Grow(additional uint64) error { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return ErrHashBuildReservationInactive - } - if additional == 0 { - return nil - } - b := r.budget - b.mu.Lock() - if r.core.state.Load() != hashBuildReservationActive { - b.mu.Unlock() - return ErrHashBuildReservationInactive - } - if b.closed || r.generation.closed { - r.generation.rejectCount++ - observeHashBuildBudget("memory", "reject", "query", additional) - err := &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: additional, Used: r.generation.used, Cap: r.generation.cap} - b.mu.Unlock() - return err - } - - _, epoch, hasProvider, cached, err := b.aggregateCapRefreshDecisionLocked(false, 0) - if cached { - if err != nil { - b.mu.Unlock() - return err - } - firstErr, aggregateRejected := r.growLocked(additional, false) - b.mu.Unlock() - if firstErr == nil { - observeHashBuildBudget("memory", "reserve", "query", additional) - observeHashBuildBudget("memory", "reserve", "cn", additional) - } - if !aggregateRejected { - return firstErr - } - } else { - b.mu.Unlock() - var refreshed bool - epoch, hasProvider, refreshed, err = b.refreshAggregateCap(false, 0) - if err != nil { - return err - } - b.mu.Lock() - firstErr, aggregateRejected := r.growLocked(additional, false) - b.mu.Unlock() - if firstErr == nil { - observeHashBuildBudget("memory", "reserve", "query", additional) - observeHashBuildBudget("memory", "reserve", "cn", additional) - } - if !aggregateRejected { - return firstErr - } - if refreshed { - hasProvider = false - } - } - - if hasProvider { - if _, _, _, err = b.refreshAggregateCap(true, epoch); err != nil { - return err - } - } - b.mu.Lock() - err, aggregateRejected := r.growLocked(additional, true) - b.mu.Unlock() - if err == nil { - observeHashBuildBudget("memory", "reserve", "query", additional) - observeHashBuildBudget("memory", "reserve", "cn", additional) - } - if aggregateRejected { - return err - } - return err + return nil, false } -// growLocked attempts one memory reservation growth. b.mu must be held. -func (r *HashBuildReservation) growLocked(additional uint64, recordAggregateReject bool) (error, bool) { - b := r.budget - g := r.generation - if r.core.state.Load() != hashBuildReservationActive { - return ErrHashBuildReservationInactive, false - } - if b.closed || g.closed { - g.rejectCount++ - observeHashBuildBudget("memory", "reject", "query", additional) - return &HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed, Requested: additional, Used: g.used, Cap: g.cap}, false - } - if b.aggregateUsed > b.aggregateCap || additional > b.aggregateCap-b.aggregateUsed { - // Defer the counter/metric until the caller knows whether a forced - // refresh can make this transient failure admissible. - if recordAggregateReject { - g.rejectCount++ - observeHashBuildBudget("memory", "reject", "cn", additional) - } - return newAdmissionError(HashBuildBudgetResourceMemory, additional, b.aggregateUsed, b.aggregateCap), true - } - if g.used > g.cap || additional > g.cap-g.used { - g.rejectCount++ - observeHashBuildBudget("memory", "reject", "query", additional) - return newAdmissionError(HashBuildBudgetResourceMemory, additional, g.used, g.cap), false - } - if r.core.size > math.MaxUint64-additional { - return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Requested: additional, Message: "hash build reservation size overflow"}, false - } - b.aggregateUsed += additional - g.used += additional - r.core.size += additional - if g.used > g.peakUsed { - g.peakUsed = g.used - } - g.reserveCount++ - return nil, false +func newAdmissionError(requested, used, cap uint64) error { + return newComponentAdmissionError( + HashBuildBudgetComponentMemory, + requested, + used, + cap, + ) } -func newAdmissionError(resource HashBuildBudgetResource, requested, used, cap uint64) error { +func newComponentAdmissionError( + component HashBuildBudgetComponent, + requested uint64, + used uint64, + cap uint64, +) error { return &HashBuildBudgetError{ Kind: HashBuildBudgetErrorAdmission, - Resource: resource, + Component: component, Requested: requested, Used: used, Cap: cap, @@ -1390,168 +1316,30 @@ func newAdmissionError(resource HashBuildBudgetResource, requested, used, cap ui } } -// HashBuildReservation is an exactly-once ownership token for one charge in -// both the CN aggregate and its generation. State transitions are atomic: -// active -> released or active -> transferred. A late release therefore -// always affects the original generation and can never decrement a newer one. -type HashBuildReservation struct { - budget *HashBuildBudget - generation *HashBuildBudgetGeneration - // core is shared by accidental token copies, keeping mutable charge and - // exactly-once state together under the budget mutex. - core *hashBuildReservationCore -} - -type hashBuildReservationCore struct { +type hashBuildSpillReservationCore struct { size uint64 state atomic.Uint32 } const ( - hashBuildReservationActive uint32 = iota - hashBuildReservationReleased - hashBuildReservationTransferred + hashBuildSpillReservationActive uint32 = iota + hashBuildSpillReservationReleased ) -// Size returns the reservation's current reconciled charge. -func (r *HashBuildReservation) Size() uint64 { - if r == nil || r.budget == nil || r.core == nil { - return 0 - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - return r.core.size -} - -// GenerationID returns the generation charged by this token. -func (r *HashBuildReservation) GenerationID() uint64 { - if r == nil || r.generation == nil { - return 0 - } - return r.generation.id -} - -// Released reports whether this token has relinquished its ownership. A -// transferred token is not released, but no longer owns the charge. -func (r *HashBuildReservation) Released() bool { - if r == nil || r.core == nil { - return true - } - if r.budget == nil { - return r.core.state.Load() != hashBuildReservationActive - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - return r.core.state.Load() != hashBuildReservationActive -} - -// Release relinquishes this token once. It returns true only for the caller -// that won the active -> released transition. -func (r *HashBuildReservation) Release() bool { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false - } - r.budget.mu.Lock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationReleased) { - r.budget.mu.Unlock() - return false - } - size := r.core.size - // The subtraction is exact for a live token. Keep a defensive branch so - // corrupted state cannot underflow and turn into an apparent huge charge. - if r.generation.used >= size { - r.generation.used -= size - } else { - r.generation.used = 0 - } - if r.budget.aggregateUsed >= size { - r.budget.aggregateUsed -= size - } else { - r.budget.aggregateUsed = 0 - } - r.generation.releaseCount++ - r.budget.mu.Unlock() - observeHashBuildBudget("memory", "release", "query", size) - observeHashBuildBudget("memory", "release", "cn", size) - return true -} - -// ReconcileDown shrinks a live charge to actual bytes. It is linearized with -// reserve/release/transfer under the owning budget mutex. Upward reconciliation -// is rejected and inactive tokens never mutate counters. -func (r *HashBuildReservation) ReconcileDown(actual uint64) (bool, error) { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false, ErrHashBuildReservationInactive - } - r.budget.mu.Lock() - if r.core.state.Load() != hashBuildReservationActive { - r.budget.mu.Unlock() - return false, ErrHashBuildReservationInactive - } - if actual > r.core.size { - r.budget.mu.Unlock() - return false, ErrHashBuildReservationUpward - } - delta := r.core.size - actual - if delta > 0 { - if r.generation.used < delta || r.budget.aggregateUsed < delta { - r.budget.mu.Unlock() - return false, ErrHashBuildReservationInactive - } - r.generation.used -= delta - r.budget.aggregateUsed -= delta - r.core.size = actual - } - r.generation.reconcileCount++ - r.budget.mu.Unlock() - if delta > 0 { - observeHashBuildBudget("memory", "reconcile", "query", delta) - observeHashBuildBudget("memory", "reconcile", "cn", delta) - } - return true, nil -} - -// Reconcile is a compatibility alias. -func (r *HashBuildReservation) Reconcile(actual uint64) (bool, error) { return r.ReconcileDown(actual) } - -// Transfer moves ownership to a fresh token exactly once. The original token -// becomes inert; releasing it after a successful transfer cannot decrement the -// budget. If Release wins the race, Transfer returns nil. -func (r *HashBuildReservation) Transfer() *HashBuildReservation { - if r == nil || r.core == nil || r.budget == nil { - return nil - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationTransferred) { - return nil - } - return &HashBuildReservation{budget: r.budget, generation: r.generation, core: &hashBuildReservationCore{size: r.core.size}} -} - -// TransferOwnership is a descriptive alias for Transfer. -func (r *HashBuildReservation) TransferOwnership() *HashBuildReservation { return r.Transfer() } - -// TransferTo is another descriptive spelling for ownership transfer. -func (r *HashBuildReservation) TransferTo() *HashBuildReservation { return r.Transfer() } - // HashBuildSpillDiskReservation owns query and CN spill-disk bytes. type HashBuildSpillDiskReservation struct { budget *HashBuildBudget generation *HashBuildBudgetGeneration - core *hashBuildReservationCore + core *hashBuildSpillReservationCore } // HashBuildSpillFDReservation owns query and CN spill file descriptors. type HashBuildSpillFDReservation struct { budget *HashBuildBudget generation *HashBuildBudgetGeneration - core *hashBuildReservationCore + core *hashBuildSpillReservationCore } -type SpillDiskReservation = HashBuildSpillDiskReservation -type SpillFDReservation = HashBuildSpillFDReservation - func (r *HashBuildSpillDiskReservation) Size() uint64 { if r == nil || r.budget == nil || r.core == nil { return 0 @@ -1568,13 +1356,6 @@ func (r *HashBuildSpillFDReservation) Size() uint64 { defer r.budget.mu.Unlock() return r.core.size } -func (r *HashBuildSpillDiskReservation) Released() bool { - return r == nil || r.core == nil || r.core.state.Load() != hashBuildReservationActive -} -func (r *HashBuildSpillFDReservation) Released() bool { - return r == nil || r.core == nil || r.core.state.Load() != hashBuildReservationActive -} - func (g *HashBuildBudgetGeneration) ReserveSpillDisk(size uint64) (*HashBuildSpillDiskReservation, error) { if g == nil || g.budget == nil { return nil, &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid} @@ -1590,22 +1371,18 @@ func (g *HashBuildBudgetGeneration) ReserveSpillDisk(size uint64) (*HashBuildSpi if b.spillDiskUsed > b.spillDiskCap || size > b.spillDiskCap-b.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "cn", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillDisk, size, b.spillDiskUsed, b.spillDiskCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, size, b.spillDiskUsed, b.spillDiskCap) } if g.spillDiskUsed > g.spillDiskCap || size > g.spillDiskCap-g.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "query", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillDisk, size, g.spillDiskUsed, g.spillDiskCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, size, g.spillDiskUsed, g.spillDiskCap) } b.spillDiskUsed += size g.spillDiskUsed += size observeHashBuildBudget("spill_disk", "reserve", "query", size) observeHashBuildBudget("spill_disk", "reserve", "cn", size) - return &HashBuildSpillDiskReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil -} - -func (g *HashBuildBudgetGeneration) ReserveSpillDiskBytes(size uint64) (*HashBuildSpillDiskReservation, error) { - return g.ReserveSpillDisk(size) + return &HashBuildSpillDiskReservation{budget: b, generation: g, core: &hashBuildSpillReservationCore{size: size}}, nil } // Grow increases one live per-file disk reservation without allocating a new @@ -1613,7 +1390,7 @@ func (g *HashBuildBudgetGeneration) ReserveSpillDiskBytes(size uint64) (*HashBui // rather than to the number of tiny batch records written to those files. func (r *HashBuildSpillDiskReservation) Grow(additional uint64) error { if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return ErrHashBuildReservationInactive + return ErrHashBuildSpillReservationInactive } if additional == 0 { return nil @@ -1622,8 +1399,8 @@ func (r *HashBuildSpillDiskReservation) Grow(additional uint64) error { g := r.generation b.mu.Lock() defer b.mu.Unlock() - if r.core.state.Load() != hashBuildReservationActive { - return ErrHashBuildReservationInactive + if r.core.state.Load() != hashBuildSpillReservationActive { + return ErrHashBuildSpillReservationInactive } if b.closed || g.closed { g.rejectCount++ @@ -1633,12 +1410,12 @@ func (r *HashBuildSpillDiskReservation) Grow(additional uint64) error { if b.spillDiskUsed > b.spillDiskCap || additional > b.spillDiskCap-b.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "cn", additional) - return newAdmissionError(HashBuildBudgetResourceSpillDisk, additional, b.spillDiskUsed, b.spillDiskCap) + return newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, additional, b.spillDiskUsed, b.spillDiskCap) } if g.spillDiskUsed > g.spillDiskCap || additional > g.spillDiskCap-g.spillDiskUsed { g.rejectCount++ observeHashBuildBudget("spill_disk", "reject", "query", additional) - return newAdmissionError(HashBuildBudgetResourceSpillDisk, additional, g.spillDiskUsed, g.spillDiskCap) + return newComponentAdmissionError(HashBuildBudgetComponentSpillDisk, additional, g.spillDiskUsed, g.spillDiskCap) } if r.core.size > math.MaxUint64-additional { return &HashBuildBudgetError{Kind: HashBuildBudgetErrorInvalid, Requested: additional, Message: "spill disk reservation size overflow"} @@ -1676,22 +1453,18 @@ func (g *HashBuildBudgetGeneration) ReserveSpillFD(size uint64) (*HashBuildSpill if b.spillFDUsed > b.spillFDCap || size > b.spillFDCap-b.spillFDUsed { g.rejectCount++ observeHashBuildBudget("spill_fd", "reject", "cn", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillFD, size, b.spillFDUsed, b.spillFDCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillFD, size, b.spillFDUsed, b.spillFDCap) } if g.spillFDUsed > g.spillFDCap || size > g.spillFDCap-g.spillFDUsed { g.rejectCount++ observeHashBuildBudget("spill_fd", "reject", "query", size) - return nil, newAdmissionError(HashBuildBudgetResourceSpillFD, size, g.spillFDUsed, g.spillFDCap) + return nil, newComponentAdmissionError(HashBuildBudgetComponentSpillFD, size, g.spillFDUsed, g.spillFDCap) } b.spillFDUsed += size g.spillFDUsed += size observeHashBuildBudget("spill_fd", "reserve", "query", size) observeHashBuildBudget("spill_fd", "reserve", "cn", size) - return &HashBuildSpillFDReservation{budget: b, generation: g, core: &hashBuildReservationCore{size: size}}, nil -} - -func (g *HashBuildBudgetGeneration) ReserveSpillFileDescriptors(size uint64) (*HashBuildSpillFDReservation, error) { - return g.ReserveSpillFD(size) + return &HashBuildSpillFDReservation{budget: b, generation: g, core: &hashBuildSpillReservationCore{size: size}}, nil } func (r *HashBuildSpillDiskReservation) Release() bool { @@ -1700,19 +1473,15 @@ func (r *HashBuildSpillDiskReservation) Release() bool { } r.budget.mu.Lock() defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationReleased) { + if !r.core.state.CompareAndSwap(hashBuildSpillReservationActive, hashBuildSpillReservationReleased) { return false } - if r.generation.spillDiskUsed >= r.core.size { - r.generation.spillDiskUsed -= r.core.size - } else { - r.generation.spillDiskUsed = 0 - } - if r.budget.spillDiskUsed >= r.core.size { - r.budget.spillDiskUsed -= r.core.size - } else { - r.budget.spillDiskUsed = 0 + if r.generation.spillDiskUsed < r.core.size || + r.budget.spillDiskUsed < r.core.size { + panic("hash build spill disk reservation release underflow") } + r.generation.spillDiskUsed -= r.core.size + r.budget.spillDiskUsed -= r.core.size observeHashBuildBudget("spill_disk", "release", "query", r.core.size) observeHashBuildBudget("spill_disk", "release", "cn", r.core.size) r.generation.releaseCount++ @@ -1725,19 +1494,15 @@ func (r *HashBuildSpillFDReservation) Release() bool { } r.budget.mu.Lock() defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationReleased) { + if !r.core.state.CompareAndSwap(hashBuildSpillReservationActive, hashBuildSpillReservationReleased) { return false } - if r.generation.spillFDUsed >= r.core.size { - r.generation.spillFDUsed -= r.core.size - } else { - r.generation.spillFDUsed = 0 - } - if r.budget.spillFDUsed >= r.core.size { - r.budget.spillFDUsed -= r.core.size - } else { - r.budget.spillFDUsed = 0 + if r.generation.spillFDUsed < r.core.size || + r.budget.spillFDUsed < r.core.size { + panic("hash build spill fd reservation release underflow") } + r.generation.spillFDUsed -= r.core.size + r.budget.spillFDUsed -= r.core.size observeHashBuildBudget("spill_fd", "release", "query", r.core.size) observeHashBuildBudget("spill_fd", "release", "cn", r.core.size) r.generation.releaseCount++ @@ -1746,20 +1511,20 @@ func (r *HashBuildSpillFDReservation) Release() bool { func (r *HashBuildSpillDiskReservation) ReconcileDown(actual uint64) (bool, error) { if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false, ErrHashBuildReservationInactive + return false, ErrHashBuildSpillReservationInactive } r.budget.mu.Lock() defer r.budget.mu.Unlock() - if r.core.state.Load() != hashBuildReservationActive { - return false, ErrHashBuildReservationInactive + if r.core.state.Load() != hashBuildSpillReservationActive { + return false, ErrHashBuildSpillReservationInactive } if actual > r.core.size { - return false, ErrHashBuildReservationUpward + return false, ErrHashBuildSpillReservationUpward } delta := r.core.size - actual if delta > 0 { if r.generation.spillDiskUsed < delta || r.budget.spillDiskUsed < delta { - return false, ErrHashBuildReservationInactive + return false, ErrHashBuildSpillReservationInactive } r.generation.spillDiskUsed -= delta r.budget.spillDiskUsed -= delta @@ -1770,71 +1535,6 @@ func (r *HashBuildSpillDiskReservation) ReconcileDown(actual uint64) (bool, erro r.generation.reconcileCount++ return true, nil } -func (r *HashBuildSpillFDReservation) ReconcileDown(actual uint64) (bool, error) { - if r == nil || r.core == nil || r.budget == nil || r.generation == nil { - return false, ErrHashBuildReservationInactive - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if r.core.state.Load() != hashBuildReservationActive { - return false, ErrHashBuildReservationInactive - } - if actual > r.core.size { - return false, ErrHashBuildReservationUpward - } - delta := r.core.size - actual - if delta > 0 { - if r.generation.spillFDUsed < delta || r.budget.spillFDUsed < delta { - return false, ErrHashBuildReservationInactive - } - r.generation.spillFDUsed -= delta - r.budget.spillFDUsed -= delta - r.core.size = actual - observeHashBuildBudget("spill_fd", "reconcile", "query", delta) - observeHashBuildBudget("spill_fd", "reconcile", "cn", delta) - } - r.generation.reconcileCount++ - return true, nil -} -func (r *HashBuildSpillDiskReservation) Reconcile(actual uint64) (bool, error) { - return r.ReconcileDown(actual) -} -func (r *HashBuildSpillFDReservation) Reconcile(actual uint64) (bool, error) { - return r.ReconcileDown(actual) -} - -func (r *HashBuildSpillDiskReservation) Transfer() *HashBuildSpillDiskReservation { - if r == nil || r.core == nil || r.budget == nil { - return nil - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationTransferred) { - return nil - } - return &HashBuildSpillDiskReservation{budget: r.budget, generation: r.generation, core: &hashBuildReservationCore{size: r.core.size}} -} -func (r *HashBuildSpillFDReservation) Transfer() *HashBuildSpillFDReservation { - if r == nil || r.core == nil || r.budget == nil { - return nil - } - r.budget.mu.Lock() - defer r.budget.mu.Unlock() - if !r.core.state.CompareAndSwap(hashBuildReservationActive, hashBuildReservationTransferred) { - return nil - } - return &HashBuildSpillFDReservation{budget: r.budget, generation: r.generation, core: &hashBuildReservationCore{size: r.core.size}} -} -func (r *HashBuildSpillDiskReservation) TransferOwnership() *HashBuildSpillDiskReservation { - return r.Transfer() -} -func (r *HashBuildSpillFDReservation) TransferOwnership() *HashBuildSpillFDReservation { - return r.Transfer() -} -func (r *HashBuildSpillDiskReservation) TransferTo() *HashBuildSpillDiskReservation { - return r.Transfer() -} -func (r *HashBuildSpillFDReservation) TransferTo() *HashBuildSpillFDReservation { return r.Transfer() } // HashBuildCeilingInputs are the finite resource sources used by // ResolveHashBuildCeiling. A zero or math.MaxUint64 source means unavailable @@ -1918,18 +1618,6 @@ func ResolveHashBuildCeiling(in HashBuildCeilingInputs) (HashBuildCeiling, error }, nil } -// ResolveHashBuildBudget is a semantic alias used by budget initialization -// callers. -func ResolveHashBuildBudget(in HashBuildCeilingInputs) (HashBuildCeiling, error) { - return ResolveHashBuildCeiling(in) -} - -// NewHashBuildBudgetFromCeiling wires a resolved ceiling into the local-CN -// aggregate/generation budget. -func NewHashBuildBudgetFromCeiling(ceiling HashBuildCeiling) (*HashBuildBudget, error) { - return NewHashBuildBudget(ceiling.CNHashCap, ceiling.QueryCap) -} - // GetHashBuildBudget returns the statement generation shared by every child // process in this BaseProcess. Different top-level processes on the same CN // charge a shared aggregate budget. diff --git a/pkg/vm/process/hashbuild_budget_test.go b/pkg/vm/process/hashbuild_budget_test.go index 118c3fdace694..31cba8e4e55bf 100644 --- a/pkg/vm/process/hashbuild_budget_test.go +++ b/pkg/vm/process/hashbuild_budget_test.go @@ -16,50 +16,65 @@ package process import ( "errors" + "fmt" "math" - "os" - "os/exec" - "runtime" - "strconv" "sync" "sync/atomic" "testing" "time" - commonmpool "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/fileservice" ) -func TestHashBuildBudgetExactLimitAndOverflow(t *testing.T) { - b, err := NewHashBuildBudget(math.MaxUint64, math.MaxUint64) - if err != nil { - t.Fatal(err) +// testPhysicalAllocation exercises the capacity controller at the same +// acquire/free boundary used by MPool without exposing a second production +// reservation API. +type testPhysicalAllocation struct { + generation *HashBuildBudgetGeneration + size uint64 + released atomic.Bool +} + +func acquireTestPhysicalAllocation( + generation *HashBuildBudgetGeneration, + size uint64, +) (*testPhysicalAllocation, error) { + if err := generation.AcquireAllocationCapacity(size); err != nil { + return nil, err } - g, err := b.OpenGeneration(1) + return &testPhysicalAllocation{generation: generation, size: size}, nil +} + +func (a *testPhysicalAllocation) Release() bool { + if a == nil || a.generation == nil || !a.released.CompareAndSwap(false, true) { + return false + } + a.generation.ReleaseAllocationCapacity(a.size) + return true +} + +func TestHashBuildBudgetPhysicalAllocationLimit(t *testing.T) { + budget := MustNewHashBuildBudget(math.MaxUint64, math.MaxUint64) + generation, err := budget.OpenGeneration(1) if err != nil { t.Fatal(err) } - tok, err := g.Reserve(math.MaxUint64) + allocation, err := acquireTestPhysicalAllocation(generation, math.MaxUint64) if err != nil { t.Fatalf("exact limit rejected: %v", err) } - if got := b.AggregateUsed(); got != math.MaxUint64 { - t.Fatalf("aggregate used = %d, want max uint64", got) - } - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("limit+1 error = %v, want admission rejection", err) + if _, err = acquireTestPhysicalAllocation(generation, 1); !errors.Is(err, ErrHashBuildBudgetAdmission) { + t.Fatalf("limit+1 error = %v", err) } - if got := b.AggregateUsed(); got != math.MaxUint64 { - t.Fatalf("failed query reservation changed aggregate: %d", got) + if budget.AggregateUsed() != math.MaxUint64 || generation.Used() != math.MaxUint64 { + t.Fatal("failed admission changed the physical allocation ledger") } - if got := g.Used(); got != math.MaxUint64 { - t.Fatalf("failed query reservation changed generation: %d", got) + if !allocation.Release() || allocation.Release() { + t.Fatal("physical allocation must release exactly once") } - if !tok.Release() || tok.Release() { - t.Fatal("release must transition exactly once") - } - if b.AggregateUsed() != 0 || g.Used() != 0 { - t.Fatalf("released reservation remains: cn=%d query=%d", b.AggregateUsed(), g.Used()) + if budget.AggregateUsed() != 0 || generation.Used() != 0 { + t.Fatal("physical allocation release leaked capacity") } } @@ -73,20 +88,19 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { tests := []struct { name string - want HashBuildBudgetResource + want HashBuildBudgetComponent call func() error }{ { name: "memory", - want: HashBuildBudgetResourceMemory, + want: HashBuildBudgetComponentMemory, call: func() error { - _, reserveErr := g.Reserve(11) - return reserveErr + return g.AcquireAllocationCapacity(11) }, }, { name: "spill disk", - want: HashBuildBudgetResourceSpillDisk, + want: HashBuildBudgetComponentSpillDisk, call: func() error { _, reserveErr := g.ReserveSpillDisk(6) return reserveErr @@ -94,7 +108,7 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { }, { name: "spill fd", - want: HashBuildBudgetResourceSpillFD, + want: HashBuildBudgetComponentSpillFD, call: func() error { _, reserveErr := g.ReserveSpillFD(2) return reserveErr @@ -107,1778 +121,445 @@ func TestHashBuildBudgetAdmissionIdentifiesResource(t *testing.T) { if err := test.call(); !errors.As(err, &budgetErr) { t.Fatalf("error=%v, want typed admission", err) } - if budgetErr.Kind != HashBuildBudgetErrorAdmission || budgetErr.Resource != test.want { + if budgetErr.Kind != HashBuildBudgetErrorAdmission || budgetErr.Component != test.want { t.Fatalf("admission kind/resource=(%d,%d), want=(%d,%d)", - budgetErr.Kind, budgetErr.Resource, HashBuildBudgetErrorAdmission, test.want) + budgetErr.Kind, budgetErr.Component, HashBuildBudgetErrorAdmission, test.want) } }) } } -func TestHashBuildBudgetQueryRejectRollsBackCN(t *testing.T) { - b := MustNewHashBuildBudget(10, 4) - g1, _ := b.OpenGeneration(1) - g2, _ := b.OpenGeneration(2) - first, err := g1.Reserve(4) +func TestHashBuildBudgetAllocationAccountIsSoleOwner(t *testing.T) { + budget := MustNewHashBuildBudget(10, 10) + generation, err := budget.OpenGeneration(1) if err != nil { t.Fatal(err) } - if _, err = g2.Reserve(7); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("error = %v, want admission rejection", err) - } - if b.AggregateUsed() != 4 || g2.Used() != 0 { - t.Fatalf("query rejection did not roll back CN: cn=%d g2=%d", b.AggregateUsed(), g2.Used()) - } - first.Release() -} - -func TestHashBuildBudgetConcurrentReserveRelease(t *testing.T) { - const workers = 32 - b := MustNewHashBuildBudget(workers, workers) - gens := make([]*HashBuildBudgetGeneration, workers) - for i := range gens { - gens[i], _ = b.OpenGeneration(uint64(i + 1)) - } - start := make(chan struct{}) - acquired := make(chan *HashBuildReservation, workers) - var wg sync.WaitGroup - for i := range gens { - wg.Add(1) - go func(g *HashBuildBudgetGeneration) { - defer wg.Done() - <-start - tok, err := g.Reserve(1) - if err == nil { - acquired <- tok - } - }(gens[i]) - } - close(start) - wg.Wait() - if len(acquired) != workers { - t.Fatalf("acquired %d reservations, want %d", len(acquired), workers) - } - if b.AggregateUsed() != workers { - t.Fatalf("aggregate used = %d, want %d", b.AggregateUsed(), workers) - } - for i := 0; i < workers; i++ { - (<-acquired).Release() + registry, err := mpool.NewAllocationAccountRegistry(1, 2) + if err != nil { + t.Fatal(err) } - if b.AggregateUsed() != 0 { - t.Fatalf("aggregate used after release = %d", b.AggregateUsed()) + account, err := registry.OpenWithController(11, generation) + if err != nil { + t.Fatal(err) } -} + mp := mpool.MustNewZero() -func TestHashBuildBudgetTransferAndClose(t *testing.T) { - b := MustNewHashBuildBudget(8, 8) - g, _ := b.OpenGeneration(7) - tok, err := g.Reserve(3) + buffer, err := mp.AllocAccounted(10, account, 1, 1) if err != nil { t.Fatal(err) } - moved := tok.Transfer() - if moved == nil || tok.Release() { - t.Fatal("transfer must make original token inert") + if generation.Used() != 10 || account.Snapshot().Used != 10 { + t.Fatal("physical allocation was not charged exactly once") } - if b.AggregateUsed() != 3 { - t.Fatalf("transfer changed charge: %d", b.AggregateUsed()) + if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is(err, ErrHashBuildBudgetAdmission) || + !errors.Is(err, mpool.ErrAllocationAccountCapacity) { + t.Fatalf("capacity error = %v", err) } - g.Close() - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed generation error = %v", err) - } - if !moved.Release() || moved.Release() { - t.Fatal("transferred token release must be exactly once") + if generation.Used() != 10 || account.Snapshot().Used != 10 || + registry.LiveAllocationMetadata() != 1 { + t.Fatal("failed allocation changed account state") } - if b.AggregateUsed() != 0 { - t.Fatalf("live token release after close leaked: %d", b.AggregateUsed()) + + generation.Close() + if _, err = mp.AllocAccounted(1, account, 1, 1); !errors.Is(err, mpool.ErrAllocationAccountSealed) { + t.Fatalf("closed generation error = %v", err) } - b.Close() - if _, err = b.OpenGeneration(8); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget error = %v", err) + mp.Free(buffer) + if generation.Used() != 0 || account.Snapshot().Used != 0 { + t.Fatal("MPool.Free did not release the sole charge") } -} - -func TestHashBuildBudgetGenerationIsolation(t *testing.T) { - b := MustNewHashBuildBudget(8, 8) - old, _ := b.OpenGeneration(1) - oldToken, _ := old.Reserve(6) - old.Close() - newGeneration, _ := b.OpenGeneration(2) - newToken, err := newGeneration.Reserve(2) - if err != nil { + account.Seal() + if _, err = registry.Finalize(account); err != nil { t.Fatal(err) } - oldToken.Release() - if newGeneration.Used() != 2 || b.AggregateUsed() != 2 { - t.Fatalf("old release affected new generation: new=%d aggregate=%d", newGeneration.Used(), b.AggregateUsed()) - } - newToken.Release() } -func TestHashBuildBudgetCapReductionFailsClosedUntilRelease(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - g1, _ := b.OpenGenerationWithCap(1, 10) - owned, err := g1.Reserve(8) +func TestHashBuildAllocationAccountRegistryBounds(t *testing.T) { + budget := MustNewHashBuildBudget(16<<10, 16<<10) + first, err := budget.OpenGeneration(1) if err != nil { t.Fatal(err) } - if err = b.UpdateAggregateCap(6); err != nil { + registry, err := first.AllocationAccountRegistry() + if err != nil { t.Fatal(err) } - g2, _ := b.OpenGenerationWithCap(2, 6) - if _, err = g2.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("cap shrink did not fail closed: %v", err) + if registry.GenerationCapacity() != hashBuildAllocationGenerationSlots || + registry.MaxAllocationMetadata() != 16<<10 { + t.Fatal("allocation registry does not follow byte-conservation bounds") } - owned.Release() - newToken, err := g2.Reserve(6) - if err != nil { - t.Fatalf("reservation after release failed: %v", err) - } - newToken.Release() -} - -func TestHashBuildReservationReconcileCopyAlias(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - g, _ := b.OpenGeneration(1) - tok, err := g.Reserve(10) + second, err := budget.OpenGeneration(2) if err != nil { t.Fatal(err) } - alias := *tok - if err := tok.Grow(2); err != nil { - t.Fatalf("grow: %v", err) + shared, err := second.AllocationAccountRegistry() + if err != nil || shared != registry { + t.Fatal("one CN budget created multiple allocation registries") } - if tok.Size() != 12 || alias.Size() != 12 || g.Used() != 12 { - t.Fatalf("alias grow diverged: token=%d alias=%d used=%d", tok.Size(), alias.Size(), g.Used()) - } - if ok, err := alias.ReconcileDown(4); !ok || err != nil { - t.Fatalf("reconcile: ok=%v err=%v", ok, err) - } - if tok.Size() != 4 || g.Used() != 4 || b.AggregateUsed() != 4 { - t.Fatalf("alias reconcile diverged: size=%d gen=%d cn=%d", tok.Size(), g.Used(), b.AggregateUsed()) - } - if _, err := tok.ReconcileDown(5); !errors.Is(err, ErrHashBuildReservationUpward) { - t.Fatalf("upward err=%v", err) - } - if !tok.Release() || alias.Release() { - t.Fatal("copy aliases must release exactly once") - } -} -func TestHashBuildReservationGrowRejectsWithoutChangingCharge(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - g, _ := b.OpenGeneration(1) - tok, err := g.Reserve(8) + large := MustNewHashBuildBudget(hashBuildAllocationMetadataMaxSlots+1, hashBuildAllocationMetadataMaxSlots+1) + largeGeneration, err := large.OpenGeneration(1) if err != nil { t.Fatal(err) } - before := g.Snapshot() - if err = tok.Grow(3); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("grow rejection=%v", err) - } - after := g.Snapshot() - if tok.Size() != 8 || after.Used != before.Used || b.AggregateUsed() != 8 { - t.Fatalf("rejected grow changed charge: token=%d generation=%d aggregate=%d", tok.Size(), after.Used, b.AggregateUsed()) - } - if after.RejectCount != before.RejectCount+1 { - t.Fatalf("reject count=%d, want %d", after.RejectCount, before.RejectCount+1) - } - tok.Release() -} - -func TestHashBuildReservationGrowHonorsInactiveClosedAndLiveCap(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - cap := uint64(20) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - b.SetAggregateCapProvider(func() (uint64, error) { return cap, nil }) - tok, err := g.Reserve(8) + largeRegistry, err := largeGeneration.AllocationAccountRegistry() if err != nil { t.Fatal(err) } - cap = 8 - now = now.Add(hashBuildBudgetCapRefreshTTL + time.Nanosecond) - if err = tok.Grow(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("live-cap grow=%v", err) - } - if tok.Size() != 8 || g.Used() != 8 || b.AggregateUsed() != 8 { - t.Fatalf("live-cap rejection changed charge") - } - tok.Release() - if err = tok.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("released grow=%v", err) + if largeRegistry.MaxAllocationMetadata() != hashBuildAllocationMetadataMaxSlots { + t.Fatal("allocation metadata exceeded its fixed headroom") } +} - cap = 20 - closed, err := g.Reserve(2) +func TestHashBuildBudgetQueryRejectRollsBackCN(t *testing.T) { + budget := MustNewHashBuildBudget(10, 4) + first, _ := budget.OpenGeneration(1) + second, _ := budget.OpenGeneration(2) + allocation, err := acquireTestPhysicalAllocation(first, 4) if err != nil { t.Fatal(err) } - g.Close() - if err = closed.Grow(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed grow=%v", err) + if _, err = acquireTestPhysicalAllocation(second, 7); !errors.Is(err, ErrHashBuildBudgetAdmission) { + t.Fatalf("query rejection error = %v", err) } - closed.Release() + if budget.AggregateUsed() != 4 || second.Used() != 0 { + t.Fatal("query rejection did not roll back the CN charge") + } + allocation.Release() } -func TestHashBuildReservationGrowConcurrentTerminalTransitions(t *testing.T) { - for i := 0; i < 100; i++ { - b := MustNewHashBuildBudget(64, 64) - g, _ := b.OpenGeneration(uint64(i + 1)) - tok, err := g.Reserve(10) +func TestHashBuildBudgetConcurrentPhysicalAllocations(t *testing.T) { + const workers = 64 + budget := MustNewHashBuildBudget(workers, workers) + allocations := make(chan *testPhysicalAllocation, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + generation, err := budget.OpenGeneration(uint64(i + 1)) if err != nil { t.Fatal(err) } - var wg sync.WaitGroup - wg.Add(3) - movedC := make(chan *HashBuildReservation, 1) - go func() { - defer wg.Done() - _ = tok.Grow(5) - }() - go func() { - defer wg.Done() - _, _ = tok.ReconcileDown(4) - }() + wg.Add(1) go func() { defer wg.Done() - movedC <- tok.Transfer() + allocation, acquireErr := acquireTestPhysicalAllocation(generation, 1) + if acquireErr != nil { + t.Errorf("acquire: %v", acquireErr) + return + } + allocations <- allocation }() - wg.Wait() - close(movedC) - tok.Release() - if moved := <-movedC; moved != nil { - moved.Release() - } - if g.Used() != 0 || b.AggregateUsed() != 0 { - t.Fatalf("iteration %d leaked charge: generation=%d aggregate=%d", i, g.Used(), b.AggregateUsed()) - } - } -} - -func TestHashBuildSpillLedgersTransferReconcile(t *testing.T) { - b := MustNewHashBuildBudget(64, 64) - g, _ := b.OpenGeneration(1) - disk, err := g.ReserveSpillDisk(100) - if err != nil { - t.Fatal(err) - } - fd, err := g.ReserveSpillFD(2) - if err != nil { - t.Fatal(err) - } - if b.SpillDiskUsed() != 100 || b.SpillFDUsed() != 2 { - t.Fatalf("used disk=%d fd=%d", b.SpillDiskUsed(), b.SpillFDUsed()) - } - if err := disk.Grow(25); err != nil { - t.Fatalf("disk grow: %v", err) } - if disk.Size() != 125 || b.SpillDiskUsed() != 125 { - t.Fatalf("grown disk token=%d used=%d", disk.Size(), b.SpillDiskUsed()) - } - if ok, err := disk.ReconcileDown(40); !ok || err != nil { - t.Fatalf("disk reconcile: %v %v", ok, err) - } - moved := fd.Transfer() - if moved == nil || fd.Release() { - t.Fatal("fd transfer") - } - g.Close() - if _, err := g.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed spill reserve=%v", err) - } - disk.Release() - moved.Release() - if b.SpillDiskUsed() != 0 || b.SpillFDUsed() != 0 { - t.Fatalf("spill leak disk=%d fd=%d", b.SpillDiskUsed(), b.SpillFDUsed()) - } -} - -func TestConfiguredSpillFDCapCushionsFirstShuffleRepartitionPeak(t *testing.T) { - const firstRepartitionPeak = uint64(16 * (64 + 64)) - if got := configuredSpillFDCap(192 << 20); got < firstRepartitionPeak { - t.Fatalf("configured spill fd cap=%d, want at least first 16-way repartition peak=%d", got, firstRepartitionPeak) + wg.Wait() + close(allocations) + if budget.AggregateUsed() != workers { + t.Fatalf("aggregate used = %d", budget.AggregateUsed()) } -} - -func TestClampSpillFDCapBoundaries(t *testing.T) { - for _, tc := range []struct { - name string - configured, processLimit uint64 - limitKnown bool - want uint64 - }{ - {name: "unknown fails closed", configured: 2048, processLimit: 1 << 20, want: 0}, - {name: "zero configured", configured: 0, processLimit: 1024, limitKnown: true, want: 0}, - {name: "below absolute headroom", configured: 2048, processLimit: 63, limitKnown: true, want: 0}, - {name: "at absolute headroom", configured: 2048, processLimit: 64, limitKnown: true, want: 0}, - {name: "one fd above headroom", configured: 2048, processLimit: 65, limitKnown: true, want: 1}, - {name: "absolute headroom dominates", configured: 2048, processLimit: 128, limitKnown: true, want: 64}, - {name: "quarter headroom dominates", configured: 2048, processLimit: 1024, limitKnown: true, want: 768}, - {name: "explicit finite cap retained", configured: 10, processLimit: 1024, limitKnown: true, want: 10}, - {name: "unlimited retains configured", configured: 2048, processLimit: math.MaxUint64, limitKnown: true, want: 2048}, - } { - t.Run(tc.name, func(t *testing.T) { - if got := clampSpillFDCap(tc.configured, tc.processLimit, tc.limitKnown); got != tc.want { - t.Fatalf("clampSpillFDCap(%d, %d, %v)=%d, want %d", - tc.configured, tc.processLimit, tc.limitKnown, got, tc.want) - } - }) + for allocation := range allocations { + allocation.Release() } -} - -func TestDefaultSpillFDCapMatchesProcessLimit(t *testing.T) { - limit, ok := processOpenFileLimit() - want := clampSpillFDCap(configuredSpillFDCap(192<<20), limit, ok) - b := MustNewHashBuildBudget(192<<20, 192<<20) - if got := b.SpillFDCap(); got != want { - t.Fatalf("spill fd cap=%d, want process-clamped cap=%d (limit=%d known=%v)", got, want, limit, ok) + if budget.AggregateUsed() != 0 { + t.Fatal("concurrent physical allocations leaked") } } -func TestHashBuildSpillFDCapUnderRLIMIT(t *testing.T) { - const ( - childEnv = "MO_HASHBUILD_RLIMIT_CHILD" - limitEnv = "MO_HASHBUILD_RLIMIT_NOFILE" - ) - if os.Getenv(childEnv) == "1" { - limit, ok := processOpenFileLimit() - if !ok { - t.Fatal("RLIMIT_NOFILE unavailable in RLIMIT child") - } - rawTarget := os.Getenv(limitEnv) - target, err := strconv.ParseUint(rawTarget, 10, 64) - if err != nil { - t.Fatalf("parse target %q: %v", rawTarget, err) - } - if limit != target { - t.Fatalf("child RLIMIT_NOFILE=%d, want %d", limit, target) - } - - configured := configuredSpillFDCap(192 << 20) - want := clampSpillFDCap(configured, limit, true) - b := MustNewHashBuildBudget(192<<20, 192<<20) - if got := b.SpillFDCap(); got != want { - t.Fatalf("child spill fd cap=%d, want %d", got, want) - } - g, err := b.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - // Simulate a budget/generation opened while the process limit was - // higher. ReserveSpillFD must sample the current RLIMIT again instead - // of trusting these stale effective caps. - b.mu.Lock() - b.spillFDCap = configured - g.spillFDCap = configured - b.mu.Unlock() - if _, err = g.ReserveSpillFD(want + 1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("RLIMIT+headroom overflow error=%v, want admission rejection", err) - } - if got := b.SpillFDCap(); got != want { - t.Fatalf("runtime preflight refreshed spill fd cap=%d, want %d", got, want) - } - token, err := g.ReserveSpillFD(want) - if err != nil { - t.Fatalf("exact safe FD cap rejected: %v", err) - } - if !token.Release() { - t.Fatal("exact safe FD reservation did not release") - } - - if err = b.SetSpillCaps(0, 10); err != nil { - t.Fatal(err) - } - wantExplicit := clampSpillFDCap(10, limit, true) - if got := b.SpillFDCap(); got != wantExplicit { - t.Fatalf("explicit finite FD cap=%d, want process-clamped %d", got, wantExplicit) - } - return - } - - switch runtime.GOOS { - case "darwin", "linux": - default: - t.Skip("RLIMIT_NOFILE subprocess is only supported on Darwin and Linux") - } - parentLimit, ok := processOpenFileLimit() - if !ok || parentLimit < hashBuildNonSpillFDHeadroom+1 { - t.Skipf("parent RLIMIT_NOFILE=%d known=%v is too small for isolated child test", parentLimit, ok) - } - target := uint64(128) - if parentLimit < target { - target = parentLimit - } - targetText := strconv.FormatUint(target, 10) - cmd := exec.Command( - "/bin/sh", "-c", - `ulimit -S -n "$MO_HASHBUILD_RLIMIT_NOFILE" && -ulimit -H -n "$MO_HASHBUILD_RLIMIT_NOFILE" && -exec "$@"`, - "sh", os.Args[0], "-test.run=^TestHashBuildSpillFDCapUnderRLIMIT$", "-test.count=1", - ) - cmd.Env = append(os.Environ(), childEnv+"=1", limitEnv+"="+targetText) - if output, err := cmd.CombinedOutput(); err != nil { - t.Fatalf("RLIMIT child failed: %v\n%s", err, output) +func TestHashBuildBudgetGenerationIsolationAndClose(t *testing.T) { + budget := MustNewHashBuildBudget(8, 8) + oldGeneration, _ := budget.OpenGeneration(1) + oldAllocation, _ := acquireTestPhysicalAllocation(oldGeneration, 6) + oldGeneration.Close() + if _, err := acquireTestPhysicalAllocation(oldGeneration, 1); !errors.Is(err, ErrHashBuildBudgetClosed) { + t.Fatalf("closed generation error = %v", err) } -} - -func TestHashBuildBudgetLiveCapProviderShrinksOpenGeneration(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGenerationWithCap(1, 10) - cap := uint64(10) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - b.SetAggregateCapProvider(func() (uint64, error) { return cap, nil }) - owned, err := g.Reserve(6) + newGeneration, _ := budget.OpenGeneration(2) + newAllocation, err := acquireTestPhysicalAllocation(newGeneration, 2) if err != nil { t.Fatal(err) } - cap = 5 - now = now.Add(hashBuildBudgetCapRefreshTTL + time.Nanosecond) - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("open generation ignored live cap shrink: %v", err) + oldAllocation.Release() + if newGeneration.Used() != 2 || budget.AggregateUsed() != 2 { + t.Fatal("old generation release affected the new generation") } - owned.Release() - token, err := g.Reserve(5) - if err != nil { - t.Fatalf("reservation at refreshed cap failed: %v", err) + newAllocation.Release() + budget.Close() + if _, err = budget.OpenGeneration(3); !errors.Is(err, ErrHashBuildBudgetClosed) { + t.Fatalf("closed budget error = %v", err) } - token.Release() } -func TestHashBuildBudgetCapProviderCachesWithinTTLAndRefreshes(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - cap := uint64(10) - b.SetAggregateCapProvider(func() (uint64, error) { +func TestHashBuildBudgetLiveCapRefresh(t *testing.T) { + budget := MustNewHashBuildBudget(10, 10) + var capValue atomic.Uint64 + capValue.Store(10) + var calls atomic.Uint64 + budget.SetAggregateCapProvider(func() (uint64, error) { calls.Add(1) - return cap, nil + return capValue.Load(), nil }) - first, err := g.Reserve(4) + budget.capRefreshTTL = time.Hour + generation, _ := budget.OpenGeneration(1) + first, err := acquireTestPhysicalAllocation(generation, 8) if err != nil { t.Fatal(err) } - if got := calls.Load(); got != 1 { - t.Fatalf("first reservation provider calls=%d, want 1", got) + if calls.Load() != 1 { + t.Fatalf("provider calls = %d", calls.Load()) } - now = now.Add(hashBuildBudgetCapRefreshTTL - time.Nanosecond) - second, err := g.Reserve(1) + capValue.Store(6) + cachedAllocation, err := acquireTestPhysicalAllocation(generation, 1) if err != nil { - t.Fatal(err) + t.Fatalf("cached cap should remain valid: %v", err) } - if got := calls.Load(); got != 1 { - t.Fatalf("TTL reservation provider calls=%d, want 1", got) - } - cap = 3 - now = now.Add(2 * time.Nanosecond) - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("expired shrink reservation=%v, want admission rejection", err) - } - if got := calls.Load(); got != 2 { - t.Fatalf("expired reservation provider calls=%d, want 2", got) + if calls.Load() != 1 { + t.Fatal("cached fast path sampled the provider") } first.Release() - second.Release() -} - -func TestHashBuildBudgetCachedFastPathSkipsRefreshGate(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = time.Hour - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - b.SetAggregateCapProvider(func() (uint64, error) { return 10, nil }) - if _, _, refreshed, err := b.refreshAggregateCap(false, 0); err != nil || !refreshed { - t.Fatalf("seed refresh: refreshed=%v err=%v", refreshed, err) - } - - b.refreshMu.Lock() - resultC := make(chan error, 1) - go func() { - _, _, refreshed, err := b.refreshAggregateCap(false, 0) - if err == nil && refreshed { - err = errors.New("cached refresh unexpectedly sampled provider") - } - resultC <- err - }() - - var err error - select { - case err = <-resultC: - b.refreshMu.Unlock() - case <-time.After(time.Second): - b.refreshMu.Unlock() - <-resultC - t.Fatal("cached refresh waited for refreshMu") + cachedAllocation.Release() + // Invalidate the cache to model the next observation interval. + budget.mu.Lock() + budget.capCached = false + budget.mu.Unlock() + if _, err = acquireTestPhysicalAllocation(generation, 7); !errors.Is(err, ErrHashBuildBudgetAdmission) { + t.Fatalf("shrunk live cap error = %v", err) } - if err != nil { - t.Fatal(err) + if budget.AggregateCap() != 6 { + t.Fatalf("aggregate cap = %d", budget.AggregateCap()) } } -func TestHashBuildBudgetUnchangedObservationSkipsRefreshGate(t *testing.T) { - inputs := HashBuildCeilingInputs{ - CgroupMemoryMax: 8 << 30, - HostMemTotal: 16 << 30, - GlobalMpoolCap: 6 << 30, - FileCacheHint: 512 << 20, - } - ceiling, err := ResolveHashBuildCeiling(inputs) +func TestHashBuildBudgetSpillLedgers(t *testing.T) { + budget := MustNewHashBuildBudget(64, 64) + generation, _ := budget.OpenGeneration(1) + disk, err := generation.ReserveSpillDisk(100) if err != nil { t.Fatal(err) } - b := MustNewHashBuildBudget(ceiling.CNHashCap, ceiling.CNHashCap) - b.installCNCapProvider(inputs) - - b.refreshMu.Lock() - doneC := make(chan struct{}) - go func() { - b.mergeObservedCNCap(inputs, ceiling.CNHashCap) - close(doneC) - }() - select { - case <-doneC: - b.refreshMu.Unlock() - case <-time.After(time.Second): - b.refreshMu.Unlock() - <-doneC - t.Fatal("unchanged CN cap observation waited for refreshMu") - } -} - -func TestHashBuildBudgetCNProviderKeepsProcessMemorySnapshot(t *testing.T) { - previousInputs := hashBuildProcessMemoryInputs - previousCap := commonmpool.GlobalCap() - previousHint := fileservice.GlobalMemoryCacheSizeHint.Swap(0) - t.Cleanup(func() { - hashBuildProcessMemoryInputs = previousInputs - commonmpool.InitCap(previousCap) - fileservice.GlobalMemoryCacheSizeHint.Store(previousHint) - }) - - hashBuildProcessMemoryInputs = HashBuildCeilingInputs{ - CgroupMemoryMax: 8 << 30, - HostMemTotal: 16 << 30, - } - commonmpool.InitCap(commonmpool.PB) - - b := MustNewHashBuildBudget(4<<30, 4<<30) - b.installCNCapProvider(HashBuildCeilingInputs{ - CgroupMemoryMax: 4 << 30, - HostMemTotal: 8 << 30, - }) - if _, err := b.sampleCNCap(); err != nil { + fd, err := generation.ReserveSpillFD(2) + if err != nil { t.Fatal(err) } - if b.liveCapInputs.CgroupMemoryMax != 8<<30 || b.liveCapInputs.HostMemTotal != 16<<30 { - t.Fatalf("physical snapshot changed: %+v", b.liveCapInputs) + if err = disk.Grow(25); err != nil || disk.Size() != 125 { + t.Fatalf("disk grow: size=%d err=%v", disk.Size(), err) } -} - -func TestHashBuildBudgetCapProviderGrowthOnAggregateReject(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - g2, _ := b.OpenGeneration(2) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - cap := uint64(10) - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return cap, nil - }) - owned, err := g.Reserve(10) - if err != nil { - t.Fatal(err) + if ok, reconcileErr := disk.ReconcileDown(40); !ok || reconcileErr != nil { + t.Fatalf("disk reconcile: ok=%v err=%v", ok, reconcileErr) } - cap = 20 - // The cached cap is still 10, so this failed aggregate admission forces a - // refresh even though the TTL has not elapsed and then succeeds at 20. - grown, err := g2.Reserve(1) - if err != nil { - t.Fatalf("growth refresh reservation=%v", err) + generation.Close() + if _, err = generation.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetClosed) { + t.Fatalf("closed spill error = %v", err) } - if got := calls.Load(); got != 2 { - t.Fatalf("growth refresh provider calls=%d, want 2", got) + disk.Release() + fd.Release() + if budget.SpillDiskUsed() != 0 || budget.SpillFDUsed() != 0 { + t.Fatal("spill reservations leaked") } - grown.Release() - owned.Release() } -func TestHashBuildBudgetCapProviderConcurrentSingleFlight(t *testing.T) { - b := MustNewHashBuildBudget(128, 128) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - var calls atomic.Int32 - started := make(chan struct{}) - release := make(chan struct{}) - b.SetAggregateCapProvider(func() (uint64, error) { - if calls.Add(1) == 1 { - close(started) - <-release - } - return 128, nil - }) - const workers = 16 - tokens := make(chan *HashBuildReservation, workers) - var wg sync.WaitGroup - for i := 0; i < workers; i++ { - wg.Add(1) - go func() { - defer wg.Done() - tok, err := g.Reserve(1) +func TestHashBuildBudgetSpillReleaseRejectsLedgerUnderflow(t *testing.T) { + for _, test := range []struct { + name string + reserve func(*HashBuildBudgetGeneration) (func() bool, error) + corrupt func(*HashBuildBudget, *HashBuildBudgetGeneration) + }{ + { + name: "disk", + reserve: func(generation *HashBuildBudgetGeneration) (func() bool, error) { + reservation, err := generation.ReserveSpillDisk(2) + return reservation.Release, err + }, + corrupt: func(budget *HashBuildBudget, generation *HashBuildBudgetGeneration) { + generation.spillDiskUsed = 1 + budget.spillDiskUsed = 1 + }, + }, + { + name: "fd", + reserve: func(generation *HashBuildBudgetGeneration) (func() bool, error) { + reservation, err := generation.ReserveSpillFD(2) + return reservation.Release, err + }, + corrupt: func(budget *HashBuildBudget, generation *HashBuildBudgetGeneration) { + generation.spillFDUsed = 1 + budget.spillFDUsed = 1 + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + budget := MustNewHashBuildBudget(64, 64) + generation, err := budget.OpenGeneration(1) if err != nil { - t.Errorf("concurrent reserve: %v", err) - return + t.Fatal(err) } - tokens <- tok - }() - } - <-started - close(release) - wg.Wait() - if got := calls.Load(); got != 1 { - t.Fatalf("concurrent provider calls=%d, want 1", got) - } - for i := 0; i < workers; i++ { - (<-tokens).Release() - } -} - -func TestHashBuildBudgetCapProviderErrorCachedFailClosed(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - want := errors.New("cgroup unavailable") - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 0, want - }) - if _, err := g.Reserve(1); !errors.Is(err, want) { - t.Fatalf("provider error=%v, want %v", err, want) - } - if _, err := g.Reserve(1); !errors.Is(err, want) { - t.Fatalf("cached provider error=%v, want %v", err, want) - } - if got := calls.Load(); got != 1 { - t.Fatalf("cached error provider calls=%d, want 1", got) - } - now = now.Add(hashBuildBudgetCapRefreshTTL + time.Nanosecond) - if _, err := g.Reserve(1); !errors.Is(err, want) { - t.Fatalf("expired provider error=%v, want %v", err, want) - } - if got := calls.Load(); got != 2 { - t.Fatalf("expired error provider calls=%d, want 2", got) - } -} - -func TestHashBuildBudgetCapProviderSharedByReserveAndGrow(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 20, nil - }) - tok, err := g.Reserve(2) - if err != nil { - t.Fatal(err) - } - if err = tok.Grow(3); err != nil { - t.Fatal(err) - } - if got := calls.Load(); got != 1 { - t.Fatalf("Reserve+Grow provider calls=%d, want 1", got) + release, err := test.reserve(generation) + if err != nil { + t.Fatal(err) + } + budget.mu.Lock() + test.corrupt(budget, generation) + budget.mu.Unlock() + defer func() { + if recover() == nil { + t.Fatal("corrupt spill ledger release did not panic") + } + }() + release() + }) } - tok.Release() } -func TestHashBuildBudgetCapProviderZeroTTLRefreshesEveryReservation(t *testing.T) { - b := MustNewHashBuildBudget(16, 16) - g, _ := b.OpenGeneration(1) - b.capRefreshTTL = 0 - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 16, nil - }) - for i := 0; i < 2; i++ { - tok, err := g.Reserve(1) - if err != nil { - t.Fatal(err) +func TestClampSpillFDCapBoundaries(t *testing.T) { + for _, test := range []struct { + configured, limit uint64 + known bool + want uint64 + }{ + {2048, 1 << 20, false, 0}, + {0, 1024, true, 0}, + {2048, 64, true, 0}, + {2048, 65, true, 1}, + {2048, 128, true, 64}, + {2048, 1024, true, 768}, + {10, 1024, true, 10}, + {2048, math.MaxUint64, true, 2048}, + } { + if got := clampSpillFDCap(test.configured, test.limit, test.known); got != test.want { + t.Fatalf("clampSpillFDCap(%d, %d, %v) = %d, want %d", + test.configured, test.limit, test.known, got, test.want) } - tok.Release() - } - if got := calls.Load(); got != 2 { - t.Fatalf("zero-TTL provider calls=%d, want 2", got) } } -func TestHashBuildBudgetUpdateAndProviderReinstallReuseCache(t *testing.T) { - b := MustNewHashBuildBudget(16, 16) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - provider := func() (uint64, error) { - calls.Add(1) - return 16, nil - } - b.SetAggregateCapProvider(provider) - first, err := g.Reserve(1) +func TestGetHashBuildBudgetInitializesAndReusesCNAggregate(t *testing.T) { + const localService = "__process_local_cn__" + hashBuildCNBudgets.Delete(localService) + t.Cleanup(func() { hashBuildCNBudgets.Delete(localService) }) + + first := &Process{Base: &BaseProcess{Lim: Limitation{Size: 2 << 20, SpillSize: 4 << 20}}} + firstGeneration, err := first.GetHashBuildBudget() if err != nil { t.Fatal(err) } - if err = b.UpdateAggregateCap(16); err != nil { - t.Fatal(err) + if firstGeneration.Cap() != 2<<20 || firstGeneration.SpillDiskCap() != 4<<20 { + t.Fatalf("first generation limits: %+v", firstGeneration.Snapshot()) } - // GetHashBuildBudget re-installs an equivalent closure before updating the - // freshly resolved cap. The update seeds the cache for this query. - b.SetAggregateCapProvider(provider) - if err = b.UpdateAggregateCap(16); err != nil { - t.Fatal(err) + if cached, cachedErr := first.GetHashBuildBudget(); cachedErr != nil || cached != firstGeneration { + t.Fatal("process generation was not cached") } - second, err := g.Reserve(1) + second := &Process{Base: &BaseProcess{Lim: Limitation{Size: 1 << 20}}} + secondGeneration, err := second.GetHashBuildBudget() if err != nil { t.Fatal(err) } - if got := calls.Load(); got != 1 { - t.Fatalf("Update+Set provider calls=%d, want 1", got) + if secondGeneration == firstGeneration || secondGeneration.budget != firstGeneration.budget || + secondGeneration.Cap() != 1<<20 { + t.Fatal("second process did not reuse the CN aggregate") } - first.Release() - second.Release() + firstGeneration.Close() + secondGeneration.Close() + firstGeneration.budget.Close() } -func TestHashBuildBudgetProviderReplacementInvalidatesCache(t *testing.T) { - b := MustNewHashBuildBudget(16, 16) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var oldCalls, newCalls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - oldCalls.Add(1) - return 16, nil +func TestResolveHashBuildCeiling(t *testing.T) { + const gib = uint64(1 << 30) + ceiling, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ + CgroupMemoryMax: 20 * gib, + HostMemTotal: 10 * gib, + GlobalMpoolCap: 30 * gib, + FileCacheHint: gib, + ProcessLimitationSize: 2 * gib, }) - owned, err := g.Reserve(1) if err != nil { t.Fatal(err) } - b.SetAggregateCapProvider(func() (uint64, error) { - newCalls.Add(1) - return 2, nil - }) - second, err := g.Reserve(1) - if err != nil { - t.Fatalf("replacement reserve=%v", err) + if ceiling.EffectiveCN != 10*gib || ceiling.Reserve != 4*gib || + ceiling.CNHashCap != 6*gib || ceiling.QueryCap != 2*gib { + t.Fatalf("ceiling = %+v", ceiling) + } + if _, err = ResolveHashBuildCeiling(HashBuildCeilingInputs{ + CgroupMemoryMax: math.MaxUint64, + }); !errors.Is(err, ErrHashBuildCeilingMissing) { + t.Fatalf("missing finite source error = %v", err) } - if oldCalls.Load() != 1 || newCalls.Load() != 1 { - t.Fatalf("provider calls old=%d new=%d, want 1,1", oldCalls.Load(), newCalls.Load()) + if small, smallErr := ResolveHashBuildCeiling(HashBuildCeilingInputs{ + HostMemTotal: 3 * gib, + FileCacheHint: 3 * gib, + }); smallErr != nil || small.CNHashCap != 3*gib/20 { + t.Fatalf("small-CN ceiling = %+v, err=%v", small, smallErr) } - owned.Release() - second.Release() } -func TestHashBuildBudgetCNSourceTurnoverRetainsRestrictiveFallback(t *testing.T) { - const gib = uint64(1 << 30) - oldInputs := HashBuildCeilingInputs{HostMemTotal: 10 * gib, CgroupMemoryMax: 20 * gib} - newInputs := HashBuildCeilingInputs{HostMemTotal: 20 * gib, CgroupMemoryMax: 10 * gib} - oldCeiling, err := ResolveHashBuildCeiling(oldInputs) - if err != nil { - t.Fatal(err) - } - newCeiling, err := ResolveHashBuildCeiling(newInputs) - if err != nil { - t.Fatal(err) - } - if oldCeiling.CNHashCap != newCeiling.CNHashCap { - t.Fatalf("test setup caps old=%d new=%d, want equal", oldCeiling.CNHashCap, newCeiling.CNHashCap) - } - b := MustNewHashBuildBudget(oldCeiling.CNHashCap, oldCeiling.CNHashCap) - b.installCNCapProvider(oldInputs) - b.mergeObservedCNCap(newInputs, newCeiling.CNHashCap) - - // A zero source sample models transient read failures before the installed - // provider gets a complete view of the turnover. The merged shared snapshot - // must retain both restrictive observations and cannot reopen the cap. - b.refreshMu.Lock() - got, err := b.resolveCNCapSample(HashBuildCeilingInputs{}) - b.refreshMu.Unlock() - if err != nil { - t.Fatal(err) - } - if got > newCeiling.CNHashCap { - t.Fatalf("fallback cap=%d, exceeds restrictive cap %d", got, newCeiling.CNHashCap) - } -} - -func TestHashBuildBudgetSlowProviderTTLStartsAfterSample(t *testing.T) { - b := MustNewHashBuildBudget(20, 20) - b.capRefreshTTL = hashBuildBudgetCapRefreshTTL - g, _ := b.OpenGeneration(1) - now := time.Unix(0, 0) - b.capNow = func() time.Time { return now } - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - now = now.Add(2 * hashBuildBudgetCapRefreshTTL) - return 20, nil - }) - first, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - second, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - if got := calls.Load(); got != 1 { - t.Fatalf("slow provider calls=%d, want freshly completed sample reused", got) - } - first.Release() - second.Release() -} - -func TestHashBuildBudgetClosedSkipsCapProvider(t *testing.T) { - b := MustNewHashBuildBudget(10, 10) - g, _ := b.OpenGeneration(1) - var calls atomic.Int32 - b.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return 10, nil - }) - b.Close() - if _, err := g.Reserve(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed reserve=%v", err) - } - if got := calls.Load(); got != 0 { - t.Fatalf("closed provider calls=%d, want 0", got) - } -} - -func TestHashBuildBudgetCompatibilityAndObservabilitySurface(t *testing.T) { - var nilBudget *HashBuildBudget - if !nilBudget.Snapshot().Closed || - nilBudget.AggregateCap() != 0 || - nilBudget.CNHashCap() != 0 || - nilBudget.QueryCap() != 0 || - nilBudget.AggregateUsed() != 0 || - nilBudget.CNHashUsed() != 0 || - nilBudget.Current() != 0 || - nilBudget.Capacity() != 0 || - !nilBudget.Closed() || - nilBudget.SpillDiskCap() != 0 || - nilBudget.SpillDiskUsed() != 0 || - nilBudget.SpillFDCap() != 0 || - nilBudget.SpillFDUsed() != 0 { - t.Fatal("nil budget accessors must report an inert closed budget") - } - nilBudget.Close() - if err := nilBudget.SetSpillCaps(1, 1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil SetSpillCaps error = %v", err) - } - - var nilGeneration *HashBuildBudgetGeneration - if nilGeneration.ID() != 0 || - nilGeneration.Cap() != 0 || - nilGeneration.QueryCap() != 0 || - nilGeneration.Capacity() != 0 || - nilGeneration.Used() != 0 || - nilGeneration.Current() != 0 || - !nilGeneration.Closed() || - nilGeneration.SpillDiskCap() != 0 || - nilGeneration.SpillDiskUsed() != 0 || - nilGeneration.SpillFDCap() != 0 || - nilGeneration.SpillFDUsed() != 0 || - !nilGeneration.Snapshot().Closed { - t.Fatal("nil generation accessors must report an inert closed generation") - } - nilGeneration.Close() - if nilGeneration.TryReserve(1) { - t.Fatal("nil generation reservation succeeded") - } - - for _, kind := range []HashBuildBudgetErrorKind{ - HashBuildBudgetErrorAdmission, - HashBuildBudgetErrorClosed, - HashBuildBudgetErrorInvalid, - HashBuildBudgetErrorCeilingMissing, - } { - err := &HashBuildBudgetError{Kind: kind} - if err.Error() == "" || err.Unwrap() == nil { - t.Fatalf("kind %d did not expose an error", kind) - } - } - var nilBudgetErr *HashBuildBudgetError - if nilBudgetErr.Error() != "" || nilBudgetErr.Unwrap() != nil || nilBudgetErr.Is(ErrHashBuildBudgetAdmission) { - t.Fatal("nil budget error must remain inert") - } - for _, tc := range []struct { - kind HashBuildBudgetErrorKind - target error - }{ - {HashBuildBudgetErrorAdmission, ErrHashBuildBudgetAdmission}, - {HashBuildBudgetErrorClosed, ErrHashBuildBudgetClosed}, - {HashBuildBudgetErrorInvalid, ErrHashBuildBudgetInvalid}, - {HashBuildBudgetErrorCeilingMissing, ErrHashBuildCeilingMissing}, - } { - if !errors.Is(&HashBuildBudgetError{Kind: tc.kind}, tc.target) { - t.Fatalf("kind %d did not match %v", tc.kind, tc.target) - } - } - if errors.Is(&HashBuildBudgetError{Kind: HashBuildBudgetErrorClosed}, - ErrHashBuildBudgetAdmission, - ) { - t.Fatal("closed budget must not match a recoverable capacity admission") - } - unknown := &HashBuildBudgetError{Kind: HashBuildBudgetErrorKind(255)} - if errors.Is(unknown, ErrHashBuildBudgetAdmission) || - !errors.Is(unknown, ErrHashBuildBudgetInvalid) { - t.Fatal("unknown error kind must remain a fatal invalid error") - } - message := &HashBuildBudgetError{Message: "explicit"} - if message.Error() != "explicit" { - t.Fatalf("explicit message = %q", message.Error()) - } - - b, err := NewHashBuildBudgetWithSpillCaps(100, 80, 200, 10) - if err != nil { - t.Fatal(err) - } - if b.AggregateCap() != 100 || b.CNHashCap() != 100 || b.Capacity() != 100 || b.QueryCap() != 80 { - t.Fatal("budget cap aliases disagree") - } - if b.SpillDiskCap() != 200 || b.SpillFDCap() != 10 { - t.Fatal("explicit spill caps were not installed") - } - if err = b.SetSpillCaps(0, 0); err != nil { - t.Fatal(err) - } - snapshot := b.Snapshot() - if snapshot.AggregateCap != 100 || snapshot.AggregateUsed != 0 || snapshot.Closed { - t.Fatalf("unexpected budget snapshot: %+v", snapshot) - } - - g, err := b.OpenGenerationWithLimits(7, 50, 100, 5) - if err != nil { - t.Fatal(err) - } - if g.ID() != 7 || g.Cap() != 50 || g.QueryCap() != 50 || g.Capacity() != 50 { - t.Fatal("generation identity or cap aliases disagree") - } - if g.SpillDiskCap() != 100 || g.SpillFDCap() != 5 || g.Current() != 0 { - t.Fatal("generation spill caps or current usage are wrong") - } - if !g.TryReserve(1) { - t.Fatal("TryReserve rejected an admissible charge") - } - token, err := g.Reserve(8) - if err != nil { - t.Fatal(err) - } - if token.GenerationID() != 7 || token.Size() != 8 || token.Released() { - t.Fatal("memory reservation accessors are inconsistent") - } - if ok, reconcileErr := token.Reconcile(6); !ok || reconcileErr != nil { - t.Fatalf("compatibility reconcile failed: ok=%v err=%v", ok, reconcileErr) - } - moved := token.TransferOwnership() - if moved == nil || !token.Released() || moved.GenerationID() != 7 { - t.Fatal("memory ownership transfer failed") - } - if !moved.Release() { - t.Fatal("transferred memory reservation did not release") - } - - disk, err := g.ReserveSpillDiskBytes(12) - if err != nil { - t.Fatal(err) - } - fd, err := g.ReserveSpillFileDescriptors(2) - if err != nil { - t.Fatal(err) - } - if disk.Size() != 12 || disk.Released() || fd.Size() != 2 || fd.Released() { - t.Fatal("spill reservation accessors are inconsistent") - } - if ok, reconcileErr := disk.Reconcile(10); !ok || reconcileErr != nil { - t.Fatalf("disk reconcile failed: ok=%v err=%v", ok, reconcileErr) - } - if ok, reconcileErr := fd.Reconcile(1); !ok || reconcileErr != nil { - t.Fatalf("fd reconcile failed: ok=%v err=%v", ok, reconcileErr) - } - movedDisk := disk.TransferTo() - movedFD := fd.TransferOwnership() - if movedDisk == nil || movedFD == nil || !disk.Released() || !fd.Released() { - t.Fatal("spill ownership transfer failed") - } - if !movedDisk.Release() || !movedFD.Release() { - t.Fatal("transferred spill reservations did not release") - } - stats := g.Stats() - if stats.ID != 7 || - g.Peak() == 0 || - g.ReserveCount() == 0 || - g.ReconcileCount() == 0 || - g.ReleaseCount() == 0 || - g.RejectCount() != 0 { - t.Fatalf("unexpected generation stats: %+v", stats) - } - - other, err := b.NewGeneration(8) - if err != nil { - t.Fatal(err) - } - other.Close() - query, err := b.OpenQueryBudget(9) - if err != nil { - t.Fatal(err) - } - query.Close() - explicit, err := b.OpenGenerationWithCapAndSpill(10, 40, 80, 4) - if err != nil { - t.Fatal(err) - } - explicit.Close() - g.Close() - b.Close() - if !b.Closed() || !g.Closed() { - t.Fatal("close accessors did not observe terminal state") - } - - ceiling, err := ResolveHashBuildBudget(HashBuildCeilingInputs{HostMemTotal: 8 << 30}) - if err != nil { - t.Fatal(err) - } - fromCeiling, err := NewHashBuildBudgetFromCeiling(ceiling) - if err != nil { - t.Fatal(err) - } - fromCeiling.Close() -} - -func TestHashBuildBudgetCompatibilityUnhappyPaths(t *testing.T) { - var nilProcess *Process - if _, err := nilProcess.GetHashBuildBudget(); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil process error = %v", err) - } - - var nilBudget *HashBuildBudget - if _, err := nilBudget.OpenGeneration(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil budget generation error = %v", err) - } - if _, err := nilBudget.OpenGenerationWithCap(1, 1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil budget explicit generation error = %v", err) - } - - b, err := NewHashBuildBudgetWithSpillCaps(100, 80, 20, 4) - if err != nil { - t.Fatal(err) - } - if _, err = b.OpenGenerationWithCap(1, 0); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("zero generation cap error = %v", err) - } - if _, err = b.OpenGenerationWithCap(1, 101); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("oversized generation cap error = %v", err) - } - - g, err := b.OpenGenerationWithSpillCaps(1, 80, 10, 2) - if err != nil { - t.Fatal(err) - } - other, err := b.OpenGenerationWithSpillCaps(2, 80, 20, 4) - if err != nil { - t.Fatal(err) - } - defaults, err := b.OpenGenerationWithSpillCaps(3, 80, 0, 0) - if err != nil { - t.Fatal(err) - } - if defaults.SpillDiskCap() != 20 || defaults.SpillFDCap() != 4 { - t.Fatal("default generation spill caps were not clamped to the CN caps") - } - defaults.Close() - - disk, err := g.ReserveSpillDisk(8) - if err != nil { - t.Fatal(err) - } - if err = disk.Grow(2); err != nil { - t.Fatal(err) - } - if err = disk.Grow(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("query disk admission error = %v", err) - } - if _, err = other.ReserveSpillDisk(11); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("CN disk admission error = %v", err) - } - if ok, reconcileErr := disk.ReconcileDown(11); ok || !errors.Is(reconcileErr, ErrHashBuildReservationUpward) { - t.Fatalf("upward disk reconcile: ok=%v err=%v", ok, reconcileErr) - } - - fd, err := g.ReserveSpillFD(1) - if err != nil { - t.Fatal(err) - } - if _, err = g.ReserveSpillFD(2); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("query FD admission error = %v", err) - } - otherFD, err := other.ReserveSpillFD(3) - if err != nil { - t.Fatal(err) - } - if _, err = other.ReserveSpillFD(1); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("CN FD admission error = %v", err) - } - if ok, reconcileErr := fd.ReconcileDown(2); ok || !errors.Is(reconcileErr, ErrHashBuildReservationUpward) { - t.Fatalf("upward FD reconcile: ok=%v err=%v", ok, reconcileErr) - } - - memory, err := g.Reserve(1) - if err != nil { - t.Fatal(err) - } - movedMemory := memory.TransferTo() - if movedMemory == nil || memory.Transfer() != nil || memory.Release() { - t.Fatal("inactive memory token accepted a second terminal transition") - } - if !movedMemory.Release() || movedMemory.Release() { - t.Fatal("memory release was not exactly once") - } - - g.Close() - if _, err = g.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed disk reservation error = %v", err) - } - if _, err = g.ReserveSpillFD(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed FD reservation error = %v", err) - } - if err = disk.Grow(1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed disk growth error = %v", err) - } - if !disk.Release() || disk.Release() || disk.Transfer() != nil { - t.Fatal("disk release was not exactly once") - } - if _, reconcileErr := disk.ReconcileDown(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("inactive disk reconcile error = %v", reconcileErr) - } - if !fd.Release() || fd.Release() || fd.Transfer() != nil { - t.Fatal("FD release was not exactly once") - } - if _, reconcileErr := fd.ReconcileDown(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("inactive FD reconcile error = %v", reconcileErr) - } - otherFD.Release() - - var nilMemory *HashBuildReservation - var nilDisk *HashBuildSpillDiskReservation - var nilFD *HashBuildSpillFDReservation - var nilGeneration *HashBuildBudgetGeneration - if nilMemory.Size() != 0 || nilMemory.GenerationID() != 0 || !nilMemory.Released() || - nilMemory.Release() || nilMemory.Transfer() != nil || nilMemory.TransferTo() != nil || - nilDisk.Size() != 0 || !nilDisk.Released() || nilDisk.Release() || - nilDisk.Transfer() != nil || nilDisk.TransferOwnership() != nil || - nilFD.Size() != 0 || !nilFD.Released() || nilFD.Release() || - nilFD.Transfer() != nil || nilFD.TransferTo() != nil { - t.Fatal("nil reservation must remain inert") - } - if _, err = nilGeneration.ReserveSpillDisk(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil disk generation error = %v", err) - } - if _, err = nilGeneration.ReserveSpillFD(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil FD generation error = %v", err) - } - if _, reconcileErr := nilMemory.Reconcile(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("nil memory reconcile error = %v", reconcileErr) - } - if _, reconcileErr := nilDisk.Reconcile(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("nil disk reconcile error = %v", reconcileErr) - } - if _, reconcileErr := nilFD.Reconcile(0); !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("nil FD reconcile error = %v", reconcileErr) - } - if err = nilMemory.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("nil memory growth error = %v", err) - } - if err = nilDisk.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("nil disk growth error = %v", err) - } - - b.Close() - if _, err = b.OpenGeneration(3); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget generation error = %v", err) - } - if _, err = b.OpenGenerationWithCap(3, 1); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget explicit generation error = %v", err) - } -} - -func TestGetHashBuildBudgetInitializesAndReusesCNAggregate(t *testing.T) { - const localService = "__process_local_cn__" - hashBuildCNBudgets.Delete(localService) - t.Cleanup(func() { hashBuildCNBudgets.Delete(localService) }) - - first := &Process{Base: &BaseProcess{Lim: Limitation{ - Size: 2 << 20, - SpillSize: 4 << 20, - }}} - firstGeneration, err := first.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if firstGeneration.Cap() != 2<<20 || firstGeneration.SpillDiskCap() != 4<<20 { - t.Fatalf("unexpected first generation limits: %+v", firstGeneration.Snapshot()) - } - cached, err := first.GetHashBuildBudget() - if err != nil || cached != firstGeneration { - t.Fatalf("process-local generation was not cached: generation=%p err=%v", cached, err) - } - - second := &Process{Base: &BaseProcess{Lim: Limitation{Size: 1 << 20}}} - secondGeneration, err := second.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if secondGeneration == firstGeneration || - secondGeneration.budget != firstGeneration.budget || - secondGeneration.Cap() != 1<<20 { - t.Fatal("second process did not reuse the CN aggregate with its own generation") - } - - aggregate := firstGeneration.budget - defaultAggregateSpillCap := aggregate.SpillDiskCap() - raisedSpillCap := defaultAggregateSpillCap + 1<<20 - third := &Process{Base: &BaseProcess{Lim: Limitation{ - Size: 1 << 20, - SpillSize: int64(raisedSpillCap), - }}} - thirdGeneration, err := third.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if thirdGeneration.SpillDiskCap() != raisedSpillCap || - aggregate.SpillDiskCap() != raisedSpillCap { - t.Fatalf("explicit spill cap was not raised at the shared ledger: generation=%d aggregate=%d want=%d", - thirdGeneration.SpillDiskCap(), aggregate.SpillDiskCap(), raisedSpillCap) - } - - lower := &Process{Base: &BaseProcess{Lim: Limitation{ - Size: 1 << 20, - SpillSize: 2 << 20, - }}} - lowerGeneration, err := lower.GetHashBuildBudget() - if err != nil { - t.Fatal(err) - } - if lowerGeneration.SpillDiskCap() != 2<<20 || - aggregate.SpillDiskCap() != raisedSpillCap { - t.Fatalf("lower per-query spill cap changed the shared ceiling: generation=%d aggregate=%d want aggregate=%d", - lowerGeneration.SpillDiskCap(), aggregate.SpillDiskCap(), raisedSpillCap) - } - - firstGeneration.Close() - secondGeneration.Close() - thirdGeneration.Close() - lowerGeneration.Close() - aggregate.Close() -} - -func TestHashBuildBudgetExplicitSpillCapConcurrentRaise(t *testing.T) { - budget := MustNewHashBuildBudget(100, 100) - t.Cleanup(budget.Close) - generation, err := budget.OpenGenerationWithSpillCaps(1, 100, 800, 1) - if err != nil { - t.Fatal(err) - } - t.Cleanup(generation.Close) - reservation, err := generation.ReserveSpillDisk(700) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { reservation.Release() }) - - caps := []uint64{801, 900, 1200, 1100} - start := make(chan struct{}) - errs := make(chan error, len(caps)) - var wg sync.WaitGroup - for _, cap := range caps { - wg.Add(1) - go func() { - defer wg.Done() - <-start - errs <- budget.raiseSpillDiskCapToExplicitLimit(cap) - }() - } - close(start) - wg.Wait() - close(errs) - for raiseErr := range errs { - if raiseErr != nil { - t.Fatal(raiseErr) - } - } - if got := budget.SpillDiskCap(); got != 1200 { - t.Fatalf("concurrent raised spill cap = %d, want 1200", got) - } - if !reservation.Release() || budget.SpillDiskUsed() != 0 { - t.Fatalf("live reservation did not release after cap growth: %+v", budget.Snapshot()) - } - - budget.Close() - if err = budget.raiseSpillDiskCapToExplicitLimit(1300); !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed budget raise error = %v, want %v", err, ErrHashBuildBudgetClosed) - } - if got := budget.SpillDiskCap(); got != 1200 { - t.Fatalf("closed budget changed spill cap to %d", got) - } -} - -func TestOpenProcessGenerationClampsStaleResolvedCapAtomically(t *testing.T) { - budget := MustNewHashBuildBudget(100, 100) - if err := budget.UpdateAggregateCap(40); err != nil { - t.Fatal(err) - } - - generation, err := budget.openProcessGeneration(1, 100, 0) - if err != nil { - t.Fatal(err) - } - if generation.Cap() != 40 { - t.Fatalf("generation cap = %d, want current aggregate cap 40", - generation.Cap()) - } - if generation.SpillDiskCap() != defaultSpillCap(40) { - t.Fatalf("spill disk cap = %d, want %d", - generation.SpillDiskCap(), defaultSpillCap(40)) - } +func TestHashBuildBudgetUsesCurrentMemoryInputs(t *testing.T) { + previous := hashBuildProcessMemoryInputs + hashBuildProcessMemoryInputs = HashBuildCeilingInputs{HostMemTotal: 8 << 30} + t.Cleanup(func() { hashBuildProcessMemoryInputs = previous }) + previousHint := fileservice.GlobalMemoryCacheSizeHint.Load() + fileservice.GlobalMemoryCacheSizeHint.Store(1 << 30) + t.Cleanup(func() { fileservice.GlobalMemoryCacheSizeHint.Store(previousHint) }) - // Explicit public configuration remains strict. Only the process path may - // clamp a ceiling sample that became stale between resolution and opening. - if _, err = budget.OpenGenerationWithCap(2, 100); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("explicit oversized generation cap returned %v", err) + inputs := hashBuildProcessMemoryInputs + inputs.FileCacheHint = uint64(fileservice.GlobalMemoryCacheSizeHint.Load()) + ceiling, err := ResolveHashBuildCeiling(inputs) + if err != nil || ceiling.CNHashCap == 0 { + t.Fatalf("current memory inputs did not produce a finite cap: %+v %v", ceiling, err) } - generation.Close() - budget.Close() } -func TestHashBuildBudgetDefensiveAndProviderFailurePaths(t *testing.T) { - for _, limits := range [][2]uint64{{0, 1}, {1, 0}, {1, 2}} { - if _, err := NewHashBuildBudget(limits[0], limits[1]); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("invalid limits %v returned %v", limits, err) - } - if _, err := NewHashBuildBudgetWithSpillCaps(limits[0], limits[1], 1, 1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("invalid spill budget limits %v returned %v", limits, err) - } - } - func() { - defer func() { - if recover() == nil { - t.Fatal("MustNewHashBuildBudget did not panic for invalid limits") - } - }() - MustNewHashBuildBudget(0, 0) - }() - - var nilBudget *HashBuildBudget - if _, _, _, err := nilBudget.refreshAggregateCap(false, 0); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil refresh error = %v", err) - } - nilBudget.SetAggregateCapProvider(func() (uint64, error) { return 1, nil }) - if err := nilBudget.UpdateAggregateCap(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("nil cap update error = %v", err) - } - - b, err := NewHashBuildBudgetWithSpillCaps(math.MaxUint64, math.MaxUint64, 10, 10) - if err != nil { - t.Fatal(err) - } - g, err := b.OpenGenerationWithSpillCaps(1, math.MaxUint64, 10, 10) - if err != nil { - t.Fatal(err) - } - if _, err = b.OpenGenerationWithSpillCaps(2, 0, 0, 0); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("invalid spill generation error = %v", err) - } - disk, err := g.ReserveSpillDisk(5) - if err != nil { - t.Fatal(err) - } - if err = b.SetSpillCaps(4, 10); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("spill cap reduction error = %v", err) - } - disk.Release() - - b.capNow = nil - if err = b.UpdateAggregateCap(math.MaxUint64); err != nil { - t.Fatal(err) - } - b.SetAggregateCapProvider(func() (uint64, error) { return 0, nil }) - if _, err = g.Reserve(1); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("zero provider ceiling error = %v", err) - } - b.capCached = true - b.capRefreshTTL = time.Hour - b.capRefreshAt = time.Now() - b.capRefreshEpoch = 2 - b.capRefreshErr = nil - called := false - b.capProvider = func() (uint64, error) { - called = true - return 1, nil - } - if _, _, refreshed, err := b.refreshAggregateCap(true, 1); err != nil || refreshed || called { - t.Fatalf("concurrent refresh was not reused: refreshed=%v called=%v err=%v", refreshed, called, err) - } - - failing, err := NewHashBuildBudget(100, 100) - if err != nil { - t.Fatal(err) - } - failingGeneration, err := failing.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - providerErr := errors.New("provider failed") - failing.SetAggregateCapProvider(func() (uint64, error) { return 0, providerErr }) - if _, err = failingGeneration.Reserve(1); !errors.Is(err, providerErr) { - t.Fatalf("reserve provider error = %v", err) - } - token := &HashBuildReservation{ - budget: failing, - generation: failingGeneration, - core: &hashBuildReservationCore{size: 1}, - } - if err = token.Grow(1); !errors.Is(err, providerErr) { - t.Fatalf("grow provider error = %v", err) - } - - forceReserve, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - forceReserve.capRefreshTTL = time.Hour - forceReserveGeneration, err := forceReserve.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - reserveCalls := 0 - forceReserve.SetAggregateCapProvider(func() (uint64, error) { - reserveCalls++ - if reserveCalls == 1 { - return 5, nil - } - return 0, providerErr - }) - seed, err := forceReserveGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - seed.Release() - if _, err = forceReserveGeneration.Reserve(6); !errors.Is(err, providerErr) { - t.Fatalf("forced reserve refresh error = %v", err) - } - - forceGrow, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - forceGrow.capRefreshTTL = time.Hour - forceGrowGeneration, err := forceGrow.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - growCalls := 0 - forceGrow.SetAggregateCapProvider(func() (uint64, error) { - growCalls++ - if growCalls == 1 { - return 5, nil +func BenchmarkHashBuildBudgetAllocationAccount(b *testing.B) { + for _, accounted := range []bool{false, true} { + mode := "unaccounted" + if accounted { + mode = "accounted" } - return 0, providerErr - }) - growToken, err := forceGrowGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = growToken.Grow(5); !errors.Is(err, providerErr) { - t.Fatalf("forced grow refresh error = %v", err) - } - growToken.Release() - - rescueGrow, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - rescueGrow.capRefreshTTL = time.Hour - rescueGeneration, err := rescueGrow.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - rescueCalls := 0 - rescueGrow.SetAggregateCapProvider(func() (uint64, error) { - rescueCalls++ - if rescueCalls == 1 { - return 5, nil + for _, size := range []int{64, 4 << 10, 64 << 10} { + b.Run(mode+"/alloc-free/"+fmt.Sprint(size), func(b *testing.B) { + mp := mpool.MustNewZero() + var generation *HashBuildBudgetGeneration + var registry *mpool.AllocationAccountRegistry + var account *mpool.AllocationAccount + if accounted { + budget := MustNewHashBuildBudget(math.MaxUint64, math.MaxUint64) + var err error + generation, err = budget.OpenGeneration(1) + if err != nil { + b.Fatal(err) + } + registry, err = mpool.NewAllocationAccountRegistry(1, 2) + if err != nil { + b.Fatal(err) + } + account, err = registry.OpenWithController(math.MaxInt64, generation) + if err != nil { + b.Fatal(err) + } + } + + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for range b.N { + var allocation []byte + var err error + if accounted { + allocation, err = mp.AllocAccounted(size, account, 1, 1) + } else { + allocation, err = mp.Alloc(size, true) + } + if err != nil { + b.Fatal(err) + } + mp.Free(allocation) + } + b.StopTimer() + if accounted { + if generation.Used() != 0 || account.Snapshot().Used != 0 { + b.Fatal("physical allocation capacity leaked") + } + if _, _, err := registry.CompleteTerminal(account); err != nil { + b.Fatal(err) + } + generation.Close() + } + }) } - return 10, nil - }) - rescueToken, err := rescueGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = rescueToken.Grow(5); err != nil || rescueToken.Size() != 6 { - t.Fatalf("forced growth rescue: size=%d err=%v", rescueToken.Size(), err) - } - rescueToken.Release() - - noProvider, err := NewHashBuildBudget(5, 5) - if err != nil { - t.Fatal(err) - } - noProviderGeneration, err := noProvider.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - noProviderToken, err := noProviderGeneration.Reserve(1) - if err != nil { - t.Fatal(err) - } - if err = noProviderToken.Grow(5); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("non-provider aggregate growth error = %v", err) - } - noProviderToken.Release() - - empty, err := NewHashBuildBudget(100, 100) - if err != nil { - t.Fatal(err) - } - if _, err = empty.resolveCNCapSample(HashBuildCeilingInputs{}); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("empty live sample error = %v", err) - } - previousCap := commonmpool.GlobalCap() - commonmpool.InitCap(2 << 30) - previousHint := fileservice.GlobalMemoryCacheSizeHint.Swap(32 << 20) - func() { - defer commonmpool.InitCap(previousCap) - defer fileservice.GlobalMemoryCacheSizeHint.Store(previousHint) - if _, err = empty.sampleCNCap(); err != nil { - t.Fatalf("live CN sample error = %v", err) - } - }() - empty.capNow = nil - empty.installCNCapProvider(HashBuildCeilingInputs{HostMemTotal: 1 << 30}) - empty.mergeObservedCNCap(HashBuildCeilingInputs{ - CgroupMemoryMax: 512 << 20, - HostMemTotal: 768 << 20, - GlobalMpoolCap: 640 << 20, - FileCacheHint: 32 << 20, - }, 50) - if empty.AggregateCap() != 50 { - t.Fatalf("merged aggregate cap = %d", empty.AggregateCap()) - } - - closedBudget, err := NewHashBuildBudget(10, 10) - if err != nil { - t.Fatal(err) - } - closedGeneration, err := closedBudget.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - closedGeneration.closed = true - closedBudget.mu.Lock() - _, err, rejected := closedGeneration.reserveLocked(1, true) - closedBudget.mu.Unlock() - if rejected || !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed reserveLocked: rejected=%v err=%v", rejected, err) - } - closedToken := &HashBuildReservation{ - budget: closedBudget, - generation: closedGeneration, - core: &hashBuildReservationCore{size: 1}, - } - if err, rejected = closedToken.growLocked(1, true); rejected || !errors.Is(err, ErrHashBuildBudgetClosed) { - t.Fatalf("closed growLocked: rejected=%v err=%v", rejected, err) - } - closedToken.core.state.Store(hashBuildReservationReleased) - if err, rejected = closedToken.growLocked(1, true); rejected || !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("inactive growLocked: rejected=%v err=%v", rejected, err) - } - budgetless := &HashBuildReservation{core: &hashBuildReservationCore{}} - if budgetless.Released() { - t.Fatal("active budgetless token reported released") - } - - overflowBudget, err := NewHashBuildBudget(math.MaxUint64, math.MaxUint64) - if err != nil { - t.Fatal(err) - } - overflowGeneration, err := overflowBudget.OpenGeneration(1) - if err != nil { - t.Fatal(err) - } - overflow := &HashBuildReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: math.MaxUint64}, - } - if err = overflow.Grow(0); err != nil { - t.Fatal(err) - } - if err = overflow.Grow(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("overflow growth error = %v", err) - } - overflowDisk := &HashBuildSpillDiskReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: math.MaxUint64}, - } - overflowBudget.spillDiskCap = math.MaxUint64 - overflowGeneration.spillDiskCap = math.MaxUint64 - if err = overflowDisk.Grow(0); err != nil { - t.Fatal(err) - } - if err = overflowDisk.Grow(1); !errors.Is(err, ErrHashBuildBudgetInvalid) { - t.Fatalf("overflow disk growth error = %v", err) - } - overflowDisk.core.state.Store(hashBuildReservationReleased) - if err = overflowDisk.Grow(1); !errors.Is(err, ErrHashBuildReservationInactive) { - t.Fatalf("inactive disk growth error = %v", err) - } - cnDiskBudget, err := NewHashBuildBudgetWithSpillCaps(100, 100, 10, 10) - if err != nil { - t.Fatal(err) - } - cnDiskFirst, err := cnDiskBudget.OpenGenerationWithSpillCaps(1, 100, 10, 10) - if err != nil { - t.Fatal(err) - } - cnDiskSecond, err := cnDiskBudget.OpenGenerationWithSpillCaps(2, 100, 10, 10) - if err != nil { - t.Fatal(err) - } - firstDisk, err := cnDiskFirst.ReserveSpillDisk(6) - if err != nil { - t.Fatal(err) - } - secondDisk, err := cnDiskSecond.ReserveSpillDisk(1) - if err != nil { - t.Fatal(err) - } - if err = secondDisk.Grow(4); !errors.Is(err, ErrHashBuildBudgetAdmission) { - t.Fatalf("CN disk growth error = %v", err) - } - firstDisk.Release() - secondDisk.Release() - - corruptMemory := &HashBuildReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: 5}, - } - if ok, reconcileErr := corruptMemory.ReconcileDown(0); ok || !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("corrupt memory reconcile: ok=%v err=%v", ok, reconcileErr) - } - if !corruptMemory.Release() || overflowBudget.AggregateUsed() != 0 || overflowGeneration.Used() != 0 { - t.Fatal("defensive memory release did not clamp corrupt counters") - } - - corruptDisk := &HashBuildSpillDiskReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: 5}, - } - if ok, reconcileErr := corruptDisk.ReconcileDown(0); ok || !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("corrupt disk reconcile: ok=%v err=%v", ok, reconcileErr) - } - if !corruptDisk.Release() { - t.Fatal("defensive disk release failed") - } - corruptFD := &HashBuildSpillFDReservation{ - budget: overflowBudget, - generation: overflowGeneration, - core: &hashBuildReservationCore{size: 5}, - } - if ok, reconcileErr := corruptFD.ReconcileDown(0); ok || !errors.Is(reconcileErr, ErrHashBuildReservationInactive) { - t.Fatalf("corrupt FD reconcile: ok=%v err=%v", ok, reconcileErr) - } - if !corruptFD.Release() { - t.Fatal("defensive FD release failed") - } - - largeHint, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ - HostMemTotal: 10 << 30, - FileCacheHint: 9 << 30, - }) - if err != nil || largeHint.RequestedReserve != 9<<30 { - t.Fatalf("large cache hint ceiling = %+v, err=%v", largeHint, err) - } - tiny, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{HostMemTotal: 128 << 20}) - if err != nil || tiny.CNHashCap == 0 { - t.Fatalf("tiny ceiling = %+v, err=%v", tiny, err) - } - if _, err = ResolveHashBuildCeiling(HashBuildCeilingInputs{HostMemTotal: 1}); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("zero resulting CN cap error = %v", err) - } -} - -func BenchmarkHashBuildBudgetReserveCachedProvider(b *testing.B) { - budget := MustNewHashBuildBudget(uint64(b.N)+1, uint64(b.N)+1) - budget.capRefreshTTL = hashBuildBudgetCapRefreshTTL - gen, _ := budget.OpenGeneration(1) - var calls atomic.Int64 - budget.SetAggregateCapProvider(func() (uint64, error) { - calls.Add(1) - return uint64(b.N) + 1, nil - }) - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - tok, err := gen.Reserve(1) - if err != nil { - b.Fatal(err) - } - tok.Release() - } - b.StopTimer() - b.ReportMetric(float64(calls.Load()), "provider-calls") -} - -func TestResolveHashBuildCeiling(t *testing.T) { - const gib = uint64(1 << 30) - got, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{ - CgroupMemoryMax: 20 * gib, - HostMemTotal: 10 * gib, - GlobalMpoolCap: 30 * gib, - FileCacheHint: gib, - ProcessLimitationSize: 2 * gib, - }) - if err != nil { - t.Fatal(err) - } - if got.EffectiveCN != 10*gib || got.RequestedReserve != 4*gib || got.Reserve != 4*gib || got.CNHashCap != 6*gib || got.QueryCap != 2*gib { - t.Fatalf("unexpected ceiling: %+v", got) - } - if _, err = ResolveHashBuildCeiling(HashBuildCeilingInputs{CgroupMemoryMax: math.MaxUint64, HostMemTotal: 0, GlobalMpoolCap: 0}); !errors.Is(err, ErrHashBuildCeilingMissing) { - t.Fatalf("missing finite source error = %v", err) - } - small, err := ResolveHashBuildCeiling(HashBuildCeilingInputs{HostMemTotal: 3 * gib, FileCacheHint: 3 * gib}) - if err != nil || small.CNHashCap != 3*gib/20 { - t.Fatalf("small-CN bounded allowance = %+v, err=%v", small, err) } } diff --git a/pkg/vm/process/hashbuild_recovery_capacity.go b/pkg/vm/process/hashbuild_recovery_capacity.go new file mode 100644 index 0000000000000..cd362393f95d8 --- /dev/null +++ b/pkg/vm/process/hashbuild_recovery_capacity.go @@ -0,0 +1,144 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package process + +import ( + "errors" + "sync" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" +) + +// HashBuildRecoveryCapacity owns query/CN headroom which physical recovery +// allocations borrow through an allocation-account capacity class. The +// physical allocation remains the sole allocation-ledger owner; borrowing +// prevents the same bytes from being charged to the shared budget twice. +type HashBuildRecoveryCapacity struct { + mu sync.Mutex + + generation *HashBuildBudgetGeneration + capacity uint64 + borrowed uint64 + closed bool +} + +func NewHashBuildRecoveryCapacity( + generation *HashBuildBudgetGeneration, +) (*HashBuildRecoveryCapacity, error) { + if generation == nil || generation.budget == nil || generation.Closed() { + return nil, ErrHashBuildBudgetInvalid + } + return &HashBuildRecoveryCapacity{generation: generation}, nil +} + +// EnsureCapacity raises the reusable recovery floor before HashBuild retains +// a source which may later need that floor to make spill progress. +func (c *HashBuildRecoveryCapacity) EnsureCapacity(target uint64) error { + if c == nil { + return ErrHashBuildBudgetInvalid + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed || c.generation == nil { + return ErrHashBuildSpillReservationInactive + } + if target <= c.capacity { + return nil + } + delta := target - c.capacity + if err := c.generation.acquireMemory(delta); err != nil { + return errors.Join(mpool.ErrAllocationAccountCapacity, err) + } + c.capacity = target + return nil +} + +func (c *HashBuildRecoveryCapacity) AcquireAllocationCapacity(size uint64) error { + if size == 0 { + return nil + } + if c == nil { + return mpool.ErrAllocationAccountInvalid + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed || c.generation == nil { + return errors.Join( + mpool.ErrAllocationAccountSealed, + ErrHashBuildSpillReservationInactive, + ) + } + if c.borrowed > c.capacity || size > c.capacity-c.borrowed { + delta := size + if c.borrowed <= c.capacity { + delta = size - (c.capacity - c.borrowed) + } + if err := c.generation.acquireMemory(delta); err != nil { + return errors.Join(mpool.ErrAllocationAccountCapacity, err) + } + c.capacity += delta + } + c.borrowed += size + return nil +} + +func (c *HashBuildRecoveryCapacity) ReleaseAllocationCapacity(size uint64) { + if size == 0 { + return + } + if c == nil { + panic("nil hash build recovery capacity") + } + c.mu.Lock() + defer c.mu.Unlock() + if size > c.borrowed { + panic("hash build recovery capacity release underflow") + } + c.borrowed -= size +} + +func (c *HashBuildRecoveryCapacity) Snapshot() (capacity, borrowed uint64) { + if c == nil { + return 0, 0 + } + c.mu.Lock() + defer c.mu.Unlock() + return c.capacity, c.borrowed +} + +// Close releases the recovery floor only after all physical borrowers have +// returned it. A failed close keeps ownership intact for terminal diagnostics. +func (c *HashBuildRecoveryCapacity) Close() error { + if c == nil { + return nil + } + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil + } + if c.borrowed != 0 { + return mpool.ErrAllocationAccountLive + } + if c.capacity != 0 { + c.generation.ReleaseAllocationCapacity(c.capacity) + } + c.capacity = 0 + c.generation = nil + c.closed = true + return nil +} + +var _ mpool.AllocationCapacityController = (*HashBuildRecoveryCapacity)(nil) diff --git a/pkg/vm/process/hashbuild_recovery_capacity_test.go b/pkg/vm/process/hashbuild_recovery_capacity_test.go new file mode 100644 index 0000000000000..da1478de058af --- /dev/null +++ b/pkg/vm/process/hashbuild_recovery_capacity_test.go @@ -0,0 +1,147 @@ +// Copyright 2026 Matrix Origin +// +// Licensed 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. + +package process + +import ( + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/stretchr/testify/require" +) + +func TestHashBuildRecoveryCapacityTransfersPhysicalCharge(t *testing.T) { + budget := MustNewHashBuildBudget(1024, 1024) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.OpenWithController(1024, generation) + require.NoError(t, err) + recovery, err := NewHashBuildRecoveryCapacity(generation) + require.NoError(t, err) + class, err := account.RegisterCapacityController(recovery) + require.NoError(t, err) + + require.NoError(t, recovery.EnsureCapacity(256)) + require.Equal(t, uint64(256), generation.Used()) + mp := mpool.MustNewZero() + recoveryBuffer, err := mp.AllocAccountedWithCapacityClass( + 128, account, 1, 1, class) + require.NoError(t, err) + require.Equal(t, uint64(128), account.Snapshot().Used) + // The physical allocation borrows the pre-admitted floor; it does not add a + // second generation charge for the same bytes. + require.Equal(t, uint64(256), generation.Used()) + + defaultBuffer, err := mp.AllocAccounted(100, account, 1, 2) + require.NoError(t, err) + require.Equal(t, uint64(228), account.Snapshot().Used) + require.Equal(t, uint64(356), generation.Used()) + require.NoError(t, recovery.EnsureCapacity(384)) + require.Equal(t, uint64(484), generation.Used()) + require.ErrorIs(t, recovery.Close(), mpool.ErrAllocationAccountLive) + require.ErrorIs( + t, account.UnregisterCapacityController(class, recovery), + mpool.ErrAllocationAccountLive, + ) + + mp.Free(recoveryBuffer) + require.Equal(t, uint64(100), account.Snapshot().Used) + require.Equal(t, uint64(484), generation.Used()) + require.NoError(t, recovery.Close()) + require.Equal(t, uint64(100), generation.Used()) + require.NoError(t, account.UnregisterCapacityController(class, recovery)) + mp.Free(defaultBuffer) + require.Zero(t, account.Snapshot().Used) + require.Zero(t, generation.Used()) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestHashBuildRecoveryCapacityRejectsUncoveredGrowth(t *testing.T) { + budget := MustNewHashBuildBudget(300, 300) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 4) + require.NoError(t, err) + account, err := registry.OpenWithController(300, generation) + require.NoError(t, err) + recovery, err := NewHashBuildRecoveryCapacity(generation) + require.NoError(t, err) + class, err := account.RegisterCapacityController(recovery) + require.NoError(t, err) + mp := mpool.MustNewZero() + defaultBuffer, err := mp.AllocAccounted(100, account, 1, 1) + require.NoError(t, err) + require.NoError(t, recovery.EnsureCapacity(200)) + + _, err = mp.AllocAccountedWithCapacityClass(201, account, 1, 2, class) + require.ErrorIs(t, err, mpool.ErrAllocationAccountCapacity) + require.Equal(t, uint64(100), account.Snapshot().Used) + require.Equal(t, uint64(300), generation.Used()) + capacity, borrowed := recovery.Snapshot() + require.Equal(t, uint64(200), capacity) + require.Zero(t, borrowed) + + require.NoError(t, recovery.Close()) + require.NoError(t, account.UnregisterCapacityController(class, recovery)) + mp.Free(defaultBuffer) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} + +func TestHashBuildRecoveryCapacityIsolatedAcrossWorkers(t *testing.T) { + budget := MustNewHashBuildBudget(512, 512) + generation, err := budget.OpenGeneration(1) + require.NoError(t, err) + registry, err := mpool.NewAllocationAccountRegistry(1, 8) + require.NoError(t, err) + account, err := registry.OpenWithController(512, generation) + require.NoError(t, err) + first, err := NewHashBuildRecoveryCapacity(generation) + require.NoError(t, err) + second, err := NewHashBuildRecoveryCapacity(generation) + require.NoError(t, err) + firstClass, err := account.RegisterCapacityController(first) + require.NoError(t, err) + secondClass, err := account.RegisterCapacityController(second) + require.NoError(t, err) + require.NotEqual(t, firstClass, secondClass) + require.NoError(t, first.EnsureCapacity(200)) + require.NoError(t, second.EnsureCapacity(200)) + + mp := mpool.MustNewZero() + ordinary, err := mp.AllocAccounted(112, account, 1, 1) + require.NoError(t, err) + require.Equal(t, uint64(512), generation.Used()) + firstBuffer, err := mp.AllocAccountedWithCapacityClass( + 200, account, 1, 2, firstClass) + require.NoError(t, err) + secondBuffer, err := mp.AllocAccountedWithCapacityClass( + 200, account, 1, 3, secondClass) + require.NoError(t, err) + require.Equal(t, uint64(512), account.Snapshot().Used) + require.Equal(t, uint64(512), generation.Used()) + + mp.Free(firstBuffer) + mp.Free(secondBuffer) + require.NoError(t, first.Close()) + require.NoError(t, second.Close()) + require.NoError(t, account.UnregisterCapacityController(firstClass, first)) + require.NoError(t, account.UnregisterCapacityController(secondClass, second)) + mp.Free(ordinary) + _, _, err = registry.CompleteTerminal(account) + require.NoError(t, err) +} diff --git a/pkg/vm/process/process.go b/pkg/vm/process/process.go index 65f83c90ff6d7..65fa44429e540 100644 --- a/pkg/vm/process/process.go +++ b/pkg/vm/process/process.go @@ -225,18 +225,46 @@ func (proc *Process) AllocVectorOfRows(typ types.Type, nele int, nsp *nulls.Null } func (proc *Process) NewBatchFromSrc(src *batch.Batch, preAllocSize int) (*batch.Batch, error) { + return proc.NewBatchFromSrcWithAllocation(src, preAllocSize, nil) +} + +// NewBatchFromSrcWithAllocation creates an empty off-heap destination whose +// first vector growth uses the supplied immutable allocation provenance. +func (proc *Process) NewBatchFromSrcWithAllocation( + src *batch.Batch, + preAllocSize int, + selection *vector.AllocationAccountSelection, +) (_ *batch.Batch, retErr error) { + if proc == nil || src == nil || preAllocSize < 0 { + return nil, mpool.ErrAllocationAccountInvalid + } bat := batch.NewOffHeapWithSize(len(src.Vecs)) + defer func() { + if retErr != nil { + bat.Clean(proc.Mp()) + } + }() bat.SetAttributes(src.Attrs) bat.Recursive = src.Recursive for i := range bat.Vecs { - v := vector.NewOffHeapVecWithType(*src.Vecs[i].GetType()) + if src.Vecs[i] == nil { + return nil, mpool.ErrAllocationAccountInvalid + } + bat.Vecs[i] = vector.NewOffHeapVecWithType(*src.Vecs[i].GetType()) + } + if selection != nil { + if err := bat.SetAllocationAccount(selection); err != nil { + return nil, err + } + } + for i := range bat.Vecs { + v := bat.Vecs[i] if v.Capacity() < preAllocSize { err := v.PreExtend(preAllocSize, proc.Mp()) if err != nil { return nil, err } } - bat.Vecs[i] = v } return bat, nil } diff --git a/proto/pipeline.proto b/proto/pipeline.proto index 9df84b6d9ef93..0339ea23ac174 100644 --- a/proto/pipeline.proto +++ b/proto/pipeline.proto @@ -645,6 +645,12 @@ message ProcessInfo { SessionLoggerInfo session_logger = 8 [(gogoproto.nullable) = false]; PrepareParamInfo prepare_params = 9 [(gogoproto.nullable) = false]; int64 affected_rows = 10; + // Planned PipelineMessage RPCs per target CN for statement-level remote + // MessageBoard and resource-account terminal ownership. + map remote_fragment_counts = 11; + // Unique physical execution attempt. Unlike the SQL statement ID, this + // changes across retries and prepared-statement executions. + bytes remote_execution_id = 12; } message SessionInfo {