rusty_common/memory/

trading_pools.rs

//! Trading-specific memory pools for zero-allocation hot paths
//!
//! This module provides specialized memory pool integrations designed for
//! high-frequency trading operations, eliminating allocations in critical paths
//! such as order processing, WebSocket message handling, and market data processing.

use super::zerocopy_pools::{
    BufferHandle, GlobalPoolManager, TypedMemoryPool, ZerocopyMemoryPool, ZerocopyPoolConfig,
};
use smallvec::SmallVec;
use smartstring::alias::String as SmartString;
use std::sync::Arc;

/// Trading-specific memory pool manager for zero-allocation operations with const generic buffer sizing
pub struct TradingPoolManager<const N: usize = 4096, const M: usize = 8> {
    /// Global pool manager for basic buffer operations
    global_pools: GlobalPoolManager,

    /// Specialized pools for trading operations
    websocket_message_pool: Arc<WebSocketMessagePool<N>>,
    order_processing_pool: Arc<OrderProcessingPool<M>>,
    market_data_pool: Arc<MarketDataPool>,
    json_processing_pool: Arc<JsonProcessingPool>,
}

impl<const N: usize, const M: usize> TradingPoolManager<N, M> {
    /// Create a new trading pool manager with optimized configurations
    pub fn new() -> Self {
        let global_pools = GlobalPoolManager::new();

        Self {
            websocket_message_pool: Arc::new(WebSocketMessagePool::new(
                global_pools.small_buffer_pool().clone(),
                global_pools.general_pool().clone(),
            )),
            order_processing_pool: Arc::new(OrderProcessingPool::new()),
            market_data_pool: Arc::new(MarketDataPool::new(
                global_pools.general_pool().clone(),
                global_pools.simd_pool().clone(),
            )),
            json_processing_pool: Arc::new(JsonProcessingPool::new(
                global_pools.small_buffer_pool().clone(),
                global_pools.general_pool().clone(),
            )),
            global_pools,
        }
    }

    /// Get WebSocket message pool
    pub const fn websocket_pool(&self) -> &Arc<WebSocketMessagePool<N>> {
        &self.websocket_message_pool
    }

    /// Get order processing pool
    pub const fn order_pool(&self) -> &Arc<OrderProcessingPool<M>> {
        &self.order_processing_pool
    }

    /// Get market data pool
    pub const fn market_data_pool(&self) -> &Arc<MarketDataPool> {
        &self.market_data_pool
    }

    /// Get JSON processing pool
    pub const fn json_pool(&self) -> &Arc<JsonProcessingPool> {
        &self.json_processing_pool
    }

    /// Get underlying global pools for direct access
    pub const fn global_pools(&self) -> &GlobalPoolManager {
        &self.global_pools
    }
}

impl<const N: usize, const M: usize> Default for TradingPoolManager<N, M> {
    fn default() -> Self {
        Self::new()
    }
}

/// WebSocket message pool for zero-allocation message processing with const generic threshold
pub struct WebSocketMessagePool<const N: usize = 4096> {
    small_buffer_pool: Arc<ZerocopyMemoryPool>, // For small messages (< N bytes)
    large_buffer_pool: Arc<ZerocopyMemoryPool>, // For large messages (< 64KB)
}

impl<const N: usize> WebSocketMessagePool<N> {
    /// Create a new WebSocket message pool with const generic threshold
    pub const fn new(
        small_pool: Arc<ZerocopyMemoryPool>,
        large_pool: Arc<ZerocopyMemoryPool>,
    ) -> Self {
        Self {
            small_buffer_pool: small_pool,
            large_buffer_pool: large_pool,
        }
    }

    /// Get a buffer for WebSocket message processing using const generic threshold
    pub fn get_message_buffer(&self, estimated_size: usize) -> Option<BufferHandle> {
        if estimated_size <= N {
            self.small_buffer_pool.allocate_buffer()
        } else {
            self.large_buffer_pool.allocate_buffer()
        }
    }

    /// Get a buffer for binary WebSocket messages
    pub fn get_binary_buffer(&self) -> Option<BufferHandle> {
        self.large_buffer_pool.allocate_buffer()
    }

    /// Get a buffer for text WebSocket messages
    pub fn get_text_buffer(&self) -> Option<BufferHandle> {
        self.small_buffer_pool.allocate_buffer()
    }
}

/// Order processing pool for zero-allocation order operations with const generic SmallVec capacity
pub struct OrderProcessingPool<const M: usize = 8> {
    // Note: ExecutionReport and Order types would need to be imported from appropriate crate
    // For now, using generic buffer and string pools
    report_buffer_pool: Arc<ZerocopyMemoryPool>,
    order_buffer_pool: Arc<ZerocopyMemoryPool>,
    string_pool: Arc<TypedMemoryPool<SmartString>>,
    order_ids: Arc<TypedMemoryPool<SmallVec<[SmartString; M]>>>,
}

impl<const M: usize> OrderProcessingPool<M> {
    /// Create a new order processing pool with const generic SmallVec capacity
    pub fn new() -> Self {
        let report_config = ZerocopyPoolConfig {
            initial_buffer_count: 128,
            max_buffer_count: 1024,
            buffer_size: 4 * 1024, // 4KB for execution reports
            simd_aligned: false,
            enable_statistics: true,
        };

        let order_config = ZerocopyPoolConfig {
            initial_buffer_count: 64,
            max_buffer_count: 512,
            buffer_size: 16 * 1024, // 16KB for order collections
            simd_aligned: false,
            enable_statistics: true,
        };

        let string_config = ZerocopyPoolConfig {
            initial_buffer_count: 256,
            max_buffer_count: 2048,
            buffer_size: std::mem::size_of::<SmartString>(),
            simd_aligned: false,
            enable_statistics: true,
        };

        let order_id_config = ZerocopyPoolConfig {
            initial_buffer_count: 128,
            max_buffer_count: 1024,
            buffer_size: std::mem::size_of::<SmallVec<[SmartString; 8]>>(),
            simd_aligned: false,
            enable_statistics: true,
        };

        Self {
            report_buffer_pool: ZerocopyMemoryPool::new(report_config),
            order_buffer_pool: ZerocopyMemoryPool::new(order_config),
            string_pool: TypedMemoryPool::new(string_config),
            order_ids: TypedMemoryPool::new(order_id_config),
        }
    }

    /// Get a buffer for execution report processing
    pub fn get_execution_report_buffer(&self) -> Option<BufferHandle> {
        self.report_buffer_pool.allocate_buffer()
    }

    /// Get a buffer for order collection processing
    pub fn get_order_collection_buffer(&self) -> Option<BufferHandle> {
        self.order_buffer_pool.allocate_buffer()
    }

    /// Get a pooled string
    pub fn get_string(
        &self,
    ) -> Option<crate::memory::zerocopy_pools::TypedBufferHandle<SmartString>> {
        self.string_pool.allocate_typed()
    }

    /// Get a pooled order ID collection with const generic capacity
    pub fn get_order_ids(
        &self,
    ) -> Option<crate::memory::zerocopy_pools::TypedBufferHandle<SmallVec<[SmartString; M]>>> {
        self.order_ids.allocate_typed()
    }
}

impl<const M: usize> Default for OrderProcessingPool<M> {
    fn default() -> Self {
        Self::new()
    }
}

/// Market data pool for zero-allocation market data processing
pub struct MarketDataPool {
    // Note: OrderBook and Trade types would need to be defined or imported
    // For now, using generic buffer pools
    orderbook_buffer_pool: Arc<ZerocopyMemoryPool>,
    trade_buffer_pool: Arc<ZerocopyMemoryPool>,
    general_buffer_pool: Arc<ZerocopyMemoryPool>,
    simd_pool: Arc<crate::memory::zerocopy_pools::SimdMemoryPool>,
}

impl MarketDataPool {
    /// Create a new market data pool
    pub fn new(
        general_pool: Arc<ZerocopyMemoryPool>,
        simd_pool: Arc<crate::memory::zerocopy_pools::SimdMemoryPool>,
    ) -> Self {
        let orderbook_config = ZerocopyPoolConfig {
            initial_buffer_count: 32,
            max_buffer_count: 256,
            buffer_size: 32 * 1024, // 32KB for orderbook data
            simd_aligned: true,
            enable_statistics: true,
        };

        let trade_config = ZerocopyPoolConfig {
            initial_buffer_count: 64,
            max_buffer_count: 512,
            buffer_size: 16 * 1024, // 16KB for trade collections
            simd_aligned: false,
            enable_statistics: true,
        };

        Self {
            orderbook_buffer_pool: ZerocopyMemoryPool::new(orderbook_config),
            trade_buffer_pool: ZerocopyMemoryPool::new(trade_config),
            general_buffer_pool: general_pool,
            simd_pool,
        }
    }

    /// Get a buffer for orderbook processing
    pub fn get_orderbook_buffer(&self) -> Option<BufferHandle> {
        self.orderbook_buffer_pool.allocate_buffer()
    }

    /// Get a buffer for trade collection processing
    pub fn get_trade_buffer(&self) -> Option<BufferHandle> {
        self.trade_buffer_pool.allocate_buffer()
    }

    /// Get a general buffer for market data processing
    pub fn get_processing_buffer(&self) -> Option<BufferHandle> {
        self.general_buffer_pool.allocate_buffer()
    }

    /// Get a SIMD buffer for vectorized market data operations
    pub fn get_simd_buffer(&self) -> Option<crate::memory::zerocopy_pools::SimdBufferHandle> {
        self.simd_pool.allocate_simd_buffer()
    }
}

/// JSON processing pool for zero-allocation JSON operations
pub struct JsonProcessingPool {
    small_buffer_pool: Arc<ZerocopyMemoryPool>, // For JSON strings < 4KB
    large_buffer_pool: Arc<ZerocopyMemoryPool>, // For JSON strings < 64KB
    serialization_threshold: usize,
}

impl JsonProcessingPool {
    /// Create a new JSON processing pool
    pub const fn new(
        small_pool: Arc<ZerocopyMemoryPool>,
        large_pool: Arc<ZerocopyMemoryPool>,
    ) -> Self {
        Self {
            small_buffer_pool: small_pool,
            large_buffer_pool: large_pool,
            serialization_threshold: 4 * 1024, // 4KB threshold
        }
    }

    /// Get a buffer for JSON serialization
    pub fn get_serialization_buffer(&self, estimated_size: usize) -> Option<BufferHandle> {
        if estimated_size <= self.serialization_threshold {
            self.small_buffer_pool.allocate_buffer()
        } else {
            self.large_buffer_pool.allocate_buffer()
        }
    }

    /// Get a buffer for JSON parsing
    pub fn get_parsing_buffer(&self) -> Option<BufferHandle> {
        self.large_buffer_pool.allocate_buffer()
    }

    /// Get a buffer for JSON value storage (generic buffer)
    pub fn get_value_buffer(&self) -> Option<BufferHandle> {
        self.small_buffer_pool.allocate_buffer()
    }
}

/// Helper trait for zero-allocation WebSocket message processing
pub trait PooledWebSocketMessage {
    /// Create a message using pooled buffers
    fn from_pooled_buffer(
        pool: &WebSocketMessagePool,
        message_type: WebSocketMessageType,
        data: &[u8],
    ) -> Option<Self>
    where
        Self: Sized;
}

/// WebSocket message types for pool optimization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WebSocketMessageType {
    /// Text-based WebSocket message (typically JSON)
    Text,
    /// Binary WebSocket message (raw bytes)
    Binary,
    /// Ping control frame for connection keep-alive
    Ping,
    /// Pong control frame in response to ping
    Pong,
    /// Close control frame to terminate connection
    Close,
}

/// Helper trait for zero-allocation order processing
/// Note: Concrete implementations should be done in crates that have access to order types
pub trait PooledOrderProcessing {
    /// Get a buffer for order processing operations
    fn get_order_processing_buffer(
        pool: &OrderProcessingPool,
        buffer_type: OrderBufferType,
    ) -> Option<BufferHandle>;
}

/// Types of order processing buffers
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrderBufferType {
    /// Buffer for processing execution reports from exchanges
    ExecutionReport,
    /// Buffer for storing collections of orders
    OrderCollection,
    /// Buffer for batch processing of order identifiers
    OrderIdBatch,
}

/// Default implementation for buffer allocation
pub struct DefaultOrderProcessing;

impl PooledOrderProcessing for DefaultOrderProcessing {
    fn get_order_processing_buffer(
        pool: &OrderProcessingPool,
        buffer_type: OrderBufferType,
    ) -> Option<BufferHandle> {
        match buffer_type {
            OrderBufferType::ExecutionReport => pool.get_execution_report_buffer(),
            OrderBufferType::OrderCollection => pool.get_order_collection_buffer(),
            OrderBufferType::OrderIdBatch => {
                // For order ID batches, use execution report buffer as they're similar size
                pool.get_execution_report_buffer()
            }
        }
    }
}

// Thread-local pool manager for maximum performance
thread_local! {
    static THREAD_LOCAL_POOLS: DefaultTradingPoolManager = DefaultTradingPoolManager::new();
}

/// Get thread-local trading pools for maximum performance
///
/// # Safety Note
///
/// This function uses the `thread_local!` macro to safely access thread-local pools.
/// The returned reference is valid for the lifetime of the closure where it's used.
/// Each thread maintains its own independent instance of `TradingPoolManager`.
pub fn with_thread_local_pools<F, R>(f: F) -> R
where
    F: FnOnce(&TradingPoolManager) -> R,
{
    THREAD_LOCAL_POOLS.with(|pools| f(pools))
}

/// Get thread-local trading pools (legacy API for compatibility)
///
/// # Deprecated
///
/// This function is deprecated due to lifetime safety concerns.
/// Use `with_thread_local_pools` instead.
#[deprecated(since = "0.1.0", note = "Use with_thread_local_pools for safer access")]
pub fn get_thread_local_pools() -> &'static DefaultTradingPoolManager {
    // Return the global pools as a safe fallback
    // This maintains API compatibility while being lifetime-safe
    get_global_trading_pools()
}

/// Global trading pool manager instance
static GLOBAL_TRADING_POOLS: std::sync::OnceLock<DefaultTradingPoolManager> =
    std::sync::OnceLock::new();

/// Get global trading pools (use for initialization and statistics)
pub fn get_global_trading_pools() -> &'static DefaultTradingPoolManager {
    GLOBAL_TRADING_POOLS.get_or_init(DefaultTradingPoolManager::new)
}

/// Initialize trading pools with custom configuration
pub fn initialize_trading_pools() -> &'static DefaultTradingPoolManager {
    get_global_trading_pools()
}

// Type aliases for backward compatibility
/// Trading pool manager with 4KB WebSocket threshold and 8-element SmallVec capacity
pub type TradingPoolManager4096 = TradingPoolManager<4096, 8>;
/// Trading pool manager with 2KB WebSocket threshold and 16-element SmallVec capacity
pub type TradingPoolManager2048 = TradingPoolManager<2048, 16>;
/// Trading pool manager with 8KB WebSocket threshold and 4-element SmallVec capacity
pub type TradingPoolManager8192 = TradingPoolManager<8192, 4>;

/// WebSocket message pool with 4KB threshold for small messages
pub type WebSocketMessagePool4096 = WebSocketMessagePool<4096>;
/// WebSocket message pool with 2KB threshold for small messages
pub type WebSocketMessagePool2048 = WebSocketMessagePool<2048>;
/// WebSocket message pool with 8KB threshold for small messages
pub type WebSocketMessagePool8192 = WebSocketMessagePool<8192>;

/// Order processing pool with SmallVec capacity of 8 elements
pub type OrderProcessingPool8 = OrderProcessingPool<8>;
/// Order processing pool with SmallVec capacity of 16 elements
pub type OrderProcessingPool16 = OrderProcessingPool<16>;
/// Order processing pool with SmallVec capacity of 4 elements
pub type OrderProcessingPool4 = OrderProcessingPool<4>;

// Default type aliases for seamless migration
pub use OrderProcessingPool8 as DefaultOrderProcessingPool;
pub use TradingPoolManager4096 as DefaultTradingPoolManager;
pub use WebSocketMessagePool4096 as DefaultWebSocketMessagePool;

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

    #[test]
    fn test_trading_pool_manager_creation() {
        let manager = DefaultTradingPoolManager::new();

        // Test that all pools are created
        assert!(manager.websocket_pool().get_text_buffer().is_some());
        assert!(manager.order_pool().get_execution_report_buffer().is_some());
        assert!(manager.market_data_pool().get_processing_buffer().is_some());
        assert!(manager.json_pool().get_parsing_buffer().is_some());
    }

    #[test]
    fn test_websocket_message_pool() {
        let manager = DefaultTradingPoolManager::new();
        let ws_pool = manager.websocket_pool();

        // Test small message buffer
        let small_buffer = ws_pool.get_message_buffer(1024);
        assert!(small_buffer.is_some());
        assert_eq!(small_buffer.unwrap().size(), 4 * 1024);

        // Test large message buffer
        let large_buffer = ws_pool.get_message_buffer(8 * 1024);
        assert!(large_buffer.is_some());
        assert_eq!(large_buffer.unwrap().size(), 64 * 1024);

        // Test specific buffer types
        assert!(ws_pool.get_binary_buffer().is_some());
        assert!(ws_pool.get_text_buffer().is_some());
    }

    #[test]
    fn test_order_processing_pool() {
        let manager = DefaultTradingPoolManager::new();
        let order_pool = manager.order_pool();

        // Test execution report buffer allocation
        let report_buffer = order_pool.get_execution_report_buffer().unwrap();
        assert_eq!(report_buffer.size(), 4 * 1024);

        // Test other allocations
        assert!(order_pool.get_order_collection_buffer().is_some());
        assert!(order_pool.get_string().is_some());
        assert!(order_pool.get_order_ids().is_some());
    }

    #[test]
    fn test_market_data_pool() {
        let manager = DefaultTradingPoolManager::new();
        let market_pool = manager.market_data_pool();

        // Test orderbook buffer allocation
        let orderbook_buffer = market_pool.get_orderbook_buffer().unwrap();
        assert_eq!(orderbook_buffer.size(), 32 * 1024);

        // Test other allocations
        assert!(market_pool.get_trade_buffer().is_some());
        assert!(market_pool.get_processing_buffer().is_some());
        assert!(market_pool.get_simd_buffer().is_some());
    }

    #[test]
    fn test_json_processing_pool() {
        let manager = DefaultTradingPoolManager::new();
        let json_pool = manager.json_pool();

        // Test buffer allocation based on size
        let small_buffer = json_pool.get_serialization_buffer(1024);
        assert!(small_buffer.is_some());
        assert_eq!(small_buffer.unwrap().size(), 4 * 1024);

        let large_buffer = json_pool.get_serialization_buffer(8 * 1024);
        assert!(large_buffer.is_some());
        assert_eq!(large_buffer.unwrap().size(), 64 * 1024);

        // Test other allocations
        assert!(json_pool.get_parsing_buffer().is_some());
        assert!(json_pool.get_value_buffer().is_some());
    }

    #[test]
    fn test_pooled_order_processing() {
        let manager = DefaultTradingPoolManager::new();
        let order_pool = manager.order_pool();

        // Test different buffer types
        let exec_buffer = DefaultOrderProcessing::get_order_processing_buffer(
            order_pool,
            OrderBufferType::ExecutionReport,
        )
        .unwrap();
        assert_eq!(exec_buffer.size(), 4 * 1024);

        let order_buffer = DefaultOrderProcessing::get_order_processing_buffer(
            order_pool,
            OrderBufferType::OrderCollection,
        )
        .unwrap();
        assert_eq!(order_buffer.size(), 16 * 1024);
    }

    #[test]
    fn test_thread_local_pools() {
        // Test with the new closure-based API
        with_thread_local_pools(|pools1| {
            with_thread_local_pools(|pools2| {
                // Should be the same instance within the same thread
                assert!(std::ptr::eq(pools1, pools2));
            });

            // Should be able to allocate from thread-local pools
            assert!(pools1.websocket_pool().get_text_buffer().is_some());
        });
    }

    #[test]
    fn test_global_pools() {
        let pools1 = get_global_trading_pools();
        let pools2 = get_global_trading_pools();

        // Should be the same instance globally
        assert!(std::ptr::eq(pools1, pools2));

        // Should be able to allocate from global pools
        assert!(pools1.order_pool().get_execution_report_buffer().is_some());
    }

    #[test]
    fn test_websocket_message_type_enum() {
        use WebSocketMessageType::*;

        let types = [Text, Binary, Ping, Pong, Close];

        // Test that all types are distinct
        for (i, &type1) in types.iter().enumerate() {
            for (j, &type2) in types.iter().enumerate() {
                if i != j {
                    assert_ne!(type1, type2);
                }
            }
        }
    }

    #[test]
    fn test_pool_buffer_reuse() {
        let manager = DefaultTradingPoolManager::new();
        let ws_pool = manager.websocket_pool();

        // Allocate and drop a buffer
        let _buffer1_ptr = {
            let buffer = ws_pool.get_text_buffer().unwrap();
            buffer.as_ptr()
        }; // Buffer is dropped here

        // Allocate another buffer - should potentially reuse the same memory
        let buffer2 = ws_pool.get_text_buffer().unwrap();

        // The test passes if we can allocate successfully (reuse is implementation detail)
        assert_eq!(buffer2.size(), 4 * 1024);
    }
}