Merge pull request #2132 from dreampiggy/fix_WebCacheOperation_use_maptable

Use a copy-weak maptable for operations stored in UIView(WebCacheOperation) category to avoid retain of operation, and also use lock to keep thread-safe

SDWebImage/UIButton+WebCache.m

@@ -146,18 +146,10 @@ static inline NSString * backgroundImageURLKeyForState(UIControlState state) {
                            completed:completedBlock];
 }
 
-- (void)sd_setImageLoadOperation:(id<SDWebImageOperation>)operation forState:(UIControlState)state {
-    [self sd_setImageLoadOperation:operation forKey:[NSString stringWithFormat:@"UIButtonImageOperation%@", @(state)]];
-}
-
 - (void)sd_cancelImageLoadForState:(UIControlState)state {
     [self sd_cancelImageLoadOperationWithKey:[NSString stringWithFormat:@"UIButtonImageOperation%@", @(state)]];
 }
 
-- (void)sd_setBackgroundImageLoadOperation:(id<SDWebImageOperation>)operation forState:(UIControlState)state {
-    [self sd_setImageLoadOperation:operation forKey:[NSString stringWithFormat:@"UIButtonBackgroundImageOperation%@", @(state)]];
-}
-
 - (void)sd_cancelBackgroundImageLoadForState:(UIControlState)state {
     [self sd_cancelImageLoadOperationWithKey:[NSString stringWithFormat:@"UIButtonBackgroundImageOperation%@", @(state)]];
 }

SDWebImage/UIImageView+WebCache.m

@@ -73,7 +73,7 @@
     [self sd_cancelCurrentAnimationImagesLoad];
     __weak __typeof(self)wself = self;
 
-    NSMutableArray<id<SDWebImageOperation>> *operationsArray = [[NSMutableArray alloc] init];
+    NSPointerArray *operationsArray = [self sd_animationOperationArray];
 
     [arrayOfURLs enumerateObjectsUsingBlock:^(NSURL *logoImageURL, NSUInteger idx, BOOL * _Nonnull stop) {
         id <SDWebImageOperation> operation = [[SDWebImageManager sharedManager] loadImageWithURL:logoImageURL options:0 progress:nil completed:^(UIImage *image, NSData *data, NSError *error, SDImageCacheType cacheType, BOOL finished, NSURL *imageURL) {
@@ -103,14 +103,40 @@
                 [sself startAnimating];
             });
         }];
-        [operationsArray addObject:operation];
+        @synchronized (self) {
+            [operationsArray addPointer:(__bridge void *)(operation)];
+        }
     }];
+}
 
-    [self sd_setImageLoadOperation:[operationsArray copy] forKey:@"UIImageViewAnimationImages"];
+static char animationLoadOperationKey;
+
+// element is weak because operation instance is retained by SDWebImageManager's runningOperations property
+// we should use lock to keep thread-safe because these method may not be acessed from main queue
+- (NSPointerArray *)sd_animationOperationArray {
+    @synchronized(self) {
+        NSPointerArray *operationsArray = objc_getAssociatedObject(self, &animationLoadOperationKey);
+        if (operationsArray) {
+            return operationsArray;
+        }
+        operationsArray = [NSPointerArray weakObjectsPointerArray];
+        objc_setAssociatedObject(self, &animationLoadOperationKey, operationsArray, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+        return operationsArray;
+    }
 }
 
 - (void)sd_cancelCurrentAnimationImagesLoad {
-    [self sd_cancelImageLoadOperationWithKey:@"UIImageViewAnimationImages"];
+    NSPointerArray *operationsArray = [self sd_animationOperationArray];
+    if (operationsArray) {
+        @synchronized (self) {
+            for (id operation in operationsArray) {
+                if ([operation conformsToProtocol:@protocol(SDWebImageOperation)]) {
+                    [operation cancel];
+                }
+            }
+            operationsArray.count = 0;
+        }
+    }
 }
 #endif
 

SDWebImage/UIView+WebCacheOperation.h

@@ -12,15 +12,17 @@
 
 #import "SDWebImageManager.h"
 
+// These methods are used to support canceling for UIView image loading, it's designed to be used internal but not external.
+// All the stored operations are weak, so it will be dalloced after image loading finished. If you need to store operations, use your own class to keep a strong reference for them.
 @interface UIView (WebCacheOperation)
 
 /**
- *  Set the image load operation (storage in a UIView based dictionary)
+ *  Set the image load operation (storage in a UIView based weak map table)
  *
  *  @param operation the operation
  *  @param key       key for storing the operation
  */
-- (void)sd_setImageLoadOperation:(nullable id)operation forKey:(nullable NSString *)key;
+- (void)sd_setImageLoadOperation:(nullable id<SDWebImageOperation>)operation forKey:(nullable NSString *)key;
 
 /**
  *  Cancel all operations for the current UIView and key

SDWebImage/UIView+WebCacheOperation.m

@@ -14,52 +14,59 @@
 
 static char loadOperationKey;
 
-typedef NSMutableDictionary<NSString *, id> SDOperationsDictionary;
+// key is copy, value is weak because operation instance is retained by SDWebImageManager's runningOperations property
+// we should use lock to keep thread-safe because these method may not be acessed from main queue
+typedef NSMapTable<NSString *, id<SDWebImageOperation>> SDOperationsDictionary;
 
 @implementation UIView (WebCacheOperation)
 
-- (SDOperationsDictionary *)operationDictionary {
-    SDOperationsDictionary *operations = objc_getAssociatedObject(self, &loadOperationKey);
-    if (operations) {
+- (SDOperationsDictionary *)sd_operationDictionary {
+    @synchronized(self) {
+        SDOperationsDictionary *operations = objc_getAssociatedObject(self, &loadOperationKey);
+        if (operations) {
+            return operations;
+        }
+        operations = [[NSMapTable alloc] initWithKeyOptions:NSMapTableCopyIn valueOptions:NSMapTableWeakMemory capacity:0];
+        objc_setAssociatedObject(self, &loadOperationKey, operations, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
         return operations;
     }
-    operations = [NSMutableDictionary dictionary];
-    objc_setAssociatedObject(self, &loadOperationKey, operations, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
-    return operations;
 }
 
-- (void)sd_setImageLoadOperation:(nullable id)operation forKey:(nullable NSString *)key {
+- (void)sd_setImageLoadOperation:(nullable id<SDWebImageOperation>)operation forKey:(nullable NSString *)key {
     if (key) {
         [self sd_cancelImageLoadOperationWithKey:key];
         if (operation) {
-            SDOperationsDictionary *operationDictionary = [self operationDictionary];
-            operationDictionary[key] = operation;
+            SDOperationsDictionary *operationDictionary = [self sd_operationDictionary];
+            @synchronized (self) {
+                [operationDictionary setObject:operation forKey:key];
+            }
         }
     }
 }
 
 - (void)sd_cancelImageLoadOperationWithKey:(nullable NSString *)key {
     // Cancel in progress downloader from queue
-    SDOperationsDictionary *operationDictionary = [self operationDictionary];
-    id operations = operationDictionary[key];
-    if (operations) {
-        if ([operations isKindOfClass:[NSArray class]]) {
-            for (id <SDWebImageOperation> operation in operations) {
-                if (operation) {
-                    [operation cancel];
-                }
-            }
-        } else if ([operations conformsToProtocol:@protocol(SDWebImageOperation)]){
-            [(id<SDWebImageOperation>) operations cancel];
+    SDOperationsDictionary *operationDictionary = [self sd_operationDictionary];
+    id<SDWebImageOperation> operation;
+    @synchronized (self) {
+        operation = [operationDictionary objectForKey:key];
+    }
+    if (operation) {
+        if ([operation conformsToProtocol:@protocol(SDWebImageOperation)]){
+            [operation cancel];
+        }
+        @synchronized (self) {
+            [operationDictionary removeObjectForKey:key];
         }
-        [operationDictionary removeObjectForKey:key];
     }
 }
 
 - (void)sd_removeImageLoadOperationWithKey:(nullable NSString *)key {
     if (key) {
-        SDOperationsDictionary *operationDictionary = [self operationDictionary];
-        [operationDictionary removeObjectForKey:key];
+        SDOperationsDictionary *operationDictionary = [self sd_operationDictionary];
+        @synchronized (self) {
+            [operationDictionary removeObjectForKey:key];
+        }
     }
 }