use std::collections::HashMap;
use std::future::{Future, ready};
use std::time::Duration;
use bytes::Bytes;
use pulsar::proto::MessageIdData;
use ruststream::{
AckError, BytesMut, HeaderMap, IncomingMessage, OutgoingMessage, Partitioned, Positioned, Str,
};
use tokio::sync::{mpsc, oneshot};
use crate::error::PulsarError;
#[cfg(feature = "testing")]
use crate::in_process::{Delivery, Settlement};
use crate::subscriber::PulsarSeeker;
pub const PARTITION_KEY_HEADER: &str = "partition-key";
#[derive(Debug)]
pub(crate) enum SettleKind {
Ack,
Nack,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PulsarPosition {
Earliest,
Latest,
MessageId(MessageIdData),
Timestamp(u64),
}
impl PulsarPosition {
#[must_use]
pub fn earliest() -> Self {
Self::Earliest
}
#[must_use]
pub fn latest() -> Self {
Self::Latest
}
#[must_use]
pub fn timestamp(millis: u64) -> Self {
Self::Timestamp(millis)
}
}
#[derive(Debug)]
pub(crate) struct SeekCmd {
pub(crate) position: PulsarPosition,
pub(crate) done: oneshot::Sender<Result<(), PulsarError>>,
}
#[derive(Debug)]
pub(crate) struct SettleCmd {
pub(crate) topic: String,
pub(crate) id: MessageIdData,
pub(crate) kind: SettleKind,
pub(crate) done: oneshot::Sender<Result<(), AckError>>,
}
#[derive(Debug)]
pub(crate) struct NackAfterCmd {
pub(crate) topic: String,
pub(crate) id: MessageIdData,
pub(crate) delay: Duration,
pub(crate) back: SettleSender,
}
#[derive(Debug)]
pub(crate) enum DriverCmd {
Settle(SettleCmd),
NackAfter(NackAfterCmd),
Seek(SeekCmd),
}
pub(crate) type SettleSender = mpsc::UnboundedSender<DriverCmd>;
#[derive(Debug)]
enum Settle {
Driver(SettleSender),
#[cfg(feature = "testing")]
InProcess(Box<Settlement>),
}
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Settle>() == size_of::<SettleSender>());
pub(crate) fn within_ack_timeout(
delay: Duration,
ack_timeout: Option<Duration>,
topic: &str,
) -> Result<(), AckError> {
match ack_timeout {
Some(timeout) if delay >= timeout => Err(AckError::Broker(Box::from(format!(
"topic '{topic}': a retry delay of {delay:?} is not shorter than the subscription's \
ack_timeout of {timeout:?}. A delayed retry is held in this process, and the \
consumer redelivers an unacknowledged message once the ack timeout elapses, so the \
message would come back early. Raise ack_timeout above the delay, or shorten it"
)))),
_ => Ok(()),
}
}
pub struct PulsarMessage {
payload: Bytes,
headers: HeaderMap,
topic: String,
id: MessageIdData,
settle: Settle,
ack_timeout: Option<Duration>,
}
impl std::fmt::Debug for PulsarMessage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PulsarMessage")
.field("topic", &self.topic)
.field("payload_len", &self.payload.len())
.finish_non_exhaustive()
}
}
impl PulsarMessage {
pub(crate) fn new(
message: &pulsar::consumer::Message<Vec<u8>>,
settle: SettleSender,
ack_timeout: Option<Duration>,
) -> Self {
let metadata = message.metadata();
let mut headers = HeaderMap::with_capacity(metadata.properties.len() + 1);
for kv in &metadata.properties {
headers.insert(kv.key.clone(), kv.value.clone());
}
if let Some(key) = &metadata.partition_key {
headers.insert(Str::from_static(PARTITION_KEY_HEADER), key.clone());
}
Self {
payload: Bytes::copy_from_slice(&message.payload.data),
headers,
topic: message.topic.clone(),
id: message.message_id().clone(),
settle: Settle::Driver(settle),
ack_timeout,
}
}
#[cfg(feature = "testing")]
pub(crate) fn in_process(
delivery: Delivery,
topic: String,
settlement: Settlement,
ack_timeout: Option<Duration>,
) -> Self {
let entry_id = u64::try_from(settlement.seq()).unwrap_or(u64::MAX);
Self {
payload: delivery.payload,
headers: delivery.headers,
topic,
id: MessageIdData {
ledger_id: 0,
entry_id,
partition: Some(-1),
..MessageIdData::default()
},
settle: Settle::InProcess(Box::new(settlement)),
ack_timeout,
}
}
#[must_use]
pub fn topic(&self) -> &str {
&self.topic
}
pub(crate) fn seeker(&self) -> PulsarSeeker {
match &self.settle {
Settle::Driver(driver) => PulsarSeeker::new(driver.clone()),
#[cfg(feature = "testing")]
Settle::InProcess(settlement) => settlement.seeker(),
}
}
}
pub(crate) async fn send_settle(
driver: SettleSender,
topic: String,
id: MessageIdData,
kind: SettleKind,
) -> Result<(), AckError> {
let (done, wait) = oneshot::channel();
driver
.send(DriverCmd::Settle(SettleCmd {
topic,
id,
kind,
done,
}))
.map_err(|_| AckError::Broker(Box::from("the subscription's driver task has shut down")))?;
wait.await
.map_err(|_| AckError::Broker(Box::from("the subscription's driver task has shut down")))?
}
impl Positioned for PulsarMessage {
type Position = PulsarPosition;
fn position(&self) -> PulsarPosition {
PulsarPosition::MessageId(self.id.clone())
}
}
impl Partitioned for PulsarMessage {
fn partition_key(&self) -> Option<&[u8]> {
self.headers.get(PARTITION_KEY_HEADER)
}
}
impl IncomingMessage for PulsarMessage {
fn payload(&self) -> &[u8] {
&self.payload
}
fn headers(&self) -> &HeaderMap {
&self.headers
}
async fn ack(self) -> Result<(), AckError> {
match self.settle {
Settle::Driver(driver) => {
send_settle(driver, self.topic, self.id, SettleKind::Ack).await
}
#[cfg(feature = "testing")]
Settle::InProcess(settlement) => {
settlement.settle();
Ok(())
}
}
}
async fn nack(self, requeue: bool) -> Result<(), AckError> {
let kind = if requeue {
SettleKind::Nack
} else {
SettleKind::Ack
};
match self.settle {
Settle::Driver(driver) => send_settle(driver, self.topic, self.id, kind).await,
#[cfg(feature = "testing")]
Settle::InProcess(settlement) => {
if requeue {
(*settlement).requeue(self.payload, self.headers)
} else {
settlement.settle();
Ok(())
}
}
}
}
fn supports_nack_after(&self) -> bool {
true
}
fn nack_after(self, delay: Duration) -> impl Future<Output = Result<(), AckError>> + Send {
if let Err(err) = within_ack_timeout(delay, self.ack_timeout, &self.topic) {
return ready(Err(err));
}
match self.settle {
Settle::Driver(driver) => {
let cmd = DriverCmd::NackAfter(NackAfterCmd {
topic: self.topic,
id: self.id,
delay,
back: driver.clone(),
});
if driver.send(cmd).is_err() {
return ready(Err(AckError::Broker(Box::from(
"the subscription's driver task has shut down",
))));
}
}
#[cfg(feature = "testing")]
Settle::InProcess(settlement) => {
return ready((*settlement).nack_after(delay, self.payload, self.headers));
}
}
ready(Ok(()))
}
fn partition_key(&self) -> Option<&[u8]> {
Partitioned::partition_key(self)
}
}
pub(crate) fn to_pulsar_message(
msg: OutgoingMessage<'_, BytesMut>,
key: Option<&str>,
) -> Result<pulsar::producer::Message, PulsarError> {
let (topic, payload, headers) = msg.into_parts();
let mut properties = HashMap::with_capacity(headers.len());
let mut partition_key = key.map(ToOwned::to_owned);
for (name, value) in headers.iter() {
let text = std::str::from_utf8(value)
.map_err(|err| PulsarError::Publish {
topic: topic.to_owned(),
source: Box::from(format!(
"header '{name}' is not UTF-8 ({err}); Pulsar carries properties as text"
)),
})?
.to_owned();
if name == PARTITION_KEY_HEADER {
partition_key.get_or_insert(text);
} else {
properties.insert(name.to_owned(), text);
}
}
Ok(pulsar::producer::Message {
payload: Vec::from(payload),
properties,
partition_key,
..Default::default()
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn partition_key_header_becomes_the_partition_key() {
let mut headers = HeaderMap::new();
headers.insert(PARTITION_KEY_HEADER, "user-42");
headers.insert("x-tenant", "acme");
let outgoing = OutgoingMessage::new("orders", b"{}".as_slice()).with_headers(headers);
let message = to_pulsar_message(outgoing, None).expect("text headers are carried");
assert_eq!(message.partition_key.as_deref(), Some("user-42"));
assert_eq!(
message.properties.get("x-tenant").map(String::as_str),
Some("acme")
);
assert!(!message.properties.contains_key(PARTITION_KEY_HEADER));
}
#[test]
fn the_call_site_key_wins_over_the_header() {
let mut headers = HeaderMap::new();
headers.insert(PARTITION_KEY_HEADER, "user-7");
let outgoing = OutgoingMessage::new("orders", b"{}".as_slice()).with_headers(headers);
let message =
to_pulsar_message(outgoing, Some("user-42")).expect("text headers are carried");
assert_eq!(message.partition_key.as_deref(), Some("user-42"));
assert!(!message.properties.contains_key(PARTITION_KEY_HEADER));
}
#[test]
fn a_delay_is_bounded_by_the_ack_timeout() {
let timeout = Duration::from_secs(30);
assert!(within_ack_timeout(Duration::from_secs(29), Some(timeout), "orders").is_ok());
assert!(within_ack_timeout(Duration::from_mins(5), None, "orders").is_ok());
let refused = within_ack_timeout(timeout, Some(timeout), "orders")
.expect_err("a delay equal to the timeout must be refused");
let message = format!("{refused:?}");
assert!(message.contains("ack_timeout"), "{message}");
}
#[test]
fn the_ack_timeout_refusal_reads_as_one_sentence() {
let timeout = Duration::from_secs(30);
let AckError::Broker(source) = within_ack_timeout(timeout, Some(timeout), "orders")
.expect_err("a delay equal to the timeout must be refused")
else {
panic!("the refusal is the broker's own");
};
assert_eq!(
source.to_string(),
"topic 'orders': a retry delay of 30s is not shorter than the subscription's \
ack_timeout of 30s. A delayed retry is held in this process, and the consumer \
redelivers an unacknowledged message once the ack timeout elapses, so the message \
would come back early. Raise ack_timeout above the delay, or shorten it",
);
}
#[test]
fn the_client_message_takes_the_buffer_the_publish_wrote() {
let buffer = BytesMut::from(&b"{\"id\":1}"[..]);
let at = buffer.as_ptr();
let outgoing = OutgoingMessage::produced("orders", buffer);
let message = to_pulsar_message(outgoing, None).expect("text headers are carried");
assert_eq!(
message.payload.as_ptr(),
at,
"the payload must be the buffer the publish wrote, not a copy of it",
);
}
#[test]
fn a_publish_that_names_no_key_leaves_unkeyed() {
let outgoing = OutgoingMessage::new("orders", b"{}".as_slice());
let message = to_pulsar_message(outgoing, None).expect("text headers are carried");
assert_eq!(message.partition_key, None);
}
#[test]
fn a_header_that_is_not_text_is_refused() {
let mut headers = HeaderMap::new();
headers.insert("x-raw", vec![0xff, b'a']);
let outgoing = OutgoingMessage::new("orders", b"{}".as_slice()).with_headers(headers);
let refused = to_pulsar_message(outgoing, None).expect_err("a binary header is refused");
assert!(
matches!(&refused, PulsarError::Publish { topic, .. } if topic == "orders"),
"{refused:?}",
);
}
}