ruststream_pulsar/
broker.rs1use std::collections::HashMap;
9use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, Ordering};
11
12use pulsar::{Authentication, Pulsar, TokioExecutor};
13use ruststream::{Broker, ConnectedBroker, DefaultPublish, DescribeServer, ServerSpec, Subscribe};
14use tokio::sync::{Mutex, OnceCell};
15
16use crate::error::{PulsarError, box_err};
17use crate::publisher::{PulsarProducer, PulsarPublish, PulsarPublisher};
18use crate::subscriber::PulsarSubscriber;
19use crate::subscription::PulsarSubscription;
20
21pub(crate) struct Core {
28 pub(crate) client: Pulsar<TokioExecutor>,
29 pub(crate) closed: AtomicBool,
30 pub(crate) producers: Mutex<HashMap<String, Arc<Mutex<PulsarProducer>>>>,
32}
33
34impl Core {
35 pub(crate) fn ensure_open(&self) -> Result<(), PulsarError> {
36 if self.closed.load(Ordering::Acquire) {
37 return Err(PulsarError::NotConnected);
38 }
39 Ok(())
40 }
41}
42
43impl std::fmt::Debug for Core {
44 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45 f.debug_struct("Core")
46 .field("closed", &self.closed.load(Ordering::Relaxed))
47 .finish_non_exhaustive()
48 }
49}
50
51pub(crate) type CoreCell = Arc<OnceCell<Arc<Core>>>;
52
53#[derive(Debug, Clone)]
69#[must_use]
70pub struct PulsarBroker {
71 url: String,
72 token: Option<String>,
73 cell: CoreCell,
75}
76
77impl PulsarBroker {
78 pub fn new(url: impl Into<String>) -> Self {
80 Self {
81 url: url.into(),
82 token: None,
83 cell: Arc::new(OnceCell::new()),
84 }
85 }
86
87 pub fn token(mut self, token: impl Into<String>) -> Self {
89 self.token = Some(token.into());
90 self
91 }
92
93 #[must_use]
95 pub fn publisher(&self) -> PulsarPublisher {
96 PulsarPublisher::new(Arc::clone(&self.cell))
97 }
98}
99
100impl Broker for PulsarBroker {
101 type Error = PulsarError;
102 type Connected = ConnectedPulsarBroker;
103
104 async fn connect(self) -> Result<Self::Connected, Self::Error> {
105 let core = self
106 .cell
107 .get_or_try_init(async || {
108 let mut builder = Pulsar::builder(self.url.clone(), TokioExecutor);
109 if let Some(token) = &self.token {
110 builder = builder.with_auth(Authentication {
111 name: "token".to_owned(),
112 data: token.clone().into_bytes(),
113 });
114 }
115 let client = builder
116 .build()
117 .await
118 .map_err(|e| PulsarError::Connect(box_err(e)))?;
119 Ok::<_, PulsarError>(Arc::new(Core {
120 client,
121 closed: AtomicBool::new(false),
122 producers: Mutex::new(HashMap::new()),
123 }))
124 })
125 .await?
126 .clone();
127 Ok(ConnectedPulsarBroker {
128 core,
129 cell: self.cell,
130 })
131 }
132}
133
134impl DescribeServer for PulsarBroker {
135 fn describe_server(&self) -> ServerSpec {
136 ServerSpec::new(
137 self.url
138 .trim_start_matches("pulsar+ssl://")
139 .trim_start_matches("pulsar://"),
140 "pulsar",
141 )
142 }
143}
144
145#[derive(Debug)]
147pub struct ConnectedPulsarBroker {
148 pub(crate) core: Arc<Core>,
149 cell: CoreCell,
151}
152
153impl ConnectedPulsarBroker {
154 #[must_use]
157 pub fn publisher(&self) -> PulsarPublisher {
158 PulsarPublisher::new(Arc::clone(&self.cell))
159 }
160
161 pub async fn subscribe_descriptor(
168 &self,
169 descriptor: PulsarSubscription,
170 ) -> Result<PulsarSubscriber, PulsarError> {
171 descriptor.validate()?;
172 self.core.ensure_open()?;
173 PulsarSubscriber::open(&self.core, descriptor).await
174 }
175}
176
177impl ConnectedBroker for ConnectedPulsarBroker {
178 type Error = PulsarError;
179 type Closed = ();
180
181 async fn shutdown(self) -> Result<(), Self::Error> {
182 self.core.closed.store(true, Ordering::Release);
183 let producers: Vec<_> = {
186 let mut map = self.core.producers.lock().await;
187 map.drain().map(|(_, producer)| producer).collect()
188 };
189 for producer in producers {
190 let mut producer = producer.lock().await;
191 if let Err(err) = Box::pin(producer.close()).await {
192 tracing::debug!(error = %err, "pulsar producer close failed");
193 }
194 }
195 Ok(())
196 }
197}
198
199impl Subscribe for ConnectedPulsarBroker {
200 type Subscriber = PulsarSubscriber;
201
202 async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
203 self.subscribe_descriptor(PulsarSubscription::new(name, "ruststream"))
206 .await
207 }
208}
209
210impl DefaultPublish for ConnectedPulsarBroker {
211 type Policy = PulsarPublish;
212}