Add a timeout to graceful termination (#846)

<!-- ELLIPSIS_HIDDEN -->



> [!IMPORTANT]
> Add a 2-second timeout for graceful server termination in
`ControllerServer` and remove unused imports in `admin.rs`.
> 
>   - **Behavior**:
> - Add `TERMINATE_TIMEOUT_DURATION` constant in `controller/mod.rs` to
set a 2-second timeout for graceful server termination.
> - Modify `terminate()` in `ControllerServer` to use
`tokio::time::timeout` for server shutdown.
> - Log a warning if the server does not terminate gracefully within the
timeout.
>   - **Misc**:
>     - Remove unused import `Stream` and `StreamExt` from `admin.rs`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=jamsocket%2Fplane&utm_source=github&utm_medium=referral)<sup>
for 02566ef. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->

Author: paulgb

plane/src/admin.rs

@@ -12,7 +12,6 @@ use crate::{
 use chrono::Duration;
 use clap::{Parser, Subcommand};
 use colored::Colorize;
-use futures_util::{Stream, StreamExt};
 use std::path::PathBuf;
 use url::Url;
 

plane/src/controller/mod.rs

@@ -58,6 +58,11 @@ mod forward_auth;
 mod proxy;
 mod terminate;
 
+/// How long to wait for the server to terminate gracefully before forcing it to shut down.
+/// We want to keep this just high enough to serve short requests. Long-lived requests
+/// will continue blocking indefinitely, which is why we need a timeout.
+const TERMINATE_TIMEOUT_DURATION: std::time::Duration = std::time::Duration::from_secs(2);
+
 #[derive(Serialize, Deserialize, Debug)]
 pub struct StatusResponse {
     pub status: String,
@@ -295,6 +300,7 @@ impl ControllerServer {
         self.heartbeat_handle.terminate().await;
 
         // Begin graceful shutdown of server.
+        tracing::info!("Initiating graceful shutdown of server");
         let Some(graceful_terminate_sender) = self.graceful_terminate_sender.take() else {
             return;
         };
@@ -307,16 +313,19 @@ impl ControllerServer {
                 return;
             };
 
-            match server_handle.await {
-                Ok(Ok(())) => {
+            match tokio::time::timeout(TERMINATE_TIMEOUT_DURATION, server_handle).await {
+                Ok(Ok(Ok(()))) => {
                     tracing::info!("Server gracefully terminated");
                 }
-                Ok(Err(err)) => {
+                Ok(Ok(Err(err))) => {
                     tracing::error!(?err, "Server error");
                 }
-                Err(err) => {
+                Ok(Err(err)) => {
                     tracing::error!(?err, "Server error");
                 }
+                Err(_) => {
+                    tracing::warn!("Server did not terminate gracefully in time, forcing shutdown");
+                }
             }
         }
     }