Merge branch 'main' into flush_tracker

Author: marvin-j97

Diff for: Cargo.toml

@@ -2,7 +2,7 @@
 name = "fjall"
 description = "LSM-based key-value storage engine"
 license = "MIT OR Apache-2.0"
-version = "2.4.4"
+version = "2.5.0"
 edition = "2021"
 rust-version = "1.74.0"
 readme = "README.md"
@@ -28,7 +28,7 @@ bytes = ["lsm-tree/bytes"]
 
 [dependencies]
 byteorder = "1.5.0"
-lsm-tree = { version = "2.4.0", default-features = false }
+lsm-tree = { version = "2.5.0", default-features = false }
 log = "0.4.21"
 std-semaphore = "0.1.0"
 tempfile = "3.10.1"

Diff for: examples/permuterm/src/main.rs

@@ -80,7 +80,7 @@ fn main() -> Result<()> {
 
     for word in WORDS {
         for term in permuterm(word) {
-            db.insert(format!("{term}#{word}"), word).unwrap();
+            db.insert(format!("{term}#{word}"), *word).unwrap();
         }
     }
 

Diff for: examples/tx-blob-cas/src/main.rs

@@ -125,7 +125,7 @@ impl Cas {
             let rc = self.decrement_ref_count(&mut tx, &prev_content_hash)?;
             if rc == 0 {
                 // No more references
-                tx.remove(&self.blobs, &prev_content_hash);
+                tx.remove(&self.blobs, prev_content_hash);
             }
         }
 
@@ -156,7 +156,7 @@ impl Cas {
             let rc = self.decrement_ref_count(&mut tx, &content_hash)?;
             if rc == 0 {
                 // No more references
-                tx.remove(&self.blobs, &content_hash);
+                tx.remove(&self.blobs, content_hash);
             }
 
             tx.commit()?;

Diff for: examples/tx-mpmc-queue/src/main.rs

@@ -63,11 +63,12 @@ fn main() -> fjall::Result<()> {
                     // Something like SingleDelete https://github.com/facebook/rocksdb/wiki/Single-Delete
                     // would be good for this type of workload
                     if let Some((key, _)) = tx.first_key_value(&tasks)? {
-                        tx.remove(&tasks, &key);
+                        let task_id = std::str::from_utf8(&key).unwrap().to_owned();
+
+                        tx.remove(&tasks, key);
 
                         tx.commit()?;
 
-                        let task_id = std::str::from_utf8(&key).unwrap();
                         println!("consumer {idx} completed task {task_id}");
 
                         counter.fetch_add(1, Relaxed);

Diff for: examples/tx-partition-move/src/main.rs

@@ -33,12 +33,13 @@ fn main() -> fjall::Result<()> {
                     // Something like SingleDelete https://github.com/facebook/rocksdb/wiki/Single-Delete
                     // would be good for this type of workload
                     if let Some((key, value)) = tx.first_key_value(&src)? {
-                        tx.remove(&src, &key);
-                        tx.insert(&dst, &key, &value);
+                        let task_id = std::str::from_utf8(&key).unwrap().to_owned();
+
+                        tx.remove(&src, key.clone());
+                        tx.insert(&dst, key, value);
 
                         tx.commit()?;
 
-                        let task_id = std::str::from_utf8(&key).unwrap();
                         println!("consumer {idx} moved {task_id}");
 
                         let ms = rng.gen_range(10..100);

Diff for: examples/tx-ssi-mpmc-queue/src/main.rs

@@ -65,13 +65,14 @@ fn main() -> fjall::Result<()> {
                     // Something like SingleDelete https://github.com/facebook/rocksdb/wiki/Single-Delete
                     // would be good for this type of workload
                     if let Some((key, _)) = tx.first_key_value(&tasks)? {
-                        tx.remove(&tasks, &key);
+                        let task_id = std::str::from_utf8(&key).unwrap().to_owned();
+
+                        tx.remove(&tasks, key);
 
                         if tx.commit()?.is_ok() {
                             counter.fetch_add(1, Relaxed);
                         }
 
-                        let task_id = std::str::from_utf8(&key).unwrap();
                         println!("consumer {idx} completed task {task_id}");
 
                         let ms = rng.gen_range(50..200);

Diff for: examples/tx-ssi-partition-move/src/main.rs

@@ -33,12 +33,13 @@ fn main() -> fjall::Result<()> {
                     // Something like SingleDelete https://github.com/facebook/rocksdb/wiki/Single-Delete
                     // would be good for this type of workload
                     if let Some((key, value)) = tx.first_key_value(&src)? {
-                        tx.remove(&src, &key);
-                        tx.insert(&dst, &key, &value);
+                        let task_id = std::str::from_utf8(&key).unwrap().to_owned();
+
+                        tx.remove(&src, key.clone());
+                        tx.insert(&dst, key, value);
 
                         tx.commit()?.ok();
 
-                        let task_id = std::str::from_utf8(&key).unwrap();
                         println!("consumer {idx} moved {task_id}");
 
                         let ms = rng.gen_range(10..100);

Diff for: src/batch/mod.rs

@@ -6,7 +6,7 @@ pub mod item;
 
 use crate::{Keyspace, PartitionHandle, PersistMode};
 use item::Item;
-use lsm_tree::{AbstractTree, ValueType};
+use lsm_tree::{AbstractTree, UserKey, UserValue, ValueType};
 use std::{collections::HashSet, sync::Arc};
 
 /// Partition key (a.k.a. column family, locality group)
@@ -52,28 +52,20 @@ impl Batch {
     }
 
     /// Inserts a key-value pair into the batch
-    pub fn insert<K: AsRef<[u8]>, V: AsRef<[u8]>>(
+    pub fn insert<K: Into<UserKey>, V: Into<UserValue>>(
         &mut self,
         p: &PartitionHandle,
         key: K,
         value: V,
     ) {
-        self.data.push(Item::new(
-            p.name.clone(),
-            key.as_ref(),
-            value.as_ref(),
-            ValueType::Value,
-        ));
+        self.data
+            .push(Item::new(p.name.clone(), key, value, ValueType::Value));
     }
 
     /// Adds a tombstone marker for a key
-    pub fn remove<K: AsRef<[u8]>>(&mut self, p: &PartitionHandle, key: K) {
-        self.data.push(Item::new(
-            p.name.clone(),
-            key.as_ref(),
-            vec![],
-            ValueType::Tombstone,
-        ));
+    pub fn remove<K: Into<UserKey>>(&mut self, p: &PartitionHandle, key: K) {
+        self.data
+            .push(Item::new(p.name.clone(), key, vec![], ValueType::Tombstone));
     }
 
     /// Commits the batch to the [`Keyspace`] atomically

Diff for: src/compaction/worker.rs

@@ -5,9 +5,14 @@
 use super::manager::CompactionManager;
 use crate::snapshot_tracker::SnapshotTracker;
 use lsm_tree::AbstractTree;
+use std::sync::atomic::AtomicUsize;
 
 /// Runs a single run of compaction.
-pub fn run(compaction_manager: &CompactionManager, snapshot_tracker: &SnapshotTracker) {
+pub fn run(
+    compaction_manager: &CompactionManager,
+    snapshot_tracker: &SnapshotTracker,
+    compaction_counter: &AtomicUsize,
+) {
     let Some(item) = compaction_manager.pop() else {
         return;
     };
@@ -21,10 +26,12 @@ pub fn run(compaction_manager: &CompactionManager, snapshot_tracker: &SnapshotTr
 
     // TODO: loop if there's more work to do
 
+    compaction_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
     if let Err(e) = item
         .tree
         .compact(strategy.inner(), snapshot_tracker.get_seqno_safe_to_gc())
     {
         log::error!("Compaction failed: {e:?}");
     };
+    compaction_counter.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
 }

Diff for: src/flush/worker.rs

@@ -141,10 +141,13 @@ pub fn run(
                         partition.name,
                         created_segments.len()
                     );
-                    flush_tracker.dequeue_tasks(&partition.name, created_segments.len());
 
+                    flush_tracker.dequeue_tasks(&partition.name, created_segments.len());
                     flush_tracker.decrement_buffer_size(memtables_size);
-                    compaction_manager.notify(partition);
+
+                    for _ in 0..parallelism {
+                        compaction_manager.notify(partition.clone());
+                    }
                 }
             }
             Err(e) => {

Diff for: src/keyspace.rs

@@ -68,6 +68,9 @@ pub struct KeyspaceInner {
     /// True if fsync failed
     pub(crate) is_poisoned: Arc<AtomicBool>,
 
+    /// Active compaction conter
+    pub(crate) active_compaction_count: Arc<AtomicUsize>,
+
     #[doc(hidden)]
     pub snapshot_tracker: Arc<SnapshotTracker>,
 }
@@ -178,6 +181,14 @@ impl Keyspace {
         self.flush_tracker.buffer_size()
     }
 
+    /// Returns the number of active compactions currently running.
+    #[doc(hidden)]
+    #[must_use]
+    pub fn active_compactions(&self) -> usize {
+        self.active_compaction_count
+            .load(std::sync::atomic::Ordering::Relaxed)
+    }
+
     /// Returns the amount of journals on disk.
     ///
     /// # Examples
@@ -546,6 +557,7 @@ impl Keyspace {
             active_background_threads: Arc::default(),
             is_poisoned: Arc::default(),
             snapshot_tracker: Arc::default(),
+            active_compaction_count: Arc::default(),
         };
 
         let keyspace = Self(Arc::new(inner));
@@ -769,6 +781,7 @@ impl Keyspace {
             active_background_threads: Arc::default(),
             is_poisoned: Arc::default(),
             snapshot_tracker: Arc::default(),
+            active_compaction_count: Arc::default(),
         };
 
         // NOTE: Lastly, fsync .fjall marker, which contains the version
@@ -848,6 +861,7 @@ impl Keyspace {
         let stop_signal = self.stop_signal.clone();
         let thread_counter = self.active_background_threads.clone();
         let snapshot_tracker = self.snapshot_tracker.clone();
+        let compaction_counter = self.active_compaction_count.clone();
 
         thread_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
 
@@ -858,7 +872,11 @@ impl Keyspace {
                     log::trace!("compaction: waiting for work");
                     compaction_manager.wait_for();
 
-                    crate::compaction::worker::run(&compaction_manager, &snapshot_tracker);
+                    crate::compaction::worker::run(
+                        &compaction_manager,
+                        &snapshot_tracker,
+                        &compaction_counter,
+                    );
                 }
 
                 log::trace!("compaction thread: exiting because keyspace is dropping");

Diff for: src/partition/mod.rs

@@ -835,15 +835,19 @@ impl PartitionHandle {
     /// # Errors
     ///
     /// Will return `Err` if an IO error occurs.
-    pub fn insert<K: AsRef<[u8]>, V: AsRef<[u8]>>(&self, key: K, value: V) -> crate::Result<()> {
+    pub fn insert<K: Into<UserKey>, V: Into<UserValue>>(
+        &self,
+        key: K,
+        value: V,
+    ) -> crate::Result<()> {
         use std::sync::atomic::Ordering;
 
         if self.is_deleted.load(Ordering::Relaxed) {
             return Err(crate::Error::PartitionDeleted);
         }
 
-        let key = key.as_ref();
-        let value = value.as_ref();
+        let key = key.into();
+        let value = value.into();
 
         let mut journal_writer = self.journal.get_writer();
 
@@ -854,7 +858,7 @@ impl PartitionHandle {
             return Err(crate::Error::Poisoned);
         }
 
-        journal_writer.write_raw(&self.name, key, value, lsm_tree::ValueType::Value, seqno)?;
+        journal_writer.write_raw(&self.name, &key, &value, lsm_tree::ValueType::Value, seqno)?;
 
         if !self.config.manual_journal_persist {
             journal_writer
@@ -914,14 +918,14 @@ impl PartitionHandle {
     /// # Errors
     ///
     /// Will return `Err` if an IO error occurs.
-    pub fn remove<K: AsRef<[u8]>>(&self, key: K) -> crate::Result<()> {
+    pub fn remove<K: Into<UserKey>>(&self, key: K) -> crate::Result<()> {
         use std::sync::atomic::Ordering;
 
         if self.is_deleted.load(Ordering::Relaxed) {
             return Err(crate::Error::PartitionDeleted);
         }
 
-        let key = key.as_ref();
+        let key = key.into();
 
         let mut journal_writer = self.journal.get_writer();
 
@@ -932,7 +936,7 @@ impl PartitionHandle {
             return Err(crate::Error::Poisoned);
         }
 
-        journal_writer.write_raw(&self.name, key, &[], lsm_tree::ValueType::Tombstone, seqno)?;
+        journal_writer.write_raw(&self.name, &key, &[], lsm_tree::ValueType::Tombstone, seqno)?;
 
         if !self.config.manual_journal_persist {
             journal_writer

Diff for: src/partition/options.rs

@@ -231,6 +231,7 @@ impl lsm_tree::coding::Decode for CreateOptions {
                     l0_threshold,
                     target_size,
                     level_ratio,
+                    ..Default::default()
                 })
             }
             1 => {

Diff for: src/tx/conflict_manager.rs

@@ -31,18 +31,18 @@ impl ConflictManager {
         }
     }
 
-    pub fn mark_read(&mut self, partition: &PartitionKey, key: &Slice) {
-        self.push_read(partition, Read::Single(key.clone()));
+    pub fn mark_read(&mut self, partition: &PartitionKey, key: Slice) {
+        self.push_read(partition, Read::Single(key));
     }
 
-    pub fn mark_conflict(&mut self, partition: &PartitionKey, key: &[u8]) {
+    pub fn mark_conflict(&mut self, partition: &PartitionKey, key: Slice) {
         if let Some(tbl) = self.conflict_keys.get_mut(partition) {
-            tbl.insert(key.into());
+            tbl.insert(key);
         } else {
             self.conflict_keys
                 .entry(partition.clone())
                 .or_default()
-                .insert(key.into());
+                .insert(key);
         }
     }