*alloc: limit scope of RCU lock

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.

1 files 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);