Ensure internal_methods_only TCMalloc variant works under sanitizers.

Sanitizers replace malloc/free, so we need `MallocExtension::GetAllocatedSize`
& friends to get a matching implementation.  If we keep
`MallocExtension_Internal_GetAllocatedSize` defined in tcmalloc.cc, we get
TCMalloc's version and do not recognize these pointers since they came from the
sanitizer.

PiperOrigin-RevId: 899185297
Change-Id: Ib830c76f5334c398bd1a0558f4325720b0c0f40b

Author: ckennelly

tcmalloc/tcmalloc.cc

@@ -1312,14 +1312,13 @@ extern "C" void MallocHook_HooksChanged() {
 using tcmalloc::tcmalloc_internal::GetOwnership;
 using tcmalloc::tcmalloc_internal::GetSize;
 
-extern "C" size_t MallocExtension_Internal_GetAllocatedSize(const void* ptr) {
+extern "C" size_t TCMalloc_Internal_GetAllocatedSize(const void* ptr) {
   TC_ASSERT(!ptr || GetOwnership(ptr) !=
                         tcmalloc::MallocExtension::Ownership::kNotOwned);
   return GetSize(ptr);
 }
 
-extern "C" size_t MallocExtension_Internal_GetEstimatedAllocatedSize(
-    size_t size) {
+extern "C" size_t TCMalloc_Internal_GetEstimatedAllocatedSize(size_t size) {
   if (ABSL_PREDICT_FALSE(!tc_globals.IsInited())) {
     return tcmalloc::tcmalloc_internal::nallocx_slow(size, 0);
   }
@@ -1333,7 +1332,7 @@ extern "C" size_t MallocExtension_Internal_GetEstimatedAllocatedSize(
   }
 }
 
-extern "C" void MallocExtension_Internal_MarkThreadBusy() {
+extern "C" void TCMalloc_Internal_MarkThreadBusy() {
   tc_globals.InitIfNecessary();
 
   if (UsePerCpuCache(tc_globals)) {
@@ -1344,6 +1343,21 @@ extern "C" void MallocExtension_Internal_MarkThreadBusy() {
   tcmalloc::tcmalloc_internal::ThreadCache::GetCache();
 }
 
+#ifndef TCMALLOC_INTERNAL_METHODS_ONLY
+extern "C" size_t MallocExtension_Internal_GetAllocatedSize(const void* ptr) {
+  return TCMalloc_Internal_GetAllocatedSize(ptr);
+}
+
+extern "C" size_t MallocExtension_Internal_GetEstimatedAllocatedSize(
+    size_t size) {
+  return TCMalloc_Internal_GetEstimatedAllocatedSize(size);
+}
+
+extern "C" void MallocExtension_Internal_MarkThreadBusy() {
+  TCMalloc_Internal_MarkThreadBusy();
+}
+#endif  // !TCMALLOC_INTERNAL_METHODS_ONLY
+
 absl::StatusOr<tcmalloc::malloc_tracing_extension::AllocatedAddressRanges>
 MallocTracingExtension_Internal_GetAllocatedAddressRanges() {
   tcmalloc::malloc_tracing_extension::AllocatedAddressRanges

tcmalloc/tcmalloc.h

@@ -281,6 +281,10 @@ static_assert(ALLOC_TOKEN_FALLBACK >= ALLOC_TOKEN_MAX);
 DECLARE_ALLOC_TOKEN_VARIANTS(ALLOC_TOKEN_FALLBACK)
 #endif
 
+size_t TCMalloc_Internal_GetAllocatedSize(const void* ptr);
+size_t TCMalloc_Internal_GetEstimatedAllocatedSize(size_t size);
+void TCMalloc_Internal_MarkThreadBusy();
+
 }  // extern "C"
 
 #endif  // TCMALLOC_TCMALLOC_H_

tcmalloc/testing/BUILD

@@ -889,6 +889,20 @@ create_tcmalloc_testsuite(
     ],
 )
 
+cc_test(
+    name = "malloc_extension_internal_methods_test",
+    srcs = ["malloc_extension_system_malloc_test.cc"],
+    copts = TCMALLOC_DEFAULT_COPTS,
+    malloc = "//tcmalloc:tcmalloc_internal_methods_only",
+    tags = ["noubsan"],
+    deps = [
+        "//tcmalloc:malloc_extension",
+        "//tcmalloc/internal:config",
+        "@com_google_absl//absl/random",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "malloc_extension_system_malloc_test",
     srcs = ["malloc_extension_system_malloc_test.cc"],
@@ -897,6 +911,7 @@ cc_test(
     tags = ["noubsan"],
     deps = [
         "//tcmalloc:malloc_extension",
+        "//tcmalloc/internal:config",
         "@com_google_absl//absl/random",
         "@com_google_googletest//:gtest_main",
     ],

tcmalloc/testing/CMakeLists.txt

@@ -803,6 +803,21 @@ tcmalloc_cc_test_variants(
     "tcmalloc::malloc_extension"
 )
 
+tcmalloc_cc_test(
+  NAME
+    tcmalloc_testing_malloc_extension_internal_methods_test
+  SRCS
+    "malloc_extension_system_malloc_test.cc"
+  DEPS
+    "GTest::gtest_main"
+    "GTest::gmock_main"
+    "GTest::gmock"
+    "absl::random_random"
+    "tcmalloc::internal_config"
+    "tcmalloc::malloc_extension"
+    "tcmalloc::tcmalloc_internal_methods_only"
+)
+
 tcmalloc_cc_test(
   NAME
     tcmalloc_testing_malloc_extension_system_malloc_test
@@ -813,6 +828,7 @@ tcmalloc_cc_test(
     "GTest::gmock_main"
     "GTest::gmock"
     "absl::random_random"
+    "tcmalloc::internal_config"
     "tcmalloc::malloc_extension"
 )
 

tcmalloc/testing/malloc_extension_system_malloc_test.cc

@@ -20,6 +20,7 @@
 
 #include "gtest/gtest.h"
 #include "absl/random/random.h"
+#include "tcmalloc/internal/config.h"
 #include "tcmalloc/malloc_extension.h"
 
 namespace tcmalloc {
@@ -86,5 +87,19 @@ TEST(MallocExtension, AllocationProfile) {
   }
 }
 
+TEST(MallocExtension, GetAllocatedSizeSystemMalloc) {
+  void* ptr = ::operator new(100);
+  std::optional<size_t> size = MallocExtension::GetAllocatedSize(ptr);
+
+  if (tcmalloc_internal::kSanitizerPresent) {
+    EXPECT_TRUE(size.has_value());
+    EXPECT_GE(*size, 100);
+  } else {
+    EXPECT_FALSE(size.has_value());
+  }
+
+  ::operator delete(ptr);
+}
+
 }  // namespace
 }  // namespace tcmalloc

tcmalloc/testing/parallel_test.cc

@@ -29,6 +29,7 @@ extern "C" {
 void* TCMallocInternalNew(size_t);
 void TCMallocInternalDelete(void*);
 void TCMallocInternalDeleteSized(void*, size_t);
+size_t TCMalloc_Internal_GetAllocatedSize(const void*);
 
 }  // extern "C"
 
@@ -53,7 +54,7 @@ struct Allocator {
       }
 
       for (void* ptr : v) {
-        ABSL_CHECK_GE(*MallocExtension::GetAllocatedSize(ptr), size);
+        ABSL_CHECK_GE(TCMalloc_Internal_GetAllocatedSize(ptr), size);
       }
 
       for (void* ptr : v) {