arminius_lib_nats/client/
request.rs1use std::time::Duration;
7
8use async_nats::{Subscriber, ToSubject, jetstream};
9use futures::StreamExt as _;
10use rand::RngExt as _;
11use serde::{Deserialize, Serialize, de::DeserializeOwned};
12
13use crate::{NatsClient, NatsError, opentelemetry::span_to_nats_headers};
14
15const NATS_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
17
18#[derive(Debug, Deserialize, Serialize)]
20pub struct NatsRequestBody<T> {
21 pub request: T,
23 pub return_subject: String,
25}
26
27impl<T: Serialize> NatsRequestBody<T> {
28 pub fn new(request: T) -> Self {
30 let number: u32 = rand::rng().random();
31 let return_subject = format!("reply.{number}");
32 NatsRequestBody { request, return_subject }
33 }
34}
35
36impl NatsClient {
37 #[tracing::instrument(skip(self, subject, request))]
39 pub async fn make_nats_request<T: Serialize>(
40 &self,
41 subject: impl ToSubject,
42 request: T,
43 ) -> Result<Subscriber, NatsError> {
44 tracing::trace!(subject = subject.to_subject().as_str(), "Sending NATS request");
45 let request = NatsRequestBody::new(request);
46 let binary = serde_json::to_string(&request)?.into_bytes();
47 let mut subscription = self.client.subscribe(request.return_subject.clone()).await?;
48 let jetstream = jetstream::new(self.client.clone());
49 let headers = span_to_nats_headers();
50 jetstream.publish_with_headers(subject, headers, binary.into()).await?.await?;
51 tracing::trace!("Successfully sent message");
52 tracing::trace!(
53 return_subject = request.return_subject,
54 "Waiting for response on NATS request"
55 );
56 if let Err(error) = subscription.unsubscribe_after(1).await {
57 tracing::error!(%request.return_subject, %error, "Failed to unsubscribe from reply subject" );
58 }
59 Ok(subscription)
60 }
61
62 #[tracing::instrument]
64 pub async fn await_nats_response<U: DeserializeOwned, E: DeserializeOwned>(
65 mut subscription: Subscriber,
66 ) -> Result<Result<U, E>, NatsError> {
67 let message = tokio::time::timeout(NATS_REQUEST_TIMEOUT, subscription.next())
68 .await
69 .map_err(|_| NatsError::Timeout)?
70 .ok_or(NatsError::ReplyStreamPrematurelyClosed)?;
71 let message: Result<U, E> = serde_json::from_slice(&message.payload)?;
72 Ok(message)
73 }
74}