Change _talloc_total_mem_internal() to ignore memory allocated from a pool when calcu...

[sfrench/samba-autobuild/.git] / lib / talloc / talloc.c

diff --git a/lib/talloc/talloc.c b/lib/talloc/talloc.c
index d0c81d3d43a4ba25defcd65b508972d4242c784a..1e25dfde4e1e998872276ef3acb1eec80606cb25 100644 (file)
--- a/lib/talloc/talloc.c
+++ b/lib/talloc/talloc.c
@@ -238,6 +238,11 @@ struct talloc_memlimit {
 static bool talloc_memlimit_check(struct talloc_memlimit *limit, size_t size);
 static bool talloc_memlimit_update(struct talloc_memlimit *limit,
                                   size_t old_size, size_t new_size);
+static void talloc_memlimit_grow(struct talloc_memlimit *limit,
+                               size_t size);
+static void talloc_memlimit_shrink(struct talloc_memlimit *limit,
+                               size_t size);
+static void talloc_memlimit_update_on_free(struct talloc_chunk *tc);
 
 typedef int (*talloc_destructor_t)(void *);
 
@@ -255,7 +260,7 @@ struct talloc_chunk {
         * if 'limit' is set it means all *new* children of the context will
         * be limited to a total aggregate size ox max_size for memory
         * allocations.
-        * cur_size is used to kep track of the current use
+        * cur_size is used to keep track of the current use
         */
        struct talloc_memlimit *limit;
 
@@ -794,7 +799,10 @@ static inline void _talloc_free_poolmem(struct talloc_chunk *tc,
                 */
                pool->hdr.c.pool = tc_pool_first_chunk(pool);
                tc_invalidate_pool(pool);
-       } else if (unlikely(pool->hdr.object_count == 0)) {
+               return;
+       }
+
+       if (unlikely(pool->hdr.object_count == 0)) {
                /*
                 * we mark the freed memory with where we called the free
                 * from. This means on a double free error we can report where
@@ -804,14 +812,23 @@ static inline void _talloc_free_poolmem(struct talloc_chunk *tc,
 
                TC_INVALIDATE_FULL_CHUNK(&pool->hdr.c);
                free(pool);
-       } else if (pool->hdr.c.pool == next_tc) {
+               return;
+       }
+
+       if (pool->hdr.c.pool == next_tc) {
                /*
                 * if pool->pool still points to end of
                 * 'tc' (which is stored in the 'next_tc' variable),
                 * we can reclaim the memory of 'tc'.
                 */
                pool->hdr.c.pool = tc;
+               return;
        }
+
+       /*
+        * Do nothing. The memory is just "wasted", waiting for the pool
+        * itself to be freed.
+        */
 }
 
 static inline void _talloc_free_children_internal(struct talloc_chunk *tc,
@@ -931,16 +948,23 @@ static inline int _talloc_free_internal(void *ptr, const char *location)
                }
 
                pool->hdr.object_count--;
-               if (unlikely(pool->hdr.object_count == 0)) {
-                       TC_INVALIDATE_FULL_CHUNK(tc);
-                       free(tc);
+
+               if (likely(pool->hdr.object_count != 0)) {
+                       return 0;
                }
-       } else if (tc->flags & TALLOC_FLAG_POOLMEM) {
-               _talloc_free_poolmem(tc, location);
-       } else {
+
                TC_INVALIDATE_FULL_CHUNK(tc);
                free(tc);
+               return 0;
        }
+
+       if (tc->flags & TALLOC_FLAG_POOLMEM) {
+               _talloc_free_poolmem(tc, location);
+               return 0;
+       }
+
+       TC_INVALIDATE_FULL_CHUNK(tc);
+       free(tc);
        return 0;
 }
 
@@ -1793,7 +1817,14 @@ static size_t _talloc_total_mem_internal(const void *ptr,
                break;
        case TOTAL_MEM_LIMIT:
                if (likely(tc->name != TALLOC_MAGIC_REFERENCE)) {
-                       total = tc->size + TC_HDR_SIZE;
+                       /*
+                        * Don't count memory allocated from a pool
+                        * when calculating limits. Only count the
+                        * pool itself.
+                        */
+                       if (!(tc->flags & TALLOC_FLAG_POOLMEM)) {
+                               total = tc->size + TC_HDR_SIZE;
+                       }
                }
                break;
        }
@@ -2537,7 +2568,7 @@ static bool talloc_memlimit_check(struct talloc_memlimit *limit, size_t size)
        for (l = limit; l != NULL; l = l->upper) {
                if (l->max_size != 0 &&
                    ((l->max_size <= l->cur_size) ||
-                    (l->max_size - l->cur_size < TC_HDR_SIZE+size))) {
+                    (l->max_size - l->cur_size < size))) {
                        return false;
                }
        }
@@ -2545,6 +2576,73 @@ static bool talloc_memlimit_check(struct talloc_memlimit *limit, size_t size)
        return true;
 }
 
+/*
+  Update memory limits when freeing a talloc_chunk.
+*/
+static void talloc_memlimit_update_on_free(struct talloc_chunk *tc)
+{
+       if (!tc->limit) {
+               return;
+       }
+
+       /*
+        * Pool entries don't count. Only the pools
+        * themselves are counted as part of the memory
+        * limits.
+        */
+       if (tc->flags & TALLOC_FLAG_POOLMEM) {
+               return;
+       }
+
+       /*
+        * If we are part of a memory limited context hierarchy
+        * we need to subtract the memory used from the counters
+        */
+
+       talloc_memlimit_shrink(tc->limit, tc->size+TC_HDR_SIZE);
+
+       if (tc->limit->parent == tc) {
+               free(tc->limit);
+       }
+
+       tc->limit = NULL;
+}
+
+/*
+  Increase memory limit accounting after a malloc/realloc.
+*/
+static void talloc_memlimit_grow(struct talloc_memlimit *limit,
+                               size_t size)
+{
+       struct talloc_memlimit *l;
+
+       for (l = limit; l != NULL; l = l->upper) {
+               size_t new_cur_size = l->cur_size + size;
+               if (new_cur_size < l->cur_size) {
+                       talloc_abort("logic error in talloc_memlimit_grow\n");
+                       return;
+               }
+               l->cur_size = new_cur_size;
+       }
+}
+
+/*
+  Decrease memory limit accounting after a free/realloc.
+*/
+static void talloc_memlimit_shrink(struct talloc_memlimit *limit,
+                               size_t size)
+{
+       struct talloc_memlimit *l;
+
+       for (l = limit; l != NULL; l = l->upper) {
+               if (l->cur_size < size) {
+                       talloc_abort("logic error in talloc_memlimit_shrink\n");
+                       return;
+               }
+               l->cur_size = l->cur_size - size;
+       }
+}
+
 static bool talloc_memlimit_update(struct talloc_memlimit *limit,
                                   size_t old_size, size_t new_size)
 {