[scudo] Drain caches when release with M_PURGE_ALL am: b226587a7f

Original change: https://googleplex-android-review.googlesource.com/c/platform/external/scudo/+/23184463

Change-Id: I5884092acb0925e46034c362a173f7992a353e1c
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

5 files changed, 54 insertions, 0 deletions

diff --git a/standalone/combined.h b/standalone/combined.h
index d5365b689aa..006605659bf 100644
--- a/standalone/combined.h
+++ b/standalone/combined.h
@@ -239,6 +239,7 @@ public:
   }
 
   TSDRegistryT *getTSDRegistry() { return &TSDRegistry; }
+  QuarantineT *getQuarantine() { return &Quarantine; }
 
   // The Cache must be provided zero-initialized.
   void initCache(CacheT *Cache) { Cache->init(&Stats, &Primary); }
@@ -254,6 +255,13 @@ public:
     TSD->getCache().destroy(&Stats);
   }
 
+  void drainCache(TSD<ThisT> *TSD) {
+    Quarantine.drainAndRecycle(&TSD->getQuarantineCache(),
+                               QuarantineCallback(*this, TSD->getCache()));
+    TSD->getCache().drain();
+  }
+  void drainCaches() { TSDRegistry.drainCaches(this); }
+
   ALWAYS_INLINE void *getHeaderTaggedPointer(void *Ptr) {
     if (!allocatorSupportsMemoryTagging<Params>())
       return Ptr;
@@ -747,6 +755,8 @@ public:
 
   void releaseToOS(ReleaseToOS ReleaseType) {
     initThreadMaybe();
+    if (ReleaseType == ReleaseToOS::ForceAll)
+      drainCaches();
     Primary.releaseToOS(ReleaseType);
     Secondary.releaseToOS();
   }
diff --git a/standalone/quarantine.h b/standalone/quarantine.h
index e65a733ced7..b5f8db0e87c 100644
--- a/standalone/quarantine.h
+++ b/standalone/quarantine.h
@@ -192,6 +192,12 @@ public:
   uptr getMaxSize() const { return atomic_load_relaxed(&MaxSize); }
   uptr getCacheSize() const { return atomic_load_relaxed(&MaxCacheSize); }
 
+  // This is supposed to be used in test only.
+  bool isEmpty() {
+    ScopedLock L(CacheMutex);
+    return Cache.getSize() == 0U;
+  }
+
   void put(CacheT *C, Callback Cb, Node *Ptr, uptr Size) {
     C->enqueue(Cb, Ptr, Size);
     if (C->getSize() > getCacheSize())
diff --git a/standalone/tests/combined_test.cpp b/standalone/tests/combined_test.cpp
index 33a309e42d6..44ba639f7aa 100644
--- a/standalone/tests/combined_test.cpp
+++ b/standalone/tests/combined_test.cpp
@@ -457,6 +457,28 @@ SCUDO_TYPED_TEST(ScudoCombinedTest, CacheDrain) NO_THREAD_SAFETY_ANALYSIS {
     TSD->unlock();
 }
 
+SCUDO_TYPED_TEST(ScudoCombinedTest, ForceCacheDrain) NO_THREAD_SAFETY_ANALYSIS {
+  auto *Allocator = this->Allocator.get();
+
+  std::vector<void *> V;
+  for (scudo::uptr I = 0; I < 64U; I++)
+    V.push_back(Allocator->allocate(
+        rand() % (TypeParam::Primary::SizeClassMap::MaxSize / 2U), Origin));
+  for (auto P : V)
+    Allocator->deallocate(P, Origin);
+
+  // `ForceAll` will also drain the caches.
+  Allocator->releaseToOS(scudo::ReleaseToOS::ForceAll);
+
+  bool UnlockRequired;
+  auto *TSD = Allocator->getTSDRegistry()->getTSDAndLock(&UnlockRequired);
+  EXPECT_TRUE(TSD->getCache().isEmpty());
+  EXPECT_EQ(TSD->getQuarantineCache().getSize(), 0U);
+  EXPECT_TRUE(Allocator->getQuarantine()->isEmpty());
+  if (UnlockRequired)
+    TSD->unlock();
+}
+
 SCUDO_TYPED_TEST(ScudoCombinedTest, ThreadedCombined) {
   std::mutex Mutex;
   std::condition_variable Cv;
diff --git a/standalone/tsd_exclusive.h b/standalone/tsd_exclusive.h
index aca9fc9b4e8..9d037731921 100644
--- a/standalone/tsd_exclusive.h
+++ b/standalone/tsd_exclusive.h
@@ -59,6 +59,13 @@ template <class Allocator> struct TSDRegistryExT {
     Initialized = false;
   }
 
+  void drainCaches(Allocator *Instance) {
+    // We don't have a way to iterate all thread local `ThreadTSD`s. Simply
+    // drain the `ThreadTSD` of current thread and `FallbackTSD`.
+    Instance->drainCache(&ThreadTSD);
+    Instance->drainCache(&FallbackTSD);
+  }
+
   ALWAYS_INLINE void initThreadMaybe(Allocator *Instance, bool MinimalInit) {
     if (LIKELY(State.InitState != ThreadState::NotInitialized))
       return;
diff --git a/standalone/tsd_shared.h b/standalone/tsd_shared.h
index e193281fc73..dcb0948ad78 100644
--- a/standalone/tsd_shared.h
+++ b/standalone/tsd_shared.h
@@ -54,6 +54,15 @@ struct TSDRegistrySharedT {
     Initialized = false;
   }
 
+  void drainCaches(Allocator *Instance) {
+    ScopedLock L(MutexTSDs);
+    for (uptr I = 0; I < NumberOfTSDs; ++I) {
+      TSDs[I].lock();
+      Instance->drainCache(&TSDs[I]);
+      TSDs[I].unlock();
+    }
+  }
+
   ALWAYS_INLINE void initThreadMaybe(Allocator *Instance,
                                      UNUSED bool MinimalInit) {
     if (LIKELY(getCurrentTSD()))