Skip to main content

cufflink_native/
nats.rs

1//! The NATS transport: a provider the project runs itself, answering its
2//! tenant's calls on their subjects.
3
4use 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";
22/// How long a stopping provider takes to answer the calls NATS delivered
23/// to it, and then to write its last replies.
24const STOP_TIMEOUT: Duration = Duration::from_secs(10);
25
26/// Where a NATS provider serves: its tenant and the NATS it connects to.
27#[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    /// A provider of the tenant `tenant` on the NATS servers `url` (a
46    /// comma-separated list), authenticated by `auth`.
47    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    /// The provider `CUFFLINK_TENANT`, `NATS_URL` (`nats://localhost:4222`
62    /// when unset) and the NATS credentials of the environment describe.
63    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
80/// A service subscribed to its calls, ready to [`serve`](Self::serve).
81pub 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    /// Answer calls until `shutdown` completes, then unsubscribe, answer the
116    /// calls already delivered and write the last replies. A call NATS
117    /// routes here while the unsubscribe is on its way is lost (the client
118    /// forgets the subscription once the unsubscribe is sent), and its
119    /// caller times out.
120    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
175/// Answer `message` on its reply subject, unless its caller no longer
176/// waits: a reply after the deadline could only race the caller's own
177/// timeout.
178async 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}