diff --git a/c/lib.cpp b/c/lib.cpp index 05f54bbb..941f8928 100644 --- a/c/lib.cpp +++ b/c/lib.cpp @@ -186,16 +186,14 @@ USEARCH_EXPORT usearch_index_t usearch_init(usearch_init_options_t* options, use using state_result_t = typename index_dense_t::state_result_t; state_result_t state = index_dense_t::make(metric, config); - if (!state) + if (!state) { *error = state.error.release(); - index_dense_t* result_ptr = new index_dense_t(std::move(state.index)); + return NULL; + } + index_dense_t* result_ptr = new (std::nothrow) index_dense_t(std::move(state.index)); if (!result_ptr) *error = "Out of memory!"; - // Let's immediately make it usable by reserving enough threads for this machine: - if (!result_ptr->try_reserve(index_limits_t())) - *error = "Out of memory when preparing contexts!"; - return result_ptr; } @@ -236,8 +234,10 @@ USEARCH_EXPORT void usearch_metadata(char const* path, usearch_init_options_t* o USEARCH_ASSERT(path && options && error && "Missing arguments"); index_dense_metadata_result_t result = index_dense_metadata_from_path(path); - if (!result) + if (!result) { *error = result.error.release(); + return; + } options->metric_kind = metric_kind_to_c(result.head.kind_metric); options->quantization = scalar_kind_to_c(result.head.kind_scalar); @@ -285,8 +285,10 @@ USEARCH_EXPORT void usearch_metadata_buffer(void const* buffer, size_t length, u USEARCH_ASSERT(buffer && length && options && error && "Missing arguments"); index_dense_metadata_result_t result = index_dense_metadata_from_buffer(memory_mapped_file_t((byte_t*)(buffer), length)); - if (!result) + if (!result) { *error = result.error.release(); + return; + } options->metric_kind = metric_kind_to_c(result.head.kind_metric); options->quantization = scalar_kind_to_c(result.head.kind_scalar); @@ -354,7 +356,8 @@ USEARCH_EXPORT void usearch_change_threads_add(usearch_index_t index, size_t thr auto& index_dense = *reinterpret_cast(index); index_limits_t limits = index_dense.limits(); limits.threads_add = threads; - index_dense.try_reserve(limits); + if (!index_dense.try_reserve(limits)) + *error = "Out of memory!"; } USEARCH_EXPORT void usearch_change_threads_search(usearch_index_t index, size_t threads, usearch_error_t* error) { @@ -362,15 +365,22 @@ USEARCH_EXPORT void usearch_change_threads_search(usearch_index_t index, size_t auto& index_dense = *reinterpret_cast(index); index_limits_t limits = index_dense.limits(); limits.threads_search = threads; - index_dense.try_reserve(limits); + if (!index_dense.try_reserve(limits)) + *error = "Out of memory!"; } USEARCH_EXPORT void usearch_change_metric_kind(usearch_index_t index, usearch_metric_kind_t kind, usearch_error_t* error) { USEARCH_ASSERT(index && error && "Missing arguments"); auto& index_dense = *reinterpret_cast(index); - index_dense.change_metric( - metric_punned_t::builtin(index_dense.dimensions(), metric_kind_to_cpp(kind), index_dense.scalar_kind())); + auto metric_punned = + metric_punned_t::builtin(index_dense.dimensions(), metric_kind_to_cpp(kind), index_dense.scalar_kind()); + if (metric_punned.missing()) { + *error = "Unsupported metric for this index's dimensions and scalar kind!"; + return; + } + if (!index_dense.try_change_metric(std::move(metric_punned))) + *error = "Failed to grow cast buffer for the new metric!"; } USEARCH_EXPORT void usearch_change_metric(usearch_index_t index, usearch_metric_t metric, void* state, @@ -384,7 +394,12 @@ USEARCH_EXPORT void usearch_change_metric(usearch_index_t index, usearch_metric_ : metric_punned_t::stateless(index_dense.dimensions(), reinterpret_cast(metric), metric_punned_signature_t::array_array_k, metric_kind_to_cpp(kind), index_dense.scalar_kind()); - index_dense.change_metric(std::move(metric_punned)); + if (metric_punned.missing()) { + *error = "Unsupported metric for this index's dimensions and scalar kind!"; + return; + } + if (!index_dense.try_change_metric(std::move(metric_punned))) + *error = "Failed to grow cast buffer for the new metric!"; } USEARCH_EXPORT void usearch_reserve(usearch_index_t index, size_t capacity, usearch_error_t* error) { diff --git a/cpp/test.cpp b/cpp/test.cpp index 91cff0b0..3ea2f86e 100644 --- a/cpp/test.cpp +++ b/cpp/test.cpp @@ -17,6 +17,7 @@ #include // `std::signal`, `SIGSEGV`, ... #include // `std::fprintf` #include // `std::_Exit` +#include // `std::numeric_limits` #include // `std::shuffle` #include // `std::default_random_engine` @@ -216,6 +217,47 @@ void test_uint40() { expect_eq(u40_default, uint40_t(0u)); } +void test_checked_size_arithmetic() { + + std::printf("Testing checked size arithmetic\n"); + std::size_t max = (std::numeric_limits::max)(); + + checked_size_result_t cast = checked_size_from_u64(42); + expect(cast); + expect_eq(cast.value, static_cast(42)); + if (sizeof(std::size_t) < sizeof(std::uint64_t)) + expect(!checked_size_from_u64((std::numeric_limits::max)())); + + checked_size_result_t sum = checked_add(max - 1, 1); + expect(sum); + expect_eq(sum.value, max); + expect(!checked_add(max, 1)); + + checked_size_result_t product = checked_mul(max / 2, 2); + expect(product); + expect_eq(product.value, max - 1); + expect(!checked_mul(max / 2 + 1, 2)); + + checked_size_result_t fused = checked_mul_add(10, 20, 30); + expect(fused); + expect_eq(fused.value, static_cast(230)); + expect(!checked_mul_add(max / 2 + 1, 2, 0)); + expect(!checked_mul_add(max / 2, 2, 2)); + + checked_size_result_t rounded = checked_round_up(17, 8); + expect(rounded); + expect_eq(rounded.value, static_cast(24)); + expect(!checked_round_up(max, 8)); + + checked_size_result_t power = checked_ceil2(17); + expect(power); + expect_eq(power.value, static_cast(32)); + checked_size_result_t zero_power = checked_ceil2(0); + expect(zero_power); + expect_eq(zero_power.value, static_cast(0)); + expect(!checked_ceil2(max)); +} + /** * @brief Tests the functionality of the custom float16_t type ensuring consistent. */ @@ -1264,6 +1306,61 @@ static void install_crash_handlers() { std::signal(signal_number, &usearch_crash_handler); } +/** + * @brief Regression test: `make(metric, config)` must return an index that is + * immediately usable - no explicit `reserve` required before `load` / + * `view` / `search`. Previously the typed graph's `{0, 0}` thread + * limits leaked into `load_from_stream`, leaving `available_threads_` + * empty and making the first `search` throw "No available threads to + * lock". + */ +void test_load_after_metric_make() { + std::printf("Testing load and view into a metric-made index\n"); + + using index_t = index_dense_gt; + std::size_t const dimensions = 32; + std::size_t const collection = 64; + + std::default_random_engine rng(7); + std::uniform_real_distribution distribution(-1.f, 1.f); + std::vector> data(collection); + for (auto& vector : data) { + 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 config(16); + + // Build an index and persist it to disk. + index_t::state_result_t built = index_t::make(metric, config); + expect(built); + expect(built.index.try_reserve(collection)); + for (std::size_t i = 0; i != collection; ++i) + expect(built.index.add(static_cast(i), data[i].data())); + char const* path = "tmp_metric_make.usearch"; + expect(built.index.save(path)); + + // Load into a fresh, metric-made index that was never explicitly reserved. + // The first `search` must not throw "No available threads to lock". + index_t::state_result_t loaded = index_t::make(metric, config); + expect(loaded); + expect(loaded.index.load(path)); + expect_eq(loaded.index.size(), collection); + std::int64_t found[8]; + expect(loaded.index.search(data[0].data(), 5).dump_to(found) != 0); + + // Same check for the memory-mapped `view` path. + index_t::state_result_t viewed = index_t::make(metric, config); + expect(viewed); + expect(viewed.index.view(path)); + expect_eq(viewed.index.size(), collection); + expect(viewed.index.search(data[0].data(), 5).dump_to(found) != 0); + + std::remove(path); +} + int main(int, char**) { install_crash_handlers(); @@ -1272,6 +1369,7 @@ int main(int, char**) { // Non-default floating-point types may result in many compilation & rounding issues. test_uint40(); + test_checked_size_arithmetic(); test_cosine(10, 10); test_cosine(10, 10); test_cosine(10, 10); @@ -1354,5 +1452,6 @@ int main(int, char**) { test_filtered_search(); test_isolate(); + test_load_after_metric_make(); return 0; } diff --git a/include/usearch/index.hpp b/include/usearch/index.hpp index 8d033244..a25c57a9 100644 --- a/include/usearch/index.hpp +++ b/include/usearch/index.hpp @@ -176,6 +176,40 @@ namespace usearch { using byte_t = char; +struct checked_size_result_t { + std::size_t value; + bool overflow; + constexpr checked_size_result_t(std::size_t value = 0, bool overflow = false) noexcept + : value(value), overflow(overflow) {} + constexpr explicit operator bool() const noexcept { return !overflow; } +}; + +constexpr checked_size_result_t checked_size_overflow() noexcept { return {0, true}; } + +constexpr checked_size_result_t checked_size_from_u64(std::uint64_t value) noexcept { + return value > static_cast((std::numeric_limits::max)()) + ? checked_size_overflow() + : checked_size_result_t{static_cast(value), false}; +} + +constexpr checked_size_result_t checked_add(std::size_t a, std::size_t b) noexcept { + return (std::numeric_limits::max)() - a < b ? checked_size_overflow() + : checked_size_result_t{a + b, false}; +} + +constexpr checked_size_result_t checked_mul(std::size_t a, std::size_t b) noexcept { + return a && b > (std::numeric_limits::max)() / a ? checked_size_overflow() + : checked_size_result_t{a * b, false}; +} + +constexpr checked_size_result_t checked_mul_add_(checked_size_result_t product, std::size_t c) noexcept { + return product ? checked_add(product.value, c) : product; +} + +constexpr checked_size_result_t checked_mul_add(std::size_t a, std::size_t b, std::size_t c) noexcept { + return checked_mul_add_(checked_mul(a, b), c); +} + template std::size_t divide_round_up(std::size_t num) noexcept { return (num + multiple_ak - 1) / multiple_ak; } @@ -184,6 +218,21 @@ inline std::size_t divide_round_up(std::size_t num, std::size_t denominator) noe return (num + denominator - 1) / denominator; } +constexpr checked_size_result_t checked_divide_round_up(std::size_t num, std::size_t denominator) noexcept { + return !denominator ? checked_size_overflow() + : (std::numeric_limits::max)() - num < denominator - 1 + ? checked_size_overflow() + : checked_size_result_t{(num + denominator - 1) / denominator, false}; +} + +constexpr checked_size_result_t checked_round_up_(checked_size_result_t quotient, std::size_t multiple) noexcept { + return quotient ? checked_mul(quotient.value, multiple) : quotient; +} + +constexpr checked_size_result_t checked_round_up(std::size_t num, std::size_t multiple) noexcept { + return checked_round_up_(checked_divide_round_up(num, multiple), multiple); +} + inline std::size_t ceil2(std::size_t v) noexcept { v--; v |= v >> 1; @@ -198,6 +247,14 @@ inline std::size_t ceil2(std::size_t v) noexcept { return v; } +inline checked_size_result_t checked_ceil2(std::size_t v) noexcept { + if (!v) + return checked_size_result_t{0, false}; + if (v > (std::size_t{1} << ((sizeof(std::size_t) * CHAR_BIT) - 1))) + return checked_size_overflow(); + return checked_size_result_t{ceil2(v), false}; +} + /// @brief Simply dereferencing misaligned pointers can be dangerous. template void misaligned_store(void* ptr, at v) noexcept { static_assert(!std::is_reference::value, "Can't store a reference"); @@ -525,9 +582,12 @@ template > class bitset_gt { count_ = 0; } - bitset_gt(std::size_t capacity) noexcept - : slots_((compressed_slot_t*)allocator_t{}.allocate(bits_slots(capacity) * sizeof(compressed_slot_t))), - count_(slots_ ? bits_slots(capacity) : 0u) { + bitset_gt(std::size_t capacity) noexcept { + checked_size_result_t slots_count = checked_divide_round_up(capacity, bits_per_slot()); + checked_size_result_t bytes = + slots_count ? checked_mul(slots_count.value, sizeof(compressed_slot_t)) : slots_count; + slots_ = bytes ? (compressed_slot_t*)allocator_t{}.allocate(bytes.value) : nullptr; + count_ = slots_ ? slots_count.value : 0u; clear(); } @@ -654,10 +714,19 @@ class striped_locks_gt { } striped_locks_gt(std::size_t threads, std::size_t connectivity) noexcept { - std::size_t desired = threads * connectivity * 4; - if (desired < 256) - desired = 256; - count_ = ceil2(desired); + checked_size_result_t desired = checked_mul(threads, connectivity); + desired = desired ? checked_mul(desired.value, std::size_t{4}) : desired; + if (!desired) { + shift_ = 64; + return; + } + + checked_size_result_t count = checked_ceil2((std::max)(desired.value, 256)); + if (!count) { + shift_ = 64; + return; + } + count_ = count.value; shift_ = 64; for (std::size_t n = count_; n > 1; n >>= 1) shift_--; @@ -665,7 +734,13 @@ class striped_locks_gt { // `cache_line_ak`-aligned address inside the allocation, regardless of // what the underlying allocator returns. constexpr std::size_t alignment_k = alignof(padded_lock_t); - raw_bytes_ = count_ * sizeof(padded_lock_t) + alignment_k; + checked_size_result_t raw_bytes = checked_mul_add(count_, sizeof(padded_lock_t), alignment_k); + if (!raw_bytes) { + count_ = 0; + shift_ = 64; + return; + } + raw_bytes_ = raw_bytes.value; raw_ = allocator_t{}.allocate(raw_bytes_); if (!raw_) { raw_bytes_ = 0; @@ -850,10 +925,14 @@ class max_heap_gt { if (new_capacity <= capacity_) return true; - new_capacity = ceil2(new_capacity); - if (new_capacity == 0) + checked_size_result_t rounded_capacity = checked_ceil2(new_capacity); + if (!rounded_capacity) return false; - new_capacity = (std::max)(new_capacity, (std::max)(capacity_ * 2u, 16u)); + checked_size_result_t doubled_capacity = checked_mul(capacity_, std::size_t{2}); + if (!doubled_capacity) + return false; + new_capacity = + (std::max)(rounded_capacity.value, (std::max)(doubled_capacity.value, 16u)); auto allocator = allocator_t{}; auto new_elements = allocator.allocate(new_capacity); if (!new_elements) @@ -1018,10 +1097,14 @@ class sorted_buffer_gt { if (new_capacity <= capacity_) return true; - new_capacity = ceil2(new_capacity); - if (new_capacity == 0) + checked_size_result_t rounded_capacity = checked_ceil2(new_capacity); + if (!rounded_capacity) + return false; + checked_size_result_t doubled_capacity = checked_mul(capacity_, std::size_t{2}); + if (!doubled_capacity) return false; - new_capacity = (std::max)(new_capacity, (std::max)(capacity_ * 2u, 16u)); + new_capacity = + (std::max)(rounded_capacity.value, (std::max)(doubled_capacity.value, 16u)); auto allocator = allocator_t{}; auto new_elements = allocator.allocate(new_capacity); if (!new_elements) @@ -1246,9 +1329,11 @@ class growing_hash_set_gt { count_ = 0; } - growing_hash_set_gt(std::size_t capacity) noexcept - : slots_((element_t*)allocator_t{}.allocate(ceil2(capacity) * sizeof(element_t))), - capacity_(slots_ ? ceil2(capacity) : 0u), count_(0u) { + growing_hash_set_gt(std::size_t capacity) noexcept : count_(0u) { + checked_size_result_t slots_count = checked_ceil2(capacity); + checked_size_result_t bytes = slots_count ? checked_mul(slots_count.value, sizeof(element_t)) : slots_count; + slots_ = bytes ? (element_t*)allocator_t{}.allocate(bytes.value) : nullptr; + capacity_ = slots_ ? slots_count.value : 0u; clear(); } @@ -1306,12 +1391,21 @@ class growing_hash_set_gt { * @return `true` if enough capacity is available, `false` if memory allocation failed. */ bool reserve(std::size_t new_capacity) noexcept { - new_capacity = (new_capacity * 5u) / 3u; + checked_size_result_t scaled_capacity = checked_mul(new_capacity, std::size_t{5}); + if (!scaled_capacity) + return false; + new_capacity = scaled_capacity.value / 3u; if (new_capacity <= capacity_) return true; - new_capacity = ceil2(new_capacity); - element_t* new_slots = (element_t*)allocator_t{}.allocate(new_capacity * sizeof(element_t)); + checked_size_result_t rounded_capacity = checked_ceil2(new_capacity); + if (!rounded_capacity) + return false; + new_capacity = rounded_capacity.value; + checked_size_result_t new_bytes = checked_mul(new_capacity, sizeof(element_t)); + if (!new_bytes) + return false; + element_t* new_slots = (element_t*)allocator_t{}.allocate(new_bytes.value); if (!new_slots) return false; @@ -1414,7 +1508,10 @@ class ring_gt { return false; // prevent data loss if (n <= capacity()) return true; - n = (std::max)(ceil2(n), 64u); + checked_size_result_t rounded_capacity = checked_ceil2(n); + if (!rounded_capacity) + return false; + n = (std::max)(rounded_capacity.value, 64u); element_t* elements = allocator_.allocate(n); if (!elements) return false; @@ -1443,8 +1540,7 @@ class ring_gt { bool try_push(element_t const& value) noexcept { if (head_ == tail_ && !empty_) return false; // `elements_` is full - - return push(value); + push(value); return true; } @@ -1504,12 +1600,22 @@ struct index_config_t { inline error_t validate() noexcept { if (connectivity == 0) connectivity = default_connectivity(); - if (connectivity_base == 0) - connectivity_base = connectivity * 2; + if (connectivity_base == 0) { + checked_size_result_t default_base = checked_mul(connectivity, std::size_t{2}); + if (!default_base) + return "Connectivity is too large"; + connectivity_base = default_base.value; + } if (connectivity < 2) return "Connectivity must be at least 2, otherwise the index degenerates into ropes"; if (connectivity_base < connectivity) return "Base layer should be at least as connected as the rest of the graph"; + checked_size_result_t neighbors_bytes = + checked_mul_add(connectivity, sizeof(std::uint64_t), sizeof(std::uint32_t)); + checked_size_result_t neighbors_base_bytes = + checked_mul_add(connectivity_base, sizeof(std::uint64_t), sizeof(std::uint32_t)); + if (!neighbors_bytes || !neighbors_base_bytes) + return "Connectivity is too large"; return {}; } @@ -1520,6 +1626,15 @@ struct index_config_t { inline bool is_valid() const noexcept { return connectivity >= 2 && connectivity_base >= connectivity; } }; +/** + * @brief Tag type selecting the "no upfront reservation" overload of + * @ref index_limits_t. Modeled after @c std::defer_lock: the + * resulting limits are all-zero and produce no allocations when + * handed to @ref index_dense_gt::try_reserve. + */ +struct unreserved_t {}; +constexpr unreserved_t unreserved{}; + /** * @brief Growth settings for the index container. * Includes the upper bound for `::members` capacity, @@ -1527,18 +1642,32 @@ struct index_config_t { */ struct index_limits_t { /// @brief Maximum number of entries in the index. - std::size_t members = 0; + std::size_t members; /// @brief Max number of threads simultaneously updating entries. - std::size_t threads_add = std::thread::hardware_concurrency(); + std::size_t threads_add; /// @brief Max number of threads simultaneously searching entries. - std::size_t threads_search = std::thread::hardware_concurrency(); + std::size_t threads_search; inline index_limits_t(std::size_t n, std::size_t t) noexcept : members(n), threads_add(t), threads_search(t) {} - inline index_limits_t(std::size_t n = 0) noexcept : index_limits_t(n, std::thread::hardware_concurrency()) {} + inline index_limits_t(std::size_t n = 0) noexcept + : index_limits_t(n, (std::max)(1, std::thread::hardware_concurrency())) {} + inline index_limits_t(unreserved_t) noexcept : members(0), threads_add(0), threads_search(0) {} /// @brief Returns the upper limit for the number of threads. inline std::size_t threads() const noexcept { return (std::max)(threads_add, threads_search); } /// @brief Returns the concurrency-level of the index - the minimum of thread counts. inline std::size_t concurrency() const noexcept { return (std::min)(threads_add, threads_search); } + /// @brief Returns a copy with zero thread counts replaced by the library default. + /// Use when carrying limits forward across operations that may have left + /// @c threads_add / @c threads_search unset (e.g. @c unreserved construction). + inline index_limits_t with_thread_defaults() const noexcept { + index_limits_t result = *this; + index_limits_t const defaults; + if (!result.threads_add) + result.threads_add = defaults.threads_add; + if (!result.threads_search) + result.threads_search = defaults.threads_search; + return result; + } }; struct index_update_config_t { @@ -2414,6 +2543,14 @@ class index_gt { /// @brief Array of thread-specific buffers for temporary data. mutable buffer_gt contexts_{}; + context_t* context_or_null_(std::size_t thread) noexcept { + return thread < contexts_.size() ? contexts_.data() + thread : nullptr; + } + + context_t const* context_or_null_(std::size_t thread) const noexcept { + return thread < contexts_.size() ? contexts_.data() + thread : nullptr; + } + public: std::size_t connectivity() const noexcept { return config_.connectivity; } std::size_t capacity() const noexcept { return nodes_capacity_; } @@ -3027,7 +3164,10 @@ class index_gt { return result.failed("Can't add to an immutable index"); // Make sure we have enough local memory to perform this request - context_t& context = contexts_[config.thread]; + context_t* context_ptr = context_or_null_(config.thread); + if (!context_ptr) + return result.failed("Reserve capacity ahead of insertions!"); + context_t& context = *context_ptr; top_candidates_t& top = context.top_candidates; next_candidates_t& next = context.next_candidates; top.clear(); @@ -3165,7 +3305,10 @@ class index_gt { compressed_slot_t updated_slot = iterator.slot_; // Make sure we have enough local memory to perform this request - context_t& context = contexts_[config.thread]; + context_t* context_ptr = context_or_null_(config.thread); + if (!context_ptr) + return result.failed("Reserve capacity ahead of updates!"); + context_t& context = *context_ptr; top_candidates_t& top = context.top_candidates; next_candidates_t& next = context.next_candidates; top.clear(); @@ -3533,26 +3676,31 @@ class index_gt { // Progress status std::size_t processed = 0; - std::size_t const total = 2 * header.size; + checked_size_result_t header_size = checked_size_from_u64(header.size); + if (!header_size) + return result.failed("Index is too large to serialize"); + checked_size_result_t total = checked_mul(std::size_t{2}, header_size.value); + if (!total) + return result.failed("Index is too large to serialize"); // Export the number of levels per node // That is both enough to estimate the overall memory consumption, // and to be able to estimate the offsets of every entry in the file. - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { node_t node = node_at_(i); level_t level = node.level(); if (!output(&level, sizeof(level))) return result.failed("Failed to serialize into stream"); - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return result.failed("Terminated by user"); } // After that dump the nodes themselves - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { span_bytes_t node_bytes = node_bytes_(node_at_(i)); if (!output(node_bytes.data(), node_bytes.size())) return result.failed("Failed to serialize into stream"); - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return result.failed("Terminated by user"); } @@ -3584,10 +3732,16 @@ class index_gt { // Allocate some dynamic memory to read all the levels using levels_allocator_t = typename dynamic_allocator_traits_t::template rebind_alloc; - buffer_gt levels(header.size); + checked_size_result_t header_size = checked_size_from_u64(header.size); + if (!header_size) + return result.failed("Index is too large"); + buffer_gt levels(header_size.value); if (!levels) return result.failed("Out of memory"); - if (!input(levels, header.size * sizeof(level_t))) + checked_size_result_t levels_bytes = checked_mul(header_size.value, sizeof(level_t)); + if (!levels_bytes) + return result.failed("Index is too large"); + if (!input(levels, levels_bytes.value)) return result.failed("Failed to pull nodes levels from the stream"); // Submit metadata @@ -3599,26 +3753,26 @@ class index_gt { pre_ = precompute_(config_); index_limits_t limits; - limits.members = header.size; + limits.members = header_size.value; limits.threads_add = (std::max)(1, old_limits.threads_add); limits.threads_search = (std::max)(1, old_limits.threads_search); if (!reserve(limits)) { reset(); return result.failed("Out of memory"); } - nodes_count_ = header.size; + nodes_count_ = header_size.value; max_level_ = static_cast(header.max_level); entry_slot_ = static_cast(header.entry_slot); // Load the nodes - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { span_bytes_t node_bytes = node_malloc_(levels[i]); if (!input(node_bytes.data(), node_bytes.size())) { reset(); return result.failed("Failed to pull nodes from the stream"); } nodes_[i] = node_t{node_bytes.data()}; - if (!progress(i + 1, header.size)) + if (!progress(i + 1, header_size.value)) return result.failed("Terminated by user"); } return {}; @@ -3765,11 +3919,14 @@ class index_gt { reset(); return result; } + checked_size_result_t header_size = checked_size_from_u64(header.size); + if (!header_size) + return result.failed("Index is too large"); // Precompute offsets of every node, but before that we need to update the configs // This could have been done with `std::exclusive_scan`, but it's only available from C++17. using offsets_allocator_t = typename dynamic_allocator_traits_t::template rebind_alloc; - buffer_gt offsets(header.size); + buffer_gt offsets(header_size.value); if (!offsets) return result.failed("Out of memory"); @@ -3781,33 +3938,47 @@ class index_gt { pre_ = precompute_(config_); misaligned_ptr_gt levels{(byte_t*)file.data() + offset + sizeof(header)}; - offsets[0u] = offset + sizeof(header) + sizeof(level_t) * header.size; - for (std::size_t i = 1; i < header.size; ++i) - offsets[i] = offsets[i - 1] + node_bytes_(levels[i - 1]); + checked_size_result_t levels_bytes = checked_mul(sizeof(level_t), header_size.value); + checked_size_result_t offset_after_header = checked_add(offset, sizeof(header)); + checked_size_result_t first_offset = levels_bytes && offset_after_header + ? checked_add(offset_after_header.value, levels_bytes.value) + : checked_size_overflow(); + if (!first_offset) + return result.failed("Index is too large"); + offsets[0u] = first_offset.value; + for (std::size_t i = 1; i < header_size.value; ++i) { + checked_size_result_t next_offset = checked_add(offsets[i - 1], node_bytes_(levels[i - 1])); + if (!next_offset) + return result.failed("Index is too large"); + offsets[i] = next_offset.value; + } - std::size_t total_bytes = offsets[header.size - 1] + node_bytes_(levels[header.size - 1]); - if (file.size() < total_bytes) { + checked_size_result_t total_bytes = + checked_add(offsets[header_size.value - 1], node_bytes_(levels[header_size.value - 1])); + if (!total_bytes) + return result.failed("Index is too large"); + if (file.size() < total_bytes.value) { reset(); return result.failed("File is corrupted and can't fit all the nodes"); } // Submit metadata and reserve memory index_limits_t limits; - limits.members = header.size; + limits.members = header_size.value; limits.threads_add = (std::max)(1, old_limits.threads_add); limits.threads_search = (std::max)(1, old_limits.threads_search); if (!reserve(limits)) { reset(); return result.failed("Out of memory"); } - nodes_count_ = header.size; + nodes_count_ = header_size.value; max_level_ = static_cast(header.max_level); entry_slot_ = static_cast(header.entry_slot); // Rapidly address all the nodes - for (std::size_t i = 0; i != header.size; ++i) { + for (std::size_t i = 0; i != header_size.value; ++i) { nodes_[i] = node_t{(byte_t*)file.data() + offsets[i]}; - if (!progress(i + 1, header.size)) + if (!progress(i + 1, header_size.value)) return result.failed("Terminated by user"); } viewed_file_ = std::move(file); @@ -3859,7 +4030,9 @@ class index_gt { // Progress status std::atomic do_tasks{true}; std::atomic processed{0}; - std::size_t const total = 3 * slots_and_levels.size(); + checked_size_result_t total = checked_mul(std::size_t{3}, slots_and_levels.size()); + if (!total) + return; // For every bottom level node, determine its parent cluster executor.dynamic(slots_and_levels.size(), [&](std::size_t thread_idx, std::size_t old_slot_as_uint) { @@ -3872,7 +4045,7 @@ class index_gt { slots_and_levels[old_slot] = {old_slot, cluster, node_at_(old_slot).level()}; ++processed; if (thread_idx == 0) - do_tasks = progress(processed.load(), total); + do_tasks = progress(processed.load(), total.value); return do_tasks.load(); }); if (!do_tasks.load()) @@ -3906,7 +4079,7 @@ class index_gt { neighbor = static_cast(old_slot_to_new[compressed_slot_t(neighbor)]); reordered_nodes[new_slot] = new_node; - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return; } @@ -3915,7 +4088,7 @@ class index_gt { slot_transition(node_at_(old_slot).ckey(), // static_cast(old_slot), // static_cast(new_slot)); - if (!progress(++processed, total)) + if (!progress(++processed, total.value)) return; } diff --git a/include/usearch/index_dense.hpp b/include/usearch/index_dense.hpp index 92e451fb..b0d6766a 100644 --- a/include/usearch/index_dense.hpp +++ b/include/usearch/index_dense.hpp @@ -276,15 +276,22 @@ inline index_dense_metadata_result_t index_dense_metadata_from_path(char const* std::uint32_t dimensions_u32[2]{0}; std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32)); - std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32); + checked_size_result_t offset_if_u32 = + checked_mul_add(std::size_t(dimensions_u32[0]), std::size_t(dimensions_u32[1]), sizeof(dimensions_u32)); std::uint64_t dimensions_u64[2]{0}; std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64)); - std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64); + checked_size_result_t rows_if_u64 = checked_size_from_u64(dimensions_u64[0]); + checked_size_result_t columns_if_u64 = checked_size_from_u64(dimensions_u64[1]); + checked_size_result_t offset_if_u64 = + rows_if_u64 && columns_if_u64 ? checked_mul_add(rows_if_u64.value, columns_if_u64.value, sizeof(dimensions_u64)) + : checked_size_overflow(); // Check if it starts with 32-bit - if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) { - if (std::fseek(file.get(), static_cast(offset_if_u32), SEEK_SET) != 0) + checked_size_result_t head_offset_if_u32 = + offset_if_u32 ? checked_add(offset_if_u32.value, sizeof(index_dense_head_buffer_t)) : offset_if_u32; + if (head_offset_if_u32 && head_offset_if_u32.value < file_size) { + if (std::fseek(file.get(), static_cast(offset_if_u32.value), SEEK_SET) != 0) return result.failed(std::strerror(errno)); read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get()); if (!read) @@ -299,8 +306,10 @@ inline index_dense_metadata_result_t index_dense_metadata_from_path(char const* } // Check if it starts with 64-bit - if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) { - if (std::fseek(file.get(), static_cast(offset_if_u64), SEEK_SET) != 0) + checked_size_result_t head_offset_if_u64 = + offset_if_u64 ? checked_add(offset_if_u64.value, sizeof(index_dense_head_buffer_t)) : offset_if_u64; + if (head_offset_if_u64 && head_offset_if_u64.value < file_size) { + if (std::fseek(file.get(), static_cast(offset_if_u64.value), SEEK_SET) != 0) return result.failed(std::strerror(errno)); read = std::fread(result.head_buffer, sizeof(index_dense_head_buffer_t), 1, file.get()); if (!read) @@ -341,15 +350,22 @@ inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_map // Check if it starts with 32-bit std::uint32_t dimensions_u32[2]{0}; std::memcpy(dimensions_u32, result.head_buffer, sizeof(dimensions_u32)); - std::size_t offset_if_u32 = std::size_t(dimensions_u32[0]) * dimensions_u32[1] + sizeof(dimensions_u32); + checked_size_result_t offset_if_u32 = + checked_mul_add(std::size_t(dimensions_u32[0]), std::size_t(dimensions_u32[1]), sizeof(dimensions_u32)); std::uint64_t dimensions_u64[2]{0}; std::memcpy(dimensions_u64, result.head_buffer, sizeof(dimensions_u64)); - std::size_t offset_if_u64 = std::size_t(dimensions_u64[0]) * dimensions_u64[1] + sizeof(dimensions_u64); + checked_size_result_t rows_if_u64 = checked_size_from_u64(dimensions_u64[0]); + checked_size_result_t columns_if_u64 = checked_size_from_u64(dimensions_u64[1]); + checked_size_result_t offset_if_u64 = + rows_if_u64 && columns_if_u64 ? checked_mul_add(rows_if_u64.value, columns_if_u64.value, sizeof(dimensions_u64)) + : checked_size_overflow(); // Check if it starts with 32-bit - if (offset_if_u32 + sizeof(index_dense_head_buffer_t) < file_size) { - std::memcpy(&result.head_buffer, file_data + offset_if_u32, sizeof(index_dense_head_buffer_t)); + checked_size_result_t head_offset_if_u32 = + offset_if_u32 ? checked_add(offset_if_u32.value, sizeof(index_dense_head_buffer_t)) : offset_if_u32; + if (head_offset_if_u32 && head_offset_if_u32.value < file_size) { + std::memcpy(&result.head_buffer, file_data + offset_if_u32.value, sizeof(index_dense_head_buffer_t)); result.config.exclude_vectors = false; result.config.use_64_bit_dimensions = false; if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0) @@ -357,8 +373,10 @@ inline index_dense_metadata_result_t index_dense_metadata_from_buffer(memory_map } // Check if it starts with 64-bit - if (offset_if_u64 + sizeof(index_dense_head_buffer_t) < file_size) { - std::memcpy(&result.head_buffer, file_data + offset_if_u64, sizeof(index_dense_head_buffer_t)); + checked_size_result_t head_offset_if_u64 = + offset_if_u64 ? checked_add(offset_if_u64.value, sizeof(index_dense_head_buffer_t)) : offset_if_u64; + if (head_offset_if_u64 && head_offset_if_u64.value < file_size) { + std::memcpy(&result.head_buffer, file_data + offset_if_u64.value, sizeof(index_dense_head_buffer_t)); result.config.exclude_vectors = false; result.config.use_64_bit_dimensions = true; if (std::memcmp(result.head_buffer, default_magic(), std::strlen(default_magic())) == 0) @@ -532,6 +550,9 @@ class index_dense_gt { : parent(other.parent), thread_id(other.thread_id), engaged(other.engaged) { other.engaged = false; } + explicit operator bool() const noexcept { + return parent.typed_ && thread_id != any_thread() && thread_id < parent.typed_->limits().threads(); + } }; public: @@ -635,16 +656,14 @@ class index_dense_gt { * @param[in] metric One of the provided or an @b ad-hoc metric, type-punned. * @param[in] config The index configuration (optional). * @param[in] free_key The key used for freed vectors (optional). + * @param[in] limits Initial reservation. Default sizes the thread pool to + * @c hardware_concurrency(); pass @c {unreserved} to skip. * @return An instance of ::index_dense_gt or error, wrapped in a `state_result_t`. - * - * ! If the `metric` isn't provided in this method, it has to be set with - * ! the `change_metric` method before the index can be used. Alternatively, - * ! if you are loading an existing index, the metric will be set automatically. */ static state_result_t make( // metric_t metric = {}, // index_dense_config_t config = {}, // - vector_key_t free_key = default_free_value()) { + vector_key_t free_key = default_free_value(), index_limits_t limits = {}) { if (metric.missing()) return state_result_t{}.failed("Metric won't be initialized!"); @@ -659,16 +678,15 @@ class index_dense_gt { index_dense_gt& index = result.index; index.config_ = config; index.free_key_ = free_key; - - // In some cases the metric is not provided, and will be set later. - if (metric) { - scalar_kind_t scalar_kind = metric.scalar_kind(); - index.casts_ = casts_punned_t::make(scalar_kind); - index.metric_ = metric; - } + index.casts_ = casts_punned_t::make(metric.scalar_kind()); + index.metric_ = metric; new (raw) index_t(config); index.typed_ = raw; + + if (!index.try_reserve(limits)) + return state_result_t{}.failed("Failed to reserve memory for the index!"); + return result; } @@ -704,7 +722,34 @@ class index_dense_gt { // The metric and its properties metric_t const& metric() const { return metric_; } - void change_metric(metric_t metric) { metric_ = std::move(metric); } + + /** + * @brief Replaces the active distance metric, resizing the per-thread cast + * buffer and rebuilding the cast dispatch table if the new metric + * changes @c bytes_per_vector() or @c scalar_kind(). + * @return @c false if the cast buffer can't be re-allocated; the metric is + * left untouched in that case so the index stays consistent. + */ + bool try_change_metric(metric_t metric) noexcept { + checked_size_result_t needed_bytes = checked_mul(limits().threads(), metric.bytes_per_vector()); + if (!needed_bytes) + return false; + if (needed_bytes.value > cast_buffer_.size()) { + cast_buffer_t new_buffer(needed_bytes.value); + if (!new_buffer) + return false; + cast_buffer_ = std::move(new_buffer); + } + casts_ = casts_punned_t::make(metric.scalar_kind()); + metric_ = std::move(metric); + return true; + } + + /// @brief Throwing counterpart of @ref try_change_metric. + void change_metric(metric_t metric) { + if (!try_change_metric(std::move(metric))) + usearch_raise_runtime_error("failed to grow cast buffer for the new metric"); + } scalar_kind_t scalar_kind() const { return metric_.scalar_kind(); } metric_kind_t metric_kind() const { return metric_.metric_kind(); } @@ -974,6 +1019,8 @@ class index_dense_gt { index_cluster_config_t cluster_config; thread_lock_t lock = thread_lock_(thread); + if (!lock) + return cluster_result_t{}.failed("Reserve capacity ahead of searches!"); cluster_config.thread = lock.thread_id; cluster_config.expansion = config_.expansion_search; metric_proxy_t metric{*this}; @@ -1033,7 +1080,10 @@ class index_dense_gt { available_threads_.push(i); // Allocate a buffer for the casted vectors. - cast_buffer_t cast_buffer(limits.threads() * metric_.bytes_per_vector()); + checked_size_result_t cast_buffer_bytes = checked_mul(limits.threads(), metric_.bytes_per_vector()); + if (!cast_buffer_bytes) + return false; + cast_buffer_t cast_buffer(cast_buffer_bytes.value); if (!cast_buffer) return false; cast_buffer_ = std::move(cast_buffer); @@ -1058,7 +1108,8 @@ class index_dense_gt { std::unique_lock free_lock(free_keys_mutex_); typed_->clear(); slot_lookup_.clear(); - vectors_lookup_.reset(); + // Tape pointers are about to be invalidated by the reset below. + std::fill(vectors_lookup_.begin(), vectors_lookup_.end(), nullptr); free_keys_.clear(); vectors_tape_allocator_.reset(); } @@ -1183,8 +1234,9 @@ class index_dense_gt { serialization_config_t config = {}, // progress_at&& progress = {}) { - // Discard all previous memory allocations of `vectors_tape_allocator_` - index_limits_t old_limits = typed_ ? typed_->limits() : index_limits_t{}; + // Preserve any explicit thread counts from the prior index state; fall + // back to library defaults when they were never set (e.g. `unreserved`). + index_limits_t new_limits = typed_ ? typed_->limits().with_thread_defaults() : index_limits_t{}; reset(); // Infer the new index size @@ -1245,9 +1297,10 @@ class index_dense_gt { config_.multi = head.multi; metric_ = metric_t::builtin(head.dimensions, head.kind_metric, head.kind_scalar); - // available_threads_.size() will be updated to old_limits.threads() later in this - // method, so use that as the number of threads to prepare for. - cast_buffer_ = cast_buffer_t(old_limits.threads() * metric_.bytes_per_vector()); + checked_size_result_t cast_buffer_bytes = checked_mul(new_limits.threads(), metric_.bytes_per_vector()); + if (!cast_buffer_bytes) + return result.failed("Failed to allocate memory for the casts"); + cast_buffer_ = cast_buffer_t(cast_buffer_bytes.value); if (!cast_buffer_) return result.failed("Failed to allocate memory for the casts"); casts_ = casts_punned_t::make(head.kind_scalar); @@ -1266,14 +1319,13 @@ class index_dense_gt { return result; if (typed_->size() != static_cast(matrix_rows)) return result.failed("Index size and the number of vectors doesn't match"); - old_limits.members = static_cast(matrix_rows); - if (!typed_->try_reserve(old_limits)) + new_limits.members = static_cast(matrix_rows); + if (!typed_->try_reserve(new_limits)) return result.failed("Failed to reserve memory for the index"); - // After the index is loaded, we may have to resize the `available_threads_` to - // match the limits of the underlying engine. + // After the index is loaded, resize `available_threads_` to match the new limits. available_threads_t available_threads; - std::size_t max_threads = old_limits.threads(); + std::size_t max_threads = new_limits.threads(); if (!available_threads.reserve(max_threads)) return result.failed("Failed to allocate memory for the available threads!"); for (std::size_t i = 0; i < max_threads; i++) @@ -1297,8 +1349,9 @@ class index_dense_gt { std::size_t offset = 0, serialization_config_t config = {}, // progress_at&& progress = {}) { - // Discard all previous memory allocations of `vectors_tape_allocator_` - index_limits_t old_limits = typed_ ? typed_->limits() : index_limits_t{}; + // Preserve any explicit thread counts from the prior index state; fall + // back to library defaults when they were never set (e.g. `unreserved`). + index_limits_t new_limits = typed_ ? typed_->limits().with_thread_defaults() : index_limits_t{}; reset(); serialization_result_t result = file.open_if_not(); @@ -1362,8 +1415,10 @@ class index_dense_gt { config_.multi = head.multi; metric_ = metric_t::builtin(head.dimensions, head.kind_metric, head.kind_scalar); // available_threads_.size() will be updated to old_limits.threads() later in this - // method, so use that as the number of threads to prepare for. - cast_buffer_ = cast_buffer_t(old_limits.threads() * metric_.bytes_per_vector()); + checked_size_result_t cast_buffer_bytes = checked_mul(new_limits.threads(), metric_.bytes_per_vector()); + if (!cast_buffer_bytes) + return result.failed("Failed to allocate memory for the casts"); + cast_buffer_ = cast_buffer_t(cast_buffer_bytes.value); if (!cast_buffer_) return result.failed("Failed to allocate memory for the casts"); casts_ = casts_punned_t::make(head.kind_scalar); @@ -1383,8 +1438,8 @@ class index_dense_gt { return result; if (typed_->size() != static_cast(matrix_rows)) return result.failed("Index size and the number of vectors doesn't match"); - old_limits.members = static_cast(matrix_rows); - if (!typed_->try_reserve(old_limits)) + new_limits.members = static_cast(matrix_rows); + if (!typed_->try_reserve(new_limits)) return result.failed("Failed to reserve memory for the index"); // Address the vectors @@ -1395,10 +1450,9 @@ class index_dense_gt { for (std::uint64_t slot = 0; slot != matrix_rows; ++slot) vectors_lookup_[slot] = (byte_t*)vectors_buffer.data() + matrix_cols * slot; - // After the index is loaded, we may have to resize the `available_threads_` to - // match the limits of the underlying engine. + // After the index is viewed, resize `available_threads_` to match the new limits. available_threads_t available_threads; - std::size_t max_threads = old_limits.threads(); + std::size_t max_threads = new_limits.threads(); if (!available_threads.reserve(max_threads)) return result.failed("Failed to allocate memory for the available threads!"); for (std::size_t i = 0; i < max_threads; i++) @@ -1663,6 +1717,11 @@ class index_dense_gt { */ labeling_result_t rename(vector_key_t from, vector_key_t to) { labeling_result_t result; + if (from == to) { + result.completed = count(from); + return result; + } + unique_lock_t lookup_lock(slot_lookup_mutex_); if (!multi() && slot_lookup_.contains(key_and_slot_t::any_slot(to))) @@ -2082,10 +2141,9 @@ class index_dense_gt { if (thread_id != any_thread()) return {*this, thread_id, false}; - available_threads_mutex_.lock(); - usearch_assert_m(available_threads_.size(), "No available threads to lock"); - available_threads_.try_pop(thread_id); - available_threads_mutex_.unlock(); + std::unique_lock lock(available_threads_mutex_); + if (!available_threads_.try_pop(thread_id)) + return {*this, any_thread(), false}; return {*this, thread_id, true}; } @@ -2106,6 +2164,8 @@ class index_dense_gt { // Cast the vector, if needed for compatibility with `metric_` thread_lock_t lock = thread_lock_(thread); + if (!lock) + return add_result_t{}.failed("Reserve capacity ahead of insertions!"); byte_t const* vector_data = reinterpret_cast(vector); { byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id; @@ -2159,6 +2219,8 @@ class index_dense_gt { // Cast the vector, if needed for compatibility with `metric_` thread_lock_t lock = thread_lock_(thread); + if (!lock) + return search_result_t{*this}.failed("Reserve capacity ahead of searches!"); byte_t const* vector_data = reinterpret_cast(vector); { byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id; @@ -2195,6 +2257,8 @@ class index_dense_gt { // Cast the vector, if needed for compatibility with `metric_` thread_lock_t lock = thread_lock_(thread); + if (!lock) + return cluster_result_t{}.failed("Reserve capacity ahead of searches!"); byte_t const* vector_data = reinterpret_cast(vector); { byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id; @@ -2219,6 +2283,8 @@ class index_dense_gt { // Cast the vector, if needed for compatibility with `metric_` thread_lock_t lock = thread_lock_(thread); + if (!lock) + return {}; byte_t const* vector_data = reinterpret_cast(vector); { byte_t* casted_data = cast_buffer_.data() + metric_.bytes_per_vector() * lock.thread_id; diff --git a/include/usearch/index_plugins.hpp b/include/usearch/index_plugins.hpp index c874fba5..e27eb859 100644 --- a/include/usearch/index_plugins.hpp +++ b/include/usearch/index_plugins.hpp @@ -1572,21 +1572,24 @@ class aligned_allocator_gt { constexpr std::size_t alignment() const { return alignment_ak; } pointer allocate(size_type length) const { - std::size_t length_bytes = alignment_ak * divide_round_up(length * sizeof(value_type)); - // Avoid overflow - if (length > length_bytes) + checked_size_result_t bytes = checked_mul(length, sizeof(value_type)); + if (!bytes) return nullptr; + checked_size_result_t length_bytes = checked_round_up(bytes.value, alignment_ak); + if (!length_bytes) + return nullptr; + std::size_t alignment = alignment_ak; #if defined(USEARCH_DEFINED_WINDOWS) - return (pointer)_aligned_malloc(length_bytes, alignment); + return (pointer)_aligned_malloc(length_bytes.value, alignment); #elif defined(USEARCH_DEFINED_APPLE) || defined(USEARCH_DEFINED_ANDROID) // Apple Clang keeps complaining that `aligned_alloc` is only available // with macOS 10.15 and newer or Android API >= 28, so let's use `posix_memalign` there. void* result = nullptr; - int status = posix_memalign(&result, alignment, length_bytes); + int status = posix_memalign(&result, alignment, length_bytes.value); return status == 0 ? (pointer)result : nullptr; #else - return (pointer)aligned_alloc(alignment, length_bytes); + return (pointer)aligned_alloc(alignment, length_bytes.value); #endif } @@ -1615,7 +1618,10 @@ class page_allocator_t { * @return A pointer to the allocated memory block, or `nullptr` if allocation fails. */ byte_t* allocate(std::size_t count_bytes) const noexcept { - count_bytes = divide_round_up(count_bytes, page_size()) * page_size(); + checked_size_result_t rounded_bytes = checked_round_up(count_bytes, page_size()); + if (!rounded_bytes) + return nullptr; + count_bytes = rounded_bytes.value; #if defined(USEARCH_DEFINED_WINDOWS) return (byte_t*)(::VirtualAlloc(NULL, count_bytes, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE)); #else @@ -1628,7 +1634,10 @@ class page_allocator_t { #if defined(USEARCH_DEFINED_WINDOWS) ::VirtualFree(page_pointer, 0, MEM_RELEASE); #else - count_bytes = divide_round_up(count_bytes, page_size()) * page_size(); + checked_size_result_t rounded_bytes = checked_round_up(count_bytes, page_size()); + if (!rounded_bytes) + return; + count_bytes = rounded_bytes.value; munmap(page_pointer, count_bytes); #endif } @@ -1724,25 +1733,39 @@ template class memory_mapping_allocator_gt { * @return A pointer to the allocated memory block, or `nullptr` if allocation fails. */ inline byte_t* allocate(std::size_t count_bytes) noexcept { - std::size_t extended_bytes = divide_round_up(count_bytes) * alignment_ak; + checked_size_result_t extended_bytes = checked_round_up(count_bytes, alignment_ak); + if (!extended_bytes) + return nullptr; std::unique_lock lock(mutex_); - if (!last_arena_ || (last_usage_ + extended_bytes >= last_capacity_)) { - std::size_t new_cap = (std::max)(last_capacity_, ceil2(extended_bytes)) * capacity_multiplier(); - byte_t* new_arena = page_allocator_t{}.allocate(new_cap); + checked_size_result_t next_usage = checked_add(last_usage_, extended_bytes.value); + if (!next_usage) + return nullptr; + if (!last_arena_ || (next_usage.value >= last_capacity_)) { + checked_size_result_t rounded_bytes = checked_ceil2(extended_bytes.value); + if (!rounded_bytes) + return nullptr; + checked_size_result_t new_cap = + checked_mul((std::max)(last_capacity_, rounded_bytes.value), capacity_multiplier()); + if (!new_cap) + return nullptr; + checked_size_result_t new_total_allocated = checked_add(total_allocated_, new_cap.value); + if (!new_total_allocated) + return nullptr; + byte_t* new_arena = page_allocator_t{}.allocate(new_cap.value); if (!new_arena) return nullptr; std::memcpy(new_arena, &last_arena_, sizeof(byte_t*)); - std::memcpy(new_arena + sizeof(byte_t*), &new_cap, sizeof(std::size_t)); + std::memcpy(new_arena + sizeof(byte_t*), &new_cap.value, sizeof(std::size_t)); wasted_space_ += total_reserved(); last_arena_ = new_arena; - last_capacity_ = new_cap; + last_capacity_ = new_cap.value; last_usage_ = head_size(); - total_allocated_ += new_cap; + total_allocated_ = new_total_allocated.value; } - wasted_space_ += extended_bytes - count_bytes; - return last_arena_ + exchange(last_usage_, last_usage_ + extended_bytes); + wasted_space_ += extended_bytes.value - count_bytes; + return last_arena_ + exchange(last_usage_, last_usage_ + extended_bytes.value); } /** @@ -1948,6 +1971,11 @@ template struct cast_to_i8_gt { for (std::size_t i = 0; i != dim; ++i) magnitude += (double)typed_input[i] * (double)typed_input[i]; magnitude = std::sqrt(magnitude); + // `!(x > 0)` also catches NaN; cast-to-int of NaN is UB. + if (!(magnitude > 0.0)) { + std::fill_n(typed_output, dim, std::int8_t{0}); + return true; + } for (std::size_t i = 0; i != dim; ++i) typed_output[i] = static_cast(usearch::clamp(typed_input[i] * 127.0 / magnitude, -127.0, 127.0)); @@ -1973,6 +2001,11 @@ template struct cast_to_u8_gt { for (std::size_t i = 0; i != dim; ++i) magnitude += (double)typed_input[i] * (double)typed_input[i]; magnitude = std::sqrt(magnitude); + // `!(x > 0)` also catches NaN; cast-to-int of NaN is UB. + if (!(magnitude > 0.0)) { + std::fill_n(typed_output, dim, std::uint8_t{0}); + return true; + } for (std::size_t i = 0; i != dim; ++i) typed_output[i] = static_cast(usearch::clamp(typed_input[i] * 255.0 / magnitude, 0.0, 255.0)); @@ -3750,7 +3783,8 @@ class flat_hash_multi_set_gt { public: std::size_t size() const noexcept { return populated_slots_; } - std::size_t capacity() const noexcept { return capacity_slots_; } + std::size_t capacity() const noexcept { return capacity_slots_ * 2u / 3u; } + std::size_t capacity_slots() const noexcept { return capacity_slots_; } flat_hash_multi_set_gt() noexcept {} ~flat_hash_multi_set_gt() noexcept { reset(); } @@ -3764,7 +3798,10 @@ class flat_hash_multi_set_gt { } // Allocate new memory - data_ = (char*)allocator_t{}.allocate(other.buckets_ * bytes_per_bucket()); + checked_size_result_t bytes = checked_mul(other.buckets_, bytes_per_bucket()); + if (!bytes) + usearch_raise_runtime_error("failed memory allocation"); + data_ = (char*)allocator_t{}.allocate(bytes.value); if (!data_) usearch_raise_runtime_error("failed memory allocation"); @@ -3804,7 +3841,10 @@ class flat_hash_multi_set_gt { allocator_t{}.deallocate(data_, buckets_ * bytes_per_bucket()); // Allocate new memory - data_ = (char*)allocator_t{}.allocate(other.buckets_ * bytes_per_bucket()); + checked_size_result_t bytes = checked_mul(other.buckets_, bytes_per_bucket()); + if (!bytes) + usearch_raise_runtime_error("failed memory allocation"); + data_ = (char*)allocator_t{}.allocate(bytes.value); if (!data_) usearch_raise_runtime_error("failed memory allocation"); @@ -3853,22 +3893,37 @@ class flat_hash_multi_set_gt { } bool try_reserve(std::size_t capacity) noexcept { - if (capacity * 3u <= capacity_slots_ * 2u) + if (capacity <= this->capacity()) return true; // Calculate new sizes - std::size_t new_slots = ceil2((capacity * 3ul) / 2ul); - std::size_t new_buckets = divide_round_up(new_slots); - new_slots = new_buckets * slots_per_bucket(); // This must be a power of two! - std::size_t new_bytes = new_buckets * bytes_per_bucket(); + checked_size_result_t scaled_capacity = checked_mul(capacity, std::size_t{3}); + if (!scaled_capacity) + return false; + checked_size_result_t slots_needed = checked_divide_round_up(scaled_capacity.value, std::size_t{2}); + if (!slots_needed) + return false; + checked_size_result_t new_slots_checked = checked_ceil2(slots_needed.value); + if (!new_slots_checked) + return false; + checked_size_result_t new_buckets_checked = + checked_divide_round_up(new_slots_checked.value, slots_per_bucket()); + if (!new_buckets_checked) + return false; + checked_size_result_t new_slots = checked_mul(new_buckets_checked.value, slots_per_bucket()); + if (!new_slots) + return false; + checked_size_result_t new_bytes = checked_mul(new_buckets_checked.value, bytes_per_bucket()); + if (!new_bytes) + return false; // Allocate new memory - char* new_data = (char*)allocator_t{}.allocate(new_bytes); + char* new_data = (char*)allocator_t{}.allocate(new_bytes.value); if (!new_data) return false; // Initialize new buckets to empty - std::memset(new_data, 0, new_bytes); + std::memset(new_data, 0, new_bytes.value); // Rehash and copy existing elements to new_data hash_t hasher; @@ -3879,7 +3934,7 @@ class flat_hash_multi_set_gt { // Rehash std::size_t hash_value = hasher(old_slot.element); - std::size_t new_slot_index = hash_value & (new_slots - 1); + std::size_t new_slot_index = hash_value & (new_slots.value - 1); // Linear probing to find an empty slot in new_data while (true) { @@ -3889,7 +3944,7 @@ class flat_hash_multi_set_gt { new_slot.header.populated |= new_slot.mask; break; } - new_slot_index = (new_slot_index + 1) & (new_slots - 1); + new_slot_index = (new_slot_index + 1) & (new_slots.value - 1); } } @@ -3897,8 +3952,8 @@ class flat_hash_multi_set_gt { if (data_) allocator_t{}.deallocate(data_, buckets_ * bytes_per_bucket()); data_ = new_data; - buckets_ = new_buckets; - capacity_slots_ = new_slots; + buckets_ = new_buckets_checked.value; + capacity_slots_ = new_slots.value; return true; } @@ -3916,16 +3971,20 @@ class flat_hash_multi_set_gt { : index_(index), parent_(parent), query_(query), equals_(equals) {} // Pre-increment: advance past tombstones and non-matching entries, - // stopping at the next matching live entry or an empty slot. + // stopping at the next matching live entry or an empty slot. When every + // slot is either live or tombstoned (no empty slot exists), the probe + // saturates after `capacity_slots_` steps and the iterator becomes + // `end()` - otherwise the loop would spin forever. equal_iterator_gt& operator++() { - do { + for (std::size_t remaining = parent_->capacity_slots_; remaining; --remaining) { index_ = (index_ + 1) & (parent_->capacity_slots_ - 1); auto slot = parent_->slot_ref(index_); bool is_empty = ~slot.header.populated & slot.mask; bool is_match = !(slot.header.deleted & slot.mask) && equals_(slot.element, query_); if (is_empty || is_match) - break; - } while (true); + return *this; + } + index_ = parent_->capacity_slots_; // saturated probe -> end() return *this; } diff --git a/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp b/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp index 510e896f..9bbe388b 100644 --- a/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +++ b/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp @@ -146,19 +146,13 @@ JNIEXPORT jlong JNICALL Java_cloud_unum_usearch_Index_c_1capacity(JNIEnv*, jclas JNIEXPORT void JNICALL Java_cloud_unum_usearch_Index_c_1reserve(JNIEnv* env, jclass, jlong c_ptr, jlong capacity, jlong threads_add, jlong threads_search) { - std::size_t t_add = static_cast(threads_add); - std::size_t t_search = static_cast(threads_search); - if (t_add == 0 || t_search == 0) { - std::size_t hc = std::thread::hardware_concurrency(); - if (hc == 0) - hc = 1; // fallback to 1 if the runtime can't report - if (t_add == 0) - t_add = hc; - if (t_search == 0) - t_search = hc; - } - index_limits_t limits(static_cast(capacity), t_add); - limits.threads_search = t_search; + // A zero from the caller means "use the library default", which is centrally + // floored at `max(1, hardware_concurrency())` inside `index_limits_t`. + index_limits_t limits = threads_add ? index_limits_t(static_cast(capacity), + static_cast(threads_add)) + : index_limits_t(static_cast(capacity)); + if (threads_search) + limits.threads_search = static_cast(threads_search); if (!reinterpret_cast(c_ptr)->try_reserve(limits)) { jclass jc = (*env).FindClass("java/lang/Error"); if (jc) diff --git a/javascript/lib.cpp b/javascript/lib.cpp index 8d87a143..173b654c 100644 --- a/javascript/lib.cpp +++ b/javascript/lib.cpp @@ -320,8 +320,10 @@ Napi::Value CompiledIndex::Remove(Napi::CallbackInfo const& ctx) { Napi::Array results = Napi::Array::New(env, length); for (std::size_t i = 0; i < length; ++i) { auto result = native_->remove(static_cast(keys[i])); - if (!result) - Napi::Error::New(ctx.Env(), result.error.release()).ThrowAsJavaScriptException(); + if (!result) { + Napi::Error::New(env, result.error.release()).ThrowAsJavaScriptException(); + return env.Null(); + } results[i] = Napi::Number::New(env, result.completed); } return results; diff --git a/python/lib.cpp b/python/lib.cpp index 645e6a09..601297d4 100644 --- a/python/lib.cpp +++ b/python/lib.cpp @@ -52,7 +52,14 @@ using progress_func_t = std::function mutex_ptr_ = std::make_unique(); + dense_index_py_t(native_t&& base) : index_dense_t(std::move(base)) {} }; struct dense_indexes_py_t { std::vector> shards_; + mutable std::unique_ptr mutex_ptr_ = std::make_unique(); void merge(std::shared_ptr shard) { shards_.push_back(shard); } std::size_t bytes_per_vector() const noexcept { return shards_.empty() ? 0 : shards_[0]->bytes_per_vector(); } @@ -85,6 +105,12 @@ struct dense_indexes_py_t { shards_.reserve(shards_.size() + paths.size()); std::mutex shards_mutex; + // Release the GIL *before* taking the per-index mutex so a Python + // thread waiting on the mutex doesn't hold the GIL - otherwise a + // worker thread in the current owner would block forever in + // `gil_scoped_acquire`. + py::gil_scoped_release release; + std::unique_lock lock(*mutex_ptr_); executor_default_t{threads}.dynamic(paths.size(), [&](std::size_t, std::size_t task_idx) { index_dense_t index = index_dense_t::make(paths[task_idx].c_str(), view); if (!index) @@ -92,6 +118,7 @@ struct dense_indexes_py_t { auto shared_index = std::make_shared(std::move(index)); std::unique_lock lock(shards_mutex); shards_.push_back(shared_index); + py::gil_scoped_acquire acquire; if (PyErr_CheckSignals() != 0) throw py::error_already_set(); return true; @@ -181,24 +208,33 @@ static void add_typed_to_index( // progress_t progress_{progress}; std::atomic processed{0}; - executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) { - dense_key_t key = *reinterpret_cast(keys_data + task_idx * keys_info.strides[0]); - scalar_at const* vector = reinterpret_cast(vectors_data + task_idx * vectors_info.strides[0]); - dense_add_result_t result = index.add(key, vector, thread_idx, force_copy); - if (!result) { - atomic_error = result.error.release(); - return false; - } - - // We don't want to check for signals from multiple threads - ++processed; - if (thread_idx == 0) - if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), vectors_count)) { - atomic_error.store("Operation has been terminated"); + { + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); + if (!index.try_reserve(index_limits_t(ceil2(index.size() + vectors_count), threads))) + throw std::invalid_argument("Out of memory!"); + executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) { + dense_key_t key = *reinterpret_cast(keys_data + task_idx * keys_info.strides[0]); + scalar_at const* vector = + reinterpret_cast(vectors_data + task_idx * vectors_info.strides[0]); + dense_add_result_t result = index.add(key, vector, thread_idx, force_copy); + if (!result) { + atomic_error = result.error.release(); return false; } - return true; - }); + + // We don't want to check for signals from multiple threads + ++processed; + if (thread_idx == 0) { + py::gil_scoped_acquire acquire; + if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), vectors_count)) { + atomic_error.store("Operation has been terminated"); + return false; + } + } + return true; + }); + } // At the end report the latest numbers, because the reporter thread may be finished earlier progress_(processed.load(), vectors_count); @@ -244,8 +280,10 @@ static void add_many_to_index( // if (!threads) threads = std::thread::hardware_concurrency(); - if (!index.try_reserve(index_limits_t(ceil2(index.size() + vectors_count), threads))) - throw std::invalid_argument("Out of memory!"); + + // `add_typed_to_index` does the `try_reserve` + executor work inside its + // own GIL-released, mutex-locked region; we just dispatch on the scalar + // kind here. // clang-format off scalar_kind_t kind = (scalar_kind != scalar_kind_t::unknown_k) @@ -285,37 +323,43 @@ static void search_typed( // if (!threads) threads = std::thread::hardware_concurrency(); - if (!index.try_reserve(index_limits_t(index.size(), threads))) - throw std::invalid_argument("Out of memory!"); // Progress status progress_t progress_{progress}; std::atomic processed{0}; atomic_error_t atomic_error{nullptr}; - executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) { - scalar_at const* vector = (scalar_at const*)(vectors_data + task_idx * vectors_info.strides[0]); - dense_search_result_t result = index.search(vector, wanted, thread_idx, exact); - if (!result) { - atomic_error = result.error.release(); - return false; - } + { + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); + if (!index.try_reserve(index_limits_t(index.size(), threads))) + throw std::invalid_argument("Out of memory!"); + executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) { + scalar_at const* vector = (scalar_at const*)(vectors_data + task_idx * vectors_info.strides[0]); + dense_search_result_t result = index.search(vector, wanted, thread_idx, exact); + if (!result) { + atomic_error = result.error.release(); + return false; + } - counts_py1d(task_idx) = - static_cast(result.dump_to(&keys_py2d(task_idx, 0), &distances_py2d(task_idx, 0), wanted)); + counts_py1d(task_idx) = + static_cast(result.dump_to(&keys_py2d(task_idx, 0), &distances_py2d(task_idx, 0), wanted)); - stats_visited_members += result.visited_members; - stats_computed_distances += result.computed_distances; + stats_visited_members += result.visited_members; + stats_computed_distances += result.computed_distances; - // We don't want to check for signals from multiple threads - ++processed; - if (thread_idx == 0) - if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), vectors_count)) { - atomic_error.store("Operation has been terminated"); - return false; + // We don't want to check for signals from multiple threads + ++processed; + if (thread_idx == 0) { + py::gil_scoped_acquire acquire; + if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), vectors_count)) { + atomic_error.store("Operation has been terminated"); + return false; + } } - return true; - }); + return true; + }); + } // At the end report the latest numbers, because the reporter thread may be finished earlier progress_(processed.load(), vectors_count); @@ -357,48 +401,54 @@ static void search_typed( // std::atomic processed{0}; atomic_error_t atomic_error{nullptr}; - executor_default_t{threads}.dynamic(indexes.shards_.size(), [&](std::size_t thread_idx, std::size_t task_idx) { - dense_index_py_t& index = *indexes.shards_[task_idx].get(); - - index_limits_t limits; - limits.members = index.size(); - limits.threads_add = 0; - limits.threads_search = 1; - if (!index.try_reserve(limits)) { - atomic_error = "Out of memory!"; - return false; - } - - for (std::size_t vector_idx = 0; vector_idx != static_cast(vectors_count); ++vector_idx) { - scalar_at const* vector = (scalar_at const*)(vectors_data + vector_idx * vectors_info.strides[0]); - dense_search_result_t result = index.search(vector, wanted, 0, exact); - if (!result) { - atomic_error = result.error.release(); + { + py::gil_scoped_release release; + std::unique_lock lock(*indexes.mutex_ptr_); + executor_default_t{threads}.dynamic(indexes.shards_.size(), [&](std::size_t thread_idx, std::size_t task_idx) { + dense_index_py_t& index = *indexes.shards_[task_idx].get(); + + index_limits_t limits; + limits.members = index.size(); + limits.threads_add = 0; + limits.threads_search = 1; + if (!index.try_reserve(limits)) { + atomic_error = "Out of memory!"; return false; } - { - auto lock = query_mutexes.lock(vector_idx); - counts_py1d(vector_idx) = static_cast(result.merge_into( // - &keys_py2d(vector_idx, 0), // - &distances_py2d(vector_idx, 0), // - static_cast(counts_py1d(vector_idx)), // - wanted)); - } + for (std::size_t vector_idx = 0; vector_idx != static_cast(vectors_count); ++vector_idx) { + scalar_at const* vector = (scalar_at const*)(vectors_data + vector_idx * vectors_info.strides[0]); + dense_search_result_t result = index.search(vector, wanted, 0, exact); + if (!result) { + atomic_error = result.error.release(); + return false; + } - stats_visited_members += result.visited_members; - stats_computed_distances += result.computed_distances; + { + auto lock = query_mutexes.lock(vector_idx); + counts_py1d(vector_idx) = static_cast(result.merge_into( // + &keys_py2d(vector_idx, 0), // + &distances_py2d(vector_idx, 0), // + static_cast(counts_py1d(vector_idx)), // + wanted)); + } - // We don't want to check for signals from multiple threads - ++processed; - if (thread_idx == 0) - if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), indexes.shards_.size())) { - atomic_error.store("Operation has been terminated"); - return false; + stats_visited_members += result.visited_members; + stats_computed_distances += result.computed_distances; + + // We don't want to check for signals from multiple threads + ++processed; + if (thread_idx == 0) { + py::gil_scoped_acquire acquire; + if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), indexes.shards_.size())) { + atomic_error.store("Operation has been terminated"); + return false; + } } - } - return true; - }); + } + return true; + }); + } // At the end report the latest numbers, because the reporter thread may be finished earlier progress_(processed.load(), indexes.shards_.size()); @@ -552,11 +602,17 @@ static py::tuple search_many_brute_force( // executor_default_t executor{threads}; exact_search_t search; - exact_search_results_t offsets_and_distances = search( // - dataset_data, dataset_count, dataset_stride, // - queries_data, queries_count, queries_stride, // - wanted, metric, executor, - [&](std::size_t passed, std::size_t total) { return PyErr_CheckSignals() == 0 && progress(passed, total); }); + exact_search_results_t offsets_and_distances; + { + py::gil_scoped_release release; + offsets_and_distances = search( // + dataset_data, dataset_count, dataset_stride, // + queries_data, queries_count, queries_stride, // + wanted, metric, executor, [&](std::size_t passed, std::size_t total) { + py::gil_scoped_acquire acquire; + return PyErr_CheckSignals() == 0 && progress(passed, total); + }); + } if (!offsets_and_distances) throw std::bad_alloc(); @@ -628,11 +684,18 @@ static py::tuple cluster_many_brute_force( // engine.max_seconds = max_seconds; engine.inertia_threshold = inertia_threshold; - kmeans_clustering_result_t result = engine( // - reinterpret_cast(dataset_info.ptr), dataset_count, dataset_stride, // - centroids.data(), wanted, dataset_dimensions * bytes_per_scalar, // - point_to_centroid_index.data(), point_to_centroid_distance.data(), dataset_kind, dataset_dimensions, executor, - [&](std::size_t passed, std::size_t total) { return PyErr_CheckSignals() == 0 && progress(passed, total); }); + kmeans_clustering_result_t result; + { + py::gil_scoped_release release; + result = engine( // + reinterpret_cast(dataset_info.ptr), dataset_count, dataset_stride, // + centroids.data(), wanted, dataset_dimensions * bytes_per_scalar, // + point_to_centroid_index.data(), point_to_centroid_distance.data(), dataset_kind, dataset_dimensions, + executor, [&](std::size_t passed, std::size_t total) { + py::gil_scoped_acquire acquire; + return PyErr_CheckSignals() == 0 && progress(passed, total); + }); + } if (!result) throw std::runtime_error(result.error.release()); @@ -720,19 +783,23 @@ static py::tuple cluster_vectors( // scalar_kind_t kind = (scalar_kind != scalar_kind_t::unknown_k) ? scalar_kind : numpy_string_to_kind(queries_info.format); - switch (kind) { - case scalar_kind_t::f64_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::f32_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::bf16_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::f16_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::e5m2_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::e4m3_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::e3m2_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::e2m3_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::i8_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::u8_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - case scalar_kind_t::b1x8_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; - default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + queries_info.format); + { + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); + switch (kind) { + case scalar_kind_t::f64_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::f32_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::bf16_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::f16_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::e5m2_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::e4m3_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::e3m2_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::e2m3_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::i8_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::u8_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + case scalar_kind_t::b1x8_k: cluster_result = index.cluster(queries_begin.as(), queries_end.as(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break; + default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + queries_info.format); + } } // clang-format on @@ -790,8 +857,13 @@ static py::tuple cluster_keys( // config.min_clusters = min_count; config.max_clusters = max_count; - dense_clustering_result_t cluster_result = - index.cluster(queries_begin, queries_end, config, keys_ptr, distances_ptr, executor, progress_t{progress}); + dense_clustering_result_t cluster_result; + { + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); + cluster_result = + index.cluster(queries_begin, queries_end, config, keys_ptr, distances_ptr, executor, progress_t{progress}); + } cluster_result.error.raise(); // Those would be set to 1 for all entries, in case of success @@ -824,7 +896,15 @@ static std::unordered_map join_index( // config.expansion = (std::max)(a.expansion_search(), b.expansion_search()); std::size_t threads = (std::min)(a.limits().threads(), b.limits().threads()); executor_default_t executor{threads}; - join_result_t result = a.join(b, config, a_to_b, b_to_a, executor, progress_t{progress}); + join_result_t result; + { + // Lock the receiver `a`; `b` is read-only from this side. Concurrent + // bidirectional `join(a, b)` and `join(b, a)` from two Python threads + // is unsupported. + py::gil_scoped_release release; + std::unique_lock lock(*a.mutex_ptr_); + result = a.join(b, config, a_to_b, b_to_a, executor, progress_t{progress}); + } forward_error(result); return a_to_b; @@ -844,9 +924,11 @@ static void compact_index(dense_index_py_t& index, std::size_t threads, progress if (!threads) threads = std::thread::hardware_concurrency(); + + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); if (!index.try_reserve(index_limits_t(index.size(), threads))) throw std::invalid_argument("Out of memory!"); - index.compact(executor_default_t{threads}, progress_t{progress}); } @@ -1266,9 +1348,11 @@ PYBIND11_MODULE(compiled, m, py::mod_gil_not_used()) { if (!threads) threads = std::thread::hardware_concurrency(); + + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); if (!index.try_reserve(index_limits_t(index.size(), threads))) throw std::invalid_argument("Out of memory!"); - index.isolate(executor_default_t{threads}); return result.completed; }, @@ -1285,9 +1369,11 @@ PYBIND11_MODULE(compiled, m, py::mod_gil_not_used()) { if (!threads) threads = std::thread::hardware_concurrency(); + + py::gil_scoped_release release; + std::unique_lock lock(*index.mutex_ptr_); if (!index.try_reserve(index_limits_t(index.size(), threads))) throw std::invalid_argument("Out of memory!"); - index.isolate(executor_default_t{threads}); return result.completed; }, diff --git a/python/scripts/test_gil_release.py b/python/scripts/test_gil_release.py new file mode 100644 index 00000000..3cd26c31 --- /dev/null +++ b/python/scripts/test_gil_release.py @@ -0,0 +1,227 @@ +"""GIL-release and progress-callback contract for the Python `Index` binding. + +USearch releases the GIL around long C++ operations so that *other* Python +work (NumPy ops, file I/O, or another Python thread doing unrelated work) can +make progress concurrently. The `Index` API itself is single-threaded from +Python's perspective - one Python thread per index at a time. + +These tests assert the GIL contract end-to-end by: + +* Spawning a background Python thread that increments a counter in a tight + loop while the main thread runs a long USearch op. If the GIL is actually + released, the counter advances meaningfully during the op. +* Validating that progress callbacks fire across the GIL-release boundary - + the callback runs from a C++ worker thread that must reacquire the GIL to + invoke the Python callable, mutate a list, and return a bool. +* Validating that returning `False` from the progress callback terminates the + operation cleanly, surfacing as a Python `RuntimeError`. +""" + +import threading +import time + +import numpy as np +import pytest + +from usearch.index import Index + + +def _background_counter(): + """Returns (start_fn, stop_fn, count_fn) for a tight-loop Python thread.""" + counter = [0] + stop = threading.Event() + + def loop(): + while not stop.is_set(): + counter[0] += 1 + + thread = threading.Thread(target=loop, daemon=True) + + def start(): + thread.start() + + def stop_and_join(): + stop.set() + thread.join() + + return start, stop_and_join, lambda: counter[0] + + +def _big_random_batch(n: int, ndim: int, seed: int = 42): + rng = np.random.default_rng(seed=seed) + keys = np.arange(n, dtype=np.uint64) + vectors = rng.standard_normal((n, ndim), dtype=np.float32) + return keys, vectors + + +# Lower bound on background-counter ticks during one short USearch op. Modern +# hardware loops the trivial `counter[0] += 1` body well over a million times +# per second; 10k is a conservative floor that distinguishes "GIL released" +# from "GIL held" without making the test slow on tiny inputs. +_GIL_TICK_FLOOR = 10_000 + + +def test_gil_released_during_add(): + start, stop_and_join, count = _background_counter() + start() + + idx = Index(ndim=64, dtype="f32") + keys, vectors = _big_random_batch(2_000, 64) + + before = count() + t0 = time.perf_counter() + idx.add(keys, vectors, threads=4) + elapsed = time.perf_counter() - t0 + after = count() + stop_and_join() + + advancement = after - before + assert advancement > _GIL_TICK_FLOOR, ( + f"GIL appears held: only {advancement:,} background ticks during a " + f"{elapsed:.3f}s add. Expected > {_GIL_TICK_FLOOR:,}." + ) + + +def test_gil_released_during_search(): + idx = Index(ndim=64, dtype="f32") + keys, vectors = _big_random_batch(1_500, 64) + idx.add(keys, vectors, threads=4) + + start, stop_and_join, count = _background_counter() + start() + + _, queries = _big_random_batch(1_000, 64, seed=7) + before = count() + t0 = time.perf_counter() + idx.search(queries, 10, threads=4) + elapsed = time.perf_counter() - t0 + after = count() + stop_and_join() + + advancement = after - before + assert advancement > _GIL_TICK_FLOOR, ( + f"GIL appears held during search: only {advancement:,} background ticks during a {elapsed:.3f}s search." + ) + + +def test_progress_callback_fires_and_completes(): + """The progress callback runs from a C++ worker thread that must reacquire + the GIL before invoking the Python callable. It must be able to mutate a + Python list and return a bool without crashing.""" + + idx = Index(ndim=64, dtype="f32") + keys, vectors = _big_random_batch(2_000, 64) + + invocations = [] + + def progress(done: int, total: int) -> bool: + invocations.append((done, total)) + return True + + idx.add(keys, vectors, threads=4, progress=progress) + + assert invocations, "progress callback was never invoked" + last_done, last_total = invocations[-1] + assert last_done == last_total == len(keys), ( + f"final progress {(last_done, last_total)} != ({len(keys)}, {len(keys)})" + ) + # Done counters should be non-decreasing across the run. + for (d_prev, _), (d_next, _) in zip(invocations, invocations[1:]): + assert d_prev <= d_next, f"progress went backwards: {d_prev} -> {d_next}" + + +def test_progress_callback_can_cancel(): + """Returning `False` from the progress callback terminates the op cleanly + and surfaces as a Python `RuntimeError` - no segfault, no UB.""" + + idx = Index(ndim=64, dtype="f32") + keys, vectors = _big_random_batch(10_000, 64) + + seen = [] + + def progress(done: int, total: int) -> bool: + seen.append(done) + # Cancel after a few progress reports so we know the path is exercised. + return len(seen) < 3 + + with pytest.raises(RuntimeError, match="terminated"): + idx.add(keys, vectors, threads=4, progress=progress) + + # Index may be partially populated; the important property is no crash and + # that the callback was actually invoked the expected number of times. + assert len(seen) >= 3 + assert len(idx) <= len(keys) + + +def test_gil_released_with_progress_callback(): + """Combined: background Python thread runs while the main thread is in + `add()` with an active progress callback. Both must work simultaneously.""" + + start, stop_and_join, count = _background_counter() + start() + + idx = Index(ndim=64, dtype="f32") + keys, vectors = _big_random_batch(2_000, 64) + + invocations = [] + + def progress(done: int, total: int) -> bool: + invocations.append((done, total)) + return True + + before = count() + idx.add(keys, vectors, threads=4, progress=progress) + after = count() + stop_and_join() + + assert after - before > _GIL_TICK_FLOOR, ( + "background thread didn't advance during add - GIL likely held while callback was active" + ) + assert invocations and invocations[-1] == (len(keys), len(keys)) + + +def test_concurrent_access_serializes_safely(): + """The documented contract is one Python thread per index; the binding + enforces it with an internal mutex so accidental concurrent access from + multiple Python threads serializes instead of crashing. Mix adds with + disjoint key ranges, searches, and lock-free getters; assert no thread + errors and that all keys from all `add` workers landed in the index.""" + + idx = Index(ndim=64, dtype="f32") + per_thread = 500 + barrier = threading.Barrier(6) + errors: list[str] = [] + + def run(target, *args): + def wrapped(): + try: + barrier.wait() + target(*args) + except Exception as e: + errors.append(f"{type(e).__name__}: {e}") + return threading.Thread(target=wrapped) + + def adder(tid: int): + rng = np.random.default_rng(seed=tid) + base = tid * per_thread + keys = np.arange(base, base + per_thread, dtype=np.uint64) + idx.add(keys, rng.standard_normal((per_thread, 64), dtype=np.float32)) + + def searcher(tid: int): + rng = np.random.default_rng(seed=1000 + tid) + for _ in range(50): + idx.search(rng.standard_normal((1, 64), dtype=np.float32), 3) + + def getters(): + for i in range(1000): + _ = i in idx + _ = len(idx) + + threads = [run(adder, t) for t in range(3)] + [run(searcher, t) for t in range(2)] + [run(getters)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, "thread errors:\n " + "\n ".join(errors) + assert len(idx) == 3 * per_thread diff --git a/rust/lib.cpp b/rust/lib.cpp index b660f3f6..5ca243c5 100644 --- a/rust/lib.cpp +++ b/rust/lib.cpp @@ -298,12 +298,7 @@ std::unique_ptr new_native_index(IndexOptions const& options) { index_dense_config_t config(options.connectivity, options.expansion_add, options.expansion_search); config.multi = options.multi; index_t index = index_t::make(metric, config); - - // Preserve constructor pre-allocation semantics (`index_limits_t{}`), but execute - // reserve after heap allocation to avoid move-induced pointer invalidation. - std::unique_ptr native = wrap(std::move(index)); - native->reserve_capacity_and_threads(0, std::thread::hardware_concurrency()); - return native; + return wrap(std::move(index)); } IndexMetadata head_to_metadata(index_dense_head_t const& head) { diff --git a/rust/lib.rs b/rust/lib.rs index 5640d702..95807483 100644 --- a/rust/lib.rs +++ b/rust/lib.rs @@ -404,14 +404,14 @@ pub mod ffi { pub fn change_expansion_search(self: &NativeIndex, n: usize); pub fn metric_kind(self: &NativeIndex) -> MetricKind; - pub fn change_metric_kind(self: &NativeIndex, metric: MetricKind); + pub fn change_metric_kind(self: &NativeIndex, metric: MetricKind) -> Result<()>; /// Changes the metric function used to calculate the distance between vectors. /// Avoids the `std::ffi::c_void` type and the `StatefulMetric` type, that the FFI /// does not support, replacing them with basic pointer-sized integer types. /// The first two arguments are the pointers to the vectors to compare, and the third /// argument is the `metric_state` propagated from the Rust layer. - pub fn change_metric(self: &NativeIndex, metric: usize, metric_state: usize); + pub fn change_metric(self: &NativeIndex, metric: usize, metric_state: usize) -> Result<()>; pub fn new_native_index(options: &IndexOptions) -> Result>; @@ -905,9 +905,7 @@ impl VectorType for f32 { Some(MetricFunction::F32Metric(metric)) => metric as *mut () as usize, _ => panic!("Expected F32Metric"), }; - index.inner.change_metric(trampoline_fn, closure_address); - - Ok(()) + index.inner.change_metric(trampoline_fn, closure_address) } } @@ -979,9 +977,7 @@ impl VectorType for i8 { Some(MetricFunction::I8Metric(metric)) => metric as *mut () as usize, _ => panic!("Expected I8Metric"), }; - index.inner.change_metric(trampoline_fn, closure_address); - - Ok(()) + index.inner.change_metric(trampoline_fn, closure_address) } } @@ -1046,9 +1042,7 @@ impl VectorType for u8 { Some(MetricFunction::U8Metric(metric)) => metric as *mut () as usize, _ => panic!("Expected U8Metric"), }; - index.inner.change_metric(trampoline_fn, closure_address); - - Ok(()) + index.inner.change_metric(trampoline_fn, closure_address) } } @@ -1120,9 +1114,7 @@ impl VectorType for f64 { Some(MetricFunction::F64Metric(metric)) => metric as *mut () as usize, _ => panic!("Expected F64Metric"), }; - index.inner.change_metric(trampoline_fn, closure_address); - - Ok(()) + index.inner.change_metric(trampoline_fn, closure_address) } } @@ -1195,9 +1187,7 @@ impl VectorType for f16 { Some(MetricFunction::F16Metric(metric)) => metric as *mut () as usize, _ => panic!("Expected F16Metric"), }; - index.inner.change_metric(trampoline_fn, closure_address); - - Ok(()) + index.inner.change_metric(trampoline_fn, closure_address) } } @@ -1270,9 +1260,7 @@ impl VectorType for b1x8 { Some(MetricFunction::B1X8Metric(metric)) => metric as *mut () as usize, _ => panic!("Expected F1X8Metric"), }; - index.inner.change_metric(trampoline_fn, closure_address); - - Ok(()) + index.inner.change_metric(trampoline_fn, closure_address) } } @@ -1409,7 +1397,7 @@ impl Index { } /// Changes the metric kind used to calculate the distance between vectors. - pub fn change_metric_kind(self: &Index, metric: ffi::MetricKind) { + pub fn change_metric_kind(self: &Index, metric: ffi::MetricKind) -> Result<(), cxx::Exception> { self.inner.change_metric_kind(metric) }