#![cfg_attr(
not(feature = "testing"),
allow(clippy::infallible_destructuring_match)
)]
use std::future::{Future, ready};
use std::sync::Arc;
use pulsar::routing_policy::RoutingPolicy;
use pulsar::{ProducerOptions, Pulsar, TokioExecutor};
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::runtime::{PublishBuilder, PublishSink};
use ruststream::{BytesMut, OutgoingMessage, PairError, PublishPolicy, Publisher, Take};
use tokio::sync::Mutex;
#[cfg(feature = "asyncapi")]
use crate::bindings::topic_channel;
use crate::broker::{ConnectedPulsarBroker, Core, CoreCell, Transport};
use crate::error::{PulsarError, box_err};
use crate::message::to_pulsar_message;
use crate::topic::PulsarTopic;
pub(crate) type PulsarProducer = pulsar::Producer<TokioExecutor>;
fn keyed_routing() -> ProducerOptions {
ProducerOptions {
routing_policy: Some(RoutingPolicy::RoundRobin),
..ProducerOptions::default()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PulsarPublishOptions {
pub partition_key: Option<String>,
}
#[derive(Clone)]
pub struct PulsarPublisher {
cell: CoreCell,
}
impl std::fmt::Debug for PulsarPublisher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PulsarPublisher").finish_non_exhaustive()
}
}
impl PulsarPublisher {
pub(crate) fn new(cell: CoreCell) -> Self {
Self { cell }
}
fn core(&self) -> Result<&Core, PulsarError> {
let core = self.cell.get().ok_or(PulsarError::NotConnected)?;
core.ensure_open()?;
Ok(core)
}
#[allow(clippy::significant_drop_tightening)]
async fn producer_for(
&self,
core: &Core,
client: &Pulsar<TokioExecutor>,
topic: &str,
) -> Result<Arc<Mutex<PulsarProducer>>, PulsarError> {
let full = PulsarTopic::parse(topic)?.as_str().to_owned();
let mut producers = core.producers.lock().await;
if let Some(producer) = producers.get(&full) {
return Ok(Arc::clone(producer));
}
let (client, named) = (client.clone(), full.clone());
let producer = core
.runtime
.spawn(async move {
Box::pin(
client
.producer()
.with_topic(named)
.with_options(keyed_routing())
.build(),
)
.await
})
.await
.map_err(|e| PulsarError::Publish {
topic: topic.to_owned(),
source: box_err(e),
})?
.map_err(|e| PulsarError::Publish {
topic: topic.to_owned(),
source: box_err(e),
})?;
let producer = Arc::new(Mutex::new(producer));
producers.insert(full, Arc::clone(&producer));
Ok(producer)
}
}
impl Publisher for PulsarPublisher {
type Payload = Take;
type Error = PulsarError;
type Options = PulsarPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_, BytesMut>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let core = self.core()?;
let size = msg.payload().len();
if size > core.max_message_size {
return Err(PulsarError::Publish {
topic: msg.name().to_owned(),
source: Box::from(format!(
"the payload is {size} bytes, over the {} the broker's max_message_size \
allows",
core.max_message_size
)),
});
}
let key = options.and_then(|options| options.partition_key.as_deref());
let client = match &core.transport {
Transport::Client(client) => client,
#[cfg(feature = "testing")]
Transport::InProcess(bus) => return bus.publish(msg, key),
};
let topic = msg.name();
let producer = Box::pin(self.producer_for(core, client, topic)).await?;
let message = to_pulsar_message(msg, key)?;
let receipt = {
let mut producer = producer.lock().await;
Box::pin(producer.send_non_blocking(message))
.await
.map_err(|e| PulsarError::Publish {
topic: topic.to_owned(),
source: box_err(e),
})?
};
receipt.await.map(|_| ()).map_err(|e| PulsarError::Publish {
topic: topic.to_owned(),
source: box_err(e),
})
}
}
pub trait PulsarPublishSteps {
#[must_use]
fn partition_key(self, key: impl Into<String>) -> Self;
}
impl<Sink, Body, Enc, Hdrs, Dest> PulsarPublishSteps for PublishBuilder<Sink, Body, Enc, Hdrs, Dest>
where
Sink: PublishSink<Options = PulsarPublishOptions>,
{
fn partition_key(mut self, key: impl Into<String>) -> Self {
self.options_mut()
.get_or_insert_with(PulsarPublishOptions::default)
.partition_key = Some(key.into());
self
}
}
#[derive(Debug, Clone, Copy, Default)]
#[must_use]
pub struct PulsarPublish;
#[cfg(feature = "asyncapi")]
impl PulsarPublish {
fn describe_channel(channel: &str) -> Bindings {
PulsarTopic::parse(channel).map_or_else(|_| Bindings::new(), |topic| topic_channel(&topic))
}
}
impl PublishPolicy<ConnectedPulsarBroker> for PulsarPublish {
type Live = PulsarPublisher;
fn pair(
self,
connected: &ConnectedPulsarBroker,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.publisher()))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
Self::describe_channel(channel)
}
}
#[cfg(all(test, feature = "testing"))]
mod tests {
use ruststream::runtime::PublishExt;
use ruststream::testing::{InProcess, TestableBroker};
use ruststream::{HeaderMap, Outgoing, Serialized};
use serde::Serialize;
use super::PulsarPublishSteps;
use crate::{ConnectedPulsarBroker, PARTITION_KEY_HEADER, PulsarBroker};
#[derive(Outgoing, Serialized)]
struct Record(&'static [u8]);
#[derive(Serialize)]
struct OrderMeta {
tenant: &'static str,
}
#[derive(Outgoing, Serialize)]
#[outgoing(name = "orders.done", headers = OrderMeta)]
struct OrderDone {
id: u64,
}
async fn connected() -> ConnectedPulsarBroker {
PulsarBroker::new("pulsar://localhost:6650")
.connect_in_process()
.await
.expect("the broker connects in process")
}
#[tokio::test]
async fn the_step_wins_over_a_call_site_header() {
let broker = connected().await;
let mut headers = HeaderMap::new();
headers.insert(PARTITION_KEY_HEADER, "user-7");
broker
.publisher()
.message(&Record(b"{}"))
.with_headers(headers)
.to("orders")
.partition_key("user-42")
.publish()
.await
.expect("publish succeeds");
let sent = broker.published("orders");
assert_eq!(
sent[0].headers().get(PARTITION_KEY_HEADER),
Some(b"user-42".as_slice())
);
}
#[tokio::test]
async fn the_step_composes_with_a_declared_header_contract() {
let broker = connected().await;
broker
.publisher()
.message(&OrderDone { id: 1 })
.with_headers(&OrderMeta { tenant: "acme" })
.partition_key("user-42")
.publish()
.await
.expect("publish succeeds");
let sent = broker.published("orders.done");
assert_eq!(sent[0].headers().get_str("tenant"), Some("acme"));
assert_eq!(
sent[0].headers().get(PARTITION_KEY_HEADER),
Some(b"user-42".as_slice())
);
}
}