From f970c497dc5ab2b885d266d59948510c430f8ec1 Mon Sep 17 00:00:00 2001
From: Qi Wang <interwq@gwu.edu>
Date: Fri, 21 Apr 2017 15:05:43 -0700
Subject: [PATCH] Implement malloc_mutex_trylock() w/ proper stats update.

---
 include/jemalloc/internal/mutex_inlines.h     | 42 +++++++++++++------
 include/jemalloc/internal/private_symbols.txt |  4 ++
 src/mutex.c                                   |  4 +-
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/include/jemalloc/internal/mutex_inlines.h b/include/jemalloc/internal/mutex_inlines.h
index 2856d844..6da21cf6 100644
--- a/include/jemalloc/internal/mutex_inlines.h
+++ b/include/jemalloc/internal/mutex_inlines.h
@@ -10,10 +10,36 @@ malloc_mutex_lock_final(malloc_mutex_t *mutex) {
 	MALLOC_MUTEX_LOCK(mutex);
 }
 
+static inline bool
+malloc_mutex_trylock_final(malloc_mutex_t *mutex) {
+	return MALLOC_MUTEX_TRYLOCK(mutex);
+}
+
+static inline void
+mutex_owner_stats_update(tsdn_t *tsdn, malloc_mutex_t *mutex) {
+	if (config_stats) {
+		mutex_prof_data_t *data = &mutex->prof_data;
+		data->n_lock_ops++;
+		if (data->prev_owner != tsdn) {
+			data->prev_owner = tsdn;
+			data->n_owner_switches++;
+		}
+	}
+}
+
 /* Trylock: return false if the lock is successfully acquired. */
 static inline bool
-malloc_mutex_trylock(malloc_mutex_t *mutex) {
-	return MALLOC_MUTEX_TRYLOCK(mutex);
+malloc_mutex_trylock(tsdn_t *tsdn, malloc_mutex_t *mutex) {
+	witness_assert_not_owner(tsdn, &mutex->witness);
+	if (isthreaded) {
+		if (malloc_mutex_trylock_final(mutex)) {
+			return true;
+		}
+		mutex_owner_stats_update(tsdn, mutex);
+	}
+	witness_lock(tsdn, &mutex->witness);
+
+	return false;
 }
 
 /* Aggregate lock prof data. */
@@ -44,18 +70,10 @@ static inline void
 malloc_mutex_lock(tsdn_t *tsdn, malloc_mutex_t *mutex) {
 	witness_assert_not_owner(tsdn, &mutex->witness);
 	if (isthreaded) {
-		if (malloc_mutex_trylock(mutex)) {
+		if (malloc_mutex_trylock_final(mutex)) {
 			malloc_mutex_lock_slow(mutex);
 		}
-		/* We own the lock now.  Update a few counters. */
-		if (config_stats) {
-			mutex_prof_data_t *data = &mutex->prof_data;
-			data->n_lock_ops++;
-			if (data->prev_owner != tsdn) {
-				data->prev_owner = tsdn;
-				data->n_owner_switches++;
-			}
-		}
+		mutex_owner_stats_update(tsdn, mutex);
 	}
 	witness_lock(tsdn, &mutex->witness);
 }
diff --git a/include/jemalloc/internal/private_symbols.txt b/include/jemalloc/internal/private_symbols.txt
index 649a689f..50590957 100644
--- a/include/jemalloc/internal/private_symbols.txt
+++ b/include/jemalloc/internal/private_symbols.txt
@@ -292,10 +292,13 @@ malloc_mutex_assert_owner
 malloc_mutex_boot
 malloc_mutex_init
 malloc_mutex_lock
+malloc_mutex_lock_final
 malloc_mutex_lock_slow
 malloc_mutex_postfork_child
 malloc_mutex_postfork_parent
 malloc_mutex_prefork
+malloc_mutex_trylock
+malloc_mutex_trylock_final
 malloc_mutex_unlock
 malloc_printf
 malloc_slow
@@ -309,6 +312,7 @@ malloc_tsd_malloc
 malloc_vcprintf
 malloc_vsnprintf
 malloc_write
+mutex_owner_stats_update
 narenas_auto
 narenas_total_get
 ncpus
diff --git a/src/mutex.c b/src/mutex.c
index 3bb5ce1d..3eec970f 100644
--- a/src/mutex.c
+++ b/src/mutex.c
@@ -81,7 +81,7 @@ malloc_mutex_lock_slow(malloc_mutex_t *mutex) {
 	int cnt = 0, max_cnt = MALLOC_MUTEX_MAX_SPIN;
 	do {
 		CPU_SPINWAIT;
-		if (!malloc_mutex_trylock(mutex)) {
+		if (!malloc_mutex_trylock_final(mutex)) {
 			data->n_spin_acquired++;
 			return;
 		}
@@ -100,7 +100,7 @@ label_spin_done:
 	uint32_t n_thds = atomic_fetch_add_u32(&data->n_waiting_thds, 1,
 	    ATOMIC_RELAXED) + 1;
 	/* One last try as above two calls may take quite some cycles. */
-	if (!malloc_mutex_trylock(mutex)) {
+	if (!malloc_mutex_trylock_final(mutex)) {
 		atomic_fetch_sub_u32(&data->n_waiting_thds, 1, ATOMIC_RELAXED);
 		data->n_spin_acquired++;
 		return;