Merge pull request #1150 from onevcat/fix/prefetcher-thread

Fix prefetcher thread

Sources/Networking/ImagePrefetcher.swift

@@ -69,7 +69,11 @@ public typealias PrefetcherSourceCompletionHandler =
 /// This is useful when you know a list of image resources and want to download them before showing. It also works with
 /// some Cocoa prefetching mechanism like table view or collection view `prefetchDataSource`, to start image downloading
 /// and caching before they display on screen.
-public class ImagePrefetcher {
+public class ImagePrefetcher: CustomStringConvertible {
+
+    public var description: String {
+        return "\(Unmanaged.passUnretained(self).toOpaque())"
+    }
 
     /// The maximum concurrent downloads to use when prefetching images. Default is 5.
     public var maxConcurrentDownloads = 5
@@ -94,7 +98,10 @@ public class ImagePrefetcher {
 
     // A manager used for prefetching. We will use the helper methods in manager.
     private let manager: KingfisherManager
-
+
+    private let pretchQueue = DispatchQueue(label: "com.onevcat.Kingfisher.ImagePrefetcher.pretchQueue")
+    private static let requestingQueue = DispatchQueue(label: "com.onevcat.Kingfisher.ImagePrefetcher.requestingQueue")
+
     private var finished: Bool {
         let totalFinished: Int = failedSources.count + skippedSources.count + completedSources.count
         return totalFinished == prefetchSources.count && tasks.isEmpty
@@ -186,7 +193,7 @@ public class ImagePrefetcher {
         // We want all callbacks from our prefetch queue, so we should ignore the callback queue in options.
         // Add our own callback dispatch queue to make sure all internal callbacks are
         // coming back in our expected queue.
-        options.callbackQueue = .untouch
+        options.callbackQueue = .dispatch(pretchQueue)
         optionsInfo = options
 
         let cache = optionsInfo.targetCache ?? .default
@@ -198,40 +205,44 @@ public class ImagePrefetcher {
     /// of assets that are required for later use in an app. This code will not try and update any UI
     /// with the results of the process.
     public func start() {
-        guard !stopped else {
-            assertionFailure("You can not restart the same prefetcher. Try to create a new prefetcher.")
-            handleComplete()
-            return
-        }
+        pretchQueue.async {
+            guard !self.stopped else {
+                assertionFailure("You can not restart the same prefetcher. Try to create a new prefetcher.")
+                self.handleComplete()
+                return
+            }
 
-        guard maxConcurrentDownloads > 0 else {
-            assertionFailure("There should be concurrent downloads value should be at least 1.")
-            handleComplete()
-            return
-        }
+            guard self.maxConcurrentDownloads > 0 else {
+                assertionFailure("There should be concurrent downloads value should be at least 1.")
+                self.handleComplete()
+                return
+            }
 
-        // Empty case.
-        guard prefetchSources.count > 0 else {
-            handleComplete()
-            return
-        }
+            // Empty case.
+            guard self.prefetchSources.count > 0 else {
+                self.handleComplete()
+                return
+            }
 
-        let initialConcurrentDownloads = min(prefetchSources.count, maxConcurrentDownloads)
-        for _ in 0 ..< initialConcurrentDownloads {
-            if let resource = self.pendingSources.popFirst() {
-                self.startPrefetching(resource)
+            let initialConcurrentDownloads = min(self.prefetchSources.count, self.maxConcurrentDownloads)
+            for _ in 0 ..< initialConcurrentDownloads {
+                if let resource = self.pendingSources.popFirst() {
+                    self.startPrefetching(resource)
+                }
             }
         }
     }
 
     /// Stops current downloading progress, and cancel any future prefetching activity that might be occuring.
     public func stop() {
-        if finished { return }
-        stopped = true
-        tasks.values.forEach { $0.cancel() }
+        pretchQueue.async {
+            if self.finished { return }
+            self.stopped = true
+            self.tasks.values.forEach { $0.cancel() }
+        }
     }
 
-    func downloadAndCache(_ source: Source) {
+    private func downloadAndCache(_ source: Source) {
 
         let downloadTaskCompletionHandler: ((Result<RetrieveImageResult, KingfisherError>) -> Void) = { result in
             self.tasks.removeValue(forKey: source.cacheKey)
@@ -253,24 +264,27 @@ public class ImagePrefetcher {
             }
         }
 
-        let downloadTask = manager.loadAndCacheImage(
-            source: source,
-            options: optionsInfo,
-            completionHandler: downloadTaskCompletionHandler)
-
+        var downloadTask: DownloadTask.WrappedTask?
+        ImagePrefetcher.requestingQueue.sync {
+            downloadTask = manager.loadAndCacheImage(
+                source: source,
+                options: optionsInfo,
+                completionHandler: downloadTaskCompletionHandler)
+        }
+
         if let downloadTask = downloadTask {
             tasks[source.cacheKey] = downloadTask
         }
     }
 
-    func append(cached source: Source) {
+    private func append(cached source: Source) {
         skippedSources.append(source)
 
         reportProgress()
         reportCompletionOrStartNext()
     }
 
-    func startPrefetching(_ source: Source)
+    private func startPrefetching(_ source: Source)
     {
         if optionsInfo.forceRefresh {
             downloadAndCache(source)
@@ -300,27 +314,45 @@ public class ImagePrefetcher {
         }
     }
 
-    func reportProgress() {
-        progressSourceBlock?(skippedSources, failedSources, completedSources)
-        progressBlock?(
-            skippedSources.compactMap { $0.asResource },
-            failedSources.compactMap { $0.asResource },
-            completedSources.compactMap { $0.asResource }
-        )
+    private func reportProgress() {
+
+        if progressBlock == nil && progressSourceBlock == nil {
+            return
+        }
+
+        let skipped = self.skippedSources
+        let failed = self.failedSources
+        let completed = self.completedSources
+        CallbackQueue.mainCurrentOrAsync.execute {
+            self.progressSourceBlock?(skipped, failed, completed)
+            self.progressBlock?(
+                skipped.compactMap { $0.asResource },
+                failed.compactMap { $0.asResource },
+                completed.compactMap { $0.asResource }
+            )
+        }
     }
 
-    func reportCompletionOrStartNext() {
-        CallbackQueue.mainAsync.execute {
-            if let resource = self.pendingSources.popFirst() {
-                self.startPrefetching(resource)
-            } else {
-                guard self.tasks.isEmpty else { return }
-                self.handleComplete()
-            }
+    private func reportCompletionOrStartNext() {
+        if let resource = self.pendingSources.popFirst() {
+            // Loose call stack for huge ammount of sources.
+            pretchQueue.async { self.startPrefetching(resource) }
+        } else {
+            guard allFinished else { return }
+            self.handleComplete()
         }
     }
+
+    var allFinished: Bool {
+        return skippedSources.count + failedSources.count + completedSources.count == prefetchSources.count
+    }
 
-    func handleComplete() {
+    private func handleComplete() {
+
+        if completionHandler == nil && completionSourceHandler == nil {
+            return
+        }
+
         // The completion handler should be called on the main thread
         CallbackQueue.mainCurrentOrAsync.execute {
             self.completionSourceHandler?(self.skippedSources, self.failedSources, self.completedSources)

Sources/Networking/SessionDataTask.swift

@@ -47,8 +47,10 @@ public class SessionDataTask {
     public let task: URLSessionDataTask
     private var callbacksStore = [CancelToken: TaskCallback]()
 
-    var callbacks: Dictionary<SessionDataTask.CancelToken, SessionDataTask.TaskCallback>.Values {
-        return callbacksStore.values
+    var callbacks: [SessionDataTask.TaskCallback] {
+        lock.lock()
+        defer { lock.unlock() }
+        return Array(callbacksStore.values)
     }
 
     private var currentToken = 0

Sources/Networking/SessionDelegate.swift

@@ -248,6 +248,6 @@ extension SessionDelegate: URLSessionDataDelegate {
             return
         }
         remove(task)
-        sessionTask.onTaskDone.call((result, Array(sessionTask.callbacks)))
+        sessionTask.onTaskDone.call((result, sessionTask.callbacks))
     }
 }

Tests/KingfisherTests/ImagePrefetcherTests.swift

@@ -252,7 +252,7 @@ class ImagePrefetcherTests: XCTestCase {
         let cache = KingfisherManager.shared.cache
         let key = testKeys[0]
 
-        cache.store(testImage, forKey: key, callbackQueue: .mainAsync) { result in
+        cache.store(testImage, forKey: key) { result in
             try! cache.memoryStorage.remove(forKey: key)
 
             XCTAssertEqual(cache.imageCachedType(forKey: key), .disk)
@@ -307,12 +307,12 @@ class ImagePrefetcherTests: XCTestCase {
             group.enter()
             let prefetcher = ImagePrefetcher(
                 resources: testURLs,
-                options: [.waitForCache])
+                options: [.cacheMemoryOnly])
             { _, _, _ in group.leave() }
             prefetcher.start()
         }
         group.notify(queue: .main) { exp.fulfill() }
-        waitForExpectations(timeout: 3, handler: nil)
+        waitForExpectations(timeout: 5, handler: nil)
     }
 
     func testPrefetchSources() {

fastlane/Fastfile

@@ -20,10 +20,7 @@ platform :ios do
   end
 
   lane :test do |options|
-    # Sometimes Travis CI would fail due to timeout. 
-    # Not sure why but just rerunning the test perfectly fixes it. :[
-    # Maybe it is related to some racing when running in CI, but it never reproducible in a local env.
-    _test(options) rescue _test(options)
+    _test(options)
   end
 
   private_lane :_test do |options|