diff --git a/src/utils/string_interning.cc b/src/utils/string_interning.cc index f2cc81740..be6c3cf2c 100644 --- a/src/utils/string_interning.cc +++ b/src/utils/string_interning.cc @@ -22,6 +22,8 @@ MemoryPool StringInternStore::memory_pool_{0}; InternedString::InternedString(absl::string_view str, bool shared) : length_(str.length()), is_shared_(shared), is_data_owner_(true) { + IsolatedMemoryScope scope { StringInternStore::memory_pool_ }; + data_ = new char[length_ + 1]; memcpy(data_, str.data(), length_); data_[length_] = '\0'; @@ -37,7 +39,7 @@ InternedString::~InternedString() { if (is_shared_) { StringInternStore::Instance().Release(this); } - + if (is_data_owner_) { delete[] data_; } else { diff --git a/testing/utils/string_interning_test.cc b/testing/utils/string_interning_test.cc index 7912954d8..631a321a4 100644 --- a/testing/utils/string_interning_test.cc +++ b/testing/utils/string_interning_test.cc @@ -126,7 +126,34 @@ TEST_F(StringInterningTest, StringInternStoreTracksMemoryInternally) { EXPECT_EQ(StringInternStore::GetMemoryUsage(), 12); interned_str.reset(); + + EXPECT_EQ(StringInternStore::GetMemoryUsage(), 0); +} + +#ifndef SAN_BUILD +TEST_F(StringInterningTest, NonInternedStringDoesAffectStoreMemory) { + static auto track_malloc_size = [](void* ptr) -> size_t { + return 21; + }; + + vmsdk::test_utils::SetTestSystemMallocSizeFunction(track_malloc_size); + + const char* test_str = "non_interned_string"; + size_t expected_size = strlen(test_str) + 1; + + int64_t initial_memory = StringInternStore::GetMemoryUsage(); + + auto non_interned = std::make_shared(test_str); + + EXPECT_EQ(StringInternStore::GetMemoryUsage(), initial_memory+21); + + non_interned.reset(); + + EXPECT_EQ(StringInternStore::GetMemoryUsage(), initial_memory); + + vmsdk::test_utils::ClearTestSystemMallocSizeFunction(); } +#endif INSTANTIATE_TEST_SUITE_P(StringInterningTests, StringInterningTest, ::testing::Values(true, false), diff --git a/vmsdk/src/memory_allocation_overrides.cc b/vmsdk/src/memory_allocation_overrides.cc index 1a1b541da..17d7677c9 100644 --- a/vmsdk/src/memory_allocation_overrides.cc +++ b/vmsdk/src/memory_allocation_overrides.cc @@ -130,13 +130,32 @@ void* PerformAndTrackAlignedAlloc(size_t align, size_t size, } return ptr; } + +namespace test_utils { + +thread_local size_t (*test_malloc_size_fn)(void*) = nullptr; + +void SetTestSystemMallocSizeFunction(size_t (*fn)(void*)) { + test_malloc_size_fn = fn; +} + +void ClearTestSystemMallocSizeFunction() { + test_malloc_size_fn = nullptr; +} + +} // namespace test_utils } // namespace vmsdk extern "C" { -// Our allocator doesn't support tracking system memory size, so we just -// return 0. +// Basically our allocator doesn't support tracking system memory size, so we just +// return 0. But if test_malloc_size_fn is set, tracking system memory size is possible. // NOLINTNEXTLINE -__attribute__((weak)) size_t empty_usable_size(void* ptr) noexcept { return 0; } +__attribute__((weak)) size_t usable_size(void* ptr) noexcept { + if (vmsdk::test_utils::test_malloc_size_fn) { + return vmsdk::test_utils::test_malloc_size_fn(ptr); + } + return 0; +} // For Valkey allocation - we need to ensure alignment by taking advantage of // jemalloc alignment properties, as there is no aligned malloc module @@ -152,7 +171,7 @@ size_t AlignSize(size_t size, int alignment = 16) { void* __wrap_malloc(size_t size) noexcept { if (!vmsdk::IsUsingValkeyAlloc()) { auto ptr = - vmsdk::PerformAndTrackMalloc(size, __real_malloc, empty_usable_size); + vmsdk::PerformAndTrackMalloc(size, __real_malloc, usable_size); vmsdk::SystemAllocTracker::GetInstance().TrackPointer(ptr); return ptr; } @@ -172,7 +191,7 @@ void __wrap_free(void* ptr) noexcept { // another DSO which doesn't have our wrapped symbols (namely libc.so). For // this reason, we bypass the tracking during the bootstrap phase. if (was_tracked || !vmsdk::IsUsingValkeyAlloc()) { - vmsdk::PerformAndTrackFree(ptr, __real_free, empty_usable_size); + vmsdk::PerformAndTrackFree(ptr, __real_free, usable_size); } else { vmsdk::PerformAndTrackFree(ptr, ValkeyModule_Free, ValkeyModule_MallocUsableSize); @@ -182,7 +201,7 @@ void __wrap_free(void* ptr) noexcept { void* __wrap_calloc(size_t __nmemb, size_t size) noexcept { if (!vmsdk::IsUsingValkeyAlloc()) { auto ptr = vmsdk::PerformAndTrackCalloc(__nmemb, size, __real_calloc, - empty_usable_size); + usable_size); vmsdk::SystemAllocTracker::GetInstance().TrackPointer(ptr); return ptr; } @@ -203,7 +222,7 @@ void* __wrap_realloc(void* ptr, size_t size) noexcept { ValkeyModule_MallocUsableSize); } else { auto new_ptr = vmsdk::PerformAndTrackRealloc(ptr, size, __real_realloc, - empty_usable_size); + usable_size); vmsdk::SystemAllocTracker::GetInstance().TrackPointer(new_ptr); return new_ptr; } @@ -212,7 +231,7 @@ void* __wrap_realloc(void* ptr, size_t size) noexcept { void* __wrap_aligned_alloc(size_t __alignment, size_t __size) noexcept { if (!vmsdk::IsUsingValkeyAlloc()) { auto ptr = vmsdk::PerformAndTrackAlignedAlloc( - __alignment, __size, __real_aligned_alloc, empty_usable_size); + __alignment, __size, __real_aligned_alloc, usable_size); vmsdk::SystemAllocTracker::GetInstance().TrackPointer(ptr); return ptr; } @@ -224,7 +243,7 @@ void* __wrap_aligned_alloc(size_t __alignment, size_t __size) noexcept { int __wrap_malloc_usable_size(void* ptr) noexcept { if (vmsdk::SystemAllocTracker::GetInstance().IsTracked(ptr)) { - return empty_usable_size(ptr); + return usable_size(ptr); } return ValkeyModule_MallocUsableSize(ptr); } diff --git a/vmsdk/src/memory_allocation_overrides.h b/vmsdk/src/memory_allocation_overrides.h index d082b0d44..6b4054fcb 100644 --- a/vmsdk/src/memory_allocation_overrides.h +++ b/vmsdk/src/memory_allocation_overrides.h @@ -36,7 +36,7 @@ WEAK_SYMBOL int (*__real_posix_memalign)(void**, size_t, // NOLINTNEXTLINE WEAK_SYMBOL void* (*__real_valloc)(size_t) = valloc; // NOLINTNEXTLINE -__attribute__((weak)) size_t empty_usable_size(void* ptr) noexcept; +__attribute__((weak)) size_t usable_size(void* ptr) noexcept; } // extern "C" // Different exception specifier between CLANG & GCC @@ -108,4 +108,23 @@ void operator delete[](void* p, std::align_val_t alignment, void operator delete[](void* p, size_t size, std::align_val_t alignment) noexcept; #endif // !SAN_BUILD + +namespace vmsdk { +namespace test_utils { + +// Set a custom malloc size function for testing purposes. +// This allows tests to provide their own implementation of malloc_usable_size +// for system allocations, enabling accurate memory tracking in tests. +// The function pointer is thread-local, so it only affects the calling thread. +// +// @param fn Function pointer that takes a void* and returns the allocated size. +// Pass nullptr to clear the test function. +void SetTestSystemMallocSizeFunction(size_t (*fn)(void*)); + +// Clear the test malloc size function, reverting to default behavior. +void ClearTestSystemMallocSizeFunction(); + +} // namespace test_utils +} // namespace vmsdk + #endif // VMSDK_SRC_MEMORY_ALLOCATION_OVERRIDES_H_