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