Skip to main content

powdb_server/
handler.rs

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