Merge pull request #42037 from jckarter/async-let-work-around-allocator-bug

IRGen: Artificially pad async let entry point contexts.

Author: jckarter

lib/IRGen/GenCall.cpp

@@ -3903,14 +3903,17 @@ void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
                                    unsigned asyncContextIndex) {
   auto &IGM = IGF.IGM;
   auto size = layout.getSize();
+  auto asyncFuncPointerVar = cast<llvm::GlobalVariable>(IGM.getAddrOfAsyncFunctionPointer(asyncFunction));
   auto asyncFuncPointer = IGF.Builder.CreateBitOrPointerCast(
-      IGM.getAddrOfAsyncFunctionPointer(asyncFunction), IGM.Int8PtrTy);
+                                           asyncFuncPointerVar, IGM.Int8PtrTy);
   auto *id = IGF.Builder.CreateIntrinsicCall(
       llvm::Intrinsic::coro_id_async,
       {llvm::ConstantInt::get(IGM.Int32Ty, size.getValue()),
        llvm::ConstantInt::get(IGM.Int32Ty, 16),
        llvm::ConstantInt::get(IGM.Int32Ty, asyncContextIndex),
        asyncFuncPointer});
+  auto inserted = IGM.AsyncCoroIDs.insert({asyncFuncPointerVar, id});
+  assert(inserted.second);
   // Call 'llvm.coro.begin', just for consistency with the normal pattern.
   // This serves as a handle that we can pass around to other intrinsics.
   auto hdl = IGF.Builder.CreateIntrinsicCall(

lib/IRGen/GenConcurrency.cpp

@@ -26,6 +26,7 @@
 #include "IRGenModule.h"
 #include "LoadableTypeInfo.h"
 #include "ScalarPairTypeInfo.h"
+#include "swift/AST/ASTContext.h"
 #include "swift/AST/ProtocolConformanceRef.h"
 #include "swift/ABI/MetadataValues.h"
 
@@ -192,6 +193,44 @@ llvm::Value *irgen::emitBuiltinStartAsyncLet(IRGenFunction &IGF,
   auto futureResultType = subs.getReplacementTypes()[0]->getCanonicalType();
   auto futureResultTypeMetadata = IGF.emitAbstractTypeMetadataRef(futureResultType);
 
+  // The concurrency runtime for older Apple OSes has a bug in task formation
+  // for `async let`s that may manifest when trying to use room in the
+  // parent task's preallocated `async let` buffer for the child task's
+  // initial task allocator slab. If targeting those older OSes, pad the
+  // context size for async let entry points to never fit in the preallocated
+  // space, so that we don't run into that bug. We leave a note on the
+  // declaration so that coroutine splitting can pad out the final context
+  // size after splitting.
+  auto deploymentAvailability
+    = AvailabilityContext::forDeploymentTarget(IGF.IGM.Context);
+  if (!deploymentAvailability.isContainedIn(
+                                   IGF.IGM.Context.getSwift57Availability())) {
+    auto taskAsyncFunctionPointer
+                = cast<llvm::GlobalVariable>(taskFunction->stripPointerCasts());
+
+    auto taskAsyncIDIter = IGF.IGM.AsyncCoroIDs.find(taskAsyncFunctionPointer);
+    assert(taskAsyncIDIter != IGF.IGM.AsyncCoroIDs.end()
+           && "async let entry point not emitted locally");
+    auto taskAsyncID = taskAsyncIDIter->second;
+    
+    // Pad out the initial context size in the async function pointer record
+    // and ID intrinsic so that it will never fit in the preallocated space.
+    uint64_t origSize = cast<llvm::ConstantInt>(taskAsyncID->getArgOperand(0))
+      ->getValue().getLimitedValue();
+    
+    uint64_t paddedSize = std::max(origSize,
+                     (NumWords_AsyncLet * IGF.IGM.getPointerSize()).getValue());
+    auto paddedSizeVal = llvm::ConstantInt::get(IGF.IGM.Int32Ty, paddedSize);
+    taskAsyncID->setArgOperand(0, paddedSizeVal);
+    
+    auto origInit = taskAsyncFunctionPointer->getInitializer();
+    auto newInit = llvm::ConstantStruct::get(
+                                   cast<llvm::StructType>(origInit->getType()),
+                                   origInit->getAggregateElement(0u),
+                                   paddedSizeVal);
+    taskAsyncFunctionPointer->setInitializer(newInit);
+  }
+  
   llvm::CallInst *call;
   if (localResultBuffer) {
     // This is @_silgen_name("swift_asyncLet_begin")

lib/IRGen/IRGenModule.h

@@ -328,10 +328,10 @@ class IRGenerator {
 
   /// The queue of IRGenModules for multi-threaded compilation.
   SmallVector<IRGenModule *, 8> Queue;
-
+  
   std::atomic<int> QueueIndex;
 
-  friend class CurrentIGMPtr;  
+  friend class CurrentIGMPtr;
 public:
   explicit IRGenerator(const IRGenOptions &opts, SILModule &module);
 
@@ -893,6 +893,10 @@ class IRGenModule {
   std::string GetObjCSectionName(StringRef Section, StringRef MachOAttributes);
   void SetCStringLiteralSection(llvm::GlobalVariable *GV, ObjCLabelType Type);
 
+  // Mapping of AsyncFunctionPointer records to their corresponding
+  // `@llvm.coro.id.async` intrinsic tag in the function implementation.
+  llvm::DenseMap<llvm::GlobalVariable*, llvm::CallInst*> AsyncCoroIDs;
+
 private:
   Size PtrSize;
   Size AtomicBoolSize;

test/IRGen/async_let_back_deploy_workaround.swift

@@ -0,0 +1,22 @@
+// RUN: %target-swift-frontend -emit-ir -target x86_64-apple-macos99.99 %s | %FileCheck --check-prefix=CHECK --check-prefix=CHECK-sans-workaround %s
+// RUN: %target-swift-frontend -emit-ir -target x86_64-apple-macos12.3 %s | %FileCheck --check-prefix=CHECK --check-prefix=CHECK-with-workaround %s
+
+// REQUIRES: OS=macosx
+ 
+// rdar://90506708: Prior to Swift 5.7, the Swift concurrency runtime had a bug
+// that led to memory corruption in cases when an `async let` child task
+// would try to use the last 16 bytes of the preallocated slab from the parent
+// to seed its own task allocator. When targeting older Apple OSes that shipped
+// with this bug in their runtime, we work around the bug by inflating the
+// initial context sizes of any async functions used as `async let` entry points
+// to ensure that the preallocated space is never used for the initial context.
+
+// CHECK: [[ASYNC_LET_ENTRY:@"\$sSiIeghHd_Sis5Error_pIegHrzo_TRTATu"]]
+// CHECK-with-workaround-SAME: = internal {{.*}} %swift.async_func_pointer <{ {{.*}}, i32 6{{[0-9][0-9]}} }>
+// CHECK-sans-workaround-SAME: = internal {{.*}} %swift.async_func_pointer <{ {{.*}}, i32 {{[0-9][0-9]}} }>
+
+// CHECK: swift_asyncLet_begin{{.*}}[[ASYNC_LET_ENTRY]]
+public func foo(x: Int, y: Int) async -> Int {
+    async let z = x + y
+    return await z
+}