use std::time::Duration;
use a2a_rs::adapter::storage::InMemoryTaskStorage;
use a2a_rs::domain::TaskState;
use a2a_rs::port::AsyncTaskLifecycle;
#[cfg(feature = "sqlx-storage")]
use a2a_rs::adapter::storage::{DatabaseConfig, SqlxTaskStorage};
fn tid(s: &str) -> a2a_rs::domain::TaskId {
s.parse().unwrap()
}
fn cid(s: &str) -> a2a_rs::domain::ContextId {
s.parse().unwrap()
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing_subscriber::fmt::init();
println!("đĒ Storage Implementation Comparison");
println!("===================================");
println!();
println!("đ§ Testing InMemory Storage");
println!("---------------------------");
let in_memory_storage = InMemoryTaskStorage::new();
run_storage_tests(&in_memory_storage, "InMemory").await?;
println!();
#[cfg(feature = "sqlx-storage")]
{
println!("đī¸ Testing SQLx Storage");
println!("-----------------------");
let config = DatabaseConfig::from_env().unwrap_or_else(|_| {
println!("âšī¸ Using default SQLite in-memory configuration");
DatabaseConfig::default()
});
println!("đ Database: {} ({:?})", config.url, config.database_type());
let sqlx_storage = SqlxTaskStorage::new(&config.url).await?;
run_storage_tests(&sqlx_storage, "SQLx").await?;
println!();
}
#[cfg(not(feature = "sqlx-storage"))]
{
println!(
"â ī¸ SQLx storage not available - compile with --features sqlite or --features postgres"
);
println!();
}
println!("⥠Performance Comparison");
println!("------------------------");
println!("đ§ InMemory Storage Performance:");
let in_memory_storage = InMemoryTaskStorage::new();
let in_memory_time = measure_performance(&in_memory_storage).await?;
println!(" Time for 100 operations: {:?}", in_memory_time);
#[cfg(feature = "sqlx-storage")]
{
println!("đī¸ SQLx Storage Performance:");
let config = DatabaseConfig::from_env().unwrap_or_default();
let sqlx_storage = SqlxTaskStorage::new(&config.url).await?;
let sqlx_time = measure_performance(&sqlx_storage).await?;
println!(" Time for 100 operations: {:?}", sqlx_time);
let ratio = sqlx_time.as_secs_f64() / in_memory_time.as_secs_f64();
println!(
" SQLx is {:.1}x {} than InMemory",
if ratio > 1.0 { ratio } else { 1.0 / ratio },
if ratio > 1.0 { "slower" } else { "faster" }
);
}
println!();
println!("đ Summary");
println!("----------");
println!("InMemory Storage:");
println!(" â
Fast performance");
println!(" â
No external dependencies");
println!(" â Data lost on restart");
println!(" â Single process only");
println!();
#[cfg(feature = "sqlx-storage")]
{
println!("SQLx Storage:");
println!(" â
Persistent data");
println!(" â
Multi-process support");
println!(" â
ACID transactions");
println!(" â
Supports SQLite, PostgreSQL, MySQL");
println!(" â Requires database setup");
println!(" â Slightly slower than in-memory");
println!();
}
println!("đĄ Use InMemory for development/testing, SQLx for production");
Ok(())
}
async fn run_storage_tests<T: AsyncTaskLifecycle>(
storage: &T,
storage_name: &str,
) -> Result<(), Box<dyn std::error::Error>> {
let task_id = format!("test-task-{}", storage_name.to_lowercase());
let task = storage.create(&tid(&task_id), &cid("test-context")).await?;
println!(
" â Created task: {} (status: {:?})",
task.id, task.status.state
);
let exists = storage.exists(&tid(&task_id)).await?;
println!(" â Task exists: {}", exists);
let updated_task = storage
.update_status(&tid(&task_id), TaskState::Working, None)
.await?;
println!(
" â Updated to Working (status: {:?})",
updated_task.status.state
);
let task_with_history = storage.get(&tid(&task_id), Some(10)).await?;
let history_count = task_with_history.history.len();
println!(" â Retrieved task with {} history entries", history_count);
let completed_task = storage
.update_status(&tid(&task_id), TaskState::Completed, None)
.await?;
println!(
" â Completed task (status: {:?})",
completed_task.status.state
);
match storage.cancel(&tid(&task_id)).await {
Ok(_) => println!(" â Unexpected: was able to cancel completed task"),
Err(_) => println!(" â Correctly prevented canceling completed task"),
}
let cancel_task_id = format!("cancel-test-{}", storage_name.to_lowercase());
storage
.create(&tid(&cancel_task_id), &cid("test-context"))
.await?;
storage
.update_status(&tid(&cancel_task_id), TaskState::Working, None)
.await?;
let canceled_task = storage.cancel(&tid(&cancel_task_id)).await?;
println!(
" â Canceled working task (status: {:?})",
canceled_task.status.state
);
Ok(())
}
async fn measure_performance<T: AsyncTaskLifecycle>(
storage: &T,
) -> Result<Duration, Box<dyn std::error::Error>> {
let start = std::time::Instant::now();
for i in 0..100 {
let task_id = format!("perf-task-{}", i);
storage.create(&tid(&task_id), &cid("perf-context")).await?;
storage
.update_status(&tid(&task_id), TaskState::Working, None)
.await?;
storage
.update_status(&tid(&task_id), TaskState::Completed, None)
.await?;
storage.get(&tid(&task_id), Some(5)).await?;
}
Ok(start.elapsed())
}