Import upstream v0.16.22, stripped
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use std::ops::Range;
|
||||
|
||||
use super::{CF_BLOBS, RocksDbStore, into_error};
|
||||
|
||||
impl RocksDbStore {
|
||||
pub(crate) async fn get_blob(
|
||||
&self,
|
||||
key: &[u8],
|
||||
range: Range<usize>,
|
||||
) -> trc::Result<Option<Vec<u8>>> {
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
db.get_pinned_cf(&db.cf_handle(CF_BLOBS).unwrap(), key)
|
||||
.map(|obj| {
|
||||
obj.map(|bytes| {
|
||||
if range.start == 0 && range.end == usize::MAX {
|
||||
bytes.to_vec()
|
||||
} else {
|
||||
bytes
|
||||
.get(range.start..std::cmp::min(bytes.len(), range.end))
|
||||
.unwrap_or_default()
|
||||
.to_vec()
|
||||
}
|
||||
})
|
||||
})
|
||||
.map_err(into_error)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> {
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
db.put_cf(&db.cf_handle(CF_BLOBS).unwrap(), key, data)
|
||||
.map_err(into_error)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> {
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
db.delete_cf(&db.cf_handle(CF_BLOBS).unwrap(), key)
|
||||
.map_err(into_error)
|
||||
.map(|_| true)
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{CF_BLOBS, RocksDbStore};
|
||||
use crate::*;
|
||||
use ::registry::schema::structs;
|
||||
use rocksdb::{
|
||||
BlockBasedOptions, Cache, ColumnFamilyDescriptor, DBCompressionType, MergeOperands,
|
||||
OptimisticTransactionDB, Options,
|
||||
};
|
||||
use std::path::PathBuf;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
const MIN_WRITE_BUFFER_SIZE: usize = 4 * 1024 * 1024;
|
||||
const MAX_WRITE_BUFFER_SIZE: usize = 64 * 1024 * 1024;
|
||||
const MIN_DB_WRITE_BUFFER_SIZE: usize = 32 * 1024 * 1024;
|
||||
const BLOOM_BITS_PER_KEY: f64 = 10.0;
|
||||
const SCAN_BLOCK_SIZE: usize = 16 * 1024;
|
||||
const CHURN_TARGET_FILE_SIZE: u64 = 16 * 1024 * 1024;
|
||||
const CHURN_DELETION_WINDOW: usize = 4096;
|
||||
const CHURN_DELETION_TRIGGER: usize = 1024;
|
||||
const CHURN_DELETION_RATIO: f64 = 0.5;
|
||||
const BYTES_PER_SYNC: u64 = 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum CfProfile {
|
||||
/// Read through `get_value` / `key_exists`, so a whole key bloom filter pays off.
|
||||
PointLookup,
|
||||
/// Read only through `iterate`, which never consults a whole key bloom filter.
|
||||
Scan,
|
||||
/// Point read and point deleted at a high rate.
|
||||
Churn,
|
||||
/// Scanned from the oldest key and point deleted once consumed, with empty values.
|
||||
Queue,
|
||||
/// Counters updated through the merge operator.
|
||||
Counter,
|
||||
/// Blob values held in RocksDB blob files.
|
||||
Blob,
|
||||
}
|
||||
|
||||
impl RocksDbStore {
|
||||
pub async fn open(config: structs::RocksDbStore) -> Result<Store, String> {
|
||||
// Create the database directory if it doesn't exist
|
||||
let idx_path: PathBuf = PathBuf::from(config.path);
|
||||
std::fs::create_dir_all(&idx_path).map_err(|err| {
|
||||
format!(
|
||||
"Failed to create database directory {}: {:?}",
|
||||
idx_path.display(),
|
||||
err
|
||||
)
|
||||
})?;
|
||||
|
||||
let cache = Cache::new_lru_cache(config.cache_size as usize);
|
||||
let write_buffer_size =
|
||||
((config.buffer_size as usize) / 4).clamp(MIN_WRITE_BUFFER_SIZE, MAX_WRITE_BUFFER_SIZE);
|
||||
let mut cfs = Vec::new();
|
||||
|
||||
// Counters
|
||||
for subspace in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
|
||||
cfs.push(ColumnFamilyDescriptor::new(
|
||||
std::str::from_utf8(&[subspace]).unwrap(),
|
||||
cf_options(CfProfile::Counter, &cache, write_buffer_size),
|
||||
));
|
||||
}
|
||||
|
||||
// Blobs
|
||||
let mut cf_opts = cf_options(CfProfile::Blob, &cache, write_buffer_size);
|
||||
cf_opts.set_enable_blob_files(true);
|
||||
cf_opts.set_min_blob_size(config.blob_size);
|
||||
cf_opts.set_enable_blob_gc(true);
|
||||
cf_opts.set_blob_gc_age_cutoff(1.0);
|
||||
cf_opts.set_blob_gc_force_threshold(0.5);
|
||||
cfs.push(ColumnFamilyDescriptor::new(CF_BLOBS, cf_opts));
|
||||
|
||||
// Other cfs
|
||||
for (subspace, profile) in [
|
||||
(SUBSPACE_INDEXES, CfProfile::Scan),
|
||||
(SUBSPACE_ACL, CfProfile::Scan),
|
||||
(SUBSPACE_TASK_QUEUE, CfProfile::Churn),
|
||||
(SUBSPACE_DELETED_ITEMS, CfProfile::Churn),
|
||||
(SUBSPACE_BLOB_LINK, CfProfile::Churn),
|
||||
(SUBSPACE_IN_MEMORY_VALUE, CfProfile::Churn),
|
||||
(SUBSPACE_PROPERTY, CfProfile::PointLookup),
|
||||
(SUBSPACE_REGISTRY, CfProfile::PointLookup),
|
||||
(SUBSPACE_QUEUE_MESSAGE, CfProfile::Churn),
|
||||
(SUBSPACE_QUEUE_EVENT, CfProfile::Queue),
|
||||
(SUBSPACE_REPORT_OUT, CfProfile::Churn),
|
||||
(SUBSPACE_REPORT_IN, CfProfile::Churn),
|
||||
(SUBSPACE_LOGS, CfProfile::Scan),
|
||||
(SUBSPACE_TELEMETRY_SPAN, CfProfile::PointLookup),
|
||||
(SUBSPACE_TELEMETRY_METRIC, CfProfile::Scan),
|
||||
(SUBSPACE_SEARCH_INDEX, CfProfile::Scan),
|
||||
(SUBSPACE_SPAM_SAMPLES, CfProfile::Churn),
|
||||
(SUBSPACE_REGISTRY_IDX, CfProfile::Scan),
|
||||
(SUBSPACE_REGISTRY_PK, CfProfile::PointLookup),
|
||||
(SUBSPACE_DIRECTORY, CfProfile::PointLookup),
|
||||
(LEGACY_SUBSPACE_BITMAP_TEXT, CfProfile::Scan),
|
||||
(LEGACY_SUBSPACE_BITMAP_TAG, CfProfile::Scan),
|
||||
] {
|
||||
cfs.push(ColumnFamilyDescriptor::new(
|
||||
std::str::from_utf8(&[subspace]).unwrap(),
|
||||
cf_options(profile, &cache, write_buffer_size),
|
||||
));
|
||||
}
|
||||
|
||||
let mut db_opts = Options::default();
|
||||
db_opts.create_missing_column_families(true);
|
||||
db_opts.create_if_missing(true);
|
||||
db_opts.set_max_background_jobs(std::cmp::max(num_cpus::get() as i32, 3));
|
||||
db_opts.increase_parallelism(std::cmp::max(num_cpus::get() as i32, 3));
|
||||
db_opts
|
||||
.set_db_write_buffer_size((config.buffer_size as usize).max(MIN_DB_WRITE_BUFFER_SIZE));
|
||||
db_opts.set_bytes_per_sync(BYTES_PER_SYNC);
|
||||
db_opts.set_wal_bytes_per_sync(BYTES_PER_SYNC);
|
||||
|
||||
Ok(Store::RocksDb(Arc::new(RocksDbStore {
|
||||
db: OptimisticTransactionDB::open_cf_descriptors(&db_opts, idx_path, cfs)
|
||||
.map_err(|err| format!("Failed to open database: {:?}", err))?
|
||||
.into(),
|
||||
worker_pool: rayon::ThreadPoolBuilder::new()
|
||||
.num_threads(std::cmp::max(
|
||||
config
|
||||
.pool_workers
|
||||
.filter(|v| *v > 0)
|
||||
.map(|v| v as usize)
|
||||
.unwrap_or_else(num_cpus::get),
|
||||
4,
|
||||
))
|
||||
.build()
|
||||
.map_err(|err| format!("Failed to build worker pool: {:?}", err))?,
|
||||
})))
|
||||
}
|
||||
|
||||
pub async fn spawn_worker<U, V>(&self, mut f: U) -> trc::Result<V>
|
||||
where
|
||||
U: FnMut() -> trc::Result<V> + Send,
|
||||
V: Sync + Send + 'static,
|
||||
{
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
self.worker_pool.scope(|s| {
|
||||
s.spawn(|_| {
|
||||
tx.send(f()).ok();
|
||||
});
|
||||
});
|
||||
|
||||
match rx.await {
|
||||
Ok(result) => result,
|
||||
Err(err) => Err(trc::EventType::Server(trc::ServerEvent::ThreadError).reason(err)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn numeric_value_merge(
|
||||
_key: &[u8],
|
||||
value: Option<&[u8]>,
|
||||
operands: &MergeOperands,
|
||||
) -> Option<Vec<u8>> {
|
||||
let mut value = if let Some(value) = value {
|
||||
i64::from_le_bytes(value.try_into().ok()?)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
for op in operands.iter() {
|
||||
value += i64::from_le_bytes(op.try_into().ok()?);
|
||||
}
|
||||
|
||||
let mut bytes = Vec::with_capacity(std::mem::size_of::<i64>());
|
||||
bytes.extend_from_slice(&value.to_le_bytes());
|
||||
Some(bytes)
|
||||
}
|
||||
|
||||
fn cf_options(profile: CfProfile, cache: &Cache, write_buffer_size: usize) -> Options {
|
||||
let mut block_opts = BlockBasedOptions::default();
|
||||
block_opts.set_block_cache(cache);
|
||||
block_opts.set_cache_index_and_filter_blocks(true);
|
||||
block_opts.set_pin_l0_filter_and_index_blocks_in_cache(true);
|
||||
|
||||
let mut opts = Options::default();
|
||||
opts.set_write_buffer_size(write_buffer_size);
|
||||
opts.set_max_write_buffer_number(4);
|
||||
|
||||
match profile {
|
||||
CfProfile::PointLookup => {
|
||||
block_opts.set_bloom_filter(BLOOM_BITS_PER_KEY, false);
|
||||
opts.set_compression_type(DBCompressionType::Lz4);
|
||||
}
|
||||
CfProfile::Scan => {
|
||||
block_opts.set_block_size(SCAN_BLOCK_SIZE);
|
||||
opts.set_compression_type(DBCompressionType::Lz4);
|
||||
}
|
||||
CfProfile::Churn => {
|
||||
block_opts.set_bloom_filter(BLOOM_BITS_PER_KEY, false);
|
||||
opts.set_compression_type(DBCompressionType::Lz4);
|
||||
opts.set_target_file_size_base(CHURN_TARGET_FILE_SIZE);
|
||||
opts.add_compact_on_deletion_collector_factory(
|
||||
CHURN_DELETION_WINDOW,
|
||||
CHURN_DELETION_TRIGGER,
|
||||
CHURN_DELETION_RATIO,
|
||||
);
|
||||
}
|
||||
CfProfile::Queue => {
|
||||
block_opts.set_block_size(SCAN_BLOCK_SIZE);
|
||||
opts.set_compression_type(DBCompressionType::None);
|
||||
opts.set_target_file_size_base(CHURN_TARGET_FILE_SIZE);
|
||||
opts.add_compact_on_deletion_collector_factory(
|
||||
CHURN_DELETION_WINDOW,
|
||||
CHURN_DELETION_TRIGGER,
|
||||
CHURN_DELETION_RATIO,
|
||||
);
|
||||
}
|
||||
CfProfile::Counter => {
|
||||
block_opts.set_bloom_filter(BLOOM_BITS_PER_KEY, false);
|
||||
opts.set_compression_type(DBCompressionType::None);
|
||||
opts.set_merge_operator_associative("merge", numeric_value_merge);
|
||||
}
|
||||
CfProfile::Blob => {
|
||||
block_opts.set_bloom_filter(BLOOM_BITS_PER_KEY, false);
|
||||
opts.set_compression_type(DBCompressionType::None);
|
||||
}
|
||||
}
|
||||
|
||||
opts.set_block_based_table_factory(&block_opts);
|
||||
|
||||
opts
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use rocksdb::{BoundColumnFamily, MultiThreaded, OptimisticTransactionDB};
|
||||
|
||||
use crate::{SUBSPACE_BLOBS, SUBSPACE_INDEXES, SUBSPACE_LOGS};
|
||||
|
||||
pub mod blob;
|
||||
pub mod main;
|
||||
pub mod read;
|
||||
pub mod write;
|
||||
|
||||
static CF_LOGS: &str = unsafe { std::str::from_utf8_unchecked(&[SUBSPACE_LOGS]) };
|
||||
static CF_INDEXES: &str = unsafe { std::str::from_utf8_unchecked(&[SUBSPACE_INDEXES]) };
|
||||
static CF_BLOBS: &str = unsafe { std::str::from_utf8_unchecked(&[SUBSPACE_BLOBS]) };
|
||||
|
||||
pub(crate) trait CfHandle {
|
||||
fn subspace_handle(&self, subspace: u8) -> Arc<BoundColumnFamily<'_>>;
|
||||
}
|
||||
|
||||
impl CfHandle for OptimisticTransactionDB<MultiThreaded> {
|
||||
#[inline(always)]
|
||||
fn subspace_handle(&self, subspace: u8) -> Arc<BoundColumnFamily<'_>> {
|
||||
let subspace = &[subspace];
|
||||
self.cf_handle(unsafe { std::str::from_utf8_unchecked(subspace) })
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RocksDbStore {
|
||||
db: Arc<OptimisticTransactionDB<MultiThreaded>>,
|
||||
worker_pool: rayon::ThreadPool,
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn into_error(err: rocksdb::Error) -> trc::Error {
|
||||
trc::StoreEvent::RocksdbError.reason(err)
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{RocksDbStore, into_error};
|
||||
use crate::{
|
||||
Deserialize, IterateParams, Key, ValueKey, backend::rocksdb::CfHandle, write::ValueClass,
|
||||
};
|
||||
use rocksdb::ReadOptions;
|
||||
|
||||
impl RocksDbStore {
|
||||
pub(crate) async fn get_value<U>(&self, key: impl Key) -> trc::Result<Option<U>>
|
||||
where
|
||||
U: Deserialize + 'static,
|
||||
{
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
let subspace = &[key.subspace()];
|
||||
let key = key.serialize(0);
|
||||
db.get_pinned_cf(
|
||||
&db.cf_handle(unsafe { std::str::from_utf8_unchecked(subspace.as_slice()) })
|
||||
.unwrap(),
|
||||
&key,
|
||||
)
|
||||
.map_err(into_error)
|
||||
.and_then(|value| {
|
||||
if let Some(value) = value {
|
||||
U::deserialize_with_key(&key, &value).map(Some)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
let subspace = &[key.subspace()];
|
||||
let key = key.serialize(0);
|
||||
db.get_pinned_cf(
|
||||
&db.cf_handle(unsafe { std::str::from_utf8_unchecked(subspace.as_slice()) })
|
||||
.unwrap(),
|
||||
&key,
|
||||
)
|
||||
.map_err(into_error)
|
||||
.map(|value| value.is_some())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn iterate<T: Key>(
|
||||
&self,
|
||||
params: IterateParams<T>,
|
||||
mut cb: impl for<'x> FnMut(&'x [u8], &'x [u8]) -> trc::Result<bool> + Sync + Send,
|
||||
) -> trc::Result<()> {
|
||||
let db = self.db.clone();
|
||||
|
||||
self.spawn_worker(move || {
|
||||
let cf = db.subspace_handle(params.begin.subspace());
|
||||
let begin = params.begin.serialize(0);
|
||||
let end = params.end.serialize(0);
|
||||
let mut upper_bound = Vec::with_capacity(end.len() + 1);
|
||||
|
||||
upper_bound.extend_from_slice(&end);
|
||||
upper_bound.push(0u8);
|
||||
|
||||
let mut read_opts = ReadOptions::default();
|
||||
read_opts.set_iterate_lower_bound(begin.as_slice());
|
||||
read_opts.set_iterate_upper_bound(upper_bound);
|
||||
|
||||
let mut it = db.raw_iterator_cf_opt(&cf, read_opts);
|
||||
if params.ascending {
|
||||
it.seek(&begin);
|
||||
} else {
|
||||
it.seek_for_prev(&end);
|
||||
}
|
||||
|
||||
while it.valid() {
|
||||
let Some(key) = it.key() else {
|
||||
break;
|
||||
};
|
||||
let value = if params.values {
|
||||
it.value().unwrap_or_default()
|
||||
} else {
|
||||
&[][..]
|
||||
};
|
||||
|
||||
if !cb(key, value)? || params.first {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if params.ascending {
|
||||
it.next();
|
||||
} else {
|
||||
it.prev();
|
||||
}
|
||||
}
|
||||
|
||||
it.status().map_err(into_error)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn get_counter(
|
||||
&self,
|
||||
key: impl Into<ValueKey<ValueClass>> + Sync + Send,
|
||||
) -> trc::Result<i64> {
|
||||
let key = key.into();
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
let cf = self.db.subspace_handle(key.subspace());
|
||||
let key = key.serialize(0);
|
||||
|
||||
db.get_pinned_cf(&cf, &key)
|
||||
.map_err(into_error)
|
||||
.and_then(|bytes| {
|
||||
Ok(if let Some(bytes) = bytes {
|
||||
i64::from_le_bytes(bytes[..].try_into().map_err(|_| {
|
||||
trc::Error::corrupted_key(&key, (&bytes[..]).into(), trc::location!())
|
||||
})?)
|
||||
} else {
|
||||
0
|
||||
})
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{CF_INDEXES, CF_LOGS, CfHandle, RocksDbStore, into_error};
|
||||
use crate::{
|
||||
Deserialize, IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER,
|
||||
SUBSPACE_QUOTA,
|
||||
backend::deserialize_i64_le,
|
||||
write::{
|
||||
AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
|
||||
ValueClass, ValueOp,
|
||||
},
|
||||
};
|
||||
use rand::RngExt;
|
||||
use rocksdb::{
|
||||
BoundColumnFamily, ErrorKind, IteratorMode, OptimisticTransactionDB,
|
||||
OptimisticTransactionOptions, WriteOptions,
|
||||
};
|
||||
use std::{
|
||||
sync::Arc,
|
||||
thread::sleep,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
impl RocksDbStore {
|
||||
pub(crate) async fn write(&self, mut batch: Batch<'_>) -> trc::Result<AssignedIds> {
|
||||
let db = self.db.clone();
|
||||
|
||||
self.spawn_worker(move || {
|
||||
let mut txn = RocksDBTransaction {
|
||||
db: &db,
|
||||
cf_indexes: db.cf_handle(CF_INDEXES).unwrap(),
|
||||
cf_logs: db.cf_handle(CF_LOGS).unwrap(),
|
||||
txn_opts: OptimisticTransactionOptions::default(),
|
||||
batch: &mut batch,
|
||||
};
|
||||
txn.txn_opts.set_snapshot(true);
|
||||
|
||||
// Begin write
|
||||
let mut retry_count = 0;
|
||||
let start = Instant::now();
|
||||
loop {
|
||||
match txn.commit() {
|
||||
Ok(result) => {
|
||||
return Ok(result);
|
||||
}
|
||||
Err(CommitError::Internal(err)) => return Err(err),
|
||||
Err(CommitError::RocksDB(err)) => match err.kind() {
|
||||
ErrorKind::Busy | ErrorKind::MergeInProgress | ErrorKind::TryAgain
|
||||
if retry_count < MAX_COMMIT_ATTEMPTS
|
||||
&& start.elapsed() < MAX_COMMIT_TIME =>
|
||||
{
|
||||
let backoff = rand::rng().random_range(50..=300);
|
||||
sleep(Duration::from_millis(backoff));
|
||||
retry_count += 1;
|
||||
}
|
||||
_ => return Err(into_error(err)),
|
||||
},
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
db.delete_range_cf(
|
||||
&db.cf_handle(std::str::from_utf8(&[from.subspace()]).unwrap())
|
||||
.unwrap(),
|
||||
from.serialize(0),
|
||||
to.serialize(0),
|
||||
)
|
||||
.map_err(into_error)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn purge_store(&self) -> trc::Result<()> {
|
||||
let db = self.db.clone();
|
||||
self.spawn_worker(move || {
|
||||
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] {
|
||||
let cf = db
|
||||
.cf_handle(std::str::from_utf8(&[subspace]).unwrap())
|
||||
.unwrap();
|
||||
|
||||
let mut delete_keys = Vec::new();
|
||||
|
||||
for row in db.iterator_cf(&cf, IteratorMode::Start) {
|
||||
let (key, value) = row.map_err(into_error)?;
|
||||
|
||||
if i64::deserialize(&value)? == 0 {
|
||||
delete_keys.push(key);
|
||||
}
|
||||
}
|
||||
|
||||
let txn_opts = OptimisticTransactionOptions::default();
|
||||
for key in delete_keys {
|
||||
let txn = db.transaction_opt(&WriteOptions::default(), &txn_opts);
|
||||
if txn
|
||||
.get_pinned_for_update_cf(&cf, &key, true)
|
||||
.map_err(into_error)?
|
||||
.map(|value| i64::deserialize(&value).map(|v| v == 0).unwrap_or(false))
|
||||
.unwrap_or(false)
|
||||
{
|
||||
txn.delete_cf(&cf, key).map_err(into_error)?;
|
||||
txn.commit().map_err(into_error)?;
|
||||
} else {
|
||||
txn.rollback().map_err(into_error)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
struct RocksDBTransaction<'x, 'y> {
|
||||
db: &'x OptimisticTransactionDB,
|
||||
cf_indexes: Arc<BoundColumnFamily<'x>>,
|
||||
cf_logs: Arc<BoundColumnFamily<'x>>,
|
||||
txn_opts: OptimisticTransactionOptions,
|
||||
batch: &'x mut Batch<'y>,
|
||||
}
|
||||
|
||||
enum CommitError {
|
||||
Internal(trc::Error),
|
||||
RocksDB(rocksdb::Error),
|
||||
}
|
||||
|
||||
impl RocksDBTransaction<'_, '_> {
|
||||
fn commit(&mut self) -> Result<AssignedIds, CommitError> {
|
||||
let mut account_id = u32::MAX;
|
||||
let mut collection = u8::MAX;
|
||||
let mut document_id = u32::MAX;
|
||||
let mut change_id = 0u64;
|
||||
let mut result = AssignedIds::default();
|
||||
let has_changes = !self.batch.changes.is_empty();
|
||||
|
||||
let txn = self
|
||||
.db
|
||||
.transaction_opt(&WriteOptions::default(), &self.txn_opts);
|
||||
|
||||
if has_changes {
|
||||
let cf = self.db.cf_handle("n").unwrap();
|
||||
for &account_id in self.batch.changes.keys() {
|
||||
let key = ValueClass::ChangeId.serialize(account_id, 0, 0, 0);
|
||||
let change_id = txn
|
||||
.get_pinned_for_update_cf(&cf, &key, true)
|
||||
.map_err(CommitError::from)
|
||||
.and_then(|bytes| {
|
||||
if let Some(bytes) = bytes {
|
||||
deserialize_i64_le(&key, &bytes)
|
||||
.map(|v| v + 1)
|
||||
.map_err(CommitError::from)
|
||||
} else {
|
||||
Ok(1)
|
||||
}
|
||||
})?;
|
||||
txn.put_cf(&cf, &key, &change_id.to_le_bytes()[..])?;
|
||||
result.push_change_id(account_id, change_id as u64);
|
||||
}
|
||||
}
|
||||
|
||||
for op in self.batch.ops.iter_mut() {
|
||||
match op {
|
||||
Operation::AccountId {
|
||||
account_id: account_id_,
|
||||
} => {
|
||||
account_id = *account_id_;
|
||||
if has_changes {
|
||||
change_id = result.set_current_change_id(account_id)?;
|
||||
}
|
||||
}
|
||||
Operation::Collection {
|
||||
collection: collection_,
|
||||
} => {
|
||||
collection = u8::from(*collection_);
|
||||
}
|
||||
Operation::DocumentId {
|
||||
document_id: document_id_,
|
||||
} => {
|
||||
document_id = *document_id_;
|
||||
}
|
||||
Operation::Value { class, op } => {
|
||||
let key = class.serialize(account_id, collection, document_id, 0);
|
||||
let cf = self.db.subspace_handle(class.subspace(collection));
|
||||
|
||||
match op {
|
||||
ValueOp::Set(value) => {
|
||||
txn.put_cf(&cf, &key, value)?;
|
||||
}
|
||||
ValueOp::SetFnc(set_op) => {
|
||||
let value = (set_op.fnc)(&set_op.params, &result)?;
|
||||
|
||||
txn.put_cf(&cf, &key, value)?;
|
||||
}
|
||||
ValueOp::MergeFnc(merge_op) => {
|
||||
let merge_result = (merge_op.fnc)(
|
||||
&merge_op.params,
|
||||
&result,
|
||||
txn.get_pinned_for_update_cf(&cf, &key, true)?.as_deref(),
|
||||
)?;
|
||||
|
||||
match merge_result {
|
||||
MergeResult::Update(value) => {
|
||||
txn.put_cf(&cf, &key, value)?;
|
||||
}
|
||||
MergeResult::Delete => {
|
||||
txn.delete_cf(&cf, &key)?;
|
||||
}
|
||||
MergeResult::Skip => (),
|
||||
}
|
||||
}
|
||||
ValueOp::AtomicAdd(by) => {
|
||||
txn.merge_cf(&cf, &key, &by.to_le_bytes()[..])?;
|
||||
}
|
||||
ValueOp::AddAndGet(by) => {
|
||||
let num = txn
|
||||
.get_pinned_for_update_cf(&cf, &key, true)
|
||||
.map_err(CommitError::from)
|
||||
.and_then(|bytes| {
|
||||
if let Some(bytes) = bytes {
|
||||
deserialize_i64_le(&key, &bytes)
|
||||
.map(|v| v + *by)
|
||||
.map_err(CommitError::from)
|
||||
} else {
|
||||
Ok(*by)
|
||||
}
|
||||
})?;
|
||||
txn.put_cf(&cf, &key, &num.to_le_bytes()[..])?;
|
||||
result.push_counter_id(num);
|
||||
}
|
||||
ValueOp::Clear => {
|
||||
txn.delete_cf(&cf, &key)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Operation::Index { field, key, set } => {
|
||||
let key = IndexKey {
|
||||
account_id,
|
||||
collection,
|
||||
document_id,
|
||||
field: *field,
|
||||
key: &*key,
|
||||
}
|
||||
.serialize(0);
|
||||
|
||||
if *set {
|
||||
txn.put_cf(&self.cf_indexes, &key, [])?;
|
||||
} else {
|
||||
txn.delete_cf(&self.cf_indexes, &key)?;
|
||||
}
|
||||
}
|
||||
Operation::Log { collection, set } => {
|
||||
let key = LogKey {
|
||||
account_id,
|
||||
collection: u8::from(*collection),
|
||||
change_id,
|
||||
}
|
||||
.serialize(0);
|
||||
|
||||
txn.put_cf(&self.cf_logs, &key, set)?;
|
||||
}
|
||||
Operation::AssertValue {
|
||||
class,
|
||||
assert_value,
|
||||
} => {
|
||||
let key = class.serialize(account_id, collection, document_id, 0);
|
||||
let cf = self.db.subspace_handle(class.subspace(collection));
|
||||
|
||||
let matches = txn
|
||||
.get_pinned_for_update_cf(&cf, &key, true)?
|
||||
.map(|value| assert_value.matches(&value))
|
||||
.unwrap_or_else(|| assert_value.is_none());
|
||||
|
||||
if !matches {
|
||||
txn.rollback()?;
|
||||
return Err(CommitError::Internal(
|
||||
trc::StoreEvent::AssertValueFailed.into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
txn.commit().map(|_| result).map_err(Into::into)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<rocksdb::Error> for CommitError {
|
||||
fn from(err: rocksdb::Error) -> Self {
|
||||
CommitError::RocksDB(err)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<trc::Error> for CommitError {
|
||||
fn from(err: trc::Error) -> Self {
|
||||
CommitError::Internal(err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user