Don't hold a lock over a continuation in test helpers (#872)

Motivation: The various 'withMumbleContinuation' APIs are supposed to be invoked synchronously with the caller. This assumption allows a lock to be acquired before the call and released from the body of the 'withMumbleContinuation' after e.g. storing the continuation. However this isn't the case and the job may be re-enqueued on the executor meaning that this is pattern is vulnerable to deadlocks. Modifications: - Rework the test helpers to avoid holding a lock when a continuation is created. - Switch to using NIOLockedValue box Result: Lower chance of deadlock

Author: glbrntt

Tests/AsyncHTTPClientTests/AsyncTestHelpers.swift

@@ -43,12 +43,11 @@ final class AsyncSequenceWriter<Element: Sendable>: AsyncSequence, @unchecked Se
 case failed(Error,CheckedContinuation<Void,Never>?)
 }
 
-privatevar_state=State.buffering(.init(),nil)
-privateletlock=NIOLock()
+privateletstate=NIOLockedValueBox<State>(.buffering([],nil))
 
 publicvarhasDemand:Bool{
-self.lock.withLock{
-switchself._state{
+self.state.withLockedValue{ state in
+switchstate{
 case.failed,.finished,.buffering:
 returnfalse
 case.waiting:
@@ -59,67 +58,132 @@ final class AsyncSequenceWriter<Element: Sendable>: AsyncSequence, @unchecked Se
 
  /// Wait until a downstream consumer has issued more demand by calling `next`.
 publicfunc demand()async{
-self.lock.lock()
+letshouldBuffer=self.state.withLockedValue{ state in
+switch state {
+case.buffering(_,.none):
+returntrue
+case.waiting:
+returnfalse
+case.buffering(_,.some),.failed(_,.some):
+preconditionFailure("Already waiting for demand. Invalid state: \(state)")
+case.finished,.failed:
+preconditionFailure("Invalid state: \(state)")
+}
+}
 
-switchself._state {
-case.buffering(let buffer,.none):
+if shouldBuffer {
 awaitwithCheckedContinuation{(continuation:CheckedContinuation<Void,Never>)in
-self._state =.buffering(buffer, continuation)
-self.lock.unlock()
+letshouldResumeContinuation=self.state.withLockedValue{ state in
+switch state {
+case.buffering(let buffer,.none):
+ state =.buffering(buffer, continuation)
+returnfalse
+case.waiting:
+returntrue
+case.buffering(_,.some),.failed(_,.some):
+preconditionFailure("Already waiting for demand. Invalid state: \(state)")
+case.finished,.failed:
+preconditionFailure("Invalid state: \(state)")
+}
+}
+
+if shouldResumeContinuation {
+ continuation.resume()
+}
 }
-
-case.waiting:
-self.lock.unlock()
-return
-
-case.buffering(_,.some),.failed(_,.some):
-letstate=self._state
-self.lock.unlock()
-preconditionFailure("Already waiting for demand. Invalid state: \(state)")
-
-case.finished,.failed:
-letstate=self._state
-self.lock.unlock()
-preconditionFailure("Invalid state: \(state)")
 }
 }
 
+privateenumNextAction{
+ /// Resume the continuation if present, and return the result if present.
+case resumeAndReturn(CheckedContinuation<Void,Never>?,Result<Element?,Error>?)
+ /// Suspend the current task and wait for the next value.
+case suspend
+}
+
 privatefunc next()asyncthrows->Element?{
-self.lock.lock()
-switchself._state {
-case.buffering(let buffer,let demandContinuation)where buffer.isEmpty:
-returntryawaitwithCheckedThrowingContinuation{ continuation in
-self._state =.waiting(continuation)
-self.lock.unlock()
- demandContinuation?.resume(returning:())
-}
+letaction:NextAction=self.state.withLockedValue{ state in
+switch state {
+case.buffering(var buffer,let demandContinuation):
+if buffer.isEmpty {
+return.suspend
+}else{
+letfirst= buffer.removeFirst()
+if first !=nil{
+ state =.buffering(buffer, demandContinuation)
+}else{
+ state =.finished
+}
+return.resumeAndReturn(nil,.success(first))
+}
+
+case.failed(let error,let demandContinuation):
+ state =.finished
+return.resumeAndReturn(demandContinuation,.failure(error))
+
+case.finished:
+return.resumeAndReturn(nil,.success(nil))
 
-case.buffering(var buffer,let demandContinuation):
-letfirst= buffer.removeFirst()
-if first !=nil{
-self._state =.buffering(buffer, demandContinuation)
-}else{
-self._state =.finished
+case.waiting:
+preconditionFailure(
+"Expected that there is always only one concurrent call to next. Invalid state: \(state)"
+)
 }
-self.lock.unlock()
-return first
+}
 
-case.failed(let error,let demandContinuation):
-self._state =.finished
-self.lock.unlock()
+switch action {
+case.resumeAndReturn(let demandContinuation,let result):
  demandContinuation?.resume()
-throw error
-
-case.finished:
-self.lock.unlock()
-returnnil
-
-case.waiting:
-letstate=self._state
-self.lock.unlock()
-preconditionFailure(
-"Expected that there is always only one concurrent call to next. Invalid state: \(state)"
-)
+returntry result?.get()
+
+case.suspend:
+ // Holding the lock here *should* be safe but because of a bug in the runtime
+ // it isn't, so drop the lock, create the continuation and then try again.
+ //
+ // See https://github.com/swiftlang/swift/issues/85668
+returntryawaitwithCheckedThrowingContinuation{
+(continuation:CheckedContinuation<Element?,anyError>)in
+letaction:NextAction=self.state.withLockedValue{ state in
+switch state {
+case.buffering(var buffer,let demandContinuation):
+if buffer.isEmpty {
+ state =.waiting(continuation)
+return.resumeAndReturn(demandContinuation,nil)
+}else{
+letfirst= buffer.removeFirst()
+if first !=nil{
+ state =.buffering(buffer, demandContinuation)
+}else{
+ state =.finished
+}
+return.resumeAndReturn(nil,.success(first))
+}
+
+case.failed(let error,let demandContinuation):
+ state =.finished
+return.resumeAndReturn(demandContinuation,.failure(error))
+
+case.finished:
+return.resumeAndReturn(nil,.success(nil))
+
+case.waiting:
+preconditionFailure(
+"Expected that there is always only one concurrent call to next. Invalid state: \(state)"
+)
+}
+}
+
+switch action {
+case.resumeAndReturn(let demandContinuation,let result):
+ demandContinuation?.resume()
+ // Resume the continuation rather than returning th result.
+iflet result {
+ continuation.resume(with: result)
+}
+case.suspend:
+preconditionFailure() // Not returned from the code above.
+}
+}
 }
 }
 
@@ -137,19 +201,19 @@ final class AsyncSequenceWriter<Element: Sendable>: AsyncSequence, @unchecked Se
 }
 
 privatefunc writeBufferOrEnd(_ element:Element?){
-letwriteAction=self.lock.withLock{()->WriteActionin
-switchself._state{
+letwriteAction=self.state.withLockedValue{state->WriteActionin
+switchstate{
 case.buffering(var buffer,let continuation):
  buffer.append(element)
-self._state=.buffering(buffer, continuation)
+state=.buffering(buffer, continuation)
 return.none
 
 case.waiting(let continuation):
-self._state=.buffering(.init(),nil)
+state=.buffering(.init(),nil)
 return.succeedContinuation(continuation, element)
 
 case.finished,.failed:
-preconditionFailure("Invalid state: \(self._state)")
+preconditionFailure("Invalid state: \(state)")
 }
 }
 
@@ -170,17 +234,17 @@ final class AsyncSequenceWriter<Element: Sendable>: AsyncSequence, @unchecked Se
  /// Drops all buffered writes and emits an error on the waiting `next`. If there is no call to `next`
  /// waiting, will emit the error on the next call to `next`.
 publicfunc fail(_ error:Error){
-leterrorAction=self.lock.withLock{()->ErrorActionin
-switchself._state{
+leterrorAction=self.state.withLockedValue{state->ErrorActionin
+switchstate{
 case.buffering(_,let demandContinuation):
-self._state=.failed(error, demandContinuation)
+state=.failed(error, demandContinuation)
 return.none
 
 case.failed,.finished:
 return.none
 
 case.waiting(let continuation):
-self._state=.finished
+state=.finished
 return.failContinuation(continuation, error)
 }
 }