diff --git a/tcmalloc/malloc_extension.cc b/tcmalloc/malloc_extension.cc index decedfbf4..dd70cc422 100644 --- a/tcmalloc/malloc_extension.cc +++ b/tcmalloc/malloc_extension.cc @@ -466,6 +466,20 @@ absl::Duration MallocExtension::GetSkipSubreleaseInterval() { #endif } +static std::atomic SoftMemoryLimitHandler_; + +void MallocExtension::SetSoftMemoryLimitHandler(SoftMemoryLimitCallback* handler) { +#if ABSL_INTERNAL_HAVE_WEAK_MALLOCEXTENSION_STUBS + SoftMemoryLimitHandler_.store(handler); +#endif +} + +MallocExtension::SoftMemoryLimitCallback* MallocExtension::GetSoftMemoryLimitHandler() { +#if ABSL_INTERNAL_HAVE_WEAK_MALLOCEXTENSION_STUBS + return SoftMemoryLimitHandler_.load(); +#endif +} + void MallocExtension::SetSkipSubreleaseInterval(absl::Duration value) { #if ABSL_INTERNAL_HAVE_WEAK_MALLOCEXTENSION_STUBS if (MallocExtension_Internal_SetSkipSubreleaseInterval == nullptr) { diff --git a/tcmalloc/malloc_extension.h b/tcmalloc/malloc_extension.h index 1c5377f4d..2fa081782 100644 --- a/tcmalloc/malloc_extension.h +++ b/tcmalloc/malloc_extension.h @@ -646,6 +646,14 @@ class MallocExtension final { // Specifies the release rate from the page heap. ProcessBackgroundActions // must be called for this to be operative. static void SetBackgroundReleaseRate(BytesPerSecond rate); + + // Specifies the pointer to a callback function to be invoked when the soft + // memory limit is reached. + using SoftMemoryLimitCallback = void(); + static void SetSoftMemoryLimitHandler(SoftMemoryLimitCallback* handler); + + // Gets the current handler of the soft memory limit event. + static SoftMemoryLimitCallback* GetSoftMemoryLimitHandler(); }; } // namespace tcmalloc diff --git a/tcmalloc/page_allocator.cc b/tcmalloc/page_allocator.cc index 2f450d8ac..f7232975c 100644 --- a/tcmalloc/page_allocator.cc +++ b/tcmalloc/page_allocator.cc @@ -138,6 +138,10 @@ void PageAllocator::ShrinkToUsageLimit(Length n) { warned = true; TC_LOG("Couldn't respect usage limit of %v and OOM is likely to follow.", limits_[kSoft]); + + if (auto* handler = MallocExtension::GetSoftMemoryLimitHandler()) { + (*handler)(); + } } bool PageAllocator::ShrinkHardBy(Length pages, LimitKind limit_kind) {