#![doc(html_no_source)] // remove it upon open-source

//! Concrete key-value backends for Lyquor persistence traits.

//!

//! `lyquor-db` implements the `lyquor_api::store` interfaces for both in-memory storage and

//! RocksDB-backed storage. Higher layers use those interfaces for simple state, versioned state,

//! artifact caches, and sequencing metadata; this crate keeps backend transactions and iteration

//! behavior below that shared storage boundary.

use std::collections::HashMap;

use std::path::Path;

use std::sync::Arc;

pub use rocksdb;

use rocksdb::{OptimisticTransactionDB, Options};

use lyquor_api::{

    anyhow::{Context, anyhow},

    log::LogClass,

    parking_lot::RwLock,

    store::{KVStore, KVStoreError, Key, SortedMapping, StoreFuture, Value},

};

use tokio::sync::Mutex;

/// In-memory key-value store used by tests and ephemeral node configurations.

pub struct MemDB(RwLock<HashMap<Bytes, Bytes>>);

impl MemDB {

    /// Create an empty in-memory store.

    pub fn new() -> Self {

        Self(RwLock::new(HashMap::new()))

    }

}

impl Default for MemDB {

    fn default() -> Self {

        Self::new()

    }

}

impl KVStore for MemDB {

    fn atomic_write<'a>(

        &'a self, changes: Box<dyn Iterator<Item = (Key, Option<Value>)> + 'a>,

    ) -> Result<(), KVStoreError> {

        let mut map = self.0.write();

        for (k, v) in 
changes16.4k
 {

            match v {

                Some(v) => {

                    map.insert(k.consolidated(), v);

                }

                None => {

                    map.remove(&k.consolidated());

                }

            }

        }

        Ok(())

    }

    fn contains(&self, key: Key) -> StoreFuture<bool> {

        Box::pin(std::future::ready(self.0.read().contains_key(&key.consolidated())))

    }

    fn get(&self, key: Key) -> StoreFuture<Result<Option<Value>, KVStoreError>> {

        Box::pin(std::future::ready(Ok(self.0.read().get(&key.consolidated()).cloned())))

    }

}

/// RocksDB-backed key-value store implementing Lyquor's shared storage trait.

#[derive(Clone)]

pub struct RocksDB(Arc<rocksdb::OptimisticTransactionDB>);

impl KVStore for RocksDB {

    fn atomic_write<'a>(

        &'a self, changes: Box<dyn Iterator<Item = (Key, Option<Value>)> + 'a>,

    ) -> Result<(), KVStoreError> {

        let mut wb = rocksdb::WriteBatchWithTransaction::default();

        for (k, v) in 
changes17
 {

            match v {

                Some(v) => wb.put(k.consolidated(), v),

                None => wb.delete(k.consolidated()),

            }

        }

        if let Err(
error0
) = self.0.write(wb).context("RocksDB write error.") {

            tracing::error!(

                operation = "write",

                class = %LogClass::NodeStorage,

                error = ?error,

                "RocksDB operation failed"

            );

            return Err(error.into());

        }

        Ok(())

    }

    fn contains(&self, key: Key) -> StoreFuture<bool> {

        let db = self.0.clone();

        let key = key.consolidated();

        Box::pin(async move {

            tokio::task::spawn_blocking(move || db.key_may_exist(key))

                .await

                .unwrap_or(false)

        })

    }

    fn get(&self, key: Key) -> StoreFuture<Result<Option<Value>, KVStoreError>> {

        let db = self.0.clone();

        let key = key.consolidated();

        Box::pin(async move {

            tokio::task::spawn_blocking(move || -> Result<Option<Value>, KVStoreError> {

                match db.get(key).context("RocksDB read error.") {

                    Ok(value) => Ok(value.map(std::convert::Into::into)),

                    Err(error) => {

                        tracing::error!(

                            operation = "read",

                            class = %LogClass::NodeStorage,

                            error = ?error,

                            "RocksDB operation failed"

                        );

                        Err(error.into())

                    }

                }

            })

            .await

            .context("RocksDB read task join error.")
?0

        })

    }

}

impl RocksDB {

    /// Open or create a RocksDB store at `path`.

    pub fn new(path: &Path) -> Result<Self, KVStoreError> {

        let mut opts = Options::default();

        opts.create_if_missing(true);

        // Compiled artifacts can be tens of megabytes. Keep them out of SST data

        // blocks so reading adjacent state keys does not decompress the artifact.

        opts.set_enable_blob_files(true);

        opts.set_min_blob_size(1024 * 1024);

        opts.set_blob_compression_type(rocksdb::DBCompressionType::Snappy);

        opts.set_enable_blob_gc(true);

        let db = match OptimisticTransactionDB::open(&opts, path).context("RocksDB open error") {

            Ok(db) => db,

            Err(error) => {

                tracing::error!(

                    path = %path.display(),

                    operation = "open",

                    class = %LogClass::NodeStorage,

                    error = ?error,

                    "RocksDB operation failed"

                );

                return Err(error.into());

            }

        };

        tracing::info!(path = %
path0
.
display0
(), "RocksDB store opened");

        Ok(Self(Arc::new(db)))

    }

}

use lyquor_api::subkey_builder;

use lyquor_primitives::{Bytes, decode_object, encode_object};

use serde::{Deserialize, Serialize};

use std::hash::Hash;

/// Persistent sorted mapping backed by a Lyquor key-value store.

pub struct SortedMappingStore<

    S: KVStore,

    K: Hash + Eq + Ord + Serialize + for<'a> Deserialize<'a> + Clone,

    V: Serialize + for<'a> Deserialize<'a> + Clone,

> {

    state: Arc<Mutex<SortedMappingState<K, V>>>,

    store: Arc<S>,

}

struct SortedMappingState<

    K: Hash + Eq + Ord + Serialize + for<'a> Deserialize<'a> + Clone,

    V: Serialize + for<'a> Deserialize<'a> + Clone,

> {

    cache: lru::LruCache<K, V>,

    max: Option<K>,

    first_idx: u64,

    next_idx: u64,

}

subkey_builder!(SortedMappingSubkey(

    ([0x00])-value(&[u8]) => Key,

    ([0x01])-index(&u64) => Key,

    ([0x02])-max() => Key,

    ([0x03])-first_idx() => Key,

    ([0x04])-next_idx() => Key,

));

impl<

    S: KVStore,

    K: Hash + Eq + Ord + Serialize + for<'a> Deserialize<'a> + Clone,

    V: Serialize + for<'a> Deserialize<'a> + Clone,

> SortedMappingStore<S, K, V>

{

    fn subkey() -> &'static SortedMappingSubkey {

        static KEY: std::sync::OnceLock<SortedMappingSubkey> = std::sync::OnceLock::new();

        KEY.get_or_init(|| 
SortedMappingSubkey::new10
(
Value::new10
().
into10
()))

    }

    /// Open a sorted mapping view over `store` with an in-memory lookup cache.

    pub async fn new(store: S, cache_size: std::num::NonZeroUsize) -> Result<Self, KVStoreError> {

        let 
max14
: 
Option<K>14
 = match store.get(Self::subkey().max()).await
?0
 {

            Some(raw) => Some(decode_object(&raw).ok_or_else(|| 
anyhow!1
("failed to decode sorted mapping max key"))
?1
),

            None => None,

        };

        let first_idx: u64 = match store.get(Self::subkey().first_idx()).await
?0
 {

            Some(raw) => decode_object(&raw).ok_or_else(|| anyhow!("failed to decode sorted mapping first index"))?,

            None => 0,

        };

        let next_idx: u64 = match store.get(Self::subkey().next_idx()).await
?0
 {

            Some(raw) => decode_object(&raw).ok_or_else(|| 
anyhow!0
("failed to decode sorted mapping next index"))
?0
,

            None => 0,

        };

        Ok(Self {

            state: Arc::new(Mutex::new(SortedMappingState {

                cache: lru::LruCache::new(cache_size),

                max,

                first_idx,

                next_idx,

            })),

            store: Arc::new(store),

        })

    }

}

impl<

    S: KVStore + 'static,

    K: Hash + Eq + Ord + Serialize + for<'a> Deserialize<'a> + Clone + Send + 'static,

    V: Serialize + for<'a> Deserialize<'a> + Clone + Send + 'static,

> SortedMapping<K, V> for SortedMappingStore<S, K, V>

{

    fn max_key(&self) -> StoreFuture<Option<K>> {

        let state = self.state.clone();

        Box::pin(async move { state.lock().await.max.clone() })

    }

    fn append(&self, key: K, value: V) -> StoreFuture<Result<bool, KVStoreError>> {

        let state = self.state.clone();

        let store = self.store.clone();

        Box::pin(async move {

            let mut state = state.lock().await;

            Self::append_new_max(store.as_ref(), &mut state, key, value)

        })

    }

    fn append_or_update_max(&self, key: K, value: V) -> StoreFuture<Result<bool, KVStoreError>> {

        let state = self.state.clone();

        let store = self.store.clone();

        Box::pin(async move {

            let mut state = state.lock().await;

            match state.max.clone() {

                Some(
max2
) if key > ma
x2
 => 
Self::append_new_max2
(
store.as_ref()2
, 
&mut state2
, 
key2
, 
value2
),

                Some(
max2
) if key == ma
x2
 => {

                    let key_raw: Bytes = encode_object(&key).into();

                    store.atomic_write(Box::new(

                        [(Self::subkey().value(&key_raw), Some(encode_object(&value).into()))].into_iter(),

                    ))
?0
;

                    state.cache.push(key.clone(), value.clone());

                    Ok(true)

                }

                Some(_) => Ok(false),

                None => Self::append_new_max(store.as_ref(), &mut state, key, value),

            }

        })

    }

    fn get_le(&self, key: K) -> StoreFuture<Result<Option<V>, KVStoreError>> {

        let state = self.state.clone();

        let store = self.store.clone();

        Box::pin(async move {

            let mut state = state.lock().await;

            let first_idx = state.first_idx;

            let next_idx = state.next_idx;

            let Some(
raw9
) = store.get(Self::subkey().index(&first_idx)).await
?0
 else {

                return Ok(None);

            };

            let mut l_key = decode_object(&raw).ok_or_else(|| 
anyhow!0
("failed to decode sorted mapping index key"))
?0
;

            if key < l_key {

                return Ok(None);

            }

            let mut l = first_idx;

            let mut r = next_idx;

            while r > l + 1 {

                let mid = (l + r) >> 1;

                let Some(raw) = store.get(Self::subkey().index(&mid)).await
?0
 else {

                    return Ok(None);

                };

                let mid_key =

                    decode_object(&raw).ok_or_else(|| 
anyhow!0
("failed to decode sorted mapping index key"))
?0
;

                if key < mid_key {

                    r = mid;

                } else {

                    l = mid;

                    l_key = mid_key;

                }

            }

            Self::get_locked(store, &mut state, l_key).await

        })

    }

    fn get(&self, key: K) -> StoreFuture<Result<Option<V>, KVStoreError>> {

        let state = self.state.clone();

        let store = self.store.clone();

        Box::pin(async move {

            let mut state = state.lock().await;

            Self::get_locked(store, &mut state, key).await

        })

    }

}

impl<

    S: KVStore + 'static,

    K: Hash + Eq + Ord + Serialize + for<'a> Deserialize<'a> + Clone + Send + 'static,

    V: Serialize + for<'a> Deserialize<'a> + Clone + Send + 'static,

> SortedMappingStore<S, K, V>

{

    fn append_new_max(store: &S, state: &mut SortedMappingState<K, V>, key: K, value: V) -> Result<bool, KVStoreError> {

        if state.max.clone().is_some_and(|max| 
key3
 <= 
max3
) {

            return Ok(false);

        }

        let idx = state.next_idx;

        let next_idx = idx + 1;

        let key_raw: Bytes = encode_object(&key).into();

        let mapping = [

            (Self::subkey().value(&key_raw), Some(encode_object(&value).into())),

            (Self::subkey().index(&idx), Some(key_raw.clone())),

            (Self::subkey().max(), Some(key_raw)),

            (Self::subkey().next_idx(), Some(encode_object(&next_idx).into())),

        ]

        .into_iter();

        store.atomic_write(Box::new(mapping))
?0
;

        state.next_idx = next_idx;

        state.cache.push(key.clone(), value);

        state.max = Some(key);

        Ok(true)

    }

    async fn get_locked(

        store: Arc<S>, state: &mut SortedMappingState<K, V>, key: K,

    ) -> Result<Option<V>, KVStoreError> {

        if let Some(
value11
) = state.cache.get(&key).cloned() {

            return Ok(Some(value));

        }

        let Some(
raw1
) = store.get(Self::subkey().value(&encode_object(&key))).await
?0
 else {

            return Ok(None);

        };

        let 
value0
: 
V0
 = decode_object(&raw).ok_or_else(|| anyhow!("failed to decode sorted mapping value"))?;

        state.cache.push(key.clone(), value.clone());

        Ok(Some(value))

    }

}

#[cfg(test)]

mod tests {

    use std::sync::Arc;

    use super::SortedMappingStore;

    use crate::{MemDB, RocksDB};

    use lyquor_api::store::{KVStore, SortedMapping};

    use lyquor_primitives::{Bytes, encode_object};

    use lyquor_test::test;

    #[test]

    fn small_state_reads_do_not_load_neighboring_artifacts() {

        use rocksdb::{PerfContext, PerfMetric, PerfStatsLevel, perf::set_perf_stats};

        let dir = tempfile::tempdir().unwrap();

        let artifact = vec![42; 4 * 1024 * 1024];

        {

            let db = RocksDB::new(dir.path()).unwrap();

            db.atomic_write(Box::new(

                [

                    (b"state:a".as_slice().into(), Some(b"metadata".to_vec().into())),

                    (b"state:z".as_slice().into(), Some(artifact.clone().into())),

                ]

                .into_iter(),

            ))

            .unwrap();

            db.0.flush().unwrap();

        }

        let db = RocksDB::new(dir.path()).unwrap();

        set_perf_stats(PerfStatsLevel::EnableCount);

        let mut perf = PerfContext::default();

        for (key, expected) in [(b"state:a", Some(b"metadata".to_vec())), (b"state:b", None)] {

            perf.reset();

            assert_eq!(db.0.get(key).unwrap(), expected);

            assert!(

                perf.metric(PerfMetric::BlockReadByte) < 64 * 1024,

                "small state lookup loaded an artifact-sized block: {}",

                perf.report(true)

            );

        }

        set_perf_stats(PerfStatsLevel::Disable);

        assert_eq!(db.0.get(b"state:z").unwrap().as_deref(), Some(artifact.as_slice()));

    }

    #[test(tokio::test)]

    async fn rocksdb_preserves_inline_and_blob_values_across_restarts() {

        let dir = tempfile::tempdir().unwrap();

        let original = vec![17; 2 * 1024 * 1024];

        let replacement = vec![23; 3 * 1024 * 1024];

        let key = b"artifact".as_slice();

        {

            let mut opts = rocksdb::Options::default();

            opts.create_if_missing(true);

            let legacy: rocksdb::OptimisticTransactionDB =

                rocksdb::OptimisticTransactionDB::open(&opts, dir.path()).unwrap();

            legacy.put(key, &original).unwrap();

            legacy.flush().unwrap();

        }

        {

            let db = RocksDB::new(dir.path()).unwrap();

            assert_eq!(db.get(key.into()).await.unwrap().unwrap().as_ref(), original);

            db.atomic_write(Box::new([(key.into(), Some(replacement.clone().into()))].into_iter()))

                .unwrap();

            db.0.flush().unwrap();

        }

        {

            let db = RocksDB::new(dir.path()).unwrap();

            assert_eq!(db.get(key.into()).await.unwrap().unwrap().as_ref(), replacement);

            db.atomic_write(Box::new([(key.into(), None)].into_iter())).unwrap();

            db.0.flush().unwrap();

        }

        let db = RocksDB::new(dir.path()).unwrap();

        assert!(db.get(key.into()).await.unwrap().is_none());

    }

    #[test(tokio::test)]

    async fn sorted_mapping_store_updates_existing_max_key() {

        let store = SortedMappingStore::new(MemDB::new(), 4.try_into().unwrap())

            .await

            .unwrap();

        assert!(store.append(7_u64, 10_u64).await.unwrap());

        assert!(!store.append(7, 12).await.unwrap());

        assert_eq!(store.get(7).await.unwrap(), Some(10));

        assert!(store.append_or_update_max(7, 11).await.unwrap());

        assert!(!store.append_or_update_max(6, 12).await.unwrap());

        assert_eq!(store.get(6).await.unwrap(), None);

        assert_eq!(store.max_key().await, Some(7));

        assert_eq!(store.get(7).await.unwrap(), Some(11));

        assert_eq!(store.get_le(7).await.unwrap(), Some(11));

    }

    #[test(tokio::test)]

    async fn sorted_mapping_rejects_corrupt_reads() {

        let db = MemDB::new();

        db.atomic_write(Box::new(

            [(

                SortedMappingStore::<MemDB, u64, u64>::subkey().max(),

                Some(Bytes::new()),

            )]

            .into_iter(),

        ))

        .unwrap();

        assert!(

            SortedMappingStore::<_, u64, u64>::new(db, 4.try_into().unwrap())

                .await

                .is_err()

        );

        let db = Arc::new(MemDB::new());

        let store = SortedMappingStore::new(db.clone(), 4.try_into().unwrap())

            .await

            .unwrap();

        assert!(store.append(7_u64, 10_u64).await.unwrap());

        let key_raw: Bytes = encode_object(&7_u64).into();

        db.atomic_write(Box::new(

            [(

                SortedMappingStore::<Arc<MemDB>, u64, u64>::subkey().value(&key_raw),

                Some(Bytes::new()),

            )]

            .into_iter(),

        ))

        .unwrap();

        let reopened: SortedMappingStore<_, u64, u64> =

            SortedMappingStore::new(db, 4.try_into().unwrap()).await.unwrap();

        assert!(reopened.get(7_u64).await.is_err());

    }

}