Custom memory patch work (based on Sean Chittenden's patch)

diff --git a/include/memcached.h b/include/memcached.h
index 05cb74a81d081aed67b895a919f0ff44786359c7..7f46a93178392b25f575d83152cfa00fc9e65b01 100644 (file)
--- a/include/memcached.h
+++ b/include/memcached.h
@@ -74,7 +74,9 @@ typedef struct memcached_string_st memcached_string_st;
 typedef struct memcached_server_st memcached_server_st;
 typedef memcached_return (*clone_func)(memcached_st *parent, memcached_st *clone);
 typedef memcached_return (*cleanup_func)(memcached_st *ptr);
-
+typedef void (*memcached_free_function)(memcached_st *ptr, void *mem);
+typedef void *(*memcached_malloc_function)(memcached_st *ptr, const size_t size);
+typedef void *(*memcached_realloc_function)(memcached_st *ptr, void *mem, const size_t size);
 
 typedef enum {
   MEMCACHED_DISTRIBUTION_MODULA,
@@ -100,6 +102,9 @@ typedef enum {
   MEMCACHED_CALLBACK_USER_DATA,
   MEMCACHED_CALLBACK_CLEANUP_FUNCTION,
   MEMCACHED_CALLBACK_CLONE_FUNCTION,
+  MEMCACHED_CALLBACK_MALLOC_FUNCTION,
+  MEMCACHED_CALLBACK_REALLOC_FUNCTION,
+  MEMCACHED_CALLBACK_FREE_FUNCTION,
 } memcached_callback;
 
 typedef enum {
@@ -213,6 +218,9 @@ struct memcached_st {
   unsigned int wheel[MEMCACHED_WHEEL_SIZE];
   clone_func on_clone;
   cleanup_func on_cleanup;
+  memcached_free_function call_free;
+  memcached_malloc_function call_malloc;
+  memcached_realloc_function call_realloc;
 #ifdef NOT_USED /* Future Use */
   uint8_t replicas;
   memcached_return warning;

diff --git a/lib/common.h b/lib/common.h
index 268109b31a2969d29fe8fb0a612280a871c3f1af..c87beb60810ff5ad968da226c4a6144ba0e44ea2 100644 (file)
--- a/lib/common.h
+++ b/lib/common.h
@@ -110,6 +110,7 @@ memcached_return value_fetch(memcached_st *ptr,
                              char *buffer,
                              memcached_result_st *result,
                              unsigned int server_key);
+void server_list_free(memcached_st *ptr, memcached_server_st *servers);
 
 
 #endif /* __COMMON_H__ */

diff --git a/lib/memcached.c b/lib/memcached.c
index 5471eba6e8b16c743f86075a867268cded2dc7ea..946af6cb8f6bdd9a4edbc0e2c8179a001eed1031 100644 (file)
--- a/lib/memcached.c
+++ b/lib/memcached.c
@@ -6,7 +6,8 @@
 memcached_st *memcached_create(memcached_st *ptr)
 {
   memcached_result_st *result_ptr;
-  if (!ptr)
+
+  if (ptr == NULL)
   {
     ptr= (memcached_st *)malloc(sizeof(memcached_st));
 
@@ -32,14 +33,19 @@ void memcached_free(memcached_st *ptr)
 {
   /* If we have anything open, lets close it now */
   memcached_quit(ptr);
-  memcached_server_list_free(ptr->hosts);
+  server_list_free(ptr, ptr->hosts);
   memcached_result_free(&ptr->result);
 
   if (ptr->on_cleanup)
     ptr->on_cleanup(ptr);
 
   if (ptr->is_allocated == MEMCACHED_ALLOCATED)
-    free(ptr);
+  {
+    if (ptr->call_free)
+      ptr->call_free(ptr, ptr);
+    else
+      free(ptr);
+  }
   else
     ptr->is_allocated= MEMCACHED_USED;
 }
@@ -94,6 +100,12 @@ memcached_st *memcached_clone(memcached_st *clone, memcached_st *ptr)
   new_clone->hash= ptr->hash;
   new_clone->user_data= ptr->user_data;
 
+  new_clone->on_clone= ptr->on_clone;
+  new_clone->on_cleanup= ptr->on_cleanup;
+  new_clone->call_free= ptr->call_free;
+  new_clone->call_malloc= ptr->call_malloc;
+  new_clone->call_realloc= ptr->call_realloc;
+
   if (ptr->on_clone)
     ptr->on_clone(ptr, new_clone);
 

diff --git a/lib/memcached_callback.c b/lib/memcached_callback.c
index 7149e276332252803f8ee35cb89114374494714f..5348183cf89e134d17d1f7daeb2d8991a9087c88 100644 (file)
--- a/lib/memcached_callback.c
+++ b/lib/memcached_callback.c
@@ -30,6 +30,24 @@ memcached_return memcached_callback_set(memcached_st *ptr,
       ptr->on_clone= func;
       break;
     }
+  case MEMCACHED_CALLBACK_MALLOC_FUNCTION:
+    {
+      memcached_malloc_function func= (memcached_malloc_function)data;
+      ptr->call_malloc= func;
+      break;
+    }
+  case MEMCACHED_CALLBACK_REALLOC_FUNCTION:
+    {
+      memcached_realloc_function func= (memcached_realloc_function)data;
+      ptr->call_realloc= func;
+      break;
+    }
+  case MEMCACHED_CALLBACK_FREE_FUNCTION:
+    {
+      memcached_free_function func= (memcached_free_function)data;
+      ptr->call_free= func;
+      break;
+    }
   default:
     return MEMCACHED_FAILURE;
   }
@@ -58,6 +76,21 @@ void *memcached_callback_get(memcached_st *ptr,
       *error= ptr->on_clone ? MEMCACHED_SUCCESS : MEMCACHED_FAILURE;
       return (void *)ptr->on_clone;
     }
+  case MEMCACHED_CALLBACK_MALLOC_FUNCTION:
+    {
+      *error= ptr->call_malloc ? MEMCACHED_SUCCESS : MEMCACHED_FAILURE;
+      return (void *)ptr->call_malloc;
+    }
+  case MEMCACHED_CALLBACK_REALLOC_FUNCTION:
+    {
+      *error= ptr->call_realloc ? MEMCACHED_SUCCESS : MEMCACHED_FAILURE;
+      return (void *)ptr->call_realloc;
+    }
+  case MEMCACHED_CALLBACK_FREE_FUNCTION:
+    {
+      *error= ptr->call_free ? MEMCACHED_SUCCESS : MEMCACHED_FAILURE;
+      return (void *)ptr->call_free;
+    }
   default:
       WATCHPOINT_ASSERT(0);
       *error= MEMCACHED_FAILURE;

diff --git a/lib/memcached_hosts.c b/lib/memcached_hosts.c
index 88c255a59d2997222b67f83107fa11cca545665f..811a230e9a3052dca3b49ef4d5b7e0c5a102e9a8 100644 (file)
--- a/lib/memcached_hosts.c
+++ b/lib/memcached_hosts.c
@@ -44,6 +44,23 @@ static void host_reset(memcached_server_st *host, char *hostname, unsigned int p
   host->sockaddr_inited= MEMCACHED_NOT_ALLOCATED;
 }
 
+void server_list_free(memcached_st *ptr, memcached_server_st *servers)
+{
+  unsigned int x;
+
+  if (servers == NULL)
+    return;
+
+  for (x= 0; x < servers->count; x++)
+    if (servers[x].address_info)
+      freeaddrinfo(servers[x].address_info);
+
+  if (ptr && ptr->call_free)
+    ptr->call_free(ptr, servers);
+  else
+    free(servers);
+}
+
 memcached_return memcached_server_push(memcached_st *ptr, memcached_server_st *list)
 {
   unsigned int x;
@@ -55,9 +72,14 @@ memcached_return memcached_server_push(memcached_st *ptr, memcached_server_st *l
 
   count= list[0].count;
 
-  new_host_list= 
-    (memcached_server_st *)realloc(ptr->hosts, 
-                                   sizeof(memcached_server_st) * (count + ptr->number_of_hosts));
+  if (ptr->call_realloc)
+    new_host_list= 
+      (memcached_server_st *)ptr->call_realloc(ptr, ptr->hosts, 
+                                               sizeof(memcached_server_st) * (count + ptr->number_of_hosts));
+  else
+    new_host_list= 
+      (memcached_server_st *)realloc(ptr->hosts, 
+                                     sizeof(memcached_server_st) * (count + ptr->number_of_hosts));
 
   if (!new_host_list)
     return MEMCACHED_MEMORY_ALLOCATION_FAILURE;
@@ -120,9 +142,13 @@ static memcached_return server_add(memcached_st *ptr, char *hostname,
   LIBMEMCACHED_MEMCACHED_SERVER_ADD_START();
 
 
-  new_host_list= (memcached_server_st *)realloc(ptr->hosts, 
-                                                sizeof(memcached_server_st) * (ptr->number_of_hosts+1));
-  if (!new_host_list)
+  if (ptr->call_realloc)
+    new_host_list= (memcached_server_st *)ptr->call_realloc(ptr, ptr->hosts, 
+                                                            sizeof(memcached_server_st) * (ptr->number_of_hosts+1));
+  else
+    new_host_list= (memcached_server_st *)realloc(ptr->hosts, 
+                                                  sizeof(memcached_server_st) * (ptr->number_of_hosts+1));
+  if (new_host_list == NULL)
     return MEMCACHED_MEMORY_ALLOCATION_FAILURE;
 
   ptr->hosts= new_host_list;
@@ -183,14 +209,5 @@ unsigned int memcached_server_list_count(memcached_server_st *ptr)
 
 void memcached_server_list_free(memcached_server_st *ptr)
 {
-  unsigned int x;
-
-  if (ptr == NULL)
-    return;
-
-  for (x= 0; x < ptr->count; x++)
-    if (ptr[x].address_info)
-      freeaddrinfo(ptr[x].address_info);
-
-  free(ptr);
+  server_list_free(NULL, ptr);
 }

diff --git a/lib/memcached_result.c b/lib/memcached_result.c
index 89de40301552ff891ca908b481068edd822f6692..45bd3975e68b2503a919282087acf40fd26aa97d 100644 (file)
--- a/lib/memcached_result.c
+++ b/lib/memcached_result.c
@@ -17,8 +17,12 @@ memcached_result_st *memcached_result_create(memcached_st *memc,
   }
   else
   {
-    ptr= (memcached_result_st *)malloc(sizeof(memcached_result_st));
-    if (!ptr)
+    if (memc->call_malloc)
+      ptr= (memcached_result_st *)memc->call_malloc(ptr->root, sizeof(memcached_result_st));
+    else
+      ptr= (memcached_result_st *)malloc(sizeof(memcached_result_st));
+
+    if (ptr == NULL)
       return NULL;
     memset(ptr, 0, sizeof(memcached_result_st));
     ptr->is_allocated= MEMCACHED_ALLOCATED;

diff --git a/lib/memcached_stats.c b/lib/memcached_stats.c
index 7c61890f086027fffb785ee3b24a4d35f8ef4a25..a818fe1fe982124ffcad5d5e09f3b70490256997 100644 (file)
--- a/lib/memcached_stats.c
+++ b/lib/memcached_stats.c
@@ -146,59 +146,69 @@ char *memcached_stat_get_value(memcached_st *ptr, memcached_stat_st *stat,
                                char *key, memcached_return *error)
 {
   char buffer[SMALL_STRING_LEN];
+  size_t length;
+  char *ret;
+
   *error= MEMCACHED_SUCCESS;
 
   if (!memcmp("pid", key, strlen("pid")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->pid);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->pid);
   else if (!memcmp("uptime", key, strlen("uptime")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->uptime);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->uptime);
   else if (!memcmp("time", key, strlen("time")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->time);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->time);
   else if (!memcmp("version", key, strlen("version")))
-    snprintf(buffer, SMALL_STRING_LEN,"%s", stat->version);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%s", stat->version);
   else if (!memcmp("pointer_size", key, strlen("pointer_size")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->pointer_size);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->pointer_size);
   else if (!memcmp("rusage_user", key, strlen("rusage_user")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u.%u", stat->rusage_user_seconds, stat->rusage_user_microseconds);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u.%u", stat->rusage_user_seconds, stat->rusage_user_microseconds);
   else if (!memcmp("rusage_system", key, strlen("rusage_system")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u.%u", stat->rusage_system_seconds, stat->rusage_system_microseconds);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u.%u", stat->rusage_system_seconds, stat->rusage_system_microseconds);
   else if (!memcmp("curr_items", key, strlen("curr_items")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->curr_items);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->curr_items);
   else if (!memcmp("total_items", key, strlen("total_items")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->total_items);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->total_items);
   else if (!memcmp("bytes", key, strlen("bytes")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->bytes);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->bytes);
   else if (!memcmp("curr_connections", key, strlen("curr_connections")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->curr_connections);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->curr_connections);
   else if (!memcmp("total_connections", key, strlen("total_connections")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->total_connections);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->total_connections);
   else if (!memcmp("connection_structures", key, strlen("connection_structures")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->connection_structures);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->connection_structures);
   else if (!memcmp("cmd_get", key, strlen("cmd_get")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->cmd_get);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->cmd_get);
   else if (!memcmp("cmd_set", key, strlen("cmd_set")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->cmd_set);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->cmd_set);
   else if (!memcmp("get_hits", key, strlen("get_hits")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->get_hits);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->get_hits);
   else if (!memcmp("get_misses", key, strlen("get_misses")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->get_misses);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->get_misses);
   else if (!memcmp("evictions", key, strlen("evictions")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->evictions);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->evictions);
   else if (!memcmp("bytes_read", key, strlen("bytes_read")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->bytes_read);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->bytes_read);
   else if (!memcmp("bytes_written", key, strlen("bytes_written")))
-    snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->bytes_written);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%llu", (unsigned long long)stat->bytes_written);
   else if (!memcmp("limit_maxbytes", key, strlen("limit_maxbytes")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->limit_maxbytes);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->limit_maxbytes);
   else if (!memcmp("threads", key, strlen("threads")))
-    snprintf(buffer, SMALL_STRING_LEN,"%u", stat->threads);
+    length= snprintf(buffer, SMALL_STRING_LEN,"%u", stat->threads);
   else
   {
     *error= MEMCACHED_NOTFOUND;
     return NULL;
   }
 
-  return strdup(buffer);
+  if (ptr->call_malloc)
+    ret= ptr->call_malloc(ptr, length + 1);
+  else
+    ret= malloc(length + 1);
+  memcpy(ret, buffer, length);
+  ret[length]= '\0';
+
+  return ret;
 }
 
 static memcached_return memcached_stats_fetch(memcached_st *ptr,
@@ -263,11 +273,19 @@ memcached_stat_st *memcached_stat(memcached_st *ptr, char *args, memcached_retur
   memcached_return rc;
   memcached_stat_st *stats;
 
-  stats= (memcached_stat_st *)malloc(sizeof(memcached_stat_st)*(ptr->number_of_hosts));
+  if (ptr->call_malloc)
+    stats= (memcached_stat_st *)ptr->call_malloc(ptr, sizeof(memcached_stat_st)*(ptr->number_of_hosts));
+  else
+    stats= (memcached_stat_st *)malloc(sizeof(memcached_stat_st)*(ptr->number_of_hosts));
+
   if (!stats)
   {
     *error= MEMCACHED_MEMORY_ALLOCATION_FAILURE;
-    free(stats);
+    if (ptr->call_free)
+      ptr->call_free(ptr, stats);
+    else
+      free(stats);
+
     return NULL;
   }
   memset(stats, 0, sizeof(memcached_stat_st)*(ptr->number_of_hosts));
@@ -312,7 +330,11 @@ char ** memcached_stat_get_keys(memcached_st *ptr, memcached_stat_st *stat,
 {
   char **list;
   size_t length= sizeof(memcached_stat_keys);
-  list= (char **)malloc(length);
+
+  if (ptr->call_malloc)
+    list= (char **)ptr->call_malloc(ptr, length);
+  else
+    list= (char **)malloc(length);
 
   if (!list)
   {
@@ -330,5 +352,14 @@ char ** memcached_stat_get_keys(memcached_st *ptr, memcached_stat_st *stat,
 
 void memcached_stat_free(memcached_st *ptr, memcached_stat_st *stat)
 {
-  free(stat);
+  if (stat == NULL)
+  {
+    WATCHPOINT_ASSERT(0); /* Be polite, but when debugging catch this as an error */
+    return;
+  }
+
+  if (ptr && ptr->call_free)
+    ptr->call_free(ptr, stat);
+  else
+    free(stat);
 }

diff --git a/lib/memcached_string.c b/lib/memcached_string.c
index 247fbe475dd0506794add134d0c15bb47f3096db..ad396118b2b878ba9ab938e9c412defadafdf648 100644 (file)
--- a/lib/memcached_string.c
+++ b/lib/memcached_string.c
@@ -18,7 +18,10 @@ memcached_return memcached_string_check(memcached_string_st *string, size_t need
     if (new_size < need)
       return MEMCACHED_MEMORY_ALLOCATION_FAILURE;
 
-    new_value= (char *)realloc(string->string, new_size);
+    if (string->root->call_realloc)
+      new_value= (char *)string->root->call_realloc(string->root, string->string, new_size);
+    else
+      new_value= (char *)realloc(string->string, new_size);
 
     if (new_value == NULL)
       return MEMCACHED_MEMORY_ALLOCATION_FAILURE;
@@ -44,8 +47,12 @@ memcached_string_st *memcached_string_create(memcached_st *ptr, memcached_string
   }
   else
   {
-    string= (memcached_string_st *)malloc(sizeof(memcached_string_st));
-    if (!string)
+    if (ptr->call_malloc)
+      string= (memcached_string_st *)ptr->call_malloc(ptr, sizeof(memcached_string_st));
+    else
+      string= (memcached_string_st *)malloc(sizeof(memcached_string_st));
+
+    if (string == NULL)
       return NULL;
     memset(string, 0, sizeof(memcached_string_st));
     string->is_allocated= MEMCACHED_ALLOCATED;
@@ -56,7 +63,11 @@ memcached_string_st *memcached_string_create(memcached_st *ptr, memcached_string
   rc=  memcached_string_check(string, initial_size);
   if (rc != MEMCACHED_SUCCESS)
   {
-    free(string);
+    if (ptr->call_free)
+      ptr->call_free(ptr, string);
+    else
+      free(string);
+
     return NULL;
   }
 
@@ -129,8 +140,12 @@ char *memcached_string_c_copy(memcached_string_st *string)
 
   WATCHPOINT_ASSERT(string->is_allocated != MEMCACHED_USED);
 
-  c_ptr= (char *)malloc((memcached_string_length(string)+1) * sizeof(char));
-  if (!c_ptr)
+  if (string->root->call_malloc)
+    c_ptr= (char *)string->root->call_malloc(string->root, (memcached_string_length(string)+1) * sizeof(char));
+  else
+    c_ptr= (char *)malloc((memcached_string_length(string)+1) * sizeof(char));
+
+  if (c_ptr == NULL)
     return NULL;
 
   memcpy(c_ptr, memcached_string_value(string), memcached_string_length(string));
@@ -153,10 +168,20 @@ void memcached_string_free(memcached_string_st *ptr)
     return;
 
   if (ptr->string)
-    free(ptr->string);
+  {
+    if (ptr->root->call_free)
+      ptr->root->call_free(ptr->root, ptr->string);
+    else
+      free(ptr->string);
+  }
 
   if (ptr->is_allocated == MEMCACHED_ALLOCATED)
-    free(ptr);
+  {
+    if (ptr->root->call_free)
+      ptr->root->call_free(ptr->root, ptr);
+    else
+      free(ptr);
+  }
   else
     ptr->is_allocated= MEMCACHED_USED;
 }

diff --git a/tests/function.c b/tests/function.c
index 6ca61d7d5addb443a35899af8f99e8a49f9e2f0b..3bff805cb0b0e1ab77b49fe4997dfa4e09798af1 100644 (file)
--- a/tests/function.c
+++ b/tests/function.c
@@ -1043,15 +1043,6 @@ uint8_t behavior_test(memcached_st *memc)
   value= memcached_behavior_get(memc, MEMCACHED_BEHAVIOR_SOCKET_RECV_SIZE);
   assert(value > 0);
 
-  {
-    int x= 5;
-    int *test_ptr;
-
-    memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USER_DATA, &x);
-    test_ptr= (int *)memcached_behavior_get(memc, MEMCACHED_BEHAVIOR_USER_DATA);
-    assert(*test_ptr == x);
-  }
-
   return 0;
 }
 
@@ -1998,6 +1989,55 @@ memcached_return pre_hash_ketama(memcached_st *memc)
 
   return MEMCACHED_SUCCESS;
 }
+void my_free(memcached_st *ptr, void *mem)
+{
+  free(mem);
+}
+
+void *my_malloc(memcached_st *ptr, const size_t size)
+{
+  return malloc(size);
+}
+
+void *my_realloc(memcached_st *ptr, void *mem, const size_t size)
+{
+  return realloc(mem, size);
+}
+
+memcached_return set_memory_alloc(memcached_st *memc)
+{
+  {
+    memcached_malloc_function test_ptr;
+    memcached_return rc;
+
+    rc= memcached_callback_set(memc, MEMCACHED_CALLBACK_MALLOC_FUNCTION, &my_malloc);
+    assert(rc == MEMCACHED_SUCCESS);
+    test_ptr= (memcached_malloc_function)memcached_callback_get(memc, MEMCACHED_CALLBACK_USER_DATA, &rc);
+    assert(test_ptr == (memcached_malloc_function)my_malloc);
+  }
+
+  {
+    memcached_realloc_function test_ptr;
+    memcached_return rc;
+
+    rc= memcached_callback_set(memc, MEMCACHED_CALLBACK_REALLOC_FUNCTION, &my_realloc);
+    assert(rc == MEMCACHED_SUCCESS);
+    test_ptr= (memcached_realloc_function)memcached_callback_get(memc, MEMCACHED_CALLBACK_USER_DATA, &rc);
+    assert(test_ptr == my_realloc);
+  }
+
+  {
+    memcached_free_function test_ptr;
+    memcached_return rc;
+
+    rc= memcached_callback_set(memc, MEMCACHED_CALLBACK_FREE_FUNCTION, my_free);
+    assert(rc == MEMCACHED_SUCCESS);
+    test_ptr= (memcached_free_function)memcached_callback_get(memc, MEMCACHED_CALLBACK_USER_DATA, &rc);
+    assert(test_ptr == my_free);
+  }
+
+  return MEMCACHED_SUCCESS;
+}
 
 memcached_return enable_consistent(memcached_st *memc)
 {
@@ -2214,6 +2254,7 @@ collection_st collection[] ={
   {"poll_timeout", poll_timeout, 0, tests},
   {"gets", enable_cas, 0, tests},
   {"consistent", enable_consistent, 0, tests},
+  {"memory_allocators", set_memory_alloc, 0, tests},
 //  {"udp", pre_udp, 0, tests},
   {"version_1_2_3", check_for_1_2_3, 0, version_1_2_3},
   {"string", 0, 0, string_tests},