Inside _talloc_free_internal(), always call talloc_memlimit_update_on_free() before...

[metze/samba/wip.git] / lib / talloc / talloc.c

diff --git a/lib/talloc/talloc.c b/lib/talloc/talloc.c
index 1e25dfde4e1e998872276ef3acb1eec80606cb25..74eca3f5e65a63194c247aa088b9ec20e24c2114 100644 (file)
--- a/lib/talloc/talloc.c
+++ b/lib/talloc/talloc.c
@@ -586,27 +586,24 @@ static inline void *__talloc(const void *context, size_t size)
                        limit = ptc->limit;
                }
 
-               if (!talloc_memlimit_check(limit, (TC_HDR_SIZE+size))) {
-                       errno = ENOMEM;
-                       return NULL;
-               }
-
                tc = talloc_alloc_pool(ptc, TC_HDR_SIZE+size);
        }
 
        if (tc == NULL) {
+               /*
+                * Only do the memlimit check/update on actual allocation.
+                */
+               if (!talloc_memlimit_check(limit, TC_HDR_SIZE + size)) {
+                       errno = ENOMEM;
+                       return NULL;
+               }
+
                tc = (struct talloc_chunk *)malloc(TC_HDR_SIZE+size);
                if (unlikely(tc == NULL)) return NULL;
                tc->flags = TALLOC_MAGIC;
                tc->pool  = NULL;
-       }
-
-       if (limit != NULL) {
-               struct talloc_memlimit *l;
 
-               for (l = limit; l != NULL; l = l->upper) {
-                       l->cur_size += TC_HDR_SIZE+size;
-               }
+               talloc_memlimit_grow(limit, TC_HDR_SIZE + size);
        }
 
        tc->limit = limit;
@@ -810,6 +807,8 @@ static inline void _talloc_free_poolmem(struct talloc_chunk *tc,
                 */
                pool->hdr.c.name = location;
 
+               talloc_memlimit_update_on_free(&pool->hdr.c);
+
                TC_INVALIDATE_FULL_CHUNK(&pool->hdr.c);
                free(pool);
                return;
@@ -910,29 +909,6 @@ static inline int _talloc_free_internal(void *ptr, const char *location)
 
        tc->flags |= TALLOC_FLAG_FREE;
 
-       /*
-        * If we are part of a memory limited context hierarchy
-        * we need to subtract the memory used from the counters
-        */
-       if (tc->limit) {
-               struct talloc_memlimit *l;
-
-               for (l = tc->limit; l != NULL; l = l->upper) {
-                       if (l->cur_size >= tc->size+TC_HDR_SIZE) {
-                               l->cur_size -= tc->size+TC_HDR_SIZE;
-                       } else {
-                               talloc_abort("cur_size memlimit counter not correct!");
-                               return 0;
-                       }
-               }
-
-               if (tc->limit->parent == tc) {
-                       free(tc->limit);
-               }
-
-               tc->limit = NULL;
-       }
-
        /* we mark the freed memory with where we called the free
         * from. This means on a double free error we can report where
         * the first free came from
@@ -953,6 +929,8 @@ static inline int _talloc_free_internal(void *ptr, const char *location)
                        return 0;
                }
 
+               talloc_memlimit_update_on_free(tc);
+
                TC_INVALIDATE_FULL_CHUNK(tc);
                free(tc);
                return 0;
@@ -963,6 +941,8 @@ static inline int _talloc_free_internal(void *ptr, const char *location)
                return 0;
        }
 
+       talloc_memlimit_update_on_free(tc);
+
        TC_INVALIDATE_FULL_CHUNK(tc);
        free(tc);
        return 0;