jcoffey-dev
7dae9b29fd
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.
530 lines
22 KiB
Rust
530 lines
22 KiB
Rust
/*
|
|
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
|
*
|
|
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
|
*/
|
|
|
|
use super::{DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, MysqlStore, into_error, is_timeout_error};
|
|
use crate::{
|
|
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
|
|
SUBSPACE_REGISTRY_IDX,
|
|
write::{
|
|
AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
|
|
ValueClass, ValueOp,
|
|
},
|
|
};
|
|
use ahash::AHashMap;
|
|
use mysql_async::{Conn, Error, IsolationLevel, TxOpts, params, prelude::Queryable};
|
|
use rand::RngExt;
|
|
use std::time::{Duration, Instant};
|
|
|
|
#[derive(Debug)]
|
|
enum CommitError {
|
|
Mysql(mysql_async::Error),
|
|
Internal(trc::Error),
|
|
//Retry,
|
|
}
|
|
|
|
impl MysqlStore {
|
|
pub(crate) async fn write(&self, mut batch: Batch<'_>) -> trc::Result<AssignedIds> {
|
|
let start = Instant::now();
|
|
let mut retry_count = 0;
|
|
let mut conn = self.conn_pool.get_conn().await.map_err(into_error)?;
|
|
|
|
loop {
|
|
let err = match self.write_trx(&mut conn, &mut batch).await {
|
|
Ok(result) => {
|
|
return Ok(result);
|
|
}
|
|
Err(err) => err,
|
|
};
|
|
|
|
let _ = conn.query_drop("ROLLBACK;").await;
|
|
|
|
match err {
|
|
CommitError::Mysql(Error::Server(err))
|
|
if [1062, 1213].contains(&err.code)
|
|
&& retry_count < MAX_COMMIT_ATTEMPTS
|
|
&& start.elapsed() < MAX_COMMIT_TIME => {}
|
|
/*CommitError::Retry => {
|
|
if retry_count > MAX_COMMIT_ATTEMPTS || start.elapsed() > MAX_COMMIT_TIME {
|
|
return Err(trc::StoreEvent::AssertValueFailed
|
|
.into_err()
|
|
.caused_by(trc::location!()));
|
|
}
|
|
}*/
|
|
CommitError::Mysql(err) => {
|
|
return Err(into_error(err));
|
|
}
|
|
CommitError::Internal(err) => {
|
|
return Err(err);
|
|
}
|
|
}
|
|
|
|
let backoff = rand::rng().random_range(50..=300);
|
|
tokio::time::sleep(Duration::from_millis(backoff)).await;
|
|
retry_count += 1;
|
|
}
|
|
}
|
|
|
|
async fn write_trx(
|
|
&self,
|
|
conn: &mut Conn,
|
|
batch: &mut Batch<'_>,
|
|
) -> Result<AssignedIds, CommitError> {
|
|
let has_changes = !batch.changes.is_empty();
|
|
let mut account_id = u32::MAX;
|
|
let mut collection = u8::MAX;
|
|
let mut document_id = u32::MAX;
|
|
let mut change_id = 0u64;
|
|
let mut asserted_values = AHashMap::new();
|
|
let mut tx_opts = TxOpts::default();
|
|
tx_opts
|
|
.with_consistent_snapshot(false)
|
|
.with_isolation_level(IsolationLevel::ReadCommitted);
|
|
let mut trx = conn.start_transaction(tx_opts).await?;
|
|
let mut result = AssignedIds::default();
|
|
|
|
if has_changes {
|
|
for &account_id in batch.changes.keys() {
|
|
let key = ValueClass::ChangeId.serialize(account_id, 0, 0, 0);
|
|
let s = trx
|
|
.prep(concat!(
|
|
"INSERT INTO n (k, v) VALUES (:k, LAST_INSERT_ID(1)) ",
|
|
"ON DUPLICATE KEY UPDATE v = LAST_INSERT_ID(v + 1)"
|
|
))
|
|
.await?;
|
|
trx.exec_drop(&s, params! {"k" => key}).await?;
|
|
let s = trx.prep("SELECT LAST_INSERT_ID()").await?;
|
|
let change_id = trx.exec_first::<i64, _, _>(&s, ()).await?.ok_or_else(|| {
|
|
mysql_async::Error::Io(mysql_async::IoError::Io(std::io::Error::other(
|
|
"LAST_INSERT_ID() did not return a value",
|
|
)))
|
|
})?;
|
|
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 key = class.serialize(account_id, collection, document_id, 0);
|
|
let subspace = class.subspace(collection);
|
|
let table = char::from(subspace);
|
|
|
|
match op {
|
|
ValueOp::Set(value) => {
|
|
if subspace != SUBSPACE_REGISTRY_IDX {
|
|
let exists = asserted_values.get(&key);
|
|
let s = if let Some(exists) = exists {
|
|
if *exists {
|
|
trx.prep(format!(
|
|
"UPDATE {} SET v = :v WHERE k = :k",
|
|
table
|
|
))
|
|
.await?
|
|
} else {
|
|
trx.prep(format!(
|
|
"INSERT INTO {} (k, v) VALUES (:k, :v)",
|
|
table
|
|
))
|
|
.await?
|
|
}
|
|
} else {
|
|
trx
|
|
.prep(
|
|
format!("INSERT INTO {} (k, v) VALUES (:k, :v) ON DUPLICATE KEY UPDATE v = VALUES(v)", table),
|
|
)
|
|
.await?
|
|
};
|
|
|
|
match trx
|
|
.exec_drop(&s, params! {"k" => key, "v" => &*value})
|
|
.await
|
|
{
|
|
Ok(_) => {
|
|
if trx.affected_rows() == 0 {
|
|
trx.rollback().await?;
|
|
return Err(trc::StoreEvent::AssertValueFailed
|
|
.into_err()
|
|
.caused_by(trc::location!())
|
|
.into());
|
|
}
|
|
}
|
|
Err(err) => {
|
|
trx.rollback().await?;
|
|
return Err(err.into());
|
|
}
|
|
}
|
|
} else {
|
|
let s = trx.prep("INSERT IGNORE INTO b (k) VALUES (?)").await?;
|
|
trx.exec_drop(&s, (key,)).await?;
|
|
}
|
|
}
|
|
ValueOp::SetFnc(set_op) => {
|
|
let value = (set_op.fnc)(&set_op.params, &result)?;
|
|
let exists = asserted_values.get(&key);
|
|
let s = if let Some(exists) = exists {
|
|
if *exists {
|
|
trx.prep(format!("UPDATE {} SET v = :v WHERE k = :k", table))
|
|
.await?
|
|
} else {
|
|
trx.prep(format!(
|
|
"INSERT INTO {} (k, v) VALUES (:k, :v)",
|
|
table
|
|
))
|
|
.await?
|
|
}
|
|
} else {
|
|
trx
|
|
.prep(
|
|
format!("INSERT INTO {} (k, v) VALUES (:k, :v) ON DUPLICATE KEY UPDATE v = VALUES(v)", table),
|
|
)
|
|
.await?
|
|
};
|
|
|
|
match trx.exec_drop(&s, params! {"k" => key, "v" => &value}).await {
|
|
Ok(_) => {
|
|
if trx.affected_rows() == 0 {
|
|
trx.rollback().await?;
|
|
return Err(trc::StoreEvent::AssertValueFailed
|
|
.into_err()
|
|
.caused_by(trc::location!())
|
|
.into());
|
|
}
|
|
}
|
|
Err(err) => {
|
|
trx.rollback().await?;
|
|
return Err(err.into());
|
|
}
|
|
}
|
|
}
|
|
ValueOp::MergeFnc(merge_op) => {
|
|
let s = trx
|
|
.prep(format!("SELECT v FROM {} WHERE k = ? FOR UPDATE", table))
|
|
.await?;
|
|
let (exists, merge_result) = trx
|
|
.exec_first::<Vec<u8>, _, _>(&s, (&key,))
|
|
.await?
|
|
.map(|bytes| {
|
|
(merge_op.fnc)(&merge_op.params, &result, Some(bytes.as_ref()))
|
|
.map(|v| (true, v))
|
|
.map_err(CommitError::from)
|
|
})
|
|
.unwrap_or_else(|| {
|
|
(merge_op.fnc)(&merge_op.params, &result, None)
|
|
.map(|v| (false, v))
|
|
.map_err(CommitError::from)
|
|
})?;
|
|
|
|
let s = if exists {
|
|
trx.prep(format!("UPDATE {} SET v = :v WHERE k = :k", table))
|
|
.await?
|
|
} else {
|
|
trx.prep(format!("INSERT INTO {} (k, v) VALUES (:k, :v)", table))
|
|
.await?
|
|
};
|
|
|
|
match merge_result {
|
|
MergeResult::Update(value) => {
|
|
if let Err(err) =
|
|
trx.exec_drop(&s, params! {"k" => key, "v" => &value}).await
|
|
{
|
|
trx.rollback().await?;
|
|
return Err(err.into());
|
|
}
|
|
}
|
|
MergeResult::Delete if exists => {
|
|
// Update asserted value
|
|
if let Some(exists) = asserted_values.get_mut(&key) {
|
|
*exists = false;
|
|
}
|
|
|
|
let s = trx
|
|
.prep(format!("DELETE FROM {} WHERE k = ?", table))
|
|
.await?;
|
|
trx.exec_drop(&s, (key,)).await?;
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
ValueOp::AtomicAdd(by) => {
|
|
if *by >= 0 {
|
|
let s = trx
|
|
.prep(format!(
|
|
concat!(
|
|
"INSERT INTO {} (k, v) VALUES (?, ?) ",
|
|
"ON DUPLICATE KEY UPDATE v = v + VALUES(v)"
|
|
),
|
|
table
|
|
))
|
|
.await?;
|
|
trx.exec_drop(&s, (key, &*by)).await?;
|
|
} else {
|
|
let s = trx
|
|
.prep(format!("UPDATE {table} SET v = v + ? WHERE k = ?"))
|
|
.await?;
|
|
trx.exec_drop(&s, (&*by, key)).await?;
|
|
}
|
|
}
|
|
ValueOp::AddAndGet(by) => {
|
|
let s = trx
|
|
.prep(format!(
|
|
concat!(
|
|
"INSERT INTO {} (k, v) VALUES (:k, LAST_INSERT_ID(:v)) ",
|
|
"ON DUPLICATE KEY UPDATE v = LAST_INSERT_ID(v + :v)"
|
|
),
|
|
table
|
|
))
|
|
.await?;
|
|
trx.exec_drop(&s, params! {"k" => key, "v" => &*by}).await?;
|
|
let s = trx.prep("SELECT LAST_INSERT_ID()").await?;
|
|
result.push_counter_id(
|
|
trx.exec_first::<i64, _, _>(&s, ()).await?.ok_or_else(|| {
|
|
mysql_async::Error::Io(mysql_async::IoError::Io(
|
|
std::io::Error::other(
|
|
"LAST_INSERT_ID() did not return a value",
|
|
),
|
|
))
|
|
})?,
|
|
);
|
|
}
|
|
ValueOp::Clear => {
|
|
// Update asserted value
|
|
if let Some(exists) = asserted_values.get_mut(&key) {
|
|
*exists = false;
|
|
}
|
|
|
|
let s = trx
|
|
.prep(format!("DELETE FROM {} WHERE k = ?", table))
|
|
.await?;
|
|
trx.exec_drop(&s, (key,)).await?;
|
|
}
|
|
}
|
|
}
|
|
Operation::Index { field, key, set } => {
|
|
let key = IndexKey {
|
|
account_id,
|
|
collection,
|
|
document_id,
|
|
field: *field,
|
|
key: &*key,
|
|
}
|
|
.serialize(0);
|
|
|
|
let s = if *set {
|
|
trx.prep("INSERT IGNORE INTO i (k) VALUES (?)").await?
|
|
} else {
|
|
trx.prep("DELETE FROM i WHERE k = ?").await?
|
|
};
|
|
trx.exec_drop(&s, (key,)).await?;
|
|
}
|
|
Operation::Log { collection, set } => {
|
|
let key = LogKey {
|
|
account_id,
|
|
collection: u8::from(*collection),
|
|
change_id,
|
|
}
|
|
.serialize(0);
|
|
|
|
let s = trx
|
|
.prep("INSERT INTO l (k, v) VALUES (?, ?) ON DUPLICATE KEY UPDATE v = VALUES(v)")
|
|
.await?;
|
|
|
|
trx.exec_drop(&s, (key, &*set)).await?;
|
|
}
|
|
Operation::AssertValue {
|
|
class,
|
|
assert_value,
|
|
} => {
|
|
let key = class.serialize(account_id, collection, document_id, 0);
|
|
let table = char::from(class.subspace(collection));
|
|
|
|
let s = trx
|
|
.prep(format!("SELECT v FROM {} WHERE k = ? FOR UPDATE", table))
|
|
.await?;
|
|
let (exists, matches) = trx
|
|
.exec_first::<Vec<u8>, _, _>(&s, (&key,))
|
|
.await?
|
|
.map(|bytes| (true, assert_value.matches(&bytes)))
|
|
.unwrap_or_else(|| (false, assert_value.is_none()));
|
|
if !matches {
|
|
trx.rollback().await?;
|
|
return Err(trc::StoreEvent::AssertValueFailed
|
|
.into_err()
|
|
.caused_by(trc::location!())
|
|
.into());
|
|
}
|
|
asserted_values.insert(key, exists);
|
|
}
|
|
}
|
|
}
|
|
|
|
trx.commit().await.map(|_| result).map_err(Into::into)
|
|
}
|
|
|
|
pub(crate) async fn purge_store(&self) -> trc::Result<()> {
|
|
let mut conn = self.conn_pool.get_conn().await.map_err(into_error)?;
|
|
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] {
|
|
purge_table(&mut conn, char::from(subspace)).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
|
|
let mut conn = self.conn_pool.get_conn().await.map_err(into_error)?;
|
|
let table = char::from(from.subspace());
|
|
let mut from = from.serialize(0);
|
|
let to = to.serialize(0);
|
|
|
|
let delete = conn
|
|
.prep(format!("DELETE FROM {table} WHERE k >= ? AND k < ?"))
|
|
.await
|
|
.map_err(into_error)?;
|
|
|
|
match conn.exec_drop(&delete, (&from, &to)).await {
|
|
Ok(_) => return Ok(()),
|
|
Err(err) if is_timeout_error(&err) => (),
|
|
Err(err) => return Err(into_error(err)),
|
|
}
|
|
|
|
let mut chunk_size = DELETE_CHUNK_SIZE;
|
|
|
|
loop {
|
|
let boundary = conn
|
|
.prep(format!(
|
|
"SELECT k FROM {table} WHERE k >= ? AND k < ? ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
|
|
))
|
|
.await
|
|
.map_err(into_error)?;
|
|
|
|
loop {
|
|
let next = match conn
|
|
.exec_first::<Vec<u8>, _, _>(&boundary, (&from, &to))
|
|
.await
|
|
{
|
|
Ok(next) => next,
|
|
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
|
|
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
|
|
break;
|
|
}
|
|
Err(err) => return Err(into_error(err)),
|
|
};
|
|
|
|
match conn
|
|
.exec_drop(&delete, (&from, next.as_ref().unwrap_or(&to)))
|
|
.await
|
|
{
|
|
Ok(_) => (),
|
|
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
|
|
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
|
|
break;
|
|
}
|
|
Err(err) => return Err(into_error(err)),
|
|
}
|
|
|
|
match next {
|
|
Some(next) => from = next,
|
|
None => return Ok(()),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn purge_table(conn: &mut Conn, table: char) -> trc::Result<()> {
|
|
let s = conn
|
|
.prep(format!("DELETE FROM {table} WHERE v = 0"))
|
|
.await
|
|
.map_err(into_error)?;
|
|
|
|
match conn.exec_drop(&s, ()).await {
|
|
Ok(_) => return Ok(()),
|
|
Err(err) if is_timeout_error(&err) => (),
|
|
Err(err) => return Err(into_error(err)),
|
|
}
|
|
|
|
let purge = conn
|
|
.prep(format!(
|
|
"DELETE FROM {table} WHERE v = 0 AND k >= ? AND k < ?"
|
|
))
|
|
.await
|
|
.map_err(into_error)?;
|
|
let purge_last = conn
|
|
.prep(format!("DELETE FROM {table} WHERE v = 0 AND k >= ?"))
|
|
.await
|
|
.map_err(into_error)?;
|
|
let mut chunk_size = DELETE_CHUNK_SIZE;
|
|
let mut from = Vec::new();
|
|
|
|
loop {
|
|
let boundary = conn
|
|
.prep(format!(
|
|
"SELECT k FROM {table} WHERE k >= ? ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
|
|
))
|
|
.await
|
|
.map_err(into_error)?;
|
|
|
|
loop {
|
|
let next = match conn.exec_first::<Vec<u8>, _, _>(&boundary, (&from,)).await {
|
|
Ok(next) => next,
|
|
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
|
|
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
|
|
break;
|
|
}
|
|
Err(err) => return Err(into_error(err)),
|
|
};
|
|
|
|
let result = match &next {
|
|
Some(next) => conn.exec_drop(&purge, (&from, next)).await,
|
|
None => conn.exec_drop(&purge_last, (&from,)).await,
|
|
};
|
|
|
|
match result {
|
|
Ok(_) => (),
|
|
Err(err) if is_timeout_error(&err) && chunk_size > MIN_DELETE_CHUNK_SIZE => {
|
|
chunk_size = (chunk_size / 2).max(MIN_DELETE_CHUNK_SIZE);
|
|
break;
|
|
}
|
|
Err(err) => return Err(into_error(err)),
|
|
}
|
|
|
|
match next {
|
|
Some(next) => from = next,
|
|
None => return Ok(()),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl From<trc::Error> for CommitError {
|
|
fn from(err: trc::Error) -> Self {
|
|
CommitError::Internal(err)
|
|
}
|
|
}
|
|
|
|
impl From<mysql_async::Error> for CommitError {
|
|
fn from(err: mysql_async::Error) -> Self {
|
|
CommitError::Mysql(err)
|
|
}
|
|
}
|