From 9f42d2c782d6ebeefa7fc152f0ca27709dedb767 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=A9=AC=E5=AE=97=E5=AF=8C?= Date: Sat, 30 Nov 2024 16:39:26 +0800 Subject: [PATCH] =?UTF-8?q?114trunk=E4=BF=AE=E5=A4=8DCVE-40066823?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 马宗富 --- .../core/prediction_manager_unittest.cc | 32 ++++++++--- .../core/prediction_model_store.cc | 54 ++++++------------- .../core/prediction_model_store.h | 28 +++------- .../core/prediction_model_store_unittest.cc | 32 +++++++---- 4 files changed, 71 insertions(+), 75 deletions(-) diff --git a/components/optimization_guide/core/prediction_manager_unittest.cc b/components/optimization_guide/core/prediction_manager_unittest.cc index 994a2cc94c..d5f1dc616f 100644 --- a/components/optimization_guide/core/prediction_manager_unittest.cc +++ b/components/optimization_guide/core/prediction_manager_unittest.cc @@ -388,6 +388,18 @@ class TestPredictionManager : public PredictionManager { OptimizationGuideLogger optimization_guide_logger_; }; +class TestPredictionModelStore : public PredictionModelStore { + public: + explicit TestPredictionModelStore(PrefService* local_state) + : local_state_(local_state) {} + + // PredictionModelStore: + PrefService* GetLocalState() const override { return local_state_; } + + private: + raw_ptr local_state_; +}; + class PredictionManagerTestBase : public ProtoDatabaseProviderTestBase { public: using StoreEntry = proto::StoreEntry; @@ -432,6 +444,12 @@ class PredictionManagerTestBase : public ProtoDatabaseProviderTestBase { prediction_manager_->SetClockForTesting(task_environment_.GetMockClock()); } + void CreateAndInitializePredictionModelStore() { + prediction_model_store_ = + std::make_unique(local_state_prefs_.get()); + prediction_model_store_->Initialize(temp_dir()); + } + std::unique_ptr CreateModelAndHostModelFeaturesStore() { // Setup the fake db and the class under test. @@ -514,6 +532,10 @@ class PredictionManagerTestBase : public ProtoDatabaseProviderTestBase { base::test::TaskEnvironment* task_environment() { return &task_environment_; } + PredictionModelStore* prediction_model_store() { + return prediction_model_store_.get(); + } + void SetComponentUpdatesPrefEnabled(bool enabled) { component_updates_enabled_ = enabled; } @@ -525,13 +547,13 @@ class PredictionManagerTestBase : public ProtoDatabaseProviderTestBase { // tsan flakes caused by other tasks running while |feature_list_| is // destroyed. base::test::ScopedFeatureList feature_list_; - std::unique_ptr prediction_model_store_; private: base::test::TaskEnvironment task_environment_{ base::test::TaskEnvironment::MainThreadType::UI, base::test::TaskEnvironment::TimeSource::MOCK_TIME}; StoreEntryMap db_store_; + std::unique_ptr prediction_model_store_; std::unique_ptr model_and_features_store_; std::unique_ptr prediction_manager_; scoped_refptr url_loader_factory_; @@ -617,9 +639,7 @@ class PredictionManagerTest : public testing::WithParamInterface, void SetUp() override { PredictionManagerTestBase::SetUp(); if (ShouldEnableInstallWideModelStore()) { - prediction_model_store_ = - PredictionModelStore::CreatePredictionModelStoreForTesting( - local_state_prefs_.get(), temp_dir()); + CreateAndInitializePredictionModelStore(); } } @@ -634,10 +654,6 @@ class PredictionManagerTest : public testing::WithParamInterface, RunUntilIdle(); } - PredictionModelStore* prediction_model_store() { - return prediction_model_store_.get(); - } - bool ShouldEnableInstallWideModelStore() const { return GetParam(); } private: diff --git a/components/optimization_guide/core/prediction_model_store.cc b/components/optimization_guide/core/prediction_model_store.cc index e9fadce71f..539816a7e7 100644 --- a/components/optimization_guide/core/prediction_model_store.cc +++ b/components/optimization_guide/core/prediction_model_store.cc @@ -105,12 +105,6 @@ void RecordModelStorageMetrics(const base::FilePath& base_store_dir) { } // namespace -// static -PredictionModelStore* PredictionModelStore::GetInstance() { - static base::NoDestructor model_store; - return model_store.get(); -} - PredictionModelStore::PredictionModelStore() : background_task_runner_(base::ThreadPool::CreateSequencedTaskRunner( {base::MayBlock(), base::TaskPriority::BEST_EFFORT})) { @@ -119,19 +113,14 @@ PredictionModelStore::PredictionModelStore() PredictionModelStore::~PredictionModelStore() = default; -void PredictionModelStore::Initialize(PrefService* local_state, - const base::FilePath& base_store_dir) { +void PredictionModelStore::Initialize(const base::FilePath& base_store_dir) { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); - DCHECK(local_state); DCHECK(!base_store_dir.empty()); // Should not be initialized already. - DCHECK(!local_state_); DCHECK(base_store_dir_.empty()); - local_state_ = local_state; base_store_dir_ = base_store_dir; - PurgeInactiveModels(); // Clean up any model files that were slated for deletion in previous @@ -142,22 +131,12 @@ void PredictionModelStore::Initialize(PrefService* local_state, FROM_HERE, base::BindOnce(&RecordModelStorageMetrics, base_store_dir_)); } -// static -std::unique_ptr -PredictionModelStore::CreatePredictionModelStoreForTesting( - PrefService* local_state, - const base::FilePath& base_store_dir) { - auto store = base::WrapUnique(new PredictionModelStore()); - store->Initialize(local_state, base_store_dir); - return store; -} - bool PredictionModelStore::HasModel( proto::OptimizationTarget optimization_target, const proto::ModelCacheKey& model_cache_key) const { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); auto metadata = ModelStoreMetadataEntry::GetModelMetadataEntryIfExists( - local_state_, optimization_target, model_cache_key); + GetLocalState(), optimization_target, model_cache_key); if (!metadata) { return false; } @@ -171,7 +150,7 @@ bool PredictionModelStore::HasModelWithVersion( int64_t version) { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); auto metadata = ModelStoreMetadataEntry::GetModelMetadataEntryIfExists( - local_state_, optimization_target, model_cache_key); + GetLocalState(), optimization_target, model_cache_key); if (!metadata) { return false; } @@ -191,7 +170,7 @@ void PredictionModelStore::LoadModel( DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); auto metadata = ModelStoreMetadataEntry::GetModelMetadataEntryIfExists( - local_state_, optimization_target, model_cache_key); + GetLocalState(), optimization_target, model_cache_key); if (!metadata) { std::move(callback).Run(nullptr); return; @@ -275,7 +254,7 @@ void PredictionModelStore::UpdateMetadataForExistingModel( if (!HasModel(optimization_target, model_cache_key)) return; - ModelStoreMetadataEntryUpdater metadata(local_state_, optimization_target, + ModelStoreMetadataEntryUpdater metadata(GetLocalState(), optimization_target, model_cache_key); auto base_model_dir = metadata.GetModelBaseDir(); DCHECK(base_store_dir_.IsParent(*base_model_dir)); @@ -299,7 +278,7 @@ void PredictionModelStore::UpdateModel( DCHECK_EQ(optimization_target, model_info.optimization_target()); DCHECK(base_store_dir_.IsParent(base_model_dir)); - ModelStoreMetadataEntryUpdater metadata(local_state_, optimization_target, + ModelStoreMetadataEntryUpdater metadata(GetLocalState(), optimization_target, model_cache_key); metadata.SetVersion(model_info.version()); metadata.SetExpiryTime( @@ -353,7 +332,7 @@ void PredictionModelStore::UpdateModelCacheKeyMapping( const proto::ModelCacheKey& server_model_cache_key) { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); ModelStoreMetadataEntryUpdater::UpdateModelCacheKeyMapping( - local_state_, optimization_target, client_model_cache_key, + GetLocalState(), optimization_target, client_model_cache_key, server_model_cache_key); } @@ -362,18 +341,18 @@ void PredictionModelStore::RemoveModel( const proto::ModelCacheKey& model_cache_key, PredictionModelStoreModelRemovalReason model_remove_reason) { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); - if (!local_state_) { + if (!GetLocalState()) { return; } RecordPredictionModelStoreModelRemovalVersionHistogram(model_remove_reason); - ModelStoreMetadataEntryUpdater metadata(local_state_, optimization_target, + ModelStoreMetadataEntryUpdater metadata(GetLocalState(), optimization_target, model_cache_key); auto base_model_dir = metadata.GetModelBaseDir(); if (base_model_dir) { DCHECK(base_store_dir_.IsParent(*base_model_dir)); ScopedDictPrefUpdate pref_update( - local_state_, prefs::localstate::kStoreFilePathsToDelete); + GetLocalState(), prefs::localstate::kStoreFilePathsToDelete); pref_update->Set(FilePathToString(*base_model_dir), true); } // Continue removing the metadata even if the model dirs does not exist. @@ -382,9 +361,10 @@ void PredictionModelStore::RemoveModel( void PredictionModelStore::PurgeInactiveModels() { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); - DCHECK(local_state_); + DCHECK(GetLocalState()); for (const auto& expired_model_dir : - ModelStoreMetadataEntryUpdater::PurgeAllInactiveMetadata(local_state_)) { + ModelStoreMetadataEntryUpdater::PurgeAllInactiveMetadata( + GetLocalState())) { DCHECK(base_store_dir_.IsParent(expired_model_dir)); // This is called at startup. So no need to schedule the deletion of the // model dirs, and instead can be deleted immediately. @@ -395,9 +375,9 @@ void PredictionModelStore::PurgeInactiveModels() { void PredictionModelStore::CleanUpOldModelFiles() { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); - DCHECK(local_state_); + DCHECK(GetLocalState()); for (const auto entry : - local_state_->GetDict(prefs::localstate::kStoreFilePathsToDelete)) { + GetLocalState()->GetDict(prefs::localstate::kStoreFilePathsToDelete)) { auto path_to_delete = StringToFilePath(entry.first); DCHECK(path_to_delete); DCHECK(base_store_dir_.IsParent(*path_to_delete)); @@ -412,13 +392,13 @@ void PredictionModelStore::CleanUpOldModelFiles() { void PredictionModelStore::OnFilePathDeleted(const std::string& path_to_delete, bool success) { DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); - DCHECK(local_state_); + DCHECK(GetLocalState()); if (!success) { // Try to delete again later. return; } - ScopedDictPrefUpdate pref_update(local_state_, + ScopedDictPrefUpdate pref_update(GetLocalState(), prefs::localstate::kStoreFilePathsToDelete); pref_update->Remove(path_to_delete); } diff --git a/components/optimization_guide/core/prediction_model_store.h b/components/optimization_guide/core/prediction_model_store.h index 72ed1cc1eb..4b45788918 100644 --- a/components/optimization_guide/core/prediction_model_store.h +++ b/components/optimization_guide/core/prediction_model_store.h @@ -8,7 +8,6 @@ #include "base/files/file_path.h" #include "base/memory/scoped_refptr.h" #include "base/memory/weak_ptr.h" -#include "base/no_destructor.h" #include "base/sequence_checker.h" #include "base/task/sequenced_task_runner.h" #include "base/values.h" @@ -34,21 +33,15 @@ class PredictionModelStore { using PredictionModelLoadedCallback = base::OnceCallback)>; - // Returns the singleton model store. - static PredictionModelStore* GetInstance(); - - static std::unique_ptr - CreatePredictionModelStoreForTesting(PrefService* local_state, - const base::FilePath& base_store_dir); + PredictionModelStore(); - // Initializes the model store with |local_state| and the |base_store_dir|. - // Model store will be usable only after it is initialized. - void Initialize(PrefService* local_state, - const base::FilePath& base_store_dir); + // Initializes the model store with |base_store_dir|. Model store will be + // usable only after it is initialized. + void Initialize(const base::FilePath& base_store_dir); PredictionModelStore(const PredictionModelStore&) = delete; PredictionModelStore& operator=(const PredictionModelStore&) = delete; - ~PredictionModelStore(); + virtual ~PredictionModelStore(); // Initializes the model store with |local_state| and the |base_store_dir|, if // initialization hasn't happened already. Model store will be usable only @@ -102,10 +95,10 @@ class PredictionModelStore { const proto::ModelCacheKey& client_model_cache_key, const proto::ModelCacheKey& server_model_cache_key); - private: - friend base::NoDestructor; + // Returns the local state that stores the prefs across all profiles. + virtual PrefService* GetLocalState() const = 0; - PredictionModelStore(); + private: // Loads the model and verifies if the model files exist and returns the // model. Otherwise nullptr is returned on any failures. @@ -148,11 +141,6 @@ class PredictionModelStore { // Invoked when model files gets deleted. void OnFilePathDeleted(const std::string& path_to_delete, bool success); - // Local state that stores the prefs across all profiles. Not owned and - // outlives |this|. - raw_ptr local_state_ GUARDED_BY_CONTEXT(sequence_checker_) = - nullptr; - // The base dir where the prediction model dirs are saved. base::FilePath base_store_dir_ GUARDED_BY_CONTEXT(sequence_checker_); diff --git a/components/optimization_guide/core/prediction_model_store_unittest.cc b/components/optimization_guide/core/prediction_model_store_unittest.cc index 26cfb7698f..9143b087a9 100644 --- a/components/optimization_guide/core/prediction_model_store_unittest.cc +++ b/components/optimization_guide/core/prediction_model_store_unittest.cc @@ -48,6 +48,18 @@ struct ModelDetail { } // namespace +class TestPredictionModelStore : public PredictionModelStore { + public: + explicit TestPredictionModelStore(PrefService* local_state) + : local_state_(local_state) {} + + // PredictionModelStore: + PrefService* GetLocalState() const override { return local_state_; } + + private: + raw_ptr local_state_; +}; + class PredictionModelStoreTest : public testing::Test { public: PredictionModelStoreTest() { @@ -64,9 +76,7 @@ class PredictionModelStoreTest : public testing::Test { ASSERT_TRUE(temp_models_dir_.CreateUniqueTempDir()); local_state_prefs_ = std::make_unique(); prefs::RegisterLocalStatePrefs(local_state_prefs_->registry()); - prediction_model_store_ = - PredictionModelStore::CreatePredictionModelStoreForTesting( - local_state_prefs_.get(), temp_models_dir_.GetPath()); + CreateAndInitializePredictionModelStore(); } void OnPredictionModelLoaded( @@ -109,6 +119,12 @@ class PredictionModelStoreTest : public testing::Test { return {model_info, base_model_dir}; } + void CreateAndInitializePredictionModelStore() { + prediction_model_store_ = + std::make_unique(local_state_prefs_.get()); + prediction_model_store_->Initialize(temp_models_dir_.GetPath()); + } + void WaitForModeLoad(proto::OptimizationTarget optimization_target, const proto::ModelCacheKey& model_cache_key) { std::unique_ptr run_loop = std::make_unique(); @@ -126,7 +142,7 @@ class PredictionModelStoreTest : public testing::Test { base::ScopedTempDir temp_models_dir_; std::unique_ptr local_state_prefs_; std::unique_ptr last_loaded_prediction_model_; - std::unique_ptr prediction_model_store_; + std::unique_ptr prediction_model_store_; }; TEST_F(PredictionModelStoreTest, BaseModelDirs) { @@ -303,9 +319,7 @@ TEST_F(PredictionModelStoreTest, ModelStorageMetrics) { RunUntilIdle(); // Recreate the model store, and that should record model storage metrics. - prediction_model_store_ = - PredictionModelStore::CreatePredictionModelStoreForTesting( - local_state_prefs_.get(), temp_models_dir_.GetPath()); + CreateAndInitializePredictionModelStore(); RunUntilIdle(); histogram_tester.ExpectUniqueSample( "OptimizationGuide.PredictionModelStore.ModelCount.PainfulPageLoad", 1, @@ -348,9 +362,7 @@ TEST_F(PredictionModelStoreTest, ExpiredModelRemoved) { base::Seconds(1)); // Recreate the store and it will remove the expired model. - prediction_model_store_ = - PredictionModelStore::CreatePredictionModelStoreForTesting( - local_state_prefs_.get(), temp_models_dir_.GetPath()); + CreateAndInitializePredictionModelStore(); RunUntilIdle(); EXPECT_FALSE(prediction_model_store_->HasModel(kTestOptimizationTargetFoo, model_cache_key)); -- Gitee