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 SAFE_ERROR_PREFIXES: &[&str] = &[
265 "table not found",
266 "column not found",
267 "parse error",
268 "type mismatch",
269 "unknown table",
270 "unknown column",
271 "unknown function",
272 "syntax error",
273 "expected",
274 "unexpected",
275 "missing",
276 "duplicate",
277 "invalid",
278 "cannot",
279 "no such",
280 "already exists",
281 "permission denied",
282 "row too large",
283 "unique constraint violation",
284 "sort input exceeds",
289 "join result exceeds",
290 "query exceeded memory budget",
291 "result too large",
292 "wal durability sync failed",
297];
298
299fn sanitize_error(e: &str) -> String {
303 let lower = e.to_lowercase();
304 for prefix in SAFE_ERROR_PREFIXES {
305 if lower.starts_with(prefix) {
306 return e.to_string();
307 }
308 }
309 "query execution error".into()
310}
311
312async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
315 let write_fut = async {
316 if msg.write_to(writer).await.is_err() {
317 return false;
318 }
319 writer.flush().await.is_ok()
320 };
321 tokio::time::timeout(WRITE_TIMEOUT, write_fut)
322 .await
323 .unwrap_or_default()
324}
325
326pub struct ConnOpts<'a> {
329 pub engine: Arc<RwLock<Engine>>,
330 pub tx_gate: TxGate,
331 pub expected_password: Option<Zeroizing<String>>,
334 pub users: Arc<UserStore>,
338 pub shutdown_rx: &'a mut watch::Receiver<bool>,
339 pub idle_timeout: Duration,
340 pub query_timeout: Duration,
341 pub rate_limiter: Option<&'a AuthRateLimiter>,
342 pub peer_addr: Option<std::net::SocketAddr>,
343 pub metrics: Arc<Metrics>,
345}
346
347type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
360
361fn execute_write_deferred(
368 engine: &Arc<RwLock<Engine>>,
369 f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
370) -> DispatchOutcome {
371 let mut eng = match engine.write() {
372 Ok(eng) => eng,
373 Err(e) => {
374 return (
375 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
376 None,
377 )
378 }
379 };
380 eng.run_with_deferred_durability(f)
381}
382
383fn dispatch_query(
384 engine: &Arc<RwLock<Engine>>,
385 query: &str,
386 principal: Option<&Principal>,
387) -> DispatchOutcome {
388 let stmt_result = parser::parse(query).map_err(|e| e.to_string());
389
390 if let Ok(stmt) = &stmt_result {
394 if let Err(e) = check_statement_permitted(principal, stmt) {
395 return (Err(e), None);
396 }
397 }
398
399 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
400 if can_try_read {
401 let res = {
402 let eng = match engine.read() {
403 Ok(eng) => eng,
404 Err(e) => {
405 return (
406 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
407 None,
408 )
409 }
410 };
411 eng.execute_powql_readonly(query)
412 };
413 match res {
414 Ok(r) => return (Ok(r), None),
415 Err(QueryError::ReadonlyNeedsWrite) => {
416 }
418 Err(e) => return (Err(e), None),
419 }
420 }
421
422 if matches!(
423 parsed_transaction_control(&stmt_result),
424 Some(TransactionControl::Rollback)
425 ) {
426 let mut eng = match engine.write() {
427 Ok(eng) => eng,
428 Err(e) => {
429 return (
430 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
431 None,
432 )
433 }
434 };
435 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
436 }
437 execute_write_deferred(engine, |eng| eng.execute_powql(query))
438}
439
440fn dispatch_sql_query(
441 engine: &Arc<RwLock<Engine>>,
442 query: &str,
443 principal: Option<&Principal>,
444) -> DispatchOutcome {
445 let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
446
447 if let Ok(stmt) = &stmt_result {
448 if let Err(e) = check_statement_permitted(principal, stmt) {
449 return (Err(e), None);
450 }
451 }
452
453 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
454 if can_try_read {
455 let res = {
456 let eng = match engine.read() {
457 Ok(eng) => eng,
458 Err(e) => {
459 return (
460 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
461 None,
462 )
463 }
464 };
465 eng.execute_sql_readonly(query)
466 };
467 match res {
468 Ok(r) => return (Ok(r), None),
469 Err(QueryError::ReadonlyNeedsWrite) => {}
470 Err(e) => return (Err(e), None),
471 }
472 }
473
474 if matches!(
475 parsed_transaction_control(&stmt_result),
476 Some(TransactionControl::Rollback)
477 ) {
478 let mut eng = match engine.write() {
479 Ok(eng) => eng,
480 Err(e) => {
481 return (
482 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
483 None,
484 )
485 }
486 };
487 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
488 }
489 execute_write_deferred(engine, |eng| eng.execute_sql(query))
490}
491
492fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
495 use powdb_query::ast::ParamValue;
496 match p {
497 WireParam::Null => ParamValue::Null,
498 WireParam::Int(v) => ParamValue::Int(*v),
499 WireParam::Float(v) => ParamValue::Float(*v),
500 WireParam::Bool(v) => ParamValue::Bool(*v),
501 WireParam::Str(s) => ParamValue::Str(s.clone()),
502 }
503}
504
505#[derive(Debug, Clone, Copy, PartialEq, Eq)]
511enum TransactionControl {
512 Begin,
513 Commit,
514 Rollback,
515}
516
517fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
518 use powdb_query::ast::Statement;
519 match stmt {
520 Statement::Begin => Some(TransactionControl::Begin),
521 Statement::Commit => Some(TransactionControl::Commit),
522 Statement::Rollback => Some(TransactionControl::Rollback),
523 _ => None,
524 }
525}
526
527fn classify_query_transaction_control(query: &str) -> Option<TransactionControl> {
528 parser::parse(query)
529 .ok()
530 .and_then(|stmt| transaction_control(&stmt))
531}
532
533fn classify_sql_transaction_control(query: &str) -> Option<TransactionControl> {
534 sql::parse_sql(query)
535 .ok()
536 .and_then(|stmt| transaction_control(&stmt))
537}
538
539fn classify_params_transaction_control(
540 query: &str,
541 params: &[WireParam],
542) -> Option<TransactionControl> {
543 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
544 parser::parse_with_params(query, &bound)
545 .ok()
546 .and_then(|stmt| transaction_control(&stmt))
547}
548
549fn parsed_transaction_control(
550 stmt_result: &Result<powdb_query::ast::Statement, String>,
551) -> Option<TransactionControl> {
552 stmt_result.as_ref().ok().and_then(transaction_control)
553}
554
555fn execute_rollback_preserving_sync_if_needed(
556 engine: &mut Engine,
557) -> Result<QueryResult, QueryError> {
558 engine.rollback_transaction_preserving_wal_archive()
559}
560
561fn dispatch_query_with_params(
562 engine: &Arc<RwLock<Engine>>,
563 query: &str,
564 params: &[WireParam],
565 principal: Option<&Principal>,
566) -> DispatchOutcome {
567 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
568
569 let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
572
573 if let Ok(stmt) = &stmt_result {
574 if let Err(e) = check_statement_permitted(principal, stmt) {
575 return (Err(e), None);
576 }
577 }
578
579 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
580 if can_try_read {
581 let res = {
582 let eng = match engine.read() {
583 Ok(eng) => eng,
584 Err(e) => {
585 return (
586 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
587 None,
588 )
589 }
590 };
591 eng.execute_powql_readonly_with_params(query, &bound)
592 };
593 match res {
594 Ok(r) => return (Ok(r), None),
595 Err(QueryError::ReadonlyNeedsWrite) => {
596 }
598 Err(e) => return (Err(e), None),
599 }
600 }
601
602 if matches!(
603 parsed_transaction_control(&stmt_result),
604 Some(TransactionControl::Rollback)
605 ) {
606 let mut eng = match engine.write() {
607 Ok(eng) => eng,
608 Err(e) => {
609 return (
610 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
611 None,
612 )
613 }
614 };
615 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
616 }
617 execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, &bound))
618}
619
620#[derive(Debug, Clone)]
621struct SyncPullRequest {
622 replica_id: String,
623 since_lsn: u64,
624 max_units: u32,
625 max_bytes: u64,
626 database_id: [u8; 16],
627 primary_generation: u64,
628 wal_format_version: u16,
629 catalog_version: u16,
630 segment_format_version: u16,
631}
632
633#[derive(Debug, Clone, Copy, PartialEq, Eq)]
634enum SyncErrorClass {
635 AuthRequired,
636 PermissionDenied,
637 InvalidReplicaId,
638 ActiveTransaction,
639 QueryExecution,
640 InvalidMaxUnits,
641 InvalidMaxBytes,
642 SyncContext,
643 StatusRead,
644 CursorLsnMismatch,
645 IdentityRead,
646 IdentityOrFormatMismatch,
647 RetainedRead,
648 RetainedUnitEncoding,
649 RetainedChunkNotApplyable,
650 LsnAheadOfRemote,
651 AckValidation,
652 AckUpdate,
653 Internal,
654}
655
656impl SyncErrorClass {
657 const fn as_label(self) -> &'static str {
658 match self {
659 Self::AuthRequired => "auth_required",
660 Self::PermissionDenied => "permission_denied",
661 Self::InvalidReplicaId => "invalid_replica_id",
662 Self::ActiveTransaction => "active_transaction",
663 Self::QueryExecution => "query_execution",
664 Self::InvalidMaxUnits => "invalid_max_units",
665 Self::InvalidMaxBytes => "invalid_max_bytes",
666 Self::SyncContext => "sync_context",
667 Self::StatusRead => "status_read",
668 Self::CursorLsnMismatch => "cursor_lsn_mismatch",
669 Self::IdentityRead => "identity_read",
670 Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
671 Self::RetainedRead => "retained_read",
672 Self::RetainedUnitEncoding => "retained_unit_encoding",
673 Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
674 Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
675 Self::AckValidation => "ack_validation",
676 Self::AckUpdate => "ack_update",
677 Self::Internal => "internal",
678 }
679 }
680}
681
682#[derive(Debug, Clone)]
683struct SyncDecision {
684 message: Message,
685 error_class: Option<SyncErrorClass>,
686}
687
688impl SyncDecision {
689 fn ok(message: Message) -> Self {
690 Self {
691 message,
692 error_class: None,
693 }
694 }
695
696 fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
697 Self {
698 message: Message::Error {
699 message: message.into(),
700 },
701 error_class: Some(class),
702 }
703 }
704}
705
706fn check_sync_protocol_permitted(
707 credential_authenticated: bool,
708 principal: Option<&Principal>,
709) -> Result<(), (SyncErrorClass, String)> {
710 if !credential_authenticated {
711 return Err((
712 SyncErrorClass::AuthRequired,
713 "sync protocol requires authentication".to_string(),
714 ));
715 }
716 if let Some(principal) = principal {
717 let allowed =
718 Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
719 if !allowed {
720 return Err((
721 SyncErrorClass::PermissionDenied,
722 format!(
723 "permission denied: role '{}' cannot use sync protocol",
724 principal.role
725 ),
726 ));
727 }
728 }
729 Ok(())
730}
731
732fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
733 if replica_id.is_empty() {
734 return Err("replica id must be non-empty".to_string());
735 }
736 if replica_id.len() > 128 {
737 return Err("replica id must be at most 128 bytes".to_string());
738 }
739 if !replica_id
740 .bytes()
741 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
742 {
743 return Err("replica id contains unsupported characters".to_string());
744 }
745 Ok(())
746}
747
748fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<(PathBuf, u64), String> {
749 let engine = engine
750 .read()
751 .map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
752 let catalog = engine.catalog();
753 Ok((catalog.data_dir().to_path_buf(), catalog.max_lsn()))
754}
755
756fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
757 match action {
758 SyncRepairAction::None => WireSyncRepairAction::None,
759 SyncRepairAction::Pull => WireSyncRepairAction::Pull,
760 SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
761 SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
762 }
763}
764
765fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
766 WireSyncStatus {
767 replica_id: status.replica_id,
768 active: status.active,
769 last_applied_lsn: status.last_applied_lsn,
770 remote_lsn: status.remote_lsn,
771 servable_lsn: status.servable_lsn,
772 unarchived_lsn: status.unarchived_lsn,
773 lag_lsn: status.lag_lsn,
774 lag_bytes: status.lag_bytes,
775 lag_ms: status.lag_ms,
776 stale: status.stale,
777 repair_action: wire_repair_action(status.repair_action),
778 last_sync_error: status.last_sync_error,
779 }
780}
781
782fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
783 WireRetainedUnit {
784 tx_id: unit.tx_id,
785 record_type: unit.record_type,
786 lsn: unit.lsn,
787 data: unit.data,
788 }
789}
790
791fn sync_operation_outcome(message: &Message) -> SyncOutcome {
792 match message {
793 Message::SyncStatusResult { .. }
794 | Message::SyncPullResult { .. }
795 | Message::SyncAckResult { .. } => SyncOutcome::Ok,
796 _ => SyncOutcome::Error,
797 }
798}
799
800fn sync_operation_label(operation: SyncOperation) -> &'static str {
801 match operation {
802 SyncOperation::Status => "status",
803 SyncOperation::Pull => "pull",
804 SyncOperation::Ack => "ack",
805 }
806}
807
808fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
809 units.iter().fold(0, |total, unit| {
810 total.saturating_add(unit.encoded_len().unwrap_or(0))
811 })
812}
813
814fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
815 match action {
816 WireSyncRepairAction::None => SyncRepairLabel::None,
817 WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
818 WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
819 WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
820 }
821}
822
823fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
824 match action {
825 WireSyncRepairAction::None => "none",
826 WireSyncRepairAction::Pull => "pull",
827 WireSyncRepairAction::AwaitArchive => "await_archive",
828 WireSyncRepairAction::Rebootstrap => "rebootstrap",
829 }
830}
831
832const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
833const FNV1A64_PRIME: u64 = 0x100000001b3;
834const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
835
836fn replica_fingerprint(replica_id: &str) -> String {
837 let mut hash = FNV1A64_OFFSET;
838 for byte in replica_id.as_bytes() {
839 hash ^= u64::from(*byte);
840 hash = hash.wrapping_mul(FNV1A64_PRIME);
841 }
842 format!("{hash:016x}")
843}
844
845fn log_replica_fingerprint(replica_id: &str) -> String {
846 if validate_wire_replica_id(replica_id).is_ok() {
847 replica_fingerprint(replica_id)
848 } else {
849 INVALID_REPLICA_FINGERPRINT.to_string()
850 }
851}
852
853#[derive(Debug, Clone)]
854struct SyncLogContext {
855 replica_fingerprint: String,
856 since_lsn: Option<u64>,
857 applied_lsn: Option<u64>,
858 observed_remote_lsn: Option<u64>,
859 max_units: Option<u32>,
860 max_bytes: Option<u64>,
861}
862
863impl SyncLogContext {
864 fn status(replica_id: &str) -> Self {
865 Self::base(replica_id)
866 }
867
868 fn pull(request: &SyncPullRequest) -> Self {
869 Self {
870 replica_fingerprint: log_replica_fingerprint(&request.replica_id),
871 since_lsn: Some(request.since_lsn),
872 applied_lsn: None,
873 observed_remote_lsn: None,
874 max_units: Some(request.max_units),
875 max_bytes: Some(request.max_bytes),
876 }
877 }
878
879 fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
880 Self {
881 replica_fingerprint: log_replica_fingerprint(replica_id),
882 since_lsn: None,
883 applied_lsn: Some(applied_lsn),
884 observed_remote_lsn: Some(observed_remote_lsn),
885 max_units: None,
886 max_bytes: None,
887 }
888 }
889
890 fn base(replica_id: &str) -> Self {
891 Self {
892 replica_fingerprint: log_replica_fingerprint(replica_id),
893 since_lsn: None,
894 applied_lsn: None,
895 observed_remote_lsn: None,
896 max_units: None,
897 max_bytes: None,
898 }
899 }
900}
901
902struct SyncExecutionContext<'a> {
903 tx_gate: TxGate,
904 connection_has_transaction: bool,
905 operation: SyncOperation,
906 log_context: SyncLogContext,
907 metrics: &'a Arc<Metrics>,
908 query_timeout: Duration,
909}
910
911fn log_sync_decision(
912 operation: SyncOperation,
913 context: &SyncLogContext,
914 elapsed: Duration,
915 decision: &SyncDecision,
916) {
917 let operation = sync_operation_label(operation);
918 let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
919 match &decision.message {
920 Message::SyncStatusResult { status } => {
921 let repair_action = wire_repair_action_label(status.repair_action);
922 if status.stale || status.repair_action != WireSyncRepairAction::None {
923 info!(
924 operation = operation,
925 replica_fingerprint = %context.replica_fingerprint,
926 remote_lsn = status.remote_lsn,
927 last_applied_lsn = ?status.last_applied_lsn,
928 servable_lsn = ?status.servable_lsn,
929 unarchived_lsn = ?status.unarchived_lsn,
930 lag_lsn = ?status.lag_lsn,
931 lag_bytes = ?status.lag_bytes,
932 lag_ms = ?status.lag_ms,
933 stale = status.stale,
934 repair_action,
935 elapsed_ms,
936 "sync decision"
937 );
938 } else {
939 debug!(
940 operation = operation,
941 replica_fingerprint = %context.replica_fingerprint,
942 remote_lsn = status.remote_lsn,
943 last_applied_lsn = ?status.last_applied_lsn,
944 repair_action,
945 elapsed_ms,
946 "sync decision"
947 );
948 }
949 }
950 Message::SyncPullResult {
951 status,
952 units,
953 has_more,
954 } => {
955 let repair_action = wire_repair_action_label(status.repair_action);
956 let units_len = units.len();
957 let payload_bytes = sync_pull_payload_bytes(units);
958 if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
959 info!(
960 operation = operation,
961 replica_fingerprint = %context.replica_fingerprint,
962 since_lsn = ?context.since_lsn,
963 max_units = ?context.max_units,
964 max_bytes = ?context.max_bytes,
965 units = units_len,
966 payload_bytes,
967 has_more = *has_more,
968 remote_lsn = status.remote_lsn,
969 last_applied_lsn = ?status.last_applied_lsn,
970 servable_lsn = ?status.servable_lsn,
971 unarchived_lsn = ?status.unarchived_lsn,
972 lag_lsn = ?status.lag_lsn,
973 stale = status.stale,
974 repair_action,
975 elapsed_ms,
976 "sync decision"
977 );
978 } else {
979 debug!(
980 operation = operation,
981 replica_fingerprint = %context.replica_fingerprint,
982 since_lsn = ?context.since_lsn,
983 units = units_len,
984 payload_bytes,
985 has_more = *has_more,
986 remote_lsn = status.remote_lsn,
987 repair_action,
988 elapsed_ms,
989 "sync decision"
990 );
991 }
992 }
993 Message::SyncAckResult {
994 previous_applied_lsn,
995 applied_lsn,
996 remote_lsn,
997 advanced,
998 status,
999 } => {
1000 let repair_action = wire_repair_action_label(status.repair_action);
1001 if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
1002 info!(
1003 operation = operation,
1004 replica_fingerprint = %context.replica_fingerprint,
1005 requested_applied_lsn = ?context.applied_lsn,
1006 observed_remote_lsn = ?context.observed_remote_lsn,
1007 previous_applied_lsn = *previous_applied_lsn,
1008 applied_lsn = *applied_lsn,
1009 remote_lsn = *remote_lsn,
1010 advanced = *advanced,
1011 stale = status.stale,
1012 repair_action,
1013 elapsed_ms,
1014 "sync decision"
1015 );
1016 } else {
1017 debug!(
1018 operation = operation,
1019 replica_fingerprint = %context.replica_fingerprint,
1020 requested_applied_lsn = ?context.applied_lsn,
1021 previous_applied_lsn = *previous_applied_lsn,
1022 applied_lsn = *applied_lsn,
1023 remote_lsn = *remote_lsn,
1024 advanced = *advanced,
1025 repair_action,
1026 elapsed_ms,
1027 "sync decision"
1028 );
1029 }
1030 }
1031 Message::Error { .. } => {
1032 warn!(
1033 operation = operation,
1034 replica_fingerprint = %context.replica_fingerprint,
1035 since_lsn = ?context.since_lsn,
1036 applied_lsn = ?context.applied_lsn,
1037 observed_remote_lsn = ?context.observed_remote_lsn,
1038 max_units = ?context.max_units,
1039 max_bytes = ?context.max_bytes,
1040 error_class = decision
1041 .error_class
1042 .unwrap_or(SyncErrorClass::Internal)
1043 .as_label(),
1044 elapsed_ms,
1045 "sync decision rejected"
1046 );
1047 }
1048 _ => {
1049 debug!(
1050 operation = operation,
1051 replica_fingerprint = %context.replica_fingerprint,
1052 elapsed_ms,
1053 "unexpected sync decision response"
1054 );
1055 }
1056 }
1057}
1058
1059fn trim_to_applyable_v1_prefix(
1060 raw_units: &mut Vec<RetainedUnit>,
1061 wire_units: &mut Vec<WireRetainedUnit>,
1062) -> Result<(), String> {
1063 let mut last_error = None;
1064 while !raw_units.is_empty() {
1065 match validate_v1_retained_units_applyable(raw_units) {
1066 Ok(()) => return Ok(()),
1067 Err(err) => {
1068 last_error = Some(err.to_string());
1069 raw_units.pop();
1070 wire_units.pop();
1071 }
1072 }
1073 }
1074 if let Some(error) = last_error {
1075 return Err(format!(
1076 "sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
1077 ));
1078 }
1079 Ok(())
1080}
1081
1082fn validate_sync_ack_applyable_boundary(
1083 data_dir: &Path,
1084 replica_id: &str,
1085 applied_lsn: u64,
1086 remote_lsn: u64,
1087) -> Result<(), String> {
1088 let status =
1089 replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
1090 let Some(previous_lsn) = status.last_applied_lsn else {
1091 return Ok(());
1092 };
1093 if applied_lsn <= previous_lsn {
1094 return Ok(());
1095 }
1096 let range_len = applied_lsn - previous_lsn;
1097 if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
1098 return Err(format!(
1099 "sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
1100 ));
1101 }
1102 let max_units =
1103 usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
1104 let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
1105 let segment_dir = retained_segments_dir(data_dir);
1106 let units = read_units_through(
1107 &segment_dir,
1108 identity.segment_identity(),
1109 previous_lsn,
1110 applied_lsn,
1111 max_units,
1112 )
1113 .map_err(|err| err.to_string())?;
1114 if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
1115 return Err(
1116 "sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
1117 );
1118 }
1119 validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
1120}
1121
1122#[cfg(test)]
1123fn dispatch_sync_status(
1124 engine: &Arc<RwLock<Engine>>,
1125 replica_id: String,
1126 credential_authenticated: bool,
1127 principal: Option<&Principal>,
1128) -> Message {
1129 dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
1130}
1131
1132fn dispatch_sync_status_decision(
1133 engine: &Arc<RwLock<Engine>>,
1134 replica_id: String,
1135 credential_authenticated: bool,
1136 principal: Option<&Principal>,
1137) -> SyncDecision {
1138 if let Err((class, message)) =
1139 check_sync_protocol_permitted(credential_authenticated, principal)
1140 {
1141 return SyncDecision::error(class, message);
1142 }
1143 if let Err(message) = validate_wire_replica_id(&replica_id) {
1144 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1145 }
1146 let (data_dir, remote_lsn) = match sync_context(engine) {
1147 Ok(context) => context,
1148 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1149 };
1150 match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
1151 Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
1152 status: wire_sync_status(status),
1153 }),
1154 Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
1155 }
1156}
1157
1158#[cfg(test)]
1159fn dispatch_sync_pull(
1160 engine: &Arc<RwLock<Engine>>,
1161 request: SyncPullRequest,
1162 credential_authenticated: bool,
1163 principal: Option<&Principal>,
1164) -> Message {
1165 dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
1166}
1167
1168fn dispatch_sync_pull_decision(
1169 engine: &Arc<RwLock<Engine>>,
1170 request: SyncPullRequest,
1171 credential_authenticated: bool,
1172 principal: Option<&Principal>,
1173) -> SyncDecision {
1174 if let Err((class, message)) =
1175 check_sync_protocol_permitted(credential_authenticated, principal)
1176 {
1177 return SyncDecision::error(class, message);
1178 }
1179 if let Err(message) = validate_wire_replica_id(&request.replica_id) {
1180 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1181 }
1182 if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
1183 return SyncDecision::error(
1184 SyncErrorClass::InvalidMaxUnits,
1185 format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
1186 );
1187 }
1188 if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
1189 return SyncDecision::error(
1190 SyncErrorClass::InvalidMaxBytes,
1191 format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
1192 );
1193 }
1194
1195 let (data_dir, remote_lsn) = match sync_context(engine) {
1196 Ok(context) => context,
1197 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1198 };
1199 let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
1200 Ok(status) => status,
1201 Err(err) => {
1202 return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
1203 }
1204 };
1205 let Some(cursor_lsn) = status.last_applied_lsn else {
1206 return SyncDecision::ok(Message::SyncPullResult {
1207 status: wire_sync_status(status),
1208 units: Vec::new(),
1209 has_more: false,
1210 });
1211 };
1212 if status.repair_action != SyncRepairAction::Pull {
1213 return SyncDecision::ok(Message::SyncPullResult {
1214 status: wire_sync_status(status),
1215 units: Vec::new(),
1216 has_more: false,
1217 });
1218 }
1219 if request.since_lsn != cursor_lsn {
1220 return SyncDecision::error(
1221 SyncErrorClass::CursorLsnMismatch,
1222 format!(
1223 "sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
1224 request.since_lsn
1225 ),
1226 );
1227 }
1228
1229 let identity = match read_identity(&data_dir) {
1230 Ok(identity) => identity,
1231 Err(err) => {
1232 return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
1233 }
1234 };
1235 let expected = identity.segment_identity();
1236 if request.database_id != expected.database_id
1237 || request.primary_generation != expected.primary_generation
1238 || request.wal_format_version != expected.wal_format_version
1239 || request.catalog_version != expected.catalog_version
1240 || request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
1241 {
1242 return SyncDecision::error(
1243 SyncErrorClass::IdentityOrFormatMismatch,
1244 "sync pull identity or format version mismatch; rebootstrap required",
1245 );
1246 }
1247
1248 let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
1249 let requested_through_lsn = request
1250 .since_lsn
1251 .saturating_add(request.max_units as u64)
1252 .min(remote_lsn)
1253 .min(status.servable_lsn.unwrap_or(request.since_lsn));
1254 let segment_dir = retained_segments_dir(&data_dir);
1255 if requested_through_lsn > request.since_lsn {
1256 if let Err(err) = validate_retained_tail_available(
1257 &segment_dir,
1258 expected,
1259 request.since_lsn,
1260 requested_through_lsn,
1261 ) {
1262 let mut rebootstrap_status = status;
1263 rebootstrap_status.stale = true;
1264 rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
1265 rebootstrap_status.last_sync_error = Some(format!(
1266 "retained history is unavailable; rebootstrap required: {err}"
1267 ));
1268 return SyncDecision::ok(Message::SyncPullResult {
1269 status: wire_sync_status(rebootstrap_status),
1270 units: Vec::new(),
1271 has_more: false,
1272 });
1273 }
1274 }
1275
1276 let raw_units = match read_units_through(
1277 &segment_dir,
1278 expected,
1279 request.since_lsn,
1280 requested_through_lsn,
1281 effective_max_units,
1282 ) {
1283 Ok(units) => units,
1284 Err(err) => {
1285 return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
1286 }
1287 };
1288
1289 let mut selected_raw = Vec::new();
1290 let mut selected = Vec::new();
1291 let mut selected_bytes = 0u64;
1292 for unit in raw_units {
1293 let wire_unit = wire_retained_unit(unit.clone());
1294 let unit_bytes = match wire_unit.encoded_len() {
1295 Ok(bytes) => bytes,
1296 Err(message) => {
1297 return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
1298 }
1299 };
1300 if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
1301 if selected.is_empty() {
1302 return SyncDecision::error(
1303 SyncErrorClass::InvalidMaxBytes,
1304 "sync pull maxBytes is too small for the next retained unit",
1305 );
1306 }
1307 break;
1308 }
1309 selected_bytes += unit_bytes;
1310 selected_raw.push(unit);
1311 selected.push(wire_unit);
1312 }
1313 if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
1314 return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
1315 }
1316
1317 let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
1318 let has_more = selected
1319 .last()
1320 .is_some_and(|unit| unit.lsn < fetchable_through_lsn);
1321 SyncDecision::ok(Message::SyncPullResult {
1322 status: wire_sync_status(status),
1323 units: selected,
1324 has_more,
1325 })
1326}
1327
1328#[cfg(test)]
1329fn dispatch_sync_ack(
1330 engine: &Arc<RwLock<Engine>>,
1331 replica_id: String,
1332 applied_lsn: u64,
1333 observed_remote_lsn: u64,
1334 credential_authenticated: bool,
1335 principal: Option<&Principal>,
1336) -> Message {
1337 dispatch_sync_ack_decision(
1338 engine,
1339 replica_id,
1340 applied_lsn,
1341 observed_remote_lsn,
1342 credential_authenticated,
1343 principal,
1344 )
1345 .message
1346}
1347
1348fn dispatch_sync_ack_decision(
1349 engine: &Arc<RwLock<Engine>>,
1350 replica_id: String,
1351 applied_lsn: u64,
1352 observed_remote_lsn: u64,
1353 credential_authenticated: bool,
1354 principal: Option<&Principal>,
1355) -> SyncDecision {
1356 if let Err((class, message)) =
1357 check_sync_protocol_permitted(credential_authenticated, principal)
1358 {
1359 return SyncDecision::error(class, message);
1360 }
1361 if let Err(message) = validate_wire_replica_id(&replica_id) {
1362 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1363 }
1364 if applied_lsn > observed_remote_lsn {
1365 return SyncDecision::error(
1366 SyncErrorClass::LsnAheadOfRemote,
1367 format!(
1368 "sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
1369 ),
1370 );
1371 }
1372 let (data_dir, remote_lsn) = match sync_context(engine) {
1373 Ok(context) => context,
1374 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1375 };
1376 if observed_remote_lsn > remote_lsn {
1377 return SyncDecision::error(
1378 SyncErrorClass::LsnAheadOfRemote,
1379 format!(
1380 "sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
1381 ),
1382 );
1383 }
1384 if let Err(message) =
1385 validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
1386 {
1387 return SyncDecision::error(SyncErrorClass::AckValidation, message);
1388 }
1389 match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
1390 Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
1391 previous_applied_lsn: summary.previous_applied_lsn,
1392 applied_lsn: summary.applied_lsn,
1393 remote_lsn: summary.remote_lsn,
1394 advanced: summary.advanced,
1395 status: wire_sync_status(summary.status),
1396 }),
1397 Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
1398 }
1399}
1400
1401async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
1402where
1403 T: Send + 'static,
1404 F: FnOnce(T) -> SyncDecision + Send + 'static,
1405{
1406 let mut handle = tokio::task::spawn_blocking(move || f(input));
1407 tokio::select! {
1408 result = &mut handle => match result {
1409 Ok(decision) => decision,
1410 Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1411 },
1412 _ = tokio::time::sleep(query_timeout) => match handle.await {
1413 Ok(decision) => decision,
1414 Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1415 },
1416 }
1417}
1418
1419async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
1420where
1421 T: Send + 'static,
1422 F: FnOnce(T) -> SyncDecision + Send + 'static,
1423{
1424 let SyncExecutionContext {
1425 tx_gate,
1426 connection_has_transaction,
1427 operation,
1428 log_context,
1429 metrics,
1430 query_timeout,
1431 } = context;
1432 let start = Instant::now();
1433 if connection_has_transaction {
1434 let decision = SyncDecision::error(
1435 SyncErrorClass::ActiveTransaction,
1436 "sync protocol is unavailable inside an active transaction",
1437 );
1438 let elapsed = start.elapsed();
1439 log_sync_decision(operation, &log_context, elapsed, &decision);
1440 metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1441 return decision.message;
1442 }
1443
1444 let permit = match tx_gate.acquire_owned().await {
1445 Ok(permit) => permit,
1446 Err(_) => {
1447 let decision =
1448 SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
1449 let elapsed = start.elapsed();
1450 log_sync_decision(operation, &log_context, elapsed, &decision);
1451 metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1452 return decision.message;
1453 }
1454 };
1455 let decision = run_blocking_sync(input, query_timeout, f).await;
1456 drop(permit);
1457 match &decision.message {
1458 Message::SyncStatusResult { status } => {
1459 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1460 }
1461 Message::SyncPullResult { status, units, .. } => {
1462 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1463 metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
1464 }
1465 Message::SyncAckResult {
1466 advanced, status, ..
1467 } => {
1468 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1469 if *advanced {
1470 metrics.inc_sync_ack_advanced();
1471 }
1472 }
1473 _ => {}
1474 }
1475 let elapsed = start.elapsed();
1476 log_sync_decision(operation, &log_context, elapsed, &decision);
1477 metrics.record_sync_operation(
1478 operation,
1479 elapsed,
1480 sync_operation_outcome(&decision.message),
1481 );
1482 decision.message
1483}
1484
1485async fn execute_wire_query(
1491 engine: Arc<RwLock<Engine>>,
1492 tx_gate: TxGate,
1493 tx_permit: &mut Option<OwnedSemaphorePermit>,
1494 query: String,
1495 principal: Option<Principal>,
1496 query_timeout: Duration,
1497 metrics: &Arc<Metrics>,
1498) -> (Message, Option<PendingDurability>) {
1499 match classify_query_transaction_control(&query) {
1500 Some(TransactionControl::Begin) => {
1501 if tx_permit.is_some() {
1502 return (
1503 Message::Error {
1504 message: sanitize_error("transaction already active"),
1505 },
1506 None,
1507 );
1508 }
1509 let permit = match tx_gate.acquire_owned().await {
1510 Ok(permit) => permit,
1511 Err(_) => {
1512 return (
1513 Message::Error {
1514 message: "query execution error".into(),
1515 },
1516 None,
1517 )
1518 }
1519 };
1520 let (response, ticket) = run_blocking_query(
1521 engine,
1522 query,
1523 principal,
1524 query_timeout,
1525 metrics,
1526 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1527 )
1528 .await;
1529 if is_success_response(&response) {
1530 *tx_permit = Some(permit);
1531 }
1532 (response, ticket)
1533 }
1534 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1535 let (response, ticket) = run_blocking_query(
1536 engine,
1537 query,
1538 principal,
1539 query_timeout,
1540 metrics,
1541 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1542 )
1543 .await;
1544 if is_success_response(&response) {
1545 tx_permit.take();
1550 }
1551 (response, ticket)
1552 }
1553 None if tx_permit.is_some() => {
1554 run_blocking_query(
1555 engine,
1556 query,
1557 principal,
1558 query_timeout,
1559 metrics,
1560 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1561 )
1562 .await
1563 }
1564 None => {
1565 let permit = match tx_gate.acquire_owned().await {
1566 Ok(permit) => permit,
1567 Err(_) => {
1568 return (
1569 Message::Error {
1570 message: "query execution error".into(),
1571 },
1572 None,
1573 )
1574 }
1575 };
1576 let out = run_blocking_query(
1577 engine,
1578 query,
1579 principal,
1580 query_timeout,
1581 metrics,
1582 |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1583 )
1584 .await;
1585 drop(permit);
1586 out
1587 }
1588 }
1589}
1590
1591async fn execute_wire_query_sql(
1592 engine: Arc<RwLock<Engine>>,
1593 tx_gate: TxGate,
1594 tx_permit: &mut Option<OwnedSemaphorePermit>,
1595 query: String,
1596 principal: Option<Principal>,
1597 query_timeout: Duration,
1598 metrics: &Arc<Metrics>,
1599) -> (Message, Option<PendingDurability>) {
1600 match classify_sql_transaction_control(&query) {
1601 Some(TransactionControl::Begin) => {
1602 if tx_permit.is_some() {
1603 return (
1604 Message::Error {
1605 message: sanitize_error("transaction already active"),
1606 },
1607 None,
1608 );
1609 }
1610 let permit = match tx_gate.acquire_owned().await {
1611 Ok(permit) => permit,
1612 Err(_) => {
1613 return (
1614 Message::Error {
1615 message: "query execution error".into(),
1616 },
1617 None,
1618 )
1619 }
1620 };
1621 let (response, ticket) = run_blocking_query(
1622 engine,
1623 query,
1624 principal,
1625 query_timeout,
1626 metrics,
1627 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1628 )
1629 .await;
1630 if is_success_response(&response) {
1631 *tx_permit = Some(permit);
1632 }
1633 (response, ticket)
1634 }
1635 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1636 let (response, ticket) = run_blocking_query(
1637 engine,
1638 query,
1639 principal,
1640 query_timeout,
1641 metrics,
1642 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1643 )
1644 .await;
1645 if is_success_response(&response) {
1646 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_sql_query(&engine, &query, principal.as_ref()),
1660 )
1661 .await
1662 }
1663 None => {
1664 let permit = match tx_gate.acquire_owned().await {
1665 Ok(permit) => permit,
1666 Err(_) => {
1667 return (
1668 Message::Error {
1669 message: "query execution error".into(),
1670 },
1671 None,
1672 )
1673 }
1674 };
1675 let out = run_blocking_query(
1676 engine,
1677 query,
1678 principal,
1679 query_timeout,
1680 metrics,
1681 |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1682 )
1683 .await;
1684 drop(permit);
1685 out
1686 }
1687 }
1688}
1689
1690#[allow(clippy::too_many_arguments)]
1694async fn execute_wire_query_with_params(
1695 engine: Arc<RwLock<Engine>>,
1696 tx_gate: TxGate,
1697 tx_permit: &mut Option<OwnedSemaphorePermit>,
1698 query: String,
1699 params: Vec<WireParam>,
1700 principal: Option<Principal>,
1701 query_timeout: Duration,
1702 metrics: &Arc<Metrics>,
1703) -> (Message, Option<PendingDurability>) {
1704 match classify_params_transaction_control(&query, ¶ms) {
1705 Some(TransactionControl::Begin) => {
1706 if tx_permit.is_some() {
1707 return (
1708 Message::Error {
1709 message: sanitize_error("transaction already active"),
1710 },
1711 None,
1712 );
1713 }
1714 let permit = match tx_gate.acquire_owned().await {
1715 Ok(permit) => permit,
1716 Err(_) => {
1717 return (
1718 Message::Error {
1719 message: "query execution error".into(),
1720 },
1721 None,
1722 )
1723 }
1724 };
1725 let (response, ticket) = run_blocking_query(
1726 engine,
1727 (query, params),
1728 principal,
1729 query_timeout,
1730 metrics,
1731 |engine, (query, params), principal| {
1732 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1733 },
1734 )
1735 .await;
1736 if is_success_response(&response) {
1737 *tx_permit = Some(permit);
1738 }
1739 (response, ticket)
1740 }
1741 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1742 let (response, ticket) = run_blocking_query(
1743 engine,
1744 (query, params),
1745 principal,
1746 query_timeout,
1747 metrics,
1748 |engine, (query, params), principal| {
1749 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1750 },
1751 )
1752 .await;
1753 if is_success_response(&response) {
1754 tx_permit.take();
1757 }
1758 (response, ticket)
1759 }
1760 None if tx_permit.is_some() => {
1761 run_blocking_query(
1762 engine,
1763 (query, params),
1764 principal,
1765 query_timeout,
1766 metrics,
1767 |engine, (query, params), principal| {
1768 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1769 },
1770 )
1771 .await
1772 }
1773 None => {
1774 let permit = match tx_gate.acquire_owned().await {
1775 Ok(permit) => permit,
1776 Err(_) => {
1777 return (
1778 Message::Error {
1779 message: "query execution error".into(),
1780 },
1781 None,
1782 )
1783 }
1784 };
1785 let out = run_blocking_query(
1786 engine,
1787 (query, params),
1788 principal,
1789 query_timeout,
1790 metrics,
1791 |engine, (query, params), principal| {
1792 dispatch_query_with_params(&engine, &query, ¶ms, principal.as_ref())
1793 },
1794 )
1795 .await;
1796 drop(permit);
1797 out
1798 }
1799 }
1800}
1801
1802struct DeferredQueryMetric {
1807 start: Instant,
1808 outcome: QueryOutcome,
1809 exceeded_timeout: bool,
1810}
1811
1812type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
1814
1815async fn run_blocking_query<T, F>(
1816 engine: Arc<RwLock<Engine>>,
1817 input: T,
1818 principal: Option<Principal>,
1819 query_timeout: Duration,
1820 metrics: &Arc<Metrics>,
1821 f: F,
1822) -> (Message, Option<PendingDurability>)
1823where
1824 T: Send + 'static,
1825 F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
1826{
1827 let _in_flight = metrics.in_flight_guard();
1828 let start = Instant::now();
1829 let mut handle = tokio::task::spawn_blocking(move || f(engine, input, principal));
1830 let mut exceeded_timeout = false;
1831 let join_result = tokio::select! {
1832 result = &mut handle => result,
1833 _ = tokio::time::sleep(query_timeout) => {
1834 exceeded_timeout = true;
1835 handle.await
1840 }
1841 };
1842
1843 let (message, ticket, outcome) = match join_result {
1844 Ok((Ok(result), ticket)) => match query_result_to_message(result) {
1845 Ok(message) => (message, ticket, QueryOutcome::Ok),
1846 Err(e) => (
1847 Message::Error {
1848 message: sanitize_error(&e.to_string()),
1849 },
1850 ticket,
1851 QueryOutcome::Error,
1852 ),
1853 },
1854 Ok((Err(e), ticket)) => {
1855 let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
1856 QueryOutcome::MemoryLimit
1857 } else {
1858 QueryOutcome::Error
1859 };
1860 (
1861 Message::Error {
1862 message: sanitize_error(&e.to_string()),
1863 },
1864 ticket,
1865 outcome,
1866 )
1867 }
1868 Err(e) => (
1869 Message::Error {
1870 message: format!("internal error: {e}"),
1871 },
1872 None,
1873 QueryOutcome::Error,
1874 ),
1875 };
1876 match ticket {
1877 Some(ticket) => (
1882 message,
1883 Some((
1884 ticket,
1885 DeferredQueryMetric {
1886 start,
1887 outcome,
1888 exceeded_timeout,
1889 },
1890 )),
1891 ),
1892 None => {
1893 if exceeded_timeout {
1894 metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
1895 } else {
1896 metrics.record_query(start.elapsed(), outcome);
1897 }
1898 (message, None)
1899 }
1900 }
1901}
1902
1903async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
1911 match tokio::task::spawn_blocking(move || ticket.wait()).await {
1912 Ok(Ok(())) => None,
1913 Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
1914 Err(e) => Some(format!("internal error: {e}")),
1915 }
1916}
1917
1918fn is_success_response(msg: &Message) -> bool {
1919 matches!(
1920 msg,
1921 Message::ResultRows { .. }
1922 | Message::ResultScalar { .. }
1923 | Message::ResultOk { .. }
1924 | Message::ResultMessage { .. }
1925 )
1926}
1927
1928fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
1929 let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref());
1930 let _ = res;
1931 if let Some(ticket) = ticket {
1934 let _ = ticket.wait();
1935 }
1936}
1937
1938pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
1939where
1940 S: AsyncRead + AsyncWrite + Unpin,
1941{
1942 let ConnOpts {
1943 engine,
1944 tx_gate,
1945 expected_password,
1946 users,
1947 shutdown_rx,
1948 idle_timeout,
1949 query_timeout,
1950 rate_limiter,
1951 peer_addr,
1952 metrics,
1953 } = opts;
1954
1955 let peer = peer_addr
1956 .map(|a| a.to_string())
1957 .unwrap_or_else(|| "unknown".into());
1958 let peer_ip = peer_addr.map(|a| a.ip());
1959
1960 let (reader, writer) = tokio::io::split(stream);
1961 let mut reader = BufReader::new(reader);
1962 let mut writer = BufWriter::new(writer);
1963
1964 let connect_msg = loop {
1969 match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
1970 Ok(Ok(Some(Message::Ping))) => {
1971 debug!(peer = %peer, "pre-auth ping");
1972 if !write_msg(&mut writer, &Message::Pong).await {
1973 return;
1974 }
1975 continue;
1976 }
1977 Ok(Ok(Some(msg))) => break msg,
1978 Ok(Ok(None)) => {
1979 debug!(peer = %peer, "client closed before CONNECT");
1980 return;
1981 }
1982 Ok(Err(e)) => {
1983 error!(peer = %peer, error = %e, "error reading CONNECT");
1984 return;
1985 }
1986 Err(_) => {
1987 warn!(peer = %peer, "idle timeout waiting for CONNECT");
1988 return;
1989 }
1990 }
1991 };
1992
1993 let principal: Option<Principal>;
1996 let credential_auth_configured = !users.is_empty() || expected_password.is_some();
1997 match connect_msg {
1998 Message::Connect {
1999 db_name,
2000 password,
2001 username,
2002 } => {
2003 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2005 if is_rate_limited(limiter, ip) {
2006 warn!(peer = %peer, "rate limited: too many auth failures");
2007 let err = Message::Error {
2008 message: "too many auth failures, try again later".into(),
2009 };
2010 write_msg(&mut writer, &err).await;
2011 return;
2012 }
2013 }
2014
2015 let outcome = authenticate_connect(
2016 &users,
2017 expected_password.as_ref().map(|p| p.as_str()),
2018 username.as_deref(),
2019 password.as_ref().map(|p| p.as_str()),
2020 );
2021
2022 match outcome {
2023 AuthOutcome::Rejected => {
2024 warn!(peer = %peer, db = %db_name, "auth rejected");
2025 metrics.inc_auth_failure();
2026 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2028 record_auth_failure(limiter, ip);
2029 }
2030 let err = Message::Error {
2031 message: "authentication failed".into(),
2032 };
2033 write_msg(&mut writer, &err).await;
2034 return;
2035 }
2036 AuthOutcome::Authenticated {
2037 principal: auth_principal,
2038 } => {
2039 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2041 clear_auth_failures(limiter, ip);
2042 }
2043 match &auth_principal {
2044 Some(p) => {
2045 info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
2046 }
2047 None => {
2048 info!(peer = %peer, db = %db_name, "client connected");
2049 }
2050 }
2051 principal = auth_principal;
2052 }
2053 }
2054
2055 let ok = Message::ConnectOk {
2056 version: env!("CARGO_PKG_VERSION").into(),
2057 };
2058 if !write_msg(&mut writer, &ok).await {
2059 return;
2060 }
2061 }
2062 _ => {
2063 warn!(peer = %peer, "first message was not CONNECT");
2064 let err = Message::Error {
2065 message: "expected CONNECT".into(),
2066 };
2067 write_msg(&mut writer, &err).await;
2068 return;
2069 }
2070 }
2071
2072 let mut tx_permit: Option<OwnedSemaphorePermit> = None;
2073 let mut carry: Option<Message> = None;
2076
2077 'conn: loop {
2079 let msg = if let Some(m) = carry.take() {
2080 m
2081 } else {
2082 tokio::select! {
2083 result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
2085 match result {
2086 Ok(Ok(Some(msg))) => msg,
2087 Ok(Ok(None)) => break,
2088 Ok(Err(e)) => {
2089 error!(peer = %peer, error = %e, "read error");
2090 break;
2091 }
2092 Err(_) => {
2093 info!(peer = %peer, "idle timeout, closing connection");
2094 let err = Message::Error { message: "idle timeout".into() };
2095 write_msg(&mut writer, &err).await;
2096 break;
2097 }
2098 }
2099 }
2100 _ = shutdown_rx.changed() => {
2102 if *shutdown_rx.borrow() {
2103 info!(peer = %peer, "server shutting down, closing connection");
2104 let err = Message::Error { message: "server shutting down".into() };
2105 write_msg(&mut writer, &err).await;
2106 break;
2107 }
2108 continue;
2109 }
2110 }
2111 };
2112
2113 if matches!(
2119 msg,
2120 Message::Query { .. } | Message::QuerySql { .. } | Message::QueryWithParams { .. }
2121 ) {
2122 const MAX_PIPELINE_BATCH: usize = 128;
2125 const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
2130
2131 enum BatchFatal {
2134 Closed,
2135 ReadError,
2136 }
2137
2138 fn approx_response_bytes(msg: &Message) -> usize {
2142 match msg {
2143 Message::ResultRows { columns, rows } => {
2144 columns.iter().map(|c| c.len() + 4).sum::<usize>()
2145 + rows
2146 .iter()
2147 .map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
2148 .sum::<usize>()
2149 }
2150 Message::ResultScalar { value } => value.len(),
2151 Message::ResultMessage { message } | Message::Error { message } => {
2152 message.len()
2153 }
2154 _ => 16,
2155 }
2156 }
2157
2158 fn complete_frame_buffered(buf: &[u8]) -> bool {
2164 buf.len() >= 6 && {
2165 let payload_len =
2166 u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
2167 buf.len() - 6 >= payload_len
2168 }
2169 }
2170
2171 let mut responses: Vec<Message> = Vec::new();
2172 let mut response_bytes: usize = 0;
2173 let mut last_ticket: Option<WalDurabilityTicket> = None;
2174 let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
2175 let mut fatal: Option<BatchFatal> = None;
2176 let mut current = msg;
2177 loop {
2178 let (response, ticket) = match current {
2179 Message::Query { query } => {
2180 if query.len() > MAX_QUERY_LENGTH {
2181 (
2182 Message::Error {
2183 message: format!(
2184 "query too large: {} bytes (max {})",
2185 query.len(),
2186 MAX_QUERY_LENGTH
2187 ),
2188 },
2189 None,
2190 )
2191 } else {
2192 debug!(peer = %peer, query = %query, "received query");
2193 execute_wire_query(
2194 engine.clone(),
2195 tx_gate.clone(),
2196 &mut tx_permit,
2197 query,
2198 principal.clone(),
2199 query_timeout,
2200 &metrics,
2201 )
2202 .await
2203 }
2204 }
2205 Message::QuerySql { query } => {
2206 if query.len() > MAX_QUERY_LENGTH {
2207 (
2208 Message::Error {
2209 message: format!(
2210 "query too large: {} bytes (max {})",
2211 query.len(),
2212 MAX_QUERY_LENGTH
2213 ),
2214 },
2215 None,
2216 )
2217 } else {
2218 debug!(peer = %peer, query = %query, "received SQL query");
2219 execute_wire_query_sql(
2220 engine.clone(),
2221 tx_gate.clone(),
2222 &mut tx_permit,
2223 query,
2224 principal.clone(),
2225 query_timeout,
2226 &metrics,
2227 )
2228 .await
2229 }
2230 }
2231 Message::QueryWithParams { query, params } => {
2232 if query.len() > MAX_QUERY_LENGTH {
2233 (
2234 Message::Error {
2235 message: format!(
2236 "query too large: {} bytes (max {})",
2237 query.len(),
2238 MAX_QUERY_LENGTH
2239 ),
2240 },
2241 None,
2242 )
2243 } else {
2244 debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
2245 execute_wire_query_with_params(
2246 engine.clone(),
2247 tx_gate.clone(),
2248 &mut tx_permit,
2249 query,
2250 params,
2251 principal.clone(),
2252 query_timeout,
2253 &metrics,
2254 )
2255 .await
2256 }
2257 }
2258 _ => unreachable!("batch loop only receives plain query frames"),
2259 };
2260 if let Some((t, m)) = ticket {
2261 last_ticket = Some(t);
2265 deferred_metrics.push(m);
2266 }
2267 response_bytes += approx_response_bytes(&response);
2268 responses.push(response);
2269
2270 if tx_permit.is_some()
2276 || responses.len() >= MAX_PIPELINE_BATCH
2277 || response_bytes >= MAX_PIPELINE_BATCH_BYTES
2278 || !complete_frame_buffered(reader.buffer())
2279 {
2280 break;
2281 }
2282 match tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)).await {
2285 Ok(Ok(Some(
2286 next @ (Message::Query { .. }
2287 | Message::QuerySql { .. }
2288 | Message::QueryWithParams { .. }),
2289 ))) => {
2290 if tx_gate.available_permits() == 0 {
2298 carry = Some(next);
2299 break;
2300 }
2301 current = next;
2302 }
2303 Ok(Ok(Some(other))) => {
2304 carry = Some(other);
2307 break;
2308 }
2309 Ok(Ok(None)) => {
2310 fatal = Some(BatchFatal::Closed);
2311 break;
2312 }
2313 Ok(Err(e)) => {
2314 error!(peer = %peer, error = %e, "read error");
2315 fatal = Some(BatchFatal::ReadError);
2316 break;
2317 }
2318 Err(_) => {
2319 error!(peer = %peer, "timeout decoding fully-buffered frame");
2322 fatal = Some(BatchFatal::ReadError);
2323 break;
2324 }
2325 }
2326 }
2327
2328 let mut durability_failed = false;
2332 if let Some(ticket) = last_ticket {
2333 if let Some(message) = settle_durability_ticket(ticket).await {
2334 durability_failed = true;
2337 for r in responses.iter_mut() {
2338 if is_success_response(r) {
2339 *r = Message::Error {
2340 message: message.clone(),
2341 };
2342 }
2343 }
2344 }
2345 }
2346 for m in deferred_metrics.drain(..) {
2347 let outcome = if m.exceeded_timeout {
2348 QueryOutcome::Timeout
2349 } else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
2350 QueryOutcome::Error
2351 } else {
2352 m.outcome
2353 };
2354 metrics.record_query(m.start.elapsed(), outcome);
2355 }
2356
2357 for r in &responses {
2358 if !write_msg(&mut writer, r).await {
2359 break 'conn;
2360 }
2361 }
2362 match fatal {
2363 None => continue,
2364 Some(BatchFatal::Closed | BatchFatal::ReadError) => break,
2365 }
2366 }
2367
2368 let response = match msg {
2369 Message::Ping => {
2370 debug!(peer = %peer, "ping");
2371 Message::Pong
2372 }
2373 Message::SyncStatus { replica_id } => {
2374 let engine = engine.clone();
2375 let principal = principal.clone();
2376 let log_context = SyncLogContext::status(&replica_id);
2377 execute_gated_sync(
2378 SyncExecutionContext {
2379 tx_gate: tx_gate.clone(),
2380 connection_has_transaction: tx_permit.is_some(),
2381 operation: SyncOperation::Status,
2382 log_context,
2383 metrics: &metrics,
2384 query_timeout,
2385 },
2386 (engine, replica_id, credential_auth_configured, principal),
2387 |(engine, replica_id, credential_authenticated, principal)| {
2388 dispatch_sync_status_decision(
2389 &engine,
2390 replica_id,
2391 credential_authenticated,
2392 principal.as_ref(),
2393 )
2394 },
2395 )
2396 .await
2397 }
2398 Message::SyncPull {
2399 replica_id,
2400 since_lsn,
2401 max_units,
2402 max_bytes,
2403 database_id,
2404 primary_generation,
2405 wal_format_version,
2406 catalog_version,
2407 segment_format_version,
2408 } => {
2409 let engine = engine.clone();
2410 let principal = principal.clone();
2411 let request = SyncPullRequest {
2412 replica_id,
2413 since_lsn,
2414 max_units,
2415 max_bytes,
2416 database_id,
2417 primary_generation,
2418 wal_format_version,
2419 catalog_version,
2420 segment_format_version,
2421 };
2422 let log_context = SyncLogContext::pull(&request);
2423 execute_gated_sync(
2424 SyncExecutionContext {
2425 tx_gate: tx_gate.clone(),
2426 connection_has_transaction: tx_permit.is_some(),
2427 operation: SyncOperation::Pull,
2428 log_context,
2429 metrics: &metrics,
2430 query_timeout,
2431 },
2432 (engine, request, credential_auth_configured, principal),
2433 |(engine, request, credential_authenticated, principal)| {
2434 dispatch_sync_pull_decision(
2435 &engine,
2436 request,
2437 credential_authenticated,
2438 principal.as_ref(),
2439 )
2440 },
2441 )
2442 .await
2443 }
2444 Message::SyncAck {
2445 replica_id,
2446 applied_lsn,
2447 remote_lsn,
2448 } => {
2449 let engine = engine.clone();
2450 let principal = principal.clone();
2451 let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
2452 execute_gated_sync(
2453 SyncExecutionContext {
2454 tx_gate: tx_gate.clone(),
2455 connection_has_transaction: tx_permit.is_some(),
2456 operation: SyncOperation::Ack,
2457 log_context,
2458 metrics: &metrics,
2459 query_timeout,
2460 },
2461 (
2462 engine,
2463 replica_id,
2464 applied_lsn,
2465 remote_lsn,
2466 credential_auth_configured,
2467 principal,
2468 ),
2469 |(
2470 engine,
2471 replica_id,
2472 applied_lsn,
2473 observed_remote_lsn,
2474 credential_authenticated,
2475 principal,
2476 )| {
2477 dispatch_sync_ack_decision(
2478 &engine,
2479 replica_id,
2480 applied_lsn,
2481 observed_remote_lsn,
2482 credential_authenticated,
2483 principal.as_ref(),
2484 )
2485 },
2486 )
2487 .await
2488 }
2489 Message::Disconnect => {
2490 debug!(peer = %peer, "received DISCONNECT");
2491 break;
2492 }
2493 _ => Message::Error {
2494 message: "unexpected message type".into(),
2495 },
2496 };
2497
2498 if !write_msg(&mut writer, &response).await {
2499 break;
2500 }
2501 }
2502
2503 if tx_permit.is_some() {
2510 let engine = engine.clone();
2511 let principal = principal.clone();
2512 let _ =
2513 tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
2514 }
2515 tx_permit.take();
2516
2517 info!(peer = %peer, "client disconnected");
2518}
2519
2520fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
2521 *total = total.saturating_add(bytes);
2522 if *total > MAX_RESPONSE_PAYLOAD_SIZE {
2523 return Err(QueryError::Execution(format!(
2524 "result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
2525 MAX_RESPONSE_PAYLOAD_SIZE
2526 )));
2527 }
2528 Ok(())
2529}
2530
2531fn query_result_to_message(result: QueryResult) -> Result<Message, QueryError> {
2532 match result {
2533 QueryResult::Rows { columns, rows } => {
2534 let mut encoded_bytes = 2usize; let mut out_columns = Vec::with_capacity(columns.len());
2536 for col in columns {
2537 charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
2538 out_columns.push(col);
2539 }
2540 charge_response_bytes(&mut encoded_bytes, 4)?; let mut str_rows = Vec::with_capacity(rows.len());
2543 for row in rows {
2544 let mut str_row = Vec::with_capacity(row.len());
2545 for value in row {
2546 let display = value_to_display(&value);
2547 charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
2548 str_row.push(display);
2549 }
2550 str_rows.push(str_row);
2551 }
2552 Ok(Message::ResultRows {
2553 columns: out_columns,
2554 rows: str_rows,
2555 })
2556 }
2557 QueryResult::Scalar(val) => Ok(Message::ResultScalar {
2558 value: value_to_display(&val),
2559 }),
2560 QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
2561 QueryResult::Created(name) => Ok(Message::ResultMessage {
2562 message: format!("type {name} created"),
2563 }),
2564 QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
2565 }
2566}
2567
2568fn value_to_display(v: &Value) -> String {
2572 v.to_wire_string()
2573}
2574
2575#[cfg(test)]
2576mod tests {
2577 use super::*;
2578 use powdb_storage::wal::WalRecordType;
2579 use powdb_sync::{
2580 write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
2581 ReplicaCursor, RetainedSegment, RetainedUnit,
2582 };
2583
2584 #[test]
2587 fn null_serializes_as_null_bareword_on_wire() {
2588 assert_eq!(value_to_display(&Value::Empty), "null");
2589 }
2590
2591 #[test]
2594 fn unique_violation_error_surfaces_to_remote_clients() {
2595 assert_eq!(
2598 sanitize_error("unique constraint violation on User.email"),
2599 "unique constraint violation on User.email"
2600 );
2601 }
2602
2603 #[test]
2604 fn internal_errors_stay_generic() {
2605 assert_eq!(
2606 sanitize_error("some internal io panic detail"),
2607 "query execution error"
2608 );
2609 }
2610
2611 #[test]
2612 fn resource_limit_errors_surface_actionable_hints() {
2613 for msg in [
2618 "sort input exceeds row limit — add a LIMIT clause",
2619 "join result exceeds row limit",
2620 "query exceeded memory budget: requested 100 bytes, limit 50 bytes",
2621 "result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
2622 ] {
2623 assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
2624 }
2625 }
2626
2627 #[test]
2628 fn oversized_result_is_rejected_before_wire_encoding() {
2629 let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
2630 let result = QueryResult::Rows {
2631 columns: vec!["payload".into()],
2632 rows: vec![vec![Value::Str(long)]],
2633 };
2634 let err = query_result_to_message(result).unwrap_err();
2635 assert!(
2636 err.to_string().starts_with("result too large"),
2637 "unexpected error: {err}"
2638 );
2639 }
2640
2641 fn parsed(q: &str) -> powdb_query::ast::Statement {
2644 parser::parse(q).unwrap()
2645 }
2646
2647 fn principal(role: &str) -> Option<Principal> {
2648 Some(Principal {
2649 name: "u".into(),
2650 role: role.into(),
2651 })
2652 }
2653
2654 #[test]
2655 fn readonly_can_read_but_not_write() {
2656 let p = principal("readonly");
2657 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2659 assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
2660 assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
2661 for q in [
2663 r#"insert User { name := "x" }"#,
2664 "User filter .id = 1 update { age := 2 }",
2665 "User filter .id = 1 delete",
2666 "drop User",
2667 "alter User add column c: str",
2668 "type T { required id: int }",
2669 "begin",
2670 "commit",
2671 "rollback",
2672 ] {
2673 let err = check_statement_permitted(p.as_ref(), &parsed(q))
2674 .expect_err(&format!("must deny: {q}"));
2675 assert!(
2676 err.to_string().contains("permission denied"),
2677 "unexpected error for {q}: {err}"
2678 );
2679 }
2680 }
2681
2682 #[test]
2683 fn readwrite_and_admin_have_full_query_access() {
2684 for role in ["readwrite", "admin"] {
2685 let p = principal(role);
2686 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2687 assert!(check_statement_permitted(
2688 p.as_ref(),
2689 &parsed(r#"insert User { name := "x" }"#)
2690 )
2691 .is_ok());
2692 assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
2693 }
2694 }
2695
2696 #[test]
2697 fn unknown_role_fails_closed_for_writes() {
2698 let p = principal("mystery");
2699 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2700 assert!(
2701 check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
2702 .is_err()
2703 );
2704 }
2705
2706 #[test]
2707 fn no_principal_means_full_access() {
2708 assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
2710 assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
2711 }
2712
2713 fn store_with_alice() -> UserStore {
2714 let mut s = UserStore::new();
2715 s.create_user("alice", "pw", "readwrite").unwrap();
2716 s
2717 }
2718
2719 #[test]
2722 fn empty_store_no_password_is_open() {
2723 let s = UserStore::new();
2724 assert_eq!(
2725 authenticate_connect(&s, None, None, None),
2726 AuthOutcome::Authenticated { principal: None }
2727 );
2728 assert_eq!(
2730 authenticate_connect(&s, None, Some("x"), Some("y")),
2731 AuthOutcome::Authenticated { principal: None }
2732 );
2733 }
2734
2735 #[test]
2736 fn empty_store_correct_shared_password_succeeds() {
2737 let s = UserStore::new();
2738 assert_eq!(
2739 authenticate_connect(&s, Some("pw"), None, Some("pw")),
2740 AuthOutcome::Authenticated { principal: None }
2741 );
2742 }
2743
2744 #[test]
2745 fn empty_store_wrong_shared_password_rejected() {
2746 let s = UserStore::new();
2747 assert_eq!(
2748 authenticate_connect(&s, Some("pw"), None, Some("bad")),
2749 AuthOutcome::Rejected
2750 );
2751 }
2752
2753 #[test]
2754 fn empty_store_missing_password_rejected_when_expected() {
2755 let s = UserStore::new();
2756 assert_eq!(
2757 authenticate_connect(&s, Some("pw"), None, None),
2758 AuthOutcome::Rejected
2759 );
2760 }
2761
2762 #[test]
2763 fn empty_store_ignores_username_for_shared_password() {
2764 let s = UserStore::new();
2767 assert_eq!(
2768 authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
2769 AuthOutcome::Authenticated { principal: None }
2770 );
2771 }
2772
2773 #[test]
2776 fn user_auth_success_binds_principal() {
2777 let s = store_with_alice();
2778 assert_eq!(
2779 authenticate_connect(&s, None, Some("alice"), Some("pw")),
2780 AuthOutcome::Authenticated {
2781 principal: Some(Principal {
2782 name: "alice".into(),
2783 role: "readwrite".into(),
2784 })
2785 }
2786 );
2787 }
2788
2789 #[test]
2790 fn user_auth_wrong_password_rejected() {
2791 let s = store_with_alice();
2792 assert_eq!(
2793 authenticate_connect(&s, None, Some("alice"), Some("bad")),
2794 AuthOutcome::Rejected
2795 );
2796 }
2797
2798 #[test]
2799 fn user_auth_unknown_user_rejected() {
2800 let s = store_with_alice();
2801 assert_eq!(
2802 authenticate_connect(&s, None, Some("mallory"), Some("pw")),
2803 AuthOutcome::Rejected
2804 );
2805 }
2806
2807 #[test]
2808 fn user_auth_missing_username_rejected() {
2809 let s = store_with_alice();
2810 assert_eq!(
2811 authenticate_connect(&s, None, None, Some("pw")),
2812 AuthOutcome::Rejected
2813 );
2814 }
2815
2816 #[test]
2817 fn user_auth_missing_password_rejected() {
2818 let s = store_with_alice();
2819 assert_eq!(
2820 authenticate_connect(&s, Some("pw"), Some("alice"), None),
2821 AuthOutcome::Rejected
2822 );
2823 }
2824
2825 #[test]
2826 fn user_auth_ignores_shared_password_when_users_present() {
2827 let s = store_with_alice();
2830 assert_eq!(
2831 authenticate_connect(&s, Some("shared"), None, Some("shared")),
2832 AuthOutcome::Rejected
2833 );
2834 }
2835
2836 #[test]
2837 fn replica_fingerprint_is_stable_and_redacted() {
2838 let replica_id = "customer-prod-replica-a";
2839 let fingerprint = replica_fingerprint(replica_id);
2840 assert_eq!(fingerprint, replica_fingerprint(replica_id));
2841 assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
2842 assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
2843 assert_eq!(fingerprint.len(), 16);
2844 assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
2845 assert!(!fingerprint.contains("customer"));
2846 assert!(!fingerprint.contains("replica"));
2847 assert!(!fingerprint.contains(replica_id));
2848 }
2849
2850 #[test]
2851 fn invalid_replica_ids_use_fixed_log_fingerprint() {
2852 assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
2853 assert_eq!(
2854 log_replica_fingerprint("customer/prod/replica"),
2855 INVALID_REPLICA_FINGERPRINT
2856 );
2857 assert_eq!(
2858 log_replica_fingerprint(&"a".repeat(4096)),
2859 INVALID_REPLICA_FINGERPRINT
2860 );
2861 }
2862
2863 #[test]
2864 fn sync_error_classes_use_bounded_labels() {
2865 assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
2866 assert_eq!(
2867 SyncErrorClass::PermissionDenied.as_label(),
2868 "permission_denied"
2869 );
2870 assert_eq!(
2871 SyncErrorClass::IdentityOrFormatMismatch.as_label(),
2872 "identity_or_format_mismatch"
2873 );
2874 assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
2875 assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
2876 }
2877
2878 fn sync_identity() -> DatabaseIdentity {
2879 DatabaseIdentity {
2880 database_id: *b"server-sync-test",
2881 primary_generation: 1,
2882 }
2883 }
2884
2885 fn retained_unit(lsn: u64) -> RetainedUnit {
2886 RetainedUnit {
2887 tx_id: 1,
2888 record_type: 4,
2889 lsn,
2890 data: lsn.to_le_bytes().to_vec(),
2891 }
2892 }
2893
2894 fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
2895 RetainedUnit {
2896 tx_id,
2897 record_type: record_type as u8,
2898 lsn,
2899 data: lsn.to_le_bytes().to_vec(),
2900 }
2901 }
2902
2903 fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
2904 let identity = sync_identity();
2905 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
2906 let units = (1..=through_lsn).map(retained_unit).collect();
2907 let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
2908 write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
2909 }
2910
2911 fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
2912 let identity = sync_identity();
2913 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
2914 let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
2915 write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
2916 }
2917
2918 fn write_sync_identity_only(data_dir: &std::path::Path) {
2919 let identity = sync_identity();
2920 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
2921 }
2922
2923 fn admin_principal() -> Principal {
2924 Principal {
2925 name: "admin".into(),
2926 role: "admin".into(),
2927 }
2928 }
2929
2930 #[test]
2931 fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
2932 let dir = tempfile::tempdir().unwrap();
2933 let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
2934
2935 match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
2936 Message::Error { message } => {
2937 assert!(message.contains("requires authentication"));
2938 }
2939 other => panic!("expected auth error, got {other:?}"),
2940 }
2941
2942 let readonly = Principal {
2943 name: "reader".into(),
2944 role: "readonly".into(),
2945 };
2946 match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
2947 Message::Error { message } => {
2948 assert!(message.contains("permission denied"));
2949 }
2950 other => panic!("expected permission error, got {other:?}"),
2951 }
2952 }
2953
2954 #[test]
2955 fn sync_status_pull_and_ack_use_server_remote_lsn() {
2956 let dir = tempfile::tempdir().unwrap();
2957 let mut engine = Engine::new(dir.path()).unwrap();
2958 engine
2959 .execute_powql("type SyncT { required id: int, v: str }")
2960 .unwrap();
2961 engine
2962 .execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
2963 .unwrap();
2964 let remote_lsn = engine.catalog().max_lsn();
2965 assert!(remote_lsn > 0);
2966 write_sync_identity_and_tail(dir.path(), remote_lsn);
2967 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
2968 .unwrap();
2969
2970 let engine = Arc::new(RwLock::new(engine));
2971 let principal = admin_principal();
2972 let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
2973 {
2974 Message::SyncStatusResult { status } => status,
2975 other => panic!("expected sync status, got {other:?}"),
2976 };
2977 assert_eq!(status.remote_lsn, remote_lsn);
2978 assert_eq!(status.servable_lsn, Some(remote_lsn));
2979 assert_eq!(status.unarchived_lsn, Some(0));
2980 assert_eq!(status.last_applied_lsn, Some(0));
2981 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
2982 assert!(status.stale);
2983
2984 let identity = sync_identity().segment_identity();
2985 let pull = SyncPullRequest {
2986 replica_id: "replica-a".into(),
2987 since_lsn: 0,
2988 max_units: MAX_SYNC_PULL_UNITS,
2989 max_bytes: MAX_SYNC_PULL_BYTES,
2990 database_id: identity.database_id,
2991 primary_generation: identity.primary_generation,
2992 wal_format_version: identity.wal_format_version,
2993 catalog_version: identity.catalog_version,
2994 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
2995 };
2996 let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
2997 Message::SyncPullResult {
2998 status,
2999 units,
3000 has_more,
3001 } => {
3002 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3003 assert!(!has_more);
3004 units
3005 }
3006 other => panic!("expected sync pull result, got {other:?}"),
3007 };
3008 assert_eq!(units.len() as u64, remote_lsn);
3009 assert_eq!(units.last().unwrap().lsn, remote_lsn);
3010
3011 let ack = match dispatch_sync_ack(
3012 &engine,
3013 "replica-a".into(),
3014 remote_lsn,
3015 remote_lsn,
3016 true,
3017 Some(&principal),
3018 ) {
3019 Message::SyncAckResult {
3020 previous_applied_lsn,
3021 applied_lsn,
3022 remote_lsn: ack_remote_lsn,
3023 advanced,
3024 status,
3025 } => {
3026 assert_eq!(previous_applied_lsn, 0);
3027 assert_eq!(applied_lsn, remote_lsn);
3028 assert_eq!(ack_remote_lsn, remote_lsn);
3029 assert!(advanced);
3030 status
3031 }
3032 other => panic!("expected sync ack result, got {other:?}"),
3033 };
3034 assert_eq!(ack.repair_action, WireSyncRepairAction::None);
3035 assert!(!ack.stale);
3036 assert_eq!(ack.lag_lsn, Some(0));
3037 }
3038
3039 #[test]
3040 fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
3041 let dir = tempfile::tempdir().unwrap();
3042 let mut engine = Engine::new(dir.path()).unwrap();
3043 engine
3044 .execute_powql("type SyncT { required id: int }")
3045 .unwrap();
3046 for id in 1..=3 {
3047 engine
3048 .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
3049 .unwrap();
3050 }
3051 let remote_lsn = engine.catalog().max_lsn();
3052 assert!(remote_lsn >= 3);
3053 write_sync_identity_and_units(
3054 dir.path(),
3055 vec![
3056 retained_unit_with(77, WalRecordType::Begin, 1),
3057 retained_unit_with(77, WalRecordType::Insert, 2),
3058 retained_unit_with(77, WalRecordType::Commit, 3),
3059 ],
3060 );
3061 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3062 .unwrap();
3063
3064 let engine = Arc::new(RwLock::new(engine));
3065 let principal = admin_principal();
3066 let identity = sync_identity().segment_identity();
3067 let cut_pull = SyncPullRequest {
3068 replica_id: "replica-a".into(),
3069 since_lsn: 0,
3070 max_units: 2,
3071 max_bytes: MAX_SYNC_PULL_BYTES,
3072 database_id: identity.database_id,
3073 primary_generation: identity.primary_generation,
3074 wal_format_version: identity.wal_format_version,
3075 catalog_version: identity.catalog_version,
3076 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3077 };
3078 match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
3079 Message::Error { message } => assert!(message.contains("cuts through transaction")),
3080 other => panic!("expected transaction-cut pull error, got {other:?}"),
3081 }
3082
3083 let cut_bytes_pull = SyncPullRequest {
3084 replica_id: "replica-a".into(),
3085 since_lsn: 0,
3086 max_units: 3,
3087 max_bytes: 58,
3088 database_id: identity.database_id,
3089 primary_generation: identity.primary_generation,
3090 wal_format_version: identity.wal_format_version,
3091 catalog_version: identity.catalog_version,
3092 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3093 };
3094 match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
3095 Message::Error { message } => assert!(message.contains("cuts through transaction")),
3096 other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
3097 }
3098
3099 let full_pull = SyncPullRequest {
3100 replica_id: "replica-a".into(),
3101 since_lsn: 0,
3102 max_units: 3,
3103 max_bytes: MAX_SYNC_PULL_BYTES,
3104 database_id: identity.database_id,
3105 primary_generation: identity.primary_generation,
3106 wal_format_version: identity.wal_format_version,
3107 catalog_version: identity.catalog_version,
3108 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3109 };
3110 match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
3111 Message::SyncPullResult { units, .. } => {
3112 assert_eq!(units.len(), 3);
3113 assert_eq!(units.last().unwrap().lsn, 3);
3114 }
3115 other => panic!("expected complete transaction pull, got {other:?}"),
3116 }
3117
3118 match dispatch_sync_ack(
3119 &engine,
3120 "replica-a".into(),
3121 2,
3122 remote_lsn,
3123 true,
3124 Some(&principal),
3125 ) {
3126 Message::Error { message } => assert!(message.contains("cuts through transaction")),
3127 other => panic!("expected transaction-cut ack error, got {other:?}"),
3128 }
3129 let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
3130 assert_eq!(cursor[0].applied_lsn, 0);
3131
3132 match dispatch_sync_ack(
3133 &engine,
3134 "replica-a".into(),
3135 3,
3136 remote_lsn,
3137 true,
3138 Some(&principal),
3139 ) {
3140 Message::SyncAckResult {
3141 previous_applied_lsn,
3142 applied_lsn,
3143 advanced,
3144 ..
3145 } => {
3146 assert_eq!(previous_applied_lsn, 0);
3147 assert_eq!(applied_lsn, 3);
3148 assert!(advanced);
3149 }
3150 other => panic!("expected complete transaction ack, got {other:?}"),
3151 }
3152 }
3153
3154 #[test]
3155 fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
3156 let dir = tempfile::tempdir().unwrap();
3157 let mut engine = Engine::new(dir.path()).unwrap();
3158 engine
3159 .execute_powql("type SyncT { required id: int }")
3160 .unwrap();
3161 for id in 1..=6 {
3162 engine
3163 .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
3164 .unwrap();
3165 }
3166 let remote_lsn = engine.catalog().max_lsn();
3167 assert!(remote_lsn >= 6);
3168 write_sync_identity_and_units(
3169 dir.path(),
3170 vec![
3171 retained_unit_with(1, WalRecordType::Begin, 1),
3172 retained_unit_with(1, WalRecordType::Insert, 2),
3173 retained_unit_with(1, WalRecordType::Commit, 3),
3174 retained_unit_with(1, WalRecordType::Begin, 4),
3175 retained_unit_with(1, WalRecordType::Insert, 5),
3176 retained_unit_with(1, WalRecordType::Commit, 6),
3177 ],
3178 );
3179 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3180 .unwrap();
3181
3182 let engine = Arc::new(RwLock::new(engine));
3183 let principal = admin_principal();
3184 let identity = sync_identity().segment_identity();
3185 let pull = SyncPullRequest {
3186 replica_id: "replica-a".into(),
3187 since_lsn: 0,
3188 max_units: 6,
3189 max_bytes: 100,
3190 database_id: identity.database_id,
3191 primary_generation: identity.primary_generation,
3192 wal_format_version: identity.wal_format_version,
3193 catalog_version: identity.catalog_version,
3194 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3195 };
3196
3197 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3198 Message::SyncPullResult {
3199 status,
3200 units,
3201 has_more,
3202 } => {
3203 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3204 assert_eq!(units.len(), 3);
3205 assert_eq!(units.last().unwrap().lsn, 3);
3206 assert!(has_more);
3207 }
3208 other => panic!("expected byte-capped applyable prefix, got {other:?}"),
3209 }
3210 }
3211
3212 #[test]
3213 fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
3214 let dir = tempfile::tempdir().unwrap();
3215 let mut engine = Engine::new(dir.path()).unwrap();
3216 engine
3217 .execute_powql("type SyncT { required id: int }")
3218 .unwrap();
3219 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3220 let remote_lsn = engine.catalog().max_lsn();
3221 assert!(remote_lsn > 0);
3222 write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
3223 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3224 .unwrap();
3225
3226 let engine = Arc::new(RwLock::new(engine));
3227 let principal = admin_principal();
3228 let identity = sync_identity().segment_identity();
3229 let pull = SyncPullRequest {
3230 replica_id: "replica-a".into(),
3231 since_lsn: 0,
3232 max_units: MAX_SYNC_PULL_UNITS,
3233 max_bytes: MAX_SYNC_PULL_BYTES,
3234 database_id: identity.database_id,
3235 primary_generation: identity.primary_generation,
3236 wal_format_version: identity.wal_format_version,
3237 catalog_version: identity.catalog_version,
3238 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3239 };
3240
3241 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3242 Message::SyncPullResult {
3243 status,
3244 units,
3245 has_more,
3246 } => {
3247 assert_eq!(status.remote_lsn, remote_lsn);
3248 assert_eq!(status.servable_lsn, Some(remote_lsn));
3249 assert_eq!(status.unarchived_lsn, Some(0));
3250 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3251 assert!(!has_more);
3252 assert_eq!(units.len() as u64, remote_lsn);
3253 assert_eq!(units.last().unwrap().lsn, remote_lsn);
3254 assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
3255 }
3256 other => panic!("expected capped sync pull result, got {other:?}"),
3257 }
3258 }
3259
3260 #[test]
3261 fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
3262 let dir = tempfile::tempdir().unwrap();
3263 let mut engine = Engine::new(dir.path()).unwrap();
3264 engine
3265 .execute_powql("type SyncT { required id: int }")
3266 .unwrap();
3267 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3268 let remote_lsn = engine.catalog().max_lsn();
3269 assert!(remote_lsn > 0);
3270 write_sync_identity_only(dir.path());
3271 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3272 .unwrap();
3273
3274 let engine = Arc::new(RwLock::new(engine));
3275 let principal = admin_principal();
3276 let identity = sync_identity().segment_identity();
3277 let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
3278 {
3279 Message::SyncStatusResult { status } => status,
3280 other => panic!("expected sync status, got {other:?}"),
3281 };
3282 assert_eq!(status.remote_lsn, remote_lsn);
3283 assert_eq!(status.servable_lsn, Some(0));
3284 assert_eq!(status.unarchived_lsn, Some(remote_lsn));
3285 assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
3286 assert!(status
3287 .last_sync_error
3288 .as_deref()
3289 .unwrap()
3290 .contains("not yet archived"));
3291
3292 let pull = SyncPullRequest {
3293 replica_id: "replica-a".into(),
3294 since_lsn: 0,
3295 max_units: MAX_SYNC_PULL_UNITS,
3296 max_bytes: MAX_SYNC_PULL_BYTES,
3297 database_id: identity.database_id,
3298 primary_generation: identity.primary_generation,
3299 wal_format_version: identity.wal_format_version,
3300 catalog_version: identity.catalog_version,
3301 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3302 };
3303 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3304 Message::SyncPullResult {
3305 status,
3306 units,
3307 has_more,
3308 } => {
3309 assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
3310 assert!(units.is_empty());
3311 assert!(!has_more);
3312 }
3313 other => panic!("expected await-archive sync pull result, got {other:?}"),
3314 }
3315 }
3316
3317 #[test]
3318 fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
3319 let dir = tempfile::tempdir().unwrap();
3320 let mut engine = Engine::new(dir.path()).unwrap();
3321 engine
3322 .execute_powql("type SyncT { required id: int }")
3323 .unwrap();
3324 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3325 engine.execute_powql("insert SyncT { id := 2 }").unwrap();
3326 let remote_lsn = engine.catalog().max_lsn();
3327 assert!(remote_lsn > 1);
3328 let servable_lsn = remote_lsn - 1;
3329 write_sync_identity_and_tail(dir.path(), servable_lsn);
3330 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3331 .unwrap();
3332
3333 let engine = Arc::new(RwLock::new(engine));
3334 let principal = admin_principal();
3335 let identity = sync_identity().segment_identity();
3336 let pull = SyncPullRequest {
3337 replica_id: "replica-a".into(),
3338 since_lsn: 0,
3339 max_units: MAX_SYNC_PULL_UNITS,
3340 max_bytes: MAX_SYNC_PULL_BYTES,
3341 database_id: identity.database_id,
3342 primary_generation: identity.primary_generation,
3343 wal_format_version: identity.wal_format_version,
3344 catalog_version: identity.catalog_version,
3345 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3346 };
3347
3348 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3349 Message::SyncPullResult {
3350 status,
3351 units,
3352 has_more,
3353 } => {
3354 assert_eq!(status.remote_lsn, remote_lsn);
3355 assert_eq!(status.servable_lsn, Some(servable_lsn));
3356 assert_eq!(status.unarchived_lsn, Some(1));
3357 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3358 assert!(!has_more);
3359 assert_eq!(units.len() as u64, servable_lsn);
3360 assert_eq!(units.last().unwrap().lsn, servable_lsn);
3361 assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
3362 }
3363 other => panic!("expected partial sync pull result, got {other:?}"),
3364 }
3365 }
3366
3367 #[test]
3368 fn sync_pull_rejects_cursor_or_format_mismatch() {
3369 let dir = tempfile::tempdir().unwrap();
3370 let mut engine = Engine::new(dir.path()).unwrap();
3371 engine
3372 .execute_powql("type SyncT { required id: int }")
3373 .unwrap();
3374 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3375 let remote_lsn = engine.catalog().max_lsn();
3376 write_sync_identity_and_tail(dir.path(), remote_lsn);
3377 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3378 .unwrap();
3379 let engine = Arc::new(RwLock::new(engine));
3380 let principal = admin_principal();
3381 let identity = sync_identity().segment_identity();
3382
3383 let wrong_cursor = SyncPullRequest {
3384 replica_id: "replica-a".into(),
3385 since_lsn: 1,
3386 max_units: 10,
3387 max_bytes: MAX_SYNC_PULL_BYTES,
3388 database_id: identity.database_id,
3389 primary_generation: identity.primary_generation,
3390 wal_format_version: identity.wal_format_version,
3391 catalog_version: identity.catalog_version,
3392 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3393 };
3394 match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
3395 Message::Error { message } => assert!(message.contains("does not match")),
3396 other => panic!("expected cursor mismatch error, got {other:?}"),
3397 }
3398
3399 let wrong_format = SyncPullRequest {
3400 replica_id: "replica-a".into(),
3401 since_lsn: 0,
3402 max_units: 10,
3403 max_bytes: MAX_SYNC_PULL_BYTES,
3404 database_id: identity.database_id,
3405 primary_generation: identity.primary_generation,
3406 wal_format_version: identity.wal_format_version,
3407 catalog_version: identity.catalog_version,
3408 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
3409 };
3410 match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
3411 Message::Error { message } => assert!(message.contains("rebootstrap required")),
3412 other => panic!("expected format mismatch error, got {other:?}"),
3413 }
3414 }
3415}