tower/spawn_ready/
service.rs

1use super::{future::ResponseFuture, SpawnReadyLayer};
2use crate::{util::ServiceExt, BoxError};
3use futures_util::future::TryFutureExt;
4use std::{
5    future::Future,
6    pin::Pin,
7    task::{ready, Context, Poll},
8};
9use tower_service::Service;
10use tracing::Instrument;
11
12/// Spawns tasks to drive an inner service to readiness.
13///
14/// See crate level documentation for more details.
15#[derive(Debug)]
16pub struct SpawnReady<S> {
17    inner: Inner<S>,
18}
19
20#[derive(Debug)]
21enum Inner<S> {
22    Service(Option<S>),
23    Future(tokio::task::JoinHandle<Result<S, BoxError>>),
24}
25
26impl<S> SpawnReady<S> {
27    /// Creates a new [`SpawnReady`] wrapping `service`.
28    pub const fn new(service: S) -> Self {
29        Self {
30            inner: Inner::Service(Some(service)),
31        }
32    }
33
34    /// Creates a layer that wraps services with [`SpawnReady`].
35    pub fn layer() -> SpawnReadyLayer {
36        SpawnReadyLayer::default()
37    }
38}
39
40impl<S> Drop for SpawnReady<S> {
41    fn drop(&mut self) {
42        if let Inner::Future(ref mut task) = self.inner {
43            task.abort();
44        }
45    }
46}
47
48impl<S, Req> Service<Req> for SpawnReady<S>
49where
50    Req: 'static,
51    S: Service<Req> + Send + 'static,
52    S::Error: Into<BoxError>,
53{
54    type Response = S::Response;
55    type Error = BoxError;
56    type Future = ResponseFuture<S::Future, S::Error>;
57
58    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
59        loop {
60            self.inner = match self.inner {
61                Inner::Service(ref mut svc) => {
62                    if let Poll::Ready(r) = svc.as_mut().expect("illegal state").poll_ready(cx) {
63                        return Poll::Ready(r.map_err(Into::into));
64                    }
65
66                    let svc = svc.take().expect("illegal state");
67                    let rx =
68                        tokio::spawn(svc.ready_oneshot().map_err(Into::into).in_current_span());
69                    Inner::Future(rx)
70                }
71                Inner::Future(ref mut fut) => {
72                    let svc = ready!(Pin::new(fut).poll(cx))??;
73                    Inner::Service(Some(svc))
74                }
75            }
76        }
77    }
78
79    fn call(&mut self, request: Req) -> Self::Future {
80        match self.inner {
81            Inner::Service(Some(ref mut svc)) => {
82                ResponseFuture::new(svc.call(request).map_err(Into::into))
83            }
84            _ => unreachable!("poll_ready must be called"),
85        }
86    }
87}
tower/spawn_ready/service.rs

tower/spawn_ready/
service.rs
