Handle SIGINT/SIGTERM gracefully

Summary:
Oomd manipulates some system state at various points, better to exit
gracefully on these signals rather than leave state polluted.

Reviewed By: brianc118

Differential Revision: D51629260

fbshipit-source-id: eaf79776453b253a086f19be4029bfbcf2851eb5

src/oomd/Main.cpp

@@ -248,6 +248,15 @@ int main(int argc, char** argv) {
   struct Oomd::IOCostCoeffs hdd_coeffs = default_hdd_coeffs;
   struct Oomd::IOCostCoeffs ssd_coeffs = default_ssd_coeffs;
 
+  sigset_t mask;
+  sigemptyset(&mask);
+  sigaddset(&mask, SIGINT);
+  sigaddset(&mask, SIGTERM);
+  if (sigprocmask(SIG_BLOCK, &mask, nullptr) == -1) {
+    perror("sigprocmask");
+    exit(EXIT_FAILURE);
+  }
+
   const char* const short_options = "hvC:w:i:f:c:lD:dr";
   option long_options[] = {
       option{"help", no_argument, nullptr, 'h'},
@@ -485,5 +494,5 @@ int main(int argc, char** argv) {
       *io_devs,
       hdd_coeffs,
       ssd_coeffs);
-  return oomd.run();
+  return oomd.run(&mask);
 }

src/oomd/Oomd.cpp

@@ -17,18 +17,19 @@
 
 #include "oomd/Oomd.h"
 
+#include <signal.h>
 #include <cmath>
 #include <functional>
 #include <thread>
 
+#include <strings.h>
 #include "oomd/CgroupContext.h"
 #include "oomd/Log.h"
 #include "oomd/dropin/FsDropInService.h"
 #include "oomd/include/Assert.h"
 #include "oomd/include/Defines.h"
 #include "oomd/util/Fs.h"
 #include "oomd/util/Util.h"
-
 namespace Oomd {
 
 Oomd::Oomd(
@@ -110,17 +111,33 @@ void Oomd::updateContext() {
   ctx_.bumpCurrentTick();
 }
 
-int Oomd::run() {
+int Oomd::run(const sigset_t* mask) {
   if (!engine_) {
     OLOG << "Could not run engine. Your config file is probably invalid\n";
     return EXIT_CANT_RECOVER;
   }
 
+  struct timespec ts {
+    .tv_sec = interval_.count(), .tv_nsec = 0,
+  };
+
   OLOG << "Running oomd";
 
   while (true) {
-    /* sleep override */
-    std::this_thread::sleep_for(interval_);
+    // sigtimedwait so if we get a SIGINT or SIGTERM we wake up and exit right
+    // away.
+    int rc;
+    do {
+      rc = sigtimedwait(mask, nullptr, &ts);
+    } while (rc == -1 && errno == EINTR);
+
+    if (rc == -1 && errno != EAGAIN) {
+      perror("sigtimedwait");
+      exit(EXIT_FAILURE);
+    } else if (rc != -1) {
+      OLOG << "Received signal " << rc << ". Exiting!";
+      return 128 + rc;
+    }
 
     if (fs_drop_in_service_) {
       fs_drop_in_service_->updateDropIns();

src/oomd/Oomd.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include <signal.h>
 #include <chrono>
 #include <memory>
 #include <string>
@@ -48,7 +49,7 @@ class Oomd {
   ~Oomd();
 
   void updateContext();
-  int run();
+  int run(const sigset_t* mask);
 
  private:
   // runtime settings