fix(storage): preserve publication headroom during collection ingest

[thweetkomputer]

Context

Fix the main-branch RESTORE OOM exposed by both architectures in PR #228's CI run. The same GroupedStreamE2e.LargeRepliesKeepSnapshotsAndDeletedHistory failure reproduces on unmodified main a6e93d3d.

Behavior before and after

The collection ingest batching probe covered decoded input, page construction, directories and retirement receipts, but omitted the allocation needed to publish the destination into the grouped side index. A cold arena may need a complete span for its first entry. The batch could therefore consume the headroom needed to finish publication and fail with OOM grouped index exceeds maxmemory.

The debugger captured a 1,163,328-byte publication allocation with only 1,140,345 bytes of worker headroom remaining. The ordinary Redis reply hides this internal reason behind the generic maxmemory error.

The probe now includes the destination index's current allocation requirement. It flushes a smaller batch when needed, allowing the existing 256 MiB test to complete. There is no fixed input-batch size limit, memory-limit increase or weakened assertion.

Implementation

• Add an owner-local PublicationAllocationBytes probe using the same arena/table sizing as PreparePublish; existing slots need no allocation.
• Include that amount with saturating arithmetic in the shared collection-ingest batching decision.
• Check that probing a cold index reports its first-span requirement without inserting a placeholder, and that an already prepared slot requires no additional capacity.
• Keep the original large Stream regression and document why its tight memory budget matters.

Design decisions and alternatives

Use the allocator's current requirement rather than a guessed constant or fraction of available memory. The probe is not a reservation across coroutine suspension; actual publication retains its admission check.

Documentation and comments

API and batching invariants are documented beside the code. Architecture is unchanged: this corrects a local headroom estimate within the existing admitted atomic-ingest flow.

Test plan

Builds completed before tests ran. Debug configuration with fault injection, matching CI:

cmake --build build-fault -j8 --target lavik lavik_grouped_ordered_write_e2e_test
cmake --build build-fault -j8 --target lavik_unit_tests
build-fault/lavik_unit_tests --gtest_filter='GroupedObjectIndexTest.*:CollectionIngestBudget.*'
build-fault/lavik_grouped_ordered_write_e2e_test build-fault/lavik --gtest_filter='GroupedStreamE2e.LargeRepliesKeepSnapshotsAndDeletedHistory'
build-fault/lavik_grouped_ordered_write_e2e_test build-fault/lavik --gtest_filter='GroupedRdbStreamE2e.RestoreStreamsFourTypesInsideAndOutsideExec:GroupedRdbStreamE2e.ZsetRestoreBatchesPreserveOrderAndRollback:GroupedRdbStreamE2e.StartupImportStreamsPagesAndSortsUnorderedZsetInput:GroupedDemotionE2e.StringListSetSortedSetGeoAndStreamRecoverCompact'

The actual invocations supplied an absolute server path and LAVIK_TEST_DATA_DIR on the test volume. Results: 40 unit tests and 5 end-to-end tests passed. pre-commit run --files passed for all five changed files; git diff --check passed.

Risk assessment

An additional owner-local lookup per decoded page and potentially earlier flushing for a cold destination. Concurrent allocations can still exhaust memory after a probe; downstream admission remains authoritative. No durable format, transaction visibility or rollback change.

Rollback plan

Revert this single commit; no data migration is needed.

Reviewer guide

Start at the ingest probe, then compare PublicationAllocationBytes with PreparePublish and the Insert helper's allocation sizing.

Follow-up work

Independent of the throughput work in #228. That PR can incorporate this fix through main after merge.

Summary by CodeRabbit

• Bug Fixes
  • Collection ingestion now includes the memory needed to publish new grouped-object entries when estimating whether a merge fits its available capacity. This gives the capacity check a more complete estimate, including cases where publication requires additional allocation.
  • Publication allocation estimates reflect whether the key already exists, helping distinguish updates from new entries without reserving memory during the check.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: f5081107-95ec-4ba3-a1cd-62c0ce9a20da

📥 Commits

Reviewing files that changed from the base of the PR and between a6e93d3 and c559341.

📒 Files selected for processing (5)

• include/lavik/storage/detail/grouped_object_index.h
• src/storage/engine/collection_ingest.cpp
• src/storage/engine/grouped_object_index.cpp
• tests/grouped_object_index_test.cpp
• tests/grouped_stream_e2e_test.cpp

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

---

Walkthrough

The change adds an allocation estimate for grouped-object publication and uses it in collection-ingest merge admission. Tests check that probing does not create an index entry and that the estimate drops to zero after publication preparation.

Changes

Publication allocation admission

Layer / File(s)	Summary
Add publication allocation probe include/lavik/storage/detail/grouped_object_index.h, src/storage/engine/grouped_object_index.cpp, tests/grouped_object_index_test.cpp	GroupedObjectIndex::PublicationAllocationBytes reports required allocation for a new key and zero for an existing key. Tests check that probing leaves the index unchanged and that the estimate is zero after PreparePublish.
Include allocation in merge admission src/storage/engine/collection_ingest.cpp, tests/grouped_stream_e2e_test.cpp	Collection-ingest merge admission adds the publication allocation estimate to its required-memory calculation. A test comment documents the memory headroom constraint during RESTORE coalescing.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to c5593

The ingest change accounts for the destination publication allocation before merging, with publication retaining its own admission check. No actionable merge-blocking risk remains.

Security Architecture Review

Security architecture risk: 🔵 Low · up to c5593

The change preserves existing memory limits and publication checks while reducing out-of-memory failures during restore. No new security weakness was identified in the inspected paths, but broader access-control and deployment behavior were not assessed.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The security-relevant resource exposure is shared memory availability during existing collection writes. The changed path accounts for additional publication cost rather than raising the memory limit or granting additional write authority. External caller permissions and broader deployment exposure remain unassessed.

Trust Boundaries and Controls

• observed — Caller-provided key identity is checked before ingest proceeds, and publication uses full-key lookup plus expected-handle comparison. These controls reject inconsistent identity and stale publication attempts; the new allocation probe does not replace them or return object payloads.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 7 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the storage fix and its purpose: preserving publication headroom during collection ingest.
Description check	✅ Passed	The description follows the required template and provides context, before-and-after behavior, implementation details, design decisions, documentation status, exact test commands and results, risks, r…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{A rabbit checks the key’s new space,
And counts the bytes before the race.
The index waits, unchanged and still,
Till publish fills the measured bill.
Then hops away with room to spare.}

_{Comment @coderabbitai help to get the list of available commands.}