[scudo][standalone] Unmap memory in tests

[cryptoad]

Summary

The more tests are added, the more we are limited by the size of the
address space on 32-bit. Implement unmapTestOnly all around (like it
is in sanitzer_common) to be able to free up some memory.
This is not intended to be a proper "destructor" for an allocator, but
allows us to not fail due to having no memory left.

[morehouse]

lib/scudo/standalone/primary32.h

93

This looks incorrect. Is I * RegionSize a valid address?

lib/scudo/standalone/tests/primary_test.cc

25–29

Should this also delete P? Same for below.

[cryptoad]

lib/scudo/standalone/primary32.h

93

It is indeed a valid address. For the 32-bit primary, the address space is split in potential regions of RegionSize bytes, that we index via a one or two level bytemap.

lib/scudo/standalone/tests/primary_test.cc

25–29

Good point thanks, I somehow thought it would still end up deleting the original pointer with a custom deleter. I will correct those.

[cryptoad]

Update the deleters to actually delete the pointer itself as well.

[morehouse]

LGTM

lib/scudo/standalone/bytemap.h

 
 template <uptr Size> class FlatByteMap {
 public:
   void initLinkerInitialized() {
     Map = reinterpret_cast<u8 *>(map(nullptr, Size, "scudo:bytemap"));
   }
   void init() { initLinkerInitialized(); }
 
+  void unmapTestOnly() { unmap(reinterpret_cast<void *>(Map), Size); }
+
   void set(uptr Index, u8 Value) {
     DCHECK_LT(Index, Size);
     DCHECK_EQ(0U, Map[Index]);
     Map[Index] = Value;
   }
   u8 operator[](uptr Index) {
     DCHECK_LT(Index, Size);
     return Map[Index];
for (uptr I = 0; I < Level1Size; I++) {
       u8 *P = get(I);
       if (!P)
         continue;
       unmap(P, Level2Size);
     }
     memset(Level1Map, 0, sizeof(atomic_uptr) * Level1Size);
   }
 
+  void unmapTestOnly() {
+    reset();
+    unmap(reinterpret_cast<void *>(Level1Map),
+          sizeof(atomic_uptr) * Level1Size);
+  }
+
   uptr size() const { return Level1Size * Level2Size; }
 
   void set(uptr Index, u8 Value) {
     DCHECK_LT(Index, Level1Size * Level2Size);
     u8 *Level2Map = getOrCreate(Index / Level2Size);
     DCHECK_EQ(0U, Level2Map[Index % Level2Size]);
     Level2Map[Index % Level2Size] = Value;
   }

lib/scudo/standalone/primary32.h

void initLinkerInitialized(s32 ReleaseToOsInterval) {
     }
     ReleaseToOsIntervalMs = ReleaseToOsInterval;
   }
   void init(s32 ReleaseToOsInterval) {
     memset(this, 0, sizeof(*this));
     initLinkerInitialized(ReleaseToOsInterval);
   }
 
+  void unmapTestOnly() {
+    while (NumberOfStashedRegions > 0)
+      unmap(reinterpret_cast<void *>(RegionsStash[--NumberOfStashedRegions]),
+            RegionSize);
+    // TODO(kostyak): unmap the TransferBatch regions as well.
+    for (uptr I = 0; I < NumRegions; I++)
+      if (PossibleRegions[I])
+        unmap(reinterpret_cast<void *>(I * RegionSize), RegionSize);
+    PossibleRegions.unmapTestOnly();
+  }
+
   TransferBatch *popBatch(CacheT *C, uptr ClassId) {
     DCHECK_LT(ClassId, NumClasses);
     SizeClassInfo *Sci = getSizeClassInfo(ClassId);
     BlockingMutexLock L(&Sci->Mutex);
     TransferBatch *B = Sci->FreeList.front();
     if (B)
       Sci->FreeList.pop_front();
     else {

lib/scudo/standalone/primary64.h

void initLinkerInitialized(s32 ReleaseToOsInterval) {
     }
     ReleaseToOsIntervalMs = ReleaseToOsInterval;
   }
   void init(s32 ReleaseToOsInterval) {
     memset(this, 0, sizeof(*this));
     initLinkerInitialized(ReleaseToOsInterval);
   }
 
+  void unmapTestOnly() {
+    unmap(reinterpret_cast<void *>(PrimaryBase), PrimarySize, UNMAP_ALL, &Data);
+    unmap(reinterpret_cast<void *>(RegionInfoArray),
+          sizeof(RegionInfo) * NumClasses);
+  }
+
   TransferBatch *popBatch(CacheT *C, uptr ClassId) {
     DCHECK_LT(ClassId, NumClasses);
     RegionInfo *Region = getRegionInfo(ClassId);
     BlockingMutexLock L(&Region->Mutex);
     TransferBatch *B = Region->FreeList.front();
     if (B)
       Region->FreeList.pop_front();
     else {

lib/scudo/standalone/tests/bytemap_test.cc

else
       EXPECT_EQ(Map[J], (J % 100) + 1);
   }
 }
 
 TEST(ScudoByteMapTest, FlatByteMap) {
   const scudo::uptr Size = 1U << 10;
   scudo::FlatByteMap<Size> Map;
   testMap(Map, Size);
+  Map.unmapTestOnly();
 }
 
 TEST(ScudoByteMapTest, TwoLevelByteMap) {
   const scudo::uptr Size1 = 1U << 6, Size2 = 1U << 12;
   scudo::TwoLevelByteMap<Size1, Size2> Map;
   testMap(Map, Size1 * Size2);
-  Map.reset();
+  Map.unmapTestOnly();
 }
 
 using TestByteMap = scudo::TwoLevelByteMap<1U << 12, 1U << 13>;
 
 struct TestByteMapParam {
   TestByteMap *Map;
   scudo::uptr Shard;
   scudo::uptr NumberOfShards;
TEST(ScudoByteMapTest, ThreadedTwoLevelByteMap) {
   for (scudo::uptr I = 0; I < NumberOfThreads; I++) {
     P[I].Map = &Map;
     P[I].Shard = I;
     P[I].NumberOfShards = NumberOfThreads;
     pthread_create(&T[I], 0, populateByteMap, &P[I]);
   }
   for (scudo::uptr I = 0; I < NumberOfThreads; I++)
     pthread_join(T[I], 0);
-  Map.reset();
+  Map.unmapTestOnly();
 }

lib/scudo/standalone/tests/primary_test.cc

 #include <thread>
 
 // Note that with small enough regions, the SizeClassAllocator64 also works on
 // 32-bit architectures. It's not something we want to encourage, but we still
 // should ensure the tests pass.
 
 template <typename Primary> static void testPrimary() {
   const scudo::uptr NumberOfAllocations = 32U;
-  std::unique_ptr<Primary> Allocator(new Primary);
+  auto Deleter = [](Primary *P) {
+    P->unmapTestOnly();
+    delete P;
+  };
+  std::unique_ptr<Primary, decltype(Deleter)> Allocator(new Primary, Deleter);
   Allocator->init(/*ReleaseToOsInterval=*/-1);
   typename Primary::CacheT Cache;
   Cache.init(nullptr, Allocator.get());
   for (scudo::uptr I = 0; I <= 16U; I++) {
     const scudo::uptr Size = 1UL << I;
     if (!Primary::canAllocate(Size))
       continue;
     const scudo::uptr ClassId = Primary::SizeClassMap::getClassIdBySize(Size);
TEST(ScudoPrimaryTest, Primary64OOM) {
   while (!Batches.empty()) {
     Allocator.pushBatch(ClassId, Batches.back());
     Batches.pop_back();
   }
   Cache.destroy(nullptr);
   Allocator.releaseToOS();
   Allocator.printStats();
   EXPECT_EQ(AllocationFailed, true);
+  Allocator.unmapTestOnly();
 }
 
 template <typename Primary> static void testIteratePrimary() {
-  std::unique_ptr<Primary> Allocator(new Primary);
+  auto Deleter = [](Primary *P) {
+    P->unmapTestOnly();
+    delete P;
+  };
+  std::unique_ptr<Primary, decltype(Deleter)> Allocator(new Primary, Deleter);
   Allocator->init(/*ReleaseToOsInterval=*/-1);
   typename Primary::CacheT Cache;
   Cache.init(nullptr, Allocator.get());
   std::vector<std::pair<scudo::uptr, void *>> V;
   for (scudo::uptr I = 0; I < 64U; I++) {
     const scudo::uptr Size = std::rand() % Primary::SizeClassMap::MaxSize;
     const scudo::uptr ClassId = Primary::SizeClassMap::getClassIdBySize(Size);
     void *P = Cache.allocate(ClassId);

 }
 
 TEST(ScudoPrimaryTest, PrimaryIterate) {
   using SizeClassMap = scudo::DefaultSizeClassMap;
   testIteratePrimary<scudo::SizeClassAllocator32<SizeClassMap, 18U>>();
   testIteratePrimary<scudo::SizeClassAllocator64<SizeClassMap, 24U>>();
 }
 
-// TODO(kostyak): reenable on 32-bit after implementing unmapTestOnly for the
-//                primary: we are running out of addressable space without.
-#if SCUDO_WORDSIZE == 64U
-
 static std::mutex Mutex;
 static std::condition_variable Cv;
 static bool Ready = false;
 
 template <typename Primary> static void performAllocations(Primary *Allocator) {
   static THREADLOCAL typename Primary::CacheT Cache;
   Cache.init(nullptr, Allocator);
   std::vector<std::pair<scudo::uptr, void *>> V;
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     while (!Ready)
       Cv.wait(Lock);
   }
   for (scudo::uptr I = 0; I < 256U; I++) {
-    const scudo::uptr Size = std::rand() % Primary::SizeClassMap::MaxSize;
+    const scudo::uptr Size = std::rand() % Primary::SizeClassMap::MaxSize / 4;
     const scudo::uptr ClassId = Primary::SizeClassMap::getClassIdBySize(Size);
     void *P = Cache.allocate(ClassId);
     if (P)
       V.push_back(std::make_pair(ClassId, P));
   }
   while (!V.empty()) {
     auto Pair = V.back();
     Cache.deallocate(Pair.first, Pair.second);
     V.pop_back();
   }
   Cache.destroy(nullptr);
 }
 
 template <typename Primary> static void testPrimaryThreaded() {
-  std::unique_ptr<Primary> Allocator(new Primary);
+  auto Deleter = [](Primary *P) {
+    P->unmapTestOnly();
+    delete P;
+  };
+  std::unique_ptr<Primary, decltype(Deleter)> Allocator(new Primary, Deleter);
   Allocator->init(/*ReleaseToOsInterval=*/-1);
   std::thread Threads[32];
   for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++)
     Threads[I] = std::thread(performAllocations<Primary>, Allocator.get());
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     Ready = true;
     Cv.notify_all();
   }
   for (auto &T : Threads)
     T.join();
   Allocator->releaseToOS();
   Allocator->printStats();
 }
 
 TEST(ScudoPrimaryTest, PrimaryThreaded) {
   using SizeClassMap = scudo::SvelteSizeClassMap;
   testPrimaryThreaded<scudo::SizeClassAllocator32<SizeClassMap, 18U>>();
-  testPrimaryThreaded<scudo::SizeClassAllocator64<SizeClassMap, 28U>>();
+  testPrimaryThreaded<scudo::SizeClassAllocator64<SizeClassMap, 24U>>();
 }
-
-#endif // SCUDO_WORDSIZE == 64U

lib/scudo/standalone/tests/tsd_test.cc

public:
 
   void initLinkerInitialized() {
     // This should only be called once by the registry.
     EXPECT_FALSE(Initialized);
     Initialized = true;
   }
   void reset() { memset(this, 0, sizeof(*this)); }
 
+  void unmapTestOnly() { TSDRegistry.unmapTestOnly(); }
   void initCache(CacheT *Cache) { memset(Cache, 0, sizeof(*Cache)); }
   void commitBack(scudo::TSD<MockAllocator> *TSD) {}
   TSDRegistryT *getTSDRegistry() { return &TSDRegistry; }
 
   bool isInitialized() { return Initialized; }
 
 private:
   bool Initialized;

 
 struct ExclusiveCaches {
   template <class Allocator>
   using TSDRegistryT = scudo::TSDRegistryExT<Allocator>;
 };
 
 TEST(ScudoTSDTest, TSDRegistryInit) {
   using AllocatorT = MockAllocator<OneCache>;
-  std::unique_ptr<AllocatorT> Allocator(new AllocatorT);
+  auto Deleter = [](AllocatorT *A) {
+    A->unmapTestOnly();
+    delete A;
+  };
+  std::unique_ptr<AllocatorT, decltype(Deleter)> Allocator(new AllocatorT,
+                                                           Deleter);
   Allocator->reset();
   EXPECT_FALSE(Allocator->isInitialized());
 
   auto Registry = Allocator->getTSDRegistry();
   Registry->initLinkerInitialized(Allocator.get());
   EXPECT_TRUE(Allocator->isInitialized());
 }
 
 template <class AllocatorT> static void testRegistry() {
-  std::unique_ptr<AllocatorT> Allocator(new AllocatorT);
+  auto Deleter = [](AllocatorT *A) {
+    A->unmapTestOnly();
+    delete A;
+  };
+  std::unique_ptr<AllocatorT, decltype(Deleter)> Allocator(new AllocatorT,
+                                                           Deleter);
   Allocator->reset();
   EXPECT_FALSE(Allocator->isInitialized());
 
   auto Registry = Allocator->getTSDRegistry();
   Registry->initThreadMaybe(Allocator.get(), /*MinimalInit=*/true);
   EXPECT_TRUE(Allocator->isInitialized());
 
   bool UnlockRequired;
template <typename AllocatorT> static void stressCache(AllocatorT *Allocator) {
   // Loop a few times to make sure that a concurrent thread isn't modifying it.
   for (scudo::uptr I = 0; I < 4096U; I++)
     EXPECT_EQ(TSD->Cache.Canary, Canary);
   if (UnlockRequired)
     TSD->unlock();
 }
 
 template <class AllocatorT> static void testRegistryThreaded() {
-  std::unique_ptr<AllocatorT> Allocator(new AllocatorT);
+  auto Deleter = [](AllocatorT *A) {
+    A->unmapTestOnly();
+    delete A;
+  };
+  std::unique_ptr<AllocatorT, decltype(Deleter)> Allocator(new AllocatorT,
+                                                           Deleter);
   Allocator->reset();
   std::thread Threads[32];
   for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++)
     Threads[I] = std::thread(stressCache<AllocatorT>, Allocator.get());
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     Ready = true;
     Cv.notify_all();

lib/scudo/standalone/tsd_exclusive.h

void initLinkerInitialized(Allocator *Instance) {
     FallbackTSD->initLinkerInitialized(Instance);
     Initialized = true;
   }
   void init(Allocator *Instance) {
     memset(this, 0, sizeof(*this));
     initLinkerInitialized(Instance);
   }
 
+  void unmapTestOnly() {
+    unmap(reinterpret_cast<void *>(FallbackTSD), sizeof(TSD<Allocator>));
+  }
+
   ALWAYS_INLINE void initThreadMaybe(Allocator *Instance, bool MinimalInit) {
     if (LIKELY(State != ThreadState::NotInitialized))
       return;
     initThread(Instance, MinimalInit);
   }
 
   ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
     if (LIKELY(State == ThreadState::Initialized)) {

lib/scudo/standalone/tsd_shared.h

void initLinkerInitialized(Allocator *Instance) {
     }
     Initialized = true;
   }
   void init(Allocator *Instance) {
     memset(this, 0, sizeof(*this));
     initLinkerInitialized(Instance);
   }
 
+  void unmapTestOnly() {
+    unmap(reinterpret_cast<void *>(TSDs),
+          sizeof(TSD<Allocator>) * NumberOfTSDs);
+  }
+
   ALWAYS_INLINE void initThreadMaybe(Allocator *Instance,
                                      UNUSED bool MinimalInit) {
     if (LIKELY(getCurrentTSD()))
       return;
     initThread(Instance);
   }
 
   ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {