apalis_core/task/
attempt.rs

1use std::sync::{
2    atomic::{AtomicUsize, Ordering},
3    Arc,
4};
5
6use serde::{Deserialize, Deserializer, Serialize, Serializer};
7
8use crate::{request::Request, service_fn::FromRequest};
9
10/// A wrapper to keep count of the attempts tried by a task
11#[derive(Debug, Clone)]
12pub struct Attempt(Arc<AtomicUsize>);
13
14// Custom serialization function
15fn serialize<S>(attempt: &Attempt, serializer: S) -> Result<S::Ok, S::Error>
16where
17    S: Serializer,
18{
19    let value = attempt.0.load(Ordering::SeqCst);
20    serializer.serialize_u64(value as u64)
21}
22
23// Custom deserialization function
24fn deserialize<'de, D>(deserializer: D) -> Result<Attempt, D::Error>
25where
26    D: Deserializer<'de>,
27{
28    let value = u64::deserialize(deserializer)?;
29    Ok(Attempt(Arc::new(AtomicUsize::new(value as usize))))
30}
31
32impl Serialize for Attempt {
33    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
34    where
35        S: Serializer,
36    {
37        serialize(self, serializer)
38    }
39}
40
41impl<'de> Deserialize<'de> for Attempt {
42    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
43    where
44        D: Deserializer<'de>,
45    {
46        deserialize(deserializer)
47    }
48}
49
50impl Default for Attempt {
51    fn default() -> Self {
52        Self(Arc::new(AtomicUsize::new(0)))
53    }
54}
55
56impl Attempt {
57    /// Build a new tracker
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    /// Build a tracker from an existing value
63    pub fn new_with_value(value: usize) -> Self {
64        Self(Arc::new(AtomicUsize::from(value)))
65    }
66
67    /// Get the current value
68    pub fn current(&self) -> usize {
69        self.0.load(std::sync::atomic::Ordering::Relaxed)
70    }
71
72    /// Increase the current value
73    pub fn increment(&self) -> usize {
74        self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
75    }
76}
77
78impl<Req, Ctx> FromRequest<Request<Req, Ctx>> for Attempt {
79    fn from_request(req: &Request<Req, Ctx>) -> Result<Self, crate::error::Error> {
80        Ok(req.parts.attempt.clone())
81    }
82}