Merge pull request #83327 from al45tair/eng/PR-156701386

[Concurrency] Fix task switch performance issue.

Author: al45tair

lib/SILOptimizer/UtilityPasses/Link.cpp

@@ -105,6 +105,8 @@ class SILLinker : public SILModuleTransform {
                            SILLinkage::HiddenExternal);
     linkUsedFunctionByName("swift_createDefaultExecutors",
                            SILLinkage::HiddenExternal);
+    linkUsedFunctionByName("swift_getDefaultExecutor",
+                           SILLinkage::HiddenExternal);
     linkEmbeddedRuntimeWitnessTables();
   }
 

stdlib/public/Concurrency/Actor.cpp

@@ -2391,13 +2391,15 @@ static bool mustSwitchToRun(SerialExecutorRef currentSerialExecutor,
   }
 
   // else, we may have to switch if the preferred task executor is different
-  auto differentTaskExecutor =
-      currentTaskExecutor.getIdentity() != newTaskExecutor.getIdentity();
-  if (differentTaskExecutor) {
-    return true;
-  }
+  if (currentTaskExecutor.getIdentity() == newTaskExecutor.getIdentity())
+    return false;
 
-  return false;
+  if (currentTaskExecutor.isUndefined())
+    currentTaskExecutor = swift_getDefaultExecutor();
+  if (newTaskExecutor.isUndefined())
+    newTaskExecutor = swift_getDefaultExecutor();
+
+  return currentTaskExecutor.getIdentity() != newTaskExecutor.getIdentity();
 }
 
 /// Given that we've assumed control of an executor on this thread,

stdlib/public/Concurrency/Executor.swift

@@ -625,6 +625,13 @@ extension MainActor {
     _createDefaultExecutorsOnce()
     return _executor!
   }
+
+  /// An unowned version of the above, for performance
+  @available(StdlibDeploymentTarget 6.2, *)
+  static var unownedExecutor: UnownedSerialExecutor {
+    _createDefaultExecutorsOnce()
+    return unsafe UnownedSerialExecutor(ordinary: _executor!)
+  }
 }
 #endif // os(WASI) || (!$Embedded && !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY)
 
@@ -643,6 +650,13 @@ extension Task where Success == Never, Failure == Never {
     _createDefaultExecutorsOnce()
     return _defaultExecutor!
   }
+
+  /// An unowned version of the above, for performance
+  @available(StdlibDeploymentTarget 6.2, *)
+  static var unownedDefaultExecutor: UnownedTaskExecutor {
+    _createDefaultExecutorsOnce()
+    return unsafe UnownedTaskExecutor(_defaultExecutor!)
+  }
 }
 
 extension Task where Success == Never, Failure == Never {

stdlib/public/Concurrency/ExecutorBridge.h

@@ -23,6 +23,9 @@ namespace swift {
 extern "C" SWIFT_CC(swift)
 SerialExecutorRef swift_getMainExecutor();
 
+extern "C" SWIFT_CC(swift)
+TaskExecutorRef swift_getDefaultExecutor();
+
 #if !SWIFT_CONCURRENCY_EMBEDDED
 extern "C" SWIFT_CC(swift)
 void *swift_createDispatchEventC(void (*handler)(void *), void *context);

stdlib/public/Concurrency/ExecutorBridge.swift

@@ -95,17 +95,23 @@ internal func _jobGetExecutorPrivateData(
 #if !SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
 @available(StdlibDeploymentTarget 6.2, *)
 @_silgen_name("swift_getMainExecutor")
-internal func _getMainExecutorAsSerialExecutor() -> (any SerialExecutor)? {
-  return MainActor.executor
+internal func _getMainExecutorAsSerialExecutor() -> UnownedSerialExecutor {
+  return unsafe MainActor.unownedExecutor
 }
 #else
 // For task-to-thread model, this is implemented in C++
 @available(StdlibDeploymentTarget 6.2, *)
 @_silgen_name("swift_getMainExecutor")
-internal func _getMainExecutorAsSerialExecutor() -> (any SerialExecutor)?
+internal func _getMainExecutorAsSerialExecutor() -> UnownedSerialExecutor
 #endif // SWIFT_STDLIB_TASK_TO_THREAD_MODEL_CONCURRENCY
 #endif // os(WASI) || !$Embedded
 
+@available(StdlibDeploymentTarget 6.2, *)
+@_silgen_name("swift_getDefaultExecutor")
+internal func _getDefaultExecutorAsTaskExecutor() -> UnownedTaskExecutor {
+  return unsafe Task.unownedDefaultExecutor
+}
+
 @available(StdlibDeploymentTarget 6.2, *)
 @_silgen_name("swift_dispatchMain")
 internal func _dispatchMain()

stdlib/public/Concurrency/ExecutorImpl.swift

@@ -47,11 +47,7 @@ internal func donateToGlobalExecutor(
 @available(SwiftStdlib 6.2, *)
 @_silgen_name("swift_task_getMainExecutorImpl")
 internal func getMainExecutor() -> UnownedSerialExecutor {
-  let executor = _getMainExecutorAsSerialExecutor()
-  if let executor {
-    return unsafe executor.asUnownedSerialExecutor()
-  }
-  return unsafe unsafeBitCast(executor, to: UnownedSerialExecutor.self)
+  return _getMainExecutorAsSerialExecutor()
 }
 
 @available(SwiftStdlib 6.2, *)

test/Concurrency/Runtime/custom_task_executor_fast_path.swift

@@ -0,0 +1,117 @@
+// RUN: %target-run-simple-swift(-Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+
+// rdar://106849189 move-only types should be supported in freestanding mode
+// UNSUPPORTED: freestanding
+
+// UNSUPPORTED: back_deployment_runtime
+// REQUIRES: concurrency_runtime
+// REQUIRES: synchronization
+// REQUIRES: libdispatch
+
+import StdlibUnittest
+import Synchronization
+import Dispatch
+
+typealias DefaultExecutorFactory = SimpleExecutorFactory
+
+struct SimpleExecutorFactory: ExecutorFactory {
+  public static var mainExecutor: any MainExecutor {
+    return SimpleMainExecutor()
+  }
+  public static var defaultExecutor: any TaskExecutor {
+    return SimpleTaskExecutor()
+  }
+}
+
+@available(SwiftStdlib 6.2, *)
+final class SimpleMainExecutor: MainExecutor, @unchecked Sendable {
+  public var isRunning: Bool = false
+
+  var shouldStop: Bool = false
+  let queue = Mutex<[UnownedJob]>([])
+
+  func enqueue(_ job: consuming ExecutorJob) {
+    let unownedJob = UnownedJob(job)
+    queue.withLock {
+      $0.append(unownedJob)
+    }
+  }
+
+  func run() throws {
+    isRunning = true
+    while !shouldStop {
+      let jobs = queue.withLock {
+        let jobs = $0
+        $0.removeAll()
+        return jobs
+      }
+      for job in jobs {
+        job.runSynchronously(on: self.asUnownedSerialExecutor())
+      }
+    }
+    isRunning = false
+  }
+
+  func stop() {
+    shouldStop = true
+  }
+}
+
+@available(SwiftStdlib 6.2, *)
+final class SimpleTaskExecutor: TaskExecutor, @unchecked Sendable {
+  let queue = Mutex<[UnownedJob]>([])
+
+  func enqueue(_ job: consuming ExecutorJob) {
+    print("Enqueued")
+    let unownedJob = UnownedJob(job)
+    queue.withLock {
+      $0.append(unownedJob)
+    }
+  }
+
+  func run() throws {
+    while true {
+      let jobs = queue.withLock {
+        let jobs = $0
+        $0.removeAll()
+        return jobs
+      }
+      for job in jobs {
+        job.runSynchronously(on: self.asUnownedTaskExecutor())
+      }
+    }
+  }
+}
+
+@available(SwiftStdlib 6.2, *)
+@main struct Main {
+  static func main() async {
+    DispatchQueue.global().async {
+      try! (Task.defaultExecutor as! SimpleTaskExecutor).run()
+    }
+
+    await withTaskGroup { group in
+      for _ in 0..<3 {
+        group.addTask() {
+          for _ in 0..<100 {
+            await withUnsafeContinuation { cont in
+              cont.resume()
+            }
+          }
+        }
+      }
+    }
+
+  }
+}
+
+// If we're on the fast path, we'll only enqueue four times (once per group)
+
+// CHECK: Enqueued
+// CHECK: Enqueued
+// CHECK: Enqueued
+// CHECK: Enqueued
+// CHECK-NOT: Enqueued