sync starter code

Signed-off-by: Alex Chi <iskyzh@gmail.com>

mini-lsm-starter/src/lsm_storage.rs

@@ -1,58 +1,151 @@
-#![allow(unused_variables)] // TODO(you): remove this lint after implementing this mod
-#![allow(dead_code)] // TODO(you): remove this lint after implementing this mod
-
+use std::collections::{BTreeSet, HashMap};
+use std::fs::File;
 use std::ops::Bound;
-use std::path::Path;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use bytes::Bytes;
-use parking_lot::RwLock;
+use parking_lot::{Mutex, RwLock};
 
 use crate::block::Block;
+use crate::compact::{
+    CompactionController, CompactionOptions, LeveledCompactionController, LeveledCompactionOptions,
+    SimpleLeveledCompactionController, SimpleLeveledCompactionOptions, TieredCompactionController,
+};
+use crate::iterators::merge_iterator::MergeIterator;
+use crate::iterators::two_merge_iterator::TwoMergeIterator;
+use crate::iterators::StorageIterator;
 use crate::lsm_iterator::{FusedIterator, LsmIterator};
-use crate::mem_table::MemTable;
-use crate::table::SsTable;
+use crate::manifest::{Manifest, ManifestRecord};
+use crate::mem_table::{map_bound, MemTable};
+use crate::table::{FileObject, SsTable, SsTableBuilder, SsTableIterator};
 
 pub type BlockCache = moka::sync::Cache<(usize, usize), Arc<Block>>;
 
 #[derive(Clone)]
-pub struct LsmStorageInner {
+pub struct LsmStorageState {
     /// The current memtable.
-    memtable: Arc<MemTable>,
-    /// Immutable memTables, from earliest to latest.
-    imm_memtables: Vec<Arc<MemTable>>,
-    /// L0 SsTables, from earliest to latest.
-    l0_sstables: Vec<Arc<SsTable>>,
-    /// L1 - L6 SsTables, sorted by key range.
-    #[allow(dead_code)]
-    levels: Vec<Vec<Arc<SsTable>>>,
-    /// The next SSTable ID.
-    next_sst_id: usize,
+    pub memtable: Arc<MemTable>,
+    /// Immutable memtables, from earliest to latest.
+    pub imm_memtables: Vec<Arc<MemTable>>,
+    /// L0 SSTs, from earliest to latest.
+    pub l0_sstables: Vec<usize>,
+    /// SsTables sorted by key range; L1 - L_max for leveled compaction, or tiers for tiered
+    /// compaction.
+    pub levels: Vec<(usize, Vec<usize>)>,
+    /// SST objects.
+    pub sstables: HashMap<usize, Arc<SsTable>>,
 }
 
-impl LsmStorageInner {
-    fn create() -> Self {
+impl LsmStorageState {
+    fn create(options: &LsmStorageOptions) -> Self {
+        let levels = match &options.compaction_options {
+            CompactionOptions::Leveled(LeveledCompactionOptions { max_levels, .. })
+            | CompactionOptions::Simple(SimpleLeveledCompactionOptions { max_levels, .. }) => (1
+                ..=*max_levels)
+                .map(|level| (level, Vec::new()))
+                .collect::<Vec<_>>(),
+            CompactionOptions::Tiered(_) | CompactionOptions::NoCompaction => Vec::new(),
+        };
         Self {
-            memtable: Arc::new(MemTable::create()),
-            imm_memtables: vec![],
-            l0_sstables: vec![],
-            levels: vec![],
-            next_sst_id: 1,
+            memtable: Arc::new(MemTable::create(0)),
+            imm_memtables: Vec::new(),
+            l0_sstables: Vec::new(),
+            levels,
+            sstables: Default::default(),
+        }
+    }
+}
+
+pub struct LsmStorageOptions {
+    pub block_size: usize,
+    pub target_sst_size: usize,
+    pub num_memtable_limit: usize,
+    pub compaction_options: CompactionOptions,
+    pub enable_wal: bool,
+}
+
+impl LsmStorageOptions {
+    pub fn default_for_week1_test() -> Self {
+        Self {
+            block_size: 4096,
+            target_sst_size: 2 << 20,
+            compaction_options: CompactionOptions::NoCompaction,
+            enable_wal: false,
+            num_memtable_limit: 3,
         }
     }
 }
 
 /// The storage interface of the LSM tree.
-pub struct LsmStorage {
-    inner: Arc<RwLock<Arc<LsmStorageInner>>>,
+pub(crate) struct LsmStorageInner {
+    pub(crate) state: Arc<RwLock<Arc<LsmStorageState>>>,
+    pub(crate) state_lock: Mutex<()>,
+    path: PathBuf,
+    pub(crate) block_cache: Arc<BlockCache>,
+    next_sst_id: AtomicUsize,
+    pub(crate) options: Arc<LsmStorageOptions>,
+    pub(crate) compaction_controller: CompactionController,
+    pub(crate) manifest: Manifest,
 }
 
-impl LsmStorage {
-    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
-        Ok(Self {
-            inner: Arc::new(RwLock::new(Arc::new(LsmStorageInner::create()))),
-        })
+pub struct MiniLsm {
+    pub(crate) inner: Arc<LsmStorageInner>,
+    compaction_notifier: crossbeam_channel::Sender<()>,
+    compaction_thread: Mutex<Option<std::thread::JoinHandle<()>>>,
+}
+
+impl Drop for MiniLsm {
+    fn drop(&mut self) {
+        self.compaction_notifier.send(()).ok();
+    }
+}
+
+impl MiniLsm {
+    pub fn close(&self) -> Result<()> {
+        unimplemented!()
+    }
+
+    pub fn open(path: impl AsRef<Path>, options: LsmStorageOptions) -> Result<Arc<Self>> {
+        unimplemented!()
+    }
+
+    pub fn get(&self, key: &[u8]) -> Result<Option<Bytes>> {
+        self.inner.get(key)
+    }
+
+    pub fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
+        self.inner.put(key, value)
+    }
+
+    pub fn delete(&self, key: &[u8]) -> Result<()> {
+        self.inner.delete(key)
+    }
+
+    pub fn scan(
+        &self,
+        lower: Bound<&[u8]>,
+        upper: Bound<&[u8]>,
+    ) -> Result<FusedIterator<LsmIterator>> {
+        self.inner.scan(lower, upper)
+    }
+
+    pub fn force_flush(&self) -> Result<()> {
+        self.inner.force_freeze_memtable()?;
+        self.inner.force_flush_next_imm_memtable()
+    }
+}
+
+impl LsmStorageInner {
+    pub(crate) fn next_sst_id(&self) -> usize {
+        self.next_sst_id
+            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
+    }
+
+    pub(crate) fn open(path: impl AsRef<Path>, options: LsmStorageOptions) -> Result<Self> {
+        unimplemented!()
     }
 
     /// Get a key from the storage. In day 7, this can be further optimized by using a bloom filter.
@@ -62,26 +155,49 @@ impl LsmStorage {
 
     /// Put a key-value pair into the storage by writing into the current memtable.
     pub fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
-        assert!(!value.is_empty(), "value cannot be empty");
-        assert!(!key.is_empty(), "key cannot be empty");
         unimplemented!()
     }
 
     /// Remove a key from the storage by writing an empty value.
-    pub fn delete(&self, _key: &[u8]) -> Result<()> {
+    pub fn delete(&self, key: &[u8]) -> Result<()> {
         unimplemented!()
     }
 
-    /// Persist data to disk.
-    pub fn sync(&self) -> Result<()> {
+    pub(crate) fn path_of_sst_static(path: impl AsRef<Path>, id: usize) -> PathBuf {
+        path.as_ref().join(format!("{:05}.sst", id))
+    }
+
+    pub(crate) fn path_of_sst(&self, id: usize) -> PathBuf {
+        Self::path_of_sst_static(&self.path, id)
+    }
+
+    pub(crate) fn path_of_wal_static(path: impl AsRef<Path>, id: usize) -> PathBuf {
+        path.as_ref().join(format!("{:05}.wal", id))
+    }
+
+    pub(crate) fn path_of_wal(&self, id: usize) -> PathBuf {
+        Self::path_of_wal_static(&self.path, id)
+    }
+
+    fn sync_dir(&self) -> Result<()> {
+        unimplemented!()
+    }
+
+    /// Force freeze the current memetable to an immutable memtable
+    pub fn force_freeze_memtable(&self) -> Result<()> {
+        unimplemented!()
+    }
+
+    /// Force flush the earliest-created immutable memtable to disk
+    pub fn force_flush_next_imm_memtable(&self) -> Result<()> {
         unimplemented!()
     }
 
     /// Create an iterator over a range of keys.
     pub fn scan(
         &self,
-        _lower: Bound<&[u8]>,
-        _upper: Bound<&[u8]>,
+        lower: Bound<&[u8]>,
+        upper: Bound<&[u8]>,
     ) -> Result<FusedIterator<LsmIterator>> {
         unimplemented!()
     }