Fix state restoration after app termination (#3127)

Task/Issue URL: https://app.asana.com/0/1202406491309510/1208083933206713/f

DuckDuckGo/Common/FileSystem/FileStore.swift

@@ -28,6 +28,7 @@ protocol FileStore {
     func hasData(at url: URL) -> Bool
     func directoryContents(at path: String) throws -> [String]
     func remove(fileAtURL url: URL)
+    func move(fileAt from: URL, to: URL)
 }
 
 extension FileStore {
@@ -91,6 +92,10 @@ final class EncryptedFileStore: FileStore {
         try? fileManager.removeItem(at: url)
     }
 
+    func move(fileAt from: URL, to: URL) {
+        try? FileManager.default.moveItem(at: from, to: to)
+    }
+
 }
 
 extension FileManager: FileStore {
@@ -125,4 +130,8 @@ extension FileManager: FileStore {
         try? removeItem(at: url)
     }
 
+    func move(fileAt from: URL, to: URL) {
+        try? moveItem(at: from, to: to)
+    }
+
 }

DuckDuckGo/StateRestoration/AppStateRestorationManager.swift

@@ -67,7 +67,8 @@ final class AppStateRestorationManager: NSObject {
             try service.restoreState(using: { coder in
                 state = try WindowsManager.restoreState(from: coder, includePinnedTabs: isCalledAtStartup)
             })
-            clearLastSessionState()
+            // rename loaded app state file
+            service.didLoadState()
         } catch CocoaError.fileReadNoSuchFile {
             // ignore
         } catch {
@@ -80,7 +81,7 @@ final class AppStateRestorationManager: NSObject {
     }
 
     func clearLastSessionState() {
-        service.removeLastSessionState()
+        service.clearState(sync: true)
     }
 
     // Cleans all stored snapshots except snapshots listed in the state
@@ -94,8 +95,10 @@ final class AppStateRestorationManager: NSObject {
     }
 
     func applicationDidFinishLaunching() {
-        let isRelaunchingAutomatically = appIsRelaunchingAutomatically
-        appIsRelaunchingAutomatically = false
+        let isRelaunchingAutomatically = self.appIsRelaunchingAutomatically
+        self.appIsRelaunchingAutomatically = false
+        // don‘t automatically restore windows if relaunched 2nd time with no recently updated app session state
+        let shouldRestorePreviousSession = self.shouldRestorePreviousSession && !service.isAppStateFileStale
         readLastSessionState(restoreWindows: shouldRestorePreviousSession || isRelaunchingAutomatically)
 
         stateChangedCancellable = WindowControllersManager.shared.stateChanged

DuckDuckGo/StateRestoration/StatePersistenceService.swift

@@ -21,12 +21,30 @@ import Foundation
 final class StatePersistenceService {
     private let fileStore: FileStore
     private let fileName: String
+    ///  The `persistentState` file is renamed to `persistentState.1` after it‘s loaded for the first time
+    ///  if no new persistentState is written during the session it will be used on the next load and renamed to `persistentState.2`
+    private var lastLoadedStateFileName: String {
+        fileName + ".1"
+    }
+    private var oldStateFileName: String {
+        fileName + ".2"
+    }
     private var lastSessionStateArchive: Data?
     private let queue = DispatchQueue(label: "StateRestorationManager.queue", qos: .background)
     private var job: DispatchWorkItem?
 
     private(set) var error: Error?
 
+    /// `false` if `persistentState` or `persistentState.1` file exists,
+    /// `true` if `persistentState.2` (old app state that was not updated after 2nd app relaunch) file exists
+    var isAppStateFileStale: Bool {
+        if fileStore.hasData(at: .persistenceLocation(for: fileName)) || fileStore.hasData(at: .persistenceLocation(for: lastLoadedStateFileName)) {
+            return false
+        } else {
+            return true
+        }
+    }
+
     init(fileStore: FileStore, fileName: String) {
         self.fileStore = fileStore
         self.fileName = fileName
@@ -49,8 +67,7 @@ final class StatePersistenceService {
 
         job?.cancel()
         job = DispatchWorkItem {
-            let location = URL.persistenceLocation(for: self.fileName)
-            self.fileStore.remove(fileAtURL: location)
+            self.performClearState()
         }
         queue.dispatch(job!, sync: sync)
     }
@@ -63,9 +80,30 @@ final class StatePersistenceService {
         lastSessionStateArchive = loadStateFromFile()
     }
 
-    func removeLastSessionState() {
+    // perform state clearing synchronously, called from `clearState(sync:)` on `StateRestorationManager.queue`
+    func performClearState() {
         lastSessionStateArchive = nil
-        fileStore.remove(fileAtURL: URL.persistenceLocation(for: self.fileName))
+        let location = URL.persistenceLocation(for: self.fileName)
+        fileStore.remove(fileAtURL: location)
+        fileStore.remove(fileAtURL: .persistenceLocation(for: self.lastLoadedStateFileName))
+        fileStore.remove(fileAtURL: .persistenceLocation(for: self.oldStateFileName))
+    }
+
+    /// rename `persistentState` to `persistentState.1` after the state was loaded
+    /// if the state was loaded from `persistentState.1`, it will be renamed to `persistentState.2`
+    /// `persistentState.2` won‘t restore windows automatically to avoid a possible crash loop
+    func didLoadState() {
+        let location = URL.persistenceLocation(for: self.fileName)
+        let location1 = URL.persistenceLocation(for: self.lastLoadedStateFileName)
+        let location2 = URL.persistenceLocation(for: self.oldStateFileName)
+        if fileStore.hasData(at: location) {
+            fileStore.remove(fileAtURL: location1)
+            fileStore.remove(fileAtURL: location2)
+            fileStore.move(fileAt: location, to: location1)
+        } else if fileStore.hasData(at: location1) {
+            fileStore.remove(fileAtURL: location2)
+            fileStore.move(fileAt: location1, to: location2)
+        }
     }
 
     @MainActor
@@ -93,12 +131,16 @@ final class StatePersistenceService {
             if !self.fileStore.persist(data, url: location) {
                 self.error = CocoaError(.fileWriteNoPermission)
             }
+            self.fileStore.remove(fileAtURL: .persistenceLocation(for: self.lastLoadedStateFileName))
+            self.fileStore.remove(fileAtURL: .persistenceLocation(for: self.oldStateFileName))
         }
         queue.dispatch(job!, sync: sync)
     }
 
     private func loadStateFromFile() -> Data? {
         fileStore.loadData(at: URL.persistenceLocation(for: self.fileName), decryptIfNeeded: false)
+        ?? fileStore.loadData(at: .persistenceLocation(for: self.lastLoadedStateFileName), decryptIfNeeded: false)
+        ?? fileStore.loadData(at: .persistenceLocation(for: self.oldStateFileName), decryptIfNeeded: false)
     }
 
     @MainActor

UnitTests/Common/FileSystem/FileStoreMock.swift

@@ -93,4 +93,11 @@ final class FileStoreMock: NSObject, FileStore {
     func remove(fileAtURL url: URL) {
         storage[url.lastPathComponent] = nil
     }
+
+    func move(fileAt from: URL, to: URL) {
+        guard storage[to.lastPathComponent] == nil,
+              let data = storage.removeValue(forKey: from.lastPathComponent) else { return }
+        storage[to.lastPathComponent] = data
+    }
+
 }

UnitTests/Common/FileSystem/StateRestorationManagerTests.swift

@@ -104,17 +104,51 @@ final class StateRestorationManagerTests: XCTestCase {
     }
 
     @MainActor
-    func testWhenLastSessionStateIsRemovedManuallyThenLastSessionCannotBeRestored() {
+    func testWhenLastSessionStateIsClearedThenLastSessionCannotBeRestored() {
         changeState("val1", 1, sync: true)
 
         srm = StatePersistenceService(fileStore: fileStore, fileName: testFileName)
         srm.loadLastSessionState()
         XCTAssertTrue(srm.canRestoreLastSessionState)
 
-        srm.removeLastSessionState()
+        srm.clearState(sync: true)
         XCTAssertFalse(srm.canRestoreLastSessionState)
     }
 
+    @MainActor
+    func testWhenSessionStateIsRestoredItCanBeRestoredAgain() {
+        changeState("val1", 1, sync: true)
+
+        srm = StatePersistenceService(fileStore: fileStore, fileName: testFileName)
+        srm.loadLastSessionState()
+        srm.didLoadState()
+        XCTAssertTrue(srm.canRestoreLastSessionState)
+
+        srm = StatePersistenceService(fileStore: fileStore, fileName: testFileName)
+        srm.loadLastSessionState()
+        srm.didLoadState()
+        XCTAssertTrue(srm.canRestoreLastSessionState)
+    }
+
+    @MainActor
+    func testWhenSameSessionStateIsRestoredTwiceItBecomesStale() {
+        changeState("val1", 1, sync: true)
+
+        srm = StatePersistenceService(fileStore: fileStore, fileName: testFileName)
+        XCTAssertFalse(srm.isAppStateFileStale)
+        srm.loadLastSessionState()
+        srm.didLoadState()
+
+        srm = StatePersistenceService(fileStore: fileStore, fileName: testFileName)
+        XCTAssertFalse(srm.isAppStateFileStale)
+        srm.loadLastSessionState()
+        srm.didLoadState()
+        XCTAssertTrue(srm.isAppStateFileStale)
+
+        srm = StatePersistenceService(fileStore: fileStore, fileName: testFileName)
+        XCTAssertTrue(srm.isAppStateFileStale)
+    }
+
     @MainActor
     func testWhenLastSessionStateIsLoadedThenChangesToStatePreserveLoadedLastSessionState() {
         changeState("lastSessionValue", 42, sync: true)
@@ -151,9 +185,13 @@ final class StateRestorationManagerTests: XCTestCase {
     func testStatePersistenceThrottlesWrites() {
         fileStore.delay = 0.1 // write operations will sleep for 100ms
         var counter = 0
-        let observer = fileStore.publisher(for: \.storage).dropFirst().sink { _ in
-            counter += 1
-        }
+
+        let observer = fileStore.publisher(for: \.storage)
+            .dropFirst()
+            .removeDuplicates()
+            .sink { _ in
+                counter += 1
+            }
 
         changeState("val1", 1)
         changeState("val2", 2)
@@ -172,9 +210,12 @@ final class StateRestorationManagerTests: XCTestCase {
         fileStore.delay = 0.01 // write operations will sleep for 100ms
 
         var counter = 0
-        let observer = fileStore.publisher(for: \.storage).dropFirst().sink { _ in
-            counter += 1
-        }
+        let observer = fileStore.publisher(for: \.storage)
+            .dropFirst()
+            .removeDuplicates()
+            .sink { _ in
+                counter += 1
+            }
 
         changeState("val1", 1, sync: true)
         changeState("val2", 2, sync: true)