viceroy_lib/
service.rs

1//! Service types.
2
3use {
4    crate::{body::Body, execute::ExecuteCtx, Error},
5    futures::future::{self, Ready},
6    hyper::{
7        http::{Request, Response},
8        server::conn::AddrStream,
9        service::Service,
10    },
11    std::{
12        convert::Infallible,
13        future::Future,
14        net::SocketAddr,
15        pin::Pin,
16        sync::Arc,
17        task::{self, Poll},
18    },
19    tracing::{event, Level},
20};
21
22/// A Viceroy service uses a Wasm module and a handler function to respond to HTTP requests.
23///
24/// This service type is used to compile a Wasm [`Module`][mod], and perform the actions necessary
25/// to initialize a [`Server`][serv] and bind the service to a local port.
26///
27/// Each time a connection is received, a [`RequestService`][req-svc] will be created, to
28/// instantiate the module and return a [`Response`][resp].
29///
30/// [mod]: https://docs.rs/wasmtime/latest/wasmtime/struct.Module.html
31/// [req-svc]: struct.RequestService.html
32/// [resp]: https://docs.rs/http/latest/http/response/struct.Response.html
33/// [serv]: https://docs.rs/hyper/latest/hyper/server/struct.Server.html
34pub struct ViceroyService {
35    ctx: Arc<ExecuteCtx>,
36}
37
38impl ViceroyService {
39    /// Create a new Viceroy service, using the given handler function and module path.
40    ///
41    /// # Example
42    ///
43    /// ```no_run
44    /// # use std::collections::HashSet;
45    /// use viceroy_lib::{Error, ExecuteCtx, ProfilingStrategy, ViceroyService};
46    /// # fn f() -> Result<(), Error> {
47    /// let adapt_core_wasm = false;
48    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None, HashSet::new(), None, Default::default(), adapt_core_wasm)?;
49    /// let svc = ViceroyService::new(ctx);
50    /// # Ok(())
51    /// # }
52    /// ```
53    pub fn new(ctx: Arc<ExecuteCtx>) -> Self {
54        Self { ctx }
55    }
56
57    /// An internal helper, create a [`RequestService`](struct.RequestService.html).
58    fn make_service(&self, remote: &AddrStream) -> RequestService {
59        RequestService::new(self.ctx.clone(), remote)
60    }
61
62    /// Bind this service to the given address and start serving responses.
63    ///
64    /// This will consume the service, using it to start a server that will execute the given module
65    /// each time a new request is sent. This function will only return if an error occurs.
66    // FIXME KTM 2020-06-22: Once `!` is stabilized, this should be `Result<!, hyper::Error>`.
67    pub async fn serve(self, addr: SocketAddr) -> Result<(), hyper::Error> {
68        let server = hyper::Server::bind(&addr).serve(self);
69        event!(Level::INFO, "Listening on http://{}", server.local_addr());
70        server.await?;
71        Ok(())
72    }
73}
74
75impl<'addr> Service<&'addr AddrStream> for ViceroyService {
76    type Response = RequestService;
77    type Error = Infallible;
78    type Future = Ready<Result<Self::Response, Self::Error>>;
79
80    fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
81        Poll::Ready(Ok(()))
82    }
83
84    fn call(&mut self, addr: &'addr AddrStream) -> Self::Future {
85        future::ok(self.make_service(addr))
86    }
87}
88
89/// A request service is responsible for handling a single request.
90///
91/// Most importantly, this structure implements the [`tower::Service`][service] trait, which allows
92/// it to be dispatched by [`ViceroyService`][viceroy] to handle a single request.
93///
94/// This object does not need to be used directly; users most likely should use
95/// [`ViceroyService::serve`][serve] to bind a service to a port, or
96/// [`ExecuteCtx::handle_request`][handle_request] to generate a response for a request when writing
97/// test cases.
98///
99/// [handle_request]: ../execute/struct.ExecuteCtx.html#method.handle_request
100/// [serve]: struct.ViceroyService.html#method.serve
101/// [service]: https://docs.rs/tower/latest/tower/trait.Service.html
102/// [viceroy]: struct.ViceroyService.html
103#[derive(Clone)]
104pub struct RequestService {
105    ctx: Arc<ExecuteCtx>,
106    local_addr: SocketAddr,
107    remote_addr: SocketAddr,
108}
109
110impl RequestService {
111    /// Create a new request service.
112    fn new(ctx: Arc<ExecuteCtx>, addr: &AddrStream) -> Self {
113        let local_addr = addr.local_addr();
114        let remote_addr = addr.remote_addr();
115
116        Self {
117            ctx,
118            local_addr,
119            remote_addr,
120        }
121    }
122}
123
124impl Service<Request<hyper::Body>> for RequestService {
125    type Response = Response<Body>;
126    type Error = Error;
127    #[allow(clippy::type_complexity)]
128    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
129
130    fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
131        Poll::Ready(Ok(()))
132    }
133
134    /// Process the request and return the response asynchronously.
135    fn call(&mut self, req: Request<hyper::Body>) -> Self::Future {
136        // Request handling currently takes ownership of the context, which is cheaply cloneable.
137        let ctx = self.ctx.clone();
138        let local = self.local_addr;
139        let remote = self.remote_addr;
140
141        // Now, use the execution context to handle the request.
142        Box::pin(async move {
143            ctx.handle_request(req, local, remote)
144                .await
145                .map(|result| result.0)
146        })
147    }
148}