Repository: apple/swift-nio
File: Tests/NIOCoreTests/AsyncSequences/NIOAsyncWriterTests.swift (Swift)

Code before:
        self.sink.setWritability(to: false)

        let suspended = expectation(description: "suspended on yield")
        self.delegate.didSuspendHandler = {
            suspended.fulfill()
        }
        let task = Task { [writer] in
            try await writer!.yield("message1")
        }
        await fulfillment(of: [suspended], timeout: 1)

        self.writer.finish()

        self.assert(suspendCallCount: 1, yieldCallCount: 0, terminateCallCount: 0)

        // We have to become writable again to unbuffer the yield
        self.sink.setWritability(to: true)

        await XCTAssertNoThrow(try await task.value)

        self.assert(suspendCallCount: 1, yieldCallCount: 1, terminateCallCount: 1)
    }

    func testWriterFinish_AndSuspendBufferedYield() async throws {
        #if compiler(>=6)
        try await withNIOThreadPoolTaskExecutor(numberOfThreads: 2) { taskExecutor in
            try await withThrowingTaskGroup(of: Void.self) { group in
                self.sink.setWritability(to: false)

                let bothSuspended = expectation(description: "suspended on both yields")
                let suspendedAgain = ConditionLock(value: false)
                self.delegate.didSuspendHandler = {
                    if self.delegate.didSuspendCallCount == 2 {
                        bothSuspended.fulfill()
                    } else if self.delegate.didSuspendCallCount > 2 {
                        suspendedAgain.lock()
                        suspendedAgain.unlock(withValue: true)
                    }
                }

                self.delegate.didYieldHandler = { _ in
                    if self.delegate.didYieldCallCount == 1 {
                        // Delay this yield until the other yield is suspended again.
                        if suspendedAgain.lock(whenValue: true, timeoutSeconds: 5) {
                            suspendedAgain.unlock()
                        } else {
                            XCTFail("Timeout while waiting for other yield to suspend again.")
                        }
                    }
                }


Review comment (suggestion):
Sorry for being late to the party but I just realized that this test could be completely re-written to make it work without any thread backed executors and fully reliable. What I wanted to achieve with this test originally was to control the order of the yields and their suspension. I think using a manual task executor would make this test 100 times more maintainable and remove any of the locks.

This is how a manual executor looks like:

```
import DequeModule
import Synchronization

@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
final class ManualTaskExecutor: TaskExecutor {
    private let jobs = Mutex<Deque<UnownedJob>>(.init())
    
    func enqueue(_ job: UnownedJob) {
        self.jobs.withLock { $0.append(job) }
    }
    
    func run() {
        while let job = self.jobs.withLock({ $0.popFirst() }) {
            job.runSynchronously(on: self.asUnownedTaskExecutor())
        }
    }
}
```

With such an executor we should be able to exactly control what child task runs at what point and control when they suspend.

Code after:
        self.sink.setWritability(to: false)

        let suspended = expectation(description: "suspended on yield")
        self.delegate.didSuspendHandler = {
            suspended.fulfill()
        }
        let task = Task { [writer] in
            try await writer!.yield("message1")
        }
        await fulfillment(of: [suspended], timeout: 1)

        self.writer.finish()

        self.assert(suspendCallCount: 1, yieldCallCount: 0, terminateCallCount: 0)

        // We have to become writable again to unbuffer the yield
        self.sink.setWritability(to: true)

        await XCTAssertNoThrow(try await task.value)

        self.assert(suspendCallCount: 1, yieldCallCount: 1, terminateCallCount: 1)
    }

    #if compiler(>=6)
    @available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *)
    func testWriterFinish_AndSuspendBufferedYield() async throws {
        try await withThrowingTaskGroup(of: Void.self) { group in
            try await withManualTaskExecutor { taskExecutor1, taskExecutor2 in
                self.sink.setWritability(to: false)

                self.delegate.didYieldHandler = { _ in
                    if self.delegate.didYieldCallCount == 1 {
                        // This is the yield of the first task. Run the second task until it suspends again
                        self.assert(suspendCallCount: 2, yieldCallCount: 1, terminateCallCount: 0)
                        taskExecutor2.runUntilQueueIsEmpty()
                        self.assert(suspendCallCount: 3, yieldCallCount: 1, terminateCallCount: 0)
                    }
                }

                group.addTask(executorPreference: taskExecutor1) { [writer] in
                    try await writer!.yield("message1")
                }
                group.addTask(executorPreference: taskExecutor2) { [writer] in
                    try await writer!.yield("message2")
                }

                // Run tasks until they are both suspended
                taskExecutor1.runUntilQueueIsEmpty()
                taskExecutor2.runUntilQueueIsEmpty()
                self.assert(suspendCallCount: 2, yieldCallCount: 0, terminateCallCount: 0)