Stages storing model features first then models for get models responses

This change updates the store in stages to ensure only one store update
in flight at a time. It prioritizes storing the model features first
and then stores the prediction models once that update is complete.

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

chrome/browser/optimization_guide/prediction/prediction_manager.cc

@@ -432,9 +432,11 @@ void PredictionManager::OnModelsAndHostFeaturesFetched(
   // from the remote Optimization Guide Service is updated.
   UpdateHostModelFeatures((*get_models_response_data)->host_model_features());
 
-  // If no models were returned, no updates to the store are required.
-  if ((*get_models_response_data)->models_size() > 0)
-    UpdatePredictionModels((*get_models_response_data)->models());
+  if ((*get_models_response_data)->models_size() > 0) {
+    // Stash the response so the models can be stored once the host
+    // model features are stored.
+    get_models_response_data_to_store_ = std::move(*get_models_response_data);
+  }
 
   // TODO(crbug/1001194): After the models and host model features are stored
   // asynchronously, the timer will be set based on the update time provided
@@ -460,6 +462,7 @@ void PredictionManager::UpdateHostModelFeatures(
           host_model_features);
     }
   }
+
   model_and_features_store_->UpdateHostModelFeatures(
       std::move(host_model_features_update_data),
       base::BindOnce(&PredictionManager::OnHostModelFeaturesStored,
@@ -504,6 +507,14 @@ void PredictionManager::OnHostModelFeaturesStored() {
   SEQUENCE_CHECKER(sequence_checker_);
   LOCAL_HISTOGRAM_BOOLEAN(
       "OptimizationGuide.PredictionManager.HostModelFeaturesStored", true);
+
+  if (get_models_response_data_to_store_ &&
+      get_models_response_data_to_store_->models_size() > 0) {
+    UpdatePredictionModels(get_models_response_data_to_store_->models());
+  }
+  // Clear any data remaining in the stored get models response.
+  get_models_response_data_to_store_.reset();
+
   // TODO(crbug/1027596): Stopping the timer can be removed once the fetch
   // callback refactor is done. Otherwise, at the start of a fetch, a timer is
   // running to handle the cases that a fetch fails but the callback is not run.
@@ -660,7 +671,11 @@ void PredictionManager::MaybeScheduleModelAndHostModelFeaturesFetch() {
   if (optimization_guide::switches::
           ShouldOverrideFetchModelsAndFeaturesTimer()) {
     SetLastModelAndFeaturesFetchAttemptTime(clock_->Now());
-    FetchModelsAndHostModelFeatures();
+    // TODO(crbug/1030358): Remove delay during tests after adding network
+    // observer to the prediction model fetcher or adding the check for offline
+    // here.
+    fetch_timer_.Start(FROM_HERE, base::TimeDelta::FromSeconds(1), this,
+                       &PredictionManager::FetchModelsAndHostModelFeatures);
   } else {
     ScheduleModelsAndHostModelFeaturesFetch();
   }

chrome/browser/optimization_guide/prediction/prediction_manager.h

@@ -295,6 +295,10 @@ class PredictionManager
   // model features from the remote Optimization Guide Service.
   std::unique_ptr<OptimizationGuideStore> model_and_features_store_;
 
+  // A stored response from a model and host model features fetch used to hold
+  // models to be stored once host model features are processed and stored.
+  std::unique_ptr<proto::GetModelsResponse> get_models_response_data_to_store_;
+
   // The URL loader factory used for fetching model and host feature updates
   // from the remote Optimization Guide Service.
   scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory_;

chrome/browser/optimization_guide/prediction/prediction_manager_unittest.cc

@@ -233,6 +233,10 @@ class TestOptimizationGuideStore : public OptimizationGuideStore {
 
   void RunInitCallback() { std::move(init_callback_).Run(); }
 
+  void RunUpdateHostModelFeaturesCallback() {
+    std::move(update_host_models_callback_).Run();
+  }
+
   void LoadPredictionModel(const EntryKey& prediction_model_entry_key,
                            PredictionModelLoadedCallback callback) override {
     model_loaded_ = true;
@@ -269,7 +273,7 @@ class TestOptimizationGuideStore : public OptimizationGuideStore {
       base::OnceClosure callback) override {
     host_model_features_update_time_ =
         *host_model_features_update_data->update_time();
-    std::move(callback).Run();
+    update_host_models_callback_ = std::move(callback);
   }
 
   void UpdatePredictionModels(
@@ -284,6 +288,7 @@ class TestOptimizationGuideStore : public OptimizationGuideStore {
   }
  private:
   base::OnceClosure init_callback_;
+  base::OnceClosure update_host_models_callback_;
   bool model_loaded_ = false;
   bool host_model_features_loaded_ = false;
 };
@@ -426,6 +431,8 @@ class PredictionManagerTest
   void SetStoreInitialized() {
     models_and_features_store()->RunInitCallback();
     RunUntilIdle();
+    // Move clock forward for any short delays added for the fetcher.
+    MoveClockForwardBy(base::TimeDelta::FromSeconds(2));
   }
 
   void MoveClockForwardBy(base::TimeDelta time_delta) {
@@ -764,6 +771,8 @@ TEST_F(PredictionManagerTest, EvaluatePredictionModelPopulatesNavData) {
 
   SetStoreInitialized();
   EXPECT_TRUE(prediction_model_fetcher()->models_fetched());
+
+  models_and_features_store()->RunUpdateHostModelFeaturesCallback();
   histogram_tester.ExpectUniqueSample(
       "OptimizationGuide.PredictionManager.HostModelFeaturesStored", true, 1);
   histogram_tester.ExpectUniqueSample(
@@ -820,6 +829,7 @@ TEST_F(PredictionManagerTest,
 
   SetStoreInitialized();
   EXPECT_TRUE(prediction_model_fetcher()->models_fetched());
+  models_and_features_store()->RunUpdateHostModelFeaturesCallback();
 
   EXPECT_EQ(OptimizationTargetDecision::kModelPredictionHoldback,
             prediction_manager()->ShouldTargetNavigation(

chrome/browser/optimization_guide/prediction/prediction_model_fetcher.cc

@@ -49,6 +49,8 @@ bool PredictionModelFetcher::FetchOptimizationGuideServiceModels(
     ModelsFetchedCallback models_fetched_callback) {
   SEQUENCE_CHECKER(sequence_checker_);
 
+  // TODO(crbug/1030358): Make the fetcher a network quality tracker observer
+  // or move this logic to the prediction manager.
   if (content::GetNetworkConnectionTracker()->IsOffline()) {
     std::move(models_fetched_callback).Run(base::nullopt);
     return false;