Implement async CanApplyOptimization support in hints manager

Note that I have not exposed the public interface yet since I want to
make sure this is sound before migrating consumers over to this method.

Bug: 1036490
Change-Id: Id44d5f09a17287a2ba28eb7593ebd41ff082bd5b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2024164Reviewed-by: Michael Crouse <mcrouse@chromium.org>
Reviewed-by: Tarun Bansal <tbansal@chromium.org>
Commit-Queue: Sophie Chang <sophiechang@chromium.org>
Cr-Commit-Position: refs/heads/master@{#736217}

chrome/browser/optimization_guide/optimization_guide_hints_manager.cc

@@ -22,6 +22,7 @@
 #include "chrome/browser/navigation_predictor/navigation_predictor_keyed_service_factory.h"
 #include "chrome/browser/optimization_guide/optimization_guide_navigation_data.h"
 #include "chrome/browser/optimization_guide/optimization_guide_permissions_util.h"
+#include "chrome/browser/optimization_guide/optimization_guide_util.h"
 #include "chrome/browser/optimization_guide/optimization_guide_web_contents_observer.h"
 #include "chrome/browser/profiles/profile.h"
 #include "components/google/core/common/google_util.h"
@@ -107,18 +108,15 @@ bool CanProcessComponentVersion(PrefService* pref_service,
   return true;
 }
 
-// Returns the page hint for the navigation, if applicable. It will use the
-// cached page hint stored in |navigation_handle| if we have already done the
-// computation to find the page hint in a previous request to the hints manager.
-// Otherwise, we will loop through the page hints in |loaded_hint| to find the
-// one that matches and store it for subsequent calls for the navigation.
-const optimization_guide::proto::PageHint* GetPageHintForNavigation(
-    content::NavigationHandle* navigation_handle,
+// Returns the page hint, if applicable. It will use the cached page hint stored
+// in |navigation_data| if provided. Otherwise, it will loop through the page
+// hints in |loaded_hint| to find the one that matches and store it in
+// |navigation_data|, if provided, for subsequent calls using that same
+// |navigation_data|.
+const optimization_guide::proto::PageHint* GetPageHint(
+    OptimizationGuideNavigationData* navigation_data,
+    const GURL& url,
     const optimization_guide::proto::Hint* loaded_hint) {
-  OptimizationGuideNavigationData* navigation_data =
-      OptimizationGuideNavigationData::GetFromNavigationHandle(
-          navigation_handle);
-
   // If we already know we had a page hint for the navigation, then just return
   // that.
   if (navigation_data && navigation_data->has_page_hint_value()) {
@@ -127,8 +125,7 @@ const optimization_guide::proto::PageHint* GetPageHintForNavigation(
 
   // We do not yet know the answer, so find the applicable page hint.
   const optimization_guide::proto::PageHint* matched_page_hint =
-      optimization_guide::FindPageHintForURL(navigation_handle->GetURL(),
-                                             loaded_hint);
+      optimization_guide::FindPageHintForURL(url, loaded_hint);
 
   if (navigation_data) {
     // Store the page hint for the next time this is called, so we do not have
@@ -595,8 +592,10 @@ void OptimizationGuideHintsManager::OnPageNavigationHintsFetched(
     base::Optional<std::unique_ptr<optimization_guide::proto::GetHintsResponse>>
         get_hints_response) {
   if (!get_hints_response.has_value() || !get_hints_response.value()) {
-    if (navigation_url)
+    if (navigation_url) {
+      PrepareToInvokeRegisteredCallbacks(*navigation_url);
       page_navigation_hints_fetchers_.erase(*navigation_url);
+    }
     return;
   }
 
@@ -625,11 +624,13 @@ void OptimizationGuideHintsManager::OnFetchedPageNavigationHintsStored(
     const base::flat_set<std::string>& page_navigation_hosts_requested) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
+  if (navigation_url) {
+    PrepareToInvokeRegisteredCallbacks(*navigation_url);
+    page_navigation_hints_fetchers_.erase(*navigation_url);
+  }
+
   for (const auto& host : page_navigation_hosts_requested)
     LoadHintForHost(host, base::DoNothing());
-
-  if (navigation_url)
-    page_navigation_hints_fetchers_.erase(*navigation_url);
 }
 
 bool OptimizationGuideHintsManager::IsHintBeingFetchedForNavigation(
@@ -835,8 +836,12 @@ OptimizationGuideHintsManager::ShouldTargetNavigation(
     const optimization_guide::proto::Hint* loaded_hint =
         hint_cache_->GetHostKeyedHintIfLoaded(host);
     const optimization_guide::proto::PageHint* matched_page_hint =
-        loaded_hint ? GetPageHintForNavigation(navigation_handle, loaded_hint)
-                    : nullptr;
+        loaded_hint
+            ? GetPageHint(
+                  OptimizationGuideNavigationData::GetFromNavigationHandle(
+                      navigation_handle),
+                  url, loaded_hint)
+            : nullptr;
 
     if (matched_page_hint && matched_page_hint->has_max_ect_trigger()) {
       max_ect_trigger = optimization_guide::ConvertProtoEffectiveConnectionType(
@@ -860,11 +865,51 @@ OptimizationGuideHintsManager::CanApplyOptimization(
     optimization_guide::OptimizationMetadata* optimization_metadata) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
+  return CanApplyOptimization(
+      OptimizationGuideNavigationData::GetFromNavigationHandle(
+          navigation_handle),
+      navigation_handle->GetURL(), optimization_type, optimization_metadata);
+}
+
+void OptimizationGuideHintsManager::CanApplyOptimizationAsync(
+    const GURL& navigation_url,
+    optimization_guide::proto::OptimizationType optimization_type,
+    OptimizationGuideDecisionCallback callback) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  optimization_guide::OptimizationMetadata metadata;
+  optimization_guide::OptimizationTypeDecision type_decision =
+      CanApplyOptimization(/*navigation_data=*/nullptr, navigation_url,
+                           optimization_type, &metadata);
+  optimization_guide::OptimizationGuideDecision decision =
+      GetOptimizationGuideDecisionFromOptimizationTypeDecision(type_decision);
+  // It's possible that a hint that applies to |navigation_url| will come in
+  // later, so only run the callback if we are sure we can apply the decision.
+  // Also return early if the decision came from an optimization filter.
+  if (decision == optimization_guide::OptimizationGuideDecision::kTrue ||
+      HasLoadedOptimizationFilter(optimization_type)) {
+    // TODO(crbug/1036490): Add UMA for what the decision was for the
+    // optimization type.
+    std::move(callback).Run(decision, metadata);
+    return;
+  }
+
+  registered_callbacks_[navigation_url][optimization_type].push_back(
+      std::move(callback));
+}
+
+optimization_guide::OptimizationTypeDecision
+OptimizationGuideHintsManager::CanApplyOptimization(
+    OptimizationGuideNavigationData* navigation_data,
+    const GURL& url,
+    optimization_guide::proto::OptimizationType optimization_type,
+    optimization_guide::OptimizationMetadata* optimization_metadata) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
   // Clear out optimization metadata if provided.
   if (optimization_metadata)
     *optimization_metadata = {};
 
-  const auto& url = navigation_handle->GetURL();
   // If the URL doesn't have a host, we cannot query the hint for it, so just
   // return early.
   if (!url.has_host())
@@ -876,13 +921,9 @@ OptimizationGuideHintsManager::CanApplyOptimization(
       hint_cache_->GetHostKeyedHintIfLoaded(host);
   bool has_hint_in_cache = hint_cache_->HasHint(host);
   const optimization_guide::proto::PageHint* matched_page_hint =
-      loaded_hint ? GetPageHintForNavigation(navigation_handle, loaded_hint)
-                  : nullptr;
+      loaded_hint ? GetPageHint(navigation_data, url, loaded_hint) : nullptr;
 
-  // Populate navigation data with hint information.
-  OptimizationGuideNavigationData* navigation_data =
-      OptimizationGuideNavigationData::GetFromNavigationHandle(
-          navigation_handle);
+  // Populate navigation data with hint information, if provided.
   if (navigation_data) {
     navigation_data->set_has_hint_after_commit(has_hint_in_cache);
 
@@ -953,6 +994,48 @@ OptimizationGuideHintsManager::CanApplyOptimization(
              : optimization_guide::OptimizationTypeDecision::kNotAllowedByHint;
 }
 
+void OptimizationGuideHintsManager::PrepareToInvokeRegisteredCallbacks(
+    const GURL& navigation_url) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  if (registered_callbacks_.find(navigation_url) == registered_callbacks_.end())
+    return;
+
+  LoadHintForHost(
+      navigation_url.host(),
+      base::BindOnce(
+          &OptimizationGuideHintsManager::OnReadyToInvokeRegisteredCallbacks,
+          ui_weak_ptr_factory_.GetWeakPtr(), navigation_url));
+}
+
+void OptimizationGuideHintsManager::OnReadyToInvokeRegisteredCallbacks(
+    const GURL& navigation_url) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  if (registered_callbacks_.find(navigation_url) ==
+      registered_callbacks_.end()) {
+    return;
+  }
+
+  for (auto& opt_type_and_callbacks :
+       registered_callbacks_.at(navigation_url)) {
+    optimization_guide::proto::OptimizationType opt_type =
+        opt_type_and_callbacks.first;
+    optimization_guide::OptimizationMetadata metadata;
+    optimization_guide::OptimizationTypeDecision type_decision =
+        CanApplyOptimization(/*navigation_data=*/nullptr, navigation_url,
+                             opt_type, &metadata);
+    // TODO(crbug/1036490): Add UMA for what the decision was for the
+    // optimization type.
+    optimization_guide::OptimizationGuideDecision decision =
+        GetOptimizationGuideDecisionFromOptimizationTypeDecision(type_decision);
+
+    for (auto& callback : opt_type_and_callbacks.second)
+      std::move(callback).Run(decision, metadata);
+  }
+  registered_callbacks_.erase(navigation_url);
+}
+
 void OptimizationGuideHintsManager::OnEffectiveConnectionTypeChanged(
     net::EffectiveConnectionType effective_connection_type) {
   current_effective_connection_type_ = effective_connection_type;
@@ -1082,6 +1165,18 @@ void OptimizationGuideHintsManager::MaybeFetchHintsForNavigation(
   }
 }
 
+void OptimizationGuideHintsManager::OnNavigationFinish(
+    const GURL& navigation_url) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  // The callbacks will be invoked when the fetch request comes back, so it
+  // will be cleaned up later.
+  if (IsHintBeingFetchedForNavigation(navigation_url))
+    return;
+
+  PrepareToInvokeRegisteredCallbacks(navigation_url);
+}
+
 void OptimizationGuideHintsManager::ClearFetchedHints() {
   hint_cache_->ClearFetchedHints();
   optimization_guide::HintsFetcher::ClearHostsSuccessfullyFetched(

chrome/browser/optimization_guide/optimization_guide_hints_manager.h

@@ -58,9 +58,14 @@ class StoreUpdateData;
 class TopHostProvider;
 }  // namespace optimization_guide
 
+class OptimizationGuideNavigationData;
 class PrefService;
 class Profile;
 
+using OptimizationGuideDecisionCallback =
+    base::OnceCallback<void(optimization_guide::OptimizationGuideDecision,
+                            const optimization_guide::OptimizationMetadata&)>;
+
 class OptimizationGuideHintsManager
     : public optimization_guide::OptimizationGuideServiceObserver,
       public network::NetworkQualityTracker::EffectiveConnectionTypeObserver,
@@ -120,6 +125,14 @@ class OptimizationGuideHintsManager
       optimization_guide::proto::OptimizationType optimization_type,
       optimization_guide::OptimizationMetadata* optimization_metadata);
 
+  // Invokes |callback| with the decision for |navigation_url| and
+  // |optimization_type|, when sufficient information has been collected by
+  // |this| to make the decision.
+  void CanApplyOptimizationAsync(
+      const GURL& navigation_url,
+      optimization_guide::proto::OptimizationType optimization_type,
+      OptimizationGuideDecisionCallback callback);
+
   // Clears fetched hints from |hint_cache_|.
   void ClearFetchedHints();
 
@@ -148,6 +161,9 @@ class OptimizationGuideHintsManager
   void OnNavigationStartOrRedirect(content::NavigationHandle* navigation_handle,
                                    base::OnceClosure callback);
 
+  // Notifies |this| that a navigation with URL |navigation_url| has finished.
+  void OnNavigationFinish(const GURL& navigation_url);
+
  private:
   FRIEND_TEST_ALL_PREFIXES(OptimizationGuideHintsManagerTest, IsGoogleURL);
   FRIEND_TEST_ALL_PREFIXES(OptimizationGuideHintsManagerFetchingTest,
@@ -291,6 +307,23 @@ class OptimizationGuideHintsManager
   void MaybeFetchHintsForNavigation(
       content::NavigationHandle* navigation_handle);
 
+  // Returns the OptimizationTypeDecision based on the given parameters.
+  // |optimization_metadata| will be populated, if applicable. If
+  // |navigation_data| is provided, some metrics will be populated within it.
+  optimization_guide::OptimizationTypeDecision CanApplyOptimization(
+      OptimizationGuideNavigationData* navigation_data,
+      const GURL& url,
+      optimization_guide::proto::OptimizationType optimization_type,
+      optimization_guide::OptimizationMetadata* optimization_metadata);
+
+  // If an entry for |navigation_url| is contained in |registered_callbacks_|,
+  // it will load the hint for |navigation_url|'s host and upon completion, will
+  // invoke the registered callbacks for |navigation_url|.
+  void PrepareToInvokeRegisteredCallbacks(const GURL& navigation_url);
+
+  // Invokes the registered callbacks for |navigation_url|, if applicable.
+  void OnReadyToInvokeRegisteredCallbacks(const GURL& navigation_url);
+
   // The OptimizationGuideService that this guide is listening to. Not owned.
   optimization_guide::OptimizationGuideService* const
       optimization_guide_service_;
@@ -320,6 +353,12 @@ class OptimizationGuideHintsManager
                  std::unique_ptr<optimization_guide::OptimizationFilter>>
       blacklist_optimization_filters_ GUARDED_BY(optimization_filters_lock_);
 
+  // A map from URL to a map of callbacks keyed by their optimization type.
+  base::flat_map<GURL,
+                 base::flat_map<optimization_guide::proto::OptimizationType,
+                                std::vector<OptimizationGuideDecisionCallback>>>
+      registered_callbacks_;
+
   // Background thread where hints processing should be performed.
   scoped_refptr<base::SequencedTaskRunner> background_task_runner_;
 

chrome/browser/optimization_guide/optimization_guide_keyed_service.cc

@@ -10,6 +10,7 @@
 #include "chrome/browser/optimization_guide/optimization_guide_navigation_data.h"
 #include "chrome/browser/optimization_guide/optimization_guide_session_statistic.h"
 #include "chrome/browser/optimization_guide/optimization_guide_top_host_provider.h"
+#include "chrome/browser/optimization_guide/optimization_guide_util.h"
 #include "chrome/browser/optimization_guide/prediction/prediction_manager.h"
 #include "chrome/browser/profiles/profile.h"
 #include "components/leveldb_proto/public/proto_database_provider.h"
@@ -21,6 +22,7 @@
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/navigation_handle.h"
 #include "content/public/browser/storage_partition.h"
+#include "url/gurl.h"
 
 namespace {
 
@@ -91,33 +93,6 @@ GetOptimizationGuideDecisionFromOptimizationTargetDecision(
   }
 }
 
-// Returns the OptimizationGuideDecision from |optimization_type_decision|.
-optimization_guide::OptimizationGuideDecision
-GetOptimizationGuideDecisionFromOptimizationTypeDecision(
-    optimization_guide::OptimizationTypeDecision optimization_type_decision) {
-  switch (optimization_type_decision) {
-    case optimization_guide::OptimizationTypeDecision::
-        kAllowedByOptimizationFilter:
-    case optimization_guide::OptimizationTypeDecision::kAllowedByHint:
-      return optimization_guide::OptimizationGuideDecision::kTrue;
-    case optimization_guide::OptimizationTypeDecision::kUnknown:
-    case optimization_guide::OptimizationTypeDecision::
-        kHadOptimizationFilterButNotLoadedInTime:
-    case optimization_guide::OptimizationTypeDecision::
-        kHadHintButNotLoadedInTime:
-    case optimization_guide::OptimizationTypeDecision::
-        kHintFetchStartedButNotAvailableInTime:
-    case optimization_guide::OptimizationTypeDecision::kDeciderNotInitialized:
-      return optimization_guide::OptimizationGuideDecision::kUnknown;
-    case optimization_guide::OptimizationTypeDecision::kNotAllowedByHint:
-    case optimization_guide::OptimizationTypeDecision::kNoMatchingPageHint:
-    case optimization_guide::OptimizationTypeDecision::kNoHintAvailable:
-    case optimization_guide::OptimizationTypeDecision::
-        kNotAllowedByOptimizationFilter:
-      return optimization_guide::OptimizationGuideDecision::kFalse;
-  }
-}
-
 }  // namespace
 
 OptimizationGuideKeyedService::OptimizationGuideKeyedService(
@@ -157,7 +132,7 @@ void OptimizationGuideKeyedService::Initialize(
   }
 }
 
-void OptimizationGuideKeyedService::MaybeLoadHintForNavigation(
+void OptimizationGuideKeyedService::OnNavigationStartOrRedirect(
     content::NavigationHandle* navigation_handle) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
@@ -167,6 +142,14 @@ void OptimizationGuideKeyedService::MaybeLoadHintForNavigation(
   }
 }
 
+void OptimizationGuideKeyedService::OnNavigationFinish(
+    const GURL& navigation_url) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  if (hints_manager_)
+    hints_manager_->OnNavigationFinish(navigation_url);
+}
+
 void OptimizationGuideKeyedService::RegisterOptimizationTypesAndTargets(
     const std::vector<optimization_guide::proto::OptimizationType>&
         optimization_types,

chrome/browser/optimization_guide/optimization_guide_keyed_service.h

@@ -35,6 +35,7 @@ class TopHostProvider;
 class PredictionManager;
 }  // namespace optimization_guide
 
+class GURL;
 class OptimizationGuideHintsManager;
 
 class OptimizationGuideKeyedService
@@ -65,9 +66,14 @@ class OptimizationGuideKeyedService
     return prediction_manager_.get();
   }
 
-  // Prompts the load of the hint for the navigation, if there is at least one
-  // optimization type registered and there is a hint available.
-  void MaybeLoadHintForNavigation(content::NavigationHandle* navigation_handle);
+  // Notifies |hints_manager_| that the navigation associated with
+  // |navigation_handle| has started or redirected.
+  void OnNavigationStartOrRedirect(
+      content::NavigationHandle* navigation_handle);
+
+  // Notifies |hints_manager_| that the navigation associated with
+  // |navigation_url| has finished.
+  void OnNavigationFinish(const GURL& navigation_url);
 
   // Clears data specific to the user.
   void ClearData();

chrome/browser/optimization_guide/optimization_guide_util.cc

@@ -31,3 +31,29 @@ bool IsHostValidToFetchFromRemoteOptimizationGuide(const std::string& host) {
   }
   return true;
 }
+
+optimization_guide::OptimizationGuideDecision
+GetOptimizationGuideDecisionFromOptimizationTypeDecision(
+    optimization_guide::OptimizationTypeDecision optimization_type_decision) {
+  switch (optimization_type_decision) {
+    case optimization_guide::OptimizationTypeDecision::
+        kAllowedByOptimizationFilter:
+    case optimization_guide::OptimizationTypeDecision::kAllowedByHint:
+      return optimization_guide::OptimizationGuideDecision::kTrue;
+    case optimization_guide::OptimizationTypeDecision::kUnknown:
+    case optimization_guide::OptimizationTypeDecision::
+        kHadOptimizationFilterButNotLoadedInTime:
+    case optimization_guide::OptimizationTypeDecision::
+        kHadHintButNotLoadedInTime:
+    case optimization_guide::OptimizationTypeDecision::
+        kHintFetchStartedButNotAvailableInTime:
+    case optimization_guide::OptimizationTypeDecision::kDeciderNotInitialized:
+      return optimization_guide::OptimizationGuideDecision::kUnknown;
+    case optimization_guide::OptimizationTypeDecision::kNotAllowedByHint:
+    case optimization_guide::OptimizationTypeDecision::kNoMatchingPageHint:
+    case optimization_guide::OptimizationTypeDecision::kNoHintAvailable:
+    case optimization_guide::OptimizationTypeDecision::
+        kNotAllowedByOptimizationFilter:
+      return optimization_guide::OptimizationGuideDecision::kFalse;
+  }
+}

chrome/browser/optimization_guide/optimization_guide_util.h

@@ -7,6 +7,8 @@
 
 #include <string>
 
+#include "components/optimization_guide/optimization_guide_decider.h"
+#include "components/optimization_guide/optimization_guide_enums.h"
 #include "components/optimization_guide/proto/models.pb.h"
 
 // Returns the string than can be used to record histograms for the optimization
@@ -20,4 +22,9 @@ std::string GetStringNameForOptimizationTarget(
 // host that is not supported by the remote optimization guide.
 bool IsHostValidToFetchFromRemoteOptimizationGuide(const std::string& host);
 
+// Returns the OptimizationGuideDecision from |optimization_type_decision|.
+optimization_guide::OptimizationGuideDecision
+GetOptimizationGuideDecisionFromOptimizationTypeDecision(
+    optimization_guide::OptimizationTypeDecision optimization_type_decision);
+
 #endif  // CHROME_BROWSER_OPTIMIZATION_GUIDE_OPTIMIZATION_GUIDE_UTIL_H_

chrome/browser/optimization_guide/optimization_guide_web_contents_observer.cc

@@ -100,7 +100,7 @@ void OptimizationGuideWebContentsObserver::DidStartNavigation(
   if (!optimization_guide_keyed_service_)
     return;
 
-  optimization_guide_keyed_service_->MaybeLoadHintForNavigation(
+  optimization_guide_keyed_service_->OnNavigationStartOrRedirect(
       navigation_handle);
   OptimizationGuideNavigationData* nav_data =
       GetOrCreateOptimizationGuideNavigationData(navigation_handle);
@@ -124,7 +124,7 @@ void OptimizationGuideWebContentsObserver::DidRedirectNavigation(
   if (!optimization_guide_keyed_service_)
     return;
 
-  optimization_guide_keyed_service_->MaybeLoadHintForNavigation(
+  optimization_guide_keyed_service_->OnNavigationStartOrRedirect(
       navigation_handle);
   OptimizationGuideNavigationData* nav_data =
       GetOrCreateOptimizationGuideNavigationData(navigation_handle);
@@ -145,12 +145,12 @@ void OptimizationGuideWebContentsObserver::DidFinishNavigation(
   // etc.) might not have navigation data associated with them, but we reduce
   // likelihood of future leaks by always trying to remove the data.
   base::ThreadTaskRunnerHandle::Get()->PostTask(
-      FROM_HERE,
-      base::BindOnce(&OptimizationGuideWebContentsObserver::
-                         FlushMetricsAndRemoveOptimizationGuideNavigationData,
-                     weak_factory_.GetWeakPtr(),
-                     navigation_handle->GetNavigationId(),
-                     navigation_handle->HasCommitted()));
+      FROM_HERE, base::BindOnce(&OptimizationGuideWebContentsObserver::
+                                    FlushMetricsAndNotifyNavigationFinish,
+                                weak_factory_.GetWeakPtr(),
+                                navigation_handle->GetNavigationId(),
+                                navigation_handle->GetURL(),
+                                navigation_handle->HasCommitted()));
 
   if (!optimization_guide_keyed_service_)
     return;
@@ -162,13 +162,19 @@ void OptimizationGuideWebContentsObserver::DidFinishNavigation(
 }
 
 void OptimizationGuideWebContentsObserver::
-    FlushMetricsAndRemoveOptimizationGuideNavigationData(int64_t navigation_id,
-                                                         bool has_committed) {
+    FlushMetricsAndNotifyNavigationFinish(int64_t navigation_id,
+                                          const GURL& navigation_url,
+                                          bool has_committed) {
   auto nav_data_iter =
       inflight_optimization_guide_navigation_datas_.find(navigation_id);
   if (nav_data_iter == inflight_optimization_guide_navigation_datas_.end())
     return;
 
+  // If we have a navigation data for it, it's probably safe to say that this
+  // URL is the main frame URL of the navigation.
+  if (optimization_guide_keyed_service_)
+    optimization_guide_keyed_service_->OnNavigationFinish(navigation_url);
+
   (nav_data_iter->second).RecordMetrics(has_committed);
 
   inflight_optimization_guide_navigation_datas_.erase(navigation_id);

chrome/browser/optimization_guide/optimization_guide_web_contents_observer.h

@@ -59,10 +59,11 @@ class OptimizationGuideWebContentsObserver
 
   // Synchronously flushes the metrics for the navigation with ID
   // |navigation_id| and then removes any data associated with it, including its
-  // entry in |inflight_optimization_guide_navigation_datas_|.
-  void FlushMetricsAndRemoveOptimizationGuideNavigationData(
-      int64_t navigation_id,
-      bool has_committed);
+  // entry in |inflight_optimization_guide_navigation_datas_|. It also notifies
+  // |optimization_guide_keyed_service_| that the navigation has finished.
+  void FlushMetricsAndNotifyNavigationFinish(int64_t navigation_id,
+                                             const GURL& navigation_url,
+                                             bool has_committed);
 
   // The data related to a given navigation ID.
   std::map<int64_t, OptimizationGuideNavigationData>