2 files changed, 33 insertions, 27 deletions

diff --git a/test/core/util/memory_counters.c b/test/core/util/memory_counters.c
index bebe94e582..fd6770bd49 100644
--- a/test/core/util/memory_counters.c
+++ b/test/core/util/memory_counters.c
@@ -39,7 +39,6 @@
 
 #include "test/core/util/memory_counters.h"
 
-static gpr_mu g_memory_mutex;
 static struct grpc_memory_counters g_memory_counters;
 static gpr_allocation_functions g_old_allocs;
 
@@ -50,12 +49,14 @@ static void guard_free(void *vptr);
 static void *guard_malloc(size_t size) {
   size_t *ptr;
   if (!size) return NULL;
-  gpr_mu_lock(&g_memory_mutex);
-  g_memory_counters.total_size_absolute += size;
-  g_memory_counters.total_size_relative += size;
-  g_memory_counters.total_allocs_absolute++;
-  g_memory_counters.total_allocs_relative++;
-  gpr_mu_unlock(&g_memory_mutex);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_size_absolute,
+                               (gpr_atm)size);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_size_relative,
+                               (gpr_atm)size);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_allocs_absolute,
+                               (gpr_atm)1);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_allocs_relative,
+                               (gpr_atm)1);
   ptr = g_old_allocs.malloc_fn(size + sizeof(size));
   *ptr++ = size;
   return ptr;
@@ -71,12 +72,14 @@ static void *guard_realloc(void *vptr, size_t size) {
     return NULL;
   }
   --ptr;
-  gpr_mu_lock(&g_memory_mutex);
-  g_memory_counters.total_size_absolute += size;
-  g_memory_counters.total_size_relative -= *ptr;
-  g_memory_counters.total_size_relative += size;
-  g_memory_counters.total_allocs_absolute++;
-  gpr_mu_unlock(&g_memory_mutex);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_size_absolute,
+                               (gpr_atm)size);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_size_relative,
+                               -(gpr_atm)*ptr);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_size_relative,
+                               (gpr_atm)size);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_allocs_absolute,
+                               (gpr_atm)1);
   ptr = g_old_allocs.realloc_fn(ptr, size + sizeof(size));
   *ptr++ = size;
   return ptr;
@@ -86,10 +89,10 @@ static void guard_free(void *vptr) {
   size_t *ptr = vptr;
   if (!vptr) return;
   --ptr;
-  gpr_mu_lock(&g_memory_mutex);
-  g_memory_counters.total_size_relative -= *ptr;
-  g_memory_counters.total_allocs_relative--;
-  gpr_mu_unlock(&g_memory_mutex);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_size_relative,
+                               -(gpr_atm)*ptr);
+  gpr_atm_no_barrier_fetch_add(&g_memory_counters.total_allocs_relative,
+                               -(gpr_atm)1);
   g_old_allocs.free_fn(ptr);
 }
 
@@ -98,20 +101,23 @@ struct gpr_allocation_functions g_guard_allocs = {guard_malloc, guard_realloc,
 
 void grpc_memory_counters_init() {
   memset(&g_memory_counters, 0, sizeof(g_memory_counters));
-  gpr_mu_init(&g_memory_mutex);
   g_old_allocs = gpr_get_allocation_functions();
   gpr_set_allocation_functions(g_guard_allocs);
 }
 
 void grpc_memory_counters_destroy() {
   gpr_set_allocation_functions(g_old_allocs);
-  gpr_mu_destroy(&g_memory_mutex);
 }
 
 struct grpc_memory_counters grpc_memory_counters_snapshot() {
   struct grpc_memory_counters counters;
-  gpr_mu_lock(&g_memory_mutex);
-  counters = g_memory_counters;
-  gpr_mu_unlock(&g_memory_mutex);
+  counters.total_size_relative =
+      gpr_atm_no_barrier_load(&g_memory_counters.total_size_relative);
+  counters.total_size_absolute =
+      gpr_atm_no_barrier_load(&g_memory_counters.total_size_absolute);
+  counters.total_allocs_relative =
+      gpr_atm_no_barrier_load(&g_memory_counters.total_allocs_relative);
+  counters.total_allocs_absolute =
+      gpr_atm_no_barrier_load(&g_memory_counters.total_allocs_absolute);
   return counters;
 }
diff --git a/test/core/util/memory_counters.h b/test/core/util/memory_counters.h
index b9b2b3adda..51487c73ce 100644
--- a/test/core/util/memory_counters.h
+++ b/test/core/util/memory_counters.h
@@ -34,13 +34,13 @@
 #ifndef GRPC_TEST_CORE_UTIL_MEMORY_COUNTERS_H
 #define GRPC_TEST_CORE_UTIL_MEMORY_COUNTERS_H
 
-#include <stddef.h>
+#include <grpc/support/atm.h>
 
 struct grpc_memory_counters {
-  size_t total_size_relative;
-  size_t total_size_absolute;
-  size_t total_allocs_relative;
-  size_t total_allocs_absolute;
+  gpr_atm total_size_relative;
+  gpr_atm total_size_absolute;
+  gpr_atm total_allocs_relative;
+  gpr_atm total_allocs_absolute;
 };
 
 void grpc_memory_counters_init();