From cd142c567835a7e0b75ebcf32eae4ca23a0f6eb0 Mon Sep 17 00:00:00 2001
From: AWoloszyn <awoloszyn@google.com>
Date: Mon, 24 Sep 2018 14:53:24 -0400
Subject: [PATCH] Linux: Be more aggressive about keeping the memory tracker
 on.

---
 core/memory_tracker/cc/memory_tracker.cpp     |  4 ++
 core/memory_tracker/cc/memory_tracker.h       |  2 +-
 .../cc/posix/memory_tracker.inc               | 55 +++++++++----------
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/core/memory_tracker/cc/memory_tracker.cpp b/core/memory_tracker/cc/memory_tracker.cpp
index b832b7728f..f7802353c2 100644
--- a/core/memory_tracker/cc/memory_tracker.cpp
+++ b/core/memory_tracker/cc/memory_tracker.cpp
@@ -68,6 +68,10 @@ std::vector<void*> DirtyPageTable::DumpAndClearAll() {
 
 template <>
 bool MemoryTracker::AddTrackingRangeImpl(void* start, size_t size) {
+  if (!EnableMemoryTrackerImpl()) {
+    return false;
+  }
+
   if (size == 0) return false;
   if (IsInRanges(reinterpret_cast<uintptr_t>(start), ranges_)) return false;
 
diff --git a/core/memory_tracker/cc/memory_tracker.h b/core/memory_tracker/cc/memory_tracker.h
index f98c8a841c..8c6b997f2d 100644
--- a/core/memory_tracker/cc/memory_tracker.h
+++ b/core/memory_tracker/cc/memory_tracker.h
@@ -131,7 +131,7 @@ inline bool IsInRanges(uintptr_t addr, std::map<uintptr_t, size_t>& ranges,
   return true;
 }
 
-// SpinLock is a spin lock implemented with atomic variable and opertions.
+// SpinLock is a spin lock implemented with atomic variable and operations.
 // Mutiple calls to Lock in a single thread will result into a deadlock.
 class SpinLock {
  public:
diff --git a/core/memory_tracker/cc/posix/memory_tracker.inc b/core/memory_tracker/cc/posix/memory_tracker.inc
index 28d4818107..fcc6fd06e5 100644
--- a/core/memory_tracker/cc/posix/memory_tracker.inc
+++ b/core/memory_tracker/cc/posix/memory_tracker.inc
@@ -14,26 +14,25 @@
  * limitations under the License.
  */
 
-
 #include <functional>
 namespace gapii {
 namespace track_memory {
 
 class PosixMemoryTracker {
-  public:
-    PosixMemoryTracker(std::function<bool(void*)> segfault_function):
-      signal_handler_registered_(false),
-      orig_action_{0},
-      handle_segfault_(segfault_function) {
-    }
-
+public:
+  PosixMemoryTracker(std::function<bool(void *)> segfault_function)
+      : orig_action_{0}, handle_segfault_(segfault_function) {}
 
-  bool IsInstalled() const { return signal_handler_registered_; }
+  bool IsInstalled() const {
+    struct sigaction orig_action = {0};
+    sigaction(SIGSEGV, nullptr, &orig_action);
+    return orig_action.sa_sigaction == &SegfaultHandlerFunction;
+  }
 
-  protected:
+protected:
   // A static wrapper of HandleSegfault() as sigaction() asks for a static
   // function.
-  static void SegfaultHandlerFunction(int sig, siginfo_t* info, void* unused);
+  static void SegfaultHandlerFunction(int sig, siginfo_t *info, void *unused);
 
   // EnableMemoryTrackerImpl calls sigaction() to register the new segfault
   // handler to the thread (and affects all the child threads), stores the
@@ -49,45 +48,45 @@ class PosixMemoryTracker {
   // othwerwise returns false.
   bool inline DisableMemoryTrackerImpl();
 
-  private:
-    bool signal_handler_registered_;  // A flag to indicate whether the signal
-                                      // handler has been registered
-    struct sigaction orig_action_;    // The original signal action for SIGSEGV
-    std::function<bool(void*)> handle_segfault_; // The function to call on a segfault
+private:
+  struct sigaction orig_action_; // The original signal action for SIGSEGV
+  std::function<bool(void *)>
+      handle_segfault_; // The function to call on a segfault
 };
 
 typedef MemoryTrackerImpl<PosixMemoryTracker> MemoryTracker;
-extern PosixMemoryTracker* unique_tracker;
+extern PosixMemoryTracker *unique_tracker;
 
-void inline PosixMemoryTracker::SegfaultHandlerFunction(int sig, siginfo_t* info, void* unused) {
+void inline PosixMemoryTracker::SegfaultHandlerFunction(int sig,
+                                                        siginfo_t *info,
+                                                        void *unused) {
   if (unique_tracker) {
     if (!unique_tracker->handle_segfault_(info->si_addr)) {
-      #ifndef NDEBUG
-        raise(SIGTRAP);
-      #endif // NDEBUG
-        (*unique_tracker->orig_action_.sa_sigaction)(sig, info, unused);
+#ifndef NDEBUG
+      raise(SIGTRAP);
+#endif // NDEBUG
+      (*unique_tracker->orig_action_.sa_sigaction)(sig, info, unused);
     }
   }
 }
 
 bool inline PosixMemoryTracker::EnableMemoryTrackerImpl() {
-  if (signal_handler_registered_) {
+  if (IsInstalled()) {
     return true;
   }
+
   unique_tracker = this;
   struct sigaction sa {
     0
   };
   sa.sa_flags = SA_SIGINFO;
   sigemptyset(&sa.sa_mask);
-  sa.sa_sigaction = SegfaultHandlerFunction;
-  signal_handler_registered_ = sigaction(SIGSEGV, &sa, &orig_action_) != 1;
-  return signal_handler_registered_;
+  sa.sa_sigaction = &SegfaultHandlerFunction;
+  return sigaction(SIGSEGV, &sa, &orig_action_) != 1;
 }
 
 bool inline PosixMemoryTracker::DisableMemoryTrackerImpl() {
-  if (signal_handler_registered_) {
-    signal_handler_registered_ = false;
+  if (IsInstalled()) {
     return sigaction(SIGSEGV, &orig_action_, nullptr) != 1;
   }
   return true;