Bugfix: potential incorrect real size of aligned allocation (StarRock…

…s#1601)

In order to get the real size of aligned allocation, MALLOCX_LG_ALIGN
should be passed as the second argument of tc_nallocx.
reference: https://www.unix.com/man-page/freebsd/3/nallocx

be/src/service/mem_hook.cpp

@@ -169,139 +169,96 @@ void operator delete[](void* p, size_t size, std::align_val_t al) noexcept {
 }
 */
 
+#ifndef BE_TEST
+#define TC_MALLOC_SIZE(ptr) tc_malloc_size(ptr)
+#define MEMORY_CONSUME_PTR(ptr) starrocks::tls_thread_status.mem_consume(tc_malloc_size(ptr))
+#define MEMORY_RELEASE_PTR(ptr) starrocks::tls_thread_status.mem_release(tc_malloc_size(ptr))
+#define MEMORY_CONSUME_SIZE(size) starrocks::tls_thread_status.mem_consume(size)
+#else
+#define TC_MALLOC_SIZE(ptr) 0
+#define MEMORY_CONSUME_PTR(ptr) (void)(ptr)
+#define MEMORY_RELEASE_PTR(ptr) (void)(ptr)
+#define MEMORY_CONSUME_SIZE(size) (void)(size)
+#endif
+
 extern "C" {
 // malloc
 void* my_alloc(size_t size) __THROW {
     void* ptr = tc_malloc(size);
-
+    // NOTE: do NOT call `tc_malloc_size` here, it may call the new operator, which in turn will
+    // call the `my_alloc`, and result in a deadloop.
 #ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        size_t actual_size = tc_nallocx(size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
-    }
+    MEMORY_CONSUME_SIZE(tc_nallocx(size, 0));
 #endif
-
     return ptr;
 }
 
 // free
 void my_free(void* p) __THROW {
-#ifndef BE_TEST
-    size_t size = tc_malloc_size(p);
-    starrocks::tls_thread_status.mem_release(size);
-#endif
-
+    MEMORY_RELEASE_PTR(p);
     tc_free(p);
 }
 
 // realloc
 void* my_realloc(void* p, size_t size) __THROW {
+    int64_t old_size = TC_MALLOC_SIZE(p);
     void* ptr = tc_realloc(p, size);
-
-#ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        int64_t old_size = 0;
-        if (p != nullptr) {
-            old_size = tc_malloc_size(p);
-        }
-        int64_t actual_size = tc_nallocx(size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size - old_size);
+    if (ptr != nullptr) {
+        int64_t new_size = TC_MALLOC_SIZE(ptr);
+        MEMORY_CONSUME_SIZE(new_size - old_size);
+    } else {
+        // nothing to do.
+        // If tc_realloc() fails the original block is left untouched; it is not freed or moved
     }
-#endif
-
     return ptr;
 }
 
 // calloc
 void* my_calloc(size_t n, size_t size) __THROW {
     void* ptr = tc_calloc(n, size);
-
-#ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        size_t actual_size = tc_nallocx(n * size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
-    }
-#endif
-
+    MEMORY_CONSUME_PTR(ptr);
     return ptr;
 }
 
 void my_cfree(void* ptr) __THROW {
-#ifndef BE_TEST
-    size_t size = tc_malloc_size(ptr);
-    starrocks::tls_thread_status.mem_release(size);
-#endif
-
+    MEMORY_RELEASE_PTR(ptr);
     tc_cfree(ptr);
 }
 
 // memalign
 void* my_memalign(size_t align, size_t size) __THROW {
     void* ptr = tc_memalign(align, size);
-
-#ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        size_t actual_size = tc_nallocx(size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
-    }
-#endif
-
+    MEMORY_CONSUME_PTR(ptr);
     return ptr;
 }
 
 // aligned_alloc
 void* my_aligned_alloc(size_t align, size_t size) __THROW {
     void* ptr = tc_memalign(align, size);
-
-#ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        size_t actual_size = tc_nallocx(size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
-    }
-#endif
-
+    MEMORY_CONSUME_PTR(ptr);
     return ptr;
 }
 
 // valloc
 void* my_valloc(size_t size) __THROW {
     void* ptr = tc_valloc(size);
-
-#ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        size_t actual_size = tc_nallocx(size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
-    }
-#endif
-
+    MEMORY_CONSUME_PTR(ptr);
     return ptr;
 }
 
 // pvalloc
 void* my_pvalloc(size_t size) __THROW {
     void* ptr = tc_pvalloc(size);
-
-#ifndef BE_TEST
-    if (LIKELY(ptr != nullptr)) {
-        size_t actual_size = tc_nallocx(size, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
-    }
-#endif
-
+    MEMORY_CONSUME_PTR(ptr);
     return ptr;
 }
 
 // posix_memalign
 int my_posix_memalign(void** r, size_t a, size_t s) __THROW {
     int ret = tc_posix_memalign(r, a, s);
-
-#ifndef BE_TEST
-    if (LIKELY(ret == 0)) {
-        size_t actual_size = tc_nallocx(s, 0);
-        starrocks::tls_thread_status.mem_consume(actual_size);
+    if (ret == 0) {
+        MEMORY_CONSUME_PTR(*r);
     }
-#endif
-
     return ret;
 }