Switch MetadataCache to use a global slab allocator.

cmake/modules/AddSwift.cmake

@@ -329,7 +329,7 @@ function(_add_variant_link_flags)
     RESULT_VAR_NAME result)
 
   if("${LFLAGS_SDK}" STREQUAL "LINUX")
-    list(APPEND result "-lpthread" "-ldl")
+    list(APPEND result "-lpthread" "-latomic" "-ldl")
   elseif("${LFLAGS_SDK}" STREQUAL "FREEBSD")
     list(APPEND result "-lpthread")
   elseif("${LFLAGS_SDK}" STREQUAL "CYGWIN")

cmake/modules/AddSwiftUnittests.cmake

@@ -42,6 +42,9 @@ function(add_swift_unittest test_dirname)
   if("${CMAKE_SYSTEM_NAME}" STREQUAL "Darwin")
     set_property(TARGET "${test_dirname}" APPEND_STRING PROPERTY
       LINK_FLAGS " -Xlinker -rpath -Xlinker ${SWIFT_LIBRARY_OUTPUT_INTDIR}/swift/macosx")
+  elseif("${CMAKE_SYSTEM_NAME}" STREQUAL "Linux")
+    set_property(TARGET "${test_dirname}" APPEND_STRING PROPERTY
+      LINK_FLAGS " -latomic")
   endif()
 
   if(SWIFT_ENABLE_GOLD_LINKER AND

include/swift/Runtime/Concurrent.h

@@ -189,8 +189,9 @@ class ConcurrentMapBase<EntryTy, false, Allocator> : protected Allocator {
     // Destroy the node's payload.
     node->~Node();
 
-    // Deallocate the node.
-    this->Deallocate(node, allocSize);
+    // Deallocate the node.  The static_cast here is required
+    // because LLVM's allocator API is insane.
+    this->Deallocate(static_cast<void*>(node), allocSize);
   }
 };
 

stdlib/public/runtime/Metadata.cpp

@@ -2809,3 +2809,99 @@ swift::swift_getGenericWitnessTable(GenericWitnessTable *genericTable,
 }
 
 uint64_t swift::RelativeDirectPointerNullPtr = 0;
+
+/***************************************************************************/
+/*** Allocator implementation **********************************************/
+/***************************************************************************/
+
+namespace {
+  struct PoolRange {
+    static constexpr uintptr_t PageSize = 16 * 1024;
+    static constexpr uintptr_t MaxPoolAllocationSize = PageSize / 2;
+
+    /// The start of the allocation.
+    char *Begin;
+
+    /// The number of bytes remaining.
+    size_t Remaining;
+  };
+}
+
+// A statically-allocated pool.  It's zero-initialized, so this
+// doesn't cost us anything in binary size.
+LLVM_ALIGNAS(alignof(void*)) static char InitialAllocationPool[64*1024];
+static std::atomic<PoolRange>
+AllocationPool{PoolRange{InitialAllocationPool,
+                         sizeof(InitialAllocationPool)}};
+
+void *MetadataAllocator::Allocate(size_t size, size_t alignment) {
+  assert(alignment <= alignof(void*));
+  assert(size % alignof(void*) == 0);
+
+  // If the size is larger than the maximum, just use malloc.
+  if (size > PoolRange::MaxPoolAllocationSize)
+    return malloc(size);
+
+  // Allocate out of the pool.
+  PoolRange curState = AllocationPool.load(std::memory_order_relaxed);
+  while (true) {
+    char *allocation;
+    PoolRange newState;
+    bool allocatedNewPage;
+
+    // Try to allocate out of the current page.
+    if (size <= curState.Remaining) {
+      allocatedNewPage = false;
+      allocation = curState.Begin;
+      newState = PoolRange{curState.Begin + size, curState.Remaining - size};
+    } else {
+      allocatedNewPage = true;
+      allocation = new char[PoolRange::PageSize];
+      newState = PoolRange{allocation + size, PoolRange::PageSize - size};
+      __asan_poison_memory_region(allocation, PoolRange::PageSize);
+    }
+
+    // Swap in the new state.
+    if (std::atomic_compare_exchange_weak_explicit(&AllocationPool,
+                                                   &curState, newState,
+                                              std::memory_order_relaxed,
+                                              std::memory_order_relaxed)) {
+      // If that succeeded, we've successfully allocated.
+      __msan_allocated_memory(allocation, size);
+      __asan_poison_memory_region(allocation, size);
+      return allocation;
+    }
+
+    // If it failed, go back to a neutral state and try again.
+    if (allocatedNewPage) {
+      delete[] allocation;
+    }
+  }
+}
+
+void MetadataAllocator::Deallocate(const void *allocation, size_t size) {
+  __asan_poison_memory_region(allocation, size);
+
+  if (size > PoolRange::MaxPoolAllocationSize) {
+    free(const_cast<void*>(allocation));
+    return;
+  }
+
+  // Check whether the allocation pool is still in the state it was in
+  // immediately after the given allocation.
+  PoolRange curState = AllocationPool.load(std::memory_order_relaxed);
+  if (reinterpret_cast<const char*>(allocation) + size != curState.Begin) {
+    return;
+  }
+
+  // Try to swap back to the pre-allocation state.  If this fails,
+  // don't bother trying again; we'll just leak the allocation.
+  PoolRange newState = { reinterpret_cast<char*>(const_cast<void*>(allocation)),
+                         curState.Remaining + size };
+  (void)
+    std::atomic_compare_exchange_strong_explicit(&AllocationPool,
+                                                 &curState, newState,
+                                                 std::memory_order_relaxed,
+                                                 std::memory_order_relaxed);
+}
+

stdlib/public/runtime/MetadataCache.h

@@ -26,11 +26,18 @@
 
 namespace swift {
 
-// For now, use malloc and free as our standard allocator for
-// metadata caches.  It might make sense in the future to take
-// advantage of the fact that we know that most allocations here
-// won't ever be deallocated.
-using MetadataAllocator = llvm::MallocAllocator;
+class MetadataAllocator : public llvm::AllocatorBase<MetadataAllocator> {
+public:
+  void Reset() {}
+
+  LLVM_ATTRIBUTE_RETURNS_NONNULL void *Allocate(size_t size, size_t alignment);
+  using AllocatorBase<MetadataAllocator>::Allocate;
+
+  void Deallocate(const void *Ptr, size_t size);
+  using AllocatorBase<MetadataAllocator>::Deallocate;
+
+  void PrintStats() const {}
+};
 
 /// A typedef for simple global caches.
 template <class EntryTy>

utils/gen-static-stdlib-link-args

@@ -62,6 +62,7 @@ function write_linkfile {
 -ldl
 -lpthread
 -lswiftCore
+-latomic
 -lswiftImageInspectionShared
 $ICU_LIBS
 -Xlinker

utils/static-executable-args.lnk

@@ -8,6 +8,7 @@
 -Xlinker
 --defsym=__import_pthread_key_create=pthread_key_create
 -lpthread
+-latomic
 -licui18n
 -licuuc
 -licudata