page_entities_model_executor_impl.h 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. // Copyright 2022 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #ifndef COMPONENTS_OPTIMIZATION_GUIDE_CORE_PAGE_ENTITIES_MODEL_EXECUTOR_IMPL_H_
  5. #define COMPONENTS_OPTIMIZATION_GUIDE_CORE_PAGE_ENTITIES_MODEL_EXECUTOR_IMPL_H_
  6. #include "base/callback_list.h"
  7. #include "base/memory/raw_ptr.h"
  8. #include "base/task/sequenced_task_runner.h"
  9. #include "base/task/task_traits.h"
  10. #include "base/task/thread_pool.h"
  11. #include "components/optimization_guide/core/entity_metadata.h"
  12. #include "components/optimization_guide/core/optimization_target_model_observer.h"
  13. #include "components/optimization_guide/core/page_entities_model_executor.h"
  14. namespace optimization_guide {
  15. class EntityAnnotatorNativeLibrary;
  16. class OptimizationGuideModelProvider;
  17. // A configuration that manages the necessary feature params required by the
  18. // PageEntitiesModelExecutor.
  19. struct PageEntitiesModelExecutorConfig {
  20. // Whether the page entities model should be reset on shutdown.
  21. bool should_reset_entity_annotator_on_shutdown = false;
  22. // Whether the path to the filters should be provided to the page entities
  23. // model.
  24. bool should_provide_filter_path = true;
  25. PageEntitiesModelExecutorConfig();
  26. PageEntitiesModelExecutorConfig(const PageEntitiesModelExecutorConfig& other);
  27. ~PageEntitiesModelExecutorConfig();
  28. };
  29. // Gets the current configuration.
  30. const PageEntitiesModelExecutorConfig& GetPageEntitiesModelExecutorConfig();
  31. // Overrides the config returned by |GetPageEntitiesModelExecutorConfig()|.
  32. void SetPageEntitiesModelExecutorConfigForTesting(
  33. const PageEntitiesModelExecutorConfig& config);
  34. // An object used to hold an entity annotator on a background thread.
  35. class EntityAnnotatorHolder {
  36. public:
  37. EntityAnnotatorHolder(
  38. scoped_refptr<base::SequencedTaskRunner> background_task_runner,
  39. scoped_refptr<base::SequencedTaskRunner> reply_task_runner,
  40. bool should_reset_entity_annotator_on_shutdown);
  41. ~EntityAnnotatorHolder();
  42. // Initializes the native library on a background thread. Will invoke
  43. // |init_callback| on |reply_task_runner_| with the max version supported for
  44. // the entity annotator on success. Otherwise, -1.
  45. void InitializeEntityAnnotatorNativeLibraryOnBackgroundThread(
  46. bool should_provide_filter_path,
  47. base::OnceCallback<void(int32_t)> init_callback);
  48. // Creates an entity annotator on the background thread and sets it to
  49. // |entity_annotator_|. Should be invoked on |background_task_runner_|.
  50. void CreateAndSetEntityAnnotatorOnBackgroundThread(
  51. const ModelInfo& model_info);
  52. // Requests for |entity_annotator_| to execute its model for |text| and map
  53. // the entities back to their metadata. Should be invoked on
  54. // |background_task_runner_|.
  55. using PageEntitiesMetadataModelExecutedCallback = base::OnceCallback<void(
  56. const absl::optional<std::vector<ScoredEntityMetadata>>&)>;
  57. void AnnotateEntitiesMetadataModelOnBackgroundThread(
  58. const std::string& text,
  59. PageEntitiesMetadataModelExecutedCallback callback);
  60. // Returns entity metadata from |entity_annotator_| for |entity_id|.
  61. // Should be invoked on |background_task_runner_|.
  62. void GetMetadataForEntityIdOnBackgroundThread(
  63. const std::string& entity_id,
  64. PageEntitiesModelExecutor::
  65. PageEntitiesModelEntityMetadataRetrievedCallback callback);
  66. // Gets the weak ptr to |this| on the background thread.
  67. base::WeakPtr<EntityAnnotatorHolder> GetBackgroundWeakPtr();
  68. private:
  69. void ResetEntityAnnotator();
  70. scoped_refptr<base::SequencedTaskRunner> background_task_runner_;
  71. scoped_refptr<base::SequencedTaskRunner> reply_task_runner_;
  72. // Whether the entity annotator should be reset on shutdown.
  73. const bool should_reset_entity_annotator_on_shutdown_;
  74. std::unique_ptr<EntityAnnotatorNativeLibrary>
  75. entity_annotator_native_library_;
  76. raw_ptr<void> entity_annotator_ = nullptr;
  77. base::WeakPtrFactory<EntityAnnotatorHolder> background_weak_ptr_factory_{
  78. this};
  79. };
  80. // Manages the loading and execution of the page entities model.
  81. class PageEntitiesModelExecutorImpl : public OptimizationTargetModelObserver,
  82. public PageEntitiesModelExecutor {
  83. public:
  84. PageEntitiesModelExecutorImpl(
  85. OptimizationGuideModelProvider* model_provider,
  86. scoped_refptr<base::SequencedTaskRunner> background_task_runner);
  87. ~PageEntitiesModelExecutorImpl() override;
  88. PageEntitiesModelExecutorImpl(const PageEntitiesModelExecutorImpl&) = delete;
  89. PageEntitiesModelExecutorImpl& operator=(
  90. const PageEntitiesModelExecutorImpl&) = delete;
  91. // PageEntitiesModelExecutor:
  92. void GetMetadataForEntityId(
  93. const std::string& entity_id,
  94. PageEntitiesModelEntityMetadataRetrievedCallback callback) override;
  95. void ExecuteModelWithInput(
  96. const std::string& text,
  97. PageEntitiesMetadataModelExecutedCallback callback) override;
  98. void AddOnModelUpdatedCallback(base::OnceClosure callback) override;
  99. absl::optional<ModelInfo> GetModelInfo() const override;
  100. // OptimizationTargetModelObserver:
  101. void OnModelUpdated(proto::OptimizationTarget optimization_target,
  102. const ModelInfo& model_info) override;
  103. private:
  104. // Invoked on the UI thread when entity annotator library has been
  105. // initialized.
  106. void OnEntityAnnotatorLibraryInitialized(
  107. OptimizationGuideModelProvider* model_provider,
  108. int32_t max_model_format_feature_flag);
  109. scoped_refptr<base::SequencedTaskRunner> background_task_runner_;
  110. // The holder used to hold the annotator used to annotate entities.
  111. std::unique_ptr<EntityAnnotatorHolder> entity_annotator_holder_;
  112. // The most recent model info given to |OnModelUpdated|.
  113. absl::optional<ModelInfo> model_info_;
  114. // Populated with callbacks if |AddOnModelUpdatedCallback| is called before a
  115. // model file is available, then is notified when |OnModelUpdated| is called.
  116. base::OnceClosureList on_model_updated_callbacks_;
  117. base::WeakPtrFactory<PageEntitiesModelExecutorImpl> weak_ptr_factory_{this};
  118. };
  119. } // namespace optimization_guide
  120. #endif // COMPONENTS_OPTIMIZATION_GUIDE_CORE_PAGE_ENTITIES_MODEL_EXECUTOR_IMPL_H_