Thread safe creation and canceling of download tasks

fixes #364

AsyncDisplayKit.xcodeproj/project.pbxproj

@@ -145,11 +145,13 @@
 		292C59A21A956527007E5DD6 /* ASRangeHandler.h in Headers */ = {isa = PBXBuildFile; fileRef = 292C599C1A956527007E5DD6 /* ASRangeHandler.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		292C59A31A956527007E5DD6 /* ASRangeHandlerRender.h in Headers */ = {isa = PBXBuildFile; fileRef = 292C599D1A956527007E5DD6 /* ASRangeHandlerRender.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		292C59A41A956527007E5DD6 /* ASRangeHandlerRender.mm in Sources */ = {isa = PBXBuildFile; fileRef = 292C599E1A956527007E5DD6 /* ASRangeHandlerRender.mm */; };
+		2967F9E21AB0A5190072E4AB /* ASBasicImageDownloaderInternal.h in Headers */ = {isa = PBXBuildFile; fileRef = 2967F9E11AB0A4CF0072E4AB /* ASBasicImageDownloaderInternal.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		296A0A2E1A9516B2005ACEAA /* ASBatchFetching.h in Headers */ = {isa = PBXBuildFile; fileRef = 296A0A2C1A9516B2005ACEAA /* ASBatchFetching.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		296A0A2F1A9516B2005ACEAA /* ASBatchFetching.m in Sources */ = {isa = PBXBuildFile; fileRef = 296A0A2D1A9516B2005ACEAA /* ASBatchFetching.m */; };
 		296A0A351A951ABF005ACEAA /* ASBatchFetchingTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 296A0A341A951ABF005ACEAA /* ASBatchFetchingTests.m */; };
 		299DA1A91A828D2900162D41 /* ASBatchContext.h in Headers */ = {isa = PBXBuildFile; fileRef = 299DA1A71A828D2900162D41 /* ASBatchContext.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		299DA1AA1A828D2900162D41 /* ASBatchContext.mm in Sources */ = {isa = PBXBuildFile; fileRef = 299DA1A81A828D2900162D41 /* ASBatchContext.mm */; };
+		29CDC2E21AAE70D000833CA4 /* ASBasicImageDownloaderContextTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 29CDC2E11AAE70D000833CA4 /* ASBasicImageDownloaderContextTests.m */; };
 		3C9C128519E616EF00E942A0 /* ASTableViewTests.m in Sources */ = {isa = PBXBuildFile; fileRef = 3C9C128419E616EF00E942A0 /* ASTableViewTests.m */; settings = {COMPILER_FLAGS = "-fno-objc-arc"; }; };
 		464052201A3F83C40061C0BA /* ASDataController.h in Headers */ = {isa = PBXBuildFile; fileRef = 464052191A3F83C40061C0BA /* ASDataController.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		464052211A3F83C40061C0BA /* ASDataController.mm in Sources */ = {isa = PBXBuildFile; fileRef = 4640521A1A3F83C40061C0BA /* ASDataController.mm */; };
@@ -302,12 +304,14 @@
 		292C599C1A956527007E5DD6 /* ASRangeHandler.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASRangeHandler.h; sourceTree = "<group>"; };
 		292C599D1A956527007E5DD6 /* ASRangeHandlerRender.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASRangeHandlerRender.h; sourceTree = "<group>"; };
 		292C599E1A956527007E5DD6 /* ASRangeHandlerRender.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASRangeHandlerRender.mm; sourceTree = "<group>"; };
+		2967F9E11AB0A4CF0072E4AB /* ASBasicImageDownloaderInternal.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = ASBasicImageDownloaderInternal.h; sourceTree = "<group>"; };
 		296A0A2C1A9516B2005ACEAA /* ASBatchFetching.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = ASBatchFetching.h; path = ../Details/ASBatchFetching.h; sourceTree = "<group>"; };
 		296A0A2D1A9516B2005ACEAA /* ASBatchFetching.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; name = ASBatchFetching.m; path = ../Details/ASBatchFetching.m; sourceTree = "<group>"; };
 		296A0A311A951715005ACEAA /* ASScrollDirection.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = ASScrollDirection.h; path = AsyncDisplayKit/Details/ASScrollDirection.h; sourceTree = SOURCE_ROOT; };
 		296A0A341A951ABF005ACEAA /* ASBatchFetchingTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASBatchFetchingTests.m; sourceTree = "<group>"; };
 		299DA1A71A828D2900162D41 /* ASBatchContext.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASBatchContext.h; sourceTree = "<group>"; };
 		299DA1A81A828D2900162D41 /* ASBatchContext.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASBatchContext.mm; sourceTree = "<group>"; };
+		29CDC2E11AAE70D000833CA4 /* ASBasicImageDownloaderContextTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASBasicImageDownloaderContextTests.m; sourceTree = "<group>"; };
 		3C9C128419E616EF00E942A0 /* ASTableViewTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = ASTableViewTests.m; sourceTree = "<group>"; };
 		464052191A3F83C40061C0BA /* ASDataController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = ASDataController.h; sourceTree = "<group>"; };
 		4640521A1A3F83C40061C0BA /* ASDataController.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = ASDataController.mm; sourceTree = "<group>"; };
@@ -446,6 +450,7 @@
 				058D0A2F195D057000B7D73C /* ASDisplayNodeTests.m */,
 				058D0A30195D057000B7D73C /* ASDisplayNodeTestsHelper.h */,
 				058D0A31195D057000B7D73C /* ASDisplayNodeTestsHelper.m */,
+				29CDC2E11AAE70D000833CA4 /* ASBasicImageDownloaderContextTests.m */,
 				052EE0651A159FEF002C6279 /* ASMultiplexImageNodeTests.m */,
 				058D0A32195D057000B7D73C /* ASMutableAttributedStringBuilderTests.m */,
 				3C9C128419E616EF00E942A0 /* ASTableViewTests.m */,
@@ -495,11 +500,11 @@
 				4640521F1A3F83C40061C0BA /* ASMultidimensionalArrayUtils.mm */,
 				058D09E8195D050800B7D73C /* ASMutableAttributedStringBuilder.h */,
 				058D09E9195D050800B7D73C /* ASMutableAttributedStringBuilder.m */,
-				292C599A1A956527007E5DD6 /* ASRangeHandlerPreload.h */,
-				292C599B1A956527007E5DD6 /* ASRangeHandlerPreload.mm */,
 				055F1A3619ABD413004DAFF1 /* ASRangeController.h */,
 				055F1A3719ABD413004DAFF1 /* ASRangeController.mm */,
 				292C599C1A956527007E5DD6 /* ASRangeHandler.h */,
+				292C599A1A956527007E5DD6 /* ASRangeHandlerPreload.h */,
+				292C599B1A956527007E5DD6 /* ASRangeHandlerPreload.mm */,
 				292C599D1A956527007E5DD6 /* ASRangeHandlerRender.h */,
 				292C599E1A956527007E5DD6 /* ASRangeHandlerRender.mm */,
 				296A0A311A951715005ACEAA /* ASScrollDirection.h */,
@@ -542,6 +547,7 @@
 			isa = PBXGroup;
 			children = (
 				296A0A2C1A9516B2005ACEAA /* ASBatchFetching.h */,
+				2967F9E11AB0A4CF0072E4AB /* ASBasicImageDownloaderInternal.h */,
 				296A0A2D1A9516B2005ACEAA /* ASBatchFetching.m */,
 				058D0A02195D050800B7D73C /* _AS-objc-internal.h */,
 				058D0A03195D050800B7D73C /* _ASCoreAnimationExtras.h */,
@@ -673,6 +679,7 @@
 				058D0A78195D05F900B7D73C /* ASDisplayNode+DebugTiming.h in Headers */,
 				058D0A79195D05F900B7D73C /* ASDisplayNode+DebugTiming.mm in Headers */,
 				058D0A7A195D05F900B7D73C /* ASDisplayNode+UIViewBridge.mm in Headers */,
+				2967F9E21AB0A5190072E4AB /* ASBasicImageDownloaderInternal.h in Headers */,
 				058D0A7B195D05F900B7D73C /* ASDisplayNodeInternal.h in Headers */,
 				058D0A7C195D05F900B7D73C /* ASImageNode+CGExtras.h in Headers */,
 				058D0A7D195D05F900B7D73C /* ASImageNode+CGExtras.m in Headers */,
@@ -857,6 +864,7 @@
 				058D0A3D195D057000B7D73C /* ASTextNodeCoreTextAdditionsTests.m in Sources */,
 				058D0A3C195D057000B7D73C /* ASMutableAttributedStringBuilderTests.m in Sources */,
 				058D0A3F195D057000B7D73C /* ASTextNodeShadowerTests.m in Sources */,
+				29CDC2E21AAE70D000833CA4 /* ASBasicImageDownloaderContextTests.m in Sources */,
 				058D0A3B195D057000B7D73C /* ASDisplayNodeTestsHelper.m in Sources */,
 				058D0A3A195D057000B7D73C /* ASDisplayNodeTests.m in Sources */,
 				052EE0661A159FEF002C6279 /* ASMultiplexImageNodeTests.m in Sources */,

AsyncDisplayKit/ASNetworkImageNode.mm

@@ -241,17 +241,11 @@
           }
         };
 
-        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
         [_cache fetchCachedImageWithURL:_URL
                           callbackQueue:dispatch_get_main_queue()
                              completion:cacheCompletion];
-        });
       } else {
-          // NSURLSessionDownloadTask will do file I/O to create a temp directory. If called on the main thread this
-          // will cause significant performance issues.
-          dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
         [self _downloadImageWithCompletion:finished];
-          });
       }
     }
   }

AsyncDisplayKit/Details/ASBasicImageDownloader.mm

@@ -12,6 +12,7 @@
 
 #import <UIKit/UIKit.h>
 
+#import "ASBasicImageDownloaderInternal.h"
 #import "ASThread.h"
 
 
@@ -19,13 +20,73 @@
 /**
  * Collection of properties associated with a download request.
  */
-@interface ASBasicImageDownloaderMetadata : NSObject
+@interface ASBasicImageDownloaderContext ()
+{
+  BOOL _invalid;
+  ASDN::RecursiveMutex _propertyLock;
+}
+
 @property (nonatomic, strong) dispatch_queue_t callbackQueue;
 @property (nonatomic, copy) void (^downloadProgressBlock)(CGFloat);
 @property (nonatomic, copy) void (^completionBlock)(CGImageRef, NSError *);
+
 @end
 
-@implementation ASBasicImageDownloaderMetadata
+@implementation ASBasicImageDownloaderContext
+
+static NSMutableDictionary *currentRequests = nil;
+static ASDN::RecursiveMutex currentRequestsLock;
+
++ (ASBasicImageDownloaderContext *)contextForURL:(NSURL *)URL
+{
+  ASDN::MutexLocker l(currentRequestsLock);
+  if (!currentRequests) {
+    currentRequests = [[NSMutableDictionary alloc] init];
+  }
+  ASBasicImageDownloaderContext *context = currentRequests[URL];
+  if (!context) {
+    context = [[ASBasicImageDownloaderContext alloc] initWithURL:URL];
+    currentRequests[URL] = context;
+  }
+  return context;
+}
+
++ (void)cancelContextWithURL:(NSURL *)URL
+{
+  ASDN::MutexLocker l(currentRequestsLock);
+  if (currentRequests) {
+    [currentRequests removeObjectForKey:URL];
+  }
+}
+
+- (instancetype)initWithURL:(NSURL *)URL
+{
+  if (self = [super init]) {
+    _URL = URL;
+  }
+  return self;
+}
+
+- (void)cancel
+{
+  ASDN::MutexLocker l(_propertyLock);
+
+  NSURLSessionTask *sessionTask = self.sessionTask;
+  if (sessionTask) {
+    [sessionTask cancel];
+    self.sessionTask = nil;
+  }
+
+  _invalid = YES;
+  [self.class cancelContextWithURL:self.URL];
+}
+
+- (BOOL)isCancelled
+{
+  ASDN::MutexLocker l(_propertyLock);
+  return _invalid;
+}
+
 @end
 
 
@@ -34,18 +95,18 @@
  * NSURLSessionDownloadTask lacks a `userInfo` property, so add this association ourselves.
  */
 @interface NSURLRequest (ASBasicImageDownloader)
-@property (nonatomic, strong) ASBasicImageDownloaderMetadata *asyncdisplaykit_metadata;
+@property (nonatomic, strong) ASBasicImageDownloaderContext *asyncdisplaykit_context;
 @end
 
 @implementation NSURLRequest (ASBasicImageDownloader)
-static const char *kMetadataKey = NSStringFromClass(ASBasicImageDownloaderMetadata.class).UTF8String;
+static const char *kContextKey = NSStringFromClass(ASBasicImageDownloaderContext.class).UTF8String;
-- (void)setAsyncdisplaykit_metadata:(ASBasicImageDownloaderMetadata *)asyncdisplaykit_metadata
+- (void)setAsyncdisplaykit_context:(ASBasicImageDownloaderContext *)asyncdisplaykit_context
 {
-  objc_setAssociatedObject(self, kMetadataKey, asyncdisplaykit_metadata, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
+  objc_setAssociatedObject(self, kContextKey, asyncdisplaykit_context, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
 }
-- (ASBasicImageDownloader *)asyncdisplaykit_metadata
+- (ASBasicImageDownloader *)asyncdisplaykit_context
 {
-  return objc_getAssociatedObject(self, kMetadataKey);
+  return objc_getAssociatedObject(self, kContextKey);
 }
 @end
 
@@ -84,21 +145,38 @@ static const char *kMetadataKey = NSStringFromClass(ASBasicImageDownloaderMetada
      downloadProgressBlock:(void (^)(CGFloat))downloadProgressBlock
                 completion:(void (^)(CGImageRef, NSError *))completion
 {
+  ASBasicImageDownloaderContext *context = [ASBasicImageDownloaderContext contextForURL:URL];
+
+  // NSURLSessionDownloadTask will do file I/O to create a temp directory. If called on the main thread this will
+  // cause significant performance issues.
+  dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+    // the downloader may have been invalidated in the time it takes to async dispatch this block
+    if ([context isCancelled]) {
+      return;
+    }
+    
     // create download task
     NSURLSessionDownloadTask *task = [_session downloadTaskWithURL:URL];
 
+    // since creating the task does disk I/O, we should check if it has been invalidated
+    if ([context isCancelled]) {
+      return;
+    }
+
     // associate metadata with it
-  ASBasicImageDownloaderMetadata *metadata = [[ASBasicImageDownloaderMetadata alloc] init];
+    context.callbackQueue = callbackQueue ?: dispatch_get_main_queue();
-  metadata.callbackQueue = callbackQueue ?: dispatch_get_main_queue();
+    context.downloadProgressBlock = downloadProgressBlock;
-  metadata.downloadProgressBlock = downloadProgressBlock;
+    context.completionBlock = completion;
-  metadata.completionBlock = completion;
+    context.sessionTask = task;
-  task.originalRequest.asyncdisplaykit_metadata = metadata;
+    task.originalRequest.asyncdisplaykit_context = context;
 
     // start downloading
     [task resume];
 
-  // return the task as an opaque cancellation token
+    context.sessionTask = task;
-  return task;
+  });
+
+  return context;
 }
 
 - (void)cancelImageDownloadForIdentifier:(id)downloadIdentifier
@@ -107,10 +185,10 @@ static const char *kMetadataKey = NSStringFromClass(ASBasicImageDownloaderMetada
     return;
   }
 
-  ASDisplayNodeAssert([downloadIdentifier isKindOfClass:NSURLSessionDownloadTask.class], @"unexpected downloadIdentifier");
+  ASDisplayNodeAssert([downloadIdentifier isKindOfClass:ASBasicImageDownloaderContext.class], @"unexpected downloadIdentifier");
-  NSURLSessionDownloadTask *task = (NSURLSessionDownloadTask *)downloadIdentifier;
+  ASBasicImageDownloaderContext *context = (ASBasicImageDownloaderContext *)downloadIdentifier;
 
-  [task cancel];
+  [context cancel];
 }
 
 
@@ -121,9 +199,9 @@ static const char *kMetadataKey = NSStringFromClass(ASBasicImageDownloaderMetada
                                       totalBytesWritten:(int64_t)totalBytesWritten
                               totalBytesExpectedToWrite:(int64_t)totalBytesExpectedToWrite
 {
-  ASBasicImageDownloaderMetadata *metadata = downloadTask.originalRequest.asyncdisplaykit_metadata;
+  ASBasicImageDownloaderContext *context = downloadTask.originalRequest.asyncdisplaykit_context;
-  if (metadata.downloadProgressBlock) {
+  if (context.downloadProgressBlock) {
-    metadata.downloadProgressBlock((CGFloat)totalBytesWritten / (CGFloat)totalBytesExpectedToWrite);
+    context.downloadProgressBlock((CGFloat)totalBytesWritten / (CGFloat)totalBytesExpectedToWrite);
   }
 }
 
@@ -131,12 +209,16 @@ static const char *kMetadataKey = NSStringFromClass(ASBasicImageDownloaderMetada
 - (void)URLSession:(NSURLSession *)session downloadTask:(NSURLSessionDownloadTask *)downloadTask
                               didFinishDownloadingToURL:(NSURL *)location
 {
+  ASBasicImageDownloaderContext *context = downloadTask.originalRequest.asyncdisplaykit_context;
+  if ([context isCancelled]) {
+    return;
+  }
+
   UIImage *image = [UIImage imageWithData:[NSData dataWithContentsOfURL:location]];
 
-  ASBasicImageDownloaderMetadata *metadata = downloadTask.originalRequest.asyncdisplaykit_metadata;
+  if (context.completionBlock) {
-  if (metadata.completionBlock) {
+    dispatch_async(context.callbackQueue, ^{
-    dispatch_async(metadata.callbackQueue, ^{
+      context.completionBlock(image.CGImage, nil);
-      metadata.completionBlock(image.CGImage, nil);
     });
   }
 }
@@ -145,10 +227,10 @@ static const char *kMetadataKey = NSStringFromClass(ASBasicImageDownloaderMetada
 - (void)URLSession:(NSURLSession *)session task:(NSURLSessionDownloadTask *)task
                            didCompleteWithError:(NSError *)error
 {
-  ASBasicImageDownloaderMetadata *metadata = task.originalRequest.asyncdisplaykit_metadata;
+  ASBasicImageDownloaderContext *context = task.originalRequest.asyncdisplaykit_context;
-  if (metadata && error) {
+  if (context && error) {
-    dispatch_async(metadata.callbackQueue, ^{
+    dispatch_async(context.callbackQueue, ^{
-      metadata.completionBlock(NULL, error);
+      context.completionBlock(NULL, error);
     });
   }
 }

AsyncDisplayKit/Details/ASImageProtocols.h

@@ -16,10 +16,12 @@
 /**
   @abstract Attempts to fetch an image with the given URL from the cache.
   @param URL The URL of the image to retrieve from the cache.
-  @param callbackQueue The queue to call `completion` on. If this value is nil, @{ref completion} will be invoked on the main-queue.
+  @param callbackQueue The queue to call `completion` on. If this value is nil, @{ref completion} will be invoked on the 
+      main-queue.
   @param completion The block to be called when the cache has either hit or missed.
   @param imageFromCache The image that was retrieved from the cache, if the image could be retrieved; nil otherwise.
-  @discussion If `URL` is nil, `completion` will be invoked immediately with a nil image.
+  @discussion If `URL` is nil, `completion` will be invoked immediately with a nil image. This method should not block
+      the calling thread as it is likely to be called from the main thread.
  */
 - (void)fetchCachedImageWithURL:(NSURL *)URL
                   callbackQueue:(dispatch_queue_t)callbackQueue
@@ -33,14 +35,18 @@
 /**
   @abstract Downloads an image with the given URL.
   @param URL The URL of the image to download.
-  @param callbackQueue The queue to call `downloadProgressBlock` and `completion` on. If this value is nil, both blocks will be invoked on the main-queue.
+  @param callbackQueue The queue to call `downloadProgressBlock` and `completion` on. If this value is nil, both blocks 
+      will be invoked on the main-queue.
   @param downloadProgressBlock The block to be invoked when the download of `URL` progresses.
   @param progress The progress of the download, in the range of (0.0, 1.0), inclusive.
   @param completion The block to be invoked when the download has completed, or has failed.
   @param image The image that was downloaded, if the image could be successfully downloaded; nil otherwise.
   @param error An error describing why the download of `URL` failed, if the download failed; nil otherwise.
-  @discussion If `URL` is nil, `completion` will be invoked immediately with a nil image and an error describing why the download failed.
+  @discussion If `URL` is nil, `completion` will be invoked immediately with a nil image and an error describing why the 
-  @result An opaque identifier to be used in canceling the download, via `cancelImageDownloadForIdentifier:`. You must retain the identifier if you wish to use it later.
+      download failed. This method is likely to be called on the main thread, so any custom implementations should make
+      sure to background any expensive download operations.
+  @result An opaque identifier to be used in canceling the download, via `cancelImageDownloadForIdentifier:`. You must 
+      retain the identifier if you wish to use it later.
  */
 - (id)downloadImageWithURL:(NSURL *)URL
              callbackQueue:(dispatch_queue_t)callbackQueue
@@ -49,7 +55,8 @@
 
 /**
   @abstract Cancels an image download.
-  @param downloadIdentifier The opaque download identifier object returned from `downloadImageWithURL:callbackQueue:downloadProgressBlock:completion:`.
+  @param downloadIdentifier The opaque download identifier object returned from 
+      `downloadImageWithURL:callbackQueue:downloadProgressBlock:completion:`.
   @discussion This method has no effect if `downloadIdentifier` is nil.
  */
 - (void)cancelImageDownloadForIdentifier:(id)downloadIdentifier;

AsyncDisplayKit/Private/ASBasicImageDownloaderInternal.h

@@ -0,0 +1,22 @@
+/* Copyright (c) 2014-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+#import "ASThread.h"
+
+
+@interface ASBasicImageDownloaderContext : NSObject
+
++ (ASBasicImageDownloaderContext *)contextForURL:(NSURL *)URL;
+
+@property (nonatomic, strong, readonly) NSURL *URL;
+@property (nonatomic, weak) NSURLSessionTask *sessionTask;
+
+- (BOOL)isCancelled;
+- (void)cancel;
+
+@end

AsyncDisplayKitTests/ASBasicImageDownloaderContextTests.m

@@ -0,0 +1,72 @@
+/* Copyright (c) 2014-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+#import <AsyncDisplayKit/ASBasicImageDownloaderInternal.h>
+#import <AsyncDisplayKit/ASBasicImageDownloader.h>
+
+#import <OCMock/OCMock.h>
+
+#import <XCTest/XCTest.h>
+
+
+@interface ASBasicImageDownloaderContextTests : XCTestCase
+
+@end
+
+@implementation ASBasicImageDownloaderContextTests
+
+- (NSURL *)randomURL
+{
+  // random URL for each test, doesn't matter that this is not really a URL
+  return [NSURL URLWithString:[NSUUID UUID].UUIDString];
+}
+
+- (void)testContextCreation
+{
+  NSURL *url = [self randomURL];
+  ASBasicImageDownloaderContext *c1 = [ASBasicImageDownloaderContext contextForURL:url];
+  ASBasicImageDownloaderContext *c2 = [ASBasicImageDownloaderContext contextForURL:url];
+  XCTAssert(c1 == c2, @"Context objects are not the same");
+}
+
+- (void)testContextInvalidation
+{
+  NSURL *url = [self randomURL];
+  ASBasicImageDownloaderContext *context = [ASBasicImageDownloaderContext contextForURL:url];
+  [context cancel];
+  XCTAssert([context isCancelled], @"Context should be cancelled");
+}
+
+- (void)testAsyncContextInvalidation
+{
+  NSURL *url = [self randomURL];
+  ASBasicImageDownloaderContext *context = [ASBasicImageDownloaderContext contextForURL:url];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"Context invalidation"];
+
+  dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+    [expectation fulfill];
+    XCTAssert([context isCancelled], @"Context should be cancelled");
+  });
+
+  [context cancel];
+  [self waitForExpectationsWithTimeout:0.1 handler:nil];
+}
+
+- (void)testContextSessionCanceled
+{
+  NSURL *url = [self randomURL];
+  id task = [OCMockObject mockForClass:[NSURLSessionTask class]];
+  ASBasicImageDownloaderContext *context = [ASBasicImageDownloaderContext contextForURL:url];
+  context.sessionTask = task;
+
+  [[task expect] cancel];
+
+  [context cancel];
+}
+
+@end