Re-establish clickhouse connection if it's idle for too long (#49)

Author: harshil-goel

src/ch_ddl.rs

@@ -28,6 +28,7 @@ use clickhouse_c::AsyncClient;
 use crate::ch_emitter::{
     ColumnMapping, EmitterConfig, EmitterError, MappingHandle, NamespaceMapping, RetryConfig,
     TableMapping, connect_client, drain_to_end_of_stream, is_retryable, quote_ident,
+    reconnect_if_idle,
 };
 use crate::shadow_catalog::{RelDescriptor, SchemaDiff, SchemaEvent};
 use crate::type_bridge::{self, ResolvedColumn};
@@ -132,6 +133,7 @@ pub struct DdlApplicator {
     /// Per-attempt cap (shares `EmitterConfig::insert_timeout`); a
     /// half-open CH socket can't park the reorder barrier past this
     query_timeout: Duration,
+    last_used: std::time::Instant,
     pub stats: DdlStats,
 }
 
@@ -160,6 +162,7 @@ impl DdlApplicator {
             conn_cfg: emitter_cfg.clone(),
             retry: emitter_cfg.retry.clone(),
             query_timeout: emitter_cfg.insert_timeout,
+            last_used: std::time::Instant::now(),
             stats: DdlStats::default(),
         })
     }
@@ -453,6 +456,7 @@ impl DdlApplicator {
         tracing::debug!(target: "walshadow::ch_ddl", sql = %sql, "applying");
         let mut attempt = 0u32;
         let mut backoff = self.retry.initial_backoff;
+        reconnect_if_idle(&mut self.client, &self.conn_cfg, self.last_used).await?;
         loop {
             let attempt_result = match tokio::time::timeout(self.query_timeout, async {
                 self.client.send_query(sql, None).await?;
@@ -468,7 +472,10 @@ impl DdlApplicator {
                 }),
             };
             match attempt_result {
-                Ok(()) => return Ok(()),
+                Ok(()) => {
+                    self.last_used = std::time::Instant::now();
+                    return Ok(());
+                }
                 Err(e) if is_retryable(&e) && attempt < self.retry.max_attempts => {
                     tracing::warn!(
                         target: "walshadow::ch_ddl",

src/ch_emitter.rs

@@ -139,6 +139,11 @@ pub struct EmitterConfig {
     /// so the inserter reconnects + resends rather than pinning the
     /// durable watermark forever. Sized far above a healthy round-trip.
     pub insert_timeout: Duration,
+    /// A CH connection idle longer than this may be half-open (NAT/LB/CH
+    /// idle-reap); reconnect before the next op instead of blocking the
+    /// full `insert_timeout` on a dead socket. Guards the start of a run,
+    /// when connections sat idle since the previous one.
+    pub idle_reconnect: Duration,
     /// Keep `_is_deleted` out of `ReplacingMergeTree`'s args so delete
     /// tombstones stay queryable instead of collapsing on FINAL. Column
     /// always emitted; off by default
@@ -153,6 +158,7 @@ pub struct EmitterConfig {
 }
 
 pub const DEFAULT_INSERT_TIMEOUT_SECS: u64 = 30;
+pub const DEFAULT_IDLE_RECONNECT_SECS: u64 = 30;
 
 /// Bounded-retry knobs. Retryable error (IO, clickhouse-c protocol,
 /// ServerException) triggers reconnect + retry up to `max_attempts`
@@ -193,6 +199,7 @@ impl Default for EmitterConfig {
             drop_table_strategy: "retain".into(),
             retry: RetryConfig::default(),
             insert_timeout: Duration::from_secs(DEFAULT_INSERT_TIMEOUT_SECS),
+            idle_reconnect: Duration::from_secs(DEFAULT_IDLE_RECONNECT_SECS),
             soft_delete: false,
             toast: crate::toast::ToastConfig::default(),
             decode_chunk_rows: crate::pipeline::decode::DECODE_CHUNK_ROWS,
@@ -629,6 +636,22 @@ pub(crate) async fn connect_client(config: &EmitterConfig) -> Result<AsyncClient
     Ok(client)
 }
 
+/// Reconnect `client` when idle since `last_used` exceeds `config.idle_reconnect`.
+/// A long-idle socket may be silently half-open; the first op on it would
+/// otherwise block the full `insert_timeout` before the reconnect-retry path
+/// notices. No-op during active streaming (connections used sub-second).
+pub(crate) async fn reconnect_if_idle(
+    client: &mut AsyncClient,
+    config: &EmitterConfig,
+    last_used: std::time::Instant,
+) -> Result<bool, EmitterError> {
+    if last_used.elapsed() >= config.idle_reconnect {
+        *client = connect_client(config).await?;
+        return Ok(true);
+    }
+    Ok(false)
+}
+
 /// Drain a CH response stream to `EndOfStream`, surfacing any
 /// `Exception` packet as [`EmitterError::ServerException`]. Used after a
 /// `send_query`/`send_data_end()` expecting no result rows (INSERT seal,

src/pipeline/inserter.rs

@@ -19,7 +19,7 @@ use tokio::task::JoinHandle;
 
 use crate::ch_emitter::{
     EmitterConfig, EmitterError, EmitterStats, append_buf, connect_client, drain_to_end_of_stream,
-    is_retryable,
+    is_retryable, reconnect_if_idle,
 };
 use crate::pipeline::Fatal;
 use crate::pipeline::ack::AckHandle;
@@ -29,6 +29,7 @@ use std::sync::atomic::Ordering;
 
 struct Inserter {
     client: AsyncClient,
+    last_used: std::time::Instant,
     alloc: Allocator,
     config: EmitterConfig,
     /// Parsed column types per table, refreshed when a batch's `schema_epoch`
@@ -74,6 +75,7 @@ impl Inserter {
         let retry = self.config.retry.clone();
         let mut attempt = 0u32;
         let mut backoff = retry.initial_backoff;
+        reconnect_if_idle(&mut self.client, &self.config, self.last_used).await?;
         loop {
             let attempt_result = match tokio::time::timeout(self.config.insert_timeout, async {
                 self.client.send_query(sql, None).await?;
@@ -93,7 +95,10 @@ impl Inserter {
                 }),
             };
             match attempt_result {
-                Ok(()) => return Ok(()),
+                Ok(()) => {
+                    self.last_used = std::time::Instant::now();
+                    return Ok(());
+                }
                 Err(e) if is_retryable(&e) && attempt < retry.max_attempts => {
                     self.stats.retries_attempted.fetch_add(1, Ordering::Relaxed);
                     attempt += 1;
@@ -177,6 +182,7 @@ pub async fn spawn_pool(
         let client = connect_client(config).await?;
         let inserter = Inserter {
             client,
+            last_used: std::time::Instant::now(),
             alloc: Allocator::global(&mimalloc::MiMalloc),
             config: config.clone(),
             asts: HashMap::new(),