Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
abstract datastore removal to make it testable
Browse files Browse the repository at this point in the history
brindy committed Nov 26, 2024
1 parent 1660429 commit 7ec9cde
Showing 4 changed files with 115 additions and 93 deletions.
55 changes: 10 additions & 45 deletions Core/WebCacheManager.swift
Original file line number Diff line number Diff line change
@@ -88,13 +88,16 @@ public class WebCacheManager: WebsiteDataManaging {
let cookieStorage: MigratableCookieStorage
let fireproofing: Fireproofing
let dataStoreIdManager: DataStoreIdManaging
let dataStoreClearingDelay: TimeInterval
let dataStoreCleaner: WebsiteDataStoreCleaning

public init(cookieStorage: MigratableCookieStorage, fireproofing: Fireproofing, dataStoreIdManager: DataStoreIdManaging, dataStoreClearingDelay: TimeInterval = 3.0) {
public init(cookieStorage: MigratableCookieStorage,
fireproofing: Fireproofing,
dataStoreIdManager: DataStoreIdManaging,
dataStoreCleaner: WebsiteDataStoreCleaning = DefaultWebsiteDataStoreCleaner()) {
self.cookieStorage = cookieStorage
self.fireproofing = fireproofing
self.dataStoreIdManager = dataStoreIdManager
self.dataStoreClearingDelay = dataStoreClearingDelay
self.dataStoreCleaner = dataStoreCleaner
}

/// The previous version saved cookies externally to the data so we can move them between containers. We now use
@@ -128,9 +131,10 @@ public class WebCacheManager: WebsiteDataManaging {

public func clear(dataStore: WKWebsiteDataStore) async {

let count = await dataStoreCleaner.countContainers()
await performMigrationIfNeeded(dataStoreIdManager: dataStoreIdManager, cookieStorage: cookieStorage, destinationStore: dataStore)
await clearData(inDataStore: dataStore, withFireproofing: fireproofing)
removeContainersIfNeeded()
dataStoreCleaner.removeAllContainersAfterDelay(previousCount: count)

}

@@ -160,47 +164,8 @@ extension WebCacheManager {
dataStoreIdManager.invalidateCurrentId()
}

private func removeContainersIfNeeded() {
// Check version here rather than on function so that we don't need complicated logic related to verison in the calling function
guard #available(iOS 17, *) else { return }

// Attempt to clean up all previous stores, but wait for a few seconds.
// If this fails, we are going to still clean them next time as WebKit keeps track of all stores for us.
Task {
try? await Task.sleep(interval: dataStoreClearingDelay)
for uuid in await WKWebsiteDataStore.allDataStoreIdentifiers {
try? await WKWebsiteDataStore.remove(forIdentifier: uuid)
}

let count = await WKWebsiteDataStore.allDataStoreIdentifiers.count
switch count {
case 0:
Pixel.fire(pixel: .debugWebsiteDataStoresCleared)

case 1:
Pixel.fire(pixel: .debugWebsiteDataStoresNotClearedOne)

default:
Pixel.fire(pixel: .debugWebsiteDataStoresNotClearedMultiple)
}
}
}

private func checkForLeftBehindDataStores(previousLeftOversCount: Int) async {
guard #available(iOS 17, *) else { return }

let params = [
"left_overs_count": "\(previousLeftOversCount)"
]

let ids = await WKWebsiteDataStore.allDataStoreIdentifiers
if ids.count > 1 {
Pixel.fire(pixel: .debugWebsiteDataStoresNotClearedMultiple, withAdditionalParameters: params)
} else if ids.count > 0 {
Pixel.fire(pixel: .debugWebsiteDataStoresNotClearedOne, withAdditionalParameters: params)
} else if previousLeftOversCount > 0 {
Pixel.fire(pixel: .debugWebsiteDataStoresCleared, withAdditionalParameters: params)
}
private func removeContainersIfNeeded(previousCount: Int) {
dataStoreCleaner.removeAllContainersAfterDelay(previousCount: previousCount)
}

private func clearData(inDataStore dataStore: WKWebsiteDataStore, withFireproofing fireproofing: Fireproofing) async {
70 changes: 70 additions & 0 deletions Core/WebsiteDataStoreCleaning.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
//
// WebsiteDataStoreCleaning.swift
// DuckDuckGo
//
// Copyright © 2024 DuckDuckGo. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

import WebKit

public protocol WebsiteDataStoreCleaning {

func countContainers() async -> Int
func removeAllContainersAfterDelay(previousCount: Int)

}

public class DefaultWebsiteDataStoreCleaner: WebsiteDataStoreCleaning {

public init() { }

public func countContainers() async -> Int {
guard #available(iOS 17, *) else { return 0 }
return await WKWebsiteDataStore.allDataStoreIdentifiers.count
}

public func removeAllContainersAfterDelay(previousCount: Int) {
guard #available(iOS 17, *) else { return }

// Attempt to clean up all previous stores, but wait for a few seconds.
// If this fails, we are going to still clean them next time as WebKit keeps track of all stores for us.
Task {
try? await Task.sleep(interval: 3.0)
for uuid in await WKWebsiteDataStore.allDataStoreIdentifiers {
try? await WKWebsiteDataStore.remove(forIdentifier: uuid)
}

await checkForLeftBehindDataStores(previousLeftOversCount: previousCount)
}
}

private func checkForLeftBehindDataStores(previousLeftOversCount: Int) async {
guard #available(iOS 17, *) else { return }

let params = [
"left_overs_count": "\(previousLeftOversCount)"
]

let ids = await WKWebsiteDataStore.allDataStoreIdentifiers
if ids.count > 1 {
Pixel.fire(pixel: .debugWebsiteDataStoresNotClearedMultiple, withAdditionalParameters: params)
} else if ids.count > 0 {
Pixel.fire(pixel: .debugWebsiteDataStoresNotClearedOne, withAdditionalParameters: params)
} else if previousLeftOversCount > 0 {
Pixel.fire(pixel: .debugWebsiteDataStoresCleared, withAdditionalParameters: params)
}
}

}
8 changes: 6 additions & 2 deletions DuckDuckGo.xcodeproj/project.pbxproj
Original file line number Diff line number Diff line change
@@ -365,9 +365,9 @@
6FEC0B852C999352006B4F6E /* FavoriteItem.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FEC0B842C999352006B4F6E /* FavoriteItem.swift */; };
6FEC0B882C999961006B4F6E /* FavoritesListInteractingAdapter.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FEC0B872C999961006B4F6E /* FavoritesListInteractingAdapter.swift */; };
6FF915822B88E07A0042AC87 /* AdAttributionFetcherTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FF915802B88E0750042AC87 /* AdAttributionFetcherTests.swift */; };
6FF9AD452CE766F700C5A406 /* NewTabPageControllerPixelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FF9AD442CE766F700C5A406 /* NewTabPageControllerPixelTests.swift */; };
6FF9AD3F2CE63DD800C5A406 /* TabSwitcherOpenDailyPixel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FF9AD3E2CE63DC200C5A406 /* TabSwitcherOpenDailyPixel.swift */; };
6FF9AD412CE6610F00C5A406 /* TabSwitcherDailyPixelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FF9AD402CE6610600C5A406 /* TabSwitcherDailyPixelTests.swift */; };
6FF9AD452CE766F700C5A406 /* NewTabPageControllerPixelTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 6FF9AD442CE766F700C5A406 /* NewTabPageControllerPixelTests.swift */; };
7B1604E82CB685B400A44EC6 /* Logger+TipKit.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B1604E72CB685B400A44EC6 /* Logger+TipKit.swift */; };
7B1604EC2CB68BDA00A44EC6 /* TipKitController+ConvenienceInitializers.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B1604EB2CB68BDA00A44EC6 /* TipKitController+ConvenienceInitializers.swift */; };
7B1604EE2CB68D2600A44EC6 /* TipKitDebugOptionsUIActionHandling.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7B1604ED2CB68D2600A44EC6 /* TipKitDebugOptionsUIActionHandling.swift */; };
@@ -545,6 +545,7 @@
85A1B3B220C6CD9900C18F15 /* MigratableCookieStorage.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85A1B3B120C6CD9900C18F15 /* MigratableCookieStorage.swift */; };
85A313972028E78A00327D00 /* release_notes.txt in Resources */ = {isa = PBXBuildFile; fileRef = 85A313962028E78A00327D00 /* release_notes.txt */; };
85A9C37920E0E00C00073340 /* HomeRow.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 85A9C37820E0E00C00073340 /* HomeRow.xcassets */; };
85AB84C12CF5DDB4007E679F /* WebsiteDataStoreCleaning.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85AB84C02CF5DDAF007E679F /* WebsiteDataStoreCleaning.swift */; };
85AD49EE2B6149110085D2D1 /* CookieStorageTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85AD49ED2B6149110085D2D1 /* CookieStorageTests.swift */; };
85AE668E2097206E0014CF04 /* NotificationView.xib in Resources */ = {isa = PBXBuildFile; fileRef = 85AE668D2097206E0014CF04 /* NotificationView.xib */; };
85AE6690209724120014CF04 /* NotificationView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 85AE668F209724120014CF04 /* NotificationView.swift */; };
@@ -1683,9 +1684,9 @@
6FEC0B842C999352006B4F6E /* FavoriteItem.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FavoriteItem.swift; sourceTree = "<group>"; };
6FEC0B872C999961006B4F6E /* FavoritesListInteractingAdapter.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FavoritesListInteractingAdapter.swift; sourceTree = "<group>"; };
6FF915802B88E0750042AC87 /* AdAttributionFetcherTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AdAttributionFetcherTests.swift; sourceTree = "<group>"; };
6FF9AD442CE766F700C5A406 /* NewTabPageControllerPixelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NewTabPageControllerPixelTests.swift; sourceTree = "<group>"; };
6FF9AD3E2CE63DC200C5A406 /* TabSwitcherOpenDailyPixel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TabSwitcherOpenDailyPixel.swift; sourceTree = "<group>"; };
6FF9AD402CE6610600C5A406 /* TabSwitcherDailyPixelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TabSwitcherDailyPixelTests.swift; sourceTree = "<group>"; };
6FF9AD442CE766F700C5A406 /* NewTabPageControllerPixelTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NewTabPageControllerPixelTests.swift; sourceTree = "<group>"; };
7B1604E72CB685B400A44EC6 /* Logger+TipKit.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Logger+TipKit.swift"; sourceTree = "<group>"; };
7B1604EB2CB68BDA00A44EC6 /* TipKitController+ConvenienceInitializers.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "TipKitController+ConvenienceInitializers.swift"; sourceTree = "<group>"; };
7B1604ED2CB68D2600A44EC6 /* TipKitDebugOptionsUIActionHandling.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TipKitDebugOptionsUIActionHandling.swift; sourceTree = "<group>"; };
@@ -1863,6 +1864,7 @@
85A313962028E78A00327D00 /* release_notes.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; name = release_notes.txt; path = fastlane/metadata/default/release_notes.txt; sourceTree = "<group>"; };
85A53EC9200D1FA20010D13F /* FileStore.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FileStore.swift; sourceTree = "<group>"; };
85A9C37820E0E00C00073340 /* HomeRow.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = HomeRow.xcassets; sourceTree = "<group>"; };
85AB84C02CF5DDAF007E679F /* WebsiteDataStoreCleaning.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WebsiteDataStoreCleaning.swift; sourceTree = "<group>"; };
85AD49ED2B6149110085D2D1 /* CookieStorageTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CookieStorageTests.swift; sourceTree = "<group>"; };
85AE668D2097206E0014CF04 /* NotificationView.xib */ = {isa = PBXFileReference; lastKnownFileType = file.xib; path = NotificationView.xib; sourceTree = "<group>"; };
85AE668F209724120014CF04 /* NotificationView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NotificationView.swift; sourceTree = "<group>"; };
@@ -6029,6 +6031,7 @@
F143C3311E4A9A6A00CFDE3A /* Web */ = {
isa = PBXGroup;
children = (
85AB84C02CF5DDAF007E679F /* WebsiteDataStoreCleaning.swift */,
85C503FA2CF0ADBD0075DF6F /* WebCacheManager+ObservationsDataClearing.swift */,
85A1B3B120C6CD9900C18F15 /* MigratableCookieStorage.swift */,
85C503F82CEF78C00075DF6F /* DataStoreIdManager.swift */,
@@ -8368,6 +8371,7 @@
85E065BC2C73A54700D73E2A /* UsageSegmentationStorage.swift in Sources */,
F143C3271E4A9A0E00CFDE3A /* Logger+Multiple.swift in Sources */,
85372447220DD103009D09CD /* UIKeyCommandExtension.swift in Sources */,
85AB84C12CF5DDB4007E679F /* WebsiteDataStoreCleaning.swift in Sources */,
85A1B3B220C6CD9900C18F15 /* MigratableCookieStorage.swift in Sources */,
9856A1992933D2EB00ACB44F /* BookmarksModelsErrorHandling.swift in Sources */,
850559D023CF647C0055C0D5 /* Fireproofing.swift in Sources */,
75 changes: 29 additions & 46 deletions DuckDuckGoTests/WebCacheManagerTests.swift
Original file line number Diff line number Diff line change
@@ -22,8 +22,6 @@ import XCTest
import WebKit
import TestUtils

// TODO DataStoreIDManager assert new install has nil id for container

class WebCacheManagerTests: XCTestCase {

let defaults = UserDefaults(suiteName: "Test")!
@@ -32,6 +30,7 @@ class WebCacheManagerTests: XCTestCase {
lazy var cookieStorage = MigratableCookieStorage(userDefaults: defaults)
lazy var fireproofing = MockFireproofing()
lazy var dataStoreIdManager = DataStoreIdManager(store: dataStoreIdStore)
let dataStoreCleaner = MockDataStoreCleaner()

override func setUp() {
super.setUp()
@@ -78,10 +77,15 @@ class WebCacheManagerTests: XCTestCase {
XCTAssertTrue(cookies.contains(where: { $0.domain == ".example.com" }))
}

// TODO create an abstraction for WKWebsiteDataStore management and pass it into the webcache manager
@MainActor
@available(iOS 17, *)
func test_WhenClearingDataAfterUsingContainer_ThenCookiesAreMigratedAndOldContainersAreRemoved() async {
// Mock having a single container so we can validate cleaning it gets called
dataStoreCleaner.countContainersReturnValue = 1

// Mock a data store id to force migration to happen
dataStoreIdStore.store = [DataStoreIdManager.Constants.currentWebContainerId.rawValue: UUID().uuidString]
dataStoreIdManager = DataStoreIdManager(store: dataStoreIdStore)

fireproofing = MockFireproofing(domains: ["example.com"])

MigratableCookieStorage.addCookies([
@@ -90,54 +94,29 @@ class WebCacheManagerTests: XCTestCase {
.make(name: "Test3", value: "Value", domain: "facebook.com"),
], defaults)

// Setup a new container and add something to it just to make it real
let uuid = await createContainer()

var dataStoreIds = await WKWebsiteDataStore.allDataStoreIdentifiers
XCTAssertEqual(1, dataStoreIds.count)

// Use the UUID we just created for the store id
dataStoreIdStore.store = [DataStoreIdManager.Constants.currentWebContainerId.rawValue: uuid.uuidString]
dataStoreIdManager = DataStoreIdManager(store: dataStoreIdStore)

let dataStore = WKWebsiteDataStore.default()
let dataStore = await WKWebsiteDataStore.default()
var cookies = await dataStore.httpCookieStore.allCookies()
XCTAssertEqual(0, cookies.count)

let webCacheManager = makeWebCacheManager()
let webCacheManager = await makeWebCacheManager()
await webCacheManager.clear(dataStore: dataStore)

try? await Task.sleep(interval: 0.3)

cookies = await dataStore.httpCookieStore.allCookies()
XCTAssertEqual(2, cookies.count)
XCTAssertTrue(cookies.contains(where: { $0.domain == "example.com" }))
XCTAssertTrue(cookies.contains(where: { $0.domain == ".example.com" }))

dataStoreIds = await WKWebsiteDataStore.allDataStoreIdentifiers
XCTAssertTrue(dataStoreIds.isEmpty)
XCTAssertEqual(1, dataStoreCleaner.removeAllContainersAfterDelayCalls.count)
XCTAssertEqual(1, dataStoreCleaner.removeAllContainersAfterDelayCalls[0])
}

@MainActor
@available(iOS 17, *)
func test_WhenClearingData_ThenOldContainersAreRemoved() async {
_ = await createContainer()
_ = await createContainer()
_ = await createContainer()
_ = await createContainer()
_ = await createContainer()

try? await Task.sleep(interval: 0.3)

var dataStoreIds = await WKWebsiteDataStore.allDataStoreIdentifiers
XCTAssertEqual(5, dataStoreIds.count)

// Mock existence of 5 containers so we can validate that cleaning it is called even without migrations
dataStoreCleaner.countContainersReturnValue = 5
await makeWebCacheManager().clear(dataStore: .default())

try? await Task.sleep(interval: 0.3)

dataStoreIds = await WKWebsiteDataStore.allDataStoreIdentifiers
XCTAssertEqual(0, dataStoreIds.count)
XCTAssertEqual(1, dataStoreCleaner.removeAllContainersAfterDelayCalls.count)
XCTAssertEqual(5, dataStoreCleaner.removeAllContainersAfterDelayCalls[0])
}

/// Temporarily disabled.
@@ -153,18 +132,22 @@ class WebCacheManagerTests: XCTestCase {
cookieStorage: cookieStorage,
fireproofing: fireproofing,
dataStoreIdManager: dataStoreIdManager,
dataStoreClearingDelay: 0.01
dataStoreCleaner: dataStoreCleaner
)
}
}

@available(iOS 17, *)
@MainActor private func createContainer() async -> UUID {
let uuid = UUID()
let containerStore = WKWebsiteDataStore(forIdentifier: uuid)
await containerStore.httpCookieStore.setCookie(.make(name: "Not", value: "Used"))
let cookies = await containerStore.httpCookieStore.allCookies()
XCTAssertEqual(1, cookies.count)
return uuid
class MockDataStoreCleaner: WebsiteDataStoreCleaning {

var countContainersReturnValue = 0
var removeAllContainersAfterDelayCalls: [Int] = []

func countContainers() async -> Int {
return countContainersReturnValue
}

func removeAllContainersAfterDelay(previousCount: Int) {
removeAllContainersAfterDelayCalls.append(previousCount)
}

}

0 comments on commit 7ec9cde

Please sign in to comment.