[Concurrency] Add environment variable for disabling async stack slab allocator.

include/swift/Runtime/EnvironmentVariables.h

@@ -14,6 +14,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#ifndef SWIFT_RUNTIME_ENVIRONMENTVARIABLES_H
+#define SWIFT_RUNTIME_ENVIRONMENTVARIABLES_H
+
 #include "swift/Threading/Once.h"
 #include "swift/shims/Visibility.h"
 
@@ -63,6 +66,12 @@ SWIFT_RUNTIME_STDLIB_SPI bool concurrencyValidateUncheckedContinuations();
 // Concurrency library can call.
 SWIFT_RUNTIME_STDLIB_SPI const char *concurrencyIsCurrentExecutorLegacyModeOverride();
 
+// Wrapper around SWIFT_DEBUG_ENABLE_TASK_SLAB_ALLOCATOR that the Concurrency
+// library can call.
+SWIFT_RUNTIME_STDLIB_SPI bool concurrencyEnableTaskSlabAllocator();
+
 } // end namespace environment
 } // end namespace runtime
 } // end namespace swift
+
+#endif // SWIFT_RUNTIME_ENVIRONMENTVARIABLES_H

stdlib/public/Concurrency/TaskAlloc.cpp

@@ -84,3 +84,8 @@ void swift::swift_job_deallocate(Job *job, void *ptr) {
 
   allocator(static_cast<AsyncTask *>(job)).dealloc(ptr);
 }
+
+#if !SWIFT_CONCURRENCY_EMBEDDED
+std::atomic<TaskAllocatorConfiguration::EnableState>
+    TaskAllocatorConfiguration::enableState = {EnableState::Uninitialized};
+#endif

stdlib/public/Concurrency/TaskPrivate.h

@@ -26,6 +26,7 @@
 #include "swift/Runtime/Atomic.h"
 #include "swift/Runtime/Concurrency.h"
 #include "swift/Runtime/DispatchShims.h"
+#include "swift/Runtime/EnvironmentVariables.h"
 #include "swift/Runtime/Error.h"
 #include "swift/Runtime/Exclusivity.h"
 #include "swift/Runtime/HeapObject.h"
@@ -744,14 +745,48 @@ class alignas(2 * sizeof(void*)) ActiveTaskStatus {
 static_assert(sizeof(ActiveTaskStatus) == ACTIVE_TASK_STATUS_SIZE,
   "ActiveTaskStatus is of incorrect size");
 
+struct TaskAllocatorConfiguration {
+#if SWIFT_CONCURRENCY_EMBEDDED
+
+  // Slab allocator is always enabled on embedded.
+  bool enableSlabAllocator() { return true; }
+
+#else
+
+  enum class EnableState : uint8_t {
+    Uninitialized,
+    Enabled,
+    Disabled,
+  };
+
+  static std::atomic<EnableState> enableState;
+
+  bool enableSlabAllocator() {
+    auto state = enableState.load(std::memory_order_relaxed);
+    if (SWIFT_UNLIKELY(state == EnableState::Uninitialized)) {
+      state = runtime::environment::concurrencyEnableTaskSlabAllocator()
+                  ? EnableState::Enabled
+                  : EnableState::Disabled;
+      enableState.store(state, std::memory_order_relaxed);
+    }
+
+    return SWIFT_UNLIKELY(state == EnableState::Enabled);
+  }
+
+#endif // SWIFT_CONCURRENCY_EMBEDDED
+};
+
 /// The size of an allocator slab. We want the full allocation to fit into a
 /// 1024-byte malloc quantum. We subtract off the slab header size, plus a
 /// little extra to stay within our limits even when there's overhead from
 /// malloc stack logging.
-static constexpr size_t SlabCapacity = 1024 - StackAllocator<0, nullptr>::slabHeaderSize() - 8;
+static constexpr size_t SlabCapacity =
+    1024 - 8 -
+    StackAllocator<0, nullptr, TaskAllocatorConfiguration>::slabHeaderSize();
 extern Metadata TaskAllocatorSlabMetadata;
 
-using TaskAllocator = StackAllocator<SlabCapacity, &TaskAllocatorSlabMetadata>;
+using TaskAllocator = StackAllocator<SlabCapacity, &TaskAllocatorSlabMetadata,
+                                     TaskAllocatorConfiguration>;
 
 /// Private storage in an AsyncTask object.
 struct AsyncTask::PrivateStorage {

stdlib/public/runtime/EnvironmentVariables.cpp

@@ -319,3 +319,7 @@ SWIFT_RUNTIME_STDLIB_SPI bool concurrencyValidateUncheckedContinuations() {
 SWIFT_RUNTIME_STDLIB_SPI const char *concurrencyIsCurrentExecutorLegacyModeOverride() {
   return runtime::environment::SWIFT_IS_CURRENT_EXECUTOR_LEGACY_MODE_OVERRIDE();
 }
+
+SWIFT_RUNTIME_STDLIB_SPI bool concurrencyEnableTaskSlabAllocator() {
+  return runtime::environment::SWIFT_DEBUG_ENABLE_TASK_SLAB_ALLOCATOR();
+}

stdlib/public/runtime/EnvironmentVariables.def

@@ -147,4 +147,9 @@ VARIABLE(SWIFT_DUMP_ACCESSIBLE_FUNCTIONS, bool, false,
          "record names, e.g. by the Distributed runtime to invoke distributed "
          "functions.")
 
+VARIABLE(SWIFT_DEBUG_ENABLE_TASK_SLAB_ALLOCATOR, bool, true,
+         "Use a slab allocator for the async stack. If not enabled, directly "
+         "call malloc/free. Direct use of malloc/free enables use of memory "
+         "debugging tools for async stack allocations.")
+
 #undef VARIABLE

stdlib/public/runtime/StackAllocator.h

@@ -57,7 +57,15 @@ namespace swift {
 ///
 /// SlabMetadataPtr specifies a fake metadata pointer to place at the beginning
 /// of slab allocations, so analysis tools can identify them.
-template <size_t SlabCapacity, Metadata *SlabMetadataPtr>
+///
+/// SlabAllocatorConfiguration allows customizing behavior. It currently allows
+/// one customization point: bool enableSlabAllocator():
+///   returns false - the stack allocator directly calls malloc/free
+///   returns true - the slab allocator is used
+/// This function MUST return the same value throughout the lifetime the stack
+/// allocator.
+template <size_t SlabCapacity, Metadata *SlabMetadataPtr,
+          typename SlabAllocatorConfiguration>
 class StackAllocator {
 private:
 
@@ -77,6 +85,8 @@ class StackAllocator {
   /// Used for unit testing.
   uint32_t numAllocatedSlabs:31;
 
+  /// The configuration object.
+  [[no_unique_address]] SlabAllocatorConfiguration configuration;
 
   /// The minimal alignment of allocated memory.
   static constexpr size_t alignment = MaximumAlignment;
@@ -227,9 +237,10 @@ class StackAllocator {
   };
 
   // Return a slab which is suitable to allocate \p size memory.
+  SWIFT_ALWAYS_INLINE
   Slab *getSlabForAllocation(size_t size) {
     Slab *slab = (lastAllocation ? lastAllocation->slab : firstSlab);
-    if (slab) {
+    if (SWIFT_LIKELY(slab)) {
       // Is there enough space in the current slab?
       if (slab->canAllocate(size))
         return slab;
@@ -249,6 +260,12 @@ class StackAllocator {
         size = std::max(size, alreadyAllocatedCapacity);
       }
     }
+
+    // This is only checked on the path that allocates a new slab, to minimize
+    // overhead when the slab allocator is enabled.
+    if (SWIFT_UNLIKELY(!configuration.enableSlabAllocator()))
+      return nullptr;
+
     size_t capacity = std::max(SlabCapacity,
                                Allocation::includingHeader(size));
     void *slabBuffer = malloc(Slab::includingHeader(capacity));
@@ -281,13 +298,15 @@ class StackAllocator {
   /// Construct a StackAllocator without a pre-allocated first slab.
   StackAllocator()
       : firstSlab(nullptr), firstSlabIsPreallocated(false),
-        numAllocatedSlabs(0) {}
+        numAllocatedSlabs(0), configuration() {}
 
   /// Construct a StackAllocator with a pre-allocated first slab.
   StackAllocator(void *firstSlabBuffer, size_t bufferCapacity) : StackAllocator() {
     // If the pre-allocated buffer can't hold a slab header, ignore it.
     if (bufferCapacity <= Slab::headerSize())
       return;
+    if (SWIFT_UNLIKELY(!configuration.enableSlabAllocator()))
+      return;
     char *start = (char *)llvm::alignAddr(firstSlabBuffer,
                                           llvm::Align(alignment));
     char *end = (char *)firstSlabBuffer + bufferCapacity;
@@ -317,6 +336,17 @@ class StackAllocator {
       size += sizeof(uintptr_t);
     size_t alignedSize = llvm::alignTo(size, llvm::Align(alignment));
     Slab *slab = getSlabForAllocation(alignedSize);
+
+    // If getSlabForAllocation returns null, that means that the slab allocator
+    // is disabled, and we should directly call malloc. We get this info
+    // indirectly rather than directly calling enableSlabAllocator() in order
+    // to minimize the overhead in the case where the slab allocator is enabled.
+    // When getSlabForAllocation gets inlined into this code, this ends up
+    // getting folded into its enableSlabAllocator() call, and the fast path
+    // where `slab` is non-null ends up with no extra conditionals at all.
+    if (SWIFT_UNLIKELY(!slab))
+      return malloc(size);
+
     Allocation *allocation = slab->allocate(alignedSize, lastAllocation);
     lastAllocation = allocation;
     assert(llvm::isAddrAligned(llvm::Align(alignment),
@@ -326,7 +356,10 @@ class StackAllocator {
 
   /// Deallocate memory \p ptr.
   void dealloc(void *ptr) {
-    if (!lastAllocation || lastAllocation->getAllocatedMemory() != ptr) {
+    if (SWIFT_UNLIKELY(!lastAllocation ||
+                       lastAllocation->getAllocatedMemory() != ptr)) {
+      if (!configuration.enableSlabAllocator())
+        return free(ptr);
       SWIFT_FATAL_ERROR(0, "freed pointer was not the last allocation");
     }
 

stdlib/toolchain/Compatibility56/include/Concurrency/TaskPrivate.h

@@ -25,7 +25,7 @@
 #include "swift/Runtime/Error.h"
 
 #define SWIFT_FATAL_ERROR swift_Concurrency_fatalError
-#include "public/runtime/StackAllocator.h"
+#include "Runtime/StackAllocator.h"
 
 namespace swift {