1use crate::deadline::CallDeadline;
5use crate::dispatch::{decode, Dispatcher};
6use async_nats::client::{PublishError, PublishErrorKind};
7use async_nats::{Client, Message, Subject, Subscriber};
8use cufflink_nats::{parse_server_list, redact_urls, NatsAuth};
9use cufflink_types::ext::{call_subject, ExtReply, QUEUE_GROUP};
10use cufflink_types::host::ext::CallErrorKind;
11use cufflink_types::slug::{is_valid_tenant_slug, TENANT_SLUG_MAX};
12use eyre::WrapErr;
13use futures_util::stream::{select_all, SelectAll};
14use futures_util::StreamExt;
15use std::fmt;
16use std::future::Future;
17use std::sync::Arc;
18use std::time::{Duration, Instant};
19use tokio_util::task::TaskTracker;
20
21const DEFAULT_URL: &str = "nats://localhost:4222";
22const STOP_TIMEOUT: Duration = Duration::from_secs(10);
25
26#[derive(Clone)]
28pub struct NatsProvider {
29 tenant: String,
30 url: String,
31 auth: NatsAuth,
32}
33
34impl fmt::Debug for NatsProvider {
35 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36 f.debug_struct("NatsProvider")
37 .field("tenant", &self.tenant)
38 .field("url", &redact_urls(&self.url))
39 .field("auth", &self.auth)
40 .finish()
41 }
42}
43
44impl NatsProvider {
45 pub fn new(
48 tenant: impl Into<String>,
49 url: impl Into<String>,
50 auth: NatsAuth,
51 ) -> eyre::Result<Self> {
52 let (tenant, url) = (tenant.into(), url.into());
53 eyre::ensure!(
54 is_valid_tenant_slug(&tenant),
55 "'{tenant}' is not a tenant slug (1 to {TENANT_SLUG_MAX} of a-z, 0-9 and '-')"
56 );
57 parse_server_list(&url)?;
58 Ok(Self { tenant, url, auth })
59 }
60
61 pub fn from_env() -> eyre::Result<Self> {
64 Self::from_vars(|name| std::env::var(name).ok().filter(|v| !v.is_empty()))
65 }
66
67 fn from_vars(lookup: impl Fn(&str) -> Option<String>) -> eyre::Result<Self> {
68 let tenant = lookup("CUFFLINK_TENANT").ok_or_else(|| {
69 eyre::eyre!("CUFFLINK_TENANT must name the tenant whose calls to serve")
70 })?;
71 let url = lookup("NATS_URL").unwrap_or_else(|| DEFAULT_URL.into());
72 Self::new(tenant, url, NatsAuth::from_vars(&lookup)?)
73 }
74
75 pub fn tenant(&self) -> &str {
76 &self.tenant
77 }
78}
79
80pub struct NatsServer {
82 client: Client,
83 tenant: Arc<str>,
84 calls: SelectAll<Subscriber>,
85 dispatcher: Arc<Dispatcher>,
86}
87
88impl NatsServer {
89 pub(crate) async fn bind(
90 service: &str,
91 provider: NatsProvider,
92 dispatcher: Dispatcher,
93 ) -> eyre::Result<Self> {
94 let client = cufflink_nats::connect(&provider.url, service, &provider.auth).await?;
95 let subscribed = dispatcher.names().map(|name| {
96 client.queue_subscribe(call_subject(&provider.tenant, name), QUEUE_GROUP.into())
97 });
98 let calls = select_all(
99 futures_util::future::try_join_all(subscribed)
100 .await
101 .wrap_err("cannot subscribe to the calls")?,
102 );
103 client
104 .flush()
105 .await
106 .wrap_err("cannot send the subscriptions to NATS")?;
107 Ok(Self {
108 client,
109 tenant: provider.tenant.into(),
110 calls,
111 dispatcher: Arc::new(dispatcher),
112 })
113 }
114
115 pub async fn serve(self, shutdown: impl Future<Output = ()>) -> eyre::Result<()> {
121 let Self {
122 client,
123 tenant,
124 mut calls,
125 dispatcher,
126 } = self;
127 tracing::info!(
128 %tenant,
129 extensions = ?dispatcher.ready().extensions,
130 "serving"
131 );
132 let tasks = TaskTracker::new();
133 let answer_in = |message: Message| {
134 tasks.spawn(answer(
135 Arc::clone(&dispatcher),
136 client.clone(),
137 Arc::clone(&tenant),
138 message,
139 ));
140 };
141 let mut shutdown = std::pin::pin!(shutdown);
142 let ended = loop {
143 tokio::select! {
144 () = &mut shutdown => break false,
145 message = calls.next() => match message {
146 Some(message) => answer_in(message),
147 None => break true,
148 },
149 }
150 };
151 for subscriber in calls.iter_mut() {
152 subscriber.drain().await.wrap_err("cannot unsubscribe")?;
153 }
154 let drained = tokio::time::timeout(
155 STOP_TIMEOUT,
156 calls.for_each(|message| {
157 answer_in(message);
158 std::future::ready(())
159 }),
160 )
161 .await;
162 tasks.close();
163 tasks.wait().await;
164 drained.wrap_err("the calls already delivered were not answered in time")?;
165 tokio::time::timeout(STOP_TIMEOUT, client.flush())
166 .await
167 .wrap_err("the last replies were not written in time")?
168 .wrap_err("cannot write the last replies")?;
169 tracing::info!("stopped");
170 eyre::ensure!(!ended, "the NATS subscriptions closed");
171 Ok(())
172 }
173}
174
175async fn answer(dispatcher: Arc<Dispatcher>, client: Client, tenant: Arc<str>, message: Message) {
179 let Some(reply_to) = message.reply else {
180 tracing::warn!(subject = %message.subject, "a call without a reply subject is dropped");
181 return;
182 };
183 let (reply, deadline) = match decode(&message.payload) {
184 Ok(request) if !request.matches(&tenant, &message.subject) => {
185 tracing::warn!(subject = %message.subject, "a call for another subject is refused");
186 (
187 ExtReply::failure(
188 CallErrorKind::Invalid.as_str(),
189 "the call does not match the subject it came on",
190 ),
191 None,
192 )
193 }
194 Ok(request) => {
195 let deadline = request.context.deadline();
196 (dispatcher.dispatch(request, deadline).await, Some(deadline))
197 }
198 Err(refused) => (refused, None),
199 };
200 if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
201 tracing::debug!(subject = %message.subject, "the caller no longer waits; no reply");
202 return;
203 }
204 publish(&client, reply_to, &reply).await;
205}
206
207async fn publish(client: &Client, reply_to: Subject, reply: &ExtReply) {
208 let sent = match send(client, reply_to.clone(), reply).await {
209 Err(e)
210 if e.downcast_ref::<PublishError>()
211 .is_some_and(|e| e.kind() == PublishErrorKind::MaxPayloadExceeded) =>
212 {
213 let too_large = CallErrorKind::TooLarge.as_str();
214 send(
215 client,
216 reply_to,
217 &ExtReply::failure(too_large, e.to_string()),
218 )
219 .await
220 }
221 sent => sent,
222 };
223 if let Err(e) = sent {
224 tracing::warn!(error = %e, "cannot send a reply");
225 }
226}
227
228async fn send(client: &Client, reply_to: Subject, reply: &ExtReply) -> eyre::Result<()> {
229 Ok(client
230 .publish(reply_to, serde_json::to_vec(reply)?.into())
231 .await?)
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237 use std::collections::HashMap;
238
239 fn vars(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
240 let map: HashMap<String, String> = pairs
241 .iter()
242 .map(|(k, v)| (k.to_string(), v.to_string()))
243 .collect();
244 move |name| map.get(name).cloned()
245 }
246
247 #[test]
248 fn a_provider_serves_the_tenant_of_its_environment() {
249 let provider = NatsProvider::from_vars(vars(&[("CUFFLINK_TENANT", "harbor")])).unwrap();
250 assert_eq!(provider.tenant(), "harbor");
251 assert_eq!(provider.url, "nats://localhost:4222");
252 assert_eq!(provider.auth, NatsAuth::Anonymous);
253 let provider = NatsProvider::from_vars(vars(&[
254 ("CUFFLINK_TENANT", "harbor"),
255 ("NATS_URL", "nats://a:4222,nats://b:4222"),
256 ("NATS_USER", "provider"),
257 ("NATS_PASSWORD", "pw"),
258 ]))
259 .unwrap();
260 assert_eq!(provider.url, "nats://a:4222,nats://b:4222");
261 assert_eq!(
262 provider.auth,
263 NatsAuth::UserPassword {
264 user: "provider".into(),
265 password: "pw".into()
266 }
267 );
268 }
269
270 #[test]
271 fn a_provider_needs_a_tenant_a_server_and_whole_credentials() {
272 for (pairs, reason) in [
273 (
274 vec![],
275 "CUFFLINK_TENANT must name the tenant whose calls to serve",
276 ),
277 (
278 vec![("CUFFLINK_TENANT", "Harbor")],
279 "'Harbor' is not a tenant slug (1 to 40 of a-z, 0-9 and '-')",
280 ),
281 (
282 vec![("CUFFLINK_TENANT", "harbor"), ("NATS_URL", " , ")],
283 "NATS_URL lists no servers",
284 ),
285 (
286 vec![("CUFFLINK_TENANT", "harbor"), ("NATS_USER", "provider")],
287 "NATS_USER and NATS_PASSWORD must be set together",
288 ),
289 ] {
290 let refused = NatsProvider::from_vars(vars(&pairs)).unwrap_err();
291 assert_eq!(refused.to_string(), reason);
292 }
293 }
294
295 #[test]
296 fn debug_never_prints_credentials() {
297 let provider = NatsProvider::new(
298 "harbor",
299 "nats://user:s3cret@nats:4222",
300 NatsAuth::NKeySeed("SUASECRET".into()),
301 )
302 .unwrap();
303 let rendered = format!("{provider:?}");
304 assert!(
305 !rendered.contains("s3cret") && !rendered.contains("SUASECRET"),
306 "{rendered}"
307 );
308 }
309}