Fetcher: avoid response deadlock (#975)

cirruslabs · Dec 16, 2024 · e27da23 · e27da23
1 parent e6a30b0
commit e27da23
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/Sources/tart/Fetcher.swift b/Sources/tart/Fetcher.swift
@@ -55,6 +55,7 @@ fileprivate class Delegate: NSObject, URLSessionDataDelegate {
     buffer = Data(capacity: Int(capacity))
 
     responseContinuation?.resume(returning: response)
+    responseContinuation = nil
     completionHandler(.allow)
   }
 
@@ -76,15 +77,20 @@ fileprivate class Delegate: NSObject, URLSessionDataDelegate {
     task: URLSessionTask,
     didCompleteWithError error: Error?
   ) {
-    if !buffer.isEmpty {
-      streamContinuation?.yield(buffer)
-      buffer.removeAll(keepingCapacity: true)
-    }
-
     if let error = error {
+      responseContinuation?.resume(throwing: error)
+      responseContinuation = nil
+
       streamContinuation?.finish(throwing: error)
+      streamContinuation = nil
     } else {
+      if !buffer.isEmpty {
+        streamContinuation?.yield(buffer)
+        buffer.removeAll(keepingCapacity: true)
+      }
+
       streamContinuation?.finish()
+      streamContinuation = nil
     }
   }
 }