diff --git a/include/dwarfs/worker_group.h b/include/dwarfs/worker_group.h index 50a20ce8..65164bf2 100644 --- a/include/dwarfs/worker_group.h +++ b/include/dwarfs/worker_group.h @@ -67,6 +67,9 @@ class worker_group { size_t size() const { return impl_->size(); } size_t queue_size() const { return impl_->queue_size(); } double get_cpu_time() const { return impl_->get_cpu_time(); } + bool set_affinity(std::vector const& cpus) { + return impl_->set_affinity(cpus); + } template bool add_job(std::packaged_task&& task) { @@ -84,6 +87,7 @@ class worker_group { virtual size_t size() const = 0; virtual size_t queue_size() const = 0; virtual double get_cpu_time() const = 0; + virtual bool set_affinity(std::vector const& cpus) = 0; }; private: diff --git a/src/dwarfs/worker_group.cpp b/src/dwarfs/worker_group.cpp index 3c04b630..36c8d1a0 100644 --- a/src/dwarfs/worker_group.cpp +++ b/src/dwarfs/worker_group.cpp @@ -32,6 +32,7 @@ #include #include +#include #include #include #include @@ -116,6 +117,10 @@ class basic_worker_group final : public worker_group::impl, private Policy { do_work(niceness > 10); }); } + +#ifndef _WIN32 + check_set_affinity_from_enviroment(group_name); +#endif } basic_worker_group(const basic_worker_group&) = delete; @@ -219,9 +224,53 @@ class basic_worker_group final : public worker_group::impl, private Policy { return t; } + bool set_affinity(std::vector const& cpus) override { + if (cpus.empty()) { + return false; + } + +#ifndef _WIN32 + std::lock_guard lock(mx_); + + cpu_set_t cpuset; + + for (auto cpu : cpus) { + CPU_SET(cpu, &cpuset); + } + + for (size_t i = 0; i < workers_.size(); ++i) { + if (auto error = pthread_setaffinity_np(workers_[i].native_handle(), + sizeof(cpu_set_t), &cpuset); + error != 0) { + return false; + } + } +#endif + + return true; + } + private: using jobs_t = std::queue; + void check_set_affinity_from_enviroment(const char* group_name) { + if (auto var = std::getenv("DWARFS_WORKER_GROUP_AFFINITY")) { + std::vector groups; + folly::split(':', var, groups); + + for (auto& group : groups) { + std::vector parts; + folly::split('=', group, parts); + + if (parts.size() == 2 && parts[0] == group_name) { + std::vector cpus; + folly::split(',', parts[1], cpus); + set_affinity(cpus); + } + } + } + } + // TODO: move out of this class static void set_thread_niceness(int niceness) { if (niceness > 0) {