External allocator support.

ccan/talloc/talloc.c

@@ -45,6 +45,7 @@
 #define TALLOC_MAGIC 0xe814ec70
 #define TALLOC_FLAG_FREE 0x01
 #define TALLOC_FLAG_LOOP 0x02
+#define TALLOC_FLAG_EXT_ALLOC 0x04
 #define TALLOC_MAGIC_REFERENCE ((const char *)1)
 
 /* by default we abort when given a bad pointer (such as when talloc_free() is called 
@@ -79,6 +80,10 @@
 static void *null_context;
 static void *autofree_context;
 
+static void *(*tc_external_alloc)(void *parent, size_t size);
+static void (*tc_external_free)(void *ptr, void *parent);
+static void *(*tc_external_realloc)(void *ptr, void *parent, size_t size);
+
 struct talloc_reference_handle {
 	struct talloc_reference_handle *next, *prev;
 	void *ptr;
@@ -181,6 +186,8 @@ const char *talloc_parent_name(const void *ptr)
 static inline void *__talloc(const void *context, size_t size)
 {
 	struct talloc_chunk *tc;
+	struct talloc_chunk *parent = NULL; /* Prevent spurious gcc warning */
+	unsigned flags = TALLOC_MAGIC;
 
 	if (unlikely(context == NULL)) {
 		context = null_context;
@@ -190,19 +197,28 @@ static inline void *__talloc(const void *context, size_t size)
 		return NULL;
 	}
 
+	if (likely(context)) {
+		parent = talloc_chunk_from_ptr(context);
+		if (unlikely(parent->flags & TALLOC_FLAG_EXT_ALLOC)) {
+			tc = tc_external_alloc(TC_PTR_FROM_CHUNK(parent),
+					       TC_HDR_SIZE+size);
+			flags |= TALLOC_FLAG_EXT_ALLOC;
+			goto alloc_done;
+		}
+	}
+
 	tc = (struct talloc_chunk *)malloc(TC_HDR_SIZE+size);
+alloc_done:
 	if (unlikely(tc == NULL)) return NULL;
 
 	tc->size = size;
-	tc->flags = TALLOC_MAGIC;
+	tc->flags = flags;
 	tc->destructor = NULL;
 	tc->child = NULL;
 	tc->name = NULL;
 	tc->refs = NULL;
 
 	if (likely(context)) {
-		struct talloc_chunk *parent = talloc_chunk_from_ptr(context);
-
 		if (parent->child) {
 			parent->child->parent = NULL;
 			tc->next = parent->child;
@@ -319,6 +335,7 @@ void *_talloc_reference(const void *context, const void *ptr)
 static inline int _talloc_free(void *ptr)
 {
 	struct talloc_chunk *tc;
+	void *oldparent = NULL;
 
 	if (unlikely(ptr == NULL)) {
 		return -1;
@@ -362,6 +379,7 @@ static inline int _talloc_free(void *ptr)
 	}
 
 	if (tc->parent) {
+		oldparent = TC_PTR_FROM_CHUNK(tc->parent);
 		_TLIST_REMOVE(tc->parent->child, tc);
 		if (tc->parent->child) {
 			tc->parent->child->parent = tc->parent;
@@ -395,7 +413,12 @@ static inline int _talloc_free(void *ptr)
 	}
 
 	tc->flags |= TALLOC_FLAG_FREE;
+
+	if (unlikely(tc->flags & TALLOC_FLAG_EXT_ALLOC))
+		tc_external_free(tc, oldparent);
+	else
 		free(tc);
+
 	return 0;
 }
 
@@ -771,6 +794,12 @@ void *_talloc_realloc(const void *context, void *ptr, size_t size, const char *n
 		return NULL;
 	}
 
+	if (unlikely(tc->flags & TALLOC_FLAG_EXT_ALLOC)) {
+		/* need to get parent before setting free flag. */
+		void *parent = talloc_parent(ptr);
+		tc->flags |= TALLOC_FLAG_FREE;
+		new_ptr = tc_external_realloc(tc, parent, size + TC_HDR_SIZE);
+	} else {
 		/* by resetting magic we catch users of the old memory */
 		tc->flags |= TALLOC_FLAG_FREE;
 
@@ -783,6 +812,8 @@ void *_talloc_realloc(const void *context, void *ptr, size_t size, const char *n
 #else
 		new_ptr = realloc(tc, size + TC_HDR_SIZE);
 #endif
+	}
+
 	if (unlikely(!new_ptr)) {	
 		tc->flags &= ~TALLOC_FLAG_FREE; 
 		return NULL; 
@@ -1401,3 +1432,24 @@ int talloc_is_parent(const void *context, const void *ptr)
 	}
 	return 0;
 }
+
+void talloc_external_enable(void *(*alloc)(void *parent, size_t size),
+			    void (*free)(void *ptr, void *parent),
+			    void *(*realloc)(void *ptr, void *parent, size_t))
+{
+	tc_external_alloc = alloc;
+	tc_external_free = free;
+	tc_external_realloc = realloc;
+}
+
+void talloc_mark_external(void *context)
+{
+	struct talloc_chunk *tc;
+
+	if (unlikely(context == NULL)) {
+		context = null_context;
+	}
+
+	tc = talloc_chunk_from_ptr(context);
+	tc->flags |= TALLOC_FLAG_EXT_ALLOC;
+}

ccan/talloc/talloc.h

@@ -928,6 +928,32 @@ size_t talloc_get_size(const void *ctx);
  */
 void *talloc_find_parent_byname(const void *ctx, const char *name);
 
+/**
+ * talloc_external_enable - set external allocators for some nodes
+ * @alloc: the malloc() equivalent
+ * @free: the free() equivalent
+ * @realloc: the realloc() equivalent
+ *
+ * talloc_mark_external() can be used to mark nodes whose children should
+ * use separate allocators.  Currently the set of allocators is global, not
+ * per-node, and is set with this function.
+ *
+ * The parent pointers is the talloc pointer of the parent.
+ */
+void talloc_external_enable(void *(*alloc)(void *parent, size_t size),
+			    void (*free)(void *ptr, void *parent),
+			    void *(*realloc)(void *ptr, void *parent, size_t));
+
+/**
+ * talloc_mark_external - children of this note must use external allocators
+ * @p: the talloc pointer
+ *
+ * This function indicates that all children (and children's children etc)
+ * should use the allocators set up wth talloc_external_enable() rather than
+ * normal malloc/free.
+ */
+void talloc_mark_external(void *ptr);
+
 /* The following definitions come from talloc.c  */
 void *_talloc(const void *context, size_t size);
 void _talloc_set_destructor(const void *ptr, int (*destructor)(void *));

ccan/talloc/test/run-external-alloc.c

+#include "talloc/talloc.c"
+#include "tap/tap.h"
+#include <assert.h>
+
+static int ext_alloc_count, ext_free_count, ext_realloc_count;
+static void *expected_parent;
+
+static void *ext_alloc(void *parent, size_t size)
+{
+	ok1(parent == expected_parent);
+	ext_alloc_count++;
+	return malloc(size);
+}
+
+static void ext_free(void *ptr, void *parent)
+{
+	ok1(parent == expected_parent);
+	ext_free_count++;
+	free(ptr);
+}
+
+static void *ext_realloc(void *ptr, void *parent, size_t size)
+{
+	ok1(parent == expected_parent);
+	ext_realloc_count++;
+	return realloc(ptr, size);
+}
+
+int main(void)
+{
+	char *p, *p2, *head;
+	plan_tests(10);
+
+	talloc_external_enable(ext_alloc, ext_free, ext_realloc);
+	head = talloc(NULL, char);
+	assert(head);
+	expected_parent = head;
+
+	talloc_mark_external(head);
+
+	p = talloc_array(head, char, 1);
+	ok1(ext_alloc_count == 1);
+	assert(p);
+
+	/* Child is also externally allocated */
+	expected_parent = p;
+	p2 = talloc(p, char);
+	ok1(ext_alloc_count == 2);
+
+	expected_parent = head;
+	p = talloc_realloc(NULL, p, char, 1000);
+	ok1(ext_realloc_count == 1);
+	assert(p);
+
+	expected_parent = p;
+	talloc_free(p2);
+	ok1(ext_free_count == 1);
+
+	expected_parent = head;
+	talloc_free(p);
+	ok1(ext_free_count == 2);
+
+	return exit_status();
+}