Add memory limiting capability to talloc

By calling talloc_set_memlimit() we can now set a max memory limit
for a whole talloc hierarchy.
ANy attempt to allocate memory beyond the max allowed for the whole
hierarchy wil cause an allocation failure.

Stealing memory correctly accounts for used memory in the old and the new
hierarchy but exceeding the memory limit in the new parent will not cause
a failure.

(cherry picked from commit a33a78c302fde61fdb7a6e71669f19be2cf5c836)

5 files changed, 327 insertions, 40 deletions

diff --git a/lib/talloc/ABI/pytalloc-util-2.0.8.sigs b/lib/talloc/ABI/pytalloc-util-2.0.8.sigs
new file mode 100644
index 00000000000..961c1a84d0b
--- /dev/null
+++ b/lib/talloc/ABI/pytalloc-util-2.0.8.sigs
@@ -0,0 +1,6 @@
+pytalloc_CObject_FromTallocPtr: PyObject *(void *)
+pytalloc_Check: int (PyObject *)
+pytalloc_GetObjectType: PyTypeObject *(void)
+pytalloc_reference_ex: PyObject *(PyTypeObject *, TALLOC_CTX *, void *)
+pytalloc_steal: PyObject *(PyTypeObject *, void *)
+pytalloc_steal_ex: PyObject *(PyTypeObject *, TALLOC_CTX *, void *)
diff --git a/lib/talloc/ABI/talloc-2.0.8.sigs b/lib/talloc/ABI/talloc-2.0.8.sigs
new file mode 100644
index 00000000000..15a9e9554fc
--- /dev/null
+++ b/lib/talloc/ABI/talloc-2.0.8.sigs
@@ -0,0 +1,63 @@
+_talloc: void *(const void *, size_t)
+_talloc_array: void *(const void *, size_t, unsigned int, const char *)
+_talloc_free: int (void *, const char *)
+_talloc_get_type_abort: void *(const void *, const char *, const char *)
+_talloc_memdup: void *(const void *, const void *, size_t, const char *)
+_talloc_move: void *(const void *, const void *)
+_talloc_realloc: void *(const void *, void *, size_t, const char *)
+_talloc_realloc_array: void *(const void *, void *, size_t, unsigned int, const char *)
+_talloc_reference_loc: void *(const void *, const void *, const char *)
+_talloc_set_destructor: void (const void *, int (*)(void *))
+_talloc_steal_loc: void *(const void *, const void *, const char *)
+_talloc_zero: void *(const void *, size_t, const char *)
+_talloc_zero_array: void *(const void *, size_t, unsigned int, const char *)
+talloc_asprintf: char *(const void *, const char *, ...)
+talloc_asprintf_append: char *(char *, const char *, ...)
+talloc_asprintf_append_buffer: char *(char *, const char *, ...)
+talloc_autofree_context: void *(void)
+talloc_check_name: void *(const void *, const char *)
+talloc_disable_null_tracking: void (void)
+talloc_enable_leak_report: void (void)
+talloc_enable_leak_report_full: void (void)
+talloc_enable_null_tracking: void (void)
+talloc_enable_null_tracking_no_autofree: void (void)
+talloc_find_parent_byname: void *(const void *, const char *)
+talloc_free_children: void (void *)
+talloc_get_name: const char *(const void *)
+talloc_get_size: size_t (const void *)
+talloc_increase_ref_count: int (const void *)
+talloc_init: void *(const char *, ...)
+talloc_is_parent: int (const void *, const void *)
+talloc_named: void *(const void *, size_t, const char *, ...)
+talloc_named_const: void *(const void *, size_t, const char *)
+talloc_parent: void *(const void *)
+talloc_parent_name: const char *(const void *)
+talloc_pool: void *(const void *, size_t)
+talloc_realloc_fn: void *(const void *, void *, size_t)
+talloc_reference_count: size_t (const void *)
+talloc_reparent: void *(const void *, const void *, const void *)
+talloc_report: void (const void *, FILE *)
+talloc_report_depth_cb: void (const void *, int, int, void (*)(const void *, int, int, int, void *), void *)
+talloc_report_depth_file: void (const void *, int, int, FILE *)
+talloc_report_full: void (const void *, FILE *)
+talloc_set_abort_fn: void (void (*)(const char *))
+talloc_set_log_fn: void (void (*)(const char *))
+talloc_set_log_stderr: void (void)
+talloc_set_memlimit: int (const void *, size_t)
+talloc_set_name: const char *(const void *, const char *, ...)
+talloc_set_name_const: void (const void *, const char *)
+talloc_show_parents: void (const void *, FILE *)
+talloc_strdup: char *(const void *, const char *)
+talloc_strdup_append: char *(char *, const char *)
+talloc_strdup_append_buffer: char *(char *, const char *)
+talloc_strndup: char *(const void *, const char *, size_t)
+talloc_strndup_append: char *(char *, const char *, size_t)
+talloc_strndup_append_buffer: char *(char *, const char *, size_t)
+talloc_total_blocks: size_t (const void *)
+talloc_total_size: size_t (const void *)
+talloc_unlink: int (const void *, void *)
+talloc_vasprintf: char *(const void *, const char *, va_list)
+talloc_vasprintf_append: char *(char *, const char *, va_list)
+talloc_vasprintf_append_buffer: char *(char *, const char *, va_list)
+talloc_version_major: int (void)
+talloc_version_minor: int (void)
diff --git a/lib/talloc/talloc.c b/lib/talloc/talloc.c
index 18ee5480957..afc44b338fa 100644
--- a/lib/talloc/talloc.c
+++ b/lib/talloc/talloc.c
@@ -71,6 +71,7 @@
 #define TALLOC_FLAG_LOOP 0x02
 #define TALLOC_FLAG_POOL 0x04		/* This is a talloc pool */
 #define TALLOC_FLAG_POOLMEM 0x08	/* This is allocated in a pool */
+
 #define TALLOC_MAGIC_REFERENCE ((const char *)1)
 
 /* by default we abort when given a bad pointer (such as when talloc_free() is called 
@@ -221,12 +222,50 @@ static struct {
 	TC_UNDEFINE_GROW_VALGRIND_CHUNK(_tc, _new_size); \
 } while (0)
 
+#define TALLOC_MEMLIMIT_CHECK(limit, size) do { \
+	struct talloc_memlimit *l; \
+	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))) { \
+			errno = ENOMEM; \
+			return NULL; \
+		} \
+	} \
+} while(0)
+
+#define TALLOC_MEMLIMIT_UPDATE(limit, o_size, n_size) do { \
+	struct talloc_memlimit *l; \
+	ssize_t d; \
+	if (o_size == 0) { \
+		d = n_size + TC_HDR_SIZE; \
+	} else { \
+		d = n_size - o_size; \
+	} \
+	for (l = limit; l != NULL; l = l->upper) { \
+		ssize_t new_size = l->cur_size + d; \
+		if (new_size < 0) { \
+			talloc_abort("cur_size memlimit counter not correct!"); \
+			errno = EINVAL; \
+			return NULL; \
+		} \
+		l->cur_size = new_size; \
+	} \
+} while(0)
+
 struct talloc_reference_handle {
 	struct talloc_reference_handle *next, *prev;
 	void *ptr;
 	const char *location;
 };
 
+struct talloc_memlimit {
+	struct talloc_chunk *parent;
+	struct talloc_memlimit *upper;
+	size_t max_size;
+	size_t cur_size;
+};
+
 typedef int (*talloc_destructor_t)(void *);
 
 struct talloc_chunk {
@@ -239,6 +278,15 @@ struct talloc_chunk {
 	unsigned flags;
 
 	/*
+	 * limit semantics:
+	 * 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
+	 */
+	struct talloc_memlimit *limit;
+
+	/*
 	 * "pool" has dual use:
 	 *
 	 * For the talloc pool itself (i.e. TALLOC_FLAG_POOL is set), "pool"
@@ -543,6 +591,7 @@ static struct talloc_chunk *talloc_alloc_pool(struct talloc_chunk *parent,
 static inline void *__talloc(const void *context, size_t size)
 {
 	struct talloc_chunk *tc = NULL;
+	struct talloc_memlimit *limit = NULL;
 
 	if (unlikely(context == NULL)) {
 		context = null_context;
@@ -553,8 +602,15 @@ static inline void *__talloc(const void *context, size_t size)
 	}
 
 	if (context != NULL) {
-		tc = talloc_alloc_pool(talloc_chunk_from_ptr(context),
-				       TC_HDR_SIZE+size);
+		struct talloc_chunk *ptc = talloc_chunk_from_ptr(context);
+
+		if (ptc->limit != NULL) {
+			limit = ptc->limit;
+		}
+
+		TALLOC_MEMLIMIT_CHECK(limit, (TC_HDR_SIZE+size));
+
+		tc = talloc_alloc_pool(ptc, TC_HDR_SIZE+size);
 	}
 
 	if (tc == NULL) {
@@ -564,6 +620,15 @@ static inline void *__talloc(const void *context, size_t size)
 		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;
+		}
+	}
+
+	tc->limit = limit;
 	tc->size = size;
 	tc->destructor = NULL;
 	tc->child = NULL;
@@ -852,6 +917,29 @@ 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 
@@ -880,6 +968,10 @@ static inline int _talloc_free_internal(void *ptr, const char *location)
 	return 0;
 }
 
+static size_t _talloc_total_limit_size(const void *ptr,
+					struct talloc_memlimit *old_limit,
+					struct talloc_memlimit *new_limit);
+
 /* 
    move a lump of memory from one talloc context to another return the
    ptr on success, or NULL if it could not be transferred.
@@ -888,6 +980,7 @@ static inline int _talloc_free_internal(void *ptr, const char *location)
 static void *_talloc_steal_internal(const void *new_ctx, const void *ptr)
 {
 	struct talloc_chunk *tc, *new_tc;
+	size_t ctx_size = 0;
 
 	if (unlikely(!ptr)) {
 		return NULL;
@@ -899,6 +992,19 @@ static void *_talloc_steal_internal(const void *new_ctx, const void *ptr)
 
 	tc = talloc_chunk_from_ptr(ptr);
 
+	if (tc->limit != NULL) {
+
+		ctx_size = _talloc_total_limit_size(ptr, NULL, NULL);
+
+		TALLOC_MEMLIMIT_UPDATE(tc->limit->upper, ctx_size, 0);
+
+		if (tc->limit->parent == tc) {
+			tc->limit->upper = NULL;
+		} else {
+			tc->limit = NULL;
+		}
+	}
+
 	if (unlikely(new_ctx == NULL)) {
 		if (tc->parent) {
 			_TLIST_REMOVE(tc->parent->child, tc);
@@ -909,7 +1015,7 @@ static void *_talloc_steal_internal(const void *new_ctx, const void *ptr)
 			if (tc->prev) tc->prev->next = tc->next;
 			if (tc->next) tc->next->prev = tc->prev;
 		}
-		
+
 		tc->parent = tc->next = tc->prev = NULL;
 		return discard_const_p(void, ptr);
 	}
@@ -935,6 +1041,19 @@ static void *_talloc_steal_internal(const void *new_ctx, const void *ptr)
 	if (new_tc->child) new_tc->child->parent = NULL;
 	_TLIST_ADD(new_tc->child, tc);
 
+	if (tc->limit || new_tc->limit) {
+		ctx_size = _talloc_total_limit_size(ptr, tc->limit,
+						    new_tc->limit);
+	}
+
+	if (new_tc->limit) {
+		struct talloc_memlimit *l;
+
+		for (l = new_tc->limit; l != NULL; l = l->upper) {
+			l->cur_size += ctx_size;
+		}
+	}
+
 	return discard_const_p(void, ptr);
 }
 
@@ -1068,7 +1187,7 @@ _PUBLIC_ int talloc_unlink(const void *context, void *ptr)
 	if (tc_c != talloc_parent_chunk(ptr)) {
 		return -1;
 	}
-	
+
 	tc_p = talloc_chunk_from_ptr(ptr);
 
 	if (tc_p->refs == NULL) {
@@ -1411,6 +1530,10 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
 		return NULL;
 	}
 
+	if (tc->limit && (size - tc->size > 0)) {
+		TALLOC_MEMLIMIT_CHECK(tc->limit, (size - tc->size));
+	}
+
 	/* handle realloc inside a talloc_pool */
 	if (unlikely(tc->flags & TALLOC_FLAG_POOLMEM)) {
 		pool_tc = (union talloc_pool_chunk *)tc->pool;
@@ -1418,7 +1541,7 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
 
 #if (ALWAYS_REALLOC == 0)
 	/* don't shrink if we have less than 1k to gain */
-	if (size < tc->size) {
+	if (size < tc->size && tc->limit == NULL) {
 		if (pool_tc) {
 			void *next_tc = tc_next_chunk(tc);
 			TC_INVALIDATE_SHRINK_CHUNK(tc, size);
@@ -1526,6 +1649,7 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
 		if (new_chunk_size == old_chunk_size) {
 			TC_UNDEFINE_GROW_CHUNK(tc, size);
 			tc->flags &= ~TALLOC_FLAG_FREE;
+			TALLOC_MEMLIMIT_UPDATE(tc->limit, tc->size, size);
 			tc->size = size;
 			return ptr;
 		}
@@ -1541,6 +1665,7 @@ _PUBLIC_ void *_talloc_realloc(const void *context, void *ptr, size_t size, cons
 			if (space_left >= space_needed) {
 				TC_UNDEFINE_GROW_CHUNK(tc, size);
 				tc->flags &= ~TALLOC_FLAG_FREE;
+				TALLOC_MEMLIMIT_UPDATE(tc->limit, tc->size, size);
 				tc->size = size;
 				pool_tc->hdr.c.pool = tc_next_chunk(tc);
 				return ptr;
@@ -1589,6 +1714,7 @@ got_new_ptr:
 		tc->next->prev = tc;
 	}
 
+	TALLOC_MEMLIMIT_UPDATE(tc->limit, tc->size, size);
 	tc->size = size;
 	_talloc_set_name_const(TC_PTR_FROM_CHUNK(tc), name);
 
@@ -1607,10 +1733,16 @@ _PUBLIC_ void *_talloc_move(const void *new_ctx, const void *_pptr)
 	return ret;
 }
 
-/*
-  return the total size of a talloc pool (subtree)
-*/
-_PUBLIC_ size_t talloc_total_size(const void *ptr)
+enum talloc_mem_count_type {
+	TOTAL_MEM_SIZE,
+	TOTAL_MEM_BLOCKS,
+	TOTAL_MEM_LIMIT,
+};
+
+static size_t _talloc_total_mem_internal(const void *ptr,
+					 enum talloc_mem_count_type type,
+					 struct talloc_memlimit *old_limit,
+					 struct talloc_memlimit *new_limit)
 {
 	size_t total = 0;
 	struct talloc_chunk *c, *tc;
@@ -1624,17 +1756,50 @@ _PUBLIC_ size_t talloc_total_size(const void *ptr)
 
 	tc = talloc_chunk_from_ptr(ptr);
 
+	if (old_limit || new_limit) {
+		if (tc->limit && tc->limit->upper == old_limit) {
+			tc->limit->upper = new_limit;
+		}
+	}
+
+	/* optimize in the memlimits case */
+	if (type == TOTAL_MEM_LIMIT &&
+	    tc->limit != NULL &&
+	    tc->limit != old_limit &&
+	    tc->limit->parent == tc) {
+		return tc->limit->cur_size;
+	}
+
 	if (tc->flags & TALLOC_FLAG_LOOP) {
 		return 0;
 	}
 
 	tc->flags |= TALLOC_FLAG_LOOP;
 
-	if (likely(tc->name != TALLOC_MAGIC_REFERENCE)) {
-		total = tc->size;
+	if (old_limit || new_limit) {
+		if (old_limit == tc->limit) {
+			tc->limit = new_limit;
+		}
 	}
-	for (c=tc->child;c;c=c->next) {
-		total += talloc_total_size(TC_PTR_FROM_CHUNK(c));
+
+	switch (type) {
+	case TOTAL_MEM_SIZE:
+		if (likely(tc->name != TALLOC_MAGIC_REFERENCE)) {
+			total = tc->size;
+		}
+		break;
+	case TOTAL_MEM_BLOCKS:
+		total++;
+		break;
+	case TOTAL_MEM_LIMIT:
+		if (likely(tc->name != TALLOC_MAGIC_REFERENCE)) {
+			total = tc->size + TC_HDR_SIZE;
+		}
+		break;
+	}
+	for (c = tc->child; c; c = c->next) {
+		total += _talloc_total_mem_internal(TC_PTR_FROM_CHUNK(c), type,
+						    old_limit, new_limit);
 	}
 
 	tc->flags &= ~TALLOC_FLAG_LOOP;
@@ -1643,36 +1808,19 @@ _PUBLIC_ size_t talloc_total_size(const void *ptr)
 }
 
 /*
+  return the total size of a talloc pool (subtree)
+*/
+_PUBLIC_ size_t talloc_total_size(const void *ptr)
+{
+	return _talloc_total_mem_internal(ptr, TOTAL_MEM_SIZE, NULL, NULL);
+}
+
+/*
   return the total number of blocks in a talloc pool (subtree)
 */
 _PUBLIC_ size_t talloc_total_blocks(const void *ptr)
 {
-	size_t total = 0;
-	struct talloc_chunk *c, *tc;
-
-	if (ptr == NULL) {
-		ptr = null_context;
-	}
-	if (ptr == NULL) {
-		return 0;
-	}
-
-	tc = talloc_chunk_from_ptr(ptr);
-
-	if (tc->flags & TALLOC_FLAG_LOOP) {
-		return 0;
-	}
-
-	tc->flags |= TALLOC_FLAG_LOOP;
-
-	total++;
-	for (c=tc->child;c;c=c->next) {
-		total += talloc_total_blocks(TC_PTR_FROM_CHUNK(c));
-	}
-
-	tc->flags &= ~TALLOC_FLAG_LOOP;
-
-	return total;
+	return _talloc_total_mem_internal(ptr, TOTAL_MEM_BLOCKS, NULL, NULL);
 }
 
 /*
@@ -1734,6 +1882,7 @@ _PUBLIC_ void talloc_report_depth_cb(const void *ptr, int depth, int max_depth,
 static void talloc_report_depth_FILE_helper(const void *ptr, int depth, int max_depth, int is_ref, void *_f)
 {
 	const char *name = talloc_get_name(ptr);
+	struct talloc_chunk *tc;
 	FILE *f = (FILE *)_f;
 
 	if (is_ref) {
@@ -1741,6 +1890,16 @@ static void talloc_report_depth_FILE_helper(const void *ptr, int depth, int max_
 		return;
 	}
 
+	tc = talloc_chunk_from_ptr(ptr);
+	if (tc->limit && tc->limit->parent == tc) {
+		fprintf(f, "%*s%-30s is a memlimit context"
+			" (max_size = %lu bytes, cur_size = %lu bytes)\n",
+			depth*4, "",
+			name,
+			(unsigned long)tc->limit->max_size,
+			(unsigned long)tc->limit->cur_size);
+	}
+
 	if (depth == 0) {
 		fprintf(f,"%stalloc report on '%s' (total %6lu bytes in %3lu blocks)\n", 
 			(max_depth < 0 ? "full " :""), name,
@@ -2359,3 +2518,43 @@ _PUBLIC_ int talloc_is_parent(const void *context, const void *ptr)
 {
 	return _talloc_is_parent(context, ptr, TALLOC_MAX_DEPTH);
 }
+
+/*
+  return the total size of memory used by this context and all children
+*/
+static size_t _talloc_total_limit_size(const void *ptr,
+					struct talloc_memlimit *old_limit,
+					struct talloc_memlimit *new_limit)
+{
+	return _talloc_total_mem_internal(ptr, TOTAL_MEM_LIMIT,
+					  old_limit, new_limit);
+}
+
+_PUBLIC_ int talloc_set_memlimit(const void *ctx, size_t max_size)
+{
+	struct talloc_chunk *tc = talloc_chunk_from_ptr(ctx);
+	struct talloc_memlimit *orig_limit;
+	struct talloc_memlimit *limit = NULL;
+
+	if (tc->limit && tc->limit->parent == tc) {
+		tc->limit->max_size = max_size;
+		return 0;
+	}
+	orig_limit = tc->limit;
+
+	limit = malloc(sizeof(struct talloc_memlimit));
+	if (limit == NULL) {
+		return 1;
+	}
+	limit->parent = tc;
+	limit->max_size = max_size;
+	limit->cur_size = _talloc_total_limit_size(ctx, tc->limit, limit);
+
+	if (orig_limit) {
+		limit->upper = orig_limit;
+	} else {
+		limit->upper = NULL;
+	}
+
+	return 0;
+}
diff --git a/lib/talloc/talloc.h b/lib/talloc/talloc.h
index e48dc09a400..f3cbcd0e7c0 100644
--- a/lib/talloc/talloc.h
+++ b/lib/talloc/talloc.h
@@ -1842,6 +1842,25 @@ void talloc_set_log_fn(void (*log_fn)(const char *message));
  */
 void talloc_set_log_stderr(void);
 
+/**
+ * @brief Set a max memory limit for the current context hierarchy
+ *	  This affects all children of this context and constrain any
+ *	  allocation in the hierarchy to never exceed the limit set.
+ *	  The limit can be removed by setting 0 (unlimited) as the
+ *	  max_size by calling the funciton again on the sam context.
+ *	  Memory limits can also be nested, meaning a hild can have
+ *	  a stricter memory limit than a parent.
+ *	  Memory limits are enforced only at memory allocation time.
+ *	  Stealing a context into a 'limited' hierarchy properly
+ *	  updates memory usage but does *not* cause failure if the
+ *	  move causes the new parent to exceed its limits. However
+ *	  any further allocation on that hierarchy will then fail.
+ *
+ * @param[in]	ctx		The talloc context to set the limit on
+ * @param[in]	max_size	The (new) max_size
+ */
+int talloc_set_memlimit(const void *ctx, size_t max_size);
+
 /* @} ******************************************************************/
 
 #if TALLOC_DEPRECATED
diff --git a/lib/talloc/wscript b/lib/talloc/wscript
index a97b5d0b047..bd7708e47d0 100644
--- a/lib/talloc/wscript
+++ b/lib/talloc/wscript
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 
 APPNAME = 'talloc'
-VERSION = '2.0.7'
+VERSION = '2.0.8'
 
 
 blddir = 'bin'