From afa489c3c5fd16bd31b2756c081c92e08937e6b7 Mon Sep 17 00:00:00 2001
From: Yinan Zhang <zyn8950@gmail.com>
Date: Fri, 21 Aug 2020 11:31:53 -0700
Subject: [PATCH] Record request size in prof info

---
 include/jemalloc/internal/arena_inlines_b.h |  5 +++--
 include/jemalloc/internal/edata.h           | 12 ++++++++++++
 include/jemalloc/internal/large_externs.h   |  2 +-
 include/jemalloc/internal/prof_inlines.h    |  4 ++--
 include/jemalloc/internal/prof_structs.h    |  2 ++
 src/large.c                                 |  4 +++-
 src/prof.c                                  |  2 +-
 7 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/include/jemalloc/internal/arena_inlines_b.h b/include/jemalloc/internal/arena_inlines_b.h
index 7971b4c7..aaef45c0 100644
--- a/include/jemalloc/internal/arena_inlines_b.h
+++ b/include/jemalloc/internal/arena_inlines_b.h
@@ -105,11 +105,12 @@ arena_prof_tctx_reset_sampled(tsd_t *tsd, const void *ptr) {
 }
 
 JEMALLOC_ALWAYS_INLINE void
-arena_prof_info_set(tsd_t *tsd, edata_t *edata, prof_tctx_t *tctx) {
+arena_prof_info_set(tsd_t *tsd, edata_t *edata, prof_tctx_t *tctx,
+    size_t size) {
 	cassert(config_prof);
 
 	assert(!edata_slab_get(edata));
-	large_prof_info_set(edata, tctx);
+	large_prof_info_set(edata, tctx, size);
 }
 
 JEMALLOC_ALWAYS_INLINE void
diff --git a/include/jemalloc/internal/edata.h b/include/jemalloc/internal/edata.h
index c0482883..11358ea1 100644
--- a/include/jemalloc/internal/edata.h
+++ b/include/jemalloc/internal/edata.h
@@ -40,6 +40,8 @@ typedef enum extent_pai_e extent_pai_t;
 struct e_prof_info_s {
 	/* Time when this was allocated. */
 	nstime_t	e_prof_alloc_time;
+	/* Allocation request size. */
+	size_t		e_prof_alloc_size;
 	/* Points to a prof_tctx_t. */
 	atomic_p_t	e_prof_tctx;
 	/*
@@ -390,6 +392,11 @@ edata_prof_alloc_time_get(const edata_t *edata) {
 	return &edata->e_prof_info.e_prof_alloc_time;
 }
 
+static inline size_t
+edata_prof_alloc_size_get(const edata_t *edata) {
+	return edata->e_prof_info.e_prof_alloc_size;
+}
+
 static inline prof_recent_t *
 edata_prof_recent_alloc_get_dont_call_directly(const edata_t *edata) {
 	return (prof_recent_t *)atomic_load_p(
@@ -526,6 +533,11 @@ edata_prof_alloc_time_set(edata_t *edata, nstime_t *t) {
 	nstime_copy(&edata->e_prof_info.e_prof_alloc_time, t);
 }
 
+static inline void
+edata_prof_alloc_size_set(edata_t *edata, size_t size) {
+	edata->e_prof_info.e_prof_alloc_size = size;
+}
+
 static inline void
 edata_prof_recent_alloc_set_dont_call_directly(edata_t *edata,
     prof_recent_t *recent_alloc) {
diff --git a/include/jemalloc/internal/large_externs.h b/include/jemalloc/internal/large_externs.h
index 27979648..8e09122d 100644
--- a/include/jemalloc/internal/large_externs.h
+++ b/include/jemalloc/internal/large_externs.h
@@ -19,6 +19,6 @@ size_t large_salloc(tsdn_t *tsdn, const edata_t *edata);
 void large_prof_info_get(tsd_t *tsd, edata_t *edata, prof_info_t *prof_info,
     bool reset_recent);
 void large_prof_tctx_reset(edata_t *edata);
-void large_prof_info_set(edata_t *edata, prof_tctx_t *tctx);
+void large_prof_info_set(edata_t *edata, prof_tctx_t *tctx, size_t size);
 
 #endif /* JEMALLOC_INTERNAL_LARGE_EXTERNS_H */
diff --git a/include/jemalloc/internal/prof_inlines.h b/include/jemalloc/internal/prof_inlines.h
index 62c56832..c76d2ae5 100644
--- a/include/jemalloc/internal/prof_inlines.h
+++ b/include/jemalloc/internal/prof_inlines.h
@@ -98,12 +98,12 @@ prof_tctx_reset_sampled(tsd_t *tsd, const void *ptr) {
 }
 
 JEMALLOC_ALWAYS_INLINE void
-prof_info_set(tsd_t *tsd, edata_t *edata, prof_tctx_t *tctx) {
+prof_info_set(tsd_t *tsd, edata_t *edata, prof_tctx_t *tctx, size_t size) {
 	cassert(config_prof);
 	assert(edata != NULL);
 	assert((uintptr_t)tctx > (uintptr_t)1U);
 
-	arena_prof_info_set(tsd, edata, tctx);
+	arena_prof_info_set(tsd, edata, tctx, size);
 }
 
 JEMALLOC_ALWAYS_INLINE bool
diff --git a/include/jemalloc/internal/prof_structs.h b/include/jemalloc/internal/prof_structs.h
index 73ac3d5c..c2a111a9 100644
--- a/include/jemalloc/internal/prof_structs.h
+++ b/include/jemalloc/internal/prof_structs.h
@@ -103,6 +103,8 @@ struct prof_info_s {
 	nstime_t		alloc_time;
 	/* Points to the prof_tctx_t corresponding to the allocation. */
 	prof_tctx_t		*alloc_tctx;
+	/* Allocation request size. */
+	size_t			alloc_size;
 };
 
 struct prof_gctx_s {
diff --git a/src/large.c b/src/large.c
index 42d2fd7d..f23839f7 100644
--- a/src/large.c
+++ b/src/large.c
@@ -281,6 +281,7 @@ large_prof_info_get(tsd_t *tsd, edata_t *edata, prof_info_t *prof_info,
 	if ((uintptr_t)alloc_tctx > (uintptr_t)1U) {
 		nstime_copy(&prof_info->alloc_time,
 		    edata_prof_alloc_time_get(edata));
+		prof_info->alloc_size = edata_prof_alloc_size_get(edata);
 		if (reset_recent) {
 			/*
 			 * Reset the pointer on the recent allocation record,
@@ -302,10 +303,11 @@ large_prof_tctx_reset(edata_t *edata) {
 }
 
 void
-large_prof_info_set(edata_t *edata, prof_tctx_t *tctx) {
+large_prof_info_set(edata_t *edata, prof_tctx_t *tctx, size_t size) {
 	nstime_t t;
 	nstime_prof_init_update(&t);
 	edata_prof_alloc_time_set(edata, &t);
+	edata_prof_alloc_size_set(edata, size);
 	edata_prof_recent_alloc_init(edata);
 	large_prof_tctx_set(edata, tctx);
 }
diff --git a/src/prof.c b/src/prof.c
index 9b651db8..258b5f2d 100644
--- a/src/prof.c
+++ b/src/prof.c
@@ -97,7 +97,7 @@ prof_malloc_sample_object(tsd_t *tsd, const void *ptr, size_t size,
 
 	edata_t *edata = emap_edata_lookup(tsd_tsdn(tsd), &arena_emap_global,
 	    ptr);
-	prof_info_set(tsd, edata, tctx);
+	prof_info_set(tsd, edata, tctx, size);
 
 	szind_t szind = sz_size2index(usize);