use crate::context::{Context, Receiver};
use crate::error::{Mecha10Error, Result};
use crate::messages::Message;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::future::Future;
use std::time::Duration;
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Request<T> {
pub correlation_id: String,
pub payload: T,
}
impl<T: Message> Message for Request<T> {}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Response<T> {
pub correlation_id: String,
pub payload: RpcResult<T>,
}
impl<T: Message> Message for Response<T> {}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RpcResult<T> {
Ok(T),
Err(RpcError),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RpcError {
pub message: String,
pub code: Option<String>,
}
pub struct RequestBuilder<'a, Req, Resp> {
ctx: &'a Context,
topic: String,
request: Request<Req>,
timeout: Duration,
_phantom: std::marker::PhantomData<Resp>,
}
impl<'a, Req, Resp> RequestBuilder<'a, Req, Resp>
where
Req: Message + Serialize + Clone,
Resp: Message + DeserializeOwned + Send + 'static,
{
pub fn timeout(mut self, duration: Duration) -> Self {
self.timeout = duration;
self
}
pub async fn execute(self) -> Result<Resp> {
let request_topic = format!("{}/request", self.topic);
let response_topic = format!("{}/response", self.topic);
let correlation_id = self.request.correlation_id.clone();
let mut responses = self.ctx.subscribe_raw::<Response<Resp>>(&response_topic).await?;
self.ctx.publish_raw(&request_topic, &self.request).await?;
let deadline = tokio::time::Instant::now() + self.timeout;
loop {
match tokio::time::timeout_at(deadline, responses.recv()).await {
Ok(Some(response)) => {
if response.correlation_id == correlation_id {
match response.payload {
RpcResult::Ok(payload) => return Ok(payload),
RpcResult::Err(err) => {
return Err(Mecha10Error::MessagingError {
message: format!("RPC error: {}", err.message),
suggestion: "Check the handler implementation for the error cause".to_string(),
});
}
}
}
continue;
}
Ok(None) => {
return Err(Mecha10Error::MessagingError {
message: "Response channel closed".to_string(),
suggestion: "Check if responder node is running".to_string(),
});
}
Err(_) => {
return Err(Mecha10Error::MessagingError {
message: format!("Request to {} timed out after {:?}", self.topic, self.timeout),
suggestion: "Check if responder node is running and responsive".to_string(),
});
}
}
}
}
}
pub trait RpcExt {
fn request<Req, Resp>(&self, topic: &str, request: &Req) -> RequestBuilder<'_, Req, Resp>
where
Req: Message + Serialize + Clone,
Resp: Message + DeserializeOwned + Send + 'static;
fn respond<Req, Resp, F, Fut>(&self, topic: &str, handler: F) -> impl Future<Output = Result<()>>
where
Req: Message + DeserializeOwned + Send + 'static,
Resp: Message + Serialize + Send + 'static,
F: Fn(Req) -> Fut + Send + 'static,
Fut: Future<Output = Result<Resp>> + Send;
fn subscribe_raw<T>(&self, topic: &str) -> impl Future<Output = Result<Receiver<T>>>
where
T: Message + DeserializeOwned + Send + 'static;
fn publish_raw<T>(&self, topic: &str, message: &T) -> impl Future<Output = Result<()>>
where
T: Message + Serialize;
}
impl RpcExt for Context {
fn request<Req, Resp>(&self, topic: &str, request: &Req) -> RequestBuilder<'_, Req, Resp>
where
Req: Message + Serialize + Clone,
Resp: Message + DeserializeOwned + Send + 'static,
{
RequestBuilder {
ctx: self,
topic: topic.to_string(),
request: Request {
correlation_id: Uuid::new_v4().to_string(),
payload: request.clone(),
},
timeout: Duration::from_secs(10), _phantom: std::marker::PhantomData,
}
}
async fn respond<Req, Resp, F, Fut>(&self, topic: &str, handler: F) -> Result<()>
where
Req: Message + DeserializeOwned + Send + 'static,
Resp: Message + Serialize + Send + 'static,
F: Fn(Req) -> Fut + Send + 'static,
Fut: Future<Output = Result<Resp>> + Send,
{
let request_topic = format!("{}/request", topic);
let response_topic = format!("{}/response", topic);
let mut requests = self.subscribe_raw::<Request<Req>>(&request_topic).await?;
let ctx = self.clone();
tokio::spawn(async move {
while let Some(request) = requests.recv().await {
let correlation_id = request.correlation_id.clone();
let result = match handler(request.payload).await {
Ok(response_payload) => RpcResult::Ok(response_payload),
Err(e) => {
tracing::error!(
correlation_id = %correlation_id,
error = %e,
"RPC handler failed"
);
RpcResult::Err(RpcError {
message: e.to_string(),
code: None,
})
}
};
let response = Response {
correlation_id: correlation_id.clone(),
payload: result,
};
if let Err(e) = ctx.publish_raw(&response_topic, &response).await {
tracing::error!(
correlation_id = %correlation_id,
topic = %response_topic,
error = %e,
"Failed to publish RPC response"
);
}
}
});
Ok(())
}
async fn subscribe_raw<T>(&self, topic: &str) -> Result<Receiver<T>>
where
T: Message + DeserializeOwned + Send + 'static,
{
Context::subscribe_raw(self, topic).await
}
async fn publish_raw<T>(&self, topic: &str, message: &T) -> Result<()>
where
T: Message + Serialize,
{
Context::publish_raw(self, topic, message).await
}
}