[Runtime] Switch MetadataCache to ConcurrentReadableHashMap. (#34307)

* [Runtime] Switch MetadataCache to ConcurrentReadableHashMap.

Use StableAddressConcurrentReadableHashMap. MetadataCacheEntry's methods for awaiting a particular state assume a stable address, where it will repeatedly examine `this` in a loop while waiting on a condition variable, so we give it a stable address to accommodate that. Some of these caches may be able to tolerate unstable addresses if this code is changed to perform the necessary table lookup each time through the loop instead. Some of them store metadata inline and we assume metadata never moves, so they'll have to stay this way.

* Have StableAddressConcurrentReadableHashMap remember the last found entry and check that before doing a more expensive lookup.

* Make a SmallMutex type that stores the mutex data out of line, and use it to get LockingConcurrentMapStorage to fit into the available space on 32-bit.

rdar://problem/70220660

Author: mikeash

include/swift/Runtime/Concurrent.h

@@ -584,7 +584,8 @@ using llvm::hash_value;
 /// process. It has no destructor, to avoid generating useless global destructor
 /// calls. The memory it allocates can be freed by calling clear() with no
 /// outstanding readers, but this won't destroy the static mutex it uses.
-template <class ElemTy> struct ConcurrentReadableHashMap {
+template <class ElemTy, class MutexTy = StaticMutex>
+struct ConcurrentReadableHashMap {
   // We use memcpy and don't call destructors. Make sure the elements will put
   // up with this.
   static_assert(std::is_trivially_copyable<ElemTy>::value,
@@ -593,6 +594,9 @@ template <class ElemTy> struct ConcurrentReadableHashMap {
                 "Elements must not have destructors (they won't be called).");
 
 private:
+  // A scoped lock type to use on MutexTy.
+  using ScopedLockTy = ScopedLockT<MutexTy, false>;
+
   /// The reciprocal of the load factor at which we expand the table. A value of
   /// 4 means that we resize at 1/4 = 75% load factor.
   static const size_t ResizeProportion = 4;
@@ -752,7 +756,7 @@ template <class ElemTy> struct ConcurrentReadableHashMap {
   std::atomic<IndexStorage *> Indices{nullptr};
 
   /// The writer lock, which must be taken before any mutation of the table.
-  StaticMutex WriterLock;
+  MutexTy WriterLock;
 
   /// The list of pointers to be freed once no readers are active.
   FreeListNode *FreeList{nullptr};
@@ -966,7 +970,7 @@ template <class ElemTy> struct ConcurrentReadableHashMap {
   /// The return value is ignored when `created` is `false`.
   template <class KeyTy, typename Call>
   void getOrInsert(KeyTy key, const Call &call) {
-    StaticScopedLock guard(WriterLock);
+    ScopedLockTy guard(WriterLock);
 
     auto *indices = Indices.load(std::memory_order_relaxed);
     if (!indices)
@@ -1018,7 +1022,7 @@ template <class ElemTy> struct ConcurrentReadableHashMap {
   /// Clear the hash table, freeing (when safe) all memory currently used for
   /// indices and elements.
   void clear() {
-    StaticScopedLock guard(WriterLock);
+    ScopedLockTy guard(WriterLock);
 
     auto *indices = Indices.load(std::memory_order_relaxed);
     auto *elements = Elements.load(std::memory_order_relaxed);
@@ -1054,28 +1058,38 @@ template <class ElemTy> struct HashMapElementWrapper {
 /// by allocating them separately and storing pointers to them. The elements of
 /// the hash table are instances of HashMapElementWrapper. A new getOrInsert
 /// method is provided that directly returns the stable element pointer.
-template <class ElemTy, class Allocator>
+template <class ElemTy, class Allocator, class MutexTy = StaticMutex>
 struct StableAddressConcurrentReadableHashMap
-    : public ConcurrentReadableHashMap<HashMapElementWrapper<ElemTy>> {
+    : public ConcurrentReadableHashMap<HashMapElementWrapper<ElemTy>, MutexTy> {
   // Implicitly trivial destructor.
   ~StableAddressConcurrentReadableHashMap() = default;
 
+  std::atomic<ElemTy *> LastFound{nullptr};
+
   /// Get or insert an element for the given key and arguments. Returns the
   /// pointer to the existing or new element, and a bool indicating whether the
   /// element was created. When false, the element already existed before the
   /// call.
   template <class KeyTy, class... ArgTys>
   std::pair<ElemTy *, bool> getOrInsert(KeyTy key, ArgTys &&...args) {
+    // See if this is a repeat of the previous search.
+    if (auto lastFound = LastFound.load(std::memory_order_acquire))
+      if (lastFound->matchesKey(key))
+        return {lastFound, false};
+
     // Optimize for the case where the value already exists.
-    if (auto wrapper = this->snapshot().find(key))
+    if (auto wrapper = this->snapshot().find(key)) {
+      LastFound.store(wrapper->Ptr, std::memory_order_relaxed);
       return {wrapper->Ptr, false};
+    }
 
     // No such element. Insert if needed. Note: another thread may have inserted
     // it in the meantime, so we still have to handle both cases!
     ElemTy *ptr = nullptr;
     bool outerCreated = false;
-    ConcurrentReadableHashMap<HashMapElementWrapper<ElemTy>>::getOrInsert(
-        key, [&](HashMapElementWrapper<ElemTy> *wrapper, bool created) {
+    ConcurrentReadableHashMap<HashMapElementWrapper<ElemTy>, MutexTy>::
+        getOrInsert(key, [&](HashMapElementWrapper<ElemTy> *wrapper,
+                             bool created) {
           if (created) {
             // Created the indirect entry. Allocate the actual storage.
             size_t allocSize =
@@ -1088,9 +1102,17 @@ struct StableAddressConcurrentReadableHashMap
           outerCreated = created;
           return true; // Keep the new entry.
         });
+    LastFound.store(ptr, std::memory_order_relaxed);
     return {ptr, outerCreated};
   }
 
+  template <class KeyTy> ElemTy *find(const KeyTy &key) {
+    auto result = this->snapshot().find(key);
+    if (!result)
+      return nullptr;
+    return result->Ptr;
+  }
+
 private:
   // Clearing would require deallocating elements, which we don't support.
   void clear() = delete;

include/swift/Runtime/Mutex.h

@@ -772,6 +772,32 @@ class StaticUnsafeMutex {
   MutexHandle Handle;
 };
 
+/// A "small" variant of a Mutex. This allocates the mutex on the heap, for
+/// places where having the mutex inline takes up too much space.
+///
+/// TODO: On OSes that provide a smaller mutex type (e.g. os_unfair_lock on
+/// Darwin), make SmallMutex use that and store it inline, or make Mutex use it
+/// and this can become a typedef there.
+class SmallMutex {
+  SmallMutex(const SmallMutex &) = delete;
+  SmallMutex &operator=(const SmallMutex &) = delete;
+  SmallMutex(SmallMutex &&) = delete;
+  SmallMutex &operator=(SmallMutex &&) = delete;
+
+public:
+  explicit SmallMutex(bool checked = false) { Ptr = new Mutex(checked); }
+  ~SmallMutex() { delete Ptr; }
+
+  void lock() { Ptr->lock(); }
+
+  void unlock() { Ptr->unlock(); }
+
+  bool try_lock() { return Ptr->try_lock(); }
+
+private:
+  Mutex *Ptr;
+};
+
 // Enforce literal requirements for static variants.
 #if SWIFT_MUTEX_SUPPORTS_CONSTEXPR
 static_assert(std::is_literal_type<StaticMutex>::value,