use crate::metrics::{Metrics, QueryOutcome, SyncOperation, SyncOutcome, SyncRepairLabel};
use crate::protocol::{Message, WireParam, WireRetainedUnit, WireSyncRepairAction, WireSyncStatus};
use powdb_auth::{Permission, Role, UserStore};
use powdb_query::executor::{is_read_only_statement, Engine, WalDurabilityTicket};
use powdb_query::parser;
use powdb_query::result::{QueryError, QueryResult};
use powdb_query::sql;
use powdb_storage::types::Value;
use powdb_sync::{
acknowledge_replica_apply, read_identity, read_units_through, replica_sync_status,
retained_segments_dir, validate_retained_tail_available, validate_v1_retained_units_applyable,
ReplicaSyncStatus, RetainedUnit, SyncRepairAction, RETAINED_SEGMENT_FORMAT_VERSION,
};
use std::collections::HashMap;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};
use tracing::{debug, error, info, warn};
use zeroize::Zeroizing;
pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
pub type TxGate = Arc<Semaphore>;
pub fn new_tx_gate() -> TxGate {
Arc::new(Semaphore::new(1))
}
const MAX_QUERY_LENGTH: usize = 1024 * 1024;
const MAX_SYNC_PULL_UNITS: u32 = 4096;
const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
#[cfg(not(test))]
const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
#[cfg(test)]
const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_AUTH_FAILURES: u32 = 5;
const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
pub fn new_rate_limiter() -> AuthRateLimiter {
Arc::new(Mutex::new(HashMap::new()))
}
fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
let now = Instant::now();
map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
if let Some((count, _)) = map.get(&ip) {
*count >= MAX_AUTH_FAILURES
} else {
false
}
}
fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
let now = Instant::now();
let entry = map.entry(ip).or_insert((0, now));
if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
*entry = (1, now);
} else {
entry.0 += 1;
}
}
fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
map.remove(&ip);
}
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
use sha2::{Digest, Sha256};
let ha = Sha256::digest(a);
let hb = Sha256::digest(b);
let mut diff = 0u8;
for (x, y) in ha.iter().zip(hb.iter()) {
diff |= x ^ y;
}
diff == 0
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Principal {
pub name: String,
pub role: String,
}
fn is_ddl_statement(stmt: &powdb_query::ast::Statement) -> bool {
use powdb_query::ast::Statement;
let inner = match stmt {
Statement::Explain(inner) => inner.as_ref(),
other => other,
};
matches!(
inner,
Statement::CreateType(_)
| Statement::AlterTable(_)
| Statement::DropTable(_)
| Statement::CreateView(_)
| Statement::DropView(_)
)
}
fn required_permission(stmt: &powdb_query::ast::Statement) -> Permission {
if is_read_only_statement(stmt) {
Permission::Read
} else if is_ddl_statement(stmt) {
Permission::Ddl
} else {
Permission::Write
}
}
fn check_statement_permitted(
principal: Option<&Principal>,
stmt: &powdb_query::ast::Statement,
) -> Result<(), QueryError> {
let Some(p) = principal else {
return Ok(());
};
if is_read_only_statement(stmt) {
return Ok(());
}
let needed = required_permission(stmt);
if Role::builtin(&p.role).is_some_and(|r| r.allows(needed)) {
return Ok(());
}
let kind = if needed == Permission::Ddl {
"schema-definition"
} else {
"write"
};
Err(QueryError::Execution(format!(
"permission denied: role '{}' cannot execute {kind} statements",
p.role
)))
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthOutcome {
Authenticated { principal: Option<Principal> },
Rejected,
}
pub fn authenticate_connect(
users: &UserStore,
expected_password: Option<&str>,
username: Option<&str>,
password: Option<&str>,
) -> AuthOutcome {
if !users.is_empty() {
let Some(name) = username else {
return AuthOutcome::Rejected;
};
let Some(candidate) = password else {
return AuthOutcome::Rejected;
};
match users.authenticate(name, candidate) {
Some(user) => AuthOutcome::Authenticated {
principal: Some(Principal {
name: user.name.clone(),
role: user.role.clone(),
}),
},
None => AuthOutcome::Rejected,
}
} else {
match expected_password {
Some(expected) => {
if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
AuthOutcome::Authenticated { principal: None }
} else {
AuthOutcome::Rejected
}
}
None => AuthOutcome::Authenticated { principal: None },
}
}
}
const DEFAULT_DB_NAME: &str = "default";
fn check_db_name(configured: Option<&str>, requested: &str) -> Result<(), String> {
if requested.is_empty() || requested == DEFAULT_DB_NAME {
return Ok(());
}
match configured {
None => Ok(()),
Some(name) if requested == name => Ok(()),
Some(name) => Err(format!(
"unknown database '{requested}'; this server serves '{name}'"
)),
}
}
const SAFE_ERROR_PREFIXES: &[&str] = &[
"table not found",
"table '",
"type '",
"column not found",
"at position",
"parse error",
"type mismatch",
"unknown table",
"unknown column",
"unknown function",
"syntax error",
"expected",
"unexpected",
"missing",
"duplicate",
"invalid",
"cannot",
"no such",
"already exists",
"permission denied",
"row too large",
"unique constraint violation",
"sort input exceeds",
"join result exceeds",
"query exceeded memory budget",
"result too large",
"wal durability sync failed",
];
fn sanitize_error(e: &str) -> String {
let lower = e.to_lowercase();
for prefix in SAFE_ERROR_PREFIXES {
if lower.starts_with(prefix) {
return e.to_string();
}
}
"query execution error".into()
}
async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
let write_fut = async {
if msg.write_to(writer).await.is_err() {
return false;
}
writer.flush().await.is_ok()
};
tokio::time::timeout(WRITE_TIMEOUT, write_fut)
.await
.unwrap_or_default()
}
pub struct ConnOpts<'a> {
pub engine: Arc<RwLock<Engine>>,
pub tx_gate: TxGate,
pub expected_password: Option<Zeroizing<String>>,
pub users: Arc<UserStore>,
pub shutdown_rx: &'a mut watch::Receiver<bool>,
pub idle_timeout: Duration,
pub query_timeout: Duration,
pub rate_limiter: Option<&'a AuthRateLimiter>,
pub peer_addr: Option<std::net::SocketAddr>,
pub metrics: Arc<Metrics>,
pub tx_wait_timeout: Duration,
pub db_name: Option<String>,
}
type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
fn execute_write_deferred(
engine: &Arc<RwLock<Engine>>,
f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
) -> DispatchOutcome {
let mut eng = match engine.write() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
eng.run_with_deferred_durability(f)
}
fn dispatch_query(
engine: &Arc<RwLock<Engine>>,
query: &str,
principal: Option<&Principal>,
) -> DispatchOutcome {
let stmt_result = parser::parse(query).map_err(|e| e.to_string());
if let Ok(stmt) = &stmt_result {
if let Err(e) = check_statement_permitted(principal, stmt) {
return (Err(e), None);
}
}
let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
if can_try_read {
let res = {
let eng = match engine.read() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
eng.execute_powql_readonly(query)
};
match res {
Ok(r) => return (Ok(r), None),
Err(QueryError::ReadonlyNeedsWrite) => {
}
Err(e) => return (Err(e), None),
}
}
if matches!(
parsed_transaction_control(&stmt_result),
Some(TransactionControl::Rollback)
) {
let mut eng = match engine.write() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
}
execute_write_deferred(engine, |eng| eng.execute_powql(query))
}
fn dispatch_sql_query(
engine: &Arc<RwLock<Engine>>,
query: &str,
principal: Option<&Principal>,
) -> DispatchOutcome {
let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
if let Ok(stmt) = &stmt_result {
if let Err(e) = check_statement_permitted(principal, stmt) {
return (Err(e), None);
}
}
let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
if can_try_read {
let res = {
let eng = match engine.read() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
eng.execute_sql_readonly(query)
};
match res {
Ok(r) => return (Ok(r), None),
Err(QueryError::ReadonlyNeedsWrite) => {}
Err(e) => return (Err(e), None),
}
}
if matches!(
parsed_transaction_control(&stmt_result),
Some(TransactionControl::Rollback)
) {
let mut eng = match engine.write() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
}
execute_write_deferred(engine, |eng| eng.execute_sql(query))
}
fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
use powdb_query::ast::ParamValue;
match p {
WireParam::Null => ParamValue::Null,
WireParam::Int(v) => ParamValue::Int(*v),
WireParam::Float(v) => ParamValue::Float(*v),
WireParam::Bool(v) => ParamValue::Bool(*v),
WireParam::Str(s) => ParamValue::Str(s.clone()),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TransactionControl {
Begin,
Commit,
Rollback,
}
fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
use powdb_query::ast::Statement;
match stmt {
Statement::Begin => Some(TransactionControl::Begin),
Statement::Commit => Some(TransactionControl::Commit),
Statement::Rollback => Some(TransactionControl::Rollback),
_ => None,
}
}
fn classify_query_transaction_control(query: &str) -> Option<TransactionControl> {
parser::parse(query)
.ok()
.and_then(|stmt| transaction_control(&stmt))
}
fn classify_sql_transaction_control(query: &str) -> Option<TransactionControl> {
sql::parse_sql(query)
.ok()
.and_then(|stmt| transaction_control(&stmt))
}
fn classify_params_transaction_control(
query: &str,
params: &[WireParam],
) -> Option<TransactionControl> {
let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
parser::parse_with_params(query, &bound)
.ok()
.and_then(|stmt| transaction_control(&stmt))
}
fn parsed_transaction_control(
stmt_result: &Result<powdb_query::ast::Statement, String>,
) -> Option<TransactionControl> {
stmt_result.as_ref().ok().and_then(transaction_control)
}
fn execute_rollback_preserving_sync_if_needed(
engine: &mut Engine,
) -> Result<QueryResult, QueryError> {
engine.rollback_transaction_preserving_wal_archive()
}
fn dispatch_query_with_params(
engine: &Arc<RwLock<Engine>>,
query: &str,
params: &[WireParam],
principal: Option<&Principal>,
) -> DispatchOutcome {
let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
if let Ok(stmt) = &stmt_result {
if let Err(e) = check_statement_permitted(principal, stmt) {
return (Err(e), None);
}
}
let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
if can_try_read {
let res = {
let eng = match engine.read() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
eng.execute_powql_readonly_with_params(query, &bound)
};
match res {
Ok(r) => return (Ok(r), None),
Err(QueryError::ReadonlyNeedsWrite) => {
}
Err(e) => return (Err(e), None),
}
}
if matches!(
parsed_transaction_control(&stmt_result),
Some(TransactionControl::Rollback)
) {
let mut eng = match engine.write() {
Ok(eng) => eng,
Err(e) => {
return (
Err(QueryError::Execution(format!("lock poisoned: {e}"))),
None,
)
}
};
return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
}
execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, &bound))
}
#[derive(Debug, Clone)]
struct SyncPullRequest {
replica_id: String,
since_lsn: u64,
max_units: u32,
max_bytes: u64,
database_id: [u8; 16],
primary_generation: u64,
wal_format_version: u16,
catalog_version: u16,
segment_format_version: u16,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SyncErrorClass {
AuthRequired,
PermissionDenied,
InvalidReplicaId,
ActiveTransaction,
QueryExecution,
InvalidMaxUnits,
InvalidMaxBytes,
SyncContext,
StatusRead,
CursorLsnMismatch,
IdentityRead,
IdentityOrFormatMismatch,
RetainedRead,
RetainedUnitEncoding,
RetainedChunkNotApplyable,
LsnAheadOfRemote,
AckValidation,
AckUpdate,
Internal,
}
impl SyncErrorClass {
const fn as_label(self) -> &'static str {
match self {
Self::AuthRequired => "auth_required",
Self::PermissionDenied => "permission_denied",
Self::InvalidReplicaId => "invalid_replica_id",
Self::ActiveTransaction => "active_transaction",
Self::QueryExecution => "query_execution",
Self::InvalidMaxUnits => "invalid_max_units",
Self::InvalidMaxBytes => "invalid_max_bytes",
Self::SyncContext => "sync_context",
Self::StatusRead => "status_read",
Self::CursorLsnMismatch => "cursor_lsn_mismatch",
Self::IdentityRead => "identity_read",
Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
Self::RetainedRead => "retained_read",
Self::RetainedUnitEncoding => "retained_unit_encoding",
Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
Self::AckValidation => "ack_validation",
Self::AckUpdate => "ack_update",
Self::Internal => "internal",
}
}
}
#[derive(Debug, Clone)]
struct SyncDecision {
message: Message,
error_class: Option<SyncErrorClass>,
}
impl SyncDecision {
fn ok(message: Message) -> Self {
Self {
message,
error_class: None,
}
}
fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
Self {
message: Message::Error {
message: message.into(),
},
error_class: Some(class),
}
}
}
fn check_sync_protocol_permitted(
credential_authenticated: bool,
principal: Option<&Principal>,
) -> Result<(), (SyncErrorClass, String)> {
if !credential_authenticated {
return Err((
SyncErrorClass::AuthRequired,
"sync protocol requires authentication".to_string(),
));
}
if let Some(principal) = principal {
let allowed =
Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
if !allowed {
return Err((
SyncErrorClass::PermissionDenied,
format!(
"permission denied: role '{}' cannot use sync protocol",
principal.role
),
));
}
}
Ok(())
}
fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
if replica_id.is_empty() {
return Err("replica id must be non-empty".to_string());
}
if replica_id.len() > 128 {
return Err("replica id must be at most 128 bytes".to_string());
}
if !replica_id
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
{
return Err("replica id contains unsupported characters".to_string());
}
Ok(())
}
fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<(PathBuf, u64), String> {
let engine = engine
.read()
.map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
let catalog = engine.catalog();
Ok((catalog.data_dir().to_path_buf(), catalog.max_lsn()))
}
fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
match action {
SyncRepairAction::None => WireSyncRepairAction::None,
SyncRepairAction::Pull => WireSyncRepairAction::Pull,
SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
}
}
fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
WireSyncStatus {
replica_id: status.replica_id,
active: status.active,
last_applied_lsn: status.last_applied_lsn,
remote_lsn: status.remote_lsn,
servable_lsn: status.servable_lsn,
unarchived_lsn: status.unarchived_lsn,
lag_lsn: status.lag_lsn,
lag_bytes: status.lag_bytes,
lag_ms: status.lag_ms,
stale: status.stale,
repair_action: wire_repair_action(status.repair_action),
last_sync_error: status.last_sync_error,
}
}
fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
WireRetainedUnit {
tx_id: unit.tx_id,
record_type: unit.record_type,
lsn: unit.lsn,
data: unit.data,
}
}
fn sync_operation_outcome(message: &Message) -> SyncOutcome {
match message {
Message::SyncStatusResult { .. }
| Message::SyncPullResult { .. }
| Message::SyncAckResult { .. } => SyncOutcome::Ok,
_ => SyncOutcome::Error,
}
}
fn sync_operation_label(operation: SyncOperation) -> &'static str {
match operation {
SyncOperation::Status => "status",
SyncOperation::Pull => "pull",
SyncOperation::Ack => "ack",
}
}
fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
units.iter().fold(0, |total, unit| {
total.saturating_add(unit.encoded_len().unwrap_or(0))
})
}
fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
match action {
WireSyncRepairAction::None => SyncRepairLabel::None,
WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
}
}
fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
match action {
WireSyncRepairAction::None => "none",
WireSyncRepairAction::Pull => "pull",
WireSyncRepairAction::AwaitArchive => "await_archive",
WireSyncRepairAction::Rebootstrap => "rebootstrap",
}
}
const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
const FNV1A64_PRIME: u64 = 0x100000001b3;
const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
fn replica_fingerprint(replica_id: &str) -> String {
let mut hash = FNV1A64_OFFSET;
for byte in replica_id.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(FNV1A64_PRIME);
}
format!("{hash:016x}")
}
fn log_replica_fingerprint(replica_id: &str) -> String {
if validate_wire_replica_id(replica_id).is_ok() {
replica_fingerprint(replica_id)
} else {
INVALID_REPLICA_FINGERPRINT.to_string()
}
}
#[derive(Debug, Clone)]
struct SyncLogContext {
replica_fingerprint: String,
since_lsn: Option<u64>,
applied_lsn: Option<u64>,
observed_remote_lsn: Option<u64>,
max_units: Option<u32>,
max_bytes: Option<u64>,
}
impl SyncLogContext {
fn status(replica_id: &str) -> Self {
Self::base(replica_id)
}
fn pull(request: &SyncPullRequest) -> Self {
Self {
replica_fingerprint: log_replica_fingerprint(&request.replica_id),
since_lsn: Some(request.since_lsn),
applied_lsn: None,
observed_remote_lsn: None,
max_units: Some(request.max_units),
max_bytes: Some(request.max_bytes),
}
}
fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
Self {
replica_fingerprint: log_replica_fingerprint(replica_id),
since_lsn: None,
applied_lsn: Some(applied_lsn),
observed_remote_lsn: Some(observed_remote_lsn),
max_units: None,
max_bytes: None,
}
}
fn base(replica_id: &str) -> Self {
Self {
replica_fingerprint: log_replica_fingerprint(replica_id),
since_lsn: None,
applied_lsn: None,
observed_remote_lsn: None,
max_units: None,
max_bytes: None,
}
}
}
struct SyncExecutionContext<'a> {
tx_gate: TxGate,
connection_has_transaction: bool,
operation: SyncOperation,
log_context: SyncLogContext,
metrics: &'a Arc<Metrics>,
query_timeout: Duration,
}
fn log_sync_decision(
operation: SyncOperation,
context: &SyncLogContext,
elapsed: Duration,
decision: &SyncDecision,
) {
let operation = sync_operation_label(operation);
let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
match &decision.message {
Message::SyncStatusResult { status } => {
let repair_action = wire_repair_action_label(status.repair_action);
if status.stale || status.repair_action != WireSyncRepairAction::None {
info!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
remote_lsn = status.remote_lsn,
last_applied_lsn = ?status.last_applied_lsn,
servable_lsn = ?status.servable_lsn,
unarchived_lsn = ?status.unarchived_lsn,
lag_lsn = ?status.lag_lsn,
lag_bytes = ?status.lag_bytes,
lag_ms = ?status.lag_ms,
stale = status.stale,
repair_action,
elapsed_ms,
"sync decision"
);
} else {
debug!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
remote_lsn = status.remote_lsn,
last_applied_lsn = ?status.last_applied_lsn,
repair_action,
elapsed_ms,
"sync decision"
);
}
}
Message::SyncPullResult {
status,
units,
has_more,
} => {
let repair_action = wire_repair_action_label(status.repair_action);
let units_len = units.len();
let payload_bytes = sync_pull_payload_bytes(units);
if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
info!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
since_lsn = ?context.since_lsn,
max_units = ?context.max_units,
max_bytes = ?context.max_bytes,
units = units_len,
payload_bytes,
has_more = *has_more,
remote_lsn = status.remote_lsn,
last_applied_lsn = ?status.last_applied_lsn,
servable_lsn = ?status.servable_lsn,
unarchived_lsn = ?status.unarchived_lsn,
lag_lsn = ?status.lag_lsn,
stale = status.stale,
repair_action,
elapsed_ms,
"sync decision"
);
} else {
debug!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
since_lsn = ?context.since_lsn,
units = units_len,
payload_bytes,
has_more = *has_more,
remote_lsn = status.remote_lsn,
repair_action,
elapsed_ms,
"sync decision"
);
}
}
Message::SyncAckResult {
previous_applied_lsn,
applied_lsn,
remote_lsn,
advanced,
status,
} => {
let repair_action = wire_repair_action_label(status.repair_action);
if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
info!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
requested_applied_lsn = ?context.applied_lsn,
observed_remote_lsn = ?context.observed_remote_lsn,
previous_applied_lsn = *previous_applied_lsn,
applied_lsn = *applied_lsn,
remote_lsn = *remote_lsn,
advanced = *advanced,
stale = status.stale,
repair_action,
elapsed_ms,
"sync decision"
);
} else {
debug!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
requested_applied_lsn = ?context.applied_lsn,
previous_applied_lsn = *previous_applied_lsn,
applied_lsn = *applied_lsn,
remote_lsn = *remote_lsn,
advanced = *advanced,
repair_action,
elapsed_ms,
"sync decision"
);
}
}
Message::Error { .. } => {
warn!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
since_lsn = ?context.since_lsn,
applied_lsn = ?context.applied_lsn,
observed_remote_lsn = ?context.observed_remote_lsn,
max_units = ?context.max_units,
max_bytes = ?context.max_bytes,
error_class = decision
.error_class
.unwrap_or(SyncErrorClass::Internal)
.as_label(),
elapsed_ms,
"sync decision rejected"
);
}
_ => {
debug!(
operation = operation,
replica_fingerprint = %context.replica_fingerprint,
elapsed_ms,
"unexpected sync decision response"
);
}
}
}
fn trim_to_applyable_v1_prefix(
raw_units: &mut Vec<RetainedUnit>,
wire_units: &mut Vec<WireRetainedUnit>,
) -> Result<(), String> {
let mut last_error = None;
while !raw_units.is_empty() {
match validate_v1_retained_units_applyable(raw_units) {
Ok(()) => return Ok(()),
Err(err) => {
last_error = Some(err.to_string());
raw_units.pop();
wire_units.pop();
}
}
}
if let Some(error) = last_error {
return Err(format!(
"sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
));
}
Ok(())
}
fn validate_sync_ack_applyable_boundary(
data_dir: &Path,
replica_id: &str,
applied_lsn: u64,
remote_lsn: u64,
) -> Result<(), String> {
let status =
replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
let Some(previous_lsn) = status.last_applied_lsn else {
return Ok(());
};
if applied_lsn <= previous_lsn {
return Ok(());
}
let range_len = applied_lsn - previous_lsn;
if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
return Err(format!(
"sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
));
}
let max_units =
usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
let segment_dir = retained_segments_dir(data_dir);
let units = read_units_through(
&segment_dir,
identity.segment_identity(),
previous_lsn,
applied_lsn,
max_units,
)
.map_err(|err| err.to_string())?;
if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
return Err(
"sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
);
}
validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
}
#[cfg(test)]
fn dispatch_sync_status(
engine: &Arc<RwLock<Engine>>,
replica_id: String,
credential_authenticated: bool,
principal: Option<&Principal>,
) -> Message {
dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
}
fn dispatch_sync_status_decision(
engine: &Arc<RwLock<Engine>>,
replica_id: String,
credential_authenticated: bool,
principal: Option<&Principal>,
) -> SyncDecision {
if let Err((class, message)) =
check_sync_protocol_permitted(credential_authenticated, principal)
{
return SyncDecision::error(class, message);
}
if let Err(message) = validate_wire_replica_id(&replica_id) {
return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
}
let (data_dir, remote_lsn) = match sync_context(engine) {
Ok(context) => context,
Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
};
match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
status: wire_sync_status(status),
}),
Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
}
}
#[cfg(test)]
fn dispatch_sync_pull(
engine: &Arc<RwLock<Engine>>,
request: SyncPullRequest,
credential_authenticated: bool,
principal: Option<&Principal>,
) -> Message {
dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
}
fn dispatch_sync_pull_decision(
engine: &Arc<RwLock<Engine>>,
request: SyncPullRequest,
credential_authenticated: bool,
principal: Option<&Principal>,
) -> SyncDecision {
if let Err((class, message)) =
check_sync_protocol_permitted(credential_authenticated, principal)
{
return SyncDecision::error(class, message);
}
if let Err(message) = validate_wire_replica_id(&request.replica_id) {
return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
}
if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
return SyncDecision::error(
SyncErrorClass::InvalidMaxUnits,
format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
);
}
if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
return SyncDecision::error(
SyncErrorClass::InvalidMaxBytes,
format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
);
}
let (data_dir, remote_lsn) = match sync_context(engine) {
Ok(context) => context,
Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
};
let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
Ok(status) => status,
Err(err) => {
return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
}
};
let Some(cursor_lsn) = status.last_applied_lsn else {
return SyncDecision::ok(Message::SyncPullResult {
status: wire_sync_status(status),
units: Vec::new(),
has_more: false,
});
};
if status.repair_action != SyncRepairAction::Pull {
return SyncDecision::ok(Message::SyncPullResult {
status: wire_sync_status(status),
units: Vec::new(),
has_more: false,
});
}
if request.since_lsn != cursor_lsn {
return SyncDecision::error(
SyncErrorClass::CursorLsnMismatch,
format!(
"sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
request.since_lsn
),
);
}
let identity = match read_identity(&data_dir) {
Ok(identity) => identity,
Err(err) => {
return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
}
};
let expected = identity.segment_identity();
if request.database_id != expected.database_id
|| request.primary_generation != expected.primary_generation
|| request.wal_format_version != expected.wal_format_version
|| request.catalog_version != expected.catalog_version
|| request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
{
return SyncDecision::error(
SyncErrorClass::IdentityOrFormatMismatch,
"sync pull identity or format version mismatch; rebootstrap required",
);
}
let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
let requested_through_lsn = request
.since_lsn
.saturating_add(request.max_units as u64)
.min(remote_lsn)
.min(status.servable_lsn.unwrap_or(request.since_lsn));
let segment_dir = retained_segments_dir(&data_dir);
if requested_through_lsn > request.since_lsn {
if let Err(err) = validate_retained_tail_available(
&segment_dir,
expected,
request.since_lsn,
requested_through_lsn,
) {
let mut rebootstrap_status = status;
rebootstrap_status.stale = true;
rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
rebootstrap_status.last_sync_error = Some(format!(
"retained history is unavailable; rebootstrap required: {err}"
));
return SyncDecision::ok(Message::SyncPullResult {
status: wire_sync_status(rebootstrap_status),
units: Vec::new(),
has_more: false,
});
}
}
let raw_units = match read_units_through(
&segment_dir,
expected,
request.since_lsn,
requested_through_lsn,
effective_max_units,
) {
Ok(units) => units,
Err(err) => {
return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
}
};
let mut selected_raw = Vec::new();
let mut selected = Vec::new();
let mut selected_bytes = 0u64;
for unit in raw_units {
let wire_unit = wire_retained_unit(unit.clone());
let unit_bytes = match wire_unit.encoded_len() {
Ok(bytes) => bytes,
Err(message) => {
return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
}
};
if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
if selected.is_empty() {
return SyncDecision::error(
SyncErrorClass::InvalidMaxBytes,
"sync pull maxBytes is too small for the next retained unit",
);
}
break;
}
selected_bytes += unit_bytes;
selected_raw.push(unit);
selected.push(wire_unit);
}
if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
}
let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
let has_more = selected
.last()
.is_some_and(|unit| unit.lsn < fetchable_through_lsn);
SyncDecision::ok(Message::SyncPullResult {
status: wire_sync_status(status),
units: selected,
has_more,
})
}
#[cfg(test)]
fn dispatch_sync_ack(
engine: &Arc<RwLock<Engine>>,
replica_id: String,
applied_lsn: u64,
observed_remote_lsn: u64,
credential_authenticated: bool,
principal: Option<&Principal>,
) -> Message {
dispatch_sync_ack_decision(
engine,
replica_id,
applied_lsn,
observed_remote_lsn,
credential_authenticated,
principal,
)
.message
}
fn dispatch_sync_ack_decision(
engine: &Arc<RwLock<Engine>>,
replica_id: String,
applied_lsn: u64,
observed_remote_lsn: u64,
credential_authenticated: bool,
principal: Option<&Principal>,
) -> SyncDecision {
if let Err((class, message)) =
check_sync_protocol_permitted(credential_authenticated, principal)
{
return SyncDecision::error(class, message);
}
if let Err(message) = validate_wire_replica_id(&replica_id) {
return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
}
if applied_lsn > observed_remote_lsn {
return SyncDecision::error(
SyncErrorClass::LsnAheadOfRemote,
format!(
"sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
),
);
}
let (data_dir, remote_lsn) = match sync_context(engine) {
Ok(context) => context,
Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
};
if observed_remote_lsn > remote_lsn {
return SyncDecision::error(
SyncErrorClass::LsnAheadOfRemote,
format!(
"sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
),
);
}
if let Err(message) =
validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
{
return SyncDecision::error(SyncErrorClass::AckValidation, message);
}
match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
previous_applied_lsn: summary.previous_applied_lsn,
applied_lsn: summary.applied_lsn,
remote_lsn: summary.remote_lsn,
advanced: summary.advanced,
status: wire_sync_status(summary.status),
}),
Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
}
}
async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
where
T: Send + 'static,
F: FnOnce(T) -> SyncDecision + Send + 'static,
{
let mut handle = tokio::task::spawn_blocking(move || f(input));
tokio::select! {
result = &mut handle => match result {
Ok(decision) => decision,
Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
},
_ = tokio::time::sleep(query_timeout) => match handle.await {
Ok(decision) => decision,
Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
},
}
}
async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
where
T: Send + 'static,
F: FnOnce(T) -> SyncDecision + Send + 'static,
{
let SyncExecutionContext {
tx_gate,
connection_has_transaction,
operation,
log_context,
metrics,
query_timeout,
} = context;
let start = Instant::now();
if connection_has_transaction {
let decision = SyncDecision::error(
SyncErrorClass::ActiveTransaction,
"sync protocol is unavailable inside an active transaction",
);
let elapsed = start.elapsed();
log_sync_decision(operation, &log_context, elapsed, &decision);
metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
return decision.message;
}
let permit = match tx_gate.acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
let decision =
SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
let elapsed = start.elapsed();
log_sync_decision(operation, &log_context, elapsed, &decision);
metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
return decision.message;
}
};
let decision = run_blocking_sync(input, query_timeout, f).await;
drop(permit);
match &decision.message {
Message::SyncStatusResult { status } => {
metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
}
Message::SyncPullResult { status, units, .. } => {
metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
}
Message::SyncAckResult {
advanced, status, ..
} => {
metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
if *advanced {
metrics.inc_sync_ack_advanced();
}
}
_ => {}
}
let elapsed = start.elapsed();
log_sync_decision(operation, &log_context, elapsed, &decision);
metrics.record_sync_operation(
operation,
elapsed,
sync_operation_outcome(&decision.message),
);
decision.message
}
async fn acquire_begin_permit(
tx_gate: &TxGate,
tx_wait_timeout: Duration,
metrics: &Arc<Metrics>,
) -> Result<OwnedSemaphorePermit, Message> {
match tokio::time::timeout(tx_wait_timeout, tx_gate.clone().acquire_owned()).await {
Ok(Ok(permit)) => Ok(permit),
Ok(Err(_)) => Err(Message::Error {
message: "query execution error".into(),
}),
Err(_) => {
metrics.inc_tx_gate_timeout();
Err(Message::Error {
message: format!(
"transaction gate timeout after {}ms waiting for concurrent transaction to complete",
tx_wait_timeout.as_millis()
),
})
}
}
}
#[allow(clippy::too_many_arguments)]
async fn execute_wire_query(
engine: Arc<RwLock<Engine>>,
tx_gate: TxGate,
tx_permit: &mut Option<OwnedSemaphorePermit>,
query: String,
principal: Option<Principal>,
query_timeout: Duration,
tx_wait_timeout: Duration,
metrics: &Arc<Metrics>,
) -> (Message, Option<PendingDurability>) {
match classify_query_transaction_control(&query) {
Some(TransactionControl::Begin) => {
if tx_permit.is_some() {
return (
Message::Error {
message: sanitize_error(
"cannot begin: a transaction is already active on this connection",
),
},
None,
);
}
let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
Ok(permit) => permit,
Err(response) => return (response, None),
};
let (response, ticket) = run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
)
.await;
if is_success_response(&response) {
*tx_permit = Some(permit);
}
(response, ticket)
}
Some(TransactionControl::Commit | TransactionControl::Rollback) => {
let (response, ticket) = run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
)
.await;
if is_success_response(&response) {
tx_permit.take();
}
(response, ticket)
}
None if tx_permit.is_some() => {
run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
)
.await
}
None => {
let permit = match tx_gate.acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
return (
Message::Error {
message: "query execution error".into(),
},
None,
)
}
};
let out = run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
)
.await;
drop(permit);
out
}
}
}
#[allow(clippy::too_many_arguments)]
async fn execute_wire_query_sql(
engine: Arc<RwLock<Engine>>,
tx_gate: TxGate,
tx_permit: &mut Option<OwnedSemaphorePermit>,
query: String,
principal: Option<Principal>,
query_timeout: Duration,
tx_wait_timeout: Duration,
metrics: &Arc<Metrics>,
) -> (Message, Option<PendingDurability>) {
match classify_sql_transaction_control(&query) {
Some(TransactionControl::Begin) => {
if tx_permit.is_some() {
return (
Message::Error {
message: sanitize_error(
"cannot begin: a transaction is already active on this connection",
),
},
None,
);
}
let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
Ok(permit) => permit,
Err(response) => return (response, None),
};
let (response, ticket) = run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
)
.await;
if is_success_response(&response) {
*tx_permit = Some(permit);
}
(response, ticket)
}
Some(TransactionControl::Commit | TransactionControl::Rollback) => {
let (response, ticket) = run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
)
.await;
if is_success_response(&response) {
tx_permit.take();
}
(response, ticket)
}
None if tx_permit.is_some() => {
run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
)
.await
}
None => {
let permit = match tx_gate.acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
return (
Message::Error {
message: "query execution error".into(),
},
None,
)
}
};
let out = run_blocking_query(
engine,
query,
principal,
query_timeout,
metrics,
|engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
)
.await;
drop(permit);
out
}
}
}
#[allow(clippy::too_many_arguments)]
async fn execute_wire_query_with_params(
engine: Arc<RwLock<Engine>>,
tx_gate: TxGate,
tx_permit: &mut Option<OwnedSemaphorePermit>,
query: String,
params: Vec<WireParam>,
principal: Option<Principal>,
query_timeout: Duration,
tx_wait_timeout: Duration,
metrics: &Arc<Metrics>,
) -> (Message, Option<PendingDurability>) {
match classify_params_transaction_control(&query, ¶ms) {
Some(TransactionControl::Begin) => {
if tx_permit.is_some() {
return (
Message::Error {
message: sanitize_error(
"cannot begin: a transaction is already active on this connection",
),
},
None,
);
}
let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
Ok(permit) => permit,
Err(response) => return (response, None),
};
let (response, ticket) = run_blocking_query(
engine,
(query, params),
principal,
query_timeout,
metrics,
|engine, (query, params), principal| {
dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
},
)
.await;
if is_success_response(&response) {
*tx_permit = Some(permit);
}
(response, ticket)
}
Some(TransactionControl::Commit | TransactionControl::Rollback) => {
let (response, ticket) = run_blocking_query(
engine,
(query, params),
principal,
query_timeout,
metrics,
|engine, (query, params), principal| {
dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
},
)
.await;
if is_success_response(&response) {
tx_permit.take();
}
(response, ticket)
}
None if tx_permit.is_some() => {
run_blocking_query(
engine,
(query, params),
principal,
query_timeout,
metrics,
|engine, (query, params), principal| {
dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
},
)
.await
}
None => {
let permit = match tx_gate.acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
return (
Message::Error {
message: "query execution error".into(),
},
None,
)
}
};
let out = run_blocking_query(
engine,
(query, params),
principal,
query_timeout,
metrics,
|engine, (query, params), principal| {
dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
},
)
.await;
drop(permit);
out
}
}
}
struct DeferredQueryMetric {
start: Instant,
outcome: QueryOutcome,
exceeded_timeout: bool,
}
type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
async fn run_blocking_query<T, F>(
engine: Arc<RwLock<Engine>>,
input: T,
principal: Option<Principal>,
query_timeout: Duration,
metrics: &Arc<Metrics>,
f: F,
) -> (Message, Option<PendingDurability>)
where
T: Send + 'static,
F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
{
let _in_flight = metrics.in_flight_guard();
let start = Instant::now();
let mut handle = tokio::task::spawn_blocking(move || f(engine, input, principal));
let mut exceeded_timeout = false;
let join_result = tokio::select! {
result = &mut handle => result,
_ = tokio::time::sleep(query_timeout) => {
exceeded_timeout = true;
handle.await
}
};
let (message, ticket, outcome) = match join_result {
Ok((Ok(result), ticket)) => match query_result_to_message(result) {
Ok(message) => (message, ticket, QueryOutcome::Ok),
Err(e) => (
Message::Error {
message: sanitize_error(&e.to_string()),
},
ticket,
QueryOutcome::Error,
),
},
Ok((Err(e), ticket)) => {
let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
QueryOutcome::MemoryLimit
} else {
QueryOutcome::Error
};
(
Message::Error {
message: sanitize_error(&e.to_string()),
},
ticket,
outcome,
)
}
Err(e) => (
Message::Error {
message: format!("internal error: {e}"),
},
None,
QueryOutcome::Error,
),
};
match ticket {
Some(ticket) => (
message,
Some((
ticket,
DeferredQueryMetric {
start,
outcome,
exceeded_timeout,
},
)),
),
None => {
if exceeded_timeout {
metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
} else {
metrics.record_query(start.elapsed(), outcome);
}
(message, None)
}
}
}
async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
match tokio::task::spawn_blocking(move || ticket.wait()).await {
Ok(Ok(())) => None,
Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
Err(e) => Some(format!("internal error: {e}")),
}
}
fn is_success_response(msg: &Message) -> bool {
matches!(
msg,
Message::ResultRows { .. }
| Message::ResultScalar { .. }
| Message::ResultOk { .. }
| Message::ResultMessage { .. }
)
}
fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref());
let _ = res;
if let Some(ticket) = ticket {
let _ = ticket.wait();
}
}
pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
where
S: AsyncRead + AsyncWrite + Unpin,
{
let ConnOpts {
engine,
tx_gate,
expected_password,
users,
shutdown_rx,
idle_timeout,
query_timeout,
rate_limiter,
peer_addr,
metrics,
tx_wait_timeout,
db_name: server_db_name,
} = opts;
let peer = peer_addr
.map(|a| a.to_string())
.unwrap_or_else(|| "unknown".into());
let peer_ip = peer_addr.map(|a| a.ip());
let (reader, writer) = tokio::io::split(stream);
let mut reader = BufReader::new(reader);
let mut writer = BufWriter::new(writer);
let connect_msg = loop {
match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
Ok(Ok(Some(Message::Ping))) => {
debug!(peer = %peer, "pre-auth ping");
if !write_msg(&mut writer, &Message::Pong).await {
return;
}
continue;
}
Ok(Ok(Some(msg))) => break msg,
Ok(Ok(None)) => {
debug!(peer = %peer, "client closed before CONNECT");
return;
}
Ok(Err(e)) => {
error!(peer = %peer, error = %e, "error reading CONNECT");
return;
}
Err(_) => {
warn!(peer = %peer, "idle timeout waiting for CONNECT");
return;
}
}
};
let principal: Option<Principal>;
let credential_auth_configured = !users.is_empty() || expected_password.is_some();
match connect_msg {
Message::Connect {
db_name,
password,
username,
} => {
if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
if is_rate_limited(limiter, ip) {
warn!(peer = %peer, "rate limited: too many auth failures");
let err = Message::Error {
message: "too many auth failures, try again later".into(),
};
write_msg(&mut writer, &err).await;
return;
}
}
let outcome = authenticate_connect(
&users,
expected_password.as_ref().map(|p| p.as_str()),
username.as_deref(),
password.as_ref().map(|p| p.as_str()),
);
match outcome {
AuthOutcome::Rejected => {
warn!(peer = %peer, db = %db_name, "auth rejected");
metrics.inc_auth_failure();
if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
record_auth_failure(limiter, ip);
}
let err = Message::Error {
message: "authentication failed".into(),
};
write_msg(&mut writer, &err).await;
return;
}
AuthOutcome::Authenticated {
principal: auth_principal,
} => {
if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
clear_auth_failures(limiter, ip);
}
match &auth_principal {
Some(p) => {
info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
}
None => {
info!(peer = %peer, db = %db_name, "client connected");
}
}
principal = auth_principal;
}
}
match check_db_name(server_db_name.as_deref(), &db_name) {
Ok(()) => {
if server_db_name.is_none() && !db_name.is_empty() && db_name != DEFAULT_DB_NAME
{
static NAMED_DB_WARNED: std::sync::atomic::AtomicBool =
std::sync::atomic::AtomicBool::new(false);
if !NAMED_DB_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
warn!(
peer = %peer, db = %db_name,
"client requested a named database but this server serves a single global database; name ignored"
);
}
}
}
Err(msg) => {
warn!(peer = %peer, db = %db_name, "rejected: unknown database");
let err = Message::Error { message: msg };
write_msg(&mut writer, &err).await;
return;
}
}
let ok = Message::ConnectOk {
version: env!("CARGO_PKG_VERSION").into(),
};
if !write_msg(&mut writer, &ok).await {
return;
}
}
_ => {
warn!(peer = %peer, "first message was not CONNECT");
let err = Message::Error {
message: "expected CONNECT".into(),
};
write_msg(&mut writer, &err).await;
return;
}
}
let mut tx_permit: Option<OwnedSemaphorePermit> = None;
let mut carry: Option<Message> = None;
'conn: loop {
let msg = if let Some(m) = carry.take() {
m
} else {
tokio::select! {
result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
match result {
Ok(Ok(Some(msg))) => msg,
Ok(Ok(None)) => break,
Ok(Err(e)) => {
error!(peer = %peer, error = %e, "read error");
break;
}
Err(_) => {
info!(peer = %peer, "idle timeout, closing connection");
let err = Message::Error { message: "idle timeout".into() };
write_msg(&mut writer, &err).await;
break;
}
}
}
_ = shutdown_rx.changed() => {
if *shutdown_rx.borrow() {
info!(peer = %peer, "server shutting down, closing connection");
let err = Message::Error { message: "server shutting down".into() };
write_msg(&mut writer, &err).await;
break;
}
continue;
}
}
};
if matches!(
msg,
Message::Query { .. } | Message::QuerySql { .. } | Message::QueryWithParams { .. }
) {
const MAX_PIPELINE_BATCH: usize = 128;
const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
enum BatchFatal {
Closed,
ReadError,
}
fn approx_response_bytes(msg: &Message) -> usize {
match msg {
Message::ResultRows { columns, rows } => {
columns.iter().map(|c| c.len() + 4).sum::<usize>()
+ rows
.iter()
.map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
.sum::<usize>()
}
Message::ResultScalar { value } => value.len(),
Message::ResultMessage { message } | Message::Error { message } => {
message.len()
}
_ => 16,
}
}
fn complete_frame_buffered(buf: &[u8]) -> bool {
buf.len() >= 6 && {
let payload_len =
u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
buf.len() - 6 >= payload_len
}
}
let mut responses: Vec<Message> = Vec::new();
let mut response_bytes: usize = 0;
let mut last_ticket: Option<WalDurabilityTicket> = None;
let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
let mut fatal: Option<BatchFatal> = None;
let mut current = msg;
loop {
let (response, ticket) = match current {
Message::Query { query } => {
if query.len() > MAX_QUERY_LENGTH {
(
Message::Error {
message: format!(
"query too large: {} bytes (max {})",
query.len(),
MAX_QUERY_LENGTH
),
},
None,
)
} else {
debug!(peer = %peer, query = %query, "received query");
execute_wire_query(
engine.clone(),
tx_gate.clone(),
&mut tx_permit,
query,
principal.clone(),
query_timeout,
tx_wait_timeout,
&metrics,
)
.await
}
}
Message::QuerySql { query } => {
if query.len() > MAX_QUERY_LENGTH {
(
Message::Error {
message: format!(
"query too large: {} bytes (max {})",
query.len(),
MAX_QUERY_LENGTH
),
},
None,
)
} else {
debug!(peer = %peer, query = %query, "received SQL query");
execute_wire_query_sql(
engine.clone(),
tx_gate.clone(),
&mut tx_permit,
query,
principal.clone(),
query_timeout,
tx_wait_timeout,
&metrics,
)
.await
}
}
Message::QueryWithParams { query, params } => {
if query.len() > MAX_QUERY_LENGTH {
(
Message::Error {
message: format!(
"query too large: {} bytes (max {})",
query.len(),
MAX_QUERY_LENGTH
),
},
None,
)
} else {
debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
execute_wire_query_with_params(
engine.clone(),
tx_gate.clone(),
&mut tx_permit,
query,
params,
principal.clone(),
query_timeout,
tx_wait_timeout,
&metrics,
)
.await
}
}
_ => unreachable!("batch loop only receives plain query frames"),
};
if let Some((t, m)) = ticket {
last_ticket = Some(t);
deferred_metrics.push(m);
}
response_bytes += approx_response_bytes(&response);
responses.push(response);
if tx_permit.is_some()
|| responses.len() >= MAX_PIPELINE_BATCH
|| response_bytes >= MAX_PIPELINE_BATCH_BYTES
|| !complete_frame_buffered(reader.buffer())
{
break;
}
match tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)).await {
Ok(Ok(Some(
next @ (Message::Query { .. }
| Message::QuerySql { .. }
| Message::QueryWithParams { .. }),
))) => {
if tx_gate.available_permits() == 0 {
carry = Some(next);
break;
}
current = next;
}
Ok(Ok(Some(other))) => {
carry = Some(other);
break;
}
Ok(Ok(None)) => {
fatal = Some(BatchFatal::Closed);
break;
}
Ok(Err(e)) => {
error!(peer = %peer, error = %e, "read error");
fatal = Some(BatchFatal::ReadError);
break;
}
Err(_) => {
error!(peer = %peer, "timeout decoding fully-buffered frame");
fatal = Some(BatchFatal::ReadError);
break;
}
}
}
let mut durability_failed = false;
if let Some(ticket) = last_ticket {
if let Some(message) = settle_durability_ticket(ticket).await {
durability_failed = true;
for r in responses.iter_mut() {
if is_success_response(r) {
*r = Message::Error {
message: message.clone(),
};
}
}
}
}
for m in deferred_metrics.drain(..) {
let outcome = if m.exceeded_timeout {
QueryOutcome::Timeout
} else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
QueryOutcome::Error
} else {
m.outcome
};
metrics.record_query(m.start.elapsed(), outcome);
}
for r in &responses {
if !write_msg(&mut writer, r).await {
break 'conn;
}
}
match fatal {
None => continue,
Some(BatchFatal::Closed | BatchFatal::ReadError) => break,
}
}
let response = match msg {
Message::Ping => {
debug!(peer = %peer, "ping");
Message::Pong
}
Message::SyncStatus { replica_id } => {
let engine = engine.clone();
let principal = principal.clone();
let log_context = SyncLogContext::status(&replica_id);
execute_gated_sync(
SyncExecutionContext {
tx_gate: tx_gate.clone(),
connection_has_transaction: tx_permit.is_some(),
operation: SyncOperation::Status,
log_context,
metrics: &metrics,
query_timeout,
},
(engine, replica_id, credential_auth_configured, principal),
|(engine, replica_id, credential_authenticated, principal)| {
dispatch_sync_status_decision(
&engine,
replica_id,
credential_authenticated,
principal.as_ref(),
)
},
)
.await
}
Message::SyncPull {
replica_id,
since_lsn,
max_units,
max_bytes,
database_id,
primary_generation,
wal_format_version,
catalog_version,
segment_format_version,
} => {
let engine = engine.clone();
let principal = principal.clone();
let request = SyncPullRequest {
replica_id,
since_lsn,
max_units,
max_bytes,
database_id,
primary_generation,
wal_format_version,
catalog_version,
segment_format_version,
};
let log_context = SyncLogContext::pull(&request);
execute_gated_sync(
SyncExecutionContext {
tx_gate: tx_gate.clone(),
connection_has_transaction: tx_permit.is_some(),
operation: SyncOperation::Pull,
log_context,
metrics: &metrics,
query_timeout,
},
(engine, request, credential_auth_configured, principal),
|(engine, request, credential_authenticated, principal)| {
dispatch_sync_pull_decision(
&engine,
request,
credential_authenticated,
principal.as_ref(),
)
},
)
.await
}
Message::SyncAck {
replica_id,
applied_lsn,
remote_lsn,
} => {
let engine = engine.clone();
let principal = principal.clone();
let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
execute_gated_sync(
SyncExecutionContext {
tx_gate: tx_gate.clone(),
connection_has_transaction: tx_permit.is_some(),
operation: SyncOperation::Ack,
log_context,
metrics: &metrics,
query_timeout,
},
(
engine,
replica_id,
applied_lsn,
remote_lsn,
credential_auth_configured,
principal,
),
|(
engine,
replica_id,
applied_lsn,
observed_remote_lsn,
credential_authenticated,
principal,
)| {
dispatch_sync_ack_decision(
&engine,
replica_id,
applied_lsn,
observed_remote_lsn,
credential_authenticated,
principal.as_ref(),
)
},
)
.await
}
Message::Disconnect => {
debug!(peer = %peer, "received DISCONNECT");
break;
}
_ => Message::Error {
message: "unexpected message type".into(),
},
};
if !write_msg(&mut writer, &response).await {
break;
}
}
if tx_permit.is_some() {
let engine = engine.clone();
let principal = principal.clone();
let _ =
tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
}
tx_permit.take();
info!(peer = %peer, "client disconnected");
}
fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
*total = total.saturating_add(bytes);
if *total > MAX_RESPONSE_PAYLOAD_SIZE {
return Err(QueryError::Execution(format!(
"result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
MAX_RESPONSE_PAYLOAD_SIZE
)));
}
Ok(())
}
fn query_result_to_message(result: QueryResult) -> Result<Message, QueryError> {
match result {
QueryResult::Rows { columns, rows } => {
let mut encoded_bytes = 2usize; let mut out_columns = Vec::with_capacity(columns.len());
for col in columns {
charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
out_columns.push(col);
}
charge_response_bytes(&mut encoded_bytes, 4)?;
let mut str_rows = Vec::with_capacity(rows.len());
for row in rows {
let mut str_row = Vec::with_capacity(row.len());
for value in row {
let display = value_to_display(&value);
charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
str_row.push(display);
}
str_rows.push(str_row);
}
Ok(Message::ResultRows {
columns: out_columns,
rows: str_rows,
})
}
QueryResult::Scalar(val) => Ok(Message::ResultScalar {
value: value_to_display(&val),
}),
QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
QueryResult::Created(name) => Ok(Message::ResultMessage {
message: format!("type {name} created"),
}),
QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
}
}
fn value_to_display(v: &Value) -> String {
v.to_wire_string()
}
#[cfg(test)]
mod tests {
use super::*;
use powdb_storage::wal::WalRecordType;
use powdb_sync::{
write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
ReplicaCursor, RetainedSegment, RetainedUnit,
};
#[test]
fn null_serializes_as_null_bareword_on_wire() {
assert_eq!(value_to_display(&Value::Empty), "null");
}
#[test]
fn unique_violation_error_surfaces_to_remote_clients() {
assert_eq!(
sanitize_error("unique constraint violation on User.email"),
"unique constraint violation on User.email"
);
}
#[test]
fn internal_errors_stay_generic() {
assert_eq!(
sanitize_error("some internal io panic detail"),
"query execution error"
);
}
#[test]
fn db_name_unpinned_accepts_any_name() {
for requested in ["", "default", "prod", "anything"] {
assert!(
check_db_name(None, requested).is_ok(),
"rejected {requested}"
);
}
}
#[test]
fn db_name_pinned_accepts_match_empty_and_default_sentinel() {
assert!(check_db_name(Some("prod"), "prod").is_ok());
assert!(check_db_name(Some("prod"), "").is_ok());
assert!(check_db_name(Some("prod"), DEFAULT_DB_NAME).is_ok());
}
#[test]
fn db_name_pinned_rejects_foreign_with_clear_message() {
let err = check_db_name(Some("prod"), "staging").unwrap_err();
assert_eq!(err, "unknown database 'staging'; this server serves 'prod'");
}
#[tokio::test]
async fn begin_permit_acquires_when_gate_is_free() {
let gate = new_tx_gate();
let metrics = Arc::new(Metrics::new());
let permit = acquire_begin_permit(&gate, Duration::from_secs(5), &metrics)
.await
.expect("should acquire a free gate");
assert_eq!(gate.available_permits(), 0, "permit must be held");
drop(permit);
assert_eq!(gate.available_permits(), 1, "permit must release on drop");
}
#[tokio::test]
async fn begin_permit_times_out_with_clear_error_and_truthful_metric() {
let gate = new_tx_gate();
let metrics = Arc::new(Metrics::new());
let _held = gate.clone().acquire_owned().await.unwrap();
let err = acquire_begin_permit(&gate, Duration::from_millis(25), &metrics)
.await
.expect_err("must time out while the gate is held");
match err {
Message::Error { message } => {
assert!(
message.contains("transaction gate timeout after 25ms"),
"unexpected message: {message}"
);
assert!(
message.contains("waiting for concurrent transaction"),
"unexpected message: {message}"
);
}
other => panic!("expected Error, got {other:?}"),
}
let rendered = metrics.render();
assert!(rendered.contains("powdb_tx_gate_timeouts_total 1"));
assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
}
#[test]
fn resource_limit_errors_surface_actionable_hints() {
for msg in [
"sort input exceeds row limit — add a LIMIT clause",
"join result exceeds row limit",
"query exceeded memory budget: requested 100 bytes, limit 50 bytes",
"result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
] {
assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
}
}
#[test]
fn oversized_result_is_rejected_before_wire_encoding() {
let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
let result = QueryResult::Rows {
columns: vec!["payload".into()],
rows: vec![vec![Value::Str(long)]],
};
let err = query_result_to_message(result).unwrap_err();
assert!(
err.to_string().starts_with("result too large"),
"unexpected error: {err}"
);
}
fn parsed(q: &str) -> powdb_query::ast::Statement {
parser::parse(q).unwrap()
}
fn principal(role: &str) -> Option<Principal> {
Some(Principal {
name: "u".into(),
role: role.into(),
})
}
#[test]
fn readonly_can_read_but_not_write() {
let p = principal("readonly");
assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
for q in [
r#"insert User { name := "x" }"#,
"User filter .id = 1 update { age := 2 }",
"User filter .id = 1 delete",
"drop User",
"alter User add column c: str",
"type T { required id: int }",
"begin",
"commit",
"rollback",
] {
let err = check_statement_permitted(p.as_ref(), &parsed(q))
.expect_err(&format!("must deny: {q}"));
assert!(
err.to_string().contains("permission denied"),
"unexpected error for {q}: {err}"
);
}
}
#[test]
fn readwrite_and_admin_have_full_query_access() {
for role in ["readwrite", "admin"] {
let p = principal(role);
assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
assert!(check_statement_permitted(
p.as_ref(),
&parsed(r#"insert User { name := "x" }"#)
)
.is_ok());
assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
}
}
#[test]
fn unknown_role_fails_closed_for_writes() {
let p = principal("mystery");
assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
assert!(
check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
.is_err()
);
}
#[test]
fn no_principal_means_full_access() {
assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
}
fn store_with_alice() -> UserStore {
let mut s = UserStore::new();
s.create_user("alice", "pw", "readwrite").unwrap();
s
}
#[test]
fn empty_store_no_password_is_open() {
let s = UserStore::new();
assert_eq!(
authenticate_connect(&s, None, None, None),
AuthOutcome::Authenticated { principal: None }
);
assert_eq!(
authenticate_connect(&s, None, Some("x"), Some("y")),
AuthOutcome::Authenticated { principal: None }
);
}
#[test]
fn empty_store_correct_shared_password_succeeds() {
let s = UserStore::new();
assert_eq!(
authenticate_connect(&s, Some("pw"), None, Some("pw")),
AuthOutcome::Authenticated { principal: None }
);
}
#[test]
fn empty_store_wrong_shared_password_rejected() {
let s = UserStore::new();
assert_eq!(
authenticate_connect(&s, Some("pw"), None, Some("bad")),
AuthOutcome::Rejected
);
}
#[test]
fn empty_store_missing_password_rejected_when_expected() {
let s = UserStore::new();
assert_eq!(
authenticate_connect(&s, Some("pw"), None, None),
AuthOutcome::Rejected
);
}
#[test]
fn empty_store_ignores_username_for_shared_password() {
let s = UserStore::new();
assert_eq!(
authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
AuthOutcome::Authenticated { principal: None }
);
}
#[test]
fn user_auth_success_binds_principal() {
let s = store_with_alice();
assert_eq!(
authenticate_connect(&s, None, Some("alice"), Some("pw")),
AuthOutcome::Authenticated {
principal: Some(Principal {
name: "alice".into(),
role: "readwrite".into(),
})
}
);
}
#[test]
fn user_auth_wrong_password_rejected() {
let s = store_with_alice();
assert_eq!(
authenticate_connect(&s, None, Some("alice"), Some("bad")),
AuthOutcome::Rejected
);
}
#[test]
fn user_auth_unknown_user_rejected() {
let s = store_with_alice();
assert_eq!(
authenticate_connect(&s, None, Some("mallory"), Some("pw")),
AuthOutcome::Rejected
);
}
#[test]
fn user_auth_missing_username_rejected() {
let s = store_with_alice();
assert_eq!(
authenticate_connect(&s, None, None, Some("pw")),
AuthOutcome::Rejected
);
}
#[test]
fn user_auth_missing_password_rejected() {
let s = store_with_alice();
assert_eq!(
authenticate_connect(&s, Some("pw"), Some("alice"), None),
AuthOutcome::Rejected
);
}
#[test]
fn user_auth_ignores_shared_password_when_users_present() {
let s = store_with_alice();
assert_eq!(
authenticate_connect(&s, Some("shared"), None, Some("shared")),
AuthOutcome::Rejected
);
}
#[test]
fn replica_fingerprint_is_stable_and_redacted() {
let replica_id = "customer-prod-replica-a";
let fingerprint = replica_fingerprint(replica_id);
assert_eq!(fingerprint, replica_fingerprint(replica_id));
assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
assert_eq!(fingerprint.len(), 16);
assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
assert!(!fingerprint.contains("customer"));
assert!(!fingerprint.contains("replica"));
assert!(!fingerprint.contains(replica_id));
}
#[test]
fn invalid_replica_ids_use_fixed_log_fingerprint() {
assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
assert_eq!(
log_replica_fingerprint("customer/prod/replica"),
INVALID_REPLICA_FINGERPRINT
);
assert_eq!(
log_replica_fingerprint(&"a".repeat(4096)),
INVALID_REPLICA_FINGERPRINT
);
}
#[test]
fn sync_error_classes_use_bounded_labels() {
assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
assert_eq!(
SyncErrorClass::PermissionDenied.as_label(),
"permission_denied"
);
assert_eq!(
SyncErrorClass::IdentityOrFormatMismatch.as_label(),
"identity_or_format_mismatch"
);
assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
}
fn sync_identity() -> DatabaseIdentity {
DatabaseIdentity {
database_id: *b"server-sync-test",
primary_generation: 1,
}
}
fn retained_unit(lsn: u64) -> RetainedUnit {
RetainedUnit {
tx_id: 1,
record_type: 4,
lsn,
data: lsn.to_le_bytes().to_vec(),
}
}
fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
RetainedUnit {
tx_id,
record_type: record_type as u8,
lsn,
data: lsn.to_le_bytes().to_vec(),
}
}
fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
let identity = sync_identity();
write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
let units = (1..=through_lsn).map(retained_unit).collect();
let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
}
fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
let identity = sync_identity();
write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
}
fn write_sync_identity_only(data_dir: &std::path::Path) {
let identity = sync_identity();
write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
}
fn admin_principal() -> Principal {
Principal {
name: "admin".into(),
role: "admin".into(),
}
}
#[test]
fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
let dir = tempfile::tempdir().unwrap();
let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
Message::Error { message } => {
assert!(message.contains("requires authentication"));
}
other => panic!("expected auth error, got {other:?}"),
}
let readonly = Principal {
name: "reader".into(),
role: "readonly".into(),
};
match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
Message::Error { message } => {
assert!(message.contains("permission denied"));
}
other => panic!("expected permission error, got {other:?}"),
}
}
#[test]
fn sync_status_pull_and_ack_use_server_remote_lsn() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int, v: str }")
.unwrap();
engine
.execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
.unwrap();
let remote_lsn = engine.catalog().max_lsn();
assert!(remote_lsn > 0);
write_sync_identity_and_tail(dir.path(), remote_lsn);
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
{
Message::SyncStatusResult { status } => status,
other => panic!("expected sync status, got {other:?}"),
};
assert_eq!(status.remote_lsn, remote_lsn);
assert_eq!(status.servable_lsn, Some(remote_lsn));
assert_eq!(status.unarchived_lsn, Some(0));
assert_eq!(status.last_applied_lsn, Some(0));
assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
assert!(status.stale);
let identity = sync_identity().segment_identity();
let pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: MAX_SYNC_PULL_UNITS,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
Message::SyncPullResult {
status,
units,
has_more,
} => {
assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
assert!(!has_more);
units
}
other => panic!("expected sync pull result, got {other:?}"),
};
assert_eq!(units.len() as u64, remote_lsn);
assert_eq!(units.last().unwrap().lsn, remote_lsn);
let ack = match dispatch_sync_ack(
&engine,
"replica-a".into(),
remote_lsn,
remote_lsn,
true,
Some(&principal),
) {
Message::SyncAckResult {
previous_applied_lsn,
applied_lsn,
remote_lsn: ack_remote_lsn,
advanced,
status,
} => {
assert_eq!(previous_applied_lsn, 0);
assert_eq!(applied_lsn, remote_lsn);
assert_eq!(ack_remote_lsn, remote_lsn);
assert!(advanced);
status
}
other => panic!("expected sync ack result, got {other:?}"),
};
assert_eq!(ack.repair_action, WireSyncRepairAction::None);
assert!(!ack.stale);
assert_eq!(ack.lag_lsn, Some(0));
}
#[test]
fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int }")
.unwrap();
for id in 1..=3 {
engine
.execute_powql(&format!("insert SyncT {{ id := {id} }}"))
.unwrap();
}
let remote_lsn = engine.catalog().max_lsn();
assert!(remote_lsn >= 3);
write_sync_identity_and_units(
dir.path(),
vec![
retained_unit_with(77, WalRecordType::Begin, 1),
retained_unit_with(77, WalRecordType::Insert, 2),
retained_unit_with(77, WalRecordType::Commit, 3),
],
);
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let identity = sync_identity().segment_identity();
let cut_pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: 2,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
Message::Error { message } => assert!(message.contains("cuts through transaction")),
other => panic!("expected transaction-cut pull error, got {other:?}"),
}
let cut_bytes_pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: 3,
max_bytes: 58,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
Message::Error { message } => assert!(message.contains("cuts through transaction")),
other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
}
let full_pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: 3,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
Message::SyncPullResult { units, .. } => {
assert_eq!(units.len(), 3);
assert_eq!(units.last().unwrap().lsn, 3);
}
other => panic!("expected complete transaction pull, got {other:?}"),
}
match dispatch_sync_ack(
&engine,
"replica-a".into(),
2,
remote_lsn,
true,
Some(&principal),
) {
Message::Error { message } => assert!(message.contains("cuts through transaction")),
other => panic!("expected transaction-cut ack error, got {other:?}"),
}
let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
assert_eq!(cursor[0].applied_lsn, 0);
match dispatch_sync_ack(
&engine,
"replica-a".into(),
3,
remote_lsn,
true,
Some(&principal),
) {
Message::SyncAckResult {
previous_applied_lsn,
applied_lsn,
advanced,
..
} => {
assert_eq!(previous_applied_lsn, 0);
assert_eq!(applied_lsn, 3);
assert!(advanced);
}
other => panic!("expected complete transaction ack, got {other:?}"),
}
}
#[test]
fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int }")
.unwrap();
for id in 1..=6 {
engine
.execute_powql(&format!("insert SyncT {{ id := {id} }}"))
.unwrap();
}
let remote_lsn = engine.catalog().max_lsn();
assert!(remote_lsn >= 6);
write_sync_identity_and_units(
dir.path(),
vec![
retained_unit_with(1, WalRecordType::Begin, 1),
retained_unit_with(1, WalRecordType::Insert, 2),
retained_unit_with(1, WalRecordType::Commit, 3),
retained_unit_with(1, WalRecordType::Begin, 4),
retained_unit_with(1, WalRecordType::Insert, 5),
retained_unit_with(1, WalRecordType::Commit, 6),
],
);
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let identity = sync_identity().segment_identity();
let pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: 6,
max_bytes: 100,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
Message::SyncPullResult {
status,
units,
has_more,
} => {
assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
assert_eq!(units.len(), 3);
assert_eq!(units.last().unwrap().lsn, 3);
assert!(has_more);
}
other => panic!("expected byte-capped applyable prefix, got {other:?}"),
}
}
#[test]
fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int }")
.unwrap();
engine.execute_powql("insert SyncT { id := 1 }").unwrap();
let remote_lsn = engine.catalog().max_lsn();
assert!(remote_lsn > 0);
write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let identity = sync_identity().segment_identity();
let pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: MAX_SYNC_PULL_UNITS,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
Message::SyncPullResult {
status,
units,
has_more,
} => {
assert_eq!(status.remote_lsn, remote_lsn);
assert_eq!(status.servable_lsn, Some(remote_lsn));
assert_eq!(status.unarchived_lsn, Some(0));
assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
assert!(!has_more);
assert_eq!(units.len() as u64, remote_lsn);
assert_eq!(units.last().unwrap().lsn, remote_lsn);
assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
}
other => panic!("expected capped sync pull result, got {other:?}"),
}
}
#[test]
fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int }")
.unwrap();
engine.execute_powql("insert SyncT { id := 1 }").unwrap();
let remote_lsn = engine.catalog().max_lsn();
assert!(remote_lsn > 0);
write_sync_identity_only(dir.path());
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let identity = sync_identity().segment_identity();
let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
{
Message::SyncStatusResult { status } => status,
other => panic!("expected sync status, got {other:?}"),
};
assert_eq!(status.remote_lsn, remote_lsn);
assert_eq!(status.servable_lsn, Some(0));
assert_eq!(status.unarchived_lsn, Some(remote_lsn));
assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
assert!(status
.last_sync_error
.as_deref()
.unwrap()
.contains("not yet archived"));
let pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: MAX_SYNC_PULL_UNITS,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
Message::SyncPullResult {
status,
units,
has_more,
} => {
assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
assert!(units.is_empty());
assert!(!has_more);
}
other => panic!("expected await-archive sync pull result, got {other:?}"),
}
}
#[test]
fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int }")
.unwrap();
engine.execute_powql("insert SyncT { id := 1 }").unwrap();
engine.execute_powql("insert SyncT { id := 2 }").unwrap();
let remote_lsn = engine.catalog().max_lsn();
assert!(remote_lsn > 1);
let servable_lsn = remote_lsn - 1;
write_sync_identity_and_tail(dir.path(), servable_lsn);
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let identity = sync_identity().segment_identity();
let pull = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: MAX_SYNC_PULL_UNITS,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
Message::SyncPullResult {
status,
units,
has_more,
} => {
assert_eq!(status.remote_lsn, remote_lsn);
assert_eq!(status.servable_lsn, Some(servable_lsn));
assert_eq!(status.unarchived_lsn, Some(1));
assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
assert!(!has_more);
assert_eq!(units.len() as u64, servable_lsn);
assert_eq!(units.last().unwrap().lsn, servable_lsn);
assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
}
other => panic!("expected partial sync pull result, got {other:?}"),
}
}
#[test]
fn sync_pull_rejects_cursor_or_format_mismatch() {
let dir = tempfile::tempdir().unwrap();
let mut engine = Engine::new(dir.path()).unwrap();
engine
.execute_powql("type SyncT { required id: int }")
.unwrap();
engine.execute_powql("insert SyncT { id := 1 }").unwrap();
let remote_lsn = engine.catalog().max_lsn();
write_sync_identity_and_tail(dir.path(), remote_lsn);
powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
.unwrap();
let engine = Arc::new(RwLock::new(engine));
let principal = admin_principal();
let identity = sync_identity().segment_identity();
let wrong_cursor = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 1,
max_units: 10,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
};
match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
Message::Error { message } => assert!(message.contains("does not match")),
other => panic!("expected cursor mismatch error, got {other:?}"),
}
let wrong_format = SyncPullRequest {
replica_id: "replica-a".into(),
since_lsn: 0,
max_units: 10,
max_bytes: MAX_SYNC_PULL_BYTES,
database_id: identity.database_id,
primary_generation: identity.primary_generation,
wal_format_version: identity.wal_format_version,
catalog_version: identity.catalog_version,
segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
};
match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
Message::Error { message } => assert!(message.contains("rebootstrap required")),
other => panic!("expected format mismatch error, got {other:?}"),
}
}
}