In-Mem 2.0

[hildebrandmw]

Introduce a second in-memory provider with the intention of replacing the current provider.

Why

The RFC outlines much of the motivation. In short, the goal here is to:

• Make the provider safe under concurrent inserts/search/deletes.
• Support proper external/internal ID translation.
• Improve test coverage
• Do so with minimal performance overhead.

The concurrency argument comes from the epoch-based-reclamation (EBR) protection scheme for internal slots. See the RFC for more details.

Known Follow-up Items

• Perf Parity: This generally has pretty good performance, but could use a little more tuning to bring it fully on-par with our current inmem index.
• Quantization: This initial prototype is lacking quantization support. Adding quantization is relatively straightforward in the primary Store, but will need some thought on how to add the reranking layer. This shouldn't be an architectural blocker, though.
• Hybrid PQ: A larger open question is how to support the max_fp_vecs feature of our current PQ implementation, which reads some full-precision and some quantized vectors during prune. Like with quantization in general, I think this is not a fundamental issue.
• Support for non-uniform sized items in slots: For this, I'm mainly thinking of multi-vectors where the number of vectors within each multi-vector can vary. In the context of multi-vectors, fast element access is less important than for traditional vectors as distance computations take considerably longer.

Suggested Reviewing Order

The majority of this PR is in a new diskann-inmem crate. To facilitate testing, this crate has an "integration-test" feature, which enables the code in diskann-inmem/src/integration. This code is an unstable public reexport of internal types meant only for consumption in the diskann-inmem/integration integration test binary.

The integration test binary is powered by diskann-benchmark-runner.

diskann-inmem

Independent Low-Level Utilities

• num.rs: Strong type utilities for byte and alignment representations.
• buffer.rs: A miri-compliant version of AlignedMemoryVectorStore. This type allows vectors/neighbors to be stored in a single larger allocation. The use of RawSlice allows slots within the Buffer to be inspected and manipulated without forcing reference materialization (which is important to prevent aliasing).
• neighbors.rs: The new version of SimpleNeighborVectorProviderAsync. This reuses the sharded-lock idea, but provides additional utility, including the ability to perform read-modify-write operations on adjacency lists.
• counters.rs: Event counters. When the "integration-test" feature is not enabled, counters become a no-op. These are enabled for testing to monitor changes.
• sharded.rs: An external-to-internal ID translation utility. The main trick with this struct is to provide utilities like Sharded::occupied_entry, which locks and returns an external/internal mapping. The proxy Entry struct is important as it verifies that such a mapping exists and provides an infallible way of deleting the mapping. This is chained with higher level operations (e.g. Provider::delete) to delete both the ID-mapping and the internal data-slot in lock-step.

Concurrency Protocol

The concurrency protocol is built upon three main layers:

• tag.rs: An atomic slot tag for controlling access to data.
• epoch.rs: The central registry where readers register and deregister. This is the crux of this PR and probably the most important file.
• store.rs: A binary blob store built on top of epoch.rs to provide the safe concurrent store for data. This provides the following operations:
  • Storage of binary data in "slots".
  • Reading of data in slots (provided by Reader).
  • Tracking on the valid/invalid state slots.
  • Finding available slots into which new data can be inserted.
  • Safe retirement of slots and eventual reclamation.

The Store in store.rs has some help from freelist.rs to accelerate locating available slots internally.

Testing: epoch.rs has unit tests with injectable delays to set up known pathological orderings. The sequencing is helped by test/sequencer.rs. In addition, test/epoch.rs includes a direct stress test for the Registry. This is particularly helpful when run under Miri, which has the ability to detect some race conditions.

A larger concurrency stress test lives in the integration-test binary. This directly tests store.rs by spinning up readers, writers, and retirers and hammers a single Store. Data is read and written into the store in a knowable pattern, allowing readers to detect torn reads, implying a race condition. For this PR, I ran the following stress test file

{
  "search_directories": [
  ],
  "output_directory": null,
  "jobs": [
    {
      "type": "store-stress",
      "content": {
        "capacity": 8192,
        "duration_secs": 600,
        "entry_bytes": 256,
        "epoch_guard_slots": 256,
        "freelist_recycle_capacity": 1024,
        "low_watermark": 4096,
        "max_ops": 50000000000,
        "readers": 32,
        "retirers": 16,
        "seed": 11935966405698895599,
        "writers": 32
      }
    }
  ]
}

with the command

cargo run --package diskann-inmem \
  --bin integration-test \
  --features integration-test \
  --release -- \
   run --input-file stress.json --output-file temp.json

The generated output was

readers:       32
writers:       32
retirers:      16
capacity:      8192
entry_bytes:   256
low_watermark: 4096
duration_secs: 600
max_ops:       50000000000
seed:          11935966405698895599
elapsed_secs:  600.001182962
reads:         110206469888
acquires_ok:   248086767
acquires_fail: 619145
retires_ok:    248082573
retires_fail:  243312644
reclaims:      108847566
transitions:   42006520
peak_live:     8131

While not a proof of correctness, this is a pretty decent stress test.

Providers

The implementation of the data provider is split into two logical pieces. The first lives in layers/ and is focused on computing distances. With this approach, I am trying to avoid the need to replicate Accessors and Strategys for each future quantization type. layers/full.rs is the full-precision implementation. One thing to note is the use of the FullPrecision marker trait from which the implementations of the whole layers API is derived for layers::Full. This allows users to include just a T: FullPrecision trait bound and really simplifies the generics upstream.

Within provider.rs - my goal here is to minimize the use of generics as much as possible. In particular for search, I use a trait object for ExpandBeam. When coupled with the layers::QueryDistance API, we can create implementations where the distance function is inlined and the number of prefetch instructions can be tailored to the data length. Fully optimizing this is still a work-in-progress.

Another thing to call out in provider.rs is the care needed for data insertion and deletion. Since external/internal ID translation is supported, we need to ensure that the translation table stays in-sync with the internal store. On insert, if we allocate an internal slot only to find the external ID already exists, we need to abort the operation rather than publish the internal slot. Similarly on delete, we first need to establish if the external/internal ID mapping exists. If so, then we can try to retire the slot. If slot retiring fails (it shouldn't, but bugs can happen) - we need to not commit the ID mapping deletion and instead return an error.

Testing

Like the concurrency stress tests, testing uses the integration-test binary. Here, the 10k YFCC dataset is used for non-trivial runs. Using the A/B functionality in diskann-benchmark-runner, we can compare against checked-in baselines. This allows us to capture rich metrics for recall, number of operations, etc. and update easily. The main logic for checking and reporting baseline mismatches is in integration/support/check.rs. The goal is to summarize all such mismatches for presentation to provide the highest signal possible.

diskann-benchmark

Integration into the benchmarks is straightforward. I elected to put everything in a single file to minimize disruption. I'm trying an approach of using diskann_benchmark_runner::Input::from_raw to separate out the deserialization types from the actual inputs, allowing richer types (e.g., a full diskann_benchmark_core::streaming::bigann::RunBook) to be loaded.

Also note how relatively simple the streaming benchmark integration is. Since the new inmem provider supports ID translation and internal slot allocation, it does not need the same level of hand-holding that the current inmem provider needs.

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (1)

rfcs/01206-inmem2.md:33

• Typo: “mechansim” should be “mechanism”.

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 6 comments.

Comments suppressed due to low confidence (2)

rfcs/01206-inmem2.md:33

• Typo in RFC text: "mechansim" should be "mechanism".
  rfcs/01206-inmem2.md:52
• Grammar: "to facilitated searches" reads incorrectly; should be "to facilitate searches".

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

diskann-inmem/integration/main.rs:101

• The panic message refers to search-directories, but the JSON key you look up is search_directories. This makes failures harder to diagnose when an input file is malformed.

            let value = obj
                .get_mut(key)
                .expect("key \"search-directories\" should exist");
            if let serde_json::Value::Array(directories) = value {

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

diskann-inmem/src/buffer.rs:35

• Buffer::new docs claim the allocation size is rounded up to the next multiple of align, but the implementation passes bytes.value() directly to Layout::from_size_align without any rounding. This is misleading for callers trying to reason about allocation size limits; the error condition is simply that the requested size exceeds what Layout accepts (in practice isize::MAX).

    /// Returns an error if the number of bytes `bytes_per_entry * entries` rounded up to
    /// the next multiple of `align` exceeds `isize::MAX`.

rfcs/01206-inmem2.md:33

• Spelling: "mechansim" should be "mechanism".

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

rfcs/01206-inmem2.md:33

• Typo: "mechansim" → "mechanism".

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

rfcs/01206-inmem2.md:33

• Spelling/wording: "mechansim" is misspelled, and "Yoloing" is informal wording for an RFC. Consider using a more precise description of the trade-off.

[Copilot]

Pull request overview

Copilot reviewed 47 out of 48 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

diskann-inmem/src/buffer.rs:35

• The Buffer::new docs mention rounding the allocation size up to the next multiple of align, but the implementation allocates exactly bytes_per_entry * entries bytes (with align only affecting alignment). The error condition should describe the actual size check.

    /// Returns an error if the number of bytes `bytes_per_entry * entries` rounded up to
    /// the next multiple of `align` exceeds `isize::MAX`.