From 1d82b12ba2c6e5885e28f8a59bc400524c13278c Mon Sep 17 00:00:00 2001
From: DreamPiggy <lizhuoli1126@126.com>
Date: Sat, 20 Jan 2018 14:44:00 +0800
Subject: [PATCH] Use a lock instead of barrier queue to avoid dispatch_sync
 blocking the main queue on race condition

---
 SDWebImage/SDWebImageDownloader.m | 99 +++++++++++++++++++------------
 1 file changed, 62 insertions(+), 37 deletions(-)

diff --git a/SDWebImage/SDWebImageDownloader.m b/SDWebImage/SDWebImageDownloader.m
index ca08ef77..a1a0e803 100644
--- a/SDWebImage/SDWebImageDownloader.m
+++ b/SDWebImage/SDWebImageDownloader.m
@@ -9,6 +9,10 @@
 #import "SDWebImageDownloader.h"
 #import "SDWebImageDownloaderOperation.h"
 
+#define LOCK(...) dispatch_semaphore_wait(self->_lock, DISPATCH_TIME_FOREVER); \
+__VA_ARGS__; \
+dispatch_semaphore_signal(self->_lock);
+
 @interface SDWebImageDownloadToken ()
 
 @property (nonatomic, weak, nullable) NSOperation<SDWebImageDownloaderOperationInterface> *downloadOperation;
@@ -36,8 +40,7 @@
 @property (assign, nonatomic, nullable) Class operationClass;
 @property (strong, nonatomic, nonnull) NSMutableDictionary<NSURL *, SDWebImageDownloaderOperation *> *URLOperations;
 @property (strong, nonatomic, nullable) SDHTTPHeadersMutableDictionary *HTTPHeaders;
-// This queue is used to serialize the handling of the network responses of all the download operation in a single queue
-@property (strong, nonatomic, nullable) dispatch_queue_t barrierQueue;
+@property (strong, nonatomic, nonnull) dispatch_semaphore_t lock; // a lock to keep the access to `URLOperations` thread-safe
 
 // The session in which data tasks will run
 @property (strong, nonatomic) NSURLSession *session;
@@ -96,7 +99,7 @@
 #else
         _HTTPHeaders = [@{@"Accept": @"image/*;q=0.8"} mutableCopy];
 #endif
-        _barrierQueue = dispatch_queue_create("com.hackemist.SDWebImageDownloaderBarrierQueue", DISPATCH_QUEUE_CONCURRENT);
+        _lock = dispatch_semaphore_create(1);
         _downloadTimeout = 15.0;
 
         [self createNewSessionWithConfiguration:sessionConfiguration];
@@ -231,13 +234,15 @@
 }
 
 - (void)cancel:(nullable SDWebImageDownloadToken *)token {
-    dispatch_barrier_async(self.barrierQueue, ^{
-        SDWebImageDownloaderOperation *operation = self.URLOperations[token.url];
-        BOOL canceled = [operation cancel:token.downloadOperationCancelToken];
-        if (canceled) {
-            [self.URLOperations removeObjectForKey:token.url];
-        }
-    });
+    NSURL *url = token.url;
+    if (!url) {
+        return;
+    }
+    SDWebImageDownloaderOperation *operation = [self operationForURL:url];
+    BOOL canceled = [operation cancel:token.downloadOperationCancelToken];
+    if (canceled) {
+        [self removeOperationForURL:url];
+    }
 }
 
 - (nullable SDWebImageDownloadToken *)addProgressCallback:(SDWebImageDownloaderProgressBlock)progressBlock
@@ -251,33 +256,24 @@
         }
         return nil;
     }
-
-    __block SDWebImageDownloadToken *token = nil;
-
-    dispatch_barrier_sync(self.barrierQueue, ^{
-        SDWebImageDownloaderOperation *operation = self.URLOperations[url];
-        if (!operation) {
-            operation = createCallback();
-            self.URLOperations[url] = operation;
-
-            __weak SDWebImageDownloaderOperation *woperation = operation;
-            operation.completionBlock = ^{
-				dispatch_barrier_sync(self.barrierQueue, ^{
-					SDWebImageDownloaderOperation *soperation = woperation;
-					if (!soperation) return;
-					if (self.URLOperations[url] == soperation) {
-						[self.URLOperations removeObjectForKey:url];
-					};
-				});
-            };
-        }
-        id downloadOperationCancelToken = [operation addHandlersForProgress:progressBlock completed:completedBlock];
-
-        token = [SDWebImageDownloadToken new];
-        token.downloadOperation = operation;
-        token.url = url;
-        token.downloadOperationCancelToken = downloadOperationCancelToken;
-    });
+    
+    SDWebImageDownloaderOperation *operation = [self operationForURL:url];
+    if (!operation) {
+        operation = createCallback();
+        [self setOperation:operation forURL:url];
+        
+        __weak typeof(self) wself = self;
+        operation.completionBlock = ^{
+            __strong typeof(wself) sself = wself;
+            [sself removeOperationForURL:url];
+        };
+    }
+    id downloadOperationCancelToken = [operation addHandlersForProgress:progressBlock completed:completedBlock];
+    
+    SDWebImageDownloadToken *token = [SDWebImageDownloadToken new];
+    token.downloadOperation = operation;
+    token.url = url;
+    token.downloadOperationCancelToken = downloadOperationCancelToken;
 
     return token;
 }
@@ -292,6 +288,35 @@
 
 #pragma mark Helper methods
 
+- (SDWebImageDownloaderOperation *)operationForURL:(NSURL *)url {
+    if (!url) {
+        return nil;
+    }
+    SDWebImageDownloaderOperation *operation;
+    LOCK({
+        operation = [self.URLOperations objectForKey:url];
+    });
+    return operation;
+}
+
+- (void)setOperation:(SDWebImageDownloaderOperation *)operation forURL:(NSURL *)url {
+    if (!operation || !url) {
+        return;
+    }
+    LOCK({
+       [self.URLOperations setObject:operation forKey:url];
+    });
+}
+
+- (void)removeOperationForURL:(NSURL *)url {
+    if (!url) {
+        return;
+    }
+    LOCK({
+        [self.URLOperations removeObjectForKey:url];
+    });
+}
+
 - (SDWebImageDownloaderOperation *)operationWithTask:(NSURLSessionTask *)task {
     SDWebImageDownloaderOperation *returnOperation = nil;
     for (SDWebImageDownloaderOperation *operation in self.downloadQueue.operations) {