From 23e72d4194172b52d26fed9f36c7af9b893f4649 Mon Sep 17 00:00:00 2001
From: Rusty Russell <rusty@rustcorp.com.au>
Date: Mon, 4 Mar 2019 19:42:57 +1030
Subject: [PATCH] timer: add hook for allocation functions.

malloc and free aren't for everyone.

Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>
---
 ccan/timer/test/run-allocator.c | 62 +++++++++++++++++++++++++++++++++
 ccan/timer/timer.c              | 31 +++++++++++++++--
 ccan/timer/timer.h              | 13 +++++++
 3 files changed, 104 insertions(+), 2 deletions(-)
 create mode 100644 ccan/timer/test/run-allocator.c

diff --git a/ccan/timer/test/run-allocator.c b/ccan/timer/test/run-allocator.c
new file mode 100644
index 00000000..75e1af06
--- /dev/null
+++ b/ccan/timer/test/run-allocator.c
@@ -0,0 +1,62 @@
+#define CCAN_TIMER_DEBUG
+/* Include the C files directly. */
+#include <ccan/timer/timer.c>
+#include <ccan/tap/tap.h>
+
+static void *test_alloc(size_t len, void *arg)
+{
+	(*(size_t *)arg)++;
+	return malloc(len);
+}
+
+static void test_free(const void *p, void *arg)
+{
+	if (p) {
+		(*(size_t *)arg)--;
+		free((void *)p);
+	}
+}
+
+static struct timemono timemono_from_nsec(unsigned long long nsec)
+{
+	struct timemono epoch = { { 0, 0 } };
+	return timemono_add(epoch, time_from_nsec(nsec));
+}
+
+int main(void)
+{
+	struct timers timers;
+	struct timer t[64];
+	size_t num_allocs = 0;
+	const struct timemono epoch = { { 0, 0 } };
+
+	plan_tests(7);
+
+	timers_set_allocator(test_alloc, test_free, &num_allocs);
+	timers_init(&timers, epoch);
+	timer_init(&t[0]);
+
+	timer_addmono(&timers, &t[0],
+		      timemono_from_nsec(TIMER_GRANULARITY << TIMER_LEVEL_BITS));
+	timers_expire(&timers, timemono_from_nsec(1));
+	ok1(num_allocs == 1);
+	timer_del(&timers, &t[0]);
+	ok1(num_allocs == 1);
+	timers_cleanup(&timers);
+	ok1(num_allocs == 0);
+
+	/* Should restore defaults */
+	timers_set_allocator(NULL, NULL, NULL);
+	ok1(timer_alloc == timer_default_alloc);
+	ok1(timer_free == timer_default_free);
+
+	timers_init(&timers, epoch);
+	timer_addmono(&timers, &t[0],
+		      timemono_from_nsec(TIMER_GRANULARITY << TIMER_LEVEL_BITS));
+	ok1(num_allocs == 0);
+	timers_cleanup(&timers);
+	ok1(num_allocs == 0);
+
+	/* This exits depending on whether all tests passed */
+	return exit_status();
+}
diff --git a/ccan/timer/timer.c b/ccan/timer/timer.c
index 84d7aa14..60741562 100644
--- a/ccan/timer/timer.c
+++ b/ccan/timer/timer.c
@@ -11,6 +11,33 @@ struct timer_level {
 	struct list_head list[PER_LEVEL];
 };
 
+static void *timer_default_alloc(size_t len, void *arg)
+{
+	return malloc(len);
+}
+
+static void timer_default_free(const void *p, void *arg)
+{
+	free((void *)p);
+}
+
+static void *(*timer_alloc)(size_t, void *) = timer_default_alloc;
+static void (*timer_free)(const void *, void *) = timer_default_free;
+static void *timer_arg;
+
+void timers_set_allocator(void *(*alloc)(size_t len, void *arg),
+			  void (*free)(const void *p, void *arg),
+			  void *arg)
+{
+	if (!alloc)
+		alloc = timer_default_alloc;
+	if (!free)
+		free = timer_default_free;
+	timer_alloc = alloc;
+	timer_free = free;
+	timer_arg = arg;
+}
+
 static uint64_t time_to_grains(struct timemono t)
 {
 	return t.ts.tv_sec * ((uint64_t)1000000000 / TIMER_GRANULARITY)
@@ -139,7 +166,7 @@ static void add_level(struct timers *timers, unsigned int level)
 	unsigned int i;
 	struct list_head from_far;
 
-	l = malloc(sizeof(*l));
+	l = timer_alloc(sizeof(*l), timer_arg);
 	if (!l)
 		return;
 
@@ -520,5 +547,5 @@ void timers_cleanup(struct timers *timers)
 	unsigned int l;
 
 	for (l = 0; l < ARRAY_SIZE(timers->level); l++)
-		free(timers->level[l]);
+		timer_free(timers->level[l], timer_arg);
 }
diff --git a/ccan/timer/timer.h b/ccan/timer/timer.h
index 5c40a3bb..5b193250 100644
--- a/ccan/timer/timer.h
+++ b/ccan/timer/timer.h
@@ -162,6 +162,19 @@ struct timer *timers_expire(struct timers *timers, struct timemono expire);
  */
 struct timers *timers_check(const struct timers *t, const char *abortstr);
 
+/**
+ * timers_set_allocator - set malloc/free functions.
+ * @alloc: allocator to use
+ * @free: unallocator to use (@p is NULL or a return from @alloc)
+ * @arg: argument to pass.
+ *
+ * This replaces the underlying malloc/free with these allocators.
+ * Setting either one to NULL restores the default allocators.
+ */
+void timers_set_allocator(void *(*alloc)(size_t len, void *arg),
+			  void (*free)(const void *p, void *arg),
+			  void *arg);
+
 #ifdef CCAN_TIMER_DEBUG
 #include <stdio.h>
 
-- 
2.39.2