From a254140e1a1d0a78f8331352116b624f7b8c6647 Mon Sep 17 00:00:00 2001
From: bwaresiak <bartek@duckduckgo.com>
Date: Tue, 26 Mar 2024 10:44:51 +0100
Subject: [PATCH] Guard against Data Store warmup crash (#2626)

---
 Core/DataStoreWarmup.swift | 20 +++++++++++++++-----
 Core/PixelEvent.swift      | 10 ++++++++--
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/Core/DataStoreWarmup.swift b/Core/DataStoreWarmup.swift
index 79087b9fc2..0ef67ac0b6 100644
--- a/Core/DataStoreWarmup.swift
+++ b/Core/DataStoreWarmup.swift
@@ -34,26 +34,36 @@ public class DataStoreWarmup {
 
 private class BlockingNavigationDelegate: NSObject, WKNavigationDelegate {
 
-    let finished = PassthroughSubject<Void, Never>()
+    var finished: PassthroughSubject? = PassthroughSubject<Void, Never>()
 
     func webView(_ webView: WKWebView, decidePolicyFor navigationAction: WKNavigationAction) async -> WKNavigationActionPolicy {
         return .allow
     }
 
     func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
-        finished.send()
+        if let finished {
+            finished.send()
+            self.finished = nil
+        } else {
+            Pixel.fire(pixel: .webKitWarmupUnexpectedDidFinish, includedParameters: [.appVersion])
+        }
     }
 
     func webViewWebContentProcessDidTerminate(_ webView: WKWebView) {
         Pixel.fire(pixel: .webKitDidTerminateDuringWarmup)
-        // We won't get a `didFinish` if the webview crashes
-        finished.send()
+
+        if let finished {
+            finished.send()
+            self.finished = nil
+        } else {
+            Pixel.fire(pixel: .webKitWarmupUnexpectedDidTerminate, includedParameters: [.appVersion])
+        }
     }
 
     var cancellable: AnyCancellable?
     func waitForLoad() async {
         await withCheckedContinuation { continuation in
-            cancellable = finished.sink { _ in
+            cancellable = finished?.sink { _ in
                 continuation.resume()
             }
         }
diff --git a/Core/PixelEvent.swift b/Core/PixelEvent.swift
index 1a882e796c..54f78f7f49 100644
--- a/Core/PixelEvent.swift
+++ b/Core/PixelEvent.swift
@@ -413,7 +413,10 @@ extension Pixel {
         case webKitDidTerminate
         case webKitTerminationDidReloadCurrentTab
         case webKitDidTerminateDuringWarmup
-        
+
+        case webKitWarmupUnexpectedDidFinish
+        case webKitWarmupUnexpectedDidTerminate
+
         case backgroundTaskSubmissionFailed
         
         case blankOverlayNotDismissed
@@ -999,7 +1002,10 @@ extension Pixel.Event {
         case .webKitDidTerminate: return "m_d_wkt"
         case .webKitDidTerminateDuringWarmup: return "m_d_webkit-terminated-during-warmup"
         case .webKitTerminationDidReloadCurrentTab: return "m_d_wktct"
-            
+
+        case .webKitWarmupUnexpectedDidFinish: return "m_d_webkit-warmup-unexpected-did-finish"
+        case .webKitWarmupUnexpectedDidTerminate: return "m_d_webkit-warmup-unexpected-did-terminate"
+
         case .backgroundTaskSubmissionFailed: return "m_bt_rf"
             
         case .blankOverlayNotDismissed: return "m_d_ovs"