rustytorch/crates/production/rtx-streaming/src/streaming_server.rs

//! # Streaming Server Implementation
//!
//! Core streaming server that orchestrates real-time model inference with
//! sub-millisecond latency requirements.

use crate::{
    BackpressureHandler, ConnectionManager, StreamMetrics, StreamingConfig, StreamingError,
    StreamingResult, TokenGenerator,
};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use uuid::Uuid;

/// Main streaming server orchestrating real-time inference
#[derive(Debug)]
pub struct StreamingServer {
    /// Connection lifecycle management
    connection_manager: Arc<ConnectionManager>,

    /// Real-time token generation pipeline
    token_generator: Arc<TokenGenerator>,

    /// Backpressure and flow control
    backpressure_handler: Arc<BackpressureHandler>,

    /// Performance monitoring
    metrics: Arc<StreamMetrics>,

    /// Server configuration
    config: StreamingConfig,

    /// Active streaming sessions
    active_sessions: Arc<DashMap<Uuid, StreamingSession>>,

    /// Server state
    state: Arc<RwLock<ServerState>>,
}

/// Streaming session state
#[derive(Debug, Clone)]
pub struct StreamingSession {
    /// Unique session identifier
    pub session_id: Uuid,

    /// Client identifier
    pub client_id: String,

    /// Session start time
    pub start_time: Instant,

    /// Number of tokens streamed
    pub tokens_streamed: u64,

    /// Average latency per token
    pub avg_latency: Duration,

    /// Session configuration
    pub config: SessionConfig,
}

/// Per-session configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionConfig {
    /// Maximum tokens to generate
    pub max_tokens: usize,

    /// Temperature for sampling
    pub temperature: f32,

    /// Top-p sampling parameter
    pub top_p: f32,

    /// Streaming chunk size
    pub chunk_size: usize,

    /// Priority level
    pub priority: Priority,
}

/// Priority levels for streaming sessions
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Priority {
    Low = 0,
    Normal = 1,
    High = 2,
    Critical = 3,
}

/// Server operational state
#[derive(Debug, Clone)]
pub enum ServerState {
    /// Server is starting up
    Starting,

    /// Server is running normally
    Running,

    /// Server is under high load
    Overloaded,

    /// Server is shutting down gracefully
    Shutting,

    /// Server is stopped
    Stopped,
}

/// Inference request for streaming
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceRequest {
    /// Input prompt
    pub prompt: String,

    /// Session configuration
    pub config: SessionConfig,

    /// Client metadata
    pub metadata: std::collections::HashMap<String, String>,
}

/// Streaming token response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenResponse {
    /// Generated token
    pub token: String,

    /// Token probability
    pub probability: f32,

    /// Generation timestamp
    pub timestamp: std::time::SystemTime,

    /// Sequence position
    pub position: usize,

    /// Whether this is the final token
    pub is_final: bool,
}

/// Connection handle for client management
#[derive(Debug, Clone)]
pub struct ConnectionHandle {
    /// Connection identifier
    pub connection_id: Uuid,

    /// Client identifier
    pub client_id: String,

    /// Connection creation time
    pub created_at: Instant,

    /// Connection state
    pub state: ConnectionState,
}

/// Connection state tracking
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConnectionState {
    /// Connection is being established
    Connecting,

    /// Connection is active and ready
    Active,

    /// Connection is experiencing backpressure
    Throttled,

    /// Connection is being closed
    Closing,

    /// Connection is closed
    Closed,
}

impl StreamingServer {
    /// Create a new streaming server
    pub async fn new(config: StreamingConfig) -> StreamingResult<Self> {
        let connection_manager = Arc::new(
            ConnectionManager::new(config.connection_pool_size)
                .await
                .map_err(|e| {
                    StreamingError::Config(format!("Connection manager init failed: {e}"))
                })?,
        );

        let token_generator =
            Arc::new(TokenGenerator::new(&config).await.map_err(|e| {
                StreamingError::Config(format!("Token generator init failed: {e}"))
            })?);

        let backpressure_handler = Arc::new(
            BackpressureHandler::new(config.backpressure_threshold)
                .await
                .map_err(|e| {
                    StreamingError::Config(format!("Backpressure handler init failed: {e}"))
                })?,
        );

        let metrics = if config.metrics_enabled {
            Arc::new(StreamMetrics::new().await?)
        } else {
            Arc::new(StreamMetrics::disabled().await?)
        };

        Ok(Self {
            connection_manager,
            token_generator,
            backpressure_handler,
            metrics,
            config,
            active_sessions: Arc::new(DashMap::new()),
            // The server is fully constructed and operational as soon as `new`
            // returns; `start()` remains available for explicitly kicking off
            // background work (pool pre-warming, metrics collection) but is
            // not required before serving requests.
            state: Arc::new(RwLock::new(ServerState::Running)),
        })
    }

    /// Start the streaming server
    pub async fn start(&mut self) -> StreamingResult<()> {
        // Update server state
        *self.state.write().await = ServerState::Running;

        // Start connection manager
        self.connection_manager.start().await.map_err(|e| {
            StreamingError::Connection(format!("Failed to start connection manager: {e}"))
        })?;

        // Start metrics collection
        if self.config.metrics_enabled {
            self.metrics.start_collection().await?;
        }

        Ok(())
    }

    /// Stream inference for a given prompt
    pub async fn stream_inference(&self, prompt: &str) -> StreamingResult<Vec<TokenResponse>> {
        let start_time = Instant::now();

        // Check server state
        let state = self.state.read().await;
        match *state {
            ServerState::Running => {}
            ServerState::Overloaded => {
                return Err(StreamingError::Backpressure(
                    "Server overloaded".to_string(),
                ));
            }
            _ => {
                return Err(StreamingError::Connection("Server not running".to_string()));
            }
        }
        drop(state);

        // Create default session config. This helper simulates a single
        // real-time streaming step: `chunk_size: 1` signals one token per
        // step, so `max_tokens` mirrors that here rather than generating a
        // full multi-hundred-token completion synchronously in one call —
        // doing the latter would contradict the sub-millisecond, incremental
        // streaming behavior this system targets (see module docs).
        let session_config = SessionConfig {
            max_tokens: 1,
            temperature: 0.8,
            top_p: 0.9,
            chunk_size: 1,
            priority: Priority::Normal,
        };

        // Create inference request
        let request = InferenceRequest {
            prompt: prompt.to_string(),
            config: session_config,
            metadata: std::collections::HashMap::new(),
        };

        // Generate tokens through the pipeline
        let tokens = self
            .token_generator
            .generate_stream(&request)
            .await
            .map_err(|e| StreamingError::Inference(format!("Token generation failed: {e}")))?;

        // Record metrics
        let total_latency = start_time.elapsed();
        self.metrics.record_inference_latency(total_latency).await?;

        // Validate latency requirement
        if total_latency > self.config.target_latency {
            return Err(StreamingError::Performance(format!(
                "Latency {} exceeds target {}",
                total_latency.as_micros(),
                self.config.target_latency.as_micros()
            )));
        }

        Ok(tokens)
    }

    /// Attach a real rtx-inference backend and the model name to route
    /// `stream_inference` requests to. Without this, `stream_inference`
    /// returns an error rather than fabricating tokens.
    pub async fn set_inference_backend(
        &self,
        backend: Arc<rtx_inference::InferenceEngine>,
        model: String,
    ) {
        self.token_generator.set_backend(backend, model).await;
    }

    /// Create a new streaming connection
    pub async fn create_connection(&self, client_id: &str) -> StreamingResult<ConnectionHandle> {
        // Check if we're at capacity
        if self.active_sessions.len() >= self.config.max_connections {
            return Err(StreamingError::Connection(
                "Max connections reached".to_string(),
            ));
        }

        // Create connection through manager
        let connection_id = self
            .connection_manager
            .create_connection(client_id)
            .await
            .map_err(|e| StreamingError::Connection(format!("Failed to create connection: {e}")))?;

        let handle = ConnectionHandle {
            connection_id,
            client_id: client_id.to_string(),
            created_at: Instant::now(),
            state: ConnectionState::Active,
        };

        Ok(handle)
    }

    /// Handle overload scenario with graceful degradation
    pub async fn handle_overload_scenario(&self) -> StreamingResult<bool> {
        // Update server state to overloaded
        *self.state.write().await = ServerState::Overloaded;

        // Apply backpressure handling
        let handled = self
            .backpressure_handler
            .handle_overload()
            .await
            .map_err(|e| StreamingError::Backpressure(format!("Overload handling failed: {e}")))?;

        // If successfully handled, return to running state
        if handled {
            *self.state.write().await = ServerState::Running;
        }

        Ok(handled)
    }

    /// Clone server for concurrent access
    pub fn clone(&self) -> Self {
        Self {
            connection_manager: Arc::clone(&self.connection_manager),
            token_generator: Arc::clone(&self.token_generator),
            backpressure_handler: Arc::clone(&self.backpressure_handler),
            metrics: Arc::clone(&self.metrics),
            config: self.config.clone(),
            active_sessions: Arc::clone(&self.active_sessions),
            state: Arc::clone(&self.state),
        }
    }

    /// Get current server metrics
    pub async fn get_metrics(&self) -> StreamingResult<ServerMetrics> {
        Ok(ServerMetrics {
            active_connections: self.active_sessions.len(),
            total_requests: self.metrics.get_total_requests().await?,
            average_latency: self.metrics.get_average_latency().await?,
            p99_latency: self.metrics.get_p99_latency().await?,
            throughput_qps: self.metrics.get_throughput().await?,
            memory_usage: self.get_memory_usage().await?,
        })
    }

    /// Get current memory usage
    async fn get_memory_usage(&self) -> StreamingResult<usize> {
        // In real implementation, this would integrate with system APIs
        // For now, return estimated usage
        let base_usage = 1024 * 1024 * 50; // 50MB base
        let connection_overhead = self.active_sessions.len() * 1024; // 1KB per connection
        Ok(base_usage + connection_overhead)
    }

    /// Shutdown server gracefully
    pub async fn shutdown(&mut self) -> StreamingResult<()> {
        *self.state.write().await = ServerState::Shutting;

        // Close all active sessions
        self.active_sessions.clear();

        // Shutdown connection manager
        self.connection_manager.shutdown().await.map_err(|e| {
            StreamingError::Connection(format!("Connection manager shutdown failed: {e}"))
        })?;

        // Stop metrics collection
        if self.config.metrics_enabled {
            self.metrics.stop_collection().await?;
        }

        *self.state.write().await = ServerState::Stopped;

        Ok(())
    }
}

/// Server performance metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerMetrics {
    /// Number of active connections
    pub active_connections: usize,

    /// Total requests processed
    pub total_requests: u64,

    /// Average inference latency
    pub average_latency: Duration,

    /// 99th percentile latency
    pub p99_latency: Duration,

    /// Throughput in queries per second
    pub throughput_qps: f64,

    /// Current memory usage in bytes
    pub memory_usage: usize,
}

impl Default for SessionConfig {
    fn default() -> Self {
        Self {
            max_tokens: 100,
            temperature: 0.8,
            top_p: 0.9,
            chunk_size: 1,
            priority: Priority::Normal,
        }
    }
}