tower/spawn_ready/
service.rs

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

tower/spawn_ready/
service.rs
