Skip to main content

arminius_lib_nats/client/
request.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Helps to send data via a request and reply pattern to other services.
6use 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
15/// Timeout before dropping a NATS request
16const NATS_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
17
18/// The body of a request to be sent over NATS.
19#[derive(Debug, Deserialize, Serialize)]
20pub struct NatsRequestBody<T> {
21	/// The actual request to serialize
22	pub request: T,
23	/// The subject over which a response will be sent
24	pub return_subject: String,
25}
26
27impl<T: Serialize> NatsRequestBody<T> {
28	/// Create a new request with a random return subject
29	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	/// Make a request over JetStream, and create a subscriber that'll wait for a response over NATS
38	#[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	/// Given a subscriber created using [`NatsClient::make_nats_request()`], get the response.
63	#[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}