Merge pull request #134 from vazarkevych/fixing-sticky-bucketing

vazarkevych · web-flow · commit f95faa9b49d5 · 2025-11-11T09:20:51.000+02:00
FeatureViewModel data racing fix
diff --git a/GrowthBookTests/FeaturesViewModelTests.swift b/GrowthBookTests/FeaturesViewModelTests.swift
@@ -22,7 +22,7 @@ class FeaturesViewModelTests: XCTestCase, FeaturesFlowDelegate {
         isSuccess = false
         isError = true
 
-        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponse, error: nil)), cachingManager: cachingManager)
+        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponse, error: nil)), cachingManager: cachingManager, forceSynchronousSave: true)
 
         viewModel.fetchFeatures(apiUrl: "")
 
@@ -35,7 +35,7 @@ class FeaturesViewModelTests: XCTestCase, FeaturesFlowDelegate {
         isSuccess = false
         isError = true
         
-        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponseEncryptedFeatures, error: nil)), cachingManager: cachingManager)
+        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponseEncryptedFeatures, error: nil)), cachingManager: cachingManager, forceSynchronousSave: true)
         
         viewModel.encryptionKey = "3tfeoyW0wlo47bDnbWDkxg=="
         viewModel.fetchFeatures(apiUrl: "")
@@ -48,7 +48,7 @@ class FeaturesViewModelTests: XCTestCase, FeaturesFlowDelegate {
         isSuccess = false
         isError = true
         
-        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponse, error: nil)), cachingManager: cachingManager)
+        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponse, error: nil)), cachingManager: cachingManager, forceSynchronousSave: true)
         
         viewModel.fetchFeatures(apiUrl: "")
         
@@ -83,7 +83,7 @@ class FeaturesViewModelTests: XCTestCase, FeaturesFlowDelegate {
         isSuccess = false
         isError = true
         
-        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponseEncryptedFeatures, error: nil)), cachingManager: cachingManager)
+        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().successResponseEncryptedFeatures, error: nil)), cachingManager: cachingManager, forceSynchronousSave: true)
         
         let encryptionKey = "3tfeoyW0wlo47bDnbWDkxg=="
         viewModel.encryptionKey = encryptionKey
@@ -110,7 +110,7 @@ class FeaturesViewModelTests: XCTestCase, FeaturesFlowDelegate {
     func testError() throws {
         isSuccess = false
         isError = true
-        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: nil, error: .failedToLoadData)), cachingManager: cachingManager)
+        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: nil, error: .failedToLoadData)), cachingManager: cachingManager, forceSynchronousSave: true)
 
         viewModel.fetchFeatures(apiUrl: "")
 
@@ -122,7 +122,7 @@ class FeaturesViewModelTests: XCTestCase, FeaturesFlowDelegate {
     func testInvalid() throws {
         isSuccess = false
         isError = true
-        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().errorResponse, error: nil)), cachingManager: cachingManager)
+        let viewModel = FeaturesViewModel(delegate: self, dataSource: FeaturesDataSource(dispatcher: MockNetworkClient(successResponse: MockResponse().errorResponse, error: nil)), cachingManager: cachingManager, forceSynchronousSave: true)
         viewModel.fetchFeatures(apiUrl: "")
 
         XCTAssertFalse(isSuccess)
diff --git a/Sources/CommonMain/Features/FeaturesViewModel.swift b/Sources/CommonMain/Features/FeaturesViewModel.swift
@@ -12,20 +12,27 @@ protocol FeaturesFlowDelegate: AnyObject {
 /// View Model for Features
 class FeaturesViewModel {
     weak var delegate: FeaturesFlowDelegate?
-    let dataSource: FeaturesDataSource
     var encryptionKey: String?
-    /// Caching Manager
-    let manager: CachingLayer
-    /// SSE Handler for background sync
-    internal var sseHandler: SSEHandler?
+    
+    private let dataSource: FeaturesDataSource
+    private let manager: CachingLayer
+    private var sseHandler: SSEHandler?
+    private let fileSaveQueue = DispatchQueue(label: "com.sdk.fileSaveQueue", qos: .utility)
+    
+    var forceSynchronousSave: Bool
         
-    init(delegate: FeaturesFlowDelegate, dataSource: FeaturesDataSource, cachingManager: CachingLayer) {
+    init(delegate: FeaturesFlowDelegate, dataSource: FeaturesDataSource, cachingManager: CachingLayer, forceSynchronousSave: Bool = false) {
         self.delegate = delegate
         self.dataSource = dataSource
         self.manager = cachingManager
+        self.forceSynchronousSave = forceSynchronousSave
         self.fetchCachedFeatures()
     }
     
+    deinit {
+        sseHandler?.disconnect()
+    }
+        
     func connectBackgroundSync(sseUrl: String) {
         guard let url = URL(string: sseUrl) else { return }
         
@@ -35,144 +42,171 @@ class FeaturesViewModel {
         let streamingUpdate = SSEHandler(url: url)
         sseHandler = streamingUpdate
         
-        streamingUpdate.addEventListener(event: "features") { [weak self] id, event, data in
-            guard let jsonData = data?.data(using: .utf8) else { return }
-            self?.prepareFeaturesData(data: jsonData)
+        streamingUpdate.addEventListener(event: "features") { [weak self] _, _, data in
+            guard let self, let jsonData = data?.data(using: .utf8) else { return }
+            self.prepareFeaturesData(data: jsonData)
         }
         streamingUpdate.connect()
         
         streamingUpdate.onDissconnect { [weak streamingUpdate] _, shouldReconnect, _ in
-            if let shouldReconnect = shouldReconnect, shouldReconnect {
+            if shouldReconnect == true {
                 streamingUpdate?.connect()
             }
         }
     }
-    
-    deinit {
-        sseHandler?.disconnect()
-    }
-    
-    private func fetchCachedFeatures(logging: Bool = false) {
-        // Check for cache data
-        if let data = manager.getContent(fileName: Constants.featureCache) {
-            let decoder = JSONDecoder()
-            if let encryptedString = String(data: data, encoding: .utf8), let encryptionKey, !encryptionKey.isEmpty {
-                let crypto: CryptoProtocol = Crypto()
-                if let features = crypto.getFeaturesFromEncryptedFeatures(encryptedString: encryptedString, encryptionKey: encryptionKey) {
-                    delegate?.featuresFetchedSuccessfully(features: features, isRemote: false)
-                } else {
-                    delegate?.featuresFetchFailed(error: .failedParsedEncryptedData, isRemote: false)
-                    if logging { logger.error("Failed get features from cached encrypted features") }
-                }
-            } else if let features = try? decoder.decode(Features.self, from: data) {
-                // Call Success Delegate with mention of data available but its not remote
-                delegate?.featuresFetchedSuccessfully(features: features, isRemote: false)
-            } else {
-                delegate?.featuresFetchFailed(error: .failedParsedData, isRemote: false)
-                if logging { logger.error("Failed parse local data") }
+        
+    func fetchFeatures(apiUrl: String?, remoteEval: Bool = false, payload: RemoteEvalParams? = nil) {
+        fetchCachedFeatures(logging: true)
+        
+        guard let apiUrl else {
+            logger.error("Missing API URL")
+            notify { $0.featuresFetchFailed(error: .failedMissingKey, isRemote: true) }
+            return
+        }
+        
+        let completion: (Result<Data, Error>) -> Void = { [weak self] result in
+            guard let self else { return }
+            switch result {
+            case .success(let data):
+                self.prepareFeaturesData(data: data)
+            case .failure(let error):
+                logger.error("Failed to fetch features: \(error.localizedDescription)")
+                self.notify { $0.featuresFetchFailed(error: .failedToLoadData, isRemote: true) }
             }
+        }
+        
+        if remoteEval {
+            dataSource.fetchRemoteEval(apiUrl: apiUrl, params: payload, fetchResult: completion)
         } else {
-            delegate?.featuresFetchFailed(error: .failedToLoadData, isRemote: false)
-            if logging { logger.info("Cache directory is empty. Nothing to fetch.") }
+            dataSource.fetchFeatures(apiUrl: apiUrl, fetchResult: completion)
         }
     }
-
-    /// Fetch Features
-    func fetchFeatures(apiUrl: String?, remoteEval: Bool = false, payload: RemoteEvalParams? = nil) {
-        // Check for cache data
-        fetchCachedFeatures(logging: true)
         
-        if let apiUrl = apiUrl {
-            if remoteEval {
-                dataSource.fetchRemoteEval(apiUrl: apiUrl, params: payload) { result in
-                    switch result {
-                    case .success(let data):
-                        self.prepareFeaturesData(data: data)
-                    case .failure(let error):
-                        self.delegate?.featuresFetchFailed(error: .failedToLoadData, isRemote: true)
-                        logger.error("Failed get features: \(error.localizedDescription)")
-                    }
-                }
+    private func fetchCachedFeatures(logging: Bool = false) {
+        guard let data = manager.getContent(fileName: Constants.featureCache) else {
+            if logging { logger.info("Cache directory is empty. Nothing to fetch.") }
+            notify { $0.featuresFetchFailed(error: .failedToLoadData, isRemote: false) }
+            return
+        }
+        
+        let decoder = JSONDecoder()
+        
+        if let encryptedString = String(data: data, encoding: .utf8),
+           let encryptionKey, !encryptionKey.isEmpty {
+            
+            let crypto: CryptoProtocol = Crypto()
+            if let features = crypto.getFeaturesFromEncryptedFeatures(encryptedString: encryptedString, encryptionKey: encryptionKey) {
+                notify { $0.featuresFetchedSuccessfully(features: features, isRemote: false) }
             } else {
-                dataSource.fetchFeatures(apiUrl: apiUrl) { result in
-                    switch result {
-                    case .success(let data):
-                        self.prepareFeaturesData(data: data)
-                    case .failure(let error):
-                        self.delegate?.featuresFetchFailed(error: .failedToLoadData, isRemote: true)
-                        logger.error("Failed get features: \(error.localizedDescription)")
-                    }
-                }
+                if logging { logger.error("Failed get features from cached encrypted features") }
+                notify { $0.featuresFetchFailed(error: .failedParsedEncryptedData, isRemote: false) }
             }
+            
+        } else if let features = try? decoder.decode(Features.self, from: data) {
+            notify { $0.featuresFetchedSuccessfully(features: features, isRemote: false) }
         } else {
-            delegate?.featuresFetchFailed(error: .failedMissingKey, isRemote: true)
-            logger.error("Failed get api URL")
+            if logging { logger.error("Failed to parse local data") }
+            notify { $0.featuresFetchFailed(error: .failedParsedData, isRemote: false) }
         }
     }
-
-    /// Cache API Response and push success event
-    func prepareFeaturesData(data: Data) {
-        // Call Success Delegate with mention of data available with remote
         
+    func prepareFeaturesData(data: Data) {
         let decoder = JSONDecoder()
-        if let jsonPetitions = try? decoder.decode(FeaturesDataModel.self, from: data) {
-            delegate?.featuresAPIModelSuccessfully(model: jsonPetitions)
-            if let encryptedString = jsonPetitions.encryptedFeatures {
-                if let encryptionKey = encryptionKey, !encryptionKey.isEmpty {
-                    let crypto: CryptoProtocol = Crypto()
-                    if let features = crypto.getFeaturesFromEncryptedFeatures(encryptedString: encryptedString, encryptionKey: encryptionKey) {
-                        if let featureData = encryptedString.data(using: .utf8) {
-                            manager.saveContent(fileName: Constants.featureCache, content: featureData)
-                        } else {
-                            logger.error("Failed encode features")
-                        }
-                        delegate?.featuresFetchedSuccessfully(features: features, isRemote: true)
-                    } else {
-                        delegate?.featuresFetchFailed(error: .failedEncryptedFeatures, isRemote: true)
-                        logger.error("Failed get features from encrypted features")
-                        return
-                    }
-                } else {
-                    delegate?.featuresFetchFailed(error: .failedMissingKey, isRemote: true)
-                    logger.error("Failed get encryption key or it's empty")
-                    return
-                }
-            } else if let features = jsonPetitions.features {
-                if let featureData = try? JSONEncoder().encode(features) {
-                    manager.saveContent(fileName: Constants.featureCache, content: featureData)
+        
+        guard let jsonPetitions = try? decoder.decode(FeaturesDataModel.self, from: data) else {
+            logger.error("Failed to decode FeaturesDataModel")
+            notify { $0.featuresFetchFailed(error: .failedParsedData, isRemote: true) }
+            return
+        }
+        
+        if let encryptedString = jsonPetitions.encryptedFeatures {
+            handleEncryptedFeatures(encryptedString: encryptedString, jsonPetitions: jsonPetitions)
+        } else if let features = jsonPetitions.features {
+            handlePlainFeatures(features, jsonPetitions: jsonPetitions)
+        } else {
+            logger.error("Missing both encrypted and plain features")
+            notify { $0.featuresFetchFailed(error: .failedMissingKey, isRemote: true) }
+        }
+    }
+        
+    private func handleEncryptedFeatures(encryptedString: String, jsonPetitions: FeaturesDataModel) {
+        guard let encryptionKey = encryptionKey, !encryptionKey.isEmpty else {
+            logger.error("Missing encryption key")
+            notify { $0.featuresFetchFailed(error: .failedMissingKey, isRemote: true) }
+            return
+        }
+        
+        let crypto = Crypto()
+        guard let features = crypto.getFeaturesFromEncryptedFeatures(encryptedString: encryptedString, encryptionKey: encryptionKey) else {
+            logger.error("Failed to decrypt features")
+            notify { $0.featuresFetchFailed(error: .failedEncryptedFeatures, isRemote: true) }
+            return
+        }
+        
+        if let featureData = encryptedString.data(using: .utf8) {
+            saveDataThreadSafe(fileName: Constants.featureCache, content: featureData)
+        } else {
+            logger.error("Failed to encode features as UTF-8")
+        }
+        
+        notify { $0.featuresFetchedSuccessfully(features: features, isRemote: true) }
+        handleSavedGroups(from: jsonPetitions)
+    }
+    
+    private func handlePlainFeatures(_ features: Features, jsonPetitions: FeaturesDataModel) {
+        if let featureData = try? JSONEncoder().encode(features) {
+            saveDataThreadSafe(fileName: Constants.featureCache, content: featureData)
+        }
+        
+        notify { $0.featuresFetchedSuccessfully(features: features, isRemote: true) }
+        handleSavedGroups(from: jsonPetitions)
+    }
+    
+    private func handleSavedGroups(from jsonPetitions: FeaturesDataModel) {
+        if let encryptedSavedGroups = jsonPetitions.encryptedSavedGroups,
+           !encryptedSavedGroups.isEmpty,
+           let encryptionKey = encryptionKey,
+           !encryptionKey.isEmpty {
+            
+            let crypto = Crypto()
+            if let savedGroups = crypto.getSavedGroupsFromEncryptedFeatures(encryptedString: encryptedSavedGroups, encryptionKey: encryptionKey) {
+                if let savedGroupsData = encryptedSavedGroups.data(using: .utf8) {
+                    saveDataThreadSafe(fileName: Constants.savedGroupsCache, content: savedGroupsData)
                 }
-                delegate?.featuresFetchedSuccessfully(features: features, isRemote: true)
+                notify { $0.savedGroupsFetchedSuccessfully(savedGroups: savedGroups, isRemote: true) }
             } else {
-                delegate?.featuresFetchFailed(error: .failedMissingKey, isRemote: true)
-                logger.error("Failed get encrypted features or it's empty")
-                return
+                logger.error("Failed to decrypt saved groups")
+                notify { $0.savedGroupsFetchFailed(error: .failedEncryptedSavedGroups, isRemote: true) }
             }
             
-            if let encryptedSavedGroups = jsonPetitions.encryptedSavedGroups, !encryptedSavedGroups.isEmpty, let encryptionKey = encryptionKey, !encryptionKey.isEmpty {
-                let crypto = Crypto()
-                if let savedGroups = crypto.getSavedGroupsFromEncryptedFeatures(encryptedString: encryptedSavedGroups, encryptionKey: encryptionKey) {
-                    if let encryptedSavedGroups = encryptedSavedGroups.data(using: .utf8) {
-                        manager.saveContent(fileName: Constants.savedGroupsCache, content: encryptedSavedGroups)
-                    } else {
-                        logger.error("Failed encode saved groups")
-                    }
-                    delegate?.savedGroupsFetchedSuccessfully(savedGroups: savedGroups, isRemote: true)
-                } else {
-                    delegate?.savedGroupsFetchFailed(error: .failedEncryptedSavedGroups, isRemote: true)
-                    logger.error("Failed get saved groups from encrypted saved groups")
-                    return
-                }
-            } else if let savedGroups = jsonPetitions.savedGroups {
-                if let savedGroupsData = try? JSONEncoder().encode(savedGroups) {
-                    manager.saveContent(fileName: Constants.savedGroupsCache, content: savedGroupsData)
-                }
-                delegate?.savedGroupsFetchedSuccessfully(savedGroups: savedGroups, isRemote: true)
+        } else if let savedGroups = jsonPetitions.savedGroups {
+            if let savedGroupsData = try? JSONEncoder().encode(savedGroups) {
+                saveDataThreadSafe(fileName: Constants.savedGroupsCache, content: savedGroupsData)
             }
-        } else {
-            delegate?.featuresFetchFailed(error: .failedParsedData, isRemote: true)
-            logger.error("Failed get features data model")
+            notify { $0.savedGroupsFetchedSuccessfully(savedGroups: savedGroups, isRemote: true) }
+        }
+    }
+        
+    private func saveDataThreadSafe(fileName: String, content: Data) {
+        if forceSynchronousSave {
+            manager.saveContent(fileName: fileName, content: content)
             return
+        } else {
+            fileSaveQueue.async { [weak self] in
+                self?.manager.saveContent(fileName: fileName, content: content)
+            }
         }
     }
         
+    private func notify(_ action: @escaping (FeaturesFlowDelegate) -> Void) {
+        if Thread.isMainThread {
+            guard let delegate = self.delegate else { return }
+            action(delegate)
+        } else {
+            DispatchQueue.main.async { [weak self] in
+                guard let self, let delegate = self.delegate else { return }
+                action(delegate)
+            }
+        }
+    }
 }
