This is an automated email from the ASF dual-hosted git repository. CritasWang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/iotdb-client-rust.git
commit 346a0e01b4023f9770424406aa85148e0b231968 Author: CritasWang <[email protected]> AuthorDate: Mon Jul 13 11:01:34 2026 +0800 V3: auto-reconnect with endpoint failover (C#-style, original error surfaces); write-redirect cache (TTL+eviction) with pool acquire_for_device routing --- src/client/mod.rs | 2 + src/client/pool.rs | 119 ++++++++++++++ src/client/redirect.rs | 334 +++++++++++++++++++++++++++++++++++++ src/client/session.rs | 435 ++++++++++++++++++++++++++++++++++++++++++++++--- src/lib.rs | 1 + 5 files changed, 864 insertions(+), 27 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 917cd13..a457c06 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -21,10 +21,12 @@ pub mod dataset; pub mod pool; +pub mod redirect; pub mod session; pub mod table_session; pub use dataset::{Row, SessionDataSet}; pub use pool::{PooledSession, SessionPool, SessionPoolConfig, TableSessionPool}; +pub use redirect::{RedirectCache, RedirectCacheStats}; pub use session::{QueryHandle, Session, SessionConfig}; pub use table_session::{TableSession, TableSessionBuilder}; diff --git a/src/client/pool.rs b/src/client/pool.rs index 2e6a864..887aff5 100644 --- a/src/client/pool.rs +++ b/src/client/pool.rs @@ -177,6 +177,43 @@ impl SessionPool { } } + /// Acquire a session for writes to `device_id`, preferring an idle + /// session already connected to the device's redirected endpoint. + /// + /// A status-400 insert response leaves a device → endpoint hint in the + /// session that saw it ([`Session::redirect_hint`]); if any idle + /// session carries such a hint **and** another idle session is + /// connected to that endpoint, that session is handed out. In every + /// other case — no hint, no matching idle session, hint expired — this + /// is exactly [`SessionPool::acquire`]. The pool never opens a new + /// connection to the hinted endpoint (Node.js-style dedicated + /// per-endpoint sessions are future work). + pub fn acquire_for_device(&self, device_id: &str) -> Result<PooledSession<'_>> { + { + let mut state = self.state.lock().expect("pool lock poisoned"); + if !state.closed { + // Any idle session may hold the hint (the one that got the + // 400), not necessarily one connected to the hinted node. + let hint = state + .idle + .iter_mut() + .find_map(|s| s.redirect_hint(device_id)); + if let Some(endpoint) = hint { + let matching = state + .idle + .iter() + .position(|s| s.is_open() && s.current_endpoint() == Some(&endpoint)); + if let Some(pos) = matching { + let session = state.idle.remove(pos).expect("index in bounds"); + drop(state); + return self.hand_out(session); + } + } + } + } + self.acquire() + } + /// Convenience: acquire a session, run one non-query statement, release. /// `USE <db>` propagates to the whole pool via the database tracking. pub fn execute_non_query(&self, sql: &str) -> Result<()> { @@ -472,6 +509,88 @@ mod tests { assert_eq!(pool.pool.config.session.sql_dialect, "table"); } + /// A local listener that accepts and immediately drops connections: + /// `Connection::open` succeeds (giving pool tests real TCP connections + /// with distinct endpoints) while any RPC on them — like the swallowed + /// `closeSession` at pool teardown — fails fast on EOF instead of + /// blocking on a reply that never comes. + fn fake_listener() -> Endpoint { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("local_addr").port(); + std::thread::spawn(move || { + for stream in listener.incoming() { + match stream { + Ok(s) => drop(s), + Err(_) => break, + } + } + }); + Endpoint::new("127.0.0.1", port) + } + + fn injected_session(endpoint: &Endpoint) -> Session { + let mut session = Session::new(dead_endpoint_config()); + let connection = + crate::connection::Connection::open(endpoint.clone(), Duration::from_millis(500)) + .expect("connect to test listener"); + session.test_inject_connection(connection); + session + } + + #[test] + fn acquire_for_device_prefers_hinted_endpoint() { + let ep_a = fake_listener(); + let ep_b = fake_listener(); + + let cfg = SessionPoolConfig { + max_size: 4, + acquire_timeout: Duration::from_millis(50), + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + + // s1 (connected to A) holds the redirect hint pointing at B; + // s2 is connected to B — the hinted target. + let mut s1 = injected_session(&ep_a); + s1.test_inject_redirect_hint("root.sg.d1", ep_b.clone()); + let s2 = injected_session(&ep_b); + pool.inject_idle(s1); + pool.inject_idle(s2); + + // Hinted device → the session on endpoint B, even though the + // session on A is first in the idle queue. + let guard = pool.acquire_for_device("root.sg.d1").unwrap(); + assert_eq!(guard.current_endpoint(), Some(&ep_b)); + drop(guard); + + // Unknown device → plain FIFO acquire (the session on A). + let guard = pool.acquire_for_device("root.sg.unknown").unwrap(); + assert_eq!(guard.current_endpoint(), Some(&ep_a)); + } + + #[test] + fn acquire_for_device_falls_back_when_no_session_matches_hint() { + let ep_a = fake_listener(); + let ep_gone = Endpoint::new("10.255.255.1", 6667); // nobody connected here + + let cfg = SessionPoolConfig { + max_size: 4, + acquire_timeout: Duration::from_millis(50), + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + let mut s1 = injected_session(&ep_a); + s1.test_inject_redirect_hint("root.sg.d1", ep_gone); + pool.inject_idle(s1); + + // Hint exists but no idle session is connected to that endpoint → + // normal acquire semantics (FIFO hand-out of s1). + let guard = pool.acquire_for_device("root.sg.d1").unwrap(); + assert_eq!(guard.current_endpoint(), Some(&ep_a)); + } + /// Live-server tests: acquire/release round-trip, reuse, and blocking /// hand-off. Skipped when no IoTDB server is reachable. #[test] diff --git a/src/client/redirect.rs b/src/client/redirect.rs new file mode 100644 index 0000000..4505504 --- /dev/null +++ b/src/client/redirect.rs @@ -0,0 +1,334 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Write-redirection cache (protocol spec §2/§3.7): a status-400 +//! (`REDIRECTION_RECOMMEND`) write **succeeded**, and its `redirectNode` +//! recommends a better endpoint for that device. This cache remembers +//! device → endpoint hints with a TTL, mirroring the Node.js SDK's +//! `RedirectCache` (TTL 300 s, oldest-entry eviction when full). +//! +//! Routing honesty: a single [`crate::Session`] holds exactly one +//! connection, so the cache does **not** reroute inserts by itself — it +//! only records hints, exposed via [`crate::Session::redirect_hint`]. +//! [`crate::SessionPool::acquire_for_device`] consults them to prefer an +//! idle session already connected to the hinted endpoint. Full per-device +//! connection routing (Node.js-style dedicated per-endpoint sessions) is +//! future work. + +use std::collections::HashMap; +use std::time::{Duration, Instant}; + +use crate::connection::Endpoint; +use crate::protocol::common::{TEndPoint, TSStatus}; + +/// Default hint lifetime, matching the Node.js SDK (300 s). +pub const DEFAULT_REDIRECT_TTL: Duration = Duration::from_secs(300); +/// Default capacity; the oldest entry is evicted when full. +pub const DEFAULT_REDIRECT_MAX_ENTRIES: usize = 1024; + +/// TTL predicate, kept as a pure function so expiry logic is testable +/// without sleeping: an entry is expired once `elapsed` exceeds `ttl`. +/// A zero `ttl` disables expiry entirely (Node.js `ttl > 0` semantics). +fn is_expired(elapsed: Duration, ttl: Duration) -> bool { + !ttl.is_zero() && elapsed > ttl +} + +struct Entry { + endpoint: Endpoint, + inserted: Instant, + /// Monotonic insertion counter — eviction order stays deterministic + /// even when two inserts land on the same `Instant`. + seq: u64, +} + +/// Snapshot of a [`RedirectCache`]'s configuration and occupancy. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RedirectCacheStats { + pub size: usize, + pub max_entries: usize, + pub ttl: Duration, +} + +/// Device → endpoint hint cache fed by status-400 `redirectNode`s. +pub struct RedirectCache { + entries: HashMap<String, Entry>, + ttl: Duration, + max_entries: usize, + seq: u64, +} + +impl Default for RedirectCache { + fn default() -> Self { + Self::new(DEFAULT_REDIRECT_TTL, DEFAULT_REDIRECT_MAX_ENTRIES) + } +} + +impl RedirectCache { + /// A cache holding up to `max_entries` hints for `ttl` each. Zero `ttl` + /// means hints never expire; zero `max_entries` disables the cache. + pub fn new(ttl: Duration, max_entries: usize) -> Self { + Self { + entries: HashMap::new(), + ttl, + max_entries, + seq: 0, + } + } + + /// The cached endpoint for `device_id`, or `None` when absent or + /// expired (expired entries are removed on the way out). + pub fn get(&mut self, device_id: &str) -> Option<Endpoint> { + self.get_at(device_id, Instant::now()) + } + + /// [`RedirectCache::get`] against an explicit "now" — the seam the TTL + /// tests use instead of sleeping. + fn get_at(&mut self, device_id: &str, now: Instant) -> Option<Endpoint> { + let entry = self.entries.get(device_id)?; + if is_expired(now.saturating_duration_since(entry.inserted), self.ttl) { + self.entries.remove(device_id); + return None; + } + Some(entry.endpoint.clone()) + } + + /// Record (or refresh) the hint for `device_id`. When the cache is full + /// and the device is new, the oldest-inserted entry is evicted first. + pub fn put(&mut self, device_id: impl Into<String>, endpoint: Endpoint) { + if self.max_entries == 0 { + return; + } + let device_id = device_id.into(); + if self.entries.len() >= self.max_entries && !self.entries.contains_key(&device_id) { + let oldest = self + .entries + .iter() + .min_by_key(|(_, e)| e.seq) + .map(|(k, _)| k.clone()); + if let Some(oldest) = oldest { + self.entries.remove(&oldest); + } + } + self.seq += 1; + self.entries.insert( + device_id, + Entry { + endpoint, + inserted: Instant::now(), + seq: self.seq, + }, + ); + } + + /// Drop all hints. + pub fn clear(&mut self) { + self.entries.clear(); + } + + pub fn stats(&self) -> RedirectCacheStats { + RedirectCacheStats { + size: self.entries.len(), + max_entries: self.max_entries, + ttl: self.ttl, + } + } +} + +/// Convert a redirect `TEndPoint` into a client [`Endpoint`], rejecting +/// nonsense (empty host, port outside `u16`). +fn endpoint_from_redirect(node: &TEndPoint) -> Option<Endpoint> { + if node.ip.is_empty() { + return None; + } + u16::try_from(node.port) + .ok() + .filter(|&p| p != 0) + .map(|port| Endpoint::new(node.ip.clone(), port)) +} + +/// Harvest redirect hints from an insert response into the cache — called +/// on the raw `TSStatus` *before* it is collapsed into a `Result`, since +/// `check_status` treats 400 as plain success and discards the node. +/// +/// A top-level `redirectNode` (single-device inserts and tablet writes) +/// applies to every device in the request; a `subStatus` list that pairs +/// 1:1 with the request's devices (multi-device `insertRecords`) is +/// harvested entry-by-entry. +pub(crate) fn record_redirects(cache: &mut RedirectCache, devices: &[&str], status: &TSStatus) { + if let Some(subs) = status.sub_status.as_deref() { + if subs.len() == devices.len() { + for (device, sub) in devices.iter().zip(subs) { + if let Some(endpoint) = sub.redirect_node.as_ref().and_then(endpoint_from_redirect) + { + cache.put(*device, endpoint); + } + } + return; + } + } + if let Some(endpoint) = status + .redirect_node + .as_ref() + .and_then(endpoint_from_redirect) + { + for device in devices { + cache.put(*device, endpoint.clone()); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ep(port: u16) -> Endpoint { + Endpoint::new("10.0.0.1", port) + } + + #[test] + fn ttl_predicate() { + let ttl = Duration::from_secs(300); + assert!(!is_expired(Duration::ZERO, ttl)); + assert!(!is_expired(ttl, ttl)); // boundary: exactly ttl is still fresh + assert!(is_expired(ttl + Duration::from_millis(1), ttl)); + // Zero TTL disables expiry. + assert!(!is_expired(Duration::from_secs(1 << 20), Duration::ZERO)); + } + + #[test] + fn put_get_roundtrip_and_clear() { + let mut cache = RedirectCache::default(); + assert_eq!(cache.get("root.sg.d1"), None); + cache.put("root.sg.d1", ep(6667)); + assert_eq!(cache.get("root.sg.d1"), Some(ep(6667))); + // Refresh overwrites. + cache.put("root.sg.d1", ep(6668)); + assert_eq!(cache.get("root.sg.d1"), Some(ep(6668))); + assert_eq!(cache.stats().size, 1); + cache.clear(); + assert_eq!(cache.stats().size, 0); + assert_eq!(cache.get("root.sg.d1"), None); + } + + #[test] + fn expired_entries_are_removed_on_get() { + let mut cache = RedirectCache::new(Duration::from_secs(300), 16); + cache.put("root.sg.d1", ep(6667)); + // Fresh "now": hit. + assert_eq!(cache.get_at("root.sg.d1", Instant::now()), Some(ep(6667))); + // A "now" past the TTL: miss, and the entry is gone afterwards. + let later = Instant::now() + Duration::from_secs(301); + assert_eq!(cache.get_at("root.sg.d1", later), None); + assert_eq!(cache.stats().size, 0); + } + + #[test] + fn zero_ttl_never_expires() { + let mut cache = RedirectCache::new(Duration::ZERO, 16); + cache.put("root.sg.d1", ep(6667)); + let far_future = Instant::now() + Duration::from_secs(1 << 20); + assert_eq!(cache.get_at("root.sg.d1", far_future), Some(ep(6667))); + } + + #[test] + fn eviction_drops_oldest_when_full() { + let mut cache = RedirectCache::new(Duration::from_secs(300), 2); + cache.put("d1", ep(1)); + cache.put("d2", ep(2)); + cache.put("d3", ep(3)); // evicts d1 (oldest insertion) + assert_eq!(cache.get("d1"), None); + assert_eq!(cache.get("d2"), Some(ep(2))); + assert_eq!(cache.get("d3"), Some(ep(3))); + assert_eq!(cache.stats().size, 2); + + // Refreshing an existing key does not evict. + cache.put("d2", ep(22)); + assert_eq!(cache.get("d2"), Some(ep(22))); + assert_eq!(cache.get("d3"), Some(ep(3))); + + // A refreshed key counts as newer: next eviction takes d3. + cache.put("d4", ep(4)); + assert_eq!(cache.get("d3"), None); + assert_eq!(cache.get("d2"), Some(ep(22))); + assert_eq!(cache.get("d4"), Some(ep(4))); + } + + #[test] + fn zero_capacity_disables_cache() { + let mut cache = RedirectCache::new(Duration::from_secs(300), 0); + cache.put("d1", ep(1)); + assert_eq!(cache.get("d1"), None); + assert_eq!( + cache.stats(), + RedirectCacheStats { + size: 0, + max_entries: 0, + ttl: Duration::from_secs(300), + } + ); + } + + fn status(code: i32) -> TSStatus { + TSStatus::new(code, None, None, None, None, None) + } + + fn tendpoint(ip: &str, port: i32) -> TEndPoint { + TEndPoint::new(ip.to_string(), port) + } + + #[test] + fn record_top_level_redirect_applies_to_all_devices() { + let mut cache = RedirectCache::default(); + let mut s = status(400); + s.redirect_node = Some(tendpoint("10.0.0.9", 6667)); + record_redirects(&mut cache, &["root.sg.d1", "root.sg.d2"], &s); + assert_eq!( + cache.get("root.sg.d1"), + Some(Endpoint::new("10.0.0.9", 6667)) + ); + assert_eq!( + cache.get("root.sg.d2"), + Some(Endpoint::new("10.0.0.9", 6667)) + ); + } + + #[test] + fn record_sub_status_redirects_per_device() { + let mut cache = RedirectCache::default(); + let mut ok = status(200); + ok.redirect_node = Some(tendpoint("10.0.0.7", 6667)); + let plain = status(200); + let mut s = status(302); + s.sub_status = Some(vec![Box::new(ok), Box::new(plain)]); + record_redirects(&mut cache, &["d1", "d2"], &s); + assert_eq!(cache.get("d1"), Some(Endpoint::new("10.0.0.7", 6667))); + assert_eq!(cache.get("d2"), None); + } + + #[test] + fn record_ignores_absent_and_invalid_redirect_nodes() { + let mut cache = RedirectCache::default(); + record_redirects(&mut cache, &["d1"], &status(200)); + let mut bad_port = status(400); + bad_port.redirect_node = Some(tendpoint("10.0.0.9", 70000)); + record_redirects(&mut cache, &["d1"], &bad_port); + let mut empty_ip = status(400); + empty_ip.redirect_node = Some(tendpoint("", 6667)); + record_redirects(&mut cache, &["d1"], &empty_ip); + assert_eq!(cache.stats().size, 0); + } +} diff --git a/src/client/session.rs b/src/client/session.rs index 265aee0..96199df 100644 --- a/src/client/session.rs +++ b/src/client/session.rs @@ -22,6 +22,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use crate::client::dataset::SessionDataSet; +use crate::client::redirect::{self, RedirectCache, RedirectCacheStats}; use crate::connection::{Connection, Endpoint}; use crate::data::record::serialize_record_values; use crate::data::{Tablet, Value}; @@ -61,6 +62,17 @@ pub struct SessionConfig { pub query_timeout_ms: i64, /// Database to select at open time (table dialect; sent as config key `db`). pub database: Option<String>, + /// Reopen the connection and retry an op once when an RPC fails at the + /// Thrift/transport level (C# `Reconnect` behavior). Default `true`. + pub enable_auto_reconnect: bool, + /// Full round-robin passes over the endpoints during a reconnect before + /// giving up (C# `RetryNum`). Default 3. + pub max_reconnect_attempts: usize, + /// Pause between reconnect passes. Default 1 s. + pub retry_interval: Duration, + /// Harvest status-400 `redirectNode` hints from insert responses into + /// the per-session [`RedirectCache`]. Default `true`. + pub enable_redirection: bool, } impl Default for SessionConfig { @@ -75,6 +87,10 @@ impl Default for SessionConfig { connect_timeout: Duration::from_secs(10), query_timeout_ms: 60_000, database: None, + enable_auto_reconnect: true, + max_reconnect_attempts: 3, + retry_interval: Duration::from_secs(1), + enable_redirection: true, } } } @@ -118,6 +134,11 @@ pub struct Session { statement_id: i64, /// Current database, tracked from `USE <db>` responses (table dialect). database: Option<String>, + /// Endpoint of the most recent (possibly dead) connection — reconnect + /// starts its round-robin here, like the C# SDK. + last_endpoint: Option<Endpoint>, + /// Device → endpoint hints harvested from status-400 insert responses. + redirect_cache: RedirectCache, } impl Session { @@ -129,6 +150,8 @@ impl Session { session_id: -1, statement_id: -1, database, + last_endpoint: None, + redirect_cache: RedirectCache::default(), } } @@ -160,6 +183,19 @@ impl Session { last_err.unwrap_or_else(|| Error::Client("no endpoints configured".into())) })?; + let (session_id, statement_id) = self.authenticate(&mut connection)?; + self.session_id = session_id; + self.statement_id = statement_id; + self.last_endpoint = Some(connection.endpoint().clone()); + self.connection = Some(connection); + Ok(()) + } + + /// Handshake on a fresh connection: `openSession` (dialect + current + /// database via the `db` config key) followed by `requestStatementId`. + /// Shared by [`Session::open`] and reconnect, so a reopened session + /// lands back in the database its `USE <db>` had selected. + fn authenticate(&self, connection: &mut Connection) -> Result<(i64, i64)> { let mut configuration = BTreeMap::new(); configuration.insert("sql_dialect".to_string(), self.config.sql_dialect.clone()); if let Some(db) = &self.database { @@ -181,14 +217,87 @@ impl Session { resp.server_protocol_version ); } - self.session_id = resp + let session_id = resp .session_id .ok_or_else(|| Error::Client("openSession response missing sessionId".into()))?; - self.statement_id = connection - .client_mut() - .request_statement_id(self.session_id)?; - self.connection = Some(connection); - Ok(()) + let statement_id = connection.client_mut().request_statement_id(session_id)?; + Ok((session_id, statement_id)) + } + + /// Reopen after a transport-level failure: drop the dead connection + /// (closing its transport), then try a **full** handshake — connect + + /// openSession + requestStatementId — round-robin over all endpoints, + /// starting at the one that just died, for up to + /// `max_reconnect_attempts` passes with `retry_interval` between passes + /// (mirroring the C# SDK's `Reconnect`). + fn reconnect(&mut self) -> Result<()> { + self.connection = None; // drop closes the old transport + let n = self.config.endpoints.len(); + if n == 0 { + return Err(Error::Client("no endpoints configured".into())); + } + let start = self + .last_endpoint + .as_ref() + .and_then(|ep| self.config.endpoints.iter().position(|e| e == ep)) + .unwrap_or(0); + let attempts = self.config.max_reconnect_attempts.max(1); + let mut last_err: Option<Error> = None; + for attempt in 0..attempts { + if attempt > 0 { + std::thread::sleep(self.config.retry_interval); + } + for i in 0..n { + let endpoint = self.config.endpoints[(start + i) % n].clone(); + let result = Connection::open(endpoint, self.config.connect_timeout).and_then( + |mut connection| { + let ids = self.authenticate(&mut connection)?; + Ok((connection, ids)) + }, + ); + match result { + Ok((connection, (session_id, statement_id))) => { + log::info!("reconnected to {}", connection.endpoint()); + self.session_id = session_id; + self.statement_id = statement_id; + self.last_endpoint = Some(connection.endpoint().clone()); + self.connection = Some(connection); + return Ok(()); + } + Err(e) => { + log::warn!("reconnect attempt {}/{attempts} failed: {e}", attempt + 1); + last_err = Some(e); + } + } + } + } + Err(last_err.unwrap_or_else(|| Error::Client("reconnect failed".into()))) + } + + /// Run an RPC op; on a transport-level failure ([`Error::Thrift`]) with + /// auto-reconnect enabled, reopen the session and retry the op exactly + /// once (C# `ExecuteClientOperationAsync`). Server status errors pass + /// through untouched. If the reconnect itself fails, the **original** + /// error is surfaced. Ops must rebuild their request inside the closure: + /// reconnecting swaps the connection and refreshes + /// `session_id`/`statement_id`. + fn with_retry<T>(&mut self, mut op: impl FnMut(&mut Self) -> Result<T>) -> Result<T> { + let original = match op(self) { + Err(e @ Error::Thrift(_)) + if self.config.enable_auto_reconnect && self.connection.is_some() => + { + e + } + other => return other, + }; + log::warn!("RPC failed at transport level ({original}); reconnecting"); + match self.reconnect() { + Ok(()) => op(self), + Err(reconnect_err) => { + log::warn!("reconnect failed ({reconnect_err}); surfacing the original error"); + Err(original) + } + } } pub fn is_open(&self) -> bool { @@ -200,6 +309,42 @@ impl Session { self.database.as_deref() } + /// The endpoint this session is currently connected to, if open. + pub fn current_endpoint(&self) -> Option<&Endpoint> { + self.connection.as_ref().map(Connection::endpoint) + } + + /// The cached redirect endpoint for `device_id`, if a status-400 insert + /// response recommended one and the hint has not expired. + /// + /// The session itself does **not** act on these hints (it holds a + /// single connection); [`crate::SessionPool::acquire_for_device`] + /// consults them to prefer a matching idle session. See + /// [`crate::client::redirect`] for the routing-honesty note. + pub fn redirect_hint(&mut self, device_id: &str) -> Option<Endpoint> { + self.redirect_cache.get(device_id) + } + + /// Occupancy/config snapshot of the redirect cache. + pub fn redirect_cache_stats(&self) -> RedirectCacheStats { + self.redirect_cache.stats() + } + + /// Drop all cached redirect hints. + pub fn clear_redirect_cache(&mut self) { + self.redirect_cache.clear(); + } + + /// Inspect an insert response for redirect hints before collapsing it + /// into a `Result` — `check_status` treats 400 as plain success and + /// would discard the recommended node. + fn check_insert_status(&mut self, devices: &[&str], status: &TSStatus) -> Result<()> { + if self.config.enable_redirection { + redirect::record_redirects(&mut self.redirect_cache, devices, status); + } + check_status(status) + } + fn connection_mut(&mut self) -> Result<&mut Connection> { self.connection .as_mut() @@ -209,11 +354,13 @@ impl Session { /// Execute a non-query statement (DDL/DML). Tracks `USE <db>` via the /// response's `database` field. pub fn execute_non_query(&mut self, sql: &str) -> Result<()> { - let req = self.statement_req(sql); - let resp = self - .connection_mut()? - .client_mut() - .execute_update_statement_v2(req)?; + let resp = self.with_retry(|session| { + let req = session.statement_req(sql); + Ok(session + .connection_mut()? + .client_mut() + .execute_update_statement_v2(req)?) + })?; check_status(&resp.status)?; if let Some(db) = resp.database { self.database = Some(db); @@ -234,11 +381,16 @@ impl Session { /// Low-level path: decoding and pagination are the caller's problem — /// prefer [`Session::execute_query`]. pub fn execute_query_raw(&mut self, sql: &str) -> Result<QueryHandle> { - let req = self.statement_req(sql); - let resp = self - .connection_mut()? - .client_mut() - .execute_query_statement_v2(req)?; + // Retry covers only this initial execute RPC: once a query id + // exists, the result set is pinned to its node and cannot be + // migrated by a reconnect (spec gotcha #13). + let resp = self.with_retry(|session| { + let req = session.statement_req(sql); + Ok(session + .connection_mut()? + .client_mut() + .execute_query_statement_v2(req)?) + })?; check_status(&resp.status)?; let query_id = resp .query_id @@ -354,8 +506,13 @@ impl Session { None, None, ); - let status = self.connection_mut()?.client_mut().insert_tablet(req)?; - check_status(&status) + let status = self.with_retry(|session| { + // Clone per attempt: a reconnect refreshes the session id. + let mut req = req.clone(); + req.session_id = session.session_id; + Ok(session.connection_mut()?.client_mut().insert_tablet(req)?) + })?; + self.check_insert_status(&[prefix_path], &status) } /// Insert one row for one device via `insertRecord`. `values[i]` pairs @@ -389,8 +546,12 @@ impl Session { None, None, ); - let status = self.connection_mut()?.client_mut().insert_record(req)?; - check_status(&status) + let status = self.with_retry(|session| { + let mut req = req.clone(); + req.session_id = session.session_id; + Ok(session.connection_mut()?.client_mut().insert_record(req)?) + })?; + self.check_insert_status(&[device_id], &status) } /// Insert one row per device via `insertRecords` (multi-device batch). @@ -449,8 +610,13 @@ impl Session { kept_timestamps, is_aligned, ); - let status = self.connection_mut()?.client_mut().insert_records(req)?; - check_status(&status) + let status = self.with_retry(|session| { + let mut req = req.clone(); + req.session_id = session.session_id; + Ok(session.connection_mut()?.client_mut().insert_records(req)?) + })?; + let devices: Vec<&str> = req.prefix_paths.iter().map(String::as_str).collect(); + self.check_insert_status(&devices, &status) } /// Insert multiple rows for one device via `insertRecordsOfOneDevice`. @@ -507,11 +673,15 @@ impl Session { timestamps, is_aligned, ); - let status = self - .connection_mut()? - .client_mut() - .insert_records_of_one_device(req)?; - check_status(&status) + let status = self.with_retry(|session| { + let mut req = req.clone(); + req.session_id = session.session_id; + Ok(session + .connection_mut()? + .client_mut() + .insert_records_of_one_device(req)?) + })?; + self.check_insert_status(&[device_id], &status) } /// [`Session::insert_record`] against an aligned device @@ -587,6 +757,23 @@ impl Drop for Session { } } +#[cfg(test)] +impl Session { + /// Test hook: install a raw connection (no handshake) so transport + /// failure paths can be exercised without a live server. + pub(crate) fn test_inject_connection(&mut self, connection: Connection) { + self.last_endpoint = Some(connection.endpoint().clone()); + self.connection = Some(connection); + self.session_id = 1; + self.statement_id = 1; + } + + /// Test hook: seed a redirect hint as if a status-400 insert had. + pub(crate) fn test_inject_redirect_hint(&mut self, device_id: &str, endpoint: Endpoint) { + self.redirect_cache.put(device_id, endpoint); + } +} + /// Stably sorts one-device record rows by timestamp — reordering the /// measurement lists in step and serializing each row's value buffer in the /// sorted order (Java `genTSInsertRecordsOfOneDeviceReq`; the server @@ -678,6 +865,12 @@ pub fn check_status(status: &TSStatus) -> Result<()> { mod tests { use super::*; + /// Serializes the live tests that create/delete databases: concurrent + /// `DELETE DATABASE` on a small single-node server can transiently + /// leave no available DataRegionGroups (server error 906) for other + /// tests' inserts. + static LIVE_DB_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + fn status(code: i32) -> TSStatus { TSStatus::new(code, None, None, None, None, None) } @@ -693,6 +886,10 @@ mod tests { assert_eq!(cfg.query_timeout_ms, 60_000); assert_eq!(cfg.connect_timeout, Duration::from_secs(10)); assert!(cfg.database.is_none()); + assert!(cfg.enable_auto_reconnect); + assert_eq!(cfg.max_reconnect_attempts, 3); + assert_eq!(cfg.retry_interval, Duration::from_secs(1)); + assert!(cfg.enable_redirection); } #[test] @@ -871,6 +1068,106 @@ mod tests { assert!(matches!(err, Error::Client(m) if m.contains("length mismatch"))); } + #[test] + fn insert_status_records_redirect_hint() { + use crate::protocol::common::TEndPoint; + let mut session = Session::new(SessionConfig::default()); + let mut s = status(400); + s.redirect_node = Some(TEndPoint::new("10.1.1.1".to_string(), 6667)); + // Status 400 is still a successful write… + assert!(session.check_insert_status(&["root.sg.d1"], &s).is_ok()); + // …and its redirect node is now cached for the device. + assert_eq!( + session.redirect_hint("root.sg.d1"), + Some(Endpoint::new("10.1.1.1", 6667)) + ); + assert_eq!(session.redirect_hint("root.sg.other"), None); + assert_eq!(session.redirect_cache_stats().size, 1); + session.clear_redirect_cache(); + assert_eq!(session.redirect_hint("root.sg.d1"), None); + + // With redirection disabled nothing is recorded (but 400 still OK). + let mut session = Session::new(SessionConfig { + enable_redirection: false, + ..Default::default() + }); + assert!(session.check_insert_status(&["root.sg.d1"], &s).is_ok()); + assert_eq!(session.redirect_hint("root.sg.d1"), None); + } + + /// Accept-then-drop listener: TCP connects succeed, but every RPC on + /// the connection dies at the Thrift level. Returns the endpoint and a + /// shared accept counter (the acceptor thread is leaked; it ends with + /// the test process). + fn accept_then_drop_listener() -> (Endpoint, std::sync::Arc<AtomicUsize>) { + use std::net::TcpListener; + use std::sync::Arc; + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("local_addr").port(); + let accepts = Arc::new(AtomicUsize::new(0)); + let counter = accepts.clone(); + std::thread::spawn(move || { + for stream in listener.incoming() { + match stream { + Ok(s) => { + counter.fetch_add(1, Ordering::SeqCst); + drop(s); // immediate close ⇒ peer reads EOF + } + Err(_) => break, + } + } + }); + (Endpoint::new("127.0.0.1", port), accepts) + } + + /// A transport-level RPC failure with auto-reconnect enabled must drive + /// `max_reconnect_attempts` full reconnect passes (visible as fresh TCP + /// connects) and, when they all fail, surface the **original** error. + #[test] + fn transport_failure_reconnects_then_surfaces_original_error() { + let (endpoint, accepts) = accept_then_drop_listener(); + let mut session = Session::new(SessionConfig { + endpoints: vec![endpoint.clone()], + connect_timeout: Duration::from_millis(500), + max_reconnect_attempts: 2, + retry_interval: Duration::from_millis(10), + ..Default::default() + }); + let connection = Connection::open(endpoint.clone(), Duration::from_millis(500)) // + .expect("connect to listener"); + session.test_inject_connection(connection); + assert_eq!(session.current_endpoint(), Some(&endpoint)); + + let err = session.execute_non_query("SHOW DATABASES").unwrap_err(); + assert!(matches!(err, Error::Thrift(_)), "got {err:?}"); + // 1 initial connection + 2 reconnect passes × 1 endpoint. Each + // failed openSession implies its connection was accepted+dropped, + // so the counter is settled once the error is back. + assert_eq!(accepts.load(Ordering::SeqCst), 3); + // The failed reconnect left the session without a connection. + assert!(!session.is_open()); + } + + /// With auto-reconnect disabled the op fails once: no reconnect + /// connects, original error surfaced directly. + #[test] + fn no_reconnect_when_disabled() { + let (endpoint, accepts) = accept_then_drop_listener(); + let mut session = Session::new(SessionConfig { + endpoints: vec![endpoint.clone()], + connect_timeout: Duration::from_millis(500), + enable_auto_reconnect: false, + ..Default::default() + }); + let connection = Connection::open(endpoint, Duration::from_millis(500)) // + .expect("connect to listener"); + session.test_inject_connection(connection); + + let err = session.execute_non_query("SHOW DATABASES").unwrap_err(); + assert!(matches!(err, Error::Thrift(_)), "got {err:?}"); + assert_eq!(accepts.load(Ordering::SeqCst), 1, "no reconnect attempts"); + } + #[test] fn open_without_endpoints_fails() { let mut session = Session::new(SessionConfig { @@ -935,6 +1232,7 @@ mod tests { } const DB: &str = "root.rusttest_date"; + let _guard = LIVE_DB_LOCK.lock().unwrap_or_else(|p| p.into_inner()); // 2026-07-13 as yyyyMMdd; deliberately not near epoch so a // days-since-epoch misinterpretation cannot coincide. const DATE_YYYYMMDD: i32 = 20260713; @@ -994,6 +1292,87 @@ mod tests { session.close().expect("close session"); } + /// V3 regression: with auto-reconnect and redirection at their default + /// (enabled) settings, normal write/read ops behave exactly as before. + /// On a single-node server no status 400 is ever issued, so the + /// redirect cache must stay empty. Skipped when no IoTDB instance is + /// reachable on localhost:6667. + #[test] + fn live_ops_with_retry_and_redirection_enabled() { + use crate::data::{tablet::Tablet, TSDataType, Value}; + use std::net::TcpStream; + if TcpStream::connect_timeout( + &"127.0.0.1:6667".parse().unwrap(), + Duration::from_millis(300), + ) + .is_err() + { + eprintln!( + "skipping live_ops_with_retry_and_redirection_enabled: \ + no IoTDB server on 127.0.0.1:6667" + ); + return; + } + + const DB: &str = "root.rusttest_retry"; + let _guard = LIVE_DB_LOCK.lock().unwrap_or_else(|p| p.into_inner()); + let cfg = SessionConfig::default(); + assert!(cfg.enable_auto_reconnect && cfg.enable_redirection); + + let mut session = Session::new(cfg); + session.open().expect("open session"); + assert_eq!( + session.current_endpoint(), + Some(&Endpoint::new("localhost", 6667)) + ); + + let _ = session.execute_non_query(&format!("DELETE DATABASE {DB}")); + session + .execute_non_query(&format!("CREATE DATABASE {DB}")) + .expect("create database"); + + // Tablet + record inserts, both passing through check_insert_status. + let mut tablet = Tablet::new( + format!("{DB}.d1"), + vec!["v".into()], + vec![TSDataType::Int32], + ) + .expect("tablet"); + tablet + .add_row(1, vec![Some(Value::Int32(7))]) + .expect("add_row"); + session.insert_tablet(&tablet).expect("insert_tablet"); + session + .insert_record( + &format!("{DB}.d1"), + 2, + vec!["v".into()], + &[Value::Int32(8)], + false, + ) + .expect("insert_record"); + + // Single node ⇒ the server never recommends a redirect. + assert_eq!(session.redirect_cache_stats().size, 0); + assert_eq!(session.redirect_hint(&format!("{DB}.d1")), None); + + let mut rows = 0; + { + let mut dataset = session + .execute_query(&format!("SELECT v FROM {DB}.d1")) + .expect("query"); + while dataset.next_row().expect("next_row").is_some() { + rows += 1; + } + } + assert_eq!(rows, 2); + + session + .execute_non_query(&format!("DELETE DATABASE {DB}")) + .expect("cleanup"); + session.close().expect("close session"); + } + /// Value-asserting live roundtrip: unsorted input, nulls, and a row count /// that is a multiple of 8 (stresses the rows/8+1 bitmap padding byte). /// Skipped when no IoTDB instance is reachable on localhost:6667. @@ -1012,6 +1391,7 @@ mod tests { } const DB: &str = "root.rusttest_readback"; + let _guard = LIVE_DB_LOCK.lock().unwrap_or_else(|p| p.into_inner()); const ROWS: i64 = 16; let mut session = Session::new(SessionConfig::default()); @@ -1116,6 +1496,7 @@ mod tests { } const DB: &str = "root.rusttest_records"; + let _guard = LIVE_DB_LOCK.lock().unwrap_or_else(|p| p.into_inner()); let read_rows = |session: &mut Session, sql: &str| -> Vec<(i64, Vec<Value>)> { let mut rows = Vec::new(); diff --git a/src/lib.rs b/src/lib.rs index 9b3d855..35a71a2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,6 +32,7 @@ pub mod protocol; pub use client::dataset::{Row, SessionDataSet}; pub use client::pool::{PooledSession, SessionPool, SessionPoolConfig, TableSessionPool}; +pub use client::redirect::{RedirectCache, RedirectCacheStats}; pub use client::session::{QueryHandle, Session, SessionConfig}; pub use client::table_session::{TableSession, TableSessionBuilder}; pub use connection::Endpoint;
