prediction_manager.cc | Explore in Territory

// Copyright 2019 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/optimization_guide/core/prediction_manager.h"

#include <memory>
#include <utility>
#include <vector>

#include "base/containers/contains.h"
#include "base/containers/fixed_flat_set.h"
#include "base/containers/flat_map.h"
#include "base/containers/flat_set.h"
#include "base/containers/flat_tree.h"
#include "base/functional/callback.h"
#include "base/metrics/histogram_functions.h"
#include "base/metrics/histogram_macros.h"
#include "base/metrics/histogram_macros_local.h"
#include "base/not_fatal_until.h"
#include "base/observer_list.h"
#include "base/path_service.h"
#include "base/sequence_checker.h"
#include "base/task/sequenced_task_runner.h"
#include "base/task/single_thread_task_runner.h"
#include "base/task/thread_pool.h"
#include "base/time/time.h"
#include "base/uuid.h"
#include "components/optimization_guide/core/model_info.h"
#include "components/optimization_guide/core/model_util.h"
#include "components/optimization_guide/core/optimization_guide_constants.h"
#include "components/optimization_guide/core/optimization_guide_enums.h"
#include "components/optimization_guide/core/optimization_guide_features.h"
#include "components/optimization_guide/core/optimization_guide_logger.h"
#include "components/optimization_guide/core/optimization_guide_permissions_util.h"
#include "components/optimization_guide/core/optimization_guide_prefs.h"
#include "components/optimization_guide/core/optimization_guide_store.h"
#include "components/optimization_guide/core/optimization_guide_switches.h"
#include "components/optimization_guide/core/optimization_guide_util.h"
#include "components/optimization_guide/core/optimization_target_model_observer.h"
#include "components/optimization_guide/core/prediction_model_download_manager.h"
#include "components/optimization_guide/core/prediction_model_fetcher_impl.h"
#include "components/optimization_guide/core/prediction_model_override.h"
#include "components/optimization_guide/core/prediction_model_store.h"
#include "components/optimization_guide/core/store_update_data.h"
#include "components/optimization_guide/optimization_guide_internals/webui/optimization_guide_internals.mojom.h"
#include "components/optimization_guide/proto/models.pb.h"
#include "components/prefs/pref_service.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "services/network/public/cpp/shared_url_loader_factory.h"

namespace optimization_guide {

namespace {

proto::ModelCacheKey GetModelCacheKey(const std::string& locale) { … }

// Util class for recording the construction and validation of a prediction
// model. The result is recorded when it goes out of scope and its destructor is
// called.
class ScopedPredictionModelConstructionAndValidationRecorder { … };

void RecordModelUpdateVersion(const proto::ModelInfo& model_info) { … }

void RecordLifecycleState(proto::OptimizationTarget optimization_target,
                          ModelDeliveryEvent event) { … }

// Returns whether models should be fetched from the
// remote Optimization Guide Service.
bool ShouldFetchModels(bool off_the_record, bool component_updates_enabled) { … }

// Returns whether the model metadata proto is on the server allowlist.
bool IsModelMetadataTypeOnServerAllowlist(const proto::Any& model_metadata) { … }

void RecordModelAvailableAtRegistration(
    proto::OptimizationTarget optimization_target,
    bool model_available_at_registration) { … }

}  // namespace

PredictionManager::ModelRegistrationInfo::ModelRegistrationInfo(
    std::optional<proto::Any> metadata)
    : … { … }

PredictionManager::ModelRegistrationInfo::~ModelRegistrationInfo() = default;

PredictionManager::PredictionManager(
    PredictionModelStore* prediction_model_store,
    scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
    PrefService* pref_service,
    bool off_the_record,
    const std::string& application_locale,
    const base::FilePath& models_dir_path,
    OptimizationGuideLogger* optimization_guide_logger,
    BackgroundDownloadServiceProvider background_download_service_provider,
    ComponentUpdatesEnabledProvider component_updates_enabled_provider)
    : … { … }

PredictionManager::~PredictionManager() { … }

void PredictionManager::Initialize(
    BackgroundDownloadServiceProvider background_download_service_provider) { … }

void PredictionManager::AddObserverForOptimizationTargetModel(
    proto::OptimizationTarget optimization_target,
    const std::optional<proto::Any>& model_metadata,
    OptimizationTargetModelObserver* observer) { … }

void PredictionManager::RemoveObserverForOptimizationTargetModel(
    proto::OptimizationTarget optimization_target,
    OptimizationTargetModelObserver* observer) { … }

base::flat_set<proto::OptimizationTarget>
PredictionManager::GetRegisteredOptimizationTargets() const { … }

void PredictionManager::SetPredictionModelFetcherForTesting(
    std::unique_ptr<PredictionModelFetcher> prediction_model_fetcher) { … }

void PredictionManager::SetPredictionModelDownloadManagerForTesting(
    std::unique_ptr<PredictionModelDownloadManager>
        prediction_model_download_manager) { … }

void PredictionManager::FetchModels() { … }

void PredictionManager::OnModelsFetched(
    const std::vector<proto::ModelInfo> models_request_info,
    std::optional<std::unique_ptr<proto::GetModelsResponse>>
        get_models_response_data) { … }

void PredictionManager::UpdateModelMetadata(
    const proto::PredictionModel& model) { … }

bool PredictionManager::ShouldDownloadNewModel(
    const proto::PredictionModel& model) const { … }

void PredictionManager::StartModelDownload(
    proto::OptimizationTarget optimization_target,
    const GURL& download_url) { … }

void PredictionManager::MaybeDownloadOrUpdatePredictionModel(
    proto::OptimizationTarget optimization_target,
    const proto::PredictionModel& get_models_response_model,
    std::unique_ptr<proto::PredictionModel> loaded_model) { … }

void PredictionManager::UpdatePredictionModels(
    const std::vector<proto::ModelInfo>& models_request_info,
    const google::protobuf::RepeatedPtrField<proto::PredictionModel>&
        prediction_models) { … }

void PredictionManager::OnModelReady(const base::FilePath& base_model_dir,
                                     const proto::PredictionModel& model) { … }

void PredictionManager::OnModelDownloadStarted(
    proto::OptimizationTarget optimization_target) { … }

void PredictionManager::OnModelDownloadFailed(
    proto::OptimizationTarget optimization_target) { … }

std::vector<optimization_guide_internals::mojom::DownloadedModelInfoPtr>
PredictionManager::GetDownloadedModelsInfoForWebUI() const { … }

void PredictionManager::NotifyObserversOfNewModel(
    proto::OptimizationTarget optimization_target,
    base::optional_ref<const ModelInfo> model_info) { … }

void PredictionManager::OnPredictionModelsStored() { … }

void PredictionManager::MaybeInitializeModelDownloads(
    download::BackgroundDownloadService* background_download_service) { … }

void PredictionManager::OnPredictionModelOverrideLoaded(
    proto::OptimizationTarget optimization_target,
    std::unique_ptr<proto::PredictionModel> prediction_model) { … }

void PredictionManager::LoadPredictionModels(
    const base::flat_set<proto::OptimizationTarget>& optimization_targets) { … }

void PredictionManager::OnLoadPredictionModel(
    proto::OptimizationTarget optimization_target,
    bool record_availability_metrics,
    std::unique_ptr<proto::PredictionModel> model) { … }

void PredictionManager::OnProcessLoadedModel(
    const proto::PredictionModel& model,
    bool success) { … }

void PredictionManager::RemoveModelFromStore(
    proto::OptimizationTarget optimization_target,
    PredictionModelStoreModelRemovalReason model_removal_reason) { … }

bool PredictionManager::ProcessAndStoreLoadedModel(
    const proto::PredictionModel& model) { … }

bool PredictionManager::ShouldUpdateStoredModelForTarget(
    proto::OptimizationTarget optimization_target,
    int64_t new_version) const { … }

void PredictionManager::StoreLoadedModelInfo(
    proto::OptimizationTarget optimization_target,
    std::unique_ptr<ModelInfo> model_info) { … }

base::FilePath PredictionManager::GetBaseModelDirForDownload(
    proto::OptimizationTarget optimization_target) { … }

void PredictionManager::OverrideTargetModelForTesting(
    proto::OptimizationTarget optimization_target,
    std::unique_ptr<ModelInfo> model_info) { … }

}  // namespace optimization_guide
chromium/components/optimization_guide/core/prediction_manager.cc
