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