Fix test crash in some circumstances

There appears to be some sort of race condition accessing thread vectors from the thread.
pabs3 · Jun 16, 2023 · 9633d32 · 9633d32
1 parent 866e60c
commit 9633d32
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/AnnService/inc/SSDServing/SSDIndex.h b/AnnService/inc/SSDServing/SSDIndex.h
@@ -111,15 +111,14 @@ namespace SPTAG {
 
                 std::vector<std::thread> threads;
 
+                NumaStrategy ns = (p_index->GetDiskIndex() != nullptr) ? NumaStrategy::SCATTER : NumaStrategy::LOCAL; // Only for SPANN, we need to avoid IO threads overlap with search threads.
+
                 SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Searching: numThread: %d, numQueries: %d.\n", p_numThreads, numQueries);
 
                 Utils::StopW sw;
 
                 for (int i = 0; i < p_numThreads; i++) { threads.emplace_back([&, i]()
                     {
-                        NumaStrategy ns = (p_index->GetDiskIndex() != nullptr) ? NumaStrategy::SCATTER : NumaStrategy::LOCAL; // Only for SPANN, we need to avoid IO threads overlap with search threads.
-                        Helper::SetThreadAffinity(i, threads[i], ns, OrderStrategy::ASC); 
-
                         Utils::StopW threadws;
                         size_t index = 0;
                         while (true)
@@ -147,6 +146,7 @@ namespace SPTAG {
                             }
                         }
                     });
+                    Helper::SetThreadAffinity(i, threads.back(), ns, OrderStrategy::ASC);
                 }
                 for (auto& thread : threads) { thread.join(); }
 

diff --git a/AnnService/src/IndexSearcher/main.cpp b/AnnService/src/IndexSearcher/main.cpp
@@ -189,13 +189,12 @@ int Process(std::shared_ptr<SearcherOptions> options, VectorIndex& index)
             std::atomic_size_t queriesSent(0);
             std::vector<std::thread> threads;
 
+            NumaStrategy ns = (index.GetIndexAlgoType() == IndexAlgoType::SPANN)? NumaStrategy::SCATTER: NumaStrategy::LOCAL; // Only for SPANN, we need to avoid IO threads overlap with search threads.
+
             auto batchstart = std::chrono::high_resolution_clock::now();
 
             for (std::uint32_t i = 0; i < options->m_threadNum; i++) { 
                 threads.emplace_back([&, i] {
-                    NumaStrategy ns = (index.GetIndexAlgoType() == IndexAlgoType::SPANN)? NumaStrategy::SCATTER: NumaStrategy::LOCAL; // Only for SPANN, we need to avoid IO threads overlap with search threads.
-                    Helper::SetThreadAffinity(i, threads[i], ns, OrderStrategy::ASC);
-
                     size_t qid = 0;
                     while (true)
                     {
@@ -213,6 +212,7 @@ int Process(std::shared_ptr<SearcherOptions> options, VectorIndex& index)
                         }
                     }
                 });
+                Helper::SetThreadAffinity(i, threads.back(), ns, OrderStrategy::ASC);
             }
             for (auto& thread : threads) { thread.join(); }