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

RocksDB stress test fixes + find_merge_thread() bugfix

Browse Source
This commit is contained in:
mdecimus
2023-12-20 17:06:32 +01:00
parent f7313eecaf
commit d4aca0a8e0
28 changed files with 819 additions and 268 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/store/Cargo.toml
View File

@@ -28,7 +28,7 @@ num_cpus = { version = "1.15.0", optional = true }
blake3 = "1.3.3"
tracing = "0.1"
lz4_flex = { version = "0.11" }
deadpool-postgres = { version = "0.11.0", optional = true }
deadpool-postgres = { version = "0.12.1", optional = true }
tokio-postgres = { version = "0.7.10", optional = true }
tokio-rustls = { version = "0.25.0", optional = true }
rustls = { version = "0.22.0", optional = true }

16
crates/store/src/backend/postgres/write.rs
View File

@@ -154,10 +154,13 @@ impl PostgresStore {
.await?
}
} else {
trx
.prepare_cached(
&format!("INSERT INTO {} (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v", table),
)
trx.prepare_cached(&format!(
concat!(
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
"ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v"
),
table
))
.await?
};
@@ -242,7 +245,10 @@ impl PostgresStore {
.serialize(false);
let s = trx
.prepare_cached("INSERT INTO l (k, v) VALUES ($1, $2) ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v")
.prepare_cached(concat!(
"INSERT INTO l (k, v) VALUES ($1, $2) ",
"ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v"
))
.await?;
trx.execute(&s, &[&key, set]).await?;
}

429
crates/store/src/backend/rocksdb/write.rs
View File

@@ -22,20 +22,27 @@
*/
use std::{
sync::Arc,
thread::sleep,
time::{Duration, Instant},
};
use rand::Rng;
use rocksdb::{Direction, ErrorKind, IteratorMode};
use roaring::RoaringBitmap;
use rocksdb::{
BoundColumnFamily, Direction, ErrorKind, IteratorMode, OptimisticTransactionDB,
OptimisticTransactionOptions, WriteOptions,
};
use super::{
bitmap::{clear_bit, set_bit},
RocksDbStore, CF_BITMAPS, CF_COUNTERS, CF_INDEXES, CF_LOGS, CF_VALUES,
};
use crate::{
write::{Batch, Operation, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME},
BitmapKey, IndexKey, Key, LogKey, ValueKey,
write::{
Batch, BitmapClass, Operation, ValueClass, ValueOp, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME,
},
BitmapKey, Deserialize, IndexKey, Key, LogKey, ValueKey,
};
impl RocksDbStore {
@@ -46,138 +53,28 @@ impl RocksDbStore {
let start = Instant::now();
let mut retry_count = 0;
let cf_bitmaps = db.cf_handle(CF_BITMAPS).unwrap();
let cf_values = db.cf_handle(CF_VALUES).unwrap();
let cf_indexes = db.cf_handle(CF_INDEXES).unwrap();
let cf_logs = db.cf_handle(CF_LOGS).unwrap();
let cf_counters = db.cf_handle(CF_COUNTERS).unwrap();
let mut txn_opts = OptimisticTransactionOptions::default();
txn_opts.set_snapshot(true);
let txn = RocksDBTransaction {
db: &db,
cf_bitmaps: db.cf_handle(CF_BITMAPS).unwrap(),
cf_values: db.cf_handle(CF_VALUES).unwrap(),
cf_indexes: db.cf_handle(CF_INDEXES).unwrap(),
cf_logs: db.cf_handle(CF_LOGS).unwrap(),
cf_counters: db.cf_handle(CF_COUNTERS).unwrap(),
txn_opts,
batch: &batch,
};
loop {
let mut account_id = u32::MAX;
let mut collection = u8::MAX;
let mut document_id = u32::MAX;
let txn = self.db.transaction();
let mut wb = txn.get_writebatch();
for op in &batch.ops {
match op {
Operation::AccountId {
account_id: account_id_,
} => {
account_id = *account_id_;
}
Operation::Collection {
collection: collection_,
} => {
collection = *collection_;
}
Operation::DocumentId {
document_id: document_id_,
} => {
document_id = *document_id_;
}
Operation::Value {
class,
op: ValueOp::Add(by),
} => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
}
.serialize(false);
wb.merge_cf(&cf_counters, &key, &by.to_le_bytes()[..]);
}
Operation::Value { class, op } => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
};
let key = key.serialize(false);
if let ValueOp::Set(value) = op {
wb.put_cf(&cf_values, &key, value);
} else {
wb.delete_cf(&cf_values, &key);
}
}
Operation::Index { field, key, set } => {
let key = IndexKey {
account_id,
collection,
document_id,
field: *field,
key,
}
.serialize(false);
if *set {
wb.put_cf(&cf_indexes, &key, []);
} else {
wb.delete_cf(&cf_indexes, &key);
}
}
Operation::Bitmap { class, set } => {
let key = BitmapKey {
account_id,
collection,
class,
block_num: 0,
}
.serialize(false);
let value = if *set {
set_bit(document_id)
} else {
clear_bit(document_id)
};
wb.merge_cf(&cf_bitmaps, key, value);
}
Operation::Log {
collection,
change_id,
set,
} => {
let key = LogKey {
account_id,
collection: *collection,
change_id: *change_id,
}
.serialize(false);
wb.put_cf(&cf_logs, &key, set);
}
Operation::AssertValue {
class,
assert_value,
} => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
};
let key = key.serialize(false);
let matches = txn
.get_cf(&cf_values, &key)?
.map(|value| assert_value.matches(&value))
.unwrap_or_else(|| assert_value.is_none());
if !matches {
return Err(crate::Error::AssertValueFailed);
}
}
}
}
match db.write(wb) {
Ok(_) => {
return Ok(());
match txn.commit() {
Ok(success) => {
return if success {
Ok(())
} else {
Err(crate::Error::AssertValueFailed)
};
}
Err(err) => match err.kind() {
ErrorKind::Busy | ErrorKind::MergeInProgress | ErrorKind::TryAgain
@@ -231,3 +128,269 @@ impl RocksDbStore {
Ok(())
}
}
struct RocksDBTransaction<'x> {
db: &'x OptimisticTransactionDB,
cf_bitmaps: Arc<BoundColumnFamily<'x>>,
cf_values: Arc<BoundColumnFamily<'x>>,
cf_indexes: Arc<BoundColumnFamily<'x>>,
cf_logs: Arc<BoundColumnFamily<'x>>,
cf_counters: Arc<BoundColumnFamily<'x>>,
txn_opts: OptimisticTransactionOptions,
batch: &'x Batch,
}
impl<'x> RocksDBTransaction<'x> {
fn commit(&self) -> Result<bool, rocksdb::Error> {
let mut account_id = u32::MAX;
let mut collection = u8::MAX;
let mut document_id = u32::MAX;
let txn = self
.db
.transaction_opt(&WriteOptions::default(), &self.txn_opts);
if !self.batch.is_atomic() {
for op in &self.batch.ops {
match op {
Operation::AccountId {
account_id: account_id_,
} => {
account_id = *account_id_;
}
Operation::Collection {
collection: collection_,
} => {
collection = *collection_;
}
Operation::DocumentId {
document_id: document_id_,
} => {
document_id = *document_id_;
}
Operation::Value {
class,
op: ValueOp::Add(by),
} => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
}
.serialize(false);
txn.merge_cf(&self.cf_counters, &key, &by.to_le_bytes()[..])?;
}
Operation::Value { class, op } => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
};
let key = key.serialize(false);
if let ValueOp::Set(value) = op {
txn.put_cf(&self.cf_values, &key, value)?;
if matches!(class, ValueClass::ReservedId) {
if let Some(bitmap) = txn
.get_pinned_cf(
&self.cf_bitmaps,
&BitmapKey {
account_id,
collection,
class: BitmapClass::DocumentIds,
block_num: 0,
}
.serialize(false),
//true,
)?
.and_then(|bytes| RoaringBitmap::deserialize(&bytes).ok())
{
if bitmap.contains(document_id) {
txn.rollback()?;
return Ok(false);
}
}
}
} else {
txn.delete_cf(&self.cf_values, &key)?;
}
}
Operation::Index { field, key, set } => {
let key = IndexKey {
account_id,
collection,
document_id,
field: *field,
key,
}
.serialize(false);
if *set {
txn.put_cf(&self.cf_indexes, &key, [])?;
} else {
txn.delete_cf(&self.cf_indexes, &key)?;
}
}
Operation::Bitmap { class, set } => {
let key = BitmapKey {
account_id,
collection,
class,
block_num: 0,
}
.serialize(false);
let value = if *set {
set_bit(document_id)
} else {
clear_bit(document_id)
};
txn.merge_cf(&self.cf_bitmaps, key, value)?;
}
Operation::Log {
collection,
change_id,
set,
} => {
let key = LogKey {
account_id,
collection: *collection,
change_id: *change_id,
}
.serialize(false);
txn.put_cf(&self.cf_logs, &key, set)?;
}
Operation::AssertValue {
class,
assert_value,
} => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
};
let key = key.serialize(false);
let matches = txn
.get_pinned_for_update_cf(&self.cf_values, &key, true)?
.map(|value| assert_value.matches(&value))
.unwrap_or_else(|| assert_value.is_none());
if !matches {
txn.rollback()?;
return Ok(false);
}
}
}
}
txn.commit().map(|_| true)
} else {
let mut wb = txn.get_writebatch();
for op in &self.batch.ops {
match op {
Operation::AccountId {
account_id: account_id_,
} => {
account_id = *account_id_;
}
Operation::Collection {
collection: collection_,
} => {
collection = *collection_;
}
Operation::DocumentId {
document_id: document_id_,
} => {
document_id = *document_id_;
}
Operation::Value {
class,
op: ValueOp::Add(by),
} => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
}
.serialize(false);
wb.merge_cf(&self.cf_counters, &key, &by.to_le_bytes()[..]);
}
Operation::Value { class, op } => {
let key = ValueKey {
account_id,
collection,
document_id,
class,
};
let key = key.serialize(false);
if let ValueOp::Set(value) = op {
wb.put_cf(&self.cf_values, &key, value);
} else {
wb.delete_cf(&self.cf_values, &key);
}
}
Operation::Index { field, key, set } => {
let key = IndexKey {
account_id,
collection,
document_id,
field: *field,
key,
}
.serialize(false);
if *set {
wb.put_cf(&self.cf_indexes, &key, []);
} else {
wb.delete_cf(&self.cf_indexes, &key);
}
}
Operation::Bitmap { class, set } => {
let key = BitmapKey {
account_id,
collection,
class,
block_num: 0,
}
.serialize(false);
let value = if *set {
set_bit(document_id)
} else {
clear_bit(document_id)
};
wb.merge_cf(&self.cf_bitmaps, key, value);
}
Operation::Log {
collection,
change_id,
set,
} => {
let key = LogKey {
account_id,
collection: *collection,
change_id: *change_id,
}
.serialize(false);
wb.put_cf(&self.cf_logs, &key, set);
}
Operation::AssertValue { .. } => unreachable!(),
}
}
self.db.write(wb).map(|_| true)
}
}
}

136
crates/store/src/dispatch/store.rs
View File

@@ -31,6 +31,12 @@ use crate::{
SUBSPACE_INDEXES, SUBSPACE_LOGS, U32_LEN,
};
#[cfg(feature = "test_mode")]
lazy_static::lazy_static! {
pub static ref BITMAPS: std::sync::Arc<parking_lot::Mutex<std::collections::HashMap<Vec<u8>, std::collections::HashSet<u32>>>> =
std::sync::Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new()));
}
impl Store {
pub async fn get_value<U>(&self, key: impl Key) -> crate::Result<Option<U>>
where
@@ -141,6 +147,86 @@ impl Store {
}
pub async fn write(&self, batch: Batch) -> crate::Result<()> {
#[cfg(feature = "test_mode")]
if std::env::var("PARANOID_WRITE").map_or(false, |v| v == "1") {
use crate::write::Operation;
let mut account_id = u32::MAX;
let mut collection = u8::MAX;
let mut document_id = u32::MAX;
let mut bitmaps = Vec::new();
for op in &batch.ops {
match op {
Operation::AccountId {
account_id: account_id_,
} => {
account_id = *account_id_;
}
Operation::Collection {
collection: collection_,
} => {
collection = *collection_;
}
Operation::DocumentId {
document_id: document_id_,
} => {
document_id = *document_id_;
}
Operation::Bitmap { class, set } => {
let key = BitmapKey {
account_id,
collection,
block_num: 0,
class,
}
.serialize(false);
bitmaps.push((key, class.clone(), document_id, *set));
}
_ => {}
}
}
match self {
#[cfg(feature = "sqlite")]
Self::SQLite(store) => store.write(batch).await,
#[cfg(feature = "foundation")]
Self::FoundationDb(store) => store.write(batch).await,
#[cfg(feature = "postgres")]
Self::PostgreSQL(store) => store.write(batch).await,
#[cfg(feature = "mysql")]
Self::MySQL(store) => store.write(batch).await,
#[cfg(feature = "rocks")]
Self::RocksDb(store) => store.write(batch).await,
}?;
for (key, class, document_id, set) in bitmaps {
let mut bitmaps = BITMAPS.lock();
let map = bitmaps.entry(key).or_default();
if set {
if !map.insert(document_id) {
println!(
concat!(
"WARNING: key {:?} already contains document {} for account ",
"{}, collection {}"
),
class, document_id, account_id, collection
);
}
} else if !map.remove(&document_id) {
println!(
concat!(
"WARNING: key {:?} does not contain document {} for account ",
"{}, collection {}"
),
class, document_id, account_id, collection
);
}
}
return Ok(());
}
match self {
#[cfg(feature = "sqlite")]
Self::SQLite(store) => store.write(batch).await,
@@ -303,6 +389,8 @@ impl Store {
.await
.unwrap();
}
BITMAPS.lock().clear();
}
#[cfg(feature = "test_mode")]
@@ -365,6 +453,8 @@ impl Store {
#[allow(unused_variables)]
pub async fn assert_is_empty(&self, blob_store: crate::BlobStore) {
use utils::codec::leb128::Leb128Iterator;
use crate::{SUBSPACE_BLOBS, SUBSPACE_COUNTERS, SUBSPACE_VALUES};
self.blob_expire_all().await;
@@ -406,9 +496,46 @@ impl Store {
return Ok(true);
}
const BM_DOCUMENT_IDS: u8 = 0;
const BM_TAG: u8 = 1 << 6;
const BM_TEXT: u8 = 1 << 7;
const TAG_TEXT: u8 = 1 << 0;
const TAG_STATIC: u8 = 1 << 1;
match key[5] {
BM_DOCUMENT_IDS => {
eprint!("Found document ids bitmap");
}
BM_TAG => {
eprint!(
"Found tagged id {} bitmap",
key[7..].iter().next_leb128::<u32>().unwrap()
);
}
TAG_TEXT => {
eprint!(
"Found tagged text {:?} bitmap",
String::from_utf8_lossy(&key[7..])
);
}
TAG_STATIC => {
eprint!("Found tagged static {} bitmap", key[7]);
}
other => {
if other & BM_TEXT == BM_TEXT {
eprint!(
"Found text hash {:?} bitmap",
String::from_utf8_lossy(&key[7..])
);
} else {
eprint!("Found unknown bitmap");
}
}
}
eprintln!(
concat!(
"Table bitmaps is not empty, account {}, collection {},",
", account {}, collection {},",
" family {}, field {}, key {:?}: {:?}"
),
u32::from_be_bytes(key[0..4].try_into().unwrap()),
@@ -420,7 +547,8 @@ impl Store {
);
}
SUBSPACE_VALUES
if key[0] >= 20
if key[0] == 3
|| key[0] >= 20
|| key.get(1..5).unwrap_or_default() == u32::MAX.to_be_bytes() =>
{
// Ignore lastId counter and ID mappings
@@ -433,7 +561,7 @@ impl Store {
SUBSPACE_INDEXES => {
eprintln!(
concat!(
"Table index is not empty, account {}, collection {}, ",
"Found index key, account {}, collection {}, ",
"document {}, property {}, value {:?}: {:?}"
),
u32::from_be_bytes(key[0..4].try_into().unwrap()),
@@ -446,7 +574,7 @@ impl Store {
}
_ => {
eprintln!(
"Table {:?} is not empty: {:?} {:?}",
"Found key in {:?}: {:?} {:?}",
char::from(subspace),
key,
value

8
crates/store/src/fts/index.rs
View File

@@ -286,6 +286,14 @@ impl Store {
{
term_index
} else {
tracing::debug!(
context = "fts_remove",
event = "not_found",
account_id = account_id,
collection = collection,
document_id = document_id,
"Term index not found"
);
return Ok(false);
};

2
crates/store/src/query/sort.rs
View File

@@ -40,7 +40,7 @@ pub struct Pagination {
anchor_offset: i32,
has_anchor: bool,
anchor_found: bool,
ids: Vec<u64>,
pub ids: Vec<u64>,
prefix_key: Option<ValueKey<ValueClass>>,
prefix_unique: bool,
}

10
crates/store/src/write/assert.rs
View File

@@ -29,11 +29,12 @@ pub struct HashedValue<T: Deserialize> {
pub inner: T,
}
#[derive(Debug, PartialEq, Eq, Hash)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AssertValue {
U32(u32),
U64(u64),
Hash(u64),
Some,
None,
}
@@ -47,6 +48,12 @@ pub trait ToAssertValue {
fn to_assert_value(&self) -> AssertValue;
}
impl ToAssertValue for AssertValue {
fn to_assert_value(&self) -> AssertValue {
*self
}
}
impl ToAssertValue for () {
fn to_assert_value(&self) -> AssertValue {
AssertValue::None
@@ -84,6 +91,7 @@ impl AssertValue {
AssertValue::U64(v) => bytes.len() == U64_LEN && u64::deserialize(bytes).unwrap() == *v,
AssertValue::Hash(v) => xxhash_rust::xxh3::xxh3_64(bytes) == *v,
AssertValue::None => false,
AssertValue::Some => true,
}
}

15
crates/store/src/write/batch.rs
View File

@@ -202,6 +202,21 @@ impl BatchBuilder {
}
}
impl Batch {
pub fn is_atomic(&self) -> bool {
!self.ops.iter().any(|op| {
matches!(
op,
Operation::AssertValue { .. }
| Operation::Value {
class: ValueClass::ReservedId,
op: ValueOp::Set(_)
}
)
})
}
}
impl Default for BatchBuilder {
fn default() -> Self {
Self::new()