Add semaphore

hulizhen · Jun 2, 2017 · 1035327 · 1035327
1 parent c676d71
commit 1035327
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 13 deletions.
diff --git a/Sources/ImageDownloader.swift b/Sources/ImageDownloader.swift
@@ -183,6 +183,7 @@ open class ImageDownloader {
 
         var downloadTaskCount = 0
         var downloadTask: RetrieveImageDownloadTask?
+        var cancelSemaphore: DispatchSemaphore?
     }
 
     // MARK: - Public property
@@ -219,6 +220,7 @@ open class ImageDownloader {
     // MARK: - Internal property
     let barrierQueue: DispatchQueue
     let processQueue: DispatchQueue
+    let cancelQueue: DispatchQueue
 
     typealias CallbackPair = (progressBlock: ImageDownloaderProgressBlock?, completionHandler: ImageDownloaderCompletionHandler?)
 
@@ -242,6 +244,7 @@ open class ImageDownloader {
 
         barrierQueue = DispatchQueue(label: "com.onevcat.Kingfisher.ImageDownloader.Barrier.\(name)", attributes: .concurrent)
         processQueue = DispatchQueue(label: "com.onevcat.Kingfisher.ImageDownloader.Process.\(name)", attributes: .concurrent)
+        cancelQueue = DispatchQueue(label: "com.onevcat.Kingfisher.ImageDownloader.Cancel.\(name)")
 
         sessionHandler = ImageDownloaderSessionHandler()
 
@@ -312,7 +315,7 @@ open class ImageDownloader {
 
                 dataTask.priority = options?.downloadPriority ?? URLSessionTask.defaultPriority
                 dataTask.resume()
-                delegate?.imageDownloader(self, willDownloadImageForURL: url, with: request)
+                self.delegate?.imageDownloader(self, willDownloadImageForURL: url, with: request)
 
                 // Hold self while the task is executing.
                 self.sessionHandler.downloadHolder = self
@@ -332,19 +335,34 @@ open class ImageDownloader {
 extension ImageDownloader {
 
     // A single key may have multiple callbacks. Only download once.
-    func setup(progressBlock: ImageDownloaderProgressBlock?, with completionHandler: ImageDownloaderCompletionHandler?, for url: URL, options: KingfisherOptionsInfo?, started: ((URLSession, ImageFetchLoad) -> Void)) {
+    func setup(progressBlock: ImageDownloaderProgressBlock?, with completionHandler: ImageDownloaderCompletionHandler?, for url: URL, options: KingfisherOptionsInfo?, started: @escaping ((URLSession, ImageFetchLoad) -> Void)) {
 
-        barrierQueue.sync(flags: .barrier) {
-            let loadObjectForURL = fetchLoads[url] ?? ImageFetchLoad()
-            let callbackPair = (progressBlock: progressBlock, completionHandler: completionHandler)
-
-            loadObjectForURL.contents.append((callbackPair, options ?? KingfisherEmptyOptionsInfo))
-
-            fetchLoads[url] = loadObjectForURL
-
-            if let session = session {
-                started(session, loadObjectForURL)
+        func prepareFetchLoad() {
+            barrierQueue.sync(flags: .barrier) {
+                let loadObjectForURL = fetchLoads[url] ?? ImageFetchLoad()
+                let callbackPair = (progressBlock: progressBlock, completionHandler: completionHandler)
+
+                loadObjectForURL.contents.append((callbackPair, options ?? KingfisherEmptyOptionsInfo))
+
+                fetchLoads[url] = loadObjectForURL
+
+                if let session = session {
+                    started(session, loadObjectForURL)
+                }
+            }
+        }
+
+        if let fetchLoad = fetchLoad(for: url), fetchLoad.downloadTaskCount == 0 {
+            if fetchLoad.cancelSemaphore == nil {
+                fetchLoad.cancelSemaphore = DispatchSemaphore(value: 0)
+            }
+            cancelQueue.async {
+                _ = fetchLoad.cancelSemaphore?.wait(timeout: .distantFuture)
+                fetchLoad.cancelSemaphore = nil
+                prepareFetchLoad()
             }
+        } else {
+            prepareFetchLoad()
         }
     }
 
@@ -464,6 +482,11 @@ class ImageDownloaderSessionHandler: NSObject, URLSessionDataDelegate, Authentic
         // We need to clean the fetch load first, before actually calling completion handler.
         cleanFetchLoad(for: url)
 
+        var leftSignal: Int
+        repeat {
+            leftSignal = fetchLoad.cancelSemaphore?.signal() ?? 0
+        } while leftSignal != 0
+
         for content in fetchLoad.contents {
             content.options.callbackDispatchQueue.safeAsync {
                 content.callback.completionHandler?(nil, error as NSError, url, nil)

diff --git a/Tests/KingfisherTests/ImageDownloaderTests.swift b/Tests/KingfisherTests/ImageDownloaderTests.swift
@@ -275,7 +275,7 @@ class ImageDownloaderTests: XCTestCase {
     }
 
     // Issue 532 https://github.com/onevcat/Kingfisher/issues/532#issuecomment-305644311
-    func _testCancelThenRestartSameDownload() {
+    func testCancelThenRestartSameDownload() {
         let expectation = self.expectation(description: "wait for downloading")
 
         let URLString = testKeys[0]