Merge pull request onevcat#1913 from onevcat/fix/callback-holding

Fix an issue that callbacks are held even when the task is removed
aybekckaya · Mar 1, 2022 · cbbeb66 · cbbeb66
2 parents 32e4acd + 40320eb
commit cbbeb66
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/Sources/Networking/SessionDataTask.swift b/Sources/Networking/SessionDataTask.swift
@@ -95,6 +95,12 @@ public class SessionDataTask {
         }
         return nil
     }
+
+    func removeAllCallbacks() -> Void {
+        lock.lock()
+        defer { lock.unlock() }
+        callbacksStore.removeAll()
+    }
 
     func resume() {
         guard !started else { return }

diff --git a/Sources/Networking/SessionDelegate.swift b/Sources/Networking/SessionDelegate.swift
@@ -105,6 +105,7 @@ open class SessionDelegate: NSObject {
         guard let url = task.originalURL else {
             return
         }
+        task.removeAllCallbacks()
         tasks[url] = nil
     }
 
@@ -256,7 +257,7 @@ extension SessionDelegate: URLSessionDataDelegate {
         guard let sessionTask = self.task(for: task) else {
             return
         }
-        remove(sessionTask)
         sessionTask.onTaskDone.call((result, sessionTask.callbacks))
+        remove(sessionTask)
     }
 }
diff --git a/Tests/KingfisherTests/KingfisherManagerTests.swift b/Tests/KingfisherTests/KingfisherManagerTests.swift
@@ -1000,6 +1000,28 @@ class KingfisherManagerTests: XCTestCase {
         coordinator.apply(.cachingOriginalImage) { called = true }
         XCTAssertEqual(coordinator.state, .done)
     }
+
+    func testCallbackClearAfterSuccess() {
+        let exp = expectation(description: #function)
+        let url = testURLs[0]
+
+        stub(url, data: testImageData)
+
+        var task: DownloadTask?
+        var called = false
+        task = manager.retrieveImage(with: url) { result in
+            XCTAssertFalse(called)
+            XCTAssertNotNil(result.value?.image)
+            if !called {
+                called = true
+                task?.cancel()
+                DispatchQueue.main.asyncAfter(deadline: .now() + 0.1) {
+                    exp.fulfill()
+                }
+            }
+        }
+        waitForExpectations(timeout: 1, handler: nil)
+    }
 
     func testCanUseCustomizeDefaultCacheSerializer() {
         let exp = expectation(description: #function)