krabka_restore/

lib.rs

//! Offline point-in-time restore of a krabka cluster.
//!
//! `krabka restore` reads a KIP-405 tiered-storage archive from object storage
//! and materializes a complete, bootable krabka data directory, replayed up to
//! a bound and verified segment by segment as it rehydrates. The bound is an
//! offset, a timestamp, or a set of exclude-record predicates, so an operator
//! can recover an event-sourced system to the state it held just before a bad
//! write. Recovery of that kind is a hand-built runbook everywhere else.
//!
//! The tool runs when the cluster does not. It reads the archive through the
//! object-store layer and formats the target through [`krabka_format`], and it
//! does not depend on the broker.
//!
//! # Stages
//!
//! A restore runs five stages, and [`restore`] drives them in order:
//!
//! 1. **Discover** lists the archive and groups the keys into segments per
//!    topic partition.
//! 2. **Verify** checks the framing and the CRC of every archived segment, and
//!    derives its true end offset and maximum timestamp.
//! 3. **Bound** compiles the operator's predicates and decides which batches
//!    and records survive.
//! 4. **Materialize** formats the target log directory and writes the
//!    partition data.
//! 5. **Report** renders what happened, as text or as JSON.
//!
//! # Limits
//!
//! A batch that the bound filters is re-encoded from the records that survive,
//! so its bytes are not identical to the archived bytes. `--exclude-key` and
//! `--exclude-header` match raw bytes and decode no payload. Without
//! `--rlmm-snapshot` a segment the old cluster had marked for deletion is
//! indistinguishable from a live one. Without `--metadata-snapshot`, topic
//! configuration, ACLs, client quotas, SCRAM credentials, and finalized
//! feature levels cannot be recovered.
//!
//! # Errors and exit codes
//!
//! The library returns [`RestoreError`]. Only [`run`] and the binary map one
//! onto an exit code, so an embedding tool keeps the structured error. See
//! [`EXIT_OK`] for the codes.

use {
    self::authenticity::SourceAuthentication,
    clap::Parser,
    krabka_remote_storage::{AuthenticatedArchive, diskless::capture::DisklessWalCapture},
    std::{
        collections::{BTreeMap, BTreeSet},
        ffi::OsString,
        fs,
        path::Path,
    },
};

mod args;
mod authenticity;
mod backend;
mod bound;
mod discover;
mod diskless;
mod error;
mod materialize;
mod report;
mod verify;

pub use self::{
    args::{
        ArchiveArgs, HeaderPattern, OffsetBound, OffsetRange, PartitionRef, RestoreArgs, TargetArgs,
    },
    backend::{ArchiveStore, object_store_config, open_archive},
    bound::{BatchDecision, Predicates, RecordDecision},
    discover::{
        ArchiveInventory, ArchiveObject, PartitionInventory, SegmentInventory,
        UNRECOGNIZED_SAMPLE_LIMIT, UnrecognizedKeys, inventory,
    },
    error::{
        EXIT_ARCHIVE_UNREADABLE, EXIT_BAD_ARGUMENTS, EXIT_DIRTY_LOG_DIR, EXIT_INTEGRITY,
        EXIT_MATERIALIZE, EXIT_OK, RestoreError,
    },
    materialize::{
        FormatTargetOutcome, SegmentOutcome, format_target, seed_rlmm_snapshot, write_segment,
    },
    report::{
        AuthenticationReport, DisklessPartitionReport, DisklessRestoreReport,
        MetadataRestoreReport, PartitionReport, ReportFormat, RestoreReport, SkippedSegment,
    },
    verify::{
        LeaderEpochStart, SegmentFacts, VerifiedSegment, verify_segment,
        verify_segment_authenticated,
    },
};

/// The restore command line.
///
/// The binary and [`run_from_args`] share it, so both accept exactly the same
/// flags.
#[derive(Parser)]
// `long_about = None` keeps the struct's rustdoc out of `--help`. Without it
// clap renders the doc comment, intra-doc links and all, as the long help.
#[command(
    name = "krabka-restore",
    version,
    about = "Rebuild a bootable krabka log directory from a KIP-405 tiered-storage archive, \
             replayed to a point in time",
    long_about = None
)]
pub struct Cli {
    /// The restore's arguments, flattened so they are top-level flags.
    #[command(flatten)]
    pub args: RestoreArgs,
}

/// Run a restore and return the process exit code.
///
/// Every failure is reported on stderr and mapped onto the exit code
/// [`RestoreError::exit_code`] gives, rather than raised.
pub async fn run(args: RestoreArgs) -> i32 {
    let format = args.report;
    match restore(&args).await {
        Ok(report) => {
            if let Some(warning) = report.metadata.warning() {
                eprintln!("krabka restore: {warning}");
            }
            match report.render(format) {
                Ok(rendered) => {
                    println!("{rendered}");
                    EXIT_OK
                }
                Err(error) => {
                    eprintln!("krabka restore: cannot render the report: {error}");
                    EXIT_MATERIALIZE
                }
            }
        }
        Err(error) => {
            eprintln!("krabka restore: {error}");
            error.exit_code()
        }
    }
}

/// Run a restore from an argv-style iterator and return the process exit code.
///
/// `--help` and `--version` render to stdout and return [`EXIT_OK`]. A
/// malformed command line renders to stderr and returns [`EXIT_BAD_ARGUMENTS`].
pub async fn run_from_args<I, T>(argv: I) -> i32
where
    I: IntoIterator<Item = T>,
    T: Into<OsString> + Clone,
{
    match Cli::try_parse_from(argv) {
        Ok(cli) => run(cli.args).await,
        Err(error) => {
            let _ = error.print();
            if error.use_stderr() {
                EXIT_BAD_ARGUMENTS
            } else {
                EXIT_OK
            }
        }
    }
}

/// Run a restore and return its report.
///
/// This is the entry point for a caller that wants the structured outcome and
/// the structured error rather than an exit code.
///
/// # Errors
///
/// Returns [`RestoreError`] for a bad bound, a bound that names a topic
/// partition the archive does not hold, a target that is not empty, an
/// archive that cannot be read, a segment that fails verification, or a target
/// that rejects a write. `--continue-on-corrupt` turns a verification failure
/// into a skipped segment in the report instead of an error.
pub async fn restore(args: &RestoreArgs) -> Result<RestoreReport, RestoreError> {
    args.validate()?;
    // Check the target before the archive scan. An operator who pointed at a
    // live data directory learns it in a second, not after a full download.
    ensure_empty_log_dir(&args.target.log_dir)?;

    let store = open_archive(args)?;
    let diskless_capture = diskless::load(args.archive.diskless_wal_capture.as_deref()).await?;
    let mut archive = inventory(&store, args).await?;
    let authenticity::AuthenticatedInputs {
        source: mut authenticated,
        capture: diskless_authentication,
        snapshots:
            authenticity::AuthenticatedSnapshots {
                metadata: metadata_bytes,
                rlmm: rlmm_snapshot,
            },
        has_classic,
        unclaimed,
    } = authenticity::authenticate_inputs(&store, args, &mut archive, diskless_capture.as_ref())
        .await?;
    let diskless_claims = diskless_authentication
        .as_ref()
        .map(|authentication| &authentication.claims);
    if let Some(snapshot) = rlmm_snapshot.as_ref().filter(|_| has_classic) {
        discover::reconcile_authenticated_snapshot(&mut archive.partitions, args, snapshot)?;
        if archive.partitions.is_empty() {
            return Err(RestoreError::EmptyArchive {
                prefix: store.prefix().unwrap_or("").to_owned(),
            });
        }
    }
    if let Some(capture) = &diskless_capture {
        diskless::add_partitions(&mut archive, capture, args)?;
    }
    if archive.partitions.is_empty() {
        return Err(RestoreError::EmptyArchive {
            prefix: store.prefix().unwrap_or("").to_owned(),
        });
    }
    let predicates = Predicates::from_args(args)?;
    let (format, diskless_next_offsets) = format_with_next_offsets(
        &store,
        args,
        &archive,
        &predicates,
        diskless_capture.as_ref(),
        diskless_claims,
        metadata_bytes.as_deref(),
    )
    .await?;
    materialize::seed_rlmm_snapshot_from(
        args,
        authenticated
            .as_ref()
            .and_then(|source| source.archive.as_ref())
            .map(AuthenticatedArchive::report),
        rlmm_snapshot.as_ref(),
    )?;

    let mut consumed_authenticated_objects = BTreeSet::new();
    if metadata_bytes.is_some() {
        consumed_authenticated_objects.insert("cluster-metadata.checkpoint".to_owned());
    }
    if rlmm_snapshot.is_some() {
        consumed_authenticated_objects.insert("rlmm-snapshot".to_owned());
    }
    let mut partitions = Vec::with_capacity(archive.partitions.len());
    let mut skipped = unclaimed;
    for entry in &archive.partitions {
        if entry.segments.is_empty() {
            continue;
        }
        let mut segments = Vec::with_capacity(entry.segments.len());
        let mut skipped_predecessor = false;
        for segment in &entry.segments {
            let verification = match &authenticated {
                Some(SourceAuthentication {
                    archive: Some(authenticated),
                    ..
                }) => {
                    verify_segment_authenticated(
                        &store,
                        &entry.partition,
                        segment,
                        authenticated.objects(),
                    )
                    .await
                }
                Some(SourceAuthentication { archive: None, .. }) | None => {
                    verify_segment(&store, &entry.partition, segment).await
                }
            };
            let verified = match verification {
                Ok(verified) => verified,
                Err(error)
                    if args.continue_on_corrupt
                        && !matches!(error, RestoreError::Authenticity { .. }) =>
                {
                    skipped.push(SkippedSegment {
                        topic: entry.partition.topic.clone(),
                        partition: entry.partition.partition,
                        segment_id: segment.segment_id,
                        reason: error.to_string(),
                    });
                    skipped_predecessor = true;
                    continue;
                }
                Err(error) => return Err(error),
            };
            if matches!(
                &authenticated,
                Some(SourceAuthentication {
                    archive: Some(_),
                    ..
                })
            ) {
                consumed_authenticated_objects
                    .extend(segment.artifacts().map(|object| object.key.to_string()));
            }
            segments.push(
                materialize::write_segment_after_skip(
                    args,
                    &entry.partition,
                    &verified,
                    &predicates,
                    skipped_predecessor,
                )
                .await?,
            );
            skipped_predecessor = false;
        }
        partitions.push(PartitionReport {
            topic: entry.partition.topic.clone(),
            partition: entry.partition.partition,
            topic_id: entry.partition.topic_id,
            segments,
        });
    }

    let diskless = match diskless_capture {
        Some(capture) => {
            let report =
                diskless::materialize(&store, args, &predicates, &capture, diskless_claims).await?;
            diskless::check_next_offsets(&report, &diskless_next_offsets)?;
            if diskless_claims.is_some() {
                consumed_authenticated_objects.extend(
                    capture
                        .partitions
                        .iter()
                        .filter(|partition| args.selects_topic(&partition.topic))
                        .flat_map(|partition| &partition.ranges)
                        .map(|range| range.object_key.clone()),
                );
            }
            Some(report)
        }
        None => None,
    };
    if let Some(SourceAuthentication { report, .. }) = &mut authenticated {
        report.objects = u64::try_from(consumed_authenticated_objects.len()).unwrap_or(u64::MAX);
    }
    Ok(RestoreReport {
        dry_run: args.dry_run,
        log_dir: args.target.log_dir.clone(),
        cluster_id: format.cluster_id,
        authentication: authenticated.map(|source| source.report),
        diskless,
        metadata: format.metadata,
        partitions,
        skipped,
    })
}

/// Format the target, seeding the restored metadata with each diskless
/// partition's next offset, and return the format outcome and those offsets.
///
/// The restored metadata starts each diskless partition's offsets at its
/// materialized end (#1524), so that end is settled before the target is
/// formatted, and [`diskless::check_next_offsets`] later holds the restore to
/// it.
async fn format_with_next_offsets(
    store: &ArchiveStore,
    args: &RestoreArgs,
    archive: &ArchiveInventory,
    predicates: &Predicates,
    capture: Option<&DisklessWalCapture>,
    claims: Option<&BTreeMap<String, krabka_remote_storage::ObjectEntry>>,
    metadata_snapshot: Option<&[u8]>,
) -> Result<(FormatTargetOutcome, Vec<diskless::DisklessNextOffset>), RestoreError> {
    let next_offsets = match capture {
        Some(capture) => diskless::next_offsets(store, predicates, capture, args, claims).await?,
        None => Vec::new(),
    };
    let format =
        materialize::format_target_with_snapshot(args, archive, metadata_snapshot, &next_offsets)
            .await?;
    Ok((format, next_offsets))
}

/// Refuse a target that already holds entries.
///
/// A restore writes a whole cluster, so it never merges into existing state.
/// An absent path is acceptable, and the formatter creates it.
fn ensure_empty_log_dir(log_dir: &Path) -> Result<(), RestoreError> {
    if !log_dir.exists() {
        return Ok(());
    }
    if fs::read_dir(log_dir)?.next().is_some() {
        return Err(RestoreError::LogDirNotEmpty(log_dir.to_path_buf()));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use {super::*, assert2::check, clap::CommandFactory as _};

    #[test]
    fn the_command_line_is_internally_consistent() { Cli::command().debug_assert(); }

    #[test]
    fn an_absent_target_is_acceptable() {
        let parent = tempfile::tempdir().expect("temp dir");
        check!(ensure_empty_log_dir(&parent.path().join("missing")).is_ok());
    }

    #[test]
    fn an_empty_target_is_acceptable() {
        let target = tempfile::tempdir().expect("temp dir");
        check!(ensure_empty_log_dir(target.path()).is_ok());
    }

    #[test]
    fn a_target_that_holds_anything_is_refused() {
        let target = tempfile::tempdir().expect("temp dir");
        fs::write(target.path().join("meta.properties"), b"version=1\n").expect("write");
        let refused = ensure_empty_log_dir(target.path());
        check!(matches!(refused, Err(RestoreError::LogDirNotEmpty(_))));
        check!(refused.expect_err("refused").exit_code() == EXIT_DIRTY_LOG_DIR);
    }

    #[tokio::test]
    async fn help_renders_and_succeeds() {
        check!(run_from_args(["krabka-restore", "--help"]).await == EXIT_OK);
    }

    #[tokio::test]
    async fn a_malformed_command_line_is_a_bad_argument() {
        check!(run_from_args(["krabka-restore", "--not-a-flag"]).await == EXIT_BAD_ARGUMENTS);
        check!(run_from_args(["krabka-restore"]).await == EXIT_BAD_ARGUMENTS);
    }
}