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, SegmentIdentity, SyncRepairAction,
13 RETAINED_SEGMENT_FORMAT_VERSION,
14};
15use std::collections::{HashMap, VecDeque};
16use std::net::IpAddr;
17use std::path::{Path, PathBuf};
18use std::sync::{Arc, Mutex, RwLock};
19use std::time::{Duration, Instant};
20use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
21use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};
22use tracing::{debug, error, info, warn};
23use zeroize::Zeroizing;
24
25pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
27
28#[derive(Clone)]
33pub struct TxGate {
34 semaphore: Arc<Semaphore>,
35 permit_count: u32,
36}
37
38pub const DEFAULT_TX_GATE_READER_PERMITS: u32 = 1024;
39
40pub fn new_tx_gate() -> TxGate {
42 new_tx_gate_with_permits(DEFAULT_TX_GATE_READER_PERMITS)
43}
44
45pub fn new_tx_gate_with_permits(permit_count: u32) -> TxGate {
51 assert!(permit_count > 0, "transaction gate requires a permit");
52 TxGate {
53 semaphore: Arc::new(Semaphore::new(permit_count as usize)),
54 permit_count,
55 }
56}
57
58impl TxGate {
59 pub fn permit_count(&self) -> u32 {
60 self.permit_count
61 }
62
63 fn available_permits(&self) -> usize {
64 self.semaphore.available_permits()
65 }
66
67 async fn acquire_many_owned(
68 self,
69 permits: u32,
70 ) -> Result<OwnedSemaphorePermit, tokio::sync::AcquireError> {
71 self.semaphore.acquire_many_owned(permits).await
72 }
73
74 fn try_acquire_many_owned(
75 self,
76 permits: u32,
77 ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
78 self.semaphore.try_acquire_many_owned(permits)
79 }
80}
81
82const MAX_QUERY_LENGTH: usize = 1024 * 1024;
84
85const MAX_WIRE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
88
89const MAX_IN_FLIGHT_READ_AHEAD_FRAMES: usize = 128;
95const MAX_IN_FLIGHT_READ_AHEAD_BYTES: usize = 1024 * 1024;
96
97const MAX_SYNC_PULL_UNITS: u32 = 4096;
99
100const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
102
103#[cfg(not(test))]
107const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
108#[cfg(test)]
109const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
110
111const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
114
115const MAX_AUTH_FAILURES: u32 = 5;
117
118const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
120
121pub fn new_rate_limiter() -> AuthRateLimiter {
123 Arc::new(Mutex::new(HashMap::new()))
124}
125
126fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
129 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
130 let now = Instant::now();
132 map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
133
134 if let Some((count, _)) = map.get(&ip) {
135 *count >= MAX_AUTH_FAILURES
136 } else {
137 false
138 }
139}
140
141fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
143 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
144 let now = Instant::now();
145 let entry = map.entry(ip).or_insert((0, now));
146 if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
148 *entry = (1, now);
149 } else {
150 entry.0 += 1;
151 }
152}
153
154fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
156 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
157 map.remove(&ip);
158}
159
160fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
163 use sha2::{Digest, Sha256};
164 let ha = Sha256::digest(a);
165 let hb = Sha256::digest(b);
166 let mut diff = 0u8;
167 for (x, y) in ha.iter().zip(hb.iter()) {
168 diff |= x ^ y;
169 }
170 diff == 0
171}
172
173#[derive(Debug, Clone, PartialEq, Eq)]
177pub struct Principal {
178 pub name: String,
179 pub role: String,
180}
181
182fn is_ddl_statement(stmt: &powdb_query::ast::Statement) -> bool {
188 use powdb_query::ast::Statement;
189 let inner = match stmt {
190 Statement::Explain(inner) => inner.as_ref(),
191 other => other,
192 };
193 matches!(
194 inner,
195 Statement::CreateType(_)
196 | Statement::AlterTable(_)
197 | Statement::DropTable(_)
198 | Statement::CreateView(_)
199 | Statement::DropView(_)
200 )
201}
202
203fn required_permission(stmt: &powdb_query::ast::Statement) -> Permission {
209 if is_read_only_statement(stmt) {
210 Permission::Read
211 } else if is_ddl_statement(stmt) {
212 Permission::Ddl
213 } else {
214 Permission::Write
215 }
216}
217
218fn check_statement_permitted(
230 principal: Option<&Principal>,
231 stmt: &powdb_query::ast::Statement,
232) -> Result<(), QueryError> {
233 let Some(p) = principal else {
234 return Ok(());
237 };
238 if is_read_only_statement(stmt) {
241 return Ok(());
242 }
243 let needed = required_permission(stmt);
244 if Role::builtin(&p.role).is_some_and(|r| r.allows(needed)) {
245 return Ok(());
246 }
247 let kind = if needed == Permission::Ddl {
248 "schema-definition"
249 } else {
250 "write"
251 };
252 Err(QueryError::Execution(format!(
253 "permission denied: role '{}' cannot execute {kind} statements",
254 p.role
255 )))
256}
257
258#[derive(Debug, Clone, PartialEq, Eq)]
260pub enum AuthOutcome {
261 Authenticated { principal: Option<Principal> },
265 Rejected,
268}
269
270pub fn authenticate_connect(
282 users: &UserStore,
283 expected_password: Option<&str>,
284 username: Option<&str>,
285 password: Option<&str>,
286) -> AuthOutcome {
287 if !users.is_empty() {
288 let Some(name) = username else {
290 return AuthOutcome::Rejected;
291 };
292 let Some(candidate) = password else {
293 return AuthOutcome::Rejected;
294 };
295 match users.authenticate(name, candidate) {
296 Some(user) => AuthOutcome::Authenticated {
297 principal: Some(Principal {
298 name: user.name.clone(),
299 role: user.role.clone(),
300 }),
301 },
302 None => AuthOutcome::Rejected,
303 }
304 } else {
305 match expected_password {
307 Some(expected) => {
308 if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
309 AuthOutcome::Authenticated { principal: None }
310 } else {
311 AuthOutcome::Rejected
312 }
313 }
314 None => AuthOutcome::Authenticated { principal: None },
315 }
316 }
317}
318
319const DEFAULT_DB_NAME: &str = "default";
323
324fn check_db_name(configured: Option<&str>, requested: &str) -> Result<(), String> {
334 if requested.is_empty() || requested == DEFAULT_DB_NAME {
335 return Ok(());
336 }
337 match configured {
338 None => Ok(()),
339 Some(name) if requested == name => Ok(()),
340 Some(name) => Err(format!(
341 "unknown database '{requested}'; this server serves '{name}'"
342 )),
343 }
344}
345
346const SAFE_ERROR_PREFIXES: &[&str] = &[
348 "table not found",
349 "table '",
353 "type '",
354 "column not found",
355 "at position",
358 "parse error",
359 "type mismatch",
360 "unknown table",
361 "unknown column",
362 "unknown function",
363 "syntax error",
364 "expected",
365 "unexpected",
366 "missing",
367 "duplicate",
368 "invalid",
369 "cannot",
370 "no such",
371 "already exists",
372 "permission denied",
373 "row too large",
374 "unique constraint violation",
375 "sort input exceeds",
380 "join result exceeds",
381 "query exceeded memory budget",
382 "result too large",
383 "wal durability sync failed",
388 "query timeout after",
392 "query cancelled",
393];
394
395fn sanitize_error(e: &str) -> String {
399 let lower = e.to_lowercase();
400 for prefix in SAFE_ERROR_PREFIXES {
401 if lower.starts_with(prefix) {
402 return e.to_string();
403 }
404 }
405 "query execution error".into()
406}
407
408async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
411 let write_fut = async {
412 if msg.write_to(writer).await.is_err() {
413 return false;
414 }
415 writer.flush().await.is_ok()
416 };
417 tokio::time::timeout(WRITE_TIMEOUT, write_fut)
418 .await
419 .unwrap_or_default()
420}
421
422pub struct ConnOpts<'a> {
425 pub engine: Arc<RwLock<Engine>>,
426 pub tx_gate: TxGate,
427 pub expected_password: Option<Zeroizing<String>>,
430 pub users: Arc<UserStore>,
434 pub shutdown_rx: &'a mut watch::Receiver<bool>,
435 pub idle_timeout: Duration,
436 pub query_timeout: Duration,
437 pub rate_limiter: Option<&'a AuthRateLimiter>,
438 pub peer_addr: Option<std::net::SocketAddr>,
439 pub metrics: Arc<Metrics>,
441 pub tx_wait_timeout: Duration,
445 pub db_name: Option<String>,
449}
450
451type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
464
465fn execute_write_deferred(
472 engine: &Arc<RwLock<Engine>>,
473 f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
474) -> DispatchOutcome {
475 let mut eng = match engine.write() {
476 Ok(eng) => eng,
477 Err(e) => {
478 return (
479 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
480 None,
481 )
482 }
483 };
484 eng.run_with_deferred_durability(f)
485}
486
487fn dispatch_query(
488 engine: &Arc<RwLock<Engine>>,
489 query: &str,
490 principal: Option<&Principal>,
491 allow_readonly_escalation: bool,
492) -> DispatchOutcome {
493 let stmt_result = parser::parse(query).map_err(|e| e.to_string());
494 dispatch_query_parsed(
495 engine,
496 query,
497 &stmt_result,
498 principal,
499 allow_readonly_escalation,
500 )
501}
502
503fn dispatch_query_parsed(
504 engine: &Arc<RwLock<Engine>>,
505 query: &str,
506 stmt_result: &Result<powdb_query::ast::Statement, String>,
507 principal: Option<&Principal>,
508 allow_readonly_escalation: bool,
509) -> DispatchOutcome {
510 if let Ok(stmt) = &stmt_result {
514 if let Err(e) = check_statement_permitted(principal, stmt) {
515 return (Err(e), None);
516 }
517 }
518
519 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
520 if can_try_read {
521 let res = {
522 let eng = match engine.read() {
523 Ok(eng) => eng,
524 Err(e) => {
525 return (
526 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
527 None,
528 )
529 }
530 };
531 eng.execute_powql_readonly(query)
532 };
533 match res {
534 Ok(r) => return (Ok(r), None),
535 Err(QueryError::ReadonlyNeedsWrite) => {
536 if !allow_readonly_escalation {
537 return (Err(QueryError::ReadonlyNeedsWrite), None);
538 }
539 }
542 Err(e) => return (Err(e), None),
543 }
544 }
545
546 if matches!(
547 parsed_transaction_control(stmt_result),
548 Some(TransactionControl::Rollback)
549 ) {
550 let mut eng = match engine.write() {
551 Ok(eng) => eng,
552 Err(e) => {
553 return (
554 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
555 None,
556 )
557 }
558 };
559 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
560 }
561 execute_write_deferred(engine, |eng| eng.execute_powql(query))
562}
563
564#[cfg(test)]
565fn dispatch_sql_query(
566 engine: &Arc<RwLock<Engine>>,
567 query: &str,
568 principal: Option<&Principal>,
569 allow_readonly_escalation: bool,
570) -> DispatchOutcome {
571 let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
572 dispatch_sql_query_parsed(
573 engine,
574 query,
575 &stmt_result,
576 principal,
577 allow_readonly_escalation,
578 )
579}
580
581fn dispatch_sql_query_parsed(
582 engine: &Arc<RwLock<Engine>>,
583 query: &str,
584 stmt_result: &Result<powdb_query::ast::Statement, String>,
585 principal: Option<&Principal>,
586 allow_readonly_escalation: bool,
587) -> DispatchOutcome {
588 if let Ok(stmt) = &stmt_result {
589 if let Err(e) = check_statement_permitted(principal, stmt) {
590 return (Err(e), None);
591 }
592 }
593
594 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
595 if can_try_read {
596 let res = {
597 let eng = match engine.read() {
598 Ok(eng) => eng,
599 Err(e) => {
600 return (
601 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
602 None,
603 )
604 }
605 };
606 eng.execute_sql_readonly(query)
607 };
608 match res {
609 Ok(r) => return (Ok(r), None),
610 Err(QueryError::ReadonlyNeedsWrite) => {
611 if !allow_readonly_escalation {
612 return (Err(QueryError::ReadonlyNeedsWrite), None);
613 }
614 }
615 Err(e) => return (Err(e), None),
616 }
617 }
618
619 if matches!(
620 parsed_transaction_control(stmt_result),
621 Some(TransactionControl::Rollback)
622 ) {
623 let mut eng = match engine.write() {
624 Ok(eng) => eng,
625 Err(e) => {
626 return (
627 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
628 None,
629 )
630 }
631 };
632 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
633 }
634 execute_write_deferred(engine, |eng| eng.execute_sql(query))
635}
636
637fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
640 use powdb_query::ast::ParamValue;
641 match p {
642 WireParam::Null => ParamValue::Null,
643 WireParam::Int(v) => ParamValue::Int(*v),
644 WireParam::Float(v) => ParamValue::Float(*v),
645 WireParam::Bool(v) => ParamValue::Bool(*v),
646 WireParam::Str(s) => ParamValue::Str(s.clone()),
647 }
648}
649
650#[derive(Debug, Clone, Copy, PartialEq, Eq)]
656enum TransactionControl {
657 Begin,
658 Commit,
659 Rollback,
660}
661
662#[derive(Debug, Clone, Copy, PartialEq, Eq)]
663enum AdmissionMode {
664 Reader,
665 Writer,
666}
667
668#[derive(Clone, Copy, Debug, PartialEq, Eq)]
669enum WireResultMode {
670 LegacyText,
671 Native,
672}
673
674fn statement_admission(stmt: &powdb_query::ast::Statement) -> AdmissionMode {
675 if is_read_only_statement(stmt) {
676 AdmissionMode::Reader
677 } else {
678 AdmissionMode::Writer
679 }
680}
681
682#[cfg(test)]
683fn classify_query_admission(query: &str) -> AdmissionMode {
684 parser::parse(query)
685 .map(|stmt| statement_admission(&stmt))
686 .unwrap_or(AdmissionMode::Writer)
687}
688
689#[cfg(test)]
690fn classify_sql_admission(query: &str) -> AdmissionMode {
691 sql::parse_sql(query)
692 .map(|stmt| statement_admission(&stmt))
693 .unwrap_or(AdmissionMode::Writer)
694}
695
696#[cfg(test)]
697fn classify_params_admission(query: &str, params: &[WireParam]) -> AdmissionMode {
698 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
699 parser::parse_with_params(query, &bound)
700 .map(|stmt| statement_admission(&stmt))
701 .unwrap_or(AdmissionMode::Writer)
702}
703
704fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
705 use powdb_query::ast::Statement;
706 match stmt {
707 Statement::Begin => Some(TransactionControl::Begin),
708 Statement::Commit => Some(TransactionControl::Commit),
709 Statement::Rollback => Some(TransactionControl::Rollback),
710 _ => None,
711 }
712}
713
714fn parsed_transaction_control(
715 stmt_result: &Result<powdb_query::ast::Statement, String>,
716) -> Option<TransactionControl> {
717 stmt_result.as_ref().ok().and_then(transaction_control)
718}
719
720fn execute_rollback_preserving_sync_if_needed(
721 engine: &mut Engine,
722) -> Result<QueryResult, QueryError> {
723 engine.rollback_transaction_preserving_wal_archive()
724}
725
726#[cfg(test)]
727fn dispatch_query_with_params(
728 engine: &Arc<RwLock<Engine>>,
729 query: &str,
730 params: &[WireParam],
731 principal: Option<&Principal>,
732 allow_readonly_escalation: bool,
733) -> DispatchOutcome {
734 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
735
736 let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
739 dispatch_query_with_bound_params_parsed(
740 engine,
741 query,
742 &bound,
743 &stmt_result,
744 principal,
745 allow_readonly_escalation,
746 )
747}
748
749fn dispatch_query_with_bound_params_parsed(
750 engine: &Arc<RwLock<Engine>>,
751 query: &str,
752 bound: &[powdb_query::ast::ParamValue],
753 stmt_result: &Result<powdb_query::ast::Statement, String>,
754 principal: Option<&Principal>,
755 allow_readonly_escalation: bool,
756) -> DispatchOutcome {
757 if let Ok(stmt) = &stmt_result {
758 if let Err(e) = check_statement_permitted(principal, stmt) {
759 return (Err(e), None);
760 }
761 }
762
763 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
764 if can_try_read {
765 let res = {
766 let eng = match engine.read() {
767 Ok(eng) => eng,
768 Err(e) => {
769 return (
770 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
771 None,
772 )
773 }
774 };
775 eng.execute_powql_readonly_with_params(query, bound)
776 };
777 match res {
778 Ok(r) => return (Ok(r), None),
779 Err(QueryError::ReadonlyNeedsWrite) => {
780 if !allow_readonly_escalation {
781 return (Err(QueryError::ReadonlyNeedsWrite), None);
782 }
783 }
784 Err(e) => return (Err(e), None),
785 }
786 }
787
788 if matches!(
789 parsed_transaction_control(stmt_result),
790 Some(TransactionControl::Rollback)
791 ) {
792 let mut eng = match engine.write() {
793 Ok(eng) => eng,
794 Err(e) => {
795 return (
796 Err(QueryError::Execution(format!("lock poisoned: {e}"))),
797 None,
798 )
799 }
800 };
801 return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
802 }
803 execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, bound))
804}
805
806#[derive(Debug, Clone)]
807struct SyncPullRequest {
808 replica_id: String,
809 since_lsn: u64,
810 max_units: u32,
811 max_bytes: u64,
812 database_id: [u8; 16],
813 primary_generation: u64,
814 wal_format_version: u16,
815 catalog_version: u16,
816 segment_format_version: u16,
817}
818
819#[derive(Debug, Clone, Copy, PartialEq, Eq)]
820enum SyncErrorClass {
821 AuthRequired,
822 PermissionDenied,
823 InvalidReplicaId,
824 ActiveTransaction,
825 QueryExecution,
826 InvalidMaxUnits,
827 InvalidMaxBytes,
828 SyncContext,
829 StatusRead,
830 CursorLsnMismatch,
831 IdentityRead,
832 IdentityOrFormatMismatch,
833 RetainedRead,
834 RetainedUnitEncoding,
835 RetainedChunkNotApplyable,
836 LsnAheadOfRemote,
837 AckValidation,
838 AckUpdate,
839 Internal,
840}
841
842impl SyncErrorClass {
843 const fn as_label(self) -> &'static str {
844 match self {
845 Self::AuthRequired => "auth_required",
846 Self::PermissionDenied => "permission_denied",
847 Self::InvalidReplicaId => "invalid_replica_id",
848 Self::ActiveTransaction => "active_transaction",
849 Self::QueryExecution => "query_execution",
850 Self::InvalidMaxUnits => "invalid_max_units",
851 Self::InvalidMaxBytes => "invalid_max_bytes",
852 Self::SyncContext => "sync_context",
853 Self::StatusRead => "status_read",
854 Self::CursorLsnMismatch => "cursor_lsn_mismatch",
855 Self::IdentityRead => "identity_read",
856 Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
857 Self::RetainedRead => "retained_read",
858 Self::RetainedUnitEncoding => "retained_unit_encoding",
859 Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
860 Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
861 Self::AckValidation => "ack_validation",
862 Self::AckUpdate => "ack_update",
863 Self::Internal => "internal",
864 }
865 }
866}
867
868#[derive(Debug, Clone)]
869struct SyncDecision {
870 message: Message,
871 error_class: Option<SyncErrorClass>,
872}
873
874impl SyncDecision {
875 fn ok(message: Message) -> Self {
876 Self {
877 message,
878 error_class: None,
879 }
880 }
881
882 fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
883 Self {
884 message: Message::Error {
885 message: message.into(),
886 },
887 error_class: Some(class),
888 }
889 }
890}
891
892fn check_sync_protocol_permitted(
893 credential_authenticated: bool,
894 principal: Option<&Principal>,
895) -> Result<(), (SyncErrorClass, String)> {
896 if !credential_authenticated {
897 return Err((
898 SyncErrorClass::AuthRequired,
899 "sync protocol requires authentication".to_string(),
900 ));
901 }
902 if let Some(principal) = principal {
903 let allowed =
904 Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
905 if !allowed {
906 return Err((
907 SyncErrorClass::PermissionDenied,
908 format!(
909 "permission denied: role '{}' cannot use sync protocol",
910 principal.role
911 ),
912 ));
913 }
914 }
915 Ok(())
916}
917
918fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
919 if replica_id.is_empty() {
920 return Err("replica id must be non-empty".to_string());
921 }
922 if replica_id.len() > 128 {
923 return Err("replica id must be at most 128 bytes".to_string());
924 }
925 if !replica_id
926 .bytes()
927 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
928 {
929 return Err("replica id contains unsupported characters".to_string());
930 }
931 Ok(())
932}
933
934fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<SyncContext, String> {
935 let engine = engine
936 .read()
937 .map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
938 let catalog = engine.catalog();
939 Ok(SyncContext {
940 data_dir: catalog.data_dir().to_path_buf(),
941 remote_lsn: catalog.max_lsn(),
942 active_catalog_version: catalog.active_catalog_version(),
943 })
944}
945
946struct SyncContext {
951 data_dir: PathBuf,
952 remote_lsn: u64,
953 active_catalog_version: u16,
954}
955
956fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
957 match action {
958 SyncRepairAction::None => WireSyncRepairAction::None,
959 SyncRepairAction::Pull => WireSyncRepairAction::Pull,
960 SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
961 SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
962 }
963}
964
965fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
966 WireSyncStatus {
967 replica_id: status.replica_id,
968 active: status.active,
969 last_applied_lsn: status.last_applied_lsn,
970 remote_lsn: status.remote_lsn,
971 servable_lsn: status.servable_lsn,
972 unarchived_lsn: status.unarchived_lsn,
973 lag_lsn: status.lag_lsn,
974 lag_bytes: status.lag_bytes,
975 lag_ms: status.lag_ms,
976 stale: status.stale,
977 repair_action: wire_repair_action(status.repair_action),
978 last_sync_error: status.last_sync_error,
979 }
980}
981
982fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
983 WireRetainedUnit {
984 tx_id: unit.tx_id,
985 record_type: unit.record_type,
986 lsn: unit.lsn,
987 data: unit.data,
988 }
989}
990
991fn sync_operation_outcome(message: &Message) -> SyncOutcome {
992 match message {
993 Message::SyncStatusResult { .. }
994 | Message::SyncPullResult { .. }
995 | Message::SyncAckResult { .. } => SyncOutcome::Ok,
996 _ => SyncOutcome::Error,
997 }
998}
999
1000fn sync_operation_label(operation: SyncOperation) -> &'static str {
1001 match operation {
1002 SyncOperation::Status => "status",
1003 SyncOperation::Pull => "pull",
1004 SyncOperation::Ack => "ack",
1005 }
1006}
1007
1008fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
1009 units.iter().fold(0, |total, unit| {
1010 total.saturating_add(unit.encoded_len().unwrap_or(0))
1011 })
1012}
1013
1014fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
1015 match action {
1016 WireSyncRepairAction::None => SyncRepairLabel::None,
1017 WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
1018 WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
1019 WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
1020 }
1021}
1022
1023fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
1024 match action {
1025 WireSyncRepairAction::None => "none",
1026 WireSyncRepairAction::Pull => "pull",
1027 WireSyncRepairAction::AwaitArchive => "await_archive",
1028 WireSyncRepairAction::Rebootstrap => "rebootstrap",
1029 }
1030}
1031
1032const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
1033const FNV1A64_PRIME: u64 = 0x100000001b3;
1034const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
1035
1036fn replica_fingerprint(replica_id: &str) -> String {
1037 let mut hash = FNV1A64_OFFSET;
1038 for byte in replica_id.as_bytes() {
1039 hash ^= u64::from(*byte);
1040 hash = hash.wrapping_mul(FNV1A64_PRIME);
1041 }
1042 format!("{hash:016x}")
1043}
1044
1045fn log_replica_fingerprint(replica_id: &str) -> String {
1046 if validate_wire_replica_id(replica_id).is_ok() {
1047 replica_fingerprint(replica_id)
1048 } else {
1049 INVALID_REPLICA_FINGERPRINT.to_string()
1050 }
1051}
1052
1053#[derive(Debug, Clone)]
1054struct SyncLogContext {
1055 replica_fingerprint: String,
1056 since_lsn: Option<u64>,
1057 applied_lsn: Option<u64>,
1058 observed_remote_lsn: Option<u64>,
1059 max_units: Option<u32>,
1060 max_bytes: Option<u64>,
1061}
1062
1063impl SyncLogContext {
1064 fn status(replica_id: &str) -> Self {
1065 Self::base(replica_id)
1066 }
1067
1068 fn pull(request: &SyncPullRequest) -> Self {
1069 Self {
1070 replica_fingerprint: log_replica_fingerprint(&request.replica_id),
1071 since_lsn: Some(request.since_lsn),
1072 applied_lsn: None,
1073 observed_remote_lsn: None,
1074 max_units: Some(request.max_units),
1075 max_bytes: Some(request.max_bytes),
1076 }
1077 }
1078
1079 fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
1080 Self {
1081 replica_fingerprint: log_replica_fingerprint(replica_id),
1082 since_lsn: None,
1083 applied_lsn: Some(applied_lsn),
1084 observed_remote_lsn: Some(observed_remote_lsn),
1085 max_units: None,
1086 max_bytes: None,
1087 }
1088 }
1089
1090 fn base(replica_id: &str) -> Self {
1091 Self {
1092 replica_fingerprint: log_replica_fingerprint(replica_id),
1093 since_lsn: None,
1094 applied_lsn: None,
1095 observed_remote_lsn: None,
1096 max_units: None,
1097 max_bytes: None,
1098 }
1099 }
1100}
1101
1102struct SyncExecutionContext<'a> {
1103 tx_gate: TxGate,
1104 connection_has_transaction: bool,
1105 operation: SyncOperation,
1106 log_context: SyncLogContext,
1107 metrics: &'a Arc<Metrics>,
1108 query_timeout: Duration,
1109}
1110
1111fn log_sync_decision(
1112 operation: SyncOperation,
1113 context: &SyncLogContext,
1114 elapsed: Duration,
1115 decision: &SyncDecision,
1116) {
1117 let operation = sync_operation_label(operation);
1118 let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
1119 match &decision.message {
1120 Message::SyncStatusResult { status } => {
1121 let repair_action = wire_repair_action_label(status.repair_action);
1122 if status.stale || status.repair_action != WireSyncRepairAction::None {
1123 info!(
1124 operation = operation,
1125 replica_fingerprint = %context.replica_fingerprint,
1126 remote_lsn = status.remote_lsn,
1127 last_applied_lsn = ?status.last_applied_lsn,
1128 servable_lsn = ?status.servable_lsn,
1129 unarchived_lsn = ?status.unarchived_lsn,
1130 lag_lsn = ?status.lag_lsn,
1131 lag_bytes = ?status.lag_bytes,
1132 lag_ms = ?status.lag_ms,
1133 stale = status.stale,
1134 repair_action,
1135 elapsed_ms,
1136 "sync decision"
1137 );
1138 } else {
1139 debug!(
1140 operation = operation,
1141 replica_fingerprint = %context.replica_fingerprint,
1142 remote_lsn = status.remote_lsn,
1143 last_applied_lsn = ?status.last_applied_lsn,
1144 repair_action,
1145 elapsed_ms,
1146 "sync decision"
1147 );
1148 }
1149 }
1150 Message::SyncPullResult {
1151 status,
1152 units,
1153 has_more,
1154 } => {
1155 let repair_action = wire_repair_action_label(status.repair_action);
1156 let units_len = units.len();
1157 let payload_bytes = sync_pull_payload_bytes(units);
1158 if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
1159 info!(
1160 operation = operation,
1161 replica_fingerprint = %context.replica_fingerprint,
1162 since_lsn = ?context.since_lsn,
1163 max_units = ?context.max_units,
1164 max_bytes = ?context.max_bytes,
1165 units = units_len,
1166 payload_bytes,
1167 has_more = *has_more,
1168 remote_lsn = status.remote_lsn,
1169 last_applied_lsn = ?status.last_applied_lsn,
1170 servable_lsn = ?status.servable_lsn,
1171 unarchived_lsn = ?status.unarchived_lsn,
1172 lag_lsn = ?status.lag_lsn,
1173 stale = status.stale,
1174 repair_action,
1175 elapsed_ms,
1176 "sync decision"
1177 );
1178 } else {
1179 debug!(
1180 operation = operation,
1181 replica_fingerprint = %context.replica_fingerprint,
1182 since_lsn = ?context.since_lsn,
1183 units = units_len,
1184 payload_bytes,
1185 has_more = *has_more,
1186 remote_lsn = status.remote_lsn,
1187 repair_action,
1188 elapsed_ms,
1189 "sync decision"
1190 );
1191 }
1192 }
1193 Message::SyncAckResult {
1194 previous_applied_lsn,
1195 applied_lsn,
1196 remote_lsn,
1197 advanced,
1198 status,
1199 } => {
1200 let repair_action = wire_repair_action_label(status.repair_action);
1201 if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
1202 info!(
1203 operation = operation,
1204 replica_fingerprint = %context.replica_fingerprint,
1205 requested_applied_lsn = ?context.applied_lsn,
1206 observed_remote_lsn = ?context.observed_remote_lsn,
1207 previous_applied_lsn = *previous_applied_lsn,
1208 applied_lsn = *applied_lsn,
1209 remote_lsn = *remote_lsn,
1210 advanced = *advanced,
1211 stale = status.stale,
1212 repair_action,
1213 elapsed_ms,
1214 "sync decision"
1215 );
1216 } else {
1217 debug!(
1218 operation = operation,
1219 replica_fingerprint = %context.replica_fingerprint,
1220 requested_applied_lsn = ?context.applied_lsn,
1221 previous_applied_lsn = *previous_applied_lsn,
1222 applied_lsn = *applied_lsn,
1223 remote_lsn = *remote_lsn,
1224 advanced = *advanced,
1225 repair_action,
1226 elapsed_ms,
1227 "sync decision"
1228 );
1229 }
1230 }
1231 Message::Error { .. } => {
1232 warn!(
1233 operation = operation,
1234 replica_fingerprint = %context.replica_fingerprint,
1235 since_lsn = ?context.since_lsn,
1236 applied_lsn = ?context.applied_lsn,
1237 observed_remote_lsn = ?context.observed_remote_lsn,
1238 max_units = ?context.max_units,
1239 max_bytes = ?context.max_bytes,
1240 error_class = decision
1241 .error_class
1242 .unwrap_or(SyncErrorClass::Internal)
1243 .as_label(),
1244 elapsed_ms,
1245 "sync decision rejected"
1246 );
1247 }
1248 _ => {
1249 debug!(
1250 operation = operation,
1251 replica_fingerprint = %context.replica_fingerprint,
1252 elapsed_ms,
1253 "unexpected sync decision response"
1254 );
1255 }
1256 }
1257}
1258
1259fn trim_to_applyable_v1_prefix(
1260 raw_units: &mut Vec<RetainedUnit>,
1261 wire_units: &mut Vec<WireRetainedUnit>,
1262) -> Result<(), String> {
1263 let mut last_error = None;
1264 while !raw_units.is_empty() {
1265 match validate_v1_retained_units_applyable(raw_units) {
1266 Ok(()) => return Ok(()),
1267 Err(err) => {
1268 last_error = Some(err.to_string());
1269 raw_units.pop();
1270 wire_units.pop();
1271 }
1272 }
1273 }
1274 if let Some(error) = last_error {
1275 return Err(format!(
1276 "sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
1277 ));
1278 }
1279 Ok(())
1280}
1281
1282fn validate_sync_ack_applyable_boundary(
1283 data_dir: &Path,
1284 replica_id: &str,
1285 applied_lsn: u64,
1286 remote_lsn: u64,
1287) -> Result<(), String> {
1288 let status =
1289 replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
1290 let Some(previous_lsn) = status.last_applied_lsn else {
1291 return Ok(());
1292 };
1293 if applied_lsn <= previous_lsn {
1294 return Ok(());
1295 }
1296 let range_len = applied_lsn - previous_lsn;
1297 if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
1298 return Err(format!(
1299 "sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
1300 ));
1301 }
1302 let max_units =
1303 usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
1304 let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
1305 let segment_dir = retained_segments_dir(data_dir);
1306 let units = read_units_through(
1307 &segment_dir,
1308 identity.segment_identity(),
1309 previous_lsn,
1310 applied_lsn,
1311 max_units,
1312 )
1313 .map_err(|err| err.to_string())?;
1314 if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
1315 return Err(
1316 "sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
1317 );
1318 }
1319 validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
1320}
1321
1322#[cfg(test)]
1323fn dispatch_sync_status(
1324 engine: &Arc<RwLock<Engine>>,
1325 replica_id: String,
1326 credential_authenticated: bool,
1327 principal: Option<&Principal>,
1328) -> Message {
1329 dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
1330}
1331
1332fn dispatch_sync_status_decision(
1333 engine: &Arc<RwLock<Engine>>,
1334 replica_id: String,
1335 credential_authenticated: bool,
1336 principal: Option<&Principal>,
1337) -> SyncDecision {
1338 if let Err((class, message)) =
1339 check_sync_protocol_permitted(credential_authenticated, principal)
1340 {
1341 return SyncDecision::error(class, message);
1342 }
1343 if let Err(message) = validate_wire_replica_id(&replica_id) {
1344 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1345 }
1346 let SyncContext {
1347 data_dir,
1348 remote_lsn,
1349 ..
1350 } = match sync_context(engine) {
1351 Ok(context) => context,
1352 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1353 };
1354 match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
1355 Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
1356 status: wire_sync_status(status),
1357 }),
1358 Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
1359 }
1360}
1361
1362#[cfg(test)]
1363fn dispatch_sync_pull(
1364 engine: &Arc<RwLock<Engine>>,
1365 request: SyncPullRequest,
1366 credential_authenticated: bool,
1367 principal: Option<&Principal>,
1368) -> Message {
1369 dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
1370}
1371
1372fn dispatch_sync_pull_decision(
1373 engine: &Arc<RwLock<Engine>>,
1374 request: SyncPullRequest,
1375 credential_authenticated: bool,
1376 principal: Option<&Principal>,
1377) -> SyncDecision {
1378 if let Err((class, message)) =
1379 check_sync_protocol_permitted(credential_authenticated, principal)
1380 {
1381 return SyncDecision::error(class, message);
1382 }
1383 if let Err(message) = validate_wire_replica_id(&request.replica_id) {
1384 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1385 }
1386 if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
1387 return SyncDecision::error(
1388 SyncErrorClass::InvalidMaxUnits,
1389 format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
1390 );
1391 }
1392 if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
1393 return SyncDecision::error(
1394 SyncErrorClass::InvalidMaxBytes,
1395 format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
1396 );
1397 }
1398
1399 let SyncContext {
1400 data_dir,
1401 remote_lsn,
1402 active_catalog_version,
1403 } = match sync_context(engine) {
1404 Ok(context) => context,
1405 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1406 };
1407 let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
1408 Ok(status) => status,
1409 Err(err) => {
1410 return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
1411 }
1412 };
1413 let Some(cursor_lsn) = status.last_applied_lsn else {
1414 return SyncDecision::ok(Message::SyncPullResult {
1415 status: wire_sync_status(status),
1416 units: Vec::new(),
1417 has_more: false,
1418 });
1419 };
1420 if status.repair_action != SyncRepairAction::Pull {
1421 return SyncDecision::ok(Message::SyncPullResult {
1422 status: wire_sync_status(status),
1423 units: Vec::new(),
1424 has_more: false,
1425 });
1426 }
1427 if request.since_lsn != cursor_lsn {
1428 return SyncDecision::error(
1429 SyncErrorClass::CursorLsnMismatch,
1430 format!(
1431 "sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
1432 request.since_lsn
1433 ),
1434 );
1435 }
1436
1437 let identity = match read_identity(&data_dir) {
1438 Ok(identity) => identity,
1439 Err(err) => {
1440 return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
1441 }
1442 };
1443 let expected = SegmentIdentity::with_catalog_version(
1447 identity.database_id,
1448 identity.primary_generation,
1449 active_catalog_version,
1450 );
1451 if request.database_id != expected.database_id
1452 || request.primary_generation != expected.primary_generation
1453 || request.wal_format_version != expected.wal_format_version
1454 || request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
1455 {
1456 return SyncDecision::error(
1457 SyncErrorClass::IdentityOrFormatMismatch,
1458 "sync pull identity or format version mismatch; rebootstrap required",
1459 );
1460 }
1461 if request.catalog_version < expected.catalog_version {
1465 return SyncDecision::error(
1466 SyncErrorClass::IdentityOrFormatMismatch,
1467 format!(
1468 "sync pull replica catalog format v{} cannot read this database's active catalog format v{}; rebootstrap with an upgraded replica required",
1469 request.catalog_version, expected.catalog_version
1470 ),
1471 );
1472 }
1473
1474 let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
1475 let requested_through_lsn = request
1476 .since_lsn
1477 .saturating_add(request.max_units as u64)
1478 .min(remote_lsn)
1479 .min(status.servable_lsn.unwrap_or(request.since_lsn));
1480 let segment_dir = retained_segments_dir(&data_dir);
1481 if requested_through_lsn > request.since_lsn {
1482 if let Err(err) = validate_retained_tail_available(
1483 &segment_dir,
1484 expected,
1485 request.since_lsn,
1486 requested_through_lsn,
1487 ) {
1488 let mut rebootstrap_status = status;
1489 rebootstrap_status.stale = true;
1490 rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
1491 rebootstrap_status.last_sync_error = Some(format!(
1492 "retained history is unavailable; rebootstrap required: {err}"
1493 ));
1494 return SyncDecision::ok(Message::SyncPullResult {
1495 status: wire_sync_status(rebootstrap_status),
1496 units: Vec::new(),
1497 has_more: false,
1498 });
1499 }
1500 }
1501
1502 let raw_units = match read_units_through(
1503 &segment_dir,
1504 expected,
1505 request.since_lsn,
1506 requested_through_lsn,
1507 effective_max_units,
1508 ) {
1509 Ok(units) => units,
1510 Err(err) => {
1511 return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
1512 }
1513 };
1514
1515 let mut selected_raw = Vec::new();
1516 let mut selected = Vec::new();
1517 let mut selected_bytes = 0u64;
1518 for unit in raw_units {
1519 let wire_unit = wire_retained_unit(unit.clone());
1520 let unit_bytes = match wire_unit.encoded_len() {
1521 Ok(bytes) => bytes,
1522 Err(message) => {
1523 return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
1524 }
1525 };
1526 if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
1527 if selected.is_empty() {
1528 return SyncDecision::error(
1529 SyncErrorClass::InvalidMaxBytes,
1530 "sync pull maxBytes is too small for the next retained unit",
1531 );
1532 }
1533 break;
1534 }
1535 selected_bytes += unit_bytes;
1536 selected_raw.push(unit);
1537 selected.push(wire_unit);
1538 }
1539 if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
1540 return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
1541 }
1542
1543 let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
1544 let has_more = selected
1545 .last()
1546 .is_some_and(|unit| unit.lsn < fetchable_through_lsn);
1547 SyncDecision::ok(Message::SyncPullResult {
1548 status: wire_sync_status(status),
1549 units: selected,
1550 has_more,
1551 })
1552}
1553
1554#[cfg(test)]
1555fn dispatch_sync_ack(
1556 engine: &Arc<RwLock<Engine>>,
1557 replica_id: String,
1558 applied_lsn: u64,
1559 observed_remote_lsn: u64,
1560 credential_authenticated: bool,
1561 principal: Option<&Principal>,
1562) -> Message {
1563 dispatch_sync_ack_decision(
1564 engine,
1565 replica_id,
1566 applied_lsn,
1567 observed_remote_lsn,
1568 credential_authenticated,
1569 principal,
1570 )
1571 .message
1572}
1573
1574fn dispatch_sync_ack_decision(
1575 engine: &Arc<RwLock<Engine>>,
1576 replica_id: String,
1577 applied_lsn: u64,
1578 observed_remote_lsn: u64,
1579 credential_authenticated: bool,
1580 principal: Option<&Principal>,
1581) -> SyncDecision {
1582 if let Err((class, message)) =
1583 check_sync_protocol_permitted(credential_authenticated, principal)
1584 {
1585 return SyncDecision::error(class, message);
1586 }
1587 if let Err(message) = validate_wire_replica_id(&replica_id) {
1588 return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1589 }
1590 if applied_lsn > observed_remote_lsn {
1591 return SyncDecision::error(
1592 SyncErrorClass::LsnAheadOfRemote,
1593 format!(
1594 "sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
1595 ),
1596 );
1597 }
1598 let SyncContext {
1599 data_dir,
1600 remote_lsn,
1601 ..
1602 } = match sync_context(engine) {
1603 Ok(context) => context,
1604 Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1605 };
1606 if observed_remote_lsn > remote_lsn {
1607 return SyncDecision::error(
1608 SyncErrorClass::LsnAheadOfRemote,
1609 format!(
1610 "sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
1611 ),
1612 );
1613 }
1614 if let Err(message) =
1615 validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
1616 {
1617 return SyncDecision::error(SyncErrorClass::AckValidation, message);
1618 }
1619 match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
1620 Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
1621 previous_applied_lsn: summary.previous_applied_lsn,
1622 applied_lsn: summary.applied_lsn,
1623 remote_lsn: summary.remote_lsn,
1624 advanced: summary.advanced,
1625 status: wire_sync_status(summary.status),
1626 }),
1627 Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
1628 }
1629}
1630
1631async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
1632where
1633 T: Send + 'static,
1634 F: FnOnce(T) -> SyncDecision + Send + 'static,
1635{
1636 let mut handle = tokio::task::spawn_blocking(move || f(input));
1637 tokio::select! {
1638 result = &mut handle => match result {
1639 Ok(decision) => decision,
1640 Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1641 },
1642 _ = tokio::time::sleep(query_timeout) => match handle.await {
1643 Ok(decision) => decision,
1644 Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1645 },
1646 }
1647}
1648
1649async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
1650where
1651 T: Send + 'static,
1652 F: FnOnce(T) -> SyncDecision + Send + 'static,
1653{
1654 let SyncExecutionContext {
1655 tx_gate,
1656 connection_has_transaction,
1657 operation,
1658 log_context,
1659 metrics,
1660 query_timeout,
1661 } = context;
1662 let start = Instant::now();
1663 if connection_has_transaction {
1664 let decision = SyncDecision::error(
1665 SyncErrorClass::ActiveTransaction,
1666 "sync protocol is unavailable inside an active transaction",
1667 );
1668 let elapsed = start.elapsed();
1669 log_sync_decision(operation, &log_context, elapsed, &decision);
1670 metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1671 return decision.message;
1672 }
1673
1674 let permit_count = tx_gate.permit_count();
1675 let permit = match tx_gate.acquire_many_owned(permit_count).await {
1676 Ok(permit) => permit,
1677 Err(_) => {
1678 let decision =
1679 SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
1680 let elapsed = start.elapsed();
1681 log_sync_decision(operation, &log_context, elapsed, &decision);
1682 metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1683 return decision.message;
1684 }
1685 };
1686 let decision = run_blocking_sync(input, query_timeout, f).await;
1687 drop(permit);
1688 match &decision.message {
1689 Message::SyncStatusResult { status } => {
1690 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1691 }
1692 Message::SyncPullResult { status, units, .. } => {
1693 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1694 metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
1695 }
1696 Message::SyncAckResult {
1697 advanced, status, ..
1698 } => {
1699 metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1700 if *advanced {
1701 metrics.inc_sync_ack_advanced();
1702 }
1703 }
1704 _ => {}
1705 }
1706 let elapsed = start.elapsed();
1707 log_sync_decision(operation, &log_context, elapsed, &decision);
1708 metrics.record_sync_operation(
1709 operation,
1710 elapsed,
1711 sync_operation_outcome(&decision.message),
1712 );
1713 decision.message
1714}
1715
1716async fn acquire_begin_permit(
1723 tx_gate: &TxGate,
1724 tx_wait_timeout: Duration,
1725 metrics: &Arc<Metrics>,
1726) -> Result<OwnedSemaphorePermit, Message> {
1727 match tokio::time::timeout(
1728 tx_wait_timeout,
1729 tx_gate.clone().acquire_many_owned(tx_gate.permit_count()),
1730 )
1731 .await
1732 {
1733 Ok(Ok(permit)) => Ok(permit),
1734 Ok(Err(_)) => Err(Message::Error {
1735 message: "query execution error".into(),
1736 }),
1737 Err(_) => {
1738 metrics.inc_tx_gate_timeout();
1739 Err(Message::Error {
1740 message: format!(
1741 "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1742 tx_wait_timeout.as_millis()
1743 ),
1744 })
1745 }
1746 }
1747}
1748
1749async fn acquire_autocommit_permit(
1759 tx_gate: &TxGate,
1760 admission: AdmissionMode,
1761 tx_wait_timeout: Duration,
1762 metrics: &Arc<Metrics>,
1763) -> Result<OwnedSemaphorePermit, Message> {
1764 let permits = match admission {
1765 AdmissionMode::Reader => 1,
1766 AdmissionMode::Writer => tx_gate.permit_count(),
1767 };
1768
1769 if let Ok(permit) = tx_gate.clone().try_acquire_many_owned(permits) {
1775 return Ok(permit);
1776 }
1777
1778 let acquire = tx_gate.clone().acquire_many_owned(permits);
1779 match tokio::time::timeout(tx_wait_timeout, acquire).await {
1780 Ok(Ok(permit)) => Ok(permit),
1781 Ok(Err(_)) => Err(Message::Error {
1782 message: "query execution error".into(),
1783 }),
1784 Err(_) => {
1785 metrics.inc_tx_gate_timeout();
1786 Err(Message::Error {
1787 message: format!(
1788 "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1789 tx_wait_timeout.as_millis()
1790 ),
1791 })
1792 }
1793 }
1794}
1795
1796#[allow(clippy::too_many_arguments)]
1812async fn run_wire_query_state_machine<Inner, D, R>(
1813 engine: Arc<RwLock<Engine>>,
1814 tx_gate: TxGate,
1815 tx_permit: &mut Option<OwnedSemaphorePermit>,
1816 parsed_query: Arc<Inner>,
1817 tx_control: Option<TransactionControl>,
1818 autocommit_admission: AdmissionMode,
1819 result_mode: WireResultMode,
1820 principal: Option<Principal>,
1821 query_timeout: Duration,
1822 query_deadline: Instant,
1823 tx_wait_timeout: Duration,
1824 metrics: &Arc<Metrics>,
1825 reader: &mut BufReader<R>,
1826 wire_read_buffer: &mut Vec<u8>,
1827 pending_messages: &mut InFlightReadAhead,
1828 dispatch: D,
1829) -> (
1830 Message,
1831 Option<PendingDurability>,
1832 Option<ConnectionTermination>,
1833)
1834where
1835 Inner: Send + Sync + 'static,
1836 D: Fn(Arc<RwLock<Engine>>, Arc<Inner>, Option<Principal>, bool) -> DispatchOutcome
1837 + Clone
1838 + Send
1839 + Sync
1840 + 'static,
1841 R: AsyncRead + Unpin,
1842{
1843 match tx_control {
1844 Some(TransactionControl::Begin) => {
1845 if tx_permit.is_some() {
1846 return (
1847 Message::Error {
1848 message: sanitize_error(
1849 "cannot begin: a transaction is already active on this connection",
1850 ),
1851 },
1852 None,
1853 None,
1854 );
1855 }
1856 let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
1857 Ok(permit) => permit,
1858 Err(response) => return (response, None, None),
1859 };
1860 let dispatch_begin = dispatch.clone();
1861 let (response, ticket, mut termination, _) = run_blocking_query(
1862 engine.clone(),
1863 Arc::clone(&parsed_query),
1864 principal.clone(),
1865 result_mode,
1866 query_timeout,
1867 query_deadline,
1868 metrics,
1869 reader,
1870 wire_read_buffer,
1871 pending_messages,
1872 move |engine, parsed_query, principal| {
1873 dispatch_begin(engine, parsed_query, principal, true)
1874 },
1875 )
1876 .await;
1877 if is_success_response(&response) {
1878 *tx_permit = Some(permit);
1879 } else if is_query_cancellation_response(&response) {
1880 *tx_permit = Some(permit);
1886 rollback_connection_transaction(engine, principal, tx_permit).await;
1887 termination = Some(ConnectionTermination::Closed);
1888 }
1889 (response, ticket, termination)
1890 }
1891 Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1892 let standalone_permit = if tx_permit.is_none() {
1893 match acquire_autocommit_permit(
1894 &tx_gate,
1895 AdmissionMode::Writer,
1896 tx_wait_timeout,
1897 metrics,
1898 )
1899 .await
1900 {
1901 Ok(permit) => Some(permit),
1902 Err(response) => return (response, None, None),
1903 }
1904 } else {
1905 None
1906 };
1907 let dispatch_commit = dispatch.clone();
1908 let (response, ticket, mut termination, _) = run_blocking_query(
1909 engine.clone(),
1910 Arc::clone(&parsed_query),
1911 principal.clone(),
1912 result_mode,
1913 query_timeout,
1914 query_deadline,
1915 metrics,
1916 reader,
1917 wire_read_buffer,
1918 pending_messages,
1919 move |engine, parsed_query, principal| {
1920 dispatch_commit(engine, parsed_query, principal, true)
1921 },
1922 )
1923 .await;
1924 if is_success_response(&response) {
1925 tx_permit.take();
1930 } else if is_query_cancellation_response(&response) {
1931 rollback_connection_transaction(engine, principal, tx_permit).await;
1932 termination = Some(ConnectionTermination::Closed);
1933 }
1934 drop(standalone_permit);
1935 (response, ticket, termination)
1936 }
1937 None if tx_permit.is_some() => {
1938 let dispatch_in_tx = dispatch.clone();
1939 let mut out = run_blocking_query(
1940 engine.clone(),
1941 Arc::clone(&parsed_query),
1942 principal.clone(),
1943 result_mode,
1944 query_timeout,
1945 query_deadline,
1946 metrics,
1947 reader,
1948 wire_read_buffer,
1949 pending_messages,
1950 move |engine, parsed_query, principal| {
1951 dispatch_in_tx(engine, parsed_query, principal, true)
1952 },
1953 )
1954 .await;
1955 if is_query_cancellation_response(&out.0) {
1956 rollback_connection_transaction(engine, principal, tx_permit).await;
1957 out.2 = Some(ConnectionTermination::Closed);
1958 }
1959 (out.0, out.1, out.2)
1960 }
1961 None => {
1962 let admission = autocommit_admission;
1963 let permit = match acquire_autocommit_permit(
1964 &tx_gate,
1965 admission,
1966 tx_wait_timeout,
1967 metrics,
1968 )
1969 .await
1970 {
1971 Ok(permit) => permit,
1972 Err(response) => return (response, None, None),
1973 };
1974 let retry_engine = Arc::clone(&engine);
1978 let retry_parsed_query = Arc::clone(&parsed_query);
1979 let retry_principal = principal.clone();
1980 let allow_readonly_escalation = admission == AdmissionMode::Writer;
1981 let dispatch_first = dispatch.clone();
1982 let mut out = run_blocking_query(
1983 engine,
1984 parsed_query,
1985 principal,
1986 result_mode,
1987 query_timeout,
1988 query_deadline,
1989 metrics,
1990 reader,
1991 wire_read_buffer,
1992 pending_messages,
1993 move |engine, parsed_query, principal| {
1994 dispatch_first(engine, parsed_query, principal, allow_readonly_escalation)
1995 },
1996 )
1997 .await;
1998 drop(permit);
1999 if out.3 {
2000 let writer_permit = match acquire_autocommit_permit(
2001 &tx_gate,
2002 AdmissionMode::Writer,
2003 tx_wait_timeout,
2004 metrics,
2005 )
2006 .await
2007 {
2008 Ok(permit) => permit,
2009 Err(response) => return (response, None, None),
2010 };
2011 let dispatch_retry = dispatch.clone();
2012 out = run_blocking_query(
2013 retry_engine,
2014 retry_parsed_query,
2015 retry_principal,
2016 result_mode,
2017 query_timeout,
2018 query_deadline,
2019 metrics,
2020 reader,
2021 wire_read_buffer,
2022 pending_messages,
2023 move |engine, parsed_query, principal| {
2024 dispatch_retry(engine, parsed_query, principal, true)
2025 },
2026 )
2027 .await;
2028 drop(writer_permit);
2029 }
2030 (out.0, out.1, out.2)
2031 }
2032 }
2033}
2034
2035#[allow(clippy::too_many_arguments)]
2038async fn execute_wire_query<R>(
2039 engine: Arc<RwLock<Engine>>,
2040 tx_gate: TxGate,
2041 tx_permit: &mut Option<OwnedSemaphorePermit>,
2042 query: String,
2043 result_mode: WireResultMode,
2044 principal: Option<Principal>,
2045 query_timeout: Duration,
2046 tx_wait_timeout: Duration,
2047 metrics: &Arc<Metrics>,
2048 reader: &mut BufReader<R>,
2049 wire_read_buffer: &mut Vec<u8>,
2050 pending_messages: &mut InFlightReadAhead,
2051) -> (
2052 Message,
2053 Option<PendingDurability>,
2054 Option<ConnectionTermination>,
2055)
2056where
2057 R: AsyncRead + Unpin,
2058{
2059 let query_deadline = Instant::now() + query_timeout;
2060 let stmt_result = parser::parse(&query).map_err(|e| e.to_string());
2065 let parsed_query = Arc::new((query, stmt_result));
2066 let tx_control = parsed_transaction_control(&parsed_query.1);
2067 let autocommit_admission = parsed_query
2068 .1
2069 .as_ref()
2070 .map(statement_admission)
2071 .unwrap_or(AdmissionMode::Writer);
2072 run_wire_query_state_machine(
2073 engine,
2074 tx_gate,
2075 tx_permit,
2076 parsed_query,
2077 tx_control,
2078 autocommit_admission,
2079 result_mode,
2080 principal,
2081 query_timeout,
2082 query_deadline,
2083 tx_wait_timeout,
2084 metrics,
2085 reader,
2086 wire_read_buffer,
2087 pending_messages,
2088 |engine,
2089 parsed_query: Arc<(String, Result<powdb_query::ast::Statement, String>)>,
2090 principal: Option<Principal>,
2091 allow| {
2092 dispatch_query_parsed(
2093 &engine,
2094 &parsed_query.0,
2095 &parsed_query.1,
2096 principal.as_ref(),
2097 allow,
2098 )
2099 },
2100 )
2101 .await
2102}
2103
2104#[allow(clippy::too_many_arguments)]
2105async fn execute_wire_query_sql<R>(
2106 engine: Arc<RwLock<Engine>>,
2107 tx_gate: TxGate,
2108 tx_permit: &mut Option<OwnedSemaphorePermit>,
2109 query: String,
2110 result_mode: WireResultMode,
2111 principal: Option<Principal>,
2112 query_timeout: Duration,
2113 tx_wait_timeout: Duration,
2114 metrics: &Arc<Metrics>,
2115 reader: &mut BufReader<R>,
2116 wire_read_buffer: &mut Vec<u8>,
2117 pending_messages: &mut InFlightReadAhead,
2118) -> (
2119 Message,
2120 Option<PendingDurability>,
2121 Option<ConnectionTermination>,
2122)
2123where
2124 R: AsyncRead + Unpin,
2125{
2126 let query_deadline = Instant::now() + query_timeout;
2127 let stmt_result = sql::parse_sql(&query).map_err(|e| e.to_string());
2128 let parsed_query = Arc::new((query, stmt_result));
2129 let tx_control = parsed_transaction_control(&parsed_query.1);
2130 let autocommit_admission = parsed_query
2131 .1
2132 .as_ref()
2133 .map(statement_admission)
2134 .unwrap_or(AdmissionMode::Writer);
2135 run_wire_query_state_machine(
2136 engine,
2137 tx_gate,
2138 tx_permit,
2139 parsed_query,
2140 tx_control,
2141 autocommit_admission,
2142 result_mode,
2143 principal,
2144 query_timeout,
2145 query_deadline,
2146 tx_wait_timeout,
2147 metrics,
2148 reader,
2149 wire_read_buffer,
2150 pending_messages,
2151 |engine,
2152 parsed_query: Arc<(String, Result<powdb_query::ast::Statement, String>)>,
2153 principal: Option<Principal>,
2154 allow| {
2155 dispatch_sql_query_parsed(
2156 &engine,
2157 &parsed_query.0,
2158 &parsed_query.1,
2159 principal.as_ref(),
2160 allow,
2161 )
2162 },
2163 )
2164 .await
2165}
2166
2167#[allow(clippy::too_many_arguments)]
2171async fn execute_wire_query_with_params<R>(
2172 engine: Arc<RwLock<Engine>>,
2173 tx_gate: TxGate,
2174 tx_permit: &mut Option<OwnedSemaphorePermit>,
2175 query: String,
2176 params: Vec<WireParam>,
2177 result_mode: WireResultMode,
2178 principal: Option<Principal>,
2179 query_timeout: Duration,
2180 tx_wait_timeout: Duration,
2181 metrics: &Arc<Metrics>,
2182 reader: &mut BufReader<R>,
2183 wire_read_buffer: &mut Vec<u8>,
2184 pending_messages: &mut InFlightReadAhead,
2185) -> (
2186 Message,
2187 Option<PendingDurability>,
2188 Option<ConnectionTermination>,
2189)
2190where
2191 R: AsyncRead + Unpin,
2192{
2193 let query_deadline = Instant::now() + query_timeout;
2194 let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
2195 let stmt_result = parser::parse_with_params(&query, &bound).map_err(|e| e.to_string());
2196 let parsed_query = Arc::new((query, bound, stmt_result));
2197 let tx_control = parsed_transaction_control(&parsed_query.2);
2198 let autocommit_admission = parsed_query
2199 .2
2200 .as_ref()
2201 .map(statement_admission)
2202 .unwrap_or(AdmissionMode::Writer);
2203 run_wire_query_state_machine(
2204 engine,
2205 tx_gate,
2206 tx_permit,
2207 parsed_query,
2208 tx_control,
2209 autocommit_admission,
2210 result_mode,
2211 principal,
2212 query_timeout,
2213 query_deadline,
2214 tx_wait_timeout,
2215 metrics,
2216 reader,
2217 wire_read_buffer,
2218 pending_messages,
2219 |engine,
2220 parsed_query: Arc<(
2221 String,
2222 Vec<powdb_query::ast::ParamValue>,
2223 Result<powdb_query::ast::Statement, String>,
2224 )>,
2225 principal: Option<Principal>,
2226 allow| {
2227 dispatch_query_with_bound_params_parsed(
2228 &engine,
2229 &parsed_query.0,
2230 &parsed_query.1,
2231 &parsed_query.2,
2232 principal.as_ref(),
2233 allow,
2234 )
2235 },
2236 )
2237 .await
2238}
2239
2240struct DeferredQueryMetric {
2245 start: Instant,
2246 outcome: QueryOutcome,
2247 exceeded_timeout: bool,
2248}
2249
2250type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
2252
2253#[allow(clippy::too_many_arguments)]
2254async fn run_blocking_query<T, F, R>(
2255 engine: Arc<RwLock<Engine>>,
2256 input: T,
2257 principal: Option<Principal>,
2258 result_mode: WireResultMode,
2259 query_timeout: Duration,
2260 query_deadline: Instant,
2261 metrics: &Arc<Metrics>,
2262 reader: &mut BufReader<R>,
2263 wire_read_buffer: &mut Vec<u8>,
2264 pending_messages: &mut InFlightReadAhead,
2265 f: F,
2266) -> (
2267 Message,
2268 Option<PendingDurability>,
2269 Option<ConnectionTermination>,
2270 bool,
2271)
2272where
2273 T: Send + 'static,
2274 F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
2275 R: AsyncRead + Unpin,
2276{
2277 let _in_flight = metrics.in_flight_guard();
2278 let start = Instant::now();
2279
2280 let timeout_ms = query_timeout.as_millis().min(u128::from(u64::MAX)) as u64;
2288 let cancel = Arc::new(powdb_query::cancel::ExecCancel::with_deadline(
2289 query_deadline,
2290 timeout_ms,
2291 ));
2292 let cancel_task = Arc::clone(&cancel);
2293 let mut handle = tokio::task::spawn_blocking(move || {
2294 let _cancel_guard = powdb_query::cancel::install(cancel_task);
2295 f(engine, input, principal)
2296 });
2297 let mut exceeded_timeout = false;
2298 let mut termination = None;
2299 let timeout = tokio::time::sleep(query_deadline.saturating_duration_since(Instant::now()));
2300 tokio::pin!(timeout);
2301 let join_result = loop {
2302 tokio::select! {
2303 result = &mut handle => break result,
2304 _ = &mut timeout => {
2305 exceeded_timeout = true;
2306 cancel.cancel(powdb_query::cancel::CancelReason::Timeout);
2311 break handle.await;
2312 }
2313 read = read_message_cancel_safe(
2314 reader,
2315 wire_read_buffer,
2316 pending_messages.remaining_bytes(),
2317 ) => {
2318 match read {
2319 Ok(Some(DecodedWireMessage { message: Message::Disconnect, .. })) => {
2320 cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2321 termination = Some(ConnectionTermination::Closed);
2322 break handle.await;
2323 }
2324 Ok(Some(frame)) => {
2325 if pending_messages.len() + 1 >= MAX_IN_FLIGHT_READ_AHEAD_FRAMES
2326 || pending_messages.wire_bytes + frame.wire_len
2327 >= MAX_IN_FLIGHT_READ_AHEAD_BYTES
2328 {
2329 cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2333 termination = Some(ConnectionTermination::ReadError);
2334 break handle.await;
2335 }
2336 pending_messages.push_back(frame);
2340 }
2341 Ok(None) => {
2342 cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2343 termination = Some(ConnectionTermination::Closed);
2344 break handle.await;
2345 }
2346 Err(_) => {
2347 cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2348 termination = Some(ConnectionTermination::ReadError);
2349 break handle.await;
2350 }
2351 }
2352 }
2353 }
2354 };
2355
2356 let (message, ticket, outcome, readonly_needs_write) = match join_result {
2357 Ok((Ok(result), ticket)) => match query_result_to_message(result, result_mode) {
2358 Ok(message) => (message, ticket, QueryOutcome::Ok, false),
2359 Err(e) => (
2360 Message::Error {
2361 message: sanitize_error(&e.to_string()),
2362 },
2363 ticket,
2364 QueryOutcome::Error,
2365 false,
2366 ),
2367 },
2368 Ok((Err(QueryError::ReadonlyNeedsWrite), ticket)) => {
2369 if exceeded_timeout {
2370 (
2371 Message::Error {
2372 message: sanitize_error(&QueryError::Timeout { timeout_ms }.to_string()),
2373 },
2374 ticket,
2375 QueryOutcome::Timeout,
2376 true,
2377 )
2378 } else {
2379 (
2380 Message::Error {
2381 message: "query execution error".into(),
2382 },
2383 ticket,
2384 QueryOutcome::Error,
2385 true,
2386 )
2387 }
2388 }
2389 Ok((Err(e), ticket)) => {
2390 if matches!(e, QueryError::Timeout { .. }) {
2394 exceeded_timeout = true;
2395 }
2396 let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
2397 QueryOutcome::MemoryLimit
2398 } else if matches!(e, QueryError::Timeout { .. }) {
2399 QueryOutcome::Timeout
2400 } else {
2401 QueryOutcome::Error
2402 };
2403 (
2404 Message::Error {
2405 message: sanitize_error(&e.to_string()),
2406 },
2407 ticket,
2408 outcome,
2409 false,
2410 )
2411 }
2412 Err(e) => (
2413 Message::Error {
2414 message: format!("internal error: {e}"),
2415 },
2416 None,
2417 QueryOutcome::Error,
2418 false,
2419 ),
2420 };
2421 let readonly_retry = readonly_needs_write && !exceeded_timeout && termination.is_none();
2422 match ticket {
2423 Some(ticket) => (
2428 message,
2429 Some((
2430 ticket,
2431 DeferredQueryMetric {
2432 start,
2433 outcome,
2434 exceeded_timeout,
2435 },
2436 )),
2437 termination,
2438 readonly_retry,
2439 ),
2440 None => {
2441 if !readonly_retry {
2442 if exceeded_timeout {
2443 metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
2444 } else {
2445 metrics.record_query(start.elapsed(), outcome);
2446 }
2447 }
2448 (message, None, termination, readonly_retry)
2449 }
2450 }
2451}
2452
2453async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
2461 match tokio::task::spawn_blocking(move || ticket.wait()).await {
2462 Ok(Ok(())) => None,
2463 Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
2464 Err(e) => Some(format!("internal error: {e}")),
2465 }
2466}
2467
2468fn is_success_response(msg: &Message) -> bool {
2469 matches!(
2470 msg,
2471 Message::ResultRows { .. }
2472 | Message::ResultScalar { .. }
2473 | Message::ResultRowsNative { .. }
2474 | Message::ResultScalarNative { .. }
2475 | Message::ResultOk { .. }
2476 | Message::ResultMessage { .. }
2477 )
2478}
2479
2480fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
2481 let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref(), true);
2482 let _ = res;
2483 if let Some(ticket) = ticket {
2486 let _ = ticket.wait();
2487 }
2488}
2489
2490fn is_query_cancellation_response(message: &Message) -> bool {
2491 matches!(
2492 message,
2493 Message::Error { message }
2494 if message.starts_with("query timeout after")
2495 || message == "query cancelled by client disconnect"
2496 )
2497}
2498
2499async fn rollback_connection_transaction(
2505 engine: Arc<RwLock<Engine>>,
2506 principal: Option<Principal>,
2507 tx_permit: &mut Option<OwnedSemaphorePermit>,
2508) {
2509 if tx_permit.is_none() {
2510 return;
2511 }
2512 let _ = tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
2513 tx_permit.take();
2514}
2515
2516struct DecodedWireMessage {
2525 message: Message,
2526 wire_len: usize,
2527}
2528
2529#[derive(Default)]
2530struct InFlightReadAhead {
2531 frames: VecDeque<DecodedWireMessage>,
2532 wire_bytes: usize,
2533}
2534
2535impl InFlightReadAhead {
2536 fn len(&self) -> usize {
2537 self.frames.len()
2538 }
2539
2540 fn is_empty(&self) -> bool {
2541 self.frames.is_empty()
2542 }
2543
2544 fn remaining_bytes(&self) -> usize {
2545 MAX_IN_FLIGHT_READ_AHEAD_BYTES.saturating_sub(self.wire_bytes)
2546 }
2547
2548 fn push_back(&mut self, frame: DecodedWireMessage) {
2549 self.wire_bytes += frame.wire_len;
2550 self.frames.push_back(frame);
2551 }
2552
2553 fn pop_front(&mut self) -> Option<Message> {
2554 let frame = self.frames.pop_front()?;
2555 self.wire_bytes -= frame.wire_len;
2556 Some(frame.message)
2557 }
2558}
2559
2560async fn read_message_cancel_safe<R>(
2561 reader: &mut BufReader<R>,
2562 buffered: &mut Vec<u8>,
2563 max_frame_len: usize,
2564) -> std::io::Result<Option<DecodedWireMessage>>
2565where
2566 R: AsyncRead + Unpin,
2567{
2568 loop {
2569 if buffered.len() >= 6 {
2570 let payload_len = u32::from_le_bytes(
2571 buffered[2..6]
2572 .try_into()
2573 .expect("four-byte wire payload length"),
2574 ) as usize;
2575 if payload_len > MAX_WIRE_PAYLOAD_SIZE {
2576 return Err(std::io::Error::new(
2577 std::io::ErrorKind::InvalidData,
2578 format!("payload too large: {payload_len} bytes (max {MAX_WIRE_PAYLOAD_SIZE})"),
2579 ));
2580 }
2581 let frame_len = 6usize.checked_add(payload_len).ok_or_else(|| {
2582 std::io::Error::new(
2583 std::io::ErrorKind::InvalidData,
2584 "wire frame length overflow",
2585 )
2586 })?;
2587 if frame_len > max_frame_len {
2588 return Err(std::io::Error::new(
2589 std::io::ErrorKind::InvalidData,
2590 format!(
2591 "wire frame exceeds the available in-flight read-ahead budget: \
2592 {frame_len} bytes (available {max_frame_len})"
2593 ),
2594 ));
2595 }
2596 if buffered.len() >= frame_len {
2597 let frame: Vec<u8> = buffered.drain(..frame_len).collect();
2598 return Message::decode(&frame)
2599 .map(|message| {
2600 Some(DecodedWireMessage {
2601 message,
2602 wire_len: frame_len,
2603 })
2604 })
2605 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e));
2606 }
2607 }
2608
2609 let mut chunk = [0u8; 8192];
2610 let wanted = if buffered.len() < 6 {
2614 6 - buffered.len()
2615 } else {
2616 let payload_len = u32::from_le_bytes(
2617 buffered[2..6]
2618 .try_into()
2619 .expect("four-byte wire payload length"),
2620 ) as usize;
2621 6usize
2622 .checked_add(payload_len)
2623 .and_then(|frame_len| frame_len.checked_sub(buffered.len()))
2624 .ok_or_else(|| {
2625 std::io::Error::new(
2626 std::io::ErrorKind::InvalidData,
2627 "invalid buffered wire frame length",
2628 )
2629 })?
2630 };
2631 let read_limit = wanted.min(chunk.len());
2632 let read = reader.read(&mut chunk[..read_limit]).await?;
2633 if read == 0 {
2634 if buffered.len() < 6 {
2635 buffered.clear();
2638 return Ok(None);
2639 }
2640 return Err(std::io::Error::new(
2641 std::io::ErrorKind::UnexpectedEof,
2642 "connection closed in the middle of a wire frame",
2643 ));
2644 }
2645 buffered.extend_from_slice(&chunk[..read]);
2646 }
2647}
2648
2649#[derive(Clone, Copy, Debug)]
2650enum ConnectionTermination {
2651 Closed,
2652 ReadError,
2653}
2654
2655pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
2656where
2657 S: AsyncRead + AsyncWrite + Unpin,
2658{
2659 let ConnOpts {
2660 engine,
2661 tx_gate,
2662 expected_password,
2663 users,
2664 shutdown_rx,
2665 idle_timeout,
2666 query_timeout,
2667 rate_limiter,
2668 peer_addr,
2669 metrics,
2670 tx_wait_timeout,
2671 db_name: server_db_name,
2672 } = opts;
2673
2674 let peer = peer_addr
2675 .map(|a| a.to_string())
2676 .unwrap_or_else(|| "unknown".into());
2677 let peer_ip = peer_addr.map(|a| a.ip());
2678
2679 let (reader, writer) = tokio::io::split(stream);
2680 let mut reader = BufReader::new(reader);
2681 let mut writer = BufWriter::new(writer);
2682
2683 let connect_msg = loop {
2688 match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
2689 Ok(Ok(Some(Message::Ping))) => {
2690 debug!(peer = %peer, "pre-auth ping");
2691 if !write_msg(&mut writer, &Message::Pong).await {
2692 return;
2693 }
2694 continue;
2695 }
2696 Ok(Ok(Some(msg))) => break msg,
2697 Ok(Ok(None)) => {
2698 debug!(peer = %peer, "client closed before CONNECT");
2699 return;
2700 }
2701 Ok(Err(e)) => {
2702 error!(peer = %peer, error = %e, "error reading CONNECT");
2703 return;
2704 }
2705 Err(_) => {
2706 warn!(peer = %peer, "idle timeout waiting for CONNECT");
2707 return;
2708 }
2709 }
2710 };
2711
2712 let principal: Option<Principal>;
2715 let credential_auth_configured = !users.is_empty() || expected_password.is_some();
2716 match connect_msg {
2717 Message::Connect {
2718 db_name,
2719 password,
2720 username,
2721 } => {
2722 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2724 if is_rate_limited(limiter, ip) {
2725 warn!(peer = %peer, "rate limited: too many auth failures");
2726 let err = Message::Error {
2727 message: "too many auth failures, try again later".into(),
2728 };
2729 write_msg(&mut writer, &err).await;
2730 return;
2731 }
2732 }
2733
2734 let outcome = authenticate_connect(
2735 &users,
2736 expected_password.as_ref().map(|p| p.as_str()),
2737 username.as_deref(),
2738 password.as_ref().map(|p| p.as_str()),
2739 );
2740
2741 match outcome {
2742 AuthOutcome::Rejected => {
2743 warn!(peer = %peer, db = %db_name, "auth rejected");
2744 metrics.inc_auth_failure();
2745 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2747 record_auth_failure(limiter, ip);
2748 }
2749 let err = Message::Error {
2750 message: "authentication failed".into(),
2751 };
2752 write_msg(&mut writer, &err).await;
2753 return;
2754 }
2755 AuthOutcome::Authenticated {
2756 principal: auth_principal,
2757 } => {
2758 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2760 clear_auth_failures(limiter, ip);
2761 }
2762 match &auth_principal {
2763 Some(p) => {
2764 info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
2765 }
2766 None => {
2767 info!(peer = %peer, db = %db_name, "client connected");
2768 }
2769 }
2770 principal = auth_principal;
2771 }
2772 }
2773
2774 match check_db_name(server_db_name.as_deref(), &db_name) {
2781 Ok(()) => {
2782 if server_db_name.is_none() && !db_name.is_empty() && db_name != DEFAULT_DB_NAME
2783 {
2784 static NAMED_DB_WARNED: std::sync::atomic::AtomicBool =
2785 std::sync::atomic::AtomicBool::new(false);
2786 if !NAMED_DB_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
2787 warn!(
2788 peer = %peer, db = %db_name,
2789 "client requested a named database but this server serves a single global database; name ignored"
2790 );
2791 }
2792 }
2793 }
2794 Err(msg) => {
2795 warn!(peer = %peer, db = %db_name, "rejected: unknown database");
2796 let err = Message::Error { message: msg };
2797 write_msg(&mut writer, &err).await;
2798 return;
2799 }
2800 }
2801
2802 let ok = Message::ConnectOk {
2803 version: env!("CARGO_PKG_VERSION").into(),
2804 };
2805 if !write_msg(&mut writer, &ok).await {
2806 return;
2807 }
2808 }
2809 _ => {
2810 warn!(peer = %peer, "first message was not CONNECT");
2811 let err = Message::Error {
2812 message: "expected CONNECT".into(),
2813 };
2814 write_msg(&mut writer, &err).await;
2815 return;
2816 }
2817 }
2818
2819 let mut tx_permit: Option<OwnedSemaphorePermit> = None;
2820 let mut wire_read_buffer = Vec::new();
2824 let mut pending_messages = InFlightReadAhead::default();
2825 let mut carry: Option<Message> = None;
2828
2829 'conn: loop {
2831 let msg = if let Some(m) = carry.take() {
2832 m
2833 } else if let Some(m) = pending_messages.pop_front() {
2834 m
2835 } else {
2836 tokio::select! {
2837 result = tokio::time::timeout(
2839 idle_timeout,
2840 read_message_cancel_safe(
2841 &mut reader,
2842 &mut wire_read_buffer,
2843 MAX_WIRE_PAYLOAD_SIZE + 6,
2844 ),
2845 ) => {
2846 match result {
2847 Ok(Ok(Some(frame))) => frame.message,
2848 Ok(Ok(None)) => break,
2849 Ok(Err(e)) => {
2850 error!(peer = %peer, error = %e, "read error");
2851 break;
2852 }
2853 Err(_) => {
2854 info!(peer = %peer, "idle timeout, closing connection");
2855 let err = Message::Error { message: "idle timeout".into() };
2856 write_msg(&mut writer, &err).await;
2857 break;
2858 }
2859 }
2860 }
2861 _ = shutdown_rx.changed() => {
2863 if *shutdown_rx.borrow() {
2864 info!(peer = %peer, "server shutting down, closing connection");
2865 let err = Message::Error { message: "server shutting down".into() };
2866 write_msg(&mut writer, &err).await;
2867 break;
2868 }
2869 continue;
2870 }
2871 }
2872 };
2873
2874 if matches!(
2880 msg,
2881 Message::Query { .. }
2882 | Message::QuerySql { .. }
2883 | Message::QueryWithParams { .. }
2884 | Message::QueryNative { .. }
2885 | Message::QuerySqlNative { .. }
2886 | Message::QueryWithParamsNative { .. }
2887 ) {
2888 const MAX_PIPELINE_BATCH: usize = 128;
2891 const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
2896
2897 fn approx_response_bytes(msg: &Message) -> usize {
2901 match msg {
2902 Message::ResultRows { columns, rows } => {
2903 columns.iter().map(|c| c.len() + 4).sum::<usize>()
2904 + rows
2905 .iter()
2906 .map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
2907 .sum::<usize>()
2908 }
2909 Message::ResultScalar { value } => value.len(),
2910 Message::ResultRowsNative { columns, rows } => {
2911 columns.iter().map(|c| c.len() + 4).sum::<usize>()
2912 + rows
2913 .iter()
2914 .map(|row| {
2915 row.iter()
2916 .map(|value| 5 + native_value_body_len(value))
2917 .sum::<usize>()
2918 })
2919 .sum::<usize>()
2920 }
2921 Message::ResultScalarNative { value } => 5 + native_value_body_len(value),
2922 Message::ResultMessage { message } | Message::Error { message } => {
2923 message.len()
2924 }
2925 _ => 16,
2926 }
2927 }
2928
2929 fn complete_frame_buffered(buf: &[u8]) -> bool {
2935 buf.len() >= 6 && {
2936 let payload_len =
2937 u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
2938 buf.len() - 6 >= payload_len
2939 }
2940 }
2941
2942 let mut responses: Vec<Message> = Vec::new();
2943 let mut response_bytes: usize = 0;
2944 let mut last_ticket: Option<WalDurabilityTicket> = None;
2945 let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
2946 let mut fatal: Option<ConnectionTermination> = None;
2947 let mut current = msg;
2948 loop {
2949 let (response, ticket, termination) = match current {
2950 Message::Query { query } => {
2951 if query.len() > MAX_QUERY_LENGTH {
2952 (
2953 Message::Error {
2954 message: format!(
2955 "query too large: {} bytes (max {})",
2956 query.len(),
2957 MAX_QUERY_LENGTH
2958 ),
2959 },
2960 None,
2961 None,
2962 )
2963 } else {
2964 debug!(peer = %peer, query = %query, "received query");
2965 execute_wire_query(
2966 engine.clone(),
2967 tx_gate.clone(),
2968 &mut tx_permit,
2969 query,
2970 WireResultMode::LegacyText,
2971 principal.clone(),
2972 query_timeout,
2973 tx_wait_timeout,
2974 &metrics,
2975 &mut reader,
2976 &mut wire_read_buffer,
2977 &mut pending_messages,
2978 )
2979 .await
2980 }
2981 }
2982 Message::QuerySql { query } => {
2983 if query.len() > MAX_QUERY_LENGTH {
2984 (
2985 Message::Error {
2986 message: format!(
2987 "query too large: {} bytes (max {})",
2988 query.len(),
2989 MAX_QUERY_LENGTH
2990 ),
2991 },
2992 None,
2993 None,
2994 )
2995 } else {
2996 debug!(peer = %peer, query = %query, "received SQL query");
2997 execute_wire_query_sql(
2998 engine.clone(),
2999 tx_gate.clone(),
3000 &mut tx_permit,
3001 query,
3002 WireResultMode::LegacyText,
3003 principal.clone(),
3004 query_timeout,
3005 tx_wait_timeout,
3006 &metrics,
3007 &mut reader,
3008 &mut wire_read_buffer,
3009 &mut pending_messages,
3010 )
3011 .await
3012 }
3013 }
3014 Message::QueryWithParams { query, params } => {
3015 if query.len() > MAX_QUERY_LENGTH {
3016 (
3017 Message::Error {
3018 message: format!(
3019 "query too large: {} bytes (max {})",
3020 query.len(),
3021 MAX_QUERY_LENGTH
3022 ),
3023 },
3024 None,
3025 None,
3026 )
3027 } else {
3028 debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
3029 execute_wire_query_with_params(
3030 engine.clone(),
3031 tx_gate.clone(),
3032 &mut tx_permit,
3033 query,
3034 params,
3035 WireResultMode::LegacyText,
3036 principal.clone(),
3037 query_timeout,
3038 tx_wait_timeout,
3039 &metrics,
3040 &mut reader,
3041 &mut wire_read_buffer,
3042 &mut pending_messages,
3043 )
3044 .await
3045 }
3046 }
3047 Message::QueryNative { query } => {
3048 if query.len() > MAX_QUERY_LENGTH {
3049 (
3050 Message::Error {
3051 message: format!(
3052 "query too large: {} bytes (max {})",
3053 query.len(),
3054 MAX_QUERY_LENGTH
3055 ),
3056 },
3057 None,
3058 None,
3059 )
3060 } else {
3061 debug!(peer = %peer, query = %query, "received native query");
3062 execute_wire_query(
3063 engine.clone(),
3064 tx_gate.clone(),
3065 &mut tx_permit,
3066 query,
3067 WireResultMode::Native,
3068 principal.clone(),
3069 query_timeout,
3070 tx_wait_timeout,
3071 &metrics,
3072 &mut reader,
3073 &mut wire_read_buffer,
3074 &mut pending_messages,
3075 )
3076 .await
3077 }
3078 }
3079 Message::QuerySqlNative { query } => {
3080 if query.len() > MAX_QUERY_LENGTH {
3081 (
3082 Message::Error {
3083 message: format!(
3084 "query too large: {} bytes (max {})",
3085 query.len(),
3086 MAX_QUERY_LENGTH
3087 ),
3088 },
3089 None,
3090 None,
3091 )
3092 } else {
3093 debug!(peer = %peer, query = %query, "received native SQL query");
3094 execute_wire_query_sql(
3095 engine.clone(),
3096 tx_gate.clone(),
3097 &mut tx_permit,
3098 query,
3099 WireResultMode::Native,
3100 principal.clone(),
3101 query_timeout,
3102 tx_wait_timeout,
3103 &metrics,
3104 &mut reader,
3105 &mut wire_read_buffer,
3106 &mut pending_messages,
3107 )
3108 .await
3109 }
3110 }
3111 Message::QueryWithParamsNative { query, params } => {
3112 if query.len() > MAX_QUERY_LENGTH {
3113 (
3114 Message::Error {
3115 message: format!(
3116 "query too large: {} bytes (max {})",
3117 query.len(),
3118 MAX_QUERY_LENGTH
3119 ),
3120 },
3121 None,
3122 None,
3123 )
3124 } else {
3125 debug!(peer = %peer, query = %query, n_params = params.len(), "received native parameterized query");
3126 execute_wire_query_with_params(
3127 engine.clone(),
3128 tx_gate.clone(),
3129 &mut tx_permit,
3130 query,
3131 params,
3132 WireResultMode::Native,
3133 principal.clone(),
3134 query_timeout,
3135 tx_wait_timeout,
3136 &metrics,
3137 &mut reader,
3138 &mut wire_read_buffer,
3139 &mut pending_messages,
3140 )
3141 .await
3142 }
3143 }
3144 _ => unreachable!("batch loop only receives plain query frames"),
3145 };
3146 if let Some((t, m)) = ticket {
3147 last_ticket = Some(t);
3151 deferred_metrics.push(m);
3152 }
3153 response_bytes += approx_response_bytes(&response);
3154 responses.push(response);
3155 if let Some(reason) = termination {
3156 fatal = Some(reason);
3157 break;
3158 }
3159
3160 if tx_permit.is_some()
3166 || responses.len() >= MAX_PIPELINE_BATCH
3167 || response_bytes >= MAX_PIPELINE_BATCH_BYTES
3168 || (pending_messages.is_empty()
3169 && !complete_frame_buffered(&wire_read_buffer)
3170 && !complete_frame_buffered(reader.buffer()))
3171 {
3172 break;
3173 }
3174 let next_message = if let Some(message) = pending_messages.pop_front() {
3177 Ok(Some(message))
3178 } else {
3179 tokio::time::timeout(
3180 idle_timeout,
3181 read_message_cancel_safe(
3182 &mut reader,
3183 &mut wire_read_buffer,
3184 MAX_WIRE_PAYLOAD_SIZE + 6,
3185 ),
3186 )
3187 .await
3188 .map(|result| result.map(|frame| frame.map(|frame| frame.message)))
3189 .unwrap_or_else(|_| {
3190 Err(std::io::Error::new(
3191 std::io::ErrorKind::TimedOut,
3192 "timeout decoding fully-buffered frame",
3193 ))
3194 })
3195 };
3196 match next_message {
3197 Ok(Some(
3198 next @ (Message::Query { .. }
3199 | Message::QuerySql { .. }
3200 | Message::QueryWithParams { .. }
3201 | Message::QueryNative { .. }
3202 | Message::QuerySqlNative { .. }
3203 | Message::QueryWithParamsNative { .. }),
3204 )) => {
3205 if tx_gate.available_permits() == 0 {
3213 carry = Some(next);
3214 break;
3215 }
3216 current = next;
3217 }
3218 Ok(Some(other)) => {
3219 carry = Some(other);
3222 break;
3223 }
3224 Ok(None) => {
3225 fatal = Some(ConnectionTermination::Closed);
3226 break;
3227 }
3228 Err(e) => {
3229 error!(peer = %peer, error = %e, "read error");
3230 fatal = Some(ConnectionTermination::ReadError);
3231 break;
3232 }
3233 }
3234 }
3235
3236 let mut durability_failed = false;
3240 if let Some(ticket) = last_ticket {
3241 if let Some(message) = settle_durability_ticket(ticket).await {
3242 durability_failed = true;
3245 for r in responses.iter_mut() {
3246 if is_success_response(r) {
3247 *r = Message::Error {
3248 message: message.clone(),
3249 };
3250 }
3251 }
3252 }
3253 }
3254 for m in deferred_metrics.drain(..) {
3255 let outcome = if m.exceeded_timeout {
3256 QueryOutcome::Timeout
3257 } else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
3258 QueryOutcome::Error
3259 } else {
3260 m.outcome
3261 };
3262 metrics.record_query(m.start.elapsed(), outcome);
3263 }
3264
3265 for r in &responses {
3266 if !write_msg(&mut writer, r).await {
3267 break 'conn;
3268 }
3269 }
3270 match fatal {
3271 None => continue,
3272 Some(ConnectionTermination::Closed | ConnectionTermination::ReadError) => break,
3273 }
3274 }
3275
3276 let response = match msg {
3277 Message::Ping => {
3278 debug!(peer = %peer, "ping");
3279 Message::Pong
3280 }
3281 Message::SyncStatus { replica_id } => {
3282 let engine = engine.clone();
3283 let principal = principal.clone();
3284 let log_context = SyncLogContext::status(&replica_id);
3285 execute_gated_sync(
3286 SyncExecutionContext {
3287 tx_gate: tx_gate.clone(),
3288 connection_has_transaction: tx_permit.is_some(),
3289 operation: SyncOperation::Status,
3290 log_context,
3291 metrics: &metrics,
3292 query_timeout,
3293 },
3294 (engine, replica_id, credential_auth_configured, principal),
3295 |(engine, replica_id, credential_authenticated, principal)| {
3296 dispatch_sync_status_decision(
3297 &engine,
3298 replica_id,
3299 credential_authenticated,
3300 principal.as_ref(),
3301 )
3302 },
3303 )
3304 .await
3305 }
3306 Message::SyncPull {
3307 replica_id,
3308 since_lsn,
3309 max_units,
3310 max_bytes,
3311 database_id,
3312 primary_generation,
3313 wal_format_version,
3314 catalog_version,
3315 segment_format_version,
3316 } => {
3317 let engine = engine.clone();
3318 let principal = principal.clone();
3319 let request = SyncPullRequest {
3320 replica_id,
3321 since_lsn,
3322 max_units,
3323 max_bytes,
3324 database_id,
3325 primary_generation,
3326 wal_format_version,
3327 catalog_version,
3328 segment_format_version,
3329 };
3330 let log_context = SyncLogContext::pull(&request);
3331 execute_gated_sync(
3332 SyncExecutionContext {
3333 tx_gate: tx_gate.clone(),
3334 connection_has_transaction: tx_permit.is_some(),
3335 operation: SyncOperation::Pull,
3336 log_context,
3337 metrics: &metrics,
3338 query_timeout,
3339 },
3340 (engine, request, credential_auth_configured, principal),
3341 |(engine, request, credential_authenticated, principal)| {
3342 dispatch_sync_pull_decision(
3343 &engine,
3344 request,
3345 credential_authenticated,
3346 principal.as_ref(),
3347 )
3348 },
3349 )
3350 .await
3351 }
3352 Message::SyncAck {
3353 replica_id,
3354 applied_lsn,
3355 remote_lsn,
3356 } => {
3357 let engine = engine.clone();
3358 let principal = principal.clone();
3359 let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
3360 execute_gated_sync(
3361 SyncExecutionContext {
3362 tx_gate: tx_gate.clone(),
3363 connection_has_transaction: tx_permit.is_some(),
3364 operation: SyncOperation::Ack,
3365 log_context,
3366 metrics: &metrics,
3367 query_timeout,
3368 },
3369 (
3370 engine,
3371 replica_id,
3372 applied_lsn,
3373 remote_lsn,
3374 credential_auth_configured,
3375 principal,
3376 ),
3377 |(
3378 engine,
3379 replica_id,
3380 applied_lsn,
3381 observed_remote_lsn,
3382 credential_authenticated,
3383 principal,
3384 )| {
3385 dispatch_sync_ack_decision(
3386 &engine,
3387 replica_id,
3388 applied_lsn,
3389 observed_remote_lsn,
3390 credential_authenticated,
3391 principal.as_ref(),
3392 )
3393 },
3394 )
3395 .await
3396 }
3397 Message::Disconnect => {
3398 debug!(peer = %peer, "received DISCONNECT");
3399 break;
3400 }
3401 _ => Message::Error {
3402 message: "unexpected message type".into(),
3403 },
3404 };
3405
3406 if !write_msg(&mut writer, &response).await {
3407 break;
3408 }
3409 }
3410
3411 if tx_permit.is_some() {
3418 let engine = engine.clone();
3419 let principal = principal.clone();
3420 let _ =
3421 tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
3422 }
3423 tx_permit.take();
3424
3425 info!(peer = %peer, "client disconnected");
3426}
3427
3428fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
3429 *total = total.saturating_add(bytes);
3430 if *total > MAX_RESPONSE_PAYLOAD_SIZE {
3431 return Err(QueryError::Execution(format!(
3432 "result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
3433 MAX_RESPONSE_PAYLOAD_SIZE
3434 )));
3435 }
3436 Ok(())
3437}
3438
3439fn native_value_body_len(value: &Value) -> usize {
3440 match value {
3441 Value::Empty => 0,
3442 Value::Int(_) | Value::Float(_) | Value::DateTime(_) => 8,
3443 Value::Bool(_) => 1,
3444 Value::Str(value) => value.len(),
3445 Value::Uuid(_) => 16,
3446 Value::Bytes(value) => value.len(),
3447 Value::Json(value) => value.len(),
3448 }
3449}
3450
3451fn query_result_to_message(
3452 result: QueryResult,
3453 result_mode: WireResultMode,
3454) -> Result<Message, QueryError> {
3455 match result {
3456 QueryResult::Rows { columns, rows } => {
3457 let mut encoded_bytes = 2usize; for col in &columns {
3459 charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
3460 }
3461 charge_response_bytes(&mut encoded_bytes, 4)?; match result_mode {
3464 WireResultMode::Native => {
3465 for row in &rows {
3466 for value in row {
3467 charge_response_bytes(
3468 &mut encoded_bytes,
3469 5 + native_value_body_len(value),
3470 )?;
3471 }
3472 }
3473 Ok(Message::ResultRowsNative { columns, rows })
3474 }
3475 WireResultMode::LegacyText => {
3476 let mut str_rows = Vec::with_capacity(rows.len());
3477 for row in rows {
3478 let mut str_row = Vec::with_capacity(row.len());
3479 for value in row {
3480 let display = value_to_display(&value);
3481 charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
3482 str_row.push(display);
3483 }
3484 str_rows.push(str_row);
3485 }
3486 Ok(Message::ResultRows {
3487 columns,
3488 rows: str_rows,
3489 })
3490 }
3491 }
3492 }
3493 QueryResult::Scalar(value) => match result_mode {
3494 WireResultMode::Native => {
3495 let mut encoded_bytes = 0;
3496 charge_response_bytes(&mut encoded_bytes, 5 + native_value_body_len(&value))?;
3497 Ok(Message::ResultScalarNative { value })
3498 }
3499 WireResultMode::LegacyText => Ok(Message::ResultScalar {
3500 value: value_to_display(&value),
3501 }),
3502 },
3503 QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
3504 QueryResult::Created(name) => Ok(Message::ResultMessage {
3505 message: format!("type {name} created"),
3506 }),
3507 QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
3508 }
3509}
3510
3511fn value_to_display(v: &Value) -> String {
3515 v.to_wire_string()
3516}
3517
3518#[cfg(test)]
3519mod tests {
3520 use super::*;
3521 use powdb_storage::wal::WalRecordType;
3522 use powdb_sync::{
3523 write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
3524 ReplicaCursor, RetainedSegment, RetainedUnit,
3525 };
3526
3527 #[test]
3530 fn null_serializes_as_null_bareword_on_wire() {
3531 assert_eq!(value_to_display(&Value::Empty), "null");
3532 }
3533
3534 #[test]
3537 fn unique_violation_error_surfaces_to_remote_clients() {
3538 assert_eq!(
3541 sanitize_error("unique constraint violation on User.email"),
3542 "unique constraint violation on User.email"
3543 );
3544 }
3545
3546 #[test]
3547 fn internal_errors_stay_generic() {
3548 assert_eq!(
3549 sanitize_error("some internal io panic detail"),
3550 "query execution error"
3551 );
3552 }
3553
3554 #[test]
3555 fn cancellation_errors_surface_to_remote_clients() {
3556 for msg in [
3560 &QueryError::Timeout { timeout_ms: 2000 }.to_string(),
3561 &QueryError::Cancelled.to_string(),
3562 ] {
3563 assert_eq!(sanitize_error(msg), *msg, "should pass through verbatim");
3564 }
3565 assert_eq!(
3567 QueryError::Timeout { timeout_ms: 2000 }.to_string(),
3568 "query timeout after 2000ms"
3569 );
3570 assert_eq!(
3571 QueryError::Cancelled.to_string(),
3572 "query cancelled by client disconnect"
3573 );
3574 }
3575
3576 #[test]
3579 fn json_cell_renders_canonical_text_on_wire() {
3580 let pj1 = powdb_storage::pj1::parse_json_text(r#"{"b":2,"a":1,"nested":{"z":true}}"#)
3585 .expect("valid JSON");
3586 let result = QueryResult::Rows {
3587 columns: vec!["doc".into()],
3588 rows: vec![vec![Value::Json(pj1.into())]],
3589 };
3590 match query_result_to_message(result, WireResultMode::LegacyText).expect("encodes") {
3591 Message::ResultRows { columns, rows } => {
3592 assert_eq!(columns, vec!["doc"]);
3593 assert_eq!(
3594 rows,
3595 vec![vec![r#"{"a":1,"b":2,"nested":{"z":true}}"#.to_string()]]
3596 );
3597 }
3598 other => panic!("expected ResultRows, got {other:?}"),
3599 }
3600 }
3601
3602 #[test]
3603 fn json_parse_error_surfaces_to_remote_clients() {
3604 for msg in [
3613 "type mismatch: invalid JSON: unexpected character 'x' at position 3",
3614 "invalid JSON: nesting exceeds depth cap 128",
3615 ] {
3616 assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
3617 }
3618 assert_eq!(
3619 sanitize_error("malformed PJ1: truncated"),
3620 "query execution error",
3621 "internal storage corruption must stay masked"
3622 );
3623 }
3624
3625 #[test]
3631 fn describe_shows_json_type_over_the_wire() {
3632 let dir = tempfile::tempdir().unwrap();
3633 let mut engine = Engine::new(dir.path()).unwrap();
3634 engine
3635 .execute_powql("type Doc { required id: int, body: json }")
3636 .expect("json column DDL should be accepted once Lane B lands");
3637 let result = engine.execute_powql("describe Doc").expect("describe runs");
3638 let msg = query_result_to_message(result, WireResultMode::LegacyText).expect("encodes");
3639 match msg {
3640 Message::ResultRows { columns, rows } => {
3641 assert_eq!(columns[1], "type");
3642 let body = rows
3644 .iter()
3645 .find(|r| r[0] == "body")
3646 .expect("body column present");
3647 assert_eq!(body[1], "json");
3648 }
3649 other => panic!("expected ResultRows, got {other:?}"),
3650 }
3651 }
3652
3653 #[test]
3656 fn db_name_unpinned_accepts_any_name() {
3657 for requested in ["", "default", "prod", "anything"] {
3658 assert!(
3659 check_db_name(None, requested).is_ok(),
3660 "rejected {requested}"
3661 );
3662 }
3663 }
3664
3665 #[test]
3666 fn db_name_pinned_accepts_match_empty_and_default_sentinel() {
3667 assert!(check_db_name(Some("prod"), "prod").is_ok());
3670 assert!(check_db_name(Some("prod"), "").is_ok());
3671 assert!(check_db_name(Some("prod"), DEFAULT_DB_NAME).is_ok());
3672 }
3673
3674 #[test]
3675 fn db_name_pinned_rejects_foreign_with_clear_message() {
3676 let err = check_db_name(Some("prod"), "staging").unwrap_err();
3677 assert_eq!(err, "unknown database 'staging'; this server serves 'prod'");
3678 }
3679
3680 #[tokio::test]
3683 async fn begin_permit_acquires_when_gate_is_free() {
3684 let gate = new_tx_gate();
3685 let metrics = Arc::new(Metrics::new());
3686 let permit = acquire_begin_permit(&gate, Duration::from_secs(5), &metrics)
3687 .await
3688 .expect("should acquire a free gate");
3689 assert_eq!(gate.available_permits(), 0, "permit must be held");
3690 drop(permit);
3691 assert_eq!(
3692 gate.available_permits(),
3693 DEFAULT_TX_GATE_READER_PERMITS as usize,
3694 "permit pool must release on drop"
3695 );
3696 }
3697
3698 #[tokio::test]
3699 async fn begin_permit_times_out_with_clear_error_and_truthful_metric() {
3700 let gate = new_tx_gate();
3701 let metrics = Arc::new(Metrics::new());
3702 let _held = gate
3704 .clone()
3705 .acquire_many_owned(DEFAULT_TX_GATE_READER_PERMITS)
3706 .await
3707 .unwrap();
3708 let err = acquire_begin_permit(&gate, Duration::from_millis(25), &metrics)
3709 .await
3710 .expect_err("must time out while the gate is held");
3711 match err {
3712 Message::Error { message } => {
3713 assert!(
3714 message.contains("transaction gate timeout after 25ms"),
3715 "unexpected message: {message}"
3716 );
3717 assert!(
3718 message.contains("waiting for concurrent transaction"),
3719 "unexpected message: {message}"
3720 );
3721 }
3722 other => panic!("expected Error, got {other:?}"),
3723 }
3724 let rendered = metrics.render();
3725 assert!(rendered.contains("powdb_tx_gate_timeouts_total 1"));
3726 assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
3728 }
3729
3730 #[test]
3731 fn admission_classification_has_query_shape_parity_and_fails_closed() {
3732 assert_eq!(
3733 classify_query_admission("User filter .id = 1"),
3734 AdmissionMode::Reader
3735 );
3736 assert_eq!(
3737 classify_sql_admission("SELECT * FROM User WHERE id = 1"),
3738 AdmissionMode::Reader
3739 );
3740 assert_eq!(
3741 classify_params_admission("User filter .id = $1", &[WireParam::Int(1)]),
3742 AdmissionMode::Reader
3743 );
3744
3745 assert_eq!(
3746 classify_query_admission("insert User { id := 1 }"),
3747 AdmissionMode::Writer
3748 );
3749 assert_eq!(
3750 classify_sql_admission("INSERT INTO User (id) VALUES (1)"),
3751 AdmissionMode::Writer
3752 );
3753 assert_eq!(
3754 classify_params_admission("insert User { id := $1 }", &[WireParam::Int(1)]),
3755 AdmissionMode::Writer
3756 );
3757 assert_eq!(
3758 classify_query_admission("this is not valid PowQL"),
3759 AdmissionMode::Writer,
3760 "uncertain statements must never enter through reader admission"
3761 );
3762 }
3763
3764 #[tokio::test]
3765 async fn writer_admission_excludes_readers() {
3766 let gate = new_tx_gate();
3767 let metrics = Arc::new(Metrics::new());
3768 let writer = acquire_begin_permit(&gate, Duration::from_secs(1), &metrics)
3769 .await
3770 .expect("writer admission");
3771 let blocked = acquire_autocommit_permit(
3772 &gate,
3773 AdmissionMode::Reader,
3774 Duration::from_millis(25),
3775 &metrics,
3776 )
3777 .await;
3778 assert!(blocked.is_err(), "reader must wait behind writer admission");
3779 drop(writer);
3780 let _reader = acquire_autocommit_permit(
3781 &gate,
3782 AdmissionMode::Reader,
3783 Duration::from_secs(1),
3784 &metrics,
3785 )
3786 .await
3787 .expect("reader must proceed after writer releases");
3788 }
3789
3790 #[tokio::test]
3791 async fn queued_writer_is_not_starved_by_later_readers() {
3792 let gate = new_tx_gate();
3793 let metrics = Arc::new(Metrics::new());
3794 let first_reader = acquire_autocommit_permit(
3795 &gate,
3796 AdmissionMode::Reader,
3797 Duration::from_secs(1),
3798 &metrics,
3799 )
3800 .await
3801 .expect("first reader admission");
3802
3803 let writer_gate = gate.clone();
3804 let writer_metrics = metrics.clone();
3805 let mut writer = tokio::spawn(async move {
3806 acquire_begin_permit(&writer_gate, Duration::from_secs(1), &writer_metrics).await
3807 });
3808 tokio::time::sleep(Duration::from_millis(10)).await;
3809
3810 let late_gate = gate.clone();
3811 let late_metrics = metrics.clone();
3812 let mut late_reader = tokio::spawn(async move {
3813 acquire_autocommit_permit(
3814 &late_gate,
3815 AdmissionMode::Reader,
3816 Duration::from_secs(1),
3817 &late_metrics,
3818 )
3819 .await
3820 });
3821 assert!(
3822 tokio::time::timeout(Duration::from_millis(25), &mut late_reader)
3823 .await
3824 .is_err(),
3825 "a later reader must queue behind the waiting writer"
3826 );
3827
3828 drop(first_reader);
3829 let writer_permit = tokio::time::timeout(Duration::from_secs(1), &mut writer)
3830 .await
3831 .expect("writer must acquire once prior readers drain")
3832 .expect("writer task")
3833 .expect("writer admission");
3834 assert!(
3835 tokio::time::timeout(Duration::from_millis(25), &mut late_reader)
3836 .await
3837 .is_err(),
3838 "writer must hold exclusive admission before the late reader"
3839 );
3840 drop(writer_permit);
3841 let _late_reader = tokio::time::timeout(Duration::from_secs(1), late_reader)
3842 .await
3843 .expect("late reader must eventually acquire")
3844 .expect("late reader task")
3845 .expect("late reader admission");
3846 }
3847
3848 fn dirty_view_engine() -> (tempfile::TempDir, Arc<RwLock<Engine>>) {
3849 let dir = tempfile::tempdir().unwrap();
3850 let mut engine = Engine::new(dir.path()).unwrap();
3851 engine
3852 .execute_powql("type Source { required id: int }")
3853 .unwrap();
3854 engine.execute_powql("insert Source { id := 1 }").unwrap();
3855 engine
3856 .execute_powql("materialize Snapshot as Source")
3857 .unwrap();
3858 engine.execute_powql("insert Source { id := 2 }").unwrap();
3859 (dir, Arc::new(RwLock::new(engine)))
3860 }
3861
3862 #[test]
3863 fn dirty_view_requests_explicit_escalation_on_every_frontend() {
3864 let (_dir, engine) = dirty_view_engine();
3865
3866 assert!(matches!(
3867 dispatch_query(&engine, "Snapshot", None, false).0,
3868 Err(QueryError::ReadonlyNeedsWrite)
3869 ));
3870 assert!(matches!(
3871 dispatch_sql_query(&engine, "SELECT * FROM Snapshot", None, false).0,
3872 Err(QueryError::ReadonlyNeedsWrite)
3873 ));
3874 assert!(matches!(
3875 dispatch_query_with_params(
3876 &engine,
3877 "Snapshot filter .id = $1",
3878 &[WireParam::Int(1)],
3879 None,
3880 false,
3881 )
3882 .0,
3883 Err(QueryError::ReadonlyNeedsWrite)
3884 ));
3885 }
3886
3887 #[tokio::test]
3888 async fn dirty_view_upgrade_waits_for_held_reader_then_records_once() {
3889 let (_dir, engine) = dirty_view_engine();
3890 let gate = new_tx_gate_with_permits(2);
3891 let metrics = Arc::new(Metrics::new());
3892 let held_reader = acquire_autocommit_permit(
3893 &gate,
3894 AdmissionMode::Reader,
3895 Duration::from_secs(1),
3896 &metrics,
3897 )
3898 .await
3899 .expect("held reader admission");
3900
3901 let (_client, server) = tokio::io::duplex(1024);
3904 let task_gate = gate.clone();
3905 let task_metrics = metrics.clone();
3906 let mut task = tokio::spawn(async move {
3907 let mut reader = BufReader::new(server);
3908 let mut wire_read_buffer = Vec::new();
3909 let mut pending_messages = InFlightReadAhead::default();
3910 let mut tx_permit = None;
3911 execute_wire_query(
3912 engine,
3913 task_gate,
3914 &mut tx_permit,
3915 "Snapshot".into(),
3916 WireResultMode::Native,
3917 None,
3918 Duration::from_secs(2),
3919 Duration::from_secs(1),
3920 &task_metrics,
3921 &mut reader,
3922 &mut wire_read_buffer,
3923 &mut pending_messages,
3924 )
3925 .await
3926 });
3927
3928 assert!(
3929 tokio::time::timeout(Duration::from_millis(100), &mut task)
3930 .await
3931 .is_err(),
3932 "dirty-view retry must wait for exclusive admission while another reader is held"
3933 );
3934 drop(held_reader);
3935
3936 let (message, ticket, termination) = tokio::time::timeout(Duration::from_secs(2), task)
3937 .await
3938 .expect("upgrade must finish after the held reader releases")
3939 .expect("query task");
3940 assert!(matches!(message, Message::ResultRowsNative { .. }));
3941 assert!(termination.is_none());
3942
3943 let (ticket, metric) = ticket.expect("view refresh must defer its WAL metric");
3944 drop(ticket);
3945 metrics.record_query(metric.start.elapsed(), metric.outcome);
3946
3947 let rendered = metrics.render();
3948 assert!(rendered.contains("powdb_queries_total{result=\"ok\"} 1"));
3949 assert!(rendered.contains("powdb_queries_total{result=\"error\"} 0"));
3950 }
3951
3952 #[tokio::test]
3953 async fn timed_out_readonly_escalation_is_not_retried_or_reported_as_generic_error() {
3954 let (_dir, engine) = dirty_view_engine();
3955 let metrics = Arc::new(Metrics::new());
3956 let (_client, server) = tokio::io::duplex(1024);
3959 let mut reader = BufReader::new(server);
3960 let mut wire_read_buffer = Vec::new();
3961 let mut pending_messages = InFlightReadAhead::default();
3962 let query_timeout = Duration::from_millis(20);
3963 let query_deadline = Instant::now() + query_timeout;
3964
3965 let (message, ticket, termination, retry) = run_blocking_query(
3966 engine,
3967 (),
3968 None,
3969 WireResultMode::Native,
3970 query_timeout,
3971 query_deadline,
3972 &metrics,
3973 &mut reader,
3974 &mut wire_read_buffer,
3975 &mut pending_messages,
3976 |_engine, (), _principal| {
3977 std::thread::sleep(Duration::from_millis(50));
3981 (Err(QueryError::ReadonlyNeedsWrite), None)
3982 },
3983 )
3984 .await;
3985
3986 assert!(ticket.is_none());
3987 assert!(termination.is_none());
3988 assert!(!retry, "a timed-out query must never be resurrected");
3989 match message {
3990 Message::Error { message } => assert!(
3991 message.contains("query timeout after 20ms"),
3992 "timeout must remain client-visible, got {message}"
3993 ),
3994 other => panic!("expected timeout error, got {other:?}"),
3995 }
3996 let rendered = metrics.render();
3997 assert!(rendered.contains("powdb_query_timeouts_total 1"));
3998 assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
3999 }
4000
4001 #[test]
4002 fn resource_limit_errors_surface_actionable_hints() {
4003 for msg in [
4008 "sort input exceeds row limit — add a LIMIT clause",
4009 "join result exceeds row limit",
4010 "query exceeded memory budget: requested 100 bytes, limit 50 bytes",
4011 "result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
4012 ] {
4013 assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
4014 }
4015 }
4016
4017 #[test]
4018 fn oversized_result_is_rejected_before_wire_encoding() {
4019 let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
4020 let result = QueryResult::Rows {
4021 columns: vec!["payload".into()],
4022 rows: vec![vec![Value::Str(long)]],
4023 };
4024 let err = query_result_to_message(result, WireResultMode::LegacyText).unwrap_err();
4025 assert!(
4026 err.to_string().starts_with("result too large"),
4027 "unexpected error: {err}"
4028 );
4029 }
4030
4031 fn parsed(q: &str) -> powdb_query::ast::Statement {
4034 parser::parse(q).unwrap()
4035 }
4036
4037 fn principal(role: &str) -> Option<Principal> {
4038 Some(Principal {
4039 name: "u".into(),
4040 role: role.into(),
4041 })
4042 }
4043
4044 #[test]
4045 fn readonly_can_read_but_not_write() {
4046 let p = principal("readonly");
4047 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4049 assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
4050 assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
4051 for q in [
4053 r#"insert User { name := "x" }"#,
4054 "User filter .id = 1 update { age := 2 }",
4055 "User filter .id = 1 delete",
4056 "drop User",
4057 "alter User add column c: str",
4058 "type T { required id: int }",
4059 "begin",
4060 "commit",
4061 "rollback",
4062 ] {
4063 let err = check_statement_permitted(p.as_ref(), &parsed(q))
4064 .expect_err(&format!("must deny: {q}"));
4065 assert!(
4066 err.to_string().contains("permission denied"),
4067 "unexpected error for {q}: {err}"
4068 );
4069 }
4070 }
4071
4072 #[test]
4073 fn readwrite_and_admin_have_full_query_access() {
4074 for role in ["readwrite", "admin"] {
4075 let p = principal(role);
4076 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4077 assert!(check_statement_permitted(
4078 p.as_ref(),
4079 &parsed(r#"insert User { name := "x" }"#)
4080 )
4081 .is_ok());
4082 assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
4083 }
4084 }
4085
4086 #[test]
4087 fn unknown_role_fails_closed_for_writes() {
4088 let p = principal("mystery");
4089 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4090 assert!(
4091 check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
4092 .is_err()
4093 );
4094 }
4095
4096 #[test]
4097 fn no_principal_means_full_access() {
4098 assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
4100 assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
4101 }
4102
4103 fn store_with_alice() -> UserStore {
4104 let mut s = UserStore::new();
4105 s.create_user("alice", "pw", "readwrite").unwrap();
4106 s
4107 }
4108
4109 #[test]
4112 fn empty_store_no_password_is_open() {
4113 let s = UserStore::new();
4114 assert_eq!(
4115 authenticate_connect(&s, None, None, None),
4116 AuthOutcome::Authenticated { principal: None }
4117 );
4118 assert_eq!(
4120 authenticate_connect(&s, None, Some("x"), Some("y")),
4121 AuthOutcome::Authenticated { principal: None }
4122 );
4123 }
4124
4125 #[test]
4126 fn empty_store_correct_shared_password_succeeds() {
4127 let s = UserStore::new();
4128 assert_eq!(
4129 authenticate_connect(&s, Some("pw"), None, Some("pw")),
4130 AuthOutcome::Authenticated { principal: None }
4131 );
4132 }
4133
4134 #[test]
4135 fn empty_store_wrong_shared_password_rejected() {
4136 let s = UserStore::new();
4137 assert_eq!(
4138 authenticate_connect(&s, Some("pw"), None, Some("bad")),
4139 AuthOutcome::Rejected
4140 );
4141 }
4142
4143 #[test]
4144 fn empty_store_missing_password_rejected_when_expected() {
4145 let s = UserStore::new();
4146 assert_eq!(
4147 authenticate_connect(&s, Some("pw"), None, None),
4148 AuthOutcome::Rejected
4149 );
4150 }
4151
4152 #[test]
4153 fn empty_store_ignores_username_for_shared_password() {
4154 let s = UserStore::new();
4157 assert_eq!(
4158 authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
4159 AuthOutcome::Authenticated { principal: None }
4160 );
4161 }
4162
4163 #[test]
4166 fn user_auth_success_binds_principal() {
4167 let s = store_with_alice();
4168 assert_eq!(
4169 authenticate_connect(&s, None, Some("alice"), Some("pw")),
4170 AuthOutcome::Authenticated {
4171 principal: Some(Principal {
4172 name: "alice".into(),
4173 role: "readwrite".into(),
4174 })
4175 }
4176 );
4177 }
4178
4179 #[test]
4180 fn user_auth_wrong_password_rejected() {
4181 let s = store_with_alice();
4182 assert_eq!(
4183 authenticate_connect(&s, None, Some("alice"), Some("bad")),
4184 AuthOutcome::Rejected
4185 );
4186 }
4187
4188 #[test]
4189 fn user_auth_unknown_user_rejected() {
4190 let s = store_with_alice();
4191 assert_eq!(
4192 authenticate_connect(&s, None, Some("mallory"), Some("pw")),
4193 AuthOutcome::Rejected
4194 );
4195 }
4196
4197 #[test]
4198 fn user_auth_missing_username_rejected() {
4199 let s = store_with_alice();
4200 assert_eq!(
4201 authenticate_connect(&s, None, None, Some("pw")),
4202 AuthOutcome::Rejected
4203 );
4204 }
4205
4206 #[test]
4207 fn user_auth_missing_password_rejected() {
4208 let s = store_with_alice();
4209 assert_eq!(
4210 authenticate_connect(&s, Some("pw"), Some("alice"), None),
4211 AuthOutcome::Rejected
4212 );
4213 }
4214
4215 #[test]
4216 fn user_auth_ignores_shared_password_when_users_present() {
4217 let s = store_with_alice();
4220 assert_eq!(
4221 authenticate_connect(&s, Some("shared"), None, Some("shared")),
4222 AuthOutcome::Rejected
4223 );
4224 }
4225
4226 #[test]
4227 fn replica_fingerprint_is_stable_and_redacted() {
4228 let replica_id = "customer-prod-replica-a";
4229 let fingerprint = replica_fingerprint(replica_id);
4230 assert_eq!(fingerprint, replica_fingerprint(replica_id));
4231 assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
4232 assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
4233 assert_eq!(fingerprint.len(), 16);
4234 assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
4235 assert!(!fingerprint.contains("customer"));
4236 assert!(!fingerprint.contains("replica"));
4237 assert!(!fingerprint.contains(replica_id));
4238 }
4239
4240 #[test]
4241 fn invalid_replica_ids_use_fixed_log_fingerprint() {
4242 assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
4243 assert_eq!(
4244 log_replica_fingerprint("customer/prod/replica"),
4245 INVALID_REPLICA_FINGERPRINT
4246 );
4247 assert_eq!(
4248 log_replica_fingerprint(&"a".repeat(4096)),
4249 INVALID_REPLICA_FINGERPRINT
4250 );
4251 }
4252
4253 #[test]
4254 fn sync_error_classes_use_bounded_labels() {
4255 assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
4256 assert_eq!(
4257 SyncErrorClass::PermissionDenied.as_label(),
4258 "permission_denied"
4259 );
4260 assert_eq!(
4261 SyncErrorClass::IdentityOrFormatMismatch.as_label(),
4262 "identity_or_format_mismatch"
4263 );
4264 assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
4265 assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
4266 }
4267
4268 fn sync_identity() -> DatabaseIdentity {
4269 DatabaseIdentity {
4270 database_id: *b"server-sync-test",
4271 primary_generation: 1,
4272 }
4273 }
4274
4275 fn retained_unit(lsn: u64) -> RetainedUnit {
4276 RetainedUnit {
4277 tx_id: 1,
4278 record_type: 4,
4279 lsn,
4280 data: lsn.to_le_bytes().to_vec(),
4281 }
4282 }
4283
4284 fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
4285 RetainedUnit {
4286 tx_id,
4287 record_type: record_type as u8,
4288 lsn,
4289 data: lsn.to_le_bytes().to_vec(),
4290 }
4291 }
4292
4293 fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
4294 let identity = sync_identity();
4295 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4296 let units = (1..=through_lsn).map(retained_unit).collect();
4297 let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
4298 write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
4299 }
4300
4301 fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
4302 let identity = sync_identity();
4303 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4304 let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
4305 write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
4306 }
4307
4308 fn write_sync_identity_only(data_dir: &std::path::Path) {
4309 let identity = sync_identity();
4310 write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4311 }
4312
4313 fn admin_principal() -> Principal {
4314 Principal {
4315 name: "admin".into(),
4316 role: "admin".into(),
4317 }
4318 }
4319
4320 #[test]
4321 fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
4322 let dir = tempfile::tempdir().unwrap();
4323 let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
4324
4325 match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
4326 Message::Error { message } => {
4327 assert!(message.contains("requires authentication"));
4328 }
4329 other => panic!("expected auth error, got {other:?}"),
4330 }
4331
4332 let readonly = Principal {
4333 name: "reader".into(),
4334 role: "readonly".into(),
4335 };
4336 match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
4337 Message::Error { message } => {
4338 assert!(message.contains("permission denied"));
4339 }
4340 other => panic!("expected permission error, got {other:?}"),
4341 }
4342 }
4343
4344 #[test]
4345 fn sync_status_pull_and_ack_use_server_remote_lsn() {
4346 let dir = tempfile::tempdir().unwrap();
4347 let mut engine = Engine::new(dir.path()).unwrap();
4348 engine
4349 .execute_powql("type SyncT { required id: int, v: str }")
4350 .unwrap();
4351 engine
4352 .execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
4353 .unwrap();
4354 let remote_lsn = engine.catalog().max_lsn();
4355 assert!(remote_lsn > 0);
4356 write_sync_identity_and_tail(dir.path(), remote_lsn);
4357 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4358 .unwrap();
4359
4360 let engine = Arc::new(RwLock::new(engine));
4361 let principal = admin_principal();
4362 let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
4363 {
4364 Message::SyncStatusResult { status } => status,
4365 other => panic!("expected sync status, got {other:?}"),
4366 };
4367 assert_eq!(status.remote_lsn, remote_lsn);
4368 assert_eq!(status.servable_lsn, Some(remote_lsn));
4369 assert_eq!(status.unarchived_lsn, Some(0));
4370 assert_eq!(status.last_applied_lsn, Some(0));
4371 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4372 assert!(status.stale);
4373
4374 let identity = sync_identity().segment_identity();
4375 let pull = SyncPullRequest {
4376 replica_id: "replica-a".into(),
4377 since_lsn: 0,
4378 max_units: MAX_SYNC_PULL_UNITS,
4379 max_bytes: MAX_SYNC_PULL_BYTES,
4380 database_id: identity.database_id,
4381 primary_generation: identity.primary_generation,
4382 wal_format_version: identity.wal_format_version,
4383 catalog_version: identity.catalog_version,
4384 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4385 };
4386 let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4387 Message::SyncPullResult {
4388 status,
4389 units,
4390 has_more,
4391 } => {
4392 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4393 assert!(!has_more);
4394 units
4395 }
4396 other => panic!("expected sync pull result, got {other:?}"),
4397 };
4398 assert_eq!(units.len() as u64, remote_lsn);
4399 assert_eq!(units.last().unwrap().lsn, remote_lsn);
4400
4401 let ack = match dispatch_sync_ack(
4402 &engine,
4403 "replica-a".into(),
4404 remote_lsn,
4405 remote_lsn,
4406 true,
4407 Some(&principal),
4408 ) {
4409 Message::SyncAckResult {
4410 previous_applied_lsn,
4411 applied_lsn,
4412 remote_lsn: ack_remote_lsn,
4413 advanced,
4414 status,
4415 } => {
4416 assert_eq!(previous_applied_lsn, 0);
4417 assert_eq!(applied_lsn, remote_lsn);
4418 assert_eq!(ack_remote_lsn, remote_lsn);
4419 assert!(advanced);
4420 status
4421 }
4422 other => panic!("expected sync ack result, got {other:?}"),
4423 };
4424 assert_eq!(ack.repair_action, WireSyncRepairAction::None);
4425 assert!(!ack.stale);
4426 assert_eq!(ack.lag_lsn, Some(0));
4427 }
4428
4429 fn seed_pullable_replica(engine: &mut Engine) -> u64 {
4430 let data_dir = engine.catalog().data_dir().to_path_buf();
4431 let remote_lsn = engine.catalog().max_lsn();
4432 assert!(remote_lsn > 0);
4433 write_sync_identity_and_tail(&data_dir, remote_lsn);
4434 powdb_sync::upsert_replica_cursor(&data_dir, ReplicaCursor::active("replica-a", 0))
4435 .unwrap();
4436 remote_lsn
4437 }
4438
4439 fn pull_request_with_catalog_version(catalog_version: u16) -> SyncPullRequest {
4440 let identity = sync_identity().segment_identity();
4441 SyncPullRequest {
4442 replica_id: "replica-a".into(),
4443 since_lsn: 0,
4444 max_units: MAX_SYNC_PULL_UNITS,
4445 max_bytes: MAX_SYNC_PULL_BYTES,
4446 database_id: identity.database_id,
4447 primary_generation: identity.primary_generation,
4448 wal_format_version: identity.wal_format_version,
4449 catalog_version,
4450 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4451 }
4452 }
4453
4454 #[test]
4455 fn fresh_database_expects_legacy_catalog_version_and_accepts_v5_replica() {
4456 use powdb_storage::catalog::{CATALOG_VERSION, LEGACY_CATALOG_VERSION};
4457
4458 let dir = tempfile::tempdir().unwrap();
4459 let mut engine = Engine::new(dir.path()).unwrap();
4460 engine
4461 .execute_powql("type Doc { required id: int, data: json }")
4462 .unwrap();
4463 engine
4464 .execute_powql(r#"insert Doc { id := 1, data := "{\"score\":20}" }"#)
4465 .unwrap();
4466 assert_eq!(
4469 engine.catalog().active_catalog_version(),
4470 LEGACY_CATALOG_VERSION
4471 );
4472 let remote_lsn = seed_pullable_replica(&mut engine);
4473
4474 let engine = Arc::new(RwLock::new(engine));
4475 let principal = admin_principal();
4476
4477 let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION);
4480 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4481 Message::SyncPullResult { units, .. } => {
4482 assert_eq!(units.len() as u64, remote_lsn);
4483 }
4484 other => panic!("expected sync pull result, got {other:?}"),
4485 }
4486
4487 let pull = pull_request_with_catalog_version(CATALOG_VERSION);
4489 assert!(matches!(
4490 dispatch_sync_pull(&engine, pull, true, Some(&principal)),
4491 Message::SyncPullResult { .. }
4492 ));
4493
4494 let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION - 1);
4497 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4498 Message::Error { message } => {
4499 assert!(message.contains("v4"), "message: {message}");
4500 assert!(message.contains("v5"), "message: {message}");
4501 assert!(
4502 message.contains("rebootstrap with an upgraded replica required"),
4503 "message: {message}"
4504 );
4505 }
4506 other => panic!("expected identity mismatch error, got {other:?}"),
4507 }
4508 }
4509
4510 #[test]
4511 fn activated_database_expects_v6_and_rejects_v5_replica() {
4512 use powdb_storage::catalog::{CATALOG_VERSION, LEGACY_CATALOG_VERSION};
4513
4514 let dir = tempfile::tempdir().unwrap();
4515 let mut engine = Engine::new(dir.path()).unwrap();
4516 engine
4517 .execute_powql("type Doc { required id: int, data: json }")
4518 .unwrap();
4519 engine
4520 .execute_powql(r#"insert Doc { id := 1, data := "{\"score\":20}" }"#)
4521 .unwrap();
4522 engine
4524 .execute_powql("alter Doc add index (.data->score)")
4525 .unwrap();
4526 assert_eq!(engine.catalog().active_catalog_version(), CATALOG_VERSION);
4527 let remote_lsn = seed_pullable_replica(&mut engine);
4528
4529 let engine = Arc::new(RwLock::new(engine));
4530 let principal = admin_principal();
4531
4532 let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION);
4535 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4536 Message::Error { message } => {
4537 assert!(message.contains("v5"), "message: {message}");
4538 assert!(message.contains("v6"), "message: {message}");
4539 assert!(
4540 message.contains("rebootstrap with an upgraded replica required"),
4541 "message: {message}"
4542 );
4543 }
4544 other => panic!("expected identity mismatch error, got {other:?}"),
4545 }
4546
4547 let pull = pull_request_with_catalog_version(CATALOG_VERSION);
4549 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4550 Message::SyncPullResult { units, .. } => {
4551 assert_eq!(units.len() as u64, remote_lsn);
4552 }
4553 other => panic!("expected sync pull result, got {other:?}"),
4554 }
4555 }
4556
4557 #[test]
4558 fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
4559 let dir = tempfile::tempdir().unwrap();
4560 let mut engine = Engine::new(dir.path()).unwrap();
4561 engine
4562 .execute_powql("type SyncT { required id: int }")
4563 .unwrap();
4564 for id in 1..=3 {
4565 engine
4566 .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
4567 .unwrap();
4568 }
4569 let remote_lsn = engine.catalog().max_lsn();
4570 assert!(remote_lsn >= 3);
4571 write_sync_identity_and_units(
4572 dir.path(),
4573 vec![
4574 retained_unit_with(77, WalRecordType::Begin, 1),
4575 retained_unit_with(77, WalRecordType::Insert, 2),
4576 retained_unit_with(77, WalRecordType::Commit, 3),
4577 ],
4578 );
4579 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4580 .unwrap();
4581
4582 let engine = Arc::new(RwLock::new(engine));
4583 let principal = admin_principal();
4584 let identity = sync_identity().segment_identity();
4585 let cut_pull = SyncPullRequest {
4586 replica_id: "replica-a".into(),
4587 since_lsn: 0,
4588 max_units: 2,
4589 max_bytes: MAX_SYNC_PULL_BYTES,
4590 database_id: identity.database_id,
4591 primary_generation: identity.primary_generation,
4592 wal_format_version: identity.wal_format_version,
4593 catalog_version: identity.catalog_version,
4594 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4595 };
4596 match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
4597 Message::Error { message } => assert!(message.contains("cuts through transaction")),
4598 other => panic!("expected transaction-cut pull error, got {other:?}"),
4599 }
4600
4601 let cut_bytes_pull = SyncPullRequest {
4602 replica_id: "replica-a".into(),
4603 since_lsn: 0,
4604 max_units: 3,
4605 max_bytes: 58,
4606 database_id: identity.database_id,
4607 primary_generation: identity.primary_generation,
4608 wal_format_version: identity.wal_format_version,
4609 catalog_version: identity.catalog_version,
4610 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4611 };
4612 match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
4613 Message::Error { message } => assert!(message.contains("cuts through transaction")),
4614 other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
4615 }
4616
4617 let full_pull = SyncPullRequest {
4618 replica_id: "replica-a".into(),
4619 since_lsn: 0,
4620 max_units: 3,
4621 max_bytes: MAX_SYNC_PULL_BYTES,
4622 database_id: identity.database_id,
4623 primary_generation: identity.primary_generation,
4624 wal_format_version: identity.wal_format_version,
4625 catalog_version: identity.catalog_version,
4626 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4627 };
4628 match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
4629 Message::SyncPullResult { units, .. } => {
4630 assert_eq!(units.len(), 3);
4631 assert_eq!(units.last().unwrap().lsn, 3);
4632 }
4633 other => panic!("expected complete transaction pull, got {other:?}"),
4634 }
4635
4636 match dispatch_sync_ack(
4637 &engine,
4638 "replica-a".into(),
4639 2,
4640 remote_lsn,
4641 true,
4642 Some(&principal),
4643 ) {
4644 Message::Error { message } => assert!(message.contains("cuts through transaction")),
4645 other => panic!("expected transaction-cut ack error, got {other:?}"),
4646 }
4647 let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
4648 assert_eq!(cursor[0].applied_lsn, 0);
4649
4650 match dispatch_sync_ack(
4651 &engine,
4652 "replica-a".into(),
4653 3,
4654 remote_lsn,
4655 true,
4656 Some(&principal),
4657 ) {
4658 Message::SyncAckResult {
4659 previous_applied_lsn,
4660 applied_lsn,
4661 advanced,
4662 ..
4663 } => {
4664 assert_eq!(previous_applied_lsn, 0);
4665 assert_eq!(applied_lsn, 3);
4666 assert!(advanced);
4667 }
4668 other => panic!("expected complete transaction ack, got {other:?}"),
4669 }
4670 }
4671
4672 #[test]
4673 fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
4674 let dir = tempfile::tempdir().unwrap();
4675 let mut engine = Engine::new(dir.path()).unwrap();
4676 engine
4677 .execute_powql("type SyncT { required id: int }")
4678 .unwrap();
4679 for id in 1..=6 {
4680 engine
4681 .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
4682 .unwrap();
4683 }
4684 let remote_lsn = engine.catalog().max_lsn();
4685 assert!(remote_lsn >= 6);
4686 write_sync_identity_and_units(
4687 dir.path(),
4688 vec![
4689 retained_unit_with(1, WalRecordType::Begin, 1),
4690 retained_unit_with(1, WalRecordType::Insert, 2),
4691 retained_unit_with(1, WalRecordType::Commit, 3),
4692 retained_unit_with(1, WalRecordType::Begin, 4),
4693 retained_unit_with(1, WalRecordType::Insert, 5),
4694 retained_unit_with(1, WalRecordType::Commit, 6),
4695 ],
4696 );
4697 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4698 .unwrap();
4699
4700 let engine = Arc::new(RwLock::new(engine));
4701 let principal = admin_principal();
4702 let identity = sync_identity().segment_identity();
4703 let pull = SyncPullRequest {
4704 replica_id: "replica-a".into(),
4705 since_lsn: 0,
4706 max_units: 6,
4707 max_bytes: 100,
4708 database_id: identity.database_id,
4709 primary_generation: identity.primary_generation,
4710 wal_format_version: identity.wal_format_version,
4711 catalog_version: identity.catalog_version,
4712 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4713 };
4714
4715 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4716 Message::SyncPullResult {
4717 status,
4718 units,
4719 has_more,
4720 } => {
4721 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4722 assert_eq!(units.len(), 3);
4723 assert_eq!(units.last().unwrap().lsn, 3);
4724 assert!(has_more);
4725 }
4726 other => panic!("expected byte-capped applyable prefix, got {other:?}"),
4727 }
4728 }
4729
4730 #[test]
4731 fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
4732 let dir = tempfile::tempdir().unwrap();
4733 let mut engine = Engine::new(dir.path()).unwrap();
4734 engine
4735 .execute_powql("type SyncT { required id: int }")
4736 .unwrap();
4737 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4738 let remote_lsn = engine.catalog().max_lsn();
4739 assert!(remote_lsn > 0);
4740 write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
4741 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4742 .unwrap();
4743
4744 let engine = Arc::new(RwLock::new(engine));
4745 let principal = admin_principal();
4746 let identity = sync_identity().segment_identity();
4747 let pull = SyncPullRequest {
4748 replica_id: "replica-a".into(),
4749 since_lsn: 0,
4750 max_units: MAX_SYNC_PULL_UNITS,
4751 max_bytes: MAX_SYNC_PULL_BYTES,
4752 database_id: identity.database_id,
4753 primary_generation: identity.primary_generation,
4754 wal_format_version: identity.wal_format_version,
4755 catalog_version: identity.catalog_version,
4756 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4757 };
4758
4759 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4760 Message::SyncPullResult {
4761 status,
4762 units,
4763 has_more,
4764 } => {
4765 assert_eq!(status.remote_lsn, remote_lsn);
4766 assert_eq!(status.servable_lsn, Some(remote_lsn));
4767 assert_eq!(status.unarchived_lsn, Some(0));
4768 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4769 assert!(!has_more);
4770 assert_eq!(units.len() as u64, remote_lsn);
4771 assert_eq!(units.last().unwrap().lsn, remote_lsn);
4772 assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
4773 }
4774 other => panic!("expected capped sync pull result, got {other:?}"),
4775 }
4776 }
4777
4778 #[test]
4779 fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
4780 let dir = tempfile::tempdir().unwrap();
4781 let mut engine = Engine::new(dir.path()).unwrap();
4782 engine
4783 .execute_powql("type SyncT { required id: int }")
4784 .unwrap();
4785 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4786 let remote_lsn = engine.catalog().max_lsn();
4787 assert!(remote_lsn > 0);
4788 write_sync_identity_only(dir.path());
4789 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4790 .unwrap();
4791
4792 let engine = Arc::new(RwLock::new(engine));
4793 let principal = admin_principal();
4794 let identity = sync_identity().segment_identity();
4795 let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
4796 {
4797 Message::SyncStatusResult { status } => status,
4798 other => panic!("expected sync status, got {other:?}"),
4799 };
4800 assert_eq!(status.remote_lsn, remote_lsn);
4801 assert_eq!(status.servable_lsn, Some(0));
4802 assert_eq!(status.unarchived_lsn, Some(remote_lsn));
4803 assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
4804 assert!(status
4805 .last_sync_error
4806 .as_deref()
4807 .unwrap()
4808 .contains("not yet archived"));
4809
4810 let pull = SyncPullRequest {
4811 replica_id: "replica-a".into(),
4812 since_lsn: 0,
4813 max_units: MAX_SYNC_PULL_UNITS,
4814 max_bytes: MAX_SYNC_PULL_BYTES,
4815 database_id: identity.database_id,
4816 primary_generation: identity.primary_generation,
4817 wal_format_version: identity.wal_format_version,
4818 catalog_version: identity.catalog_version,
4819 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4820 };
4821 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4822 Message::SyncPullResult {
4823 status,
4824 units,
4825 has_more,
4826 } => {
4827 assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
4828 assert!(units.is_empty());
4829 assert!(!has_more);
4830 }
4831 other => panic!("expected await-archive sync pull result, got {other:?}"),
4832 }
4833 }
4834
4835 #[test]
4836 fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
4837 let dir = tempfile::tempdir().unwrap();
4838 let mut engine = Engine::new(dir.path()).unwrap();
4839 engine
4840 .execute_powql("type SyncT { required id: int }")
4841 .unwrap();
4842 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4843 engine.execute_powql("insert SyncT { id := 2 }").unwrap();
4844 let remote_lsn = engine.catalog().max_lsn();
4845 assert!(remote_lsn > 1);
4846 let servable_lsn = remote_lsn - 1;
4847 write_sync_identity_and_tail(dir.path(), servable_lsn);
4848 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4849 .unwrap();
4850
4851 let engine = Arc::new(RwLock::new(engine));
4852 let principal = admin_principal();
4853 let identity = sync_identity().segment_identity();
4854 let pull = SyncPullRequest {
4855 replica_id: "replica-a".into(),
4856 since_lsn: 0,
4857 max_units: MAX_SYNC_PULL_UNITS,
4858 max_bytes: MAX_SYNC_PULL_BYTES,
4859 database_id: identity.database_id,
4860 primary_generation: identity.primary_generation,
4861 wal_format_version: identity.wal_format_version,
4862 catalog_version: identity.catalog_version,
4863 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4864 };
4865
4866 match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4867 Message::SyncPullResult {
4868 status,
4869 units,
4870 has_more,
4871 } => {
4872 assert_eq!(status.remote_lsn, remote_lsn);
4873 assert_eq!(status.servable_lsn, Some(servable_lsn));
4874 assert_eq!(status.unarchived_lsn, Some(1));
4875 assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4876 assert!(!has_more);
4877 assert_eq!(units.len() as u64, servable_lsn);
4878 assert_eq!(units.last().unwrap().lsn, servable_lsn);
4879 assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
4880 }
4881 other => panic!("expected partial sync pull result, got {other:?}"),
4882 }
4883 }
4884
4885 #[test]
4886 fn sync_pull_rejects_cursor_or_format_mismatch() {
4887 let dir = tempfile::tempdir().unwrap();
4888 let mut engine = Engine::new(dir.path()).unwrap();
4889 engine
4890 .execute_powql("type SyncT { required id: int }")
4891 .unwrap();
4892 engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4893 let remote_lsn = engine.catalog().max_lsn();
4894 write_sync_identity_and_tail(dir.path(), remote_lsn);
4895 powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4896 .unwrap();
4897 let engine = Arc::new(RwLock::new(engine));
4898 let principal = admin_principal();
4899 let identity = sync_identity().segment_identity();
4900
4901 let wrong_cursor = SyncPullRequest {
4902 replica_id: "replica-a".into(),
4903 since_lsn: 1,
4904 max_units: 10,
4905 max_bytes: MAX_SYNC_PULL_BYTES,
4906 database_id: identity.database_id,
4907 primary_generation: identity.primary_generation,
4908 wal_format_version: identity.wal_format_version,
4909 catalog_version: identity.catalog_version,
4910 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4911 };
4912 match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
4913 Message::Error { message } => assert!(message.contains("does not match")),
4914 other => panic!("expected cursor mismatch error, got {other:?}"),
4915 }
4916
4917 let wrong_format = SyncPullRequest {
4918 replica_id: "replica-a".into(),
4919 since_lsn: 0,
4920 max_units: 10,
4921 max_bytes: MAX_SYNC_PULL_BYTES,
4922 database_id: identity.database_id,
4923 primary_generation: identity.primary_generation,
4924 wal_format_version: identity.wal_format_version,
4925 catalog_version: identity.catalog_version,
4926 segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
4927 };
4928 match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
4929 Message::Error { message } => assert!(message.contains("rebootstrap required")),
4930 other => panic!("expected format mismatch error, got {other:?}"),
4931 }
4932 }
4933}