async_nats/service/
mod.rs

1// Copyright 2020-2023 The NATS Authors
2// Licensed under the Apache License, Version 2.0 (the "License");
3// you may not use this file except in compliance with the License.
4// You may obtain a copy of the License at
5//
6// http://www.apache.org/licenses/LICENSE-2.0
7//
8// Unless required by applicable law or agreed to in writing, software
9// distributed under the License is distributed on an "AS IS" BASIS,
10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11// See the License for the specific language governing permissions and
12// limitations under the License.
13
14pub mod error;
15
16use std::{
17    collections::HashMap,
18    fmt::Display,
19    pin::Pin,
20    sync::{Arc, Mutex},
21    time::{Duration, Instant},
22};
23
24use bytes::Bytes;
25pub mod endpoint;
26use crate::datetime::{self, rfc3339, DateTime};
27use futures_util::{
28    stream::{self, SelectAll},
29    Future, StreamExt,
30};
31use regex::Regex;
32use serde::{Deserialize, Serialize};
33use std::sync::LazyLock;
34use tokio::{sync::broadcast::Sender, task::JoinHandle};
35use tracing::debug;
36
37use crate::{
38    client::PublishErrorKind, Client, Error, HeaderMap, Message, PublishError, Subscriber,
39};
40
41use self::endpoint::Endpoint;
42
43const SERVICE_API_PREFIX: &str = "$SRV";
44const DEFAULT_QUEUE_GROUP: &str = "q";
45pub const NATS_SERVICE_ERROR: &str = "Nats-Service-Error";
46pub const NATS_SERVICE_ERROR_CODE: &str = "Nats-Service-Error-Code";
47
48// uses recommended semver validation expression from
49// https://semver.org/#is-there-a-suggested-regular-expression-regex-to-check-a-semver-string
50static SEMVER: LazyLock<Regex> = LazyLock::new(|| {
51    Regex::new(r"^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-((?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(?:\.(?:0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*))?(?:\+([0-9a-zA-Z-]+(?:\.[0-9a-zA-Z-]+)*))?$")
52        .unwrap()
53});
54// From ADR-33: Name can only have A-Z, a-z, 0-9, dash, underscore.
55static NAME: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^[A-Za-z0-9\-_]+$").unwrap());
56
57/// Represents state for all endpoints.
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub(crate) struct Endpoints {
60    pub(crate) endpoints: HashMap<String, endpoint::Inner>,
61}
62
63/// Response for `PING` requests.
64#[derive(Serialize, Deserialize)]
65pub struct PingResponse {
66    /// Response type.
67    #[serde(rename = "type")]
68    pub kind: String,
69    /// Service name.
70    pub name: String,
71    /// Service id.
72    pub id: String,
73    /// Service version.
74    pub version: String,
75    /// Additional metadata
76    #[serde(default, deserialize_with = "endpoint::null_meta_as_default")]
77    pub metadata: HashMap<String, String>,
78}
79
80/// Response for `STATS` requests.
81#[derive(Serialize, Deserialize)]
82pub struct Stats {
83    /// Response type.
84    #[serde(rename = "type")]
85    pub kind: String,
86    /// Service name.
87    pub name: String,
88    /// Service id.
89    pub id: String,
90    // Service version.
91    pub version: String,
92    #[serde(with = "rfc3339")]
93    pub started: DateTime,
94    /// Statistics of all endpoints.
95    pub endpoints: Vec<endpoint::Stats>,
96}
97
98/// Information about service instance.
99/// Service name.
100#[derive(Serialize, Deserialize, Debug, Clone)]
101pub struct Info {
102    /// Response type.
103    #[serde(rename = "type")]
104    pub kind: String,
105    /// Service name.
106    pub name: String,
107    /// Service id.
108    pub id: String,
109    /// Service description.
110    pub description: String,
111    /// Service version.
112    pub version: String,
113    /// Additional metadata
114    #[serde(default, deserialize_with = "endpoint::null_meta_as_default")]
115    pub metadata: HashMap<String, String>,
116    /// Info about all service endpoints.
117    pub endpoints: Vec<endpoint::Info>,
118}
119
120/// Configuration of the [Service].
121#[derive(Serialize, Deserialize, Debug)]
122pub struct Config {
123    /// Really the kind of the service. Shared by all the services that have the same name.
124    /// This name can only have A-Z, a-z, 0-9, dash, underscore
125    pub name: String,
126    /// a human-readable description about the service
127    pub description: Option<String>,
128    /// A SemVer valid service version.
129    pub version: String,
130    /// Custom handler for providing the `EndpointStats.data` value.
131    #[serde(skip)]
132    pub stats_handler: Option<StatsHandler>,
133    /// Additional service metadata
134    pub metadata: Option<HashMap<String, String>>,
135    /// Custom queue group config
136    pub queue_group: Option<String>,
137}
138
139pub struct ServiceBuilder {
140    client: Client,
141    description: Option<String>,
142    stats_handler: Option<StatsHandler>,
143    metadata: Option<HashMap<String, String>>,
144    queue_group: Option<String>,
145}
146
147impl ServiceBuilder {
148    fn new(client: Client) -> Self {
149        Self {
150            client,
151            description: None,
152            stats_handler: None,
153            metadata: None,
154            queue_group: None,
155        }
156    }
157
158    /// Description for the service.
159    pub fn description<S: ToString>(mut self, description: S) -> Self {
160        self.description = Some(description.to_string());
161        self
162    }
163
164    /// Handler for custom service statistics.
165    pub fn stats_handler<F>(mut self, handler: F) -> Self
166    where
167        F: FnMut(String, endpoint::Stats) -> serde_json::Value + Send + Sync + 'static,
168    {
169        self.stats_handler = Some(StatsHandler(Box::new(handler)));
170        self
171    }
172
173    /// Additional service metadata.
174    pub fn metadata(mut self, metadata: HashMap<String, String>) -> Self {
175        self.metadata = Some(metadata);
176        self
177    }
178
179    /// Custom queue group. Default is `q`.
180    pub fn queue_group<S: ToString>(mut self, queue_group: S) -> Self {
181        self.queue_group = Some(queue_group.to_string());
182        self
183    }
184
185    /// Starts the service with configured options.
186    pub async fn start<N: ToString, V: ToString>(
187        self,
188        name: N,
189        version: V,
190    ) -> Result<Service, Error> {
191        Service::add(
192            self.client,
193            Config {
194                name: name.to_string(),
195                version: version.to_string(),
196                description: self.description,
197                stats_handler: self.stats_handler,
198                metadata: self.metadata,
199                queue_group: self.queue_group,
200            },
201        )
202        .await
203    }
204}
205
206/// Verbs that can be used to acquire information from the services.
207pub enum Verb {
208    Ping,
209    Stats,
210    Info,
211    Schema,
212}
213
214impl Display for Verb {
215    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
216        match self {
217            Verb::Ping => write!(f, "PING"),
218            Verb::Stats => write!(f, "STATS"),
219            Verb::Info => write!(f, "INFO"),
220            Verb::Schema => write!(f, "SCHEMA"),
221        }
222    }
223}
224
225pub trait ServiceExt {
226    type Output: Future<Output = Result<Service, crate::Error>>;
227
228    /// Adds a Service instance.
229    ///
230    /// # Examples
231    ///
232    /// ```no_run
233    /// # #[tokio::main]
234    /// # async fn main() -> Result<(), async_nats::Error> {
235    /// use async_nats::service::ServiceExt;
236    /// use futures_util::StreamExt;
237    /// let client = async_nats::connect("demo.nats.io").await?;
238    /// let mut service = client
239    ///     .add_service(async_nats::service::Config {
240    ///         name: "generator".to_string(),
241    ///         version: "1.0.0".to_string(),
242    ///         description: None,
243    ///         stats_handler: None,
244    ///         metadata: None,
245    ///         queue_group: None,
246    ///     })
247    ///     .await?;
248    ///
249    /// let mut endpoint = service.endpoint("get").await?;
250    ///
251    /// if let Some(request) = endpoint.next().await {
252    ///     request.respond(Ok("hello".into())).await?;
253    /// }
254    ///
255    /// # Ok(())
256    /// # }
257    /// ```
258    fn add_service(&self, config: Config) -> Self::Output;
259
260    /// Returns Service instance builder.
261    ///
262    /// # Examples
263    ///
264    /// ```no_run
265    /// # #[tokio::main]
266    /// # async fn main() -> Result<(), async_nats::Error> {
267    /// use async_nats::service::ServiceExt;
268    /// use futures_util::StreamExt;
269    /// let client = async_nats::connect("demo.nats.io").await?;
270    /// let mut service = client
271    ///     .service_builder()
272    ///     .description("some service")
273    ///     .stats_handler(|endpoint, stats| serde_json::json!({ "endpoint": endpoint }))
274    ///     .start("products", "1.0.0")
275    ///     .await?;
276    ///
277    /// let mut endpoint = service.endpoint("get").await?;
278    ///
279    /// if let Some(request) = endpoint.next().await {
280    ///     request.respond(Ok("hello".into())).await?;
281    /// }
282    /// # Ok(())
283    /// # }
284    /// ```
285    fn service_builder(&self) -> ServiceBuilder;
286}
287
288impl ServiceExt for Client {
289    type Output = Pin<Box<dyn Future<Output = Result<Service, crate::Error>> + Send>>;
290
291    fn add_service(&self, config: Config) -> Self::Output {
292        let client = self.clone();
293        Box::pin(async { Service::add(client, config).await })
294    }
295
296    fn service_builder(&self) -> ServiceBuilder {
297        ServiceBuilder::new(self.clone())
298    }
299}
300
301/// Service instance.
302///
303/// # Examples
304///
305/// ```no_run
306/// # #[tokio::main]
307/// # async fn main() -> Result<(), async_nats::Error> {
308/// use async_nats::service::ServiceExt;
309/// use futures_util::StreamExt;
310/// let client = async_nats::connect("demo.nats.io").await?;
311/// let mut service = client.service_builder().start("generator", "1.0.0").await?;
312/// let mut endpoint = service.endpoint("get").await?;
313///
314/// if let Some(request) = endpoint.next().await {
315///     request.respond(Ok("hello".into())).await?;
316/// }
317///
318/// # Ok(())
319/// # }
320/// ```
321#[derive(Debug)]
322pub struct Service {
323    endpoints_state: Arc<Mutex<Endpoints>>,
324    info: Info,
325    client: Client,
326    handle: JoinHandle<Result<(), Error>>,
327    shutdown_tx: Sender<()>,
328    subjects: Arc<Mutex<Vec<String>>>,
329    queue_group: String,
330}
331
332impl Service {
333    async fn add(client: Client, config: Config) -> Result<Service, Error> {
334        // validate service version semver string.
335        if !SEMVER.is_match(config.version.as_str()) {
336            return Err(Box::new(std::io::Error::new(
337                std::io::ErrorKind::InvalidInput,
338                "service version is not a valid semver string",
339            )));
340        }
341        // validate service name.
342        if !NAME.is_match(config.name.as_str()) {
343            return Err(Box::new(std::io::Error::new(
344                std::io::ErrorKind::InvalidInput,
345                "service name is not a valid string (only A-Z, a-z, 0-9, _, - are allowed)",
346            )));
347        }
348        let endpoints_state = Arc::new(Mutex::new(Endpoints {
349            endpoints: HashMap::new(),
350        }));
351
352        let queue_group = config
353            .queue_group
354            .unwrap_or(DEFAULT_QUEUE_GROUP.to_string());
355        let id = crate::id_generator::next();
356        let started = datetime::now();
357        let subjects = Arc::new(Mutex::new(Vec::new()));
358        let info = Info {
359            kind: "io.nats.micro.v1.info_response".to_string(),
360            name: config.name.clone(),
361            id: id.clone(),
362            description: config.description.clone().unwrap_or_default(),
363            version: config.version.clone(),
364            metadata: config.metadata.clone().unwrap_or_default(),
365            endpoints: Vec::new(),
366        };
367
368        let (shutdown_tx, _) = tokio::sync::broadcast::channel(1);
369
370        // create subscriptions for all verbs.
371        let mut pings =
372            verb_subscription(client.clone(), Verb::Ping, config.name.clone(), id.clone()).await?;
373        let mut infos =
374            verb_subscription(client.clone(), Verb::Info, config.name.clone(), id.clone()).await?;
375        let mut stats =
376            verb_subscription(client.clone(), Verb::Stats, config.name.clone(), id.clone()).await?;
377
378        // Start a task for handling verbs subscriptions.
379        let handle = tokio::task::spawn({
380            let mut stats_callback = config.stats_handler;
381            let info = info.clone();
382            let endpoints_state = endpoints_state.clone();
383            let client = client.clone();
384            async move {
385                loop {
386                    tokio::select! {
387                        Some(ping) = pings.next() => {
388                            let pong = serde_json::to_vec(&PingResponse{
389                                kind: "io.nats.micro.v1.ping_response".to_string(),
390                                name: info.name.clone(),
391                                id: info.id.clone(),
392                                version: info.version.clone(),
393                                metadata: info.metadata.clone(),
394                            })?;
395                            client.publish(ping.reply.unwrap(), pong.into()).await?;
396                        },
397                        Some(info_request) = infos.next() => {
398                            let info = info.clone();
399
400                            let endpoints: Vec<endpoint::Info> = {
401                                endpoints_state.lock().unwrap().endpoints.values().map(|value| {
402                                    endpoint::Info {
403                                        name: value.name.to_owned(),
404                                        subject: value.subject.to_owned(),
405                                        queue_group: value.queue_group.to_owned(),
406                                        metadata: value.metadata.to_owned()
407                                    }
408                                }).collect()
409                            };
410                            let info = Info {
411                                endpoints,
412                                ..info
413                            };
414                            let info_json = serde_json::to_vec(&info).map(Bytes::from)?;
415                            client.publish(info_request.reply.unwrap(), info_json.clone()).await?;
416                        },
417                        Some(stats_request) = stats.next() => {
418                            if let Some(stats_callback) = stats_callback.as_mut() {
419                                let mut endpoint_stats_locked = endpoints_state.lock().unwrap();
420                                for (key, value) in &mut endpoint_stats_locked.endpoints {
421                                    let data = stats_callback.0(key.to_string(), value.clone().into());
422                                    value.data = Some(data);
423                                }
424                            }
425                            let stats = serde_json::to_vec(&Stats {
426                                kind: "io.nats.micro.v1.stats_response".to_string(),
427                                name: info.name.clone(),
428                                id: info.id.clone(),
429                                version: info.version.clone(),
430                                started,
431                                endpoints: endpoints_state.lock().unwrap().endpoints.values().cloned().map(Into::into).collect(),
432                            })?;
433                            client.publish(stats_request.reply.unwrap(), stats.into()).await?;
434                        },
435                        else => break,
436                    }
437                }
438                Ok(())
439            }
440        });
441        Ok(Service {
442            endpoints_state,
443            info,
444            client,
445            handle,
446            shutdown_tx,
447            subjects,
448            queue_group,
449        })
450    }
451    /// Stops this instance of the [Service].
452    /// If there are more instances of [Services][Service] with the same name, the [Service] will
453    /// be scaled down by one instance. If it was the only running instance, it will effectively
454    /// remove the service entirely.
455    pub async fn stop(self) -> Result<(), Error> {
456        self.shutdown_tx.send(())?;
457        self.handle.abort();
458        Ok(())
459    }
460
461    /// Resets [Stats] of the [Service] instance.
462    pub async fn reset(&mut self) {
463        for value in self.endpoints_state.lock().unwrap().endpoints.values_mut() {
464            value.errors = 0;
465            value.processing_time = Duration::default();
466            value.requests = 0;
467            value.average_processing_time = Duration::default();
468        }
469    }
470
471    /// Returns [Stats] for this service instance.
472    pub async fn stats(&self) -> HashMap<String, endpoint::Stats> {
473        self.endpoints_state
474            .lock()
475            .unwrap()
476            .endpoints
477            .iter()
478            .map(|(key, value)| (key.to_owned(), value.to_owned().into()))
479            .collect()
480    }
481
482    /// Returns [Info] for this service instance.
483    pub async fn info(&self) -> Info {
484        self.info.clone()
485    }
486
487    /// Creates a group for endpoints under common prefix.
488    ///
489    /// # Examples
490    ///
491    /// ```no_run
492    /// # #[tokio::main]
493    /// # async fn main() -> Result<(), async_nats::Error> {
494    /// use async_nats::service::ServiceExt;
495    /// let client = async_nats::connect("demo.nats.io").await?;
496    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
497    ///
498    /// let v1 = service.group("v1");
499    /// let products = v1.endpoint("products").await?;
500    /// # Ok(())
501    /// # }
502    /// ```
503    pub fn group<S: ToString>(&self, prefix: S) -> Group {
504        self.group_with_queue_group(prefix, self.queue_group.clone())
505    }
506
507    /// Creates a group for endpoints under common prefix with custom queue group.
508    ///
509    /// # Examples
510    ///
511    /// ```no_run
512    /// # #[tokio::main]
513    /// # async fn main() -> Result<(), async_nats::Error> {
514    /// use async_nats::service::ServiceExt;
515    /// let client = async_nats::connect("demo.nats.io").await?;
516    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
517    ///
518    /// let v1 = service.group("v1");
519    /// let products = v1.endpoint("products").await?;
520    /// # Ok(())
521    /// # }
522    /// ```
523    pub fn group_with_queue_group<S: ToString, Z: ToString>(
524        &self,
525        prefix: S,
526        queue_group: Z,
527    ) -> Group {
528        Group {
529            subjects: self.subjects.clone(),
530            prefix: prefix.to_string(),
531            stats: self.endpoints_state.clone(),
532            client: self.client.clone(),
533            shutdown_tx: self.shutdown_tx.clone(),
534            queue_group: queue_group.to_string(),
535        }
536    }
537
538    /// Builder for customized [Endpoint] creation.
539    ///
540    /// # Examples
541    ///
542    /// ```no_run
543    /// # #[tokio::main]
544    /// # async fn main() -> Result<(), async_nats::Error> {
545    /// use async_nats::service::ServiceExt;
546    /// let client = async_nats::connect("demo.nats.io").await?;
547    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
548    ///
549    /// let products = service
550    ///     .endpoint_builder()
551    ///     .name("api")
552    ///     .add("products")
553    ///     .await?;
554    /// # Ok(())
555    /// # }
556    /// ```
557    pub fn endpoint_builder(&self) -> EndpointBuilder {
558        EndpointBuilder::new(
559            self.client.clone(),
560            self.endpoints_state.clone(),
561            self.shutdown_tx.clone(),
562            self.subjects.clone(),
563            self.queue_group.clone(),
564        )
565    }
566
567    /// Adds a new endpoint to the [Service].
568    ///
569    /// # Examples
570    ///
571    /// ```no_run
572    /// # #[tokio::main]
573    /// # async fn main() -> Result<(), async_nats::Error> {
574    /// use async_nats::service::ServiceExt;
575    /// let client = async_nats::connect("demo.nats.io").await?;
576    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
577    ///
578    /// let products = service.endpoint("products").await?;
579    /// # Ok(())
580    /// # }
581    /// ```
582    pub async fn endpoint<S: ToString>(&self, subject: S) -> Result<Endpoint, Error> {
583        EndpointBuilder::new(
584            self.client.clone(),
585            self.endpoints_state.clone(),
586            self.shutdown_tx.clone(),
587            self.subjects.clone(),
588            self.queue_group.clone(),
589        )
590        .add(subject)
591        .await
592    }
593}
594
595pub struct Group {
596    prefix: String,
597    stats: Arc<Mutex<Endpoints>>,
598    client: Client,
599    shutdown_tx: Sender<()>,
600    subjects: Arc<Mutex<Vec<String>>>,
601    queue_group: String,
602}
603
604impl Group {
605    /// Creates a group for [Endpoints][Endpoint] under common prefix.
606    ///
607    /// # Examples
608    ///
609    /// ```no_run
610    /// # #[tokio::main]
611    /// # async fn main() -> Result<(), async_nats::Error> {
612    /// use async_nats::service::ServiceExt;
613    /// let client = async_nats::connect("demo.nats.io").await?;
614    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
615    ///
616    /// let v1 = service.group("v1");
617    /// let products = v1.endpoint("products").await?;
618    /// # Ok(())
619    /// # }
620    /// ```
621    pub fn group<S: ToString>(&self, prefix: S) -> Group {
622        self.group_with_queue_group(prefix, self.queue_group.clone())
623    }
624
625    /// Creates a group for [Endpoints][Endpoint] under common prefix with custom queue group.
626    ///
627    /// # Examples
628    ///
629    /// ```no_run
630    /// # #[tokio::main]
631    /// # async fn main() -> Result<(), async_nats::Error> {
632    /// use async_nats::service::ServiceExt;
633    /// let client = async_nats::connect("demo.nats.io").await?;
634    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
635    ///
636    /// let v1 = service.group("v1");
637    /// let products = v1.endpoint("products").await?;
638    /// # Ok(())
639    /// # }
640    /// ```
641    pub fn group_with_queue_group<S: ToString, Z: ToString>(
642        &self,
643        prefix: S,
644        queue_group: Z,
645    ) -> Group {
646        Group {
647            prefix: format!("{}.{}", self.prefix, prefix.to_string()),
648            stats: self.stats.clone(),
649            client: self.client.clone(),
650            shutdown_tx: self.shutdown_tx.clone(),
651            subjects: self.subjects.clone(),
652            queue_group: queue_group.to_string(),
653        }
654    }
655
656    /// Adds a new endpoint to the [Service] under current [Group]
657    ///
658    /// # Examples
659    ///
660    /// ```no_run
661    /// # #[tokio::main]
662    /// # async fn main() -> Result<(), async_nats::Error> {
663    /// use async_nats::service::ServiceExt;
664    /// let client = async_nats::connect("demo.nats.io").await?;
665    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
666    /// let v1 = service.group("v1");
667    ///
668    /// let products = v1.endpoint("products").await?;
669    /// # Ok(())
670    /// # }
671    /// ```
672    pub async fn endpoint<S: ToString>(&self, subject: S) -> Result<Endpoint, Error> {
673        let endpoint = self.endpoint_builder();
674        endpoint.add(subject.to_string()).await
675    }
676
677    /// Builder for customized [Endpoint] creation under current [Group]
678    ///
679    /// # Examples
680    ///
681    /// ```no_run
682    /// # #[tokio::main]
683    /// # async fn main() -> Result<(), async_nats::Error> {
684    /// use async_nats::service::ServiceExt;
685    /// let client = async_nats::connect("demo.nats.io").await?;
686    /// let mut service = client.service_builder().start("service", "1.0.0").await?;
687    /// let v1 = service.group("v1");
688    ///
689    /// let products = v1.endpoint_builder().name("api").add("products").await?;
690    /// # Ok(())
691    /// # }
692    /// ```
693    pub fn endpoint_builder(&self) -> EndpointBuilder {
694        let mut endpoint = EndpointBuilder::new(
695            self.client.clone(),
696            self.stats.clone(),
697            self.shutdown_tx.clone(),
698            self.subjects.clone(),
699            self.queue_group.clone(),
700        );
701        endpoint.prefix = Some(self.prefix.clone());
702        endpoint
703    }
704}
705
706async fn verb_subscription(
707    client: Client,
708    verb: Verb,
709    name: String,
710    id: String,
711) -> Result<stream::Fuse<SelectAll<Subscriber>>, Error> {
712    let verb_all = client
713        .subscribe(format!("{SERVICE_API_PREFIX}.{verb}"))
714        .await?;
715    let verb_name = client
716        .subscribe(format!("{SERVICE_API_PREFIX}.{verb}.{name}"))
717        .await?;
718    let verb_id = client
719        .subscribe(format!("{SERVICE_API_PREFIX}.{verb}.{name}.{id}"))
720        .await?;
721    Ok(stream::select_all([verb_all, verb_id, verb_name]).fuse())
722}
723
724type ShutdownReceiverFuture = Pin<
725    Box<dyn Future<Output = Result<(), tokio::sync::broadcast::error::RecvError>> + Send + Sync>,
726>;
727
728/// Request returned by [Service] [Stream][futures_util::Stream].
729#[derive(Debug)]
730pub struct Request {
731    issued: Instant,
732    client: Client,
733    pub message: Message,
734    endpoint: String,
735    stats: Arc<Mutex<Endpoints>>,
736}
737
738impl Request {
739    /// Sends response for the request.
740    ///
741    /// # Examples
742    ///
743    /// ```no_run
744    /// # #[tokio::main]
745    /// # async fn main() -> Result<(), async_nats::Error> {
746    /// use async_nats::service::ServiceExt;
747    /// use futures_util::StreamExt;
748    /// # let client = async_nats::connect("demo.nats.io").await?;
749    /// # let mut service = client
750    /// #    .service_builder().start("serviceA", "1.0.0.1").await?;
751    /// let mut endpoint = service.endpoint("endpoint").await?;
752    /// let request = endpoint.next().await.unwrap();
753    /// request.respond(Ok("hello".into())).await?;
754    /// # Ok(())
755    /// # }
756    /// ```
757    pub async fn respond(&self, response: Result<Bytes, error::Error>) -> Result<(), PublishError> {
758        self.respond_with_headers(response, HeaderMap::new()).await
759    }
760
761    /// Sends response for the request with headers.
762    ///
763    /// On error responses, [Nats-Service-Error][NATS_SERVICE_ERROR] and
764    /// [Nats-Service-Error-Code][NATS_SERVICE_ERROR_CODE] are always set from the provided
765    /// [`error::Error`]. If the provided [HeaderMap] already contains values for either
766    /// of those headers, they will be overridden. All other user-supplied headers
767    /// are preserved.
768    ///
769    /// # Examples
770    ///
771    /// ```no_run
772    /// # #[tokio::main]
773    /// # async fn main() -> Result<(), async_nats::Error> {
774    /// use async_nats::service::ServiceExt;
775    /// use futures_util::StreamExt;
776    /// # let client = async_nats::connect("demo.nats.io").await?;
777    /// # let mut service = client
778    /// #    .service_builder().start("serviceA", "1.0.0.1").await?;
779    /// let mut endpoint = service.endpoint("endpoint").await?;
780    /// let request = endpoint.next().await.unwrap();
781    /// let mut headers = async_nats::HeaderMap::new();
782    /// headers.insert("x-success", "true");
783    /// request
784    ///     .respond_with_headers(Ok("hello".into()), headers)
785    ///     .await?;
786    /// # Ok(())
787    /// # }
788    /// ```
789    pub async fn respond_with_headers(
790        &self,
791        response: Result<Bytes, error::Error>,
792        mut headers: HeaderMap,
793    ) -> Result<(), PublishError> {
794        let reply = match self.message.reply.clone() {
795            None => {
796                return Err(PublishError::with_source(
797                    PublishErrorKind::InvalidSubject,
798                    "Request is missing reply subject to respond to",
799                ))
800            }
801            Some(subject) => subject,
802        };
803        let result = match response {
804            Ok(payload) => {
805                if headers.is_empty() {
806                    self.client.publish(reply, payload).await
807                } else {
808                    self.client
809                        .publish_with_headers(reply, headers, payload)
810                        .await
811                }
812            }
813            Err(err) => {
814                self.stats
815                    .lock()
816                    .unwrap()
817                    .endpoints
818                    .entry(self.endpoint.clone())
819                    .and_modify(|stats| {
820                        stats.last_error = Some(err.clone());
821                        stats.errors += 1;
822                    })
823                    .or_default();
824                headers.insert(NATS_SERVICE_ERROR, err.status.as_str());
825                headers.insert(NATS_SERVICE_ERROR_CODE, err.code.to_string().as_str());
826                self.client
827                    .publish_with_headers(reply, headers, "".into())
828                    .await
829            }
830        };
831        let elapsed = self.issued.elapsed();
832        let mut stats = self.stats.lock().unwrap();
833        let stats = stats.endpoints.get_mut(self.endpoint.as_str()).unwrap();
834        stats.requests += 1;
835        stats.processing_time += elapsed;
836        stats.average_processing_time = {
837            let avg_nanos = (stats.processing_time.as_nanos() / stats.requests as u128) as u64;
838            Duration::from_nanos(avg_nanos)
839        };
840        result
841    }
842}
843
844#[derive(Debug)]
845pub struct EndpointBuilder {
846    client: Client,
847    stats: Arc<Mutex<Endpoints>>,
848    shutdown_tx: Sender<()>,
849    name: Option<String>,
850    metadata: Option<HashMap<String, String>>,
851    subjects: Arc<Mutex<Vec<String>>>,
852    queue_group: String,
853    prefix: Option<String>,
854}
855
856impl EndpointBuilder {
857    fn new(
858        client: Client,
859        stats: Arc<Mutex<Endpoints>>,
860        shutdown_tx: Sender<()>,
861        subjects: Arc<Mutex<Vec<String>>>,
862        queue_group: String,
863    ) -> EndpointBuilder {
864        EndpointBuilder {
865            client,
866            stats,
867            subjects,
868            shutdown_tx,
869            name: None,
870            metadata: None,
871            queue_group,
872            prefix: None,
873        }
874    }
875
876    /// Name of the [Endpoint]. By default, the subject of the endpoint is used.
877    pub fn name<S: ToString>(mut self, name: S) -> EndpointBuilder {
878        self.name = Some(name.to_string());
879        self
880    }
881
882    /// Metadata specific for the [Endpoint].
883    pub fn metadata(mut self, metadata: HashMap<String, String>) -> EndpointBuilder {
884        self.metadata = Some(metadata);
885        self
886    }
887
888    /// Custom queue group for the [Endpoint]. Otherwise, it will be derived from group or service.
889    pub fn queue_group<S: ToString>(mut self, queue_group: S) -> EndpointBuilder {
890        self.queue_group = queue_group.to_string();
891        self
892    }
893
894    /// Finalizes the builder and adds the [Endpoint].
895    pub async fn add<S: ToString>(self, subject: S) -> Result<Endpoint, Error> {
896        let mut subject = subject.to_string();
897        if let Some(prefix) = self.prefix {
898            subject = format!("{prefix}.{subject}");
899        }
900        let endpoint_name = self.name.clone().unwrap_or_else(|| subject.clone());
901        let name = self
902            .name
903            .clone()
904            .unwrap_or_else(|| subject.clone().replace('.', "-"));
905        let requests = self
906            .client
907            .queue_subscribe(subject.to_owned(), self.queue_group.to_string())
908            .await?;
909        debug!("created service for endpoint {subject}");
910
911        let shutdown_rx = self.shutdown_tx.subscribe();
912
913        let mut stats = self.stats.lock().unwrap();
914        stats
915            .endpoints
916            .entry(endpoint_name.clone())
917            .or_insert(endpoint::Inner {
918                name,
919                subject: subject.clone(),
920                metadata: self.metadata.unwrap_or_default(),
921                queue_group: self.queue_group.clone(),
922                ..Default::default()
923            });
924        self.subjects.lock().unwrap().push(subject.clone());
925        Ok(Endpoint {
926            requests,
927            stats: self.stats.clone(),
928            client: self.client.clone(),
929            endpoint: endpoint_name,
930            shutdown: Some(shutdown_rx),
931            shutdown_future: None,
932        })
933    }
934}
935
936pub struct StatsHandler(pub Box<dyn FnMut(String, endpoint::Stats) -> serde_json::Value + Send>);
937
938impl std::fmt::Debug for StatsHandler {
939    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
940        write!(f, "Stats handler")
941    }
942}
943
944#[cfg(test)]
945mod tests {
946    use super::*;
947
948    #[tokio::test]
949    async fn test_group_with_queue_group() {
950        let server = nats_server::run_basic_server();
951        let client = crate::connect(server.client_url()).await.unwrap();
952
953        let group = Group {
954            prefix: "test".to_string(),
955            stats: Arc::new(Mutex::new(Endpoints {
956                endpoints: HashMap::new(),
957            })),
958            client,
959            shutdown_tx: tokio::sync::broadcast::channel(1).0,
960            subjects: Arc::new(Mutex::new(vec![])),
961            queue_group: "default".to_string(),
962        };
963
964        let new_group = group.group_with_queue_group("v1", "custom_queue");
965
966        assert_eq!(new_group.prefix, "test.v1");
967        assert_eq!(new_group.queue_group, "custom_queue");
968    }
969
970    #[tokio::test]
971    async fn test_respond_with_headers_overrides_error_headers() {
972        let server = nats_server::run_basic_server();
973        let client = crate::connect(server.client_url()).await.unwrap();
974
975        let service = client
976            .service_builder()
977            .start("test-service", "1.0.0")
978            .await
979            .unwrap();
980
981        let subject = "test.subject";
982        let mut endpoint = service.endpoint(subject).await.unwrap();
983
984        let handler = async {
985            if let Some(request) = endpoint.next().await {
986                let mut resp_headers = HeaderMap::new();
987                resp_headers.insert("x-success", "false");
988                resp_headers.insert(NATS_SERVICE_ERROR, "user-supplied-value");
989                resp_headers.insert(NATS_SERVICE_ERROR_CODE, "999");
990
991                let err = error::Error {
992                    status: "internal-error".to_string(),
993                    code: 500,
994                };
995
996                request
997                    .respond_with_headers(Err(err), resp_headers)
998                    .await
999                    .expect("failed to send response");
1000            }
1001        };
1002
1003        let requester = crate::connect(server.client_url()).await.unwrap();
1004        let request_fut = async { requester.request(subject, "".into()).await.unwrap() };
1005
1006        let (_, resp) = tokio::join!(handler, request_fut);
1007
1008        let headers = resp.headers.expect("expected headers on reply");
1009        assert_eq!(headers.get("x-success").unwrap().as_str(), "false");
1010        assert_eq!(
1011            headers.get(NATS_SERVICE_ERROR).unwrap().as_str(),
1012            "internal-error"
1013        );
1014        assert_eq!(
1015            headers.get(NATS_SERVICE_ERROR_CODE).unwrap().as_str(),
1016            "500"
1017        );
1018    }
1019
1020    #[tokio::test]
1021    async fn test_respond_with_headers_preserves_headers_on_success() {
1022        let server = nats_server::run_basic_server();
1023        let client = crate::connect(server.client_url()).await.unwrap();
1024
1025        let service = client
1026            .service_builder()
1027            .start("test-service", "1.0.0")
1028            .await
1029            .unwrap();
1030
1031        let subject = "test.subject";
1032        let mut endpoint = service.endpoint(subject).await.unwrap();
1033
1034        let handler = async {
1035            if let Some(request) = endpoint.next().await {
1036                let mut resp_headers = HeaderMap::new();
1037                resp_headers.insert("x-success", "false");
1038                resp_headers.insert("x-request-id", "req-123");
1039                resp_headers.insert(NATS_SERVICE_ERROR, "user-supplied-value");
1040                resp_headers.insert(NATS_SERVICE_ERROR_CODE, "999");
1041
1042                request
1043                    .respond_with_headers(Ok("ok".into()), resp_headers)
1044                    .await
1045                    .unwrap();
1046            }
1047        };
1048
1049        let requester = crate::connect(server.client_url()).await.unwrap();
1050        let request_fut = async { requester.request(subject, "".into()).await.unwrap() };
1051
1052        let (_, resp) = tokio::join!(handler, request_fut);
1053
1054        let headers = resp.headers.expect("expected headers on reply");
1055        assert_eq!(headers.get("x-success").unwrap().as_str(), "false");
1056        assert_eq!(headers.get("x-request-id").unwrap().as_str(), "req-123");
1057        assert_eq!(
1058            headers.get(NATS_SERVICE_ERROR).unwrap().as_str(),
1059            "user-supplied-value"
1060        );
1061        assert_eq!(
1062            headers.get(NATS_SERVICE_ERROR_CODE).unwrap().as_str(),
1063            "999"
1064        );
1065    }
1066}
async_nats/service/mod.rs

async_nats/service/
mod.rs
