1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use std::marker::PhantomData;

use futures::{future::BoxFuture, FutureExt};
use tower::{Layer, Service};

use crate::{context::HasJobContext, error::BoxDynError, request::JobRequest, worker::WorkerId};

/// An error occurred while trying to acknowledge a message.
#[derive(Debug, thiserror::Error)]
pub enum AckError {
    /// Acknowledgement failed
    #[error("Acknowledgement failed {0}")]
    NoAck(#[source] BoxDynError),
}

/// A trait for acknowledging successful job processing
#[async_trait::async_trait]
pub trait Ack<J> {
    /// The data to fetch from context to allow acknowledgement
    type Acknowledger;
    /// Acknowledges successful processing of the given request
    async fn ack(&self, worker_id: &WorkerId, data: &Self::Acknowledger) -> Result<(), AckError>;
}

/// A layer that acknowledges a job completed successfully
#[derive(Debug)]
pub struct AckLayer<A: Ack<J>, J> {
    ack: A,
    job_type: PhantomData<J>,
    worker_id: WorkerId,
}

impl<A: Ack<J>, J> AckLayer<A, J> {
    /// Build a new [AckLayer] for a job
    pub fn new(ack: A, worker_id: WorkerId) -> Self {
        Self {
            ack,
            job_type: PhantomData,
            worker_id,
        }
    }
}

impl<A, J, S> Layer<S> for AckLayer<A, J>
where
    S: Service<JobRequest<J>> + Send + 'static,
    S::Error: std::error::Error + Send + Sync + 'static,
    S::Future: Send + 'static,
    A: Ack<J> + Clone + Send + Sync + 'static,
{
    type Service = AckService<S, A, J>;

    fn layer(&self, service: S) -> Self::Service {
        AckService {
            service,
            ack: self.ack.clone(),
            job_type: PhantomData,
            worker_id: self.worker_id.clone(),
        }
    }
}

/// The underlying service for an [AckLayer]
#[derive(Debug)]
pub struct AckService<SV, A, J> {
    service: SV,
    ack: A,
    job_type: PhantomData<J>,
    worker_id: WorkerId,
}

impl<SV, A, J> Service<JobRequest<J>> for AckService<SV, A, J>
where
    SV: Service<JobRequest<J>> + Send + Sync + 'static,
    SV::Error: std::error::Error + Send + Sync + 'static,
    <SV as Service<JobRequest<J>>>::Future: std::marker::Send + 'static,
    A: Ack<J> + Send + 'static + Clone + Send + Sync,
    J: 'static,
    <SV as Service<JobRequest<J>>>::Response: std::marker::Send,
    <A as Ack<J>>::Acknowledger: Sync + Send + Clone,
{
    type Response = SV::Response;
    type Error = SV::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, request: JobRequest<J>) -> Self::Future {
        let ack = self.ack.clone();
        let worker_id = self.worker_id.clone();
        let data = request
            .context()
            .data_opt::<<A as Ack<J>>::Acknowledger>()
            .cloned();

        let fut = self.service.call(request);
        let fut_with_ack = async move {
            let res = fut.await;
            if let Some(data) = data {
                if let Err(e) = ack.ack(&worker_id.clone(), &data).await {
                    tracing::warn!("Acknowledgement Failed: {}", e);
                }
            } else {
                tracing::warn!(
                    "Acknowledgement could not be called due to missing ack data in context : {}",
                    &std::any::type_name::<<A as Ack<J>>::Acknowledger>()
                );
            }
            res
        };
        fut_with_ack.boxed()
    }
}