filip/Stalwart
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

FDB blob storage support

Browse Source
This commit is contained in:
mdecimus
2023-11-28 10:40:30 +01:00
parent 0168c1dca8
commit a1f7ad891a
6 changed files with 214 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
crates/store/src/backend/foundationdb/blob.rs
View File

@@ -23,22 +23,138 @@
use std::ops::Range;
use foundationdb::{options::StreamingMode, FdbError, KeySelector, RangeOption};
use futures::StreamExt;
use crate::{write::key::KeySerializer, Error, BLOB_HASH_LEN, SUBSPACE_BLOB_DATA};
use super::FdbStore;
const MAX_BLOCK_SIZE: usize = 100000;
impl FdbStore {
pub(crate) async fn get_blob(
&self,
key: &[u8],
range: Range<u32>,
) -> crate::Result<Option<Vec<u8>>> {
todo!()
let block_start = range.start as usize / MAX_BLOCK_SIZE;
let bytes_start = range.start as usize % MAX_BLOCK_SIZE;
let block_end = (range.end as usize / MAX_BLOCK_SIZE) + 1;
let begin = KeySerializer::new(key.len() + 3)
.write(SUBSPACE_BLOB_DATA)
.write(key)
.write(block_start as u16)
.finalize();
let end = KeySerializer::new(key.len() + 3)
.write(SUBSPACE_BLOB_DATA)
.write(key)
.write(block_end as u16)
.finalize();
let key_len = begin.len();
let trx = self.db.create_trx()?;
let mut values = trx.get_ranges(
RangeOption {
begin: KeySelector::first_greater_or_equal(begin),
end: KeySelector::first_greater_or_equal(end),
mode: StreamingMode::WantAll,
reverse: false,
..RangeOption::default()
},
true,
);
let mut blob_data: Option<Vec<u8>> = None;
let blob_range = (range.end - range.start) as usize;
'outer: while let Some(values) = values.next().await {
for value in values? {
let key = value.key();
if key.len() == key_len {
let value = value.value();
if let Some(blob_data) = &mut blob_data {
blob_data.extend_from_slice(
value
.get(
..std::cmp::min(
blob_range.saturating_sub(blob_data.len()),
value.len(),
),
)
.unwrap_or(&[]),
);
if blob_data.len() == blob_range {
break 'outer;
}
} else {
let blob_size = if blob_range <= (5 * (1 << 20)) {
blob_range
} else if value.len() == MAX_BLOCK_SIZE {
MAX_BLOCK_SIZE * 2
} else {
value.len()
};
let mut blob_data_ = Vec::with_capacity(blob_size);
blob_data_.extend_from_slice(
value
.get(
bytes_start
..std::cmp::min(bytes_start + blob_range, value.len()),
)
.unwrap_or(&[]),
);
if blob_data_.len() == blob_range {
return Ok(Some(blob_data_));
}
blob_data = blob_data_.into();
}
}
}
}
Ok(blob_data)
}
pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> crate::Result<()> {
todo!()
for (chunk_pos, chunk_bytes) in data.chunks(MAX_BLOCK_SIZE).enumerate() {
let trx = self.db.create_trx()?;
trx.set(
&KeySerializer::new(key.len() + 3)
.write(SUBSPACE_BLOB_DATA)
.write(key)
.write(chunk_pos as u16)
.finalize(),
chunk_bytes,
);
trx.commit()
.await
.map_err(|err| Error::from(FdbError::from(err)))?;
}
Ok(())
}
pub(crate) async fn delete_blob(&self, key: &[u8]) -> crate::Result<bool> {
todo!()
if key.len() < BLOB_HASH_LEN {
return Ok(false);
}
let trx = self.db.create_trx()?;
trx.clear_range(
&KeySerializer::new(key.len() + 3)
.write(SUBSPACE_BLOB_DATA)
.write(key)
.write(0u16)
.finalize(),
&KeySerializer::new(key.len() + 3)
.write(SUBSPACE_BLOB_DATA)
.write(key)
.write(u16::MAX)
.finalize(),
);
match trx.commit().await {
Ok(_) => Ok(true),
Err(err) => Err(FdbError::from(err).into()),
}
}
}

2
crates/store/src/backend/foundationdb/read.rs
View File

@@ -351,7 +351,7 @@ impl FdbStore {
);
}
SUBSPACE_LOGS => {
delete_keys.push(key.to_vec());
delete_keys.push(key_.to_vec());
}
_ => panic!("Invalid key found in database: {key:?} for subspace {subspace}"),

9
crates/store/src/lib.rs
View File

@@ -213,3 +213,12 @@ impl From<Store> for FtsStore {
Self::Store(store)
}
}
impl From<Store> for BlobStore {
fn from(store: Store) -> Self {
match store {
Store::SQLite(store) => Self::Sqlite(store),
Store::FoundationDb(store) => Self::FoundationDb(store),
}
}
}

20
tests/src/store/assign_id.rs
View File

@@ -50,7 +50,7 @@ async fn test_1(db: Store) {
for id in 0..100 {
handles.push({
let db = db.clone();
tokio::spawn(async move { db.assign_document_id(0, 0).await })
tokio::spawn(async move { db.assign_document_id(0, u8::MAX).await })
});
expected_ids.insert(id);
}
@@ -80,9 +80,9 @@ async fn test_2(db: Store) {
// Create document ids and try reassigning
let mut expected_ids = AHashSet::new();
let mut batch = BatchBuilder::new();
batch.with_account_id(0).with_collection(0);
batch.with_account_id(0).with_collection(u8::MAX);
for pos in 0..100 {
let id = db.assign_document_id(0, 0).await.unwrap();
let id = db.assign_document_id(0, u8::MAX).await.unwrap();
if pos % 2 == 0 {
batch.create_document(id);
} else {
@@ -95,14 +95,14 @@ async fn test_2(db: Store) {
tokio::time::sleep(Duration::from_secs(3)).await;
for _ in 0..expected_ids.len() {
let id = db.assign_document_id(0, 0).await.unwrap();
let id = db.assign_document_id(0, u8::MAX).await.unwrap();
assert!(
expected_ids.remove(&id),
"already assigned or invalid: {id}"
);
}
assert_eq!(db.assign_document_id(0, 0).await.unwrap(), 100);
assert_eq!(db.assign_document_id(0, 0).await.unwrap(), 101);
assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 100);
assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 101);
db.destroy().await;
}
@@ -111,7 +111,7 @@ async fn test_3(db: Store) {
// Try reassigning deleted ids
let mut expected_ids = AHashSet::new();
let mut batch = BatchBuilder::new();
batch.with_account_id(0).with_collection(0);
batch.with_account_id(0).with_collection(u8::MAX);
for id in 0..100 {
if id % 2 == 0 {
batch.create_document(id);
@@ -121,14 +121,14 @@ async fn test_3(db: Store) {
}
db.write(batch.build()).await.unwrap();
for _ in 0..expected_ids.len() {
let id = db.assign_document_id(0, 0).await.unwrap();
let id = db.assign_document_id(0, u8::MAX).await.unwrap();
assert!(
expected_ids.remove(&id),
"already assigned or invalid: {id}"
);
}
assert_eq!(db.assign_document_id(0, 0).await.unwrap(), 100);
assert_eq!(db.assign_document_id(0, 0).await.unwrap(), 101);
assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 100);
assert_eq!(db.assign_document_id(0, u8::MAX).await.unwrap(), 101);
db.destroy().await;
}

82
tests/src/store/blob.rs
View File

@@ -52,7 +52,6 @@ path = "{TMP}/db.db?mode=rwc"
#[tokio::test]
pub async fn blob_tests() {
let temp_dir = TempDir::new("blob_tests", true);
let mut blob_store = None;
for (store_id, store_cfg) in [("s3", CONFIG_S3), ("fs", CONFIG_LOCAL)] {
let config =
@@ -65,20 +64,24 @@ pub async fn blob_tests() {
_ => unreachable!(),
};
println!("Testing store {}...", store_id);
println!("Testing blob store {}...", store_id);
test_store(blob_store_.clone()).await;
blob_store = Some(blob_store_);
}
let blob_store = blob_store.unwrap();
// Start SQLite store
// Init store
let store: Store = SqliteStore::open(
//let store: Store = FdbStore::open(
&Config::new(&CONFIG_DB.replace("{TMP}", temp_dir.path.as_path().to_str().unwrap()))
.unwrap(),
)
.await
.unwrap()
.into();
store.destroy().await;
// Test internal blob store
let blob_store: BlobStore = store.clone().into();
test_store(blob_store.clone()).await;
// Blob hash exists
let hash = BlobHash::from(b"abc".as_slice());
@@ -387,17 +390,76 @@ pub async fn blob_tests() {
}
async fn test_store(store: BlobStore) {
// Test small blob
const DATA: &[u8] = b"Lorem ipsum dolor sit amet, consectetur adipiscing elit. Fusce erat nisl, dignissim a porttitor id, varius nec arcu. Sed mauris.";
let hash = BlobHash::from(DATA);
store.put_blob(b"abc", DATA).await.unwrap();
store.put_blob(hash.as_slice(), DATA).await.unwrap();
assert_eq!(
String::from_utf8(store.get_blob(b"abc", 0..u32::MAX).await.unwrap().unwrap()).unwrap(),
String::from_utf8(
store
.get_blob(hash.as_slice(), 0..u32::MAX)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(DATA).unwrap()
);
assert_eq!(
String::from_utf8(store.get_blob(b"abc", 11..57).await.unwrap().unwrap()).unwrap(),
String::from_utf8(
store
.get_blob(hash.as_slice(), 11..57)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(&DATA[11..57]).unwrap()
);
assert!(store.delete_blob(b"abc").await.unwrap());
assert!(store.get_blob(b"abc", 0..u32::MAX).await.unwrap().is_none());
assert!(store.delete_blob(hash.as_slice()).await.unwrap());
assert!(store
.get_blob(hash.as_slice(), 0..u32::MAX)
.await
.unwrap()
.is_none());
// Test large blob
let mut data = Vec::with_capacity(50 * 1024 * 1024);
while data.len() < 50 * 1024 * 1024 {
data.extend_from_slice(DATA);
let marker = format!(" [{}] ", data.len());
data.extend_from_slice(marker.as_bytes());
}
let hash = BlobHash::from(&data);
store.put_blob(hash.as_slice(), &data).await.unwrap();
assert_eq!(
String::from_utf8(
store
.get_blob(hash.as_slice(), 0..u32::MAX)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(&data).unwrap()
);
assert_eq!(
String::from_utf8(
store
.get_blob(hash.as_slice(), 3000111..4000999)
.await
.unwrap()
.unwrap()
)
.unwrap(),
std::str::from_utf8(&data[3000111..4000999]).unwrap()
);
assert!(store.delete_blob(hash.as_slice()).await.unwrap());
assert!(store
.get_blob(hash.as_slice(), 0..u32::MAX)
.await
.unwrap()
.is_none());
}

5
tests/src/store/mod.rs
View File

@@ -29,7 +29,7 @@ use std::io::Read;
use ::store::Store;
use store::backend::sqlite::SqliteStore;
use store::backend::{foundationdb::FdbStore, sqlite::SqliteStore};
use utils::config::Config;
pub struct TempDir {
@@ -49,7 +49,8 @@ pub async fn store_tests() {
temp_dir.path.display(),
temp_dir.path.display()
);
let db: Store = SqliteStore::open(&Config::new(&config_file).unwrap())
//let db: Store = SqliteStore::open(&Config::new(&config_file).unwrap())
let db: Store = FdbStore::open(&Config::new(&config_file).unwrap())
.await
.unwrap()
.into();