diff --git a/.bleep b/.bleep index c89ef1de8..6eb87407b 100644 --- a/.bleep +++ b/.bleep @@ -1 +1 @@ -5a1cf681f7e2691687623b60387a88076493015f \ No newline at end of file +faf3f5b20af8cf4f1c8038ba48d44880674ca9ec \ No newline at end of file diff --git a/.cargo/audit.toml b/.cargo/audit.toml index 746e58a9d..e8a430531 100644 --- a/.cargo/audit.toml +++ b/.cargo/audit.toml @@ -1,7 +1,12 @@ +# Advisories against rustls-webpki 0.101.7, pulled in transitively by +# aws-sdk-s3-transfer-manager through the legacy rustls 0.21 chain used by +# dial9's worker-s3 feature. No patch exists in 0.101.x; not reachable in +# our usage because this is TLS client use only and does not parse CRLs. +# Remove once the upstream aws-s3-transfer-manager-rs fix ships. + [advisories] ignore = [ - # This came from the prometheus crate's protobuf encoder. - # We don't use the protobuf encoder, only the text one. - # https://rustsec.org/advisories/RUSTSEC-2024-0437 - "RUSTSEC-2024-0437", + "RUSTSEC-2026-0098", # rustls-webpki: URI name constraints incorrectly accepted + "RUSTSEC-2026-0099", # rustls-webpki: name constraints accepted for wildcard certs + "RUSTSEC-2026-0104", # rustls-webpki: reachable panic in CRL parsing ] diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 22a4c4581..adf2b014c 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -38,12 +38,17 @@ jobs: - name: Run cargo fmt run: cargo fmt --all -- --check + - name: Run cargo check + run: cargo check --workspace + - name: Run cargo test + if: matrix.toolchain != '1.84.0' run: cargo test --verbose --lib --bins --tests --no-fail-fast # Need to run doc tests separately. # (https://github.com/rust-lang/cargo/issues/6669) - name: Run cargo doc test + if: matrix.toolchain != '1.84.0' run: cargo test --verbose --doc - name: Run cargo clippy diff --git a/.github/workflows/semgrep.yml b/.github/workflows/semgrep.yml index b40314b36..3ae3dd574 100644 --- a/.github/workflows/semgrep.yml +++ b/.github/workflows/semgrep.yml @@ -1,24 +1,30 @@ +name: Semgrep OSS scan on: pull_request: {} + push: + branches: [main, master] workflow_dispatch: {} - push: - branches: - - main - - master schedule: - - cron: '0 0 * * *' -name: Semgrep config + - cron: '0 0 15 * *' +concurrency: + group: semgrep-${{ github.event_name }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true +permissions: + contents: read jobs: semgrep: - name: semgrep/ci - runs-on: ubuntu-latest - env: - SEMGREP_APP_TOKEN: ${{ secrets.SEMGREP_APP_TOKEN }} - SEMGREP_URL: https://cloudflare.semgrep.dev - SEMGREP_APP_URL: https://cloudflare.semgrep.dev - SEMGREP_VERSION_CHECK_URL: https://cloudflare.semgrep.dev/api/check-version - container: - image: returntocorp/semgrep + name: semgrep-oss + runs-on: ubuntu-slim steps: - - uses: actions/checkout@v4 - - run: semgrep ci + - uses: actions/checkout@v5 + with: + fetch-depth: 1 + - id: cache-semgrep + uses: actions/cache@v5 + with: + path: ~/.local + key: semgrep-1.160.0-${{ runner.os }} + - if: steps.cache-semgrep.outputs.cache-hit != 'true' + run: pip install --user semgrep==1.160.0 + - run: echo "$HOME/.local/bin" >> "$GITHUB_PATH" + - run: semgrep scan --config=auto diff --git a/.gitignore b/.gitignore index abc8cf514..a8fa88f99 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ dhat-heap.json .vscode .idea .cover -bleeper.user.toml \ No newline at end of file +bleeper.user.toml +.DS_Store diff --git a/Cargo.toml b/Cargo.toml index d3c8603be..6a0df2d9e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ members = [ "pingora-ketama", "pingora-load-balancing", "pingora-memory-cache", + "pingora-prometheus", "tinyufo", ] @@ -38,10 +39,14 @@ bytes = "1.0" derivative = "2.2.0" http = "1" log = "0.4" +futures = "0.3" h2 = ">=0.4.11" once_cell = "1" lru = "0.16.3" ahash = ">=0.8.9" +[workspace.lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)'] } + [profile.bench] debug = true diff --git a/docs/quick_start.md b/docs/quick_start.md index 329a5be4a..731248665 100644 --- a/docs/quick_start.md +++ b/docs/quick_start.md @@ -19,7 +19,8 @@ cargo new load_balancer In your project's `cargo.toml` file add the following to your dependencies ``` async-trait="0.1" -pingora = { version = "0.3", features = [ "lb" ] } +pingora = { version = "0.8.0", features = ["openssl", "lb"] } + ``` ### Create a pingora server @@ -321,4 +322,4 @@ The full code for this example is available in this repository under Other examples that you may find helpful are also available here [pingora-proxy/examples/](../pingora-proxy/examples/) -[pingora/examples](../pingora/examples/) \ No newline at end of file +[pingora/examples](../pingora/examples/) diff --git a/docs/user_guide/conf.md b/docs/user_guide/conf.md index 1f55859ea..435e5b12f 100644 --- a/docs/user_guide/conf.md +++ b/docs/user_guide/conf.md @@ -23,12 +23,91 @@ group: webusers | threads | number of threads per service | number | | user | the user the pingora server should be run under after daemonization | string | | group | the group the pingora server should be run under after daemonization | string | +| working_directory | the working directory for the daemonized process | string | | client_bind_to_ipv4 | source IPv4 addresses to bind to when connecting to server | list of string | | client_bind_to_ipv6 | source IPv6 addresses to bind to when connecting to server| list of string | | ca_file | The path to the root CA file | string | | s2n_config_cache_size | The maximum number of unique s2n configs to cache. A value of 0 disables the cache. Default: 10 (s2n-tls only) | number | | work_stealing | Enable work stealing runtime (default true). See Pingora runtime (WIP) section for more info | bool | -| upstream_keepalive_pool_size | The number of total connections to keep in the connection pool | number | +| runtime_enable_alt_timer | Enable Tokio's experimental alternative timer on work-stealing service runtimes. Requires building with `--cfg tokio_unstable`. Ignored when `work_stealing` is disabled. Default: `false` | bool | +| fast_timeout_to_tokio_threshold_seconds | Timeout durations greater than this value use Tokio's native timeout instead of Pingora's fast timeout. Default: `900`. Set to `null` to disable the Tokio fallback. | number | +| runtime_metrics_poll_time_histogram | Enable Tokio poll-time histograms on service runtimes. Requires building with `--cfg tokio_unstable`; adds two timestamp reads to every task poll. Default: `false` | bool | +| runtime_metrics_poll_time_histogram_scale | Bucket scale for Tokio poll-time histograms. Valid values: `linear`, `log`. Ignored unless `runtime_metrics_poll_time_histogram` is enabled. | string | +| runtime_metrics_poll_time_histogram_resolution_micros | Width of the first Tokio poll-time histogram bucket in microseconds. Must be greater than 0. Ignored unless `runtime_metrics_poll_time_histogram` is enabled. | number | +| runtime_metrics_poll_time_histogram_buckets | Number of Tokio poll-time histogram buckets. Must be greater than 0 and at most 1024. Memory usage scales with runtimes × workers × buckets. Ignored unless `runtime_metrics_poll_time_histogram` is enabled. | number | +| upstream_keepalive_pool_size | The number of idle upstream connections to keep per tokio worker. The pool's effective ceiling is `upstream_keepalive_pool_size × threads`. Eviction is globally consistent across workers. | number | +| daemon_wait_for_ready | When `true` and `daemon` is `true`, the parent process waits for the daemon to signal readiness (via `SIGUSR1`) before exiting. This causes systemd to delay sending `SIGQUIT` to the old process until the new instance is fully bootstrapped. Default: `false` | bool | +| daemon_ready_timeout_seconds | How long (in seconds) the parent waits for the daemon to signal readiness when `daemon_wait_for_ready` is `true`. If the daemon does not signal in time the parent exits with a non-zero code, causing systemd to abort the reload. Default: `600` | number | +| daemon_notify_timeout_seconds | How long (in seconds) the daemon retries sending `SIGUSR1` to the parent when the attempt fails with a permission error. This covers the brief window after the fork where the parent has not yet dropped its UID to match the daemon. Default: `60` | number | + +## dial9 + +dial9 Tokio runtime telemetry is configured programmatically, not through +the YAML configuration file. This avoids applying experimental telemetry to +every service runtime and lets services provide non-serializable options such +as a pre-built S3 client. + +dial9 is only available when Pingora is built with the `dial9` feature and +`--cfg tokio_unstable`. Services can override the global runtime options with +`runtime_opts_override()`: + +```rust +use pingora::server::{Dial9RuntimeOpts, RuntimeOpts}; +use pingora::services::Service; + +struct MyService; + +impl Service for MyService { + fn name(&self) -> &str { + "my-service" + } + + fn runtime_opts_override(&self, global: &RuntimeOpts) -> Option { + let mut opts = global.clone(); + opts.dial9 = Some( + Dial9RuntimeOpts::new("/var/lib/pingora/dial9/my-service/trace.bin") + .with_max_file_size(100 * 1024 * 1024) + .with_max_total_size(512 * 1024 * 1024), + ); + Some(opts) + } +} +``` + +When built with the `dial9-worker-s3` feature, sealed trace segments can also +be uploaded to an S3-compatible bucket: + +```rust +use pingora::server::{Dial9RuntimeOpts, Dial9S3UploadOpts, RuntimeOpts}; +use pingora::services::Service; + +struct MyService { + s3_client: aws_sdk_s3::Client, +} + +impl Service for MyService { + fn name(&self) -> &str { + "my-service" + } + + fn runtime_opts_override(&self, global: &RuntimeOpts) -> Option { + let mut opts = global.clone(); + opts.dial9 = Some( + Dial9RuntimeOpts::new("/var/lib/pingora/dial9/my-service/trace.bin") + .with_s3_upload( + Dial9S3UploadOpts::new("my-trace-bucket", "my-service") + .with_prefix("traces/my-service") + .with_region("us-east-1") + .with_client(self.s3_client.clone()), + ), + ); + Some(opts) + } +} +``` + +The S3 client is optional. When omitted, dial9 uses the AWS SDK default +configuration chain and its bucket-region detection. ## Extension Any unknown settings will be ignored. This allows extending the conf file to add and pass user defined settings. See User defined configuration section. diff --git a/docs/user_guide/modify_filter.md b/docs/user_guide/modify_filter.md index 3e5378fb9..a833fc27a 100644 --- a/docs/user_guide/modify_filter.md +++ b/docs/user_guide/modify_filter.md @@ -123,8 +123,7 @@ impl ProxyHttp for MyGateway { fn main() { ... - let mut prometheus_service_http = - pingora::services::listening::Service::prometheus_http_service(); + let mut prometheus_service_http = pingora_prometheus::prometheus_http_service(); prometheus_service_http.add_tcp("127.0.0.1:6192"); my_server.add_service(prometheus_service_http); diff --git a/docs/user_guide/phase.md b/docs/user_guide/phase.md index 3c80f9138..23bfb7bc0 100644 --- a/docs/user_guide/phase.md +++ b/docs/user_guide/phase.md @@ -29,7 +29,8 @@ Pingora-proxy allows users to insert arbitrary logic into the life of a request. upstream_request_filter --> request_body_filter; request_body_filter --> SendReq{{IO: send request to upstream}}; SendReq-->RecvResp{{IO: read response from upstream}}; - RecvResp-->upstream_response_filter-->response_filter-->upstream_response_body_filter-->response_body_filter-->logging-->endreq("request done"); + RecvResp-.feature: adjust_upstream_modules.->adjust_upstream_modules; + adjust_upstream_modules-->upstream_response_filter-->response_filter-->upstream_response_body_filter-->response_body_filter-->logging-->endreq("request done"); fail_to_connect --can retry-->upstream_peer; fail_to_connect --can't retry-->fail_to_proxy--send error response-->logging; @@ -92,6 +93,11 @@ If the error is not retry-able, the request will end. ### `upstream_request_filter()` This phase is to modify requests before sending to upstream. +### `adjust_upstream_modules()` _(feature: `adjust_upstream_modules`)_ +This phase is triggered when the upstream response header arrives, before upstream modules (such as `upstream_compression`) process it. + +Use this to configure upstream module behavior based on the response header, e.g. setting a dictionary for dictionary-based content encoding. The response header is provided as an immutable reference; to modify the response header itself, use `upstream_response_filter()` instead. + ### `upstream_response_filter()/upstream_response_body_filter()/upstream_response_trailer_filter()` This phase is triggered after an upstream response header/body/trailer is received. @@ -129,6 +135,13 @@ This is also not a phase, but another callback. `fail_to_proxy()` errors are automatically logged in the error log, but users may not be interested in every error. For example, downstream errors are logged if the client disconnects early, but these errors can become noisy if users are mainly interested in observing upstream issues. This callback can inspect the error and returns true or false. If true, the error will not be written to the log. +### `suppress_proxy_warn_log()` +This is also not a phase, but another callback. + +This experimental callback can suppress proxy warning logs that do not reach `fail_to_proxy()`, such as retryable proxy upstream failures and downstream errors ignored while cache fill continues. The callback receives a `ProxyWarnLogContext` so users can distinguish these warning contexts. Final proxy errors are still handled by `suppress_error_log()`. + +This API may change or be removed until indicated otherwise. Suppressing retry warning logs can remove the only per-retry audit record, so users should provide alternative observability, such as metrics or logs from this hook. + ### Cache filters To be documented diff --git a/docs/user_guide/phase_chart.md b/docs/user_guide/phase_chart.md index 949887243..b915f9509 100644 --- a/docs/user_guide/phase_chart.md +++ b/docs/user_guide/phase_chart.md @@ -14,7 +14,8 @@ Pingora proxy phases without caching upstream_request_filter --> request_body_filter; request_body_filter --> SendReq{{IO: send request to upstream}}; SendReq-->RecvResp{{IO: read response from upstream}}; - RecvResp-->upstream_response_filter-->response_filter-->upstream_response_body_filter-->response_body_filter-->logging-->endreq("request done"); + RecvResp-.feature: adjust_upstream_modules.->adjust_upstream_modules; + adjust_upstream_modules-->upstream_response_filter-->response_filter-->upstream_response_body_filter-->response_body_filter-->logging-->endreq("request done"); fail_to_connect --can retry-->upstream_peer; fail_to_connect --can't retry-->fail_to_proxy--send error response-->logging; diff --git a/docs/user_guide/prom.md b/docs/user_guide/prom.md index f248b0e68..1e83c0f3c 100644 --- a/docs/user_guide/prom.md +++ b/docs/user_guide/prom.md @@ -1,10 +1,21 @@ # Prometheus -Pingora has a built-in prometheus HTTP metric server for scraping. +The [`pingora-prometheus`](https://docs.rs/pingora-prometheus) crate provides a +Prometheus HTTP metrics server for scraping. + +## Adding the Dependency + +Add `pingora-prometheus` to your `Cargo.toml`: + +```toml +pingora-prometheus = "0.8.0" +``` + +## Setting up a Prometheus Metrics Endpoint ```rust ... - let mut prometheus_service_http = Service::prometheus_http_service(); + let mut prometheus_service_http = pingora_prometheus::prometheus_http_service(); prometheus_service_http.add_tcp("0.0.0.0:1234"); my_server.add_service(prometheus_service_http); my_server.run_forever(); diff --git a/docs/user_guide/rate_limiter.md b/docs/user_guide/rate_limiter.md index 31a6b5a91..554d23a66 100644 --- a/docs/user_guide/rate_limiter.md +++ b/docs/user_guide/rate_limiter.md @@ -5,8 +5,8 @@ Pingora provides a crate `pingora-limits` which provides a simple and easy to us 1. Add the following dependencies to your `Cargo.toml`: ```toml async-trait="0.1" - pingora = { version = "0.3", features = [ "lb" ] } - pingora-limits = "0.3.0" + pingora = { version = "0.8", features = [ "lb", "openssl" ] } + pingora-limits = "0.8.0" once_cell = "1.19.0" ``` 2. Declare a global rate limiter map to store the rate limiter for each client. In this example, we use `appid`. @@ -20,6 +20,7 @@ Pingora provides a crate `pingora-limits` which provides a simple and easy to us ```rust use async_trait::async_trait; use once_cell::sync::Lazy; +use pingora::http::ResponseHeader; use pingora::prelude::*; use pingora_limits::rate::Rate; use std::sync::Arc; diff --git a/pingora-boringssl/src/ext.rs b/pingora-boringssl/src/ext.rs index 256e4ac5b..2dfcfef2e 100644 --- a/pingora-boringssl/src/ext.rs +++ b/pingora-boringssl/src/ext.rs @@ -146,6 +146,20 @@ pub fn clear_error_stack() { let _ = ErrorStack::get(); } +/// Export keying material from a TLS connection +/// +/// Derives keying material for application use in accordance with RFC 5705. +/// +/// See [SSL_export_keying_material](https://commondatastorage.googleapis.com/chromium-boringssl-docs/ssl.h.html#SSL_export_keying_material). +pub fn ssl_export_keying_material( + ssl: &SslRef, + out: &mut [u8], + label: &str, + context: Option<&[u8]>, +) -> Result<(), ErrorStack> { + ssl.export_keying_material(out, label, context) +} + /// Create a new [Ssl] from &[SslAcceptor] /// /// This function is needed because [Ssl::new()] doesn't take `&SslContextRef` like openssl-rs diff --git a/pingora-cache/Cargo.toml b/pingora-cache/Cargo.toml index 401d827ce..bfbebc514 100644 --- a/pingora-cache/Cargo.toml +++ b/pingora-cache/Cargo.toml @@ -37,8 +37,8 @@ httpdate = "1.0.2" log = { workspace = true } async-trait = { workspace = true } parking_lot = "0.12" -cf-rustracing = "1.0" -cf-rustracing-jaeger = "1.0" +cf-rustracing = { version = "1.0", optional = true } +cf-rustracing-jaeger = { version = "1.0", optional = true } rmp = "0.8.14" tokio = { workspace = true } lru = { workspace = true } @@ -49,11 +49,13 @@ strum = { version = "0.26", features = ["derive"] } rand = "0.8" [dev-dependencies] +cf-rustracing = "1.0" +cf-rustracing-jaeger = "1.0" tokio-test = "0.4" tokio = { workspace = true, features = ["fs"] } env_logger = "0.11" dhat = "0" -futures = "0.3" +futures = { workspace = true } [[bench]] name = "simple_lru_memory" @@ -73,3 +75,4 @@ openssl = ["pingora-core/openssl"] boringssl = ["pingora-core/boringssl"] rustls = ["pingora-core/rustls"] s2n = ["pingora-core/s2n"] +trace = ["dep:cf-rustracing", "dep:cf-rustracing-jaeger"] diff --git a/pingora-cache/src/cache_control.rs b/pingora-cache/src/cache_control.rs index 98af7fbbc..f8612080a 100644 --- a/pingora-cache/src/cache_control.rs +++ b/pingora-cache/src/cache_control.rs @@ -92,6 +92,51 @@ impl DirectiveValue { } } } + + /// Parse the [DirectiveValue] as delta seconds, permitting a fractional component. + /// + /// Values with a fractional part are rounded down (floored) to the nearest + /// non-negative integer: e.g. `1.9` -> `1`, `0.5` -> `0`. This is useful + /// for compatibility with upstreams that emit fractional ttls (which + /// RFC 9111 strict integer parsing rejects). + /// + /// Integer parsing is attempted first, so strictly-numeric values (including + /// overflow-capped values and quoted integers) behave identically to + /// [Self::parse_as_delta_seconds]. Negative, non-finite, or non-numeric + /// values still return an error. + /// + /// `"`s are ignored. The value is capped to [DELTA_SECONDS_OVERFLOW_VALUE]. + pub fn parse_as_delta_seconds_floor(&self) -> Result { + // UTF-8 validate once; on non-UTF8 input, propagate the same error as + // [Self::parse_as_delta_seconds]. + let s = self.parse_as_str()?; + match s.parse::() { + Ok(value) => Ok(value), + Err(e) if e.kind() == &IntErrorKind::PosOverflow => Ok(DELTA_SECONDS_OVERFLOW_VALUE), + Err(int_err) => { + // Fall back to parsing as a non-negative finite float and floor. + // On any failure, return an error equivalent to the strict + // [Self::parse_as_delta_seconds] u32-parse error. + match s.parse::() { + Ok(f) if f.is_finite() && f >= 0.0 => { + if f >= DELTA_SECONDS_OVERFLOW_VALUE as f64 { + Ok(DELTA_SECONDS_OVERFLOW_VALUE) + } else { + // Safe cast: `f` is finite, non-negative, and strictly + // less than `DELTA_SECONDS_OVERFLOW_VALUE` (i32::MAX), + // which fits in `u32` after flooring. + Ok(f.floor() as u32) + } + } + _ => Error::e_because( + ErrorType::InternalError, + "could not parse value as u32", + int_err, + ), + } + } + } + } } /// An ordered map to store cache control key value pairs. @@ -102,6 +147,15 @@ pub type DirectiveMap = IndexMap>; pub struct CacheControl { /// The parsed directives pub directives: DirectiveMap, + /// When set, delta-seconds directives (`max-age`, `s-maxage`, + /// `stale-while-revalidate`, `stale-if-error`) accept fractional values + /// and round them down to the nearest non-negative integer. + /// + /// Defaults to `false`, matching RFC 9111 strict integer parsing. Enable + /// via [CacheControl::with_float_seconds] (or by assigning the field + /// directly) for contexts that need to interoperate with upstreams that + /// emit fractional ttls. + pub allow_float_seconds: bool, } /// Cacheability calculated from cache control. @@ -198,7 +252,20 @@ impl CacheControl { directives.insert(key.unwrap(), value); } } - Some(CacheControl { directives }) + Some(CacheControl { + directives, + allow_float_seconds: false, + }) + } + + /// Builder setter: enable fractional delta-seconds parsing. + /// + /// See [CacheControl::allow_float_seconds] for semantics. Returns `self` + /// so it can be chained onto a parser call, e.g. + /// `CacheControl::from_resp_headers(&resp).map(|cc| cc.with_float_seconds())`. + pub fn with_float_seconds(mut self) -> Self { + self.allow_float_seconds = true; + self } /// Parse from the given header name in `headers` @@ -282,7 +349,12 @@ impl CacheControl { fn parse_delta_seconds(&self, key: &str) -> Result> { if let Some(Some(dir_value)) = self.directives.get(key) { - Ok(Some(dir_value.parse_as_delta_seconds()?)) + let value = if self.allow_float_seconds { + dir_value.parse_as_delta_seconds_floor()? + } else { + dir_value.parse_as_delta_seconds()? + }; + Ok(Some(value)) } else { Ok(None) } @@ -869,4 +941,168 @@ mod tests { let cc = CacheControl::from_req_headers(&req).unwrap(); assert!(cc.only_if_cached()) } + + #[test] + fn test_parse_as_delta_seconds_floor() { + // Integer values behave identically to the strict parser + let v = DirectiveValue(b"10".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 10); + + let v = DirectiveValue(b"\"10\"".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 10); + + // Quoted fractional values are unwrapped by parse_as_str and floored. + let v = DirectiveValue(b"\"1.5\"".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 1); + + let v = DirectiveValue(b"0".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 0); + + // Integer positive overflow is still capped + let v = DirectiveValue(b"99999999999999999999".to_vec()); + assert_eq!( + v.parse_as_delta_seconds_floor().unwrap(), + DELTA_SECONDS_OVERFLOW_VALUE + ); + + // Floats are floored + let v = DirectiveValue(b"1.5".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 1); + + let v = DirectiveValue(b"1.9".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 1); + + let v = DirectiveValue(b"0.5".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 0); + + let v = DirectiveValue(b"3600.0".to_vec()); + assert_eq!(v.parse_as_delta_seconds_floor().unwrap(), 3600); + + // Float positive overflow is capped + let v = DirectiveValue(b"99999999999.5".to_vec()); + assert_eq!( + v.parse_as_delta_seconds_floor().unwrap(), + DELTA_SECONDS_OVERFLOW_VALUE + ); + + // Negative values are rejected (matches strict behavior) + assert!(DirectiveValue(b"-1".to_vec()) + .parse_as_delta_seconds_floor() + .is_err()); + assert!(DirectiveValue(b"-1.5".to_vec()) + .parse_as_delta_seconds_floor() + .is_err()); + + // Non-finite / non-numeric values are rejected + assert!(DirectiveValue(b"NaN".to_vec()) + .parse_as_delta_seconds_floor() + .is_err()); + assert!(DirectiveValue(b"inf".to_vec()) + .parse_as_delta_seconds_floor() + .is_err()); + assert!(DirectiveValue(b"abc".to_vec()) + .parse_as_delta_seconds_floor() + .is_err()); + + // Non-UTF8 bytes are rejected with the same utf-8 error as the strict parser. + let v = DirectiveValue(b"ba\xFFr".to_vec()); + assert_eq!( + v.parse_as_delta_seconds_floor() + .unwrap_err() + .context + .unwrap() + .to_string(), + "could not parse value as utf8", + ); + } + + #[test] + fn test_cache_control_allow_float_seconds_non_utf8_value() { + // Non-UTF8 bytes inside `max-age` should still produce the utf-8 error + // when the float-permitting flag is on, matching the strict parser. + let mut resp = response::Builder::new().body(()).unwrap(); + resp.headers_mut().insert( + CACHE_CONTROL, + HeaderValue::from_bytes(b"max-age=ba\xFFr").unwrap(), + ); + let (parts, _) = resp.into_parts(); + let cc = CacheControl::from_resp_headers(&parts) + .unwrap() + .with_float_seconds(); + assert_eq!( + cc.max_age().unwrap_err().context.unwrap().to_string(), + "could not parse value as utf8", + ); + } + + #[test] + fn test_cache_control_allow_float_seconds_default_off() { + // Default (strict) parsing: fractional values produce an error, and + // [InterpretCacheControl::fresh_duration] returns None, matching the + // pre-existing behavior. + let resp = build_response(CACHE_CONTROL, "max-age=10.7"); + let cc = CacheControl::from_resp_headers(&resp).unwrap(); + assert!(!cc.allow_float_seconds); + assert!(cc.max_age().is_err()); + assert!(cc.fresh_duration().is_none()); + } + + #[test] + fn test_cache_control_with_float_seconds() { + // `max-age` with a fractional value is floored when the flag is on. + let resp = build_response(CACHE_CONTROL, "max-age=10.7"); + let cc = CacheControl::from_resp_headers(&resp) + .unwrap() + .with_float_seconds(); + assert!(cc.allow_float_seconds); + assert_eq!(cc.max_age().unwrap().unwrap(), 10); + assert_eq!(cc.fresh_duration().unwrap(), Duration::from_secs(10)); + + // `s-maxage` still wins over `max-age` and is also floored. + let resp = build_response(CACHE_CONTROL, "s-maxage=3600.99, max-age=1800"); + let cc = CacheControl::from_resp_headers(&resp) + .unwrap() + .with_float_seconds(); + assert_eq!(cc.s_maxage().unwrap().unwrap(), 3600); + assert_eq!(cc.fresh_duration().unwrap(), Duration::from_secs(3600)); + + // `stale-while-revalidate` and `stale-if-error` also pick up flooring. + let resp = build_response( + CACHE_CONTROL, + "max-age=10, stale-while-revalidate=60.5, stale-if-error=30.9", + ); + let cc = CacheControl::from_resp_headers(&resp) + .unwrap() + .with_float_seconds(); + assert_eq!(cc.stale_while_revalidate().unwrap().unwrap(), 60); + assert_eq!(cc.stale_if_error().unwrap().unwrap(), 30); + assert_eq!( + cc.serve_stale_while_revalidate_duration().unwrap(), + Duration::from_secs(60) + ); + assert_eq!( + cc.serve_stale_if_error_duration().unwrap(), + Duration::from_secs(30) + ); + + // Integer values are unaffected when the flag is on. + let resp = build_response(CACHE_CONTROL, "max-age=12345"); + let cc = CacheControl::from_resp_headers(&resp) + .unwrap() + .with_float_seconds(); + assert_eq!(cc.fresh_duration().unwrap(), Duration::from_secs(12345)); + + // Invalid (non-numeric, negative) values still fail to parse under the flag. + let resp = build_response(CACHE_CONTROL, "max-age=abc"); + let cc = CacheControl::from_resp_headers(&resp) + .unwrap() + .with_float_seconds(); + assert!(cc.max_age().is_err()); + + let resp = build_response(CACHE_CONTROL, "max-age=-1.5"); + let cc = CacheControl::from_resp_headers(&resp) + .unwrap() + .with_float_seconds(); + assert!(cc.max_age().is_err()); + } } diff --git a/pingora-cache/src/eviction/lru.rs b/pingora-cache/src/eviction/lru.rs index d241ee69d..aa157eb07 100644 --- a/pingora-cache/src/eviction/lru.rs +++ b/pingora-cache/src/eviction/lru.rs @@ -19,7 +19,7 @@ use crate::key::CompactCacheKey; use async_trait::async_trait; use log::{info, warn}; -use pingora_error::{BError, ErrorType::*, OrErr, Result}; +use pingora_error::{BError, Error, ErrorType::*, OrErr, Result}; use pingora_lru::Lru; use rand::Rng; use serde::de::SeqAccess; @@ -62,6 +62,16 @@ impl Manager { Manager(Lru::with_capacity_and_watermark(limit, capacity, watermark)) } + /// Return the current total cache weight limit. + pub fn weight_limit(&self) -> usize { + self.0.weight_limit() + } + + /// Set the total cache weight limit used by future eviction decisions. + pub fn set_weight_limit(&self, limit: usize) { + self.0.set_weight_limit(limit); + } + /// Get the number of shards pub fn shards(&self) -> usize { self.0.shards() @@ -72,7 +82,9 @@ impl Manager { self.0.shard_weight(shard) } - /// Get the number of items in a specific shard + /// Get the number of items in a specific shard. Best-effort + /// lock-free read; see [`pingora_lru::Lru::shard_len`] for the + /// consistency semantics. pub fn shard_len(&self, shard: usize) -> usize { self.0.shard_len(shard) } @@ -85,6 +97,15 @@ impl Manager { (u64key(key) % N as u64) as usize } + /// Peek at the least-recently-used key in the given shard without evicting it. + /// + /// Returns the cache key at the LRU tail of the shard, or `None` if empty. + /// Useful for reporting the eviction frontier (the age of the next item + /// that would be evicted). + pub fn peek_lru(&self, shard: usize) -> Option { + self.0.peek_lru(shard).map(|(key, _weight)| key) + } + /// Serialize the given shard pub fn serialize_shard(&self, shard: usize) -> Result> { use rmp_serde::encode::Serializer; @@ -246,10 +267,21 @@ impl EvictionManager for Manager { .await .or_err(InternalError, "async blocking IO failure")??; + // Per-shard errors are isolated so a single failing shard does not abort + // the entire save and leave the remaining shards stale on disk. + let mut saved_shards = 0usize; + let mut failed_shards = 0usize; for i in 0..N { - let data = self.serialize_shard(i)?; + let data = match self.serialize_shard(i) { + Ok(d) => d, + Err(e) => { + warn!("Failed to serialize shard {i}: {e}. Skipping shard."); + failed_shards += 1; + continue; + } + }; let dir_path = dir_path.to_owned(); - tokio::task::spawn_blocking(move || { + let result = tokio::task::spawn_blocking(move || { let dir_path = Path::new(&dir_path); let final_path = dir_path.join(format!("{}.{i}", FILE_NAME)); // create a temporary filename using a randomized u32 hash to minimize the chance of multiple writers writing to the same tmp file @@ -272,47 +304,104 @@ impl EvictionManager for Manager { ) }) }) - .await - .or_err(InternalError, "async blocking IO failure")??; + .await; + + match result { + Ok(Ok(())) => saved_shards += 1, + Ok(Err(e)) => { + warn!("Failed to save shard {i}: {e}. Skipping shard."); + failed_shards += 1; + } + Err(join_err) => { + warn!( + "Failed to save shard {i}: async blocking IO failure {join_err}. Skipping shard." + ); + failed_shards += 1; + } + } + } + + if failed_shards == 0 { + info!("Successfully saved {saved_shards}/{N} shards."); + Ok(()) + } else if failed_shards == N { + Error::e_explain( + InternalError, + format!("All {N} shards failed to save; see prior warnings for per-shard causes."), + ) + } else { + warn!( + "Saved {saved_shards}/{N} shards; {failed_shards} shards failed. Persisted cache state may be incomplete." + ); + Ok(()) } - Ok(()) } async fn load(&self, dir_path: &str) -> Result<()> { - // TODO: check the saved shards so that we load all the save files - let mut loaded_shards = 0; + // Per-shard errors are isolated so a single failing shard does not abort + // the entire load. A missing shard file is treated as an empty shard. + let mut loaded_shards = 0usize; + let mut missing_shards = 0usize; + let mut error_shards = 0usize; for i in 0..N { let dir_path = dir_path.to_owned(); - let data = tokio::task::spawn_blocking(move || { + let read_result = tokio::task::spawn_blocking(move || { let file_path = Path::new(&dir_path).join(format!("{}.{i}", FILE_NAME)); - let mut file = File::open(&file_path) - .or_err_with(InternalError, || err_str_path("fail to open", &file_path))?; + let mut file = match File::open(&file_path) { + Ok(f) => f, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(e) => { + return Err(e).or_err_with(InternalError, || { + err_str_path("fail to open", &file_path) + }); + } + }; let mut buffer = Vec::with_capacity(8192); file.read_to_end(&mut buffer) .or_err_with(InternalError, || { err_str_path("fail to read from", &file_path) })?; - Ok::, BError>(buffer) + Ok::>, BError>(Some(buffer)) }) - .await - .or_err(InternalError, "async blocking IO failure")??; + .await; + + let data = match read_result { + Ok(Ok(Some(buf))) => buf, + Ok(Ok(None)) => { + missing_shards += 1; + continue; + } + Ok(Err(e)) => { + warn!("Failed to load shard {i}: {e}. Skipping shard."); + error_shards += 1; + continue; + } + Err(join_err) => { + warn!( + "Failed to load shard {i}: async blocking IO failure {join_err}. Skipping shard." + ); + error_shards += 1; + continue; + } + }; if let Err(e) = self.deserialize_shard(&data) { - warn!("Failed to deserialize shard {}: {}. Skipping shard.", i, e); - continue; // Skip shard and move onto the next one + warn!("Failed to deserialize shard {i}: {e}. Skipping shard."); + error_shards += 1; + continue; } loaded_shards += 1; } - // Log how many shards were successfully loaded - if loaded_shards < N { - warn!( - "Only loaded {}/{} shards. Cache may be incomplete.", - loaded_shards, N - ) + if loaded_shards == N { + info!("Successfully loaded {loaded_shards}/{N} shards."); + } else if loaded_shards == 0 && error_shards == 0 { + info!("No persisted LRU shards found. Cache will start empty."); } else { - info!("Successfully loaded {}/{} shards.", loaded_shards, N) + warn!( + "Loaded {loaded_shards}/{N} shards (missing: {missing_shards}, errored: {error_shards}). Cache may be incomplete." + ); } cleanup_temp_files(dir_path); @@ -411,6 +500,26 @@ mod test { assert_eq!(v[1], key2); } + #[test] + fn test_set_weight_limit() { + let lru = Manager::<1>::with_capacity(4, 10); + let until = SystemTime::now(); + let key1 = CacheKey::new("", "a", "1").to_compact(); + let key2 = CacheKey::new("", "b", "1").to_compact(); + let key3 = CacheKey::new("", "c", "1").to_compact(); + assert_eq!(lru.weight_limit(), 4); + assert!(lru.admit(key1.clone(), 1, until).is_empty()); + assert!(lru.admit(key2.clone(), 1, until).is_empty()); + + lru.set_weight_limit(1); + assert_eq!(lru.weight_limit(), 1); + let evicted = lru.admit(key3, 1, until); + assert_eq!(evicted, vec![key1, key2]); + + lru.set_weight_limit(10); + assert_eq!(lru.weight_limit(), 10); + } + #[test] fn test_access() { let lru = Manager::<1>::with_capacity(4, 10); @@ -579,6 +688,215 @@ mod test { assert_eq!(ser1, lru2.serialize_shard(1).unwrap()); } + #[tokio::test] + async fn test_load_no_shards() { + // Loading from an empty directory should succeed with an empty LRU. + let test_dir = "/tmp/test_lru_no_shards"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + + let lru = Manager::<4>::with_capacity(10, 10); + lru.load(test_dir).await.unwrap(); + assert_eq!(lru.total_items(), 0); + assert_eq!(lru.total_size(), 0); + + std::fs::remove_dir_all(test_dir).unwrap(); + } + + #[tokio::test] + async fn test_load_partial_shards() { + // A subset of shard files is missing on disk. Load should succeed and + // populate the LRU from only the shards that exist; missing shards are + // treated as empty rather than aborting the load. + let test_dir = "/tmp/test_lru_partial_shards"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + + let until = SystemTime::now(); + let src = Manager::<4>::with_capacity(100, 100); + for i in 0..16 { + src.admit( + CacheKey::new("", format!("k{i}"), "1").to_compact(), + 1, + until, + ); + } + src.save(test_dir).await.unwrap(); + let baseline = src.total_items(); + assert!(baseline > 0); + + // Remove half the shard files. + std::fs::remove_file(format!("{test_dir}/lru.data.1")).unwrap(); + std::fs::remove_file(format!("{test_dir}/lru.data.3")).unwrap(); + + let dst = Manager::<4>::with_capacity(100, 100); + dst.load(test_dir).await.unwrap(); + // Some entries loaded, but fewer than the baseline. + assert!(dst.total_items() > 0); + assert!(dst.total_items() < baseline); + + std::fs::remove_dir_all(test_dir).unwrap(); + } + + #[tokio::test] + async fn test_save_partial_failure_continues() { + // A per-shard rename failure must not abort the whole save. We simulate + // by pre-creating a directory at one shard's final path so that shard's + // atomic rename fails while the others succeed. + let test_dir = "/tmp/test_lru_save_partial_fail"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + std::fs::create_dir(format!("{test_dir}/lru.data.1")).unwrap(); + + let until = SystemTime::now(); + let src = Manager::<4>::with_capacity(100, 100); + for i in 0..16 { + src.admit( + CacheKey::new("", format!("k{i}"), "1").to_compact(), + 1, + until, + ); + } + src.save(test_dir).await.unwrap(); + + // Shards 0, 2, 3 should be regular files; shard 1 is still a directory. + for i in [0, 2, 3] { + let p = format!("{test_dir}/lru.data.{i}"); + let meta = std::fs::metadata(&p).unwrap(); + assert!(meta.is_file(), "shard {i} should be a regular file"); + assert!(meta.len() > 0, "shard {i} should be non-empty"); + } + assert!(std::fs::metadata(format!("{test_dir}/lru.data.1")) + .unwrap() + .is_dir()); + + std::fs::remove_dir(format!("{test_dir}/lru.data.1")).unwrap(); + std::fs::remove_dir_all(test_dir).unwrap(); + } + + #[tokio::test] + async fn test_save_total_failure_returns_err() { + // If every shard fails to save, the function must return Err so callers + // can alarm on the catastrophic case. + let test_dir = "/tmp/test_lru_save_total_fail"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + for i in 0..4 { + std::fs::create_dir(format!("{test_dir}/lru.data.{i}")).unwrap(); + } + + let until = SystemTime::now(); + let src = Manager::<4>::with_capacity(100, 100); + src.admit(CacheKey::new("", "k", "1").to_compact(), 1, until); + + let err = src.save(test_dir).await.unwrap_err(); + assert!( + err.to_string().contains("All 4 shards failed to save"), + "unexpected error message: {err}" + ); + + for i in 0..4 { + std::fs::remove_dir(format!("{test_dir}/lru.data.{i}")).unwrap(); + } + std::fs::remove_dir_all(test_dir).unwrap(); + } + + #[tokio::test] + async fn test_load_unreadable_shard_continues() { + // A shard path that opens successfully but fails to read (here a + // directory at the shard path) should not abort load of other shards. + let test_dir = "/tmp/test_lru_unreadable_shard"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + + let until = SystemTime::now(); + let src = Manager::<4>::with_capacity(100, 100); + for i in 0..16 { + src.admit( + CacheKey::new("", format!("k{i}"), "1").to_compact(), + 1, + until, + ); + } + src.save(test_dir).await.unwrap(); + + // Replace one shard with a directory so File::open succeeds but + // read_to_end returns an error. + std::fs::remove_file(format!("{test_dir}/lru.data.2")).unwrap(); + std::fs::create_dir(format!("{test_dir}/lru.data.2")).unwrap(); + + let dst = Manager::<4>::with_capacity(100, 100); + dst.load(test_dir).await.unwrap(); + // Entries from the other 3 shards still loaded. + assert!(dst.total_items() > 0); + + std::fs::remove_dir(format!("{test_dir}/lru.data.2")).unwrap(); + std::fs::remove_dir_all(test_dir).unwrap(); + } + + #[tokio::test] + async fn test_load_corrupt_shard_continues() { + // A corrupt shard file should not abort load of the remaining shards. + let test_dir = "/tmp/test_lru_corrupt_shard"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + + let until = SystemTime::now(); + let src = Manager::<4>::with_capacity(100, 100); + for i in 0..16 { + src.admit( + CacheKey::new("", format!("k{i}"), "1").to_compact(), + 1, + until, + ); + } + src.save(test_dir).await.unwrap(); + + // Truncate one shard to non-empty garbage so deserialize fails. + std::fs::write(format!("{test_dir}/lru.data.2"), b"not valid msgpack").unwrap(); + + let dst = Manager::<4>::with_capacity(100, 100); + dst.load(test_dir).await.unwrap(); + // We should have loaded entries from the other 3 shards. + assert!(dst.total_items() > 0); + + std::fs::remove_dir_all(test_dir).unwrap(); + } + + #[tokio::test] + async fn test_save_then_load_roundtrip_with_remove() { + // After load with missing shards, a subsequent save must persist every + // shard file again so the next load is complete. + let test_dir = "/tmp/test_lru_save_after_partial_load"; + let _ = std::fs::remove_dir_all(test_dir); + std::fs::create_dir_all(test_dir).unwrap(); + + let until = SystemTime::now(); + let src = Manager::<4>::with_capacity(100, 100); + for i in 0..16 { + src.admit( + CacheKey::new("", format!("k{i}"), "1").to_compact(), + 1, + until, + ); + } + src.save(test_dir).await.unwrap(); + std::fs::remove_file(format!("{test_dir}/lru.data.1")).unwrap(); + + let dst = Manager::<4>::with_capacity(100, 100); + dst.load(test_dir).await.unwrap(); + dst.save(test_dir).await.unwrap(); + + for i in 0..4 { + assert!( + Path::new(&format!("{test_dir}/lru.data.{i}")).exists(), + "shard {i} should exist after save" + ); + } + + std::fs::remove_dir_all(test_dir).unwrap(); + } + #[tokio::test] async fn test_temp_file_cleanup() { let test_dir = "/tmp/test_lru_cleanup"; @@ -614,4 +932,42 @@ mod test { // Cleanup test directory std::fs::remove_dir_all(dir_path).unwrap(); } + + #[test] + fn test_peek_lru() { + let lru = Manager::<1>::with_capacity(20, 20); + let until = SystemTime::now(); + + // empty shard returns None + assert!(lru.peek_lru(0).is_none()); + + let key1 = CacheKey::new("", "a", "1").to_compact(); + lru.admit(key1.clone(), 1, until); + // single item: it's both the head and the tail + assert_eq!(lru.peek_lru(0).unwrap(), key1); + + // admit more keys to push key1 to the tail + let key2 = CacheKey::new("", "b", "1").to_compact(); + lru.admit(key2.clone(), 1, until); + for i in 0..5 { + lru.admit( + CacheKey::new("", format!("f{i}"), "1").to_compact(), + 1, + until, + ); + } + // key1 is the LRU tail (admitted first) + assert_eq!(lru.peek_lru(0).unwrap(), key1); + + // promote key1 — now key2 becomes the tail + lru.access(&key1, 1, until); + assert_eq!(lru.peek_lru(0).unwrap(), key2); + + // peek_lru should not remove the item + assert_eq!(lru.peek_lru(0).unwrap(), key2); + assert!(lru.peek(&key2)); + + // out-of-bounds shard returns None + assert!(lru.peek_lru(999).is_none()); + } } diff --git a/pingora-cache/src/lib.rs b/pingora-cache/src/lib.rs index 867cff086..6d13409ba 100644 --- a/pingora-cache/src/lib.rs +++ b/pingora-cache/src/lib.rs @@ -16,7 +16,6 @@ #![allow(clippy::new_without_default)] -use cf_rustracing::tag::Tag; use http::{method::Method, request::Parts as ReqHeader, response::Parts as RespHeader}; use key::{CacheHashKey, CompactCacheKey, HashBinary}; use lock::WritePermit; @@ -27,7 +26,7 @@ use pingora_timeout::timeout; use std::time::{Duration, Instant, SystemTime}; use storage::MissFinishType; use strum::IntoStaticStr; -use trace::{CacheTraceCTX, Span}; +use trace::{CacheTraceCTX, Span, Tag}; pub mod cache_control; pub mod eviction; @@ -435,6 +434,7 @@ impl HttpCache { self.phase = CachePhase::Disabled(reason); self.release_write_lock(reason); // enabled_ctx will be cleared out + #[cfg_attr(not(feature = "trace"), allow(unused_mut))] let mut inner_enabled = self .inner_mut() .enabled_ctx @@ -1049,6 +1049,7 @@ impl HttpCache { inner_enabled.meta.replace(meta); + #[cfg_attr(not(feature = "trace"), allow(unused_mut))] let mut span = inner_enabled.traces.child("update_meta"); let result = inner_enabled .storage @@ -1291,6 +1292,7 @@ impl HttpCache { .enabled_ctx .as_mut() .expect("Cache enabled on cache_lookup"); + #[cfg_attr(not(feature = "trace"), allow(unused_mut))] let mut span = inner_enabled.traces.child("lookup"); let key = inner.key.as_ref().unwrap(); // safe, this phase should have cache key let now = Instant::now(); @@ -1446,6 +1448,7 @@ impl HttpCache { /// Check [Self::is_cache_locked()], panic if this request doesn't have a read lock. pub async fn cache_lock_wait(&mut self) -> LockStatus { let inner_enabled = self.inner_enabled_mut(); + #[cfg_attr(not(feature = "trace"), allow(unused_mut))] let mut span = inner_enabled.traces.child("cache_lock"); // should always call is_cache_locked() before this function, which should guarantee that // the inner cache has a read lock and lock ctx @@ -1535,6 +1538,7 @@ impl HttpCache { }) } + #[cfg_attr(not(feature = "trace"), allow(unused_mut))] async fn purge_impl( storage: &'static (dyn storage::Storage + Sync), eviction: Option<&'static (dyn eviction::EvictionManager + Sync)>, diff --git a/pingora-cache/src/lock.rs b/pingora-cache/src/lock.rs index 5633b09cb..102b3380b 100644 --- a/pingora-cache/src/lock.rs +++ b/pingora-cache/src/lock.rs @@ -14,8 +14,8 @@ //! Cache lock +use crate::trace::{Span, Tag}; use crate::{hashtable::ConcurrentHashTable, key::CacheHashKey, CacheKey}; -use crate::{Span, Tag}; use http::Extensions; use pingora_timeout::timeout; diff --git a/pingora-cache/src/memory.rs b/pingora-cache/src/memory.rs index 6ab57c808..e6e12acf8 100644 --- a/pingora-cache/src/memory.rs +++ b/pingora-cache/src/memory.rs @@ -426,7 +426,7 @@ impl Storage for MemCache { #[cfg(test)] mod test { use super::*; - use cf_rustracing::span::Span; + use crate::trace::Span; use once_cell::sync::Lazy; fn gen_meta() -> CacheMeta { diff --git a/pingora-cache/src/put.rs b/pingora-cache/src/put.rs index fbbbb70e7..942650558 100644 --- a/pingora-cache/src/put.rs +++ b/pingora-cache/src/put.rs @@ -84,6 +84,7 @@ impl CachePutCtx { } async fn put_header(&mut self, meta: CacheMeta) -> Result<()> { + #[cfg_attr(not(feature = "trace"), allow(unused_mut))] let mut trace = self.trace.child("cache put header", |o| o.start()); let miss_handler = self .storage @@ -239,7 +240,7 @@ impl CachePutCtx { #[cfg(test)] mod test { use super::*; - use cf_rustracing::span::Span; + use crate::trace::Span; use once_cell::sync::Lazy; struct TestCachePut(); diff --git a/pingora-cache/src/trace.rs b/pingora-cache/src/trace.rs index f27929a2c..e8ab85cfd 100644 --- a/pingora-cache/src/trace.rs +++ b/pingora-cache/src/trace.rs @@ -13,26 +13,129 @@ // limitations under the License. //! Distributed tracing helpers +//! +//! When the `trace` feature is enabled, this module re-exports the real +//! [`cf_rustracing`]/[`cf_rustracing_jaeger`] span types. +//! +//! When the `trace` feature is **disabled**, lightweight no-op shim types are +//! provided instead so that the rest of the crate compiles without pulling in +//! the tracing dependencies. -use cf_rustracing_jaeger::span::SpanContextState; use std::time::SystemTime; use crate::{CacheMeta, CachePhase, HitStatus}; -pub use cf_rustracing::tag::Tag; +// --------------------------------------------------------------------------- +// Real tracing implementation (feature = "trace") +// --------------------------------------------------------------------------- +#[cfg(feature = "trace")] +mod real { + pub use cf_rustracing::tag::Tag; -pub type Span = cf_rustracing::span::Span; -pub type SpanHandle = cf_rustracing::span::SpanHandle; + use cf_rustracing_jaeger::span::SpanContextState; -#[derive(Debug)] -pub(crate) struct CacheTraceCTX { - // parent span - pub cache_span: Span, - // only spans across multiple calls need to store here - pub miss_span: Span, - pub hit_span: Span, + pub type Span = cf_rustracing::span::Span; + pub type SpanHandle = cf_rustracing::span::SpanHandle; } +#[cfg(feature = "trace")] +pub use real::*; + +// --------------------------------------------------------------------------- +// No-op shim types (feature = "trace" disabled) +// --------------------------------------------------------------------------- +#[cfg(not(feature = "trace"))] +mod noop { + /// A no-op replacement for [`cf_rustracing::tag::Tag`]. + #[derive(Debug)] + pub struct Tag { + _priv: (), + } + + impl Tag { + /// Create a no-op tag. All arguments are ignored. + #[inline] + pub fn new(_name: N, _value: V) -> Self { + Tag { _priv: () } + } + } + + /// A no-op replacement for a rustracing `Span`. + #[derive(Debug)] + pub struct Span { + _priv: (), + } + + impl Span { + /// Return an inactive (no-op) span. + #[inline] + pub fn inactive() -> Self { + Span { _priv: () } + } + + /// Return a no-op handle. + #[inline] + pub fn handle(&self) -> SpanHandle { + SpanHandle { _priv: () } + } + + /// No-op: create a child span. + #[inline] + pub fn child(&self, _name: &'static str, _f: F) -> Span + where + F: FnOnce(SpanOptionsPlaceholder) -> SpanOptionsPlaceholder, + { + Span::inactive() + } + + /// No-op: set a single tag via a closure. + #[inline] + pub fn set_tag Tag>(&self, _f: F) {} + + /// No-op: set multiple tags via a closure. + #[inline] + pub fn set_tags(&self, _f: F) + where + F: FnOnce() -> I, + I: IntoIterator, + { + } + + /// No-op: set a finish time. + #[inline] + pub fn set_finish_time std::time::SystemTime>(&self, _f: F) {} + } + + /// Placeholder type used in [`Span::child`] closure signatures so that + /// existing call-sites like `span.child("name", |o| o.start())` compile. + #[doc(hidden)] + pub struct SpanOptionsPlaceholder { + _priv: (), + } + + impl SpanOptionsPlaceholder { + /// No-op: mirrors `SpanOptions::start()`. + #[inline] + pub fn start(self) -> Self { + self + } + } + + /// A no-op replacement for a rustracing `SpanHandle`. + #[derive(Debug)] + pub struct SpanHandle { + _priv: (), + } +} + +#[cfg(not(feature = "trace"))] +pub use noop::*; + +// --------------------------------------------------------------------------- +// Shared helpers (work with both real and no-op types) +// --------------------------------------------------------------------------- + +/// Tag a span with metadata from a [`CacheMeta`]. pub fn tag_span_with_meta(span: &mut Span, meta: &CacheMeta) { fn ts2epoch(ts: SystemTime) -> f64 { ts.duration_since(SystemTime::UNIX_EPOCH) @@ -55,6 +158,15 @@ pub fn tag_span_with_meta(span: &mut Span, meta: &CacheMeta) { }); } +#[derive(Debug)] +pub(crate) struct CacheTraceCTX { + // parent span + pub cache_span: Span, + // only spans across multiple calls need to store here + pub miss_span: Span, + pub hit_span: Span, +} + impl CacheTraceCTX { pub fn new() -> Self { CacheTraceCTX { diff --git a/pingora-cache/src/variance.rs b/pingora-cache/src/variance.rs index cce8160f7..fa2633f23 100644 --- a/pingora-cache/src/variance.rs +++ b/pingora-cache/src/variance.rs @@ -35,6 +35,12 @@ impl<'a> VarianceBuilder<'a> { self.values.insert(name.into(), Cow::Owned(value)); } + /// Move String name and byte string value to the variance key. Not sensitive to insertion order. + /// Useful when both the name and value are generated at runtime. + pub fn add_owned_name_value(&mut self, name: String, value: Vec) { + self.values.insert(Cow::Owned(name), Cow::Owned(value)); + } + /// Check whether this variance key actually has variance, or just refers to the root asset pub fn has_variance(&self) -> bool { !self.values.is_empty() diff --git a/pingora-core/Cargo.toml b/pingora-core/Cargo.toml index b6cd261d1..914a6047b 100644 --- a/pingora-core/Cargo.toml +++ b/pingora-core/Cargo.toml @@ -31,7 +31,7 @@ pingora-s2n = { version = "0.8.0", path = "../pingora-s2n", optional = true } bstr = { workspace = true } tokio = { workspace = true, features = ["net", "rt-multi-thread", "signal"] } tokio-stream = { workspace = true } -futures = "0.3" +futures = { workspace = true } async-trait = { workspace = true } httparse = { workspace = true } bytes = { workspace = true } @@ -47,7 +47,6 @@ strum = "0.26.2" strum_macros = "0.26.2" libc = "0.2.70" chrono = { version = "~0.4.31", features = ["alloc"], default-features = false } -prometheus = "0.13" sentry = { version = "0.36", features = [ "backtrace", "contexts", @@ -75,8 +74,9 @@ lru = { workspace = true, optional = true } daggy = "0.8" [target.'cfg(unix)'.dependencies] -daemonize = "0.5.0" -nix = "~0.24.3" +flurry = "0.5" +daemonix = "0.1.0" +nix = { version = "~0.31.1", features = ["socket", "net", "fs", "uio"] } [target.'cfg(windows)'.dependencies] windows-sys = { version = "0.59.0", features = ["Win32_Networking_WinSock"] } @@ -85,19 +85,26 @@ windows-sys = { version = "0.59.0", features = ["Win32_Networking_WinSock"] } h2 = { workspace = true, features = ["unstable"] } tokio-stream = { version = "0.1", features = ["full"] } env_logger = "0.11" -reqwest = { version = "0.11", features = [ +reqwest = { version = "0.12", features = [ "rustls-tls", + "http2", ], default-features = false } -hyper = "0.14" +hyper = { version = "1", features = ["client", "http1", "http2"] } +hyper-util = { version = "0.1", features = ["client-legacy", "http1", "http2"] } +http-body-util = "0.1" rstest = "0.23.0" rustls = "0.23" [target.'cfg(unix)'.dev-dependencies] -hyperlocal = "0.8" +hyperlocal = "0.9" jemallocator = "0.5" [features] default = [] +# Requires building with `--cfg tokio_unstable`. +dial9 = ["pingora-runtime/dial9"] +dial9-worker-s3 = ["dial9", "pingora-runtime/dial9-worker-s3"] +dial9-cpu-profiling = ["dial9", "pingora-runtime/dial9-cpu-profiling"] openssl = ["pingora-openssl", "openssl_derived"] boringssl = ["pingora-boringssl", "openssl_derived"] rustls = ["pingora-rustls", "any_tls", "dep:x509-parser", "ouroboros"] diff --git a/pingora-core/src/apps/http_app.rs b/pingora-core/src/apps/http_app.rs index f511012c0..0cd6e17d9 100644 --- a/pingora-core/src/apps/http_app.rs +++ b/pingora-core/src/apps/http_app.rs @@ -99,7 +99,7 @@ where } let persistent_settings = HttpPersistentSettings::for_session(&http); match http.finish().await { - Ok(c) => c.map(|s| ReusedHttpStream::new(s, Some(persistent_settings))), + Ok(c) => c.map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)), Err(e) => { error!("HTTP server fails to finish the request: {e}"); None @@ -207,7 +207,7 @@ where } let persistent_settings = HttpPersistentSettings::for_session(&http); match http.finish().await { - Ok(c) => c.map(|s| ReusedHttpStream::new(s, Some(persistent_settings))), + Ok(c) => c.map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)), Err(e) => { error!("HTTP server fails to finish the request: {e}"); None diff --git a/pingora-core/src/apps/mod.rs b/pingora-core/src/apps/mod.rs index d751fbcce..de8c5d1df 100644 --- a/pingora-core/src/apps/mod.rs +++ b/pingora-core/src/apps/mod.rs @@ -15,16 +15,16 @@ //! The abstraction and implementation interface for service application logic pub mod http_app; -pub mod prometheus_http_app; use crate::server::ShutdownWatch; use async_trait::async_trait; +use bytes::BytesMut; use log::{debug, error}; -use std::future::poll_fn; +use std::any::Any; use std::sync::Arc; use crate::protocols::http::v2::server; -use crate::protocols::http::ServerSession; +use crate::protocols::http::{ReusableHttpStream, ServerSession}; use crate::protocols::Digest; use crate::protocols::Stream; use crate::protocols::ALPN; @@ -58,7 +58,7 @@ pub trait ServerApp { async fn cleanup(&self) {} } #[non_exhaustive] -#[derive(Default)] +#[derive(Clone, Debug, Default)] /// HTTP Server options that control how the server handles some transport types. pub struct HttpServerOptions { /// Allow HTTP/2 for plaintext. @@ -84,10 +84,32 @@ pub struct HttpServerOptions { pub keepalive_request_limit: Option, } -#[derive(Debug, Clone)] +/// Settings persisted across HTTP/1.x keepalive requests on the same downstream connection. +/// +/// In addition to framework-managed keepalive parameters, this struct can carry an optional +/// user-defined context via [`set_user_context`](Self::set_user_context). The proxy layer +/// populates this through `ProxyHttp::persist_connection_context` +/// and delivers it to the next request through `ProxyHttp::on_connection_reuse`. +/// +/// Also carries pipelined-prefix bytes when the caller has opted into HTTP/1.1 +/// pipelining on the previous session. See +/// [`Self::set_pipelined_prefix`] and the +/// [`HttpSession::set_pipelining_enabled`](crate::protocols::http::v1::server::HttpSession::set_pipelining_enabled) +/// docs for the RFC 9112 §9.3.2 semantics. +#[derive(Debug)] pub struct HttpPersistentSettings { keepalive_timeout: Option, keepalive_reuses_remaining: Option, + /// User-defined context to carry to the next request on this connection. + user_context: Option>, + /// Bytes read past the end of the previous request's body, to be parsed + /// as the next pipelined request on the reused connection. + pipelined_prefix: Option, + /// Whether HTTP/1.1 pipelining was enabled on the previous session; + /// propagates to the next session so the proxy-level opt-in sticks + /// across keepalive reuses without the adopter having to re-enable + /// it on every request. + pipelining_enabled: bool, } impl HttpPersistentSettings { @@ -95,13 +117,37 @@ impl HttpPersistentSettings { HttpPersistentSettings { keepalive_timeout: session.get_keepalive(), keepalive_reuses_remaining: session.get_keepalive_reuses_remaining(), + user_context: None, + pipelined_prefix: None, + pipelining_enabled: session.pipelining_enabled(), } } + /// Set a user-defined context to be carried to the next request on this connection. + pub fn set_user_context(&mut self, ctx: Box) { + self.user_context = Some(ctx); + } + + /// Take the user-defined context, if any. + pub fn take_user_context(&mut self) -> Option> { + self.user_context.take() + } + + /// Set pipelined-prefix bytes to be fed to the next session on this + /// connection. Called by the proxy layer when HTTP/1.1 pipelining is + /// enabled on the current session and overread bytes were present at + /// reuse time. + pub fn set_pipelined_prefix(&mut self, prefix: BytesMut) { + self.pipelined_prefix = Some(prefix); + } + pub fn apply_to_session(self, session: &mut ServerSession) { let Self { keepalive_timeout, mut keepalive_reuses_remaining, + user_context, + pipelined_prefix, + pipelining_enabled, } = self; // Reduce the number of times the connection for this session can be @@ -112,6 +158,18 @@ impl HttpPersistentSettings { session.set_keepalive(keepalive_timeout); session.set_keepalive_reuses_remaining(keepalive_reuses_remaining); + + // Carry user context into the session for the proxy layer to consume + session.set_connection_user_context(user_context); + + // Replay pipelining opt-in so it stays on across keepalive reuses. + session.set_pipelining_enabled(pipelining_enabled); + + // Feed any pipelined prefix bytes to the new session's request parser + // so they are treated as the start of the next request. + if let Some(prefix) = pipelined_prefix { + session.set_pipelined_prefix(prefix); + } } } @@ -129,6 +187,19 @@ impl ReusedHttpStream { } } + /// Build a reusable HTTP stream from a finished session, preserving any + /// pipelined prefix bytes in the persistent settings for the next request. + pub fn from_reusable_stream( + reusable: ReusableHttpStream, + mut persistent_settings: HttpPersistentSettings, + ) -> Self { + let (stream, pipelined_prefix) = reusable.into_parts(); + if let Some(prefix) = pipelined_prefix { + persistent_settings.set_pipelined_prefix(prefix); + } + Self::new(stream, Some(persistent_settings)) + } + pub fn consume(self) -> (Stream, Option) { (self.stream, self.persistent_settings) } @@ -139,9 +210,11 @@ impl ReusedHttpStream { pub trait HttpServerApp { /// Similar to the [`ServerApp`], this function is called whenever a new HTTP session is established. /// - /// After successful processing, [`ServerSession::finish()`] can be called to return an optionally reusable - /// connection back to the service. The caller needs to make sure that the connection is in a reusable state - /// i.e., no error or incomplete read or write headers or bodies. Otherwise a `None` should be returned. + /// After successful processing, [`ServerSession::finish()`] can be + /// called to return an optionally reusable connection back to the service. + /// The caller needs to make sure that the connection is in a reusable state + /// i.e., no error or incomplete read or write headers or bodies. Otherwise + /// a `None` should be returned. async fn process_new_http( self: &Arc, mut session: ServerSession, @@ -193,8 +266,13 @@ where .as_ref() .map_or(false, |o| o.force_custom); + // h2c is for cleartext connections; on TLS, ALPN handles protocol negotiation. + // Otherwise, h2c stays true on TLS streams, forcing HTTP/1.1 clients into HTTP/2 + if stream.get_ssl_digest().is_some() { + h2c = false; + } // try to read h2 preface - if h2c && !custom { + else if h2c && !custom { let mut buf = [0u8; H2_PREFACE.len()]; let peeked = stream .try_peek(&mut buf) @@ -222,8 +300,7 @@ where }); let h2_options = self.h2_options(); - let h2_conn = server::handshake(stream, h2_options).await; - let mut h2_conn = match h2_conn { + let h2_conn = match server::handshake(stream, h2_options).await { Err(e) => { error!("H2 handshake error {e}"); return None; @@ -231,36 +308,21 @@ where Ok(c) => c, }; - let mut shutdown = shutdown.clone(); - loop { - // this loop ends when the client decides to close the h2 conn - // TODO: add a timeout? - let h2_stream = tokio::select! { - _ = shutdown.changed() => { - h2_conn.graceful_shutdown(); - let _ = poll_fn(|cx| h2_conn.poll_closed(cx)) - .await.map_err(|e| error!("H2 error waiting for shutdown {e}")); - return None; - } - h2_stream = server::HttpSession::from_h2_conn(&mut h2_conn, digest.clone()) => h2_stream - }; - let h2_stream = match h2_stream { - Err(e) => { - // It is common for the client to just disconnect TCP without properly - // closing H2. So we don't log the errors here - debug!("H2 error when accepting new stream {e}"); - return None; - } - Ok(s) => s?, // None means the connection is ready to be closed - }; - let app = self.clone(); - let shutdown = shutdown.clone(); + // The accept-loop body — including the graceful-shutdown state + // machine — lives in `server::accept_downstream_sessions` so that + // the same code path is exercised by tests in `protocols::http::v2`. + let app = self.clone(); + let shutdown_for_session = shutdown.clone(); + server::accept_downstream_sessions(h2_conn, digest, shutdown.clone(), |h2_stream| { + let app = app.clone(); + let shutdown = shutdown_for_session.clone(); pingora_runtime::current_handle().spawn(async move { // Note, `PersistentSettings` not currently relevant for h2 app.process_new_http(ServerSession::new_http2(h2_stream), &shutdown) .await; }); - } + }) + .await; } else if custom || matches!(stream.selected_alpn_proto(), Some(ALPN::Custom(_))) { return self.clone().process_custom_session(stream, shutdown).await; } else { @@ -295,3 +357,56 @@ where self.http_cleanup().await; } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio_test::io::Builder; + + #[test] + fn test_persistent_settings_user_context_roundtrip() { + // Create a mock H1 session + let mock_io = Builder::new().build(); + let mut session = ServerSession::new_http1(Box::new(mock_io)); + session.set_keepalive(Some(60)); + + // Snapshot settings (no user context yet) + let mut settings = HttpPersistentSettings::for_session(&session); + assert!(settings.take_user_context().is_none()); + + // Set user context + settings.set_user_context(Box::new(123u64)); + + // Apply to a fresh session -- user context should transfer + let mock_io2 = Builder::new().build(); + let mut session2 = ServerSession::new_http1(Box::new(mock_io2)); + settings.apply_to_session(&mut session2); + + // The user context should now be on the session + let ctx = session2.take_connection_user_context(); + assert!(ctx.is_some()); + let val = ctx.unwrap().downcast::().unwrap(); + assert_eq!(*val, 123u64); + + // Keepalive should also have been applied + assert_eq!(session2.get_keepalive(), Some(60)); + } + + #[test] + fn test_persistent_settings_no_user_context_by_default() { + let mock_io = Builder::new().build(); + let mut session = ServerSession::new_http1(Box::new(mock_io)); + session.set_keepalive(Some(30)); + + let settings = HttpPersistentSettings::for_session(&session); + + let mock_io2 = Builder::new().build(); + let mut session2 = ServerSession::new_http1(Box::new(mock_io2)); + settings.apply_to_session(&mut session2); + + // No user context should be present + assert!(session2.take_connection_user_context().is_none()); + // Keepalive should still work + assert_eq!(session2.get_keepalive(), Some(30)); + } +} diff --git a/pingora-core/src/apps/prometheus_http_app.rs b/pingora-core/src/apps/prometheus_http_app.rs deleted file mode 100644 index ed8a217a1..000000000 --- a/pingora-core/src/apps/prometheus_http_app.rs +++ /dev/null @@ -1,60 +0,0 @@ -// Copyright 2026 Cloudflare, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! An HTTP application that reports Prometheus metrics. - -use async_trait::async_trait; -use http::Response; -use prometheus::{Encoder, TextEncoder}; - -use super::http_app::HttpServer; -use crate::apps::http_app::ServeHttp; -use crate::modules::http::compression::ResponseCompressionBuilder; -use crate::protocols::http::ServerSession; - -/// An HTTP application that reports Prometheus metrics. -/// -/// This application will report all the [static metrics](https://docs.rs/prometheus/latest/prometheus/index.html#static-metrics) -/// collected via the [Prometheus](https://docs.rs/prometheus/) crate; -pub struct PrometheusHttpApp; - -#[async_trait] -impl ServeHttp for PrometheusHttpApp { - async fn response(&self, _http_session: &mut ServerSession) -> Response> { - let encoder = TextEncoder::new(); - let metric_families = prometheus::gather(); - let mut buffer = vec![]; - encoder.encode(&metric_families, &mut buffer).unwrap(); - Response::builder() - .status(200) - .header(http::header::CONTENT_TYPE, encoder.format_type()) - .header(http::header::CONTENT_LENGTH, buffer.len()) - .body(buffer) - .unwrap() - } -} - -/// The [HttpServer] for [PrometheusHttpApp] -/// -/// This type provides the functionality of [PrometheusHttpApp] with compression enabled -pub type PrometheusServer = HttpServer; - -impl PrometheusServer { - pub fn new() -> Self { - let mut server = Self::new_app(PrometheusHttpApp); - // enable gzip level 7 compression - server.add_module(ResponseCompressionBuilder::enable(7)); - server - } -} diff --git a/pingora-core/src/connectors/http/mod.rs b/pingora-core/src/connectors/http/mod.rs index 2545cf7cb..5a671ef2a 100644 --- a/pingora-core/src/connectors/http/mod.rs +++ b/pingora-core/src/connectors/http/mod.rs @@ -21,6 +21,8 @@ use crate::protocols::http::client::HttpSession; use crate::protocols::http::v1::client::HttpSession as Http1Session; use crate::upstreams::peer::Peer; use pingora_error::Result; +use std::sync::atomic::AtomicU64; +use std::sync::Arc; use std::time::Duration; pub mod custom; @@ -151,6 +153,17 @@ where pub fn prefer_h1(&self, peer: &impl Peer) { self.h2.prefer_h1(peer); } + + /// Return the number of times a pooled connection was found to contain + /// unexpected data from the server. + pub fn unexpected_data_connection_count(&self) -> u64 { + self.h1.unexpected_data_connection_count() + } + + /// Return a shared reference to the unexpected data connection counter for periodic metric reporting. + pub fn unexpected_data_connection_counter(&self) -> Arc { + self.h1.unexpected_data_connection_counter() + } } #[cfg(test)] diff --git a/pingora-core/src/connectors/http/v1.rs b/pingora-core/src/connectors/http/v1.rs index 62ecfcb6c..ab04b2f64 100644 --- a/pingora-core/src/connectors/http/v1.rs +++ b/pingora-core/src/connectors/http/v1.rs @@ -17,6 +17,8 @@ use crate::protocols::http::v1::client::HttpSession; use crate::upstreams::peer::Peer; use pingora_error::Result; +use std::sync::atomic::AtomicU64; +use std::sync::Arc; use std::time::Duration; pub struct Connector { @@ -60,6 +62,17 @@ impl Connector { .release_stream(stream, peer.reuse_hash(), idle_timeout); } } + + /// Return the number of times a pooled connection was found to contain + /// unexpected data from the server. + pub fn unexpected_data_connection_count(&self) -> u64 { + self.transport.unexpected_data_connection_count() + } + + /// Return a shared reference to the unexpected data connection counter for periodic metric reporting. + pub fn unexpected_data_connection_counter(&self) -> Arc { + self.transport.unexpected_data_connection_counter() + } } #[cfg(test)] diff --git a/pingora-core/src/connectors/http/v2.rs b/pingora-core/src/connectors/http/v2.rs index c8e804d46..207109bd5 100644 --- a/pingora-core/src/connectors/http/v2.rs +++ b/pingora-core/src/connectors/http/v2.rs @@ -13,7 +13,7 @@ // limitations under the License. use super::HttpSession; -use crate::connectors::{ConnectorOptions, TransportConnector}; +use crate::connectors::{ConnectorOptions, IdleConnection, PoolCallback, TransportConnector}; use crate::protocols::http::custom::client::Session; use crate::protocols::http::v1::client::HttpSession as Http1Session; use crate::protocols::http::v2::client::{drive_connection, Http2Session}; @@ -24,9 +24,10 @@ use bytes::Bytes; use h2::client::SendRequest; use log::debug; use parking_lot::{Mutex, RwLock}; -use pingora_error::{Error, ErrorType::*, OrErr, Result}; +use pingora_error::{Error, ErrorType::*, OkOrErr, OrErr, Result}; use pingora_pool::{ConnectionMeta, ConnectionPool, PoolNode}; use std::collections::HashMap; +use std::io::ErrorKind; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -128,6 +129,14 @@ impl ConnectionRef { self.0.shutting_down.load(Ordering::Relaxed) } + /// Mark this connection for shutdown. + /// + /// No new streams will be created on this connection and + /// it will be discarded once all active streams are released. + pub fn mark_shutdown(&self) { + self.0.shutting_down.store(true, Ordering::Relaxed); + } + // spawn a stream if more stream is allowed, otherwise return Ok(None) pub async fn spawn_stream(&self) -> Result> { // Atomically check if the current_stream is over the limit @@ -144,16 +153,24 @@ impl ConnectionRef { Err(e) => { // fail to create the stream, reset the counter self.0.current_streams.fetch_sub(1, Ordering::SeqCst); - // Remote sends GOAWAY(NO_ERROR): graceful shutdown: this connection no longer - // accepts new streams. We can still try to create new connection. - if e.root_cause() + + // Check for graceful shutdown conditions where we can retry with a new connection + let is_graceful_shutdown = e + .root_cause() .downcast_ref::() .map(|e| { - e.is_go_away() && e.is_remote() && e.reason() == Some(h2::Reason::NO_ERROR) + // Remote sends GOAWAY(NO_ERROR): graceful shutdown + (e.is_go_away() && e.is_remote() && e.reason() == Some(h2::Reason::NO_ERROR)) + // Or broken pipe wrapped inside an h2::Error: stream closed unexpectedly + || (e.is_io() + && e.get_io() + .map(|io| io.kind() == ErrorKind::BrokenPipe) + .unwrap_or(false)) }) - .unwrap_or(false) - { - self.0.shutting_down.store(true, Ordering::Relaxed); + .unwrap_or(false); + + if is_graceful_shutdown { + self.mark_shutdown(); Ok(None) } else { Err(e) @@ -264,6 +281,7 @@ pub struct Connector { idle_pool: Arc>, // the pool of h2 connections that have ongoing streams in_use_pool: InUsePool, + pool_callback: Option, } impl Connector { @@ -272,11 +290,15 @@ impl Connector { let pool_size = options .as_ref() .map_or(DEFAULT_POOL_SIZE, |o| o.keepalive_pool_size); + let pool_callback = options + .as_ref() + .and_then(|o| o.keepalive_pool_callback.clone()); // connection offload is handled by the [TransportConnector] Connector { transport: TransportConnector::new(options), idle_pool: Arc::new(ConnectionPool::new(pool_size)), in_use_pool: InUsePool::new(), + pool_callback, } } @@ -326,12 +348,17 @@ impl Connector { // the caller that the server speaks h2c } } - let max_h2_stream = peer.get_peer_options().map_or(1, |o| o.max_h2_streams); - let conn = handshake(stream, max_h2_stream, peer.h2_ping_interval()).await?; - let h2_stream = conn - .spawn_stream() - .await? - .expect("newly created connections should have at least one free stream"); + let peer_options = peer.get_peer_options(); + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = peer_options.map_or(1, |o| o.max_h2_streams); + settings.ping_interval = peer.h2_ping_interval(); + settings.stream_window_size = peer_options.and_then(|o| o.h2_stream_window_size); + settings.connection_window_size = peer_options.and_then(|o| o.h2_connection_window_size); + let conn = handshake(stream, settings).await?; + let h2_stream = conn.spawn_stream().await?.or_err( + H2Error, + "newly created connection has no free streams (server may have sent GOAWAY)", + )?; if conn.more_streams_allowed() { self.in_use_pool.insert(peer.reuse_hash(), conn); } @@ -434,11 +461,25 @@ impl Connector { }; let closed = conn.0.closed.clone(); let (notify_evicted, watch_use) = self.idle_pool.put(&meta, conn); + let idle_meta = IdleConnection::new(meta); let pool = self.idle_pool.clone(); //clone the arc + let keepalive_pool_callback = self.pool_callback.clone(); let rt = pingora_runtime::current_handle(); rt.spawn(async move { - pool.idle_timeout(&meta, idle_timeout, notify_evicted, closed, watch_use) - .await; + if pool + .idle_timeout( + &idle_meta.connection, + idle_timeout, + notify_evicted, + closed, + watch_use, + ) + .await + { + if let Some(callback) = keepalive_pool_callback { + callback(idle_meta.elapsed()); + } + } }); } else { self.in_use_pool.insert(reuse_hash, conn); @@ -467,19 +508,84 @@ impl Connector { // 8 Mbytes = 80 Mbytes X 100ms, which should be enough for most links. const H2_WINDOW_SIZE: u32 = 1 << 23; -pub async fn handshake( - stream: Stream, - max_streams: usize, - h2_ping_interval: Option, -) -> Result { +/// Maximum allowed H2 window size per [RFC 9113 §6.9.1](https://datatracker.ietf.org/doc/html/rfc9113#section-6.9.1-7). +const H2_MAX_WINDOW_SIZE: u32 = (1u32 << 31) - 1; + +/// Settings for HTTP/2 handshake. +/// +/// # Example +/// +/// ```rust,ignore +/// use pingora_core::connectors::http::v2::{handshake, H2HandshakeSettings}; +/// +/// // With custom window sizes +/// let mut settings = H2HandshakeSettings::new(); +/// settings.max_streams = 100; +/// settings.stream_window_size = Some(1 << 20); // 1MiB +/// settings.connection_window_size = Some(1 << 24); // 16MiB +/// let conn = handshake(stream, settings).await?; +/// ``` +#[derive(Debug, Clone, Default)] +#[non_exhaustive] +pub struct H2HandshakeSettings { + /// The maximum number of concurrent streams allowed on this connection. + pub max_streams: usize, + /// Optional interval for sending H2 ping frames to keep the connection alive. + pub ping_interval: Option, + /// Optional initial per-stream receive window size in bytes. + /// If `None`, the default of 8MB is used. + pub stream_window_size: Option, + /// Optional initial connection-level receive window size in bytes. + /// If `None`, the default of 8MB is used. + pub connection_window_size: Option, +} + +impl H2HandshakeSettings { + /// Create a new `H2HandshakeSettings` with all defaults. + pub fn new() -> Self { + Self::default() + } +} + +/// Perform an HTTP/2 handshake on the given stream with the given settings. +pub async fn handshake(stream: Stream, settings: H2HandshakeSettings) -> Result { use h2::client::Builder; use pingora_runtime::current_handle; + let max_streams = settings.max_streams; + // Safe guard: new_http_session() assumes there should be at least one free stream if max_streams == 0 { return Error::e_explain(H2Error, "zero max_stream configured"); } + // Validate window sizes against RFC 9113 §6.9.1 limit + // https://datatracker.ietf.org/doc/html/rfc9113#section-6.9.1-7 + if settings + .stream_window_size + .is_some_and(|w| w == 0 || w > H2_MAX_WINDOW_SIZE) + { + return Error::e_explain( + H2Error, + format!( + "stream_window_size must be between 1 and {} (2^31-1)", + H2_MAX_WINDOW_SIZE + ), + ); + } + if settings + .connection_window_size + .is_some_and(|w| w == 0 || w > H2_MAX_WINDOW_SIZE) + { + return Error::e_explain( + H2Error, + format!( + "connection_window_size must be between 1 and {} (2^31-1)", + H2_MAX_WINDOW_SIZE + ), + ); + } + let id = stream.id(); let digest = Digest { // NOTE: this field is always false because the digest is shared across all streams @@ -490,16 +596,16 @@ pub async fn handshake( proxy_digest: stream.get_proxy_digest(), socket_digest: stream.get_socket_digest(), }; - // TODO: make these configurable + let stream_window = settings.stream_window_size.unwrap_or(H2_WINDOW_SIZE); + let conn_window = settings.connection_window_size.unwrap_or(H2_WINDOW_SIZE); let (send_req, connection) = Builder::new() .enable_push(false) .initial_max_send_streams(max_streams) // The limit for the server. Server push is not allowed, so this value doesn't matter .max_concurrent_streams(1) .max_frame_size(64 * 1024) // advise server to send larger frames - .initial_window_size(H2_WINDOW_SIZE) - // should this be max_streams * H2_WINDOW_SIZE? - .initial_connection_window_size(H2_WINDOW_SIZE) + .initial_window_size(stream_window) + .initial_connection_window_size(conn_window) .handshake(stream) .await .or_err(HandshakeError, "during H2 handshake")?; @@ -521,7 +627,7 @@ pub async fn handshake( connection, id, closed_tx, - h2_ping_interval, + settings.ping_interval, ping_timeout_clone, ) .await; @@ -541,6 +647,9 @@ pub async fn handshake( mod tests { use super::*; use crate::upstreams::peer::HttpPeer; + use bytes::Bytes; + use http::{Response, StatusCode}; + use pingora_http::RequestHeader; #[tokio::test] #[cfg(feature = "any_tls")] @@ -674,6 +783,63 @@ mod tests { assert_eq!(id, h2_5.conn.id()); } + /// `spawn_stream` must return `Ok(None)` and mark the connection as shutting + /// down when the underlying I/O channel is closed (BrokenPipe). This + /// exercises the BrokenPipe branch of `spawn_stream` directly without going + /// through the full proxy stack. + #[tokio::test] + async fn test_spawn_stream_broken_pipe_marks_shutdown() { + let (client_io, server_io) = tokio::io::duplex(65536); + let (send_req, connection) = h2::client::handshake(client_io).await.unwrap(); + let (closed_tx, closed_rx) = watch::channel(false); + let ping_timeout = Arc::new(AtomicBool::new(false)); + let conn = ConnectionRef::new(send_req, closed_rx, ping_timeout, 0, 10, Digest::default()); + + // Drive the H2 client connection task in the background. + // When the connection terminates it will signal via closed_tx. + let conn_handle = tokio::spawn(async move { + let _ = connection.await; + // Signal that the connection task has finished. + let _ = closed_tx.send(true); + }); + + // Complete the server-side H2 handshake, then drop the server connection. + // Dropping server_conn closes the write end of the duplex, so the client + // connection task will read EOF and terminate with BrokenPipe. + let server_conn = h2::server::handshake(server_io).await.unwrap(); + drop(server_conn); + + // Wait until the client connection task has fully processed the EOF. + conn_handle.await.unwrap(); + + // spawn_stream must detect BrokenPipe, mark shutdown, and return Ok(None) + // so the caller can retry on a fresh connection rather than propagating the error. + let result = conn.spawn_stream().await; + assert!(result.is_ok(), "expected Ok(None), got Err"); + assert!(result.unwrap().is_none(), "expected None stream"); + assert!( + conn.is_shutting_down(), + "connection should be marked as shutting down" + ); + } + + #[tokio::test] + async fn test_mark_shutdown_prevents_new_streams() { + let (client_io, _server_io) = tokio::io::duplex(65536); + let (send_req, _connection) = h2::client::handshake(client_io).await.unwrap(); + let (_closed_tx, closed_rx) = watch::channel(false); + let ping_timeout = Arc::new(AtomicBool::new(false)); + let conn = ConnectionRef::new(send_req, closed_rx, ping_timeout, 0, 10, Digest::default()); + + assert!(conn.more_streams_allowed()); + assert!(!conn.is_shutting_down()); + + conn.mark_shutdown(); + + assert!(conn.is_shutting_down()); + assert!(!conn.more_streams_allowed()); + } + #[cfg(all(feature = "any_tls", unix))] #[tokio::test] async fn test_h2_reuse_rejects_fd_mismatch() { @@ -744,4 +910,142 @@ mod tests { .unwrap() .is_none()); } + + #[tokio::test] + async fn test_h2_handshake_settings_validation() { + use super::H2HandshakeSettings; + + // Test zero stream window size is rejected + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = 100; + settings.stream_window_size = Some(0); + let (client, _server) = tokio::io::duplex(65536); + match handshake(Box::new(client), settings).await { + Err(e) => assert!( + e.to_string() + .contains("stream_window_size must be between 1"), + "Unexpected error: {}", + e + ), + Ok(_) => panic!("Expected error for stream_window_size = 0"), + } + + // Test zero connection window size is rejected + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = 100; + settings.connection_window_size = Some(0); + let (client, _server) = tokio::io::duplex(65536); + match handshake(Box::new(client), settings).await { + Err(e) => assert!( + e.to_string() + .contains("connection_window_size must be between 1"), + "Unexpected error: {}", + e + ), + Ok(_) => panic!("Expected error for connection_window_size = 0"), + } + + // Test exceeding max stream window size is rejected + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = 100; + settings.stream_window_size = Some(super::H2_MAX_WINDOW_SIZE + 1); + let (client, _server) = tokio::io::duplex(65536); + match handshake(Box::new(client), settings).await { + Err(e) => assert!( + e.to_string() + .contains("stream_window_size must be between 1"), + "Unexpected error: {}", + e + ), + Ok(_) => panic!("Expected error for stream_window_size > max"), + } + + // Test exceeding max connection window size is rejected + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = 100; + settings.connection_window_size = Some(super::H2_MAX_WINDOW_SIZE + 1); + let (client, _server) = tokio::io::duplex(65536); + match handshake(Box::new(client), settings).await { + Err(e) => assert!( + e.to_string() + .contains("connection_window_size must be between 1"), + "Unexpected error: {}", + e + ), + Ok(_) => panic!("Expected error for connection_window_size > max"), + } + } + + #[tokio::test] + async fn test_h2_handshake_custom_window_sizes() { + // Test that valid custom window sizes are accepted and handshake succeeds + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = 100; + settings.stream_window_size = Some(1 << 20); // 1MiB + settings.connection_window_size = Some(1 << 24); // 16MiB + + let (client, server) = tokio::io::duplex(65536); + + // Spawn server side + tokio::spawn(async move { + let mut server_conn = h2::server::handshake(server).await.unwrap(); + if let Some(result) = server_conn.accept().await { + let (_request, mut respond) = result.unwrap(); + let resp = Response::builder().status(StatusCode::OK).body(()).unwrap(); + let mut stream = respond.send_response(resp, false).unwrap(); + stream.send_data(Bytes::from("ok"), true).unwrap(); + server_conn.graceful_shutdown(); + } + // Drive the server connection until the client closes + while let Some(_res) = server_conn.accept().await {} + }); + + // Client side - should succeed with custom window sizes + let conn = handshake(Box::new(client), settings).await.unwrap(); + + // Verify we can spawn a stream and complete a request/response cycle + let mut stream = conn.spawn_stream().await.unwrap().unwrap(); + let mut request = RequestHeader::build("GET", b"/", None).unwrap(); + request + .insert_header(http::header::HOST, "example.com") + .unwrap(); + stream + .write_request_header(Box::new(request), true) + .unwrap(); + + stream.read_response_header().await.unwrap(); + assert_eq!(stream.response_header().unwrap().status, 200); + } + + /// `spawn_stream()` must return `Ok(None)` when the server sends + /// GOAWAY(NO_ERROR) before any streams are opened. + #[tokio::test] + async fn test_spawn_stream_goaway_no_error_returns_none() { + let (client_io, server_io) = tokio::io::duplex(65536); + let (send_req, connection) = h2::client::handshake(client_io).await.unwrap(); + let (closed_tx, closed_rx) = watch::channel(false); + let ping_timeout = Arc::new(AtomicBool::new(false)); + let conn = ConnectionRef::new(send_req, closed_rx, ping_timeout, 0, 10, Digest::default()); + + let conn_handle = tokio::spawn(async move { + let _ = connection.await; + let _ = closed_tx.send(true); + }); + + let mut server_conn = h2::server::handshake(server_io).await.unwrap(); + server_conn.graceful_shutdown(); + let _ = server_conn.accept().await; + drop(server_conn); + + conn_handle.await.unwrap(); + + let result = conn.spawn_stream().await; + assert!( + result.is_ok(), + "expected Ok(None), got Err: {:?}", + result.as_ref().err() + ); + assert!(result.unwrap().is_none()); + assert!(conn.is_shutting_down()); + } } diff --git a/pingora-core/src/connectors/l4.rs b/pingora-core/src/connectors/l4.rs index bd7439d4b..d275030fd 100644 --- a/pingora-core/src/connectors/l4.rs +++ b/pingora-core/src/connectors/l4.rs @@ -313,15 +313,9 @@ mod tests { use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; - use tokio::io::AsyncWriteExt; use tokio::time::sleep; - /// Some of the tests below are flaky when making new connections to mock - /// servers. The servers are simple tokio listeners, so failures there are - /// not indicative of real errors. This function will retry the peer/server - /// in increasing intervals until it either succeeds in connecting or a long - /// timeout expires (max 10sec) - #[cfg(unix)] + #[cfg(target_os = "linux")] async fn wait_for_peer

(peer: &P) where P: Peer + Send + Sync, @@ -394,11 +388,17 @@ mod tests { #[tokio::test] async fn test_conn_timeout() { - // 192.0.2.1 is effectively a blackhole + // 192.0.2.1 is TEST-NET-1 (RFC 5737) — SYN packets are silently + // dropped on Linux, producing ConnectTimedout. On macOS the kernel + // may instead return ENETUNREACH (ConnectNoRoute). let mut peer = BasicPeer::new("192.0.2.1:79"); - peer.options.connection_timeout = Some(std::time::Duration::from_millis(1)); //1ms - let new_session = connect(&peer, None).await; - assert_eq!(new_session.unwrap_err().etype(), &ConnectTimedout) + peer.options.connection_timeout = Some(Duration::from_millis(1)); + let err = connect(&peer, None).await.unwrap_err(); + assert!( + err.etype() == &ConnectTimedout || err.etype() == &ConnectNoRoute, + "unexpected error type: {:?}", + err.etype() + ); } #[tokio::test] @@ -412,7 +412,7 @@ mod tests { let move_flag = Arc::clone(&flag); peer.options.upstream_tcp_sock_tweak_hook = Some(Arc::new(move |_| { - move_flag.fetch_xor(true, Ordering::SeqCst); + move_flag.fetch_not(Ordering::SeqCst); Ok(()) })); @@ -537,6 +537,7 @@ mod tests { // one-off mock server async fn mock_inet_connect_server() -> u16 { + use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/pingora-core/src/connectors/mod.rs b/pingora-core/src/connectors/mod.rs index 3e3c1c462..60acb7b40 100644 --- a/pingora-core/src/connectors/mod.rs +++ b/pingora-core/src/connectors/mod.rs @@ -37,10 +37,35 @@ use pingora_error::{Error, ErrorType::*, OrErr, Result}; use pingora_pool::{ConnectionMeta, ConnectionPool}; use std::collections::HashMap; use std::net::SocketAddr; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::time::{Duration, Instant}; use tls::TlsConnector; use tokio::sync::Mutex; +#[derive(Clone, Debug)] +pub(crate) struct IdleConnection { + connection: ConnectionMeta, + idle_since: Instant, +} + +impl IdleConnection { + pub(crate) fn new(connection: ConnectionMeta) -> Self { + Self { + connection, + idle_since: Instant::now(), + } + } + + fn elapsed(&self) -> Duration { + self.idle_since.elapsed() + } +} + +/// Callback invoked when an idle upstream connection leaves the keep-alive pool +/// without being reused. +pub type PoolCallback = Arc; + /// The options to configure a [TransportConnector] #[derive(Clone)] pub struct ConnectorOptions { @@ -64,7 +89,11 @@ pub struct ConnectorOptions { /// env variable. This can be used by tools like Wireshark to decrypt traffic /// for debugging purposes. pub debug_ssl_keylog: bool, - /// How many connections to keepalive + /// Effective global cap for the keepalive pool. Derived from + /// `server_conf.upstream_keepalive_pool_size * server_conf.threads` + /// in [`Self::from_server_conf`] so that operator-facing config keeps + /// its per-worker meaning even though the pool itself now uses a single + /// global LRU. pub keepalive_pool_size: usize, /// Optionally offload the connection establishment to dedicated thread pools /// @@ -79,6 +108,9 @@ pub struct ConnectorOptions { pub bind_to_v4: Vec, /// Bind to any of the given source IPv4 addresses pub bind_to_v6: Vec, + /// Optional callback for observing how long upstream connections stayed idle + /// before leaving the keep-alive pool without reuse. + pub keepalive_pool_callback: Option, } impl ConnectorOptions { @@ -115,10 +147,13 @@ impl ConnectorOptions { #[cfg(feature = "s2n")] s2n_config_cache_size: server_conf.s2n_config_cache_size, debug_ssl_keylog: server_conf.upstream_debug_ssl_keylog, - keepalive_pool_size: server_conf.upstream_keepalive_pool_size, + keepalive_pool_size: server_conf + .upstream_keepalive_pool_size + .saturating_mul(server_conf.threads.max(1)), offload_threadpool, bind_to_v4, bind_to_v6, + keepalive_pool_callback: None, } } @@ -134,6 +169,7 @@ impl ConnectorOptions { offload_threadpool: None, bind_to_v4: vec![], bind_to_v6: vec![], + keepalive_pool_callback: None, } } } @@ -146,6 +182,10 @@ pub struct TransportConnector { bind_to_v4: Vec, bind_to_v6: Vec, preferred_http_version: PreferredHttpVersion, + /// Wrapped in `Arc` so external consumers (e.g. proxy services) can clone a reference + /// for periodic metric reporting without needing access to the connector itself. + unexpected_data_conn_count: Arc, + keepalive_pool_callback: Option, } const DEFAULT_POOL_SIZE: usize = 128; @@ -165,6 +205,9 @@ impl TransportConnector { let bind_to_v6 = options .as_ref() .map_or_else(Vec::new, |o| o.bind_to_v6.clone()); + let keepalive_pool_callback = options + .as_ref() + .and_then(|o| o.keepalive_pool_callback.clone()); TransportConnector { tls_ctx: tls::Connector::new(options), connection_pool: Arc::new(ConnectionPool::new(pool_size)), @@ -172,6 +215,8 @@ impl TransportConnector { bind_to_v4, bind_to_v6, preferred_http_version: PreferredHttpVersion::new(), + unexpected_data_conn_count: Arc::new(AtomicU64::new(0)), + keepalive_pool_callback, } } @@ -212,7 +257,9 @@ impl TransportConnector { // test_reusable_stream: we assume server would never actively send data // first on an idle stream. #[cfg(unix)] - if peer.matches_fd(stream.id()) && test_reusable_stream(&mut stream) { + if peer.matches_fd(stream.id()) + && test_reusable_stream(&mut stream, &self.unexpected_data_conn_count) + { Some(stream) } else { None @@ -227,7 +274,10 @@ impl TransportConnector { } } if peer.matches_sock(WrappedRawSocket(stream.id() as RawSocket)) - && test_reusable_stream(&mut stream) + && test_reusable_stream( + &mut stream, + &self.unexpected_data_conn_count, + ) { Some(stream) } else { @@ -259,9 +309,9 @@ impl TransportConnector { &self, mut stream: Stream, key: u64, // usually peer.reuse_hash() - idle_timeout: Option, + idle_timeout: Option, ) { - if !test_reusable_stream(&mut stream) { + if !test_reusable_stream(&mut stream, &self.unexpected_data_conn_count) { return; } let id = stream.id(); @@ -270,11 +320,25 @@ impl TransportConnector { let stream = Arc::new(Mutex::new(stream)); let locked_stream = stream.clone().try_lock_owned().unwrap(); // safe as we just created it let (notify_close, watch_use) = self.connection_pool.put(&meta, stream); + let idle_meta = IdleConnection::new(meta); let pool = self.connection_pool.clone(); //clone the arc + let keepalive_pool_callback = self.keepalive_pool_callback.clone(); let rt = pingora_runtime::current_handle(); rt.spawn(async move { - pool.idle_poll(locked_stream, &meta, idle_timeout, notify_close, watch_use) - .await; + if pool + .idle_poll( + locked_stream, + &idle_meta.connection, + idle_timeout, + notify_close, + watch_use, + ) + .await + { + if let Some(callback) = keepalive_pool_callback { + callback(idle_meta.elapsed()); + } + } }); } @@ -301,6 +365,21 @@ impl TransportConnector { pub fn prefer_h1(&self, peer: &impl Peer) { self.preferred_http_version.add(peer, 1); } + + /// Return the number of times a pooled connection was found to contain unexpected data + /// from the server. + pub fn unexpected_data_connection_count(&self) -> u64 { + self.unexpected_data_conn_count.load(Ordering::Relaxed) + } + + /// Return a shared reference to the unexpected data connection counter. + /// + /// This allows external consumers (e.g. proxy services) to clone the `Arc` and + /// periodically read the counter for metric reporting without needing ongoing + /// access to the connector. + pub fn unexpected_data_connection_counter(&self) -> Arc { + self.unexpected_data_conn_count.clone() + } } // Perform the actual L4 and tls connection steps while respecting the peer's @@ -376,7 +455,7 @@ use futures::future::FutureExt; use tokio::io::AsyncReadExt; /// Test whether a stream is already closed or not reusable (server sent unexpected data) -fn test_reusable_stream(stream: &mut Stream) -> bool { +fn test_reusable_stream(stream: &mut Stream, unexpected_data_conn_count: &AtomicU64) -> bool { let mut buf = [0; 1]; // tokio::task::unconstrained because now_or_never may yield None when the future is ready let result = tokio::task::unconstrained(stream.read(&mut buf[..])).now_or_never(); @@ -387,6 +466,7 @@ fn test_reusable_stream(stream: &mut Stream) -> bool { debug!("Idle connection is closed"); } else { warn!("Unexpected data read in idle connection"); + unexpected_data_conn_count.fetch_add(1, Ordering::Relaxed); } } Err(e) => { @@ -482,15 +562,14 @@ pub(crate) mod test_utils { #[cfg(test)] #[cfg(feature = "any_tls")] mod tests { + use std::time::Duration; + use pingora_error::ErrorType; use tls::Connector; use super::*; use crate::upstreams::peer::BasicPeer; - // 192.0.2.1 is effectively a black hole - const BLACK_HOLE: &str = "192.0.2.1:79"; - #[tokio::test] async fn test_connect() { let connector = TransportConnector::new(None); @@ -547,15 +626,23 @@ mod tests { server_handle.await.unwrap(); } + // 192.0.2.1 is TEST-NET-1 (RFC 5737) — SYN packets are silently + // dropped on Linux, producing ConnectTimedout. On macOS the kernel + // may instead return ENETUNREACH (ConnectNoRoute). + const BLACKHOLE: &str = "192.0.2.1:79"; + async fn do_test_conn_timeout(conf: Option) { let connector = TransportConnector::new(conf); - let mut peer = BasicPeer::new(BLACK_HOLE); - peer.options.connection_timeout = Some(std::time::Duration::from_millis(1)); - let stream = connector.new_stream(&peer).await; - match stream { - Ok(_) => panic!("should throw an error"), - Err(e) => assert_eq!(e.etype(), &ConnectTimedout), - } + let mut peer = BasicPeer::new(BLACKHOLE); + peer.options.connection_timeout = Some(Duration::from_millis(1)); + let Err(e) = connector.new_stream(&peer).await else { + panic!("should throw an error"); + }; + assert!( + e.etype() == &ConnectTimedout || e.etype() == &ConnectNoRoute, + "unexpected error type: {:?}", + e.etype() + ); } #[tokio::test] @@ -580,13 +667,21 @@ mod tests { let stream = connector.new_stream(&peer).await; let error = stream.unwrap_err(); - // XXX: some systems will allow the socket to bind and connect without error, only to timeout - assert!(error.etype() == &ConnectError || error.etype() == &ConnectTimedout) + // The exact error varies by platform: Linux may return ConnectError, + // some systems time out (ConnectTimedout), and macOS/others may + // return ConnectNoRoute (ENETUNREACH) for unreachable addresses. + assert!( + error.etype() == &ConnectError + || error.etype() == &ConnectTimedout + || error.etype() == &ConnectNoRoute, + "unexpected error type: {:?}", + error.etype() + ) } /// Helper function for testing error handling in the `do_connect` function. - /// This assumes that the connection will fail to on the peer and returns - /// the decomposed error type and message + /// This assumes that the connection will fail on the peer and returns + /// the decomposed error type and message. async fn get_do_connect_failure_with_peer(peer: &BasicPeer) -> (ErrorType, String) { let tls_connector = Connector::new(None); let stream = do_connect(peer, None, None, &tls_connector.ctx).await; @@ -604,27 +699,65 @@ mod tests { #[tokio::test] async fn test_do_connect_with_total_timeout() { - let mut peer = BasicPeer::new(BLACK_HOLE); - peer.options.total_connection_timeout = Some(std::time::Duration::from_millis(1)); + let mut peer = BasicPeer::new(BLACKHOLE); + peer.options.total_connection_timeout = Some(Duration::from_millis(1)); let (etype, context) = get_do_connect_failure_with_peer(&peer).await; - assert_eq!(etype, ConnectTimedout); - assert!(context.contains("total-connection timeout")); + assert!( + etype == ConnectTimedout || etype == ConnectNoRoute, + "unexpected error type: {etype:?}" + ); + if etype == ConnectTimedout { + assert!(context.contains("total-connection timeout")); + } } #[tokio::test] async fn test_tls_connect_timeout_supersedes_total() { - let mut peer = BasicPeer::new(BLACK_HOLE); - peer.options.total_connection_timeout = Some(std::time::Duration::from_millis(10)); - peer.options.connection_timeout = Some(std::time::Duration::from_millis(1)); + let mut peer = BasicPeer::new(BLACKHOLE); + peer.options.total_connection_timeout = Some(Duration::from_millis(10)); + peer.options.connection_timeout = Some(Duration::from_millis(1)); let (etype, context) = get_do_connect_failure_with_peer(&peer).await; - assert_eq!(etype, ConnectTimedout); - assert!(!context.contains("total-connection timeout")); + assert!( + etype == ConnectTimedout || etype == ConnectNoRoute, + "unexpected error type: {etype:?}" + ); + if etype == ConnectTimedout { + assert!(!context.contains("total-connection timeout")); + } } #[tokio::test] async fn test_do_connect_without_total_timeout() { - let peer = BasicPeer::new(BLACK_HOLE); + let peer = BasicPeer::new(BLACKHOLE); let (etype, context) = get_do_connect_failure_with_peer(&peer).await; assert!(etype != ConnectTimedout || !context.contains("total-connection timeout")); } + + #[tokio::test] + async fn test_unexpected_data_connection_count_increments() { + // Create a duplex stream where we control both ends + let (mut server, client) = tokio::io::duplex(64); + + let counter = AtomicU64::new(0); + let mut stream: Stream = Box::new(client); + + // With no data available, the stream should be considered reusable + assert!(test_reusable_stream(&mut stream, &counter)); + assert_eq!(counter.load(Ordering::Relaxed), 0); + + // Write unexpected data from the server side + use tokio::io::AsyncWriteExt; + server.write_all(b"unexpected").await.unwrap(); + + // Give the data a moment to be buffered + tokio::task::yield_now().await; + + // Now test_reusable_stream should detect the unexpected data + assert!(!test_reusable_stream(&mut stream, &counter)); + assert_eq!( + counter.load(Ordering::Relaxed), + 1, + "unexpected_data_connection_count should have incremented" + ); + } } diff --git a/pingora-core/src/connectors/tls/boringssl_openssl/mod.rs b/pingora-core/src/connectors/tls/boringssl_openssl/mod.rs index 9bb3a5a68..8585deef7 100644 --- a/pingora-core/src/connectors/tls/boringssl_openssl/mod.rs +++ b/pingora-core/src/connectors/tls/boringssl_openssl/mod.rs @@ -193,7 +193,7 @@ where } } - if let Some(curve) = peer.get_peer_options().and_then(|o| o.curves) { + if let Some(curve) = peer.get_peer_options().and_then(|o| o.curves.as_deref()) { ssl_set_groups_list(&mut ssl_conf, curve).or_err(InternalError, "invalid curves")?; } diff --git a/pingora-core/src/connectors/tls/mod.rs b/pingora-core/src/connectors/tls/mod.rs index c49be80b1..5d75789b4 100644 --- a/pingora-core/src/connectors/tls/mod.rs +++ b/pingora-core/src/connectors/tls/mod.rs @@ -41,7 +41,7 @@ pub use rustls::*; /// > characters only letters, digits, and hyphen. There are also some /// > restrictions on the length. Labels must be 63 characters or less. /// - https://datatracker.ietf.org/doc/html/rfc1034#section-3.5 -#[cfg(feature = "any_tls")] +#[cfg(any(feature = "openssl_derived", feature = "rustls"))] pub fn replace_leftmost_underscore(sni: &str) -> Option { // wildcard is only leftmost label if let Some((leftmost, rest)) = sni.split_once('.') { @@ -56,7 +56,7 @@ pub fn replace_leftmost_underscore(sni: &str) -> Option { None } -#[cfg(feature = "any_tls")] +#[cfg(any(feature = "openssl_derived", feature = "rustls"))] #[cfg(test)] mod tests { use super::*; diff --git a/pingora-core/src/connectors/tls/rustls/mod.rs b/pingora-core/src/connectors/tls/rustls/mod.rs index ff3759296..23e3a307e 100644 --- a/pingora-core/src/connectors/tls/rustls/mod.rs +++ b/pingora-core/src/connectors/tls/rustls/mod.rs @@ -61,6 +61,9 @@ impl TlsConnector { where Self: Sized, { + // rustls 0.23+ requires an explicit CryptoProvider. + pingora_rustls::install_default_crypto_provider(); + // NOTE: Rustls only supports TLS 1.2 & 1.3 // TODO: currently using Rustls defaults @@ -135,9 +138,28 @@ where { let config = &tls_ctx.config; - // TODO: setup CA/verify cert store from peer - // peer.get_ca() returns None by default. It must be replaced by the - // implementation of `peer` + // Build per-peer CA store if provided + let peer_ca_store: Option> = if let Some(ca_list) = peer.get_ca() { + if ca_list.is_empty() { + return Error::e_explain(InvalidCert, "per-peer CA list is empty"); + } + let mut ca_store = RootCertStore::empty(); + for ca_cert in &**ca_list { + let cert_der = CertificateDer::from(ca_cert); + ca_store.add(cert_der).or_err( + InvalidCert, + "Failed to add per-peer CA certificate to root store", + )?; + } + + Some(Arc::new(ca_store)) + } else { + None + }; + + // Determine effective CA store for this connection + let effective_ca_store = peer_ca_store.as_ref().unwrap_or(&tls_ctx.ca_certs); + let key_pair = peer.get_client_cert_key(); let mut updated_config_opt: Option = match key_pair { None => None, @@ -164,7 +186,7 @@ where &version::TLS12, &version::TLS13, ]) - .with_root_certificates(Arc::clone(&tls_ctx.ca_certs)); + .with_root_certificates(Arc::clone(effective_ca_store)); debug!("added root ca certificates"); let mut updated_config = builder.with_client_auth_cert(certs, private_key).or_err( @@ -177,6 +199,17 @@ where } }; + // Ensure config is updated if per-peer CA is set but no client cert + if peer_ca_store.is_some() && updated_config_opt.is_none() { + let mut updated_config = + RusTlsClientConfig::builder_with_protocol_versions(&[&version::TLS12, &version::TLS13]) + .with_root_certificates(Arc::clone(effective_ca_store)) + .with_no_client_auth(); + + updated_config.key_log = Arc::clone(&config.key_log); + updated_config_opt = Some(updated_config); + } + if let Some(alpn) = alpn_override.as_ref().or(peer.get_alpn()) { let alpn_protocols = alpn.to_wire_protocols(); if let Some(updated_config) = updated_config_opt.as_mut() { @@ -212,7 +245,7 @@ where // Builds the custom_verifier when verification_mode is set. if let Some(mode) = verification_mode { - let delegate = WebPkiServerVerifier::builder(Arc::clone(&tls_ctx.ca_certs)) + let delegate = WebPkiServerVerifier::builder(Arc::clone(effective_ca_store)) .build() .or_err(InvalidCert, "Failed to build WebPkiServerVerifier")?; diff --git a/pingora-core/src/connectors/tls/s2n/mod.rs b/pingora-core/src/connectors/tls/s2n/mod.rs index fbfdd7e73..5deb44b44 100644 --- a/pingora-core/src/connectors/tls/s2n/mod.rs +++ b/pingora-core/src/connectors/tls/s2n/mod.rs @@ -83,11 +83,11 @@ impl TlsConnector { if self.config_cache.is_some() { let config_hash = config_options.config_hash(); if let Some(config) = self.load_config_from_cache(config_hash) { - return Ok(config); + Ok(config) } else { let config = create_s2n_config(&self.options, config_options)?; self.put_config_in_cache(config_hash, config.clone()); - return Ok(config); + Ok(config) } } else { create_s2n_config(&self.options, config_options) @@ -116,14 +116,14 @@ impl TlsConnector { let mut cache_size = DEFAULT_CONFIG_CACHE_SIZE; if let Some(opts) = options { if let Some(cache_size_config) = opts.s2n_config_cache_size { - if cache_size_config <= 0 { + if cache_size_config == 0 { return None; } else { cache_size = NonZero::new(cache_size_config).unwrap(); } } } - return Some(Arc::new(Mutex::new(LruCache::new(cache_size)))); + Some(Arc::new(Mutex::new(LruCache::new(cache_size)))) } } @@ -145,7 +145,7 @@ where let config = tls_ctx.load_config(config_options)?; let connection_builder = S2NConnectionBuilder { - config: config, + config, psk_config: peer.get_psk().cloned(), security_policy: Some(security_policy.clone()), }; @@ -178,7 +178,7 @@ fn create_s2n_config( if let Some(conf) = connector_options.as_ref() { if let Some(ca_file_path) = conf.ca_file.as_ref() { - let ca_pem = load_pem_file(&ca_file_path)?; + let ca_pem = load_pem_file(ca_file_path)?; builder .trust_pem(&ca_pem) .or_err(InternalError, "failed to load ca cert")?; @@ -210,7 +210,7 @@ fn create_s2n_config( if let Some(client_cert_key) = config_options.client_cert_key { builder - .load_pem(&client_cert_key.raw_pem(), &client_cert_key.key()) + .load_pem(client_cert_key.raw_pem(), client_cert_key.key()) .or_err(InternalError, "invalid peer client cert or key")?; } @@ -243,9 +243,9 @@ fn create_s2n_config( )?; } - Ok(builder + builder .build() - .or_err(InternalError, "failed to build s2n config")?) + .or_err(InternalError, "failed to build s2n config") } #[derive(Clone)] diff --git a/pingora-core/src/listeners/l4.rs b/pingora-core/src/listeners/l4.rs index 1c0052f89..5532b635e 100644 --- a/pingora-core/src/listeners/l4.rs +++ b/pingora-core/src/listeners/l4.rs @@ -44,6 +44,20 @@ use crate::protocols::GetSocketDigest; use crate::protocols::TcpKeepalive; #[cfg(unix)] use crate::server::ListenFds; +#[cfg(unix)] +use std::sync::LazyLock; + +/// Per-address async lock map for serializing the check-bind-insert sequence +/// in [`ListenerEndpointBuilder::listen`]. +/// +/// With `ListenFds` using a synchronous `parking_lot::Mutex`, the lock cannot +/// be held across `bind().await`. This global map ensures that only one task at +/// a time can be in the process of looking up, binding, and inserting a given +/// address — preventing two concurrent callers from both seeing "not found" and +/// racing to bind the same address. +#[cfg(unix)] +static BIND_LOCKS: LazyLock>>> = + LazyLock::new(flurry::HashMap::new); const TCP_LISTENER_MAX_TRY: usize = 30; const TCP_LISTENER_TRY_STEP: Duration = Duration::from_secs(1); @@ -326,19 +340,38 @@ impl ListenerEndpointBuilder { let listener = if let Some(fds_table) = fds { let addr_str = listen_addr.as_ref(); - // consider make this mutex std::sync::Mutex or OnceCell - let mut table = fds_table.lock().await; + // Acquire a per-address async lock so that only one task at a + // time can go through the check-bind-insert sequence for a given + // address. The flurry guard is dropped before the await so its + // !Send pointer does not cross an await point. + let addr_lock = { + let guard = BIND_LOCKS.pin(); + match guard.get(addr_str) { + Some(existing) => existing.clone(), + None => { + let new_lock = Arc::new(tokio::sync::Mutex::new(())); + match guard.try_insert(addr_str.to_string(), new_lock.clone()) { + Ok(inserted) => inserted.clone(), + Err(e) => e.current.clone(), + } + } + } + }; + let _guard = addr_lock.lock().await; + + let existing_fd = fds_table.lock().get(addr_str).copied(); - if let Some(fd) = table.get(addr_str) { - from_raw_fd(&listen_addr, *fd)? + if let Some(fd) = existing_fd { + from_raw_fd(&listen_addr, fd)? } else { - // not found let listener = bind(&listen_addr).await?; - table.add(addr_str.to_string(), listener.as_raw_fd()); + fds_table + .lock() + .add(addr_str.to_string(), listener.as_raw_fd()); listener } } else { - // not found, no fd table + // no fd table bind(&listen_addr).await? }; @@ -386,6 +419,15 @@ impl ListenerEndpoint { self.listen_addr.as_ref() } + /// Return the local address this endpoint is bound to. + /// + /// Useful when the listener was bound to port 0 (OS-assigned) to + /// discover the actual port. + #[cfg(test)] + pub fn local_addr(&self) -> Option { + self.listener.local_addr() + } + fn apply_stream_settings(&self, stream: &mut Stream) -> Result<()> { // settings are applied based on whether the underlying stream supports it stream.set_nodelay()?; @@ -470,11 +512,9 @@ mod test { #[tokio::test] async fn test_listen_tcp() { - let addr = "127.0.0.1:7100"; - let mut builder = ListenerEndpoint::builder(); - builder.listen_addr(ServerAddress::Tcp(addr.into(), None)); + builder.listen_addr(ServerAddress::Tcp("127.0.0.1:0".into(), None)); #[cfg(unix)] let listener = builder.listen(None).await.unwrap(); @@ -482,6 +522,8 @@ mod test { #[cfg(windows)] let listener = builder.listen().await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { // just try to accept once listener.accept().await.unwrap(); @@ -500,7 +542,7 @@ mod test { let mut builder = ListenerEndpoint::builder(); - builder.listen_addr(ServerAddress::Tcp("[::]:7101".into(), sock_opt)); + builder.listen_addr(ServerAddress::Tcp("[::]:0".into(), sock_opt)); #[cfg(unix)] let listener = builder.listen(None).await.unwrap(); @@ -508,15 +550,17 @@ mod test { #[cfg(windows)] let listener = builder.listen().await.unwrap(); + let port = listener.local_addr().unwrap().port(); + tokio::spawn(async move { // just try to accept twice listener.accept().await.unwrap(); listener.accept().await.unwrap(); }); - tokio::net::TcpStream::connect("127.0.0.1:7101") + tokio::net::TcpStream::connect(format!("127.0.0.1:{port}")) .await .expect_err("cannot connect to v4 addr"); - tokio::net::TcpStream::connect("[::1]:7101") + tokio::net::TcpStream::connect(format!("[::1]:{port}")) .await .expect("can connect to v6 addr"); } diff --git a/pingora-core/src/listeners/mod.rs b/pingora-core/src/listeners/mod.rs index e44f17359..db7997655 100644 --- a/pingora-core/src/listeners/mod.rs +++ b/pingora-core/src/listeners/mod.rs @@ -86,6 +86,9 @@ use std::{any::Any, fs::Permissions, sync::Arc}; use l4::{ListenerEndpoint, Stream as L4Stream}; use tls::{Acceptor, TlsSettings}; +pub use crate::protocols::l4::stream::{ + L4BufferSettings, DEFAULT_L4_READ_BUFFER_SIZE, DEFAULT_L4_WRITE_BUFFER_SIZE, +}; pub use crate::protocols::tls::ALPN; use crate::protocols::GetSocketDigest; pub use l4::{ServerAddress, TcpSocketOptions}; @@ -118,11 +121,51 @@ pub trait TlsAccept { pub type TlsAcceptCallbacks = Box; +/// Callback for processing raw bytes before TLS handshake. +/// +/// This trait allows applications to read and process data from the raw TCP stream +/// before the TLS handshake occurs. This is useful for protocols like HAProxy's +/// PROXY protocol, which sends client address information before TLS. +/// +/// # Example +/// +/// ```rust,ignore +/// use pingora_core::listeners::PreTlsProcess; +/// use pingora_core::protocols::l4::stream::Stream as L4Stream; +/// use async_trait::async_trait; +/// +/// struct ProxyProtocolHandler; +/// +/// #[async_trait] +/// impl PreTlsProcess for ProxyProtocolHandler { +/// async fn process(&self, stream: &mut L4Stream) -> pingora_error::Result<()> { +/// // Read PROXY protocol header, update socket digest, etc. +/// Ok(()) +/// } +/// } +/// ``` +#[async_trait] +pub trait PreTlsProcess: Send + Sync { + /// Process the raw stream before TLS handshake. + /// + /// The implementation can read bytes from the stream (e.g., PROXY protocol header) + /// and update the stream's socket digest with parsed information such as the + /// real client address. + /// + /// If this method returns an error, the connection will be dropped. + async fn process(&self, stream: &mut L4Stream) -> Result<()>; +} + +/// Type alias for a boxed pre-TLS processor. +pub type PreTlsCallback = Arc; + struct TransportStackBuilder { l4: ServerAddress, tls: Option, + l4_buffer: L4BufferSettings, #[cfg(feature = "connection_filter")] connection_filter: Option>, + pre_tls_callback: Option, } impl TransportStackBuilder { @@ -148,14 +191,96 @@ impl TransportStackBuilder { Ok(TransportStack { l4, tls: self.tls.take().map(|tls| Arc::new(tls.build())), + l4_buffer: self.l4_buffer, + pre_tls_callback: self.pre_tls_callback.clone(), }) } } +/// Configuration for one listening endpoint. +/// +/// This configures the endpoint address and endpoint-specific transport +/// settings such as [`TcpSocketOptions`], [`TlsSettings`], and L4 +/// [`BufStream`](tokio::io::BufStream) buffer sizes. +pub struct ListenerConfig { + l4: ServerAddress, + tls: Option, + l4_buffer: L4BufferSettings, +} + +impl ListenerConfig { + /// Create a TCP listening endpoint config. + pub fn tcp(addr: impl Into) -> Self { + Self { + l4: ServerAddress::Tcp(addr.into(), None), + tls: None, + l4_buffer: L4BufferSettings::default(), + } + } + + /// Create a Unix domain socket listening endpoint config. + #[cfg(unix)] + pub fn uds(addr: impl Into) -> Self { + Self { + l4: ServerAddress::Uds(addr.into(), None), + tls: None, + l4_buffer: L4BufferSettings::default(), + } + } + + /// Set TCP socket options for this endpoint. + /// + /// # Panics + /// + /// Panics if this endpoint is not TCP. + #[track_caller] + pub fn tcp_socket_options(mut self, options: TcpSocketOptions) -> Self { + match &mut self.l4 { + ServerAddress::Tcp(_, opt) => *opt = Some(options), + #[cfg(unix)] + ServerAddress::Uds(_, _) => { + panic!("TCP socket options can only be set on TCP endpoints") + } + } + self + } + + /// Set Unix domain socket permissions for this endpoint. + /// + /// # Panics + /// + /// Panics if this endpoint is not a Unix domain socket. + #[cfg(unix)] + #[track_caller] + pub fn permissions(mut self, permissions: Permissions) -> Self { + match &mut self.l4 { + ServerAddress::Uds(_, perm) => *perm = Some(permissions), + ServerAddress::Tcp(_, _) => { + panic!("Unix domain socket permissions can only be set on UDS endpoints") + } + } + self + } + + /// Set TLS settings for this endpoint. + pub fn tls(mut self, settings: TlsSettings) -> Self { + self.tls = Some(settings); + self + } + + /// Set L4 `BufStream` buffer sizes for this endpoint. + pub fn l4_buffer(mut self, settings: L4BufferSettings) -> Self { + self.l4_buffer = settings; + self + } +} + #[derive(Clone)] pub(crate) struct TransportStack { l4: ListenerEndpoint, tls: Option>, + l4_buffer: L4BufferSettings, + pre_tls_callback: Option, } impl TransportStack { @@ -168,6 +293,8 @@ impl TransportStack { Ok(UninitializedStream { l4: stream, tls: self.tls.clone(), + l4_buffer: self.l4_buffer, + pre_tls_callback: self.pre_tls_callback.clone(), }) } @@ -179,12 +306,19 @@ impl TransportStack { pub(crate) struct UninitializedStream { l4: L4Stream, tls: Option>, + l4_buffer: L4BufferSettings, + pre_tls_callback: Option, } impl UninitializedStream { pub async fn handshake(mut self) -> Result { - self.l4.set_buffer(); + self.l4.set_buffer(self.l4_buffer); if let Some(tls) = self.tls { + // Process pre-TLS data if a callback is configured (e.g., PROXY protocol) + if let Some(ref callback) = self.pre_tls_callback { + callback.process(&mut self.l4).await?; + } + let tls_stream = tls.tls_handshake(self.l4).await?; Ok(Box::new(tls_stream)) } else { @@ -205,6 +339,7 @@ pub struct Listeners { stacks: Vec, #[cfg(feature = "connection_filter")] connection_filter: Option>, + pre_tls_callback: Option, } impl Listeners { @@ -214,6 +349,7 @@ impl Listeners { stacks: vec![], #[cfg(feature = "connection_filter")] connection_filter: None, + pre_tls_callback: None, } } /// Create a new [`Listeners`] with a TCP server endpoint from the given string. @@ -243,18 +379,22 @@ impl Listeners { /// Add a TCP endpoint to `self`. pub fn add_tcp(&mut self, addr: &str) { - self.add_address(ServerAddress::Tcp(addr.into(), None)); + self.add_listener(ListenerConfig::tcp(addr)); } /// Add a TCP endpoint to `self`, with the given [`TcpSocketOptions`]. pub fn add_tcp_with_settings(&mut self, addr: &str, sock_opt: TcpSocketOptions) { - self.add_address(ServerAddress::Tcp(addr.into(), Some(sock_opt))); + self.add_listener(ListenerConfig::tcp(addr).tcp_socket_options(sock_opt)); } /// Add a Unix domain socket endpoint to `self`. #[cfg(unix)] pub fn add_uds(&mut self, addr: &str, perm: Option) { - self.add_address(ServerAddress::Uds(addr.into(), perm)); + let endpoint = perm.map_or_else( + || ListenerConfig::uds(addr), + |perm| ListenerConfig::uds(addr).permissions(perm), + ); + self.add_listener(endpoint); } /// Add a TLS endpoint to `self` with the [Mozilla Intermediate](https://wiki.mozilla.org/Security/Server_Side_TLS#Intermediate_compatibility_.28recommended.29) @@ -272,7 +412,11 @@ impl Listeners { sock_opt: Option, settings: TlsSettings, ) { - self.add_endpoint(ServerAddress::Tcp(addr.into(), sock_opt), Some(settings)); + let mut endpoint = ListenerConfig::tcp(addr).tls(settings); + if let Some(sock_opt) = sock_opt { + endpoint = endpoint.tcp_socket_options(sock_opt); + } + self.add_listener(endpoint); } /// Add the given [`ServerAddress`] to `self`. @@ -294,13 +438,57 @@ impl Listeners { } } - /// Add the given [`ServerAddress`] to `self` with the given [`TlsSettings`] if provided + /// Add the given listener endpoint to `self`. + pub fn add_listener(&mut self, endpoint: ListenerConfig) { + let ListenerConfig { l4, tls, l4_buffer } = endpoint; + self.stacks.push(TransportStackBuilder { + l4, + tls, + l4_buffer, + #[cfg(feature = "connection_filter")] + connection_filter: self.connection_filter.clone(), + pre_tls_callback: self.pre_tls_callback.clone(), + }); + } + + /// Set a pre-TLS callback for all endpoints in this listener collection. + /// + /// The callback will be invoked after TCP accept but before the TLS handshake, + /// allowing the application to read and process data such as PROXY protocol + /// headers that arrive before TLS. + /// + /// # Example + /// + /// ```rust,ignore + /// use pingora_core::listeners::{Listeners, PreTlsProcess}; + /// use std::sync::Arc; + /// + /// let callback = Arc::new(MyProxyProtocolHandler::new()); + /// let mut listeners = Listeners::new(); + /// listeners.set_pre_tls_callback(callback); + /// listeners.add_tls("0.0.0.0:443", "cert.pem", "key.pem")?; + /// ``` + pub fn set_pre_tls_callback(&mut self, callback: PreTlsCallback) { + log::debug!("Setting pre-TLS callback on Listeners"); + + // Store the callback for future endpoints + self.pre_tls_callback = Some(callback.clone()); + + // Apply to existing stacks + for stack in &mut self.stacks { + stack.pre_tls_callback = Some(callback.clone()); + } + } + + /// Add the given [`ServerAddress`] to `self` with the given [`TlsSettings`] if provided. pub fn add_endpoint(&mut self, l4: ServerAddress, tls: Option) { self.stacks.push(TransportStackBuilder { l4, tls, + l4_buffer: L4BufferSettings::default(), #[cfg(feature = "connection_filter")] connection_filter: self.connection_filter.clone(), + pre_tls_callback: self.pre_tls_callback.clone(), }) } @@ -337,14 +525,11 @@ mod test { #[cfg(feature = "any_tls")] use tokio::io::AsyncWriteExt; use tokio::net::TcpStream; - use tokio::time::{sleep, Duration}; #[tokio::test] async fn test_listen_tcp() { - let addr1 = "127.0.0.1:7101"; - let addr2 = "127.0.0.1:7102"; - let mut listeners = Listeners::tcp(addr1); - listeners.add_tcp(addr2); + let mut listeners = Listeners::tcp("127.0.0.1:0"); + listeners.add_tcp("127.0.0.1:0"); let listeners = listeners .build( @@ -355,6 +540,10 @@ mod test { .unwrap(); assert_eq!(listeners.len(), 2); + let addrs: Vec<_> = listeners + .iter() + .map(|s| s.l4.local_addr().unwrap()) + .collect(); for listener in listeners { tokio::spawn(async move { // just try to accept once @@ -363,11 +552,77 @@ mod test { }); } - // make sure the above starts before the lines below - sleep(Duration::from_millis(10)).await; + // The listeners are already bound (port resolved during build()), + // so the kernel accepts connections into the backlog immediately. + // No readiness wait needed — connect will succeed as soon as the + // OS has completed the TCP handshake. + TcpStream::connect(addrs[0]).await.unwrap(); + TcpStream::connect(addrs[1]).await.unwrap(); + } + + #[test] + fn test_add_listener_config_tcp_l4_buffer() { + let mut listeners = Listeners::new(); + let tcp_options = TcpSocketOptions { + dscp: Some(10), + ..Default::default() + }; + let l4_buffer = L4BufferSettings { + read: Some(0), + write: None, + }; + + listeners.add_listener( + ListenerConfig::tcp("127.0.0.1:7107") + .tcp_socket_options(tcp_options) + .l4_buffer(l4_buffer), + ); + + assert_eq!(listeners.stacks.len(), 1); + assert_eq!(listeners.stacks[0].l4_buffer, l4_buffer); + assert_eq!(listeners.stacks[0].l4_buffer.read_capacity(), 0); + assert_eq!( + listeners.stacks[0].l4_buffer.write_capacity(), + DEFAULT_L4_WRITE_BUFFER_SIZE + ); + + match &listeners.stacks[0].l4 { + ServerAddress::Tcp(addr, Some(options)) => { + assert_eq!(addr, "127.0.0.1:7107"); + assert_eq!(options.dscp, Some(10)); + } + other => panic!("unexpected listener address: {other:?}"), + } + } + + #[cfg(unix)] + #[test] + fn test_add_listener_config_uds_l4_buffer() { + let mut listeners = Listeners::new(); + let l4_buffer = L4BufferSettings::unbuffered(); + + listeners.add_listener(ListenerConfig::uds("/tmp/test_builder_uds").l4_buffer(l4_buffer)); - TcpStream::connect(addr1).await.unwrap(); - TcpStream::connect(addr2).await.unwrap(); + assert_eq!(listeners.stacks.len(), 1); + assert_eq!(listeners.stacks[0].l4_buffer, l4_buffer); + assert_eq!(listeners.stacks[0].l4_buffer.read_capacity(), 0); + assert_eq!(listeners.stacks[0].l4_buffer.write_capacity(), 0); + + match &listeners.stacks[0].l4 { + ServerAddress::Uds(addr, None) => assert_eq!(addr, "/tmp/test_builder_uds"), + other => panic!("unexpected listener address: {other:?}"), + } + } + + #[test] + fn test_l4_buffer_settings_defaults_per_direction() { + let l4_buffer = L4BufferSettings { + read: None, + write: Some(0), + }; + + assert_eq!(l4_buffer.read_capacity(), DEFAULT_L4_READ_BUFFER_SIZE); + assert_eq!(l4_buffer.write_capacity(), 0); } #[tokio::test] @@ -400,9 +655,8 @@ mod test { .await .unwrap(); }); - // make sure the above starts before the lines below - sleep(Duration::from_millis(10)).await; - + // The listener is already bound, so the kernel accepts connections + // into the backlog immediately. No readiness wait needed. let client = reqwest::Client::builder() .danger_accept_invalid_certs(true) .build() diff --git a/pingora-core/src/listeners/tls/rustls/mod.rs b/pingora-core/src/listeners/tls/rustls/mod.rs index 0ca94d514..e7376fc01 100644 --- a/pingora-core/src/listeners/tls/rustls/mod.rs +++ b/pingora-core/src/listeners/tls/rustls/mod.rs @@ -48,6 +48,9 @@ impl TlsSettings { /// /// Todo: Return a result instead of panicking XD pub fn build(self) -> Acceptor { + // rustls 0.23+ requires an explicit CryptoProvider. + pingora_rustls::install_default_crypto_provider(); + let Ok(Some((certs, key))) = load_certs_and_key_files(&self.cert_path, &self.key_path) else { panic!( diff --git a/pingora-core/src/listeners/tls/s2n/mod.rs b/pingora-core/src/listeners/tls/s2n/mod.rs index ed689445f..af547bbe9 100644 --- a/pingora-core/src/listeners/tls/s2n/mod.rs +++ b/pingora-core/src/listeners/tls/s2n/mod.rs @@ -89,7 +89,7 @@ impl TlsSettings { let config = builder.build().unwrap(); let connection_builder = S2NConnectionBuilder { - config: config, + config, psk_config: self.psk_config.clone(), security_policy: Some(policy.clone()), }; @@ -105,7 +105,7 @@ impl TlsSettings { self.set_alpn(ALPN::H2H1); } - fn set_alpn(&mut self, alpn: ALPN) { + pub fn set_alpn(&mut self, alpn: ALPN) { self.alpn = Some(alpn); } diff --git a/pingora-core/src/protocols/http/compression/mod.rs b/pingora-core/src/protocols/http/compression/mod.rs index 9e8e96fba..301de0fbe 100644 --- a/pingora-core/src/protocols/http/compression/mod.rs +++ b/pingora-core/src/protocols/http/compression/mod.rs @@ -303,9 +303,18 @@ impl ResponseCompressionCtx { Action::Compress(algorithm) => { let idx = algorithm.index(); let compressor = match algorithm { - Algorithm::Dcz => dictionary.as_ref().and_then(|d| { - algorithm.maybe_compressor_with_dictionary(levels[idx], d) - }), + Algorithm::Dcz => { + // RFC 9842: dictionary-compressed responses vary on + // Available-Dictionary so caches don't serve this variant + // to clients with a different or missing dictionary. + let enc = dictionary.as_ref().and_then(|d| { + algorithm.maybe_compressor_with_dictionary(levels[idx], d) + }); + if enc.is_some() { + add_vary_header(resp, &AVAILABLE_DICTIONARY); + } + enc + } _ => algorithm.compressor(levels[idx]), }; (compressor, preserve_etag[idx]) @@ -780,6 +789,13 @@ fn compressible(resp: &ResponseHeader) -> bool { } } +/// Header name for the Available-Dictionary request header ([RFC 9842]). +/// TODO: Replace with http::header when available. +/// +/// [RFC 9842]: https://datatracker.ietf.org/doc/html/rfc9842 +static AVAILABLE_DICTIONARY: http::HeaderName = + http::HeaderName::from_static("available-dictionary"); + // add Vary header with the specified value or extend an existing Vary header value fn add_vary_header(resp: &mut ResponseHeader, value: &http::header::HeaderName) { use http::header::{HeaderValue, VARY}; @@ -1055,6 +1071,11 @@ mod tests_dictionary_compression { resp.headers.get("content-encoding").unwrap().as_bytes(), b"dcz" ); + // RFC 9842: DCZ responses must vary on Available-Dictionary. + assert!(resp.headers.get_all("vary").iter().any(|v| v + .as_bytes() + .split(|b| *b == b',') + .any(|t| t.trim_ascii().eq_ignore_ascii_case(b"available-dictionary")))); let input = Bytes::from_static(b"The quick brown fox jumps over the lazy dog again."); let compressed = ctx.response_body_filter(Some(&input), true).unwrap(); @@ -1080,6 +1101,11 @@ mod tests_dictionary_compression { // no dictionary set, no compression applied assert!(resp.headers.get("content-encoding").is_none()); + // No compression → no Vary: available-dictionary. + assert!(!resp.headers.get_all("vary").iter().any(|v| v + .as_bytes() + .split(|b| *b == b',') + .any(|t| t.trim_ascii().eq_ignore_ascii_case(b"available-dictionary")))); } #[test] @@ -1099,6 +1125,11 @@ mod tests_dictionary_compression { // dcz first but no dictionary, no automatic fallback assert!(resp.headers.get("content-encoding").is_none()); + // No compression → no Vary: available-dictionary. + assert!(!resp.headers.get_all("vary").iter().any(|v| v + .as_bytes() + .split(|b| *b == b',') + .any(|t| t.trim_ascii().eq_ignore_ascii_case(b"available-dictionary")))); } #[test] @@ -1152,6 +1183,11 @@ mod tests_dictionary_compression { resp.headers.get("transfer-encoding").unwrap().as_bytes(), b"chunked" ); + // RFC 9842: DCZ responses must vary on Available-Dictionary. + assert!(resp.headers.get_all("vary").iter().any(|v| v + .as_bytes() + .split(|b| *b == b',') + .any(|t| t.trim_ascii().eq_ignore_ascii_case(b"available-dictionary")))); let chunk1 = Bytes::from_static(b"First chunk. "); let output1 = ctx.response_body_filter(Some(&chunk1), false); @@ -1166,4 +1202,33 @@ mod tests_dictionary_compression { assert_eq!(total_in, chunk1.len() + chunk2.len()); assert!(total_out > 0); } + + #[test] + fn regular_compression_no_available_dictionary_vary() { + // Gzip compression should produce Vary: Accept-Encoding but NOT + // Vary: available-dictionary. + let mut ctx = ResponseCompressionCtx::new(3, false, false); + + let mut req = RequestHeader::build("GET", b"/page.html", None).unwrap(); + req.insert_header("accept-encoding", "gzip").unwrap(); + ctx.request_filter(&req); + + let mut resp = ResponseHeader::build(200, None).unwrap(); + resp.insert_header("content-type", "text/html").unwrap(); + resp.insert_header("content-length", "1000").unwrap(); + ctx.response_header_filter(&mut resp, false); + + assert_eq!( + resp.headers.get("content-encoding").unwrap().as_bytes(), + b"gzip" + ); + assert!(resp.headers.get_all("vary").iter().any(|v| v + .as_bytes() + .split(|b| *b == b',') + .any(|t| t.trim_ascii().eq_ignore_ascii_case(b"accept-encoding")))); + assert!(!resp.headers.get_all("vary").iter().any(|v| v + .as_bytes() + .split(|b| *b == b',') + .any(|t| t.trim_ascii().eq_ignore_ascii_case(b"available-dictionary")))); + } } diff --git a/pingora-core/src/protocols/http/mod.rs b/pingora-core/src/protocols/http/mod.rs index f5bc729d3..42d55da7c 100644 --- a/pingora-core/src/protocols/http/mod.rs +++ b/pingora-core/src/protocols/http/mod.rs @@ -27,7 +27,7 @@ pub mod subrequest; pub mod v1; pub mod v2; -pub use server::Session as ServerSession; +pub use server::{ReusableHttpStream, Session as ServerSession}; /// The Pingora server name string pub const SERVER_NAME: &[u8; 7] = b"Pingora"; diff --git a/pingora-core/src/protocols/http/server.rs b/pingora-core/src/protocols/http/server.rs index 035a65cc3..a4519494c 100644 --- a/pingora-core/src/protocols/http/server.rs +++ b/pingora-core/src/protocols/http/server.rs @@ -18,17 +18,40 @@ use super::custom::server::Session as SessionCustom; use super::error_resp; use super::subrequest::server::HttpSession as SessionSubrequest; use super::v1::server::HttpSession as SessionV1; -use super::v2::server::HttpSession as SessionV2; +use super::v2::server::{HttpSession as SessionV2, Idle}; use super::HttpTask; use crate::custom_session; use crate::protocols::{Digest, SocketAddr, Stream}; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; use http::HeaderValue; use http::{header::AsHeaderName, HeaderMap}; use pingora_error::{Error, Result}; use pingora_http::{RequestHeader, ResponseHeader}; +use std::any::Any; use std::time::Duration; +/// A reusable HTTP/1.x stream and bytes already read for the next request. +#[derive(Debug)] +pub struct ReusableHttpStream { + stream: Stream, + pipelined_prefix: Option, +} + +impl ReusableHttpStream { + pub(crate) fn new(stream: Stream, pipelined_prefix: Option) -> Self { + Self { + stream, + pipelined_prefix, + } + } + + /// Split the reusable connection into its underlying stream and optional + /// bytes already read for the next pipelined request. + pub fn into_parts(self) -> (Stream, Option) { + (self.stream, self.pipelined_prefix) + } +} + /// HTTP server session object for both HTTP/1.x and HTTP/2 pub enum Session { H1(SessionV1), @@ -226,11 +249,13 @@ impl Session { } } - /// Finish the life of this request. - /// For H1, if connection reuse is supported, a Some(Stream) will be returned, otherwise None. + /// Finish the life of this request and return a reusable stream, if any. + /// + /// For H1, if connection reuse is supported, a reusable stream will be returned, + /// otherwise None. /// For H2, always return None because H2 stream is not reusable. /// For subrequests, there is no true underlying stream to return. - pub async fn finish(self) -> Result> { + pub async fn finish(self) -> Result> { match self { Self::H1(mut s) => { // need to flush body due to buffering @@ -317,6 +342,26 @@ impl Session { } } + /// Set user-defined context to carry across requests on the same keepalive connection. + /// + /// Only applicable for HTTP/1.x connections; noop for h2, subrequest, and custom sessions. + pub fn set_connection_user_context(&mut self, ctx: Option>) { + if let Self::H1(s) = self { + s.set_connection_user_context(ctx); + } + } + + /// Take the user-defined context from the previous request on this keepalive connection. + /// + /// Returns `None` for h2, subrequest, and custom sessions, or if no context was persisted. + pub fn take_connection_user_context(&mut self) -> Option> { + if let Self::H1(s) = self { + s.take_connection_user_context() + } else { + None + } + } + /// Sets the downstream read timeout. This will trigger if we're unable /// to read from the stream after `timeout`. /// @@ -436,6 +481,22 @@ impl Session { } } + /// Controls behaviour when the client closes the connection after the request body. + /// + /// When **enabled** (default), a client close is returned as a `ConnectionClosed` + /// error so the proxy aborts immediately. When **disabled**, `read_body_or_idle` + /// stays pending so the proxy can finish delivering the upstream response. + /// + /// Only meaningful for H1 (TCP). Noop for H2/subrequest/custom. + pub fn set_abort_on_close(&mut self, abort: bool) { + match self { + Self::H1(s) => s.set_abort_on_close(abort), + Self::H2(_) => {} + Self::Subrequest(_) => {} + Self::Custom(_) => {} + } + } + /// Return a digest of the request including the method, path and Host header // TODO: make this use a `Formatter` pub fn request_summary(&self) -> String { @@ -471,6 +532,19 @@ impl Session { } } + /// Give up the H2 stream with a custom reason. + /// + /// For H2, this sends a `RST_STREAM` frame with the specified reason. + /// For H1, subrequests, and custom sessions, this is a no-op since they don't support + /// stream reset reasons. + /// + /// See [`super::v2::server::HttpSession::shutdown_with_reason`] for available reasons. + pub fn shutdown_with_reason(&mut self, reason: h2::Reason) { + if let Self::H2(s) = self { + s.shutdown_with_reason(reason); + } + } + pub fn to_h1_raw(&self) -> Bytes { match self { Self::H1(s) => s.get_headers_raw_bytes(), @@ -608,6 +682,19 @@ impl Session { } } + /// Return a future that waits for the client to abort this H2 stream without + /// reading any body data. + /// + /// For HTTP/2 this resolves when the client resets the stream (RST_STREAM) or the + /// stream errors. Other protocols have no out-of-band abort signal, so this + /// returns `None` for them. + pub fn watch_h2_stream_reset(&mut self) -> Option> { + match self { + Self::H2(s) => Some(s.idle()), + _ => None, + } + } + pub fn as_http1(&self) -> Option<&SessionV1> { match self { Self::H1(s) => Some(s), @@ -774,4 +861,106 @@ impl Session { Self::Custom(_) => None, } } + + /// Check if this session supports the cancel-safe proxy task API. + /// + /// Currently supported by HTTP/1.x and Subrequest server sessions; + /// toggled per-session via [`set_proxy_tasks_enabled`](Self::set_proxy_tasks_enabled). + pub fn supports_proxy_task_api(&self) -> bool { + match self { + Self::H1(s) => s.proxy_tasks_enabled(), + Self::Subrequest(s) => s.proxy_tasks_enabled(), + Self::H2(_) => false, + Self::Custom(_) => false, + } + } + + /// Enable or disable the cancel-safe proxy task API for this session. + pub fn set_proxy_tasks_enabled(&mut self, enabled: bool) { + match self { + Self::H1(s) => s.set_proxy_tasks_enabled(enabled), + Self::Subrequest(s) => s.set_proxy_tasks_enabled(enabled), + Self::H2(_) => {} + Self::Custom(_) => {} + } + } + + /// Whether HTTP/1.1 request pipelining is enabled for this session. + /// + /// Always false for H2 / Subrequest / Custom (pipelining is an H/1.1-only + /// concept). For H1, see + /// [`HttpSession::set_pipelining_enabled`](crate::protocols::http::v1::server::HttpSession::set_pipelining_enabled). + pub fn pipelining_enabled(&self) -> bool { + match self { + Self::H1(s) => s.pipelining_enabled(), + _ => false, + } + } + + /// Enable or disable HTTP/1.1 request pipelining on this session. + /// + /// No-op for H2 / Subrequest / Custom. See + /// [`HttpSession::set_pipelining_enabled`](crate::protocols::http::v1::server::HttpSession::set_pipelining_enabled) + /// for semantics. + pub fn set_pipelining_enabled(&mut self, enabled: bool) { + if let Self::H1(s) = self { + s.set_pipelining_enabled(enabled); + } + } + + /// Set pipelined bytes to be parsed as the start of this session's request. + /// + /// No-op for non-H1 sessions. See + /// [`HttpSession::set_pipelined_prefix`](crate::protocols::http::v1::server::HttpSession::set_pipelined_prefix) + /// for the lifecycle. + pub fn set_pipelined_prefix(&mut self, prefix: BytesMut) { + if let Self::H1(s) = self { + s.set_pipelined_prefix(prefix); + } + } + + /// Queue a downstream proxy task for cancel-safe writing. + /// + /// # Panics + /// Panics if called on a session that doesn't support the proxy task API. + /// Check [`supports_proxy_task_api`](Self::supports_proxy_task_api) first, + /// or use `write_response_header()` / `write_response_body()` for other + /// session types. + pub fn send_downstream_proxy_task(&mut self, task: HttpTask) { + match self { + Self::H1(s) => s.send_proxy_task(task), + Self::H2(_) => panic!("H2 proxy task API not yet implemented"), + Self::Subrequest(s) => s.send_proxy_task(task), + Self::Custom(_) => panic!("Custom proxy task API not yet implemented"), + } + } + + /// Check if there are pending downstream proxy tasks queued for writing. + /// + /// Returns false for sessions that don't support the proxy task API. + pub fn has_pending_downstream_proxy_tasks(&self) -> bool { + match self { + Self::H1(s) => s.has_pending_proxy_tasks(), + Self::H2(_) => false, // TODO: implement for H2 + Self::Subrequest(s) => s.has_pending_proxy_tasks(), + Self::Custom(_) => false, // TODO: implement for custom + } + } + + /// Write all queued downstream proxy tasks in a cancel-safe manner. + /// Returns `Ok(true)` if this was the end of the response stream. + /// + /// # Panics + /// Panics if called on a session that doesn't support the proxy task API. + /// Check [`supports_proxy_task_api`](Self::supports_proxy_task_api) first, + /// or use `write_response_header()` / `write_response_body()` for other + /// session types. + pub async fn write_downstream_proxy_tasks(&mut self) -> Result { + match self { + Self::H1(s) => s.write_proxy_tasks().await, + Self::H2(_) => panic!("H2 proxy task API not yet implemented"), + Self::Subrequest(s) => s.write_proxy_tasks().await, + Self::Custom(_) => panic!("Custom proxy task API not yet implemented"), + } + } } diff --git a/pingora-core/src/protocols/http/subrequest/server.rs b/pingora-core/src/protocols/http/subrequest/server.rs index c91dbf916..938c8c97a 100644 --- a/pingora-core/src/protocols/http/subrequest/server.rs +++ b/pingora-core/src/protocols/http/subrequest/server.rs @@ -36,13 +36,14 @@ use bytes::Bytes; use http::HeaderValue; use http::{header, header::AsHeaderName, HeaderMap, Method}; use log::{debug, trace, warn}; -use pingora_error::{Error, ErrorType::*, OkOrErr, Result}; +use pingora_error::{Error, ErrorType::*, OkOrErr, OrErr, Result}; use pingora_http::{RequestHeader, ResponseHeader}; use pingora_timeout::timeout; +use std::collections::VecDeque; use std::time::Duration; use tokio::sync::{mpsc, oneshot}; -use super::body::{BodyReader, BodyWriter}; +use super::body::{BodyMode, BodyReader, BodyWriter, PREMATURE_BODY_END}; use crate::protocols::http::{ body_buffer::FixedBuffer, server::Session as GenericHttpSession, @@ -53,6 +54,47 @@ use crate::protocols::http::{ }; use crate::protocols::{Digest, SocketAddr}; +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)] +enum StreamEndState { + #[default] + Open, + FinishRequired, + Finished, +} + +/// State for the cancel-safe proxy task write API. +#[derive(Default)] +struct ProxyTaskState { + /// Tasks remain queued until `Permit::send` commits them to the channel. + tasks: VecDeque, + /// Final `Done` is waiting for channel capacity. + finish_in_progress: bool, + /// End-of-stream observed from already-consumed tasks. + stream_end: StreamEndState, +} + +impl ProxyTaskState { + fn require_finish(&mut self) { + self.stream_end = StreamEndState::FinishRequired; + } + + fn mark_finished(&mut self) { + self.stream_end = StreamEndState::Finished; + } + + fn clear_stream_end(&mut self) { + self.stream_end = StreamEndState::Open; + } + + fn finish_required(&self) -> bool { + self.stream_end == StreamEndState::FinishRequired + } + + fn finished(&self) -> bool { + self.stream_end == StreamEndState::Finished + } +} + /// The HTTP server session pub struct HttpSession { // these are only options because we allow dropping them separately on shutdown @@ -76,6 +118,11 @@ pub struct HttpSession { // TODO: likely doesn't need to be a separate bool when/if moving away from dummy SessionV1 clear_request_body_headers: bool, digest: Option>, + /// Whether the cancel-safe proxy task API is enabled for this session. + /// Defaults to `false`. Toggle via [`set_proxy_tasks_enabled`](Self::set_proxy_tasks_enabled). + proxy_tasks_enabled: bool, + /// Cancel-safe proxy task state. + proxy_task_state: ProxyTaskState, } /// A handle to the subrequest session itself to interact or read from it. @@ -133,6 +180,8 @@ impl HttpSession { upgraded: false, clear_request_body_headers: false, digest: digest.map(Box::new), + proxy_tasks_enabled: false, + proxy_task_state: ProxyTaskState::default(), }, SubrequestHandle { tx: downstream_tx, @@ -310,12 +359,24 @@ impl HttpSession { // XXX: don't add additional downstream headers, unlike h1, subreq is mostly treated as a pipe - // Allow informational header (excluding 101) to pass through without affecting the state - // of the request + let mut upgrade_ok: Option = None; if header.status == 101 || !header.status.is_informational() { - // reset request body to done for incomplete upgrade handshakes - if let Some(upgrade_ok) = self.is_upgrade(&header) { - if upgrade_ok { + upgrade_ok = self.is_upgrade(&header); + } + + // TODO propagate h2 end + debug!("send response header (subrequest)"); + match self + .tx + .as_mut() + .expect("tx valid before shutdown") + .send(HttpTask::Header(header.clone(), false)) + .await + { + Ok(()) => { + self.init_body_writer(&header); + self.response_written = Some(*header); + if let Some(true) = upgrade_ok { debug!("ok upgrade handshake"); // For ws we use HTTP1_0 do_read_body_until_closed // @@ -342,25 +403,10 @@ impl HttpSession { // TODO: this has no effect resetting the body counter of TE chunked self.body_reader.convert_to_close_delimited(); } - } else { + } else if upgrade_ok == Some(false) { debug!("bad upgrade handshake!"); // continue to read body as-is, this is now just a regular request } - } - self.init_body_writer(&header); - } - - // TODO propagate h2 end - debug!("send response header (subrequest)"); - match self - .tx - .as_mut() - .expect("tx valid before shutdown") - .send(HttpTask::Header(header.clone(), false)) - .await - { - Ok(()) => { - self.response_written = Some(*header); Ok(()) } Err(e) => Error::e_because(WriteError, "writing response header", e), @@ -390,40 +436,12 @@ impl HttpSession { } fn init_body_writer(&mut self, header: &ResponseHeader) { - use http::StatusCode; - /* the following responses don't have body 204, 304, and HEAD */ - if matches!( - header.status, - StatusCode::NO_CONTENT | StatusCode::NOT_MODIFIED - ) || self.get_method() == Some(&Method::HEAD) - { - self.body_writer.init_content_length(0); - return; - } - - if header.status.is_informational() && header.status != StatusCode::SWITCHING_PROTOCOLS { - // 1xx response, not enough to init body - return; - } - - if self.is_upgrade(header) == Some(true) { - self.body_writer.init_close_delimited(); - } else if is_chunked_encoding_from_headers(&header.headers) { - // transfer-encoding takes priority over content-length - self.body_writer.init_close_delimited(); - } else { - let content_length = - header_value_content_length(header.headers.get(http::header::CONTENT_LENGTH)); - match content_length { - Some(length) => { - self.body_writer.init_content_length(length); - } - None => { - /* TODO: 1. connection: keepalive cannot be used, - 2. mark connection must be closed */ - self.body_writer.init_close_delimited(); - } - } + if let Some(mode) = body_mode_for_header( + header, + self.get_method(), + self.is_upgrade(header) == Some(true), + ) { + apply_body_mode(&mut self.body_writer, mode); } } @@ -800,6 +818,428 @@ impl HttpSession { ) .await } + + // Cancel-safe proxy task API. Unlike v1's partial `AsyncWrite` state + // machines, subrequest uses mpsc `reserve()` + synchronous `Permit::send`. + + /// Whether the cancel-safe proxy task API is enabled for this session. + pub fn proxy_tasks_enabled(&self) -> bool { + self.proxy_tasks_enabled + } + + /// Enable or disable the cancel-safe proxy task API for this session. + pub fn set_proxy_tasks_enabled(&mut self, enabled: bool) { + self.proxy_tasks_enabled = enabled; + } + + /// Queue a proxy task for cancel-safe writing. + pub fn send_proxy_task(&mut self, task: HttpTask) { + self.proxy_task_state.tasks.push_back(task); + } + + /// Whether there are pending proxy tasks queued for writing. + pub fn has_pending_proxy_tasks(&self) -> bool { + self.proxy_task_state.finish_in_progress + || self.proxy_task_state.finish_required() + || !self.proxy_task_state.tasks.is_empty() + } + + /// Write queued proxy tasks. Cancelling while waiting for channel capacity + /// leaves the current task in `proxy_task_state.tasks`. + pub async fn write_proxy_tasks(&mut self) -> Result { + loop { + if self.proxy_task_state.finished() { + self.proxy_task_state.tasks.clear(); + return Ok(true); + } + + if self.proxy_task_state.finish_in_progress { + self.finish_proxy_task().await?; + self.proxy_task_state.mark_finished(); + return Ok(true); + } + + let Some(front) = self.proxy_task_state.tasks.front() else { + break; + }; + + // Tasks with no underlying channel send: handle synchronously + // without reserving a permit. + match front { + HttpTask::Done => { + self.proxy_task_state.tasks.pop_front(); + self.proxy_task_state.require_finish(); + continue; + } + HttpTask::Failed(_) => { + let HttpTask::Failed(e) = self + .proxy_task_state + .tasks + .pop_front() + .expect("queue had a Failed task at the front") + else { + unreachable!() + }; + self.proxy_task_state.clear_stream_end(); + return Err(e); + } + _ => {} + } + + if let HttpTask::Header(_, header_end) = front { + let already_sent = match self.response_written.as_ref() { + Some(resp) => !resp.status.is_informational() || self.upgraded, + None => false, + }; + if already_sent { + warn!("Respond header is already sent, cannot send again (subrequest, proxy task)"); + if *header_end { + self.proxy_task_state.require_finish(); + } + self.proxy_task_state.tasks.pop_front(); + continue; + } + } + + if let Some((upgraded_task, body_end, no_data)) = match front { + HttpTask::Body(data, end) => { + Some((false, *end, data.as_ref().is_none_or(|d| d.is_empty()))) + } + HttpTask::UpgradedBody(data, end) => { + Some((true, *end, data.as_ref().is_none_or(|d| d.is_empty()))) + } + _ => None, + } { + if upgraded_task != self.upgraded { + if upgraded_task { + panic!("Unexpected UpgradedBody task received on un-upgraded downstream session (subrequest, proxy task)"); + } else { + panic!("Unexpected Body task received on upgraded downstream session (subrequest, proxy task)"); + } + } + + if body_end { + self.proxy_task_state.require_finish(); + } + + if no_data { + self.proxy_task_state.tasks.pop_front(); + continue; + } + + match self.body_writer.body_mode { + BodyMode::Complete(_) => { + self.proxy_task_state.tasks.pop_front(); + continue; + } + BodyMode::ContentLength(total, written) if written >= total => { + self.proxy_task_state.tasks.pop_front(); + continue; + } + BodyMode::ToSelect => { + self.proxy_task_state.tasks.pop_front(); + self.proxy_task_state.clear_stream_end(); + return Error::e_explain( + InternalError, + "subrequest body proxy task before header is sent", + ); + } + _ => {} + } + } + + // `reserve()` is the only cancellation point; the queued task is + // popped only after a permit is acquired. + let tx_ref = self + .tx + .as_ref() + .ok_or_else(|| Error::explain(InternalError, "subrequest tx already shut down"))?; + let permit = match self.write_timeout { + Some(t) => match timeout(t, tx_ref.reserve()).await { + Ok(res) => res.or_err(WriteError, "subrequest channel closed")?, + Err(_) => { + return Error::e_explain( + WriteTimedout, + format!("reserving subrequest channel slot, timeout: {t:?}"), + ); + } + }, + None => tx_ref + .reserve() + .await + .or_err(WriteError, "subrequest channel closed")?, + }; + + // From here until `permit.send`, no `.await`; dispatch is atomic. + let task = self + .proxy_task_state + .tasks + .pop_front() + .expect("queue non-empty"); + + match task { + HttpTask::Header(header, hdr_end) => { + if hdr_end { + self.proxy_task_state.require_finish(); + } + + let upgrade_ok = if header.status == 101 || !header.status.is_informational() { + let outcome = self.v1_inner.is_upgrade(&header); + let mode = body_mode_for_header( + &header, + self.v1_inner.get_method(), + outcome == Some(true), + ); + (outcome, mode) + } else { + (None, None) + }; + + permit.send(HttpTask::Header(header.clone(), false)); + + if let Some(mode) = upgrade_ok.1 { + apply_body_mode(&mut self.body_writer, mode); + } + self.response_written = Some(*header); + if let Some(true) = upgrade_ok.0 { + debug!("ok upgrade handshake (subrequest, proxy task)"); + self.upgraded = true; + if self.body_reader.need_init() { + self.init_body_reader(); + } else { + self.body_reader.convert_to_close_delimited(); + } + } else if upgrade_ok.0 == Some(false) { + debug!("bad upgrade handshake! (subrequest, proxy task)"); + } + } + + HttpTask::Body(data, end) => { + if end { + self.proxy_task_state.require_finish(); + } + dispatch_body_inline( + &mut self.body_writer, + &mut self.body_bytes_sent, + self.upgraded, + data, + /* upgraded_task = */ false, + permit, + )?; + } + HttpTask::UpgradedBody(data, end) => { + if end { + self.proxy_task_state.require_finish(); + } + dispatch_body_inline( + &mut self.body_writer, + &mut self.body_bytes_sent, + self.upgraded, + data, + /* upgraded_task = */ true, + permit, + )?; + } + + HttpTask::Trailer(trailers) => { + permit.send(HttpTask::Trailer(trailers)); + self.proxy_task_state.require_finish(); + } + + HttpTask::Done | HttpTask::Failed(_) => { + unreachable!("Done/Failed are handled above without reserving a permit") + } + } + } + + // Match `response_duplex_vec`: finish whenever any task signalled EOS. + if self.proxy_task_state.finish_required() || self.body_writer.finished() { + self.finish_proxy_task().await?; + self.proxy_task_state.mark_finished(); + return Ok(true); + } + + Ok(self.body_writer.finished()) + } + + async fn finish_proxy_task(&mut self) -> Result<()> { + if matches!( + &self.body_writer.body_mode, + BodyMode::Complete(_) | BodyMode::ToSelect + ) { + self.maybe_force_close_body_reader(); + self.proxy_task_state.finish_in_progress = false; + return Ok(()); + } + + if let BodyMode::ContentLength(total, written) = self.body_writer.body_mode { + if written < total { + self.body_writer.body_mode = BodyMode::Complete(written); + self.proxy_task_state.finish_in_progress = false; + self.proxy_task_state.clear_stream_end(); + return Error::e_explain( + PREMATURE_BODY_END, + format!( + "Content-length: {total} bytes written: {written} (subrequest, proxy task)" + ), + ); + } + } + + self.proxy_task_state.finish_in_progress = true; + self.dispatch_finish().await?; + self.proxy_task_state.finish_in_progress = false; + Ok(()) + } + + /// Dispatch the final `HttpTask::Done`, mirroring `body_writer::finish`. + async fn dispatch_finish(&mut self) -> Result<()> { + // Reserve cancel-safely, then synchronously update body_mode and send. + let tx_ref = self + .tx + .as_ref() + .ok_or_else(|| Error::explain(InternalError, "subrequest tx already shut down"))?; + let permit = match self.write_timeout { + Some(t) => match timeout(t, tx_ref.reserve()).await { + Ok(res) => res.or_err(WriteError, "subrequest channel closed")?, + Err(_) => { + return Error::e_explain( + WriteTimedout, + format!("reserving subrequest channel slot for finish, timeout: {t:?}"), + ); + } + }, + None => tx_ref + .reserve() + .await + .or_err(WriteError, "subrequest channel closed")?, + }; + + match self.body_writer.body_mode { + BodyMode::ContentLength(_total, written) => { + self.body_writer.body_mode = BodyMode::Complete(written); + permit.send(HttpTask::Done); + } + BodyMode::UntilClose(written) => { + self.body_writer.body_mode = BodyMode::Complete(written); + permit.send(HttpTask::Done); + } + BodyMode::Complete(_) => { + unreachable!("no-op body modes are handled before reserve") + } + BodyMode::ToSelect => { + unreachable!("no-op body modes are handled before reserve") + } + } + self.maybe_force_close_body_reader(); + Ok(()) + } +} + +fn body_mode_for_header( + header: &ResponseHeader, + method: Option<&Method>, + is_upgrade_ok: bool, +) -> Option { + use http::StatusCode; + if header.status.is_informational() && header.status != StatusCode::SWITCHING_PROTOCOLS { + return None; + } + if matches!( + header.status, + StatusCode::NO_CONTENT | StatusCode::NOT_MODIFIED + ) || method == Some(&Method::HEAD) + { + return Some(BodyMode::ContentLength(0, 0)); + } + if is_upgrade_ok || is_chunked_encoding_from_headers(&header.headers) { + Some(BodyMode::UntilClose(0)) + } else { + let content_length = + header_value_content_length(header.headers.get(http::header::CONTENT_LENGTH)); + match content_length { + Some(length) => Some(BodyMode::ContentLength(length, 0)), + None => Some(BodyMode::UntilClose(0)), + } + } +} + +fn apply_body_mode(body_writer: &mut BodyWriter, mode: BodyMode) { + match mode { + BodyMode::ContentLength(total, 0) => body_writer.init_content_length(total), + BodyMode::UntilClose(0) => body_writer.init_close_delimited(), + _ => body_writer.body_mode = mode, + } +} + +/// Body dispatch variant that avoids borrowing the whole session while a +/// channel permit borrows `self.tx`. +fn dispatch_body_inline( + body_writer: &mut BodyWriter, + body_bytes_sent: &mut usize, + upgraded: bool, + data: Option, + upgraded_task: bool, + permit: mpsc::Permit<'_, HttpTask>, +) -> Result<()> { + if upgraded_task != upgraded { + if upgraded_task { + panic!("Unexpected UpgradedBody task received on un-upgraded downstream session (subrequest, proxy task)"); + } else { + panic!("Unexpected Body task received on upgraded downstream session (subrequest, proxy task)"); + } + } + + let Some(d) = data else { + drop(permit); + return Ok(()); + }; + if d.is_empty() { + drop(permit); + return Ok(()); + } + + let (to_count, next_mode) = match &body_writer.body_mode { + BodyMode::ContentLength(total, written) => { + if written >= total { + drop(permit); + return Ok(()); + } + let remaining = *total - *written; + let to_write = if remaining < d.len() { + warn!("Trying to write data over content-length (subrequest, proxy task): {total}"); + remaining + } else { + d.len() + }; + ( + to_write, + BodyMode::ContentLength(*total, *written + to_write), + ) + } + BodyMode::UntilClose(written) => (d.len(), BodyMode::UntilClose(*written + d.len())), + BodyMode::Complete(_) => { + drop(permit); + return Ok(()); + } + BodyMode::ToSelect => { + drop(permit); + return Error::e_explain( + InternalError, + "subrequest body proxy task before header is sent", + ); + } + }; + + let to_send = if to_count < d.len() { + d.slice(..to_count) + } else { + d + }; + permit.send(HttpTask::Body(Some(to_send), false)); + body_writer.body_mode = next_mode; + *body_bytes_sent += to_count; + Ok(()) } #[cfg(test)] @@ -817,25 +1257,49 @@ mod tests_stream { let _ = env_logger::builder().is_test(true).try_init(); } + fn test_header(status: StatusCode) -> ResponseHeader { + ResponseHeader::build(status, None) + .expect("test status code should build a response header") + } + + fn recv_task(rx: &mut mpsc::Receiver) -> HttpTask { + rx.try_recv() + .expect("expected subrequest output task to be queued") + } + async fn session_from_input(input: &[u8]) -> (HttpSession, SubrequestHandle) { let mock_io = Builder::new().read(input).build(); let mut http_stream = GenericHttpSession::new_http1(Box::new(mock_io)); - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); let (mut http_stream, handle) = HttpSession::new_from_session(&http_stream); - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); (http_stream, handle) } - async fn build_upgrade_req(upgrade: &str, conn: &str) -> (HttpSession, SubrequestHandle) { + pub(super) async fn build_upgrade_req( + upgrade: &str, + conn: &str, + ) -> (HttpSession, SubrequestHandle) { let input = format!("GET / HTTP/1.1\r\nHost: pingora.org\r\nUpgrade: {upgrade}\r\nConnection: {conn}\r\n\r\n"); session_from_input(input.as_bytes()).await } - async fn build_req() -> (HttpSession, SubrequestHandle) { + pub(super) async fn build_req() -> (HttpSession, SubrequestHandle) { let input = "GET / HTTP/1.1\r\nHost: pingora.org\r\n\r\n".to_string(); session_from_input(input.as_bytes()).await } + pub(super) async fn build_head_req() -> (HttpSession, SubrequestHandle) { + let input = "HEAD / HTTP/1.1\r\nHost: pingora.org\r\n\r\n".to_string(); + session_from_input(input.as_bytes()).await + } + #[tokio::test] async fn read_basic() { init_log(); @@ -880,12 +1344,12 @@ mod tests_stream { async fn read_upgrade_req_with_1xx_response() { let (mut http_stream, _handle) = build_upgrade_req("websocket", "upgrade").await; assert!(http_stream.is_upgrade_req()); - let mut response = ResponseHeader::build(StatusCode::CONTINUE, None).unwrap(); + let mut response = test_header(StatusCode::CONTINUE); response.set_version(http::Version::HTTP_11); http_stream .write_response_header(Box::new(response)) .await - .unwrap(); + .expect("test operation should succeed"); // 100 won't affect body state assert!(http_stream.is_body_done()); } @@ -893,13 +1357,15 @@ mod tests_stream { #[tokio::test] async fn write() { let (mut http_stream, mut handle) = build_req().await; - let mut new_response = ResponseHeader::build(StatusCode::OK, None).unwrap(); - new_response.append_header("Foo", "Bar").unwrap(); + let mut new_response = test_header(StatusCode::OK); + new_response + .append_header("Foo", "Bar") + .expect("test operation should succeed"); http_stream .write_response_header_ref(&new_response) .await - .unwrap(); - match handle.rx.try_recv().unwrap() { + .expect("test operation should succeed"); + match recv_task(&mut handle.rx) { HttpTask::Header(header, end) => { assert_eq!(header.status, StatusCode::OK); assert_eq!(header.headers["foo"], "Bar"); @@ -912,12 +1378,12 @@ mod tests_stream { #[tokio::test] async fn write_informational() { let (mut http_stream, mut handle) = build_req().await; - let response_100 = ResponseHeader::build(StatusCode::CONTINUE, None).unwrap(); + let response_100 = test_header(StatusCode::CONTINUE); http_stream .write_response_header_ref(&response_100) .await - .unwrap(); - match handle.rx.try_recv().unwrap() { + .expect("test operation should succeed"); + match recv_task(&mut handle.rx) { HttpTask::Header(header, end) => { assert_eq!(header.status, StatusCode::CONTINUE); assert!(!end); @@ -925,12 +1391,12 @@ mod tests_stream { t => panic!("unexpected task {t:?}"), } - let response_200 = ResponseHeader::build(StatusCode::OK, None).unwrap(); + let response_200 = test_header(StatusCode::OK); http_stream .write_response_header_ref(&response_200) .await - .unwrap(); - match handle.rx.try_recv().unwrap() { + .expect("test operation should succeed"); + match recv_task(&mut handle.rx) { HttpTask::Header(header, end) => { assert_eq!(header.status, StatusCode::OK); assert!(!end); @@ -942,15 +1408,16 @@ mod tests_stream { #[tokio::test] async fn write_101_switching_protocol() { let (mut http_stream, mut handle) = build_upgrade_req("WebSocket", "Upgrade").await; - let mut response_101 = - ResponseHeader::build(StatusCode::SWITCHING_PROTOCOLS, None).unwrap(); - response_101.append_header("Foo", "Bar").unwrap(); + let mut response_101 = test_header(StatusCode::SWITCHING_PROTOCOLS); + response_101 + .append_header("Foo", "Bar") + .expect("test operation should succeed"); http_stream .write_response_header_ref(&response_101) .await - .unwrap(); + .expect("test operation should succeed"); - match handle.rx.try_recv().unwrap() { + match recv_task(&mut handle.rx) { HttpTask::Header(header, end) => { assert_eq!(header.status, StatusCode::SWITCHING_PROTOCOLS); assert!(!end); @@ -964,16 +1431,17 @@ mod tests_stream { .write_body(wire_body.clone()) .await .unwrap() - .unwrap(); + .expect("test operation should succeed"); assert_eq!(wire_body.len(), n); // this write should be ignored - let response_502 = ResponseHeader::build(StatusCode::BAD_GATEWAY, None).unwrap(); + let response_502 = ResponseHeader::build(StatusCode::BAD_GATEWAY, None) + .expect("test operation should succeed"); http_stream .write_response_header_ref(&response_502) .await - .unwrap(); + .expect("test operation should succeed"); - match handle.rx.try_recv().unwrap() { + match recv_task(&mut handle.rx) { HttpTask::Body(body, _end) => { assert_eq!(body.unwrap().len(), n); } @@ -989,12 +1457,14 @@ mod tests_stream { async fn write_body_cl() { let (mut http_stream, _handle) = build_req().await; let wire_body = Bytes::from(&b"a"[..]); - let mut new_response = ResponseHeader::build(StatusCode::OK, None).unwrap(); - new_response.append_header("Content-Length", "1").unwrap(); + let mut new_response = test_header(StatusCode::OK); + new_response + .append_header("Content-Length", "1") + .expect("test operation should succeed"); http_stream .write_response_header_ref(&new_response) .await - .unwrap(); + .expect("test operation should succeed"); assert_eq!( http_stream.body_writer.body_mode, BodyMode::ContentLength(1, 0) @@ -1003,29 +1473,37 @@ mod tests_stream { .write_body(wire_body.clone()) .await .unwrap() - .unwrap(); + .expect("test operation should succeed"); assert_eq!(wire_body.len(), n); - let n = http_stream.finish().await.unwrap().unwrap(); + let n = http_stream + .finish() + .await + .expect("test async operation should succeed") + .expect("test operation should succeed"); assert_eq!(wire_body.len(), n); } #[tokio::test] async fn write_body_until_close() { let (mut http_stream, _handle) = build_req().await; - let new_response = ResponseHeader::build(StatusCode::OK, None).unwrap(); + let new_response = test_header(StatusCode::OK); http_stream .write_response_header_ref(&new_response) .await - .unwrap(); + .expect("test operation should succeed"); assert_eq!(http_stream.body_writer.body_mode, BodyMode::UntilClose(0)); let wire_body = Bytes::from(&b"PAYLOAD"[..]); let n = http_stream .write_body(wire_body.clone()) .await .unwrap() - .unwrap(); + .expect("test operation should succeed"); assert_eq!(wire_body.len(), n); - let n = http_stream.finish().await.unwrap().unwrap(); + let n = http_stream + .finish() + .await + .expect("test async operation should succeed") + .expect("test operation should succeed"); assert_eq!(wire_body.len(), n); } @@ -1037,17 +1515,27 @@ mod tests_stream { let input3 = b"abc"; let mock_io = Builder::new().read(&input1[..]).read(&input2[..]).build(); let mut http_stream = GenericHttpSession::new_http1(Box::new(mock_io)); - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); let (mut http_stream, handle) = HttpSession::new_from_session(&http_stream); - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); handle .tx .send(HttpTask::Body(Some(Bytes::from(&input3[..])), false)) .await - .unwrap(); + .expect("test operation should succeed"); assert_eq!(http_stream.get_path(), &b"/a?q=b%20c"[..]); - let res = http_stream.read_body().await.unwrap().unwrap(); + let res = http_stream + .read_body() + .await + .expect("test async operation should succeed") + .expect("test operation should succeed"); assert_eq!(res, &input3[..]); assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(3)); } @@ -1056,22 +1544,27 @@ mod tests_stream { async fn test_write_body_write_timeout() { let (mut http_stream, _handle) = build_req().await; http_stream.write_timeout = Some(Duration::from_millis(100)); - let mut new_response = ResponseHeader::build(StatusCode::OK, None).unwrap(); - new_response.append_header("Content-Length", "10").unwrap(); + let mut new_response = test_header(StatusCode::OK); + new_response + .append_header("Content-Length", "10") + .expect("test operation should succeed"); http_stream .write_response_header_ref(&new_response) .await - .unwrap(); + .expect("test operation should succeed"); let body_write_buf = Bytes::from(&b"abc"[..]); http_stream .write_body(body_write_buf.clone()) .await - .unwrap(); + .expect("test operation should succeed"); http_stream .write_body(body_write_buf.clone()) .await - .unwrap(); - http_stream.write_body(body_write_buf).await.unwrap(); + .expect("test operation should succeed"); + http_stream + .write_body(body_write_buf) + .await + .expect("test async operation should succeed"); // channel full let last_body = Bytes::from(&b"a"[..]); let res = http_stream.write_body(last_body).await; @@ -1081,8 +1574,11 @@ mod tests_stream { #[tokio::test] async fn test_write_continue_resp() { let (mut http_stream, mut handle) = build_req().await; - http_stream.write_continue_response().await.unwrap(); - match handle.rx.try_recv().unwrap() { + http_stream + .write_continue_response() + .await + .expect("test async operation should succeed"); + match recv_task(&mut handle.rx) { HttpTask::Header(header, end) => { assert_eq!(header.status, StatusCode::CONTINUE); assert!(!end); @@ -1095,7 +1591,10 @@ mod tests_stream { let mock_io = Builder::new().read(input).build(); let mut http_stream = GenericHttpSession::new_http1(Box::new(mock_io)); // Read the request in v1 inner session to set up headers properly - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); let (http_stream, handle) = HttpSession::new_from_session(&http_stream); (http_stream, handle) } @@ -1157,9 +1656,15 @@ mod tests_stream { async fn build_upgrade_req_with_body(header: &[u8]) -> (HttpSession, SubrequestHandle) { let mock_io = Builder::new().read(header).build(); let mut http_stream = GenericHttpSession::new_http1(Box::new(mock_io)); - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); let (mut http_stream, handle) = HttpSession::new_from_session(&http_stream); - http_stream.read_request().await.unwrap(); + http_stream + .read_request() + .await + .expect("test async operation should succeed"); (http_stream, handle) } @@ -1179,10 +1684,14 @@ mod tests_stream { .tx .send(HttpTask::Body(Some(Bytes::from(POST_BODY_DATA)), true)) .await - .unwrap(); + .expect("test operation should succeed"); let mut buf = vec![]; - while let Some(b) = http_stream.read_body_bytes().await.unwrap() { + while let Some(b) = http_stream + .read_body_bytes() + .await + .expect("test async operation should succeed") + { buf.put_slice(&b); } assert_eq!(buf, POST_BODY_DATA); @@ -1191,12 +1700,12 @@ mod tests_stream { assert!(http_stream.is_body_done()); - let mut response = ResponseHeader::build(StatusCode::SWITCHING_PROTOCOLS, None).unwrap(); + let mut response = test_header(StatusCode::SWITCHING_PROTOCOLS); response.set_version(http::Version::HTTP_11); http_stream .write_response_header(Box::new(response)) .await - .unwrap(); + .expect("test operation should succeed"); // body reader type switches assert!(!http_stream.is_body_done()); @@ -1206,15 +1715,1128 @@ mod tests_stream { .tx .send(HttpTask::Body(Some(Bytes::from(&ws_data[..])), false)) .await - .unwrap(); + .expect("test operation should succeed"); - let buf = http_stream.read_body_bytes().await.unwrap().unwrap(); + let buf = http_stream + .read_body_bytes() + .await + .expect("test async operation should succeed") + .expect("test operation should succeed"); assert_eq!(buf, ws_data.as_slice()); assert!(!http_stream.is_body_done()); // EOF ends body drop(handle.tx); - assert!(http_stream.read_body_bytes().await.unwrap().is_none()); + assert!(http_stream + .read_body_bytes() + .await + .expect("test async operation should succeed") + .is_none()); assert!(http_stream.is_body_done()); } } + +#[cfg(test)] +mod test_proxy_tasks { + //! Cancel-safe proxy task API tests for the subrequest server session. + + use super::tests_stream::{build_head_req, build_req, build_upgrade_req}; + use super::*; + use http::StatusCode; + + fn init_log() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + fn test_header(status: StatusCode) -> ResponseHeader { + ResponseHeader::build(status, None) + .expect("test status code should build a response header") + } + + fn recv_task(rx: &mut mpsc::Receiver) -> HttpTask { + rx.try_recv() + .expect("expected subrequest output task to be queued") + } + + fn assert_rx_empty(rx: &mut mpsc::Receiver) { + assert!(matches!( + rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + + /// Dropping a blocked `send` does not deliver its value. + #[tokio::test(start_paused = true)] + async fn test_tokio_mpsc_send_cancel_drops_value() { + let (tx, mut rx) = mpsc::channel::(1); + tx.send(1) + .await + .expect("test async operation should succeed"); + let send_fut = tx.send(2); + tokio::pin!(send_fut); + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(10)) => {} + _ = &mut send_fut => panic!("expected the timer to win"), + }; + assert_eq!(rx.recv().await, Some(1)); + assert_eq!(rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)); + } + + /// Dropping a blocked `reserve` does not consume capacity. + #[tokio::test(start_paused = true)] + async fn test_tokio_mpsc_reserve_cancel_releases_slot() { + let (tx, mut rx) = mpsc::channel::(1); + tx.send(1) + .await + .expect("test async operation should succeed"); + let reserve_fut = tx.reserve(); + tokio::pin!(reserve_fut); + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(10)) => {} + _ = &mut reserve_fut => panic!("expected the timer to win"), + }; + assert_eq!(rx.recv().await, Some(1)); + assert_eq!(rx.try_recv(), Err(mpsc::error::TryRecvError::Empty)); + } + + #[tokio::test] + async fn test_send_proxy_task_and_write() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + assert!(session.proxy_tasks_enabled()); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "5") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("hello")), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!(end); + assert!(!session.has_pending_proxy_tasks()); + assert_eq!(session.body_bytes_sent(), 5); + + match recv_task(&mut handle.rx) { + HttpTask::Header(h, false) => assert_eq!(h.status, StatusCode::OK), + t => panic!("expected Header, got {t:?}"), + } + match recv_task(&mut handle.rx) { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"hello"), + t => panic!("expected Body, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_informational_head_does_not_init_body_writer() { + init_log(); + + let (mut regular, mut regular_handle) = build_head_req().await; + regular + .write_response_header(Box::new(test_header(StatusCode::CONTINUE))) + .await + .expect("regular informational header write should succeed"); + assert_eq!(regular.body_writer.body_mode, BodyMode::ToSelect); + assert!(matches!( + recv_task(&mut regular_handle.rx), + HttpTask::Header(..) + )); + + let (mut proxy, mut proxy_handle) = build_head_req().await; + proxy.set_proxy_tasks_enabled(true); + proxy.send_proxy_task(HttpTask::Header( + Box::new(test_header(StatusCode::CONTINUE)), + false, + )); + assert!(!proxy + .write_proxy_tasks() + .await + .expect("proxy task write should succeed")); + assert_eq!(proxy.body_writer.body_mode, BodyMode::ToSelect); + assert!(matches!( + recv_task(&mut proxy_handle.rx), + HttpTask::Header(..) + )); + } + + #[tokio::test(start_paused = true)] + async fn test_proxy_task_with_timeout() { + init_log(); + // Do not drain `handle.rx` before the first write; the 5th response task blocks on capacity. + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + session.set_write_timeout(Some(Duration::from_millis(50))); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + for i in 0..5 { + session.send_proxy_task(HttpTask::Body( + Some(Bytes::from(format!("body-{i}"))), + i == 4, + )); + } + + let err = session + .write_proxy_tasks() + .await + .expect_err("full subrequest output channel should time out"); + assert_eq!(err.etype(), &WriteTimedout); + assert!(session.has_pending_proxy_tasks()); + + let mut delivered = Vec::new(); + while let Ok(task) = handle.rx.try_recv() { + delivered.push(task); + } + assert_eq!(delivered.len(), 4); + assert!(matches!(delivered[0], HttpTask::Header(..))); + + session.set_write_timeout(None); + let end = session + .write_proxy_tasks() + .await + .expect("retry after freeing channel capacity should complete"); + assert!(end); + while let Ok(task) = handle.rx.try_recv() { + delivered.push(task); + } + + assert_eq!(delivered.len(), 7); + assert!(matches!(delivered.last(), Some(HttpTask::Done))); + let body_count = delivered + .iter() + .filter(|t| matches!(t, HttpTask::Body(..))) + .count(); + assert_eq!(body_count, 5); + } + + #[tokio::test] + async fn test_proxy_task_channel_closed_errors() { + init_log(); + let (mut session, handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + drop(handle.rx); + + session.send_proxy_task(HttpTask::Header( + Box::new(test_header(StatusCode::OK)), + false, + )); + let err = session + .write_proxy_tasks() + .await + .expect_err("closed subrequest output channel should error"); + assert_eq!(err.etype(), &WriteError); + assert!(session.has_pending_proxy_tasks()); + } + + /// Repeatedly cancel while blocked on channel capacity, then verify the + /// receiver sees each queued task exactly once and in order. + #[tokio::test(start_paused = true)] + async fn test_proxy_task_cancel_safety() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "5") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + for i in 0..4 { + session.send_proxy_task(HttpTask::Body( + Some(Bytes::from(vec![b'A' + i as u8; 1])), + false, + )); + } + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("E")), true)); + + let mut cancel_count = 0; + let mut delivered: Vec = Vec::new(); + loop { + if !session.has_pending_proxy_tasks() { + break; + } + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => { + cancel_count += 1; + while let Ok(task) = handle.rx.try_recv() { + delivered.push(task); + } + } + result = session.write_proxy_tasks() => { + result.expect("test operation should succeed"); + } + } + } + + assert!( + cancel_count >= 1, + "expected at least one cancellation during cancel-safe write, got {cancel_count}" + ); + + assert_eq!(session.proxy_task_state.tasks.len(), 0); + + while let Ok(task) = handle.rx.try_recv() { + delivered.push(task); + } + + assert!(matches!(delivered[0], HttpTask::Header(_, false))); + let mut body_bytes = Vec::new(); + let mut saw_done = false; + for task in &delivered[1..] { + match task { + HttpTask::Body(Some(b), false) => body_bytes.extend_from_slice(b), + HttpTask::Done => { + assert!(!saw_done, "Done delivered more than once"); + saw_done = true; + } + t => panic!("unexpected task in delivery: {t:?}"), + } + } + assert!(saw_done, "expected Done to be delivered"); + assert_eq!( + body_bytes, b"ABCDE", + "body chunks must arrive exactly once, in order" + ); + + assert_eq!(session.body_bytes_sent(), 5); + } + + /// `was_upgraded()` must remain false if the 101 send is cancelled before + /// it reaches the subrequest channel. + #[tokio::test(start_paused = true)] + async fn test_proxy_task_upgrade_consistency() { + init_log(); + let (mut session, mut handle) = build_upgrade_req("websocket", "Upgrade").await; + assert!(session.is_upgrade_req()); + session.set_proxy_tasks_enabled(true); + + // Four 1xx headers fill the upstream channel; the 101 then blocks. + for _ in 0..4 { + session.send_proxy_task(HttpTask::Header( + Box::new(test_header(StatusCode::CONTINUE)), + false, + )); + } + let mut h101 = test_header(StatusCode::SWITCHING_PROTOCOLS); + h101.set_version(http::Version::HTTP_11); + h101.insert_header("upgrade", "websocket") + .expect("test operation should succeed"); + h101.insert_header("connection", "Upgrade") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(h101), false)); + + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => {} + _ = session.write_proxy_tasks() => panic!("expected reserve to be cancelled"), + }; + + assert!( + !session.was_upgraded(), + "was_upgraded must remain false until the 101 send actually completes" + ); + + for _ in 0..4 { + recv_task(&mut handle.rx); + } + session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!( + session.was_upgraded(), + "after the 101 send completes, was_upgraded must be true" + ); + } + + /// Same upgrade consistency check for the regular `write_response_header` + /// path, which also awaits on the subrequest output channel. + #[tokio::test(start_paused = true)] + async fn test_write_response_header_upgrade_cancel_consistency() { + init_log(); + let (mut session, mut handle) = build_upgrade_req("websocket", "Upgrade").await; + + for _ in 0..4 { + session + .write_response_header(Box::new(test_header(StatusCode::CONTINUE))) + .await + .expect("test operation should succeed"); + } + + let mut h101 = test_header(StatusCode::SWITCHING_PROTOCOLS); + h101.set_version(http::Version::HTTP_11); + h101.insert_header("upgrade", "websocket") + .expect("test operation should succeed"); + h101.insert_header("connection", "Upgrade") + .expect("test operation should succeed"); + + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => {} + _ = session.write_response_header(Box::new(h101.clone())) => { + panic!("expected header send to be cancelled") + } + }; + assert!(!session.was_upgraded()); + assert_eq!(session.body_writer.body_mode, BodyMode::ToSelect); + + for _ in 0..4 { + recv_task(&mut handle.rx); + } + session + .write_response_header(Box::new(h101)) + .await + .expect("test async operation should succeed"); + assert!(session.was_upgraded()); + } + + /// Trailers are dispatched correctly through `write_proxy_tasks`. + /// Matching regular `response_duplex_vec`, a final `Done` follows Trailer. + #[tokio::test] + async fn test_proxy_task_trailers() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("hi")), false)); + let mut trailers = http::HeaderMap::new(); + trailers.insert("x-final", http::HeaderValue::from_static("done")); + session.send_proxy_task(HttpTask::Trailer(Some(Box::new(trailers)))); + + let end = session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!(end); + + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Body(..))); + match recv_task(&mut handle.rx) { + HttpTask::Trailer(Some(t)) => { + assert_eq!( + t.get("x-final").expect("test trailer should be present"), + "done" + ); + } + t => panic!("expected Trailer, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_trailer_before_content_length_complete_errors() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "5") + .expect("test content-length header is valid"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + let mut trailers = HeaderMap::new(); + trailers.insert("x-final", http::HeaderValue::from_static("done")); + session.send_proxy_task(HttpTask::Trailer(Some(Box::new(trailers)))); + + let err = session + .write_proxy_tasks() + .await + .expect_err("trailers before content-length body completion should error"); + assert_eq!(err.etype(), &PREMATURE_BODY_END); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Trailer(..))); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test(start_paused = true)] + async fn test_proxy_task_trailer_cancel_safety() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + for _ in 0..4 { + session.send_proxy_task(HttpTask::Header( + Box::new(test_header(StatusCode::CONTINUE)), + false, + )); + } + let mut trailers = HeaderMap::new(); + trailers.insert("x-final", http::HeaderValue::from_static("done")); + session.send_proxy_task(HttpTask::Trailer(Some(Box::new(trailers)))); + + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => {} + _ = session.write_proxy_tasks() => panic!("expected trailer reserve to be cancelled"), + }; + + let mut prefix = Vec::new(); + while let Ok(task) = handle.rx.try_recv() { + prefix.push(task); + } + assert_eq!(prefix.len(), 4); + assert!(prefix.iter().all(|t| matches!(t, HttpTask::Header(..)))); + let end = session + .write_proxy_tasks() + .await + .expect("resume after trailer cancellation should complete"); + assert!(end); + + match recv_task(&mut handle.rx) { + HttpTask::Trailer(Some(t)) => { + assert_eq!(t.get("x-final").expect("trailer present"), "done") + } + t => panic!("expected Trailer after resume, got {t:?}"), + } + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test(start_paused = true)] + async fn test_proxy_task_trailer_cancel_safety_after_body() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + for _ in 0..3 { + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("x")), false)); + } + let mut trailers = HeaderMap::new(); + trailers.insert("x-final", http::HeaderValue::from_static("done")); + session.send_proxy_task(HttpTask::Trailer(Some(Box::new(trailers)))); + + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => {} + _ = session.write_proxy_tasks() => panic!("expected trailer reserve to be cancelled"), + }; + + let mut prefix = Vec::new(); + while let Ok(task) = handle.rx.try_recv() { + prefix.push(task); + } + assert_eq!(prefix.len(), 4); + assert!(matches!(prefix[0], HttpTask::Header(..))); + assert!(prefix[1..].iter().all(|t| matches!(t, HttpTask::Body(..)))); + let end = session + .write_proxy_tasks() + .await + .expect("resume after body trailer cancellation should complete"); + assert!(end); + + match recv_task(&mut handle.rx) { + HttpTask::Trailer(Some(t)) => { + assert_eq!(t.get("x-final").expect("trailer present"), "done") + } + t => panic!("expected Trailer after resume, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + } + + /// `body_bytes_sent` is only incremented after the synchronous + /// `Permit::send`, not on a cancelled `reserve().await`. + #[tokio::test(start_paused = true)] + async fn test_proxy_task_body_counter_no_double_count() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "12") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("AAAA")), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("BBBB")), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("CCCC")), true)); + + let mut received_body_bytes = Vec::new(); + loop { + if !session.has_pending_proxy_tasks() { + break; + } + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => { + while let Ok(task) = handle.rx.try_recv() { + if let HttpTask::Body(Some(b), _) = &task { + received_body_bytes.extend_from_slice(b); + } + } + assert!( + session.body_bytes_sent() <= received_body_bytes.len(), + "body_bytes_sent ({}) must not exceed bytes actually delivered ({})", + session.body_bytes_sent(), + received_body_bytes.len(), + ); + } + result = session.write_proxy_tasks() => { + result.expect("test operation should succeed"); + } + } + } + + while let Ok(task) = handle.rx.try_recv() { + if let HttpTask::Body(Some(b), _) = &task { + received_body_bytes.extend_from_slice(b); + } + } + + assert_eq!(session.body_bytes_sent(), 12); + assert_eq!(&received_body_bytes[..], b"AAAABBBBCCCC"); + } + + /// Cancelling while reserving capacity for the final Done must leave the + /// finish operation resumable. + #[tokio::test(start_paused = true)] + async fn test_proxy_task_finish_cancel_safety() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("a")), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("b")), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("c")), true)); + + // Header + three bodies fill the 4-slot channel. The final Done + // reserve blocks and is cancelled. + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => {} + _ = session.write_proxy_tasks() => panic!("expected finish reserve to be cancelled"), + }; + assert!(session.proxy_task_state.finish_in_progress); + + let first = recv_task(&mut handle.rx); + assert!(matches!(first, HttpTask::Header(..))); + + // Only one slot is available. Resuming must emit exactly one Done and + // return without trying to reserve a second slot. + let end = tokio::time::timeout(Duration::from_millis(5), session.write_proxy_tasks()) + .await + .expect("resume should not need a second channel slot") + .expect("test operation should succeed"); + assert!(end); + + let mut delivered = Vec::new(); + while let Ok(task) = handle.rx.try_recv() { + delivered.push(task); + } + assert_eq!(delivered.len(), 4); + assert!(matches!(delivered.last(), Some(HttpTask::Done))); + assert_eq!( + delivered + .iter() + .filter(|t| matches!(t, HttpTask::Done)) + .count(), + 1 + ); + } + + #[tokio::test] + async fn test_proxy_task_done_only_noops_without_channel() { + init_log(); + let (mut session, _handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + session.send_proxy_task(HttpTask::Done); + session.shutdown(); + + let end = session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!(end); + } + + #[tokio::test] + async fn test_proxy_task_done_only_noops_with_live_channel() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + session.send_proxy_task(HttpTask::Done); + + let end = session + .write_proxy_tasks() + .await + .expect("Done-only proxy task should complete without channel output"); + assert!(end); + assert!(!session.has_pending_proxy_tasks()); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_header_only_end() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::NO_CONTENT); + session.send_proxy_task(HttpTask::Header(Box::new(header), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!(end); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_head_response_drops_body() { + init_log(); + let (mut session, mut handle) = build_head_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "10") + .expect("test content-length header is valid"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("not-sent")), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("HEAD proxy task response should complete"); + assert!(end); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + assert_eq!(session.body_bytes_sent(), 0); + } + + #[tokio::test(start_paused = true)] + async fn test_proxy_task_duplicate_final_header_does_not_reserve() { + init_log(); + let (mut session, _handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + + for _ in 0..3 { + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("x")), false)); + } + let duplicate = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(duplicate), false)); + + tokio::time::timeout(Duration::from_millis(5), session.write_proxy_tasks()) + .await + .expect("duplicate final header should be dropped without reserving") + .expect("test operation should succeed"); + assert!(!session.has_pending_proxy_tasks()); + } + + #[tokio::test] + async fn test_proxy_task_duplicate_final_header_preserves_end() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + recv_task(&mut handle.rx); + + let duplicate = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(duplicate), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!(end); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_failed_propagates_without_sending_later_tasks() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + session.send_proxy_task(HttpTask::Failed(Error::explain(InternalError, "boom"))); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("late")), true)); + + let err = session + .write_proxy_tasks() + .await + .expect_err("Failed proxy task should propagate error"); + assert_eq!(err.etype(), &InternalError); + assert!(session.has_pending_proxy_tasks()); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_failed_clears_sticky_eos() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(None, true)); + session.send_proxy_task(HttpTask::Failed(Error::explain(InternalError, "boom"))); + + let err = session + .write_proxy_tasks() + .await + .expect_err("Failed after EOS should still propagate error"); + assert_eq!(err.etype(), &InternalError); + assert!(!session.has_pending_proxy_tasks()); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + assert_rx_empty(&mut handle.rx); + + let end = session + .write_proxy_tasks() + .await + .expect("retry after Failed should not emit stale Done"); + assert!(!end); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_body_before_header_errors() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("body")), true)); + + let err = session + .write_proxy_tasks() + .await + .expect_err("body before response header should be rejected"); + assert_eq!(err.etype(), &InternalError); + assert!(!session.has_pending_proxy_tasks()); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test(start_paused = true)] + async fn test_proxy_task_none_body_does_not_reserve() { + init_log(); + let (mut session, _handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + for _ in 0..3 { + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("x")), false)); + } + session.send_proxy_task(HttpTask::Body(None, false)); + + tokio::time::timeout(Duration::from_millis(5), session.write_proxy_tasks()) + .await + .expect("Body(None) should not reserve channel capacity") + .expect("test operation should succeed"); + assert!(!session.has_pending_proxy_tasks()); + } + + #[tokio::test(start_paused = true)] + async fn test_proxy_task_no_data_eos_survives_cancellation() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + for _ in 0..3 { + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("x")), false)); + } + session.send_proxy_task(HttpTask::Body(None, true)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::new()), true)); + // This later body blocks on the full channel after the no-data EOS + // task has been consumed. The sticky EOS flag must survive that cancel. + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("y")), false)); + + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(5)) => {} + _ = session.write_proxy_tasks() => panic!("expected reserve after no-data EOS to be cancelled"), + }; + + let mut prefix = Vec::new(); + while let Ok(task) = handle.rx.try_recv() { + prefix.push(task); + } + assert_eq!(prefix.len(), 4); + assert!(matches!(prefix[0], HttpTask::Header(..))); + assert!(prefix[1..].iter().all(|t| matches!(t, HttpTask::Body(..)))); + let end = session + .write_proxy_tasks() + .await + .expect("resume after no-data EOS cancellation should complete"); + assert!(end); + + let mut delivered = Vec::new(); + while let Ok(task) = handle.rx.try_recv() { + delivered.push(task); + } + assert_eq!(delivered.len(), 2); + match &delivered[0] { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"y"), + t => panic!("expected Body(y) after resume, got {t:?}"), + } + assert!(matches!(delivered[1], HttpTask::Done)); + assert_eq!(session.body_bytes_sent(), 4); + } + + #[tokio::test] + async fn test_proxy_task_content_length_overrun_truncates() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "3") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("abcdef")), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + assert!(end); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + match recv_task(&mut handle.rx) { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"abc"), + t => panic!("expected truncated Body, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + assert_eq!(session.body_bytes_sent(), 3); + } + + #[tokio::test] + async fn test_proxy_task_exact_content_length_without_end_sends_done() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "3") + .expect("test content-length header is valid"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("abc")), false)); + + let end = session + .write_proxy_tasks() + .await + .expect("exact content-length proxy task response should complete"); + assert!(end); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + match recv_task(&mut handle.rx) { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"abc"), + t => panic!("expected exact content-length Body, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + async fn test_proxy_task_late_tasks_after_finished_are_dropped() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "3") + .expect("test content-length header is valid"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("abc")), true)); + session + .write_proxy_tasks() + .await + .expect("initial response should complete"); + + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Body(..))); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + + let mut trailers = HeaderMap::new(); + trailers.insert("x-late", http::HeaderValue::from_static("ignored")); + session.send_proxy_task(HttpTask::Trailer(Some(Box::new(trailers)))); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("late")), true)); + session.send_proxy_task(HttpTask::Failed(Error::explain(InternalError, "late"))); + + let end = session + .write_proxy_tasks() + .await + .expect("late tasks after finished stream should be dropped"); + assert!(end); + assert_rx_empty(&mut handle.rx); + assert!(!session.has_pending_proxy_tasks()); + } + + #[tokio::test] + async fn test_proxy_task_chunked_header_uses_until_close() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("transfer-encoding", "chunked") + .expect("test transfer-encoding header is valid"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("chunk")), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("chunked proxy task response should complete"); + assert!(end); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + match recv_task(&mut handle.rx) { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"chunk"), + t => panic!("expected chunked body task, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_eq!(session.body_bytes_sent(), 5); + } + + #[tokio::test] + async fn test_proxy_task_premature_content_length_errors_before_done() { + init_log(); + let (mut session, mut handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + let mut header = test_header(StatusCode::OK); + header + .insert_header("content-length", "5") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("hi")), true)); + + let err = session + .write_proxy_tasks() + .await + .expect_err("premature content-length should error before Done"); + assert_eq!(err.etype(), &PREMATURE_BODY_END); + assert_eq!(session.body_writer.body_mode, BodyMode::Complete(2)); + assert!(!session.has_pending_proxy_tasks()); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + match recv_task(&mut handle.rx) { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"hi"), + t => panic!("expected body before premature end, got {t:?}"), + } + assert_rx_empty(&mut handle.rx); + } + + #[tokio::test] + #[should_panic( + expected = "Unexpected UpgradedBody task received on un-upgraded downstream session" + )] + async fn test_upgraded_body_on_non_upgraded_session_panics() { + init_log(); + let (mut session, _handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + assert!(!session.was_upgraded()); + + let header = test_header(StatusCode::OK); + session.send_proxy_task(HttpTask::Header(Box::new(header), false)); + session.send_proxy_task(HttpTask::UpgradedBody(Some(Bytes::from("ws")), true)); + + let _ = session.write_proxy_tasks().await; + } + + #[tokio::test] + #[should_panic(expected = "Unexpected Body task received on upgraded downstream session")] + async fn test_body_on_upgraded_session_panics() { + init_log(); + let (mut session, _handle) = build_upgrade_req("websocket", "Upgrade").await; + session.set_proxy_tasks_enabled(true); + + let mut h101 = test_header(StatusCode::SWITCHING_PROTOCOLS); + h101.set_version(http::Version::HTTP_11); + h101.insert_header("upgrade", "websocket") + .expect("test operation should succeed"); + h101.insert_header("connection", "Upgrade") + .expect("test operation should succeed"); + session.send_proxy_task(HttpTask::Header(Box::new(h101), false)); + session + .write_proxy_tasks() + .await + .expect("test async operation should succeed"); + + session.send_proxy_task(HttpTask::Body(Some(Bytes::from("plain")), true)); + let _ = session.write_proxy_tasks().await; + } + + #[tokio::test] + async fn test_proxy_task_upgraded_body_happy_path() { + init_log(); + let (mut session, mut handle) = build_upgrade_req("websocket", "Upgrade").await; + session.set_proxy_tasks_enabled(true); + + let mut h101 = test_header(StatusCode::SWITCHING_PROTOCOLS); + h101.set_version(http::Version::HTTP_11); + h101.insert_header("upgrade", "websocket") + .expect("test upgrade header is valid"); + h101.insert_header("connection", "Upgrade") + .expect("test connection header is valid"); + session.send_proxy_task(HttpTask::Header(Box::new(h101), false)); + session.send_proxy_task(HttpTask::UpgradedBody(Some(Bytes::from("ws")), true)); + + let end = session + .write_proxy_tasks() + .await + .expect("upgraded proxy task response should complete"); + assert!(end); + assert!(session.was_upgraded()); + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Header(..))); + match recv_task(&mut handle.rx) { + HttpTask::Body(Some(b), false) => assert_eq!(&b[..], b"ws"), + t => panic!("expected upgraded body task, got {t:?}"), + } + assert!(matches!(recv_task(&mut handle.rx), HttpTask::Done)); + assert_eq!(session.body_bytes_sent(), 2); + } + + #[tokio::test] + #[should_panic( + expected = "Unexpected UpgradedBody task received on un-upgraded downstream session" + )] + async fn test_upgraded_body_on_non_upgraded_session_panics_while_full() { + init_log(); + let (mut session, _handle) = build_req().await; + session.set_proxy_tasks_enabled(true); + + for _ in 0..4 { + session.send_proxy_task(HttpTask::Header( + Box::new(test_header(StatusCode::CONTINUE)), + false, + )); + } + session.send_proxy_task(HttpTask::UpgradedBody(Some(Bytes::from("ws")), true)); + let _ = session.write_proxy_tasks().await; + } +} diff --git a/pingora-core/src/protocols/http/v1/body.rs b/pingora-core/src/protocols/http/v1/body.rs index 72899257c..460137c6c 100644 --- a/pingora-core/src/protocols/http/v1/body.rs +++ b/pingora-core/src/protocols/http/v1/body.rs @@ -20,9 +20,11 @@ use pingora_error::{ OrErr, Result, }; use std::fmt::Debug; +use std::future::Future; +use std::pin::{pin, Pin}; +use std::task::{ready, Context, Poll}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; -use crate::protocols::l4::stream::AsyncWriteVec; use crate::utils::BufRef; // TODO: make this dynamically adjusted @@ -249,6 +251,47 @@ impl BodyReader { self.get_body_overread().is_some_and(|b| !b.is_empty()) } + /// Take ownership of the overread bytes, leaving `None` in their place. + /// + /// Overread bytes are bytes that were read from the stream but are beyond the + /// end of the current request's body — i.e. they belong to a pipelined next + /// request on the same connection. Callers that support HTTP/1.1 pipelining + /// extract these bytes here and feed them to the next session's request + /// parser via [`HttpSession::set_pipelined_prefix`](super::server::HttpSession::set_pipelined_prefix). + pub fn take_body_overread(&mut self) -> Option { + self.body_buf_overread.take() + } + + /// Append bytes to the overread buffer from outside the body-parsing path. + /// + /// The body reader's overread buffer is normally populated by + /// [`Self::init_content_length`] (for zero-length bodies) and + /// [`Self::finish_body_buf`] (for sized bodies) when the stream read + /// that completed the current request also pulled in bytes from the + /// next pipelined request. That path covers the "both requests in one + /// read" shape. + /// + /// The "second request arrives in a separate read after the first + /// request's body is already done" shape is different: the body + /// reader never sees those bytes — they land on the idle-branch read + /// in [`super::server::HttpSession::read_body_or_idle`]. When + /// pipelining is enabled, that caller stashes the idle read here so + /// a single downstream overread surface covers both shapes and + /// [`Self::take_body_overread`] returns them uniformly. + pub fn push_body_overread(&mut self, bytes: &[u8]) { + if bytes.is_empty() { + return; + } + match self.body_buf_overread.as_mut() { + Some(buf) => buf.extend_from_slice(bytes), + None => { + let mut buf = BytesMut::with_capacity(bytes.len()); + buf.extend_from_slice(bytes); + self.body_buf_overread = Some(buf); + } + } + } + pub fn body_done(&self) -> bool { matches!(self.body_state, PS::Complete(_) | PS::Done(_)) } @@ -905,14 +948,188 @@ pub enum BodyMode { type BM = BodyMode; +// ============================================================================ +// Cancel-safe body writing types +// ============================================================================ + +impl BodyMode { + /// Extract `(total, written)` from `ContentLength`, panicking on mismatch. + fn expect_content_length(&self) -> (usize, usize) { + match self { + BodyMode::ContentLength(total, written) => (*total, *written), + _ => panic!("wrong body mode: expected ContentLength, got {:?}", self), + } + } + + /// Extract `written` from `ChunkedEncoding`, panicking on mismatch. + fn expect_chunked(&self) -> usize { + match self { + BodyMode::ChunkedEncoding(written) => *written, + _ => panic!("wrong body mode: expected ChunkedEncoding, got {:?}", self), + } + } + + /// Extract `written` from `UntilClose`, panicking on mismatch. + fn expect_until_close(&self) -> usize { + match self { + BodyMode::UntilClose(written) => *written, + _ => panic!("wrong body mode: expected UntilClose, got {:?}", self), + } + } +} + +/// Type alias for the chunked encoding buffer chain +type ChunkedBuf = bytes::buf::Chain, &'static [u8]>; + +enum WriteBuf { + /// Simple bytes buffer + Simple(Bytes), + /// Chained buffer for chunked encoding or other complex writes + Chained(C), +} + +// Implement Buf for WriteBuf to delegate to the inner buffer +impl Buf for WriteBuf { + fn remaining(&self) -> usize { + match self { + WriteBuf::Simple(b) => b.remaining(), + WriteBuf::Chained(c) => c.remaining(), + } + } + + fn chunk(&self) -> &[u8] { + match self { + WriteBuf::Simple(b) => b.chunk(), + WriteBuf::Chained(c) => c.chunk(), + } + } + + fn advance(&mut self, cnt: usize) { + match self { + WriteBuf::Simple(b) => b.advance(cnt), + WriteBuf::Chained(c) => c.advance(cnt), + } + } +} + +enum WriteState { + /// No write in progress + Idle, + /// Writing data (original size, bytes remaining to write) + Writing(usize, WriteBuf), + /// Flushing after write (original size to return) + Flushing(usize), + /// Write complete (bytes written in this task) + Done(usize), + /// Write timed out - cannot be reused + TimedOut, +} + +// Custom Debug implementation since we can't derive it with futures +impl std::fmt::Debug for WriteState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + WriteState::Idle => write!(f, "Idle"), + WriteState::Writing(size, _buf) => { + write!(f, "Writing(size: {})", size) + } + WriteState::Flushing(size) => write!(f, "Flushing(size: {})", size), + WriteState::Done(size) => write!(f, "Done(size: {})", size), + WriteState::TimedOut => write!(f, "TimedOut"), + } + } +} + +enum FinishWriteState { + /// No finish task queued + NotStarted, + /// Finish queued but not started yet + Idle, + /// Writing last chunk marker (for chunked encoding) + WritingLastChunk(WriteBuf), + /// Flushing after writing last chunk + Flushing, + /// Finish complete + Done, +} + +// Custom Debug implementation since WriteBuf doesn't implement Debug +impl std::fmt::Debug for FinishWriteState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FinishWriteState::NotStarted => write!(f, "NotStarted"), + FinishWriteState::Idle => write!(f, "Idle"), + FinishWriteState::WritingLastChunk(_) => write!(f, "WritingLastChunk"), + FinishWriteState::Flushing => write!(f, "Flushing"), + FinishWriteState::Done => write!(f, "Done"), + } + } +} + +/// Internal state for the cancel-safe body write state machine. +/// +/// Tracks the pending body bytes, write progress +/// (idle → writing → flushing → done), and an optional timeout. +struct SendBodyState { + /// Application bytes queued to be written + pending_bytes: Option, + /// Current write state for cancel-safe operations + write_state: WriteState, + /// Timeout duration for this write task + timeout_duration: Option, + /// Timeout future (only created if write returns Pending) + timeout_fut: Option + Send + Sync>>>, +} + +impl std::fmt::Debug for SendBodyState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SendBodyState") + .field("pending_bytes", &self.pending_bytes) + .field("write_state", &self.write_state) + .field("timeout_duration", &self.timeout_duration) + .field( + "timeout_fut", + &self.timeout_fut.as_ref().map(|_| "Some(Future)"), + ) + .finish() + } +} + +impl SendBodyState { + fn new() -> Self { + SendBodyState { + pending_bytes: None, + write_state: WriteState::Idle, + timeout_duration: None, + timeout_fut: None, + } + } +} + +/// Tracks how response body bytes are framed and written to the wire. +/// +/// Supports both a legacy async API (`write_body` / `finish`) and a cancel-safe +/// task API that can be driven inside a `tokio::select!` loop without losing +/// write progress. pub struct BodyWriter { pub body_mode: BodyMode, + // Boxed to reduce inline size. Only used by the cancel-safe proxy task API. + send_body_state: Box, + send_finish_state: FinishWriteState, +} + +impl Default for BodyWriter { + fn default() -> Self { + Self::new() + } } impl BodyWriter { pub fn new() -> Self { BodyWriter { body_mode: BM::ToSelect, + send_body_state: Box::new(SendBodyState::new()), + send_finish_state: FinishWriteState::NotStarted, } } @@ -1019,7 +1236,7 @@ impl BodyWriter { let chuck_size_buf = format!("{:X}\r\n", chunk_size); let mut output_buf = Bytes::from(chuck_size_buf).chain(buf).chain(&b"\r\n"[..]); stream - .write_vec_all(&mut output_buf) + .write_all_buf(&mut output_buf) .await .or_err(WriteError, "while writing body")?; stream.flush().await.or_err(WriteError, "flushing body")?; @@ -1109,6 +1326,549 @@ impl BodyWriter { _ => panic!("wrong body mode: {:?}", self.body_mode), } } + + // ======================================================================== + // Cancel-safe body task API + // ======================================================================== + + #[cfg(test)] + pub fn has_pending_body_task(&self) -> bool { + self.send_body_state.pending_bytes.is_some() + || !matches!( + self.send_body_state.write_state, + WriteState::Idle | WriteState::Done(_) | WriteState::TimedOut + ) + } + + /// Queue application bytes as a body write task with an optional timeout. + /// This is a non-async function that just saves the bytes. + /// Call `write_current_body_task()` to actually perform the write. + /// + /// The timeout, if provided, will be enforced internally across all + /// write attempts, even if the write is cancelled and resumed via `tokio::select!`. + pub fn send_body_task(&mut self, bytes: Bytes, timeout: Option) { + assert!( + matches!( + self.send_body_state.write_state, + WriteState::Idle | WriteState::Done(_) + ), + "send_body_task called while previous task is still in progress: {:?}", + self.send_body_state.write_state + ); + self.send_body_state.pending_bytes = Some(bytes); + self.send_body_state.write_state = WriteState::Idle; + self.send_body_state.timeout_duration = timeout; + self.send_body_state.timeout_fut = None; + } + + /// Writes the current queued body task to the stream. + /// + /// ## Cancel-safety + /// + /// This function can be safely used in a `tokio::select!` loop. + /// Returns `Ok(Some(bytes_written))` when complete, `Ok(None)` if no bytes to write. + pub async fn write_current_body_task(&mut self, stream: &mut S) -> Result> + where + S: AsyncWrite + Unpin + Send, + { + // Use poll_fn to wrap our poll-based implementation + std::future::poll_fn(|cx| self.poll_write_current_body_task(cx, Pin::new(stream))).await + } + + /// Poll-based implementation for writing body tasks. + /// This is the core implementation that maintains state across cancellations. + fn poll_write_current_body_task( + &mut self, + cx: &mut Context<'_>, + stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + // Check if already timed out - don't allow reuse + if matches!(self.send_body_state.write_state, WriteState::TimedOut) { + return Poll::Ready(Error::e_explain( + WriteTimedout, + "write task previously timed out", + )); + } + + // Lazy timeout optimization: Poll write first, create timeout only if needed. + // + // This follows the pattern from `pingora_timeout::Timeout` to avoid allocating + // and registering timeout futures when writes complete immediately (the common case). + // + // Fast path: Write completes → return immediately, no timeout future created + // Slow path: Write blocks → lazily create timeout future and poll both + + // First, try the write operation + // Dispatch to the appropriate body mode handler + let result = match self.body_mode { + BM::Complete(_) => Poll::Ready(Ok(None)), + BM::ContentLength(_, _) => self.poll_write_content_length_body_task(cx, stream), + BM::ChunkedEncoding(_) => self.poll_write_chunked_body_task(cx, stream), + BM::UntilClose(_) => self.poll_write_until_close_body_task(cx, stream), + BM::ToSelect => Poll::Ready(Ok(None)), + }; + + // If write completed immediately, return without ever creating/polling timeout + if result.is_ready() { + return result; + } + + // Write returned Pending - lazily create and check timeout if duration is set + if let Some(duration) = self.send_body_state.timeout_duration { + let timeout = self.send_body_state.timeout_fut.get_or_insert_with(|| { + Box::pin(pingora_timeout::sleep(duration)) + as std::pin::Pin + Send + Sync>> + }); + + if timeout.as_mut().poll(cx).is_ready() { + // Timeout fired! Mark state as timed out and clear the timeout future + self.send_body_state.write_state = WriteState::TimedOut; + self.send_body_state.timeout_fut = None; + return Poll::Ready(Error::e_explain( + WriteTimedout, + "writing body task timed out", + )); + } + } + + // Both write and timeout are pending + Poll::Pending + } + + // ======================================================================== + // Cancel-safe finish task API + // ======================================================================== + + #[cfg(test)] + pub fn has_pending_finish_task(&self) -> bool { + !matches!( + self.send_finish_state, + FinishWriteState::NotStarted | FinishWriteState::Done + ) + } + + /// Queue a finish operation as a task. + /// This is a non-async function that just marks the finish as pending. + /// Call `write_current_finish_task()` to actually perform the finish. + /// + /// This API is stateful and cancel-safe - use it when you need to finish + /// the body in a `tokio::select!` loop or other cancellable context. + pub fn send_finish_task(&mut self) { + self.send_finish_state = FinishWriteState::Idle; + } + + /// Async function that performs the current queued finish task on the stream. + /// This function is cancel-safe and can be called in a `tokio::select!` loop. + /// Returns `Ok(Some(bytes_written))` when complete, `Ok(None)` if already complete. + /// + /// This API is stateful - it tracks progress across cancellations and can be + /// safely resumed after being dropped mid-execution. + pub async fn write_current_finish_task(&mut self, stream: &mut S) -> Result> + where + S: AsyncWrite + Unpin + Send, + { + // Use poll_fn to wrap our poll-based implementation + std::future::poll_fn(|cx| self.poll_write_current_finish_task(cx, Pin::new(stream))).await + } + + /// Poll-based implementation for finish tasks. + /// This is the core implementation that maintains state across cancellations. + fn poll_write_current_finish_task( + &mut self, + cx: &mut Context<'_>, + stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + // If no finish queued, return None + if matches!( + self.send_finish_state, + FinishWriteState::NotStarted | FinishWriteState::Done + ) { + return Poll::Ready(Ok(None)); + } + + // Route to body-mode-specific implementation + match self.body_mode { + BM::Complete(_) => Poll::Ready(Ok(None)), + BM::ContentLength(_, _) => self.poll_finish_content_length_task(cx, stream), + BM::ChunkedEncoding(_) => self.poll_finish_chunked_task(cx, stream), + BM::UntilClose(_) => self.poll_finish_until_close_task(cx, stream), + BM::ToSelect => Poll::Ready(Ok(None)), + } + } + + /// Finish content-length body - just validates and updates state. + /// No I/O needed since body write tasks already flushed after the last write. + fn poll_finish_content_length_task( + &mut self, + _cx: &mut Context<'_>, + _stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + let written = match self.body_mode { + BM::ContentLength(total, w) => { + if w < total { + self.send_finish_state = FinishWriteState::Done; + return Poll::Ready(Error::e_explain( + PREMATURE_BODY_END, + format!("Content-length: {total} bytes written: {w}"), + )); + } + w + } + _ => panic!("wrong body mode: {:?}", self.body_mode), + }; + + // All bytes written - just update state to Complete + self.body_mode = BM::Complete(written); + self.send_finish_state = FinishWriteState::Done; + Poll::Ready(Ok(Some(written))) + } + + /// Poll-based helper to finish chunked encoding body + fn poll_finish_chunked_task( + &mut self, + cx: &mut Context<'_>, + mut stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + let written = match self.body_mode { + BM::ChunkedEncoding(w) => w, + _ => panic!("wrong body mode: {:?}", self.body_mode), + }; + + loop { + match &mut self.send_finish_state { + FinishWriteState::Idle => { + // Start writing last chunk marker "0\r\n\r\n" + let buf = WriteBuf::Simple(Bytes::from_static(&LAST_CHUNK[..])); + self.send_finish_state = FinishWriteState::WritingLastChunk(buf); + } + FinishWriteState::WritingLastChunk(buf) => { + // Re-create tokio's stateless `write_all_buf` future per + // poll; progress is carried by `buf`. + ready!(pin!(stream.as_mut().get_mut().write_all_buf(buf)).poll(cx)) + .map_err(|e| Error::because(WriteError, "while writing last chunk", e))?; + + // All bytes written, move to flushing state + self.send_finish_state = FinishWriteState::Flushing; + } + FinishWriteState::Flushing => { + // Poll flush + ready!(stream.as_mut().poll_flush(cx)) + .map_err(|e| Error::because(WriteError, "flushing after last chunk", e))?; + + // Flush complete! Update body_mode and mark done + self.body_mode = BM::Complete(written); + self.send_finish_state = FinishWriteState::Done; + return Poll::Ready(Ok(Some(written))); + } + FinishWriteState::Done => { + unreachable!( + "Done state should have been handled in poll_write_current_finish_task" + ) + } + FinishWriteState::NotStarted => { + unreachable!("NotStarted state should have been handled in poll_write_current_finish_task") + } + } + } + } + + /// Finish until-close body - just updates state. + /// No I/O needed since body write tasks already flushed after each write. + fn poll_finish_until_close_task( + &mut self, + _cx: &mut Context<'_>, + _stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + let written = match self.body_mode { + BM::UntilClose(w) => w, + _ => panic!("wrong body mode: {:?}", self.body_mode), + }; + + // Just update state to Complete + self.body_mode = BM::Complete(written); + self.send_finish_state = FinishWriteState::Done; + Poll::Ready(Ok(Some(written))) + } + + // ======================================================================== + // Internal helpers + // ======================================================================== + + /// Internal helper to poll a body task that writes in content-length mode + /// and flushes at end. + fn poll_write_content_length_body_task( + &mut self, + cx: &mut Context<'_>, + mut stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + // Move to Writing state if we're Idle + if matches!(self.send_body_state.write_state, WriteState::Idle) { + if let Some(mut bytes) = self.send_body_state.pending_bytes.take() { + let (total, written) = self.body_mode.expect_content_length(); + + // Check if we've already written everything + if written >= total { + self.send_body_state.write_state = WriteState::Done(0); + return Poll::Ready(Ok(None)); + } + + let original_size = bytes.len(); + let remaining = total - written; + + // Truncate bytes if they exceed content-length + if original_size > remaining { + warn!( + "Trying to write {} bytes over content-length: {}, truncating to {}", + original_size, total, remaining + ); + bytes.truncate(remaining); + } + + let bytes_to_write = bytes.len(); + self.send_body_state.write_state = + WriteState::Writing(bytes_to_write, WriteBuf::Simple(bytes)); + } else { + self.send_body_state.write_state = WriteState::Done(0); + return Poll::Ready(Ok(None)); + } + } + + // Handle Writing state - do the write, transition to Flushing or Done + if let WriteState::Writing(size, ref mut buf) = &mut self.send_body_state.write_state { + let bytes_written = *size; + + // Attempt write. Re-create tokio's stateless `write_all_buf` + // future per poll; progress is carried by `buf`. + match ready!(pin!(stream.as_mut().get_mut().write_all_buf(buf)).poll(cx)) { + Ok(()) => { + // Write completed - update body_mode to track bytes written + let (total, written) = self.body_mode.expect_content_length(); + self.body_mode = BM::ContentLength(total, written + bytes_written); + + if written + bytes_written >= total { + // All content-length bytes written, flush needed + self.send_body_state.write_state = WriteState::Flushing(bytes_written); + } else { + // More bytes to come, no flush needed + self.send_body_state.write_state = WriteState::Done(bytes_written); + } + } + Err(e) => { + return Poll::Ready(Error::e_because(WriteError, "while writing body", e)) + } + } + } + + // Handle Flushing state - do the flush, transition to Done + if let WriteState::Flushing(size) = self.send_body_state.write_state { + let bytes_written = size; + + // Attempt flush + match ready!(stream.poll_flush(cx)) { + Ok(()) => { + // Flush completed - transition to Done + self.send_body_state.write_state = WriteState::Done(bytes_written); + } + Err(e) => return Poll::Ready(Error::e_because(WriteError, "flushing body", e)), + } + } + + // Return based on final state + match self.send_body_state.write_state { + WriteState::Done(size) => { + self.send_body_state.timeout_fut = None; + Poll::Ready(Ok(Some(size))) + } + WriteState::TimedOut => Poll::Ready(Error::e_explain( + WriteTimedout, + "write task previously timed out", + )), + WriteState::Writing(..) | WriteState::Flushing(..) => { + unreachable!("Writing/Flushing states should have been handled above or returned Pending via ready!") + } + WriteState::Idle => { + unreachable!("Idle state should have been handled in setup") + } + } + } + + /// Poll-based implementation for chunked encoding mode + fn poll_write_chunked_body_task( + &mut self, + cx: &mut Context<'_>, + mut stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + // Move to Writing state if we're Idle + if matches!(self.send_body_state.write_state, WriteState::Idle) { + if let Some(bytes) = self.send_body_state.pending_bytes.take() { + let application_bytes_size = bytes.len(); + + // Format the chunk: size\r\ndata\r\n + let chunk_size_header = format!("{:X}\r\n", application_bytes_size); + let output_buf = Bytes::from(chunk_size_header) + .chain(bytes) + .chain(&b"\r\n"[..]); + + // Store the chained buffer directly to avoid copying + self.send_body_state.write_state = + WriteState::Writing(application_bytes_size, WriteBuf::Chained(output_buf)); + } else { + self.send_body_state.write_state = WriteState::Done(0); + return Poll::Ready(Ok(None)); + } + } + + // Handle Writing state - do the write using vectored I/O, transition to Flushing + if let WriteState::Writing(size, ref mut buf) = &mut self.send_body_state.write_state { + let bytes_written = *size; + + // Attempt vectored write for chained buffer (chunk size + data + CRLF). + // Re-create tokio's stateless `write_all_buf` future per poll; it + // opportunistically uses vectored writes when the stream supports + // them. Progress is carried by `buf`. + match ready!(pin!(stream.as_mut().get_mut().write_all_buf(buf)).poll(cx)) { + Ok(()) => { + // Write completed - update body_mode with application bytes (not wire bytes) + let written = self.body_mode.expect_chunked(); + self.body_mode = BM::ChunkedEncoding(written + bytes_written); + + // Chunked encoding always flushes + self.send_body_state.write_state = WriteState::Flushing(bytes_written); + } + Err(e) => { + return Poll::Ready(Error::e_because(WriteError, "while writing body", e)) + } + } + } + + // Handle Flushing state - do the flush, transition to Done + if let WriteState::Flushing(size) = self.send_body_state.write_state { + let bytes_written = size; + + // Attempt flush + match ready!(stream.poll_flush(cx)) { + Ok(()) => { + // Flush completed - transition to Done + self.send_body_state.write_state = WriteState::Done(bytes_written); + } + Err(e) => return Poll::Ready(Error::e_because(WriteError, "flushing body", e)), + } + } + + // Return based on final state + match self.send_body_state.write_state { + WriteState::Done(size) => { + self.send_body_state.timeout_fut = None; + Poll::Ready(Ok(Some(size))) + } + WriteState::TimedOut => Poll::Ready(Error::e_explain( + WriteTimedout, + "write task previously timed out", + )), + WriteState::Writing(..) | WriteState::Flushing(..) => { + unreachable!("Writing/Flushing states should have been handled above or returned Pending via ready!") + } + WriteState::Idle => { + unreachable!("Idle state should have been handled in setup") + } + } + } + + /// Poll-based implementation for UntilClose (close-delimited) body mode + fn poll_write_until_close_body_task( + &mut self, + cx: &mut Context<'_>, + mut stream: Pin<&mut S>, + ) -> Poll>> + where + S: AsyncWrite + Unpin + Send, + { + // Move to Writing state if we're Idle + if matches!(self.send_body_state.write_state, WriteState::Idle) { + if let Some(bytes) = self.send_body_state.pending_bytes.take() { + let original_size = bytes.len(); + self.send_body_state.write_state = + WriteState::Writing(original_size, WriteBuf::Simple(bytes)); + } else { + self.send_body_state.write_state = WriteState::Done(0); + return Poll::Ready(Ok(None)); + } + } + + // Handle Writing state - do the write, transition to Flushing + if let WriteState::Writing(size, ref mut buf) = &mut self.send_body_state.write_state { + let bytes_written = *size; + + // Attempt write. Re-create tokio's stateless `write_all_buf` + // future per poll; progress is carried by `buf`. + match ready!(pin!(stream.as_mut().get_mut().write_all_buf(buf)).poll(cx)) { + Ok(()) => { + // Write completed - update body_mode to track bytes written + let written = self.body_mode.expect_until_close(); + self.body_mode = BM::UntilClose(written + bytes_written); + + // Close-delimited mode always flushes + self.send_body_state.write_state = WriteState::Flushing(bytes_written); + } + Err(e) => { + return Poll::Ready(Error::e_because(WriteError, "while writing body", e)) + } + } + } + + // Handle Flushing state - do the flush, transition to Done + if let WriteState::Flushing(size) = self.send_body_state.write_state { + let bytes_written = size; + + // Attempt flush + match ready!(stream.poll_flush(cx)) { + Ok(()) => { + // Flush completed - transition to Done + self.send_body_state.write_state = WriteState::Done(bytes_written); + } + Err(e) => return Poll::Ready(Error::e_because(WriteError, "flushing body", e)), + } + } + + // Return based on final state + match self.send_body_state.write_state { + WriteState::Done(size) => { + self.send_body_state.timeout_fut = None; + Poll::Ready(Ok(Some(size))) + } + WriteState::TimedOut => Poll::Ready(Error::e_explain( + WriteTimedout, + "write task previously timed out", + )), + WriteState::Writing(..) | WriteState::Flushing(..) => { + unreachable!("Writing/Flushing states should have been handled above or returned Pending via ready!") + } + WriteState::Idle => { + unreachable!("Idle state should have been handled in setup") + } + } + } } #[cfg(test)] @@ -1717,25 +2477,33 @@ mod tests { let res = body_reader.read_body(&mut mock_io).await.unwrap().unwrap(); assert_eq!(res, BufRef::new(0, 0)); assert_eq!(body_reader.body_state, ParseState::Chunked(0, 0, 2, 2)); - let res = body_reader.read_body(&mut mock_io).await.unwrap().unwrap(); - assert_eq!(res, BufRef::new(3, 1)); // input1 concat input2 - assert_eq!(&input2[1..2], body_reader.get_body(&res)); - assert_eq!(body_reader.body_state, ParseState::Chunked(1, 6, 11, 0)); - let res = body_reader.read_body(&mut mock_io).await.unwrap(); - assert_eq!(res, None); - assert_eq!(body_reader.body_state, ParseState::Complete(1)); - assert_eq!(body_reader.get_body_overread(), None); + let _res = body_reader.read_body(&mut mock_io).await.unwrap().unwrap(); } #[tokio::test] - async fn read_with_body_partial_head_terminal_crlf() { + async fn read_with_body_partial_head_chunk_incomplete() { init_log(); let input1 = b"1\r"; - let input2 = b"\na\r\n0\r\n\r"; - let input3 = b"\n"; - let mut mock_io = Builder::new() - .read(&input1[..]) - .read(&input2[..]) + let mut mock_io = Builder::new().read(&input1[..]).build(); + let mut body_reader = BodyReader::new(false); + body_reader.init_chunked(b""); + let res = body_reader.read_body(&mut mock_io).await.unwrap().unwrap(); + assert_eq!(res, BufRef::new(0, 0)); + assert_eq!(body_reader.body_state, ParseState::Chunked(0, 0, 2, 2)); + let res = body_reader.read_body(&mut mock_io).await; + assert!(res.is_err()); + assert_eq!(body_reader.body_state, ParseState::Done(0)); + } + + #[tokio::test] + async fn read_with_body_partial_head_terminal_crlf() { + init_log(); + let input1 = b"1\r"; + let input2 = b"\na\r\n0\r\n\r"; + let input3 = b"\n"; + let mut mock_io = Builder::new() + .read(&input1[..]) + .read(&input2[..]) .read(&input3[..]) .build(); let mut body_reader = BodyReader::new(false); @@ -1925,21 +2693,6 @@ mod tests { assert_eq!(body_reader.get_body_overread(), Some(&b"abc"[..])); } - #[tokio::test] - async fn read_with_body_partial_head_chunk_incomplete() { - init_log(); - let input1 = b"1\r"; - let mut mock_io = Builder::new().read(&input1[..]).build(); - let mut body_reader = BodyReader::new(false); - body_reader.init_chunked(b""); - let res = body_reader.read_body(&mut mock_io).await.unwrap().unwrap(); - assert_eq!(res, BufRef::new(0, 0)); - assert_eq!(body_reader.body_state, ParseState::Chunked(0, 0, 2, 2)); - let res = body_reader.read_body(&mut mock_io).await; - assert!(res.is_err()); - assert_eq!(body_reader.body_state, ParseState::Done(0)); - } - #[tokio::test] async fn read_with_body_trailers() { init_log(); @@ -2319,7 +3072,7 @@ mod tests { } #[tokio::test] - async fn write_body_http10() { + async fn write_body_until_close() { init_log(); let data = b"a"; let mut mock_io = Builder::new().write(&data[..]).write(&data[..]).build(); @@ -2345,3 +3098,1019 @@ mod tests { assert_eq!(body_writer.body_mode, BodyMode::Complete(2)); } } + +#[cfg(test)] +mod test_body_task_api { + use super::*; + use crate::protocols::http::v1::test_util::FlushTrackingMock; + use tokio_test::io::Builder; + + // Cancel-safety tests use tokio::select! to race a short sleep against a mock + // I/O wait, simulating cancellation. We use #[tokio::test(start_paused = true)] + // on these tests so that tokio auto-advances time deterministically rather than + // relying on wall-clock timing. + + fn init_log() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + #[tokio::test] + async fn test_has_pending_body_task() { + init_log(); + let data = b"test data"; + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + // Initially should have no pending task + assert!(!body_writer.has_pending_body_task()); + + // After queuing bytes, should have pending task + body_writer.send_body_task(Bytes::from_static(data), None); + assert!(body_writer.has_pending_body_task()); + } + + #[tokio::test(start_paused = true)] + async fn cancel_safe_content_length_write() { + init_log(); + let data = b"Hello, World!"; + + // Create a mock stream that will block to allow cancellation + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(data) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + // Queue the bytes to write + body_writer.send_body_task(Bytes::from_static(data), None); + + // Use tokio::select! loop - keep looping until write completes + let mut cancel_count = 0; + let mut total_bytes_written = 0; + + loop { + // Break if no pending writes + if !body_writer.has_pending_body_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + // Timeout fires first, cancelling the write + cancel_count += 1; + } + result = body_writer.write_current_body_task(&mut mock_io) => { + // Write completed + assert!(result.is_ok(), "Write should succeed"); + if let Ok(Some(n)) = result { + total_bytes_written += n; + } + } + } + } + + assert!( + cancel_count > 0, + "At least one cancellation should have occurred" + ); + assert_eq!( + total_bytes_written, + data.len(), + "Should have written all application bytes" + ); + assert_eq!( + body_writer.body_mode, + BodyMode::ContentLength(data.len(), data.len()) + ); + + // Now test finish() in a select loop as well + let mut mock_io_finish = Builder::new().build(); + + loop { + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(5)) => { + // Allow cancellation attempts + } + result = body_writer.finish(&mut mock_io_finish) => { + assert!(result.is_ok()); + break; + } + } + } + + assert_eq!(body_writer.body_mode, BodyMode::Complete(data.len())); + } + + #[tokio::test(start_paused = true)] + async fn cancel_safe_chunked_write() { + init_log(); + let data = b"abcdefghij"; + let expected_output = b"A\r\nabcdefghij\r\n"; + + // Mock stream that blocks to allow cancellation + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(expected_output) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_chunked(); + + // Queue bytes + body_writer.send_body_task(Bytes::from_static(data), None); + + // Use select loop - keep looping until write completes + let mut cancel_count = 0; + let mut total_bytes_written = 0; + + loop { + if !body_writer.has_pending_body_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + cancel_count += 1; + } + result = body_writer.write_current_body_task(&mut mock_io) => { + assert!(result.is_ok()); + if let Ok(Some(n)) = result { + total_bytes_written += n; + } + } + } + } + + assert!(cancel_count > 0, "Should have cancelled at least once"); + assert_eq!( + total_bytes_written, + data.len(), + "Should have written all application bytes" + ); + assert_eq!(body_writer.body_mode, BodyMode::ChunkedEncoding(data.len())); + + // Test finish() with select loop - must write terminating chunk + let mut mock_io_finish = Builder::new() + .wait(std::time::Duration::from_millis(50)) + .write(&LAST_CHUNK[..]) // Expect 0\r\n\r\n + .build(); + + loop { + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(5)) => {} + result = body_writer.finish(&mut mock_io_finish) => { + assert!(result.is_ok()); + break; + } + } + } + + assert_eq!(body_writer.body_mode, BodyMode::Complete(data.len())); + } + + #[tokio::test(start_paused = true)] + async fn cancel_safe_until_close_write() { + init_log(); + let data = b"test data"; + + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(data) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_close_delimited(); + + body_writer.send_body_task(Bytes::from_static(data), None); + + // Use select loop - keep looping until write completes + let mut cancel_count = 0; + let mut total_bytes_written = 0; + + loop { + if !body_writer.has_pending_body_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + cancel_count += 1; + } + result = body_writer.write_current_body_task(&mut mock_io) => { + assert!(result.is_ok()); + if let Ok(Some(n)) = result { + total_bytes_written += n; + } + } + } + } + + assert!(cancel_count > 0); + assert_eq!( + total_bytes_written, + data.len(), + "Should have written all application bytes" + ); + + // Test finish() with select loop + let mut mock_io_finish = Builder::new().build(); + + loop { + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(5)) => {} + result = body_writer.finish(&mut mock_io_finish) => { + assert!(result.is_ok()); + break; + } + } + } + } + + #[tokio::test(start_paused = true)] + async fn cancel_safe_multiple_cancellations() { + init_log(); + let data = b"Long test data that requires multiple writes"; + + // Create a mock that blocks multiple times + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(50)) + .write(&data[..15]) + .wait(std::time::Duration::from_millis(50)) + .write(&data[15..30]) + .wait(std::time::Duration::from_millis(50)) + .write(&data[30..]) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + body_writer.send_body_task(Bytes::from_static(data), None); + + // Loop until write completes, allowing cancellations + let mut cancel_count = 0; + let mut total_bytes_written = 0; + + loop { + if !body_writer.has_pending_body_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(5)) => { + cancel_count += 1; + } + result = body_writer.write_current_body_task(&mut mock_io) => { + assert!(result.is_ok()); + if let Ok(Some(n)) = result { + total_bytes_written += n; + } + } + } + } + + assert!(cancel_count >= 2, "Should have multiple cancellations"); + assert_eq!( + total_bytes_written, + data.len(), + "Should have written all application bytes" + ); + + // Test finish with select loop + let mut mock_io_finish = Builder::new().build(); + + loop { + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(5)) => {} + result = body_writer.finish(&mut mock_io_finish) => { + assert!(result.is_ok()); + break; + } + } + } + } + + #[tokio::test(start_paused = true)] + async fn cancel_safe_partial_writes() { + init_log(); + let data = b"12345678901234567890"; // 20 bytes + + // Simulate partial writes with blocking + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(50)) + .write(&data[..7]) + .wait(std::time::Duration::from_millis(50)) + .write(&data[7..14]) + .wait(std::time::Duration::from_millis(50)) + .write(&data[14..]) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + body_writer.send_body_task(Bytes::from_static(data), None); + + let mut cancel_count = 0; + let mut total_bytes_written = 0; + + loop { + if !body_writer.has_pending_body_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + cancel_count += 1; + } + result = body_writer.write_current_body_task(&mut mock_io) => { + assert!(result.is_ok()); + if let Ok(Some(n)) = result { + total_bytes_written += n; + } + } + } + } + + assert!(cancel_count > 0); + assert_eq!( + total_bytes_written, + data.len(), + "Should have written all application bytes" + ); + + // Test finish in select loop + let mut mock_io_finish = Builder::new() + .wait(std::time::Duration::from_millis(30)) + .build(); + + loop { + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(5)) => {} + result = body_writer.finish(&mut mock_io_finish) => { + assert!(result.is_ok(), "Finish should succeed after cancel-safe writes"); + break; + } + } + } + } + + #[tokio::test] + async fn test_task_write_timeout() { + init_log(); + let data = b"test data"; + + // Create a mock that blocks forever + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_secs(1000)) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + // Queue the task with a timeout + body_writer.send_body_task( + Bytes::from_static(data), + Some(std::time::Duration::from_millis(50)), + ); + + // The write should timeout + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_err(), "Write should timeout"); + + // Check that it's a timeout error + if let Err(e) = result { + assert_eq!(e.etype(), &WriteTimedout); + } + } + + // Even if the user's select! cancels the write, the internal timeout + // should continue counting across cancellations. + #[tokio::test] + async fn test_task_timeout_persists_across_cancellations() { + init_log(); + let data = b"test data"; + + // Create a mock that blocks for a while + // Since timeout is 100ms and this waits 200ms, the write should never happen + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(200)) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + // Queue the task with a 100ms timeout + body_writer.send_body_task( + Bytes::from_static(data), + Some(std::time::Duration::from_millis(100)), + ); + + let mut attempts = 0; + let mut timedout = false; + + // Try to write in a loop, but cancel early each time + // The timeout should still fire even though we're cancelling + loop { + if !body_writer.has_pending_body_task() { + break; + } + + attempts += 1; + + tokio::select! { + // Cancel after just 10ms each time + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + // Cancelled by our select, continue looping + continue; + } + result = body_writer.write_current_body_task(&mut mock_io) => { + match result { + Ok(_) => { + // Write succeeded before timeout + break; + } + Err(e) if e.etype() == &WriteTimedout => { + // Timeout fired! + timedout = true; + break; + } + Err(e) => { + panic!("Unexpected error: {:?}", e); + } + } + } + } + } + + assert!(timedout, "Timeout should have fired despite cancellations"); + assert!( + attempts >= 5, + "Should have had multiple attempts before timeout" + ); + } + + #[tokio::test] + async fn test_task_write_succeeds_within_timeout() { + init_log(); + let data = b"Hello, World!"; + + // Create a mock that completes quickly + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(20)) + .write(data) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + // Queue with a generous timeout + body_writer.send_body_task( + Bytes::from_static(data), + Some(std::time::Duration::from_millis(500)), + ); + + // Write should succeed + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_ok(), "Write should succeed: {:?}", result); + assert_eq!(result.unwrap(), Some(data.len())); + } + + #[tokio::test] + async fn test_task_write_no_timeout() { + init_log(); + let data = b"test data"; + + // Create a mock that takes a bit of time but eventually succeeds + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(data) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + // Queue without timeout + body_writer.send_body_task(Bytes::from_static(data), None); + + // Write should eventually succeed + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_ok(), "Write should succeed without timeout"); + assert_eq!(result.unwrap(), Some(data.len())); + } + + #[tokio::test] + async fn test_task_chunked_write_timeout() { + init_log(); + let data = b"chunked data"; + + // Create a mock that blocks + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_secs(1000)) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_chunked(); + + // Queue with short timeout + body_writer.send_body_task( + Bytes::from_static(data), + Some(std::time::Duration::from_millis(50)), + ); + + // Should timeout + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_err()); + if let Err(e) = result { + assert_eq!(e.etype(), &WriteTimedout); + } + } + + #[tokio::test] + async fn test_task_timeout_reset_on_new_task() { + init_log(); + let data1 = b"first"; + let data2 = b"second"; + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data1.len() + data2.len()); + + // Queue first task with short timeout + body_writer.send_body_task( + Bytes::from_static(data1), + Some(std::time::Duration::from_millis(50)), + ); + + // Wait a bit but don't let it timeout yet + tokio::time::sleep(std::time::Duration::from_millis(30)).await; + + // Queue a new task with a longer timeout + // This should reset/replace the timeout + body_writer.send_body_task( + Bytes::from_static(data2), + Some(std::time::Duration::from_millis(500)), + ); + + // Create a mock that takes some time + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(data2) + .build(); + + // The second write should succeed with its own timeout + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!( + result.is_ok(), + "Second task should succeed with new timeout" + ); + } + + #[tokio::test] + async fn test_task_timeout_with_partial_writes() { + init_log(); + let data1 = b"first"; + let data2 = b"second"; + let data3 = b"third"; + + // Mock that writes data1 quickly, data2 with delay, data3 blocks forever + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(10)) + .write(data1) + .wait(std::time::Duration::from_millis(40)) + .write(data2) + .wait(std::time::Duration::from_secs(1000)) + .build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data1.len() + data2.len() + data3.len()); + + let mut total_written = 0; + + // First write - should succeed within timeout + body_writer.send_body_task( + Bytes::from_static(data1), + Some(std::time::Duration::from_millis(100)), + ); + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_ok()); + total_written += result.unwrap().unwrap(); + + // Second write - should succeed within timeout + body_writer.send_body_task( + Bytes::from_static(data2), + Some(std::time::Duration::from_millis(100)), + ); + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_ok()); + total_written += result.unwrap().unwrap(); + + // Third write - should timeout + body_writer.send_body_task( + Bytes::from_static(data3), + Some(std::time::Duration::from_millis(50)), + ); + let result = body_writer.write_current_body_task(&mut mock_io).await; + assert!(result.is_err()); + assert_eq!(result.unwrap_err().etype(), &WriteTimedout); + + // We should have written data1 and data2 but not data3 + assert_eq!(total_written, data1.len() + data2.len()); + assert!( + total_written < data1.len() + data2.len() + data3.len(), + "Should not have written all data" + ); + } + + // Cancel-safe finish task for chunked encoding: send_finish_task() queues + // the terminating chunk, write_current_finish_task() writes it and can be + // cancelled and resumed in a select! loop. + // Verifies that the finish flushes the stream exactly once. + #[tokio::test(start_paused = true)] + async fn cancel_safe_finish_task_chunked() { + init_log(); + + let data = Bytes::from("hello"); + let expected_chunk = b"5\r\nhello\r\n"; + + let mock_io = Builder::new().write(expected_chunk).build(); + let (mut flush_mock, flush_count) = FlushTrackingMock::new(mock_io); + + let mut body_writer = BodyWriter::new(); + body_writer.init_chunked(); + + // Write body data via task API + body_writer.send_body_task(data, None); + body_writer + .write_current_body_task(&mut flush_mock) + .await + .unwrap(); + + // Chunked body writes always flush after each chunk + assert_eq!( + FlushTrackingMock::flush_count(&flush_count), + 1, + "Chunked body data write should flush once" + ); + + // Queue the finish task + body_writer.send_finish_task(); + assert!(body_writer.has_pending_finish_task()); + + // Write the finish in a select! loop with cancellations + let mock_io_finish = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(b"0\r\n\r\n") + .build(); + let (mut flush_mock_finish, flush_count_finish) = FlushTrackingMock::new(mock_io_finish); + + let mut cancel_count = 0; + + loop { + if !body_writer.has_pending_finish_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + cancel_count += 1; + } + result = body_writer.write_current_finish_task(&mut flush_mock_finish) => { + assert!(result.is_ok()); + break; + } + } + } + + assert!(cancel_count > 0, "Should have cancelled at least once"); + assert!(matches!(body_writer.body_mode, BodyMode::Complete(_))); + assert_eq!( + FlushTrackingMock::flush_count(&flush_count_finish), + 1, + "Chunked finish should flush exactly once" + ); + } + + // Finish task for content-length is a no-op (no terminating chunk needed), + // but it should still transition body_mode to Complete. + // Verifies that no flush occurs (content-length finish has no I/O). + #[tokio::test] + async fn finish_task_content_length() { + init_log(); + + let data = b"hello"; + let mock_io = Builder::new().write(data).build(); + let (mut flush_mock, flush_count) = FlushTrackingMock::new(mock_io); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(data.len()); + + body_writer.send_body_task(Bytes::from_static(data), None); + body_writer + .write_current_body_task(&mut flush_mock) + .await + .unwrap(); + + // Content-length body write flushes when all bytes are written + assert_eq!( + FlushTrackingMock::flush_count(&flush_count), + 1, + "Content-length body write should flush once (all bytes written)" + ); + + body_writer.send_finish_task(); + let mock_io_finish = Builder::new().build(); + let (mut flush_mock_finish, flush_count_finish) = FlushTrackingMock::new(mock_io_finish); + let result = body_writer + .write_current_finish_task(&mut flush_mock_finish) + .await; + assert!(result.is_ok()); + assert!(matches!(body_writer.body_mode, BodyMode::Complete(_))); + assert_eq!( + FlushTrackingMock::flush_count(&flush_count_finish), + 0, + "Content-length finish should not flush (no I/O needed)" + ); + } + + // Verifies that body_mode byte tracking is correct when writing + // content-length body in multiple chunks. Each intermediate chunk + // does not trigger a flush; the body_mode must still accumulate + // bytes correctly so that finish_task succeeds. + #[tokio::test] + async fn content_length_body_mode_tracks_across_chunks() { + init_log(); + + let chunk1 = b"Hello"; + let chunk2 = b", World!"; + let total_len = chunk1.len() + chunk2.len(); // 13 + + // Mock expects both writes; the final write triggers a flush internally + let mut mock_io = Builder::new().write(chunk1).write(chunk2).build(); + + let mut body_writer = BodyWriter::new(); + body_writer.init_content_length(total_len); + + // Write first chunk (intermediate, no flush expected) + body_writer.send_body_task(Bytes::from_static(chunk1), None); + let result = body_writer + .write_current_body_task(&mut mock_io) + .await + .unwrap(); + assert_eq!(result, Some(chunk1.len())); + assert!( + !body_writer.finished(), + "Should not be finished after first chunk" + ); + + // Verify body_mode tracks the bytes from the first chunk + assert!( + matches!(body_writer.body_mode, BodyMode::ContentLength(total, written) + if total == total_len && written == chunk1.len()), + "body_mode should reflect bytes written so far, got: {:?}", + body_writer.body_mode + ); + + // Write second chunk (final, completes content-length) + body_writer.send_body_task(Bytes::from_static(chunk2), None); + let result = body_writer + .write_current_body_task(&mut mock_io) + .await + .unwrap(); + assert_eq!(result, Some(chunk2.len())); + assert!( + body_writer.finished(), + "Should be finished after all bytes written" + ); + + // Finish should succeed since all content-length bytes were written + body_writer.send_finish_task(); + let mut mock_io_finish = Builder::new().build(); + let result = body_writer + .write_current_finish_task(&mut mock_io_finish) + .await; + assert!( + result.is_ok(), + "finish_task should succeed when all content-length bytes written" + ); + assert!(matches!(body_writer.body_mode, BodyMode::Complete(_))); + } +} + +#[cfg(test)] +mod test_poll_body_writer { + //! Tests that drive the cancel-safe poll-based body writers directly + //! via [`futures::task::noop_waker`] and an in-memory [`AsyncWrite`] + //! mock — no tokio runtime required. + //! + //! These exercise the `write_all_buf` call sites in + //! [`BodyWriter::poll_write_current_body_task`] and + //! [`BodyWriter::poll_write_current_finish_task`]: re-creating tokio's + //! stateless future on each poll and verifying that progress carried + //! by the caller-owned [`Buf`](bytes::Buf) is preserved across + //! [`Poll::Pending`] returns. + use super::*; + use bytes::Bytes; + use futures::task::noop_waker; + use std::collections::VecDeque; + use std::io; + use std::pin::Pin; + use std::task::{Context, Poll}; + use tokio::io::AsyncWrite; + + /// Programmable in-memory [`AsyncWrite`]. Each scripted entry dictates + /// the outcome of one `poll_write` / `poll_write_vectored` call. + struct MockWriter { + written: Vec, + next: VecDeque>>, + vectored: bool, + } + + impl MockWriter { + fn new(vectored: bool, results: Vec>>) -> Self { + Self { + written: Vec::new(), + next: results.into(), + vectored, + } + } + } + + impl AsyncWrite for MockWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let r = self + .next + .pop_front() + .expect("MockWriter::poll_write: out of scripted results"); + if let Poll::Ready(Ok(n)) = &r { + self.written.extend_from_slice(&buf[..*n]); + } + r + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + let r = self + .next + .pop_front() + .expect("MockWriter::poll_write_vectored: out of scripted results"); + if let Poll::Ready(Ok(n)) = &r { + let mut remaining = *n; + for slice in bufs { + if remaining == 0 { + break; + } + let take = slice.len().min(remaining); + self.written.extend_from_slice(&slice[..take]); + remaining -= take; + } + } + r + } + + fn is_write_vectored(&self) -> bool { + self.vectored + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + /// Drive a poll closure to completion under a [`noop_waker`] context. + /// + /// Panics after [`MAX_POLLS`] iterations to avoid silently hanging the + /// test runner when a mock is misconfigured and never returns + /// [`Poll::Ready`]. + fn drive(mut poll_fn: impl FnMut(&mut Context<'_>) -> Poll) -> R { + const MAX_POLLS: usize = 1024; + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + for _ in 0..MAX_POLLS { + if let Poll::Ready(v) = poll_fn(&mut cx) { + return v; + } + } + panic!("drive(): poll closure did not complete after {MAX_POLLS} iterations"); + } + + #[test] + fn content_length_write_completes() { + let data = b"hello"; + let mut mock = MockWriter::new(false, vec![Poll::Ready(Ok(5))]); + let mut bw = BodyWriter::new(); + bw.init_content_length(data.len()); + bw.send_body_task(Bytes::from_static(data), None); + + let res = drive(|cx| bw.poll_write_current_body_task(cx, Pin::new(&mut mock))); + assert!(matches!(res, Ok(Some(5)))); + assert_eq!(mock.written, data); + assert!(matches!(bw.body_mode, BodyMode::ContentLength(5, 5))); + } + + #[test] + fn content_length_write_resumes_after_pending_and_short_write() { + // Pending once, then 2-byte short write, then the remaining 3 bytes. + // Exercises the per-poll re-creation of tokio's WriteAllBuf future + // with progress carried by the BodyWriter-owned buffer. + let data = b"hello"; + let mut mock = MockWriter::new( + false, + vec![Poll::Pending, Poll::Ready(Ok(2)), Poll::Ready(Ok(3))], + ); + let mut bw = BodyWriter::new(); + bw.init_content_length(data.len()); + bw.send_body_task(Bytes::from_static(data), None); + + let res = drive(|cx| bw.poll_write_current_body_task(cx, Pin::new(&mut mock))); + assert!(matches!(res, Ok(Some(5)))); + assert_eq!(mock.written, data); + } + + #[test] + fn chunked_write_uses_vectored_path() { + // Chunked body emits "5\r\nhello\r\n" = 10 bytes as a 3-chunk + // WriteBuf::Chained. Note: the WriteBuf type uses the default + // Buf::chunks_vectored impl, which only ever populates one IoSlice, + // so tokio's write_all_buf issues 3 separate single-slice + // poll_write_vectored calls (one per inner Bytes chunk). + let payload = b"hello"; + let mut mock = MockWriter::new( + true, + vec![ + Poll::Ready(Ok(3)), // "5\r\n" + Poll::Ready(Ok(5)), // "hello" + Poll::Ready(Ok(2)), // "\r\n" + ], + ); + let mut bw = BodyWriter::new(); + bw.init_chunked(); + bw.send_body_task(Bytes::from_static(payload), None); + + let res = drive(|cx| bw.poll_write_current_body_task(cx, Pin::new(&mut mock))); + assert!(matches!(res, Ok(Some(5)))); // application bytes + assert_eq!(mock.written, b"5\r\nhello\r\n"); + } + + #[test] + fn chunked_write_resumes_on_short_write() { + // Exercise the restart path: short-write on the first chunk + // (Ok(2) for "5\r\n"), then completion (Ok(1) for "\n"), then the + // remaining two inner chunks. + let payload = b"hello"; + let mut mock = MockWriter::new( + true, + vec![ + Poll::Ready(Ok(2)), // partial "5\r" + Poll::Ready(Ok(1)), // remaining "\n" + Poll::Ready(Ok(5)), // "hello" + Poll::Ready(Ok(2)), // "\r\n" + ], + ); + let mut bw = BodyWriter::new(); + bw.init_chunked(); + bw.send_body_task(Bytes::from_static(payload), None); + + let res = drive(|cx| bw.poll_write_current_body_task(cx, Pin::new(&mut mock))); + assert!(matches!(res, Ok(Some(5)))); + assert_eq!(mock.written, b"5\r\nhello\r\n"); + } + + #[test] + fn until_close_write_completes() { + let data = b"abcdef"; + let mut mock = MockWriter::new(false, vec![Poll::Ready(Ok(6))]); + let mut bw = BodyWriter::new(); + bw.init_close_delimited(); + bw.send_body_task(Bytes::from_static(data), None); + + let res = drive(|cx| bw.poll_write_current_body_task(cx, Pin::new(&mut mock))); + assert!(matches!(res, Ok(Some(6)))); + assert_eq!(mock.written, data); + } + + #[test] + fn finish_chunked_writes_terminator() { + // After a chunked body, finish() must emit the "0\r\n\r\n" terminator. + let mut mock = MockWriter::new(true, vec![Poll::Ready(Ok(5))]); + let mut bw = BodyWriter::new(); + bw.init_chunked(); + // Pretend we already wrote 5 application bytes. + bw.body_mode = BodyMode::ChunkedEncoding(5); + bw.send_finish_task(); + + let res = drive(|cx| bw.poll_write_current_finish_task(cx, Pin::new(&mut mock))); + assert!(res.is_ok()); + assert_eq!(mock.written, b"0\r\n\r\n"); + assert!(matches!(bw.body_mode, BodyMode::Complete(5))); + } +} diff --git a/pingora-core/src/protocols/http/v1/client.rs b/pingora-core/src/protocols/http/v1/client.rs index 5f9e46107..a60aad1f3 100644 --- a/pingora-core/src/protocols/http/v1/client.rs +++ b/pingora-core/src/protocols/http/v1/client.rs @@ -625,13 +625,6 @@ impl HttpSession { // follow https://datatracker.ietf.org/doc/html/rfc9112#section-6.3 let preread_body = self.preread_body.as_ref().unwrap().get(&self.buf[..]); - if let Some(req) = self.request_written.as_ref() { - if req.method == http::method::Method::HEAD { - self.body_reader.init_content_length(0, preread_body); - return; - } - } - let upgraded = if let Some(code) = self.get_status() { match code.as_u16() { 101 => self.is_upgrade_req(), @@ -650,6 +643,13 @@ impl HttpSession { false }; + if let Some(req) = self.request_written.as_ref() { + if req.method == http::method::Method::HEAD { + self.body_reader.init_content_length(0, preread_body); + return; + } + } + if upgraded { self.body_reader.init_close_delimited(preread_body); self.close_delimited_resp = true; @@ -2224,6 +2224,283 @@ hello"; http_stream.respect_keepalive(); assert!(!http_stream.will_keepalive()); } + + #[tokio::test] + async fn read_informational_head_request() { + init_log(); + // HEAD request that receives 100 Continue followed by 200 OK + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input1 = b"HTTP/1.1 100 Continue\r\n\r\n"; + let input2 = b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n"; + + let mock_io = Builder::new() + .write(&wire[..]) + .read(&input1[..]) + .read(&input2[..]) + .build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + // Write HEAD request + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + // Read 100 Continue + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 100); + assert!(!eob, "100 Continue for HEAD should not signal end of body"); + } + _ => { + panic!("task should be informational header") + } + } + + // Read final 200 OK + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 200); + assert!(eob, "HEAD 200 response should signal end of body"); + } + _ => { + panic!("task should be final header") + } + } + + // Body reader should be Complete(0) for HEAD + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + } + + #[tokio::test] + async fn read_informational_multiple_head_request() { + init_log(); + // HEAD request that receives 100 Continue, 103 Early Hints, then 200 OK + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input = b"HTTP/1.1 100 Continue\r\n\r\nHTTP/1.1 103 Early Hints\r\n\r\nHTTP/1.1 200 OK\r\nContent-Length: 50\r\n\r\n"; + + let mock_io = Builder::new().write(&wire[..]).read(&input[..]).build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + // Read 100 Continue + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 100); + assert!(!eob, "100 Continue for HEAD should not signal end of body"); + } + _ => { + panic!("task should be 100 header") + } + } + + // Read 103 Early Hints + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 103); + assert!( + !eob, + "103 Early Hints for HEAD should not signal end of body" + ); + } + _ => { + panic!("task should be 103 header") + } + } + + // Read 200 OK — end of body + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 200); + assert!(eob, "HEAD 200 response should signal end of body"); + } + _ => { + panic!("task should be final header") + } + } + + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + } + + #[tokio::test] + async fn read_basic_head() { + init_log(); + // Basic HEAD + 200 + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input = b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n"; + + let mock_io = Builder::new().write(&wire[..]).read(&input[..]).build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 200); + assert!(eob, "HEAD 200 should be end of body"); + } + _ => { + panic!("task should be header") + } + } + + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + + // Keepalive should work for a properly-framed HEAD response + http_stream.respect_keepalive(); + assert!(http_stream.will_keepalive()); + } + + #[tokio::test] + async fn read_head_informational_keepalive() { + init_log(); + // HEAD + 100 Continue + 200 OK, then verify keepalive is preserved. + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input1 = b"HTTP/1.1 100 Continue\r\n\r\n"; + let input2 = b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n"; + + let mock_io = Builder::new() + .write(&wire[..]) + .read(&input1[..]) + .read(&input2[..]) + .build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + // 100 Continue + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 100); + assert!(!eob); + } + _ => panic!("task should be informational header"), + } + + // 200 OK + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 200); + assert!(eob); + } + _ => panic!("task should be final header"), + } + + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + + // Keepalive must still work after the 100 + 200 sequence + http_stream.respect_keepalive(); + assert!(http_stream.will_keepalive()); + } + + #[tokio::test] + async fn read_head_204() { + init_log(); + // HEAD + 204 No Content + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input = b"HTTP/1.1 204 No Content\r\n\r\n"; + + let mock_io = Builder::new().write(&wire[..]).read(&input[..]).build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 204); + assert!(eob, "HEAD 204 should be end of body"); + } + _ => panic!("task should be header"), + } + + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + } + + #[tokio::test] + async fn read_head_304() { + init_log(); + // HEAD + 304 Not Modified + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input = b"HTTP/1.1 304 Not Modified\r\nContent-Length: 100\r\n\r\n"; + + let mock_io = Builder::new().write(&wire[..]).read(&input[..]).build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 304); + assert!(eob, "HEAD 304 should be end of body"); + } + _ => panic!("task should be header"), + } + + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + } + + #[tokio::test] + async fn read_head_101_non_upgrade() { + init_log(); + // HEAD + 101 where the request is not an upgrade request. + // Contrived, but verifies the new code path: 101 check fires first, + // is_upgrade_req() returns false, then HEAD check fires. + let wire = b"HEAD / HTTP/1.1\r\n\r\n"; + let input = b"HTTP/1.1 101 Switching Protocols\r\n\r\n"; + + let mock_io = Builder::new().write(&wire[..]).read(&input[..]).build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + + let new_request = RequestHeader::build("HEAD", b"/", None).unwrap(); + http_stream + .write_request_header(Box::new(new_request)) + .await + .unwrap(); + + let task = http_stream.read_response_task().await.unwrap(); + match task { + HttpTask::Header(h, eob) => { + assert_eq!(h.status, 101); + // HEAD without Upgrade headers → not an upgrade, body is "done" + assert!(eob, "HEAD 101 (non-upgrade) should be end of body"); + } + _ => panic!("task should be header"), + } + + assert_eq!(http_stream.body_reader.body_state, ParseState::Complete(0)); + } } #[cfg(test)] diff --git a/pingora-core/src/protocols/http/v1/header.rs b/pingora-core/src/protocols/http/v1/header.rs new file mode 100644 index 000000000..9daaad328 --- /dev/null +++ b/pingora-core/src/protocols/http/v1/header.rs @@ -0,0 +1,460 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Cancel-safe header writing for HTTP/1.x + +use bytes::Bytes; +use pingora_error::{Error, ErrorType::*, Result}; +use std::future::Future; +use std::pin::{pin, Pin}; +use std::task::{ready, Context, Poll}; +use tokio::io::{AsyncWrite, AsyncWriteExt}; + +enum HeaderWriteState { + /// No write in progress + Idle, + /// Writing header bytes (original size, buffer) + Writing(usize, Bytes), + /// Flushing after write (original size to return) + Flushing(usize), + /// Write complete + Done, + /// Write timed out - cannot be reused + TimedOut, +} + +// Custom Debug implementation +impl std::fmt::Debug for HeaderWriteState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HeaderWriteState::Idle => write!(f, "Idle"), + HeaderWriteState::Writing(size, _) => write!(f, "Writing(size: {})", size), + HeaderWriteState::Flushing(size) => write!(f, "Flushing(size: {})", size), + HeaderWriteState::Done => write!(f, "Done"), + HeaderWriteState::TimedOut => write!(f, "TimedOut"), + } + } +} + +/// Internal state for the cancel-safe header write state machine. +/// +/// Tracks the pending header bytes, write progress (idle → writing → flushing → done), +/// and an optional timeout that is lazily created on the first `Pending` poll. +struct SendHeaderState { + /// Serialized header bytes ready to be written + pending_header: Option, + /// Whether to flush after writing + should_flush: bool, + /// Current write state + write_state: HeaderWriteState, + /// Timeout duration for this write task + timeout_duration: Option, + /// Timeout future (only created if write returns Pending) + timeout_fut: Option + Send + Sync>>>, +} + +// Custom Debug implementation since timeout_fut doesn't implement Debug +impl std::fmt::Debug for SendHeaderState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SendHeaderState") + .field("pending_header", &self.pending_header) + .field("should_flush", &self.should_flush) + .field("write_state", &self.write_state) + .field("timeout_duration", &self.timeout_duration) + .field( + "timeout_fut", + &self.timeout_fut.as_ref().map(|_| "Some(Future)"), + ) + .finish() + } +} + +impl SendHeaderState { + fn new() -> Self { + SendHeaderState { + pending_header: None, + should_flush: false, + write_state: HeaderWriteState::Idle, + timeout_duration: None, + timeout_fut: None, + } + } +} + +/// Cancel-safe header writer for HTTP/1.x response headers. +/// +/// This writer allows response headers to be written to a downstream connection +/// inside a `tokio::select!` loop without losing progress. If the write is +/// cancelled (e.g. because another branch of the select fires first), the +/// partially-written state is preserved and will be resumed on the next call to +/// [`write_current_header_task`](Self::write_current_header_task). +/// +/// ## Usage +/// +/// 1. Call [`send_header_task`](Self::send_header_task) with pre-serialized +/// header bytes, a flush flag, and an optional timeout. +/// 2. Await [`write_current_header_task`](Self::write_current_header_task) +/// (possibly inside `tokio::select!`). The method returns `Ok(bytes_written)` +/// on success. +/// +/// A timeout, if set, is enforced *across* cancellations — the clock keeps +/// ticking even when the future is dropped and re-polled. +pub struct HeaderWriter { + // Boxed to reduce inline size. Only used by the cancel-safe proxy task API. + send_header_state: Box, +} + +impl Default for HeaderWriter { + fn default() -> Self { + Self::new() + } +} + +impl HeaderWriter { + pub fn new() -> Self { + HeaderWriter { + send_header_state: Box::new(SendHeaderState::new()), + } + } + + #[cfg(test)] + pub fn has_pending_header_task(&self) -> bool { + self.send_header_state.pending_header.is_some() + || !matches!( + self.send_header_state.write_state, + HeaderWriteState::Idle | HeaderWriteState::Done | HeaderWriteState::TimedOut + ) + } + + /// Queue serialized header bytes as a write task with an optional timeout. + /// This is a non-async function that just saves the bytes. + /// Call [`write_current_header_task`](Self::write_current_header_task) to actually perform the write. + pub fn send_header_task( + &mut self, + header_bytes: Bytes, + should_flush: bool, + timeout: Option, + ) { + assert!( + matches!( + self.send_header_state.write_state, + HeaderWriteState::Idle | HeaderWriteState::Done + ), + "send_header_task called while previous task is still in progress: {:?}", + self.send_header_state.write_state + ); + self.send_header_state.pending_header = Some(header_bytes); + self.send_header_state.should_flush = should_flush; + self.send_header_state.write_state = HeaderWriteState::Idle; + self.send_header_state.timeout_duration = timeout; + self.send_header_state.timeout_fut = None; + } + + /// Async function that writes the current queued header task to the stream. + /// This function is cancel-safe and can be called in a `tokio::select!` loop. + /// Returns `Ok(bytes_written)` when complete, `Ok(0)` if no bytes to write. + pub async fn write_current_header_task(&mut self, stream: &mut S) -> Result + where + S: AsyncWrite + Unpin + Send, + { + std::future::poll_fn(|cx| self.poll_write_current_header_task(cx, Pin::new(stream))).await + } + + /// Poll-based implementation for writing the current header task. + fn poll_write_current_header_task( + &mut self, + cx: &mut Context<'_>, + stream: Pin<&mut S>, + ) -> Poll> + where + S: AsyncWrite + Unpin + Send, + { + // Check if already timed out - don't allow reuse + if matches!( + self.send_header_state.write_state, + HeaderWriteState::TimedOut + ) { + return Poll::Ready(Error::e_explain( + WriteTimedout, + "header write task previously timed out", + )); + } + + // First, try the write operation + match self.poll_do_write_header_and_flush(cx, stream) { + Poll::Ready(Ok(size)) => { + // Write completed! Clear timeout and return + if matches!(self.send_header_state.write_state, HeaderWriteState::Done) { + self.send_header_state.timeout_fut = None; + } + return Poll::Ready(Ok(size)); + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => { + // Write is pending - now check timeout + } + } + + // Lazy timeout optimization: Polls write first, creates timeout only if needed. + // This follows the pattern from `pingora_timeout::Timeout` to avoid allocating + // timeout futures when writes complete immediately (the common case). + if let Some(duration) = self.send_header_state.timeout_duration { + let timeout = self.send_header_state.timeout_fut.get_or_insert_with(|| { + Box::pin(pingora_timeout::sleep(duration)) + as std::pin::Pin + Send + Sync>> + }); + + if timeout.as_mut().poll(cx).is_ready() { + // Timeout fired! + self.send_header_state.write_state = HeaderWriteState::TimedOut; + self.send_header_state.timeout_fut = None; + return Poll::Ready(Error::e_explain( + WriteTimedout, + "writing header task timed out", + )); + } + } + + // Both write and timeout are pending + Poll::Pending + } + + /// Poll-based helper to write header bytes and optionally flush. + /// Handles state transitions explicitly. + fn poll_do_write_header_and_flush( + &mut self, + cx: &mut Context<'_>, + mut stream: Pin<&mut S>, + ) -> Poll> + where + S: AsyncWrite + Unpin + Send, + { + // Handle Idle state - take pending header and transition to Writing + if matches!(self.send_header_state.write_state, HeaderWriteState::Idle) { + if let Some(header_bytes) = self.send_header_state.pending_header.take() { + let size = header_bytes.len(); + self.send_header_state.write_state = HeaderWriteState::Writing(size, header_bytes); + } else { + // No pending header + self.send_header_state.write_state = HeaderWriteState::Done; + return Poll::Ready(Ok(0)); + } + } + + // Write if in Writing state + if let HeaderWriteState::Writing(original_size, ref mut buf) = + self.send_header_state.write_state + { + let size = original_size; + // Re-create tokio's stateless `write_all_buf` future per poll; + // progress is carried by `buf` (the caller-owned `Bytes`). + ready!(pin!(stream.as_mut().get_mut().write_all_buf(buf)).poll(cx)) + .map_err(|e| Error::because(WriteError, "writing response header", e))?; + + // Write complete - transition to next state + if self.send_header_state.should_flush { + self.send_header_state.write_state = HeaderWriteState::Flushing(size); + } else { + self.send_header_state.write_state = HeaderWriteState::Done; + return Poll::Ready(Ok(size)); + } + } + + // Handle the state after writing (or if we started in a non-Writing state) + match self.send_header_state.write_state { + HeaderWriteState::Flushing(size) => { + ready!(stream.as_mut().poll_flush(cx)) + .map_err(|e| Error::because(WriteError, "flushing response header", e))?; + // Flush complete - transition to Done + self.send_header_state.write_state = HeaderWriteState::Done; + Poll::Ready(Ok(size)) + } + HeaderWriteState::Done => Poll::Ready(Ok(0)), + HeaderWriteState::TimedOut => Poll::Ready(Error::e_explain( + WriteTimedout, + "header write task previously timed out", + )), + HeaderWriteState::Idle => { + unreachable!("Idle state should have been handled above") + } + HeaderWriteState::Writing(..) => { + unreachable!("Writing state should have been handled above") + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocols::http::v1::test_util::FlushTrackingMock; + use tokio_test::io::Builder; + + fn init_log() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + #[tokio::test] + async fn test_simple_header_write_no_flush() { + init_log(); + let header_data = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n"; + + let mock_io = Builder::new().write(header_data).build(); + let (mut flush_mock, flush_count) = FlushTrackingMock::new(mock_io); + + let mut header_writer = HeaderWriter::new(); + header_writer.send_header_task(Bytes::from_static(header_data), false, None); + + let result = header_writer + .write_current_header_task(&mut flush_mock) + .await; + assert!(result.is_ok()); + assert_eq!(result.unwrap(), header_data.len()); + assert_eq!( + FlushTrackingMock::flush_count(&flush_count), + 0, + "should_flush=false should not flush" + ); + } + + #[tokio::test] + async fn test_header_write_with_flush() { + init_log(); + let header_data = b"HTTP/1.1 200 OK\r\n\r\n"; + + let mock_io = Builder::new().write(header_data).build(); + let (mut flush_mock, flush_count) = FlushTrackingMock::new(mock_io); + + let mut header_writer = HeaderWriter::new(); + header_writer.send_header_task(Bytes::from_static(header_data), true, None); + + let result = header_writer + .write_current_header_task(&mut flush_mock) + .await; + assert!(result.is_ok()); + assert_eq!(result.unwrap(), header_data.len()); + assert_eq!( + FlushTrackingMock::flush_count(&flush_count), + 1, + "should_flush=true should flush exactly once" + ); + } + + // Uses start_paused for deterministic timer-based cancellation in select! + #[tokio::test(start_paused = true)] + async fn test_cancel_safe_header_write() { + init_log(); + let header_data = b"HTTP/1.1 200 OK\r\nServer: pingora\r\n\r\n"; + + // Mock that blocks to allow cancellation + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(100)) + .write(header_data) + .build(); + + let mut header_writer = HeaderWriter::new(); + header_writer.send_header_task(Bytes::from_static(header_data), false, None); + + let mut cancel_count = 0; + + loop { + if !header_writer.has_pending_header_task() { + break; + } + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + cancel_count += 1; + } + result = header_writer.write_current_header_task(&mut mock_io) => { + assert!(result.is_ok()); + assert_eq!(result.unwrap(), header_data.len()); + break; + } + } + } + + assert!(cancel_count > 0, "Should have cancelled at least once"); + } + + #[tokio::test] + async fn test_header_write_timeout() { + init_log(); + let header_data = b"HTTP/1.1 200 OK\r\n\r\n"; + + // Mock that blocks forever + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_secs(1000)) + .build(); + + let mut header_writer = HeaderWriter::new(); + header_writer.send_header_task( + Bytes::from_static(header_data), + false, + Some(std::time::Duration::from_millis(50)), + ); + + let result = header_writer.write_current_header_task(&mut mock_io).await; + assert!(result.is_err()); + assert_eq!(result.unwrap_err().etype(), &WriteTimedout); + } + + #[tokio::test] + async fn test_header_write_timeout_persists() { + init_log(); + let header_data = b"HTTP/1.1 200 OK\r\n\r\n"; + + // Mock that blocks for a while + let mut mock_io = Builder::new() + .wait(std::time::Duration::from_millis(200)) + .build(); + + let mut header_writer = HeaderWriter::new(); + header_writer.send_header_task( + Bytes::from_static(header_data), + false, + Some(std::time::Duration::from_millis(100)), + ); + + let mut attempts = 0; + let mut timedout = false; + + loop { + if !header_writer.has_pending_header_task() { + break; + } + + attempts += 1; + + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => { + continue; + } + result = header_writer.write_current_header_task(&mut mock_io) => { + match result { + Ok(_) => break, + Err(e) if e.etype() == &WriteTimedout => { + timedout = true; + break; + } + Err(e) => panic!("Unexpected error: {:?}", e), + } + } + } + } + + assert!(timedout, "Timeout should have fired"); + assert!(attempts >= 5, "Should have had multiple attempts"); + } +} diff --git a/pingora-core/src/protocols/http/v1/mod.rs b/pingora-core/src/protocols/http/v1/mod.rs index 196024917..6f085a70f 100644 --- a/pingora-core/src/protocols/http/v1/mod.rs +++ b/pingora-core/src/protocols/http/v1/mod.rs @@ -17,4 +17,110 @@ pub(crate) mod body; pub mod client; pub mod common; +pub(crate) mod header; pub mod server; + +/// Test utilities shared across HTTP/1.x unit tests +#[cfg(test)] +pub(crate) mod test_util { + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::task::{Context, Poll}; + use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + use tokio_test::io::Mock; + + /// A wrapper around [`Mock`] that counts flush calls. + /// + /// `tokio_test::io::Mock`'s `poll_flush` always returns `Ready(Ok(()))`, + /// so we can't detect flush calls via mock alone. This wrapper counts them. + #[derive(Debug)] + pub(crate) struct FlushTrackingMock { + inner: Mock, + flush_count: Arc, + } + + impl FlushTrackingMock { + pub(crate) fn new(mock: Mock) -> (Self, Arc) { + let flush_count = Arc::new(AtomicUsize::new(0)); + ( + FlushTrackingMock { + inner: mock, + flush_count: flush_count.clone(), + }, + flush_count, + ) + } + + pub(crate) fn flush_count(counter: &Arc) -> usize { + counter.load(Ordering::Relaxed) + } + } + + impl AsyncRead for FlushTrackingMock { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_read(cx, buf) + } + } + + impl AsyncWrite for FlushTrackingMock { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_write(cx, buf) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let result = Pin::new(&mut this.inner).poll_flush(cx); + if let Poll::Ready(Ok(())) = &result { + this.flush_count.fetch_add(1, Ordering::Relaxed); + } + result + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } + } + + // Implement IO-required traits so FlushTrackingMock can be used as Box + // in HttpSession tests (server.rs). + use crate::protocols::{ + raw_connect::ProxyDigest, GetProxyDigest, GetSocketDigest, GetTimingDigest, Peek, Shutdown, + SocketDigest, Ssl, TimingDigest, UniqueID, UniqueIDType, + }; + + #[async_trait::async_trait] + impl Shutdown for FlushTrackingMock { + async fn shutdown(&mut self) -> () {} + } + impl UniqueID for FlushTrackingMock { + fn id(&self) -> UniqueIDType { + 0 + } + } + impl Ssl for FlushTrackingMock {} + impl GetTimingDigest for FlushTrackingMock { + fn get_timing_digest(&self) -> Vec> { + vec![] + } + } + impl GetProxyDigest for FlushTrackingMock { + fn get_proxy_digest(&self) -> Option> { + None + } + } + impl GetSocketDigest for FlushTrackingMock { + fn get_socket_digest(&self) -> Option> { + None + } + } + impl Peek for FlushTrackingMock {} +} diff --git a/pingora-core/src/protocols/http/v1/server.rs b/pingora-core/src/protocols/http/v1/server.rs index b071e6fde..45837c1da 100644 --- a/pingora-core/src/protocols/http/v1/server.rs +++ b/pingora-core/src/protocols/http/v1/server.rs @@ -27,15 +27,49 @@ use pingora_error::{Error, ErrorType::*, OrErr, Result}; use pingora_http::{IntoCaseHeaderName, RequestHeader, ResponseHeader}; use pingora_timeout::timeout; use regex::bytes::Regex; +use std::any::Any; +use std::collections::VecDeque; use std::time::Duration; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::body::{BodyReader, BodyWriter}; use super::common::*; -use crate::protocols::http::{body_buffer::FixedBuffer, date, HttpTask}; +use super::header::HeaderWriter; +use crate::protocols::http::{body_buffer::FixedBuffer, date, HttpTask, ReusableHttpStream}; use crate::protocols::{Digest, SocketAddr, Stream}; use crate::utils::{BufRef, KVRef}; +/// Tracks which writer is currently processing a task. +/// +/// This enables resuming writes after cancellation. Each variant stores the +/// minimal data needed for cleanup after write completes. +#[derive(Debug)] +enum ProxyTaskWriter { + /// Currently writing a header task. + /// Stores: (header for `response_written`, end_stream flag) + WritingHeader(Box, bool), + /// Currently writing a body task (`Body` or `UpgradedBody`). + /// Stores: (end_stream flag) + WritingBody(bool), + /// Currently finishing the body (writing last chunk + flush). + FinishingBody, +} + +/// State for the cancel-safe proxy task write API. +#[derive(Default)] +struct ProxyTaskState { + /// Lazily initialized — `HeaderWriter::new()` heap-allocates. + header_writer: Option, + tasks: VecDeque, + current_writer: Option, +} + +impl ProxyTaskState { + fn header_writer(&mut self) -> &mut HeaderWriter { + self.header_writer.get_or_insert_with(HeaderWriter::new) + } +} + /// The HTTP 1.x server session pub struct HttpSession { underlying_stream: Stream, @@ -51,6 +85,8 @@ pub struct HttpSession { body_reader: BodyReader, /// A state machine to track how to write the response body body_writer: BodyWriter, + /// Cancel-safe proxy task state. + proxy_task_state: ProxyTaskState, /// An internal buffer to buf multiple body writes to reduce the underlying syscalls body_write_buf: BytesMut, /// Track how many application (not on the wire) body bytes already sent @@ -86,6 +122,38 @@ pub struct HttpSession { /// Number of times the upstream connection associated with this session can be reused /// after this session ends keepalive_reuses_remaining: Option, + /// User-defined context carried across requests on the same keepalive connection. + /// Set by [`HttpPersistentSettings::apply_to_session`](crate::apps::HttpPersistentSettings::apply_to_session), + /// consumed by the proxy layer via [`take_connection_user_context`](Self::take_connection_user_context). + connection_user_context: Option>, + /// Whether the client has closed the TCP connection (sent FIN / read returned 0). + half_closed: bool, + /// When true (default), a client close after the request body is surfaced as a + /// `ConnectionClosed` error so the proxy aborts immediately. When false, the + /// close is tolerated and `read_body_or_idle` stays pending so the proxy can + /// finish delivering the upstream response (RFC 9112 Section 9.6). + abort_on_close: bool, + /// Whether the cancel-safe proxy task API is enabled for this session. + /// Defaults to false. Can be enabled via [`set_proxy_tasks_enabled`](Self::set_proxy_tasks_enabled). + proxy_tasks_enabled: bool, + /// Whether HTTP/1.1 request pipelining is enabled for this session. + /// Defaults to false. Can be enabled via [`set_pipelining_enabled`](Self::set_pipelining_enabled). + /// See [`Self::set_pipelining_enabled`] for RFC 9112 §9.3.2 semantics. + pipelining_enabled: bool, + /// Pipelined bytes from the previous request on the same keep-alive connection, + /// to be parsed as the start of this session's request. Consumed on the first + /// call to [`Self::read_request`]. Set via [`Self::set_pipelined_prefix`] after + /// the previous session's [`BodyReader::take_body_overread`] yielded bytes. + pipelined_prefix: Option, + /// Set once the idle-branch of [`Self::read_body_or_idle`] has read the + /// first bytes of a pipelined next request and pushed them onto the body + /// reader's overread surface. Further idle polls on the same request + /// return pending instead of re-reading the stream, so the body-pump + /// `tokio::select!` loop can exit via its other branches while the + /// stashed bytes travel through `reuse()` + + /// [`super::super::HttpPersistentSettings`] into the next session. + /// Scoped narrowly so it cannot affect FIN / `abort_on_close` semantics. + pipelined_idle_bytes_stashed: bool, } impl HttpSession { @@ -108,6 +176,7 @@ impl HttpSession { preread_body: None, body_reader: BodyReader::new(false), body_writer: BodyWriter::new(), + proxy_task_state: ProxyTaskState::default(), body_write_buf: BytesMut::new(), keepalive_timeout: KeepaliveStatus::Off, update_resp_headers: true, @@ -126,6 +195,76 @@ impl HttpSession { // default on to avoid rejecting requests after body as pipelined close_on_response_before_downstream_finish: true, keepalive_reuses_remaining: None, + connection_user_context: None, + half_closed: false, + abort_on_close: true, + proxy_tasks_enabled: false, + pipelining_enabled: false, + pipelined_prefix: None, + pipelined_idle_bytes_stashed: false, + } + } + + async fn read_request_buf( + &mut self, + buf: &mut BytesMut, + already_read: usize, + ) -> Result> { + let read_result = { + let read_event = self.underlying_stream.read_buf(buf); + match self.keepalive_timeout { + KeepaliveStatus::Timeout(d) => match timeout(d, read_event).await { + Ok(res) => res, + Err(e) => { + debug!("keepalive timeout {d:?} reached, {e}"); + return Ok(None); + } + }, + KeepaliveStatus::Infinite => { + // FIXME: this should only apply to reads between requests + read_event.await + } + KeepaliveStatus::Off => match self.read_timeout { + Some(t) => match timeout(t, read_event).await { + Ok(res) => res, + Err(e) => { + debug!("read timeout {t:?} reached, {e}"); + return Error::e_explain(ReadTimedout, format!("timeout: {t:?}")); + } + }, + None => read_event.await, + }, + } + }; + + match read_result { + Ok(n_read) => { + if n_read == 0 { + if already_read > 0 { + Error::e_explain( + ConnectionClosed, + format!( + "while reading request headers, bytes already read: {}", + already_read + ), + ) + } else { + /* common when client decides to close a keepalived session */ + debug!("Client prematurely closed connection with 0 byte sent"); + Ok(None) + } + } else { + Ok(Some(n_read)) + } + } + Err(e) => { + if already_read > 0 { + Error::e_because(ReadError, "while reading request headers", e) + } else { + /* nothing harmful since we have not ready any thing yet */ + Ok(None) + } + } } } @@ -138,6 +277,24 @@ impl HttpSession { self.buf.clear(); let mut buf = BytesMut::with_capacity(INIT_HEADER_BUF_SIZE); let mut already_read: usize = 0; + // If the caller (e.g. the proxy layer completing a pipelined request on + // a reused keep-alive connection) handed us bytes that were read past + // the end of the previous request's body, pre-fill our parse buffer so + // the header parser sees them as the start of this request. The loop + // below tries to parse first when we already have pipelined bytes — + // a pipelined prefix can contain a complete request header, in which + // case we must NOT issue another stream read (which would block). + let mut skip_next_read = false; + if let Some(prefix) = self + .pipelined_prefix + .take() + .filter(|prefix| !prefix.is_empty()) + { + buf.reserve(prefix.len()); + buf.extend_from_slice(&prefix); + already_read = prefix.len(); + skip_next_read = true; + } loop { if already_read > MAX_HEADER_SIZE { /* NOTE: this check only blocks second read. The first large read is allowed @@ -149,61 +306,19 @@ impl HttpSession { ); } - let read_result = { - let read_event = self.underlying_stream.read_buf(&mut buf); - match self.keepalive_timeout { - KeepaliveStatus::Timeout(d) => match timeout(d, read_event).await { - Ok(res) => res, - Err(e) => { - debug!("keepalive timeout {d:?} reached, {e}"); - return Ok(None); - } - }, - KeepaliveStatus::Infinite => { - // FIXME: this should only apply to reads between requests - read_event.await - } - KeepaliveStatus::Off => match self.read_timeout { - Some(t) => match timeout(t, read_event).await { - Ok(res) => res, - Err(e) => { - debug!("read timeout {t:?} reached, {e}"); - return Error::e_explain(ReadTimedout, format!("timeout: {t:?}")); - } - }, - None => read_event.await, - }, - } - }; - let n = match read_result { - Ok(n_read) => { - if n_read == 0 { - if already_read > 0 { - return Error::e_explain( - ConnectionClosed, - format!( - "while reading request headers, bytes already read: {}", - already_read - ), - ); - } else { - /* common when client decides to close a keepalived session */ - debug!("Client prematurely closed connection with 0 byte sent"); - return Ok(None); - } - } - n_read - } - - Err(e) => { - if already_read > 0 { - return Error::e_because(ReadError, "while reading request headers", e); - } - /* nothing harmful since we have not ready any thing yet */ - return Ok(None); - } - }; - already_read += n; + // On the first iteration after a pipelined prefix was injected, + // attempt to parse what we already have before issuing a stream + // read. If the prefix contains a complete request header, a + // subsequent read_buf() would block for data that may never come + // (the client already pipelined everything it had to send for + // this request and is waiting for our response). + if skip_next_read { + skip_next_read = false; + } else if let Some(n) = self.read_request_buf(&mut buf, already_read).await? { + already_read += n; + } else { + return Ok(None); + } // Use loop as GOTO to retry escaped request buffer, not a real loop loop { @@ -276,6 +391,7 @@ impl HttpSession { self.request_header = Some(request_header); self.body_reader.reinit(); + self.pipelined_idle_bytes_stashed = false; self.response_written = None; self.respect_keepalive(); @@ -495,101 +611,12 @@ impl HttpSession { /// Write the response header to the client. /// This function can be called more than once to send 1xx informational headers excluding 101. pub async fn write_response_header(&mut self, mut header: Box) -> Result<()> { - if header.status.is_informational() && self.ignore_info_resp(header.status.into()) { - debug!("ignoring informational headers"); + // Prepare header (handle upgrades, set headers, initialize body writer, serialize to bytes) + let Some((write_buf, flush)) = self.prepare_response_header(&mut header)? else { + // Header already sent or should be ignored return Ok(()); - } - - if let Some(resp) = self.response_written.as_ref() { - if !resp.status.is_informational() || self.upgraded { - warn!("Respond header is already sent, cannot send again"); - return Ok(()); - } - } - - // if body unfinished, or request header was not finished reading - if self.close_on_response_before_downstream_finish - && (self.request_header.is_none() || !self.is_body_done()) - { - debug!("set connection close before downstream finish"); - self.set_keepalive(None); - } - - // no need to add these headers to 1xx responses - if !header.status.is_informational() && self.update_resp_headers { - /* update headers */ - header.insert_header(header::DATE, date::get_cached_date())?; - - // TODO: make these lazy static - let connection_value = if self.will_keepalive() { - "keep-alive" - } else { - "close" - }; - header.insert_header(header::CONNECTION, connection_value)?; - } - - if header.status == 101 { - // make sure the connection is closed at the end when 101/upgrade is used - self.set_keepalive(None); - } - - // Allow informational header (excluding 101) to pass through without affecting the state - // of the request - if header.status == 101 || !header.status.is_informational() { - // reset request body to done for incomplete upgrade handshakes - if let Some(upgrade_ok) = self.is_upgrade(&header) { - if upgrade_ok { - debug!("ok upgrade handshake"); - // For ws we use HTTP1_0 do_read_body_until_closed - // - // On ws close the initiator sends a close frame and - // then waits for a response from the peer, once it receives - // a response it closes the conn. After receiving a - // control frame indicating the connection should be closed, - // a peer discards any further data received. - // https://www.rfc-editor.org/rfc/rfc6455#section-1.4 - self.upgraded = true; - // Now that the upgrade was successful, we need to change - // how we interpret the rest of the body as pass-through. - if self.body_reader.need_init() { - self.init_body_reader(); - } else { - // already initialized - // immediately start reading the rest of the body as upgraded - // (in practice most upgraded requests shouldn't have any body) - // - // TODO: https://datatracker.ietf.org/doc/html/rfc9110#name-upgrade - // the most spec-compliant behavior is to switch interpretation - // after sending the former body, - // we immediately switch interpretation to match nginx - self.body_reader.convert_to_close_delimited(); - } - } else { - // this was a request that requested Upgrade, - // but upstream did not comply - debug!("bad upgrade handshake!"); - // continue to read body as-is, this is now just a regular request - } - } - self.init_body_writer(&header); - } - - // Defense-in-depth: if response body is close-delimited, mark session - // as un-reusable - if self.body_writer.is_close_delimited() { - self.set_keepalive(None); - } - - // Don't have to flush response with content length because it is less - // likely to be real time communication. So do flush when - // 1.1xx response: client needs to see it before the rest of response - // 2.No content length: the response could be generated in real time - let flush = header.status.is_informational() - || header.headers.get(header::CONTENT_LENGTH).is_none(); + }; - let mut write_buf = BytesMut::with_capacity(INIT_HEADER_BUF_SIZE); - http_resp_header_to_buf(&header, &mut write_buf).unwrap(); match self.underlying_stream.write_all(&write_buf).await { Ok(()) => { // flush the stream if 1xx header or there is no response body @@ -600,7 +627,6 @@ impl HttpSession { .or_err(WriteError, "flushing response header")?; } self.response_written = Some(header); - self.body_bytes_sent += write_buf.len(); Ok(()) } Err(e) => Error::e_because(WriteError, "writing response header", e), @@ -662,6 +688,24 @@ impl HttpSession { self.keepalive_reuses_remaining } + /// Set user-defined context to carry across requests on the same keepalive connection. + /// + /// This is typically called by + /// [`HttpPersistentSettings::apply_to_session`](crate::apps::HttpPersistentSettings::apply_to_session) + /// during the keepalive reuse loop. The proxy layer consumes it via + /// [`take_connection_user_context`](Self::take_connection_user_context). + pub fn set_connection_user_context(&mut self, ctx: Option>) { + self.connection_user_context = ctx; + } + + /// Take the user-defined context from the previous request on this keepalive connection. + /// + /// Returns `None` if this is the first request on the connection or if no context was + /// persisted by the previous request. + pub fn take_connection_user_context(&mut self) -> Option> { + self.connection_user_context.take() + } + /// Return whether the session will be keepalived for connection reuse. pub fn will_keepalive(&self) -> bool { !matches!( @@ -726,6 +770,117 @@ impl HttpSession { } } + /// Prepare response header for writing: handle upgrades, set headers, initialize body writer. + /// This contains all the synchronous logic that should happen before writing the header. + /// Returns Ok(Some((bytes, should_flush))) if the header should be written, Ok(None) if should skip. + fn prepare_response_header( + &mut self, + header: &mut ResponseHeader, + ) -> Result> { + // Check if we should ignore informational responses + if header.status.is_informational() && self.ignore_info_resp(header.status.into()) { + debug!("ignoring informational headers"); + return Ok(None); + } + + // Check if we already sent a response header + if let Some(ref resp) = self.response_written { + if !resp.status.is_informational() || self.upgraded { + warn!("Respond header is already sent, cannot send again"); + return Ok(None); + } + } + + // if body unfinished, or request header was not finished reading + if self.close_on_response_before_downstream_finish + && (self.request_header.is_none() || !self.is_body_done()) + { + debug!("set connection close before downstream finish"); + self.set_keepalive(None); + } + + // no need to add these headers to 1xx responses + if !header.status.is_informational() && self.update_resp_headers { + /* update headers */ + header.insert_header(header::DATE, date::get_cached_date())?; + + // TODO: make these lazy static + let connection_value = if self.will_keepalive() { + "keep-alive" + } else { + "close" + }; + header.insert_header(header::CONNECTION, connection_value)?; + } + + if header.status == 101 { + // make sure the connection is closed at the end when 101/upgrade is used + self.set_keepalive(None); + } + + // Allow informational header (excluding 101) to pass through without affecting the state + // of the request + if header.status == 101 || !header.status.is_informational() { + // reset request body to done for incomplete upgrade handshakes + if let Some(upgrade_ok) = self.is_upgrade(header) { + if upgrade_ok { + debug!("ok upgrade handshake"); + // For ws we use HTTP1_0 do_read_body_until_closed + // + // On ws close the initiator sends a close frame and + // then waits for a response from the peer, once it receives + // a response it closes the conn. After receiving a + // control frame indicating the connection should be closed, + // a peer discards any further data received. + // https://www.rfc-editor.org/rfc/rfc6455#section-1.4 + self.upgraded = true; + // Now that the upgrade was successful, we need to change + // how we interpret the rest of the body as pass-through. + if self.body_reader.need_init() { + self.init_body_reader(); + } else { + // already initialized + // immediately start reading the rest of the body as upgraded + // (in practice most upgraded requests shouldn't have any body) + // + // TODO: https://datatracker.ietf.org/doc/html/rfc9110#name-upgrade + // the most spec-compliant behavior is to switch interpretation + // after sending the former body, + // we immediately switch interpretation to match nginx + self.body_reader.convert_to_close_delimited(); + } + } else { + // this was a request that requested Upgrade, + // but upstream did not comply + debug!("bad upgrade handshake!"); + // continue to read body as-is, this is now just a regular request + } + } + self.init_body_writer(header); + } + + // Defense-in-depth: if response body is close-delimited, mark session + // as un-reusable + if self.body_writer.is_close_delimited() { + self.set_keepalive(None); + } + + // Serialize header to bytes + let mut write_buf = BytesMut::with_capacity(INIT_HEADER_BUF_SIZE); + http_resp_header_to_buf(header, &mut write_buf) + .map_err(|_| Error::explain(WriteError, "serializing response header"))?; + + // Determine if we should flush + // Don't have to flush response with content length because it is less + // likely to be real time communication. So do flush when + // 1. 1xx response: client needs to see it before the rest of response + // 2. No content length: the response could be generated in real time + let should_flush = header.status.is_informational() + || header.headers.get(header::CONTENT_LENGTH).is_none(); + + Ok(Some((write_buf.freeze(), should_flush))) + } + fn init_body_writer(&mut self, header: &ResponseHeader) { use http::StatusCode; /* the following responses don't have body 204, 304, and HEAD */ @@ -781,6 +936,66 @@ impl HttpSession { } } + /// Whether the cancel-safe proxy task API is enabled for this session. + pub fn proxy_tasks_enabled(&self) -> bool { + self.proxy_tasks_enabled + } + + /// Enable or disable the cancel-safe proxy task API for this session. + pub fn set_proxy_tasks_enabled(&mut self, enabled: bool) { + self.proxy_tasks_enabled = enabled; + } + + /// Whether HTTP/1.1 request pipelining is enabled for this session. + pub fn pipelining_enabled(&self) -> bool { + self.pipelining_enabled + } + + /// Enable or disable HTTP/1.1 request pipelining on this session. + /// + /// When enabled, if the client pipelines requests on a single keep-alive + /// connection (sends request N+1 before reading response N), the proxy will + /// serve each request sequentially with responses in request order as + /// required by RFC 9112 §9.3.2. Each pipelined request still goes through + /// independent upstream selection; only the downstream connection is reused. + /// + /// When disabled (default), pipelined bytes received alongside request N + /// cause the session to be marked un-reusable: response N is still + /// delivered, the connection closes, and request N+1 is dropped. Clients + /// are expected to detect the close and retry on a fresh connection per + /// RFC 9112 §9.3.2. + /// + /// Sequential dispatch only: response N must be fully written before + /// request N+1 begins processing. No parallel pipelining. + pub fn set_pipelining_enabled(&mut self, enabled: bool) { + self.pipelining_enabled = enabled; + } + + /// Set pipelined bytes to be parsed as the start of this session's request. + /// + /// Called by the proxy layer when continuing a keep-alive connection whose + /// previous session yielded overread bytes. The prefix is consumed on the + /// first [`Self::read_request`] call; the parser treats the prefix + any + /// further stream reads as the next request's header + body bytes. + pub fn set_pipelined_prefix(&mut self, prefix: BytesMut) { + debug_assert!( + self.pipelined_prefix.is_none(), + "pipelined prefix already set" + ); + self.pipelined_prefix = Some(prefix); + } + + /// Take ownership of bytes read past the end of this session's request + /// body. When non-empty, those bytes are the start of a pipelined + /// follow-up request on the same keep-alive connection and should be + /// fed to the next session via [`Self::set_pipelined_prefix`]. + /// + /// Returns `None` when no overread is present. After this call, the + /// session's body-reader no longer holds the bytes. + pub(crate) fn take_body_overread(&mut self) -> Option { + self.body_reader.take_body_overread() + } + async fn do_write_body_buf(&mut self) -> Result> { // Don't flush empty chunks, they are considered end of body for chunks if self.body_write_buf.is_empty() { @@ -925,31 +1140,115 @@ impl HttpSession { /// This function will (async) block forever until the client closes the connection. pub async fn idle(&mut self) -> Result { - // NOTE: this implementation breaks http pipelining, ideally we need poll_error - // NOTE: buf cannot be empty, openssl-rs read() requires none empty buf. - let mut buf: [u8; 1] = [0; 1]; - self.underlying_stream - .read(&mut buf) + // OpenSSL read requires a non-empty buffer. Keep this probe at one byte + // so idle-style reads consume at most one byte before returning control. + self.read_idle_probe("during HTTP idle state") .await - .or_err(ReadError, "during HTTP idle state") + .map(|(_, read)| read) } - /// This function will return body bytes (same as [`Self::read_body_bytes()`]), but after + async fn read_idle_probe(&mut self, context: &'static str) -> Result<([u8; 1], usize)> { + let mut probe = [0; 1]; + let read = self + .underlying_stream + .read(&mut probe) + .await + .or_err(ReadError, context)?; + Ok((probe, read)) + } + + /// This function will return body bytes (same as [`Self::read_body_bytes()`]), but after /// the client body finishes (`Ok(None)` is returned), calling this function again will block /// forever, same as [`Self::idle()`]. + /// + /// By default (`abort_on_close = true`), if the client closes the connection + /// (sends TCP FIN, i.e. `read == 0`) after the request body is complete, a + /// `ConnectionClosed` error is returned. + /// + /// When `abort_on_close` is **disabled**, the close is tolerated: the future stays + /// pending so the proxy can finish delivering the upstream response via the write + /// path (per RFC 9112 Section 9.6). A true disconnect (RST) will be caught later + /// when the response write fails. + /// + /// Note that this marks the connection as half-closed if FIN is detected. If this function + /// is called after the connection is already marked half-closed and `abort_on_close` is + /// **disabled**, then it will pend forever. pub async fn read_body_or_idle(&mut self, no_body_expected: bool) -> Result> { if no_body_expected || self.is_body_done() { + if self.half_closed { + if self.abort_on_close { + return Error::e_explain( + ConnectionClosed, + if self.response_written.is_none() { + "Prematurely before response header is sent" + } else { + "Prematurely before response body is complete" + }, + ); + } + return std::future::pending().await; + } + // When pipelining is enabled and an earlier idle read already + // stashed the next request's bytes as overread, any further + // poll of this function on the same request must not read + // the stream again (the proxy's body-pump `select!` loop + // will call back into here repeatedly until its other + // branches resolve the request). Go straight to pending. + // `abort_on_close` and the FIN handling above stay untouched + // — this branch is exclusive to the pipelining case where + // bytes (not FIN) arrived on the idle poll. + if self.pipelining_enabled && self.pipelined_idle_bytes_stashed { + return std::future::pending().await; + } // XXX: account for upgraded body reader change, if the read half split from the write half - let read = self.idle().await?; + let (probe, read) = self + .read_idle_probe("during HTTP body-or-idle state") + .await?; if read == 0 { - Error::e_explain( - ConnectionClosed, - if self.response_written.is_none() { - "Prematurely before response header is sent" - } else { - "Prematurely before response body is complete" - }, - ) + self.half_closed = true; + self.set_keepalive(None); + if self.abort_on_close { + Error::e_explain( + ConnectionClosed, + if self.response_written.is_none() { + "Prematurely before response header is sent" + } else { + "Prematurely before response body is complete" + }, + ) + } else { + debug!("downstream closed (FIN), keeping write side open"); + // If the connection is fully closed, writing the response side + // will fail. + std::future::pending().await + } + } else if self.pipelining_enabled { + // The read bytes are the start of a pipelined next + // request on this keep-alive connection (RFC 9112 + // §9.3.2). Stash them on the body reader's overread + // surface so the existing `take_body_overread` + + // `HttpPersistentSettings` extraction path picks them + // up at `reuse()` time and feeds them to the next + // session via `set_pipelined_prefix`. + // + // Returning pending (rather than `Ok(None)` or an + // error) signals the body-pump `tokio::select!` loop + // that the downstream has no more body work to do on + // this request — the loop exits naturally via the + // upstream-response-done / response-write-done + // branches, and `finish()` runs its standard pipelining + // extraction. The read == 0 FIN path above is unchanged; this + // branch only handles a non-zero idle read that belongs to the + // next pipelined request, so it leaves `half_closed` and + // `abort_on_close` untouched. + // Keep the stash and flag update adjacent and synchronous. + // Once the prefix byte is handed to the overread path, the + // flag prevents later idle polls for this request from + // reading the stream again. + self.body_reader.push_body_overread(&probe[..read]); + self.pipelined_idle_bytes_stashed = true; + debug!("pipelined request bytes stashed as overread ({read} bytes)"); + std::future::pending().await } else { Error::e_explain(ConnectError, "Sent data after end of body") } @@ -958,6 +1257,11 @@ impl HttpSession { } } + /// Whether the client has half-closed the TCP connection. + pub fn is_half_closed(&self) -> bool { + self.half_closed + } + /// Return the raw bytes of the request header. pub fn get_headers_raw_bytes(&self) -> Bytes { self.raw_header.as_ref().unwrap().get_bytes(&self.buf) @@ -1055,6 +1359,18 @@ impl HttpSession { self.close_on_response_before_downstream_finish = close; } + /// Controls behaviour when the client closes the connection after the request body. + /// + /// When **enabled** (default), a client close is returned as a `ConnectionClosed` + /// error so the proxy aborts immediately. + /// + /// When **disabled**, `read_body_or_idle` stays pending on a client close so the + /// proxy can finish delivering the upstream response (RFC 9112 Section 9.6). A true + /// disconnect (RST) will surface later when the response write fails. + pub fn set_abort_on_close(&mut self, abort: bool) { + self.abort_on_close = abort; + } + /// Return the [Digest] of the connection. pub fn digest(&self) -> &Digest { &self.digest @@ -1081,32 +1397,35 @@ impl HttpSession { .map(|d| d.local_addr())? } - /// Consume `self`, if the connection can be reused, the underlying stream will be returned - /// to be fed to the next [`Self::new()`]. This drains any remaining request body if it hasn't - /// yet been read and the stream is reusable. + /// Consume `self`, if the connection can be reused, the underlying stream and any pipelined + /// prefix bytes will be returned to be fed to the next [`Self::new()`]. This drains any + /// remaining request body if it hasn't yet been read and the stream is reusable. /// /// The next session can just call [`Self::read_request()`]. /// /// If the connection cannot be reused, the underlying stream will be closed and `None` will be /// returned. If there was an error while draining any remaining request body that error will /// be returned. - pub async fn reuse(mut self) -> Result> { + pub async fn reuse(mut self) -> Result> { if !self.will_keepalive() { debug!("HTTP shutdown connection"); self.shutdown().await; Ok(None) } else { self.drain_request_body().await?; - // XXX: currently pipelined requests are not properly read without - // pipelining support, and pingora 400s if pipelined requests are sent - // in the middle of another request. - // We will mark the connection as un-reusable so it may be closed, - // the pipelined request left unread, and the client can attempt to resend - if self.body_reader.has_bytes_overread() { + if self.body_reader.has_bytes_overread() && !self.pipelining_enabled { debug!("bytes overread on request, disallowing reuse"); Ok(None) } else { - Ok(Some(self.underlying_stream)) + let pipelined_prefix = self + .pipelining_enabled + .then(|| self.take_body_overread()) + .flatten() + .filter(|prefix| !prefix.is_empty()); + Ok(Some(ReusableHttpStream::new( + self.underlying_stream, + pipelined_prefix, + ))) } } } @@ -1235,6 +1554,152 @@ impl HttpSession { Ok(end_stream || self.body_writer.finished()) } + /// Queue a proxy task for cancel-safe writing with the current write_timeout. + /// The task will be written when `write_proxy_tasks()` is called. + /// + /// A write canceled mid-operation can be resumed via `write_proxy_tasks()`. + pub fn send_proxy_task(&mut self, task: HttpTask) { + self.proxy_task_state.tasks.push_back(task); + } + + /// Check if there are pending proxy tasks queued for writing. + pub fn has_pending_proxy_tasks(&self) -> bool { + self.proxy_task_state.current_writer.is_some() || !self.proxy_task_state.tasks.is_empty() + } + + /// Write all queued proxy tasks (response `HttpTask`s from `send_proxy_task`) + /// in a cancel-safe manner. + /// + /// If cancelled mid-write, the next call will resume the in-progress write. + /// + /// Returns `Ok(true)` if this was the end of the response stream. + // Leverages the cancel-safe `HeaderWriter` and `BodyWriter` primitives. + // TODO: we can do the same for the non-cancel-safe APIs. + pub async fn write_proxy_tasks(&mut self) -> Result { + let mut end_stream = false; + + // TODO: buffer body data like response_duplex_vec + loop { + // - Resume any in-progress write + if let Some(ref writer_state) = self.proxy_task_state.current_writer { + match writer_state { + ProxyTaskWriter::WritingHeader(_, _) => { + let _bytes_written = self + .proxy_task_state + .header_writer() + .write_current_header_task(&mut self.underlying_stream) + .await + .map_err(|e| e.into_down())?; + } + ProxyTaskWriter::WritingBody(_) => { + let written = self + .body_writer + .write_current_body_task(&mut self.underlying_stream) + .await + .map_err(|e| e.into_down())?; + if let Some(n) = written { + self.body_bytes_sent += n; + } + } + ProxyTaskWriter::FinishingBody => { + self.body_writer + .write_current_finish_task(&mut self.underlying_stream) + .await + .map_err(|e| e.into_down())?; + } + } + + match self + .proxy_task_state + .current_writer + .take() + .expect("writer state present") + { + ProxyTaskWriter::WritingHeader(header, end) => { + self.response_written = Some(header); + end_stream = end; + } + ProxyTaskWriter::WritingBody(end) => { + end_stream = end; + } + ProxyTaskWriter::FinishingBody => { + end_stream = true; + self.maybe_force_close_body_reader(); + break; // fine to break after finish, no tasks should be queued after + } + } + continue; + } + + // - Send tasks, set state. + // Pop next task + let Some(task) = self.proxy_task_state.tasks.pop_front() else { + if end_stream { + self.body_writer.send_finish_task(); + self.proxy_task_state.current_writer = Some(ProxyTaskWriter::FinishingBody); + continue; + } + break; + }; + + match task { + HttpTask::Header(mut header, end) => { + let Some((write_buf, should_flush)) = + self.prepare_response_header(&mut header)? + else { + end_stream = end; + continue; + }; + // header only responses will want to flush + let flush = should_flush || self.body_writer.finished(); + self.proxy_task_state + .header_writer() + .send_header_task(write_buf, flush, None); + self.proxy_task_state.current_writer = + Some(ProxyTaskWriter::WritingHeader(header, end)); + } + HttpTask::Body(ref data, end) => { + if self.upgraded { + panic!("Unexpected Body task received on upgraded downstream session"); + } + if let Some(d) = data.as_ref() { + if !d.is_empty() { + let body_timeout = self.write_timeout(d.len()); + self.body_writer.send_body_task(d.clone(), body_timeout); + self.proxy_task_state.current_writer = + Some(ProxyTaskWriter::WritingBody(end)); + continue; + } + } + end_stream = end; + } + HttpTask::UpgradedBody(ref data, end) => { + if !self.upgraded { + panic!("Unexpected UpgradedBody task received on un-upgraded downstream session"); + } + if let Some(d) = data.as_ref() { + if !d.is_empty() { + let body_timeout = self.write_timeout(d.len()); + self.body_writer.send_body_task(d.clone(), body_timeout); + self.proxy_task_state.current_writer = + Some(ProxyTaskWriter::WritingBody(end)); + continue; + } + } + end_stream = end; + } + HttpTask::Trailer(_) | HttpTask::Done => { + end_stream = true; + } + HttpTask::Failed(e) => { + return Err(e); + } + } + } + + Ok(end_stream || self.body_writer.finished()) + } + /// Get the reference of the [Stream] that this HTTP session is operating upon. pub fn stream(&self) -> &Stream { &self.underlying_stream @@ -2296,6 +2761,30 @@ mod tests_stream { assert_eq!(wire_body.len(), n); } + #[tokio::test] + async fn body_bytes_sent_excludes_response_header() { + let read_wire = b"GET / HTTP/1.1\r\n\r\n"; + let wire_header = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n"; + let wire_body = b"hello"; + let mock_io = Builder::new() + .read(read_wire) + .write(wire_header) + .write(wire_body) + .build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.read_request().await.unwrap(); + let mut new_response = ResponseHeader::build(StatusCode::OK, None).unwrap(); + new_response.append_header("Content-Length", "5").unwrap(); + http_stream.update_resp_headers = false; + http_stream + .write_response_header(Box::new(new_response)) + .await + .unwrap(); + assert_eq!(http_stream.body_bytes_sent(), 0); + http_stream.write_body(wire_body).await.unwrap(); + assert_eq!(http_stream.body_bytes_sent(), wire_body.len()); + } + #[tokio::test] async fn write_body_http10() { let read_wire = b"GET / HTTP/1.1\r\n\r\n"; @@ -2700,6 +3189,37 @@ Content-Length: 5\r\n\ let reused = http_stream.reuse().await.unwrap(); assert!(reused.is_none()); } + + #[test] + fn test_connection_user_context_set_and_take() { + let mock_io = Builder::new().build(); + let mut session = HttpSession::new(Box::new(mock_io)); + + // Initially no context + assert!(session.take_connection_user_context().is_none()); + + // Set a context + session.set_connection_user_context(Some(Box::new(42u64))); + + // Take it back + let ctx = session.take_connection_user_context(); + assert!(ctx.is_some()); + let val = ctx.unwrap().downcast::().unwrap(); + assert_eq!(*val, 42u64); + + // After take, it's gone + assert!(session.take_connection_user_context().is_none()); + } + + #[test] + fn test_connection_user_context_set_none_clears() { + let mock_io = Builder::new().build(); + let mut session = HttpSession::new(Box::new(mock_io)); + + session.set_connection_user_context(Some(Box::new("hello".to_string()))); + session.set_connection_user_context(None); + assert!(session.take_connection_user_context().is_none()); + } } #[cfg(test)] @@ -2736,25 +3256,30 @@ mod test_sync { } #[cfg(test)] -mod test_timeouts { +mod test_proxy_tasks { use super::*; + use http::StatusCode; use std::future::IntoFuture; use tokio_test::io::{Builder, Mock}; - /// An upper limit for any read within any test to prevent tests from hanging forever if - /// an internal read call never returns, etc. + fn init_log() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + // An upper limit for any read within any test to prevent tests from hanging forever if + // an internal read call never returns, etc. const TEST_MAX_WAIT_FOR_READ: Duration = Duration::from_secs(3); - /// The duration of 600 seconds is chosen to be "effectively forever" for the purpose of testing + // The duration of 600 seconds is chosen to be "effectively forever" for the purpose of testing const TEST_FOREVER_DURATION: Duration = Duration::from_secs(600); - /// The read_timeout to use, when we want to test that a read operation times out + // The read_timeout to use, when we want to test that a read operation times out const TEST_READ_TIMEOUT: Duration = Duration::from_secs(1); #[derive(Debug)] struct ReadBlockedForeverError; - /// Returns a client stream that will "never" send any bytes / return from a read operation + // Returns a client stream that will "never" send any bytes / return from a read operation fn mocked_blocking_headers_forever_stream() -> Box { Box::new(Builder::new().wait(TEST_FOREVER_DURATION).build()) } @@ -2771,8 +3296,8 @@ mod test_timeouts { ) } - /// Helper function to test a read operation with a tokio timeout - /// to prevent tests from hanging forever in case of a bug + // Helper function to test a read operation with a tokio timeout + // to prevent tests from hanging forever in case of a bug async fn test_read_with_tokio_timeout( read_future: F, ) -> Result>, ReadBlockedForeverError> @@ -2816,6 +3341,352 @@ mod test_timeouts { assert!(res.is_ok()); assert_eq!(res.unwrap().unwrap_err().etype(), &ReadTimedout); } + + #[tokio::test] + async fn test_send_proxy_task_and_write() { + init_log(); + + // We need to know exact bytes that will be written + // "HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nhello" + let expected_header = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n"; + let expected_body = b"hello"; + + let mock_io = Builder::new() + .write(expected_header) + .write(expected_body) + .build(); + + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.update_resp_headers = false; // Disable automatic headers + + // Queue header task + let mut header = ResponseHeader::build(StatusCode::OK, Some(5)).unwrap(); + header.insert_header("Content-Length", "5").unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), false)); + + // Queue body task + http_stream.send_proxy_task(HttpTask::Body(Some(Bytes::from("hello")), true)); + + // Write all tasks + let end_stream = http_stream.write_proxy_tasks().await.unwrap(); + assert!(end_stream); + } + + #[tokio::test] + async fn test_proxy_task_with_timeout() { + init_log(); + + let expected_header = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n"; + let expected_body = b"hello"; + + let mock_io = Builder::new() + .write(expected_header) + .write(expected_body) + .build(); + + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.update_resp_headers = false; + http_stream.write_timeout = Some(Duration::from_secs(1)); // Set write timeout + + // Queue tasks + let mut header = ResponseHeader::build(StatusCode::OK, Some(5)).unwrap(); + header.insert_header("Content-Length", "5").unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), false)); + http_stream.send_proxy_task(HttpTask::Body(Some(Bytes::from("hello")), true)); + + // Verify initial state + assert_eq!( + http_stream.body_bytes_sent(), + 0, + "Should start with 0 bytes sent" + ); + + // Write all tasks with timeout + let end_stream = http_stream.write_proxy_tasks().await.unwrap(); + assert!(end_stream); + + // Verify body bytes were counted correctly (not double counted) + assert_eq!( + http_stream.body_bytes_sent(), + 5, + "Should count exactly 5 bytes (application level), not double counted" + ); + } + + // Test that write_proxy_tasks is cancel-safe: if the future is dropped mid-execution, + // unwritten tasks should remain in the queue. + #[tokio::test] + async fn test_proxy_task_cancel_safety() { + init_log(); + + let expected_header = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n"; + // First chunk: "5\r\nhello\r\n" + let expected_chunk1 = b"5\r\nhello\r\n"; + + // Create a mock IO that will write the header and first chunk, + // but will block indefinitely on the second chunk + let mock_io = Builder::new() + .write(expected_header) + .write(expected_chunk1) + .wait(Duration::from_secs(999)) // This will cause timeout + .build(); + + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.update_resp_headers = false; + http_stream.write_timeout = Some(Duration::from_millis(100)); + + // Queue 3 tasks: header + 2 body chunks + let mut header = ResponseHeader::build(StatusCode::OK, None).unwrap(); + header + .insert_header("Transfer-Encoding", "chunked") + .unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), false)); + http_stream.send_proxy_task(HttpTask::Body(Some(Bytes::from("hello")), false)); + http_stream.send_proxy_task(HttpTask::Body(Some(Bytes::from("world")), true)); + + // Verify we have 3 tasks queued + assert_eq!(http_stream.proxy_task_state.tasks.len(), 3); + + // Try to write all tasks - this should timeout while writing the second body chunk + let result = http_stream.write_proxy_tasks().await; + assert!(result.is_err()); + assert_eq!(result.unwrap_err().etype(), &WriteTimedout); + + // With the refactored cancel-safe design: + // - First task (header) was written successfully and removed from queue + // - Second task (first body "hello") was removed and sent to BodyWriter, write succeeded, state cleared + // - Third task (second body "world") was removed and sent to BodyWriter, timed out mid-write + // - The in-progress write state is tracked in current_writer, NOT in the queue + assert_eq!( + http_stream.proxy_task_state.tasks.len(), + 0, + "Queue should be empty - tasks are owned by writers once sent" + ); + + // The task being written should be tracked in current_writer + assert!( + matches!( + http_stream.proxy_task_state.current_writer, + Some(ProxyTaskWriter::WritingBody(_)) + ), + "Should be mid-write of body task - writer owns the 'world' task state" + ); + + // Verify body_bytes_sent only counts the successfully written "hello" (5 bytes) + // not the timed-out "world" + assert_eq!( + http_stream.body_bytes_sent(), + 5, + "Should only count the 5 bytes from 'hello', not the incomplete 'world' write" + ); + + // On next call to write_proxy_tasks(), Step 1 will resume the "world" write + } + + use crate::protocols::http::v1::test_util::FlushTrackingMock; + + // Test that write_continue_response can be called before write_proxy_tasks + // and both work correctly together. + #[tokio::test] + async fn test_continue_response_before_proxy_tasks() { + init_log(); + + // Expected bytes written: + // 1. 100 Continue response + // 2. 200 OK response header + // 3. Body data + let expected_continue = b"HTTP/1.1 100 Continue\r\n\r\n"; + let expected_header = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n"; + let expected_body = b"hello"; + + let mock_io = Builder::new() + .write(expected_continue) + .write(expected_header) + .write(expected_body) + .build(); + + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.update_resp_headers = false; // Disable automatic headers + + // First, write the 100 Continue response + http_stream.write_continue_response().await.unwrap(); + + // Verify that 100 Continue was recorded + assert!( + http_stream.response_written().is_some(), + "100 Continue should be recorded in response_written" + ); + assert_eq!( + http_stream.response_written().unwrap().status, + StatusCode::CONTINUE, + "Should have recorded 100 Continue" + ); + + // Now queue the actual response using proxy tasks + let mut header = ResponseHeader::build(StatusCode::OK, Some(5)).unwrap(); + header.insert_header("Content-Length", "5").unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), false)); + http_stream.send_proxy_task(HttpTask::Body(Some(Bytes::from("hello")), true)); + + // Write all proxy tasks + let end_stream = http_stream.write_proxy_tasks().await.unwrap(); + assert!(end_stream, "Should indicate end of stream"); + + // Verify final response is 200 OK, not 100 Continue + assert_eq!( + http_stream.response_written().unwrap().status, + StatusCode::OK, + "Final response should be 200 OK, overwriting 100 Continue" + ); + } + + #[tokio::test] + async fn test_head_response_with_content_length_flushes() { + init_log(); + + // HEAD request line + headers + let request = b"HEAD / HTTP/1.1\r\nHost: example.com\r\n\r\n"; + let expected_header = b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n"; + + let mock_io = Builder::new().read(request).write(expected_header).build(); + let (flush_mock, flush_count) = FlushTrackingMock::new(mock_io); + let mut http_stream = HttpSession::new(Box::new(flush_mock)); + http_stream.update_resp_headers = false; + + // Read the HEAD request + http_stream.read_request().await.unwrap(); + assert_eq!(http_stream.get_method(), Some(&Method::HEAD)); + + // Queue header with Content-Length (body will be empty for HEAD) + let mut header = ResponseHeader::build(StatusCode::OK, Some(2)).unwrap(); + header.insert_header("Content-Length", "100").unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), true)); + + let flush_before = FlushTrackingMock::flush_count(&flush_count); + let end_stream = http_stream.write_proxy_tasks().await.unwrap(); + let flush_after = FlushTrackingMock::flush_count(&flush_count); + + assert!(end_stream, "HEAD response should be end of stream"); + assert!( + flush_after > flush_before, + "Should flush after writing HEAD response header with Content-Length \ + (body_writer.finished() is true). Got flush_before={flush_before}, \ + flush_after={flush_after}" + ); + } + + #[tokio::test] + async fn test_204_response_with_content_length_flushes() { + init_log(); + + let request = b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"; + let expected_header = b"HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n"; + + let mock_io = Builder::new().read(request).write(expected_header).build(); + let (flush_mock, flush_count) = FlushTrackingMock::new(mock_io); + let mut http_stream = HttpSession::new(Box::new(flush_mock)); + http_stream.update_resp_headers = false; + + http_stream.read_request().await.unwrap(); + + let mut header = ResponseHeader::build(StatusCode::NO_CONTENT, Some(2)).unwrap(); + header.insert_header("Content-Length", "0").unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), true)); + + let flush_before = FlushTrackingMock::flush_count(&flush_count); + let end_stream = http_stream.write_proxy_tasks().await.unwrap(); + let flush_after = FlushTrackingMock::flush_count(&flush_count); + + assert!(end_stream, "204 response should be end of stream"); + assert!( + flush_after > flush_before, + "Should flush after writing 204 response header with Content-Length \ + (body_writer.finished() is true). Got flush_before={flush_before}, \ + flush_after={flush_after}" + ); + } + + #[tokio::test] + #[should_panic( + expected = "Unexpected UpgradedBody task received on un-upgraded downstream session" + )] + async fn test_upgraded_body_on_non_upgraded_session_panics() { + init_log(); + + let request = b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"; + let expected_header = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n"; + // UpgradedBody on a non-upgraded session should panic before writing, + // but if the bug exists, BodyWriter would encode it as a chunk: + let expected_chunk = b"5\r\nhello\r\n"; + let expected_finish = b"0\r\n\r\n"; + + let mock_io = Builder::new() + .read(request) + .write(expected_header) + // If the panic check is missing, the body gets written as a chunk + .write(expected_chunk) + .write(expected_finish) + .build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.update_resp_headers = false; + + http_stream.read_request().await.unwrap(); + assert!( + !http_stream.was_upgraded(), + "Session should NOT be upgraded" + ); + + // Queue a normal header + let mut header = ResponseHeader::build(StatusCode::OK, Some(2)).unwrap(); + header + .insert_header("Transfer-Encoding", "chunked") + .unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), false)); + + // Queue an UpgradedBody task on a non-upgraded session — should panic + http_stream.send_proxy_task(HttpTask::UpgradedBody(Some(Bytes::from("hello")), true)); + + // This should panic before/during the body write + let _ = http_stream.write_proxy_tasks().await; + } + + #[tokio::test] + #[should_panic(expected = "Unexpected Body task received on upgraded downstream session")] + async fn test_body_on_upgraded_session_panics() { + init_log(); + + // Upgrade request + let request = + b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n"; + // 101 Switching Protocols response + let expected_header = + b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n"; + // If the panic check is missing, Body data would be written raw (close-delimited) + let expected_body = b"hello"; + + let mock_io = Builder::new() + .read(request) + .write(expected_header) + .write(expected_body) + .build(); + let mut http_stream = HttpSession::new(Box::new(mock_io)); + http_stream.update_resp_headers = false; + + http_stream.read_request().await.unwrap(); + + // Queue 101 header to complete the upgrade + let mut header = ResponseHeader::build(StatusCode::SWITCHING_PROTOCOLS, Some(3)).unwrap(); + header.insert_header("Upgrade", "websocket").unwrap(); + header.insert_header("Connection", "Upgrade").unwrap(); + http_stream.send_proxy_task(HttpTask::Header(Box::new(header), false)); + + // Queue a regular Body task on what will be an upgraded session — should panic + http_stream.send_proxy_task(HttpTask::Body(Some(Bytes::from("hello")), true)); + + // This should panic (after writing the header, session becomes upgraded, + // then the Body task should be rejected) + let _ = http_stream.write_proxy_tasks().await; + } } #[cfg(test)] @@ -2925,3 +3796,494 @@ mod test_overread { assert!(reused.is_none()); } } + +#[cfg(test)] +mod test_abort_on_close { + use super::*; + use pingora_error::ErrorType; + use tokio_test::io::Builder; + + fn init_log() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + /// Helper: create an HttpSession whose request has been read and body is done, + /// with the mock stream returning EOF on the next read (simulating client FIN). + async fn session_with_eof() -> HttpSession { + let request = b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n"; + let mock_io = Builder::new().read(&request[..]).build(); + let mut s = HttpSession::new(Box::new(mock_io)); + s.read_request().await.unwrap(); + s + } + + #[tokio::test] + async fn default_abort_on_close_returns_error() { + init_log(); + let mut s = session_with_eof().await; + + assert!(s.abort_on_close); + let err = s.read_body_or_idle(true).await.unwrap_err(); + assert_eq!(*err.etype(), ErrorType::ConnectionClosed); + assert!(s.is_half_closed()); + } + + #[tokio::test] + async fn abort_on_close_false_stays_pending() { + init_log(); + let mut s = session_with_eof().await; + s.set_abort_on_close(false); + + let result = tokio::time::timeout( + std::time::Duration::from_millis(50), + s.read_body_or_idle(true), + ) + .await; + + assert!(result.is_err(), "expected timeout (pending), got a result"); + assert!(s.is_half_closed()); + } + + #[tokio::test] + async fn abort_on_close_error_message_before_response() { + init_log(); + let mut s = session_with_eof().await; + + assert!(s.response_written().is_none()); + let err = s.read_body_or_idle(true).await.unwrap_err(); + let msg = format!("{err}"); + assert!( + msg.contains("Prematurely before response header is sent"), + "unexpected error message: {msg}" + ); + } + + #[tokio::test] + async fn abort_on_close_error_message_after_response_header() { + init_log(); + let mut s = session_with_eof().await; + + // Simulate that a response header has already been sent. + let resp = ResponseHeader::build(200, None).unwrap(); + s.response_written = Some(Box::new(resp)); + let err = s.read_body_or_idle(true).await.unwrap_err(); + let msg = format!("{err}"); + assert!( + msg.contains("Prematurely before response body is complete"), + "unexpected error message: {msg}" + ); + } + + #[tokio::test] + async fn no_body_expected_false_reads_body_then_idles() { + init_log(); + let request = b"POST / HTTP/1.1\r\nHost: localhost\r\nContent-Length: 3\r\n\r\n"; + let mock_io = Builder::new().read(&request[..]).read(b"abc").build(); + let mut s = HttpSession::new(Box::new(mock_io)); + s.read_request().await.unwrap(); + + // 1) no_body_expected = false should still read request body while not done. + let body = s.read_body_or_idle(false).await.unwrap().unwrap(); + assert_eq!(body.as_ref(), b"abc"); + assert!(s.is_body_done()); + + // 2) Once body is naturally done, it transitions to idle behavior on the next call. + let err = s.read_body_or_idle(false).await.unwrap_err(); + assert_eq!(*err.etype(), ErrorType::ConnectionClosed); + let msg = format!("{err}"); + assert!( + msg.contains("Prematurely before response header is sent"), + "unexpected error message: {msg}" + ); + } + + #[tokio::test] + async fn set_abort_on_close_toggles() { + init_log(); + let mut s = session_with_eof().await; + + assert!(s.abort_on_close); + s.set_abort_on_close(false); + assert!(!s.abort_on_close); + s.set_abort_on_close(true); + assert!(s.abort_on_close); + } +} + +#[cfg(test)] +mod test_pipelining { + //! Tests for HTTP/1.1 request pipelining support (RFC 9112 §9.3.2). + //! + //! Pipelining is an opt-in behavior: when enabled via + //! [`HttpSession::set_pipelining_enabled`], the session tolerates + //! overread bytes on reuse (they belong to the next request) and a new + //! session can have them fed in via [`HttpSession::set_pipelined_prefix`]. + //! + //! When disabled (default), overread bytes cause [`HttpSession::reuse`] + //! to return `Ok(None)` so the connection closes — the historical + //! behavior preserved for callers that do not opt in. + + use super::*; + use rstest::rstest; + use tokio_test::io::Builder; + + fn init_log() { + let _ = env_logger::builder().is_test(true).try_init(); + } + + /// Default state: pipelining is off. + #[tokio::test] + async fn pipelining_disabled_by_default() { + init_log(); + let mock_io = Builder::new().build(); + let s = HttpSession::new(Box::new(mock_io)); + assert!(!s.pipelining_enabled()); + } + + /// Toggling the pipelining flag is round-trippable. + #[tokio::test] + async fn set_pipelining_enabled_toggles() { + init_log(); + let mock_io = Builder::new().build(); + let mut s = HttpSession::new(Box::new(mock_io)); + assert!(!s.pipelining_enabled()); + s.set_pipelining_enabled(true); + assert!(s.pipelining_enabled()); + s.set_pipelining_enabled(false); + assert!(!s.pipelining_enabled()); + } + + /// When pipelining is disabled (default), overread bytes must cause + /// reuse to return `None`. Pipelining opt-in must not regress that + /// compatibility behavior. + #[rstest] + #[case(true)] // pipelining explicitly off + #[case(false)] // pipelining flag never set + #[tokio::test] + async fn reuse_rejects_overread_when_pipelining_disabled(#[case] explicit_off: bool) { + init_log(); + let request = + b"GET / HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\npipelined_next"; + let mock_io = Builder::new().read(request).build(); + let mut s = HttpSession::new(Box::new(mock_io)); + if explicit_off { + s.set_pipelining_enabled(false); + } + s.read_request().await.unwrap(); + // Overread is captured when body reading initializes — poll the + // body to trigger the init_content_length path. + let _ = s.read_body_bytes().await.unwrap(); + assert!(s.body_reader.has_bytes_overread()); + let reused = s.reuse().await.unwrap(); + assert!( + reused.is_none(), + "reuse must return None without pipelining" + ); + } + + /// When pipelining is enabled and overread bytes are present, + /// reuse returns both the stream and the extracted prefix. + #[tokio::test] + async fn reuse_allows_overread_when_pipelining_enabled() { + init_log(); + let request = + b"GET / HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\npipelined_next"; + let mock_io = Builder::new().read(request).build(); + let mut s = HttpSession::new(Box::new(mock_io)); + s.set_pipelining_enabled(true); + s.read_request().await.unwrap(); + let _ = s.read_body_bytes().await.unwrap(); + assert!(s.body_reader.has_bytes_overread()); + + let reused = s.reuse().await.unwrap().expect("connection reusable"); + let (_stream, prefix) = reused.into_parts(); + let prefix = prefix.expect("overread must be returned as pipelined prefix"); + assert_eq!(prefix.as_ref(), b"pipelined_next"); + } + + /// Same-read pipelining with no prior body poll still extracts the prefix. + #[tokio::test] + async fn reuse_extracts_prefix_without_body_poll() { + init_log(); + let req1 = b"GET /one HTTP/1.1\r\nHost: pingora.org\r\n\r\n"; + let req2 = b"GET /two HTTP/1.1\r\nHost: pingora.org\r\n\r\n"; + let mut combined = Vec::with_capacity(req1.len() + req2.len()); + combined.extend_from_slice(req1); + combined.extend_from_slice(req2); + + let mock_io = Builder::new().read(&combined).build(); + let mut a = HttpSession::new(Box::new(mock_io)); + a.set_pipelining_enabled(true); + a.read_request().await.unwrap(); + assert_eq!(a.req_header().uri.path(), "/one"); + + let reused = a.reuse().await.unwrap().expect("connection reusable"); + let (stream, prefix) = reused.into_parts(); + let prefix = prefix.expect("pipelined prefix must be extracted during reuse"); + assert_eq!(prefix.as_ref(), req2); + + let mut b = HttpSession::new(stream); + b.set_pipelining_enabled(true); + b.set_pipelined_prefix(prefix); + b.read_request() + .await + .unwrap() + .expect("pipelined request must parse"); + assert_eq!(b.req_header().uri.path(), "/two"); + } + + /// Content-Length: 0 has the same extraction requirement as absent length. + #[tokio::test] + async fn reuse_extracts_content_length_zero_prefix_without_body_poll() { + init_log(); + let req1 = b"GET /one HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + let req2 = b"GET /two HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + let mut combined = Vec::with_capacity(req1.len() + req2.len()); + combined.extend_from_slice(req1); + combined.extend_from_slice(req2); + + let mock_io = Builder::new().read(&combined).build(); + let mut a = HttpSession::new(Box::new(mock_io)); + a.set_pipelining_enabled(true); + a.read_request().await.unwrap(); + assert_eq!(a.req_header().uri.path(), "/one"); + + let reused = a.reuse().await.unwrap().expect("connection reusable"); + let (_stream, prefix) = reused.into_parts(); + let prefix = prefix.expect("pipelined prefix must be extracted during reuse"); + assert_eq!(prefix.as_ref(), req2); + } + + /// The new session parses the pipelined prefix as the start of a + /// request without issuing any stream read — the mock_io allows no + /// reads, so if read_request() tried to pull from the stream it would + /// panic. This is the essential pipelining property: a prefix that + /// already contains a complete request is parsed without waiting for + /// additional bytes. + #[tokio::test] + async fn read_request_consumes_complete_prefix_without_stream_read() { + init_log(); + let prefix = b"GET /two HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + // Mock IO that would panic on any read — ensures the parse is + // wholly satisfied by the pipelined prefix. + let mock_io = Builder::new().build(); + let mut s = HttpSession::new(Box::new(mock_io)); + s.set_pipelined_prefix(BytesMut::from(&prefix[..])); + let n = s + .read_request() + .await + .unwrap() + .expect("request must parse from prefix alone"); + assert!(n > 0); + assert_eq!(s.req_header().uri.path(), "/two"); + } + + /// When the prefix is only the beginning of a request, read_request() + /// continues to read from the stream to complete the header. + #[tokio::test] + async fn read_request_falls_through_to_stream_for_partial_prefix() { + init_log(); + let prefix = b"GET /two HTTP/1.1\r\nHost: "; + let rest = b"pingora.org\r\nContent-Length: 0\r\n\r\n"; + let mock_io = Builder::new().read(rest).build(); + let mut s = HttpSession::new(Box::new(mock_io)); + s.set_pipelined_prefix(BytesMut::from(&prefix[..])); + let n = s + .read_request() + .await + .unwrap() + .expect("request must parse across prefix + stream"); + assert!(n > 0); + assert_eq!(s.req_header().uri.path(), "/two"); + } + + /// Body-pump path: request 2's bytes arrive in a SEPARATE read + /// after request 1 has been fully consumed. The proxy's body-pump + /// loop polls the downstream socket via + /// [`HttpSession::read_body_or_idle`]`(true)` while request 1's + /// response is still being written. The idle branch at + /// `read_body_or_idle` currently raises + /// `ConnectError("Sent data after end of body")` when the idle + /// read returns > 0 bytes — which is exactly the shape pipelining + /// traffic takes when requests span TCP segment boundaries. + /// + /// This covers the two-segment pipelining case: request 2's bytes + /// arrive during the proxy's idle poll, not during request 1's body + /// read. The reuse() overread path (already covered by the tests + /// above) never fires because request 2's bytes were never in + /// `body_buf_overread` to begin with. + /// + /// When pipelining is enabled on the session, this branch must + /// NOT raise `ConnectError`. Instead, the byte(s) read by + /// `idle()` must be stashed so the reuse() path can hand them + /// to the next session via the standard `take_body_overread` + /// extractor. `idle()` uses a 1-byte probe buffer, so the + /// overread surface will typically hold 1 byte per idle poll — + /// the remaining bytes of request 2 stay on the underlying + /// stream and are read by the next session's `read_request` + /// (which seeds itself with the pipelined prefix and continues + /// reading from the stream to complete the header). + #[tokio::test] + async fn idle_read_stashes_bytes_when_pipelining_enabled() { + init_log(); + let req1 = b"GET /one HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + // Only the first byte of req2 is queued — the idle-branch + // read in `read_body_or_idle` uses a 1-byte probe buffer, + // so that's all it will consume. The rest of req2 would + // live on the kernel socket buffer in real traffic and be + // drained by the next session. + let req2_first = b"G"; + + // No `.wait(...)` between the two reads — we want the + // second read to be immediately available once the first + // consumer polls. `tokio-test::io::Builder` delivers reads + // one poll at a time regardless, which is what models a + // TCP segment boundary for our purposes. + let mock_io = Builder::new().read(&req1[..]).read(&req2_first[..]).build(); + + let mut s = HttpSession::new(Box::new(mock_io)); + s.set_pipelining_enabled(true); + + // Consume request 1 fully. Body is zero-length so body_done + // is true; no overread is captured in body_buf_overread + // because req2's bytes were NOT in the same read as req1. + s.read_request().await.unwrap(); + assert_eq!(s.req_header().uri.path(), "/one"); + let _ = s.read_body_bytes().await.unwrap(); + assert!(s.is_body_done()); + assert!( + !s.body_reader.has_bytes_overread(), + "precondition: req2 must arrive in a separate read, not as overread on req1" + ); + + // This is the proxy's body-pump poll. Post-fix, the idle + // branch reads the byte, pushes it to the body reader's + // overread surface, and stays pending — signaling the + // body-pump `select!` loop that the downstream has no more + // body activity to wait on (the loop exits via its other + // branches when the upstream response completes). + // + // We assert the *causal* invariant, not a wall-clock one: + // poll the future repeatedly, yielding between polls to + // let the mock I/O stack drain, until either (a) it + // resolves (which is a failure — it MUST stay pending) or + // (b) we observe enough bookkeeping progress to know the + // idle read has completed. The proxy_tasks channel via + // `proxy_tasks_rx` isn't wired in this test, so "enough + // progress" is signaled by tracking `poll_count` alone; + // the actual overread presence is asserted after the + // future is dropped. + // + // Scope the future in an async block so its borrow on `s` + // ends when we exit the block — the body-reader check + // needs a fresh borrow. + { + let fut = s.read_body_or_idle(true); + tokio::pin!(fut); + // Drive the future forward a bounded number of times. + // Under the fix it will always stay Pending; a broken + // fix resolves Ready in the first few polls. + for _ in 0..10 { + match futures::poll!(fut.as_mut()) { + std::task::Poll::Pending => { + tokio::task::yield_now().await; + } + std::task::Poll::Ready(Err(e)) => panic!( + "read_body_or_idle(true) must not raise an error when \ + pipelining is enabled and the idle read returns > 0 bytes \ + (those bytes are the start of pipelined request 2, not \ + illegal trailing body). Got error: {e:?}" + ), + std::task::Poll::Ready(Ok(body)) => panic!( + "read_body_or_idle(true) must stay pending after stashing \ + pipelined bytes (the body-pump `select!` exits via its \ + other branches). Got body: {body:?}" + ), + } + } + // Future still pending — exit the scope, which drops + // `fut` and releases the mutable borrow on `s`. + } + + // The byte must be extractable as overread, so the + // standard reuse() + HttpPersistentSettings pipeline can + // hand it to the next session. + let overread = s + .take_body_overread() + .expect("pipelined request 2 byte must be retrievable as overread"); + assert_eq!( + overread.as_ref(), + req2_first, + "stashed bytes must be the idle-read probe byte from request 2" + ); + } + + /// Symmetric to the test above: pipelining OFF means the idle + /// branch still raises `ConnectError` as it did pre-patch. This + /// preserves upstream behavior for non-adopters. + #[tokio::test] + async fn idle_read_still_raises_when_pipelining_disabled() { + init_log(); + let req1 = b"GET /one HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + // Single byte of req2 — idle-branch read uses a 1-byte probe + // buffer, error path fires, mock is fully drained. + let req2_first = b"G"; + + let mock_io = Builder::new().read(&req1[..]).read(&req2_first[..]).build(); + + let mut s = HttpSession::new(Box::new(mock_io)); + // Leave pipelining at the default (off). + s.read_request().await.unwrap(); + let _ = s.read_body_bytes().await.unwrap(); + assert!(s.is_body_done()); + + let err = s + .read_body_or_idle(true) + .await + .expect_err("pipelining off: idle read > 0 must raise ConnectError"); + assert_eq!( + *err.etype(), + pingora_error::ErrorType::ConnectError, + "non-adopter callers must still see ConnectError on surplus idle bytes" + ); + } + + /// End-to-end: session A finishes with overread, bytes are extracted, + /// session B consumes them via set_pipelined_prefix and parses the + /// pipelined request without reading from the (empty) stream. + #[tokio::test] + async fn pipelined_request_chain_end_to_end() { + init_log(); + + // Session A: read request 1 with pipelined request 2 bytes appended. + let req1 = b"GET /one HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + let req2 = b"GET /two HTTP/1.1\r\nHost: pingora.org\r\nContent-Length: 0\r\n\r\n"; + let mut combined = Vec::with_capacity(req1.len() + req2.len()); + combined.extend_from_slice(req1); + combined.extend_from_slice(req2); + + let mock_io_a = Builder::new().read(&combined).build(); + let mut a = HttpSession::new(Box::new(mock_io_a)); + a.set_pipelining_enabled(true); + a.read_request().await.unwrap(); + assert_eq!(a.req_header().uri.path(), "/one"); + // Poll the body to trigger init_content_length which captures + // the bytes past Content-Length: 0 as overread. + let _ = a.read_body_bytes().await.unwrap(); + assert!(a.body_reader.has_bytes_overread()); + + let overread = a.take_body_overread().expect("overread present"); + + // Session B: construct with an empty stream (pipelined prefix is + // everything we need), feed the overread, parse the next request. + let mock_io_b = Builder::new().build(); + let mut b = HttpSession::new(Box::new(mock_io_b)); + b.set_pipelining_enabled(true); + b.set_pipelined_prefix(overread); + b.read_request() + .await + .unwrap() + .expect("pipelined request must parse"); + assert_eq!(b.req_header().uri.path(), "/two"); + } +} diff --git a/pingora-core/src/protocols/http/v2/mod.rs b/pingora-core/src/protocols/http/v2/mod.rs index 017118079..8f664c9d3 100644 --- a/pingora-core/src/protocols/http/v2/mod.rs +++ b/pingora-core/src/protocols/http/v2/mod.rs @@ -93,13 +93,14 @@ mod test { use h2::frame::*; use http::{HeaderMap, Method, Uri}; use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream}; + use tokio::sync::{oneshot, watch}; use tokio_stream::StreamExt; use pingora_http::{RequestHeader, ResponseHeader}; use pingora_timeout::sleep; use crate::protocols::{ - http::v2::server::{handshake, HttpSession}, + http::v2::server::{self, handshake, HttpSession}, Digest, }; @@ -111,7 +112,10 @@ mod test { // Client handles.push(tokio::spawn(async move { - let conn = crate::connectors::http::v2::handshake(Box::new(client), 500, None) + use crate::connectors::http::v2::H2HandshakeSettings; + let mut settings = H2HandshakeSettings::new(); + settings.max_streams = 500; + let conn = crate::connectors::http::v2::handshake(Box::new(client), settings) .await .unwrap(); @@ -271,4 +275,350 @@ mod test { assert!(handle.await.is_ok()); } } + + #[tokio::test] + async fn test_graceful_shutdown_processes_inflight_stream() { + // HEADERS arrive on the server after the shutdown signal + // fires, but before the client has observed GOAWAY. + let (mut client, server) = duplex(65536); + // Use channels for deterministic timing. + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let (write_headers_tx, write_headers_rx) = oneshot::channel::<()>(); + + let client_handle = tokio::spawn(async move { + client + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") + .await + .unwrap(); + let mut codec: h2::Codec = h2::Codec::new(client); + codec.send(Settings::default().into()).await.unwrap(); + codec.send(Settings::ack().into()).await.unwrap(); + + // Wait until the test has triggered shutdown on the server before + // sending HEADERS. See the function-level comment for why. + write_headers_rx.await.unwrap(); + + let mut headers = Headers::new( + 1.into(), + Pseudo::request( + Method::GET, + Uri::from_static("https://one.one.one.one/"), + None, + ), + HeaderMap::new(), + ); + headers.set_end_headers(); + headers.set_end_stream(); + codec.send(headers.into()).await.unwrap(); + + let mut saw_response = false; + let mut saw_goaway = false; + let _ = pingora_timeout::timeout(Duration::from_secs(5), async { + while let Some(frame) = codec.next().await { + match frame.unwrap() { + h2::frame::Frame::Headers(_) => { + saw_response = true; + } + h2::frame::Frame::GoAway(_) => { + saw_goaway = true; + } + _ => {} + } + if saw_response && saw_goaway { + break; + } + } + }) + .await; + + assert!(saw_response, "expected response for stream 1"); + assert!(saw_goaway, "expected at least one GOAWAY frame"); + }); + + let connection = handshake(Box::new(server), None).await.unwrap(); + let digest = Arc::new(Digest::default()); + + let trigger = tokio::spawn(async move { + sleep(Duration::from_millis(50)).await; + shutdown_tx.send(true).unwrap(); + // Wait long enough that the server task is guaranteed to have + // observed the shutdown signal and committed to its post-shutdown + // path before putting anything on the wire. + sleep(Duration::from_millis(50)).await; + write_headers_tx.send(()).unwrap(); + }); + + let mut session_handles = vec![]; + server::accept_downstream_sessions(connection, digest, shutdown_rx, |mut session| { + session_handles.push(tokio::spawn(async move { + let req = session.req_header(); + assert_eq!(req.method, Method::GET); + let resp = Box::new(ResponseHeader::build(200, None).unwrap()); + session.write_response_header(resp, true).unwrap(); + })); + }) + .await; + + trigger.await.unwrap(); + assert_eq!( + session_handles.len(), + 1, + "expected stream 1 to be surfaced after shutdown_initiated" + ); + for h in session_handles { + h.await.unwrap(); + } + client_handle.await.unwrap(); + } + + #[tokio::test] + async fn test_graceful_shutdown_processes_post_goaway_stream() { + // Client opens stream 1 after it has observed the + // server's GOAWAY frame. Stream 1 is below the GOAWAY(MAX) + // last_stream_id, so per RFC 9113 §6.8 the server must still process + // it. + let (mut client, server) = duplex(65536); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + + let client_handle = tokio::spawn(async move { + client + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") + .await + .unwrap(); + let mut codec: h2::Codec = h2::Codec::new(client); + codec.send(Settings::default().into()).await.unwrap(); + codec.send(Settings::ack().into()).await.unwrap(); + + // Block until the server's GOAWAY is observed. HEADERS for + // stream 1 are sent strictly after this point. + let mut saw_goaway_before_headers = false; + while let Some(frame) = codec.next().await { + if matches!(frame.unwrap(), h2::frame::Frame::GoAway(_)) { + saw_goaway_before_headers = true; + break; + } + } + assert!( + saw_goaway_before_headers, + "expected GOAWAY before sending HEADERS for stream 1" + ); + + let mut headers = Headers::new( + 1.into(), + Pseudo::request( + Method::GET, + Uri::from_static("https://one.one.one.one/"), + None, + ), + HeaderMap::new(), + ); + headers.set_end_headers(); + headers.set_end_stream(); + codec.send(headers.into()).await.unwrap(); + + let mut saw_response_for_stream_1 = false; + let _ = pingora_timeout::timeout(Duration::from_secs(5), async { + while let Some(frame) = codec.next().await { + if let Ok(h2::frame::Frame::Headers(h)) = frame { + if h.stream_id() == 1u32 { + saw_response_for_stream_1 = true; + break; + } + } + } + }) + .await; + assert!( + saw_response_for_stream_1, + "expected response on stream 1 after GOAWAY", + ); + }); + + let connection = handshake(Box::new(server), None).await.unwrap(); + let digest = Arc::new(Digest::default()); + + let trigger = tokio::spawn(async move { + sleep(Duration::from_millis(50)).await; + shutdown_tx.send(true).unwrap(); + }); + + let mut session_handles = vec![]; + server::accept_downstream_sessions(connection, digest, shutdown_rx, |mut session| { + session_handles.push(tokio::spawn(async move { + let resp = Box::new(ResponseHeader::build(200, None).unwrap()); + session.write_response_header(resp, true).unwrap(); + })); + }) + .await; + + trigger.await.unwrap(); + assert_eq!( + session_handles.len(), + 1, + "expected exactly one stream surfaced after GOAWAY" + ); + for h in session_handles { + h.await.unwrap(); + } + client_handle.await.unwrap(); + } + + #[tokio::test] + async fn test_graceful_shutdown_idle_connection_exits_promptly() { + let (mut client, server) = duplex(65536); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + + let client_handle = tokio::spawn(async move { + client + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") + .await + .unwrap(); + let mut codec: h2::Codec = h2::Codec::new(client); + codec.send(Settings::default().into()).await.unwrap(); + codec.send(Settings::ack().into()).await.unwrap(); + + // Wait for the server's GOAWAY, then drop the codec to close the + // connection so the accept loop can exit. + let mut saw_goaway = false; + let _ = pingora_timeout::timeout(Duration::from_secs(3), async { + while let Some(frame) = codec.next().await { + if matches!(frame.unwrap(), h2::frame::Frame::GoAway(_)) { + saw_goaway = true; + break; + } + } + }) + .await; + assert!(saw_goaway, "expected GOAWAY"); + }); + + let connection = handshake(Box::new(server), None).await.unwrap(); + let digest = Arc::new(Digest::default()); + + let trigger = tokio::spawn(async move { + sleep(Duration::from_millis(20)).await; + shutdown_tx.send(true).unwrap(); + }); + + let result = pingora_timeout::timeout( + Duration::from_secs(2), + server::accept_downstream_sessions(connection, digest, shutdown_rx, |_session| { + panic!("did not expect any sessions on an idle connection"); + }), + ) + .await; + assert!(result.is_ok(), "accept loop hung after shutdown"); + + trigger.await.unwrap(); + client_handle.await.unwrap(); + } + + #[tokio::test] + async fn test_graceful_shutdown_refuses_stream_above_last_stream_id() { + // After the server commits to a final last_stream_id and emits the + // closing GOAWAY, any stream the client tries to open above that id + // must be refused. The accept loop must not surface it and must exit + // cleanly via the `Ok(None)` arm of `from_h2_conn`. + let (mut client, server) = duplex(65536); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + + let client_handle = tokio::spawn(async move { + client + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") + .await + .unwrap(); + let mut codec: h2::Codec = h2::Codec::new(client); + codec.send(Settings::default().into()).await.unwrap(); + codec.send(Settings::ack().into()).await.unwrap(); + + // Open stream 1 so the server will commit last_stream_id >= 1 + // when it eventually emits its closing GOAWAY. + let mut headers = Headers::new( + 1.into(), + Pseudo::request( + Method::GET, + Uri::from_static("https://one.one.one.one"), + None, + ), + HeaderMap::new(), + ); + headers.set_end_headers(); + headers.set_end_stream(); + codec.send(headers.into()).await.unwrap(); + + // Drain frames from the server. Break once we've seen the + // response and at least one GOAWAY so the test doesn't race + // its own outer timeout while waiting on a quiet codec. + let mut saw_response = false; + let mut saw_goaway = false; + let _ = pingora_timeout::timeout(Duration::from_secs(3), async { + while let Some(frame) = codec.next().await { + match frame { + Ok(h2::frame::Frame::Headers(h)) if h.stream_id() == 1 => { + saw_response = true; + } + Ok(h2::frame::Frame::GoAway(_)) => { + saw_goaway = true; + } + Ok(_) => {} + Err(_) => break, + } + if saw_response && saw_goaway { + break; + } + } + }) + .await; + assert!(saw_response, "expected response for stream 1"); + assert!(saw_goaway, "expected at least one GOAWAY frame"); + + // Try to open stream 3 (above last_stream_id). The send may + // succeed locally (duplex buffer) or fail (peer half closed); + // either way the server-side codec must not surface the stream. + let mut headers = Headers::new( + 3.into(), + Pseudo::request( + Method::GET, + Uri::from_static("https://one.one.one.one"), + None, + ), + HeaderMap::new(), + ); + headers.set_end_headers(); + headers.set_end_stream(); + let _ = codec.send(headers.into()).await; + }); + + let connection = handshake(Box::new(server), None).await.unwrap(); + let digest = Arc::new(Digest::default()); + + let trigger = tokio::spawn(async move { + sleep(Duration::from_millis(50)).await; + shutdown_tx.send(true).unwrap(); + }); + + let mut session_handles = vec![]; + let result = pingora_timeout::timeout( + Duration::from_secs(5), + server::accept_downstream_sessions(connection, digest, shutdown_rx, |mut session| { + session_handles.push(tokio::spawn(async move { + let resp = Box::new(ResponseHeader::build(200, None).unwrap()); + session.write_response_header(resp, true).unwrap(); + })); + }), + ) + .await; + assert!(result.is_ok(), "accept loop hung after shutdown"); + assert_eq!( + session_handles.len(), + 1, + "only stream 1 may be surfaced; streams above last_stream_id must be refused" + ); + + trigger.await.unwrap(); + for h in session_handles { + h.await.unwrap(); + } + client_handle.await.unwrap(); + } } diff --git a/pingora-core/src/protocols/http/v2/server.rs b/pingora-core/src/protocols/http/v2/server.rs index 363b7357d..ce426b5ae 100644 --- a/pingora-core/src/protocols/http/v2/server.rs +++ b/pingora-core/src/protocols/http/v2/server.rs @@ -34,6 +34,7 @@ use crate::protocols::http::date::get_cached_date; use crate::protocols::http::v1::client::http_req_header_to_wire; use crate::protocols::http::HttpTask; use crate::protocols::{Digest, SocketAddr, Stream}; +use crate::server::ShutdownWatch; use crate::{Error, ErrorType, OrErr, Result}; const BODY_BUF_LIMIT: usize = 1024 * 64; @@ -63,6 +64,77 @@ pub async fn handshake(io: Stream, options: Option) -> Result( + mut conn: H2Connection, + digest: Arc, + mut shutdown: ShutdownWatch, + mut on_session: F, +) where + F: FnMut(HttpSession), +{ + let mut shutdown_initiated = false; + loop { + let h2_stream = if shutdown_initiated { + HttpSession::from_h2_conn(&mut conn, digest.clone()).await + } else { + tokio::select! { + // Poll the shutdown signal first so a concurrent signal is + // observed deterministically. `from_h2_conn` is cancel-safe + // and is polled again on the next iteration. + biased; + _ = shutdown.changed() => { + conn.graceful_shutdown(); + shutdown_initiated = true; + continue; + } + h2_stream = HttpSession::from_h2_conn(&mut conn, digest.clone()) => h2_stream, + } + }; + match h2_stream { + Err(e) => { + // It is common for the client to just disconnect TCP without + // properly closing H2. So we don't log the errors here + debug!("H2 error when accepting new stream {e}"); + return; + } + // None means the connection is ready to be closed + Ok(None) => return, + Ok(Some(session)) => on_session(session), + } + } +} + use futures::task::Context; use futures::task::Poll; use std::pin::Pin; @@ -477,10 +549,23 @@ impl HttpSession { /// Give up the stream abruptly. /// - /// This will send a `INTERNAL_ERROR` stream error to the client + /// This will send an `INTERNAL_ERROR` stream error to the client. pub fn shutdown(&mut self) { + self.shutdown_with_reason(h2::Reason::INTERNAL_ERROR); + } + + /// Give up the stream abruptly with a custom reason. + /// + /// This will send a `RST_STREAM` frame with the given reason to the client. + /// + /// Useful reasons include: + /// - [`h2::Reason::HTTP_1_1_REQUIRED`] - Signal to the client that HTTP/1.1 should be used + /// instead. Per RFC 7540 §9.1.2, clients should retry the request over HTTP/1.1. + /// - [`h2::Reason::CANCEL`] - Indicate the stream is no longer needed. + /// - [`h2::Reason::REFUSED_STREAM`] - Indicate the stream was refused before processing. + pub fn shutdown_with_reason(&mut self, reason: h2::Reason) { if !self.ended { - self.send_response.send_reset(h2::Reason::INTERNAL_ERROR); + self.send_response.send_reset(reason); } } diff --git a/pingora-core/src/protocols/l4/listener.rs b/pingora-core/src/protocols/l4/listener.rs index 7d00005e1..a6055267a 100644 --- a/pingora-core/src/protocols/l4/listener.rs +++ b/pingora-core/src/protocols/l4/listener.rs @@ -67,6 +67,20 @@ impl AsRawSocket for Listener { } impl Listener { + /// Return the local address this listener is bound to. + /// + /// For TCP listeners this is the resolved address (including the + /// OS-assigned port when the listener was bound to port 0). + /// Returns `None` for non-TCP listeners (e.g. Unix domain sockets). + #[cfg(test)] + pub fn local_addr(&self) -> Option { + match self { + Self::Tcp(l) => l.local_addr().ok(), + #[cfg(unix)] + Self::Unix(_) => None, + } + } + /// Accept a connection from the listening endpoint pub async fn accept(&self) -> io::Result { match &self { diff --git a/pingora-core/src/protocols/l4/socket.rs b/pingora-core/src/protocols/l4/socket.rs index 46decd2ff..fd098c348 100644 --- a/pingora-core/src/protocols/l4/socket.rs +++ b/pingora-core/src/protocols/l4/socket.rs @@ -66,7 +66,7 @@ impl SocketAddr { fn from_sockaddr_storage(sock: &SockaddrStorage) -> Option { if let Some(v4) = sock.as_sockaddr_in() { return Some(SocketAddr::Inet(StdSockAddr::V4( - std::net::SocketAddrV4::new(v4.ip().into(), v4.port()), + std::net::SocketAddrV4::new(v4.ip(), v4.port()), ))); } else if let Some(v6) = sock.as_sockaddr_in6() { return Some(SocketAddr::Inet(StdSockAddr::V6( diff --git a/pingora-core/src/protocols/l4/stream.rs b/pingora-core/src/protocols/l4/stream.rs index 4aa70f705..0f89db570 100644 --- a/pingora-core/src/protocols/l4/stream.rs +++ b/pingora-core/src/protocols/l4/stream.rs @@ -179,7 +179,7 @@ impl RawStreamWrapper { #[cfg(target_os = "linux")] enable_rx_ts: false, #[cfg(target_os = "linux")] - reusable_cmsg_space: nix::cmsg_space!(nix::sys::time::TimeSpec), + reusable_cmsg_space: nix::cmsg_space!(nix::sys::socket::Timestamps), } } @@ -242,7 +242,8 @@ impl AsyncRead for RawStreamWrapper { as *mut [u8]) }; let mut iov = [IoSliceMut::new(b)]; - rs_wrapper.reusable_cmsg_space.clear(); + + rs_wrapper.reusable_cmsg_space.fill(0); match s.try_io(Interest::READABLE, || { recvmsg::( @@ -255,7 +256,7 @@ impl AsyncRead for RawStreamWrapper { }) { Ok(r) => { if let Some(ControlMessageOwned::ScmTimestampsns(rtime)) = r - .cmsgs() + .cmsgs()? .find(|i| matches!(i, ControlMessageOwned::ScmTimestampsns(_))) { // The returned timestamp is a real (i.e. not monotonic) timestamp @@ -354,14 +355,69 @@ impl AsRawSocket for RawStreamWrapper { } } -// Large read buffering helps reducing syscalls with little trade-off -// Ssl layer always does "small" reads in 16k (TLS record size) so L4 read buffer helps a lot. -const BUF_READ_SIZE: usize = 64 * 1024; -// Small write buf to match MSS. Too large write buf delays real time communication. -// This buffering effectively implements something similar to Nagle's algorithm. -// The benefit is that user space can control when to flush, where Nagle's can't be controlled. -// And userspace buffering reduce both syscalls and small packets. -const BUF_WRITE_SIZE: usize = 1460; +/// The default L4 read buffer size. +/// +/// Large read buffering helps reducing syscalls with little trade-off. The SSL +/// layer always does "small" reads in 16k chunks (TLS record size), so L4 read +/// buffering helps a lot. +pub const DEFAULT_L4_READ_BUFFER_SIZE: usize = 64 * 1024; + +/// The default L4 write buffer size. +/// +/// Small write buffering matches a typical MSS. Too large a write buffer delays +/// real-time communication. This buffering effectively implements something +/// similar to Nagle's algorithm, but user space can control when to flush. +pub const DEFAULT_L4_WRITE_BUFFER_SIZE: usize = 1460; + +/// L4 [`BufStream`] buffer sizing. +/// +/// Leaving either side as `None` preserves Pingora's default for that side. +/// Setting either side to `Some(0)` disables `BufStream` buffering for that +/// direction. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct L4BufferSettings { + /// Read buffer size in bytes. `None` uses [`DEFAULT_L4_READ_BUFFER_SIZE`]. + pub read: Option, + /// Write buffer size in bytes. `None` uses [`DEFAULT_L4_WRITE_BUFFER_SIZE`]. + pub write: Option, +} + +impl L4BufferSettings { + /// Create settings with both read and write buffer sizes set explicitly. + pub fn new(read: usize, write: usize) -> Self { + Self { + read: Some(read), + write: Some(write), + } + } + + /// Create settings that disable both read and write `BufStream` buffering. + pub fn unbuffered() -> Self { + Self::new(0, 0) + } + + /// Set the read buffer size. + pub fn read(mut self, read: usize) -> Self { + self.read = Some(read); + self + } + + /// Set the write buffer size. + pub fn write(mut self, write: usize) -> Self { + self.write = Some(write); + self + } + + /// Resolved read buffer size after applying defaults. + pub fn read_capacity(&self) -> usize { + self.read.unwrap_or(DEFAULT_L4_READ_BUFFER_SIZE) + } + + /// Resolved write buffer size after applying defaults. + pub fn write_capacity(&self) -> usize { + self.write.unwrap_or(DEFAULT_L4_WRITE_BUFFER_SIZE) + } +} // NOTE: with writer buffering, users need to call flush() to make sure the data is actually // sent. Otherwise data could be stuck in the buffer forever or get lost when stream is closed. @@ -434,7 +490,7 @@ impl Stream { if let RawStream::Tcp(s) = &self.stream_mut().get_mut().stream { let timestamp_options = TimestampingFlag::SOF_TIMESTAMPING_RX_SOFTWARE | TimestampingFlag::SOF_TIMESTAMPING_SOFTWARE; - setsockopt(s.as_raw_fd(), sockopt::Timestamping, ×tamp_options) + setsockopt(&s, sockopt::Timestamping, ×tamp_options) .or_err(InternalError, "failed to set SOF_TIMESTAMPING_RX_SOFTWARE")?; self.stream_mut().get_mut().enable_rx_ts(true); } @@ -447,8 +503,11 @@ impl Stream { Ok(()) } - /// Put Some data back to the head of the stream to be read again - pub(crate) fn rewind(&mut self, data: &[u8]) { + /// Put some data back to the head of the stream to be read again. + /// + /// This is useful when you've read data to detect a protocol (e.g., PROXY protocol) + /// but the data wasn't what you expected, so you need to "unread" it. + pub fn rewind(&mut self, data: &[u8]) { if !data.is_empty() { self.rewind_read_buf.push(data.to_vec()); } @@ -456,13 +515,18 @@ impl Stream { /// Set the buffer of BufStream /// It is only set later because of the malloc overhead in critical accept() path - pub(crate) fn set_buffer(&mut self) { + pub(crate) fn set_buffer(&mut self, buffer: L4BufferSettings) { use std::mem; // Since BufStream doesn't provide an API to adjust the buf directly, // we take the raw stream out of it and put it in a new BufStream with the size we want let stream = mem::take(&mut self.stream); - let stream = - stream.map(|s| BufStream::with_capacity(BUF_READ_SIZE, BUF_WRITE_SIZE, s.into_inner())); + let stream = stream.map(|s| { + BufStream::with_capacity( + buffer.read_capacity(), + buffer.write_capacity(), + s.into_inner(), + ) + }); let _ = mem::replace(&mut self.stream, stream); } } @@ -734,129 +798,6 @@ impl AsyncWrite for Stream { } } -pub mod async_write_vec { - use bytes::Buf; - use futures::ready; - use std::future::Future; - use std::io::IoSlice; - use std::pin::Pin; - use std::task::{Context, Poll}; - use tokio::io; - use tokio::io::AsyncWrite; - - /* - the missing write_buf https://github.com/tokio-rs/tokio/pull/3156#issuecomment-738207409 - https://github.com/tokio-rs/tokio/issues/2610 - In general vectored write is lost when accessing the trait object: Box - */ - - #[must_use = "futures do nothing unless you `.await` or poll them"] - pub struct WriteVec<'a, W, B> { - writer: &'a mut W, - buf: &'a mut B, - } - - #[must_use = "futures do nothing unless you `.await` or poll them"] - pub struct WriteVecAll<'a, W, B> { - writer: &'a mut W, - buf: &'a mut B, - } - - pub trait AsyncWriteVec { - fn poll_write_vec( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - _buf: &mut B, - ) -> Poll>; - - fn write_vec<'a, B>(&'a mut self, src: &'a mut B) -> WriteVec<'a, Self, B> - where - Self: Sized, - B: Buf, - { - WriteVec { - writer: self, - buf: src, - } - } - - fn write_vec_all<'a, B>(&'a mut self, src: &'a mut B) -> WriteVecAll<'a, Self, B> - where - Self: Sized, - B: Buf, - { - WriteVecAll { - writer: self, - buf: src, - } - } - } - - impl Future for WriteVec<'_, W, B> - where - W: AsyncWriteVec + Unpin, - B: Buf, - { - type Output = io::Result; - - fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll> { - let me = &mut *self; - Pin::new(&mut *me.writer).poll_write_vec(ctx, me.buf) - } - } - - impl Future for WriteVecAll<'_, W, B> - where - W: AsyncWriteVec + Unpin, - B: Buf, - { - type Output = io::Result<()>; - - fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll> { - let me = &mut *self; - while me.buf.has_remaining() { - let n = ready!(Pin::new(&mut *me.writer).poll_write_vec(ctx, me.buf))?; - if n == 0 { - return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); - } - } - Poll::Ready(Ok(())) - } - } - - /* from https://github.com/tokio-rs/tokio/blob/master/tokio-util/src/lib.rs#L177 */ - impl AsyncWriteVec for T - where - T: AsyncWrite, - { - fn poll_write_vec( - self: Pin<&mut Self>, - ctx: &mut Context, - buf: &mut B, - ) -> Poll> { - const MAX_BUFS: usize = 64; - - if !buf.has_remaining() { - return Poll::Ready(Ok(0)); - } - - let n = if self.is_write_vectored() { - let mut slices = [IoSlice::new(&[]); MAX_BUFS]; - let cnt = buf.chunks_vectored(&mut slices); - ready!(self.poll_write_vectored(ctx, &slices[..cnt]))? - } else { - ready!(self.poll_write(ctx, buf.chunk()))? - }; - - buf.advance(n); - - Poll::Ready(Ok(n)) - } - } -} - -pub use async_write_vec::AsyncWriteVec; - #[derive(Debug)] struct AccumulatedDuration { total: Duration, diff --git a/pingora-core/src/protocols/tls/s2n/stream.rs b/pingora-core/src/protocols/tls/s2n/stream.rs index 059718eae..3f12ea449 100644 --- a/pingora-core/src/protocols/tls/s2n/stream.rs +++ b/pingora-core/src/protocols/tls/s2n/stream.rs @@ -123,8 +123,9 @@ where T: AsyncRead + AsyncWrite + std::marker::Unpin, { pub fn from_s2n_stream(stream: S2NTlsStream>) -> TlsStream { - let mut timing: TimingDigest = Default::default(); - timing.established_ts = SystemTime::now(); + let timing = TimingDigest { + established_ts: SystemTime::now(), + }; let digest = Some(Arc::new(SslDigest::from_stream(Some(&stream)))); TlsStream { stream, diff --git a/pingora-core/src/server/bootstrap_services.rs b/pingora-core/src/server/bootstrap_services.rs index 10df272f0..74c81d79c 100644 --- a/pingora-core/src/server/bootstrap_services.rs +++ b/pingora-core/src/server/bootstrap_services.rs @@ -18,13 +18,17 @@ use async_trait::async_trait; use log::{debug, error, info}; use parking_lot::Mutex; use std::sync::Arc; -use tokio::sync::{broadcast, Mutex as TokioMutex}; +use tokio::sync::broadcast; #[cfg(feature = "sentry")] use sentry::ClientOptions; +#[cfg(unix)] +use crate::server::daemon::notify_parent_ready_for_fds; #[cfg(unix)] use crate::server::ListenFds; +#[cfg(unix)] +use std::time::Duration; use crate::{ prelude::Opt, @@ -32,19 +36,16 @@ use crate::{ services::{background::BackgroundService, ServiceReadyNotifier}, }; +/// Default timeout for retrying `SIGUSR1` to the parent when it fails with `EPERM`. +#[cfg(unix)] +const DEFAULT_DAEMON_NOTIFY_TIMEOUT: Duration = Duration::from_secs(60); + /// Service that allows the bootstrap process to be delayed until after /// dependencies are ready pub struct BootstrapService { inner: Arc>, } -/// Sentry is typically started as part of the bootstrap process, but if the -/// bootstrap service is used, we want to initialize Sentry before anything else -/// to make sure errors are captured. -pub struct SentryInitService { - inner: Arc>, -} - impl BootstrapService { pub fn new(inner: &Arc>) -> Self { BootstrapService { @@ -53,14 +54,6 @@ impl BootstrapService { } } -impl SentryInitService { - pub fn new(inner: &Arc>) -> Self { - SentryInitService { - inner: Arc::clone(inner), - } - } -} - /// Encapsulation of the data needed to bootstrap the server pub struct Bootstrap { completed: bool, @@ -73,7 +66,17 @@ pub struct Bootstrap { execution_phase_watch: broadcast::Sender, #[cfg(unix)] - listen_fds: Option, + listen_fds: ListenFds, + + /// PID of the original parent process to notify via `SIGUSR1` after bootstrap completes. + /// Set when [`ServerConf::daemon_wait_for_ready`] is `true`. + #[cfg(unix)] + notify_parent_pid: Option, + + /// How long to keep retrying `SIGUSR1` to the parent when it fails with `EPERM`. + /// See [`ServerConf::daemon_notify_timeout_seconds`]. + #[cfg(unix)] + daemon_notify_timeout: std::time::Duration, #[cfg(feature = "sentry")] #[cfg_attr(docsrs, doc(cfg(feature = "sentry")))] @@ -82,6 +85,14 @@ pub struct Bootstrap { /// Panics and other events sentry captures will be sent to this DSN **only /// in release mode** pub sentry: Option, + + /// The Sentry [`ClientInitGuard`](sentry::ClientInitGuard) returned by + /// [`sentry::init`]. + /// + /// This guard must be kept alive for the lifetime of the server, because + /// dropping it flushes and disables the Sentry client. + #[cfg(all(not(debug_assertions), feature = "sentry"))] + sentry_guard: Option, } impl Bootstrap { @@ -102,11 +113,20 @@ impl Bootstrap { upgrade, upgrade_sock, #[cfg(unix)] - listen_fds: None, + listen_fds: Arc::new(Mutex::new(Fds::new())), + #[cfg(unix)] + notify_parent_pid: None, + #[cfg(unix)] + daemon_notify_timeout: conf + .daemon_notify_timeout_seconds + .map(|n| Duration::from_secs(n.get())) + .unwrap_or(DEFAULT_DAEMON_NOTIFY_TIMEOUT), execution_phase_watch: execution_phase_watch.clone(), completed: false, #[cfg(feature = "sentry")] sentry: None, + #[cfg(all(not(debug_assertions), feature = "sentry"))] + sentry_guard: None, } } @@ -115,13 +135,33 @@ impl Bootstrap { self.sentry = sentry_config; } - /// Start sentry based on the configured options. To prevent multiple - /// initializations, this function will consume the sentry configuration - /// stored in the bootstrap - fn start_sentry(&mut self) { - // Only init sentry in release builds - #[cfg(all(not(debug_assertions), feature = "sentry"))] - let _guard = self.sentry.take().map(|opts| sentry::init(opts)); + /// Store the parent process PID to notify via `SIGUSR1` after bootstrap completes. + /// Only relevant when [`ServerConf::daemon_wait_for_ready`] is `true`. + #[cfg(unix)] + pub fn set_notify_parent_pid(&mut self, pid: u32) { + self.notify_parent_pid = Some(pid); + } + + /// Initialize the Sentry client from the configured [`ClientOptions`] and + /// store the resulting guard. + /// + /// The [`ClientOptions`] are preserved (not consumed) so that sentry can be + /// re-initialized after daemonization, when the transport thread spawned by + /// the previous [`sentry::init`] call is lost due to `fork()`. + /// + /// The resulting [`sentry::ClientInitGuard`] is stored in `self` so that it + /// lives as long as the [`Bootstrap`] (and therefore the + /// [`Server`](super::Server)), keeping the Sentry client active for the + /// lifetime of the process. + /// + /// Sentry is only initialized in release builds; in debug builds this is a + /// no-op. + #[cfg(feature = "sentry")] + pub(super) fn start_sentry(&mut self) { + #[cfg(not(debug_assertions))] + { + self.sentry_guard = self.sentry.as_ref().map(|opts| sentry::init(opts.clone())); + } } pub fn bootstrap(&mut self) { @@ -136,14 +176,34 @@ impl Bootstrap { .send(ExecutionPhase::Bootstrap) .ok(); - self.start_sentry(); + // Temporarily initialize sentry if it isn't already active, so that + // errors during fd loading are captured. If sentry was already + // initialized with a persistent guard by `Server::run()` (as in + // the `bootstrap_as_a_service` path), we skip this to avoid + // clobbering the persistent guard. + #[cfg(all(not(debug_assertions), feature = "sentry"))] + let _guard = if self.sentry_guard.is_none() { + self.sentry.as_ref().map(|opts| sentry::init(opts.clone())) + } else { + None + }; if self.test { info!("Server Test passed, exiting"); std::process::exit(0); } - // load fds + // Notify the parent process that it can exit. It might seem like we should load the file + // descriptors from the old process first, but the purpose of this notification is to + // release the parent so that the process managing it (e.g. systemd) can continue and send + // a quit signal to the old process. That quit signal is required before the old process + // will start trying to send its file descriptors to us — so if we called load_fds first, + // we would be guaranteeing a timeout. + #[cfg(unix)] + if let Some(pid) = self.notify_parent_pid { + notify_parent_ready_for_fds(pid, self.daemon_notify_timeout); + } + #[cfg(unix)] match self.load_fds(self.upgrade) { Ok(_) => { @@ -168,17 +228,18 @@ impl Bootstrap { #[cfg(unix)] fn load_fds(&mut self, upgrade: bool) -> Result<(), nix::Error> { - let mut fds = Fds::new(); if upgrade { debug!("Trying to receive socks"); - fds.get_from_sock(self.upgrade_sock.as_str())? + let mut fds = Fds::new(); + fds.get_from_sock(self.upgrade_sock.as_str())?; + // Mutate through the existing Arc so all clones held by services see the update. + *self.listen_fds.lock() = fds; } - self.listen_fds = Some(Arc::new(TokioMutex::new(fds))); Ok(()) } #[cfg(unix)] - pub fn get_fds(&self) -> Option { + pub fn get_fds(&self) -> ListenFds { self.listen_fds.clone() } } @@ -194,15 +255,3 @@ impl BackgroundService for BootstrapService { notifier.notify_ready(); } } - -#[async_trait] -impl BackgroundService for SentryInitService { - async fn start_with_ready_notifier( - &self, - _shutdown: ShutdownWatch, - notifier: ServiceReadyNotifier, - ) { - self.inner.lock().start_sentry(); - notifier.notify_ready(); - } -} diff --git a/pingora-core/src/server/configuration/mod.rs b/pingora-core/src/server/configuration/mod.rs index 020c90fbe..1b0171106 100644 --- a/pingora-core/src/server/configuration/mod.rs +++ b/pingora-core/src/server/configuration/mod.rs @@ -22,12 +22,18 @@ use clap::Parser; use log::{debug, trace}; use pingora_error::{Error, ErrorType::*, OrErr, Result}; +pub use pingora_runtime::RuntimeMetricsPollTimeHistogramScale; +use pingora_runtime::{RuntimeMetricsOpts, RuntimeOpts}; use serde::{Deserialize, Serialize}; use std::ffi::OsString; use std::fs; +use std::num::NonZeroU64; +use std::path::PathBuf; +use std::time::Duration; // default maximum upstream retries for retry-able proxy errors const DEFAULT_MAX_RETRIES: usize = 16; +const MAX_RUNTIME_METRICS_POLL_TIME_HISTOGRAM_BUCKETS: usize = 1024; /// The configuration file /// @@ -59,12 +65,22 @@ pub struct ServerConf { pub user: Option, /// Similar to `user`, the group this process should switch to. pub group: Option, + /// Working directory for the daemonized process. + /// + /// Only applied when `daemon` is `true`; set this to start the daemon from a known cwd. + // TODO: other OS path options should likely be `PathBuf` as well. + pub working_directory: Option, /// How many threads **each** service should get. The threads are not shared across services. pub threads: usize, /// Number of listener tasks to use per fd. This allows for parallel accepts. pub listener_tasks_per_fd: usize, /// Allow work stealing between threads of the same service. Default `true`. pub work_stealing: bool, + /// Enable Tokio's experimental alternative timer on work-stealing service runtimes. + /// + /// Requires building with `--cfg tokio_unstable`. Ignored when + /// [`Self::work_stealing`] is disabled. + pub runtime_enable_alt_timer: bool, /// The path to CA file the SSL library should use. If empty, the default trust store location /// defined by the SSL library will be used. pub ca_file: Option, @@ -116,6 +132,80 @@ pub struct ServerConf { /// The retry interval is 1 second between attempts. /// If not set, defaults to 5 retries. pub upgrade_sock_connect_accept_max_retries: Option, + /// The maximum number of threads in each runtime's blocking thread pool. + /// + /// The blocking pool handles [`tokio::task::spawn_blocking`] tasks. + /// When not set, the tokio default (512) is used. + pub max_blocking_threads: Option, + /// How long, in seconds, idle blocking threads are kept alive before being shut down. + /// + /// When not set, the tokio default (10 seconds) is used. + pub blocking_threads_ttl_seconds: Option, + /// Timeout durations greater than this threshold use Tokio's native timeout instead of + /// Pingora's fast timeout. + /// + /// This avoids retaining long-duration cancelled timers in Pingora's shared timer map until + /// their original deadline. When not set, defaults to 900 seconds (15 minutes). Set to `null` + /// to disable the Tokio fallback. + pub fast_timeout_to_tokio_threshold_seconds: Option, + /// Enable Tokio's poll-time histogram on runtimes created by this server. + /// + /// This adds two timestamp reads to every task poll, so it should be + /// enabled deliberately when investigating runtime latency. Requires + /// building with `--cfg tokio_unstable`. + pub runtime_metrics_poll_time_histogram: bool, + /// Bucket scale for Tokio's poll-time histogram. + /// + /// Ignored unless [`Self::runtime_metrics_poll_time_histogram`] is enabled. + pub runtime_metrics_poll_time_histogram_scale: Option, + /// Width of the first Tokio poll-time histogram bucket in microseconds. + /// + /// Ignored unless [`Self::runtime_metrics_poll_time_histogram`] is enabled. + pub runtime_metrics_poll_time_histogram_resolution_micros: Option, + /// Number of Tokio poll-time histogram buckets. + /// + /// Ignored unless [`Self::runtime_metrics_poll_time_histogram`] is enabled. Memory usage + /// scales with runtimes × workers × buckets, so values above 1024 are rejected. + pub runtime_metrics_poll_time_histogram_buckets: Option, + /// When `daemon` is `true`, controls whether the parent process of the daemon fork waits for + /// the child to signal readiness before exiting. + /// + /// When `false` (default), the parent exits immediately after the daemon fork, matching the + /// traditional daemonization behavior. Systemd will consider the service started as soon as + /// the parent exits, which may be before the child has finished bootstrapping. + /// + /// When `true`, the parent waits (up to [`Self::daemon_ready_timeout_seconds`]) for the child + /// to send `SIGUSR1` after bootstrap completes. This causes systemd to delay any subsequent + /// steps (such as sending `SIGQUIT` to the old process) until the new instance is fully ready + /// to serve traffic. If the child does not signal in time, the parent exits with a non-zero + /// exit code, causing systemd to abort the reload. + pub daemon_wait_for_ready: bool, + /// Timeout in seconds for the parent process to wait for the child to signal readiness during + /// daemonization when [`Self::daemon_wait_for_ready`] is `true`. + /// + /// If the child does not send `SIGUSR1` within this timeout, the parent exits with a non-zero + /// exit code. + /// + /// Defaults to 600 seconds (10 minutes). + pub daemon_ready_timeout_seconds: Option, + /// How long the child process will keep retrying `SIGUSR1` to the parent when the signal + /// fails with a permission error (`EPERM`) during daemonization. + /// + /// After the daemon fork, the parent always drops its credentials to the configured user and + /// group (see [`Self::user`], [`Self::group`]). Because the privilege drop happens after the + /// fork, there is a small window where the child may attempt to signal the parent before the + /// parent has finished changing its credentials. During this window the kernel will reject the + /// signal with `EPERM` because the child and parent are running as different users. The child + /// retries every 100 ms until this timeout elapses. + /// + /// In practice this window is very small, so the default of 60 seconds is far more than + /// enough to account for it. + /// + /// Only retries on `EPERM`; any other error (e.g. `ESRCH` — parent no longer exists) is + /// treated as fatal and logged without retrying. + /// + /// Defaults to 60 seconds. + pub daemon_notify_timeout_seconds: Option, } impl Default for ServerConf { @@ -134,9 +224,11 @@ impl Default for ServerConf { upgrade_sock: "/tmp/pingora_upgrade.sock".to_string(), user: None, group: None, + working_directory: None, threads: 1, listener_tasks_per_fd: 1, work_stealing: true, + runtime_enable_alt_timer: false, upstream_keepalive_pool_size: 128, upstream_connect_offload_threadpools: None, upstream_connect_offload_thread_per_pool: None, @@ -144,6 +236,18 @@ impl Default for ServerConf { graceful_shutdown_timeout_seconds: None, max_retries: DEFAULT_MAX_RETRIES, upgrade_sock_connect_accept_max_retries: None, + max_blocking_threads: None, + blocking_threads_ttl_seconds: None, + fast_timeout_to_tokio_threshold_seconds: Some( + pingora_timeout::fast_timeout::DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD.as_secs(), + ), + runtime_metrics_poll_time_histogram: false, + runtime_metrics_poll_time_histogram_scale: None, + runtime_metrics_poll_time_histogram_resolution_micros: None, + runtime_metrics_poll_time_histogram_buckets: None, + daemon_ready_timeout_seconds: None, + daemon_wait_for_ready: false, + daemon_notify_timeout_seconds: None, } } } @@ -248,10 +352,52 @@ impl ServerConf { } pub fn validate(self) -> Result { - // TODO: do the validation + if self.max_blocking_threads == Some(0) { + return Error::e_explain(ReadError, "max_blocking_threads must be greater than zero"); + } + if self.runtime_metrics_poll_time_histogram_resolution_micros == Some(0) { + return Error::e_explain( + ReadError, + "runtime_metrics_poll_time_histogram_resolution_micros must be greater than zero", + ); + } + if self.runtime_metrics_poll_time_histogram_buckets == Some(0) { + return Error::e_explain( + ReadError, + "runtime_metrics_poll_time_histogram_buckets must be greater than zero", + ); + } + if self + .runtime_metrics_poll_time_histogram_buckets + .is_some_and(|buckets| buckets > MAX_RUNTIME_METRICS_POLL_TIME_HISTOGRAM_BUCKETS) + { + return Error::e_explain( + ReadError, + format!( + "runtime_metrics_poll_time_histogram_buckets must be at most {MAX_RUNTIME_METRICS_POLL_TIME_HISTOGRAM_BUCKETS}" + ), + ); + } Ok(self) } + /// Build the default runtime options derived from this server configuration. + pub fn runtime_opts(&self) -> RuntimeOpts { + RuntimeOpts { + metrics: RuntimeMetricsOpts { + poll_time_histogram: self.runtime_metrics_poll_time_histogram, + poll_time_histogram_scale: self.runtime_metrics_poll_time_histogram_scale, + poll_time_histogram_resolution: self + .runtime_metrics_poll_time_histogram_resolution_micros + .map(Duration::from_micros), + poll_time_histogram_buckets: self.runtime_metrics_poll_time_histogram_buckets, + }, + enable_alt_timer: self.runtime_enable_alt_timer, + #[cfg(feature = "dial9")] + dial9: None, + } + } + pub fn merge_with_opt(&mut self, opt: &Opt) { if opt.daemon { self.daemon = true; @@ -301,9 +447,11 @@ mod tests { upgrade_sock: "".to_string(), user: None, group: None, + working_directory: None, threads: 1, listener_tasks_per_fd: 1, work_stealing: true, + runtime_enable_alt_timer: false, upstream_keepalive_pool_size: 4, upstream_connect_offload_threadpools: None, upstream_connect_offload_thread_per_pool: None, @@ -311,6 +459,18 @@ mod tests { graceful_shutdown_timeout_seconds: None, max_retries: 1, upgrade_sock_connect_accept_max_retries: None, + max_blocking_threads: None, + blocking_threads_ttl_seconds: None, + fast_timeout_to_tokio_threshold_seconds: Some( + pingora_timeout::fast_timeout::DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD.as_secs(), + ), + runtime_metrics_poll_time_histogram: false, + runtime_metrics_poll_time_histogram_scale: None, + runtime_metrics_poll_time_histogram_resolution_micros: None, + runtime_metrics_poll_time_histogram_buckets: None, + daemon_ready_timeout_seconds: None, + daemon_wait_for_ready: false, + daemon_notify_timeout_seconds: None, }; // cargo test -- --nocapture not_a_test_i_cannot_write_yaml_by_hand println!("{}", conf.to_yaml()); @@ -348,5 +508,173 @@ version: 1 assert_eq!(1, conf.version); assert_eq!(DEFAULT_MAX_RETRIES, conf.max_retries); assert_eq!("/tmp/pingora.pid", conf.pid_file); + assert_eq!( + Some(pingora_timeout::fast_timeout::DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD.as_secs()), + conf.fast_timeout_to_tokio_threshold_seconds + ); + } + + #[test] + fn test_runtime_enable_alt_timer_config() { + init_log(); + let conf_str = r#" +--- +version: 1 +runtime_enable_alt_timer: true + "#; + + let conf = ServerConf::from_yaml(conf_str).unwrap(); + assert!(conf.runtime_enable_alt_timer); + } + + #[test] + fn test_runtime_opts_from_config() { + init_log(); + let conf_str = r#" +--- +version: 1 +runtime_enable_alt_timer: true +runtime_metrics_poll_time_histogram: true +runtime_metrics_poll_time_histogram_scale: log +runtime_metrics_poll_time_histogram_resolution_micros: 20 +runtime_metrics_poll_time_histogram_buckets: 16 + "#; + + let conf = ServerConf::from_yaml(conf_str).unwrap(); + let opts = conf.runtime_opts(); + assert!(opts.enable_alt_timer); + assert!(opts.metrics.poll_time_histogram); + assert_eq!( + Some(RuntimeMetricsPollTimeHistogramScale::Log), + opts.metrics.poll_time_histogram_scale + ); + assert_eq!( + Some(Duration::from_micros(20)), + opts.metrics.poll_time_histogram_resolution + ); + assert_eq!(Some(16), opts.metrics.poll_time_histogram_buckets); + #[cfg(feature = "dial9")] + assert!(opts.dial9.is_none()); + } + + #[test] + fn test_working_directory_deserializes_from_yaml_string() { + init_log(); + let conf_str = r#" +--- +version: 1 +daemon: true +working_directory: /var/lib/pingora + "#; + + let conf = ServerConf::from_yaml(conf_str).unwrap(); + assert_eq!( + conf.working_directory.as_deref(), + Some(std::path::Path::new("/var/lib/pingora")) + ); + + let yaml = serde_yaml::to_value(&conf).unwrap(); + assert_eq!( + yaml.get("working_directory"), + Some(&serde_yaml::Value::String("/var/lib/pingora".to_string())) + ); + } + + #[test] + fn test_zero_max_blocking_threads_is_rejected() { + init_log(); + let conf_str = r#" +--- +version: 1 +max_blocking_threads: 0 + "#; + let result = ServerConf::from_yaml(conf_str); + assert!( + result.is_err(), + "max_blocking_threads: 0 should fail validation" + ); + } + + #[test] + fn test_valid_max_blocking_threads() { + init_log(); + let conf_str = r#" +--- +version: 1 +max_blocking_threads: 64 +blocking_threads_ttl_seconds: 30 + "#; + let conf = ServerConf::from_yaml(conf_str).unwrap(); + assert_eq!(Some(64), conf.max_blocking_threads); + assert_eq!(Some(30), conf.blocking_threads_ttl_seconds); + } + + #[test] + fn test_fast_timeout_to_tokio_threshold_config() { + init_log(); + let conf_str = r#" +--- +version: 1 +fast_timeout_to_tokio_threshold_seconds: 120 + "#; + let conf = ServerConf::from_yaml(conf_str).unwrap(); + assert_eq!(Some(120), conf.fast_timeout_to_tokio_threshold_seconds); + } + + #[test] + fn test_fast_timeout_to_tokio_threshold_can_be_disabled() { + init_log(); + let conf_str = r#" +--- +version: 1 +fast_timeout_to_tokio_threshold_seconds: + "#; + let conf = ServerConf::from_yaml(conf_str).unwrap(); + assert_eq!(None, conf.fast_timeout_to_tokio_threshold_seconds); + } + + #[test] + fn test_runtime_poll_time_histogram_config() { + init_log(); + let conf_str = r#" +--- +version: 1 +runtime_metrics_poll_time_histogram: true +runtime_metrics_poll_time_histogram_scale: log +runtime_metrics_poll_time_histogram_resolution_micros: 20 +runtime_metrics_poll_time_histogram_buckets: 16 + "#; + + let conf = ServerConf::from_yaml(conf_str).unwrap(); + assert!(conf.runtime_metrics_poll_time_histogram); + assert_eq!( + Some(RuntimeMetricsPollTimeHistogramScale::Log), + conf.runtime_metrics_poll_time_histogram_scale + ); + assert_eq!( + Some(20), + conf.runtime_metrics_poll_time_histogram_resolution_micros + ); + assert_eq!(Some(16), conf.runtime_metrics_poll_time_histogram_buckets); + } + + #[test] + fn test_runtime_poll_time_histogram_bucket_limit() { + init_log(); + let conf_str = format!( + r#" +--- +version: 1 +runtime_metrics_poll_time_histogram: true +runtime_metrics_poll_time_histogram_buckets: {} + "#, + MAX_RUNTIME_METRICS_POLL_TIME_HISTOGRAM_BUCKETS + 1 + ); + + let result = ServerConf::from_yaml(&conf_str); + assert!( + result.is_err(), + "excessive runtime_metrics_poll_time_histogram_buckets should fail validation" + ); } } diff --git a/pingora-core/src/server/daemon.rs b/pingora-core/src/server/daemon.rs index 7381fc936..54037788f 100644 --- a/pingora-core/src/server/daemon.rs +++ b/pingora-core/src/server/daemon.rs @@ -12,18 +12,71 @@ // See the License for the specific language governing permissions and // limitations under the License. -use daemonize::{Daemonize, Stdio}; -use log::{debug, error}; +use daemonix::{Daemonize, Outcome, Stdio}; +use log::{debug, error, info}; +use pingora_error::{Error, ErrorType, OrErr, Result}; use std::ffi::CString; use std::fs::{self, OpenOptions}; use std::os::unix::prelude::OpenOptionsExt; use std::path::Path; +use std::process; +use std::thread; +use std::time::{Duration, Instant}; use crate::server::configuration::ServerConf; +/// Error returned by [`send_signal`]. +#[derive(Debug)] +pub(crate) enum SignalError { + /// The caller does not have permission to send the signal to the target process (`EPERM`). + PermissionDenied, + /// Any other error from `kill(2)`. Contains the raw `errno` value. + OtherSignalError(i32), +} + +impl std::fmt::Display for SignalError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + SignalError::PermissionDenied => write!(f, "permission denied (EPERM)"), + SignalError::OtherSignalError(errno) => { + write!(f, "kill failed with errno {errno}") + } + } + } +} + +/// Send `signal` to the process identified by `pid`. +/// +/// Returns `Ok(())` on success. On failure, maps `errno` to [`SignalError`]: +/// - `EPERM` → [`SignalError::PermissionDenied`] +/// - anything else → [`SignalError::OtherSignalError`] containing the raw errno value. +fn send_signal(pid: libc::pid_t, signal: libc::c_int) -> Result<(), SignalError> { + // SAFETY: `kill(2)` is safe to call with any pid/signal combination — invalid values + // simply return an error via errno rather than causing undefined behavior. + let ret = unsafe { libc::kill(pid, signal) }; + if ret == 0 { + return Ok(()); + } + let errno = std::io::Error::last_os_error().raw_os_error().unwrap_or(-1); + if errno == libc::EPERM { + Err(SignalError::PermissionDenied) + } else { + Err(SignalError::OtherSignalError(errno)) + } +} + // Utilities to daemonize a pingora server, i.e. run the process in the background, possibly // under a different running user and/or group. +/// Default timeout for the parent to wait for the daemon child to signal readiness. +const DEFAULT_DAEMON_READY_TIMEOUT: Duration = Duration::from_secs(600); + +/// How long to sleep between `SIGUSR1` send attempts when `EPERM` is returned. +const NOTIFY_RETRY_INTERVAL: Duration = Duration::from_millis(100); + +/// How long to sleep between pid-file liveness checks in the async wait loop. +const LIVENESS_CHECK_INTERVAL: Duration = Duration::from_millis(100); + // XXX: this operation should have been done when the old service is exiting. // Now the new pid file just kick the old one out of the way fn move_old_pid(path: &str) { @@ -45,7 +98,15 @@ fn move_old_pid(path: &str) { } } +/// # Safety +/// +/// `name` must be a valid, null-terminated C string. The returned `gid_t` is read from the +/// `passwd` struct returned by `getpwnam(3)`, which points to a static buffer that may be +/// overwritten by subsequent calls to `getpwnam` or `getpwuid`. The caller must not hold the +/// pointer across such calls. unsafe fn gid_for_username(name: &CString) -> Option { + // SAFETY: `name` is a valid CString; `getpwnam` returns a pointer to a static buffer + // or null. We read `pw_gid` immediately and do not retain the pointer. let passwd = libc::getpwnam(name.as_ptr() as *const libc::c_char); if !passwd.is_null() { return Some((*passwd).pw_gid); @@ -53,15 +114,287 @@ unsafe fn gid_for_username(name: &CString) -> Option { None } +/// Drop the parent process's UID to the user specified in [`ServerConf::user`]. +/// +/// The kernel only permits a process to send a signal to another if they share the same UID (or +/// the sender is root). Since the daemon child sends `SIGUSR1` to the parent to signal readiness, +/// the parent must be running as the same UID as the child by the time that signal arrives — +/// otherwise the kernel will reject it with `EPERM`. +/// +/// This function is called in the `Outcome::Parent` path immediately after `execute()` returns, +/// before the parent enters its readiness wait loop, so the parent's UID matches the child's as +/// quickly as possible after the fork. +/// +/// Only the UID is changed; the GID is left as-is. Signal permission checks are based on UID, +/// so changing the GID is not necessary for this purpose. +/// +/// Logs an error and continues if the user cannot be resolved or `setuid` fails — the parent +/// is short-lived and about to exit, so a failed privilege drop is non-fatal. The child's +/// `EPERM` retry window (see [`ServerConf::daemon_notify_timeout_seconds`]) exists precisely to +/// cover the small gap between the fork and the parent completing this UID change. +fn drop_privileges_in_parent(conf: &ServerConf) -> Result<()> { + let Some(user) = conf.user.as_ref() else { + return Ok(()); + }; + + let user_cstr = CString::new(user.as_str()).or_err_with(ErrorType::Custom("Daemon"), || { + format!("drop_privileges_in_parent: user '{user}' invalid") + })?; + + // SAFETY: `user_cstr` is a valid CString. `getpwnam` returns a pointer to a static + // buffer or null. We read `pw_uid` immediately and do not retain the pointer. + let passwd = unsafe { libc::getpwnam(user_cstr.as_ptr() as *const libc::c_char) }; + if passwd.is_null() { + return Error::e_explain( + ErrorType::Custom("Daemon"), + format!("drop_privileges_in_parent: user '{user}' not found"), + ); + } + + // SAFETY: `passwd` was checked for null above. We dereference it once to read `pw_uid`. + let uid = unsafe { (*passwd).pw_uid }; + // SAFETY: `setuid(2)` is safe to call with any uid — invalid values return an error. + let ret = unsafe { libc::setuid(uid) }; + if ret == 0 { + Ok(()) + } else { + Error::e_explain( + ErrorType::Custom("Daemon"), + format!( + "drop_privileges_in_parent: setuid({uid}) failed: {}", + std::io::Error::last_os_error() + ), + ) + } +} + +/// Outcome of calling [`daemonize`]. +/// +/// When [`ServerConf::daemon_wait_for_ready`] is `true`, the child process must call +/// [`notify_parent_ready_for_fds`] after bootstrap completes to unblock the parent's wait loop. +pub struct DaemonizeResult { + /// The PID of the original parent process to notify via `SIGUSR1` after bootstrap completes. + /// + /// `Some` when [`ServerConf::daemon_wait_for_ready`] is `true`, `None` otherwise. + pub notify_parent_pid: Option, +} + /// Start a server instance as a daemon. -#[cfg(unix)] -pub fn daemonize(conf: &ServerConf) { - // TODO: customize working dir +/// +/// Both code paths use [`daemonize::Daemonize::execute()`] rather than calling `fork()` directly. +/// `execute()` returns an [`Outcome`] to the caller in each process rather than having the parent +/// exit inside the crate, which gives us the opportunity to run additional logic in the parent +/// before it exits. +/// +/// When [`ServerConf::daemon_wait_for_ready`] is `false` (the default), the parent exits +/// immediately — matching the behavior of `start()`. +/// +/// When `daemon_wait_for_ready` is `true`, the parent registers a `SIGUSR1` handler before +/// forking, then waits (in a sleep loop polling the pid file and the signal flag) for up to +/// [`ServerConf::daemon_ready_timeout_seconds`] (default 600 s) for the grandchild to send +/// `SIGUSR1`. On success the parent exits with code 0. On timeout, or if the daemon process +/// exits before signaling, the parent exits with code 1, causing systemd to abort the reload. +/// +/// Returns a [`DaemonizeResult`] that is only meaningful to the child process. The parent always +/// exits before returning. +pub fn daemonize(conf: &ServerConf) -> DaemonizeResult { + // Capture the parent PID before forking so it can be passed to the grandchild. The + // grandchild sends SIGUSR1 to this PID after bootstrap completes. + let parent_pid = if conf.daemon_wait_for_ready { + Some(process::id()) + } else { + None + }; + + move_old_pid(&conf.pid_file); + + match build_daemonize(conf).execute() { + Outcome::Parent(result) => { + result.unwrap_or_else(|e| panic!("Daemonize failed: {e}")); + } + Outcome::Child(result) => { + result.unwrap_or_else(|e| panic!("Daemonize child setup failed: {e}")); + return DaemonizeResult { + notify_parent_pid: parent_pid, + }; + } + } + + if conf.daemon_wait_for_ready { + // Drop root privileges before waiting so the parent does not linger as root. + if let Err(e) = drop_privileges_in_parent(conf) { + error!("drop_privileges_in_parent failed: {e}"); + + // Exiting the parent process should be fine because if downgrading + // the user's privileges fails here, it will fail in the child and + // the child will exit too + process::exit(1); + } + + let timeout = conf + .daemon_ready_timeout_seconds + .map(|n| Duration::from_secs(n.get())) + .unwrap_or(DEFAULT_DAEMON_READY_TIMEOUT); + + info!( + "Waiting up to {:?} for daemon to signal readiness via SIGUSR1", + timeout + ); + wait_for_ready_or_exit(&conf.pid_file, timeout); + } + + process::exit(0); +} + +/// Build a single-threaded tokio runtime for the parent's signal wait loop. +/// +/// The parent process is short-lived and only needs to wait for a signal and check the pid file. +/// A current-thread runtime avoids spawning worker threads in a process that is about to exit. +fn build_parent_runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build tokio runtime for parent signal wait") +} + +/// Wait for the daemon grandchild to send `SIGUSR1`, up to `timeout`. +/// +/// Uses a local tokio runtime with [`tokio::signal::unix`] to listen for `SIGUSR1` instead of +/// raw signal handlers and polling loops. The daemon's PID is checked periodically via the pid +/// file — if the process exits before signaling, the parent aborts. +/// +/// Exits the process directly: +/// - exit code 0 if `SIGUSR1` is received (daemon is ready). +/// - exit code 1 if `timeout` elapses (daemon took too long). +/// - exit code 1 if the pid file exists and the process is no longer running. +fn wait_for_ready_or_exit(pid_file: &str, timeout: Duration) { + let rt = build_parent_runtime(); + let pid_file = pid_file.to_owned(); + + rt.block_on(async move { + use tokio::signal::unix::{signal, SignalKind}; + use tokio::time::{interval, timeout as tokio_timeout}; + + let mut sigusr1 = + signal(SignalKind::user_defined1()).expect("failed to register SIGUSR1 listener"); + + let mut liveness_check = interval(LIVENESS_CHECK_INTERVAL); + let mut daemon_pid: Option = None; + + let result = tokio_timeout(timeout, async { + loop { + tokio::select! { + _ = sigusr1.recv() => { + info!("Daemon signaled readiness, parent exiting"); + return; + } + _ = liveness_check.tick() => { + if daemon_pid.is_none() { + daemon_pid = try_read_pid_file(&pid_file); + } + if let Some(pid) = daemon_pid { + if !process_is_running(pid) { + error!( + "Daemon process (pid {pid}) is no longer running \ + before signaling readiness, aborting" + ); + process::exit(1); + } + } + } + } + } + }) + .await; + + if result.is_err() { + error!("Daemon did not signal readiness within {timeout:?}, aborting"); + process::exit(1); + } + }); +} + +/// Notify the parent process that the daemon is ready to serve traffic by sending `SIGUSR1`. +/// +/// Should be called by the daemon process after bootstrap is complete when +/// [`ServerConf::daemon_wait_for_ready`] is `true`. `parent_pid` is the PID of the original +/// process captured before the fork and stored in [`DaemonizeResult::notify_parent_pid`]. +/// +/// `SIGUSR1` sets an atomic flag that the parent's wait loop checks, causing it to exit with +/// code 0 and allowing systemd to proceed with the next step of the service reload. +/// +/// If `kill(2)` returns `EPERM` — which can happen transiently when the child's UID has just +/// been changed by `setuid` and the kernel hasn't yet updated the credential check — the +/// function sleeps for [`NOTIFY_RETRY_INTERVAL`] (100 ms) and retries until `notify_timeout` +/// elapses, at which point it logs an error and returns. Any other error (e.g. `ESRCH`, +/// meaning the parent no longer exists) is logged and the function returns immediately without +/// retrying. +pub fn notify_parent_ready_for_fds(parent_pid: u32, notify_timeout: Duration) { + let parent_pid = parent_pid as libc::pid_t; + info!( + "Sending SIGUSR1 to parent process (pid {}) to signal daemon readiness", + parent_pid + ); + + let start = Instant::now(); + + while start.elapsed() < notify_timeout { + match send_signal(parent_pid, libc::SIGUSR1) { + Ok(()) => return, + Err(SignalError::PermissionDenied) => { + debug!( + "Permission denied sending SIGUSR1 to parent (pid {}), retrying in {:?}", + parent_pid, NOTIFY_RETRY_INTERVAL + ); + thread::sleep(NOTIFY_RETRY_INTERVAL); + } + Err(SignalError::OtherSignalError(errno)) => { + error!( + "Failed to send SIGUSR1 to parent (pid {}): errno {errno}", + parent_pid + ); + return; + } + } + } + + error!( + "Permission denied sending SIGUSR1 to parent (pid {}), giving up after {:?}", + parent_pid, notify_timeout + ); +} + +/// Try to read a PID from `pid_file`. Returns `None` if the file does not exist or cannot be +/// parsed. +fn try_read_pid_file(pid_file: &str) -> Option { + fs::read_to_string(pid_file) + .ok() + .and_then(|c| c.trim().parse().ok()) +} + +/// Returns `true` if a process with `pid` is currently running. +fn process_is_running(pid: libc::pid_t) -> bool { + // Signal 0 does not send a signal; it just checks whether the process exists and whether + // we have permission to signal it. EPERM (no permission) is not possible here because + // drop_privileges_in_parent guarantees the parent has already dropped to the same user as + // the daemon child before this function is called. + send_signal(pid, 0).is_ok() +} + +/// Build a [`Daemonize`] instance configured from `conf`, without calling `start()` or +/// `execute()`. The caller is responsible for driving execution. +fn build_daemonize(conf: &ServerConf) -> Daemonize<()> { let daemonize = Daemonize::new() .umask(0o007) // allow same group to access files but not everyone else .pid_file(&conf.pid_file); + let daemonize = if let Some(working_directory) = conf.working_directory.as_ref() { + daemonize.working_directory(working_directory) + } else { + daemonize + }; + let daemonize = if let Some(error_log) = conf.error_log.as_ref() { let err = OpenOptions::new() .append(true) @@ -82,6 +415,7 @@ pub fn daemonize(conf: &ServerConf) { Some(user) => { let user_cstr = CString::new(user.as_str()).unwrap(); + // SAFETY: `user_cstr` is a valid CString. See `gid_for_username` safety docs. #[cfg(target_os = "macos")] let group_id = unsafe { gid_for_username(&user_cstr).map(|gid| gid as i32) }; #[cfg(target_os = "freebsd")] @@ -92,7 +426,8 @@ pub fn daemonize(conf: &ServerConf) { daemonize .privileged_action(move || { if let Some(gid) = group_id { - // Set the supplemental group privileges for the child process. + // SAFETY: `user_cstr` is a valid CString captured by the closure. + // `initgroups(3)` is safe to call with a valid username and gid. unsafe { libc::initgroups(user_cstr.as_ptr() as *const libc::c_char, gid); } @@ -104,12 +439,8 @@ pub fn daemonize(conf: &ServerConf) { None => daemonize, }; - let daemonize = match conf.group.as_ref() { + match conf.group.as_ref() { Some(group) => daemonize.group(group.as_str()), None => daemonize, - }; - - move_old_pid(&conf.pid_file); - - daemonize.start().unwrap(); // hard crash when fail + } } diff --git a/pingora-core/src/server/mod.rs b/pingora-core/src/server/mod.rs index 406c0d0cd..0956b4475 100644 --- a/pingora-core/src/server/mod.rs +++ b/pingora-core/src/server/mod.rs @@ -27,7 +27,14 @@ use daemon::daemonize; use daggy::NodeIndex; use log::{debug, error, info, warn}; use parking_lot::Mutex; -use pingora_runtime::Runtime; +#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))] +pub use pingora_runtime::Dial9S3UploadOpts; +use pingora_runtime::{BlockingPoolOpts, Runtime, RuntimeBuilder}; +#[cfg(feature = "dial9")] +pub use pingora_runtime::{ + Dial9RuntimeOpts, DEFAULT_DIAL9_MAX_FILE_SIZE, DEFAULT_DIAL9_MAX_TOTAL_SIZE, +}; +pub use pingora_runtime::{RuntimeMetricsOpts, RuntimeOpts}; use pingora_timeout::fast_timeout; #[cfg(feature = "sentry")] use sentry::ClientOptions; @@ -36,11 +43,11 @@ use std::thread; use std::time::SystemTime; #[cfg(unix)] use tokio::signal::unix; -use tokio::sync::{broadcast, watch, Mutex as TokioMutex}; +use tokio::sync::{broadcast, watch}; use tokio::time::{sleep, Duration}; use crate::prelude::background_service; -use crate::server::bootstrap_services::{Bootstrap, BootstrapService, SentryInitService}; +use crate::server::bootstrap_services::{Bootstrap, BootstrapService}; use crate::services::{ DependencyGraph, ServiceHandle, ServiceReadyNotifier, ServiceReadyWatch, ServiceWithDependents, }; @@ -117,7 +124,7 @@ pub enum ExecutionPhase { /// to shutdown pub type ShutdownWatch = watch::Receiver; #[cfg(unix)] -pub type ListenFds = Arc>; +pub type ListenFds = Arc>; /// The type of shutdown process that has been requested. #[derive(Debug)] @@ -197,10 +204,6 @@ impl Default for RunArgs { /// services (see [crate::services]). The server object handles signals, reading configuration, /// zero downtime upgrade and error reporting. pub struct Server { - // This is a way to add services that have to be run before any others - // without requiring dependencies to be set directly - init_services: Vec>, - services: HashMap, shutdown_watch: watch::Sender, // TODO: we many want to drop this copy to let sender call closed() @@ -276,43 +279,47 @@ impl Server { .send(ExecutionPhase::GracefulUpgradeTransferringFds) .ok(); - if let Some(fds) = self.listen_fds() { - let fds = fds.lock().await; - info!("Trying to send socks"); - // XXX: this is blocking IO - match fds.send_to_sock(self.configuration.as_ref().upgrade_sock.as_str()) { - Ok(_) => { - info!("listener sockets sent"); - } - Err(e) => { - error!("Unable to send listener sockets to new process: {e}"); - // sentry log error on fd send failure - #[cfg(all(not(debug_assertions), feature = "sentry"))] - sentry::capture_error(&e); + let sent_fds = { + let fds = self.listen_fds(); + let fds = fds.lock(); + if fds.is_empty() { + info!("No socks to send, shutting down."); + false + } else { + info!("Trying to send socks"); + match fds.send_to_sock(self.configuration.as_ref().upgrade_sock.as_str()) { + Ok(_) => { + info!("listener sockets sent"); + } + Err(e) => { + error!("Unable to send listener sockets to new process: {e}"); + #[cfg(all(not(debug_assertions), feature = "sentry"))] + sentry::capture_error(&e); + } } + true } + }; + if sent_fds { self.execution_phase_watch .send(ExecutionPhase::GracefulUpgradeCloseTimeout) .ok(); sleep(Duration::from_secs(CLOSE_TIMEOUT)).await; - info!("Broadcasting graceful shutdown"); - // gracefully exiting - match self.shutdown_watch.send(true) { - Ok(_) => { - info!("Graceful shutdown started!"); - } - Err(e) => { - error!("Graceful shutdown broadcast failed: {e}"); - // switch to fast shutdown - return ShutdownType::Graceful; - } + } + info!("Broadcasting graceful shutdown"); + // gracefully exiting + match self.shutdown_watch.send(true) { + Ok(_) => { + info!("Graceful shutdown started!"); + } + Err(e) => { + error!("Graceful shutdown broadcast failed: {e}"); + // switch to fast shutdown + return ShutdownType::Graceful; } - info!("Broadcast graceful shutdown complete"); - ShutdownType::Graceful - } else { - info!("No socks to send, shutting down."); - ShutdownType::Graceful } + info!("Broadcast graceful shutdown complete"); + ShutdownType::Graceful } } } @@ -364,25 +371,33 @@ impl Server { /// Get the configured file descriptors for listening #[cfg(unix)] - fn listen_fds(&self) -> Option { + fn listen_fds(&self) -> ListenFds { self.bootstrap.lock().get_fds() } #[allow(clippy::too_many_arguments)] fn run_service( mut service: Box, - #[cfg(unix)] fds: Option, + #[cfg(unix)] fds: ListenFds, shutdown: ShutdownWatch, threads: usize, work_stealing: bool, listeners_per_fd: usize, ready_notifier: ServiceReadyNotifier, dependency_watches: Vec, + blocking_opts: BlockingPoolOpts, + runtime_opts: RuntimeOpts, ) -> Runtime // NOTE: we need to keep the runtime outside async since // otherwise the runtime will be dropped. { - let service_runtime = Server::create_runtime(service.name(), threads, work_stealing); + let service_runtime = Server::create_runtime( + service.name(), + threads, + work_stealing, + blocking_opts, + runtime_opts, + ); let service_name = service.name().to_string(); service_runtime.get_handle().spawn(async move { // Wait for all dependencies to be ready @@ -408,7 +423,7 @@ impl Server { service .start_service( #[cfg(unix)] - fds, + Some(fds), shutdown, listeners_per_fd, ready_notifier, @@ -446,7 +461,6 @@ impl Server { Server { services: Default::default(), - init_services: Default::default(), shutdown_watch: tx, shutdown_recv: rx, execution_phase_watch, @@ -495,7 +509,6 @@ impl Server { Ok(Server { services: Default::default(), - init_services: Default::default(), shutdown_watch: tx, shutdown_recv: rx, execution_phase_watch, @@ -506,31 +519,6 @@ impl Server { }) } - /// Add a service that all other services will wait on before starting. - fn add_init_service(&mut self, service: impl ServiceWithDependents + 'static) { - let boxed_service = Box::new(service); - self.init_services.push(boxed_service); - } - - /// Add the init services as dependencies for all existing services - fn apply_init_service_dependencies(&mut self) { - let services = self - .services - .values() - .map(|service| service.service_handle.clone()) - .collect::>(); - let global_deps = self - .init_services - .drain(..) - .collect::>() - .into_iter() - .map(|dep| self.add_boxed_service(dep)) - .collect::>(); - for service in services { - service.add_dependencies(&global_deps); - } - } - /// Add a service to this server. /// /// Returns a [`ServiceHandle`] that can be used to declare dependencies. @@ -610,20 +598,10 @@ impl Server { /// /// The created service will handle the zero-downtime upgrade from an older version of the server /// to this one. It will try to get all its listening sockets in order to take them over. - /// - /// Other bootstrapping functionality like sentry initialization will also be handled, but as a - /// service that will complete before any other service starts. pub fn bootstrap_as_a_service(&mut self) -> ServiceHandle { let bootstrap_service = background_service("Bootstrap Service", BootstrapService::new(&self.bootstrap)); - let sentry_service = background_service( - "Sentry Init Service", - SentryInitService::new(&self.bootstrap), - ); - - self.add_init_service(sentry_service); - self.add_service(bootstrap_service) } @@ -651,8 +629,6 @@ impl Server { /// Instead it will either start the daemon process and exit, or panic /// if daemonization fails. pub fn run(mut self, run_args: RunArgs) { - self.apply_init_service_dependencies(); - info!("Server starting"); let conf = self.configuration.as_ref(); @@ -661,8 +637,13 @@ impl Server { if conf.daemon { info!("Daemonizing the server"); fast_timeout::pause_for_fork(); - daemonize(&self.configuration); + let daemonize_result = daemonize(&self.configuration); fast_timeout::unpause(); + // If daemon_wait_for_ready is enabled, pass the parent PID to bootstrap so it + // can send SIGUSR1 to the parent after bootstrap completes. + if let Some(pid) = daemonize_result.notify_parent_pid { + self.bootstrap.lock().set_notify_parent_pid(pid); + } } #[cfg(windows)] @@ -670,6 +651,30 @@ impl Server { panic!("Daemonizing under windows is not supported"); } + let blocking_opts = BlockingPoolOpts { + max_threads: conf.max_blocking_threads, + thread_keep_alive: conf.blocking_threads_ttl_seconds.map(Duration::from_secs), + }; + let runtime_opts = conf.runtime_opts(); + if conf.runtime_enable_alt_timer && !conf.work_stealing { + warn!("runtime_enable_alt_timer is ignored when work_stealing is disabled"); + } + // This global timeout threshold is intended to be configured once during server startup, + // before service runtimes begin creating timeout futures. + fast_timeout::set_fast_timeout_to_tokio_threshold( + conf.fast_timeout_to_tokio_threshold_seconds + .map(Duration::from_secs), + ); + + // Initialize (or re-initialize) sentry and persist the guard for + // the lifetime of the server. When daemonizing, the transport + // thread spawned by any earlier `sentry::init` during + // `bootstrap()` is lost after `fork()`, so a fresh init in the + // child process is required. In non-daemon mode this is the + // authoritative initialization that keeps sentry active. + #[cfg(feature = "sentry")] + self.bootstrap.lock().start_sentry(); + // Holds tuples of runtimes and their service name. let mut runtimes: Vec<(Runtime, String)> = Vec::new(); @@ -704,6 +709,10 @@ impl Server { let threads = wrapper.service.threads().unwrap_or(conf.threads); let name = wrapper.service.name().to_string(); + let service_runtime_opts = wrapper + .service + .runtime_opts_override(&runtime_opts) + .unwrap_or_else(|| runtime_opts.clone()); // Extract dependency watches from the ServiceHandle let dependencies = self @@ -743,13 +752,21 @@ impl Server { self.configuration.listener_tasks_per_fd, ready_notifier, dependency_watches, + blocking_opts.clone(), + service_runtime_opts, ); runtimes.push((runtime, name)); } // blocked on main loop so that it runs forever // Only work steal runtime can use block_on() - let server_runtime = Server::create_runtime("Server", 1, true); + let server_runtime = Server::create_runtime( + "Server", + 1, + true, + BlockingPoolOpts::default(), + RuntimeOpts::default(), + ); #[cfg(unix)] let shutdown_type = server_runtime .get_handle() @@ -818,11 +835,17 @@ impl Server { .ok(); } - fn create_runtime(name: &str, threads: usize, work_steal: bool) -> Runtime { - if work_steal { - Runtime::new_steal(threads, name) - } else { - Runtime::new_no_steal(threads, name) - } + fn create_runtime( + name: &str, + threads: usize, + work_steal: bool, + blocking_opts: BlockingPoolOpts, + runtime_opts: RuntimeOpts, + ) -> Runtime { + RuntimeBuilder::new(threads, name) + .work_steal(work_steal) + .blocking_pool_opts(blocking_opts) + .runtime_opts(runtime_opts) + .build() } } diff --git a/pingora-core/src/server/transfer_fd/mod.rs b/pingora-core/src/server/transfer_fd/mod.rs index 3f852aec1..b4f49cca5 100644 --- a/pingora-core/src/server/transfer_fd/mod.rs +++ b/pingora-core/src/server/transfer_fd/mod.rs @@ -16,7 +16,7 @@ use log::{debug, error, warn}; use nix::errno::Errno; #[cfg(target_os = "linux")] -use nix::sys::socket::{self, AddressFamily, RecvMsg, SockFlag, SockType, UnixAddr}; +use nix::sys::socket::{self, AddressFamily, Backlog, RecvMsg, SockFlag, SockType, UnixAddr}; #[cfg(target_os = "linux")] use nix::sys::stat; use nix::{Error, NixPath}; @@ -24,6 +24,8 @@ use std::collections::HashMap; use std::io::Write; #[cfg(target_os = "linux")] use std::io::{IoSlice, IoSliceMut}; +#[cfg(target_os = "linux")] +use std::os::fd::{AsRawFd, BorrowedFd}; use std::os::unix::io::RawFd; #[cfg(target_os = "linux")] use std::{thread, time}; @@ -50,6 +52,10 @@ impl Fds { self.map.get(bind) } + pub fn is_empty(&self) -> bool { + self.map.is_empty() + } + pub fn serialize(&self) -> (Vec, Vec) { self.map.iter().map(|(key, val)| (key.clone(), val)).unzip() } @@ -127,20 +133,27 @@ where // TODO: warn if exist but not able to unlink } }; - socket::bind(listen_fd, &unix_addr).unwrap(); + socket::bind(listen_fd.as_raw_fd(), &unix_addr).unwrap(); /* sock is created before we change user, need to give permission */ stat::fchmodat( - None, + // SAFETY: AT_FDCWD is a well-defined POSIX sentinel constant used by *at() syscalls + // to indicate the current working directory. It is not a real file descriptor and does + // not require ownership or lifetime guarantees. + unsafe { BorrowedFd::borrow_raw(libc::AT_FDCWD) }, path, stat::Mode::from_bits_truncate(0o666), stat::FchmodatFlags::FollowSymlink, ) .unwrap(); - socket::listen(listen_fd, 8).unwrap(); + socket::listen( + &listen_fd, + Backlog::new(8).expect("8 is well within SOMAXCONN"), + ) + .unwrap(); - let fd = match accept_with_retry_timeout(listen_fd, max_retry) { + let fd = match accept_with_retry_timeout(listen_fd.as_raw_fd(), max_retry) { Ok(fd) => fd, Err(e) => { error!("Giving up reading socket from: {path}, error: {e:?}"); @@ -163,7 +176,7 @@ where .unwrap(); let mut fds: Vec = Vec::new(); - for cmsg in msg.cmsgs() { + for cmsg in msg.cmsgs()? { if let socket::ControlMessageOwned::ScmRights(mut vec_fds) = cmsg { fds.append(&mut vec_fds) } else { @@ -250,7 +263,7 @@ where let mut nonblocking_polls = 0; let conn_result: Result = loop { - match socket::connect(send_fd, &unix_addr) { + match socket::connect(send_fd.as_raw_fd(), &unix_addr) { Ok(_) => break Ok(0), Err(e) => match e { /* If the new process hasn't created the upgrade sock we'll get an ENOENT. @@ -295,7 +308,7 @@ where let cmsg = [scm; 1]; loop { match socket::sendmsg( - send_fd, + send_fd.as_raw_fd(), &io_vec, &cmsg, socket::MsgFlags::empty(), @@ -347,6 +360,8 @@ where #[cfg(test)] #[cfg(target_os = "linux")] mod tests { + use std::os::fd::AsRawFd; + use super::*; use log::{debug, error}; @@ -415,7 +430,7 @@ mod tests { assert_eq!(1, buf[31]); }); - let fds = vec![dumb_fd]; + let fds = vec![dumb_fd.as_raw_fd()]; let buf: [u8; 128] = [1; 128]; match send_fds_to(fds, &buf, "/tmp/pingora_fds_receive.sock", None) { Ok(sent) => { @@ -442,7 +457,7 @@ mod tests { None, ) .unwrap(); - fds.add(key1.clone(), dumb_fd1); + fds.add(key1.clone(), dumb_fd1.as_raw_fd()); let key2 = "1.1.1.1:443".to_string(); let dumb_fd2 = socket::socket( AddressFamily::Unix, @@ -451,7 +466,7 @@ mod tests { None, ) .unwrap(); - fds.add(key2.clone(), dumb_fd2); + fds.add(key2.clone(), dumb_fd2.as_raw_fd()); let child = thread::spawn(move || { let mut fds2 = Fds::new(); @@ -478,7 +493,7 @@ mod tests { ) .unwrap(); - let fds = vec![dumb_fd]; + let fds = vec![dumb_fd.as_raw_fd()]; let buf: [u8; 32] = [1; 32]; // Try to send with a custom max_retries of 2 diff --git a/pingora-core/src/services/listening.rs b/pingora-core/src/services/listening.rs index 4be5c4d95..3afb98618 100644 --- a/pingora-core/src/services/listening.rs +++ b/pingora-core/src/services/listening.rs @@ -23,12 +23,12 @@ use crate::listeners::tls::TlsSettings; #[cfg(feature = "connection_filter")] use crate::listeners::AcceptAllFilter; use crate::listeners::{ - ConnectionFilter, Listeners, ServerAddress, TcpSocketOptions, TransportStack, + ConnectionFilter, ListenerConfig, Listeners, ServerAddress, TcpSocketOptions, TransportStack, }; use crate::protocols::Stream; #[cfg(unix)] use crate::server::ListenFds; -use crate::server::ShutdownWatch; +use crate::server::{RuntimeOpts, ShutdownWatch}; use crate::services::Service as ServiceTrait; use async_trait::async_trait; @@ -40,6 +40,9 @@ use std::fs::Permissions; use std::sync::Arc; use std::time::Duration; +/// Override the runtime options used to run a listening service. +pub type RuntimeOptsOverride = Arc Option + Send + Sync>; + /// The type of service that is associated with a list of listening endpoints and a particular application pub struct Service { name: String, @@ -47,6 +50,7 @@ pub struct Service { app_logic: Option, /// The number of preferred threads. `None` to follow global setting. pub threads: Option, + runtime_opts_override: Option, #[cfg(feature = "connection_filter")] connection_filter: Arc, } @@ -59,6 +63,7 @@ impl Service { listeners: Listeners::new(), app_logic: Some(app_logic), threads: None, + runtime_opts_override: None, #[cfg(feature = "connection_filter")] connection_filter: Arc::new(AcceptAllFilter), } @@ -72,11 +77,19 @@ impl Service { listeners, app_logic: Some(app_logic), threads: None, + runtime_opts_override: None, #[cfg(feature = "connection_filter")] connection_filter: Arc::new(AcceptAllFilter), } } + /// Set a runtime options override for this service. + /// + /// Returning [`None`] from the override uses the global runtime options. + pub fn set_runtime_opts_override(&mut self, override_fn: RuntimeOptsOverride) { + self.runtime_opts_override = Some(override_fn); + } + /// Set a custom connection filter for this service. /// /// The connection filter will be applied to all incoming connections @@ -123,6 +136,11 @@ impl Service { self.listeners.add_tcp(addr); } + /// Add a listening endpoint configured by a [`ListenerConfig`]. + pub fn add_listener(&mut self, endpoint: ListenerConfig) { + self.listeners.add_listener(endpoint); + } + /// Add a TCP listening endpoint with the given [`TcpSocketOptions`]. pub fn add_tcp_with_settings(&mut self, addr: &str, sock_opt: TcpSocketOptions) { self.listeners.add_tcp_with_settings(addr, sock_opt); @@ -308,18 +326,10 @@ impl ServiceTrait for Service { fn threads(&self) -> Option { self.threads } -} -use crate::apps::prometheus_http_app::PrometheusServer; - -impl Service { - /// The Prometheus HTTP server - /// - /// The HTTP server endpoint that reports Prometheus metrics collected in the entire service - pub fn prometheus_http_service() -> Self { - Service::new( - "Prometheus metric HTTP".to_string(), - PrometheusServer::new(), - ) + fn runtime_opts_override(&self, global: &RuntimeOpts) -> Option { + self.runtime_opts_override + .as_ref() + .and_then(|override_fn| override_fn(global)) } } diff --git a/pingora-core/src/services/mod.rs b/pingora-core/src/services/mod.rs index 7c4504280..d2be29794 100644 --- a/pingora-core/src/services/mod.rs +++ b/pingora-core/src/services/mod.rs @@ -34,6 +34,7 @@ use tokio::sync::watch; #[cfg(unix)] use crate::server::ListenFds; +use crate::server::RuntimeOpts; use crate::server::ShutdownWatch; pub mod background; @@ -321,6 +322,15 @@ pub trait ServiceWithDependents: Send + Sync { None } + /// Override the runtime options for this service. + /// + /// Returning [`None`] uses the global runtime options derived from + /// [`ServerConf`](crate::server::configuration::ServerConf). + fn runtime_opts_override(&self, global: &RuntimeOpts) -> Option { + let _ = global; + None + } + /// This is currently called to inform the service about the delay it /// experienced from between waiting on its dependencies. Default behavior /// is to log the time. @@ -371,6 +381,10 @@ where S::threads(self) } + fn runtime_opts_override(&self, global: &RuntimeOpts) -> Option { + S::runtime_opts_override(self, global) + } + fn on_startup_delay(&self, time_waited: Duration) { S::on_startup_delay(self, time_waited) } @@ -412,6 +426,15 @@ pub trait Service: Sync + Send { None } + /// Override the runtime options for this service. + /// + /// Returning [`None`] uses the global runtime options derived from + /// [`ServerConf`](crate::server::configuration::ServerConf). + fn runtime_opts_override(&self, global: &RuntimeOpts) -> Option { + let _ = global; + None + } + /// This is currently called to inform the service about the delay it /// experienced from between waiting on its dependencies. Default behavior /// is to log the time. diff --git a/pingora-core/src/tls/mod.rs b/pingora-core/src/tls/mod.rs deleted file mode 100644 index 277b5b409..000000000 --- a/pingora-core/src/tls/mod.rs +++ /dev/null @@ -1,806 +0,0 @@ -// Copyright 2024 Cloudflare, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -//! This module contains a dummy TLS implementation for the scenarios where real TLS -//! implementations are unavailable. - -macro_rules! impl_display { - ($ty:ty) => { - impl std::fmt::Display for $ty { - fn fmt(&self, _f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { - Ok(()) - } - } - }; -} - -macro_rules! impl_deref { - ($from:ty => $to:ty) => { - impl std::ops::Deref for $from { - type Target = $to; - fn deref(&self) -> &$to { - panic!("Not implemented"); - } - } - impl std::ops::DerefMut for $from { - fn deref_mut(&mut self) -> &mut $to { - panic!("Not implemented"); - } - } - }; -} - -pub mod ssl { - use super::error::ErrorStack; - use super::x509::verify::X509VerifyParamRef; - use super::x509::{X509VerifyResult, X509}; - - /// An error returned from an ALPN selection callback. - pub struct AlpnError; - impl AlpnError { - /// Terminate the handshake with a fatal alert. - pub const ALERT_FATAL: AlpnError = Self {}; - - /// Do not select a protocol, but continue the handshake. - pub const NOACK: AlpnError = Self {}; - } - - /// A type which allows for configuration of a client-side TLS session before connection. - pub struct ConnectConfiguration; - impl_deref! {ConnectConfiguration => SslRef} - impl ConnectConfiguration { - /// Configures the use of Server Name Indication (SNI) when connecting. - pub fn set_use_server_name_indication(&mut self, _use_sni: bool) { - panic!("Not implemented"); - } - - /// Configures the use of hostname verification when connecting. - pub fn set_verify_hostname(&mut self, _verify_hostname: bool) { - panic!("Not implemented"); - } - - /// Returns an `Ssl` configured to connect to the provided domain. - pub fn into_ssl(self, _domain: &str) -> Result { - panic!("Not implemented"); - } - - /// Like `SslContextBuilder::set_verify`. - pub fn set_verify(&mut self, _mode: SslVerifyMode) { - panic!("Not implemented"); - } - - /// Like `SslContextBuilder::set_alpn_protos`. - pub fn set_alpn_protos(&mut self, _protocols: &[u8]) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Returns a mutable reference to the X509 verification configuration. - pub fn param_mut(&mut self) -> &mut X509VerifyParamRef { - panic!("Not implemented"); - } - } - - /// An SSL error. - #[derive(Debug)] - pub struct Error; - impl_display!(Error); - impl Error { - pub fn code(&self) -> ErrorCode { - panic!("Not implemented"); - } - } - - /// An error code returned from SSL functions. - #[derive(PartialEq)] - pub struct ErrorCode(i32); - impl ErrorCode { - /// An error occurred in the SSL library. - pub const SSL: ErrorCode = Self(0); - } - - /// An identifier of a session name type. - pub struct NameType; - impl NameType { - pub const HOST_NAME: NameType = Self {}; - } - - /// The state of an SSL/TLS session. - pub struct Ssl; - impl Ssl { - /// Creates a new `Ssl`. - pub fn new(_ctx: &SslContextRef) -> Result { - panic!("Not implemented"); - } - } - impl_deref! {Ssl => SslRef} - - /// A type which wraps server-side streams in a TLS session. - pub struct SslAcceptor; - impl SslAcceptor { - /// Creates a new builder configured to connect to non-legacy clients. This should - /// generally be considered a reasonable default choice. - pub fn mozilla_intermediate_v5( - _method: SslMethod, - ) -> Result { - panic!("Not implemented"); - } - } - - /// A builder for `SslAcceptor`s. - pub struct SslAcceptorBuilder; - impl SslAcceptorBuilder { - /// Consumes the builder, returning a `SslAcceptor`. - pub fn build(self) -> SslAcceptor { - panic!("Not implemented"); - } - - /// Sets the callback used by a server to select a protocol for Application Layer Protocol - /// Negotiation (ALPN). - pub fn set_alpn_select_callback(&mut self, _callback: F) - where - F: for<'a> Fn(&mut SslRef, &'a [u8]) -> Result<&'a [u8], AlpnError> - + 'static - + Sync - + Send, - { - panic!("Not implemented"); - } - - /// Loads a certificate chain from a file. - pub fn set_certificate_chain_file>( - &mut self, - _file: P, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Loads the private key from a file. - pub fn set_private_key_file>( - &mut self, - _file: P, - _file_type: SslFiletype, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Sets the maximum supported protocol version. - pub fn set_max_proto_version( - &mut self, - _version: Option, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - } - - /// Reference to an [`SslCipher`]. - pub struct SslCipherRef; - impl SslCipherRef { - /// Returns the name of the cipher. - pub fn name(&self) -> &'static str { - panic!("Not implemented"); - } - } - - /// A type which wraps client-side streams in a TLS session. - pub struct SslConnector; - impl SslConnector { - /// Creates a new builder for TLS connections. - pub fn builder(_method: SslMethod) -> Result { - panic!("Not implemented"); - } - - /// Returns a structure allowing for configuration of a single TLS session before connection. - pub fn configure(&self) -> Result { - panic!("Not implemented"); - } - - /// Returns a shared reference to the inner raw `SslContext`. - pub fn context(&self) -> &SslContextRef { - panic!("Not implemented"); - } - } - - /// A builder for `SslConnector`s. - pub struct SslConnectorBuilder; - impl SslConnectorBuilder { - /// Consumes the builder, returning an `SslConnector`. - pub fn build(self) -> SslConnector { - panic!("Not implemented"); - } - - /// Sets the list of supported ciphers for protocols before TLSv1.3. - pub fn set_cipher_list(&mut self, _cipher_list: &str) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Sets the context’s supported signature algorithms. - pub fn set_sigalgs_list(&mut self, _sigalgs: &str) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Sets the minimum supported protocol version. - pub fn set_min_proto_version( - &mut self, - _version: Option, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Sets the maximum supported protocol version. - pub fn set_max_proto_version( - &mut self, - _version: Option, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Use the default locations of trusted certificates for verification. - pub fn set_default_verify_paths(&mut self) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Loads trusted root certificates from a file. - pub fn set_ca_file>( - &mut self, - _file: P, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Loads a leaf certificate from a file. - pub fn set_certificate_file>( - &mut self, - _file: P, - _file_type: SslFiletype, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Loads the private key from a file. - pub fn set_private_key_file>( - &mut self, - _file: P, - _file_type: SslFiletype, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Sets the TLS key logging callback. - pub fn set_keylog_callback(&mut self, _callback: F) - where - F: Fn(&SslRef, &str) + 'static + Sync + Send, - { - panic!("Not implemented"); - } - } - - /// A context object for TLS streams. - pub struct SslContext; - impl SslContext { - /// Creates a new builder object for an `SslContext`. - pub fn builder(_method: SslMethod) -> Result { - panic!("Not implemented"); - } - } - impl_deref! {SslContext => SslContextRef} - - /// A builder for `SslContext`s. - pub struct SslContextBuilder; - impl SslContextBuilder { - /// Consumes the builder, returning a new `SslContext`. - pub fn build(self) -> SslContext { - panic!("Not implemented"); - } - } - - /// Reference to [`SslContext`] - pub struct SslContextRef; - - /// An identifier of the format of a certificate or key file. - pub struct SslFiletype; - impl SslFiletype { - /// The PEM format. - pub const PEM: SslFiletype = Self {}; - } - - /// A type specifying the kind of protocol an `SslContext`` will speak. - pub struct SslMethod; - impl SslMethod { - /// Support all versions of the TLS protocol. - pub fn tls() -> SslMethod { - panic!("Not implemented"); - } - } - - /// Reference to an [`Ssl`]. - pub struct SslRef; - impl SslRef { - /// Like [`SslContextBuilder::set_verify`]. - pub fn set_verify(&mut self, _mode: SslVerifyMode) { - panic!("Not implemented"); - } - - /// Returns the current cipher if the session is active. - pub fn current_cipher(&self) -> Option<&SslCipherRef> { - panic!("Not implemented"); - } - - /// Sets the host name to be sent to the server for Server Name Indication (SNI). - pub fn set_hostname(&mut self, _hostname: &str) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Returns the peer’s certificate, if present. - pub fn peer_certificate(&self) -> Option { - panic!("Not implemented"); - } - - /// Returns the certificate verification result. - pub fn verify_result(&self) -> X509VerifyResult { - panic!("Not implemented"); - } - - /// Returns a string describing the protocol version of the session. - pub fn version_str(&self) -> &'static str { - panic!("Not implemented"); - } - - /// Returns the protocol selected via Application Layer Protocol Negotiation (ALPN). - pub fn selected_alpn_protocol(&self) -> Option<&[u8]> { - panic!("Not implemented"); - } - - /// Returns the servername sent by the client via Server Name Indication (SNI). - pub fn servername(&self, _type_: NameType) -> Option<&str> { - panic!("Not implemented"); - } - } - - /// Options controlling the behavior of certificate verification. - pub struct SslVerifyMode; - impl SslVerifyMode { - /// Verifies that the peer’s certificate is trusted. - pub const PEER: Self = Self {}; - - /// Disables verification of the peer’s certificate. - pub const NONE: Self = Self {}; - } - - /// An SSL/TLS protocol version. - pub struct SslVersion; - impl SslVersion { - /// TLSv1.0 - pub const TLS1: SslVersion = Self {}; - - /// TLSv1.2 - pub const TLS1_2: SslVersion = Self {}; - - /// TLSv1.3 - pub const TLS1_3: SslVersion = Self {}; - } - - /// A standard implementation of protocol selection for Application Layer Protocol Negotiation - /// (ALPN). - pub fn select_next_proto<'a>(_server: &[u8], _client: &'a [u8]) -> Option<&'a [u8]> { - panic!("Not implemented"); - } -} - -pub mod ssl_sys { - pub const X509_V_OK: i32 = 0; - pub const X509_V_ERR_INVALID_CALL: i32 = 69; -} - -pub mod error { - use super::ssl::Error; - - /// Collection of [`Errors`] from OpenSSL. - #[derive(Debug)] - pub struct ErrorStack; - impl_display!(ErrorStack); - impl std::error::Error for ErrorStack {} - impl ErrorStack { - /// Returns the contents of the OpenSSL error stack. - pub fn get() -> ErrorStack { - panic!("Not implemented"); - } - - /// Returns the errors in the stack. - pub fn errors(&self) -> &[Error] { - panic!("Not implemented"); - } - } -} - -pub mod x509 { - use super::asn1::{Asn1IntegerRef, Asn1StringRef, Asn1TimeRef}; - use super::error::ErrorStack; - use super::hash::{DigestBytes, MessageDigest}; - use super::nid::Nid; - - /// An `X509` public key certificate. - #[derive(Debug, Clone)] - pub struct X509; - impl_deref! {X509 => X509Ref} - impl X509 { - /// Deserializes a PEM-encoded X509 structure. - pub fn from_pem(_pem: &[u8]) -> Result { - panic!("Not implemented"); - } - } - - /// A type to destructure and examine an `X509Name`. - pub struct X509NameEntries<'a> { - marker: std::marker::PhantomData<&'a ()>, - } - impl<'a> Iterator for X509NameEntries<'a> { - type Item = &'a X509NameEntryRef; - fn next(&mut self) -> Option<&'a X509NameEntryRef> { - panic!("Not implemented"); - } - } - - /// Reference to `X509NameEntry`. - pub struct X509NameEntryRef; - impl X509NameEntryRef { - pub fn data(&self) -> &Asn1StringRef { - panic!("Not implemented"); - } - } - - /// Reference to `X509Name`. - pub struct X509NameRef; - impl X509NameRef { - /// Returns the name entries by the nid. - pub fn entries_by_nid(&self, _nid: Nid) -> X509NameEntries<'_> { - panic!("Not implemented"); - } - } - - /// Reference to `X509`. - pub struct X509Ref; - impl X509Ref { - /// Returns this certificate’s subject name. - pub fn subject_name(&self) -> &X509NameRef { - panic!("Not implemented"); - } - - /// Returns a digest of the DER representation of the certificate. - pub fn digest(&self, _hash_type: MessageDigest) -> Result { - panic!("Not implemented"); - } - - /// Returns the certificate’s Not After validity period. - pub fn not_after(&self) -> &Asn1TimeRef { - panic!("Not implemented"); - } - - /// Returns this certificate’s serial number. - pub fn serial_number(&self) -> &Asn1IntegerRef { - panic!("Not implemented"); - } - } - - /// The result of peer certificate verification. - pub struct X509VerifyResult; - impl X509VerifyResult { - /// Return the integer representation of an `X509VerifyResult`. - pub fn as_raw(&self) -> i32 { - panic!("Not implemented"); - } - } - - pub mod store { - use super::super::error::ErrorStack; - use super::X509; - - /// A builder type used to construct an `X509Store`. - pub struct X509StoreBuilder; - impl X509StoreBuilder { - /// Returns a builder for a certificate store.. - pub fn new() -> Result { - panic!("Not implemented"); - } - - /// Constructs the `X509Store`. - pub fn build(self) -> X509Store { - panic!("Not implemented"); - } - - /// Adds a certificate to the certificate store. - pub fn add_cert(&mut self, _cert: X509) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - } - - /// A certificate store to hold trusted X509 certificates. - pub struct X509Store; - impl_deref! {X509Store => X509StoreRef} - - /// Reference to an `X509Store`. - pub struct X509StoreRef; - } - - pub mod verify { - /// Reference to `X509VerifyParam`. - pub struct X509VerifyParamRef; - } -} - -pub mod nid { - /// A numerical identifier for an OpenSSL object. - pub struct Nid; - impl Nid { - pub const COMMONNAME: Nid = Self {}; - pub const ORGANIZATIONNAME: Nid = Self {}; - pub const ORGANIZATIONALUNITNAME: Nid = Self {}; - } -} - -pub mod pkey { - use super::error::ErrorStack; - - /// A public or private key. - #[derive(Clone)] - pub struct PKey { - marker: std::marker::PhantomData, - } - impl std::ops::Deref for PKey { - type Target = PKeyRef; - fn deref(&self) -> &PKeyRef { - panic!("Not implemented"); - } - } - impl std::ops::DerefMut for PKey { - fn deref_mut(&mut self) -> &mut PKeyRef { - panic!("Not implemented"); - } - } - impl PKey { - pub fn private_key_from_pem(_pem: &[u8]) -> Result, ErrorStack> { - panic!("Not implemented"); - } - } - - /// Reference to `PKey`. - pub struct PKeyRef { - marker: std::marker::PhantomData, - } - - /// A tag type indicating that a key has private components. - #[derive(Clone)] - pub enum Private {} - unsafe impl HasPrivate for Private {} - - /// A trait indicating that a key has private components. - pub unsafe trait HasPrivate {} -} - -pub mod hash { - /// A message digest algorithm. - pub struct MessageDigest; - impl MessageDigest { - pub fn sha256() -> MessageDigest { - panic!("Not implemented"); - } - } - - /// The resulting bytes of a digest. - pub struct DigestBytes; - impl AsRef<[u8]> for DigestBytes { - fn as_ref(&self) -> &[u8] { - panic!("Not implemented"); - } - } -} - -pub mod asn1 { - use super::bn::BigNum; - use super::error::ErrorStack; - - /// A reference to an `Asn1Integer`. - pub struct Asn1IntegerRef; - impl Asn1IntegerRef { - /// Converts the integer to a `BigNum`. - pub fn to_bn(&self) -> Result { - panic!("Not implemented"); - } - } - - /// A reference to an `Asn1String`. - pub struct Asn1StringRef; - impl Asn1StringRef { - pub fn as_utf8(&self) -> Result<&str, ErrorStack> { - panic!("Not implemented"); - } - } - - /// Reference to an `Asn1Time` - pub struct Asn1TimeRef; - impl_display! {Asn1TimeRef} -} - -pub mod bn { - use super::error::ErrorStack; - - /// Dynamically sized large number implementation - pub struct BigNum; - impl BigNum { - /// Returns a hexadecimal string representation of `self`. - pub fn to_hex_str(&self) -> Result<&str, ErrorStack> { - panic!("Not implemented"); - } - } -} - -pub mod ext { - use super::error::ErrorStack; - use super::pkey::{HasPrivate, PKeyRef}; - use super::ssl::{Ssl, SslAcceptor, SslRef}; - use super::x509::store::X509StoreRef; - use super::x509::verify::X509VerifyParamRef; - use super::x509::X509Ref; - - /// Add name as an additional reference identifier that can match the peer's certificate - pub fn add_host(_verify_param: &mut X509VerifyParamRef, _host: &str) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Set the verify cert store of `_ssl` - pub fn ssl_set_verify_cert_store( - _ssl: &mut SslRef, - _cert_store: &X509StoreRef, - ) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Load the certificate into `_ssl` - pub fn ssl_use_certificate(_ssl: &mut SslRef, _cert: &X509Ref) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Load the private key into `_ssl` - pub fn ssl_use_private_key(_ssl: &mut SslRef, _key: &PKeyRef) -> Result<(), ErrorStack> - where - T: HasPrivate, - { - panic!("Not implemented"); - } - - /// Clear the error stack - pub fn clear_error_stack() {} - - /// Create a new [Ssl] from &[SslAcceptor] - pub fn ssl_from_acceptor(_acceptor: &SslAcceptor) -> Result { - panic!("Not implemented"); - } - - /// Suspend the TLS handshake when a certificate is needed. - pub fn suspend_when_need_ssl_cert(_ssl: &mut SslRef) { - panic!("Not implemented"); - } - - /// Unblock a TLS handshake after the certificate is set. - pub fn unblock_ssl_cert(_ssl: &mut SslRef) { - panic!("Not implemented"); - } - - /// Whether the TLS error is SSL_ERROR_WANT_X509_LOOKUP - pub fn is_suspended_for_cert(_error: &super::ssl::Error) -> bool { - panic!("Not implemented"); - } - - /// Add the certificate into the cert chain of `_ssl` - pub fn ssl_add_chain_cert(_ssl: &mut SslRef, _cert: &X509Ref) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Set renegotiation - pub fn ssl_set_renegotiate_mode_freely(_ssl: &mut SslRef) {} - - /// Set the curves/groups of `_ssl` - pub fn ssl_set_groups_list(_ssl: &mut SslRef, _groups: &str) -> Result<(), ErrorStack> { - panic!("Not implemented"); - } - - /// Sets whether a second keyshare to be sent in client hello when PQ is used. - pub fn ssl_use_second_key_share(_ssl: &mut SslRef, _enabled: bool) {} - - /// Get a mutable SslRef ouf of SslRef, which is a missing functionality even when holding &mut SslStream - /// # Safety - pub unsafe fn ssl_mut(_ssl: &SslRef) -> &mut SslRef { - panic!("Not implemented"); - } -} - -pub mod tokio_ssl { - use std::pin::Pin; - use std::task::{Context, Poll}; - use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; - - use super::error::ErrorStack; - use super::ssl::{Error, Ssl, SslRef}; - - /// A TLS session over a stream. - #[derive(Debug)] - pub struct SslStream { - marker: std::marker::PhantomData, - } - impl SslStream { - /// Creates a new `SslStream`. - pub fn new(_ssl: Ssl, _stream: S) -> Result { - panic!("Not implemented"); - } - - /// Initiates a client-side TLS handshake. - pub async fn connect(self: Pin<&mut Self>) -> Result<(), Error> { - panic!("Not implemented"); - } - - /// Initiates a server-side TLS handshake. - pub async fn accept(self: Pin<&mut Self>) -> Result<(), Error> { - panic!("Not implemented"); - } - - /// Returns a shared reference to the `Ssl` object associated with this stream. - pub fn ssl(&self) -> &SslRef { - panic!("Not implemented"); - } - - /// Returns a shared reference to the underlying stream. - pub fn get_ref(&self) -> &S { - panic!("Not implemented"); - } - - /// Returns a mutable reference to the underlying stream. - pub fn get_mut(&mut self) -> &mut S { - panic!("Not implemented"); - } - } - impl AsyncRead for SslStream - where - S: AsyncRead + AsyncWrite, - { - fn poll_read( - self: Pin<&mut Self>, - _ctx: &mut Context<'_>, - _buf: &mut ReadBuf<'_>, - ) -> Poll> { - panic!("Not implemented"); - } - } - impl AsyncWrite for SslStream - where - S: AsyncRead + AsyncWrite, - { - fn poll_write( - self: Pin<&mut Self>, - _ctx: &mut Context<'_>, - _buf: &[u8], - ) -> Poll> { - panic!("Not implemented"); - } - - fn poll_flush(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll> { - panic!("Not implemented"); - } - - fn poll_shutdown( - self: Pin<&mut Self>, - _ctx: &mut Context<'_>, - ) -> Poll> { - panic!("Not implemented"); - } - } -} diff --git a/pingora-core/src/upstreams/peer.rs b/pingora-core/src/upstreams/peer.rs index c9ae0a66d..c7f5e40ca 100644 --- a/pingora-core/src/upstreams/peer.rs +++ b/pingora-core/src/upstreams/peer.rs @@ -33,6 +33,7 @@ use pingora_error::{ }; #[cfg(feature = "s2n")] use pingora_s2n::S2NPolicy; +use std::borrow::Cow; use std::collections::BTreeMap; use std::fmt::{Display, Formatter, Result as FmtResult}; use std::hash::{Hash, Hasher}; @@ -431,8 +432,14 @@ pub struct PeerOptions { pub s2n_security_policy: Option, #[cfg(feature = "s2n")] pub max_blinding_delay: Option, - // how many concurrent h2 stream are allowed in the same connection + /// How many concurrent h2 streams are allowed in the same connection. pub max_h2_streams: usize, + /// Initial per-stream H2 receive window size in bytes. + /// If `None`, the default of 8MB is used. + pub h2_stream_window_size: Option, + /// Initial connection-level H2 receive window size in bytes. + /// If `None`, the default of 8MB is used. + pub h2_connection_window_size: Option, /// Allow invalid Content-Length in HTTP/1 responses (non-RFC compliant). /// /// When enabled, invalid Content-Length responses are treated as close-delimited responses. @@ -441,16 +448,16 @@ pub struct PeerOptions { /// It exists primarily for compatibility with legacy servers that send malformed headers. pub allow_h1_response_invalid_content_length: bool, pub extra_proxy_headers: BTreeMap>, - // The list of curve the tls connection should advertise - // if `None`, the default curves will be used - pub curves: Option<&'static str>, - // see ssl_use_second_key_share + /// The list of curves the tls connection should advertise + /// if `None`, the default curves will be used + pub curves: Option>, + /// see ssl_use_second_key_share pub second_keyshare: bool, - // whether to enable TCP fast open + /// whether to enable TCP fast open pub tcp_fast_open: bool, - // use Arc because Clone is required but not allowed in trait object + /// use Arc because Clone is required but not allowed in trait object pub tracer: Option, - // A custom L4 connector to use to establish new L4 connections + /// A custom L4 connector to use to establish new L4 connections pub custom_l4: Option>, #[derivative(Debug = "ignore")] pub upstream_tcp_sock_tweak_hook: @@ -494,6 +501,8 @@ impl PeerOptions { #[cfg(feature = "s2n")] max_blinding_delay: None, max_h2_streams: 1, + h2_stream_window_size: None, + h2_connection_window_size: None, allow_h1_response_invalid_content_length: false, extra_proxy_headers: BTreeMap::new(), curves: None, @@ -685,6 +694,12 @@ impl Hash for HttpPeer { self.group_key.hash(state); // max h2 stream settings self.options.max_h2_streams.hash(state); + // h2_stream_window_size and h2_connection_window_size are intentionally excluded + // from the reuse hash for now. These are per-connection settings applied at handshake + // time and may be revisited alongside other h2 settings that could be dynamically + // adjusted over the lifetime of a connection. + self.options.curves.hash(state); + self.options.second_keyshare.hash(state); } } diff --git a/pingora-core/src/utils/tls/rustls.rs b/pingora-core/src/utils/tls/rustls.rs index 429b37243..d4e4e5c92 100644 --- a/pingora-core/src/utils/tls/rustls.rs +++ b/pingora-core/src/utils/tls/rustls.rs @@ -101,17 +101,17 @@ pub struct CertKey { certificates: Vec, } -#[self_referencing] +#[self_referencing(pub_extras)] #[derive(Debug)] pub struct WrappedX509 { - raw_cert: Vec, + pub raw_cert: Vec, #[borrows(raw_cert)] #[covariant] - cert: X509Certificate<'this>, + pub cert: X509Certificate<'this>, } -fn parse_x509(raw_cert: &C) -> X509Certificate<'_> +pub fn parse_x509(raw_cert: &C) -> X509Certificate<'_> where C: AsRef<[u8]>, { diff --git a/pingora-core/tests/bootstrap_as_a_service.rs b/pingora-core/tests/bootstrap_as_a_service.rs new file mode 100644 index 000000000..fa88ff205 --- /dev/null +++ b/pingora-core/tests/bootstrap_as_a_service.rs @@ -0,0 +1,136 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Integration tests for `bootstrap_as_a_service`. +//! +//! Verifies that when `bootstrap_as_a_service()` dependencies are declared, the +//! `BootstrapComplete` execution phase is not reached until all dependency services have +//! finished their initialization work. + +use async_trait::async_trait; +use pingora_core::server::ShutdownWatch; +use pingora_core::server::{configuration::ServerConf, ExecutionPhase, RunArgs, Server}; +use pingora_core::services::background::{background_service, BackgroundService}; +use pingora_core::services::ServiceReadyNotifier; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +/// A background service that sets a flag when it completes, after an optional delay. +/// +/// Signals readiness only after completing its initialization work so that dependent +/// services (like `BootstrapService`) cannot start until this service is truly done. +struct TrackableService { + delay: Duration, + completed: Arc, +} + +#[async_trait] +impl BackgroundService for TrackableService { + async fn start_with_ready_notifier( + &self, + _shutdown: ShutdownWatch, + ready_notifier: ServiceReadyNotifier, + ) { + if !self.delay.is_zero() { + tokio::time::sleep(self.delay).await; + } + self.completed.store(true, Ordering::SeqCst); + // Signal readiness only after work is done — this is what the dependency + // mechanism waits on before allowing BootstrapService to proceed. + ready_notifier.notify_ready(); + } +} + +/// Verifies that `bootstrap_as_a_service` does not reach `BootstrapComplete` until all +/// declared dependency services have finished their initialization work. +#[test] +fn test_bootstrap_waits_for_dependencies() { + let conf = ServerConf { + grace_period_seconds: Some(1), + graceful_shutdown_timeout_seconds: Some(1), + ..Default::default() + }; + + let mut server = Server::new_with_opt_and_conf(None, conf); + let mut phase = server.watch_execution_phase(); + + // Two dependency services with delays. The second (150 ms) sets the pace. + let dep1_done = Arc::new(AtomicBool::new(false)); + let dep2_done = Arc::new(AtomicBool::new(false)); + + let dep1_handle = server.add_service(background_service( + "dep1", + TrackableService { + delay: Duration::from_millis(50), + completed: dep1_done.clone(), + }, + )); + let dep2_handle = server.add_service(background_service( + "dep2", + TrackableService { + delay: Duration::from_millis(150), + completed: dep2_done.clone(), + }, + )); + + // BootstrapService must not reach BootstrapComplete until dep1 and dep2 are done. + let bootstrap_handle = server.bootstrap_as_a_service(); + bootstrap_handle.add_dependencies([&dep1_handle, &dep2_handle]); + + // When using bootstrap_as_a_service, do NOT call server.bootstrap() separately — + // the BootstrapService runs as a background service during run(), and emits + // BootstrapComplete only after all its declared dependencies are ready. + let _join = std::thread::spawn(move || { + server.run(RunArgs::default()); + }); + + let mut received_bootstrap = false; + let mut received_bootstrap_complete = false; + + // Collect phases until BootstrapComplete is seen. Running may arrive + // before or after Bootstrap/BootstrapComplete since main_loop starts + // concurrently with the service runtimes. + loop { + match phase.blocking_recv() { + Ok(ExecutionPhase::Bootstrap) => { + received_bootstrap = true; + } + Ok(ExecutionPhase::BootstrapComplete) => { + // Both dependencies must have set their flags before bootstrap completes. + assert!( + dep1_done.load(Ordering::SeqCst), + "dep1 should be done before BootstrapComplete" + ); + assert!( + dep2_done.load(Ordering::SeqCst), + "dep2 should be done before BootstrapComplete" + ); + received_bootstrap_complete = true; + break; + } + Ok(_) => {} + Err(_) => break, + } + } + + assert!(received_bootstrap, "should have seen Bootstrap phase"); + assert!( + received_bootstrap_complete, + "should have seen BootstrapComplete phase" + ); + + // Shut down cleanly. + std::process::exit(0); +} diff --git a/pingora-core/tests/test_basic.rs b/pingora-core/tests/test_basic.rs index 0c9f87f9b..ae6ee810b 100644 --- a/pingora-core/tests/test_basic.rs +++ b/pingora-core/tests/test_basic.rs @@ -14,6 +14,8 @@ mod utils; +#[cfg(all(unix, feature = "any_tls"))] +use hyper_util::client::legacy::Client; #[cfg(all(unix, feature = "any_tls"))] use hyperlocal::{UnixClientExt, Uri}; @@ -55,8 +57,49 @@ async fn test_https_http2() { async fn test_uds() { utils::init(); let url = Uri::new("/tmp/echo.sock", "/").into(); - let client = hyper::Client::unix(); + let client: Client> = + Client::unix(); let res = client.get(url).await.unwrap(); assert_eq!(res.status(), reqwest::StatusCode::OK); } + +#[cfg(feature = "any_tls")] +#[tokio::test] +async fn test_h1_tls_with_h2c_enabled() { + utils::init(); + + let client = reqwest::Client::builder() + .danger_accept_invalid_certs(true) + .http1_only() + .build() + .unwrap(); + + let res = client.get("https://127.0.0.1:6161").send().await.unwrap(); + assert_eq!(res.status(), reqwest::StatusCode::OK); + assert_eq!(res.version(), reqwest::Version::HTTP_11); +} + +#[cfg(feature = "any_tls")] +#[tokio::test] +async fn test_h2_tls_with_h2c_enabled() { + utils::init(); + + let client = reqwest::Client::builder() + .danger_accept_invalid_certs(true) + .build() + .unwrap(); + + let res = client.get("https://127.0.0.1:6161").send().await.unwrap(); + assert_eq!(res.status(), reqwest::StatusCode::OK); + assert_eq!(res.version(), reqwest::Version::HTTP_2); +} + +#[tokio::test] +async fn test_h2c_tcp_still_works() { + utils::init(); + + let res = reqwest::get("http://127.0.0.1:6160").await.unwrap(); + assert_eq!(res.status(), reqwest::StatusCode::OK); + assert_eq!(res.version(), reqwest::Version::HTTP_11); +} diff --git a/pingora-core/tests/utils/mod.rs b/pingora-core/tests/utils/mod.rs index a5016c0b3..a8d526b7e 100644 --- a/pingora-core/tests/utils/mod.rs +++ b/pingora-core/tests/utils/mod.rs @@ -89,7 +89,20 @@ fn entry_point(opt: Option) { let echo_service_http = Service::with_listeners("Echo Service HTTP".to_string(), listeners, EchoApp); + // Echo service with h2c enabled + TLS listener (for testing h2c + TLS interaction) + let mut h2c_tls_settings = + pingora_core::listeners::tls::TlsSettings::intermediate(&cert_path, &key_path).unwrap(); + h2c_tls_settings.enable_h2(); + let mut h2c_listeners = Listeners::tcp("0.0.0.0:6160"); + h2c_listeners.add_tls_with_settings("0.0.0.0:6161", None, h2c_tls_settings); + let mut h2c_app = pingora_core::apps::http_app::HttpServer::new_app(EchoApp); + h2c_app.server_options.get_or_insert_default().h2c = true; + + let echo_service_h2c = + Service::with_listeners("Echo Service H2C".to_string(), h2c_listeners, h2c_app); + my_server.add_service(echo_service_http); + my_server.add_service(echo_service_h2c); my_server.run_forever(); } diff --git a/pingora-http/src/lib.rs b/pingora-http/src/lib.rs index 954be81b1..bdca3c65e 100644 --- a/pingora-http/src/lib.rs +++ b/pingora-http/src/lib.rs @@ -771,7 +771,7 @@ mod tests { #[test] fn test_single_header() { - let mut req = RequestHeader::build("GET", b"\\", None).unwrap(); + let mut req = RequestHeader::build("GET", b"/", None).unwrap(); req.insert_header("foo", "bar").unwrap(); req.insert_header("FoO", "Bar").unwrap(); let mut buf: Vec = vec![]; @@ -833,7 +833,7 @@ mod tests { #[test] fn test_multiple_header() { - let mut req = RequestHeader::build("GET", b"\\", None).unwrap(); + let mut req = RequestHeader::build("GET", b"/", None).unwrap(); req.append_header("FoO", "Bar").unwrap(); req.append_header("fOO", "bar").unwrap(); req.append_header("BAZ", "baR").unwrap(); diff --git a/pingora-load-balancing/Cargo.toml b/pingora-load-balancing/Cargo.toml index d6f5d41e4..f9785be05 100644 --- a/pingora-load-balancing/Cargo.toml +++ b/pingora-load-balancing/Cargo.toml @@ -27,7 +27,7 @@ arc-swap = "1" fnv = "1" rand = "0.8" tokio = { workspace = true } -futures = "0" +futures = { workspace = true } log = { workspace = true } http = { workspace = true } derivative.workspace = true diff --git a/pingora-load-balancing/src/health_check.rs b/pingora-load-balancing/src/health_check.rs index 5e97fb369..3292385ab 100644 --- a/pingora-load-balancing/src/health_check.rs +++ b/pingora-load-balancing/src/health_check.rs @@ -289,15 +289,16 @@ where let session = self.connector.get_http_session(&peer).await?; let mut session = session.0; + + session.set_write_timeout(peer.options.write_timeout); + let req = Box::new(self.req.clone()); session.write_request_header(req).await?; session.finish_request_body().await?; custom_session!(session.finish_custom().await?); - if let Some(read_timeout) = peer.options.read_timeout { - session.set_read_timeout(Some(read_timeout)); - } + session.set_read_timeout(peer.options.read_timeout); session.read_response_header().await?; diff --git a/pingora-load-balancing/src/lib.rs b/pingora-load-balancing/src/lib.rs index 0e1bc6e5b..bf6a5d3bc 100644 --- a/pingora-load-balancing/src/lib.rs +++ b/pingora-load-balancing/src/lib.rs @@ -32,7 +32,7 @@ use std::hash::{Hash, Hasher}; use std::io::Result as IoResult; use std::net::ToSocketAddrs; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; mod background; pub mod discovery; @@ -303,6 +303,15 @@ impl Backends { } } +/// Timing information from the most recent [`LoadBalancer::update`] call. +#[derive(Debug, Clone, Copy)] +pub struct UpdateTimings { + /// Time spent in [`ServiceDiscovery::discover`]. + pub discovery_duration: Duration, + /// Time spent building the selection algorithm and storing the updated backends. + pub build_duration: Duration, +} + /// A [LoadBalancer] instance contains the service discovery, health check and backend selection /// all together. /// @@ -317,6 +326,11 @@ where config: Option, + /// Timing information from the most recent [`update`](Self::update) call. + /// + /// `None` until the first successful update completes. + last_update_timing: ArcSwap>, + /// How frequent the health check logic (if set) should run. /// /// If `None`, the health check logic will only run once at the beginning. @@ -366,6 +380,7 @@ where backends, selector, config: config_opt, + last_update_timing: ArcSwap::new(Arc::new(None)), health_check_frequency: None, update_frequency: None, parallel_health_check: false, @@ -381,18 +396,37 @@ where /// /// This function will be called every `update_frequency` if this [LoadBalancer] instance /// is running as a background service. + /// + /// On success, the timing information from this call is stored and can be + /// retrieved via [`last_update_timing`](Self::last_update_timing). pub async fn update(&self) -> Result<()> { + use std::sync::atomic::{AtomicU64, Ordering::Relaxed}; + + let build_nanos = AtomicU64::new(0); + let total_start = Instant::now(); + self.backends .update(|backends| { + let build_start = Instant::now(); let selector = if let Some(config) = &self.config { S::build_with_config(&backends, config) } else { S::build(&backends) }; - - self.selector.store(Arc::new(selector)) + self.selector.store(Arc::new(selector)); + build_nanos.store(build_start.elapsed().as_nanos() as u64, Relaxed); }) - .await + .await?; + + let total = total_start.elapsed(); + let build = Duration::from_nanos(build_nanos.load(Relaxed)); + + self.last_update_timing.store(Arc::new(Some(UpdateTimings { + discovery_duration: total.saturating_sub(build), + build_duration: build, + }))); + + Ok(()) } /// Return the first healthy [Backend] according to the selection algorithm and the @@ -442,6 +476,13 @@ where pub fn backends(&self) -> &Backends { &self.backends } + + /// Return the timing information from the most recent successful [`update`](Self::update) call. + /// + /// Returns `None` if [`update`](Self::update) has never completed successfully. + pub fn last_update_timing(&self) -> Option { + **self.last_update_timing.load() + } } #[cfg(test)] @@ -602,6 +643,39 @@ mod test { assert!(!backends.ready(&bad)); } + #[tokio::test] + async fn test_lb_update_stores_timing() { + let discovery = discovery::Static::default(); + let b1 = Backend::new("1.1.1.1:80").unwrap(); + let b2 = Backend::new("1.0.0.1:80").unwrap(); + discovery.add(b1.clone()); + discovery.add(b2.clone()); + + let lb = LoadBalancer::::from_backends(Backends::new(Box::new( + discovery, + ))); + + // Before first update, timing should be None + assert!(lb.last_update_timing().is_none()); + + lb.update().await.unwrap(); + + // After update, timing should be populated + let timing = lb + .last_update_timing() + .expect("timing should be Some after update"); + assert!(timing.discovery_duration > Duration::ZERO); + assert!(timing.build_duration > Duration::ZERO); + + // Backends should be populated + let backend = lb.backends().get_backend(); + assert!(backend.contains(&b1)); + assert!(backend.contains(&b2)); + + // Selection should work + assert!(lb.select(b"test", 10).is_some()); + } + mod thread_safety { use super::*; diff --git a/pingora-lru/benches/bench_lru.rs b/pingora-lru/benches/bench_lru.rs index c0bdc7761..02dadec82 100644 --- a/pingora-lru/benches/bench_lru.rs +++ b/pingora-lru/benches/bench_lru.rs @@ -12,137 +12,237 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! Benchmark for `Lru::promote()` vs `Lru::promote_top_n()`. +//! +//! Tests both small (original) and production-scale LRU sizes to show how +//! the `promote_top_n` optimization behaves at different scales. +//! +//! Run with: `cargo bench -p pingora-lru --bench bench_lru` +//! +//! ## Results (Apple M3 Max, 2026-04-03) +//! +//! Benchmark tiers simulate production-level data sizes (items/shard +//! ranging from ~100 to ~500K) with both uniform-hot and heavy-hitter +//! access patterns. +//! +//! ### 8-threaded — uniform hot set (10% of items are 100x hotter) +//! +//! | Items/shard | promote | top_n(0) | top_n(3) | top_n(10) | top_n(50) | top_n(100) | +//! |-------------|------------|----------|----------|-----------|-----------|------------| +//! | 10 (orig) | 366ns | 476ns | 271ns | **164ns** | 164ns | 164ns | +//! | 100K (typ) | **457ns** | 480ns | 437ns | 520ns | 1227ns | 2394ns | +//! +//! ### 8-threaded — heavy hitters (10 or 100 items are 10,000x hotter) +//! +//! | Items/shard | promote | top_n(0) | top_n(3) | top_n(10) | top_n(50) | top_n(100) | +//! |-----------------|------------|----------|----------|-----------|-----------|------------| +//! | 100K, 10 hot | **649ns** | 688ns | 652ns | 773ns | 1811ns | 3534ns | +//! | 100K, 100 hot | **607ns** | 632ns | 607ns | 716ns | 1493ns | 2759ns | +//! +//! ### Single-threaded — uniform hot set (10% of items are 100x hotter) +//! +//! | Items/shard | promote | top_n(0) | top_n(3) | top_n(10) | top_n(50) | top_n(100) | +//! |-------------|------------|----------|----------|-----------|-----------|------------| +//! | 10 (orig) | 22ns | 20ns | 29ns | **30ns** | 30ns | 30ns | +//! | 100K (typ) | **297ns** | 306ns | 314ns | 332ns | 663ns | 1092ns | +//! +//! **Conclusions**: +//! +//! - `promote_top_n(0)` is strictly worse than `promote()` — it takes a +//! wasted read lock before falling through to the write lock every time. +//! +//! - `promote_top_n(n)` for n > 0 only wins at the original small scale +//! (10 items/shard) where the threshold covers the entire shard. +//! +//! - Even with heavy-hitter patterns (10 items at 10,000x weight), +//! `promote()` ties or wins at production scale. With 10 hot items +//! across 32 shards, most shards have 0-1 hot items, so the read-lock +//! scan is wasted on the majority of cold-item accesses. +//! +//! - At production scale (~100K+ items/shard), plain `promote()` is fastest +//! regardless of access pattern. + use rand::distributions::WeightedIndex; use rand::prelude::*; -use std::sync::Arc; +use std::sync::{Arc, Barrier}; use std::thread; use std::time::Instant; -// Non-uniform distributions, 100 items, 10 of them are 100x more likely to appear -const WEIGHTS: &[usize] = &[ - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 100, 100, 100, - 100, 100, 100, 100, 100, 100, 100, -]; - const ITERATIONS: usize = 5_000_000; const THREADS: usize = 8; -fn main() { - let lru = parking_lot::Mutex::new(lru::LruCache::::unbounded()); - - let plru = pingora_lru::Lru::<(), 10>::with_capacity(1000, 100); - // populate first, then we bench access/promotion - for i in 0..WEIGHTS.len() { - lru.lock().put(i as u64, ()); - } - for i in 0..WEIGHTS.len() { - plru.admit(i as u64, (), 1); +/// Build a weight distribution where the first `hot_count` items have +/// `hot_weight`x the access probability. +fn make_weights(n: usize, hot_count: usize, hot_weight: usize) -> Vec { + let mut weights = vec![1usize; n]; + for w in weights.iter_mut().take(hot_count) { + *w = hot_weight; } + weights +} - // single thread - let mut rng = thread_rng(); - let dist = WeightedIndex::new(WEIGHTS).unwrap(); +fn bench_config(label: &str, items: usize, shards: usize, hot_pct: usize, hot_weight: usize) { + let hot_count = items * hot_pct / 100; + bench_config_abs(label, items, shards, hot_count, hot_weight); +} - let before = Instant::now(); - for _ in 0..ITERATIONS { - lru.lock().get(&(dist.sample(&mut rng) as u64)); - } - let elapsed = before.elapsed(); +fn bench_config_abs(label: &str, items: usize, shards: usize, hot_count: usize, hot_weight: usize) { println!( - "lru promote total {elapsed:?}, {:?} avg per operation", - elapsed / ITERATIONS as u32 + "\n=== {label}: {items} items, {shards} shards ({} per shard), \ + {hot_count} items are {hot_weight}x hotter ===", + items / shards ); - let before = Instant::now(); - for _ in 0..ITERATIONS { - plru.promote(dist.sample(&mut rng) as u64); - } - let elapsed = before.elapsed(); - println!( - "pingora lru promote total {elapsed:?}, {:?} avg per operation", - elapsed / ITERATIONS as u32 - ); + let weights = make_weights(items, hot_count, hot_weight); + let dist = Arc::new(WeightedIndex::new(&weights).unwrap()); - let before = Instant::now(); - for _ in 0..ITERATIONS { - plru.promote_top_n(dist.sample(&mut rng) as u64, 10); + match shards { + 10 => bench_shards::<10>(items, &dist), + 32 => bench_shards::<32>(items, &dist), + _ => panic!("unsupported shard count: {shards}"), } - let elapsed = before.elapsed(); - println!( - "pingora lru promote_top_10 total {elapsed:?}, {:?} avg per operation", - elapsed / ITERATIONS as u32 - ); +} - // concurrent - - let lru = Arc::new(lru); - let mut handlers = vec![]; - for i in 0..THREADS { - let lru = lru.clone(); - let handler = thread::spawn(move || { - let mut rng = thread_rng(); - let dist = WeightedIndex::new(WEIGHTS).unwrap(); - let before = Instant::now(); - for _ in 0..ITERATIONS { - lru.lock().get(&(dist.sample(&mut rng) as u64)); - } - let elapsed = before.elapsed(); - println!( - "lru promote total {elapsed:?}, {:?} avg per operation thread {i}", - elapsed / ITERATIONS as u32 - ); - }); - handlers.push(handler); +/// Populate a fresh LRU with `items` entries. +fn make_lru(items: usize) -> pingora_lru::Lru<(), N> { + let lru = pingora_lru::Lru::<(), N>::with_capacity(items, items / N); + for i in 0..items { + lru.admit(i as u64, (), 1); } - for thread in handlers { - thread.join().unwrap(); + lru +} + +fn bench_shards(items: usize, dist: &Arc>) { + // Each variant gets a fresh LRU to avoid state contamination from + // prior runs warming hot items to the head. + + // --- Single-threaded --- + println!(" Single-threaded:"); + { + let lru = make_lru::(items); + let mut rng = thread_rng(); + let before = Instant::now(); + for _ in 0..ITERATIONS { + lru.promote(dist.sample(&mut rng) as u64); + } + let elapsed = before.elapsed(); + println!( + " promote: {elapsed:?} total, {:?} avg", + elapsed / ITERATIONS as u32, + ); } - let plru = Arc::new(plru); - - let mut handlers = vec![]; - for i in 0..THREADS { - let plru = plru.clone(); - let handler = thread::spawn(move || { - let mut rng = thread_rng(); - let dist = WeightedIndex::new(WEIGHTS).unwrap(); - let before = Instant::now(); - for _ in 0..ITERATIONS { - plru.promote(dist.sample(&mut rng) as u64); - } - let elapsed = before.elapsed(); - println!( - "pingora lru promote total {elapsed:?}, {:?} avg per operation thread {i}", - elapsed / ITERATIONS as u32 - ); - }); - handlers.push(handler); + for top_n in [0, 3, 10, 50, 100] { + let lru = make_lru::(items); + let mut rng = thread_rng(); + let before = Instant::now(); + for _ in 0..ITERATIONS { + lru.promote_top_n(dist.sample(&mut rng) as u64, top_n); + } + let elapsed = before.elapsed(); + println!( + " promote_top_{top_n:<3} {elapsed:?} total, {:?} avg", + elapsed / ITERATIONS as u32, + ); } - for thread in handlers { - thread.join().unwrap(); + + // --- Multi-threaded --- + println!(" {THREADS}-threaded:"); + + { + let lru = Arc::new(make_lru::(items)); + let barrier = Arc::new(Barrier::new(THREADS)); + let mut handlers = vec![]; + for _ in 0..THREADS { + let lru = lru.clone(); + let dist = Arc::clone(dist); + let barrier = barrier.clone(); + handlers.push(thread::spawn(move || { + let mut rng = thread_rng(); + barrier.wait(); + let before = Instant::now(); + for _ in 0..ITERATIONS { + lru.promote(dist.sample(&mut rng) as u64); + } + before.elapsed() + })); + } + let elapsed: Vec<_> = handlers.into_iter().map(|h| h.join().unwrap()).collect(); + let avg = elapsed.iter().sum::() / THREADS as u32; + println!( + " promote: avg {avg:?}, {:?} avg per op", + avg / ITERATIONS as u32, + ); } - let mut handlers = vec![]; - for i in 0..THREADS { - let plru = plru.clone(); - let handler = thread::spawn(move || { - let mut rng = thread_rng(); - let dist = WeightedIndex::new(WEIGHTS).unwrap(); - let before = Instant::now(); - for _ in 0..ITERATIONS { - plru.promote_top_n(dist.sample(&mut rng) as u64, 10); - } - let elapsed = before.elapsed(); - println!( - "pingora lru promote_top_10 total {elapsed:?}, {:?} avg per operation thread {i}", - elapsed / ITERATIONS as u32 - ); - }); - handlers.push(handler); + for top_n in [0, 3, 10, 50, 100] { + let lru = Arc::new(make_lru::(items)); + let barrier = Arc::new(Barrier::new(THREADS)); + let mut handlers = vec![]; + for _ in 0..THREADS { + let lru = lru.clone(); + let dist = Arc::clone(dist); + let barrier = barrier.clone(); + handlers.push(thread::spawn(move || { + let mut rng = thread_rng(); + barrier.wait(); + let before = Instant::now(); + for _ in 0..ITERATIONS { + lru.promote_top_n(dist.sample(&mut rng) as u64, top_n); + } + before.elapsed() + })); + } + let elapsed: Vec<_> = handlers.into_iter().map(|h| h.join().unwrap()).collect(); + let avg = elapsed.iter().sum::() / THREADS as u32; + println!( + " promote_top_{top_n:<3} avg {avg:?}, {:?} avg per op", + avg / ITERATIONS as u32, + ); } - for thread in handlers { - thread.join().unwrap(); +} + +fn main() { + // Benchmark tiers to simulate production-level data sizes: + // Small = original bench scale (10 items/shard) + // Typical = ~100K items/shard (3.2M total across 32 shards) + // Large = ~500K items/shard (16M total) — gated behind + // BENCH_LARGE=1 to avoid OOM on CI runners (~1.5GB heap) + // + // Note: the Typical tier allocates ~150MB per make_lru() call. With + // multiple variants (promote + 5 top_n values) × configs, total peak + // memory is ~1GB. Well within CI limits but notable for constrained machines. + + // Original benchmark scale (100 items, 10 shards = 10 per shard) + bench_config("Small (original bench scale)", 100, 10, 10, 100); + + // Typical (~100K items/shard), 10% hot + bench_config("Typical (100K/shard, 10% hot)", 3_200_000, 32, 10, 100); + + // Typical (~100K items/shard), heavy-hitter: only 10 items dominate + // Simulates viral content / popular API endpoints where a handful of + // assets receive the vast majority of traffic. + bench_config_abs( + "Typical (100K/shard, 10 heavy hitters)", + 3_200_000, + 32, + 10, + 10_000, + ); + + // Typical (~100K items/shard), moderate hot set: 100 items dominate + bench_config_abs( + "Typical (100K/shard, 100 heavy hitters)", + 3_200_000, + 32, + 100, + 10_000, + ); + + // Large (~500K items/shard, ~1.5GB heap) + if std::env::var("BENCH_LARGE").is_ok() { + bench_config("Large (500K/shard, 10% hot)", 16_000_000, 32, 10, 100); + } else { + println!("\n=== Skipping large bench (set BENCH_LARGE=1 to enable) ==="); } } diff --git a/pingora-lru/src/lib.rs b/pingora-lru/src/lib.rs index 23728c4f3..386dc173a 100644 --- a/pingora-lru/src/lib.rs +++ b/pingora-lru/src/lib.rs @@ -25,13 +25,18 @@ use linked_list::{LinkedList, LinkedListIter}; use hashbrown::HashMap; use parking_lot::RwLock; +use rand::Rng; use std::sync::atomic::{AtomicUsize, Ordering}; /// The LRU with `N` shards pub struct Lru { units: [RwLock>; N], + /// Lock-free `Relaxed` shadow of each shard's item count, backing + /// [`Lru::shard_len`] and the P2C selection in [`Lru::evict_to_limit`]. + /// Maintained alongside [`Lru::len`] at every count-mutating site. + shard_lens: [AtomicUsize; N], weight: AtomicUsize, - weight_limit: usize, + weight_limit: AtomicUsize, len_watermark: Option, len: AtomicUsize, evicted_weight: AtomicUsize, @@ -59,13 +64,18 @@ impl Lru { ) -> Self { // use the unsafe code from ArrayVec just to init the array let mut units = arrayvec::ArrayVec::<_, N>::new(); + let mut shard_lens = arrayvec::ArrayVec::<_, N>::new(); for _ in 0..N { units.push(RwLock::new(LruUnit::with_capacity(capacity))); + shard_lens.push(AtomicUsize::new(0)); } Lru { units: units.into_inner().map_err(|_| "").unwrap(), + shard_lens: shard_lens + .into_inner() + .expect("shard_lens ArrayVec filled with exactly N elements"), weight: AtomicUsize::new(0), - weight_limit, + weight_limit: AtomicUsize::new(weight_limit), len_watermark, len: AtomicUsize::new(0), evicted_weight: AtomicUsize::new(0), @@ -73,6 +83,33 @@ impl Lru { } } + /// Return the current total weight limit. + pub fn weight_limit(&self) -> usize { + self.weight_limit.load(Ordering::Relaxed) + } + + /// Set the total weight limit used by [`Self::evict_to_limit`]. + pub fn set_weight_limit(&self, weight_limit: usize) { + self.weight_limit.store(weight_limit, Ordering::Relaxed); + } + + /// Increment item-count bookkeeping for `shard`. Both atomics use + /// `Relaxed`; called while holding the shard write lock so that + /// `len` and `shard_lens[shard]` advance in lockstep. + #[inline] + fn incr_count(&self, shard: usize) { + self.len.fetch_add(1, Ordering::Relaxed); + self.shard_lens[shard].fetch_add(1, Ordering::Relaxed); + } + + /// Decrement item-count bookkeeping for `shard`. See + /// [`Self::incr_count`]. + #[inline] + fn decr_count(&self, shard: usize) { + self.len.fetch_sub(1, Ordering::Relaxed); + self.shard_lens[shard].fetch_sub(1, Ordering::Relaxed); + } + /// Admit the key value to the [Lru] /// /// Return the shard index which the asset is added to @@ -91,7 +128,7 @@ impl Lru { self.weight.fetch_sub(old_weight, Ordering::Relaxed); } else { // Assume old_weight == 0 means a new item is admitted - self.len.fetch_add(1, Ordering::Relaxed); + self.incr_count(shard); } } shard @@ -128,9 +165,18 @@ impl Lru { /// Promote to the top n of the LRU /// - /// This function is a bit more efficient in terms of reducing lock contention because it - /// will acquire a write lock only if the key is outside top n but only acquires a read lock - /// when the key is already in the top n. + /// This function acquires a read lock first to check if the key is already + /// in the top `n` positions. If so, it returns early without a write lock. + /// Otherwise it falls through to a write lock for the actual promotion. + /// + /// **Performance note**: this optimization only helps when `n` covers a + /// significant fraction of the shard. At production scale (~100K+ items + /// per shard), hot items are rarely in the top N positions, so the + /// read-lock scan is usually wasted work that adds latency without + /// reducing contention. Benchmarks (`cargo bench --bench bench_lru`) + /// show that plain [`promote()`](Self::promote) is faster at scale. + /// Consider using `promote()` directly unless profiling shows a clear + /// benefit for your workload. /// /// Return false if the item doesn't exist pub fn promote_top_n(&self, key: u64, top: usize) -> bool { @@ -141,14 +187,23 @@ impl Lru { unit.write().access(key) } - /// Evict at most one item from the given shard + /// Evict at most one item from the given shard, identified by the + /// hash-like `shard` seed (mapped into `0..N` via `% N`). /// - /// Return the evicted asset and its size if there is anything to evict + /// Return the evicted asset and its size if there is anything to evict. pub fn evict_shard(&self, shard: u64) -> Option<(T, usize)> { - let evicted = self.units[get_shard(shard, N)].write().evict(); + self.evict_shard_at(get_shard(shard, N)) + } + + /// Evict at most one item from the shard at index `shard` (in `0..N`). + /// Internal entry point that skips the `% N` round-trip in + /// [`Self::evict_shard`]. + fn evict_shard_at(&self, shard: usize) -> Option<(T, usize)> { + assert!(shard < N); + let evicted = self.units[shard].write().evict(); if let Some((_, weight)) = evicted.as_ref() { self.weight.fetch_sub(*weight, Ordering::Relaxed); - self.len.fetch_sub(1, Ordering::Relaxed); + self.decr_count(shard); self.evicted_weight.fetch_add(*weight, Ordering::Relaxed); self.evicted_len.fetch_add(1, Ordering::Relaxed); } @@ -159,43 +214,93 @@ impl Lru { /// /// Return a list of evicted items. /// - /// The evicted items are randomly selected from all the shards. + /// Each iteration selects the shard to evict from using the "power of two + /// choices" strategy: two shards are picked uniformly at random and the + /// one with more items is chosen (see + /// ). This biases + /// eviction toward longer shards and drives [`Self::shard_len`] toward a + /// uniform distribution, which keeps per-shard serialization cost (e.g. + /// `pingora_cache::eviction::lru::Manager::serialize_shard`) bounded. + /// + /// Selection is by item count, not weight, even when eviction is + /// triggered by `weight_limit`. With heavily skewed item weights this + /// may evict more items than a weight-biased policy to reach the same + /// total weight — the tradeoff is intentional in favor of bounded + /// per-shard serialization cost. + /// + /// O(1) per iteration in the common case. If the chosen shard is + /// empty when we acquire its write lock (the Relaxed shadow may + /// not always reflect actual emptiness, and P2C may tie-break to + /// an empty shard when all shadow lengths are equal), we linearly + /// probe successive shard indices until one yields an item or we + /// wrap back to the starting shard — at which point every shard + /// was observed empty and we exit. Bounded by at most N probes + /// per outer iteration. pub fn evict_to_limit(&self) -> Vec<(T, usize)> { + self.evict_to_limit_with_rng(&mut rand::thread_rng()) + } + + /// Internal entry point for [`Self::evict_to_limit`] that lets tests + /// inject a seeded RNG for deterministic P2C selection. + fn evict_to_limit_with_rng(&self, rng: &mut R) -> Vec<(T, usize)> { let mut evicted = vec![]; let mut initial_weight = self.weight(); let mut initial_len = self.len(); - let mut shard_seed = rand::random(); // start from a random shard - let mut empty_shard = 0; - - // Entries can be admitted or removed from the LRU by others during the loop below - // Track initial size not to over evict due to entries admitted after the loop starts - // self.weight() / self.len() is also used not to over evict - // due to entries already removed by others - while ((initial_weight > self.weight_limit && self.weight() > self.weight_limit) + + // Transient over-limit weight can persist until the next + // admit/increment_weight call, which is acceptable because the + // next admission will re-trigger eviction. + let weight_limit = self.weight_limit(); + while (initial_weight > weight_limit && self.weight() > weight_limit) || self .len_watermark - .is_some_and(|w| initial_len > w && self.len() > w)) - && empty_shard < N + .is_some_and(|w| initial_len > w && self.len() > w) { - if let Some(i) = self.evict_shard(shard_seed) { - initial_weight -= i.1; - initial_len = initial_len.saturating_sub(1); - evicted.push(i) + // Power of two choices: pick the longer of two random shards. + // N == 1 short-circuits the redundant second roll. + let start = if N <= 1 { + 0 } else { - empty_shard += 1; + let a = rng.gen_range(0..N); + let b = rng.gen_range(0..N); + if self.shard_len(a) >= self.shard_len(b) { + a + } else { + b + } + }; + // Try the chosen shard first; on a miss (empty or raced), + // linearly probe successive indices. Wrapping back to + // `start` means every shard was observed empty, so we exit. + let mut shard = start; + let evicted_one = loop { + if let Some(item) = self.evict_shard_at(shard) { + break Some(item); + } + shard = (shard + 1) % N; + if shard == start { + break None; + } + }; + match evicted_one { + Some(i) => { + initial_weight = initial_weight.saturating_sub(i.1); + initial_len = initial_len.saturating_sub(1); + evicted.push(i); + } + None => break, } - // move on to the next shard - shard_seed += 1; } evicted } /// Remove the given asset. pub fn remove(&self, key: u64) -> Option<(T, usize)> { - let removed = self.units[get_shard(key, N)].write().remove(key); + let shard = get_shard(key, N); + let removed = self.units[shard].write().remove(key); if let Some((_, weight)) = removed.as_ref() { self.weight.fetch_sub(*weight, Ordering::Relaxed); - self.len.fetch_sub(1, Ordering::Relaxed); + self.decr_count(shard); } removed } @@ -204,12 +309,10 @@ impl Lru { /// /// Useful to recreate an LRU in most-to-least order pub fn insert_tail(&self, key: u64, data: T, weight: usize) -> bool { - if self.units[get_shard(key, N)] - .write() - .insert_tail(key, data, weight) - { + let shard = get_shard(key, N); + if self.units[shard].write().insert_tail(key, data, weight) { self.weight.fetch_add(weight, Ordering::Relaxed); - self.len.fetch_add(1, Ordering::Relaxed); + self.incr_count(shard); true } else { false @@ -226,7 +329,25 @@ impl Lru { self.units[get_shard(key, N)].read().peek_weight(key) } + /// Peek at the least-recently-used item in the given shard without removing it. + /// + /// Returns a clone of the data and the weight, or `None` if the shard is empty + /// or `shard >= N`. + pub fn peek_lru(&self, shard: usize) -> Option<(T, usize)> + where + T: Clone, + { + self.units + .get(shard)? + .read() + .peek_lru() + .map(|(data, weight)| (data.clone(), weight)) + } + /// Return the current total weight. + /// + /// Lock-free `Relaxed` load. Best-effort: not synchronized with + /// concurrent admissions or evictions on other threads. pub fn weight(&self) -> usize { self.weight.load(Ordering::Relaxed) } @@ -261,9 +382,15 @@ impl Lru { N } - /// Get the number of items inside a shard + /// Get the number of items inside a shard. + /// + /// Lock-free `Relaxed` load from a per-shard atomic shadow. Best-effort: + /// there is no cross-thread ordering between this and [`Self::len`], and + /// `Σ shard_len(i)` is not guaranteed to equal [`Self::len`] at any + /// given instant. Suitable for eviction-balance heuristics and + /// observability; not suitable for synchronization. pub fn shard_len(&self, shard: usize) -> usize { - self.units[shard].read().len() + self.shard_lens[shard].load(Ordering::Relaxed) } /// Get the weight (total size) inside a shard @@ -374,6 +501,19 @@ impl LruUnit { (node.data, node.weight) }) } + + /// Peek at the least-recently-used item without removing it. + /// + /// Returns a reference to the data and weight of the tail item, or `None` + /// if empty. + pub fn peek_lru(&self) -> Option<(&T, usize)> { + self.order + .tail() + .and_then(|idx| self.order.peek(idx)) + .and_then(|key| self.lookup_table.get(&key)) + .map(|node| (&node.data, node.weight)) + } + // TODO: scan the tail up to K elements to decide which ones to evict pub fn remove(&mut self, key: u64) -> Option<(T, usize)> { @@ -400,6 +540,7 @@ impl LruUnit { true } + #[cfg(test)] pub fn len(&self) -> usize { assert_eq!(self.lookup_table.len(), self.order.len()); self.lookup_table.len() @@ -583,7 +724,6 @@ mod test_lru { assert_eq!(lru.len(), 6); let evicted = lru.evict_to_limit(); - // NOTE: there is a low chance this test would fail see the TODO in evict_to_limit assert_eq!(lru.weight(), 6); assert_eq!(lru.len(), 3); assert_eq!(lru.evicted_weight(), 6); @@ -678,6 +818,138 @@ mod test_lru { assert!(!lru.insert_tail(6, 6, 7)); } + #[test] + fn test_evict_to_limit_p2c_bias() { + use rand::rngs::StdRng; + use rand::SeedableRng; + + // Shard 0 starts with 50 items, shard 1 with 10 (all weight 1). + // weight_limit=30 forces 30 evictions. P2C-by-length should pick + // shard 0 (the longer one) most of the time, driving toward + // balance. Expected share from shard 0: P2C ≈ 0.75 (P(shard 0) + // = 3/4 per pick while it stays longer), uniform ≈ 0.50, + // always-shortest ≈ 0.67 (capped by shard 1's 10 items), + // always-longest ≈ 1.0. The (0.65..0.95) window distinguishes + // P2C from uniform; the upper bound catches a degenerate + // always-longest regression. + const TRIALS: u64 = 50; + let mut total_from_shard0 = 0usize; + let mut total_evicted = 0usize; + + for seed in 0..TRIALS { + let lru = Lru::::with_capacity(30, 64); + for k in 0..50u64 { + // even keys → shard 0 + lru.admit(k * 2, k * 2, 1); + } + for k in 0..10u64 { + // odd keys → shard 1 + lru.admit(k * 2 + 1, k * 2 + 1, 1); + } + assert_eq!(lru.weight(), 60); + + let mut rng = StdRng::seed_from_u64(seed); + let evicted = lru.evict_to_limit_with_rng(&mut rng); + assert!( + lru.weight() <= 30, + "post-eviction weight {} exceeds limit", + lru.weight() + ); + total_from_shard0 += evicted.iter().filter(|(k, _)| k % 2 == 0).count(); + total_evicted += evicted.len(); + } + + assert!(total_evicted > 1000, "too few evictions: {total_evicted}"); + let share = total_from_shard0 as f64 / total_evicted as f64; + assert!( + (0.65..0.95).contains(&share), + "expected shard-0 eviction share in 0.65..0.95 (P2C ≈ 0.75); got {share}" + ); + } + + #[test] + fn test_evict_to_limit_break_on_empty_shards_over_limit() { + // Force `weight` above the limit while every shard is empty + // (simulating bookkeeping skew). The linear probe must wrap + // around all N shards and exit cleanly. + let lru = Lru::::with_capacity(10, 16); + lru.weight.fetch_add(100, Ordering::Relaxed); + assert_eq!(lru.evict_to_limit().len(), 0); + } + + #[test] + fn test_watermark_eviction_with_zero_weight_items() { + // All items have weight 0 so the weight-limit guard never fires; + // only the length watermark drives eviction. P2C-by-length should + // still reach the watermark regardless of weight values. + let lru = Lru::::with_capacity_and_watermark(usize::MAX / 2, 10, Some(2)); + for k in 0..6u64 { + lru.insert_tail(k, k, 0); + } + assert_eq!(lru.len(), 6); + assert_eq!(lru.weight(), 0); + let evicted = lru.evict_to_limit(); + assert_eq!(lru.len(), 2); + assert_eq!(evicted.len(), 4); + } + + #[test] + fn test_evict_to_limit_with_mostly_empty_shards() { + // 7/8 shards empty: both random rolls land on empty shards ~77% + // of the time, exercising the linear-probe fallback heavily. + let lru = Lru::::with_capacity(2, 16); + for k in 0..8u64 { + // multiples of 8 hash to shard 0 + lru.admit(k * 8, k * 8, 1); + } + assert_eq!(lru.weight(), 8); + + let evicted = lru.evict_to_limit(); + assert_eq!(lru.weight(), 2); + assert_eq!(evicted.len(), 6); + assert!(evicted.iter().all(|(k, _)| k % 8 == 0)); + } + + #[test] + fn test_evict_to_limit_below_limit_returns_immediately() { + // Smoke test: outer guard short-circuits when already under limit. + let lru = Lru::::with_capacity(0, 16); + assert_eq!(lru.evict_to_limit().len(), 0); + } + + #[test] + fn test_evict_to_limit_n1() { + // N=1 is a trivial special case in the selection logic; ensure + // basic eviction still works. + let lru = Lru::::with_capacity(2, 16); + for k in 0..5u64 { + lru.admit(k, k, 1); + } + assert_eq!(lru.weight(), 5); + let evicted = lru.evict_to_limit(); + assert_eq!(lru.weight(), 2); + assert_eq!(evicted.len(), 3); + } + + #[test] + fn test_set_weight_limit_affects_eviction() { + let lru = Lru::::with_capacity(10, 16); + for k in 0..5u64 { + lru.admit(k, k, 2); + } + assert_eq!(lru.weight(), 10); + assert_eq!(lru.weight_limit(), 10); + + lru.set_weight_limit(4); + assert_eq!(lru.weight_limit(), 4); + let evicted = lru.evict_to_limit(); + assert_eq!(lru.weight(), 4); + assert_eq!(evicted.len(), 3); + + lru.set_weight_limit(20); + assert_eq!(lru.evict_to_limit().len(), 0); + } + #[test] fn test_watermark_eviction() { const WEIGHT_LIMIT: usize = usize::MAX / 2; @@ -696,6 +968,29 @@ mod test_lru { assert_eq!(evicted.len(), 2); assert_eq!(lru.evicted_len(), 2); } + + #[test] + fn test_peek_lru() { + let lru = Lru::::with_capacity(10, 10); + + // empty shard + assert!(lru.peek_lru(0).is_none()); + + lru.admit(1, 10, 1); + assert_eq!(lru.peek_lru(0).unwrap(), (10, 1)); + + lru.admit(2, 20, 2); + // key 1 is LRU tail + assert_eq!(lru.peek_lru(0).unwrap(), (10, 1)); + + // promote key 1 + lru.promote(1); + // key 2 is now LRU tail + assert_eq!(lru.peek_lru(0).unwrap(), (20, 2)); + + // out-of-bounds returns None + assert!(lru.peek_lru(999).is_none()); + } } #[cfg(test)] @@ -865,4 +1160,33 @@ mod test_lru_unit { assert_eq!(lru.used_weight(), 1 + 3 + 4 + 5); assert_lru(&lru, &[2, 3, 4, 5]); } + + #[test] + fn test_peek_lru() { + let mut lru = LruUnit::with_capacity(10); + + // empty returns None + assert!(lru.peek_lru().is_none()); + + // single item is both head and tail + lru.admit(1, 10, 1); + let (data, weight) = lru.peek_lru().unwrap(); + assert_eq!(*data, 10); + assert_eq!(weight, 1); + + // second admission pushes first to tail + lru.admit(2, 20, 2); + let (data, _) = lru.peek_lru().unwrap(); + assert_eq!(*data, 10); // key 1 is LRU tail + + // promote key 1 — now key 2 is tail + lru.access(1); + let (data, _) = lru.peek_lru().unwrap(); + assert_eq!(*data, 20); // key 2 is now LRU tail + + // peek doesn't remove + assert!(lru.peek_lru().is_some()); + assert!(lru.peek(1).is_some()); + assert!(lru.peek(2).is_some()); + } } diff --git a/pingora-openssl/src/ext.rs b/pingora-openssl/src/ext.rs index 18e0fdfed..43c5f7fbd 100644 --- a/pingora-openssl/src/ext.rs +++ b/pingora-openssl/src/ext.rs @@ -166,6 +166,20 @@ pub fn clear_error_stack() { let _ = ErrorStack::get(); } +/// Export keying material from a TLS connection +/// +/// Derives keying material for application use in accordance with RFC 5705. +/// +/// See [SSL_export_keying_material](https://www.openssl.org/docs/man1.1.1/man3/SSL_export_keying_material.html). +pub fn ssl_export_keying_material( + ssl: &SslRef, + out: &mut [u8], + label: &str, + context: Option<&[u8]>, +) -> Result<(), ErrorStack> { + ssl.export_keying_material(out, label, context) +} + /// Create a new [Ssl] from &[SslAcceptor] /// /// this function is to unify the interface between this crate and [`pingora-boringssl`](https://docs.rs/pingora-boringssl) @@ -228,4 +242,19 @@ mod tests { // Invalid input (contains null byte) assert!(ssl_set_groups_list(ssl_ref, "P-256\0P-384").is_err()); } + + #[test] + fn test_ssl_export_keying_material_exists() { + // This test verifies that ssl_export_keying_material function exists + // and has the correct signature. Actual functional testing requires + // an established TLS connection. + let ctx_builder = SslContextBuilder::new(SslMethod::tls()).unwrap(); + let ssl = Ssl::new(&ctx_builder.build()).unwrap(); + let ssl_ref = &ssl; + let mut out = [0u8; 32]; + + // This will fail since there's no established connection, but verifies + // the function signature is correct + let _ = ssl_export_keying_material(ssl_ref, &mut out, "test", None); + } } diff --git a/pingora-pool/Cargo.toml b/pingora-pool/Cargo.toml index 5d841a4c6..b1a39c501 100644 --- a/pingora-pool/Cargo.toml +++ b/pingora-pool/Cargo.toml @@ -18,12 +18,20 @@ path = "src/lib.rs" [dependencies] tokio = { workspace = true, features = ["sync", "io-util"] } -thread_local = "1.0" lru = { workspace = true } log = { workspace = true } parking_lot = "0.12" +dashmap = "5" crossbeam-queue = "0.3" +futures = { workspace = true } pingora-timeout = { version = "0.8.0", path = "../pingora-timeout" } [dev-dependencies] tokio-test = "0.4" + +[lints] +workspace = true + +[[bench]] +name = "connection_pool" +harness = false diff --git a/pingora-pool/benches/connection_pool.rs b/pingora-pool/benches/connection_pool.rs new file mode 100644 index 000000000..58a85386d --- /dev/null +++ b/pingora-pool/benches/connection_pool.rs @@ -0,0 +1,274 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Connection-pool microbenchmarks. +//! +//! Run with: `cargo bench -p pingora-pool --bench connection_pool` +//! +//! These benchmarks include both no-eviction scaling cases (each thread +//! reuses its own key) and an eviction-pressure case (more unique keys than +//! pool capacity) so that eviction code paths actually get exercised under +//! concurrency. +//! +//! The LRU-only benchmark includes the LRU module directly with +//! `#[path = "../src/lru.rs"]` because `Lru` is `pub(crate)` and not part of +//! the public `pingora-pool` API. The shard counts used here are imported +//! from that module so they cannot drift from the implementation. + +#[allow(dead_code, unused_imports)] +#[path = "../src/lru.rs"] +mod lru; + +use lru::{Lru, N_SHARDS as LRU_SHARDS}; +use pingora_pool::{ConnectionMeta, ConnectionPool}; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::hint::black_box; +use std::sync::{Arc, Barrier}; +use std::thread; +use std::time::{Duration, Instant}; + +const POOL_SIZE: usize = 1024; +const ROUND_TRIP_ITERATIONS: usize = 1_000_000; +const EVICTION_ITERATIONS: usize = 100_000; +const THREADED_ITERATIONS: usize = 250_000; +const LRU_ITERATIONS: usize = 250_000; +const THREAD_COUNTS: [usize; 5] = [1, 2, 4, 8, 16]; +/// For eviction-pressure benchmarks, each thread cycles through this many +/// distinct keys so that put() routinely overflows the pool capacity. +/// 2x the per-thread cap on master so master's thread-local LRU evicts on +/// roughly half of every thread's puts; on the sharded branch the same +/// workload pushes total unique keys far above the global cap, so eviction +/// fires constantly there too. +const EVICTION_KEYS_PER_THREAD: u64 = (POOL_SIZE as u64) * 2; + +fn print_result(name: &str, elapsed: Duration, iterations: usize) { + println!( + "{name:<32} {elapsed:?} total, {:?} avg/op", + elapsed / iterations as u32 + ); +} + +fn bench_round_trip() { + let pool = ConnectionPool::new(POOL_SIZE); + let meta = ConnectionMeta::new(1, 1); + pool.put(&meta, 1usize); + + let before = Instant::now(); + for _ in 0..ROUND_TRIP_ITERATIONS { + let conn = pool.get(&meta.key).unwrap(); + black_box(conn); + pool.put(&meta, conn); + } + print_result( + "single-thread get+put", + before.elapsed(), + ROUND_TRIP_ITERATIONS, + ); +} + +fn bench_eviction_pressure() { + let pool = ConnectionPool::new(POOL_SIZE); + let before = Instant::now(); + + for id in 0..EVICTION_ITERATIONS { + let meta = ConnectionMeta::new(id as u64, id as _); + pool.put(&meta, black_box(id)); + } + + print_result( + "put with eviction pressure", + before.elapsed(), + EVICTION_ITERATIONS, + ); +} + +/// Use one distinct pool key per worker. +/// +/// This intentionally avoids claiming a specific pool-map shard layout: the +/// connection pool may use a concurrent map whose internal shard selection is +/// not part of this crate's API. +fn pool_distinct_keys(threads: usize) -> Vec { + let keys: Vec<_> = (0..threads as u64).collect(); + assert_eq!(keys.len(), threads); + keys +} + +fn bench_pool_pattern(label: &str, keys: Vec) { + let threads = keys.len(); + let pool = Arc::new(ConnectionPool::new(POOL_SIZE)); + let barrier = Arc::new(Barrier::new(threads)); + let mut handles = Vec::with_capacity(threads); + + for (worker, key) in keys.into_iter().enumerate() { + let pool = pool.clone(); + let barrier = barrier.clone(); + handles.push(thread::spawn(move || { + let meta = ConnectionMeta::new(key, worker as _); + pool.put(&meta, worker); + barrier.wait(); + + // Timing starts after barrier release so warmup setup doesn't + // count against the measured loop. + let before = Instant::now(); + for _ in 0..THREADED_ITERATIONS { + let conn = pool.get(&meta.key).unwrap(); + black_box(conn); + pool.put(&meta, conn); + } + before.elapsed() + })); + } + + let elapsed: Duration = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .sum(); + print_result( + &format!("pool {label} {threads:>2} threads"), + elapsed / threads as u32, + THREADED_ITERATIONS, + ); +} + +/// Concurrent benchmark that intentionally exceeds pool capacity so the LRU +/// eviction path runs while threads contend on the pool. +fn bench_pool_eviction_pressure(threads: usize) { + let pool = Arc::new(ConnectionPool::new(POOL_SIZE)); + let barrier = Arc::new(Barrier::new(threads)); + let mut handles = Vec::with_capacity(threads); + + for worker in 0..threads { + let pool = pool.clone(); + let barrier = barrier.clone(); + handles.push(thread::spawn(move || { + let base = (worker as u64) * EVICTION_KEYS_PER_THREAD; + barrier.wait(); + let before = Instant::now(); + for i in 0..THREADED_ITERATIONS { + let key = base + (i as u64 % EVICTION_KEYS_PER_THREAD); + let meta = ConnectionMeta::new(key, key as _); + pool.put(&meta, black_box(worker)); + } + before.elapsed() + })); + } + + let elapsed: Duration = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .sum(); + print_result( + &format!("pool eviction {threads:>2} threads"), + elapsed / threads as u32, + THREADED_ITERATIONS, + ); +} + +fn bench_pool_scaling() { + println!("ConnectionPool get+put scaling"); + for threads in THREAD_COUNTS { + bench_pool_pattern("distinct-keys", pool_distinct_keys(threads)); + } + + println!("ConnectionPool eviction-pressure scaling"); + for threads in THREAD_COUNTS { + bench_pool_eviction_pressure(threads); + } +} + +fn lru_shard(key: &i32) -> usize { + let mut hasher = DefaultHasher::new(); + key.hash(&mut hasher); + hasher.finish() as usize % LRU_SHARDS +} + +fn keys_spread_across_shards(threads: usize) -> Vec { + let mut keys = Vec::with_capacity(threads); + let mut used = [false; LRU_SHARDS]; + let mut key = 0; + while keys.len() < threads { + let shard = lru_shard(&key); + if !used[shard] { + used[shard] = true; + keys.push(key); + } + key += 1; + } + assert_eq!(keys.len(), threads); + assert_eq!(used.iter().filter(|used| **used).count(), threads); + keys +} + +fn keys_on_one_shard(threads: usize) -> Vec { + let mut keys = Vec::with_capacity(threads); + let mut key = 0; + while keys.len() < threads { + if lru_shard(&key) == 0 { + keys.push(key); + } + key += 1; + } + assert_eq!(keys.len(), threads); + assert!(keys.iter().all(|key| lru_shard(key) == 0)); + keys +} + +fn bench_lru_pattern(label: &str, keys: Vec) { + let threads = keys.len(); + let lru = Arc::new(Lru::new(POOL_SIZE)); + let barrier = Arc::new(Barrier::new(threads)); + let mut handles = Vec::with_capacity(threads); + + for key in keys { + let lru = lru.clone(); + let barrier = barrier.clone(); + handles.push(thread::spawn(move || { + barrier.wait(); + let before = Instant::now(); + for _ in 0..LRU_ITERATIONS { + let (_, evicted) = lru.add(key, ()); + black_box(evicted); + black_box(lru.pop(&key)); + } + before.elapsed() + })); + } + + let elapsed: Duration = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .sum(); + print_result( + &format!("lru {label} {threads:>2} threads"), + elapsed / threads as u32, + LRU_ITERATIONS, + ); +} + +fn bench_lru_scaling() { + println!("LRU put+pop scaling"); + for threads in THREAD_COUNTS { + bench_lru_pattern("spread", keys_spread_across_shards(threads)); + bench_lru_pattern("1-shard", keys_on_one_shard(threads)); + } +} + +fn main() { + println!("pingora-pool benchmarks"); + bench_round_trip(); + bench_eviction_pressure(); + bench_pool_scaling(); + bench_lru_scaling(); +} diff --git a/pingora-pool/src/connection.rs b/pingora-pool/src/connection.rs index 47dad5324..9f4600ba7 100644 --- a/pingora-pool/src/connection.rs +++ b/pingora-pool/src/connection.rs @@ -14,8 +14,9 @@ //! Generic connection pooling +use dashmap::DashMap; use log::{debug, warn}; -use parking_lot::{Mutex, RwLock}; +use parking_lot::Mutex; use pingora_timeout::{sleep, timeout}; use std::collections::HashMap; use std::io; @@ -25,6 +26,7 @@ use tokio::io::{AsyncRead, AsyncReadExt}; use tokio::sync::{oneshot, watch, Notify, OwnedMutexGuard}; use super::lru::Lru; +use futures::FutureExt; type GroupKey = u64; #[cfg(unix)] @@ -176,48 +178,40 @@ impl PoolNode { } } +type Pool = PoolNode>; + /// Connection pool /// /// [ConnectionPool] holds reusable connections. A reusable connection is released to this pool to /// be picked up by another user/request. pub struct ConnectionPool { - // TODO: n-way pools to reduce lock contention - pool: RwLock>>>>, + // Concurrent per-key pool index; each value handles per-key connection storage. + pools: DashMap>>, lru: Lru, } impl ConnectionPool { - /// Create a new [ConnectionPool] with a size limit. + /// Create a new [ConnectionPool] with a global size limit. /// - /// When a connection is released to this pool, the least recently used connection will be dropped. + /// When a connection is released to this pool and total occupancy is at + /// or above `size`, the least recently used connection is dropped. pub fn new(size: usize) -> Self { ConnectionPool { - pool: RwLock::new(HashMap::with_capacity(size)), // this is oversized since some connections will have the same key + pools: DashMap::with_capacity(size), lru: Lru::new(size), } } - /* get or create and insert a pool node for the hash key */ - fn get_pool_node(&self, key: GroupKey) -> Arc>> { - { - let pool = self.pool.read(); - if let Some(v) = pool.get(&key) { - return (*v).clone(); - } - } // read lock released here - - { - // write lock section - let mut pool = self.pool.write(); - // check again since another task might have already added it - if let Some(v) = pool.get(&key) { - return (*v).clone(); - } - let node = Arc::new(PoolNode::new()); - let node_ret = node.clone(); - pool.insert(key, node); // TODO: check dup - node_ret - } + /// Insert a connection under `key` while the DashMap entry guard is held. + /// + /// Holding the guard through [`PoolNode::insert`] prevents empty-node cleanup + /// from removing the map entry between looking it up and repopulating it. + fn insert_pool_connection(&self, key: GroupKey, id: ID, connection: PoolConnection) { + let pool_node = self + .pools + .entry(key) + .or_insert_with(|| Arc::new(PoolNode::new())); + pool_node.insert(id, connection); } /// Attempt to remove an empty [`PoolNode`] entry from the pool `HashMap`. @@ -231,39 +225,32 @@ impl ConnectionPool { /// removing a node that was concurrently repopulated between the caller's /// initial `is_empty()` hint and this write-lock acquisition. /// - /// # Race window - /// - /// There is a narrow window where another thread could have called - /// [`get_pool_node`] (obtaining a clone of the `Arc`) just before - /// we remove the entry. If that thread then inserts a connection into the - /// now-orphaned node, the connection is dropped when the last `Arc` reference - /// goes away. This is benign: the `oneshot::Sender` inside the dropped - /// `PoolConnection` is also dropped, which resolves the corresponding - /// `watch_use` receiver in `idle_poll`/`idle_timeout`, causing a clean exit. - /// The next request to the same upstream simply creates a fresh connection. - /// This trade-off matches the existing concurrency model of the pool and is - /// consistent with how hyper-util and Go's `net/http` handle this case. + /// Insertions go through [`Self::insert_pool_connection`], which holds the + /// DashMap entry guard until the connection is in the node. That prevents + /// this cleanup from removing a node between an inserter's entry lookup and + /// its [`PoolNode::insert`] call. fn try_remove_empty_node(&self, key: GroupKey) { - let mut pool = self.pool.write(); - if let Some(node) = pool.get(&key) { + if let Some(node) = self.pools.get(&key) { if node.is_empty() { - pool.remove(&key); + // Release the DashMap read guard before remove_if() acquires + // mutable access to the same shard. Re-check emptiness in the + // predicate because another thread may repopulate the node in + // between dropping this guard and attempting removal. + drop(node); + self.pools.remove_if(&key, |_, node| node.is_empty()); } } } // only remove from the pool because lru already removed it fn pop_evicted(&self, meta: &ConnectionMeta) { - let pool_node = { - let pool = self.pool.read(); - match pool.get(&meta.key) { - Some(v) => (*v).clone(), - None => { - warn!("Fail to get pool node for {:?}", meta); - return; - } // nothing to pop, should return error? - } - }; // read lock released here + let pool_node = match self.pools.get(&meta.key) { + Some(v) => v.value().clone(), + None => { + warn!("Fail to get pool node for {meta:?}"); + return; + } // nothing to pop, should return error? + }; pool_node.remove(meta.id); debug!("evict fd: {} from key {}", meta.id, meta.key); @@ -285,13 +272,10 @@ impl ConnectionPool { /// Get a connection from this pool under the same group key pub fn get(&self, key: &GroupKey) -> Option { - let pool_node = { - let pool = self.pool.read(); - match pool.get(key) { - Some(v) => (*v).clone(), - None => return None, - } - }; // read lock released here + let pool_node = match self.pools.get(key) { + Some(v) => v.value().clone(), + None => return None, + }; if let Some((id, connection)) = pool_node.get_any() { self.lru.pop(&id); // the notified is not needed @@ -314,21 +298,20 @@ impl ConnectionPool { /// Release a connection to this pool for reuse /// - /// - The returned [`Arc`] will notify any listen when the connection is evicted from the pool. + /// - The returned [`Arc`] will notify any listener when the connection is evicted from the pool. /// - The returned [`oneshot::Receiver`] will notify when the connection is being picked up by [Self::get()]. pub fn put( &self, meta: &ConnectionMeta, connection: S, ) -> (Arc, oneshot::Receiver) { - let (notify_close, replaced) = self.lru.add(meta.id, meta.clone()); - if let Some(meta) = replaced { - self.pop_evicted(&meta); - }; - let pool_node = self.get_pool_node(meta.key); + let (notify_close, evicted) = self.lru.add(meta.id, meta.clone()); + for meta in &evicted { + self.pop_evicted(meta); + } let (notify_use, watch_use) = oneshot::channel(); let connection = PoolConnection::new(notify_use, connection); - pool_node.insert(meta.id, connection); + self.insert_pool_connection(meta.key, meta.id, connection); (notify_close, watch_use) } @@ -338,6 +321,8 @@ impl ConnectionPool { /// remove it from the pool and drop the connection. /// /// If the connection is reused via [Self::get()] or being evicted, this function will just exit. + /// + /// Returns `true` if the connection was evicted from the pool, and `false` otherwise. pub async fn idle_poll( &self, connection: OwnedMutexGuard, @@ -345,44 +330,65 @@ impl ConnectionPool { timeout: Option, notify_evicted: Arc, watch_use: oneshot::Receiver, - ) where + ) -> bool + where Stream: AsyncRead + Unpin + Send, { + // Reuse this same Notified future in the watch_use branch: notify_one() + // may deliver the wakeup to an already-polled future, so creating a new + // notified() future after watch_use resolves could miss the eviction. + let evicted = notify_evicted.notified(); + tokio::pin!(evicted); + let read_result = tokio::select! { biased; - _ = watch_use => { - debug!("idle connection is being picked up"); - return + event = watch_use => { + return match event { + Ok(_) => { + debug!("idle connection is being picked up"); + false + } + // `watch_use` also resolves when the sender is dropped. + // During LRU eviction, pop_evicted() removes the + // PoolConnection, dropping the sender after notify_evicted + // has been signaled. Keep this biased branch first for the + // common reuse path, but confirm the eviction signal before + // classifying sender drop as eviction. + Err(_) => evicted.now_or_never().is_some(), + }; }, - _ = notify_evicted.notified() => { + _ = &mut evicted => { debug!("idle connection is being evicted"); // TODO: gracefully close the connection? - return - } + return true + }, read_result = read_with_timeout(connection , timeout) => read_result }; match read_result { Ok(n) => { if n > 0 { - warn!("Data received on idle client connection, close it") + warn!("Data received on idle client connection, close it"); } else { - debug!("Peer closed the idle connection or timeout") + debug!("Peer closed the idle connection or timeout"); } } Err(e) => { debug!("error with the idle connection, close it {:?}", e); } - } + }; // connection terminated from either peer or timer self.pop_closed(meta); + false } /// Passively wait to close the connection after the timeout /// /// If this connection is not being picked up or evicted before the timeout is reach, this /// function will remove it from the pool and close the connection. + /// + /// Returns `true` if the connection was evicted from the pool, and `false` otherwise. pub async fn idle_timeout( &self, meta: &ConnectionMeta, @@ -390,27 +396,48 @@ impl ConnectionPool { notify_evicted: Arc, mut notify_closed: watch::Receiver, watch_use: oneshot::Receiver, - ) { + ) -> bool { + // Reuse this same Notified future in the watch_use branch: notify_one() + // may deliver the wakeup to an already-polled future, so creating a new + // notified() future after watch_use resolves could miss the eviction. + let evicted = notify_evicted.notified(); + tokio::pin!(evicted); + tokio::select! { biased; - _ = watch_use => { - debug!("idle connection is being picked up"); + event = watch_use => { + match event { + Ok(_) => { + debug!("idle connection is being picked up"); + false + } + // `watch_use` also resolves when the sender is dropped. + // During LRU eviction, pop_evicted() removes the + // PoolConnection, dropping the sender after notify_evicted + // has been signaled. Keep this biased branch first for the + // common reuse path, but confirm the eviction signal before + // classifying sender drop as eviction. + Err(_) => evicted.now_or_never().is_some(), + } }, - _ = notify_evicted.notified() => { + _ = &mut evicted => { debug!("idle connection is being evicted"); // TODO: gracefully close the connection? - } + true + }, _ = notify_closed.changed() => { // assume always changed from false to true debug!("idle connection is being closed"); self.pop_closed(meta); + false } // async expression is evaluated if timeout is None but it's never polled, set it to MAX _ = sleep(timeout.unwrap_or(Duration::MAX)), if timeout.is_some() => { debug!("idle connection is being evicted"); self.pop_closed(meta); + false } - }; + } } } @@ -442,6 +469,14 @@ mod tests { use tokio::sync::Mutex as AsyncMutex; use tokio_test::io::{Builder, Mock}; + fn pool_len(pool: &ConnectionPool) -> usize { + pool.pools.len() + } + + fn pool_contains(pool: &ConnectionPool, key: GroupKey) -> bool { + pool.pools.contains_key(&key) + } + #[tokio::test] async fn test_lookup() { let meta1 = ConnectionMeta::new(101, 1); @@ -533,8 +568,8 @@ mod tests { let closed_item = tokio::select! { _ = cp.idle_poll(mock_io1.try_lock_owned().unwrap(), &meta1, None, c1, u1) => {debug!("notifier1"); 1}, - _ = cp.idle_poll(mock_io2.try_lock_owned().unwrap(), &meta1, None, c2, u2) => {debug!("notifier2"); 2}, - _ = cp.idle_poll(mock_io3.try_lock_owned().unwrap(), &meta1, None, c3, u3) => {debug!("notifier3"); 3}, + _ = cp.idle_poll(mock_io2.try_lock_owned().unwrap(), &meta2, None, c2, u2) => {debug!("notifier2"); 2}, + _ = cp.idle_poll(mock_io3.try_lock_owned().unwrap(), &meta3, None, c3, u3) => {debug!("notifier3"); 3}, }; assert_eq!(closed_item, 1); @@ -563,8 +598,8 @@ mod tests { let closed_item = tokio::select! { _ = cp.idle_poll(mock_io1.try_lock_owned().unwrap(), &meta1, Some(Duration::from_secs(1)), c1, u1) => {debug!("notifier1"); 1}, - _ = cp.idle_poll(mock_io2.try_lock_owned().unwrap(), &meta1, Some(Duration::from_secs(2)), c2, u2) => {debug!("notifier2"); 2}, - _ = cp.idle_poll(mock_io3.try_lock_owned().unwrap(), &meta1, Some(Duration::from_secs(3)), c3, u3) => {debug!("notifier3"); 3}, + _ = cp.idle_poll(mock_io2.try_lock_owned().unwrap(), &meta2, Some(Duration::from_secs(2)), c2, u2) => {debug!("notifier2"); 2}, + _ = cp.idle_poll(mock_io3.try_lock_owned().unwrap(), &meta3, Some(Duration::from_secs(3)), c3, u3) => {debug!("notifier3"); 3}, }; assert_eq!(closed_item, 1); @@ -593,8 +628,8 @@ mod tests { let closed_item = tokio::select! { _ = cp.idle_poll(mock_io1.try_lock_owned().unwrap(), &meta1, None, c1, u1) => {debug!("notifier1"); 1}, - _ = cp.idle_poll(mock_io2.try_lock_owned().unwrap(), &meta1, None, c2, u2) => {debug!("notifier2"); 2}, - _ = cp.idle_poll(mock_io3.try_lock_owned().unwrap(), &meta1, None, c3, u3) => {debug!("notifier3"); 3}, + _ = cp.idle_poll(mock_io2.try_lock_owned().unwrap(), &meta2, None, c2, u2) => {debug!("notifier2"); 2}, + _ = cp.idle_poll(mock_io3.try_lock_owned().unwrap(), &meta3, None, c3, u3) => {debug!("notifier3"); 3}, }; assert_eq!(closed_item, 1); @@ -602,6 +637,340 @@ mod tests { assert!(cp.get(&meta1.key).is_none()) // mock_io1 should already be removed by idle_poll } + #[tokio::test] + async fn test_idle_poll_reports_notify_evicted() { + let meta1 = ConnectionMeta::new(101, 1); + let mock_io1 = Arc::new(AsyncMutex::new( + Builder::new().wait(Duration::from_secs(99)).build(), + )); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta1, mock_io1.clone()); + notify_evicted.notify_one(); + + let evicted = cp + .idle_poll( + mock_io1.try_lock_owned().unwrap(), + &meta1, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!(evicted, "notify_evicted should report eviction"); + } + + #[tokio::test] + async fn test_idle_poll_reports_lru_eviction_after_pool_remove() { + let meta1 = ConnectionMeta::new(101, 1); + let mock_io1 = Arc::new(AsyncMutex::new( + Builder::new().wait(Duration::from_secs(99)).build(), + )); + let meta2 = ConnectionMeta::new(202, 2); + let mock_io2 = Arc::new(AsyncMutex::new( + Builder::new().wait(Duration::from_secs(99)).build(), + )); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta1, mock_io1.clone()); + cp.put(&meta2, mock_io2); + + let evicted = cp + .idle_poll( + mock_io1.try_lock_owned().unwrap(), + &meta1, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!(evicted, "LRU eviction should report eviction"); + } + + #[tokio::test] + async fn test_idle_poll_reports_sender_drop_without_notify_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let mock_io = Arc::new(AsyncMutex::new( + Builder::new().wait(Duration::from_secs(99)).build(), + )); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta, mock_io.clone()); + cp.pop_closed(&meta); + + let evicted = cp + .idle_poll( + mock_io.try_lock_owned().unwrap(), + &meta, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!( + !evicted, + "sender drop without notify should not report eviction" + ); + } + + #[tokio::test] + async fn test_idle_poll_reports_reuse_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let mock_io = Arc::new(AsyncMutex::new( + Builder::new().wait(Duration::from_secs(99)).build(), + )); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta, mock_io.clone()); + assert!(cp.get(&meta.key).is_some()); + + let evicted = cp + .idle_poll( + mock_io.try_lock_owned().unwrap(), + &meta, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!(!evicted, "reused connection should not report eviction"); + } + + #[tokio::test] + async fn test_idle_poll_reports_peer_close_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let mock_io = Arc::new(AsyncMutex::new(Builder::new().read(b"").build())); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta, mock_io.clone()); + + let evicted = cp + .idle_poll( + mock_io.try_lock_owned().unwrap(), + &meta, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!(!evicted, "peer close should not report eviction"); + assert!(cp.get(&meta.key).is_none()); + } + + #[tokio::test] + async fn test_idle_poll_reports_unexpected_data_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let mock_io = Arc::new(AsyncMutex::new(Builder::new().read(b"x").build())); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta, mock_io.clone()); + + let evicted = cp + .idle_poll( + mock_io.try_lock_owned().unwrap(), + &meta, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!(!evicted, "unexpected data should not report eviction"); + assert!(cp.get(&meta.key).is_none()); + } + + #[tokio::test] + async fn test_idle_poll_reports_read_error_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let mock_io = Arc::new(AsyncMutex::new( + Builder::new() + .read_error(io::Error::other("read failed")) + .build(), + )); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta, mock_io.clone()); + + let evicted = cp + .idle_poll( + mock_io.try_lock_owned().unwrap(), + &meta, + None, + notify_evicted, + watch_use, + ) + .await; + + assert!(!evicted, "read error should not report eviction"); + assert!(cp.get(&meta.key).is_none()); + } + + #[tokio::test] + async fn test_idle_poll_reports_timeout_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let mock_io = Arc::new(AsyncMutex::new( + Builder::new().wait(Duration::from_secs(99)).build(), + )); + let cp: ConnectionPool>> = ConnectionPool::new(1); + + let (notify_evicted, watch_use) = cp.put(&meta, mock_io.clone()); + + let evicted = cp + .idle_poll( + mock_io.try_lock_owned().unwrap(), + &meta, + Some(Duration::from_millis(10)), + notify_evicted, + watch_use, + ) + .await; + + assert!(!evicted, "idle poll timeout should not report eviction"); + assert!(cp.get(&meta.key).is_none()); + } + + #[tokio::test] + async fn test_idle_timeout_reports_timeout_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let cp: ConnectionPool = ConnectionPool::new(1); + let (notify_evicted, watch_use) = cp.put(&meta, "v1".to_string()); + let (_notify_closed, notify_closed_rx) = watch::channel(false); + + let evicted = cp + .idle_timeout( + &meta, + Some(Duration::from_millis(10)), + notify_evicted, + notify_closed_rx, + watch_use, + ) + .await; + + assert!(!evicted, "idle timeout should not report eviction"); + assert!(cp.get(&meta.key).is_none()); + } + + #[tokio::test] + async fn test_idle_timeout_reports_reuse_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let cp: ConnectionPool = ConnectionPool::new(1); + let (notify_evicted, watch_use) = cp.put(&meta, "v1".to_string()); + let (_notify_closed, notify_closed_rx) = watch::channel(false); + + assert_eq!(cp.get(&meta.key), Some("v1".to_string())); + + let evicted = cp + .idle_timeout(&meta, None, notify_evicted, notify_closed_rx, watch_use) + .await; + + assert!(!evicted, "reused connection should not report eviction"); + } + + #[tokio::test] + async fn test_idle_timeout_reports_notify_evicted() { + let meta = ConnectionMeta::new(101, 1); + let cp: ConnectionPool = ConnectionPool::new(1); + let (notify_evicted, watch_use) = cp.put(&meta, "v1".to_string()); + let (_notify_closed, notify_closed_rx) = watch::channel(false); + + notify_evicted.notify_one(); + + let evicted = cp + .idle_timeout(&meta, None, notify_evicted, notify_closed_rx, watch_use) + .await; + + assert!(evicted, "notify_evicted should report eviction"); + } + + #[tokio::test] + async fn test_idle_timeout_reports_lru_eviction_after_pool_remove() { + let meta1 = ConnectionMeta::new(101, 1); + let meta2 = ConnectionMeta::new(202, 2); + let cp: ConnectionPool = ConnectionPool::new(1); + let (notify_evicted, watch_use) = cp.put(&meta1, "v1".to_string()); + let (_notify_closed, notify_closed_rx) = watch::channel(false); + + cp.put(&meta2, "v2".to_string()); + + let evicted = cp + .idle_timeout(&meta1, None, notify_evicted, notify_closed_rx, watch_use) + .await; + + assert!(evicted, "LRU eviction should report eviction"); + } + + #[tokio::test] + async fn test_idle_timeout_reports_lru_eviction_after_notify_registered() { + let meta1 = ConnectionMeta::new(101, 1); + let meta2 = ConnectionMeta::new(202, 2); + let cp = Arc::new(ConnectionPool::new(1)); + let (notify_evicted, watch_use) = cp.put(&meta1, "v1".to_string()); + let (_notify_closed, notify_closed_rx) = watch::channel(false); + + let idle_cp = cp.clone(); + let idle_meta = meta1.clone(); + let idle_task = tokio::spawn(async move { + idle_cp + .idle_timeout( + &idle_meta, + None, + notify_evicted, + notify_closed_rx, + watch_use, + ) + .await + }); + + tokio::task::yield_now().await; + cp.put(&meta2, "v2".to_string()); + + assert!( + idle_task.await.unwrap(), + "LRU eviction should report eviction after notify future was registered" + ); + } + + #[tokio::test] + async fn test_idle_timeout_reports_sender_drop_without_notify_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let cp: ConnectionPool = ConnectionPool::new(1); + let (notify_evicted, watch_use) = cp.put(&meta, "v1".to_string()); + let (_notify_closed, notify_closed_rx) = watch::channel(false); + + cp.pop_closed(&meta); + + let evicted = cp + .idle_timeout(&meta, None, notify_evicted, notify_closed_rx, watch_use) + .await; + + assert!( + !evicted, + "sender drop without notify should not report eviction" + ); + } + + #[tokio::test] + async fn test_idle_timeout_reports_notify_closed_not_evicted() { + let meta = ConnectionMeta::new(101, 1); + let cp: ConnectionPool = ConnectionPool::new(1); + let (notify_evicted, watch_use) = cp.put(&meta, "v1".to_string()); + let (notify_closed, notify_closed_rx) = watch::channel(false); + + notify_closed.send(true).unwrap(); + + let evicted = cp + .idle_timeout(&meta, None, notify_evicted, notify_closed_rx, watch_use) + .await; + + assert!(!evicted, "notify_closed should not report eviction"); + assert!(cp.get(&meta.key).is_none()); + } + #[test] fn test_pool_node_is_empty() { let node: PoolNode = PoolNode::new(); @@ -648,12 +1017,12 @@ mod tests { let cp: ConnectionPool = ConnectionPool::new(2); cp.put(&meta, "v1".to_string()); - assert_eq!(cp.pool.read().len(), 1, "pool should have 1 node"); + assert_eq!(pool_len(&cp), 1, "pool should have 1 node"); cp.pop_closed(&meta); assert_eq!( - cp.pool.read().len(), + pool_len(&cp), 0, "empty PoolNode should be removed after pop_closed" ); @@ -669,13 +1038,13 @@ mod tests { let cp: ConnectionPool = ConnectionPool::new(2); cp.put(&meta, "v1".to_string()); - assert_eq!(cp.pool.read().len(), 1); + assert_eq!(pool_len(&cp), 1); let conn = cp.get(&meta.key); assert!(conn.is_some()); assert_eq!( - cp.pool.read().len(), + pool_len(&cp), 0, "empty PoolNode should be removed after get() takes the last connection" ); @@ -694,11 +1063,11 @@ mod tests { // Remove both connections via pop_closed, but the first pop_closed // won't remove the node since meta2 is still there. cp.pop_closed(&meta1); - assert_eq!(cp.pool.read().len(), 1, "node should still exist"); + assert_eq!(pool_len(&cp), 1, "node should still exist"); cp.pop_closed(&meta2); assert_eq!( - cp.pool.read().len(), + pool_len(&cp), 0, "node should be removed after last connection is popped" ); @@ -717,10 +1086,10 @@ mod tests { cp.pop_closed(&meta1); assert!( - cp.pool.read().contains_key(&101), + pool_contains(&cp, 101), "node should still exist because meta2's connection is still in it" ); - assert_eq!(cp.pool.read().len(), 1); + assert_eq!(pool_len(&cp), 1); // The remaining connection should still be retrievable let conn = cp.get(&meta1.key); @@ -736,18 +1105,18 @@ mod tests { cp.put(&meta_a, "a".to_string()); cp.put(&meta_b, "b".to_string()); - assert_eq!(cp.pool.read().len(), 2); + assert_eq!(pool_len(&cp), 2); // Remove all connections for key 101 cp.pop_closed(&meta_a); assert_eq!( - cp.pool.read().len(), + pool_len(&cp), 1, "only key 101's empty node should be removed" ); - assert!(!cp.pool.read().contains_key(&101), "key 101 should be gone"); - assert!(cp.pool.read().contains_key(&202), "key 202 should remain"); + assert!(!pool_contains(&cp, 101), "key 101 should be gone"); + assert!(pool_contains(&cp, 202), "key 202 should remain"); // key 202's connection should still be retrievable let conn = cp.get(&meta_b.key); @@ -763,16 +1132,47 @@ mod tests { let cp: ConnectionPool = ConnectionPool::new(1); cp.put(&meta1, "v1".to_string()); - assert_eq!(cp.pool.read().len(), 1); + assert_eq!(pool_len(&cp), 1); // This put evicts meta1 (LRU size = 1), making key 101's node empty. cp.put(&meta2, "v2".to_string()); assert!( - !cp.pool.read().contains_key(&101), + !pool_contains(&cp, 101), "key 101's empty node should be removed after its only connection was evicted" ); - assert!(cp.pool.read().contains_key(&202)); + assert!(pool_contains(&cp, 202)); + } + + #[test] + fn test_concurrent_empty_node_cleanup_does_not_orphan_put() { + const KEY: GroupKey = 101; + let cp = Arc::new(ConnectionPool::new(2_000)); + let start = Arc::new(std::sync::Barrier::new(2)); + + let cleanup_cp = cp.clone(); + let cleanup_start = start.clone(); + let cleanup = std::thread::spawn(move || { + cleanup_start.wait(); + for _ in 0..10_000 { + cleanup_cp.try_remove_empty_node(KEY); + std::thread::yield_now(); + } + }); + + start.wait(); + for id in 1..=1_000 { + let value = format!("v{id}"); + cp.put(&ConnectionMeta::new(KEY, id), value.clone()); + assert_eq!( + cp.get(&KEY), + Some(value), + "put connection should remain reachable during empty-node cleanup" + ); + std::thread::yield_now(); + } + + cleanup.join().unwrap(); } #[tokio::test] @@ -784,18 +1184,18 @@ mod tests { cp.put(&meta1, "first".to_string()); cp.pop_closed(&meta1); - assert_eq!(cp.pool.read().len(), 0, "node should be cleaned up"); + assert_eq!(pool_len(&cp), 0, "node should be cleaned up"); // Re-insert for the same key let meta2 = ConnectionMeta::new(101, 2); cp.put(&meta2, "second".to_string()); - assert_eq!(cp.pool.read().len(), 1); + assert_eq!(pool_len(&cp), 1); let conn = cp.get(&meta2.key); assert_eq!(conn, Some("second".to_string())); assert_eq!( - cp.pool.read().len(), + pool_len(&cp), 0, "node should be cleaned up again after get" ); diff --git a/pingora-pool/src/lru.rs b/pingora-pool/src/lru.rs index c6a72d8ac..f128af07b 100644 --- a/pingora-pool/src/lru.rs +++ b/pingora-pool/src/lru.rs @@ -12,18 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use core::hash::Hash; +use core::hash::{Hash, Hasher}; use lru::LruCache; -use parking_lot::RwLock; -use std::cell::RefCell; -use std::sync::atomic::{AtomicBool, Ordering::Relaxed}; +use parking_lot::Mutex; +use std::collections::hash_map::DefaultHasher; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering::Relaxed}; use std::sync::Arc; -use thread_local::ThreadLocal; use tokio::sync::Notify; pub struct Node { pub close_notifier: Arc, pub meta: T, + order: u64, } impl Node { @@ -31,6 +31,7 @@ impl Node { Node { close_notifier: Arc::new(Notify::new()), meta, + order: 0, } } @@ -39,12 +40,28 @@ impl Node { } } +/// Number of independent LRU shards. Independent of the `N_SHARDS` constant +/// in `crate::connection`. `pub` only so the in-crate benchmark can use it. +pub const N_SHARDS: usize = 16; + +/// Bounded, key-sharded LRU used to age out idle connections. +/// +/// Each key is hashed into one of [`N_SHARDS`] shards, so the same key +/// always lands in the same shard. `put` and `pop` each take one shard +/// lock on the fast path. When occupancy exceeds `size`, one caller drains +/// the overflow globally in a single batch. pub struct Lru where K: Send, T: Send, { - lru: RwLock>>>>, + /// Per-shard caches. Individually unbounded; global cap enforced via `len`. + lrus: [Mutex>>; N_SHARDS], + /// Total live entries across shards. Lower `order` = older entry. + len: AtomicUsize, + /// Monotonic stamp assigned at insert; used to compare age across shards. + order: AtomicU64, + /// Global capacity. Eviction runs when `len` exceeds this. size: usize, drain: AtomicBool, } @@ -56,75 +73,121 @@ where { pub fn new(size: usize) -> Self { Lru { - lru: RwLock::new(ThreadLocal::new()), + lrus: std::array::from_fn(|_| Mutex::new(LruCache::unbounded())), + len: AtomicUsize::new(0), + order: AtomicU64::new(0), size, drain: AtomicBool::new(false), } } - // put a node in and return the meta of the replaced node - pub fn put(&self, key: K, value: Node) -> Option { - if self.drain.load(Relaxed) { - value.notify_close(); // sort of hack to simulate being evicted right away - return None; + #[inline] + fn shard(&self, key: &K) -> &Mutex>> { + let mut hasher = DefaultHasher::new(); + key.hash(&mut hasher); + &self.lrus[(hasher.finish() as usize) % N_SHARDS] + } + + /// Drain to `size - headroom` by globally oldest `order`. + /// + /// Holds every shard lock for the duration to serialize concurrent + /// evictors and prevent victim-selection races. Returns the metas of + /// evicted nodes; `notify_close` has fired on each before return. + fn evict_lru(&self) -> Vec { + // Fast path: another evictor may have already drained the overflow. + if self.len.load(Relaxed) <= self.size { + return Vec::new(); } - let lru = self.lru.read(); /* read lock */ - let lru_cache = &mut *(lru - .get_or(|| RefCell::new(LruCache::unbounded())) - .borrow_mut()); - lru_cache.put(key, value); - if lru_cache.len() > self.size { - match lru_cache.pop_lru() { - Some((_, v)) => { - // TODO: drop the lock here? - v.notify_close(); - return Some(v.meta); + + let headroom = self.evict_headroom(); + let drain_target = self.size.saturating_sub(headroom); + let mut lrus: Vec<_> = self.lrus.iter().map(|lru| lru.lock()).collect(); + // A single over-capacity insert usually drains headroom + 1 entries. + // Concurrent inserts can grow the batch further, so this is only a hint. + let mut evicted_metas = Vec::with_capacity(headroom + 1); + while self.len.load(Relaxed) > drain_target { + let mut oldest = None; + for (index, lru) in lrus.iter().enumerate() { + if let Some((_, node)) = lru.peek_lru() { + if oldest.is_none_or(|(_, order)| node.order < order) { + oldest = Some((index, node.order)); + } } - None => return None, } + + let Some((index, _)) = oldest else { break }; + let Some((_, evicted)) = lrus[index].pop_lru() else { + break; + }; + self.len.fetch_sub(1, Relaxed); + evicted.notify_close(); + evicted_metas.push(evicted.meta); + } + evicted_metas + } + + /// Headroom below `size`. Amortizes the all-shards-lock cost over + /// several subsequent over-capacity inserts. Capped at [`N_SHARDS`] + /// and `size / 4` so tiny pools keep tight semantics. + #[inline] + fn evict_headroom(&self) -> usize { + N_SHARDS.min(self.size / 4) + } + + /// Insert `value` for `key`. Returns evicted metas, oldest first. + /// + /// `len` is only incremented on new inserts (not replacements), so + /// over-capacity eviction tracks live entries. + pub fn put(&self, key: K, mut value: Node) -> Vec { + if self.drain.load(Relaxed) { + value.notify_close(); // sort of hack to simulate being evicted right away + return Vec::new(); + } + + value.order = self.order.fetch_add(1, Relaxed); + if self.shard(&key).lock().put(key, value).is_some() { + // replaced + return Vec::new(); + } + if self.len.fetch_add(1, Relaxed) + 1 > self.size { + return self.evict_lru(); } - None - /* read lock dropped */ + Vec::new() } - pub fn add(&self, key: K, meta: T) -> (Arc, Option) { + pub fn add(&self, key: K, meta: T) -> (Arc, Vec) { let node = Node::new(meta); let notifier = node.close_notifier.clone(); - // TODO: check if the key is already in it (notifier, self.put(key, node)) } pub fn pop(&self, key: &K) -> Option> { - let lru = self.lru.read(); /* read lock */ - let lru_cache = &mut *(lru - .get_or(|| RefCell::new(LruCache::unbounded())) - .borrow_mut()); - lru_cache.pop(key) - /* read lock dropped */ + let popped = self.shard(key).lock().pop(key); + if popped.is_some() { + self.len.fetch_sub(1, Relaxed); + } + popped } #[allow(dead_code)] pub fn drain(&self) { self.drain.store(true, Relaxed); - /* drain need to go through all the local lru cache objects - * acquire an exclusive write lock to make it safe */ - let mut lru = self.lru.write(); /* write lock */ - let lru_cache_iter = lru.iter_mut(); - for lru_cache_rc in lru_cache_iter { - let mut lru_cache = lru_cache_rc.borrow_mut(); + for lru in &self.lrus { + let mut lru_cache = lru.lock(); for (_, item) in lru_cache.iter() { item.notify_close(); } lru_cache.clear(); } - /* write lock dropped */ + self.len.store(0, Relaxed); } } #[cfg(test)] mod tests { use super::*; + use futures::FutureExt; use log::debug; #[tokio::test] @@ -158,6 +221,22 @@ mod tests { assert_eq!(closed_item, 2); } + #[tokio::test] + async fn test_replaced_node_is_not_treated_as_eviction() { + let pool: Lru = Lru::new(2); + let (replaced_notifier, evicted) = pool.add(1, 10); + assert!(evicted.is_empty()); + + let (_, evicted) = pool.add(1, 20); + assert!(evicted.is_empty()); + assert_eq!(pool.pop(&1).unwrap().meta, 20); + + assert!( + replaced_notifier.notified().now_or_never().is_none(), + "replaced node should not notify close" + ); + } + #[tokio::test] async fn test_drain() { let pool: Lru = Lru::new(4); diff --git a/pingora-prometheus/Cargo.toml b/pingora-prometheus/Cargo.toml new file mode 100644 index 000000000..d97012131 --- /dev/null +++ b/pingora-prometheus/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "pingora-prometheus" +version = "0.8.0" +authors = ["Pingora Team at Cloudflare "] +license = "Apache-2.0" +edition = "2021" +repository = "https://github.com/cloudflare/pingora" +categories = ["asynchronous", "network-programming"] +keywords = ["async", "http", "prometheus", "pingora"] +description = """ +A Prometheus metrics HTTP server for pingora services. +""" + +[lib] +name = "pingora_prometheus" +path = "src/lib.rs" + +[dependencies] +pingora-core = { version = "0.8.0", path = "../pingora-core", default-features = false } +prometheus = "0.14" +async-trait = { workspace = true } +http = { workspace = true } diff --git a/pingora-prometheus/src/lib.rs b/pingora-prometheus/src/lib.rs new file mode 100644 index 000000000..cfd90b898 --- /dev/null +++ b/pingora-prometheus/src/lib.rs @@ -0,0 +1,131 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#![warn(clippy::all)] + +//! A Prometheus metrics HTTP server for [pingora](https://docs.rs/pingora) services. +//! +//! This crate provides [`PrometheusHttpApp`] and [`PrometheusServer`], which serve +//! all [static metrics](https://docs.rs/prometheus/latest/prometheus/index.html#static-metrics) +//! collected via the [`prometheus`] crate as an HTTP endpoint. +//! +//! # Example +//! +//! ```rust,ignore +//! use pingora_core::services::listening::Service; +//! use pingora_prometheus::new_prometheus_server; +//! +//! let mut prometheus_service = Service::new( +//! "Prometheus HTTP".to_string(), +//! new_prometheus_server(), +//! ); +//! prometheus_service.add_tcp("127.0.0.1:6150"); +//! server.add_service(prometheus_service); +//! ``` +//! +//! Or use the convenience function: +//! +//! ```rust,ignore +//! let mut prometheus_service = pingora_prometheus::prometheus_http_service(); +//! prometheus_service.add_tcp("127.0.0.1:6150"); +//! server.add_service(prometheus_service); +//! ``` + +use async_trait::async_trait; +use http::Response; +use prometheus::{Encoder, TextEncoder}; + +use pingora_core::apps::http_app::{HttpServer, ServeHttp}; +use pingora_core::modules::http::compression::ResponseCompressionBuilder; +use pingora_core::protocols::http::ServerSession; +use pingora_core::services::listening::Service; + +/// Re-export of the [`prometheus`] crate. +/// +/// Use this re-export to ensure your metrics are registered in the same +/// global registry that [`PrometheusHttpApp`] gathers from, avoiding +/// version mismatches that would cause metrics to silently not appear. +/// +/// # Example +/// +/// ```rust,ignore +/// use pingora_prometheus::prometheus::{self, register_int_counter, IntCounter}; +/// use once_cell::sync::Lazy; +/// +/// static REQUESTS: Lazy = Lazy::new(|| { +/// register_int_counter!("requests_total", "Total requests").unwrap() +/// }); +/// ``` +pub use prometheus; + +/// An HTTP application that reports Prometheus metrics. +/// +/// This application will report all the [static metrics](https://docs.rs/prometheus/latest/prometheus/index.html#static-metrics) +/// collected via the [Prometheus](https://docs.rs/prometheus/) crate. +/// +/// Currently serves metrics on all request paths. By convention, Prometheus +/// scrapers expect metrics at `/metrics`. Since this app is typically bound +/// to a dedicated listener address, this works in practice, but callers +/// should be aware of this if sharing the listener with other routes. +// TODO: consider restricting to `/metrics` and returning 404 for other paths +pub struct PrometheusHttpApp; + +#[async_trait] +impl ServeHttp for PrometheusHttpApp { + async fn response(&self, _http_session: &mut ServerSession) -> Response> { + let encoder = TextEncoder::new(); + let metric_families = prometheus::gather(); + let mut buffer = vec![]; + encoder.encode(&metric_families, &mut buffer).unwrap(); + Response::builder() + .status(200) + .header(http::header::CONTENT_TYPE, encoder.format_type()) + .header(http::header::CONTENT_LENGTH, buffer.len()) + .body(buffer) + .unwrap() + } +} + +/// The [`HttpServer`] for [`PrometheusHttpApp`]. +/// +/// This type provides the functionality of [`PrometheusHttpApp`] with gzip +/// compression enabled (level 7). +pub type PrometheusServer = HttpServer; + +/// Create a new [`PrometheusServer`] with compression enabled. +pub fn new_prometheus_server() -> PrometheusServer { + let mut server = PrometheusServer::new_app(PrometheusHttpApp); + // enable gzip level 7 compression + server.add_module(ResponseCompressionBuilder::enable(7)); + server +} + +/// Create a Prometheus HTTP [`Service`] ready to have endpoints added. +/// +/// This is a convenience function that creates a [`Service`] wrapping a +/// [`PrometheusServer`] with compression enabled. +/// +/// # Example +/// +/// ```rust,ignore +/// let mut prometheus_service = pingora_prometheus::prometheus_http_service(); +/// prometheus_service.add_tcp("127.0.0.1:6150"); +/// server.add_service(prometheus_service); +/// ``` +pub fn prometheus_http_service() -> Service { + Service::new( + "Prometheus metric HTTP".to_string(), + new_prometheus_server(), + ) +} diff --git a/pingora-proxy/Cargo.toml b/pingora-proxy/Cargo.toml index c685b8c41..d01e8b57b 100644 --- a/pingora-proxy/Cargo.toml +++ b/pingora-proxy/Cargo.toml @@ -25,7 +25,7 @@ pingora-cache = { version = "0.8.0", path = "../pingora-cache", default-features tokio = { workspace = true, features = ["macros", "net"] } pingora-http = { version = "0.8.0", path = "../pingora-http" } http = { workspace = true } -futures = "0.3" +futures = { workspace = true } bytes = { workspace = true } async-trait = { workspace = true } log = { workspace = true } @@ -36,17 +36,21 @@ regex = "1" rand = "0.8" [dev-dependencies] -reqwest = { version = "0.11", features = [ +reqwest = { version = "0.12", features = [ "gzip", "rustls-tls", + "http2", ], default-features = false } httparse = { workspace = true } tokio-test = "0.4" env_logger = "0.11" -hyper = "0.14" -tokio-tungstenite = "0.20.1" +hyper = { version = "1", features = ["client", "http1", "http2"] } +hyper-util = { version = "0.1", features = ["client-legacy", "http1", "http2"] } +http-body-util = "0.1" +tokio-tungstenite = "0.26" pingora-limits = { version = "0.8.0", path = "../pingora-limits" } pingora-load-balancing = { version = "0.8.0", path = "../pingora-load-balancing", default-features=false } +pingora-prometheus = { version = "0.8.0", path = "../pingora-prometheus" } prometheus = "0" futures-util = "0.3" serde = { version = "1.0", features = ["derive"] } @@ -54,7 +58,7 @@ serde_json = "1.0" serde_yaml = "0.9" [target.'cfg(unix)'.dev-dependencies] -hyperlocal = "0.8" +hyperlocal = "0.9" [features] default = [] @@ -69,7 +73,10 @@ s2n = ["pingora-core/s2n", "pingora-cache/s2n", "any_tls"] openssl_derived = ["any_tls"] any_tls = [] sentry = ["pingora-core/sentry"] +patched_http1 = ["pingora-core/patched_http1"] +upstream_modules = [] connection_filter = ["pingora-core/connection_filter"] +trace = ["pingora-cache/trace"] [[example]] name = "connection_filter" diff --git a/pingora-proxy/examples/connection_filter.rs b/pingora-proxy/examples/connection_filter.rs index 1c346c6fa..67bedcf1b 100644 --- a/pingora-proxy/examples/connection_filter.rs +++ b/pingora-proxy/examples/connection_filter.rs @@ -49,8 +49,11 @@ struct BlockAllFilter; #[async_trait] impl ConnectionFilter for BlockAllFilter { - async fn should_accept(&self, addr: &std::net::SocketAddr) -> bool { - info!("BLOCKING connection from {} (BlockAllFilter active)", addr); + async fn should_accept(&self, addr: Option<&std::net::SocketAddr>) -> bool { + info!( + "BLOCKING connection from {:?} (BlockAllFilter active)", + addr + ); false } } diff --git a/pingora-proxy/examples/gateway.rs b/pingora-proxy/examples/gateway.rs index 83b7c1caf..79c1646cc 100644 --- a/pingora-proxy/examples/gateway.rs +++ b/pingora-proxy/examples/gateway.rs @@ -129,8 +129,7 @@ fn main() { my_proxy.add_tcp("0.0.0.0:6191"); my_server.add_service(my_proxy); - let mut prometheus_service_http = - pingora_core::services::listening::Service::prometheus_http_service(); + let mut prometheus_service_http = pingora_prometheus::prometheus_http_service(); prometheus_service_http.add_tcp("127.0.0.1:6192"); my_server.add_service(prometheus_service_http); diff --git a/pingora-proxy/examples/pipelining.rs b/pingora-proxy/examples/pipelining.rs new file mode 100644 index 000000000..1d596200d --- /dev/null +++ b/pingora-proxy/examples/pipelining.rs @@ -0,0 +1,70 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use async_trait::async_trait; + +use pingora_core::server::configuration::Opt; +use pingora_core::server::Server; +use pingora_core::upstreams::peer::HttpPeer; +use pingora_core::Result; +use pingora_proxy::{ProxyHttp, Session}; + +pub struct PipelinedGateway; + +#[async_trait] +impl ProxyHttp for PipelinedGateway { + type CTX = (); + fn new_ctx(&self) -> Self::CTX {} + + async fn early_request_filter( + &self, + session: &mut Session, + _ctx: &mut Self::CTX, + ) -> Result<()> { + // Opt in once per session; persists across keep-alive reuses. + session.set_pipelining_enabled(true); + Ok(()) + } + + async fn upstream_peer( + &self, + _session: &mut Session, + _ctx: &mut Self::CTX, + ) -> Result> { + let peer = HttpPeer::new(("httpbin.org", 80), false, "httpbin.org".into()); + Ok(Box::new(peer)) + } +} + +// RUST_LOG=INFO cargo run --example pipelining +// +// Two pipelined GETs on one connection (expect two `HTTP/1.1 200` lines): +// printf 'GET /get HTTP/1.1\r\nHost: httpbin.org\r\n\r\nGET /get HTTP/1.1\r\nHost: httpbin.org\r\n\r\n' \ +// | ncat --no-shutdown localhost 6191 \ +// | grep -oE 'HTTP/1.1 [0-9]{3}' + +fn main() { + env_logger::init(); + + let opt = Opt::parse_args(); + let mut my_server = Server::new(Some(opt)).unwrap(); + my_server.bootstrap(); + + let mut my_proxy = + pingora_proxy::http_proxy_service(&my_server.configuration, PipelinedGateway); + my_proxy.add_tcp("0.0.0.0:6191"); + my_server.add_service(my_proxy); + + my_server.run_forever(); +} diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs index f89f53d33..2773cfe03 100644 --- a/pingora-proxy/src/lib.rs +++ b/pingora-proxy/src/lib.rs @@ -46,7 +46,7 @@ use pingora_http::{RequestHeader, ResponseHeader}; use std::fmt::Debug; use std::str; use std::sync::{ - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, Arc, }; use std::time::Duration; @@ -72,7 +72,7 @@ use pingora_core::protocols::http::SERVER_NAME; use pingora_core::protocols::Stream; use pingora_core::protocols::{Digest, UniqueID}; use pingora_core::server::configuration::ServerConf; -use pingora_core::server::ShutdownWatch; +use pingora_core::server::{RuntimeOpts, ShutdownWatch}; use pingora_core::upstreams::peer::{HttpPeer, Peer}; use pingora_error::{Error, ErrorSource, ErrorType::*, OrErr, Result}; @@ -91,10 +91,10 @@ use subrequest::{BodyMode, Ctx as SubrequestCtx}; pub use proxy_cache::range_filter::{range_header_filter, MultiRangeInfo, RangeType}; pub use proxy_purge::PurgeStatus; -pub use proxy_trait::{FailToProxy, ProxyHttp}; +pub use proxy_trait::{FailToProxy, ProxyHttp, ProxyWarnLogContext}; pub mod prelude { - pub use crate::{http_proxy, http_proxy_service, ProxyHttp, Session}; + pub use crate::{http_proxy, http_proxy_service, ProxyHttp, ProxyWarnLogContext, Session}; } pub type ProcessCustomSession = Arc< @@ -119,6 +119,8 @@ where pub server_options: Option, pub h2_options: Option, pub downstream_modules: HttpModules, + #[cfg(feature = "upstream_modules")] + pub upstream_modules: HttpModules, max_retries: usize, process_custom_session: Option>, } @@ -153,6 +155,8 @@ impl HttpProxy { server_options: None, h2_options: None, downstream_modules: HttpModules::new(), + #[cfg(feature = "upstream_modules")] + upstream_modules: HttpModules::new(), max_retries: conf.max_retries, process_custom_session: None, } @@ -169,13 +173,15 @@ where connector: C, on_custom: Option>, server_options: Option, + client_options: Option, ) -> Self where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync, { - let client_upstream = - Connector::new_custom(Some(ConnectorOptions::from_server_conf(&conf)), connector); + let client_options = + client_options.unwrap_or_else(|| ConnectorOptions::from_server_conf(&conf)); + let client_upstream = Connector::new_custom(Some(client_options), connector); HttpProxy { inner, @@ -184,12 +190,25 @@ where shutdown_flag: Arc::new(AtomicBool::new(false)), server_options, downstream_modules: HttpModules::new(), + #[cfg(feature = "upstream_modules")] + upstream_modules: HttpModules::new(), max_retries: conf.max_retries, process_custom_session: on_custom, h2_options: None, } } + /// Return the number of times a pooled upstream connection was found to contain + /// unexpected data from the server. + pub fn unexpected_data_connection_count(&self) -> u64 { + self.client_upstream.unexpected_data_connection_count() + } + + /// Return a shared reference to the unexpected data connection counter for periodic metric reporting. + pub fn unexpected_data_connection_counter(&self) -> Arc { + self.client_upstream.unexpected_data_connection_counter() + } + /// Initialize the downstream modules for this proxy. /// /// This method must be called after creating an [`HttpProxy`] with [`HttpProxy::new()`] @@ -204,6 +223,8 @@ where { self.inner .init_downstream_modules(&mut self.downstream_modules); + #[cfg(feature = "upstream_modules")] + self.inner.init_upstream_modules(&mut self.upstream_modules); } async fn handle_new_request( @@ -417,14 +438,17 @@ where if reuse { // TODO: log error - let persistent_settings = HttpPersistentSettings::for_session(&session); + let mut persistent_settings = HttpPersistentSettings::for_session(&session); + if let Some(uc) = self.inner.persist_connection_context(&session, ctx) { + persistent_settings.set_user_context(uc); + } session .downstream_session .finish() .await .ok() .flatten() - .map(|s| ReusedHttpStream::new(s, Some(persistent_settings))) + .map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)) } else { None } @@ -461,6 +485,10 @@ pub struct Session { pub subrequest_spawner: Option, // Downstream filter modules pub downstream_modules_ctx: HttpModuleCtx, + /// Upstream filter modules. These run before `upstream_compression` and see the raw + /// (pre-compression) upstream response body. + #[cfg(feature = "upstream_modules")] + pub upstream_modules_ctx: HttpModuleCtx, /// Upstream response body bytes received (payload only). Set by proxy layer. /// TODO: move this into an upstream session digest for future fields. upstream_body_bytes_received: usize, @@ -474,6 +502,7 @@ impl Session { fn new( downstream_session: impl Into>, downstream_modules: &HttpModules, + #[cfg(feature = "upstream_modules")] upstream_modules: &HttpModules, shutdown_flag: Arc, ) -> Self { Session { @@ -486,6 +515,8 @@ impl Session { subrequest_ctx: None, subrequest_spawner: None, // optionally set later on downstream_modules_ctx: downstream_modules.build_ctx(), + #[cfg(feature = "upstream_modules")] + upstream_modules_ctx: upstream_modules.build_ctx(), upstream_body_bytes_received: 0, upstream_write_pending_time: Duration::ZERO, shutdown_flag, @@ -501,6 +532,8 @@ impl Session { Self::new( Box::new(HttpSession::new_http1(stream)), &modules, + #[cfg(feature = "upstream_modules")] + &HttpModules::new(), Arc::new(AtomicBool::new(false)), ) } @@ -513,10 +546,47 @@ impl Session { Self::new( Box::new(HttpSession::new_http1(stream)), downstream_modules, + #[cfg(feature = "upstream_modules")] + &HttpModules::new(), Arc::new(AtomicBool::new(false)), ) } + /// Run upstream module filters on the given [`HttpTask`]. + /// + /// Upstream modules process each task **before** `upstream_compression` and + /// see the raw (pre-compression) upstream response. Like the downstream + /// module path, `response_trailer_filter` and `response_done_filter` return + /// values are converted to body tasks when present. + #[cfg(feature = "upstream_modules")] + pub async fn upstream_modules_filter_task(&mut self, t: &mut HttpTask) -> Result<()> { + match t { + HttpTask::Header(header, eos) => { + self.upstream_modules_ctx + .response_header_filter(header, *eos) + .await?; + } + HttpTask::Body(body, eos) | HttpTask::UpgradedBody(body, eos) => { + self.upstream_modules_ctx.response_body_filter(body, *eos)?; + } + HttpTask::Trailer(trailers) => { + if let Some(buf) = self + .upstream_modules_ctx + .response_trailer_filter(trailers)? + { + *t = HttpTask::Body(Some(buf), true); + } + } + HttpTask::Done => { + if let Some(buf) = self.upstream_modules_ctx.response_done_filter()? { + *t = HttpTask::Body(Some(buf), true); + } + } + HttpTask::Failed(_) => {} + } + Ok(()) + } + pub fn as_downstream_mut(&mut self) -> &mut HttpSession { &mut self.downstream_session } @@ -584,57 +654,115 @@ impl Session { .await } - pub async fn write_response_tasks(&mut self, mut tasks: Vec) -> Result { - let mut seen_upgraded = self.was_upgraded(); - for task in tasks.iter_mut() { - match task { - HttpTask::Header(resp, end) => { - self.downstream_modules_ctx - .response_header_filter(resp, *end) - .await?; - } - HttpTask::Body(data, end) => { - self.downstream_modules_ctx - .response_body_filter(data, *end)?; - } - HttpTask::UpgradedBody(data, end) => { - seen_upgraded = true; - self.downstream_modules_ctx - .response_body_filter(data, *end)?; - } - HttpTask::Trailer(trailers) => { - if let Some(buf) = self - .downstream_modules_ctx - .response_trailer_filter(trailers)? - { - // Write the trailers into the body if the filter - // returns a buffer. - // - // Note, this will not work if end of stream has already - // been seen or we've written content-length bytes. - // (Trailers should never come after upgraded body) - *task = HttpTask::Body(Some(buf), true); - } - } - HttpTask::Done => { - // `Done` can be sent in certain response paths to mark end - // of response if not already done via trailers or body with - // end flag set. - // If the filter returns body bytes on Done, - // write them into the response. + // Run downstream module response filters on a single task, updating + // `seen_upgraded` to track whether an upgrade has been seen. Used by both + // `send_downstream_proxy_task` and `write_response_tasks`. + async fn downstream_response_task_filter( + &mut self, + task: &mut HttpTask, + seen_upgraded: &mut bool, + ) -> Result<()> { + match task { + HttpTask::Header(resp, end) => { + self.downstream_modules_ctx + .response_header_filter(resp, *end) + .await?; + } + HttpTask::Body(data, end) => { + self.downstream_modules_ctx + .response_body_filter(data, *end)?; + } + HttpTask::UpgradedBody(data, end) => { + *seen_upgraded = true; + self.downstream_modules_ctx + .response_body_filter(data, *end)?; + } + HttpTask::Trailer(trailers) => { + if let Some(buf) = self + .downstream_modules_ctx + .response_trailer_filter(trailers)? + { + // Write the trailers into the body if the filter + // returns a buffer. // // Note, this will not work if end of stream has already // been seen or we've written content-length bytes. - if let Some(buf) = self.downstream_modules_ctx.response_done_filter()? { - if seen_upgraded { - *task = HttpTask::UpgradedBody(Some(buf), true); - } else { - *task = HttpTask::Body(Some(buf), true); - } + // (Trailers should never come after upgraded body) + *task = HttpTask::Body(Some(buf), true); + } + } + HttpTask::Done => { + // `Done` can be sent in certain response paths to mark end + // of response if not already done via trailers or body with + // end flag set. + // If the filter returns body bytes on Done, + // write them into the response. + // + // Note, this will not work if end of stream has already + // been seen or we've written content-length bytes. + if let Some(buf) = self.downstream_modules_ctx.response_done_filter()? { + if *seen_upgraded { + *task = HttpTask::UpgradedBody(Some(buf), true); + } else { + *task = HttpTask::Body(Some(buf), true); } } - _ => { /* Failed */ } } + _ => { /* Failed */ } + } + Ok(()) + } + + /// Queue a downstream proxy task for cancel-safe writing after running + /// downstream module filters. This allows decoupling cache writes from + /// downstream writes. + /// + /// Only works with sessions that support the proxy task API (currently H1). + /// + /// # Panics + /// Panics if the session doesn't support the proxy task API. + /// Use `write_response_tasks()` for sessions that don't support the proxy task API. + pub async fn send_downstream_proxy_task(&mut self, mut task: HttpTask) -> Result<()> { + let mut seen_upgraded = self.was_upgraded(); + self.downstream_response_task_filter(&mut task, &mut seen_upgraded) + .await?; + self.downstream_session.send_downstream_proxy_task(task); + Ok(()) + } + + /// Enable or disable the cancel-safe proxy task API for this session. + /// + /// When disabled, the proxy falls back to the blocking `write_response_tasks` + /// path. This can be called from request filters to opt out on a per-request + /// basis. + pub fn set_proxy_tasks_enabled(&mut self, enabled: bool) { + self.downstream_session.set_proxy_tasks_enabled(enabled); + } + + /// Check if there are pending downstream tasks queued for writing. + /// Used for backpressure - don't queue more cache tasks if we have pending writes. + /// Returns false for sessions that don't support the proxy task API. + pub fn has_pending_downstream_tasks(&self) -> bool { + self.downstream_session.supports_proxy_task_api() + && self.downstream_session.has_pending_downstream_proxy_tasks() + } + + /// Write all queued downstream proxy tasks. This is cancel-safe and can be called + /// in a select! loop while waiting for upstream tasks. + /// For sessions that don't support the proxy task API, this is a no-op. + pub async fn write_downstream_proxy_tasks(&mut self) -> Result { + if self.downstream_session.supports_proxy_task_api() { + self.downstream_session.write_downstream_proxy_tasks().await + } else { + Ok(false) + } + } + + pub async fn write_response_tasks(&mut self, mut tasks: Vec) -> Result { + let mut seen_upgraded = self.was_upgraded(); + for task in tasks.iter_mut() { + self.downstream_response_task_filter(task, &mut seen_upgraded) + .await?; } self.downstream_session.response_duplex_vec(tasks).await } @@ -785,14 +913,17 @@ where // TODO: log error self.inner.logging(&mut session, None, &mut ctx).await; self.cleanup_sub_req(&mut session); - let persistent_settings = HttpPersistentSettings::for_session(&session); + let mut persistent_settings = HttpPersistentSettings::for_session(&session); + if let Some(uc) = self.inner.persist_connection_context(&session, &ctx) { + persistent_settings.set_user_context(uc); + } return session .downstream_session .finish() .await .ok() .flatten() - .map(|s| ReusedHttpStream::new(s, Some(persistent_settings))); + .map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)); } /* else continue */ } @@ -876,18 +1007,27 @@ where match e { Some(error) => { let retry = error.retry(); + // only log error that will be retried here, the final error will be logged below + if retry + && !self.inner.suppress_proxy_warn_log( + &session, + &ctx, + &error, + ProxyWarnLogContext::UpstreamRetry, + ) + { + warn!( + "Fail to proxy: {}, tries: {}, retry: {}, {}", + error, + retries, + retry, + self.inner.request_summary(&session, &ctx) + ); + } proxy_error = Some(error); if !retry { break; } - // only log error that will be retried here, the final error will be logged below - warn!( - "Fail to proxy: {}, tries: {}, retry: {}, {}", - proxy_error.as_ref().unwrap(), - retries, - retry, - self.inner.request_summary(&session, &ctx) - ); } None => { proxy_error = None; @@ -931,7 +1071,7 @@ where if !self.inner.suppress_error_log(&session, &ctx, e) { error!( "Fail to proxy: {}, status: {}, tries: {}, retry: {}, {}", - final_error.as_ref().unwrap(), + e, res.error_code, retries, false, // we never retry here @@ -971,14 +1111,17 @@ where session.downstream_session.on_proxy_failure(e); if res.can_reuse_downstream { - let persistent_settings = HttpPersistentSettings::for_session(&session); + let mut persistent_settings = HttpPersistentSettings::for_session(&session); + if let Some(uc) = self.inner.persist_connection_context(&session, ctx) { + persistent_settings.set_user_context(uc); + } session .downstream_session .finish() .await .ok() .flatten() - .map(|s| ReusedHttpStream::new(s, Some(persistent_settings))) + .map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)) } else { None } @@ -1021,6 +1164,8 @@ where Some(downstream_session) => Session::new( downstream_session, &self.downstream_modules, + #[cfg(feature = "upstream_modules")] + &self.upstream_modules, self.shutdown_flag.clone(), ), None => return, // bad request @@ -1135,9 +1280,12 @@ where { async fn process_new_http( self: &Arc, - session: HttpSession, + mut session: HttpSession, shutdown: &ShutdownWatch, ) -> Option { + // Extract user context from the previous request before the session is moved into the Box + let prev_user_ctx = session.take_connection_user_context(); + let session = Box::new(session); // TODO: keepalive pool, use stack @@ -1145,6 +1293,8 @@ where Some(downstream_session) => Session::new( downstream_session, &self.downstream_modules, + #[cfg(feature = "upstream_modules")] + &self.upstream_modules, self.shutdown_flag.clone(), ), None => return None, // bad request @@ -1155,7 +1305,14 @@ where session.set_keepalive(None); } - let ctx = self.inner.new_ctx(); + let mut ctx = self.inner.new_ctx(); + + // Deliver user context from the previous request on this reused connection + if let Some(prev_ctx) = prev_user_ctx { + self.inner + .on_connection_reuse(&mut session, &mut ctx, prev_ctx); + } + self.process_request(session, ctx).await } @@ -1190,7 +1347,7 @@ where // TODO implement h2_options } -use pingora_core::services::listening::Service; +use pingora_core::services::listening::{RuntimeOptsOverride, Service}; /// Create an [`HttpProxy`] without wrapping it in a [`Service`]. /// @@ -1272,7 +1429,8 @@ where SV::CTX: Send + Sync + 'static, C: custom::Connector, { - let mut proxy = HttpProxy::new_custom(inner, conf.clone(), connector, Some(on_custom), None); + let mut proxy = + HttpProxy::new_custom(inner, conf.clone(), connector, Some(on_custom), None, None); proxy.handle_init_modules(); Service::new(name.to_string(), proxy) @@ -1295,6 +1453,8 @@ where connector: C, custom: Option>, server_options: Option, + client_options: Option, + runtime_opts_override: Option, } impl ProxyServiceBuilder @@ -1319,6 +1479,8 @@ where connector: (), custom: None, server_options: None, + client_options: None, + runtime_opts_override: None, } } } @@ -1353,6 +1515,8 @@ where inner, name, server_options, + client_options, + runtime_opts_override, .. } = self; ProxyServiceBuilder { @@ -1362,9 +1526,19 @@ where connector, custom: Some(on_custom), server_options, + client_options, + runtime_opts_override, } } + /// Set the upstream client connector options for the [ProxyServiceBuilder]. + /// + /// Returns a new [ProxyServiceBuilder] with the upstream client connector options set. + pub fn client_options(mut self, options: ConnectorOptions) -> Self { + self.client_options = Some(options); + self + } + /// Set the server options for the [ProxyServiceBuilder]. /// /// Returns a new [ProxyServiceBuilder] with the server options set. @@ -1373,6 +1547,17 @@ where self } + /// Set a runtime options override for the [Service] built by this builder. + /// + /// Returning [`None`] from the override uses the global runtime options. + pub fn runtime_opts_override(mut self, override_fn: F) -> Self + where + F: Fn(&RuntimeOpts) -> Option + Send + Sync + 'static, + { + self.runtime_opts_override = Some(Arc::new(override_fn)); + self + } + /// Builds a new [Service] from the [ProxyServiceBuilder]. /// /// This function takes ownership of the [ProxyServiceBuilder] and returns a new [Service] with @@ -1387,11 +1572,24 @@ where connector, custom, server_options, + client_options, + runtime_opts_override, } = self; - let mut proxy = HttpProxy::new_custom(inner, conf, connector, custom, server_options); + let mut proxy = HttpProxy::new_custom( + inner, + conf, + connector, + custom, + server_options, + client_options, + ); proxy.handle_init_modules(); - Service::new(name, proxy) + let mut service = Service::new(name, proxy); + if let Some(runtime_opts_override) = runtime_opts_override { + service.set_runtime_opts_override(runtime_opts_override); + } + service } } diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs index 43b2ace95..fecf6fbd2 100644 --- a/pingora-proxy/src/proxy_cache.rs +++ b/pingora-proxy/src/proxy_cache.rs @@ -487,6 +487,7 @@ where } } + // No enabled() guard: no concurrent upstream can disable cache here. if let Err(e) = session.cache.finish_hit_handler().await { warn!("Error during finish_hit_handler: {}", e); } @@ -1276,7 +1277,7 @@ pub mod range_filter { pub ranges: Vec>, pub boundary: String, total_length: usize, - content_type: Option, + pub content_type: Option, } impl MultiRangeInfo { diff --git a/pingora-proxy/src/proxy_common.rs b/pingora-proxy/src/proxy_common.rs index e1d36f699..6c40760c2 100644 --- a/pingora-proxy/src/proxy_common.rs +++ b/pingora-proxy/src/proxy_common.rs @@ -1,3 +1,17 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + /// Possible downstream states during request multiplexing #[derive(Debug, Clone, Copy)] pub(crate) enum DownstreamStateMachine { @@ -36,19 +50,28 @@ impl DownstreamStateMachine { matches!(self, Self::Errored) } - /// Move the state machine to Finished state if `set` is true + /// Move the state machine to Finished state if `set` is true. + /// + /// No-op when the current state is [`Errored`](Self::Errored) — once errored the + /// downstream connection must not be reused, and late upstream chunks arriving + /// via `rx.recv()` must not overwrite that decision. pub fn maybe_finished(&mut self, set: bool) { - if set { + if set && !self.is_errored() { *self = Self::ReadingFinished } } - /// Reset if we should continue reading from the downstream again. - /// Only used with upgraded connections when body mode changes. + /// Reset to [`Reading`](Self::Reading) for upgraded connections when body mode changes. + /// + /// No-op when the current state is [`Errored`](Self::Errored). pub fn reset(&mut self) { - *self = Self::Reading; + if !self.is_errored() { + *self = Self::Reading; + } } + /// Transition to [`Errored`](Self::Errored). This is a terminal state: once entered, + /// no other state transition is permitted and the connection must not be reused. pub fn to_errored(&mut self) { *self = Self::Errored } @@ -97,3 +120,61 @@ impl ResponseStateMachine { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normal_lifecycle() { + let mut ds = DownstreamStateMachine::new(false); + assert!(ds.is_reading()); + assert!(ds.can_poll()); + assert!(!ds.is_errored()); + + ds.maybe_finished(true); + assert!(!ds.is_reading()); + assert!(ds.is_done()); + assert!(ds.can_poll()); // ReadingFinished still allows polling (for idle) + assert!(!ds.is_errored()); + } + + #[test] + fn errored_is_terminal() { + let mut ds = DownstreamStateMachine::new(false); + ds.to_errored(); + assert!(ds.is_errored()); + assert!(!ds.can_poll()); + assert!(ds.is_done()); + } + + /// `maybe_finished(false)` is always a no-op regardless of state. + #[test] + fn maybe_finished_false_is_noop() { + let mut ds = DownstreamStateMachine::new(false); + ds.to_errored(); + ds.maybe_finished(false); // must not panic + assert!(ds.is_errored()); + assert!(!ds.can_poll()); + } + + /// `maybe_finished(true)` on `Errored` is a no-op — `Errored` is terminal. + #[test] + fn maybe_finished_true_noop_on_errored() { + let mut ds = DownstreamStateMachine::new(false); + ds.to_errored(); + ds.maybe_finished(true); // must not overwrite Errored + assert!(ds.is_errored()); + assert!(!ds.can_poll()); + } + + /// `reset()` on `Errored` is a no-op — `Errored` is terminal. + #[test] + fn reset_noop_on_errored() { + let mut ds = DownstreamStateMachine::new(false); + ds.to_errored(); + ds.reset(); // must not overwrite Errored + assert!(ds.is_errored()); + assert!(!ds.can_poll()); + } +} diff --git a/pingora-proxy/src/proxy_custom.rs b/pingora-proxy/src/proxy_custom.rs index 630791115..fc0b690c1 100644 --- a/pingora-proxy/src/proxy_custom.rs +++ b/pingora-proxy/src/proxy_custom.rs @@ -257,7 +257,90 @@ where } } - // returns whether server (downstream) session can be reused + #[allow(clippy::too_many_arguments)] + async fn process_upstream_tasks_custom( + &self, + session: &mut Session, + ctx: &mut SV::CTX, + initial_task: HttpTask, + rx: &mut mpsc::Receiver, + serve_from_cache: &mut ServeFromCache, + range_body_filter: &mut proxy_cache::range_filter::RangeBodyFilter, + response_state: &mut ResponseStateMachine, + ) -> Result> + where + SV: ProxyHttp + Send + Sync, + SV::CTX: Send + Sync, + { + if serve_from_cache.should_discard_upstream() { + // just drain, do we need to do anything else? + return Ok(None); + } + + // Batch: pull as many tasks as we can from rx + let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); + tasks.push(initial_task); + while let Ok(task) = rx.try_recv() { + tasks.push(task); + } + + /* run filters before sending to downstream */ + let mut filtered_tasks = Vec::with_capacity(TASK_BUFFER_SIZE); + for mut t in tasks { + if self.revalidate_or_stale(session, &mut t, ctx).await { + serve_from_cache.enable(); + response_state.enable_cached_response(); + // skip downstream filtering entirely as the 304 will not be sent + break; + } + #[cfg(feature = "upstream_modules")] + if let HttpTask::Header(header, end_of_stream) = &t { + self.inner + .adjust_upstream_modules(session, header, *end_of_stream, ctx) + .await?; + } + #[cfg(feature = "upstream_modules")] + session.upstream_modules_filter_task(&mut t).await?; + session.upstream_compression.response_filter(&mut t); + // check error and abort + // otherwise the error is surfaced via write_response_tasks() + if !serve_from_cache.should_send_to_downstream() { + if let HttpTask::Failed(e) = t { + return Err(e); + } + } + filtered_tasks.push( + self.custom_response_filter( + session, + t, + ctx, + serve_from_cache, + range_body_filter, + false, + ) + .await?, + ); + if serve_from_cache.is_miss_header() { + response_state.enable_cached_response(); + } + } + + if !serve_from_cache.should_send_to_downstream() { + // TODO: need to derive response_done from filtered_tasks in case downstream failed already + return Ok(None); + } + + let response_done = session.write_response_tasks(filtered_tasks).await?; + + Ok(Some(response_done)) + } + + // TODO: pre-existing inconsistency with proxy_h1/proxy_h2 to address in a follow-up: + // upstream task rx.recv() branch is missing + // downstream_state.maybe_finished(session.is_body_done()) after processing. proxy_h1 has + // this because upgrade responses can force the body done — since custom upstreams can + // serve H1 downstreams that support upgrades, the same may be needed here. + // Returns whether server (downstream) session can be reused #[allow(clippy::too_many_arguments)] async fn custom_bidirection_down_to_up( &self, @@ -303,6 +386,8 @@ where let mut serve_from_cache = ServeFromCache::new(); let mut range_body_filter = proxy_cache::range_filter::RangeBodyFilter::new(); + let mut next_upstream_task: Option = None; + let mut upstream_custom = true; let mut downstream_custom = true; @@ -329,11 +414,18 @@ where if wait_for_cache_fill { // ignore downstream error so that upstream can continue to write cache downstream_state.to_errored(); - warn!( - "Downstream Error ignored during caching: {}, {}", - e, - self.inner.request_summary(session, ctx) - ); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } continue; } else { return Err(e.into_down()); @@ -361,93 +453,159 @@ where }; }, - task = rx.recv(), if !response_state.upstream_done() => { - debug!("upstream event"); - + // Handle buffered upstream task from previous iteration + task = async { next_upstream_task.take() }, if next_upstream_task.is_some() => { + debug!("buffered upstream event: {:?}", task); if let Some(t) = task { - debug!("upstream event custom: {:?}", t); - if serve_from_cache.should_discard_upstream() { - // just drain, do we need to do anything else? - continue; - } - // pull as many tasks as we can - let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - tasks.push(t); - while let Ok(task) = rx.try_recv() { - tasks.push(task); - } - - /* run filters before sending to downstream */ - let mut filtered_tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - for mut t in tasks { - if self.revalidate_or_stale(session, &mut t, ctx).await { - serve_from_cache.enable(); - response_state.enable_cached_response(); - // skip downstream filtering entirely as the 304 will not be sent - break; - } - session.upstream_compression.response_filter(&mut t); - // check error and abort - // otherwise the error is surfaced via write_response_tasks() - if !serve_from_cache.should_send_to_downstream() { - if let HttpTask::Failed(e) = t { - return Err(e); - } - } - filtered_tasks.push( - self.custom_response_filter(session, t, ctx, - &mut serve_from_cache, - &mut range_body_filter, false).await?); - if serve_from_cache.is_miss_header() { - response_state.enable_cached_response(); - } - } - - if !serve_from_cache.should_send_to_downstream() { - // TODO: need to derive response_done from filtered_tasks in case downstream failed already + let Some(response_done) = self.process_upstream_tasks_custom( + session, + ctx, + t, + &mut rx, + &mut serve_from_cache, + &mut range_body_filter, + &mut response_state, + ).await? else { + // nothing sent downstream e.g. serve_from_cache continue; - } + }; + response_state.maybe_set_upstream_done(response_done); + } else { + debug!("empty upstream event"); + response_state.maybe_set_upstream_done(true); + } + }, + task = rx.recv(), if !response_state.upstream_done() && next_upstream_task.is_none() => { + debug!("upstream event: {:?}", task); + if let Some(t) = task { let upgraded = session.was_upgraded(); - let response_done = session.write_response_tasks(filtered_tasks).await?; + let Some(response_done) = self.process_upstream_tasks_custom( + session, + ctx, + t, + &mut rx, + &mut serve_from_cache, + &mut range_body_filter, + &mut response_state, + ).await? else { + // nothing sent downstream e.g. serve_from_cache + continue; + }; if !upgraded && session.was_upgraded() && downstream_state.can_poll() { // just upgraded, the downstream state should be reset to continue to // poll body trace!("reset downstream state on upgrade"); downstream_state.reset(); } - response_state.maybe_set_upstream_done(response_done); } else { debug!("empty upstream event"); response_state.maybe_set_upstream_done(true); } - } + }, task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded), - if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => { + if !response_state.cached_done() + && !downstream_state.is_errored() + && serve_from_cache.is_on() + && !session.has_pending_downstream_tasks() => { // backpressure: don't queue if pending writes + let task = self.custom_response_filter(session, task?, ctx, &mut serve_from_cache, &mut range_body_filter, true).await?; - match session.write_response_tasks(vec![task]).await { - Ok(b) => response_state.maybe_set_cache_done(b), - Err(e) => if serve_from_cache.is_miss() { - // give up writing to downstream but wait for upstream cache write to finish - downstream_state.to_errored(); - response_state.maybe_set_cache_done(true); - warn!( - "Downstream Error ignored during caching: {}, {}", - e, - self.inner.request_summary(session, ctx) - ); - continue; - } else { - return Err(e); + + if session.downstream_session.supports_proxy_task_api() { + session.send_downstream_proxy_task(task).await?; + } else { + match session.write_response_tasks(vec![task]).await { + Ok(b) => response_state.maybe_set_cache_done(b), + Err(e) => if serve_from_cache.is_miss() { + // give up writing to downstream but wait for upstream cache write to finish + downstream_state.to_errored(); + response_state.maybe_set_cache_done(true); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + session.downstream_session.on_proxy_failure(e); + continue; + } else { + return Err(e); + } + } + // A storage error can disable cache between cached_done + // being set and here; see the same guard in proxy_h1.rs. + if response_state.cached_done() && session.cache.enabled() { + if let Err(e) = session.cache.finish_hit_handler().await { + warn!("Error during finish_hit_handler: {}", e); + } } } - if response_state.cached_done() { - if let Err(e) = session.cache.finish_hit_handler().await { - warn!("Error during finish_hit_handler: {}", e); + } + + // Write queued downstream proxy tasks while also polling for upstream tasks. + // This allows cache writes to continue even when downstream is stalled. + // + // "Gate" branch: ready(()) resolves immediately, so the guard controls + // whether we enter. This is not a busy-loop because every path through + // the inner select either (a) drains all pending tasks via + // write_downstream_proxy_tasks (making the guard false), (b) stores an + // upstream task in next_upstream_task (making the guard false), or + // (c) blocks on real I/O inside the nested select. + _ = std::future::ready(()), if session.has_pending_downstream_tasks() && next_upstream_task.is_none() => { + tokio::select! { + // Try to write downstream proxy tasks (cancel-safe) + write_result = session.write_downstream_proxy_tasks() => { + match write_result { + Ok(end) => { + response_state.maybe_set_cache_done(end); + // See enabled() guard comment above. + if response_state.cached_done() && session.cache.enabled() { + if let Err(e) = session.cache.finish_hit_handler().await { + warn!("Error during finish_hit_handler: {}", e); + } + } + } + Err(e) => if serve_from_cache.is_miss() { + // give up writing to downstream but wait for upstream cache write to finish + downstream_state.to_errored(); + response_state.maybe_set_cache_done(true); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream write error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + session.downstream_session.on_proxy_failure(e); + } else { + return Err(e); + } + } + } + + // Also poll for upstream tasks - if we get one, cancel the write and handle it. + upstream_task = rx.recv(), if !response_state.upstream_done() && serve_from_cache.is_on() && next_upstream_task.is_none() => { + if let Some(t) = upstream_task { + next_upstream_task = Some(t); + continue; + } else { + response_state.maybe_set_upstream_done(true); + } } } } diff --git a/pingora-proxy/src/proxy_h1.rs b/pingora-proxy/src/proxy_h1.rs index 9f04289c2..ec3d293f0 100644 --- a/pingora-proxy/src/proxy_h1.rs +++ b/pingora-proxy/src/proxy_h1.rs @@ -267,6 +267,83 @@ where Ok(()) } + #[allow(clippy::too_many_arguments)] + async fn process_upstream_tasks( + &self, + session: &mut Session, + ctx: &mut SV::CTX, + initial_task: HttpTask, + rx: &mut mpsc::Receiver, + serve_from_cache: &mut ServeFromCache, + range_body_filter: &mut proxy_cache::range_filter::RangeBodyFilter, + response_state: &mut ResponseStateMachine, + ) -> Result> + where + SV: ProxyHttp + Send + Sync, + SV::CTX: Send + Sync, + { + if serve_from_cache.should_discard_upstream() { + // just drain, do we need to do anything else? + return Ok(None); + } + + // Batch: pull as many tasks as we can from rx + let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); + tasks.push(initial_task); + // tokio::task::unconstrained because now_or_never may yield None when the future is ready + while let Some(maybe_task) = tokio::task::unconstrained(rx.recv()).now_or_never() { + debug!("upstream event now: {:?}", maybe_task); + if let Some(t) = maybe_task { + tasks.push(t); + } else { + break; // upstream closed + } + } + + /* run filters before sending to downstream */ + let mut filtered_tasks = Vec::with_capacity(TASK_BUFFER_SIZE); + for mut t in tasks { + if self.revalidate_or_stale(session, &mut t, ctx).await { + serve_from_cache.enable(); + response_state.enable_cached_response(); + // skip downstream filtering entirely as the 304 will not be sent + break; + } + #[cfg(feature = "upstream_modules")] + if let HttpTask::Header(header, end_of_stream) = &t { + self.inner + .adjust_upstream_modules(session, header, *end_of_stream, ctx) + .await?; + } + #[cfg(feature = "upstream_modules")] + session.upstream_modules_filter_task(&mut t).await?; + session.upstream_compression.response_filter(&mut t); + let task = self + .h1_response_filter(session, t, ctx, serve_from_cache, range_body_filter, false) + .await?; + if serve_from_cache.is_miss_header() { + response_state.enable_cached_response(); + } + // check error and abort + // otherwise the error is surfaced via write_response_tasks() + if !serve_from_cache.should_send_to_downstream() { + if let HttpTask::Failed(e) = task { + return Err(e); + } + } + filtered_tasks.push(task); + } + + if !serve_from_cache.should_send_to_downstream() { + // TODO: need to derive response_done from filtered_tasks in case downstream failed already + return Ok(None); + } + + let response_done = session.write_response_tasks(filtered_tasks).await?; + + Ok(Some(response_done)) + } + // todo use this function to replace bidirection_1to2() // returns whether this server (downstream) session can be reused async fn proxy_handle_downstream( @@ -329,6 +406,8 @@ where let mut serve_from_cache = proxy_cache::ServeFromCache::new(); let mut range_body_filter = proxy_cache::range_filter::RangeBodyFilter::new(); + let mut next_upstream_task: Option = None; + /* duplex mode without caching * Read body from downstream while reading response from upstream * If response is done, only read body from downstream @@ -388,11 +467,18 @@ where if wait_for_cache_fill { // ignore downstream error so that upstream can continue to write cache downstream_state.to_errored(); - warn!( - "Downstream Error ignored during caching: {}, {}", - e, - self.inner.request_summary(session, ctx) - ); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } // This will not be treated as a final error, but we should signal to // downstream session regardless session.downstream_session.on_proxy_failure(e); @@ -424,68 +510,56 @@ where // If tx is closed, the upstream has already finished its job. downstream_state.maybe_finished(tx.is_closed()); debug!("waiting for permit {send_permit:?}, upstream closed {}", tx.is_closed()); - /* No permit, wait on more capacity to avoid starving. + /* No permit, wait on more capacity to avoid starving. * Otherwise this select only blocks on rx, which might send no data * before the entire body is uploaded. * once more capacity arrives we just loop back */ }, - task = rx.recv(), if !response_state.upstream_done() => { - debug!("upstream event: {:?}", task); + // Handle buffered upstream task from previous iteration + task = async { next_upstream_task.take() }, if next_upstream_task.is_some() => { + debug!("buffered upstream event: {:?}", task); if let Some(t) = task { - if serve_from_cache.should_discard_upstream() { - // just drain, do we need to do anything else? - continue; - } - // pull as many tasks as we can - let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - tasks.push(t); - // tokio::task::unconstrained because now_or_never may yield None when the future is ready - while let Some(maybe_task) = tokio::task::unconstrained(rx.recv()).now_or_never() { - debug!("upstream event now: {:?}", maybe_task); - if let Some(t) = maybe_task { - tasks.push(t); - } else { - break; // upstream closed - } - } - - /* run filters before sending to downstream */ - let mut filtered_tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - for mut t in tasks { - if self.revalidate_or_stale(session, &mut t, ctx).await { - serve_from_cache.enable(); - response_state.enable_cached_response(); - // skip downstream filtering entirely as the 304 will not be sent - break; - } - session.upstream_compression.response_filter(&mut t); - let task = self.h1_response_filter(session, t, ctx, - &mut serve_from_cache, - &mut range_body_filter, false).await?; - if serve_from_cache.is_miss_header() { - response_state.enable_cached_response(); - } - // check error and abort - // otherwise the error is surfaced via write_response_tasks() - if !serve_from_cache.should_send_to_downstream() { - if let HttpTask::Failed(e) = task { - return Err(e); - } - } - filtered_tasks.push(task); - } - - if !serve_from_cache.should_send_to_downstream() { - // TODO: need to derive response_done from filtered_tasks in case downstream failed already + let Some(response_done) = self.process_upstream_tasks( + session, + ctx, + t, + &mut rx, + &mut serve_from_cache, + &mut range_body_filter, + &mut response_state, + ).await? else { + // nothing sent downstream e.g. serve_from_cache continue; - } + }; + response_state.maybe_set_upstream_done(response_done); + // unsuccessful upgrade response may force the request done + downstream_state.maybe_finished(session.is_body_done()); + } else { + debug!("empty upstream event"); + response_state.maybe_set_upstream_done(true); + } + }, - // set to downstream + task = rx.recv(), if !response_state.upstream_done() && next_upstream_task.is_none() => { + debug!("upstream event: {:?}", task); + if let Some(t) = task { let upgraded = session.was_upgraded(); - let response_done = session.write_response_tasks(filtered_tasks).await?; + let Some(response_done) = self.process_upstream_tasks( + session, + ctx, + t, + &mut rx, + &mut serve_from_cache, + &mut range_body_filter, + &mut response_state, + ).await? else { + // nothing sent downstream e.g. serve_from_cache + continue; + }; if !upgraded && session.was_upgraded() && downstream_state.can_poll() { + // TODO: write can happen async now // just upgraded, the downstream state should be reset to continue to // poll body trace!("reset downstream state on upgrade"); @@ -502,35 +576,114 @@ where }, task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded), - if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => { + if !response_state.cached_done() + && !downstream_state.is_errored() + && serve_from_cache.is_on() + && !session.has_pending_downstream_tasks() => { // backpressure: don't queue if pending writes let task = self.h1_response_filter(session, task?, ctx, &mut serve_from_cache, &mut range_body_filter, true).await?; debug!("serve_from_cache task {task:?}"); - match session.write_response_tasks(vec![task]).await { - Ok(b) => response_state.maybe_set_cache_done(b), - Err(e) => if serve_from_cache.is_miss() { - // give up writing to downstream but wait for upstream cache write to finish - downstream_state.to_errored(); - response_state.maybe_set_cache_done(true); - warn!( - "Downstream Error ignored during caching: {}, {}", - e, - self.inner.request_summary(session, ctx) - ); - // This will not be treated as a final error, but we should signal to - // downstream session regardless - session.downstream_session.on_proxy_failure(e); - continue; - } else { - return Err(e); + if session.downstream_session.supports_proxy_task_api() { + session.send_downstream_proxy_task(task).await?; + } else { + match session.write_response_tasks(vec![task]).await { + Ok(b) => response_state.maybe_set_cache_done(b), + Err(e) => if serve_from_cache.is_miss() { + // give up writing to downstream but wait for upstream cache write to finish + downstream_state.to_errored(); + response_state.maybe_set_cache_done(true); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + // This will not be treated as a final error, but we should signal to + // downstream session regardless + session.downstream_session.on_proxy_failure(e); + continue; + } else { + return Err(e); + } + } + // A storage error can disable cache between cached_done + // being set and here; disable() drops the enabled_ctx so + // finish_hit_handler would panic without this guard. + if response_state.cached_done() && session.cache.enabled() { + if let Err(e) = session.cache.finish_hit_handler().await { + warn!("Error during finish_hit_handler: {}", e); + } } } - if response_state.cached_done() { - if let Err(e) = session.cache.finish_hit_handler().await { - warn!("Error during finish_hit_handler: {}", e); + } + + // Write queued downstream proxy tasks while also polling for upstream tasks. + // This allows cache writes to continue even when downstream is stalled. + // + // "Gate" branch: ready(()) resolves immediately, so the guard controls + // whether we enter. This is not a busy-loop because every path through + // the inner select either (a) drains all pending tasks via + // write_downstream_proxy_tasks (making the guard false), (b) stores an + // upstream task in next_upstream_task (making the guard false), or + // (c) blocks on real I/O inside the nested select. + _ = std::future::ready(()), if session.has_pending_downstream_tasks() && next_upstream_task.is_none() => { + tokio::select! { + // Try to write downstream proxy tasks (cancel-safe) + write_result = session.write_downstream_proxy_tasks() => { + match write_result { + Ok(end) => { + response_state.maybe_set_cache_done(end); + // See enabled() guard comment above. + if response_state.cached_done() && session.cache.enabled() { + if let Err(e) = session.cache.finish_hit_handler().await { + warn!("Error during finish_hit_handler: {}", e); + } + } + } + Err(e) => if serve_from_cache.is_miss() { + // give up writing to downstream but wait for upstream cache write to finish + downstream_state.to_errored(); + response_state.maybe_set_cache_done(true); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream write error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + // This will not be treated as a final error, but we should signal to + // downstream session regardless + session.downstream_session.on_proxy_failure(e); + } else { + return Err(e); + } + } + } + + // Also poll for upstream tasks - if we get one, cancel the write and handle it. + // Only poll if there is no buffered task already waiting to be processed. + upstream_task = rx.recv(), if !response_state.upstream_done() && serve_from_cache.is_on() && next_upstream_task.is_none() => { + if let Some(t) = upstream_task { + // Store this upstream task to be processed next iteration + next_upstream_task = Some(t); + continue; + } else { + response_state.maybe_set_upstream_done(true); + } } } } diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs index 808da5bc5..b5ff83731 100644 --- a/pingora-proxy/src/proxy_h2.rs +++ b/pingora-proxy/src/proxy_h2.rs @@ -89,7 +89,7 @@ where { let mut req = session.req_header().clone(); - if req.version != Version::HTTP_2 { + if req.version != Version::HTTP_2 || session.downstream_session.is_custom() { /* remove H1 specific headers */ // https://github.com/hyperium/h2/blob/d3b9f1e36aadc1a7a6804e2f8e86d3fe4a244b4f/src/proto/streams/send.rs#L72 req.remove_header(&http::header::TRANSFER_ENCODING); @@ -213,11 +213,21 @@ where match ret { Ok((downstream_can_reuse, _upstream)) => (downstream_can_reuse, None), Err(e) => { + let upstream_read_timeout = + e.esource == ErrorSource::Upstream && matches!(e.etype, ReadTimedout); + let downstream_error = e.esource == ErrorSource::Downstream; // On application level upstream read timeouts, send RST_STREAM CANCEL, - // we know we have not received END_STREAM at this point since we read timed out + // we know we have not received END_STREAM at this point since we read timed out. + // Also cancel the upstream stream when downstream goes away/resets so the + // upstream peer can release the stream promptly. // TODO: implement for write timeouts? - if e.esource == ErrorSource::Upstream && matches!(e.etype, ReadTimedout) { + if upstream_read_timeout || downstream_error { client_body.send_reset(h2::Reason::CANCEL); + if upstream_read_timeout { + // Mark the underlying H2 connection for shutdown so it's not used + // for new streams in case it is hung. + client_session.conn.mark_shutdown(); + } } (false, Some(e)) } @@ -262,6 +272,89 @@ where (server_session_reuse, error) } + #[allow(clippy::too_many_arguments)] + async fn process_upstream_tasks_h2( + &self, + session: &mut Session, + ctx: &mut SV::CTX, + initial_task: HttpTask, + rx: &mut mpsc::Receiver, + serve_from_cache: &mut ServeFromCache, + range_body_filter: &mut proxy_cache::range_filter::RangeBodyFilter, + response_state: &mut ResponseStateMachine, + ) -> Result> + where + SV: ProxyHttp + Send + Sync, + SV::CTX: Send + Sync, + { + if serve_from_cache.should_discard_upstream() { + // just drain, do we need to do anything else? + return Ok(None); + } + + // Batch: pull as many tasks as we can from rx + let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); + tasks.push(initial_task); + // tokio::task::unconstrained because now_or_never may yield None when the future is ready + while let Some(maybe_task) = tokio::task::unconstrained(rx.recv()).now_or_never() { + if let Some(t) = maybe_task { + tasks.push(t); + } else { + break; // upstream closed + } + } + + /* run filters before sending to downstream */ + let mut filtered_tasks = Vec::with_capacity(TASK_BUFFER_SIZE); + for mut t in tasks { + if self.revalidate_or_stale(session, &mut t, ctx).await { + serve_from_cache.enable(); + response_state.enable_cached_response(); + // skip downstream filtering entirely as the 304 will not be sent + break; + } + #[cfg(feature = "upstream_modules")] + if let HttpTask::Header(header, end_of_stream) = &t { + self.inner + .adjust_upstream_modules(session, header, *end_of_stream, ctx) + .await?; + } + #[cfg(feature = "upstream_modules")] + session.upstream_modules_filter_task(&mut t).await?; + session.upstream_compression.response_filter(&mut t); + // check error and abort + // otherwise the error is surfaced via write_response_tasks() + if !serve_from_cache.should_send_to_downstream() { + if let HttpTask::Failed(e) = t { + return Err(e); + } + } + filtered_tasks.push( + self.h2_response_filter( + session, + t, + ctx, + serve_from_cache, + range_body_filter, + false, + ) + .await?, + ); + if serve_from_cache.is_miss_header() { + response_state.enable_cached_response(); + } + } + + if !serve_from_cache.should_send_to_downstream() { + // TODO: need to derive response_done from filtered_tasks in case downstream failed already + return Ok(None); + } + + let response_done = session.write_response_tasks(filtered_tasks).await?; + + Ok(Some(response_done)) + } + // returns whether server (downstream) session can be reused async fn bidirection_down_to_up( &self, @@ -319,6 +412,8 @@ where let mut serve_from_cache = ServeFromCache::new(); let mut range_body_filter = proxy_cache::range_filter::RangeBodyFilter::new(); + let mut next_upstream_task: Option = None; + /* duplex mode * see the Same function for h1 for more comments */ @@ -356,11 +451,18 @@ where if wait_for_cache_fill { // ignore downstream error so that upstream can continue to write cache downstream_state.to_errored(); - warn!( - "Downstream Error ignored during caching: {}, {}", - e, - self.inner.request_summary(session, ctx) - ); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } // This will not be treated as a final error, but we should signal to // downstream session regardless session.downstream_session.on_proxy_failure(e); @@ -375,6 +477,35 @@ where Ok(request_done) => { downstream_state.maybe_finished(request_done); }, + Err(e) if e.esource == ErrorSource::Downstream => { + // Downstream reset/errored while the upstream write was blocked + // (e.g. on upstream flow control). Same policy as the read error + // handling above: ignore the downstream error if the upstream + // response is being admitted to cache, otherwise fail so the + // downstream stream handles are dropped promptly. + let wait_for_cache_fill = (!serve_from_cache.is_on() && support_cache_partial_read) + || serve_from_cache.is_miss(); + if !wait_for_cache_fill { + return Err(e); + } + // ignore downstream error so that upstream can continue to write cache + downstream_state.to_errored(); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + // This will not be treated as a final error, but we should signal to + // downstream session anyway. + session.downstream_session.on_proxy_failure(e); + }, Err(e) => { // mark request done, attempt to drain receive warn!("Upstream h2 body send error: {e}"); @@ -385,57 +516,47 @@ where }; }, - task = rx.recv(), if !response_state.upstream_done() => { + // Handle buffered upstream task from previous iteration + task = async { next_upstream_task.take() }, if next_upstream_task.is_some() => { + debug!("buffered upstream event: {:?}", task); if let Some(t) = task { - debug!("upstream event: {:?}", t); - if serve_from_cache.should_discard_upstream() { - // just drain, do we need to do anything else? - continue; - } - // pull as many tasks as we can - let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - tasks.push(t); - // tokio::task::unconstrained because now_or_never may yield None when the future is ready - while let Some(maybe_task) = tokio::task::unconstrained(rx.recv()).now_or_never() { - if let Some(t) = maybe_task { - tasks.push(t); - } else { - break - } - } - - /* run filters before sending to downstream */ - let mut filtered_tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - for mut t in tasks { - if self.revalidate_or_stale(session, &mut t, ctx).await { - serve_from_cache.enable(); - response_state.enable_cached_response(); - // skip downstream filtering entirely as the 304 will not be sent - break; - } - session.upstream_compression.response_filter(&mut t); - // check error and abort - // otherwise the error is surfaced via write_response_tasks() - if !serve_from_cache.should_send_to_downstream() { - if let HttpTask::Failed(e) = t { - return Err(e); - } - } - filtered_tasks.push( - self.h2_response_filter(session, t, ctx, - &mut serve_from_cache, - &mut range_body_filter, false).await?); - if serve_from_cache.is_miss_header() { - response_state.enable_cached_response(); - } - } - - if !serve_from_cache.should_send_to_downstream() { - // TODO: need to derive response_done from filtered_tasks in case downstream failed already + let Some(response_done) = self.process_upstream_tasks_h2( + session, + ctx, + t, + &mut rx, + &mut serve_from_cache, + &mut range_body_filter, + &mut response_state, + ).await? else { + // nothing sent downstream e.g. serve_from_cache continue; + }; + if session.was_upgraded() { + return Error::e_explain(H2Error, "upgraded while proxying to h2 session"); } + response_state.maybe_set_upstream_done(response_done); + } else { + debug!("empty upstream event"); + response_state.maybe_set_upstream_done(true); + } + }, - let response_done = session.write_response_tasks(filtered_tasks).await?; + task = rx.recv(), if !response_state.upstream_done() && next_upstream_task.is_none() => { + debug!("upstream event: {:?}", task); + if let Some(t) = task { + let Some(response_done) = self.process_upstream_tasks_h2( + session, + ctx, + t, + &mut rx, + &mut serve_from_cache, + &mut range_body_filter, + &mut response_state, + ).await? else { + // nothing sent downstream e.g. serve_from_cache + continue; + }; if session.was_upgraded() { // it is very weird if the downstream session decides to upgrade // since the client h2 session cannot, return an error on this case @@ -446,37 +567,113 @@ where debug!("empty upstream event"); response_state.maybe_set_upstream_done(true); } - } + }, task = serve_from_cache.next_http_task(&mut session.cache, &mut range_body_filter, upgraded), - if !response_state.cached_done() && !downstream_state.is_errored() && serve_from_cache.is_on() => { + if !response_state.cached_done() + && !downstream_state.is_errored() + && serve_from_cache.is_on() + && !session.has_pending_downstream_tasks() => { // backpressure: don't queue if pending writes + let task = self.h2_response_filter(session, task?, ctx, &mut serve_from_cache, &mut range_body_filter, true).await?; debug!("serve_from_cache task {task:?}"); - match session.write_response_tasks(vec![task]).await { - Ok(b) => response_state.maybe_set_cache_done(b), - Err(e) => if serve_from_cache.is_miss() { - // give up writing to downstream but wait for upstream cache write to finish - downstream_state.to_errored(); - response_state.maybe_set_cache_done(true); - warn!( - "Downstream Error ignored during caching: {}, {}", - e, - self.inner.request_summary(session, ctx) - ); - // This will not be treated as a final error, but we should signal to - // downstream session regardless - session.downstream_session.on_proxy_failure(e); - continue; - } else { - return Err(e); + if session.downstream_session.supports_proxy_task_api() { + session.send_downstream_proxy_task(task).await?; + } else { + match session.write_response_tasks(vec![task]).await { + Ok(b) => response_state.maybe_set_cache_done(b), + Err(e) => if serve_from_cache.is_miss() { + // give up writing to downstream but wait for upstream cache write to finish + downstream_state.to_errored(); + response_state.maybe_set_cache_done(true); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream Error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + // This will not be treated as a final error, but we should signal to + // downstream session regardless + session.downstream_session.on_proxy_failure(e); + continue; + } else { + return Err(e); + } + } + // A storage error can disable cache between cached_done + // being set and here; see the same guard in proxy_h1.rs. + if response_state.cached_done() && session.cache.enabled() { + if let Err(e) = session.cache.finish_hit_handler().await { + warn!("Error during finish_hit_handler: {}", e); + } } } - if response_state.cached_done() { - if let Err(e) = session.cache.finish_hit_handler().await { - warn!("Error during finish_hit_handler: {}", e); + } + + // Write queued downstream proxy tasks while also polling for upstream tasks. + // This allows cache writes to continue even when downstream is stalled. + // + // "Gate" branch: ready(()) resolves immediately, so the guard controls + // whether we enter. This is not a busy-loop because every path through + // the inner select either (a) drains all pending tasks via + // write_downstream_proxy_tasks (making the guard false), (b) stores an + // upstream task in next_upstream_task (making the guard false), or + // (c) blocks on real I/O inside the nested select. + _ = std::future::ready(()), if session.has_pending_downstream_tasks() && next_upstream_task.is_none() => { + tokio::select! { + // Try to write downstream proxy tasks (cancel-safe) + write_result = session.write_downstream_proxy_tasks() => { + match write_result { + Ok(end) => { + response_state.maybe_set_cache_done(end); + // See disabled() guard comment above. + // See enabled() guard comment above. + if response_state.cached_done() && session.cache.enabled() { + if let Err(e) = session.cache.finish_hit_handler().await { + warn!("Error during finish_hit_handler: {}", e); + } + } + } + Err(e) => if serve_from_cache.is_miss() { + // give up writing to downstream but wait for upstream cache write to finish + downstream_state.to_errored(); + response_state.maybe_set_cache_done(true); + if !self.inner.suppress_proxy_warn_log( + session, + ctx, + &e, + ProxyWarnLogContext::DownstreamCache, + ) { + warn!( + "Downstream write error ignored during caching: {}, {}", + e, + self.inner.request_summary(session, ctx) + ); + } + session.downstream_session.on_proxy_failure(e); + } else { + return Err(e); + } + } + } + + // Also poll for upstream tasks - if we get one, cancel the write and handle it. + upstream_task = rx.recv(), if !response_state.upstream_done() && serve_from_cache.is_on() && next_upstream_task.is_none() => { + if let Some(t) = upstream_task { + next_upstream_task = Some(t); + continue; + } else { + response_state.maybe_set_upstream_done(true); + } } } } @@ -736,15 +933,42 @@ where return Ok(false); } - if let Some(data) = data { - debug!("Write {} bytes body to h2 upstream", data.len()); - write_body(client_body, data, end_of_body, write_timeout) - .await - .map_err(|e| e.into_up())?; + let (data, end) = match data { + Some(data) => { + debug!("Write {} bytes body to h2 upstream", data.len()); + (data, end_of_body) + } + None => { + debug!("Read downstream body done"); + /* send a standalone END_STREAM flag */ + (Bytes::new(), true) + } + }; + + /* For H2 downstreams, race the upstream write against a downstream stream + * reset. A write blocked on upstream flow control would otherwise keep the + * downstream stream handles referenced while a downstream RST_STREAM goes + * unobserved, pinning the downstream connection window credit until the + * write completes. */ + if let Some(reset) = session.downstream_session.watch_h2_stream_reset() { + tokio::select! { + biased; + res = write_body(client_body, data, end, write_timeout) => { + res.map_err(|e| e.into_up())?; + } + reset = reset => { + return match reset { + Ok(reason) => Error::e_explain( + H2Error, + format!("downstream reset stream (reason: {reason}) while writing body to upstream"), + ), + Err(e) => Err(e), + } + .map_err(|e| e.into_down()); + } + } } else { - debug!("Read downstream body done"); - /* send a standalone END_STREAM flag */ - write_body(client_body, Bytes::new(), true, write_timeout) + write_body(client_body, data, end, write_timeout) .await .map_err(|e| e.into_up())?; } @@ -806,12 +1030,25 @@ pub(crate) async fn pipe_up_to_down_response( } } - while let Some(chunk) = client - .read_response_body() - .await - .map_err(|e| e.into_up()) - .transpose() - { + // Read body from H2 upstream, racing each read against tx.closed(). + // + // When proxying an H2 upstream response with Content-Length to an H1 downstream, + // bidirection_down_to_up() may determine the response is complete (all Content-Length + // bytes written) and exit before the H2 stream signals END_STREAM. This drops the + // receiving end (rx) of the channel. Without this race, read_response_body() would + // block until the H2 stream eventually ends (e.g. via trailers or read_timeout), + // while the downstream side (which could be H1) is in theory already done. + loop { + let chunk = tokio::select! { + biased; + body = client.read_response_body() => { + body.map_err(|e| e.into_up()).transpose() + } + _ = tx.closed() => None, + }; + let Some(chunk) = chunk else { + break; + }; let data = match chunk { Ok(d) => d, Err(e) => { @@ -834,10 +1071,10 @@ pub(crate) async fn pipe_up_to_down_response( .send(HttpTask::Body(Some(data), eos)) .await .or_err(InternalError, "sending h2 body to pipe"); - // If the if the response with content-length is sent to an HTTP1 downstream, + // If the response with content-length is sent to an HTTP1 downstream, // bidirection_down_to_up() could decide that the body has finished and exit without // waiting for this function to signal the eos. In this case tx being closed is not - // an sign of error. It should happen if the only thing left for the h2 to send is + // a sign of error. It should happen if the only thing left for the h2 to send is // an empty data frame with eos set. if sent.is_err() && eos && empty { return Ok(()); @@ -852,12 +1089,26 @@ pub(crate) async fn pipe_up_to_down_response( } } - // attempt to get trailers - let trailers = match client.read_trailers().await { - Ok(t) => t, - Err(e) => { - // Similar to above, push the error to downstream and then quit - let _ = tx.send(HttpTask::Failed(e.into_up())).await; + // If the channel is already closed, downstream is finished + // TODO: note that this does skip trailers/done, but downstream + // has already finished so no more is in theory necessary to send + if tx.is_closed() { + return Ok(()); + } + + // attempt to get trailers, racing against channel close + let trailers = tokio::select! { + biased; + t = client.read_trailers() => { + match t { + Ok(t) => t, + Err(e) => { + let _ = tx.send(HttpTask::Failed(e.into_up())).await; + return Ok(()); + } + } + } + _ = tx.closed() => { return Ok(()); } }; diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs index d5a3efde4..d5f911211 100644 --- a/pingora-proxy/src/proxy_trait.rs +++ b/pingora-proxy/src/proxy_trait.rs @@ -19,8 +19,25 @@ use pingora_cache::{ RespCacheable::{self, *}, }; use proxy_cache::range_filter::{self}; +use std::any::Any; use std::time::Duration; +/// Context for proxy warning logs that can be suppressed by +/// [`ProxyHttp::suppress_proxy_warn_log`]. +/// +/// These contexts are distinct from final proxy errors, which are handled by +/// [`ProxyHttp::suppress_error_log`]. +/// +/// Experimental: this API may change or be removed until indicated otherwise. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[non_exhaustive] +pub enum ProxyWarnLogContext { + /// A proxy upstream attempt failed with a retryable error. + UpstreamRetry, + /// A downstream error was ignored so cache fill could continue. + DownstreamCache, +} + /// The interface to control the HTTP proxy /// /// The methods in [ProxyHttp] are filters/callbacks which will be performed on all requests at their @@ -56,6 +73,23 @@ pub trait ProxyHttp { modules.add_module(ResponseCompressionBuilder::enable(0)); } + /// Set up upstream modules. + /// + /// In this phase, users can add [HttpModules] that will process upstream responses + /// **before** `upstream_compression`. This is the correct place to register modules + /// that need to observe the raw (pre-compression) upstream response body, such as + /// a dictionary store for shared dictionary compression. + /// + /// Upstream modules are ordered by [`HttpModuleBuilder::order()`]: higher values run + /// first. They are invoked on each upstream response task (header, body, trailers) + /// before `upstream_compression` processes the task. + /// + /// By default this method does nothing. + /// + /// This method requires the `upstream_modules` feature to be enabled. + #[cfg(feature = "upstream_modules")] + fn init_upstream_modules(&self, _modules: &mut HttpModules) {} + /// Handle the incoming request. /// /// In this phase, users can parse, validate, rate limit, perform access control and/or @@ -292,6 +326,39 @@ pub trait ProxyHttp { Ok(()) } + /// Adjust upstream modules before they process the response header. + /// + /// This filter is called when the upstream response header arrives, before upstream modules + /// (such as `upstream_compression`) run their response header filter. Use this to configure + /// module behavior based on the response, e.g. setting a dictionary for dictionary-based + /// content encoding. + /// + /// This filter may be called more than once per request if the upstream sends informational + /// (1xx) response headers before the final response. Implementations can check + /// [`upstream_response.status.is_informational()`](http::StatusCode::is_informational) to + /// distinguish informational headers from the final response if needed. + /// + /// `end_of_stream` indicates whether the response header is also the end of the response + /// (e.g. for HEAD responses or 304s with no body). + /// + /// The response header is provided as an immutable reference. To modify the response header + /// itself, use [`Self::upstream_response_filter()`] instead. + /// + /// This filter requires the `upstream_modules` feature to be enabled. + #[cfg(feature = "upstream_modules")] + async fn adjust_upstream_modules( + &self, + _session: &mut Session, + _upstream_response: &ResponseHeader, + _end_of_stream: bool, + _ctx: &mut Self::CTX, + ) -> Result<()> + where + Self::CTX: Send + Sync, + { + Ok(()) + } + /// Modify the response header from the upstream /// /// The modification is before caching, so any change here will be stored in the cache if enabled. @@ -438,11 +505,70 @@ pub trait ProxyHttp { { } + /// Called after [`Self::logging`] when the downstream connection will be reused for another + /// HTTP/1.x keepalive request. The returned value, if any, will be carried to the next + /// request on this connection and delivered via [`Self::on_connection_reuse`]. + /// + /// Use this to persist debugging or timing information across keepalive requests. + /// This is only called for HTTP/1.x keepalive connections, not for HTTP/2. + /// It is also called on error paths when the downstream connection is eligible for reuse. + /// + /// The default implementation returns `None` (no context persisted). + fn persist_connection_context( + &self, + _session: &Session, + _ctx: &Self::CTX, + ) -> Option> { + None + } + + /// Called at the start of a new request on a reused HTTP/1.x keepalive connection, + /// before [`Self::early_request_filter`]. The `prev_ctx` argument is the value returned + /// by [`Self::persist_connection_context`] from the previous request on this connection. + /// + /// This is only called for HTTP/1.x keepalive connections, not for HTTP/2. + /// It is not called when `persist_connection_context` returned `None` on the previous request. + /// + /// Use this to transfer state from the previous request into the new request's context. + fn on_connection_reuse( + &self, + _session: &mut Session, + _ctx: &mut Self::CTX, + _prev_ctx: Box, + ) { + } + /// A value of true means that the log message will be suppressed. The default value is false. + /// + /// See also: [`Self::suppress_proxy_warn_log`]. fn suppress_error_log(&self, _session: &Session, _ctx: &Self::CTX, _error: &Error) -> bool { false } + /// A value of true means that the proxy warning log message will be suppressed. + /// The default value is false. + /// + /// This hook currently applies to retryable proxy upstream failures and downstream errors + /// ignored while cache fill continues. Final proxy errors are still handled by + /// [`Self::suppress_error_log`]. + /// + /// Suppressing retry warning logs can remove the only per-retry audit record. Callers that + /// suppress these logs should provide alternative observability, such as metrics or logs in + /// their implementation of this hook. + /// + /// This hook runs inline on retry and cache-error paths, so implementations should be cheap. + /// + /// Experimental: this API may change or be removed until indicated otherwise. + fn suppress_proxy_warn_log( + &self, + _session: &Session, + _ctx: &Self::CTX, + _error: &Error, + _context: ProxyWarnLogContext, + ) -> bool { + false + } + /// This filter is called when there is an error **after** a connection is established (or reused) /// to the upstream. fn error_while_proxy( diff --git a/pingora-proxy/src/subrequest/pipe.rs b/pingora-proxy/src/subrequest/pipe.rs index 6dd4a57ec..7845d4dc2 100644 --- a/pingora-proxy/src/subrequest/pipe.rs +++ b/pingora-proxy/src/subrequest/pipe.rs @@ -42,16 +42,37 @@ pub enum InputBodyType { SaveBody(usize), } -/// Context struct as a result of subrequest piping. -#[derive(Clone)] +/// Outcome of [`pipe_subrequest`]. +#[derive(Debug, Default)] pub struct PipeSubrequestState { - /// The saved (captured) body from the main session. + /// Captured body from the main session. pub saved_body: Option, + /// Did the subrequest produce a response header? Checked before the task + /// filter runs, so a filtered-out header still counts. + pub header_received: bool, + /// The spawned subrequest task handle. Always set after spawn. Caller is + /// responsible for awaiting/inspecting state. + pub join_handle: Option>, + /// The receiving half of the pipe channel. When the coordinator exits + /// `pipe_subrequest` before the subrequest task finishes writing, this + /// receiver must be kept alive and drained alongside the join handle; + /// otherwise dropping it breaks the pipe and prevents the writer from + /// completing its cache-write lifecycle. + pub pipe_rx: Option>, } impl PipeSubrequestState { - fn new() -> PipeSubrequestState { - PipeSubrequestState { saved_body: None } + /// Creates a snapshot for error reporting, excluding the join handle. + /// Moves `pipe_rx` into the snapshot so the receiver stays alive through + /// the error path and is not dropped when `self` is cleaned up. + /// Used by [`map_pipe_err`] to capture state at the point of failure. + pub fn snapshot_for_error(&mut self) -> Self { + PipeSubrequestState { + saved_body: self.saved_body.clone(), + header_received: self.header_received, + join_handle: None, + pipe_rx: self.pipe_rx.take(), + } } } @@ -79,9 +100,9 @@ impl PipeSubrequestError { fn map_pipe_err>>( result: Result, from_subreq: bool, - state: &PipeSubrequestState, + state: &mut PipeSubrequestState, ) -> Result { - result.map_err(|e| PipeSubrequestError::new(e, from_subreq, state.clone())) + result.map_err(|e| PipeSubrequestError::new(e, from_subreq, state.snapshot_for_error())) } #[derive(Debug, Clone)] @@ -182,13 +203,13 @@ where }; let mut downstream_state = DownstreamStateMachine::new(no_body_input); - let mut state = PipeSubrequestState::new(); - state.saved_body = saved_body; + let mut state = PipeSubrequestState { + saved_body, + ..Default::default() + }; - // Have the subrequest remove all body-related headers if no body will be sent - // TODO: we could also await the join handle, but subrequest may be running logging phase - // also the full run() may also await cache fill if downstream fails - let _join_handle = tokio::spawn(async move { + // Remove headers if no body. + let join_handle = tokio::spawn(async move { if no_body_input { subrequest .session_mut() @@ -196,10 +217,14 @@ where .expect("PreparedSubrequest must be subrequest") .clear_request_body_headers(); } - subrequest.run().await + let _ = subrequest.run().await; }); + state.join_handle = Some(join_handle); let tx = subrequest_handle.tx; - let mut rx = subrequest_handle.rx; + // Move rx into state immediately so it survives all exit paths (early `?` + // returns, errors, and the normal success path). The select loop borrows it + // back via `state.pipe_rx.as_mut().expect(...)`. + state.pipe_rx = Some(subrequest_handle.rx); let mut wants_body = false; let mut wants_body_rx_err = false; @@ -216,20 +241,27 @@ where .or_err(InternalError, "try_reserve() body pipe for subrequest"); tokio::select! { - task = rx.recv(), if !response_state.upstream_done() => { + task = state.pipe_rx.as_mut().expect("pipe_rx always set after spawn").recv(), if !response_state.upstream_done() => { debug!("upstream event: {:?}", task); if let Some(t) = task { + // Did the subrequest get headers? + if matches!(&t, HttpTask::Header(..)) { + state.header_received = true; + } // pull as many tasks as we can const TASK_BUFFER_SIZE: usize = 4; let mut tasks = Vec::with_capacity(TASK_BUFFER_SIZE); - let task = map_pipe_err(task_filter(t), false, &state)?; + let task = map_pipe_err(task_filter(t), false, &mut state)?; if let Some(filtered) = task { tasks.push(filtered); } // tokio::task::unconstrained because now_or_never may yield None when the future is ready - while let Some(maybe_task) = tokio::task::unconstrained(rx.recv()).now_or_never() { + while let Some(maybe_task) = tokio::task::unconstrained(state.pipe_rx.as_mut().expect("pipe_rx always set after spawn").recv()).now_or_never() { if let Some(t) = maybe_task { - let task = map_pipe_err(task_filter(t), false, &state)?; + if matches!(&t, HttpTask::Header(..)) { + state.header_received = true; + } + let task = map_pipe_err(task_filter(t), false, &mut state)?; if let Some(filtered) = task { tasks.push(filtered); } @@ -239,7 +271,7 @@ where } // FIXME: if one of these tasks is Failed(e), the session will return that // error; in this case, the error is actually from the subreq - let response_done = map_pipe_err(session.write_response_tasks(tasks).await, false, &state)?; + let response_done = map_pipe_err(session.write_response_tasks(tasks).await, false, &mut state)?; // NOTE: technically it is the downstream whose response state has finished here // we consider the subrequest's work done however @@ -248,9 +280,7 @@ where // (can only happen with a real session, TODO to allow with preset body) downstream_state.maybe_finished(!use_preset_body && session.is_body_done()); } else { - // quite possible that the subrequest may be finished, though the main session - // is not - we still must exit in this case - debug!("empty upstream event"); + debug!("upstream channel closed early"); response_state.maybe_set_upstream_done(true); } }, @@ -291,7 +321,7 @@ where // this is the first subrequest // send the body debug!("downstream event: main body for subrequest"); - let body = map_pipe_err(body.map_err(|e| e.into_down()), false, &state)?; + let body = map_pipe_err(body.map_err(|e| e.into_down()), false, &mut state)?; // If the request is websocket, `None` body means the request is closed. // Set the response to be done as well so that the request completes normally. @@ -307,7 +337,7 @@ where state.saved_body.as_mut(), send_permit.expect("checked is_ok()"), ) - .await, false, &state)?; + .await, false, &mut state)?; downstream_state.maybe_finished(request_done); @@ -328,7 +358,7 @@ where is_body_done, None, send_permit.expect("checked is_ok()"), - ), false, &state)?; + ), false, &mut state)?; downstream_state.maybe_finished(request_done); }, @@ -397,3 +427,64 @@ fn do_send_body_to_pipe( Ok(end_of_body) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::subrequest::Ctx as SubrequestCtx; + use crate::{Session, Subrequest, SubrequestSpawner}; + use async_trait::async_trait; + use pingora_core::protocols::http::ServerSession as HttpSession; + use std::sync::Arc; + + /// Drops session without producing output — channels close, rx returns None. + struct NoopApp; + + #[async_trait] + impl Subrequest for NoopApp { + async fn process_subrequest( + self: Arc, + _session: Box, + _ctx: Box, + ) { + } + } + + async fn mock_session() -> Session { + let input = b"GET / HTTP/1.1\r\nHost: test\r\n\r\n"; + let mock_io = tokio_test::io::Builder::new().read(&input[..]).build(); + let mut session = Session::new_h1(Box::new(mock_io) as pingora_core::protocols::Stream); + session + .downstream_session + .read_request() + .await + .expect("mock request should parse"); + session + } + + #[tokio::test] + async fn no_header_received_when_subrequest_exits_silently() { + let mut session = mock_session().await; + + let spawner = SubrequestSpawner::new(Arc::new(NoopApp)); + let ctx = SubrequestCtx::builder().body_mode(BodyMode::NoBody).build(); + let (subrequest, handle) = spawner.create_subrequest(session.as_downstream(), ctx); + + let result = pipe_subrequest( + &mut session, + subrequest, + handle, + |task| Ok(Some(task)), + InputBodyType::Preset(InputBody::NoBody), + ) + .await; + + let state = + result.unwrap_or_else(|e| panic!("pipe should return Ok, not Err: {:?}", e.error)); + assert!( + !state.header_received, + "no header should have been received from the no-op subrequest" + ); + assert!(state.join_handle.is_some(), "task handle should be set"); + } +} diff --git a/pingora-proxy/tests/test_basic.rs b/pingora-proxy/tests/test_basic.rs index 77303fc30..85a3bfc12 100644 --- a/pingora-proxy/tests/test_basic.rs +++ b/pingora-proxy/tests/test_basic.rs @@ -17,14 +17,18 @@ mod utils; use bytes::Bytes; use h2::client; use http::Request; -use hyper::{body::HttpBody, header::HeaderValue, Body, Client}; +use http_body_util::BodyExt; +use hyper_util::client::legacy::Client; #[cfg(unix)] use hyperlocal::{UnixClientExt, Uri}; use reqwest::{header, StatusCode}; +#[cfg(feature = "patched_http1")] use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use utils::server_utils::init; +use utils::server_utils::{ + init, reset_suppress_proxy_warn_log_calls, suppress_proxy_warn_log_calls, +}; fn is_specified_port(port: u16) -> bool { (1..65535).contains(&port) @@ -161,21 +165,21 @@ async fn test_h2_to_h2() { async fn test_h2c_to_h2c() { init(); - let client = hyper::client::Client::builder() + let client = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()) .http2_only(true) - .build_http(); + .build_http::>(); - let mut req = hyper::Request::builder() + let mut req = http::Request::builder() .uri("http://127.0.0.1:6146") - .body(Body::empty()) + .body(http_body_util::Empty::::new()) .unwrap(); req.headers_mut() - .insert("x-h2", HeaderValue::from_bytes(b"true").unwrap()); + .insert("x-h2", http::HeaderValue::from_bytes(b"true").unwrap()); let res = client.request(req).await.unwrap(); assert_eq!(res.status(), reqwest::StatusCode::OK); assert_eq!(res.version(), reqwest::Version::HTTP_2); - let body = res.into_body().data().await.unwrap().unwrap(); + let body = res.into_body().collect().await.unwrap().to_bytes(); assert_eq!(body.as_ref(), b"Hello World!\n"); } @@ -183,21 +187,21 @@ async fn test_h2c_to_h2c() { async fn test_h1_on_h2c_port() { init(); - let client = hyper::client::Client::builder() + let client = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()) .http2_only(false) - .build_http(); + .build_http::>(); - let mut req = hyper::Request::builder() + let mut req = http::Request::builder() .uri("http://127.0.0.1:6146") - .body(Body::empty()) + .body(http_body_util::Empty::::new()) .unwrap(); req.headers_mut() - .insert("x-h2", HeaderValue::from_bytes(b"true").unwrap()); + .insert("x-h2", http::HeaderValue::from_bytes(b"true").unwrap()); let res = client.request(req).await.unwrap(); assert_eq!(res.status(), reqwest::StatusCode::OK); assert_eq!(res.version(), reqwest::Version::HTTP_11); - let body = res.into_body().data().await.unwrap().unwrap(); + let body = res.into_body().collect().await.unwrap().to_bytes(); assert_eq!(body.as_ref(), b"Hello World!\n"); } @@ -303,7 +307,7 @@ async fn test_h2_head() { async fn test_simple_proxy_uds() { init(); let url = Uri::new("/tmp/pingora_proxy.sock", "/").into(); - let client = Client::unix(); + let client: Client> = Client::unix(); let res = client.get(url).await.unwrap(); @@ -324,7 +328,10 @@ async fn test_simple_proxy_uds() { assert_eq!(sockaddr.ip().to_string(), "127.0.0.2"); assert!(is_specified_port(sockaddr.port())); - let body = hyper::body::to_bytes(body).await.unwrap(); + let body = http_body_util::BodyExt::collect(body) + .await + .unwrap() + .to_bytes(); assert_eq!(body.as_ref(), b"Hello World!\n"); } @@ -380,15 +387,18 @@ async fn test_dropped_conn_get() { assert_eq!(res.status(), StatusCode::OK); } + reset_suppress_proxy_warn_log_calls(); let res = client .get("http://127.0.0.1:6147/bad_lb") .header("x-port", port) + .header("x-test-suppress-proxy-warn-log", "true") .send() .await .unwrap(); // retry gives 200 assert_eq!(res.status(), StatusCode::OK); + assert!(suppress_proxy_warn_log_calls() > 0); let body = res.text().await.unwrap(); assert_eq!(body, "dog!\n"); } @@ -746,6 +756,9 @@ async fn test_connect_close() { assert_eq!(body, "Hello World!\n"); } +// Authority-form CONNECT request targets require patched HTTP/1 parsing until +// general request-target form support is available. +#[cfg(feature = "patched_http1")] #[tokio::test] async fn test_connect_proxying_disallowed_h1() { init(); @@ -784,6 +797,7 @@ async fn test_connect_proxying_disallowed_h2() { } } +#[cfg(feature = "patched_http1")] #[tokio::test] async fn test_connect_proxying_allowed_h1() { init(); @@ -1116,3 +1130,273 @@ async fn test_103_die() { let res = reqwest::get("http://127.0.0.1:6147/103-die").await.unwrap(); assert_eq!(res.status(), StatusCode::BAD_GATEWAY); } + +// A downstream RST_STREAM must be observed even while the proxy is blocked writing the +// request body to the upstream (parked on h2 flow control). Otherwise the stream is held +// open as a zombie: its handles stay referenced and the downstream connection-window +// credit is never released. On catching the RST the proxy should also cancel the +// upstream stream promptly. +#[tokio::test] +async fn test_h2_downstream_rst_while_upstream_write_blocked() { + use std::future::poll_fn; + use std::time::Duration; + + init(); + + // An h2c upstream that accepts one request but never reads its body and never + // sends window updates, so the proxy's upstream write blocks on flow control. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_port = listener.local_addr().unwrap().port(); + let (reset_tx, reset_rx) = tokio::sync::oneshot::channel::(); + + tokio::spawn(async move { + let (io, _) = listener.accept().await.unwrap(); + let mut conn = h2::server::Builder::new() + // tiny stream window so the proxy's write parks quickly + .initial_window_size(1024) + .handshake::<_, Bytes>(io) + .await + .unwrap(); + let (req, mut send_response) = conn.accept().await.unwrap().unwrap(); + // hold the body reader without reading it: no window updates are granted + let _body = req.into_body(); + + // keep driving the connection in the background + tokio::spawn(async move { + while let Some(res) = conn.accept().await { + if res.is_err() { + break; + } + } + }); + + // wait for the proxy to reset our stream + let reason = poll_fn(|cx| send_response.poll_reset(cx)).await.unwrap(); + let _ = reset_tx.send(reason); + }); + + // h2c downstream client to the proxy + let tcp = TcpStream::connect("127.0.0.1:6146").await.unwrap(); + let (mut client, conn) = client::handshake(tcp).await.unwrap(); + tokio::spawn(async move { + // ignore errors: the proxy may tear the connection down after the RST + let _ = conn.await; + }); + + let req = Request::builder() + .method("POST") + .uri("http://127.0.0.1:6146/") + .header("x-h2", "true") + .header("x-port", upstream_port.to_string()) + .body(()) + .unwrap(); + + let (_response, mut req_body) = client.send_request(req, false).unwrap(); + + // Push body until the proxy stops granting capacity, meaning it is no longer + // reading the downstream body because its upstream write is parked on flow control. + let mut sent = 0usize; + while sent < 512 * 1024 { + req_body.reserve_capacity(16 * 1024); + let granted = match tokio::time::timeout( + Duration::from_millis(500), + poll_fn(|cx| req_body.poll_capacity(cx)), + ) + .await + { + Ok(Some(Ok(n))) => n, + Ok(other) => panic!("downstream send capacity error: {other:?}"), + // no new capacity for a while: the proxy is parked on the upstream write + Err(_) => break, + }; + req_body + .send_data(Bytes::from(vec![0u8; granted]), false) + .unwrap(); + sent += granted; + } + // we must have at least filled the upstream stream window for the write to park + assert!(sent >= 1024, "only sent {sent} bytes"); + + // reset the stream while the proxy is blocked writing upstream + req_body.send_reset(h2::Reason::CANCEL); + + // the proxy should catch the RST promptly (not hang on the blocked write) + // and cancel the upstream stream + let reason = tokio::time::timeout(Duration::from_secs(5), reset_rx) + .await + .expect("proxy did not cancel the upstream stream after the downstream RST") + .expect("upstream watcher task died before observing a reset"); + assert_eq!(reason, h2::Reason::CANCEL); +} + +// Same as above, but with caching enabled and a cacheable upstream response mid-admission: +// the downstream RST caught during the blocked upstream write must be ignored (like +// downstream read errors are during caching) so the cache fill can complete. +#[tokio::test] +async fn test_h2_downstream_rst_during_cache_fill() { + use std::future::poll_fn; + use std::time::Duration; + + init(); + + let uri = "http://127.0.0.1:6154/test_h2_downstream_rst_during_cache_fill"; + + // An h2c upstream that never reads the request body (so the proxy's upstream write + // parks on flow control) but streams a cacheable response: first chunk right away, + // final chunk + EOS only after the test signals that downstream sent its RST. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_port = listener.local_addr().unwrap().port(); + let (finish_tx, finish_rx) = tokio::sync::oneshot::channel::<()>(); + let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>(); + + tokio::spawn(async move { + let (io, _) = listener.accept().await.unwrap(); + let mut conn = h2::server::Builder::new() + // tiny stream window so the proxy's request body write parks quickly + .initial_window_size(1024) + .handshake::<_, Bytes>(io) + .await + .unwrap(); + let (req, mut send_response) = conn.accept().await.unwrap().unwrap(); + // hold the body reader without reading it: no window updates are granted + let body = req.into_body(); + + // keep driving the connection in the background + tokio::spawn(async move { + while let Some(res) = conn.accept().await { + if res.is_err() { + break; + } + } + }); + + let resp = http::Response::builder() + .status(200) + .header("cache-control", "max-age=30") + .body(()) + .unwrap(); + let mut resp_body = send_response.send_response(resp, false).unwrap(); + resp_body + .send_data(Bytes::from_static(b"hello "), false) + .unwrap(); + + // hold the end of the response so the cache fill is still in progress + // when the downstream resets + finish_rx.await.unwrap(); + resp_body + .send_data(Bytes::from_static(b"world!"), true) + .unwrap(); + + // Keep the stream handles alive until the test is done. Dropping them here + // would make h2 send an implicit RST_STREAM(NO_ERROR) (response complete + // without consuming the request body), which the proxy's upstream read can + // observe before the clean EOS and abort the cache admission. + let _ = done_rx.await; + drop(body); + drop(resp_body); + }); + + // h2c downstream client to the caching proxy service + let tcp = TcpStream::connect("127.0.0.1:6154").await.unwrap(); + let (mut client, conn) = client::handshake(tcp).await.unwrap(); + tokio::spawn(async move { + let _ = conn.await; + }); + + let req = Request::builder() + .method("POST") + .uri(uri) + .header("x-h2", "true") + .header("x-port", upstream_port.to_string()) + .body(()) + .unwrap(); + let (response, mut req_body) = client.send_request(req, false).unwrap(); + + // a small first body chunk, fits the upstream window + req_body.reserve_capacity(16); + let granted = poll_fn(|cx| req_body.poll_capacity(cx)) + .await + .unwrap() + .unwrap(); + assert!(granted >= 16); + req_body + .send_data(Bytes::from_static(b"upload.........."), false) + .unwrap(); + + // wait for the response header and first chunk: the miss admission + // (cache fill) is now in progress + let (head, mut resp_body) = response.await.unwrap().into_parts(); + assert_eq!(head.status, 200); + assert_eq!(head.headers.get("x-cache-status").unwrap(), "miss"); + let chunk = resp_body.data().await.unwrap().unwrap(); + assert_eq!(&chunk[..], b"hello "); + let _ = resp_body.flow_control().release_capacity(chunk.len()); + + // flood the request body until the proxy stops granting capacity, i.e. it is + // parked writing to the upstream whose window is exhausted + let mut sent = 0usize; + while sent < 512 * 1024 { + req_body.reserve_capacity(16 * 1024); + let granted = match tokio::time::timeout( + Duration::from_millis(500), + poll_fn(|cx| req_body.poll_capacity(cx)), + ) + .await + { + Ok(Some(Ok(n))) => n, + Ok(other) => panic!("downstream send capacity error: {other:?}"), + Err(_) => break, + }; + req_body + .send_data(Bytes::from(vec![0u8; granted]), false) + .unwrap(); + sent += granted; + } + assert!(sent >= 1024, "only sent {sent} bytes"); + + // reset the stream while the proxy is blocked writing upstream; + // since a cache fill is in progress, the proxy should swallow this error + // and keep admitting the upstream response + req_body.send_reset(h2::Reason::CANCEL); + + // give the proxy a moment to observe the RST, then let the upstream finish + // the response + tokio::time::sleep(Duration::from_millis(200)).await; + finish_tx.send(()).unwrap(); + + // let the fill complete + tokio::time::sleep(Duration::from_secs(1)).await; + + // the object must now be fully in cache: a new request is a hit with the + // complete body and does not need the (single use) upstream + let tcp = TcpStream::connect("127.0.0.1:6154").await.unwrap(); + let (mut client, conn) = client::handshake(tcp).await.unwrap(); + tokio::spawn(async move { + let _ = conn.await; + }); + let req = Request::builder() + .method("GET") + .uri(uri) + .header("x-port", upstream_port.to_string()) + .body(()) + .unwrap(); + let (response, _) = client.send_request(req, true).unwrap(); + let (head, mut resp_body) = tokio::time::timeout(Duration::from_secs(5), response) + .await + .expect("no response for the second request") + .unwrap() + .into_parts(); + assert_eq!(head.status, 200); + assert_eq!(head.headers.get("x-cache-status").unwrap(), "hit"); + + let mut body = Vec::new(); + while let Some(chunk) = resp_body.data().await { + let chunk = chunk.unwrap(); + let _ = resp_body.flow_control().release_capacity(chunk.len()); + body.extend_from_slice(&chunk); + } + assert_eq!(body, b"hello world!"); + + // release the upstream's stream handles + let _ = done_tx.send(()); +} diff --git a/pingora-proxy/tests/test_upstream.rs b/pingora-proxy/tests/test_upstream.rs index d76c294d4..862009e5a 100644 --- a/pingora-proxy/tests/test_upstream.rs +++ b/pingora-proxy/tests/test_upstream.rs @@ -17,9 +17,11 @@ mod utils; use utils::server_utils::init; use utils::websocket::{WS_ECHO, WS_ECHO_RAW}; +use bytes::Bytes; use futures::{SinkExt, StreamExt}; +use http::header::{HeaderName, HeaderValue}; +use http_body_util::BodyExt; use pingora_http::ResponseHeader; -use reqwest::header::{HeaderName, HeaderValue}; use reqwest::{StatusCode, Version}; use std::time::{Duration, Instant}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -70,6 +72,25 @@ async fn test_connection_die() { assert!(body.is_err()); } +// This test is ignored because it has a fundamental timing dependency. +// +// The nginx origin sends a 200 response and flushes it, then sleeps 1s +// and kills the connection (RST). The test expects the client to always +// receive the 200 before the connection dies. +// +// This fails under CI load because the 15MB request body takes longer +// than 1s to write. The proxy's select! loop is busy writing body chunks +// upstream and can't read the 200 response concurrently. When the 1s +// expires and nginx sends a TCP RST, the RST discards all buffered data +// (including the 200) per TCP semantics. The proxy then sees an upstream +// error and resets the client connection. +// +// The underlying issue is that TCP RST discards unread buffered data, +// so the 200 response is lost even though it was sent before the RST. +// Fixing this would require the proxy to read the response before or +// concurrently with the body write completing, which is a deeper +// architectural change. +#[ignore] #[tokio::test] async fn test_upload_connection_die() { init(); @@ -181,7 +202,7 @@ async fn test_ws_server_ends_conn() { ws_stream.close(None).await.unwrap(); let msg = ws_stream.next().await.unwrap().unwrap(); // assert echo - assert_eq!("test", msg.into_text().unwrap()); + assert_eq!(msg.into_text().unwrap(), "test"); let msg = ws_stream.next().await.unwrap().unwrap(); // assert graceful close assert!(matches!(msg, Message::Close(None))); @@ -363,22 +384,22 @@ async fn test_upgrade_body_after_101() { #[tokio::test] async fn test_download_timeout() { init(); - use hyper::body::HttpBody; use tokio::time::sleep; - let client = hyper::Client::new(); - let uri: hyper::Uri = "http://127.0.0.1:6147/download_large/".parse().unwrap(); - let req = hyper::Request::builder() + let client = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()) + .build_http::>(); + let uri: http::Uri = "http://127.0.0.1:6147/download_large/".parse().unwrap(); + let req = http::Request::builder() .uri(uri) .header("x-write-timeout", "1") - .body(hyper::Body::empty()) + .body(http_body_util::Empty::::new()) .unwrap(); let mut res = client.request(req).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); let mut err = false; sleep(Duration::from_secs(2)).await; - while let Some(chunk) = res.body_mut().data().await { + while let Some(chunk) = res.body_mut().frame().await { if chunk.is_err() { err = true; } @@ -389,31 +410,217 @@ async fn test_download_timeout() { #[tokio::test] async fn test_download_timeout_min_rate() { init(); - use hyper::body::HttpBody; use tokio::time::sleep; - let client = hyper::Client::new(); - let uri: hyper::Uri = "http://127.0.0.1:6147/download/".parse().unwrap(); - let req = hyper::Request::builder() + let client = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()) + .build_http::>(); + let uri: http::Uri = "http://127.0.0.1:6147/download/".parse().unwrap(); + let req = http::Request::builder() .uri(uri) .header("x-write-timeout", "1") .header("x-min-rate", "10000") - .body(hyper::Body::empty()) + .body(http_body_util::Empty::::new()) .unwrap(); let mut res = client.request(req).await.unwrap(); assert_eq!(res.status(), StatusCode::OK); let mut err = false; sleep(Duration::from_secs(2)).await; - while let Some(chunk) = res.body_mut().data().await { + while let Some(chunk) = res.body_mut().frame().await { if chunk.is_err() { err = true; } } - // no error as write timeout is overridden by min rate assert!(!err); } +// When an H2 origin sends all Content-Length bytes in a DATA frame but times out +// before sending END_STREAM, a subsequent downstream h1 request may be blocked +#[tokio::test] +async fn test_h2_upstream_no_end_stream_read_timeout() { + init(); + + // Spawn a custom H2 origin: + // Request 1: sends Content-Length body WITHOUT END_STREAM, then goes quiet + // Request 2+: responds instantly with body + END_STREAM + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin_port = listener.local_addr().unwrap().port(); + + tokio::spawn(async move { + loop { + let Ok((tcp, _addr)) = listener.accept().await else { + break; + }; + tokio::spawn(async move { + let mut conn = match h2::server::handshake(tcp).await { + Ok(c) => c, + Err(e) => { + eprintln!("h2 handshake error: {e}"); + return; + } + }; + + let request_count = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)); + while let Some(result) = conn.accept().await { + let (request, mut respond) = match result { + Ok(r) => r, + Err(e) => { + eprintln!("h2 accept error: {e}"); + return; + } + }; + let _ = request; + let count = request_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1; + + // Spawn handler so conn.accept() keeps driving the h2 connection + tokio::spawn(async move { + let resp = http::Response::builder() + .status(200) + .header(http::header::CONTENT_LENGTH, "11") + .body(()) + .unwrap(); + + if count == 1 { + // Request 1: send body WITHOUT end_of_stream, then go quiet. + let mut send_stream = respond.send_response(resp, false).unwrap(); + send_stream + .send_data(bytes::Bytes::from("hello world"), false) + .unwrap(); + // Hold the stream open — simulates an origin that sent all CL + // bytes but hasn't closed the stream. + tokio::time::sleep(Duration::from_secs(30)).await; + } else { + // Request 2+: respond instantly with body + END_STREAM + let mut send_stream = respond.send_response(resp, false).unwrap(); + send_stream + .send_data(bytes::Bytes::from("hello world"), true) + .unwrap(); + } + }); + } + }); + } + }); + + // The listener was bound before the spawn (line 426), so the kernel + // is already accepting connections into the backlog. No readiness + // wait needed. + let client = reqwest::Client::new(); + let url = "http://127.0.0.1:6147/test"; + + let resp1 = client + .get(url) + .header("x-port", origin_port.to_string()) + .header("x-h2", "true") + .header("x-read-timeout-ms", "4000") + .send() + .await + .unwrap(); + assert_eq!(resp1.status(), StatusCode::OK); + assert_eq!(resp1.text().await.unwrap(), "hello world"); + + // Request 2: reqwest reuses the H1 connection + // but if blocked / stalled, the proxy can't read this request + // until read_timeout fires (~4s) + let start = Instant::now(); + let resp2 = timeout( + Duration::from_secs(10), + client + .get(url) + .header("x-port", origin_port.to_string()) + .header("x-h2", "true") + .send(), + ) + .await; + let elapsed = start.elapsed(); + + match resp2 { + Ok(Ok(resp)) => { + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(resp.text().await.unwrap(), "hello world"); + assert!( + elapsed < Duration::from_secs(2), + "Second request on reused H1 connection took {elapsed:?}, \ + expected < 2s (may be blocked on H2 upstream to end stream)" + ); + } + Ok(Err(e)) => { + panic!("Second request failed: {e}"); + } + Err(_) => { + panic!("Second request timed out after 10s."); + } + } +} + +/// Mock origin that sends 100 Continue then a final response for any request. +/// Returns the port the server is listening on. +async fn mock_100_continue_server() -> u16 { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + + tokio::spawn(async move { + if let Ok((mut stream, _addr)) = listener.accept().await { + // Read the request (just drain it) + let mut buf = [0u8; 4096]; + let _ = stream.read(&mut buf).await.unwrap(); + + // Send 100 Continue + stream + .write_all(b"HTTP/1.1 100 Continue\r\n\r\n") + .await + .unwrap(); + // Small delay so the client reads the 100 separately + tokio::time::sleep(Duration::from_millis(100)).await; + + // Send final 200 OK with Content-Length (but no body for HEAD) + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 42\r\n\r\n") + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + } + }); + + port +} + +#[tokio::test] +async fn test_head_with_100_continue() { + init(); + + let port = mock_100_continue_server().await; + + let mut stream = TcpStream::connect("127.0.0.1:6147").await.unwrap(); + stream + .write_all( + format!("HEAD / HTTP/1.1\r\nHost: localhost\r\nx-port: {port}\r\n\r\n").as_bytes(), + ) + .await + .unwrap(); + + // Read through any 1xx until we get the final (non-1xx) response + let result = timeout(Duration::from_secs(5), async { + let mut resp; + let mut body; + loop { + (resp, body) = read_response_header(&mut stream).await; + if resp.status.as_u16() >= 200 { + return (resp, body); + } + } + }) + .await + .expect("should not time out waiting for final response"); + + let (resp, body) = result; + assert_eq!(resp.status.as_u16(), 200); + // HEAD responses have no body even with Content-Length + assert!(body.is_empty(), "HEAD response should have no body"); +} + mod test_cache { use super::*; use std::str::FromStr; @@ -1193,35 +1400,37 @@ mod test_cache { // set up a one-off mock server // (warp / hyper don't have custom 1xx sending capabilities yet) - async fn mock_1xx_server(port: u16, cc_header: &str) { - use tokio::io::AsyncWriteExt; - - let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{}", port)) - .await - .unwrap(); - if let Ok((mut stream, _addr)) = listener.accept().await { - stream.write_all(b"HTTP/1.1 103 Early Hints\r\nLink: ; rel=preconnect\r\n\r\n").await.unwrap(); - // wait a bit so that the client can read - sleep(Duration::from_millis(100)).await; - stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: 5\r\nCache-Control: {}\r\n\r\nhello", cc_header).as_bytes()).await.unwrap(); - sleep(Duration::from_millis(100)).await; - } + // One-shot mock server that sends a 103 Early Hints then a final 200. + // Binds to port 0 (OS-assigned) and returns the actual port via a + // oneshot channel once the listener is ready. + fn spawn_mock_1xx_server(cc_header: &'static str) -> tokio::sync::oneshot::Receiver { + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + use tokio::io::AsyncWriteExt; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let _ = tx.send(port); // signal: port is bound + if let Ok((mut stream, _addr)) = listener.accept().await { + stream.write_all(b"HTTP/1.1 103 Early Hints\r\nLink: ; rel=preconnect\r\n\r\n").await.unwrap(); + sleep(Duration::from_millis(100)).await; + stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: 5\r\nCache-Control: {}\r\n\r\nhello", cc_header).as_bytes()).await.unwrap(); + sleep(Duration::from_millis(100)).await; + } + }); + rx } init(); let url = "http://127.0.0.1:6148/unique/test_1xx_caching"; - tokio::spawn(async { - mock_1xx_server(6151, "max-age=5").await; - }); - // wait for server to start - sleep(Duration::from_millis(100)).await; + let port = spawn_mock_1xx_server("max-age=5").await.unwrap(); let client = reqwest::Client::new(); let res = client .get(url) - .header("x-port", "6151") + .header("x-port", port.to_string()) .send() .await .unwrap(); @@ -1230,9 +1439,10 @@ mod test_cache { assert_eq!(headers["x-cache-status"], "miss"); assert_eq!(res.text().await.unwrap(), "hello"); + // Second request to the same URL should be a cache hit (no server needed) let res = client .get(url) - .header("x-port", "6151") + .header("x-port", port.to_string()) .send() .await .unwrap(); @@ -1244,15 +1454,11 @@ mod test_cache { // 1xx shouldn't interfere with bypass let url = "http://127.0.0.1:6148/unique/test_1xx_bypass"; - tokio::spawn(async { - mock_1xx_server(6152, "private, no-store").await; - }); - // wait for server to start - sleep(Duration::from_millis(100)).await; + let port = spawn_mock_1xx_server("private, no-store").await.unwrap(); let res = client .get(url) - .header("x-port", "6152") + .header("x-port", port.to_string()) .send() .await .unwrap(); @@ -1262,16 +1468,11 @@ mod test_cache { assert_eq!(res.text().await.unwrap(), "hello"); // restart the one-off server - still uncacheable - sleep(Duration::from_millis(100)).await; - tokio::spawn(async { - mock_1xx_server(6152, "private, no-store").await; - }); - // wait for server to start - sleep(Duration::from_millis(100)).await; + let port = spawn_mock_1xx_server("private, no-store").await.unwrap(); let res = client .get(url) - .header("x-port", "6152") + .header("x-port", port.to_string()) .send() .await .unwrap(); @@ -2732,6 +2933,154 @@ mod test_cache { assert_eq!(res.text().await.unwrap(), "hello world"); } + #[tokio::test] + #[ignore = "flaky in CI due to timing/resource contention"] + async fn test_caching_when_downstream_stalls() { + use std::net::ToSocketAddrs; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpStream; + + init(); + let url = "http://127.0.0.1:6148/unique/test_caching_when_downstream_stalls/download/"; + + // Connection 1: read 10KiB then stall, holding the cache lock while + // the proxy populates cache from upstream. + let slow_task = tokio::spawn(async move { + let addr = "127.0.0.1:6148".to_socket_addrs().unwrap().next().unwrap(); + let mut stream = TcpStream::connect(&addr).await.unwrap(); + + let request = concat!( + "GET /unique/test_caching_when_downstream_stalls/download/ HTTP/1.1\r\n", + "Host: 127.0.0.1:6148\r\n", + "x-lock: true\r\n", + "x-set-cache-control: public, max-age=60\r\n", + "\r\n", + ); + stream.write_all(request.as_bytes()).await.unwrap(); + + let mut buf = [0; 10 * 1024]; + let mut b = &mut buf[..]; + while !b.is_empty() { + let n = stream.read(b).await.unwrap(); + b = &mut b[n..] + } + + // Hold the stalled connection open long enough + sleep(Duration::from_secs(10)).await; + }); + + // Give connection 1 time to acquire the cache lock + sleep(Duration::from_secs(1)).await; + + // Connection 2: should get a cache hit once the proxy finishes + // populating cache from upstream (independent of stall). + let start = tokio::time::Instant::now(); + let res = reqwest::Client::new() + .get(url) + .header("x-lock", "true") + .header("x-set-cache-control", "public, max-age=60") + .timeout(Duration::from_secs(8)) + .send() + .await + .unwrap(); + + assert_eq!(res.status(), StatusCode::OK); + let headers = res.headers(); + assert_eq!(headers["x-cache-status"], "hit"); + + // If the cache was populated fast enough (before connection 2 arrived), + // there is no lock contention and x-cache-lock-time-ms is absent. + // If there was contention, the wait should be short. + if let Some(lock_ms) = headers.get("x-cache-lock-time-ms") { + let ms: u64 = lock_ms.to_str().unwrap().parse().unwrap(); + assert!( + ms < 2000, + "lock wait {ms}ms should be well under the 2s timeout" + ); + } + + assert_eq!( + res.text().await.unwrap(), + String::from("A").repeat(4 * 1024 * 1024) + ); + + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_secs(5), + "second request took {elapsed:?}, should be fast" + ); + + // Don't wait for the slow connection + slow_task.abort(); + } + + // Same as test_caching_when_downstream_stalls but the proxy connects + // to the origin over H2 (via the x-h2 header). + // + #[tokio::test] + #[ignore = "flaky in CI due to timing/resource contention"] + async fn test_caching_h2_upstream_when_downstream_stalls() { + use std::net::ToSocketAddrs; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpStream; + + init(); + let url = "http://127.0.0.1:6148/unique/test_caching_h2_upstream_when_downstream_stalls/download/"; + + let slow_task = tokio::spawn(async move { + let addr = "127.0.0.1:6148".to_socket_addrs().unwrap().next().unwrap(); + let mut stream = TcpStream::connect(&addr).await.unwrap(); + + let request = concat!( + "GET /unique/test_caching_h2_upstream_when_downstream_stalls/download/ HTTP/1.1\r\n", + "Host: 127.0.0.1:6148\r\n", + "x-h2: true\r\n", + "x-lock: true\r\n", + "x-set-cache-control: public, max-age=60\r\n", + "\r\n", + ); + stream.write_all(request.as_bytes()).await.unwrap(); + + let mut buf = [0; 10 * 1024]; + let mut b = &mut buf[..]; + while !b.is_empty() { + let n = stream.read(b).await.unwrap(); + b = &mut b[n..] + } + + sleep(Duration::from_secs(10)).await; + }); + + sleep(Duration::from_secs(1)).await; + + let start = tokio::time::Instant::now(); + let res = reqwest::Client::new() + .get(url) + .header("x-h2", "true") + .header("x-lock", "true") + .header("x-set-cache-control", "public, max-age=60") + .timeout(Duration::from_secs(8)) + .send() + .await + .unwrap(); + + assert_eq!(res.status(), StatusCode::OK); + let headers = res.headers(); + assert_eq!(headers["x-cache-status"], "hit"); + assert_eq!( + res.text().await.unwrap(), + String::from("A").repeat(4 * 1024 * 1024) + ); + + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_secs(5), + "second request took {elapsed:?}, should be fast (upstream-speed-bound)" + ); + + slow_task.abort(); + } + async fn send_vary_req_with_headers_with_dups( url: &str, vary_field: &str, @@ -3355,4 +3704,145 @@ mod test_cache { assert_eq!(headers["x-cache-status"], "hit"); assert_eq!(res.text().await.unwrap(), "hello world"); } + + // Ignored until H2 downstream gets the proxy task API + // (write_response_tasks blocks on flow control today). + // multi_thread needed for h2 connection driver tasks. + #[tokio::test(flavor = "multi_thread")] + #[ignore] + async fn test_cache_h2_downstream_stalls() { + init(); + + use h2::client; + use http::Request; + use tokio::net::TcpStream; + use tokio::time::{timeout, Duration}; + + // Step 1: Connection 1 - Open h2 connection to h2c cache proxy (port 6154) and STALL + let tcp1 = TcpStream::connect("127.0.0.1:6154").await.unwrap(); + let (mut h2_client1, h2_conn1) = client::handshake(tcp1).await.unwrap(); + + tokio::spawn(async move { + if let Err(e) = h2_conn1.await { + eprintln!("H2 connection 1 error: {:?}", e); + } + }); + + // Request the cached resource on connection 1 + let request1 = Request::builder() + .uri("http://127.0.0.1/unique/test_h2_stall/download/") + .body(()) + .unwrap(); + + let (response1, _) = h2_client1.send_request(request1, true).unwrap(); + let response1 = response1.await.unwrap(); + assert_eq!(response1.status(), 200); + assert_eq!(response1.headers()["x-cache-status"], "miss"); + + let mut body1 = response1.into_body(); + + // Read first chunk but don't release flow control to stall connection 1 + let first_chunk = body1.data().await.unwrap().unwrap(); + assert!(!first_chunk.is_empty()); + + // Connection 2 - While conn 1 is stalled, try to get the same cached resource + let tcp2 = TcpStream::connect("127.0.0.1:6154").await.unwrap(); + let (mut h2_client2, h2_conn2) = client::handshake(tcp2).await.unwrap(); + + tokio::spawn(async move { + if let Err(e) = h2_conn2.await { + eprintln!("H2 connection 2 error: {:?}", e); + } + }); + + let request2 = Request::builder() + .uri("http://127.0.0.1/unique/test_h2_stall/download/") + .body(()) + .unwrap(); + + let (response2, _) = h2_client2.send_request(request2, true).unwrap(); + + // Try to read, proxy should not be blocked + let response2 = match timeout(Duration::from_secs(5), response2).await { + Ok(Ok(resp)) => resp, + Ok(Err(e)) => panic!("Connection 2 failed: {:?}", e), + Err(_) => panic!("Connection 2 timed out - proxy blocked without proxy task API!"), + }; + + assert_eq!(response2.status(), 200); + assert_eq!(response2.headers()["x-cache-status"], "hit"); + + // Read full response from connection 2 + let mut body2 = response2.into_body(); + let mut received2 = Vec::new(); + while let Some(Ok(chunk)) = timeout(Duration::from_secs(5), body2.data()) + .await + .expect("should not time out waiting for data") + { + let len = chunk.len(); + received2.extend_from_slice(&chunk); + body2.flow_control().release_capacity(len).unwrap(); + } + + assert_eq!( + received2.len(), + 4 * 1024 * 1024, + "Connection 2 should receive full cached response" + ); + + // Clean up: unstall connection 1 + body1 + .flow_control() + .release_capacity(first_chunk.len()) + .unwrap(); + } + + // Test cache population from H2 upstream origin with H1 downstream. + #[tokio::test] + async fn test_cache_upstream_h2_downstream_h1() { + init(); + + let test_url = "http://127.0.0.1:6148/unique/test_h2_upstream/download/"; + + // Step 1: Populate cache from H2 origin (cache miss) + let client = reqwest::Client::new(); + let res = client + .get(test_url) + .header("x-h2", "true") + .header("x-lock", "true") + .header("x-set-cache-control", "public, max-age=60") + .send() + .await + .unwrap(); + + assert_eq!(res.status(), 200); + assert_eq!(res.headers()["x-cache-status"], "miss"); + assert_eq!(res.headers()["origin-http2"], "h2c"); + + let body = res.bytes().await.unwrap(); + assert_eq!( + body.len(), + 4 * 1024 * 1024, + "Should receive full 4MB response" + ); + + // Step 2: Request again and verify cache hit + let res = client + .get(test_url) + .header("x-h2", "true") + .header("x-set-cache-control", "public, max-age=60") + .send() + .await + .unwrap(); + + assert_eq!(res.status(), 200); + assert_eq!(res.headers()["x-cache-status"], "hit"); + + let body = res.bytes().await.unwrap(); + assert_eq!( + body.len(), + 4 * 1024 * 1024, + "Should receive full 4MB from cache" + ); + } } diff --git a/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf b/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf index f19c974cb..969695eb9 100644 --- a/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf +++ b/pingora-proxy/tests/utils/conf/origin/conf/nginx.conf @@ -311,7 +311,6 @@ http { location /download/ { content_by_lua_block { - ngx.req.read_body() local body = string.rep("A", 4194304) ngx.header["Content-Length"] = #body ngx.print(body) diff --git a/pingora-proxy/tests/utils/mock_origin.rs b/pingora-proxy/tests/utils/mock_origin.rs index 74840e195..fa5a327fa 100644 --- a/pingora-proxy/tests/utils/mock_origin.rs +++ b/pingora-proxy/tests/utils/mock_origin.rs @@ -59,7 +59,22 @@ fn init() -> bool { .output() .unwrap(); }); - // wait until the server is up - thread::sleep(time::Duration::from_secs(2)); - true + // Wait until openresty is accepting connections, then give it a moment + // to finish worker initialization. + let deadline = time::Instant::now() + time::Duration::from_secs(10); + while time::Instant::now() < deadline { + if std::net::TcpStream::connect_timeout( + &"127.0.0.1:8000".parse().unwrap(), + time::Duration::from_millis(100), + ) + .is_ok() + { + // Port is listening; allow a brief window for workers to finish + // initializing before tests start sending real requests. + thread::sleep(time::Duration::from_millis(500)); + return true; + } + thread::sleep(time::Duration::from_millis(50)); + } + panic!("mock origin (openresty) failed to start within 10s"); } diff --git a/pingora-proxy/tests/utils/server_utils.rs b/pingora-proxy/tests/utils/server_utils.rs index 0df71336d..b2ad2bbc5 100644 --- a/pingora-proxy/tests/utils/server_utils.rs +++ b/pingora-proxy/tests/utils/server_utils.rs @@ -43,9 +43,12 @@ use pingora_core::upstreams::peer::HttpPeer; use pingora_core::utils::tls::CertKey; use pingora_error::{Error, ErrorSource, ErrorType::*, Result}; use pingora_http::{RequestHeader, ResponseHeader}; -use pingora_proxy::{FailToProxy, ProxyHttp, Session}; +use pingora_proxy::{FailToProxy, ProxyHttp, ProxyWarnLogContext, Session}; use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; use std::thread; use std::time::{Duration, SystemTime}; @@ -241,6 +244,16 @@ impl ProxyHttp for ExampleProxyHttps { pub struct ExampleProxyHttp {} +static SUPPRESS_PROXY_WARN_LOG_CALLS: AtomicUsize = AtomicUsize::new(0); + +pub fn reset_suppress_proxy_warn_log_calls() { + SUPPRESS_PROXY_WARN_LOG_CALLS.store(0, Ordering::Relaxed); +} + +pub fn suppress_proxy_warn_log_calls() -> usize { + SUPPRESS_PROXY_WARN_LOG_CALLS.load(Ordering::Relaxed) +} + #[async_trait] impl ProxyHttp for ExampleProxyHttp { type CTX = CTX; @@ -253,8 +266,19 @@ impl ProxyHttp for ExampleProxyHttp { session: &mut Session, _ctx: &mut Self::CTX, ) -> Result<()> { - let req = session.req_header(); - let downstream_compression = req.headers.get("x-downstream-compression").is_some(); + let proxy_tasks_enabled = session + .req_header() + .headers + .get("x-proxy-tasks-enabled") + .is_some(); + if proxy_tasks_enabled { + session.downstream_session.set_proxy_tasks_enabled(true); + } + let downstream_compression = session + .req_header() + .headers + .get("x-downstream-compression") + .is_some(); if downstream_compression { session .downstream_modules_ctx @@ -346,6 +370,15 @@ impl ProxyHttp for ExampleProxyHttp { peer.options.set_http_version(2, 2); } + if let Some(ms) = req + .headers + .get("x-read-timeout-ms") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + { + peer.options.read_timeout = Some(std::time::Duration::from_millis(ms)); + } + Ok(peer) } @@ -361,6 +394,23 @@ impl ProxyHttp for ExampleProxyHttp { ) -> Result<()> { connected_to_upstream_common(reused, digest, ctx) } + + fn suppress_proxy_warn_log( + &self, + session: &Session, + _ctx: &Self::CTX, + _error: &Error, + context: ProxyWarnLogContext, + ) -> bool { + if session.get_header_bytes("x-test-suppress-proxy-warn-log") == b"true" + && context == ProxyWarnLogContext::UpstreamRetry + { + SUPPRESS_PROXY_WARN_LOG_CALLS.fetch_add(1, Ordering::Relaxed); + true + } else { + false + } + } } static CACHE_BACKEND: Lazy = Lazy::new(MemCache::new); @@ -649,6 +699,12 @@ impl ProxyHttp for ExampleProxyCache { upstream_response.remove_header(&CONTENT_LENGTH); upstream_response.remove_header(&TRANSFER_ENCODING); } + // Allow tests to inject Cache-Control into the upstream response + if let Some(cc) = session.req_header().headers.get("x-set-cache-control") { + upstream_response + .insert_header(http::header::CACHE_CONTROL, cc) + .unwrap(); + } Ok(()) } @@ -814,6 +870,15 @@ fn test_main() { pingora_proxy::http_proxy_service(&my_server.configuration, ExampleProxyCache {}); proxy_service_cache.add_tcp("0.0.0.0:6148"); + // H2C-enabled cache proxy on port 6154 + let mut proxy_service_cache_h2c = + pingora_proxy::http_proxy_service(&my_server.configuration, ExampleProxyCache {}); + let cache_h2c_logic = proxy_service_cache_h2c.app_logic_mut().unwrap(); + let mut cache_h2c_options = HttpServerOptions::default(); + cache_h2c_options.h2c = true; + cache_h2c_logic.server_options = Some(cache_h2c_options); + proxy_service_cache_h2c.add_tcp("0.0.0.0:6154"); + #[cfg(feature = "any_tls")] { let cert_path = format!("{}/tests/keys/server.crt", env!("CARGO_MANIFEST_DIR")); @@ -830,6 +895,7 @@ fn test_main() { Box::new(proxy_service_http), Box::new(proxy_service_http_connect), Box::new(proxy_service_cache), + Box::new(proxy_service_cache_h2c), ]; if let Some(proxy_service_https) = proxy_service_https_opt { @@ -864,16 +930,28 @@ pub struct PskTlsServer { #[cfg(feature = "s2n")] impl PskTlsServer { pub fn start() -> Self { - let server_handle = thread::spawn(|| { + use std::sync::mpsc; + use std::time::Duration; + + // Use a channel to wait for the server to bind its port. + // A TCP probe can't be used here because the TLS acceptor would + // try to handshake the probe connection, fail, and panic. + let (tx, rx) = mpsc::channel(); + let server_handle = thread::spawn(move || { let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(Self::run_server()); + rt.block_on(Self::run_server(tx)); }); + + // Wait up to 10s for the server to signal it has bound the port. + rx.recv_timeout(Duration::from_secs(10)) + .expect("PSK TLS server failed to start within 10s"); + PskTlsServer { handle: server_handle, } } - async fn run_server() { + async fn run_server(ready_tx: std::sync::mpsc::Sender<()>) { use pingora_core::{protocols::tls::S2NConnectionBuilder, tls::TlsAcceptor}; use pingora_core::{ protocols::tls::{Psk, PskConfig, PskType}, @@ -890,6 +968,8 @@ impl PskTlsServer { let addr: std::net::SocketAddr = "127.0.0.1:6151".parse().unwrap(); let listener = TcpListener::bind(addr).await.unwrap(); + let _ = ready_tx.send(()); // signal: port is bound + let mut config_builder = Config::builder(); unsafe { config_builder.disable_x509_verification(); @@ -906,12 +986,20 @@ impl PskTlsServer { let acceptor = TlsAcceptor::new(connection_builder); loop { - use tokio::{io::AsyncWriteExt, net::tcp}; + use tokio::io::AsyncWriteExt; let (tcp_stream, _) = listener.accept().await.unwrap(); - let mut stream = acceptor.clone().accept(tcp_stream).await.unwrap(); + // Don't panic on handshake failure — a stale connection or probe + // shouldn't take down the server for subsequent real connections. + let mut stream = match acceptor.clone().accept(tcp_stream).await { + Ok(s) => s, + Err(e) => { + log::warn!("PSK TLS server: handshake failed: {e}"); + continue; + } + }; let response = b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\nhello"; - stream.write(response).await.unwrap(); - stream.shutdown().await; + let _ = stream.write_all(response).await; + let _ = stream.shutdown().await; } } } diff --git a/pingora-runtime/Cargo.toml b/pingora-runtime/Cargo.toml index 5de4f26b1..de25fff05 100644 --- a/pingora-runtime/Cargo.toml +++ b/pingora-runtime/Cargo.toml @@ -16,11 +16,18 @@ Multithreaded Tokio runtime with the option of disabling work stealing. name = "pingora_runtime" path = "src/lib.rs" +[lints] +workspace = true + [dependencies] rand = "0.8" tokio = { workspace = true, features = ["rt-multi-thread", "sync", "time"] } once_cell = { workspace = true } +serde = { version = "1.0", features = ["derive"] } thread_local = "1" +log = { workspace = true } +dial9-tokio-telemetry = { version = "0.3.10", optional = true } +aws-sdk-s3 = { version = "1", optional = true } [dev-dependencies] tokio = { workspace = true, features = ["io-util", "net"] } @@ -28,3 +35,13 @@ tokio = { workspace = true, features = ["io-util", "net"] } [[bench]] name = "hello" harness = false + +[features] +default = [] +dial9 = ["dep:dial9-tokio-telemetry"] +dial9-worker-s3 = [ + "dial9", + "dep:aws-sdk-s3", + "dial9-tokio-telemetry/worker-s3", +] +dial9-cpu-profiling = ["dial9", "dial9-tokio-telemetry/cpu-profiling"] diff --git a/pingora-runtime/src/lib.rs b/pingora-runtime/src/lib.rs index a0468f4fc..7d617198f 100644 --- a/pingora-runtime/src/lib.rs +++ b/pingora-runtime/src/lib.rs @@ -25,6 +25,9 @@ use once_cell::sync::{Lazy, OnceCell}; use rand::Rng; +use serde::{Deserialize, Serialize}; +#[cfg(feature = "dial9")] +use std::path::PathBuf; use std::sync::Arc; use std::thread::JoinHandle; use std::time::Duration; @@ -32,32 +35,476 @@ use thread_local::ThreadLocal; use tokio::runtime::{Builder, Handle}; use tokio::sync::oneshot::{channel, Sender}; +/// Default maximum size of a dial9 trace segment file. +#[cfg(feature = "dial9")] +pub const DEFAULT_DIAL9_MAX_FILE_SIZE: u64 = 100 * 1024 * 1024; +/// Default maximum bytes retained locally by dial9. +#[cfg(feature = "dial9")] +pub const DEFAULT_DIAL9_MAX_TOTAL_SIZE: u64 = 512 * 1024 * 1024; + +/// Configuration options for the blocking thread pool used by the runtime. +/// +/// These options control the behavior of the blocking thread pool that handles +/// [`tokio::task::spawn_blocking`] tasks. +#[derive(Debug, Clone, Default)] +pub struct BlockingPoolOpts { + /// The maximum number of threads in the blocking thread pool. + /// + /// When not set, the tokio default (512) is used. + pub max_threads: Option, + /// The duration that idle blocking threads are kept alive before being shut down. + /// + /// When not set, the tokio default (10 seconds) is used. + pub thread_keep_alive: Option, +} + +/// Configuration options for runtime metrics collection. +#[derive(Debug, Clone, Default)] +pub struct RuntimeMetricsOpts { + /// Enable Tokio's poll-time histogram on the runtime. + /// + /// This must be configured before the runtime is built. Enabling it adds + /// two timestamp reads to every task poll. + pub poll_time_histogram: bool, + /// Histogram bucket scale for Tokio's poll-time histogram. + pub poll_time_histogram_scale: Option, + /// Width of the first histogram bucket. + pub poll_time_histogram_resolution: Option, + /// Number of histogram buckets. Memory usage scales with runtimes × workers × buckets. + pub poll_time_histogram_buckets: Option, +} + +/// Configuration options for a Tokio runtime. +#[derive(Debug, Clone, Default)] +pub struct RuntimeOpts { + /// Options for runtime metrics collection. + pub metrics: RuntimeMetricsOpts, + /// Enable Tokio's experimental alternative timer. + /// + /// This requires building with `--cfg tokio_unstable` and only applies to + /// Tokio's multi-threaded runtime. + pub enable_alt_timer: bool, + /// Options for dial9 Tokio telemetry. + #[cfg(feature = "dial9")] + pub dial9: Option, +} + +/// Configuration options for dial9 Tokio telemetry. +#[cfg(feature = "dial9")] +#[derive(Debug, Clone)] +pub struct Dial9RuntimeOpts { + /// Trace output path after server configuration defaults are applied. + pub trace_path: PathBuf, + /// Rotate trace segments after this many bytes. + pub max_file_size: u64, + /// Maximum bytes retained on local disk. + pub max_total_size: u64, + /// Wall-clock trace rotation period. + pub rotation_period: Option, + /// Enable dial9 task spawn/terminate tracking. + pub task_tracking: bool, + /// Upload sealed trace segments to S3-compatible storage. + #[cfg(feature = "dial9-worker-s3")] + pub s3_upload: Option, +} + +#[cfg(feature = "dial9")] +impl Dial9RuntimeOpts { + /// Create dial9 runtime options using Pingora's dial9 defaults. + pub fn new(trace_path: impl Into) -> Self { + Self { + trace_path: trace_path.into(), + max_file_size: DEFAULT_DIAL9_MAX_FILE_SIZE, + max_total_size: DEFAULT_DIAL9_MAX_TOTAL_SIZE, + rotation_period: None, + task_tracking: true, + #[cfg(feature = "dial9-worker-s3")] + s3_upload: None, + } + } + + /// Set the maximum size of each trace segment file. + pub fn with_max_file_size(mut self, max_file_size: u64) -> Self { + self.max_file_size = max_file_size; + self + } + + /// Set the maximum bytes retained on local disk. + pub fn with_max_total_size(mut self, max_total_size: u64) -> Self { + self.max_total_size = max_total_size; + self + } + + /// Set the wall-clock trace rotation period. + pub fn with_rotation_period(mut self, rotation_period: Duration) -> Self { + self.rotation_period = Some(rotation_period); + self + } + + /// Enable or disable dial9 task spawn/terminate tracking. + pub fn with_task_tracking(mut self, task_tracking: bool) -> Self { + self.task_tracking = task_tracking; + self + } + + /// Set S3-compatible upload options for sealed trace segments. + #[cfg(feature = "dial9-worker-s3")] + pub fn with_s3_upload(mut self, s3_upload: Dial9S3UploadOpts) -> Self { + self.s3_upload = Some(s3_upload); + self + } +} + +/// Configuration options for dial9 S3-compatible trace uploads. +#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))] +#[derive(Debug, Clone)] +pub struct Dial9S3UploadOpts { + /// S3 bucket that receives sealed trace segments. + pub bucket: String, + /// Service name included in uploaded object keys. + pub service_name: String, + /// Optional key prefix. + pub prefix: Option, + /// Optional region override. + pub region: Option, + /// Optional pre-built S3 client for custom credentials or endpoints. + pub client: Option, +} + +#[cfg(all(feature = "dial9", feature = "dial9-worker-s3"))] +impl Dial9S3UploadOpts { + /// Create S3-compatible upload options. + pub fn new(bucket: impl Into, service_name: impl Into) -> Self { + Self { + bucket: bucket.into(), + service_name: service_name.into(), + prefix: None, + region: None, + client: None, + } + } + + /// Set the object key prefix. + pub fn with_prefix(mut self, prefix: impl Into) -> Self { + self.prefix = Some(prefix.into()); + self + } + + /// Set the AWS region override. + pub fn with_region(mut self, region: impl Into) -> Self { + self.region = Some(region.into()); + self + } + + /// Set a pre-built S3 client for custom credentials or endpoints. + pub fn with_client(mut self, client: aws_sdk_s3::Client) -> Self { + self.client = Some(client); + self + } +} + +/// Bucket scale for Tokio's poll-time histogram. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RuntimeMetricsPollTimeHistogramScale { + /// Equal-width buckets. + Linear, + /// Buckets double in width at each step. + Log, +} + /// Pingora async multi-threaded runtime /// /// The `Steal` flavor is effectively tokio multi-threaded runtime. /// /// The `NoSteal` flavor is backed by multiple tokio single-threaded runtime. pub enum Runtime { - Steal(tokio::runtime::Runtime), + Steal { + runtime: tokio::runtime::Runtime, + #[cfg(feature = "dial9")] + dial9_guard: Option, + }, NoSteal(NoStealRuntime), } +/// Apply [`BlockingPoolOpts`] to a tokio [`Builder`]. +fn apply_blocking_opts(builder: &mut Builder, opts: &BlockingPoolOpts) { + if let Some(max) = opts.max_threads { + builder.max_blocking_threads(max); + } + if let Some(ttl) = opts.thread_keep_alive { + builder.thread_keep_alive(ttl); + } +} + +/// Apply [`RuntimeMetricsOpts`] to a tokio [`Builder`]. +// The replacement `metrics_poll_time_histogram_configuration` API is not used +// here so this crate can continue to compile against older Tokio 1.x versions +// selected by downstream applications while still honoring these knobs in +// tokio-unstable builds. +#[allow(deprecated)] +fn apply_metrics_opts(builder: &mut Builder, opts: &RuntimeMetricsOpts) { + #[cfg(tokio_unstable)] + if opts.poll_time_histogram { + builder.enable_metrics_poll_time_histogram(); + + if let Some(scale) = opts.poll_time_histogram_scale { + builder.metrics_poll_count_histogram_scale(match scale { + RuntimeMetricsPollTimeHistogramScale::Linear => { + tokio::runtime::HistogramScale::Linear + } + RuntimeMetricsPollTimeHistogramScale::Log => tokio::runtime::HistogramScale::Log, + }); + } + if let Some(resolution) = opts + .poll_time_histogram_resolution + .filter(|resolution| !resolution.is_zero()) + { + builder.metrics_poll_count_histogram_resolution(resolution); + } + if let Some(buckets) = opts + .poll_time_histogram_buckets + .filter(|buckets| *buckets > 0) + { + builder.metrics_poll_count_histogram_buckets(buckets); + } + } + + #[cfg(not(tokio_unstable))] + let _ = (builder, opts); +} + +/// Apply timer options from [`RuntimeOpts`] to a tokio [`Builder`]. +fn apply_timer_opts(builder: &mut Builder, opts: &RuntimeOpts) { + #[cfg(tokio_unstable)] + if opts.enable_alt_timer { + builder.enable_alt_timer(); + } + + #[cfg(not(tokio_unstable))] + let _ = (builder, opts); +} + +#[cfg(feature = "dial9")] +fn build_dial9_runtime( + builder: Builder, + runtime_name: &str, + opts: &Dial9RuntimeOpts, +) -> std::io::Result<( + tokio::runtime::Runtime, + dial9_tokio_telemetry::telemetry::TelemetryGuard, +)> { + use dial9_tokio_telemetry::telemetry::{RotatingWriter, TracedRuntime}; + use std::io::{Error, ErrorKind}; + + if opts.max_file_size == 0 { + return Err(Error::new( + ErrorKind::InvalidInput, + "dial9 max_file_size must be greater than zero", + )); + } + if opts.max_total_size == 0 { + return Err(Error::new( + ErrorKind::InvalidInput, + "dial9 max_total_size must be greater than zero", + )); + } + if opts.max_file_size > opts.max_total_size { + return Err(Error::new( + ErrorKind::InvalidInput, + "dial9 max_file_size must be less than or equal to max_total_size", + )); + } + + if let Some(parent) = opts.trace_path.parent() { + std::fs::create_dir_all(parent)?; + } + + let writer = RotatingWriter::builder() + .base_path(opts.trace_path.clone()) + .max_file_size(opts.max_file_size) + .max_total_size(opts.max_total_size) + .maybe_rotation_period(opts.rotation_period) + .build()?; + + let traced = TracedRuntime::builder() + .with_trace_path(opts.trace_path.clone()) + .with_runtime_name(runtime_name) + .with_task_tracking(opts.task_tracking); + + #[cfg(feature = "dial9-worker-s3")] + if let Some(s3_upload) = &opts.s3_upload { + if s3_upload.bucket.trim().is_empty() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "dial9 s3 bucket must not be empty", + )); + } + if s3_upload.service_name.trim().is_empty() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "dial9 s3 service_name must not be empty", + )); + } + let s3_config = dial9_tokio_telemetry::background_task::s3::S3Config::builder() + .bucket(s3_upload.bucket.clone()) + .service_name(s3_upload.service_name.clone()) + .maybe_prefix(s3_upload.prefix.clone()) + .maybe_region(s3_upload.region.clone()); + let traced = traced.with_s3_uploader(s3_config.build()); + if let Some(client) = s3_upload.client.clone() { + return traced + .with_s3_client(client) + .build_and_start(builder, writer); + } + return traced.build_and_start(builder, writer); + } + + traced.build_and_start(builder, writer) +} + +/// Builder for constructing a [`Runtime`]. +/// +/// # Example +/// +/// ``` +/// use pingora_runtime::{RuntimeBuilder, BlockingPoolOpts}; +/// use std::time::Duration; +/// +/// let rt = RuntimeBuilder::new(4, "my-service") +/// .blocking_pool_opts(BlockingPoolOpts { +/// max_threads: Some(64), +/// thread_keep_alive: Some(Duration::from_secs(30)), +/// }) +/// .build(); +/// ``` +pub struct RuntimeBuilder { + threads: usize, + name: String, + work_steal: bool, + blocking_pool_opts: BlockingPoolOpts, + runtime_opts: RuntimeOpts, +} + +impl RuntimeBuilder { + /// Create a new builder with the given number of worker threads and runtime name. + /// + /// Work stealing is enabled by default. + pub fn new(threads: usize, name: &str) -> Self { + Self { + threads, + name: name.to_string(), + work_steal: true, + blocking_pool_opts: BlockingPoolOpts::default(), + runtime_opts: RuntimeOpts::default(), + } + } + + /// Set whether work stealing is enabled. + /// + /// When `true` (the default), a tokio multi-thread runtime is used. + /// When `false`, a pool of single-threaded tokio runtimes is used instead. + pub fn work_steal(mut self, enabled: bool) -> Self { + self.work_steal = enabled; + self + } + + /// Set the [`BlockingPoolOpts`] for the runtime's blocking thread pool. + pub fn blocking_pool_opts(mut self, opts: BlockingPoolOpts) -> Self { + self.blocking_pool_opts = opts; + self + } + + /// Set the [`RuntimeMetricsOpts`] for the runtime. + pub fn metrics_opts(mut self, opts: RuntimeMetricsOpts) -> Self { + self.runtime_opts.metrics = opts; + self + } + + /// Set the [`RuntimeOpts`] for the runtime. + pub fn runtime_opts(mut self, opts: RuntimeOpts) -> Self { + self.runtime_opts = opts; + self + } + + /// Set whether Tokio's experimental alternative timer is enabled. + /// + /// This requires building with `--cfg tokio_unstable` and only applies to + /// work-stealing runtimes. + pub fn enable_alt_timer(mut self, enabled: bool) -> Self { + self.runtime_opts.enable_alt_timer = enabled; + self + } + + fn build_work_stealing_tokio_builder(&self) -> Builder { + let mut builder = Builder::new_multi_thread(); + builder + .enable_all() + .worker_threads(self.threads) + .thread_name(&self.name); + apply_blocking_opts(&mut builder, &self.blocking_pool_opts); + apply_metrics_opts(&mut builder, &self.runtime_opts.metrics); + apply_timer_opts(&mut builder, &self.runtime_opts); + builder + } + + /// Build the [`Runtime`]. + pub fn build(self) -> Runtime { + if self.work_steal { + let mut builder = self.build_work_stealing_tokio_builder(); + #[cfg(feature = "dial9")] + let dial9_guard = if let Some(dial9_opts) = &self.runtime_opts.dial9 { + let runtime_name = self.name.clone(); + match build_dial9_runtime(builder, &runtime_name, dial9_opts) { + Ok((runtime, guard)) => { + return Runtime::Steal { + runtime, + dial9_guard: Some(guard), + }; + } + Err(e) => { + log::warn!( + "failed to initialize dial9 runtime telemetry for {runtime_name}: {e}" + ); + builder = self.build_work_stealing_tokio_builder(); + None + } + } + } else { + None + }; + let runtime = builder + .build() + .expect("failed to build work-stealing Tokio runtime"); + Runtime::Steal { + runtime, + #[cfg(feature = "dial9")] + dial9_guard, + } + } else { + #[cfg(feature = "dial9")] + if self.runtime_opts.dial9.is_some() { + log::warn!("dial9 runtime telemetry is ignored when work stealing is disabled"); + } + Runtime::NoSteal(NoStealRuntime::new( + self.threads, + &self.name, + self.blocking_pool_opts, + self.runtime_opts, + )) + } + } +} + impl Runtime { /// Create a `Steal` flavor runtime. This just a regular tokio runtime pub fn new_steal(threads: usize, name: &str) -> Self { - Self::Steal( - Builder::new_multi_thread() - .enable_all() - .worker_threads(threads) - .thread_name(name) - .build() - .unwrap(), - ) + RuntimeBuilder::new(threads, name).build() } /// Create a `NoSteal` flavor runtime. This is backed by multiple tokio current-thread runtime pub fn new_no_steal(threads: usize, name: &str) -> Self { - Self::NoSteal(NoStealRuntime::new(threads, name)) + RuntimeBuilder::new(threads, name).work_steal(false).build() } /// Return the &[Handle] of the [Runtime]. @@ -67,7 +514,7 @@ impl Runtime { /// for each async task. pub fn get_handle(&self) -> &Handle { match self { - Self::Steal(r) => r.handle(), + Self::Steal { runtime, .. } => runtime.handle(), Self::NoSteal(r) => r.get_runtime(), } } @@ -76,7 +523,15 @@ impl Runtime { /// all runtimes exit. pub fn shutdown_timeout(self, timeout: Duration) { match self { - Self::Steal(r) => r.shutdown_timeout(timeout), + Self::Steal { + runtime, + #[cfg(feature = "dial9")] + dial9_guard, + } => { + #[cfg(feature = "dial9")] + drop(dial9_guard); + runtime.shutdown_timeout(timeout); + } Self::NoSteal(r) => r.shutdown_timeout(timeout), } } @@ -109,6 +564,8 @@ type Pools = Arc>>; pub struct NoStealRuntime { threads: usize, name: String, + blocking_opts: BlockingPoolOpts, + runtime_opts: RuntimeOpts, // Lazily init the runtimes so that they are created after pingora // daemonize itself. Otherwise the runtime threads are lost. pools: Pools, @@ -116,12 +573,19 @@ pub struct NoStealRuntime { } impl NoStealRuntime { - /// Create a new [NoStealRuntime]. Panic if `threads` is 0 - pub fn new(threads: usize, name: &str) -> Self { + /// Create a new [`NoStealRuntime`] with blocking pool options. Panic if `threads` is 0. + pub fn new( + threads: usize, + name: &str, + blocking_opts: BlockingPoolOpts, + runtime_opts: RuntimeOpts, + ) -> Self { assert!(threads != 0); NoStealRuntime { threads, name: name.to_string(), + blocking_opts, + runtime_opts, pools: Arc::new(OnceCell::new()), controls: OnceCell::new(), } @@ -131,7 +595,13 @@ impl NoStealRuntime { let mut pools = Vec::with_capacity(self.threads); let mut controls = Vec::with_capacity(self.threads); for _ in 0..self.threads { - let rt = Builder::new_current_thread().enable_all().build().unwrap(); + let mut builder = Builder::new_current_thread(); + builder.enable_all(); + apply_blocking_opts(&mut builder, &self.blocking_opts); + apply_metrics_opts(&mut builder, &self.runtime_opts.metrics); + let rt = builder + .build() + .expect("failed to build no-steal Tokio runtime worker"); let handler = rt.handle().clone(); let (tx, rx) = channel::(); let pools_ref = self.pools.clone(); diff --git a/pingora-rustls/Cargo.toml b/pingora-rustls/Cargo.toml index efa377bf6..51cd00ff7 100644 --- a/pingora-rustls/Cargo.toml +++ b/pingora-rustls/Cargo.toml @@ -18,7 +18,7 @@ path = "src/lib.rs" log = "0.4.21" pingora-error = { version = "0.8.0", path = "../pingora-error"} ring = "0.17.12" -rustls = "0.23.12" +rustls = { version = "0.23.12", features = ["ring"] } rustls-native-certs = "0.7.1" rustls-pemfile = "2.1.2" rustls-pki-types = "1.7.0" diff --git a/pingora-rustls/src/lib.rs b/pingora-rustls/src/lib.rs index 097a8da5a..deb0c88bb 100644 --- a/pingora-rustls/src/lib.rs +++ b/pingora-rustls/src/lib.rs @@ -28,9 +28,19 @@ use pingora_error::{Error, ErrorType, OrErr, Result}; pub use rustls::server::danger::{ClientCertVerified, ClientCertVerifier}; pub use rustls::server::{ClientCertVerifierBuilder, WebPkiClientVerifier}; pub use rustls::{ - client::WebPkiServerVerifier, version, CertificateError, ClientConfig, DigitallySignedStruct, - Error as RusTlsError, KeyLogFile, RootCertStore, ServerConfig, SignatureScheme, Stream, + client::WebPkiServerVerifier, crypto::CryptoProvider, version, CertificateError, ClientConfig, + DigitallySignedStruct, Error as RusTlsError, KeyLogFile, RootCertStore, ServerConfig, + SignatureScheme, Stream, }; + +/// Install the default `ring` CryptoProvider for rustls. +/// +/// rustls 0.23+ requires an explicit provider. This function installs `ring` +/// as the process-level default. Safe to call multiple times — subsequent +/// calls are no-ops. +pub fn install_default_crypto_provider() { + let _ = CryptoProvider::install_default(rustls::crypto::ring::default_provider()); +} pub use rustls_native_certs::load_native_certs; use rustls_pemfile::Item; pub use rustls_pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; diff --git a/pingora-s2n/src/ext.rs b/pingora-s2n/src/ext.rs new file mode 100644 index 000000000..c3bfc2c86 --- /dev/null +++ b/pingora-s2n/src/ext.rs @@ -0,0 +1,57 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Extended functionalities for s2n-tls + +use s2n_tls::connection::Connection; +use s2n_tls::error::Error; + +/// Export keying material from a TLS connection +/// +/// Derives keying material for application use in accordance with RFC 5705. +/// +/// Note: Currently only available with TLS 1.3 connections. +/// +/// See [tls_exporter](https://docs.rs/s2n-tls/latest/s2n_tls/connection/struct.Connection.html#method.tls_exporter). +pub fn ssl_export_keying_material( + conn: &Connection, + out: &mut [u8], + label: &str, + context: Option<&[u8]>, +) -> Result<(), Error> { + let context_bytes = context.unwrap_or(&[]); + conn.tls_exporter(label.as_bytes(), context_bytes, out) +} + +#[cfg(test)] +mod tests { + use super::*; + use s2n_tls::config::Builder; + use s2n_tls::enums::Mode; + + #[test] + fn test_ssl_export_keying_material_exists() { + // This test verifies that ssl_export_keying_material function exists + // and has the correct signature. Actual functional testing requires + // an established TLS connection. + let config = Builder::new().build().unwrap(); + let mut conn = s2n_tls::connection::Connection::new(Mode::Client); + conn.set_config(config).unwrap(); + let mut out = [0u8; 32]; + + // This will fail since there's no established connection, but verifies + // the function signature is correct + let _ = ssl_export_keying_material(&conn, &mut out, "test", None); + } +} diff --git a/pingora-s2n/src/lib.rs b/pingora-s2n/src/lib.rs index aef1cef35..e4444ed63 100644 --- a/pingora-s2n/src/lib.rs +++ b/pingora-s2n/src/lib.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod ext; + use pingora_error::{Error, ErrorType, Result}; use std::fs; diff --git a/pingora-timeout/src/fast_timeout.rs b/pingora-timeout/src/fast_timeout.rs index 27535e11e..fed336429 100644 --- a/pingora-timeout/src/fast_timeout.rs +++ b/pingora-timeout/src/fast_timeout.rs @@ -27,6 +27,7 @@ use super::timer::*; use super::*; use once_cell::sync::Lazy; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; static TIMER_MANAGER: Lazy> = Lazy::new(|| { @@ -35,6 +36,15 @@ static TIMER_MANAGER: Lazy> = Lazy::new(|| { tm }); +/// Default duration above which [`fast_timeout()`] falls back to Tokio's timeout. +pub const DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD: Duration = Duration::from_secs(15 * 60); +const DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS: u64 = + DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD.as_secs(); +const FAST_TIMEOUT_TO_TOKIO_DISABLED: u64 = u64::MAX; + +static FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS: AtomicU64 = + AtomicU64::new(DEFAULT_FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS); + fn check_clock_thread(tm: &Arc) { if tm.should_i_start_clock() { std::thread::Builder::new() @@ -44,6 +54,35 @@ fn check_clock_thread(tm: &Arc) { } } +fn should_use_tokio_timeout(duration: Duration) -> bool { + let threshold_secs = FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS.load(Ordering::Relaxed); + threshold_secs != FAST_TIMEOUT_TO_TOKIO_DISABLED && duration.as_secs() > threshold_secs +} + +/// Set the duration above which [`fast_timeout()`] falls back to Tokio's timeout. +/// +/// When set to `None`, [`fast_timeout()`] never falls back to Tokio's timeout. +/// +/// Long Pingora fast-timeout timers are kept in [`TimerManager`] until +/// their deadline expires. Tokio's timeout has cancellation cleanup, so it is a +/// better fit for long durations that are often cancelled before expiry. +pub fn set_fast_timeout_to_tokio_threshold(threshold: Option) { + FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS.store( + threshold.map_or(FAST_TIMEOUT_TO_TOKIO_DISABLED, |threshold| { + threshold.as_secs() + }), + Ordering::Relaxed, + ); +} + +/// Return the duration above which [`fast_timeout()`] falls back to Tokio's timeout. +pub fn fast_timeout_to_tokio_threshold() -> Option { + match FAST_TIMEOUT_TO_TOKIO_THRESHOLD_SECS.load(Ordering::Relaxed) { + FAST_TIMEOUT_TO_TOKIO_DISABLED => None, + threshold_secs => Some(Duration::from_secs(threshold_secs)), + } +} + /// The timeout generated by [fast_timeout()]. /// /// Users don't need to interact with this object. @@ -59,13 +98,56 @@ impl ToTimeout for FastTimeout { } } +enum TimeoutKind { + Fast, + Tokio, +} + +/// The timeout selector generated by [fast_timeout()]. +/// +/// Users don't need to interact with this object. +pub struct FastOrTokioTimeout { + duration: Duration, + kind: TimeoutKind, +} + +impl FastOrTokioTimeout { + fn use_tokio(&self) -> bool { + matches!(self.kind, TimeoutKind::Tokio) + } +} + +impl ToTimeout for FastOrTokioTimeout { + fn timeout(&self) -> Pin + Send + Sync>> { + match self.kind { + TimeoutKind::Fast => FastTimeout::create(self.duration).timeout(), + TimeoutKind::Tokio => TokioTimeout::create(self.duration).timeout(), + } + } + + fn create(duration: Duration) -> Self { + Self { + duration, + kind: if should_use_tokio_timeout(duration) { + TimeoutKind::Tokio + } else { + TimeoutKind::Fast + }, + } + } +} + /// Similar to [tokio::time::timeout] but more efficient. -pub fn fast_timeout(duration: Duration, future: T) -> Timeout +pub fn fast_timeout(duration: Duration, future: T) -> Timeout where T: Future, { - check_clock_thread(&TIMER_MANAGER); - Timeout::new_with_delay(future, duration) + let callback = FastOrTokioTimeout::create(duration); + if !callback.use_tokio() { + check_clock_thread(&TIMER_MANAGER); + } + + Timeout::new_with_callback(future, callback) } /// Similar to [tokio::time::sleep] but more efficient. @@ -115,7 +197,7 @@ mod tests { tokio_sleep(Duration::from_secs(1)).await; 1 }; - let to = fast_timeout(Duration::from_secs(1000), fut); + let to = fast_timeout(Duration::from_secs(60), fut); assert_eq!(to.await.unwrap(), 1) } @@ -125,7 +207,35 @@ mod tests { fast_sleep(Duration::from_secs(1)).await; 1 }; - let to = fast_timeout(Duration::from_secs(1000), fut); + let to = fast_timeout(Duration::from_secs(60), fut); assert_eq!(to.await.unwrap(), 1) } + + #[test] + fn test_default_fast_timeout_to_tokio_threshold() { + assert_eq!( + Some(Duration::from_secs(15 * 60)), + fast_timeout_to_tokio_threshold() + ); + } + + #[test] + fn test_fast_timeout_uses_fast_below_threshold() { + let callback = FastOrTokioTimeout::create(Duration::from_secs(1)); + assert!(!callback.use_tokio()); + } + + #[test] + fn test_fast_timeout_uses_fast_at_threshold() { + let callback = FastOrTokioTimeout::create(fast_timeout_to_tokio_threshold().unwrap()); + assert!(!callback.use_tokio()); + } + + #[test] + fn test_fast_timeout_uses_tokio_above_threshold() { + let callback = FastOrTokioTimeout::create( + fast_timeout_to_tokio_threshold().unwrap() + Duration::from_secs(1), + ); + assert!(callback.use_tokio()); + } } diff --git a/pingora-timeout/src/lib.rs b/pingora-timeout/src/lib.rs index 707f7be86..318348a5d 100644 --- a/pingora-timeout/src/lib.rs +++ b/pingora-timeout/src/lib.rs @@ -109,10 +109,14 @@ where F: ToTimeout, { pub(crate) fn new_with_delay(value: T, d: Duration) -> Timeout { + Self::new_with_callback(value, F::create(d)) + } + + pub(crate) fn new_with_callback(value: T, callback: F) -> Timeout { Timeout { value, delay: None, - callback: F::create(d), + callback, } } } diff --git a/pingora/Cargo.toml b/pingora/Cargo.toml index cb16664e1..d10ce3093 100644 --- a/pingora/Cargo.toml +++ b/pingora/Cargo.toml @@ -29,6 +29,7 @@ pingora-load-balancing = { version = "0.8.0", path = "../pingora-load-balancing" pingora-proxy = { version = "0.8.0", path = "../pingora-proxy", optional = true, default-features = false } pingora-cache = { version = "0.8.0", path = "../pingora-cache", optional = true, default-features = false } + # Only used for documenting features, but doesn't work in any other dependency # group :( document-features = { version = "0.2.10", optional = true } @@ -37,18 +38,21 @@ document-features = { version = "0.2.10", optional = true } clap = { version = "4.5", features = ["derive"] } tokio = { workspace = true, features = ["rt-multi-thread", "signal"] } env_logger = "0.11" -reqwest = { version = "0.11", features = ["rustls"], default-features = false } -hyper = "0.14" +reqwest = { version = "0.12", features = ["rustls-tls", "http2"], default-features = false } +hyper = { version = "1", features = ["client", "http1", "http2"] } +hyper-util = { version = "0.1", features = ["client-legacy", "http1", "http2"] } +http-body-util = "0.1" async-trait = { workspace = true } http = { workspace = true } log = { workspace = true } -prometheus = "0.13" +pingora-prometheus = { version = "0.8.0", path = "../pingora-prometheus" } +prometheus = "0.14" once_cell = { workspace = true } bytes = { workspace = true } regex = "1" [target.'cfg(unix)'.dev-dependencies] -hyperlocal = "0.8" +hyperlocal = "0.9" jemallocator = "0.5" [features] @@ -126,6 +130,24 @@ time = [] ## Enable sentry for error notifications sentry = ["pingora-core/sentry"] +## Enable dial9 Tokio runtime telemetry configuration. +## +## Requires building with `--cfg tokio_unstable`. +dial9 = ["pingora-core/dial9"] + +## Enable dial9 S3-compatible trace segment upload support. +dial9-worker-s3 = ["dial9", "pingora-core/dial9-worker-s3"] + +## Enable dial9 CPU profiling support. +dial9-cpu-profiling = ["dial9", "pingora-core/dial9-cpu-profiling"] + +## Enable upstream modules: the `adjust_upstream_modules` callback, the +## `upstream_modules_ctx` on Session, and `init_upstream_modules` on ProxyHttp. +## +## Allows registering custom upstream modules that process response tasks +## before `upstream_compression`, and configuring them. +upstream_modules = ["pingora-proxy?/upstream_modules"] + ## Enable pre-TLS connection filtering connection_filter = [ "pingora-core/connection_filter", @@ -136,7 +158,7 @@ connection_filter = [ # These features are intentionally not documented openssl_derived = ["any_tls"] any_tls = [] -patched_http1 = ["pingora-core/patched_http1"] +patched_http1 = ["pingora-core/patched_http1", "pingora-proxy?/patched_http1"] document-features = [ "dep:document-features", "proxy", @@ -144,5 +166,9 @@ document-features = [ "cache", "time", "sentry", + "dial9", + "dial9-worker-s3", + "dial9-cpu-profiling", "connection_filter" ] +trace = ["pingora-cache?/trace", "pingora-proxy?/trace"] diff --git a/pingora/examples/graceful_upgrade.rs b/pingora/examples/graceful_upgrade.rs new file mode 100644 index 000000000..5a64ff7fd --- /dev/null +++ b/pingora/examples/graceful_upgrade.rs @@ -0,0 +1,186 @@ +// Copyright 2026 Cloudflare, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! # Graceful Upgrade Example +//! +//! Demonstrates the `daemon_wait_for_ready` feature, which coordinates graceful process upgrades +//! by ensuring the new process is fully bootstrapped before the old one begins shutting down. +//! +//! ## Background +//! +//! In a standard daemonized pingora service, the parent process exits immediately after the +//! daemon fork. During a graceful upgrade, the process manager sends SIGQUIT to the old process +//! as soon as the new process's parent exits — potentially before the new process has finished +//! initializing its backends, consistent hash rings, or other state. This can cause a brief +//! window of 502s. +//! +//! With `daemon_wait_for_ready = true`, the parent instead waits for the daemon to send SIGUSR1 +//! before exiting. The process manager only proceeds to stop the old process once the new one +//! signals that it is ready to serve traffic. +//! +//! ## Service startup order +//! +//! This example sets up the following dependency chain: +//! +//! ```text +//! BackendDiscoveryService HashRingService +//! \ / +//! \ / +//! BootstrapService (socket transfer + SIGUSR1 to parent) +//! ``` +//! +//! The bootstrap service — which handles transferring listening sockets from the old process and +//! sending SIGUSR1 to the parent to signal readiness — only runs after both slow initialization +//! services have completed. This ensures the parent never exits until the new process is truly +//! ready to serve traffic. +//! +//! ## Usage +//! +//! ```bash +//! # Run interactively (no daemonization) +//! cargo run --example graceful_upgrade -p pingora +//! +//! # Run as a daemon +//! cargo run --example graceful_upgrade -p pingora -- -d +//! +//! # Graceful upgrade of a running daemon instance +//! cargo run --example graceful_upgrade -p pingora -- -d -u +//! ``` + +use async_trait::async_trait; +use bytes::Bytes; +use clap::Parser; +use http::{Response, StatusCode}; +use log::info; +use std::num::NonZeroU64; +use std::time::Duration; +use tokio::time::sleep; + +use pingora::apps::http_app::ServeHttp; +use pingora::prelude::Opt; +use pingora::protocols::http::ServerSession; +use pingora::server::configuration::ServerConf; +use pingora::server::{Server, ShutdownWatch}; +use pingora::services::background::{background_service, BackgroundService}; +use pingora::services::listening::Service as ListeningService; + +/// Simulates slow backend discovery — e.g. resolving upstream endpoints from a service registry. +pub struct BackendDiscoveryService; + +#[async_trait] +impl BackgroundService for BackendDiscoveryService { + async fn start(&self, _shutdown: ShutdownWatch) { + info!("BackendDiscoveryService: discovering backends..."); + sleep(Duration::from_secs(2)).await; + info!("BackendDiscoveryService: backends ready"); + } +} + +/// Simulates slow consistent hash ring construction. Runs in parallel with +/// `BackendDiscoveryService`; bootstrap waits for both to complete. +pub struct HashRingService; + +#[async_trait] +impl BackgroundService for HashRingService { + async fn start(&self, _shutdown: ShutdownWatch) { + info!("HashRingService: building consistent hash ring..."); + sleep(Duration::from_secs(3)).await; + info!("HashRingService: hash ring ready"); + } +} + +/// A minimal HTTP service that responds to every request with 200 OK. +/// +/// Accepts an optional `sleep` query parameter specifying how many seconds to wait before +/// responding (e.g. `GET /?sleep=20`). This makes in-flight requests easy to observe during a +/// graceful upgrade: a request with a long sleep that arrives just before the upgrade begins will +/// still be running when the new process starts up, demonstrating that the old process keeps +/// serving until all connections are drained. +pub struct HelloApp; + +#[async_trait] +impl ServeHttp for HelloApp { + async fn response(&self, http_stream: &mut ServerSession) -> Response> { + let delay_secs = http_stream + .req_header() + .uri + .query() + .and_then(|q| { + q.split('&').find_map(|pair| { + let (key, val) = pair.split_once('=')?; + if key == "sleep" { + val.parse::().ok() + } else { + None + } + }) + }) + .unwrap_or(0); + + if delay_secs > 0 { + sleep(Duration::from_secs(delay_secs)).await; + } + + let body = Bytes::from("hello from graceful_upgrade example\n"); + Response::builder() + .status(StatusCode::OK) + .header(http::header::CONTENT_TYPE, "text/plain") + .header(http::header::CONTENT_LENGTH, body.len()) + .body(body.to_vec()) + .unwrap() + } +} + +fn main() { + env_logger::init(); + + let opt = Some(Opt::parse()); + + // Build a ServerConf with daemon_wait_for_ready enabled. + // + // When the server is started with -d (daemon mode), the parent process waits for SIGUSR1 + // before exiting. The daemon sends SIGUSR1 only after the bootstrap service completes — + // which in this example means after both slow services have signaled readiness. + let conf = ServerConf { + daemon: true, + daemon_wait_for_ready: true, + daemon_ready_timeout_seconds: NonZeroU64::new(60), + ..ServerConf::default() + }; + + let mut server = Server::new_with_opt_and_conf(opt, conf); + + // Add the slow initialization services and retain their handles so bootstrap can depend + // on them. Both run in parallel; the slowest (HashRingService at 3s) sets the pace. + let backend_handle = server.add_service(background_service( + "backend_discovery", + BackendDiscoveryService, + )); + let hash_ring_handle = server.add_service(background_service("hash_ring", HashRingService)); + + // bootstrap_as_a_service() registers the bootstrap service (socket transfer from the old + // process + SIGUSR1 to the parent) and returns its ServiceHandle. Declaring the slow + // services as dependencies ensures bootstrap only runs once both are ready. + let bootstrap_handle = server.bootstrap_as_a_service(); + bootstrap_handle.add_dependencies([&backend_handle, &hash_ring_handle]); + + let mut http_service = ListeningService::new("hello_http".to_string(), HelloApp); + http_service.add_tcp("0.0.0.0:8000"); + + server + .add_service(http_service) + .add_dependency(backend_handle); + + server.run_forever(); +} diff --git a/pingora/examples/server.rs b/pingora/examples/server.rs index 0a055acc0..37a246cd4 100644 --- a/pingora/examples/server.rs +++ b/pingora/examples/server.rs @@ -20,7 +20,6 @@ use pingora::protocols::TcpKeepalive; use pingora::server::configuration::Opt; use pingora::server::{Server, ShutdownWatch}; use pingora::services::background::{background_service, BackgroundService}; -use pingora::services::listening::Service as ListeningService; use pingora::services::ServiceWithDependents; use async_trait::async_trait; @@ -186,7 +185,7 @@ pub fn main() { &key_path, ); - let mut prometheus_service_http = ListeningService::prometheus_http_service(); + let mut prometheus_service_http = pingora_prometheus::prometheus_http_service(); prometheus_service_http.add_tcp("127.0.0.1:6150"); let background_service = background_service("example", ExampleBackgroundService {}); diff --git a/pingora/tests/pingora_conf.yaml b/pingora/tests/pingora_conf.yaml deleted file mode 100644 index c21ae15a1..000000000 --- a/pingora/tests/pingora_conf.yaml +++ /dev/null @@ -1,5 +0,0 @@ ---- -version: 1 -client_bind_to_ipv4: - - 127.0.0.2 -ca_file: tests/keys/server.crt \ No newline at end of file