refactor(core): extract shared StorageBase from RolloutStore - #215
Merged
Conversation
First step of lance-format#214. Lifts the schema-agnostic half of `RolloutStore` into a new `store_base` module so `ContextStore` and `DatagenStore` can be migrated onto it in follow-ups, instead of each carrying its own forked copy. `StorageBase` owns everything that is not schema-specific: - open/create, and the `load_with_options` / `create_with_options` / `DatasetNotFound`-fallback triad that is currently written out four times across the tree - the resident `ShardWriter` lifecycle: lazy open, fence detection and single retry, identity-checked invalidation, `close()`, and the `Drop` fallback that seals an evicted store's memtable - `put` (durable append, no seal) and `flush` (seal + drain), with the add/flush latency and error metrics - the WAL merge: prepare/commit split, surgical generation drain reusing the shard's current epoch, post-drain directory deletion, and both the count and time triggers - compaction with `defer_index_remap`, ZoneMap index management, MemWAL index initialization, shard-manifest discovery, and the LSM scanner `RolloutStore` is now the rollout *schema* plus encode/decode and the typed read APIs, delegating storage to its `base` field. Behavior is unchanged -- this commit is a pure extraction, and rollout is the reference implementation the other two stores will be moved onto. Also deduplicates `align_batch_to_schema`, `derive_shard_id`, `is_fenced_error`, `is_not_found_error`, `ListSource` and `PreparedMerge` into `store_base`; `rollout_store` re-exports them so the public API is unchanged. `DatagenStore` now imports them from `store_base` rather than reaching into `rollout_store`. Net -1145 lines. Full workspace test suite passes unchanged (178 core lib tests, WAL merge concurrency and generation-cleanup integration tests), clippy is clean, and rustdoc warnings drop from 13 to 7. Co-Authored-By: Claude <noreply@anthropic.com>
This was referenced Jul 27, 2026
beinan
added a commit
that referenced
this pull request
Jul 28, 2026
Second step of #214, after #215. `ContextStore` now delegates its storage layer to `StorageBase` instead of carrying its own forked implementation. ## Bugs this closes **WAL merge, which this store never had.** Flushed generations accumulated under `_mem_wal/` forever and every read unioned all of them, so read cost grew without bound in the number of writes. This was the single biggest gap between the three stores. Now exposed as `maybe_merge_wal` (count trigger), `cleanup_wal` (time trigger) and `pending_wal_generations` — the same machinery rollout uses. **Compaction dropped `storage_options`.** Both `compact` and `create_id_index` reloaded via a bare `Dataset::open(uri)`, silently discarding credentials — broken on any authenticated object store. Both now reload through `StorageBase::load_with_options`. **Compaction didn't defer index remap.** The base table carries a fieldless MemWAL index and Lance's inline remap panics on it (`"An index existed with no fields"`). **A writer per `add`.** `write_entries` opened, `put` and `close`d a `ShardWriter` on every call, paying a cold DNS resolution + TCP/TLS handshake + epoch claim each time. Now one resident writer with fence-retry. **`add` is `&self`.** The writer lives behind the base's mutex, so the server takes a read lock and concurrent appends stop serializing on the store's `RwLock`. ## `add` visibility is preserved, not silently weakened This was the one real design fork, and it's worth being explicit since the two stores genuinely differed: - `ContextStore.add` did `put` + `close`, and `close` seals — so rows were **visible on return**, and `validate_unique_ids`, `upsert_by_external_id` and the tombstone path all depend on that read-your-write. - `RolloutStore.add` deliberately gave visibility up: `put` only, with the flush sweeper (30s) or `?flush=true` providing it. Rather than pick one and break the other, `StorageBase` gains a `seal_on_put` switch. Rollout keeps `false`; `ContextStore` defaults to `seal_on_add: true`, preserving today's behavior **and its correctness**. Callers appending trusted batches that don't use the read-back paths can opt into rollout's throughput profile with `seal_on_add: false`. ## Sharding now matches rollout Writes go to the shard owned by this writer instance (`ContextStoreOptions::shard_id`) rather than one derived per `(bot_id, session_id)`. One instance owns exactly one shard, so the epoch-fencing invariant holds by construction and merge has a single well-defined shard to drain. `derive_region_id` is removed. This is a data-layout change: existing rows stay readable (the read path scans all shards from the manifests), but new writes land in a different shard. ## `ContextStore` is no longer `Clone` It owns a `StorageBase`, which owns the resident writer and a `Drop` that seals it — so it must have exactly one owner, or two writers race for one shard. Both clone sites now open their own handle: - the background compactor, which only rewrites base-table fragments and would otherwise contend with the write path - Python's `Context.fork`, which branches the in-memory context over the same dataset The compactor's reopen needed one bit of care: `open_inner` can spawn the compaction task, which reopens via `open_inner` — a mutual recursion between two `async fn`s that the compiler cannot prove `Send`, even though it's unreachable at runtime. Broken with a boxed `dyn Future` in a free function (`open_for_compaction`), with a comment explaining why it can't just be a flag. ## Verification Three new tests pin the behavior that was at risk: - `add_is_visible_on_return_by_default` — the read-your-write contract the uniqueness/upsert/tombstone paths rely on - `deferred_seal_makes_add_durable_but_invisible_until_flush` — mirrors rollout's equivalent test - `wal_generations_merge_into_the_base_table` — generations are reclaimed and rows stay readable Full workspace suite passes (180 core lib tests, up from 177), `cargo clippy --workspace --all-targets` clean.⚠️ **Not verified:** the Python binding tests. `maturin` isn't available in this environment, so `Context.fork`'s signature change is compile-checked only. Worth running `python/tests/test_context.py` and `test_compaction.py` before merge. ## Follow-ups `DatagenStore` is the last store to migrate (it has no compaction at all today), then the arbitrary-schema layer from Part 2 of #214. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude <noreply@anthropic.com>
beinan
added a commit
that referenced
this pull request
Jul 28, 2026
Third and last store of #214 Part 1, after #215 and #216. All three fixed-schema stores now share one storage implementation. Stacked on #216 — review that first; this diff is `datagen_store.rs` only. ## Bugs this closes **Compaction, which this store had none of.** Every WAL merge appends a fragment to the base table, and datagen seals once per checkpoint batch, so fragments grew without bound and scans degraded with them. There was no `compact_files` call anywhere in the file. Now `compact` / `should_compact` / `compaction_stats`, with `defer_index_remap` set — the fieldless MemWAL index panics Lance's inline remap. **Scalar index, which it also had none of.** `DatasetIndexExt` was imported solely for `load_indices` in `mem_wal_index_present`. `create_event_id_index` now builds a ZoneMap on `event_id`; point lookups by event id previously always scanned. **`load_with_options` dropped storage options.** It fell back to a bare `Dataset::open` whenever no options were passed, and never attached a session — so every flushed-generation read allocated Lance's 6 GiB default cache, keyed per generation URI. ## Deleted duplication `merge_own_shard` was a near-line-for-line copy of rollout's, and buffered every flushed generation fully in memory. Gone, along with `merge_own_shard_if_ready`, the writer lifecycle (`write_with_resident_writer` / `put_seal_drain` / `ensure_write_writer`), `ensure_mem_wal`, `mem_wal_index_present`, shard-snapshot discovery, `flushed_generation_uri`, `open_flushed_dataset`, `load_with_options`, `create_with_options`, the `Drop` impl, and a second copy of `is_fenced_error`. Net **-269 lines** in this file. ## Behavior preserved exactly `seal_on_put: true`. Each append is one complete checkpoint batch that must be durable *and* readable before the writer moves on, so a crash cannot expose a partially checkpointed step. This is the store's existing `put` → `force_seal_active` → `wait_for_flush_drain` behavior, unchanged. `key_column` is `event_id`, not `id`. Event ids are derived deterministically (UUIDv5 over item/checkpoint/ordinal), which is what makes a retried append idempotent under LSM dedup. The base already took the key column as a parameter, so it needed no change to accommodate this. ## Verification New test `compaction_folds_merged_fragments_and_preserves_reads` pins behavior that could not exist before: fragments accumulate across merges, compaction reduces the count, and every event stays readable afterward. - 181 core lib tests pass (up from 180) - 182 Python tests pass - `cargo clippy --workspace --all-targets` clean ## What's left in #214 Part 1 is done after this — all three stores share `StorageBase`. Part 2 (arbitrary user-defined schemas with a required `id`) builds on it: the base already validates `key_column` against the schema at open, which is the enforcement point that work needs. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
First step of #214.
Lifts the schema-agnostic half of
RolloutStoreinto a newstore_basemodule, soContextStoreandDatagenStorecan be migrated onto it in follow-ups instead of each carrying its own forked copy of the same machinery.Rollout is the reference. This PR is a pure extraction — no behavior change. The point is to establish the shared implementation with rollout's semantics before the other two stores are moved onto it, since those migrations are where the actual behavior changes land (
ContextStoregaining WAL merge and seal-on-flush,DatagenStoregaining compaction and an index).What
StorageBaseownsload_with_options/create_with_options/DatasetNotFound-fallback triad that is currently written out four times across the treeclose(), and theDropfallback that seals an evicted store's memtableput(durable append, no seal) andflush(seal + drain), with the add/flush latency and error metricsdefer_index_remap, ZoneMap index management, MemWAL index initialization, shard-manifest discovery, and the LSM scannerThe two things the base needs to stay schema-agnostic are passed via
StorageBaseOptions:key_column(the LSM merge key and indexed column) and an optionallatest_schemafor merge-time additive evolution.key_columnis validated against the schema at open — which is also the hook Part 2 of #214 will use to enforce the requiredidfield.What stays in
RolloutStoreThe rollout schema:
rollout_schema(),records_to_batch/batch_to_rollout_records,RolloutFilters, the SQL path, and the typed read APIs. Storage is delegated to itsbasefield.Also deduplicated
align_batch_to_schema,derive_shard_id,is_fenced_error,is_not_found_error,ListSourceandPreparedMergeall move intostore_base.rollout_storere-exports them so the public API is unchanged, andDatagenStorenow imports them fromstore_baserather than reaching intorollout_store— removing the odd dependency where the datagen store pulled helpers out of the rollout store.Verification
wal_merge_concurrencyandwal_merge_generation_cleanupintegration tests — the ones that actually exercise the epoch/fence/drain invariantscargo clippy --workspace --all-targetscleanFollow-ups (from #214)
Migrate
ContextStore(biggest behavior delta: it currently never merges WAL and never seals on write) andDatagenStore(gains compaction, which it has none of today) onto this base, then build the arbitrary-schema layer on top.🤖 Generated with Claude Code