diff --git a/scripts/base/frameworks/supervisor/api.zeek b/scripts/base/frameworks/supervisor/api.zeek index ed330d5a3f..d335a04543 100644 --- a/scripts/base/frameworks/supervisor/api.zeek +++ b/scripts/base/frameworks/supervisor/api.zeek @@ -24,6 +24,7 @@ export { interface: string &optional; directory: string &optional; scripts: vector of string &default = vector(); + cpu_affinity: int &optional; cluster: table[string] of ClusterEndpoint &default=table(); # TODO: separate node config fields from status fields ? diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 24b991d88d..c82580387f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -204,6 +204,7 @@ set(MAIN_SRCS net_util.cc util.cc module_util.cc + zeek-affinity.cc Anon.cc Attr.cc Base64.cc diff --git a/src/Supervisor.cc b/src/Supervisor.cc index a719b83115..1aef520686 100644 --- a/src/Supervisor.cc +++ b/src/Supervisor.cc @@ -732,6 +732,11 @@ Supervisor::Node Supervisor::Node::FromRecord(const RecordVal* node) if ( directory_val ) rval.directory = directory_val->AsString()->CheckString(); + auto affinity_val = node->Lookup("cpu_affinity"); + + if ( affinity_val ) + rval.cpu_affinity = affinity_val->AsInt(); + auto scripts_val = node->Lookup("scripts")->AsVectorVal(); for ( auto i = 0; i < scripts_val->Size(); ++i ) @@ -779,7 +784,10 @@ Supervisor::Node Supervisor::Node::FromJSON(std::string_view json) rval.interface = *it; if ( auto it = j.find("directory"); it != j.end() ) - rval.directory= *it; + rval.directory = *it; + + if ( auto it = j.find("cpu_affinity"); it != j.end() ) + rval.cpu_affinity = *it; auto scripts = j["scripts"]; @@ -833,6 +841,9 @@ IntrusivePtr Supervisor::Node::ToRecord() const if ( directory ) rval->Assign(rt->FieldOffset("directory"), new StringVal(*directory)); + if ( cpu_affinity ) + rval->Assign(rt->FieldOffset("cpu_affinity"), val_mgr->GetInt(*cpu_affinity)); + auto st = BifType::Record::Supervisor::Node->FieldType("scripts"); auto scripts_val = new VectorVal(st->AsVectorType()); rval->Assign(rt->FieldOffset("scripts"), scripts_val); diff --git a/src/Supervisor.h b/src/Supervisor.h index 70cc89e3a6..623afc36a0 100644 --- a/src/Supervisor.h +++ b/src/Supervisor.h @@ -47,6 +47,7 @@ public: std::string name; std::optional interface; std::optional directory; + std::optional cpu_affinity; std::vector scripts; std::map cluster; diff --git a/src/main.cc b/src/main.cc index f5fcb2a07e..021c5207aa 100644 --- a/src/main.cc +++ b/src/main.cc @@ -78,6 +78,8 @@ extern "C" { #include "setsignal.h" }; +#include "zeek-affinity.h" + #ifdef USE_PERFTOOLS_DEBUG HeapLeakChecker* heap_checker = 0; int perftools_leaks = 0; @@ -944,7 +946,7 @@ int main(int argc, char** argv) { if ( setenv("CLUSTER_NODE", node_name.data(), true) == -1 ) { - fprintf(stderr, "cluster node %s failed to setenv: %s\n", + fprintf(stderr, "node '%s' failed to setenv: %s\n", node_name.data(), strerror(errno)); exit(1); } @@ -954,7 +956,7 @@ int main(int argc, char** argv) { if ( chdir(zeek::supervised_node->directory->data()) ) { - fprintf(stderr, "supervised node %s failed to chdir to %s: %s\n", + fprintf(stderr, "node '%s' failed to chdir to %s: %s\n", node_name.data(), zeek::supervised_node->directory->data(), strerror(errno)); @@ -962,6 +964,15 @@ int main(int argc, char** argv) } } + if ( zeek::supervised_node->cpu_affinity ) + { + auto res = zeek::set_affinity(*zeek::supervised_node->cpu_affinity); + + if ( ! res ) + fprintf(stderr, "node '%s' failed to set CPU affinity: %s\n", + node_name.data(), strerror(errno)); + } + for ( const auto& s : zeek::supervised_node->scripts ) options.scripts_to_load.emplace_back(s); } diff --git a/src/zeek-affinity.cc b/src/zeek-affinity.cc new file mode 100644 index 0000000000..7fef4c203a --- /dev/null +++ b/src/zeek-affinity.cc @@ -0,0 +1,51 @@ +// See the file "COPYING" in the main distribution directory for copyright. + +#if defined(__linux__) + +#if !defined(_GNU_SOURCE) +#define _GNU_SOURCE +#endif + +#include + +namespace zeek { +bool set_affinity(int core_number) + { + cpu_set_t cpus; + CPU_ZERO(&cpus); + CPU_SET(core_number, &cpus); + auto res = sched_setaffinity(0, sizeof(cpus), &cpus); + return res == 0; + } +} // namespace zeek + +#elif defined(__FreeBSD__) + +#include +#include + +namespace zeek { +bool set_affinity(int core_number) + { + cpuset_t cpus; + CPU_ZERO(&cpus); + CPU_SET(core_number, &cpus); + auto res = cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, -1, + sizeof(cpus), &cpus); + return res == 0; + } +} // namespace zeek + +#else + +#include + +namespace zeek { +bool set_affinity(int core_number) + { + errno = ENOTSUP; + return false; + } +} // namespace zeek + +#endif diff --git a/src/zeek-affinity.h b/src/zeek-affinity.h new file mode 100644 index 0000000000..ce4bdd9c49 --- /dev/null +++ b/src/zeek-affinity.h @@ -0,0 +1,15 @@ +// See the file "COPYING" in the main distribution directory for copyright. + +namespace zeek { + +/** + * Set the process affinity to a given CPU. Currently only supported on + * Linux and FreeBSD. + * @param core_number the core to which this process should set its affinity. + * Cores are typically numbered 0..N. + * @return true if the affinity is successfully set and false if not with + * errno additionally being set to indicate the reason. + */ +bool set_affinity(int core_number); + +} // namespace zeek