Signed-off-by: Alex Chi Z <iskyzh@gmail.com>
This commit is contained in:
Alex Chi Z
2024-01-30 14:48:03 +08:00
parent 9eab75ec1a
commit acc3c959aa
9 changed files with 345 additions and 24 deletions

View File

@@ -243,15 +243,11 @@ impl MiniLsm {
}))
}
pub fn new_txn(&self) -> Result<Arc<Transaction>> {
self.inner.new_txn()
}
pub fn get(&self, key: &[u8]) -> Result<Option<Bytes>> {
self.inner.get(key)
}
pub fn write_batch<T: AsRef<[u8]>>(&self, batch: &[WriteBatchRecord<T>]) -> Result<u64> {
pub fn write_batch<T: AsRef<[u8]>>(&self, batch: &[WriteBatchRecord<T>]) -> Result<()> {
self.inner.write_batch(batch)
}
@@ -267,6 +263,10 @@ impl MiniLsm {
self.inner.sync()
}
pub fn new_txn(&self) -> Result<Arc<Transaction>> {
self.inner.new_txn()
}
pub fn scan(&self, lower: Bound<&[u8]>, upper: Bound<&[u8]>) -> Result<TxnIterator> {
self.inner.scan(lower, upper)
}
@@ -302,8 +302,6 @@ impl LsmStorageInner {
self.manifest.as_ref().unwrap()
}
/// Start the storage engine by either loading an existing directory or creating a new one if the directory does
/// not exist.
/// Start the storage engine by either loading an existing directory or creating a new one if the directory does
/// not exist.
pub(crate) fn open(path: impl AsRef<Path>, options: LsmStorageOptions) -> Result<Self> {
@@ -443,9 +441,6 @@ impl LsmStorageInner {
self.state.read().memtable.sync_wal()
}
pub fn new_txn(self: &Arc<Self>) -> Result<Arc<Transaction>> {
Ok(self.mvcc().new_txn(self.clone(), self.options.serializable))
}
/// Get a key from the storage. In day 7, this can be further optimized by using a bloom filter.
pub fn get(self: &Arc<Self>, key: &[u8]) -> Result<Option<Bytes>> {
let txn = self.mvcc().new_txn(self.clone(), self.options.serializable);
@@ -531,7 +526,7 @@ impl LsmStorageInner {
Ok(None)
}
pub fn write_batch<T: AsRef<[u8]>>(&self, batch: &[WriteBatchRecord<T>]) -> Result<u64> {
pub fn write_batch_inner<T: AsRef<[u8]>>(&self, batch: &[WriteBatchRecord<T>]) -> Result<u64> {
let _lck = self.mvcc().write_lock.lock();
let ts = self.mvcc().latest_commit_ts() + 1;
for record in batch {
@@ -566,10 +561,33 @@ impl LsmStorageInner {
Ok(ts)
}
pub fn write_batch<T: AsRef<[u8]>>(
self: &Arc<Self>,
batch: &[WriteBatchRecord<T>],
) -> Result<()> {
if !self.options.serializable {
self.write_batch_inner(batch)?;
} else {
let txn = self.mvcc().new_txn(self.clone(), self.options.serializable);
for record in batch {
match record {
WriteBatchRecord::Del(key) => {
txn.delete(key.as_ref());
}
WriteBatchRecord::Put(key, value) => {
txn.put(key.as_ref(), value.as_ref());
}
}
}
txn.commit()?;
}
Ok(())
}
/// Put a key-value pair into the storage by writing into the current memtable.
pub fn put(self: &Arc<Self>, key: &[u8], value: &[u8]) -> Result<()> {
if !self.options.serializable {
self.write_batch(&[WriteBatchRecord::Put(key, value)])?;
self.write_batch_inner(&[WriteBatchRecord::Put(key, value)])?;
} else {
let txn = self.mvcc().new_txn(self.clone(), self.options.serializable);
txn.put(key, value);
@@ -581,7 +599,7 @@ impl LsmStorageInner {
/// Remove a key from the storage by writing an empty value.
pub fn delete(self: &Arc<Self>, key: &[u8]) -> Result<()> {
if !self.options.serializable {
self.write_batch(&[WriteBatchRecord::Del(key)])?;
self.write_batch_inner(&[WriteBatchRecord::Del(key)])?;
} else {
let txn = self.mvcc().new_txn(self.clone(), self.options.serializable);
txn.delete(key);
@@ -720,6 +738,10 @@ impl LsmStorageInner {
Ok(())
}
pub fn new_txn(self: &Arc<Self>) -> Result<Arc<Transaction>> {
Ok(self.mvcc().new_txn(self.clone(), self.options.serializable))
}
/// Create an iterator over a range of keys.
pub fn scan<'a>(
self: &'a Arc<Self>,

View File

@@ -1,3 +1,6 @@
#![allow(unused_variables)] // TODO(you): remove this lint after implementing this mod
#![allow(dead_code)] // TODO(you): remove this lint after implementing this mod
pub mod txn;
pub mod watermark;
@@ -23,6 +26,7 @@ pub(crate) struct CommittedTxnData {
pub(crate) struct LsmMvccInner {
pub(crate) write_lock: Mutex<()>,
pub(crate) commit_lock: Mutex<()>,
pub(crate) ts: Arc<Mutex<(u64, Watermark)>>,
pub(crate) committed_txns: Arc<Mutex<BTreeMap<u64, CommittedTxnData>>>,
}
@@ -31,6 +35,7 @@ impl LsmMvccInner {
pub fn new(initial_ts: u64) -> Self {
Self {
write_lock: Mutex::new(()),
commit_lock: Mutex::new(()),
ts: Arc::new(Mutex::new((initial_ts, Watermark::new()))),
committed_txns: Arc::new(Mutex::new(BTreeMap::new())),
}

View File

@@ -7,7 +7,7 @@ use std::{
},
};
use anyhow::Result;
use anyhow::{bail, Result};
use bytes::Bytes;
use crossbeam_skiplist::{map::Entry, SkipMap};
use ouroboros::self_referencing;
@@ -18,6 +18,7 @@ use crate::{
lsm_iterator::{FusedIterator, LsmIterator},
lsm_storage::{LsmStorageInner, WriteBatchRecord},
mem_table::map_bound,
mvcc::CommittedTxnData,
};
pub struct Transaction {
@@ -34,6 +35,11 @@ impl Transaction {
if self.committed.load(Ordering::SeqCst) {
panic!("cannot operate on committed txn!");
}
if let Some(guard) = &self.key_hashes {
let mut guard = guard.lock();
let (_, read_set) = &mut *guard;
read_set.insert(farmhash::hash32(key));
}
if let Some(entry) = self.local_storage.get(key) {
if entry.value().is_empty() {
return Ok(None);
@@ -75,7 +81,7 @@ impl Transaction {
if let Some(key_hashes) = &self.key_hashes {
let mut key_hashes = key_hashes.lock();
let (write_hashes, _) = &mut *key_hashes;
write_hashes.insert(crc32fast::hash(key));
write_hashes.insert(farmhash::hash32(key));
}
}
@@ -88,7 +94,7 @@ impl Transaction {
if let Some(key_hashes) = &self.key_hashes {
let mut key_hashes = key_hashes.lock();
let (write_hashes, _) = &mut *key_hashes;
write_hashes.insert(crc32fast::hash(key));
write_hashes.insert(farmhash::hash32(key));
}
}
@@ -96,6 +102,29 @@ impl Transaction {
self.committed
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.expect("cannot operate on committed txn!");
let _commit_lock = self.inner.mvcc().commit_lock.lock();
let serializability_check;
if let Some(guard) = &self.key_hashes {
let guard = guard.lock();
let (write_set, read_set) = &*guard;
println!(
"commit txn: write_set: {:?}, read_set: {:?}",
write_set, read_set
);
if !write_set.is_empty() {
let committed_txns = self.inner.mvcc().committed_txns.lock();
for (_, txn_data) in committed_txns.range((self.read_ts + 1)..) {
for key_hash in read_set {
if txn_data.key_hashes.contains(key_hash) {
bail!("serializable check failed");
}
}
}
}
serializability_check = true;
} else {
serializability_check = false;
}
let batch = self
.local_storage
.iter()
@@ -107,7 +136,32 @@ impl Transaction {
}
})
.collect::<Vec<_>>();
self.inner.write_batch(&batch)?;
let ts = self.inner.write_batch_inner(&batch)?;
if serializability_check {
let mut committed_txns = self.inner.mvcc().committed_txns.lock();
let mut key_hashes = self.key_hashes.as_ref().unwrap().lock();
let (write_set, _) = &mut *key_hashes;
let old_data = committed_txns.insert(
ts,
CommittedTxnData {
key_hashes: std::mem::take(write_set),
read_ts: self.read_ts,
commit_ts: ts,
},
);
assert!(old_data.is_none());
// remove unneeded txn data
let watermark = self.inner.mvcc().watermark();
while let Some(entry) = committed_txns.first_entry() {
if *entry.key() < watermark {
entry.remove();
} else {
break;
}
}
}
Ok(())
}
}
@@ -164,7 +218,7 @@ impl StorageIterator for TxnLocalIterator {
}
pub struct TxnIterator {
_txn: Arc<Transaction>,
txn: Arc<Transaction>,
iter: TwoMergeIterator<TxnLocalIterator, FusedIterator<LsmIterator>>,
}
@@ -173,8 +227,11 @@ impl TxnIterator {
txn: Arc<Transaction>,
iter: TwoMergeIterator<TxnLocalIterator, FusedIterator<LsmIterator>>,
) -> Result<Self> {
let mut iter = Self { _txn: txn, iter };
let mut iter = Self { txn, iter };
iter.skip_deletes()?;
if iter.is_valid() {
iter.add_to_read_set(iter.key());
}
Ok(iter)
}
@@ -184,6 +241,14 @@ impl TxnIterator {
}
Ok(())
}
fn add_to_read_set(&self, key: &[u8]) {
if let Some(guard) = &self.txn.key_hashes {
let mut guard = guard.lock();
let (_, read_set) = &mut *guard;
read_set.insert(farmhash::hash32(key));
}
}
}
impl StorageIterator for TxnIterator {
@@ -204,6 +269,9 @@ impl StorageIterator for TxnIterator {
fn next(&mut self) -> Result<()> {
self.iter.next()?;
self.skip_deletes()?;
if self.is_valid() {
self.add_to_read_set(self.key());
}
Ok(())
}

View File

@@ -17,3 +17,4 @@ mod week3_day2;
mod week3_day3;
mod week3_day4;
mod week3_day5;
mod week3_day6;

View File

@@ -46,6 +46,8 @@ fn test_txn_integration() {
],
);
let txn4 = storage.new_txn().unwrap();
assert_eq!(txn4.get(b"test1").unwrap(), Some(Bytes::from("233")));
assert_eq!(txn4.get(b"test2").unwrap(), Some(Bytes::from("233")));
check_lsm_iter_result_by_key(
&mut txn4.scan(Bound::Unbounded, Bound::Unbounded).unwrap(),
vec![
@@ -53,4 +55,21 @@ fn test_txn_integration() {
(Bytes::from("test2"), Bytes::from("233")),
],
);
txn4.put(b"test2", b"2333");
assert_eq!(txn4.get(b"test1").unwrap(), Some(Bytes::from("233")));
assert_eq!(txn4.get(b"test2").unwrap(), Some(Bytes::from("2333")));
check_lsm_iter_result_by_key(
&mut txn4.scan(Bound::Unbounded, Bound::Unbounded).unwrap(),
vec![
(Bytes::from("test1"), Bytes::from("233")),
(Bytes::from("test2"), Bytes::from("2333")),
],
);
txn4.delete(b"test2");
assert_eq!(txn4.get(b"test1").unwrap(), Some(Bytes::from("233")));
assert_eq!(txn4.get(b"test2").unwrap(), None);
check_lsm_iter_result_by_key(
&mut txn4.scan(Bound::Unbounded, Bound::Unbounded).unwrap(),
vec![(Bytes::from("test1"), Bytes::from("233"))],
);
}

View File

@@ -0,0 +1,108 @@
use std::ops::Bound;
use bytes::Bytes;
use tempfile::tempdir;
use crate::{
compact::CompactionOptions,
iterators::StorageIterator,
lsm_storage::{LsmStorageOptions, MiniLsm},
};
#[test]
fn test_serializable_1() {
let dir = tempdir().unwrap();
let mut options = LsmStorageOptions::default_for_week2_test(CompactionOptions::NoCompaction);
options.serializable = true;
let storage = MiniLsm::open(&dir, options.clone()).unwrap();
storage.put(b"key1", b"1").unwrap();
storage.put(b"key2", b"2").unwrap();
let txn1 = storage.new_txn().unwrap();
let txn2 = storage.new_txn().unwrap();
txn1.put(b"key1", &txn1.get(b"key2").unwrap().unwrap());
txn2.put(b"key2", &txn2.get(b"key1").unwrap().unwrap());
txn1.commit().unwrap();
assert!(txn2.commit().is_err());
drop(txn2);
assert_eq!(storage.get(b"key1").unwrap(), Some(Bytes::from("2")));
assert_eq!(storage.get(b"key2").unwrap(), Some(Bytes::from("2")));
}
#[test]
fn test_serializable_2() {
let dir = tempdir().unwrap();
let mut options = LsmStorageOptions::default_for_week2_test(CompactionOptions::NoCompaction);
options.serializable = true;
let storage = MiniLsm::open(&dir, options.clone()).unwrap();
let txn1 = storage.new_txn().unwrap();
let txn2 = storage.new_txn().unwrap();
txn1.put(b"key1", b"1");
txn2.put(b"key1", b"2");
txn1.commit().unwrap();
txn2.commit().unwrap();
assert_eq!(storage.get(b"key1").unwrap(), Some(Bytes::from("2")));
}
#[test]
fn test_serializable_3_ts_range() {
let dir = tempdir().unwrap();
let mut options = LsmStorageOptions::default_for_week2_test(CompactionOptions::NoCompaction);
options.serializable = true;
let storage = MiniLsm::open(&dir, options.clone()).unwrap();
storage.put(b"key1", b"1").unwrap();
storage.put(b"key2", b"2").unwrap();
let txn1 = storage.new_txn().unwrap();
txn1.put(b"key1", &txn1.get(b"key2").unwrap().unwrap());
txn1.commit().unwrap();
let txn2 = storage.new_txn().unwrap();
txn2.put(b"key2", &txn2.get(b"key1").unwrap().unwrap());
txn2.commit().unwrap();
drop(txn2);
assert_eq!(storage.get(b"key1").unwrap(), Some(Bytes::from("2")));
assert_eq!(storage.get(b"key2").unwrap(), Some(Bytes::from("2")));
}
#[test]
fn test_serializable_4_scan() {
let dir = tempdir().unwrap();
let mut options = LsmStorageOptions::default_for_week2_test(CompactionOptions::NoCompaction);
options.serializable = true;
let storage = MiniLsm::open(&dir, options.clone()).unwrap();
storage.put(b"key1", b"1").unwrap();
storage.put(b"key2", b"2").unwrap();
let txn1 = storage.new_txn().unwrap();
let txn2 = storage.new_txn().unwrap();
txn1.put(b"key1", &txn1.get(b"key2").unwrap().unwrap());
txn1.commit().unwrap();
let mut iter = txn2.scan(Bound::Unbounded, Bound::Unbounded).unwrap();
while iter.is_valid() {
iter.next().unwrap();
}
txn2.put(b"key2", b"1");
assert!(txn2.commit().is_err());
drop(txn2);
assert_eq!(storage.get(b"key1").unwrap(), Some(Bytes::from("2")));
assert_eq!(storage.get(b"key2").unwrap(), Some(Bytes::from("2")));
}
#[test]
fn test_serializable_5_read_only() {
let dir = tempdir().unwrap();
let mut options = LsmStorageOptions::default_for_week2_test(CompactionOptions::NoCompaction);
options.serializable = true;
let storage = MiniLsm::open(&dir, options.clone()).unwrap();
storage.put(b"key1", b"1").unwrap();
storage.put(b"key2", b"2").unwrap();
let txn1 = storage.new_txn().unwrap();
txn1.put(b"key1", &txn1.get(b"key2").unwrap().unwrap());
txn1.commit().unwrap();
let txn2 = storage.new_txn().unwrap();
txn2.get(b"key1").unwrap().unwrap();
let mut iter = txn2.scan(Bound::Unbounded, Bound::Unbounded).unwrap();
while iter.is_valid() {
iter.next().unwrap();
}
txn2.commit().unwrap();
assert_eq!(storage.get(b"key1").unwrap(), Some(Bytes::from("2")));
assert_eq!(storage.get(b"key2").unwrap(), Some(Bytes::from("2")));
}