From bd39e0b29bc81e8b350235b51e871c49b19d5996 Mon Sep 17 00:00:00 2001
From: "xianliang.li" <xianliang.li@zilliz.com>
Date: Tue, 13 May 2025 20:31:31 +0800
Subject: [PATCH] share the pool externally to ensure that other tasks can be
 scheduled

Signed-off-by: xianliang.li <xianliang.li@zilliz.com>
---
 cmake/libs/libcardinal.cmake        |  2 +-
 include/knowhere/comp/task.h        | 10 ++++++++++
 include/knowhere/comp/thread_pool.h |  5 +++++
 python/knowhere/knowhere.i          |  5 ++---
 src/common/thread/thread.cc         | 11 +++++++++++
 5 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/cmake/libs/libcardinal.cmake b/cmake/libs/libcardinal.cmake
index 7c980de84..8d63155c6 100644
--- a/cmake/libs/libcardinal.cmake
+++ b/cmake/libs/libcardinal.cmake
@@ -1,5 +1,5 @@
 # Use short SHA1 as version
-set(CARDINAL_VERSION v2.5.7 )
+set(CARDINAL_VERSION v2.5.8)
 set(CARDINAL_REPO_URL "https://github.com/zilliztech/cardinal.git")
 
 set(CARDINAL_REPO_DIR "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/cardinal")
diff --git a/include/knowhere/comp/task.h b/include/knowhere/comp/task.h
index 2c24342a8..676d87cee 100644
--- a/include/knowhere/comp/task.h
+++ b/include/knowhere/comp/task.h
@@ -12,6 +12,9 @@
 #include <functional>
 #include <memory>
 #include <vector>
+
+#include "folly/executors/CPUThreadPoolExecutor.h"
+
 namespace knowhere {
 
 void
@@ -26,6 +29,13 @@ size_t
 GetSearchThreadPoolSize();
 size_t
 GetBuildThreadPoolSize();
+
+folly::CPUThreadPoolExecutor&
+GetSearchThreadPool();
+
+folly::CPUThreadPoolExecutor&
+GetBuildThreadPool();
+
 class ThreadPool {
  public:
     class ScopedOmpSetter {
diff --git a/include/knowhere/comp/thread_pool.h b/include/knowhere/comp/thread_pool.h
index ee2e8b58a..78aa7d865 100644
--- a/include/knowhere/comp/thread_pool.h
+++ b/include/knowhere/comp/thread_pool.h
@@ -141,6 +141,11 @@ class ThreadPool {
         return pool_.getPendingTaskCount();
     }
 
+    folly::CPUThreadPoolExecutor&
+    GetPool() {
+        return pool_;
+    }
+
     void
     SetNumThreads(uint32_t num_threads) {
         if (num_threads == 0) {
diff --git a/python/knowhere/knowhere.i b/python/knowhere/knowhere.i
index abd113f3e..fb2c8e7a5 100644
--- a/python/knowhere/knowhere.i
+++ b/python/knowhere/knowhere.i
@@ -37,7 +37,6 @@ typedef uint64_t size_t;
 #include <knowhere/comp/brute_force.h>
 #include <knowhere/comp/knowhere_config.h>
 #include <knowhere/comp/local_file_manager.h>
-#include <knowhere/comp/task.h>
 #include <knowhere/comp/index_param.h>
 #include <fstream>
 #include <string>
@@ -691,12 +690,12 @@ SetSimdType(const std::string type) {
 
 void
 SetBuildThreadPool(uint32_t num_threads) {
-    knowhere::InitBuildThreadPool(num_threads);
+    knowhere::KnowhereConfig::SetBuildThreadPoolSize(num_threads);
 }
 
 void
 SetSearchThreadPool(uint32_t num_threads) {
-    knowhere::InitSearchThreadPool(num_threads);
+    knowhere::KnowhereConfig::SetSearchThreadPoolSize(num_threads);
 }
 
 %}
diff --git a/src/common/thread/thread.cc b/src/common/thread/thread.cc
index fa1709d96..39a93e888 100644
--- a/src/common/thread/thread.cc
+++ b/src/common/thread/thread.cc
@@ -18,6 +18,7 @@
 #include <thread>
 #include <utility>
 
+#include "folly/executors/CPUThreadPoolExecutor.h"
 #include "knowhere/comp/thread_pool.h"
 
 namespace knowhere {
@@ -57,6 +58,16 @@ InitBuildThreadPool(uint32_t num_threads) {
     ThreadPool::InitGlobalBuildThreadPool(num_threads);
 }
 
+folly::CPUThreadPoolExecutor&
+GetBuildThreadPool() {
+    return ThreadPool::GetGlobalBuildThreadPool()->GetPool();
+}
+
+folly::CPUThreadPoolExecutor&
+GetSearchThreadPool() {
+    return ThreadPool::GetGlobalSearchThreadPool()->GetPool();
+}
+
 void
 InitSearchThreadPool(uint32_t num_threads) {
     ThreadPool::InitGlobalSearchThreadPool(num_threads);