Added ability to override mem alloc funcs

1 files changed, 23 insertions, 3 deletions

diff --git a/src/core/support/alloc.c b/src/core/support/alloc.c
index bfcb77956b..2550b2476d 100644
--- a/src/core/support/alloc.c
+++ b/src/core/support/alloc.c
@@ -34,13 +34,33 @@
 #include <grpc/support/alloc.h>
 
 #include <stdlib.h>
+#include <grpc/support/log.h>
 #include <grpc/support/port_platform.h>
 #include "src/core/profiling/timers.h"
 
+static void *(*malloc_fn)(size_t size) = malloc;
+static void *(*realloc_fn)(void *ptr, size_t size) = realloc;
+static void (*free_fn)(void *ptr) = free;
+
+void gpr_set_malloc(void *(*custom_malloc)(size_t size)) {
+  GPR_ASSERT(custom_malloc != NULL);
+  malloc_fn = custom_malloc;
+}
+
+void gpr_set_realloc(void *(*custom_realloc)(void *ptr, size_t size)) {
+  GPR_ASSERT(custom_realloc != NULL);
+  realloc_fn = custom_realloc;
+}
+
+void gpr_set_free(void (*custom_free)(void *ptr)) {
+  GPR_ASSERT(custom_free != NULL);
+  free_fn = custom_free;
+}
+
 void *gpr_malloc(size_t size) {
   void *p;
   GPR_TIMER_BEGIN("gpr_malloc", 0);
-  p = malloc(size);
+  p = malloc_fn(size);
   if (!p) {
     abort();
   }
@@ -50,13 +70,13 @@ void *gpr_malloc(size_t size) {
 
 void gpr_free(void *p) {
   GPR_TIMER_BEGIN("gpr_free", 0);
-  free(p);
+  free_fn(p);
   GPR_TIMER_END("gpr_free", 0);
 }
 
 void *gpr_realloc(void *p, size_t size) {
   GPR_TIMER_BEGIN("gpr_realloc", 0);
-  p = realloc(p, size);
+  p = realloc_fn(p, size);
   if (!p) {
     abort();
   }