1use crate::metrics::{Metrics, QueryOutcome, SyncOperation, SyncOutcome, SyncRepairLabel};
2use crate::protocol::{Message, WireParam, WireRetainedUnit, WireSyncRepairAction, WireSyncStatus};
3use powdb_auth::{Permission, Role, UserStore};
4use powdb_query::executor::{is_read_only_statement, Engine, WalDurabilityTicket};
5use powdb_query::parser;
6use powdb_query::result::{QueryError, QueryResult};
7use powdb_query::sql;
8use powdb_storage::types::Value;
9use powdb_sync::{
10 acknowledge_replica_apply, read_identity, read_units_through, replica_sync_status,
11 retained_segments_dir, validate_retained_tail_available, validate_v1_retained_units_applyable,
12 ReplicaSyncStatus, RetainedUnit, SyncRepairAction, RETAINED_SEGMENT_FORMAT_VERSION,
13};
14use std::collections::HashMap;
15use std::net::IpAddr;
16use std::path::{Path, PathBuf};
17use std::sync::{Arc, Mutex, RwLock};
18use std::time::{Duration, Instant};
19use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
20use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};
21use tracing::{debug, error, info, warn};
22use zeroize::Zeroizing;
23
24pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
26
27pub type TxGate = Arc<Semaphore>;
32
33pub fn new_tx_gate() -> TxGate {
35 Arc::new(Semaphore::new(1))
36}
37
38const MAX_QUERY_LENGTH: usize = 1024 * 1024;
40
41const MAX_SYNC_PULL_UNITS: u32 = 4096;
43
44const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
46
47#[cfg(not(test))]
51const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
52#[cfg(test)]
53const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
54
55const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
58
59const MAX_AUTH_FAILURES: u32 = 5;
61
62const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
64
65pub fn new_rate_limiter() -> AuthRateLimiter {
67 Arc::new(Mutex::new(HashMap::new()))
68}
69
70fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
73 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
74 let now = Instant::now();
76 map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
77
78 if let Some((count, _)) = map.get(&ip) {
79 *count >= MAX_AUTH_FAILURES
80 } else {
81 false
82 }
83}
84
85fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
87 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
88 let now = Instant::now();
89 let entry = map.entry(ip).or_insert((0, now));
90 if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
92 *entry = (1, now);
93 } else {
94 entry.0 += 1;
95 }
96}
97
98fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
100 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
101 map.remove(&ip);
102}
103
104fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
107 use sha2::{Digest, Sha256};
108 let ha = Sha256::digest(a);
109 let hb = Sha256::digest(b);
110 let mut diff = 0u8;
111 for (x, y) in ha.iter().zip(hb.iter()) {
112 diff |= x ^ y;
113 }
114 diff == 0
115}
116
117#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct Principal {
122 pub name: String,
123 pub role: String,
124}
125
126fn is_ddl_statement(stmt: &powdb_query::ast::Statement) -> bool {
132 use powdb_query::ast::Statement;
133 let inner = match stmt {
134 Statement::Explain(inner) => inner.as_ref(),
135 other => other,
136 };
137 matches!(
138 inner,
139 Statement::CreateType(_)
140 | Statement::AlterTable(_)
141 | Statement::DropTable(_)
142 | Statement::CreateView(_)
143 | Statement::DropView(_)
144 )
145}
146
147fn required_permission(stmt: &powdb_query::ast::Statement) -> Permission {
153 if is_read_only_statement(stmt) {
154 Permission::Read
155 } else if is_ddl_statement(stmt) {
156 Permission::Ddl
157 } else {
158 Permission::Write
159 }
160}
161
162fn check_statement_permitted(
174 principal: Option<&Principal>,
175 stmt: &powdb_query::ast::Statement,
176) -> Result<(), QueryError> {
177 let Some(p) = principal else {
178 return Ok(());
181 };
182 if is_read_only_statement(stmt) {
185 return Ok(());
186 }
187 let needed = required_permission(stmt);
188 if Role::builtin(&p.role).is_some_and(|r| r.allows(needed)) {
189 return Ok(());
190 }
191 let kind = if needed == Permission::Ddl {
192 "schema-definition"
193 } else {
194 "write"
195 };
196 Err(QueryError::Execution(format!(
197 "permission denied: role '{}' cannot execute {kind} statements",
198 p.role
199 )))
200}
201
202#[derive(Debug, Clone, PartialEq, Eq)]
204pub enum AuthOutcome {
205 Authenticated { principal: Option<Principal> },
209 Rejected,
212}
213
214pub fn authenticate_connect(
226 users: &UserStore,
227 expected_password: Option<&str>,
228 username: Option<&str>,
229 password: Option<&str>,
230) -> AuthOutcome {
231 if !users.is_empty() {
232 let Some(name) = username else {
234 return AuthOutcome::Rejected;
235 };
236 let Some(candidate) = password else {
237 return AuthOutcome::Rejected;
238 };
239 match users.authenticate(name, candidate) {
240 Some(user) => AuthOutcome::Authenticated {
241 principal: Some(Principal {
242 name: user.name.clone(),
243 role: user.role.clone(),
244 }),
245 },
246 None => AuthOutcome::Rejected,
247 }
248 } else {
249 match expected_password {
251 Some(expected) => {
252 if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
253 AuthOutcome::Authenticated { principal: None }
254 } else {
255 AuthOutcome::Rejected
256 }
257 }
258 None => AuthOutcome::Authenticated { principal: None },
259 }
260 }
261}
262
263const DEFAULT_DB_NAME: &str = "default";
267
268fn check_db_name(configured: Option<&str>, requested: &str) -> Result<(), String> {
278 if requested.is_empty() || requested == DEFAULT_DB_NAME {
279 return Ok(());
280 }
281 match configured {
282 None => Ok(()),
283 Some(name) if requested == name => Ok(()),
284 Some(name) => Err(format!(
285 "unknown database '{requested}'; this server serves '{name}'"
286 )),
287 }
288}
289
290const SAFE_ERROR_PREFIXES: &[&str] = &[
292 "table not found",
293 "table '",
297 "type '",
298 "column not found",
299 "at position",
302 "parse error",
303 "type mismatch",
304 "unknown table",
305 "unknown column",
306 "unknown function",
307 "syntax error",
308 "expected",
309 "unexpected",
310 "missing",
311 "duplicate",
312 "invalid",
313 "cannot",
314 "no such",
315 "already exists",
316 "permission denied",
317 "row too large",
318 "unique constraint violation",
319 "sort input exceeds",
324 "join result exceeds",
325 "query exceeded memory budget",
326 "result too large",
327 "wal durability sync failed",
332];
333
334fn sanitize_error(e: &str) -> String {
338 let lower = e.to_lowercase();
339 for prefix in SAFE_ERROR_PREFIXES {
340 if lower.starts_with(prefix) {
341 return e.to_string();
342 }
343 }
344 "query execution error".into()
345}
346
347async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
350 let write_fut = async {
351 if msg.write_to(writer).await.is_err() {
352 return false;
353 }
354 writer.flush().await.is_ok()
355 };
356 tokio::time::timeout(WRITE_TIMEOUT, write_fut)
357 .await
358 .unwrap_or_default()
359}
360
361pub struct ConnOpts<'a> {
364 pub engine: Arc<RwLock<Engine>>,
365 pub tx_gate: TxGate,
366 pub expected_password: Option<Zeroizing<String>>,
369 pub users: Arc<UserStore>,
373 pub shutdown_rx: &'a mut watch::Receiver<bool>,
374 pub idle_timeout: Duration,
375 pub query_timeout: Duration,
376 pub rate_limiter: Option<&'a AuthRateLimiter>,
377 pub peer_addr: Option<std::net::SocketAddr>,
378 pub metrics: Arc<Metrics>,
380 pub tx_wait_timeout: Duration,
384 pub db_name: Option<String>,
388}
389
390type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
403
404fn execute_write_deferred(
411 engine: &Arc<RwLock<Engine>>,
412 f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
413) -> DispatchOutcome {
414 let mut eng = match engine.write() {
415 Ok(eng) => eng,
416 Err(e) => {
417 return (
418 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
419 None,
420 )
421 }
422 };
423 eng.run_with_deferred_durability(f)
424}
425
426fn dispatch_query(
427 engine: &Arc<RwLock<Engine>>,
428 query: &str,
429 principal: Option<&Principal>,
430) -> DispatchOutcome {
431 let stmt_result = parser::parse(query).map_err(|e| e.to_string());
432
433 if let Ok(stmt) = &stmt_result {
437 if let Err(e) = check_statement_permitted(principal, stmt) {
438 return (Err(e), None);
439 }
440 }
441
442 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
443 if can_try_read {
444 let res = {
445 let eng = match engine.read() {
446 Ok(eng) => eng,
447 Err(e) => {
448 return (
449 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
450 None,
451 )
452 }
453 };
454 eng.execute_powql_readonly(query)
455 };
456 match res {
457 Ok(r) => return (Ok(r), None),
458 Err(QueryError::ReadonlyNeedsWrite) => {
459 }
461 Err(e) => return (Err(e), None),
462 }
463 }
464
465 if matches!(
466 parsed_transaction_control(&stmt_result),
467 Some(TransactionControl::Rollback)
468 ) {
469 let mut eng = match engine.write() {
470 Ok(eng) => eng,
471 Err(e) => {
472 return (
473 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
474 None,
475 )
476 }
477 };
478 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
479 }
480 execute_write_deferred(engine, |eng| eng.execute_powql(query))
481}
482
483fn dispatch_sql_query(
484 engine: &Arc<RwLock<Engine>>,
485 query: &str,
486 principal: Option<&Principal>,
487) -> DispatchOutcome {
488 let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
489
490 if let Ok(stmt) = &stmt_result {
491 if let Err(e) = check_statement_permitted(principal, stmt) {
492 return (Err(e), None);
493 }
494 }
495
496 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
497 if can_try_read {
498 let res = {
499 let eng = match engine.read() {
500 Ok(eng) => eng,
501 Err(e) => {
502 return (
503 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
504 None,
505 )
506 }
507 };
508 eng.execute_sql_readonly(query)
509 };
510 match res {
511 Ok(r) => return (Ok(r), None),
512 Err(QueryError::ReadonlyNeedsWrite) => {}
513 Err(e) => return (Err(e), None),
514 }
515 }
516
517 if matches!(
518 parsed_transaction_control(&stmt_result),
519 Some(TransactionControl::Rollback)
520 ) {
521 let mut eng = match engine.write() {
522 Ok(eng) => eng,
523 Err(e) => {
524 return (
525 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
526 None,
527 )
528 }
529 };
530 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
531 }
532 execute_write_deferred(engine, |eng| eng.execute_sql(query))
533}
534
535fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
538 use powdb_query::ast::ParamValue;
539 match p {
540 WireParam::Null => ParamValue::Null,
541 WireParam::Int(v) => ParamValue::Int(*v),
542 WireParam::Float(v) => ParamValue::Float(*v),
543 WireParam::Bool(v) => ParamValue::Bool(*v),
544 WireParam::Str(s) => ParamValue::Str(s.clone()),
545 }
546}
547
548#[derive(Debug, Clone, Copy, PartialEq, Eq)]
554enum TransactionControl {
555 Begin,
556 Commit,
557 Rollback,
558}
559
560fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
561 use powdb_query::ast::Statement;
562 match stmt {
563 Statement::Begin => Some(TransactionControl::Begin),
564 Statement::Commit => Some(TransactionControl::Commit),
565 Statement::Rollback => Some(TransactionControl::Rollback),
566 _ => None,
567 }
568}
569
570fn classify_query_transaction_control(query: &str) -> Option<TransactionControl> {
571 parser::parse(query)
572 .ok()
573 .and_then(|stmt| transaction_control(&stmt))
574}
575
576fn classify_sql_transaction_control(query: &str) -> Option<TransactionControl> {
577 sql::parse_sql(query)
578 .ok()
579 .and_then(|stmt| transaction_control(&stmt))
580}
581
582fn classify_params_transaction_control(
583 query: &str,
584 params: &[WireParam],
585) -> Option<TransactionControl> {
586 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
587 parser::parse_with_params(query, &bound)
588 .ok()
589 .and_then(|stmt| transaction_control(&stmt))
590}
591
592fn parsed_transaction_control(
593 stmt_result: &Result<powdb_query::ast::Statement, String>,
594) -> Option<TransactionControl> {
595 stmt_result.as_ref().ok().and_then(transaction_control)
596}
597
598fn execute_rollback_preserving_sync_if_needed(
599 engine: &mut Engine,
600) -> Result<QueryResult, QueryError> {
601 engine.rollback_transaction_preserving_wal_archive()
602}
603
604fn dispatch_query_with_params(
605 engine: &Arc<RwLock<Engine>>,
606 query: &str,
607 params: &[WireParam],
608 principal: Option<&Principal>,
609) -> DispatchOutcome {
610 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
611
612 let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
615
616 if let Ok(stmt) = &stmt_result {
617 if let Err(e) = check_statement_permitted(principal, stmt) {
618 return (Err(e), None);
619 }
620 }
621
622 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
623 if can_try_read {
624 let res = {
625 let eng = match engine.read() {
626 Ok(eng) => eng,
627 Err(e) => {
628 return (
629 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
630 None,
631 )
632 }
633 };
634 eng.execute_powql_readonly_with_params(query, &bound)
635 };
636 match res {
637 Ok(r) => return (Ok(r), None),
638 Err(QueryError::ReadonlyNeedsWrite) => {
639 }
641 Err(e) => return (Err(e), None),
642 }
643 }
644
645 if matches!(
646 parsed_transaction_control(&stmt_result),
647 Some(TransactionControl::Rollback)
648 ) {
649 let mut eng = match engine.write() {
650 Ok(eng) => eng,
651 Err(e) => {
652 return (
653 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
654 None,
655 )
656 }
657 };
658 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
659 }
660 execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, &bound))
661}
662
663#[derive(Debug, Clone)]
664struct SyncPullRequest {
665 replica_id: String,
666 since_lsn: u64,
667 max_units: u32,
668 max_bytes: u64,
669 database_id: [u8; 16],
670 primary_generation: u64,
671 wal_format_version: u16,
672 catalog_version: u16,
673 segment_format_version: u16,
674}
675
676#[derive(Debug, Clone, Copy, PartialEq, Eq)]
677enum SyncErrorClass {
678 AuthRequired,
679 PermissionDenied,
680 InvalidReplicaId,
681 ActiveTransaction,
682 QueryExecution,
683 InvalidMaxUnits,
684 InvalidMaxBytes,
685 SyncContext,
686 StatusRead,
687 CursorLsnMismatch,
688 IdentityRead,
689 IdentityOrFormatMismatch,
690 RetainedRead,
691 RetainedUnitEncoding,
692 RetainedChunkNotApplyable,
693 LsnAheadOfRemote,
694 AckValidation,
695 AckUpdate,
696 Internal,
697}
698
699impl SyncErrorClass {
700 const fn as_label(self) -> &'static str {
701 match self {
702 Self::AuthRequired => "auth_required",
703 Self::PermissionDenied => "permission_denied",
704 Self::InvalidReplicaId => "invalid_replica_id",
705 Self::ActiveTransaction => "active_transaction",
706 Self::QueryExecution => "query_execution",
707 Self::InvalidMaxUnits => "invalid_max_units",
708 Self::InvalidMaxBytes => "invalid_max_bytes",
709 Self::SyncContext => "sync_context",
710 Self::StatusRead => "status_read",
711 Self::CursorLsnMismatch => "cursor_lsn_mismatch",
712 Self::IdentityRead => "identity_read",
713 Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
714 Self::RetainedRead => "retained_read",
715 Self::RetainedUnitEncoding => "retained_unit_encoding",
716 Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
717 Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
718 Self::AckValidation => "ack_validation",
719 Self::AckUpdate => "ack_update",
720 Self::Internal => "internal",
721 }
722 }
723}
724
725#[derive(Debug, Clone)]
726struct SyncDecision {
727 message: Message,
728 error_class: Option<SyncErrorClass>,
729}
730
731impl SyncDecision {
732 fn ok(message: Message) -> Self {
733 Self {
734 message,
735 error_class: None,
736 }
737 }
738
739 fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
740 Self {
741 message: Message::Error {
742 message: message.into(),
743 },
744 error_class: Some(class),
745 }
746 }
747}
748
749fn check_sync_protocol_permitted(
750 credential_authenticated: bool,
751 principal: Option<&Principal>,
752) -> Result<(), (SyncErrorClass, String)> {
753 if !credential_authenticated {
754 return Err((
755 SyncErrorClass::AuthRequired,
756 "sync protocol requires authentication".to_string(),
757 ));
758 }
759 if let Some(principal) = principal {
760 let allowed =
761 Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
762 if !allowed {
763 return Err((
764 SyncErrorClass::PermissionDenied,
765 format!(
766 "permission denied: role '{}' cannot use sync protocol",
767 principal.role
768 ),
769 ));
770 }
771 }
772 Ok(())
773}
774
775fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
776 if replica_id.is_empty() {
777 return Err("replica id must be non-empty".to_string());
778 }
779 if replica_id.len() > 128 {
780 return Err("replica id must be at most 128 bytes".to_string());
781 }
782 if !replica_id
783 .bytes()
784 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
785 {
786 return Err("replica id contains unsupported characters".to_string());
787 }
788 Ok(())
789}
790
791fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<(PathBuf, u64), String> {
792 let engine = engine
793 .read()
794 .map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
795 let catalog = engine.catalog();
796 Ok((catalog.data_dir().to_path_buf(), catalog.max_lsn()))
797}
798
799fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
800 match action {
801 SyncRepairAction::None => WireSyncRepairAction::None,
802 SyncRepairAction::Pull => WireSyncRepairAction::Pull,
803 SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
804 SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
805 }
806}
807
808fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
809 WireSyncStatus {
810 replica_id: status.replica_id,
811 active: status.active,
812 last_applied_lsn: status.last_applied_lsn,
813 remote_lsn: status.remote_lsn,
814 servable_lsn: status.servable_lsn,
815 unarchived_lsn: status.unarchived_lsn,
816 lag_lsn: status.lag_lsn,
817 lag_bytes: status.lag_bytes,
818 lag_ms: status.lag_ms,
819 stale: status.stale,
820 repair_action: wire_repair_action(status.repair_action),
821 last_sync_error: status.last_sync_error,
822 }
823}
824
825fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
826 WireRetainedUnit {
827 tx_id: unit.tx_id,
828 record_type: unit.record_type,
829 lsn: unit.lsn,
830 data: unit.data,
831 }
832}
833
834fn sync_operation_outcome(message: &Message) -> SyncOutcome {
835 match message {
836 Message::SyncStatusResult { .. }
837 | Message::SyncPullResult { .. }
838 | Message::SyncAckResult { .. } => SyncOutcome::Ok,
839 _ => SyncOutcome::Error,
840 }
841}
842
843fn sync_operation_label(operation: SyncOperation) -> &'static str {
844 match operation {
845 SyncOperation::Status => "status",
846 SyncOperation::Pull => "pull",
847 SyncOperation::Ack => "ack",
848 }
849}
850
851fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
852 units.iter().fold(0, |total, unit| {
853 total.saturating_add(unit.encoded_len().unwrap_or(0))
854 })
855}
856
857fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
858 match action {
859 WireSyncRepairAction::None => SyncRepairLabel::None,
860 WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
861 WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
862 WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
863 }
864}
865
866fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
867 match action {
868 WireSyncRepairAction::None => "none",
869 WireSyncRepairAction::Pull => "pull",
870 WireSyncRepairAction::AwaitArchive => "await_archive",
871 WireSyncRepairAction::Rebootstrap => "rebootstrap",
872 }
873}
874
875const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
876const FNV1A64_PRIME: u64 = 0x100000001b3;
877const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
878
879fn replica_fingerprint(replica_id: &str) -> String {
880 let mut hash = FNV1A64_OFFSET;
881 for byte in replica_id.as_bytes() {
882 hash ^= u64::from(*byte);
883 hash = hash.wrapping_mul(FNV1A64_PRIME);
884 }
885 format!("{hash:016x}")
886}
887
888fn log_replica_fingerprint(replica_id: &str) -> String {
889 if validate_wire_replica_id(replica_id).is_ok() {
890 replica_fingerprint(replica_id)
891 } else {
892 INVALID_REPLICA_FINGERPRINT.to_string()
893 }
894}
895
896#[derive(Debug, Clone)]
897struct SyncLogContext {
898 replica_fingerprint: String,
899 since_lsn: Option<u64>,
900 applied_lsn: Option<u64>,
901 observed_remote_lsn: Option<u64>,
902 max_units: Option<u32>,
903 max_bytes: Option<u64>,
904}
905
906impl SyncLogContext {
907 fn status(replica_id: &str) -> Self {
908 Self::base(replica_id)
909 }
910
911 fn pull(request: &SyncPullRequest) -> Self {
912 Self {
913 replica_fingerprint: log_replica_fingerprint(&request.replica_id),
914 since_lsn: Some(request.since_lsn),
915 applied_lsn: None,
916 observed_remote_lsn: None,
917 max_units: Some(request.max_units),
918 max_bytes: Some(request.max_bytes),
919 }
920 }
921
922 fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
923 Self {
924 replica_fingerprint: log_replica_fingerprint(replica_id),
925 since_lsn: None,
926 applied_lsn: Some(applied_lsn),
927 observed_remote_lsn: Some(observed_remote_lsn),
928 max_units: None,
929 max_bytes: None,
930 }
931 }
932
933 fn base(replica_id: &str) -> Self {
934 Self {
935 replica_fingerprint: log_replica_fingerprint(replica_id),
936 since_lsn: None,
937 applied_lsn: None,
938 observed_remote_lsn: None,
939 max_units: None,
940 max_bytes: None,
941 }
942 }
943}
944
945struct SyncExecutionContext<'a> {
946 tx_gate: TxGate,
947 connection_has_transaction: bool,
948 operation: SyncOperation,
949 log_context: SyncLogContext,
950 metrics: &'a Arc<Metrics>,
951 query_timeout: Duration,
952}
953
954fn log_sync_decision(
955 operation: SyncOperation,
956 context: &SyncLogContext,
957 elapsed: Duration,
958 decision: &SyncDecision,
959) {
960 let operation = sync_operation_label(operation);
961 let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
962 match &decision.message {
963 Message::SyncStatusResult { status } => {
964 let repair_action = wire_repair_action_label(status.repair_action);
965 if status.stale || status.repair_action != WireSyncRepairAction::None {
966 info!(
967 operation = operation,
968 replica_fingerprint = %context.replica_fingerprint,
969 remote_lsn = status.remote_lsn,
970 last_applied_lsn = ?status.last_applied_lsn,
971 servable_lsn = ?status.servable_lsn,
972 unarchived_lsn = ?status.unarchived_lsn,
973 lag_lsn = ?status.lag_lsn,
974 lag_bytes = ?status.lag_bytes,
975 lag_ms = ?status.lag_ms,
976 stale = status.stale,
977 repair_action,
978 elapsed_ms,
979 "sync decision"
980 );
981 } else {
982 debug!(
983 operation = operation,
984 replica_fingerprint = %context.replica_fingerprint,
985 remote_lsn = status.remote_lsn,
986 last_applied_lsn = ?status.last_applied_lsn,
987 repair_action,
988 elapsed_ms,
989 "sync decision"
990 );
991 }
992 }
993 Message::SyncPullResult {
994 status,
995 units,
996 has_more,
997 } => {
998 let repair_action = wire_repair_action_label(status.repair_action);
999 let units_len = units.len();
1000 let payload_bytes = sync_pull_payload_bytes(units);
1001 if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
1002 info!(
1003 operation = operation,
1004 replica_fingerprint = %context.replica_fingerprint,
1005 since_lsn = ?context.since_lsn,
1006 max_units = ?context.max_units,
1007 max_bytes = ?context.max_bytes,
1008 units = units_len,
1009 payload_bytes,
1010 has_more = *has_more,
1011 remote_lsn = status.remote_lsn,
1012 last_applied_lsn = ?status.last_applied_lsn,
1013 servable_lsn = ?status.servable_lsn,
1014 unarchived_lsn = ?status.unarchived_lsn,
1015 lag_lsn = ?status.lag_lsn,
1016 stale = status.stale,
1017 repair_action,
1018 elapsed_ms,
1019 "sync decision"
1020 );
1021 } else {
1022 debug!(
1023 operation = operation,
1024 replica_fingerprint = %context.replica_fingerprint,
1025 since_lsn = ?context.since_lsn,
1026 units = units_len,
1027 payload_bytes,
1028 has_more = *has_more,
1029 remote_lsn = status.remote_lsn,
1030 repair_action,
1031 elapsed_ms,
1032 "sync decision"
1033 );
1034 }
1035 }
1036 Message::SyncAckResult {
1037 previous_applied_lsn,
1038 applied_lsn,
1039 remote_lsn,
1040 advanced,
1041 status,
1042 } => {
1043 let repair_action = wire_repair_action_label(status.repair_action);
1044 if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
1045 info!(
1046 operation = operation,
1047 replica_fingerprint = %context.replica_fingerprint,
1048 requested_applied_lsn = ?context.applied_lsn,
1049 observed_remote_lsn = ?context.observed_remote_lsn,
1050 previous_applied_lsn = *previous_applied_lsn,
1051 applied_lsn = *applied_lsn,
1052 remote_lsn = *remote_lsn,
1053 advanced = *advanced,
1054 stale = status.stale,
1055 repair_action,
1056 elapsed_ms,
1057 "sync decision"
1058 );
1059 } else {
1060 debug!(
1061 operation = operation,
1062 replica_fingerprint = %context.replica_fingerprint,
1063 requested_applied_lsn = ?context.applied_lsn,
1064 previous_applied_lsn = *previous_applied_lsn,
1065 applied_lsn = *applied_lsn,
1066 remote_lsn = *remote_lsn,
1067 advanced = *advanced,
1068 repair_action,
1069 elapsed_ms,
1070 "sync decision"
1071 );
1072 }
1073 }
1074 Message::Error { .. } => {
1075 warn!(
1076 operation = operation,
1077 replica_fingerprint = %context.replica_fingerprint,
1078 since_lsn = ?context.since_lsn,
1079 applied_lsn = ?context.applied_lsn,
1080 observed_remote_lsn = ?context.observed_remote_lsn,
1081 max_units = ?context.max_units,
1082 max_bytes = ?context.max_bytes,
1083 error_class = decision
1084 .error_class
1085 .unwrap_or(SyncErrorClass::Internal)
1086 .as_label(),
1087 elapsed_ms,
1088 "sync decision rejected"
1089 );
1090 }
1091 _ => {
1092 debug!(
1093 operation = operation,
1094 replica_fingerprint = %context.replica_fingerprint,
1095 elapsed_ms,
1096 "unexpected sync decision response"
1097 );
1098 }
1099 }
1100}
1101
1102fn trim_to_applyable_v1_prefix(
1103 raw_units: &mut Vec<RetainedUnit>,
1104 wire_units: &mut Vec<WireRetainedUnit>,
1105) -> Result<(), String> {
1106 let mut last_error = None;
1107 while !raw_units.is_empty() {
1108 match validate_v1_retained_units_applyable(raw_units) {
1109 Ok(()) => return Ok(()),
1110 Err(err) => {
1111 last_error = Some(err.to_string());
1112 raw_units.pop();
1113 wire_units.pop();
1114 }
1115 }
1116 }
1117 if let Some(error) = last_error {
1118 return Err(format!(
1119 "sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
1120 ));
1121 }
1122 Ok(())
1123}
1124
1125fn validate_sync_ack_applyable_boundary(
1126 data_dir: &Path,
1127 replica_id: &str,
1128 applied_lsn: u64,
1129 remote_lsn: u64,
1130) -> Result<(), String> {
1131 let status =
1132 replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
1133 let Some(previous_lsn) = status.last_applied_lsn else {
1134 return Ok(());
1135 };
1136 if applied_lsn <= previous_lsn {
1137 return Ok(());
1138 }
1139 let range_len = applied_lsn - previous_lsn;
1140 if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
1141 return Err(format!(
1142 "sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
1143 ));
1144 }
1145 let max_units =
1146 usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
1147 let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
1148 let segment_dir = retained_segments_dir(data_dir);
1149 let units = read_units_through(
1150 &segment_dir,
1151 identity.segment_identity(),
1152 previous_lsn,
1153 applied_lsn,
1154 max_units,
1155 )
1156 .map_err(|err| err.to_string())?;
1157 if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
1158 return Err(
1159 "sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
1160 );
1161 }
1162 validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
1163}
1164
1165#[cfg(test)]
1166fn dispatch_sync_status(
1167 engine: &Arc<RwLock<Engine>>,
1168 replica_id: String,
1169 credential_authenticated: bool,
1170 principal: Option<&Principal>,
1171) -> Message {
1172 dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
1173}
1174
1175fn dispatch_sync_status_decision(
1176 engine: &Arc<RwLock<Engine>>,
1177 replica_id: String,
1178 credential_authenticated: bool,
1179 principal: Option<&Principal>,
1180) -> SyncDecision {
1181 if let Err((class, message)) =
1182 check_sync_protocol_permitted(credential_authenticated, principal)
1183 {
1184 return SyncDecision::error(class, message);
1185 }
1186 if let Err(message) = validate_wire_replica_id(&replica_id) {
1187 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1188 }
1189 let (data_dir, remote_lsn) = match sync_context(engine) {
1190 Ok(context) => context,
1191 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1192 };
1193 match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
1194 Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
1195 status: wire_sync_status(status),
1196 }),
1197 Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
1198 }
1199}
1200
1201#[cfg(test)]
1202fn dispatch_sync_pull(
1203 engine: &Arc<RwLock<Engine>>,
1204 request: SyncPullRequest,
1205 credential_authenticated: bool,
1206 principal: Option<&Principal>,
1207) -> Message {
1208 dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
1209}
1210
1211fn dispatch_sync_pull_decision(
1212 engine: &Arc<RwLock<Engine>>,
1213 request: SyncPullRequest,
1214 credential_authenticated: bool,
1215 principal: Option<&Principal>,
1216) -> SyncDecision {
1217 if let Err((class, message)) =
1218 check_sync_protocol_permitted(credential_authenticated, principal)
1219 {
1220 return SyncDecision::error(class, message);
1221 }
1222 if let Err(message) = validate_wire_replica_id(&request.replica_id) {
1223 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1224 }
1225 if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
1226 return SyncDecision::error(
1227 SyncErrorClass::InvalidMaxUnits,
1228 format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
1229 );
1230 }
1231 if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
1232 return SyncDecision::error(
1233 SyncErrorClass::InvalidMaxBytes,
1234 format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
1235 );
1236 }
1237
1238 let (data_dir, remote_lsn) = match sync_context(engine) {
1239 Ok(context) => context,
1240 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1241 };
1242 let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
1243 Ok(status) => status,
1244 Err(err) => {
1245 return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
1246 }
1247 };
1248 let Some(cursor_lsn) = status.last_applied_lsn else {
1249 return SyncDecision::ok(Message::SyncPullResult {
1250 status: wire_sync_status(status),
1251 units: Vec::new(),
1252 has_more: false,
1253 });
1254 };
1255 if status.repair_action != SyncRepairAction::Pull {
1256 return SyncDecision::ok(Message::SyncPullResult {
1257 status: wire_sync_status(status),
1258 units: Vec::new(),
1259 has_more: false,
1260 });
1261 }
1262 if request.since_lsn != cursor_lsn {
1263 return SyncDecision::error(
1264 SyncErrorClass::CursorLsnMismatch,
1265 format!(
1266 "sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
1267 request.since_lsn
1268 ),
1269 );
1270 }
1271
1272 let identity = match read_identity(&data_dir) {
1273 Ok(identity) => identity,
1274 Err(err) => {
1275 return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
1276 }
1277 };
1278 let expected = identity.segment_identity();
1279 if request.database_id != expected.database_id
1280 || request.primary_generation != expected.primary_generation
1281 || request.wal_format_version != expected.wal_format_version
1282 || request.catalog_version != expected.catalog_version
1283 || request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
1284 {
1285 return SyncDecision::error(
1286 SyncErrorClass::IdentityOrFormatMismatch,
1287 "sync pull identity or format version mismatch; rebootstrap required",
1288 );
1289 }
1290
1291 let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
1292 let requested_through_lsn = request
1293 .since_lsn
1294 .saturating_add(request.max_units as u64)
1295 .min(remote_lsn)
1296 .min(status.servable_lsn.unwrap_or(request.since_lsn));
1297 let segment_dir = retained_segments_dir(&data_dir);
1298 if requested_through_lsn > request.since_lsn {
1299 if let Err(err) = validate_retained_tail_available(
1300 &segment_dir,
1301 expected,
1302 request.since_lsn,
1303 requested_through_lsn,
1304 ) {
1305 let mut rebootstrap_status = status;
1306 rebootstrap_status.stale = true;
1307 rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
1308 rebootstrap_status.last_sync_error = Some(format!(
1309 "retained history is unavailable; rebootstrap required: {err}"
1310 ));
1311 return SyncDecision::ok(Message::SyncPullResult {
1312 status: wire_sync_status(rebootstrap_status),
1313 units: Vec::new(),
1314 has_more: false,
1315 });
1316 }
1317 }
1318
1319 let raw_units = match read_units_through(
1320 &segment_dir,
1321 expected,
1322 request.since_lsn,
1323 requested_through_lsn,
1324 effective_max_units,
1325 ) {
1326 Ok(units) => units,
1327 Err(err) => {
1328 return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
1329 }
1330 };
1331
1332 let mut selected_raw = Vec::new();
1333 let mut selected = Vec::new();
1334 let mut selected_bytes = 0u64;
1335 for unit in raw_units {
1336 let wire_unit = wire_retained_unit(unit.clone());
1337 let unit_bytes = match wire_unit.encoded_len() {
1338 Ok(bytes) => bytes,
1339 Err(message) => {
1340 return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
1341 }
1342 };
1343 if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
1344 if selected.is_empty() {
1345 return SyncDecision::error(
1346 SyncErrorClass::InvalidMaxBytes,
1347 "sync pull maxBytes is too small for the next retained unit",
1348 );
1349 }
1350 break;
1351 }
1352 selected_bytes += unit_bytes;
1353 selected_raw.push(unit);
1354 selected.push(wire_unit);
1355 }
1356 if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
1357 return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
1358 }
1359
1360 let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
1361 let has_more = selected
1362 .last()
1363 .is_some_and(|unit| unit.lsn < fetchable_through_lsn);
1364 SyncDecision::ok(Message::SyncPullResult {
1365 status: wire_sync_status(status),
1366 units: selected,
1367 has_more,
1368 })
1369}
1370
1371#[cfg(test)]
1372fn dispatch_sync_ack(
1373 engine: &Arc<RwLock<Engine>>,
1374 replica_id: String,
1375 applied_lsn: u64,
1376 observed_remote_lsn: u64,
1377 credential_authenticated: bool,
1378 principal: Option<&Principal>,
1379) -> Message {
1380 dispatch_sync_ack_decision(
1381 engine,
1382 replica_id,
1383 applied_lsn,
1384 observed_remote_lsn,
1385 credential_authenticated,
1386 principal,
1387 )
1388 .message
1389}
1390
1391fn dispatch_sync_ack_decision(
1392 engine: &Arc<RwLock<Engine>>,
1393 replica_id: String,
1394 applied_lsn: u64,
1395 observed_remote_lsn: u64,
1396 credential_authenticated: bool,
1397 principal: Option<&Principal>,
1398) -> SyncDecision {
1399 if let Err((class, message)) =
1400 check_sync_protocol_permitted(credential_authenticated, principal)
1401 {
1402 return SyncDecision::error(class, message);
1403 }
1404 if let Err(message) = validate_wire_replica_id(&replica_id) {
1405 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1406 }
1407 if applied_lsn > observed_remote_lsn {
1408 return SyncDecision::error(
1409 SyncErrorClass::LsnAheadOfRemote,
1410 format!(
1411 "sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
1412 ),
1413 );
1414 }
1415 let (data_dir, remote_lsn) = match sync_context(engine) {
1416 Ok(context) => context,
1417 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1418 };
1419 if observed_remote_lsn > remote_lsn {
1420 return SyncDecision::error(
1421 SyncErrorClass::LsnAheadOfRemote,
1422 format!(
1423 "sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
1424 ),
1425 );
1426 }
1427 if let Err(message) =
1428 validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
1429 {
1430 return SyncDecision::error(SyncErrorClass::AckValidation, message);
1431 }
1432 match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
1433 Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
1434 previous_applied_lsn: summary.previous_applied_lsn,
1435 applied_lsn: summary.applied_lsn,
1436 remote_lsn: summary.remote_lsn,
1437 advanced: summary.advanced,
1438 status: wire_sync_status(summary.status),
1439 }),
1440 Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
1441 }
1442}
1443
1444async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
1445where
1446 T: Send + 'static,
1447 F: FnOnce(T) -> SyncDecision + Send + 'static,
1448{
1449 let mut handle = tokio::task::spawn_blocking(move || f(input));
1450 tokio::select! {
1451 result = &mut handle => match result {
1452 Ok(decision) => decision,
1453 Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1454 },
1455 _ = tokio::time::sleep(query_timeout) => match handle.await {
1456 Ok(decision) => decision,
1457 Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1458 },
1459 }
1460}
1461
1462async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
1463where
1464 T: Send + 'static,
1465 F: FnOnce(T) -> SyncDecision + Send + 'static,
1466{
1467 let SyncExecutionContext {
1468 tx_gate,
1469 connection_has_transaction,
1470 operation,
1471 log_context,
1472 metrics,
1473 query_timeout,
1474 } = context;
1475 let start = Instant::now();
1476 if connection_has_transaction {
1477 let decision = SyncDecision::error(
1478 SyncErrorClass::ActiveTransaction,
1479 "sync protocol is unavailable inside an active transaction",
1480 );
1481 let elapsed = start.elapsed();
1482 log_sync_decision(operation, &log_context, elapsed, &decision);
1483 metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1484 return decision.message;
1485 }
1486
1487 let permit = match tx_gate.acquire_owned().await {
1488 Ok(permit) => permit,
1489 Err(_) => {
1490 let decision =
1491 SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
1492 let elapsed = start.elapsed();
1493 log_sync_decision(operation, &log_context, elapsed, &decision);
1494 metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1495 return decision.message;
1496 }
1497 };
1498 let decision = run_blocking_sync(input, query_timeout, f).await;
1499 drop(permit);
1500 match &decision.message {
1501 Message::SyncStatusResult { status } => {
1502 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1503 }
1504 Message::SyncPullResult { status, units, .. } => {
1505 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1506 metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
1507 }
1508 Message::SyncAckResult {
1509 advanced, status, ..
1510 } => {
1511 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1512 if *advanced {
1513 metrics.inc_sync_ack_advanced();
1514 }
1515 }
1516 _ => {}
1517 }
1518 let elapsed = start.elapsed();
1519 log_sync_decision(operation, &log_context, elapsed, &decision);
1520 metrics.record_sync_operation(
1521 operation,
1522 elapsed,
1523 sync_operation_outcome(&decision.message),
1524 );
1525 decision.message
1526}
1527
1528async fn acquire_begin_permit(
1535 tx_gate: &TxGate,
1536 tx_wait_timeout: Duration,
1537 metrics: &Arc<Metrics>,
1538) -> Result<OwnedSemaphorePermit, Message> {
1539 match tokio::time::timeout(tx_wait_timeout, tx_gate.clone().acquire_owned()).await {
1540 Ok(Ok(permit)) => Ok(permit),
1541 Ok(Err(_)) => Err(Message::Error {
1542 message: "query execution error".into(),
1543 }),
1544 Err(_) => {
1545 metrics.inc_tx_gate_timeout();
1546 Err(Message::Error {
1547 message: format!(
1548 "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1549 tx_wait_timeout.as_millis()
1550 ),
1551 })
1552 }
1553 }
1554}
1555
1556async fn acquire_autocommit_permit(
1566 tx_gate: &TxGate,
1567 tx_wait_timeout: Duration,
1568 metrics: &Arc<Metrics>,
1569) -> Result<OwnedSemaphorePermit, Message> {
1570 match tokio::time::timeout(tx_wait_timeout, tx_gate.clone().acquire_owned()).await {
1571 Ok(Ok(permit)) => Ok(permit),
1572 Ok(Err(_)) => Err(Message::Error {
1573 message: "query execution error".into(),
1574 }),
1575 Err(_) => {
1576 metrics.inc_tx_gate_timeout();
1577 Err(Message::Error {
1578 message: format!(
1579 "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1580 tx_wait_timeout.as_millis()
1581 ),
1582 })
1583 }
1584 }
1585}
1586
1587#[allow(clippy::too_many_arguments)]
1593async fn execute_wire_query(
1594 engine: Arc<RwLock<Engine>>,
1595 tx_gate: TxGate,
1596 tx_permit: &mut Option<OwnedSemaphorePermit>,
1597 query: String,
1598 principal: Option<Principal>,
1599 query_timeout: Duration,
1600 tx_wait_timeout: Duration,
1601 metrics: &Arc<Metrics>,
1602) -> (Message, Option<PendingDurability>) {
1603 match classify_query_transaction_control(&query) {
1604 Some(TransactionControl::Begin) => {
1605 if tx_permit.is_some() {
1606 return (
1607 Message::Error {
1608 message: sanitize_error(
1609 "cannot begin: a transaction is already active on this connection",
1610 ),
1611 },
1612 None,
1613 );
1614 }
1615 let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
1616 Ok(permit) => permit,
1617 Err(response) => return (response, None),
1618 };
1619 let (response, ticket) = run_blocking_query(
1620 engine,
1621 query,
1622 principal,
1623 query_timeout,
1624 metrics,
1625 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1626 )
1627 .await;
1628 if is_success_response(&response) {
1629 *tx_permit = Some(permit);
1630 }
1631 (response, ticket)
1632 }
1633 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1634 let (response, ticket) = run_blocking_query(
1635 engine,
1636 query,
1637 principal,
1638 query_timeout,
1639 metrics,
1640 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1641 )
1642 .await;
1643 if is_success_response(&response) {
1644 tx_permit.take();
1649 }
1650 (response, ticket)
1651 }
1652 None if tx_permit.is_some() => {
1653 run_blocking_query(
1654 engine,
1655 query,
1656 principal,
1657 query_timeout,
1658 metrics,
1659 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1660 )
1661 .await
1662 }
1663 None => {
1664 let permit = match acquire_autocommit_permit(&tx_gate, tx_wait_timeout, metrics).await {
1665 Ok(permit) => permit,
1666 Err(response) => return (response, None),
1667 };
1668 let out = run_blocking_query(
1669 engine,
1670 query,
1671 principal,
1672 query_timeout,
1673 metrics,
1674 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1675 )
1676 .await;
1677 drop(permit);
1678 out
1679 }
1680 }
1681}
1682
1683#[allow(clippy::too_many_arguments)]
1684async fn execute_wire_query_sql(
1685 engine: Arc<RwLock<Engine>>,
1686 tx_gate: TxGate,
1687 tx_permit: &mut Option<OwnedSemaphorePermit>,
1688 query: String,
1689 principal: Option<Principal>,
1690 query_timeout: Duration,
1691 tx_wait_timeout: Duration,
1692 metrics: &Arc<Metrics>,
1693) -> (Message, Option<PendingDurability>) {
1694 match classify_sql_transaction_control(&query) {
1695 Some(TransactionControl::Begin) => {
1696 if tx_permit.is_some() {
1697 return (
1698 Message::Error {
1699 message: sanitize_error(
1700 "cannot begin: a transaction is already active on this connection",
1701 ),
1702 },
1703 None,
1704 );
1705 }
1706 let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
1707 Ok(permit) => permit,
1708 Err(response) => return (response, None),
1709 };
1710 let (response, ticket) = run_blocking_query(
1711 engine,
1712 query,
1713 principal,
1714 query_timeout,
1715 metrics,
1716 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1717 )
1718 .await;
1719 if is_success_response(&response) {
1720 *tx_permit = Some(permit);
1721 }
1722 (response, ticket)
1723 }
1724 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1725 let (response, ticket) = run_blocking_query(
1726 engine,
1727 query,
1728 principal,
1729 query_timeout,
1730 metrics,
1731 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1732 )
1733 .await;
1734 if is_success_response(&response) {
1735 tx_permit.take();
1738 }
1739 (response, ticket)
1740 }
1741 None if tx_permit.is_some() => {
1742 run_blocking_query(
1743 engine,
1744 query,
1745 principal,
1746 query_timeout,
1747 metrics,
1748 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1749 )
1750 .await
1751 }
1752 None => {
1753 let permit = match acquire_autocommit_permit(&tx_gate, tx_wait_timeout, metrics).await {
1754 Ok(permit) => permit,
1755 Err(response) => return (response, None),
1756 };
1757 let out = run_blocking_query(
1758 engine,
1759 query,
1760 principal,
1761 query_timeout,
1762 metrics,
1763 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1764 )
1765 .await;
1766 drop(permit);
1767 out
1768 }
1769 }
1770}
1771
1772#[allow(clippy::too_many_arguments)]
1776async fn execute_wire_query_with_params(
1777 engine: Arc<RwLock<Engine>>,
1778 tx_gate: TxGate,
1779 tx_permit: &mut Option<OwnedSemaphorePermit>,
1780 query: String,
1781 params: Vec<WireParam>,
1782 principal: Option<Principal>,
1783 query_timeout: Duration,
1784 tx_wait_timeout: Duration,
1785 metrics: &Arc<Metrics>,
1786) -> (Message, Option<PendingDurability>) {
1787 match classify_params_transaction_control(&query, ¶ms) {
1788 Some(TransactionControl::Begin) => {
1789 if tx_permit.is_some() {
1790 return (
1791 Message::Error {
1792 message: sanitize_error(
1793 "cannot begin: a transaction is already active on this connection",
1794 ),
1795 },
1796 None,
1797 );
1798 }
1799 let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
1800 Ok(permit) => permit,
1801 Err(response) => return (response, None),
1802 };
1803 let (response, ticket) = run_blocking_query(
1804 engine,
1805 (query, params),
1806 principal,
1807 query_timeout,
1808 metrics,
1809 |engine, (query, params), principal| {
1810 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1811 },
1812 )
1813 .await;
1814 if is_success_response(&response) {
1815 *tx_permit = Some(permit);
1816 }
1817 (response, ticket)
1818 }
1819 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1820 let (response, ticket) = run_blocking_query(
1821 engine,
1822 (query, params),
1823 principal,
1824 query_timeout,
1825 metrics,
1826 |engine, (query, params), principal| {
1827 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1828 },
1829 )
1830 .await;
1831 if is_success_response(&response) {
1832 tx_permit.take();
1835 }
1836 (response, ticket)
1837 }
1838 None if tx_permit.is_some() => {
1839 run_blocking_query(
1840 engine,
1841 (query, params),
1842 principal,
1843 query_timeout,
1844 metrics,
1845 |engine, (query, params), principal| {
1846 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1847 },
1848 )
1849 .await
1850 }
1851 None => {
1852 let permit = match acquire_autocommit_permit(&tx_gate, tx_wait_timeout, metrics).await {
1853 Ok(permit) => permit,
1854 Err(response) => return (response, None),
1855 };
1856 let out = run_blocking_query(
1857 engine,
1858 (query, params),
1859 principal,
1860 query_timeout,
1861 metrics,
1862 |engine, (query, params), principal| {
1863 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1864 },
1865 )
1866 .await;
1867 drop(permit);
1868 out
1869 }
1870 }
1871}
1872
1873struct DeferredQueryMetric {
1878 start: Instant,
1879 outcome: QueryOutcome,
1880 exceeded_timeout: bool,
1881}
1882
1883type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
1885
1886async fn run_blocking_query<T, F>(
1887 engine: Arc<RwLock<Engine>>,
1888 input: T,
1889 principal: Option<Principal>,
1890 query_timeout: Duration,
1891 metrics: &Arc<Metrics>,
1892 f: F,
1893) -> (Message, Option<PendingDurability>)
1894where
1895 T: Send + 'static,
1896 F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
1897{
1898 let _in_flight = metrics.in_flight_guard();
1899 let start = Instant::now();
1900 let mut handle = tokio::task::spawn_blocking(move || f(engine, input, principal));
1901 let mut exceeded_timeout = false;
1902 let join_result = tokio::select! {
1903 result = &mut handle => result,
1904 _ = tokio::time::sleep(query_timeout) => {
1905 exceeded_timeout = true;
1906 handle.await
1911 }
1912 };
1913
1914 let (message, ticket, outcome) = match join_result {
1915 Ok((Ok(result), ticket)) => match query_result_to_message(result) {
1916 Ok(message) => (message, ticket, QueryOutcome::Ok),
1917 Err(e) => (
1918 Message::Error {
1919 message: sanitize_error(&e.to_string()),
1920 },
1921 ticket,
1922 QueryOutcome::Error,
1923 ),
1924 },
1925 Ok((Err(e), ticket)) => {
1926 let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
1927 QueryOutcome::MemoryLimit
1928 } else {
1929 QueryOutcome::Error
1930 };
1931 (
1932 Message::Error {
1933 message: sanitize_error(&e.to_string()),
1934 },
1935 ticket,
1936 outcome,
1937 )
1938 }
1939 Err(e) => (
1940 Message::Error {
1941 message: format!("internal error: {e}"),
1942 },
1943 None,
1944 QueryOutcome::Error,
1945 ),
1946 };
1947 match ticket {
1948 Some(ticket) => (
1953 message,
1954 Some((
1955 ticket,
1956 DeferredQueryMetric {
1957 start,
1958 outcome,
1959 exceeded_timeout,
1960 },
1961 )),
1962 ),
1963 None => {
1964 if exceeded_timeout {
1965 metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
1966 } else {
1967 metrics.record_query(start.elapsed(), outcome);
1968 }
1969 (message, None)
1970 }
1971 }
1972}
1973
1974async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
1982 match tokio::task::spawn_blocking(move || ticket.wait()).await {
1983 Ok(Ok(())) => None,
1984 Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
1985 Err(e) => Some(format!("internal error: {e}")),
1986 }
1987}
1988
1989fn is_success_response(msg: &Message) -> bool {
1990 matches!(
1991 msg,
1992 Message::ResultRows { .. }
1993 | Message::ResultScalar { .. }
1994 | Message::ResultOk { .. }
1995 | Message::ResultMessage { .. }
1996 )
1997}
1998
1999fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
2000 let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref());
2001 let _ = res;
2002 if let Some(ticket) = ticket {
2005 let _ = ticket.wait();
2006 }
2007}
2008
2009pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
2010where
2011 S: AsyncRead + AsyncWrite + Unpin,
2012{
2013 let ConnOpts {
2014 engine,
2015 tx_gate,
2016 expected_password,
2017 users,
2018 shutdown_rx,
2019 idle_timeout,
2020 query_timeout,
2021 rate_limiter,
2022 peer_addr,
2023 metrics,
2024 tx_wait_timeout,
2025 db_name: server_db_name,
2026 } = opts;
2027
2028 let peer = peer_addr
2029 .map(|a| a.to_string())
2030 .unwrap_or_else(|| "unknown".into());
2031 let peer_ip = peer_addr.map(|a| a.ip());
2032
2033 let (reader, writer) = tokio::io::split(stream);
2034 let mut reader = BufReader::new(reader);
2035 let mut writer = BufWriter::new(writer);
2036
2037 let connect_msg = loop {
2042 match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
2043 Ok(Ok(Some(Message::Ping))) => {
2044 debug!(peer = %peer, "pre-auth ping");
2045 if !write_msg(&mut writer, &Message::Pong).await {
2046 return;
2047 }
2048 continue;
2049 }
2050 Ok(Ok(Some(msg))) => break msg,
2051 Ok(Ok(None)) => {
2052 debug!(peer = %peer, "client closed before CONNECT");
2053 return;
2054 }
2055 Ok(Err(e)) => {
2056 error!(peer = %peer, error = %e, "error reading CONNECT");
2057 return;
2058 }
2059 Err(_) => {
2060 warn!(peer = %peer, "idle timeout waiting for CONNECT");
2061 return;
2062 }
2063 }
2064 };
2065
2066 let principal: Option<Principal>;
2069 let credential_auth_configured = !users.is_empty() || expected_password.is_some();
2070 match connect_msg {
2071 Message::Connect {
2072 db_name,
2073 password,
2074 username,
2075 } => {
2076 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2078 if is_rate_limited(limiter, ip) {
2079 warn!(peer = %peer, "rate limited: too many auth failures");
2080 let err = Message::Error {
2081 message: "too many auth failures, try again later".into(),
2082 };
2083 write_msg(&mut writer, &err).await;
2084 return;
2085 }
2086 }
2087
2088 let outcome = authenticate_connect(
2089 &users,
2090 expected_password.as_ref().map(|p| p.as_str()),
2091 username.as_deref(),
2092 password.as_ref().map(|p| p.as_str()),
2093 );
2094
2095 match outcome {
2096 AuthOutcome::Rejected => {
2097 warn!(peer = %peer, db = %db_name, "auth rejected");
2098 metrics.inc_auth_failure();
2099 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2101 record_auth_failure(limiter, ip);
2102 }
2103 let err = Message::Error {
2104 message: "authentication failed".into(),
2105 };
2106 write_msg(&mut writer, &err).await;
2107 return;
2108 }
2109 AuthOutcome::Authenticated {
2110 principal: auth_principal,
2111 } => {
2112 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2114 clear_auth_failures(limiter, ip);
2115 }
2116 match &auth_principal {
2117 Some(p) => {
2118 info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
2119 }
2120 None => {
2121 info!(peer = %peer, db = %db_name, "client connected");
2122 }
2123 }
2124 principal = auth_principal;
2125 }
2126 }
2127
2128 match check_db_name(server_db_name.as_deref(), &db_name) {
2135 Ok(()) => {
2136 if server_db_name.is_none() && !db_name.is_empty() && db_name != DEFAULT_DB_NAME
2137 {
2138 static NAMED_DB_WARNED: std::sync::atomic::AtomicBool =
2139 std::sync::atomic::AtomicBool::new(false);
2140 if !NAMED_DB_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
2141 warn!(
2142 peer = %peer, db = %db_name,
2143 "client requested a named database but this server serves a single global database; name ignored"
2144 );
2145 }
2146 }
2147 }
2148 Err(msg) => {
2149 warn!(peer = %peer, db = %db_name, "rejected: unknown database");
2150 let err = Message::Error { message: msg };
2151 write_msg(&mut writer, &err).await;
2152 return;
2153 }
2154 }
2155
2156 let ok = Message::ConnectOk {
2157 version: env!("CARGO_PKG_VERSION").into(),
2158 };
2159 if !write_msg(&mut writer, &ok).await {
2160 return;
2161 }
2162 }
2163 _ => {
2164 warn!(peer = %peer, "first message was not CONNECT");
2165 let err = Message::Error {
2166 message: "expected CONNECT".into(),
2167 };
2168 write_msg(&mut writer, &err).await;
2169 return;
2170 }
2171 }
2172
2173 let mut tx_permit: Option<OwnedSemaphorePermit> = None;
2174 let mut carry: Option<Message> = None;
2177
2178 'conn: loop {
2180 let msg = if let Some(m) = carry.take() {
2181 m
2182 } else {
2183 tokio::select! {
2184 result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
2186 match result {
2187 Ok(Ok(Some(msg))) => msg,
2188 Ok(Ok(None)) => break,
2189 Ok(Err(e)) => {
2190 error!(peer = %peer, error = %e, "read error");
2191 break;
2192 }
2193 Err(_) => {
2194 info!(peer = %peer, "idle timeout, closing connection");
2195 let err = Message::Error { message: "idle timeout".into() };
2196 write_msg(&mut writer, &err).await;
2197 break;
2198 }
2199 }
2200 }
2201 _ = shutdown_rx.changed() => {
2203 if *shutdown_rx.borrow() {
2204 info!(peer = %peer, "server shutting down, closing connection");
2205 let err = Message::Error { message: "server shutting down".into() };
2206 write_msg(&mut writer, &err).await;
2207 break;
2208 }
2209 continue;
2210 }
2211 }
2212 };
2213
2214 if matches!(
2220 msg,
2221 Message::Query { .. } | Message::QuerySql { .. } | Message::QueryWithParams { .. }
2222 ) {
2223 const MAX_PIPELINE_BATCH: usize = 128;
2226 const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
2231
2232 enum BatchFatal {
2235 Closed,
2236 ReadError,
2237 }
2238
2239 fn approx_response_bytes(msg: &Message) -> usize {
2243 match msg {
2244 Message::ResultRows { columns, rows } => {
2245 columns.iter().map(|c| c.len() + 4).sum::<usize>()
2246 + rows
2247 .iter()
2248 .map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
2249 .sum::<usize>()
2250 }
2251 Message::ResultScalar { value } => value.len(),
2252 Message::ResultMessage { message } | Message::Error { message } => {
2253 message.len()
2254 }
2255 _ => 16,
2256 }
2257 }
2258
2259 fn complete_frame_buffered(buf: &[u8]) -> bool {
2265 buf.len() >= 6 && {
2266 let payload_len =
2267 u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
2268 buf.len() - 6 >= payload_len
2269 }
2270 }
2271
2272 let mut responses: Vec<Message> = Vec::new();
2273 let mut response_bytes: usize = 0;
2274 let mut last_ticket: Option<WalDurabilityTicket> = None;
2275 let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
2276 let mut fatal: Option<BatchFatal> = None;
2277 let mut current = msg;
2278 loop {
2279 let (response, ticket) = match current {
2280 Message::Query { query } => {
2281 if query.len() > MAX_QUERY_LENGTH {
2282 (
2283 Message::Error {
2284 message: format!(
2285 "query too large: {} bytes (max {})",
2286 query.len(),
2287 MAX_QUERY_LENGTH
2288 ),
2289 },
2290 None,
2291 )
2292 } else {
2293 debug!(peer = %peer, query = %query, "received query");
2294 execute_wire_query(
2295 engine.clone(),
2296 tx_gate.clone(),
2297 &mut tx_permit,
2298 query,
2299 principal.clone(),
2300 query_timeout,
2301 tx_wait_timeout,
2302 &metrics,
2303 )
2304 .await
2305 }
2306 }
2307 Message::QuerySql { query } => {
2308 if query.len() > MAX_QUERY_LENGTH {
2309 (
2310 Message::Error {
2311 message: format!(
2312 "query too large: {} bytes (max {})",
2313 query.len(),
2314 MAX_QUERY_LENGTH
2315 ),
2316 },
2317 None,
2318 )
2319 } else {
2320 debug!(peer = %peer, query = %query, "received SQL query");
2321 execute_wire_query_sql(
2322 engine.clone(),
2323 tx_gate.clone(),
2324 &mut tx_permit,
2325 query,
2326 principal.clone(),
2327 query_timeout,
2328 tx_wait_timeout,
2329 &metrics,
2330 )
2331 .await
2332 }
2333 }
2334 Message::QueryWithParams { query, params } => {
2335 if query.len() > MAX_QUERY_LENGTH {
2336 (
2337 Message::Error {
2338 message: format!(
2339 "query too large: {} bytes (max {})",
2340 query.len(),
2341 MAX_QUERY_LENGTH
2342 ),
2343 },
2344 None,
2345 )
2346 } else {
2347 debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
2348 execute_wire_query_with_params(
2349 engine.clone(),
2350 tx_gate.clone(),
2351 &mut tx_permit,
2352 query,
2353 params,
2354 principal.clone(),
2355 query_timeout,
2356 tx_wait_timeout,
2357 &metrics,
2358 )
2359 .await
2360 }
2361 }
2362 _ => unreachable!("batch loop only receives plain query frames"),
2363 };
2364 if let Some((t, m)) = ticket {
2365 last_ticket = Some(t);
2369 deferred_metrics.push(m);
2370 }
2371 response_bytes += approx_response_bytes(&response);
2372 responses.push(response);
2373
2374 if tx_permit.is_some()
2380 || responses.len() >= MAX_PIPELINE_BATCH
2381 || response_bytes >= MAX_PIPELINE_BATCH_BYTES
2382 || !complete_frame_buffered(reader.buffer())
2383 {
2384 break;
2385 }
2386 match tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)).await {
2389 Ok(Ok(Some(
2390 next @ (Message::Query { .. }
2391 | Message::QuerySql { .. }
2392 | Message::QueryWithParams { .. }),
2393 ))) => {
2394 if tx_gate.available_permits() == 0 {
2402 carry = Some(next);
2403 break;
2404 }
2405 current = next;
2406 }
2407 Ok(Ok(Some(other))) => {
2408 carry = Some(other);
2411 break;
2412 }
2413 Ok(Ok(None)) => {
2414 fatal = Some(BatchFatal::Closed);
2415 break;
2416 }
2417 Ok(Err(e)) => {
2418 error!(peer = %peer, error = %e, "read error");
2419 fatal = Some(BatchFatal::ReadError);
2420 break;
2421 }
2422 Err(_) => {
2423 error!(peer = %peer, "timeout decoding fully-buffered frame");
2426 fatal = Some(BatchFatal::ReadError);
2427 break;
2428 }
2429 }
2430 }
2431
2432 let mut durability_failed = false;
2436 if let Some(ticket) = last_ticket {
2437 if let Some(message) = settle_durability_ticket(ticket).await {
2438 durability_failed = true;
2441 for r in responses.iter_mut() {
2442 if is_success_response(r) {
2443 *r = Message::Error {
2444 message: message.clone(),
2445 };
2446 }
2447 }
2448 }
2449 }
2450 for m in deferred_metrics.drain(..) {
2451 let outcome = if m.exceeded_timeout {
2452 QueryOutcome::Timeout
2453 } else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
2454 QueryOutcome::Error
2455 } else {
2456 m.outcome
2457 };
2458 metrics.record_query(m.start.elapsed(), outcome);
2459 }
2460
2461 for r in &responses {
2462 if !write_msg(&mut writer, r).await {
2463 break 'conn;
2464 }
2465 }
2466 match fatal {
2467 None => continue,
2468 Some(BatchFatal::Closed | BatchFatal::ReadError) => break,
2469 }
2470 }
2471
2472 let response = match msg {
2473 Message::Ping => {
2474 debug!(peer = %peer, "ping");
2475 Message::Pong
2476 }
2477 Message::SyncStatus { replica_id } => {
2478 let engine = engine.clone();
2479 let principal = principal.clone();
2480 let log_context = SyncLogContext::status(&replica_id);
2481 execute_gated_sync(
2482 SyncExecutionContext {
2483 tx_gate: tx_gate.clone(),
2484 connection_has_transaction: tx_permit.is_some(),
2485 operation: SyncOperation::Status,
2486 log_context,
2487 metrics: &metrics,
2488 query_timeout,
2489 },
2490 (engine, replica_id, credential_auth_configured, principal),
2491 |(engine, replica_id, credential_authenticated, principal)| {
2492 dispatch_sync_status_decision(
2493 &engine,
2494 replica_id,
2495 credential_authenticated,
2496 principal.as_ref(),
2497 )
2498 },
2499 )
2500 .await
2501 }
2502 Message::SyncPull {
2503 replica_id,
2504 since_lsn,
2505 max_units,
2506 max_bytes,
2507 database_id,
2508 primary_generation,
2509 wal_format_version,
2510 catalog_version,
2511 segment_format_version,
2512 } => {
2513 let engine = engine.clone();
2514 let principal = principal.clone();
2515 let request = SyncPullRequest {
2516 replica_id,
2517 since_lsn,
2518 max_units,
2519 max_bytes,
2520 database_id,
2521 primary_generation,
2522 wal_format_version,
2523 catalog_version,
2524 segment_format_version,
2525 };
2526 let log_context = SyncLogContext::pull(&request);
2527 execute_gated_sync(
2528 SyncExecutionContext {
2529 tx_gate: tx_gate.clone(),
2530 connection_has_transaction: tx_permit.is_some(),
2531 operation: SyncOperation::Pull,
2532 log_context,
2533 metrics: &metrics,
2534 query_timeout,
2535 },
2536 (engine, request, credential_auth_configured, principal),
2537 |(engine, request, credential_authenticated, principal)| {
2538 dispatch_sync_pull_decision(
2539 &engine,
2540 request,
2541 credential_authenticated,
2542 principal.as_ref(),
2543 )
2544 },
2545 )
2546 .await
2547 }
2548 Message::SyncAck {
2549 replica_id,
2550 applied_lsn,
2551 remote_lsn,
2552 } => {
2553 let engine = engine.clone();
2554 let principal = principal.clone();
2555 let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
2556 execute_gated_sync(
2557 SyncExecutionContext {
2558 tx_gate: tx_gate.clone(),
2559 connection_has_transaction: tx_permit.is_some(),
2560 operation: SyncOperation::Ack,
2561 log_context,
2562 metrics: &metrics,
2563 query_timeout,
2564 },
2565 (
2566 engine,
2567 replica_id,
2568 applied_lsn,
2569 remote_lsn,
2570 credential_auth_configured,
2571 principal,
2572 ),
2573 |(
2574 engine,
2575 replica_id,
2576 applied_lsn,
2577 observed_remote_lsn,
2578 credential_authenticated,
2579 principal,
2580 )| {
2581 dispatch_sync_ack_decision(
2582 &engine,
2583 replica_id,
2584 applied_lsn,
2585 observed_remote_lsn,
2586 credential_authenticated,
2587 principal.as_ref(),
2588 )
2589 },
2590 )
2591 .await
2592 }
2593 Message::Disconnect => {
2594 debug!(peer = %peer, "received DISCONNECT");
2595 break;
2596 }
2597 _ => Message::Error {
2598 message: "unexpected message type".into(),
2599 },
2600 };
2601
2602 if !write_msg(&mut writer, &response).await {
2603 break;
2604 }
2605 }
2606
2607 if tx_permit.is_some() {
2614 let engine = engine.clone();
2615 let principal = principal.clone();
2616 let _ =
2617 tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
2618 }
2619 tx_permit.take();
2620
2621 info!(peer = %peer, "client disconnected");
2622}
2623
2624fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
2625 *total = total.saturating_add(bytes);
2626 if *total > MAX_RESPONSE_PAYLOAD_SIZE {
2627 return Err(QueryError::Execution(format!(
2628 "result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
2629 MAX_RESPONSE_PAYLOAD_SIZE
2630 )));
2631 }
2632 Ok(())
2633}
2634
2635fn query_result_to_message(result: QueryResult) -> Result<Message, QueryError> {
2636 match result {
2637 QueryResult::Rows { columns, rows } => {
2638 let mut encoded_bytes = 2usize; let mut out_columns = Vec::with_capacity(columns.len());
2640 for col in columns {
2641 charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
2642 out_columns.push(col);
2643 }
2644 charge_response_bytes(&mut encoded_bytes, 4)?; let mut str_rows = Vec::with_capacity(rows.len());
2647 for row in rows {
2648 let mut str_row = Vec::with_capacity(row.len());
2649 for value in row {
2650 let display = value_to_display(&value);
2651 charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
2652 str_row.push(display);
2653 }
2654 str_rows.push(str_row);
2655 }
2656 Ok(Message::ResultRows {
2657 columns: out_columns,
2658 rows: str_rows,
2659 })
2660 }
2661 QueryResult::Scalar(val) => Ok(Message::ResultScalar {
2662 value: value_to_display(&val),
2663 }),
2664 QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
2665 QueryResult::Created(name) => Ok(Message::ResultMessage {
2666 message: format!("type {name} created"),
2667 }),
2668 QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
2669 }
2670}
2671
2672fn value_to_display(v: &Value) -> String {
2676 v.to_wire_string()
2677}
2678
2679#[cfg(test)]
2680mod tests {
2681 use super::*;
2682 use powdb_storage::wal::WalRecordType;
2683 use powdb_sync::{
2684 write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
2685 ReplicaCursor, RetainedSegment, RetainedUnit,
2686 };
2687
2688 #[test]
2691 fn null_serializes_as_null_bareword_on_wire() {
2692 assert_eq!(value_to_display(&Value::Empty), "null");
2693 }
2694
2695 #[test]
2698 fn unique_violation_error_surfaces_to_remote_clients() {
2699 assert_eq!(
2702 sanitize_error("unique constraint violation on User.email"),
2703 "unique constraint violation on User.email"
2704 );
2705 }
2706
2707 #[test]
2708 fn internal_errors_stay_generic() {
2709 assert_eq!(
2710 sanitize_error("some internal io panic detail"),
2711 "query execution error"
2712 );
2713 }
2714
2715 #[test]
2718 fn json_cell_renders_canonical_text_on_wire() {
2719 let pj1 = powdb_storage::pj1::parse_json_text(r#"{"b":2,"a":1,"nested":{"z":true}}"#)
2724 .expect("valid JSON");
2725 let result = QueryResult::Rows {
2726 columns: vec!["doc".into()],
2727 rows: vec![vec![Value::Json(pj1.into())]],
2728 };
2729 match query_result_to_message(result).expect("encodes") {
2730 Message::ResultRows { columns, rows } => {
2731 assert_eq!(columns, vec!["doc"]);
2732 assert_eq!(
2733 rows,
2734 vec![vec![r#"{"a":1,"b":2,"nested":{"z":true}}"#.to_string()]]
2735 );
2736 }
2737 other => panic!("expected ResultRows, got {other:?}"),
2738 }
2739 }
2740
2741 #[test]
2742 fn json_parse_error_surfaces_to_remote_clients() {
2743 for msg in [
2752 "type mismatch: invalid JSON: unexpected character 'x' at position 3",
2753 "invalid JSON: nesting exceeds depth cap 128",
2754 ] {
2755 assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
2756 }
2757 assert_eq!(
2758 sanitize_error("malformed PJ1: truncated"),
2759 "query execution error",
2760 "internal storage corruption must stay masked"
2761 );
2762 }
2763
2764 #[test]
2770 fn describe_shows_json_type_over_the_wire() {
2771 let dir = tempfile::tempdir().unwrap();
2772 let mut engine = Engine::new(dir.path()).unwrap();
2773 engine
2774 .execute_powql("type Doc { required id: int, body: json }")
2775 .expect("json column DDL should be accepted once Lane B lands");
2776 let result = engine.execute_powql("describe Doc").expect("describe runs");
2777 let msg = query_result_to_message(result).expect("encodes");
2778 match msg {
2779 Message::ResultRows { columns, rows } => {
2780 assert_eq!(columns[1], "type");
2781 let body = rows
2783 .iter()
2784 .find(|r| r[0] == "body")
2785 .expect("body column present");
2786 assert_eq!(body[1], "json");
2787 }
2788 other => panic!("expected ResultRows, got {other:?}"),
2789 }
2790 }
2791
2792 #[test]
2795 fn db_name_unpinned_accepts_any_name() {
2796 for requested in ["", "default", "prod", "anything"] {
2797 assert!(
2798 check_db_name(None, requested).is_ok(),
2799 "rejected {requested}"
2800 );
2801 }
2802 }
2803
2804 #[test]
2805 fn db_name_pinned_accepts_match_empty_and_default_sentinel() {
2806 assert!(check_db_name(Some("prod"), "prod").is_ok());
2809 assert!(check_db_name(Some("prod"), "").is_ok());
2810 assert!(check_db_name(Some("prod"), DEFAULT_DB_NAME).is_ok());
2811 }
2812
2813 #[test]
2814 fn db_name_pinned_rejects_foreign_with_clear_message() {
2815 let err = check_db_name(Some("prod"), "staging").unwrap_err();
2816 assert_eq!(err, "unknown database 'staging'; this server serves 'prod'");
2817 }
2818
2819 #[tokio::test]
2822 async fn begin_permit_acquires_when_gate_is_free() {
2823 let gate = new_tx_gate();
2824 let metrics = Arc::new(Metrics::new());
2825 let permit = acquire_begin_permit(&gate, Duration::from_secs(5), &metrics)
2826 .await
2827 .expect("should acquire a free gate");
2828 assert_eq!(gate.available_permits(), 0, "permit must be held");
2829 drop(permit);
2830 assert_eq!(gate.available_permits(), 1, "permit must release on drop");
2831 }
2832
2833 #[tokio::test]
2834 async fn begin_permit_times_out_with_clear_error_and_truthful_metric() {
2835 let gate = new_tx_gate();
2836 let metrics = Arc::new(Metrics::new());
2837 let _held = gate.clone().acquire_owned().await.unwrap();
2839 let err = acquire_begin_permit(&gate, Duration::from_millis(25), &metrics)
2840 .await
2841 .expect_err("must time out while the gate is held");
2842 match err {
2843 Message::Error { message } => {
2844 assert!(
2845 message.contains("transaction gate timeout after 25ms"),
2846 "unexpected message: {message}"
2847 );
2848 assert!(
2849 message.contains("waiting for concurrent transaction"),
2850 "unexpected message: {message}"
2851 );
2852 }
2853 other => panic!("expected Error, got {other:?}"),
2854 }
2855 let rendered = metrics.render();
2856 assert!(rendered.contains("powdb_tx_gate_timeouts_total 1"));
2857 assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
2859 }
2860
2861 #[test]
2862 fn resource_limit_errors_surface_actionable_hints() {
2863 for msg in [
2868 "sort input exceeds row limit — add a LIMIT clause",
2869 "join result exceeds row limit",
2870 "query exceeded memory budget: requested 100 bytes, limit 50 bytes",
2871 "result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
2872 ] {
2873 assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
2874 }
2875 }
2876
2877 #[test]
2878 fn oversized_result_is_rejected_before_wire_encoding() {
2879 let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
2880 let result = QueryResult::Rows {
2881 columns: vec!["payload".into()],
2882 rows: vec![vec![Value::Str(long)]],
2883 };
2884 let err = query_result_to_message(result).unwrap_err();
2885 assert!(
2886 err.to_string().starts_with("result too large"),
2887 "unexpected error: {err}"
2888 );
2889 }
2890
2891 fn parsed(q: &str) -> powdb_query::ast::Statement {
2894 parser::parse(q).unwrap()
2895 }
2896
2897 fn principal(role: &str) -> Option<Principal> {
2898 Some(Principal {
2899 name: "u".into(),
2900 role: role.into(),
2901 })
2902 }
2903
2904 #[test]
2905 fn readonly_can_read_but_not_write() {
2906 let p = principal("readonly");
2907 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2909 assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
2910 assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
2911 for q in [
2913 r#"insert User { name := "x" }"#,
2914 "User filter .id = 1 update { age := 2 }",
2915 "User filter .id = 1 delete",
2916 "drop User",
2917 "alter User add column c: str",
2918 "type T { required id: int }",
2919 "begin",
2920 "commit",
2921 "rollback",
2922 ] {
2923 let err = check_statement_permitted(p.as_ref(), &parsed(q))
2924 .expect_err(&format!("must deny: {q}"));
2925 assert!(
2926 err.to_string().contains("permission denied"),
2927 "unexpected error for {q}: {err}"
2928 );
2929 }
2930 }
2931
2932 #[test]
2933 fn readwrite_and_admin_have_full_query_access() {
2934 for role in ["readwrite", "admin"] {
2935 let p = principal(role);
2936 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2937 assert!(check_statement_permitted(
2938 p.as_ref(),
2939 &parsed(r#"insert User { name := "x" }"#)
2940 )
2941 .is_ok());
2942 assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
2943 }
2944 }
2945
2946 #[test]
2947 fn unknown_role_fails_closed_for_writes() {
2948 let p = principal("mystery");
2949 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2950 assert!(
2951 check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
2952 .is_err()
2953 );
2954 }
2955
2956 #[test]
2957 fn no_principal_means_full_access() {
2958 assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
2960 assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
2961 }
2962
2963 fn store_with_alice() -> UserStore {
2964 let mut s = UserStore::new();
2965 s.create_user("alice", "pw", "readwrite").unwrap();
2966 s
2967 }
2968
2969 #[test]
2972 fn empty_store_no_password_is_open() {
2973 let s = UserStore::new();
2974 assert_eq!(
2975 authenticate_connect(&s, None, None, None),
2976 AuthOutcome::Authenticated { principal: None }
2977 );
2978 assert_eq!(
2980 authenticate_connect(&s, None, Some("x"), Some("y")),
2981 AuthOutcome::Authenticated { principal: None }
2982 );
2983 }
2984
2985 #[test]
2986 fn empty_store_correct_shared_password_succeeds() {
2987 let s = UserStore::new();
2988 assert_eq!(
2989 authenticate_connect(&s, Some("pw"), None, Some("pw")),
2990 AuthOutcome::Authenticated { principal: None }
2991 );
2992 }
2993
2994 #[test]
2995 fn empty_store_wrong_shared_password_rejected() {
2996 let s = UserStore::new();
2997 assert_eq!(
2998 authenticate_connect(&s, Some("pw"), None, Some("bad")),
2999 AuthOutcome::Rejected
3000 );
3001 }
3002
3003 #[test]
3004 fn empty_store_missing_password_rejected_when_expected() {
3005 let s = UserStore::new();
3006 assert_eq!(
3007 authenticate_connect(&s, Some("pw"), None, None),
3008 AuthOutcome::Rejected
3009 );
3010 }
3011
3012 #[test]
3013 fn empty_store_ignores_username_for_shared_password() {
3014 let s = UserStore::new();
3017 assert_eq!(
3018 authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
3019 AuthOutcome::Authenticated { principal: None }
3020 );
3021 }
3022
3023 #[test]
3026 fn user_auth_success_binds_principal() {
3027 let s = store_with_alice();
3028 assert_eq!(
3029 authenticate_connect(&s, None, Some("alice"), Some("pw")),
3030 AuthOutcome::Authenticated {
3031 principal: Some(Principal {
3032 name: "alice".into(),
3033 role: "readwrite".into(),
3034 })
3035 }
3036 );
3037 }
3038
3039 #[test]
3040 fn user_auth_wrong_password_rejected() {
3041 let s = store_with_alice();
3042 assert_eq!(
3043 authenticate_connect(&s, None, Some("alice"), Some("bad")),
3044 AuthOutcome::Rejected
3045 );
3046 }
3047
3048 #[test]
3049 fn user_auth_unknown_user_rejected() {
3050 let s = store_with_alice();
3051 assert_eq!(
3052 authenticate_connect(&s, None, Some("mallory"), Some("pw")),
3053 AuthOutcome::Rejected
3054 );
3055 }
3056
3057 #[test]
3058 fn user_auth_missing_username_rejected() {
3059 let s = store_with_alice();
3060 assert_eq!(
3061 authenticate_connect(&s, None, None, Some("pw")),
3062 AuthOutcome::Rejected
3063 );
3064 }
3065
3066 #[test]
3067 fn user_auth_missing_password_rejected() {
3068 let s = store_with_alice();
3069 assert_eq!(
3070 authenticate_connect(&s, Some("pw"), Some("alice"), None),
3071 AuthOutcome::Rejected
3072 );
3073 }
3074
3075 #[test]
3076 fn user_auth_ignores_shared_password_when_users_present() {
3077 let s = store_with_alice();
3080 assert_eq!(
3081 authenticate_connect(&s, Some("shared"), None, Some("shared")),
3082 AuthOutcome::Rejected
3083 );
3084 }
3085
3086 #[test]
3087 fn replica_fingerprint_is_stable_and_redacted() {
3088 let replica_id = "customer-prod-replica-a";
3089 let fingerprint = replica_fingerprint(replica_id);
3090 assert_eq!(fingerprint, replica_fingerprint(replica_id));
3091 assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
3092 assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
3093 assert_eq!(fingerprint.len(), 16);
3094 assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
3095 assert!(!fingerprint.contains("customer"));
3096 assert!(!fingerprint.contains("replica"));
3097 assert!(!fingerprint.contains(replica_id));
3098 }
3099
3100 #[test]
3101 fn invalid_replica_ids_use_fixed_log_fingerprint() {
3102 assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
3103 assert_eq!(
3104 log_replica_fingerprint("customer/prod/replica"),
3105 INVALID_REPLICA_FINGERPRINT
3106 );
3107 assert_eq!(
3108 log_replica_fingerprint(&"a".repeat(4096)),
3109 INVALID_REPLICA_FINGERPRINT
3110 );
3111 }
3112
3113 #[test]
3114 fn sync_error_classes_use_bounded_labels() {
3115 assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
3116 assert_eq!(
3117 SyncErrorClass::PermissionDenied.as_label(),
3118 "permission_denied"
3119 );
3120 assert_eq!(
3121 SyncErrorClass::IdentityOrFormatMismatch.as_label(),
3122 "identity_or_format_mismatch"
3123 );
3124 assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
3125 assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
3126 }
3127
3128 fn sync_identity() -> DatabaseIdentity {
3129 DatabaseIdentity {
3130 database_id: *b"server-sync-test",
3131 primary_generation: 1,
3132 }
3133 }
3134
3135 fn retained_unit(lsn: u64) -> RetainedUnit {
3136 RetainedUnit {
3137 tx_id: 1,
3138 record_type: 4,
3139 lsn,
3140 data: lsn.to_le_bytes().to_vec(),
3141 }
3142 }
3143
3144 fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
3145 RetainedUnit {
3146 tx_id,
3147 record_type: record_type as u8,
3148 lsn,
3149 data: lsn.to_le_bytes().to_vec(),
3150 }
3151 }
3152
3153 fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
3154 let identity = sync_identity();
3155 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
3156 let units = (1..=through_lsn).map(retained_unit).collect();
3157 let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
3158 write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
3159 }
3160
3161 fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
3162 let identity = sync_identity();
3163 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
3164 let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
3165 write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
3166 }
3167
3168 fn write_sync_identity_only(data_dir: &std::path::Path) {
3169 let identity = sync_identity();
3170 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
3171 }
3172
3173 fn admin_principal() -> Principal {
3174 Principal {
3175 name: "admin".into(),
3176 role: "admin".into(),
3177 }
3178 }
3179
3180 #[test]
3181 fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
3182 let dir = tempfile::tempdir().unwrap();
3183 let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
3184
3185 match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
3186 Message::Error { message } => {
3187 assert!(message.contains("requires authentication"));
3188 }
3189 other => panic!("expected auth error, got {other:?}"),
3190 }
3191
3192 let readonly = Principal {
3193 name: "reader".into(),
3194 role: "readonly".into(),
3195 };
3196 match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
3197 Message::Error { message } => {
3198 assert!(message.contains("permission denied"));
3199 }
3200 other => panic!("expected permission error, got {other:?}"),
3201 }
3202 }
3203
3204 #[test]
3205 fn sync_status_pull_and_ack_use_server_remote_lsn() {
3206 let dir = tempfile::tempdir().unwrap();
3207 let mut engine = Engine::new(dir.path()).unwrap();
3208 engine
3209 .execute_powql("type SyncT { required id: int, v: str }")
3210 .unwrap();
3211 engine
3212 .execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
3213 .unwrap();
3214 let remote_lsn = engine.catalog().max_lsn();
3215 assert!(remote_lsn > 0);
3216 write_sync_identity_and_tail(dir.path(), remote_lsn);
3217 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3218 .unwrap();
3219
3220 let engine = Arc::new(RwLock::new(engine));
3221 let principal = admin_principal();
3222 let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
3223 {
3224 Message::SyncStatusResult { status } => status,
3225 other => panic!("expected sync status, got {other:?}"),
3226 };
3227 assert_eq!(status.remote_lsn, remote_lsn);
3228 assert_eq!(status.servable_lsn, Some(remote_lsn));
3229 assert_eq!(status.unarchived_lsn, Some(0));
3230 assert_eq!(status.last_applied_lsn, Some(0));
3231 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3232 assert!(status.stale);
3233
3234 let identity = sync_identity().segment_identity();
3235 let pull = SyncPullRequest {
3236 replica_id: "replica-a".into(),
3237 since_lsn: 0,
3238 max_units: MAX_SYNC_PULL_UNITS,
3239 max_bytes: MAX_SYNC_PULL_BYTES,
3240 database_id: identity.database_id,
3241 primary_generation: identity.primary_generation,
3242 wal_format_version: identity.wal_format_version,
3243 catalog_version: identity.catalog_version,
3244 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3245 };
3246 let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3247 Message::SyncPullResult {
3248 status,
3249 units,
3250 has_more,
3251 } => {
3252 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3253 assert!(!has_more);
3254 units
3255 }
3256 other => panic!("expected sync pull result, got {other:?}"),
3257 };
3258 assert_eq!(units.len() as u64, remote_lsn);
3259 assert_eq!(units.last().unwrap().lsn, remote_lsn);
3260
3261 let ack = match dispatch_sync_ack(
3262 &engine,
3263 "replica-a".into(),
3264 remote_lsn,
3265 remote_lsn,
3266 true,
3267 Some(&principal),
3268 ) {
3269 Message::SyncAckResult {
3270 previous_applied_lsn,
3271 applied_lsn,
3272 remote_lsn: ack_remote_lsn,
3273 advanced,
3274 status,
3275 } => {
3276 assert_eq!(previous_applied_lsn, 0);
3277 assert_eq!(applied_lsn, remote_lsn);
3278 assert_eq!(ack_remote_lsn, remote_lsn);
3279 assert!(advanced);
3280 status
3281 }
3282 other => panic!("expected sync ack result, got {other:?}"),
3283 };
3284 assert_eq!(ack.repair_action, WireSyncRepairAction::None);
3285 assert!(!ack.stale);
3286 assert_eq!(ack.lag_lsn, Some(0));
3287 }
3288
3289 #[test]
3290 fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
3291 let dir = tempfile::tempdir().unwrap();
3292 let mut engine = Engine::new(dir.path()).unwrap();
3293 engine
3294 .execute_powql("type SyncT { required id: int }")
3295 .unwrap();
3296 for id in 1..=3 {
3297 engine
3298 .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
3299 .unwrap();
3300 }
3301 let remote_lsn = engine.catalog().max_lsn();
3302 assert!(remote_lsn >= 3);
3303 write_sync_identity_and_units(
3304 dir.path(),
3305 vec![
3306 retained_unit_with(77, WalRecordType::Begin, 1),
3307 retained_unit_with(77, WalRecordType::Insert, 2),
3308 retained_unit_with(77, WalRecordType::Commit, 3),
3309 ],
3310 );
3311 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3312 .unwrap();
3313
3314 let engine = Arc::new(RwLock::new(engine));
3315 let principal = admin_principal();
3316 let identity = sync_identity().segment_identity();
3317 let cut_pull = SyncPullRequest {
3318 replica_id: "replica-a".into(),
3319 since_lsn: 0,
3320 max_units: 2,
3321 max_bytes: MAX_SYNC_PULL_BYTES,
3322 database_id: identity.database_id,
3323 primary_generation: identity.primary_generation,
3324 wal_format_version: identity.wal_format_version,
3325 catalog_version: identity.catalog_version,
3326 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3327 };
3328 match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
3329 Message::Error { message } => assert!(message.contains("cuts through transaction")),
3330 other => panic!("expected transaction-cut pull error, got {other:?}"),
3331 }
3332
3333 let cut_bytes_pull = SyncPullRequest {
3334 replica_id: "replica-a".into(),
3335 since_lsn: 0,
3336 max_units: 3,
3337 max_bytes: 58,
3338 database_id: identity.database_id,
3339 primary_generation: identity.primary_generation,
3340 wal_format_version: identity.wal_format_version,
3341 catalog_version: identity.catalog_version,
3342 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3343 };
3344 match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
3345 Message::Error { message } => assert!(message.contains("cuts through transaction")),
3346 other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
3347 }
3348
3349 let full_pull = SyncPullRequest {
3350 replica_id: "replica-a".into(),
3351 since_lsn: 0,
3352 max_units: 3,
3353 max_bytes: MAX_SYNC_PULL_BYTES,
3354 database_id: identity.database_id,
3355 primary_generation: identity.primary_generation,
3356 wal_format_version: identity.wal_format_version,
3357 catalog_version: identity.catalog_version,
3358 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3359 };
3360 match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
3361 Message::SyncPullResult { units, .. } => {
3362 assert_eq!(units.len(), 3);
3363 assert_eq!(units.last().unwrap().lsn, 3);
3364 }
3365 other => panic!("expected complete transaction pull, got {other:?}"),
3366 }
3367
3368 match dispatch_sync_ack(
3369 &engine,
3370 "replica-a".into(),
3371 2,
3372 remote_lsn,
3373 true,
3374 Some(&principal),
3375 ) {
3376 Message::Error { message } => assert!(message.contains("cuts through transaction")),
3377 other => panic!("expected transaction-cut ack error, got {other:?}"),
3378 }
3379 let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
3380 assert_eq!(cursor[0].applied_lsn, 0);
3381
3382 match dispatch_sync_ack(
3383 &engine,
3384 "replica-a".into(),
3385 3,
3386 remote_lsn,
3387 true,
3388 Some(&principal),
3389 ) {
3390 Message::SyncAckResult {
3391 previous_applied_lsn,
3392 applied_lsn,
3393 advanced,
3394 ..
3395 } => {
3396 assert_eq!(previous_applied_lsn, 0);
3397 assert_eq!(applied_lsn, 3);
3398 assert!(advanced);
3399 }
3400 other => panic!("expected complete transaction ack, got {other:?}"),
3401 }
3402 }
3403
3404 #[test]
3405 fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
3406 let dir = tempfile::tempdir().unwrap();
3407 let mut engine = Engine::new(dir.path()).unwrap();
3408 engine
3409 .execute_powql("type SyncT { required id: int }")
3410 .unwrap();
3411 for id in 1..=6 {
3412 engine
3413 .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
3414 .unwrap();
3415 }
3416 let remote_lsn = engine.catalog().max_lsn();
3417 assert!(remote_lsn >= 6);
3418 write_sync_identity_and_units(
3419 dir.path(),
3420 vec![
3421 retained_unit_with(1, WalRecordType::Begin, 1),
3422 retained_unit_with(1, WalRecordType::Insert, 2),
3423 retained_unit_with(1, WalRecordType::Commit, 3),
3424 retained_unit_with(1, WalRecordType::Begin, 4),
3425 retained_unit_with(1, WalRecordType::Insert, 5),
3426 retained_unit_with(1, WalRecordType::Commit, 6),
3427 ],
3428 );
3429 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3430 .unwrap();
3431
3432 let engine = Arc::new(RwLock::new(engine));
3433 let principal = admin_principal();
3434 let identity = sync_identity().segment_identity();
3435 let pull = SyncPullRequest {
3436 replica_id: "replica-a".into(),
3437 since_lsn: 0,
3438 max_units: 6,
3439 max_bytes: 100,
3440 database_id: identity.database_id,
3441 primary_generation: identity.primary_generation,
3442 wal_format_version: identity.wal_format_version,
3443 catalog_version: identity.catalog_version,
3444 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3445 };
3446
3447 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3448 Message::SyncPullResult {
3449 status,
3450 units,
3451 has_more,
3452 } => {
3453 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3454 assert_eq!(units.len(), 3);
3455 assert_eq!(units.last().unwrap().lsn, 3);
3456 assert!(has_more);
3457 }
3458 other => panic!("expected byte-capped applyable prefix, got {other:?}"),
3459 }
3460 }
3461
3462 #[test]
3463 fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
3464 let dir = tempfile::tempdir().unwrap();
3465 let mut engine = Engine::new(dir.path()).unwrap();
3466 engine
3467 .execute_powql("type SyncT { required id: int }")
3468 .unwrap();
3469 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3470 let remote_lsn = engine.catalog().max_lsn();
3471 assert!(remote_lsn > 0);
3472 write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
3473 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3474 .unwrap();
3475
3476 let engine = Arc::new(RwLock::new(engine));
3477 let principal = admin_principal();
3478 let identity = sync_identity().segment_identity();
3479 let pull = SyncPullRequest {
3480 replica_id: "replica-a".into(),
3481 since_lsn: 0,
3482 max_units: MAX_SYNC_PULL_UNITS,
3483 max_bytes: MAX_SYNC_PULL_BYTES,
3484 database_id: identity.database_id,
3485 primary_generation: identity.primary_generation,
3486 wal_format_version: identity.wal_format_version,
3487 catalog_version: identity.catalog_version,
3488 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3489 };
3490
3491 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3492 Message::SyncPullResult {
3493 status,
3494 units,
3495 has_more,
3496 } => {
3497 assert_eq!(status.remote_lsn, remote_lsn);
3498 assert_eq!(status.servable_lsn, Some(remote_lsn));
3499 assert_eq!(status.unarchived_lsn, Some(0));
3500 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3501 assert!(!has_more);
3502 assert_eq!(units.len() as u64, remote_lsn);
3503 assert_eq!(units.last().unwrap().lsn, remote_lsn);
3504 assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
3505 }
3506 other => panic!("expected capped sync pull result, got {other:?}"),
3507 }
3508 }
3509
3510 #[test]
3511 fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
3512 let dir = tempfile::tempdir().unwrap();
3513 let mut engine = Engine::new(dir.path()).unwrap();
3514 engine
3515 .execute_powql("type SyncT { required id: int }")
3516 .unwrap();
3517 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3518 let remote_lsn = engine.catalog().max_lsn();
3519 assert!(remote_lsn > 0);
3520 write_sync_identity_only(dir.path());
3521 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3522 .unwrap();
3523
3524 let engine = Arc::new(RwLock::new(engine));
3525 let principal = admin_principal();
3526 let identity = sync_identity().segment_identity();
3527 let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
3528 {
3529 Message::SyncStatusResult { status } => status,
3530 other => panic!("expected sync status, got {other:?}"),
3531 };
3532 assert_eq!(status.remote_lsn, remote_lsn);
3533 assert_eq!(status.servable_lsn, Some(0));
3534 assert_eq!(status.unarchived_lsn, Some(remote_lsn));
3535 assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
3536 assert!(status
3537 .last_sync_error
3538 .as_deref()
3539 .unwrap()
3540 .contains("not yet archived"));
3541
3542 let pull = SyncPullRequest {
3543 replica_id: "replica-a".into(),
3544 since_lsn: 0,
3545 max_units: MAX_SYNC_PULL_UNITS,
3546 max_bytes: MAX_SYNC_PULL_BYTES,
3547 database_id: identity.database_id,
3548 primary_generation: identity.primary_generation,
3549 wal_format_version: identity.wal_format_version,
3550 catalog_version: identity.catalog_version,
3551 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3552 };
3553 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3554 Message::SyncPullResult {
3555 status,
3556 units,
3557 has_more,
3558 } => {
3559 assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
3560 assert!(units.is_empty());
3561 assert!(!has_more);
3562 }
3563 other => panic!("expected await-archive sync pull result, got {other:?}"),
3564 }
3565 }
3566
3567 #[test]
3568 fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
3569 let dir = tempfile::tempdir().unwrap();
3570 let mut engine = Engine::new(dir.path()).unwrap();
3571 engine
3572 .execute_powql("type SyncT { required id: int }")
3573 .unwrap();
3574 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3575 engine.execute_powql("insert SyncT { id := 2 }").unwrap();
3576 let remote_lsn = engine.catalog().max_lsn();
3577 assert!(remote_lsn > 1);
3578 let servable_lsn = remote_lsn - 1;
3579 write_sync_identity_and_tail(dir.path(), servable_lsn);
3580 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3581 .unwrap();
3582
3583 let engine = Arc::new(RwLock::new(engine));
3584 let principal = admin_principal();
3585 let identity = sync_identity().segment_identity();
3586 let pull = SyncPullRequest {
3587 replica_id: "replica-a".into(),
3588 since_lsn: 0,
3589 max_units: MAX_SYNC_PULL_UNITS,
3590 max_bytes: MAX_SYNC_PULL_BYTES,
3591 database_id: identity.database_id,
3592 primary_generation: identity.primary_generation,
3593 wal_format_version: identity.wal_format_version,
3594 catalog_version: identity.catalog_version,
3595 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3596 };
3597
3598 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3599 Message::SyncPullResult {
3600 status,
3601 units,
3602 has_more,
3603 } => {
3604 assert_eq!(status.remote_lsn, remote_lsn);
3605 assert_eq!(status.servable_lsn, Some(servable_lsn));
3606 assert_eq!(status.unarchived_lsn, Some(1));
3607 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3608 assert!(!has_more);
3609 assert_eq!(units.len() as u64, servable_lsn);
3610 assert_eq!(units.last().unwrap().lsn, servable_lsn);
3611 assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
3612 }
3613 other => panic!("expected partial sync pull result, got {other:?}"),
3614 }
3615 }
3616
3617 #[test]
3618 fn sync_pull_rejects_cursor_or_format_mismatch() {
3619 let dir = tempfile::tempdir().unwrap();
3620 let mut engine = Engine::new(dir.path()).unwrap();
3621 engine
3622 .execute_powql("type SyncT { required id: int }")
3623 .unwrap();
3624 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3625 let remote_lsn = engine.catalog().max_lsn();
3626 write_sync_identity_and_tail(dir.path(), remote_lsn);
3627 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3628 .unwrap();
3629 let engine = Arc::new(RwLock::new(engine));
3630 let principal = admin_principal();
3631 let identity = sync_identity().segment_identity();
3632
3633 let wrong_cursor = SyncPullRequest {
3634 replica_id: "replica-a".into(),
3635 since_lsn: 1,
3636 max_units: 10,
3637 max_bytes: MAX_SYNC_PULL_BYTES,
3638 database_id: identity.database_id,
3639 primary_generation: identity.primary_generation,
3640 wal_format_version: identity.wal_format_version,
3641 catalog_version: identity.catalog_version,
3642 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3643 };
3644 match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
3645 Message::Error { message } => assert!(message.contains("does not match")),
3646 other => panic!("expected cursor mismatch error, got {other:?}"),
3647 }
3648
3649 let wrong_format = SyncPullRequest {
3650 replica_id: "replica-a".into(),
3651 since_lsn: 0,
3652 max_units: 10,
3653 max_bytes: MAX_SYNC_PULL_BYTES,
3654 database_id: identity.database_id,
3655 primary_generation: identity.primary_generation,
3656 wal_format_version: identity.wal_format_version,
3657 catalog_version: identity.catalog_version,
3658 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
3659 };
3660 match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
3661 Message::Error { message } => assert!(message.contains("rebootstrap required")),
3662 other => panic!("expected format mismatch error, got {other:?}"),
3663 }
3664 }
3665}