2 files changed, 4 insertions, 16 deletions
diff --git a/src/tests/vnet/session/tcp_echo.c b/src/tests/vnet/session/tcp_echo.c
index 6c99294099b..6e41e4b2c1a 100644
--- a/src/tests/vnet/session/tcp_echo.c
+++ b/src/tests/vnet/session/tcp_echo.c
@@ -1409,21 +1409,14 @@ main (int argc, char **argv)
   echo_main_t *em = &echo_main;
   unformat_input_t _argv, *a = &_argv;
   u8 *chroot_prefix;
-  u8 *heap, *uri = 0;
+  u8 *uri = 0;
   u8 *bind_uri = (u8 *) "tcp://0.0.0.0/1234";
   u8 *connect_uri = (u8 *) "tcp://6.0.1.1/1234";
   u64 bytes_to_send = 64 << 10, mbytes;
   char *app_name;
   u32 tmp;
-  mheap_t *h;
 
-  clib_mem_init (0, 256 << 20);
-
-  heap = clib_mem_get_per_cpu_heap ();
-  h = mheap_header (heap);
-
-  /* make the main heap thread-safe */
-  h->flags |= MHEAP_FLAG_THREAD_SAFE;
+  clib_mem_init_thread_safe (0, 256 << 20);
 
   memset (em, 0, sizeof (*em));
   em->session_index_by_vpp_handles = hash_create (0, sizeof (uword));
diff --git a/src/tests/vnet/session/udp_echo.c b/src/tests/vnet/session/udp_echo.c
index 2e176ff4b7b..5691ca4f042 100644
--- a/src/tests/vnet/session/udp_echo.c
+++ b/src/tests/vnet/session/udp_echo.c
@@ -1207,16 +1207,11 @@ main (int argc, char **argv)
   app_session_t *session;
   u8 *chroot_prefix;
   char *app_name;
-  mheap_t *h;
-  u8 *heap;
   u32 tmp;
   int i;
 
-  clib_mem_init (0, 256 << 20);
-  heap = clib_mem_get_per_cpu_heap ();
-  h = mheap_header (heap);
-  /* make the main heap thread-safe */
-  h->flags |= MHEAP_FLAG_THREAD_SAFE;
+  clib_mem_init_thread_safe (0, 256 << 20);
+
   svm_fifo_segment_main_init (0x200000000ULL, 20);
 
   vec_validate (utm->rx_buf, 8192);