Fix memory issue on DBP webviews (#1785)

Task/Issue URL: https://app.asana.com/0/1203581873609357/1205763744134003/f **Description**: Fix memory issue on DBP webviews after the operations are ran
duckduckgo · Nov 1, 2023 · 94a8a3d · 94a8a3d
1 parent 14d504e
commit 94a8a3d
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 10 deletions.
diff --git a/...s/DataBrokerProtection/Sources/DataBrokerProtection/CCF/DataBrokerProtectionFeature.swift b/...s/DataBrokerProtection/Sources/DataBrokerProtection/CCF/DataBrokerProtectionFeature.swift
@@ -43,7 +43,7 @@ enum CCFReceivedMethodName: String {
 struct DataBrokerProtectionFeature: Subfeature {
     var messageOriginPolicy: MessageOriginPolicy = .all
     var featureName: String = "brokerProtection"
-    var broker: UserScriptMessageBroker? // This broker is not related to DBP brokers. It's just a name we inherit from Subfeature
+    weak var broker: UserScriptMessageBroker? // This broker is not related to DBP brokers. It's just a name we inherit from Subfeature
 
     weak var delegate: CCFCommunicationDelegate?
 

diff --git a/...ges/DataBrokerProtection/Sources/DataBrokerProtection/CCF/DataBrokerProtectionUtils.swift b/...ges/DataBrokerProtection/Sources/DataBrokerProtection/CCF/DataBrokerProtectionUtils.swift
@@ -19,18 +19,19 @@ import Foundation
 import WebKit
 import BrowserServicesKit
 import UserScript
+import Common
 
 @MainActor
 final class DataBrokerUserContentController: WKUserContentController {
 
-    let dataBrokerUserScripts: DataBrokerUserScript
+    var dataBrokerUserScripts: DataBrokerUserScript?
 
     init(with privacyConfigurationManager: PrivacyConfigurationManaging, prefs: ContentScopeProperties, delegate: CCFCommunicationDelegate) {
         dataBrokerUserScripts = DataBrokerUserScript(privacyConfig: privacyConfigurationManager, prefs: prefs, delegate: delegate)
 
         super.init()
 
-        dataBrokerUserScripts.userScripts.forEach {
+        dataBrokerUserScripts?.userScripts.forEach {
             let userScript = $0.makeWKUserScriptSync()
             self.installUserScripts([userScript], handlers: [$0])
         }
@@ -44,6 +45,20 @@ final class DataBrokerUserContentController: WKUserContentController {
         handlers.forEach { self.addHandler($0) }
         wkUserScripts.forEach(self.addUserScript)
     }
+
+    public func cleanUpBeforeClosing() {
+        os_log("Cleaning up DBP user scripts", log: .dataBrokerProtection)
+
+        self.removeAllUserScripts()
+        self.removeAllScriptMessageHandlers()
+
+        self.removeAllContentRuleLists()
+        dataBrokerUserScripts = nil
+    }
+
+    deinit {
+        os_log("DataBrokerUserContentController Deinit", log: .dataBrokerProtection)
+    }
 }
 
 @MainActor

diff --git a/LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/CCF/WebViewHandler.swift b/LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/CCF/WebViewHandler.swift
@@ -37,9 +37,9 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
 
     private let isFakeBroker: Bool
     private let webViewConfiguration: WKWebViewConfiguration
-    private let userContentController: DataBrokerUserContentController?
+    private var userContentController: DataBrokerUserContentController?
 
-    private var webView: WKWebView?
+    private var webView: WebView?
     private var window: NSWindow?
 
     init(privacyConfig: PrivacyConfigurationManaging, prefs: ContentScopeProperties, delegate: CCFCommunicationDelegate, isFakeBroker: Bool = false) {
@@ -56,7 +56,7 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
     }
 
     func initializeWebView(showWebView: Bool) async {
-        webView = WKWebView(frame: CGRect(origin: .zero, size: CGSize(width: 1024, height: 1024)), configuration: webViewConfiguration)
+        webView = WebView(frame: CGRect(origin: .zero, size: CGSize(width: 1024, height: 1024)), configuration: webViewConfiguration)
         webView?.navigationDelegate = self
 
         if showWebView {
@@ -80,10 +80,19 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
 
     func finish() {
         os_log("WebViewHandler finished", log: .action)
+
         webView?.stopLoading()
+        userContentController?.cleanUpBeforeClosing()
+
+        userContentController = nil
+        webView?.navigationDelegate = nil
         webView = nil
     }
 
+    deinit {
+        print("WebViewHandler Deinit")
+    }
+
     func waitForWebViewLoad(timeoutInSeconds: Int = 0) async throws {
         try await withCheckedThrowingContinuation { continuation in
             self.activeContinuation = continuation
@@ -103,7 +112,7 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
     func execute(action: Action, data: CCFRequestData) {
         os_log("Executing action: %{public}@", log: .action, String(describing: action.actionType.rawValue))
 
-        userContentController?.dataBrokerUserScripts.dataBrokerFeature.pushAction(
+        userContentController?.dataBrokerUserScripts?.dataBrokerFeature.pushAction(
             method: .onActionReceived,
             webView: self.webView!,
             params: Params(state: ActionRequest(action: action, data: data))
@@ -159,3 +168,11 @@ extension DataBrokerProtectionWebViewHandler: WKNavigationDelegate {
         }
     }
 }
+
+private class WebView: WKWebView {
+
+    deinit {
+        configuration.userContentController.removeAllUserScripts()
+        os_log("DBP WebView Deinit", log: .action)
+    }
+}
diff --git a/...erProtection/Sources/DataBrokerProtection/Operations/DataBrokerOperationsCollection.swift b/...erProtection/Sources/DataBrokerProtection/Operations/DataBrokerOperationsCollection.swift
@@ -149,12 +149,14 @@ final class DataBrokerOperationsCollection: Operation {
                     guard let self = self else { return false }
                     return !self.isCancelled
                 })
-                os_log("Finished operation: %{public}@", log: .dataBrokerProtection, String(describing: id.uuidString))
 
                 if let sleepInterval = intervalBetweenOperations {
                     os_log("Waiting...: %{public}f", log: .dataBrokerProtection, sleepInterval)
                     try await Task.sleep(nanoseconds: UInt64(sleepInterval) * 1_000_000_000)
                 }
+
+                finish()
+
             } catch {
                 os_log("Error: %{public}@", log: .dataBrokerProtection, error.localizedDescription)
                 if let error = error as? DataBrokerProtectionError,
@@ -177,5 +179,7 @@ final class DataBrokerOperationsCollection: Operation {
 
         didChangeValue(forKey: #keyPath(isExecuting))
         didChangeValue(forKey: #keyPath(isFinished))
+
+        os_log("Finished operation: %{public}@", log: .dataBrokerProtection, String(describing: id.uuidString))
     }
 }
diff --git a/...okerProtection/Sources/DataBrokerProtection/Scheduler/DataBrokerProtectionProcessor.swift b/...okerProtection/Sources/DataBrokerProtection/Scheduler/DataBrokerProtectionProcessor.swift
@@ -69,7 +69,7 @@ final class DataBrokerProtectionProcessor {
         runOperations(operationType: .optOut,
                       priorityDate: nil,
                       showWebView: showWebView) {
-            os_log("Scans done", log: .dataBrokerProtection)
+            os_log("Optouts done", log: .dataBrokerProtection)
             completion?()
         }
     }
@@ -154,4 +154,8 @@ final class DataBrokerProtectionProcessor {
 
         return collections
     }
+
+    deinit {
+        os_log("Deinit DataBrokerProtectionProcessor", log: .dataBrokerProtection)
+    }
 }
diff --git a/...ckages/DataBrokerProtection/Sources/DataBrokerProtection/UI/DBPUICommunicationLayer.swift b/...ckages/DataBrokerProtection/Sources/DataBrokerProtection/UI/DBPUICommunicationLayer.swift
@@ -62,7 +62,7 @@ enum DBPUISendableMethodName: String {
 struct DBPUICommunicationLayer: Subfeature {
     var messageOriginPolicy: MessageOriginPolicy = .all
     var featureName: String = "dbpuiCommunication"
-    var broker: UserScriptMessageBroker?
+    weak var broker: UserScriptMessageBroker?
 
     weak var delegate: DBPUICommunicationDelegate?