arminius_lib_nats/
requester.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Request-reply over NATS: typed command dispatch for the server side, and
6//! fire-and-wait request for the client side.
7
8use std::{marker::PhantomData, ops::Deref, time::Duration};
9
10use arminius_lib_message_bus::{Command, CommandError, Error, ReceivedCommand, Requester};
11use async_nats::jetstream;
12use futures::StreamExt as _;
13use serde::Serialize;
14
15use crate::{
16	NatsError, client::NatsClient, headers::set_reply_subject, opentelemetry::span_to_nats_headers,
17};
18
19/// Timeout before dropping a NATS request.
20const REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
21
22/// Internal server-side reply handle, used by [`NatsReceivedCommand`].
23struct ReplyHandle<R> {
24	/// Trace context propagated from the incoming request.
25	context: opentelemetry::Context,
26	/// Subject the reply is published to.
27	return_subject: String,
28	/// The original request message, acknowledged before the reply goes out.
29	message: jetstream::Message,
30	/// NatsClient used to publish the reply.
31	client: async_nats::Client,
32	/// Marker binding the reply type `R` without holding a value.
33	_response: PhantomData<fn() -> R>,
34}
35
36impl<R> std::fmt::Debug for ReplyHandle<R> {
37	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38		f.debug_struct("ReplyHandle")
39			.field("return_subject", &self.return_subject)
40			.finish_non_exhaustive()
41	}
42}
43
44impl<R: Serialize> ReplyHandle<R> {
45	/// Acknowledges the original request message.
46	pub(crate) async fn ack(&self) -> Result<(), NatsError> {
47		self.message.ack().await.map_err(NatsError::NatsError)
48	}
49
50	/// Acknowledges the original request message, then sends the reply.
51	pub(crate) async fn reply(&self, response: R) -> Result<(), NatsError> {
52		self.ack().await?;
53		let bytes = serde_json::to_vec(&response)?;
54		self.client.publish(self.return_subject.clone(), bytes.into()).await?;
55		Ok(())
56	}
57
58	/// Provides access to the trace context propagated from the incoming
59	/// request.
60	pub(crate) fn context(&self) -> &opentelemetry::Context {
61		&self.context
62	}
63}
64
65/// Server-side wrapper around a received command, implementing
66/// [`ReceivedCommand<C>`].
67pub struct NatsReceivedCommand<C: Command> {
68	/// The deserialized command payload.
69	command: C,
70	/// Handle used to publish the reply and acknowledge the message.
71	handle: ReplyHandle<C::Reply>,
72}
73
74impl<C> std::fmt::Debug for NatsReceivedCommand<C>
75where
76	C: Command + std::fmt::Debug,
77{
78	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79		f.debug_struct("NatsReceivedCommand")
80			.field("command", &self.command)
81			.field("handle", &self.handle)
82			.finish()
83	}
84}
85
86impl<C: Command> Deref for NatsReceivedCommand<C> {
87	type Target = C;
88	fn deref(&self) -> &C {
89		&self.command
90	}
91}
92
93impl<C: Command> ReceivedCommand<C> for NatsReceivedCommand<C> {
94	async fn ack(&self) -> Result<(), Error> {
95		self.handle.ack().await.map_err(|error| Error::SendFailed { error: error.to_string() })
96	}
97
98	async fn reply(&self, response: C::Reply) -> Result<(), Error> {
99		self.handle
100			.reply(response)
101			.await
102			.map_err(|error| Error::SendFailed { error: error.to_string() })
103	}
104
105	fn context(&self) -> &opentelemetry::Context {
106		self.handle.context()
107	}
108}
109
110impl<C: Command> NatsReceivedCommand<C> {
111	/// Constructs a `NatsReceivedCommand<C>` from its raw parts. Called by
112	/// `responder.rs`.
113	pub(crate) fn new(
114		command: C,
115		context: opentelemetry::Context,
116		return_subject: String,
117		message: jetstream::Message,
118		client: async_nats::Client,
119	) -> Self {
120		NatsReceivedCommand {
121			command,
122			handle: ReplyHandle {
123				context,
124				return_subject,
125				message,
126				client,
127				_response: PhantomData,
128			},
129		}
130	}
131}
132
133/// Hidden helpers used by `responder.rs`.
134#[doc(hidden)]
135pub mod __private {
136	use async_nats::jetstream;
137
138	/// Extracts the OTel context from NATS message headers.
139	pub fn extract_otel_context(message: &jetstream::Message) -> opentelemetry::Context {
140		crate::opentelemetry::extract_otel_context(message)
141	}
142}
143
144impl Requester for NatsClient {
145	/// Sends `cmd` over JetStream and awaits a typed reply on a core-NATS
146	/// inbox.
147	///
148	/// Encodes the command payload with [`Command::into_payload`], ensures the
149	/// command's JetStream stream exists, then publishes to the subject from
150	/// [`Command::subject`] with the reply subject carried in the header.
151	/// Ensuring the stream means a command sent before its responder is up is
152	/// durably captured rather than dropped.
153	async fn request<C: Command + Send>(&self, cmd: C) -> Result<C::Reply, Error> {
154		let subject = cmd.subject().into_owned();
155		let return_subject = self.client.new_inbox();
156		let payload = cmd
157			.into_payload()
158			.map_err(|error: CommandError| Error::SendFailed { error: error.to_string() })?;
159
160		let mut subscription = self
161			.client
162			.subscribe(return_subject.clone())
163			.await
164			.map_err(|error| Error::SendFailed { error: error.to_string() })?;
165
166		// Ensure the command stream exists so the request is durably captured
167		// even if the responder has not started yet (mirrors
168		// `Publisher::publish`).
169		self.ensure_command_stream(C::SUBJECT_FILTER)
170			.await
171			.map_err(|error| Error::SendFailed { error: error.to_string() })?;
172
173		let mut headers = span_to_nats_headers();
174		set_reply_subject(&mut headers, &return_subject);
175		self.jetstream
176			.publish_with_headers(subject.clone(), headers, payload.into())
177			.await
178			.map_err(|error| Error::SendFailed { error: error.to_string() })?
179			.await
180			.map_err(|error| Error::SendFailed { error: error.to_string() })?;
181
182		if let Err(error) = subscription.unsubscribe_after(1).await {
183			tracing::error!(
184				return_subject,
185				%error,
186				"Failed to unsubscribe from command reply subject"
187			);
188		}
189
190		let message = tokio::time::timeout(REQUEST_TIMEOUT, subscription.next())
191			.await
192			.map_err(|_| Error::ReceptionFailed {
193				subject: subject.clone(),
194				error: "request timed out".into(),
195			})?
196			.ok_or_else(|| Error::ReceptionFailed {
197				subject: subject.clone(),
198				error: "reply stream prematurely closed".into(),
199			})?;
200
201		serde_json::from_slice::<C::Reply>(&message.payload)
202			.map_err(|error| Error::ReceptionFailed { subject, error: error.to_string() })
203	}
204}
arminius_lib_nats/requester.rs

arminius_lib_nats/
requester.rs
