/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use bookmarks::Bookmarks;
use changeset_fetcher::ArcChangesetFetcher;
use cloned::cloned;
use context::CoreContext;
use futures::compat::Stream01CompatExt;
use futures::future::BoxFuture;
use futures::pin_mut;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::FutureExt;
use futures::TryFutureExt;
use futures_ext::future::spawn_controlled;
use futures_ext::future::ControlledHandle;
use futures_ext::future::FbTryFutureExt;
use futures_ext::future::TryShared;
use futures_stats::TimedFutureExt;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use parking_lot::Mutex;
use rand::Rng;
use revset::AncestorsNodeStream;
use stats::prelude::*;
use tokio::sync::Notify;
use tokio::sync::RwLock;

use crate::dag::ops::DagAddHeads;
use crate::dag::VertexListWithOptions;
use crate::idmap::IdMap;
use crate::parents::FetchParents;
use crate::read_only::ReadOnlySegmentedChangelog;
use crate::segmented_changelog_delegate;
use crate::update::server_namedag;
use crate::update::vertexlist_from_seedheads;
use crate::update::SeedHead;
use crate::update::ServerNameDag;
use crate::CloneData;
use crate::CloneHints;
use crate::InProcessIdDag;
use crate::Location;
use crate::MismatchedHeadsError;
use crate::SegmentedChangelog;

define_stats! {
    prefix = "mononoke.segmented_changelog.ondemand";
    location_to_changeset_id: timeseries(Sum),
    changeset_id_to_location: timeseries(Sum),
    missing_notification_handle: timeseries(Sum),
}

mod need_update {
    use stats::prelude::*;
    // The stats that are not per repo could be described as redundant. In situations where we are
    // debugging however it's nice to have the aggregation ready available to diagnose instances or
    // regions.
    define_stats! {
        prefix = "mononoke.segmented_changelog.need_update";
        count: timeseries(Sum),
        tries:
            histogram(1, 0, 30, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
        duration_ms:
            histogram(1000, 0, 60_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
        count_per_repo: dynamic_timeseries("{}.count", (repo_id: i32); Sum),
        tries_per_repo: dynamic_histogram(
            "{}.tries", (repo_id: i32);
            1, 0, 30, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99
        ),
        duration_ms_per_repo: dynamic_histogram(
            "{}.duration_ms", (repo_id: i32);
            1000, 0, 60_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99
        ),
    }
}

mod actual_update {
    use stats::prelude::*;
    define_stats! {
        prefix = "mononoke.segmented_changelog.update";
        count: timeseries(Sum),
        failure: timeseries(Sum),
        success: timeseries(Sum),
        duration_ms:
            histogram(1000, 0, 60_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
        count_per_repo: dynamic_timeseries("{}.count", (repo_id: i32); Sum),
        failure_per_repo: dynamic_timeseries("{}.failure", (repo_id: i32); Sum),
        success_per_repo: dynamic_timeseries("{}.success", (repo_id: i32); Sum),
        duration_ms_per_repo: dynamic_histogram(
            "{}.duration_ms", (repo_id: i32);
            1000, 0, 60_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99
        ),
    }
}

pub struct OnDemandUpdateSegmentedChangelog {
    repo_id: RepositoryId,
    namedag: Arc<RwLock<ServerNameDag>>,
    changeset_fetcher: ArcChangesetFetcher,
    bookmarks: Arc<dyn Bookmarks>,
    seed_heads: Vec<SeedHead>,
    clone_hints: Option<CloneHints>,
    ongoing_update: Arc<Mutex<Option<TryShared<BoxFuture<'static, Result<()>>>>>>,
}

impl OnDemandUpdateSegmentedChangelog {
    pub fn new(
        ctx: CoreContext,
        repo_id: RepositoryId,
        iddag: InProcessIdDag,
        idmap: Arc<dyn IdMap>,
        changeset_fetcher: ArcChangesetFetcher,
        bookmarks: Arc<dyn Bookmarks>,
        seed_heads: Vec<SeedHead>,
        clone_hints: Option<CloneHints>,
    ) -> Result<Self> {
        let namedag = server_namedag(ctx, iddag, idmap)?;
        let namedag = Arc::new(RwLock::new(namedag));
        Ok(Self {
            repo_id,
            namedag,
            changeset_fetcher,
            bookmarks,
            seed_heads,
            clone_hints,
            ongoing_update: Arc::new(Mutex::new(None)),
        })
    }

    pub fn with_periodic_update_to_master_bookmark(
        self: Arc<Self>,
        ctx: &CoreContext,
        period: Duration,
    ) -> PeriodicUpdateSegmentedChangelog {
        PeriodicUpdateSegmentedChangelog::new(ctx, self, period)
    }

    // Updating the Dag has 3 phases:
    // * loading the data that is required for the update;
    // * updating the IdMap;
    // * updating the IdDag;
    //
    // The Dag can function well for serving requests as long as the commits involved have been
    // built so we want to have easy read access to both the IdMap and the IdDag. The IdMap is a
    // very simple structure and because it's described as an Arc<dyn IdMap> we push the update
    // locking logic to the storage.  The IdDag is a complicated structure that we ask to update
    // itself. Those functions take mutable references. Updating the storage of the iddag to hide
    // the complexities of locking is more difficult. We deal with the IdDag directly by wrapping
    // it in a RwLock. The RwLock allows for easy read access which we expect to be the predominant
    // access pattern.
    //
    // Updates to the dag are not completely stable so racing updates can have conflicting results.
    // In case of conflics one of the update processes would have to restart. It's easier to reason
    // about the process if we just allow one "thread" to start an update process. The update
    // process is locked by a sync mutex. The "threads" that fail the race to update are asked to
    // wait until the ongoing update is complete. The waiters will poll on a shared future that
    // tracks the ongoing dag update. After the update is complete the waiters will go back to
    // checking if the data they have is available in the dag. It is possible that the dag is
    // updated in between determining that the an update is needed and acquiring the ongoing_update
    // lock. This is fine because the update building process checks the state of dag before the
    // dag and updates only what is necessary if necessary.
    /// Update the Dag to incorporate the commit pointed to by head.
    /// Returns true if it performed an actual update false if it simply waited for ongoing update
    /// to finish.
    async fn try_update(&self, ctx: &CoreContext, heads: &VertexListWithOptions) -> Result<bool> {
        let to_wait = {
            let mut ongoing_update = self.ongoing_update.lock();

            if let Some(fut) = &*ongoing_update {
                fut.clone().map(|_| Ok(false)).boxed()
            } else {
                cloned!(
                    ctx,
                    heads,
                    self.repo_id,
                    self.namedag,
                    self.changeset_fetcher
                );
                let task_ongoing_update = self.ongoing_update.clone();
                let update_task = async move {
                    let result =
                        the_actual_update(ctx, repo_id, namedag, changeset_fetcher, &heads).await;
                    let mut ongoing_update = task_ongoing_update.lock();
                    *ongoing_update = None;
                    result
                }
                .boxed()
                .try_shared();

                *ongoing_update = Some(update_task.clone());

                update_task.map_ok(|_| true).boxed()
            }
        };
        Ok(to_wait.await?)
    }

    async fn build_up_to_vertex_list(
        &self,
        ctx: &CoreContext,
        list: &VertexListWithOptions,
    ) -> Result<()> {
        let update_loop = async {
            let mut tries: i64 = 0;
            loop {
                tries += 1;
                if self.try_update(ctx, list).await? {
                    return Ok(tries);
                }
            }
        };
        let (stats, ret) = update_loop.timed().await;
        match ret {
            Ok(tries) if tries > 0 => {
                need_update::STATS::count.add_value(1);
                need_update::STATS::count_per_repo.add_value(1, (self.repo_id.id(),));

                need_update::STATS::tries.add_value(tries);
                need_update::STATS::tries_per_repo.add_value(tries, (self.repo_id.id(),));

                need_update::STATS::duration_ms.add_value(stats.completion_time.as_millis() as i64);
                need_update::STATS::duration_ms_per_repo.add_value(
                    stats.completion_time.as_millis() as i64,
                    (self.repo_id.id(),),
                );
                let mut scuba = ctx.scuba().clone();
                scuba.add_future_stats(&stats);
                scuba.add("repo_id", self.repo_id.id());
                scuba.add("tries", tries);
                scuba.log_with_msg("segmented_changelog_need_update", None);
            }
            _ => {}
        }
        ret.map(|_| ())
    }

    async fn build_up_to_bookmark(&self, ctx: &CoreContext) -> Result<()> {
        let vertex_list =
            vertexlist_from_seedheads(ctx, &self.seed_heads, self.bookmarks.as_ref()).await?;
        self.build_up_to_vertex_list(ctx, &vertex_list).await
    }

    async fn are_descendants_of_known_commtis(
        &self,
        ctx: &CoreContext,
        heads: &[ChangesetId],
    ) -> Result<bool> {
        let changeset_fetcher = self.changeset_fetcher.clone();
        let id_map = self.namedag.read().await.map().clone_idmap();
        let max_commits = tunables::tunables()
            .segmented_changelog_client_max_commits_to_traverse()
            .unwrap_or_default();
        for cs_id in heads {
            let ancestors =
                AncestorsNodeStream::new(ctx.clone(), &changeset_fetcher, *cs_id).compat();
            let ids = ancestors
                .take(max_commits as usize)
                .try_filter_map(|cs_id| {
                    cloned!(ctx, id_map);
                    async move { id_map.find_dag_id(&ctx, cs_id).await }
                });
            pin_mut!(ids);
            if ids.try_next().await?.is_none() {
                return Ok(false);
            }
        }
        Ok(true)
    }

    async fn build_up_to_client_heads(
        &self,
        ctx: &CoreContext,
        client_heads: &[ChangesetId],
    ) -> Result<()> {
        if self
            .are_descendants_of_known_commtis(ctx, client_heads)
            .await?
        {
            let client_heads = client_heads.iter().map(SeedHead::from);
            let mut seed_heads = self.seed_heads.clone();
            seed_heads.extend(client_heads);
            let vertex_list =
                vertexlist_from_seedheads(ctx, &seed_heads, self.bookmarks.as_ref()).await?;
            self.build_up_to_vertex_list(ctx, &vertex_list).await
        } else {
            self.build_up_to_bookmark(ctx).await
        }
    }

    async fn are_heads_assigned(&self, ctx: &CoreContext, heads: &[ChangesetId]) -> Result<bool> {
        let namedag = self.namedag.read().await;
        let idmap_wrapper = namedag.map();
        let dag_id_map = idmap_wrapper
            .clone_idmap()
            .find_many_dag_ids(ctx, heads.to_vec())
            .await?;
        if dag_id_map.len() != heads.len() {
            // Maybe heads have duplicated items? Double check.
            let mut heads = heads.to_vec();
            heads.sort_unstable();
            heads.dedup();
            if dag_id_map.len() != heads.len() {
                return Ok(false);
            }
        }
        // It is safer to check that the dag_ids we got are also in the iddag.
        for (_cs_id, dag_id) in dag_id_map {
            if !namedag.dag().contains_id(dag_id)? {
                return Ok(false);
            }
        }
        Ok(true)
    }
}

async fn the_actual_update(
    ctx: CoreContext,
    repo_id: RepositoryId,
    namedag: Arc<RwLock<ServerNameDag>>,
    changeset_fetcher: ArcChangesetFetcher,
    heads: &VertexListWithOptions,
) -> Result<()> {
    let monitored = async {
        let mut namedag = namedag.write().await;
        let parent_fetcher = FetchParents::new(ctx.clone(), changeset_fetcher);

        namedag.add_heads(&parent_fetcher, heads).await?;
        namedag.map().flush_writes().await?;
        Ok(())
    };
    actual_update::STATS::count.add_value(1);
    actual_update::STATS::count_per_repo.add_value(1, (repo_id.id(),));
    let (stats, ret) = monitored.timed().await;
    actual_update::STATS::duration_ms.add_value(stats.completion_time.as_millis() as i64);
    actual_update::STATS::duration_ms_per_repo
        .add_value(stats.completion_time.as_millis() as i64, (repo_id.id(),));
    if ret.is_ok() {
        actual_update::STATS::success.add_value(1);
        actual_update::STATS::success_per_repo.add_value(1, (repo_id.id(),));
    } else {
        actual_update::STATS::failure.add_value(1);
        actual_update::STATS::failure_per_repo.add_value(1, (repo_id.id(),));
    }
    let mut scuba = ctx.scuba().clone();
    scuba.add_future_stats(&stats);
    scuba.add("repo_id", repo_id.id());
    scuba.add("success", ret.is_ok());
    let msg = ret.as_ref().err().map(|err| format!("{:?}", err));
    scuba.log_with_msg("segmented_changelog_update", msg);
    ret
}

#[async_trait]
impl SegmentedChangelog for OnDemandUpdateSegmentedChangelog {
    async fn location_to_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        location: Location<ChangesetId>,
        count: u64,
    ) -> Result<Vec<ChangesetId>> {
        STATS::location_to_changeset_id.add_value(1);
        // Location descendant may not be the ideal entry to build up to, it could be a good idea
        // to have client_head here too.
        self.build_up_to_heads(ctx, &[location.descendant])
            .await
            .context("error while getting an up to date dag")?;
        let namedag = self.namedag.read().await;
        let read_dag = ReadOnlySegmentedChangelog::new(namedag.dag(), namedag.map().clone_idmap());
        read_dag
            .location_to_many_changeset_ids(ctx, location, count)
            .await
    }

    async fn many_changeset_ids_to_locations(
        &self,
        ctx: &CoreContext,
        master_heads: Vec<ChangesetId>,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, Result<Location<ChangesetId>>>> {
        STATS::changeset_id_to_location.add_value(1);
        self.build_up_to_heads(ctx, &master_heads)
            .await
            .context("error while getting an up to date dag")?;
        let namedag = self.namedag.read().await;
        let read_dag = ReadOnlySegmentedChangelog::new(namedag.dag(), namedag.map().clone_idmap());
        read_dag
            .many_changeset_ids_to_locations(ctx, master_heads, cs_ids)
            .await
    }

    async fn clone_data(
        &self,
        ctx: &CoreContext,
    ) -> Result<(CloneData<ChangesetId>, HashMap<ChangesetId, HgChangesetId>)> {
        let namedag = self.namedag.read().await;
        let read_dag = ReadOnlySegmentedChangelog::new(namedag.dag(), namedag.map().clone_idmap());
        let hints = if let (Some(clone_hints), Some(idmap_version)) = (
            self.clone_hints.as_ref(),
            namedag.map().as_inner().idmap_version(),
        ) {
            clone_hints.fetch(ctx, idmap_version).await?
        } else {
            HashMap::new()
        };
        read_dag.clone_data_with_hints(ctx, hints).await
    }

    async fn pull_data(
        &self,
        ctx: &CoreContext,
        common: Vec<ChangesetId>,
        missing: Vec<ChangesetId>,
    ) -> Result<CloneData<ChangesetId>> {
        let heads: Vec<_> = common.iter().chain(missing.iter()).cloned().collect();
        self.build_up_to_heads(ctx, &heads)
            .await
            .context("error while getting an up to date dag")?;
        let namedag = self.namedag.read().await;
        let read_dag = ReadOnlySegmentedChangelog::new(namedag.dag(), namedag.map().clone_idmap());
        read_dag.pull_data(ctx, common, missing).await
    }

    async fn disabled(&self, _ctx: &CoreContext) -> Result<bool> {
        Ok(false)
    }
    async fn is_ancestor(
        &self,
        ctx: &CoreContext,
        ancestor: ChangesetId,
        descendant: ChangesetId,
    ) -> Result<Option<bool>> {
        let namedag = self.namedag.read().await;
        let read_dag = ReadOnlySegmentedChangelog::new(namedag.dag(), namedag.map().clone_idmap());
        read_dag.is_ancestor(ctx, ancestor, descendant).await
    }

    async fn build_up_to_heads(&self, ctx: &CoreContext, heads: &[ChangesetId]) -> Result<bool> {
        if !self.are_heads_assigned(ctx, heads).await? {
            self.build_up_to_client_heads(ctx, heads).await?;
            // The IdDag has two groups, the MASTER group and the NON_MASTER group. The MASTER
            // group is reserved for commits that can be "lazy" client-side e.g. ancestors of
            // the master bookmark. The NON_MASTER group can contain all other changesets e.g.
            // local commits. At the moment server-side we only handle updating the MASTER
            // group. Note for the future. We should pay attention to potential races between
            // a changeset being used and the bookmark being updated.
            if !self.are_heads_assigned(ctx, heads).await? {
                let err = MismatchedHeadsError::new(self.repo_id, heads.to_vec());
                return Err(err.into());
            }
        }
        Ok(true)
    }
}

pub struct PeriodicUpdateSegmentedChangelog {
    on_demand_update_sc: Arc<OnDemandUpdateSegmentedChangelog>,
    _handle: ControlledHandle,
    #[allow(dead_code)] // useful for testing
    notify: Arc<Notify>,
}

impl PeriodicUpdateSegmentedChangelog {
    pub fn new(
        ctx: &CoreContext,
        on_demand_update_sc: Arc<OnDemandUpdateSegmentedChangelog>,
        period: Duration,
    ) -> Self {
        let notify = Arc::new(Notify::new());
        let _handle = spawn_controlled({
            let ctx = ctx.clone();
            let my_dag = Arc::clone(&on_demand_update_sc);
            let notify = Arc::clone(&notify);
            async move {
                // jitter is here so not all repos try to update at the same time
                let jitter = rand::thread_rng().gen_range(Duration::from_secs(0)..period);
                // there's lots of warmup happenning at server startup so wait at least a period
                // before starting to update
                tokio::time::sleep(period + jitter).await;

                let mut interval = tokio::time::interval(period);
                loop {
                    // Reset the scuba sequence counter for each loop iteration.
                    let ctx = ctx.with_mutated_scuba(|scuba| scuba.with_seq("seq"));
                    let _ = interval.tick().await;
                    if let Err(err) = my_dag.build_up_to_bookmark(&ctx).await {
                        slog::error!(
                            ctx.logger(),
                            "failed to update segmented changelog dag: {:?}",
                            err
                        );
                    }
                    notify.notify_waiters();
                }
            }
        });
        Self {
            on_demand_update_sc,
            _handle,
            notify,
        }
    }

    #[cfg(test)]
    pub async fn wait_for_update(&self) {
        self.notify.notified().await;
    }
}

segmented_changelog_delegate!(
    PeriodicUpdateSegmentedChangelog,
    |&self, ctx: &CoreContext| { &self.on_demand_update_sc }
);


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use mononoke_types::ChangesetId;

use crate::idmap::IdMap;
use crate::read_only::ReadOnlySegmentedChangelog;
use crate::segmented_changelog_delegate;
use crate::CloneData;
use crate::InProcessIdDag;
use crate::Location;
use crate::SegmentedChangelog;

// We call it owned because the iddag is owned.
pub struct OwnedSegmentedChangelog {
    pub(crate) iddag: InProcessIdDag,
    pub(crate) idmap: Arc<dyn IdMap>,
}

impl OwnedSegmentedChangelog {
    pub fn new(iddag: InProcessIdDag, idmap: Arc<dyn IdMap>) -> Self {
        Self { iddag, idmap }
    }
}

segmented_changelog_delegate!(OwnedSegmentedChangelog, |&self, ctx: &CoreContext| {
    ReadOnlySegmentedChangelog::new(&self.iddag, self.idmap.clone())
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use changeset_fetcher::ArcChangesetFetcher;
use context::CoreContext;

use crate::dag::errors::BackendError;
use crate::dag::namedag::MemNameDag;
use crate::dag::ops::Parents;
use crate::dag::Result;
use crate::dag::VertexName;
use crate::idmap::cs_id_from_vertex_name;
use crate::idmap::vertex_name_from_cs_id;

pub struct FetchParents {
    ctx: CoreContext,
    changeset_fetcher: ArcChangesetFetcher,
}

impl FetchParents {
    pub fn new(ctx: CoreContext, changeset_fetcher: ArcChangesetFetcher) -> Self {
        Self {
            ctx,
            changeset_fetcher,
        }
    }
}

#[async_trait::async_trait]
impl Parents for FetchParents {
    async fn parent_names(&self, name: VertexName) -> Result<Vec<VertexName>> {
        let cs_id = cs_id_from_vertex_name(&name);
        let parents = self
            .changeset_fetcher
            .get_parents(&self.ctx, cs_id)
            .await
            .map_err(BackendError::from)?;

        Ok(parents.iter().map(vertex_name_from_cs_id).collect())
    }

    async fn hint_subdag_for_insertion(&self, _heads: &[VertexName]) -> Result<MemNameDag> {
        // No dirty scope here, so always return empty
        Ok(MemNameDag::new())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use futures::future::abortable;
use futures::future::AbortHandle;
use mononoke_types::ChangesetId;
use rand::Rng;
use reloader::Loader;
use reloader::Reloader;
use slog::info;
use tokio::sync::Notify;
use tunables::tunables;

use crate::manager::SegmentedChangelogManager;
use crate::segmented_changelog_delegate;
use crate::types::SegmentedChangelogVersion;
use crate::CloneData;
use crate::Location;
use crate::SegmentedChangelog;

struct SegmentedChangelogLoader {
    manager: SegmentedChangelogManager,
    ctx: CoreContext,
    last_loaded_version: Option<SegmentedChangelogVersion>,
}

type LoadedSegmentedChangelog = Arc<dyn SegmentedChangelog + Send + Sync>;

#[async_trait]
impl Loader<LoadedSegmentedChangelog> for SegmentedChangelogLoader {
    async fn load(&mut self) -> Result<Option<LoadedSegmentedChangelog>> {
        let (sc, sc_version) = self.manager.load(&self.ctx).await?;
        self.last_loaded_version = Some(sc_version);
        Ok(Some(sc))
    }

    async fn needs_reload(&mut self) -> Result<bool> {
        Ok(self.last_loaded_version != Some(self.manager.latest_version(&self.ctx).await?))
    }
}

pub struct PeriodicReloadSegmentedChangelog {
    reloader: Reloader<LoadedSegmentedChangelog>,
    abort_handle: AbortHandle,
}

impl PeriodicReloadSegmentedChangelog {
    pub async fn start<L: Loader<LoadedSegmentedChangelog> + Send + Sync + 'static>(
        ctx: &CoreContext,
        period: Duration,
        loader: L,
        name: String,
    ) -> Result<Self> {
        let force_reload_notify = Arc::new(Notify::new());

        let ctx_clone = ctx.clone();
        let force_reload_notify_clone = force_reload_notify.clone();

        // This is a future to trigger force reload of segmented changelog
        let fut = async move {
            let mut force_reload_val = tunables().by_repo_segmented_changelog_force_reload(&name);
            loop {
                let mut jitter = tunables()
                    .segmented_changelog_force_reload_jitter_secs()
                    .unwrap_or_default();
                if jitter <= 0 {
                    jitter = 30;
                }
                let jitter = rand::thread_rng().gen_range(
                    Duration::from_secs(0)..Duration::from_secs(jitter.try_into().unwrap()),
                );
                tokio::time::sleep(jitter).await;

                let new_force_reload_val =
                    tunables().by_repo_segmented_changelog_force_reload(&name);
                if force_reload_val != new_force_reload_val {
                    info!(ctx_clone.logger(), "force reloading segmented changelog");
                    force_reload_notify_clone.notify_waiters();
                    force_reload_val = new_force_reload_val;
                }
            }
        };

        let (fut, abort_handle) = abortable(fut);
        tokio::spawn(fut);

        Ok(Self {
            reloader: Reloader::reload_periodically_with_skew_and_force_reload(
                ctx.clone(),
                period,
                loader,
                force_reload_notify,
            )
            .await?,
            abort_handle,
        })
    }

    pub async fn start_from_manager(
        ctx: &CoreContext,
        period: Duration,
        manager: SegmentedChangelogManager,
        name: String,
    ) -> Result<Self> {
        Self::start(
            ctx,
            period,
            SegmentedChangelogLoader {
                manager,
                ctx: ctx.clone(),
                last_loaded_version: None,
            },
            name,
        )
        .await
    }

    #[cfg(test)]
    pub async fn wait_for_update(&self) {
        self.reloader.wait_for_update().await
    }
}

segmented_changelog_delegate!(
    PeriodicReloadSegmentedChangelog,
    |&self, ctx: &CoreContext| { self.reloader.load() }
);

impl Drop for PeriodicReloadSegmentedChangelog {
    fn drop(&mut self) {
        self.abort_handle.abort()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::sync::Arc;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use context::CoreContext;
use futures::prelude::*;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use stats::prelude::*;

use crate::idmap::IdMap;
use crate::CloneData;
use crate::DagId;
use crate::DagIdSet;
use crate::FirstAncestorConstraint;
use crate::Group;
use crate::InProcessIdDag;
use crate::Location;
use crate::SegmentedChangelog;

const IDMAP_CHANGESET_FETCH_BATCH: usize = 500;

define_stats! {
    prefix = "mononoke.segmented_changelog.read_only";
    location_to_changeset_id: timeseries(Sum),
}

pub struct ReadOnlySegmentedChangelog<'a> {
    pub(crate) iddag: &'a InProcessIdDag,
    pub(crate) idmap: Arc<dyn IdMap>,
}

#[async_trait]
impl<'a> SegmentedChangelog for ReadOnlySegmentedChangelog<'a> {
    async fn location_to_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        location: Location<ChangesetId>,
        count: u64,
    ) -> Result<Vec<ChangesetId>> {
        STATS::location_to_changeset_id.add_value(1);
        let location = location
            .and_then_descendant(|hgid| self.idmap.get_dag_id(ctx, hgid))
            .await?;
        self.known_location_to_many_changeset_ids(ctx, location, count)
            .await
    }

    async fn many_changeset_ids_to_locations(
        &self,
        ctx: &CoreContext,
        master_heads: Vec<ChangesetId>,
        cs_ids: Vec<ChangesetId>,
    ) -> Result<HashMap<ChangesetId, Result<Location<ChangesetId>>>> {
        let (master_head_dag_ids, cs_to_dag_id) = futures::try_join!(
            self.idmap.find_many_dag_ids(ctx, master_heads.clone()),
            self.idmap.find_many_dag_ids(ctx, cs_ids),
        )
        .context("failed fetching changeset to dag_id translations")?;
        if master_head_dag_ids.is_empty() {
            // When the client has multiple heads, we are content with the server finding only one
            // of the heads. This situation comes up when master moves backwards.  The server may
            // be reseeded after that and will not have multiple heads. The client then may have
            // multiple heads and we will have to treat the heads that are not found as non master
            // heads.
            bail!(
                "failed to find idmap entries for all commits listed in \
                the master heads list: {:?}",
                master_heads
            );
        }
        let constraints = FirstAncestorConstraint::KnownUniversally {
            heads: DagIdSet::from_spans(master_head_dag_ids.into_values()),
        };
        let cs_to_vlocation: HashMap<ChangesetId, Result<Option<Location<_>>>> = cs_to_dag_id
            .into_iter()
            .map(|(cs_id, dag_id)| {
                let result = self
                    .iddag
                    .to_first_ancestor_nth(dag_id, constraints.clone());
                let cs_id_result = match result
                {
                    // Preserve error message in server response by flatten the error.
                    Err(e) => Err(format_err!(
                        "failed to compute the common descendant and distance for {} with heads {:?}: {:?}",
                        cs_id,
                        &master_heads,
                        e
                    )),
                    Ok(Some((v, dist))) => Ok(Some(Location::new(v, dist))),
                    Ok(None) => Ok(None),
                };
                (cs_id, cs_id_result)
            })
            .collect();
        let common_cs_ids = {
            let to_fetch = cs_to_vlocation
                .values()
                .filter_map(|l| match l {
                    Ok(Some(l)) => Some(l.descendant),
                    _ => None,
                })
                .collect();
            self.idmap
                .find_many_changeset_ids(ctx, to_fetch)
                .await
                .context("failed fetching dag_id to changeset translations")?
        };
        let locations: HashMap<ChangesetId, Result<Location<_>>> = cs_to_vlocation
            .into_iter()
            .filter_map(|(cs, cs_result)| {
                let cs_result = match cs_result {
                    Ok(Some(location)) => Some(location.try_map_descendant(|dag_id| {
                        common_cs_ids.get(&dag_id).cloned().ok_or_else(|| {
                            format_err!("failed to find dag_id translation for {}", dag_id)
                        })
                    })),
                    Ok(None) => None,
                    Err(e) => Some(Err(e)),
                };
                cs_result.map(|r| (cs, r))
            })
            .collect();
        Ok(locations)
    }

    async fn clone_data(
        &self,
        ctx: &CoreContext,
    ) -> Result<(CloneData<ChangesetId>, HashMap<ChangesetId, HgChangesetId>)> {
        self.clone_data_with_hints(ctx, HashMap::new()).await
    }

    async fn pull_data(
        &self,
        ctx: &CoreContext,
        common: Vec<ChangesetId>,
        missing: Vec<ChangesetId>,
    ) -> Result<CloneData<ChangesetId>> {
        let all_cs_ids: Vec<_> = common.iter().chain(missing.iter()).cloned().collect();
        let request_ids = self
            .idmap
            .find_many_dag_ids(ctx, all_cs_ids.clone())
            .await?;

        let common_ids = common
            .iter()
            .map(|i| {
                request_ids
                    .get(i)
                    .copied()
                    .ok_or_else(|| format_err!("common head {} not found", i))
            })
            .collect::<Result<Vec<_>>>()?;
        let missing_ids = missing
            .iter()
            .map(|i| {
                request_ids
                    .get(i)
                    .copied()
                    .ok_or_else(|| format_err!("missing head {} not found", i))
            })
            .collect::<Result<Vec<_>>>()?;

        let common_id_set = DagIdSet::from_spans(common_ids.into_iter());
        let missing_id_set = DagIdSet::from_spans(missing_ids.into_iter());
        let common_ancestors_id_set = self.iddag.ancestors(common_id_set)?;
        let missing_ancestors_id_set = self.iddag.ancestors(missing_id_set)?;

        let missing_id_set = missing_ancestors_id_set.difference(&common_ancestors_id_set);
        let flat_segments = self.iddag.idset_to_flat_segments(missing_id_set)?;

        let ids = flat_segments.parents_head_and_roots().into_iter().collect();

        let idmap = self
            .idmap
            .find_many_changeset_ids(ctx, ids)
            .await
            .context("error retrieving mappings for parents_head_and_roots")?;

        let pull_data = CloneData {
            flat_segments,
            idmap: idmap.into_iter().collect(),
        };
        Ok(pull_data)
    }

    /// Test if `ancestor` is an ancestor of `descendant`.
    /// Returns None in case segmented changelog doesn't know about either of those commit.
    async fn is_ancestor(
        &self,
        ctx: &CoreContext,
        ancestor: ChangesetId,
        descendant: ChangesetId,
    ) -> Result<Option<bool>> {
        let request_ids = self
            .idmap
            .find_many_dag_ids_maybe_stale(ctx, vec![ancestor, descendant])
            .await?;
        let ancestor_id = if let Some(ancestor_id) = request_ids.get(&ancestor) {
            ancestor_id
        } else {
            return Ok(None);
        };
        let descendant_id = if let Some(descendant_id) = request_ids.get(&descendant) {
            descendant_id
        } else {
            return Ok(None);
        };

        // Even though the ids exist, our local DAG might not have them.
        let all = self.iddag.all()?;
        if !all.contains(descendant_id.clone()) || !all.contains(ancestor_id.clone()) {
            return Ok(None);
        }

        Ok(Some(self.iddag.is_ancestor(*ancestor_id, *descendant_id)?))
    }

    async fn disabled(&self, _ctx: &CoreContext) -> Result<bool> {
        Ok(false)
    }
}

impl<'a> ReadOnlySegmentedChangelog<'a> {
    pub fn new(iddag: &'a InProcessIdDag, idmap: Arc<dyn IdMap>) -> Self {
        Self { iddag, idmap }
    }

    pub(crate) async fn known_location_to_many_changeset_ids(
        &self,
        ctx: &CoreContext,
        location: Location<DagId>,
        count: u64,
    ) -> Result<Vec<ChangesetId>> {
        STATS::location_to_changeset_id.add_value(1);
        let mut dist_ancestor_dag_id = self
            .iddag
            .first_ancestor_nth(location.descendant, location.distance)
            .with_context(|| format!("failed to compute location origin for {:?}", location))?;
        let mut dag_ids = vec![dist_ancestor_dag_id];
        for _ in 1..count {
            let parents = self
                .iddag
                .parent_ids(dist_ancestor_dag_id)
                .with_context(|| format!("looking up parents ids for {}", dist_ancestor_dag_id))?;
            if parents.len() != 1 {
                return Err(format_err!(
                    "invalid request: changeset with dag_id {} does not have {} single parent ancestors",
                    location.descendant,
                    location.distance + count - 1
                ));
            }
            dist_ancestor_dag_id = parents[0];
            dag_ids.push(dist_ancestor_dag_id);
        }
        let changeset_futures = dag_ids
            .into_iter()
            .map(|dag_id| self.idmap.get_changeset_id(ctx, dag_id));
        stream::iter(changeset_futures)
            .buffered(IDMAP_CHANGESET_FETCH_BATCH)
            .try_collect()
            .await
    }

    pub(crate) async fn clone_data_with_hints(
        &self,
        ctx: &CoreContext,
        mut hints: HashMap<DagId, (ChangesetId, HgChangesetId)>,
    ) -> Result<(CloneData<ChangesetId>, HashMap<ChangesetId, HgChangesetId>)> {
        let group = Group::MASTER;
        let flat_segments = self
            .iddag
            .flat_segments(group)
            .context("error during flat segment retrieval")?;
        let (idmap, hints) = {
            let universal_ids: Vec<_> = self
                .iddag
                .universal_ids()
                .context("error computing universal ids")?
                .into_iter()
                .collect();
            let mut to_fetch: Vec<_> =
                Vec::with_capacity(universal_ids.len().saturating_sub(hints.len()));

            let mut idmap = BTreeMap::new();
            let mut output_hints = HashMap::with_capacity(hints.len());

            for id in universal_ids {
                if let Some((cs, hgcs)) = hints.remove(&id) {
                    idmap.insert(id, cs);
                    output_hints.insert(cs, hgcs);
                } else {
                    to_fetch.push(id);
                }
            }

            if !to_fetch.is_empty() {
                let fetched_idmap = self
                    .idmap
                    .find_many_changeset_ids(ctx, to_fetch)
                    .await
                    .context("error retrieving mappings for dag universal ids")?;
                idmap.extend(fetched_idmap);
            }

            (idmap, output_hints)
        };
        let clone_data = CloneData {
            flat_segments,
            idmap,
        };
        Ok((clone_data, hints))
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use mononoke_types::hash::Blake2;
use sql::mysql_async::from_value_opt;
use sql::mysql_async::prelude::ConvIr;
use sql::mysql_async::prelude::FromValue;
use sql::mysql_async::FromValueError;
use sql::mysql_async::Value;

use crate::types::IdDagVersion;
use crate::types::IdMapVersion;

type FromValueResult<T> = Result<T, FromValueError>;

impl From<IdDagVersion> for Value {
    fn from(iddag_version: IdDagVersion) -> Self {
        Value::Bytes(iddag_version.0.as_ref().into())
    }
}

impl ConvIr<IdDagVersion> for Blake2 {
    fn new(v: Value) -> FromValueResult<Self> {
        match v {
            Value::Bytes(bytes) => {
                Blake2::from_bytes(&bytes).map_err(move |_| FromValueError(Value::Bytes(bytes)))
            }
            v => Err(FromValueError(v)),
        }
    }

    fn commit(self) -> IdDagVersion {
        IdDagVersion(self)
    }

    fn rollback(self) -> Value {
        Value::Bytes(self.as_ref().into())
    }
}

impl FromValue for IdDagVersion {
    type Intermediate = Blake2;
}

impl From<IdMapVersion> for Value {
    fn from(idmap_version: IdMapVersion) -> Self {
        Value::UInt(idmap_version.0)
    }
}

impl ConvIr<IdMapVersion> for IdMapVersion {
    fn new(v: Value) -> FromValueResult<Self> {
        Ok(IdMapVersion(from_value_opt(v)?))
    }

    fn commit(self) -> Self {
        self
    }

    fn rollback(self) -> Value {
        self.into()
    }
}

impl FromValue for IdMapVersion {
    type Intermediate = IdMapVersion;
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobrepo::BlobRepo;
use blobstore::Blobstore;
use blobstore_factory::MetadataSqlFactory;
use blobstore_factory::ReadOnlyStorage;
use bonsai_hg_mapping::BonsaiHgMapping;
use bonsai_hg_mapping::BonsaiHgMappingArc;
use bookmarks::Bookmarks;
use bookmarks::BookmarksArc;
use bulkops::Direction;
use bulkops::PublicChangesetBulkFetch;
use caching_ext::CacheHandlerFactory;
use changeset_fetcher::ChangesetFetcher;
use changeset_fetcher::PrefetchedChangesetsFetcher;
use changesets::ChangesetEntry;
use changesets::ChangesetsArc;
use context::CoreContext;
use futures::stream;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_stats::TimedFutureExt;
use metaconfig_types::MetadataDatabaseConfig;
use mononoke_types::Generation;
use mononoke_types::RepositoryId;
use phases::PhasesArc;
use repo_blobstore::RepoBlobstoreArc;
use repo_identity::RepoIdentityRef;
use slog::debug;
use slog::error;
use slog::info;
use sql_ext::facebook::MyAdmin;
use sql_ext::facebook::MysqlOptions;
use sql_ext::replication::NoReplicaLagMonitor;
use sql_ext::replication::ReplicaLagMonitor;
use stats::prelude::*;
use tunables::tunables;

use crate::dag::ops::DagAddHeads;
use crate::dag::DagAlgorithm;
use crate::iddag::IdDagSaveStore;
use crate::idmap::cs_id_from_vertex_name;
use crate::idmap::CacheHandlers;
use crate::idmap::IdMapFactory;
use crate::owned::OwnedSegmentedChangelog;
use crate::parents::FetchParents;
use crate::types::IdMapVersion;
use crate::types::SegmentedChangelogVersion;
use crate::update::server_namedag;
use crate::update::vertexlist_from_seedheads;
use crate::update::SeedHead;
use crate::version_store::SegmentedChangelogVersionStore;
use crate::CloneHints;
use crate::InProcessIdDag;
use crate::SegmentedChangelogSqlConnections;

define_stats! {
    prefix = "mononoke.segmented_changelog.tailer.update";
    count: timeseries(Sum),
    failure: timeseries(Sum),
    success: timeseries(Sum),
    duration_ms:
        histogram(1000, 0, 60_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99),
    count_per_repo: dynamic_timeseries("{}.count", (repo_id: i32); Sum),
    failure_per_repo: dynamic_timeseries("{}.failure", (repo_id: i32); Sum),
    success_per_repo: dynamic_timeseries("{}.success", (repo_id: i32); Sum),
    duration_ms_per_repo: dynamic_histogram(
        "{}.duration_ms", (repo_id: i32);
        1000, 0, 60_000, Average, Sum, Count; P 5; P 25; P 50; P 75; P 95; P 97; P 99
    ),
}

const DEFAULT_LOG_SAMPLING_RATE: usize = 5000;

#[derive(Clone, Copy, PartialEq)]
pub enum OperationMode {
    ContinousIncrementalUpdate(Duration),
    ForceReseed,
    SingleIncrementalUpdate,
}

pub struct SegmentedChangelogTailer {
    repo_id: RepositoryId,
    changeset_fetcher: Arc<PrefetchedChangesetsFetcher>,
    bulk_fetch: Arc<PublicChangesetBulkFetch>,
    bookmarks: Arc<dyn Bookmarks>,
    seed_heads: Vec<SeedHead>,
    sc_version_store: SegmentedChangelogVersionStore,
    iddag_save_store: IdDagSaveStore,
    idmap_factory: IdMapFactory,
    clone_hints: CloneHints,
    bonsai_hg_mapping: Arc<dyn BonsaiHgMapping>,
}

impl SegmentedChangelogTailer {
    pub fn new(
        repo_id: RepositoryId,
        connections: SegmentedChangelogSqlConnections,
        replica_lag_monitor: Arc<dyn ReplicaLagMonitor>,
        changeset_fetcher: Arc<PrefetchedChangesetsFetcher>,
        bulk_fetch: Arc<PublicChangesetBulkFetch>,
        bonsai_hg_mapping: Arc<dyn BonsaiHgMapping>,
        blobstore: Arc<dyn Blobstore>,
        bookmarks: Arc<dyn Bookmarks>,
        seed_heads: Vec<SeedHead>,
        cache_handler_factory: Option<CacheHandlerFactory>,
    ) -> Self {
        let clone_hints = CloneHints::new(connections.0.clone(), repo_id, blobstore.clone());
        let sc_version_store = SegmentedChangelogVersionStore::new(connections.0.clone(), repo_id);
        let iddag_save_store = IdDagSaveStore::new(repo_id, blobstore);
        let mut idmap_factory = IdMapFactory::new(connections.0, replica_lag_monitor, repo_id);
        if let Some(cache_handler_factory) = cache_handler_factory {
            let cache_handlers = CacheHandlers::new(cache_handler_factory);
            idmap_factory = idmap_factory.with_cache_handlers(cache_handlers);
        }
        Self {
            repo_id,
            changeset_fetcher,
            bulk_fetch,
            bookmarks,
            seed_heads,
            sc_version_store,
            iddag_save_store,
            idmap_factory,
            clone_hints,
            bonsai_hg_mapping,
        }
    }

    pub async fn build_from(
        ctx: &CoreContext,
        blobrepo: &BlobRepo,
        storage_config_metadata: &MetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
        seed_heads: Vec<SeedHead>,
        prefetched_commits: impl Stream<Item = Result<ChangesetEntry, Error>>,
        cache_handler_factory: Option<CacheHandlerFactory>,
    ) -> Result<Self> {
        let repo_id = blobrepo.repo_identity().id();

        let db_address = match storage_config_metadata {
            MetadataDatabaseConfig::Local(_) => None,
            MetadataDatabaseConfig::Remote(remote_config) => {
                Some(remote_config.primary.db_address.clone())
            }
        };
        let replica_lag_monitor: Arc<dyn ReplicaLagMonitor> = match db_address {
            None => Arc::new(NoReplicaLagMonitor()),
            Some(address) => {
                let my_admin = MyAdmin::new(ctx.fb).context("building myadmin client")?;
                Arc::new(my_admin.single_shard_lag_monitor(address))
            }
        };

        let sql_factory = MetadataSqlFactory::new(
            ctx.fb,
            storage_config_metadata.clone(),
            mysql_options.clone(),
            ReadOnlyStorage(false),
        )
        .await
        .with_context(|| format!("constructing metadata sql factory for repo {}", repo_id))?;

        let segmented_changelog_sql_connections = sql_factory
            .open::<SegmentedChangelogSqlConnections>()
            .await
            .with_context(|| {
                format!(
                    "error constructing segmented changelog sql connections for repo {}",
                    repo_id
                )
            })?;

        let changeset_fetcher = Arc::new(
            PrefetchedChangesetsFetcher::new(
                repo_id,
                blobrepo.changesets_arc(),
                prefetched_commits,
            )
            .await?,
        );

        let bulk_fetcher = Arc::new(PublicChangesetBulkFetch::new(
            blobrepo.changesets_arc(),
            blobrepo.phases_arc(),
        ));

        let bonsai_hg_mapping = blobrepo.bonsai_hg_mapping_arc();

        Ok(SegmentedChangelogTailer::new(
            repo_id,
            segmented_changelog_sql_connections,
            replica_lag_monitor,
            changeset_fetcher,
            bulk_fetcher,
            bonsai_hg_mapping,
            blobrepo.repo_blobstore_arc(),
            blobrepo.bookmarks_arc(),
            seed_heads,
            cache_handler_factory,
        ))
    }

    pub async fn run(&self, ctx: &CoreContext, mode: OperationMode) {
        STATS::success.add_value(0);
        STATS::success_per_repo.add_value(0, (self.repo_id.id(),));

        let mut interval = if let OperationMode::ContinousIncrementalUpdate(period) = mode {
            Some(tokio::time::interval(period))
        } else {
            None
        };

        loop {
            if let Some(interval) = &mut interval {
                let _ = interval.tick().await;
            }
            debug!(ctx.logger(), "woke up to update");

            STATS::count.add_value(1);
            STATS::count_per_repo.add_value(1, (self.repo_id.id(),));

            let (stats, update_result) = self
                .once(ctx, mode == OperationMode::ForceReseed)
                .timed()
                .await;

            STATS::duration_ms.add_value(stats.completion_time.as_millis() as i64);
            STATS::duration_ms_per_repo.add_value(
                stats.completion_time.as_millis() as i64,
                (self.repo_id.id(),),
            );

            let mut scuba = ctx.scuba().clone();
            scuba.add_future_stats(&stats);
            scuba.add("repo_id", self.repo_id.id());
            scuba.add("success", update_result.is_ok());

            let msg = match update_result {
                Ok(_) => {
                    STATS::success.add_value(1);
                    STATS::success_per_repo.add_value(1, (self.repo_id.id(),));
                    None
                }
                Err(err) => {
                    STATS::failure.add_value(1);
                    STATS::failure_per_repo.add_value(1, (self.repo_id.id(),));
                    error!(
                        ctx.logger(),
                        "failed to incrementally update segmented changelog: {:?}", err
                    );
                    Some(format!("{:?}", err))
                }
            };
            scuba.log_with_msg("segmented_changelog_tailer_update", msg);
            if mode == OperationMode::SingleIncrementalUpdate || mode == OperationMode::ForceReseed
            {
                break;
            }
        }
    }

    pub async fn once(
        &self,
        ctx: &CoreContext,
        force_reseed: bool,
    ) -> Result<OwnedSegmentedChangelog> {
        info!(
            ctx.logger(),
            "starting incremental update to segmented changelog",
        );

        let (seeding, idmap_version, iddag) = {
            let sc_version = self.sc_version_store.get(ctx).await.with_context(|| {
                format!(
                    "error loading segmented changelog version for repo {}",
                    self.repo_id
                )
            })?;

            match sc_version {
                Some(sc_version) => {
                    if force_reseed {
                        (
                            true,
                            sc_version.idmap_version.bump(),
                            InProcessIdDag::new_in_process(),
                        )
                    } else {
                        let iddag = self
                            .iddag_save_store
                            .load(ctx, sc_version.iddag_version)
                            .await
                            .with_context(|| {
                                format!("failed to load iddag for repo {}", self.repo_id)
                            })?;
                        (false, sc_version.idmap_version, iddag)
                    }
                }
                None => (true, IdMapVersion(1), InProcessIdDag::new_in_process()),
            }
        };

        if let Ok(set) = iddag.all() {
            info!(
                ctx.logger(),
                "iddag initialized, it covers {} ids",
                set.count(),
            );
        }
        let idmap = self.idmap_factory.for_writer(ctx, idmap_version);

        let mut namedag = server_namedag(ctx.clone(), iddag, idmap)?;

        let heads =
            vertexlist_from_seedheads(ctx, &self.seed_heads, self.bookmarks.as_ref()).await?;

        let head_commits: Vec<_> = namedag
            .heads(namedag.master_group().await?)
            .await?
            .iter()
            .await?
            .map_ok(|name| cs_id_from_vertex_name(&name))
            .try_collect()
            .await?;

        let changeset_fetcher = {
            let namedag_max_gen = stream::iter(head_commits.iter().map(Ok::<_, Error>))
                .try_fold(0, {
                    let fetcher = &self.changeset_fetcher;
                    move |max, cs_id| async move {
                        let gen = fetcher.get_generation_number(ctx, *cs_id).await?;
                        Ok(max.max(gen.value()))
                    }
                })
                .await?;
            let heads_min_gen = stream::iter(
                heads
                    .vertexes()
                    .iter()
                    .map(|name| Ok::<_, Error>(cs_id_from_vertex_name(name))),
            )
            .try_fold(Generation::max_gen().value(), {
                let fetcher = &self.changeset_fetcher;
                move |min, cs_id| async move {
                    let gen = fetcher.get_generation_number(ctx, cs_id).await?;
                    Ok(min.min(gen.value()))
                }
            })
            .await?;

            if heads_min_gen.saturating_sub(namedag_max_gen) > 1000 {
                let repo_bounds = self
                    .bulk_fetch
                    .get_repo_bounds_after_commits(ctx, head_commits)
                    .await?;
                info!(ctx.logger(), "prefetching changeset entries",);
                let mut counter = 0usize;
                // This has the potential to cause OOM by fetching a large
                // chunk of the repo
                let missing = self
                    .bulk_fetch
                    .fetch_bounded(ctx, Direction::NewestFirst, Some(repo_bounds))
                    .map(|res| {
                        counter += 1;
                        let sampling_rate = tunables()
                            .segmented_changelog_tailer_log_sampling_rate()
                            .unwrap_or_default();
                        let sampling_rate = if sampling_rate <= 0 {
                            DEFAULT_LOG_SAMPLING_RATE
                        } else {
                            sampling_rate as usize
                        };
                        if counter % sampling_rate == 0 {
                            info!(
                                ctx.logger(),
                                "fetched {} changeset entries in total", counter,
                            );
                        }
                        res
                    });
                Arc::new(self.changeset_fetcher.clone_with_extension(missing).await?)
            } else {
                self.changeset_fetcher.clone()
            }
        };

        let parent_fetcher = FetchParents::new(ctx.clone(), changeset_fetcher);

        info!(ctx.logger(), "starting the actual update");
        // Note on memory use: we do not flush the changes out in the middle
        // of writing to the IdMap.
        // Thus, if OOMs happen here, the IdMap may need to flush writes to the DB
        // at interesting points.
        let changed = namedag.add_heads(&parent_fetcher, &heads).await?;

        self.clone_hints
            .add_hints(
                ctx,
                &namedag,
                idmap_version,
                self.bonsai_hg_mapping.as_ref(),
            )
            .await?;

        let (idmap, iddag) = namedag.into_idmap_dag();
        let idmap = idmap.finish().await?;

        if !changed {
            info!(
                ctx.logger(),
                "segmented changelog already up to date, skipping update to iddag",
            );
            let owned = OwnedSegmentedChangelog::new(iddag, idmap);
            return Ok(owned);
        }

        info!(ctx.logger(), "IdMap updated, IdDag updated",);

        // Save the IdDag
        let iddag_version = self
            .iddag_save_store
            .save(ctx, &iddag)
            .await
            .with_context(|| format!("error saving iddag for repo {}", self.repo_id))?;

        // Update SegmentedChangelogVersion
        let sc_version = SegmentedChangelogVersion::new(iddag_version, idmap_version);
        if seeding {
            self.sc_version_store
                .set(ctx, sc_version)
                .await
                .with_context(|| {
                    format!(
                        "error updating segmented changelog version store for repo {}",
                        self.repo_id
                    )
                })?;
            info!(ctx.logger(), "successfully seeded segmented changelog",);
        } else {
            self.sc_version_store
                .update(ctx, sc_version)
                .await
                .with_context(|| {
                    format!(
                        "error updating segmented changelog version store for repo {}",
                        self.repo_id
                    )
                })?;
            info!(
                ctx.logger(),
                "successful incremental update to segmented changelog",
            );
        }

        let owned = OwnedSegmentedChangelog::new(iddag, idmap);
        Ok(owned)
    }
}