Skip to main content

apalis_core/task/
attempt.rs

1//! A thread-safe tracker for counting the number of attempts made by a task.
2//!
3//! The `Attempt` struct wraps an atomic counter, allowing concurrent increment and retrieval of the attempt count. It is designed to be used within the Apalis job/task system, enabling tasks to keep track of how many times they have been retried or executed.
4//!
5//! Features:
6//! - Thread-safe increment and retrieval of attempt count.
7//! - Integration with apalis `FromRequest` trait for extracting attempt information from a task context.
8//! - Optional (via the `serde` feature) serialization and deserialization support for persisting or transmitting attempt state.
9use std::{
10    convert::Infallible,
11    sync::{Arc, atomic::AtomicUsize},
12};
13
14use crate::{task::Task, task::from_request::FromRequest};
15
16/// A wrapper to keep count of the attempts tried by a task
17#[derive(Debug, Clone)]
18pub struct Attempt(Arc<AtomicUsize>);
19
20impl Default for Attempt {
21    fn default() -> Self {
22        Self(Arc::new(AtomicUsize::new(0)))
23    }
24}
25
26impl Attempt {
27    /// Build a new tracker
28    #[must_use]
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    /// Build a tracker from an existing value
34    #[must_use]
35    pub fn new_with_value(value: usize) -> Self {
36        Self(Arc::new(AtomicUsize::from(value)))
37    }
38
39    /// Get the current value
40    #[must_use]
41    pub fn current(&self) -> usize {
42        self.0.load(std::sync::atomic::Ordering::SeqCst)
43    }
44
45    /// Increase the current value and returns the new value
46    #[must_use]
47    pub(crate) fn increment(&self) -> usize {
48        self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
49    }
50}
51
52impl<Args> FromRequest<Task<Args>> for Attempt
53where
54    Args: Sync,
55{
56    type Error = Infallible;
57    async fn from_request(task: &Task<Args>) -> Result<Self, Self::Error> {
58        Ok(task.raw_attempt().clone())
59    }
60}
61
62#[cfg(feature = "serde")]
63mod serde_impl {
64    use std::sync::atomic::Ordering;
65
66    use serde::{Deserialize, Deserializer, Serialize, Serializer};
67
68    use super::*;
69
70    // Custom serialization function
71    fn serialize<S>(attempt: &Attempt, serializer: S) -> Result<S::Ok, S::Error>
72    where
73        S: Serializer,
74    {
75        let value = attempt.0.load(Ordering::SeqCst);
76        serializer.serialize_u64(value as u64)
77    }
78
79    // Custom deserialization function
80    fn deserialize<'de, D>(deserializer: D) -> Result<Attempt, D::Error>
81    where
82        D: Deserializer<'de>,
83    {
84        let value = u64::deserialize(deserializer)?;
85        Ok(Attempt(Arc::new(AtomicUsize::new(value as usize))))
86    }
87
88    impl Serialize for Attempt {
89        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
90        where
91            S: Serializer,
92        {
93            serialize(self, serializer)
94        }
95    }
96
97    impl<'de> Deserialize<'de> for Attempt {
98        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
99        where
100            D: Deserializer<'de>,
101        {
102            deserialize(deserializer)
103        }
104    }
105}