ImageDownloader.swift

final class ImageDownloader {

    typealias CompletionHandler = (UIImage?) -> Void

    private struct ResponseHandler {
        let dataTask: URLSessionDataTask
        var completionHandlers: [CompletionHandler]
    }

    private let lock = NSLock()

    private let dataTasks = NSMutableOrderedSet()
    private var activeDataTaskCount = 0
    private var responseHandlers: [URL: ResponseHandler] = [:]

    private let urlSession: URLSession
    private let maxActiveDataTaskCount: Int

    init(memoryCapacity: Int = 100 * 1024 * 1024,
         diskCapacity: Int = 150 * 1024 * 1024,
         maxActiveDataTaskCount: Int = 10) {

        let configuration = URLSessionConfiguration.default
        configuration.urlCache = URLCache(memoryCapacity: memoryCapacity,
                                          diskCapacity: diskCapacity,
                                          diskPath: "com.dpgao.ImageCache")
        configuration.requestCachePolicy = .returnCacheDataElseLoad
        self.urlSession = URLSession(configuration: configuration)

        self.maxActiveDataTaskCount = maxActiveDataTaskCount
    }

    private func enqueue(_ dataTask: URLSessionTask) {
        dataTask.resume()
        dataTasks.insert(dataTask, at: 0)
        activeDataTaskCount += 1
        while activeDataTaskCount > maxActiveDataTaskCount {
            activeDataTaskCount -= 1
            (dataTasks[activeDataTaskCount] as! URLSessionTask).suspend()
        }
    }

    private func remove(_ dataTask: URLSessionDataTask) {
        if dataTask.state != .suspended {
            activeDataTaskCount -= 1
        }
        dataTasks.remove(dataTask)
        while activeDataTaskCount < min(dataTasks.count, maxActiveDataTaskCount) {
            (dataTasks[activeDataTaskCount] as! URLSessionTask).resume()
            activeDataTaskCount += 1
        }
    }

    private func promote(_ dataTask: URLSessionDataTask) {
        if dataTask.state == .suspended {
            dataTasks.remove(dataTask)
            enqueue(dataTask)
        }
    }

    func download(from url: URL, completionHandler: @escaping CompletionHandler) {
        lock.lock()
        let handler: ResponseHandler
        if var responseHandler = responseHandlers.removeValue(forKey: url) {
            promote(responseHandler.dataTask)
            responseHandler.completionHandlers.append(completionHandler)
            handler = responseHandler
        } else {
            let dataTask = urlSession.dataTask(with: url) { data, _, _ in
                self.lock.lock()
                let responseHandler = self.responseHandlers.removeValue(forKey: url)!
                self.remove(responseHandler.dataTask)
                self.lock.unlock()
                let image = data.flatMap(UIImage.init)
                responseHandler.completionHandlers.forEach { $0(image) }
            }
            enqueue(dataTask)
            handler = ResponseHandler(dataTask: dataTask, completionHandlers: [completionHandler])
        }
        responseHandlers[url] = handler
        lock.unlock()
    }
}