Commit 2f7fdc20 authored by Dan Melnic's avatar Dan Melnic Committed by Facebook GitHub Bot

Add support for Subprocess to call sched_setaffinity

Summary: Add support for Subprocess to call sched_setaffinity

Reviewed By: yfeldblum

Differential Revision: D29722725

fbshipit-source-id: b0d4577bf3caaeb65137c8168bd27e1f402969da
parent d26d241b
......@@ -540,6 +540,14 @@ int Subprocess::prepareChild(
}
}
#ifdef __linux__
// Best effort
if (options.cpuSet_.hasValue()) {
const auto& cpuSet = options.cpuSet_.value();
::sched_setaffinity(0, sizeof(cpuSet), &cpuSet);
}
#endif
// We don't have to explicitly close the parent's end of all pipes,
// as they all have the FD_CLOEXEC flag set and will be closed at
// exec time.
......
......@@ -475,6 +475,13 @@ class Subprocess {
}
#endif
#if defined(__linux__)
Options& setCpuSet(const cpu_set_t& cpuSet) {
cpuSet_ = cpuSet;
return *this;
}
#endif
private:
typedef boost::container::flat_map<int, int> FdMap;
FdMap fdActions_;
......@@ -493,6 +500,9 @@ class Subprocess {
// none means `vfork()` instead of a custom `clone()`
// Optional<> is used because value of '0' means do clone without any flags.
Optional<clone_flags_t> cloneFlags_;
#endif
#if defined(__linux__)
Optional<cpu_set_t> cpuSet_;
#endif
};
......
......@@ -348,6 +348,27 @@ TEST(SimpleSubprocessTest, DetachExecFails) {
"/no/such/file");
}
TEST(SimpleSubprocessTest, Affinity) {
#ifdef __linux__
cpu_set_t cpuSet0;
CPU_ZERO(&cpuSet0);
CPU_SET(1, &cpuSet0);
CPU_SET(2, &cpuSet0);
CPU_SET(3, &cpuSet0);
Subprocess::Options options;
Subprocess proc(
std::vector<std::string>{"/bin/sleep", "5"}, options.setCpuSet(cpuSet0));
EXPECT_NE(proc.pid(), -1);
cpu_set_t cpuSet1;
CPU_ZERO(&cpuSet1);
auto ret = ::sched_getaffinity(proc.pid(), sizeof(cpu_set_t), &cpuSet1);
CHECK_EQ(ret, 0);
CHECK_EQ(::memcmp(&cpuSet0, &cpuSet1, sizeof(cpu_set_t)), 0);
auto retCode = proc.waitOrTerminateOrKill(1s, 1s);
EXPECT_TRUE(retCode.killed());
#endif // __linux__
}
TEST(SimpleSubprocessTest, FromExistingProcess) {
// Manually fork a child process using fork() without exec(), and test waiting
// for it using the Subprocess API in the parent process.
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment