use std::collections::HashMap;
use std::num::{NonZeroU32, NonZeroUsize};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use pulsar::{Authentication, OperationRetryOptions, Pulsar, TokioExecutor};
#[cfg(feature = "testing")]
use ruststream::testing::{Coordinator, InProcess, TestableBroker};
use ruststream::{
Broker, BrokerMoves, ConnectedBroker, DeclareRetryError, DefaultPublish, DescribeServer,
RetryDeclaration, ServerSpec, Subscribe,
};
#[cfg(feature = "testing")]
use ruststream::{BytesMut, OutgoingMessage, RawMessage};
use tokio::runtime::Handle;
use tokio::sync::{Mutex, OnceCell};
use crate::error::{PulsarError, box_err};
#[cfg(feature = "testing")]
use crate::in_process::{self, Bus, Route, Subscriptions};
use crate::publisher::{PulsarProducer, PulsarPublish, PulsarPublisher};
use crate::subscriber::PulsarSubscriber;
use crate::subscription::{DeclaredRetries, PulsarSubscription};
pub(crate) const DEFAULT_MAX_MESSAGE_SIZE: usize = 5 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct OperationRetries {
attempts: Option<NonZeroU32>,
delay: Duration,
timeout: Duration,
}
const DEFAULT_RETRY_DELAY: Duration = Duration::from_secs(5);
const DEFAULT_OPERATION_TIMEOUT: Duration = Duration::from_secs(30);
impl OperationRetries {
pub const fn unbounded() -> Self {
Self {
attempts: None,
delay: DEFAULT_RETRY_DELAY,
timeout: DEFAULT_OPERATION_TIMEOUT,
}
}
pub const fn attempts(attempts: NonZeroU32) -> Self {
Self {
attempts: Some(attempts),
..Self::unbounded()
}
}
pub const fn delay(mut self, delay: Duration) -> Self {
self.delay = delay;
self
}
pub const fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
}
impl Default for OperationRetries {
fn default() -> Self {
Self::unbounded()
}
}
impl From<OperationRetries> for OperationRetryOptions {
fn from(retries: OperationRetries) -> Self {
Self {
operation_timeout: retries.timeout,
retry_delay: retries.delay,
max_retries: retries.attempts.map(|attempts| attempts.get() - 1),
}
}
}
pub(crate) enum Transport {
Client(Pulsar<TokioExecutor>),
#[cfg(feature = "testing")]
InProcess(Arc<Bus>),
}
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Transport>() == size_of::<Pulsar<TokioExecutor>>());
pub(crate) struct Core {
pub(crate) transport: Transport,
pub(crate) closed: AtomicBool,
pub(crate) producers: Mutex<HashMap<String, Arc<Mutex<PulsarProducer>>>>,
default_subscription: Option<String>,
pub(crate) declared_retries: DeclaredRetries,
pub(crate) max_message_size: usize,
pub(crate) runtime: Handle,
#[cfg(feature = "testing")]
opened: Subscriptions,
}
impl Core {
fn new(transport: Transport, broker: &PulsarBroker, runtime: Handle) -> Self {
Self {
transport,
runtime,
closed: AtomicBool::new(false),
producers: Mutex::new(HashMap::new()),
default_subscription: broker.default_subscription.clone(),
declared_retries: DeclaredRetries::default(),
max_message_size: broker.max_message_size,
#[cfg(feature = "testing")]
opened: Subscriptions::default(),
}
}
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>,
default_subscription: Option<String>,
retries: Option<OperationRetries>,
max_message_size: usize,
cell: CoreCell,
}
impl PulsarBroker {
pub fn new(url: impl Into<String>) -> Self {
Self {
url: url.into(),
token: None,
default_subscription: None,
retries: None,
max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
cell: Arc::new(OnceCell::new()),
}
}
pub fn token(mut self, token: impl Into<String>) -> Self {
self.token = Some(token.into());
self
}
pub fn default_subscription(mut self, subscription: impl Into<String>) -> Self {
self.default_subscription = Some(subscription.into());
self
}
pub fn operation_retries(mut self, retries: OperationRetries) -> Self {
self.retries = Some(retries);
self
}
pub fn max_message_size(mut self, bytes: NonZeroUsize) -> Self {
self.max_message_size = bytes.get();
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(),
});
}
if let Some(retries) = self.retries.clone() {
builder = builder.with_operation_retry_options(retries.into());
}
let client = builder
.build()
.await
.map_err(|e| PulsarError::Connect(box_err(e)))?;
Ok::<_, PulsarError>(Arc::new(Core::new(
Transport::Client(client),
&self,
Handle::current(),
)))
})
.await?
.clone();
Ok(ConnectedPulsarBroker {
core,
cell: self.cell,
})
}
}
#[cfg(feature = "testing")]
impl InProcess for PulsarBroker {
async fn connect_in_process(self) -> Result<Self::Connected, Self::Error> {
in_process::check_url(&self.url)?;
let core = self
.cell
.get_or_init(async || {
let runtime = Handle::current();
Arc::new(Core::new(
Transport::InProcess(Bus::new(runtime.clone())),
&self,
runtime,
))
})
.await
.clone();
Ok(ConnectedPulsarBroker {
core,
cell: self.cell,
})
}
}
#[cfg(feature = "testing")]
ruststream::register_testable_broker!(PulsarBroker);
impl DescribeServer for PulsarBroker {
fn describe_server(&self) -> ServerSpec {
ServerSpec::from_url(&self.url, "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()?;
#[cfg(feature = "testing")]
let (route, name, subscription) = (
Route::of(&descriptor)?,
descriptor.source_name().to_owned(),
descriptor.subscription().to_owned(),
);
#[cfg(feature = "testing")]
let recorded = route.clone();
let subscriber = match &self.core.transport {
Transport::Client(client) => {
let (client, core) = (client.clone(), Arc::clone(&self.core));
let topic = descriptor.display_topic();
self.core
.runtime
.spawn(async move { PulsarSubscriber::open(&client, descriptor, &core).await })
.await
.map_err(|err| PulsarError::Subscribe {
topic,
source: box_err(err),
})??
}
#[cfg(feature = "testing")]
Transport::InProcess(bus) => {
let (subscriber, consumer) = in_process::subscribe(bus, descriptor, route)?;
self.core
.opened
.record(name, subscription, recorded, Some(consumer));
return Ok(subscriber);
}
};
#[cfg(feature = "testing")]
self.core.opened.record(name, subscription, recorded, None);
Ok(subscriber)
}
}
impl ConnectedBroker for ConnectedPulsarBroker {
type Error = PulsarError;
type Closed = ();
async fn shutdown(self) -> Result<(), Self::Error> {
self.core.closed.store(true, Ordering::Release);
#[cfg(feature = "testing")]
if let Transport::InProcess(bus) = &self.core.transport {
bus.close();
}
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;
type Copies = BrokerMoves;
async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
let subscription = self.core.default_subscription.as_deref().ok_or_else(|| {
PulsarError::Invalid(format!(
"bare topic name '{name}' has no subscription to join: set \
`PulsarBroker::default_subscription(..)`, or mount \
`PulsarSubscription::new(topic, subscription)`"
))
})?;
let descriptor = self.core.declared_retries.subscription(name, subscription);
self.subscribe_descriptor(descriptor).await
}
fn declare_retry(
&self,
name: &str,
declaration: &RetryDeclaration,
) -> Result<(), DeclareRetryError> {
self.core.declared_retries.declare(name, declaration)
}
}
impl DefaultPublish for ConnectedPulsarBroker {
type Policy = PulsarPublish;
}
#[cfg(feature = "testing")]
impl TestableBroker for ConnectedPulsarBroker {
fn install_coordinator(&self, coordinator: Coordinator) {
if let Transport::InProcess(bus) = &self.core.transport {
bus.install(coordinator);
}
}
fn inject(&self, message: OutgoingMessage<'_>) {
let bus = self.bus("inject");
if self.core.ensure_open().is_err() {
return;
}
let name = message.name();
let owned = OutgoingMessage::produced(name, BytesMut::from(message.payload()))
.with_headers(message.headers().clone());
if let Err(err) = bus.publish(owned, None) {
panic!("the injected message to {name:?} is not one a server takes: {err}");
}
}
fn published(&self, name: &str) -> Vec<RawMessage> {
self.bus("published").published(name)
}
fn routes(&self, destination: &str, subscriptions: &[&str]) -> Vec<usize> {
let router = match &self.core.transport {
Transport::InProcess(bus) => Some(bus.router()),
Transport::Client(_) => None,
};
self.core.opened.routes(destination, subscriptions, router)
}
}
#[cfg(feature = "testing")]
impl ConnectedPulsarBroker {
fn bus(&self, what: &str) -> &Arc<Bus> {
match &self.core.transport {
Transport::InProcess(bus) => bus,
Transport::Client(_) => panic!(
"TestableBroker::{what} reached a broker connected with `connect`; the harness \
drives the transport `connect_in_process` produces"
),
}
}
}
#[cfg(test)]
mod tests {
#[cfg(feature = "testing")]
use ruststream::Publisher;
use ruststream::nonzero;
use super::*;
#[test]
fn the_description_carries_the_address_alone() {
for (url, address) in [
("pulsar://broker:6650", "broker:6650"),
("pulsar+ssl://broker:6651", "broker:6651"),
("pulsar://broker", "broker"),
("pulsar://broker:6650/", "broker:6650"),
("pulsar://admin:s3cret@broker:6650", "broker:6650"),
("pulsar+ssl://admin:s3cret@broker", "broker"),
("pulsar://admin:p@ssw0rd@broker:6650", "broker:6650"),
] {
assert_eq!(
PulsarBroker::new(url).describe_server().host.as_deref(),
Some(address),
"url {url}",
);
}
}
#[test]
fn a_bound_in_tries_becomes_the_retries_that_follow_the_first() {
let bounded = OperationRetryOptions::from(OperationRetries::attempts(nonzero!(3u32)));
assert_eq!(bounded.max_retries, Some(2));
let once = OperationRetryOptions::from(OperationRetries::attempts(nonzero!(1u32)));
assert_eq!(
once.max_retries,
Some(0),
"one try must leave the first answer final",
);
let unbounded = OperationRetryOptions::from(OperationRetries::unbounded());
assert_eq!(unbounded.max_retries, None);
assert_eq!(unbounded.retry_delay, DEFAULT_RETRY_DELAY);
assert_eq!(unbounded.operation_timeout, DEFAULT_OPERATION_TIMEOUT);
}
#[cfg(feature = "testing")]
#[tokio::test]
async fn an_address_connect_refuses_is_refused_in_process() {
for url in ["not a url", "pulsar://"] {
let refused = PulsarBroker::new(url).connect_in_process().await;
assert!(
matches!(refused, Err(PulsarError::Connect(_))),
"{url:?} must be refused, got {refused:?}",
);
}
}
#[cfg(feature = "testing")]
#[tokio::test]
async fn an_early_publisher_publishes_in_process() {
let broker = PulsarBroker::new("pulsar://localhost:6650");
let early = broker.publisher();
let connected = broker
.connect_in_process()
.await
.expect("connects in process");
early
.publish(
OutgoingMessage::produced("orders", BytesMut::from(&b"{}"[..])),
None,
)
.await
.expect("the early publisher reaches the in-process transport");
assert_eq!(connected.published("orders").len(), 1);
connected.shutdown().await.expect("shutdown");
let late = early
.publish(
OutgoingMessage::produced("orders", BytesMut::from(&b"{}"[..])),
None,
)
.await;
assert!(matches!(late, Err(PulsarError::NotConnected)), "{late:?}");
}
#[cfg(feature = "testing")]
#[tokio::test]
async fn a_delivery_held_across_shutdown_is_not_handed_back_in_process() {
use futures::StreamExt;
use ruststream::{IncomingMessage, Subscriber};
let connected = PulsarBroker::new("pulsar://localhost:6650")
.connect_in_process()
.await
.expect("connects in process");
let mut subscriber = connected
.subscribe_descriptor(PulsarSubscription::new("orders", "workers"))
.await
.expect("subscribes");
let publisher = connected.publisher();
for _ in 0..2 {
publisher
.publish(
OutgoingMessage::produced("orders", BytesMut::from(&b"{}"[..])),
None,
)
.await
.expect("publishes");
}
let (requeued, delayed) = {
let mut stream = subscriber.stream();
let mut next = async || {
stream
.next()
.await
.expect("the subscription is open")
.expect("a delivery")
};
(next().await, next().await)
};
connected.shutdown().await.expect("shutdown");
let requeue = requeued.nack(true).await;
assert!(requeue.is_err(), "{requeue:?}");
let delay = delayed.nack_after(Duration::from_secs(1)).await;
assert!(delay.is_err(), "{delay:?}");
}
#[cfg(feature = "testing")]
#[tokio::test]
async fn what_a_closed_consumer_left_unsettled_reaches_the_next_one() {
use futures::StreamExt;
use ruststream::{IncomingMessage, Subscriber};
let connected = PulsarBroker::new("pulsar://localhost:6650")
.connect_in_process()
.await
.expect("connects in process");
let descriptor = PulsarSubscription::new("orders", "workers");
let publisher = connected.publisher();
let mut first = connected
.subscribe_descriptor(descriptor.clone())
.await
.expect("subscribes");
for payload in [b"dropped".as_slice(), b"requeued", b"queued"] {
publisher
.publish(
OutgoingMessage::produced("orders", BytesMut::from(payload)),
None,
)
.await
.expect("publishes");
}
let requeued = {
let mut stream = first.stream();
let dropped = stream.next().await.expect("open").expect("a delivery");
drop(dropped);
stream.next().await.expect("open").expect("a delivery")
};
drop(first);
requeued
.nack(true)
.await
.expect("a closed consumer still hands a delivery back");
let mut next = connected
.subscribe_descriptor(descriptor)
.await
.expect("subscribes again");
let mut stream = next.stream();
let mut back = Vec::new();
for _ in 0..3 {
let delivery = stream.next().await.expect("open").expect("a delivery");
back.push(delivery.payload().to_vec());
delivery.ack().await.expect("ack");
}
back.sort();
assert_eq!(
back,
[
b"dropped".to_vec(),
b"queued".to_vec(),
b"requeued".to_vec()
]
);
}
#[cfg(feature = "testing")]
#[tokio::test]
async fn a_payload_over_the_limit_is_refused() {
let connected = PulsarBroker::new("pulsar://localhost:6650")
.max_message_size(nonzero!(16usize))
.connect_in_process()
.await
.expect("connects in process");
let publisher = connected.publisher();
let publish = async |size: usize| {
publisher
.publish(
OutgoingMessage::produced("orders", BytesMut::from(vec![0; size].as_slice())),
None,
)
.await
};
publish(16)
.await
.expect("a payload at the limit is published");
let refused = publish(17).await;
assert!(
matches!(refused, Err(PulsarError::Publish { .. })),
"{refused:?}"
);
assert_eq!(connected.published("orders").len(), 1);
}
#[test]
fn credentials_never_reach_the_description() {
let described = PulsarBroker::new("pulsar://admin:s3cret@broker:6650").describe_server();
let host = described.host.expect("a URL with a host describes one");
assert!(!host.contains('@'), "userinfo survived in {host:?}");
assert!(!host.contains("admin"), "user name survived in {host:?}");
assert!(!host.contains("s3cret"), "password survived in {host:?}");
}
}