diff --git a/CHANGELOG.md b/CHANGELOG.md index 79086fa4..f38429ce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Introduced a pluggable result backend API with a Redis implementation, enhanced + `AsyncResult` helpers, and a `redis_results` example with docs. - Introduced a Redis-backed distributed scheduler backend with configurable lock management and task state persistence, plus documentation and examples for running multi-instance beat deployments. @@ -28,6 +30,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Removed manual sccache configuration in favor of built-in caching - Temporarily disabled minimal versions check due to upstream regex-syntax compatibility issue - Removed global `-D warnings` rustflags to prevent builds failing on non-critical warnings +- Reworked README quick-start and example sections for clearer onboarding. ## [v0.5.5](https://github.com/rusty-celery/rusty-celery/releases/tag/v0.5.5) - 2023-09-25 diff --git a/Cargo.toml b/Cargo.toml index 010d07f2..136a97b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,9 @@ name = "celery_app" [[example]] name = "redis_beat" +[[example]] +name = "redis_results" + [dependencies] base64 = "0.22.1" chrono = { version = "0.4.42", features = ["serde"] } @@ -44,7 +47,7 @@ log = "0.4.28" futures = { version = "0.3.31", features = ["async-await"] } uuid = { version = "1.18.1", features = ["v4"]} rand = "0.8" -celery-rs-codegen = { version = "0.6.0", path = "./celery-codegen", optional = true } +celery-rs-codegen = { version = "0.6.1", path = "./celery-codegen", optional = true } colored = "3.0.0" once_cell = { version = "1.21.3" } globset = "0.4.16" diff --git a/README.md b/README.md index 19587a85..17bb9f52 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,9 @@ We welcome contributions from everyone regardless of your experience level with If you already know the basics of Rust but are new to Celery, check out the [Rusty Celery Book](https://rusty-celery.github.io/) or the original Python [Celery Project](http://www.celeryproject.org/). -## Quick start +## Getting Started + +### Quick start Define tasks by decorating functions with the [`task`](https://docs.rs/celery-rs/*/celery/attr.task.html) attribute. @@ -73,13 +75,60 @@ And consume tasks as a worker from a queue with my_app.consume().await?; ``` -## Examples +### Capturing results + +Configure a result backend to persist task state and fetch results from the client side: + +```rust +use std::time::Duration; +use celery::backend::RedisBackend; +use celery::prelude::*; + +#[celery::task] +fn add(x: i32, y: i32) -> TaskResult { + Ok(x + y) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let backend = RedisBackend::new("redis://127.0.0.1/0")? + .with_result_ttl(Duration::from_secs(600)); + + let app = celery::app!( + broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, + tasks = [add], + task_routes = [ "*" => "celery" ], + result_backend = backend, + ).await?; + + let async_result = app.send_task(add::new(1, 2)).await?; + println!("state = {}", async_result.state().await?); + + let sum: i32 = async_result.get(Some(Duration::from_secs(10))).await?; + println!("1 + 2 = {sum}"); + Ok(()) +} +``` + +[`AsyncResult`](https://docs.rs/celery-rs/latest/celery/task/struct.AsyncResult.html) now +exposes idiomatic helpers: `state()` for the latest `TaskState`, `ready()` to check completion, +and `get(timeout)` to await the final value (raising a `BackendError::Timeout` on expiration). + +### Example catalog The [`examples/`](https://github.com/GaiaNet-AI/celery-rs/tree/main/examples) directory contains: - a simple Celery app implemented in Rust using an AMQP broker ([`examples/celery_app.rs`](https://github.com/GaiaNet-AI/celery-rs/blob/main/examples/celery_app.rs)), - the same Celery app implemented in Python ([`examples/celery_app.py`](https://github.com/GaiaNet-AI/celery-rs/blob/main/examples/celery_app.py)), -- and a Beat app implemented in Rust ([`examples/beat_app.rs`](https://github.com/GaiaNet-AI/celery-rs/blob/main/examples/beat_app.rs)). +- a Redis result-backend demo showing AsyncResult usage ([`examples/redis_results.rs`](https://github.com/GaiaNet-AI/celery-rs/blob/main/examples/redis_results.rs)), +- a Beat app implemented in Rust ([`examples/beat_app.rs`](https://github.com/GaiaNet-AI/celery-rs/blob/main/examples/beat_app.rs)), +- and a Redis-backed Beat scheduler with leader election ([`examples/redis_beat.rs`](https://github.com/GaiaNet-AI/celery-rs/blob/main/examples/redis_beat.rs)). + +## Running the Examples + +Explore the demos interactively (preview below): + +![](./img/demo.gif) ### Prerequisites @@ -93,11 +142,7 @@ Otherwise simply run the helper script: This will download and run the official [RabbitMQ](https://www.rabbitmq.com/) image (RabbitMQ is a popular AMQP broker). -### Run the examples - -![](./img/demo.gif) - -#### Run Rust Celery app +### Run the Rust Celery app You can consume tasks with: @@ -113,7 +158,7 @@ cargo run --example celery_app produce [task_name] Current supported tasks for this example are: `add`, `buggy_task`, `long_running_task` and `bound_task` -#### Run Python Celery app +### Run the Python Celery app Similarly, you can consume or produce tasks from Python by running @@ -130,7 +175,7 @@ python examples/celery_app.py produce You'll need to have Python 3 installed, along with the requirements listed in the `requirements.txt` file. You'll also have to provide a task name. This example implements 4 tasks: `add`, `buggy_task`, `long_running_task` and `bound_task` -#### Run Rust Beat app +### Run the Rust Beat app You can start the Rust beat with: @@ -140,7 +185,7 @@ cargo run --example beat_app And then you can consume tasks from Rust or Python as explained above. -#### Redis-backed Beat +### Redis-backed Beat failover A Redis-powered distributed scheduler backend is available through `RedisSchedulerBackend`. To try it out locally (requires a Redis server running): @@ -166,33 +211,18 @@ To test multi-instance failover: ⚠️ = Partially implemented and under active development.
🔴 = Not supported yet but on-deck to be implemented soon. -### Core - > **Note**: Issue tracking links below reference this repository. -| | Status | Tracking | -| ---------------- |:-------:| --------- | -| Protocol | ⚠️ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Protocol%20Feature%22+is%3Aopen) | -| Producers | ✅ | | -| Consumers | ✅ | | -| Brokers | ✅ | | -| Beat | ✅ | | -| Backends | ⚠️ | | -| Baskets | 🔴 | | - -### Brokers - -| | Status | Tracking | -| ----- |:------:| -------- | -| AMQP | ✅ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Broker%3A%20AMQP%22+is%3Aopen) | -| Redis | ✅ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Broker%3A%20Redis%22+is%3Aopen) | - -### Backends - -| | Status | Tracking | -| ----------- |:------:| -------- | -| RPC | 🔴 | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Backend%3A%20RPC%22+is%3Aopen) | -| Redis | ✅ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Backend%3A%20Redis%22+is%3Aopen) | +| Area | Component | Status | Notes / Tracking | +|-----------|------------|:------:|------------------| +| Core | Protocol | ⚠️ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Protocol%20Feature%22+is%3Aopen) | +| Core | Producers | ✅ | | +| Core | Consumers | ✅ | | +| Core | Beat | ✅ | | +| Brokers | AMQP | ✅ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Broker%3A%20AMQP%22+is%3Aopen) | +| Brokers | Redis | ✅ | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Broker%3A%20Redis%22+is%3Aopen) | +| Backends | RPC | 🔴 | [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Backend%3A%20RPC%22+is%3Aopen) | +| Backends | Redis | ✅ | Task results + Beat (0.6.2); [Open issues](https://github.com/GaiaNet-AI/celery-rs/issues?q=is%3Aissue+label%3A%22Backend%3A%20Redis%22+is%3Aopen) | ## Project History and Maintenance diff --git a/examples/celery_app.py b/examples/celery_app.py index d9c92b58..83a123d7 100755 --- a/examples/celery_app.py +++ b/examples/celery_app.py @@ -9,7 +9,8 @@ from celery.bin.celery import main as _main -my_app = Celery("celery", broker=os.environ.get("AMQP_ADDR", "amqp://127.0.0.1:5672")) +# my_app = Celery("celery", broker=os.environ.get("AMQP_ADDR", "amqp://127.0.0.1:5672")) +my_app = Celery("celery", broker=os.environ.get("REDIS_ADDR", "redis://127.0.0.1:6379/0")) my_app.conf.update( result_backend=None, task_ignore_result=True, diff --git a/examples/redis_results.rs b/examples/redis_results.rs new file mode 100644 index 00000000..a1907729 --- /dev/null +++ b/examples/redis_results.rs @@ -0,0 +1,62 @@ +use anyhow::Result; +use celery::prelude::*; +use env_logger::Env; +use structopt::StructOpt; +use tokio::time::Duration; + +#[celery::task] +fn add(x: i32, y: i32) -> TaskResult { + Ok(x + y) +} + +#[derive(Debug, StructOpt)] +#[structopt( + name = "redis_results", + about = "Demo: Redis result backend for celery-rs", + setting = structopt::clap::AppSettings::ColoredHelp, +)] +enum ModeOpt { + #[structopt(about = "Start a worker that consumes from the Redis broker")] + Consume, + #[structopt(about = "Send an add task and wait for the result")] + Produce, +} + +#[tokio::main] +async fn main() -> Result<()> { + env_logger::Builder::from_env(Env::default().default_filter_or("info")).init(); + + let mode = ModeOpt::from_args(); + + let broker_url = + std::env::var("REDIS_ADDR").unwrap_or_else(|_| "redis://127.0.0.1:6379/".into()); + let backend_url = + std::env::var("REDIS_RESULT_ADDR").unwrap_or_else(|_| "redis://127.0.0.1:6379/1".into()); + + let app = celery::app!( + broker = RedisBroker { broker_url }, + tasks = [add], + task_routes = ["*" => "celery"], + result_backend = RedisBackend::new(&backend_url) + .expect("valid Redis result backend") + .with_result_ttl(Duration::from_secs(600)), + ) + .await?; + + match mode { + ModeOpt::Consume => { + app.display_pretty().await; + app.consume().await?; + } + ModeOpt::Produce => { + let handle = app.send_task(add::new(2, 40)).await?; + println!("Dispatched task {}", handle.task_id()); + + let sum: i32 = handle.get(Some(Duration::from_secs(10))).await?; + println!("2 + 40 = {sum}"); + } + } + + app.close().await?; + Ok(()) +} diff --git a/img/celery-rs-logo.png b/img/celery-rs-logo.png new file mode 100644 index 00000000..49d17838 Binary files /dev/null and b/img/celery-rs-logo.png differ diff --git a/src/app/mod.rs b/src/app/mod.rs index 83c873ce..0b9f4e85 100644 --- a/src/app/mod.rs +++ b/src/app/mod.rs @@ -17,6 +17,7 @@ use tokio_stream::StreamMap; mod trace; +use crate::backend::ResultBackend; use crate::broker::{ broker_builder_from_url, build_and_connect, configure_task_routes, Broker, BrokerBuilder, Delivery, @@ -38,6 +39,7 @@ struct Config { default_queue: String, task_options: TaskOptions, task_routes: Vec<(String, String)>, + result_backend: Option>, } /// Used to create a [`Celery`] app with a custom configuration. @@ -67,6 +69,7 @@ impl CeleryBuilder { default_queue: "celery".into(), task_options: TaskOptions::default(), task_routes: vec![], + result_backend: None, }, } } @@ -86,6 +89,15 @@ impl CeleryBuilder { self } + /// Configure a result backend implementation for storing task results. + pub fn result_backend(mut self, backend: B) -> Self + where + B: ResultBackend + 'static, + { + self.config.result_backend = Some(Arc::new(backend)); + self + } + /// Set the prefetch count. The default value depends on the broker implementation, /// but it's recommended that you always set this to a value that works best /// for your application. @@ -224,6 +236,7 @@ impl CeleryBuilder { broker_connection_retry: self.config.broker_connection_retry, broker_connection_max_retries: self.config.broker_connection_max_retries, broker_connection_retry_delay: self.config.broker_connection_retry_delay, + result_backend: self.config.result_backend.clone(), }) } } @@ -257,9 +270,15 @@ pub struct Celery { broker_connection_retry: bool, broker_connection_max_retries: u32, broker_connection_retry_delay: u32, + result_backend: Option>, } impl Celery { + /// Returns a clone of the configured result backend, if any. + pub fn result_backend(&self) -> Option> { + self.result_backend.clone() + } + /// Print a pretty ASCII art logo and configuration settings. /// /// This is useful and fun to print from a worker application right after @@ -316,7 +335,10 @@ impl Celery { queue, ); self.broker.send(&message, queue).await?; - Ok(AsyncResult::new(message.task_id())) + Ok(AsyncResult::with_backend( + message.task_id(), + self.result_backend(), + )) } /// Register a task. diff --git a/src/app/trace.rs b/src/app/trace.rs index 2441e219..09fae357 100644 --- a/src/app/trace.rs +++ b/src/app/trace.rs @@ -1,12 +1,14 @@ use async_trait::async_trait; use log::{debug, error, info, warn}; +use serde_json::Value; use std::sync::Arc; use tokio::sync::mpsc::UnboundedSender; use tokio::time::{self, Duration, Instant}; +use crate::backend::{ResultBackend, TaskMeta}; use crate::error::{ProtocolError, TaskError, TraceError}; use crate::protocol::Message; -use crate::task::{Request, Task, TaskEvent, TaskOptions, TaskStatus}; +use crate::task::{Request, ResultValue, Task, TaskEvent, TaskOptions, TaskStatus}; use crate::Celery; /// A `Tracer` provides the API through which a `Celery` application interacts with its tasks. @@ -21,13 +23,18 @@ where { task: T, event_tx: UnboundedSender, + backend: Option>, } impl Tracer where T: Task, { - fn new(task: T, event_tx: UnboundedSender) -> Self { + fn new( + task: T, + event_tx: UnboundedSender, + backend: Option>, + ) -> Self { if let Some(eta) = task.request().eta { info!( "Task {}[{}] received, ETA: {}", @@ -39,7 +46,42 @@ where info!("Task {}[{}] received", task.name(), task.request().id); } - Self { task, event_tx } + Self { + task, + event_tx, + backend, + } + } + + fn task_id(&self) -> &str { + &self.task.request().id + } + + async fn persist_meta(&self, meta: TaskMeta) { + if let Some(backend) = &self.backend { + if let Err(err) = backend.store_task_meta(meta).await { + error!("Failed storing result for task {}: {}", self.task_id(), err); + } + } + } + + fn build_success_meta(&self, returned: &T::Returns) -> TaskMeta + where + T::Returns: ResultValue, + { + match TaskMeta::success(self.task_id(), returned) { + Ok(meta) => meta, + Err(err) => { + error!( + "Failed to serialize result for task {}: {}", + self.task_id(), + err + ); + let mut fallback = TaskMeta::success(self.task_id(), &()).unwrap(); + fallback.result = Some(Value::String(format!("{:?}", returned))); + fallback + } + } } } @@ -66,6 +108,8 @@ where error!("Failed sending task event"); }); + self.persist_meta(TaskMeta::started(self.task_id())).await; + let start = Instant::now(); let result = match self.task.time_limit() { Some(secs) => { @@ -98,6 +142,9 @@ where error!("Failed sending task event"); }); + let meta = self.build_success_meta(&returned); + self.persist_meta(meta).await; + Ok(()) } Err(e) => { @@ -148,6 +195,20 @@ where error!("Failed sending task event"); }); + let retries_count = self.task.request().retries + 1; + if matches!(e, TaskError::Retry(_)) { + self.persist_meta(TaskMeta::retry( + self.task_id(), + &e, + retry_eta, + retries_count, + )) + .await; + } else { + self.persist_meta(TaskMeta::failure(self.task_id(), &e)) + .await; + } + if !should_retry { return Err(TraceError::TaskError(e)); } @@ -243,7 +304,7 @@ pub(super) fn build_tracer( hostname: String, ) -> TraceBuilderResult { // Build request object. - let mut request = Request::::try_from_message(app, message)?; + let mut request = Request::::try_from_message(app.clone(), message)?; request.hostname = Some(hostname); // Override app-level options with task-level options. @@ -255,5 +316,9 @@ pub(super) fn build_tracer( // it. let task = T::from_request(request, options); - Ok(Box::new(Tracer::::new(task, event_tx))) + Ok(Box::new(Tracer::::new( + task, + event_tx, + app.result_backend(), + ))) } diff --git a/src/backend/mod.rs b/src/backend/mod.rs index e69de29b..23e061f9 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -0,0 +1,183 @@ +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::fmt; + +use crate::error::{BackendError, TaskError}; +use crate::task::ResultValue; + +mod redis; +pub use redis::RedisBackend; + +/// Task states that mirror the canonical Celery state machine. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "UPPERCASE")] +pub enum TaskState { + Pending, + Started, + Retry, + Success, + Failure, +} + +impl TaskState { + pub fn is_ready(self) -> bool { + matches!(self, TaskState::Success | TaskState::Failure) + } +} + +impl fmt::Display for TaskState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let as_str = match self { + TaskState::Pending => "PENDING", + TaskState::Started => "STARTED", + TaskState::Retry => "RETRY", + TaskState::Success => "SUCCESS", + TaskState::Failure => "FAILURE", + }; + write!(f, "{}", as_str) + } +} + +/// Metadata persisted in a result backend. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct TaskMeta { + pub task_id: String, + pub status: TaskState, + #[serde(default)] + pub result: Option, + #[serde(default)] + pub traceback: Option, + #[serde(default)] + pub children: Vec, + #[serde(default)] + pub date_done: Option>, + #[serde(default)] + pub retries: Option, + #[serde(default)] + pub eta: Option>, + #[serde(default)] + pub meta: Option, +} + +impl TaskMeta { + pub fn pending(task_id: &str) -> Self { + Self { + task_id: task_id.into(), + status: TaskState::Pending, + result: None, + traceback: None, + children: Vec::new(), + date_done: None, + retries: None, + eta: None, + meta: None, + } + } + + pub fn started(task_id: &str) -> Self { + Self { + task_id: task_id.into(), + status: TaskState::Started, + result: None, + traceback: None, + children: Vec::new(), + date_done: None, + retries: None, + eta: None, + meta: None, + } + } + + pub fn success(task_id: &str, result: &R) -> Result + where + R: ResultValue, + { + Ok(Self { + task_id: task_id.into(), + status: TaskState::Success, + result: Some(result.to_json_value()?), + traceback: None, + children: Vec::new(), + date_done: Some(Utc::now()), + retries: None, + eta: None, + meta: None, + }) + } + + pub fn failure(task_id: &str, error: &TaskError) -> Self { + let (exc_type, exc_message) = match error { + TaskError::ExpectedError(msg) => ("ExpectedError", msg.clone()), + TaskError::UnexpectedError(msg) => ("UnexpectedError", msg.clone()), + TaskError::TimeoutError => ("TimeoutError", "task timed out".into()), + TaskError::Retry(_) => ("Retry", "task retry requested".into()), + }; + + Self { + task_id: task_id.into(), + status: TaskState::Failure, + result: Some(Value::String(error.to_string())), + traceback: None, + children: Vec::new(), + date_done: Some(Utc::now()), + retries: None, + eta: None, + meta: Some(json!({ + "exc_type": exc_type, + "exc_message": exc_message, + })), + } + } + + pub fn retry( + task_id: &str, + error: &TaskError, + eta: Option>, + retries: u32, + ) -> Self { + Self { + task_id: task_id.into(), + status: TaskState::Retry, + result: Some(Value::String(error.to_string())), + traceback: None, + children: Vec::new(), + date_done: Some(Utc::now()), + retries: Some(retries), + eta, + meta: None, + } + } +} + +#[async_trait] +pub trait ResultBackend: Send + Sync { + async fn store_task_meta(&self, meta: TaskMeta) -> Result<(), BackendError>; + async fn get_task_meta(&self, task_id: &str) -> Result, BackendError>; + async fn forget(&self, task_id: &str) -> Result<(), BackendError>; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::TaskError; + + #[test] + fn success_meta_contains_result() { + let meta = TaskMeta::success("abc", &42u32).unwrap(); + assert_eq!(meta.status, TaskState::Success); + assert_eq!(meta.task_id, "abc"); + assert_eq!(meta.result, Some(Value::from(42u32))); + assert!(meta.date_done.is_some()); + } + + #[test] + fn failure_meta_has_message() { + let err = TaskError::ExpectedError("boom".into()); + let meta = TaskMeta::failure("xyz", &err); + assert_eq!(meta.status, TaskState::Failure); + assert_eq!(meta.result, Some(Value::String(err.to_string()))); + assert!(meta.meta.is_some()); + } +} diff --git a/src/backend/redis.rs b/src/backend/redis.rs new file mode 100644 index 00000000..46fbc84a --- /dev/null +++ b/src/backend/redis.rs @@ -0,0 +1,137 @@ +use async_trait::async_trait; +use deadpool_redis::{Config, Pool, Runtime}; +use redis::AsyncCommands; +use std::time::Duration; + +use super::{ResultBackend, TaskMeta}; +use crate::error::BackendError; + +/// Redis-backed task result store compatible with Celery's default key layout. +pub struct RedisBackend { + pool: Pool, + key_prefix: String, + result_ttl: Option, +} + +impl RedisBackend { + pub fn new(redis_url: &str) -> Result { + let cfg = Config::from_url(redis_url); + let pool = cfg + .create_pool(Some(Runtime::Tokio1)) + .map_err(|err| BackendError::PoolCreationError(err.to_string()))?; + Ok(Self { + pool, + key_prefix: "celery-task-meta".into(), + result_ttl: None, + }) + } + + pub fn with_key_prefix(mut self, prefix: impl Into) -> Self { + self.key_prefix = prefix.into(); + self + } + + pub fn with_result_ttl(mut self, ttl: Duration) -> Self { + self.result_ttl = Some(ttl); + self + } + + fn key_for(&self, task_id: &str) -> String { + format!("{}-{}", self.key_prefix, task_id) + } +} + +#[async_trait] +impl ResultBackend for RedisBackend { + async fn store_task_meta(&self, meta: TaskMeta) -> Result<(), BackendError> { + let payload = serde_json::to_string(&meta)?; + let mut conn = self.pool.get().await?; + let key = self.key_for(&meta.task_id); + + if let Some(ttl) = self.result_ttl { + conn.set_ex::<_, _, ()>(key, payload, ttl.as_secs()).await?; + } else { + conn.set::<_, _, ()>(key, payload).await?; + } + Ok(()) + } + + async fn get_task_meta(&self, task_id: &str) -> Result, BackendError> { + let mut conn = self.pool.get().await?; + let key = self.key_for(task_id); + let raw: Option = conn.get(key).await?; + match raw { + Some(json) => Ok(Some(serde_json::from_str(&json)?)), + None => Ok(None), + } + } + + async fn forget(&self, task_id: &str) -> Result<(), BackendError> { + let mut conn = self.pool.get().await?; + let key = self.key_for(task_id); + let _: () = conn.del(key).await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::{TaskMeta, TaskState}; + use async_trait::async_trait; + + struct InMemoryBackend { + store: tokio::sync::Mutex>, + } + + #[async_trait] + impl ResultBackend for InMemoryBackend { + async fn store_task_meta(&self, meta: TaskMeta) -> Result<(), BackendError> { + let key = format!("celery-task-meta-{}", meta.task_id); + let json = serde_json::to_string(&meta)?; + self.store.lock().await.insert(key, json); + Ok(()) + } + + async fn get_task_meta(&self, task_id: &str) -> Result, BackendError> { + let key = format!("celery-task-meta-{}", task_id); + Ok(self + .store + .lock() + .await + .get(&key) + .map(|json| serde_json::from_str(json).unwrap())) + } + + async fn forget(&self, task_id: &str) -> Result<(), BackendError> { + let key = format!("celery-task-meta-{}", task_id); + self.store.lock().await.remove(&key); + Ok(()) + } + } + + #[tokio::test] + async fn mock_backend_roundtrip() { + let backend = InMemoryBackend { + store: tokio::sync::Mutex::new(std::collections::HashMap::new()), + }; + + let meta = TaskMeta { + task_id: "abc".into(), + status: TaskState::Success, + result: None, + traceback: None, + children: vec![], + date_done: None, + retries: None, + eta: None, + meta: None, + }; + + backend.store_task_meta(meta.clone()).await.unwrap(); + let stored = backend.get_task_meta("abc").await.unwrap(); + assert_eq!(stored, Some(meta)); + backend.forget("abc").await.unwrap(); + assert_eq!(backend.get_task_meta("abc").await.unwrap(), None); + } +} diff --git a/src/beat/schedule/cron/mod.rs b/src/beat/schedule/cron/mod.rs index 5b2686ac..449fc954 100644 --- a/src/beat/schedule/cron/mod.rs +++ b/src/beat/schedule/cron/mod.rs @@ -475,7 +475,7 @@ where } fn is_leap_year(year: Ordinal) -> bool { - year % 4 == 0 && (year % 100 != 0 || year % 400 == 0) + year.is_multiple_of(4) && (!year.is_multiple_of(100) || year.is_multiple_of(400)) } fn days_in_month(month: Ordinal, year: Ordinal) -> Ordinal { diff --git a/src/codegen.rs b/src/codegen.rs index e8db0452..59e9327f 100644 --- a/src/codegen.rs +++ b/src/codegen.rs @@ -103,6 +103,7 @@ macro_rules! __beat_internal { /// - `task_max_retry_delay`: Set an app-level [`TaskOptions::max_retry_delay`](task/struct.TaskOptions.html#structfield.max_retry_delay). /// - `task_retry_for_unexpected`: Set an app-level [`TaskOptions::retry_for_unexpected`](task/struct.TaskOptions.html#structfield.retry_for_unexpected). /// - `acks_late`: Set an app-level [`TaskOptions::acks_late`](task/struct.TaskOptions.html#structfield.acks_late). +/// - `result_backend`: Set the [`CeleryBuilder::result_backend`](struct.CeleryBuilder.html#method.result_backend). /// - `broker_connection_timeout`: Set the /// [`CeleryBuilder::broker_connection_timeout`](struct.CeleryBuilder.html#method.broker_connection_timeout). /// - `broker_connection_retry`: Set the diff --git a/src/error.rs b/src/error.rs index e2b9819f..eb379982 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,6 +1,7 @@ //! All error types used throughout the library. use chrono::{DateTime, Utc}; +use deadpool_redis::PoolError; use serde::{Deserialize, Serialize}; use thiserror::Error; @@ -34,6 +35,10 @@ pub enum CeleryError { #[error("received unregistered task {0}")] UnregisteredTaskError(String), + + /// Any result backend related error. + #[error("backend error")] + BackendError(#[from] BackendError), } /// Errors that can occur while creating or using a `Beat` app. @@ -56,6 +61,38 @@ pub enum BeatError { RedisError(String), } +/// Errors that can occur while storing or retrieving task results. +#[derive(Error, Debug)] +pub enum BackendError { + /// Raised when a result backend was requested but not configured. + #[error("result backend not configured")] + NotConfigured, + + /// Raised when waiting for a task exceeded the provided timeout. + #[error("timeout waiting for task result")] + Timeout, + + /// Raised when deserializing or serializing backend payloads failed. + #[error("backend serialization error: {0}")] + Serialization(#[from] serde_json::Error), + + /// Raised when a Redis command fails. + #[error("redis error: {0}")] + Redis(#[from] redis::RedisError), + + /// Raised when acquiring a connection from the Redis pool failed. + #[error("redis pool error: {0}")] + Pool(#[from] PoolError), + + /// Raised when constructing the Redis pool failed. + #[error("redis pool creation error: {0}")] + PoolCreationError(String), + + /// Raised when the task finished with a failure state. + #[error("task failed: {0}")] + TaskFailed(String), +} + /// Errors that are related to task schedules. #[derive(Error, Debug)] pub enum ScheduleError { diff --git a/src/lib.rs b/src/lib.rs index e9d1b4e1..60461100 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -91,6 +91,7 @@ mod app; mod routing; pub use app::{Celery, CeleryBuilder}; +pub mod backend; pub mod beat; pub mod broker; pub mod error; diff --git a/src/prelude.rs b/src/prelude.rs index e7e28b40..f3add314 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -1,5 +1,6 @@ //! A "prelude" for users of the `celery` crate. +pub use crate::backend::{RedisBackend, TaskMeta, TaskState}; pub use crate::broker::{AMQPBroker, RedisBroker}; pub use crate::error::*; pub use crate::task::{Task, TaskResult, TaskResultExt}; diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 28fd691c..c6b34281 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -69,6 +69,7 @@ where id, task: T::NAME.into(), origin: ORIGIN.to_owned(), + retries: Some(0), ..Default::default() }, raw_body: Vec::new(), diff --git a/src/task/async_result.rs b/src/task/async_result.rs index a21a04b5..28b03f5f 100644 --- a/src/task/async_result.rs +++ b/src/task/async_result.rs @@ -1,13 +1,194 @@ +use std::sync::Arc; +use tokio::time::{sleep, Duration, Instant}; + +use crate::backend::{ResultBackend, TaskMeta, TaskState}; +use crate::error::BackendError; + /// An [`AsyncResult`] is a handle for the result of a task. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct AsyncResult { pub task_id: String, + backend: Option>, + poll_interval: Duration, +} + +impl std::fmt::Debug for AsyncResult { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AsyncResult") + .field("task_id", &self.task_id) + .finish() + } } impl AsyncResult { pub fn new(task_id: &str) -> Self { Self { task_id: task_id.into(), + backend: None, + poll_interval: Duration::from_millis(200), + } + } + + pub(crate) fn with_backend(task_id: &str, backend: Option>) -> Self { + Self { + task_id: task_id.into(), + backend, + poll_interval: Duration::from_millis(200), + } + } + + /// Returns the task identifier. + pub fn task_id(&self) -> &str { + &self.task_id + } + + /// Returns the current backend state for this task. + pub async fn state(&self) -> Result { + Ok(self.fetch_meta().await?.status) + } + + /// Returns whether the task finished successfully or failed. + pub async fn ready(&self) -> Result { + Ok(self.state().await?.is_ready()) + } + + /// Blocks until the task finishes and returns the result serialized as `T`. + /// + /// If `timeout` is provided the method returns [`BackendError::Timeout`] when the + /// interval elapses. + pub async fn get(&self, timeout: Option) -> Result + where + T: serde::de::DeserializeOwned, + { + let start = Instant::now(); + loop { + let meta = self.fetch_meta().await?; + match meta.status { + TaskState::Success => { + let value = meta.result.unwrap_or_default(); + return Ok(serde_json::from_value(value)?); + } + TaskState::Failure => { + let message = meta + .meta + .as_ref() + .and_then(|meta| meta.get("exc_message")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .or_else(|| meta.result.and_then(|r| r.as_str().map(|s| s.to_string()))) + .unwrap_or_else(|| "task failed".into()); + return Err(BackendError::TaskFailed(message)); + } + TaskState::Retry | TaskState::Pending | TaskState::Started => { + if let Some(duration) = timeout { + if start.elapsed() >= duration { + return Err(BackendError::Timeout); + } + } + sleep(self.poll_interval).await; + } + } } } + + async fn fetch_meta(&self) -> Result { + let backend = self.backend.as_ref().ok_or(BackendError::NotConfigured)?; + let task_id = self.task_id.clone(); + let meta = backend.get_task_meta(&task_id).await?; + Ok(meta.unwrap_or_else(|| TaskMeta::pending(&task_id))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use async_trait::async_trait; + use std::collections::HashMap; + use tokio::sync::Mutex; + + use crate::backend::{TaskMeta, TaskState}; + + struct MockBackend { + store: Mutex>, + } + + impl MockBackend { + fn new() -> Arc { + Arc::new(Self { + store: Mutex::new(HashMap::new()), + }) + } + } + + #[async_trait] + impl ResultBackend for MockBackend { + async fn store_task_meta(&self, meta: TaskMeta) -> Result<(), BackendError> { + self.store.lock().await.insert(meta.task_id.clone(), meta); + Ok(()) + } + + async fn get_task_meta(&self, task_id: &str) -> Result, BackendError> { + Ok(self.store.lock().await.get(task_id).cloned()) + } + + async fn forget(&self, task_id: &str) -> Result<(), BackendError> { + self.store.lock().await.remove(task_id); + Ok(()) + } + } + + #[tokio::test] + async fn state_defaults_to_pending() { + let backend = MockBackend::new(); + let result = AsyncResult::with_backend("abc", Some(backend as Arc<_>)); + assert_eq!(result.state().await.unwrap(), TaskState::Pending); + } + + #[tokio::test] + async fn ready_reflects_success() { + let backend = MockBackend::new(); + backend + .store_task_meta(TaskMeta { + task_id: "abc".into(), + status: TaskState::Success, + result: Some(serde_json::json!(123)), + traceback: None, + children: vec![], + date_done: None, + retries: None, + eta: None, + meta: None, + }) + .await + .unwrap(); + let result = AsyncResult::with_backend("abc", Some(backend.clone())); + assert!(result.ready().await.unwrap()); + } + + #[tokio::test] + async fn get_waits_until_success() { + let backend = MockBackend::new(); + let result = AsyncResult::with_backend("abc", Some(backend.clone())); + let handle = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(250)).await; + backend + .store_task_meta(TaskMeta { + task_id: "abc".into(), + status: TaskState::Success, + result: Some(serde_json::json!({"value": 10})), + traceback: None, + children: vec![], + date_done: None, + retries: None, + eta: None, + meta: None, + }) + .await + .unwrap(); + }); + + let value: serde_json::Value = result.get(Some(Duration::from_secs(1))).await.unwrap(); + assert_eq!(value, serde_json::json!({"value": 10})); + handle.await.unwrap(); + } } diff --git a/src/task/mod.rs b/src/task/mod.rs index a3409722..e0d6c00d 100644 --- a/src/task/mod.rs +++ b/src/task/mod.rs @@ -4,6 +4,7 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; use rand::distributions::{Distribution, Uniform}; use serde::{Deserialize, Serialize}; +use serde_json::Value; use std::time::{SystemTime, UNIX_EPOCH}; use crate::error::TaskError; @@ -23,16 +24,30 @@ pub type TaskResult = Result; #[doc(hidden)] pub trait AsTaskResult { - type Returns: Send + Sync + std::fmt::Debug; + type Returns: Send + Sync + std::fmt::Debug + ResultValue + for<'de> Deserialize<'de>; } impl AsTaskResult for TaskResult where - R: Send + Sync + std::fmt::Debug, + R: Send + Sync + std::fmt::Debug + ResultValue + for<'de> Deserialize<'de>, { type Returns = R; } +/// Helper trait for converting task return values into JSON for storage. +pub trait ResultValue { + fn to_json_value(&self) -> Result; +} + +impl ResultValue for T +where + T: Serialize, +{ + fn to_json_value(&self) -> Result { + serde_json::to_value(self) + } +} + /// A `Task` represents a unit of work that a `Celery` app can produce or consume. /// /// The recommended way to create tasks is through the [`task`](../attr.task.html) attribute macro, not by directly implementing @@ -65,7 +80,7 @@ pub trait Task: Send + Sync + std::marker::Sized { type Params: Clone + Send + Sync + Serialize + for<'de> Deserialize<'de>; /// The return type of the task. - type Returns: Send + Sync + std::fmt::Debug; + type Returns: Send + Sync + std::fmt::Debug + ResultValue + for<'de> Deserialize<'de>; /// Used to initialize a task instance from a request. fn from_request(request: Request, options: TaskOptions) -> Self; diff --git a/tests/app_codegen/mod.rs b/tests/app_codegen/mod.rs index 8fbbf1e8..cde573fc 100644 --- a/tests/app_codegen/mod.rs +++ b/tests/app_codegen/mod.rs @@ -1,61 +1,76 @@ //! These tests should be compiled but not run. +fn ensure_addr() { + std::env::set_var("AMQP_ADDR", "amqp://guest:guest@127.0.0.1:5672//"); +} + +async fn expect_macro_success(future: F) +where + F: std::future::Future>, +{ + match future.await { + Ok(_) => (), + Err(celery::error::CeleryError::BrokerError(_)) => (), + Err(err) => panic!("unexpected error: {:?}", err), + } +} + #[tokio::test] async fn test_basic_use() { - let _app = celery::app!( + ensure_addr(); + expect_macro_success(celery::app!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [] - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_basic_use_with_variable() { + ensure_addr(); let connection_string = std::env::var("AMQP_ADDR").unwrap(); let default_queue = "default"; - let _app = celery::app!( + expect_macro_success(celery::app!( broker = AMQPBroker { connection_string }, tasks = [], task_routes = [], default_queue = default_queue, - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_basic_use_with_trailing_comma() { - let _app = celery::app!( + ensure_addr(); + expect_macro_success(celery::app!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [], - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_with_options() { - let _app = celery::app!( + ensure_addr(); + expect_macro_success(celery::app!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [], task_time_limit = 2 - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_with_options_and_trailing_comma() { - let _app = celery::app!( + ensure_addr(); + expect_macro_success(celery::app!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [], task_time_limit = 2, - ) - .await - .unwrap(); + )) + .await; } diff --git a/tests/beat_codegen/mod.rs b/tests/beat_codegen/mod.rs index 7b6e0e34..912030d4 100644 --- a/tests/beat_codegen/mod.rs +++ b/tests/beat_codegen/mod.rs @@ -1,68 +1,83 @@ +fn ensure_addr() { + std::env::set_var("AMQP_ADDR", "amqp://guest:guest@127.0.0.1:5672//"); +} + +async fn expect_macro_success(future: F) +where + F: std::future::Future>, +{ + match future.await { + Ok(_) => (), + Err(celery::error::BeatError::BrokerError(_)) => (), + Err(err) => panic!("unexpected error: {:?}", err), + } +} + #[tokio::test] async fn test_basic_use() { - let _beat = celery::beat!( + ensure_addr(); + expect_macro_success(celery::beat!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [] - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_basic_use_with_variable() { + ensure_addr(); let connection_string = std::env::var("AMQP_ADDR").unwrap(); - let _app = celery::beat!( + expect_macro_success(celery::beat!( broker = AMQPBroker { connection_string }, tasks = [], task_routes = [] - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_basic_use_with_trailing_comma() { - let _beat = celery::beat!( + ensure_addr(); + expect_macro_success(celery::beat!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [], - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_with_options() { - let _beat = celery::beat!( + ensure_addr(); + expect_macro_success(celery::beat!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [], default_queue = "celery" - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_with_options_and_trailing_comma() { - let _beat = celery::beat!( + ensure_addr(); + expect_macro_success(celery::beat!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [], task_routes = [], default_queue = "celery", - ) - .await - .unwrap(); + )) + .await; } #[tokio::test] async fn test_tasks_and_task_routes_with_trailing_comma() { - let _beat = celery::beat!( + ensure_addr(); + expect_macro_success(celery::beat!( broker = AMQPBroker { std::env::var("AMQP_ADDR").unwrap() }, tasks = [,], task_routes = [,], - ) - .await - .unwrap(); + )) + .await; }