use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use pulsar::{Authentication, Pulsar, TokioExecutor};
use ruststream::{Broker, ConnectedBroker, DefaultPublish, DescribeServer, ServerSpec, Subscribe};
use tokio::sync::{Mutex, OnceCell};
use crate::error::{PulsarError, box_err};
use crate::publisher::{PulsarProducer, PulsarPublish, PulsarPublisher};
use crate::subscriber::PulsarSubscriber;
use crate::subscription::PulsarSubscription;
pub(crate) struct Core {
pub(crate) client: Pulsar<TokioExecutor>,
pub(crate) closed: AtomicBool,
pub(crate) producers: Mutex<HashMap<String, Arc<Mutex<PulsarProducer>>>>,
}
impl Core {
pub(crate) fn ensure_open(&self) -> Result<(), PulsarError> {
if self.closed.load(Ordering::Acquire) {
return Err(PulsarError::NotConnected);
}
Ok(())
}
}
impl std::fmt::Debug for Core {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Core")
.field("closed", &self.closed.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
pub(crate) type CoreCell = Arc<OnceCell<Arc<Core>>>;
#[derive(Debug, Clone)]
#[must_use]
pub struct PulsarBroker {
url: String,
token: Option<String>,
cell: CoreCell,
}
impl PulsarBroker {
pub fn new(url: impl Into<String>) -> Self {
Self {
url: url.into(),
token: None,
cell: Arc::new(OnceCell::new()),
}
}
pub fn token(mut self, token: impl Into<String>) -> Self {
self.token = Some(token.into());
self
}
#[must_use]
pub fn publisher(&self) -> PulsarPublisher {
PulsarPublisher::new(Arc::clone(&self.cell))
}
}
impl Broker for PulsarBroker {
type Error = PulsarError;
type Connected = ConnectedPulsarBroker;
async fn connect(self) -> Result<Self::Connected, Self::Error> {
let core = self
.cell
.get_or_try_init(async || {
let mut builder = Pulsar::builder(self.url.clone(), TokioExecutor);
if let Some(token) = &self.token {
builder = builder.with_auth(Authentication {
name: "token".to_owned(),
data: token.clone().into_bytes(),
});
}
let client = builder
.build()
.await
.map_err(|e| PulsarError::Connect(box_err(e)))?;
Ok::<_, PulsarError>(Arc::new(Core {
client,
closed: AtomicBool::new(false),
producers: Mutex::new(HashMap::new()),
}))
})
.await?
.clone();
Ok(ConnectedPulsarBroker {
core,
cell: self.cell,
})
}
}
impl DescribeServer for PulsarBroker {
fn describe_server(&self) -> ServerSpec {
ServerSpec::new(
self.url
.trim_start_matches("pulsar+ssl://")
.trim_start_matches("pulsar://"),
"pulsar",
)
}
}
#[derive(Debug)]
pub struct ConnectedPulsarBroker {
pub(crate) core: Arc<Core>,
cell: CoreCell,
}
impl ConnectedPulsarBroker {
#[must_use]
pub fn publisher(&self) -> PulsarPublisher {
PulsarPublisher::new(Arc::clone(&self.cell))
}
pub async fn subscribe_descriptor(
&self,
descriptor: PulsarSubscription,
) -> Result<PulsarSubscriber, PulsarError> {
descriptor.validate()?;
self.core.ensure_open()?;
PulsarSubscriber::open(&self.core, descriptor).await
}
}
impl ConnectedBroker for ConnectedPulsarBroker {
type Error = PulsarError;
type Closed = ();
async fn shutdown(self) -> Result<(), Self::Error> {
self.core.closed.store(true, Ordering::Release);
let producers: Vec<_> = {
let mut map = self.core.producers.lock().await;
map.drain().map(|(_, producer)| producer).collect()
};
for producer in producers {
let mut producer = producer.lock().await;
if let Err(err) = Box::pin(producer.close()).await {
tracing::debug!(error = %err, "pulsar producer close failed");
}
}
Ok(())
}
}
impl Subscribe for ConnectedPulsarBroker {
type Subscriber = PulsarSubscriber;
async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
self.subscribe_descriptor(PulsarSubscription::new(name, "ruststream"))
.await
}
}
impl DefaultPublish for ConnectedPulsarBroker {
type Policy = PulsarPublish;
}