perfect-cli cleans up on SIGINT, fixed a problem where cpu_set would silently fail

include/perfect/cpu_set.hpp

@@ -12,18 +12,10 @@
 #include <string>
 #include <vector>
 
+#include "detail/fs.hpp"
 #include "init.hpp"
 #include "result.hpp"
 
-#define SUCCESS_OR_RETURN(stmt) \
-{\
-  Result _ret; \
-  _ret = (stmt); \
-if (_ret != Result::SUCCESS) {\
-  return _ret;\
-}\
-}
-
 std::set<int> operator-(const std::set<int> &lhs, const std::set<int> &rhs) {
   std::set<int> result;
   for (auto e : lhs) {
@@ -34,6 +26,17 @@ std::set<int> operator-(const std::set<int> &lhs, const std::set<int> &rhs) {
   return result;
 }
 
+// intersection
+std::set<int> operator&(const std::set<int> &lhs, const std::set<int> &rhs) {
+  std::set<int> result;
+  for (auto e : lhs) {
+    if (1 == rhs.count(e)) {
+      result.insert(e);
+    }
+  }
+  return result;
+}
+
 std::string remove_space(const std::string &s) {
   std::string result;
 
@@ -86,7 +89,6 @@ std::set<int> parse_token(const std::string &token) {
 }
 
 std::set<int> parse_cpuset(const std::string &s) {
-  // std::cerr << "parse_cpuset: parsing '" << s << "'\n";
   std::set<int> result;
 
   std::string token;
@@ -109,11 +111,12 @@ namespace perfect {
 class CpuSet {
 public:
   std::string path_;
-  std::set<int> cpus_;
-  std::set<int> mems_;
   CpuSet *parent_;
 
-  // make sure cpuset is initialized
+  CpuSet() : path_(""), parent_(nullptr) {}
+  CpuSet(const CpuSet &other) : path_(other.path_), parent_(other.parent_) {}
+
+  // make sure cpuset system is initialized
   static Result init() {
 
     // check for "nodev cpuset" in /proc/filesystems
@@ -162,23 +165,24 @@ public:
     return Result::SUCCESS;
   }
 
-  std::string get_raw_cpus() {
+  std::string get_raw_cpus() const {
-    std::ifstream is(path_ + "/cpuset.cpus");
+    std::string path = path_ + "/cpuset.cpus";
+    std::ifstream is(path);
     std::stringstream ss;
     ss << is.rdbuf();
     return remove_space(ss.str());
   }
 
-  std::string get_raw_mems() {
+  std::string get_raw_mems() const {
     std::ifstream is(path_ + "/cpuset.mems");
     std::stringstream ss;
     ss << is.rdbuf();
     return remove_space(ss.str());
   }
 
-  std::set<int> get_cpus() { return parse_cpuset(get_raw_cpus()); }
+  std::set<int> get_cpus() const { return parse_cpuset(get_raw_cpus()); }
 
-  std::set<int> get_mems() { return parse_cpuset(get_raw_mems()); }
+  std::set<int> get_mems() const { return parse_cpuset(get_raw_mems()); }
 
   // migrate the caller task from this cpu set to another
   Result migrate_self_to(CpuSet &other) {
@@ -193,11 +197,12 @@ public:
     std::string line;
     while (std::getline(is, line)) {
       line = remove_space(line);
+
       if (std::to_string(self) == line) {
-        // std::cerr << "migrating self task " << line << " to " << other.path
+        // std::cerr << "migrating self task " << line << " to " << other.path_
         //            << "\n";
-        other.write_task(line);
+        pid_t pid = std::stoi(line);
-        return Result::SUCCESS;
+        return other.write_task(pid);
       }
     }
     return Result::NO_TASK;
@@ -205,46 +210,58 @@ public:
 
   // migrate tasks in this cpu set to another
   Result migrate_tasks_to(CpuSet &other) {
+    // other must have cpus and mems
+    auto s = other.get_cpus();
+
+    assert(!other.get_cpus().empty());
+    assert(!other.get_mems().empty());
+
     // enable memory migration in other
-    SUCCESS_OR_RETURN(other.enable_memory_migration());
+    PERFECT_SUCCESS_OR_RETURN(other.enable_memory_migration());
 
     // read this tasks and write each line to other.tasks
     std::ifstream is(path_ + "/tasks");
     std::string line;
     while (std::getline(is, line)) {
-      // std::cerr << "migrating task " << line << " to " << other.path << "\n";
+      pid_t pid = std::stoi(line);
-      other.write_task(line);
+      // std::cerr << "migrating task " << pid << " to " << other.path_ << "\n";
+      Result result = other.write_task(pid);
+      if (Result::ERRNO_INVALID == result) {
+        // std::cerr << "task " << pid << " is unmovable\n";
+      } else {
+        PERFECT_SUCCESS_OR_RETURN(result);
+      }
     }
-
     return Result::SUCCESS;
   }
 
   Result enable_memory_migration() {
-    std::ofstream ofs(path_ + "/" + "cpuset.memory_migrate");
+    return detail::write_str(path_ + "/cpuset.memory_migrate", "1");
-    ofs << "1";
-    ofs.close();
-    if (ofs.fail()) {
-    switch (errno) {
-    case EACCES:
-      return Result::NO_PERMISSION;
-    case ENOENT:
-      return Result::NOT_SUPPORTED;
-    default:
-      return Result::UNKNOWN;
-    }
-  }
-    return Result::SUCCESS;
   }
 
-  void write_task(const std::string &task) {
+  Result write_task(pid_t pid) {
-    // write `task` to path/tasks
+    return detail::write_str(path_ + "/tasks", std::to_string(pid) + "\n");
-    std::ofstream os(path_ + "/tasks");
+  }
-    os << task << "\n";
+
+  static Result get_affinity(std::set<int> &cpus, pid_t pid) {
+    cpu_set_t mask;
+    CPU_ZERO(&mask);
+    if (sched_getaffinity(pid, sizeof(mask), &mask)) {
+      return from_errno(errno);
+    }
+
+    cpus.clear();
+    for (int i = 0; i < CPU_SETSIZE; ++i) {
+      if
+        CPU_ISSET(i, &mask) { cpus.insert(i); }
+    }
+
+    return Result::SUCCESS;
   }
 
   // object representing the root CPU set
   static Result get_root(CpuSet &root) {
-    SUCCESS_OR_RETURN(CpuSet::init());
+    PERFECT_SUCCESS_OR_RETURN(CpuSet::init());
     root.path_ = "/dev/cpuset";
     root.parent_ = nullptr;
     return Result::SUCCESS;
@@ -256,7 +273,7 @@ public:
   Result make_child(CpuSet &child, const std::string &name) {
 
     if (mkdir((path_ + "/" + name).c_str(),
-              S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH)) {
+              S_IRUSR | S_IWUSR | S_IRGRP | S_IXGRP | S_IROTH | S_IXOTH)) {
       switch (errno) {
       case EEXIST: {
         // okay
@@ -264,8 +281,6 @@ public:
       }
       case EACCES:
         return Result::NO_PERMISSION;
-      case ENOENT:
-      case EROFS:
       default:
         return Result::UNKNOWN;
       }
@@ -276,6 +291,8 @@ public:
     return Result::SUCCESS;
   }
 
+  std::vector<CpuSet> get_children() { assert(false && "unimplemented"); }
+
   Result enable_cpu(const int cpu) {
     std::set<int> cpus = get_cpus();
     cpus.insert(cpu);
@@ -290,30 +307,28 @@ public:
     return write_cpus(finalCpus);
   }
 
-  // FIXME: check error
   Result write_cpus(std::set<int> cpus) {
-    std::ofstream os(path_ + "/cpuset.cpus");
+    std::string str;
     bool comma = false;
     for (auto cpu : cpus) {
       if (comma)
-        os << ",";
+        str += ",";
-      os << cpu << "-" << cpu;
+      str += std::to_string(cpu) + "-" + std::to_string(cpu);
       comma = true;
     }
-    return Result::SUCCESS;
+    return detail::write_str(path_ + "/cpuset.cpus", str);
   }
 
-  // FIXME: check write
   Result write_mems(std::set<int> mems) {
-    std::ofstream os(path_ + "/cpuset.mems");
+    std::string str;
     bool comma = false;
     for (auto mem : mems) {
       if (comma)
-        os << ",";
+        str += ",";
-      os << mem << "-" << mem;
+      str += std::to_string(mem) + "-" + std::to_string(mem);
       comma = true;
     }
-    return Result::SUCCESS;
+    return detail::write_str(path_ + "/cpuset.mems", str);
   }
 
   Result enable_mem(const int mem) {
@@ -331,26 +346,35 @@ public:
   }
 
   Result destroy() {
+
+    // already destroyed
+    if (!detail::path_exists(path_)) {
+      return Result::SUCCESS;
+    }
+
     // remove all child cpu sets
 
     // move all attached processes back to parent
-    assert(parent_);
+    assert(parent_ && "should not call destroy on root cpuset");
-    migrate_tasks_to(*parent_);
+    PERFECT_SUCCESS_OR_RETURN(migrate_tasks_to(*parent_));
 
     // remove with rmdir
+    Result result = Result::UNKNOWN;
     if (rmdir(path_.c_str())) {
       switch (errno) {
+      case ENOENT:
+        // already gone
+        result = Result::SUCCESS;
+        break;
       default:
         std::cerr << "unhandled error in rmdir: " << strerror(errno) << "\n";
-        return Result::UNKNOWN;
+        result = Result::UNKNOWN;
       }
     }
 
     path_ = "";
-    return Result::SUCCESS;
+    return result;
   }
-
-
 };
 
 std::ostream &operator<<(std::ostream &s, const CpuSet &c) {

include/perfect/detail/fs.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cstring>
 #include <fstream>
 #include <string>
 
@@ -32,15 +33,20 @@ Result write_str(const std::string &path, const std::string &val) {
   if (ofs.fail()) {
     switch (errno) {
     case EACCES:
-    std::cerr << "EACCES when writing to " << path << "\n";
+      // std::cerr << "EACCES when writing to " << path << "\n";
       return Result::NO_PERMISSION;
     case EPERM:
-      std::cerr << "EPERM when writing to " << path << "\n";
+      // std::cerr << "EPERM when writing to " << path << "\n";
       return Result::NO_PERMISSION;
     case ENOENT:
-    std::cerr << "ENOENT when writing to " << path << "\n";
+      // std::cerr << "ENOENT when writing to " << path << "\n";
       return Result::NOT_SUPPORTED;
+    case EINVAL:
+      // std::cerr << "EINVAL when writing to " << path << "\n";
+      return Result::ERRNO_INVALID;
     default:
+      std::cerr << strerror(errno) << " when writing " << val << " to " << path
+                << "\n";
       return Result::UNKNOWN;
     }
   }

include/perfect/detail/os/linux.hpp

@@ -118,8 +118,17 @@ Result set_high_priority() {
   return Result::SUCCESS;
 }
 
-
+// disable all but one SMT thread for all CPUs the calling process can run on
-
+Result disable_smt() {
+  return Result::NOT_SUPPORTED;
 }
 
+// enable SMT for all CPUs the calling process can run on
+Result enable_smt() {
+return Result::NOT_SUPPORTED;
+}
+
+
+} // namespace detail
+
 } // namespace perfect

include/perfect/result.hpp

@@ -22,7 +22,7 @@ enum class Result {
   NO_PERMISSION,
   NOT_SUPPORTED,
   NO_TASK,
-
+  ERRNO_INVALID,
   NVML_NO_PERMISSION,
   NVML_NOT_SUPPORTED,
   NVML_UNINITIALIZED,
@@ -88,6 +88,8 @@ const char *get_string(const Result &result) {
     return "unsupported operation";
   case Result::NO_TASK:
     return "no such task";
+  case Result::ERRNO_INVALID:
+    return "errno EINVAL";
   case Result::UNKNOWN:
     return "unknown error";
   case Result::NVML_NOT_SUPPORTED:

tools/perfect.cpp

@@ -1,6 +1,7 @@
 #include <cassert>
 #include <cerrno>
 #include <chrono>
+#include <functional>
 #include <iostream>
 #include <string>
 #include <thread>
@@ -9,6 +10,7 @@
 #ifdef __linux__
 #include <fcntl.h>
 #include <pwd.h>
+#include <signal.h>
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <sys/wait.h>
@@ -28,6 +30,26 @@
 #include "perfect/os_perf.hpp"
 #include "perfect/priority.hpp"
 
+typedef std::function<perfect::Result()> CleanupFn;
+std::vector<CleanupFn> cleanups;
+
+// restore the system state to how we found it
+void cleanup(int dummy) {
+  (void)dummy;
+  std::cerr << "caught ctrl-c\n";
+
+  // unregister our handler
+  signal(SIGINT, SIG_DFL);
+  std::cerr << "cleaning up\n";
+  std::cerr << "ctrl-c again to quit\n";
+
+  for (auto f : cleanups) {
+    perfect::Result result = f();
+  }
+
+  exit(EXIT_FAILURE);
+}
+
 // argv should be null-terminated
 // outf and errf are file descriptors to where stdout and stderr should be
 // redirected write stdout to out and stderr to err, if not null
@@ -145,6 +167,9 @@ int fork_child(char *const *argv, int outf, int errf) {
 }
 
 int main(int argc, char **argv) {
+
+  signal(SIGINT, cleanup);
+
   using namespace clipp;
 
   size_t numUnshielded = 0;
@@ -186,7 +211,9 @@ int main(int argc, char **argv) {
                     maxOsPerf = false;
                   }),
                   option("--aslr").set(aslr, true).doc("enable ASLR"),
-                  option("--no-priority").set(highPriority, false).doc("don't set high priority"),
+                  option("--no-priority")
+                      .set(highPriority, false)
+                      .doc("don't set high priority"),
                   option("--cpu-turbo").doc("enable CPU turbo").call([&]() {
                     cpuTurbo = true;
                   }),
@@ -283,34 +310,44 @@ int main(int argc, char **argv) {
   }
 
   // handle CPU shielding
-  perfect::CpuSet shielded, unshielded;
+  perfect::CpuSet root, shielded, unshielded;
   if (numShielded) {
     std::cerr << "shielding " << numShielded << " cpus\n";
 
-    perfect::CpuSet root;
     PERFECT(perfect::CpuSet::get_root(root));
     PERFECT(root.make_child(shielded, "shielded"));
     PERFECT(root.make_child(unshielded, "unshielded"));
 
     std::cerr << "enable memory\n";
     PERFECT(shielded.enable_mem(0));
-    PERFECT(shielded.enable_mem(0));
+    PERFECT(unshielded.enable_mem(0));
 
     std::cerr << "enable cpus\n";
     size_t i = 0;
-    for (; i < numShielded; ++i) {
+    for (; i < cpus.size() - numShielded; ++i) {
-      std::cerr << "shield cpu " << cpus[i] << "\n";
-      shielded.enable_cpu(cpus[i]);
-    }
-    for (; i < cpus.size(); ++i) {
       std::cerr << "unshield cpu " << cpus[i] << "\n";
       unshielded.enable_cpu(cpus[i]);
     }
+    for (; i < cpus.size(); ++i) {
+      std::cerr << "shield cpu " << cpus[i] << "\n";
+      shielded.enable_cpu(cpus[i]);
+    }
 
     std::cerr << "migrate self\n";
     PERFECT(root.migrate_self_to(shielded));
-    std::cerr << "migrate other\n";
+    std::cerr << "migrate other (1/2)\n";
     PERFECT(root.migrate_tasks_to(unshielded));
+    // some tasks may have been spawned by unmigrated tasks while we migrated
+    std::cerr << "migrate other (2/2)\n";
+    PERFECT(root.migrate_tasks_to(unshielded));
+
+    cleanups.push_back(CleanupFn([&] {
+      std::cerr << "cleanup: shielded cpu set\n";
+      shielded.destroy();
+      std::cerr << "cleanup: unshielded cpu set\n";
+      unshielded.destroy();
+      return perfect::Result::SUCCESS;
+    }));
   }
 
   // handle aslr
@@ -330,6 +367,11 @@ int main(int argc, char **argv) {
       std::cerr << "enabling cpu turbo\n";
       PERFECT(perfect::enable_cpu_turbo());
     }
+
+    cleanups.push_back(CleanupFn([&] {
+      std::cerr << "cleanup: restore CPU turbo state\n";
+      return perfect::set_cpu_turbo_state(cpuTurboState);
+    }));
   }
 
   // handle governor
@@ -342,6 +384,11 @@ int main(int argc, char **argv) {
         PERFECT(perfect::os_perf_state_maximum(cpu));
       }
     }
+
+    cleanups.push_back(CleanupFn([&] {
+      std::cerr << "cleanup: os governor\n";
+      return perfect::set_os_perf_state(osPerfState);
+    }));
   }
 
   if (highPriority) {