Add TaskExecutor conformance to EventLoops

Author: fabianfett

Sources/NIOCore/EventLoop+SerialExecutor.swift

@@ -51,7 +51,7 @@ extension NIOSerialEventLoopExecutor {
 /// This type is not recommended for use because it risks problems with unowned
 /// executors. Adopters are recommended to conform their own event loop
 /// types to `SerialExecutor`.
-final class NIODefaultSerialEventLoopExecutor {
+final class NIODefaultEventLoopExecutor {
     @usableFromInline
     let loop: EventLoop
 
@@ -62,7 +62,7 @@ final class NIODefaultSerialEventLoopExecutor {
 }
 
 @available(macOS 14.0, iOS 17.0, watchOS 10.0, tvOS 17.0, *)
-extension NIODefaultSerialEventLoopExecutor: SerialExecutor {
+extension NIODefaultEventLoopExecutor: SerialExecutor {
     @inlinable
     public func enqueue(_ job: consuming ExecutorJob) {
         self.loop.enqueue(job)
@@ -71,12 +71,42 @@ extension NIODefaultSerialEventLoopExecutor: SerialExecutor {
     @inlinable
     public func asUnownedSerialExecutor() -> UnownedSerialExecutor {
         UnownedSerialExecutor(complexEquality: self)
-
     }
 
     @inlinable
-    public func isSameExclusiveExecutionContext(other: NIODefaultSerialEventLoopExecutor) -> Bool {
+    public func isSameExclusiveExecutionContext(other: NIODefaultEventLoopExecutor) -> Bool {
         self.loop === other.loop
     }
 }
 #endif
+
+#if compiler(>=6.0)
+/// A helper protocol that can be mixed in to a NIO ``EventLoop`` to provide an
+/// automatic conformance to `TaskExecutor`.
+///
+/// Implementers of `EventLoop` should consider conforming to this protocol as
+/// well on Swift 6.0 and later.
+@available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+public protocol NIOTaskEventLoopExecutor: NIOSerialEventLoopExecutor & TaskExecutor { }
+
+@available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+extension NIOTaskEventLoopExecutor {
+    @inlinable
+    func asUnownedTaskExecutor() -> UnownedTaskExecutor {
+        UnownedTaskExecutor(ordinary: self)
+    }
+
+    @inlinable
+    public var taskExecutor: any TaskExecutor {
+        self
+    }
+}
+
+@available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+extension NIODefaultEventLoopExecutor: TaskExecutor {
+    @inlinable
+    public func asUnownedTaskExecutor() -> UnownedTaskExecutor {
+        UnownedTaskExecutor(ordinary: self)
+    }
+}
+#endif

Sources/NIOCore/EventLoop.swift

@@ -383,7 +383,7 @@ extension EventLoop {
     #if compiler(>=5.9)
     @available(macOS 14.0, iOS 17.0, watchOS 10.0, tvOS 17.0, *)
     public var executor: any SerialExecutor {
-        NIODefaultSerialEventLoopExecutor(self)
+        NIODefaultEventLoopExecutor(self)
     }
 
     @inlinable
@@ -398,6 +398,13 @@ extension EventLoop {
         }
     }
     #endif
+
+    #if compiler(>=6.0)
+    @available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+    public var taskExecutor: any TaskExecutor {
+        NIODefaultEventLoopExecutor(self)
+    }
+    #endif
 }
 
 extension EventLoopGroup {

Sources/NIOPosix/SelectableEventLoop.swift

@@ -883,3 +883,9 @@ internal func assertExpression(_ body: () -> Bool) {
         return body()
     }())
 }
+
+// MARK: TaskExecutor conformance
+#if compiler(>=6.0)
+@available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+extension SelectableEventLoop: NIOTaskEventLoopExecutor { }
+#endif

Tests/NIOPosixTests/TaskExecutorTests.swift

@@ -0,0 +1,51 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+import NIOCore
+import NIOEmbedded
+import NIOPosix
+import XCTest
+
+final class TaskExecutorTests: XCTestCase {
+    @available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+    func testBasicExecutorFitsOnEventLoop_MTELG() async throws {
+        #if compiler(>=6.0)
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 2)
+        defer {
+            try! group.syncShutdownGracefully()
+        }
+        let loops = Array(group.makeIterator())
+        await withTaskGroup(of: Void.self) { taskGroup in
+            taskGroup.addTask(executorPreference: loops[0].taskExecutor) {
+                loops[0].assertInEventLoop()
+                loops[1].assertNotInEventLoop()
+
+                withUnsafeCurrentTask { task in
+                    // this currently fails on macOS
+                    XCTAssertEqual(task?.unownedTaskExecutor, loops[0].taskExecutor.asUnownedTaskExecutor())
+                }
+            }
+
+            taskGroup.addTask(executorPreference: loops[1].taskExecutor) {
+                loops[0].assertNotInEventLoop()
+                loops[1].assertInEventLoop()
+
+                withUnsafeCurrentTask { task in
+                    // this currently fails on macOS
+                    XCTAssertEqual(task?.unownedTaskExecutor, loops[1].taskExecutor.asUnownedTaskExecutor())
+                }
+            }
+        }
+        #endif
+    }
+}