From 93271ed29c94f78b0dd842265a9bc66aa5d762b2 Mon Sep 17 00:00:00 2001 From: CritasWang Date: Tue, 25 Aug 2026 17:46:32 +0800 Subject: [PATCH] Fix socket timeouts, connection lifecycle and session pool accounting - Bound every blocking socket read/write after the TCP handshake with a new socket_timeout (SO_RCVTIMEO/SO_SNDTIMEO via socket2), applied before the TLS handshake: RPC reads, drop-time closeSession and the TLS handshake can no longer block forever; SO_KEEPALIVE is enabled. connect_timeout is now one total budget per endpoint across all resolved addresses. - Fail over at the handshake level in Session::open (connect + openSession + requestStatementId), mirroring reconnect(). - Mark a connection broken on transport-level failures so is_open() stops lying; pools discard the session instead of handing it out again, and fetch_results (which cannot be retried) also marks it. - SessionPool: spend the acquire_timeout budget on growth/hand-out failures instead of failing instantly, decrement live under the state lock (no lost Condvar wakeups), wake waiters from close() before the blocking closeSession calls, and treat acquire_timeout=Duration::MAX as "wait without deadline" instead of panicking. - Skip the reconnect pacing sleeps for pooled sessions so a pool slot is not held for the full reconnect walk. - Redirect: normalize endpoint matching (case/brackets/loopback), prefer the newest hint when idle sessions hold conflicting hints (process-wide insertion seq), make cache TTL/capacity configurable, and expose TableSessionPool::acquire_for_device. - Add regression tests using silent/fake listeners for every fix above. Closes #4, closes #5, closes #6, closes #7, closes #8, closes #9, closes #10. --- Cargo.toml | 2 + README.md | 23 +++ README_ZH.md | 19 ++ src/client/pool.rs | 289 ++++++++++++++++++++++++++---- src/client/redirect.rs | 26 ++- src/client/session.rs | 338 +++++++++++++++++++++++++++++++++--- src/client/table_session.rs | 17 ++ src/connection/mod.rs | 224 ++++++++++++++++++++++-- 8 files changed, 859 insertions(+), 79 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 5b58379..c95d2d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,8 @@ thrift = "0.23" byteorder = "1.5" chrono = "0.4" log = "0.4" +# SO_RCVTIMEO/SO_SNDTIMEO + SO_KEEPALIVE on the blocking TcpStream. +socket2 = "0.5" # thrift 0.23 allows any uuid 1.x; 1.21+ raises MSRV to Rust 1.85. uuid = "=1.20.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"], optional = true } diff --git a/README.md b/README.md index 93e607b..eaaf898 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,29 @@ cargo run --example table_session cargo run --example session_pool ``` +## Timeouts & liveness + +`connect_timeout` bounds the TCP connect per endpoint attempt. `socket_timeout` +(default 60 s) bounds **every blocking socket read/write after that** — the TLS +handshake, every RPC read, and the best-effort `closeSession` sent when a +session is dropped — so a peer that accepts the connection and then goes silent +cannot hang the client forever. Set it to `None` to restore the old unbounded +blocking behaviour (zero is treated the same as `None`). + +```rust +let config = SessionConfig { + connect_timeout: std::time::Duration::from_secs(10), + socket_timeout: Some(std::time::Duration::from_secs(30)), + ..Default::default() +}; +// or: TableSession::builder().connect_timeout(..).socket_timeout(..) +``` + +On a transport-level failure with `enable_auto_reconnect` (default `true`) the +session reconnects and retries the operation once; without auto-reconnect the +connection is marked broken and pools discard the session instead of handing it +out again. + ## TLS & RPC compression **RPC compression** (IoTDB's term for the Thrift *compact protocol*) is a plain config flag: diff --git a/README_ZH.md b/README_ZH.md index 3da4c98..88942cd 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -172,6 +172,25 @@ cargo run --example table_session cargo run --example session_pool ``` +## 超时与连接活性 + +`connect_timeout` 限定每个端点的 TCP 建连时长。`socket_timeout`（默认 60 s）限定建连之后的 +**每一次阻塞式 socket 读写**——包括 TLS 握手、所有 RPC 读取，以及会话销毁时发送的尽力而为 +`closeSession`——因此"接受连接后保持沉默"的对端不会再让客户端无限阻塞。设为 `None` +可恢复旧的无限阻塞行为（0 与 `None` 等价）。 + +```rust +let config = SessionConfig { + connect_timeout: std::time::Duration::from_secs(10), + socket_timeout: Some(std::time::Duration::from_secs(30)), + ..Default::default() +}; +// 或：TableSession::builder().connect_timeout(..).socket_timeout(..) +``` + +传输层失败时，若 `enable_auto_reconnect`（默认 `true`）开启，会话会重连并重试该操作一次； +未开启自动重连时，连接会被标记为已损坏，连接池将丢弃该会话而不是再次发放。 + ## TLS 与 RPC 压缩 **RPC 压缩**（IoTDB 术语，实为 Thrift *compact 协议*）只是一个配置开关： diff --git a/src/client/pool.rs b/src/client/pool.rs index 3df10aa..df501ab 100644 --- a/src/client/pool.rs +++ b/src/client/pool.rs @@ -204,12 +204,27 @@ impl SessionPool { /// the acquire retried; a fresh session is opened while under /// `max_size`. pub fn acquire(&self) -> Result> { - let deadline = Instant::now() + self.config.acquire_timeout; + // `checked_add` makes `acquire_timeout = Duration::MAX` mean + // "wait without a deadline" instead of panicking on Instant overflow. + let deadline = Instant::now().checked_add(self.config.acquire_timeout); + let mut last_err: Option = None; let mut state = self.state.lock().expect("pool lock poisoned"); loop { if state.closed { return Err(Error::Client("session pool is closed".into())); } + if let Some(deadline) = deadline { + if Instant::now() >= deadline { + return Err(last_err.unwrap_or_else(|| { + Error::Client(format!( + "pool exhausted: no session available within {:?} ({} live, max {})", + self.config.acquire_timeout, + self.live.load(Ordering::Relaxed), + self.config.max_size + )) + })); + } + } let expired = self.sweep_idle(&mut state); if !expired.is_empty() { drop(state); @@ -220,12 +235,25 @@ impl SessionPool { continue; // re-check closed/idle after re-locking } // Idle session available → validate liveness, evict the dead. - while let Some(entry) = state.idle.pop_front() { + if let Some(entry) = state.idle.pop_front() { if entry.session.is_open() { drop(state); - return self.hand_out(entry.session); + match self.hand_out(entry.session) { + Ok(guard) => return Ok(guard), + Err(e) => { + // Account under the lock, then keep spending the + // acquire_timeout budget on the remaining idle + // candidates instead of failing instantly. + last_err = Some(e); + state = self.state.lock().expect("pool lock poisoned"); + self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + continue; + } + } } self.live.fetch_sub(1, Ordering::Relaxed); + continue; // dead session discarded; retry } // Below capacity → grow lazily. Count the slot while still // holding the lock so concurrent acquires cannot overshoot. @@ -233,29 +261,57 @@ impl SessionPool { self.live.fetch_add(1, Ordering::Relaxed); drop(state); match self.open_session() { - Ok(session) => return self.hand_out(session), + Ok(session) => match self.hand_out(session) { + Ok(guard) => return Ok(guard), + Err(e) => { + let state = self.state.lock().expect("pool lock poisoned"); + self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + drop(state); + // A broken USE replay is a terminal error for + // this acquire: report it directly. + return Err(e); + } + }, Err(e) => { + last_err = Some(e); + state = self.state.lock().expect("pool lock poisoned"); self.live.fetch_sub(1, Ordering::Relaxed); self.available.notify_one(); - return Err(e); + // Spend the remaining acquire_timeout budget waiting + // for a release instead of failing immediately; the + // loop re-checks closed/idle/deadline at the top. + state = self.wait_for_change(state, deadline); } } + continue; } // At capacity → wait for a release, bounded by the deadline. - let now = Instant::now(); - if now >= deadline { - return Err(Error::Client(format!( - "pool exhausted: no session available within {:?} ({} live, max {})", - self.config.acquire_timeout, - self.live.load(Ordering::Relaxed), - self.config.max_size - ))); + state = self.wait_for_change(state, deadline); + } + } + + /// Park on the availability Condvar until a release or (when a deadline + /// is set) the deadline passes. The returned guard still holds the state + /// lock so the caller re-evaluates the predicate safely. + fn wait_for_change<'a>( + &self, + state: std::sync::MutexGuard<'a, PoolState>, + deadline: Option, + ) -> std::sync::MutexGuard<'a, PoolState> { + match deadline { + Some(deadline) => { + let now = Instant::now(); + if now >= deadline { + return state; + } + let (guard, _) = self + .available + .wait_timeout(state, deadline - now) + .expect("pool lock poisoned"); + guard } - let (guard, _) = self - .available - .wait_timeout(state, deadline - now) - .expect("pool lock poisoned"); - state = guard; + None => self.available.wait(state).expect("pool lock poisoned"), } } @@ -271,24 +327,45 @@ impl SessionPool { /// 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> { - { + let matched = { let mut state = self.state.lock().expect("pool lock poisoned"); - if !state.closed { + if state.closed { + None + } else { // Any idle session may hold the hint (the one that got the // 400), not necessarily one connected to the hinted node. + // When several sessions hold conflicting hints, prefer the + // newest (highest cache seq). let hint = state .idle .iter_mut() - .find_map(|e| e.session.redirect_hint(device_id)); - if let Some(endpoint) = hint { - let matching = state.idle.iter().position(|e| { - e.session.is_open() && e.session.current_endpoint() == Some(&endpoint) - }); - if let Some(pos) = matching { - let entry = state.idle.remove(pos).expect("index in bounds"); - drop(state); - return self.hand_out(entry.session); - } + .filter_map(|e| e.session.redirect_hint_with_seq(device_id)) + .max_by_key(|(_, seq)| *seq); + match hint { + Some((endpoint, _)) => state + .idle + .iter() + .position(|e| { + e.session.is_open() + && e.session + .current_endpoint() + .is_some_and(|current| current.equivalent(&endpoint)) + }) + .map(|pos| state.idle.remove(pos).expect("index in bounds").session), + None => None, + } + } + }; + if let Some(session) = matched { + match self.hand_out(session) { + Ok(guard) => return Ok(guard), + Err(_) => { + // Account under the lock, then fall back to the normal + // acquire path (the hinted hand-out may still succeed). + let state = self.state.lock().expect("pool lock poisoned"); + self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + drop(state); } } } @@ -305,16 +382,20 @@ impl SessionPool { /// sessions. Sessions currently handed out are closed when their guards /// drop. pub fn close(&self) { + // Set the flag, decrement `live` and wake every waiter under the + // same lock the waiters re-check the predicate under (no lost + // Condvar wakeups); the blocking closeSession RPCs run after the + // lock is released, bounded by the session's socket_timeout. let drained = { let mut state = self.state.lock().expect("pool lock poisoned"); state.closed = true; + self.live.fetch_sub(state.idle.len(), Ordering::Relaxed); + self.available.notify_all(); std::mem::take(&mut state.idle) }; - self.live.fetch_sub(drained.len(), Ordering::Relaxed); for mut entry in drained { let _ = entry.session.close(); } - self.available.notify_all(); } fn open_session(&self) -> Result { @@ -329,19 +410,22 @@ impl SessionPool { /// Final step of acquire: sync the session onto the pool's current /// database before handing it out. On USE failure the session is - /// discarded, not returned to the pool. + /// discarded, not returned to the pool; the caller re-takes the state + /// lock to decrement `live` and notify, keeping the decrement and the + /// predicate waiters re-evaluate under the same lock. fn hand_out(&self, mut session: Session) -> Result> { let pool_db = self.database.lock().expect("pool lock poisoned").clone(); if let Some(db) = pool_db { if session.database() != Some(db.as_str()) { if let Err(e) = session.execute_non_query(&format!("USE {db}")) { let _ = session.close(); - self.live.fetch_sub(1, Ordering::Relaxed); - self.available.notify_one(); return Err(e); } } } + // Pooled sessions skip the reconnect pacing sleeps so a pool slot is + // never held for the full C#-style reconnect walk. + session.mark_pooled(); Ok(PooledSession { pool: self, session: Some(session), @@ -352,8 +436,10 @@ impl SessionPool { /// live ones update the pool database and go back to the idle queue. fn release(&self, session: Session) { if !session.is_open() { + let state = self.state.lock().expect("pool lock poisoned"); self.live.fetch_sub(1, Ordering::Relaxed); self.available.notify_one(); + drop(state); return; } if let Some(db) = session.database() { @@ -364,19 +450,20 @@ impl SessionPool { } let mut state = self.state.lock().expect("pool lock poisoned"); if state.closed { - drop(state); self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + drop(state); let mut session = session; let _ = session.close(); } else { state.idle.push_back(IdleEntry::new(session)); let expired = self.sweep_idle(&mut state); + self.available.notify_one(); drop(state); for mut session in expired { let _ = session.close(); } } - self.available.notify_one(); } /// Test hook: push a pre-built session (possibly dead) into the idle @@ -454,6 +541,14 @@ impl TableSessionPool { self.pool.acquire() } + /// Acquire a table-dialect session for writes to `device_id`, + /// preferring an idle session already connected to the device's + /// redirected endpoint — the table-model counterpart of + /// [`SessionPool::acquire_for_device`]. + pub fn acquire_for_device(&self, device_id: &str) -> Result> { + self.pool.acquire_for_device(device_id) + } + /// Convenience: acquire, run one non-query statement, release. pub fn execute_non_query(&self, sql: &str) -> Result<()> { self.pool.execute_non_query(sql) @@ -505,6 +600,91 @@ mod tests { assert_eq!(cfg.idle_sweep_interval, Duration::from_secs(30)); } + /// F5: a failed growth attempt must spend the acquire_timeout budget + /// (waiting for a release) instead of failing instantly — the old + /// behaviour returned the connect error in ~200us against a 300ms budget. + #[test] + fn growth_failure_spends_acquire_timeout_budget() { + let cfg = SessionPoolConfig { + max_size: 1, + acquire_timeout: Duration::from_millis(300), + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + let started = Instant::now(); + match pool.acquire() { + Err(Error::Thrift(_)) => {} + other => panic!("expected thrift connect error, got {other:?}"), + } + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_millis(250), + "acquire failed in {elapsed:?} instead of spending the budget" + ); + assert_eq!(pool.live_count(), 0); + } + + /// F6: waiters blocked on a full pool must be woken by `close()` + /// promptly — the wakeup cannot depend on the blocking closeSession RPCs + /// that run afterwards (the old code notified only after them). + #[test] + fn close_wakes_waiters_promptly() { + let ep = fake_listener(); + let cfg = SessionPoolConfig { + max_size: 1, + acquire_timeout: Duration::from_secs(10), + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = std::sync::Arc::new(SessionPool::new(cfg).unwrap()); + pool.inject_idle(injected_session(&ep)); + let _held = pool.acquire().unwrap(); + + let (tx, rx) = std::sync::mpsc::channel(); + let waiter_pool = std::sync::Arc::clone(&pool); + let waiter = std::thread::spawn(move || { + let _ = tx.send(waiter_pool.acquire().map(|_| ())); + }); + std::thread::sleep(Duration::from_millis(100)); + pool.close(); + match rx.recv_timeout(Duration::from_secs(2)) { + Ok(Err(Error::Client(msg))) => assert!(msg.contains("closed"), "{msg}"), + other => panic!("waiter should get the closed-pool error promptly, got {other:?}"), + } + waiter.join().expect("waiter thread"); + } + + /// F13: `acquire_timeout = Duration::MAX` must mean "wait without a + /// deadline", not panic on `Instant + Duration` overflow. The close + /// below is what releases the waiter. + #[test] + fn acquire_timeout_max_waits_until_closed_without_panicking() { + let ep = fake_listener(); + let cfg = SessionPoolConfig { + max_size: 1, + acquire_timeout: Duration::MAX, + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = std::sync::Arc::new(SessionPool::new(cfg).unwrap()); + pool.inject_idle(injected_session(&ep)); + let _held = pool.acquire().unwrap(); + + let (tx, rx) = std::sync::mpsc::channel(); + let waiter_pool = std::sync::Arc::clone(&pool); + let waiter = std::thread::spawn(move || { + let _ = tx.send(waiter_pool.acquire().map(|_| ())); + }); + std::thread::sleep(Duration::from_millis(100)); + pool.close(); + match rx.recv_timeout(Duration::from_secs(2)) { + Ok(Err(Error::Client(msg))) => assert!(msg.contains("closed"), "{msg}"), + other => panic!("Duration::MAX waiter should wake on close, got {other:?}"), + } + waiter.join().expect("waiter thread"); + } + #[test] fn min_greater_than_max_is_rejected() { let cfg = SessionPoolConfig { @@ -665,6 +845,41 @@ mod tests { assert_eq!(guard.current_endpoint(), Some(&ep_a)); } + /// F12: when idle sessions hold conflicting hints for one device, the + /// **newest** hint wins, not the session nearest the queue head. + #[test] + fn acquire_for_device_prefers_newest_conflicting_hint() { + let ep_x = fake_listener(); + let ep_y = fake_listener(); + let ep_b = fake_listener(); + let ep_c = 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(); + + // Older hint (inserted first, process-wide seq 1) points at B. + let mut s_old_hint = injected_session(&ep_x); + s_old_hint.test_inject_redirect_hint("root.sg.d1", ep_b.clone()); + // Newer hint (seq 2) points at C. Both hint-holders precede the + // hinted sessions in the idle queue, so queue position would pick B. + let mut s_new_hint = injected_session(&ep_y); + s_new_hint.test_inject_redirect_hint("root.sg.d1", ep_c.clone()); + let s_on_b = injected_session(&ep_b); + let s_on_c = injected_session(&ep_c); + pool.inject_idle(s_old_hint); + pool.inject_idle(s_new_hint); + pool.inject_idle(s_on_b); + pool.inject_idle(s_on_c); + + let guard = pool.acquire_for_device("root.sg.d1").unwrap(); + assert_eq!(guard.current_endpoint(), Some(&ep_c)); + } + #[test] fn acquire_for_device_falls_back_when_no_session_matches_hint() { let ep_a = fake_listener(); diff --git a/src/client/redirect.rs b/src/client/redirect.rs index 4505504..7e52a72 100644 --- a/src/client/redirect.rs +++ b/src/client/redirect.rs @@ -30,6 +30,7 @@ //! future work. use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use crate::connection::Endpoint; @@ -40,6 +41,11 @@ 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; +/// Process-wide insertion counter: `seq` values stay comparable across +/// per-session caches, so a pool can tell which of two conflicting hints +/// for one device is newer. +static NEXT_SEQ: AtomicU64 = AtomicU64::new(0); + /// 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). @@ -68,7 +74,6 @@ pub struct RedirectCache { entries: HashMap, ttl: Duration, max_entries: usize, - seq: u64, } impl Default for RedirectCache { @@ -85,7 +90,6 @@ impl RedirectCache { entries: HashMap::new(), ttl, max_entries, - seq: 0, } } @@ -95,15 +99,27 @@ impl RedirectCache { self.get_at(device_id, Instant::now()) } + /// Like [`RedirectCache::get`], plus the insertion sequence — higher + /// is newer. Pools use it to prefer the newest hint when idle sessions + /// hold conflicting hints for one device. + pub fn get_with_seq(&mut self, device_id: &str) -> Option<(Endpoint, u64)> { + self.get_at_with_seq(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 { + self.get_at_with_seq(device_id, now) + .map(|(endpoint, _)| endpoint) + } + + fn get_at_with_seq(&mut self, device_id: &str, now: Instant) -> Option<(Endpoint, u64)> { 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()) + Some((entry.endpoint.clone(), entry.seq)) } /// Record (or refresh) the hint for `device_id`. When the cache is full @@ -123,13 +139,13 @@ impl RedirectCache { self.entries.remove(&oldest); } } - self.seq += 1; + let seq = NEXT_SEQ.fetch_add(1, Ordering::Relaxed); self.entries.insert( device_id, Entry { endpoint, inserted: Instant::now(), - seq: self.seq, + seq, }, ); } diff --git a/src/client/session.rs b/src/client/session.rs index 91480c4..153e78a 100644 --- a/src/client/session.rs +++ b/src/client/session.rs @@ -56,10 +56,21 @@ pub struct SessionConfig { pub sql_dialect: String, pub fetch_size: i32, pub zone_id: String, - /// TCP connect timeout per endpoint attempt. + /// Total TCP connect timeout per endpoint attempt. pub connect_timeout: Duration, + /// Client-side bound on each blocking socket read/write + /// (SO_RCVTIMEO/SO_SNDTIMEO), applied after the TCP connect and before + /// the TLS handshake: it bounds the handshake, every RPC read and the + /// best-effort drop-time `closeSession`. `None` restores unbounded + /// blocking. Default 60 s. + pub socket_timeout: Option, /// Per-query server-side timeout in milliseconds (request body field). pub query_timeout_ms: i64, + /// Redirect-hint cache TTL. Default 300 s (Node.js `RedirectCache`). + pub redirect_cache_ttl: Duration, + /// Redirect-hint cache capacity (oldest entry evicted when full); + /// `0` disables the cache. Default 1024. + pub redirect_cache_max_entries: usize, /// Database to select at open time (table dialect; sent as config key `db`). pub database: Option, /// Reopen the connection and retry an op once when an RPC fails at the @@ -116,7 +127,10 @@ impl Default for SessionConfig { fetch_size: 1024, zone_id: "UTC+8".into(), connect_timeout: Duration::from_secs(10), + socket_timeout: Some(Duration::from_secs(60)), query_timeout_ms: 60_000, + redirect_cache_ttl: redirect::DEFAULT_REDIRECT_TTL, + redirect_cache_max_entries: redirect::DEFAULT_REDIRECT_MAX_ENTRIES, database: None, enable_auto_reconnect: true, max_reconnect_attempts: 3, @@ -155,6 +169,7 @@ impl SessionConfig { pub fn connection_options(&self) -> ConnectionOptions { ConnectionOptions { connect_timeout: self.connect_timeout, + socket_timeout: self.socket_timeout, protocol: if self.enable_rpc_compression { RpcProtocol::Compact } else { @@ -204,11 +219,21 @@ pub struct Session { last_endpoint: Option, /// Device → endpoint hints harvested from status-400 insert responses. redirect_cache: RedirectCache, + /// A transport-level RPC failure was observed and the connection is + /// presumed desynchronized; with auto-reconnect off there is no reopen + /// path, so `is_open` reports false and pools discard the session. + broken: bool, + /// This session is handed out from a pool: reconnect skips the + /// between-attempt sleeps so a pool slot is not held for the full + /// C#-style pacing (pool.rs). + pooled: bool, } impl Session { pub fn new(config: SessionConfig) -> Self { let database = config.database.clone(); + let redirect_cache = + RedirectCache::new(config.redirect_cache_ttl, config.redirect_cache_max_entries); Self { config, connection: None, @@ -216,7 +241,9 @@ impl Session { statement_id: -1, database, last_endpoint: None, - redirect_cache: RedirectCache::default(), + redirect_cache, + broken: false, + pooled: false, } } @@ -233,28 +260,33 @@ impl Session { let start = ENDPOINT_START_INDEX.fetch_add(1, Ordering::Relaxed) % n; let options = self.config.connection_options(); - let mut connection = None; let mut last_err: Option = None; + // Fail over at the *full* handshake level, not just TCP: a node that + // accepts the connection but rejects openSession must not fail the + // whole open while other nodes are healthy (reconnect() has always + // retried connect + authenticate together). for i in 0..n { let endpoint = self.config.endpoints[(start + i) % n].clone(); - match Connection::open(endpoint, &options) { - Ok(c) => { - connection = Some(c); - break; + let result = Connection::open(endpoint.clone(), &options).and_then(|mut connection| { + let ids = self.authenticate(&mut connection)?; + Ok((connection, ids)) + }); + match result { + Ok((connection, (session_id, statement_id))) => { + self.session_id = session_id; + self.statement_id = statement_id; + self.last_endpoint = Some(connection.endpoint().clone()); + self.connection = Some(connection); + self.broken = false; + return Ok(()); + } + Err(e) => { + log::warn!("open against {endpoint} failed: {e}"); + last_err = Some(e); } - Err(e) => last_err = Some(e), } } - let mut connection = connection.ok_or_else(|| { - 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(()) + Err(last_err.unwrap_or_else(|| Error::Client("no endpoints configured".into()))) } /// Handshake on a fresh connection: `openSession` (dialect + current @@ -298,6 +330,7 @@ impl Session { /// (mirroring the C# SDK's `Reconnect`). fn reconnect(&mut self) -> Result<()> { self.connection = None; // drop closes the old transport + self.broken = false; let n = self.config.endpoints.len(); if n == 0 { return Err(Error::Client("no endpoints configured".into())); @@ -311,7 +344,10 @@ impl Session { let options = self.config.connection_options(); let mut last_err: Option = None; for attempt in 0..attempts { - if attempt > 0 { + // Pooled sessions hold their pool slot during reconnect; skip + // the pacing sleeps so a full reconnect walk cannot monopolize + // the slot for tens of seconds. + if attempt > 0 && !self.pooled { std::thread::sleep(self.config.retry_interval); } for i in 0..n { @@ -327,6 +363,7 @@ impl Session { self.statement_id = statement_id; self.last_endpoint = Some(connection.endpoint().clone()); self.connection = Some(connection); + self.broken = false; return Ok(()); } Err(e) => { @@ -353,11 +390,26 @@ impl Session { { e } + Err(e @ Error::Thrift(_)) => { + // No reconnect path: the connection may be desynchronized + // (a socket timeout abandons a half-read frame; a + // frame-too-large rejection abandons an undrained body). + // Mark it broken so is_open() turns false and pools + // discard the session instead of handing it out again. + self.broken = true; + return Err(e); + } other => return other, }; log::warn!("RPC failed at transport level ({original}); reconnecting"); match self.reconnect() { - Ok(()) => op(self), + Ok(()) => match op(self) { + Err(e @ Error::Thrift(_)) => { + self.broken = true; + Err(e) + } + other => other, + }, Err(reconnect_err) => { log::warn!("reconnect failed ({reconnect_err}); surfacing the original error"); Err(original) @@ -366,7 +418,7 @@ impl Session { } pub fn is_open(&self) -> bool { - self.connection.is_some() + self.connection.is_some() && !self.broken } /// The database currently selected on this session, if any. @@ -379,6 +431,12 @@ impl Session { self.connection.as_ref().map(Connection::endpoint) } + /// Mark this session as pool-owned: reconnect skips the pacing sleeps + /// so a pool slot is not held for the full C#-style reconnect walk. + pub(crate) fn mark_pooled(&mut self) { + self.pooled = true; + } + /// The cached redirect endpoint for `device_id`, if a status-400 insert /// response recommended one and the hint has not expired. /// @@ -390,6 +448,13 @@ impl Session { self.redirect_cache.get(device_id) } + /// Like [`Session::redirect_hint`], but also returns the cache's + /// insertion sequence so pools can prefer the newest hint when idle + /// sessions hold conflicting hints for one device. + pub(crate) fn redirect_hint_with_seq(&mut self, device_id: &str) -> Option<(Endpoint, u64)> { + self.redirect_cache.get_with_seq(device_id) + } + /// Occupancy/config snapshot of the redirect cache. pub fn redirect_cache_stats(&self) -> RedirectCacheStats { self.redirect_cache.stats() @@ -484,7 +549,17 @@ impl Session { self.config.query_timeout_ms, self.statement_id, ); - let resp = self.connection_mut()?.client_mut().fetch_results_v2(req)?; + let resp = match self.connection_mut()?.client_mut().fetch_results_v2(req) { + Ok(resp) => resp, + Err(e) => { + // fetch_results deliberately bypasses with_retry (spec + // gotcha #13: the result set is pinned to its node), so a + // transport-level failure here must still mark the + // connection broken for the pool. + self.broken = true; + return Err(Error::from(e)); + } + }; check_status(&resp.status)?; if !resp.has_result_set { return Ok((Vec::new(), false)); @@ -1020,6 +1095,197 @@ mod tests { TSStatus::new(code, None, None, None, None, None) } + /// A fake IoTDB handshake server for testing `Session::open` failover + /// without a live server: real Thrift framed messages, where the handler + /// can reject `openSession` on demand. + mod fake_auth_server { + use crate::connection::Endpoint; + use crate::protocol::client::*; + use crate::protocol::common; + use crate::protocol::common::TSStatus; + use std::collections::BTreeMap; + use std::net::TcpListener; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use thrift::protocol::{TBinaryInputProtocol, TBinaryOutputProtocol}; + use thrift::server::TProcessor; + use thrift::transport::{ + TFramedReadTransport, TFramedWriteTransport, TIoChannel, TTcpChannel, + }; + + struct FakeAuthHandler { + fail_open_session: bool, + } + + fn ok_status() -> TSStatus { + TSStatus::new(200, None, None, None, None, None) + } + + fn reject_status() -> TSStatus { + TSStatus::new( + 1000, + Some("openSession rejected by test".into()), + None, + None, + None, + None, + ) + } + + /// Compact impl for the handler methods the fake server never + /// dispatches: a declarative macro keeps this test module small. + macro_rules! fake_unimplemented_handlers { + ($( $name:ident ( $( $arg:ident : $ty:ty ),* ) : $ret:ty ; )*) => { + $( + fn $name(&self $(, $arg: $ty)*) -> thrift::Result<$ret> { + #[allow(unused_parens)] + let _ = ($( $arg ),*); + unimplemented!( + "{} is not exercised by the fake handshake listener", + stringify!($name) + ) + } + )* + }; + } + + impl IClientRPCServiceSyncHandler for FakeAuthHandler { + fn handle_open_session( + &self, + _req: TSOpenSessionReq, + ) -> thrift::Result { + Ok(TSOpenSessionResp::new( + if self.fail_open_session { + reject_status() + } else { + ok_status() + }, + TSProtocolVersion::IotdbServiceProtocolV3, + if self.fail_open_session { + None:: + } else { + Some(1_i64) + }, + None::>, + )) + } + + fn handle_request_statement_id(&self, _session_id: i64) -> thrift::Result { + Ok(1) + } + + fn handle_close_session(&self, _req: TSCloseSessionReq) -> thrift::Result { + Ok(ok_status()) + } + + fake_unimplemented_handlers! { + handle_execute_query_statement_v2(req: TSExecuteStatementReq): TSExecuteStatementResp; + handle_execute_update_statement_v2(req: TSExecuteStatementReq): TSExecuteStatementResp; + handle_execute_statement_v2(req: TSExecuteStatementReq): TSExecuteStatementResp; + handle_execute_raw_data_query_v2(req: TSRawDataQueryReq): TSExecuteStatementResp; + handle_execute_last_data_query_v2(req: TSLastDataQueryReq): TSExecuteStatementResp; + handle_execute_fast_last_data_query_for_one_prefix_path(req: TSFastLastDataQueryForOnePrefixPathReq): TSExecuteStatementResp; + handle_execute_fast_last_data_query_for_one_device_v2(req: TSFastLastDataQueryForOneDeviceReq): TSExecuteStatementResp; + handle_execute_aggregation_query_v2(req: TSAggregationQueryReq): TSExecuteStatementResp; + handle_fetch_results_v2(req: TSFetchResultsReq): TSFetchResultsResp; + handle_execute_statement(req: TSExecuteStatementReq): TSExecuteStatementResp; + handle_execute_batch_statement(req: TSExecuteBatchStatementReq): common::TSStatus; + handle_execute_query_statement(req: TSExecuteStatementReq): TSExecuteStatementResp; + handle_execute_update_statement(req: TSExecuteStatementReq): TSExecuteStatementResp; + handle_fetch_results(req: TSFetchResultsReq): TSFetchResultsResp; + handle_fetch_metadata(req: TSFetchMetadataReq): TSFetchMetadataResp; + handle_cancel_operation(req: TSCancelOperationReq): common::TSStatus; + handle_close_operation(req: TSCloseOperationReq): common::TSStatus; + handle_prepare_statement(req: TSPrepareReq): TSPrepareResp; + handle_execute_prepared_statement(req: TSExecutePreparedReq): TSExecuteStatementResp; + handle_deallocate_prepared_statement(req: TSDeallocatePreparedReq): common::TSStatus; + handle_get_time_zone(session_id: i64): TSGetTimeZoneResp; + handle_set_time_zone(req: TSSetTimeZoneReq): common::TSStatus; + handle_get_properties(): ServerProperties; + handle_set_storage_group(session_id: i64, storage_group: String): common::TSStatus; + handle_create_timeseries(req: TSCreateTimeseriesReq): common::TSStatus; + handle_create_aligned_timeseries(req: TSCreateAlignedTimeseriesReq): common::TSStatus; + handle_create_multi_timeseries(req: TSCreateMultiTimeseriesReq): common::TSStatus; + handle_delete_timeseries(session_id: i64, path: Vec): common::TSStatus; + handle_delete_storage_groups(session_id: i64, storage_group: Vec): common::TSStatus; + handle_insert_record(req: TSInsertRecordReq): common::TSStatus; + handle_insert_string_record(req: TSInsertStringRecordReq): common::TSStatus; + handle_insert_tablet(req: TSInsertTabletReq): common::TSStatus; + handle_insert_tablets(req: TSInsertTabletsReq): common::TSStatus; + handle_insert_records(req: TSInsertRecordsReq): common::TSStatus; + handle_insert_records_of_one_device(req: TSInsertRecordsOfOneDeviceReq): common::TSStatus; + handle_insert_string_records_of_one_device(req: TSInsertStringRecordsOfOneDeviceReq): common::TSStatus; + handle_insert_string_records(req: TSInsertStringRecordsReq): common::TSStatus; + handle_test_insert_tablet(req: TSInsertTabletReq): common::TSStatus; + handle_test_insert_tablets(req: TSInsertTabletsReq): common::TSStatus; + handle_test_insert_record(req: TSInsertRecordReq): common::TSStatus; + handle_test_insert_string_record(req: TSInsertStringRecordReq): common::TSStatus; + handle_test_insert_records(req: TSInsertRecordsReq): common::TSStatus; + handle_test_insert_records_of_one_device(req: TSInsertRecordsOfOneDeviceReq): common::TSStatus; + handle_test_insert_string_records(req: TSInsertStringRecordsReq): common::TSStatus; + handle_delete_data(req: TSDeleteDataReq): common::TSStatus; + handle_execute_raw_data_query(req: TSRawDataQueryReq): TSExecuteStatementResp; + handle_execute_last_data_query(req: TSLastDataQueryReq): TSExecuteStatementResp; + handle_execute_aggregation_query(req: TSAggregationQueryReq): TSExecuteStatementResp; + handle_create_schema_template(req: TSCreateSchemaTemplateReq): common::TSStatus; + handle_append_schema_template(req: TSAppendSchemaTemplateReq): common::TSStatus; + handle_prune_schema_template(req: TSPruneSchemaTemplateReq): common::TSStatus; + handle_query_schema_template(req: TSQueryTemplateReq): TSQueryTemplateResp; + handle_show_configuration_template(): common::TShowConfigurationTemplateResp; + handle_show_configuration(node_id: i32): common::TShowConfigurationResp; + handle_set_schema_template(req: TSSetSchemaTemplateReq): common::TSStatus; + handle_unset_schema_template(req: TSUnsetSchemaTemplateReq): common::TSStatus; + handle_drop_schema_template(req: TSDropSchemaTemplateReq): common::TSStatus; + handle_create_timeseries_using_schema_template(req: TCreateTimeseriesUsingSchemaTemplateReq): common::TSStatus; + handle_handshake(info: TSyncIdentityInfo): common::TSStatus; + handle_send_pipe_data(buff: Vec): common::TSStatus; + handle_send_file(meta_info: TSyncTransportMetaInfo, buff: Vec): common::TSStatus; + handle_pipe_transfer(req: TPipeTransferReq): TPipeTransferResp; + handle_pipe_subscribe(req: TPipeSubscribeReq): TPipeSubscribeResp; + handle_get_backup_configuration(): TSBackupConfigurationResp; + handle_fetch_all_connections_info(): TSConnectionInfoResp; + handle_test_connection_empty_r_p_c(): common::TSStatus; + } + } + /// A listener that answers `openSession` + `requestStatementId` + /// (and `closeSession`) with real framed Thrift messages. When + /// `fail_first_connection` is set, the **first** connection gets a + /// rejected `openSession`; every later connection succeeds. The + /// acceptor thread is leaked; it ends with the test process. + pub fn auth_listener(fail_first_connection: bool) -> Endpoint { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("local_addr").port(); + let connections = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&connections); + std::thread::spawn(move || { + for stream in listener.incoming() { + let Ok(stream) = stream else { break }; + let index = counter.fetch_add(1, Ordering::Relaxed); + let fail_this = fail_first_connection && index == 0; + let channel = TTcpChannel::with_stream(stream); + let (read_half, write_half) = channel.split().expect("split channel"); + let mut i_prot = + TBinaryInputProtocol::new(TFramedReadTransport::new(read_half), true); + let mut o_prot = + TBinaryOutputProtocol::new(TFramedWriteTransport::new(write_half), true); + let processor = IClientRPCServiceSyncProcessor::new(FakeAuthHandler { + fail_open_session: fail_this, + }); + // authenticate() = openSession (+ requestStatementId on + // success); closeSession is handled the same way when the + // test closes the session. + if processor.process(&mut i_prot, &mut o_prot).is_err() { + break; + } + if !fail_this && processor.process(&mut i_prot, &mut o_prot).is_err() { + break; + } + } + }); + Endpoint::new("127.0.0.1", port) + } + } + #[test] fn default_config() { let cfg = SessionConfig::default(); @@ -1030,6 +1296,9 @@ mod tests { assert_eq!(cfg.fetch_size, 1024); assert_eq!(cfg.query_timeout_ms, 60_000); assert_eq!(cfg.connect_timeout, Duration::from_secs(10)); + assert_eq!(cfg.socket_timeout, Some(Duration::from_secs(60))); + assert_eq!(cfg.redirect_cache_ttl, Duration::from_secs(300)); + assert_eq!(cfg.redirect_cache_max_entries, 1024); assert!(cfg.database.is_none()); assert!(cfg.enable_auto_reconnect); assert_eq!(cfg.max_reconnect_attempts, 3); @@ -1056,6 +1325,7 @@ mod tests { let cfg = SessionConfig::default(); let options = cfg.connection_options(); assert_eq!(options.connect_timeout, cfg.connect_timeout); + assert_eq!(options.socket_timeout, cfg.socket_timeout); assert_eq!(options.protocol, RpcProtocol::Binary); #[cfg(feature = "tls")] assert!(options.tls.is_none()); @@ -1487,6 +1757,30 @@ mod tests { 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"); + // F7: without a reconnect path the desynchronized connection is + // marked broken, so is_open() stops lying and pools discard the + // session instead of handing it out again. + assert!(!session.is_open()); + } + + /// F4 regression: failover covers the **handshake**, not just the TCP + /// connect. One fake server backs both endpoint entries and rejects + /// `openSession` on its first connection: whichever entry is tried + /// first fails, and the second entry must complete the handshake. + #[test] + fn open_fails_over_to_next_endpoint_when_authentication_fails() { + let endpoint = fake_auth_server::auth_listener(true); + let mut session = Session::new(SessionConfig { + endpoints: vec![endpoint.clone(), endpoint.clone()], + connect_timeout: Duration::from_millis(500), + socket_timeout: Some(Duration::from_millis(500)), + ..Default::default() + }); + session + .open() + .expect("the second endpoint entry must take over after the rejected openSession"); + assert_eq!(session.current_endpoint(), Some(&endpoint)); + let _ = session.close(); } #[test] diff --git a/src/client/table_session.rs b/src/client/table_session.rs index b708e26..02fb409 100644 --- a/src/client/table_session.rs +++ b/src/client/table_session.rs @@ -102,6 +102,23 @@ impl TableSessionBuilder { self } + /// Client-side bound on each socket read/write (SO_RCVTIMEO/SO_SNDTIMEO), + /// applied after the TCP connect and before the TLS handshake. + /// See [`SessionConfig::socket_timeout`]. + pub fn socket_timeout(mut self, timeout: Duration) -> Self { + self.config.socket_timeout = Some(timeout); + self + } + + /// Redirect-hint cache TTL and capacity. + /// See [`SessionConfig::redirect_cache_ttl`] / + /// [`SessionConfig::redirect_cache_max_entries`]. + pub fn redirect_cache(mut self, ttl: Duration, max_entries: usize) -> Self { + self.config.redirect_cache_ttl = ttl; + self.config.redirect_cache_max_entries = max_entries; + self + } + pub fn query_timeout_ms(mut self, timeout_ms: i64) -> Self { self.config.query_timeout_ms = timeout_ms; self diff --git a/src/connection/mod.rs b/src/connection/mod.rs index 7bcdf81..ed1406f 100644 --- a/src/connection/mod.rs +++ b/src/connection/mod.rs @@ -26,7 +26,9 @@ use std::io::{Read, Write}; use std::net::{TcpStream, ToSocketAddrs}; #[cfg(feature = "tls")] use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; + +use socket2::SockRef; #[cfg(feature = "tls")] use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier}; @@ -97,11 +99,19 @@ pub struct TlsOptions { pub client_key_path: Option, } -/// How to open a [`Connection`]: timeout, wire protocol, optional TLS. +/// How to open a [`Connection`]: timeouts, wire protocol, optional TLS. #[derive(Debug, Clone)] pub struct ConnectionOptions { - /// TCP connect timeout per endpoint attempt. Default 10 s. + /// Total TCP connect timeout per endpoint attempt (shared across every + /// resolved address of that endpoint). Default 10 s. pub connect_timeout: Duration, + /// Client-side bound on each blocking socket read/write (SO_RCVTIMEO / + /// SO_SNDTIMEO), applied after the TCP connect and **before** the TLS + /// handshake, so it bounds the handshake, every RPC read and the + /// best-effort drop-time `closeSession` alike. `None` restores the + /// old unbounded blocking behaviour. Default 60 s (matches the default + /// server-side `query_timeout_ms`). + pub socket_timeout: Option, /// Wire protocol; must match the server (see [`RpcProtocol`]). pub protocol: RpcProtocol, /// Wrap the TCP stream in TLS before the Thrift transports. @@ -113,6 +123,7 @@ impl Default for ConnectionOptions { fn default() -> Self { Self { connect_timeout: Duration::from_secs(10), + socket_timeout: Some(Duration::from_secs(60)), protocol: RpcProtocol::Binary, #[cfg(feature = "tls")] tls: None, @@ -167,6 +178,32 @@ impl Endpoint { } Ok(Self::new(host, port)) } + + /// Loose equality for redirect-hint matching: the port must match and + /// the host compares case-insensitively after trimming and stripping + /// IPv6 brackets; loopback spellings (`localhost`, `127.x.y.z`, + /// `::1`) are equivalent to each other. General hostname-vs-IP + /// resolution would need DNS and is deliberately not done here. + pub fn equivalent(&self, other: &Self) -> bool { + self.port == other.port + && (normalized_host(&self.host) == normalized_host(&other.host) + || (is_loopback_host(&self.host) && is_loopback_host(&other.host))) + } +} + +fn normalized_host(host: &str) -> String { + host.trim() + .trim_start_matches('[') + .trim_end_matches(']') + .trim_end_matches('.') + .to_ascii_lowercase() +} + +fn is_loopback_host(host: &str) -> bool { + host == "localhost" + || host + .parse::() + .is_ok_and(|ip| ip.is_loopback()) } impl std::fmt::Display for Endpoint { @@ -189,15 +226,15 @@ pub struct Connection { impl Connection { /// Establish a TCP connection to `endpoint` (bounded by - /// `options.connect_timeout`), optionally wrap it in TLS, and stack - /// framed transport + the selected protocol on top. + /// `options.connect_timeout`; every subsequent socket read/write is + /// bounded by `options.socket_timeout`), optionally wrap it in TLS, and + /// stack framed transport + the selected protocol on top. pub fn open(endpoint: Endpoint, options: &ConnectionOptions) -> Result { - let stream = connect_stream(&endpoint, options.connect_timeout)?; - stream.set_nodelay(true).map_err(thrift::Error::from)?; + let stream = connect_stream(&endpoint, options.connect_timeout, options.socket_timeout)?; #[cfg(feature = "tls")] if let Some(tls) = &options.tls { - let stream = tls_handshake(&endpoint, stream, tls)?; + let stream = tls_handshake(&endpoint, stream, tls, options.socket_timeout)?; let shared = SharedTlsStream::new(stream); let (input, output) = build_protocols(shared.clone(), shared, options.protocol); return Ok(Self { @@ -257,14 +294,38 @@ where } } -/// Resolve the endpoint and try each resolved address with the connect timeout. -fn connect_stream(endpoint: &Endpoint, connect_timeout: Duration) -> Result { +/// Resolve the endpoint and try each resolved address, sharing one total +/// `connect_timeout` budget across all of them (a multi-address hostname +/// must not multiply the configured bound), then apply `socket_timeout` +/// and TCP keepalive before handing the stream up. +fn connect_stream( + endpoint: &Endpoint, + connect_timeout: Duration, + socket_timeout: Option, +) -> Result { let addrs = (endpoint.host.as_str(), endpoint.port) .to_socket_addrs() .map_err(thrift::Error::from)?; + // A single deadline for the whole endpoint attempt. `checked_add` + // treats Duration::MAX as "no bound" instead of panicking on overflow. + let deadline = Instant::now().checked_add(connect_timeout); let mut last_err: Option = None; for addr in addrs { - match TcpStream::connect_timeout(&addr, connect_timeout) { + let attempt_timeout = match deadline { + Some(deadline) => { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + last_err = Some(std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!("connect to {endpoint} timed out"), + )); + break; + } + remaining + } + None => connect_timeout, + }; + match connect_one(addr, attempt_timeout, socket_timeout) { Ok(stream) => return Ok(stream), Err(e) => last_err = Some(e), } @@ -275,12 +336,40 @@ fn connect_stream(endpoint: &Endpoint, connect_timeout: Duration) -> Result, +) -> std::io::Result { + // std's connect_timeout keeps the established cross-platform connect + // semantics; SockRef then applies the options on the existing socket. + let stream = TcpStream::connect_timeout(&addr, connect_timeout)?; + let socket = SockRef::from(&stream); + socket.set_nodelay(true)?; + socket.set_keepalive(true)?; + // A zero duration would select non-blocking mode on some platforms; + // treat it as "no timeout" instead (None is the documented way). + if let Some(timeout) = socket_timeout.filter(|timeout| !timeout.is_zero()) { + socket.set_read_timeout(Some(timeout))?; + socket.set_write_timeout(Some(timeout))?; + } + Ok(stream) +} + +/// Run the TLS handshake over an established TCP stream. The stream +/// already carries the socket-level read/write timeout, so a peer that +/// accepts and then stalls the handshake fails here instead of blocking +/// forever. #[cfg(feature = "tls")] fn tls_handshake( endpoint: &Endpoint, mut stream: TcpStream, tls: &TlsOptions, + socket_timeout: Option, ) -> Result> { let config = tls_client_config(tls)?; let domain = tls.domain_override.as_deref().unwrap_or(&endpoint.host); @@ -289,9 +378,19 @@ fn tls_handshake( let mut connection = ClientConnection::new(config, server_name).map_err(|e| Error::Tls(e.to_string()))?; - connection - .complete_io(&mut stream) - .map_err(|e| Error::Tls(e.to_string()))?; + connection.complete_io(&mut stream).map_err(|e| { + // A blocking socket with SO_RCVTIMEO reports WouldBlock/TimedOut + // when the peer stalls the handshake: turn it into an actionable + // error instead of leaking the OS error kind. + if matches!( + e.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut + ) { + let bound = socket_timeout.map_or("no socket timeout".into(), |t| format!("{t:?}")); + return Error::Tls(format!("TLS handshake timed out ({bound})")); + } + Error::Tls(e.to_string()) + })?; if connection.is_handshaking() { return Err(Error::Tls("TLS handshake did not complete".into())); } @@ -547,15 +646,71 @@ mod tests { ); } + /// F11: redirect-hint matching must tolerate case/bracket/loopback + /// spelling differences, but still require the same port. + #[test] + fn endpoint_equivalent_normalizes_and_matches_loopback() { + assert!(Endpoint::new("LOCALHOST", 6667).equivalent(&Endpoint::new("localhost", 6667))); + assert!(Endpoint::new("localhost", 6667).equivalent(&Endpoint::new("127.0.0.1", 6667))); + assert!(Endpoint::new("127.0.0.1", 6667).equivalent(&Endpoint::new("::1", 6667))); + assert!(Endpoint::new("[::1]", 6667).equivalent(&Endpoint::new("::1", 6667))); + assert!(!Endpoint::new("localhost", 6667).equivalent(&Endpoint::new("localhost", 6668))); + assert!( + !Endpoint::new("localhost", 6667).equivalent(&Endpoint::new("iotdb.example.com", 6667)) + ); + } + #[test] fn default_options_are_binary_no_tls() { let options = ConnectionOptions::default(); assert_eq!(options.connect_timeout, Duration::from_secs(10)); + assert_eq!(options.socket_timeout, Some(Duration::from_secs(60))); assert_eq!(options.protocol, RpcProtocol::Binary); #[cfg(feature = "tls")] assert!(options.tls.is_none()); } + /// A local listener that accepts connections and then stays silent — + /// the equivalent of a peer that finished the TCP handshake and never + /// replies (GC pause, dropped firewall state, accepting LB). The accept + /// thread parks while holding the stream so it never sends EOF. + pub(super) fn silent_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 || { + let (_stream, _) = listener.accept().expect("accept"); + std::thread::park(); + }); + Endpoint::new("127.0.0.1", port) + } + + /// `socket_timeout` must bound reads after the TCP handshake: against + /// a peer that accepts and never replies, the RPC returns a Thrift + /// error around the configured bound instead of blocking forever. + #[test] + fn socket_timeout_bounds_reads_after_handshake() { + use crate::protocol::client::TIClientRPCServiceSyncClient; + + let endpoint = silent_listener(); + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + socket_timeout: Some(Duration::from_millis(300)), + ..Default::default() + }; + let mut connection = Connection::open(endpoint, &options).expect("TCP connect succeeds"); + let started = Instant::now(); + let err = connection + .client_mut() + .request_statement_id(1) + .expect_err("silent peer must not block forever"); + assert!(matches!(err, thrift::Error::Transport(_)), "got {err:?}"); + let elapsed = started.elapsed(); + assert!( + elapsed < Duration::from_secs(2), + "read took {elapsed:?}, socket timeout not applied" + ); + } + /// A local listener that accepts and immediately drops connections, so /// `Connection::open` (which issues no RPC) succeeds for any protocol. pub(super) fn accept_then_drop_listener() -> Endpoint { @@ -744,6 +899,7 @@ mod tls_tests { domain_override: Some("localhost".into()), ..Default::default() }), + ..Default::default() }; let mut connection = Connection::open(endpoint, &options).expect("TLS handshake"); assert_eq!(connection.protocol(), RpcProtocol::Binary); @@ -771,6 +927,7 @@ mod tls_tests { domain_override: Some("localhost".into()), ..Default::default() }), + ..Default::default() }; let err = match Connection::open(endpoint, &options) { Ok(_) => panic!("untrusted self-signed cert must fail the handshake"), @@ -791,6 +948,7 @@ mod tls_tests { accept_invalid_certs: true, ..Default::default() }), + ..Default::default() }; let connection = Connection::open(endpoint, &options).expect("TLS handshake"); assert_eq!(connection.protocol(), RpcProtocol::Compact); @@ -810,6 +968,7 @@ mod tls_tests { accept_invalid_certs: true, ..Default::default() }), + ..Default::default() }; assert!( Connection::open(endpoint, &options).is_err(), @@ -824,6 +983,34 @@ mod tls_tests { ); } + /// A socket timeout bounds the TLS handshake itself: against a peer + /// that accepts and then stays silent, `Connection::open` fails with a + /// TLS timeout error around the configured bound instead of blocking + /// forever. + #[test] + fn tls_handshake_times_out_against_silent_peer() { + let endpoint = super::tests::silent_listener(); + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + socket_timeout: Some(Duration::from_millis(300)), + protocol: RpcProtocol::Binary, + tls: Some(TlsOptions { + accept_invalid_certs: true, + ..Default::default() + }), + }; + let started = Instant::now(); + let err = match Connection::open(endpoint, &options) { + Ok(_) => panic!("silent peer must not complete a TLS handshake"), + Err(e) => e, + }; + assert!(matches!(err, Error::Tls(_)), "got {err:?}"); + assert!( + started.elapsed() < Duration::from_secs(2), + "TLS handshake was not bounded by socket_timeout" + ); + } + /// Dispatch pair against the *same kind* of plain (non-TLS) endpoint: /// `tls: None` opens fine (plain TCP), `tls: Some(..)` — even with /// certificate verification disabled — dies in the handshake with a @@ -838,6 +1025,7 @@ mod tls_tests { connect_timeout: Duration::from_millis(500), protocol: RpcProtocol::Binary, tls: None, + ..Default::default() }; Connection::open(endpoint.clone(), &plain).expect("plain open against plain listener"); @@ -848,6 +1036,7 @@ mod tls_tests { accept_invalid_certs: true, ..Default::default() }), + ..Default::default() }; let err = match Connection::open(endpoint, &tls) { Ok(_) => panic!("TLS handshake against a plain endpoint must fail"), @@ -872,6 +1061,7 @@ mod tls_tests { client_key_path: Some(fixture("client-key.pem")), ..Default::default() }), + ..Default::default() }; let connection = Connection::open(endpoint, &options).expect("TLS handshake with identity"); assert_eq!(connection.protocol(), RpcProtocol::Binary); @@ -899,6 +1089,7 @@ mod tls_tests { client_key_path: key.clone(), ..Default::default() }), + ..Default::default() }; let err = match Connection::open(endpoint.clone(), &options) { Ok(_) => panic!("half a client identity must fail"), @@ -924,6 +1115,7 @@ mod tls_tests { client_key_path: Some(fixture("does-not-exist-key.pem")), ..Default::default() }), + ..Default::default() }; let err = match Connection::open(endpoint, &options) { Ok(_) => panic!("missing client key must fail"), @@ -949,6 +1141,7 @@ mod tls_tests { client_key_path: Some(fixture("client-cert.pem")), // not a key ..Default::default() }), + ..Default::default() }; let err = match Connection::open(endpoint, &options) { Ok(_) => panic!("a certificate is not a private key"), @@ -968,6 +1161,7 @@ mod tls_tests { ca_cert_path: Some(fixture("does-not-exist.pem")), ..Default::default() }), + ..Default::default() }; let err = match Connection::open(endpoint, &options) { Ok(_) => panic!("missing CA file must fail"),