chromium/components/segmentation_platform/internal/execution/model_manager_impl.cc

// Copyright 2023 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/segmentation_platform/internal/execution/model_manager_impl.h"

#include <map>
#include <memory>
#include <optional>

#include "base/functional/bind.h"
#include "base/functional/callback_helpers.h"
#include "base/time/clock.h"
#include "base/time/time.h"
#include "base/trace_event/typed_macros.h"
#include "components/segmentation_platform/internal/execution/model_manager.h"
#include "components/segmentation_platform/internal/metadata/metadata_utils.h"
#include "components/segmentation_platform/internal/proto/model_prediction.pb.h"
#include "components/segmentation_platform/internal/stats.h"
#include "components/segmentation_platform/public/model_provider.h"
#include "components/segmentation_platform/public/proto/model_metadata.pb.h"
#include "components/segmentation_platform/public/proto/segmentation_platform.pb.h"

namespace optimization_guide {
class OptimizationGuideModelProvider;
}  // namespace optimization_guide

namespace segmentation_platform {

ModelManagerImpl::ModelManagerImpl(
    const base::flat_set<SegmentId>& segment_ids,
    ModelProviderFactory* model_provider_factory,
    base::Clock* clock,
    SegmentInfoDatabase* segment_database,
    const SegmentationModelUpdatedCallback& model_updated_callback)
    : … { … }

void ModelManagerImpl::Initialize() { … }

ModelManagerImpl::~ModelManagerImpl() = default;

ModelProvider* ModelManagerImpl::GetModelProvider(
    proto::SegmentId segment_id,
    proto::ModelSource model_source) { … }

void ModelManagerImpl::SetSegmentationModelUpdatedCallbackForTesting(
    ModelManager::SegmentationModelUpdatedCallback model_updated_callback) { … }

void ModelManagerImpl::OnSegmentationModelUpdated(
    proto::ModelSource model_source,
    proto::SegmentId segment_id,
    std::optional<proto::SegmentationModelMetadata> metadata,
    int64_t model_version) { … }

void ModelManagerImpl::OnSegmentInfoFetchedForModelUpdate(
    proto::SegmentId segment_id,
    proto::ModelSource model_source,
    proto::SegmentationModelMetadata metadata,
    int64_t model_version,
    const proto::SegmentInfo* old_segment_info) { … }

void ModelManagerImpl::OnSegmentInfoDeleted(SegmentId segment_id,
                                            proto::ModelSource model_source,
                                            int64_t deleted_version,
                                            bool success) { … }

void ModelManagerImpl::OnUpdatedSegmentInfoStored(
    proto::SegmentInfo segment_info,
    std::optional<int64_t> old_model_version,
    bool success) { … }

}  // namespace segmentation_platform