From f5d47ccd90d6776252b4ddc5b2ea36eccaa15430 Mon Sep 17 00:00:00 2001
From: robot-clickhouse <robot-clickhouse@users.noreply.github.com>
Date: Wed, 18 Dec 2024 17:08:49 +0000
Subject: [PATCH] Backport #73517 to 24.10: Fix race in
 `MergeTreeIndexVectorSimilarity`

---
 .../MergeTreeIndexVectorSimilarity.cpp          | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/src/Storages/MergeTree/MergeTreeIndexVectorSimilarity.cpp b/src/Storages/MergeTree/MergeTreeIndexVectorSimilarity.cpp
index 5a725922e14..4a9982784e2 100644
--- a/src/Storages/MergeTree/MergeTreeIndexVectorSimilarity.cpp
+++ b/src/Storages/MergeTree/MergeTreeIndexVectorSimilarity.cpp
@@ -7,6 +7,7 @@
 #include <Common/formatReadable.h>
 #include <Common/getNumberOfCPUCoresToUse.h>
 #include <Common/logger_useful.h>
+#include <Common/threadPoolCallbackRunner.h>
 #include <Common/typeid_cast.h>
 #include <Core/Field.h>
 #include <Core/ServerSettings.h>
@@ -292,16 +293,9 @@ void updateImpl(const ColumnArray * column_array, const ColumnArray::Offsets & c
     /// indexes are build simultaneously (e.g. multiple merges run at the same time).
     auto & thread_pool = Context::getGlobalContextInstance()->getBuildVectorSimilarityIndexThreadPool();
 
-    auto add_vector_to_index = [&](USearchIndex::vector_key_t key, size_t row, ThreadGroupPtr thread_group)
+    ThreadPoolCallbackRunnerLocal<void> runner(thread_pool, "VectorSimIndex");
+    auto add_vector_to_index = [&](USearchIndex::vector_key_t key, size_t row)
     {
-        SCOPE_EXIT_SAFE(
-            if (thread_group)
-                CurrentThread::detachFromGroupIfNotDetached();
-        );
-
-        if (thread_group)
-            CurrentThread::attachToGroupIfDetached(thread_group);
-
         /// add is thread-safe
         auto result = index->add(key, &column_array_data_float_data[column_array_offsets[row - 1]]);
         if (!result)
@@ -319,11 +313,10 @@ void updateImpl(const ColumnArray * column_array, const ColumnArray::Offsets & c
     for (size_t row = 0; row < rows; ++row)
     {
         auto key = static_cast<USearchIndex::vector_key_t>(index_size + row);
-        auto task = [group = CurrentThread::getGroup(), &add_vector_to_index, key, row] { add_vector_to_index(key, row, group); };
-        thread_pool.scheduleOrThrowOnError(task);
+        runner([&add_vector_to_index, key, row] { add_vector_to_index(key, row); });
     }
 
-    thread_pool.wait();
+    runner.waitForAllToFinishAndRethrowFirstError();
 }
 
 }