hft_pipeline.py

"""
High-Frequency Trading Data Pipeline
=====================================
A high-performance Python framework for processing market data with sub-millisecond latency.
Designed to handle 10,000+ records/minute with 99.9% uptime using asyncio and multiprocessing.

Author: [Your Name]
Date: November 2024
"""

import asyncio
import multiprocessing as mp
from multiprocessing import Queue, Manager
import time
import json
import random
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field, asdict
from collections import deque, defaultdict
import numpy as np
from enum import Enum
import logging
import signal
import sys
import heapq
from concurrent.futures import ProcessPoolExecutor
import threading

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s.%(msecs)03d - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# ============================================================================
# Data Models
# ============================================================================

class OrderType(Enum):
    """Order types for the order book"""
    LIMIT = "LIMIT"
    MARKET = "MARKET"
    STOP = "STOP"
    STOP_LIMIT = "STOP_LIMIT"

class Side(Enum):
    """Order side (buy/sell)"""
    BUY = "BUY"
    SELL = "SELL"

class AssetClass(Enum):
    """Asset classes supported"""
    EQUITY = "EQUITY"
    DERIVATIVE = "DERIVATIVE"
    OPTION = "OPTION"
    FUTURE = "FUTURE"

@dataclass
class MarketData:
    """Market data message structure"""
    timestamp: float
    symbol: str
    asset_class: str
    bid_price: float
    ask_price: float
    bid_size: int
    ask_size: int
    last_price: float
    volume: int
    sequence_num: int
    
    def to_dict(self) -> dict:
        return asdict(self)

@dataclass
class Order:
    """Order structure for order book"""
    order_id: str
    timestamp: float
    symbol: str
    side: str
    order_type: str
    price: float
    quantity: int
    remaining_quantity: int = field(init=False)
    
    def __post_init__(self):
        self.remaining_quantity = self.quantity
    
    def __lt__(self, other):
        """For priority queue ordering"""
        if self.side == Side.BUY.value:
            return self.price > other.price  # Higher price = higher priority for buys
        return self.price < other.price  # Lower price = higher priority for sells

@dataclass
class Trade:
    """Executed trade structure"""
    trade_id: str
    timestamp: float
    symbol: str
    price: float
    quantity: int
    buyer_order_id: str
    seller_order_id: str
    
# ============================================================================
# Order Book Engine
# ============================================================================

class OrderBook:
    """
    High-performance order book implementation with O(log n) operations
    """
    def __init__(self, symbol: str):
        self.symbol = symbol
        self.bids = []  # Max heap (using negative prices)
        self.asks = []  # Min heap
        self.orders = {}  # order_id -> Order mapping
        self.last_trade_price = None
        self.last_update_time = time.time()
        self.trade_count = 0
        
    def add_order(self, order: Order) -> Optional[List[Trade]]:
        """Add order to book and attempt matching"""
        trades = []
        
        if order.order_type == OrderType.MARKET.value:
            trades = self._match_market_order(order)
        else:
            trades = self._match_limit_order(order)
            if order.remaining_quantity > 0:
                self._add_to_book(order)
        
        self.last_update_time = time.time()
        return trades
    
    def _match_market_order(self, order: Order) -> List[Trade]:
        """Match market order against opposite side"""
        trades = []
        opposite_book = self.asks if order.side == Side.BUY.value else self.bids
        
        while opposite_book and order.remaining_quantity > 0:
            best_order = heapq.heappop(opposite_book)
            if order.side == Side.BUY.value:
                best_order.price = -best_order.price  # Restore actual price
            
            trade = self._execute_trade(order, best_order)
            trades.append(trade)
            
            if best_order.remaining_quantity > 0:
                heapq.heappush(opposite_book, best_order)
        
        return trades
    
    def _match_limit_order(self, order: Order) -> List[Trade]:
        """Match limit order if crossable"""
        trades = []
        opposite_book = self.asks if order.side == Side.BUY.value else self.bids
        
        while opposite_book and order.remaining_quantity > 0:
            best_order = opposite_book[0]
            actual_price = -best_order.price if order.side == Side.BUY.value else best_order.price
            
            # Check if orders cross
            if order.side == Side.BUY.value and order.price < actual_price:
                break
            elif order.side == Side.SELL.value and order.price > actual_price:
                break
            
            best_order = heapq.heappop(opposite_book)
            if order.side == Side.BUY.value:
                best_order.price = -best_order.price
            
            trade = self._execute_trade(order, best_order)
            trades.append(trade)
            
            if best_order.remaining_quantity > 0:
                heapq.heappush(opposite_book, best_order)
        
        return trades
    
    def _execute_trade(self, aggressive_order: Order, passive_order: Order) -> Trade:
        """Execute trade between two orders"""
        trade_qty = min(aggressive_order.remaining_quantity, passive_order.remaining_quantity)
        trade_price = passive_order.price
        
        aggressive_order.remaining_quantity -= trade_qty
        passive_order.remaining_quantity -= trade_qty
        
        self.trade_count += 1
        self.last_trade_price = trade_price
        
        return Trade(
            trade_id=f"T{self.trade_count:08d}",
            timestamp=time.time(),
            symbol=self.symbol,
            price=trade_price,
            quantity=trade_qty,
            buyer_order_id=aggressive_order.order_id if aggressive_order.side == Side.BUY.value else passive_order.order_id,
            seller_order_id=aggressive_order.order_id if aggressive_order.side == Side.SELL.value else passive_order.order_id
        )
    
    def _add_to_book(self, order: Order):
        """Add order to appropriate side of book"""
        self.orders[order.order_id] = order
        
        if order.side == Side.BUY.value:
            order.price = -order.price  # Negate for max heap
            heapq.heappush(self.bids, order)
        else:
            heapq.heappush(self.asks, order)
    
    def get_best_bid_ask(self) -> Tuple[Optional[float], Optional[float]]:
        """Get current best bid and ask prices"""
        best_bid = -self.bids[0].price if self.bids else None
        best_ask = self.asks[0].price if self.asks else None
        return best_bid, best_ask
    
    def get_spread(self) -> Optional[float]:
        """Calculate current bid-ask spread"""
        bid, ask = self.get_best_bid_ask()
        return ask - bid if bid and ask else None

# ============================================================================
# Market Data Generator
# ============================================================================

class MarketDataGenerator:
    """
    Simulates realistic market data feed for testing
    """
    def __init__(self, symbols: List[str], base_rate: int = 1000):
        self.symbols = symbols
        self.base_rate = base_rate  # messages per second
        self.sequence_nums = defaultdict(int)
        self.prices = {symbol: 100.0 + random.random() * 100 for symbol in symbols}
        self.volumes = {symbol: 0 for symbol in symbols}
        
    async def generate_stream(self) -> MarketData:
        """Generate continuous stream of market data"""
        while True:
            symbol = random.choice(self.symbols)
            self.sequence_nums[symbol] += 1
            
            # Random walk for price
            price_change = random.gauss(0, 0.01)
            self.prices[symbol] *= (1 + price_change)
            
            # Generate bid/ask around mid price
            spread = random.uniform(0.01, 0.05)
            bid_price = self.prices[symbol] - spread/2
            ask_price = self.prices[symbol] + spread/2
            
            # Random volumes
            bid_size = random.randint(100, 10000)
            ask_size = random.randint(100, 10000)
            trade_volume = random.randint(0, 1000)
            self.volumes[symbol] += trade_volume
            
            # Determine asset class
            if symbol.endswith('_OPT'):
                asset_class = AssetClass.OPTION.value
            elif symbol.endswith('_FUT'):
                asset_class = AssetClass.FUTURE.value
            else:
                asset_class = AssetClass.EQUITY.value
            
            data = MarketData(
                timestamp=time.time(),
                symbol=symbol,
                asset_class=asset_class,
                bid_price=round(bid_price, 4),
                ask_price=round(ask_price, 4),
                bid_size=bid_size,
                ask_size=ask_size,
                last_price=round(self.prices[symbol], 4),
                volume=self.volumes[symbol],
                sequence_num=self.sequence_nums[symbol]
            )
            
            # Control rate
            await asyncio.sleep(1.0 / self.base_rate)
            yield data

# ============================================================================
# Data Processing Pipeline
# ============================================================================

class DataProcessor:
    """
    Process market data with various analytics
    """
    def __init__(self):
        self.processed_count = 0
        self.error_count = 0
        self.latencies = deque(maxlen=10000)
        
    def process(self, data: MarketData) -> Dict:
        """Process individual market data point"""
        start_time = time.time()
        
        try:
            # Simulate processing
            result = {
                'symbol': data.symbol,
                'mid_price': (data.bid_price + data.ask_price) / 2,
                'spread': data.ask_price - data.bid_price,
                'spread_bps': ((data.ask_price - data.bid_price) / data.bid_price) * 10000,
                'bid_ask_ratio': data.bid_size / data.ask_size if data.ask_size > 0 else 0,
                'timestamp': data.timestamp,
                'processed_at': time.time()
            }
            
            self.processed_count += 1
            
        except Exception as e:
            logger.error(f"Processing error: {e}")
            self.error_count += 1
            result = None
        
        # Track latency
        latency = (time.time() - start_time) * 1000  # Convert to milliseconds
        self.latencies.append(latency)
        
        return result
    
    def get_stats(self) -> Dict:
        """Get processing statistics"""
        if self.latencies:
            latencies_array = np.array(self.latencies)
            return {
                'processed_count': self.processed_count,
                'error_count': self.error_count,
                'error_rate': self.error_count / max(self.processed_count, 1),
                'avg_latency_ms': np.mean(latencies_array),
                'p50_latency_ms': np.percentile(latencies_array, 50),
                'p95_latency_ms': np.percentile(latencies_array, 95),
                'p99_latency_ms': np.percentile(latencies_array, 99),
                'max_latency_ms': np.max(latencies_array)
            }
        return {}

# ============================================================================
# Async Pipeline Manager
# ============================================================================

class AsyncPipeline:
    """
    Manages async data pipeline with multiple consumers
    """
    def __init__(self, num_workers: int = 4):
        self.num_workers = num_workers
        self.queue = asyncio.Queue(maxsize=10000)
        self.processors = [DataProcessor() for _ in range(num_workers)]
        self.running = True
        self.start_time = None
        self.message_count = 0
        
    async def producer(self, generator: MarketDataGenerator):
        """Producer coroutine - generates market data"""
        async for data in generator.generate_stream():
            if not self.running:
                break
            
            try:
                await asyncio.wait_for(self.queue.put(data), timeout=0.001)
                self.message_count += 1
            except asyncio.TimeoutError:
                logger.warning("Queue full, dropping message")
    
    async def consumer(self, processor_id: int):
        """Consumer coroutine - processes market data"""
        processor = self.processors[processor_id]
        
        while self.running:
            try:
                data = await asyncio.wait_for(self.queue.get(), timeout=1.0)
                result = processor.process(data)
                
                # In production, would send to downstream systems
                if result and self.message_count % 1000 == 0:
                    logger.debug(f"Processed: {result['symbol']} @ {result['mid_price']:.4f}")
                    
            except asyncio.TimeoutError:
                continue
            except Exception as e:
                logger.error(f"Consumer {processor_id} error: {e}")
    
    async def monitor(self):
        """Monitor coroutine - reports statistics"""
        while self.running:
            await asyncio.sleep(5)
            
            if self.start_time:
                elapsed = time.time() - self.start_time
                rate = self.message_count / elapsed if elapsed > 0 else 0
                
                # Aggregate stats from all processors
                total_processed = sum(p.processed_count for p in self.processors)
                total_errors = sum(p.error_count for p in self.processors)
                all_latencies = []
                for p in self.processors:
                    all_latencies.extend(list(p.latencies))
                
                if all_latencies:
                    latencies_array = np.array(all_latencies)
                    
                    logger.info(f"""
=== Pipeline Statistics ===
Messages Generated: {self.message_count:,}
Messages Processed: {total_processed:,}
Messages/sec: {rate:.1f}
Queue Size: {self.queue.qsize()}
Error Rate: {total_errors/max(total_processed, 1):.4%}
Latency P50: {np.percentile(latencies_array, 50):.3f}ms
Latency P99: {np.percentile(latencies_array, 99):.3f}ms
Uptime: {elapsed:.1f}s
===========================
                    """.strip())
    
    async def run(self, symbols: List[str], duration: int = 60):
        """Run the pipeline for specified duration"""
        self.start_time = time.time()
        generator = MarketDataGenerator(symbols, base_rate=200)  # 200 msgs/sec per symbol
        
        # Create tasks
        tasks = []
        tasks.append(asyncio.create_task(self.producer(generator)))
        
        for i in range(self.num_workers):
            tasks.append(asyncio.create_task(self.consumer(i)))
        
        tasks.append(asyncio.create_task(self.monitor()))
        
        # Run for specified duration
        await asyncio.sleep(duration)
        
        # Shutdown
        self.running = False
        await asyncio.gather(*tasks, return_exceptions=True)
        
        # Final stats
        logger.info(f"Pipeline completed. Total messages: {self.message_count:,}")

# ============================================================================
# Multiprocessing Pipeline
# ============================================================================

class MultiprocessPipeline:
    """
    Distributed processing using multiprocessing for CPU-intensive tasks
    """
    def __init__(self, num_processes: int = mp.cpu_count()):
        self.num_processes = num_processes
        self.manager = Manager()
        self.input_queue = self.manager.Queue(maxsize=10000)
        self.output_queue = self.manager.Queue(maxsize=10000)
        self.stats = self.manager.dict()
        self.running = self.manager.Value('b', True)
        
    @staticmethod
    def worker_process(worker_id: int, input_queue: Queue, output_queue: Queue, 
                       stats: dict, running):
        """Worker process for CPU-intensive processing"""
        processor = DataProcessor()
        processed = 0
        
        while running.value:
            try:
                data = input_queue.get(timeout=1)
                result = processor.process(data)
                
                if result:
                    output_queue.put(result)
                    processed += 1
                    
                    # Update shared stats periodically
                    if processed % 100 == 0:
                        stats[f'worker_{worker_id}_processed'] = processed
                        
            except:
                continue
        
        stats[f'worker_{worker_id}_final'] = processed
    
    def run(self, symbols: List[str], duration: int = 60):
        """Run multiprocess pipeline"""
        processes = []
        
        # Start worker processes
        for i in range(self.num_processes):
            p = mp.Process(
                target=self.worker_process,
                args=(i, self.input_queue, self.output_queue, self.stats, self.running)
            )
            p.start()
            processes.append(p)
            logger.info(f"Started worker process {i}")
        
        # Generate test data
        start_time = time.time()
        message_count = 0
        
        while time.time() - start_time < duration:
            for symbol in symbols:
                data = MarketData(
                    timestamp=time.time(),
                    symbol=symbol,
                    asset_class=AssetClass.EQUITY.value,
                    bid_price=100.0 + random.random(),
                    ask_price=100.1 + random.random(),
                    bid_size=1000,
                    ask_size=1000,
                    last_price=100.05,
                    volume=10000,
                    sequence_num=message_count
                )
                
                try:
                    self.input_queue.put(data, timeout=0.001)
                    message_count += 1
                except:
                    pass
                
                time.sleep(0.001)  # Control rate
        
        # Shutdown
        self.running.value = False
        for p in processes:
            p.join(timeout=5)
        
        logger.info(f"Multiprocess pipeline completed. Messages: {message_count:,}")
        logger.info(f"Worker stats: {dict(self.stats)}")

# ============================================================================
# Main Pipeline Controller
# ============================================================================

class HFTPipeline:
    """
    Main controller for the HFT data pipeline
    """
    def __init__(self):
        self.symbols = [
            'AAPL', 'MSFT', 'GOOGL', 'AMZN', 'TSLA',  # Equities
            'SPY_OPT', 'QQQ_OPT', 'IWM_OPT',  # Options
            'ES_FUT', 'NQ_FUT', 'CL_FUT'  # Futures
        ]
        self.order_books = {symbol: OrderBook(symbol) for symbol in self.symbols}
        self.shutdown_event = threading.Event()
        
    def signal_handler(self, signum, frame):
        """Handle shutdown signals gracefully"""
        logger.info("Shutdown signal received")
        self.shutdown_event.set()
        sys.exit(0)
    
    async def run_async_demo(self):
        """Run async pipeline demonstration"""
        logger.info("Starting Async Pipeline Demo")
        logger.info(f"Processing {len(self.symbols)} symbols")
        logger.info("Target: 10,000+ messages/minute")
        
        pipeline = AsyncPipeline(num_workers=4)
        await pipeline.run(self.symbols, duration=30)
    
    def run_multiprocess_demo(self):
        """Run multiprocessing pipeline demonstration"""
        logger.info("Starting Multiprocess Pipeline Demo")
        logger.info(f"Using {mp.cpu_count()} CPU cores")
        
        pipeline = MultiprocessPipeline()
        pipeline.run(self.symbols, duration=30)
    
    def run_order_book_demo(self):
        """Demonstrate order book processing"""
        logger.info("Starting Order Book Demo")
        
        book = self.order_books['AAPL']
        orders_processed = 0
        trades_executed = 0
        
        start_time = time.time()
        
        # Generate and process orders
        for _ in range(10000):
            order = Order(
                order_id=f"O{orders_processed:08d}",
                timestamp=time.time(),
                symbol='AAPL',
                side=random.choice([Side.BUY.value, Side.SELL.value]),
                order_type=random.choice([OrderType.LIMIT.value, OrderType.MARKET.value]),
                price=100.0 + random.gauss(0, 1),
                quantity=random.randint(100, 1000)
            )
            
            trades = book.add_order(order)
            orders_processed += 1
            trades_executed += len(trades)
        
        elapsed = time.time() - start_time
        
        logger.info(f"""
=== Order Book Performance ===
Orders Processed: {orders_processed:,}
Trades Executed: {trades_executed:,}
Orders/sec: {orders_processed/elapsed:.1f}
Avg Latency: {elapsed/orders_processed*1000:.3f}ms
Best Bid/Ask: {book.get_best_bid_ask()}
Spread: {book.get_spread():.4f}
==============================
        """.strip())
    
    def run(self, mode: str = 'all'):
        """
        Run the HFT pipeline in specified mode
        
        Args:
            mode: 'async', 'multiprocess', 'orderbook', or 'all'
        """
        # Setup signal handlers
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        
        logger.info("""
╔════════════════════════════════════════════════════════════╗
║     High-Frequency Trading Data Pipeline                      ║
║     ============================================             ║
║     Performance Targets:                                      ║
║     - 10,000+ messages/minute                                 ║
║     - Sub-millisecond latency                                 ║
║     - 99.9% uptime                                           ║
╚════════════════════════════════════════════════════════════╝
        """.strip())
        
        if mode in ['async', 'all']:
            asyncio.run(self.run_async_demo())
            
        if mode in ['multiprocess', 'all']:
            self.run_multiprocess_demo()
            
        if mode in ['orderbook', 'all']:
            self.run_order_book_demo()
        
        logger.info("Pipeline demonstration completed successfully!")

# ============================================================================
# Entry Point
# ============================================================================

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='High-Frequency Trading Data Pipeline')
    parser.add_argument('--mode', choices=['async', 'multiprocess', 'orderbook', 'all'],
                       default='all', help='Pipeline mode to run')
    parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
    
    args = parser.parse_args()
    
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    pipeline = HFTPipeline()
    pipeline.run(mode=args.mode)