[PATCH 5/7] *alloc: limit scope of RCU lock

[Eric Wong <e@80x24.org>]

From: Eric Wong <mwrap-perl@80x24.org>

We can delay taking the RCU lock until the real_malloc() call
succeeds.  This gives us accurate stats in case of ENOMEM
failures and reduces our LoC a bit, too.
---
 mwrap_core.h | 62 ++++++++++++++++++++++++----------------------------
 1 file changed, 28 insertions(+), 34 deletions(-)

diff --git a/mwrap_core.h b/mwrap_core.h
index 14e5d79..3161243 100644
--- a/mwrap_core.h
+++ b/mwrap_core.h
@@ -509,8 +509,6 @@ static void
 alloc_insert_rcu(struct src_loc *l, struct alloc_hdr *h, size_t size,
 		void *real, size_t generation)
 {
-	/* we need src_loc to remain alive for the duration of this call */
-	if (!h) return;
 	h->size = size;
 	h->real = real;
 	h->as.live.loc = l;
@@ -540,11 +538,8 @@ static bool is_power_of_two(size_t n)
 static int
 mwrap_memalign(void **pp, size_t alignment, size_t size, struct src_loc *sl)
 {
-	struct src_loc *l;
-	struct alloc_hdr *h;
 	void *real;
 	size_t asize;
-	size_t generation = 0;
 	size_t d = alignment / sizeof(void*);
 	size_t r = alignment % sizeof(void*);
 
@@ -563,18 +558,18 @@ mwrap_memalign(void **pp, size_t alignment, size_t size, struct src_loc *sl)
 	    __builtin_add_overflow(asize, sizeof(struct alloc_hdr), &asize))
 		return ENOMEM;
 
-	l = update_stats_rcu_lock(&generation, size, sl);
-
 	real = real_malloc(asize);
 	if (real) {
 		void *p = hdr2ptr(real);
 		if (!ptr_is_aligned(p, alignment))
 			p = ptr_align(p, alignment);
-		h = ptr2hdr(p);
-		alloc_insert_rcu(l, h, size, real, generation);
+		struct alloc_hdr *h = ptr2hdr(p);
+		size_t gen = 0;
+		struct src_loc *l = update_stats_rcu_lock(&gen, size, sl);
+		alloc_insert_rcu(l, h, size, real, gen);
+		update_stats_rcu_unlock(l);
 		*pp = p;
 	}
-	update_stats_rcu_unlock(l);
 
 	return real ? 0 : ENOMEM;
 }
@@ -652,16 +647,16 @@ void *malloc(size_t size)
 	if (__builtin_add_overflow(size, sizeof(struct alloc_hdr), &asize))
 		goto enomem;
 
-	size_t generation = 0;
-	SRC_LOC_BT(bt);
-	struct src_loc *l = update_stats_rcu_lock(&generation, size, &bt.sl);
-	struct alloc_hdr *h;
-	void *p = h = real_malloc(asize);
-	if (h) {
-		alloc_insert_rcu(l, h, size, h, generation);
+	void *p = real_malloc(asize);
+	if (p) {
+		size_t gen = 0;
+		SRC_LOC_BT(bt);
+		struct alloc_hdr *h = p;
+		struct src_loc *l = update_stats_rcu_lock(&gen, size, &bt.sl);
+		alloc_insert_rcu(l, h, size, h, gen);
+		update_stats_rcu_unlock(l);
 		p = hdr2ptr(h);
 	}
-	update_stats_rcu_unlock(l);
 	if (p) return p;
 enomem:
 	errno = ENOMEM;
@@ -671,22 +666,21 @@ enomem:
 void *calloc(size_t nmemb, size_t size)
 {
 	size_t asize;
-	size_t generation = 0;
 
 	if (__builtin_mul_overflow(size, nmemb, &size))
 		goto enomem;
 	if (__builtin_add_overflow(size, sizeof(struct alloc_hdr), &asize))
 		goto enomem;
-	struct alloc_hdr *h;
-	SRC_LOC_BT(bt);
-	struct src_loc *l = update_stats_rcu_lock(&generation, size, &bt.sl);
-	void *p = h = real_malloc(asize);
+	void *p = real_malloc(asize);
 	if (p) {
-		alloc_insert_rcu(l, h, size, h, generation);
-		p = hdr2ptr(h);
+		size_t gen = 0;
+		struct alloc_hdr *h = p;
+		SRC_LOC_BT(bt);
+		struct src_loc *l = update_stats_rcu_lock(&gen, size, &bt.sl);
+		alloc_insert_rcu(l, h, size, h, gen);
+		update_stats_rcu_unlock(l);
+		return memset(hdr2ptr(h), 0, size);
 	}
-	update_stats_rcu_unlock(l);
-	if (p) return memset(p, 0, size);
 enomem:
 	errno = ENOMEM;
 	return 0;
@@ -704,16 +698,16 @@ void *realloc(void *ptr, size_t size)
 		errno = ENOMEM;
 		return 0;
 	}
-	struct alloc_hdr *h;
-	size_t generation = 0;
-	SRC_LOC_BT(bt);
-	struct src_loc *l = update_stats_rcu_lock(&generation, size, &bt.sl);
-	void *p = h = real_malloc(asize);
+	void *p = real_malloc(asize);
 	if (p) {
-		alloc_insert_rcu(l, h, size, h, generation);
+		size_t gen = 0;
+		struct alloc_hdr *h = p;
+		SRC_LOC_BT(bt);
+		struct src_loc *l = update_stats_rcu_lock(&gen, size, &bt.sl);
+		alloc_insert_rcu(l, h, size, h, gen);
+		update_stats_rcu_unlock(l);
 		p = hdr2ptr(h);
 	}
-	update_stats_rcu_unlock(l);
 
 	if (ptr && p) {
 		struct alloc_hdr *old = ptr2hdr(ptr);