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