Move progress / completion handler logic to the operation

SDWebImage/SDWebImageDownloader.m

@@ -10,13 +10,10 @@
 #import "SDWebImageDownloaderOperation.h"
 #import <ImageIO/ImageIO.h>
 
-static NSString *const kProgressCallbackKey = @"progress";
-static NSString *const kCompletedCallbackKey = @"completed";
-
 @interface _SDWebImageDownloaderToken : NSObject
 
 @property (nonatomic, strong) NSURL *url;
-@property (nonatomic, strong) id callbacks;
+@property (nonatomic, strong) id downloadOperationCancelToken;
 
 @end
 
@@ -28,7 +25,6 @@ static NSString *const kCompletedCallbackKey = @"completed";
 @property (strong, nonatomic) NSOperationQueue *downloadQueue;
 @property (weak, nonatomic) NSOperation *lastAddedOperation;
 @property (assign, nonatomic) Class operationClass;
-@property (strong, nonatomic) NSMutableDictionary *URLCallbacks;
 @property (strong, nonatomic) NSMutableDictionary *URLOperations;
 @property (strong, nonatomic) NSMutableDictionary *HTTPHeaders;
 // This queue is used to serialize the handling of the network responses of all the download operation in a single queue
@@ -77,7 +73,6 @@ static NSString *const kCompletedCallbackKey = @"completed";
         _executionOrder = SDWebImageDownloaderFIFOExecutionOrder;
         _downloadQueue = [NSOperationQueue new];
         _downloadQueue.maxConcurrentOperationCount = 6;
-        _URLCallbacks = [NSMutableDictionary new];
         _URLOperations = [NSMutableDictionary new];
 #ifdef SD_WEBP
         _HTTPHeaders = [@{@"Accept": @"image/webp,image/*;q=0.8"} mutableCopy];
@@ -125,7 +120,6 @@ static NSString *const kCompletedCallbackKey = @"completed";
 }
 
 - (id)downloadImageWithURL:(NSURL *)url options:(SDWebImageDownloaderOptions)options progress:(SDWebImageDownloaderProgressBlock)progressBlock completed:(SDWebImageDownloaderCompletedBlock)completedBlock {
-    __block SDWebImageDownloaderOperation *operation;
     __weak SDWebImageDownloader *wself = self;
 
     return [self addProgressCallback:progressBlock completedBlock:completedBlock forURL:url createCallback:^SDWebImageDownloaderOperation *{
@@ -144,44 +138,7 @@ static NSString *const kCompletedCallbackKey = @"completed";
         else {
             request.allHTTPHeaderFields = wself.HTTPHeaders;
         }
-        operation = [[wself.operationClass alloc] initWithRequest:request
+        SDWebImageDownloaderOperation *operation = [[wself.operationClass alloc] initWithRequest:request options:options];
-                                                          options:options
-                                                         progress:^(NSInteger receivedSize, NSInteger expectedSize) {
-                                                             SDWebImageDownloader *sself = wself;
-                                                             if (!sself) return;
-                                                             __block NSArray *callbacksForURL;
-                                                             dispatch_sync(sself.barrierQueue, ^{
-                                                                 callbacksForURL = [sself.URLCallbacks[url] copy];
-                                                             });
-                                                             for (NSDictionary *callbacks in callbacksForURL) {
-                                                                 dispatch_async(dispatch_get_main_queue(), ^{
-                                                                     SDWebImageDownloaderProgressBlock callback = callbacks[kProgressCallbackKey];
-                                                                     if (callback) callback(receivedSize, expectedSize);
-                                                                 });
-                                                             }
-                                                         }
-                                                        completed:^(UIImage *image, NSData *data, NSError *error, BOOL finished) {
-                                                            SDWebImageDownloader *sself = wself;
-                                                            if (!sself) return;
-                                                            __block NSArray *callbacksForURL;
-                                                            dispatch_barrier_sync(sself.barrierQueue, ^{
-                                                                callbacksForURL = [sself.URLCallbacks[url] copy];
-                                                                if (finished) {
-                                                                    [sself.URLCallbacks removeObjectForKey:url];
-                                                                }
-                                                            });
-                                                            for (NSDictionary *callbacks in callbacksForURL) {
-                                                                SDWebImageDownloaderCompletedBlock callback = callbacks[kCompletedCallbackKey];
-                                                                if (callback) callback(image, data, error, finished);
-                                                            }
-                                                        }
-                                                        cancelled:^{
-                                                            SDWebImageDownloader *sself = wself;
-                                                            if (!sself) return;
-                                                            dispatch_barrier_async(sself.barrierQueue, ^{
-                                                                [sself.URLCallbacks removeObjectForKey:url];
-                                                            });
-                                                        }];
         operation.shouldDecompressImages = wself.shouldDecompressImages;
         
         if (wself.urlCredential) {
@@ -212,18 +169,12 @@ static NSString *const kCompletedCallbackKey = @"completed";
         return;
     }
 
-    dispatch_barrier_async(self.barrierQueue, ^{
+    _SDWebImageDownloaderToken *typedToken = (_SDWebImageDownloaderToken *)token;
-        _SDWebImageDownloaderToken *typedToken = (_SDWebImageDownloaderToken *)token;
+    SDWebImageDownloaderOperation *operation = self.URLOperations[typedToken.url];
-
+    BOOL canceled = [operation cancel:typedToken.downloadOperationCancelToken];
-        NSMutableArray *callbacksForURL = self.URLCallbacks[typedToken.url];
+    if (canceled) {
-        [callbacksForURL removeObjectIdenticalTo:typedToken.callbacks];
+        [self.URLOperations removeObjectForKey:typedToken.url];
-
+    }
-        // If this was the last set of callbacks, then cancel the operation
-        if (callbacksForURL.count == 0) {
-            SDWebImageDownloaderOperation *operation = self.URLOperations[typedToken.url];
-            [operation cancel];
-        }
-    });
 }
 
 - (id)addProgressCallback:(SDWebImageDownloaderProgressBlock)progressBlock completedBlock:(SDWebImageDownloaderCompletedBlock)completedBlock forURL:(NSURL *)url createCallback:(SDWebImageDownloaderOperation *(^)())createCallback {
@@ -235,31 +186,25 @@ static NSString *const kCompletedCallbackKey = @"completed";
         return nil;
     }
 
-    __block _SDWebImageDownloaderToken *token = nil;
+    SDWebImageDownloaderOperation *operation = self.URLOperations[url];
+    if (!operation) {
+        operation = createCallback();
+        self.URLOperations[url] = operation;
 
-    dispatch_barrier_sync(self.barrierQueue, ^{
+        __weak SDWebImageDownloaderOperation *woperation = operation;
-        if (!self.URLCallbacks[url]) {
+        operation.completionBlock = ^{
-            self.URLCallbacks[url] = [NSMutableArray new];
+          SDWebImageDownloaderOperation *soperation = woperation;
-        }
+          if (!soperation) return;
+          if (self.URLOperations[url] == soperation) {
+              [self.URLOperations removeObjectForKey:url];
+          };
+        };
+    }
+    id downloadOperationCancelToken = [operation addHandlersForProgress:progressBlock completed:completedBlock];
 
-        // Handle single download of simultaneous download request for the same URL
+    _SDWebImageDownloaderToken *token = [_SDWebImageDownloaderToken new];
-        NSMutableArray *callbacksForURL = self.URLCallbacks[url];
+    token.url = url;
-        NSMutableDictionary *callbacks = [NSMutableDictionary new];
+    token.downloadOperationCancelToken = downloadOperationCancelToken;
-        if (progressBlock) callbacks[kProgressCallbackKey] = [progressBlock copy];
-        if (completedBlock) callbacks[kCompletedCallbackKey] = [completedBlock copy];
-        [callbacksForURL addObject:callbacks];
-        self.URLCallbacks[url] = callbacksForURL;
-
-        SDWebImageDownloaderOperation *operation = self.URLOperations[url];
-        if (!operation) {
-            operation = createCallback();
-            self.URLOperations[url] = operation;
-        }
-
-        token = [_SDWebImageDownloaderToken new];
-        token.url = url;
-        token.callbacks = callbacks;
-    });
 
     return token;
 }

SDWebImage/SDWebImageDownloaderOperation.h

@@ -61,18 +61,33 @@ extern NSString *const SDWebImageDownloadFinishNotification;
  *
  *  @param request        the URL request
  *  @param options        downloader options
- *  @param progressBlock  the block executed when a new chunk of data arrives. 
- *                        @note the progress block is executed on a background queue
- *  @param completedBlock the block executed when the download is done. 
- *                        @note the completed block is executed on the main queue for success. If errors are found, there is a chance the block will be executed on a background queue
- *  @param cancelBlock    the block executed if the download (operation) is cancelled
  *
  *  @return the initialized instance
  */
 - (id)initWithRequest:(NSURLRequest *)request
-              options:(SDWebImageDownloaderOptions)options
+              options:(SDWebImageDownloaderOptions)options;
-             progress:(SDWebImageDownloaderProgressBlock)progressBlock
+
-            completed:(SDWebImageDownloaderCompletedBlock)completedBlock
+/**
-            cancelled:(SDWebImageNoParamsBlock)cancelBlock;
+ *  Adds handlers for progress and completion. Returns a tokent that can be passed to -cancel: to cancel this set of
+ *  callbacks.
+ *
+ *  @param progressBlock  the block executed when a new chunk of data arrives.
+ *                        @note the progress block is executed on a background queue
+ *  @param completedBlock the block executed when the download is done.
+ *                        @note the completed block is executed on the main queue for success. If errors are found, there is a chance the block will be executed on a background queue
+ *
+ *  @return the token to use to cancel this set of handlers
+ */
+- (id)addHandlersForProgress:(SDWebImageDownloaderProgressBlock)progressBlock
+                   completed:(SDWebImageDownloaderCompletedBlock)completedBlock;
+
+/**
+ *  Cancels a set of callbacks. Once all callbacks are canceled, the operation is cancelled.
+ *
+ *  @param token the token representing a set of callbacks to cancel
+ *
+ *  @return YES if the operation was stopped because this was the last token to be canceled. NO otherwise.
+ */
+- (BOOL)cancel:(id)token;
 
 @end

SDWebImage/SDWebImageDownloaderOperation.m

@@ -17,11 +17,12 @@ NSString *const SDWebImageDownloadReceiveResponseNotification = @"SDWebImageDown
 NSString *const SDWebImageDownloadStopNotification = @"SDWebImageDownloadStopNotification";
 NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinishNotification";
 
+static NSString *const kProgressCallbackKey = @"progress";
+static NSString *const kCompletedCallbackKey = @"completed";
+
 @interface SDWebImageDownloaderOperation () <NSURLConnectionDataDelegate>
 
-@property (copy, nonatomic) SDWebImageDownloaderProgressBlock progressBlock;
+@property (strong, nonatomic) NSMutableArray *callbackBlocks;
-@property (copy, nonatomic) SDWebImageDownloaderCompletedBlock completedBlock;
-@property (copy, nonatomic) SDWebImageNoParamsBlock cancelBlock;
 
 @property (assign, nonatomic, getter = isExecuting) BOOL executing;
 @property (assign, nonatomic, getter = isFinished) BOOL finished;
@@ -45,18 +46,13 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
 @synthesize finished = _finished;
 
 - (id)initWithRequest:(NSURLRequest *)request
-              options:(SDWebImageDownloaderOptions)options
+              options:(SDWebImageDownloaderOptions)options {
-             progress:(SDWebImageDownloaderProgressBlock)progressBlock
-            completed:(SDWebImageDownloaderCompletedBlock)completedBlock
-            cancelled:(SDWebImageNoParamsBlock)cancelBlock {
     if ((self = [super init])) {
         _request = request;
         _shouldDecompressImages = YES;
         _shouldUseCredentialStorage = YES;
         _options = options;
-        _progressBlock = [progressBlock copy];
+        _callbackBlocks = [NSMutableArray new];
-        _completedBlock = [completedBlock copy];
-        _cancelBlock = [cancelBlock copy];
         _executing = NO;
         _finished = NO;
         _expectedSize = 0;
@@ -65,6 +61,24 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
     return self;
 }
 
+- (id)addHandlersForProgress:(SDWebImageDownloaderProgressBlock)progressBlock
+                   completed:(SDWebImageDownloaderCompletedBlock)completedBlock {
+    NSMutableDictionary *callbacks = [NSMutableDictionary new];
+    if (progressBlock) callbacks[kProgressCallbackKey] = [progressBlock copy];
+    if (completedBlock) callbacks[kCompletedCallbackKey] = [completedBlock copy];
+    [self.callbackBlocks addObject:callbacks];
+    return callbacks;
+}
+
+- (BOOL)cancel:(id)token {
+    [self.callbackBlocks removeObjectIdenticalTo:token];
+    if (self.callbackBlocks.count == 0) {
+        [self cancel];
+        return YES;
+    }
+    return NO;
+}
+
 - (void)start {
     @synchronized (self) {
         if (self.isCancelled) {
@@ -100,8 +114,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
     [self.connection start];
 
     if (self.connection) {
-        if (self.progressBlock) {
+        for (SDWebImageDownloaderProgressBlock progressBlock in [self.callbackBlocks valueForKey:kProgressCallbackKey]) {
-            self.progressBlock(0, NSURLResponseUnknownLength);
+            progressBlock(0, NSURLResponseUnknownLength);
         }
         dispatch_async(dispatch_get_main_queue(), ^{
             [[NSNotificationCenter defaultCenter] postNotificationName:SDWebImageDownloadStartNotification object:self];
@@ -123,8 +137,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
         }
     }
     else {
-        if (self.completedBlock) {
+        for (SDWebImageDownloaderCompletedBlock completedBlock in [self.callbackBlocks valueForKey:kCompletedCallbackKey]) {
-            self.completedBlock(nil, nil, [NSError errorWithDomain:NSURLErrorDomain code:0 userInfo:@{NSLocalizedDescriptionKey : @"Connection can't be initialized"}], YES);
+            completedBlock(nil, nil, [NSError errorWithDomain:NSURLErrorDomain code:0 userInfo:@{NSLocalizedDescriptionKey : @"Connection can't be initialized"}], YES);
         }
     }
 
@@ -161,7 +175,6 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
 - (void)cancelInternal {
     if (self.isFinished) return;
     [super cancel];
-    if (self.cancelBlock) self.cancelBlock();
 
     if (self.connection) {
         [self.connection cancel];
@@ -185,9 +198,7 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
 }
 
 - (void)reset {
-    self.cancelBlock = nil;
+    [self.callbackBlocks removeAllObjects];
-    self.completedBlock = nil;
-    self.progressBlock = nil;
     self.connection = nil;
     self.imageData = nil;
     self.thread = nil;
@@ -217,8 +228,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
     if (![response respondsToSelector:@selector(statusCode)] || ([((NSHTTPURLResponse *)response) statusCode] < 400 && [((NSHTTPURLResponse *)response) statusCode] != 304)) {
         NSInteger expected = response.expectedContentLength > 0 ? (NSInteger)response.expectedContentLength : 0;
         self.expectedSize = expected;
-        if (self.progressBlock) {
+        for (SDWebImageDownloaderProgressBlock progressBlock in [self.callbackBlocks valueForKey:kProgressCallbackKey]) {
-            self.progressBlock(0, expected);
+            progressBlock(0, expected);
         }
 
         self.imageData = [[NSMutableData alloc] initWithCapacity:expected];
@@ -241,8 +252,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
             [[NSNotificationCenter defaultCenter] postNotificationName:SDWebImageDownloadStopNotification object:self];
         });
 
-        if (self.completedBlock) {
+        for (SDWebImageDownloaderCompletedBlock completedBlock in [self.callbackBlocks valueForKey:kCompletedCallbackKey]) {
-            self.completedBlock(nil, nil, [NSError errorWithDomain:NSURLErrorDomain code:[((NSHTTPURLResponse *)response) statusCode] userInfo:nil], YES);
+            completedBlock(nil, nil, [NSError errorWithDomain:NSURLErrorDomain code:[((NSHTTPURLResponse *)response) statusCode] userInfo:nil], YES);
         }
         CFRunLoopStop(CFRunLoopGetCurrent());
         [self done];
@@ -252,7 +263,7 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
 - (void)connection:(NSURLConnection *)connection didReceiveData:(NSData *)data {
     [self.imageData appendData:data];
 
-    if ((self.options & SDWebImageDownloaderProgressiveDownload) && self.expectedSize > 0 && self.completedBlock) {
+    if ((self.options & SDWebImageDownloaderProgressiveDownload) && self.expectedSize > 0) {
         // The following code is from http://www.cocoaintheshell.com/2011/05/progressive-images-download-imageio/
         // Thanks to the author @Nyx0uf
 
@@ -319,8 +330,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
                 }
                 CGImageRelease(partialImageRef);
                 dispatch_main_sync_safe(^{
-                    if (self.completedBlock) {
+                    for (SDWebImageDownloaderCompletedBlock completedBlock in [self.callbackBlocks valueForKey:kCompletedCallbackKey]) {
-                        self.completedBlock(image, nil, nil, NO);
+                        completedBlock(image, nil, nil, NO);
                     }
                 });
             }
@@ -329,8 +340,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
         CFRelease(imageSource);
     }
 
-    if (self.progressBlock) {
+    for (SDWebImageDownloaderProgressBlock progressBlock in [self.callbackBlocks valueForKey:kProgressCallbackKey]) {
-        self.progressBlock(self.imageData.length, self.expectedSize);
+        progressBlock(self.imageData.length, self.expectedSize);
     }
 }
 
@@ -362,7 +373,7 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
 }
 
 - (void)connectionDidFinishLoading:(NSURLConnection *)aConnection {
-    SDWebImageDownloaderCompletedBlock completionBlock = self.completedBlock;
+    NSArray *completionBlocks = [[self.callbackBlocks valueForKey:kCompletedCallbackKey] copy];
     @synchronized(self) {
         CFRunLoopStop(CFRunLoopGetCurrent());
         self.thread = nil;
@@ -377,7 +388,7 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
         responseFromCached = NO;
     }
     
-    if (completionBlock) {
+    for (SDWebImageDownloaderCompletedBlock completionBlock in completionBlocks) {
         if (self.options & SDWebImageDownloaderIgnoreCachedResponse && responseFromCached) {
             completionBlock(nil, nil, nil, YES);
         } else if (self.imageData) {
@@ -401,7 +412,6 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
             completionBlock(nil, nil, [NSError errorWithDomain:SDWebImageErrorDomain code:0 userInfo:@{NSLocalizedDescriptionKey : @"Image data is nil"}], YES);
         }
     }
-    self.completionBlock = nil;
     [self done];
 }
 
@@ -415,8 +425,8 @@ NSString *const SDWebImageDownloadFinishNotification = @"SDWebImageDownloadFinis
         });
     }
 
-    if (self.completedBlock) {
+    for (SDWebImageDownloaderCompletedBlock completedBlock in [self.callbackBlocks valueForKey:kCompletedCallbackKey]) {
-        self.completedBlock(nil, nil, error, YES);
+        completedBlock(nil, nil, error, YES);
     }
     self.completionBlock = nil;
     [self done];

Tests/Tests/SDWebImageDownloaderTests.m

@@ -32,6 +32,20 @@
     [super tearDown];
 }
 
+- (BOOL)spinRunLoopWithTimeout:(NSTimeInterval)timeout untilBlockIsTrue:(BOOL(^)())block {
+  CFTimeInterval timeoutDate = CACurrentMediaTime() + 5.;
+  while (true) {
+      if (block()) {
+          return YES;
+      }
+      if (CACurrentMediaTime() > timeoutDate) {
+          return NO;
+      }
+      CFRunLoopRunInMode(kCFRunLoopDefaultMode, 1., true);
+  }
+  return NO;
+}
+
 - (void)testThatDownloadingSameURLTwiceAndCancellingFirstWorks {
     NSURL *imageURL = [NSURL URLWithString:@"http://static2.dmcdn.net/static/video/656/177/44771656:jpeg_preview_small.jpg?20120509154705"];
 
@@ -54,13 +68,40 @@
 
     [[SDWebImageDownloader sharedDownloader] cancel:token1];
 
-    CFTimeInterval timeoutDate = CACurrentMediaTime() + 5.;
+    success = [self spinRunLoopWithTimeout:5. untilBlockIsTrue:^BOOL{
-    while (true) {
+        return success;
-        if (CACurrentMediaTime() > timeoutDate || success) {
+    }];
-            break;
+
-        }
+    if (!success) {
-        CFRunLoopRunInMode(kCFRunLoopDefaultMode, 1., true);
+        XCTFail(@"Failed to download image");
     }
+}
+
+- (void)testThatCancelingDownloadThenRequestingAgainWorks {
+    NSURL *imageURL = [NSURL URLWithString:@"http://static2.dmcdn.net/static/video/656/177/44771656:jpeg_preview_small.jpg?20120509154705"];
+
+    id token1 = [[SDWebImageDownloader sharedDownloader] downloadImageWithURL:imageURL
+                                                                      options:0
+                                                                     progress:nil
+                                                                    completed:^(UIImage *image, NSData *data, NSError *error, BOOL finished) {
+                                                                        XCTFail(@"Shouldn't have completed here.");
+                                                                    }];
+    expect(token1).toNot.beNil();
+
+    [[SDWebImageDownloader sharedDownloader] cancel:token1];
+
+    __block BOOL success = NO;
+    id token2 = [[SDWebImageDownloader sharedDownloader] downloadImageWithURL:imageURL
+                                                                      options:0
+                                                                     progress:nil
+                                                                    completed:^(UIImage *image, NSData *data, NSError *error, BOOL finished) {
+                                                                        success = YES;
+                                                                    }];
+    expect(token2).toNot.beNil();
+
+    success = [self spinRunLoopWithTimeout:5. untilBlockIsTrue:^BOOL{
+        return success;
+    }];
 
     if (!success) {
         XCTFail(@"Failed to download image");