server: add a global request timeout

ENVIRONMENT_VARIABLES.md

@@ -36,6 +36,7 @@ The diom server accepts the following environment variables:
 | `$DIOM_EPHEMERAL_DB_FILENAME` | Filename under the directory specified in `path`. |
 | `$DIOM_EPHEMERAL_DB_PATH` | Directory in which this database is stored |
 | `$DIOM_FSYNC_MODE` | When fsyncing, should we use fsync(2) or fdatasync(2) |
+| `$DIOM_GLOBAL_TIMEOUT_MS` | Hard timeout for all client requests |
 | `$DIOM_JWT_ALGORITHM` | JWT signature algorithm |
 | `$DIOM_JWT_AUDIENCE` | Expected `aud` values. When set, the token must contain one of these values in its `aud` claim. When absent, `aud` is not validated. |
 | `$DIOM_JWT_ISSUER` | Expected `iss` values. When set, the token's `iss` claim must match one of these values. When absent, `iss` is not validated. |

config.defaults.toml

@@ -31,6 +31,9 @@ environment = "dev"
 # When fsyncing, should we use fsync(2) or fdatasync(2)
 fsync_mode = "sync-data"
 
+# Hard timeout for all client requests
+global_timeout_ms = 1000
+
 # The address to listen on
 listen_address = "[::]:8624"
 

crates/diom-error/src/lib.rs

@@ -95,6 +95,10 @@ impl Error {
         })
     }
 
+    pub fn request_timeout() -> Self {
+        Self::new(ErrorType::RequestTimeout)
+    }
+
     pub fn shutting_down() -> Self {
         Self::new(ErrorType::ShuttingDown)
     }
@@ -128,6 +132,11 @@ impl Error {
                 Some("NOT_READY".to_owned()),
                 None,
             ),
+            ErrorType::RequestTimeout => (
+                StatusCode::GATEWAY_TIMEOUT,
+                Some("REQUEST_TIMEOUT".to_owned()),
+                None,
+            ),
             ErrorType::ShuttingDown => (
                 StatusCode::SERVICE_UNAVAILABLE,
                 Some("SHUTTING_DOWN".to_owned()),
@@ -204,6 +213,11 @@ impl IntoResponse for Error {
                 MsgPackOrJson(json!({"code": "NOT_READY", "detail": message})),
             )
                 .into_response(),
+            ErrorType::RequestTimeout => (
+                StatusCode::GATEWAY_TIMEOUT,
+                MsgPackOrJson(json!({"code": "REQUEST_TIMEOUT", "detail": ""})),
+            )
+                .into_response(),
             ErrorType::ShuttingDown => (
                 StatusCode::SERVICE_UNAVAILABLE,
                 MsgPackOrJson(json!({"code": "SHUTTING_DOWN", "detail": "server shutting down"})),
@@ -277,8 +291,12 @@ pub enum ErrorType {
         detail: Option<String>,
     },
 
+    RequestTimeout,
+
     /// The operation cannot proceed because the server is not yet ready
-    NotReady { message: String },
+    NotReady {
+        message: String,
+    },
 
     /// The operation cannot proceed because the server is shutting down
     ShuttingDown,
@@ -294,6 +312,7 @@ impl fmt::Display for ErrorType {
             Self::Authentication(s) => write!(f, "authn {s}"),
             Self::Authorization(s) => write!(f, "authz {s}"),
             Self::ShuttingDown => write!(f, "shutting_down"),
+            Self::RequestTimeout => write!(f, "request timed out "),
             Self::Operation { detail, status, .. } => {
                 if let Some(detail) = detail {
                     detail.fmt(f)

crates/test-utils/src/server.rs

@@ -333,6 +333,7 @@ pub fn default_server_config(workdir: &Path) -> ConfigurationInner {
         fsync_mode: FsyncMode::SyncData,
         admin_token: Some(TEST_ADMIN_TOKEN.to_string()),
         jwt: Default::default(),
+        global_timeout: NonZeroDurationMs::from_millis(5000).unwrap(),
     }
 }
 

src/cfg/defaults.rs

@@ -111,3 +111,7 @@ pub(super) const fn background_cleanup_interval() -> NonZeroDurationMs {
 pub(super) fn default_database_size() -> MemorySize {
     MemorySize::Percent(20)
 }
+
+pub(super) const fn global_timeout() -> NonZeroDurationMs {
+    NonZeroDurationMs::from_secs(10).unwrap()
+}

src/cfg/mod.rs

@@ -612,6 +612,10 @@ pub struct ConfigurationInner {
     #[env_overridable(nest_with_prefix("OPENTELEMETRY"))]
     #[dumpable_config(nest)]
     pub opentelemetry: OpenTelemetryConfig,
+
+    #[serde(default = "defaults::global_timeout", rename = "global_timeout_ms")]
+    /// Hard timeout for all client requests
+    pub global_timeout: NonZeroDurationMs,
 }
 
 impl ConfigurationInner {

src/lib.rs

@@ -13,9 +13,9 @@ use std::{
 use aide::axum::ApiRouter;
 use axum::{
     Extension,
-    extract::DefaultBodyLimit,
+    extract::{DefaultBodyLimit, State},
     middleware,
-    response::IntoResponse as _,
+    response::{IntoResponse as _, Response},
     serve::{Listener, ListenerExt as _},
 };
 use diom_authorization::{AccessRuleList, Permissions};
@@ -24,6 +24,7 @@ use diom_error::Error;
 use diom_msgs::TopicPublishNotifier;
 use diom_proto::{InternalClient, InternalRequest, InternalRequestError};
 use fjall_utils::{Databases, ReadonlyDatabases};
+use futures_util::FutureExt;
 use opentelemetry::metrics::Meter;
 use serde::{Serialize, de::DeserializeOwned};
 use tokio::{
@@ -106,7 +107,7 @@ pub struct AppState {
     pub(crate) topic_publish_notifier: TopicPublishNotifier,
 }
 
-fn handle_panic(err: Box<dyn std::any::Any + Send + 'static>) -> axum::response::Response {
+fn handle_panic(err: Box<dyn std::any::Any + Send + 'static>) -> Response {
     if let Some(err) = err.downcast_ref::<String>() {
         tracing::error!(?err, "Unhandled panic");
     } else if let Some(err) = err.downcast_ref::<&'static str>() {
@@ -117,6 +118,15 @@ fn handle_panic(err: Box<dyn std::any::Any + Send + 'static>) -> axum::response:
     Error::internal("unhandled internal panic").into_response()
 }
 
+fn timeout_layer(
+    State(timeout): State<Duration>,
+    request: axum::extract::Request,
+    next: middleware::Next,
+) -> impl Future<Output = Response> {
+    tokio::time::timeout(timeout, next.run(request))
+        .map(|r| r.unwrap_or_else(|_| Error::request_timeout().into_response()))
+}
+
 async fn axum_tcp_listener(
     listener: Option<TcpListener>,
     listen_address: SocketAddr,
@@ -187,9 +197,11 @@ async fn run_internal(
     api_router: axum::Router,
     mut internal_req_rx: mpsc::Receiver<InternalRequest>,
     request_metrics: RequestMetrics,
+    cfg: Configuration,
 ) {
     let svc = api_router.layer((
         trace_layer(),
+        middleware::from_fn_with_state(cfg.global_timeout.into(), timeout_layer),
         CatchPanicLayer::custom(handle_panic),
         middleware::from_fn_with_state(
             request_metrics,
@@ -290,7 +302,7 @@ async fn fail_until_bootstrapped(
     path: axum::extract::MatchedPath,
     request: axum::extract::Request,
     next: middleware::Next,
-) -> axum::response::Response {
+) -> Response {
     let is_admin_route = path.as_str().starts_with("/api/v1.admin.cluster.");
     if !(is_admin_route || BOOTSTRAPPED.load(Ordering::Relaxed)) {
         return Error::not_ready("this node has not yet finished bootstrapping").into_response();
@@ -407,6 +419,7 @@ pub async fn run_with_listeners(
         api_router.clone(),
         internal_req_rx,
         request_metrics.with_connection_type(ConnectionType::Internal),
+        cfg.clone(),
     ));
 
     openapi::postprocess_spec(&mut openapi);
@@ -415,6 +428,7 @@ pub async fn run_with_listeners(
     let svc = router
         .layer((
             trace_layer(),
+            middleware::from_fn_with_state(cfg.global_timeout.into(), timeout_layer),
             CatchPanicLayer::custom(handle_panic),
             middleware::from_fn(core::cluster::middleware::capture_log_id),
             middleware::from_fn(diom_proto::capture_accept_hdr),