pub use self::ZmqRpcPublish as Publish;
pub mod prelude {
pub use ruststream::prelude::*;
pub use crate::endpoint::ZmqEndpoint;
pub use ruststream::RequestReply;
pub use super::{Publish, ZmqRpc};
}
use std::fmt;
use std::future::{Future, ready};
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use bytes::Bytes;
use futures::Stream;
use futures::lock::Mutex;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::{
Broker, BytesMut, ConnectedBroker, DefaultPublish, DescribeServer, NamedCopies,
OutgoingMessage, PairError, PublishPolicy, Publisher, RequestReply, ServerSpec, Str, Subscribe,
Subscriber, Take,
};
use tokio::sync::OnceCell;
use zeromq::prelude::*;
use zeromq::util::PeerIdentity;
use zeromq::{DealerSocket, RouterSendHalf, RouterSocket, SocketOptions};
#[cfg(feature = "asyncapi")]
use crate::bindings::{self, SocketPair};
#[cfg(feature = "testing")]
use crate::common::DeliveryReceiver;
use crate::common::{
DEFAULT_READ_AHEAD, DriverHandle, Lifecycle, SharedLifecycle, WireSubscriber, delivery_channel,
send_with_retry,
};
use crate::endpoint::{Endpoint, EndpointRole, ZmqEndpoint};
use crate::error::ZmqError;
#[cfg(feature = "testing")]
use crate::in_process::{self, Bus};
use crate::message::ZmqMessage;
use crate::wire;
pub(crate) const REPLY_PREFIX: &str = "zmq-reply:";
pub(crate) const REPLY_TO_HEADER: &str = "reply-to";
pub(crate) const CORRELATION_ID_HEADER: &str = "correlation-id";
#[cfg(feature = "asyncapi")]
pub(crate) const REPLY_ADDRESS_LOCATION: &str = "$message.header#/reply-to";
pub(crate) fn reply_address(identity: &[u8]) -> String {
format!("{REPLY_PREFIX}{}", hex_encode(identity))
}
pub(crate) fn new_correlation_id() -> String {
format!("req-{}-{}", std::process::id(), hex_encode(&rand_suffix()))
}
fn hex_encode(bytes: &[u8]) -> String {
use std::fmt::Write as _;
let mut out = String::with_capacity(bytes.len() * 2);
for b in bytes {
let _ = write!(out, "{b:02x}");
}
out
}
fn hex_decode(text: &str) -> Option<Vec<u8>> {
if !text.len().is_multiple_of(2) {
return None;
}
(0..text.len())
.step_by(2)
.map(|i| u8::from_str_radix(&text[i..i + 2], 16).ok())
.collect()
}
#[derive(Debug, Clone)]
#[must_use]
pub struct ZmqRpc {
endpoint: Endpoint,
read_ahead: NonZeroUsize,
cell: Arc<OnceCell<RpcShared>>,
}
#[derive(Clone)]
pub(crate) struct RpcShared {
lifecycle: SharedLifecycle,
router_tx: Arc<OnceCell<ReplyRoute>>,
}
#[derive(Clone)]
enum ReplyRoute {
Router(Arc<Mutex<Option<RouterSendHalf>>>),
#[cfg(feature = "testing")]
InProcess(Arc<Bus>),
}
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<ReplyRoute>() == size_of::<Arc<Mutex<Option<RouterSendHalf>>>>());
impl fmt::Debug for RpcShared {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RpcShared").finish_non_exhaustive()
}
}
impl ZmqRpc {
pub fn new<Role: EndpointRole>(endpoint: ZmqEndpoint<Role>) -> Self {
Self {
endpoint: endpoint.into_inner(),
read_ahead: DEFAULT_READ_AHEAD,
cell: Arc::new(OnceCell::new()),
}
}
pub const fn read_ahead(mut self, read_ahead: NonZeroUsize) -> Self {
self.read_ahead = read_ahead;
self
}
#[must_use]
pub fn publisher(&self) -> ZmqRpcPublisher {
ZmqRpcPublisher {
cell: Arc::clone(&self.cell),
}
}
}
impl Broker for ZmqRpc {
type Error = ZmqError;
type Connected = ConnectedZmqRpc;
async fn connect(self) -> Result<Self::Connected, Self::Error> {
self.connect_with(Lifecycle::new).await
}
}
impl ZmqRpc {
pub(crate) async fn connect_with(
self,
lifecycle: fn(Endpoint) -> Lifecycle,
) -> Result<ConnectedZmqRpc, ZmqError> {
let shared = self
.cell
.get_or_try_init(async || {
self.endpoint.validate()?;
Ok::<_, ZmqError>(RpcShared {
lifecycle: Arc::new(lifecycle(self.endpoint.clone())),
router_tx: Arc::new(OnceCell::new()),
})
})
.await?
.clone();
Ok(ConnectedZmqRpc {
shared,
read_ahead: self.read_ahead,
cell: self.cell,
})
}
}
impl DescribeServer for ZmqRpc {
fn describe_server(&self) -> ServerSpec {
self.endpoint.server_spec()
}
}
#[derive(Debug)]
pub struct ConnectedZmqRpc {
shared: RpcShared,
read_ahead: NonZeroUsize,
cell: Arc<OnceCell<RpcShared>>,
}
impl ConnectedZmqRpc {
#[cfg(feature = "testing")]
pub(crate) fn lifecycle(&self) -> &Lifecycle {
&self.shared.lifecycle
}
#[must_use]
pub fn bound_address(&self) -> Option<String> {
self.shared.lifecycle.bound_address()
}
#[must_use]
pub fn publisher(&self) -> ZmqRpcPublisher {
ZmqRpcPublisher {
cell: Arc::clone(&self.cell),
}
}
}
impl ConnectedBroker for ConnectedZmqRpc {
type Error = ZmqError;
type Closed = ();
fn shutdown(self) -> impl Future<Output = Result<(), Self::Error>> {
self.shared.lifecycle.closed.store(true, Ordering::Release);
ready(Ok(()))
}
}
pub struct ZmqRpcSubscriber {
name: String,
inner: WireSubscriber,
}
impl fmt::Debug for ZmqRpcSubscriber {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ZmqRpcSubscriber")
.field("name", &self.name)
.finish_non_exhaustive()
}
}
impl Subscriber for ZmqRpcSubscriber {
type Message = ZmqMessage;
type Error = ZmqError;
fn stream(&mut self) -> impl Stream<Item = Result<ZmqMessage, ZmqError>> + Send + '_ {
self.inner.stream()
}
}
impl Subscribe for ConnectedZmqRpc {
type Subscriber = ZmqRpcSubscriber;
type Copies = NamedCopies;
async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
self.shared.lifecycle.ensure_open()?;
#[cfg(feature = "testing")]
if let Some(bus) = self.shared.lifecycle.in_process_bus() {
let slot = self.shared.lifecycle.attach_in_process(name).await?;
let (rx, registration) = bus.subscribe(name, read_request);
let _ = self
.shared
.router_tx
.set(ReplyRoute::InProcess(Arc::clone(bus)));
return Ok(ZmqRpcSubscriber {
name: name.to_owned(),
inner: WireSubscriber {
rx: DeliveryReceiver::InProcess(rx),
_driver: DriverHandle::InProcess(registration, slot),
},
});
}
let lifecycle = Arc::clone(&self.shared.lifecycle);
let subscription = name.to_owned();
let (socket, slot) = self
.shared
.lifecycle
.on_runtime(async move {
let mut socket = RouterSocket::new();
let slot = lifecycle
.attach_receiver(&mut socket, &subscription)
.await?;
Ok((socket, slot))
})
.await?;
let (send_half, mut recv_half) = socket.split();
let route = self
.shared
.router_tx
.get_or_init(async || ReplyRoute::Router(Arc::new(Mutex::new(None))))
.await;
match route {
ReplyRoute::Router(router) => *router.lock().await = Some(send_half),
#[cfg(feature = "testing")]
ReplyRoute::InProcess(_) => {}
}
let (tx, rx) = delivery_channel(self.read_ahead);
let task = self.shared.lifecycle.spawn(async move {
loop {
match recv_half.recv().await {
Ok(message) => {
let Some(item) = read_request(message) else {
continue;
};
if tx.send(item).await.is_err() {
break;
}
}
Err(err) => {
if tx
.send(Err(ZmqError::Receive(err.to_string())))
.await
.is_err()
{
break;
}
}
}
}
});
Ok(ZmqRpcSubscriber {
name: name.to_owned(),
inner: WireSubscriber {
rx,
_driver: DriverHandle::Task(task, slot),
},
})
}
}
pub(crate) fn read_request(message: zeromq::ZmqMessage) -> Option<Result<ZmqMessage, ZmqError>> {
let mut frames = message.into_vecdeque();
let identity = frames.pop_front()?;
let rest: Result<zeromq::ZmqMessage, _> = frames.try_into();
let Ok(rest) = rest else {
return Some(Err(ZmqError::Wire(
"a request needs name and payload frames".into(),
)));
};
Some(wire::decode(rest).map(|(name, mut headers, payload)| {
headers.insert(Str::from_static(REPLY_TO_HEADER), reply_address(&identity));
ZmqMessage::new(name, headers, payload)
}))
}
#[derive(Clone)]
pub struct ZmqRpcPublisher {
cell: Arc<OnceCell<RpcShared>>,
}
impl fmt::Debug for ZmqRpcPublisher {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ZmqRpcPublisher").finish_non_exhaustive()
}
}
impl ZmqRpcPublisher {
fn shared(&self) -> Result<&RpcShared, ZmqError> {
let shared = self.cell.get().ok_or(ZmqError::NotConnected)?;
shared.lifecycle.ensure_open()?;
Ok(shared)
}
}
impl Publisher for ZmqRpcPublisher {
type Payload = Take;
type Error = ZmqError;
type Options = ();
async fn publish(
&self,
msg: OutgoingMessage<'_, BytesMut>,
_options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let shared = self.shared()?;
let Some(identity_hex) = msg.name().strip_prefix(REPLY_PREFIX) else {
return Err(ZmqError::Send {
name: msg.name().to_owned(),
reason: format!(
"the rpc publisher routes '{REPLY_PREFIX}...' replies; use request() for outbound requests"
),
});
};
let identity = hex_decode(identity_hex).ok_or_else(|| ZmqError::Send {
name: msg.name().to_owned(),
reason: "malformed reply address".to_owned(),
})?;
let router = shared.router_tx.get().ok_or_else(|| ZmqError::Send {
name: msg.name().to_owned(),
reason: "no responder subscription is attached".to_owned(),
})?;
let (name, payload, headers) = msg.into_parts();
let message = wire::encode_addressed_to(
name,
Bytes::from(identity),
"reply",
&headers,
payload.freeze(),
)?;
match router {
ReplyRoute::Router(router) => {
let mut guard = router.lock().await;
let Some(half) = guard.as_mut() else {
return Err(ZmqError::Send {
name: name.to_owned(),
reason: "no responder subscription is attached".to_owned(),
});
};
let sent = half.send(message).await;
drop(guard);
sent.map_err(|e| ZmqError::Send {
name: name.to_owned(),
reason: e.to_string(),
})
}
#[cfg(feature = "testing")]
ReplyRoute::InProcess(bus) => bus.reply(name, message),
}
}
}
impl RequestReply for ZmqRpcPublisher {
type Reply = ZmqMessage;
async fn request(
&self,
msg: OutgoingMessage<'_, BytesMut>,
timeout: Duration,
) -> Result<Self::Reply, Self::Error> {
let shared = self.shared()?;
#[cfg(feature = "testing")]
if let Some(bus) = shared.lifecycle.in_process_bus() {
return request_in_process(bus, &shared.lifecycle, msg, timeout).await;
}
let address = shared.lifecycle.send_target().address().to_owned();
let mut options = SocketOptions::default();
options.peer_identity(PeerIdentity::new());
options.connect_timeout(timeout.min(Duration::from_secs(10)));
let mut dealer = DealerSocket::with_options(options);
dealer
.connect(&address)
.await
.map_err(|e| ZmqError::Endpoint {
endpoint: address.clone(),
source: crate::error::box_err(e),
})?;
let correlation = msg
.headers()
.correlation_id()
.map_or_else(new_correlation_id, str::to_owned);
let (name, payload, mut headers) = msg.into_parts();
headers.insert(Str::from_static(CORRELATION_ID_HEADER), correlation.clone());
let request = wire::encode_to(name, name, &headers, payload.freeze())?;
send_with_retry(&mut dealer, name, request).await?;
let exchange = async {
loop {
let reply = dealer
.recv()
.await
.map_err(|e| ZmqError::Receive(e.to_string()))?;
let (name, reply_headers, payload) = wire::decode(reply)?;
if reply_headers.correlation_id() == Some(correlation.as_str()) {
return Ok(ZmqMessage::new(name, reply_headers, payload));
}
}
};
tokio::time::timeout(timeout, exchange)
.await
.unwrap_or(Err(ZmqError::RequestTimeout))
}
}
#[cfg(feature = "testing")]
async fn request_in_process(
bus: &Bus,
lifecycle: &Lifecycle,
msg: OutgoingMessage<'_, BytesMut>,
timeout: Duration,
) -> Result<ZmqMessage, ZmqError> {
let correlation = msg
.headers()
.correlation_id()
.map_or_else(new_correlation_id, str::to_owned);
let (name, payload, mut headers) = msg.into_parts();
headers.insert(Str::from_static(CORRELATION_ID_HEADER), correlation.clone());
let request = wire::encode_to(name, name, &headers, payload.freeze())?;
let local = lifecycle.local_listener().is_some();
in_process::request(bus, local, name, request, &correlation, timeout).await
}
#[cfg(feature = "testing")]
pub(crate) fn peer_identity() -> Bytes {
Bytes::copy_from_slice(&rand_suffix())
}
fn rand_suffix() -> [u8; 8] {
use std::sync::atomic::AtomicU64;
static SEQ: AtomicU64 = AtomicU64::new(0);
let seq = SEQ.fetch_add(1, Ordering::Relaxed);
seq.to_be_bytes()
}
#[derive(Debug, Clone, Copy, Default)]
#[must_use]
pub struct ZmqRpcPublish;
impl PublishPolicy<ConnectedZmqRpc> for ZmqRpcPublish {
type Live = ZmqRpcPublisher;
fn pair(
self,
connected: &ConnectedZmqRpc,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.publisher()))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
bindings::channel(SocketPair::DealerRouter, channel)
}
#[cfg(feature = "asyncapi")]
fn reply_address_location(&self) -> Option<&'static str> {
Some(REPLY_ADDRESS_LOCATION)
}
}
impl DefaultPublish for ConnectedZmqRpc {
type Policy = ZmqRpcPublish;
}
#[cfg(test)]
mod tests {
use ruststream::{BatchSubscriber, Subscriber};
use super::ZmqRpcSubscriber;
use crate::ZmqSubscriber;
fn batches<S: BatchSubscriber>() {}
fn delivers<S: Subscriber>() {}
#[test]
fn the_one_way_patterns_batch_and_the_responder_only_delivers() {
batches::<ZmqSubscriber>();
delivers::<ZmqRpcSubscriber>();
}
}