saluki_io/net/server/
grpc.rs

1use std::{convert::Infallible, time::Duration};
2
3use async_trait::async_trait;
4use http::Request;
5use saluki_common::sync::shutdown::ShutdownHandle;
6use saluki_core::runtime::{InitializationError, ShutdownStrategy, Supervisable, SupervisorFuture};
7use saluki_error::ErrorContext as _;
8use tokio::{pin, select, sync::oneshot, time::timeout};
9use tonic::{
10    body::Body,
11    server::NamedService,
12    service::Routes,
13    transport::{server::TcpIncoming, Server},
14};
15use tower::Service;
16use tracing::warn;
17
18#[cfg(unix)]
19use crate::net::unix::{ensure_unix_socket_free, set_unix_socket_write_only};
20use crate::net::ListenAddress;
21
22/// A gRPC server.
23///
24/// Allows serving multiple gRPC services from a single endpoint.
25///
26/// This type is a thin wrapper over helper types from `tonic` and `axum`, and principally is meant to provide an opaque
27/// gRPC server implementation that operates correctly when run under supervision. As such, this type can't be manually
28/// served: it is only usable by adding it to a supervisor.
29///
30/// # Supervision
31///
32/// The listen address is bound during initialization, so a failure to bind is raised before the supervised worker
33/// starts running, and a restart rebinds.
34///
35/// The server will attempt to gracefully shutdown existing connections when the parent supervisor signals shutdown.
36/// This will cause the worker to utilize the maximum allowable grace period during shutdown: it will attempt to take as
37/// long as necessary to gracefully shutdown existing connections, bounded only by the parent supervisor.
38pub struct GrpcServer {
39    listen_addr: ListenAddress,
40    routes: Option<Routes>,
41    graceful_shutdown_timeout: Option<Duration>,
42}
43
44impl GrpcServer {
45    /// Creates an empty server with no attached services, configured to listen on the given address.
46    pub fn new(listen_addr: ListenAddress) -> Self {
47        Self {
48            listen_addr,
49            routes: None,
50            graceful_shutdown_timeout: None,
51        }
52    }
53
54    /// Sets the graceful shutdown timeout for this server.
55    ///
56    /// During shutdown, the server will for all in-flight connections to complete before ultimately completing itself.
57    /// When no timeout is specified, this will lead to the worker taking the maximum allowable time to shutdown if
58    /// connections are blocked or otherwise "stuck." Setting an explicit graceful shutdown timeout will cause the
59    /// worker to bound how long it waits for in-flight connections to shutdown before forcefully completing and moving
60    /// on.
61    ///
62    /// Defaults to no timeout (wait as long as allowed).
63    pub fn with_graceful_shutdown_timeout(mut self, timeout: Duration) -> Self {
64        self.graceful_shutdown_timeout = Some(timeout);
65        self
66    }
67
68    /// Adds a new service to this server.
69    pub fn add_service<S>(mut self, svc: S) -> Self
70    where
71        S: Service<Request<Body>, Error = Infallible> + NamedService + Clone + Send + Sync + 'static,
72        S::Response: axum::response::IntoResponse,
73        S::Future: Send + 'static,
74    {
75        let routes = self.routes.take().unwrap_or_default().add_service(svc);
76
77        Self {
78            routes: Some(routes),
79            ..self
80        }
81    }
82}
83
84#[async_trait]
85impl Supervisable for GrpcServer {
86    fn name(&self) -> &str {
87        "grpc_server"
88    }
89
90    fn shutdown_strategy(&self) -> ShutdownStrategy {
91        // Utilize the maximum allowable grace period to give connections a chance to gracefully shutdown.
92        ShutdownStrategy::Graceful(Duration::MAX)
93    }
94
95    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
96        let routes = self.routes.clone().unwrap_or_default();
97        let shutdown_timeout = self.graceful_shutdown_timeout.unwrap_or(Duration::MAX);
98
99        match &self.listen_addr {
100            ListenAddress::Tcp(addr) => {
101                let listener = TcpIncoming::bind(*addr)
102                    .with_error_context(|| format!("Failed to bind listener for gRPC server ({}).", addr))?;
103
104                Ok(Box::pin(async move {
105                    let (drain_tx, drain_rx) = oneshot::channel();
106                    let serve = Server::default().serve_with_incoming_shutdown(routes, listener, async move {
107                        let _ = drain_rx.await;
108                    });
109
110                    pin!(serve, process_shutdown);
111
112                    select! {
113                        result = &mut serve => result.error_context("Failed to serve gRPC server."),
114
115                        _ = &mut process_shutdown => {
116                            let _ = drain_tx.send(());
117
118                            match timeout(shutdown_timeout, serve).await {
119                                Ok(Ok(())) => Ok(()),
120                                Ok(Err(e)) => Err(e).error_context("Failed to serve gRPC server."),
121                                Err(_) => {
122                                    warn!("Failed to gracefully drain gRPC connections.");
123                                    Ok(())
124                                },
125                            }
126                        },
127                    }
128                }))
129            }
130            #[cfg(unix)]
131            ListenAddress::Unix(path) => {
132                let path = path.clone();
133                ensure_unix_socket_free(&path)
134                    .await
135                    .with_error_context(|| format!("Failed to clear gRPC Unix socket '{}'.", path.display()))?;
136                let listener = tokio::net::UnixListener::bind(&path)
137                    .with_error_context(|| format!("Failed to bind gRPC Unix listener on '{}'.", path.display()))?;
138                set_unix_socket_write_only(&path).await.with_error_context(|| {
139                    format!("Failed to set permissions on gRPC Unix socket '{}'.", path.display())
140                })?;
141                let incoming = tokio_stream::wrappers::UnixListenerStream::new(listener);
142
143                Ok(Box::pin(async move {
144                    let (drain_tx, drain_rx) = oneshot::channel();
145                    let serve = Server::default().serve_with_incoming_shutdown(routes, incoming, async move {
146                        let _ = drain_rx.await;
147                    });
148
149                    pin!(serve, process_shutdown);
150
151                    select! {
152                        result = &mut serve => result.error_context("Failed to serve gRPC server."),
153
154                        _ = &mut process_shutdown => {
155                            let _ = drain_tx.send(());
156
157                            match timeout(shutdown_timeout, serve).await {
158                                Ok(Ok(())) => Ok(()),
159                                Ok(Err(e)) => Err(e).error_context("Failed to serve gRPC server."),
160                                Err(_) => {
161                                    warn!("Failed to gracefully drain gRPC connections.");
162                                    Ok(())
163                                },
164                            }
165                        },
166                    }
167                }))
168            }
169            _ => Err(InitializationError::Failed {
170                source: saluki_error::generic_error!("gRPC endpoint must be a TCP or Unix address."),
171            }),
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use std::net::{SocketAddr, TcpListener as StdTcpListener};
179
180    use saluki_common::sync::shutdown::ShutdownCoordinator;
181    use tokio::io::AsyncWriteExt as _;
182    use tokio::net::TcpStream;
183    use tokio::time::timeout;
184
185    use super::*;
186
187    /// Bound on any server await in these tests, so a hang fails rather than stalling the suite.
188    const TEST_TIMEOUT: Duration = Duration::from_secs(10);
189
190    /// Reserves a loopback port and releases it, yielding an address a server can bind.
191    fn free_local_addr() -> SocketAddr {
192        let listener = StdTcpListener::bind("127.0.0.1:0").expect("should bind an ephemeral port");
193        let addr = listener.local_addr().expect("should have a local address");
194        drop(listener);
195        addr
196    }
197
198    #[tokio::test]
199    async fn binds_during_initialization() {
200        // The listener is bound by `initialize`, so the port is taken before anything serves. That is what makes a bind
201        // failure a non-restartable initialization error rather than a runtime one.
202        let addr = free_local_addr();
203        let run = GrpcServer::new(ListenAddress::Tcp(addr))
204            .initialize(ShutdownHandle::noop())
205            .await
206            .expect("should initialize");
207
208        assert!(
209            StdTcpListener::bind(addr).is_err(),
210            "initialization should have bound {addr} before the worker future ran"
211        );
212
213        drop(run);
214    }
215
216    #[tokio::test]
217    async fn bind_failure_is_an_initialization_error() {
218        let addr = free_local_addr();
219        let _held = StdTcpListener::bind(addr).expect("should hold the address");
220
221        match GrpcServer::new(ListenAddress::Tcp(addr))
222            .initialize(ShutdownHandle::noop())
223            .await
224        {
225            Ok(_) => panic!("initialization should have failed to bind {addr}"),
226            Err(e) => {
227                let error = e.to_string();
228                assert!(error.contains("Failed to bind listener"), "unexpected error: {error}");
229            }
230        }
231    }
232
233    #[tokio::test]
234    async fn releases_its_port_once_the_worker_finishes() {
235        let addr = free_local_addr();
236        let mut coordinator = ShutdownCoordinator::default();
237        let run = GrpcServer::new(ListenAddress::Tcp(addr))
238            .initialize(coordinator.register())
239            .await
240            .expect("should initialize");
241
242        coordinator.shutdown();
243        timeout(TEST_TIMEOUT, run)
244            .await
245            .expect("server should stop on shutdown")
246            .expect("server should stop cleanly");
247
248        assert!(
249            StdTcpListener::bind(addr).is_ok(),
250            "the server should have released {addr} when its worker finished"
251        );
252    }
253
254    #[tokio::test]
255    async fn an_idle_peer_does_not_wedge_the_drain() {
256        // `tonic` waits for every connection to close and imposes no bound of its own, so a peer that connects and then
257        // does nothing would otherwise hold shutdown open indefinitely.
258        let addr = free_local_addr();
259        let mut coordinator = ShutdownCoordinator::default();
260        let run = GrpcServer::new(ListenAddress::Tcp(addr))
261            .with_graceful_shutdown_timeout(Duration::from_secs(1))
262            .initialize(coordinator.register())
263            .await
264            .expect("should initialize");
265        let run = tokio::spawn(run);
266
267        let mut stream = TcpStream::connect(addr).await.expect("should connect");
268        stream.write_all(b"PRI * HTTP/2.0\r\n").await.expect("should write");
269        stream.flush().await.expect("should flush");
270        tokio::time::sleep(Duration::from_millis(100)).await;
271
272        coordinator.shutdown();
273        timeout(TEST_TIMEOUT, run)
274            .await
275            .expect("server should finish draining rather than waiting on an idle peer")
276            .expect("server task should not panic")
277            .expect("server should stop cleanly");
278    }
279}
saluki_io/net/server/grpc.rs

saluki_io/net/server/
grpc.rs
