xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(thread-pool): make CPU affinity opt-in (#623)

This commit is contained in:
luoxiaojian 2026-07-29 15:30:04 +08:00 committed by GitHub
parent 31be25598a
commit 1ad6df2539
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
25 changed files with 435 additions and 74 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
python/tests/detail/test_db_config.py
View File

@ -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)

3
python/zvec/zvec.py
View File

@ -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.

103
src/ailego/parallel/thread_pool.cc
View File

@ -12,33 +12,64 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <algorithm>
#include <zvec/ailego/logger/logger.h>
#include <zvec/ailego/parallel/thread_pool.h>
#if (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__)
#if defined(__linux__) && !defined(__ANDROID__)
#include <pthread.h>
static inline void BindThreads(std::vector<std::thread> &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<std::thread> &pool) {
cpu_set_t allowed_mask;
if (!GetAllowedCpuMask(&allowed_mask)) {
return;
}
std::vector<int> 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<std::thread> &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<std::thread> &) {}
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
} // namespace zvec

2
src/binding/python/model/common/python_config.cc
View File

@ -242,4 +242,4 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) {
}
} // namespace zvec
} // namespace zvec

10
src/core/interface/index.cc
View File

@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <algorithm>
#include <magic_enum/magic_enum.hpp>
#include <zvec/core/framework/index_error.h>
#include <zvec/core/framework/index_storage.h>
@ -1013,19 +1014,24 @@ int Index::Merge(const std::vector<Index::Pointer> &indexes,
}
// must declare here to ensure its lifespan can cover reducer->reduce()
std::unique_ptr<ailego::ThreadPool> 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<uint32_t>(effective_write_concurrency, options.pool->count());
} else {
local_thread_pool =
std::make_unique<ailego::ThreadPool>(options.write_concurrency);
std::make_unique<ailego::ThreadPool>(options.write_concurrency, false);
reducer->set_thread_pool(local_thread_pool.get());
effective_write_concurrency =
static_cast<uint32_t>(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;

58
src/db/common/cgroup_util.cc
View File

@ -13,6 +13,9 @@
// limitations under the License.
#include "db/common/cgroup_util.h"
#include <charconv>
#include <sstream>
#include <system_error>
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<uint64_t>::max() && quota != 0 &&
period > 0) {
cpu_cores_ =
static_cast<int>(std::ceil(static_cast<double>(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<uint64_t>((std::numeric_limits<int>::max)())) {
return false;
}
*cpu_cores = static_cast<int>(cores);
return true;
}
void CgroupUtil::updateMemoryLimit() {
if (readMemoryCgroup()) {
return;
@ -491,4 +527,4 @@ double CgroupUtil::calculateMacOSCpuUsage() {
}
#endif
} // namespace zvec
} // namespace zvec

7
src/db/common/cgroup_util.h
View File

@ -17,7 +17,6 @@
#include <fstream>
#include <iostream>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
@ -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
} // namespace zvec

6
src/db/common/config.cc
View File

@ -35,10 +35,12 @@ GlobalConfig::ConfigData::ConfigData()
DEFAULT_MEMORY_LIMIT_RATIO),
log_config(std::make_shared<ConsoleLogConfig>()),
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
} // namespace zvec

10
src/db/common/global_resource.cc
View File

@ -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());
});

2
src/db/index/segment/segment.cc
View File

@ -1661,7 +1661,7 @@ Result<VectorColumnIndexer::Ptr> 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<uint32_t>(merge_options.pool->count());
} else {
merge_options.write_concurrency = concurrency;
}

2
src/db/index/segment/segment_helper.cc
View File

@ -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<uint32_t>(merge_options.pool->count());
} else {
merge_options.write_concurrency = concurrency;
}

25
src/include/zvec/ailego/parallel/thread_pool.h
View File

@ -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<std::mutex> lock(queue_mutex_);
stopping_ = true;
work_cond_.notify_all();
}

1
src/include/zvec/core/framework/index_threads.h
View File

@ -14,6 +14,7 @@
#pragma once
#include <algorithm>
#include <thread>
#include <utility>
#include <zvec/ailego/parallel/thread_pool.h>

14
src/include/zvec/db/config.h
View File

@ -93,6 +93,8 @@ class ZVEC_API GlobalConfig : public ailego::Singleton<GlobalConfig> {
// 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<GlobalConfig> {
// 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<GlobalConfig> {
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<GlobalConfig> {
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;

88
tests/ailego/parallel/thread_pool_test.cc
View File

@ -12,14 +12,98 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <algorithm>
#include <chrono>
#include <iostream>
#include <limits>
#include <memory>
#include <thread>
#include <type_traits>
#include <gtest/gtest.h>
#include <zvec/ailego/parallel/thread_pool.h>
#if defined(__linux__) && !defined(__ANDROID__)
#include <pthread.h>
#endif
using namespace zvec::ailego;
static_assert(!std::is_constructible<ThreadPool, uint32_t>::value,
"ThreadPool thread count must not be ambiguous with binding");
static_assert(!std::is_constructible<ThreadPool, bool>::value,
"ThreadPool binding must require an explicit thread count");
static_assert(std::is_constructible<ThreadPool, uint32_t, bool>::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<uint32_t>::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<ThreadPool>()) {}
@ -38,7 +122,9 @@ struct A {
};
struct B {
B(void) : pool(std::make_shared<ThreadPool>(true)) {}
B(void)
: pool(std::make_shared<ThreadPool>(
std::max(std::thread::hardware_concurrency(), 1u), true)) {}
std::string ThreadMain(uint32_t &num) {
aaa.pool->enqueue(

15
tests/core/interface/index_interface_test.cc
View File

@ -16,6 +16,7 @@
#include <cmath>
#include <functional>
#include <iostream>
#include <limits>
#include <numeric>
#include <random>
#include <unordered_map>
@ -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<uint32_t>::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<uint32_t>::max)();
merge_options.pool = &pool;
ASSERT_TRUE(0 ==
index3->Merge({index1, index2}, filter, merge_options));
ASSERT_TRUE(2 == index3->GetDocCount());
{
VectorDataBuffer fetched_vector_data;

45
tests/db/common/cgroup_util_test.cc Normal file
View File

@ -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 <gtest/gtest.h>
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);
}

13
tests/db/common/config_test.cc
View File

@ -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);
}
}

46
tests/db/common/global_resource_test.cc Normal file
View File

@ -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 <algorithm>
#include <thread>
#include <gtest/gtest.h>
#include <zvec/db/config.h>
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());
}

6
tools/core/bench.cc
View File

@ -35,12 +35,13 @@ class Bench {
retrieval_mode_{retrieval_mode},
filter_mode_{filter_mode} {
if (threads_ == 0) {
pool_ = make_shared<ThreadPool>(false);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(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<ThreadPool>(false);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(threads_, false);
threads_ = pool_->count();
cout << "Using thread pool count[" << threads_ << "]" << endl;
}
if (batch_count_ < 1) {

6
tools/core/bench_original.cc
View File

@ -69,12 +69,13 @@ class Bench {
retrieval_mode_{retrieval_mode},
filter_mode_{filter_mode} {
if (threads_ == 0) {
pool_ = make_shared<ThreadPool>(false);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(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<ThreadPool>(false);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(threads_, false);
threads_ = pool_->count();
cout << "Using thread pool count[" << threads_ << "]" << endl;
}
if (batch_count_ < 1) {

2
tools/core/local_builder.cc
View File

@ -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<uint32_t>(pool.count());
std::atomic<size_t> 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<uint32_t>(pool.count());
std::atomic<size_t> finished{0};
int errcode = 0;
std::mutex mutex;

2
tools/core/local_builder_original.cc
View File

@ -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<uint32_t>(pool.count());
std::atomic<size_t> 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<uint32_t>(pool.count());
std::atomic<size_t> finished{0};
int errcode = 0;
std::mutex mutex;

16
tools/core/recall.cc
View File

@ -36,12 +36,13 @@ class Recall {
batch_count_(batch_count),
filter_mode_{filter_mode} {
if (threads_ == 0) {
pool_ = make_shared<ThreadPool>(true);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(threads_, true);
pool_ = make_shared<ThreadPool>(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<ThreadPool>(true, threads_);
pool_ = make_shared<ThreadPool>(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<ThreadPool>(true);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(threads_, true);
pool_ = make_shared<ThreadPool>(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<ThreadPool>(true, threads_);
pool_ = make_shared<ThreadPool>(threads_, false);
threads_ = pool_->count();
cout << "Query size too small, resize thread pool count[" << threads_
<< "]" << endl;
}

16
tools/core/recall_original.cc
View File

@ -76,12 +76,13 @@ class Recall {
batch_count_(batch_count),
filter_mode_{filter_mode} {
if (threads_ == 0) {
pool_ = make_shared<ThreadPool>(true);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(threads_, true);
pool_ = make_shared<ThreadPool>(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<ThreadPool>(true, threads_);
pool_ = make_shared<ThreadPool>(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<ThreadPool>(true);
pool_ = make_shared<ThreadPool>();
threads_ = pool_->count();
cout << "Using cpu count as thread pool count[" << threads_ << "]"
<< endl;
} else {
pool_ = make_shared<ThreadPool>(threads_, true);
pool_ = make_shared<ThreadPool>(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<ThreadPool>(true, threads_);
pool_ = make_shared<ThreadPool>(threads_, false);
threads_ = pool_->count();
cout << "Query size too small, resize thread pool count[" << threads_
<< "]" << endl;
}