Configuring Timeout for Graceful Shutdown

[yhl25]

Version

latest

Description

I'm currently working with tonic for implementing gRPC services in Rust, and I've encountered a situation where the server might not shut down gracefully due to long-running or unresponsive RPC connections.

I tried this code:

I have updated the bi-directional streaming example with custom shutdown logic, where the server shutdown is expected to be triggered based on a signal. This works as expected until a stream connection remains active indefinitely, causing the server to not shutdown.

async fn shutdown_signal(
    mut shutdown_on_err: mpsc::Receiver<()>,
) {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install SIGINT handler");
    };

    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    let shutdown_on_err_future = async {
        shutdown_on_err.recv().await;
    };

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
        _ = shutdown_on_err_future => {},
    }
    println!("shutdown signal received");
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
    let shutdown_signal = shutdown_signal(shutdown_rx);

    tokio::spawn(async move {
        sleep(Duration::from_secs(30)).await;
        println!("Shutting down the server after 30s");
        shutdown_tx.send(()).await.unwrap();
    });

    let server = EchoServer {};
    Server::builder()
        .add_service(pb::echo_server::EchoServer::new(server))
        .serve_with_shutdown("[::1]:50051".to_socket_addrs().unwrap().next().unwrap(), shutdown_signal)
        .await
        .unwrap();

    Ok(())
}

However, if some RPC is abnormally pending for a long time, the server has to wait. Can I set a timeout for the graceful shutdown?