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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! Tourniquet integration with the [celery](https://docs.rs/celery) library.

use std::fmt::Debug;

use crate::{Connector, Next, RoundRobin};
use async_trait::async_trait;
use celery::{
    broker::{AMQPBroker, Broker},
    error::BrokerError::BadRoutingPattern,
    error::CeleryError::{self, *},
    task::{Signature, Task},
    Celery,
};
#[cfg(feature = "trace")]
use tracing::{
    field::{display, Empty},
    instrument, Span,
};

impl Next for CeleryError {
    fn is_next(&self) -> bool {
        match self {
            BrokerError(BadRoutingPattern(_)) => false,
            BrokerError(_) | IoError(_) | ProtocolError(_) => true,
            NoQueueToConsume
            | ForcedShutdown
            | TaskRegistrationError(_)
            | UnregisteredTaskError(_) => false,
        }
    }
}

/// Ready to use connector for Celery.
///
/// Please refer to
/// [celery's documentation](https://docs.rs/celery/^0.3/celery/struct.CeleryBuilder.html)
/// for more information about the fields of this structure.
///
/// Note that this is a basic connector for celery, exposing only the most common options used by
/// celery producers, not consumers.  Please create your own connector should you need finer
/// grained control over the created celery client.
pub struct CeleryConnector<'a> {
    pub name: &'a str,
    pub default_queue: Option<&'a str>,
    pub routes: &'a [(&'a str, &'a str)],
    pub connection_timeout: Option<u32>,
}

impl<'a> Default for CeleryConnector<'a> {
    fn default() -> Self {
        Self { name: "celery", default_queue: None, routes: &[], connection_timeout: None }
    }
}

#[async_trait]
impl<'a> Connector<String, Celery<AMQPBroker>, CeleryError> for CeleryConnector<'a> {
    #[cfg_attr(feature = "trace", tracing::instrument(skip(self), err))]
    async fn connect(&self, url: &String) -> Result<Celery<AMQPBroker>, CeleryError> {
        let mut builder = Celery::<AMQPBroker>::builder(self.name, url.as_ref());

        if let Some(queue) = self.default_queue {
            builder = builder.default_queue(queue);
        }
        for (pattern, queue) in self.routes {
            builder = builder.task_route(*pattern, *queue);
        }
        if let Some(timeout) = self.connection_timeout {
            builder = builder.broker_connection_timeout(timeout);
        }

        builder.build().await
    }
}

impl<SvcSrc, B, Conn> RoundRobin<SvcSrc, Celery<B>, CeleryError, Conn>
where
    SvcSrc: Debug,
    B: Broker + 'static,
    Conn: Connector<SvcSrc, Celery<B>, CeleryError>,
{
    /// Send a Celery task.
    ///
    /// The `task_gen` argument returns a signature for each attempt, should each attempt hold a
    /// different value (e.g. trace id, attempt id, timestamp, ...).
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use celery::task::TaskResult;
    /// # use tourniquet::{RoundRobin, conn::CeleryConnector};
    /// #
    /// #[celery::task]
    /// async fn do_work(work: String) -> TaskResult<()> {
    ///     // Some work
    /// # println!("{}", work);
    ///     Ok(())
    /// }
    ///
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let rr = RoundRobin::new(
    ///     vec!["amqp://rabbit01:5672/".to_owned(), "amqp://rabbit02:5672".to_owned()],
    ///     CeleryConnector { name: "rr", routes: &[("*", "my_route")], ..Default::default() },
    /// );
    ///
    /// # let work = "foo".to_owned();
    /// rr.send_task(|| do_work::new(work.clone())).await.expect("Failed to send task");
    /// # Ok(())
    /// # }
    /// ```
    #[cfg_attr(
        feature = "trace",
        instrument(
            fields(task_name = display(Signature::<T>::task_name()), task_id = Empty),
            skip(self, task_gen),
            err,
        ),
    )]
    pub async fn send_task<T, F>(&self, task_gen: F) -> Result<String, CeleryError>
    where
        T: Task + 'static,
        F: Fn() -> Signature<T>,
    {
        #[cfg(feature = "trace")]
        tracing::info!("Sending task {}", Signature::<T>::task_name());

        let task_gen = &task_gen;
        let task = self.run(|celery| async move { celery.send_task(task_gen()).await }).await?;

        #[cfg(feature = "trace")]
        Span::current().record("task_id", &display(&task.task_id));

        Ok(task.task_id)
    }
}