fix(prometheus): make remote write timeout retryable (#8639)

* fix(prometheus): make remote write timeout retryable

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(prometheus): enforce pending row timeout budget

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(prometheus): skip pending-row timeout fallback when batcher is disabled

PendingRowsBatcher::try_new returns None when max_batch_rows,
max_concurrent_flushes, worker_channel_capacity or max_inflight_requests
is zero, meaning remote writes bypass batching entirely. The timeout
fallback predicate now mirrors these enablement conditions so the HTTP
timeout is not raised when no request can wait for a pending-row flush.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(prometheus): skip pending-row timeout fallback in async batch mode

With PENDING_ROWS_BATCH_SYNC=false, pending-row submissions return right
after enqueue and no request waits for a flush, so raising the global
HTTP timeout only delays unrelated routes. Export the batch sync mode
predicate from the servers crate and consult it in the frontend's
effective_http_options so the fallback is skipped in asynchronous mode.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

---------

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

config/config.md

@@ -27,7 +27,7 @@
 | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute the runtime for global write operations. |
 | `http` | -- | -- | The HTTP server options. |
 | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
-| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout. |
+| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.<br/>When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the<br/>`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. |
 | `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
 | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>This allows browser to access http APIs without CORS restrictions |
 | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |
@@ -251,7 +251,7 @@
 | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute the runtime for global write operations. |
 | `http` | -- | -- | The HTTP server options. |
 | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. |
-| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout. |
+| `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.<br/>When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the<br/>`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. |
 | `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
 | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>This allows browser to access http APIs without CORS restrictions |
 | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. |

config/frontend.example.toml

@@ -49,6 +49,8 @@ default_column_prefix = "greptime"
 ## The address to bind the HTTP server.
 addr = "127.0.0.1:4000"
 ## HTTP request timeout. Set to 0 to disable timeout.
+## When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
+## `prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value.
 timeout = "0s"
 ## HTTP request body limit.
 ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.

config/standalone.example.toml

@@ -62,6 +62,8 @@ max_concurrent_queries = 0
 ## The address to bind the HTTP server.
 addr = "127.0.0.1:4000"
 ## HTTP request timeout. Set to 0 to disable timeout.
+## When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
+## `prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value.
 timeout = "0s"
 ## HTTP request body limit.
 ## The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.

src/frontend/src/server.rs

@@ -14,6 +14,7 @@
 
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::time::Duration;
 
 use auth::UserProviderRef;
 use axum::extract::{Request, State};
@@ -21,7 +22,7 @@ use axum::middleware::Next;
 use axum::response::IntoResponse;
 use common_base::Plugins;
 use common_config::Configurable;
-use common_telemetry::info;
+use common_telemetry::{info, warn};
 use meta_client::MetaClientOptions;
 use servers::error::Error as ServerError;
 use servers::grpc::builder::GrpcServerBuilder;
@@ -32,12 +33,12 @@ use servers::grpc::{GrpcOptions, GrpcServer};
 use servers::http::event::LogValidatorRef;
 use servers::http::result::error_result::ErrorResponse;
 use servers::http::utils::router::RouterConfigurator;
-use servers::http::{HttpServer, HttpServerBuilder};
+use servers::http::{HttpOptions, HttpServer, HttpServerBuilder};
 use servers::interceptor::LogIngestInterceptorRef;
 use servers::metrics_handler::MetricsHandler;
 use servers::mysql::server::{MysqlServer, MysqlSpawnConfig, MysqlSpawnRef};
 use servers::otel_arrow::OtelArrowServiceHandler;
-use servers::pending_rows_batcher::PendingRowsBatcher;
+use servers::pending_rows_batcher::{PendingRowsBatcher, pending_rows_batch_sync_enabled};
 use servers::postgres::PostgresServer;
 use servers::request_memory_limiter::ServerMemoryLimiter;
 use servers::server::{Server, ServerHandlers};
@@ -102,7 +103,7 @@ where
         opts: &FrontendOptions,
         request_memory_limiter: ServerMemoryLimiter,
     ) -> HttpServerBuilder {
-        let mut builder = HttpServerBuilder::new(opts.http.clone())
+        let mut builder = HttpServerBuilder::new(effective_http_options(opts))
             .with_memory_limiter(request_memory_limiter)
             .with_sql_handler(self.instance.clone());
 
@@ -396,6 +397,133 @@ where
     }
 }
 
+fn effective_http_options(opts: &FrontendOptions) -> HttpOptions {
+    effective_http_options_with_sync(opts, pending_rows_batch_sync_enabled())
+}
+
+fn effective_http_options_with_sync(opts: &FrontendOptions, batch_sync: bool) -> HttpOptions {
+    let mut http = opts.http.clone();
+    let flush_interval = opts.prom_store.pending_rows_flush_interval;
+    let fallback_timeout = flush_interval.saturating_add(Duration::from_secs(1));
+    // In asynchronous batch mode submissions return right after enqueue and
+    // no request waits for a pending-row flush, so the timeout must not be
+    // raised either.
+    if !opts.prom_store.pending_rows_batching_enabled()
+        || !batch_sync
+        || http.timeout.is_zero()
+        || http.timeout > fallback_timeout
+    {
+        return http;
+    }
+
+    let configured_timeout = http.timeout;
+    http.timeout = fallback_timeout;
+    warn!(
+        ?configured_timeout,
+        ?flush_interval,
+        ?fallback_timeout,
+        "HTTP request timeout is not longer than the pending-row timeout fallback; using the fallback"
+    );
+    http
+}
+
 fn parse_addr(addr: &str) -> Result<SocketAddr> {
     addr.parse().context(error::ParseAddrSnafu { addr })
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::*;
+
+    #[test]
+    fn test_effective_http_timeout_for_pending_rows() {
+        let cases = [
+            ("disabled timeout", 0, 5000, true, true, 0),
+            ("disabled prom store", 1000, 5000, false, true, 1000),
+            ("disabled metric engine", 1000, 5000, true, false, 1000),
+            ("disabled batching", 1000, 0, true, true, 1000),
+            ("timeout below flush interval", 4000, 5000, true, true, 6000),
+            (
+                "timeout equals flush interval",
+                5000,
+                5000,
+                true,
+                true,
+                6000,
+            ),
+            ("timeout below fallback", 5500, 5000, true, true, 6000),
+            ("timeout equals fallback", 6000, 5000, true, true, 6000),
+            ("timeout above fallback", 7000, 5000, true, true, 7000),
+        ];
+
+        for (name, timeout, flush_interval, enable, with_metric_engine, expected) in cases {
+            let mut opts = FrontendOptions::default();
+            opts.http.timeout = Duration::from_millis(timeout);
+            opts.prom_store.pending_rows_flush_interval = Duration::from_millis(flush_interval);
+            opts.prom_store.enable = enable;
+            opts.prom_store.with_metric_engine = with_metric_engine;
+
+            assert_eq!(
+                Duration::from_millis(expected),
+                effective_http_options_with_sync(&opts, true).timeout,
+                "{name}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_effective_http_timeout_skips_fallback_in_async_batch_mode() {
+        // With `PENDING_ROWS_BATCH_SYNC=false`, submissions return right after
+        // enqueue and no request waits for a pending-row flush, so the
+        // timeout must not be raised.
+        let mut opts = FrontendOptions::default();
+        opts.http.timeout = Duration::from_millis(1000);
+        opts.prom_store.pending_rows_flush_interval = Duration::from_millis(5000);
+
+        assert_eq!(
+            Duration::from_millis(1000),
+            effective_http_options_with_sync(&opts, false).timeout,
+        );
+        assert_eq!(
+            Duration::from_millis(6000),
+            effective_http_options_with_sync(&opts, true).timeout,
+        );
+    }
+
+    #[test]
+    fn test_effective_http_timeout_skips_fallback_when_batcher_disabled() {
+        // Mirrors the conditions under which `PendingRowsBatcher::try_new`
+        // returns `None`; in these cases no request can wait for a pending-row
+        // flush, so the timeout must not be raised.
+        type KnobMutator = fn(&mut FrontendOptions);
+        let cases: [(&str, KnobMutator); 4] = [
+            ("zero max_batch_rows", |opts| {
+                opts.prom_store.max_batch_rows = 0
+            }),
+            ("zero max_concurrent_flushes", |opts| {
+                opts.prom_store.max_concurrent_flushes = 0
+            }),
+            ("zero worker_channel_capacity", |opts| {
+                opts.prom_store.worker_channel_capacity = 0
+            }),
+            ("zero max_inflight_requests", |opts| {
+                opts.prom_store.max_inflight_requests = 0
+            }),
+        ];
+
+        for (name, disable_batcher) in cases {
+            let mut opts = FrontendOptions::default();
+            opts.http.timeout = Duration::from_millis(1000);
+            opts.prom_store.pending_rows_flush_interval = Duration::from_millis(5000);
+            disable_batcher(&mut opts);
+
+            assert_eq!(
+                Duration::from_millis(1000),
+                effective_http_options_with_sync(&opts, true).timeout,
+                "{name}"
+            );
+        }
+    }
+}

src/frontend/src/service_config/prom_store.rs

@@ -56,6 +56,22 @@ fn default_flow_notification_queue_capacity() -> NonZeroUsize {
     NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN)
 }
 
+impl PromStoreOptions {
+    /// Returns whether the pending rows batcher can be enabled with these
+    /// options. Mirrors the enablement conditions of
+    /// `PendingRowsBatcher::try_new` in the servers crate, which returns
+    /// `None` when any of these knobs is zero.
+    pub fn pending_rows_batching_enabled(&self) -> bool {
+        self.enable
+            && self.with_metric_engine
+            && !self.pending_rows_flush_interval.is_zero()
+            && self.max_batch_rows > 0
+            && self.max_concurrent_flushes > 0
+            && self.worker_channel_capacity > 0
+            && self.max_inflight_requests > 0
+    }
+}
+
 impl Default for PromStoreOptions {
     fn default() -> Self {
         Self {

src/servers/src/http.rs

@@ -836,7 +836,12 @@ impl HttpServer {
     /// Callers should call this method after [HttpServer::make_app()].
     pub fn build(&self, router: Router) -> Result<Router> {
         let timeout_layer = if self.options.timeout != Duration::default() {
-            Some(ServiceBuilder::new().layer(DynamicTimeoutLayer::new(self.options.timeout)))
+            Some(
+                ServiceBuilder::new().layer(
+                    DynamicTimeoutLayer::new(self.options.timeout)
+                        .with_status_code_fn(Self::request_timeout_status_code),
+                ),
+            )
         } else {
             info!("HTTP server timeout is disabled");
             None
@@ -941,6 +946,14 @@ impl HttpServer {
             ))
     }
 
+    fn request_timeout_status_code(request: &Request) -> HttpStatusCode {
+        if request.uri().path() == "/v1/prometheus/write" {
+            HttpStatusCode::GATEWAY_TIMEOUT
+        } else {
+            HttpStatusCode::REQUEST_TIMEOUT
+        }
+    }
+
     fn route_metrics<S>(metrics_handler: MetricsHandler) -> Router<S> {
         Router::new()
             .route("/metrics", routing::get(handler::metrics))
@@ -1366,9 +1379,8 @@ mod test {
 
     use arrow_ipc::reader::StreamReader;
     use arrow_schema::DataType;
-    use axum::handler::Handler;
     use axum::http::StatusCode;
-    use axum::routing::get;
+    use axum::routing::{get, post};
     use common_query::{Output, OutputData};
     use common_recordbatch::RecordBatches;
     use datafusion_expr::LogicalPlan;
@@ -1429,10 +1441,6 @@ mod test {
         }
     }
 
-    fn timeout() -> DynamicTimeoutLayer {
-        DynamicTimeoutLayer::new(Duration::from_millis(10))
-    }
-
     async fn forever() {
         pending().await
     }
@@ -1446,10 +1454,11 @@ mod test {
         let server = HttpServerBuilder::new(options)
             .with_sql_handler(instance.clone())
             .build();
-        server.build(server.make_app()).unwrap().route(
-            "/test/timeout",
-            get(forever.layer(ServiceBuilder::new().layer(timeout()))),
-        )
+        let app = server
+            .make_app()
+            .route("/test/timeout", get(forever))
+            .route("/v1/prometheus/write", post(forever));
+        server.build(app).unwrap()
     }
 
     #[tokio::test]
@@ -1603,11 +1612,18 @@ mod test {
         common_telemetry::init_default_ut_logging();
 
         let (tx, _rx) = mpsc::channel(100);
-        let app = make_test_app(tx);
+        let options = HttpOptions {
+            timeout: Duration::from_millis(10),
+            ..Default::default()
+        };
+        let app = make_test_app_custom(tx, options);
         let client = TestClient::new(app).await;
         let res = client.get("/test/timeout").send().await;
         assert_eq!(res.status(), StatusCode::REQUEST_TIMEOUT);
 
+        let res = client.post("/v1/prometheus/write").send().await;
+        assert_eq!(res.status(), StatusCode::GATEWAY_TIMEOUT);
+
         let now = Instant::now();
         let res = client
             .get("/test/timeout")

src/servers/src/http/timeout.rs

@@ -39,11 +39,16 @@ pub struct ResponseFuture<T> {
     inner: T,
     #[pin]
     sleep: Sleep,
+    status_code: StatusCode,
 }
 
 impl<T> ResponseFuture<T> {
-    pub(crate) fn new(inner: T, sleep: Sleep) -> Self {
-        ResponseFuture { inner, sleep }
+    pub(crate) fn new(inner: T, sleep: Sleep, status_code: StatusCode) -> Self {
+        ResponseFuture {
+            inner,
+            sleep,
+            status_code,
+        }
     }
 }
 
@@ -58,7 +63,7 @@ where
 
         if this.sleep.poll(cx).is_ready() {
             let mut res = Response::default();
-            *res.status_mut() = StatusCode::REQUEST_TIMEOUT;
+            *res.status_mut() = *this.status_code;
             return Poll::Ready(Ok(res));
         }
 
@@ -72,12 +77,22 @@ where
 #[derive(Debug, Clone)]
 pub struct DynamicTimeoutLayer {
     default_timeout: Duration,
+    status_code_fn: fn(&Request<Body>) -> StatusCode,
 }
 
 impl DynamicTimeoutLayer {
     /// Create a timeout from a duration
     pub fn new(default_timeout: Duration) -> Self {
-        DynamicTimeoutLayer { default_timeout }
+        DynamicTimeoutLayer {
+            default_timeout,
+            status_code_fn: |_| StatusCode::REQUEST_TIMEOUT,
+        }
+    }
+
+    /// Sets a function that selects the timeout response status for each request.
+    pub fn with_status_code_fn(mut self, status_code_fn: fn(&Request<Body>) -> StatusCode) -> Self {
+        self.status_code_fn = status_code_fn;
+        self
     }
 }
 
@@ -85,7 +100,7 @@ impl<S> Layer<S> for DynamicTimeoutLayer {
     type Service = DynamicTimeout<S>;
 
     fn layer(&self, service: S) -> Self::Service {
-        DynamicTimeout::new(service, self.default_timeout)
+        DynamicTimeout::new(service, self.default_timeout, self.status_code_fn)
     }
 }
 
@@ -94,14 +109,20 @@ impl<S> Layer<S> for DynamicTimeoutLayer {
 pub struct DynamicTimeout<S> {
     inner: S,
     default_timeout: Duration,
+    status_code_fn: fn(&Request<Body>) -> StatusCode,
 }
 
 impl<S> DynamicTimeout<S> {
     /// Create a new [`DynamicTimeout`] with the given timeout
-    pub fn new(inner: S, default_timeout: Duration) -> Self {
+    pub fn new(
+        inner: S,
+        default_timeout: Duration,
+        status_code_fn: fn(&Request<Body>) -> StatusCode,
+    ) -> Self {
         DynamicTimeout {
             inner,
             default_timeout,
+            status_code_fn,
         }
     }
 }
@@ -122,6 +143,7 @@ where
     }
 
     fn call(&mut self, request: Request<Body>) -> Self::Future {
+        let status_code = (self.status_code_fn)(&request);
         let timeout = request
             .headers()
             .get(GREPTIME_DB_HEADER_TIMEOUT)
@@ -137,10 +159,10 @@ where
         if timeout.is_zero() {
             // 30 years. See `Instant::far_future`.
             let far_future = Instant::now() + Duration::from_secs(86400 * 365 * 30);
-            ResponseFuture::new(response, tokio::time::sleep_until(far_future))
+            ResponseFuture::new(response, tokio::time::sleep_until(far_future), status_code)
         } else {
             let sleep = tokio::time::sleep(timeout);
-            ResponseFuture::new(response, sleep)
+            ResponseFuture::new(response, sleep, status_code)
         }
     }
 }

src/servers/src/pending_rows_batcher.rs

@@ -64,6 +64,21 @@ use crate::prom_row_builder::{
 const PHYSICAL_TABLE_KEY: &str = "physical_table";
 /// Whether wait for ingestion result before reply to client.
 const PENDING_ROWS_BATCH_SYNC_ENV: &str = "PENDING_ROWS_BATCH_SYNC";
+
+/// Returns whether pending-row batch submissions wait for the flush result
+/// before replying to the client (synchronous mode), controlled by the
+/// `PENDING_ROWS_BATCH_SYNC` environment variable and defaulting to `true`.
+///
+/// Callers that reason about how long a remote write request may block (e.g.
+/// the frontend HTTP timeout fallback) must consult this instead of
+/// duplicating the env lookup.
+pub fn pending_rows_batch_sync_enabled() -> bool {
+    std::env::var(PENDING_ROWS_BATCH_SYNC_ENV)
+        .ok()
+        .as_deref()
+        .and_then(|v| v.parse::<bool>().ok())
+        .unwrap_or(true)
+}
 const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3;
 const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3;
 const MAX_CONCURRENT_FLOW_NOTIFICATIONS: usize = 8;
@@ -372,11 +387,7 @@ impl PendingRowsBatcher {
         }
 
         let (shutdown, _) = broadcast::channel(1);
-        let pending_rows_batch_sync = std::env::var(PENDING_ROWS_BATCH_SYNC_ENV)
-            .ok()
-            .as_deref()
-            .and_then(|v| v.parse::<bool>().ok())
-            .unwrap_or(true);
+        let pending_rows_batch_sync = pending_rows_batch_sync_enabled();
         let workers = Arc::new(DashMap::new());
         PENDING_WORKERS.set(workers.len() as i64);
         let (flow_notification_tx, flow_notification_rx) =