kevy-embedded 7.0.0

//! Keyspace sharding for the embedded store: routing + per-shard persistence
//! bring-up (load / migrate). Each shard is a fully independent
//! `kevy_store::Store` + AOF behind its own lock (shared-nothing), so
//! concurrent access on different shards never contends. `n == 1` is the
//! original single-shard layout (one snapshot + one AOF under the configured
//! filenames, `dump-0.rdb` / `aof-0.aof` by default). `n > 1` keeps per-shard
//! `aof-{i}.aof` + a `shards.meta` recording the count; the first open at
//! `n > 1` re-shards a legacy single AOF into per-shard files.
//!
//! Dir interop with the `kevy` server: a single-shard dir is byte-identical
//! to the server's 1-thread layout, and `n == 1` records `shards.meta` too
//! so neither side needs inference.

use std::io;
#[cfg(feature = "persist")]
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};
#[cfg(feature = "persist")]
use std::time::Instant;

use kevy_hash::KevyHash;
#[cfg(feature = "persist")]
use kevy_persist::reshard::{ShardLayout, commit_reshard, merge_sources, recover_journal};
#[cfg(feature = "persist")]
use kevy_persist::{Aof, Routing, ShardsMeta, layout, layout::infer_files_n};
use kevy_store::Store as Keyspace;

use crate::config::Config;
#[cfg(feature = "persist")]
use crate::metric::KevyMetric;
use crate::metric::OpenReport;
#[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
pub(crate) use crate::shard_tier::{resolve_tier_budget, tier_tick_upkeep};
use crate::store::Inner;

/// Route a key to its shard. `n == 1` short-circuits to 0; power-of-two `n`
/// uses a mask, else modulo. Same `KevyHash` the server's shard router uses.
#[inline]
pub(crate) fn shard_idx(key: &[u8], n: usize) -> usize {
    if n == 1 {
        return 0;
    }
    let h = key.kevy_hash() as usize;
    if n.is_power_of_two() { h & (n - 1) } else { h % n }
}

/// The embedded store's file layout for [`kevy_persist::reshard`]: the
/// standard per-shard names (`dump-{i}.rdb` / `aof-{i}.aof`). `n == 1`
/// coincides with shard 0's names, which are also the legacy single-file
/// names — old single-file dirs load without migration.
#[cfg(feature = "persist")]
struct EmbLayout;

#[cfg(feature = "persist")]
impl ShardLayout for EmbLayout {
    fn snapshot_path(&self, dir: &Path, i: usize, _n: usize) -> PathBuf {
        layout::snapshot_path(dir, i)
    }
    fn aof_path(&self, dir: &Path, i: usize, _n: usize) -> PathBuf {
        layout::aof_path(dir, i)
    }
}

fn fresh_keyspace(config: &Config) -> Keyspace {
    let mut s = Keyspace::new();
    s.set_max_memory(config.maxmemory, config.eviction_policy);
    // no cached clock: the reaper refreshes it only once a tick, so a read
    // could see a key up to a tick past its deadline; a key without a TTL
    // reads no clock either way, one with a TTL reads it fresh
    s
}

/// Build the `n` shard `Inner`s for `config`, loading / migrating persistence.
/// The `bus` lives on shard 0 (pub/sub is process-wide, not sharded); other
/// shards get an idle bus that is never touched.
/// The built shards, the open's report, and the newest catalog frame the
/// restore met.
pub(crate) type Built = (Vec<Arc<RwLock<Inner>>>, OpenReport, Option<kevy_resp::Argv>);

pub(crate) fn build_shards(config: &Config) -> io::Result<Built> {
    let n = config.shards.max(1);
    #[allow(unused_mut)] // mut is the persist path's (load/reshard) need
    let mut stores: Vec<Keyspace> = (0..n).map(|_| fresh_keyspace(config)).collect();

    // Without the `persist` feature the build is always pure in-memory.
    #[cfg(not(feature = "persist"))]
    return Ok((into_inners_mem(stores), OpenReport::default(), None));

    #[cfg(feature = "persist")]
    build_shards_persist(config, n, stores)
}

/// Persistence bring-up half of [`build_shards`]: load / migrate the
/// on-disk layout, then open each shard's live AOF.
#[cfg(feature = "persist")]
fn build_shards_persist(config: &Config, n: usize, mut stores: Vec<Keyspace>) -> io::Result<Built> {
    let Some(dir) = crate::store_wire::disk_dir(config).cloned() else {
        // Pure in-memory: no persistence, no AOF — and no disk for a
        // cold tier: tiering config on a mem-only store is a named
        // refusal, never a silent ignore.
        #[cfg(feature = "tier")]
        if config.tier_budget.is_some() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "tiering requires a disk data dir (with_persist); a memory-only store has no cold tier",
            ));
        }
        let inners = into_inners(stores, (0..n).map(|_| None).collect());
        return Ok((inners, OpenReport::default(), None));
    };
    std::fs::create_dir_all(&dir)?;
    // Complete (or safely discard) a reshard a crash interrupted, before
    // reading the layout — same roll-forward the server runtime does.
    recover_journal(&dir, &EmbLayout)?;
    // Tiering comes up BEFORE the load: replay drives entries
    // into the hot map, so a dataset bigger than the budget must spill
    // inline as it replays — TierState has to exist by then. The vlog
    // wipe-at-open contract precedes the refill, so a reopen never
    // double-counts.
    enable_tiering(config, &dir, &mut stores)?;
    let mut catalog = None;
    let (mut report, walked) = load_or_reshard(&dir, config, n, &mut stores, &mut catalog)?;
    let mut aofs = open_live_aofs(config, &dir, &walked, &mut report)?;
    prepare_aofs(config, &dir, &mut stores, &mut aofs, &mut report, &mut catalog)?;
    Ok((into_inners(stores, aofs), report, catalog))
}

/// Settle each AOF's staging ring and set how it takes appends, sweep the
/// segments no frame claimed, then anchor each AOF's growth-rule baseline.
#[cfg(feature = "persist")]
fn prepare_aofs(
    config: &Config,
    dir: &Path,
    stores: &mut [Keyspace],
    aofs: &mut [Option<Aof>],
    report: &mut OpenReport,
    catalog: &mut Option<kevy_resp::Argv>,
) -> io::Result<()> {
    #[cfg(target_arch = "wasm32")]
    let _ = (dir, &report, &catalog);
    #[cfg(not(target_arch = "wasm32"))]
    for (i, (aof, store)) in aofs.iter_mut().zip(stores.iter_mut()).enumerate() {
        if let Some(aof) = aof {
            crate::shard_restore::open_stage(dir, config, i, store, aof, report, catalog)?;
        }
        // Only now is a sealed segment with no SEGMENTED frame an orphan:
        // the ring may have owed the log that frame, and a segment swept
        // before it lands leaves the frame naming nothing.
        store.sweep_orphan_row_segs();
    }
    // Anchor each AOF's growth-rule baseline to the live image's estimated
    // rewrite size. `Aof::open` can only baseline at the file size, which
    // for a short-lived process re-opening the same directory resets the
    // growth ratio every run — the log then grows without bound while the
    // live keyspace stays tiny (a CLI reusing one store directory paid
    // 100 MB of file for 3 live keys). The estimate serialises every key,
    // so it runs only where the baseline can change a decision.
    let policy = config.rewrite_policy();
    for (aof, store) in aofs.iter_mut().zip(stores.iter()) {
        if let Some(aof) = aof
            && policy.baseline_matters(aof.size_bytes())
        {
            aof.anchor_rewrite_baseline(kevy_persist::estimate_rewrite_size(store));
        }
    }
    Ok(())
}

/// Open each shard's live AOF for append (if persistence is on). The
/// open repairs (quarantines + truncates) any dropped tail replay just
/// tolerated — the quarantine paths land in `report`. `walked[i]` is where shard
/// i's replay stopped, when it dropped nothing.
#[cfg(feature = "persist")]
fn open_live_aofs(
    config: &Config,
    dir: &Path,
    walked: &[Option<u64>],
    report: &mut OpenReport,
) -> io::Result<Vec<Option<Aof>>> {
    let aofs: Vec<Option<Aof>> = if config.aof {
        walked
            .iter()
            .enumerate()
            .map(|(i, &whole)| {
                Aof::open_after_replay(
                    &layout::aof_path(dir, i),
                    config.appendfsync,
                    config.replay_mode,
                    whole,
                )
                .map(Some)
            })
            .collect::<io::Result<_>>()?
    } else {
        walked.iter().map(|_| None).collect()
    };
    for aof in aofs.iter().flatten() {
        if let Some(q) = aof.open_quarantine() {
            report.quarantine_paths.push(q.to_path_buf());
        }
    }
    Ok(aofs)
}

/// Tiering bring-up: per-shard vlog at
/// `<dir>/tier/<i>`, opened (wiped — the vlog is per-boot disposable)
/// BEFORE replay, so in-replay demotion can spill a
/// bigger-than-budget dataset as it loads. The configured budget
/// resolves here (auto/percent probe the memory bound; both refuse by
/// name on failure) and is the whole store's — each shard gets an
/// even slice.
#[cfg(feature = "persist")]
fn enable_tiering(config: &Config, dir: &Path, stores: &mut [Keyspace]) -> io::Result<()> {
    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
    if config.tier_budget.is_some() {
        let per_shard = resolve_tier_budget(config, stores.len())?;
        for (i, store) in stores.iter_mut().enumerate() {
            store.enable_tiering(&dir.join("tier").join(i.to_string()), per_shard)?;
            store.set_tier_max_spill(config.max_spill_value);
        }
    }
    #[cfg(not(all(feature = "tier", not(target_arch = "wasm32"))))]
    let _ = (config, dir, stores);
    Ok(())
}

/// Read the shard layout meta and either load in place (same layout)
/// or re-shard (losslessly) into the configured `n`.
#[cfg(feature = "persist")]
fn load_or_reshard(
    dir: &Path,
    config: &Config,
    n: usize,
    stores: &mut [Keyspace],
    catalog: &mut Option<kevy_resp::Argv>,
) -> io::Result<(OpenReport, Vec<Option<u64>>)> {
    let meta_path = layout::shards_meta_path(dir);
    let prev = ShardsMeta::read(&meta_path);
    // The embedded store always routes by KevyHash; a dir written by a
    // slots-routing server re-shards (losslessly) on first embedded open.
    let same_layout = match prev {
        Some(m) => m.n == n && m.routing == Routing::KevyHash,
        // No meta + n==1: the single-file layout — unless the files say
        // multi-shard (a meta-less pre-1.5 server dir). Loading only
        // shard 0 of those silently dropped (k-1)/k of the keyspace.
        None => n == 1 && infer_files_n(dir) <= 1,
    };

    if same_layout {
        let loaded = load_in_place(dir, config, stores, catalog)?;
        ShardsMeta::new(n, Routing::KevyHash).write(&meta_path)?;
        return Ok(loaded);
    }
    {
        let src_n = prev.map(|m| m.n).or_else(|| {
            let k = infer_files_n(dir);
            (k > 1).then_some(k)
        });
        // The commit also records the new layout — including n == 1: a
        // stale meta from a larger prior n would otherwise trigger a second
        // re-shard next open, whose sources were already renamed to
        // `.premigration` (the shrink-to-one open would come up empty).
        Ok((reshard(dir, config, n, src_n, stores, catalog)?, vec![None; n]))
    }
}

/// Same-layout load: each shard reads its own snapshot + AOF directly.
#[cfg(feature = "persist")]
fn load_in_place(
    dir: &Path,
    config: &Config,
    stores: &mut [Keyspace],
    catalog: &mut Option<kevy_resp::Argv>,
) -> io::Result<(OpenReport, Vec<Option<u64>>)> {
    let mut report = OpenReport::default();
    let start = Instant::now();
    let mut walked = Vec::with_capacity(stores.len());
    for (i, store) in stores.iter_mut().enumerate() {
        let restored =
            crate::shard_restore::restore_one_shard(dir, config, i, store, &mut report, catalog);
        walked.push(restored?);
    }
    report.elapsed_ms = start.elapsed().as_millis() as u64;
    emit_replay(config, &report);
    Ok((report, walked))
}

/// Re-shard: load every source file into one temp keyspace, redistribute
/// each key to its target shard, then hand the crash-safe commit to
/// [`kevy_persist::reshard`] — per-shard snapshots land under `.reshard`
/// temp names, a durable journal marks the commit point, and only then are
/// the sources backed up (`.premigration.<nanos>`) and the temps finalized.
/// A crash at any point either leaves the old layout intact or is rolled
/// forward by `build_shards`' recovery on the next open. Each shard's fresh
/// Before the layout changes, each old shard's staging ring hands what it
/// owes to that shard's AOF, which the merge then reads; the ring goes.
#[cfg(all(feature = "persist", not(target_arch = "wasm32")))]
fn settle_stages(dir: &Path, config: &Config, src_n: usize) -> io::Result<()> {
    for i in 0..src_n {
        let stage = layout::stage_path(dir, i);
        if stage.exists() {
            let mut aof = Aof::open(&layout::aof_path(dir, i), config.appendfsync)?;
            aof.settle_stage(&stage, |_| {})?;
        }
    }
    Ok(())
}

/// AOF opens after this returns; the snapshot is the full migrated state.
#[cfg(feature = "persist")]
fn reshard(
    dir: &Path,
    config: &Config,
    n: usize,
    prev_n: Option<usize>,
    stores: &mut [Keyspace],
    catalog: &mut Option<kevy_resp::Argv>,
) -> io::Result<OpenReport> {
    let lay = EmbLayout;
    // Source layout: prior shard files, or a legacy single AOF/snapshot.
    let src_n = prev_n.unwrap_or(1);
    #[cfg(not(target_arch = "wasm32"))]
    settle_stages(dir, config, src_n)?;
    let (temp, report) = merge_into_temp(dir, config, src_n, catalog)?;
    redistribute(&temp, n, stores);
    let target = ShardsMeta::new(n, Routing::KevyHash);
    commit_reshard(dir, src_n, target, stores, catalog.as_ref(), &lay)?;
    // The merge scratch vlog is dead once the temp keyspace is gone.
    // The attribute rides the same cfg as the code: a module-level one
    // is unfulfilled in every build where this block is compiled out,
    // and `expect` reports that as an error — correctly.
    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
    #[expect(clippy::let_underscore_must_use, reason = "the scratch dir is already dead")]
    if config.tier_budget.is_some() {
        drop(temp);
        let _ = std::fs::remove_dir_all(dir.join("tier").join(".reshard-merge"));
    }
    Ok(report)
}

/// Merge every source file into one temp keyspace and report the replay
/// metrics. B11 on the migration path: a merged dataset bigger than the
/// tier budget must not OOM either — the temp keyspace tiers into a
/// scratch vlog dir (wiped by `reshard` after redistribution) and the
/// merge demotes inline exactly like boot replay does.
#[cfg(feature = "persist")]
fn merge_into_temp(
    dir: &Path,
    config: &Config,
    src_n: usize,
    catalog: &mut Option<kevy_resp::Argv>,
) -> io::Result<(Keyspace, OpenReport)> {
    let lay = EmbLayout;
    let mut temp = fresh_keyspace(config);
    #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
    if config.tier_budget.is_some() {
        // The merge temp holds the whole keyspace — full store budget.
        let budget = resolve_tier_budget(config, 1)?;
        temp.enable_tiering(&dir.join("tier").join(".reshard-merge"), budget)?;
        temp.set_tier_max_spill(config.max_spill_value);
    }
    let mut total_cmds = 0u64;
    let start = Instant::now();
    merge_sources(dir, src_n, &lay, &mut temp, |store, args| {
        if crate::shard_restore::is_catalog(&args) {
            crate::shard_restore::keep_newest(catalog, args);
            return;
        }
        total_cmds += 1;
        crate::replay::apply(store, &args);
        if total_cmds.is_multiple_of(kevy_persist::REPLAY_DEMOTE_INTERVAL) {
            store.demote_to_watermark();
        }
    })?;
    // Replay-metric byte count: the source AOFs (sizes read before the
    // commit renames them away). Per-file drop/corrupt detail is not
    // plumbed through the merge path — a migrating open reports totals.
    let total_bytes = (0..src_n)
        .map(|i| lay.aof_path(dir, i, src_n))
        .filter_map(|p| std::fs::metadata(p).ok())
        .map(|m| m.len())
        .sum();
    let report = OpenReport {
        replayed_commands: total_cmds,
        replayed_bytes: total_bytes,
        elapsed_ms: start.elapsed().as_millis() as u64,
        ..OpenReport::default()
    };
    emit_replay(config, &report);
    Ok((temp, report))
}

/// Redistribute the merged keyspace into the target shards. A cold
/// stub names the TEMP keyspace's vlog — a foreign log the target
/// cannot read — so the source side materializes before shipping
/// (`load_value`'s Cold arm is unreachable by this contract); the
/// target demotes inline to stay under its own budget (B11).
#[cfg(feature = "persist")]
fn redistribute(temp: &Keyspace, n: usize, stores: &mut [Keyspace]) {
    temp.snapshot_each(|key, value, ttl_ms| {
        let hot;
        let value = match temp.materialize_cold(key, value) {
            Some(v) => {
                hot = v;
                &hot
            }
            None => value,
        };
        let target = &mut stores[shard_idx(key, n)];
        target.load_value(key, value, ttl_ms);
        target.try_demote_after_write();
    });
}

#[cfg(feature = "persist")]
fn emit_replay(config: &Config, report: &OpenReport) {
    if let Some(sink) = &config.metric_sink {
        sink.emit(KevyMetric::Replay {
            commands: report.replayed_commands,
            bytes: report.replayed_bytes + report.dropped_bytes,
            elapsed_ms: report.elapsed_ms,
            dropped_bytes: report.dropped_bytes,
            corrupt: report.corrupt,
        });
    }
}

#[cfg(feature = "persist")]
fn into_inners(stores: Vec<Keyspace>, aofs: Vec<Option<Aof>>) -> Vec<Arc<RwLock<Inner>>> {
    stores
        .into_iter()
        .zip(aofs)
        .map(|(store, aof)| Arc::new(RwLock::new(Inner::new(store, aof))))
        .collect()
}

/// In-memory-only `Inner` build (the whole story without `persist`).
#[cfg(not(feature = "persist"))]
fn into_inners_mem(stores: Vec<Keyspace>) -> Vec<Arc<RwLock<Inner>>> {
    stores.into_iter().map(|store| Arc::new(RwLock::new(Inner::new(store)))).collect()
}