diff --git a/cpp/test.cpp b/cpp/test.cpp index 3ea2f86e..15744346 100644 --- a/cpp/test.cpp +++ b/cpp/test.cpp @@ -55,9 +55,11 @@ #define SZ_USE_X86_AVX512 0 // Sanitizers hate AVX512 #include // Levenshtein distance implementation +#include #include #include #include +#include using namespace unum::usearch; using namespace unum; @@ -1306,6 +1308,412 @@ static void install_crash_handlers() { std::signal(signal_number, &usearch_crash_handler); } +/** + * @brief Resident-set size of the current process in bytes, or 0 if it + * cannot be read (non-Linux). Reads `/proc/self/statm`, whose second + * field is the resident page count. + */ +std::size_t current_rss_bytes() { +#if defined(__linux__) + std::FILE* file = std::fopen("/proc/self/statm", "r"); + if (!file) + return 0; + unsigned long total_pages = 0, resident_pages = 0; + int scanned = std::fscanf(file, "%lu %lu", &total_pages, &resident_pages); + std::fclose(file); + return scanned == 2 ? static_cast(resident_pages) * 4096u : 0; +#else + return 0; +#endif +} + +/** + * @brief Exercises ::global_rebuild_gt: a non-blocking, interruptible global + * rebuild of an HNSW index that persists the result to disk. Checks + * that the reconstructed graph keeps recall close to the original, + * and that the rebuild does @b not double the process RSS - the + * shadow shares vectors with the primary and only rebuilds the graph. + */ +void test_global_rebuild() { + std::printf("Testing global rebuild\n"); + + using index_t = index_dense_gt; + std::size_t const dimensions = 256; // large enough that vector RAM is a + std::size_t const collection = 12000; // clear, measurable share of the RSS + std::size_t const queries = 200; + std::size_t const wanted = 10; + + std::default_random_engine rng(42); + std::uniform_real_distribution distribution(-1.f, 1.f); + auto make_vector = [&](std::vector& vector) { + vector.resize(dimensions); + for (auto& value : vector) + value = distribution(rng); + }; + + // Clustered dataset: uniformly random vectors in a high dimension all look + // equidistant (curse of dimensionality), so nearest-neighbor search has no + // signal. Drawing vectors as a centroid plus small noise gives genuine + // neighborhood structure - and thus a meaningful recall to compare. + std::size_t const clusters = 64; + std::vector> centroids(clusters); + for (auto& centroid : centroids) + make_vector(centroid); + auto make_clustered = [&](std::vector& vector, std::size_t cluster) { + vector.resize(dimensions); + for (std::size_t d = 0; d != dimensions; ++d) + vector[d] = centroids[cluster][d] + 0.15f * distribution(rng); + }; + + std::vector> data(collection), query(queries); + for (std::size_t i = 0; i != collection; ++i) + make_clustered(data[i], i % clusters); + for (std::size_t i = 0; i != queries; ++i) + make_clustered(query[i], i % clusters); + + // Brute-force cosine ground truth: the top-`wanted` keys per query. + auto cosine_distance = [&](float const* a, float const* b) { + double dot = 0, norm_a = 0, norm_b = 0; + for (std::size_t i = 0; i != dimensions; ++i) + dot += a[i] * b[i], norm_a += a[i] * a[i], norm_b += b[i] * b[i]; + double denominator = std::sqrt(norm_a) * std::sqrt(norm_b); + return denominator > 0 ? 1.0 - dot / denominator : 1.0; + }; + std::vector> truth(queries); + for (std::size_t q = 0; q != queries; ++q) { + std::vector> ranked(collection); + for (std::size_t i = 0; i != collection; ++i) + ranked[i] = {cosine_distance(query[q].data(), data[i].data()), static_cast(i)}; + std::partial_sort(ranked.begin(), ranked.begin() + wanted, ranked.end()); + for (std::size_t k = 0; k != wanted; ++k) + truth[q].push_back(ranked[k].second); + } + + // Recall@`wanted` of an index against the brute-force ground truth. + auto recall_of = [&](index_t& index) { + std::size_t hits = 0; + for (std::size_t q = 0; q != queries; ++q) { + std::int64_t found[32]; + std::size_t count = index.search(query[q].data(), wanted).dump_to(found); + for (std::size_t k = 0; k != count; ++k) + for (std::size_t t = 0; t != wanted; ++t) + if (found[k] == truth[q][t]) { + ++hits; + break; + } + } + return double(hits) / double(queries * wanted); + }; + + // Build the live primary index. + metric_punned_t metric(dimensions, metric_kind_t::cos_k, scalar_kind()); + index_dense_config_t config(16); + index_t::state_result_t primary_result = index_t::make(metric, config); + expect(primary_result); + index_t& primary = primary_result.index; + expect(primary.try_reserve(collection + 64)); + for (std::size_t i = 0; i != collection; ++i) + expect(primary.add(static_cast(i), data[i].data())); + + double recall_baseline = recall_of(primary); + std::printf("- baseline recall@%zu: %.3f\n", wanted, recall_baseline); + expect(recall_baseline > 0.5); + + // Drive a non-blocking global rebuild, interleaving live mutations between + // the budgeted steps. + using rebuild_t = global_rebuild_gt; + rebuild_t rebuild(primary, 128); + char const* path = "tmp_global_rebuild.usearch"; + + // Pre-place a sentinel file at the destination. A crash-safe rebuild must + // leave it byte-for-byte intact until the very end - it streams into a + // temporary file and publishes only via an atomic rename - so a process + // kill mid-rebuild can never corrupt the previous index. + std::vector const sentinel(96, '\x7e'); + { + std::FILE* sentinel_file = std::fopen(path, "wb"); + expect(sentinel_file != nullptr); + std::fwrite(sentinel.data(), 1, sentinel.size(), sentinel_file); + std::fclose(sentinel_file); + } + auto read_path = [&]() { + std::vector bytes; + std::FILE* file = std::fopen(path, "rb"); + if (!file) + return bytes; + char buffer[256]; + for (std::size_t n; (n = std::fread(buffer, 1, sizeof(buffer), file)) > 0;) + bytes.insert(bytes.end(), buffer, buffer + n); + std::fclose(file); + return bytes; + }; + + // Watch the process RSS across the rebuild: with the shadow sharing the + // primary's vectors it must not balloon by a whole vector-set. + std::size_t const primary_vectors_bytes = primary.memory_stats().vectors_allocated; + std::size_t const rss_before = current_rss_bytes(); + std::size_t rss_peak = rss_before; + + expect(rebuild.begin(path)); + + std::vector extra_a, extra_b; + make_vector(extra_a); + make_vector(extra_b); + std::int64_t const new_key_a = static_cast(collection) + 1; + std::int64_t const new_key_b = static_cast(collection) + 2; + std::int64_t const removed_key = 7; + bool side_ops_done = false; + bool shadow_checked = false; + while (rebuild.active()) { + expect(rebuild.step()); + rss_peak = (std::max)(rss_peak, current_rss_bytes()); + if (rebuild.phase() == rebuild_t::phase_saving_k && !shadow_checked) { + // Migration is complete: the shadow holds a fully rebuilt graph, + // but - thanks to the zero-copy migration - not one duplicated + // vector. Its nodes alias the primary's vector storage. + expect(rebuild.shadow() != nullptr); + index_t::memory_stats_t shadow_memory = rebuild.shadow()->memory_stats(); + expect_eq(shadow_memory.vectors_allocated, std::size_t(0)); + expect(shadow_memory.graph_allocated > 0); + std::printf("- shadow memory: graph %zu, vectors %zu (primary vectors %zu)\n", + shadow_memory.graph_allocated, shadow_memory.vectors_allocated, + primary.memory_stats().vectors_allocated); + // Mid-rebuild, with the shadow already being streamed, the + // destination still holds the untouched sentinel - the bytes go + // to `.tmp`, not to `path`. + expect(read_path() == sentinel); + shadow_checked = true; + } + if (!side_ops_done) { + // Adds hit the primary right away; the remove is tombstoned until + // the rebuild finishes, so the on-disk snapshot stays untouched. + expect(rebuild.add(new_key_a, extra_a.data())); + expect(rebuild.add(new_key_b, extra_b.data())); + expect(rebuild.remove(removed_key)); + expect_eq(rebuild.deferred_remove_count(), 1ul); + expect(primary.contains(removed_key)); // Still deferred. + side_ops_done = true; + } + } + expect(rebuild.finished()); + expect(side_ops_done); + expect(shadow_checked); + expect(rebuild.shadow() == nullptr); // released at completion + // Completion atomically replaced the sentinel with the real index. + expect(read_path() != sentinel); + + // RSS must not double during the rebuild. The shadow shares the primary's + // vectors and only rebuilds the graph, so its peak overhead is one graph - + // comfortably under a full vector-set duplicate, which is what a naive + // clone-everything rebuild would have added. + if (rss_before) { + std::printf("- RSS: before %zu, peak %zu (primary vectors %zu)\n", // + rss_before, rss_peak, primary_vectors_bytes); + expect(rss_peak < rss_before * 3 / 2); // grew by well under 50% + expect(rss_peak - rss_before < primary_vectors_bytes); // vectors not duplicated + } + + // The tombstoned removal is applied now; the mid-rebuild adds are live. + expect(!primary.contains(removed_key)); + expect(primary.contains(new_key_a)); + expect(primary.contains(new_key_b)); + + // Load the persisted snapshot and confirm the rebuilt graph's accuracy. + index_t::state_result_t loaded_result = index_t::make(path); + expect(loaded_result); + index_t& loaded = loaded_result.index; + // The snapshot is the `begin()` generation: every original vector, and + // none of the keys added once the rebuild was already in flight. + expect_eq(loaded.size(), collection); + expect(loaded.contains(removed_key)); + expect(!loaded.contains(new_key_a)); + + double recall_rebuilt = recall_of(loaded); + std::printf("- rebuilt recall@%zu: %.3f\n", wanted, recall_rebuilt); + // A global rebuild reconstructs the graph from scratch, so recall shifts + // slightly; it must not drop materially below the original. + expect(recall_rebuilt > recall_baseline - 0.05); + + std::remove(path); +} + +/** + * @brief Exercises ::persistent_index_gt - snapshot + WAL durability around + * `index_dense_gt`. Checks fresh open + add + search; clean reopen + * replays the WAL and recovers state; manual checkpoint advances the + * generation and retires the old WAL; truncating the WAL tail still + * recovers everything but the lost record (crash-resilience). + */ +void test_persistent_index() { + std::printf("Testing persistent index (snapshot + WAL)\n"); + + using index_t = index_dense_gt; + using pi_t = persistent_index_gt; + std::size_t const dimensions = 64; + std::size_t const initial_count = 200; + std::size_t const after_reopen_count = 50; + char const* base = "tmp_pi_test"; + + // Clean up any leftover from a previous run so the test is hermetic. + std::remove((std::string(base) + ".manifest").c_str()); + for (int g = 0; g < 4; ++g) { + std::remove((std::string(base) + "." + std::to_string(g) + ".snapshot").c_str()); + std::remove((std::string(base) + "." + std::to_string(g) + ".wal").c_str()); + } + + std::default_random_engine rng(1729); + std::uniform_real_distribution distribution(-1.f, 1.f); + auto make_vec = [&](std::vector& vector) { + vector.resize(dimensions); + for (auto& value : vector) + value = distribution(rng); + }; + + metric_punned_t metric(dimensions, metric_kind_t::cos_k, scalar_kind()); + index_dense_config_t index_config(16); + pi_t::config_t pi_config; + pi_config.checkpoint_after_ops = 10'000'000; // no auto-checkpoint in this test + pi_config.initial_capacity = initial_count + after_reopen_count + 64; + + // ---- 1. Fresh open: add `initial_count` vectors, then close. ---- + std::vector> data; + { + pi_t::open_result_t opened = pi_t::open(base, metric, index_config, pi_config); + expect(opened); + pi_t& pi = *opened.index; + expect_eq(pi.size(), std::size_t(0)); + expect_eq(pi.generation(), std::uint64_t(0)); + data.resize(initial_count); + for (std::size_t i = 0; i != initial_count; ++i) { + make_vec(data[i]); + expect(pi.add(static_cast(i), data[i].data())); + } + expect_eq(pi.size(), initial_count); + } + + // ---- 2. Reopen: snapshot + WAL replay must restore the full state. ---- + { + pi_t::open_result_t reopened = pi_t::open(base, metric, index_config, pi_config); + expect(reopened); + pi_t& pi = *reopened.index; + expect_eq(pi.size(), initial_count); + expect_eq(pi.generation(), std::uint64_t(0)); // still gen 0 + // The recovered vectors must be the original bytes. + std::vector probe(dimensions); + for (std::size_t i = 0; i < initial_count; i += 17) { + expect_eq(pi.get(static_cast(i), probe.data()), std::size_t(1)); + expect(std::equal(data[i].begin(), data[i].end(), probe.data())); + } + // A search returns sane results too. + std::int64_t found[8]; + std::size_t n = pi.search(data[0].data(), 5).dump_to(found); + expect(n != 0); + } + + // ---- 3. Reopen, add more, manual checkpoint -> generation increments. ---- + { + pi_t::open_result_t reopened = pi_t::open(base, metric, index_config, pi_config); + expect(reopened); + pi_t& pi = *reopened.index; + data.resize(initial_count + after_reopen_count); + for (std::size_t i = initial_count; i != initial_count + after_reopen_count; ++i) { + make_vec(data[i]); + expect(pi.add(static_cast(i), data[i].data())); + } + // Remove a few keys so the checkpoint actually compacts something. + for (std::size_t i = 0; i < 10; ++i) + expect_eq(pi.remove(static_cast(i)).completed, std::size_t(1)); + + auto err = pi.checkpoint(); + expect(!err); + expect(pi.generation() == 1); + expect(!pi.checkpoint_in_flight()); + + // Old generation's files must be gone, new ones must exist. + std::FILE* old_snap = std::fopen((std::string(base) + ".0.snapshot").c_str(), "rb"); + expect(old_snap == nullptr); + std::FILE* new_snap = std::fopen((std::string(base) + ".1.snapshot").c_str(), "rb"); + expect(new_snap != nullptr); + if (new_snap) + std::fclose(new_snap); + } + + // ---- 4. Reopen post-checkpoint: snapshot.1 + (empty) wal.1 round-trips. ---- + std::size_t const expected_size = initial_count + after_reopen_count - 10; + { + pi_t::open_result_t reopened = pi_t::open(base, metric, index_config, pi_config); + expect(reopened); + pi_t& pi = *reopened.index; + expect_eq(pi.generation(), std::uint64_t(1)); + expect_eq(pi.size(), expected_size); + std::vector probe(dimensions); + // Surviving key still recovers the exact original bytes. + std::int64_t survivor = static_cast(initial_count + 3); + expect(pi.contains(survivor)); + expect_eq(pi.get(survivor, probe.data()), std::size_t(1)); + expect(std::equal(data[survivor].begin(), data[survivor].end(), probe.data())); + // Removed key stays gone. + expect(!pi.contains(0)); + + // Append one more op so we have something to truncate-test below. + std::vector extra; + make_vec(extra); + expect(pi.add(9'999'999, extra.data())); + } + + // ---- 5. Crash-recover from a torn WAL tail. ---- + // Truncate the active WAL by a few bytes - simulates a torn final record + // (a partial write that did not flush). Recovery must drop the torn tail + // and load everything else cleanly. + { + std::string wal_path = std::string(base) + ".1.wal"; + std::FILE* file = std::fopen(wal_path.c_str(), "rb"); + expect(file != nullptr); + std::fseek(file, 0, SEEK_END); + long size = std::ftell(file); + std::fclose(file); + // Drop the last 5 bytes - guaranteed to mangle the trailing record. + expect(size > 5); + int rc = ::truncate(wal_path.c_str(), size - 5); + expect(rc == 0); + + pi_t::open_result_t reopened = pi_t::open(base, metric, index_config, pi_config); + expect(reopened); + pi_t& pi = *reopened.index; + // The trailing add (key 9'999'999) was torn off; everything else stays. + expect(!pi.contains(9'999'999)); + expect_eq(pi.size(), expected_size); + } + + // ---- 6. Auto-checkpoint fires when ops cross the configured threshold. ---- + { + pi_t::config_t small_cfg = pi_config; + small_cfg.checkpoint_after_ops = 50; + small_cfg.checkpoint_step_budget = 64; + pi_t::open_result_t reopened = pi_t::open(base, metric, index_config, small_cfg); + expect(reopened); + pi_t& pi = *reopened.index; + std::uint64_t starting_gen = pi.generation(); + std::vector probe(dimensions); + for (std::size_t i = 0; i != 200; ++i) { + make_vec(probe); + expect(pi.add(static_cast(2'000'000 + i), probe.data())); + } + // Auto-trigger must have started a checkpoint mid-way through the 200 + // adds; finish any in-flight one and check the generation moved on. + auto err = pi.checkpoint(); + expect(!err); + expect(pi.generation() > starting_gen); + } + + // Cleanup + std::remove((std::string(base) + ".manifest").c_str()); + for (int g = 0; g < 4; ++g) { + std::remove((std::string(base) + "." + std::to_string(g) + ".snapshot").c_str()); + std::remove((std::string(base) + "." + std::to_string(g) + ".wal").c_str()); + } +} + /** * @brief Regression test: `make(metric, config)` must return an index that is * immediately usable - no explicit `reserve` required before `load` / @@ -1453,5 +1861,7 @@ int main(int, char**) { test_filtered_search(); test_isolate(); test_load_after_metric_make(); + test_persistent_index(); + test_global_rebuild(); return 0; } diff --git a/include/usearch/global_rebuild.hpp b/include/usearch/global_rebuild.hpp new file mode 100644 index 00000000..8c6bcfbf --- /dev/null +++ b/include/usearch/global_rebuild.hpp @@ -0,0 +1,232 @@ +/** + * @file global_rebuild.hpp + * @author Mikhail Chichvarin + * @brief Non-blocking @b global-rebuild orchestrator for `index_dense_gt`. + * + * Persists the HNSW index to disk without a stop-the-world `save`. Rebuilds + * it into a fresh "shadow" peer by re-insertion, then streams that frozen + * shadow to disk through the resumable `save_to_stream`. Both phases run in + * budgeted steps; the live "primary" keeps serving reads and writes. + * + * 1. `phase_migrating` - re-insert the key-set captured at `begin()` + * into the shadow, with `copy_vector = false` so shadow nodes alias + * the primary's vector bytes (extra RAM during a rebuild is one + * graph, not a full clone). + * 2. `phase_saving` - stream the now-frozen shadow into `.tmp`. + * 3. `phase_done` - atomic `rename` onto ``, release the shadow, + * replay tombstoned removes. + * + * Concurrent mutations: `add` always hits the primary; `remove` during a + * rebuild is tombstoned and replayed at completion, so the on-disk snapshot + * equals the `begin()` generation exactly. + * + * Crash safety: writes go through `.tmp` and become canonical only on + * the `rename` (atomic on POSIX). A crash at any point leaves either the + * previous file or the new file, never a truncated one. This is not resume- + * across-restart: the continuation cursor lives in RAM, so a killed rebuild + * is restarted from `begin`, not continued. + */ +#ifndef UNUM_USEARCH_GLOBAL_REBUILD_HPP +#define UNUM_USEARCH_GLOBAL_REBUILD_HPP + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace unum { +namespace usearch { + +/// @brief Orchestrates an interruptible, non-blocking global rebuild of a +/// dense index, persisting a freshly reconstructed copy to disk. +/// `scalar_at` must match the index's stored scalar kind - the +/// migration reinterprets the primary's raw vector bytes as it. +template // +class global_rebuild_gt { + public: + using index_t = index_at; + using scalar_t = scalar_at; + using vector_key_t = typename index_t::vector_key_t; + using add_result_t = typename index_t::add_result_t; + using labeling_result_t = typename index_t::labeling_result_t; + using search_result_t = typename index_t::search_result_t; + + enum phase_t { + phase_idle_k = 0, + phase_migrating_k = 1, + phase_saving_k = 2, + phase_done_k = 3, + }; + + struct result_t { + error_t error{}; + explicit operator bool() const noexcept { return !error; } + result_t failed(error_t message) noexcept { + error = std::move(message); + return std::move(*this); + } + }; + + private: + index_t* primary_ = nullptr; + std::unique_ptr shadow_; + std::size_t budget_ = 256; + phase_t phase_ = phase_idle_k; + + std::vector migration_keys_; + std::size_t migration_cursor_ = 0; + + // Written into `.tmp`, renamed onto `` only on completion. + std::string final_path_; + std::string temp_path_; + output_file_t file_{nullptr}; + index_dense_serialized_state_t save_state_; + + std::vector deferred_removes_; + + public: + /// @param[in] step_budget Vectors migrated, or vectors/nodes serialized, + /// per `step()`. Smaller = shorter pauses. + explicit global_rebuild_gt(index_t& primary, std::size_t step_budget = 256) noexcept + : primary_(&primary), budget_(step_budget ? step_budget : 1) {} + + global_rebuild_gt(global_rebuild_gt const&) = delete; + global_rebuild_gt& operator=(global_rebuild_gt const&) = delete; + + ~global_rebuild_gt() { + // Abandoned mid-rebuild: the destination was never touched, so just + // discard the temp file. + if (active()) { + file_.close(); + std::remove(temp_path_.c_str()); + } + } + + phase_t phase() const noexcept { return phase_; } + bool active() const noexcept { return phase_ == phase_migrating_k || phase_ == phase_saving_k; } + bool finished() const noexcept { return phase_ == phase_done_k; } + std::size_t deferred_remove_count() const noexcept { return deferred_removes_.size(); } + /// @brief Populated during the rebuild, released (null) at `phase_done_k`. + /// Vectors alias the primary's storage - do not outlive it. + index_t const* shadow() const noexcept { return shadow_.get(); } + + /// @brief Always routed to the primary, never blocked. + template + add_result_t add(vector_key_t key, scalar_other_at const* vector) { return primary_->add(key, vector); } + + /// @brief Tombstoned during an active rebuild, replayed at completion, + /// so the on-disk snapshot stays equal to the `begin()` generation. + labeling_result_t remove(vector_key_t key) { + if (!active()) + return primary_->remove(key); + deferred_removes_.push_back(key); + labeling_result_t result; + result.completed = 1; + return result; + } + + /// @brief Begin a global rebuild, persisting the result to @p path. + result_t begin(char const* path) { + result_t result; + if (active()) + return result.failed("A global rebuild is already in flight"); + // Zero-copy migration reinterprets the primary's stored bytes as + // `scalar_t`, so the adapter's scalar type must match. + if (primary_->scalar_kind() != unum::usearch::scalar_kind()) + return result.failed("Adapter scalar type must match the index's stored scalar kind"); + + std::size_t live = primary_->size(); + migration_keys_.resize(live); + if (live) + primary_->export_keys(migration_keys_.data(), 0, live); + migration_cursor_ = 0; + + typename index_t::copy_result_t forked = primary_->fork(); + if (!forked) + return result.failed(std::move(forked.error)); + shadow_.reset(new (std::nothrow) index_t(std::move(forked.index))); + if (!shadow_) + return result.failed("Out of memory for the shadow index"); + if (live && !shadow_->try_reserve(live)) + return result.failed("Failed to reserve the shadow index"); + + final_path_ = path; + temp_path_ = final_path_ + ".tmp"; + file_ = output_file_t(temp_path_.c_str()); + serialization_result_t io = file_.open_if_not(); + if (!io) + return result.failed(std::move(io.error)); + + save_state_ = index_dense_serialized_state_t{}; + deferred_removes_.clear(); + phase_ = phase_migrating_k; + return result; + } + + /// @brief Advance the rebuild by one budgeted chunk. Idempotent when + /// idle or finished. Closes the file and replays tombstones on + /// the step that completes the save. + result_t step() { + result_t result; + + if (phase_ == phase_migrating_k) { + std::size_t migrated = 0; + while (migrated < budget_ && migration_cursor_ < migration_keys_.size()) { + vector_key_t key = migration_keys_[migration_cursor_++]; + byte_t const* vector = primary_->vector_data(key); + if (!vector) + continue; // key already gone (a deferred-remove race-with-self) + // Zero-copy: shadow aliases primary's vector bytes. + add_result_t added = shadow_->add(key, reinterpret_cast(vector), + index_t::any_thread(), /*copy_vector=*/false); + if (!added) + return result.failed(std::move(added.error)); + ++migrated; + } + if (migration_cursor_ >= migration_keys_.size()) + phase_ = phase_saving_k; + return result; + } + + if (phase_ == phase_saving_k) { + serialization_result_t io; + serialization_result_t saved = shadow_->save_to_stream( + [&](void const* buffer, std::size_t length) { + io = file_.write(buffer, length); + return !!io; + }, + save_state_, budget_); + if (!saved) + return result.failed(std::move(saved.error)); + if (save_state_.done()) { + file_.close(); + // Atomic publish on POSIX; Windows `rename` cannot overwrite, + // hence the fallback with its tiny non-atomic window there. + if (std::rename(temp_path_.c_str(), final_path_.c_str()) != 0) { + std::remove(final_path_.c_str()); + if (std::rename(temp_path_.c_str(), final_path_.c_str()) != 0) + return result.failed("Failed to publish the rebuilt index file"); + } + for (std::size_t i = 0; i != deferred_removes_.size(); ++i) + primary_->remove(deferred_removes_[i]); + // Shadow nodes alias the primary's vectors - don't outlive it. + shadow_.reset(); + phase_ = phase_done_k; + } + return result; + } + + return result; + } +}; + +} // namespace usearch +} // namespace unum + +#endif // UNUM_USEARCH_GLOBAL_REBUILD_HPP diff --git a/include/usearch/index.hpp b/include/usearch/index.hpp index b4e12be1..aca8f7dd 100644 --- a/include/usearch/index.hpp +++ b/include/usearch/index.hpp @@ -2124,6 +2124,38 @@ struct index_serialized_header_t { std::uint64_t entry_slot = 0; }; +/** + * @brief Resumable cursor for @b incremental, interruptible serialization. + * + * The plain `save_to_stream` writes the whole index in one blocking call. To + * persist a large graph to disk without a stop-the-world pause, the resumable + * `save_to_stream` overload writes the graph in bounded chunks: each call + * emits up to a caller-provided budget of nodes and then returns, leaving this + * cursor pointing at the next node. Feed the same cursor back in to continue. + * + * The node count is @b frozen into `total` on the first call - the caller is + * responsible for not structurally mutating the graph (adding nodes, changing + * node levels) while a save is in flight. Appending brand-new nodes past + * `total` is harmless: they are simply not part of this snapshot. + */ +struct index_serialized_state_t { + enum stage_t : std::uint8_t { + stage_header_k = 0, + stage_levels_k = 1, + stage_nodes_k = 2, + stage_done_k = 3, + }; + /// @brief Which section of the file is being written next. + stage_t stage = stage_header_k; + /// @brief Index of the next node to emit within the current stage. + std::uint64_t cursor = 0; + /// @brief Node count captured when the save began; frozen for its lifetime. + std::uint64_t total = 0; + + bool begun() const noexcept { return stage != stage_header_k; } + bool done() const noexcept { return stage == stage_done_k; } +}; + using default_key_t = std::uint64_t; using default_slot_t = std::uint32_t; using default_distance_t = float; @@ -3707,6 +3739,75 @@ class index_gt { return {}; } + /** + * @brief Resumable, interruptible variant of `save_to_stream`. + * + * Writes at most @p node_budget nodes worth of data per call, then returns + * while updating @p state. Keep calling with the same @p state until + * `state.done()` becomes true. The output layout is byte-identical to the + * blocking `save_to_stream`, so the resulting file loads with the regular + * `load_from_stream`. + * + * The graph must @b not be structurally mutated between calls (see + * ::index_serialized_state_t). Appending new nodes past the frozen count + * is allowed and simply excluded from this snapshot. + */ + template + serialization_result_t save_to_stream(output_callback_at&& output, index_serialized_state_t& state, + std::size_t node_budget) const noexcept { + + serialization_result_t result; + std::size_t budget = node_budget; + + // Stage 1: the fixed-size header. Freezes the node count for the + // remainder of the save. + if (state.stage == index_serialized_state_t::stage_header_k) { + index_serialized_header_t header; + header.size = nodes_count_; + header.connectivity = config_.connectivity; + header.connectivity_base = config_.connectivity_base; + header.max_level = max_level_; + header.entry_slot = entry_slot_; + if (!output(&header, sizeof(header))) + return result.failed("Failed to serialize the header into stream"); + state.total = header.size; + state.cursor = 0; + state.stage = index_serialized_state_t::stage_levels_k; + } + + // Stage 2: one `level_t` per node, enough to size every node on load. + while (state.stage == index_serialized_state_t::stage_levels_k) { + if (state.cursor == state.total) { + state.cursor = 0; + state.stage = index_serialized_state_t::stage_nodes_k; + continue; + } + if (!budget) + return result; + node_t node = node_at_(static_cast(state.cursor)); + level_t level = node.level(); + if (!output(&level, sizeof(level))) + return result.failed("Failed to serialize into stream"); + ++state.cursor, --budget; + } + + // Stage 3: the node tapes themselves. + while (state.stage == index_serialized_state_t::stage_nodes_k) { + if (state.cursor == state.total) { + state.stage = index_serialized_state_t::stage_done_k; + continue; + } + if (!budget) + return result; + span_bytes_t node_bytes = node_bytes_(node_at_(static_cast(state.cursor))); + if (!output(node_bytes.data(), node_bytes.size())) + return result.failed("Failed to serialize into stream"); + ++state.cursor, --budget; + } + + return result; + } + /** * @brief Symmetric to `save_from_stream`, pulls data from a stream. */ diff --git a/include/usearch/index_dense.hpp b/include/usearch/index_dense.hpp index b0d6766a..1c0d6455 100644 --- a/include/usearch/index_dense.hpp +++ b/include/usearch/index_dense.hpp @@ -386,6 +386,44 @@ inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_map return result.failed("Not a dense USearch index!"); } +/** + * @brief Resumable cursor for @b incremental, interruptible serialization of + * an `index_dense_gt`. + * + * Mirrors ::index_serialized_state_t but covers all three sections of a dense + * index file: the vector matrix, the dense metadata header, and the embedded + * HNSW graph. The resumable `save_to_stream` overload writes a bounded chunk + * per call and updates this cursor; loop until `done()` is true. + * + * The vector count and the live/deleted tallies are @b frozen on the first + * call so that the matrix, the header, and the graph all agree even if the + * index keeps accepting brand-new keys (in higher slots) meanwhile. + */ +struct index_dense_serialized_state_t { + enum stage_t : std::uint8_t { + stage_dimensions_k = 0, + stage_vectors_k = 1, + stage_head_k = 2, + stage_graph_k = 3, + stage_done_k = 4, + }; + /// @brief Which section of the file is being written next. + stage_t stage = stage_dimensions_k; + /// @brief Index of the next vector row to emit during ::stage_vectors_k. + std::uint64_t cursor = 0; + /// @brief Vector / node count frozen when the save began. + std::uint64_t total = 0; + /// @brief Live (non-deleted) vector count frozen when the save began. + std::uint64_t count_present = 0; + /// @brief Bytes per stored vector, frozen when the save began. + std::uint64_t bytes_per_vector = 0; + /// @brief Resumable cursor for the embedded HNSW graph. + index_serialized_state_t graph; + + bool begun() const noexcept { return stage != stage_dimensions_k; } + bool done() const noexcept { return stage == stage_done_k; } +}; + /** * @brief Oversimplified type-punned index for equidimensional vectors * with automatic @b down-casting, hardware-specific @b SIMD metrics, @@ -813,6 +851,27 @@ class index_dense_gt { return typed_->level_of((*matching_slots.first).slot); } + /** + * @brief Returns a pointer to the @b stored vector bytes for @p key, or + * `nullptr` if the key is absent. + * + * The pointer aliases the index's own vector tape: it stays valid only + * until that entry is removed, the index is compacted, or the index is + * destroyed, and the bytes are in the index's native scalar layout + * (see `scalar_kind`). Useful for a @b zero-copy migration into a peer + * index of the same configuration - the peer can reference these bytes + * via `add(..., copy_vector = false)` instead of duplicating them. For a + * multi-index the first matching entry is returned. + */ + byte_t const* vector_data(vector_key_t key) const { + usearch_assert_m(config().enable_key_lookups, "Key lookups are disabled"); + shared_lock_t lookup_lock(slot_lookup_mutex_); + auto matching_slots = slot_lookup_.equal_range(key_and_slot_t::any_slot(key)); + if (matching_slots.first == matching_slots.second) + return nullptr; + return vectors_lookup_[(*matching_slots.first).slot]; + } + dynamic_allocator_t const& allocator() const { return typed_->dynamic_allocator(); } vector_key_t const& free_key() const { return free_key_; } @@ -1209,6 +1268,111 @@ class index_dense_gt { return typed_->save_to_stream(std::forward(output), std::forward(progress)); } + /** + * @brief Resumable, interruptible variant of `save_to_stream`. + * + * Writes the dense index to @p output in bounded chunks: each call emits + * at most @p budget vectors (during the matrix stage) or nodes (during the + * graph stage), then returns while updating @p state. Keep calling with + * the same @p state until `state.done()` is true. The byte layout matches + * the blocking `save_to_stream`, so the file loads with the regular + * `load_from_stream` / `load`. + * + * The index must not be @b structurally mutated between calls: no removed + * slot may be recycled and no node level may change. Appending brand-new + * keys is fine - they land in higher slots and are excluded from the + * frozen snapshot. The ::global_rebuild_gt adapter enforces this for the + * shadow index it persists. + */ + template + serialization_result_t save_to_stream(output_callback_at&& output, // + index_dense_serialized_state_t& state, // + std::size_t budget, // + serialization_config_t config = {}) const { + + serialization_result_t result; + + // Stage 1: the matrix dimensions. Freezes the snapshot's vector count, + // live tally, and per-vector stride for every later stage. + if (state.stage == index_dense_serialized_state_t::stage_dimensions_k) { + state.total = typed_->size(); + state.count_present = size(); + state.bytes_per_vector = metric_.bytes_per_vector(); + state.cursor = 0; + if (!config.exclude_vectors) { + if (!config.use_64_bit_dimensions) { + std::uint32_t dimensions[2]; + dimensions[0] = static_cast(state.total); + dimensions[1] = static_cast(state.bytes_per_vector); + if (!output(&dimensions, sizeof(dimensions))) + return result.failed("Failed to serialize into stream"); + } else { + std::uint64_t dimensions[2]; + dimensions[0] = state.total; + dimensions[1] = state.bytes_per_vector; + if (!output(&dimensions, sizeof(dimensions))) + return result.failed("Failed to serialize into stream"); + } + state.stage = index_dense_serialized_state_t::stage_vectors_k; + } else + state.stage = index_dense_serialized_state_t::stage_head_k; + } + + // Stage 2: the vector matrix, one frozen row at a time. + while (state.stage == index_dense_serialized_state_t::stage_vectors_k) { + if (state.cursor == state.total) { + state.stage = index_dense_serialized_state_t::stage_head_k; + continue; + } + if (!budget) + return result; + byte_t* vector = vectors_lookup_[state.cursor]; + if (!output(vector, state.bytes_per_vector)) + return result.failed("Failed to serialize into stream"); + ++state.cursor, --budget; + } + + // Stage 3: the dense metadata header, using the frozen tallies. + if (state.stage == index_dense_serialized_state_t::stage_head_k) { + index_dense_head_buffer_t buffer; + std::memset(buffer, 0, sizeof(buffer)); + index_dense_head_t head{buffer}; + std::memcpy(buffer, default_magic(), std::strlen(default_magic())); + + using version_t = index_dense_head_t::version_t; + head.version_major = static_cast(USEARCH_VERSION_MAJOR); + head.version_minor = static_cast(USEARCH_VERSION_MINOR); + head.version_patch = static_cast(USEARCH_VERSION_PATCH); + + head.kind_metric = metric_.metric_kind(); + head.kind_scalar = metric_.scalar_kind(); + head.kind_key = unum::usearch::scalar_kind(); + head.kind_compressed_slot = unum::usearch::scalar_kind(); + + head.count_present = state.count_present; + head.count_deleted = state.total - state.count_present; + head.dimensions = dimensions(); + head.multi = multi(); + + if (!output(&buffer, sizeof(buffer))) + return result.failed("Failed to serialize into stream"); + state.stage = index_dense_serialized_state_t::stage_graph_k; + } + + // Stage 4: the embedded HNSW graph, delegated to the resumable + // `index_gt::save_to_stream`, which carries its own sub-cursor. + if (state.stage == index_dense_serialized_state_t::stage_graph_k) { + serialization_result_t graph_result = + typed_->save_to_stream(std::forward(output), state.graph, budget); + if (!graph_result) + return graph_result; + if (state.graph.done()) + state.stage = index_dense_serialized_state_t::stage_done_k; + } + + return result; + } + /** * @brief Estimate the binary length (in bytes) of the serialized index. */ diff --git a/include/usearch/persistent_index.hpp b/include/usearch/persistent_index.hpp new file mode 100644 index 00000000..0bdf7d5e --- /dev/null +++ b/include/usearch/persistent_index.hpp @@ -0,0 +1,529 @@ +/** + * @file persistent_index.hpp + * @author Mikhail Chichvarin + * @brief Snapshot + WAL persistence wrapper around `index_dense_gt`. + * + * On-disk (base name ``): + * ..snapshot - full index file (written by `global_rebuild_gt`, + * crash-safe via temp + rename); + * ..wal - append-only log, framed + * `[u32 len][u32 crc32][u8 op][key][vec]`; + * .manifest - atomic pointer to the current generation `g`. + * + * Order is WAL-first: append the record, then apply to RAM. A crash between + * the two replays the record on recovery, so committed ops are not lost + * beyond what is still in the OS page cache. No `fsync` - durability is + * best-effort vs hard crashes, traded for throughput. + * + * Auto-checkpoint when ops since last cross `checkpoint_after_ops`: + * begin - under `wal_mutex_`, drain in-flight index applies, snapshot + * the WAL offset as a watermark, `rebuild.begin(snapshot.)`; + * step - each mutating op drives one budgeted `rebuild.step()` + * (try-locked, so only one stepper at a time); + * complete - under `wal_mutex_`, splice `wal.[watermark..end]` onto + * a fresh `wal.` (header + suffix), atomically commit + * the manifest, retire the old snapshot and WAL. + * Both barriers are ~ms. A crash before the manifest commit leaves the + * previous generation canonical. + */ +#ifndef UNUM_USEARCH_PERSISTENT_INDEX_HPP +#define UNUM_USEARCH_PERSISTENT_INDEX_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace unum { +namespace usearch { + +/// @brief Durable wrapper around `index_dense_gt`: snapshot + WAL with +/// non-blocking auto-checkpoints via `global_rebuild_gt`. +template // +class persistent_index_gt { + public: + using index_t = index_at; + using scalar_t = scalar_at; + using vector_key_t = typename index_t::vector_key_t; + using add_result_t = typename index_t::add_result_t; + using labeling_result_t = typename index_t::labeling_result_t; + using search_result_t = typename index_t::search_result_t; + using rebuild_t = global_rebuild_gt; + using metric_t = typename index_t::metric_t; + + struct config_t { + std::size_t checkpoint_after_ops = 100'000; + std::size_t checkpoint_step_budget = 256; + std::size_t initial_capacity = 1024; + }; + + struct open_result_t { + std::unique_ptr index; + error_t error; + explicit operator bool() const noexcept { return !error; } + }; + + private: + std::string base_path_; + config_t config_; + index_t index_; + rebuild_t rebuild_; + + std::uint64_t generation_ = 0; + std::uint64_t dimensions_ = 0; + std::uint32_t scalar_kind_ = 0; + std::uint32_t metric_kind_ = 0; + std::size_t vector_bytes_ = 0; + + // WAL state. All file writes happen under `wal_mutex_`. + std::mutex wal_mutex_; + file_t wal_file_; + std::string wal_path_; + std::uint64_t wal_bytes_ = 0; + std::vector record_buffer_; + + std::atomic ops_since_checkpoint_{0}; + std::atomic in_flight_{0}; + std::mutex checkpoint_mutex_; + std::atomic checkpoint_active_{false}; + std::uint64_t checkpoint_watermark_ = 0; + std::uint64_t checkpoint_gen_ = 0; + error_t checkpoint_error_{}; + + persistent_index_gt(char const* path, config_t config) + : base_path_(path), config_(config), rebuild_(index_, config.checkpoint_step_budget) {} + + public: + persistent_index_gt(persistent_index_gt const&) = delete; + persistent_index_gt& operator=(persistent_index_gt const&) = delete; + + ~persistent_index_gt() { + // Best-effort flush; the OS will write the rest of the page cache out. + // An in-flight checkpoint is just dropped - its temp file is cleaned + // by `rebuild_`'s dtor; the manifest still points at the previous + // generation and recovery is consistent. `wal_file_` closes itself. + if (wal_file_) + std::fflush(wal_file_.get()); + } + + /// @brief Open (or recover) a persistent index at @p path. + static open_result_t open(char const* path, metric_t metric, index_dense_config_t index_config, + config_t config = {}) { + open_result_t result; + std::unique_ptr self(new (std::nothrow) persistent_index_gt(path, config)); + if (!self) + return {nullptr, error_t("Out of memory for persistent_index")}; + error_t err = self->open_(std::move(metric), index_config); + if (err) { + result.error = std::move(err); + return result; + } + result.index = std::move(self); + return result; + } + + add_result_t add(vector_key_t key, scalar_t const* vector) { + // WAL-first: a crash between WAL append and index apply replays the + // record on recovery, so a committed-to-WAL op is never lost. + { + std::unique_lock lock(wal_mutex_); + if (!wal_file_) + return add_result_t{}.failed("WAL is not open"); + append_record_(persistent_op_add_k, key, reinterpret_cast(vector)); + in_flight_.fetch_add(1, std::memory_order_acq_rel); + ops_since_checkpoint_.fetch_add(1, std::memory_order_relaxed); + } + add_result_t r = index_.add(key, vector); + in_flight_.fetch_sub(1, std::memory_order_acq_rel); + maybe_drive_checkpoint_(); + return r; + } + + labeling_result_t remove(vector_key_t key) { + { + std::unique_lock lock(wal_mutex_); + if (!wal_file_) { + labeling_result_t r; + return r.failed("WAL is not open"); + } + append_record_(persistent_op_remove_k, key, nullptr); + in_flight_.fetch_add(1, std::memory_order_acq_rel); + ops_since_checkpoint_.fetch_add(1, std::memory_order_relaxed); + } + labeling_result_t r = index_.remove(key); + in_flight_.fetch_sub(1, std::memory_order_acq_rel); + maybe_drive_checkpoint_(); + return r; + } + + search_result_t search(scalar_t const* query, std::size_t wanted) const { return index_.search(query, wanted); } + std::size_t get(vector_key_t key, scalar_t* out, std::size_t count = 1) const { + return index_.get(key, out, count); + } + bool contains(vector_key_t key) const { return index_.contains(key); } + std::size_t size() const noexcept { return index_.size(); } + std::uint64_t generation() const noexcept { return generation_; } + bool checkpoint_in_flight() const noexcept { return checkpoint_active_.load(std::memory_order_acquire); } + /// @brief Drive a checkpoint to completion (starts one if needed). + error_t checkpoint() { + if (!checkpoint_active_.load(std::memory_order_acquire)) + start_checkpoint_(); + if (checkpoint_error_) + return std::move(checkpoint_error_); + while (checkpoint_active_.load(std::memory_order_acquire)) { + drive_checkpoint_step_(); + if (checkpoint_error_) + return std::move(checkpoint_error_); + std::this_thread::yield(); + } + return {}; + } + + private: + std::string gen_path_(char const* suffix, std::uint64_t g) const { + return base_path_ + "." + std::to_string(g) + suffix; + } + std::string manifest_path_() const { return base_path_ + ".manifest"; } + + /// @brief Atomically replace the manifest via temp + rename. On POSIX + /// the rename is atomic; the Windows fallback has a tiny window. + error_t write_manifest_(std::uint64_t generation) { + persistent_manifest_t manifest{}; + std::memcpy(manifest.magic, persistent_manifest_magic_k, 4); + manifest.format_version = persistent_wal_version_k; + manifest.generation = generation; + std::string path = manifest_path_(); + std::string tmp = path + ".tmp"; + std::size_t written; + { + file_t file{std::fopen(tmp.c_str(), "wb")}; + if (!file) + return error_t("Failed to create manifest temp file"); + written = std::fwrite(&manifest, 1, sizeof(manifest), file.get()); + } + if (written != sizeof(manifest)) { + std::remove(tmp.c_str()); + return error_t("Short write on manifest temp file"); + } + if (std::rename(tmp.c_str(), path.c_str()) != 0) { + std::remove(path.c_str()); + if (std::rename(tmp.c_str(), path.c_str()) != 0) + return error_t("Failed to rename manifest"); + } + return {}; + } + + void refresh_metadata_() { + dimensions_ = index_.dimensions(); + scalar_kind_ = static_cast(index_.scalar_kind()); + metric_kind_ = static_cast(index_.metric_kind()); + vector_bytes_ = dimensions_ * sizeof(scalar_t); + } + + /// @brief Build a WAL header reflecting current metadata. + persistent_wal_header_t make_wal_header_() const { + persistent_wal_header_t header{}; + std::memcpy(header.magic, persistent_wal_magic_k, 4); + header.format_version = persistent_wal_version_k; + header.dimensions = dimensions_; + header.scalar_kind = scalar_kind_; + header.metric_kind = metric_kind_; + return header; + } + + /// @brief Open @p path in append mode, learning its current size. + error_t open_wal_for_append_(std::string const& path) { + file_t file{std::fopen(path.c_str(), "ab")}; + if (!file) + return error_t("Failed to open WAL for append"); + long size = -1; + if (file_t probe{std::fopen(path.c_str(), "rb")}) { + std::fseek(probe.get(), 0, SEEK_END); + size = std::ftell(probe.get()); + } + if (size < 0) + return error_t("Failed to stat WAL file"); + wal_file_ = std::move(file); + wal_path_ = path; + wal_bytes_ = static_cast(size); + record_buffer_.reserve(1 + sizeof(vector_key_t) + vector_bytes_); + return {}; + } + + /// @brief Build and append one record. Caller holds `wal_mutex_`. + void append_record_(persistent_op_t op, vector_key_t key, char const* vector_or_null) { + std::uint32_t payload_len = static_cast( // + 1 + sizeof(vector_key_t) + (vector_or_null ? vector_bytes_ : 0)); + record_buffer_.resize(payload_len); + record_buffer_[0] = static_cast(op); + std::memcpy(&record_buffer_[1], &key, sizeof(vector_key_t)); + if (vector_or_null) + std::memcpy(&record_buffer_[1 + sizeof(vector_key_t)], vector_or_null, vector_bytes_); + std::uint32_t crc = crc32_ieee(record_buffer_.data(), payload_len); + std::FILE* out = wal_file_.get(); + std::fwrite(&payload_len, sizeof(payload_len), 1, out); + std::fwrite(&crc, sizeof(crc), 1, out); + std::fwrite(record_buffer_.data(), 1, payload_len, out); + wal_bytes_ += sizeof(payload_len) + sizeof(crc) + payload_len; + } + + /// @brief Either recover from an existing manifest or initialize fresh. + error_t open_(metric_t metric, index_dense_config_t index_config) { + // Probe for an existing manifest. Absence = fresh open. + std::uint64_t generation = 0; + bool manifest_exists = false; + if (file_t file{std::fopen(manifest_path_().c_str(), "rb")}) { + persistent_manifest_t manifest{}; + if (std::fread(&manifest, 1, sizeof(manifest), file.get()) != sizeof(manifest)) + return error_t("Manifest is truncated"); + if (std::memcmp(manifest.magic, persistent_manifest_magic_k, 4) != 0) + return error_t("Manifest magic mismatch"); + generation = manifest.generation; + manifest_exists = true; + } + + if (!manifest_exists) { + // Fresh: empty index, snapshot.0, wal.0 (just header), manifest=0. + typename index_t::state_result_t made = index_t::make(std::move(metric), index_config); + if (!made) + return std::move(made.error); + index_ = std::move(made.index); + if (!index_.try_reserve(config_.initial_capacity)) + return error_t("Failed to reserve initial capacity"); + refresh_metadata_(); + generation_ = 0; + + auto saved = index_.save(gen_path_(".snapshot", 0).c_str()); + if (!saved) + return std::move(saved.error); + + std::string wal_path = gen_path_(".wal", 0); + { + file_t file{std::fopen(wal_path.c_str(), "wb")}; + if (!file) + return error_t("Failed to create WAL file"); + persistent_wal_header_t header = make_wal_header_(); + std::fwrite(&header, sizeof(header), 1, file.get()); + std::fflush(file.get()); + } + + if (error_t err = open_wal_for_append_(wal_path)) + return err; + return write_manifest_(0); + } + + // Recovery: load snapshot, reserve generously, replay WAL. + generation_ = generation; + std::string snapshot_path = gen_path_(".snapshot", generation_); + typename index_t::state_result_t loaded = index_t::make(snapshot_path.c_str()); + if (!loaded) + return std::move(loaded.error); + index_ = std::move(loaded.index); + refresh_metadata_(); + + std::string wal_path = gen_path_(".wal", generation_); + std::uint64_t wal_size = 0; + if (file_t probe{std::fopen(wal_path.c_str(), "rb")}) { + std::fseek(probe.get(), 0, SEEK_END); + long s = std::ftell(probe.get()); + if (s > 0) + wal_size = static_cast(s); + } + // Upper bound on records that could fit in `wal_size`, used to size + // capacity before replay (over-reserve is harmless). + std::size_t wal_records_estimate = 0; + if (wal_size > sizeof(persistent_wal_header_t)) + wal_records_estimate = (wal_size - sizeof(persistent_wal_header_t)) / (8 + 1 + sizeof(vector_key_t)); + if (!index_.try_reserve(index_.size() + wal_records_estimate + config_.initial_capacity)) + return error_t("Failed to reserve capacity for replay"); + + // Replay every well-framed WAL record onto `index_`, stopping at the + // first incomplete or crc-mismatched one. The header is validated + // against the loaded snapshot's metadata. + file_t wal{std::fopen(wal_path.c_str(), "rb")}; + if (!wal) + return error_t("Failed to open WAL for replay"); + persistent_wal_header_t header{}; + if (std::fread(&header, 1, sizeof(header), wal.get()) != sizeof(header)) + return error_t("WAL header truncated"); + if (std::memcmp(header.magic, persistent_wal_magic_k, 4) != 0) + return error_t("WAL magic mismatch"); + if (header.dimensions != dimensions_ || header.scalar_kind != scalar_kind_ || + header.metric_kind != metric_kind_) + return error_t("WAL metadata does not match the snapshot"); + + std::uint32_t const max_payload = static_cast(1 + sizeof(vector_key_t) + vector_bytes_); + std::vector buffer; + buffer.reserve(max_payload); + while (true) { + std::uint32_t len = 0, crc = 0; + if (std::fread(&len, 1, sizeof(len), wal.get()) != sizeof(len)) + break; + if (std::fread(&crc, 1, sizeof(crc), wal.get()) != sizeof(crc)) + break; + if (len < 1 + sizeof(vector_key_t) || len > max_payload) + break; + buffer.resize(len); + if (std::fread(buffer.data(), 1, len, wal.get()) != len) + break; + if (crc32_ieee(buffer.data(), len) != crc) + break; + + // Replay is single-threaded: `contains` keeps add idempotent. + persistent_op_t op = static_cast(static_cast(buffer[0])); + vector_key_t key{}; + std::memcpy(&key, &buffer[1], sizeof(vector_key_t)); + if (op == persistent_op_add_k) { + if (!index_.contains(key)) { + scalar_t const* vector = reinterpret_cast(&buffer[1 + sizeof(vector_key_t)]); + auto r = index_.add(key, vector); + if (!r) + return std::move(r.error); + } + } else if (op == persistent_op_remove_k) + index_.remove(key); // tolerant of an absent key + else + break; // unknown op + } + return open_wal_for_append_(wal_path); + } + + /// @brief Auto-trigger + step-driver, called from mutating paths. + void maybe_drive_checkpoint_() { + if (!checkpoint_active_.load(std::memory_order_acquire)) { + if (ops_since_checkpoint_.load(std::memory_order_relaxed) >= config_.checkpoint_after_ops) + start_checkpoint_(); + } + if (checkpoint_active_.load(std::memory_order_acquire)) + drive_checkpoint_step_(); + } + + /// @brief Begin barrier: drain in-flight applies under `wal_mutex_`, so + /// the snapshot covers exactly the WAL prefix up to `wal_bytes_`, + /// then `rebuild.begin`. + void start_checkpoint_() { + std::unique_lock lock(wal_mutex_); + if (checkpoint_active_.load(std::memory_order_acquire)) + return; + + // Holding `wal_mutex_` keeps new appends out; in-flight only drops. + while (in_flight_.load(std::memory_order_acquire) > 0) + std::this_thread::yield(); + + checkpoint_watermark_ = wal_bytes_; + checkpoint_gen_ = generation_ + 1; + std::string snapshot_path = gen_path_(".snapshot", checkpoint_gen_); + typename rebuild_t::result_t br = rebuild_.begin(snapshot_path.c_str()); + if (!br) { + checkpoint_error_ = std::move(br.error); + return; + } + checkpoint_active_.store(true, std::memory_order_release); + } + + /// @brief Drives one rebuild step; on the step that finishes the rebuild, + /// splices the WAL suffix onto a fresh `wal.` and atomically + /// commits the manifest. Held briefly under `wal_mutex_` to keep + /// the suffix size stable. + void drive_checkpoint_step_() { + std::unique_lock lock(checkpoint_mutex_, std::try_to_lock); + if (!lock.owns_lock()) + return; + if (!checkpoint_active_.load(std::memory_order_acquire)) + return; + auto sr = rebuild_.step(); + if (!sr) { + checkpoint_error_ = std::move(sr.error); + return; + } + if (!rebuild_.finished()) + return; + + // Finish the checkpoint: copy the WAL suffix onto a fresh generation, + // commit the manifest atomically, retire the old files. + std::unique_lock wal_lock(wal_mutex_); + if (wal_file_) + std::fflush(wal_file_.get()); + + std::string old_wal = wal_path_; + std::uint64_t suffix_start = checkpoint_watermark_; + std::uint64_t suffix_end = wal_bytes_; + std::string new_wal = gen_path_(".wal", checkpoint_gen_); + + { + file_t dst{std::fopen(new_wal.c_str(), "wb")}; + if (!dst) { + checkpoint_error_ = error_t("Failed to create new-gen WAL"); + return; + } + persistent_wal_header_t header = make_wal_header_(); + std::fwrite(&header, sizeof(header), 1, dst.get()); + + if (suffix_end > suffix_start) { + file_t src{std::fopen(old_wal.c_str(), "rb")}; + if (!src) { + std::remove(new_wal.c_str()); + checkpoint_error_ = error_t("Failed to open old WAL for suffix copy"); + return; + } + std::fseek(src.get(), static_cast(suffix_start), SEEK_SET); + char buffer[64 * 1024]; + std::uint64_t remaining = suffix_end - suffix_start; + while (remaining > 0) { + std::size_t want = + remaining < sizeof(buffer) ? static_cast(remaining) : sizeof(buffer); + std::size_t got = std::fread(buffer, 1, want, src.get()); + if (got == 0) + break; + std::fwrite(buffer, 1, got, dst.get()); + remaining -= got; + } + } + std::fflush(dst.get()); + } + + // Swap the active WAL handle BEFORE the manifest commit: a crash here + // still leaves the manifest pointing at the old generation and the + // old WAL intact for recovery. + wal_file_ = file_t{std::fopen(new_wal.c_str(), "ab")}; + if (!wal_file_) { + checkpoint_error_ = error_t("Failed to reopen new WAL for append"); + return; + } + wal_path_ = new_wal; + wal_bytes_ = sizeof(persistent_wal_header_t) + (suffix_end - suffix_start); + + // Atomic commit: from here the new generation is canonical. + if (error_t err = write_manifest_(checkpoint_gen_)) { + checkpoint_error_ = std::move(err); + return; + } + + // Retire the old generation. A crash before these removes leaves + // orphans but recovery is correct - manifest already moved on. + std::remove(old_wal.c_str()); + std::remove(gen_path_(".snapshot", generation_).c_str()); + + generation_ = checkpoint_gen_; + ops_since_checkpoint_.store(0, std::memory_order_relaxed); + checkpoint_active_.store(false, std::memory_order_release); + } +}; + +} // namespace usearch +} // namespace unum + +#endif // UNUM_USEARCH_PERSISTENT_INDEX_HPP diff --git a/include/usearch/persistent_index_helpers.hpp b/include/usearch/persistent_index_helpers.hpp new file mode 100644 index 00000000..478f9d0d --- /dev/null +++ b/include/usearch/persistent_index_helpers.hpp @@ -0,0 +1,67 @@ +/** + * @file persistent_index_helpers.hpp + * @author Mikhail Chichvarin + * @brief Helper utilities for `persistent_index_gt`: CRC32, RAII `FILE*`, + * WAL/manifest record headers and op codes. + */ +#ifndef UNUM_USEARCH_PERSISTENT_INDEX_HELPERS_HPP +#define UNUM_USEARCH_PERSISTENT_INDEX_HELPERS_HPP + +#include +#include +#include +#include + +namespace unum { +namespace usearch { + +/// @brief Bytewise CRC32 (IEEE polynomial). Records are short; no table. +inline std::uint32_t crc32_ieee(void const* data, std::size_t bytes) noexcept { + std::uint8_t const* p = static_cast(data); + std::uint32_t crc = 0xFFFFFFFFu; + for (std::size_t i = 0; i != bytes; ++i) { + crc ^= p[i]; + for (int k = 0; k != 8; ++k) + crc = (crc >> 1) ^ ((crc & 1u) ? 0xEDB88320u : 0u); + } + return ~crc; +} + +enum persistent_op_t : std::uint8_t { + persistent_op_add_k = 1, + persistent_op_remove_k = 2, +}; + +/// @brief RAII wrapper over a C `FILE*`: closes on scope exit, drops every +/// `std::fclose(file); return error_t(...)` paired cleanup. +struct file_closer_t { + void operator()(std::FILE* f) const noexcept { + if (f) + std::fclose(f); + } +}; +using file_t = std::unique_ptr; + +/// @brief Self-describing header at the start of every WAL file. +struct persistent_wal_header_t { + char magic[4]; // "uwal" + std::uint32_t format_version; + std::uint64_t dimensions; + std::uint32_t scalar_kind; + std::uint32_t metric_kind; +}; +static constexpr char const* persistent_wal_magic_k = "uwal"; +static constexpr std::uint32_t persistent_wal_version_k = 1; + +/// @brief Atomic pointer to the current durable generation. +struct persistent_manifest_t { + char magic[4]; // "umft" + std::uint32_t format_version; + std::uint64_t generation; +}; +static constexpr char const* persistent_manifest_magic_k = "umft"; + +} // namespace usearch +} // namespace unum + +#endif // UNUM_USEARCH_PERSISTENT_INDEX_HELPERS_HPP