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,156 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{FdbStore, MAX_VALUE_SIZE};
|
||||
use crate::{
|
||||
IterateParams, SUBSPACE_BLOBS,
|
||||
backend::foundationdb::into_error,
|
||||
write::{AnyKey, key::KeySerializer},
|
||||
};
|
||||
use std::ops::Range;
|
||||
use trc::AddContext;
|
||||
use types::blob_hash::BLOB_HASH_LEN;
|
||||
|
||||
impl FdbStore {
|
||||
pub(crate) async fn get_blob(
|
||||
&self,
|
||||
key: &[u8],
|
||||
range: Range<usize>,
|
||||
) -> trc::Result<Option<Vec<u8>>> {
|
||||
let block_start = range.start / MAX_VALUE_SIZE;
|
||||
let bytes_start = range.start % MAX_VALUE_SIZE;
|
||||
let block_end = (range.end / MAX_VALUE_SIZE) + 1;
|
||||
|
||||
let begin = KeySerializer::new(key.len() + 2)
|
||||
.write(key)
|
||||
.write(block_start as u16)
|
||||
.finalize();
|
||||
let end = KeySerializer::new(key.len() + 2)
|
||||
.write(key)
|
||||
.write(block_end as u16)
|
||||
.finalize();
|
||||
let key_len = begin.len();
|
||||
|
||||
let mut blob_data: Option<Vec<u8>> = None;
|
||||
let blob_range = range.end - range.start;
|
||||
|
||||
self.iterate(
|
||||
IterateParams::new(
|
||||
AnyKey {
|
||||
subspace: SUBSPACE_BLOBS,
|
||||
key: begin,
|
||||
},
|
||||
AnyKey {
|
||||
subspace: SUBSPACE_BLOBS,
|
||||
key: end,
|
||||
},
|
||||
),
|
||||
|key, value| {
|
||||
if key.len() == key_len {
|
||||
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 {
|
||||
return Ok(false);
|
||||
}
|
||||
} else {
|
||||
let blob_size = if blob_range <= (5 * (1 << 20)) {
|
||||
blob_range
|
||||
} else if value.len() == MAX_VALUE_SIZE {
|
||||
MAX_VALUE_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(&[]),
|
||||
);
|
||||
let is_done = blob_data_.len() == blob_range;
|
||||
blob_data = blob_data_.into();
|
||||
if is_done {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(true)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.caused_by(trc::location!())?;
|
||||
|
||||
Ok(blob_data)
|
||||
}
|
||||
|
||||
pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> trc::Result<()> {
|
||||
const N_CHUNKS: usize = (1 << 5) - 1;
|
||||
let last_chunk = std::cmp::max(
|
||||
(data.len() / MAX_VALUE_SIZE)
|
||||
+ if !data.len().is_multiple_of(MAX_VALUE_SIZE) {
|
||||
1
|
||||
} else {
|
||||
0
|
||||
},
|
||||
1,
|
||||
) - 1;
|
||||
let mut trx = self.db.create_trx().map_err(into_error)?;
|
||||
|
||||
for (chunk_pos, chunk_bytes) in data.chunks(MAX_VALUE_SIZE).enumerate() {
|
||||
trx.set(
|
||||
&KeySerializer::new(key.len() + 3)
|
||||
.write(SUBSPACE_BLOBS)
|
||||
.write(key)
|
||||
.write(chunk_pos as u16)
|
||||
.finalize(),
|
||||
chunk_bytes,
|
||||
);
|
||||
if chunk_pos == last_chunk || (chunk_pos > 0 && chunk_pos % N_CHUNKS == 0) {
|
||||
self.commit(trx, false).await?;
|
||||
if chunk_pos < last_chunk {
|
||||
trx = self.db.create_trx().map_err(into_error)?;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_blob(&self, key: &[u8]) -> trc::Result<bool> {
|
||||
if key.len() < BLOB_HASH_LEN {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let trx = self.db.create_trx().map_err(into_error)?;
|
||||
trx.clear_range(
|
||||
&KeySerializer::new(key.len() + 3)
|
||||
.write(SUBSPACE_BLOBS)
|
||||
.write(key)
|
||||
.write(0u16)
|
||||
.finalize(),
|
||||
&KeySerializer::new(key.len() + 3)
|
||||
.write(SUBSPACE_BLOBS)
|
||||
.write(key)
|
||||
.write(u16::MAX)
|
||||
.finalize(),
|
||||
);
|
||||
|
||||
self.commit(trx, false).await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::FdbStore;
|
||||
use crate::Store;
|
||||
use foundationdb::{Database, api, api::NetworkAutoStop, options::DatabaseOption};
|
||||
use parking_lot::Mutex;
|
||||
use registry::schema::structs;
|
||||
use std::sync::Arc;
|
||||
|
||||
static FDB_NETWORK: Mutex<Option<NetworkAutoStop>> = Mutex::new(None);
|
||||
|
||||
impl FdbStore {
|
||||
pub async fn open(config: structs::FoundationDbStore) -> Result<Store, String> {
|
||||
{
|
||||
let mut guard = FDB_NETWORK.lock();
|
||||
if guard.is_none() {
|
||||
let network = unsafe {
|
||||
api::FdbApiBuilder::default()
|
||||
.build()
|
||||
.map_err(|err| format!("Failed to boot FoundationDB: {err:?}"))?
|
||||
.boot()
|
||||
.map_err(|err| format!("Failed to boot FoundationDB: {err:?}"))?
|
||||
};
|
||||
*guard = Some(network);
|
||||
}
|
||||
}
|
||||
|
||||
let db = Database::new(config.cluster_file.as_deref())
|
||||
.map_err(|err| format!("Failed to create FoundationDB database: {err:?}"))?;
|
||||
|
||||
if let Some(value) = config.transaction_timeout {
|
||||
db.set_option(DatabaseOption::TransactionTimeout(
|
||||
value.into_inner().as_millis() as i32,
|
||||
))
|
||||
.map_err(|err| format!("Failed to set option: {err:?}"))?;
|
||||
}
|
||||
if let Some(value) = config.transaction_retry_limit {
|
||||
db.set_option(DatabaseOption::TransactionRetryLimit(value as i32))
|
||||
.map_err(|err| format!("Failed to set option: {err:?}"))?;
|
||||
}
|
||||
if let Some(value) = config.transaction_retry_delay {
|
||||
db.set_option(DatabaseOption::TransactionMaxRetryDelay(
|
||||
value.into_inner().as_millis() as i32,
|
||||
))
|
||||
.map_err(|err| format!("Failed to set option: {err:?}"))?;
|
||||
}
|
||||
if let Some(value) = config.machine_id {
|
||||
db.set_option(DatabaseOption::MachineId(value))
|
||||
.map_err(|err| format!("Failed to set option: {err:?}"))?;
|
||||
}
|
||||
if let Some(value) = config.datacenter_id {
|
||||
db.set_option(DatabaseOption::DatacenterId(value))
|
||||
.map_err(|err| format!("Failed to set option: {err:?}"))?;
|
||||
}
|
||||
|
||||
Ok(Store::FoundationDb(Arc::new(Self {
|
||||
db,
|
||||
version: Default::default(),
|
||||
})))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use foundationdb::{Database, FdbError};
|
||||
use std::{
|
||||
sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
pub mod blob;
|
||||
pub mod main;
|
||||
pub mod read;
|
||||
pub mod write;
|
||||
|
||||
const MAX_VALUE_SIZE: usize = 100000;
|
||||
|
||||
const REFRESH_READ_VERSION_AFTER: Duration = Duration::from_secs(1);
|
||||
const MAX_READ_VERSION_AGE: Duration = Duration::from_secs(4);
|
||||
|
||||
pub struct FdbStore {
|
||||
db: Database,
|
||||
version: ReadVersion,
|
||||
}
|
||||
|
||||
pub(crate) struct ReadVersion {
|
||||
base: Instant,
|
||||
version: AtomicI64,
|
||||
obtained: AtomicU64,
|
||||
refreshing: AtomicBool,
|
||||
}
|
||||
|
||||
impl ReadVersion {
|
||||
fn now(&self) -> u64 {
|
||||
self.base.elapsed().as_nanos() as u64
|
||||
}
|
||||
|
||||
fn current(&self) -> i64 {
|
||||
self.version.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
fn age(&self) -> u64 {
|
||||
self.now()
|
||||
.saturating_sub(self.obtained.load(Ordering::Acquire))
|
||||
}
|
||||
|
||||
fn store_max(&self, version: i64) {
|
||||
let mut current = self.version.load(Ordering::Relaxed);
|
||||
while version > current {
|
||||
match self.version.compare_exchange_weak(
|
||||
current,
|
||||
version,
|
||||
Ordering::Release,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => break,
|
||||
Err(actual) => current = actual,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn refreshed(&self, version: i64) {
|
||||
self.store_max(version);
|
||||
self.obtained.store(self.now(), Ordering::Release);
|
||||
}
|
||||
|
||||
fn raise_floor(&self, version: i64) {
|
||||
self.store_max(version);
|
||||
}
|
||||
|
||||
fn expire(&self) {
|
||||
self.obtained.store(0, Ordering::Release);
|
||||
}
|
||||
|
||||
fn try_begin_refresh(&self) -> Option<RefreshGuard<'_>> {
|
||||
if self
|
||||
.refreshing
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
|
||||
.is_ok()
|
||||
{
|
||||
Some(RefreshGuard(&self.refreshing))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ReadVersion {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
base: Instant::now(),
|
||||
version: AtomicI64::new(0),
|
||||
obtained: AtomicU64::new(0),
|
||||
refreshing: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct RefreshGuard<'a>(&'a AtomicBool);
|
||||
|
||||
impl Drop for RefreshGuard<'_> {
|
||||
fn drop(&mut self) {
|
||||
self.0.store(false, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn into_error(error: FdbError) -> trc::Error {
|
||||
trc::StoreEvent::FoundationdbError
|
||||
.reason(error.message())
|
||||
.ctx(trc::Key::Code, error.code())
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{
|
||||
FdbStore, MAX_READ_VERSION_AGE, MAX_VALUE_SIZE, REFRESH_READ_VERSION_AFTER, into_error,
|
||||
};
|
||||
use crate::{
|
||||
Deserialize, IterateParams, Key, ValueKey, WITH_SUBSPACE,
|
||||
backend::deserialize_i64_le,
|
||||
write::{MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, ValueClass, key::KeySerializer},
|
||||
};
|
||||
use foundationdb::{
|
||||
FdbError, KeySelector, RangeOption, Transaction,
|
||||
future::FdbSlice,
|
||||
options::{self},
|
||||
};
|
||||
use futures::TryStreamExt;
|
||||
use std::time::Instant;
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) enum ChunkedValue {
|
||||
Single(FdbSlice),
|
||||
Chunked { n_chunks: u8, bytes: Vec<u8> },
|
||||
None,
|
||||
}
|
||||
|
||||
struct ChunkedValueCollector {
|
||||
key: Vec<u8>,
|
||||
bytes: Vec<u8>,
|
||||
}
|
||||
|
||||
impl FdbStore {
|
||||
pub(crate) async fn get_value<U>(&self, key: impl Key) -> trc::Result<Option<U>>
|
||||
where
|
||||
U: Deserialize,
|
||||
{
|
||||
let key = key.serialize(WITH_SUBSPACE);
|
||||
let mut retry_count = 0;
|
||||
let start = Instant::now();
|
||||
|
||||
loop {
|
||||
let trx = self.read_trx().await?;
|
||||
|
||||
match read_chunked_value(&key, &trx, true).await {
|
||||
Ok(ChunkedValue::Single(bytes)) => {
|
||||
return U::deserialize_with_key(key.get(1..).unwrap_or_default(), &bytes)
|
||||
.map(Some);
|
||||
}
|
||||
Ok(ChunkedValue::Chunked { bytes, .. }) => {
|
||||
return U::deserialize_owned_with_key(key.get(1..).unwrap_or_default(), bytes)
|
||||
.map(Some);
|
||||
}
|
||||
Ok(ChunkedValue::None) => return Ok(None),
|
||||
Err(err) => {
|
||||
self.on_read_error(trx, err, &mut retry_count, start)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn key_exists(&self, key: impl Key) -> trc::Result<bool> {
|
||||
let key = key.serialize(WITH_SUBSPACE);
|
||||
let mut retry_count = 0;
|
||||
let start = Instant::now();
|
||||
|
||||
loop {
|
||||
let trx = self.read_trx().await?;
|
||||
|
||||
match read_chunked_value(&key, &trx, true).await {
|
||||
Ok(ChunkedValue::Single(_) | ChunkedValue::Chunked { .. }) => return Ok(true),
|
||||
Ok(ChunkedValue::None) => return Ok(false),
|
||||
Err(err) => {
|
||||
self.on_read_error(trx, err, &mut retry_count, start)
|
||||
.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 begin = params.begin.serialize(WITH_SUBSPACE);
|
||||
let end = params.end.serialize(WITH_SUBSPACE);
|
||||
let mut retry_count = 0;
|
||||
let start = Instant::now();
|
||||
|
||||
if !params.first {
|
||||
let mut last_key = vec![];
|
||||
let mut chunked_key: Option<ChunkedValueCollector> = None;
|
||||
|
||||
'outer: loop {
|
||||
let begin_selector = if last_key.is_empty() {
|
||||
KeySelector::first_greater_or_equal(&begin)
|
||||
} else {
|
||||
KeySelector::first_greater_than(&last_key)
|
||||
};
|
||||
|
||||
let trx = self.read_trx().await?;
|
||||
let mut values = trx.get_ranges(
|
||||
RangeOption {
|
||||
begin: begin_selector,
|
||||
end: KeySelector::first_greater_than(&end),
|
||||
mode: options::StreamingMode::WantAll,
|
||||
reverse: !params.ascending,
|
||||
..Default::default()
|
||||
},
|
||||
true,
|
||||
);
|
||||
|
||||
let mut last_key_ = vec![];
|
||||
loop {
|
||||
match values.try_next().await {
|
||||
Ok(Some(values)) => {
|
||||
let mut key = &[] as &[u8];
|
||||
for value in values.iter() {
|
||||
key = value.key();
|
||||
|
||||
// Check whether we are collecting a chunked value
|
||||
let cb_key = key.get(1..).unwrap_or_default();
|
||||
let cb_value = value.value();
|
||||
|
||||
if let Some(chunk) = &mut chunked_key {
|
||||
if chunk.key.len() + 1 == cb_key.len()
|
||||
&& cb_key[..chunk.key.len()] == chunk.key[..]
|
||||
{
|
||||
// This is a chunk of the current value
|
||||
if params.values {
|
||||
chunk.bytes.extend_from_slice(cb_value);
|
||||
}
|
||||
continue;
|
||||
} else {
|
||||
// Return collected chunked value
|
||||
if !cb(&chunk.key, &chunk.bytes)? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Reset collector
|
||||
chunked_key = None;
|
||||
}
|
||||
}
|
||||
|
||||
if cb_value.len() < MAX_VALUE_SIZE {
|
||||
if !cb(cb_key, cb_value)? {
|
||||
return Ok(());
|
||||
}
|
||||
} else {
|
||||
// Start collecting chunked value
|
||||
chunked_key = Some(ChunkedValueCollector {
|
||||
key: cb_key.to_vec(),
|
||||
bytes: if params.values {
|
||||
cb_value.to_vec()
|
||||
} else {
|
||||
Vec::new()
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
if values.more() {
|
||||
last_key_ = key.to_vec();
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
// Return any chunked value collected
|
||||
if let Some(chunked_key) = chunked_key.take() {
|
||||
cb(&chunked_key.key, &chunked_key.bytes)?;
|
||||
}
|
||||
|
||||
break 'outer;
|
||||
}
|
||||
Err(e) => {
|
||||
drop(values);
|
||||
if e.code() == 1007 && !last_key_.is_empty() {
|
||||
// Transaction is too old to perform reads or be committed
|
||||
last_key = last_key_;
|
||||
continue 'outer;
|
||||
} else if e.is_retryable()
|
||||
&& retry_count < MAX_COMMIT_ATTEMPTS
|
||||
&& start.elapsed() < MAX_COMMIT_TIME
|
||||
{
|
||||
// Transient error such as a cached read version ahead of lagging
|
||||
// storage servers (code 1009); resume from the last key read,
|
||||
// refresh the read version and back off before retrying.
|
||||
if !last_key_.is_empty() {
|
||||
last_key = last_key_;
|
||||
}
|
||||
self.version.expire();
|
||||
trx.on_error(e).await.map_err(into_error)?;
|
||||
retry_count += 1;
|
||||
continue 'outer;
|
||||
} else {
|
||||
return Err(into_error(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
loop {
|
||||
let trx = self.read_trx().await?;
|
||||
let mut values = trx.get_ranges_keyvalues(
|
||||
RangeOption {
|
||||
begin: KeySelector::first_greater_or_equal(&begin),
|
||||
end: KeySelector::first_greater_than(&end),
|
||||
mode: options::StreamingMode::Small,
|
||||
reverse: !params.ascending,
|
||||
..Default::default()
|
||||
},
|
||||
true,
|
||||
);
|
||||
|
||||
match values.try_next().await {
|
||||
Ok(Some(value)) => {
|
||||
cb(value.key().get(1..).unwrap_or_default(), value.value())?;
|
||||
break;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) => {
|
||||
drop(values);
|
||||
self.on_read_error(trx, e, &mut retry_count, start).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn get_counter(
|
||||
&self,
|
||||
key: impl Into<ValueKey<ValueClass>> + Sync + Send,
|
||||
) -> trc::Result<i64> {
|
||||
let key = key.into().serialize(WITH_SUBSPACE);
|
||||
let mut retry_count = 0;
|
||||
let start = Instant::now();
|
||||
|
||||
loop {
|
||||
let trx = self.read_trx().await?;
|
||||
match trx.get(&key, true).await {
|
||||
Ok(Some(bytes)) => return deserialize_i64_le(&key, &bytes),
|
||||
Ok(None) => return Ok(0),
|
||||
Err(e) => {
|
||||
self.on_read_error(trx, e, &mut retry_count, start).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn on_read_error(
|
||||
&self,
|
||||
trx: Transaction,
|
||||
err: FdbError,
|
||||
retry_count: &mut u32,
|
||||
start: Instant,
|
||||
) -> trc::Result<()> {
|
||||
if err.is_retryable()
|
||||
&& *retry_count < MAX_COMMIT_ATTEMPTS
|
||||
&& start.elapsed() < MAX_COMMIT_TIME
|
||||
{
|
||||
// The cached read version may be ahead of lagging storage servers under heavy write
|
||||
// load (code 1009); expire it so the retry obtains a fresh read version, then let
|
||||
// FoundationDB back off before retrying.
|
||||
self.version.expire();
|
||||
trx.on_error(err).await.map_err(into_error)?;
|
||||
*retry_count += 1;
|
||||
Ok(())
|
||||
} else {
|
||||
Err(into_error(err))
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_trx(&self) -> trc::Result<Transaction> {
|
||||
let trx = self.db.create_trx().map_err(into_error)?;
|
||||
let version = self.version.current();
|
||||
let age = self.version.age();
|
||||
|
||||
if version != 0 && age < MAX_READ_VERSION_AGE.as_nanos() as u64 {
|
||||
if age >= REFRESH_READ_VERSION_AFTER.as_nanos() as u64
|
||||
&& let Some(_guard) = self.version.try_begin_refresh()
|
||||
{
|
||||
let read_version = trx.get_read_version().await.map_err(into_error)?;
|
||||
self.version.refreshed(read_version);
|
||||
} else {
|
||||
trx.set_read_version(version);
|
||||
}
|
||||
} else {
|
||||
let read_version = trx.get_read_version().await.map_err(into_error)?;
|
||||
self.version.refreshed(read_version);
|
||||
}
|
||||
|
||||
Ok(trx)
|
||||
}
|
||||
|
||||
pub(crate) fn invalidate_read_snapshot(&self) {
|
||||
self.version.expire();
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_chunked_value(
|
||||
key: &[u8],
|
||||
trx: &Transaction,
|
||||
snapshot: bool,
|
||||
) -> Result<ChunkedValue, FdbError> {
|
||||
if let Some(bytes) = trx.get(key, snapshot).await? {
|
||||
if bytes.len() < MAX_VALUE_SIZE {
|
||||
Ok(ChunkedValue::Single(bytes))
|
||||
} else {
|
||||
let mut value = Vec::with_capacity(bytes.len() * 2);
|
||||
value.extend_from_slice(&bytes);
|
||||
let mut key = KeySerializer::new(key.len() + 1)
|
||||
.write(key)
|
||||
.write(0u8)
|
||||
.finalize();
|
||||
|
||||
while let Some(bytes) = trx.get(&key, snapshot).await? {
|
||||
value.extend_from_slice(&bytes);
|
||||
*key.last_mut().unwrap() += 1;
|
||||
}
|
||||
|
||||
Ok(ChunkedValue::Chunked {
|
||||
bytes: value,
|
||||
n_chunks: *key.last().unwrap(),
|
||||
})
|
||||
}
|
||||
} else {
|
||||
Ok(ChunkedValue::None)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,436 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{
|
||||
FdbStore, MAX_VALUE_SIZE, into_error,
|
||||
read::{ChunkedValue, read_chunked_value},
|
||||
};
|
||||
use crate::{
|
||||
backend::deserialize_i64_le,
|
||||
write::{
|
||||
AssignedIds, Batch, IndexPropertyClass, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult,
|
||||
Operation, QueueClass, RegistryClass, SearchIndexType, TaskQueueClass, TelemetryClass,
|
||||
ValueClass, ValueOp, key::KeySerializer,
|
||||
},
|
||||
*,
|
||||
};
|
||||
use foundationdb::{
|
||||
FdbError, KeySelector, RangeOption, Transaction,
|
||||
options::{self, MutationType},
|
||||
};
|
||||
use futures::TryStreamExt;
|
||||
use rand::RngExt;
|
||||
use std::{
|
||||
borrow::Cow,
|
||||
cmp::Ordering,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
use trc::AddContext;
|
||||
|
||||
impl FdbStore {
|
||||
pub(crate) async fn write(&self, batch: Batch<'_>) -> trc::Result<AssignedIds> {
|
||||
let start = Instant::now();
|
||||
let mut retry_count = 0;
|
||||
let has_changes = !batch.changes.is_empty();
|
||||
|
||||
loop {
|
||||
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 trx = self.db.create_trx().map_err(into_error)?;
|
||||
|
||||
if has_changes {
|
||||
for &account_id in batch.changes.keys() {
|
||||
debug_assert!(account_id != u32::MAX);
|
||||
let key = ValueClass::ChangeId.serialize(account_id, 0, 0, WITH_SUBSPACE);
|
||||
let change_id =
|
||||
if let Some(bytes) = trx.get(&key, false).await.map_err(into_error)? {
|
||||
deserialize_i64_le(&key, &bytes)? + 1
|
||||
} else {
|
||||
1
|
||||
};
|
||||
trx.set(&key, &change_id.to_le_bytes()[..]);
|
||||
result.push_change_id(account_id, change_id as u64);
|
||||
}
|
||||
}
|
||||
|
||||
for op in 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 mut key =
|
||||
class.serialize(account_id, collection, document_id, WITH_SUBSPACE);
|
||||
|
||||
match op {
|
||||
ValueOp::Set(value) => {
|
||||
if !chunk_value(&trx, &mut key, value, class, None).await {
|
||||
trx.cancel();
|
||||
return Err(trc::StoreEvent::FoundationdbError
|
||||
.ctx(trc::Key::Reason, "Value is too large"));
|
||||
}
|
||||
}
|
||||
ValueOp::SetFnc(set_op) => {
|
||||
let value = (set_op.fnc)(&set_op.params, &result)?;
|
||||
if !chunk_value(&trx, &mut key, &value, class, None).await {
|
||||
trx.cancel();
|
||||
return Err(trc::StoreEvent::FoundationdbError
|
||||
.ctx(trc::Key::Reason, "Value is too large"));
|
||||
}
|
||||
}
|
||||
ValueOp::MergeFnc(merge_op) => {
|
||||
let (merge_result, prev_num_chunks) =
|
||||
match read_chunked_value(&key, &trx, false)
|
||||
.await
|
||||
.map_err(into_error)
|
||||
.caused_by(trc::location!())?
|
||||
{
|
||||
ChunkedValue::Single(slice) => (
|
||||
(merge_op.fnc)(
|
||||
&merge_op.params,
|
||||
&result,
|
||||
Some(slice.as_ref()),
|
||||
)?,
|
||||
1,
|
||||
),
|
||||
ChunkedValue::Chunked { bytes, n_chunks } => (
|
||||
(merge_op.fnc)(
|
||||
&merge_op.params,
|
||||
&result,
|
||||
Some(bytes.as_ref()),
|
||||
)?,
|
||||
n_chunks as usize + 1,
|
||||
),
|
||||
ChunkedValue::None => {
|
||||
((merge_op.fnc)(&merge_op.params, &result, None)?, 0)
|
||||
}
|
||||
};
|
||||
|
||||
match merge_result {
|
||||
MergeResult::Update(value) => {
|
||||
if !chunk_value(
|
||||
&trx,
|
||||
&mut key,
|
||||
&value,
|
||||
class,
|
||||
Some(prev_num_chunks),
|
||||
)
|
||||
.await
|
||||
{
|
||||
trx.cancel();
|
||||
return Err(trc::StoreEvent::FoundationdbError
|
||||
.ctx(trc::Key::Reason, "Value is too large"));
|
||||
}
|
||||
}
|
||||
MergeResult::Delete => {
|
||||
if prev_num_chunks > 1 {
|
||||
clear_chunks(&trx, &key, None).await;
|
||||
} else {
|
||||
trx.clear(&key);
|
||||
}
|
||||
}
|
||||
MergeResult::Skip => (),
|
||||
}
|
||||
}
|
||||
ValueOp::AtomicAdd(by) => {
|
||||
trx.atomic_op(&key, &by.to_le_bytes()[..], MutationType::Add);
|
||||
}
|
||||
ValueOp::AddAndGet(by) => {
|
||||
let num = if let Some(bytes) =
|
||||
trx.get(&key, false).await.map_err(into_error)?
|
||||
{
|
||||
deserialize_i64_le(&key, &bytes)? + *by
|
||||
} else {
|
||||
*by
|
||||
};
|
||||
trx.set(&key, &num.to_le_bytes()[..]);
|
||||
result.push_counter_id(num);
|
||||
}
|
||||
ValueOp::Clear => {
|
||||
if is_chunked_value(key[0], class) {
|
||||
clear_chunks(&trx, &key, None).await;
|
||||
} else {
|
||||
trx.clear(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Operation::Index { field, key, set } => {
|
||||
let key = IndexKey {
|
||||
account_id,
|
||||
collection,
|
||||
document_id,
|
||||
field: *field,
|
||||
key: &*key,
|
||||
}
|
||||
.serialize(WITH_SUBSPACE);
|
||||
|
||||
if *set {
|
||||
trx.set(&key, &[]);
|
||||
} else {
|
||||
trx.clear(&key);
|
||||
}
|
||||
}
|
||||
Operation::Log { collection, set } => {
|
||||
let key = LogKey {
|
||||
account_id,
|
||||
collection: u8::from(*collection),
|
||||
change_id,
|
||||
}
|
||||
.serialize(WITH_SUBSPACE);
|
||||
|
||||
trx.set(&key, set);
|
||||
}
|
||||
Operation::AssertValue {
|
||||
class,
|
||||
assert_value,
|
||||
} => {
|
||||
let key =
|
||||
class.serialize(account_id, collection, document_id, WITH_SUBSPACE);
|
||||
|
||||
let matches = match read_chunked_value(&key, &trx, false).await {
|
||||
Ok(ChunkedValue::Single(bytes)) => assert_value.matches(bytes.as_ref()),
|
||||
Ok(ChunkedValue::Chunked { bytes, .. }) => {
|
||||
assert_value.matches(bytes.as_ref())
|
||||
}
|
||||
Ok(ChunkedValue::None) => assert_value.is_none(),
|
||||
Err(_) => false,
|
||||
};
|
||||
|
||||
if !matches {
|
||||
trx.cancel();
|
||||
return Err(trc::StoreEvent::AssertValueFailed.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if self
|
||||
.commit(
|
||||
trx,
|
||||
retry_count < MAX_COMMIT_ATTEMPTS && start.elapsed() < MAX_COMMIT_TIME,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(result);
|
||||
} else {
|
||||
let backoff = rand::rng().random_range(50..=100);
|
||||
tokio::time::sleep(Duration::from_millis(backoff)).await;
|
||||
retry_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn commit(&self, trx: Transaction, will_retry: bool) -> trc::Result<bool> {
|
||||
match trx.commit().await {
|
||||
Ok(result) => {
|
||||
let commit_version = result.committed_version().map_err(into_error)?;
|
||||
self.version.raise_floor(commit_version);
|
||||
Ok(true)
|
||||
}
|
||||
Err(err) => {
|
||||
if will_retry {
|
||||
err.on_error().await.map_err(into_error)?;
|
||||
Ok(false)
|
||||
} else {
|
||||
Err(into_error(FdbError::from(err)))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn purge_store(&self) -> trc::Result<()> {
|
||||
// Obtain all zero counters
|
||||
let mut delete_keys = Vec::new();
|
||||
for subspace in [SUBSPACE_COUNTER, SUBSPACE_QUOTA, SUBSPACE_IN_MEMORY_COUNTER] {
|
||||
let trx = self.db.create_trx().map_err(into_error)?;
|
||||
let from_key = [subspace, 0u8];
|
||||
let to_key = [subspace, u8::MAX, u8::MAX, u8::MAX, u8::MAX, u8::MAX];
|
||||
|
||||
let mut values = trx.get_ranges_keyvalues(
|
||||
RangeOption {
|
||||
begin: KeySelector::first_greater_or_equal(&from_key[..]),
|
||||
end: KeySelector::first_greater_or_equal(&to_key[..]),
|
||||
mode: options::StreamingMode::WantAll,
|
||||
reverse: false,
|
||||
..Default::default()
|
||||
},
|
||||
true,
|
||||
);
|
||||
|
||||
while let Some(value) = values.try_next().await.map_err(into_error)? {
|
||||
if value.value().iter().all(|byte| *byte == 0) {
|
||||
delete_keys.push(value.key().to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if delete_keys.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Delete keys
|
||||
let integer = 0i64.to_le_bytes();
|
||||
for chunk in delete_keys.chunks(1024) {
|
||||
let mut retry_count = 0;
|
||||
loop {
|
||||
let trx = self.db.create_trx().map_err(into_error)?;
|
||||
for key in chunk {
|
||||
trx.atomic_op(key, &integer, MutationType::CompareAndClear);
|
||||
}
|
||||
|
||||
if self.commit(trx, retry_count < MAX_COMMIT_ATTEMPTS).await? {
|
||||
break;
|
||||
} else {
|
||||
retry_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
|
||||
let from = from.serialize(WITH_SUBSPACE);
|
||||
let to = to.serialize(WITH_SUBSPACE);
|
||||
|
||||
let trx = self.db.create_trx().map_err(into_error)?;
|
||||
trx.clear_range(&from, &to);
|
||||
self.commit(trx, false).await.map(|_| ())
|
||||
}
|
||||
}
|
||||
|
||||
fn is_chunked_subspace(subspace: u8) -> bool {
|
||||
matches!(
|
||||
subspace,
|
||||
crate::SUBSPACE_PROPERTY
|
||||
| crate::SUBSPACE_SEARCH_INDEX
|
||||
| crate::SUBSPACE_QUEUE_MESSAGE
|
||||
| crate::SUBSPACE_TASK_QUEUE
|
||||
| crate::SUBSPACE_DIRECTORY
|
||||
| crate::SUBSPACE_REGISTRY
|
||||
| crate::SUBSPACE_DELETED_ITEMS
|
||||
| crate::SUBSPACE_SPAM_SAMPLES
|
||||
| crate::SUBSPACE_REPORT_IN
|
||||
| crate::SUBSPACE_REPORT_OUT
|
||||
| crate::SUBSPACE_TELEMETRY_SPAN
|
||||
)
|
||||
}
|
||||
|
||||
fn is_chunked_value(subspace: u8, class: &ValueClass) -> bool {
|
||||
is_chunked_subspace(subspace)
|
||||
&& match class {
|
||||
ValueClass::Property(_)
|
||||
| ValueClass::IndexProperty(IndexPropertyClass::Hash { .. })
|
||||
| ValueClass::Registry(RegistryClass::Item { .. })
|
||||
| ValueClass::Queue(QueueClass::Message(_))
|
||||
| ValueClass::TaskQueue(TaskQueueClass::Task { .. })
|
||||
| ValueClass::Telemetry(TelemetryClass::Span(_)) => true,
|
||||
ValueClass::SearchIndex(index) => matches!(index.typ, SearchIndexType::Document),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
async fn clear_chunks(trx: &Transaction, key: &[u8], from_chunk: Option<u8>) {
|
||||
let to = KeySerializer::new(key.len() + 1)
|
||||
.write(key)
|
||||
.write(u8::MAX)
|
||||
.finalize();
|
||||
let from = match from_chunk {
|
||||
Some(from_chunk) => Cow::Owned(
|
||||
KeySerializer::new(key.len() + 1)
|
||||
.write(key)
|
||||
.write(from_chunk)
|
||||
.finalize(),
|
||||
),
|
||||
None => Cow::Borrowed(key),
|
||||
};
|
||||
|
||||
#[cfg(debug_assertions)]
|
||||
{
|
||||
let mut chunks = trx.get_ranges_keyvalues(
|
||||
RangeOption {
|
||||
begin: KeySelector::first_greater_or_equal(from.as_ref()),
|
||||
end: KeySelector::first_greater_or_equal(to.as_slice()),
|
||||
mode: options::StreamingMode::WantAll,
|
||||
..Default::default()
|
||||
},
|
||||
true,
|
||||
);
|
||||
|
||||
while let Ok(Some(chunk)) = chunks.try_next().await {
|
||||
let found = chunk.key();
|
||||
debug_assert!(
|
||||
found.len() == key.len() + 1 || found == key,
|
||||
"chunk range of {key:?} holds foreign key {found:?}, clearing it would destroy data"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
trx.clear_range(from.as_ref(), &to);
|
||||
}
|
||||
|
||||
async fn chunk_value(
|
||||
trx: &Transaction,
|
||||
key: &mut Vec<u8>,
|
||||
value: &[u8],
|
||||
class: &ValueClass,
|
||||
prev_num_chunks: Option<usize>,
|
||||
) -> bool {
|
||||
let num_chunks = if value.len() > MAX_VALUE_SIZE {
|
||||
value.len().div_ceil(MAX_VALUE_SIZE)
|
||||
} else {
|
||||
1
|
||||
};
|
||||
|
||||
if num_chunks > u8::MAX as usize {
|
||||
return false;
|
||||
}
|
||||
|
||||
if is_chunked_value(key[0], class)
|
||||
&& prev_num_chunks.is_none_or(|prev_num_chunks| prev_num_chunks > num_chunks)
|
||||
{
|
||||
clear_chunks(trx, key, Some((num_chunks - 1) as u8)).await;
|
||||
}
|
||||
|
||||
if value.len() > MAX_VALUE_SIZE {
|
||||
for (pos, chunk) in value.chunks(MAX_VALUE_SIZE).enumerate() {
|
||||
match pos.cmp(&1) {
|
||||
Ordering::Less => {}
|
||||
Ordering::Equal => {
|
||||
key.push(0);
|
||||
}
|
||||
Ordering::Greater => {
|
||||
*key.last_mut().unwrap() += 1;
|
||||
}
|
||||
}
|
||||
trx.set(key, chunk);
|
||||
}
|
||||
} else {
|
||||
trx.set(key, value);
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
Reference in New Issue
Block a user