From 4fc5c4fbac156c9f44452d3f30216451711dfa18 Mon Sep 17 00:00:00 2001
From: David Carlier <devnexen@gmail.com>
Date: Wed, 1 Jun 2022 22:04:11 +0100
Subject: [PATCH] New configure option '--enable-pageid' for Linux

The option makes jemalloc use prctl with PR_SET_VMA to tag memory mappings with
"jemalloc_pg" or "jemalloc_pg_overcommit". This allows to easily identify
jemalloc's mappings in /proc/<pid>/maps. PR_SET_VMA is only available in Linux
5.17 and above.
---
 configure.ac                                  | 25 +++++++++++++++++
 .../internal/jemalloc_internal_defs.h.in      |  6 ++++
 src/pages.c                                   | 28 +++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/configure.ac b/configure.ac
index 917d9a80..0ae579ee 100644
--- a/configure.ac
+++ b/configure.ac
@@ -2067,6 +2067,14 @@ if test "x$have_memcntl" = "x1" ; then
   AC_DEFINE([JEMALLOC_HAVE_MEMCNTL], [ ], [ ])
 fi
 
+AC_CHECK_FUNC([prctl],
+	      [have_prctl="1"],
+	      [have_prctl="0"],
+	      )
+if test "x$have_prctl" = "x1" ; then
+  AC_DEFINE([JEMALLOC_HAVE_PRCTL], [ ], [ ])
+fi
+
 dnl Disable lazy locking by default.
 AC_ARG_ENABLE([lazy_lock],
   [AS_HELP_STRING([--enable-lazy-lock],
@@ -2435,6 +2443,22 @@ else
   AC_DEFINE([JEMALLOC_TLS_MODEL], [ ], [ ])
 fi
 
+dnl Do not compile with debugging by default.
+AC_ARG_ENABLE([pageid],
+  [AS_HELP_STRING([--enable-pageid],
+                  [Enable named pages])],
+[if test "x$enable_pageid" = "xno" ; then
+  enable_pageid="0"
+else
+  enable_pageid="1"
+fi
+],
+[enable_pageid="0"]
+)
+if test "x$enable_pageid" = "x1" ; then
+  AC_DEFINE([JEMALLOC_PAGEID], [ ], [ ])
+fi
+
 dnl ============================================================================
 dnl Enable background threads if possible.
 
@@ -2691,5 +2715,6 @@ AC_MSG_RESULT([xmalloc            : ${enable_xmalloc}])
 AC_MSG_RESULT([log                : ${enable_log}])
 AC_MSG_RESULT([lazy_lock          : ${enable_lazy_lock}])
 AC_MSG_RESULT([cache-oblivious    : ${enable_cache_oblivious}])
+AC_MSG_RESULT([pageid             : ${enable_pageid}])
 AC_MSG_RESULT([cxx                : ${enable_cxx}])
 AC_MSG_RESULT([===============================================================================])
diff --git a/include/jemalloc/internal/jemalloc_internal_defs.h.in b/include/jemalloc/internal/jemalloc_internal_defs.h.in
index 888ef470..6dbd8780 100644
--- a/include/jemalloc/internal/jemalloc_internal_defs.h.in
+++ b/include/jemalloc/internal/jemalloc_internal_defs.h.in
@@ -162,6 +162,12 @@
 /* Use gcc intrinsics for profile backtracing if defined. */
 #undef JEMALLOC_PROF_GCC
 
+/* JEMALLOC_PAGEID enabled page id */
+#undef JEMALLOC_PAGEID
+
+/* JEMALLOC_HAVE_PRCTL checks prctl */
+#undef JEMALLOC_HAVE_PRCTL
+
 /*
  * JEMALLOC_DSS enables use of sbrk(2) to allocate extents from the data storage
  * segment (DSS).
diff --git a/src/pages.c b/src/pages.c
index 8c83a7de..b672e4de 100644
--- a/src/pages.c
+++ b/src/pages.c
@@ -21,6 +21,13 @@
 #else
 #define PAGES_FD_TAG -1
 #endif
+#ifdef JEMALLOC_HAVE_PRCTL
+#include <sys/prctl.h>
+#ifndef PR_SET_VMA
+#define PR_SET_VMA 0x53564d41
+#define PR_SET_VMA_ANON_NAME 0
+#endif
+#endif
 
 /******************************************************************************/
 /* Data. */
@@ -98,6 +105,22 @@ static int madvise_MADV_DONTNEED_zeroes_pages()
 }
 #endif
 
+#ifdef JEMALLOC_PAGEID
+static int os_page_id(void *addr, size_t size, const char *name)
+{
+#ifdef JEMALLOC_HAVE_PRCTL
+	/*
+	 * While parsing `/proc/<pid>/maps` file, the block could appear as
+	 * 7f4836000000-7f4836800000 rw-p 00000000 00:00 0 [anon:jemalloc_pg_overcommit]`
+	 */
+	return prctl(PR_SET_VMA, PR_SET_VMA_ANON_NAME, (uintptr_t)addr, size,
+	    (uintptr_t)name);
+#else
+	return 0;
+#endif
+}
+#endif
+
 /******************************************************************************/
 /*
  * Function prototypes for static functions that are referenced prior to
@@ -162,6 +185,11 @@ os_pages_map(void *addr, size_t size, size_t alignment, bool *commit) {
 #endif
 	assert(ret == NULL || (addr == NULL && ret != addr) || (addr != NULL &&
 	    ret == addr));
+#ifdef JEMALLOC_PAGEID
+	int n = os_page_id(ret, size,
+	    os_overcommits ? "jemalloc_pg_overcommit" : "jemalloc_pg");
+	assert(n == 0 || (n == -1 && get_errno() == EINVAL));
+#endif
 	return ret;
 }