[compiler-rt] [scudo] Add EnableMultiRegions mode (PR #98076)

[via llvm-commits]

================
@@ -555,18 +715,208 @@ template <typename Config> class SizeClassAllocator64 {
     bool isPopulatingFreeList GUARDED_BY(FLLock) = false;
   };
   struct RegionInfo : UnpaddedRegionInfo {
+    // This is only used when `Config::getEnableMultiRegions` is enabled and is
+    // guarded by the mutex in `RegionInfoManager`.
+    RegionInfo *Next = nullptr;
     char Padding[SCUDO_CACHE_LINE_SIZE -
-                 (sizeof(UnpaddedRegionInfo) % SCUDO_CACHE_LINE_SIZE)] = {};
+                 ((sizeof(UnpaddedRegionInfo) + sizeof(RegionInfo *)) %
+                  SCUDO_CACHE_LINE_SIZE)] = {};
   };
   static_assert(sizeof(RegionInfo) % SCUDO_CACHE_LINE_SIZE == 0, "");
 
-  RegionInfo *getRegionInfo(uptr ClassId) {
-    DCHECK_LT(ClassId, NumClasses);
-    return &RegionInfoArray[ClassId];
-  }
+  template <bool IsMultiRegions = false> struct RegionInfoAlloc {
+    RegionInfo *allocate() {
+      UNREACHABLE("RegionInfo is statically allocated");
+    }
 
-  uptr getRegionBaseByClassId(uptr ClassId) {
-    RegionInfo *Region = getRegionInfo(ClassId);
+    void verifyTheNumberOfAllocatedRegionInfo(uptr NumRegionInfo) {
+      DCHECK_EQ(NumRegionInfo, NumClasses);
+    }
+  };
+
+  template <> struct RegionInfoAlloc</*isMultiRegions=*/true> {
+    RegionInfo *allocate() {
+      ScopedLock L(M);
+      return S.pop();
+    }
+
+    void verifyTheNumberOfAllocatedRegionInfo(uptr NumRegionInfo) {
+      ScopedLock L(M);
+      DCHECK_EQ(NumRegionInfo, S.Size);
+    }
+
+    HybridMutex M;
+    // According to the following,
+    //   DR1351: If the brace-or-equal-initializer of a non-static data
+    //   member invokes a defaulted default constructor of its class or of an
+    //   enclosing class in a potentially evaluated subexpression, the program
+    //   is ill-formed.
+    // So we have to `outline` the `Size`/`Array` into another struct `Storage`.
+    struct Storage {
+      RegionInfo *pop() {
+        if (Size == NumEntries)
+          return nullptr;
+        return &Array[Size++];
+      }
+      // The amount memory used by this allocator is about (NumEntries *
+      // RegionSize). For example, region with size 256 KB will have 2GB space
+      // available.
+      // TODO(chiahungduan): Consider having this configurable.
+      static constexpr uptr NumEntries = 1UL << 13;
+      uptr Size = 0;
+      alignas(SCUDO_CACHE_LINE_SIZE) RegionInfo Array[NumEntries];
+    } S GUARDED_BY(M);
+  };
+
+  template <bool IsMultiRegions = false> struct RegionInfoInterface {
+    struct RegionInfoIter {
+      RegionInfoIter(RegionInfo *Region) : CurRegionInfo(Region) {}
+      RegionInfo *operator->() { return CurRegionInfo; }
+      RegionInfoIter &operator++() {
+        CurRegionInfo = nullptr;
+        return *this;
+      }
+      RegionInfo *get() { return CurRegionInfo; }
+      bool end() { return CurRegionInfo == nullptr; }
+      RegionInfo *CurRegionInfo = nullptr;
+    };
+
+    void init(UNUSED RegionInfoAlloc<IsMultiRegions> &Allocator) {
+      // The RegionInfo storage is statically initialized.
+    }
+
+    ALWAYS_INLINE RegionInfo *getCurRegionInfo(uptr ClassId) {
+      DCHECK_LT(ClassId, NumClasses);
+      return &RegionInfoArray[ClassId];
+    }
+    ALWAYS_INLINE RegionInfoIter getRegionInfoIter(uptr ClassId) {
+      return RegionInfoIter(getCurRegionInfo(ClassId));
+    }
+
+    void pushRegionInfo(UNUSED RegionInfo *Region, UNUSED uptr ClassId) {
+      UNREACHABLE("Only MultiRegions supports this operation\n");
+    }
+    // TODO: Review the uses of `getRawRegionInfoArray` and
+    // `getRawRegionInfoArraySize` and see if we can deprecate them.
+    RegionInfo *getRawRegionInfoArray() { return RegionInfoArray; }
+    uptr getRawRegionInfoArraySize() {
+      return static_cast<uptr>(sizeof(RegionInfoArray));
+    }
+    void shuffle(u32 *Seed) {
+      scudo::shuffle(RegionInfoArray, NumClasses, Seed);
+    }
+
+    alignas(SCUDO_CACHE_LINE_SIZE) RegionInfo RegionInfoArray[NumClasses];
+  };
+
+  template <> struct RegionInfoInterface</*isMultiRegions=*/true> {
+    struct RegionInfoIter {
+      RegionInfoIter(RegionInfo *Region, HybridMutex &RegionInfoListLock)
+          : CurRegionInfo(Region), M(RegionInfoListLock) {}
+      RegionInfo *operator->() { return CurRegionInfo; }
+      RegionInfoIter &operator++() {
+        ScopedLock L(M);
+        CurRegionInfo = CurRegionInfo->Next;
+        return *this;
+      }
+      RegionInfo *get() { return CurRegionInfo; }
+      bool end() { return CurRegionInfo == nullptr; }
+      RegionInfo *CurRegionInfo = nullptr;
+      HybridMutex &M;
+    };
+
+    void init(RegionInfoAlloc</*isMultiRegions=*/true> &Allocator) {
+      for (uptr I = 0; I < NumClasses; I++) {
+        RegionInfo *Region = Allocator.allocate();
+        LowestAddrRegionInfo[I].P = Region;
+        CurrentRegionInfo[I].P = Region;
+      }
+    }
+
+    // Return the last pushed RegionInfo. For one size class, the current
+    // RegionInfo is responsible for the page mapping and the other RegionInfos
+    // will have been exhausted already.
+    ALWAYS_INLINE RegionInfo *getCurRegionInfo(uptr ClassId) {
+      DCHECK_LT(ClassId, NumClasses);
+      return CurrentRegionInfo[ClassId].P;
+    }
+
+    ALWAYS_INLINE RegionInfoIter getRegionInfoIter(uptr ClassId) {
+      return RegionInfoIter(LowestAddrRegionInfo[ClassId].P,
+                            RegionInfoLock[ClassId]);
+    }
+
+    // RegionInfos for the same size class will be stored in the order of base
+    // address. Which means every RegionInfo visiting will be starting from
+    // lowest address and which aligns with how pointer grouping works.
+    void pushRegionInfo(RegionInfo *Region, uptr ClassId)
+        REQUIRES(Region->MMLock) {
+      DCHECK_LT(ClassId, NumClasses);
+      DCHECK(Region->MemMapInfo.MemMap.isAllocated());
+
+      // The creation of new region requires holding the MMLock of current
+      // region to ensure only one thread is allocating the new region.
+      CurrentRegionInfo[ClassId].P->MMLock.assertHeld();
+
+      RegionInfo *RegionCursor = LowestAddrRegionInfo[ClassId].P;
+      DCHECK_NE(RegionCursor, nullptr);
+
+      ScopedLock L(RegionInfoLock[ClassId]);
+
+      if (Region->RegionBeg < RegionCursor->RegionBeg) {
+        Region->Next = RegionCursor;
+        LowestAddrRegionInfo[ClassId].P = Region;
+      } else {
+        while (RegionCursor->Next != nullptr &&
+               Region->RegionBeg > RegionCursor->Next->RegionBeg) {
+          RegionCursor = RegionCursor->Next;
+        }
+
+        Region->Next = RegionCursor->Next;
+        RegionCursor->Next = Region;
+      }
+
+      if (SCUDO_DEBUG) {
+        RegionInfo *R = LowestAddrRegionInfo[ClassId].P;
+        while (R->Next != nullptr) {
+          DCHECK_LT(R->RegionBeg, R->Next->RegionBeg);
+          R = R->Next;
+        }
+      }
+
+      CurrentRegionInfo[ClassId].P = Region;
+    }
+    // Multiple-Regions doesn't support this.
+    RegionInfo *getRawRegionInfoArray() { return nullptr; }
+    uptr getRawRegionInfoArraySize() { return 0; }
+
+    void shuffle(u32 *Seed) {
+      if (SCUDO_DEBUG) {
+        // We don't support shuffling two arrays with same randomness. This is
+        // supposed to be done at the initialization stage so that we can simply
+        // update the `LowestAddrRegionInfo` by copying the `CurrentRegionInfo`.
+        for (uptr I = 0; I < NumClasses; ++I)
+          CHECK_EQ(CurrentRegionInfo[I].P, LowestAddrRegionInfo[I].P);
+      }
+      scudo::shuffle(CurrentRegionInfo, NumClasses, Seed);
+      memcpy(CurrentRegionInfo, LowestAddrRegionInfo,
+             sizeof(RegionInfoPointer) * NumClasses);
+    }
+
+    // Scudo requires the data member constant initializable. Array of raw
+    // pointers doesn't meet the condition. Therefore, wrap the pointer in the
----------------
ChiaHungDuan wrote:

Done.

https://github.com/llvm/llvm-project/pull/98076