apalis_postgres/
pubsub.rs

1use apalis_core::backend::future::BoxSyncFuture;
2use futures::{Stream, StreamExt, TryStreamExt, stream::BoxStream};
3use serde::Deserialize;
4use sqlx::postgres::{PgListener, PgNotification};
5use std::{
6    pin::Pin,
7    sync::Mutex,
8    task::{Context, Poll},
9};
10
11use crate::{PgTaskId, error::Error};
12
13/// A standalone listener for `apalis::job::insert`
14pub struct Pubsub {
15    state: State,
16    listener: Option<Mutex<BoxStream<'static, Result<PgNotification, Error>>>>,
17    pool: sqlx::PgPool,
18    namespace: String,
19}
20
21impl Clone for Pubsub {
22    fn clone(&self) -> Self {
23        Self {
24            state: State::Starting,
25            listener: None,
26            pool: self.pool.clone(),
27            namespace: self.namespace.clone(),
28        }
29    }
30}
31
32impl Pubsub {
33    pub fn new(pool: sqlx::PgPool, namespace: String) -> Self {
34        Self {
35            state: State::Starting,
36            pool,
37            namespace,
38            listener: None,
39        }
40    }
41}
42
43enum State {
44    Starting,
45
46    Connecting {
47        fut: BoxSyncFuture<Result<PgListener, Error>>,
48    },
49
50    Listening,
51    Closed,
52}
53
54/// A new event emitted when a new job is added
55#[derive(Debug, Deserialize)]
56pub struct InsertEvent {
57    pub job_type: String,
58    pub id: PgTaskId,
59}
60
61impl Stream for Pubsub {
62    type Item = Result<PgTaskId, Error>;
63
64    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
65        let this = self.get_mut();
66
67        loop {
68            match &mut this.state {
69                State::Starting => {
70                    this.state = State::Connecting {
71                        fut: {
72                            let pool = this.pool.clone();
73                            let fut = Box::pin(async move {
74                                let mut listener = PgListener::connect_with(&pool).await?;
75
76                                listener.listen("apalis::job::insert").await?;
77
78                                Ok(listener)
79                            });
80                            BoxSyncFuture::new(fut)
81                        },
82                    }
83                }
84                State::Connecting { fut } => {
85                    let listener = match fut.poll_unpin(cx) {
86                        Poll::Pending => return Poll::Pending,
87
88                        Poll::Ready(Err(err)) => {
89                            this.state = State::Closed;
90                            return Poll::Ready(Some(Err(err)));
91                        }
92
93                        Poll::Ready(Ok(listener)) => listener,
94                    };
95
96                    this.listener = Some(Mutex::new(
97                        listener.into_stream().map_err(|e| e.into()).boxed(),
98                    ));
99
100                    this.state = State::Listening;
101                }
102
103                State::Listening => {
104                    let listener = this.listener.as_mut().expect("listener initialized");
105
106                    match listener.get_mut().unwrap().as_mut().poll_next(cx) {
107                        Poll::Pending => return Poll::Pending,
108
109                        Poll::Ready(None) => {
110                            this.state = State::Closed;
111                            return Poll::Ready(None);
112                        }
113
114                        Poll::Ready(Some(Err(_))) => {
115                            continue;
116                        }
117
118                        Poll::Ready(Some(Ok(notification))) => {
119                            let Ok(ev) =
120                                serde_json::from_str::<InsertEvent>(notification.payload())
121                            else {
122                                continue;
123                            };
124
125                            if ev.job_type != this.namespace {
126                                continue;
127                            }
128
129                            return Poll::Ready(Some(Ok(ev.id)));
130                        }
131                    }
132                }
133                State::Closed => {
134                    return Poll::Ready(None);
135                }
136            }
137        }
138    }
139}
apalis_postgres/pubsub.rs

apalis_postgres/
pubsub.rs
