How to enable WKURLSchemeHandler to do work off main thread?

Question

I am trying to get WKURLSchemeHandler to serve video files for when a WebView uses a custom url scheme. I realize that didReceive(data) can be called multiple times so I have figured out how to load my video file in chunks and send it back.

The problem is that all of this work is being done on the main thread. I cannot find an example of how to successfully get this done on a background thread. All examples of WKURLSchemeHandler that I can find including WWDC presentation video here (near the end of the video) are all so basic. None of them show how to handle large file let alone how to push work off of main thread.

And if I simply wrap everything inside a DispatchQueue.global(qos: .background).async {...} then my app crashes b/c WebView throws an unmanaged exception with the error this task has already been stopped !

Anyone knows how to successfully do this?

Answer 1

I finally figured it out. I can't believe how difficult this was. No wonder Apple hasn't released any samples around this. Here's my code:

// This is based on "Customized Loading in WKWebView" WWDC video (near the end of the
// video) at https://developer.apple.com/videos/play/wwdc2017/220 and A LOT of trial
// and error to figure out how to push work to background thread.
//
// To better understand how WKURLSchemeTask (and internally WebURLSchemeTask) works
// you can refer to the source code of WebURLSchemeTask at
// https://github.com/WebKit/WebKit/blob/main/Source/WebKit/UIProcess/WebURLSchemeTask.cpp
//
// Looking at that source code you can see that a call to any of the internals of
// WebURLSchemeTask (which is made through WKURLSchemeTask) is expected to be on the
// main thread, as you can see by the ASSERT(RunLoop::isMain()) statements at the
// beginning of pretty much every function and property getters. I'm not sure why Apple
// has decided to do these on the main thread since that would result in a blocked UI
// thread if we need to return large responses/files. At the very least they should have
// allowed for calls to come back on any thread and internally pass them to the main
// thread so that developers wouldn't have to write thread-synchronization code over and
// over every time they want to use WKURLSchemeHandler.
//
// The solution to pushing things off main thread is rather cumbersome. We need to call
// into DispatchQueue.global(qos: .background).async {...} but also manually ensure that
// everything is synchronized between the main and bg thread. We also manually need to
// keep track of the stopped tasks b/c a WKURLSchemeTask does not have any properties that
// we could query to see if it has stopped. If we respond to a WKURLSchemeTask that has
// stopped then an unmanaged exception is thrown which Swift cannot catch and the entire
// app will crash.
public class MyURLSchemeHandler: NSObject, WKURLSchemeHandler {
    private var stoppedTaskURLs: [URLRequest] = []

    public func webView(_ webView: WKWebView, start urlSchemeTask: WKURLSchemeTask) {
        let request = urlSchemeTask.request
        guard let requestUrl = request.url else { return }
        
        DispatchQueue.global(qos: .background).async { [weak self] in
            guard let strongSelf = self, requestUrl.scheme == "my-video-url-scheme" else {
                return
            }

            let filePath = requestUrl.absoluteString
            if let fileHandle = FileHandle(forReadingAtPath: filePath) {
                // video files can be very large in size, so read them in chuncks.
                let chunkSize = 1024 * 1024 // 1Mb
                let response = URLResponse(url: requestUrl,
                                           mimeType: "video/mp4",
                                           expectedContentLength: chunkSize,
                                           textEncodingName: nil)
                strongSelf.postResponse(to: urlSchemeTask, response: response)
                var data = fileHandle.readData(ofLength: chunkSize) // get the first chunk
                while (!data.isEmpty && !strongSelf.hasTaskStopped(urlSchemeTask)) {
                    strongSelf.postResponse(to: urlSchemeTask, data: data)
                    data = fileHandle.readData(ofLength: chunkSize) // get the next chunk
                }
                fileHandle.closeFile()
                strongSelf.postFinished(to: urlSchemeTask)
            } else {
                strongSelf.postFailed(
                    to: urlSchemeTask,
                    error: NSError(domain: "Failed to fetch resource",
                                   code: 0,
                                   userInfo: nil))
            }
            
            // remove the task from the list of stopped tasks (if it is there)
            // since we're done with it anyway
            strongSelf.stoppedTaskURLs = strongSelf.stoppedTaskURLs.filter{$0 != request}
        }
    }
    
    public func webView(_ webView: WKWebView, stop urlSchemeTask: WKURLSchemeTask) {
        if (!self.hasTaskStopped(urlSchemeTask)) {
            self.stoppedTaskURLs.append(urlSchemeTask.request)
        }
    }
    
    private func hasTaskStopped(_ urlSchemeTask: WKURLSchemeTask) -> Bool {
        return self.stoppedTaskURLs.contains{$0 == urlSchemeTask.request}
    }
    
    private func postResponse(to urlSchemeTask: WKURLSchemeTask,  response: URLResponse) {
        post(to: urlSchemeTask, action: {urlSchemeTask.didReceive(response)})
    }
    
    private func postResponse(to urlSchemeTask: WKURLSchemeTask,  data: Data) {
        post(to: urlSchemeTask, action: {urlSchemeTask.didReceive(data)})
    }
    
    private func postFinished(to urlSchemeTask: WKURLSchemeTask) {
        post(to: urlSchemeTask, action: {urlSchemeTask.didFinish()})
    }
    
    private func postFailed(to urlSchemeTask: WKURLSchemeTask, error: NSError) {
        post(to: urlSchemeTask, action: {urlSchemeTask.didFailWithError(error)})
    }
    
    private func post(to urlSchemeTask: WKURLSchemeTask, action: @escaping () -> Void) {
        let group = DispatchGroup()
        group.enter()
        DispatchQueue.main.async { [weak self] in
            if (self?.hasTaskStopped(urlSchemeTask) == false) {
                action()
            }
            group.leave()
        }
        group.wait()
    }
}