From ff3a3e827d2741721d1ac3307da91ce99113cba8 Mon Sep 17 00:00:00 2001
From: zhongwuzw <zhongwuzw@qq.com>
Date: Mon, 18 Mar 2019 12:01:46 +0800
Subject: [PATCH] Fix prefetch cancellation issues

---
 SDWebImage/Private/SDAsyncBlockOperation.m | 22 ++++++++++++++----
 SDWebImage/SDWebImagePrefetcher.m          | 27 ++++++++++++++++------
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/SDWebImage/Private/SDAsyncBlockOperation.m b/SDWebImage/Private/SDAsyncBlockOperation.m
index 5e3e0fe5..8862ef8e 100644
--- a/SDWebImage/Private/SDAsyncBlockOperation.m
+++ b/SDWebImage/Private/SDAsyncBlockOperation.m
@@ -10,14 +10,17 @@
 
 @interface SDAsyncBlockOperation ()
 
-@property (nonatomic, assign) BOOL isExecuting;
-@property (nonatomic, assign) BOOL isFinished;
+@property (assign, nonatomic, getter = isExecuting) BOOL executing;
+@property (assign, nonatomic, getter = isFinished) BOOL finished;
 @property (nonatomic, copy, nonnull) SDAsyncBlock executionBlock;
 
 @end
 
 @implementation SDAsyncBlockOperation
 
+@synthesize executing = _executing;
+@synthesize finished = _finished;
+
 - (nonnull instancetype)initWithBlock:(nonnull SDAsyncBlock)block {
     self = [super init];
     if (self) {
@@ -32,8 +35,12 @@
 }
 
 - (void)start {
+    if (self.isCancelled) {
+        return;
+    }
+    
     [self willChangeValueForKey:@"isExecuting"];
-    self.isExecuting = YES;
+    self.executing = YES;
     [self didChangeValueForKey:@"isExecuting"];
     
     if (self.executionBlock) {
@@ -43,11 +50,16 @@
     }
 }
 
+- (void)cancel {
+    [super cancel];
+    [self complete];
+}
+
 - (void)complete {
     [self willChangeValueForKey:@"isExecuting"];
     [self willChangeValueForKey:@"isFinished"];
-    self.isExecuting = NO;
-    self.isFinished = YES;
+    self.executing = NO;
+    self.finished = YES;
     [self didChangeValueForKey:@"isExecuting"];
     [self didChangeValueForKey:@"isFinished"];
 }
diff --git a/SDWebImage/SDWebImagePrefetcher.m b/SDWebImage/SDWebImagePrefetcher.m
index 82cb1054..ddf532e7 100644
--- a/SDWebImage/SDWebImagePrefetcher.m
+++ b/SDWebImage/SDWebImagePrefetcher.m
@@ -22,7 +22,8 @@
 }
 
 @property (nonatomic, copy, readwrite) NSArray<NSURL *> *urls;
-@property (nonatomic, strong) NSPointerArray *operations;
+@property (nonatomic, strong) NSPointerArray *loadOperations;
+@property (nonatomic, strong) NSPointerArray *prefetchOperations;
 @property (nonatomic, weak) SDWebImagePrefetcher *prefetcher;
 @property (nonatomic, copy, nullable) SDWebImagePrefetcherCompletionBlock completionBlock;
 @property (nonatomic, copy, nullable) SDWebImagePrefetcherProgressBlock progressBlock;
@@ -93,7 +94,8 @@
     token->_finishedCount = 0;
     token->_totalCount = token.urls.count;
     atomic_flag_clear(&(token->_isAllFinished));
-    token.operations = [NSPointerArray weakObjectsPointerArray];
+    token.loadOperations = [NSPointerArray weakObjectsPointerArray];
+    token.prefetchOperations = [NSPointerArray weakObjectsPointerArray];
     token.progressBlock = progressBlock;
     token.completionBlock = completionBlock;
     [self addRunningToken:token];
@@ -103,12 +105,12 @@
 }
 
 - (void)startPrefetchWithToken:(SDWebImagePrefetchToken * _Nonnull)token {
-    NSPointerArray *operations = token.operations;
+    NSPointerArray *operations = token.loadOperations;
     for (NSURL *url in token.urls) {
         __weak typeof(self) wself = self;
         SDAsyncBlockOperation *prefetchOperation = [SDAsyncBlockOperation blockOperationWithBlock:^(SDAsyncBlockOperation * _Nonnull asyncOperation) {
              __strong typeof(wself) strongSelf = wself;
-             if (!strongSelf) {
+             if (!strongSelf || asyncOperation.isCancelled) {
                  return;
              }
              id<SDWebImageOperation> operation = [strongSelf.manager loadImageWithURL:url options:strongSelf.options context:strongSelf.context progress:nil completed:^(UIImage * _Nullable image, NSData * _Nullable data, NSError * _Nullable error, SDImageCacheType cacheType, BOOL finished, NSURL * _Nullable imageURL) {
@@ -119,7 +121,6 @@
                  if (!finished) {
                      return;
                  }
-                 [asyncOperation complete];
                  
                  atomic_fetch_add_explicit(&(token->_finishedCount), 1, memory_order_relaxed);
                  if (error) {
@@ -137,11 +138,15 @@
                          [sself removeRunningToken:token];
                      }
                  }
+                 [asyncOperation complete];
              }];
             @synchronized (token) {
                 [operations addPointer:(__bridge void *)operation];
             }
         }];
+        @synchronized (token) {
+            [token.prefetchOperations addPointer:(__bridge void *)prefetchOperation];
+        }
         [self.prefetchQueue addOperation:prefetchOperation];
     }
 }
@@ -256,12 +261,20 @@
 
 - (void)cancel {
     @synchronized (self) {
-        for (id operation in self.operations) {
+        [self.prefetchOperations compact];
+        for (id operation in self.prefetchOperations) {
+            SDAsyncBlockOperation *asyncBlockOperation = operation;
+            [asyncBlockOperation cancel];
+        }
+        self.prefetchOperations.count = 0;
+        
+        [self.loadOperations compact];
+        for (id operation in self.loadOperations) {
             if ([operation conformsToProtocol:@protocol(SDWebImageOperation)]) {
                 [operation cancel];
             }
         }
-        self.operations.count = 0;
+        self.loadOperations.count = 0;
     }
     self.completionBlock = nil;
     self.progressBlock = nil;