concurrency: make the startAsyncLet closure no-escaping

The closure does not escape the startAsyncLet - endAsyncLet scope. Even though it's (potentially) running on a different thread.

The substantial change in the runtime is to not call swift_release on the closure context if it's a non-escaping closure.

Author: eeckstein

include/swift/ABI/Task.h

@@ -582,13 +582,13 @@ class FutureAsyncContext : public AsyncContext {
 /// This matches the ABI of a closure `() async throws -> ()`
 using AsyncVoidClosureEntryPoint =
   SWIFT_CC(swiftasync)
-  void (SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT HeapObject *);
+  void (SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT void *);
 
 /// This matches the ABI of a closure `<T>() async throws -> T`
 using AsyncGenericClosureEntryPoint =
     SWIFT_CC(swiftasync)
     void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT HeapObject *);
+         SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT void *);
 
 /// This matches the ABI of the resume function of a closure
 ///  `() async throws -> ()`.
@@ -602,7 +602,7 @@ class AsyncContextPrefix {
   // passing the closure context instead of via the async context)
   AsyncVoidClosureEntryPoint *__ptrauth_swift_task_resume_function
       asyncEntryPoint;
-   HeapObject *closureContext;
+  void *closureContext;
   SwiftError *errorResult;
 };
 
@@ -615,7 +615,7 @@ class FutureAsyncContextPrefix {
   // passing the closure context instead of via the async context)
   AsyncGenericClosureEntryPoint *__ptrauth_swift_task_resume_function
       asyncEntryPoint;
-  HeapObject *closureContext;
+  void *closureContext;
   SwiftError *errorResult;
 };
 

include/swift/Runtime/Concurrency.h

@@ -281,15 +281,15 @@ bool swift_taskGroup_isEmpty(TaskGroup *group);
 /// \code
 /// func swift_asyncLet_start<T>(
 ///     _ alet: Builtin.RawPointer,
-///     operation: __owned @Sendable @escaping () async throws -> T
+///     operation: __owned @Sendable () async throws -> T
 /// )
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_asyncLet_start(
     AsyncLet *alet,
     const Metadata *futureResultType,
     void *closureEntryPoint,
-    HeapObject * /* +1 */ closureContext);
+    void *closureContext);
 
 /// This matches the ABI of a closure `<T>(Builtin.RawPointer) async -> T`
 using AsyncLetWaitSignature =

include/swift/Runtime/RuntimeFunctions.def

@@ -1682,7 +1682,7 @@ FUNCTION(DefaultActorDeallocateResilient,
 ///     AsyncLet *alet,
 ///     const Metadata *futureResultType,
 ///     void *closureEntryPoint,
-///     HeapObject *closureContext
+///     OpaqueValue *closureContext
 /// );
 FUNCTION(AsyncLetStart,
          swift_asyncLet_start, SwiftCC,
@@ -1691,7 +1691,7 @@ FUNCTION(AsyncLetStart,
          ARGS(SwiftAsyncLetPtrTy, // AsyncLet*, alias for Int8PtrTy
               TypeMetadataPtrTy,  // futureResultType
               Int8PtrTy,          // closureEntry
-              RefCountedPtrTy,    // closureContext
+              OpaquePtrTy,        // closureContext
          ),
          ATTRS(NoUnwind, ArgMemOnly))
 

lib/AST/Builtins.cpp

@@ -1496,7 +1496,7 @@ static ValueDecl *getStartAsyncLet(ASTContext &ctx, Identifier id) {
   builder.addParameter(makeConcrete(OptionalType::get(ctx.TheRawPointerType)));
 
   // operation async function pointer: () async throws -> T
-  auto extInfo = ASTExtInfoBuilder().withAsync().withThrows().build();
+  auto extInfo = ASTExtInfoBuilder().withAsync().withThrows().withNoEscape().build();
   builder.addParameter(
       makeConcrete(FunctionType::get({ }, genericParam, extInfo)));
 

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -100,9 +100,10 @@ OVERRIDE_TASK(task_create_group_future_common, AsyncTaskAndContext, , , ,
               (JobFlags flags, TaskGroup *group,
                const Metadata *futureResultType,
                FutureAsyncSignature::FunctionType *function,
-               HeapObject */* +1 */ closureContext, size_t initialContextSize),
+               void *closureContext, bool owningClosureContext,
+               size_t initialContextSize),
               (flags, group, futureResultType, function, closureContext,
-               initialContextSize))
+               owningClosureContext, initialContextSize))
 
 OVERRIDE_TASK(task_future_wait, void, SWIFT_EXPORT_FROM(swift_Concurrency),
               SWIFT_CC(swiftasync), swift::,
@@ -155,7 +156,7 @@ OVERRIDE_ASYNC_LET(asyncLet_start, void,
                    (AsyncLet *alet,
                     const Metadata *futureResultType,
                     void *closureEntryPoint,
-                    HeapObject */* +1 */ closureContext
+                    void *closureContext
                    ),
                    (alet, futureResultType,
                     closureEntryPoint, closureContext))

stdlib/public/Concurrency/AsyncLet.cpp

@@ -97,15 +97,15 @@ SWIFT_CC(swift)
 static void swift_asyncLet_startImpl(AsyncLet *alet,
                                      const Metadata *futureResultType,
                                      void *closureEntryPoint,
-                                     HeapObject * /* +1 */ closureContext) {
+                                     void *closureContext) {
   AsyncTask *parent = swift_task_getCurrent();
   assert(parent && "async-let cannot be created without parent task");
 
   auto flags = JobFlags(JobKind::Task, parent->getPriority());
   flags.task_setIsFuture(true);
   flags.task_setIsChildTask(true);
 
-  auto childTaskAndContext = swift_task_create_future(
+  auto childTaskAndContext = swift_task_create_future_no_escaping(
       flags,
       futureResultType,
       closureEntryPoint,

stdlib/public/Concurrency/AsyncLet.swift

@@ -20,7 +20,7 @@ import Swift
 @_silgen_name("swift_asyncLet_start")
 public func _asyncLetStart<T>(
   asyncLet: Builtin.RawPointer,
-  operation: __owned @Sendable @escaping () async throws -> T
+  operation: @Sendable () async throws -> T
 )
 
 /// Similar to _taskFutureGet but for AsyncLet

stdlib/public/Concurrency/Task.cpp

@@ -301,7 +301,7 @@ static void completeTaskWithClosure(SWIFT_ASYNC_CONTEXT AsyncContext *context,
   auto asyncContextPrefix = reinterpret_cast<AsyncContextPrefix *>(
       reinterpret_cast<char *>(context) - sizeof(AsyncContextPrefix));
 
-  swift_release(asyncContextPrefix->closureContext);
+  swift_release((HeapObject *)asyncContextPrefix->closureContext);
 
   // Clean up the rest of the task.
   return completeTask(context, error);
@@ -336,7 +336,7 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
     JobFlags flags, TaskGroup *group,
     const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,
-    HeapObject * /* +1 */ closureContext,
+    void *closureContext, bool owningClosureContext,
     size_t initialContextSize) {
   assert((futureResultType != nullptr) == flags.task_isFuture());
   assert(!flags.task_isFuture() ||
@@ -472,7 +472,8 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
   // be is the final hop.  Store a signed null instead.
   initialContext->Parent = nullptr;
   initialContext->ResumeParent = reinterpret_cast<TaskContinuationFunction *>(
-      closureContext ? &completeTaskWithClosure : &completeTask);
+      (closureContext && owningClosureContext) ? &completeTaskWithClosure :
+                                                 &completeTask);
   initialContext->Flags = AsyncContextKind::Ordinary;
   initialContext->Flags.setShouldNotDeallocateInCallee(true);
 
@@ -493,7 +494,8 @@ static AsyncTaskAndContext swift_task_create_group_future_commonImpl(
 static AsyncTaskAndContext swift_task_create_group_future_common(
     JobFlags flags, TaskGroup *group, const Metadata *futureResultType,
     FutureAsyncSignature::FunctionType *function,
-    HeapObject * /* +1 */ closureContext, size_t initialContextSize);
+    void *closureContext, bool owningClosureContext,
+    size_t initialContextSize);
 
 AsyncTaskAndContext
 swift::swift_task_create_f(JobFlags flags,
@@ -522,6 +524,7 @@ AsyncTaskAndContext swift::swift_task_create_group_future_f(
   return swift_task_create_group_future_common(flags, group,
                                                futureResultType,
                                                function, nullptr,
+                                               /*owningClosureContext*/ false,
                                                initialContextSize);
 }
 
@@ -557,6 +560,26 @@ AsyncTaskAndContext swift::swift_task_create_future(JobFlags flags,
   return swift_task_create_group_future_common(
       flags, nullptr, futureResultType,
       taskEntry, closureContext,
+      /*owningClosureContext*/ true,
+      initialContextSize);
+}
+
+AsyncTaskAndContext swift::swift_task_create_future_no_escaping(JobFlags flags,
+                     const Metadata *futureResultType,
+                     void *closureEntry,
+                     void *closureContext) {
+  FutureAsyncSignature::FunctionType *taskEntry;
+  size_t initialContextSize;
+  std::tie(taskEntry, initialContextSize)
+    = getAsyncClosureEntryPointAndContextSize<
+      FutureAsyncSignature,
+      SpecialPointerAuthDiscriminators::AsyncFutureFunction
+    >(closureEntry, (HeapObject *)closureContext);
+
+  return swift_task_create_group_future_common(
+      flags, nullptr, futureResultType,
+      taskEntry, closureContext,
+      /*owningClosureContext*/ false,
       initialContextSize);
 }
 
@@ -576,6 +599,7 @@ swift::swift_task_create_group_future(
   return swift_task_create_group_future_common(
       flags, group, futureResultType,
       taskEntry, closureContext,
+      /*owningClosureContext*/ true,
       initialContextSize);
 }
 

stdlib/public/Concurrency/TaskPrivate.h

@@ -55,6 +55,11 @@ void runJobInEstablishedExecutorContext(Job *job);
 /// Clear the active task reference for the current thread.
 AsyncTask *_swift_task_clearCurrent();
 
+AsyncTaskAndContext swift_task_create_future_no_escaping(JobFlags flags,
+                     const Metadata *futureResultType,
+                     void *closureEntry,
+                     void *closureContext);
+
 #if defined(SWIFT_STDLIB_SINGLE_THREADED_RUNTIME)
 #define SWIFT_CONCURRENCY_COOPERATIVE_GLOBAL_EXECUTOR 1
 #else

test/SILGen/async_let.swift

@@ -21,8 +21,9 @@ func testAsyncLetInt() async -> Int {
   // CHECK: [[THICK_CLOSURE:%.*]] = thin_to_thick_function [[CLOSURE]] : $@convention(thin) @Sendable @async () -> Int to $@Sendable @async @callee_guaranteed () -> Int
   // CHECK: [[REABSTRACT_THUNK:%.*]] = function_ref @$sSiIeghHd_Sis5Error_pIegHrzo_TR : $@convention(thin) @async (@guaranteed @Sendable @async @callee_guaranteed () -> Int) -> (@out Int, @error Error)
   // CHECK: [[REABSTRACT_CLOSURE:%.*]] = partial_apply [callee_guaranteed] [[REABSTRACT_THUNK]]([[THICK_CLOSURE]]) : $@convention(thin) @async (@guaranteed @Sendable @async @callee_guaranteed () -> Int) -> (@out Int, @error Error)
-  // CHECK: [[CLOSURE_ARG:%.*]] = convert_function [[REABSTRACT_CLOSURE]] : $@async @callee_guaranteed () -> (@out Int, @error Error) to $@async @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error Error) for <Int>
-  // CHECK: [[ASYNC_LET_START:%.*]] = builtin "startAsyncLet"<Int>([[CLOSURE_ARG]] : $@async @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error Error) for <Int>) : $Builtin.RawPointer
+  // CHECK: [[ESCAPING_CLOSURE:%.*]] = convert_function [[REABSTRACT_CLOSURE]] : $@async @callee_guaranteed () -> (@out Int, @error Error) to $@async @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error Error) for <Int>
+  // CHECK: [[CLOSURE_ARG:%.*]] = convert_escape_to_noescape [not_guaranteed] [[ESCAPING_CLOSURE]] : $@async @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error Error) for <Int> to $@noescape @async @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error Error) for <Int>
+  // CHECK: [[ASYNC_LET_START:%.*]] = builtin "startAsyncLet"<Int>([[CLOSURE_ARG]] : $@noescape @async @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error Error) for <Int>) : $Builtin.RawPointer
   async let i = await getInt()
 
   // CHECK: [[ASYNC_LET_GET:%.*]] = function_ref @swift_asyncLet_wait : $@convention(thin) @async <τ_0_0> (Builtin.RawPointer) -> @out τ_0_0