Implement waiting cache for original image

Sources/General/KingfisherManager.swift

@@ -333,6 +333,10 @@ public class KingfisherManager {
         {
             switch result {
             case .success(let value):
+                let needToCacheOriginalImage = options.cacheOriginalImage &&
+                                               options.processor != DefaultImageProcessor.default
+                let coordinator = CacheCallbackCoordinator(
+                    shouldWaitForCache: options.waitForCache, shouldCacheOriginal: needToCacheOriginalImage)
                 // Add image to cache.
                 let targetCache = options.targetCache ?? self.cache
                 targetCache.store(
@@ -343,7 +347,7 @@ public class KingfisherManager {
                     toDisk: !options.cacheMemoryOnly)
                 {
                     _ in
-                    if options.waitForCache {
+                    coordinator.apply(.cachingImage) {
                         let result = RetrieveImageResult(
                             image: value.image,
                             cacheType: .none,
@@ -355,8 +359,7 @@ public class KingfisherManager {
                 }
 
                 // Add original image to cache if necessary.
-                let needToCacheOriginalImage = options.cacheOriginalImage &&
-                    options.processor != DefaultImageProcessor.default
+
                 if needToCacheOriginalImage {
                     let originalCache = options.originalCache ?? targetCache
                     originalCache.storeToDisk(
@@ -364,9 +367,21 @@ public class KingfisherManager {
                         forKey: source.cacheKey,
                         processorIdentifier: DefaultImageProcessor.default.identifier,
                         expiration: options.diskCacheExpiration)
+                    {
+                        _ in
+                        coordinator.apply(.cachingOriginalImage) {
+                            let result = RetrieveImageResult(
+                                image: value.image,
+                                cacheType: .none,
+                                source: source,
+                                originalSource: context.originalSource
+                            )
+                            completionHandler?(.success(result))
+                        }
+                    }
                 }
 
-                if !options.waitForCache {
+                coordinator.apply(.cacheInitiated) {
                     let result = RetrieveImageResult(
                         image: value.image,
                         cacheType: .none,
@@ -496,6 +511,10 @@ public class KingfisherManager {
 
                             var cacheOptions = options
                             cacheOptions.callbackQueue = .untouch
+
+                            let coordinator = CacheCallbackCoordinator(
+                                shouldWaitForCache: options.waitForCache, shouldCacheOriginal: false)
+
                             targetCache.store(
                                 processedImage,
                                 forKey: key,
@@ -503,7 +522,7 @@ public class KingfisherManager {
                                 toDisk: !options.cacheMemoryOnly)
                             {
                                 _ in
-                                if options.waitForCache {
+                                coordinator.apply(.cachingImage) {
                                     let value = RetrieveImageResult(
                                         image: processedImage,
                                         cacheType: .none,
@@ -514,7 +533,7 @@ public class KingfisherManager {
                                 }
                             }
 
-                            if !options.waitForCache {
+                            coordinator.apply(.cacheInitiated) {
                                 let value = RetrieveImageResult(
                                     image: processedImage,
                                     cacheType: .none,
@@ -566,3 +585,65 @@ struct RetrievingContext {
         return propagationErrors
     }
 }
+
+class CacheCallbackCoordinator {
+
+    enum State {
+        case idle
+        case imageCached
+        case originalImageCached
+        case done
+    }
+
+    enum Action {
+        case cacheInitiated
+        case cachingImage
+        case cachingOriginalImage
+    }
+
+    private let shouldWaitForCache: Bool
+    private let shouldCacheOriginal: Bool
+
+    private (set) var state: State = .idle
+
+    init(shouldWaitForCache: Bool, shouldCacheOriginal: Bool) {
+        self.shouldWaitForCache = shouldWaitForCache
+        self.shouldCacheOriginal = shouldCacheOriginal
+    }
+
+    func apply(_ action: Action, trigger: () -> Void) {
+        switch (state, action) {
+        case (.done, _):
+            break
+
+        // From .idle
+        case (.idle, .cacheInitiated):
+            if !shouldWaitForCache {
+                state = .done
+                trigger()
+            }
+        case (.idle, .cachingImage):
+            if shouldCacheOriginal {
+                state = .imageCached
+            } else {
+                state = .done
+                trigger()
+            }
+        case (.idle, .cachingOriginalImage):
+            state = .originalImageCached
+
+        // From .imageCached
+        case (.imageCached, .cachingOriginalImage):
+            state = .done
+            trigger()
+
+        // From .originalImageCached
+        case (.originalImageCached, .cachingImage):
+            state = .done
+            trigger()
+
+        default:
+            assertionFailure("This case should not happen in CacheCallbackCoordinator: \(state) - \(action)")
+        }
+    }
+}

Sources/General/KingfisherOptionsInfo.swift

@@ -84,10 +84,10 @@ public enum KingfisherOptionsInfoItem {
     /// See `.transition` option for more.
     case forceTransition
     
-    ///  If set, Kingfisher will only cache the value in memory but not in disk.
+    /// If set, Kingfisher will only cache the value in memory but not in disk.
     case cacheMemoryOnly
     
-    ///  If set, Kingfisher will wait for caching operation to be completed before calling the completion block.
+    /// If set, Kingfisher will wait for caching operation to be completed before calling the completion block.
     case waitForCache
     
     /// If set, Kingfisher will only try to retrieve the image from cache, but not from network. If the image is

Tests/KingfisherTests/KingfisherManagerTests.swift

@@ -417,11 +417,9 @@ class KingfisherManagerTests: XCTestCase {
                 options: [.processor(p), .cacheOriginalImage, .originalCache(originalCache)])
             {
                 result in
-                delay(0.2) { // .waitForCache only works for regular cache, not for original cache.
-                    let originalCached = originalCache.imageCachedType(forKey: url.cacheKey)
-                    XCTAssertEqual(originalCached, .disk)
-                    exp.fulfill()
-                }
+                let originalCached = originalCache.imageCachedType(forKey: url.cacheKey)
+                XCTAssertEqual(originalCached, .disk)
+                exp.fulfill()
             }
         }
         waitForExpectations(timeout: 3, handler: nil)
@@ -868,6 +866,57 @@ class KingfisherManagerTests: XCTestCase {
 
         waitForExpectations(timeout: 1, handler: nil)
     }
+
+    func testCacheCallbackCoordinatorStateChanging() {
+        var coordinator = CacheCallbackCoordinator(
+            shouldWaitForCache: false, shouldCacheOriginal: false)
+        var called = false
+        coordinator.apply(.cacheInitiated) {
+            called = true
+        }
+        XCTAssertTrue(called)
+        XCTAssertEqual(coordinator.state, .done)
+        coordinator.apply(.cachingImage) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .done)
+
+        coordinator = CacheCallbackCoordinator(
+            shouldWaitForCache: true, shouldCacheOriginal: false)
+        called = false
+        coordinator.apply(.cacheInitiated) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .idle)
+        coordinator.apply(.cachingImage) {
+            called = true
+        }
+        XCTAssertTrue(called)
+        XCTAssertEqual(coordinator.state, .done)
+
+        coordinator = CacheCallbackCoordinator(
+            shouldWaitForCache: false, shouldCacheOriginal: true)
+        coordinator.apply(.cacheInitiated) {
+            called = true
+        }
+        XCTAssertEqual(coordinator.state, .done)
+        coordinator.apply(.cachingOriginalImage) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .done)
+
+        coordinator = CacheCallbackCoordinator(
+            shouldWaitForCache: true, shouldCacheOriginal: true)
+        coordinator.apply(.cacheInitiated) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .idle)
+        coordinator.apply(.cachingOriginalImage) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .originalImageCached)
+        coordinator.apply(.cachingImage) { called = true }
+        XCTAssertEqual(coordinator.state, .done)
+
+        coordinator = CacheCallbackCoordinator(
+            shouldWaitForCache: true, shouldCacheOriginal: true)
+        coordinator.apply(.cacheInitiated) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .idle)
+        coordinator.apply(.cachingImage) { XCTFail() }
+        XCTAssertEqual(coordinator.state, .imageCached)
+        coordinator.apply(.cachingOriginalImage) { called = true }
+        XCTAssertEqual(coordinator.state, .done)
+    }
 }
 
 class SimpleProcessor: ImageProcessor {