kevy_embedded/

store.rs

//! [`Store`] — the embedded entry point. Wraps `kevy_store::Store` with
//! per-shard locks (for cross-thread access), optional AOF auto-logging, an
//! optional background TTL reaper, and an in-process pub/sub bus.

use std::io;
use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};

use kevy_persist::Argv;
use kevy_store::ExpireStats;

use crate::config::Config;
use crate::shard::{build_shards, shard_idx};

pub use crate::store_inner::WeakStore;
pub(crate) use crate::store_inner::{DropGuard, Inner};

/// The keyspace shards (`hash(key) % n`), each a fully independent
/// `kevy_store::Store` + AOF behind its own lock. `n == 1` (the default) is a
/// one-element vec = the original single-lock store.
pub(crate) type Shards = Arc<Vec<Arc<RwLock<Inner>>>>;

/// The embedded keyspace.
///
/// **`Store` is `Clone`** (since v1.1.0). A clone is a cheap `Arc` bump:
/// every clone reaches the same underlying shards + AOF + reaper + pub/sub
/// bus. The reaper thread is joined and each shard's AOF is flushed exactly
/// once, when the **last** clone is dropped.
///
/// ```
/// use kevy_embedded::{Config, Store};
///
/// # fn main() -> std::io::Result<()> {
/// let s = Store::open(Config::default().with_ttl_reaper_manual())?;
/// let s2 = s.clone();
/// std::thread::spawn(move || {
///     s2.set(b"from-thread", b"v").unwrap();
/// }).join().unwrap();
/// assert_eq!(s.get(b"from-thread")?, Some(b"v".to_vec()));
/// # Ok(())
/// # }
/// ```
///
/// Every method takes `&self`. Sharding (see [`Config::with_shards`]) lets a
/// multi-threaded consumer scale across cores; pub/sub is process-wide
/// (handled on shard 0).
#[derive(Clone)]
pub struct Store {
    pub(crate) shards: Shards,
    /// Shared drop guard: signals + joins reaper and flushes AOFs when the
    /// LAST `Store` clone (or `Subscription`) holding a strong ref drops.
    pub(crate) guard: Arc<DropGuard>,
    pub(crate) config: Config,
    /// v2.3 CDC feed handle (read API side); shards carry clones for
    /// the write side. `None` = feed off (or wasm).
    #[cfg(not(target_arch = "wasm32"))]
    pub(crate) feed: Option<std::sync::Arc<Mutex<kevy_replicate::feed::FeedSource>>>,
    /// v2.4 blocking-pop wake channel (always present; writers pay one
    /// Relaxed load while nobody blocks).
    pub(crate) blocker: Arc<crate::ops_blocking::Blocker>,
    /// v2.5 index registry (catalog + version).
    pub(crate) indexes: Arc<crate::ops_index::IndexReg>,
    /// v2.6 view registry.
    pub(crate) views: Arc<crate::ops_view::ViewReg>,
}

impl Store {
    /// Open an embedded keyspace per `config`.
    ///
    /// - Pure in-memory when `config.data_dir` is `None`.
    /// - With persistence: each shard loads its snapshot then replays its AOF
    ///   (`config.shards > 1` re-shards a legacy single AOF on first open).
    /// - Spawns a background TTL reaper thread when
    ///   `config.ttl_reaper == Background` (the default).
    /// - When `config.replica_upstream = Some("host:port")`, spawns a
    ///   background thread that streams replication frames from the
    ///   named primary and applies them to this store; local writes are
    ///   rejected with `READONLY` (see [`Self::open_replica`]).
    pub fn open(config: Config) -> io::Result<Self> {
        Self::open_inner(config)
    }

    /// Answer one RESP request against this store using the SAME
    /// read-only verb whitelist the embedded RESP listener serves
    /// (`Config::with_resp_listener`). The reply is appended to `out`
    /// as raw RESP bytes; write verbs answer `-ERR` like the listener
    /// does. This is the programmatic face of the listener — tooling
    /// (e.g. `kevy-cli --embed`) inspects a store without a socket.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn dispatch_readonly(&self, argv: &[Vec<u8>], out: &mut Vec<u8>) {
        crate::listener::verbs_dispatch(self, argv, out);
    }

    fn open_inner(config: Config) -> io::Result<Self> {
        let shards: Shards = Arc::new(build_shards(&config)?);
        let (reaper_stop, reaper_join) = crate::reaper::spawn_reaper(&config, &shards)?;
        #[cfg(not(target_arch = "wasm32"))]
        let replica_runner = crate::replica_glue::spawn_replica_runner(&config, &shards);
        #[cfg(not(target_arch = "wasm32"))]
        let replica_source = spawn_writer_source(&config, &shards)?;
        #[cfg(not(target_arch = "wasm32"))]
        let feed = Store::feed_open(&config)?;
        #[cfg(not(target_arch = "wasm32"))]
        if let Some(f) = &feed {
            for shard in shards.iter() {
                let mut g = lock_write(shard);
                g.feed = Some(f.clone());
            }
        }
        let (blocker, indexes, views) = wire_registries(&shards);
        let guard = Arc::new(DropGuard {
            reaper_stop,
            reaper_join: Mutex::new(reaper_join),
            shards_for_flush: shards.clone(),
            #[cfg(not(target_arch = "wasm32"))]
            replica_runner,
            #[cfg(not(target_arch = "wasm32"))]
            feed_close: match (&feed, &config.data_dir) {
                (Some(f), Some(d)) => Some((f.clone(), d.clone())),
                _ => None,
            },
            #[cfg(not(target_arch = "wasm32"))]
            replica_source,
        });
        let store = Store {
            shards,
            guard,
            config,
            #[cfg(not(target_arch = "wasm32"))]
            feed,
            blocker,
            indexes,
            views,
        };
        store.idx_boot();
        store.view_boot();
        #[cfg(not(target_arch = "wasm32"))]
        if let Some(addr) = store.config.resp_listener {
            crate::listener::spawn(addr, store.downgrade())?;
        }
        Ok(store)
    }

    /// Convenience constructor for an embed-as-read-replica store
    /// streaming writes from `upstream` (`"host:port"` of a kevy
    /// server's replication listener).
    ///
    /// The replica:
    /// - has its local AOF force-disabled (the upstream stream is the
    ///   source of truth; replica AOF would diverge and double-apply
    ///   on restart);
    /// - rejects every local write with a `READONLY` `io::Error`
    ///   (you can still call read APIs concurrently);
    /// - reconnects with exponential backoff on disconnect, resuming
    ///   from the last applied offset;
    /// - gets a process-unique `replica_id` so an open / drop / reopen
    ///   cycle within the primary's reconnect window does not look like
    ///   the same slot from the primary's POV (which would evict
    ///   backlog frames the new embed still needs from offset 0).
    ///   Override via [`Config::with_replica_id`] when you specifically
    ///   want the slot to be re-claimed across restarts.
    ///
    /// For full builder control (custom replica id, backoff bounds,
    /// snapshot dir, etc.) use [`Self::open`] with
    /// [`Config::with_replica_upstream`] + the related setters
    /// instead.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn open_replica(upstream: impl Into<String>) -> io::Result<Self> {
        let cfg = Config::default()
            .without_aof()
            .with_replica_id(crate::replica_glue::fresh_replica_id())
            .with_replica_upstream(upstream);
        Self::open(cfg)
    }

    /// `true` when this store was opened against a replication
    /// upstream — local writes are rejected with `READONLY`.
    pub fn is_replica(&self) -> bool {
        self.config.replica_upstream.is_some()
    }

    /// Retarget this replica at a new primary URL (`host:port`). The
    /// runner picks up the change on its next connect — which is
    /// forced now by `shutdown`ing the current socket clone, so the
    /// retarget lands within `Config::replica_reconnect_min` (default
    /// 100 ms) of this call.
    ///
    /// Returns `Err` with `ErrorKind::InvalidInput` when this store is
    /// not a replica (no upstream was configured at open). Application
    /// code typically drives this from a `kevy-elect` failover signal —
    /// see [`docs/cluster.md`](https://github.com/goliajp/kevy/blob/develop/docs/cluster.md).
    /// `kevy-embedded` itself stays elect-protocol-agnostic; the
    /// integration glue lives in the application.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn set_replica_upstream(&self, new_upstream: impl Into<String>) -> io::Result<()> {
        if !self.is_replica() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "set_replica_upstream called on a non-replica store",
            ));
        }
        let Some(runner) = self.guard.replica_runner.as_ref() else {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "replica runner is not active (open was racy?)",
            ));
        };
        runner.set_upstream(new_upstream.into());
        Ok(())
    }

    /// The active config (a clone — modifying it has no effect on the
    /// running store). Useful for introspection / `INFO`-style telemetry.
    pub fn config(&self) -> &Config {
        &self.config
    }

    // ---- escape hatches -------------------------------------------------

    /// Run `f` against the underlying `kevy_store::Store` under its lock. Use
    /// for direct access to methods this crate hasn't wrapped. The closure can
    /// mutate, but *does not auto-log to the AOF* — call [`Self::log`] yourself
    /// if the mutation must survive a crash.
    ///
    /// **Sharded stores:** this targets shard 0 only. Use [`Self::with_key`]
    /// to reach the shard owning a specific key.
    pub fn with<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut kevy_store::Store) -> R,
    {
        let mut g = self.lock();
        f(&mut g.store)
    }

    /// Like [`Self::with`] but targets the shard that owns `key`.
    pub fn with_key<F, R>(&self, key: &[u8], f: F) -> R
    where
        F: FnOnce(&mut kevy_store::Store) -> R,
    {
        let mut g = self.wshard(key);
        f(&mut g.store)
    }

    /// `KEYS` / `SCAN`-glob across **every shard** — the cross-shard
    /// replacement for `with(|s| s.collect_keys(pat, lim))`, which only sees
    /// shard 0 once sharding is on. Behaves identically to `with(...)` when
    /// `shard_count() == 1`. `limit` bounds the *total* returned across shards.
    /// Takes a read lock per shard (concurrent-safe).
    pub fn collect_keys(&self, pattern: Option<&[u8]>, limit: Option<usize>) -> Vec<Vec<u8>> {
        let mut out = Vec::new();
        for shard in self.shards.iter() {
            if limit.is_some_and(|l| out.len() >= l) {
                break;
            }
            let remaining = limit.map(|l| l - out.len());
            out.extend(lock_read(shard).store.collect_keys(pattern, remaining));
        }
        out
    }

    /// Run `f` against **each shard's** underlying `kevy_store::Store` (in
    /// shard-index order) — the cross-shard escape hatch. The caller assembles
    /// the merged result. Pairs with [`Self::shard_count`]. For a single key,
    /// prefer [`Self::with_key`]; for a glob scan, prefer [`Self::collect_keys`].
    pub fn for_each_shard<F: FnMut(&mut kevy_store::Store)>(&self, mut f: F) {
        for shard in self.shards.iter() {
            f(&mut lock_write(shard).store);
        }
    }

    /// Number of keyspace shards (`== Config::shards`).
    #[inline]
    pub fn shard_count(&self) -> usize {
        self.shards.len()
    }

    /// Append a raw RESP-frame argument list to the shard owning its key's
    /// AOF. No-op when persistence is disabled.
    pub fn log(&self, parts: &[&[u8]]) -> io::Result<()> {
        let mut g = match parts.get(1) {
            Some(key) => self.wshard(key),
            None => self.lock(),
        };
        if let Some(aof) = &mut g.aof {
            let argv = Argv::from(parts.iter().map(|p| p.to_vec()).collect::<Vec<_>>());
            aof.append(&argv)?;
        }
        Ok(())
    }

    // ---- maintenance ----------------------------------------------------

    /// Run one TTL-reaper tick across every shard. Required call cadence in
    /// `Manual` mode (~10×/s to match Redis `hz=10`). Returns the summed stats.
    pub fn tick(&self) -> ExpireStats {
        let mut total = ExpireStats::default();
        for shard in self.shards.iter() {
            let stats = {
                let mut g = lock_write(shard);
                g.store.tick_expire(self.config.reaper_samples, self.config.reaper_max_rounds)
            };
            total.sampled += stats.sampled;
            total.expired += stats.expired;
            // Auto-rewrite rides the caller-driven tick in Manual mode; the
            // non-blocking path releases the lock for the disk spill.
            crate::reaper::concurrent_auto_rewrite(
                shard,
                self.config.auto_aof_rewrite_pct,
                self.config.auto_aof_rewrite_min_size,
                self.config.metric_sink.as_ref(),
            );
        }
        total
    }

    // Durability methods (`rewrite_aof`, `save_snapshot`) live in
    // `crate::store_persist` to keep this file under the 500-LOC
    // project ceiling.
    // Data-type methods live in `crate::ops` / `crate::info`.

    /// Crate-internal: clone shard 0's handle for a `Subscription`'s bus.
    pub(crate) fn inner_handle(&self) -> Arc<RwLock<Inner>> {
        self.shards[0].clone()
    }

    /// Crate-internal: clone the shared `Arc<DropGuard>`.
    pub(crate) fn guard_handle(&self) -> Arc<DropGuard> {
        self.guard.clone()
    }

    fn shard_for(&self, key: &[u8]) -> &Arc<RwLock<Inner>> {
        &self.shards[shard_idx(key, self.shards.len())]
    }

    /// Write-lock the shard owning `key`.
    pub(crate) fn wshard(&self, key: &[u8]) -> RwLockWriteGuard<'_, Inner> {
        lock_write(self.shard_for(key))
    }

    /// Read-lock the shard owning `key` (GET fast path — concurrent readers
    /// across shards run in parallel).
    pub(crate) fn rshard(&self, key: &[u8]) -> RwLockReadGuard<'_, Inner> {
        lock_read(self.shard_for(key))
    }

    /// Write-lock shard 0 — pub/sub bus + keyless escape hatches.
    pub(crate) fn lock(&self) -> RwLockWriteGuard<'_, Inner> {
        lock_write(&self.shards[0])
    }

    /// Run `f` over every shard's write guard, summing a `usize` (DBSIZE etc.).
    pub(crate) fn sum_shards<F: Fn(&mut Inner) -> usize>(&self, f: F) -> usize {
        self.shards.iter().map(|s| f(&mut lock_write(s))).sum()
    }

    /// Run `f` over every shard's write guard, summing a `u64`.
    pub(crate) fn sum_shards_u64<F: Fn(&mut Inner) -> u64>(&self, f: F) -> u64 {
        self.shards.iter().map(|s| f(&mut lock_write(s))).sum()
    }

    /// Run a fallible `f` over every shard (mutating, e.g. FLUSHALL).
    pub(crate) fn try_for_each_shard<F: FnMut(&mut Inner) -> io::Result<()>>(
        &self,
        mut f: F,
    ) -> io::Result<()> {
        for s in self.shards.iter() {
            f(&mut lock_write(s))?;
        }
        Ok(())
    }
}


pub(crate) use crate::store_glue::{commit_write, lock_read, lock_write, store_err};

/// Spawn the embed-as-writer replication source (v3.2) when
/// `Config::embed_writer_listen_addr` is set, wiring the shared source
/// into every shard's `Inner` so `commit_write` pushes mutations into
/// the backlog inline (done once at open under the shard's write lock;
/// reads of `Inner::writer_source` afterwards are uncontended).
///
/// The snapshot provider freezes every shard's COW view under the
/// source lock (so ack_offset and the frozen keyspace are one point in
/// time — writes between the two would replay twice otherwise), then
/// serializes outside the locks via the persist writer.
#[cfg(not(target_arch = "wasm32"))]
fn spawn_writer_source(
    config: &Config,
    shards: &Shards,
) -> io::Result<Option<crate::replica_source::ReplicaSource>> {
    let Some(addr) = config.embed_writer_listen_addr.as_ref() else {
        return Ok(None);
    };
    let shards_for_snap: Shards = Arc::clone(shards);
    let snapshot: crate::replica_source::SnapshotProvider =
        Arc::new(move || crate::replica_source::freeze_and_serialize(&shards_for_snap));
    let rs = crate::replica_source::ReplicaSource::spawn(
        addr,
        config.embed_writer_backlog_bytes,
        snapshot,
    )?;
    let shared = rs.shared_source();
    for shard in shards.iter() {
        let mut g = lock_write(shard);
        g.writer_source = Some(shared.clone());
    }
    Ok(Some(rs))
}

/// Create the store-level registries (blocking-pop waker, index +
/// view catalogs) and hand every shard's `Inner` a clone.
#[allow(clippy::type_complexity)]
fn wire_registries(
    shards: &Shards,
) -> (
    Arc<crate::ops_blocking::Blocker>,
    Arc<crate::ops_index::IndexReg>,
    Arc<crate::ops_view::ViewReg>,
) {
    let blocker = Arc::new(crate::ops_blocking::Blocker::new());
    let indexes = Arc::new(crate::ops_index::IndexReg::default());
    let views = Arc::new(crate::ops_view::ViewReg::default());
    for shard in shards.iter() {
        let mut g = lock_write(shard);
        g.blocker = Some(blocker.clone());
        g.idx_reg = Some(indexes.clone());
        g.view_reg = Some(views.clone());
    }
    (blocker, indexes, views)
}


#[cfg(test)]
#[path = "store_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "store_tests_shard.rs"]
mod tests_shard;
#[cfg(test)]
#[path = "store_tests_p2.rs"]
mod tests_p2;
#[cfg(test)]
#[path = "store_tests_p3.rs"]
mod tests_p3;
#[cfg(test)]
#[path = "store_tests_bitmap.rs"]
mod tests_bitmap;
#[cfg(test)]
#[path = "store_tests_bonus.rs"]
mod tests_bonus;
#[cfg(test)]
#[path = "store_tests_scan.rs"]
mod tests_scan;
#[cfg(test)]
#[path = "store_tests_atomic.rs"]
mod tests_atomic;
#[cfg(test)]
#[path = "store_tests_more.rs"]
mod tests_more;
#[cfg(test)]
#[path = "store_tests_keyspace.rs"]
mod tests_keyspace;
#[cfg(test)]
#[path = "store_tests_atomic_all.rs"]
mod tests_atomic_all;
#[cfg(test)]
#[path = "store_tests_replay_all.rs"]
mod tests_replay_all;
#[cfg(test)]
#[path = "store_tests_op_table.rs"]
mod tests_op_table;