feat(state_store): automatically resize LMDB map when full

[slliland]

Summary

This PR adds automatic LMDB map resizing when a write transaction fails with MDB_MAP_FULL.

Previously, users had to increase COCOINDEX_LMDB_MAP_SIZE and restart the process manually. With this change, the configured value remains the initial map size, but the state store can recover automatically by growing the map and retrying the failed transaction.

Closes #2108.

Changes

When a write fails with heed::MdbError::MapFull, the state store now:

1. Drops the failed write transaction.
2. Acquires an exclusive resize guard.
3. Reads the current map size from Env::info().
4. Doubles the size with overflow checking and page-size alignment.
5. Resizes the LMDB environment.
6. Retries the complete transaction batch.

Transaction bodies were changed from FnOnce to Fn + Send + Sync so they can be retried. Affected call sites now clone or share their inputs instead of consuming them.

A shared TxnCoordinator coordinates transactions and resizing:

• read and write transactions hold a shared guard;
• resizing requires an exclusive guard;
• ReadTxn keeps its guard alive for the lifetime of the underlying LMDB transaction.

clear_staged_tracking, the remaining ordinary direct-write path, was also moved through the batcher so it uses the same resize-and-retry behavior.

Motivation

LMDB requires a fixed map size. Once that limit is reached, writes fail with MDB_MAP_FULL, even when additional disk space is available.

Automatically growing the map avoids requiring users to estimate future storage needs, update the environment variable, and restart CocoIndex when the initial allocation becomes too small.

Testing

Added tests covering:

• automatic growth from a deliberately small map;
• successful retry and data verification after MDB_MAP_FULL;
• resizing while an active reader is present;
• map-size page alignment;
• clear_staged_tracking with matching, non-matching, and missing tokens.

Validation performed:

prek run --all-files --show-diff-on-failure

All project checks passed, including formatting, type checking, Rust and Python tests, example builds, dependency validation, and generated documentation checks.

The focused state-store test suite also passes with 14 tests:

cargo test -p cocoindex_core state_store

Breaking changes

None. COCOINDEX_LMDB_MAP_SIZE keeps its existing meaning as the initial LMDB map size, and this PR does not change the user-facing API or configuration format.

Follow-up considerations

A few related areas were considered but intentionally kept separate from this change:

• Cross-process MDB_MAP_RESIZED handling could be added if multiple processes need to share the same LMDB environment.
• A configurable maximum map size could provide a clearer failure policy in disk-constrained environments.
• Low-frequency environment-management writes such as create_app_store and drop_app could also use the same retry mechanism for complete coverage.

These are useful follow-ups but they are not required for the main MDB_MAP_FULL recovery path addressed here.

[georgeh0]

Thanks for the PR!

A quick question: if the MDB size is already increased from the default. In the next run (after a new process starts), what will happen?

[georgeh0]

This is an Arc<RwLock<()>>, which detaches the lock from the actual underlying resource (the db_env) it wants to protect. Is it possible to use a Arc<RwLock<heed::Env<heed::WithoutTls>>> instead?

[slliland]

Thanks! I used Arc<RwLock<()>> because the lock coordinates transaction lifetimes with Env::resize(), rather than protecting ownership of the Env handle itself.
Since heed::Env is cloneable and already shared across the state store, wrapping one handle in Arc<RwLock> would only be effective if every LMDB access were also refactored to go through that lock. In the current design, transactions hold a shared guard for their full lifetime, while resize takes the exclusive guard.
One possible improvement would be a CoordinatedEnv wrapper that keeps Env private and exposes only guarded operations. Do you think that would be preferable here, or is the current barrier-based approach acceptable for this change?

[georgeh0]

Makes sense.

Now my only remaining suggestion is to directly use the original type Arc<RwLock<()>> instead of TxnCoordinator here. I feel this is clearer: it's used as a coordinator but only within the context of StorageInner, which is already clear by the field name.

[slliland]

Yes I removed the TxnCoordinator alias and use Arc<tokio::sync::RwLock<()>> directly on StorageInner.coord (and TxnRunner)

[slliland]

After the map auto-resizes, the committed data stays on disk. On the next run, we still open LMDB with the configured lmdb_map_size, and LMDB raises the effective map size to at least the space already used by the existing database. That may be smaller than the previous process's map limit, but it is enough to read the persisted data.

If it fills up again, we resize from the current env.info().map_size, not from the original default. So the user does not need to update COCOINDEX_LMDB_MAP_SIZE manually.

[georgeh0]

I don't think we need two separate branches here.

Dropping for wtxn and read_guard happen any way (by RAII) and they needs to happen for both cases no matter if the error is caused by map full. A ? may be just sufficient?

[slliland]

I simplified this to use ?. RAII now drops both the write transaction and the read guard on any error, while the outer retry loop still handles MDB_MAP_FULL.

[georgeh0]

Similar here, I think we don't need separate out the is_map_full() branch. A ? may be just sufficient?

[slliland]

I removed the separate is_map_full handling here as well and now propagate the commit error with ?.

[georgeh0]

Makes sense.

Now my only remaining suggestion is to directly use the original type Arc<RwLock<()>> instead of TxnCoordinator here. I feel this is clearer: it's used as a coordinator but only within the context of StorageInner, which is already clear by the field name.

[georgeh0]

After the map auto-resizes, the committed data stays on disk. On the next run, we still open LMDB with the configured lmdb_map_size, and LMDB raises the effective map size to at least the space already used by the existing database. That may be smaller than the previous process's map limit, but it is enough to read the persisted data.

If it fills up again, we resize from the current env.info().map_size, not from the original default. So the user does not need to update COCOINDEX_LMDB_MAP_SIZE manually.

During past experience, I roughly remembered LMDB may complain on DB open time if the initial size given to it is too small. But it may be on an older version and may not be very accurate.

Did you happen to get a chance to really test it out (either manual running, or by a written test)? Thanks!

[badmonster0]

@slliland wonder how did the project go? please keep us posted if you have any questions!

[slliland]

@slliland wonder how did the project go? please keep us posted if you have any questions!

I’m currently doing a final review based on the comments. I should be able to submit an updated version within the next day

[slliland]

After the map auto-resizes, the committed data stays on disk. On the next run, we still open LMDB with the configured lmdb_map_size, and LMDB raises the effective map size to at least the space already used by the existing database. That may be smaller than the previous process's map limit, but it is enough to read the persisted data.
If it fills up again, we resize from the current env.info().map_size, not from the original default. So the user does not need to update COCOINDEX_LMDB_MAP_SIZE manually.

During past experience, I roughly remembered LMDB may complain on DB open time if the initial size given to it is too small. But it may be on an older version and may not be very accurate.

Did you happen to get a chance to really test it out (either manual running, or by a written test)? Thanks!

Yes I tested this locally with a two-process setup. The first process opened the environment with a deliberately small 256 KiB map, wrote 64 × 16 KiB values through Storage::run_txn, triggered auto-resize, committed, and exited. The map had grown to 4 MiB by the end of that process.

The second process then reopened the same database path through Storage::new using the original 256 KiB config. It opened successfully with an effective map size of about 2.1 MiB, and the committed keys were still readable.

So the reopened map was smaller than the previous process’s 4 MiB limit, but still large enough for the persisted data. I couldn’t reproduce an open-time failure with the current heed/LMDB version. I didn’t include the automated subprocess test in this PR to keep the test surface small, but I’m happy to add it if you think CI coverage would be useful.

[georgeh0]

After the map auto-resizes, the committed data stays on disk. On the next run, we still open LMDB with the configured lmdb_map_size, and LMDB raises the effective map size to at least the space already used by the existing database. That may be smaller than the previous process's map limit, but it is enough to read the persisted data.
If it fills up again, we resize from the current env.info().map_size, not from the original default. So the user does not need to update COCOINDEX_LMDB_MAP_SIZE manually.

During past experience, I roughly remembered LMDB may complain on DB open time if the initial size given to it is too small. But it may be on an older version and may not be very accurate.
Did you happen to get a chance to really test it out (either manual running, or by a written test)? Thanks!

Yes I tested this locally with a two-process setup. The first process opened the environment with a deliberately small 256 KiB map, wrote 64 × 16 KiB values through Storage::run_txn, triggered auto-resize, committed, and exited. The map had grown to 4 MiB by the end of that process.

The second process then reopened the same database path through Storage::new using the original 256 KiB config. It opened successfully with an effective map size of about 2.1 MiB, and the committed keys were still readable.

So the reopened map was smaller than the previous process’s 4 MiB limit, but still large enough for the persisted data. I couldn’t reproduce an open-time failure with the current heed/LMDB version. I didn’t include the automated subprocess test in this PR to keep the test surface small, but I’m happy to add it if you think CI coverage would be useful.

Thanks for verifying carefully!

PR looks good! Merging!