From 98b56ab23dd4d3dc826f06906e6c51c9c9d4d52a Mon Sep 17 00:00:00 2001
From: Qi Wang <interwq@gwu.edu>
Date: Tue, 27 Nov 2018 12:38:47 -0800
Subject: [PATCH] Store the bin shard selection in TSD.

This avoids having to choose bin shard on the fly, also will allow flexible bin
binding for each thread.
---
 include/jemalloc/internal/bin.h       |  4 +---
 include/jemalloc/internal/bin_types.h | 17 +++++++++++++++++
 include/jemalloc/internal/tsd.h       |  5 +++--
 src/arena.c                           |  3 +--
 src/jemalloc.c                        | 11 ++++++++---
 5 files changed, 30 insertions(+), 10 deletions(-)
 create mode 100644 include/jemalloc/internal/bin_types.h

diff --git a/include/jemalloc/internal/bin.h b/include/jemalloc/internal/bin.h
index baa0acf3..f542c882 100644
--- a/include/jemalloc/internal/bin.h
+++ b/include/jemalloc/internal/bin.h
@@ -2,14 +2,12 @@
 #define JEMALLOC_INTERNAL_BIN_H
 
 #include "jemalloc/internal/bin_stats.h"
+#include "jemalloc/internal/bin_types.h"
 #include "jemalloc/internal/extent_types.h"
 #include "jemalloc/internal/extent_structs.h"
 #include "jemalloc/internal/mutex.h"
 #include "jemalloc/internal/sc.h"
 
-#define BIN_SHARDS_MAX (1 << EXTENT_BITS_BINSHARD_WIDTH)
-#define N_BIN_SHARDS_DEFAULT 1
-
 /*
  * A bin contains a set of extents that are currently being used for slab
  * allocations.
diff --git a/include/jemalloc/internal/bin_types.h b/include/jemalloc/internal/bin_types.h
new file mode 100644
index 00000000..3533606b
--- /dev/null
+++ b/include/jemalloc/internal/bin_types.h
@@ -0,0 +1,17 @@
+#ifndef JEMALLOC_INTERNAL_BIN_TYPES_H
+#define JEMALLOC_INTERNAL_BIN_TYPES_H
+
+#include "jemalloc/internal/sc.h"
+
+#define BIN_SHARDS_MAX (1 << EXTENT_BITS_BINSHARD_WIDTH)
+#define N_BIN_SHARDS_DEFAULT 1
+
+/* Used in TSD static initializer only. Real init in arena_bind(). */
+#define TSD_BINSHARDS_ZERO_INITIALIZER {{UINT8_MAX}}
+
+typedef struct tsd_binshards_s tsd_binshards_t;
+struct tsd_binshards_s {
+	uint8_t binshard[SC_NBINS];
+};
+
+#endif /* JEMALLOC_INTERNAL_BIN_TYPES_H */
diff --git a/include/jemalloc/internal/tsd.h b/include/jemalloc/internal/tsd.h
index 4dc2274a..00a9500b 100644
--- a/include/jemalloc/internal/tsd.h
+++ b/include/jemalloc/internal/tsd.h
@@ -3,6 +3,7 @@
 
 #include "jemalloc/internal/arena_types.h"
 #include "jemalloc/internal/assert.h"
+#include "jemalloc/internal/bin_types.h"
 #include "jemalloc/internal/jemalloc_internal_externs.h"
 #include "jemalloc/internal/prof_types.h"
 #include "jemalloc/internal/ql.h"
@@ -74,7 +75,7 @@ typedef void (*test_callback_t)(int *);
     O(iarena,			arena_t *,		arena_t *)	\
     O(arena,			arena_t *,		arena_t *)	\
     O(arenas_tdata,		arena_tdata_t *,	arena_tdata_t *)\
-    O(binshard,			unsigned,		unsigned)	\
+    O(binshards,		tsd_binshards_t,	tsd_binshards_t)\
     O(tcache,			tcache_t,		tcache_t)	\
     O(witness_tsd,              witness_tsd_t,		witness_tsdn_t)	\
     MALLOC_TEST_TSD
@@ -94,7 +95,7 @@ typedef void (*test_callback_t)(int *);
     NULL,								\
     NULL,								\
     NULL,								\
-    ((unsigned)-1),							\
+    TSD_BINSHARDS_ZERO_INITIALIZER,					\
     TCACHE_ZERO_INITIALIZER,						\
     WITNESS_TSD_INITIALIZER						\
     MALLOC_TEST_TSD_INITIALIZER						\
diff --git a/src/arena.c b/src/arena.c
index 7017bd7a..d34de859 100644
--- a/src/arena.c
+++ b/src/arena.c
@@ -1349,8 +1349,7 @@ arena_bin_choose_lock(tsdn_t *tsdn, arena_t *arena, szind_t binind,
 	if (tsdn_null(tsdn) || tsd_arena_get(tsdn_tsd(tsdn)) == NULL) {
 		*binshard = 0;
 	} else {
-		*binshard = tsd_binshard_get(tsdn_tsd(tsdn)) %
-		    bin_infos[binind].n_shards;
+		*binshard = tsd_binshardsp_get(tsdn_tsd(tsdn))->binshard[binind];
 	}
 	assert(*binshard < bin_infos[binind].n_shards);
 	bin = &arena->bins[binind].bin_shards[*binshard];
diff --git a/src/jemalloc.c b/src/jemalloc.c
index 1f7ed2eb..1620d0d3 100644
--- a/src/jemalloc.c
+++ b/src/jemalloc.c
@@ -379,9 +379,14 @@ arena_bind(tsd_t *tsd, unsigned ind, bool internal) {
 		tsd_iarena_set(tsd, arena);
 	} else {
 		tsd_arena_set(tsd, arena);
-		unsigned binshard = atomic_fetch_add_u(&arena->binshard_next, 1,
-		    ATOMIC_RELAXED) % BIN_SHARDS_MAX;
-		tsd_binshard_set(tsd, binshard);
+		unsigned shard = atomic_fetch_add_u(&arena->binshard_next, 1,
+		    ATOMIC_RELAXED);
+		tsd_binshards_t *bins = tsd_binshardsp_get(tsd);
+		for (unsigned i = 0; i < SC_NBINS; i++) {
+			assert(bin_infos[i].n_shards > 0 &&
+			    bin_infos[i].n_shards <= BIN_SHARDS_MAX);
+			bins->binshard[i] = shard % bin_infos[i].n_shards;
+		}
 	}
 }