From cdef8686c745d3fe6b35f95d5d7dcce478a79645 Mon Sep 17 00:00:00 2001 From: Jalin Wang Date: Fri, 26 Jun 2026 16:35:24 +0800 Subject: [PATCH] fix(diskann, hnsw sparse builder): lost-wakeup race in builder progress loops (#530) --- .../algorithm/cluster/multi_chunk_cluster.cc | 36 ++++++++------- .../algorithm/cluster/multi_chunk_cluster.h | 16 +++++++ src/core/algorithm/diskann/diskann_builder.cc | 40 ++++++++++------- .../hnsw_sparse/hnsw_sparse_builder.cc | 19 ++++---- .../algorithm/diskann/diskann_builder_test.cc | 44 +++++++++++++++++++ 5 files changed, 115 insertions(+), 40 deletions(-) diff --git a/src/core/algorithm/cluster/multi_chunk_cluster.cc b/src/core/algorithm/cluster/multi_chunk_cluster.cc index 4f12256..c194796 100644 --- a/src/core/algorithm/cluster/multi_chunk_cluster.cc +++ b/src/core/algorithm/cluster/multi_chunk_cluster.cc @@ -201,16 +201,18 @@ int MultiChunkClusterAlgorithm::cluster(IndexThreads::Pointer threads, threads->count(), ¢s, &finished)); } - while (!task_group->is_finished()) { + { std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to cluster while waiting finish"); - return errcode_; + while (finished.load() < chunk_count_) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to cluster while waiting finish"); + return errcode_; + } + LOG_INFO("Finish Chunk Count %zu, Finished Percent %.3f%%", + finished.load(), finished.load() * 100.0f / chunk_count_); } - LOG_INFO("Finish Chunk Count %zu, Finished Percent %.3f%%", finished.load(), - finished.load() * 100.0f / chunk_count_); } if (error_.load(std::memory_order_acquire)) { @@ -284,16 +286,18 @@ int MultiChunkClusterAlgorithm::label(IndexThreads::Pointer threads, threads->count(), cents, out, &finished)); } - while (!task_group->is_finished()) { + { std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to cluster while waiting finish"); - return errcode_; + while (finished.load() < features_count) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to cluster while waiting finish"); + return errcode_; + } + LOG_INFO("Finish label cnt %zu, finished percent %.3f%%", finished.load(), + finished.load() * 100.0f / features_count); } - LOG_INFO("Finish label cnt %zu, finished percent %.3f%%", finished.load(), - finished.load() * 100.0f / features_count); } if (error_.load(std::memory_order_acquire)) { diff --git a/src/core/algorithm/cluster/multi_chunk_cluster.h b/src/core/algorithm/cluster/multi_chunk_cluster.h index 1c82073..19096d7 100644 --- a/src/core/algorithm/cluster/multi_chunk_cluster.h +++ b/src/core/algorithm/cluster/multi_chunk_cluster.h @@ -231,6 +231,10 @@ void MultiChunkNumericalAlgorithm::do_cluster( chunk, algorithm.centroids().count(), features_->count(), cost); (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } } return; @@ -267,6 +271,10 @@ void MultiChunkNumericalAlgorithm::do_label( } (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } } } @@ -372,6 +380,10 @@ void MultiChunkNumericalInnerProductAlgorithm::do_cluster( chunk, algorithm.centroids().count(), features_->count(), cost); (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } } } @@ -406,6 +418,10 @@ void MultiChunkNumericalInnerProductAlgorithm::do_label( } (*finished)++; + { + std::lock_guard lk(mutex_); + cond_.notify_one(); + } } } diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index 0265e63..e344844 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -267,17 +267,21 @@ int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) { threads->count(), &finished)); } - while (!task_group->is_finished()) { + { std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; + while (finished.load() < entity_.doc_cnt()) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to build index while waiting finish"); + return errcode_; + } + LOG_INFO("Built cnt %zu, finished percent %.3f%%", + (size_t)finished.load(), + finished.load() * 100.0f / entity_.doc_cnt()); } - LOG_INFO("Built cnt %zu, finished percent %.3f%%", (size_t)finished.load(), - finished.load() * 100.0f / entity_.doc_cnt()); } + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to build index while waiting finish"); return errcode_; @@ -300,17 +304,21 @@ int DiskAnnBuilder::prune_internal(IndexThreads::Pointer threads) { threads->count(), &finished)); } - while (!task_group->is_finished()) { + { std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to purne index while waiting finish"); - return errcode_; + while (finished.load() < entity_.doc_cnt()) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to prune index while waiting finish"); + return errcode_; + } + LOG_INFO("Prune cnt %zu, finished percent %.3f%%", + (size_t)finished.load(), + finished.load() * 100.0f / entity_.doc_cnt()); } - LOG_INFO("Prune cnt %zu, finished percent %.3f%%", (size_t)finished.load(), - finished.load() * 100.0f / entity_.doc_cnt()); } + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to prune index while waiting finish"); return errcode_; diff --git a/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc index 0d8d27e..4ca8f89 100644 --- a/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc +++ b/src/core/algorithm/hnsw_sparse/hnsw_sparse_builder.cc @@ -307,17 +307,20 @@ int HnswSparseBuilder::build_graph(IndexThreads::Pointer threads, i, threads->count(), &finished)); } - while (!task_group->is_finished()) { + { std::unique_lock lk(mutex_); - cond_.wait_until(lk, std::chrono::system_clock::now() + - std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; + while (finished.load() < entity_.doc_cnt()) { + cond_.wait_until(lk, std::chrono::system_clock::now() + + std::chrono::seconds(check_interval_secs_)); + if (error_.load(std::memory_order_acquire)) { + LOG_ERROR("Failed to build index while waiting finish"); + return errcode_; + } + LOG_INFO("Built cnt %u, finished percent %.3f%%", finished.load(), + finished.load() * 100.0f / entity_.doc_cnt()); } - LOG_INFO("Built cnt %u, finished percent %.3f%%", finished.load(), - finished.load() * 100.0f / entity_.doc_cnt()); } + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to build index while waiting finish"); return errcode_; diff --git a/tests/core/algorithm/diskann/diskann_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 098765c..55b89d4 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -99,6 +100,49 @@ TEST_F(DiskAnnBuilderTest, TestGeneral) { ASSERT_GT(stats.built_costtime(), 0UL); } +// Regression test: building a small DiskAnn index must complete quickly. +// A lost-wakeup bug in the condition-variable progress loops previously caused +// 15–30 second stalls during train/build on small datasets because +// notify_one() was either missing or racing against a wrong predicate. +TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) { + constexpr size_t kSmallDim = 4; + constexpr size_t kSmallDocCnt = 12; + + auto meta = make_shared(IndexMeta::DataType::DT_FP32, kSmallDim); + meta->set_metric("SquaredEuclidean", 0, Params()); + + IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(builder, nullptr); + + auto holder = make_shared>( + kSmallDim); + for (size_t i = 0; i < kSmallDocCnt; ++i) { + NumericalVector vec(kSmallDim, static_cast(i)); + ASSERT_TRUE(holder->emplace(i, vec)); + } + + Params params; + params.set("zvec.diskann.builder.max_degree", 32); + params.set("zvec.diskann.builder.list_size", 50); + params.set("zvec.diskann.builder.max_pq_chunk_num", 2); + params.set("zvec.diskann.builder.threads", 4); + + ASSERT_EQ(0, builder->init(*meta, params)); + + auto t0 = std::chrono::steady_clock::now(); + ASSERT_EQ(0, builder->train(holder)); + ASSERT_EQ(0, builder->build(holder)); + auto t1 = std::chrono::steady_clock::now(); + + auto elapsed_ms = + std::chrono::duration_cast(t1 - t0).count(); + // Before the fix, this took 15–30 seconds. After the fix, it should + // complete in well under 5 seconds even on slow CI machines. + EXPECT_LT(elapsed_ms, 5000) + << "DiskAnn build with " << kSmallDocCnt << " vectors took " << elapsed_ms + << " ms — likely a lost-wakeup regression in progress loops."; +} + // DiskAnn is now exposed implicitly: no caller ever invokes a // ``LoadDiskAnnPlugin`` / ``IsLibAioAvailable`` API (those were removed from // the public surface together with ``zvec.load_diskann_plugin()`` in Python).