Track memory usage. Add new CMake option TRACK_MEMORY_USAGE (It is ON by default).

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/CMakeLists.txt

@@ -6,6 +6,8 @@ set(LEAN_VERSION_MINOR 1)
 set(CMAKE_COLOR_MAKEFILE ON)
 enable_testing()
 
+option(TRACK_MEMORY_USAGE "TRACK_MEMORY_USAGE" ON)
+
 # Added for CTest
 INCLUDE(CTest)
 CONFIGURE_FILE(${LEAN_SOURCE_DIR}/CTestCustom.cmake.in
@@ -57,6 +59,11 @@ set(EXTRA_LIBS ${EXTRA_LIBS} ${GMP_LIBRARIES})
 find_package(MPFR 3.1.0)
 set(EXTRA_LIBS ${EXTRA_LIBS} ${MPFR_LIBRARIES})
 
+# TRACK_MEMORY_USAGE
+if("${TRACK_MEMORY_USAGE}" MATCHES "ON")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D LEAN_TRACK_MEMORY")
+endif()
+
 # tcmalloc
 option(TCMALLOC "TCMALLOC" ON)
 if("${TCMALLOC}" MATCHES "ON")
@@ -64,6 +71,7 @@ if("${TCMALLOC}" MATCHES "ON")
   if(${TCMALLOC_FOUND})
     set(EXTRA_LIBS ${EXTRA_LIBS} ${TCMALLOC_LIBRARIES})
     message(STATUS "Using tcmalloc.")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D HAS_TCMALLOC")
   else()
     message(WARNING "FAILED to find tcmalloc, using standard malloc.")
   endif()
@@ -80,6 +88,12 @@ if("${READLINE}" MATCHES "ON")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DLEAN_USE_READLINE")
 endif()
 
+# Check malloc_usable_size
+if(NOT "${TCMALLOC_FOUND}" AND "${TRACK_MEMORY_USAGE}" MATCHES "ON")
+  find_package(MallocUsableSize)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I ${MALLOC_DIR} -D HAS_MALLOC_USABLE_SIZE")
+endif()
+
 include_directories(${LEAN_SOURCE_DIR})
 
 add_subdirectory(util)

src/cmake/Modules/CheckMallocUsableSize.cc

@@ -0,0 +1,25 @@
+/*
+Copyright (c) 2013 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+
+#include <stdio.h>
+#include <string.h>
+#include <malloc.h>
+
+#define MAX_SZ 1024
+
+int main() {
+    for (unsigned i = 1; i < MAX_SZ; i++) {
+        void * p = malloc(i);
+        size_t r = malloc_usable_size(p);
+        if (r < i || r > 2*(i + 8)) {
+            fprintf(stderr, "Unexpected malloc_usable_size behavior");
+            return 0;
+        }
+        free(p);
+    }
+    return 1;
+}

src/cmake/Modules/FindMallocUsableSize.cmake

@@ -0,0 +1,16 @@
+find_path(MALLOC_DIR NAMES malloc.h )
+
+try_run(MUS_CHECK MUS_CHECK_BUILD
+  ${LEAN_BINARY_DIR}${CMAKE_FILES_DIRECTORY}/CMakeTmp
+  ${LEAN_SOURCE_DIR}/cmake/Modules/CheckMallocUsableSize.cc
+  CMAKE_FLAGS -DINCLUDE_DIRECTORIES=${MALLOC_DIR}
+  RUN_OUTPUT_VARIABLE MUS_TRY_OUT)
+
+if("${MUS_CHECK_BUILD}" MATCHES "TRUE" AND "${MUS_CHECK}" MATCHES "0")
+  message(STATUS "Found malloc_usable_size")
+else()
+  message(STATUS "Usable malloc_usable_size was not detected")
+endif()
+
+
+

src/util/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_library(util trace.cpp debug.cpp name.cpp exception.cpp
-  hash.cpp escaped.cpp bit_tricks.cpp safe_arith.cpp ascii.cpp)
+  hash.cpp escaped.cpp bit_tricks.cpp safe_arith.cpp ascii.cpp
+  memory.cpp)

src/util/memory.cpp

@@ -0,0 +1,146 @@
+/*
+Copyright (c) 2013 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#include <new>
+#include <cstdlib>
+
+#if !defined(LEAN_TRACK_MEMORY)
+
+namespace lean {
+size_t     get_allocated_memory() {
+    return 0;
+}
+
+long long  get_thread_allocated_memory() {
+    return 0;
+}
+
+void * malloc(size_t sz)  {
+    void * r = ::malloc(sz);
+    if (r || sz == 0)
+        return r;
+    else
+        throw std::bad_alloc();
+}
+
+void * realloc(void * ptr, size_t sz) {
+    void * r = ::realloc(ptr, sz);
+    if (r || sz == 0)
+        return r;
+    else
+        throw std::bad_alloc();
+}
+
+void free(void * ptr) {
+    ::free(ptr);
+}
+}
+
+#else
+#include <atomic>
+
+#if defined(HAS_MALLOC_USABLE_SIZE)
+#include <malloc.h>
+namespace lean {
+inline size_t malloc_size(void * ptr)             { return malloc_usable_size(ptr); }
+inline void * malloc_core(size_t sz)              { return ::malloc(sz); }
+inline void * realloc_core(void * ptr, size_t sz) { return realloc(ptr, sz); }
+inline void   free_core(void * ptr)               { free(ptr); }
+}
+// REMARK: We commented the following piece of code because tc_malloc_size is hanging
+
+// #elif defined(HAS_TCMALLOC)
+// #include <gperftools/tcmalloc.h>
+// namespace lean {
+// inline size_t malloc_size(void * ptr)             { return tc_malloc_size(ptr); }
+// inline void * malloc_core(size_t sz)              { return tc_malloc(sz); }
+// inline void * realloc_core(void * ptr, size_t sz) { return tc_realloc(ptr, sz); }
+// inline void   free_core(void * ptr)               { tc_free(ptr); }
+// }
+#else
+namespace lean {
+void * save_alloc_size(void * ptr, size_t sz) {
+    if (ptr) {
+        size_t * r = static_cast<size_t*>(ptr);
+        *r = sz;
+        return r + 1;
+    } else {
+        return ptr;
+    }
+}
+inline size_t malloc_size(void * ptr)             { return static_cast<size_t*>(ptr)[-1]; }
+inline void * malloc_core(size_t sz)              { return save_alloc_size(::malloc(sz + sizeof(size_t)), sz); }
+inline void * realloc_core(void * ptr, size_t sz) { return save_alloc_size(::realloc(static_cast<size_t*>(ptr) - 1, sz + sizeof(size_t)), sz); }
+inline void   free_core(void * ptr)               { ::free(static_cast<size_t*>(ptr) - 1); }
+}
+#endif
+
+namespace lean {
+class alloc_info {
+    std::atomic<size_t> m_size;
+public:
+    alloc_info():m_size(0) {}
+    ~alloc_info() {}
+    void inc(size_t sz) { m_size += sz; }
+    void dec(size_t sz) { m_size -= sz; }
+    size_t size() const { return m_size; }
+};
+
+class thread_alloc_info {
+    size_t m_size; // It can be negative
+public:
+    thread_alloc_info():m_size(0) {}
+    void inc(size_t sz) { m_size += sz; }
+    void dec(size_t sz) { m_size -= sz; }
+    long long size() const { return static_cast<long long>(m_size); }
+};
+
+static alloc_info                     g_global_memory;
+static thread_local thread_alloc_info g_thread_memory;
+
+size_t     get_allocated_memory() { return g_global_memory.size(); }
+long long  get_thread_allocated_memory() { return g_thread_memory.size(); }
+
+void * malloc(size_t sz)  {
+    void * r = malloc_core(sz);
+    if (r || sz == 0) {
+        size_t rsz = malloc_size(r);
+        g_global_memory.inc(rsz);
+        g_thread_memory.inc(rsz);
+        return r;
+    } else {
+        throw std::bad_alloc();
+    }
+}
+
+void * realloc(void * ptr, size_t sz) {
+    size_t old_sz = malloc_size(ptr);
+    g_global_memory.dec(old_sz);
+    g_global_memory.inc(sz);
+    g_thread_memory.dec(old_sz);
+    g_thread_memory.inc(sz);
+    void * r = realloc_core(ptr, sz);
+    if (r || sz == 0)
+        return r;
+    else
+        throw std::bad_alloc();
+}
+
+void free(void * ptr) {
+    if (ptr) {
+        size_t sz = malloc_size(ptr);
+        g_global_memory.dec(sz);
+        g_thread_memory.dec(sz);
+    }
+    free_core(ptr);
+}
+}
+
+void* operator new(std::size_t sz) throw(std::bad_alloc) { return lean::malloc(sz); }
+void  operator delete(void * ptr) throw() { return lean::free(ptr); }
+void* operator new[](std::size_t sz) throw(std::bad_alloc) { return lean::malloc(sz); }
+void  operator delete[](void * ptr) throw() { return lean::free(ptr); }
+#endif

src/util/memory.h

@@ -0,0 +1,19 @@
+/*
+Copyright (c) 2013 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#pragma once
+
+namespace lean {
+size_t get_allocated_memory();
+long long get_thread_allocated_memory();
+void * malloc(size_t sz);
+void * realloc(void * ptr, size_t sz);
+void free(void * ptr);
+}
+
+
+
+

src/util/numerics/gmp_init.cpp

@@ -5,28 +5,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 */
 #include <gmp.h>
-#include <new>
-#include <cstdlib>
+#include "util/memory.h"
 
-extern "C" void * cxx_malloc(size_t size) {
-    void * p = malloc(size);
-    if (p != 0 || size == 0)
-        return p;
-
-    throw std::bad_alloc();
-}
-
-extern "C" void * cxx_realloc(void * q, size_t, size_t new_size) {
-    void* p = realloc(q, new_size);
-    if (p != 0 || new_size == 0)
-        return p;
-
-    throw std::bad_alloc();
-}
-
-extern "C" void cxx_free(void * p, size_t) {
-    free(p);
-}
+extern "C" void * cxx_malloc(size_t size) { return lean::malloc(size); }
+extern "C" void * cxx_realloc(void * q, size_t, size_t new_size) { return lean::realloc(q, new_size); }
+extern "C" void cxx_free(void * p, size_t) { return lean::free(p); }
 
 /**
    \brief Auxiliary class for initializing the GMP memory allocation