Integrate prediction manager into optimization guide keyed service

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

chrome/browser/optimization_guide/optimization_guide_hints_manager.cc

@@ -693,7 +693,7 @@ void OptimizationGuideHintsManager::OnHintLoaded(
 }
 
 void OptimizationGuideHintsManager::RegisterOptimizationTypes(
-    std::vector<optimization_guide::proto::OptimizationType>
+    const std::vector<optimization_guide::proto::OptimizationType>&
         optimization_types) {
   bool should_load_new_optimization_filter = false;
   for (const auto optimization_type : optimization_types) {

chrome/browser/optimization_guide/optimization_guide_hints_manager.h

@@ -89,7 +89,7 @@ class OptimizationGuideHintsManager
   // Registers the optimization types that have the potential for hints to be
   // called by consumers of the Optimization Guide.
   void RegisterOptimizationTypes(
-      std::vector<optimization_guide::proto::OptimizationType>
+      const std::vector<optimization_guide::proto::OptimizationType>&
           optimization_types);
 
   // Returns whether there have been any optimization types registered.

chrome/browser/optimization_guide/optimization_guide_keyed_service.cc

@@ -11,6 +11,7 @@
 #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_web_contents_observer.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"
 #include "components/optimization_guide/command_line_top_host_provider.h"
@@ -143,6 +144,8 @@ void OptimizationGuideKeyedService::Initialize(
       database_provider, top_host_provider_.get(),
       content::BrowserContext::GetDefaultStoragePartition(profile)
           ->GetURLLoaderFactoryForBrowserProcess());
+  prediction_manager_ =
+      std::make_unique<optimization_guide::PredictionManager>();
 }
 
 void OptimizationGuideKeyedService::MaybeLoadHintForNavigation(
@@ -155,12 +158,15 @@ void OptimizationGuideKeyedService::MaybeLoadHintForNavigation(
   }
 }
 
-void OptimizationGuideKeyedService::RegisterOptimizationTypes(
-    std::vector<optimization_guide::proto::OptimizationType>
-        optimization_types) {
+void OptimizationGuideKeyedService::RegisterOptimizationTypesAndTargets(
+    const std::vector<optimization_guide::proto::OptimizationType>&
+        optimization_types,
+    const std::vector<optimization_guide::proto::OptimizationTarget>&
+        optimization_targets) {
   DCHECK(hints_manager_);
-
+  DCHECK(prediction_manager_);
   hints_manager_->RegisterOptimizationTypes(optimization_types);
+  prediction_manager_->RegisterOptimizationTargets(optimization_targets);
 }
 
 optimization_guide::OptimizationGuideDecision
@@ -198,10 +204,6 @@ void OptimizationGuideKeyedService::Shutdown() {
 }
 
 void OptimizationGuideKeyedService::UpdateSessionFCP(base::TimeDelta fcp) {
-  // TODO(crbug/1001194): This will be passed to the
-  // OptimizationGuidePredictionManager that will own the FCP session
-  // statistics.
-  if (!session_fcp_)
-    session_fcp_ = std::make_unique<OptimizationGuideSessionStatistic>();
-  session_fcp_->AddSample(static_cast<float>(fcp.InMilliseconds()));
+  if (prediction_manager_)
+    prediction_manager_->UpdateFCPSessionStatistics(fcp);
 }

chrome/browser/optimization_guide/optimization_guide_keyed_service.h

@@ -31,10 +31,10 @@ class ProtoDatabaseProvider;
 namespace optimization_guide {
 class OptimizationGuideService;
 class TopHostProvider;
+class PredictionManager;
 }  // namespace optimization_guide
 
 class OptimizationGuideHintsManager;
-class OptimizationGuideSessionStatistic;
 
 class OptimizationGuideKeyedService
     : public KeyedService,
@@ -60,6 +60,10 @@ class OptimizationGuideKeyedService
     return top_host_provider_.get();
   }
 
+  optimization_guide::PredictionManager* GetPredictionManager() {
+    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);
@@ -68,9 +72,11 @@ class OptimizationGuideKeyedService
   void ClearData();
 
   // optimization_guide::OptimizationGuideDecider implementation:
-  void RegisterOptimizationTypes(
-      std::vector<optimization_guide::proto::OptimizationType>
-          optimization_types) override;
+  void RegisterOptimizationTypesAndTargets(
+      const std::vector<optimization_guide::proto::OptimizationType>&
+          optimization_types,
+      const std::vector<optimization_guide::proto::OptimizationTarget>&
+          optimization_targets) override;
   optimization_guide::OptimizationGuideDecision CanApplyOptimization(
       content::NavigationHandle* navigation_handle,
       optimization_guide::proto::OptimizationTarget optimization_target,
@@ -80,21 +86,18 @@ class OptimizationGuideKeyedService
   // KeyedService implementation:
   void Shutdown() override;
 
-  // Update |session_fcp_| with the provided fcp value.
+  // Updates |prediction_manager_| with the provided fcp value.
   void UpdateSessionFCP(base::TimeDelta fcp);
 
-  OptimizationGuideSessionStatistic* GetSessionFCPForTesting() const {
-    return session_fcp_.get();
-  }
-
  private:
   content::BrowserContext* browser_context_;
 
   // Manages the storing, loading, and fetching of hints.
   std::unique_ptr<OptimizationGuideHintsManager> hints_manager_;
 
-  // The current session's FCP statistics for HTTP/HTTPS navigations.
-  std::unique_ptr<OptimizationGuideSessionStatistic> session_fcp_;
+  // Manages the storing, loading, and evaluating of optimization target
+  // prediction models.
+  std::unique_ptr<optimization_guide::PredictionManager> prediction_manager_;
 
   // The top host provider to use for fetching information for the user's top
   // hosts. Will be null if the user has not consented to this type of browser

chrome/browser/optimization_guide/optimization_guide_keyed_service_browsertest.cc

@@ -181,7 +181,8 @@ class OptimizationGuideKeyedServiceBrowserTest
 
   void RegisterWithKeyedService() {
     OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
-        ->RegisterOptimizationTypes({optimization_guide::proto::NOSCRIPT});
+        ->RegisterOptimizationTypesAndTargets(
+            {optimization_guide::proto::NOSCRIPT}, /*optimization_targets=*/{});
   }
 
   void PushHintsComponentAndWaitForCompletion() {

chrome/browser/optimization_guide/prediction/prediction_manager.h

@@ -52,6 +52,10 @@ class PredictionManager
   // Updates |session_fcp_| and |previous_fcp_| with |fcp|.
   void UpdateFCPSessionStatistics(base::TimeDelta fcp);
 
+  OptimizationGuideSessionStatistic* GetFCPSessionStatisticsForTesting() const {
+    return const_cast<OptimizationGuideSessionStatistic*>(&session_fcp_);
+  }
+
   // network::NetworkQualityTracker::EffectiveConnectionTypeObserver
   // implementation:
   void OnEffectiveConnectionTypeChanged(

chrome/browser/optimization_guide/prediction/prediction_manager_browsertest.cc

@@ -9,6 +9,7 @@
 #include "chrome/browser/optimization_guide/optimization_guide_keyed_service.h"
 #include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
 #include "chrome/browser/optimization_guide/optimization_guide_session_statistic.h"
+#include "chrome/browser/optimization_guide/prediction/prediction_manager.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/ui/browser.h"
 #include "chrome/test/base/in_process_browser_test.h"
@@ -50,7 +51,9 @@ class PredictionManagerBrowserTest : public InProcessBrowserTest {
 
   void RegisterWithKeyedService() {
     OptimizationGuideKeyedServiceFactory::GetForProfile(browser()->profile())
-        ->RegisterOptimizationTypes({optimization_guide::proto::NOSCRIPT});
+        ->RegisterOptimizationTypesAndTargets(
+            {optimization_guide::proto::NOSCRIPT},
+            {optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD});
   }
 
   std::unique_ptr<page_load_metrics::PageLoadMetricsTestWaiter>
@@ -84,8 +87,9 @@ IN_PROC_BROWSER_TEST_F(PredictionManagerBrowserTest,
   ui_test_utils::NavigateToURL(browser(), https_url_with_content());
   waiter->Wait();
 
-  OptimizationGuideSessionStatistic* session_fcp =
-      keyed_service->GetSessionFCPForTesting();
+  const OptimizationGuideSessionStatistic* session_fcp =
+      keyed_service->GetPredictionManager()
+          ->GetFCPSessionStatisticsForTesting();
   EXPECT_TRUE(session_fcp);
   EXPECT_EQ(1u, session_fcp->GetNumberOfSamples());
 }
@@ -102,8 +106,9 @@ IN_PROC_BROWSER_TEST_F(PredictionManagerBrowserTest,
   ui_test_utils::NavigateToURL(browser(), https_url_with_content());
   waiter->Wait();
 
-  OptimizationGuideSessionStatistic* session_fcp =
-      keyed_service->GetSessionFCPForTesting();
+  const OptimizationGuideSessionStatistic* session_fcp =
+      keyed_service->GetPredictionManager()
+          ->GetFCPSessionStatisticsForTesting();
   float current_mean = session_fcp->GetMean();
 
   waiter = CreatePageLoadMetricsTestWaiter();

components/optimization_guide/optimization_guide_decider.h

@@ -39,10 +39,12 @@ struct OptimizationMetadata {
 
 class OptimizationGuideDecider {
  public:
-  // Registers the optimization types that intend to be queried during the
-  // session.
-  virtual void RegisterOptimizationTypes(
-      std::vector<proto::OptimizationType> optimization_types) = 0;
+  // Registers the optimization types and targets that intend to be queried
+  // during the session. It is expected for this to be called after the browser
+  // has been initialized.
+  virtual void RegisterOptimizationTypesAndTargets(
+      const std::vector<proto::OptimizationType>& optimization_types,
+      const std::vector<proto::OptimizationTarget>& optimization_targets) = 0;
 
   // Returns whether the current conditions match |optimization_target| and
   // |optimization_type| can be applied for the URL associated with

components/previews/content/previews_optimization_guide_decider.cc

@@ -89,10 +89,11 @@ PreviewsOptimizationGuideDecider::PreviewsOptimizationGuideDecider(
       registered_optimization_types_(GetOptimizationTypesToRegister()) {
   DCHECK(optimization_guide_decider_);
 
-  optimization_guide_decider_->RegisterOptimizationTypes(
+  optimization_guide_decider_->RegisterOptimizationTypesAndTargets(
       std::vector<optimization_guide::proto::OptimizationType>(
           registered_optimization_types_.begin(),
-          registered_optimization_types_.end()));
+          registered_optimization_types_.end()),
+      {optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD});
 }
 
 PreviewsOptimizationGuideDecider::~PreviewsOptimizationGuideDecider() = default;

components/previews/content/previews_optimization_guide_decider_unittest.cc

@@ -29,12 +29,18 @@ class TestOptimizationGuideDecider
   TestOptimizationGuideDecider() = default;
   ~TestOptimizationGuideDecider() override = default;
 
-  void RegisterOptimizationTypes(
-      std::vector<optimization_guide::proto::OptimizationType>
-          optimization_types) override {
+  void RegisterOptimizationTypesAndTargets(
+      const std::vector<optimization_guide::proto::OptimizationType>&
+          optimization_types,
+      const std::vector<optimization_guide::proto::OptimizationTarget>&
+          optimization_targets) override {
     registered_optimization_types_ =
         base::flat_set<optimization_guide::proto::OptimizationType>(
             optimization_types.begin(), optimization_types.end());
+
+    registered_optimization_targets_ =
+        base::flat_set<optimization_guide::proto::OptimizationTarget>(
+            optimization_targets.begin(), optimization_targets.end());
   }
 
   // Returns the optimization types registered with the Optimization Guide
@@ -44,6 +50,13 @@ class TestOptimizationGuideDecider
     return registered_optimization_types_;
   }
 
+  // Returns the optimization targets registered with the Optimization Guide
+  // Decider.
+  base::flat_set<optimization_guide::proto::OptimizationTarget>
+  registered_optimization_targets() {
+    return registered_optimization_targets_;
+  }
+
   optimization_guide::OptimizationGuideDecision CanApplyOptimization(
       content::NavigationHandle* navigation_handle,
       optimization_guide::proto::OptimizationTarget optimization_target,
@@ -81,6 +94,11 @@ class TestOptimizationGuideDecider
   base::flat_set<optimization_guide::proto::OptimizationType>
       registered_optimization_types_;
 
+  // The optimization targets that were registered with the Optimization Guide
+  // Decider.
+  base::flat_set<optimization_guide::proto::OptimizationTarget>
+      registered_optimization_targets_;
+
   std::map<std::tuple<GURL, optimization_guide::proto::OptimizationType>,
            std::tuple<optimization_guide::OptimizationGuideDecision,
                       optimization_guide::OptimizationMetadata>>
@@ -158,7 +176,7 @@ class PreviewsOptimizationGuideDeciderTest : public testing::Test {
 };
 
 TEST_F(PreviewsOptimizationGuideDeciderTest,
-       InitializationRegistersCorrectOptimizationTypes) {
+       InitializationRegistersCorrectOptimizationTypesAndTargets) {
   base::test::ScopedFeatureList scoped_feature_list;
   scoped_feature_list.InitWithFeatures(
       {previews::features::kLitePageServerPreviews,
@@ -187,6 +205,17 @@ TEST_F(PreviewsOptimizationGuideDeciderTest,
   EXPECT_TRUE(registered_optimization_types.find(
                   optimization_guide::proto::RESOURCE_LOADING) !=
               registered_optimization_types.end());
+
+  // We expect that the PAINFUL_PAGE_LOAD optimization target is always
+  // registered.
+  base::flat_set<optimization_guide::proto::OptimizationTarget>
+      registered_optimization_targets =
+          optimization_guide_decider()->registered_optimization_targets();
+  EXPECT_EQ(1u, registered_optimization_targets.size());
+  EXPECT_TRUE(
+      registered_optimization_targets.find(
+          optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD) !=
+      registered_optimization_targets.end());
 }
 
 TEST_F(PreviewsOptimizationGuideDeciderTest,
@@ -218,6 +247,17 @@ TEST_F(PreviewsOptimizationGuideDeciderTest,
   EXPECT_EQ(registered_optimization_types.find(
                 optimization_guide::proto::RESOURCE_LOADING),
             registered_optimization_types.end());
+
+  // We expect that the PAINFUL_PAGE_LOAD optimization target is always
+  // registered.
+  base::flat_set<optimization_guide::proto::OptimizationTarget>
+      registered_optimization_targets =
+          optimization_guide_decider()->registered_optimization_targets();
+  EXPECT_EQ(1u, registered_optimization_targets.size());
+  EXPECT_TRUE(
+      registered_optimization_targets.find(
+          optimization_guide::proto::OPTIMIZATION_TARGET_PAINFUL_PAGE_LOAD) !=
+      registered_optimization_targets.end());
 }
 
 TEST_F(PreviewsOptimizationGuideDeciderTest,