From 773f39429f8b29bf4f70572830b5636fb857de21 Mon Sep 17 00:00:00 2001 From: Mert Buran Date: Thu, 3 Feb 2022 14:55:53 +0100 Subject: [PATCH 1/2] RUMM-1931 stopTrackingDatadogEvents added to avoid memory leaks Without removing ScriptMessageHandlers explicitly, they stay alive in memory even after deallocation of the webview --- .../WKUserContentController+Datadog.swift | 16 ++++++++++ ...WKUserContentController+DatadogTests.swift | 31 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift b/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift index 10922285cd..336eee974d 100644 --- a/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift +++ b/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift @@ -27,6 +27,22 @@ public extension WKUserContentController { addDatadogMessageHandler(allowedWebViewHosts: hosts, hostsSanitizer: HostsSanitizer()) } + /// Disables Datadog iOS SDK and Datadog Browser SDK integration. + /// + /// Removes Datadog's ScriptMessageHandler and UserScript from the caller. + /// _NOTE:_ This method **must** be called when the webview can be deinitialized. + func stopTrackingDatadogEvents() { + removeScriptMessageHandler(forName: DatadogMessageHandler.name) + + let nonDatadogUserScripts = userScripts.filter { + return !$0.source.starts(with: Self.jsCodePrefix) + } + removeAllUserScripts() + nonDatadogUserScripts.forEach { + addUserScript($0) + } + } + internal func addDatadogMessageHandler(allowedWebViewHosts: Set, hostsSanitizer: HostsSanitizing) { guard !isTracking else { userLogger.warn("`trackDatadogEvents(in:)` was called more than once for the same WebView. Second call will be ignored. Make sure you call it only once.") diff --git a/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift b/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift index 22181cfeb3..cabbc9bbc3 100644 --- a/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift +++ b/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift @@ -23,6 +23,12 @@ final class DDUserContentController: WKUserContentController { } } +final class MockMessageHandler: NSObject, WKScriptMessageHandler { + func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) { + return + } +} + final class MockScriptMessage: WKScriptMessage { let mockBody: Any @@ -96,6 +102,31 @@ class WKUserContentController_DatadogTests: XCTestCase { XCTAssertEqual(recordedLogMessages, Array(repeating: "`trackDatadogEvents(in:)` was called more than once for the same WebView. Second call will be ignored. Make sure you call it only once.", count: multipleTimes - 1)) } + func testWhenStoppingTracking_itKeepsNonDatadogComponents() throws { + let controller = DDUserContentController() + + controller.trackDatadogEvents(in: []) + + let componentCount = 10 + for i in 0.. Date: Tue, 8 Feb 2022 15:03:43 +0100 Subject: [PATCH 2/2] RUMM-1931 PR comments addressed --- .../WebViewTrackingFixtureViewController.swift | 17 +++++++++++------ .../WKUserContentController+Datadog.swift | 2 +- .../WKUserContentController+DatadogTests.swift | 4 +--- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/Datadog/Example/Scenarios/WebView/WebViewTrackingFixtureViewController.swift b/Datadog/Example/Scenarios/WebView/WebViewTrackingFixtureViewController.swift index 2e5de7355d..92d0b05b8f 100644 --- a/Datadog/Example/Scenarios/WebView/WebViewTrackingFixtureViewController.swift +++ b/Datadog/Example/Scenarios/WebView/WebViewTrackingFixtureViewController.swift @@ -27,15 +27,20 @@ class ShopistWebviewViewController: UIViewController { override func viewDidLoad() { super.viewDidLoad() - let controller = WKUserContentController() - controller.trackDatadogEvents(in: ["shopist.io"]) - let config = WKWebViewConfiguration() - config.userContentController = controller - - webView = WKWebView(frame: UIScreen.main.bounds, configuration: config) + webView = WKWebView(frame: UIScreen.main.bounds, configuration: WKWebViewConfiguration()) view.addSubview(webView) } + override func viewWillAppear(_ animated: Bool) { + super.viewWillAppear(animated) + webView.configuration.userContentController.trackDatadogEvents(in: ["shopist.io"]) + } + + override func viewWillDisappear(_ animated: Bool) { + super.viewWillDisappear(animated) + webView.configuration.userContentController.stopTrackingDatadogEvents() + } + override func viewDidAppear(_ animated: Bool) { super.viewDidAppear(animated) webView.load(request) diff --git a/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift b/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift index 336eee974d..7b270b1841 100644 --- a/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift +++ b/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift @@ -30,7 +30,7 @@ public extension WKUserContentController { /// Disables Datadog iOS SDK and Datadog Browser SDK integration. /// /// Removes Datadog's ScriptMessageHandler and UserScript from the caller. - /// _NOTE:_ This method **must** be called when the webview can be deinitialized. + /// - Note: This method **must** be called when the webview can be deinitialized. func stopTrackingDatadogEvents() { removeScriptMessageHandler(forName: DatadogMessageHandler.name) diff --git a/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift b/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift index cabbc9bbc3..62d1eb953b 100644 --- a/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift +++ b/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift @@ -24,9 +24,7 @@ final class DDUserContentController: WKUserContentController { } final class MockMessageHandler: NSObject, WKScriptMessageHandler { - func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) { - return - } + func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) { } } final class MockScriptMessage: WKScriptMessage {