[Swift]WKUserContentController.add(_:name:)を使うときに気をつけること

ざっくりまとめ。

  • WKUserContentController.add(_:name:)を使うとwebとアプリ側でやりとりができる
  • 注意しないとメモリーリークするぞ
    • add したなら remove もセットで実装しましょう

です。

追加したら解除まで!

環境

  • Swfit v4.2
  • Xcode v10.1

やりたいこと

WKWebView をつかって web を表示なんかしていると...
web 側とアプリ側で連携したくなることがあります。
特に自分で用意した web とかだと特に。

WKUserContentController.add(_:name:)

WKUserContentController を使うことでソレが可能になります。

class ViewController {
    func viewDidLoad() {
        let config = WKWebViewConfiguration()
        let controller = WKUserContentController()
        controller.add(self, name: "handle")
        config.userContentController = controller
        let webView = WKWebView(frame: .zero, configuration: config)
        view.addSubView(webView)
    }
}
extension ViewController: WKScriptMessageHandler {
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
        print(message.name) // handle が吐かれる
        print(message.body) // xxx が吐かれる
    }
}

js 側では webkit.messageHandlers.handle.postMessage("xxx") こうやると先の WKScriptMessageHandler の箇所が呼ばれますっと

メモリーリーク

上記のサンプルだと、メモリーリークします。
画面を閉じるなりしても、 ViewController.deinit が呼ばれないのです。
というのも、add(_:name:) したときに self を渡していることで、循環参照が発生してしまっているのです。
なので、画面を閉じても解放されないというわけです。
これを解放するようにするためには、WKUserContentController.remove(forName:) を呼ぶ必要があるのです。

web とアプリ側で連携するためには!?とかいてるほとんどのとこで、この補足がないのでついつい忘れがちですよね...😅

appear / disappear で制御する

かといって、deinitがよばれないわけですから、、。
ぱっと思いつく方法としては、appear / disappear での add / remove です

class ViewController {
    var webView: WKWebView?
    func viewDidLoad() {
        let config = WKWebViewConfiguration()
        let controller = WKUserContentController()
//        controller.add(self, name: "handle") <- ここではやらずに
        config.userContentController = controller
        webView = WKWebView(frame: .zero, configuration: config)
        view.addSubView(webView!)
    }
    
    func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)
        webView?.configuration.userContentController.add(self, name: "handle")
    }
    
    func viewWillDisappear(_ animated: Bool) {
        super.viewWillDisappear(animated)
        webView?.configuration.userContentController.remove(forName: "handle")
    }
}

こういう感じで書いておくと、画面を閉じたあとしっかり解放されます。
ただ、このケースだと、他の画面に遷移した場合(この画面が裏に回った場合)に
web 側からなにか飛ばされてきても拾えなくなってしまいます。
(むしろ、裏に回ってるときは拾いたくないという仕様であればこれで万事okです)

それを避けるために isBeingPresented とか isMovingFromParent とかそこらへんのプロパティを駆使することになりそうですけれど...。
なんというか appear / disappear の両方でそのプロパティをみながら制御するというのがちょっと... 🤔
みたいな気持ちにもなります。

サブクラス化してみる

ということで、画面状態のプロパティを見ることなく、
画面が解放される然るべきときにちゃんと解放されるようにしたい..というのを目指してサブクラスをつくってみます

class UserContentController: WKUserContentController {
    weak var delegate: UserContentControllerDelegate?
    init() {
        super.init()
        add(self, name: "handle")
    }
}
extension UserContentController: WKScriptMessageHandler {
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
        switch message.name {
        case "handle":
            delegate?.didReceiveHandleMessage(body: message.body)
        default: break
        }
    }
}

protocol UserContentControllerDelegate: NSObjectPoolProtocol {
    func didReceiveHandleMessage(_ name: String, body: Any)
}

こういう感じ。
それで ViewController 側では WKScriptMessageHandler のかわりに UserContentControllerDelegate を利用して、
かつ、 delegate も設定してあげる。
そうすると、 ViewController と UserContentController の間には循環参照は発生しないので、画面が閉じられて然るべきときにちゃんと解放されるようになります。
ViewController の方は...

今度は UserContentController の中で循環参照が発生してしまっているので、
ViewController は解放されても UserContentController が残り続ける結果に。
ということで、もう少し手を加えて...
解放する用のメソッドを用意してあげて、無事に呼ばれるようになった ViewController.deinit でそのメソッドを呼んであげる。

class ViewController {
    var webView: WKWebView?
    func viewDidLoad() {
        let config = WKWebViewConfiguration()
        let controller = UserContentController()
        config.userContentController = controller
        config.delegate = self
        webView = WKWebView(frame: .zero, configuration: config)
        view.addSubView(webView)
    }
    
    deinit {
        if let controller = webView.configuration.userContentController as? UserContentController {
            controller.invalidate() // 解放されるように呼んであげる
        }
    }
}
extension ViewController: UserContentControllerDelegate {
    func didReceiveHandleMessage(body: Any?) {
    }
}

class UserContentController: WKUserContentController {
    weak var delegate: UserContentControllerDelegate?
    init() {
        super.init()
        add(self, name: "handle")
    }
    func invalidate() {
        remove(forName: "handle")
    }
}
extension UserContentController: WKScriptMessageHandler {
    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
        switch message.name {
        case "handle":
            delegate?.didReceiveHandleMessage(body: message.body)
        default:
            print("unknown message: " + message.name)
        }
    }
}

protocol UserContentControllerDelegate: NSObjectPoolProtocol {
    func didReceiveHandleMessage(body: Any?)
}

フルでかくと最終的にはこういう感じ。
これで、 ViewController も UserContentController もちゃんと解放されるようになります💪

補足

ツールとかなにかをいれることで手助けに!

今回僕も実装時には、remove(forName:) のほうをしっかり(?)忘れていたのですが、すぐに気づくことが出来たのです。
(メモリーリークって実害が顕著に出てくるまで気づきにくかったりするのですが...)

そのきっかけになったのが...

です。感謝感謝🙏
こういうツールを入れておくと気づきが早くなるのでおすすめです!

WKWebViewの不思議

userContentController を使うにあたって、WKWebViewConfiguration とかも用意したりしてるんですけど、
でも config の方は自分でやる必要あるんだろうか?🤔とふと思って

let webView = WKWebView(frame: .zero)
let controller = UserContentController()
webView.configuration.userContentController = controller

というふうに、WKWebView初期化時ではなくて、あとからセットするように書いてみたんです。
これでもエラーなどはでることなくビルドできます。。
ただし、実行してみると... deinit でよばれるはずの UserContentController.invalidate() が呼ばれないんですよね。
なぜなら、if let controller = webView?.configuration.controller as? UserContentController が nilになるので...。
なんか、あとから設定したのは中で握りつぶされてるっぽいのです。。
type をみてみると、 UserContentController をセットしたはずなのに WKUserContentController だったので...
だとしたら、ビルド失敗するなり、実行時にワーニングを出すなりしてほしいところですけど...🤔

参考

WKUserContentController