kou commented on code in PR #46034:
URL: https://github.com/apache/arrow/pull/46034#discussion_r2030744846
##########
cpp/src/arrow/util/thread_pool_test.cc:
##########
@@ -1039,35 +1040,63 @@ TEST(TestGlobalThreadPool, Capacity) {
// Exercise default capacity heuristic
ASSERT_OK(DelEnvVar("OMP_NUM_THREADS"));
ASSERT_OK(DelEnvVar("OMP_THREAD_LIMIT"));
+
+#ifdef __linux__
+ int expected_capacity = arrow::internal::GetAffinityCpuCount();
+#else
int hw_capacity = std::thread::hardware_concurrency();
- ASSERT_EQ(ThreadPool::DefaultCapacity(), hw_capacity);
+ int expected_capacity = hw_capacity;
+#endif
+
+ ASSERT_EQ(ThreadPool::DefaultCapacity(), expected_capacity);
Review Comment:
How about this?
```suggestion
ASSERT_LT(ThreadPool::DefaultCapacity(), hw_capacity);
```
##########
cpp/src/arrow/util/affinity.h:
##########
@@ -0,0 +1,26 @@
+#pragma once
+
+#include <thread>
+
+#ifdef __linux__
+#include <sched.h>
+#include <unistd.h>
+#endif
+
+namespace arrow {
+namespace internal {
+
+// Returns the number of CPUs the current process is allowed to use.
+// Falls back to std::thread::hardware_concurrency() if affinity is not
available.
+inline int GetAffinityCpuCount() {
+#ifdef __linux__
+ cpu_set_t mask;
+ if (sched_getaffinity(0, sizeof(mask), &mask) == 0) {
+ return CPU_COUNT(&mask);
+ }
+#endif
+ return std::thread::hardware_concurrency();
+}
Review Comment:
How about moving this to `cpp/src/arrow/util/cpu_info.cc`?
##########
cpp/src/arrow/util/cpu_info.cc:
##########
@@ -513,7 +514,7 @@ struct CpuInfo::Impl {
OsRetrieveCacheSize(&cache_sizes);
OsRetrieveCpuInfo(&hardware_flags, &vendor, &model_name);
original_hardware_flags = hardware_flags;
- num_cores =
std::max(static_cast<int>(std::thread::hardware_concurrency()), 1);
+ num_cores = std::max(GetAffinityCpuCount(), 1);
Review Comment:
How about not changing this and adding a new method
(`num_affinity_cores()`?) instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
