Set jemalloc dirty_decay_ms and muzzy_decay_ms to 0 by default, add function to set the values to something else

wesm · wesm · commit daa541605e7c · 2019-10-21T12:48:30.000-05:00
diff --git a/cpp/src/arrow/memory_pool.cc b/cpp/src/arrow/memory_pool.cc
@@ -46,7 +46,7 @@
 // See discussion in https://github.com/jemalloc/jemalloc/issues/1621
 
 #ifdef NDEBUG
-const char* je_arrow_malloc_conf = "oversize_threshold:0";
+const char* je_arrow_malloc_conf = "oversize_threshold:0,dirty_decay_ms:0,muzzy_decay_ms:0";
 #else
 // In debug mode, add memory poisoning on alloc / free
 const char* je_arrow_malloc_conf = "oversize_threshold:0,junk:true";
@@ -389,6 +389,30 @@ MemoryPool* default_memory_pool() {
 #endif
 }
 
+#define RETURN_IF_JEMALLOC_ERROR(ERR)                  \
+  do {                                                 \
+    if (err != 0) {                                    \
+      return Status::UnknownError(std::strerror(ERR)); \
+    }                                                  \
+  } while (0)
+
+Status jemalloc_set_decay_ms(int ms) {
+#ifdef ARROW_JEMALLOC
+  ssize_t decay_time_ms = static_cast<ssize_t>(ms);
+
+  int err = mallctl("arenas.dirty_decay_ms", nullptr, nullptr, &decay_time_ms,
+                    sizeof(decay_time_ms));
+  RETURN_IF_JEMALLOC_ERROR(err);
+  err = mallctl("arenas.muzzy_decay_ms", nullptr, nullptr, &decay_time_ms,
+                sizeof(decay_time_ms));
+  RETURN_IF_JEMALLOC_ERROR(err);
+
+  return Status::OK();
+#else
+  return Status::Invalid("jemalloc support is not built");
+#endif
+}
+
 ///////////////////////////////////////////////////////////////////////
 // LoggingMemoryPool implementation
 
diff --git a/cpp/src/arrow/memory_pool.h b/cpp/src/arrow/memory_pool.h
@@ -160,6 +160,16 @@ ARROW_EXPORT MemoryPool* system_memory_pool();
 /// May return NotImplemented if jemalloc is not available.
 ARROW_EXPORT Status jemalloc_memory_pool(MemoryPool** out);
 
+/// \brief Set jemalloc memory page purging behavior for future-created arenas
+/// to the indicated number of milliseconds. See dirty_decay_ms and
+/// muzzy_decay_ms options in jemalloc for a description of what these do. The
+/// default is configured to 0 which releases memory more aggressively to the
+/// operating system. If you have a long-running application that uses a lot of
+/// memory, you may wish to set this to a higher value. The jemalloc default is
+/// 10000 (10 seconds)
+ARROW_EXPORT
+Status jemalloc_set_decay_ms(int ms);
+
 /// Return a process-wide memory pool based on mimalloc.
 ///
 /// May return NotImplemented if mimalloc is not available.
diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py
@@ -108,7 +108,8 @@ def parse_git(root, **kwargs):
 from pyarrow.lib import (MemoryPool, LoggingMemoryPool, ProxyMemoryPool,
                          total_allocated_bytes, set_memory_pool,
                          default_memory_pool, logging_memory_pool,
-                         proxy_memory_pool, log_memory_allocations)
+                         proxy_memory_pool, log_memory_allocations,
+                         jemalloc_set_decay_ms)
 
 # I/O
 from pyarrow.lib import (HdfsFile, NativeFile, PythonFile,
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
@@ -260,6 +260,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
     cdef CMemoryPool* c_default_memory_pool" arrow::default_memory_pool"()
 
+    CStatus c_jemalloc_set_decay_ms" arrow::jemalloc_set_decay_ms"(int ms)
+
     cdef cppclass CListType" arrow::ListType"(CDataType):
         CListType(const shared_ptr[CDataType]& value_type)
         CListType(const shared_ptr[CField]& field)
diff --git a/python/pyarrow/memory.pxi b/python/pyarrow/memory.pxi
@@ -151,3 +151,21 @@ def total_allocated_bytes():
     """
     cdef CMemoryPool* pool = c_get_memory_pool()
     return pool.bytes_allocated()
+
+
+def jemalloc_set_decay_ms(decay_ms):
+    """
+    Set arenas.dirty_decay_ms and arenas.muzzy_decay_ms to indicated number of
+    milliseconds. A value of 0 (the default) results in dirty / muzzy memory
+    pages being released right away to the OS, while a higher value will result
+    in a time-based decay. See the jemalloc docs for more information
+
+    It's best to set this at the start of your application.
+
+    Parameters
+    ----------
+    decay_ms : int
+        Number of milliseconds to set for jemalloc decay conf parameters. Note
+        that this change will only affect future memory arenas
+    """
+    check_status(c_jemalloc_set_decay_ms(decay_ms))
