Fix threading issue between using a Semaphore and async/await

Sources/BrowserServicesKit/ContentBlocking/ContentBlockerRulesManager.swift

@@ -242,26 +242,20 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
         Logger.contentBlocking.debug("Lookup compiled rules")
         prepareSourceManagers()
         let initialCompilationTask = LookupRulesTask(sourceManagers: Array(sourceManagers.values))
-        let mutex = DispatchSemaphore(value: 0)
 
-        Task {
-            do {
-                try await initialCompilationTask.lookupCachedRulesLists()
-            } catch {
-                Logger.contentBlocking.debug("❌ Lookup failed: \(error.localizedDescription, privacy: .public)")
-            }
-            mutex.signal()
-        }
-        // We want to confine Compilation work to WorkQueue, so we wait to come back from async Task
-        mutex.wait()
+        let result: [LookupRulesTask.LookupResult]
+
+        do {
+            result = try initialCompilationTask.lookupCachedRulesLists()
 
-        if let result = initialCompilationTask.result {
             let rules = result.map(Rules.init(compilationResult:))
-            Logger.contentBlocking.debug("🟩 Found \(rules.count, privacy: .public) rules")
+            Logger.contentBlocking.debug("🟩 Lookup Found \(rules.count, privacy: .public) rules")
             applyRules(rules)
             return true
+        } catch {
+            Logger.contentBlocking.debug("❌ Lookup failed: \(error.localizedDescription, privacy: .public)")
+            return false
         }
-        return false
     }
 
     /*
@@ -273,18 +267,10 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
 
         let initialCompilationTask = LastCompiledRulesLookupTask(sourceRules: rulesSource.contentBlockerRulesLists,
                                                                  lastCompiledRules: lastCompiledRules)
-        let mutex = DispatchSemaphore(value: 0)
-        Task {
-            try? await initialCompilationTask.fetchCachedRulesLists()
-            mutex.signal()
-        }
-        // We want to confine Compilation work to WorkQueue, so we wait to come back from async Task
-        mutex.wait()
-
-        let rulesFound = initialCompilationTask.getFetchedRules()
+        let rules = initialCompilationTask.fetchCachedRulesLists()
 
-        if let rulesFound {
-            applyRules(rulesFound)
+        if let rules {
+            applyRules(rules)
         } else {
             lock.lock()
             state = .idle
@@ -294,7 +280,7 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
         // No matter if rules were found or not, we need to schedule recompilation, after all
         scheduleCompilation()
 
-        return rulesFound != nil
+        return rules != nil
     }
 
     private func prepareSourceManagers() {

Sources/BrowserServicesKit/ContentBlocking/ContentBlockingRulesLastCompiledRulesLookupTask.swift

@@ -41,32 +41,49 @@ extension ContentBlockerRulesManager {
             self.lastCompiledRules = lastCompiledRules
         }
 
-        func fetchCachedRulesLists() async throws {
+        func fetchCachedRulesLists() -> [Rules]? {
             let sourceRulesNames = sourceRules.map { $0.name }
             let filteredBySourceLastCompiledRules = lastCompiledRules.filter { sourceRulesNames.contains($0.name) }
 
             guard filteredBySourceLastCompiledRules.count == sourceRules.count else {
                 // We should only load rule lists from cache, in case we can match every one of these
-                throw WKError(.contentRuleListStoreLookUpFailed)
+                return nil
             }
 
             var result: [CachedRulesList] = []
+            let group = DispatchGroup()
+
             for rules in filteredBySourceLastCompiledRules {
-                guard let ruleList = try await Task(operation: { @MainActor in
-                    try await WKContentRuleListStore.default().contentRuleList(forIdentifier: rules.identifier.stringValue)
-                }).value else { throw WKError(.contentRuleListStoreLookUpFailed) }
-
-                result.append(CachedRulesList(name: rules.name,
-                                              rulesList: ruleList,
-                                              tds: rules.trackerData,
-                                              rulesIdentifier: rules.identifier))
+                group.enter()
+
+                DispatchQueue.main.async {
+                    // This needs to be called from the main thread.
+                    WKContentRuleListStore.default().lookUpContentRuleList(forIdentifier: rules.identifier.stringValue) { ruleList, error in
+                        guard let ruleList, error == nil else {
+                            group.leave()
+                            return
+                        }
+
+                        result.append(CachedRulesList(name: rules.name,
+                                                      rulesList: ruleList,
+                                                      tds: rules.trackerData,
+                                                      rulesIdentifier: rules.identifier))
+                        group.leave()
+                    }
+                }
             }
-            self.result = result
+
+            let operationResult = group.wait(timeout: .now() + 6)
+
+            guard operationResult == .success, result.count == filteredBySourceLastCompiledRules.count else {
+                return nil
+            }
+
+            return getRules(from: result)
         }
 
-        public func getFetchedRules() -> [Rules]? {
-            guard let result else { return nil }
-            return result.map {
+        public func getRules(from cached: [CachedRulesList]) -> [Rules] {
+            return cached.map {
                 let surrogateTDS = ContentBlockerRulesManager.extractSurrogates(from: $0.tds)
                 let encodedData = try? JSONEncoder().encode(surrogateTDS)
                 let encodedTrackerData = String(data: encodedData!, encoding: .utf8)!

Sources/BrowserServicesKit/ContentBlocking/ContentBlockingRulesLookupTask.swift

@@ -28,31 +28,45 @@ extension ContentBlockerRulesManager {
 
         private let sourceManagers: [ContentBlockerRulesSourceManager]
 
-        public private(set) var result: [LookupResult]?
-
         init(sourceManagers: [ContentBlockerRulesSourceManager]) {
             self.sourceManagers = sourceManagers
         }
 
-        func lookupCachedRulesLists() async throws {
+        func lookupCachedRulesLists() throws -> [LookupResult]  {
+
+            let models = sourceManagers.compactMap { $0.makeModel() }
+            if models.count != sourceManagers.count {
+                // We should only load rule lists, in case we can match every one of the expected ones
+                throw WKError(.contentRuleListStoreLookUpFailed)
+            }
 
             var result = [LookupResult]()
-            for sourceManager in sourceManagers {
-                guard let model = sourceManager.makeModel() else {
-                    throw WKError(.contentRuleListStoreLookUpFailed)
-                }
+            let group = DispatchGroup()
+
+            for model in models {
+                group.enter()
 
-                guard let ruleList = try await Task(operation: { @MainActor in
-                    try await WKContentRuleListStore.default().contentRuleList(forIdentifier: model.rulesIdentifier.stringValue)
-                }).value else {
-                    // All lists must be found for this to be considered successful
-                    throw WKError(.contentRuleListStoreLookUpFailed)
+                DispatchQueue.main.async {
+                    // This needs to be called from the main thread.
+                    WKContentRuleListStore.default().lookUpContentRuleList(forIdentifier: model.rulesIdentifier.stringValue) { ruleList, error in
+                        guard let ruleList, error == nil else {
+                            group.leave()
+                            return
+                        }
+
+                        result.append((ruleList, model))
+                        group.leave()
+                    }
                 }
+            }
+
+            let operationResult = group.wait(timeout: .now() + 6)
 
-                result.append((ruleList, model))
+            guard operationResult == .success, result.count == models.count else {
+                throw WKError(.contentRuleListStoreLookUpFailed)
             }
-            self.result = result
-        }
 
+            return result
+        }
     }
 }