use async_trait::async_trait;
use sled::Db;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info, trace, warn};
use super::backend::StateBackend;
use super::compression::{compress, decompress, CompressionConfig};
use crate::error::{StateError, StateResult};
#[derive(Debug, Clone, Default)]
pub struct SledBackendStats {
pub get_count: u64,
pub put_count: u64,
pub delete_count: u64,
pub flush_count: u64,
pub bytes_written: u64,
pub bytes_read: u64,
pub bytes_compressed: u64,
pub bytes_decompressed: u64,
pub bytes_before_compression: u64,
pub bytes_after_compression: u64,
}
#[derive(Debug, Clone)]
pub struct SledConfig {
pub path: PathBuf,
pub cache_capacity: u64,
pub compression: CompressionConfig,
pub flush_every: u64,
pub segment_size: usize,
}
impl Default for SledConfig {
fn default() -> Self {
Self {
path: PathBuf::from("/tmp/sled_state"),
cache_capacity: 128 * 1024 * 1024, compression: CompressionConfig::balanced(),
flush_every: 1000,
segment_size: 8 * 1024 * 1024, }
}
}
impl SledConfig {
pub fn new<P: AsRef<Path>>(path: P) -> Self {
Self {
path: path.as_ref().to_path_buf(),
..Default::default()
}
}
pub fn with_cache_capacity(mut self, capacity: u64) -> Self {
self.cache_capacity = capacity;
self
}
pub fn with_compression(mut self, config: CompressionConfig) -> Self {
self.compression = config;
self
}
pub fn without_compression(mut self) -> Self {
self.compression = CompressionConfig::none();
self
}
pub fn with_fast_compression(mut self) -> Self {
self.compression = CompressionConfig::fast();
self
}
pub fn with_balanced_compression(mut self) -> Self {
self.compression = CompressionConfig::balanced();
self
}
pub fn with_best_compression(mut self) -> Self {
self.compression = CompressionConfig::best_size();
self
}
pub fn with_flush_every(mut self, count: u64) -> Self {
self.flush_every = count;
self
}
}
pub struct SledStateBackend {
db: Arc<Db>,
config: SledConfig,
stats: Arc<RwLock<SledBackendStats>>,
op_counter: Arc<RwLock<u64>>,
}
impl SledStateBackend {
pub async fn open(config: SledConfig) -> StateResult<Self> {
info!("Opening Sled database at {:?}", config.path);
if let Some(parent) = config.path.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Failed to create directory: {}", e),
}
})?;
}
let db = sled::Config::new()
.path(&config.path)
.cache_capacity(config.cache_capacity)
.use_compression(false) .segment_size(config.segment_size)
.open()
.map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Failed to open database: {}", e),
})?;
info!(
"Sled database opened successfully, {} keys",
db.len()
);
Ok(Self {
db: Arc::new(db),
config,
stats: Arc::new(RwLock::new(SledBackendStats::default())),
op_counter: Arc::new(RwLock::new(0)),
})
}
pub async fn temporary() -> StateResult<Self> {
let temp_dir = std::env::temp_dir().join(format!("sled_test_{}", uuid::Uuid::new_v4()));
Self::open(SledConfig::new(temp_dir)).await
}
pub async fn flush(&self) -> StateResult<usize> {
trace!("Flushing Sled database");
let flushed = tokio::task::block_in_place(|| {
self.db.flush().map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Flush failed: {}", e),
})
})?;
let mut stats = self.stats.write().await;
stats.flush_count += 1;
debug!("Flushed {} bytes to disk", flushed);
Ok(flushed)
}
pub async fn stats(&self) -> SledBackendStats {
self.stats.read().await.clone()
}
pub async fn compression_ratio(&self) -> Option<f64> {
let stats = self.stats.read().await;
if stats.bytes_before_compression > 0 {
Some((stats.bytes_after_compression as f64 / stats.bytes_before_compression as f64) * 100.0)
} else {
None
}
}
pub async fn space_saved(&self) -> u64 {
let stats = self.stats.read().await;
stats.bytes_before_compression.saturating_sub(stats.bytes_after_compression)
}
pub async fn size_on_disk(&self) -> StateResult<u64> {
let path = &self.config.path;
let mut total_size = 0u64;
let entries = tokio::fs::read_dir(path).await.map_err(|e| {
StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Failed to read directory: {}", e),
}
})?;
let mut entries = entries;
while let Some(entry) = entries.next_entry().await.map_err(|e| {
StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Failed to read entry: {}", e),
}
})? {
if let Ok(metadata) = entry.metadata().await {
total_size += metadata.len();
}
}
Ok(total_size)
}
pub async fn checkpoint(&self) -> StateResult<Vec<(Vec<u8>, Vec<u8>)>> {
info!("Creating Sled checkpoint");
let checkpoint = tokio::task::block_in_place(|| {
self.db
.iter()
.map(|result| {
result
.map(|(k, v)| (k.to_vec(), v.to_vec()))
.map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Checkpoint iteration failed: {}", e),
})
})
.collect::<StateResult<Vec<_>>>()
})?;
info!("Checkpoint created with {} entries", checkpoint.len());
Ok(checkpoint)
}
pub async fn restore(&self, checkpoint: Vec<(Vec<u8>, Vec<u8>)>) -> StateResult<()> {
info!("Restoring Sled from checkpoint with {} entries", checkpoint.len());
tokio::task::block_in_place(|| {
self.db.clear().map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Clear failed: {}", e),
})?;
for (key, value) in checkpoint {
self.db.insert(key, value).map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Insert failed during restore: {}", e),
})?;
}
self.db.flush().map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Flush failed after restore: {}", e),
})?;
Ok(())
})?;
info!("Restore completed successfully");
Ok(())
}
async fn maybe_flush(&self) -> StateResult<()> {
let mut counter = self.op_counter.write().await;
*counter += 1;
if *counter >= self.config.flush_every {
*counter = 0;
drop(counter); self.flush().await?;
}
Ok(())
}
}
#[async_trait]
impl StateBackend for SledStateBackend {
async fn get(&self, key: &[u8]) -> StateResult<Option<Vec<u8>>> {
trace!("Getting key from Sled: {:?}", key);
let result = tokio::task::block_in_place(|| {
self.db.get(key).map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Get failed: {}", e),
})
})?;
let mut stats = self.stats.write().await;
stats.get_count += 1;
if let Some(compressed_value) = result {
let compressed_len = compressed_value.len() as u64;
stats.bytes_read += compressed_len;
let decompressed = decompress(&compressed_value, &self.config.compression)
.map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Decompression failed: {}", e),
})?;
let decompressed_len = decompressed.len() as u64;
stats.bytes_decompressed += decompressed_len;
trace!(
"Decompressed {} bytes to {} bytes (ratio: {:.2}%)",
compressed_len,
decompressed_len,
(compressed_len as f64 / decompressed_len as f64) * 100.0
);
Ok(Some(decompressed))
} else {
Ok(None)
}
}
async fn put(&self, key: &[u8], value: &[u8]) -> StateResult<()> {
trace!("Putting key to Sled: {:?}, value size: {} bytes", key, value.len());
let original_len = value.len() as u64;
let compressed = compress(value, &self.config.compression)
.map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Compression failed: {}", e),
})?;
let compressed_len = compressed.len() as u64;
trace!(
"Compressed {} bytes to {} bytes (ratio: {:.2}%)",
original_len,
compressed_len,
(compressed_len as f64 / original_len as f64) * 100.0
);
tokio::task::block_in_place(|| {
self.db.insert(key, compressed.as_slice()).map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Put failed: {}", e),
})
})?;
let mut stats = self.stats.write().await;
stats.put_count += 1;
stats.bytes_written += (key.len() as u64 + compressed_len);
stats.bytes_compressed += original_len;
stats.bytes_before_compression += original_len;
stats.bytes_after_compression += compressed_len;
drop(stats);
self.maybe_flush().await?;
Ok(())
}
async fn delete(&self, key: &[u8]) -> StateResult<()> {
trace!("Deleting key from Sled: {:?}", key);
tokio::task::block_in_place(|| {
self.db.remove(key).map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Delete failed: {}", e),
})
})?;
let mut stats = self.stats.write().await;
stats.delete_count += 1;
drop(stats);
self.maybe_flush().await?;
Ok(())
}
async fn list_keys(&self, prefix: &[u8]) -> StateResult<Vec<Vec<u8>>> {
trace!("Listing keys with prefix: {:?}", prefix);
let keys = tokio::task::block_in_place(|| {
if prefix.is_empty() {
self.db
.iter()
.keys()
.map(|result| {
result
.map(|k| k.to_vec())
.map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Key iteration failed: {}", e),
})
})
.collect::<StateResult<Vec<_>>>()
} else {
self.db
.scan_prefix(prefix)
.keys()
.map(|result| {
result
.map(|k| k.to_vec())
.map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Prefix scan failed: {}", e),
})
})
.collect::<StateResult<Vec<_>>>()
}
})?;
Ok(keys)
}
async fn clear(&self) -> StateResult<()> {
debug!("Clearing Sled database");
tokio::task::block_in_place(|| {
self.db.clear().map_err(|e| StateError::StorageError {
backend_type: "sled".to_string(),
details: format!("Clear failed: {}", e),
})
})?;
self.flush().await?;
Ok(())
}
async fn count(&self) -> StateResult<usize> {
Ok(self.db.len())
}
async fn contains(&self, key: &[u8]) -> StateResult<bool> {
Ok(tokio::task::block_in_place(|| {
self.db.contains_key(key).unwrap_or(false)
}))
}
}
impl Drop for SledStateBackend {
fn drop(&mut self) {
if let Err(e) = self.db.flush() {
error!("Failed to flush Sled database on drop: {}", e);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::state::backend::tests::*;
async fn create_test_backend() -> SledStateBackend {
SledStateBackend::temporary().await.unwrap()
}
#[tokio::test]
async fn test_sled_backend_basic() {
let backend = create_test_backend().await;
test_backend_basic_ops(backend).await;
}
#[tokio::test]
async fn test_sled_backend_list_keys() {
let backend = create_test_backend().await;
test_backend_list_keys(backend).await;
}
#[tokio::test]
async fn test_sled_backend_clear() {
let backend = create_test_backend().await;
test_backend_clear(backend).await;
}
#[tokio::test]
async fn test_sled_backend_contains() {
let backend = create_test_backend().await;
test_backend_contains(backend).await;
}
#[tokio::test]
async fn test_sled_persistence() {
let temp_dir = std::env::temp_dir().join(format!("sled_persist_test_{}", uuid::Uuid::new_v4()));
let config = SledConfig::new(&temp_dir);
{
let backend = SledStateBackend::open(config.clone()).await.unwrap();
backend.put(b"key1", b"value1").await.unwrap();
backend.put(b"key2", b"value2").await.unwrap();
backend.flush().await.unwrap();
}
{
let backend = SledStateBackend::open(config).await.unwrap();
assert_eq!(backend.count().await.unwrap(), 2);
assert_eq!(
backend.get(b"key1").await.unwrap(),
Some(b"value1".to_vec())
);
assert_eq!(
backend.get(b"key2").await.unwrap(),
Some(b"value2".to_vec())
);
}
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
#[tokio::test]
async fn test_checkpoint_restore() {
let backend = create_test_backend().await;
backend.put(b"key1", b"value1").await.unwrap();
backend.put(b"key2", b"value2").await.unwrap();
let checkpoint = backend.checkpoint().await.unwrap();
assert_eq!(checkpoint.len(), 2);
backend.clear().await.unwrap();
assert_eq!(backend.count().await.unwrap(), 0);
backend.restore(checkpoint).await.unwrap();
assert_eq!(backend.count().await.unwrap(), 2);
assert_eq!(
backend.get(b"key1").await.unwrap(),
Some(b"value1".to_vec())
);
}
#[tokio::test]
async fn test_stats() {
let backend = create_test_backend().await;
backend.put(b"key1", b"value1").await.unwrap();
backend.put(b"key2", b"value2").await.unwrap();
backend.get(b"key1").await.unwrap();
backend.delete(b"key1").await.unwrap();
let stats = backend.stats().await;
assert_eq!(stats.put_count, 2);
assert_eq!(stats.get_count, 1);
assert_eq!(stats.delete_count, 1);
assert!(stats.bytes_written > 0);
assert!(stats.bytes_read > 0);
}
#[tokio::test]
async fn test_size_on_disk() {
let backend = create_test_backend().await;
let initial_size = backend.size_on_disk().await.unwrap();
for i in 0..100 {
let key = format!("key:{}", i).into_bytes();
let value = vec![0u8; 1000]; backend.put(&key, &value).await.unwrap();
}
backend.flush().await.unwrap();
let final_size = backend.size_on_disk().await.unwrap();
assert!(final_size > initial_size);
}
#[tokio::test]
async fn test_concurrent_access() {
use std::sync::Arc;
let backend = Arc::new(create_test_backend().await);
let mut handles = vec![];
for i in 0..10 {
let backend = Arc::clone(&backend);
let handle = tokio::spawn(async move {
for j in 0..10 {
let key = format!("key:{}:{}", i, j).into_bytes();
let value = format!("value:{}:{}", i, j).into_bytes();
backend.put(&key, &value).await.unwrap();
}
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
backend.flush().await.unwrap();
assert_eq!(backend.count().await.unwrap(), 100);
}
#[tokio::test]
async fn test_compression_roundtrip() {
let backend = create_test_backend().await;
let test_data = b"This is some test data that should be compressed";
backend.put(b"test_key", test_data).await.unwrap();
let result = backend.get(b"test_key").await.unwrap();
assert_eq!(result, Some(test_data.to_vec()));
}
#[tokio::test]
async fn test_compression_algorithms() {
let test_data = vec![b'x'; 10000];
{
let temp_dir = std::env::temp_dir().join(format!("sled_lz4_test_{}", uuid::Uuid::new_v4()));
let config = SledConfig::new(&temp_dir).with_fast_compression();
let backend = SledStateBackend::open(config).await.unwrap();
backend.put(b"key", &test_data).await.unwrap();
let result = backend.get(b"key").await.unwrap().unwrap();
assert_eq!(result, test_data);
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
{
let temp_dir = std::env::temp_dir().join(format!("sled_snappy_test_{}", uuid::Uuid::new_v4()));
let config = SledConfig::new(&temp_dir).with_balanced_compression();
let backend = SledStateBackend::open(config).await.unwrap();
backend.put(b"key", &test_data).await.unwrap();
let result = backend.get(b"key").await.unwrap().unwrap();
assert_eq!(result, test_data);
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
{
let temp_dir = std::env::temp_dir().join(format!("sled_zstd_test_{}", uuid::Uuid::new_v4()));
let config = SledConfig::new(&temp_dir).with_best_compression();
let backend = SledStateBackend::open(config).await.unwrap();
backend.put(b"key", &test_data).await.unwrap();
let result = backend.get(b"key").await.unwrap().unwrap();
assert_eq!(result, test_data);
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
{
let temp_dir = std::env::temp_dir().join(format!("sled_none_test_{}", uuid::Uuid::new_v4()));
let config = SledConfig::new(&temp_dir).without_compression();
let backend = SledStateBackend::open(config).await.unwrap();
backend.put(b"key", &test_data).await.unwrap();
let result = backend.get(b"key").await.unwrap().unwrap();
assert_eq!(result, test_data);
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
}
#[tokio::test]
async fn test_compression_stats() {
let backend = create_test_backend().await;
let data = vec![b'A'; 10000]; backend.put(b"key", &data).await.unwrap();
let stats = backend.stats().await;
assert!(stats.bytes_compressed > 0);
assert!(stats.bytes_before_compression >= 10000);
assert!(stats.bytes_after_compression > 0);
assert!(stats.bytes_after_compression < stats.bytes_before_compression);
}
#[tokio::test]
async fn test_compression_ratio() {
let backend = create_test_backend().await;
let data = vec![0u8; 100000]; backend.put(b"key", &data).await.unwrap();
let ratio = backend.compression_ratio().await;
assert!(ratio.is_some());
let ratio_value = ratio.unwrap();
assert!(ratio_value > 0.0);
assert!(ratio_value < 100.0);
println!("Compression ratio: {:.2}%", ratio_value);
}
#[tokio::test]
async fn test_space_saved() {
let backend = create_test_backend().await;
let data = vec![b'X'; 50000]; backend.put(b"key", &data).await.unwrap();
let space_saved = backend.space_saved().await;
assert!(space_saved > 0);
println!("Space saved: {} bytes", space_saved);
}
#[tokio::test]
async fn test_large_data_compression() {
let backend = create_test_backend().await;
let large_data = vec![b'Z'; 1024 * 1024];
backend.put(b"large_key", &large_data).await.unwrap();
let result = backend.get(b"large_key").await.unwrap().unwrap();
assert_eq!(result.len(), large_data.len());
assert_eq!(result, large_data);
let stats = backend.stats().await;
let compression_ratio = (stats.bytes_after_compression as f64
/ stats.bytes_before_compression as f64) * 100.0;
println!("Large data compression ratio: {:.2}%", compression_ratio);
assert!(compression_ratio < 10.0, "Compression ratio should be less than 10% for repeated data");
}
#[tokio::test]
async fn test_mixed_data_compression() {
let backend = create_test_backend().await;
let compressible = vec![b'A'; 5000];
let random = (0..5000).map(|i| (i % 256) as u8).collect::<Vec<_>>();
backend.put(b"compressible", &compressible).await.unwrap();
backend.put(b"random", &random).await.unwrap();
assert_eq!(
backend.get(b"compressible").await.unwrap().unwrap(),
compressible
);
assert_eq!(
backend.get(b"random").await.unwrap().unwrap(),
random
);
let stats = backend.stats().await;
assert!(stats.bytes_compressed > 0);
assert!(stats.bytes_decompressed > 0);
}
#[tokio::test]
async fn test_compression_persistence() {
let temp_dir = std::env::temp_dir().join(format!("sled_compress_persist_{}", uuid::Uuid::new_v4()));
let config = SledConfig::new(&temp_dir).with_best_compression();
let test_data = vec![b'B'; 20000];
{
let backend = SledStateBackend::open(config.clone()).await.unwrap();
backend.put(b"key", &test_data).await.unwrap();
backend.flush().await.unwrap();
}
{
let backend = SledStateBackend::open(config).await.unwrap();
let result = backend.get(b"key").await.unwrap().unwrap();
assert_eq!(result, test_data);
}
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
#[tokio::test]
async fn test_compression_with_custom_config() {
use super::compression::{CompressionAlgorithm, CompressionConfig};
let temp_dir = std::env::temp_dir().join(format!("sled_custom_compress_{}", uuid::Uuid::new_v4()));
let compression_config = CompressionConfig {
algorithm: CompressionAlgorithm::Zstd,
level: 15,
threshold: 512,
enable_stats: true,
};
let config = SledConfig::new(&temp_dir).with_compression(compression_config);
let backend = SledStateBackend::open(config).await.unwrap();
let small_data = vec![b'S'; 256];
backend.put(b"small", &small_data).await.unwrap();
let large_data = vec![b'L'; 10000];
backend.put(b"large", &large_data).await.unwrap();
assert_eq!(
backend.get(b"small").await.unwrap().unwrap(),
small_data
);
assert_eq!(
backend.get(b"large").await.unwrap().unwrap(),
large_data
);
tokio::fs::remove_dir_all(&temp_dir).await.ok();
}
#[tokio::test]
async fn test_empty_value_compression() {
let backend = create_test_backend().await;
backend.put(b"empty", b"").await.unwrap();
let result = backend.get(b"empty").await.unwrap();
assert_eq!(result, Some(vec![]));
}
#[tokio::test]
async fn test_compression_stats_accuracy() {
let backend = create_test_backend().await;
let data1 = vec![b'1'; 5000];
let data2 = vec![b'2'; 8000];
let data3 = vec![b'3'; 12000];
backend.put(b"key1", &data1).await.unwrap();
backend.put(b"key2", &data2).await.unwrap();
backend.put(b"key3", &data3).await.unwrap();
let stats = backend.stats().await;
assert!(stats.bytes_before_compression >= 25000);
assert!(stats.bytes_after_compression < stats.bytes_before_compression);
backend.get(b"key1").await.unwrap();
backend.get(b"key2").await.unwrap();
backend.get(b"key3").await.unwrap();
let stats = backend.stats().await;
assert!(stats.bytes_decompressed >= 25000);
}
}