RUMM-1931 stopTrackingDatadogEvents added to avoid memory leaks

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)

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<String>, 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.")

Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift

@@ -23,6 +23,10 @@ final class DDUserContentController: WKUserContentController {
     }
 }
 
+final class MockMessageHandler: NSObject, WKScriptMessageHandler {
+    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) { }
+}
+
 final class MockScriptMessage: WKScriptMessage {
     let mockBody: Any
 
@@ -96,6 +100,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..<componentCount {
+            let userScript = WKUserScript(
+                source: String.mockRandom(),
+                injectionTime: (i % 2 == 0 ? .atDocumentStart : .atDocumentEnd),
+                forMainFrameOnly: i % 2 == 0
+            )
+            controller.addUserScript(userScript)
+            controller.add(MockMessageHandler(), name: String.mockRandom())
+        }
+
+        XCTAssertEqual(controller.userScripts.count, componentCount + 1)
+        XCTAssertEqual(controller.messageHandlers.count, componentCount + 1)
+
+        controller.stopTrackingDatadogEvents()
+
+        XCTAssertEqual(controller.userScripts.count, componentCount)
+        XCTAssertEqual(controller.messageHandlers.count, componentCount)
+    }
+
     func testItLogsInvalidWebMessages() throws {
         let previousUserLogger = userLogger
         defer { userLogger = previousUserLogger }