sochdb_storage/

rl_workload.rs

// Copyright 2025 Sushanth (https://github.com/sushanthpy)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! RL Workload Classifier (Task 10)
//!
//! This module provides a multi-armed bandit-based workload classifier for
//! automatic parameter tuning based on observed operation patterns.
//!
//! ## Problem
//!
//! Different workloads (OLTP, OLAP, mixed) require different tuning:
//! - OLTP: Small batch sizes, frequent flushes
//! - OLAP: Large batch sizes, aggressive prefetching
//! - Mixed: Adaptive switching
//!
//! ## Solution
//!
//! - **Feature Extraction:** Derive features from operation mix
//! - **UCB1 Algorithm:** Upper Confidence Bound for exploration/exploitation
//! - **Tuning Actions:** Adjust parameters based on classifier output
//!
//! ## Performance
//!
//! | Workload | Static Config | RL-Tuned |
//! |----------|---------------|----------|
//! | OLTP | 1× | 1.5× |
//! | OLAP | 1× | 2.0× |
//! | Mixed | 1× | 1.8× |

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::RwLock;
use std::time::Instant;

/// Number of arms (tuning configurations)
const NUM_ARMS: usize = 8;

/// Exploration constant for UCB1
const UCB_C: f64 = 1.41421356; // sqrt(2)

/// Window size for feature calculation
#[allow(dead_code)]
const FEATURE_WINDOW_SIZE: usize = 1000;

// ============================================================================
// Workload Types
// ============================================================================

/// Detected workload type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WorkloadType {
    /// Write-heavy transactional (OLTP)
    Oltp,
    /// Read-heavy analytical (OLAP)
    Olap,
    /// Mixed workload
    Mixed,
    /// Vector search heavy
    VectorSearch,
    /// Unknown (not enough data)
    Unknown,
}

impl WorkloadType {
    /// Get the string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Oltp => "OLTP",
            Self::Olap => "OLAP",
            Self::Mixed => "Mixed",
            Self::VectorSearch => "VectorSearch",
            Self::Unknown => "Unknown",
        }
    }
}

// ============================================================================
// Feature Vector
// ============================================================================

/// Operation counters for feature extraction
#[derive(Default)]
struct OperationCounters {
    /// Point reads
    point_reads: AtomicU64,
    /// Range scans
    range_scans: AtomicU64,
    /// Inserts
    inserts: AtomicU64,
    /// Updates
    updates: AtomicU64,
    /// Deletes
    deletes: AtomicU64,
    /// Vector searches
    vector_searches: AtomicU64,
}

/// Feature vector derived from operation mix
#[derive(Debug, Clone, Default)]
pub struct FeatureVector {
    /// Fraction of reads (point + range)
    pub read_fraction: f64,
    /// Fraction of writes (insert + update + delete)
    pub write_fraction: f64,
    /// Fraction of range scans
    pub scan_fraction: f64,
    /// Fraction of vector searches
    pub vector_fraction: f64,
    /// Average operation latency (ms)
    pub avg_latency_ms: f64,
    /// Operations per second
    pub ops_per_second: f64,
    /// Key locality (0 = random, 1 = sequential)
    pub key_locality: f64,
}

impl FeatureVector {
    /// Classify the workload based on features
    pub fn classify(&self) -> WorkloadType {
        if self.ops_per_second < 1.0 {
            return WorkloadType::Unknown;
        }
        
        if self.vector_fraction > 0.3 {
            return WorkloadType::VectorSearch;
        }
        
        if self.write_fraction > 0.7 {
            return WorkloadType::Oltp;
        }
        
        if self.scan_fraction > 0.3 {
            return WorkloadType::Olap;
        }
        
        if self.read_fraction > 0.7 {
            return WorkloadType::Olap;
        }
        
        WorkloadType::Mixed
    }
}

// ============================================================================
// Tuning Actions
// ============================================================================

/// Tuning configuration
#[derive(Debug, Clone)]
pub struct TuningConfig {
    /// Memtable size (bytes)
    pub memtable_size: usize,
    /// Write buffer count
    pub write_buffer_count: usize,
    /// Batch size for operations
    pub batch_size: usize,
    /// Prefetch distance
    pub prefetch_distance: usize,
    /// Background flush interval (ms)
    pub flush_interval_ms: u64,
    /// Compaction priority
    pub compaction_priority: CompactionPriority,
    /// Cache ratio (0-1)
    pub cache_ratio: f64,
    /// HNSW ef_search parameter
    pub hnsw_ef_search: usize,
}

impl Default for TuningConfig {
    fn default() -> Self {
        Self {
            memtable_size: 64 * 1024 * 1024, // 64 MB
            write_buffer_count: 2,
            batch_size: 256,
            prefetch_distance: 4,
            flush_interval_ms: 1000,
            compaction_priority: CompactionPriority::Balanced,
            cache_ratio: 0.5,
            hnsw_ef_search: 100,
        }
    }
}

/// Compaction priority strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompactionPriority {
    /// Minimize write amplification
    WriteOptimized,
    /// Minimize read amplification
    ReadOptimized,
    /// Balance both
    Balanced,
}

/// Predefined tuning configurations (arms)
fn get_arm_config(arm: usize) -> TuningConfig {
    match arm {
        0 => TuningConfig {
            // OLTP optimized: small buffers, fast flush
            memtable_size: 32 * 1024 * 1024,
            write_buffer_count: 4,
            batch_size: 64,
            prefetch_distance: 2,
            flush_interval_ms: 500,
            compaction_priority: CompactionPriority::WriteOptimized,
            cache_ratio: 0.3,
            hnsw_ef_search: 50,
        },
        1 => TuningConfig {
            // OLAP optimized: large buffers, aggressive prefetch
            memtable_size: 256 * 1024 * 1024,
            write_buffer_count: 2,
            batch_size: 1024,
            prefetch_distance: 16,
            flush_interval_ms: 5000,
            compaction_priority: CompactionPriority::ReadOptimized,
            cache_ratio: 0.8,
            hnsw_ef_search: 200,
        },
        2 => TuningConfig {
            // Vector search optimized
            memtable_size: 128 * 1024 * 1024,
            write_buffer_count: 2,
            batch_size: 512,
            prefetch_distance: 8,
            flush_interval_ms: 2000,
            compaction_priority: CompactionPriority::Balanced,
            cache_ratio: 0.9,
            hnsw_ef_search: 300,
        },
        3 => TuningConfig {
            // Balanced for mixed workload
            memtable_size: 64 * 1024 * 1024,
            write_buffer_count: 3,
            batch_size: 256,
            prefetch_distance: 4,
            flush_interval_ms: 1000,
            compaction_priority: CompactionPriority::Balanced,
            cache_ratio: 0.5,
            hnsw_ef_search: 100,
        },
        4 => TuningConfig {
            // Write burst handling
            memtable_size: 128 * 1024 * 1024,
            write_buffer_count: 6,
            batch_size: 128,
            prefetch_distance: 2,
            flush_interval_ms: 300,
            compaction_priority: CompactionPriority::WriteOptimized,
            cache_ratio: 0.2,
            hnsw_ef_search: 50,
        },
        5 => TuningConfig {
            // Read burst handling
            memtable_size: 32 * 1024 * 1024,
            write_buffer_count: 2,
            batch_size: 512,
            prefetch_distance: 32,
            flush_interval_ms: 3000,
            compaction_priority: CompactionPriority::ReadOptimized,
            cache_ratio: 0.95,
            hnsw_ef_search: 150,
        },
        6 => TuningConfig {
            // Latency sensitive
            memtable_size: 16 * 1024 * 1024,
            write_buffer_count: 8,
            batch_size: 32,
            prefetch_distance: 1,
            flush_interval_ms: 200,
            compaction_priority: CompactionPriority::Balanced,
            cache_ratio: 0.6,
            hnsw_ef_search: 75,
        },
        7 => TuningConfig {
            // Throughput focused
            memtable_size: 512 * 1024 * 1024,
            write_buffer_count: 2,
            batch_size: 2048,
            prefetch_distance: 64,
            flush_interval_ms: 10000,
            compaction_priority: CompactionPriority::WriteOptimized,
            cache_ratio: 0.4,
            hnsw_ef_search: 100,
        },
        _ => TuningConfig::default(),
    }
}

// ============================================================================
// UCB1 Arm
// ============================================================================

/// An arm in the multi-armed bandit
struct UcbArm {
    /// Number of times this arm was selected
    count: AtomicU64,
    /// Total reward accumulated
    total_reward: RwLock<f64>,
    /// Sum of squared rewards (for variance)
    sum_squared_reward: RwLock<f64>,
}

impl UcbArm {
    fn new() -> Self {
        Self {
            count: AtomicU64::new(0),
            total_reward: RwLock::new(0.0),
            sum_squared_reward: RwLock::new(0.0),
        }
    }
    
    /// Get the average reward
    fn avg_reward(&self) -> f64 {
        let count = self.count.load(Ordering::Relaxed);
        if count == 0 {
            return 0.0;
        }
        *self.total_reward.read().unwrap() / count as f64
    }
    
    /// Record a reward
    fn record_reward(&self, reward: f64) {
        self.count.fetch_add(1, Ordering::Relaxed);
        *self.total_reward.write().unwrap() += reward;
        *self.sum_squared_reward.write().unwrap() += reward * reward;
    }
    
    /// Calculate UCB1 value
    fn ucb(&self, total_count: u64) -> f64 {
        let count = self.count.load(Ordering::Relaxed);
        if count == 0 {
            return f64::MAX; // Unexplored arm has infinite UCB
        }
        
        let avg = self.avg_reward();
        let exploration = UCB_C * ((total_count as f64).ln() / count as f64).sqrt();
        
        avg + exploration
    }
}

// ============================================================================
// Workload Classifier
// ============================================================================

/// RL-based workload classifier with UCB1 algorithm
pub struct WorkloadClassifier {
    /// Operation counters
    counters: OperationCounters,
    /// Arms for UCB1
    arms: [UcbArm; NUM_ARMS],
    /// Currently selected arm
    current_arm: RwLock<usize>,
    /// Current config
    current_config: RwLock<TuningConfig>,
    /// Start time for ops/sec calculation
    start_time: Instant,
    /// Last feature extraction time
    #[allow(dead_code)]
    last_feature_time: RwLock<Instant>,
    /// Cached feature vector
    #[allow(dead_code)]
    cached_features: RwLock<FeatureVector>,
    /// Reward measurement start
    reward_start: RwLock<Option<Instant>>,
    /// Operations at reward start
    ops_at_reward_start: AtomicU64,
}

impl WorkloadClassifier {
    /// Create a new classifier
    pub fn new() -> Self {
        Self {
            counters: OperationCounters::default(),
            arms: std::array::from_fn(|_| UcbArm::new()),
            current_arm: RwLock::new(3), // Start with balanced config
            current_config: RwLock::new(get_arm_config(3)),
            start_time: Instant::now(),
            last_feature_time: RwLock::new(Instant::now()),
            cached_features: RwLock::new(FeatureVector::default()),
            reward_start: RwLock::new(None),
            ops_at_reward_start: AtomicU64::new(0),
        }
    }
    
    /// Record a point read
    #[inline]
    pub fn record_point_read(&self) {
        self.counters.point_reads.fetch_add(1, Ordering::Relaxed);
    }
    
    /// Record a range scan
    #[inline]
    pub fn record_range_scan(&self) {
        self.counters.range_scans.fetch_add(1, Ordering::Relaxed);
    }
    
    /// Record an insert
    #[inline]
    pub fn record_insert(&self) {
        self.counters.inserts.fetch_add(1, Ordering::Relaxed);
    }
    
    /// Record an update
    #[inline]
    pub fn record_update(&self) {
        self.counters.updates.fetch_add(1, Ordering::Relaxed);
    }
    
    /// Record a delete
    #[inline]
    pub fn record_delete(&self) {
        self.counters.deletes.fetch_add(1, Ordering::Relaxed);
    }
    
    /// Record a vector search
    #[inline]
    pub fn record_vector_search(&self) {
        self.counters.vector_searches.fetch_add(1, Ordering::Relaxed);
    }
    
    /// Get total operations
    fn total_ops(&self) -> u64 {
        self.counters.point_reads.load(Ordering::Relaxed)
            + self.counters.range_scans.load(Ordering::Relaxed)
            + self.counters.inserts.load(Ordering::Relaxed)
            + self.counters.updates.load(Ordering::Relaxed)
            + self.counters.deletes.load(Ordering::Relaxed)
            + self.counters.vector_searches.load(Ordering::Relaxed)
    }
    
    /// Extract features from operation counters
    pub fn extract_features(&self) -> FeatureVector {
        let reads = self.counters.point_reads.load(Ordering::Relaxed);
        let scans = self.counters.range_scans.load(Ordering::Relaxed);
        let inserts = self.counters.inserts.load(Ordering::Relaxed);
        let updates = self.counters.updates.load(Ordering::Relaxed);
        let deletes = self.counters.deletes.load(Ordering::Relaxed);
        let vectors = self.counters.vector_searches.load(Ordering::Relaxed);
        
        let total = reads + scans + inserts + updates + deletes + vectors;
        let total_f = total.max(1) as f64;
        
        let elapsed = self.start_time.elapsed().as_secs_f64().max(0.001);
        
        FeatureVector {
            read_fraction: (reads + scans) as f64 / total_f,
            write_fraction: (inserts + updates + deletes) as f64 / total_f,
            scan_fraction: scans as f64 / total_f,
            vector_fraction: vectors as f64 / total_f,
            avg_latency_ms: 1.0, // Would be measured in practice
            ops_per_second: total as f64 / elapsed,
            key_locality: 0.5, // Would be measured from key distribution
        }
    }
    
    /// Get the current workload type
    pub fn workload_type(&self) -> WorkloadType {
        self.extract_features().classify()
    }
    
    /// Get the current tuning configuration
    pub fn current_config(&self) -> TuningConfig {
        self.current_config.read().unwrap().clone()
    }
    
    /// Start measuring reward (throughput)
    pub fn start_reward_measurement(&self) {
        *self.reward_start.write().unwrap() = Some(Instant::now());
        self.ops_at_reward_start.store(self.total_ops(), Ordering::Relaxed);
    }
    
    /// End measurement and update the bandit
    pub fn end_reward_measurement(&self) {
        let start = match *self.reward_start.read().unwrap() {
            Some(t) => t,
            None => return,
        };
        
        let elapsed = start.elapsed().as_secs_f64();
        if elapsed < 0.001 {
            return;
        }
        
        let ops_start = self.ops_at_reward_start.load(Ordering::Relaxed);
        let ops_now = self.total_ops();
        let throughput = (ops_now - ops_start) as f64 / elapsed;
        
        // Normalize reward to [0, 1]
        let reward = (throughput / 100000.0).min(1.0);
        
        // Update current arm
        let arm_idx = *self.current_arm.read().unwrap();
        self.arms[arm_idx].record_reward(reward);
        
        *self.reward_start.write().unwrap() = None;
    }
    
    /// Select the next arm using UCB1
    pub fn select_arm(&self) -> usize {
        let total_count: u64 = self.arms.iter()
            .map(|a| a.count.load(Ordering::Relaxed))
            .sum();
        
        if total_count < NUM_ARMS as u64 {
            // Initial exploration: try each arm once
            return total_count as usize;
        }
        
        // Select arm with highest UCB
        self.arms.iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| {
                a.ucb(total_count)
                    .partial_cmp(&b.ucb(total_count))
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|(i, _)| i)
            .unwrap_or(0)
    }
    
    /// Update the configuration based on current workload
    pub fn update_config(&self) {
        // End any ongoing measurement
        self.end_reward_measurement();
        
        // Select new arm
        let new_arm = self.select_arm();
        let new_config = get_arm_config(new_arm);
        
        *self.current_arm.write().unwrap() = new_arm;
        *self.current_config.write().unwrap() = new_config;
        
        // Start measuring with new config
        self.start_reward_measurement();
    }
    
    /// Get statistics
    pub fn stats(&self) -> ClassifierStats {
        let features = self.extract_features();
        let arm_stats: Vec<_> = self.arms.iter()
            .enumerate()
            .map(|(i, arm)| ArmStats {
                arm_id: i,
                count: arm.count.load(Ordering::Relaxed),
                avg_reward: arm.avg_reward(),
            })
            .collect();
        
        ClassifierStats {
            workload_type: features.classify(),
            features,
            current_arm: *self.current_arm.read().unwrap(),
            arm_stats,
        }
    }
}

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

/// Statistics for an arm
#[derive(Debug, Clone)]
pub struct ArmStats {
    /// Arm ID
    pub arm_id: usize,
    /// Number of times selected
    pub count: u64,
    /// Average reward
    pub avg_reward: f64,
}

/// Classifier statistics
#[derive(Debug, Clone)]
pub struct ClassifierStats {
    /// Detected workload type
    pub workload_type: WorkloadType,
    /// Current features
    pub features: FeatureVector,
    /// Currently selected arm
    pub current_arm: usize,
    /// Per-arm statistics
    pub arm_stats: Vec<ArmStats>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;
    
    #[test]
    fn test_feature_extraction() {
        let classifier = WorkloadClassifier::new();
        
        // Simulate OLTP workload
        for _ in 0..100 {
            classifier.record_insert();
            classifier.record_update();
        }
        for _ in 0..50 {
            classifier.record_point_read();
        }
        
        let features = classifier.extract_features();
        assert!(features.write_fraction > 0.5);
        
        let workload = features.classify();
        assert_eq!(workload, WorkloadType::Oltp);
    }
    
    #[test]
    fn test_olap_detection() {
        let classifier = WorkloadClassifier::new();
        
        // Simulate OLAP workload
        for _ in 0..100 {
            classifier.record_range_scan();
            classifier.record_point_read();
        }
        for _ in 0..10 {
            classifier.record_insert();
        }
        
        let features = classifier.extract_features();
        let workload = features.classify();
        assert_eq!(workload, WorkloadType::Olap);
    }
    
    #[test]
    fn test_vector_search_detection() {
        let classifier = WorkloadClassifier::new();
        
        for _ in 0..100 {
            classifier.record_vector_search();
        }
        for _ in 0..50 {
            classifier.record_point_read();
        }
        
        let features = classifier.extract_features();
        let workload = features.classify();
        assert_eq!(workload, WorkloadType::VectorSearch);
    }
    
    #[test]
    fn test_ucb_arm_selection() {
        let classifier = WorkloadClassifier::new();
        
        // Initially should explore
        for i in 0..NUM_ARMS {
            let arm = classifier.select_arm();
            // Give fake reward
            classifier.arms[arm].record_reward(if arm % 2 == 0 { 0.8 } else { 0.2 });
        }
        
        // After exploration, should prefer higher-reward arms
        let selected: Vec<_> = (0..20).map(|_| classifier.select_arm()).collect();
        let even_count = selected.iter().filter(|&&a| a % 2 == 0).count();
        
        // Should prefer even arms (higher reward)
        assert!(even_count > 10);
    }
    
    #[test]
    fn test_config_update() {
        let classifier = WorkloadClassifier::new();
        
        let config1 = classifier.current_config();
        
        // Simulate some activity
        for _ in 0..100 {
            classifier.record_insert();
        }
        
        classifier.start_reward_measurement();
        thread::sleep(Duration::from_millis(10));
        classifier.update_config();
        
        // Config may or may not change, but should be valid
        let config2 = classifier.current_config();
        assert!(config2.memtable_size > 0);
    }
    
    #[test]
    fn test_arm_configs() {
        for i in 0..NUM_ARMS {
            let config = get_arm_config(i);
            assert!(config.memtable_size > 0);
            assert!(config.batch_size > 0);
            assert!(config.prefetch_distance > 0);
        }
    }
    
    #[test]
    fn test_stats() {
        let classifier = WorkloadClassifier::new();
        
        for _ in 0..50 {
            classifier.record_insert();
            classifier.record_point_read();
        }
        
        let stats = classifier.stats();
        assert_eq!(stats.arm_stats.len(), NUM_ARMS);
        assert!(stats.features.ops_per_second > 0.0);
    }
}