From 38cda690ddd6c9db0321b06abaa4d19f884326b6 Mon Sep 17 00:00:00 2001
From: Jason Evans <je@facebook.com>
Date: Wed, 14 Apr 2010 11:24:45 -0700
Subject: [PATCH] Fix profiling regression caused by bugfix.

Properly set the context associated with each allocated object, even
when the object is not sampled.

Remove debug print code that slipped in.
---
 jemalloc/src/prof.c | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/jemalloc/src/prof.c b/jemalloc/src/prof.c
index c13bc04f..93904b8c 100644
--- a/jemalloc/src/prof.c
+++ b/jemalloc/src/prof.c
@@ -651,10 +651,11 @@ prof_malloc(const void *ptr, prof_thr_cnt_t *cnt)
 
 	assert(ptr != NULL);
 
-	prof_ctx_set(ptr, cnt->ctx);
 	prof_sample_accum_update(size);
 
 	if ((uintptr_t)cnt > (uintptr_t)1U) {
+		prof_ctx_set(ptr, cnt->ctx);
+
 		cnt->epoch++;
 		/*********/
 		mb_write();
@@ -670,7 +671,8 @@ prof_malloc(const void *ptr, prof_thr_cnt_t *cnt)
 		/*********/
 		mb_write();
 		/*********/
-	}
+	} else
+		prof_ctx_set(ptr, (prof_ctx_t *)(uintptr_t)1U);
 }
 
 void
@@ -680,10 +682,8 @@ prof_realloc(const void *ptr, prof_thr_cnt_t *cnt, const void *old_ptr,
 	size_t size = isalloc(ptr);
 	prof_thr_cnt_t *told_cnt;
 
-	if (ptr != NULL) {
-		prof_ctx_set(ptr, cnt->ctx);
+	if (ptr != NULL)
 		prof_sample_accum_update(size);
-	}
 
 	if ((uintptr_t)old_ctx > (uintptr_t)1U) {
 		told_cnt = prof_lookup(old_ctx->bt);
@@ -692,7 +692,6 @@ prof_realloc(const void *ptr, prof_thr_cnt_t *cnt, const void *old_ptr,
 			 * It's too late to propagate OOM for this realloc(),
 			 * so operate directly on old_cnt->ctx->cnt_merged.
 			 */
-			malloc_printf("XXX BANG A\n");
 			malloc_mutex_lock(&old_ctx->lock);
 			old_ctx->cnt_merged.curobjs--;
 			old_ctx->cnt_merged.curbytes -= old_size;
@@ -704,8 +703,11 @@ prof_realloc(const void *ptr, prof_thr_cnt_t *cnt, const void *old_ptr,
 
 	if ((uintptr_t)told_cnt > (uintptr_t)1U)
 		told_cnt->epoch++;
-	if ((uintptr_t)cnt > (uintptr_t)1U)
+	if ((uintptr_t)cnt > (uintptr_t)1U) {
+		prof_ctx_set(ptr, cnt->ctx);
 		cnt->epoch++;
+	} else
+		prof_ctx_set(ptr, (prof_ctx_t *)(uintptr_t)1U);
 	/*********/
 	mb_write();
 	/*********/
@@ -758,7 +760,6 @@ prof_free(const void *ptr)
 			 * OOM during free() cannot be propagated, so operate
 			 * directly on cnt->ctx->cnt_merged.
 			 */
-			malloc_printf("XXX BANG B\n");
 			malloc_mutex_lock(&ctx->lock);
 			ctx->cnt_merged.curobjs--;
 			ctx->cnt_merged.curbytes -= size;