From a1f7ad891ac2d74ef4051b0aa4858a8fbd194f7b Mon Sep 17 00:00:00 2001 From: mdecimus Date: Tue, 28 Nov 2023 10:40:30 +0100 Subject: [PATCH] FDB blob storage support --- crates/store/src/backend/foundationdb/blob.rs | 122 +++++++++++++++++- crates/store/src/backend/foundationdb/read.rs | 2 +- crates/store/src/lib.rs | 9 ++ tests/src/store/assign_id.rs | 20 +-- tests/src/store/blob.rs | 82 ++++++++++-- tests/src/store/mod.rs | 5 +- 6 files changed, 214 insertions(+), 26 deletions(-) diff --git a/crates/store/src/backend/foundationdb/blob.rs b/crates/store/src/backend/foundationdb/blob.rs index f416a373..08a64584 100644 --- a/crates/store/src/backend/foundationdb/blob.rs +++ b/crates/store/src/backend/foundationdb/blob.rs @@ -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, ) -> crate::Result>> { - 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> = 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 { - 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()), + } } } diff --git a/crates/store/src/backend/foundationdb/read.rs b/crates/store/src/backend/foundationdb/read.rs index 3f0e72ea..8dffd31f 100644 --- a/crates/store/src/backend/foundationdb/read.rs +++ b/crates/store/src/backend/foundationdb/read.rs @@ -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}"), diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index 698d6886..2d848be3 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -213,3 +213,12 @@ impl From for FtsStore { Self::Store(store) } } + +impl From for BlobStore { + fn from(store: Store) -> Self { + match store { + Store::SQLite(store) => Self::Sqlite(store), + Store::FoundationDb(store) => Self::FoundationDb(store), + } + } +} diff --git a/tests/src/store/assign_id.rs b/tests/src/store/assign_id.rs index 7cb36b1e..0d589bcb 100644 --- a/tests/src/store/assign_id.rs +++ b/tests/src/store/assign_id.rs @@ -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; } diff --git a/tests/src/store/blob.rs b/tests/src/store/blob.rs index 2526f358..9bb233e6 100644 --- a/tests/src/store/blob.rs +++ b/tests/src/store/blob.rs @@ -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()); } diff --git a/tests/src/store/mod.rs b/tests/src/store/mod.rs index 23fc146b..aca2c777 100644 --- a/tests/src/store/mod.rs +++ b/tests/src/store/mod.rs @@ -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();