Skip to main content

questdb/egress/
reader.rs

1/*******************************************************************************
2 *     ___                  _   ____  ____
3 *    / _ \ _   _  ___  ___| |_|  _ \| __ )
4 *   | | | | | | |/ _ \/ __| __| | | |  _ \
5 *   | |_| | |_| |  __/\__ \ |_| |_| | |_) |
6 *    \__\_\\__,_|\___||___/\__|____/|____/
7 *
8 *  Copyright (c) 2014-2019 Appsicle
9 *  Copyright (c) 2019-2025 QuestDB
10 *
11 *  Licensed under the Apache License, Version 2.0 (the "License");
12 *  you may not use this file except in compliance with the License.
13 *  You may obtain a copy of the License at
14 *
15 *  http://www.apache.org/licenses/LICENSE-2.0
16 *
17 *  Unless required by applicable law or agreed to in writing, software
18 *  distributed under the License is distributed on an "AS IS" BASIS,
19 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 *  See the License for the specific language governing permissions and
21 *  limitations under the License.
22 *
23 ******************************************************************************/
24
25//! `Reader` (per-connection) + `Cursor` (per-query) public API.
26//!
27//! Each `Reader` allows at most one in-flight cursor at a time
28//! (runtime-checked, not type-encoded). `Cursor::cancel()` issues a
29//! CANCEL frame and drains until the terminal frame, leaving the
30//! Reader reusable. Dropping a cursor before it has reached a
31//! terminal closes the underlying WebSocket: subsequent operations
32//! on the Reader fail at the transport layer (open a fresh Reader to
33//! recover). Call `Cursor::cancel()` (or read until `next_batch()`
34//! returns `None`) before drop if you want to keep the existing
35//! connection alive.
36//!
37//! The `sync-reader-qwp-ws` feature gate is applied at the module
38//! declaration in `egress/mod.rs`; an inner `#![cfg(...)]` here would
39//! duplicate that gate (clippy::duplicated_attributes) without
40//! changing what's compiled.
41
42use std::net::Ipv4Addr;
43use std::sync::Arc;
44use std::sync::atomic::{AtomicU64, Ordering};
45use std::time::Duration;
46
47use bytes::{Bytes, BytesMut};
48
49use crate::egress::binds::{Bind, SimpleNullKind};
50use crate::egress::column::ColumnView;
51use crate::egress::config::{Endpoint, ReaderConfig, Target};
52use crate::egress::decoder::DecodedBatch;
53use crate::egress::decoder::ZstdScratch;
54use crate::egress::query_request::{
55    QUERY_FLAG_RESET_DICT, QueryRequest, QueryRequestBuilder, REQUEST_ID_OFFSET,
56};
57use crate::egress::schema::Schema;
58use crate::egress::server_event::UpgradeReject;
59use crate::egress::server_event::{ServerEvent, ServerInfo, ServerRole, decode_frame};
60use crate::egress::symbol_dict::SymbolDict;
61use crate::egress::tracker::HostHealthTracker;
62use crate::egress::transport::{CLOSE_TIMEOUT, WRITE_TIMEOUT, WsTransport};
63use crate::egress::wire::capabilities::has_query_flags;
64use crate::egress::wire::header::HEADER_LEN;
65use crate::egress::wire::msg_kind::MsgKind;
66use crate::egress::wire::varint;
67use crate::error::{Error, ErrorCode, Result, fmt};
68
69// ---------------------------------------------------------------------------
70// Reader
71// ---------------------------------------------------------------------------
72
73/// Diagnostic counters shared between a [`Reader`] and its FFI handle.
74///
75/// Held by the Reader via [`Arc`] so the FFI surface can clone it once
76/// at handle-construction time and serve stat reads thereafter without
77/// touching the `UnsafeCell<Reader>` that holds the Reader. That
78/// decouples counter reads from the Reader's borrow stack: a stat
79/// getter no longer synthesises a `&Reader` while a laundered
80/// `&mut Reader` (held by an in-flight `ReaderQuery` / `Cursor`) is
81/// still on the stack — eliminating the aliasing question entirely.
82///
83/// All four counters are `Relaxed` — pure counters with no associated
84/// happens-before requirement.
85#[derive(Debug, Default)]
86pub struct ReaderStats {
87    /// Total wire bytes (frame header + payload) read off the
88    /// transport since this connection was opened.
89    pub bytes_received: AtomicU64,
90    /// Total bytes granted to the server via CREDIT (`0x15`) frames
91    /// since this connection was opened.
92    pub credit_granted_total: AtomicU64,
93    /// Nanoseconds spent in `transport.read_frame()` since this
94    /// connection was opened. Saturates at `u64::MAX`.
95    pub read_ns: AtomicU64,
96    /// Nanoseconds spent in `decode_frame()` since this connection
97    /// was opened. Saturates at `u64::MAX`.
98    pub decode_ns: AtomicU64,
99}
100
101/// Per-connection reader. Owns the WebSocket transport and the
102/// connection-scoped symbol dictionary.
103pub struct Reader {
104    /// Snapshot of the config used to open this connection. Owned (not
105    /// borrowed) because the cursor's failover machinery needs to outlive
106    /// the original `from_config` call and reach back into the address
107    /// list / failover knobs after the user has dropped their builder.
108    ///
109    /// Wrapped in [`Arc`] so reconnect attempts share a single
110    /// allocation: each attempt would otherwise deep-clone the addr
111    /// vec, the path string, and the boxed auth payload — with
112    /// `failover_max_attempts` up to `1024`, that's hundreds of
113    /// allocations per failure event. Reference-count bumps are free
114    /// in comparison.
115    cfg: Arc<ReaderConfig>,
116    /// Index into [`ReaderConfig::addrs`] this connection is bound to.
117    /// Updated on mid-query failover so the cursor walks the list in the
118    /// right order ("skip the failed one first") on the next failure.
119    addr_idx: usize,
120    /// Live WS transport. `Option` only so that mid-query failover
121    /// can take the dead transport out via [`Option::take`] (releasing
122    /// its TCP FD) **before** sleeping on the backoff. Outside of the
123    /// brief reconnect window inside [`Reader::reconnect_with_failover`],
124    /// this is always `Some`. Use [`Reader::transport`] /
125    /// [`Reader::transport_mut`] to access — they assert this invariant.
126    transport: Option<WsTransport>,
127    dict: SymbolDict,
128    /// Schema for the in-flight query. Populated from the first
129    /// `RESULT_BATCH` (`batch_seq == 0`) and reused by continuation
130    /// batches; `ReaderQuery::execute` clears it at query start and the
131    /// reconnect path clears it on failover so a replayed query re-reads
132    /// it from the new node's batch 0. A single slot suffices because a
133    /// `Reader` runs one cursor at a time; pipelined `request_id`s would
134    /// need a map keyed by request id.
135    query_schema: Option<Schema>,
136    next_request_id: i64,
137    cursor_active: bool,
138    /// Server's `SERVER_INFO` (`0x18`), captured eagerly during connect.
139    /// The single QWP version always sends it as the first frame, so this
140    /// is `Some` outside the brief reconnect window; multi-addr role
141    /// filtering uses it to dismiss endpoints whose role doesn't match
142    /// `target`.
143    server_info: Option<ServerInfo>,
144    /// Diagnostic counters (`bytes_received`, `credit_granted_total`,
145    /// `read_ns`, `decode_ns`) shared with the FFI handle via `Arc` so
146    /// that monitoring-thread stat reads can be served without ever
147    /// touching the `UnsafeCell<Reader>` that the FFI uses to hold this
148    /// `Reader`. Decoupling the counters from the Reader's borrow stack
149    /// removes the aliasing question of "what happens when a stat
150    /// getter synthesises a `&Reader` while a laundered `&mut Reader`
151    /// is in flight": the stat getter doesn't touch the Reader at all.
152    ///
153    /// The one-thread-at-a-time rule that governs the rest of the
154    /// Reader API is intentionally relaxed for these counters and
155    /// `reset_timing`: their getters take `&self`, touch only atomics,
156    /// and may be invoked concurrently from a monitoring thread while
157    /// another thread is driving a cursor. Every other accessor
158    /// (`current_addr`, `server_info`, `server_version`) reads
159    /// non-atomic state and remains bound by the one-thread-at-a-time
160    /// contract — racing them with an in-flight cursor is undefined
161    /// behaviour. `Relaxed` is sufficient: these are pure counters with
162    /// no associated happens-before requirement on other state.
163    stats: Arc<ReaderStats>,
164    /// Reusable zstd decompressor + output buffer. Keeps a persistent
165    /// `ZSTD_DCtx` across batches (so we don't pay context init per
166    /// `RESULT_BATCH`) and a `Vec<u8>` whose allocation is reused as
167    /// successive frames decompress through it.
168    zstd_scratch: ZstdScratch,
169    /// Per-client host-health tracker shared across the initial connect
170    /// and every mid-query reconnect. Implements the failover.md §2
171    /// priority lattice — endpoints are picked by (state tier × zone
172    /// tier × index), not by round-robin rotation; see
173    /// [`HostHealthTracker`]. Classifications accumulate across
174    /// Executes; only the round-attempted bits reset between walks.
175    /// Lives on the Reader so long-lived clients converge on the
176    /// healthiest endpoint over time.
177    tracker: HostHealthTracker,
178    /// Per-Reader PRNG for failover backoff jitter. Egress backoff
179    /// uses **full-jitter** `[0, base)` per failover.md §3.1 — a
180    /// query client is single-user and benefits from the lowest
181    /// expected recovery time. Lives on the Reader so the state
182    /// persists across reconnect cycles within a single Reader's
183    /// lifetime.
184    failover_rng: FailoverRng,
185}
186
187// Compile-time pin for the cross-thread contract the FFI and the public
188// Rust API both depend on: `Reader` may be migrated to a worker thread
189// while a monitoring thread reads `bytes_received` / `read_ns` /
190// `decode_ns` / `credit_granted_total` via the `Arc<ReaderStats>`.
191//
192// Without this assertion, a future field addition (`Rc<…>`, `RefCell<…>`,
193// `MutexGuard<'static, …>`, a custom `!Send`/`!Sync` type) would silently
194// flip Reader off `Send`/`Sync` and the PR description's claim that
195// "the reader handle may be migrated between threads" would turn false
196// without any signal — runtime tests would keep passing because nothing
197// actually exercises the migration. Pinning it here makes the bound
198// load-bearing: a regression breaks compilation.
199const _: fn() = || {
200    fn assert_send_sync<T: Send + Sync>() {}
201    assert_send_sync::<Reader>();
202    assert_send_sync::<ReaderStats>();
203    assert_send_sync::<HostHealthTracker>();
204};
205
206// Query and cursor handles may be migrated between threads, provided the
207// caller establishes a happens-before edge and never accesses a handle
208// concurrently. Keep this assertion next to Reader's stronger Send + Sync
209// assertion so a future non-Send field cannot silently narrow that contract.
210const _: fn() = || {
211    fn assert_send<T: Send>() {}
212    assert_send::<crate::egress::ReaderQuery<'_>>();
213    assert_send::<crate::egress::Cursor<'_>>();
214    #[cfg(feature = "arrow-egress")]
215    assert_send::<crate::egress::arrow::CursorRecordBatchReader<'_, '_>>();
216    #[cfg(feature = "polars-egress")]
217    assert_send::<crate::egress::arrow::polars::CursorPolarsIter<'_, '_>>();
218};
219
220impl Reader {
221    /// Open a new connection from a connect string.
222    pub fn from_conf<T: AsRef<str>>(conf: T) -> Result<Self> {
223        let cfg = ReaderConfig::from_conf(conf)?;
224        Self::from_config(&cfg)
225    }
226
227    /// Open a new connection from the config string stored in the
228    /// `QDB_CLIENT_CONF` environment variable. Format matches [`Reader::from_conf`].
229    pub fn from_env() -> Result<Self> {
230        let conf = std::env::var("QDB_CLIENT_CONF").map_err(|e| match e {
231            std::env::VarError::NotPresent => {
232                fmt!(ConfigError, "Environment variable QDB_CLIENT_CONF not set.")
233            }
234            std::env::VarError::NotUnicode(_) => fmt!(
235                InvalidUtf8,
236                "Environment variable QDB_CLIENT_CONF is set but its value is not valid UTF-8."
237            ),
238        })?;
239        Self::from_conf(conf)
240    }
241
242    /// Walk `cfg.addrs` via the per-client host-health tracker, opening
243    /// the highest-priority unattempted endpoint and eagerly consuming
244    /// the `SERVER_INFO` frame. Accepts the first endpoint whose role
245    /// matches `cfg.target`. Returns:
246    ///
247    /// - `RoleMismatch` if every endpoint connected but none advertised
248    ///   a matching role (last-seen role surfaced in the message).
249    /// - `AuthError` if at least one endpoint 401/403'd and every other
250    ///   endpoint failed too (per-endpoint accumulation lets the message
251    ///   name every endpoint that rejected credentials).
252    /// - `SocketError` if every endpoint failed at the transport layer
253    ///   (refused / timed out / TLS error / etc.).
254    /// - whatever the last attempt returned otherwise.
255    ///
256    /// The initial connect deliberately does **not** apply the egress
257    /// failover backoff schedule — it walks every address once and
258    /// reports back. Mid-query failover (via [`Cursor::next_batch`]) is
259    /// what uses `failover_backoff_*` to space retries.
260    ///
261    /// The tracker is constructed fresh here, so every host starts at
262    /// `Unknown` state and the priority-based pick degenerates to the
263    /// user-supplied `addr=` order. From this Reader onward, the
264    /// tracker accumulates classifications across Executes per the
265    /// failover.md §2 priority lattice.
266    pub fn from_config(cfg: &ReaderConfig) -> Result<Self> {
267        // Re-run cap and consistency checks. `from_conf` validated at
268        // parse time, but `ReaderConfig`'s `pub` fields can be mutated
269        // post-parse (`#[non_exhaustive]` blocks struct-literal
270        // construction, not field assignment), so a caller could
271        // otherwise sneak a `failover_backoff_max_ms = u64::MAX` past
272        // the parse-time hard cap and induce multi-day `thread::sleep`s
273        // during a failover storm.
274        cfg.validate()?;
275        // Single deep clone at the API boundary. Every subsequent
276        // reconnect attempt — initial walk, mid-query failover —
277        // shares the same allocation via `Arc::clone`.
278        let cfg = Arc::new(cfg.clone());
279        // Wire the `zone=` knob and the `target=primary` flag into the
280        // tracker. Per failover.md §2, `target=primary` collapses every
281        // host's zone tier to `Same` regardless of `zone=` — writers
282        // must be followed across zones — so we pass the bool through.
283        // Comparison against `SERVER_INFO.zone_id` / `X-QuestDB-Zone`
284        // is case-insensitive and lives inside `HostHealthTracker`.
285        let mut tracker = HostHealthTracker::new(
286            cfg.addrs.len(),
287            cfg.zone.as_deref(),
288            matches!(cfg.target, Target::Primary),
289        );
290        let walk = walk_via_tracker(
291            &mut tracker,
292            &cfg,
293            // Initial connect: no fall-through reset — every host
294            // starts at `Unknown`, so a single pass exhausts the list.
295            // Failover.md §2.2 / spec §11.9.3: the retry-after-reset
296            // pass is only meaningful when classifications have
297            // accumulated, which doesn't happen on a fresh tracker.
298            false,
299            // Spec §6 / §11.9.3 WalkTracker pseudocode: `AuthError`
300            // is terminal — credentials are cluster-wide, retrying
301            // every host floods server logs without recovery. Matches
302            // the Java reference's `connect()` which rethrows on
303            // `QwpAuthFailedException` immediately.
304            &[
305                ErrorCode::ConfigError,
306                ErrorCode::UnsupportedServer,
307                ErrorCode::AuthError,
308            ],
309        )?;
310        Ok(Reader {
311            cfg,
312            addr_idx: walk.session.idx,
313            transport: Some(walk.session.transport),
314            dict: SymbolDict::new(),
315            query_schema: None,
316            next_request_id: 1,
317            cursor_active: false,
318            server_info: walk.session.server_info,
319            stats: Arc::new(ReaderStats::default()),
320            zstd_scratch: ZstdScratch::new(),
321            tracker,
322            failover_rng: FailoverRng::new(),
323        })
324    }
325
326    /// Open a single endpoint by index. Used by [`walk_via_tracker`] on
327    /// both initial connect and mid-query failover. On success, returns
328    /// a [`TransportSession`] holding the bound socket plus the
329    /// `SERVER_INFO` (when applicable); the caller decides whether to
330    /// wrap it in a fresh `Reader` (initial connect) or splice into an
331    /// existing one (reconnect). On role mismatch, a `RoleMismatch`
332    /// error carrying the observed role + zone via `UpgradeReject` is
333    /// surfaced so the tracker can classify identically to a `421`
334    /// upgrade reject.
335    fn connect_endpoint(cfg: &ReaderConfig, idx: usize) -> Result<TransportSession> {
336        let mut transport = WsTransport::connect_to(cfg, idx).map_err(|e| {
337            // Prepend the endpoint so a connect/handshake/auth failure
338            // names the host it came from. Without this, aggregated
339            // multi-endpoint diagnostics surface only the tungstenite
340            // message ("HTTP error: 401") with no way to tell which
341            // endpoint refused.
342            let endpoint = &cfg.addrs[idx];
343            let mut annotated = Error::new(e.code(), format!("endpoint {}: {}", endpoint, e.msg()));
344            if let Some(r) = e.upgrade_reject() {
345                annotated = annotated.with_upgrade_reject(r.clone());
346            }
347            if let Some(info) = e.server_info() {
348                annotated = annotated.with_server_info(info.clone());
349            }
350            annotated
351        })?;
352        let server_info = if transport.server_version() >= 1 {
353            Some(read_server_info_frame(
354                &mut transport,
355                Duration::from_millis(cfg.server_info_timeout_ms),
356            )?)
357        } else {
358            None
359        };
360        if !matches!(cfg.target, Target::Any) {
361            match server_info.as_ref() {
362                None => {
363                    // No SERVER_INFO was supplied, so there's no wire role
364                    // to match against `target`. Surface a plain
365                    // `RoleMismatch` without `UpgradeReject` — there's no
366                    // role or zone to attach. With the single QWP version
367                    // (which always sends SERVER_INFO) this is unreachable
368                    // for a conformant server; it remains as a guard.
369                    return Err(fmt!(
370                        RoleMismatch,
371                        "endpoint {} supplied no SERVER_INFO and cannot match target={:?}",
372                        idx,
373                        cfg.target
374                    ));
375                }
376                Some(info) if !target_matches(cfg.target, info.role) => {
377                    // The endpoint advertised a role that doesn't match `target=`.
378                    // Attach `UpgradeReject` carrying the advertised role
379                    // and zone so the host-health tracker classifies
380                    // identically to a `421+role` response — same
381                    // semantics, same data payload, regardless of which
382                    // surface the rejection arrived on.
383                    //
384                    // Also attach the full `SERVER_INFO` so callers can
385                    // see the cluster/node identity of the last endpoint
386                    // that refused (wire-egress.md §11.9.3): `epoch`,
387                    // `cluster_id`, `node_id`, `capabilities`,
388                    // `server_wall_ns` — none of which fit on
389                    // `UpgradeReject`. Lets operators distinguish "no
390                    // endpoint matched target=" from "all endpoints
391                    // unreachable".
392                    let role = info.role;
393                    let role_name = role.as_str();
394                    let reject =
395                        UpgradeReject::new(role.as_u8(), role_name.clone(), info.zone_id.clone());
396                    return Err(Error::new(
397                        ErrorCode::RoleMismatch,
398                        format!(
399                            "endpoint {} role={} cluster={:?} does not match target={:?}",
400                            idx, role_name, info.cluster_id, cfg.target,
401                        ),
402                    )
403                    .with_upgrade_reject(reject)
404                    .with_server_info(info.clone()));
405                }
406                _ => {}
407            }
408        }
409        Ok(TransportSession {
410            idx,
411            transport,
412            server_info,
413        })
414    }
415
416    /// Reconnect this Reader in place after a mid-query transport
417    /// failure. Walks the configured endpoint list via the per-client
418    /// [`HostHealthTracker`] (failover.md §2 priority lattice — Healthy
419    /// → Unknown → TransientReject → TransportError → TopologyReject;
420    /// same-zone preferred when zone is configured). On success, the
421    /// old transport has been closed, the new transport + `SERVER_INFO`
422    /// are bound, the symbol dict and per-query schema are reset to
423    /// empty, and `addr_idx` reflects the new endpoint. The caller must
424    /// re-issue the
425    /// `QUERY_REQUEST` with a freshly-allocated `request_id`.
426    ///
427    /// The `failed_idx` argument is the address index that just failed
428    /// — `record_mid_stream_failure` demotes it from `Healthy` to
429    /// `TransportError` so the tracker won't reach for it first on the
430    /// next walk.
431    /// `budget` is cursor-owned and spans the whole `Execute()` call:
432    /// reconnect rounds, backoff growth, and the wall-clock deadline do
433    /// not reset after a successful replay.
434    ///
435    /// `on_attempt` is invoked once per reconnect round right before
436    /// the `walk_via_tracker` dial runs (after the configured
437    /// post-failure backoff sleep, so the wall-clock cost of the
438    /// backoff is included in the elapsed measurement the caller
439    /// derives). Passed by `&mut dyn` instead of generic `impl FnMut`
440    /// so adding the hook doesn't monomorphise this large function per
441    /// call site — there is one non-trivial caller
442    /// (`Cursor::failover_reconnect_and_replay`).
443    fn reconnect_with_failover(
444        &mut self,
445        failed_idx: usize,
446        budget: &mut FailoverBudget,
447        on_attempt: &mut dyn FnMut(u32),
448    ) -> Result<u32> {
449        let cfg = Arc::clone(&self.cfg);
450        let mut last_err: Option<Error> = None;
451        let mut deadline_exhausted = false;
452        // Spec invariant (failover.md §2.3): mid-stream demote MUST run
453        // before the next `begin_round(forget=true)` — reversing the
454        // order would let sticky-Healthy preserve the just-failed host
455        // as priority pick. `walk_via_tracker` only calls
456        // `begin_round(true)` on the fall-through reset, never before
457        // the first `pick_next`, but the demote still has to land
458        // before any walk so the first `pick_next` skips the dead host.
459        self.tracker.record_mid_stream_failure(failed_idx);
460        // Drop the dead transport entirely **before** sleeping on the
461        // backoff. `Drop for WsTransport` already issues a fire-and-
462        // forget WS Close, so the explicit `drop(dead)` is what
463        // releases the underlying TCP FD. Without this `take`, every
464        // reconnect attempt against a dead cluster would hold the
465        // dead FD for the whole
466        // `failover_max_attempts × failover_backoff_max_ms` window.
467        if let Some(dead) = self.transport.take() {
468            drop(dead);
469        }
470        // Cumulative dial count across every outer attempt's walk.
471        // `FailoverResetEvent.attempts` carries this back to the user so
472        // long-running diagnostics see real dial pressure, not just the
473        // attempt index that landed.
474        let mut total_dials: u32 = 0;
475        // Per-failure reconnect counter — i.e. how many
476        // `walk_via_tracker` rounds this call fired. Distinct from the
477        // cursor-level budget because this call can enter with only
478        // part of the original per-Execute budget left.
479        let mut attempts_made: u32 = 0;
480        loop {
481            match budget.before_reconnect_round(&cfg, &mut self.failover_rng) {
482                Ok(()) => {}
483                Err(FailoverBudgetStop::AttemptsExhausted) => break,
484                Err(FailoverBudgetStop::DeadlineExhausted) => {
485                    deadline_exhausted = true;
486                    break;
487                }
488            }
489            // Count the attempt only after the shared budget gate above
490            // has let us through; otherwise we'd over-report attempts
491            // in exhaustion messages.
492            attempts_made = attempts_made.saturating_add(1);
493            // Fire the per-attempt hook *after* the deadline gate (so
494            // the count we report matches the one the exhaustion errors
495            // report) and *before* the dial (so observers see "about
496            // to dial attempt N for this failure" rather than
497            // retroactive "dial N finished"). Pass the 1-based attempt
498            // number; the caller already knows the trigger and start
499            // time.
500            on_attempt(attempts_made);
501            match walk_via_tracker(
502                &mut self.tracker,
503                &cfg,
504                // Per failover.md §11.9.3, the WalkTracker fall-through
505                // reset pass is for reconnects only — gives stale
506                // `TransientReject` / `TopologyReject` hosts from prior
507                // outages another shot before declaring the walk failed.
508                true,
509                // Spec §6: AuthError is terminal during reconnect
510                // (cluster-wide credentials problem; retrying every
511                // host floods server logs without recovery). Initial
512                // connect accumulates instead — see `from_config`.
513                &[
514                    ErrorCode::ConfigError,
515                    ErrorCode::UnsupportedServer,
516                    ErrorCode::AuthError,
517                ],
518            ) {
519                Ok(walk) => {
520                    total_dials = total_dials.saturating_add(walk.dials);
521                    // Splice the new transport state into self, keeping
522                    // the counters callers query
523                    // (`bytes_received`, `credit_granted_total`,
524                    // `read_ns`, `decode_ns`, `next_request_id`).
525                    self.transport = Some(walk.session.transport);
526                    self.server_info = walk.session.server_info;
527                    self.dict = SymbolDict::new();
528                    self.query_schema = None;
529                    self.addr_idx = walk.session.idx;
530                    return Ok(total_dials);
531                }
532                Err(e) => match e.code() {
533                    code if !is_failover_eligible(code) => {
534                        // Hard error (auth, config, unsupported server,
535                        // etc.). Don't keep bouncing — these will fail
536                        // identically on every endpoint.
537                        return Err(e);
538                    }
539                    _ => {
540                        warn_on_protocol_error_failover(&e, "reconnect walk");
541                        last_err = Some(e);
542                    }
543                },
544            }
545        }
546        if deadline_exhausted {
547            let last_msg = last_err
548                .as_ref()
549                .map(|e| e.msg().to_string())
550                .unwrap_or_else(|| "<no error captured>".to_string());
551            return Err(fmt!(
552                SocketError,
553                "failover wall-clock budget exhausted (failover_max_duration_ms={}) after {} attempt(s); last error: {}",
554                cfg.failover_max_duration_ms,
555                attempts_made,
556                last_msg
557            ));
558        }
559        Err(last_err.unwrap_or_else(|| {
560            // Report the attempts this call actually ran; the
561            // cursor-level cap may have been partly spent by earlier
562            // successful failovers in the same Execute.
563            fmt!(
564                SocketError,
565                "failover exhausted after {} attempts",
566                attempts_made
567            )
568        }))
569    }
570
571    /// The endpoint this connection is currently bound to. Borrowed
572    /// from the configured address list, so the borrow lives as long
573    /// as `&self`. Stable across connect-string reorderings, unlike
574    /// the (deliberately not exposed) underlying address-list index.
575    pub fn current_addr(&self) -> &Endpoint {
576        &self.cfg.addrs[self.addr_idx]
577    }
578
579    /// Mutable access to the live transport. Returns `SocketError`
580    /// when the transport is `None`, which happens after the connection
581    /// was torn down: either a cursor was dropped before being fully
582    /// read (drop closes the WebSocket — see [`Cursor`]'s docs), or a
583    /// mid-query failover exhausted its retry budget. Either way the
584    /// Reader is "poisoned"; the fix is to drain cursors (`next_batch()`
585    /// until `None`) or `Cursor::cancel()` before dropping them, or to
586    /// open a fresh Reader. Inside `reconnect_with_failover` the transport
587    /// is only briefly absent (between dropping the dead one and splicing
588    /// in a new one); that path uses `self.transport` directly and never
589    /// goes through this accessor.
590    fn transport_mut(&mut self) -> Result<&mut WsTransport> {
591        self.transport.as_mut().ok_or_else(|| {
592            fmt!(
593                SocketError,
594                "Reader connection is closed and cannot be reused: a cursor was dropped before being \
595                 fully read, or a mid-query failover exhausted its retry budget. To keep the \
596                 connection reusable, drain each cursor (call next_batch() until it returns None) or \
597                 call cursor.cancel() before dropping it; otherwise open a fresh Reader."
598            )
599        })
600    }
601
602    /// Read access to the live transport. See [`Reader::transport_mut`].
603    fn transport_ref(&self) -> Result<&WsTransport> {
604        self.transport.as_ref().ok_or_else(|| {
605            fmt!(
606                SocketError,
607                "Reader connection is closed and cannot be reused: a cursor was dropped before being \
608                 fully read, or a mid-query failover exhausted its retry budget. To keep the \
609                 connection reusable, drain each cursor (call next_batch() until it returns None) or \
610                 call cursor.cancel() before dropping it; otherwise open a fresh Reader."
611            )
612        })
613    }
614
615    /// Allocate the next `request_id`, skipping `0` and negatives on
616    /// wrap. `0` is the server-side sentinel for "no active streaming
617    /// request" and must never be used by the client.
618    fn alloc_request_id(&mut self) -> i64 {
619        let id = self.next_request_id;
620        let next = self.next_request_id.wrapping_add(1);
621        self.next_request_id = if next <= 0 { 1 } else { next };
622        id
623    }
624
625    /// Total wire bytes (frame header + payload) read off the transport
626    /// since this connection was opened. Useful for benchmarking the
627    /// effective throughput a query produces.
628    pub fn bytes_received(&self) -> u64 {
629        self.stats.bytes_received.load(Ordering::Relaxed)
630    }
631
632    /// `true` when the underlying transport has been torn down (mid-stream
633    /// cursor abandonment, fatal socket error, role-mismatch failover that
634    /// couldn't find a replacement). Pool return paths should treat such a
635    /// reader as must-close.
636    pub fn transport_torn_down(&self) -> bool {
637        self.transport.is_none()
638    }
639
640    /// Total bytes granted to the server via CREDIT (`0x15`) frames
641    /// since this connection was opened. Useful for verifying that
642    /// flow-control replenishment behaves as expected — in particular,
643    /// that `Cursor::cancel()` doesn't continue topping up the server's
644    /// budget while draining frames it's about to discard.
645    pub fn credit_granted_total(&self) -> u64 {
646        self.stats.credit_granted_total.load(Ordering::Relaxed)
647    }
648
649    /// Diagnostic accumulator (nanoseconds): time spent in
650    /// `transport.read_frame()`. Saturates at `u64::MAX` (~584 years).
651    /// Reset to zero by [`Reader::reset_timing`].
652    pub fn read_ns(&self) -> u64 {
653        self.stats.read_ns.load(Ordering::Relaxed)
654    }
655    /// Diagnostic accumulator (nanoseconds): time spent in
656    /// `decode_frame()`. Saturates at `u64::MAX`.
657    /// Reset to zero by [`Reader::reset_timing`].
658    pub fn decode_ns(&self) -> u64 {
659        self.stats.decode_ns.load(Ordering::Relaxed)
660    }
661    /// Reset both `read_ns` and `decode_ns` accumulators to zero.
662    pub fn reset_timing(&self) {
663        self.stats.read_ns.store(0, Ordering::Relaxed);
664        self.stats.decode_ns.store(0, Ordering::Relaxed);
665    }
666
667    /// Borrow the shared diagnostic counters. The FFI clones this at
668    /// `qwp_reader_from_conf` time so its stat getters can read the
669    /// counters without touching the `UnsafeCell<Reader>` that holds
670    /// this Reader — eliminating the aliasing question of "what
671    /// happens when a stat getter synthesises a `&Reader` while a
672    /// laundered `&mut Reader` is in flight."
673    pub fn stats(&self) -> &Arc<ReaderStats> {
674        &self.stats
675    }
676
677    /// `SERVER_INFO` (`0x18`) captured at connect time. `None` only while
678    /// a reconnect is in flight; the single QWP version always supplies it.
679    pub fn server_info(&self) -> Option<&ServerInfo> {
680        self.server_info.as_ref()
681    }
682
683    /// Negotiated QWP version this connection is using. Returns
684    /// `SocketError` when the Reader is poisoned after a failed
685    /// mid-query failover.
686    pub fn server_version(&self) -> Result<u8> {
687        Ok(self.transport_ref()?.server_version())
688    }
689
690    /// Connection-scoped symbol dictionary.
691    pub fn symbol_dict(&self) -> &SymbolDict {
692        &self.dict
693    }
694
695    /// Begin building a parametrised query. The returned `ReaderQuery`
696    /// exclusively borrows the reader; only one in-flight cursor at a
697    /// time. Append binds in placeholder order, then call `.execute()`.
698    pub fn prepare<S: Into<String>>(&mut self, sql: S) -> ReaderQuery<'_> {
699        ReaderQuery {
700            reader: self,
701            builder: QueryRequest::builder(sql),
702            reset_symbol_dict: false,
703            on_failover_reset: None,
704            on_failover_progress: None,
705        }
706    }
707
708    /// Execute a SQL statement with no binds and return a streaming
709    /// cursor. Convenience for `self.prepare(sql).execute()`.
710    pub fn execute<S: Into<String>>(&mut self, sql: S) -> Result<Cursor<'_>> {
711        self.prepare(sql).execute()
712    }
713}
714
715// ---------------------------------------------------------------------------
716// Query builder
717// ---------------------------------------------------------------------------
718
719/// Notification delivered to the [`ReaderQuery::on_failover_reset`]
720/// callback right before replayed batches start arriving on a new
721/// connection. Mirrors the Java `onFailoverReset(newNode)` contract:
722/// the user-side handler is responsible for discarding any rows it
723/// had accumulated from the previous (now-dead) connection, since the
724/// query restarts from `batch_seq=0` against the new endpoint.
725///
726/// Marked `#[non_exhaustive]` so we can add fields without breaking
727/// downstream pattern matches.
728#[derive(Debug, Clone)]
729#[non_exhaustive]
730pub struct FailoverResetEvent {
731    /// Endpoint that just failed. Use `failed_addr.host` /
732    /// `failed_addr.port` directly; the [`Endpoint`] struct replaces
733    /// the older `(String, u16)` tuple.
734    ///
735    /// The address-list index is deliberately not exposed: indices
736    /// are brittle if the connect string is reordered between runs,
737    /// and the endpoint host/port is stable.
738    pub failed_addr: Endpoint,
739    /// Endpoint of the new connection.
740    pub new_addr: Endpoint,
741    /// `SERVER_INFO` of the new endpoint (`None` only if the server
742    /// omitted it).
743    pub new_server_info: Option<ServerInfo>,
744    /// Newly-allocated `request_id` the cursor will receive frames for
745    /// from now on. Different from `Cursor::request_id` *before* the
746    /// failover.
747    pub new_request_id: i64,
748    /// Count of reconnect dials the current failover cycle burned
749    /// before this success. `1` means the first reconnect dial
750    /// succeeded and its replay write went through cleanly. Larger
751    /// values mean earlier dials in this cycle missed (rotating
752    /// through endpoints) before one landed. Pairs with
753    /// [`elapsed`](Self::elapsed) — both measure the same failover
754    /// event.
755    pub attempts: u32,
756    /// The error that triggered this failover (the failure of the
757    /// previous connection). The full error — code + message — is
758    /// preserved so callers can both route on the [`ErrorCode`] (for
759    /// metrics / categorization) and log the raw message (for
760    /// diagnostics: `errno` text on `SocketError`, peer info on
761    /// `TlsError`, decode-site detail on `ProtocolError`, etc.). Use
762    /// [`Error::code`] to extract just the category.
763    ///
764    /// Without this, the cause-of-death of the previous connection is
765    /// lost forever once failover succeeds — it's not re-surfaced as
766    /// `Err` anywhere else in the cursor's API.
767    pub trigger: Error,
768    /// Wall-clock time spent reconnecting (sleep + dial + handshake +
769    /// SERVER_INFO read). Excludes the time from the cursor's last
770    /// successful read until the failure was observed.
771    pub elapsed: std::time::Duration,
772}
773
774/// Boxed user callback type for failover-reset notifications.
775type FailoverResetCallback<'r> = Box<dyn FnMut(&FailoverResetEvent) + Send + 'r>;
776
777/// Phase discriminant on [`FailoverProgressEvent`].
778///
779/// The same callback fires for every phase of a mid-query failover —
780/// from the moment the cursor's connection dies through to either a
781/// successful reconnect or an exhausted retry budget. Operators can
782/// route on the phase to feed SLO dashboards ("disconnected for N
783/// seconds" alerts), per-attempt retry telemetry, or a one-shot
784/// "gave up" notifier.
785///
786/// Marked `#[non_exhaustive]` so we can add phases (e.g. a hypothetical
787/// `Cancelled` for cancel-during-failover races) without breaking
788/// downstream matches.
789#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
790#[non_exhaustive]
791pub enum FailoverPhase {
792    /// The cursor's connection just died. Fires once, *before* the
793    /// retry loop runs.
794    Disconnected = 0,
795    /// A reconnect dial is about to be attempted. Fires once per
796    /// outer-loop iteration of the retry walk, *after* the inter-
797    /// attempt backoff sleep has elapsed.
798    Retrying = 1,
799    /// A reconnect succeeded; replayed batches will start arriving on
800    /// the new connection. Fires immediately *before* the
801    /// [`ReaderQuery::on_failover_reset`] callback (when both are
802    /// installed) so a single sink sees the entire lifecycle.
803    Reset = 2,
804    /// The retry budget is exhausted. The cursor is terminal; the
805    /// error returned to the caller is in
806    /// [`FailoverProgressEvent::final_error`].
807    GaveUp = 3,
808}
809
810/// Notification delivered to the
811/// [`ReaderQuery::on_failover_progress`] callback at each transition
812/// of a mid-query failover lifecycle. See [`FailoverPhase`] for the
813/// per-variant semantics.
814///
815/// Several fields are populated only in certain phases — see the
816/// per-field docs. Marked `#[non_exhaustive]` so we can add fields
817/// without breaking downstream pattern matches.
818#[derive(Debug, Clone)]
819#[non_exhaustive]
820pub struct FailoverProgressEvent {
821    /// Which lifecycle phase fired this event.
822    pub phase: FailoverPhase,
823    /// Endpoint that died. Set on every phase — even `Reset` keeps it
824    /// so a single sink can correlate the failed/new pair without
825    /// remembering state across calls.
826    pub failed_addr: Endpoint,
827    /// New endpoint the cursor is now bound to. `Some` only on
828    /// [`FailoverPhase::Reset`].
829    pub new_addr: Option<Endpoint>,
830    /// `SERVER_INFO` of the new endpoint. `Some` only on
831    /// [`FailoverPhase::Reset`].
832    pub new_server_info: Option<ServerInfo>,
833    /// Newly-allocated `request_id`. `Some` only on
834    /// [`FailoverPhase::Reset`].
835    pub new_request_id: Option<i64>,
836    /// 1-based attempt counter:
837    ///
838    /// - `0` on `Disconnected` (no attempt yet).
839    /// - `N ≥ 1` on `Retrying` for the Nth dial.
840    /// - On `Reset`, the attempt that landed.
841    /// - On `GaveUp`, the total number of attempts burned. May be `0`
842    ///   when the wall-clock deadline was already exhausted before any
843    ///   walk fired.
844    pub attempt: u32,
845    /// The error that triggered the failover (the original
846    /// cause-of-death of the previous connection). Preserved across
847    /// every phase so subscribers see consistent context regardless of
848    /// when they latch on.
849    pub trigger: Error,
850    /// Wall-clock time since the disconnect was observed (the start of
851    /// the failover cycle). Monotonically non-decreasing across phases
852    /// of the same event.
853    pub elapsed: std::time::Duration,
854    /// Final error returned to the caller. `Some` only on
855    /// [`FailoverPhase::GaveUp`]; this is the value the next call to
856    /// [`Cursor::next_batch`] (or `add_credit`) will surface.
857    pub final_error: Option<Error>,
858}
859
860/// Boxed user callback type for failover-progress notifications.
861type FailoverProgressCallback<'r> = Box<dyn FnMut(&FailoverProgressEvent) + Send + 'r>;
862
863/// Borrows a `Reader` exclusively while the query is being constructed and
864/// (eventually) the cursor is live.
865///
866/// `ReaderQuery` is [`Send`], but not safe for concurrent access. It may be
867/// moved to another thread after an explicit happens-before hand-off. Any
868/// installed failover callback must therefore also be [`Send`]; it runs on
869/// whichever thread subsequently drives the cursor.
870#[must_use = "ReaderQuery does nothing until you call .execute(); dropping it discards \
871              the prepared SQL and any binds without sending a QUERY_REQUEST"]
872pub struct ReaderQuery<'r> {
873    reader: &'r mut Reader,
874    builder: QueryRequestBuilder,
875    /// Request a query-scoped SYMBOL dict reset; translated to a
876    /// `query_flags` trailer at [`Self::execute`] iff the server advertised
877    /// `CAP_QUERY_FLAGS`.
878    reset_symbol_dict: bool,
879    /// Optional handler called every time the cursor reconnects after a
880    /// transport-level failure (see [`FailoverResetEvent`]).
881    on_failover_reset: Option<FailoverResetCallback<'r>>,
882    /// Optional progress handler invoked at every phase of a mid-query
883    /// failover lifecycle — see [`FailoverProgressEvent`] /
884    /// [`FailoverPhase`].
885    on_failover_progress: Option<FailoverProgressCallback<'r>>,
886}
887
888macro_rules! bind_method {
889    ($name:ident, $($arg:ident : $ty:ty),*) => {
890        pub fn $name(mut self, $($arg : $ty),*) -> Self {
891            // Manually re-assign because QueryRequestBuilder consumes self.
892            self.builder = self.builder.$name($($arg),*);
893            self
894        }
895    };
896}
897
898impl<'r> ReaderQuery<'r> {
899    /// Override the `initial_credit` (bytes; `0` = unbounded).
900    pub fn initial_credit(mut self, credit: u64) -> Self {
901        self.builder = self.builder.initial_credit(credit);
902        self
903    }
904
905    /// Request a query-scoped SYMBOL dict: the server resets the connection
906    /// dict before streaming this query so it never inherits symbols from
907    /// earlier queries on the same connection. Silently no-op against a server
908    /// that does not advertise `CAP_QUERY_FLAGS`.
909    pub fn reset_symbol_dict(mut self, reset: bool) -> Self {
910        self.reset_symbol_dict = reset;
911        self
912    }
913
914    /// Install a callback fired every time the cursor's underlying
915    /// connection is replaced via mid-query failover. The closure
916    /// receives a [`FailoverResetEvent`] describing the new endpoint and
917    /// runs *before* any replayed `RESULT_BATCH` arrives — the
918    /// user-side handler must use this signal to discard rows it had
919    /// accumulated from the previous (now-dead) connection. The query
920    /// restarts from `batch_seq=0` against the new endpoint with a
921    /// fresh `request_id`.
922    ///
923    /// **Installing this callback is the caller's opt-in to "I will
924    /// handle replay-after-data-delivered correctly."** Without it,
925    /// [`Cursor::next_batch`] refuses to fail over once any batch has
926    /// been yielded — returning
927    /// [`crate::ErrorCode::FailoverWouldDuplicate`]
928    /// instead — to avoid silently doubling up rows in the caller's
929    /// accumulator. Initial-connect failover (before any batch is
930    /// yielded) is transparent and does not require this callback.
931    ///
932    /// Calling this method twice on the same `ReaderQuery` **replaces**
933    /// the previous closure — only the most recent callback is invoked.
934    /// The callback must be [`Send`]: a query/cursor may be handed to
935    /// another thread, and the callback then runs and is dropped on that
936    /// destination thread. This bound is required even if the caller never
937    /// migrates the handle.
938    ///
939    /// Mirrors the Java client's `onFailoverReset(newNode)` contract.
940    ///
941    /// # Panics from the callback
942    ///
943    /// The callback is invoked synchronously from inside
944    /// [`Cursor::next_batch`] (specifically, from the failover-replay
945    /// path). If the callback panics, the unwind propagates through
946    /// `next_batch` to the caller. The cursor's [`Drop`] still runs,
947    /// which closes the WebSocket cleanly, so no resources are leaked
948    /// — but the `Cursor` is gone. There is no "swallow and resume"
949    /// behavior; treat a panicking callback as a bug and either
950    /// `catch_unwind` inside the callback yourself or ensure the
951    /// callback is panic-free. The C FFI binding wraps the callback in
952    /// `catch_unwind` + `abort()` (panics across the C boundary are
953    /// undefined behavior); the pure-Rust API leaves them as normal
954    /// unwinds.
955    ///
956    /// ```no_run
957    /// use std::sync::{Arc, Mutex};
958    /// use questdb::egress::{FailoverResetEvent, Reader};
959    ///
960    /// # fn ex() -> questdb::Result<()> {
961    /// let mut reader = Reader::from_conf(
962    ///     "ws::addr=db-a:9000,db-b:9000;target=primary",
963    /// )?;
964    /// // The handler accumulates rows in a buffer shared with the
965    /// // callback. On failover the callback discards what was buffered
966    /// // — the replayed query restarts at `batch_seq=0` against the
967    /// // new endpoint, so anything already pushed would otherwise
968    /// // double up.
969    /// let rows: Arc<Mutex<Vec<i64>>> = Arc::new(Mutex::new(Vec::new()));
970    /// let rows_for_cb = Arc::clone(&rows);
971    /// let mut cursor = reader
972    ///     .prepare("select x from t order by ts")
973    ///     .on_failover_reset(move |ev: &FailoverResetEvent| {
974    ///         eprintln!(
975    ///             "failover: {} → {} after {} attempt(s) ({:?}, trigger={:?}: {})",
976    ///             ev.failed_addr, ev.new_addr,
977    ///             ev.attempts, ev.elapsed,
978    ///             ev.trigger.code(), ev.trigger.msg(),
979    ///         );
980    ///         rows_for_cb.lock().unwrap().clear();
981    ///     })
982    ///     .execute()?;
983    /// while let Some(_batch) = cursor.next_batch()? {
984    ///     // ... project `_batch` into `rows.lock().unwrap()` ...
985    /// }
986    /// # let _ = rows; Ok(())
987    /// # }
988    /// ```
989    pub fn on_failover_reset<F>(mut self, callback: F) -> Self
990    where
991        F: FnMut(&FailoverResetEvent) + Send + 'r,
992    {
993        self.on_failover_reset = Some(Box::new(callback));
994        self
995    }
996
997    /// Install a callback fired at every phase of a mid-query failover
998    /// lifecycle: `Disconnected` when the cursor's connection dies,
999    /// `Retrying` before each reconnect dial attempt, `Reset` after a
1000    /// successful failover (immediately before
1001    /// [`Self::on_failover_reset`] runs), and `GaveUp` when the retry
1002    /// budget is exhausted.
1003    ///
1004    /// This callback is observational: installing it does **not** authorize
1005    /// replay after a batch has already reached the caller. Install
1006    /// [`Self::on_failover_reset`] as well when the caller can discard partial
1007    /// results safely. Without a reset callback, a post-delivery failure still
1008    /// returns [`ErrorCode::FailoverWouldDuplicate`].
1009    ///
1010    /// Calling this method twice on the same `ReaderQuery` **replaces**
1011    /// the previous closure — only the most recent callback is invoked.
1012    /// The callback must be [`Send`]: a query/cursor may be handed to
1013    /// another thread, and the callback then runs and is dropped on that
1014    /// destination thread. This bound is required even if the caller never
1015    /// migrates the handle.
1016    ///
1017    /// # Reentrancy
1018    ///
1019    /// The callback is invoked synchronously on the cursor's drive
1020    /// thread, while [`Cursor::next_batch`] (or `add_credit`) is
1021    /// mid-mutation of the underlying `Reader`. The same contract as
1022    /// [`Self::on_failover_reset`] applies:
1023    ///
1024    /// - **Must not** call back into the originating reader, query, or
1025    ///   cursor — including read-only stat getters.
1026    /// - **Must not** panic / `longjmp` / unwind across the boundary
1027    ///   (the FFI trampoline `catch_unwind` + `abort`s on escape).
1028    /// - **Must not** block indefinitely — every batch read, CREDIT
1029    ///   grant, and cancel waits until the callback returns.
1030    pub fn on_failover_progress<F>(mut self, callback: F) -> Self
1031    where
1032        F: FnMut(&FailoverProgressEvent) + Send + 'r,
1033    {
1034        self.on_failover_progress = Some(Box::new(callback));
1035        self
1036    }
1037
1038    /// Append a typed bind parameter.
1039    pub fn bind(mut self, value: Bind) -> Self {
1040        self.builder = self.builder.bind(value);
1041        self
1042    }
1043
1044    bind_method!(bind_null, kind: SimpleNullKind);
1045    bind_method!(bind_bool, v: bool);
1046    bind_method!(bind_i8, v: i8);
1047    bind_method!(bind_i16, v: i16);
1048    bind_method!(bind_i32, v: i32);
1049    bind_method!(bind_i64, v: i64);
1050    bind_method!(bind_f32, v: f32);
1051    bind_method!(bind_f64, v: f64);
1052    bind_method!(bind_timestamp_micros, v: i64);
1053    bind_method!(bind_timestamp_nanos, v: i64);
1054    bind_method!(bind_date_millis, v: i64);
1055    bind_method!(bind_uuid, v: [u8; 16]);
1056    bind_method!(bind_long256, v: [u8; 32]);
1057    bind_method!(bind_char, v: u16);
1058    bind_method!(bind_ipv4, v: Ipv4Addr);
1059
1060    pub fn bind_varchar<S: Into<String>>(mut self, v: S) -> Self {
1061        self.builder = self.builder.bind_varchar(v);
1062        self
1063    }
1064
1065    pub fn bind_decimal64(mut self, value: i64, scale: i8) -> Self {
1066        self.builder = self.builder.bind_decimal64(value, scale);
1067        self
1068    }
1069
1070    pub fn bind_decimal128(mut self, value: i128, scale: i8) -> Self {
1071        self.builder = self.builder.bind_decimal128(value, scale);
1072        self
1073    }
1074
1075    pub fn bind_decimal256(mut self, bytes: [u8; 32], scale: i8) -> Self {
1076        self.builder = self.builder.bind_decimal256(bytes, scale);
1077        self
1078    }
1079
1080    pub fn bind_geohash(mut self, value: u64, precision_bits: u8) -> Self {
1081        self.builder = self.builder.bind_geohash(value, precision_bits);
1082        self
1083    }
1084
1085    pub fn bind_binary<B: Into<Vec<u8>>>(mut self, v: B) -> Self {
1086        self.builder = self.builder.bind_binary(v);
1087        self
1088    }
1089
1090    pub fn bind_null_varchar(mut self) -> Self {
1091        self.builder = self.builder.bind_null_varchar();
1092        self
1093    }
1094
1095    pub fn bind_null_binary(mut self) -> Self {
1096        self.builder = self.builder.bind_null_binary();
1097        self
1098    }
1099
1100    pub fn bind_null_decimal64(mut self, scale: i8) -> Self {
1101        self.builder = self.builder.bind_null_decimal64(scale);
1102        self
1103    }
1104
1105    pub fn bind_null_decimal128(mut self, scale: i8) -> Self {
1106        self.builder = self.builder.bind_null_decimal128(scale);
1107        self
1108    }
1109
1110    pub fn bind_null_decimal256(mut self, scale: i8) -> Self {
1111        self.builder = self.builder.bind_null_decimal256(scale);
1112        self
1113    }
1114
1115    pub fn bind_null_geohash(mut self, precision_bits: u8) -> Self {
1116        self.builder = self.builder.bind_null_geohash(precision_bits);
1117        self
1118    }
1119
1120    /// Send the QUERY_REQUEST and return a streaming `Cursor`.
1121    pub fn execute(self) -> Result<Cursor<'r>> {
1122        if self.reader.cursor_active {
1123            return Err(fmt!(
1124                InvalidApiCall,
1125                "another cursor is already in flight on this connection (only one cursor at a time per Reader)"
1126            ));
1127        }
1128        let request_id = self.reader.alloc_request_id();
1129        // The schema rides the first RESULT_BATCH (batch_seq == 0) of each
1130        // query; clear any schema left from the prior query so a stray
1131        // continuation batch can't bind rows to a stale schema.
1132        self.reader.query_schema = None;
1133        // Cap-gate the query_flags trailer: only emit it when the server
1134        // advertised CAP_QUERY_FLAGS, so an older server sees the baseline
1135        // QUERY_REQUEST layout and the reset request silently degrades.
1136        let server_supports_query_flags = self
1137            .reader
1138            .server_info()
1139            .map(|info| has_query_flags(info.capabilities))
1140            .unwrap_or(false);
1141        let query_flags = if self.reset_symbol_dict && server_supports_query_flags {
1142            QUERY_FLAG_RESET_DICT
1143        } else {
1144            0
1145        };
1146        let req = self
1147            .builder
1148            .request_id(request_id)
1149            .query_flags(query_flags)
1150            .build()?;
1151        let credit_enabled = req.initial_credit() > 0;
1152        // Encode the QUERY_REQUEST once and stash the bytes on the
1153        // cursor. Mid-query failover replays the query by patching
1154        // the 8-byte `request_id` span in place and writing the same
1155        // buffer again — no builder clone, no bind clone, no
1156        // re-encode. The wire layout is:
1157        //   [0]   MsgKind::QueryRequest (1 byte)
1158        //   [1..9] request_id (i64 LE, 8 bytes)
1159        //   [9..]  varint sql_len, sql, varint initial_credit,
1160        //          varint binds_len, encoded binds...
1161        // Encoding can fail (e.g. an unsupported bind kind) — that
1162        // failure surfaces here and the cursor never starts.
1163        let mut encoded_request = Vec::with_capacity(64);
1164        req.encode(&mut encoded_request)?;
1165        // Layout invariant guard, runtime-checked in release too: the
1166        // failover-replay path patches `[REQUEST_ID_OFFSET..+8]` of
1167        // this buffer with a fresh request_id on every reconnect. If
1168        // `QueryRequest::encode` ever changes the prefix (adds a
1169        // length header, version byte, different MsgKind), patching
1170        // the wrong offset would silently corrupt every replayed
1171        // request — and the corruption surfaces as a `ProtocolError`
1172        // which is itself failover-eligible, so the cursor would
1173        // burn its retry budget bouncing through the cluster with
1174        // bad bytes. Fail loudly at execute() time instead.
1175        if encoded_request.len() < REQUEST_ID_OFFSET + 8
1176            || encoded_request[0] != MsgKind::QueryRequest.as_u8()
1177        {
1178            return Err(fmt!(
1179                ProtocolError,
1180                "QUERY_REQUEST encoding layout invariant violated (len={}, first={:?})",
1181                encoded_request.len(),
1182                encoded_request.first().copied(),
1183            ));
1184        }
1185        debug_assert_eq!(
1186            i64::from_le_bytes(
1187                encoded_request[REQUEST_ID_OFFSET..REQUEST_ID_OFFSET + 8]
1188                    .try_into()
1189                    .expect("length checked above"),
1190            ),
1191            request_id,
1192            "request_id at byte offset {} doesn't match the value just encoded",
1193            REQUEST_ID_OFFSET,
1194        );
1195        // Wrap the encoded request as Bytes once. `Bytes::from(Vec)` is
1196        // a zero-copy move; cloning a Bytes is a refcount bump so the
1197        // initial write and the stashed copy share one allocation.
1198        let encoded_request: Bytes = encoded_request.into();
1199        self.reader
1200            .transport_mut()?
1201            .write_message(encoded_request.clone())?;
1202
1203        self.reader.cursor_active = true;
1204        let failover_budget = FailoverBudget::new(&self.reader.cfg);
1205        Ok(Cursor {
1206            reader: self.reader,
1207            request_id,
1208            last_batch: None,
1209            terminal: None,
1210            credit_enabled,
1211            cancelling: false,
1212            done: false,
1213            terminal_error: None,
1214            encoded_request,
1215            on_failover_reset: self.on_failover_reset,
1216            on_failover_progress: self.on_failover_progress,
1217            failover_budget,
1218            failover_resets: 0,
1219            decode_failover_rounds: 0,
1220            stale_plan_retries: 0,
1221            data_delivered: false,
1222            #[cfg(feature = "arrow-egress")]
1223            drifted_batch: None,
1224            #[cfg(feature = "arrow-egress")]
1225            sym_values: crate::egress::arrow::SymbolValuesCache::default(),
1226            #[cfg(feature = "arrow-egress")]
1227            sym_scratch: crate::egress::arrow::SymbolBuildScratch::default(),
1228            #[cfg(feature = "polars-egress")]
1229            symbol_registry: None,
1230            #[cfg(feature = "polars-egress")]
1231            symbol_delta_modes: Vec::new(),
1232        })
1233    }
1234}
1235
1236/// Patch the request_id span of a stashed `QUERY_REQUEST` payload in
1237/// place and return it as fresh `Bytes`.
1238///
1239/// Fast path: `Bytes::try_into_mut` recovers the underlying `BytesMut`
1240/// zero-copy when the buffer is uniquely owned (the previous
1241/// `write_message` clone has been dropped). Patching mutates 8 bytes in
1242/// place, then `BytesMut::freeze` returns to `Bytes` zero-copy. The
1243/// multi-MB bind payload is never copied across reconnects.
1244///
1245/// Slow path: tungstenite still holds a reference (e.g., a partial write
1246/// flushed only after this routine ran). `try_into_mut` returns the
1247/// original `Bytes` back via `Err`; we fall back to a one-time
1248/// allocate-and-copy via `Bytes::copy_from_slice`. Same cost as the
1249/// pre-fix code, but unreachable in the steady state where every
1250/// `write_message` returns with the WS frame fully flushed.
1251fn patch_request_id(buf: Bytes, new_rid: i64) -> Bytes {
1252    let mut buf = match buf.try_into_mut() {
1253        Ok(buf_mut) => buf_mut,
1254        Err(shared) => BytesMut::from(&shared[..]),
1255    };
1256    buf[REQUEST_ID_OFFSET..REQUEST_ID_OFFSET + 8].copy_from_slice(&new_rid.to_le_bytes());
1257    buf.freeze()
1258}
1259
1260/// Bounded read timeout applied to the underlying TCP stream for the
1261/// duration of [`Cursor::cancel`]'s post-CANCEL drain.
1262///
1263/// Without this, a stuck-but-not-RST'd peer that stops sending bytes
1264/// after we deliver the CANCEL frame would block the drain
1265/// indefinitely. The drain consumes whatever batches the server
1266/// already had in flight plus the terminal QUERY_ERROR; under healthy
1267/// operation each frame arrives within milliseconds. 30 s is far past
1268/// any realistic batch transit and short enough that an unresponsive
1269/// peer surfaces a clear error rather than appearing to hang.
1270const CANCEL_DRAIN_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
1271
1272/// Dedicated per-Execute cap on failover rounds triggered by a *frame
1273/// decode* failure, as opposed to a raw transport (socket / TLS / WS)
1274/// read failure.
1275///
1276/// A decode failure is ambiguous: it can be transient wire corruption
1277/// (a truncated WS frame, a malformed varint emitted by a dying
1278/// endpoint) — which a single reconnect to a fresh connection cures —
1279/// or a *deterministic* protocol violation (unknown `MsgKind`,
1280/// mismatched lengths, a version-mismatched frame) that every replay
1281/// reproduces byte-for-byte. Both surface as `ProtocolError`, so they
1282/// are indistinguishable at the [`ErrorCode`] level.
1283///
1284/// Routing decode failures through the *full* per-Execute failover
1285/// budget lets a deterministically-corrupting server drive a
1286/// reconnect -> replay -> re-corrupt loop that drains every reconnect
1287/// round (`failover_max_attempts - 1`, up to 1023) — burning the
1288/// backoff schedule and emitting one warning per round — before the
1289/// identical decode error is finally surfaced. Capping decode-driven
1290/// replays at a small constant preserves recovery from a one-off
1291/// transient blip (one reconnect rules it out) while failing fast on
1292/// deterministic corruption. Raw read failures are unaffected and keep
1293/// the full budget.
1294const MAX_DECODE_FAILOVER_ROUNDS: u32 = 1;
1295
1296/// Bound on transparent re-issues of a query the server rejected with the
1297/// transient stale-cached-plan `INTERNAL_ERROR` (see [`is_stale_plan_error`]).
1298///
1299/// Each retry is a same-connection resend that makes the server recompile
1300/// the query against the table's *current* metadata; a single retry clears
1301/// the realistic one-`ALTER`-in-flight race. The cap stops a table under
1302/// relentless concurrent schema churn from looping forever — once spent, the
1303/// stale-plan error surfaces like any other server error. Retries are
1304/// naturally paced by the server's recompile + round-trip latency, so the
1305/// loop never busy-spins and no client-side sleep is needed.
1306const MAX_STALE_PLAN_RETRIES: u32 = 15;
1307
1308/// Classifies the origin of a mid-query stream failure routed through
1309/// [`Cursor::failover_after_stream_failure`]. Decode failures get a
1310/// small dedicated replay cap ([`MAX_DECODE_FAILOVER_ROUNDS`]); raw
1311/// transport read failures keep the full per-Execute budget.
1312#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1313enum StreamFailureKind {
1314    /// A raw transport read failed (socket closed, TLS reset, truncated
1315    /// WS frame at the transport layer).
1316    Read,
1317    /// A complete frame was read but `decode_frame` rejected it.
1318    Decode,
1319}
1320
1321impl StreamFailureKind {
1322    /// Human-readable context for the failover warning.
1323    fn context(self) -> &'static str {
1324        match self {
1325            StreamFailureKind::Read => "mid-query frame read",
1326            StreamFailureKind::Decode => "mid-query frame decode",
1327        }
1328    }
1329}
1330
1331// ---------------------------------------------------------------------------
1332// Cursor + BatchView
1333// ---------------------------------------------------------------------------
1334
1335#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1336enum FailoverBudgetStop {
1337    AttemptsExhausted,
1338    DeadlineExhausted,
1339}
1340
1341/// Mutable per-Execute failover budget.
1342///
1343/// A cursor may fail over, replay, then fail again before the query
1344/// reaches its terminal frame. The public knobs are per Execute, not
1345/// per outage, so reconnect rounds, backoff growth, and the wall-clock
1346/// deadline all live here on the cursor rather than inside
1347/// `Reader::reconnect_with_failover`.
1348struct FailoverBudget {
1349    reconnect_rounds_remaining: u32,
1350    next_backoff_ms: u64,
1351    deadline: Option<std::time::Instant>,
1352}
1353
1354impl FailoverBudget {
1355    fn new(cfg: &ReaderConfig) -> Self {
1356        let deadline = if cfg.failover_max_duration_ms == 0 {
1357            None
1358        } else {
1359            Some(std::time::Instant::now() + Duration::from_millis(cfg.failover_max_duration_ms))
1360        };
1361        Self {
1362            reconnect_rounds_remaining: cfg.failover_reconnect_rounds(),
1363            next_backoff_ms: cfg.failover_backoff_initial_ms,
1364            deadline,
1365        }
1366    }
1367
1368    fn advance_backoff(&mut self, max_backoff_ms: u64) {
1369        if self.next_backoff_ms > 0 {
1370            self.next_backoff_ms = self.next_backoff_ms.saturating_mul(2).min(max_backoff_ms);
1371        }
1372    }
1373
1374    fn before_reconnect_round(
1375        &mut self,
1376        cfg: &ReaderConfig,
1377        rng: &mut FailoverRng,
1378    ) -> std::result::Result<(), FailoverBudgetStop> {
1379        if self.reconnect_rounds_remaining == 0 {
1380            return Err(FailoverBudgetStop::AttemptsExhausted);
1381        }
1382
1383        // Failover.md §11.9 + §3.1: the first reconnect sleeps the
1384        // configured initial backoff, then subsequent rounds grow the
1385        // base exponentially up to the configured max. A zero initial
1386        // backoff is the documented "no sleeps" sentinel.
1387        let jittered_ms = rng.full_jitter_ms(self.next_backoff_ms);
1388        let sleep_dur = match self.deadline {
1389            Some(dl) => match dl.checked_duration_since(std::time::Instant::now()) {
1390                Some(remaining) if !remaining.is_zero() => {
1391                    std::cmp::min(Duration::from_millis(jittered_ms), remaining)
1392                }
1393                _ => return Err(FailoverBudgetStop::DeadlineExhausted),
1394            },
1395            None => Duration::from_millis(jittered_ms),
1396        };
1397        std::thread::sleep(sleep_dur);
1398        self.advance_backoff(cfg.failover_backoff_max_ms);
1399        self.reconnect_rounds_remaining = self.reconnect_rounds_remaining.saturating_sub(1);
1400        Ok(())
1401    }
1402}
1403
1404/// Reason the stream ended. Surfaced via [`Cursor::terminal`] once
1405/// `next_batch` returns `None`.
1406///
1407/// `#[non_exhaustive]` because future protocol revisions may add
1408/// terminal kinds (e.g. server-side timeouts).
1409#[derive(Debug, Clone)]
1410#[non_exhaustive]
1411pub enum Terminal {
1412    /// `RESULT_END` (`0x12`).
1413    End { final_seq: u64, total_rows: u64 },
1414    /// `EXEC_DONE` (`0x16`) — non-SELECT acknowledgement.
1415    ExecDone { op_type: u8, rows_affected: u64 },
1416}
1417
1418/// Streaming cursor over `RESULT_BATCH` frames.
1419///
1420/// `next_batch` advances the stream by one batch, returning `None` once a
1421/// terminal frame arrives (which is then accessible via [`Cursor::terminal`]).
1422/// `cancel` sends a `CANCEL` frame and drains until the server's terminal.
1423///
1424/// `Cursor` is [`Send`], but not safe for concurrent access. It may be moved
1425/// to another thread after an explicit happens-before hand-off. Failover
1426/// callbacks run on whichever thread drives the cursor.
1427#[must_use = "Cursor must be drained via next_batch() or cancelled via cancel(); \
1428              dropping mid-stream sends a best-effort CANCEL and closes the WebSocket, \
1429              tearing down the connection for the next query on this Reader"]
1430pub struct Cursor<'r> {
1431    reader: &'r mut Reader,
1432    request_id: i64,
1433    last_batch: Option<DecodedBatch>,
1434    terminal: Option<Terminal>,
1435    /// Pre-encoded `QUERY_REQUEST` payload from `execute()`, stashed
1436    /// so the cursor can resend the same query on a fresh connection
1437    /// after mid-query failover. The 8-byte `request_id` lives at
1438    /// `[REQUEST_ID_OFFSET..REQUEST_ID_OFFSET + 8]`; replay recovers
1439    /// `BytesMut` via [`Bytes::try_into_mut`], overwrites that span
1440    /// with a freshly-allocated id, and re-freezes — so the multi-MB
1441    /// `Bind::Binary` / `Bind::Varchar` payload is never copied
1442    /// across reconnects, only the 8-byte request_id span is mutated
1443    /// in place.
1444    encoded_request: Bytes,
1445    /// User callback fired right before replayed batches arrive on a
1446    /// new connection. See [`ReaderQuery::on_failover_reset`].
1447    on_failover_reset: Option<FailoverResetCallback<'r>>,
1448    /// User callback fired at every phase of a mid-query failover
1449    /// lifecycle. See [`ReaderQuery::on_failover_progress`].
1450    on_failover_progress: Option<FailoverProgressCallback<'r>>,
1451    /// Shared per-Execute budget for mid-query failover. This spans
1452    /// every reconnect in the cursor's life; it must not reset after a
1453    /// successful replay.
1454    failover_budget: FailoverBudget,
1455    /// Number of successful failover resets observed by this cursor
1456    /// since `execute()`. Useful for tests and for asserting the
1457    /// query did not silently restart under the user's feet.
1458    failover_resets: u32,
1459    /// Count of failover rounds this cursor has triggered specifically
1460    /// from a *frame decode* failure (as opposed to a raw transport
1461    /// read failure). Capped at [`MAX_DECODE_FAILOVER_ROUNDS`] so a
1462    /// deterministically-corrupting server can't drive a
1463    /// reconnect/replay loop that drains the whole per-Execute budget;
1464    /// once the cap is hit the decode error is surfaced terminally.
1465    decode_failover_rounds: u32,
1466    /// Count of transparent same-connection query re-issues this cursor has
1467    /// performed in response to the server's transient stale-cached-plan
1468    /// `INTERNAL_ERROR` (see [`Cursor::next_batch`] and
1469    /// [`is_stale_plan_error`]). Capped at [`MAX_STALE_PLAN_RETRIES`]; stays
1470    /// `0` on the happy path. Distinct from `failover_resets` — no reconnect
1471    /// is involved, the recompile happens on the existing healthy connection.
1472    stale_plan_retries: u32,
1473    /// Sticky: set the first time a `RESULT_BATCH` is yielded to the
1474    /// caller and never reset. Drives the safety check in
1475    /// [`Cursor::next_batch`] that refuses mid-query failover when no
1476    /// [`ReaderQuery::on_failover_reset`] callback is installed —
1477    /// silently replaying after the caller already received rows
1478    /// would deliver duplicates the caller has no way to detect.
1479    /// Distinct from `last_batch.is_some()`, which is cleared at the
1480    /// start of every replay; this flag must NOT reset, because the
1481    /// hazard is "the caller saw data at some point during this
1482    /// query," not "on the current connection."
1483    data_delivered: bool,
1484    /// `true` when the QUERY_REQUEST set `initial_credit > 0`. The
1485    /// cursor then auto-emits a CREDIT (`0x15`) frame after each
1486    /// RESULT_BATCH consumed, replenishing the server's per-request
1487    /// budget by exactly the wire size of the batch we just received
1488    /// (12-byte header + payload).
1489    credit_enabled: bool,
1490    /// Set once `cancel()` has written its CANCEL frame and entered the
1491    /// drain loop. Suppresses auto-credit replenishment for the rest of
1492    /// the cursor's life so the server's budget is allowed to drain to
1493    /// zero — this is the backpressure that hastens the post-cancel
1494    /// terminal. Without this, every drained batch would top the budget
1495    /// back up and the server could keep streaming at full rate until
1496    /// it finally observed the CANCEL on its input socket.
1497    cancelling: bool,
1498    /// Set once any terminal frame has been observed for this cursor:
1499    /// `RESULT_END`, `EXEC_DONE`, or `QUERY_ERROR` (including the
1500    /// `STATUS_CANCELLED` reply to `cancel()`). Also set on the
1501    /// failover-give-up path and on every other error-terminal in
1502    /// `next_batch`. Drives the early return in `next_batch()` so a
1503    /// follow-up call doesn't try to read another frame off a server
1504    /// that has already finished with this `request_id`. `terminal`
1505    /// (the public lifecycle accessor) only stores the success
1506    /// terminals; error terminals are stashed in `terminal_error`
1507    /// instead and re-raised from any subsequent `next_batch` /
1508    /// `add_credit` call so a transient-retry caller can't mistake
1509    /// an errored cursor for a clean RESULT_END.
1510    done: bool,
1511    /// `Some(err)` iff the cursor terminated with an error (failover
1512    /// give-up, server `QUERY_ERROR`, decode failure, stale-rid, etc).
1513    /// Clone-replayed by every public method that would otherwise
1514    /// short-circuit on `self.done` — without this, the first call
1515    /// surfaces the error and every subsequent call returns
1516    /// `Ok(None)`, looking indistinguishable from a clean RESULT_END
1517    /// to a caller with a retry-on-transient-error loop.
1518    ///
1519    /// Captured at most once (the first error wins) so a follow-up
1520    /// failure during teardown can't overwrite the originating cause.
1521    terminal_error: Option<Error>,
1522    /// A batch decoded but not handed out because its Arrow schema drifted
1523    /// from the pinned one, parked for replay on the next `next_arrow_batch*`
1524    /// call so the rows are recoverable rather than dropped.
1525    #[cfg(feature = "arrow-egress")]
1526    drifted_batch: Option<DecodedBatch>,
1527    /// Connection-dict SYMBOL values array, interned once per cursor and reused
1528    /// across batches until the dict grows (see [`SymbolValuesCache`]).
1529    #[cfg(feature = "arrow-egress")]
1530    sym_values: crate::egress::arrow::SymbolValuesCache,
1531    #[cfg(feature = "arrow-egress")]
1532    sym_scratch: crate::egress::arrow::SymbolBuildScratch,
1533    /// Per-cursor SYMBOL → polars `Categories`, interned once and grown
1534    /// incrementally across batches (see [`SymbolRegistry`]).
1535    #[cfg(feature = "polars-egress")]
1536    symbol_registry: Option<crate::egress::arrow::polars::SymbolRegistry>,
1537    /// `[i]` = column `i` is a delta-mode SYMBOL, captured per batch from the
1538    /// `DecodedBatch` before it is assembled.
1539    #[cfg(feature = "polars-egress")]
1540    symbol_delta_modes: Vec<bool>,
1541}
1542
1543/// Borrow-free outcome of `next_batch_inner`. The wrapper in
1544/// `next_batch` matches on this and constructs the public `BatchView`
1545/// (which holds borrows into `self`) only in the `HaveBatch` arm —
1546/// keeping the inner result borrow-free is what lets the `Err` arm
1547/// mutate `self.terminal_error` to stash the cursor-killing error
1548/// for replay on subsequent calls.
1549enum NextOutcome {
1550    HaveBatch,
1551    Done,
1552}
1553
1554impl<'r> Cursor<'r> {
1555    pub fn request_id(&self) -> i64 {
1556        self.request_id
1557    }
1558
1559    /// `Some` after a `RESULT_END` or `EXEC_DONE` has been observed.
1560    pub fn terminal(&self) -> Option<&Terminal> {
1561        self.terminal.as_ref()
1562    }
1563
1564    /// Whether dropping this cursor leaves its reader connection reusable.
1565    pub fn connection_reusable(&self) -> bool {
1566        self.done && !self.reader.transport_torn_down()
1567    }
1568
1569    /// Pass-through to [`Reader::credit_granted_total`]. Exists so
1570    /// callers holding the cursor's mutable borrow on the reader can
1571    /// still observe the connection-level CREDIT-bytes counter.
1572    pub fn credit_granted_total(&self) -> u64 {
1573        self.reader
1574            .stats
1575            .credit_granted_total
1576            .load(Ordering::Relaxed)
1577    }
1578
1579    /// Advance the cursor by one batch. Returns `Ok(None)` when the stream
1580    /// has terminated (success). `QUERY_ERROR` becomes `Err`.
1581    ///
1582    /// On a transport-level failure (socket close, TLS error, WS
1583    /// framing error), the cursor will reconnect to the next address
1584    /// in the configured list (with exponential backoff and a bounded
1585    /// retry budget — see `failover_*` config keys), replay the
1586    /// `QUERY_REQUEST` with a fresh `request_id`, and resume from
1587    /// `batch_seq=0` on the new connection. The user-side handler is
1588    /// notified before any replayed batches arrive via the
1589    /// [`ReaderQuery::on_failover_reset`] callback. If failover is
1590    /// disabled (`failover=off`) or the retry budget is exhausted,
1591    /// the failure is surfaced as the underlying error.
1592    ///
1593    /// **Silent-duplicate guard.** If a batch has already been
1594    /// yielded to the caller and no `on_failover_reset` callback was
1595    /// installed, the cursor refuses to fail over and returns
1596    /// [`crate::ErrorCode::FailoverWouldDuplicate`]
1597    /// instead. Replay would otherwise re-deliver rows the caller
1598    /// already consumed — with no signal — because the server
1599    /// restarts streaming from `batch_seq=0` on the new connection.
1600    /// Install the callback (and discard partial state on each
1601    /// invocation) to opt in to seeing replays; otherwise re-execute
1602    /// the query from scratch when this error fires. Failover that
1603    /// happens before the first batch is yielded — including initial
1604    /// connect failover — is unaffected and remains transparent.
1605    ///
1606    /// Failover-eligible decode errors (malformed payload, bad varint,
1607    /// zstd corruption) use the same reconnect-and-replay path as
1608    /// transport failures. Replaying after rows were already yielded is
1609    /// still blocked unless the caller installed a replay-aware callback,
1610    /// since the server restarts streaming from `batch_seq=0`.
1611    ///
1612    /// **Blocking time during failover.** When failover is engaged,
1613    /// this method blocks the calling thread for the duration of the
1614    /// reconnect cycle: each attempt sleeps the configured backoff
1615    /// (capped by `failover_backoff_max_ms`), then dials, handshakes,
1616    /// and reads `SERVER_INFO` against the next endpoint. The
1617    /// worst-case wall-clock blocking time is approximately
1618    /// `2 × (failover_max_attempts - 1) × failover_backoff_max_ms`
1619    /// plus per-attempt connect+handshake overhead — with the
1620    /// parse-time caps that's up to ~2 hours. There is no per-call timeout or
1621    /// AtomicBool cancel hook; use `on_failover_progress` for observability.
1622    /// If you need bounded latency, set `failover_max_attempts` and
1623    /// `failover_backoff_max_ms` to values appropriate for your SLA, or set
1624    /// `failover=off` and handle reconnect at the application layer.
1625    pub fn next_batch(&mut self) -> Result<Option<BatchView<'_>>> {
1626        // Replay-on-terminal guard. If the cursor previously terminated
1627        // with an error, surface that error on every subsequent call
1628        // rather than collapsing to `Ok(None)` (which is the clean-EOF
1629        // signal — a retry-on-transient-error caller would silently
1630        // treat an incomplete result set as complete).
1631        if self.done {
1632            return match self.terminal_error.as_ref() {
1633                Some(e) => Err(e.clone()),
1634                None => Ok(None),
1635            };
1636        }
1637        // Inner returns a borrow-free discriminant so the borrow
1638        // checker can split the lifetime — the Err arm needs to
1639        // mutate `self.terminal_error`, which it can't if the
1640        // inner result still holds a reference into `self`.
1641        // Capture is conditioned on `self.done` (set by every
1642        // error-terminal path, either directly or via
1643        // `terminate_with_close`) and on `terminal_error.is_none()`
1644        // so the FIRST cause wins — a follow-up teardown failure
1645        // can't overwrite the originating error.
1646        match self.next_batch_inner() {
1647            Ok(NextOutcome::HaveBatch) => {
1648                // `next_batch_inner` populates `last_batch` (via `.insert`)
1649                // and verifies `query_schema` is `Some` before returning
1650                // `HaveBatch`, so both are present here. Re-check with the
1651                // inner's *soft* pattern rather than `.expect()`: a panic
1652                // would abort the whole process across the FFI boundary
1653                // (`panic=abort`), so a future refactor that breaks the
1654                // invariant must surface a terminal `ProtocolError`, not
1655                // kill the host.
1656                if self.last_batch.is_none() || self.reader.query_schema.is_none() {
1657                    let err = fmt!(
1658                        ProtocolError,
1659                        "internal invariant: next_batch produced a batch without a decoded view or schema"
1660                    );
1661                    self.terminate_with_close();
1662                    if self.done && self.terminal_error.is_none() {
1663                        self.terminal_error = Some(err.clone());
1664                    }
1665                    return Err(err);
1666                }
1667                Ok(Some(BatchView {
1668                    decoded: self.last_batch.as_ref().unwrap(),
1669                    dict: &self.reader.dict,
1670                    schema: self.reader.query_schema.as_ref().unwrap(),
1671                }))
1672            }
1673            Ok(NextOutcome::Done) => Ok(None),
1674            Err(e) => {
1675                if self.done && self.terminal_error.is_none() {
1676                    self.terminal_error = Some(e.clone());
1677                }
1678                Err(e)
1679            }
1680        }
1681    }
1682
1683    /// Wrap this cursor as an Arrow [`RecordBatchReader`]. Blocks until
1684    /// the first `RESULT_BATCH` is decoded, then snapshots its schema.
1685    /// Mid-stream schema drift poisons the adapter; re-wrap to resume.
1686    /// Returns [`ErrorCode::NoSchema`] if the stream terminates before
1687    /// any batch is produced.
1688    ///
1689    /// [`RecordBatchReader`]: arrow::array::RecordBatchReader
1690    /// [`ErrorCode::NoSchema`]: crate::ErrorCode::NoSchema
1691    #[cfg(feature = "arrow-egress")]
1692    pub fn as_arrow_reader<'c>(
1693        &'c mut self,
1694    ) -> Result<crate::egress::arrow::CursorRecordBatchReader<'r, 'c>> {
1695        crate::egress::arrow::CursorRecordBatchReader::new(self)
1696    }
1697
1698    /// Eagerly drain every batch and return them together with the
1699    /// pinned Arrow schema. Symmetric with
1700    /// [`Cursor::fetch_all_polars`](crate::egress::Cursor::fetch_all_polars).
1701    /// Errors as [`ErrorCode::NoSchema`] if the stream ends without
1702    /// producing a batch; surfaces drift as
1703    /// [`ErrorCode::SchemaDrift`].
1704    ///
1705    /// [`ErrorCode::NoSchema`]: crate::ErrorCode::NoSchema
1706    /// [`ErrorCode::SchemaDrift`]: crate::ErrorCode::SchemaDrift
1707    #[cfg(feature = "arrow-egress")]
1708    pub fn fetch_all_arrow(
1709        &mut self,
1710    ) -> Result<(arrow::datatypes::SchemaRef, Vec<arrow::array::RecordBatch>)> {
1711        // Materialise-whole: nothing leaves the library until the full
1712        // result is built, so a mid-query failover can re-read it
1713        // transparently. Opt into replay and discard the partial
1714        // accumulation when the cursor reports a reset.
1715        self.enable_internal_replay();
1716        let mut reader = self.as_arrow_reader()?;
1717        let mut resets_seen = reader.failover_resets();
1718        let mut batches: Vec<arrow::array::RecordBatch> = Vec::new();
1719        loop {
1720            // Manual drive (not `for`/`by_ref`) so the reset counter can be
1721            // polled between batches without holding an iterator borrow.
1722            let Some(item) = reader.next() else { break };
1723            let rb = item.map_err(|e| {
1724                crate::egress::arrow::try_downcast_questdb(&e)
1725                    .cloned()
1726                    .unwrap_or_else(|| fmt!(ArrowExport, "{}", e))
1727            })?;
1728            let resets_now = reader.failover_resets();
1729            if resets_now != resets_seen {
1730                resets_seen = resets_now;
1731                batches.clear();
1732            }
1733            batches.push(rb);
1734        }
1735        Ok((reader.schema(), batches))
1736    }
1737
1738    /// Drift-checked iterator over Polars [`DataFrame`](polars::frame::DataFrame)s,
1739    /// one per QWP batch. Snapshots the first batch's Arrow schema
1740    /// and yields `Err(SchemaDrift)` then terminates if a
1741    /// later batch diverges. Returns `Err(NoSchema)` if the stream
1742    /// ends before any batch is produced.
1743    ///
1744    /// Use this in preference to a `while let Some(df) = cursor.next_polars()?`
1745    /// loop when you care about schema consistency mid-stream.
1746    #[cfg(feature = "polars-egress")]
1747    pub fn iter_polars<'c>(&'c mut self) -> Result<crate::egress::arrow::CursorPolarsIter<'r, 'c>> {
1748        crate::egress::arrow::CursorPolarsIter::new(self)
1749    }
1750
1751    /// Next batch as an Arrow [`RecordBatch`](arrow::array::RecordBatch).
1752    /// `Ok(None)` on stream end; replays terminal errors like
1753    /// [`Cursor::next_batch`]. No drift check — use
1754    /// [`Cursor::as_arrow_reader`] for that.
1755    #[cfg(feature = "arrow-egress")]
1756    pub fn next_arrow_batch(&mut self) -> Result<Option<arrow::array::RecordBatch>> {
1757        self.next_arrow_batch_inner(None, false)
1758    }
1759
1760    #[cfg(feature = "arrow-egress")]
1761    #[doc(hidden)]
1762    pub fn next_arrow_batch_inner(
1763        &mut self,
1764        expected_schema: Option<&arrow::datatypes::SchemaRef>,
1765        compact: bool,
1766    ) -> Result<Option<arrow::array::RecordBatch>> {
1767        use crate::egress::arrow::{batch_arrow_schema, batch_to_record_batch_with, schemas_equal};
1768        use std::sync::Arc;
1769
1770        if self.done {
1771            return match self.terminal_error.as_ref() {
1772                Some(e) => Err(e.clone()),
1773                None => Ok(None),
1774            };
1775        }
1776        // Replay a batch that drifted on a previous call before reading a new
1777        // frame; its transport side effects already ran, so skip
1778        // `next_batch_inner`.
1779        let decoded = if let Some(stashed) = self.drifted_batch.take() {
1780            stashed
1781        } else {
1782            let outcome = match self.next_batch_inner() {
1783                Ok(o) => o,
1784                Err(e) => {
1785                    if self.done && self.terminal_error.is_none() {
1786                        self.terminal_error = Some(e.clone());
1787                    }
1788                    return Err(e);
1789                }
1790            };
1791            match outcome {
1792                NextOutcome::Done => return Ok(None),
1793                // `next_batch_inner` populates `last_batch` before returning
1794                // `HaveBatch`; re-check softly rather than `.expect()`, since a
1795                // panic would abort the whole process across the FFI boundary
1796                // (`panic=abort`) if a future refactor broke the invariant.
1797                NextOutcome::HaveBatch => match self.last_batch.take() {
1798                    Some(b) => b,
1799                    None => {
1800                        let e = fmt!(
1801                            ProtocolError,
1802                            "internal invariant: next_batch produced a batch without a decoded view"
1803                        );
1804                        self.stash_arrow_terminal_error(&e);
1805                        return Err(e);
1806                    }
1807                },
1808            }
1809        };
1810        let egress_schema = match self.reader.query_schema.as_ref() {
1811            Some(s) => s.clone(),
1812            None => {
1813                let e = fmt!(
1814                    ProtocolError,
1815                    "internal invariant: next_batch produced a batch without a decoded schema"
1816                );
1817                self.stash_arrow_terminal_error(&e);
1818                return Err(e);
1819            }
1820        };
1821        let arrow_schema = match batch_arrow_schema(&egress_schema, &decoded) {
1822            Ok(s) => Arc::new(s),
1823            Err(e) => {
1824                self.stash_arrow_terminal_error(&e);
1825                return Err(e);
1826            }
1827        };
1828        if let Some(expected) = expected_schema
1829            && !schemas_equal(expected.as_ref(), arrow_schema.as_ref())
1830        {
1831            let e = fmt!(
1832                SchemaDrift,
1833                "mid-stream Arrow schema drift: expected schema differs from batch_seq={}",
1834                decoded.batch_seq
1835            );
1836            // Keep the batch so its rows stay retrievable via
1837            // `Cursor::next_arrow_batch` rather than dropped.
1838            self.drifted_batch = Some(decoded);
1839            return Err(e);
1840        }
1841        #[cfg(feature = "polars-egress")]
1842        {
1843            self.symbol_delta_modes.clear();
1844            self.symbol_delta_modes
1845                .extend(decoded.columns.iter().map(|c| {
1846                    matches!(
1847                        c,
1848                        crate::egress::decoder::DecodedColumn::Symbol {
1849                            local_dict: None,
1850                            ..
1851                        }
1852                    )
1853                }));
1854        }
1855        match batch_to_record_batch_with(
1856            arrow_schema,
1857            &egress_schema,
1858            decoded,
1859            &self.reader.dict,
1860            &mut self.sym_values,
1861            if compact {
1862                Some(&mut self.sym_scratch)
1863            } else {
1864                None
1865            },
1866        ) {
1867            Ok(rb) => Ok(Some(rb)),
1868            Err(e) => {
1869                self.stash_arrow_terminal_error(&e);
1870                Err(e)
1871            }
1872        }
1873    }
1874
1875    #[cfg(feature = "polars-egress")]
1876    pub(crate) fn symbol_registry_synced(
1877        &mut self,
1878    ) -> Result<&crate::egress::arrow::polars::SymbolRegistry> {
1879        let reg = self
1880            .symbol_registry
1881            .get_or_insert_with(crate::egress::arrow::polars::SymbolRegistry::new);
1882        reg.sync(&self.reader.dict)?;
1883        Ok(reg)
1884    }
1885
1886    #[cfg(feature = "polars-egress")]
1887    pub(crate) fn symbol_delta_modes(&self) -> &[bool] {
1888        &self.symbol_delta_modes
1889    }
1890
1891    // Replay-contract stash for fatal errors that bypass `next_batch_inner`
1892    // (missing/invalid Arrow schema, `batch_to_record_batch`): marks the
1893    // cursor terminal so the error replays on every later call instead of
1894    // silently advancing. Schema drift is NOT terminal — it leaves the cursor
1895    // live and parks the drifted batch in `drifted_batch` so the caller can
1896    // re-snapshot and retrieve those rows on the next call (see
1897    // `next_arrow_batch_inner`).
1898    #[cfg(feature = "arrow-egress")]
1899    fn stash_arrow_terminal_error(&mut self, err: &Error) {
1900        self.done = true;
1901        if self.terminal_error.is_none() {
1902            self.terminal_error = Some(err.clone());
1903        }
1904    }
1905
1906    fn next_batch_inner(&mut self) -> Result<NextOutcome> {
1907        loop {
1908            // Transport read: a failure here (socket closed, TLS
1909            // reset, truncated WS frame) is what failover is for.
1910            let (header, payload) = match self.read_frame_raw() {
1911                Ok(hp) => hp,
1912                Err(e) => {
1913                    self.failover_after_stream_failure(e, StreamFailureKind::Read)?;
1914                    continue;
1915                }
1916            };
1917            // Capture wire size BEFORE the decode consumes the header.
1918            let wire_bytes = HEADER_LEN as u64 + header.payload_length as u64;
1919            // Decode failures can be the symptom of a dying endpoint that
1920            // managed to emit one complete-but-corrupt WS frame. Route
1921            // failover-eligible errors through the same replay machinery as
1922            // raw read failures; deterministic codes such as
1923            // UnsupportedServer remain terminal via `is_failover_eligible`.
1924            // Unlike a raw read failure, a decode failure gets only a small
1925            // dedicated replay cap (`MAX_DECODE_FAILOVER_ROUNDS`) so a
1926            // deterministically-corrupting server can't drive a
1927            // reconnect/replay loop that drains the whole per-Execute budget
1928            // — see `failover_after_stream_failure`.
1929            let t1 = std::time::Instant::now();
1930            let decode_result = decode_frame(
1931                header,
1932                &payload,
1933                &mut self.reader.dict,
1934                &mut self.reader.query_schema,
1935                &mut self.reader.zstd_scratch,
1936            );
1937            // Account for decode time on both arms — the error path is
1938            // rare and terminal, but skipping the sample makes the
1939            // metric subtly biased toward "successful decodes are slow."
1940            self.reader.stats.decode_ns.fetch_add(
1941                u64::try_from(t1.elapsed().as_nanos()).unwrap_or(u64::MAX),
1942                Ordering::Relaxed,
1943            );
1944            let event = match decode_result {
1945                Ok(ev) => ev,
1946                Err(e) => {
1947                    self.failover_after_stream_failure(e, StreamFailureKind::Decode)?;
1948                    continue;
1949                }
1950            };
1951            match event {
1952                ServerEvent::Batch(b) => {
1953                    if b.request_id != self.request_id {
1954                        let err = fmt!(
1955                            ProtocolError,
1956                            "RESULT_BATCH request_id {} != cursor {}",
1957                            b.request_id,
1958                            self.request_id
1959                        );
1960                        // Stale-rid frames mean the server is still
1961                        // streaming for an old request — keep reading
1962                        // would only deepen the corruption.
1963                        self.terminate_with_close();
1964                        return Err(err);
1965                    }
1966                    // Replenish the server's per-request byte budget for
1967                    // the bytes we just took off the wire. The wire bytes
1968                    // are no longer pinned in our buffer; sending CREDIT
1969                    // here matches the server's "release on drain" policy.
1970                    //
1971                    // Suppress replenishment once `cancel()` has started
1972                    // draining: topping the server's budget back up while
1973                    // we're throwing the bytes away defeats the very
1974                    // backpressure that should be hastening cancellation.
1975                    if self.credit_enabled
1976                        && !self.cancelling
1977                        && let Err(e) = self.send_credit_frame(wire_bytes)
1978                    {
1979                        // A failed credit write means the transport
1980                        // just died. Surface it as a hard cursor
1981                        // failure rather than leaving the cursor
1982                        // "active" (which would let the next
1983                        // `next_batch` call silently failover and
1984                        // mask the credit-write error from the user).
1985                        self.terminate_with_close();
1986                        return Err(e);
1987                    }
1988                    // decode_result_batch guarantees `query_schema` is
1989                    // populated on Ok (batch_seq == 0 sets it; > 0 errors
1990                    // when it's absent). Defensive check rather than an
1991                    // `.expect()` so an internal-invariant violation can't
1992                    // abort the process across the FFI boundary.
1993                    if self.reader.query_schema.is_none() {
1994                        let err = fmt!(ProtocolError, "RESULT_BATCH decoded without a schema");
1995                        self.terminate_with_close();
1996                        return Err(err);
1997                    }
1998                    let last = self.last_batch.insert(b);
1999                    // Latch sticky `data_delivered` BEFORE yielding the
2000                    // batch view — a subsequent failover-eligible read
2001                    // error must see the latch already set, since by
2002                    // that point the caller has consumed at least one
2003                    // row from this query.
2004                    self.data_delivered = true;
2005                    // BatchView construction is hoisted to `next_batch`
2006                    // (the wrapper) so the inner returns a borrow-free
2007                    // discriminant; the wrapper re-acquires the borrows
2008                    // on `last_batch`, `dict`, and `query_schema` itself.
2009                    // `last` is still in scope here only for the side
2010                    // effects (insert + data_delivered).
2011                    let _ = last;
2012                    return Ok(NextOutcome::HaveBatch);
2013                }
2014                ServerEvent::End {
2015                    request_id,
2016                    final_seq,
2017                    total_rows,
2018                } => {
2019                    if let Err(e) = self.check_rid(request_id, "RESULT_END") {
2020                        self.terminate_with_close();
2021                        return Err(e);
2022                    }
2023                    self.terminal = Some(Terminal::End {
2024                        final_seq,
2025                        total_rows,
2026                    });
2027                    self.reader.cursor_active = false;
2028                    self.done = true;
2029                    return Ok(NextOutcome::Done);
2030                }
2031                ServerEvent::ExecDone {
2032                    request_id,
2033                    op_type,
2034                    rows_affected,
2035                } => {
2036                    if let Err(e) = self.check_rid(request_id, "EXEC_DONE") {
2037                        self.terminate_with_close();
2038                        return Err(e);
2039                    }
2040                    self.terminal = Some(Terminal::ExecDone {
2041                        op_type,
2042                        rows_affected,
2043                    });
2044                    self.reader.cursor_active = false;
2045                    self.done = true;
2046                    return Ok(NextOutcome::Done);
2047                }
2048                ServerEvent::Error {
2049                    request_id,
2050                    status,
2051                    message,
2052                } => {
2053                    if let Err(e) = self.check_rid(request_id, "QUERY_ERROR") {
2054                        self.terminate_with_close();
2055                        return Err(e);
2056                    }
2057                    // Transparent recovery from the transient stale-cached-plan
2058                    // fault. An async `ALTER COLUMN TYPE` bumps the table's
2059                    // metadata version between this query's server-side
2060                    // compilation and its execution, so the server rejects its
2061                    // own cached plan with `INTERNAL_ERROR`. The recompile on
2062                    // the very next execution succeeds — this is exactly how
2063                    // QuestDB's PGWire / REST endpoints self-heal, and that
2064                    // friction must never leak to the caller (it is not
2065                    // something a user can act on, so surfacing it is pure
2066                    // noise).
2067                    //
2068                    // Replaying is safe only before any row was handed to the
2069                    // caller (`!data_delivered`): the fault is a compile-time
2070                    // error that fires before `batch_seq == 0`, so in practice
2071                    // the guard always holds — but it is load-bearing, because
2072                    // replaying after delivery would re-stream rows the caller
2073                    // already consumed. The connection is healthy (the
2074                    // `QUERY_ERROR` is terminal only for *this* request_id), so
2075                    // unlike failover we re-issue on the same connection with a
2076                    // fresh request_id instead of reconnecting. `cancelling`
2077                    // suppresses the retry so a concurrent `cancel()` wins.
2078                    if !self.cancelling
2079                        && !self.data_delivered
2080                        && self.stale_plan_retries < MAX_STALE_PLAN_RETRIES
2081                        && is_stale_plan_error(status, &message)
2082                    {
2083                        self.stale_plan_retries = self.stale_plan_retries.saturating_add(1);
2084                        match self.replay_query_same_connection() {
2085                            Ok(()) => continue,
2086                            Err(e) => {
2087                                self.reader.cursor_active = false;
2088                                self.done = true;
2089                                return Err(e);
2090                            }
2091                        }
2092                    }
2093                    self.reader.cursor_active = false;
2094                    self.done = true;
2095                    return Err(map_server_status(status, message));
2096                }
2097                ServerEvent::CacheReset { .. } => {
2098                    // `decode_frame` already cleared the connection dict.
2099                    self.reset_symbol_caches();
2100                    continue;
2101                }
2102                ServerEvent::ServerInfo(_) => {
2103                    // State already mutated by decode_frame; keep reading.
2104                    continue;
2105                }
2106            }
2107        }
2108    }
2109
2110    /// Number of successful failover reconnects this cursor has
2111    /// observed since `execute()`. Useful for tests asserting the
2112    /// query did or did not silently restart.
2113    pub fn failover_resets(&self) -> u32 {
2114        self.failover_resets
2115    }
2116
2117    /// Number of times this cursor transparently re-issued its query on the
2118    /// current connection after the server reported the transient
2119    /// stale-cached-plan `INTERNAL_ERROR` (see [`Cursor::next_batch`]).
2120    /// Stays `0` on the happy path; exposed for tests and diagnostics that
2121    /// want to confirm the self-heal fired (and how often) without the
2122    /// caller ever seeing the underlying error.
2123    pub fn stale_plan_retries(&self) -> u32 {
2124        self.stale_plan_retries
2125    }
2126
2127    /// Opt this cursor into transparent mid-query replay from the
2128    /// materialise-whole adapters (`fetch_all_polars`, `fetch_all_arrow`).
2129    /// Those adapters hold the entire result internally and discard their
2130    /// accumulator on a reset (tracked via [`Cursor::failover_resets`]), so
2131    /// replay-from-`batch_seq 0` re-reads the whole result exactly once —
2132    /// nothing has left the library. Installing a no-op reset callback is
2133    /// what clears the silent-duplicate guard in `next_batch_inner` (see
2134    /// [`would_silently_duplicate`]); the streaming entry points
2135    /// (`iter_polars`, `next_polars`, `next_arrow_batch`) deliberately do
2136    /// **not** call this, so a batch already yielded to the caller still
2137    /// surfaces [`ErrorCode::FailoverWouldDuplicate`].
2138    ///
2139    /// Leaves a user-installed callback in place: if the caller already
2140    /// opted into replays, that contract wins.
2141    #[cfg(feature = "arrow-egress")]
2142    pub(crate) fn enable_internal_replay(&mut self) {
2143        if self.on_failover_reset.is_none() {
2144            self.on_failover_reset = Some(Box::new(|_: &FailoverResetEvent| {}));
2145        }
2146    }
2147
2148    /// The endpoint the cursor's underlying connection is currently
2149    /// bound to. While the cursor is live the `Reader` is mutably
2150    /// borrowed, so [`Reader::current_addr`] is unreachable from
2151    /// user code — this is the in-cursor accessor for "which
2152    /// endpoint did the last batch come from?". After mid-query
2153    /// failover, this reflects the new endpoint (matching the
2154    /// `new_addr` from the most recent
2155    /// [`crate::egress::FailoverResetEvent`]).
2156    pub fn current_addr(&self) -> &Endpoint {
2157        self.reader.current_addr()
2158    }
2159
2160    /// Negotiated QWP version of the cursor's underlying connection. The
2161    /// in-cursor accessor for [`Reader::server_version`], unreachable from
2162    /// user code while the cursor holds the `Reader`'s mutable borrow.
2163    /// Reflects the renegotiated version after mid-query failover.
2164    pub fn server_version(&self) -> Result<u8> {
2165        self.reader.server_version()
2166    }
2167
2168    /// `SERVER_INFO` of the cursor's currently connected endpoint;
2169    /// `None` only while a reconnect is in flight (the single QWP
2170    /// version always supplies it). The in-cursor accessor for
2171    /// [`Reader::server_info`], unreachable from user code while the
2172    /// cursor holds the `Reader`'s mutable borrow. Reflects the new
2173    /// endpoint after mid-query failover.
2174    pub fn server_info(&self) -> Option<&ServerInfo> {
2175        self.reader.server_info()
2176    }
2177
2178    /// Read one raw frame (header + payload) off the transport, with
2179    /// no decode. Errors here are transport-level (socket closed,
2180    /// truncated WS frame, TLS reset, etc.). Decoding is deliberately
2181    /// NOT done here — the caller decides whether decode failures are
2182    /// failover-eligible too.
2183    fn read_frame_raw(
2184        &mut self,
2185    ) -> Result<(crate::egress::wire::header::FrameHeader, bytes::Bytes)> {
2186        let t0 = std::time::Instant::now();
2187        let (header, payload) = self.reader.transport_mut()?.read_frame()?;
2188        self.reader.stats.read_ns.fetch_add(
2189            u64::try_from(t0.elapsed().as_nanos()).unwrap_or(u64::MAX),
2190            Ordering::Relaxed,
2191        );
2192        let wire_bytes = HEADER_LEN as u64 + header.payload_length as u64;
2193        self.reader
2194            .stats
2195            .bytes_received
2196            .fetch_add(wire_bytes, Ordering::Relaxed);
2197        Ok((header, payload))
2198    }
2199
2200    /// Shared failover gate for failures observed while consuming a query
2201    /// stream. This covers raw transport reads and failover-eligible decode
2202    /// errors so both surfaces obey the same cancellation, duplicate-delivery,
2203    /// callback, budget, and endpoint-tracker rules.
2204    fn failover_after_stream_failure(&mut self, e: Error, kind: StreamFailureKind) -> Result<()> {
2205        if self.cancelling || !self.reader.cfg.failover || !is_failover_eligible(e.code()) {
2206            // Match every other terminal path in this loop: tear down the
2207            // WS so the cursor's flags stay coherent with the transport
2208            // state, with no half-cooked cursors that defer cleanup to
2209            // `Reader::Drop`.
2210            self.terminate_with_close();
2211            return Err(e);
2212        }
2213        // Silent-duplicate guard. If at least one batch was already yielded
2214        // to the caller and they didn't install a reset callback,
2215        // replay would deliver those rows again with no signal — see
2216        // `ErrorCode::FailoverWouldDuplicate`. The exact-once contract is
2217        // "rows surface to the caller at most once unless they explicitly
2218        // opt in to seeing replays."
2219        //
2220        // The trigger error `e` is preserved in the message so the caller
2221        // still learns *why* the cursor died; diagnostics shouldn't get
2222        // worse just because we re-classified the surface.
2223        if would_silently_duplicate(self.data_delivered, self.on_failover_reset.is_some()) {
2224            let err = fmt!(
2225                FailoverWouldDuplicate,
2226                "mid-query failover would replay rows already delivered to the caller \
2227                 (install on_failover_reset to authorize replay); \
2228                 cursor terminated. Trigger: {} ({:?})",
2229                e.msg(),
2230                e.code()
2231            );
2232            self.terminate_with_close();
2233            return Err(err);
2234        }
2235        // Decode-driven replays get a small dedicated cap. A decode
2236        // failure can be transient wire corruption (one reconnect to a
2237        // fresh connection cures it) or a deterministic protocol
2238        // violation (every replay reproduces it byte-for-byte). The two
2239        // are indistinguishable at the `ErrorCode` level, so we allow a
2240        // bounded number of decode-triggered replays to recover the
2241        // transient case, then surface the decode error terminally
2242        // rather than draining the full per-Execute failover budget
2243        // (and emitting one warning per round) against a server that
2244        // will just re-corrupt the replayed query forever. Raw transport
2245        // read failures are unaffected and keep the full budget.
2246        if kind == StreamFailureKind::Decode {
2247            if self.decode_failover_rounds >= MAX_DECODE_FAILOVER_ROUNDS {
2248                self.terminate_with_close();
2249                return Err(e);
2250            }
2251            self.decode_failover_rounds = self.decode_failover_rounds.saturating_add(1);
2252        }
2253        warn_on_protocol_error_failover(&e, kind.context());
2254        self.failover_reconnect_and_replay(e)
2255    }
2256
2257    /// Re-issue the stashed `QUERY_REQUEST` on the *current* connection with
2258    /// a fresh `request_id`. Used to transparently recover from the
2259    /// transient stale-cached-plan `INTERNAL_ERROR`: the connection is
2260    /// healthy (the `QUERY_ERROR` was terminal only for the old
2261    /// request_id), so unlike [`Cursor::failover_reconnect_and_replay`] this
2262    /// does NOT reconnect. It patches the 8-byte request_id span in place,
2263    /// clears the per-query schema (mirroring [`ReaderQuery::execute`] so a
2264    /// stale schema can't bind the replayed rows), drops any half-built
2265    /// batch view, and resends the same bytes verbatim — no builder/bind
2266    /// clone, no re-encode. The server recompiles against the table's
2267    /// current metadata and streams afresh from `batch_seq == 0`.
2268    ///
2269    /// The connection-scoped symbol dict is deliberately *not* reset: this
2270    /// is a sequential query on the same connection (just like a second
2271    /// `execute()`), so the dict — and the per-cursor caches keyed on it —
2272    /// remain valid. `cursor_active` stays `true`; the cursor is still live.
2273    fn replay_query_same_connection(&mut self) -> Result<()> {
2274        let new_rid = self.reader.alloc_request_id();
2275        self.request_id = new_rid;
2276        self.encoded_request = patch_request_id(std::mem::take(&mut self.encoded_request), new_rid);
2277        // Mirror execute(): the schema rides batch_seq==0 of the new query.
2278        self.reader.query_schema = None;
2279        self.last_batch = None;
2280        // Any parked drift-replay batch belonged to the rejected attempt.
2281        #[cfg(feature = "arrow-egress")]
2282        {
2283            self.drifted_batch = None;
2284        }
2285        self.reader
2286            .transport_mut()
2287            .and_then(|t| t.write_message(self.encoded_request.clone()))
2288    }
2289
2290    /// Drop the per-cursor SYMBOL caches keyed on the connection dict.
2291    /// Must be called whenever `self.dict` is replaced, otherwise a
2292    /// re-grown dict can alias stale interned values/codes.
2293    fn reset_symbol_caches(&mut self) {
2294        #[cfg(feature = "arrow-egress")]
2295        {
2296            self.sym_values = crate::egress::arrow::SymbolValuesCache::default();
2297            self.sym_scratch = crate::egress::arrow::SymbolBuildScratch::default();
2298        }
2299        #[cfg(feature = "polars-egress")]
2300        {
2301            self.symbol_registry = None;
2302        }
2303    }
2304
2305    /// Mid-query failover: the underlying connection just died with
2306    /// `trigger`. Walk the address list (skipping the failed endpoint
2307    /// first), with exponential backoff, until a fresh connection is
2308    /// established; then reset the cursor for replay (new
2309    /// `request_id`, cleared `last_batch`), re-encode the original
2310    /// `QUERY_REQUEST`, and notify the user-side handler so it can
2311    /// discard accumulated rows. On exhausted budget or hard error,
2312    /// the cursor is marked terminal and the failure is propagated.
2313    fn failover_reconnect_and_replay(&mut self, trigger: Error) -> Result<()> {
2314        let mut trigger = trigger;
2315        loop {
2316            let started = std::time::Instant::now();
2317            let failed_idx = self.reader.addr_idx;
2318            // Snapshot the failing endpoint before reconnect mutates
2319            // `addr_idx` — `FailoverResetEvent` reports it back to the user.
2320            let failed_addr = self.reader.cfg.addrs[failed_idx].clone();
2321
2322            // Phase: Disconnected. Fires before the retry loop runs so an
2323            // SLO dashboard sees the outage *now*, not retroactively when
2324            // a reconnect lands or the budget exhausts.
2325            if let Some(cb) = self.on_failover_progress.as_mut() {
2326                let event = FailoverProgressEvent {
2327                    phase: FailoverPhase::Disconnected,
2328                    failed_addr: failed_addr.clone(),
2329                    new_addr: None,
2330                    new_server_info: None,
2331                    new_request_id: None,
2332                    attempt: 0,
2333                    trigger: trigger.clone(),
2334                    elapsed: started.elapsed(),
2335                    final_error: None,
2336                };
2337                cb(&event);
2338            }
2339
2340            // Phase: Retrying. The closure fires once per outer-loop
2341            // iteration of `reconnect_with_failover`. We split the borrow
2342            // on `self` so the closure can mutate the progress callback
2343            // while `reader.reconnect_with_failover` holds a `&mut Reader`.
2344            // `last_attempt` is tracked outside the closure so the GaveUp
2345            // event can report the final attempt count even when the
2346            // reconnect loop breaks out via the wall-clock-deadline path
2347            // (which doesn't surface the count in its `Err`).
2348            let mut last_attempt: u32 = 0;
2349            let reconnect_result = {
2350                let Self {
2351                    reader,
2352                    on_failover_progress,
2353                    failover_budget,
2354                    ..
2355                } = self;
2356                let failed_addr_ref = &failed_addr;
2357                let trigger_ref = &trigger;
2358                reader.reconnect_with_failover(failed_idx, failover_budget, &mut |attempt: u32| {
2359                    last_attempt = attempt;
2360                    if let Some(cb) = on_failover_progress.as_mut() {
2361                        let event = FailoverProgressEvent {
2362                            phase: FailoverPhase::Retrying,
2363                            failed_addr: failed_addr_ref.clone(),
2364                            new_addr: None,
2365                            new_server_info: None,
2366                            new_request_id: None,
2367                            attempt,
2368                            trigger: trigger_ref.clone(),
2369                            elapsed: started.elapsed(),
2370                            final_error: None,
2371                        };
2372                        cb(&event);
2373                    }
2374                })
2375            };
2376            let attempts = match reconnect_result {
2377                Ok(n) => n,
2378                Err(e) => {
2379                    // Phase: GaveUp. Fire before mutating state / returning
2380                    // so the callback sees the cursor in its
2381                    // about-to-be-terminal form and can correlate against
2382                    // the error the caller is about to receive via
2383                    // `next_batch`.
2384                    if let Some(cb) = self.on_failover_progress.as_mut() {
2385                        let event = FailoverProgressEvent {
2386                            phase: FailoverPhase::GaveUp,
2387                            failed_addr: failed_addr.clone(),
2388                            new_addr: None,
2389                            new_server_info: None,
2390                            new_request_id: None,
2391                            attempt: last_attempt,
2392                            trigger: trigger.clone(),
2393                            elapsed: started.elapsed(),
2394                            final_error: Some(e.clone()),
2395                        };
2396                        cb(&event);
2397                    }
2398                    self.reader.cursor_active = false;
2399                    self.done = true;
2400                    // Surface the most diagnostic error. The original
2401                    // `trigger` is almost always a generic transport
2402                    // failure (socket close, decode error). Anything
2403                    // specific the reconnect saw — auth rejected, role
2404                    // mismatched on every endpoint, config-level issue —
2405                    // tells the user *what to fix* and should win over
2406                    // the original cause-of-death.
2407                    return Err(if prefer_over_trigger(e.code()) {
2408                        e
2409                    } else {
2410                        trigger
2411                    });
2412                }
2413            };
2414            // Reset connection-scoped state. The new connection has its
2415            // own (empty) dict and per-query schema already (set up by
2416            // `connect_endpoint`). Drop any in-flight batch buffer so we
2417            // don't accidentally surface a stale view.
2418            self.last_batch = None;
2419            // The parked drift-replay batch belongs to the old stream.
2420            #[cfg(feature = "arrow-egress")]
2421            {
2422                self.drifted_batch = None;
2423            }
2424            // The new connection installed a fresh empty dict; the SYMBOL
2425            // caches keyed on the old one would otherwise alias stale values.
2426            self.reset_symbol_caches();
2427            // Allocate a fresh request_id and re-issue the same
2428            // QUERY_REQUEST bytes. The cursor stashed the encoded
2429            // payload at `execute()` time; here we patch the 8-byte
2430            // request_id span in place and write the buffer
2431            // verbatim. No builder clone, no Bind clone, no
2432            // re-encode — and crucially no memcpy of the body
2433            // either: the previous `write_message` call has dropped
2434            // its `Bytes` clone, so this clone is uniquely owned and
2435            // `try_into_mut` recovers the underlying `BytesMut`
2436            // zero-copy. With `failover_max_attempts` up to `1024`
2437            // and queries that may carry multi-MB `Bind::Binary`
2438            // payloads, this is the difference between a few bytes
2439            // and gigabytes of churn per failure event.
2440            let new_rid = self.reader.alloc_request_id();
2441            self.request_id = new_rid;
2442            self.encoded_request =
2443                patch_request_id(std::mem::take(&mut self.encoded_request), new_rid);
2444            match self
2445                .reader
2446                .transport_mut()
2447                .and_then(|t| t.write_message(self.encoded_request.clone()))
2448            {
2449                Ok(()) => {
2450                    self.failover_resets = self.failover_resets.saturating_add(1);
2451                    let new_addr = self.reader.cfg.addrs[self.reader.addr_idx].clone();
2452                    let new_server_info = self.reader.server_info.clone();
2453                    // Report the successful reconnect to telemetry first, then
2454                    // invoke the reset hook that lets the caller discard its
2455                    // partial result before any replayed batch is delivered.
2456                    if let Some(cb) = self.on_failover_progress.as_mut() {
2457                        let event = FailoverProgressEvent {
2458                            phase: FailoverPhase::Reset,
2459                            failed_addr: failed_addr.clone(),
2460                            new_addr: Some(new_addr.clone()),
2461                            new_server_info: new_server_info.clone(),
2462                            new_request_id: Some(new_rid),
2463                            attempt: attempts,
2464                            trigger: trigger.clone(),
2465                            elapsed: started.elapsed(),
2466                            final_error: None,
2467                        };
2468                        cb(&event);
2469                    }
2470                    if let Some(cb) = self.on_failover_reset.as_mut() {
2471                        let event = FailoverResetEvent {
2472                            failed_addr,
2473                            new_addr,
2474                            new_server_info,
2475                            new_request_id: new_rid,
2476                            attempts,
2477                            trigger,
2478                            elapsed: started.elapsed(),
2479                        };
2480                        cb(&event);
2481                    }
2482                    return Ok(());
2483                }
2484                Err(e) => {
2485                    // The freshly reconnected socket died while sending the
2486                    // replayed QUERY_REQUEST. That failed replay is the next
2487                    // Execute attempt in Java's model, so it has already spent
2488                    // the reconnect round that got us here. If the same
2489                    // cursor-owned budget still has room, feed the write error
2490                    // back through the same reconnect loop; do not invent a
2491                    // separate write-retry schedule.
2492                    warn_on_protocol_error_failover(&e, "replay query write");
2493                    if !self.reader.cfg.failover || !is_failover_eligible(e.code()) {
2494                        if let Some(cb) = self.on_failover_progress.as_mut() {
2495                            let event = FailoverProgressEvent {
2496                                phase: FailoverPhase::GaveUp,
2497                                failed_addr: failed_addr.clone(),
2498                                new_addr: None,
2499                                new_server_info: None,
2500                                new_request_id: None,
2501                                attempt: attempts,
2502                                trigger: trigger.clone(),
2503                                elapsed: started.elapsed(),
2504                                final_error: Some(e.clone()),
2505                            };
2506                            cb(&event);
2507                        }
2508                        if let Some(dead) = self.reader.transport.take() {
2509                            drop(dead);
2510                        }
2511                        self.reader.cursor_active = false;
2512                        self.done = true;
2513                        return Err(e);
2514                    }
2515                    trigger = e;
2516                    continue;
2517                }
2518            }
2519        }
2520    }
2521
2522    /// Send a CANCEL frame and drain until the server emits a terminal
2523    /// frame for this request.
2524    ///
2525    /// Blocking, but bounded. The CANCEL write inherits the transport's
2526    /// `WRITE_TIMEOUT`; immediately after the CANCEL is accepted by
2527    /// the kernel send buffer, the read timeout is tightened to
2528    /// `CANCEL_DRAIN_READ_TIMEOUT` and the write timeout to
2529    /// `CLOSE_TIMEOUT` for the duration of the credit-nudge + drain.
2530    /// That bounds the worst-case latency at one `WRITE_TIMEOUT`
2531    /// (CANCEL) + `CLOSE_TIMEOUT` (nudge) + `CANCEL_DRAIN_READ_TIMEOUT`
2532    /// (drain) — installing the drain bounds before the nudge avoids
2533    /// a second `WRITE_TIMEOUT` window on a stuck TLS peer. If the
2534    /// CANCEL write itself fails, the transport is torn down before
2535    /// the error is returned so the cursor's flags and the underlying
2536    /// connection state are left coherent.
2537    pub fn cancel(&mut self) -> Result<()> {
2538        if self.done {
2539            return Ok(());
2540        }
2541        // Record the user's intent to cancel BEFORE attempting any
2542        // network write. If the CANCEL write (or the credit-nudge
2543        // write) fails because the transport just died, a subsequent
2544        // `next_batch` MUST NOT failover-replay the query — the user
2545        // explicitly asked to cancel it. The failover guard in
2546        // `next_batch` is keyed on `self.cancelling`; setting it after
2547        // the writes leaves a window where a failed write returns
2548        // `Err` with `cancelling=false`, and the next `next_batch`
2549        // call would silently reconnect to another endpoint and run
2550        // the query the user just cancelled.
2551        //
2552        // Side benefit (which used to be the only purpose of setting
2553        // this flag): from this point on the cursor stops topping up
2554        // the server's credit window, so the remaining budget bleeds
2555        // off and the server stops generating new batches behind the
2556        // cancel.
2557        self.cancelling = true;
2558        let mut payload = Vec::with_capacity(9);
2559        payload.push(MsgKind::Cancel.as_u8());
2560        payload.extend_from_slice(&self.request_id.to_le_bytes());
2561
2562        // Capture the CANCEL write error explicitly: a `?` here would
2563        // leave `cancelling=true, done=false, transport=Some(broken)`,
2564        // and the half-broken transport would only be cleaned up when
2565        // `Reader::Drop` ran. Tearing it down here keeps the cursor's
2566        // flags and the transport in lockstep with the other terminal
2567        // paths in `next_batch`.
2568        let write_outcome = match self.reader.transport_mut() {
2569            Ok(t) => t.write_message(Bytes::from(payload)),
2570            Err(e) => Err(e),
2571        };
2572        if let Err(e) = write_outcome {
2573            self.terminate_with_close();
2574            return Err(e);
2575        }
2576        // Bound the drain reads AND the credit-nudge write before
2577        // anything else can block. tungstenite's `read()` is otherwise
2578        // a pure blocking syscall, and a stuck-but-not-RST'd TLS peer
2579        // whose kernel send buffer is still draining can absorb the
2580        // credit-nudge write for the full `WRITE_TIMEOUT` (60 s)
2581        // before the drain timeout would otherwise have a chance to
2582        // fire. Tightening to `CLOSE_TIMEOUT` here caps the worst-case
2583        // cancel() latency at `WRITE_TIMEOUT` (CANCEL) + `CLOSE_TIMEOUT`
2584        // (nudge) + `CANCEL_DRAIN_READ_TIMEOUT` (drain) instead of
2585        // 2 × `WRITE_TIMEOUT` + drain.
2586        if let Some(t) = self.reader.transport.as_mut() {
2587            t.set_read_timeout(Some(CANCEL_DRAIN_READ_TIMEOUT));
2588            t.set_write_timeout(Some(CLOSE_TIMEOUT));
2589        }
2590
2591        // Wake the server in case it's already credit-suspended. The
2592        // server's `handleCancel` only sets a flag; the cancel takes
2593        // effect when `streamResults` is next re-entered, which on a
2594        // credit-suspended stream happens only via `handleCredit`. A
2595        // 1-byte top-up is enough — `streamResults` checks the cancel
2596        // flag before the credit check, so the abort path fires
2597        // immediately and emits the terminal QUERY_ERROR. Without this
2598        // nudge a `cancel()` against a credit-suspended server would
2599        // deadlock.
2600        // Best-effort: the CANCEL frame has already been accepted by
2601        // the server, so reporting the credit-nudge failure as the
2602        // user-visible result of `cancel()` would mislead — the user
2603        // would see "cancel failed" while the cancellation is in
2604        // fact under way. If the nudge write fails (transport just
2605        // died) the drain loop below will pick up the same transport
2606        // failure and either route through failover or terminate the
2607        // cursor (depending on `cancelling`, which we already set).
2608        // If the nudge succeeds the drain proceeds normally. Either
2609        // way, swallowing the error here gives the user the truthful
2610        // signal: the cancellation request was delivered.
2611        if self.credit_enabled {
2612            // No-accounting variant: this 1-byte nudge exists only to
2613            // unstick a credit-suspended server so it can deliver the
2614            // QUERY_ERROR for our CANCEL. Bumping
2615            // `stats.credit_granted_total` here would violate the
2616            // counter's documented purpose ("cancel doesn't continue
2617            // topping up the server's budget"). See
2618            // `write_credit_frame_raw`.
2619            let _ = self.write_credit_frame_raw(1);
2620        }
2621
2622        // Drain until any terminal frame (RESULT_END / EXEC_DONE /
2623        // QUERY_ERROR including STATUS_CANCELLED) — swallow batches
2624        // between CANCEL and the server's acknowledgement. `done` is
2625        // the right guard here, not `terminal`: an error terminal
2626        // sets `done` but leaves `terminal` as `None`.
2627        let mut drain_result: Result<()> = Ok(());
2628        while !self.done {
2629            match self.next_batch() {
2630                Ok(Some(_)) => {} // discarded
2631                Ok(None) => break,
2632                Err(e) => {
2633                    if matches!(e.code(), crate::ErrorCode::Cancelled) {
2634                        break;
2635                    }
2636                    drain_result = Err(e);
2637                    break;
2638                }
2639            }
2640        }
2641
2642        // Restore timeouts if the connection survived.
2643        if let Some(t) = self.reader.transport.as_mut() {
2644            t.set_read_timeout(None);
2645            t.set_write_timeout(Some(WRITE_TIMEOUT));
2646        }
2647
2648        drain_result
2649    }
2650
2651    /// Manually grant the server `additional_bytes` of read budget on
2652    /// this cursor's request. Useful when the user wants a larger
2653    /// outstanding window than the per-batch auto-replenishment would
2654    /// give them, or when initial_credit was 0 but the user changes
2655    /// their mind mid-stream.
2656    ///
2657    /// Mirrors [`Self::next_batch`]'s failover policy: a transport-
2658    /// class write failure on the current connection triggers a
2659    /// reconnect-and-replay (when the connect string declares
2660    /// failover endpoints), after which the credit frame is re-sent
2661    /// on the new connection so the user's grant is preserved. If the
2662    /// reconnect fails or the failure is not failover-eligible
2663    /// (auth/config/protocol), the cursor is torn down so a follow-up
2664    /// `next_batch` sees a dead cursor instead of silently failing
2665    /// over.
2666    pub fn add_credit(&mut self, additional_bytes: u64) -> Result<()> {
2667        if self.done {
2668            return Err(match self.terminal_error.as_ref() {
2669                Some(e) => e.clone(),
2670                None => fmt!(InvalidApiCall, "cursor is terminal; add_credit not allowed"),
2671            });
2672        }
2673        let first_err = match self.send_credit_frame(additional_bytes) {
2674            Ok(()) => return Ok(()),
2675            Err(e) => e,
2676        };
2677        if self.cancelling || !self.reader.cfg.failover || !is_failover_eligible(first_err.code()) {
2678            self.terminate_with_close();
2679            return Err(first_err);
2680        }
2681        // Mirrors the silent-duplicate guard in `next_batch`. Once data
2682        // has been delivered to the caller without an
2683        // `on_failover_reset` callback, a reconnect-and-replay would
2684        // re-deliver those rows with no signal — violating the
2685        // exact-once contract. The trigger error is preserved in the
2686        // message so the caller still learns why the cursor died.
2687        if would_silently_duplicate(self.data_delivered, self.on_failover_reset.is_some()) {
2688            let err = fmt!(
2689                FailoverWouldDuplicate,
2690                "mid-query failover would replay rows already delivered to the caller \
2691                 (install on_failover_reset to authorize replay); \
2692                 cursor terminated. Trigger: {} ({:?})",
2693                first_err.msg(),
2694                first_err.code()
2695            );
2696            self.terminate_with_close();
2697            return Err(err);
2698        }
2699        warn_on_protocol_error_failover(&first_err, "add_credit write");
2700        self.failover_reconnect_and_replay(first_err)?;
2701        // Replay succeeded; the user's grant intent applies to the new
2702        // request now in flight. Re-send on the new connection. If
2703        // *that* fails too, treat it as a sticky terminal failure
2704        // rather than recursing — one failover per user call keeps the
2705        // latency bound predictable.
2706        match self.send_credit_frame(additional_bytes) {
2707            Ok(()) => Ok(()),
2708            Err(e) => {
2709                self.terminate_with_close();
2710                Err(e)
2711            }
2712        }
2713    }
2714
2715    fn send_credit_frame(&mut self, additional_bytes: u64) -> Result<()> {
2716        self.write_credit_frame_raw(additional_bytes)?;
2717        self.reader
2718            .stats
2719            .credit_granted_total
2720            .fetch_add(additional_bytes, Ordering::Relaxed);
2721        Ok(())
2722    }
2723
2724    /// Wire-only CREDIT emit, **without** bumping
2725    /// `stats.credit_granted_total`. Used by `cancel()`'s wake nudge so
2726    /// the counter's documented invariant — "`cancel()` doesn't
2727    /// continue topping up the server's budget" — holds exactly,
2728    /// without a "modulo the 1-byte cancel nudge" caveat. Every other
2729    /// CREDIT path goes through `send_credit_frame` and is accounted for.
2730    fn write_credit_frame_raw(&mut self, additional_bytes: u64) -> Result<()> {
2731        let mut payload = Vec::with_capacity(16);
2732        payload.push(MsgKind::Credit.as_u8());
2733        payload.extend_from_slice(&self.request_id.to_le_bytes());
2734        varint::encode_u64(additional_bytes, &mut payload);
2735        self.reader
2736            .transport_mut()?
2737            .write_message(Bytes::from(payload))?;
2738        Ok(())
2739    }
2740
2741    fn check_rid(&self, got: i64, what: &str) -> Result<()> {
2742        if got != self.request_id {
2743            return Err(fmt!(
2744                ProtocolError,
2745                "{} request_id {} != cursor {}",
2746                what,
2747                got,
2748                self.request_id
2749            ));
2750        }
2751        Ok(())
2752    }
2753
2754    /// Mark the cursor terminal and tear down the underlying WS
2755    /// transport. Used on every irrecoverable post-read error path in
2756    /// `next_batch` so the cursor's `cursor_active` / `done` flags
2757    /// and the transport are always left coherent — no half-cooked
2758    /// cursors that rely on `Drop` to clean up, and no stale frames
2759    /// left buffered for a follow-up `Reader::prepare()` to pick up.
2760    ///
2761    /// `take()` + explicit `drop` matches `reconnect_with_failover`'s
2762    /// pattern: `close_in_place` issues the WS Close frame but leaves
2763    /// the `WsTransport` (and its TCP `FD` + tungstenite read/write
2764    /// buffers) alive until the value is dropped. Leaving the dead
2765    /// transport in `self.reader.transport = Some(_)` would pin the
2766    /// FD and several MiB of buffers until the entire `Reader` is
2767    /// dropped — a bounded but real leak per terminated cursor.
2768    /// Taking ownership and dropping here releases both immediately.
2769    fn terminate_with_close(&mut self) {
2770        if let Some(mut t) = self.reader.transport.take() {
2771            t.close_in_place();
2772            drop(t);
2773        }
2774        self.reader.cursor_active = false;
2775        self.done = true;
2776    }
2777}
2778
2779impl Drop for Cursor<'_> {
2780    fn drop(&mut self) {
2781        // `cursor_active` is cleared by `next_batch()` on every terminal
2782        // path (RESULT_END, EXEC_DONE, QUERY_ERROR) and by `cancel()`
2783        // once it's drained. If it's still set at drop time, this cursor
2784        // was abandoned mid-stream: query frames are still en route on
2785        // the WS, and reusing the Reader for a new query would let the
2786        // next cursor pick them up and trip the request_id check.
2787        //
2788        // Send a best-effort CANCEL frame before tearing the WebSocket
2789        // down. Without this, the server keeps streaming `RESULT_BATCH`
2790        // frames for the abandoned request until it observes the WS
2791        // close — holding dictionary + schema + flow-control state for
2792        // a request the user no longer cares about. The CANCEL gets the
2793        // server to release that state immediately. `try_write_cancel`
2794        // tightens the write timeout so a stuck peer can't hold this
2795        // dropping thread for the full `WRITE_TIMEOUT`, and swallows
2796        // every error: Drop has nowhere to surface them.
2797        //
2798        // Defensive: while the cursor invariant says transport is
2799        // `Some` whenever `cursor_active` is true (the failover
2800        // paths clear `cursor_active` whenever they leave the
2801        // transport `None`), `Drop` should never panic.
2802        if self.reader.cursor_active {
2803            if let Some(mut t) = self.reader.transport.take() {
2804                if !self.cancelling {
2805                    t.try_write_cancel(self.request_id);
2806                }
2807                t.close_in_place();
2808                drop(t);
2809            }
2810            self.reader.cursor_active = false;
2811        }
2812    }
2813}
2814
2815/// Borrowed view over the most recently decoded batch.
2816#[must_use = "BatchView is a borrowed projection; dropping it without iterating \
2817              the rows or calling its accessors throws away the just-decoded batch"]
2818pub struct BatchView<'c> {
2819    decoded: &'c DecodedBatch,
2820    dict: &'c SymbolDict,
2821    schema: &'c Schema,
2822}
2823
2824impl<'c> BatchView<'c> {
2825    pub fn request_id(&self) -> i64 {
2826        self.decoded.request_id
2827    }
2828
2829    pub fn batch_seq(&self) -> u64 {
2830        self.decoded.batch_seq
2831    }
2832
2833    /// Per-batch wire flags from the frame header. Useful for asserting
2834    /// that compression / Gorilla paths were actually exercised.
2835    pub fn flags(&self) -> u8 {
2836        self.decoded.flags
2837    }
2838
2839    pub fn schema(&self) -> &'c Schema {
2840        self.schema
2841    }
2842
2843    pub fn row_count(&self) -> usize {
2844        self.decoded.row_count
2845    }
2846
2847    pub fn column_count(&self) -> usize {
2848        self.decoded.columns.len()
2849    }
2850
2851    /// Project a single column to a typed view.
2852    pub fn column(&self, idx: usize) -> Result<ColumnView<'_>> {
2853        self.decoded.column_view(idx, self.dict)
2854    }
2855
2856    /// Connection-scoped symbol dictionary backing every SYMBOL column
2857    /// in this batch; a `SymbolColumn`'s codes index into it.
2858    pub fn dict(&self) -> &'c SymbolDict {
2859        self.dict
2860    }
2861}
2862
2863/// Predicate for the failover trigger filter. Mirrors the Java
2864/// reference's "transport-level terminal failure" classification: any
2865/// failure that's plausibly fixable by reconnecting to a different
2866/// endpoint, but not failures that signal a hard problem (auth, bad
2867/// SQL, malformed binds, role-mismatch on a single-node config) which
2868/// would just bounce off every endpoint identically.
2869/// Predicate gating the silent-duplicate guard in
2870/// [`Cursor::next_batch`]: returns `true` when a mid-query failover
2871/// would silently re-deliver rows the caller has already consumed.
2872///
2873/// Replay restarts at `batch_seq=0` against the new endpoint, so the
2874/// caller's accumulator would see every previously-yielded row again.
2875/// The opt-in for "I will discard partial state on each replay" is installing
2876/// [`ReaderQuery::on_failover_reset`], which fires immediately before the first
2877/// replayed batch arrives on the new connection. The progress callback is
2878/// telemetry-only and does not authorize replay. Without a reset hook, the
2879/// only safe response is to terminate the cursor and let the caller re-execute
2880/// from scratch.
2881///
2882/// Extracted as a free function so the truth table is unit-testable
2883/// without needing a live transport.
2884fn would_silently_duplicate(data_delivered: bool, has_reset_callback: bool) -> bool {
2885    data_delivered && !has_reset_callback
2886}
2887
2888fn is_failover_eligible(code: ErrorCode) -> bool {
2889    matches!(
2890        code,
2891        ErrorCode::SocketError
2892            | ErrorCode::ConnectTimeout
2893            | ErrorCode::HandshakeError
2894            | ErrorCode::TlsError
2895            | ErrorCode::ProtocolError
2896            | ErrorCode::CouldNotResolveAddr
2897            // RoleMismatch is "soft" for failover purposes: we just
2898            // skip this endpoint and try the next one (counting against
2899            // the budget). The eventual surfaced error is RoleMismatch
2900            // if the budget exhausts entirely on mismatching nodes.
2901            | ErrorCode::RoleMismatch
2902    )
2903}
2904
2905/// `ProtocolError` is failover-eligible because it most often signals
2906/// transient wire-frame corruption (truncated WS frame, malformed
2907/// varint mid-stream) that a fresh connection will recover from. The
2908/// same code, however, also fires on deterministic protocol bugs
2909/// (unknown `MsgKind`, mismatched lengths) — and the silent-duplicate
2910/// guard in [`Cursor::next_batch`] only blocks replay when *no*
2911/// `on_failover_reset` callback is installed. With a callback set,
2912/// replay proceeds even for deterministic violations.
2913///
2914/// Emit a warning whenever a `ProtocolError` actually triggers failover
2915/// so operators can spot masked corruption. Routed through the `log`
2916/// facade (level `warn`) rather than unconditional `eprintln!`: with no
2917/// logger installed this is a no-op, so a deterministically-corrupting
2918/// server can't spam the process's stderr, while operators who want the
2919/// signal install a logger and filter by target/level. The
2920/// decode-replay cap (`MAX_DECODE_FAILOVER_ROUNDS`) independently bounds
2921/// how many times this can fire per Execute for decode-triggered
2922/// failover.
2923fn warn_on_protocol_error_failover(err: &Error, context: &str) {
2924    if err.code() == ErrorCode::ProtocolError {
2925        log::warn!(
2926            "ProtocolError triggered failover ({}): {} — \
2927             reconnecting may mask transient wire-frame corruption \
2928             (truncated frames, malformed varints) or a deterministic \
2929             protocol violation; check server logs if this recurs.",
2930            context,
2931            err.msg()
2932        );
2933    }
2934}
2935
2936/// Errors that carry more diagnostic value than a generic transport
2937/// `trigger` (the cause-of-death of the previous connection). When the
2938/// failover loop surfaces one of these, the user should see *that*,
2939/// not the original socket close — these tell the user *what to fix*
2940/// (credentials, cluster topology, server version, config, TLS / WS
2941/// handshake), whereas the trigger just says "the network broke at
2942/// some point."
2943///
2944/// `HandshakeError` and `TlsError` are preferred for the same reason
2945/// as `AuthError`: when every reachable endpoint rejects the WS
2946/// upgrade or fails certificate validation, the original
2947/// `SocketError` trigger ("connection dropped") is far less
2948/// actionable than the handshake/cert message that actually names
2949/// the problem.
2950fn prefer_over_trigger(code: ErrorCode) -> bool {
2951    matches!(
2952        code,
2953        ErrorCode::AuthError
2954            | ErrorCode::RoleMismatch
2955            | ErrorCode::ConfigError
2956            | ErrorCode::UnsupportedServer
2957            | ErrorCode::HandshakeError
2958            | ErrorCode::TlsError
2959    )
2960}
2961
2962/// Splitmix64 PRNG state for failover backoff jitter. Lives on the
2963/// `Reader`; each instance gets a distinct seed at construction time.
2964/// Splitmix64 is the simplest non-trivial 64-bit generator with good
2965/// statistical properties for this use case (uniform draws over small
2966/// integer ranges); avoids pulling `rand` into the `sync-reader-qwp-ws`
2967/// feature.
2968///
2969/// The state is mutated on every draw. Splitmix64 is full-period
2970/// (cycles through all 2^64 values), so deterministic seeding is fine
2971/// — the only requirement is that draws within a single reconnect
2972/// round are uncorrelated.
2973#[derive(Debug)]
2974pub(crate) struct FailoverRng {
2975    state: u64,
2976}
2977
2978impl FailoverRng {
2979    /// Seed from process time + a per-process monotonic counter so two
2980    /// Readers built in the same nanosecond still get distinct streams.
2981    pub(crate) fn new() -> Self {
2982        use std::sync::atomic::{AtomicU64, Ordering};
2983        static COUNTER: AtomicU64 = AtomicU64::new(0);
2984        let now_ns = std::time::SystemTime::now()
2985            .duration_since(std::time::UNIX_EPOCH)
2986            .map(|d| d.as_nanos() as u64)
2987            .unwrap_or(0);
2988        let bump = COUNTER.fetch_add(1, Ordering::Relaxed);
2989        // XOR-mix the two so neither's alone determines the seed —
2990        // SystemTime can be coarse on some platforms; the counter
2991        // alone would make collisions across processes likely.
2992        Self {
2993            state: now_ns ^ bump.wrapping_mul(0x9E37_79B9_7F4A_7C15),
2994        }
2995    }
2996
2997    /// Splitmix64 step. Returns a uniformly-distributed `u64`.
2998    fn next_u64(&mut self) -> u64 {
2999        self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
3000        let mut z = self.state;
3001        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
3002        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
3003        z ^ (z >> 31)
3004    }
3005
3006    /// Full-jitter draw per failover.md §3.1: `FullJitter(base) =
3007    /// uniform_long[0, base)`. Returns the random milliseconds to
3008    /// sleep before the next reconnect attempt. `base = 0` returns 0
3009    /// (sleeping for zero is a no-op).
3010    pub(crate) fn full_jitter_ms(&mut self, base: u64) -> u64 {
3011        if base == 0 {
3012            return 0;
3013        }
3014        // Modulo is safe: the bias for tiny `base` against the 2^64
3015        // value space is far below the resolution we care about for a
3016        // backoff jitter (sub-microsecond bias on a millisecond
3017        // schedule).
3018        self.next_u64() % base
3019    }
3020}
3021
3022/// Per the Java reference (`QwpQueryClient.matchesTarget`):
3023/// `STANDALONE` counts as `PRIMARY` so single-node OSS deployments work
3024/// with `target=primary`.
3025fn target_matches(target: Target, role: ServerRole) -> bool {
3026    match target {
3027        Target::Any => true,
3028        Target::Primary => matches!(
3029            role,
3030            ServerRole::Primary | ServerRole::PrimaryCatchup | ServerRole::Standalone
3031        ),
3032        Target::Replica => matches!(role, ServerRole::Replica),
3033    }
3034}
3035
3036/// Bound socket + decoded `SERVER_INFO` for one endpoint. Internal
3037/// intermediate produced by [`Reader::connect_endpoint`] and consumed
3038/// by [`walk_via_tracker`] / [`Reader::from_config`] /
3039/// [`Reader::reconnect_with_failover`].
3040struct TransportSession {
3041    idx: usize,
3042    transport: WsTransport,
3043    server_info: Option<ServerInfo>,
3044}
3045
3046/// Result of a successful tracker walk.
3047struct WalkOutcome {
3048    session: TransportSession,
3049    /// Number of `connect_endpoint` calls the walk made before
3050    /// landing on a successful endpoint. Includes failed picks before
3051    /// the success. The `FailoverResetEvent.attempts` field carries this
3052    /// value back to the user (cumulative across outer reconnect
3053    /// cycles).
3054    dials: u32,
3055}
3056
3057/// Walk the tracker until either an endpoint accepts or the round is
3058/// exhausted. Shared between [`Reader::from_config`] (initial connect)
3059/// and [`Reader::reconnect_with_failover`] (mid-query failover).
3060///
3061/// `allow_reset_pass`: when `true`, on exhaustion call
3062/// `tracker.begin_round(forget=true)` once and walk the list one more
3063/// time (failover.md §11.9.3). Initial connect passes `false` (the
3064/// tracker is fresh — every host already starts at `Unknown` and a
3065/// second pass would be a no-op anyway).
3066///
3067/// `terminal_codes`: error codes that abort the walk immediately
3068/// rather than being recorded into the tracker. Both callers pass
3069/// `[ConfigError, UnsupportedServer, AuthError]` — `AuthError` is
3070/// cluster-wide (credentials don't differ per host); the others are
3071/// build-level (client built without a feature the server requires)
3072/// or config-level (bad URL / unresolved name). Retrying every host
3073/// against any of these floods server logs without recovery, so the
3074/// walk bails on the first occurrence per spec §6 / §11.9.3.
3075fn walk_via_tracker(
3076    tracker: &mut HostHealthTracker,
3077    cfg: &Arc<ReaderConfig>,
3078    allow_reset_pass: bool,
3079    terminal_codes: &[ErrorCode],
3080) -> Result<WalkOutcome> {
3081    // Reset the within-round attempted bits. Topology classifications
3082    // accumulated by prior Executes are preserved (the within-outage
3083    // reset per failover.md §11.9.2). The fall-through pass below is
3084    // what re-evaluates stale classifications.
3085    tracker.begin_round(false);
3086    let mut last_role_mismatch: Option<Error> = None;
3087    let mut last_transport_err: Option<Error> = None;
3088    let mut retried_after_reset = false;
3089    let mut dials: u32 = 0;
3090    loop {
3091        let idx = match tracker.pick_next() {
3092            Some(i) => i,
3093            None => {
3094                if allow_reset_pass && !retried_after_reset {
3095                    // Failover.md §11.9.3 fall-through reset: give
3096                    // stale `TransientReject` / `TopologyReject` hosts
3097                    // from prior outages another shot before declaring
3098                    // the entire walk failed. Only one reset, then fail.
3099                    tracker.begin_round(true);
3100                    retried_after_reset = true;
3101                    continue;
3102                }
3103                break;
3104            }
3105        };
3106        dials = dials.saturating_add(1);
3107        match Reader::connect_endpoint(cfg.as_ref(), idx) {
3108            Ok(session) => {
3109                // Update zone tier from `SERVER_INFO.zone_id` when the
3110                // server advertised one (gated by `CAP_ZONE`). `record_zone`
3111                // with `None`/empty is a no-op, so passing the field
3112                // unconditionally is safe even when the server advertised
3113                // no zone (CAP_ZONE=0).
3114                if let Some(info) = session.server_info.as_ref() {
3115                    tracker.record_zone(idx, info.zone_id.as_deref());
3116                }
3117                tracker.record_success(idx);
3118                return Ok(WalkOutcome { session, dials });
3119            }
3120            Err(e) => {
3121                let code = e.code();
3122                if terminal_codes.contains(&code) {
3123                    // Hard error (config, unsupported server, auth).
3124                    // Bail out before recording into the tracker;
3125                    // there's no point preserving classifications when
3126                    // the walk is about to fail outright.
3127                    return Err(e);
3128                }
3129                match code {
3130                    ErrorCode::RoleMismatch => {
3131                        // Pull the role/zone bytes out of `UpgradeReject`
3132                        // (set by both the SERVER_INFO target-mismatch path
3133                        // and the `421 + X-QuestDB-Role` upgrade-reject path
3134                        // in transport.rs). A mismatch with no
3135                        // `UpgradeReject` (the no-SERVER_INFO guard)
3136                        // defaults to topological.
3137                        let reject = e.upgrade_reject();
3138                        let transient = reject.is_some_and(|r| r.is_transient());
3139                        if let Some(r) = reject {
3140                            tracker.record_zone(idx, r.zone.as_deref());
3141                        }
3142                        tracker.record_role_reject(idx, transient);
3143                        last_role_mismatch = Some(e);
3144                    }
3145                    _ => {
3146                        tracker.record_transport_error(idx);
3147                        last_transport_err = Some(e);
3148                    }
3149                }
3150            }
3151        }
3152    }
3153    // Walk exhausted (and reset pass, if any, exhausted too). Prefer
3154    // surfacing the last RoleMismatch (carries `UpgradeReject` with the
3155    // advertised role + zone, useful for diagnosing "no endpoint
3156    // matched target=") over a generic transport flop.
3157    if let Some(e) = last_role_mismatch {
3158        return Err(e);
3159    }
3160    Err(last_transport_err
3161        .unwrap_or_else(|| fmt!(SocketError, "all {} endpoints unreachable", cfg.addrs.len())))
3162}
3163
3164/// Read one frame off a fresh transport and expect `SERVER_INFO`.
3165/// Called once per successful upgrade. Uses throwaway dict / schema /
3166/// zstd scratch since `SERVER_INFO` itself
3167/// never carries symbols, schemas, or compressed payload — those state
3168/// machines only kick in once the Reader is assembled and starts
3169/// pulling `RESULT_BATCH` frames.
3170///
3171/// Bounded by `timeout` (sourced from
3172/// [`ReaderConfig::server_info_timeout_ms`], default 5 s per
3173/// failover.md §1.1). The `auth_timeout_ms` knob covers the HTTP
3174/// upgrade-response read only, and a server that accepts the upgrade
3175/// but then never sends the `SERVER_INFO` binary frame would
3176/// otherwise stall the connect indefinitely. The timeout is applied
3177/// as a TCP read deadline; on expiry the underlying read surfaces as
3178/// an `io::ErrorKind::WouldBlock` / `TimedOut` and tungstenite
3179/// renders it as `Error::Io` — which the transport mapper classifies
3180/// as `SocketError` (failover-eligible so the walk continues to the
3181/// next host).
3182///
3183/// The deadline is cleared on the way out so subsequent
3184/// `Cursor::next_batch` reads (which can legitimately block for as
3185/// long as the server takes to plan and execute the query) aren't
3186/// subject to it.
3187fn read_server_info_frame(transport: &mut WsTransport, timeout: Duration) -> Result<ServerInfo> {
3188    transport.set_read_timeout(Some(timeout));
3189    let result = transport.read_frame();
3190    transport.set_read_timeout(None);
3191    let (header, payload) = result?;
3192    let mut dict = SymbolDict::new();
3193    let mut query_schema: Option<Schema> = None;
3194    let mut zstd_scratch = ZstdScratch::new();
3195    let event = decode_frame(
3196        header,
3197        &payload,
3198        &mut dict,
3199        &mut query_schema,
3200        &mut zstd_scratch,
3201    )?;
3202    match event {
3203        ServerEvent::ServerInfo(info) => Ok(info),
3204        other => Err(fmt!(
3205            ProtocolError,
3206            "expected SERVER_INFO as the first frame, got {:?}",
3207            std::mem::discriminant(&other)
3208        )),
3209    }
3210}
3211
3212/// Substrings QuestDB uses for the transient "the cached query plan no
3213/// longer matches the table's current metadata version" condition. The
3214/// server raises it (as `INTERNAL_ERROR`, `0x06`) when an async
3215/// `ALTER TABLE ... ALTER COLUMN TYPE` bumps a table's metadata version
3216/// between a SELECT's compilation and its execution. The wire has no
3217/// dedicated "retryable" status — every transient server fault is folded
3218/// into `INTERNAL_ERROR` — so the condition is identified by message text,
3219/// matched case-insensitively. Kept lowercase so the match is a plain
3220/// `contains` against the lowercased server message.
3221const STALE_PLAN_PATTERNS: [&str; 2] = [
3222    "cached query plan cannot be used",
3223    "table schema has changed",
3224];
3225
3226/// True when a server `QUERY_ERROR` is the transient stale-cached-plan
3227/// fault that [`Cursor::next_batch`] recovers from by transparently
3228/// re-issuing the query. Gated on `INTERNAL_ERROR` so a genuinely
3229/// different error that merely echoes the text in its message can't be
3230/// silently swallowed.
3231fn is_stale_plan_error(status: crate::egress::wire::msg_kind::StatusCode, message: &str) -> bool {
3232    use crate::egress::wire::msg_kind::StatusCode as S;
3233    if status != S::InternalError {
3234        return false;
3235    }
3236    let lower = message.to_ascii_lowercase();
3237    STALE_PLAN_PATTERNS.iter().any(|p| lower.contains(p))
3238}
3239
3240fn map_server_status(
3241    status: crate::egress::wire::msg_kind::StatusCode,
3242    message: String,
3243) -> crate::Error {
3244    use crate::ErrorCode as C;
3245    use crate::egress::wire::msg_kind::StatusCode as S;
3246    let code = match status {
3247        S::SchemaMismatch => C::ServerSchemaMismatch,
3248        S::ParseError => C::ServerParseError,
3249        S::InternalError => C::ServerInternalError,
3250        S::SecurityError => C::ServerSecurityError,
3251        S::Cancelled => C::Cancelled,
3252        S::LimitExceeded => C::ServerLimitExceeded,
3253    };
3254    crate::Error::new(code, message)
3255}
3256
3257#[cfg(test)]
3258mod tests {
3259    use super::*;
3260
3261    /// `ReaderStats` lives behind `Arc` so the FFI handle can clone it
3262    /// once and read counters without touching the `UnsafeCell<Reader>`
3263    /// that owns the Reader. This test pins the contract that writes
3264    /// through one clone are observable through any other — the
3265    /// premise the FFI relies on for `_bytes_received` / `_read_ns` /
3266    /// etc. to return up-to-date values without crossing the cell.
3267    #[test]
3268    fn reader_stats_arc_clones_share_storage() {
3269        let stats = Arc::new(ReaderStats::default());
3270        let alias = Arc::clone(&stats);
3271        stats.bytes_received.fetch_add(42, Ordering::Relaxed);
3272        stats.credit_granted_total.fetch_add(7, Ordering::Relaxed);
3273        stats.read_ns.fetch_add(1_000, Ordering::Relaxed);
3274        stats.decode_ns.fetch_add(500, Ordering::Relaxed);
3275        assert_eq!(alias.bytes_received.load(Ordering::Relaxed), 42);
3276        assert_eq!(alias.credit_granted_total.load(Ordering::Relaxed), 7);
3277        assert_eq!(alias.read_ns.load(Ordering::Relaxed), 1_000);
3278        assert_eq!(alias.decode_ns.load(Ordering::Relaxed), 500);
3279        // Reset via the inner Reader's API is visible through the
3280        // FFI's clone too (the contract of `qwp_reader_reset_timing`).
3281        alias.read_ns.store(0, Ordering::Relaxed);
3282        alias.decode_ns.store(0, Ordering::Relaxed);
3283        assert_eq!(stats.read_ns.load(Ordering::Relaxed), 0);
3284        assert_eq!(stats.decode_ns.load(Ordering::Relaxed), 0);
3285    }
3286
3287    /// Anchors `REQUEST_ID_OFFSET` to the actual `QueryRequest::encode`
3288    /// output. The failover-replay path in `Cursor::failover_reconnect_and_replay`
3289    /// patches `[REQUEST_ID_OFFSET..+8]` of the stashed encoded request
3290    /// to substitute a fresh request_id; if `encode` ever grows a prefix,
3291    /// the constant must move with it. This test fails red on any layout
3292    /// drift before the runtime guard in `execute()` would.
3293    #[test]
3294    fn request_id_offset_matches_encoder_layout() {
3295        const RID: i64 = 0x0123_4567_89AB_CDEF;
3296        let req = QueryRequest::builder("SELECT 1")
3297            .request_id(RID)
3298            .build()
3299            .expect("build");
3300        let mut buf = Vec::new();
3301        req.encode(&mut buf).expect("encode");
3302
3303        assert!(buf.len() >= REQUEST_ID_OFFSET + 8);
3304        assert_eq!(buf[0], MsgKind::QueryRequest.as_u8());
3305        let mut id_bytes = [0u8; 8];
3306        id_bytes.copy_from_slice(&buf[REQUEST_ID_OFFSET..REQUEST_ID_OFFSET + 8]);
3307        assert_eq!(i64::from_le_bytes(id_bytes), RID);
3308    }
3309
3310    /// Confirm `patch_request_id` mutates the request_id span and
3311    /// preserves every other byte, on both the unique-owner fast path
3312    /// and the shared-owner fallback path. This is what makes
3313    /// failover-replay zero-copy on the body: the multi-MB tail must
3314    /// be byte-identical to the original after a patch.
3315    #[test]
3316    fn patch_request_id_preserves_body_and_updates_id() {
3317        const OLD_RID: i64 = 0x1111_2222_3333_4444;
3318        const NEW_RID: i64 = 0x5555_6666_7777_8888;
3319        // Build a realistic encoded request so the test exercises the
3320        // same layout the production replay path patches.
3321        let req = QueryRequest::builder("SELECT * FROM big_table WHERE x > $1")
3322            .request_id(OLD_RID)
3323            .build()
3324            .expect("build");
3325        let mut original = Vec::with_capacity(64);
3326        req.encode(&mut original).expect("encode");
3327        let original = Bytes::from(original);
3328
3329        // Unique-owner fast path: only this Bytes references the buffer,
3330        // so try_into_mut succeeds and the patch is in-place.
3331        let patched = patch_request_id(original.clone(), NEW_RID);
3332        // The cloned `original` we kept around drops at scope end; the
3333        // call above received its own clone which write_message would
3334        // consume. Verify the returned Bytes carries the new id.
3335        assert_eq!(patched[0], MsgKind::QueryRequest.as_u8());
3336        let mut id_bytes = [0u8; 8];
3337        id_bytes.copy_from_slice(&patched[REQUEST_ID_OFFSET..REQUEST_ID_OFFSET + 8]);
3338        assert_eq!(i64::from_le_bytes(id_bytes), NEW_RID);
3339        // Body before and after the request_id span is byte-identical.
3340        assert_eq!(
3341            &patched[..REQUEST_ID_OFFSET],
3342            &original[..REQUEST_ID_OFFSET]
3343        );
3344        assert_eq!(
3345            &patched[REQUEST_ID_OFFSET + 8..],
3346            &original[REQUEST_ID_OFFSET + 8..]
3347        );
3348
3349        // Shared-owner fallback: hold an extra clone alive across the
3350        // call so try_into_mut returns Err and patch_request_id falls
3351        // back to BytesMut::from(&shared[..]). Same correctness.
3352        let _hold = patched.clone();
3353        let patched_again = patch_request_id(patched, OLD_RID);
3354        let mut id_bytes = [0u8; 8];
3355        id_bytes.copy_from_slice(&patched_again[REQUEST_ID_OFFSET..REQUEST_ID_OFFSET + 8]);
3356        assert_eq!(i64::from_le_bytes(id_bytes), OLD_RID);
3357    }
3358
3359    /// Exhaustively pin `is_failover_eligible` against every
3360    /// `ErrorCode` variant. The function is a single `matches!` arm
3361    /// today; this guards against (a) silently dropping an arm
3362    /// during a refactor, (b) accidentally promoting a hard error
3363    /// (auth, config) into the eligible set, which would make the
3364    /// failover loop bounce off identical-failure endpoints. Adding
3365    /// a new `ErrorCode` variant later forces this test to be
3366    /// updated — that's the point.
3367    #[test]
3368    fn is_failover_eligible_matrix() {
3369        use ErrorCode::*;
3370        // Eligible: every transport-level failure that may differ
3371        // between endpoints, plus RoleMismatch (soft skip).
3372        for code in [
3373            SocketError,
3374            ConnectTimeout,
3375            HandshakeError,
3376            TlsError,
3377            ProtocolError,
3378            CouldNotResolveAddr,
3379            RoleMismatch,
3380        ] {
3381            assert!(
3382                is_failover_eligible(code),
3383                "{:?} must be failover-eligible",
3384                code
3385            );
3386        }
3387        // Not eligible: failures that signal a hard problem
3388        // (credentials, config, server build) which would fail
3389        // identically on every endpoint, OR are client-side
3390        // validation errors / server-reported terminals that aren't
3391        // about transport.
3392        for code in [
3393            ConfigError,
3394            InvalidApiCall,
3395            AuthError,
3396            UnsupportedServer,
3397            InvalidUtf8,
3398            InvalidBind,
3399            ServerSchemaMismatch,
3400            ServerParseError,
3401            ServerInternalError,
3402            ServerSecurityError,
3403            LimitExceeded,
3404            ServerLimitExceeded,
3405            Cancelled,
3406        ] {
3407            assert!(
3408                !is_failover_eligible(code),
3409                "{:?} must NOT be failover-eligible",
3410                code
3411            );
3412        }
3413    }
3414
3415    /// Pin the `StatusCode` → `ErrorCode` mapping. Every server-reported
3416    /// terminal status maps to a distinct `ErrorCode`; a refactor that
3417    /// merges two arms (e.g. lumps `LimitExceeded` and `InternalError`
3418    /// together) would silently swallow useful per-status discrimination.
3419    /// Adding a new `StatusCode` variant later forces this test to be
3420    /// updated — that's the point.
3421    #[test]
3422    fn map_server_status_matrix() {
3423        use crate::egress::wire::msg_kind::StatusCode as S;
3424        use ErrorCode as C;
3425
3426        let cases: &[(S, C)] = &[
3427            (S::SchemaMismatch, C::ServerSchemaMismatch),
3428            (S::ParseError, C::ServerParseError),
3429            (S::InternalError, C::ServerInternalError),
3430            (S::SecurityError, C::ServerSecurityError),
3431            (S::Cancelled, C::Cancelled),
3432            (S::LimitExceeded, C::ServerLimitExceeded),
3433        ];
3434
3435        for (status, expected_code) in cases {
3436            let err = map_server_status(*status, "msg".to_string());
3437            assert_eq!(
3438                err.code(),
3439                *expected_code,
3440                "status {:?} should map to {:?}",
3441                status,
3442                expected_code
3443            );
3444            assert_eq!(err.msg(), "msg");
3445        }
3446
3447        // Sanity: each ErrorCode in the table is unique. If two
3448        // statuses ever collapse to the same code, this assertion
3449        // surfaces it — the matrix above could be wrong-but-passing if
3450        // both sides changed in lockstep.
3451        let mut seen = std::collections::HashSet::new();
3452        for (_, code) in cases {
3453            assert!(
3454                seen.insert(*code),
3455                "ErrorCode {:?} mapped from two distinct StatusCode values",
3456                code
3457            );
3458        }
3459    }
3460
3461    /// Pin `prefer_over_trigger`: the failover loop surfaces these
3462    /// codes in place of the original transport `trigger` because
3463    /// they tell the user *what to fix* (credentials, topology,
3464    /// server build, config). Bouncing through the matrix locks the
3465    /// predicate so a refactor that drops `UnsupportedServer` or
3466    /// `ConfigError` from the preferred set goes red.
3467    #[test]
3468    fn prefer_over_trigger_matrix() {
3469        use ErrorCode::*;
3470        for code in [
3471            AuthError,
3472            RoleMismatch,
3473            ConfigError,
3474            UnsupportedServer,
3475            HandshakeError,
3476            TlsError,
3477        ] {
3478            assert!(
3479                prefer_over_trigger(code),
3480                "{:?} must be preferred over the trigger",
3481                code
3482            );
3483        }
3484        // Generic transport flops, decode failures, and client-side
3485        // validation errors are NOT more diagnostic than the trigger
3486        // — keep the original cause-of-death in those cases.
3487        for code in [
3488            SocketError,
3489            ProtocolError,
3490            CouldNotResolveAddr,
3491            InvalidApiCall,
3492            InvalidUtf8,
3493            InvalidBind,
3494            ServerInternalError,
3495            Cancelled,
3496        ] {
3497            assert!(
3498                !prefer_over_trigger(code),
3499                "{:?} must NOT be preferred over the trigger",
3500                code
3501            );
3502        }
3503    }
3504
3505    /// Pin the exponential base schedule without measuring wall-clock
3506    /// time. Socket setup and scheduler delays are unrelated to the
3507    /// configured backoff, so elapsed-time integration assertions can
3508    /// fail even when this progression is correct.
3509    #[test]
3510    fn failover_budget_backoff_base_grows_and_caps() {
3511        let mut budget = FailoverBudget {
3512            reconnect_rounds_remaining: 9,
3513            next_backoff_ms: 10,
3514            deadline: None,
3515        };
3516        let observed: [u64; 9] = std::array::from_fn(|_| {
3517            let base = budget.next_backoff_ms;
3518            budget.advance_backoff(20);
3519            base
3520        });
3521        assert_eq!(observed, [10, 20, 20, 20, 20, 20, 20, 20, 20]);
3522
3523        let mut disabled = FailoverBudget {
3524            reconnect_rounds_remaining: 1,
3525            next_backoff_ms: 0,
3526            deadline: None,
3527        };
3528        disabled.advance_backoff(20);
3529        assert_eq!(disabled.next_backoff_ms, 0);
3530    }
3531
3532    /// Exercise the production call path without using elapsed time as
3533    /// an oracle. The configured bases keep the real sleeps at 0–1 ms;
3534    /// scheduler oversleep can delay the test but cannot change its
3535    /// state assertions.
3536    #[test]
3537    fn before_reconnect_round_applies_configured_backoff_cap() {
3538        let cfg = ReaderConfig::from_conf(concat!(
3539            "ws::addr=localhost:9000;",
3540            "failover_max_attempts=4;",
3541            "failover_backoff_initial_ms=1;",
3542            "failover_backoff_max_ms=2"
3543        ))
3544        .unwrap();
3545        let mut budget = FailoverBudget::new(&cfg);
3546        let mut rng = FailoverRng { state: 0 };
3547
3548        let observed: [u64; 3] = std::array::from_fn(|_| {
3549            budget.before_reconnect_round(&cfg, &mut rng).unwrap();
3550            budget.next_backoff_ms
3551        });
3552        assert_eq!(observed, [2, 2, 2]);
3553        assert_eq!(budget.reconnect_rounds_remaining, 0);
3554        assert_eq!(
3555            budget.before_reconnect_round(&cfg, &mut rng),
3556            Err(FailoverBudgetStop::AttemptsExhausted)
3557        );
3558
3559        let disabled_cfg = ReaderConfig::from_conf(concat!(
3560            "ws::addr=localhost:9000;",
3561            "failover_max_attempts=2;",
3562            "failover_backoff_initial_ms=0;",
3563            "failover_backoff_max_ms=2"
3564        ))
3565        .unwrap();
3566        let mut disabled = FailoverBudget::new(&disabled_cfg);
3567        let mut disabled_rng = FailoverRng { state: 0 };
3568        disabled
3569            .before_reconnect_round(&disabled_cfg, &mut disabled_rng)
3570            .unwrap();
3571        assert_eq!(disabled.next_backoff_ms, 0);
3572        assert_eq!(disabled.reconnect_rounds_remaining, 0);
3573    }
3574
3575    /// `base = 0` MUST return 0 without touching the splitmix state.
3576    /// A backoff of zero is the documented "sleep is a no-op" sentinel
3577    /// and the caller passes it whenever `failover_backoff_initial_ms`
3578    /// has been driven to zero by repeated doubling under saturation.
3579    #[test]
3580    fn full_jitter_ms_zero_base_returns_zero() {
3581        let mut rng = FailoverRng::new();
3582        for _ in 0..32 {
3583            assert_eq!(rng.full_jitter_ms(0), 0);
3584        }
3585    }
3586
3587    /// Every draw lies in `[0, base)` — the full-jitter contract from
3588    /// failover.md §3.1. SF ingress uses a different scheme (centered
3589    /// jitter, `[base/2, 3*base/2)`, in `qwp_ws_driver.rs`); this test
3590    /// pins the egress full-jitter contract, which — unlike that — may
3591    /// wait near zero. 10k samples per base across several bases
3592    /// (powers of two, near-`u32::MAX`, and primes that exercise the
3593    /// `% base` reduction) catches both off-by-one and signed/unsigned
3594    /// mix-ups.
3595    #[test]
3596    fn full_jitter_ms_draws_are_in_range() {
3597        let mut rng = FailoverRng::new();
3598        for &base in &[1u64, 2, 80, 100, 1_000, 65_537, u32::MAX as u64] {
3599            for _ in 0..10_000 {
3600                let d = rng.full_jitter_ms(base);
3601                assert!(
3602                    d < base,
3603                    "full_jitter_ms({}) returned {}, which is >= base \
3604                     (full-jitter draws must be in [0, base))",
3605                    base,
3606                    d
3607                );
3608            }
3609        }
3610    }
3611
3612    /// The draws span the full `[0, base)` range, not a clamped sub-
3613    /// interval. With `base = 100` and 10k samples drawn from a
3614    /// Splitmix64-derived uniform, statistical guarantees are
3615    /// effectively certain: P(no sample < 10) = (0.9)^10000 ≈ 10^-457,
3616    /// and likewise for >= 90. A regression to a constant or a
3617    /// half-range clamp would fail one of the two assertions
3618    /// deterministically. This replaces the prior wall-clock-based
3619    /// `failover_backoff_uses_full_jitter` test, which had to drown
3620    /// scheduler noise out of an integration measurement.
3621    #[test]
3622    fn full_jitter_ms_distribution_covers_full_range() {
3623        let mut rng = FailoverRng::new();
3624        let mut saw_low = false;
3625        let mut saw_high = false;
3626        for _ in 0..10_000 {
3627            let d = rng.full_jitter_ms(100);
3628            if d < 10 {
3629                saw_low = true;
3630            }
3631            if d >= 90 {
3632                saw_high = true;
3633            }
3634            if saw_low && saw_high {
3635                break;
3636            }
3637        }
3638        assert!(
3639            saw_low,
3640            "expected at least one draw < 10 out of 10k samples"
3641        );
3642        assert!(
3643            saw_high,
3644            "expected at least one draw >= 90 out of 10k samples"
3645        );
3646    }
3647
3648    /// Truth-table coverage for the silent-duplicate guard.
3649    ///
3650    /// The four input combinations cover every reachable cursor state
3651    /// at the moment a failover-eligible transport error fires.
3652    /// Only `on_failover_reset` is replay-aware. The progress callback is
3653    /// telemetry-only and does not affect this predicate.
3654    ///
3655    /// | data_delivered | reset callback installed | refuses replay? |
3656    /// |----------------|--------------------------|-----------------|
3657    /// | false          | false                     | no — initial-connect-style failover, transparent |
3658    /// | false          | true                      | no — caller will be notified anyway |
3659    /// | true           | false                     | **YES** — silent duplicates would otherwise reach the caller |
3660    /// | true           | true                      | no — caller opted in to replays |
3661    ///
3662    /// A regression that flipped the predicate (e.g. inverted the
3663    /// callback check or removed the data-delivered latch) would fail
3664    /// at least one row of this matrix.
3665    #[test]
3666    fn would_silently_duplicate_truth_table() {
3667        // No data yet — failover is always safe, regardless of reset hook.
3668        assert!(!would_silently_duplicate(false, false));
3669        assert!(!would_silently_duplicate(false, true));
3670        // Data already delivered — only the reset hook unlocks replay.
3671        assert!(would_silently_duplicate(true, false));
3672        assert!(!would_silently_duplicate(true, true));
3673    }
3674}