Merge pull request #344 from MartinNowack/feat_malloc

Add support for tcmalloc

4 files changed, 48 insertions, 40 deletions

diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index 854754b0..33cee9f9 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -2437,6 +2437,39 @@ void Executor::bindModuleConstants() {
   }
 }
 
+void Executor::checkMemoryUsage() {
+  if (!MaxMemory)
+    return;
+  if ((stats::instructions & 0xFFFF) == 0) {
+    // We need to avoid calling GetTotalMallocUsage() often because it
+    // is O(elts on freelist). This is really bad since we start
+    // to pummel the freelist once we hit the memory cap.
+    unsigned mbs = util::GetTotalMallocUsage() >> 20;
+    if (mbs > MaxMemory) {
+      if (mbs > MaxMemory + 100) {
+        // just guess at how many to kill
+        unsigned numStates = states.size();
+        unsigned toKill = std::max(1U, numStates - numStates * MaxMemory / mbs);
+        klee_warning("killing %d states (over memory cap)", toKill);
+        std::vector<ExecutionState *> arr(states.begin(), states.end());
+        for (unsigned i = 0, N = arr.size(); N && i < toKill; ++i, --N) {
+          unsigned idx = rand() % N;
+          // Make two pulls to try and not hit a state that
+          // covered new code.
+          if (arr[idx]->coveredNew)
+            idx = rand() % N;
+
+          std::swap(arr[idx], arr[N - 1]);
+          terminateStateEarly(*arr[N - 1], "Memory limit exceeded.");
+        }
+      }
+      atMemoryLimit = true;
+    } else {
+      atMemoryLimit = false;
+    }
+  }
+}
+
 void Executor::run(ExecutionState &initialState) {
   bindModuleConstants();
 
@@ -2522,39 +2555,7 @@ void Executor::run(ExecutionState &initialState) {
     executeInstruction(state, ki);
     processTimers(&state, MaxInstructionTime);
 
-    if (MaxMemory) {
-      if ((stats::instructions & 0xFFFF) == 0) {
-        // We need to avoid calling GetMallocUsage() often because it
-        // is O(elts on freelist). This is really bad since we start
-        // to pummel the freelist once we hit the memory cap.
-        unsigned mbs = util::GetTotalMallocUsage() >> 20;
-        if (mbs > MaxMemory) {
-          if (mbs > MaxMemory + 100) {
-            // just guess at how many to kill
-            unsigned numStates = states.size();
-            unsigned toKill = std::max(1U, numStates - numStates*MaxMemory/mbs);
-
-            klee_warning("killing %d states (over memory cap)", toKill);
-
-            std::vector<ExecutionState*> arr(states.begin(), states.end());
-            for (unsigned i=0,N=arr.size(); N && i<toKill; ++i,--N) {
-              unsigned idx = rand() % N;
-
-              // Make two pulls to try and not hit a state that
-              // covered new code.
-              if (arr[idx]->coveredNew)
-                idx = rand() % N;
-
-              std::swap(arr[idx], arr[N-1]);
-              terminateStateEarly(*arr[N-1], "Memory limit exceeded.");
-            }
-          }
-          atMemoryLimit = true;
-        } else {
-          atMemoryLimit = false;
-        }
-      }
-    }
+    checkMemoryUsage();
 
     updateStates(&state);
   }
diff --git a/lib/Core/Executor.h b/lib/Core/Executor.h
index 919d2124..8bfa278a 100644
--- a/lib/Core/Executor.h
+++ b/lib/Core/Executor.h
@@ -399,7 +399,8 @@ private:
   void initTimers();
   void processTimers(ExecutionState *current,
                      double maxInstTime);
-                
+  void checkMemoryUsage();
+
 public:
   Executor(const InterpreterOptions &opts, InterpreterHandler *ie);
   virtual ~Executor();
diff --git a/lib/Core/StatsTracker.cpp b/lib/Core/StatsTracker.cpp
index 2e107fb3..9995b7e2 100644
--- a/lib/Core/StatsTracker.cpp
+++ b/lib/Core/StatsTracker.cpp
@@ -16,6 +16,7 @@
 #include "klee/Internal/Module/KModule.h"
 #include "klee/Internal/Module/KInstruction.h"
 #include "klee/Internal/Support/ModuleUtil.h"
+#include "klee/Internal/System/MemoryUsage.h"
 #include "klee/Internal/System/Time.h"
 #include "klee/Internal/Support/ErrorHandling.h"
 #include "klee/SolverStats.h"
@@ -404,11 +405,7 @@ void StatsTracker::writeStatsLine() {
              << "," << numBranches
              << "," << util::getUserTime()
              << "," << executor.states.size()
-#if LLVM_VERSION_CODE > LLVM_VERSION(3, 2)
-             << "," << sys::Process::GetMallocUsage()
-#else
-             << "," << sys::Process::GetTotalMemoryUsage()
-#endif
+             << "," << util::GetTotalMallocUsage()
              << "," << stats::queries
              << "," << stats::queryConstructs
              << "," << 0 // was numObjects
diff --git a/lib/Support/MemoryUsage.cpp b/lib/Support/MemoryUsage.cpp
index 32a7eb3b..a9f4026d 100644
--- a/lib/Support/MemoryUsage.cpp
+++ b/lib/Support/MemoryUsage.cpp
@@ -11,6 +11,10 @@
 
 #include "klee/Config/config.h"
 
+#ifdef HAVE_GPERFTOOLS_MALLOC_EXTENSION_H
+#include "gperftools/malloc_extension.h"
+#endif
+
 #ifdef HAVE_MALLINFO
 #include <malloc.h>
 #endif
@@ -21,7 +25,12 @@
 using namespace klee;
 
 size_t util::GetTotalMallocUsage() {
-#ifdef HAVE_MALLINFO
+#ifdef HAVE_GPERFTOOLS_MALLOC_EXTENSION_H
+  uint64_t value;
+  MallocExtension::instance()->GetNumericProperty(
+      "generic.current_allocated_bytes", &value);
+  return value;
+#elif HAVE_MALLINFO
   struct mallinfo mi = ::mallinfo();
   // The malloc implementation in glibc (pmalloc2)
   // does not include mmap()'ed memory in mi.uordblks