performance.py

"""
Performance optimizations for Render deployment
Includes caching, connection pooling, and memory optimization
"""
import asyncio
import hashlib
import time
from typing import Any, Optional
from functools import lru_cache
import httpx

class PerformanceCache:
    """Simple in-memory cache with TTL"""
    
    def __init__(self, ttl: int = 300):
        self.cache = {}
        self.ttl = ttl
    
    def get(self, key: str) -> Optional[Any]:
        """Get value from cache if not expired"""
        if key in self.cache:
            value, timestamp = self.cache[key]
            if time.time() - timestamp < self.ttl:
                return value
            else:
                del self.cache[key]
        return None
    
    def set(self, key: str, value: Any):
        """Set value in cache with timestamp"""
        self.cache[key] = (value, time.time())
    
    def clear_expired(self):
        """Clear expired entries"""
        current_time = time.time()
        expired_keys = [
            key for key, (_, timestamp) in self.cache.items()
            if current_time - timestamp >= self.ttl
        ]
        for key in expired_keys:
            del self.cache[key]
    
    def get_stats(self):
        """Get cache statistics"""
        return {
            'total_entries': len(self.cache),
            'memory_mb': len(str(self.cache)) / 1024 / 1024
        }

# Global cache instance
_cache = PerformanceCache(ttl=300)

def cache_key(url: str, scanner_name: str) -> str:
    """Generate cache key for URL + scanner combination"""
    key_str = f"{scanner_name}:{url}"
    return hashlib.md5(key_str.encode()).hexdigest()

def get_cached_result(url: str, scanner_name: str) -> Optional[Any]:
    """Get cached scan result"""
    key = cache_key(url, scanner_name)
    return _cache.get(key)

def set_cached_result(url: str, scanner_name: str, result: Any):
    """Cache scan result"""
    key = cache_key(url, scanner_name)
    _cache.set(key, result)

def clear_old_cache():
    """Clear expired cache entries"""
    _cache.clear_expired()

class ConnectionPool:
    """Shared HTTP connection pool for all scanners"""
    
    def __init__(self, pool_size: int = 100, keepalive_expiry: float = 5.0):
        self.limits = httpx.Limits(
            max_connections=pool_size,
            max_keepalive_connections=pool_size // 2,
            keepalive_expiry=keepalive_expiry
        )
        self._client = None
    
    async def get_client(self) -> httpx.AsyncClient:
        """Get or create shared HTTP client"""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                limits=self.limits,
                timeout=httpx.Timeout(15.0, connect=5.0),
                http2=True,
                follow_redirects=False
            )
        return self._client
    
    async def close(self):
        """Close HTTP client"""
        if self._client and not self._client.is_closed:
            await self._client.aclose()

# Global connection pool
_pool = ConnectionPool()

async def get_http_client() -> httpx.AsyncClient:
    """Get shared HTTP client from pool"""
    return await _pool.get_client()

async def close_http_pool():
    """Close HTTP connection pool"""
    await _pool.close()

@lru_cache(maxsize=1000)
def cached_url_parse(url: str) -> dict:
    """Cache URL parsing results"""
    from urllib.parse import urlparse
    parsed = urlparse(url)
    return {
        'scheme': parsed.scheme,
        'netloc': parsed.netloc,
        'hostname': parsed.hostname,
        'port': parsed.port,
        'path': parsed.path
    }

class BatchProcessor:
    """Batch multiple requests to reduce overhead"""
    
    def __init__(self, batch_size: int = 5, timeout: float = 1.0):
        self.batch_size = batch_size
        self.timeout = timeout
        self.queue = []
        self.lock = asyncio.Lock()
    
    async def add(self, item: Any) -> list:
        """Add item to batch and process when full"""
        async with self.lock:
            self.queue.append(item)
            
            if len(self.queue) >= self.batch_size:
                batch = self.queue[:self.batch_size]
                self.queue = self.queue[self.batch_size:]
                return batch
        
        return []
    
    async def flush(self) -> list:
        """Process remaining items"""
        async with self.lock:
            batch = self.queue.copy()
            self.queue.clear()
            return batch

class MemoryOptimizer:
    """Optimize memory usage"""
    
    @staticmethod
    def truncate_html(html: str, max_size: int = 100000) -> str:
        """Truncate large HTML responses"""
        if len(html) > max_size:
            return html[:max_size]
        return html
    
    @staticmethod
    def compress_result(result: dict) -> dict:
        """Remove verbose fields from results"""
        if isinstance(result, dict):
            # Remove raw HTML/text from results
            result.pop('raw_html', None)
            result.pop('raw_response', None)
            result.pop('full_headers', None)
            
            # Truncate long strings
            for key, value in result.items():
                if isinstance(value, str) and len(value) > 1000:
                    result[key] = value[:1000] + '...'
        
        return result

async def warm_up():
    """Warm up services to prevent cold starts"""
    try:
        # Warm up HTTP client
        client = await get_http_client()
        
        # Warm up cache
        _cache.get('warmup')
        
        return True
    except:
        return False

def get_performance_stats() -> dict:
    """Get performance statistics"""
    return {
        'cache': _cache.get_stats(),
        'pool': {
            'active': _pool._client is not None and not _pool._client.is_closed
        }
    }