diff --git a/python/tests/detail/test_db_config.py b/python/tests/detail/test_db_config.py index 16203a8..7125843 100644 --- a/python/tests/detail/test_db_config.py +++ b/python/tests/detail/test_db_config.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import inspect import pytest import tempfile import os @@ -149,6 +150,16 @@ class TestDbConfigMemoryLimitValidation: class TestDbConfigThreadValidation: + def test_thread_binding_options_are_not_exposed(self): + parameters = inspect.signature(zvec.init).parameters + assert "query_thread_binding" not in parameters + assert "optimize_thread_binding" not in parameters + + with pytest.raises(TypeError): + zvec.init(query_thread_binding=True) + with pytest.raises(TypeError): + zvec.init(optimize_thread_binding=True) + @run_in_subprocess def test_query_threads(self): zvec.init(query_threads=1) diff --git a/python/zvec/zvec.py b/python/zvec/zvec.py index 0046a49..bd2bd32 100644 --- a/python/zvec/zvec.py +++ b/python/zvec/zvec.py @@ -81,7 +81,8 @@ def init( Must be ≥ 1 if provided. optimize_threads (Optional[int], optional): Threads for background tasks (e.g., compaction, indexing). - If ``None``, defaults to same as ``query_threads`` or CPU count. + If ``None`` (default), uses the same environment-aware default as + ``query_threads``. invert_to_forward_scan_ratio (Optional[float], optional): Threshold to switch from inverted index to full forward scan. Range: [0.0, 1.0]. Higher → more aggressive index skipping. diff --git a/src/ailego/parallel/thread_pool.cc b/src/ailego/parallel/thread_pool.cc index 54a8970..09da09e 100644 --- a/src/ailego/parallel/thread_pool.cc +++ b/src/ailego/parallel/thread_pool.cc @@ -12,33 +12,64 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include -#if (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__) +#if defined(__linux__) && !defined(__ANDROID__) #include -static inline void BindThreads(std::vector &pool) { - uint32_t hc = std::thread::hardware_concurrency(); - if (hc > 1) { - cpu_set_t mask; +static inline bool GetAllowedCpuMask(cpu_set_t *mask) { + CPU_ZERO(mask); + const int error = pthread_getaffinity_np(pthread_self(), sizeof(*mask), mask); + if (error != 0) { + LOG_WARN("Failed to get thread affinity mask, error[%d]", error); + return false; + } + return true; +} - for (size_t i = 0u; i < pool.size(); ++i) { - CPU_ZERO(&mask); - CPU_SET(i % hc, &mask); - pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); +static inline void BindThreads(std::vector &pool) { + cpu_set_t allowed_mask; + if (!GetAllowedCpuMask(&allowed_mask)) { + return; + } + + std::vector allowed_cpus; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &allowed_mask)) { + allowed_cpus.push_back(cpu); + } + } + if (allowed_cpus.empty()) { + LOG_WARN("Cannot bind thread pool: affinity mask has no allowed CPUs"); + return; + } + + cpu_set_t target_mask; + for (size_t i = 0u; i < pool.size(); ++i) { + CPU_ZERO(&target_mask); + CPU_SET(allowed_cpus[i % allowed_cpus.size()], &target_mask); + const int error = pthread_setaffinity_np(pool[i].native_handle(), + sizeof(target_mask), &target_mask); + if (error != 0) { + LOG_WARN("Failed to bind thread pool worker[%zu], error[%d]", i, error); } } } static inline void UnbindThreads(std::vector &pool) { - cpu_set_t mask; - CPU_ZERO(&mask); - - for (size_t i = 0u; i < CPU_SETSIZE; ++i) { - CPU_SET(i, &mask); + cpu_set_t allowed_mask; + if (!GetAllowedCpuMask(&allowed_mask)) { + return; } + for (size_t i = 0u; i < pool.size(); ++i) { - pthread_setaffinity_np(pool[i].native_handle(), sizeof(mask), &mask); + const int error = pthread_setaffinity_np( + pool[i].native_handle(), sizeof(allowed_mask), &allowed_mask); + if (error != 0) { + LOG_WARN("Failed to unbind thread pool worker[%zu], error[%d]", i, error); + } } } #else @@ -49,12 +80,44 @@ static inline void UnbindThreads(std::vector &) {} namespace zvec { namespace ailego { +namespace { + +uint32_t SafeWorkerCount() noexcept { + return std::max(std::thread::hardware_concurrency(), 1u); +} + +} // namespace + +ThreadPool::ThreadPool() : ThreadPool(SafeWorkerCount(), false) {} + ThreadPool::ThreadPool(uint32_t size, bool binding) { - for (uint32_t i = 0u; i < size; ++i) { - pool_.emplace_back(&ThreadPool::worker, this); + const uint32_t max_size = SafeWorkerCount(); + const uint32_t safe_size = std::min(std::max(size, 1u), max_size); + if (safe_size != size) { + LOG_WARN( + "ThreadPool worker count[%u] is outside supported range[1, %u], " + "clamped to %u", + size, max_size, safe_size); } - if (binding) { - this->bind(); + + try { + // Avoid vector reallocations after workers have started. If thread creation + // still fails, stop and join the workers already created before rethrowing. + pool_.reserve(safe_size); + for (uint32_t i = 0u; i < safe_size; ++i) { + pool_.emplace_back(&ThreadPool::worker, this); + } + if (binding) { + this->bind(); + } + } catch (...) { + this->stop(); + for (auto &worker : pool_) { + if (worker.joinable()) { + worker.join(); + } + } + throw; } } @@ -130,4 +193,4 @@ bool ThreadPool::picking(ThreadPool::Task *task) { } } // namespace ailego -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index d6ad1f4..a9305da 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -242,4 +242,4 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 56bf9f7..73dd795 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include @@ -1013,19 +1014,24 @@ int Index::Merge(const std::vector &indexes, } // must declare here to ensure its lifespan can cover reducer->reduce() std::unique_ptr local_thread_pool = nullptr; + uint32_t effective_write_concurrency = options.write_concurrency; if (options.pool != nullptr) { reducer->set_thread_pool(options.pool); + effective_write_concurrency = + std::min(effective_write_concurrency, options.pool->count()); } else { local_thread_pool = - std::make_unique(options.write_concurrency); + std::make_unique(options.write_concurrency, false); reducer->set_thread_pool(local_thread_pool.get()); + effective_write_concurrency = + static_cast(local_thread_pool->count()); } ailego::Params reducer_params; reducer_params.set(core::PARAM_MIXED_STREAMER_REDUCER_ENABLE_PK_REWRITE, true); reducer_params.set(core::PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, - options.write_concurrency); + effective_write_concurrency); if (reducer->init(reducer_params) != 0) { LOG_ERROR("Failed to init reducer"); return core::IndexError_Runtime; diff --git a/src/db/common/cgroup_util.cc b/src/db/common/cgroup_util.cc index 0e137d7..fbc78e8 100644 --- a/src/db/common/cgroup_util.cc +++ b/src/db/common/cgroup_util.cc @@ -13,6 +13,9 @@ // limitations under the License. #include "db/common/cgroup_util.h" +#include +#include +#include namespace zvec { @@ -114,19 +117,15 @@ bool CgroupUtil::readCpuCgroup() { // cgroup v2 std::ifstream file("/sys/fs/cgroup/cpu.max"); if (file.is_open()) { - uint64_t quota, period; - char slash; - file >> quota >> slash >> period; - file.close(); + std::string cpu_max; + std::getline(file, cpu_max); - if (quota != std::numeric_limits::max() && quota != 0 && - period > 0) { - cpu_cores_ = - static_cast(std::ceil(static_cast(quota) / period)); + int cpu_cores = 0; + if (parseCpuMax(cpu_max, &cpu_cores)) { + cpu_cores_ = cpu_cores; return true; - } else { - return false; } + return false; } // cgroup v1 @@ -150,6 +149,43 @@ bool CgroupUtil::readCpuCgroup() { return false; } +bool CgroupUtil::parseCpuMax(const std::string &cpu_max, int *cpu_cores) { + if (cpu_cores == nullptr) { + return false; + } + + std::istringstream stream(cpu_max); + std::string quota_token; + uint64_t period = 0; + if (!(stream >> quota_token >> period) || period == 0) { + return false; + } + + std::string trailing_token; + if (stream >> trailing_token) { + return false; + } + if (quota_token == "max") { + return false; + } + + uint64_t quota = 0; + const auto result = std::from_chars( + quota_token.data(), quota_token.data() + quota_token.size(), quota); + if (result.ec != std::errc{} || + result.ptr != quota_token.data() + quota_token.size() || quota == 0) { + return false; + } + + const uint64_t cores = quota / period + (quota % period != 0 ? 1 : 0); + if (cores > static_cast((std::numeric_limits::max)())) { + return false; + } + + *cpu_cores = static_cast(cores); + return true; +} + void CgroupUtil::updateMemoryLimit() { if (readMemoryCgroup()) { return; @@ -491,4 +527,4 @@ double CgroupUtil::calculateMacOSCpuUsage() { } #endif -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/common/cgroup_util.h b/src/db/common/cgroup_util.h index 759e231..997c6bd 100644 --- a/src/db/common/cgroup_util.h +++ b/src/db/common/cgroup_util.h @@ -17,7 +17,6 @@ #include #include #include -#include #include #include #include @@ -45,6 +44,10 @@ class CgroupUtil { static int getCpuLimit(); static uint64_t getMemoryLimit(); + // Parse a cgroup v2 cpu.max value. Returns false for an unlimited or + // malformed value so callers can fall back to the host CPU count. + static bool parseCpuMax(const std::string &cpu_max, int *cpu_cores); + // Static methods to get other resources static double getCpuUsage(); static uint64_t getMemoryUsage(); @@ -101,4 +104,4 @@ class CgroupUtil { #endif }; -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/common/config.cc b/src/db/common/config.cc index 57eaae8..14b53a6 100644 --- a/src/db/common/config.cc +++ b/src/db/common/config.cc @@ -35,10 +35,12 @@ GlobalConfig::ConfigData::ConfigData() DEFAULT_MEMORY_LIMIT_RATIO), log_config(std::make_shared()), query_thread_count(CgroupUtil::getCpuLimit()), + query_thread_binding(false), invert_to_forward_scan_ratio(0.9), brute_force_by_keys_ratio(0.1), fts_brute_force_by_keys_ratio(0.05), - optimize_thread_count(CgroupUtil::getCpuLimit()), + optimize_thread_count(query_thread_count), + optimize_thread_binding(false), jieba_dict_dir() {} Status GlobalConfig::Validate(const ConfigData &config) const { @@ -165,4 +167,4 @@ uint64_t GlobalConfig::memory_limit_bytes() const noexcept { FACTORY_REGISTER_LOGGER(AppendLogger); -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/common/global_resource.cc b/src/db/common/global_resource.cc index 711bf3d..62d1966 100644 --- a/src/db/common/global_resource.cc +++ b/src/db/common/global_resource.cc @@ -20,11 +20,13 @@ namespace zvec { void GlobalResource::initialize() { static std::once_flag flag; - std::call_once(flag, [this]() mutable { - this->query_thread_pool_.reset( - new ailego::ThreadPool(GlobalConfig::Instance().query_thread_count())); + std::call_once(flag, [this]() { + this->query_thread_pool_.reset(new ailego::ThreadPool( + GlobalConfig::Instance().query_thread_count(), + GlobalConfig::Instance().query_thread_binding())); this->optimize_thread_pool_.reset(new ailego::ThreadPool( - GlobalConfig::Instance().optimize_thread_count())); + GlobalConfig::Instance().optimize_thread_count(), + GlobalConfig::Instance().optimize_thread_binding())); zvec::ailego::MemoryLimitPool::get_instance().init( GlobalConfig::Instance().memory_limit_bytes()); }); diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 984d736..36a3b5e 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -1661,7 +1661,7 @@ Result SegmentImpl::merge_vector_indexer( if (concurrency == 0) { merge_options.pool = GlobalResource::Instance().optimize_thread_pool(); merge_options.write_concurrency = - GlobalConfig::Instance().optimize_thread_count(); + static_cast(merge_options.pool->count()); } else { merge_options.write_concurrency = concurrency; } diff --git a/src/db/index/segment/segment_helper.cc b/src/db/index/segment/segment_helper.cc index 1238a48..f1d5881 100644 --- a/src/db/index/segment/segment_helper.cc +++ b/src/db/index/segment/segment_helper.cc @@ -797,7 +797,7 @@ Status SegmentHelper::MergeWithOptionalReuse( if (concurrency == 0) { merge_options.pool = GlobalResource::Instance().optimize_thread_pool(); merge_options.write_concurrency = - GlobalConfig::Instance().optimize_thread_count(); + static_cast(merge_options.pool->count()); } else { merge_options.write_concurrency = concurrency; } diff --git a/src/include/zvec/ailego/parallel/thread_pool.h b/src/include/zvec/ailego/parallel/thread_pool.h index a7b63ed..1b122b6 100644 --- a/src/include/zvec/ailego/parallel/thread_pool.h +++ b/src/include/zvec/ailego/parallel/thread_pool.h @@ -134,16 +134,23 @@ class ZVEC_AILEGO_API ThreadPool { std::condition_variable cond_{}; }; - //! Constructor + /** + * Create a thread pool. + * + * The requested worker count is clamped to + * [1, std::thread::hardware_concurrency()] (with a hardware fallback of 1). + * This guarantees that the pool can always execute queued work while + * preventing excessive worker creation when callers provide an invalid or + * unreasonably large size. Requests to oversubscribe the available hardware + * are therefore reduced. + * + * @param size Requested worker count. + * @param binding Whether to bind workers to allowed CPU cores. + */ explicit ThreadPool(uint32_t size, bool binding); //! Constructor - explicit ThreadPool(bool binding) - : ThreadPool{std::max(std::thread::hardware_concurrency(), 1u), binding} { - } - - //! Constructor - ThreadPool(void) : ThreadPool{false} {} + ThreadPool(void); //! Destructor ~ThreadPool(void) { @@ -164,9 +171,9 @@ class ZVEC_AILEGO_API ThreadPool { //! Stop all threads void stop(void) { - // Set stop flag as ture, then wake all threads - stopping_ = true; + // Set the stop flag while holding the same lock used by workers. std::lock_guard lock(queue_mutex_); + stopping_ = true; work_cond_.notify_all(); } diff --git a/src/include/zvec/core/framework/index_threads.h b/src/include/zvec/core/framework/index_threads.h index 38e85b7..c0986cf 100644 --- a/src/include/zvec/core/framework/index_threads.h +++ b/src/include/zvec/core/framework/index_threads.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include diff --git a/src/include/zvec/db/config.h b/src/include/zvec/db/config.h index 4bc4f5c..ff17570 100644 --- a/src/include/zvec/db/config.h +++ b/src/include/zvec/db/config.h @@ -93,6 +93,8 @@ class ZVEC_API GlobalConfig : public ailego::Singleton { // query uint32_t query_thread_count; + // CPU binding is opt-in at the DB layer. + bool query_thread_binding; float invert_to_forward_scan_ratio; float brute_force_by_keys_ratio; // Independent from brute_force_by_keys_ratio: per-candidate FTS cost @@ -101,6 +103,8 @@ class ZVEC_API GlobalConfig : public ailego::Singleton { // optimize uint32_t optimize_thread_count; + // CPU binding is opt-in at the DB layer. + bool optimize_thread_binding; // FTS jieba tokenizer default dict dir (lowest-priority fallback; // per-field config > ZVEC_JIEBA_DICT_DIR > this). Empty by default. @@ -166,6 +170,11 @@ class ZVEC_API GlobalConfig : public ailego::Singleton { return config_.query_thread_count; } + //! Query thread binding + bool query_thread_binding() const noexcept { + return config_.query_thread_binding; + } + //! Invert to forward scan ratio float invert_to_forward_scan_ratio() const noexcept { return config_.invert_to_forward_scan_ratio; @@ -187,6 +196,11 @@ class ZVEC_API GlobalConfig : public ailego::Singleton { return config_.optimize_thread_count; } + //! Optimize thread binding + bool optimize_thread_binding() const noexcept { + return config_.optimize_thread_binding; + } + //! Effective jieba dict dir. Thread-safe. std::string jieba_dict_dir() const; diff --git a/tests/ailego/parallel/thread_pool_test.cc b/tests/ailego/parallel/thread_pool_test.cc index 70459f5..fb196f9 100644 --- a/tests/ailego/parallel/thread_pool_test.cc +++ b/tests/ailego/parallel/thread_pool_test.cc @@ -12,14 +12,98 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include +#include #include +#include +#include #include #include +#if defined(__linux__) && !defined(__ANDROID__) +#include +#endif + using namespace zvec::ailego; +static_assert(!std::is_constructible::value, + "ThreadPool thread count must not be ambiguous with binding"); +static_assert(!std::is_constructible::value, + "ThreadPool binding must require an explicit thread count"); +static_assert(std::is_constructible::value, + "ThreadPool must accept an explicit count and binding"); + +TEST(ThreadPool, ConstructorBehavior) { + const auto hardware_concurrency = + std::max(std::thread::hardware_concurrency(), 1u); + + ThreadPool default_pool; + EXPECT_EQ(hardware_concurrency, default_pool.count()); + + ThreadPool single_unbound_pool(1, false); + EXPECT_EQ(1u, single_unbound_pool.count()); + + const uint32_t bound_count = std::min(hardware_concurrency, 2u); + ThreadPool bound_pool(bound_count, true); + EXPECT_EQ(bound_count, bound_pool.count()); +} + +TEST(ThreadPool, CapsExcessiveWorkerCount) { + const auto hardware_concurrency = + std::max(std::thread::hardware_concurrency(), 1u); + ThreadPool pool((std::numeric_limits::max)(), false); + EXPECT_EQ(hardware_concurrency, pool.count()); +} + +TEST(ThreadPool, EnsuresAtLeastOneWorker) { + ThreadPool pool(0, false); + EXPECT_EQ(1u, pool.count()); + + bool executed = false; + pool.execute_and_wait([&executed]() { executed = true; }); + EXPECT_TRUE(executed); +} + +#if defined(__linux__) && !defined(__ANDROID__) +TEST(ThreadPool, BindingRespectsCallerAffinityMask) { + cpu_set_t caller_mask; + CPU_ZERO(&caller_mask); + ASSERT_EQ(0, pthread_getaffinity_np(pthread_self(), sizeof(caller_mask), + &caller_mask)); + ASSERT_GT(CPU_COUNT(&caller_mask), 0); + + ThreadPool pool(1, true); + cpu_set_t bound_mask; + CPU_ZERO(&bound_mask); + int get_affinity_error = -1; + pool.execute_and_wait([&]() { + get_affinity_error = + pthread_getaffinity_np(pthread_self(), sizeof(bound_mask), &bound_mask); + }); + + ASSERT_EQ(0, get_affinity_error); + ASSERT_EQ(1, CPU_COUNT(&bound_mask)); + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &bound_mask)) { + EXPECT_TRUE(CPU_ISSET(cpu, &caller_mask)); + } + } + + pool.unbind(); + cpu_set_t unbound_mask; + CPU_ZERO(&unbound_mask); + pool.execute_and_wait([&]() { + get_affinity_error = pthread_getaffinity_np( + pthread_self(), sizeof(unbound_mask), &unbound_mask); + }); + + ASSERT_EQ(0, get_affinity_error); + EXPECT_TRUE(CPU_EQUAL(&caller_mask, &unbound_mask)); +} +#endif + struct A { A(void) : pool(std::make_shared()) {} @@ -38,7 +122,9 @@ struct A { }; struct B { - B(void) : pool(std::make_shared(true)) {} + B(void) + : pool(std::make_shared( + std::max(std::thread::hardware_concurrency(), 1u), true)) {} std::string ThreadMain(uint32_t &num) { aaa.pool->enqueue( diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 4453756..26f199c 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -748,7 +749,11 @@ TEST(IndexInterface, Merge) { { // test reduce auto index3 = create_index_func(param_target, index_name + "3"); ASSERT_NE(nullptr, index3); - ASSERT_TRUE(0 == index3->Merge({index1, index2}, IndexFilter())); + MergeOptions merge_options; + merge_options.write_concurrency = + (std::numeric_limits::max)(); + ASSERT_TRUE(0 == index3->Merge({index1, index2}, IndexFilter(), + merge_options)); ASSERT_TRUE(3 == index3->GetDocCount()); { VectorDataBuffer fetched_vector_data; @@ -777,7 +782,13 @@ TEST(IndexInterface, Merge) { ASSERT_NE(nullptr, index3); auto filter = IndexFilter(); filter.set([](uint64_t key) { return key == 0; }); // TODO: uint32? - ASSERT_TRUE(0 == index3->Merge({index1, index2}, filter)); + zvec::ailego::ThreadPool pool(1, false); + MergeOptions merge_options; + merge_options.write_concurrency = + (std::numeric_limits::max)(); + merge_options.pool = &pool; + ASSERT_TRUE(0 == + index3->Merge({index1, index2}, filter, merge_options)); ASSERT_TRUE(2 == index3->GetDocCount()); { VectorDataBuffer fetched_vector_data; diff --git a/tests/db/common/cgroup_util_test.cc b/tests/db/common/cgroup_util_test.cc new file mode 100644 index 0000000..e011f8c --- /dev/null +++ b/tests/db/common/cgroup_util_test.cc @@ -0,0 +1,45 @@ +// Copyright 2025-present the zvec project +// +// 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. + +#include "db/common/cgroup_util.h" +#include + +using namespace zvec; + +TEST(CgroupUtil, ParseCpuMaxQuota) { + int cpu_cores = 0; + + ASSERT_TRUE(CgroupUtil::parseCpuMax("100000 100000", &cpu_cores)); + EXPECT_EQ(1, cpu_cores); + + ASSERT_TRUE(CgroupUtil::parseCpuMax("150000 100000\n", &cpu_cores)); + EXPECT_EQ(2, cpu_cores); + + ASSERT_TRUE(CgroupUtil::parseCpuMax("400000 100000", &cpu_cores)); + EXPECT_EQ(4, cpu_cores); +} + +TEST(CgroupUtil, ParseCpuMaxUnlimitedOrInvalid) { + int cpu_cores = 7; + + EXPECT_FALSE(CgroupUtil::parseCpuMax("max 100000", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000/100000", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000 0", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("0 100000", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000 100000 trailing", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("invalid", &cpu_cores)); + EXPECT_FALSE(CgroupUtil::parseCpuMax("100000 100000", nullptr)); + + EXPECT_EQ(7, cpu_cores); +} diff --git a/tests/db/common/config_test.cc b/tests/db/common/config_test.cc index d86e160..be5f6ed 100644 --- a/tests/db/common/config_test.cc +++ b/tests/db/common/config_test.cc @@ -29,6 +29,15 @@ class ConfigTest : public ::testing::Test { } }; +TEST_F(ConfigTest, ThreadConfigDataDefaults) { + GlobalConfig::ConfigData config; + + ASSERT_GT(config.query_thread_count, 0u); + ASSERT_EQ(config.query_thread_count, config.optimize_thread_count); + ASSERT_FALSE(config.query_thread_binding); + ASSERT_FALSE(config.optimize_thread_binding); +} + TEST_F(ConfigTest, InitializeWithDefaultConfig) { GlobalConfig::ConfigData config; @@ -42,10 +51,12 @@ TEST_F(ConfigTest, InitializeWithDefaultConfig) { GlobalConfig::LogLevel::kWarn); ASSERT_EQ(GlobalConfig::Instance().log_type(), "ConsoleLogger"); ASSERT_GT(GlobalConfig::Instance().query_thread_count(), 0); + ASSERT_FALSE(GlobalConfig::Instance().query_thread_binding()); ASSERT_EQ(GlobalConfig::Instance().invert_to_forward_scan_ratio(), 0.9f); ASSERT_EQ(GlobalConfig::Instance().brute_force_by_keys_ratio(), 0.1f); ASSERT_EQ(GlobalConfig::Instance().fts_brute_force_by_keys_ratio(), 0.05f); ASSERT_GT(GlobalConfig::Instance().optimize_thread_count(), 0); + ASSERT_FALSE(GlobalConfig::Instance().optimize_thread_binding()); } TEST_F(ConfigTest, InitializeWithCustomConsoleLogConfig) { @@ -243,4 +254,4 @@ TEST_F(ConfigTest, JiebaDictDirSetterIsIndependentOfInitialize) { ASSERT_EQ(GlobalConfig::Instance().jieba_dict_dir(), ""); GlobalConfig::Instance().set_default_jieba_dict_dir(saved); -} \ No newline at end of file +} diff --git a/tests/db/common/global_resource_test.cc b/tests/db/common/global_resource_test.cc new file mode 100644 index 0000000..921c1f4 --- /dev/null +++ b/tests/db/common/global_resource_test.cc @@ -0,0 +1,46 @@ +// Copyright 2025-present the zvec project +// +// 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. + +#include "db/common/global_resource.h" +#include +#include +#include +#include + +using namespace zvec; + +TEST(GlobalResource, UsesEffectiveCountsAndDisablesBindingByDefault) { + GlobalConfig::ConfigData config; + config.query_thread_count = 2; + config.optimize_thread_count = 1; + EXPECT_FALSE(config.query_thread_binding); + EXPECT_FALSE(config.optimize_thread_binding); + + const auto status = GlobalConfig::Instance().Initialize(config); + ASSERT_TRUE(status.ok()) << status.message(); + EXPECT_FALSE(GlobalConfig::Instance().query_thread_binding()); + EXPECT_FALSE(GlobalConfig::Instance().optimize_thread_binding()); + + const auto max_workers = + std::max(std::thread::hardware_concurrency(), 1u); + const auto expected_query_workers = std::min( + GlobalConfig::Instance().query_thread_count(), max_workers); + const auto expected_optimize_workers = std::min( + GlobalConfig::Instance().optimize_thread_count(), max_workers); + + EXPECT_EQ(expected_query_workers, + GlobalResource::Instance().query_thread_pool()->count()); + EXPECT_EQ(expected_optimize_workers, + GlobalResource::Instance().optimize_thread_pool()->count()); +} diff --git a/tools/core/bench.cc b/tools/core/bench.cc index 2587dfe..2d35881 100644 --- a/tools/core/bench.cc +++ b/tools/core/bench.cc @@ -35,12 +35,13 @@ class Bench { retrieval_mode_{retrieval_mode}, filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { @@ -339,12 +340,13 @@ class SparseBench { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { diff --git a/tools/core/bench_original.cc b/tools/core/bench_original.cc index c787b88..bee3c1e 100644 --- a/tools/core/bench_original.cc +++ b/tools/core/bench_original.cc @@ -69,12 +69,13 @@ class Bench { retrieval_mode_{retrieval_mode}, filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { @@ -380,12 +381,13 @@ class SparseBench { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(false); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { diff --git a/tools/core/local_builder.cc b/tools/core/local_builder.cc index f8ea9ca..6ae7697 100644 --- a/tools/core/local_builder.cc +++ b/tools/core/local_builder.cc @@ -290,6 +290,7 @@ int do_build_sparse_by_streamer(IndexStreamer::Pointer &streamer, uint32_t thread_count) { int ret; ailego::ThreadPool pool(thread_count, false); + thread_count = static_cast(pool.count()); std::atomic finished{0}; int errcode = 0; std::mutex mutex; @@ -478,6 +479,7 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, const IndexStorage::Pointer &storage = nullptr) { int ret; ailego::ThreadPool pool(thread_count, false); + thread_count = static_cast(pool.count()); std::atomic finished{0}; int errcode = 0; std::mutex mutex; diff --git a/tools/core/local_builder_original.cc b/tools/core/local_builder_original.cc index 77c8513..7a6111d 100644 --- a/tools/core/local_builder_original.cc +++ b/tools/core/local_builder_original.cc @@ -283,6 +283,7 @@ int do_build_sparse_by_streamer(IndexStreamer::Pointer &streamer, uint32_t thread_count) { int ret; ailego::ThreadPool pool(thread_count, false); + thread_count = static_cast(pool.count()); std::atomic finished{0}; int errcode = 0; std::mutex mutex; @@ -437,6 +438,7 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, uint32_t thread_count, RetrievalMode retrieval_mode) { int ret; ailego::ThreadPool pool(thread_count, false); + thread_count = static_cast(pool.count()); std::atomic finished{0}; int errcode = 0; std::mutex mutex; diff --git a/tools/core/recall.cc b/tools/core/recall.cc index d5be191..4980eee 100644 --- a/tools/core/recall.cc +++ b/tools/core/recall.cc @@ -36,12 +36,13 @@ class Recall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { - pool_ = make_shared(threads_, true); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { @@ -89,7 +90,8 @@ class Recall { if (batch_queries_.size() < threads_) { threads_ = batch_queries_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; } @@ -738,12 +740,13 @@ class SparseRecall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { - pool_ = make_shared(threads_, true); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { @@ -820,7 +823,8 @@ class SparseRecall { if (batch_sparse_counts_.size() < threads_) { threads_ = batch_sparse_counts_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; } diff --git a/tools/core/recall_original.cc b/tools/core/recall_original.cc index 997ef5a..22adbc5 100644 --- a/tools/core/recall_original.cc +++ b/tools/core/recall_original.cc @@ -76,12 +76,13 @@ class Recall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { - pool_ = make_shared(threads_, true); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { @@ -127,7 +128,8 @@ class Recall { if (batch_queries_.size() < threads_) { threads_ = batch_queries_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; } @@ -903,12 +905,13 @@ class SparseRecall { batch_count_(batch_count), filter_mode_{filter_mode} { if (threads_ == 0) { - pool_ = make_shared(true); + pool_ = make_shared(); threads_ = pool_->count(); cout << "Using cpu count as thread pool count[" << threads_ << "]" << endl; } else { - pool_ = make_shared(threads_, true); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Using thread pool count[" << threads_ << "]" << endl; } if (batch_count_ < 1) { @@ -985,7 +988,8 @@ class SparseRecall { if (batch_sparse_counts_.size() < threads_) { threads_ = batch_sparse_counts_.size(); - pool_ = make_shared(true, threads_); + pool_ = make_shared(threads_, false); + threads_ = pool_->count(); cout << "Query size too small, resize thread pool count[" << threads_ << "]" << endl; }