Place models and host model features from a fetch in the store.

This change persists the model and host model features received from a
fetch to the remote Optimization Guide Service by places them within
the store owned by the PredictionManager.

A future change will add browser tests and metrics to confirm expected
behavior of prediction manager, fetcher, and models.

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

chrome/browser/optimization_guide/prediction/prediction_manager.cc

@@ -29,6 +29,7 @@
 #include "components/optimization_guide/optimization_guide_prefs.h"
 #include "components/optimization_guide/optimization_guide_store.h"
 #include "components/optimization_guide/optimization_guide_switches.h"
+#include "components/optimization_guide/store_update_data.h"
 #include "components/optimization_guide/top_host_provider.h"
 #include "components/prefs/pref_service.h"
 #include "content/public/browser/navigation_handle.h"
@@ -62,9 +63,6 @@ constexpr base::TimeDelta kFetchRetryDelay = base::TimeDelta::FromMinutes(16);
 // The amount of time to wait after a successful fetch of models and host model
 // features before requesting an update from the remote Optimization Guide
 // Service.
-//
-// TODO(mcrouse): Remove this when the models and features are placed into the
-// store and rely on the store to provide the update time.
 constexpr base::TimeDelta kUpdateModelsAndFeaturesDelay =
     base::TimeDelta::FromHours(24);
 
@@ -421,31 +419,15 @@ void PredictionManager::OnModelsAndHostFeaturesFetched(
   if (!get_models_response_data)
     return;
 
-  // TODO(crbug/1001194): Asynchronously store the models and host model
-  // features within the persistent store.
-
-  // The set of host model features that the models in
-  // |get_models_response_data| required in order to be evaluated. Every host
-  // model feature returned should contain all the features for the models
-  // currently supported.
-  base::flat_set<std::string> host_model_features;
-  if ((*get_models_response_data)->host_model_features_size() > 0) {
-    UpdateHostModelFeatures((*get_models_response_data)->host_model_features());
-
-    host_model_features.reserve((*get_models_response_data)
-                                    ->host_model_features(0)
-                                    .model_features_size());
-    for (const auto& model_feature :
-         (*get_models_response_data)->host_model_features(0).model_features()) {
-      if (model_feature.has_feature_name())
-        host_model_features.insert(model_feature.feature_name());
-    }
-  }
+  // Update host model features, even if empty so the store metadata
+  // that contains the update time for new models and features to be fetched
+  // 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) {
-    UpdatePredictionModels((*get_models_response_data)->mutable_models(),
-                           host_model_features);
-  }
   // TODO(crbug/1001194): After the models and host model features are stored
   // asynchronously, the timer will be set based on the update time provided
   // by the store.
@@ -459,8 +441,21 @@ void PredictionManager::UpdateHostModelFeatures(
     const google::protobuf::RepeatedPtrField<proto::HostModelFeatures>&
         host_model_features) {
   SEQUENCE_CHECKER(sequence_checker_);
-  for (const auto& host_model_features : host_model_features)
-    ProcessAndStoreHostModelFeatures(host_model_features);
+  std::unique_ptr<StoreUpdateData> host_model_features_update_data =
+      StoreUpdateData::CreateHostModelFeaturesStoreUpdateData(
+          /*update_time=*/clock_->Now() + kUpdateModelsAndFeaturesDelay,
+          /*expiry_time=*/clock_->Now() +
+              features::StoredHostModelFeaturesFreshnessDuration());
+  for (const auto& host_model_features : host_model_features) {
+    if (ProcessAndStoreHostModelFeatures(host_model_features)) {
+      host_model_features_update_data->CopyHostModelFeaturesIntoUpdateData(
+          host_model_features);
+    }
+  }
+  model_and_features_store_->UpdateHostModelFeatures(
+      std::move(host_model_features_update_data),
+      base::BindOnce(&PredictionManager::OnHostModelFeaturesStored,
+                     ui_weak_ptr_factory_.GetWeakPtr()));
 }
 
 std::unique_ptr<PredictionModel> PredictionManager::CreatePredictionModel(
@@ -471,13 +466,41 @@ std::unique_ptr<PredictionModel> PredictionManager::CreatePredictionModel(
 }
 
 void PredictionManager::UpdatePredictionModels(
-    google::protobuf::RepeatedPtrField<proto::PredictionModel>*
-        prediction_models,
-    const base::flat_set<std::string>& host_model_features) {
+    const google::protobuf::RepeatedPtrField<proto::PredictionModel>&
+        prediction_models) {
+  SEQUENCE_CHECKER(sequence_checker_);
+  std::unique_ptr<StoreUpdateData> prediction_model_update_data =
+      StoreUpdateData::CreatePredictionModelStoreUpdateData();
+  bool models_to_store = false;
+  for (const auto& model : prediction_models) {
+    if (ProcessAndStorePredictionModel(model)) {
+      prediction_model_update_data->CopyPredictionModelIntoUpdateData(model);
+      models_to_store = true;
+    }
+  }
+  if (models_to_store) {
+    model_and_features_store_->UpdatePredictionModels(
+        std::move(prediction_model_update_data),
+        base::BindOnce(&PredictionManager::OnPredictionModelsStored,
+                       ui_weak_ptr_factory_.GetWeakPtr()));
+  }
+}
+
+void PredictionManager::OnPredictionModelsStored() {
   SEQUENCE_CHECKER(sequence_checker_);
-  std::unique_ptr<PredictionModel> prediction_model;
-  for (auto& model : *prediction_models)
-    ProcessAndStorePredictionModel(model);
+  LOCAL_HISTOGRAM_BOOLEAN(
+      "OptimizationGuide.PredictionManager.PredictionModelsStored", true);
+}
+
+void PredictionManager::OnHostModelFeaturesStored() {
+  SEQUENCE_CHECKER(sequence_checker_);
+  LOCAL_HISTOGRAM_BOOLEAN(
+      "OptimizationGuide.PredictionManager.HostModelFeaturesStored", true);
+  // 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.
+  fetch_timer_.Stop();
+  ScheduleModelsAndHostModelFeaturesFetch();
 }
 
 void PredictionManager::OnStoreInitialized() {
@@ -557,42 +580,42 @@ void PredictionManager::OnLoadPredictionModel(
     ProcessAndStorePredictionModel(*model);
 }
 
-void PredictionManager::ProcessAndStorePredictionModel(
+bool PredictionManager::ProcessAndStorePredictionModel(
     const proto::PredictionModel& model) {
   SEQUENCE_CHECKER(sequence_checker_);
   if (!model.model_info().has_optimization_target())
-    return;
+    return false;
   if (!registered_optimization_targets_.contains(
           model.model_info().optimization_target())) {
-    return;
+    return false;
   }
 
   std::unique_ptr<PredictionModel> prediction_model =
       CreatePredictionModel(model);
   if (!prediction_model)
-    return;
+    return false;
 
   auto it = optimization_target_prediction_model_map_.find(
       model.model_info().optimization_target());
   if (it == optimization_target_prediction_model_map_.end()) {
     optimization_target_prediction_model_map_.emplace(
         model.model_info().optimization_target(), std::move(prediction_model));
-    return;
+    return true;
   }
   if (it->second->GetVersion() != prediction_model->GetVersion()) {
     it->second = std::move(prediction_model);
-    return;
+    return true;
   }
-  return;
+  return false;
 }
 
-void PredictionManager::ProcessAndStoreHostModelFeatures(
+bool PredictionManager::ProcessAndStoreHostModelFeatures(
     const proto::HostModelFeatures& host_model_features) {
   SEQUENCE_CHECKER(sequence_checker_);
   if (!host_model_features.has_host())
-    return;
+    return false;
   if (host_model_features.model_features_size() == 0)
-    return;
+    return false;
 
   base::flat_map<std::string, float> model_features_for_host;
   model_features_for_host.reserve(host_model_features.model_features_size());
@@ -618,11 +641,11 @@ void PredictionManager::ProcessAndStoreHostModelFeatures(
     }
   }
   if (model_features_for_host.size() == 0)
-    return;
+    return false;
 
   host_model_features_map_[host_model_features.host()] =
       model_features_for_host;
-  return;
+  return true;
 }
 
 void PredictionManager::MaybeScheduleModelAndHostModelFeaturesFetch() {

chrome/browser/optimization_guide/prediction/prediction_manager.h

@@ -140,17 +140,19 @@ class PredictionManager
   virtual std::unique_ptr<PredictionModel> CreatePredictionModel(
       const proto::PredictionModel& model) const;
 
-  // Process |host_model_features| to be stored in |host_model_features_map|.
+  // Process |host_model_features| to be stored in memory in the host model
+  // features map for immediate use and asynchronously write them to the model
+  // and features store to be persisted.
   void UpdateHostModelFeatures(
       const google::protobuf::RepeatedPtrField<proto::HostModelFeatures>&
           host_model_features);
 
-  // Process |prediction_models| to be stored in
-  // |optimization_target_prediction_model_map_|.
+  // Process |prediction_models| to be stored in the in memory optimization
+  // target prediction model map for immediate use and asynchronously write the
+  // models to the model and features store to be persisted.
   void UpdatePredictionModels(
-      google::protobuf::RepeatedPtrField<proto::PredictionModel>*
-          prediction_models,
-      const base::flat_set<std::string>& host_model_features);
+      const google::protobuf::RepeatedPtrField<proto::PredictionModel>&
+          prediction_models);
 
  private:
   // Called on construction to register optimization targets, initialize the
@@ -181,7 +183,9 @@ class PredictionManager
   // Callback when the models and host model features have been fetched from the
   // remote Optimization Guide Service and are ready for parsing. Processes the
   // prediction models and the host model features in the response and stores
-  // them for use.
+  // them for use. The metadata entry containing the time that updates should be
+  // fetched from the remote Optimization Guide Service is updated, even when
+  // the response is empty.
   void OnModelsAndHostFeaturesFetched(
       base::Optional<std::unique_ptr<proto::GetModelsResponse>>
           get_models_response_data);
@@ -192,6 +196,16 @@ class PredictionManager
   // true.
   void OnStoreInitialized();
 
+  // Callback run after prediction models are stored in
+  // |model_and_features_store_|.
+  void OnPredictionModelsStored();
+
+  // Callback run after host model features are stored in
+  // |model_and_features_store_|. |fetch_timer_| is stopped and the timer is
+  // rescheduled based on when models and host model features should be fetched
+  // again.
+  void OnHostModelFeaturesStored();
+
   // Request the store to load all the host model features it contains. This
   // must be completed before any prediction models can be loaded from the
   // store.
@@ -218,13 +232,15 @@ class PredictionManager
       std::unique_ptr<proto::PredictionModel> prediction_model);
 
   // Process |model| into a PredictionModel object and store it in the
-  // |optimization_target_prediction_model_map_|.
-  void ProcessAndStorePredictionModel(const proto::PredictionModel& model);
+  // |optimization_target_prediction_model_map_|. Return true if a prediction
+  // model object was created and successfully stored, otherwise false.
+  bool ProcessAndStorePredictionModel(const proto::PredictionModel& model);
 
   // Process |host_model_features| from the into host model features
   // usable by the PredictionManager. The processed host model features are
-  // stored in |host_model_features_map_|.
-  void ProcessAndStoreHostModelFeatures(
+  // stored in |host_model_features_map_|. Return true if host model features
+  // can be constructed and successfully stored, otherwise, return false.
+  bool ProcessAndStoreHostModelFeatures(
       const proto::HostModelFeatures& host_model_features);
 
   // Return the time when a prediction model and host model features fetch was

components/optimization_guide/optimization_guide_store.h

@@ -178,8 +178,9 @@ class OptimizationGuideStore {
   std::unique_ptr<StoreUpdateData> CreateUpdateDataForPredictionModels() const;
 
   // Updates the prediction models contained in the store. The callback is run
-  // asynchronously after the database stores the prediction models.
-  void UpdatePredictionModels(
+  // asynchronously after the database stores the prediction models. Virtualized
+  // for testing.
+  virtual void UpdatePredictionModels(
       std::unique_ptr<StoreUpdateData> prediction_models_update_data,
       base::OnceClosure callback);
 
@@ -212,7 +213,8 @@ class OptimizationGuideStore {
 
   // Updates the host model features contained in the store. The callback is run
   // asynchronously after the database stores the host model features.
-  void UpdateHostModelFeatures(
+  // Virtualized for testing.
+  virtual void UpdateHostModelFeatures(
       std::unique_ptr<StoreUpdateData> host_model_features_update_data,
       base::OnceClosure callback);
 
@@ -248,6 +250,7 @@ class OptimizationGuideStore {
  private:
   friend class OptimizationGuideStoreTest;
   friend class StoreUpdateData;
+  friend class TestOptimizationGuideStore;
 
   using EntryKeyPrefix = std::string;
   using EntryKeySet = base::flat_set<EntryKey>;

components/optimization_guide/optimization_guide_store_unittest.cc

@@ -171,8 +171,7 @@ class OptimizationGuideStoreTest : public testing::Test {
         prediction_model = CreatePredictionModel();
     prediction_model->mutable_model_info()->set_optimization_target(
         optimization_target);
-    update_data->MovePredictionModelIntoUpdateData(
-        std::move(*prediction_model));
+    update_data->CopyPredictionModelIntoUpdateData(*prediction_model);
   }
 
   // Moves |host_model_features_count| into |update_data|.
@@ -186,8 +185,7 @@ class OptimizationGuideStoreTest : public testing::Test {
       model_feature->set_feature_name("host_feat1");
       model_feature->set_double_value(2.0);
       host_model_features.set_host(host_suffix);
-      update_data->MoveHostModelFeaturesIntoUpdateData(
-          std::move(host_model_features));
+      update_data->CopyHostModelFeaturesIntoUpdateData(host_model_features);
     }
   }
 

components/optimization_guide/store_update_data.cc

@@ -157,8 +157,8 @@ void StoreUpdateData::MoveHintIntoUpdateData(proto::Hint&& hint) {
                                  std::move(entry_proto));
 }
 
-void StoreUpdateData::MoveHostModelFeaturesIntoUpdateData(
-    proto::HostModelFeatures&& host_model_features) {
+void StoreUpdateData::CopyHostModelFeaturesIntoUpdateData(
+    const proto::HostModelFeatures& host_model_features) {
   // All future modifications must be made by the same thread. Note, |this| may
   // have been constructed on another thread.
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
@@ -174,14 +174,13 @@ void StoreUpdateData::MoveHostModelFeaturesIntoUpdateData(
       OptimizationGuideStore::StoreEntryType::kHostModelFeatures));
   entry_proto.set_expiry_time_secs(
       expiry_time_->ToDeltaSinceWindowsEpoch().InSeconds());
-  entry_proto.set_allocated_host_model_features(
-      new proto::HostModelFeatures(std::move(host_model_features)));
+  entry_proto.mutable_host_model_features()->CopyFrom(host_model_features);
   entries_to_save_->emplace_back(std::move(host_model_features_entry_key),
                                  std::move(entry_proto));
 }
 
-void StoreUpdateData::MovePredictionModelIntoUpdateData(
-    proto::PredictionModel&& prediction_model) {
+void StoreUpdateData::CopyPredictionModelIntoUpdateData(
+    const proto::PredictionModel& prediction_model) {
   // All future modifications must be made by the same thread. Note, |this| may
   // have been constructed on another thread.
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
@@ -196,8 +195,7 @@ void StoreUpdateData::MovePredictionModelIntoUpdateData(
   proto::StoreEntry entry_proto;
   entry_proto.set_entry_type(static_cast<proto::StoreEntryType>(
       OptimizationGuideStore::StoreEntryType::kPredictionModel));
-  entry_proto.set_allocated_prediction_model(
-      new proto::PredictionModel(std::move(prediction_model)));
+  entry_proto.mutable_prediction_model()->CopyFrom(prediction_model);
   entries_to_save_->emplace_back(std::move(prediction_model_entry_key),
                                  std::move(entry_proto));
 }

components/optimization_guide/store_update_data.h

@@ -65,15 +65,13 @@ class StoreUpdateData {
   // called, |hint| is no longer valid.
   void MoveHintIntoUpdateData(proto::Hint&& hint);
 
-  // Moves |host_model_features| in this update data. Afterwards,
-  // |host_model_features| is no longer valid.
-  void MoveHostModelFeaturesIntoUpdateData(
-      proto::HostModelFeatures&& host_model_features);
-
-  // Moves |prediction_model| in this update data. Afterwards,
-  // |prediction_model| is no longer valid.
-  void MovePredictionModelIntoUpdateData(
-      proto::PredictionModel&& prediction_model);
+  // Copies |host_model_features| into this update data.
+  void CopyHostModelFeaturesIntoUpdateData(
+      const proto::HostModelFeatures& host_model_features);
+
+  // Copies |prediction_model| into this update data.
+  void CopyPredictionModelIntoUpdateData(
+      const proto::PredictionModel& prediction_model);
 
   // Returns the store entry updates along with ownership to them.
   std::unique_ptr<EntryVector> TakeUpdateEntries();

components/optimization_guide/store_update_data_unittest.cc

@@ -82,8 +82,7 @@ TEST(StoreUpdateDataTest, BuildPredictionModelUpdateData) {
 
   std::unique_ptr<StoreUpdateData> prediction_model_update =
       StoreUpdateData::CreatePredictionModelStoreUpdateData();
-  prediction_model_update->MovePredictionModelIntoUpdateData(
-      std::move(prediction_model));
+  prediction_model_update->CopyPredictionModelIntoUpdateData(prediction_model);
   EXPECT_FALSE(prediction_model_update->component_version().has_value());
   EXPECT_FALSE(prediction_model_update->update_time().has_value());
   // Verify there is 1 store entry.
@@ -106,7 +105,7 @@ TEST(StoreUpdateDataTest, BuildHostModelFeaturesUpdateData) {
           host_model_features_update_time +
               optimization_guide::features::
                   StoredHostModelFeaturesFreshnessDuration());
-  host_model_features_update->MoveHostModelFeaturesIntoUpdateData(
+  host_model_features_update->CopyHostModelFeaturesIntoUpdateData(
       std::move(host_model_features));
   EXPECT_FALSE(host_model_features_update->component_version().has_value());
   EXPECT_TRUE(host_model_features_update->update_time().has_value());