Fix code review issues: race conditions, API, and docs

Correctness fixes:
- Fix race condition in shard creation by using atomic counter for IDs
  instead of calculating from vector length
- Fix subscription state inconsistency: rollback subscription from state
  when send fails, allowing retry without unsubscribe first

API improvements:
- Add is_running() method and prevent double start()
- Rename subscribe_simple to subscribe_all with better docs
- Add config validations for circuit_breaker_threshold > 0 and
  data_timeout >= ping_interval

Documentation:
- Add thread safety docs to ShardManager and Metrics
- Document all MetricsSnapshot fields
- Add ConfigError::InvalidConnection variant

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

Author: KaboomFox

examples/polymarket.rs

@@ -143,7 +143,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .map(|i| format!("token-{}", i))
         .collect();
 
-    manager.subscribe_simple(new_tokens).await?;
+    manager.subscribe_all(new_tokens).await?;
     info!("Subscribed to 10 more tokens, total: {}", manager.total_subscriptions());
 
     // Run for a while

src/config.rs

@@ -114,6 +114,19 @@ impl ShardManagerConfigBuilder {
             ));
         }
 
+        if self.config.health.data_timeout < self.config.health.ping_interval {
+            return Err(ConfigError::InvalidHealth(
+                "data_timeout should be >= ping_interval to avoid false positives".to_string(),
+            ));
+        }
+
+        // Validate connection config
+        if self.config.connection.circuit_breaker_threshold == 0 {
+            return Err(ConfigError::InvalidConnection(
+                "circuit_breaker_threshold must be > 0".to_string(),
+            ));
+        }
+
         // Validate max subscriptions
         if let Some(0) = self.config.max_subscriptions_per_shard {
             return Err(ConfigError::InvalidSubscriptionLimit(
@@ -134,6 +147,9 @@ pub enum ConfigError {
     /// Invalid health configuration
     #[error("Invalid health configuration: {0}")]
     InvalidHealth(String),
+    /// Invalid connection configuration
+    #[error("Invalid connection configuration: {0}")]
+    InvalidConnection(String),
     /// Invalid subscription limit
     #[error("Invalid subscription limit: {0}")]
     InvalidSubscriptionLimit(String),

src/connection.rs

@@ -63,6 +63,7 @@ impl<H: WebSocketHandler> Connection<H> {
     }
 
     /// Create a new connection with a ready signal for hot switchover
+    #[allow(clippy::too_many_arguments)]
     pub fn with_ready_signal(
         shard_id: usize,
         handler: Arc<H>,

src/manager.rs

@@ -8,6 +8,7 @@ use futures_util::FutureExt;
 use parking_lot::RwLock;
 use std::collections::{HashMap, HashSet};
 use std::panic::AssertUnwindSafe;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::JoinHandle;
@@ -23,13 +24,21 @@ type ShardUnsubData<S> = (usize, Vec<S>, mpsc::Sender<ConnectionCommand>, usize)
 /// Timeout for waiting for new connection during hot switchover
 const HOT_SWITCHOVER_TIMEOUT: Duration = Duration::from_secs(10);
 
-/// Manages multiple WebSocket shards with auto-rebalancing and hot switchover
+/// Manages multiple WebSocket shards with auto-rebalancing and hot switchover.
+///
+/// # Thread Safety
+///
+/// `ShardManager` is `Send + Sync` and all methods can be safely called from
+/// multiple tasks concurrently. Internal state is protected by `parking_lot::RwLock`
+/// which does not poison on panic.
 pub struct ShardManager<H: WebSocketHandler> {
     handler: Arc<H>,
     config: ShardManagerConfig,
     metrics: Arc<Metrics>,
     state: Arc<RwLock<ManagerState<H::Subscription>>>,
     shard_handles: RwLock<Vec<JoinHandle<()>>>,
+    /// Monotonically increasing counter for shard IDs to prevent race conditions
+    next_shard_id: AtomicUsize,
 }
 
 struct ManagerState<S: Clone + Eq + std::hash::Hash> {
@@ -65,6 +74,7 @@ impl<H: WebSocketHandler> ShardManager<H> {
             metrics: Arc::new(Metrics::new()),
             state: Arc::new(RwLock::new(ManagerState::default())),
             shard_handles: RwLock::new(Vec::new()),
+            next_shard_id: AtomicUsize::new(0),
         }
     }
 
@@ -73,12 +83,32 @@ impl<H: WebSocketHandler> ShardManager<H> {
         self.metrics.clone()
     }
 
+    /// Check if the manager is currently running
+    pub fn is_running(&self) -> bool {
+        self.state.read().is_running
+    }
+
     /// Start the shard manager
     ///
     /// This will create the initial shards based on existing subscriptions
     /// from the handler. If there are no initial subscriptions, no shards
     /// are created until the first subscription is added (lazy creation).
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the manager is already running or if configuration
+    /// is invalid.
     pub async fn start(&self) -> Result<(), Error> {
+        // Check if already running
+        {
+            let state = self.state.read();
+            if state.is_running {
+                return Err(Error::Handler(
+                    "ShardManager is already running".to_string(),
+                ));
+            }
+        }
+
         let subscriptions = self.handler.subscriptions();
         let max_per_shard = self.max_per_shard();
 
@@ -103,8 +133,9 @@ impl<H: WebSocketHandler> ShardManager<H> {
             max_per_shard
         );
 
-        // Create shards
-        for shard_id in 0..shard_count {
+        // Create shards using atomic counter for IDs
+        for _ in 0..shard_count {
+            let shard_id = self.next_shard_id.fetch_add(1, Ordering::SeqCst);
             self.create_shard(shard_id).await?;
         }
 
@@ -199,6 +230,9 @@ impl<H: WebSocketHandler> ShardManager<H> {
             state.shards_being_created.clear();
         }
 
+        // Reset shard ID counter for potential restart
+        self.next_shard_id.store(0, Ordering::SeqCst);
+
         info!("ShardManager stopped");
         Ok(())
     }
@@ -270,8 +304,8 @@ impl<H: WebSocketHandler> ShardManager<H> {
                     Ok((shard_id, tx, sub_count))
                 }
                 None if self.config.auto_rebalance => {
-                    // Need to create new shard - reserve the ID to prevent race conditions
-                    let new_id = state.shards.len() + state.shards_being_created.len();
+                    // Need to create new shard - use atomic counter for thread-safe ID assignment
+                    let new_id = self.next_shard_id.fetch_add(1, Ordering::SeqCst);
                     state.shards_being_created.insert(new_id);
                     Err(new_id)
                 }
@@ -328,13 +362,30 @@ impl<H: WebSocketHandler> ShardManager<H> {
             .update_shard(shard_id, |s| s.subscription_count = sub_count);
 
         // Send subscribe message outside the lock
-        if let Some(msg) = self.handler.subscription_message(&[item]) {
+        if let Some(msg) = self.handler.subscription_message(std::slice::from_ref(&item)) {
             if let Err(e) = command_tx.send(ConnectionCommand::Send(msg)).await {
                 self.metrics.record_subscription_send_failed();
                 warn!(
                     "[SHARD-{}] Failed to send subscription message: {}",
                     shard_id, e
                 );
+
+                // Rollback: remove subscription from state so retry is possible
+                {
+                    let mut state = self.state.write();
+                    if let Some(shard) = state.shards.get_mut(shard_id) {
+                        shard.remove_subscription(&item);
+                    }
+                    state.subscription_to_shard.remove(&item);
+                }
+
+                // Update metrics to reflect rollback
+                let new_count = {
+                    let state = self.state.read();
+                    state.shards.get(shard_id).map(|s| s.subscription_count()).unwrap_or(0)
+                };
+                self.metrics.update_shard(shard_id, |s| s.subscription_count = new_count);
+
                 return SubscribeResult::SendFailed {
                     shard_id,
                     error: e.to_string(),
@@ -345,10 +396,16 @@ impl<H: WebSocketHandler> ShardManager<H> {
         SubscribeResult::Success { shard_id }
     }
 
-    /// Convenience method that returns only the shard IDs of successful subscriptions
+    /// Subscribe to items and return affected shard IDs.
+    ///
+    /// This is a convenience method that returns the shard IDs of successful
+    /// subscriptions. Use [`subscribe`] instead if you need per-item error details.
+    ///
+    /// # Errors
     ///
-    /// This is a simpler API for cases where you don't need per-item error handling.
-    pub async fn subscribe_simple(&self, items: Vec<H::Subscription>) -> Result<Vec<usize>, Error> {
+    /// Returns an error only if all subscriptions fail. Partial success returns `Ok`
+    /// with the shard IDs that were affected.
+    pub async fn subscribe_all(&self, items: Vec<H::Subscription>) -> Result<Vec<usize>, Error> {
         let results = self.subscribe(items).await;
         let mut affected_shards_set = HashSet::new();
         let mut had_failure = false;

src/metrics.rs

@@ -2,12 +2,18 @@ use parking_lot::RwLock;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
 
-/// Metrics for observability
+/// Metrics for observability.
 ///
 /// This struct provides counters and gauges for monitoring WebSocket health.
 /// Use `snapshot()` to get a point-in-time view of all metrics, or use
 /// individual getter methods for specific values.
 ///
+/// # Thread Safety
+///
+/// `Metrics` is `Send + Sync` and all methods are safe to call from multiple
+/// tasks concurrently. Counters use atomic operations, and per-shard metrics
+/// are protected by `parking_lot::RwLock`.
+///
 /// # Example
 /// ```ignore
 /// let metrics = manager.metrics();
@@ -319,24 +325,43 @@ impl Metrics {
     }
 }
 
-/// A point-in-time snapshot of all metrics
+/// A point-in-time snapshot of all metrics.
+///
+/// Use [`Metrics::snapshot()`] to get a consistent view of all metrics at once.
+/// This is the recommended way to export metrics to monitoring systems.
 #[derive(Debug, Clone)]
 pub struct MetricsSnapshot {
+    /// Total number of WebSocket connections established (including reconnections)
     pub connections_total: u64,
+    /// Total number of reconnection attempts after disconnection
     pub reconnections_total: u64,
+    /// Total number of WebSocket messages received across all shards
     pub messages_received_total: u64,
+    /// Total number of WebSocket messages sent across all shards
     pub messages_sent_total: u64,
+    /// Total number of errors encountered (connection failures, panics, etc.)
     pub errors_total: u64,
+    /// Total number of WebSocket ping frames sent for health monitoring
     pub pings_sent_total: u64,
+    /// Total number of WebSocket pong frames received
     pub pongs_received_total: u64,
+    /// Total number of health check failures (pong timeouts, data timeouts)
     pub health_failures_total: u64,
+    /// Total number of shard rebalancing operations (new shard creation)
     pub rebalances_total: u64,
+    /// Total number of hot switchover operations initiated
     pub hot_switchovers_total: u64,
+    /// Total number of hot switchover operations that failed
     pub hot_switchover_failures_total: u64,
+    /// Total number of subscription message send failures
     pub subscription_send_failures_total: u64,
+    /// Total number of times the circuit breaker tripped due to consecutive failures
     pub circuit_breaker_trips_total: u64,
+    /// Current number of active (connected) shards
     pub active_connections: usize,
+    /// Current total number of subscriptions across all shards
     pub total_subscriptions: usize,
+    /// Per-shard metrics with detailed connection state
     pub shards: Vec<ShardMetrics>,
 }