From e72b84cec44852dd76365cb7e1bf691b56a8adfc Mon Sep 17 00:00:00 2001
From: Alexey Makhalov <amakhalov@vmware.com>
Date: Tue, 29 Aug 2017 21:10:08 +0000
Subject: [PATCH 2/2] malloc arena fix

---
 elf/dl-tunables.list |  5 ++++
 malloc/arena.c       | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 malloc/malloc.c      | 31 +++++++++++++++++++++++++
 malloc/malloc.h      |  1 +
 4 files changed, 101 insertions(+)

diff --git a/elf/dl-tunables.list b/elf/dl-tunables.list
index c188c6a..15a1a14 100644
--- a/elf/dl-tunables.list
+++ b/elf/dl-tunables.list
@@ -76,6 +76,11 @@ glibc {
       minval: 1
       security_level: SXID_IGNORE
     }
+    arena_stickiness {
+      type: SIZE_T
+      env_alias: MALLOC_ARENA_STICKINESS
+      security_level: SXID_IGNORE
+    }
     tcache_max {
       type: SIZE_T
       security_level: SXID_ERASE
diff --git a/malloc/arena.c b/malloc/arena.c
index dc14fae..f0edf2b 100644
--- a/malloc/arena.c
+++ b/malloc/arena.c
@@ -63,6 +63,12 @@ typedef struct _heap_info
   char pad[-6 * SIZE_SZ & MALLOC_ALIGN_MASK];
 } heap_info;
 
+typedef struct _arena_tracker
+{
+  mstate arena; /* Arena most recently tracked for growth. */
+  size_t growth; /* Current size in bytes. */
+} arena_tracker;
+
 /* Get a compile-time error if the heap_info padding is not correct
    to make alignment work as expected in sYSMALLOc.  */
 extern int sanity_check_heap_info_alignment[(sizeof (heap_info)
@@ -73,6 +79,8 @@ extern int sanity_check_heap_info_alignment[(sizeof (heap_info)
 
 static __thread mstate thread_arena attribute_tls_model_ie;
 
+static __thread arena_tracker thread_arena_tracker attribute_tls_model_ie;
+
 /* Arena free list.  free_list_lock synchronizes access to the
    free_list variable below, and the next_free and attached_threads
    members of struct malloc_state objects.  No other locks must be
@@ -236,6 +244,7 @@ TUNABLE_CALLBACK_FNDECL (set_perturb_byte, int32_t)
 TUNABLE_CALLBACK_FNDECL (set_trim_threshold, size_t)
 TUNABLE_CALLBACK_FNDECL (set_arena_max, size_t)
 TUNABLE_CALLBACK_FNDECL (set_arena_test, size_t)
+TUNABLE_CALLBACK_FNDECL (set_arena_stickiness, size_t)
 #if USE_TCACHE
 TUNABLE_CALLBACK_FNDECL (set_tcache_max, size_t)
 TUNABLE_CALLBACK_FNDECL (set_tcache_count, size_t)
@@ -327,6 +336,7 @@ ptmalloc_init (void)
   TUNABLE_GET (mmap_max, int32_t, TUNABLE_CALLBACK (set_mmaps_max));
   TUNABLE_GET (arena_max, size_t, TUNABLE_CALLBACK (set_arena_max));
   TUNABLE_GET (arena_test, size_t, TUNABLE_CALLBACK (set_arena_test));
+  TUNABLE_GET (arena_stickiness, size_t, TUNABLE_CALLBACK (set_arena_stickiness));
 #if USE_TCACHE
   TUNABLE_GET (tcache_max, size_t, TUNABLE_CALLBACK (set_tcache_max));
   TUNABLE_GET (tcache_count, size_t, TUNABLE_CALLBACK (set_tcache_count));
@@ -392,6 +402,13 @@ ptmalloc_init (void)
                     __libc_mallopt (M_MMAP_THRESHOLD, atoi (&envline[16]));
                 }
               break;
+            case 16:
+              if (!__builtin_expect (__libc_enable_secure, 0))
+                {
+                  if (memcmp (envline, "ARENA_STICKINESS", 16) == 0)
+                    __libc_mallopt (M_ARENA_STICKINESS, atoi (&envline[17]));
+                }
+              break;
             default:
               break;
             }
@@ -974,6 +991,53 @@ arena_get_retry (mstate ar_ptr, size_t bytes)
   return ar_ptr;
 }
 
+static void
+internal_function
+arena_stickiness_track_alloc (void *victim)
+{
+  if (!victim || chunk_is_mmapped (mem2chunk (victim)))
+    return;
+
+  if (thread_arena_tracker.arena != arena_for_chunk (mem2chunk (victim))) {
+    thread_arena_tracker.growth = 0;
+    thread_arena_tracker.arena = arena_for_chunk (mem2chunk (victim));
+  } else {
+    thread_arena_tracker.growth += chunksize (mem2chunk (victim));
+    if (thread_arena_tracker.growth >= mp_.arena_stickiness) {
+      /* Swtich thread to the next arena */
+      mstate replaced_arena = thread_arena;
+      mstate next_to_use = replaced_arena->next;
+
+      __libc_lock_lock (free_list_lock);
+      detach_arena (replaced_arena);
+#if 0
+      /* If this was the last attached thread for this arena, put the
+	 arena on the free list.  */
+      if (replaced_arena->attached_threads == 0)
+	{
+	  replaced_arena->next_free = free_list;
+	  free_list = replaced_arena;
+	}
+#endif
+      if (next_to_use->attached_threads == 0)
+        remove_from_free_list (next_to_use);
+      ++next_to_use->attached_threads;
+
+      __libc_lock_unlock (free_list_lock);
+      thread_arena = next_to_use;
+    }
+  }
+}
+
+/* chunk must be valid and not mmaped.  */
+static void
+internal_function
+arena_stickiness_track_free (mchunkptr chunk)
+{
+  if (thread_arena_tracker.arena == arena_for_chunk (chunk))
+    thread_arena_tracker.growth -= chunksize (chunk);
+}
+
 static void __attribute__ ((section ("__libc_thread_freeres_fn")))
 arena_thread_freeres (void)
 {
diff --git a/malloc/malloc.c b/malloc/malloc.c
index 54e406b..29787a5 100644
--- a/malloc/malloc.c
+++ b/malloc/malloc.c
@@ -1723,6 +1723,7 @@ struct malloc_par
   INTERNAL_SIZE_T mmap_threshold;
   INTERNAL_SIZE_T arena_test;
   INTERNAL_SIZE_T arena_max;
+  INTERNAL_SIZE_T arena_stickiness;
 
   /* Memory map support */
   int n_mmaps;
@@ -1787,6 +1788,7 @@ static struct malloc_par mp_ =
   .mmap_threshold = DEFAULT_MMAP_THRESHOLD,
   .trim_threshold = DEFAULT_TRIM_THRESHOLD,
 #define NARENAS_FROM_NCORES(n) ((n) * (sizeof (long) == 4 ? 2 : 8))
+  .arena_stickiness = 0,
   .arena_test = NARENAS_FROM_NCORES (1)
 #if USE_TCACHE
   ,
@@ -3083,6 +3085,10 @@ __libc_malloc (size_t bytes)
 
   assert (!victim || chunk_is_mmapped (mem2chunk (victim)) ||
           ar_ptr == arena_for_chunk (mem2chunk (victim)));
+
+  if (mp_.arena_stickiness > 0)
+    arena_stickiness_track_alloc (victim);
+
   return victim;
 }
 libc_hidden_def (__libc_malloc)
@@ -3126,6 +3132,9 @@ __libc_free (void *mem)
 
   MAYBE_INIT_TCACHE ();
 
+  if (mp_.arena_stickiness > 0)
+    arena_stickiness_track_free (p);
+
   ar_ptr = arena_for_chunk (p);
   _int_free (ar_ptr, p, 0);
 }
@@ -3226,6 +3235,8 @@ __libc_realloc (void *oldmem, size_t bytes)
       return newmem;
     }
 
+  if (mp_.arena_stickiness > 0)
+    arena_stickiness_track_free (oldp);
   __libc_lock_lock (ar_ptr->mutex);
 
   newp = _int_realloc (ar_ptr, oldp, oldsize, nb);
@@ -3234,6 +3245,9 @@ __libc_realloc (void *oldmem, size_t bytes)
   assert (!newp || chunk_is_mmapped (mem2chunk (newp)) ||
           ar_ptr == arena_for_chunk (mem2chunk (newp)));
 
+  if (mp_.arena_stickiness > 0)
+    arena_stickiness_track_alloc (newp);
+
   if (newp == NULL)
     {
       /* Try harder to allocate memory in other arenas.  */
@@ -3452,6 +3466,9 @@ __libc_calloc (size_t n, size_t elem_size)
       return mem;
     }
 
+  if (mp_.arena_stickiness > 0)
+    arena_stickiness_track_alloc (mem);
+
   csz = chunksize (p);
 
 #if MORECORE_CLEARS
@@ -5145,6 +5162,15 @@ do_set_arena_max (size_t value)
   return 1;
 }
 
+static inline int
+__always_inline
+do_set_arena_stickiness (size_t value)
+{
+  LIBC_PROBE (memory_mallopt_arena_stickiness, 2, value, mp_.arena_stickiness);
+  mp_.arena_stickiness = value;
+  return 1;
+}
+
 #if USE_TCACHE
 static inline int
 __always_inline
@@ -5237,6 +5263,11 @@ __libc_mallopt (int param_number, int value)
       if (value > 0)
 	do_set_arena_max (value);
       break;
+
+    case M_ARENA_STICKINESS:
+      if (value > 0)
+	do_set_arena_stickiness (value);
+      break;
     }
   __libc_lock_unlock (av->mutex);
   return res;
diff --git a/malloc/malloc.h b/malloc/malloc.h
index 339ab64..31bdb44 100644
--- a/malloc/malloc.h
+++ b/malloc/malloc.h
@@ -121,6 +121,7 @@ extern struct mallinfo mallinfo (void) __THROW;
 #define M_PERTURB           -6
 #define M_ARENA_TEST        -7
 #define M_ARENA_MAX         -8
+#define M_ARENA_STICKINESS  -9
 
 /* General SVID/XPG interface to tunable parameters. */
 extern int mallopt (int __param, int __val) __THROW;
-- 
2.9.3