373 lines
16 KiB
Rust
373 lines
16 KiB
Rust
/*
|
|
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
|
|
*
|
|
* 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, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
|
|
RegistryClass, TaskQueueClass, TelemetryClass, ValueClass, ValueOp, key::KeySerializer,
|
|
},
|
|
*,
|
|
};
|
|
use foundationdb::{
|
|
FdbError, KeySelector, RangeOption, Transaction,
|
|
options::{self, MutationType},
|
|
};
|
|
use futures::TryStreamExt;
|
|
use rand::Rng;
|
|
use std::{
|
|
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) {
|
|
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) {
|
|
trx.cancel();
|
|
return Err(trc::StoreEvent::FoundationdbError
|
|
.ctx(trc::Key::Reason, "Value is too large"));
|
|
}
|
|
}
|
|
ValueOp::MergeFnc(merge_op) => {
|
|
let (merge_result, is_chunked) =
|
|
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()),
|
|
)?,
|
|
false,
|
|
),
|
|
ChunkedValue::Chunked { bytes, .. } => (
|
|
(merge_op.fnc)(
|
|
&merge_op.params,
|
|
&result,
|
|
Some(bytes.as_ref()),
|
|
)?,
|
|
true,
|
|
),
|
|
ChunkedValue::None => (
|
|
(merge_op.fnc)(&merge_op.params, &result, None)?,
|
|
false,
|
|
),
|
|
};
|
|
|
|
match merge_result {
|
|
MergeResult::Update(value) => {
|
|
if !chunk_value(&trx, &mut key, &value) {
|
|
trx.cancel();
|
|
return Err(trc::StoreEvent::FoundationdbError
|
|
.ctx(trc::Key::Reason, "Value is too large"));
|
|
}
|
|
}
|
|
MergeResult::Delete => {
|
|
if is_chunked {
|
|
trx.clear_range(
|
|
&key,
|
|
&KeySerializer::new(key.len() + 1)
|
|
.write(key.as_slice())
|
|
.write(u8::MAX)
|
|
.finalize(),
|
|
);
|
|
} 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 matches!(
|
|
key[0],
|
|
SUBSPACE_TASK_QUEUE
|
|
| SUBSPACE_IN_MEMORY_VALUE
|
|
| SUBSPACE_PROPERTY
|
|
| SUBSPACE_QUEUE_MESSAGE
|
|
| SUBSPACE_REPORT_OUT
|
|
| SUBSPACE_REPORT_IN
|
|
| SUBSPACE_TELEMETRY_SPAN
|
|
| SUBSPACE_SEARCH_INDEX
|
|
| SUBSPACE_LOGS
|
|
) && matches!(
|
|
class,
|
|
ValueClass::Property(_)
|
|
| ValueClass::Queue(_)
|
|
| ValueClass::Registry(RegistryClass::Item { .. })
|
|
| ValueClass::ShareNotification { .. }
|
|
| ValueClass::Telemetry(TelemetryClass::Metric { .. })
|
|
| ValueClass::TaskQueue(TaskQueueClass::Task { .. })
|
|
| ValueClass::InMemory(_)
|
|
) {
|
|
// Clear range for potentially chunked values to avoid leaving orphaned chunks
|
|
trx.clear_range(
|
|
&key,
|
|
&KeySerializer::new(key.len() + 1)
|
|
.write(key.as_slice())
|
|
.write(u8::MAX)
|
|
.finalize(),
|
|
);
|
|
} 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 chunk_value(trx: &Transaction, key: &mut Vec<u8>, value: &[u8]) -> bool {
|
|
if !value.is_empty() && 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 => {
|
|
if pos < u8::MAX as usize {
|
|
*key.last_mut().unwrap() += 1;
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
trx.set(key, chunk);
|
|
}
|
|
} else {
|
|
trx.set(key, value.as_ref());
|
|
}
|
|
|
|
true
|
|
}
|