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}