talloc: allow replacement allocator

This allows us to both allocators which handle failure themselves, and
allocators which insert failures.

diff --git a/ccan/talloc/talloc.c b/ccan/talloc/talloc.c
index 63423a50f5ec05e73aa1816910fdc82860fbb2d0..0c6bcda955165e5cb088ec9dadd9da5863c95676 100644 (file)
--- a/ccan/talloc/talloc.c
+++ b/ccan/talloc/talloc.c
@@ -82,6 +82,10 @@
 static void *null_context;
 static pid_t *autofree_context;
 
+static void *(*tc_malloc)(size_t size) = malloc;
+static void (*tc_free)(void *ptr) = free;
+static void *(*tc_realloc)(void *ptr, size_t size) = realloc;
+
 static void *(*tc_external_realloc)(const void *parent, void *ptr, size_t size);
 static void (*tc_lock)(const void *ctx);
 static void (*tc_unlock)(void);
@@ -283,7 +287,7 @@ static inline void *__talloc(const void *context, size_t size)
                }
        }
 
-       tc = (struct talloc_chunk *)malloc(TC_HDR_SIZE+size);
+       tc = (struct talloc_chunk *)tc_malloc(TC_HDR_SIZE+size);
 alloc_done:
        return init_talloc(parent, tc, size, external);
 }
@@ -571,7 +575,7 @@ static inline int _talloc_free(const void *ptr)
        if (unlikely(tc->flags & TALLOC_FLAG_EXT_ALLOC))
                tc_external_realloc(oldparent, tc, 0);
        else
-               free(tc);
+               tc_free(tc);
 
        return 0;
 }
@@ -923,13 +927,13 @@ void *_talloc_realloc(const void *context, void *ptr, size_t size, const char *n
                tc->flags |= TALLOC_FLAG_FREE;
 
 #if ALWAYS_REALLOC
-               new_ptr = malloc(size + TC_HDR_SIZE);
+               new_ptr = tc_malloc(size + TC_HDR_SIZE);
                if (new_ptr) {
                        memcpy(new_ptr, tc, tc->size + TC_HDR_SIZE);
-                       free(tc);
+                       tc_free(tc);
                }
 #else
-               new_ptr = realloc(tc, size + TC_HDR_SIZE);
+               new_ptr = tc_realloc(tc, size + TC_HDR_SIZE);
 #endif
        }
 
@@ -1613,6 +1617,15 @@ int talloc_is_parent(const void *context, const void *ptr)
        return ret;
 }
 
+void talloc_set_allocator(void *(*malloc)(size_t size),
+                         void (*free)(void *ptr),
+                         void *(*realloc)(void *ptr, size_t size))
+{
+       tc_malloc = malloc;
+       tc_free = free;
+       tc_realloc = realloc;
+}
+
 void *talloc_add_external(const void *ctx,
                          void *(*realloc)(const void *, void *, size_t),
                          void (*lock)(const void *p),

diff --git a/ccan/talloc/talloc.h b/ccan/talloc/talloc.h
index d27c68922c9c8f0d98a0b5fe76d3ea4b8093d3fb..5263b9fed269261f833b2e9d55765a7242ca964b 100644 (file)
--- a/ccan/talloc/talloc.h
+++ b/ccan/talloc/talloc.h
@@ -957,6 +957,19 @@ size_t talloc_get_size(const void *ctx);
  */
 void *talloc_find_parent_byname(const void *ctx, const char *name);
 
+/**
+ * talloc_set_allocator - set the allocations function(s) for talloc.
+ * @malloc: the malloc function
+ * @free: the free function
+ * @realloc: the realloc function
+ *
+ * Instead of using the standard malloc, free and realloc, talloc will use
+ * these replacements.  @realloc will never be called with size 0 or ptr NULL.
+ */
+void talloc_set_allocator(void *(*malloc)(size_t size),
+                         void (*free)(void *ptr),
+                         void *(*realloc)(void *ptr, size_t size));
+
 /**
  * talloc_add_external - create an externally allocated node
  * @ctx: the parent

diff --git a/ccan/talloc/test/run-set_allocator.c b/ccan/talloc/test/run-set_allocator.c
new file mode 100644 (file)
index 0000000..e485c62
--- /dev/null
+++ b/ccan/talloc/test/run-set_allocator.c
@@ -0,0 +1,57 @@
+#include <ccan/failtest/failtest_override.h>
+#include <ccan/talloc/talloc.c>
+#include <stdbool.h>
+#include <ccan/tap/tap.h>
+#include <ccan/failtest/failtest.h>
+
+static unsigned my_malloc_count, my_free_count, my_realloc_count;
+
+static void *my_malloc(size_t size)
+{
+       my_malloc_count++;
+       return malloc(size);
+}
+
+static void my_free(void *ptr)
+{
+       my_free_count++;
+       free(ptr);
+}
+
+static void *my_realloc(void *ptr, size_t size)
+{
+       my_realloc_count++;
+       ok1(ptr);
+       ok1(size);
+       return realloc(ptr, size);
+}
+
+int main(int argc, char *argv[])
+{
+       int *p1, *p2;
+
+       plan_tests(14);
+       failtest_init(argc, argv);
+       talloc_set_allocator(my_malloc, my_free, my_realloc);
+       p1 = talloc_array(NULL, int, 10);
+       ok1(my_malloc_count == 1);
+       ok1(my_free_count == 0);
+       ok1(my_realloc_count == 0);
+
+       p1 = talloc_realloc(NULL, p1, int, 10000);
+       ok1(my_malloc_count == 1);
+       ok1(my_free_count == 0);
+       ok1(my_realloc_count == 1);
+
+       p2 = talloc(p1, int);
+       ok1(my_malloc_count == 2);
+       ok1(my_free_count == 0);
+       ok1(my_realloc_count == 1);
+
+       talloc_free(p1);
+       ok1(my_malloc_count == 2);
+       ok1(my_free_count == 2);
+       ok1(my_realloc_count == 1);
+
+       failtest_exit(exit_status());
+}