diff --git a/include/svs/core/data/simple.h b/include/svs/core/data/simple.h index 819a330d8..a21ff7ddd 100644 --- a/include/svs/core/data/simple.h +++ b/include/svs/core/data/simple.h @@ -678,6 +678,23 @@ template class Blocked { template inline constexpr bool is_blocked_v = false; template inline constexpr bool is_blocked_v> = true; +template struct remove_blocking { + using type = Alloc; +}; +template struct remove_blocking> { + using type = Alloc; +}; +template using remove_blocked_t = typename remove_blocking::type; + +template +constexpr remove_blocked_t remove_blocked(const Alloc& alloc) { + if constexpr (is_blocked_v) { + return alloc.get_allocator(); + } else { + return alloc; + } +} + } // namespace data namespace lib::detail { @@ -696,6 +713,9 @@ class SimpleData> { public: ///// Static Members + /// The static dimensionality of the underlying data. + static constexpr size_t extent = Extent; + /// /// Default block size in bytes. /// diff --git a/include/svs/extensions/vamana/scalar.h b/include/svs/extensions/vamana/scalar.h index 021dec305..a88d15c11 100644 --- a/include/svs/extensions/vamana/scalar.h +++ b/include/svs/extensions/vamana/scalar.h @@ -94,4 +94,109 @@ VamanaBuildAdaptor svs_invoke( return VamanaBuildAdaptor{adapt_for_self(data, distance)}; } +template +class TransactionData { + public: + using points_type = Points; + using const_point_type = typename Points::const_value_type; + + using value_type = typename Data::value_type; + using const_value_type = typename Data::const_value_type; + + TransactionData( + const Data& data, + const Points& points, + DataPoints data_points, + std::span slots + ) + : data_(data) + , points_(points) + , data_points_(std::move(data_points)) + , slots_(slots.begin(), slots.end()) { + assert(std::is_sorted(slots_.begin(), slots_.end()) && "Slots must be sorted"); + } + + const Data& get_data() const { return data_; } + + // Get the index of the point corresponding to the given index' id. + size_t get_point_index(size_t id) const { + // find id in slots + auto it = std::lower_bound(slots_.begin(), slots_.end(), id); + if (it != slots_.end() && *it == id) { + return std::distance(slots_.begin(), it); + } else { + throw std::out_of_range("Index not found in slots"); + } + } + + const_point_type get_point(size_t id) const { + return points_.get_datum(get_point_index(id)); + } + + const points_type& get_points() const { return points_; } + size_t num_points() const { return points_.size(); } + size_t get_slot(size_t i) const { return slots_[i]; } + + size_t size() const { + return std::max(data_.size(), slots_.empty() ? 0 : slots_.back() + 1); + } + size_t dimensions() const { return data_.dimensions(); } + const_value_type get_datum(size_t i) const { + // find id in slots + auto it = std::lower_bound(slots_.begin(), slots_.end(), i); + if (it != slots_.end() && *it == i) { + return data_points_.get_datum(std::distance(slots_.begin(), it)); + } else { + return data_.get_datum(i); + } + } + void prefetch(size_t) const {} // data_.prefetch(i); } + + template void copy_points(TargetData& target_data) const { + assert( + &target_data == &data_ && "Target data must be the same as the original data" + ); + // auto new_size = std::max_element(slots_.begin(), slots_.end()) + 1; + // assuming, slots_ is ordered, we can just take the last element and add 1 + auto new_size = slots_.back() + 1; + if (new_size > target_data.size()) { + target_data.resize(new_size); + } + for (size_t i = 0; i < points_.size(); ++i) { + target_data.set_datum(slots_[i], points_.get_datum(i)); + } + } + + private: + const Data& data_; + const Points& points_; + const DataPoints data_points_; + std::vector slots_; +}; + +template +auto svs_invoke( + svs::tag_t, + const Data& data, + const Points& points, + std::span slots, + Pool& pool +) { + using element_type = typename Data::element_type; + using points_type = SQDataset< + element_type, + Data::extent, + data::remove_blocked_t>; + using compressor_type = + detail::Compressor; + + const auto scale = data.get_scale(); + const auto bias = data.get_bias(); + auto compressor = compressor_type{scale, bias}; + auto compressed = compressor(points, pool, data::remove_blocked(data.get_allocator())); + return TransactionData( + data, points, points_type{std::move(compressed), scale, bias}, slots + ); +} + } // namespace svs::quantization::scalar diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 5f0ce7c16..9f1a29e77 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -107,6 +107,183 @@ class ValidBuilder { const std::vector& status_; }; +/// @brief transactional graph wrapper for a graph that allows for transactional updates to +/// the graph structure without modifying the underlying graph. +/// @tparam Graph The type of the underlying immutable graph. +template struct TransactionGraph { + using index_type = typename Graph::index_type; + using value_type = typename Graph::value_type; + using const_value_type = typename Graph::const_value_type; + using reference = typename Graph::reference; + using const_reference = typename Graph::const_reference; + using updated_nodes_type = std::vector>; + using slots_map_type = std::vector; //, index_type>; + static constexpr index_type invalid_index = static_cast(-1); + + private: + const Graph& graph_; + const size_t num_nodes_; + const size_t new_size_; + updated_nodes_type updated_nodes_; + std::unique_ptr nodes_mutex_; + slots_map_type slots_map_; + size_t num_slots_ = 0; + + static inline bool has_edge_in_span(const_reference neighbors, index_type to) { + return std::find(neighbors.begin(), neighbors.end(), to) != neighbors.end(); + } + + index_type get_slot_index_unsafe(index_type id) const { return slots_map_[id]; } + + index_type get_slot_index(index_type id) const { + std::shared_lock lock(*nodes_mutex_); + return get_slot_index_unsafe(id); + } + + index_type get_or_add_slot_unsafe(index_type id) { + if (slots_map_[id] != invalid_index) { + return slots_map_[id]; + } + index_type new_slot = num_slots_++; + slots_map_[id] = new_slot; + updated_nodes_.emplace_back(); + return new_slot; + } + + index_type get_or_add_slot(index_type id) { + index_type slot_index = get_slot_index(id); + if (slot_index != invalid_index) { + return slot_index; + } + std::lock_guard lock(*nodes_mutex_); + return get_or_add_slot_unsafe(id); + } + + index_type get_or_load_slot(index_type id) { + index_type slot_index = get_slot_index(id); + if (slot_index != invalid_index) { + return slot_index; + } + std::lock_guard lock(*nodes_mutex_); + slot_index = get_slot_index_unsafe(id); + if (slot_index != invalid_index) { + return slot_index; + } + + slot_index = num_slots_++; + slots_map_[id] = slot_index; + auto node = graph_.get_node(id); + updated_nodes_.emplace_back(node.begin(), node.end()); + return slot_index; + } + + public: + TransactionGraph(const Graph& graph, std::span slots) + : graph_(graph) + , num_nodes_(graph.n_nodes()) + , new_size_(std::max(num_nodes_, slots.empty() ? 0 : slots.back() + 1)) + , nodes_mutex_(std::make_unique()) + , slots_map_(new_size_, invalid_index) + , num_slots_(slots.size()) { + updated_nodes_.reserve(std::min(slots.size() * graph_.max_degree(), new_size_)); + updated_nodes_.resize(slots.size()); + for (size_t i = 0; i < slots.size(); ++i) { + slots_map_[slots[i]] = i; + } + } + + size_t n_nodes() const { return new_size_; } + + const_reference get_node(index_type id) const { + index_type slot_index = get_slot_index(id); + if (slot_index != invalid_index) { + return updated_nodes_[slot_index]; + } else { + return graph_.get_node(id); + } + } + + bool has_edge(index_type from, index_type to) const { + index_type slot_index = get_slot_index(from); + if (slot_index != invalid_index) { + return has_edge_in_span(updated_nodes_[slot_index], to); + } else { + return graph_.has_edge(from, to); + } + } + + void clear_node(index_type id) { + index_type slot_index = get_or_add_slot(id); + updated_nodes_[slot_index].clear(); + } + + void replace_node(index_type id, std::vector neighbors) { + replace_node(id, std::span(neighbors)); + } + + void replace_node(index_type id, std::span neighbors) { + auto slot_index = get_or_add_slot(id); + updated_nodes_[slot_index].assign(neighbors.begin(), neighbors.end()); + + for (auto neighbor : neighbors) { + get_or_load_slot(neighbor); + } + } + + size_t get_node_degree(index_type id) const { + auto slot_index = get_slot_index(id); + if (slot_index != invalid_index) { + return updated_nodes_[slot_index].size(); + } else { + return graph_.get_node_degree(id); + } + } + + size_t max_degree() const { return graph_.max_degree(); } + + size_t add_edge(index_type from, index_type to) { + if (from == to) { + return get_node_degree(from); + ; + } + if (from >= new_size_ || to >= new_size_) { + throw std::out_of_range("Index out of range in TransactionalGraph::add_edge"); + } + auto slot_index = get_or_load_slot(from); + auto& updated_node = updated_nodes_[slot_index]; + if (updated_node.size() >= graph_.max_degree()) { + return updated_node.size(); + } + if (!has_edge_in_span(updated_node, to)) { + updated_node.push_back(to); + } + return updated_node.size(); + } + + void prefetch_node(index_type id) const { + auto slot_index = get_slot_index(id); + if (slot_index == invalid_index) { + graph_.prefetch_node(id); + } + } + + template void copy_edges(G& target_graph) const { + assert( + &target_graph == &graph_ && "Target graph must be same as the original graph" + ); + if (new_size_ > target_graph.n_nodes()) { + target_graph.unsafe_resize(new_size_); + } + for (size_t id = 0; id < slots_map_.size(); ++id) { + index_type slot_index = get_slot_index_unsafe(id); + if (slot_index != invalid_index) { + auto& src = updated_nodes_[slot_index]; + target_graph.replace_node(id, src); + } + } + } +}; + template class MutableVamanaIndex { friend class MultiMutableVamanaIndex; @@ -156,7 +333,8 @@ class MutableVamanaIndex { // Thread local data structures. distance_type distance_; - threads::ThreadPoolHandle threadpool_; + mutable threads::ThreadPoolHandle + threadpool_; // TODO: Remove mutable and const_casts in the code. lib::ReadWriteProtected search_parameters_; // Configurations @@ -611,26 +789,10 @@ class MutableVamanaIndex { ); } - /// - /// @brief Add the points with the given external IDs to the dataset. - // - /// When `delete_entries` is called, a soft deletion is performed, marking the entries - /// as `deleted`. When `consolidate` is called, the state of these deleted entries - /// becomes `empty`. When `add_points` is called with the `reuse_empty` flag enabled, - /// the memory is scanned from the beginning to locate and fill these empty entries with - /// new points. - /// - /// @param points Dataset of points to add. - /// @param external_ids The external IDs of the corresponding points. Must be a - /// container implementing forward iteration. - /// @param reuse_empty A flag that determines whether to reuse empty entries that may - /// exist after deletion and consolidation. When enabled, scan from the beginning to - /// find and fill these empty entries when adding new points. - /// template - std::vector add_points( + auto add_points_compute_changes( const Points& points, const ExternalIds& external_ids, bool reuse_empty = false - ) { + ) const { const size_t num_points = points.size(); const size_t num_ids = external_ids.size(); if (num_points != num_ids) { @@ -664,30 +826,16 @@ class MutableVamanaIndex { size_t needed = num_points - slots.size(); size_t current_size = data_.size(); size_t new_size = current_size + needed; - data_.resize(new_size); - - // Graph resizing marked as un-safe because graph contain internal references - // and thus it's not a good idea to go around shrinking the graph without care. - // - // However, we are only growing here, so resizing will not change any - // invariants. - graph_.unsafe_resize(new_size); - status_.resize(new_size, SlotMetadata::Empty); - // Append the correct number of extra slots. - threads::UnitRange extra_points{current_size, current_size + needed}; + threads::UnitRange extra_points{current_size, new_size}; slots.insert(slots.end(), extra_points.begin(), extra_points.end()); } assert(slots.size() == num_points); - // Try to update the id translation now that we have internal ids. - // If this fails, we still haven't mutated the index data structure so we're safe - // to throw an exception. - translator_.insert(external_ids, slots); - - // Copy the given points into the data and clear the adjacency lists for the graph. - copy_points(points, slots); - clear_lists(slots); + // Create transaction data wrapper + auto data_wrapper = + extensions::transaction_data_builder(data_, points, slots, threadpool_); + TransactionGraph graph_wrapper{graph_, slots}; // Patch in the new neighbors. auto parameters = VamanaBuildParameters{ @@ -702,8 +850,8 @@ class MutableVamanaIndex { auto prefetch_parameters = GreedySearchPrefetchParameters{sp.prefetch_lookahead_, sp.prefetch_step_}; VamanaBuilder builder{ - graph_, - data_, + graph_wrapper, + data_wrapper, distance_, parameters, threadpool_, @@ -711,8 +859,35 @@ class MutableVamanaIndex { logger_, logging::Level::Trace}; builder.construct(alpha_, entry_point(), slots, logging::Level::Trace, logger_); + return std::make_tuple( + std::move(data_wrapper), std::move(graph_wrapper), std::move(slots) + ); + } + + template + std::vector add_points_commit( + const ExternalIds& external_ids, + const std::tuple, std::vector>& + changes + ) { + auto& [data_wrapper, graph_wrapper, slots] = changes; + + // Try to update the id translation now that we have internal ids. + // If this fails, we still haven't mutated the index data structure so we're safe + // to throw an exception. + translator_.insert(external_ids, slots); + + // Update data_ + data_wrapper.copy_points(data_); + graph_wrapper.copy_edges(graph_); + // Mark all added entries as valid. - for (const auto& i : slots) { + auto new_size = data_.size(); + assert(graph_.n_nodes() == new_size); + assert(status_.size() <= new_size); + status_.resize(new_size, SlotMetadata::Valid); + for (auto i : slots) { + assert(i < new_size); status_[i] = SlotMetadata::Valid; } @@ -722,6 +897,30 @@ class MutableVamanaIndex { return slots; } + /// + /// @brief Add the points with the given external IDs to the dataset. + // + /// When `delete_entries` is called, a soft deletion is performed, marking the entries + /// as `deleted`. When `consolidate` is called, the state of these deleted entries + /// becomes `empty`. When `add_points` is called with the `reuse_empty` flag enabled, + /// the memory is scanned from the beginning to locate and fill these empty entries with + /// new points. + /// + /// @param points Dataset of points to add. + /// @param external_ids The external IDs of the corresponding points. Must be a + /// container implementing forward iteration. + /// @param reuse_empty A flag that determines whether to reuse empty entries that may + /// exist after deletion and consolidation. When enabled, scan from the beginning to + /// find and fill these empty entries when adding new points. + /// + template + std::vector add_points( + const Points& points, const ExternalIds& external_ids, bool reuse_empty = false + ) { + auto changes = add_points_compute_changes(points, external_ids, reuse_empty); + return add_points_commit(external_ids, changes); + } + /// /// Delete all IDs stored in the random-access container `ids`. /// diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 96ffeaa06..711d5e5b0 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -27,6 +27,9 @@ #include "svs/lib/preprocessor.h" #include "svs/lib/threads.h" +#include +#include + namespace svs::index::vamana::extensions { ///// @@ -661,4 +664,238 @@ double svs_invoke( return static_cast(dist); } +template < + data::ImmutableMemoryDataset Data, + data::ImmutableMemoryDataset Points, + typename DataPoints> +class TransactionData { + public: + static constexpr size_t out_of_points_id = std::numeric_limits::max(); + + using points_type = Points; + using const_point_type = typename Points::const_value_type; + + using value_type = typename Data::value_type; + using const_value_type = typename Data::const_value_type; + + TransactionData( + const Data& data, + const Points& points, + DataPoints data_points, + std::span slots + ) + : data_(data) + , points_(points) + , data_points_(std::move(data_points)) + , slots_(slots.begin(), slots.end()) { + // slots to be divided to 2 parts: 1: slots within data_.size() and 2: slots with + // index >= data_.size() appropriate info should be stored in fields to allow point + // index computation in form: if id >= data_.size() then compute points index in + // form: id - data_.size() + out_of_data_slots_begin else compute points index in + // form: std::lower_bound(slots_.begin(), slots_.end(), id) - slots_.begin() + out_of_data_slots_begin_ = + std::lower_bound(slots_.begin(), slots_.end(), data_.size()) - slots_.begin(); + assert(std::is_sorted(slots_.begin(), slots_.end()) && "Slots must be sorted"); + } + + const Data& get_data() const { return data_; } + + // Get the index of the point corresponding to the given index' id. + size_t get_point_index(size_t id) const { + // find id in slots + if (id >= data_.size()) [[likely]] { + auto point_index = id - data_.size() + out_of_data_slots_begin_; + assert(point_index < points_.size() && "Point index out of bounds"); + if (point_index < points_.size()) [[likely]] { + return point_index; + } + } else { + auto it = std::lower_bound(slots_.begin(), slots_.end(), id); + if (it != slots_.end() && *it == id) { + return std::distance(slots_.begin(), it); + } + } + return out_of_points_id; + } + + const_point_type get_point(size_t id) const { + auto point_index = get_point_index(id); + if (point_index == out_of_points_id) { + throw std::out_of_range("ID not found in transaction data points"); + } + return points_.get_datum(point_index); + } + + const points_type& get_points() const { return points_; } + size_t num_points() const { return points_.size(); } + size_t get_slot(size_t i) const { return slots_[i]; } + + size_t size() const { + return std::max(data_.size(), slots_.empty() ? 0 : slots_.back() + 1); + } + size_t dimensions() const { return data_.dimensions(); } + const_value_type get_datum(size_t i) const { + auto point_index = get_point_index(i); + if (point_index != out_of_points_id) { + return data_points_.get_datum(point_index); + } else { + return data_.get_datum(i); + } + } + void prefetch(size_t i) const { + auto point_index = get_point_index(i); + if (point_index != out_of_points_id) { + data_points_.prefetch(point_index); + } else { + data_.prefetch(i); + } + } + + template void copy_points(TargetData& target_data) const { + assert( + &target_data == &data_ && "Target data must be the same as the original data" + ); + auto new_size = slots_.back() + 1; + if (new_size > target_data.size()) { + target_data.resize(new_size); + } + for (size_t i = 0; i < points_.size(); ++i) { + target_data.set_datum(slots_[i], points_.get_datum(i)); + } + } + + private: + const Data& data_; + const Points& points_; + const DataPoints data_points_; + std::vector slots_; + size_t out_of_data_slots_begin_; +}; + +struct TransactionDataBuilder { + template < + data::ImmutableMemoryDataset Data, + data::ImmutableMemoryDataset Points, + threads::ThreadPool Pool> + auto operator()( + const Data& data, const Points& points, std::span slots, Pool& pool + ) const { + return svs_invoke(*this, data, points, slots, pool); + } +}; + +inline constexpr TransactionDataBuilder transaction_data_builder{}; + +template +auto svs_invoke( + svs::tag_t, + const Data& data, + const Points& points, + std::span slots, + Pool& pool +) { + assert(points.size() == slots.size() && "Points and slots must have the same size"); + using data_element_type = typename Data::element_type; + constexpr size_t extent = Data::extent; + using allocator_type = typename data::remove_blocked_t; + using data_points_type = data::SimpleData; + + auto data_points = data_points_type{ + points.size(), + svs::lib::MaybeStatic(data.dimensions()), + data::remove_blocked(data.get_allocator())}; + svs::threads::parallel_for( + pool, + svs::threads::StaticPartition(points.size()), + [&](auto is, auto SVS_UNUSED(tid)) { + for (auto i : is) { + data_points.set_datum(i, points.get_datum(i)); + } + } + ); + return TransactionData(data, points, std::move(data_points), slots); +} + +/// @brief Implementation for transaction dataset/vamana build adaptors. +template +struct TransactionBuildAdaptor { + public: + using inner_adaptor_type = InnerAdaptor; + // using transaction_data_type = TransactionData; + + TransactionBuildAdaptor( + /*const TransactionData& data, */ Distance search_distance, + InnerAdaptor inner_adaptor + ) + : /*data_(data) + , */ + search_distance_(std::move(search_distance)) + , inner_adaptor_(std::move(inner_adaptor)) {} + + using search_distance_type = Distance; + using general_distance_type = typename InnerAdaptor::general_distance_type; + + // Helper template + template using const_point_type = typename Data::const_point_type; + + template + const_point_type access_query_for_graph_search(const Data& data, size_t i) const { + return data.get_point(i); + } + + auto graph_search_accessor() const { return inner_adaptor_.graph_search_accessor(); } + + search_distance_type& graph_search_distance() { return search_distance_; } + + general_distance_type& general_distance() { return inner_adaptor_.general_distance(); } + + auto general_accessor() const { return inner_adaptor_.general_accessor(); } + + template + SVS_FORCE_INLINE const const_point_type& modify_post_search_query( + const Data& SVS_UNUSED(data), + size_t SVS_UNUSED(i), + const const_point_type& pre_search_query + ) const { + return pre_search_query; + } + + static constexpr bool refix_argument_after_search = + inner_adaptor_type::refix_argument_after_search; + + template //, typename Q> + SVS_FORCE_INLINE Neighbor post_search_modify( + const Data& data, + general_distance_type& SVS_UNUSED(d), + const const_point_type& query, + // const Q& query, + const N& n + ) const { + if constexpr (RebuildNeighbour) { + auto id = n.id(); + return Neighbor( + id, distance::compute(search_distance_, query, data.get_datum(id)) + ); + } else { + return n; + } + } + + private: + // const transaction_data_type& data_; + Distance search_distance_; + inner_adaptor_type inner_adaptor_; +}; + +// Default to the ``TransactionBuildAdaptor``. +template +auto svs_invoke( + svs::tag_t, + const TransactionData& data, + const Distance& distance +) { + return TransactionBuildAdaptor{ + /*data, */ distance, build_adaptor(data.get_data(), distance)}; +} + } // namespace svs::index::vamana::extensions diff --git a/include/svs/index/vamana/multi.h b/include/svs/index/vamana/multi.h index b7f3d9b53..2be07cf7c 100644 --- a/include/svs/index/vamana/multi.h +++ b/include/svs/index/vamana/multi.h @@ -183,16 +183,33 @@ class MultiMutableVamanaIndex { template void prepare_added_id_by_label(const Labels& labels, std::vector& adds) { - for (const auto l : labels) { + adds = make_added_ids(labels); + return commit_added_id_by_label(labels, adds); + } + + template + std::vector make_added_ids(const Labels& labels) const { + std::vector adds(labels.size()); + std::iota(adds.begin(), adds.end(), counter_); + return adds; + } + + template + void commit_added_id_by_label( + const Labels& labels, const std::vector& adds + ) { + assert(labels.size() == adds.size()); + for (size_t i = 0; i < labels.size(); ++i) { + const auto& l = labels[i]; if (label_to_external_.find(l) == label_to_external_.end()) { label_to_external_.insert({l, std::vector{}}); } - size_t new_external_id = counter_++; + size_t new_external_id = adds[i]; label_to_external_[l].push_back(new_external_id); external_to_label_.insert({new_external_id, l}); - adds.push_back(new_external_id); } + counter_ += labels.size(); } public: @@ -357,21 +374,38 @@ class MultiMutableVamanaIndex { template std::vector add_points(const Points& points, const Labels& labels, bool reuse_empty = false) { + auto changes = add_points_compute_changes(points, labels, reuse_empty); + return add_points_commit(labels, std::move(changes)); + } + + template + auto add_points_compute_changes( + const Points& points, const Labels& labels, bool reuse_empty = false + ) const { const size_t num_points = points.size(); const size_t num_labels = labels.size(); if (num_points != num_labels) { throw ANNEXCEPTION( - "Number of points ({}) not equal to the number of external ids ({})!", + "Number of points ({}) not equal to the number of labels ({})!", num_points, num_labels ); } - std::vector adds; - adds.reserve(num_labels); - prepare_added_id_by_label(labels, adds); - index_->add_points(points, adds, reuse_empty); + auto adds = make_added_ids(labels); + return std::make_tuple( + std::move(adds), index_->add_points_compute_changes(points, adds, reuse_empty) + ); + } + template + std::vector add_points_commit( + const Labels& labels, + const std::tuple, InternalChanges>& changes + ) { + auto& [adds, internal_changes] = changes; + index_->add_points_commit(adds, internal_changes); + commit_added_id_by_label(labels, adds); return adds; } diff --git a/include/svs/lib/concurrency/upgradable_mutex.h b/include/svs/lib/concurrency/upgradable_mutex.h new file mode 100644 index 000000000..71d404f19 --- /dev/null +++ b/include/svs/lib/concurrency/upgradable_mutex.h @@ -0,0 +1,441 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/lib/exception.h" + +#include +#include +#include + +namespace svs::lib { + +/// +/// @brief A mutex supporting shared, upgrade, and exclusive ownership. +/// +/// In addition to the shared/exclusive ownership offered by ``std::shared_mutex``, this +/// mutex introduces a third, intermediate "upgrade" ownership level: +/// +/// - **Shared** ownership may be held by any number of threads simultaneously. Compatible +/// with other shared owners and with a single upgrade owner. +/// - **Upgrade** ownership may be held by at most one thread at a time. It is compatible +/// with any number of shared owners but excludes other upgrade owners and exclusive +/// owners. An upgrade owner may atomically transition to exclusive ownership via +/// ``unlock_upgrade_and_lock`` (or the ``UpgradeToUniqueLock`` helper). +/// - **Exclusive** ownership may be held by at most one thread at a time and is +/// incompatible with every other kind of ownership. +/// +/// The typical use case is a reader that decides, while holding a shared/upgrade lock, that +/// it must perform a write. Because at most one upgrade owner can exist, such an owner can +/// upgrade to exclusive ownership without releasing the lock and racing with another +/// writer. +/// +/// This type satisfies the C++ named requirements *Lockable*, *SharedLockable*, and +/// *SharedTimedLockable* is intentionally not provided (no timed operations are exposed). +/// The interface is compatible with ``std::unique_lock`` (exclusive) and +/// ``std::shared_lock`` (shared). Upgrade ownership is managed via ``UpgradeLock``. +/// +class UpgradableMutex { + public: + UpgradableMutex() = default; + ~UpgradableMutex() = default; + + // Non-copyable and non-movable. + UpgradableMutex(const UpgradableMutex&) = delete; + UpgradableMutex& operator=(const UpgradableMutex&) = delete; + UpgradableMutex(UpgradableMutex&&) = delete; + UpgradableMutex& operator=(UpgradableMutex&&) = delete; + + ///// Exclusive ownership + + /// @brief Acquire exclusive ownership, blocking until it is available. + void lock() { + std::unique_lock lock{mutex_}; + // Wait until no other thread holds exclusive or upgrade ownership, then reserve + // exclusive ownership by setting the write flag. This prevents new shared/upgrade + // owners from being admitted. + write_gate_.wait(lock, [this]() { + return (state_ & (write_entered_ | upgrade_entered_)) == 0; + }); + state_ |= write_entered_; + // Drain any remaining shared owners. + reader_gate_.wait(lock, [this]() { return (state_ & n_readers_) == 0; }); + } + + /// @brief Attempt to acquire exclusive ownership without blocking. + /// @returns ``true`` if ownership was acquired, ``false`` otherwise. + bool try_lock() { + std::lock_guard lock{mutex_}; + if (state_ != 0) { + return false; + } + state_ = write_entered_; + return true; + } + + /// @brief Release exclusive ownership. + void unlock() { + { + std::lock_guard lock{mutex_}; + state_ = 0; + } + write_gate_.notify_all(); + } + + ///// Shared ownership + + /// @brief Acquire shared ownership, blocking until it is available. + void lock_shared() { + std::unique_lock lock{mutex_}; + write_gate_.wait(lock, [this]() { + return (state_ & write_entered_) == 0 && (state_ & n_readers_) != n_readers_; + }); + set_readers(readers() + 1); + } + + /// @brief Attempt to acquire shared ownership without blocking. + /// @returns ``true`` if ownership was acquired, ``false`` otherwise. + bool try_lock_shared() { + std::lock_guard lock{mutex_}; + if ((state_ & write_entered_) != 0 || (state_ & n_readers_) == n_readers_) { + return false; + } + set_readers(readers() + 1); + return true; + } + + /// @brief Release shared ownership. + void unlock_shared() { + // Determine which gate to signal before releasing the mutex so the decision is + // made atomically with the state update, but the actual notify happens after the + // lock is released. This avoids the notified thread immediately re-contending on + // mutex_ before it can make progress. + bool notify_readers = false; + bool notify_writers = false; + { + std::lock_guard lock{mutex_}; + unsigned num_readers = readers() - 1; + set_readers(num_readers); + if ((state_ & write_entered_) != 0) { + // A writer is waiting to drain the readers. + notify_readers = (num_readers == 0); + } else if (num_readers == n_readers_ - 1) { + // We just freed a reader slot that a would-be reader may be waiting on. + notify_writers = true; + } + } + if (notify_readers) { + reader_gate_.notify_one(); + } else if (notify_writers) { + write_gate_.notify_one(); + } + } + + ///// Upgrade ownership + + /// @brief Acquire upgrade ownership, blocking until it is available. + void lock_upgrade() { + std::unique_lock lock{mutex_}; + write_gate_.wait(lock, [this]() { + return (state_ & (write_entered_ | upgrade_entered_)) == 0 && + (state_ & n_readers_) != n_readers_; + }); + state_ |= upgrade_entered_; + set_readers(readers() + 1); + } + + /// @brief Attempt to acquire upgrade ownership without blocking. + /// @returns ``true`` if ownership was acquired, ``false`` otherwise. + bool try_lock_upgrade() { + std::lock_guard lock{mutex_}; + if ((state_ & (write_entered_ | upgrade_entered_)) != 0 || + (state_ & n_readers_) == n_readers_) { + return false; + } + state_ |= upgrade_entered_; + set_readers(readers() + 1); + return true; + } + + /// @brief Release upgrade ownership. + void unlock_upgrade() { + { + std::lock_guard lock{mutex_}; + // Clear the upgrade flag and decrement the reader count in one write to + // avoid a redundant read-modify-write on state_. + state_ = (state_ & ~(upgrade_entered_ | n_readers_)) | + ((readers() - 1) & n_readers_); + } + write_gate_.notify_all(); + } + + ///// Ownership transitions + + /// @brief Atomically convert upgrade ownership into exclusive ownership. + /// + /// Blocks until all other shared owners have released their shared ownership. The + /// calling thread must already hold upgrade ownership. + void unlock_upgrade_and_lock() { + std::unique_lock lock{mutex_}; + // Drop our own reader slot and the upgrade flag, then reserve exclusive ownership. + unsigned num_readers = readers() - 1; + state_ &= ~upgrade_entered_; + state_ |= write_entered_; + set_readers(num_readers); + // Wait for the remaining shared owners to drain. + reader_gate_.wait(lock, [this]() { return (state_ & n_readers_) == 0; }); + } + + /// @brief Attempt to convert upgrade ownership into exclusive ownership without + /// blocking. + /// + /// Succeeds only if the calling thread is the sole owner of the mutex. + /// @returns ``true`` if the transition succeeded, ``false`` otherwise. + bool try_unlock_upgrade_and_lock() { + std::lock_guard lock{mutex_}; + // Success requires that the only reader is our own upgrade ownership. + if (state_ != (upgrade_entered_ | 1U)) { + return false; + } + state_ = write_entered_; + return true; + } + + /// @brief Atomically convert exclusive ownership into upgrade ownership. + /// + /// Does not block. The calling thread must already hold exclusive ownership. + void unlock_and_lock_upgrade() { + { + std::lock_guard lock{mutex_}; + state_ = upgrade_entered_ | 1U; + } + write_gate_.notify_all(); + } + + /// @brief Atomically convert upgrade ownership into shared ownership. + /// + /// Does not block. The calling thread must already hold upgrade ownership. + void unlock_upgrade_and_lock_shared() { + { + std::lock_guard lock{mutex_}; + state_ &= ~upgrade_entered_; + } + // While upgrade ownership was held, write_entered_ was necessarily 0, so no + // shared readers were waiting on write_gate_. Only upgrade and exclusive waiters + // were blocked by upgrade_entered_, and at most one of them can succeed. + write_gate_.notify_one(); + } + + /// @brief Atomically convert exclusive ownership into shared ownership. + /// + /// Does not block. The calling thread must already hold exclusive ownership. + void unlock_and_lock_shared() { + { + std::lock_guard lock{mutex_}; + state_ = 1U; + } + write_gate_.notify_all(); + } + + private: + // The high bit encodes exclusive ("write") ownership, the next-highest bit encodes + // upgrade ownership, and the remaining bits count the number of shared owners (an + // upgrade owner also counts as a shared owner). + static constexpr unsigned write_entered_ = 1U << (sizeof(unsigned) * CHAR_BIT - 1); + static constexpr unsigned upgrade_entered_ = write_entered_ >> 1; + static constexpr unsigned n_readers_ = ~(write_entered_ | upgrade_entered_); + + unsigned readers() const { return state_ & n_readers_; } + void set_readers(unsigned num_readers) { + state_ = (state_ & ~n_readers_) | (num_readers & n_readers_); + } + + std::mutex mutex_{}; + // Notifies threads waiting to acquire shared, upgrade, or exclusive ownership. + std::condition_variable write_gate_{}; + // Notifies an exclusive/upgrading owner waiting for shared owners to drain. + std::condition_variable reader_gate_{}; + unsigned state_ = 0; +}; + +/// +/// @brief RAII wrapper managing upgrade ownership of an ``UpgradableMutex``. +/// +/// Analogous to ``std::shared_lock`` and ``std::unique_lock``, but for the upgrade +/// ownership level. Supports deferred, try-to-lock, and adopt-lock construction and is +/// movable but not copyable. +/// +/// @tparam Mutex The mutex type. Must provide ``lock_upgrade``, ``try_lock_upgrade``, and +/// ``unlock_upgrade``. +/// +template class UpgradeLock { + public: + using mutex_type = Mutex; + + /// @brief Construct without an associated mutex. Owns no lock. + UpgradeLock() noexcept = default; + + /// @brief Acquire upgrade ownership of ``m``, blocking until available. + explicit UpgradeLock(mutex_type& m) + : mutex_{&m} + , owns_{true} { + mutex_->lock_upgrade(); + } + + /// @brief Associate with ``m`` without acquiring ownership. + UpgradeLock(mutex_type& m, std::defer_lock_t) noexcept + : mutex_{&m} + , owns_{false} {} + + /// @brief Attempt to acquire upgrade ownership of ``m`` without blocking. + UpgradeLock(mutex_type& m, std::try_to_lock_t) + : mutex_{&m} + , owns_{m.try_lock_upgrade()} {} + + /// @brief Assume that the calling thread already holds upgrade ownership of ``m``. + UpgradeLock(mutex_type& m, std::adopt_lock_t) + : mutex_{&m} + , owns_{true} {} + + UpgradeLock(const UpgradeLock&) = delete; + UpgradeLock& operator=(const UpgradeLock&) = delete; + + /// @brief Move construct, transferring ownership from ``other``. + UpgradeLock(UpgradeLock&& other) noexcept + : mutex_{other.mutex_} + , owns_{other.owns_} { + other.mutex_ = nullptr; + other.owns_ = false; + } + + /// @brief Move assign, releasing any currently held ownership first. + UpgradeLock& operator=(UpgradeLock&& other) noexcept { + if (this != &other) { + if (owns_) { + mutex_->unlock_upgrade(); + } + mutex_ = other.mutex_; + owns_ = other.owns_; + other.mutex_ = nullptr; + other.owns_ = false; + } + return *this; + } + + /// @brief Release upgrade ownership if held. + ~UpgradeLock() { + if (owns_) { + mutex_->unlock_upgrade(); + } + } + + /// @brief Acquire upgrade ownership, blocking until available. + void lock() { + validate(); + mutex_->lock_upgrade(); + owns_ = true; + } + + /// @brief Attempt to acquire upgrade ownership without blocking. + /// @returns ``true`` if ownership was acquired, ``false`` otherwise. + bool try_lock() { + validate(); + owns_ = mutex_->try_lock_upgrade(); + return owns_; + } + + /// @brief Release upgrade ownership. + void unlock() { + if (!owns_) { + throw ANNEXCEPTION("Attempting to unlock an UpgradeLock that owns no lock!"); + } + mutex_->unlock_upgrade(); + owns_ = false; + } + + /// @brief Disassociate from the mutex without releasing ownership. + /// @returns A pointer to the associated mutex, or ``nullptr`` if none. + mutex_type* release() noexcept { + mutex_type* result = mutex_; + mutex_ = nullptr; + owns_ = false; + return result; + } + + /// @brief Return the associated mutex, or ``nullptr`` if none. + mutex_type* mutex() const noexcept { return mutex_; } + + /// @brief Return ``true`` if this object owns upgrade ownership of its mutex. + bool owns_lock() const noexcept { return owns_; } + + /// @brief Return ``true`` if this object owns upgrade ownership of its mutex. + explicit operator bool() const noexcept { return owns_; } + + private: + void validate() const { + if (mutex_ == nullptr) { + throw ANNEXCEPTION("UpgradeLock has no associated mutex!"); + } + if (owns_) { + throw ANNEXCEPTION("UpgradeLock already owns its lock!"); + } + } + + mutex_type* mutex_ = nullptr; + bool owns_ = false; +}; + +/// +/// @brief Scoped helper that temporarily promotes an ``UpgradeLock`` to exclusive +/// ownership. +/// +/// On construction, the associated mutex is transitioned from upgrade ownership to +/// exclusive ownership (blocking until other shared owners release). On destruction, the +/// mutex is transitioned back to upgrade ownership. The referenced ``UpgradeLock`` must own +/// upgrade ownership for the lifetime of this object. +/// +/// @tparam Mutex The mutex type. Must provide ``unlock_upgrade_and_lock`` and +/// ``unlock_and_lock_upgrade``. +/// +template class UpgradeToUniqueLock { + public: + using mutex_type = Mutex; + + /// @brief Promote ``upgrade_lock`` to exclusive ownership. + explicit UpgradeToUniqueLock(UpgradeLock& upgrade_lock) + : source_{&upgrade_lock} { + if (!source_->owns_lock()) { + throw ANNEXCEPTION( + "Cannot promote an UpgradeLock that does not own upgrade ownership!" + ); + } + source_->mutex()->unlock_upgrade_and_lock(); + } + + UpgradeToUniqueLock(const UpgradeToUniqueLock&) = delete; + UpgradeToUniqueLock& operator=(const UpgradeToUniqueLock&) = delete; + UpgradeToUniqueLock(UpgradeToUniqueLock&&) = delete; + UpgradeToUniqueLock& operator=(UpgradeToUniqueLock&&) = delete; + + /// @brief Demote back to upgrade ownership. + ~UpgradeToUniqueLock() { source_->mutex()->unlock_and_lock_upgrade(); } + + private: + UpgradeLock* source_; +}; + +} // namespace svs::lib diff --git a/include/svs/quantization/scalar/scalar.h b/include/svs/quantization/scalar/scalar.h index 1ac48ca34..a807b1d67 100644 --- a/include/svs/quantization/scalar/scalar.h +++ b/include/svs/quantization/scalar/scalar.h @@ -391,6 +391,7 @@ class SQDataset { , bias_(bias) , data_{std::move(data)} {} + const allocator_type& get_allocator() const { return data_.get_allocator(); } size_t size() const { return data_.size(); } size_t dimensions() const { return data_.dimensions(); } size_t element_size() const { return sizeof(element_type) * dimensions(); } @@ -588,6 +589,15 @@ template class DecompressionAdaptor { inner_.fix_argument(view()); } + // template + // requires (!std::is_integral_v) + // void fix_argument(std::span left) { + // // Convert the query to float + // decompressed_.resize(left.size()); + // auto decompressed_view = data::SimpleDataView(decompressed_.data(), 1, + // left.size()); decompressed_view.set_datum(0, left); inner_.fix_argument(view()); + // } + template float compute(const Right& right) const { return distance::compute(inner_, view(), right); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8c812d35a..cf0f54c24 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -92,6 +92,7 @@ set(TEST_SOURCES ${TEST_DIR}/svs/lib/version.cpp ${TEST_DIR}/svs/lib/uuid.cpp ${TEST_DIR}/svs/lib/concurrency/readwrite_protected.cpp + ${TEST_DIR}/svs/lib/concurrency/upgradable_mutex.cpp ${TEST_DIR}/svs/lib/avx_detection.cpp # Third Party ${TEST_DIR}/svs/third-party/fmt.cpp diff --git a/tests/integration/vamana/scalar_iterator.cpp b/tests/integration/vamana/scalar_iterator.cpp index 027275a64..f8b6593a8 100644 --- a/tests/integration/vamana/scalar_iterator.cpp +++ b/tests/integration/vamana/scalar_iterator.cpp @@ -15,6 +15,7 @@ */ // svs +#include "svs/extensions/vamana/scalar.h" #include "svs/index/vamana/iterator.h" #include "svs/quantization/scalar/scalar.h" @@ -273,7 +274,9 @@ void dynamic_index_with_iterator(const Distance& distance, Data data) { } } -CATCH_TEST_CASE("SQDataset Vamana Iterator", "[integration][vamana][iterator][scalar]") { +CATCH_TEST_CASE( + "SQDataset Vamana Iterator", "[integration][vamana][iterator][scalar][dynamic]" +) { auto dist = svs::distance::DistanceL2(); auto original = test_dataset::data_f32(); constexpr size_t E = 128; diff --git a/tests/svs/index/vamana/dynamic_index_2.cpp b/tests/svs/index/vamana/dynamic_index_2.cpp index e590ae547..075e06dc5 100644 --- a/tests/svs/index/vamana/dynamic_index_2.cpp +++ b/tests/svs/index/vamana/dynamic_index_2.cpp @@ -248,7 +248,7 @@ void test_loop( } } -CATCH_TEST_CASE("Testing Graph Index", "[graph_index][dynamic_index]") { +CATCH_TEST_CASE("Testing Graph Index", "[graph_index][vamana][dynamic]") { // Set hyper parameters here const size_t max_degree = 64; #if defined(NDEBUG) diff --git a/tests/svs/index/vamana/iterator.cpp b/tests/svs/index/vamana/iterator.cpp index 2b8d5e52f..954dfc470 100644 --- a/tests/svs/index/vamana/iterator.cpp +++ b/tests/svs/index/vamana/iterator.cpp @@ -243,7 +243,7 @@ struct DynamicChecker { } // namespace -CATCH_TEST_CASE("Vamana Iterator", "[index][vamana][iterator]") { +CATCH_TEST_CASE("Vamana Iterator", "[index][vamana][iterator][dynamic]") { // This tests the general behavior of the iterator for correctness. // It is not concerned with whether the returned neighbors are accurate. // diff --git a/tests/svs/index/vamana/multi.cpp b/tests/svs/index/vamana/multi.cpp index 470313f96..14ffa33fe 100644 --- a/tests/svs/index/vamana/multi.cpp +++ b/tests/svs/index/vamana/multi.cpp @@ -49,7 +49,7 @@ template float pick_alpha(Distance SVS_UNUSED(dist)) { CATCH_TEMPLATE_TEST_CASE( "Multi-vector dynamic vamana index", - "[long][index][vamana][multi]", + "[long][index][vamana][multi][dynamic]", svs::DistanceL2, svs::DistanceIP, svs::DistanceCosineSimilarity diff --git a/tests/svs/lib/concurrency/upgradable_mutex.cpp b/tests/svs/lib/concurrency/upgradable_mutex.cpp new file mode 100644 index 000000000..46ab241ec --- /dev/null +++ b/tests/svs/lib/concurrency/upgradable_mutex.cpp @@ -0,0 +1,256 @@ +/* + * Copyright 2024 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/lib/concurrency/upgradable_mutex.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include +#include +#include + +namespace { +using svs::lib::UpgradableMutex; +using svs::lib::UpgradeLock; +using svs::lib::UpgradeToUniqueLock; +} // namespace + +CATCH_TEST_CASE("Upgradable Mutex - Basic Ownership", "[core][utils][upgradable_mutex]") { + CATCH_SECTION("Exclusive try_lock is exclusive") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock()); + // No other ownership can be acquired while exclusively held. + CATCH_REQUIRE_FALSE(mutex.try_lock()); + CATCH_REQUIRE_FALSE(mutex.try_lock_shared()); + CATCH_REQUIRE_FALSE(mutex.try_lock_upgrade()); + mutex.unlock(); + // After release, ownership is available again. + CATCH_REQUIRE(mutex.try_lock()); + mutex.unlock(); + } + + CATCH_SECTION("Multiple shared owners are allowed") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock_shared()); + CATCH_REQUIRE(mutex.try_lock_shared()); + // Exclusive ownership is blocked while shared owners exist. + CATCH_REQUIRE_FALSE(mutex.try_lock()); + mutex.unlock_shared(); + mutex.unlock_shared(); + CATCH_REQUIRE(mutex.try_lock()); + mutex.unlock(); + } + + CATCH_SECTION("Upgrade coexists with shared but not with upgrade or exclusive") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock_upgrade()); + // Shared readers are still allowed alongside an upgrade owner. + CATCH_REQUIRE(mutex.try_lock_shared()); + // A second upgrade owner is not allowed. + CATCH_REQUIRE_FALSE(mutex.try_lock_upgrade()); + // Exclusive ownership is blocked. + CATCH_REQUIRE_FALSE(mutex.try_lock()); + mutex.unlock_shared(); + mutex.unlock_upgrade(); + CATCH_REQUIRE(mutex.try_lock()); + mutex.unlock(); + } +} + +CATCH_TEST_CASE("Upgradable Mutex - Transitions", "[core][utils][upgradable_mutex]") { + CATCH_SECTION("try_unlock_upgrade_and_lock succeeds when sole owner") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock_upgrade()); + CATCH_REQUIRE(mutex.try_unlock_upgrade_and_lock()); + // Now exclusively held. + CATCH_REQUIRE_FALSE(mutex.try_lock_shared()); + mutex.unlock(); + } + + CATCH_SECTION("try_unlock_upgrade_and_lock fails with concurrent shared owner") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock_upgrade()); + CATCH_REQUIRE(mutex.try_lock_shared()); + CATCH_REQUIRE_FALSE(mutex.try_unlock_upgrade_and_lock()); + mutex.unlock_shared(); + // With the shared owner gone, the upgrade can now be promoted. + CATCH_REQUIRE(mutex.try_unlock_upgrade_and_lock()); + mutex.unlock(); + } + + CATCH_SECTION("unlock_upgrade_and_lock blocks until shared owners drain") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock_upgrade()); + CATCH_REQUIRE(mutex.try_lock_shared()); + + std::atomic promoted{false}; + std::thread promoter{[&]() { + mutex.unlock_upgrade_and_lock(); + promoted.store(true); + }}; + + // The promoter must wait for the shared owner to release. + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + CATCH_REQUIRE_FALSE(promoted.load()); + + mutex.unlock_shared(); + promoter.join(); + CATCH_REQUIRE(promoted.load()); + mutex.unlock(); + } + + CATCH_SECTION("unlock_and_lock_upgrade downgrades exclusive to upgrade") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock()); + mutex.unlock_and_lock_upgrade(); + // Shared readers are now admitted again. + CATCH_REQUIRE(mutex.try_lock_shared()); + mutex.unlock_shared(); + mutex.unlock_upgrade(); + } + + CATCH_SECTION("unlock_and_lock_shared downgrades exclusive to shared") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock()); + mutex.unlock_and_lock_shared(); + CATCH_REQUIRE(mutex.try_lock_shared()); + mutex.unlock_shared(); + mutex.unlock_shared(); + } + + CATCH_SECTION("unlock_upgrade_and_lock_shared downgrades upgrade to shared") { + UpgradableMutex mutex{}; + CATCH_REQUIRE(mutex.try_lock_upgrade()); + mutex.unlock_upgrade_and_lock_shared(); + // Another upgrade owner may now enter. + CATCH_REQUIRE(mutex.try_lock_upgrade()); + mutex.unlock_upgrade(); + mutex.unlock_shared(); + } +} + +CATCH_TEST_CASE( + "Upgradable Mutex - Standard Lock Helpers", "[core][utils][upgradable_mutex]" +) { + CATCH_SECTION("std::shared_lock and std::unique_lock interoperate") { + UpgradableMutex mutex{}; + { + std::shared_lock shared{mutex}; + CATCH_REQUIRE(shared.owns_lock()); + CATCH_REQUIRE_FALSE(mutex.try_lock()); + } + { + std::unique_lock unique{mutex}; + CATCH_REQUIRE(unique.owns_lock()); + CATCH_REQUIRE_FALSE(mutex.try_lock_shared()); + } + CATCH_REQUIRE(mutex.try_lock()); + mutex.unlock(); + } +} + +CATCH_TEST_CASE("UpgradeLock helper", "[core][utils][upgradable_mutex]") { + CATCH_SECTION("Scoped acquisition and release") { + UpgradableMutex mutex{}; + { + UpgradeLock lock{mutex}; + CATCH_REQUIRE(lock.owns_lock()); + CATCH_REQUIRE(static_cast(lock)); + CATCH_REQUIRE(lock.mutex() == &mutex); + CATCH_REQUIRE_FALSE(mutex.try_lock_upgrade()); + } + CATCH_REQUIRE(mutex.try_lock_upgrade()); + mutex.unlock_upgrade(); + } + + CATCH_SECTION("Deferred and try construction") { + UpgradableMutex mutex{}; + UpgradeLock lock{mutex, std::defer_lock}; + CATCH_REQUIRE_FALSE(lock.owns_lock()); + CATCH_REQUIRE(lock.try_lock()); + CATCH_REQUIRE(lock.owns_lock()); + lock.unlock(); + CATCH_REQUIRE_FALSE(lock.owns_lock()); + lock.lock(); + CATCH_REQUIRE(lock.owns_lock()); + } + + CATCH_SECTION("Move semantics transfer ownership") { + UpgradableMutex mutex{}; + UpgradeLock lock{mutex}; + UpgradeLock moved{std::move(lock)}; + CATCH_REQUIRE(moved.owns_lock()); + CATCH_REQUIRE_FALSE(lock.owns_lock()); + CATCH_REQUIRE(lock.mutex() == nullptr); + } + + CATCH_SECTION("Unlocking an unowned lock throws") { + UpgradableMutex mutex{}; + UpgradeLock lock{mutex, std::defer_lock}; + CATCH_REQUIRE_THROWS(lock.unlock()); + } +} + +CATCH_TEST_CASE("UpgradeToUniqueLock helper", "[core][utils][upgradable_mutex]") { + CATCH_SECTION("Promotes and restores upgrade ownership") { + UpgradableMutex mutex{}; + UpgradeLock lock{mutex}; + { + UpgradeToUniqueLock unique{lock}; + // While promoted, no shared readers may enter. + CATCH_REQUIRE_FALSE(mutex.try_lock_shared()); + } + // After the scope, upgrade ownership is restored and shared readers are allowed. + CATCH_REQUIRE(mutex.try_lock_shared()); + mutex.unlock_shared(); + } + + CATCH_SECTION("Throws when the source lock owns nothing") { + UpgradableMutex mutex{}; + UpgradeLock lock{mutex, std::defer_lock}; + CATCH_REQUIRE_THROWS(UpgradeToUniqueLock{lock}); + } +} + +CATCH_TEST_CASE("Upgradable Mutex - Concurrent Stress", "[core][utils][upgradable_mutex]") { + UpgradableMutex mutex{}; + int shared_value = 0; + constexpr int num_writers = 4; + constexpr int increments_per_writer = 1000; + + std::vector writers{}; + for (int i = 0; i < num_writers; ++i) { + writers.emplace_back([&]() { + for (int j = 0; j < increments_per_writer; ++j) { + UpgradeLock lock{mutex}; + // Read phase under upgrade ownership (shared readers may coexist). + int observed = shared_value; + // Decide to write; promote to exclusive to mutate safely. + UpgradeToUniqueLock unique{lock}; + shared_value = observed + 1; + } + }); + } + for (auto& writer : writers) { + writer.join(); + } + CATCH_REQUIRE(shared_value == num_writers * increments_per_writer); +}