herodb/benches/common/backends.rs

// benches/common/backends.rs
use herodb::storage::Storage;
use herodb::storage_sled::SledStorage;
use herodb::storage_trait::StorageBackend;
use std::sync::Arc;
use tempfile::TempDir;

/// Backend type identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendType {
    Redb,
    Sled,
}

impl BackendType {
    pub fn name(&self) -> &'static str {
        match self {
            BackendType::Redb => "redb",
            BackendType::Sled => "sled",
        }
    }

    pub fn all() -> Vec<BackendType> {
        vec![BackendType::Redb, BackendType::Sled]
    }
}

/// Wrapper for benchmark backends with automatic cleanup
pub struct BenchmarkBackend {
    pub storage: Arc<dyn StorageBackend>,
    pub backend_type: BackendType,
    _temp_dir: TempDir, // Kept for automatic cleanup
}

impl BenchmarkBackend {
    /// Create a new redb backend for benchmarking
    pub fn new_redb() -> Result<Self, Box<dyn std::error::Error>> {
        let temp_dir = TempDir::new()?;
        let db_path = temp_dir.path().join("bench.db");
        let storage = Storage::new(db_path, false, None)?;
        
        Ok(Self {
            storage: Arc::new(storage),
            backend_type: BackendType::Redb,
            _temp_dir: temp_dir,
        })
    }

    /// Create a new sled backend for benchmarking
    pub fn new_sled() -> Result<Self, Box<dyn std::error::Error>> {
        let temp_dir = TempDir::new()?;
        let db_path = temp_dir.path().join("bench.sled");
        let storage = SledStorage::new(db_path, false, None)?;
        
        Ok(Self {
            storage: Arc::new(storage),
            backend_type: BackendType::Sled,
            _temp_dir: temp_dir,
        })
    }

    /// Create a backend of the specified type
    pub fn new(backend_type: BackendType) -> Result<Self, Box<dyn std::error::Error>> {
        match backend_type {
            BackendType::Redb => Self::new_redb(),
            BackendType::Sled => Self::new_sled(),
        }
    }

    /// Get the backend name for display
    pub fn name(&self) -> &'static str {
        self.backend_type.name()
    }

    /// Pre-populate the backend with test data
    pub fn populate_strings(&self, data: &[(String, String)]) -> Result<(), Box<dyn std::error::Error>> {
        for (key, value) in data {
            self.storage.set(key.clone(), value.clone())?;
        }
        Ok(())
    }

    /// Pre-populate with hash data
    pub fn populate_hashes(&self, data: &[(String, Vec<(String, String)>)]) -> Result<(), Box<dyn std::error::Error>> {
        for (key, fields) in data {
            self.storage.hset(key, fields.clone())?;
        }
        Ok(())
    }

    /// Pre-populate with list data
    pub fn populate_lists(&self, data: &[(String, Vec<String>)]) -> Result<(), Box<dyn std::error::Error>> {
        for (key, elements) in data {
            self.storage.rpush(key, elements.clone())?;
        }
        Ok(())
    }

    /// Clear all data from the backend
    pub fn clear(&self) -> Result<(), Box<dyn std::error::Error>> {
        self.storage.flushdb()?;
        Ok(())
    }

    /// Get the number of keys in the database
    pub fn dbsize(&self) -> Result<i64, Box<dyn std::error::Error>> {
        Ok(self.storage.dbsize()?)
    }
}

/// Helper function to create and populate a backend for read benchmarks
pub fn setup_populated_backend(
    backend_type: BackendType,
    num_keys: usize,
    value_size: usize,
) -> Result<BenchmarkBackend, Box<dyn std::error::Error>> {
    use super::DataGenerator;
    
    let backend = BenchmarkBackend::new(backend_type)?;
    let mut generator = DataGenerator::new(42);
    let data = generator.generate_string_pairs(num_keys, value_size);
    backend.populate_strings(&data)?;
    
    Ok(backend)
}

/// Helper function to create and populate a backend with hash data
pub fn setup_populated_backend_hashes(
    backend_type: BackendType,
    num_hashes: usize,
    fields_per_hash: usize,
    value_size: usize,
) -> Result<BenchmarkBackend, Box<dyn std::error::Error>> {
    use super::DataGenerator;
    
    let backend = BenchmarkBackend::new(backend_type)?;
    let mut generator = DataGenerator::new(42);
    let data = generator.generate_hash_data(num_hashes, fields_per_hash, value_size);
    backend.populate_hashes(&data)?;
    
    Ok(backend)
}

/// Helper function to create and populate a backend with list data
pub fn setup_populated_backend_lists(
    backend_type: BackendType,
    num_lists: usize,
    elements_per_list: usize,
    element_size: usize,
) -> Result<BenchmarkBackend, Box<dyn std::error::Error>> {
    use super::DataGenerator;
    
    let backend = BenchmarkBackend::new(backend_type)?;
    let mut generator = DataGenerator::new(42);
    let data = generator.generate_list_data(num_lists, elements_per_list, element_size);
    backend.populate_lists(&data)?;
    
    Ok(backend)
}

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

    #[test]
    fn test_backend_creation() {
        let redb = BenchmarkBackend::new_redb();
        assert!(redb.is_ok());

        let sled = BenchmarkBackend::new_sled();
        assert!(sled.is_ok());
    }

    #[test]
    fn test_backend_populate() {
        let backend = BenchmarkBackend::new_redb().unwrap();
        let data = vec![
            ("key1".to_string(), "value1".to_string()),
            ("key2".to_string(), "value2".to_string()),
        ];
        
        backend.populate_strings(&data).unwrap();
        assert_eq!(backend.dbsize().unwrap(), 2);
    }

    #[test]
    fn test_backend_clear() {
        let backend = BenchmarkBackend::new_redb().unwrap();
        let data = vec![("key1".to_string(), "value1".to_string())];
        
        backend.populate_strings(&data).unwrap();
        assert_eq!(backend.dbsize().unwrap(), 1);
        
        backend.clear().unwrap();
        assert_eq!(backend.dbsize().unwrap(), 0);
    }
}