3 files changed, 18 insertions, 3 deletions
diff --git a/lib/tdb2/test/external-agent.c b/lib/tdb2/test/external-agent.c
index 23874777da..01c7106fba 100644
--- a/lib/tdb2/test/external-agent.c
+++ b/lib/tdb2/test/external-agent.c
@@ -17,6 +17,8 @@
 
 static struct tdb_context *tdb;
 
+void (*external_agent_free)(void *) = free;
+
 static enum TDB_ERROR clear_if_first(int fd, void *arg)
 {
 /* We hold a lock offset 4 always, so we can tell if anyone is holding it.
@@ -100,10 +102,10 @@ static enum agent_return do_operation(enum operation op, const char *name)
 			ret = OTHER_FAILURE;
 		} else if (!tdb_deq(data, k)) {
 			ret = OTHER_FAILURE;
-			free(data.dptr);
+			external_agent_free(data.dptr);
 		} else {
 			ret = SUCCESS;
-			free(data.dptr);
+			external_agent_free(data.dptr);
 		}
 		break;
 	case STORE:
diff --git a/lib/tdb2/test/external-agent.h b/lib/tdb2/test/external-agent.h
index 9eada10750..cf9789b669 100644
--- a/lib/tdb2/test/external-agent.h
+++ b/lib/tdb2/test/external-agent.h
@@ -35,6 +35,9 @@ enum agent_return external_agent_operation(struct agent *handle,
 					   enum operation op,
 					   const char *name);
 
+/* Hook into free() on tdb_data in external agent. */
+void (*external_agent_free)(void *);
+
 /* Mapping enum -> string. */
 const char *agent_return_name(enum agent_return ret);
 const char *operation_name(enum operation op);
diff --git a/lib/tdb2/test/run-57-die-during-transaction.c b/lib/tdb2/test/run-57-die-during-transaction.c
index 6e9cb1e3e7..42102ae732 100644
--- a/lib/tdb2/test/run-57-die-during-transaction.c
+++ b/lib/tdb2/test/run-57-die-during-transaction.c
@@ -15,8 +15,9 @@ static int ftruncate_check(int fd, off_t length);
 
 /* There's a malloc inside transaction_setup_recovery, and valgrind complains
  * when we longjmp and leak it. */
-#define MAX_ALLOCATIONS 200
+#define MAX_ALLOCATIONS 10
 static void *allocated[MAX_ALLOCATIONS];
+static unsigned max_alloc = 0;
 
 static void *malloc_noleak(size_t len)
 {
@@ -25,6 +26,10 @@ static void *malloc_noleak(size_t len)
 	for (i = 0; i < MAX_ALLOCATIONS; i++)
 		if (!allocated[i]) {
 			allocated[i] = malloc(len);
+			if (i > max_alloc) {
+				max_alloc = i;
+				diag("max_alloc: %i", max_alloc);
+			}
 			return allocated[i];
 		}
 	diag("Too many allocations!");
@@ -37,6 +42,10 @@ static void *realloc_noleak(void *p, size_t size)
 
 	for (i = 0; i < MAX_ALLOCATIONS; i++) {
 		if (allocated[i] == p) {
+			if (i > max_alloc) {
+				max_alloc = i;
+				diag("max_alloc: %i", max_alloc);
+			}
 			return allocated[i] = realloc(p, size);
 		}
 	}
@@ -270,6 +279,7 @@ int main(int argc, char *argv[])
 	plan_tests(24);
 	unlock_callback = maybe_die;
 
+	external_agent_free = free_noleak;
 	agent = prepare_external_agent();
 	if (!agent)
 		err(1, "preparing agent");