WKWebView causes my view controller to leak

My view controller displays a WKWebView. I installed a message handler, a cool Web Kit feature that allows my code to be notified from inside the web page:

override func viewDidAppear(animated: Bool) {
    super.viewDidAppear(animated)
    let url = // ...
    self.wv.loadRequest(NSURLRequest(URL:url))
    self.wv.configuration.userContentController.addScriptMessageHandler(
        self, name: "dummy")
}

func userContentController(userContentController: WKUserContentController,
    didReceiveScriptMessage message: WKScriptMessage) {
        // ...
}

So far so good, but now I've discovered that my view controller is leaking - when it is supposed to be deallocated, it isn't:

deinit {
    println("dealloc") // never called
}

It appears that merely installing myself as a message handler causes a retain cycle and hence a leak!


Correct as usual, King Friday. It turns out that the WKUserContentController retains its message handler. This makes a certain amount of sense, since it could hardly send a message to its message handler if its message handler had ceased to exist. It's parallel to the way a CAAnimation retains its delegate, for example.

However, it also causes a retain cycle, because the WKUserContentController itself is leaking. That doesn't matter much on its own (it's only 16K), but the retain cycle and leak of the view controller are bad.

My workaround is to interpose a trampoline object between the WKUserContentController and the message handler. The trampoline object has only a weak reference to the real message handler, so there's no retain cycle. Here's the trampoline object:

class LeakAvoider : NSObject, WKScriptMessageHandler {
    weak var delegate : WKScriptMessageHandler?
    init(delegate:WKScriptMessageHandler) {
        self.delegate = delegate
        super.init()
    }
    func userContentController(userContentController: WKUserContentController,
        didReceiveScriptMessage message: WKScriptMessage) {
            self.delegate?.userContentController(
                userContentController, didReceiveScriptMessage: message)
    }
}

Now when we install the message handler, we install the trampoline object instead of self:

self.wv.configuration.userContentController.addScriptMessageHandler(
    LeakAvoider(delegate:self), name: "dummy")

It works! Now deinit is called, proving that there is no leak. It looks like this shouldn't work, because we created our LeakAvoider object and never held a reference to it; but remember, the WKUserContentController itself is retaining it, so there's no problem.

For completeness, now that deinit is called, you can uninstall the message handler there, though I don't think this is actually necessary:

deinit {
    println("dealloc")
    self.wv.stopLoading()
    self.wv.configuration.userContentController.removeScriptMessageHandlerForName("dummy")
}

The leak is caused by userContentController.addScriptMessageHandler(self, name: "handlerName") which will keep a reference to the message handler self.

To prevent leaks, simply remove the message handler via userContentController.removeScriptMessageHandlerForName("handlerName") when you no longer need it. If you add the addScriptMessageHandler at viewDidAppear, its a good idea to remove it in viewDidDisappear.


The solution posted by matt is just what's needed. Thought I'd translate it to objective-c code

@interface WeakScriptMessageDelegate : NSObject<WKScriptMessageHandler>

@property (nonatomic, weak) id<WKScriptMessageHandler> scriptDelegate;

- (instancetype)initWithDelegate:(id<WKScriptMessageHandler>)scriptDelegate;

@end

@implementation WeakScriptMessageDelegate

- (instancetype)initWithDelegate:(id<WKScriptMessageHandler>)scriptDelegate
{
    self = [super init];
    if (self) {
        _scriptDelegate = scriptDelegate;
    }
    return self;
}

- (void)userContentController:(WKUserContentController *)userContentController didReceiveScriptMessage:(WKScriptMessage *)message
{
    [self.scriptDelegate userContentController:userContentController didReceiveScriptMessage:message];
}

@end

Then make use of it like this:

WKUserContentController *userContentController = [[WKUserContentController alloc] init];    
[userContentController addScriptMessageHandler:[[WeakScriptMessageDelegate alloc] initWithDelegate:self] name:@"name"];

I've also noted that you also need to remove the message handler(s) during teardown, otherwise the handler(s) will still live on (even if everything else about the webview is deallocated):

WKUserContentController *controller = 
self.webView.configuration.userContentController;

[controller removeScriptMessageHandlerForName:@"message"];

Details

  • Swift 5.1
  • Xcode 11.6 (11E708)

Solution

based on Matt's answer

protocol ScriptMessageHandlerDelegate: class {
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage)
}

class ScriptMessageHandler: NSObject, WKScriptMessageHandler {

    deinit { print("____ DEINITED: \(self)") }
    private var configuration: WKWebViewConfiguration!
    private weak var delegate: ScriptMessageHandlerDelegate?
    private var scriptNamesSet = Set<String>()

    init(configuration: WKWebViewConfiguration, delegate: ScriptMessageHandlerDelegate) {
        self.configuration = configuration
        self.delegate = delegate
        super.init()
    }

    func deinitHandler() {
        scriptNamesSet.forEach { configuration.userContentController.removeScriptMessageHandler(forName: $0) }
        configuration = nil
    }
    
    func registerScriptHandling(scriptNames: [String]) {
        for scriptName in scriptNames {
            if scriptNamesSet.contains(scriptName) { continue }
            configuration.userContentController.add(self, name: scriptName)
            scriptNamesSet.insert(scriptName)
        }
    }

    func userContentController(_ userContentController: WKUserContentController,
                               didReceive message: WKScriptMessage) {
        delegate?.userContentController(userContentController, didReceive: message)
    }
}

Full Sample

Do not forget to paste the Solution code here

import UIKit
import WebKit

class ViewController: UIViewController {

    override func viewDidLoad() {
        super.viewDidLoad()
        let button = UIButton(frame: CGRect(x: 0, y: 0, width: 100, height: 40))
        button.setTitle("WebView", for: .normal)
        view.addSubview(button)
        button.center = view.center
        button.addTarget(self, action: #selector(touchedUpInsed(button:)), for: .touchUpInside)
        button.setTitleColor(.blue, for: .normal)
    }
    
    @objc func touchedUpInsed(button: UIButton) {
        let viewController = WebViewController()
        present(viewController, animated: true, completion: nil)
    }
}

class WebViewController: UIViewController {

    private weak var webView: WKWebView!
    private var scriptMessageHandler: ScriptMessageHandler!
    private let url = URL(string: "http://google.com")!
    deinit {
        scriptMessageHandler.deinitHandler()
        print("____ DEINITED: \(self)")
    }

    override func viewDidLoad() {
        super.viewDidLoad()
        let configuration = WKWebViewConfiguration()
        scriptMessageHandler = ScriptMessageHandler(configuration: configuration, delegate: self)
        let scriptName = "GetUrlAtDocumentStart"
        scriptMessageHandler.registerScriptHandling(scriptNames: [scriptName])

        let jsScript = "webkit.messageHandlers.\(scriptName).postMessage(document.URL)"
        let script = WKUserScript(source: jsScript, injectionTime: .atDocumentStart, forMainFrameOnly: true)
        configuration.userContentController.addUserScript(script)
        
        let webView = WKWebView(frame: .zero, configuration: configuration)
        self.view.addSubview(webView)
        self.webView = webView
        webView.translatesAutoresizingMaskIntoConstraints = false
        webView.topAnchor.constraint(equalTo: view.topAnchor).isActive = true
        webView.leftAnchor.constraint(equalTo: view.leftAnchor).isActive = true
        view.bottomAnchor.constraint(equalTo: webView.bottomAnchor).isActive = true
        view.rightAnchor.constraint(equalTo: webView.rightAnchor).isActive = true
        webView.load(URLRequest(url: url))
    }
}

extension WebViewController: ScriptMessageHandlerDelegate {
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
        print("received \"\(message.body)\" from \"\(message.name)\" script")
    }
}

Info.plist

add in your Info.plist transport security setting

<key>NSAppTransportSecurity</key>
<dict>
    <key>NSAllowsArbitraryLoads</key>
    <true/>
</dict>