use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct JobContext {
pub job_id: Uuid,
pub job_type: String,
pub retry_count: u32,
pub metadata: HashMap<String, String>,
}
impl JobContext {
pub fn new(job_id: Uuid, job_type: String) -> Self {
Self {
job_id,
job_type,
retry_count: 0,
metadata: HashMap::new(),
}
}
pub fn with_retry_count(mut self, count: u32) -> Self {
self.retry_count = count;
self
}
pub fn add_metadata(mut self, key: String, value: String) -> Self {
self.metadata.insert(key, value);
self
}
}
pub type JobResult = Result<(), Box<dyn std::error::Error + Send + Sync>>;
#[async_trait]
pub trait Job: Send + Sync + Serialize + for<'de> Deserialize<'de> {
async fn execute(&self, ctx: JobContext) -> JobResult;
fn job_type(&self) -> &str;
fn max_retries(&self) -> Option<u32> {
None
}
fn timeout_seconds(&self) -> Option<u64> {
None
}
async fn before_execute(&self, _ctx: &JobContext) -> JobResult {
Ok(())
}
async fn after_execute(&self, _ctx: &JobContext) -> JobResult {
Ok(())
}
async fn on_failure(&self, _ctx: &JobContext, _error: &dyn std::error::Error) {}
}
pub struct JobRegistry {
handlers: Arc<tokio::sync::RwLock<HashMap<String, Box<dyn JobHandler>>>>,
}
impl JobRegistry {
pub fn new() -> Self {
Self {
handlers: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
}
}
pub async fn register<J: Job + 'static>(&self, job_type: &str) {
let mut handlers = self.handlers.write().await;
handlers.insert(
job_type.to_string(),
Box::new(TypedJobHandler::<J>::new()),
);
tracing::info!(job_type = %job_type, "Registered job handler");
}
pub async fn execute(
&self,
job_type: &str,
payload: serde_json::Value,
ctx: JobContext,
) -> JobResult {
let handlers = self.handlers.read().await;
if let Some(handler) = handlers.get(job_type) {
handler.handle(payload, ctx).await
} else {
Err(format!("No handler registered for job type: {}", job_type).into())
}
}
}
impl Default for JobRegistry {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
trait JobHandler: Send + Sync {
async fn handle(&self, payload: serde_json::Value, ctx: JobContext) -> JobResult;
}
struct TypedJobHandler<J: Job> {
_phantom: std::marker::PhantomData<J>,
}
impl<J: Job> TypedJobHandler<J> {
fn new() -> Self {
Self {
_phantom: std::marker::PhantomData,
}
}
}
#[async_trait]
impl<J: Job + 'static> JobHandler for TypedJobHandler<J> {
async fn handle(&self, payload: serde_json::Value, ctx: JobContext) -> JobResult {
let job: J = serde_json::from_value(payload)
.map_err(|e| format!("Failed to deserialize job: {}", e))?;
job.before_execute(&ctx).await?;
let result = job.execute(ctx.clone()).await;
match &result {
Ok(_) => {
job.after_execute(&ctx).await?;
}
Err(e) => {
job.on_failure(&ctx, e.as_ref()).await;
}
}
result
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Serialize, Deserialize)]
struct TestJob {
message: String,
}
#[async_trait]
impl Job for TestJob {
async fn execute(&self, _ctx: JobContext) -> JobResult {
println!("Executing: {}", self.message);
Ok(())
}
fn job_type(&self) -> &str {
"test_job"
}
}
#[tokio::test]
async fn test_job_registry() {
let registry = JobRegistry::new();
registry.register::<TestJob>("test_job").await;
let payload = serde_json::json!({
"message": "Hello, World!"
});
let ctx = JobContext::new(Uuid::new_v4(), "test_job".to_string());
let result = registry.execute("test_job", payload, ctx).await;
assert!(result.is_ok());
}
}