use futures::{SinkExt, StreamExt, future};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use thiserror::Error;
use tokio::net::TcpStream;
use tokio::sync::{Mutex, broadcast, mpsc, oneshot};
use tokio_util::codec::Framed;
use crate::codec::{StompCodec, StompItem};
use crate::frame::Frame;
use crate::parser::DEFAULT_MAX_FRAME_SIZE;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Heartbeat {
pub send_ms: u32,
pub receive_ms: u32,
}
impl Heartbeat {
pub fn new(send_ms: u32, receive_ms: u32) -> Self {
Self {
send_ms,
receive_ms,
}
}
pub fn disabled() -> Self {
Self::new(0, 0)
}
pub fn from_duration(interval: Duration) -> Self {
let ms = interval.as_millis().min(u32::MAX as u128) as u32;
Self::new(ms, ms)
}
}
impl Default for Heartbeat {
fn default() -> Self {
Self::new(10000, 10000)
}
}
impl std::fmt::Display for Heartbeat {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{},{}", self.send_ms, self.receive_ms)
}
}
#[derive(Clone)]
pub(crate) struct SubscriptionEntry {
pub(crate) id: String,
pub(crate) sender: mpsc::Sender<Frame>,
pub(crate) ack: String,
pub(crate) headers: Vec<(String, String)>,
}
pub(crate) type Subscriptions = HashMap<String, Vec<SubscriptionEntry>>;
pub(crate) type PendingMap = HashMap<String, VecDeque<(String, Frame)>>;
pub(crate) type ResubEntry = (String, String, String, Vec<(String, String)>);
pub(crate) type PendingReceipts = HashMap<String, oneshot::Sender<Result<(), ServerError>>>;
#[derive(Error, Debug)]
pub enum ConnError {
#[error("io error: {0}")]
Io(#[from] std::io::Error),
#[error("protocol error: {0}")]
Protocol(String),
#[error("receipt timeout: no RECEIPT received for '{0}' within timeout")]
ReceiptTimeout(String),
#[error("server rejected connection: {0}")]
ServerRejected(ServerError),
#[error("frame rejected: {0}")]
FrameRejected(ServerError),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServerError {
pub message: String,
pub body: Option<String>,
pub receipt_id: Option<String>,
pub frame: Frame,
}
impl ServerError {
pub fn from_frame(frame: Frame) -> Self {
let message = frame
.get_header("message")
.unwrap_or("unknown error")
.to_string();
let body = if frame.body.is_empty() {
None
} else {
String::from_utf8(frame.body.clone()).ok()
};
let receipt_id = frame.get_header("receipt-id").map(|s| s.to_string());
Self {
message,
body,
receipt_id,
frame,
}
}
}
impl std::fmt::Display for ServerError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "STOMP server error: {}", self.message)?;
if let Some(body) = &self.body {
write!(f, " - {}", body)?;
}
Ok(())
}
}
impl std::error::Error for ServerError {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ReceivedFrame {
Frame(Frame),
Error(ServerError),
}
impl ReceivedFrame {
pub fn is_error(&self) -> bool {
matches!(self, ReceivedFrame::Error(_))
}
pub fn is_frame(&self) -> bool {
matches!(self, ReceivedFrame::Frame(_))
}
pub fn into_frame(self) -> Option<Frame> {
match self {
ReceivedFrame::Frame(f) => Some(f),
ReceivedFrame::Error(_) => None,
}
}
pub fn into_error(self) -> Option<ServerError> {
match self {
ReceivedFrame::Frame(_) => None,
ReceivedFrame::Error(e) => Some(e),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AckMode {
Auto,
Client,
ClientIndividual,
}
impl AckMode {
fn as_str(&self) -> &'static str {
match self {
AckMode::Auto => "auto",
AckMode::Client => "client",
AckMode::ClientIndividual => "client-individual",
}
}
}
#[derive(Clone, Default)]
pub struct ConnectOptions {
pub accept_version: Option<String>,
pub client_id: Option<String>,
pub host: Option<String>,
pub headers: Vec<(String, String)>,
pub heartbeat_tx: Option<mpsc::Sender<()>>,
pub disconnect_timeout: Option<Duration>,
pub connect_timeout: Option<Duration>,
pub max_frame_size: Option<usize>,
}
impl std::fmt::Debug for ConnectOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConnectOptions")
.field("accept_version", &self.accept_version)
.field("client_id", &self.client_id)
.field("host", &self.host)
.field("headers", &self.headers)
.field(
"heartbeat_tx",
&self.heartbeat_tx.as_ref().map(|_| "Some(...)"),
)
.field("disconnect_timeout", &self.disconnect_timeout)
.field("connect_timeout", &self.connect_timeout)
.field("max_frame_size", &self.max_frame_size)
.finish()
}
}
impl ConnectOptions {
pub fn new() -> Self {
Self::default()
}
pub fn accept_version(mut self, version: impl Into<String>) -> Self {
self.accept_version = Some(version.into());
self
}
pub fn client_id(mut self, id: impl Into<String>) -> Self {
self.client_id = Some(id.into());
self
}
pub fn host(mut self, host: impl Into<String>) -> Self {
self.host = Some(host.into());
self
}
pub fn header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.push((key.into(), value.into()));
self
}
pub fn disconnect_timeout(mut self, timeout: Duration) -> Self {
self.disconnect_timeout = Some(timeout);
self
}
pub fn connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = Some(timeout);
self
}
pub fn max_frame_size(mut self, max_frame_size: usize) -> Self {
self.max_frame_size = Some(max_frame_size);
self
}
pub fn heartbeat_notify(mut self, tx: mpsc::Sender<()>) -> Self {
self.heartbeat_tx = Some(tx);
self
}
}
pub fn parse_heartbeat_header(header: &str) -> (u64, u64) {
let mut parts = header.split(',');
let cx = parts
.next()
.and_then(|s| s.trim().parse::<u64>().ok())
.unwrap_or(0);
let cy = parts
.next()
.and_then(|s| s.trim().parse::<u64>().ok())
.unwrap_or(0);
(cx, cy)
}
pub fn negotiate_heartbeats(
client_out: u64,
client_in: u64,
server_out: u64,
server_in: u64,
) -> (Option<Duration>, Option<Duration>) {
let negotiated_out_ms = std::cmp::max(client_out, server_in);
let negotiated_in_ms = std::cmp::max(client_in, server_out);
let outgoing = if negotiated_out_ms == 0 {
None
} else {
Some(Duration::from_millis(negotiated_out_ms))
};
let incoming = if negotiated_in_ms == 0 {
None
} else {
Some(Duration::from_millis(negotiated_in_ms))
};
(outgoing, incoming)
}
fn extract_destination_from_error(frame: &Frame) -> Option<String> {
if let Some(dest) = frame.get_header("destination") {
return Some(dest.to_string());
}
let message = frame.get_header("message").unwrap_or("");
let body = String::from_utf8_lossy(&frame.body);
let text = format!("{} {}", message, body);
for prefix in ["/topic/", "/queue/"] {
if let Some(start) = text.find(prefix) {
let rest = &text[start..];
let end = rest
.find(|c: char| c.is_whitespace() || c == ',' || c == '"' || c == '\'')
.unwrap_or(rest.len());
if end > prefix.len() {
return Some(rest[..end].to_string());
}
}
}
None
}
fn extract_subscription_id_from_error(frame: &Frame) -> Option<String> {
let message = frame.get_header("message").unwrap_or("");
let body = String::from_utf8_lossy(&frame.body);
let text = format!("{} {}", message, body);
if let Some(idx) = text.to_lowercase().find("subscription ") {
let rest = &text[idx + 13..]; let end = rest
.find(|c: char| c.is_whitespace() || c == ',' || c == '"' || c == '\'')
.unwrap_or(rest.len());
if end > 0 {
return Some(rest[..end].to_string());
}
}
None
}
async fn lookup_destination_by_sub_id(
sub_id: &str,
subscriptions: &Arc<Mutex<Subscriptions>>,
) -> Option<String> {
let map = subscriptions.lock().await;
for (dest, entries) in map.iter() {
for entry in entries {
if entry.id == sub_id {
return Some(dest.clone());
}
}
}
None
}
fn deliver_and_keep(entry: &SubscriptionEntry, frame: &Frame) -> bool {
match entry.sender.try_send(frame.clone()) {
Ok(()) => true,
Err(mpsc::error::TrySendError::Full(_)) => true,
Err(mpsc::error::TrySendError::Closed(_)) => false,
}
}
#[derive(Clone)]
pub struct Connection {
outbound_tx: mpsc::Sender<StompItem>,
inbound_rx: Arc<Mutex<mpsc::Receiver<Frame>>>,
shutdown_tx: broadcast::Sender<()>,
subscriptions: Arc<Mutex<Subscriptions>>,
sub_id_counter: Arc<AtomicU64>,
pending: Arc<Mutex<PendingMap>>,
pending_receipts: Arc<Mutex<PendingReceipts>>,
disconnect_timeout: Duration,
}
#[derive(Debug)]
pub struct ReceiptHandle {
receipt_id: String,
rx: oneshot::Receiver<Result<(), ServerError>>,
pending_receipts: Arc<Mutex<PendingReceipts>>,
}
impl ReceiptHandle {
pub fn receipt_id(&self) -> &str {
&self.receipt_id
}
pub async fn wait(self, timeout: Duration) -> Result<(), ConnError> {
let ReceiptHandle {
receipt_id,
rx,
pending_receipts,
} = self;
match tokio::time::timeout(timeout, rx).await {
Ok(Ok(Ok(()))) => Ok(()),
Ok(Ok(Err(err))) => Err(ConnError::FrameRejected(err)),
Ok(Err(_)) => {
Err(ConnError::Protocol(
"receipt channel closed unexpectedly".into(),
))
}
Err(_) => {
let mut receipts = pending_receipts.lock().await;
receipts.remove(&receipt_id);
Err(ConnError::ReceiptTimeout(receipt_id))
}
}
}
}
impl Connection {
pub const NO_HEARTBEAT: &'static str = "0,0";
pub const DEFAULT_HEARTBEAT: &'static str = "10000,10000";
pub const DEFAULT_DISCONNECT_TIMEOUT: Duration = Duration::from_secs(5);
pub async fn connect(
addr: &str,
login: &str,
passcode: &str,
client_hb: &str,
) -> Result<Self, ConnError> {
Self::connect_with_options(addr, login, passcode, client_hb, ConnectOptions::default())
.await
}
pub async fn connect_with_options(
addr: &str,
login: &str,
passcode: &str,
client_hb: &str,
options: ConnectOptions,
) -> Result<Self, ConnError> {
let (out_tx, mut out_rx) = mpsc::channel::<StompItem>(32);
let (in_tx, in_rx) = mpsc::channel::<Frame>(32);
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let sub_id_counter = Arc::new(AtomicU64::new(1));
let (shutdown_tx, _) = broadcast::channel::<()>(1);
let pending: Arc<Mutex<PendingMap>> = Arc::new(Mutex::new(HashMap::new()));
let pending_clone = pending.clone();
let pending_receipts: Arc<Mutex<PendingReceipts>> = Arc::new(Mutex::new(HashMap::new()));
let pending_receipts_clone = pending_receipts.clone();
let addr = addr.to_string();
let login = login.to_string();
let passcode = passcode.to_string();
let client_hb = client_hb.to_string();
let accept_version = options.accept_version.unwrap_or_else(|| "1.2".to_string());
let host = options.host.unwrap_or_else(|| "/".to_string());
let client_id = options.client_id;
let custom_headers = options.headers;
let heartbeat_notify_tx = options.heartbeat_tx;
let connect_timeout = options.connect_timeout;
let max_frame_size = options.max_frame_size.unwrap_or(DEFAULT_MAX_FRAME_SIZE);
let mut backoff_secs: u64 = 1;
let mut last_err: Option<ConnError> = None;
let attempt = async {
loop {
let stream = match TcpStream::connect(&addr).await {
Ok(s) => s,
Err(e) => {
tracing::warn!(
addr = %addr,
error = %e,
backoff_secs,
"initial connect failed, retrying in {}s",
backoff_secs,
);
last_err = Some(ConnError::Io(e));
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(30);
continue;
}
};
let mut framed =
Framed::new(stream, StompCodec::with_max_frame_size(max_frame_size));
let connect = Self::build_connect_frame(
&accept_version,
&host,
&login,
&passcode,
&client_hb,
&client_id,
&custom_headers,
);
if let Err(e) = framed.send(StompItem::Frame(connect)).await {
tracing::warn!(
addr = %addr,
error = %e,
backoff_secs,
"failed to send CONNECT frame, retrying in {}s",
backoff_secs,
);
last_err = Some(ConnError::Io(e));
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(30);
continue;
}
match Self::await_connected_response(&mut framed).await {
Ok(server_hb) => {
tracing::info!(addr = %addr, "connected to broker");
let (cx, cy) = parse_heartbeat_header(&client_hb);
let (sx, sy) = parse_heartbeat_header(&server_hb);
let (si, ri) = negotiate_heartbeats(cx, cy, sx, sy);
return Ok::<_, ConnError>((framed, si, ri));
}
Err(e @ ConnError::ServerRejected(_)) => {
return Err(e);
}
Err(e) => {
tracing::warn!(
addr = %addr,
error = %e,
backoff_secs,
"handshake failed, retrying in {}s",
backoff_secs,
);
last_err = Some(e);
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(30);
continue;
}
}
}
};
let (framed, send_interval, recv_interval) = match connect_timeout {
Some(timeout) => match tokio::time::timeout(timeout, attempt).await {
Ok(result) => result?,
Err(_elapsed) => {
return Err(last_err.unwrap_or_else(|| {
ConnError::Io(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("connect to {} timed out after {:?}", addr, timeout),
))
}));
}
},
None => attempt.await?,
};
let mut shutdown_sub = shutdown_tx.subscribe();
let subscriptions_clone = subscriptions.clone();
tokio::spawn(async move {
let mut backoff_secs: u64 = 1;
let mut current_framed = Some(framed);
let mut current_send_interval = send_interval;
let mut current_recv_interval = recv_interval;
let mut first_iteration = true;
let mut subscription_errors: HashMap<String, u32> = HashMap::new();
let mut abandoned_sub_ids: std::collections::HashSet<String> =
std::collections::HashSet::new();
const SUBSCRIPTION_ERROR_THRESHOLD: u32 = 3;
loop {
tokio::select! {
biased;
_ = shutdown_sub.recv() => break,
_ = future::ready(()) => {},
}
let framed = if let Some(f) = current_framed.take() {
f
} else {
match TcpStream::connect(&addr).await {
Ok(stream) => {
let mut framed = Framed::new(
stream,
StompCodec::with_max_frame_size(max_frame_size),
);
let connect = Self::build_connect_frame(
&accept_version,
&host,
&login,
&passcode,
&client_hb,
&client_id,
&custom_headers,
);
if let Err(e) = framed.send(StompItem::Frame(connect)).await {
tracing::warn!(
addr = %addr,
error = %e,
backoff_secs,
"reconnect: failed to send CONNECT frame, retrying in {}s",
backoff_secs,
);
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(30);
continue;
}
match Self::await_connected_response(&mut framed).await {
Ok(server_hb) => {
tracing::info!(addr = %addr, "reconnected to broker");
let (cx, cy) = parse_heartbeat_header(&client_hb);
let (sx, sy) = parse_heartbeat_header(&server_hb);
let (si, ri) = negotiate_heartbeats(cx, cy, sx, sy);
current_send_interval = si;
current_recv_interval = ri;
framed
}
Err(e) => {
tracing::warn!(
addr = %addr,
error = %e,
backoff_secs,
"reconnect: handshake failed, retrying in {}s",
backoff_secs,
);
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(30);
continue;
}
}
}
Err(e) => {
tracing::warn!(
addr = %addr,
error = %e,
backoff_secs,
"reconnect: broker unreachable, retrying in {}s",
backoff_secs,
);
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(30);
continue;
}
}
};
let (send_interval, recv_interval) = (current_send_interval, current_recv_interval);
let last_received = Arc::new(AtomicU64::new(current_millis()));
let writer_last_sent = Arc::new(AtomicU64::new(current_millis()));
let (mut sink, mut stream) = framed.split();
let in_tx = in_tx.clone();
let subscriptions = subscriptions_clone.clone();
{
let mut p = pending_clone.lock().await;
p.clear();
}
if first_iteration {
first_iteration = false;
} else {
let subs_snapshot: Vec<ResubEntry> = {
let map = subscriptions.lock().await;
let mut v: Vec<ResubEntry> = Vec::new();
for (dest, vec) in map.iter() {
for entry in vec.iter() {
v.push((
dest.clone(),
entry.id.clone(),
entry.ack.clone(),
entry.headers.clone(),
));
}
}
v
};
for (dest, id, ack, headers) in subs_snapshot {
let mut sf = Frame::new("SUBSCRIBE");
sf = sf
.header("id", &id)
.header("destination", &dest)
.header("ack", &ack);
for (k, v) in headers {
sf = sf.header(&k, &v);
}
if let Err(e) = sink.send(StompItem::Frame(sf)).await {
tracing::warn!(
destination = %dest,
subscription_id = %id,
error = %e,
"failed to resubscribe after reconnect",
);
}
}
}
let mut hb_tick = match send_interval {
Some(d) => tokio::time::interval(d),
None => tokio::time::interval(Duration::from_secs(86400)),
};
let watchdog_half = recv_interval.map(|d| d / 2);
let conn_start = tokio::time::Instant::now();
let mut shutting_down = false;
'conn: loop {
tokio::select! {
_ = shutdown_sub.recv() => { let _ = sink.close().await; shutting_down = true; break 'conn; }
maybe = out_rx.recv() => {
match maybe {
Some(item) => if sink.send(item).await.is_err() { break 'conn } else { writer_last_sent.store(current_millis(), Ordering::SeqCst); }
None => break 'conn,
}
}
item = stream.next() => {
match item {
Some(Ok(StompItem::Heartbeat)) => {
last_received.store(current_millis(), Ordering::SeqCst);
if let Some(ref tx) = heartbeat_notify_tx {
let _ = tx.try_send(());
}
}
Some(Ok(StompItem::Frame(f))) => {
last_received.store(current_millis(), Ordering::SeqCst);
if f.command == "MESSAGE" {
let mut dest_opt: Option<String> = None;
let mut sub_opt: Option<String> = None;
let mut msg_id_opt: Option<String> = None;
for (k, v) in &f.headers {
let kl = k.to_lowercase();
if kl == "destination" {
dest_opt = Some(v.clone());
} else if kl == "subscription" {
sub_opt = Some(v.clone());
} else if kl == "message-id" {
msg_id_opt = Some(v.clone());
}
}
let mut need_pending = false;
if let Some(sub_id) = &sub_opt {
let map = subscriptions.lock().await;
for vec in map.values() {
for entry in vec.iter() {
if &entry.id == sub_id && entry.ack != "auto" {
need_pending = true;
}
}
}
} else if let Some(dest) = &dest_opt {
let map = subscriptions.lock().await;
if let Some(vec) = map.get(dest) {
for entry in vec.iter() {
if entry.ack != "auto" {
need_pending = true;
break;
}
}
}
}
if let Some(msg_id) = msg_id_opt.clone().filter(|_| need_pending) {
if let Some(sub_id) = &sub_opt {
let mut p = pending_clone.lock().await;
let q = p
.entry(sub_id.clone())
.or_insert_with(VecDeque::new);
q.push_back((msg_id.clone(), f.clone()));
} else if let Some(dest) = &dest_opt {
let map = subscriptions.lock().await;
if let Some(vec) = map.get(dest) {
let mut p = pending_clone.lock().await;
for entry in vec.iter() {
let q = p
.entry(entry.id.clone())
.or_insert_with(VecDeque::new);
q.push_back((msg_id.clone(), f.clone()));
}
}
}
}
if let Some(sub_id) = sub_opt {
let mut map = subscriptions.lock().await;
for vec in map.values_mut() {
vec.retain(|entry| {
entry.id != sub_id || deliver_and_keep(entry, &f)
});
}
map.retain(|_, vec| !vec.is_empty());
} else if let Some(dest) = dest_opt {
let mut map = subscriptions.lock().await;
if let Some(vec) = map.get_mut(&dest) {
vec.retain(|entry| deliver_and_keep(entry, &f));
if vec.is_empty() {
map.remove(&dest);
}
}
}
} else if f.command == "RECEIPT" {
if let Some(receipt_id) = f.get_header("receipt-id") {
let mut receipts = pending_receipts_clone.lock().await;
if let Some(sender) = receipts.remove(receipt_id) {
let _ = sender.send(Ok(()));
}
}
continue;
} else if f.command == "ERROR" {
if let Some(receipt_id) = f.get_header("receipt-id") {
let mut receipts =
pending_receipts_clone.lock().await;
if let Some(sender) = receipts.remove(receipt_id) {
let _ = sender
.send(Err(ServerError::from_frame(f.clone())));
}
}
let sub_id = extract_subscription_id_from_error(&f);
if let Some(ref id) = sub_id
&& abandoned_sub_ids.contains(id)
{
continue;
}
let dest = if let Some(d) = extract_destination_from_error(&f)
{
Some(d)
} else if let Some(ref id) = sub_id {
lookup_destination_by_sub_id(id, &subscriptions).await
} else {
None
};
if let Some(dest) = dest {
let count = {
let c = subscription_errors
.entry(dest.clone())
.or_insert(0);
*c += 1;
*c
};
if count >= SUBSCRIPTION_ERROR_THRESHOLD {
let mut map = subscriptions.lock().await;
if map.remove(&dest).is_some() {
if let Some(id) = sub_id {
abandoned_sub_ids.insert(id);
}
let msg = format!(
"Subscription abandoned: {} errors for {}",
count, dest
);
let abandon_frame = Frame::new("ERROR")
.header("message", &msg)
.header("destination", &dest)
.header("x-abandoned", "true");
let _ = in_tx.send(abandon_frame).await;
}
}
}
}
let _ = in_tx.send(f).await;
}
Some(Err(_)) | None => break 'conn,
}
}
_ = hb_tick.tick() => {
if let Some(dur) = send_interval {
let last = writer_last_sent.load(Ordering::SeqCst);
if current_millis().saturating_sub(last) >= dur.as_millis() as u64 {
if sink.send(StompItem::Heartbeat).await.is_err() { break 'conn; }
writer_last_sent.store(current_millis(), Ordering::SeqCst);
}
}
}
_ = async { if let Some(interval) = watchdog_half { tokio::time::sleep(interval).await } else { future::pending::<()>().await } } => {
if let Some(recv_dur) = recv_interval {
let last = last_received.load(Ordering::SeqCst);
if current_millis().saturating_sub(last) > (recv_dur.as_millis() as u64 * 2) {
let _ = sink.close().await; break 'conn;
}
}
}
}
}
if shutting_down || shutdown_sub.try_recv().is_ok() {
break;
}
let stable_duration = conn_start.elapsed();
if stable_duration >= Duration::from_secs(backoff_secs.max(5)) {
backoff_secs = 1;
tracing::info!(
addr = %addr,
stable_secs = stable_duration.as_secs(),
"connection dropped after stable session, reconnecting in 1s",
);
} else {
backoff_secs = (backoff_secs * 2).min(30);
tracing::warn!(
addr = %addr,
stable_secs = stable_duration.as_secs(),
backoff_secs,
"connection dropped quickly, reconnecting in {}s",
backoff_secs,
);
}
tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
}
});
Ok(Connection {
outbound_tx: out_tx,
inbound_rx: Arc::new(Mutex::new(in_rx)),
shutdown_tx,
subscriptions,
sub_id_counter,
pending,
pending_receipts,
disconnect_timeout: options
.disconnect_timeout
.unwrap_or(Self::DEFAULT_DISCONNECT_TIMEOUT),
})
}
fn build_connect_frame(
accept_version: &str,
host: &str,
login: &str,
passcode: &str,
heartbeat: &str,
client_id: &Option<String>,
custom_headers: &[(String, String)],
) -> Frame {
let mut connect = Frame::new("CONNECT")
.header("accept-version", accept_version)
.header("host", host)
.header("login", login)
.header("passcode", passcode)
.header("heart-beat", heartbeat);
if let Some(id) = client_id {
connect = connect.header("client-id", id);
}
let reserved = [
"accept-version",
"host",
"login",
"passcode",
"heart-beat",
"client-id",
];
for (k, v) in custom_headers {
if !reserved.contains(&k.to_lowercase().as_str()) {
connect = connect.header(k, v);
}
}
connect
}
async fn await_connected_response(
framed: &mut Framed<TcpStream, StompCodec>,
) -> Result<String, ConnError> {
loop {
match framed.next().await {
Some(Ok(StompItem::Frame(f))) => {
if f.command == "CONNECTED" {
let server_hb = f.get_header("heart-beat").unwrap_or("0,0").to_string();
return Ok(server_hb);
} else if f.command == "ERROR" {
return Err(ConnError::ServerRejected(ServerError::from_frame(f)));
}
}
Some(Ok(StompItem::Heartbeat)) => {
continue;
}
Some(Err(e)) => {
return Err(ConnError::Io(e));
}
None => {
return Err(ConnError::Protocol(
"connection closed before CONNECTED received".to_string(),
));
}
}
}
}
pub async fn send(&self, destination: &str, body: impl AsRef<str>) -> Result<(), ConnError> {
let frame = Frame::new("SEND")
.header("destination", destination)
.set_body(body.as_ref().as_bytes().to_vec());
self.send_frame(frame).await
}
pub async fn send_frame(&self, frame: Frame) -> Result<(), ConnError> {
self.outbound_tx
.send(StompItem::Frame(frame))
.await
.map_err(|_| ConnError::Protocol("send channel closed".into()))
}
fn generate_receipt_id() -> String {
static RECEIPT_COUNTER: AtomicU64 = AtomicU64::new(1);
format!("rcpt-{}", RECEIPT_COUNTER.fetch_add(1, Ordering::SeqCst))
}
pub async fn send_frame_with_receipt(&self, frame: Frame) -> Result<ReceiptHandle, ConnError> {
let receipt_id = Self::generate_receipt_id();
let (tx, rx) = oneshot::channel();
{
let mut receipts = self.pending_receipts.lock().await;
receipts.insert(receipt_id.clone(), tx);
}
let mut frame = frame;
frame
.headers
.retain(|(key, _)| !key.eq_ignore_ascii_case("receipt"));
let frame_with_receipt = frame.receipt(&receipt_id);
if let Err(err) = self.send_frame(frame_with_receipt).await {
let mut receipts = self.pending_receipts.lock().await;
receipts.remove(&receipt_id);
return Err(err);
}
Ok(ReceiptHandle {
receipt_id,
rx,
pending_receipts: self.pending_receipts.clone(),
})
}
pub async fn send_frame_confirmed(
&self,
frame: Frame,
timeout: Duration,
) -> Result<(), ConnError> {
self.send_frame_with_receipt(frame)
.await?
.wait(timeout)
.await
}
pub async fn subscribe_with_headers(
&self,
destination: &str,
ack: AckMode,
extra_headers: Vec<(String, String)>,
) -> Result<crate::subscription::Subscription, ConnError> {
let id = self
.sub_id_counter
.fetch_add(1, Ordering::SeqCst)
.to_string();
let (tx, rx) = mpsc::channel::<Frame>(16);
{
let mut map = self.subscriptions.lock().await;
map.entry(destination.to_string())
.or_insert_with(Vec::new)
.push(SubscriptionEntry {
id: id.clone(),
sender: tx.clone(),
ack: ack.as_str().to_string(),
headers: extra_headers.clone(),
});
}
let mut f = Frame::new("SUBSCRIBE");
f = f
.header("id", &id)
.header("destination", destination)
.header("ack", ack.as_str());
for (k, v) in &extra_headers {
f = f.header(k, v);
}
self.outbound_tx
.send(StompItem::Frame(f))
.await
.map_err(|_| ConnError::Protocol("send channel closed".into()))?;
Ok(crate::subscription::Subscription::new(
id,
destination.to_string(),
rx,
self.clone(),
))
}
pub async fn subscribe(
&self,
destination: &str,
ack: AckMode,
) -> Result<crate::subscription::Subscription, ConnError> {
self.subscribe_with_headers(destination, ack, Vec::new())
.await
}
pub async fn subscribe_with_options(
&self,
destination: &str,
ack: AckMode,
options: crate::subscription::SubscriptionOptions,
) -> Result<crate::subscription::Subscription, ConnError> {
self.subscribe_with_headers(destination, ack, options.headers)
.await
}
pub async fn unsubscribe(&self, subscription_id: &str) -> Result<(), ConnError> {
let mut found = false;
{
let mut map = self.subscriptions.lock().await;
let mut remove_keys: Vec<String> = Vec::new();
for (dest, vec) in map.iter_mut() {
if let Some(pos) = vec.iter().position(|entry| entry.id == subscription_id) {
vec.remove(pos);
found = true;
}
if vec.is_empty() {
remove_keys.push(dest.clone());
}
}
for k in remove_keys {
map.remove(&k);
}
}
if !found {
return Err(ConnError::Protocol("subscription id not found".into()));
}
let mut f = Frame::new("UNSUBSCRIBE");
f = f.header("id", subscription_id);
self.outbound_tx
.send(StompItem::Frame(f))
.await
.map_err(|_| ConnError::Protocol("send channel closed".into()))?;
Ok(())
}
pub(crate) fn unsubscribe_best_effort(&self, subscription_id: &str) {
if let Ok(mut map) = self.subscriptions.try_lock() {
for vec in map.values_mut() {
vec.retain(|entry| entry.id != subscription_id);
}
map.retain(|_, vec| !vec.is_empty());
}
let f = Frame::new("UNSUBSCRIBE").header("id", subscription_id);
let _ = self.outbound_tx.try_send(StompItem::Frame(f));
}
#[allow(clippy::collapsible_if, clippy::collapsible_else_if)]
pub async fn ack(&self, subscription_id: &str, message_id: &str) -> Result<(), ConnError> {
let mut removed_any = false;
{
let mut p = self.pending.lock().await;
if let Some(queue) = p.get_mut(subscription_id) {
if let Some(pos) = queue.iter().position(|(mid, _)| mid == message_id) {
let mut ack_mode = "client".to_string();
{
let map = self.subscriptions.lock().await;
'outer: for vec in map.values() {
for entry in vec.iter() {
if entry.id == subscription_id {
ack_mode = entry.ack.clone();
break 'outer;
}
}
}
}
if ack_mode == "client" {
for _ in 0..=pos {
queue.pop_front();
removed_any = true;
}
} else if queue.remove(pos).is_some() {
removed_any = true;
}
if queue.is_empty() {
p.remove(subscription_id);
}
}
}
}
let mut f = Frame::new("ACK");
f = f
.header("id", message_id)
.header("subscription", subscription_id);
self.outbound_tx
.send(StompItem::Frame(f))
.await
.map_err(|_| ConnError::Protocol("send channel closed".into()))?;
let _ = removed_any;
Ok(())
}
#[allow(clippy::collapsible_if, clippy::collapsible_else_if)]
pub async fn nack(&self, subscription_id: &str, message_id: &str) -> Result<(), ConnError> {
let mut removed_any = false;
{
let mut p = self.pending.lock().await;
if let Some(queue) = p.get_mut(subscription_id) {
if let Some(pos) = queue.iter().position(|(mid, _)| mid == message_id) {
let mut ack_mode = "client".to_string();
{
let map = self.subscriptions.lock().await;
'outer2: for vec in map.values() {
for entry in vec.iter() {
if entry.id == subscription_id {
ack_mode = entry.ack.clone();
break 'outer2;
}
}
}
}
if ack_mode == "client" {
for _ in 0..=pos {
queue.pop_front();
removed_any = true;
}
} else if queue.remove(pos).is_some() {
removed_any = true;
}
if queue.is_empty() {
p.remove(subscription_id);
}
}
}
}
let mut f = Frame::new("NACK");
f = f
.header("id", message_id)
.header("subscription", subscription_id);
self.outbound_tx
.send(StompItem::Frame(f))
.await
.map_err(|_| ConnError::Protocol("send channel closed".into()))?;
let _ = removed_any;
Ok(())
}
async fn send_transaction_frame(
&self,
command: &str,
transaction_id: &str,
) -> Result<(), ConnError> {
let f = Frame::new(command).header("transaction", transaction_id);
self.outbound_tx
.send(StompItem::Frame(f))
.await
.map_err(|_| ConnError::Protocol("send channel closed".into()))
}
pub async fn begin(&self, transaction_id: &str) -> Result<(), ConnError> {
self.send_transaction_frame("BEGIN", transaction_id).await
}
pub async fn commit(&self, transaction_id: &str) -> Result<(), ConnError> {
self.send_transaction_frame("COMMIT", transaction_id).await
}
pub async fn abort(&self, transaction_id: &str) -> Result<(), ConnError> {
self.send_transaction_frame("ABORT", transaction_id).await
}
pub async fn next_frame(&self) -> Option<ReceivedFrame> {
let mut rx = self.inbound_rx.lock().await;
let frame = rx.recv().await?;
if frame.command == "ERROR" {
Some(ReceivedFrame::Error(ServerError::from_frame(frame)))
} else {
Some(ReceivedFrame::Frame(frame))
}
}
pub async fn close(self) -> Result<(), ConnError> {
let result = match self.send_frame_with_receipt(Frame::new("DISCONNECT")).await {
Ok(handle) => handle.wait(self.disconnect_timeout).await,
Err(err) => Err(err),
};
let _ = self.shutdown_tx.send(());
result
}
}
fn current_millis() -> u64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::thread;
use tokio::sync::mpsc;
fn make_message(
message_id: &str,
subscription: Option<&str>,
destination: Option<&str>,
) -> Frame {
let mut f = Frame::new("MESSAGE");
f = f.header("message-id", message_id);
if let Some(s) = subscription {
f = f.header("subscription", s);
}
if let Some(d) = destination {
f = f.header("destination", d);
}
f
}
fn read_stomp_frame(stream: &mut TcpStream) -> String {
let mut bytes = Vec::new();
let mut byte = [0u8; 1];
loop {
stream.read_exact(&mut byte).unwrap();
bytes.push(byte[0]);
if byte[0] == 0 {
break;
}
}
String::from_utf8(bytes).unwrap()
}
async fn connect_for_test(addr: &str) -> Connection {
Connection::connect_with_options(
addr,
"guest",
"guest",
"0,0",
ConnectOptions::default().disconnect_timeout(Duration::from_millis(50)),
)
.await
.unwrap()
}
fn header_value<'a>(frame: &'a str, name: &str) -> &'a str {
frame
.lines()
.find_map(|line| {
let (key, value) = line.split_once(':')?;
key.eq_ignore_ascii_case(name).then_some(value)
})
.unwrap()
}
fn start_receipt_rejection_server(
message: &'static str,
body: &'static str,
response_delay: Duration,
) -> (SocketAddr, thread::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let handle = thread::spawn(move || {
if let Ok((mut stream, _)) = listener.accept() {
let _connect = read_stomp_frame(&mut stream);
stream
.write_all(b"CONNECTED\nversion:1.2\nheart-beat:0,0\n\n\0")
.unwrap();
stream.flush().unwrap();
let send_frame = read_stomp_frame(&mut stream);
let receipt_id = header_value(&send_frame, "receipt").to_string();
thread::sleep(response_delay);
let error_frame = format!(
"ERROR\nreceipt-id:{}\nmessage:{}\n\n{}\0",
receipt_id, message, body
);
stream.write_all(error_frame.as_bytes()).unwrap();
stream.flush().unwrap();
thread::sleep(Duration::from_millis(100));
}
});
(addr, handle)
}
fn start_receipt_success_server(
response_delay: Duration,
) -> (SocketAddr, thread::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let handle = thread::spawn(move || {
if let Ok((mut stream, _)) = listener.accept() {
let _connect = read_stomp_frame(&mut stream);
stream
.write_all(b"CONNECTED\nversion:1.2\nheart-beat:0,0\n\n\0")
.unwrap();
stream.flush().unwrap();
let send_frame = read_stomp_frame(&mut stream);
let receipt_id = header_value(&send_frame, "receipt").to_string();
thread::sleep(response_delay);
let receipt_frame = format!("RECEIPT\nreceipt-id:{}\n\n\0", receipt_id);
stream.write_all(receipt_frame.as_bytes()).unwrap();
stream.flush().unwrap();
thread::sleep(Duration::from_millis(100));
}
});
(addr, handle)
}
#[tokio::test]
async fn connect_timeout_bounds_an_unreachable_broker() {
let addr = {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
listener.local_addr().unwrap()
};
let start = std::time::Instant::now();
let result = Connection::connect_with_options(
&addr.to_string(),
"guest",
"guest",
"0,0",
ConnectOptions::default().connect_timeout(Duration::from_millis(200)),
)
.await;
let elapsed = start.elapsed();
let err = result.err();
assert!(
matches!(err, Some(ConnError::Io(_))),
"expected an Io error, got {:?}",
err
);
assert!(
elapsed < Duration::from_secs(1),
"connect should have given up near the 200ms bound, took {:?}",
elapsed
);
}
#[tokio::test]
async fn connect_timeout_surfaces_the_last_handshake_error() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let handle = thread::spawn(move || {
if let Ok((mut stream, _)) = listener.accept() {
let _connect = read_stomp_frame(&mut stream);
}
});
let result = Connection::connect_with_options(
&addr.to_string(),
"guest",
"guest",
"0,0",
ConnectOptions::default().connect_timeout(Duration::from_millis(200)),
)
.await;
let err = result.err();
assert!(
matches!(err, Some(ConnError::Protocol(_))),
"expected the handshake Protocol error, got {:?}",
err
);
handle.join().unwrap();
}
#[tokio::test]
async fn connect_timeout_does_not_disturb_a_reachable_broker() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let handle = thread::spawn(move || {
if let Ok((mut stream, _)) = listener.accept() {
let _connect = read_stomp_frame(&mut stream);
stream
.write_all(b"CONNECTED\nversion:1.2\nheart-beat:0,0\n\n\0")
.unwrap();
stream.flush().unwrap();
thread::sleep(Duration::from_millis(100));
}
});
let result = Connection::connect_with_options(
&addr.to_string(),
"guest",
"guest",
"0,0",
ConnectOptions::default()
.connect_timeout(Duration::from_secs(5))
.disconnect_timeout(Duration::from_millis(50)),
)
.await;
assert!(
result.is_ok(),
"expected a connection, got error {:?}",
result.as_ref().err()
);
handle.join().unwrap();
}
#[test]
fn deliver_and_keep_keeps_full_prunes_closed() {
let (tx, _rx) = mpsc::channel::<Frame>(1);
tx.try_send(Frame::new("MESSAGE")).unwrap();
let full = SubscriptionEntry {
id: "1".into(),
sender: tx,
ack: "auto".into(),
headers: vec![],
};
assert!(
deliver_and_keep(&full, &Frame::new("MESSAGE")),
"a full (slow-consumer) channel must be kept, not pruned"
);
let (tx, rx) = mpsc::channel::<Frame>(1);
drop(rx);
let closed = SubscriptionEntry {
id: "2".into(),
sender: tx,
ack: "auto".into(),
headers: vec![],
};
assert!(
!deliver_and_keep(&closed, &Frame::new("MESSAGE")),
"a closed channel (dropped receiver) must be pruned"
);
}
#[tokio::test]
async fn receipt_arriving_before_the_wait_is_not_lost() {
let (addr, server) = start_receipt_success_server(Duration::ZERO);
let addr = addr.to_string();
let conn = connect_for_test(&addr).await;
let handle = conn
.send_frame_with_receipt(
Frame::new("SEND")
.header("destination", "/queue/fast")
.set_body(b"payload".to_vec()),
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
handle
.wait(Duration::from_secs(2))
.await
.expect("RECEIPT delivered before the wait began must still resolve it");
let _ = conn.close().await;
server.join().unwrap();
}
#[tokio::test]
async fn caller_supplied_receipt_header_is_replaced_not_appended() {
let (addr, server) = start_receipt_success_server(Duration::ZERO);
let addr = addr.to_string();
let conn = connect_for_test(&addr).await;
let handle = conn
.send_frame_with_receipt(
Frame::new("SEND")
.header("destination", "/queue/test")
.header("Receipt", "caller-set")
.set_body(b"payload".to_vec()),
)
.await
.unwrap();
assert_ne!(handle.receipt_id(), "caller-set");
handle
.wait(Duration::from_secs(2))
.await
.expect("generated receipt id must be the only one on the wire");
let _ = conn.close().await;
server.join().unwrap();
}
#[tokio::test]
async fn error_with_receipt_id_rejects_confirmed_send_and_stays_inbound() {
let (addr, server) =
start_receipt_rejection_server("publish denied", "not allowed", Duration::ZERO);
let addr = addr.to_string();
let conn = connect_for_test(&addr).await;
let result = conn
.send_frame_confirmed(
Frame::new("SEND")
.header("destination", "/queue/forbidden")
.set_body(b"payload".to_vec()),
Duration::from_secs(2),
)
.await;
let rejected_receipt_id = match result {
Err(ConnError::FrameRejected(err)) => {
assert_eq!(err.message, "publish denied");
assert_eq!(err.body, Some("not allowed".to_string()));
err.receipt_id
.expect("frame rejection should include receipt-id")
}
other => panic!("expected FrameRejected, got {:?}", other),
};
let received = tokio::time::timeout(Duration::from_secs(1), conn.next_frame())
.await
.unwrap();
match received {
Some(ReceivedFrame::Error(err)) => {
assert_eq!(err.message, "publish denied");
assert_eq!(err.body, Some("not allowed".to_string()));
assert_eq!(err.receipt_id, Some(rejected_receipt_id));
}
other => panic!("expected forwarded ERROR frame, got {:?}", other),
}
let _ = conn.close().await;
server.join().unwrap();
}
#[tokio::test]
async fn wait_for_receipt_reports_frame_rejection_and_keeps_error_inbound() {
let (addr, server) = start_receipt_rejection_server(
"receipt rejected",
"permission denied",
Duration::from_millis(100),
);
let addr = addr.to_string();
let conn = connect_for_test(&addr).await;
let handle = conn
.send_frame_with_receipt(
Frame::new("SEND")
.header("destination", "/queue/forbidden")
.set_body(b"payload".to_vec()),
)
.await
.unwrap();
let receipt_id = handle.receipt_id().to_string();
let result = handle.wait(Duration::from_secs(2)).await;
match result {
Err(ConnError::FrameRejected(err)) => {
assert_eq!(err.message, "receipt rejected");
assert_eq!(err.body, Some("permission denied".to_string()));
assert_eq!(err.receipt_id, Some(receipt_id.clone()));
}
other => panic!("expected FrameRejected, got {:?}", other),
}
let received = tokio::time::timeout(Duration::from_secs(1), conn.next_frame())
.await
.unwrap();
match received {
Some(ReceivedFrame::Error(err)) => {
assert_eq!(err.message, "receipt rejected");
assert_eq!(err.body, Some("permission denied".to_string()));
assert_eq!(err.receipt_id, Some(receipt_id));
}
other => panic!("expected forwarded ERROR frame, got {:?}", other),
}
let _ = conn.close().await;
server.join().unwrap();
}
#[tokio::test]
async fn test_cumulative_ack_removes_prefix() {
let (out_tx, mut out_rx) = mpsc::channel::<StompItem>(8);
let (_in_tx, in_rx) = mpsc::channel::<Frame>(8);
let (shutdown_tx, _) = broadcast::channel::<()>(1);
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let pending: Arc<Mutex<PendingMap>> = Arc::new(Mutex::new(HashMap::new()));
let sub_id_counter = Arc::new(AtomicU64::new(1));
let (sub_sender, _sub_rx) = mpsc::channel::<Frame>(4);
{
let mut map = subscriptions.lock().await;
map.insert(
"/queue/x".to_string(),
vec![SubscriptionEntry {
id: "s1".to_string(),
sender: sub_sender,
ack: "client".to_string(),
headers: Vec::new(),
}],
);
}
{
let mut p = pending.lock().await;
let mut q = VecDeque::new();
q.push_back((
"m1".to_string(),
make_message("m1", Some("s1"), Some("/queue/x")),
));
q.push_back((
"m2".to_string(),
make_message("m2", Some("s1"), Some("/queue/x")),
));
q.push_back((
"m3".to_string(),
make_message("m3", Some("s1"), Some("/queue/x")),
));
p.insert("s1".to_string(), q);
}
let conn = Connection {
outbound_tx: out_tx,
inbound_rx: Arc::new(Mutex::new(in_rx)),
shutdown_tx,
subscriptions: subscriptions.clone(),
sub_id_counter,
pending: pending.clone(),
pending_receipts: Arc::new(Mutex::new(HashMap::new())),
disconnect_timeout: Connection::DEFAULT_DISCONNECT_TIMEOUT,
};
conn.ack("s1", "m2").await.expect("ack failed");
{
let p = pending.lock().await;
let q = p.get("s1").expect("missing s1");
assert_eq!(q.len(), 1);
assert_eq!(q.front().unwrap().0, "m3");
}
if let Some(item) = out_rx.recv().await {
match item {
StompItem::Frame(f) => assert_eq!(f.command, "ACK"),
_ => panic!("expected frame"),
}
} else {
panic!("no outbound frame sent")
}
}
#[tokio::test]
async fn test_client_individual_ack_removes_only_one() {
let (out_tx, mut out_rx) = mpsc::channel::<StompItem>(8);
let (_in_tx, in_rx) = mpsc::channel::<Frame>(8);
let (shutdown_tx, _) = broadcast::channel::<()>(1);
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let pending: Arc<Mutex<PendingMap>> = Arc::new(Mutex::new(HashMap::new()));
let sub_id_counter = Arc::new(AtomicU64::new(1));
let (sub_sender, _sub_rx) = mpsc::channel::<Frame>(4);
{
let mut map = subscriptions.lock().await;
map.insert(
"/queue/y".to_string(),
vec![SubscriptionEntry {
id: "s2".to_string(),
sender: sub_sender,
ack: "client-individual".to_string(),
headers: Vec::new(),
}],
);
}
{
let mut p = pending.lock().await;
let mut q = VecDeque::new();
q.push_back((
"a".to_string(),
make_message("a", Some("s2"), Some("/queue/y")),
));
q.push_back((
"b".to_string(),
make_message("b", Some("s2"), Some("/queue/y")),
));
q.push_back((
"c".to_string(),
make_message("c", Some("s2"), Some("/queue/y")),
));
p.insert("s2".to_string(), q);
}
let conn = Connection {
outbound_tx: out_tx,
inbound_rx: Arc::new(Mutex::new(in_rx)),
shutdown_tx,
subscriptions: subscriptions.clone(),
sub_id_counter,
pending: pending.clone(),
pending_receipts: Arc::new(Mutex::new(HashMap::new())),
disconnect_timeout: Connection::DEFAULT_DISCONNECT_TIMEOUT,
};
conn.ack("s2", "b").await.expect("ack failed");
{
let p = pending.lock().await;
let q = p.get("s2").expect("missing s2");
assert_eq!(q.len(), 2);
assert_eq!(q[0].0, "a");
assert_eq!(q[1].0, "c");
}
if let Some(item) = out_rx.recv().await {
match item {
StompItem::Frame(f) => assert_eq!(f.command, "ACK"),
_ => panic!("expected frame"),
}
} else {
panic!("no outbound frame sent")
}
}
#[tokio::test]
async fn test_subscription_receive_delivers_message() {
let (out_tx, _out_rx) = mpsc::channel::<StompItem>(8);
let (_in_tx, in_rx) = mpsc::channel::<Frame>(8);
let (shutdown_tx, _) = broadcast::channel::<()>(1);
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let pending: Arc<Mutex<PendingMap>> = Arc::new(Mutex::new(HashMap::new()));
let sub_id_counter = Arc::new(AtomicU64::new(1));
let conn = Connection {
outbound_tx: out_tx,
inbound_rx: Arc::new(Mutex::new(in_rx)),
shutdown_tx,
subscriptions: subscriptions.clone(),
sub_id_counter,
pending: pending.clone(),
pending_receipts: Arc::new(Mutex::new(HashMap::new())),
disconnect_timeout: Connection::DEFAULT_DISCONNECT_TIMEOUT,
};
let subscription = conn
.subscribe("/queue/test", AckMode::Auto)
.await
.expect("subscribe failed");
{
let map = conn.subscriptions.lock().await;
let vec = map.get("/queue/test").expect("missing subscription vec");
let sender = &vec[0].sender;
let f = make_message("m1", Some(&vec[0].id), Some("/queue/test"));
sender.try_send(f).expect("send to subscription failed");
}
let mut rx = subscription.into_receiver();
if let Some(received) = rx.recv().await {
assert_eq!(received.command, "MESSAGE");
let mut found = false;
for (k, _v) in &received.headers {
if k.to_lowercase() == "message-id" {
found = true;
break;
}
}
assert!(found, "message-id header missing");
} else {
panic!("no message received on subscription")
}
}
#[tokio::test]
async fn test_subscription_ack_removes_pending_and_sends_ack() {
let (out_tx, mut out_rx) = mpsc::channel::<StompItem>(8);
let (_in_tx, in_rx) = mpsc::channel::<Frame>(8);
let (shutdown_tx, _) = broadcast::channel::<()>(1);
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let pending: Arc<Mutex<PendingMap>> = Arc::new(Mutex::new(HashMap::new()));
let sub_id_counter = Arc::new(AtomicU64::new(1));
let conn = Connection {
outbound_tx: out_tx,
inbound_rx: Arc::new(Mutex::new(in_rx)),
shutdown_tx,
subscriptions: subscriptions.clone(),
sub_id_counter,
pending: pending.clone(),
pending_receipts: Arc::new(Mutex::new(HashMap::new())),
disconnect_timeout: Connection::DEFAULT_DISCONNECT_TIMEOUT,
};
let subscription = conn
.subscribe("/queue/ack", AckMode::Client)
.await
.expect("subscribe failed");
let sub_id = subscription.id().to_string();
while out_rx.try_recv().is_ok() {}
{
let mut p = conn.pending.lock().await;
let mut q = VecDeque::new();
q.push_back((
"mid-1".to_string(),
make_message("mid-1", Some(&sub_id), Some("/queue/ack")),
));
p.insert(sub_id.clone(), q);
}
subscription.ack("mid-1").await.expect("ack failed");
{
let p = conn.pending.lock().await;
assert!(p.get(&sub_id).is_none() || p.get(&sub_id).unwrap().is_empty());
}
if let Some(item) = out_rx.recv().await {
match item {
StompItem::Frame(f) => assert_eq!(f.command, "ACK"),
_ => panic!("expected frame"),
}
} else {
panic!("no outbound frame sent")
}
}
fn setup_test_connection() -> (Connection, mpsc::Receiver<StompItem>) {
let (out_tx, out_rx) = mpsc::channel::<StompItem>(8);
let (_in_tx, in_rx) = mpsc::channel::<Frame>(8);
let (shutdown_tx, _) = broadcast::channel::<()>(1);
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let pending: Arc<Mutex<PendingMap>> = Arc::new(Mutex::new(HashMap::new()));
let sub_id_counter = Arc::new(AtomicU64::new(1));
let conn = Connection {
outbound_tx: out_tx,
inbound_rx: Arc::new(Mutex::new(in_rx)),
shutdown_tx,
subscriptions,
sub_id_counter,
pending,
pending_receipts: Arc::new(Mutex::new(HashMap::new())),
disconnect_timeout: Connection::DEFAULT_DISCONNECT_TIMEOUT,
};
(conn, out_rx)
}
fn verify_transaction_frame(frame: Frame, expected_command: &str, expected_tx_id: &str) {
assert_eq!(frame.command, expected_command);
assert!(
frame
.headers
.iter()
.any(|(k, v)| k == "transaction" && v == expected_tx_id),
"transaction header with id '{}' not found",
expected_tx_id
);
}
#[tokio::test]
async fn test_begin_transaction_sends_frame() {
let (conn, mut out_rx) = setup_test_connection();
conn.begin("tx1").await.expect("begin failed");
if let Some(StompItem::Frame(f)) = out_rx.recv().await {
verify_transaction_frame(f, "BEGIN", "tx1");
} else {
panic!("no outbound frame sent")
}
}
#[tokio::test]
async fn test_commit_transaction_sends_frame() {
let (conn, mut out_rx) = setup_test_connection();
conn.commit("tx1").await.expect("commit failed");
if let Some(StompItem::Frame(f)) = out_rx.recv().await {
verify_transaction_frame(f, "COMMIT", "tx1");
} else {
panic!("no outbound frame sent")
}
}
#[tokio::test]
async fn test_abort_transaction_sends_frame() {
let (conn, mut out_rx) = setup_test_connection();
conn.abort("tx1").await.expect("abort failed");
if let Some(StompItem::Frame(f)) = out_rx.recv().await {
verify_transaction_frame(f, "ABORT", "tx1");
} else {
panic!("no outbound frame sent")
}
}
#[tokio::test]
async fn test_send_convenience_produces_correct_frame() {
let (conn, mut out_rx) = setup_test_connection();
conn.send("/queue/events", "hello world")
.await
.expect("send failed");
if let Some(StompItem::Frame(f)) = out_rx.recv().await {
assert_eq!(f.command, "SEND");
assert_eq!(f.get_header("destination"), Some("/queue/events"));
assert_eq!(f.body, b"hello world");
} else {
panic!("no outbound frame sent")
}
}
#[test]
fn test_extract_destination_from_error_header() {
let frame = Frame::new("ERROR")
.header("message", "AMQ339016: Error creating STOMP subscription")
.header("destination", "/topic/test.restricted");
let dest = extract_destination_from_error(&frame);
assert_eq!(dest, Some("/topic/test.restricted".to_string()));
}
#[test]
fn test_extract_destination_from_error_message() {
let frame = Frame::new("ERROR").header(
"message",
"AMQ339016: Error creating subscription for /topic/test.restricted",
);
let dest = extract_destination_from_error(&frame);
assert_eq!(dest, Some("/topic/test.restricted".to_string()));
}
#[test]
fn test_extract_destination_from_error_body() {
let frame = Frame::new("ERROR")
.header("message", "AMQ339016: Error creating subscription")
.set_body(b"User guest is not authorized for /queue/orders".to_vec());
let dest = extract_destination_from_error(&frame);
assert_eq!(dest, Some("/queue/orders".to_string()));
}
#[test]
fn test_extract_destination_from_error_none() {
let frame = Frame::new("ERROR").header("message", "Generic error without destination info");
let dest = extract_destination_from_error(&frame);
assert_eq!(dest, None);
}
#[test]
fn test_extract_destination_from_error_with_trailing_punct() {
let frame = Frame::new("ERROR").header(
"message",
"Error for /topic/events, please check permissions",
);
let dest = extract_destination_from_error(&frame);
assert_eq!(dest, Some("/topic/events".to_string()));
}
#[test]
fn test_extract_subscription_id_from_error_artemis_format() {
let frame =
Frame::new("ERROR").header("message", "AMQ339016 Error creating subscription 1");
let sub_id = extract_subscription_id_from_error(&frame);
assert_eq!(sub_id, Some("1".to_string()));
}
#[test]
fn test_extract_subscription_id_from_error_numeric() {
let frame = Frame::new("ERROR").header("message", "Error for subscription 123 on server");
let sub_id = extract_subscription_id_from_error(&frame);
assert_eq!(sub_id, Some("123".to_string()));
}
#[test]
fn test_extract_subscription_id_from_error_none() {
let frame = Frame::new("ERROR").header("message", "Generic connection error");
let sub_id = extract_subscription_id_from_error(&frame);
assert_eq!(sub_id, None);
}
#[tokio::test]
async fn test_lookup_destination_by_sub_id() {
let subscriptions: Arc<Mutex<Subscriptions>> = Arc::new(Mutex::new(HashMap::new()));
let (sender, _rx) = mpsc::channel::<Frame>(4);
{
let mut map = subscriptions.lock().await;
map.insert(
"/topic/test.restricted".to_string(),
vec![SubscriptionEntry {
id: "1".to_string(),
sender,
ack: "auto".to_string(),
headers: Vec::new(),
}],
);
}
let dest = lookup_destination_by_sub_id("1", &subscriptions).await;
assert_eq!(dest, Some("/topic/test.restricted".to_string()));
let dest = lookup_destination_by_sub_id("999", &subscriptions).await;
assert_eq!(dest, None);
}
}