UIImageView subclass to deal with MJPEG streams

MJPEGImageView.swift

//
//  MJPEGImageView.swift
//  TestMJPEG
//
//  Created by Nico Ameghino on 4/23/16.
//  Copyright © 2016 Nicolas Ameghino. All rights reserved.
//

import UIKit

fileprivate extension Data {
    static let JPEGStartMarker = Data(bytes: [0xFF, 0xD8])
    static let JPEGEndMarker = Data(bytes: [0xFF, 0xD9])
}

class MJPEGImageView: UIImageView {
    
    var isStreaming: Bool { return dataHandler?.isStreaming ?? false }
    
    var placeholderImage: UIImage? {
        didSet {
            setPlaceholderImage()
        }
    }
    private var dataHandler: MJPEGDataHandler?
    
    private func setup() {
        setPlaceholderImage()
    }
    
    private func setPlaceholderImage() {
        DispatchQueue.main.async { [unowned self] in
            self.image = self.placeholderImage
        }
    }
    
    func start(with url: URL) {
        if isStreaming {
            stop()
        }
        dataHandler = MJPEGDataHandler(url: url)
        dataHandler?.delegate = self
        dataHandler?.start()
    }
    
    func stop() {
        dataHandler?.stop()
        setPlaceholderImage()
    }
}

extension MJPEGImageView: MJPEGDataHandlerDelegate {
    fileprivate func handler(_ handler: MJPEGDataHandler, didGet newImage: UIImage) {
        DispatchQueue.main.async { [unowned self] in
            self.image = newImage
        }
    }
    
    fileprivate func handler(_ handler: MJPEGDataHandler, didFailWith error: Error?) {
        NSLog("MJPEG data handler error: \(error)")
    }
}

fileprivate protocol MJPEGDataHandlerDelegate: class {
    func handler(_ handler: MJPEGDataHandler, didGet newImage: UIImage)
    func handler(_ handler: MJPEGDataHandler, didFailWith error: Error?)
}

fileprivate class MJPEGDataHandler: NSObject {
    
    weak var delegate: MJPEGDataHandlerDelegate?
    
    private(set) var isStreaming: Bool = false
    fileprivate var shouldProcessNewFrame: Bool = true
    
    // Queue to process the URL session delegate calls
    private let queue: OperationQueue = {
        let q = OperationQueue()
        q.name = "mjpeg.stream"
        q.maxConcurrentOperationCount = 1
        return q
    }()
    
    // Queue to synchronize buffer access
    fileprivate let synchronizationQueue: OperationQueue = {
        let q = OperationQueue()
        q.name = "mjpeg.sync"
        q.maxConcurrentOperationCount = 1
        return q
    }()
    
    private(set) var session: URLSession!
    private(set) var task: URLSessionDataTask!
    private(set) var url: URL?
    
    fileprivate var buffer = Data()
    
    init(url: URL) {
        super.init()
        self.url = url
        session = URLSession(configuration: .ephemeral, delegate: self, delegateQueue: queue)
        #if DEBUG
            NSLog("mjpeg data handler - initialized")
        #endif
    }
    
    func start() {
        #if DEBUG
            NSLog("mjpeg data handler - starting")
        #endif
        shouldProcessNewFrame = true
        isStreaming = true
        guard let url = url else {
            NSLog("MJPEG data handler error: attempted to start an instance without a URL")
            return
        }
        task = session.dataTask(with: url)
        task.resume()
        #if DEBUG
            NSLog("mjpeg data handler - started")
        #endif
        
    }
    
    func stop() {
        shouldProcessNewFrame = false
        isStreaming = false
        session.invalidateAndCancel()
        synchronizationQueue.cancelAllOperations()
        #if DEBUG
            NSLog("mjpeg data handler - stopped")
        #endif
        
    }
}

extension MJPEGDataHandler: URLSessionDataDelegate {
    fileprivate func urlSession(_ session: URLSession, didBecomeInvalidWithError error: Error?) {
        #if DEBUG
            NSLog("mjpeg data handler - failed: \(error.debugDescription)")
        #endif
        
        delegate?.handler(self, didFailWith: error)
    }
    
    fileprivate func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse, completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
        
        #if DEBUG
            NSLog("mjpeg data handler - response received")
            NSLog("\n\(response)")
        #endif
        
        
        #if DEBUG
            NSLog("mjpeg data handler - enqueued data access (new buffer)")
        #endif
        synchronizationQueue.addOperation { [unowned self] in
            #if DEBUG
                NSLog("mjpeg data handler - dequeued data access (new buffer)")
            #endif
            self.buffer = Data()
        }
        
        completionHandler(.allow)
    }
    
    fileprivate func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        #if DEBUG
            NSLog("mjpeg data handler - enqueued data access (add bytes)")
        #endif
        synchronizationQueue.addOperation { [unowned self] in
            guard self.shouldProcessNewFrame else { return }
            
            #if DEBUG
                NSLog("mjpeg data handler - dequeued data access (add bytes)")
            #endif
            
            
            self.buffer.append(data)
            
            
            #if DEBUG
                NSLog("mjpeg data handler - buffer size is now \(self.buffer.count) bytes (\(self.buffer.count / 1024) k - \(self.buffer.count / (1024 * 1024)) m))")
            #endif
            
            
            if self.buffer.count > 50 * 1024 { // if buffer > 50K, something went wrong, nuke and start over. -nico
                self.buffer = Data()
                #if DEBUG
                    NSLog("mjpeg data handler - buffer 'full', resetting")
                #endif
            }
            
            guard
                let start = self.buffer.range(of: .JPEGStartMarker),
                let end = self.buffer.range(of: .JPEGEndMarker)
                else {
                    #if DEBUG
                        NSLog("mjpeg data handler - no jpeg markers found, bailing out")
                    #endif
                    
                    return
            }
            
            if end.upperBound < start.lowerBound {
                return
            }
            
            let frameRange: Range<Int> = start.lowerBound..<end.upperBound
            let frameBytes = self.buffer.subdata(in: frameRange)
            guard let frame = UIImage(data: frameBytes) else { return }
            self.buffer.removeFirst(end.upperBound)
            #if DEBUG
                NSLog("mjpeg data handler - new frame, notifying delegate")
            #endif
            self.delegate?.handler(self, didGet: frame)
            
        }
    }
}