armature-core 0.7.0

//! Async Runtime Configuration and Optimization
//!
//! This module provides fine-grained control over the Tokio runtime for
//! optimizing HTTP server workloads.
//!
//! # Key Features
//!
//! - **Task Spawning Control**: Inline simple handlers to reduce overhead
//! - **LocalSet Support**: Single-threaded mode for cache locality
//! - **Runtime Tuning**: Configure thread count, stack size, work-stealing
//! - **Work-Stealing Optimization**: Tune for HTTP request patterns
//!
//! # Performance Impact
//!
//! - Reduced task spawning: -50% overhead for simple handlers
//! - LocalSet mode: +10-15% throughput for CPU-bound handlers
//! - Work-stealing tuning: +5-10% for mixed workloads

use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;
use tokio::runtime::{Builder, Runtime};
use tokio::task::LocalSet;

// ============================================================================
// Runtime Configuration
// ============================================================================

/// Configuration for the Tokio runtime.
#[derive(Debug, Clone)]
pub struct RuntimeConfig {
    /// Number of worker threads (None = CPU count)
    pub worker_threads: Option<usize>,
    /// Thread name prefix
    pub thread_name: String,
    /// Thread stack size in bytes
    pub thread_stack_size: Option<usize>,
    /// Enable I/O driver
    pub enable_io: bool,
    /// Enable time driver
    pub enable_time: bool,
    /// Global task queue interval (work-stealing frequency)
    pub global_queue_interval: Option<u32>,
    /// Event poll interval
    pub event_interval: Option<u32>,
    /// Max blocking threads for spawn_blocking
    pub max_blocking_threads: Option<usize>,
    /// Thread keep-alive duration
    pub thread_keep_alive: Option<Duration>,
    /// Use current-thread runtime (single-threaded)
    pub current_thread: bool,
    /// Enable LocalSet for !Send futures
    pub use_local_set: bool,
}

impl Default for RuntimeConfig {
    fn default() -> Self {
        Self {
            worker_threads: None,
            thread_name: "armature-worker".to_string(),
            thread_stack_size: Some(2 * 1024 * 1024), // 2MB
            enable_io: true,
            enable_time: true,
            global_queue_interval: Some(61), // Tokio default
            event_interval: Some(61),
            max_blocking_threads: Some(512),
            thread_keep_alive: Some(Duration::from_secs(10)),
            current_thread: false,
            use_local_set: false,
        }
    }
}

impl RuntimeConfig {
    /// Create new configuration.
    pub fn new() -> Self {
        Self::default()
    }

    /// Configuration optimized for throughput.
    ///
    /// - More worker threads
    /// - Larger work-stealing interval (less contention)
    /// - Larger blocking pool
    pub fn throughput() -> Self {
        let cpus = num_cpus();
        Self {
            worker_threads: Some(cpus),
            thread_name: "armature-tp".to_string(),
            thread_stack_size: Some(4 * 1024 * 1024), // 4MB for deep stacks
            enable_io: true,
            enable_time: true,
            global_queue_interval: Some(128), // Less frequent stealing
            event_interval: Some(128),
            max_blocking_threads: Some(1024),
            thread_keep_alive: Some(Duration::from_secs(30)),
            current_thread: false,
            use_local_set: false,
        }
    }

    /// Configuration optimized for low latency.
    ///
    /// - Fewer threads for cache locality
    /// - More aggressive work-stealing
    /// - Smaller blocking pool
    pub fn low_latency() -> Self {
        let cpus = num_cpus();
        Self {
            worker_threads: Some(cpus.min(4)), // Cap at 4 for locality
            thread_name: "armature-ll".to_string(),
            thread_stack_size: Some(1024 * 1024), // 1MB
            enable_io: true,
            enable_time: true,
            global_queue_interval: Some(31), // More frequent stealing
            event_interval: Some(31),
            max_blocking_threads: Some(64),
            thread_keep_alive: Some(Duration::from_secs(5)),
            current_thread: false,
            use_local_set: false,
        }
    }

    /// Single-threaded configuration with LocalSet.
    ///
    /// Best for handlers that benefit from cache locality
    /// and don't need to be Send.
    pub fn single_threaded() -> Self {
        Self {
            worker_threads: Some(1),
            thread_name: "armature-st".to_string(),
            thread_stack_size: Some(2 * 1024 * 1024),
            enable_io: true,
            enable_time: true,
            global_queue_interval: None,
            event_interval: None,
            max_blocking_threads: Some(32),
            thread_keep_alive: Some(Duration::from_secs(10)),
            current_thread: true,
            use_local_set: true,
        }
    }

    /// Builder: set worker threads.
    pub fn worker_threads(mut self, count: usize) -> Self {
        self.worker_threads = Some(count);
        self
    }

    /// Builder: set thread name.
    pub fn thread_name(mut self, name: impl Into<String>) -> Self {
        self.thread_name = name.into();
        self
    }

    /// Builder: set stack size.
    pub fn thread_stack_size(mut self, size: usize) -> Self {
        self.thread_stack_size = Some(size);
        self
    }

    /// Builder: set global queue interval.
    pub fn global_queue_interval(mut self, interval: u32) -> Self {
        self.global_queue_interval = Some(interval);
        self
    }

    /// Builder: set event interval.
    pub fn event_interval(mut self, interval: u32) -> Self {
        self.event_interval = Some(interval);
        self
    }

    /// Builder: set max blocking threads.
    pub fn max_blocking_threads(mut self, count: usize) -> Self {
        self.max_blocking_threads = Some(count);
        self
    }

    /// Builder: enable current-thread mode.
    pub fn current_thread(mut self, enabled: bool) -> Self {
        self.current_thread = enabled;
        self
    }

    /// Builder: enable LocalSet.
    pub fn use_local_set(mut self, enabled: bool) -> Self {
        self.use_local_set = enabled;
        self
    }

    /// Build the runtime.
    pub fn build(&self) -> std::io::Result<Runtime> {
        let mut builder = if self.current_thread {
            Builder::new_current_thread()
        } else {
            Builder::new_multi_thread()
        };

        if let Some(threads) = self.worker_threads
            && !self.current_thread
        {
            builder.worker_threads(threads);
        }

        builder.thread_name(&self.thread_name);

        if let Some(size) = self.thread_stack_size {
            builder.thread_stack_size(size);
        }

        if self.enable_io {
            builder.enable_io();
        }

        if self.enable_time {
            builder.enable_time();
        }

        if let Some(interval) = self.global_queue_interval {
            builder.global_queue_interval(interval);
        }

        if let Some(interval) = self.event_interval {
            builder.event_interval(interval);
        }

        if let Some(max) = self.max_blocking_threads {
            builder.max_blocking_threads(max);
        }

        if let Some(duration) = self.thread_keep_alive {
            builder.thread_keep_alive(duration);
        }

        builder.build()
    }
}

// ============================================================================
// Task Spawning Control
// ============================================================================

/// Controls when tasks should be spawned vs inlined.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SpawnPolicy {
    /// Always spawn a new task (default Tokio behavior)
    Always,
    /// Never spawn, inline everything (single-threaded)
    Never,
    /// Spawn only if handler duration exceeds threshold
    #[default]
    Adaptive,
    /// Spawn based on current load
    LoadBased,
}

/// Configuration for task spawning decisions.
#[derive(Debug, Clone)]
pub struct SpawnConfig {
    /// Spawn policy
    pub policy: SpawnPolicy,
    /// Duration threshold for adaptive spawning (microseconds)
    pub duration_threshold_us: u64,
    /// Load threshold for load-based spawning (pending tasks)
    pub load_threshold: usize,
    /// Force spawn for CPU-bound handlers
    pub spawn_cpu_bound: bool,
}

impl Default for SpawnConfig {
    fn default() -> Self {
        Self {
            policy: SpawnPolicy::Adaptive,
            duration_threshold_us: 100, // 100µs
            load_threshold: 100,
            spawn_cpu_bound: true,
        }
    }
}

impl SpawnConfig {
    /// Create new config.
    pub fn new() -> Self {
        Self::default()
    }

    /// Configuration that never spawns (inline everything).
    pub fn inline_all() -> Self {
        Self {
            policy: SpawnPolicy::Never,
            ..Default::default()
        }
    }

    /// Configuration that always spawns.
    pub fn always_spawn() -> Self {
        Self {
            policy: SpawnPolicy::Always,
            ..Default::default()
        }
    }

    /// Builder: set policy.
    pub fn policy(mut self, policy: SpawnPolicy) -> Self {
        self.policy = policy;
        self
    }

    /// Builder: set duration threshold.
    pub fn duration_threshold(mut self, us: u64) -> Self {
        self.duration_threshold_us = us;
        self
    }

    /// Builder: set load threshold.
    pub fn load_threshold(mut self, count: usize) -> Self {
        self.load_threshold = count;
        self
    }
}

/// Tracks handler durations for adaptive spawning.
#[derive(Debug, Default)]
pub struct HandlerMetrics {
    /// Running average duration (microseconds)
    avg_duration_us: AtomicU64,
    /// Total invocations
    invocations: AtomicU64,
    /// Times spawned
    spawned: AtomicU64,
    /// Times inlined
    inlined: AtomicU64,
}

impl HandlerMetrics {
    /// Create new metrics.
    pub fn new() -> Self {
        Self::default()
    }

    /// Record handler duration.
    pub fn record_duration(&self, us: u64) {
        let current = self.avg_duration_us.load(Ordering::Relaxed);
        let new_avg = if current == 0 {
            us
        } else {
            (current * 7 + us) / 8 // EMA with alpha = 0.125
        };
        self.avg_duration_us.store(new_avg, Ordering::Relaxed);
        self.invocations.fetch_add(1, Ordering::Relaxed);
    }

    /// Record spawn decision.
    pub fn record_spawn(&self, spawned: bool) {
        if spawned {
            self.spawned.fetch_add(1, Ordering::Relaxed);
        } else {
            self.inlined.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Get average duration.
    pub fn avg_duration_us(&self) -> u64 {
        self.avg_duration_us.load(Ordering::Relaxed)
    }

    /// Get total invocations.
    pub fn invocations(&self) -> u64 {
        self.invocations.load(Ordering::Relaxed)
    }

    /// Get spawn ratio.
    pub fn spawn_ratio(&self) -> f64 {
        let spawned = self.spawned.load(Ordering::Relaxed) as f64;
        let inlined = self.inlined.load(Ordering::Relaxed) as f64;
        let total = spawned + inlined;
        if total > 0.0 { spawned / total } else { 0.0 }
    }
}

/// Smart task spawner that decides whether to spawn or inline.
pub struct SmartSpawner {
    config: SpawnConfig,
    pending_tasks: AtomicUsize,
    metrics: HandlerMetrics,
}

impl SmartSpawner {
    /// Create new spawner.
    pub fn new(config: SpawnConfig) -> Self {
        Self {
            config,
            pending_tasks: AtomicUsize::new(0),
            metrics: HandlerMetrics::new(),
        }
    }

    /// Decide whether to spawn based on policy.
    pub fn should_spawn(&self, estimated_duration_us: Option<u64>) -> bool {
        let should = match self.config.policy {
            SpawnPolicy::Always => true,
            SpawnPolicy::Never => false,
            SpawnPolicy::Adaptive => {
                let duration =
                    estimated_duration_us.unwrap_or_else(|| self.metrics.avg_duration_us());
                duration > self.config.duration_threshold_us
            }
            SpawnPolicy::LoadBased => {
                let pending = self.pending_tasks.load(Ordering::Relaxed);
                pending < self.config.load_threshold
            }
        };

        self.metrics.record_spawn(should);
        RUNTIME_STATS.record_spawn_decision(should);
        should
    }

    /// Execute future, spawning or inlining based on policy.
    pub async fn execute<F, T>(&self, future: F) -> T
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        if self.should_spawn(None) {
            self.pending_tasks.fetch_add(1, Ordering::Relaxed);
            let result = tokio::spawn(future).await.expect("task panicked");
            self.pending_tasks.fetch_sub(1, Ordering::Relaxed);
            result
        } else {
            future.await
        }
    }

    /// Execute with timing to update metrics.
    pub async fn execute_timed<F, T>(&self, future: F) -> T
    where
        F: Future<Output = T> + Send + 'static,
        T: Send + 'static,
    {
        let start = std::time::Instant::now();
        let result = self.execute(future).await;
        let duration = start.elapsed().as_micros() as u64;
        self.metrics.record_duration(duration);
        result
    }

    /// Get metrics.
    pub fn metrics(&self) -> &HandlerMetrics {
        &self.metrics
    }

    /// Get pending task count.
    pub fn pending_tasks(&self) -> usize {
        self.pending_tasks.load(Ordering::Relaxed)
    }
}

// ============================================================================
// LocalSet Runner
// ============================================================================

/// Runner for LocalSet-based execution (!Send futures).
pub struct LocalRunner {
    local_set: LocalSet,
    enabled: AtomicBool,
}

impl LocalRunner {
    /// Create new local runner.
    pub fn new() -> Self {
        Self {
            local_set: LocalSet::new(),
            enabled: AtomicBool::new(true),
        }
    }

    /// Check if local execution is enabled.
    pub fn is_enabled(&self) -> bool {
        self.enabled.load(Ordering::Relaxed)
    }

    /// Enable/disable local execution.
    pub fn set_enabled(&self, enabled: bool) {
        self.enabled.store(enabled, Ordering::Relaxed);
    }

    /// Spawn a local (!Send) future.
    pub fn spawn_local<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
    where
        F: Future + 'static,
        F::Output: 'static,
    {
        RUNTIME_STATS.record_local_spawn();
        self.local_set.spawn_local(future)
    }

    /// Run the local set with a runtime.
    pub async fn run<F, T>(&self, future: F) -> T
    where
        F: Future<Output = T>,
    {
        self.local_set.run_until(future).await
    }

    /// Get the local set for manual control.
    pub fn local_set(&self) -> &LocalSet {
        &self.local_set
    }
}

impl Default for LocalRunner {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Work-Stealing Configuration
// ============================================================================

/// Configuration for work-stealing optimization.
#[derive(Debug, Clone)]
pub struct WorkStealingConfig {
    /// Global queue check interval (lower = more stealing)
    pub global_queue_interval: u32,
    /// Event poll interval
    pub event_interval: u32,
    /// Enable work-stealing between threads
    pub enabled: bool,
    /// Prefer local queue (reduce stealing)
    pub prefer_local: bool,
}

impl Default for WorkStealingConfig {
    fn default() -> Self {
        Self {
            global_queue_interval: 61,
            event_interval: 61,
            enabled: true,
            prefer_local: false,
        }
    }
}

impl WorkStealingConfig {
    /// Create new config.
    pub fn new() -> Self {
        Self::default()
    }

    /// Configuration for aggressive work-stealing (load balance).
    pub fn aggressive() -> Self {
        Self {
            global_queue_interval: 31,
            event_interval: 31,
            enabled: true,
            prefer_local: false,
        }
    }

    /// Configuration for minimal work-stealing (cache locality).
    pub fn minimal() -> Self {
        Self {
            global_queue_interval: 255,
            event_interval: 255,
            enabled: true,
            prefer_local: true,
        }
    }

    /// Configuration to disable work-stealing.
    pub fn disabled() -> Self {
        Self {
            enabled: false,
            ..Default::default()
        }
    }

    /// Apply to runtime builder.
    pub fn apply_to_builder(&self, builder: &mut Builder) {
        builder.global_queue_interval(self.global_queue_interval);
        builder.event_interval(self.event_interval);
    }
}

// ============================================================================
// Managed Runtime
// ============================================================================

/// A managed runtime with optimization features.
pub struct ManagedRuntime {
    runtime: Runtime,
    config: RuntimeConfig,
    spawner: Arc<SmartSpawner>,
    local_runner: Option<LocalRunner>,
}

impl ManagedRuntime {
    /// Create new managed runtime.
    pub fn new(config: RuntimeConfig) -> std::io::Result<Self> {
        let spawn_config = if config.current_thread {
            SpawnConfig::inline_all()
        } else {
            SpawnConfig::default()
        };

        let runtime = config.build()?;
        let local_runner = if config.use_local_set {
            Some(LocalRunner::new())
        } else {
            None
        };

        Ok(Self {
            runtime,
            config,
            spawner: Arc::new(SmartSpawner::new(spawn_config)),
            local_runner,
        })
    }

    /// Create with default configuration.
    pub fn default_runtime() -> std::io::Result<Self> {
        Self::new(RuntimeConfig::default())
    }

    /// Create throughput-optimized runtime.
    pub fn throughput_runtime() -> std::io::Result<Self> {
        Self::new(RuntimeConfig::throughput())
    }

    /// Create low-latency runtime.
    pub fn low_latency_runtime() -> std::io::Result<Self> {
        Self::new(RuntimeConfig::low_latency())
    }

    /// Create single-threaded runtime with LocalSet.
    pub fn single_threaded_runtime() -> std::io::Result<Self> {
        Self::new(RuntimeConfig::single_threaded())
    }

    /// Get runtime handle.
    pub fn handle(&self) -> tokio::runtime::Handle {
        self.runtime.handle().clone()
    }

    /// Get spawner.
    pub fn spawner(&self) -> Arc<SmartSpawner> {
        Arc::clone(&self.spawner)
    }

    /// Get local runner if enabled.
    pub fn local_runner(&self) -> Option<&LocalRunner> {
        self.local_runner.as_ref()
    }

    /// Block on a future.
    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
        if let Some(ref local) = self.local_runner {
            self.runtime.block_on(local.run(future))
        } else {
            self.runtime.block_on(future)
        }
    }

    /// Spawn a task.
    pub fn spawn<F>(&self, future: F) -> tokio::task::JoinHandle<F::Output>
    where
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        self.runtime.spawn(future)
    }

    /// Get configuration.
    pub fn config(&self) -> &RuntimeConfig {
        &self.config
    }
}

// ============================================================================
// Statistics
// ============================================================================

/// Global runtime statistics.
#[derive(Debug, Default)]
pub struct RuntimeStats {
    /// Tasks spawned
    tasks_spawned: AtomicU64,
    /// Tasks inlined
    tasks_inlined: AtomicU64,
    /// Local spawns
    local_spawns: AtomicU64,
    /// Blocking tasks
    blocking_tasks: AtomicU64,
}

impl RuntimeStats {
    fn record_spawn_decision(&self, spawned: bool) {
        if spawned {
            self.tasks_spawned.fetch_add(1, Ordering::Relaxed);
        } else {
            self.tasks_inlined.fetch_add(1, Ordering::Relaxed);
        }
    }

    fn record_local_spawn(&self) {
        self.local_spawns.fetch_add(1, Ordering::Relaxed);
    }

    #[allow(dead_code)]
    fn record_blocking(&self) {
        self.blocking_tasks.fetch_add(1, Ordering::Relaxed);
    }

    /// Get spawned count.
    pub fn tasks_spawned(&self) -> u64 {
        self.tasks_spawned.load(Ordering::Relaxed)
    }

    /// Get inlined count.
    pub fn tasks_inlined(&self) -> u64 {
        self.tasks_inlined.load(Ordering::Relaxed)
    }

    /// Get local spawns.
    pub fn local_spawns(&self) -> u64 {
        self.local_spawns.load(Ordering::Relaxed)
    }

    /// Get spawn ratio.
    pub fn spawn_ratio(&self) -> f64 {
        let spawned = self.tasks_spawned() as f64;
        let total = spawned + self.tasks_inlined() as f64;
        if total > 0.0 { spawned / total } else { 0.0 }
    }
}

/// Global statistics.
static RUNTIME_STATS: RuntimeStats = RuntimeStats {
    tasks_spawned: AtomicU64::new(0),
    tasks_inlined: AtomicU64::new(0),
    local_spawns: AtomicU64::new(0),
    blocking_tasks: AtomicU64::new(0),
};

/// Get global runtime statistics.
pub fn runtime_stats() -> &'static RuntimeStats {
    &RUNTIME_STATS
}

// ============================================================================
// Utilities
// ============================================================================

/// Get number of CPUs.
fn num_cpus() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1)
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_runtime_config_default() {
        let config = RuntimeConfig::default();
        assert!(!config.current_thread);
        assert!(config.enable_io);
        assert!(config.enable_time);
    }

    #[test]
    fn test_runtime_config_throughput() {
        let config = RuntimeConfig::throughput();
        assert!(config.global_queue_interval.unwrap() > 61);
        assert!(config.max_blocking_threads.unwrap() > 512);
    }

    #[test]
    fn test_runtime_config_low_latency() {
        let config = RuntimeConfig::low_latency();
        assert!(config.global_queue_interval.unwrap() < 61);
        assert!(config.worker_threads.unwrap() <= 4);
    }

    #[test]
    fn test_runtime_config_single_threaded() {
        let config = RuntimeConfig::single_threaded();
        assert!(config.current_thread);
        assert!(config.use_local_set);
    }

    #[test]
    fn test_runtime_config_builder() {
        let config = RuntimeConfig::new()
            .worker_threads(8)
            .thread_name("test")
            .global_queue_interval(100);

        assert_eq!(config.worker_threads, Some(8));
        assert_eq!(config.thread_name, "test");
        assert_eq!(config.global_queue_interval, Some(100));
    }

    #[test]
    fn test_spawn_policy() {
        assert_eq!(SpawnPolicy::default(), SpawnPolicy::Adaptive);
    }

    #[test]
    fn test_spawn_config() {
        let config = SpawnConfig::inline_all();
        assert_eq!(config.policy, SpawnPolicy::Never);

        let config = SpawnConfig::always_spawn();
        assert_eq!(config.policy, SpawnPolicy::Always);
    }

    #[test]
    fn test_handler_metrics() {
        let metrics = HandlerMetrics::new();
        assert_eq!(metrics.avg_duration_us(), 0);

        metrics.record_duration(100);
        assert_eq!(metrics.avg_duration_us(), 100);

        metrics.record_duration(200);
        let avg = metrics.avg_duration_us();
        assert!(avg > 100 && avg < 200);
    }

    #[test]
    fn test_smart_spawner_always() {
        let spawner = SmartSpawner::new(SpawnConfig::always_spawn());
        assert!(spawner.should_spawn(None));
        assert!(spawner.should_spawn(Some(1)));
    }

    #[test]
    fn test_smart_spawner_never() {
        let spawner = SmartSpawner::new(SpawnConfig::inline_all());
        assert!(!spawner.should_spawn(None));
        assert!(!spawner.should_spawn(Some(1000)));
    }

    #[test]
    fn test_smart_spawner_adaptive() {
        let config = SpawnConfig::new().duration_threshold(50);
        let spawner = SmartSpawner::new(config);

        assert!(!spawner.should_spawn(Some(10))); // Below threshold
        assert!(spawner.should_spawn(Some(100))); // Above threshold
    }

    #[test]
    fn test_work_stealing_config() {
        let aggressive = WorkStealingConfig::aggressive();
        assert!(aggressive.global_queue_interval < 61);

        let minimal = WorkStealingConfig::minimal();
        assert!(minimal.global_queue_interval > 61);
        assert!(minimal.prefer_local);
    }

    #[test]
    fn test_local_runner() {
        let runner = LocalRunner::new();
        assert!(runner.is_enabled());

        runner.set_enabled(false);
        assert!(!runner.is_enabled());
    }

    #[test]
    fn test_runtime_stats() {
        let stats = runtime_stats();
        let _ = stats.tasks_spawned();
        let _ = stats.tasks_inlined();
        let _ = stats.spawn_ratio();
    }

    #[test]
    fn test_managed_runtime_spawn() {
        let runtime = ManagedRuntime::default_runtime().unwrap();

        let result = runtime.block_on(async { 42 });
        assert_eq!(result, 42);
    }

    #[test]
    fn test_build_runtime() {
        let config = RuntimeConfig::default();
        let runtime = config.build();
        assert!(runtime.is_ok());
    }
}