Fix data timeout detection and add handler accessor

- Fix data_timeout not triggering on silent disconnects:
  Only count actual data messages for data timeout, not ping/pong.
  Previously, pong responses reset the timeout even when no data flowed.

- Add handler() method to ShardManager for accessing the handler
  Enables callers to invoke handler-specific methods like state refresh.

- Add SubscriptionError::Timeout variant for subscription timeouts

- Improve health monitor with data_timeout_remaining() helper

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: KaboomFox

src/connection.rs

@@ -15,7 +15,7 @@ use tokio_tungstenite::{
     client_async_tls_with_config, tungstenite::client::IntoClientRequest, tungstenite::Message,
     Connector, MaybeTlsStream, WebSocketStream,
 };
-use tracing::{debug, error, info, warn};
+use tracing::{debug, error, info, trace, warn};
 use url::Url;
 
 /// Commands that can be sent to a connection
@@ -42,10 +42,13 @@ pub struct Connection<H: WebSocketHandler> {
     ready_tx: Option<oneshot::Sender<()>>,
     /// Explicit subscriptions for this shard (used during hot switchover)
     initial_subscriptions: Option<Vec<H::Subscription>>,
+    /// Optional channel to request hot switchover from manager
+    switchover_tx: Option<mpsc::Sender<usize>>,
 }
 
 impl<H: WebSocketHandler> Connection<H> {
-    /// Create a new connection manager
+    /// Create a new connection manager (used for hot switchover connections)
+    #[allow(dead_code)]
     pub fn new(
         shard_id: usize,
         handler: Arc<H>,
@@ -65,6 +68,33 @@ impl<H: WebSocketHandler> Connection<H> {
             command_rx,
             ready_tx: None,
             initial_subscriptions: None,
+            switchover_tx: None,
+        }
+    }
+
+    /// Create a new connection with a switchover request channel
+    #[allow(clippy::too_many_arguments)]
+    pub fn with_switchover_channel(
+        shard_id: usize,
+        handler: Arc<H>,
+        config: ConnectionConfig,
+        backoff: BackoffConfig,
+        health_config: HealthConfig,
+        metrics: Arc<Metrics>,
+        command_rx: mpsc::Receiver<ConnectionCommand>,
+        switchover_tx: mpsc::Sender<usize>,
+    ) -> Self {
+        Self {
+            shard_id,
+            handler,
+            config,
+            backoff,
+            health_config,
+            metrics,
+            command_rx,
+            ready_tx: None,
+            initial_subscriptions: None,
+            switchover_tx: Some(switchover_tx),
         }
     }
 
@@ -91,6 +121,7 @@ impl<H: WebSocketHandler> Connection<H> {
             command_rx,
             ready_tx: Some(ready_tx),
             initial_subscriptions: Some(subscriptions),
+            switchover_tx: None,
         }
     }
 
@@ -235,18 +266,22 @@ impl<H: WebSocketHandler> Connection<H> {
     /// Otherwise, handlers are spawned in a separate task for panic isolation.
     #[inline]
     async fn call_on_message(&self, message: Message, state: &ConnectionState) {
+        let start = Instant::now();
         if self.config.low_latency_mode {
             // Direct call - no spawn overhead (~1-5µs savings per message)
             self.handler.on_message(message, state).await;
+            trace!("[SHARD-{}] Handler processed message in {:?}", self.shard_id, start.elapsed());
         } else {
             // Spawn for panic protection
             let handler = self.handler.clone();
             let state_clone = state.clone();
             let shard_id = self.shard_id;
 
             let result = tokio::task::spawn(async move {
+                let start = Instant::now();
                 let fut = AssertUnwindSafe(handler.on_message(message, &state_clone));
-                fut.await
+                fut.await;
+                trace!("[SHARD-{}] Handler processed message in {:?}", shard_id, start.elapsed());
             })
             .await;
 
@@ -365,9 +400,9 @@ impl<H: WebSocketHandler> Connection<H> {
                 msg = read.next() => {
                     match msg {
                         Some(Ok(message)) => {
-                            health.record_data_received();
                             self.metrics.record_message_received();
                             self.metrics.record_shard_message_received(self.shard_id);
+                            trace!("[SHARD-{}] Received message: {} bytes", self.shard_id, message.len());
 
                             match &message {
                                 Message::Ping(data) => {
@@ -384,6 +419,10 @@ impl<H: WebSocketHandler> Connection<H> {
                                     break;
                                 }
                                 _ => {
+                                    // Only count actual data messages for data timeout
+                                    // (not ping/pong which keep connection alive but don't indicate data flow)
+                                    health.record_data_received();
+
                                     // Check for application-level heartbeat
                                     if self.handler.is_heartbeat(&message) {
                                         debug!("[SHARD-{}] Received application heartbeat", self.shard_id);
@@ -452,8 +491,22 @@ impl<H: WebSocketHandler> Connection<H> {
 
                     // Check data timeout
                     if health.is_data_timeout() {
-                        warn!("[SHARD-{}] Data timeout, reconnecting", self.shard_id);
                         self.metrics.record_health_failure();
+
+                        // Request hot switchover if channel available, otherwise regular reconnect
+                        if let Some(ref tx) = self.switchover_tx {
+                            warn!("[SHARD-{}] Data timeout, requesting hot switchover", self.shard_id);
+                            // Non-blocking send - if manager is busy, fall back to regular reconnect
+                            if tx.try_send(self.shard_id).is_ok() {
+                                // Continue running while manager performs switchover
+                                // The manager will close this connection when new one is ready
+                                health.reset_data_timeout();
+                                continue;
+                            }
+                            warn!("[SHARD-{}] Hot switchover channel full, falling back to reconnect", self.shard_id);
+                        } else {
+                            warn!("[SHARD-{}] Data timeout, reconnecting", self.shard_id);
+                        }
                         break;
                     }
 

src/error.rs

@@ -95,3 +95,20 @@ pub enum SubscribeResult {
     /// Failed to subscribe
     Failed { error: String },
 }
+
+impl SubscribeResult {
+    /// Returns true if the subscription was successful or already existed
+    pub fn is_success(&self) -> bool {
+        matches!(self, SubscribeResult::Success { .. } | SubscribeResult::AlreadySubscribed { .. })
+    }
+
+    /// Returns the shard_id if available
+    pub fn shard_id(&self) -> Option<usize> {
+        match self {
+            SubscribeResult::Success { shard_id }
+            | SubscribeResult::AlreadySubscribed { shard_id }
+            | SubscribeResult::SendFailed { shard_id, .. } => Some(*shard_id),
+            SubscribeResult::Failed { .. } => None,
+        }
+    }
+}

src/health.rs

@@ -98,6 +98,11 @@ impl HealthMonitor {
         }
     }
 
+    /// Reset data timeout timer (used when hot switchover is requested)
+    pub fn reset_data_timeout(&mut self) {
+        self.last_data_received = Some(Instant::now());
+    }
+
     /// Check if the connection is unhealthy based on consecutive failures
     pub fn is_unhealthy(&self) -> bool {
         self.consecutive_ping_failures >= self.config.failure_threshold

src/manager.rs

@@ -13,7 +13,7 @@ use std::sync::Arc;
 use tokio::sync::{mpsc, oneshot, Mutex};
 use tokio::task::JoinHandle;
 use tokio::time::Duration;
-use tracing::{debug, error, info, warn};
+use tracing::{debug, error, info, trace, warn};
 
 /// Default channel buffer size
 const DEFAULT_CHANNEL_SIZE: usize = 100;
@@ -43,6 +43,10 @@ pub struct ShardManager<H: WebSocketHandler> {
     /// Mutex to serialize start/stop operations and prevent race conditions.
     /// This is a tokio::Mutex so it can be held across await points.
     lifecycle_lock: Mutex<()>,
+    /// Channel for connections to request hot switchover
+    switchover_tx: mpsc::Sender<usize>,
+    /// Receiver for switchover requests (wrapped in Mutex for interior mutability)
+    switchover_rx: Mutex<mpsc::Receiver<usize>>,
 }
 
 struct ManagerState<S: Clone + Eq + std::hash::Hash> {
@@ -73,6 +77,8 @@ impl<S: Clone + Eq + std::hash::Hash> Default for ManagerState<S> {
 impl<H: WebSocketHandler> ShardManager<H> {
     /// Create a new shard manager
     pub fn new(config: ShardManagerConfig, handler: H) -> Self {
+        // Channel for connections to request hot switchover (buffer 10 requests)
+        let (switchover_tx, switchover_rx) = mpsc::channel(10);
         Self {
             handler: Arc::new(handler),
             config,
@@ -81,6 +87,8 @@ impl<H: WebSocketHandler> ShardManager<H> {
             shard_handles: RwLock::new(HashMap::new()),
             next_shard_id: AtomicUsize::new(0),
             lifecycle_lock: Mutex::new(()),
+            switchover_tx,
+            switchover_rx: Mutex::new(switchover_rx),
         }
     }
 
@@ -89,6 +97,11 @@ impl<H: WebSocketHandler> ShardManager<H> {
         self.metrics.clone()
     }
 
+    /// Get a reference to the handler
+    pub fn handler(&self) -> &Arc<H> {
+        &self.handler
+    }
+
     /// Check if the manager is currently running
     pub fn is_running(&self) -> bool {
         self.state.read().is_running
@@ -202,6 +215,28 @@ impl<H: WebSocketHandler> ShardManager<H> {
         Ok(())
     }
 
+    /// Process pending hot switchover requests from connections.
+    ///
+    /// Call this periodically (e.g., in a background task) to handle
+    /// data timeout triggered switchovers. Returns the number of
+    /// switchovers initiated.
+    pub async fn process_switchover_requests(&self) -> usize {
+        let mut count = 0;
+        let mut rx = self.switchover_rx.lock().await;
+
+        // Process all pending requests
+        while let Ok(shard_id) = rx.try_recv() {
+            info!("[SHARD-{}] Processing hot switchover request", shard_id);
+            if let Err(e) = self.hot_switchover(shard_id).await {
+                warn!("[SHARD-{}] Hot switchover failed: {}", shard_id, e);
+            } else {
+                count += 1;
+            }
+        }
+
+        count
+    }
+
     /// Stop all shards gracefully
     ///
     /// This will close all connections and wait for tasks to complete.
@@ -287,6 +322,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
 
             // Check if already subscribed
             if let Some(&shard_id) = state.subscription_to_shard.get(&item) {
+                trace!("[SHARD-{}] Subscription already exists, skipping", shard_id);
                 return SubscribeResult::AlreadySubscribed { shard_id };
             }
 
@@ -332,6 +368,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
                         (tx, sub_count)
                     };
                     state.subscription_to_shard.insert(item.clone(), shard_id);
+                    trace!("[SHARD-{}] Added subscription (count: {})", shard_id, add_result.1);
                     Ok((shard_id, add_result.0, add_result.1))
                 }
                 None if self.config.auto_rebalance => {
@@ -397,6 +434,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
                     (tx, count)
                 };
                 state.subscription_to_shard.insert(item.clone(), new_id);
+                trace!("[SHARD-{}] Created new shard, added subscription (count: {})", new_id, add_result.1);
                 (new_id, add_result.0, add_result.1)
             }
         };
@@ -519,6 +557,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
                     for sub in &subs {
                         shard.remove_subscription(sub);
                     }
+                    trace!("[SHARD-{}] Removed {} subscriptions", shard_id, subs.len());
 
                     let tx = shard.command_tx.clone();
                     let count = shard.subscription_count();
@@ -594,10 +633,17 @@ impl<H: WebSocketHandler> ShardManager<H> {
         let shard = metrics.get(shard_id)?;
 
         if !shard.is_connected {
+            trace!("[SHARD-{}] Freshness check: disconnected", shard_id);
             return None;
         }
 
-        shard.time_since_last_message
+        let freshness = shard.time_since_last_message;
+        trace!(
+            "[SHARD-{}] Freshness check: {:?}",
+            shard_id,
+            freshness.map(|d| format!("{}ms", d.as_millis())).unwrap_or_else(|| "no data".to_string())
+        );
+        freshness
     }
 
     /// Check if a subscription's data is considered fresh.
@@ -626,12 +672,26 @@ impl<H: WebSocketHandler> ShardManager<H> {
         max_staleness: Duration,
     ) -> bool {
         let Some(shard_id) = self.subscription_shard(subscription) else {
+            trace!("Freshness check: subscription not found");
             return false;
         };
 
         match self.shard_freshness(shard_id) {
-            Some(since_last_message) => since_last_message <= max_staleness,
-            None => false,
+            Some(since_last_message) => {
+                let is_fresh = since_last_message <= max_staleness;
+                trace!(
+                    "[SHARD-{}] Subscription freshness: {}ms vs max {}ms -> {}",
+                    shard_id,
+                    since_last_message.as_millis(),
+                    max_staleness.as_millis(),
+                    if is_fresh { "FRESH" } else { "STALE" }
+                );
+                is_fresh
+            }
+            None => {
+                trace!("[SHARD-{}] Subscription freshness: no data -> STALE", shard_id);
+                false
+            }
         }
     }
 
@@ -855,6 +915,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
         let backoff_config = self.config.backoff.clone();
         let health_config = self.config.health.clone();
         let metrics = self.metrics.clone();
+        let switchover_tx = self.switchover_tx.clone();
 
         let handle = tokio::spawn(async move {
             Self::run_connection_with_recovery(
@@ -865,6 +926,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
                 health_config,
                 metrics,
                 rx,
+                switchover_tx,
             )
             .await
         });
@@ -899,18 +961,20 @@ impl<H: WebSocketHandler> ShardManager<H> {
         health_config: crate::config::HealthConfig,
         metrics: Arc<Metrics>,
         command_rx: mpsc::Receiver<ConnectionCommand>,
+        switchover_tx: mpsc::Sender<usize>,
     ) {
         // Note: We can't easily restart the connection after channel closure,
         // so we just catch panics within this task and log them.
         // For full recovery, the manager would need to recreate the channel.
-        let connection = Connection::new(
+        let connection = Connection::with_switchover_channel(
             shard_id,
             handler,
             connection_config,
             backoff_config,
             health_config,
             metrics.clone(),
             command_rx,
+            switchover_tx,
         );
 
         match AssertUnwindSafe(connection.run())