From 836bcb24a88e7a9976a6ec3c9c5240d0dcf1c0ea Mon Sep 17 00:00:00 2001
From: Bartek Waresiak <bartek@duckduckgo.com>
Date: Thu, 10 Oct 2024 12:23:37 +0200
Subject: [PATCH 1/2] Add error handling to contrainer removal

---
 Core/WebCacheManager.swift | 39 ++++++++++++++++++++++++++++++++------
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/Core/WebCacheManager.swift b/Core/WebCacheManager.swift
index c43462a36c..1b7dcefd11 100644
--- a/Core/WebCacheManager.swift
+++ b/Core/WebCacheManager.swift
@@ -82,14 +82,28 @@ public class WebCacheManager {
                       dataStoreIdManager: DataStoreIdManaging = DataStoreIdManager.shared) async {
 
         var cookiesToUpdate = [HTTPCookie]()
+        var leftoverContainerIDs = [UUID]()
         if #available(iOS 17, *) {
-            cookiesToUpdate += await containerBasedClearing(storeIdManager: dataStoreIdManager) ?? []
+            let result = await containerBasedClearing(storeIdManager: dataStoreIdManager)
+            cookiesToUpdate += result.cookies
+            leftoverContainerIDs = result.leftoverContainerIDs
         }
 
         // Perform legacy clearing to migrate to new container
         cookiesToUpdate += await legacyDataClearing() ?? []
 
         cookieStorage.updateCookies(cookiesToUpdate, keepingPreservedLogins: logins)
+
+        // Attempt to clean up leftover stores again after a delay
+        // This should not be a problem as these containers are not supposed to be used anymore.
+        // If this fails, we are going to still clean them next time as WebKit keeps track of all stores for us.
+        if #available(iOS 17, *), !leftoverContainerIDs.isEmpty {
+            DispatchQueue.main.asyncAfter(deadline: .now() + 3) {
+                for uuid in leftoverContainerIDs {
+                    WKWebsiteDataStore.remove(forIdentifier: uuid, completionHandler: { _ in })
+                }
+            }
+        }
     }
 
 }
@@ -113,18 +127,27 @@ extension WebCacheManager {
     }
 
     @available(iOS 17, *)
-    private func containerBasedClearing(storeIdManager: DataStoreIdManaging) async -> [HTTPCookie]? {
+    private func containerBasedClearing(storeIdManager: DataStoreIdManaging) async -> (cookies: [HTTPCookie],
+                                                                                       leftoverContainerIDs: [UUID]) {
         guard let containerId = storeIdManager.currentId else {
             storeIdManager.invalidateCurrentIdAndAllocateNew()
-            return []
+            return ([], [])
         }
         storeIdManager.invalidateCurrentIdAndAllocateNew()
 
+        var leftoverContainerIDs = [UUID]()
+
         var dataStore: WKWebsiteDataStore? = WKWebsiteDataStore(forIdentifier: containerId)
-        let cookies = await dataStore?.httpCookieStore.allCookies()
+        let cookies = await dataStore?.httpCookieStore.allCookies() ?? []
         dataStore = nil
 
         var uuids = await WKWebsiteDataStore.allDataStoreIdentifiers
+
+        // There may be a timing issue related to fetching current container, so append previous UUID as a precaution
+        if uuids.firstIndex(of: containerId) == nil {
+            uuids.append(containerId)
+        }
+
         if let newContainerID = storeIdManager.currentId,
             let newIdIndex = uuids.firstIndex(of: newContainerID) {
             assertionFailure("Attempted to cleanup current Data Store")
@@ -133,11 +156,15 @@ extension WebCacheManager {
 
         let previousLeftOversCount = max(0, uuids.count - 1) // -1 because one store is expected to be cleared
         for uuid in uuids {
-            try? await WKWebsiteDataStore.remove(forIdentifier: uuid)
+            do {
+                try await WKWebsiteDataStore.remove(forIdentifier: uuid)
+            } catch {
+                leftoverContainerIDs.append(uuid)
+            }
         }
         await checkForLeftBehindDataStores(previousLeftOversCount: previousLeftOversCount)
 
-        return cookies
+        return (cookies, leftoverContainerIDs)
     }
 
     private func legacyDataClearing() async -> [HTTPCookie]? {

From 264865f1a85bd2373ce5f379de4933d720abcdca Mon Sep 17 00:00:00 2001
From: Bartek Waresiak <bartek@duckduckgo.com>
Date: Mon, 14 Oct 2024 15:42:13 +0200
Subject: [PATCH 2/2] Update async code

---
 Core/WebCacheManager.swift | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/Core/WebCacheManager.swift b/Core/WebCacheManager.swift
index 1b7dcefd11..b6655f82b6 100644
--- a/Core/WebCacheManager.swift
+++ b/Core/WebCacheManager.swift
@@ -98,9 +98,10 @@ public class WebCacheManager {
         // This should not be a problem as these containers are not supposed to be used anymore.
         // If this fails, we are going to still clean them next time as WebKit keeps track of all stores for us.
         if #available(iOS 17, *), !leftoverContainerIDs.isEmpty {
-            DispatchQueue.main.asyncAfter(deadline: .now() + 3) {
+            Task {
+                try? await Task.sleep(for: .seconds(3))
                 for uuid in leftoverContainerIDs {
-                    WKWebsiteDataStore.remove(forIdentifier: uuid, completionHandler: { _ in })
+                    try? await WKWebsiteDataStore.remove(forIdentifier: uuid)
                 }
             }
         }