[scudo] Drain caches when release with M_PURGE_ALL

[Chia-hungDuan]

Summary

This will drain both quarantine and local caches.

[cferris]

LGTM

compiler-rt/lib/scudo/standalone/combined.h

 #ifdef GWP_ASAN_HOOKS
     if (getFlags()->GWP_ASAN_InstallSignalHandlers)
       gwp_asan::segv_handler::uninstallSignalHandlers();
     GuardedAlloc.uninitTestOnly();
 #endif // GWP_ASAN_HOOKS
   }
 
   TSDRegistryT *getTSDRegistry() { return &TSDRegistry; }
+  QuarantineT *getQuarantine() { return &Quarantine; }
 
   // The Cache must be provided zero-initialized.
   void initCache(CacheT *Cache) { Cache->init(&Stats, &Primary); }
 
   // Release the resources used by a TSD, which involves:
   // - draining the local quarantine cache to the global quarantine;
   // - releasing the cached pointers back to the Primary;
   // - unlinking the local stats from the global ones (destroying the cache does
   //   the last two items).
   void commitBack(TSD<ThisT> *TSD) {
     Quarantine.drain(&TSD->getQuarantineCache(),
                      QuarantineCallback(*this, TSD->getCache()));
     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;
     auto UntaggedPtr = untagPointer(Ptr);
     if (UntaggedPtr != Ptr)
       return UntaggedPtr;
     // Secondary, or pointer allocated while memory tagging is unsupported or
     // disabled. The tag mismatch is okay in the latter case because tags will
#endif
   void printStats() {
     ScopedString Str;
     getStats(&Str);
     Str.output();
   }
 
   void releaseToOS(ReleaseToOS ReleaseType) {
     initThreadMaybe();
+    if (ReleaseType == ReleaseToOS::ForceAll)
+      drainCaches();
     Primary.releaseToOS(ReleaseType);
     Secondary.releaseToOS();
   }
 
   // Iterate over all chunks and call a callback for all busy chunks located
   // within the provided memory range. Said callback must not use this allocator
   // or a deadlock can ensue. This fits Android's malloc_iterate() needs.
   void iterateOverChunks(uptr Base, uptr Size, iterate_callback Callback,

compiler-rt/lib/scudo/standalone/quarantine.h

void init(uptr Size, uptr CacheSize) NO_THREAD_SAFETY_ANALYSIS {
     atomic_store_relaxed(&MaxCacheSize, CacheSize);
 
     Cache.init();
   }
 
   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())
       drain(C, Cb);
   }
 
   void NOINLINE drain(CacheT *C, Callback Cb) EXCLUDES(CacheMutex) {
     bool needRecycle = false;

compiler-rt/lib/scudo/standalone/tests/combined_test.cpp

SCUDO_TYPED_TEST(ScudoCombinedTest, CacheDrain) NO_THREAD_SAFETY_ANALYSIS {
   auto *TSD = Allocator->getTSDRegistry()->getTSDAndLock(&UnlockRequired);
   EXPECT_TRUE(!TSD->getCache().isEmpty());
   TSD->getCache().drain();
   EXPECT_TRUE(TSD->getCache().isEmpty());
   if (UnlockRequired)
     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;
   bool Ready = false;
   auto *Allocator = this->Allocator.get();
   std::thread Threads[32];
   for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++)
     Threads[I] = std::thread([&]() {

compiler-rt/lib/scudo/standalone/tsd_exclusive.h

void unmapTestOnly(Allocator *Instance) EXCLUDES(Mutex) {
     PThreadKey = {};
     FallbackTSD.commitBack(Instance);
     FallbackTSD = {};
     State = {};
     ScopedLock L(Mutex);
     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;
     initThread(Instance, MinimalInit);
   }
 
   // TODO(chiahungduan): Consider removing the argument `UnlockRequired` by
   // embedding the logic into TSD or always locking the TSD. It will enable us

compiler-rt/lib/scudo/standalone/tsd_shared.h

for (u32 I = 0; I < TSDsArraySize; I++) {
       TSDs[I].commitBack(Instance);
       TSDs[I] = {};
     }
     setCurrentTSD(nullptr);
     ScopedLock L(Mutex);
     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()))
       return;
     initThread(Instance);
   }
 
   // TSDs is an array of locks and which is not supported for marking