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,543 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
|
||||
*
|
||||
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
||||
*/
|
||||
|
||||
use super::{PostgresStore, into_error, is_timeout_error};
|
||||
use crate::{
|
||||
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
|
||||
SUBSPACE_REGISTRY_IDX,
|
||||
backend::postgres::{DELETE_CHUNK_SIZE, MIN_DELETE_CHUNK_SIZE, into_pool_error},
|
||||
write::{
|
||||
AssignedIds, Batch, MAX_COMMIT_ATTEMPTS, MAX_COMMIT_TIME, MergeResult, Operation,
|
||||
ValueClass, ValueOp,
|
||||
},
|
||||
};
|
||||
use ahash::AHashMap;
|
||||
use deadpool_postgres::Object;
|
||||
use rand::RngExt;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio_postgres::{IsolationLevel, error::SqlState};
|
||||
|
||||
#[derive(Debug)]
|
||||
enum CommitError {
|
||||
Postgres(tokio_postgres::Error),
|
||||
Internal(trc::Error),
|
||||
//Retry,
|
||||
}
|
||||
|
||||
impl PostgresStore {
|
||||
pub(crate) async fn write(&self, mut batch: Batch<'_>) -> trc::Result<AssignedIds> {
|
||||
let mut conn = self.conn_pool.get().await.map_err(into_pool_error)?;
|
||||
let start = Instant::now();
|
||||
let mut retry_count = 0;
|
||||
|
||||
loop {
|
||||
match self.write_trx(&mut conn, &mut batch).await {
|
||||
Ok(result) => {
|
||||
return Ok(result);
|
||||
}
|
||||
Err(err) => {
|
||||
match err {
|
||||
CommitError::Postgres(err) => match err.code() {
|
||||
Some(
|
||||
&SqlState::T_R_SERIALIZATION_FAILURE
|
||||
| &SqlState::T_R_DEADLOCK_DETECTED,
|
||||
) if retry_count < MAX_COMMIT_ATTEMPTS
|
||||
&& start.elapsed() < MAX_COMMIT_TIME => {}
|
||||
Some(&SqlState::UNIQUE_VIOLATION) => {
|
||||
return Err(trc::StoreEvent::AssertValueFailed
|
||||
.into_err()
|
||||
.reason("Unique violation")
|
||||
.caused_by(trc::location!()));
|
||||
}
|
||||
_ => return Err(into_error(err)),
|
||||
},
|
||||
CommitError::Internal(err) => return Err(err),
|
||||
/*CommitError::Retry => {
|
||||
if retry_count > MAX_COMMIT_ATTEMPTS
|
||||
|| start.elapsed() > MAX_COMMIT_TIME
|
||||
{
|
||||
return Err(trc::StoreEvent::AssertValueFailed
|
||||
.into_err()
|
||||
.caused_by(trc::location!()));
|
||||
}
|
||||
}*/
|
||||
}
|
||||
|
||||
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 Object,
|
||||
batch: &mut Batch<'_>,
|
||||
) -> 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 asserted_values = AHashMap::new();
|
||||
let trx = conn
|
||||
.build_transaction()
|
||||
.isolation_level(IsolationLevel::ReadCommitted)
|
||||
.start()
|
||||
.await?;
|
||||
let mut result = AssignedIds::default();
|
||||
let has_changes = !batch.changes.is_empty();
|
||||
|
||||
if has_changes {
|
||||
for &account_id in batch.changes.keys() {
|
||||
let key = ValueClass::ChangeId.serialize(account_id, 0, 0, 0);
|
||||
let s = trx
|
||||
.prepare_cached(concat!(
|
||||
"INSERT INTO n (k, v) VALUES ($1, 1) ",
|
||||
"ON CONFLICT(k) DO UPDATE SET v = n.v + 1 RETURNING v"
|
||||
))
|
||||
.await?;
|
||||
let change_id = trx
|
||||
.query_one(&s, &[&key])
|
||||
.await
|
||||
.and_then(|row| row.try_get::<_, i64>(0))?;
|
||||
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 s = if let Some(exists) = asserted_values.get(&key) {
|
||||
if *exists {
|
||||
trx.prepare_cached(&format!(
|
||||
"UPDATE {} SET v = $2 WHERE k = $1",
|
||||
table
|
||||
))
|
||||
.await?
|
||||
} else {
|
||||
trx.prepare_cached(&format!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2)",
|
||||
table
|
||||
))
|
||||
.await?
|
||||
}
|
||||
} else {
|
||||
trx.prepare_cached(&format!(
|
||||
concat!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
|
||||
"ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v"
|
||||
),
|
||||
table
|
||||
))
|
||||
.await?
|
||||
};
|
||||
|
||||
if trx.execute(&s, &[&key, &(*value)]).await? == 0 {
|
||||
return Err(trc::StoreEvent::AssertValueFailed
|
||||
.into_err()
|
||||
.caused_by(trc::location!())
|
||||
.into());
|
||||
}
|
||||
} else {
|
||||
let s = trx
|
||||
.prepare_cached(
|
||||
"INSERT INTO b (k) VALUES ($1) ON CONFLICT (k) DO NOTHING",
|
||||
)
|
||||
.await?;
|
||||
trx.execute(&s, &[&key]).await?;
|
||||
}
|
||||
}
|
||||
ValueOp::SetFnc(set_op) => {
|
||||
let value = (set_op.fnc)(&set_op.params, &result)?;
|
||||
|
||||
let s = if let Some(exists) = asserted_values.get(&key) {
|
||||
if *exists {
|
||||
trx.prepare_cached(&format!(
|
||||
"UPDATE {} SET v = $2 WHERE k = $1",
|
||||
table
|
||||
))
|
||||
.await?
|
||||
} else {
|
||||
trx.prepare_cached(&format!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2)",
|
||||
table
|
||||
))
|
||||
.await?
|
||||
}
|
||||
} else {
|
||||
trx.prepare_cached(&format!(
|
||||
concat!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
|
||||
"ON CONFLICT (k) DO UPDATE SET v = EXCLUDED.v"
|
||||
),
|
||||
table
|
||||
))
|
||||
.await?
|
||||
};
|
||||
|
||||
if trx.execute(&s, &[&key, &value]).await? == 0 {
|
||||
return Err(trc::StoreEvent::AssertValueFailed
|
||||
.into_err()
|
||||
.caused_by(trc::location!())
|
||||
.into());
|
||||
}
|
||||
}
|
||||
ValueOp::MergeFnc(merge_op) => {
|
||||
let s = trx
|
||||
.prepare_cached(&format!(
|
||||
"SELECT v FROM {} WHERE k = $1 FOR UPDATE",
|
||||
table
|
||||
))
|
||||
.await?;
|
||||
let (exists, merge_result) = trx
|
||||
.query_opt(&s, &[&key])
|
||||
.await?
|
||||
.map(|row| {
|
||||
row.try_get::<_, &[u8]>(0)
|
||||
.map_err(CommitError::from)
|
||||
.and_then(|v| {
|
||||
(merge_op.fnc)(&merge_op.params, &result, Some(v))
|
||||
.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)
|
||||
})?;
|
||||
|
||||
match merge_result {
|
||||
MergeResult::Update(value) => {
|
||||
let s = if exists {
|
||||
trx.prepare_cached(&format!(
|
||||
"UPDATE {} SET v = $2 WHERE k = $1",
|
||||
table
|
||||
))
|
||||
.await?
|
||||
} else {
|
||||
trx.prepare_cached(&format!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2)",
|
||||
table
|
||||
))
|
||||
.await?
|
||||
};
|
||||
|
||||
trx.execute(&s, &[&key, &value]).await?;
|
||||
}
|
||||
MergeResult::Delete if exists => {
|
||||
let s = trx
|
||||
.prepare_cached(&format!(
|
||||
"DELETE FROM {} WHERE k = $1",
|
||||
table
|
||||
))
|
||||
.await?;
|
||||
trx.execute(&s, &[&key]).await?;
|
||||
|
||||
// Update asserted value
|
||||
if let Some(exists) = asserted_values.get_mut(&key) {
|
||||
*exists = false;
|
||||
}
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
}
|
||||
ValueOp::AtomicAdd(by) => {
|
||||
if *by >= 0 {
|
||||
let s = trx
|
||||
.prepare_cached(&format!(
|
||||
concat!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
|
||||
"ON CONFLICT(k) DO UPDATE SET v = {}.v + EXCLUDED.v"
|
||||
),
|
||||
table, table
|
||||
))
|
||||
.await?;
|
||||
trx.execute(&s, &[&key, &*by]).await?;
|
||||
} else {
|
||||
let s = trx
|
||||
.prepare_cached(&format!(
|
||||
"UPDATE {table} SET v = v + $1 WHERE k = $2"
|
||||
))
|
||||
.await?;
|
||||
trx.execute(&s, &[&*by, &key]).await?;
|
||||
}
|
||||
}
|
||||
ValueOp::AddAndGet(by) => {
|
||||
let s = trx
|
||||
.prepare_cached(&format!(
|
||||
concat!(
|
||||
"INSERT INTO {} (k, v) VALUES ($1, $2) ",
|
||||
"ON CONFLICT(k) DO UPDATE SET v = {}.v + EXCLUDED.v RETURNING v"
|
||||
),
|
||||
table, table
|
||||
))
|
||||
.await?;
|
||||
result.push_counter_id(
|
||||
trx.query_one(&s, &[&key, &*by])
|
||||
.await
|
||||
.and_then(|row| row.try_get::<_, i64>(0))?,
|
||||
);
|
||||
}
|
||||
ValueOp::Clear => {
|
||||
let s = trx
|
||||
.prepare_cached(&format!("DELETE FROM {} WHERE k = $1", table))
|
||||
.await?;
|
||||
trx.execute(&s, &[&key]).await?;
|
||||
|
||||
// Update asserted value
|
||||
if let Some(exists) = asserted_values.get_mut(&key) {
|
||||
*exists = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Operation::Index { field, key, set } => {
|
||||
let key = IndexKey {
|
||||
account_id,
|
||||
collection,
|
||||
document_id,
|
||||
field: *field,
|
||||
key: &*key,
|
||||
}
|
||||
.serialize(0);
|
||||
|
||||
let s = if *set {
|
||||
trx.prepare_cached(
|
||||
"INSERT INTO i (k) VALUES ($1) ON CONFLICT (k) DO NOTHING",
|
||||
)
|
||||
.await?
|
||||
} else {
|
||||
trx.prepare_cached("DELETE FROM i WHERE k = $1").await?
|
||||
};
|
||||
trx.execute(&s, &[&key]).await?;
|
||||
}
|
||||
Operation::Log { collection, set } => {
|
||||
let key = LogKey {
|
||||
account_id,
|
||||
collection: u8::from(*collection),
|
||||
change_id,
|
||||
}
|
||||
.serialize(0);
|
||||
|
||||
let s = trx
|
||||
.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?;
|
||||
}
|
||||
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
|
||||
.prepare_cached(&format!("SELECT v FROM {} WHERE k = $1 FOR UPDATE", table))
|
||||
.await?;
|
||||
let (exists, matches) = trx
|
||||
.query_opt(&s, &[&key])
|
||||
.await?
|
||||
.map(|row| {
|
||||
row.try_get::<_, &[u8]>(0)
|
||||
.map_or((true, false), |v| (true, assert_value.matches(v)))
|
||||
})
|
||||
.unwrap_or_else(|| (false, assert_value.is_none()));
|
||||
if !matches {
|
||||
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 conn = self.conn_pool.get().await.map_err(into_pool_error)?;
|
||||
|
||||
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] {
|
||||
purge_table(&conn, char::from(subspace)).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
|
||||
let conn = self.conn_pool.get().await.map_err(into_pool_error)?;
|
||||
let table = char::from(from.subspace());
|
||||
let mut from = from.serialize(0);
|
||||
let to = to.serialize(0);
|
||||
|
||||
let delete = conn
|
||||
.prepare_cached(&format!("DELETE FROM {table} WHERE k >= $1 AND k < $2"))
|
||||
.await
|
||||
.map_err(into_error)?;
|
||||
|
||||
match conn.execute(&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
|
||||
.prepare_cached(&format!(
|
||||
"SELECT k FROM {table} WHERE k >= $1 AND k < $2 ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
|
||||
))
|
||||
.await
|
||||
.map_err(into_error)?;
|
||||
|
||||
loop {
|
||||
let next = match conn.query_opt(&boundary, &[&from, &to]).await {
|
||||
Ok(next) => match next {
|
||||
Some(row) => Some(row.try_get::<_, Vec<u8>>(0).map_err(into_error)?),
|
||||
None => None,
|
||||
},
|
||||
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
|
||||
.execute(&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: &Object, table: char) -> trc::Result<()> {
|
||||
let s = conn
|
||||
.prepare_cached(&format!("DELETE FROM {table} WHERE v = 0"))
|
||||
.await
|
||||
.map_err(into_error)?;
|
||||
|
||||
match conn.execute(&s, &[]).await {
|
||||
Ok(_) => return Ok(()),
|
||||
Err(err) if is_timeout_error(&err) => (),
|
||||
Err(err) => return Err(into_error(err)),
|
||||
}
|
||||
|
||||
let purge = conn
|
||||
.prepare_cached(&format!(
|
||||
"DELETE FROM {table} WHERE v = 0 AND k >= $1 AND k < $2"
|
||||
))
|
||||
.await
|
||||
.map_err(into_error)?;
|
||||
let purge_last = conn
|
||||
.prepare_cached(&format!("DELETE FROM {table} WHERE v = 0 AND k >= $1"))
|
||||
.await
|
||||
.map_err(into_error)?;
|
||||
let mut chunk_size = DELETE_CHUNK_SIZE;
|
||||
let mut from = Vec::new();
|
||||
|
||||
loop {
|
||||
let boundary = conn
|
||||
.prepare_cached(&format!(
|
||||
"SELECT k FROM {table} WHERE k >= $1 ORDER BY k ASC LIMIT 1 OFFSET {chunk_size}"
|
||||
))
|
||||
.await
|
||||
.map_err(into_error)?;
|
||||
|
||||
loop {
|
||||
let next = match conn.query_opt(&boundary, &[&from]).await {
|
||||
Ok(next) => match next {
|
||||
Some(row) => Some(row.try_get::<_, Vec<u8>>(0).map_err(into_error)?),
|
||||
None => None,
|
||||
},
|
||||
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.execute(&purge, &[&from, next]).await,
|
||||
None => conn.execute(&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<tokio_postgres::Error> for CommitError {
|
||||
fn from(err: tokio_postgres::Error) -> Self {
|
||||
CommitError::Postgres(err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user