Pass alloc_ctx down profiling path.

With this change, when profiling is enabled, we avoid doing redundant rtree
lookups. Also changed dalloc_atx_t to alloc_atx_t, as it's now used on
allocation path as well (to speed up profiling).

5 files changed, 67 insertions, 47 deletions

diff --git a/include/jemalloc/internal/arena_inlines_b.h b/include/jemalloc/internal/arena_inlines_b.h
index 382289e..8c76e0b 100644
--- a/include/jemalloc/internal/arena_inlines_b.h
+++ b/include/jemalloc/internal/arena_inlines_b.h
@@ -3,9 +3,10 @@
 
 #ifndef JEMALLOC_ENABLE_INLINE
 szind_t arena_bin_index(arena_t *arena, arena_bin_t *bin);
-prof_tctx_t *arena_prof_tctx_get(tsdn_t *tsdn, const void *ptr);
+prof_tctx_t *arena_prof_tctx_get(tsdn_t *tsdn, const void *ptr,
+    alloc_ctx_t *ctx);
 void arena_prof_tctx_set(tsdn_t *tsdn, const void *ptr, size_t usize,
-    prof_tctx_t *tctx);
+    alloc_ctx_t *ctx, prof_tctx_t *tctx);
 void arena_prof_tctx_reset(tsdn_t *tsdn, const void *ptr, prof_tctx_t *tctx);
 void arena_decay_ticks(tsdn_t *tsdn, arena_t *arena, unsigned nticks);
 void arena_decay_tick(tsdn_t *tsdn, arena_t *arena);
@@ -16,10 +17,10 @@ size_t arena_salloc(tsdn_t *tsdn, const void *ptr);
 size_t arena_vsalloc(tsdn_t *tsdn, const void *ptr);
 void arena_dalloc_no_tcache(tsdn_t *tsdn, void *ptr);
 void arena_dalloc(tsdn_t *tsdn, void *ptr, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool slow_path);
+    alloc_ctx_t *alloc_ctx, bool slow_path);
 void arena_sdalloc_no_tcache(tsdn_t *tsdn, void *ptr, size_t size);
 void arena_sdalloc(tsdn_t *tsdn, void *ptr, size_t size, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool slow_path);
+    alloc_ctx_t *alloc_ctx, bool slow_path);
 #endif
 
 #if (defined(JEMALLOC_ENABLE_INLINE) || defined(JEMALLOC_ARENA_C_))
@@ -30,27 +31,41 @@ arena_bin_index(arena_t *arena, arena_bin_t *bin) {
 	return binind;
 }
 
-JEMALLOC_INLINE prof_tctx_t *
-arena_prof_tctx_get(tsdn_t *tsdn, const void *ptr) {
+JEMALLOC_ALWAYS_INLINE prof_tctx_t *
+arena_prof_tctx_get(tsdn_t *tsdn, const void *ptr, alloc_ctx_t *alloc_ctx) {
 	cassert(config_prof);
 	assert(ptr != NULL);
 
-	const extent_t *extent = iealloc(tsdn, ptr);
-	if (unlikely(!extent_slab_get(extent))) {
-		return large_prof_tctx_get(tsdn, extent);
+	/* Static check. */
+	if (alloc_ctx == NULL) {
+		const extent_t *extent = iealloc(tsdn, ptr);
+		if (unlikely(!extent_slab_get(extent))) {
+			return large_prof_tctx_get(tsdn, extent);
+		}
+	} else {
+		if (unlikely(!alloc_ctx->slab)) {
+			return large_prof_tctx_get(tsdn, iealloc(tsdn, ptr));
+		}
 	}
 	return (prof_tctx_t *)(uintptr_t)1U;
 }
 
-JEMALLOC_INLINE void
+JEMALLOC_ALWAYS_INLINE void
 arena_prof_tctx_set(tsdn_t *tsdn, const void *ptr, size_t usize,
-    prof_tctx_t *tctx) {
+    alloc_ctx_t *alloc_ctx, prof_tctx_t *tctx) {
 	cassert(config_prof);
 	assert(ptr != NULL);
 
-	extent_t *extent = iealloc(tsdn, ptr);
-	if (unlikely(!extent_slab_get(extent))) {
-		large_prof_tctx_set(tsdn, extent, tctx);
+	/* Static check. */
+	if (alloc_ctx == NULL) {
+		extent_t *extent = iealloc(tsdn, ptr);
+		if (unlikely(!extent_slab_get(extent))) {
+			large_prof_tctx_set(tsdn, extent, tctx);
+		}
+	} else {
+		if (unlikely(!alloc_ctx->slab)) {
+			large_prof_tctx_set(tsdn, iealloc(tsdn, ptr), tctx);
+		}
 	}
 }
 
@@ -196,7 +211,7 @@ arena_dalloc_no_tcache(tsdn_t *tsdn, void *ptr) {
 
 JEMALLOC_ALWAYS_INLINE void
 arena_dalloc(tsdn_t *tsdn, void *ptr, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool slow_path) {
+    alloc_ctx_t *alloc_ctx, bool slow_path) {
 	assert(!tsdn_null(tsdn) || tcache == NULL);
 	assert(ptr != NULL);
 
@@ -208,9 +223,9 @@ arena_dalloc(tsdn_t *tsdn, void *ptr, tcache_t *tcache,
 	szind_t szind;
 	bool slab;
 	rtree_ctx_t *rtree_ctx;
-	if (dalloc_ctx != NULL) {
-		szind = dalloc_ctx->szind;
-		slab = dalloc_ctx->slab;
+	if (alloc_ctx != NULL) {
+		szind = alloc_ctx->szind;
+		slab = alloc_ctx->slab;
 		assert(szind != NSIZES);
 	} else {
 		rtree_ctx = tsd_rtree_ctx(tsdn_tsd(tsdn));
@@ -293,7 +308,7 @@ arena_sdalloc_no_tcache(tsdn_t *tsdn, void *ptr, size_t size) {
 
 JEMALLOC_ALWAYS_INLINE void
 arena_sdalloc(tsdn_t *tsdn, void *ptr, size_t size, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool slow_path) {
+    alloc_ctx_t *alloc_ctx, bool slow_path) {
 	assert(!tsdn_null(tsdn) || tcache == NULL);
 	assert(ptr != NULL);
 	assert(size <= LARGE_MAXCLASS);
@@ -305,9 +320,9 @@ arena_sdalloc(tsdn_t *tsdn, void *ptr, size_t size, tcache_t *tcache,
 
 	szind_t szind;
 	bool slab;
-	UNUSED dalloc_ctx_t local_ctx;
+	UNUSED alloc_ctx_t local_ctx;
 	if (config_prof && opt_prof) {
-		if (dalloc_ctx == NULL) {
+		if (alloc_ctx == NULL) {
 			/* Uncommon case and should be a static check. */
 			rtree_ctx_t rtree_ctx_fallback;
 			rtree_ctx_t *rtree_ctx = tsdn_rtree_ctx(tsdn,
@@ -316,10 +331,10 @@ arena_sdalloc(tsdn_t *tsdn, void *ptr, size_t size, tcache_t *tcache,
 			    (uintptr_t)ptr, true, &local_ctx.szind,
 			    &local_ctx.slab);
 			assert(local_ctx.szind == size2index(size));
-			dalloc_ctx = &local_ctx;
+			alloc_ctx = &local_ctx;
 		}
-		slab = dalloc_ctx->slab;
-		szind = dalloc_ctx->szind;
+		slab = alloc_ctx->slab;
+		szind = alloc_ctx->szind;
 	} else {
 		/*
 		 * There is no risk of being confused by a promoted sampled
diff --git a/include/jemalloc/internal/arena_structs_b.h b/include/jemalloc/internal/arena_structs_b.h
index 1f4fb6b..1370b53 100644
--- a/include/jemalloc/internal/arena_structs_b.h
+++ b/include/jemalloc/internal/arena_structs_b.h
@@ -260,8 +260,8 @@ struct arena_tdata_s {
 	ticker_t		decay_ticker;
 };
 
-/* Used to pass rtree lookup context down the deallocation path. */
-struct dalloc_ctx_s {
+/* Used to pass rtree lookup context down the path. */
+struct alloc_ctx_s {
 	szind_t szind;
 	bool slab;
 };
diff --git a/include/jemalloc/internal/arena_types.h b/include/jemalloc/internal/arena_types.h
index 435b930..e243aab 100644
--- a/include/jemalloc/internal/arena_types.h
+++ b/include/jemalloc/internal/arena_types.h
@@ -19,7 +19,7 @@ typedef struct arena_decay_s arena_decay_t;
 typedef struct arena_bin_s arena_bin_t;
 typedef struct arena_s arena_t;
 typedef struct arena_tdata_s arena_tdata_t;
-typedef struct dalloc_ctx_s dalloc_ctx_t;
+typedef struct alloc_ctx_s alloc_ctx_t;
 
 typedef enum {
 	percpu_arena_disabled = 0,
diff --git a/include/jemalloc/internal/jemalloc_internal_inlines_c.h b/include/jemalloc/internal/jemalloc_internal_inlines_c.h
index 4fb3424..7884a20 100644
--- a/include/jemalloc/internal/jemalloc_internal_inlines_c.h
+++ b/include/jemalloc/internal/jemalloc_internal_inlines_c.h
@@ -15,10 +15,10 @@ void *ipalloct(tsdn_t *tsdn, size_t usize, size_t alignment, bool zero,
 void *ipalloc(tsd_t *tsd, size_t usize, size_t alignment, bool zero);
 size_t ivsalloc(tsdn_t *tsdn, const void *ptr);
 void idalloctm(tsdn_t *tsdn, void *ptr, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool is_internal, bool slow_path);
+    alloc_ctx_t *alloc_ctx, bool is_internal, bool slow_path);
 void idalloc(tsd_t *tsd, void *ptr);
 void isdalloct(tsdn_t *tsdn, void *ptr, size_t size, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool slow_path);
+    alloc_ctx_t *alloc_ctx, bool slow_path);
 void *iralloct_realign(tsdn_t *tsdn, void *ptr, size_t oldsize, size_t size,
     size_t extra, size_t alignment, bool zero, tcache_t *tcache,
     arena_t *arena);
@@ -107,7 +107,7 @@ ivsalloc(tsdn_t *tsdn, const void *ptr) {
 }
 
 JEMALLOC_ALWAYS_INLINE void
-idalloctm(tsdn_t *tsdn, void *ptr, tcache_t *tcache, dalloc_ctx_t *dalloc_ctx,
+idalloctm(tsdn_t *tsdn, void *ptr, tcache_t *tcache, alloc_ctx_t *alloc_ctx,
     bool is_internal, bool slow_path) {
 	assert(ptr != NULL);
 	assert(!is_internal || tcache == NULL);
@@ -120,7 +120,7 @@ idalloctm(tsdn_t *tsdn, void *ptr, tcache_t *tcache, dalloc_ctx_t *dalloc_ctx,
 	if (!is_internal && *tsd_reentrancy_levelp_get(tsdn_tsd(tsdn)) != 0) {
 		tcache = NULL;
 	}
-	arena_dalloc(tsdn, ptr, tcache, dalloc_ctx, slow_path);
+	arena_dalloc(tsdn, ptr, tcache, alloc_ctx, slow_path);
 }
 
 JEMALLOC_ALWAYS_INLINE void
@@ -130,9 +130,9 @@ idalloc(tsd_t *tsd, void *ptr) {
 
 JEMALLOC_ALWAYS_INLINE void
 isdalloct(tsdn_t *tsdn, void *ptr, size_t size, tcache_t *tcache,
-    dalloc_ctx_t *dalloc_ctx, bool slow_path) {
+    alloc_ctx_t *alloc_ctx, bool slow_path) {
 	witness_assert_depth_to_rank(tsdn, WITNESS_RANK_CORE, 0);
-	arena_sdalloc(tsdn, ptr, size, tcache, dalloc_ctx, slow_path);
+	arena_sdalloc(tsdn, ptr, size, tcache, alloc_ctx, slow_path);
 }
 
 JEMALLOC_ALWAYS_INLINE void *
diff --git a/include/jemalloc/internal/prof_inlines_b.h b/include/jemalloc/internal/prof_inlines_b.h
index 5ee72c5..6a79c01 100644
--- a/include/jemalloc/internal/prof_inlines_b.h
+++ b/include/jemalloc/internal/prof_inlines_b.h
@@ -5,20 +5,22 @@
 bool	prof_active_get_unlocked(void);
 bool	prof_gdump_get_unlocked(void);
 prof_tdata_t	*prof_tdata_get(tsd_t *tsd, bool create);
-prof_tctx_t	*prof_tctx_get(tsdn_t *tsdn, const void *ptr);
+prof_tctx_t	*prof_tctx_get(tsdn_t *tsdn, const void *ptr,
+    alloc_ctx_t *alloc_ctx);
 void	prof_tctx_set(tsdn_t *tsdn, const void *ptr, size_t usize,
-    prof_tctx_t *tctx);
+    alloc_ctx_t *alloc_ctx, prof_tctx_t *tctx);
 void	prof_tctx_reset(tsdn_t *tsdn, const void *ptr, prof_tctx_t *tctx);
 bool	prof_sample_accum_update(tsd_t *tsd, size_t usize, bool update,
     prof_tdata_t **tdata_out);
 prof_tctx_t	*prof_alloc_prep(tsd_t *tsd, size_t usize, bool prof_active,
     bool update);
 void	prof_malloc(tsdn_t *tsdn, const void *ptr, size_t usize,
-    prof_tctx_t *tctx);
+    alloc_ctx_t *alloc_ctx, prof_tctx_t *tctx);
 void	prof_realloc(tsd_t *tsd, const void *ptr, size_t usize,
     prof_tctx_t *tctx, bool prof_active, bool updated, const void *old_ptr,
     size_t old_usize, prof_tctx_t *old_tctx);
-void	prof_free(tsd_t *tsd, const void *ptr, size_t usize);
+void	prof_free(tsd_t *tsd, const void *ptr, size_t usize,
+    alloc_ctx_t *alloc_ctx);
 #endif
 
 #if (defined(JEMALLOC_ENABLE_INLINE) || defined(JEMALLOC_PROF_C_))
@@ -67,19 +69,20 @@ prof_tdata_get(tsd_t *tsd, bool create) {
 }
 
 JEMALLOC_ALWAYS_INLINE prof_tctx_t *
-prof_tctx_get(tsdn_t *tsdn, const void *ptr) {
+prof_tctx_get(tsdn_t *tsdn, const void *ptr, alloc_ctx_t *alloc_ctx) {
 	cassert(config_prof);
 	assert(ptr != NULL);
 
-	return arena_prof_tctx_get(tsdn, ptr);
+	return arena_prof_tctx_get(tsdn, ptr, alloc_ctx);
 }
 
 JEMALLOC_ALWAYS_INLINE void
-prof_tctx_set(tsdn_t *tsdn, const void *ptr, size_t usize, prof_tctx_t *tctx) {
+prof_tctx_set(tsdn_t *tsdn, const void *ptr, size_t usize,
+    alloc_ctx_t *alloc_ctx, prof_tctx_t *tctx) {
 	cassert(config_prof);
 	assert(ptr != NULL);
 
-	arena_prof_tctx_set(tsdn, ptr, usize, tctx);
+	arena_prof_tctx_set(tsdn, ptr, usize, alloc_ctx, tctx);
 }
 
 JEMALLOC_ALWAYS_INLINE void
@@ -145,7 +148,8 @@ prof_alloc_prep(tsd_t *tsd, size_t usize, bool prof_active, bool update) {
 }
 
 JEMALLOC_ALWAYS_INLINE void
-prof_malloc(tsdn_t *tsdn, const void *ptr, size_t usize, prof_tctx_t *tctx) {
+prof_malloc(tsdn_t *tsdn, const void *ptr, size_t usize, alloc_ctx_t *alloc_ctx,
+    prof_tctx_t *tctx) {
 	cassert(config_prof);
 	assert(ptr != NULL);
 	assert(usize == isalloc(tsdn, ptr));
@@ -153,7 +157,8 @@ prof_malloc(tsdn_t *tsdn, const void *ptr, size_t usize, prof_tctx_t *tctx) {
 	if (unlikely((uintptr_t)tctx > (uintptr_t)1U)) {
 		prof_malloc_sample_object(tsdn, ptr, usize, tctx);
 	} else {
-		prof_tctx_set(tsdn, ptr, usize, (prof_tctx_t *)(uintptr_t)1U);
+		prof_tctx_set(tsdn, ptr, usize, alloc_ctx,
+		    (prof_tctx_t *)(uintptr_t)1U);
 	}
 }
 
@@ -188,7 +193,7 @@ prof_realloc(tsd_t *tsd, const void *ptr, size_t usize, prof_tctx_t *tctx,
 	if (unlikely(sampled)) {
 		prof_malloc_sample_object(tsd_tsdn(tsd), ptr, usize, tctx);
 	} else if (moved) {
-		prof_tctx_set(tsd_tsdn(tsd), ptr, usize,
+		prof_tctx_set(tsd_tsdn(tsd), ptr, usize, NULL,
 		    (prof_tctx_t *)(uintptr_t)1U);
 	} else if (unlikely(old_sampled)) {
 		/*
@@ -199,7 +204,7 @@ prof_realloc(tsd_t *tsd, const void *ptr, size_t usize, prof_tctx_t *tctx,
 		 */
 		prof_tctx_reset(tsd_tsdn(tsd), ptr, tctx);
 	} else {
-		assert((uintptr_t)prof_tctx_get(tsd_tsdn(tsd), ptr) ==
+		assert((uintptr_t)prof_tctx_get(tsd_tsdn(tsd), ptr, NULL) ==
 		    (uintptr_t)1U);
 	}
 
@@ -216,8 +221,8 @@ prof_realloc(tsd_t *tsd, const void *ptr, size_t usize, prof_tctx_t *tctx,
 }
 
 JEMALLOC_ALWAYS_INLINE void
-prof_free(tsd_t *tsd, const void *ptr, size_t usize) {
-	prof_tctx_t *tctx = prof_tctx_get(tsd_tsdn(tsd), ptr);
+prof_free(tsd_t *tsd, const void *ptr, size_t usize, alloc_ctx_t *alloc_ctx) {
+	prof_tctx_t *tctx = prof_tctx_get(tsd_tsdn(tsd), ptr, alloc_ctx);
 
 	cassert(config_prof);
 	assert(usize == isalloc(tsd_tsdn(tsd), ptr));