ntdb: allocator attribute.

This is designed to allow us to make ntdb_context (and NTDB_DATA returned
from ntdb_fetch) a talloc pointer.  But it can also be used for any other
alternate allocator.

Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>

15 files changed, 311 insertions, 92 deletions

diff --git a/lib/ntdb/check.c b/lib/ntdb/check.c
index cff6e0345e..723e7b11bf 100644
--- a/lib/ntdb/check.c
+++ b/lib/ntdb/check.c
@@ -20,9 +20,17 @@
 #include <ccan/asearch/asearch.h>
 
 /* We keep an ordered array of offsets. */
-static bool append(ntdb_off_t **arr, size_t *num, ntdb_off_t off)
+static bool append(struct ntdb_context *ntdb,
+		   ntdb_off_t **arr, size_t *num, ntdb_off_t off)
 {
-	ntdb_off_t *new = realloc(*arr, (*num + 1) * sizeof(ntdb_off_t));
+	ntdb_off_t *new;
+
+	if (*num == 0) {
+		new = ntdb->alloc_fn(ntdb, sizeof(ntdb_off_t), ntdb->alloc_data);
+	} else {
+		new = ntdb->expand_fn(*arr, (*num + 1) * sizeof(ntdb_off_t),
+				  ntdb->alloc_data);
+	}
 	if (!new)
 		return false;
 	new[(*num)++] = off;
@@ -708,7 +716,7 @@ static enum NTDB_ERROR check_linear(struct ntdb_context *ntdb,
 			}
 			/* This record should be in free lists. */
 			if (frec_ftable(&rec.f) != NTDB_FTABLE_NONE
-			    && !append(fr, num_free, off)) {
+			    && !append(ntdb, fr, num_free, off)) {
 				return ntdb_logerr(ntdb, NTDB_ERR_OOM,
 						  NTDB_LOG_ERROR,
 						  "ntdb_check: tracking %zu'th"
@@ -722,7 +730,7 @@ static enum NTDB_ERROR check_linear(struct ntdb_context *ntdb,
 			uint64_t klen, dlen, extra;
 
 			/* This record is used! */
-			if (!append(used, num_used, off)) {
+			if (!append(ntdb, used, num_used, off)) {
 				return ntdb_logerr(ntdb, NTDB_ERR_OOM,
 						  NTDB_LOG_ERROR,
 						  "ntdb_check: tracking %zu'th"
@@ -858,7 +866,7 @@ _PUBLIC_ enum NTDB_ERROR ntdb_check_(struct ntdb_context *ntdb,
 out:
 	ntdb_allrecord_unlock(ntdb, F_RDLCK);
 	ntdb_unlock_expand(ntdb, F_RDLCK);
-	free(fr);
-	free(used);
+	ntdb->free_fn(fr, ntdb->alloc_data);
+	ntdb->free_fn(used, ntdb->alloc_data);
 	return ecode;
 }
diff --git a/lib/ntdb/io.c b/lib/ntdb/io.c
index e1518062b1..b20141cc8f 100644
--- a/lib/ntdb/io.c
+++ b/lib/ntdb/io.c
@@ -120,7 +120,7 @@ static enum NTDB_ERROR ntdb_oob(struct ntdb_context *ntdb,
 
 		ntdb_logerr(ntdb, NTDB_ERR_IO, NTDB_LOG_ERROR,
 			   "ntdb_oob len %lld beyond internal"
-			   " malloc size %lld",
+			   " alloc size %lld",
 			   (long long)(off + len),
 			   (long long)ntdb->file->map_size);
 		return NTDB_ERR_IO;
@@ -336,7 +336,7 @@ enum NTDB_ERROR ntdb_write_convert(struct ntdb_context *ntdb, ntdb_off_t off,
 	enum NTDB_ERROR ecode;
 
 	if (unlikely((ntdb->flags & NTDB_CONVERT))) {
-		void *conv = malloc(len);
+		void *conv = ntdb->alloc_fn(ntdb, len, ntdb->alloc_data);
 		if (!conv) {
 			return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 					  "ntdb_write: no memory converting"
@@ -344,8 +344,8 @@ enum NTDB_ERROR ntdb_write_convert(struct ntdb_context *ntdb, ntdb_off_t off,
 		}
 		memcpy(conv, rec, len);
 		ecode = ntdb->io->twrite(ntdb, off,
-					ntdb_convert(ntdb, conv, len), len);
-		free(conv);
+					 ntdb_convert(ntdb, conv, len), len);
+		ntdb->free_fn(conv, ntdb->alloc_data);
 	} else {
 		ecode = ntdb->io->twrite(ntdb, off, rec, len);
 	}
@@ -388,16 +388,17 @@ static void *_ntdb_alloc_read(struct ntdb_context *ntdb, ntdb_off_t offset,
 	enum NTDB_ERROR ecode;
 
 	/* some systems don't like zero length malloc */
-	buf = malloc(prefix + len ? prefix + len : 1);
+	buf = ntdb->alloc_fn(ntdb, prefix + len ? prefix + len : 1,
+			  ntdb->alloc_data);
 	if (!buf) {
 		ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_USE_ERROR,
-			   "ntdb_alloc_read malloc failed len=%zu",
+			   "ntdb_alloc_read alloc failed len=%zu",
 			   (size_t)(prefix + len));
 		return NTDB_ERR_PTR(NTDB_ERR_OOM);
 	} else {
 		ecode = ntdb->io->tread(ntdb, offset, buf+prefix, len);
 		if (unlikely(ecode != NTDB_SUCCESS)) {
-			free(buf);
+			ntdb->free_fn(buf, ntdb->alloc_data);
 			return NTDB_ERR_PTR(ecode);
 		}
 	}
@@ -448,8 +449,9 @@ static enum NTDB_ERROR ntdb_expand_file(struct ntdb_context *ntdb,
 	}
 
 	if (ntdb->flags & NTDB_INTERNAL) {
-		char *new = realloc(ntdb->file->map_ptr,
-				    ntdb->file->map_size + addition);
+		char *new = ntdb->expand_fn(ntdb->file->map_ptr,
+					ntdb->file->map_size + addition,
+					ntdb->alloc_data);
 		if (!new) {
 			return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 					  "No memory to expand database");
@@ -566,7 +568,7 @@ void ntdb_access_release(struct ntdb_context *ntdb, const void *p)
 	if (hp) {
 		hdr = *hp;
 		*hp = hdr->next;
-		free(hdr);
+		ntdb->free_fn(hdr, ntdb->alloc_data);
 	} else
 		ntdb->direct_access--;
 }
@@ -583,7 +585,7 @@ enum NTDB_ERROR ntdb_access_commit(struct ntdb_context *ntdb, void *p)
 		else
 			ecode = ntdb_write(ntdb, hdr->off, p, hdr->len);
 		*hp = hdr->next;
-		free(hdr);
+		ntdb->free_fn(hdr, ntdb->alloc_data);
 	} else {
 		ntdb->direct_access--;
 		ecode = NTDB_SUCCESS;
diff --git a/lib/ntdb/lock.c b/lib/ntdb/lock.c
index 625a2c4934..95824db742 100644
--- a/lib/ntdb/lock.c
+++ b/lib/ntdb/lock.c
@@ -387,10 +387,16 @@ static enum NTDB_ERROR ntdb_nest_lock(struct ntdb_context *ntdb,
 				  "ntdb_nest_lock: already have a hash lock?");
 	}
 #endif
-
-	new_lck = (struct ntdb_lock *)realloc(
-		ntdb->file->lockrecs,
-		sizeof(*ntdb->file->lockrecs) * (ntdb->file->num_lockrecs+1));
+	if (ntdb->file->lockrecs == NULL) {
+		new_lck = ntdb->alloc_fn(ntdb->file, sizeof(*ntdb->file->lockrecs),
+				     ntdb->alloc_data);
+	} else {
+		new_lck = (struct ntdb_lock *)ntdb->expand_fn(
+			ntdb->file->lockrecs,
+			sizeof(*ntdb->file->lockrecs)
+			* (ntdb->file->num_lockrecs+1),
+			ntdb->alloc_data);
+	}
 	if (new_lck == NULL) {
 		return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 				  "ntdb_nest_lock:"
diff --git a/lib/ntdb/ntdb.c b/lib/ntdb/ntdb.c
index df9fe709b2..a74e0f4b78 100644
--- a/lib/ntdb/ntdb.c
+++ b/lib/ntdb/ntdb.c
@@ -204,7 +204,8 @@ _PUBLIC_ enum NTDB_ERROR ntdb_append(struct ntdb_context *ntdb,
 		}
 
 		/* Slow path. */
-		newdata = malloc(key.dsize + old_dlen + dbuf.dsize);
+		newdata = ntdb->alloc_fn(ntdb, key.dsize + old_dlen + dbuf.dsize,
+				     ntdb->alloc_data);
 		if (!newdata) {
 			ecode = ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 					   "ntdb_append:"
@@ -230,7 +231,7 @@ _PUBLIC_ enum NTDB_ERROR ntdb_append(struct ntdb_context *ntdb,
 	ecode = replace_data(ntdb, &h, key, new_dbuf, off, old_room, true);
 
 out_free_newdata:
-	free(newdata);
+	ntdb->free_fn(newdata, ntdb->alloc_data);
 out:
 	ntdb_unlock_hashes(ntdb, h.hlock_start, h.hlock_range, F_WRLCK);
 	return ecode;
@@ -454,16 +455,20 @@ enum NTDB_ERROR COLD ntdb_logerr(struct ntdb_context *ntdb,
 		return ecode;
 
 	va_start(ap, fmt);
-	len = vasprintf(&message, fmt, ap);
+	len = vsnprintf(NULL, 0, fmt, ap);
 	va_end(ap);
 
-	if (len < 0) {
+	message = ntdb->alloc_fn(ntdb, len + 1, ntdb->alloc_data);
+	if (!message) {
 		ntdb->log_fn(ntdb, NTDB_LOG_ERROR, NTDB_ERR_OOM,
 			    "out of memory formatting message:", ntdb->log_data);
 		ntdb->log_fn(ntdb, level, ecode, fmt, ntdb->log_data);
 	} else {
+		va_start(ap, fmt);
+		vsnprintf(message, len+1, fmt, ap);
+		va_end(ap);
 		ntdb->log_fn(ntdb, level, ecode, message, ntdb->log_data);
-		free(message);
+		ntdb->free_fn(message, ntdb->alloc_data);
 	}
 	errno = saved_errno;
 	return ecode;
diff --git a/lib/ntdb/ntdb.h b/lib/ntdb/ntdb.h
index 1a011fddf5..87f3f5bbfa 100644
--- a/lib/ntdb/ntdb.h
+++ b/lib/ntdb/ntdb.h
@@ -633,6 +633,7 @@ enum ntdb_attribute_type {
 	NTDB_ATTRIBUTE_STATS = 3,
 	NTDB_ATTRIBUTE_OPENHOOK = 4,
 	NTDB_ATTRIBUTE_FLOCK = 5,
+	NTDB_ATTRIBUTE_ALLOCATOR = 6
 };
 
 /**
@@ -865,6 +866,30 @@ struct ntdb_attribute_flock {
 };
 
 /**
+ * struct ntdb_attribute_allocator - allocator for ntdb to use.
+ *
+ * You can replace malloc/free with your own allocation functions.
+ * The allocator takes an "owner" pointer, which is either NULL (for
+ * the initial struct ntdb_context and struct ntdb_file), or a
+ * previously allocated pointer.  This is useful for relationship
+ * tracking, such as the talloc library.
+ *
+ * The expand function is realloc, but only ever used to expand an
+ * existing allocation.
+ *
+ * Be careful mixing allocators: two ntdb_contexts which have the same file
+ * open will share the same struct ntdb_file.  This may be allocated by one
+ * ntdb's allocator, and freed by the other.
+ */
+struct ntdb_attribute_allocator {
+	struct ntdb_attribute_base base; /* .attr = NTDB_ATTRIBUTE_ALLOCATOR */
+	void *(*alloc)(const void *owner, size_t len, void *priv_data);
+	void *(*expand)(void *old, size_t newlen, void *priv_data);
+	void (*free)(void *old, void *priv_data);
+	void *priv_data;
+};
+
+/**
  * union ntdb_attribute - ntdb attributes.
  *
  * This represents all the known attributes.
@@ -872,7 +897,8 @@ struct ntdb_attribute_flock {
  * See also:
  *	struct ntdb_attribute_log, struct ntdb_attribute_hash,
  *	struct ntdb_attribute_seed, struct ntdb_attribute_stats,
- *	struct ntdb_attribute_openhook, struct ntdb_attribute_flock.
+ *	struct ntdb_attribute_openhook, struct ntdb_attribute_flock,
+ *	struct ntdb_attribute_allocator alloc.
  */
 union ntdb_attribute {
 	struct ntdb_attribute_base base;
@@ -882,6 +908,7 @@ union ntdb_attribute {
 	struct ntdb_attribute_stats stats;
 	struct ntdb_attribute_openhook openhook;
 	struct ntdb_attribute_flock flock;
+	struct ntdb_attribute_allocator alloc;
 };
 
 #ifdef  __cplusplus
diff --git a/lib/ntdb/open.c b/lib/ntdb/open.c
index 01a0928074..dc473d2680 100644
--- a/lib/ntdb/open.c
+++ b/lib/ntdb/open.c
@@ -126,7 +126,7 @@ static enum NTDB_ERROR ntdb_new_database(struct ntdb_context *ntdb,
 	/* Always make db a multiple of NTDB_PGSIZE */
 	dbsize = (sizeof(*newdb) + NTDB_PGSIZE-1) & ~(NTDB_PGSIZE-1);
 	remaindersize = dbsize - sizeof(*newdb);
-	newdb = malloc(dbsize);
+	newdb = ntdb->alloc_fn(ntdb, dbsize, ntdb->alloc_data);
 	if (!newdb) {
 		return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 				   "ntdb_new_database: failed to allocate");
@@ -218,13 +218,13 @@ static enum NTDB_ERROR ntdb_new_database(struct ntdb_context *ntdb,
 	}
 
 out:
-	free(newdb);
+	ntdb->free_fn(newdb, ntdb->alloc_data);
 	return ecode;
 }
 
 static enum NTDB_ERROR ntdb_new_file(struct ntdb_context *ntdb)
 {
-	ntdb->file = malloc(sizeof(*ntdb->file));
+	ntdb->file = ntdb->alloc_fn(NULL, sizeof(*ntdb->file), ntdb->alloc_data);
 	if (!ntdb->file)
 		return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 				  "ntdb_open: cannot alloc ntdb_file structure");
@@ -266,6 +266,12 @@ _PUBLIC_ enum NTDB_ERROR ntdb_set_attribute(struct ntdb_context *ntdb,
 		ntdb->unlock_fn = attr->flock.unlock;
 		ntdb->lock_data = attr->flock.data;
 		break;
+	case NTDB_ATTRIBUTE_ALLOCATOR:
+		ntdb->alloc_fn = attr->alloc.alloc;
+		ntdb->expand_fn = attr->alloc.expand;
+		ntdb->free_fn = attr->alloc.free;
+		ntdb->alloc_data = attr->alloc.priv_data;
+		break;
 	default:
 		return ntdb_logerr(ntdb, NTDB_ERR_EINVAL,
 				   NTDB_LOG_USE_ERROR,
@@ -311,6 +317,12 @@ _PUBLIC_ enum NTDB_ERROR ntdb_get_attribute(struct ntdb_context *ntdb,
 		attr->flock.unlock = ntdb->unlock_fn;
 		attr->flock.data = ntdb->lock_data;
 		break;
+	case NTDB_ATTRIBUTE_ALLOCATOR:
+		attr->alloc.alloc = ntdb->alloc_fn;
+		attr->alloc.expand = ntdb->expand_fn;
+		attr->alloc.free = ntdb->free_fn;
+		attr->alloc.priv_data = ntdb->alloc_data;
+		break;
 	default:
 		return ntdb_logerr(ntdb, NTDB_ERR_EINVAL,
 				   NTDB_LOG_USE_ERROR,
@@ -406,9 +418,40 @@ static enum NTDB_ERROR capabilities_ok(struct ntdb_context *ntdb,
 	return ecode;
 }
 
+static void *default_alloc(const void *owner, size_t len, void *priv_data)
+{
+	return malloc(len);
+}
+
+static void *default_expand(void *ptr, size_t len, void *priv_data)
+{
+	return realloc(ptr, len);
+}
+
+static void default_free(void *ptr, void *priv_data)
+{
+	free(ptr);
+}
+
+/* First allocation needs manual search of attributes. */
+static struct ntdb_context *alloc_ntdb(const union ntdb_attribute *attr,
+				       const char *name)
+{
+	size_t len = sizeof(struct ntdb_context) + strlen(name) + 1;
+
+	while (attr) {
+		if  (attr->base.attr == NTDB_ATTRIBUTE_ALLOCATOR) {
+			return attr->alloc.alloc(NULL, len,
+						 attr->alloc.priv_data);
+		}
+		attr = attr->base.next;
+	}
+	return default_alloc(NULL, len, NULL);
+}
+
 _PUBLIC_ struct ntdb_context *ntdb_open(const char *name, int ntdb_flags,
-			     int open_flags, mode_t mode,
-			     union ntdb_attribute *attr)
+					int open_flags, mode_t mode,
+					union ntdb_attribute *attr)
 {
 	struct ntdb_context *ntdb;
 	struct stat st;
@@ -422,7 +465,7 @@ _PUBLIC_ struct ntdb_context *ntdb_open(const char *name, int ntdb_flags,
 	enum NTDB_ERROR ecode;
 	int openlock;
 
-	ntdb = malloc(sizeof(*ntdb) + strlen(name) + 1);
+	ntdb = alloc_ntdb(attr, name);
 	if (!ntdb) {
 		/* Can't log this */
 		errno = ENOMEM;
@@ -441,6 +484,9 @@ _PUBLIC_ struct ntdb_context *ntdb_open(const char *name, int ntdb_flags,
 	memset(&ntdb->stats, 0, sizeof(ntdb->stats));
 	ntdb->stats.base.attr = NTDB_ATTRIBUTE_STATS;
 	ntdb->stats.size = sizeof(ntdb->stats);
+	ntdb->alloc_fn = default_alloc;
+	ntdb->expand_fn = default_expand;
+	ntdb->free_fn = default_free;
 
 	while (attr) {
 		switch (attr->base.attr) {
@@ -738,7 +784,8 @@ fail_errno:
 			assert(ntdb->file->num_lockrecs == 0);
 			if (ntdb->file->map_ptr) {
 				if (ntdb->flags & NTDB_INTERNAL) {
-					free(ntdb->file->map_ptr);
+					ntdb->free_fn(ntdb->file->map_ptr,
+						      ntdb->alloc_data);
 				} else
 					ntdb_munmap(ntdb->file);
 			}
@@ -746,12 +793,12 @@ fail_errno:
 				ntdb_logerr(ntdb, NTDB_ERR_IO, NTDB_LOG_ERROR,
 					   "ntdb_open: failed to close ntdb fd"
 					   " on error: %s", strerror(errno));
-			free(ntdb->file->lockrecs);
-			free(ntdb->file);
+			ntdb->free_fn(ntdb->file->lockrecs, ntdb->alloc_data);
+			ntdb->free_fn(ntdb->file, ntdb->alloc_data);
 		}
 	}
 
-	free(ntdb);
+	ntdb->free_fn(ntdb, ntdb->alloc_data);
 	errno = saved_errno;
 	return NULL;
 }
@@ -769,7 +816,7 @@ _PUBLIC_ int ntdb_close(struct ntdb_context *ntdb)
 
 	if (ntdb->file->map_ptr) {
 		if (ntdb->flags & NTDB_INTERNAL)
-			free(ntdb->file->map_ptr);
+			ntdb->free_fn(ntdb->file->map_ptr, ntdb->alloc_data);
 		else
 			ntdb_munmap(ntdb->file);
 	}
@@ -777,8 +824,8 @@ _PUBLIC_ int ntdb_close(struct ntdb_context *ntdb)
 		ntdb_lock_cleanup(ntdb);
 		if (--ntdb->file->refcnt == 0) {
 			ret = close(ntdb->file->fd);
-			free(ntdb->file->lockrecs);
-			free(ntdb->file);
+			ntdb->free_fn(ntdb->file->lockrecs, ntdb->alloc_data);
+			ntdb->free_fn(ntdb->file, ntdb->alloc_data);
 		}
 	}
 
@@ -793,7 +840,7 @@ _PUBLIC_ int ntdb_close(struct ntdb_context *ntdb)
 #ifdef NTDB_TRACE
 	close(ntdb->tracefd);
 #endif
-	free(ntdb);
+	ntdb->free_fn(ntdb, ntdb->alloc_data);
 
 	return ret;
 }
diff --git a/lib/ntdb/private.h b/lib/ntdb/private.h
index 01f2b45e9b..ee8eeb76af 100644
--- a/lib/ntdb/private.h
+++ b/lib/ntdb/private.h
@@ -596,6 +596,12 @@ struct ntdb_context {
 	void *hash_data;
 	uint64_t hash_seed;
 
+	/* Allocate and free functions. */
+	void *(*alloc_fn)(const void *owner, size_t len, void *priv_data);
+	void *(*expand_fn)(void *old, size_t newlen, void *priv_data);
+	void (*free_fn)(void *old, void *priv_data);
+	void *alloc_data;
+
 	/* Our open hook, if any. */
 	enum NTDB_ERROR (*openhook)(int fd, void *data);
 	void *openhook_data;
diff --git a/lib/ntdb/summary.c b/lib/ntdb/summary.c
index 571d48ff4d..f5313bec55 100644
--- a/lib/ntdb/summary.c
+++ b/lib/ntdb/summary.c
@@ -261,7 +261,7 @@ _PUBLIC_ enum NTDB_ERROR ntdb_summary(struct ntdb_context *ntdb,
 		+ num_caps * (strlen(CAPABILITY_FORMAT) + 20
 			      + strlen(" (uncheckable,read-only)"));
 
-	*summary = malloc(len);
+	*summary = ntdb->alloc_fn(ntdb, len, ntdb->alloc_data);
 	if (!*summary) {
 		ecode = ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 				   "ntdb_summary: failed to allocate string");
@@ -309,20 +309,20 @@ _PUBLIC_ enum NTDB_ERROR ntdb_summary(struct ntdb_context *ntdb,
 	add_capabilities(ntdb, *summary);
 
 unlock:
-	free(hashesg);
-	free(freeg);
-	free(keysg);
-	free(datag);
-	free(extrag);
-	free(uncoalg);
-	free(hashes);
-	free(freet);
-	free(keys);
-	free(data);
-	free(extra);
-	free(uncoal);
-	free(ftables);
-	free(chains);
+	ntdb->free_fn(hashesg, ntdb->alloc_data);
+	ntdb->free_fn(freeg, ntdb->alloc_data);
+	ntdb->free_fn(keysg, ntdb->alloc_data);
+	ntdb->free_fn(datag, ntdb->alloc_data);
+	ntdb->free_fn(extrag, ntdb->alloc_data);
+	ntdb->free_fn(uncoalg, ntdb->alloc_data);
+	ntdb->free_fn(hashes, ntdb->alloc_data);
+	ntdb->free_fn(freet, ntdb->alloc_data);
+	ntdb->free_fn(keys, ntdb->alloc_data);
+	ntdb->free_fn(data, ntdb->alloc_data);
+	ntdb->free_fn(extra, ntdb->alloc_data);
+	ntdb->free_fn(uncoal, ntdb->alloc_data);
+	ntdb->free_fn(ftables, ntdb->alloc_data);
+	ntdb->free_fn(chains, ntdb->alloc_data);
 
 	ntdb_allrecord_unlock(ntdb, F_RDLCK);
 	ntdb_unlock_expand(ntdb, F_RDLCK);
diff --git a/lib/ntdb/test/api-20-alloc-attr.c b/lib/ntdb/test/api-20-alloc-attr.c
new file mode 100644
index 0000000000..d5c7e718bc
--- /dev/null
+++ b/lib/ntdb/test/api-20-alloc-attr.c
@@ -0,0 +1,108 @@
+#include "config.h"
+#include "ntdb.h"
+#include "tap-interface.h"
+#include <ccan/hash/hash.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <assert.h>
+
+#include "logging.h"
+
+static const struct ntdb_context *curr_ntdb;
+static const struct ntdb_file *curr_file;
+
+static int owner_null_count,
+	owner_weird_count, alloc_count, free_count, expand_count;
+
+static void *test_alloc(const void *owner, size_t len, void *priv_data)
+{
+	void *ret;
+
+	if (!owner) {
+		owner_null_count++;
+	} else if (owner != curr_ntdb && owner != curr_file) {
+		owner_weird_count++;
+	}
+
+	alloc_count++;
+	ret = malloc(len);
+
+	/* The first time, this is the current ntdb, next is
+	 * for the file struct. */
+	if (!owner) {
+		if (!curr_ntdb) {
+			curr_ntdb = ret;
+		} else if (!curr_file) {
+			curr_file = ret;
+		}
+	}
+	assert(priv_data == &owner_weird_count);
+	return ret;
+}
+
+static void *test_expand(void *old, size_t newlen, void *priv_data)
+{
+	expand_count++;
+
+	assert(priv_data == &owner_weird_count);
+	return realloc(old, newlen);
+}
+
+static void test_free(void *old, void *priv_data)
+{
+	assert(priv_data == &owner_weird_count);
+	if (old) {
+		free_count++;
+	}
+	free(old);
+}
+
+int main(int argc, char *argv[])
+{
+	unsigned int i, j;
+	union ntdb_attribute alloc_attr;
+	struct ntdb_context *ntdb;
+	int flags[] = { NTDB_INTERNAL, NTDB_DEFAULT, NTDB_NOMMAP,
+			NTDB_INTERNAL|NTDB_CONVERT, NTDB_CONVERT,
+			NTDB_NOMMAP|NTDB_CONVERT };
+	NTDB_DATA key = { (unsigned char *)&j, sizeof(j) };
+	NTDB_DATA data = { (unsigned char *)&j, sizeof(j) };
+
+	alloc_attr.base.next = &tap_log_attr;
+	alloc_attr.base.attr = NTDB_ATTRIBUTE_ALLOCATOR;
+
+	alloc_attr.alloc.alloc = test_alloc;
+	alloc_attr.alloc.expand = test_expand;
+	alloc_attr.alloc.free = test_free;
+	alloc_attr.alloc.priv_data = &owner_weird_count;
+
+	plan_tests(sizeof(flags) / sizeof(flags[0]) * (1 + 500 * 3 + 4) + 1);
+
+	for (i = 0; i < sizeof(flags) / sizeof(flags[0]); i++) {
+		curr_ntdb = NULL;
+		curr_file = NULL;
+		ntdb = ntdb_open("run-12-store.ntdb", flags[i],
+			       O_RDWR|O_CREAT|O_TRUNC, 0600, &alloc_attr);
+		ok1(ntdb);
+		if (!ntdb)
+			continue;
+
+		for (j = 0; j < 500; j++) {
+			NTDB_DATA d = { NULL, 0 }; /* Bogus GCC warning */
+			ok1(ntdb_store(ntdb, key, data, NTDB_REPLACE) == 0);
+			ok1(ntdb_fetch(ntdb, key, &d) == NTDB_SUCCESS);
+			ok1(ntdb_deq(d, data));
+			test_free(d.dptr, &owner_weird_count);
+		}
+		ntdb_close(ntdb);
+
+		ok1(owner_null_count == 2+i*2);
+		ok1(owner_weird_count == 0);
+		ok1(alloc_count == free_count);
+		ok1(expand_count != 0);
+	}
+
+	ok1(tap_log_messages == 0);
+	return exit_status();
+}
diff --git a/lib/ntdb/test/failtest_helper.c b/lib/ntdb/test/failtest_helper.c
index cc110919c3..45b24512e9 100644
--- a/lib/ntdb/test/failtest_helper.c
+++ b/lib/ntdb/test/failtest_helper.c
@@ -39,6 +39,7 @@ static bool is_unlock(const struct failtest_call *call)
 bool exit_check_log(struct tlist_calls *history)
 {
 	const struct failtest_call *i;
+	unsigned int malloc_count = 0;
 
 	tlist_for_each(history, i, list) {
 		if (!i->fail)
@@ -52,8 +53,11 @@ bool exit_check_log(struct tlist_calls *history)
 			continue;
 
 		/* Initial allocation of ntdb doesn't log. */
-		if (failmatch(i, INITIAL_NTDB_MALLOC))
-			continue;
+		if (i->type == FAILTEST_MALLOC) {
+			if (malloc_count++ == 0) {
+				continue;
+			}
+		}
 
 		/* We don't block "failures" on non-blocking locks. */
 		if (is_nonblocking_lock(i))
@@ -77,8 +81,7 @@ block_repeat_failures(struct tlist_calls *history)
 	if (failtest_suppress)
 		return FAIL_DONT_FAIL;
 
-	if (failmatch(last, INITIAL_NTDB_MALLOC)
-	    || failmatch(last, URANDOM_OPEN)
+	if (failmatch(last, URANDOM_OPEN)
 	    || failmatch(last, URANDOM_READ)) {
 		return FAIL_PROBE;
 	}
diff --git a/lib/ntdb/test/failtest_helper.h b/lib/ntdb/test/failtest_helper.h
index 5f9166d8d5..8d1c374515 100644
--- a/lib/ntdb/test/failtest_helper.h
+++ b/lib/ntdb/test/failtest_helper.h
@@ -4,7 +4,6 @@
 #include <stdbool.h>
 
 /* FIXME: Check these! */
-#define INITIAL_NTDB_MALLOC	"open.c", 425, FAILTEST_MALLOC
 #define URANDOM_OPEN		"open.c", 62, FAILTEST_OPEN
 #define URANDOM_READ		"open.c", 42, FAILTEST_READ
 
diff --git a/lib/ntdb/test/run-57-die-during-transaction.c b/lib/ntdb/test/run-57-die-during-transaction.c
index c508854f01..b6ce57590b 100644
--- a/lib/ntdb/test/run-57-die-during-transaction.c
+++ b/lib/ntdb/test/run-57-die-during-transaction.c
@@ -78,7 +78,7 @@ static void free_all(void)
 }
 
 #define malloc malloc_noleak
-#define free free_noleak
+#define free(x) free_noleak(x)
 #define realloc realloc_noleak
 
 #include "ntdb-source.h"
diff --git a/lib/ntdb/transaction.c b/lib/ntdb/transaction.c
index b03bbbdb8a..11dd7b4af5 100644
--- a/lib/ntdb/transaction.c
+++ b/lib/ntdb/transaction.c
@@ -26,7 +26,7 @@
 
 #include "private.h"
 #include <assert.h>
-#define SAFE_FREE(x) do { if ((x) != NULL) {free((void *)x); (x)=NULL;} } while(0)
+#define SAFE_FREE(ntdb, x) do { if ((x) != NULL) {ntdb->free_fn((void *)x, ntdb->alloc_data); (x)=NULL;} } while(0)
 
 /*
   transaction design:
@@ -213,12 +213,12 @@ static enum NTDB_ERROR transaction_write(struct ntdb_context *ntdb, ntdb_off_t o
 		uint8_t **new_blocks;
 		/* expand the blocks array */
 		if (ntdb->transaction->blocks == NULL) {
-			new_blocks = (uint8_t **)malloc(
-				(blk+1)*sizeof(uint8_t *));
+			new_blocks = (uint8_t **)ntdb->alloc_fn(ntdb,
+				    (blk+1)*sizeof(uint8_t *), ntdb->alloc_data);
 		} else {
-			new_blocks = (uint8_t **)realloc(
+			new_blocks = (uint8_t **)ntdb->expand_fn(
 				ntdb->transaction->blocks,
-				(blk+1)*sizeof(uint8_t *));
+				(blk+1)*sizeof(uint8_t *), ntdb->alloc_data);
 		}
 		if (new_blocks == NULL) {
 			ecode = ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
@@ -234,13 +234,16 @@ static enum NTDB_ERROR transaction_write(struct ntdb_context *ntdb, ntdb_off_t o
 
 	/* allocate and fill a block? */
 	if (ntdb->transaction->blocks[blk] == NULL) {
-		ntdb->transaction->blocks[blk] = (uint8_t *)calloc(NTDB_PGSIZE, 1);
+		ntdb->transaction->blocks[blk] = (uint8_t *)
+			ntdb->alloc_fn(ntdb->transaction->blocks, NTDB_PGSIZE,
+				   ntdb->alloc_data);
 		if (ntdb->transaction->blocks[blk] == NULL) {
 			ecode = ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 					   "transaction_write:"
 					   " failed to allocate");
 			goto fail;
 		}
+		memset(ntdb->transaction->blocks[blk], 0, NTDB_PGSIZE);
 		if (ntdb->transaction->old_map_size > blk * NTDB_PGSIZE) {
 			ntdb_len_t len2 = NTDB_PGSIZE;
 			if (len2 + (blk * NTDB_PGSIZE) > ntdb->transaction->old_map_size) {
@@ -257,7 +260,7 @@ static enum NTDB_ERROR transaction_write(struct ntdb_context *ntdb, ntdb_off_t o
 						   " failed to"
 						   " read old block: %s",
 						   strerror(errno));
-				SAFE_FREE(ntdb->transaction->blocks[blk]);
+				SAFE_FREE(ntdb, ntdb->transaction->blocks[blk]);
 				goto fail;
 			}
 		}
@@ -450,10 +453,11 @@ static void _ntdb_transaction_cancel(struct ntdb_context *ntdb)
 	/* free all the transaction blocks */
 	for (i=0;i<ntdb->transaction->num_blocks;i++) {
 		if (ntdb->transaction->blocks[i] != NULL) {
-			free(ntdb->transaction->blocks[i]);
+			ntdb->free_fn(ntdb->transaction->blocks[i],
+				      ntdb->alloc_data);
 		}
 	}
-	SAFE_FREE(ntdb->transaction->blocks);
+	SAFE_FREE(ntdb, ntdb->transaction->blocks);
 
 	if (ntdb->transaction->magic_offset) {
 		const struct ntdb_methods *methods = ntdb->transaction->io_methods;
@@ -484,7 +488,7 @@ static void _ntdb_transaction_cancel(struct ntdb_context *ntdb)
 	if (ntdb_has_open_lock(ntdb))
 		ntdb_unlock_open(ntdb, F_WRLCK);
 
-	SAFE_FREE(ntdb->transaction);
+	SAFE_FREE(ntdb, ntdb->transaction);
 }
 
 /*
@@ -536,20 +540,22 @@ _PUBLIC_ enum NTDB_ERROR ntdb_transaction_start(struct ntdb_context *ntdb)
 	}
 
 	ntdb->transaction = (struct ntdb_transaction *)
-		calloc(sizeof(struct ntdb_transaction), 1);
+		ntdb->alloc_fn(ntdb, sizeof(struct ntdb_transaction),
+			       ntdb->alloc_data);
 	if (ntdb->transaction == NULL) {
 		return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 				   "ntdb_transaction_start:"
 				   " cannot allocate");
 	}
+	memset(ntdb->transaction, 0, sizeof(*ntdb->transaction));
 
 	/* get the transaction write lock. This is a blocking lock. As
 	   discussed with Volker, there are a number of ways we could
 	   make this async, which we will probably do in the future */
 	ecode = ntdb_transaction_lock(ntdb, F_WRLCK);
 	if (ecode != NTDB_SUCCESS) {
-		SAFE_FREE(ntdb->transaction->blocks);
-		SAFE_FREE(ntdb->transaction);
+		SAFE_FREE(ntdb, ntdb->transaction->blocks);
+		SAFE_FREE(ntdb, ntdb->transaction);
 		return ecode;
 	}
 
@@ -573,8 +579,8 @@ _PUBLIC_ enum NTDB_ERROR ntdb_transaction_start(struct ntdb_context *ntdb)
 
 fail_allrecord_lock:
 	ntdb_transaction_unlock(ntdb, F_WRLCK);
-	SAFE_FREE(ntdb->transaction->blocks);
-	SAFE_FREE(ntdb->transaction);
+	SAFE_FREE(ntdb, ntdb->transaction->blocks);
+	SAFE_FREE(ntdb, ntdb->transaction);
 	return ecode;
 }
 
@@ -690,7 +696,8 @@ static struct ntdb_recovery_record *alloc_recovery(struct ntdb_context *ntdb,
 	unsigned char *p;
 	const struct ntdb_methods *old_methods = ntdb->io;
 
-	rec = malloc(sizeof(*rec) + ntdb_recovery_size(ntdb));
+	rec = ntdb->alloc_fn(ntdb, sizeof(*rec) + ntdb_recovery_size(ntdb),
+			 ntdb->alloc_data);
 	if (!rec) {
 		ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 			   "transaction_setup_recovery:"
@@ -764,7 +771,7 @@ static struct ntdb_recovery_record *alloc_recovery(struct ntdb_context *ntdb,
 	return rec;
 
 fail:
-	free(rec);
+	ntdb->free_fn(rec, ntdb->alloc_data);
 	ntdb->io = old_methods;
 	return NTDB_ERR_PTR(ecode);
 }
@@ -847,7 +854,7 @@ static enum NTDB_ERROR transaction_setup_recovery(struct ntdb_context *ntdb)
 
 	ecode = ntdb_recovery_area(ntdb, methods, &recovery_off, recovery);
 	if (ecode) {
-		free(recovery);
+		ntdb->free_fn(recovery, ntdb->alloc_data);
 		return ecode;
 	}
 
@@ -859,7 +866,7 @@ static enum NTDB_ERROR transaction_setup_recovery(struct ntdb_context *ntdb)
 						sizeof(*recovery)
 						+ recovery->max_len,
 						NTDB_LOCK_WAIT, true);
-			free(recovery);
+			ntdb->free_fn(recovery, ntdb->alloc_data);
 			if (ecode != NTDB_SUCCESS) {
 				return ntdb_logerr(ntdb, ecode, NTDB_LOG_ERROR,
 						  "ntdb_recovery_allocate:"
@@ -876,7 +883,7 @@ static enum NTDB_ERROR transaction_setup_recovery(struct ntdb_context *ntdb)
 		recovery_off = create_recovery_area(ntdb, recovery_size,
 						    recovery);
 		if (NTDB_OFF_IS_ERR(recovery_off)) {
-			free(recovery);
+			ntdb->free_fn(recovery, ntdb->alloc_data);
 			return NTDB_OFF_TO_ERR(recovery_off);
 		}
 	}
@@ -891,14 +898,14 @@ static enum NTDB_ERROR transaction_setup_recovery(struct ntdb_context *ntdb)
 	ecode = methods->twrite(ntdb, recovery_off, recovery,
 				sizeof(*recovery) + recovery_size);
 	if (ecode != NTDB_SUCCESS) {
-		free(recovery);
+		ntdb->free_fn(recovery, ntdb->alloc_data);
 		return ntdb_logerr(ntdb, ecode, NTDB_LOG_ERROR,
 				  "ntdb_transaction_setup_recovery:"
 				  " failed to write recovery data");
 	}
 	transaction_write_existing(ntdb, recovery_off, recovery, recovery_size);
 
-	free(recovery);
+	ntdb->free_fn(recovery, ntdb->alloc_data);
 
 	/* as we don't have ordered writes, we have to sync the recovery
 	   data before we update the magic to indicate that the recovery
@@ -1075,10 +1082,10 @@ _PUBLIC_ enum NTDB_ERROR ntdb_transaction_commit(struct ntdb_context *ntdb)
 
 			return ecode;
 		}
-		SAFE_FREE(ntdb->transaction->blocks[i]);
+		SAFE_FREE(ntdb, ntdb->transaction->blocks[i]);
 	}
 
-	SAFE_FREE(ntdb->transaction->blocks);
+	SAFE_FREE(ntdb, ntdb->transaction->blocks);
 	ntdb->transaction->num_blocks = 0;
 
 	/* ensure the new data is on disk */
@@ -1158,7 +1165,7 @@ enum NTDB_ERROR ntdb_transaction_recover(struct ntdb_context *ntdb)
 
 	recovery_eof = rec.eof;
 
-	data = (unsigned char *)malloc(rec.len);
+	data = (unsigned char *)ntdb->alloc_fn(ntdb, rec.len, ntdb->alloc_data);
 	if (data == NULL) {
 		return ntdb_logerr(ntdb, NTDB_ERR_OOM, NTDB_LOG_ERROR,
 				  "ntdb_transaction_recover:"
@@ -1186,7 +1193,7 @@ enum NTDB_ERROR ntdb_transaction_recover(struct ntdb_context *ntdb)
 
 		ecode = ntdb->io->twrite(ntdb, ofs, p, len);
 		if (ecode != NTDB_SUCCESS) {
-			free(data);
+			ntdb->free_fn(data, ntdb->alloc_data);
 			return ntdb_logerr(ntdb, ecode, NTDB_LOG_ERROR,
 					  "ntdb_transaction_recover:"
 					  " failed to recover %zu bytes"
@@ -1196,7 +1203,7 @@ enum NTDB_ERROR ntdb_transaction_recover(struct ntdb_context *ntdb)
 		p += len;
 	}
 
-	free(data);
+	ntdb->free_fn(data, ntdb->alloc_data);
 
 	ecode = transaction_sync(ntdb, 0, ntdb->file->map_size);
 	if (ecode != NTDB_SUCCESS) {
diff --git a/lib/ntdb/traverse.c b/lib/ntdb/traverse.c
index 45478755bd..ee1e1006db 100644
--- a/lib/ntdb/traverse.c
+++ b/lib/ntdb/traverse.c
@@ -36,10 +36,10 @@ _PUBLIC_ int64_t ntdb_traverse_(struct ntdb_context *ntdb,
 
 		count++;
 		if (fn && fn(ntdb, k, d, p)) {
-			free(k.dptr);
+			ntdb->free_fn(k.dptr, ntdb->alloc_data);
 			return count;
 		}
-		free(k.dptr);
+		ntdb->free_fn(k.dptr, ntdb->alloc_data);
 	}
 
 	if (ecode != NTDB_ERR_NOEXIST) {
@@ -63,7 +63,7 @@ _PUBLIC_ enum NTDB_ERROR ntdb_nextkey(struct ntdb_context *ntdb, NTDB_DATA *key)
 	struct ntdb_used_record rec;
 
 	tinfo.prev = find_and_lock(ntdb, *key, F_RDLCK, &h, &rec, &tinfo);
-	free(key->dptr);
+	ntdb->free_fn(key->dptr, ntdb->alloc_data);
 	if (NTDB_OFF_IS_ERR(tinfo.prev)) {
 		return NTDB_OFF_TO_ERR(tinfo.prev);
 	}
diff --git a/lib/ntdb/wscript b/lib/ntdb/wscript
index fd9c96b59b..0a90b46c80 100644
--- a/lib/ntdb/wscript
+++ b/lib/ntdb/wscript
@@ -69,6 +69,7 @@ def configure(conf):
                                 'test/api-13-delete.c',
                                 'test/api-14-exists.c',
                                 'test/api-16-wipe_all.c',
+                                'test/api-20-alloc-attr.c',
                                 'test/api-21-parse_record.c',
                                 'test/api-55-transaction.c',
                                 'test/api-80-tdb_fd.c',