sui_rpc_store/indexer/

restore.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Entry point for bulk-loading the [`RpcStoreSchema`]'s
//! derived-index CFs from a [`RestoreSource`].
//!
//! Registers the three live-object-derivable index pipelines
//! ([`ObjectByOwner`], [`ObjectByType`], [`Balance`]) plus the
//! [`ObjectVersionByCheckpoint`] and [`PackageVersions`] floor rows —
//! and, when the caller's [`RestoreLayer`] opts in, the raw
//! [`Objects`] CF — against a single [`RestoreDriver`] and returns a
//! [`Service`] driving the restore through to completion. Once
//! finished, every registered
//! pipeline's `__restore` row is `Complete` and its `__watermark`
//! row is set to the source's target, so the regular
//! [`Indexer::add_pipelines`] path will accept them for tip
//! indexing.
//!
//! Restoration is run separately from tip indexing — open the
//! database, call [`restore_indexes`] to populate the indexes,
//! then construct an [`Indexer`] over the same store to start
//! tip-following.
//!
//! [`Indexer`]: crate::Indexer
//! [`Indexer::add_pipelines`]: crate::Indexer::add_pipelines

use std::sync::Arc;

use anyhow::Context as _;
use sui_consistent_store::Batch;
use sui_consistent_store::ChainId;
use sui_consistent_store::Db;
use sui_consistent_store::FrameworkSchema;
use sui_consistent_store::PipelineTaskKey;
use sui_consistent_store::Watermark;
use sui_consistent_store::restore::RestoreDriver;
use sui_consistent_store::restore::RestoreDriverConfig;
use sui_consistent_store::restore::RestoreSource;
use sui_consistent_store::restore::metrics::RestoreMetrics;
use sui_consistent_store::restore_state;
use sui_futures::service::Service;
use sui_indexer_alt_framework::pipeline::Processor;
use sui_types::storage::ObjectStore;
use sui_types::sui_system_state::SuiSystemStateTrait;
use sui_types::sui_system_state::get_sui_system_state;
use tracing::info;
use tracing::warn;

use crate::RestoreLayer;
use crate::RpcStoreReader;
use crate::RpcStoreSchema;
use crate::indexer::balance::Balance;
use crate::indexer::checkpoint_contents::CheckpointContents;
use crate::indexer::checkpoint_seq_by_digest::CheckpointSeqByDigest;
use crate::indexer::checkpoint_summary::CheckpointSummary;
use crate::indexer::effects::Effects;
use crate::indexer::epochs::Epochs;
use crate::indexer::event_bitmap::EventBitmap;
use crate::indexer::events::Events;
use crate::indexer::object_by_owner::ObjectByOwner;
use crate::indexer::object_by_type::ObjectByType;
use crate::indexer::object_version_by_checkpoint::ObjectVersionByCheckpoint;
use crate::indexer::objects::Objects;
use crate::indexer::package_versions::PackageVersions;
use crate::indexer::transaction_bitmap::TransactionBitmap;
use crate::indexer::transactions::Transactions;
use crate::indexer::tx_metadata_by_seq::TxMetadataBySeq;
use crate::indexer::tx_seq_by_digest::TxSeqByDigest;
use crate::schema::epochs;
use crate::schema::primitives::U64Be;
use crate::schema::pruning_watermark;

/// The embedded fullnode's **live cohort**: the pipelines that
/// [`restore_indexes`] bulk-loads and that are restored to the
/// perpetual store's tip `T`, then follow live from there. They are
/// bounded by the live object set, so a snapshot restore reproduces
/// them exactly.
///
/// Matches the live half of
/// [`PipelineLayer::embedded`](crate::config::PipelineLayer::embedded);
/// the `embedded_registers_only_cohort_pipelines` test pins the two
/// together.
pub const LIVE_COHORT: &[&str] = &[ObjectByOwner::NAME, ObjectByType::NAME, Balance::NAME];

/// The embedded fullnode's **history cohort**: the pipelines seeded to
/// the lowest available checkpoint `L` and backfilled upward from the
/// perpetual store, then followed live.
///
/// Most cannot be reconstructed from a live-object snapshot at all --
/// they record ledger history (`tx_seq` <-> digest maps, the
/// transaction and event bitmaps) and per-epoch metadata (`epochs`) --
/// so they are seeded, never restored.
///
/// `object_version_by_checkpoint` and `package_versions` are the
/// exceptions: when some prefix of the chain has been pruned
/// (`L > 0`) they are *both* restored and backfilled.
/// [`restore_indexes`] bulk-loads their floor rows at the tip `T` (the
/// versions live in the snapshot but predate the available window, so a
/// checkpoint-bounded read treats them as having always existed), and
/// the history seed then rewinds their `__watermark` to `L-1` so they
/// also backfill the per-checkpoint detail over `(L, T]` --
/// `object_version_by_checkpoint`'s per-checkpoint changes, and
/// `package_versions`'s real publish checkpoint for versions published
/// in the window. The embedded bootstrap runs the restore before the
/// seed, so the `L-1` watermark wins. When nothing has been pruned
/// (`L == 0`) the backfill covers the whole chain and rebuilds both
/// CFs in full, so the embedded bootstrap skips restoring them
/// entirely ([`RestoreLayer::live_only`]) -- restoring them would
/// waste bulk-load work and stamp `__watermark = T` rows that would
/// make them skip `(0, T]`.
///
/// [`RestoreLayer::live_only`]: crate::RestoreLayer::live_only
///
/// Matches the history half of
/// [`PipelineLayer::embedded`](crate::config::PipelineLayer::embedded).
pub const HISTORY_COHORT: &[&str] = &[
    Epochs::NAME,
    ObjectVersionByCheckpoint::NAME,
    PackageVersions::NAME,
    TxSeqByDigest::NAME,
    TxMetadataBySeq::NAME,
    TransactionBitmap::NAME,
    EventBitmap::NAME,
];

/// The two [`HISTORY_COHORT`] members whose floor rows the bulk
/// restore also loads when a prefix of the chain has been pruned
/// (`layer.history_floors`; see [`HISTORY_COHORT`]'s docs).
const HISTORY_FLOORS: &[&str] = &[ObjectVersionByCheckpoint::NAME, PackageVersions::NAME];

/// Register every [`Restore`]-implementing pipeline opted in by
/// `layer` on a [`RestoreDriver`] bound to `db` / `schema` and
/// `source`, then run the resulting [`Service`].
///
/// The live-cohort pipelines are always registered.
/// `object_version_by_checkpoint` and `package_versions` -- history-
/// cohort members whose floor rows are bulk-loaded from the live set
/// (the history seed separately rewinds their watermarks so they also
/// backfill `(L, T]`) -- are registered only when
/// `layer.history_floors` is set; a caller whose history cohort
/// replays from genesis leaves them to the backfill. The raw
/// [`Objects`] pipeline is only registered when `layer.objects` is
/// set. The returned `Service`'s primary task completes once every
/// registered pipeline transitions to [`RestoreState::Complete`].
///
/// [`Restore`]: sui_consistent_store::Restore
/// [`RestoreState::Complete`]: sui_consistent_store::restore_state::Complete
pub fn restore_indexes<Src: RestoreSource>(
    db: Db,
    schema: Arc<RpcStoreSchema>,
    source: Src,
    config: RestoreDriverConfig,
    layer: RestoreLayer,
    metrics: Arc<RestoreMetrics>,
) -> anyhow::Result<Service> {
    // Capture the anchor before the driver consumes `source`: the
    // checkpoint-pinned object index attributes every restored live
    // object to it.
    let target_checkpoint = source.target_checkpoint();
    let mut driver = RestoreDriver::new(db, schema, source, config, metrics);
    if layer.history_floors {
        // History-cohort members, but their floor rows are restored
        // from the live set; the embedded history seed later rewinds
        // their watermarks so they also backfill `(L, T]`.
        driver.register(ObjectVersionByCheckpoint::for_restore(target_checkpoint))?;
        driver.register(PackageVersions)?;
    }
    driver.register(ObjectByOwner)?;
    driver.register(ObjectByType)?;
    driver.register(Balance)?;
    if layer.objects {
        driver.register(Objects)?;
    }
    driver.run()
}

/// After [`restore_indexes`] returns, prime the framework state of
/// every pipeline that the restore did *not* cover so tip indexing
/// resumes from `target_watermark.checkpoint_hi_inclusive + 1`
/// across the board instead of replaying from genesis for the
/// raw-chain-data and bitmap pipelines.
///
/// Specifically, for every pipeline not in `layer`'s restored
/// set, writes:
///
/// - `__watermark = target_watermark` — the framework's
///   tip-resume reads this and starts at
///   `checkpoint_hi_inclusive + 1`.
/// - `__chain_id = target_chain_id` — pins the pipeline to the
///   chain the snapshot was taken from, matching what
///   [`restore_indexes`]'s finalize step already wrote for the
///   restored pipelines.
///
/// Also writes the singleton `pruning_watermark` so
/// `available_range` queries and the bitmap CFs' compaction
/// filters reflect that data only starts at the post-restore
/// floor (`tx_seq_lo = target_watermark.tx_hi`,
/// `checkpoint_lo = checkpoint_hi_inclusive + 1`).
///
/// Idempotent: re-running after a successful restore overwrites
/// the unrestored pipelines' watermarks with the same values and
/// re-writes the pruning row.
pub fn floor_unrestored_pipelines(
    db: &Db,
    schema: &Arc<RpcStoreSchema>,
    target_watermark: Watermark,
    target_chain_id: ChainId,
    layer: &RestoreLayer,
) -> anyhow::Result<()> {
    // The set of pipelines `restore_indexes` registered for this
    // `layer` -- their watermarks were written by the restore driver's
    // finalize step and must not be clobbered here.
    let mut restored: Vec<&'static str> =
        vec![ObjectByOwner::NAME, ObjectByType::NAME, Balance::NAME];
    if layer.history_floors {
        restored.push(ObjectVersionByCheckpoint::NAME);
        restored.push(PackageVersions::NAME);
    }
    if layer.objects {
        restored.push(Objects::NAME);
    }

    // Every rpc-store pipeline. Kept exhaustive so any new
    // pipeline added to `PipelineLayer` needs an explicit
    // decision here about whether it's a restore-time pipeline
    // or a tip-only one.
    let all: &[&'static str] = &[
        Epochs::NAME,
        CheckpointSummary::NAME,
        CheckpointContents::NAME,
        CheckpointSeqByDigest::NAME,
        Transactions::NAME,
        TxSeqByDigest::NAME,
        TxMetadataBySeq::NAME,
        Effects::NAME,
        Events::NAME,
        Objects::NAME,
        ObjectVersionByCheckpoint::NAME,
        ObjectByOwner::NAME,
        ObjectByType::NAME,
        Balance::NAME,
        PackageVersions::NAME,
        TransactionBitmap::NAME,
        EventBitmap::NAME,
    ];

    // Use the owned `FrameworkSchema` over `Db` (rather than the
    // borrowed view from `Db::framework`) so the `DbMap`s line up
    // with `Batch::put`'s `R = Db` expectation.
    let framework = FrameworkSchema::new(db.clone());
    let mut batch = db.batch();
    for name in all.iter().filter(|n| !restored.contains(n)) {
        let key = PipelineTaskKey::new(*name);
        batch
            .put(&framework.watermarks, &key, &target_watermark)
            .with_context(|| format!("stage __watermark for {name:?}"))?;
        batch
            .put(&framework.chain_ids, &key, &target_chain_id)
            .with_context(|| format!("stage __chain_id for {name:?}"))?;
    }

    let (k, v) = pruning_watermark::store(&pruning_watermark::Watermarks {
        tx_seq_lo: target_watermark.tx_hi,
        checkpoint_lo: target_watermark.checkpoint_hi_inclusive.saturating_add(1),
    });
    batch
        .put(&schema.pruning_watermark, &k, &v)
        .context("stage pruning_watermark row")?;

    // Seed the `epochs` row for the epoch the snapshot lands in. The
    // chain advanced to it at the anchor's end-of-epoch checkpoint,
    // but the `epochs` pipeline only emits a start record while
    // processing such a checkpoint, which tip indexing skips on
    // resume (it starts at anchor + 1). Without this seed, the
    // current epoch's row would never get its protocol version, gas
    // price, start timestamp, start checkpoint, or system state.
    // Requires the raw objects to read the on-chain `SuiSystemState`,
    // so it only runs when the `objects` CF was restored; failure to
    // read it is logged and skipped rather than failing the restore.
    if layer.objects {
        let reader = RpcStoreReader::new(db.clone(), schema.clone());
        let start_checkpoint = target_watermark.checkpoint_hi_inclusive.saturating_add(1);
        match seed_current_epoch_start(schema, &reader, Some(start_checkpoint), &mut batch) {
            Ok(epoch) => info!(
                epoch,
                start_checkpoint, "seeded start record for restore epoch"
            ),
            Err(e) => warn!(
                error = %e,
                "could not seed the restore epoch's start record; get_epoch / \
                 get_committee / Move type-layout resolution for the current epoch \
                 will be unavailable until the next epoch boundary",
            ),
        }
    }

    batch.commit().context("commit floor batch")?;

    // The watermark is durable; publish it to this database's
    // bitmap compaction filters.
    schema.set_pruning_floor(target_watermark.tx_hi);

    Ok(())
}

/// Stage a start record for the epoch reflected by the on-chain
/// `SuiSystemState` in `objects`, keyed by that epoch.
///
/// The `epochs` pipeline derives a start record only from an
/// end-of-epoch checkpoint's `epoch_info`, which a restore-then-tip
/// flow never processes (tip indexing resumes at `anchor + 1`), so a
/// freshly restored database has no start record for the epoch it
/// landed in. This reconstructs that record straight from the
/// restored object set: protocol version, reference gas price, and
/// epoch-start timestamp come from the `SuiSystemState`, the BCS of
/// which is stored so `get_committee` and Move type-layout resolution
/// work too.
///
/// `start_checkpoint` is supplied by the caller and may be `None`.
/// The formal-snapshot restore lands at an epoch boundary, so it
/// passes `Some(anchor + 1)`. The embedded-fullnode restore lands at
/// a *mid-epoch* tip, so the epoch's first checkpoint is unknown and
/// it passes `None`; the upward backfill fills `start_checkpoint` in
/// later if that boundary falls within the available range.
///
/// Stages a merge into `batch`; the caller commits. Returns the epoch
/// that was seeded.
pub fn seed_current_epoch_start(
    schema: &RpcStoreSchema,
    objects: &dyn ObjectStore,
    start_checkpoint: Option<u64>,
    batch: &mut Batch,
) -> anyhow::Result<u64> {
    let system_state =
        get_sui_system_state(objects).context("read SuiSystemState from restored objects")?;
    let epoch = system_state.epoch();
    let system_state_bcs = bcs::to_bytes(&system_state).context("bcs encode SuiSystemState")?;
    batch
        .merge(
            &schema.epochs,
            &U64Be(epoch),
            &epochs::start(
                system_state.protocol_version(),
                system_state.reference_gas_price(),
                system_state.epoch_start_timestamp_ms(),
                start_checkpoint,
                Some(system_state_bcs),
            ),
        )
        .context("stage epoch start seed")?;
    Ok(epoch)
}

/// Seed the framework state for the embedded fullnode's
/// [`HISTORY_COHORT`] after [`restore_indexes`] has bulk-loaded the
/// [`LIVE_COHORT`].
///
/// The live cohort resumes from the restore target `T` (written by
/// the restore driver's finalize step). The history cohort is *not*
/// restored; instead each of its pipelines is seeded to
/// `history_watermark` — the lowest available checkpoint `L` in the
/// perpetual store — so tip indexing backfills `(L, T]` from the
/// perpetual store and then follows live. For each history pipeline
/// this writes:
///
/// - `__watermark = history_watermark` — the framework resumes at
///   `history_watermark.checkpoint_hi_inclusive + 1`.
/// - `__chain_id = chain_id` — pins the chain, matching what the
///   restore driver wrote for the live cohort.
///
/// When `objects` is supplied, also seeds the current epoch's
/// `epochs` row from its on-chain `SuiSystemState` — a *partial*
/// start record without `start_checkpoint` (see
/// [`seed_current_epoch_start`]) — so `get_epoch` / `get_committee`
/// and Move type-layout resolution work immediately after restore
/// rather than only once the backfill reaches the epoch's boundary.
/// `objects` is read through the [`ObjectStore`] trait, so the
/// embedded caller passes the validator's perpetual store directly
/// (this crate stays free of any `sui-core` dependency).
///
/// Stamps the singleton `pruning_watermark` at the lowest available
/// checkpoint `L` — the first checkpoint the backfill will index,
/// `history_watermark.checkpoint_hi_inclusive + 1` — with
/// `tx_seq_lo = history_watermark.tx_hi` (the first `tx_seq` that
/// checkpoint contributes). This records that nothing below `L` is
/// available and sets the bitmap compaction filter's floor. The
/// backfill only ever writes `tx_seq` at or above the floor, so the
/// filter drops nothing it produces. (The *upper* bound of history
/// availability while a backfill is in progress is a separate
/// concern, handled elsewhere.)
///
/// Idempotent: re-running overwrites the same rows. Does not touch
/// the live cohort or the deactivated (perpetual-store-served) CFs.
pub fn seed_history_cohort(
    db: &Db,
    schema: &RpcStoreSchema,
    history_watermark: Watermark,
    chain_id: ChainId,
    objects: Option<&dyn ObjectStore>,
) -> anyhow::Result<()> {
    let framework = FrameworkSchema::new(db.clone());
    let mut batch = db.batch();

    for name in HISTORY_COHORT {
        let key = PipelineTaskKey::new(*name);
        batch
            .put(&framework.watermarks, &key, &history_watermark)
            .with_context(|| format!("stage __watermark for {name:?}"))?;
        batch
            .put(&framework.chain_ids, &key, &chain_id)
            .with_context(|| format!("stage __chain_id for {name:?}"))?;
    }

    // Stamp the pruning floor at the lowest available checkpoint `L`
    // (the first checkpoint the backfill will index). `tx_seq_lo` is
    // the tx count through the seed point, i.e. the first `tx_seq`
    // checkpoint `L` contributes, so the bitmap compaction filter
    // drops nothing the backfill produces.
    let tx_seq_lo = history_watermark.tx_hi;
    let (k, v) = pruning_watermark::store(&pruning_watermark::Watermarks {
        tx_seq_lo,
        checkpoint_lo: history_watermark.checkpoint_hi_inclusive.saturating_add(1),
    });
    batch
        .put(&schema.pruning_watermark, &k, &v)
        .context("stage pruning_watermark row")?;

    if let Some(objects) = objects {
        // Mid-epoch restore: the epoch's first checkpoint precedes the
        // tip, so seed a partial start record (no `start_checkpoint`).
        match seed_current_epoch_start(schema, objects, None, &mut batch) {
            Ok(epoch) => info!(epoch, "seeded partial start record for the current epoch"),
            Err(e) => warn!(
                error = %e,
                "could not seed the current epoch's start record; get_epoch / \
                 get_committee / Move type-layout resolution for the current epoch \
                 will be unavailable until the backfill reaches its boundary",
            ),
        }
    }

    batch.commit().context("commit history-cohort seed batch")?;

    // The watermark is durable; publish it to this database's
    // bitmap compaction filters.
    schema.set_pruning_floor(tx_seq_lo);

    Ok(())
}

/// Whether any embedded-cohort pipeline has a bulk restore in
/// progress: an `__restore` row in the `InProgress` state, meaning a
/// restore crashed mid-run and its per-shard cursors are resumable.
///
/// Tip indexing never writes `__restore` rows, and the restore
/// driver's finalize flips every registered row to `Complete` in one
/// atomic batch, so an `InProgress` row is unambiguous. The embedded
/// bootstrap uses this to tell a resumable restore apart from other
/// states that merely lack live watermarks -- a fresh store, or a
/// from-genesis run that crashed between the live pipelines' first
/// commits -- which must be cleared before restoring, because the
/// bulk restore merges and puts on top of whatever rows exist.
pub fn restore_in_progress(db: &Db) -> anyhow::Result<bool> {
    let framework = db.framework();
    for name in LIVE_COHORT.iter().chain(HISTORY_COHORT) {
        let key = PipelineTaskKey::new(*name);
        let state = framework
            .restore
            .get(&key)
            .with_context(|| format!("reading restore state for {name}"))?;
        if let Some(state) = state
            && matches!(state.state, Some(restore_state::State::InProgress(_)))
        {
            return Ok(true);
        }
    }
    Ok(false)
}

/// Whether the persisted history-cohort state is exactly the
/// post-restore, pre-seed state: both history-floor pipelines
/// (`object_version_by_checkpoint` and `package_versions`) carry the
/// watermark the restore driver's finalize stamped, while no other
/// history pipeline has ever committed.
///
/// This is the only state the embedded bootstrap may repair with
/// [`seed_history_cohort`] alone (the crash window between the
/// restore's finalize and the history seed): the floor pipelines'
/// restored rows anchor checkpoint-pinned reads at the target, and
/// the seed's rewind makes the backfill cover everything else in
/// `(L, T]`. Any other history coverage below the available floor
/// sits behind a pruned gap the backfill can never fill -- seeding
/// over it would advertise the range above the floor as served while
/// stale sub-floor rows resolve checkpoint-pinned reads to pre-gap
/// versions -- so such a store must be cleared and re-restored
/// instead.
pub fn history_seed_pending(db: &Db) -> anyhow::Result<bool> {
    let framework = db.framework();
    for name in HISTORY_COHORT {
        let key = PipelineTaskKey::new(*name);
        let watermark = framework
            .watermarks
            .get(&key)
            .with_context(|| format!("reading watermark for {name}"))?;
        if watermark.is_some() != HISTORY_FLOORS.contains(name) {
            return Ok(false);
        }
    }
    Ok(true)
}

#[cfg(test)]
mod tests {
    use async_trait::async_trait;
    use bytes::Bytes;
    use futures::StreamExt;
    use futures::stream;
    use futures::stream::BoxStream;
    use sui_consistent_store::ChainId;
    use sui_consistent_store::Db;
    use sui_consistent_store::DbOptions;
    use sui_consistent_store::PipelineTaskKey;
    use sui_consistent_store::Watermark;
    use sui_consistent_store::restore::RestoreChunk;
    use sui_consistent_store::restore_state;
    use sui_indexer_alt_framework::pipeline::Processor;
    use sui_types::base_types::ObjectID;
    use sui_types::base_types::SuiAddress;
    use sui_types::object::Object;

    use super::*;
    use crate::RpcStoreSchema;
    use crate::indexer::objects::Objects;
    use crate::schema::object_by_owner::OwnerKind;

    /// Minimal [`RestoreSource`] that wraps a `Vec<RestoreChunk>`
    /// and uses the 4-byte BE chunk index as cursor. Lets us
    /// drive the end-to-end pipeline registration / commit path
    /// without standing up a real snapshot.
    struct VecSource {
        target: u64,
        chain_id: ChainId,
        chunks: Vec<RestoreChunk>,
    }

    impl VecSource {
        fn from_objects(target: u64, chain_id: ChainId, objects: Vec<Vec<Object>>) -> Self {
            let chunks = objects
                .into_iter()
                .enumerate()
                .map(|(i, objs)| RestoreChunk {
                    objects: objs,
                    cursor: Bytes::copy_from_slice(&(i as u32).to_be_bytes()),
                })
                .collect();
            Self {
                target,
                chain_id,
                chunks,
            }
        }
    }

    #[async_trait]
    impl RestoreSource for VecSource {
        fn target_checkpoint(&self) -> u64 {
            self.target
        }

        fn target_chain_id(&self) -> ChainId {
            self.chain_id
        }

        fn shards(&self) -> u32 {
            1
        }

        fn stream(
            &self,
            shard_id: u32,
            cursor: Option<Bytes>,
        ) -> BoxStream<'_, anyhow::Result<RestoreChunk>> {
            assert_eq!(shard_id, 0);
            let resume_after = cursor.map(|c| {
                let mut buf = [0u8; 4];
                buf.copy_from_slice(&c[..4]);
                u32::from_be_bytes(buf)
            });
            let chunks: Vec<_> = self
                .chunks
                .iter()
                .enumerate()
                .filter_map(|(i, chunk)| {
                    let i = i as u32;
                    if let Some(after) = resume_after
                        && i <= after
                    {
                        None
                    } else {
                        Some(Ok(RestoreChunk {
                            objects: chunk.objects.clone(),
                            cursor: chunk.cursor.clone(),
                        }))
                    }
                })
                .collect();
            stream::iter(chunks).boxed()
        }
    }

    /// End-to-end: drive a handful of address-owned objects
    /// through every registered pipeline. Verifies that the
    /// rows we expect end up in `object_version_by_checkpoint` and
    /// `object_by_owner`, and that every pipeline's
    /// `__restore` / `__watermark` rows are set up for the
    /// tip-indexer to take over.
    #[tokio::test]
    async fn restore_indexes_populates_schema_and_finalises() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let schema = Arc::new(schema);

        let owner = SuiAddress::random_for_testing_only();
        let objects: Vec<Object> = (1..=4u8)
            .map(|i| Object::with_id_owner_for_testing(ObjectID::from_single_byte(i), owner))
            .collect();

        let chain_id = ChainId([7u8; 32]);
        let source = VecSource::from_objects(123, chain_id, vec![objects.clone()]);

        restore_indexes(
            db.clone(),
            schema.clone(),
            source,
            RestoreDriverConfig::default(),
            RestoreLayer::indexes_only(),
            RestoreMetrics::new(None, &prometheus::Registry::new()),
        )
        .unwrap()
        .shutdown()
        .await
        .unwrap();

        // Each object's restore floor row landed in the checkpoint-pinned
        // index at the anchor checkpoint (123).
        for o in &objects {
            assert_eq!(
                schema
                    .get_object_version_at_checkpoint(o.id(), 123)
                    .unwrap(),
                Some(o.version()),
            );
        }

        // Owner index has every object under the same
        // AddressOwner(owner) key.
        let owned: Vec<(OwnerKind, ObjectID)> = schema
            .iter_objects_owned_by_address(owner)
            .unwrap()
            .map(Result::unwrap)
            .map(|(key, _v)| (key.kind, key.object_id))
            .collect();
        let mut got_ids: Vec<_> = owned.iter().map(|(_, id)| *id).collect();
        got_ids.sort();
        let mut expected_ids: Vec<_> = objects.iter().map(|o| o.id()).collect();
        expected_ids.sort();
        assert_eq!(got_ids, expected_ids);
        for (kind, _) in &owned {
            assert!(matches!(kind, OwnerKind::AddressOwner(addr) if *addr == owner));
        }

        // `indexes_only` did not register `objects`, so the
        // `(id, version)` CF stays empty.
        for o in &objects {
            assert_eq!(schema.get_object_by_key(o.id(), o.version()).unwrap(), None,);
        }

        // Every registered pipeline finished and has __restore
        // Complete, __watermark, and __chain_id all set. `objects`
        // was not registered with `indexes_only`, so it has no
        // __restore row at all.
        for name in [
            ObjectVersionByCheckpoint::NAME,
            ObjectByOwner::NAME,
            ObjectByType::NAME,
            Balance::NAME,
            PackageVersions::NAME,
        ] {
            let key = PipelineTaskKey::new(name);
            let state = db.framework().restore.get(&key).unwrap().unwrap();
            match state.state.unwrap() {
                restore_state::State::Complete(c) => assert_eq!(c.restored_at, 123),
                other => panic!("expected Complete, got {other:?}"),
            }
            let wm = db.framework().watermarks.get(&key).unwrap().unwrap();
            assert_eq!(wm, Watermark::for_checkpoint(123));
            let pinned_chain_id = db.framework().chain_ids.get(&key).unwrap().unwrap();
            assert_eq!(pinned_chain_id, chain_id);
        }
        let objects_key = PipelineTaskKey::new(Objects::NAME);
        assert!(
            db.framework().restore.get(&objects_key).unwrap().is_none(),
            "indexes_only should leave the objects pipeline unregistered",
        );

        // A finalized `indexes_only` restore is exactly the state the
        // history seed repairs.
        assert!(!restore_in_progress(&db).unwrap());
        assert!(history_seed_pending(&db).unwrap());
    }

    /// `floor_unrestored_pipelines` writes a `__watermark` /
    /// `__chain_id` row for every pipeline outside the restored
    /// set and stamps the singleton `pruning_watermark` so the
    /// available range tracks the post-restore floor.
    #[test]
    fn floor_unrestored_pipelines_writes_watermarks_for_tip_only_pipelines() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let schema = Arc::new(schema);

        let chain_id = ChainId([42u8; 32]);
        let target = Watermark {
            epoch_hi_inclusive: 7,
            checkpoint_hi_inclusive: 1_000,
            tx_hi: 5_000,
            timestamp_ms_hi_inclusive: 1_700_000_000_000,
        };

        floor_unrestored_pipelines(&db, &schema, target, chain_id, &RestoreLayer::all()).unwrap();

        // Sample raw-chain-data / bitmap pipelines that the
        // formal-snapshot path doesn't cover — every one of them
        // should be primed with the target watermark + chain id.
        for name in [
            Epochs::NAME,
            CheckpointSummary::NAME,
            CheckpointContents::NAME,
            CheckpointSeqByDigest::NAME,
            Transactions::NAME,
            TxSeqByDigest::NAME,
            TxMetadataBySeq::NAME,
            Effects::NAME,
            Events::NAME,
            TransactionBitmap::NAME,
            EventBitmap::NAME,
        ] {
            let key = PipelineTaskKey::new(name);
            assert_eq!(
                db.framework().watermarks.get(&key).unwrap(),
                Some(target),
                "{name} should have the post-restore watermark",
            );
            assert_eq!(
                db.framework().chain_ids.get(&key).unwrap(),
                Some(chain_id),
                "{name} should pin the restored chain id",
            );
        }

        // Restored pipelines are left to whatever the restore
        // driver wrote (here: nothing, since we didn't actually
        // run a restore in this test). The helper must not
        // clobber them.
        for name in [
            ObjectVersionByCheckpoint::NAME,
            ObjectByOwner::NAME,
            ObjectByType::NAME,
            Balance::NAME,
            PackageVersions::NAME,
            Objects::NAME,
        ] {
            let key = PipelineTaskKey::new(name);
            assert!(
                db.framework().watermarks.get(&key).unwrap().is_none(),
                "{name} watermark should be untouched by the floor helper",
            );
            assert!(
                db.framework().chain_ids.get(&key).unwrap().is_none(),
                "{name} chain id should be untouched by the floor helper",
            );
        }

        // Pruning singleton reflects the post-restore floor: tx
        // ids and checkpoint sequences below this row aren't
        // available in the new database.
        assert_eq!(
            schema.get_pruning_watermarks().unwrap(),
            Some(crate::schema::pruning_watermark::Watermarks {
                tx_seq_lo: target.tx_hi,
                checkpoint_lo: target.checkpoint_hi_inclusive + 1,
            }),
        );
    }

    /// With `RestoreLayer::indexes_only`, the `objects` pipeline
    /// is *not* in the restored set, so the floor helper primes
    /// it the same way it does the raw-chain-data pipelines.
    #[test]
    fn floor_unrestored_pipelines_includes_objects_when_layer_skips_it() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let schema = Arc::new(schema);

        let chain_id = ChainId([11u8; 32]);
        let target = Watermark::for_checkpoint(42);

        floor_unrestored_pipelines(
            &db,
            &schema,
            target,
            chain_id,
            &RestoreLayer::indexes_only(),
        )
        .unwrap();

        let key = PipelineTaskKey::new(Objects::NAME);
        assert_eq!(db.framework().watermarks.get(&key).unwrap(), Some(target),);
        assert_eq!(db.framework().chain_ids.get(&key).unwrap(), Some(chain_id));
    }

    /// `RestoreLayer::all` additionally registers the `objects`
    /// pipeline, so every restored live object lands in the
    /// `(id, version)` CF and the pipeline itself transitions to
    /// `Complete`.
    #[tokio::test]
    async fn restore_indexes_with_objects_layer_populates_objects_cf() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let schema = Arc::new(schema);

        let owner = SuiAddress::random_for_testing_only();
        let objects: Vec<Object> = (1..=4u8)
            .map(|i| Object::with_id_owner_for_testing(ObjectID::from_single_byte(i), owner))
            .collect();

        let chain_id = ChainId([9u8; 32]);
        let source = VecSource::from_objects(123, chain_id, vec![objects.clone()]);

        restore_indexes(
            db.clone(),
            schema.clone(),
            source,
            RestoreDriverConfig::default(),
            RestoreLayer::all(),
            RestoreMetrics::new(None, &prometheus::Registry::new()),
        )
        .unwrap()
        .shutdown()
        .await
        .unwrap();

        // Every object lands at its current version in `objects`.
        for o in &objects {
            assert_eq!(
                schema.get_object_by_key(o.id(), o.version()).unwrap(),
                Some(o.clone()),
            );
        }

        // The `objects` pipeline's __restore / __watermark /
        // __chain_id rows all match the source target.
        let key = PipelineTaskKey::new(Objects::NAME);
        let state = db.framework().restore.get(&key).unwrap().unwrap();
        match state.state.unwrap() {
            restore_state::State::Complete(c) => assert_eq!(c.restored_at, 123),
            other => panic!("expected Complete, got {other:?}"),
        }
        assert_eq!(
            db.framework().watermarks.get(&key).unwrap().unwrap(),
            Watermark::for_checkpoint(123),
        );
        assert_eq!(
            db.framework().chain_ids.get(&key).unwrap().unwrap(),
            chain_id,
        );
    }

    /// `seed_history_cohort` primes every history-cohort pipeline with
    /// the seed watermark and the chain id, stamps the pruning floor at
    /// the lowest available checkpoint `L`, and leaves the live cohort
    /// untouched (the restore driver owns it).
    #[test]
    fn seed_history_cohort_seeds_history_watermarks_and_pruning_floor() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();

        let chain_id = ChainId([5u8; 32]);
        // Seed point: committed through checkpoint 999 / tx 5000, so the
        // backfill resumes at (and the floor is) checkpoint 1000.
        let seed = Watermark {
            checkpoint_hi_inclusive: 999,
            tx_hi: 5_000,
            ..Watermark::default()
        };
        seed_history_cohort(&db, &schema, seed, chain_id, None).unwrap();

        // History cohort resumes from the seed point, pinned to the
        // chain.
        for name in HISTORY_COHORT {
            let key = PipelineTaskKey::new(*name);
            assert_eq!(
                db.framework().watermarks.get(&key).unwrap(),
                Some(seed),
                "{name} should resume from the seed point",
            );
            assert_eq!(
                db.framework().chain_ids.get(&key).unwrap(),
                Some(chain_id),
                "{name} should pin the chain id",
            );
        }

        // Live cohort is the restore driver's responsibility — the
        // history seed must not touch it.
        for name in LIVE_COHORT {
            let key = PipelineTaskKey::new(*name);
            assert!(
                db.framework().watermarks.get(&key).unwrap().is_none(),
                "{name} watermark must be untouched by the history seed",
            );
        }

        // Pruning floor sits at the lowest available checkpoint
        // (seed + 1) and the seed point's tx count.
        assert_eq!(
            schema.get_pruning_watermarks().unwrap(),
            Some(crate::schema::pruning_watermark::Watermarks {
                tx_seq_lo: 5_000,
                checkpoint_lo: 1_000,
            }),
        );
    }

    /// `RestoreLayer::live_only` registers only the live-cohort
    /// pipelines: `object_version_by_checkpoint` and
    /// `package_versions` get no restore state, no watermark, no
    /// chain id, and no floor rows, so a from-genesis backfill owns
    /// them outright and never skips `(0, T]`.
    #[tokio::test]
    async fn restore_indexes_live_only_skips_history_floors() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let schema = Arc::new(schema);

        let owner = SuiAddress::random_for_testing_only();
        let objects: Vec<Object> = (1..=4u8)
            .map(|i| Object::with_id_owner_for_testing(ObjectID::from_single_byte(i), owner))
            .collect();

        let chain_id = ChainId([8u8; 32]);
        let source = VecSource::from_objects(123, chain_id, vec![objects.clone()]);

        restore_indexes(
            db.clone(),
            schema.clone(),
            source,
            RestoreDriverConfig::default(),
            RestoreLayer::live_only(),
            RestoreMetrics::new(None, &prometheus::Registry::new()),
        )
        .unwrap()
        .shutdown()
        .await
        .unwrap();

        // The live cohort restored and finalized as usual.
        for name in LIVE_COHORT {
            let key = PipelineTaskKey::new(*name);
            let state = db.framework().restore.get(&key).unwrap().unwrap();
            assert!(
                matches!(
                    state.state.unwrap(),
                    restore_state::State::Complete(c) if c.restored_at == 123,
                ),
                "{name} should complete its restore",
            );
            assert_eq!(
                db.framework().watermarks.get(&key).unwrap(),
                Some(Watermark::for_checkpoint(123)),
                "{name} should resume from the restore target",
            );
        }

        // The history-floor pipelines were never registered: no
        // restore state, no watermark (so tip indexing backfills them
        // from genesis), and no floor rows in their CFs.
        for name in [ObjectVersionByCheckpoint::NAME, PackageVersions::NAME] {
            let key = PipelineTaskKey::new(name);
            assert!(
                db.framework().restore.get(&key).unwrap().is_none(),
                "{name} should have no restore state under live_only",
            );
            assert!(
                db.framework().watermarks.get(&key).unwrap().is_none(),
                "{name} should have no watermark under live_only",
            );
        }
        for o in &objects {
            assert_eq!(
                schema
                    .get_object_version_at_checkpoint(o.id(), 123)
                    .unwrap(),
                None,
                "live_only must write no restore floor rows",
            );
        }

        // A `live_only` restore leaves nothing to seed: the whole
        // history cohort (floors included) belongs to the
        // from-genesis backfill.
        assert!(!history_seed_pending(&db).unwrap());
    }

    /// The bootstrap state predicates track the restore and seed
    /// lifecycle: a fresh store reports neither, a mid-run restore
    /// reports `restore_in_progress`, a finalized `indexes_only`
    /// restore reports `history_seed_pending` until the history seed
    /// runs, and after the seed both are off.
    #[tokio::test]
    async fn state_predicates_track_restore_and_seed_lifecycle() {
        let dir = tempfile::tempdir().unwrap();
        let (db, schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let schema = Arc::new(schema);

        // Fresh store: nothing in progress, nothing to seed.
        assert!(!restore_in_progress(&db).unwrap());
        assert!(!history_seed_pending(&db).unwrap());

        // A registered-but-unfinished restore (the state a mid-run
        // crash leaves behind): an `InProgress` row with no watermark.
        let framework = FrameworkSchema::new(db.clone());
        let in_progress = sui_consistent_store::RestoreState::default().with_in_progress(
            restore_state::InProgress {
                target_checkpoint: 123,
                shards: Default::default(),
            },
        );
        let mut batch = db.batch();
        batch
            .put(
                &framework.restore,
                &PipelineTaskKey::new(Balance::NAME),
                &in_progress,
            )
            .unwrap();
        batch.commit().unwrap();
        assert!(restore_in_progress(&db).unwrap());
        assert!(!history_seed_pending(&db).unwrap());

        // Run the restore to completion at the same target: every
        // registered row flips to `Complete` and the store enters the
        // post-restore, pre-seed state.
        let owner = SuiAddress::random_for_testing_only();
        let objects = vec![Object::with_id_owner_for_testing(
            ObjectID::from_single_byte(1),
            owner,
        )];
        let source = VecSource::from_objects(123, ChainId([3u8; 32]), vec![objects]);
        restore_indexes(
            db.clone(),
            schema.clone(),
            source,
            RestoreDriverConfig::default(),
            RestoreLayer::indexes_only(),
            RestoreMetrics::new(None, &prometheus::Registry::new()),
        )
        .unwrap()
        .shutdown()
        .await
        .unwrap();
        assert!(!restore_in_progress(&db).unwrap());
        assert!(history_seed_pending(&db).unwrap());

        // The history seed stamps every history pipeline, closing the
        // pre-seed window.
        seed_history_cohort(
            &db,
            &schema,
            Watermark::for_checkpoint(99),
            ChainId([3u8; 32]),
            None,
        )
        .unwrap();
        assert!(!restore_in_progress(&db).unwrap());
        assert!(!history_seed_pending(&db).unwrap());
    }

    /// `history_seed_pending` requires exactly the post-restore
    /// watermark shape: floor pipelines watermarked, everything else
    /// in the history cohort unwatermarked. Partial coverage in
    /// either direction is not seedable.
    #[test]
    fn history_seed_pending_rejects_partial_history_coverage() {
        let dir = tempfile::tempdir().unwrap();
        let (db, _schema) = Db::open::<RpcStoreSchema>(dir.path(), DbOptions::default()).unwrap();
        let framework = FrameworkSchema::new(db.clone());

        let stamp = |name: &'static str| {
            let mut batch = db.batch();
            batch
                .put(
                    &framework.watermarks,
                    &PipelineTaskKey::new(name),
                    &Watermark::for_checkpoint(123),
                )
                .unwrap();
            batch.commit().unwrap();
        };

        // A non-floor history pipeline with prior committed coverage
        // (and no floor watermarks) is not the pre-seed state.
        stamp(Epochs::NAME);
        assert!(!history_seed_pending(&db).unwrap());

        // Both floors watermarked on top of that: still committed
        // non-floor coverage, still not seedable.
        stamp(ObjectVersionByCheckpoint::NAME);
        stamp(PackageVersions::NAME);
        assert!(!history_seed_pending(&db).unwrap());
    }

    /// The live and history cohorts are disjoint and each has the
    /// expected size. (Their union being exactly the embedded layer's
    /// enabled set is pinned by `embedded_registers_only_cohort_pipelines`.)
    #[test]
    fn cohorts_are_disjoint() {
        let live: std::collections::BTreeSet<_> = LIVE_COHORT.iter().collect();
        let history: std::collections::BTreeSet<_> = HISTORY_COHORT.iter().collect();
        assert!(live.is_disjoint(&history), "cohorts must not overlap");
        assert_eq!(live.len(), 3);
        assert_eq!(history.len(), 7);
    }
}