Merge pull request #1016 from AnthonyMDev/whenAll-memory-leaks

Fix memory leaks in Promise whenAll
apollographql · Feb 13, 2020 · 802624e · 802624e
2 parents 995219c + 73ec450
commit 802624e
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 2 deletions.
diff --git a/Sources/Apollo/Promise.swift b/Sources/Apollo/Promise.swift
@@ -3,6 +3,7 @@ import Dispatch
 func whenAll<Value>(_ promises: [Promise<Value>], notifyOn queue: DispatchQueue = .global()) -> Promise<[Value]> {
   return Promise { (fulfill, reject) in
     let group = DispatchGroup()
+    var rejected = false
 
     for promise in promises {
       group.enter()
@@ -11,11 +12,15 @@ func whenAll<Value>(_ promises: [Promise<Value>], notifyOn queue: DispatchQueue
         group.leave()
       }.catch { error in
         reject(error)
+        rejected = true
+        group.leave()
       }
     }
 
     group.notify(queue: queue) {
-      fulfill(promises.map { $0.result!.value! })
+      if !rejected {
+        fulfill(promises.map { $0.result!.value! })
+      }
     }
   }
 }

diff --git a/Sources/Apollo/ResultOrPromise.swift b/Sources/Apollo/ResultOrPromise.swift
@@ -19,6 +19,7 @@ func whenAll<Value>(_ resultsOrPromises: [ResultOrPromise<Value>], notifyOn queu
 
   return .promise(Promise { (fulfill, reject) in
     let group = DispatchGroup()
+    var rejected = false
 
     for resultOrPromise in resultsOrPromises {
       group.enter()
@@ -27,11 +28,15 @@ func whenAll<Value>(_ resultsOrPromises: [ResultOrPromise<Value>], notifyOn queu
         group.leave()
       }.catch { error in
         reject(error)
+        rejected = true
+        group.leave()
       }
     }
 
     group.notify(queue: queue) {
-      fulfill(resultsOrPromises.map { $0.result!.value! })
+      if !rejected {
+        fulfill(resultsOrPromises.map { $0.result!.value! })
+      }
     }
   })
 }

diff --git a/Tests/ApolloTests/PromiseTests.swift b/Tests/ApolloTests/PromiseTests.swift
@@ -249,4 +249,34 @@ class PromiseTests: XCTestCase {
 
     waitForExpectations(timeout: 1)
   }
+
+  func testWhenAllRejectsWhenAnyOfThePromisesRejectsAsync_doesNotCreateMemoryLeak() throws {
+    var executeReject: ((Error) -> Void)?
+
+    var promises: [Promise<String>] = [Promise(fulfilled: "foo"),
+                                       Promise<String> { _, reject in executeReject = reject },
+                                       Promise(fulfilled: "bar")]
+    weak var weakPromise0 = promises[0]
+    weak var weakPromise1 = promises[1]
+    weak var weakPromise2 = promises[2]
+
+    let expectation = self.expectation(description: "whenAll catch handler invoked")
+
+    whenAll(promises).catch { error in
+      XCTAssert(error is TestError)
+
+      expectation.fulfill()
+    }
+
+    promises = []
+    executeReject?(TestError())
+    executeReject = nil
+
+    waitForExpectations(timeout: 1)
+
+    XCTAssertNil(weakPromise0)
+    XCTAssertNil(weakPromise1)
+    XCTAssertNil(weakPromise2)
+  }
+
 }