use std::time::Duration;
use tokio::sync::broadcast;
use crate::network::NetworkProvider;
use crate::node::{NamespacedMessage, Node, NodeError};
#[derive(Debug, thiserror::Error)]
pub enum RequestError {
#[error("send failed: {0}")]
Send(NodeError),
#[error("timed out waiting for reply")]
Timeout,
#[error("subscription channel closed")]
ChannelClosed,
}
pub async fn send_and_wait<N, F, R>(
node: &Node<N>,
peer_id: &str,
namespace: &str,
msg_type: &str,
payload: &serde_json::Value,
timeout: Duration,
filter: F,
) -> Result<R, RequestError>
where
N: NetworkProvider + 'static,
F: Fn(&NamespacedMessage) -> Option<R>,
{
let mut rx = node.subscribe(namespace);
node.send_typed(peer_id, namespace, msg_type, payload)
.await
.map_err(RequestError::Send)?;
let deadline = tokio::time::Instant::now() + timeout;
loop {
let remaining = deadline.duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return Err(RequestError::Timeout);
}
match tokio::time::timeout(remaining, rx.recv()).await {
Ok(Ok(msg)) => {
if let Some(result) = filter(&msg) {
return Ok(result);
}
}
Ok(Err(broadcast::error::RecvError::Lagged(n))) => {
tracing::warn!(missed = n, "request_reply: receiver lagged");
}
Ok(Err(broadcast::error::RecvError::Closed)) => {
return Err(RequestError::ChannelClosed);
}
Err(_elapsed) => {
return Err(RequestError::Timeout);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use tokio::time::Duration;
use crate::envelope::codec::JsonCodec;
use crate::envelope::EnvelopeCodec;
use crate::network::*;
use crate::session::PeerRegistry;
use crate::transport::websocket::WebSocketTransport;
use crate::transport::WsConfig;
struct MockNetworkProvider {
identity: NodeIdentity,
local_addr: PeerAddr,
peer_event_tx: tokio::sync::broadcast::Sender<NetworkPeerEvent>,
mock_peers: Arc<tokio::sync::RwLock<Vec<NetworkPeer>>>,
}
impl MockNetworkProvider {
fn new(id: &str) -> Self {
let (peer_event_tx, _) = tokio::sync::broadcast::channel(64);
Self {
identity: NodeIdentity {
id: id.to_string(),
hostname: format!("truffle-test-{id}"),
name: format!("Test Node {id}"),
dns_name: None,
ip: Some("127.0.0.1".parse().unwrap()),
},
local_addr: PeerAddr {
ip: Some("127.0.0.1".parse().unwrap()),
hostname: format!("truffle-test-{id}"),
dns_name: None,
},
peer_event_tx,
mock_peers: Arc::new(tokio::sync::RwLock::new(Vec::new())),
}
}
fn event_sender(&self) -> tokio::sync::broadcast::Sender<NetworkPeerEvent> {
self.peer_event_tx.clone()
}
}
impl NetworkProvider for MockNetworkProvider {
fn local_identity(&self) -> NodeIdentity { self.identity.clone() }
fn local_addr(&self) -> PeerAddr { self.local_addr.clone() }
fn peer_events(&self) -> tokio::sync::broadcast::Receiver<NetworkPeerEvent> {
self.peer_event_tx.subscribe()
}
async fn start(&mut self) -> Result<(), NetworkError> { Ok(()) }
async fn stop(&mut self) -> Result<(), NetworkError> { Ok(()) }
async fn peers(&self) -> Vec<NetworkPeer> { self.mock_peers.read().await.clone() }
async fn dial_tcp(&self, _addr: &str, _port: u16) -> Result<tokio::net::TcpStream, NetworkError> {
Err(NetworkError::DialFailed("mock".into()))
}
async fn listen_tcp(&self, _port: u16) -> Result<NetworkTcpListener, NetworkError> {
Err(NetworkError::ListenFailed("mock".into()))
}
async fn unlisten_tcp(&self, _port: u16) -> Result<(), NetworkError> { Ok(()) }
async fn bind_udp(&self, _port: u16) -> Result<NetworkUdpSocket, NetworkError> {
Err(NetworkError::NotRunning)
}
async fn ping(&self, _addr: &str) -> Result<PingResult, NetworkError> {
Ok(PingResult { latency: Duration::from_millis(1), connection: "direct".to_string(), peer_addr: None })
}
async fn health(&self) -> HealthInfo {
HealthInfo { state: "running".to_string(), healthy: true, ..Default::default() }
}
}
fn ws_config(port: u16) -> WsConfig {
WsConfig {
port,
ping_interval: Duration::from_secs(300),
pong_timeout: Duration::from_secs(300),
..Default::default()
}
}
async fn random_port() -> u16 {
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
l.local_addr().unwrap().port()
}
fn make_loopback_peer(id: &str) -> NetworkPeer {
NetworkPeer {
id: id.to_string(),
hostname: format!("truffle-test-{id}"),
ip: "127.0.0.1".parse().unwrap(),
online: true,
cur_addr: Some("127.0.0.1:41641".to_string()),
relay: None,
os: Some("linux".to_string()),
last_seen: Some("2026-03-25T12:00:00Z".to_string()),
key_expiry: None,
dns_name: None,
}
}
async fn make_test_node(
id: &str,
ws_port: u16,
) -> (
Node<MockNetworkProvider>,
tokio::sync::broadcast::Sender<NetworkPeerEvent>,
) {
let provider = MockNetworkProvider::new(id);
let event_tx = provider.event_sender();
let network = Arc::new(provider);
let ws_transport = Arc::new(WebSocketTransport::new(network.clone(), ws_config(ws_port)));
let session = Arc::new(PeerRegistry::new(network.clone(), ws_transport));
session.start().await;
let codec: Arc<dyn EnvelopeCodec> = Arc::new(JsonCodec);
let node = Node::from_parts(network, session, codec);
(node, event_tx)
}
#[tokio::test]
async fn test_send_and_wait_timeout() {
let ws_port = random_port().await;
let (node, event_tx) = make_test_node("node-a", ws_port).await;
let _ = event_tx.send(NetworkPeerEvent::Joined(make_loopback_peer("peer-b")));
tokio::time::sleep(Duration::from_millis(50)).await;
let result = send_and_wait(
&node,
"peer-b",
"test-ns",
"ping",
&serde_json::json!({}),
Duration::from_millis(100),
|_msg| -> Option<()> { None },
)
.await;
assert!(matches!(result, Err(RequestError::Timeout) | Err(RequestError::Send(_))));
}
#[tokio::test]
async fn test_send_and_wait_send_failure_unknown_peer() {
let ws_port = random_port().await;
let (node, _event_tx) = make_test_node("node-a", ws_port).await;
let result = send_and_wait(
&node,
"nonexistent",
"test-ns",
"ping",
&serde_json::json!({}),
Duration::from_secs(1),
|_msg| -> Option<()> { Some(()) },
)
.await;
assert!(matches!(result, Err(RequestError::Send(_))));
}
#[tokio::test]
async fn test_send_and_wait_filter_matches() {
let msg = NamespacedMessage {
from: "peer-b".to_string(),
namespace: "test-ns".to_string(),
msg_type: "response".to_string(),
payload: serde_json::json!({"value": 42}),
timestamp: None,
};
let filter = |m: &NamespacedMessage| -> Option<i64> {
if m.msg_type == "response" {
m.payload.get("value")?.as_i64()
} else {
None
}
};
assert_eq!(filter(&msg), Some(42));
let non_matching = NamespacedMessage {
from: "peer-b".to_string(),
namespace: "test-ns".to_string(),
msg_type: "other".to_string(),
payload: serde_json::json!({}),
timestamp: None,
};
assert_eq!(filter(&non_matching), None);
}
}