pub use crate::api::{AuraEngine, Engine, EngineBuilder, Options, Snapshot};
pub mod types {
pub use crate::storage::{Key, Value, ValuePointer, Entry, Batch, Range};
pub use crate::config::Config;
pub use crate::error::{Error, Result};
}
#[derive(Debug, Clone)]
pub struct EngineStats {
pub total_operations: u64,
pub reads: u64,
pub writes: u64,
pub deletes: u64,
pub memtable_size: usize,
pub wal_size: usize,
pub vlog_size: usize,
pub sst_files: u64,
pub write_amplification: f64,
pub read_amplification: f64,
}
impl Default for EngineStats {
fn default() -> Self {
Self {
total_operations: 0,
reads: 0,
writes: 0,
deletes: 0,
memtable_size: 0,
wal_size: 0,
vlog_size: 0,
sst_files: 0,
write_amplification: 1.0,
read_amplification: 1.0,
}
}
}
#[derive(Debug, Clone)]
pub enum EngineStatus {
Starting,
Running,
ReadOnly,
ShuttingDown,
Error(String),
}
#[derive(Debug, Clone)]
pub struct EngineInfo {
pub version: String,
pub status: EngineStatus,
pub stats: EngineStats,
pub config_summary: String,
}
impl EngineInfo {
pub fn new() -> Self {
Self {
version: env!("CARGO_PKG_VERSION").to_string(),
status: EngineStatus::Starting,
stats: EngineStats::default(),
config_summary: "Default configuration".to_string(),
}
}
}
#[derive(Debug, Clone)]
pub struct HealthCheck {
pub healthy: bool,
pub message: String,
pub timestamp: u64,
}
impl HealthCheck {
pub fn new(healthy: bool, message: String) -> Self {
Self {
healthy,
message,
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64,
}
}
pub fn healthy(message: String) -> Self {
Self::new(true, message)
}
pub fn unhealthy(message: String) -> Self {
Self::new(false, message)
}
}
#[async_trait::async_trait]
pub trait EngineExt: Engine {
async fn info(&self) -> Result<EngineInfo>;
async fn stats(&self) -> Result<EngineStats>;
async fn health_check(&self) -> Result<HealthCheck>;
async fn compact(&self) -> Result<()>;
async fn backup(&self, path: &std::path::Path) -> Result<()>;
async fn restore(&self, path: &std::path::Path) -> Result<()>;
}
#[async_trait::async_trait]
impl EngineExt for AuraEngine {
async fn info(&self) -> Result<EngineInfo> {
let mut info = EngineInfo::new();
info.status = EngineStatus::Running;
info.config_summary = format!("DB: {:?}, WAL: {:?}, VLog: {:?}",
self.config.db_path,
self.config.wal.wal_path,
self.config.value_log.vlog_path
);
Ok(info)
}
async fn stats(&self) -> Result<EngineStats> {
let mut stats = EngineStats::default();
let memtable = self.memtable.read();
stats.memtable_size = memtable.memory_usage();
Ok(stats)
}
async fn health_check(&self) -> Result<HealthCheck> {
if *self.closed.read() {
return Ok(HealthCheck::unhealthy("Engine is closed".to_string()));
}
let memtable = self.memtable.read();
if memtable.is_empty() {
}
Ok(HealthCheck::healthy("Engine is healthy".to_string()))
}
async fn compact(&self) -> Result<()> {
self.flush_memtable().await?;
Ok(())
}
async fn backup(&self, _path: &std::path::Path) -> Result<()> {
Err(crate::error::Error::Unknown("Backup not implemented yet".to_string()))
}
async fn restore(&self, _path: &std::path::Path) -> Result<()> {
Err(crate::error::Error::Unknown("Restore not implemented yet".to_string()))
}
}
pub struct AdvancedEngineBuilder {
config: crate::config::Config,
}
impl AdvancedEngineBuilder {
pub fn new() -> Self {
Self {
config: crate::config::Config::default(),
}
}
pub fn with_db_path(mut self, path: std::path::PathBuf) -> Self {
self.config.db_path = path;
self
}
pub fn with_wal_config(mut self, wal_config: crate::config::WalConfig) -> Self {
self.config.wal = wal_config;
self
}
pub fn with_vlog_config(mut self, vlog_config: crate::config::ValueLogConfig) -> Self {
self.config.value_log = vlog_config;
self
}
pub fn with_memtable_config(mut self, memtable_config: crate::config::MemtableConfig) -> Self {
self.config.memtable = memtable_config;
self
}
pub fn with_sst_config(mut self, sst_config: crate::config::SstConfig) -> Self {
self.config.sst = sst_config;
self
}
pub fn with_compaction_config(mut self, compaction_config: crate::config::CompactionConfig) -> Self {
self.config.compaction = compaction_config;
self
}
pub fn with_cache_config(mut self, cache_config: crate::config::CacheConfig) -> Self {
self.config.cache = cache_config;
self
}
pub fn with_learned_index_config(mut self, learned_index_config: crate::config::LearnedIndexConfig) -> Self {
self.config.learned_index = learned_index_config;
self
}
pub fn with_rl_agent_config(mut self, rl_agent_config: crate::config::RlAgentConfig) -> Self {
self.config.rl_agent = rl_agent_config;
self
}
pub fn with_performance_config(mut self, performance_config: crate::config::PerformanceConfig) -> Self {
self.config.performance = performance_config;
self
}
pub async fn build(self) -> Result<AuraEngine> {
AuraEngine::new(self.config).await
}
}
impl Default for AdvancedEngineBuilder {
fn default() -> Self {
Self::new()
}
}
pub async fn create_engine(db_path: std::path::PathBuf) -> Result<AuraEngine> {
EngineBuilder::new()
.with_db_path(db_path)
.build()
.await
}
pub async fn create_engine_with_config(config: crate::config::Config) -> Result<AuraEngine> {
AuraEngine::new(config).await
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_engine_ext() {
let temp_dir = tempdir().unwrap();
let engine = create_engine(temp_dir.path().to_path_buf()).await.unwrap();
let info = engine.info().await.unwrap();
assert_eq!(info.version, env!("CARGO_PKG_VERSION"));
let stats = engine.stats().await.unwrap();
assert_eq!(stats.total_operations, 0);
let health = engine.health_check().await.unwrap();
assert!(health.healthy);
engine.close().await.unwrap();
}
#[tokio::test]
async fn test_advanced_builder() {
let temp_dir = tempdir().unwrap();
let engine = AdvancedEngineBuilder::new()
.with_db_path(temp_dir.path().to_path_buf())
.build()
.await
.unwrap();
assert!(engine.info().await.is_ok());
engine.close().await.unwrap();
}
}