Switch to using UnsafeContinuations in NIOAsyncWriter and NIOAsyncSequenceProducer

Benchmarks/Thresholds/5.10/NIOPosixBenchmarks.TCPEchoAsyncChannel.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 1317015
-}
+  "mallocCountTotal" : 82420
+}

Benchmarks/Thresholds/5.9/NIOPosixBenchmarks.TCPEchoAsyncChannel.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 1317022
+  "mallocCountTotal" : 82426
 }

Benchmarks/Thresholds/main/NIOPosixBenchmarks.TCPEchoAsyncChannel.p90.json

@@ -1,3 +1,3 @@
 {
-  "mallocCountTotal" : 1317015
+  "mallocCountTotal" : 82420
 }

Sources/NIOCore/AsyncChannel/AsyncChannelOutboundWriterHandler.swift

@@ -170,25 +170,34 @@ extension NIOAsyncChannelOutboundWriterHandler {
 
         @inlinable
         func didYield(contentsOf sequence: Deque<OutboundOut>) {
-            // This always called from an async context, so we must loop-hop.
-            self.eventLoop.execute {
+            if self.eventLoop.inEventLoop {
                 self.handler._didYield(sequence: sequence)
+            } else {
+                self.eventLoop.execute {
+                    self.handler._didYield(sequence: sequence)
+                }
             }
         }
 
         @inlinable
         func didYield(_ element: OutboundOut) {
-            // This always called from an async context, so we must loop-hop.
-            self.eventLoop.execute {
+            if self.eventLoop.inEventLoop {
                 self.handler._didYield(element: element)
+            } else {
+                self.eventLoop.execute {
+                    self.handler._didYield(element: element)
+                }
             }
         }
 
         @inlinable
         func didTerminate(error: Error?) {
-            // This always called from an async context, so we must loop-hop.
-            self.eventLoop.execute {
+            if self.eventLoop.inEventLoop {
                 self.handler._didTerminate(error: error)
+            } else {
+                self.eventLoop.execute {
+                    self.handler._didTerminate(error: error)
+                }
             }
         }
     }

Sources/NIOCore/AsyncSequences/NIOThrowingAsyncSequenceProducer.swift

@@ -533,7 +533,7 @@ extension NIOThrowingAsyncSequenceProducer {
                 case .suspendTask:
                     // It is safe to hold the lock across this method
                     // since the closure is guaranteed to be run straight away
-                    return try await withCheckedThrowingContinuation { continuation in
+                    return try await withUnsafeThrowingContinuation { continuation in
                         let action = self._stateMachine.next(for: continuation)
 
                         switch action {
@@ -606,7 +606,7 @@ extension NIOThrowingAsyncSequenceProducer {
             case streaming(
                 backPressureStrategy: Strategy,
                 buffer: Deque<Element>,
-                continuation: CheckedContinuation<Element?, Error>?,
+                continuation: UnsafeContinuation<Element?, Error>?,
                 hasOutstandingDemand: Bool,
                 iteratorInitialized: Bool
             )
@@ -787,20 +787,20 @@ extension NIOThrowingAsyncSequenceProducer {
             /// Indicates that the continuation should be resumed and
             /// ``NIOThrowingAsyncSequenceProducer/Source/YieldResult/produceMore`` should be returned.
             case resumeContinuationAndReturnProduceMore(
-                continuation: CheckedContinuation<Element?, Error>,
+                continuation: UnsafeContinuation<Element?, Error>,
                 element: Element
             )
             /// Indicates that the continuation should be resumed and
             /// ``NIOThrowingAsyncSequenceProducer/Source/YieldResult/stopProducing`` should be returned.
             case resumeContinuationAndReturnStopProducing(
-                continuation: CheckedContinuation<Element?, Error>,
+                continuation: UnsafeContinuation<Element?, Error>,
                 element: Element
             )
             /// Indicates that the yielded elements have been dropped.
             case returnDropped
 
             @usableFromInline
-            init(shouldProduceMore: Bool, continuationAndElement: (CheckedContinuation<Element?, Error>, Element)? = nil) {
+            init(shouldProduceMore: Bool, continuationAndElement: (UnsafeContinuation<Element?, Error>, Element)? = nil) {
                 switch (shouldProduceMore, continuationAndElement) {
                 case (true, .none):
                     self = .returnProduceMore
@@ -902,7 +902,7 @@ extension NIOThrowingAsyncSequenceProducer {
         enum FinishAction {
             /// Indicates that the continuation should be resumed with `nil` and
             /// that ``NIOAsyncSequenceProducerDelegate/didTerminate()`` should be called.
-            case resumeContinuationWithFailureAndCallDidTerminate(CheckedContinuation<Element?, Error>, Failure?)
+            case resumeContinuationWithFailureAndCallDidTerminate(UnsafeContinuation<Element?, Error>, Failure?)
             /// Indicates that nothing should be done.
             case none
         }
@@ -956,7 +956,7 @@ extension NIOThrowingAsyncSequenceProducer {
             case callDidTerminate
             /// Indicates that the continuation should be resumed with a `CancellationError` and
             /// that ``NIOAsyncSequenceProducerDelegate/didTerminate()`` should be called.
-            case resumeContinuationWithCancellationErrorAndCallDidTerminate(CheckedContinuation<Element?, Error>)
+            case resumeContinuationWithCancellationErrorAndCallDidTerminate(UnsafeContinuation<Element?, Error>)
             /// Indicates that nothing should be done.
             case none
         }
@@ -1129,7 +1129,7 @@ extension NIOThrowingAsyncSequenceProducer {
         }
 
         @inlinable
-        mutating func next(for continuation: CheckedContinuation<Element?, Error>) -> NextForContinuationAction {
+        mutating func next(for continuation: UnsafeContinuation<Element?, Error>) -> NextForContinuationAction {
             switch self._state {
             case .initial:
                 // We are transitioning away from the initial state in `next()`

Tests/NIOCoreTests/AsyncChannel/AsyncChannelTests.swift

@@ -85,7 +85,7 @@ final class AsyncChannelTests: XCTestCase {
             inboundReader = wrapped.inboundStream
 
             try await channel.testingEventLoop.executeInContext {
-                XCTAssertEqual(0, closeRecorder.outboundCloses)
+                XCTAssertEqual(1, closeRecorder.outboundCloses)
             }
         }
 
@@ -159,7 +159,7 @@ final class AsyncChannelTests: XCTestCase {
                 inboundReader = wrapped.inboundStream
 
                 try await channel.testingEventLoop.executeInContext {
-                    XCTAssertEqual(0, closeRecorder.allCloses)
+                    XCTAssertEqual(1, closeRecorder.allCloses)
                 }
             }
 

Tests/NIOCoreTests/AsyncSequences/NIOAsyncWriterTests.swift

@@ -15,25 +15,32 @@
 import DequeModule
 import NIOCore
 import XCTest
+import NIOConcurrencyHelpers
 
 private struct SomeError: Error, Hashable {}
 
 private final class MockAsyncWriterDelegate: NIOAsyncWriterSinkDelegate, @unchecked Sendable {
     typealias Element = String
 
-    var didYieldCallCount = 0
+    var _didYieldCallCount = NIOLockedValueBox(0)
+    var didYieldCallCount: Int {
+        self._didYieldCallCount.withLockedValue { $0 }
+    }
     var didYieldHandler: ((Deque<String>) -> Void)?
     func didYield(contentsOf sequence: Deque<String>) {
-        self.didYieldCallCount += 1
+        self._didYieldCallCount.withLockedValue { $0 += 1 }
         if let didYieldHandler = self.didYieldHandler {
             didYieldHandler(sequence)
         }
     }
 
-    var didTerminateCallCount = 0
+    var _didTerminateCallCount = NIOLockedValueBox(0)
+    var didTerminateCallCount: Int {
+        self._didTerminateCallCount.withLockedValue { $0 }
+    }
     var didTerminateHandler: ((Error?) -> Void)?
     func didTerminate(error: Error?) {
-        self.didTerminateCallCount += 1
+        self._didTerminateCallCount.withLockedValue { $0 += 1 }
         if let didTerminateHandler = self.didTerminateHandler {
             didTerminateHandler(error)
         }