nntp_proxy/

proxy.rs

//! NNTP Proxy implementation
//!
//! This module contains the main `NntpProxy` struct which orchestrates
//! connection handling, routing, and resource management.

use anyhow::{Context, Result};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tracing::{debug, info, warn};

use crate::auth::AuthHandler;
use crate::cache::{HybridCacheConfig, UnifiedCache};
use crate::config::{Config, RoutingMode, Server};
use crate::constants::buffer::{POOL, POOL_COUNT};
use crate::metrics::{ConnectionStatsAggregator, MetricsCollector};
use crate::network::NetworkOptimizer;
use crate::pool::{BufferPool, ConnectionProvider, DeadpoolConnectionProvider, prewarm_pools};
use crate::router;
use crate::session::ClientSession;
use crate::types::ClientAddress;
use crate::types::{self, BufferSize, TransferMetrics};

/// Builder for constructing an `NntpProxy` with optional configuration overrides
///
/// # Examples
///
/// Basic usage with defaults:
/// ```no_run
/// # #[tokio::main]
/// # async fn main() -> anyhow::Result<()> {
/// use nntp_proxy::{NntpProxyBuilder, Config, RoutingMode};
/// use nntp_proxy::config::load_config;
///
/// let config = load_config("config.toml")?;
/// let proxy = NntpProxyBuilder::new(config)
///     .with_routing_mode(RoutingMode::Hybrid)
///     .build()
///     .await?;
/// # Ok(())
/// # }
/// ```
///
/// With custom buffer pool size:
/// ```no_run
/// # #[tokio::main]
/// # async fn main() -> anyhow::Result<()> {
/// use nntp_proxy::{NntpProxyBuilder, Config, RoutingMode};
/// use nntp_proxy::config::load_config;
///
/// let config = load_config("config.toml")?;
/// let proxy = NntpProxyBuilder::new(config)
///     .with_routing_mode(RoutingMode::PerCommand)
///     .with_buffer_pool_size(512 * 1024)  // 512KB buffers
///     .with_buffer_pool_count(64)         // 64 buffers
///     .build()
///     .await?;
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
pub struct NntpProxyBuilder {
    config: Config,
    routing_mode: RoutingMode,
    buffer_size: Option<usize>,
    buffer_count: Option<usize>,
}

impl NntpProxyBuilder {
    /// Create a new builder with the given configuration
    ///
    /// The routing mode defaults to `Standard` (1:1) mode.
    #[must_use]
    pub fn new(config: Config) -> Self {
        Self {
            config,
            routing_mode: RoutingMode::Stateful,
            buffer_size: None,
            buffer_count: None,
        }
    }

    /// Set the routing mode
    ///
    /// Available modes:
    /// - `Standard`: 1:1 client-to-backend mapping (default)
    /// - `PerCommand`: Each command routes to a different backend
    /// - `Hybrid`: Starts in per-command mode, switches to stateful when needed
    #[must_use]
    pub fn with_routing_mode(mut self, mode: RoutingMode) -> Self {
        self.routing_mode = mode;
        self
    }

    /// Override the default buffer pool size (256KB)
    ///
    /// This affects the size of each buffer in the pool. Larger buffers
    /// can improve throughput for large article transfers but use more memory.
    #[must_use]
    pub fn with_buffer_pool_size(mut self, size: usize) -> Self {
        self.buffer_size = Some(size);
        self
    }

    /// Override the default buffer pool count (32)
    ///
    /// This affects how many buffers are pre-allocated. Should roughly match
    /// the expected number of concurrent connections.
    #[must_use]
    pub fn with_buffer_pool_count(mut self, count: usize) -> Self {
        self.buffer_count = Some(count);
        self
    }

    /// Build the `NntpProxy` instance
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No servers are configured
    /// - Connection providers cannot be created
    /// - Buffer size is zero
    /// - Hybrid cache initialization fails (if disk cache configured)
    pub async fn build(self) -> Result<NntpProxy> {
        if self.config.servers.is_empty() {
            anyhow::bail!("No servers configured in configuration");
        }

        // Use provided values or defaults
        let buffer_size = self.buffer_size.unwrap_or(POOL);
        let buffer_count = self.buffer_count.unwrap_or(POOL_COUNT);

        // Create deadpool connection providers for each server
        let connection_providers: Result<Vec<DeadpoolConnectionProvider>> = self
            .config
            .servers
            .iter()
            .map(|server| {
                info!(
                    "Configuring deadpool connection provider for '{}'",
                    server.name
                );
                DeadpoolConnectionProvider::from_server_config(server)
            })
            .collect();

        let connection_providers = connection_providers?;

        let buffer_pool = BufferPool::new(
            BufferSize::try_new(buffer_size)
                .map_err(|_| anyhow::anyhow!("Buffer size must be non-zero"))?,
            buffer_count,
        );

        // Create metrics collector (before moving servers)
        let metrics = MetricsCollector::new(self.config.servers.len());

        let servers = Arc::new(self.config.servers);

        // Create backend selector and add all backends
        let router = Arc::new({
            use types::BackendId;
            let backend_strategy = self.config.proxy.backend_selection;
            connection_providers.iter().enumerate().fold(
                router::BackendSelector::with_strategy(backend_strategy),
                |mut r, (idx, provider)| {
                    let backend_id = BackendId::from_index(idx);
                    r.add_backend(
                        backend_id,
                        servers[idx].name.clone(),
                        provider.clone(),
                        servers[idx].tier,
                    );
                    r
                },
            )
        });

        // Create auth handler from config
        let auth_handler = {
            let all_users: Vec<(String, String)> = self
                .config
                .client_auth
                .all_users()
                .into_iter()
                .map(|(u, p)| (u.to_string(), p.to_string()))
                .collect();

            if all_users.is_empty() {
                Arc::new(AuthHandler::default())
            } else {
                Arc::new(AuthHandler::with_users(all_users).with_context(|| {
                    "Invalid authentication configuration. \
                             If you set username/password in config, they cannot be empty. \
                             Remove them entirely to disable authentication."
                })?)
            }
        };

        // Create article cache (always enabled for availability tracking)
        // If max_capacity=0, only tracks which backends have articles (no content caching)
        let (cache, cache_articles) = if let Some(cache_config) = &self.config.cache {
            let capacity = cache_config.max_capacity.as_u64();
            let cache_articles = cache_config.cache_articles;

            // Check if disk cache is configured
            let cache = if let Some(disk_config) = &cache_config.disk {
                // Create hybrid cache with disk tier
                let hybrid_config = HybridCacheConfig {
                    memory_capacity: capacity,
                    disk_path: disk_config.path.clone(),
                    disk_capacity: disk_config.capacity.as_u64(),
                    shards: disk_config.shards,
                    compression: disk_config.compression,
                    ttl: cache_config.ttl,
                    cache_articles,
                };

                info!(
                    "Initializing hybrid cache: memory={}MB, disk={}GB at {:?}",
                    capacity / (1024 * 1024),
                    disk_config.capacity.as_u64() / (1024 * 1024 * 1024),
                    disk_config.path
                );

                Arc::new(
                    UnifiedCache::hybrid(hybrid_config)
                        .await
                        .context("Failed to initialize hybrid disk cache")?,
                )
            } else {
                // Memory-only cache
                Arc::new(UnifiedCache::memory(
                    capacity,
                    cache_config.ttl,
                    cache_articles,
                ))
            };

            if capacity > 0 {
                if cache_articles {
                    info!(
                        "Article cache enabled: max_capacity={}, ttl={}s (full caching)",
                        cache_config.max_capacity,
                        cache_config.ttl.as_secs()
                    );
                } else {
                    info!(
                        "Article cache enabled: max_capacity={}, ttl={}s (availability-only, bodies not cached)",
                        cache_config.max_capacity,
                        cache_config.ttl.as_secs()
                    );
                }
            } else {
                info!("Backend availability tracking enabled (cache disabled, capacity=0)");
            }
            (cache, cache_articles)
        } else {
            debug!("Cache not configured, using availability-only mode (capacity=0)");
            (
                Arc::new(UnifiedCache::memory(0, Duration::from_secs(3600), false)),
                true,
            )
        };

        // Extract adaptive_precheck from cache config (default: false)
        let adaptive_precheck = self
            .config
            .cache
            .as_ref()
            .map(|c| c.adaptive_precheck)
            .unwrap_or(false);

        let start_instant = Instant::now();

        Ok(NntpProxy {
            servers,
            router,
            connection_providers,
            buffer_pool,
            routing_mode: self.routing_mode,
            auth_handler,
            metrics,
            connection_stats: ConnectionStatsAggregator::new(),
            cache,
            cache_articles,
            adaptive_precheck,
            last_activity_nanos: Arc::new(AtomicU64::new(0)),
            active_clients: Arc::new(AtomicUsize::new(0)),
            start_instant,
        })
    }

    /// Build the `NntpProxy` instance (synchronous version)
    ///
    /// This version only supports memory-only cache. If disk cache is configured
    /// in the config, it will be ignored and a warning will be logged.
    /// Use `build()` for full disk cache support.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No servers are configured
    /// - Connection providers cannot be created
    /// - Buffer size is zero
    pub fn build_sync(self) -> Result<NntpProxy> {
        if self.config.servers.is_empty() {
            anyhow::bail!("No servers configured in configuration");
        }

        // Use provided values or defaults
        let buffer_size = self.buffer_size.unwrap_or(POOL);
        let buffer_count = self.buffer_count.unwrap_or(POOL_COUNT);

        // Create deadpool connection providers for each server
        let connection_providers: Result<Vec<DeadpoolConnectionProvider>> = self
            .config
            .servers
            .iter()
            .map(|server| {
                info!(
                    "Configuring deadpool connection provider for '{}'",
                    server.name
                );
                DeadpoolConnectionProvider::from_server_config(server)
            })
            .collect();

        let connection_providers = connection_providers?;

        let buffer_pool = BufferPool::new(
            BufferSize::try_new(buffer_size)
                .map_err(|_| anyhow::anyhow!("Buffer size must be non-zero"))?,
            buffer_count,
        );

        // Create metrics collector (before moving servers)
        let metrics = MetricsCollector::new(self.config.servers.len());

        let servers = Arc::new(self.config.servers);

        // Create backend selector and add all backends
        let router = Arc::new({
            use types::BackendId;
            let backend_strategy = self.config.proxy.backend_selection;
            connection_providers.iter().enumerate().fold(
                router::BackendSelector::with_strategy(backend_strategy),
                |mut r, (idx, provider)| {
                    let backend_id = BackendId::from_index(idx);
                    r.add_backend(
                        backend_id,
                        servers[idx].name.clone(),
                        provider.clone(),
                        servers[idx].tier,
                    );
                    r
                },
            )
        });

        // Create auth handler from config
        let auth_handler = {
            let all_users: Vec<(String, String)> = self
                .config
                .client_auth
                .all_users()
                .into_iter()
                .map(|(u, p)| (u.to_string(), p.to_string()))
                .collect();

            if all_users.is_empty() {
                Arc::new(AuthHandler::default())
            } else {
                Arc::new(AuthHandler::with_users(all_users).with_context(|| {
                    "Invalid authentication configuration. \
                             If you set username/password in config, they cannot be empty. \
                             Remove them entirely to disable authentication."
                })?)
            }
        };

        // Create article cache (memory-only in sync version)
        let (cache, cache_articles) = if let Some(cache_config) = &self.config.cache {
            let capacity = cache_config.max_capacity.as_u64();
            let cache_articles = cache_config.cache_articles;

            // Warn if disk cache is configured but we're using sync build
            if cache_config.disk.is_some() {
                warn!(
                    "Disk cache configured but build_sync() called - using memory-only cache. Use build() for disk cache support."
                );
            }

            // Memory-only cache
            let cache = Arc::new(UnifiedCache::memory(
                capacity,
                cache_config.ttl,
                cache_articles,
            ));

            if capacity > 0 {
                if cache_articles {
                    info!(
                        "Article cache enabled: max_capacity={}, ttl={}s (full caching)",
                        cache_config.max_capacity,
                        cache_config.ttl.as_secs()
                    );
                } else {
                    info!(
                        "Article cache enabled: max_capacity={}, ttl={}s (availability-only, bodies not cached)",
                        cache_config.max_capacity,
                        cache_config.ttl.as_secs()
                    );
                }
            } else {
                info!("Backend availability tracking enabled (cache disabled, capacity=0)");
            }
            (cache, cache_articles)
        } else {
            debug!("Cache not configured, using availability-only mode (capacity=0)");
            (
                Arc::new(UnifiedCache::memory(0, Duration::from_secs(3600), false)),
                true,
            )
        };

        // Extract adaptive_precheck from cache config (default: false)
        let adaptive_precheck = self
            .config
            .cache
            .as_ref()
            .map(|c| c.adaptive_precheck)
            .unwrap_or(false);

        let start_instant = Instant::now();

        Ok(NntpProxy {
            servers,
            router,
            connection_providers,
            buffer_pool,
            routing_mode: self.routing_mode,
            auth_handler,
            metrics,
            connection_stats: ConnectionStatsAggregator::new(),
            cache,
            cache_articles,
            adaptive_precheck,
            last_activity_nanos: Arc::new(AtomicU64::new(0)),
            active_clients: Arc::new(AtomicUsize::new(0)),
            start_instant,
        })
    }
}

#[derive(Debug, Clone)]
pub struct NntpProxy {
    servers: Arc<Vec<Server>>,
    /// Backend selector for round-robin load balancing
    router: Arc<router::BackendSelector>,
    /// Connection providers per server - easily swappable implementation
    connection_providers: Vec<DeadpoolConnectionProvider>,
    /// Buffer pool for I/O operations
    buffer_pool: BufferPool,
    /// Routing mode (Standard, PerCommand, or Hybrid)
    routing_mode: RoutingMode,
    /// Authentication handler for client auth interception
    auth_handler: Arc<AuthHandler>,
    /// Metrics collector for TUI and monitoring (always enabled)
    metrics: MetricsCollector,
    /// Connection statistics aggregator (reduces log spam)
    connection_stats: ConnectionStatsAggregator,
    /// Article cache (always present - tracks backend availability even with capacity=0)
    cache: Arc<UnifiedCache>,
    /// Whether to cache article bodies (config-driven)
    cache_articles: bool,
    /// Whether to use adaptive availability prechecking for STAT/HEAD
    adaptive_precheck: bool,
    /// Timestamp (as epoch nanos) when last client disconnected (for idle detection)
    /// Uses epoch nanos since Instant isn't shareable across threads easily
    last_activity_nanos: Arc<AtomicU64>,
    /// Number of currently active client connections
    active_clients: Arc<AtomicUsize>,
    /// Reference instant for converting nanos to duration
    start_instant: Instant,
}

/// Classify an error as a client disconnect (broken pipe/connection reset)
///
/// Returns true for errors that indicate the client disconnected normally,
/// which should be logged at DEBUG level rather than WARN.
///
/// # Examples
///
/// ```
/// use std::io::{Error, ErrorKind};
/// use nntp_proxy::is_client_disconnect_error;
///
/// let broken_pipe = Error::from(ErrorKind::BrokenPipe);
/// let wrapped = anyhow::Error::from(broken_pipe);
/// assert!(is_client_disconnect_error(&wrapped));
///
/// let other_error = anyhow::anyhow!("some other error");
/// assert!(!is_client_disconnect_error(&other_error));
/// ```
#[inline]
pub fn is_client_disconnect_error(e: &anyhow::Error) -> bool {
    crate::session::error_classification::ErrorClassifier::is_client_disconnect(e)
}

impl NntpProxy {
    // Helper methods for session management

    /// Idle timeout after which pools are cleared when a new client connects (5 minutes)
    const IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);

    #[inline]
    fn record_connection_opened(&self) {
        self.metrics.connection_opened();
    }

    #[inline]
    fn record_connection_closed(&self) {
        self.metrics.connection_closed();
    }

    /// Increment active client count
    ///
    /// Call this when a new client connection is accepted.
    #[inline]
    fn increment_active_clients(&self) {
        self.active_clients.fetch_add(1, Ordering::Relaxed);
    }

    /// Decrement active client count and update last activity timestamp
    ///
    /// Call this when a client connection closes.
    #[inline]
    fn decrement_active_clients(&self) {
        let prev = self.active_clients.fetch_sub(1, Ordering::Relaxed);

        // When last client disconnects, record the timestamp
        if prev == 1 {
            let nanos = self.start_instant.elapsed().as_nanos() as u64;
            self.last_activity_nanos.store(nanos, Ordering::Relaxed);
        }
    }

    /// Check if pools should be cleared due to idle timeout
    ///
    /// Returns true if pools were cleared.
    /// Pools are cleared when:
    /// 1. No clients are currently active
    /// 2. The last activity was more than IDLE_TIMEOUT ago
    ///
    /// This prevents stale connections from accumulating during overnight idle periods.
    fn check_and_clear_stale_pools(&self) -> bool {
        // Fast path: if there are active clients, pools are in use
        if self.active_clients.load(Ordering::Relaxed) > 0 {
            return false;
        }

        let last_activity_nanos = self.last_activity_nanos.load(Ordering::Relaxed);

        // If never been active, no need to clear
        if last_activity_nanos == 0 {
            return false;
        }

        let last_activity = Duration::from_nanos(last_activity_nanos);
        let now = self.start_instant.elapsed();
        let idle_duration = now.saturating_sub(last_activity);

        if idle_duration > Self::IDLE_TIMEOUT {
            info!(
                idle_secs = idle_duration.as_secs(),
                pool_count = self.connection_providers.len(),
                "Clearing stale pool connections after idle timeout"
            );

            for provider in &self.connection_providers {
                provider.clear_idle_connections();
            }

            true
        } else {
            false
        }
    }

    /// Build a session with standard configuration (conditionally enables metrics)
    fn build_session(
        &self,
        client_addr: ClientAddress,
        router: Option<Arc<router::BackendSelector>>,
        routing_mode: RoutingMode,
        cache: Arc<UnifiedCache>,
    ) -> ClientSession {
        // Start with base builder
        let builder = ClientSession::builder(
            client_addr,
            self.buffer_pool.clone(),
            self.auth_handler.clone(),
            self.metrics.clone(),
        )
        .with_routing_mode(routing_mode)
        .with_connection_stats(self.connection_stats.clone())
        .with_cache(cache)
        .with_cache_articles(self.cache_articles)
        .with_adaptive_precheck(self.adaptive_precheck);

        // Apply optional router
        let builder = match router {
            Some(r) => builder.with_router(r),
            None => builder,
        };

        builder.build()
    }

    /// Log session completion and record stats
    fn log_session_completion(
        &self,
        client_addr: ClientAddress,
        session_id: &str,
        session: &ClientSession,
        routing_mode: crate::config::RoutingMode,
        metrics: &types::TransferMetrics,
    ) {
        self.connection_stats
            .record_disconnection(session.username().as_deref(), routing_mode.short_name());

        debug!(
            "Session {} [{}] ↑{} ↓{}",
            client_addr,
            session_id,
            crate::formatting::format_bytes(metrics.client_to_backend.as_u64()),
            crate::formatting::format_bytes(metrics.backend_to_client.as_u64())
        );
    }

    /// Create a new `NntpProxy` with the given configuration and routing mode
    ///
    /// This is a convenience method that uses the builder internally.
    /// For more control over configuration, use [`NntpProxy::builder`].
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use nntp_proxy::{NntpProxy, Config, RoutingMode};
    /// # use nntp_proxy::config::load_config;
    /// # #[tokio::main]
    /// # async fn main() -> anyhow::Result<()> {
    /// let config = load_config("config.toml")?;
    /// let proxy = NntpProxy::new(config, RoutingMode::Hybrid).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn new(config: Config, routing_mode: RoutingMode) -> Result<Self> {
        NntpProxyBuilder::new(config)
            .with_routing_mode(routing_mode)
            .build()
            .await
    }

    /// Create a new `NntpProxy` synchronously (memory-only cache)
    ///
    /// This version only supports memory-only cache. If disk cache is configured,
    /// it will be ignored. Use [`NntpProxy::new`] for full disk cache support.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use nntp_proxy::{NntpProxy, Config, RoutingMode};
    /// # use nntp_proxy::config::load_config;
    /// # fn main() -> anyhow::Result<()> {
    /// let config = load_config("config.toml")?;
    /// let proxy = NntpProxy::new_sync(config, RoutingMode::Hybrid)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new_sync(config: Config, routing_mode: RoutingMode) -> Result<Self> {
        NntpProxyBuilder::new(config)
            .with_routing_mode(routing_mode)
            .build_sync()
    }

    /// Create a builder for more fine-grained control over proxy configuration
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use nntp_proxy::{NntpProxy, Config, RoutingMode};
    /// # use nntp_proxy::config::load_config;
    /// # #[tokio::main]
    /// # async fn main() -> anyhow::Result<()> {
    /// let config = load_config("config.toml")?;
    /// let proxy = NntpProxy::builder(config)
    ///     .with_routing_mode(RoutingMode::Hybrid)
    ///     .with_buffer_pool_size(512 * 1024)
    ///     .build()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    #[must_use]
    pub fn builder(config: Config) -> NntpProxyBuilder {
        NntpProxyBuilder::new(config)
    }

    /// Prewarm all connection pools before accepting clients
    /// Creates all connections concurrently and returns when ready
    pub async fn prewarm_connections(&self) -> Result<()> {
        prewarm_pools(&self.connection_providers, &self.servers).await
    }

    /// Gracefully shutdown all connection pools
    pub async fn graceful_shutdown(&self) {
        info!("Initiating graceful shutdown...");

        // Flush cache writes to disk before shutting down
        // With WriteOnInsertion, writes are enqueued async - close() waits for completion
        info!("Flushing disk cache writes...");
        if let Err(e) = self.cache.close().await {
            warn!("Error closing cache: {}", e);
        }

        info!("Shutting down connection pools...");
        for provider in &self.connection_providers {
            provider.graceful_shutdown().await;
        }

        info!("All connection pools have been shut down gracefully");
    }

    /// Get the list of servers
    #[must_use]
    #[inline]
    pub fn servers(&self) -> &[Server] {
        &self.servers
    }

    /// Get the router
    #[must_use]
    #[inline]
    pub fn router(&self) -> &Arc<router::BackendSelector> {
        &self.router
    }

    /// Get the connection providers
    #[must_use]
    #[inline]
    pub fn connection_providers(&self) -> &[DeadpoolConnectionProvider] {
        &self.connection_providers
    }

    /// Get the buffer pool
    #[must_use]
    #[inline]
    pub fn buffer_pool(&self) -> &BufferPool {
        &self.buffer_pool
    }

    /// Get the article cache (always present - capacity 0 if not configured)
    #[must_use]
    #[inline]
    pub fn cache(&self) -> &Arc<UnifiedCache> {
        &self.cache
    }

    /// Get the metrics collector
    #[must_use]
    #[inline]
    pub fn metrics(&self) -> &MetricsCollector {
        &self.metrics
    }

    /// Get connection stats aggregator
    #[must_use]
    #[inline]
    pub fn connection_stats(&self) -> &ConnectionStatsAggregator {
        &self.connection_stats
    }

    /// Log backend routing selection
    #[inline]
    fn log_routing_selection(
        &self,
        client_addr: ClientAddress,
        backend_id: crate::types::BackendId,
        server: &Server,
    ) {
        info!(
            "Routing client {} to backend {:?} ({}:{})",
            client_addr, backend_id, server.host, server.port
        );
    }

    /// Log connection pool status for monitoring
    #[inline]
    fn log_pool_status(&self, server_idx: usize) {
        let pool_status = self.connection_providers[server_idx].status();
        debug!(
            "Pool status for {}: {}/{} available, {} created",
            self.servers[server_idx].name,
            pool_status.available,
            pool_status.max_size,
            pool_status.created
        );
    }

    /// Prepare stateful connection - route, greet, optimize
    async fn prepare_stateful_connection(
        &self,
        client_stream: &mut TcpStream,
        client_addr: ClientAddress,
    ) -> Result<crate::types::BackendId> {
        self.record_connection_opened();

        let client_id = types::ClientId::new();
        let backend_id = self.router.route_command(client_id, "")?;
        let server_idx = backend_id.as_index();

        self.log_routing_selection(client_addr, backend_id, &self.servers[server_idx]);
        self.send_greeting(client_stream, client_addr).await?;
        self.log_pool_status(server_idx);
        self.apply_tcp_optimizations(client_stream);

        Ok(backend_id)
    }

    /// Prepare per-command connection - record, greet, optimize
    async fn prepare_per_command_connection(
        &self,
        client_stream: &mut TcpStream,
        client_addr: ClientAddress,
    ) -> Result<()> {
        self.record_connection_opened();
        self.send_greeting(client_stream, client_addr).await?;
        self.apply_tcp_optimizations(client_stream);
        Ok(())
    }

    /// Create session with router and cache configuration
    #[inline]
    fn create_session(
        &self,
        client_addr: ClientAddress,
        router: Option<Arc<crate::router::BackendSelector>>,
    ) -> ClientSession {
        self.build_session(client_addr, router, self.routing_mode, self.cache.clone())
    }

    /// Generate short session ID for logging
    #[inline]
    fn generate_session_id(&self, session: &ClientSession) -> String {
        crate::formatting::short_id(session.client_id().as_uuid())
    }

    /// Send greeting to client
    #[inline]
    async fn send_greeting(
        &self,
        client_stream: &mut TcpStream,
        client_addr: ClientAddress,
    ) -> Result<()> {
        crate::protocol::send_proxy_greeting(client_stream, client_addr).await
    }

    /// Apply TCP optimizations to client socket
    #[inline]
    fn apply_tcp_optimizations(&self, client_stream: &TcpStream) {
        use crate::network::TcpOptimizer;
        TcpOptimizer::new(client_stream)
            .optimize()
            .map_err(|e| debug!("Failed to optimize client socket: {}", e))
            .ok();
    }

    /// Get display name for current routing mode
    #[inline]
    fn routing_mode_display_name(&self) -> &'static str {
        if self.cache.entry_count() > 0 {
            "caching"
        } else {
            "per-command"
        }
    }

    /// Finalize stateful session with metrics and cleanup
    fn finalize_stateful_session(
        &self,
        metrics: Result<TransferMetrics>,
        client_addr: ClientAddress,
        session_id: &str,
        session: &ClientSession,
        backend_id: crate::types::BackendId,
    ) -> Result<()> {
        self.record_connection_if_unauthenticated(session);
        self.router.complete_command(backend_id);
        self.record_session_metrics(metrics, client_addr, session_id, session, Some(backend_id))?;
        self.record_connection_closed();
        Ok(())
    }

    /// Finalize per-command session with logging and cleanup
    fn finalize_per_command_session(
        &self,
        metrics: Result<TransferMetrics>,
        client_addr: ClientAddress,
        session_id: &str,
        session: &ClientSession,
    ) -> Result<()> {
        self.record_session_metrics(metrics, client_addr, session_id, session, None)?;
        self.record_connection_closed();
        Ok(())
    }

    /// Record connection for unauthenticated sessions only
    #[inline]
    fn record_connection_if_unauthenticated(&self, session: &ClientSession) {
        if !self.auth_handler.is_enabled() || session.username().is_none() {
            let mode = self.session_mode_label(session.mode());
            self.connection_stats
                .record_connection(session.username().as_deref(), mode);
        }
    }

    /// Record session metrics and log completion or errors
    fn record_session_metrics(
        &self,
        metrics: Result<TransferMetrics>,
        client_addr: ClientAddress,
        session_id: &str,
        session: &ClientSession,
        backend_id: Option<crate::types::BackendId>,
    ) -> Result<()> {
        match metrics {
            Ok(m) => {
                self.log_session_completion(
                    client_addr,
                    session_id,
                    session,
                    self.routing_mode,
                    &m,
                );

                if let Some(bid) = backend_id {
                    self.metrics
                        .record_client_to_backend_bytes_for(bid, m.client_to_backend.as_u64());
                    self.metrics
                        .record_backend_to_client_bytes_for(bid, m.backend_to_client.as_u64());
                }
                Ok(())
            }
            Err(e) => {
                if let Some(bid) = backend_id {
                    self.metrics.record_error(bid);
                }

                // Only log non-client-disconnect errors (avoid spam from normal disconnects)
                if !is_client_disconnect_error(&e) {
                    warn!("Session error for client {}: {:?}", client_addr, e);
                }
                Err(e)
            }
        }
    }

    /// Get session mode label for logging
    #[inline]
    fn session_mode_label(&self, session_mode: crate::session::SessionMode) -> &'static str {
        use crate::session::SessionMode;
        match (session_mode, self.routing_mode) {
            (SessionMode::PerCommand, _) => "per-command",
            (SessionMode::Stateful, RoutingMode::Stateful) => "standard",
            (SessionMode::Stateful, RoutingMode::Hybrid) => "hybrid",
            (SessionMode::Stateful, _) => "stateful",
        }
    }

    pub async fn handle_client(
        &self,
        mut client_stream: TcpStream,
        client_addr: ClientAddress,
    ) -> Result<()> {
        debug!("New client connection from {}", client_addr);

        // Check for stale pools before handling (lazy recreation after idle)
        self.check_and_clear_stale_pools();
        self.increment_active_clients();

        let result = async {
            let backend_id = self
                .prepare_stateful_connection(&mut client_stream, client_addr)
                .await?;
            let server_idx = backend_id.as_index();

            let session = self.create_session(client_addr, None);
            let session_id = self.generate_session_id(&session);

            debug!("Starting stateful session for client {}", client_addr);

            let metrics = session
                .handle_stateful_session(
                    client_stream,
                    backend_id,
                    &self.connection_providers[server_idx],
                    &self.servers[server_idx].name,
                )
                .await;

            self.finalize_stateful_session(metrics, client_addr, &session_id, &session, backend_id)
        }
        .await;

        self.decrement_active_clients();
        result
    }

    /// Handle client connection using per-command routing mode
    ///
    /// This creates a session with the router, allowing commands from this client
    /// to be routed to different backends based on load balancing.  
    pub async fn handle_client_per_command_routing(
        &self,
        client_stream: TcpStream,
        client_addr: ClientAddress,
    ) -> Result<()> {
        // Check for stale pools before handling (lazy recreation after idle)
        self.check_and_clear_stale_pools();
        self.increment_active_clients();

        let result = self
            .handle_per_command_client(client_stream, client_addr)
            .await;

        self.decrement_active_clients();
        result
    }

    /// Handle a per-command routing session
    async fn handle_per_command_client(
        &self,
        mut client_stream: TcpStream,
        client_addr: ClientAddress,
    ) -> Result<()> {
        let mode_label = self.routing_mode_display_name();
        debug!(
            "New {} routing client connection from {}",
            mode_label, client_addr
        );

        self.prepare_per_command_connection(&mut client_stream, client_addr)
            .await?;

        let session = self.create_session(client_addr, Some(self.router.clone()));
        let session_id = self.generate_session_id(&session);

        let metrics = session
            .handle_per_command_routing(client_stream)
            .await
            .with_context(|| {
                format!(
                    "{} routing session failed for {} [{}]",
                    mode_label, client_addr, session_id
                )
            });

        self.finalize_per_command_session(metrics, client_addr, &session_id, &session)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;

    fn create_test_config() -> Config {
        use crate::config::{health_check_max_per_cycle, health_check_pool_timeout};
        use crate::types::{HostName, MaxConnections, Port, ServerName};
        Config {
            servers: vec![
                Server {
                    host: HostName::try_new("server1.example.com".to_string()).unwrap(),
                    port: Port::try_new(119).unwrap(),
                    name: ServerName::try_new("Test Server 1".to_string()).unwrap(),
                    username: None,
                    password: None,
                    max_connections: MaxConnections::try_new(5).unwrap(),
                    use_tls: false,
                    tls_verify_cert: true,
                    tls_cert_path: None,
                    connection_keepalive: None,
                    health_check_max_per_cycle: health_check_max_per_cycle(),
                    health_check_pool_timeout: health_check_pool_timeout(),
                    tier: 0,
                },
                Server {
                    host: HostName::try_new("server2.example.com".to_string()).unwrap(),
                    port: Port::try_new(119).unwrap(),
                    name: ServerName::try_new("Test Server 2".to_string()).unwrap(),
                    username: None,
                    password: None,
                    max_connections: MaxConnections::try_new(8).unwrap(),
                    use_tls: false,
                    tls_verify_cert: true,
                    tls_cert_path: None,
                    connection_keepalive: None,
                    health_check_max_per_cycle: health_check_max_per_cycle(),
                    health_check_pool_timeout: health_check_pool_timeout(),
                    tier: 0,
                },
                Server {
                    host: HostName::try_new("server3.example.com".to_string()).unwrap(),
                    port: Port::try_new(119).unwrap(),
                    name: ServerName::try_new("Test Server 3".to_string()).unwrap(),
                    username: None,
                    password: None,
                    max_connections: MaxConnections::try_new(12).unwrap(),
                    use_tls: false,
                    tls_verify_cert: true,
                    tls_cert_path: None,
                    connection_keepalive: None,
                    health_check_max_per_cycle: health_check_max_per_cycle(),
                    health_check_pool_timeout: health_check_pool_timeout(),
                    tier: 0,
                },
            ],
            ..Default::default()
        }
    }

    #[test]
    fn test_proxy_creation_with_servers() {
        let config = create_test_config();
        let proxy = Arc::new(
            NntpProxy::new_sync(config, RoutingMode::Stateful).expect("Failed to create proxy"),
        );

        assert_eq!(proxy.servers().len(), 3);
        assert_eq!(proxy.servers()[0].name.as_str(), "Test Server 1");
    }

    #[test]
    fn test_proxy_creation_with_empty_servers() {
        let config = Config {
            servers: vec![],
            ..Default::default()
        };
        let result = NntpProxy::new_sync(config, RoutingMode::Stateful);

        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("No servers configured")
        );
    }

    #[test]
    fn test_proxy_has_router() {
        let config = create_test_config();
        let proxy = Arc::new(
            NntpProxy::new_sync(config, RoutingMode::Stateful).expect("Failed to create proxy"),
        );

        // Proxy should have a router with backends
        assert_eq!(proxy.router.backend_count(), 3);
    }

    #[test]
    fn test_builder_basic_usage() {
        let config = create_test_config();
        let proxy = NntpProxy::builder(config)
            .build_sync()
            .expect("Failed to build proxy");

        assert_eq!(proxy.servers().len(), 3);
        assert_eq!(proxy.router.backend_count(), 3);
    }

    #[test]
    fn test_builder_with_routing_mode() {
        let config = create_test_config();
        let proxy = NntpProxy::builder(config)
            .with_routing_mode(RoutingMode::PerCommand)
            .build_sync()
            .expect("Failed to build proxy");

        assert_eq!(proxy.servers().len(), 3);
    }

    #[test]
    fn test_builder_with_custom_buffer_pool() {
        let config = create_test_config();
        let proxy = NntpProxy::builder(config)
            .with_buffer_pool_size(512 * 1024)
            .with_buffer_pool_count(64)
            .build_sync()
            .expect("Failed to build proxy");

        assert_eq!(proxy.servers().len(), 3);
        // Pool size and count are used internally but not exposed for verification
    }

    #[test]
    fn test_builder_with_all_options() {
        let config = create_test_config();
        let proxy = NntpProxy::builder(config)
            .with_routing_mode(RoutingMode::Hybrid)
            .with_buffer_pool_size(1024 * 1024)
            .with_buffer_pool_count(16)
            .build_sync()
            .expect("Failed to build proxy");

        assert_eq!(proxy.servers().len(), 3);
        assert_eq!(proxy.router.backend_count(), 3);
    }

    #[test]
    fn test_builder_empty_servers_error() {
        let config = Config {
            servers: vec![],
            ..Default::default()
        };
        let result = NntpProxy::builder(config).build_sync();

        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("No servers configured")
        );
    }

    #[test]
    fn test_backward_compatibility_new() {
        // Ensure NntpProxy::new_sync() still works (it uses builder internally)
        let config = create_test_config();
        let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful)
            .expect("Failed to create proxy with new_sync()");

        assert_eq!(proxy.servers().len(), 3);
        assert_eq!(proxy.router.backend_count(), 3);
    }

    // Tests for is_client_disconnect_error function
    mod error_classification {
        use super::*;
        use std::io::{Error, ErrorKind};

        #[test]
        fn test_broken_pipe_is_client_disconnect() {
            let io_err = Error::from(ErrorKind::BrokenPipe);
            let err = anyhow::Error::from(io_err);
            assert!(is_client_disconnect_error(&err));
        }

        #[test]
        fn test_connection_reset_is_client_disconnect() {
            let io_err = Error::from(ErrorKind::ConnectionReset);
            let err = anyhow::Error::from(io_err);
            assert!(is_client_disconnect_error(&err));
        }

        #[test]
        fn test_other_io_errors_not_client_disconnect() {
            let error_kinds = vec![
                ErrorKind::NotFound,
                ErrorKind::PermissionDenied,
                ErrorKind::ConnectionRefused,
                ErrorKind::ConnectionAborted,
                ErrorKind::AddrInUse,
                ErrorKind::AddrNotAvailable,
                ErrorKind::TimedOut,
                ErrorKind::Interrupted,
                ErrorKind::UnexpectedEof,
                ErrorKind::WouldBlock,
            ];

            for kind in error_kinds {
                let io_err = Error::from(kind);
                let err = anyhow::Error::from(io_err);
                assert!(
                    !is_client_disconnect_error(&err),
                    "{:?} should not be classified as client disconnect",
                    kind
                );
            }
        }

        #[test]
        fn test_non_io_error_not_client_disconnect() {
            let err = anyhow::anyhow!("generic error message");
            assert!(!is_client_disconnect_error(&err));
        }

        #[test]
        fn test_wrapped_broken_pipe_error() {
            let io_err = Error::from(ErrorKind::BrokenPipe);
            let err = anyhow::Error::from(io_err).context("failed to write to client");
            assert!(is_client_disconnect_error(&err));
        }

        #[test]
        fn test_wrapped_connection_reset_error() {
            let io_err = Error::from(ErrorKind::ConnectionReset);
            let err = anyhow::Error::from(io_err).context("failed to read from client");
            assert!(is_client_disconnect_error(&err));
        }

        #[test]
        fn test_deeply_wrapped_error() {
            let io_err = Error::from(ErrorKind::BrokenPipe);
            let err = anyhow::Error::from(io_err)
                .context("inner context")
                .context("outer context");
            assert!(is_client_disconnect_error(&err));
        }

        #[test]
        fn test_custom_io_error_message() {
            let io_err = Error::new(ErrorKind::BrokenPipe, "custom broken pipe");
            let err = anyhow::Error::from(io_err);
            assert!(is_client_disconnect_error(&err));
        }
    }

    // Tests for new DRY helper methods
    mod helper_methods {
        use super::*;
        use crate::session::SessionMode;

        #[test]
        fn test_session_mode_label_per_command() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::PerCommand).unwrap();

            let label = proxy.session_mode_label(SessionMode::PerCommand);
            assert_eq!(label, "per-command");
        }

        #[test]
        fn test_session_mode_label_stateful_standard() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            let label = proxy.session_mode_label(SessionMode::Stateful);
            assert_eq!(label, "standard");
        }

        #[test]
        fn test_session_mode_label_stateful_hybrid() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Hybrid).unwrap();

            let label = proxy.session_mode_label(SessionMode::Stateful);
            assert_eq!(label, "hybrid");
        }

        #[test]
        fn test_routing_mode_display_name_caching() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::PerCommand).unwrap();

            // Empty cache (default 0 capacity) should return "per-command"
            assert_eq!(proxy.routing_mode_display_name(), "per-command");
        }

        #[test]
        fn test_generate_session_id_format() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            let session = proxy.create_session(
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                None,
            );

            let session_id = proxy.generate_session_id(&session);

            // Should be a short UUID (8 characters)
            assert_eq!(session_id.len(), 8);
        }

        #[test]
        fn test_create_session_without_router() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            let session = proxy.create_session(
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                None,
            );

            // Session should be created successfully
            assert_eq!(session.mode(), SessionMode::Stateful);
        }

        #[test]
        fn test_create_session_with_router() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::PerCommand).unwrap();

            let session = proxy.create_session(
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                Some(proxy.router.clone()),
            );

            // Session mode depends on router presence and config
            // With router, it should be per-command
            assert_eq!(session.mode(), SessionMode::PerCommand);
        }

        #[test]
        fn test_record_connection_if_unauthenticated_no_auth() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            let session = proxy.create_session(
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                None,
            );

            // Should not panic
            proxy.record_connection_if_unauthenticated(&session);

            // Connection should be recorded (we can't easily verify without exposing internals)
        }

        #[test]
        fn test_record_session_metrics_success() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            let session = proxy.create_session(
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                None,
            );
            let session_id = proxy.generate_session_id(&session);

            let metrics = TransferMetrics {
                client_to_backend: crate::types::ClientToBackendBytes::new(1024),
                backend_to_client: crate::types::BackendToClientBytes::new(2048),
            };

            let result = proxy.record_session_metrics(
                Ok(metrics),
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                &session_id,
                &session,
                Some(crate::types::BackendId::from_index(0)),
            );

            assert!(result.is_ok());
        }

        #[test]
        fn test_record_session_metrics_error() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            let session = proxy.create_session(
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                None,
            );
            let session_id = proxy.generate_session_id(&session);

            let result = proxy.record_session_metrics(
                Err(anyhow::anyhow!("test error")),
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap()),
                &session_id,
                &session,
                Some(crate::types::BackendId::from_index(0)),
            );

            assert!(result.is_err());
            assert_eq!(result.unwrap_err().to_string(), "test error");
        }

        #[tokio::test]
        async fn test_prepare_per_command_connection() {
            let config = create_test_config();
            let proxy = Arc::new(
                NntpProxy::new(config, RoutingMode::PerCommand)
                    .await
                    .unwrap(),
            );

            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            // Spawn a simple acceptor that reads greeting
            tokio::spawn(async move {
                let (stream, _) = listener.accept().await.unwrap();
                let mut buf = [0u8; 1024];
                let _ = stream.try_read(&mut buf); // Read greeting
            });

            let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
            let client_addr = ClientAddress::from(stream.peer_addr().unwrap());

            let result = proxy
                .prepare_per_command_connection(&mut stream, client_addr)
                .await;
            assert!(result.is_ok());
        }

        #[test]
        fn test_routing_mode_display_name_empty_cache() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Hybrid).unwrap();

            let _empty_cache = Arc::new(crate::cache::UnifiedCache::memory(
                100,
                std::time::Duration::from_secs(3600),
                false,
            ));
            assert_eq!(proxy.routing_mode_display_name(), "per-command");
        }

        #[test]
        fn test_log_routing_selection() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            let backend_id = crate::types::BackendId::from_index(0);
            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());

            // Should not panic
            proxy.log_routing_selection(client_addr, backend_id, &proxy.servers()[0]);
        }

        #[test]
        fn test_log_pool_status() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            // Should not panic
            proxy.log_pool_status(0);
        }

        #[tokio::test]
        async fn test_apply_tcp_optimizations() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new(config, RoutingMode::Stateful).await.unwrap());

            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            tokio::spawn(async move {
                let (_stream, _) = listener.accept().await.unwrap();
            });

            let stream = tokio::net::TcpStream::connect(addr).await.unwrap();

            // Should not panic
            proxy.apply_tcp_optimizations(&stream);
        }

        #[test]
        fn test_session_mode_labels_all_combinations() {
            use crate::session::SessionMode;

            // PerCommand mode
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config.clone(), RoutingMode::PerCommand).unwrap();
            assert_eq!(
                proxy.session_mode_label(SessionMode::PerCommand),
                "per-command"
            );

            // Stateful mode
            let proxy = NntpProxy::new_sync(config.clone(), RoutingMode::Stateful).unwrap();
            assert_eq!(proxy.session_mode_label(SessionMode::Stateful), "standard");

            // Hybrid mode
            let proxy = NntpProxy::new_sync(config, RoutingMode::Hybrid).unwrap();
            assert_eq!(proxy.session_mode_label(SessionMode::Stateful), "hybrid");
        }

        #[test]
        fn test_finalize_stateful_session_success() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());
            let session = proxy.create_session(client_addr, None);
            let session_id = proxy.generate_session_id(&session);
            let backend_id = crate::types::BackendId::from_index(0);

            let metrics = TransferMetrics {
                client_to_backend: crate::types::ClientToBackendBytes::new(512),
                backend_to_client: crate::types::BackendToClientBytes::new(1024),
            };

            let result = proxy.finalize_stateful_session(
                Ok(metrics),
                client_addr,
                &session_id,
                &session,
                backend_id,
            );

            assert!(result.is_ok());
        }

        #[test]
        fn test_finalize_stateful_session_error() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap());

            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());
            let session = proxy.create_session(client_addr, None);
            let session_id = proxy.generate_session_id(&session);
            let backend_id = crate::types::BackendId::from_index(0);

            let result = proxy.finalize_stateful_session(
                Err(anyhow::anyhow!("connection error")),
                client_addr,
                &session_id,
                &session,
                backend_id,
            );

            assert!(result.is_err());
        }

        #[test]
        fn test_finalize_per_command_session_success() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::PerCommand).unwrap());

            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());
            let session = proxy.create_session(client_addr, Some(proxy.router.clone()));
            let session_id = proxy.generate_session_id(&session);

            let metrics = TransferMetrics {
                client_to_backend: crate::types::ClientToBackendBytes::new(256),
                backend_to_client: crate::types::BackendToClientBytes::new(512),
            };

            let result =
                proxy.finalize_per_command_session(Ok(metrics), client_addr, &session_id, &session);

            assert!(result.is_ok());
        }

        #[test]
        fn test_finalize_per_command_session_error() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::PerCommand).unwrap());

            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());
            let session = proxy.create_session(client_addr, Some(proxy.router.clone()));
            let session_id = proxy.generate_session_id(&session);

            let result = proxy.finalize_per_command_session(
                Err(anyhow::anyhow!("session failed")),
                client_addr,
                &session_id,
                &session,
            );

            assert!(result.is_err());
        }

        #[test]
        fn test_record_session_metrics_without_backend() {
            let config = create_test_config();
            let proxy = Arc::new(NntpProxy::new_sync(config, RoutingMode::PerCommand).unwrap());

            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());
            let session = proxy.create_session(client_addr, Some(proxy.router.clone()));
            let session_id = proxy.generate_session_id(&session);

            let metrics = TransferMetrics {
                client_to_backend: crate::types::ClientToBackendBytes::new(128),
                backend_to_client: crate::types::BackendToClientBytes::new(256),
            };

            // Without backend_id (per-command mode)
            let result =
                proxy.record_session_metrics(Ok(metrics), client_addr, &session_id, &session, None);

            assert!(result.is_ok());
        }

        #[test]
        fn test_generate_session_id_uniqueness() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            let client_addr =
                ClientAddress::from("127.0.0.1:12345".parse::<std::net::SocketAddr>().unwrap());

            let session1 = proxy.create_session(client_addr, None);
            let session2 = proxy.create_session(client_addr, None);

            let id1 = proxy.generate_session_id(&session1);
            let id2 = proxy.generate_session_id(&session2);

            // Should generate different IDs for different sessions
            assert_ne!(id1, id2);
        }
    }

    mod idle_tracking {
        use super::*;
        use std::time::Duration;

        #[test]
        fn test_idle_timeout_constant() {
            // Ensure the idle timeout is 5 minutes
            assert_eq!(NntpProxy::IDLE_TIMEOUT, Duration::from_secs(5 * 60));
        }

        #[test]
        fn test_active_clients_increment_decrement() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            assert_eq!(proxy.active_clients.load(Ordering::Relaxed), 0);

            proxy.increment_active_clients();
            assert_eq!(proxy.active_clients.load(Ordering::Relaxed), 1);

            proxy.increment_active_clients();
            assert_eq!(proxy.active_clients.load(Ordering::Relaxed), 2);

            proxy.decrement_active_clients();
            assert_eq!(proxy.active_clients.load(Ordering::Relaxed), 1);

            proxy.decrement_active_clients();
            assert_eq!(proxy.active_clients.load(Ordering::Relaxed), 0);
        }

        #[test]
        fn test_last_activity_updated_on_last_client_disconnect() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            // Initially no activity
            assert_eq!(proxy.last_activity_nanos.load(Ordering::Relaxed), 0);

            // Connect two clients
            proxy.increment_active_clients();
            proxy.increment_active_clients();

            // First client disconnects - should not update timestamp
            proxy.decrement_active_clients();
            assert_eq!(proxy.last_activity_nanos.load(Ordering::Relaxed), 0);

            // Second (last) client disconnects - should update timestamp
            proxy.decrement_active_clients();
            assert!(proxy.last_activity_nanos.load(Ordering::Relaxed) > 0);
        }

        #[test]
        fn test_check_and_clear_skips_when_clients_active() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            // Simulate a past activity
            proxy.last_activity_nanos.store(1, Ordering::Relaxed);

            // With active clients, should not clear
            proxy.increment_active_clients();
            let cleared = proxy.check_and_clear_stale_pools();
            assert!(!cleared);
        }

        #[test]
        fn test_check_and_clear_skips_when_never_active() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            // No prior activity (last_activity_nanos = 0)
            let cleared = proxy.check_and_clear_stale_pools();
            assert!(!cleared);
        }

        #[test]
        fn test_check_and_clear_skips_when_recently_active() {
            let config = create_test_config();
            let proxy = NntpProxy::new_sync(config, RoutingMode::Stateful).unwrap();

            // Connect and disconnect to set timestamp
            proxy.increment_active_clients();
            proxy.decrement_active_clients();

            // Should not clear - just disconnected (within timeout)
            let cleared = proxy.check_and_clear_stale_pools();
            assert!(!cleared);
        }
    }
}