use std::sync::Arc;
use crate::tasks::{TaskError, TaskOutcome, TaskRecord, TaskStatus, store::AbstractTaskStore};
#[derive(Clone)]
pub struct MemoryTaskStore {
tasks: Arc<tokio::sync::RwLock<Vec<TaskRecord>>>,
batch_size: usize,
lease_duration: chrono::Duration,
}
impl MemoryTaskStore {
pub fn new(batch_size: usize) -> Self {
Self::with_lease_duration(batch_size, std::time::Duration::from_secs(300))
}
pub fn with_lease_duration(batch_size: usize, lease_duration: std::time::Duration) -> Self {
Self {
tasks: Arc::new(tokio::sync::RwLock::new(Vec::new())),
batch_size: batch_size.max(1),
lease_duration: chrono::Duration::from_std(lease_duration).unwrap_or_default(),
}
}
pub async fn task_count(&self) -> usize {
self.tasks.read().await.len()
}
pub async fn tasks(&self) -> Vec<TaskRecord> {
self.tasks.read().await.clone()
}
}
fn is_active_status(status: TaskStatus) -> bool {
matches!(
status,
TaskStatus::Pending | TaskStatus::Running | TaskStatus::Suspended
)
}
impl AbstractTaskStore for MemoryTaskStore {
async fn claim_tasks(&self, runner_id: &str) -> Result<Vec<TaskRecord>, TaskError> {
let mut tasks = self.tasks.write().await;
let now = chrono::Utc::now();
let mut ready_indices: Vec<_> = tasks
.iter()
.enumerate()
.filter(|(_, task)| {
(task.status == TaskStatus::Pending
&& task.ready_at.map(|ts| ts <= now).unwrap_or(true))
|| (task.status == TaskStatus::Running
&& task.leased_until.is_some_and(|ts| ts <= now))
})
.map(|(i, task)| {
let ready_at = if task.status == TaskStatus::Running {
task.leased_until.unwrap_or(now)
} else {
task.ready_at.unwrap_or(now)
};
(i, ready_at, task.created_at)
})
.collect();
ready_indices.sort_by_key(|(_, ready_at, created_at)| (*ready_at, *created_at));
let claimed = ready_indices
.iter()
.take(self.batch_size)
.map(|(i, _, _)| {
let task = &mut tasks[*i];
task.status = TaskStatus::Running;
task.locked_by = Some(runner_id.to_string());
let lease_duration = task
.lease_duration_ms
.map(chrono::Duration::milliseconds)
.unwrap_or(self.lease_duration);
task.leased_until = Some(now + lease_duration);
task.updated_at = now;
task.clone()
})
.collect();
Ok(claimed)
}
async fn commit_outcome(
&self,
task_id: uuid::Uuid,
runner_id: &str,
outcome: TaskOutcome,
) -> Result<(), TaskError> {
let mut tasks = self.tasks.write().await;
let Some(task) = tasks.iter_mut().find(|task| {
task.id == task_id
&& task.status == TaskStatus::Running
&& task.locked_by.as_deref() == Some(runner_id)
}) else {
return Ok(());
};
let now = chrono::Utc::now();
match outcome {
TaskOutcome::Complete { result } => {
task.status = TaskStatus::Succeeded;
task.result = Some(result);
task.completed_at = Some(now);
task.ready_at = None;
}
TaskOutcome::Suspend {
topic,
state,
output,
} => {
task.status = TaskStatus::Suspended;
task.resume_topic = Some(topic);
task.state = Some(state);
task.output = output;
task.ready_at = None;
}
TaskOutcome::Sleep { state, delay } => {
task.status = TaskStatus::Pending;
task.state = Some(state);
task.ready_at = Some(now + chrono::Duration::from_std(delay).unwrap_or_default());
}
TaskOutcome::Retry { delay, error } => {
task.attempts += 1;
task.last_error = Some(error);
if task.max_attempts.is_some_and(|max| task.attempts >= max) {
task.status = TaskStatus::Failed;
task.ready_at = None;
task.completed_at = Some(now);
} else {
task.status = TaskStatus::Pending;
let retry_delay = delay
.map(chrono::Duration::from_std)
.transpose()
.unwrap_or_default()
.unwrap_or_else(|| {
task.retry_delay_ms
.map(chrono::Duration::milliseconds)
.unwrap_or_default()
});
task.ready_at = Some(now + retry_delay);
}
}
TaskOutcome::Fail { error } => {
task.status = TaskStatus::Failed;
task.last_error = Some(error);
task.ready_at = None;
task.completed_at = Some(now);
}
}
task.locked_by = None;
task.leased_until = None;
task.updated_at = now;
Ok(())
}
async fn store_task(&self, record: TaskRecord) -> Result<(), TaskError> {
let mut tasks = self.tasks.write().await;
if tasks.iter().any(|task| task.id == record.id) {
return Err(TaskError::AlreadyExists(record.id.to_string()));
}
if let Some(identity) = record.identity.as_deref()
&& is_active_status(record.status)
&& tasks.iter().any(|task| {
task.identity.as_deref() == Some(identity) && is_active_status(task.status)
})
{
return Err(TaskError::IdentityError);
}
tasks.push(record);
Ok(())
}
async fn resume(&self, topic: &str, input: String) -> Result<u64, TaskError> {
let mut tasks = self.tasks.write().await;
let now = chrono::Utc::now();
let mut count = 0;
for task in tasks.iter_mut() {
if task.status == TaskStatus::Suspended && task.resume_topic.as_deref() == Some(topic) {
task.status = TaskStatus::Pending;
task.resume_topic = None;
task.resume_input = Some(input.clone());
task.ready_at = Some(now);
task.updated_at = now;
count += 1;
}
}
Ok(count)
}
async fn run_migrations(&self) -> Result<(), TaskError> {
Ok(())
}
}