From d20d023c0e9f09be4b40d3c7d12785a4a0dfa878 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Mon, 26 Oct 2026 09:52:57 -0500 Subject: [PATCH 1/2] fix(ingest): bound JSON flattening rows and expanded bytes --- docs/json-flattening-limits.md | 75 + src/cli.rs | 32 + src/handlers/http/ingest.rs | 29 + src/handlers/http/kinesis.rs | 285 ++- src/handlers/http/logstream.rs | 7 +- src/handlers/http/modal/utils/ingest_utils.rs | 30 +- src/option.rs | 86 + src/otel/logs.rs | 508 +++++- src/storage/field_stats.rs | 114 +- src/utils/json/flatten.rs | 1571 +++++++++++++++-- src/utils/json/mod.rs | 300 +++- tools/flatten-bench/.gitignore | 2 + tools/flatten-bench/Cargo.lock | 436 +++++ tools/flatten-bench/Cargo.toml | 32 + tools/flatten-bench/README.md | 79 + tools/flatten-bench/RESULTS.md | 155 ++ tools/flatten-bench/build.rs | 71 + tools/flatten-bench/http_smoke.py | 706 ++++++++ tools/flatten-bench/src/alloc_counter.rs | 134 ++ tools/flatten-bench/src/api.rs | 175 ++ tools/flatten-bench/src/bin/flatten_bench.rs | 384 ++++ tools/flatten-bench/src/fixtures.rs | 225 +++ tools/flatten-bench/src/lib.rs | 86 + tools/flatten-bench/tests/alloc_limits.rs | 477 +++++ 24 files changed, 5805 insertions(+), 194 deletions(-) create mode 100644 docs/json-flattening-limits.md create mode 100644 tools/flatten-bench/.gitignore create mode 100644 tools/flatten-bench/Cargo.lock create mode 100644 tools/flatten-bench/Cargo.toml create mode 100644 tools/flatten-bench/README.md create mode 100644 tools/flatten-bench/RESULTS.md create mode 100644 tools/flatten-bench/build.rs create mode 100644 tools/flatten-bench/http_smoke.py create mode 100644 tools/flatten-bench/src/alloc_counter.rs create mode 100644 tools/flatten-bench/src/api.rs create mode 100644 tools/flatten-bench/src/bin/flatten_bench.rs create mode 100644 tools/flatten-bench/src/fixtures.rs create mode 100644 tools/flatten-bench/src/lib.rs create mode 100644 tools/flatten-bench/tests/alloc_limits.rs diff --git a/docs/json-flattening-limits.md b/docs/json-flattening-limits.md new file mode 100644 index 000000000..04bc08b91 --- /dev/null +++ b/docs/json-flattening-limits.md @@ -0,0 +1,75 @@ +# JSON flattening limits + +Generic flattening turns every nested array of an event into rows: an event becomes one row +for each combination of the values of its arrays. Sibling arrays multiply, so +`{"a": [0..999], "b": [0..999], "c": [0..999]}`, about 11 KiB of JSON, would expand into +10^9 rows. Three limits cap that expansion. Parseable checks them before it builds any row, +and rejects the whole request with `400 Bad Request` when one is exceeded. Nothing from a +rejected request is ingested. + +| Environment variable | Flag | Default | Limit | +| --------------------------------- | --------------------------------- | -------------------- | --------------------------------------------------------- | +| `P_MAX_FLATTEN_ROWS_PER_EVENT` | `--max-flatten-rows-per-event` | `1000` | Rows a single event may expand into | +| `P_MAX_FLATTEN_ROWS_PER_REQUEST` | `--max-flatten-rows-per-request` | `10000` | Rows all events of one request may expand into together | +| `P_MAX_FLATTEN_BYTES_PER_REQUEST` | `--max-flatten-bytes-per-request` | `67108864` (64 MiB) | Bytes all those rows may add up to, measured as compact JSON | + +Each value must be a positive integer. A limit cannot be set to zero or switched off. + +The byte limit counts the size of the expanded rows as compact JSON, including keys, +quotes, escapes and delimiters. It also counts what is copied into every row afterwards: +Kinesis request metadata and OTEL body key prefixes (see below). It is not a measure of the +memory a request uses. + +After generic flattening, every row is flattened structurally: nested objects become fields +named after the keys on their path, joined with `_`. A parent's name is repeated in every +field below it, so `{"<128 character name>": {"f0": 0, ..., "f31": 31}, "r": [0, 1, 2, 3]}` +expands into 4 rows of about 400 bytes, but each of them structurally flattened takes about +4.4 KiB. Every event is therefore charged the larger of the two: the compact JSON of its +generic rows, or of those rows structurally flattened. Both are computed in one pass over +the event, before any row or field name is built. The structurally flattened size counts +every scalar as a field of its own, even when its joined name collides with another field +(for example `{"a_b": 1, "a": {"b": 2}}`, where one `a_b` overwrites the other), so with +such collisions it is an overestimate. Empty objects add no field. + +## What counts as an event + +The limits apply only where generic flattening runs: + +- **JSON and custom sources** (`/api/v1/ingest`, `/api/v1/logstream/{name}`) on datasets using + schema version V1. Every object in the request body is one event, after any nested + top-level arrays are spliced in. Bodies nested deeper than `P_MAX_FLATTEN_LEVEL` skip + generic flattening, so these limits do not apply to them. +- **Kinesis**: every decoded `data` record is one event, even when it holds a JSON array. + Each of its rows also gets the request's `requestId` and `timestamp`, so their compact + JSON size, with keys and delimiters, is counted once per row, on top of the larger of the + generic and structurally flattened sizes. It is counted in full even for rows that + already have these fields. +- **OTEL logs**: a log record whose body (or a body field) is a string holding a JSON object + has that object flattened. All flattened bodies of one log record together are one event. + This counts every intermediate row, even though the rows end up merged into the single + log record. Each top-level field of a flattened body is renamed to `{body key}_{field}` + before the merge, so the body key, with its escapes and the `_`, is counted again for + every such field in every row. Structurally flattening the record afterwards repeats the + body key in the name of every field nested below those too, `{body key}_{field}_{name}`. + A body is charged the larger of its generic rows with the renamed top-level fields, or + of its structurally flattened rows with the body key in every field name. This counts + the rows of a body as if they were all kept, though only the last value of each field + ends up in the record. +- **Schema detection** (`/api/v1/logstream/schema/detect`) counts the request body the same + way as JSON ingestion, structurally flattened rows included. + +Every request has a budget of its own. All events of a request share that budget, across +custom partition elements, Kinesis records, and OTEL resources and scopes. OTEL metrics and +traces, Kafka, and V0 datasets do not use generic flattening, so these limits do not apply +to them. + +Internal dataset statistics (`pstats`) use the configured byte limit. Their row limits come +from the size of the stats payload, which is the field count times `P_MAX_FIELD_STATISTICS`. + +## When a request is rejected + +The error message names the limit that was exceeded and its value. To get under the limit, +reduce the number or length of the nested arrays in each event, or split the batch into +smaller requests. You can also raise the limit. + +`tools/flatten-bench/http_smoke.py` checks these responses against a built server binary. diff --git a/src/cli.rs b/src/cli.rs index 5a2dbb9e9..08a7c5a3b 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -36,6 +36,9 @@ pub const DEFAULT_USERNAME: &str = "admin"; pub const DEFAULT_PASSWORD: &str = "admin"; pub const DATASET_FIELD_COUNT_LIMIT: usize = 1000; +pub const DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT: usize = 1000; +pub const DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST: usize = 10_000; +pub const DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST: usize = 64 * 1024 * 1024; const DEFAULT_PARQUET_METADATA_CACHE_SIZE: usize = 50 * 1024 * 1024; // 50 MB #[derive(Parser)] #[command( @@ -640,6 +643,35 @@ pub struct Options { )] pub max_event_payload_size: usize, + // generic flattening expands nested arrays into the cartesian product of their values, + // these bound that expansion for JSON/custom sources, Kinesis and OTEL log bodies + #[arg( + long, + env = "P_MAX_FLATTEN_ROWS_PER_EVENT", + default_value_t = DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT, + value_parser = validation::validate_flatten_limit, + help = "Maximum number of rows a single event may expand into when its nested arrays are flattened" + )] + pub max_flatten_rows_per_event: usize, + + #[arg( + long, + env = "P_MAX_FLATTEN_ROWS_PER_REQUEST", + default_value_t = DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST, + value_parser = validation::validate_flatten_limit, + help = "Maximum number of rows all events of an ingest request may expand into when their nested arrays are flattened" + )] + pub max_flatten_rows_per_request: usize, + + #[arg( + long, + env = "P_MAX_FLATTEN_BYTES_PER_REQUEST", + default_value_t = DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST, + value_parser = validation::validate_flatten_limit, + help = "Maximum total size in bytes of the rows, measured as compact JSON, that all events of an ingest request may expand into when their nested arrays are flattened. This is not a bound on the memory used" + )] + pub max_flatten_bytes_per_request: usize, + #[arg( long, env = "P_NORMALISE_OTEL_PAYLOAD", diff --git a/src/handlers/http/ingest.rs b/src/handlers/http/ingest.rs index 2f79ba105..8b4efbacb 100644 --- a/src/handlers/http/ingest.rs +++ b/src/handlers/http/ingest.rs @@ -562,6 +562,8 @@ impl actix_web::ResponseError for PostError { | CreateStream(CreateStreamError::StreamNameValidation(_)) | OtelNotSupported(_) => StatusCode::BAD_REQUEST, + JsonFlattenError(e) if e.is_limit_exceeded() => StatusCode::BAD_REQUEST, + Event(_) | CreateStream(_) | CustomError(_) @@ -1129,4 +1131,31 @@ mod tests { &Float64Array::from(vec![None, None, None, Some(2.0)]) ); } + + #[test] + fn flatten_limit_errors_are_bad_requests() { + use actix_web::{ResponseError, http::StatusCode}; + + use super::PostError; + use crate::utils::json::flatten::JsonFlattenError; + + for limit in [ + JsonFlattenError::RowsPerEventLimitExceeded(1000), + JsonFlattenError::RowsPerRequestLimitExceeded(10_000), + JsonFlattenError::BytesPerRequestLimitExceeded(64), + ] { + let message = limit.to_string(); + let error = PostError::JsonFlattenError(limit); + assert_eq!(error.status_code(), StatusCode::BAD_REQUEST); + assert!(error.to_string().contains(&message)); + } + // Limit errors from the JSON and Kinesis paths arrive wrapped in anyhow + let error = PostError::Invalid(JsonFlattenError::RowsPerEventLimitExceeded(1).into()); + assert_eq!(error.status_code(), StatusCode::BAD_REQUEST); + // Other flatten errors keep their status + assert_eq!( + PostError::JsonFlattenError(JsonFlattenError::ExpectedObjectInArray).status_code(), + StatusCode::INTERNAL_SERVER_ERROR + ); + } } diff --git a/src/handlers/http/kinesis.rs b/src/handlers/http/kinesis.rs index 6fc713773..38b845dbf 100644 --- a/src/handlers/http/kinesis.rs +++ b/src/handlers/http/kinesis.rs @@ -21,7 +21,11 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use std::str; -use crate::utils::json::flatten::{generic_flattening, has_more_than_max_allowed_levels}; +use crate::parseable::PARSEABLE; +use crate::utils::json::flatten::{ + FlattenBudget, FlattenCost, JsonFlattenError, expand_generic_flattening, flatten_pipeline_cost, + has_more_than_levels, string_json_len, +}; #[derive(Serialize, Deserialize, Debug)] #[serde(rename_all = "camelCase")] @@ -47,6 +51,8 @@ struct Data { // ] // } // The data field is base64 encoded JSON (there can be multiple data fields), and there is a requestId and timestamp field. +// Each decoded data field is one event for `budget`, even when it holds an array, charged +// for its generic rows or, when larger, those rows structurally flattened, plus the metadata. // Kinesis logs are flattened to the following format: // { // "CHANGE": 3.16, @@ -59,8 +65,21 @@ struct Data { // "requestId": "b858288a-f5d8-4181-a746-3f3dd716be8a", // "timestamp": "1704964113659" // } -pub fn flatten_kinesis_logs(message: Message) -> Result, anyhow::Error> { +pub fn flatten_kinesis_logs( + message: Message, + budget: &mut FlattenBudget, +) -> Result, anyhow::Error> { + flatten_kinesis_logs_inner(message, budget, PARSEABLE.options.event_flatten_level) +} + +fn flatten_kinesis_logs_inner( + message: Message, + budget: &mut FlattenBudget, + max_flatten_level: usize, +) -> Result, anyhow::Error> { let mut vec_kinesis_json = Vec::new(); + let timestamp = message.timestamp.to_string(); + let metadata_bytes = metadata_json_len(&message.request_id, ×tamp); for record in message.records.iter() { let bytes = STANDARD.decode(record.data.clone())?; @@ -70,8 +89,25 @@ pub fn flatten_kinesis_logs(message: Message) -> Result, anyhow::Erro // If it has less than or equal to the allowed levels, we flatten it. // If it has more than the allowed levels, we just push it as is // without flattening or modifying it. - if !has_more_than_max_allowed_levels(&json, 1) { - let flattened_json_arr = generic_flattening(&json)?; + if !has_more_than_levels(&json, 1, max_flatten_level) { + // `push_logs` structurally flattens the rows later, charge that now too + let cost = flatten_pipeline_cost(&json, budget.limits())?.charged(); + // Every row gets its own copy of the metadata, charge it before building any + let limits = budget.limits(); + let metadata = metadata_bytes + .and_then(|bytes| bytes.checked_mul(cost.rows)) + .ok_or(JsonFlattenError::BytesPerRequestLimitExceeded( + limits.max_bytes_per_request, + ))?; + let cost = cost.add( + FlattenCost { + rows: 0, + bytes: metadata, + }, + limits, + )?; + budget.charge_event(cost)?; + let flattened_json_arr = expand_generic_flattening(&json)?; for flattened_json in flattened_json_arr { let mut kinesis_json: Map = serde_json::from_value(flattened_json)?; @@ -79,10 +115,7 @@ pub fn flatten_kinesis_logs(message: Message) -> Result, anyhow::Erro "requestId".to_owned(), Value::String(message.request_id.clone()), ); - kinesis_json.insert( - "timestamp".to_owned(), - Value::String(message.timestamp.to_string()), - ); + kinesis_json.insert("timestamp".to_owned(), Value::String(timestamp.clone())); vec_kinesis_json.push(Value::Object(kinesis_json)); } @@ -113,3 +146,239 @@ pub fn flatten_kinesis_logs(message: Message) -> Result, anyhow::Erro Ok(vec_kinesis_json) } + +/// Compact JSON bytes the `requestId` and `timestamp` fields add to a row, with their +/// delimiters. Charged in full even for rows that already have these fields, which get +/// overwritten, or are empty and need one comma less. +fn metadata_json_len(request_id: &str, timestamp: &str) -> Option { + const FIXED: usize = r#","requestId":,"timestamp":"#.len(); + string_json_len(request_id)? + .checked_add(string_json_len(timestamp)?)? + .checked_add(FIXED) +} + +#[cfg(test)] +mod tests { + use base64::{Engine as _, engine::general_purpose::STANDARD}; + use serde_json::{Value, json}; + + use super::{Data, Message, flatten_kinesis_logs_inner}; + use crate::utils::json::flatten::{FlattenBudget, FlattenLimits, JsonFlattenError}; + + fn message(records: &[Value]) -> Message { + message_with_id(records, "id") + } + + fn message_with_id(records: &[Value], request_id: &str) -> Message { + Message { + records: records + .iter() + .map(|record| Data { + data: STANDARD.encode(record.to_string()), + }) + .collect(), + request_id: request_id.to_owned(), + timestamp: 1, + } + } + + fn limits(rows_per_event: usize, rows_per_request: usize) -> FlattenBudget { + FlattenBudget::new(FlattenLimits { + max_rows_per_event: rows_per_event, + max_rows_per_request: rows_per_request, + max_bytes_per_request: 1 << 20, + }) + } + + fn limit_error(result: anyhow::Result>) -> JsonFlattenError { + result + .unwrap_err() + .downcast::() + .expect("a flatten limit error") + } + + // 2 + 3 rows, from elements that are small on their own + fn record() -> Value { + json!([{"a": [1, 2]}, {"b": [1, 2, 3]}]) + } + + #[test] + fn decoded_record_is_one_event() { + let mut budget = limits(5, 10); + let rows = flatten_kinesis_logs_inner(message(&[record()]), &mut budget, 10).unwrap(); + assert_eq!(rows.len(), 5); + assert_eq!( + rows[0], + json!({"a": 1, "requestId": "id", "timestamp": "1"}) + ); + assert_eq!(budget.used().rows, 5); + + assert!(matches!( + limit_error(flatten_kinesis_logs_inner( + message(&[record()]), + &mut limits(4, 10), + 10 + )), + JsonFlattenError::RowsPerEventLimitExceeded(4) + )); + } + + #[test] + fn records_share_the_request_budget() { + let mut budget = limits(5, 10); + let rows = + flatten_kinesis_logs_inner(message(&[record(), record()]), &mut budget, 10).unwrap(); + assert_eq!(rows.len(), 10); + + assert!(matches!( + limit_error(flatten_kinesis_logs_inner( + message(&[record(), record()]), + &mut limits(5, 9), + 10 + )), + JsonFlattenError::RowsPerRequestLimitExceeded(9) + )); + } + + #[test] + fn records_nested_too_deep_are_kept_and_not_charged() { + let deep = json!({"a": {"b": [1, 2, 3]}}); + let mut budget = limits(1, 1); + let rows = flatten_kinesis_logs_inner(message(std::slice::from_ref(&deep)), &mut budget, 1) + .unwrap(); + assert_eq!(rows, vec![deep]); + assert_eq!(budget.used().rows, 0); + } + + fn byte_limit(bytes: usize) -> FlattenBudget { + FlattenBudget::new(FlattenLimits { + max_rows_per_event: 1000, + max_rows_per_request: 10_000, + max_bytes_per_request: bytes, + }) + } + + #[test] + fn metadata_of_every_row_is_charged_exactly() { + // Escapes in the request id are charged as serialized + let records = [record(), json!({"c": "x\n\"y"}), json!({})]; + let id = "id-\"quoted\"\t"; + let mut budget = byte_limit(1 << 20); + let rows = + flatten_kinesis_logs_inner(message_with_id(&records, id), &mut budget, 10).unwrap(); + assert_eq!(rows.len(), 7); + let bytes: usize = rows + .iter() + .map(|row| serde_json::to_vec(row).unwrap().len()) + .sum(); + // The empty row is charged one comma that it does not get + assert_eq!(budget.used().bytes, bytes + 1); + + let used = budget.used().bytes; + flatten_kinesis_logs_inner(message_with_id(&records, id), &mut byte_limit(used), 10) + .unwrap(); + assert!(matches!( + limit_error(flatten_kinesis_logs_inner( + message_with_id(&records, id), + &mut byte_limit(used - 1), + 10 + )), + JsonFlattenError::BytesPerRequestLimitExceeded(limit) if limit == used - 1 + )); + } + + #[test] + fn large_metadata_repeated_over_many_rows_is_rejected() { + // 10 records of 100 small rows each: about 7 KiB of rows, but 1000 copies of a + // 4 KiB request id would be 4 MiB + let records: Vec = (0..10) + .map(|_| json!({"a": (0..100).collect::>()})) + .collect(); + let id = "r".repeat(4096); + let mut budget = byte_limit(64 * 1024); + assert!(matches!( + limit_error(flatten_kinesis_logs_inner( + message_with_id(&records, &id), + &mut budget, + 10 + )), + JsonFlattenError::BytesPerRequestLimitExceeded(65536) + )); + // A single record is over the limit by itself, so nothing was charged + assert_eq!(budget.used().bytes, 0); + // The same rows pass with a short request id + flatten_kinesis_logs_inner(message(&records), &mut byte_limit(64 * 1024), 10).unwrap(); + } + + #[test] + fn metadata_bytes_add_up_over_records() { + // Each record fits in the byte limit on its own, but not both of them + let records = [json!({"a": [1, 2, 3]}), json!({"a": [1, 2, 3]})]; + let id = "r".repeat(100); + let mut one = byte_limit(1 << 20); + flatten_kinesis_logs_inner(message_with_id(&records[..1], &id), &mut one, 10).unwrap(); + let per_record = one.used().bytes; + assert!(per_record > 3 * 100); + + let mut budget = byte_limit(2 * per_record - 1); + assert!(matches!( + limit_error(flatten_kinesis_logs_inner( + message_with_id(&records, &id), + &mut budget, + 10 + )), + JsonFlattenError::BytesPerRequestLimitExceeded(_) + )); + assert_eq!(budget.used().bytes, per_record); + flatten_kinesis_logs_inner( + message_with_id(&records, &id), + &mut byte_limit(2 * per_record), + 10, + ) + .unwrap(); + } + + #[test] + fn structurally_flattened_field_names_are_charged() { + // 4 rows of 32 fields below a long parent, whose name `push_logs` later copies into + // every field: the generic rows and metadata fit in 4096 bytes, those fields do not + let long_parent = |parent: &str| { + let fields: serde_json::Map = + (0..32).map(|i| (format!("f{i}"), json!(i))).collect(); + let mut map = serde_json::Map::new(); + map.insert(parent.to_owned(), Value::Object(fields)); + map.insert("r".to_owned(), json!([0, 1, 2, 3])); + Value::Object(map) + }; + let long = long_parent(&"p".repeat(128)); + let generic = + crate::utils::json::flatten::generic_flattening_cost(&long, &FlattenLimits::DEFAULT) + .unwrap(); + assert!(generic.bytes + 4 * r#","requestId":"id","timestamp":"1""#.len() <= 4096); + let mut budget = byte_limit(4096); + assert!(matches!( + limit_error(flatten_kinesis_logs_inner( + message(std::slice::from_ref(&long)), + &mut budget, + 10 + )), + JsonFlattenError::BytesPerRequestLimitExceeded(4096) + )); + assert_eq!(budget.used().bytes, 0); + + // With a short parent the structurally flattened rows and their metadata are charged + let mut budget = byte_limit(4096); + let rows = + flatten_kinesis_logs_inner(message(&[long_parent("p")]), &mut budget, 10).unwrap(); + let bytes: usize = rows + .iter() + .map(|row| { + let mut row = row.clone(); + crate::utils::json::flatten::flatten(&mut row, "_", None, None, None, false) + .unwrap(); + serde_json::to_vec(&row).unwrap().len() + }) + .sum(); + assert_eq!(budget.used().bytes, bytes); + } +} diff --git a/src/handlers/http/logstream.rs b/src/handlers/http/logstream.rs index 82c58bd70..2db9b1bea 100644 --- a/src/handlers/http/logstream.rs +++ b/src/handlers/http/logstream.rs @@ -33,7 +33,8 @@ use crate::tenants::TenantNotFound; use crate::utils::actix::extract_session_key_from_req; use crate::utils::get_tenant_id_from_request; use crate::utils::json::flatten::{ - self, convert_to_array, generic_flattening, has_more_than_max_allowed_levels, + self, FlattenBudget, FlattenLimits, convert_to_array, has_more_than_max_allowed_levels, + pipeline_flattening, }; use crate::{LOCK_EXPECT, stats, validator}; @@ -120,7 +121,9 @@ pub async fn detect_schema(Json(json): Json) -> Result match convert_to_array(flattened) { Ok(array) => array, Err(e) => { diff --git a/src/handlers/http/modal/utils/ingest_utils.rs b/src/handlers/http/modal/utils/ingest_utils.rs index fea950875..d5acb7ec4 100644 --- a/src/handlers/http/modal/utils/ingest_utils.rs +++ b/src/handlers/http/modal/utils/ingest_utils.rs @@ -50,7 +50,10 @@ use crate::{ storage::StreamType, utils::{ get_tenant_id_from_request, - json::{convert_array_to_object, flatten::convert_to_array}, + json::{ + convert_array_to_object, + flatten::{FlattenBudget, FlattenLimits, convert_to_array}, + }, }, }; @@ -238,12 +241,14 @@ pub fn flatten_and_push_logs( let tenant_str = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); // Verify the dataset fields count verify_dataset_fields_count(stream_name, tenant_id)?; + // Shared by every event of this request + let mut budget = FlattenBudget::new(FlattenLimits::from_config()); match log_source { LogSource::Kinesis => { //custom flattening required for Amazon Kinesis let message: Message = serde_json::from_value(json)?; - let flattened_kinesis_data = flatten_kinesis_logs(message)?; + let flattened_kinesis_data = flatten_kinesis_logs(message, &mut budget)?; let record = convert_to_array(flattened_kinesis_data)?; push_logs( stream_name, @@ -253,12 +258,13 @@ pub fn flatten_and_push_logs( time_partition, telemetry_type, tenant_id, + &mut budget, )?; } LogSource::OtelLogs => { let logs: LogsData = deserialize_otel_json(json, PARSEABLE.options.normalise_otel_payload)?; - let records = flatten_otel_logs(&logs, tenant_str); + let records = flatten_otel_logs(&logs, tenant_str, &mut budget)?; if !records.is_empty() { push_logs( stream_name, @@ -268,6 +274,7 @@ pub fn flatten_and_push_logs( None, telemetry_type, tenant_id, + &mut budget, )?; } } @@ -284,6 +291,7 @@ pub fn flatten_and_push_logs( None, telemetry_type, tenant_id, + &mut budget, )?; } } @@ -300,6 +308,7 @@ pub fn flatten_and_push_logs( None, telemetry_type, tenant_id, + &mut budget, )?; } } @@ -311,6 +320,7 @@ pub fn flatten_and_push_logs( time_partition, telemetry_type, tenant_id, + &mut budget, )?, } @@ -320,9 +330,19 @@ pub fn flatten_and_push_logs( #[instrument( name = "push_logs", level = "info", - skip(json, log_source, p_custom_fields, time_partition, telemetry_type, tenant_id), + skip( + json, + log_source, + p_custom_fields, + time_partition, + telemetry_type, + tenant_id, + budget + ), fields(stream_name, record_count = tracing::field::Empty) )] +// `budget` is shared with the other parts of the request, so it can not be folded into them +#[allow(clippy::too_many_arguments)] #[cfg_attr(feature = "hotpath", hotpath::measure)] pub fn push_logs( stream_name: &str, @@ -332,6 +352,7 @@ pub fn push_logs( time_partition: Option, telemetry_type: TelemetryType, tenant_id: &Option, + budget: &mut FlattenBudget, ) -> Result<(), PostError> { let stream = PARSEABLE.get_stream(stream_name, tenant_id)?; let time_partition_limit = stream.get_time_partition_limit(); @@ -349,6 +370,7 @@ pub fn push_logs( custom_partition.as_ref(), schema_version, log_source, + budget, )?; if data.is_empty() { diff --git a/src/option.rs b/src/option.rs index cdfad83db..f879a90d5 100644 --- a/src/option.rs +++ b/src/option.rs @@ -272,6 +272,16 @@ pub mod validation { Ok(query.to_owned()) } + pub fn validate_flatten_limit(s: &str) -> Result { + match s.parse::() { + Ok(limit) if limit > 0 => Ok(limit), + _ => Err(format!( + "It should be a positive integer no larger than {}", + usize::MAX + )), + } + } + pub fn validate_payload_size(s: &str) -> Result { const MIN_SIZE: usize = 100; // 100 bytes const MAX_SIZE: usize = 100 * 1024 * 1024; // 100 MB @@ -309,3 +319,79 @@ mod ui_tests { assert!(validation::ui_root("other").is_err()); } } + +#[cfg(test)] +mod flatten_limit_tests { + use clap::{CommandFactory, FromArgMatches}; + + use super::validation; + use crate::cli::Options; + use crate::utils::json::flatten::FlattenLimits; + + // Parses `args` with the env fallbacks of every argument turned off, so that P_* + // variables set where the tests run do not change the result + fn parse(args: &[&str]) -> Result { + let mut matches = Options::command() + .mut_args(|arg| arg.env(None::<&'static str>)) + .try_get_matches_from(std::iter::once("parseable").chain(args.iter().copied()))?; + Options::from_arg_matches_mut(&mut matches) + } + + #[test] + fn flatten_limits_must_be_positive_integers() { + assert_eq!(validation::validate_flatten_limit("1"), Ok(1)); + assert_eq!( + validation::validate_flatten_limit(&usize::MAX.to_string()), + Ok(usize::MAX) + ); + for invalid in ["0", "-1", "", "1.5", "ten", "18446744073709551616000"] { + assert!( + validation::validate_flatten_limit(invalid).is_err(), + "{invalid}" + ); + } + } + + #[test] + fn flatten_limits_default_and_override() { + let options = parse(&[]).unwrap(); + assert_eq!( + FlattenLimits::from_options(&options), + FlattenLimits::DEFAULT + ); + assert_eq!( + FlattenLimits::DEFAULT, + FlattenLimits { + max_rows_per_event: 1000, + max_rows_per_request: 10_000, + max_bytes_per_request: 67_108_864, + } + ); + + let options = parse(&[ + "--max-flatten-rows-per-event", + "5", + "--max-flatten-rows-per-request", + "50", + "--max-flatten-bytes-per-request", + "500", + ]) + .unwrap(); + assert_eq!( + FlattenLimits::from_options(&options), + FlattenLimits { + max_rows_per_event: 5, + max_rows_per_request: 50, + max_bytes_per_request: 500, + } + ); + + for flag in [ + "--max-flatten-rows-per-event", + "--max-flatten-rows-per-request", + "--max-flatten-bytes-per-request", + ] { + assert!(parse(&[flag, "0"]).is_err()); + } + } +} diff --git a/src/otel/logs.rs b/src/otel/logs.rs index eb53bb39a..2d7af8379 100644 --- a/src/otel/logs.rs +++ b/src/otel/logs.rs @@ -19,7 +19,10 @@ use super::otel_utils::collect_json_from_values; use super::otel_utils::convert_epoch_nano_to_timestamp; use super::otel_utils::insert_attributes; use crate::metrics::increment_logs_collected_by_date; -use crate::utils::json::flatten::generic_flattening; +use crate::utils::json::flatten::{ + FlattenBudget, FlattenCost, JsonFlattenError, expand_generic_flattening, flatten_pipeline_cost, + string_json_len, +}; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::logs::v1::LogRecord; use opentelemetry_proto::tonic::logs::v1::LogsData; @@ -170,7 +173,11 @@ fn category_from_body(body_str: &str) -> &'static str { /// this function flattens the `LogRecord` object /// and returns a `Map` of the flattened json /// this function is called recursively for each log record object in the otel logs -pub fn flatten_log_record(log_record: &LogRecord) -> Map { +/// The generic flattening of all JSON object bodies of the record is one event for `budget` +pub fn flatten_log_record( + log_record: &LogRecord, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { let mut log_record_json: Map = Map::new(); log_record_json.insert( "time_unix_nano".to_string(), @@ -195,17 +202,59 @@ pub fn flatten_log_record(log_record: &LogRecord) -> Map { if log_record.body.is_some() { let body = &log_record.body; let body_json = collect_json_from_values(body, "body"); - for (key, value) in &body_json { + // Values that are strings which can be parsed as JSON objects get their fields extracted + let parsed_bodies: Vec> = body_json + .values() + .map(|value| match value { + Value::String(s) => serde_json::from_str::(s) + .ok() + .filter(Value::is_object), + _ => None, + }) + .collect(); + // The rows of every parsed body are merged into this record, but they are all built + // first, so charge all of them before building any + let limits = budget.limits(); + let mut body_cost = FlattenCost::default(); + for (key, parsed) in body_json.keys().zip(&parsed_bodies) { + let Some(parsed @ Value::Object(fields)) = parsed else { + continue; + }; + let cost = flatten_pipeline_cost(parsed, limits)?; + // Every row holds each root field of the body once, and each of them gets a + // copy of the body key, as `{key}_{field}`. Structurally flattening the record + // later (in `push_logs`) copies the key into the name of every field below a + // root field too, as `{key}_{field}_{name}`. Charge the larger of the two. + let prefix = string_json_len(key).map(|quoted| quoted - 2 + "_".len()); + let generic = prefix + .and_then(|prefix| prefix.checked_mul(fields.len())) + .and_then(|bytes| bytes.checked_mul(cost.rows)) + .and_then(|bytes| bytes.checked_add(cost.generic_bytes)); + let structural = prefix + .and_then(|prefix| prefix.checked_mul(cost.structural_fields)) + .and_then(|bytes| bytes.checked_add(cost.structural_bytes)); + let bytes = generic + .zip(structural) + .map(|(generic, structural)| generic.max(structural)) + .ok_or(JsonFlattenError::BytesPerRequestLimitExceeded( + limits.max_bytes_per_request, + ))?; + body_cost = body_cost.add( + FlattenCost { + rows: cost.rows, + bytes, + }, + limits, + )?; + } + budget.charge_event(body_cost)?; + + for ((key, value), parsed) in body_json.iter().zip(&parsed_bodies) { // Always insert the original body field as is log_record_json.insert(key.clone(), value.clone()); - // If value is a string that can be parsed as JSON object, extract its fields - if let Value::String(s) = value - && let Ok(parsed) = serde_json::from_str::(s) - && parsed.is_object() - && let Ok(flattened_values) = generic_flattening(&parsed) - { - for flattened_value in flattened_values { + if let Some(parsed) = parsed { + for flattened_value in expand_generic_flattening(parsed)? { if let Value::Object(flattened_obj) = flattened_value { for (inner_key, inner_value) in flattened_obj { let prefixed_key = format!("{key}_{inner_key}"); @@ -255,12 +304,15 @@ pub fn flatten_log_record(log_record: &LogRecord) -> Map { Value::String(hex::encode(&log_record.trace_id)), ); - log_record_json + Ok(log_record_json) } /// this function flattens the `ScopeLogs` object /// and returns a `Vec` of `Map` of the flattened json -fn flatten_scope_log(scope_log: &ScopeLogs, tenant_id: &str) -> Vec> { +fn flatten_scope_log( + scope_log: &ScopeLogs, + budget: &mut FlattenBudget, +) -> Result>, JsonFlattenError> { let mut vec_scope_log_json = Vec::new(); let mut scope_log_json = Map::new(); if let Some(scope) = &scope_log.scope { @@ -281,16 +333,13 @@ fn flatten_scope_log(scope_log: &ScopeLogs, tenant_id: &str) -> Vec( get_scope_logs: fn(&T) -> &[ScopeLogs], get_schema_url: fn(&T) -> &str, tenant_id: &str, -) -> Vec + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> where T: std::fmt::Debug, { let mut vec_otel_json = Vec::new(); + let mut logs_collected: u64 = 0; for resource_log in resource_logs { let mut resource_log_json = Map::new(); @@ -322,7 +373,8 @@ where let scope_logs = get_scope_logs(resource_log); for scope_log in scope_logs { - vec_resource_logs_json.extend(flatten_scope_log(scope_log, tenant_id)); + vec_resource_logs_json.extend(flatten_scope_log(scope_log, budget)?); + logs_collected += scope_log.log_records.len() as u64; } resource_log_json.insert( @@ -336,36 +388,56 @@ where } } - vec_otel_json + // Only count the logs of a request that was converted in full + let date = chrono::Utc::now().date_naive().to_string(); + increment_logs_collected_by_date(logs_collected, &date, tenant_id); + + Ok(vec_otel_json) } -pub fn flatten_otel_protobuf(message: &ExportLogsServiceRequest, tenant_id: &str) -> Vec { +/// Same as `flatten_otel_logs`, for the protobuf request +pub fn flatten_otel_protobuf( + message: &ExportLogsServiceRequest, + tenant_id: &str, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { process_resource_logs( &message.resource_logs, |record| record.resource.as_ref(), |record| &record.scope_logs, |record| &record.schema_url, tenant_id, + budget, ) } /// this function performs the custom flattening of the otel logs /// and returns a `Vec` of `Value::Object` of the flattened json -pub fn flatten_otel_logs(message: &LogsData, tenant_id: &str) -> Vec { +/// Fails, with nothing returned, once a log body is over the limits of `budget`, which is +/// shared by all records of the request +pub fn flatten_otel_logs( + message: &LogsData, + tenant_id: &str, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { process_resource_logs( &message.resource_logs, |record| record.resource.as_ref(), |record| &record.scope_logs, |record| &record.schema_url, tenant_id, + budget, ) } #[cfg(test)] mod tests { - use super::{flatten_log_record, normalize_severity}; - use opentelemetry_proto::tonic::common::v1::{AnyValue, any_value}; - use opentelemetry_proto::tonic::logs::v1::LogRecord; + use super::{flatten_log_record, flatten_otel_logs, normalize_severity}; + use crate::utils::json::flatten::{ + FlattenBudget, FlattenCost, FlattenLimits, JsonFlattenError, generic_flattening_with_limits, + }; + use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, KeyValueList, any_value}; + use opentelemetry_proto::tonic::logs::v1::{LogRecord, LogsData, ResourceLogs, ScopeLogs}; use serde_json::json; fn flatten_record( @@ -373,14 +445,18 @@ mod tests { text: &str, body: &str, ) -> serde_json::Map { - flatten_log_record(&LogRecord { - severity_number: number, - severity_text: text.to_string(), - body: Some(AnyValue { - value: Some(any_value::Value::StringValue(body.to_string())), - }), - ..Default::default() - }) + flatten_log_record( + &LogRecord { + severity_number: number, + severity_text: text.to_string(), + body: Some(AnyValue { + value: Some(any_value::Value::StringValue(body.to_string())), + }), + ..Default::default() + }, + &mut FlattenBudget::new(FlattenLimits::DEFAULT), + ) + .unwrap() } #[test] @@ -461,4 +537,370 @@ mod tests { fn preserves_provided_severity_fields() { assert_eq!(normalize_severity(18, "error"), (18, "error".to_string())); } + + fn limits(rows_per_event: usize, rows_per_request: usize) -> FlattenBudget { + FlattenBudget::new(FlattenLimits { + max_rows_per_event: rows_per_event, + max_rows_per_request: rows_per_request, + max_bytes_per_request: 1 << 20, + }) + } + + fn string_body(body: &str) -> Option { + Some(AnyValue { + value: Some(any_value::Value::StringValue(body.to_string())), + }) + } + + fn record(body: Option) -> LogRecord { + LogRecord { + severity_number: 9, + body, + ..Default::default() + } + } + + // Two resources, the first with two scopes, holding one record each + fn logs(body: &str) -> LogsData { + let scope = || ScopeLogs { + log_records: vec![record(string_body(body))], + ..Default::default() + }; + LogsData { + resource_logs: vec![ + ResourceLogs { + scope_logs: vec![scope(), scope()], + ..Default::default() + }, + ResourceLogs { + scope_logs: vec![scope()], + ..Default::default() + }, + ], + } + } + + const SIX_ROW_BODY: &str = r#"{"a": [1, 2, 3], "b": [true, false]}"#; + + #[test] + fn json_body_rows_are_merged_into_the_record_and_charged() { + let mut budget = limits(6, 6); + let flattened = + flatten_log_record(&record(string_body(SIX_ROW_BODY)), &mut budget).unwrap(); + // The last row wins, as before + assert_eq!(flattened["body"], json!(SIX_ROW_BODY)); + assert_eq!(flattened["body_a"], json!(3)); + assert_eq!(flattened["body_b"], json!(false)); + assert_eq!(budget.used().rows, 6); + + assert!(matches!( + flatten_log_record(&record(string_body(SIX_ROW_BODY)), &mut limits(5, 6)), + Err(JsonFlattenError::RowsPerEventLimitExceeded(5)) + )); + } + + #[test] + fn all_json_bodies_of_a_record_are_one_event() { + let kv = |key: &str, body: &str| KeyValue { + key: key.to_string(), + value: string_body(body), + ..Default::default() + }; + let body = Some(AnyValue { + value: Some(any_value::Value::KvlistValue(KeyValueList { + values: vec![ + kv("x", r#"{"a": [1, 2, 3]}"#), + kv("plain", "not json"), + kv("y", r#"{"b": [1, 2, 3]}"#), + ], + })), + }); + let mut budget = limits(6, 6); + let flattened = flatten_log_record(&record(body.clone()), &mut budget).unwrap(); + assert_eq!(flattened["x_a"], json!(3)); + assert_eq!(flattened["y_b"], json!(3)); + assert_eq!(flattened["plain"], json!("not json")); + assert_eq!(budget.used().rows, 6); + // Each body fits on its own, but not both together + assert!(matches!( + flatten_log_record(&record(body), &mut limits(5, 6)), + Err(JsonFlattenError::RowsPerEventLimitExceeded(5)) + )); + } + + #[test] + fn bodies_that_are_not_json_objects_are_kept_and_not_charged() { + for body in ["not json", "[1, 2, 3]", "12", r#"{"a": [1, 2"#] { + let mut budget = limits(1, 1); + let flattened = flatten_log_record(&record(string_body(body)), &mut budget).unwrap(); + assert_eq!(flattened["body"], json!(body)); + assert_eq!(budget.used().rows, 0, "body: {body}"); + } + let mut budget = limits(1, 1); + flatten_log_record(&record(None), &mut budget).unwrap(); + assert_eq!(budget.used().rows, 0); + } + + #[test] + fn records_of_all_scopes_and_resources_share_the_request_budget() { + let mut budget = limits(6, 18); + assert_eq!( + flatten_otel_logs(&logs(SIX_ROW_BODY), "tenant", &mut budget) + .unwrap() + .len(), + 3 + ); + assert_eq!(budget.used().rows, 18); + // The third record, in the second resource, goes over the request limit + assert!(matches!( + flatten_otel_logs(&logs(SIX_ROW_BODY), "tenant", &mut limits(6, 17)), + Err(JsonFlattenError::RowsPerRequestLimitExceeded(17)) + )); + // The second record, in the second scope of the first resource, goes over it too + assert!(matches!( + flatten_otel_logs(&logs(SIX_ROW_BODY), "tenant", &mut limits(6, 11)), + Err(JsonFlattenError::RowsPerRequestLimitExceeded(11)) + )); + } + + #[test] + fn logs_collected_metric_counts_only_converted_requests() { + use crate::metrics::TOTAL_LOGS_COLLECTED_BY_DATE; + + // A tenant of its own, as the counter is global + let tenant = "flatten-limit-metric-tenant"; + let collected = || { + let date = chrono::Utc::now().date_naive().to_string(); + TOTAL_LOGS_COLLECTED_BY_DATE + .with_label_values(&["all", &date, tenant]) + .get() + }; + + let before = collected(); + // Fails on the third record, after the scopes of the first resource were converted + assert!(flatten_otel_logs(&logs(SIX_ROW_BODY), tenant, &mut limits(6, 17)).is_err()); + assert_eq!(collected(), before); + + flatten_otel_logs(&logs(SIX_ROW_BODY), tenant, &mut limits(6, 18)).unwrap(); + assert_eq!(collected(), before + 3); + } + + fn byte_limit(bytes: usize) -> FlattenBudget { + FlattenBudget::new(FlattenLimits { + max_rows_per_event: 1000, + max_rows_per_request: 10_000, + max_bytes_per_request: bytes, + }) + } + + fn kv_body(entries: &[(&str, &str)]) -> Option { + Some(AnyValue { + value: Some(any_value::Value::KvlistValue(KeyValueList { + values: entries + .iter() + .map(|(key, body)| KeyValue { + key: key.to_string(), + value: string_body(body), + ..Default::default() + }) + .collect(), + })), + }) + } + + // Compact JSON bytes of the rows of `body`, with their fields renamed as they are merged + // into the record + fn prefixed_rows_len(key: &str, body: &str) -> usize { + let parsed: serde_json::Value = serde_json::from_str(body).unwrap(); + generic_flattening_with_limits(&parsed, FlattenLimits::DEFAULT) + .unwrap() + .iter() + .map(|row| { + let prefixed: serde_json::Map = row + .as_object() + .unwrap() + .iter() + .map(|(field, value)| (format!("{key}_{field}"), value.clone())) + .collect(); + serde_json::to_vec(&prefixed).unwrap().len() + }) + .sum() + } + + // One row of 32 short scalar fields, about 300 bytes before the key is copied into them + fn wide_body() -> String { + let fields: serde_json::Map = + (0..32).map(|i| (format!("f{i}"), json!(i))).collect(); + serde_json::Value::Object(fields).to_string() + } + + #[test] + fn long_body_key_copied_into_every_field_is_charged() { + let body = wide_body(); + let long_key = "k".repeat(128); + // 32 copies of the key alone are over the limit, the parsed body is not + let mut budget = byte_limit(4096); + assert!(matches!( + flatten_log_record( + &record(kv_body(&[(long_key.as_str(), body.as_str())])), + &mut budget + ), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + )); + assert_eq!(budget.used(), FlattenCost::default()); + + let flattened = + flatten_log_record(&record(kv_body(&[("k", body.as_str())])), &mut budget).unwrap(); + assert_eq!(flattened["k_f31"], json!(31)); + assert_eq!(budget.used().rows, 1); + assert_eq!(budget.used().bytes, prefixed_rows_len("k", &body)); + } + + #[test] + fn body_key_prefixes_are_charged_exactly() { + // Two rows of two fields each, under a key with escapes + let key = "k\"\t\u{1}"; + let body = r#"{"a": [1, 2], "b": "x"}"#; + let other = r#"{"c": {"d": [true, null]}}"#; + let body_record = || record(kv_body(&[(key, body), ("plain", "text"), ("o", other)])); + let expected = prefixed_rows_len(key, body) + prefixed_rows_len("o", other); + + let mut budget = byte_limit(expected); + let flattened = flatten_log_record(&body_record(), &mut budget).unwrap(); + assert_eq!(flattened[&format!("{key}_a")], json!(2)); + assert_eq!(flattened["o_c"], json!({"d": null})); + assert_eq!( + budget.used(), + FlattenCost { + rows: 4, + bytes: expected + } + ); + + let mut budget = byte_limit(expected - 1); + assert!(matches!( + flatten_log_record(&body_record(), &mut budget), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(limit)) if limit == expected - 1 + )); + assert_eq!(budget.used(), FlattenCost::default()); + } + + #[test] + fn body_key_prefixes_add_up_over_records() { + let body = wide_body(); + let key = "k".repeat(16); + let per_record = prefixed_rows_len(&key, &body); + let mut budget = byte_limit(2 * per_record - 1); + flatten_log_record( + &record(kv_body(&[(key.as_str(), body.as_str())])), + &mut budget, + ) + .unwrap(); + assert!(matches!( + flatten_log_record( + &record(kv_body(&[(key.as_str(), body.as_str())])), + &mut budget + ), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(_)) + )); + assert_eq!(budget.used().bytes, per_record); + } + + #[test] + fn logs_collected_metric_is_unchanged_when_a_body_key_is_over_the_budget() { + use crate::metrics::TOTAL_LOGS_COLLECTED_BY_DATE; + + let tenant = "flatten-limit-prefix-metric-tenant"; + let collected = || { + let date = chrono::Utc::now().date_naive().to_string(); + TOTAL_LOGS_COLLECTED_BY_DATE + .with_label_values(&["all", &date, tenant]) + .get() + }; + let body = wide_body(); + let long_key = "k".repeat(128); + let scope = |key: &str| ScopeLogs { + log_records: vec![record(kv_body(&[(key, body.as_str())]))], + ..Default::default() + }; + // The first resource fits, the long key in the second one does not + let logs = LogsData { + resource_logs: vec![ + ResourceLogs { + scope_logs: vec![scope("k")], + ..Default::default() + }, + ResourceLogs { + scope_logs: vec![scope(&long_key)], + ..Default::default() + }, + ], + }; + + let before = collected(); + assert!(matches!( + flatten_otel_logs(&logs, tenant, &mut byte_limit(4096)), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + )); + assert_eq!(collected(), before); + } + + // Compact JSON bytes of the rows of `body` merged into the record under `key`, once + // `push_logs` structurally flattens them + fn structural_prefixed_rows_len(key: &str, body: &str) -> usize { + let parsed: serde_json::Value = serde_json::from_str(body).unwrap(); + generic_flattening_with_limits(&parsed, FlattenLimits::DEFAULT) + .unwrap() + .iter() + .map(|row| { + let mut prefixed = serde_json::Value::Object( + row.as_object() + .unwrap() + .iter() + .map(|(field, value)| (format!("{key}_{field}"), value.clone())) + .collect(), + ); + crate::utils::json::flatten::flatten(&mut prefixed, "_", None, None, None, false) + .unwrap(); + serde_json::to_vec(&prefixed).unwrap().len() + }) + .sum() + } + + // 4 rows of the 32 fields of `wide_body` below a parent + fn nested_body() -> String { + format!(r#"{{"p": {}, "r": [0, 1, 2, 3]}}"#, wide_body()) + } + + #[test] + fn body_key_copied_into_nested_field_names_is_charged() { + let body = nested_body(); + let long_key = "k".repeat(128); + // The generic rows and their root field prefixes fit, but structurally flattening + // the record names all 32 fields below `p` after the long key + assert!(prefixed_rows_len(&long_key, &body) <= 4096); + let mut budget = byte_limit(4096); + assert!(matches!( + flatten_log_record( + &record(kv_body(&[(long_key.as_str(), body.as_str())])), + &mut budget + ), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + )); + assert_eq!(budget.used(), FlattenCost::default()); + + // With a short key, the structurally flattened rows are the larger, and charged + let expected = structural_prefixed_rows_len("k", &body); + assert!(expected > prefixed_rows_len("k", &body)); + let flattened = + flatten_log_record(&record(kv_body(&[("k", body.as_str())])), &mut budget).unwrap(); + assert_eq!(flattened["k_p"]["f31"], json!(31)); + assert_eq!( + budget.used(), + FlattenCost { + rows: 4, + bytes: expected + } + ); + } } diff --git a/src/storage/field_stats.rs b/src/storage/field_stats.rs index 3f81d2dbe..141453856 100644 --- a/src/storage/field_stats.rs +++ b/src/storage/field_stats.rs @@ -40,6 +40,8 @@ use crate::utils::create_intracluster_auth_headermap; use crate::utils::get_tenant_id_from_request; use crate::utils::get_user_from_request; use crate::utils::json::apply_generic_flattening_for_partition; +use crate::utils::json::flatten::FlattenBudget; +use crate::utils::json::flatten::FlattenLimits; use actix_web::HttpRequest; use actix_web::HttpResponse; use actix_web::Responder; @@ -174,11 +176,17 @@ async fn calculate_field_stats_inner( let stats_value = serde_json::to_value(&stats).map_err(|e| ObjectStorageError::Invalid(e.into()))?; ensure_dataset_stats_stream(tenant_id).await?; + let limits = field_stats_flatten_limits( + stats.field_stats.len(), + max_field_statistics, + FlattenLimits::from_config().max_bytes_per_request, + ); let vec_json = apply_generic_flattening_for_partition( stats_value, None, None, Some(&DATASET_STATS_CUSTOM_PARTITION.to_string()), + &mut FlattenBudget::new(limits), )?; let mut p_custom_fields = HashMap::new(); p_custom_fields.insert(USER_AGENT_KEY.to_string(), "parseable".to_string()); @@ -210,6 +218,25 @@ async fn calculate_field_stats_inner( Ok(stats_calculated) } +/// Generic flattening limits for one stats payload. It expands into a row per distinct value +/// of each field, and there are at most `max_field_statistics` of those per field, so the +/// row limits follow from its shape instead of the limits for ingested events, which a wide +/// dataset would exceed +fn field_stats_flatten_limits( + field_count: usize, + max_field_statistics: usize, + max_bytes: usize, +) -> FlattenLimits { + let max_rows = field_count + .max(1) + .saturating_mul(max_field_statistics.max(1)); + FlattenLimits { + max_rows_per_event: max_rows, + max_rows_per_request: max_rows, + max_bytes_per_request: max_bytes, + } +} + fn log_stats_calculation_time( stream_name: &str, parquet_path: &Path, @@ -1111,9 +1138,12 @@ mod tests { use temp_dir::TempDir; use crate::storage::field_stats::{ - FieldCountState, NULL_VALUE, build_stats_sql, - collect_all_field_stats_from_parquet_blocking, collect_field_counts, + DATASET_STATS_CUSTOM_PARTITION, DatasetStats, DistinctStat, FieldCountState, FieldStat, + NULL_VALUE, build_stats_sql, collect_all_field_stats_from_parquet_blocking, + collect_field_counts, field_stats_flatten_limits, }; + use crate::utils::json::apply_generic_flattening_for_partition; + use crate::utils::json::flatten::FlattenBudget; async fn create_test_parquet_with_data() -> (TempDir, std::path::PathBuf) { let temp_dir = TempDir::new().unwrap(); @@ -1494,4 +1524,84 @@ mod tests { ROWS as f64 / per_run.as_secs_f64() ); } + + fn dataset_stats(distinct_per_field: &[usize]) -> serde_json::Value { + let field_stats = distinct_per_field + .iter() + .enumerate() + .map(|(field, &distinct)| FieldStat { + field_name: format!("field{field}"), + count: 10, + distinct_count: distinct as i64, + distinct_stats: (0..distinct) + .map(|value| DistinctStat { + distinct_value: format!("value{value}"), + count: 1, + }) + .collect(), + }) + .collect(); + serde_json::to_value(DatasetStats { + dataset_name: "dataset".to_owned(), + stats_id: "stats".to_owned(), + field_stats, + }) + .unwrap() + } + + fn flatten_stats( + stats: serde_json::Value, + field_count: usize, + max_field_statistics: usize, + ) -> anyhow::Result> { + apply_generic_flattening_for_partition( + stats, + None, + None, + Some(&DATASET_STATS_CUSTOM_PARTITION.to_string()), + &mut FlattenBudget::new(field_stats_flatten_limits( + field_count, + max_field_statistics, + 1 << 30, + )), + ) + } + + #[test] + fn widest_stats_payload_fits_its_flatten_limits() { + // Far more rows than the limits for ingested events allow + let fields = 300; + let max_field_statistics = 50; + let rows = flatten_stats( + dataset_stats(&[max_field_statistics; 300]), + fields, + max_field_statistics, + ) + .unwrap(); + assert_eq!(rows.len(), fields * max_field_statistics); + assert_eq!( + rows[0]["field_stats_distinct_stats_distinct_value"], + "value0" + ); + + // A field without distinct values still takes a row + let rows = flatten_stats(dataset_stats(&[0, 2]), 2, 2).unwrap(); + assert_eq!(rows.len(), 3); + let rows = flatten_stats(dataset_stats(&[0]), 1, 0).unwrap(); + assert_eq!(rows.len(), 1); + } + + #[test] + fn stats_payload_flatten_limits_are_bounded() { + assert_eq!(field_stats_flatten_limits(3, 4, 7).max_rows_per_event, 12); + assert_eq!(field_stats_flatten_limits(3, 4, 7).max_rows_per_request, 12); + assert_eq!(field_stats_flatten_limits(3, 4, 7).max_bytes_per_request, 7); + assert_eq!(field_stats_flatten_limits(0, 0, 7).max_rows_per_event, 1); + assert_eq!( + field_stats_flatten_limits(usize::MAX, 2, 7).max_rows_per_event, + usize::MAX + ); + // More distinct values than the payload can hold are rejected + assert!(flatten_stats(dataset_stats(&[3, 2]), 2, 2).is_err()); + } } diff --git a/src/utils/json/flatten.rs b/src/utils/json/flatten.rs index 55b4b9d4c..b7c2609ad 100644 --- a/src/utils/json/flatten.rs +++ b/src/utils/json/flatten.rs @@ -26,6 +26,10 @@ use serde_json::value::Value; use thiserror::Error; +use crate::cli::{ + DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST, DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT, + DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST, Options, +}; use crate::parseable::PARSEABLE; // Global variable to track the first timestamp encountered during validation @@ -59,6 +63,31 @@ pub enum JsonFlattenError { ExpectedObjectInArray, #[error("Found non-object element while flattening array of objects")] NonObjectInArray, + #[error( + "Flattening a single event would produce more than {0} rows (P_MAX_FLATTEN_ROWS_PER_EVENT). Reduce the number or size of nested arrays in the event, or split it into several events" + )] + RowsPerEventLimitExceeded(usize), + #[error( + "Flattening the request would produce more than {0} rows (P_MAX_FLATTEN_ROWS_PER_REQUEST). Split the batch into smaller requests, or reduce the number or size of nested arrays in its events" + )] + RowsPerRequestLimitExceeded(usize), + #[error( + "Flattening the request would produce more than {0} bytes of JSON (P_MAX_FLATTEN_BYTES_PER_REQUEST). Split the batch into smaller requests, or reduce the number or size of nested arrays in its events" + )] + BytesPerRequestLimitExceeded(usize), +} + +impl JsonFlattenError { + /// True for the errors of an event or request that generic flattening would expand + /// beyond the configured `FlattenLimits` + pub fn is_limit_exceeded(&self) -> bool { + matches!( + self, + Self::RowsPerEventLimitExceeded(_) + | Self::RowsPerRequestLimitExceeded(_) + | Self::BytesPerRequestLimitExceeded(_) + ) + } } // Recursively flattens JSON objects and arrays, e.g. with the separator `.`, starting from the TOP @@ -310,87 +339,710 @@ pub fn flatten_array_objects( Ok(()) } +/// Caps on the work generic flattening may do, see `FlattenBudget` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FlattenLimits { + /// Rows a single event may expand into + pub max_rows_per_event: usize, + /// Rows all events of a request may expand into together + pub max_rows_per_request: usize, + /// Compact JSON bytes of all rows a request may expand into together. This measures the + /// expanded JSON, not the memory actually used to hold it + pub max_bytes_per_request: usize, +} + +impl FlattenLimits { + pub const DEFAULT: Self = Self { + max_rows_per_event: DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT, + max_rows_per_request: DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST, + max_bytes_per_request: DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST, + }; + + pub fn from_options(options: &Options) -> Self { + Self { + max_rows_per_event: options.max_flatten_rows_per_event, + max_rows_per_request: options.max_flatten_rows_per_request, + max_bytes_per_request: options.max_flatten_bytes_per_request, + } + } + + /// Limits configured for this server + pub fn from_config() -> Self { + Self::from_options(&PARSEABLE.options) + } + + fn check_event_rows(&self, rows: Option) -> Result { + rows.filter(|&rows| rows <= self.max_rows_per_event).ok_or( + JsonFlattenError::RowsPerEventLimitExceeded(self.max_rows_per_event), + ) + } + + fn check_bytes(&self, bytes: Option) -> Result { + bytes + .filter(|&bytes| bytes <= self.max_bytes_per_request) + .ok_or(JsonFlattenError::BytesPerRequestLimitExceeded( + self.max_bytes_per_request, + )) + } +} + +/// Rows generic flattening produces, and the sum of their compact JSON lengths +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct FlattenCost { + pub rows: usize, + pub bytes: usize, +} + +impl FlattenCost { + /// Adds up the costs of two parts of one event, failing like `generic_flattening_cost` + /// once the total is over the limits + pub fn add(self, other: Self, limits: &FlattenLimits) -> Result { + Ok(Self { + rows: limits.check_event_rows(self.rows.checked_add(other.rows))?, + bytes: limits.check_bytes(self.bytes.checked_add(other.bytes))?, + }) + } +} + +/// Tracks the generic flattening work of one request (or internally generated payload) +/// against its `FlattenLimits`. Create one per request and pass it to every event of it. +#[derive(Debug, Clone)] +pub struct FlattenBudget { + limits: FlattenLimits, + used: FlattenCost, +} + +impl FlattenBudget { + pub fn new(limits: FlattenLimits) -> Self { + Self { + limits, + used: FlattenCost::default(), + } + } + + pub fn limits(&self) -> &FlattenLimits { + &self.limits + } + + pub fn used(&self) -> FlattenCost { + self.used + } + + /// Charges the cost of one event, leaving the budget unchanged when that is over a limit + pub fn charge_event(&mut self, cost: FlattenCost) -> Result<(), JsonFlattenError> { + let limits = &self.limits; + limits.check_event_rows(Some(cost.rows))?; + let rows = self + .used + .rows + .checked_add(cost.rows) + .filter(|&rows| rows <= limits.max_rows_per_request) + .ok_or(JsonFlattenError::RowsPerRequestLimitExceeded( + limits.max_rows_per_request, + ))?; + let bytes = limits.check_bytes(self.used.bytes.checked_add(cost.bytes))?; + self.used = FlattenCost { rows, bytes }; + Ok(()) + } +} + /// Recursively flattens a JSON value. /// - If the value is an array, it flattens all elements of the array. /// - If the value is an object, it flattens all nested objects and arrays. -/// - If the JSON value is heavily nested (with more than 4 levels of hierarchy), returns error /// - Otherwise, it returns the value itself in a vector. /// +/// Every element left after splicing nested top level arrays is one event charged to +/// `budget`, see `charge_generic_flattening`. Nothing is built when a limit is exceeded. +/// Callers skip this for values nested deeper than `P_MAX_FLATTEN_LEVEL`. Callers that +/// structurally flatten the rows afterwards use `pipeline_flattening` instead. +/// /// Examples: /// 1. `{"a": 1}` ~> `[{"a": 1}]` /// 2. `[{"a": 1}, {"b": 2}]` ~> `[{"a": 1}, {"b": 2}]` /// 3. `[{"a": [{"b": 1}, {"c": 2}]}]` ~> `[{"a": {"b": 1)}}, {"a": {"c": 2)}}]` /// 4. `{"a": [{"b": 1}, {"c": 2}], "d": {"e": 4}}` ~> `[{"a": {"b":1}, "d": {"e":4}}, {"a": {"c":2}, "d": {"e":4}}]` -/// 5. `{"a":{"b":{"c":{"d":{"e":["a","b"]}}}}}` ~> returns error - heavily nested, cannot flatten this JSON -pub fn generic_flattening(value: &Value) -> Result, JsonFlattenError> { +pub fn generic_flattening( + value: &Value, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { + flattening::(value, budget) +} + +/// `generic_flattening` with a budget of its own +pub fn generic_flattening_with_limits( + value: &Value, + limits: FlattenLimits, +) -> Result, JsonFlattenError> { + generic_flattening(value, &mut FlattenBudget::new(limits)) +} + +/// Charges `budget` for `generic_flattening(value)`, treating every element left after +/// splicing nested top level arrays as one event +pub fn charge_generic_flattening( + value: &Value, + budget: &mut FlattenBudget, +) -> Result<(), JsonFlattenError> { + charge_events::(value, budget) +} + +/// Counts the rows `generic_flattening` builds for `value` taken as a single event, and +/// the sum of their compact JSON lengths, without building or cloning anything. Fails as +/// soon as the rows are over `max_rows_per_event` or the bytes over `max_bytes_per_request`, +/// including when counting them would overflow. +pub fn generic_flattening_cost( + value: &Value, + limits: &FlattenLimits, +) -> Result { + let expansion = value_cost::(value, false, limits)?; + Ok(FlattenCost { + rows: expansion.rows, + bytes: expansion.bytes, + }) +} + +/// Builds the rows of generic flattening without charging any budget. Only call it for a +/// value whose `generic_flattening_cost` (or `flatten_pipeline_cost`) was charged first. +pub(crate) fn expand_generic_flattening(value: &Value) -> Result, JsonFlattenError> { generic_flattening_inner(value, false) } +/// `generic_flattening(&value, budget)` that consumes `value`, charged to `budget` the same +/// way. An object, or an array of objects, with no arrays nested anywhere inside is left +/// unchanged by generic flattening, so it is returned as is instead of being copied, and +/// checked and charged in a single walk. +pub fn generic_flattening_owned( + value: Value, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { + flattening_owned::(value, budget) +} + +/// What generic flattening followed by structural flattening (`flatten` with the `_` +/// separator, which every server caller of generic flattening runs on its rows) costs for +/// one event, see `flatten_pipeline_cost` +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct PipelineCost { + /// Rows generic flattening builds. Structural flattening keeps them one for one. + pub rows: usize, + /// Compact JSON bytes of the generic rows, exactly as `generic_flattening_cost` counts them + pub generic_bytes: usize, + /// Compact JSON bytes of the generic rows once structurally flattened. Every scalar of a + /// row is counted as a field of its own, under its joined name, even when that name + /// collides with another field and overwrites it, so collisions make this an overcount. + pub structural_bytes: usize, + /// Fields of the structurally flattened rows that are objects, collisions included + pub structural_fields: usize, +} + +impl PipelineCost { + /// What to charge for the event: its rows, with the larger of the bytes of its two + /// representations + pub fn charged(&self) -> FlattenCost { + FlattenCost { + rows: self.rows, + bytes: self.generic_bytes.max(self.structural_bytes), + } + } +} + +/// `generic_flattening_cost` that also counts, in the same walk, the rows structurally +/// flattened afterwards, whose field names repeat the names of their parents. Never builds +/// a row or a field name. Fails like `generic_flattening_cost`, and also as soon as the +/// structurally flattened bytes are over `max_bytes_per_request`. +pub fn flatten_pipeline_cost( + value: &Value, + limits: &FlattenLimits, +) -> Result { + let expansion = value_cost::(value, false, limits)?; + Ok(PipelineCost { + rows: expansion.rows, + generic_bytes: expansion.bytes, + structural_bytes: limits.check_bytes(expansion.structural_bytes())?, + structural_fields: expansion.keyed_fields, + }) +} + +/// `charge_generic_flattening` charging every event `flatten_pipeline_cost(event).charged()` +pub fn charge_flatten_pipeline( + value: &Value, + budget: &mut FlattenBudget, +) -> Result<(), JsonFlattenError> { + charge_events::(value, budget) +} + +/// `generic_flattening` for callers that structurally flatten the rows afterwards: the same +/// rows, charged with `charge_flatten_pipeline` +pub fn pipeline_flattening( + value: &Value, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { + flattening::(value, budget) +} + +/// `generic_flattening_owned` for callers that structurally flatten the rows afterwards: +/// the same rows, charged with `charge_flatten_pipeline`. The single walk of the array-free +/// fast path counts the structurally flattened bytes as well. +pub fn pipeline_flattening_owned( + value: Value, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { + flattening_owned::(value, budget) +} + +// `STRUCTURAL` selects the charge of `charge_flatten_pipeline` over the one of +// `charge_generic_flattening`, here and below +fn flattening( + value: &Value, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { + charge_events::(value, budget)?; + expand_generic_flattening(value) +} + +fn charge_events( + value: &Value, + budget: &mut FlattenBudget, +) -> Result<(), JsonFlattenError> { + match value { + Value::Array(arr) => arr + .iter() + .try_for_each(|item| charge_events::(item, budget)), + event => { + let cost = if STRUCTURAL { + flatten_pipeline_cost(event, budget.limits())?.charged() + } else { + generic_flattening_cost(event, budget.limits())? + }; + budget.charge_event(cost) + } + } +} + +fn flattening_owned( + value: Value, + budget: &mut FlattenBudget, +) -> Result, JsonFlattenError> { + match value { + Value::Object(map) => match array_free_charge::(&map) { + Some(len) => { + charge_array_free(len, budget)?; + Ok(vec![Value::Object(map)]) + } + None => flattening::(&Value::Object(map), budget), + }, + Value::Array(arr) => { + for (index, item) in arr.iter().enumerate() { + let len = match item { + Value::Object(map) => array_free_charge::(map), + _ => None, + }; + let Some(len) = len else { + // The events before this one are charged already + for item in &arr[index..] { + charge_events::(item, budget)?; + } + return expand_generic_flattening(&Value::Array(arr)); + }; + charge_array_free(len, budget)?; + } + Ok(arr) + } + value => flattening::(&value, budget), + } +} + +// Charges an array-free object of `len` bytes, failing the way `charge_events` does for +// it: its cost walk only fails on the bytes of such an object (with a positive row limit), +// and only once they are over the limit +fn charge_array_free(len: usize, budget: &mut FlattenBudget) -> Result<(), JsonFlattenError> { + budget.limits.check_bytes(Some(len))?; + budget.charge_event(FlattenCost { + rows: 1, + bytes: len, + }) +} + +// Bytes `charge_events` charges for an array-free object: its compact JSON length, or the +// larger of that and its structurally flattened length. `None` when an array is nested +// anywhere inside it (or a length overflows), leaving it to the full cost walk. +fn array_free_charge(map: &Map) -> Option { + let cost = array_free_cost::(map)?; + if !STRUCTURAL { + return Some(cost.json_len); + } + // A single row that is an object, see `Expansion::structural_bytes` + let structural = if cost.fields == 0 { + "{}".len() + } else { + cost.fields + .checked_mul(4)? + .checked_add(1)? + .checked_add(cost.field_bytes)? + }; + Some(cost.json_len.max(structural)) +} + +// An array-free object, and the fields structural flattening makes of it, counted like +// `Expansion::fields` and `Expansion::field_bytes` (only when `STRUCTURAL`) +struct ArrayFree { + json_len: usize, + fields: usize, + field_bytes: usize, +} + +fn array_free_cost(map: &Map) -> Option { + // `{`, `}`, the commas and a `:` per key. A map entry takes more memory than that, so + // this cannot overflow + let mut len = 2 + map.len().saturating_sub(1) + map.len(); + let mut fields = 0usize; + let mut field_bytes = 0usize; + for (key, value) in map { + let key_len = string_json_len(key)?; + let value_len = match value { + Value::Array(_) => return None, + Value::Object(nested) => { + let nested = array_free_cost::(nested)?; + if STRUCTURAL { + // Every field below is named `{key}_{name}` + fields = fields.checked_add(nested.fields)?; + field_bytes = nested + .fields + .checked_mul(key_len - 2 + 1)? + .checked_add(nested.field_bytes)? + .checked_add(field_bytes)?; + } + nested.json_len + } + scalar => { + let value_len = scalar_json_len(scalar)?; + if STRUCTURAL { + fields = fields.checked_add(1)?; + field_bytes = field_bytes + .checked_add(key_len - 2)? + .checked_add(value_len)?; + } + value_len + } + }; + len = len.checked_add(key_len)?.checked_add(value_len)?; + } + Some(ArrayFree { + json_len: len, + fields, + field_bytes, + }) +} + +// The rows a value expands into under generic flattening: as alternatives for the key +// holding it when inside an object, or as the rows of an event at the top. Besides the rows +// and their compact JSON bytes, the walk counts, when `STRUCTURAL`, what structural +// flattening makes of those rows: a field per scalar in a row, named after the keys on the +// path to it joined with `_`, with empty objects dropped. Field names are counted without +// building them: a name below a key is that key's escaped content, a `_` and the name +// below it, and escaping a joined name escapes each part on its own. +#[derive(Clone, Copy, Default)] +struct Expansion { + rows: usize, + bytes: usize, + // The fields below are only counted when `STRUCTURAL` + /// Scalars over all rows, each a field once structurally flattened + fields: usize, + /// Those of them inside an object, whose name gets a `_` and the name of the key holding + /// the value; the others are scalar rows of their own + keyed_fields: usize, + /// Escaped name content (below the key holding the value) and value bytes of all fields + field_bytes: usize, + /// Rows with no field, which structural flattening leaves as `{}` + empty_rows: usize, +} + +impl Expansion { + fn scalar(len: usize) -> Self { + Self { + rows: 1, + bytes: len, + fields: 1, + keyed_fields: 0, + field_bytes: len, + empty_rows: 0, + } + } + + // All rows of both, as the items of an array + fn sum( + self, + other: Self, + limits: &FlattenLimits, + ) -> Result { + let rows = limits.check_event_rows(self.rows.checked_add(other.rows))?; + let bytes = limits.check_bytes(self.bytes.checked_add(other.bytes))?; + if !STRUCTURAL { + return Ok(Self { + rows, + bytes, + ..Self::default() + }); + } + // Every field takes a byte at least for its value, so `field_bytes` bounds the + // field counts, and rows bound the empty rows + let field_bytes = limits.check_bytes(self.field_bytes.checked_add(other.field_bytes))?; + Ok(Self { + rows, + bytes, + fields: self.fields + other.fields, + keyed_fields: self.keyed_fields + other.keyed_fields, + field_bytes, + empty_rows: self.empty_rows + other.empty_rows, + }) + } + + // Compact JSON bytes of the rows once structurally flattened, with a `{}` around each + // row that is an object, a `,` between its fields, and `"` `"` `:` around each name. + // A row that is a scalar is its only, unkeyed field, and stays as it is. + fn structural_bytes(&self) -> Option { + let scalar_rows = self.fields.checked_sub(self.keyed_fields)?; + let object_rows = self.rows.checked_sub(scalar_rows)?; + // `{}` and the commas of the rows with fields, `{}` of the empty ones + object_rows + .checked_add(self.empty_rows)? + .checked_add(self.keyed_fields.checked_mul(4)?)? + .checked_add(self.field_bytes) + } +} + +// Mirrors `generic_flattening_inner`. No branch there yields an empty set of rows for an +// object or for a value inside one, so every partial count below is a lower bound of the +// final one and going over a limit (or overflowing) on the way is final. Partial structural +// bytes are lower bounds of the final structural bytes too, which are charged when larger +// than the generic ones. +fn value_cost( + value: &Value, + normalize_empty_array: bool, + limits: &FlattenLimits, +) -> Result { + match value { + Value::Array(arr) if arr.is_empty() && normalize_empty_array => { + Ok(Expansion::scalar(NULL_JSON_LEN)) + } + Value::Array(arr) => arr.iter().try_fold(Expansion::default(), |total, item| { + total.sum::( + value_cost::(item, normalize_empty_array, limits)?, + limits, + ) + }), + Value::Object(map) => object_cost::(map, limits), + scalar => Ok(Expansion::scalar( + limits.check_bytes(scalar_json_len(scalar))?, + )), + } +} + +// Every row of an object holds each of its keys once, with one of the alternatives that +// key's value expands into, so the rows are the product of the alternatives per key +fn object_cost( + map: &Map, + limits: &FlattenLimits, +) -> Result { + // `{`, `}`, the commas and every `"key":` + let mut row_overhead = Some(2 + map.len().saturating_sub(1)); + let mut rows = 1usize; + // Bytes of the values in all rows so far + let mut value_bytes = 0usize; + // Structurally flattened fields of all rows so far, see `Expansion` + let mut fields = 0usize; + let mut field_bytes = 0usize; + let mut empty_rows = 1usize; + for (key, value) in map { + let key_len = string_json_len(key); + row_overhead = row_overhead + .zip(key_len) + .and_then(|(overhead, key)| overhead.checked_add(key)?.checked_add(1)); + let alternatives = match value { + // Arrays use the normalization of array items, which objects ignore + Value::Array(_) | Value::Object(_) => value_cost::(value, true, limits)?, + scalar => Expansion::scalar(limits.check_bytes(scalar_json_len(scalar))?), + }; + let previous_rows = rows; + if matches!(value, Value::Array(_) | Value::Object(_)) { + // Each row so far is repeated for every alternative, and each alternative + // appears once per row so far + value_bytes = limits.check_bytes( + value_bytes + .checked_mul(alternatives.rows) + .zip(alternatives.bytes.checked_mul(rows)) + .and_then(|(old, new)| old.checked_add(new)), + )?; + rows = limits.check_event_rows(rows.checked_mul(alternatives.rows))?; + } else { + // A scalar is the only alternative, so it appears once in each row + value_bytes = limits.check_bytes( + alternatives + .bytes + .checked_mul(rows) + .and_then(|bytes| value_bytes.checked_add(bytes)), + )?; + } + if STRUCTURAL { + // Under this key a field that is the value itself is named `{key}`, a field + // below it `{key}_{name}` + let placed = key_len + .and_then(|len| alternatives.fields.checked_mul(len - 2)) + .and_then(|names| names.checked_add(alternatives.keyed_fields)) + .and_then(|names| names.checked_add(alternatives.field_bytes)); + // Repeated like the value bytes above + field_bytes = limits.check_bytes( + field_bytes + .checked_mul(alternatives.rows) + .zip(placed.and_then(|placed| placed.checked_mul(previous_rows))) + .and_then(|(old, new)| old.checked_add(new)), + )?; + // Every field takes a byte at least for its value, so these are no more than + // the `field_bytes` just checked, and the empty rows no more than the rows + fields = fields * alternatives.rows + alternatives.fields * previous_rows; + empty_rows *= alternatives.empty_rows; + } + } + let bytes = row_overhead + .and_then(|overhead| rows.checked_mul(overhead)) + .and_then(|overhead| overhead.checked_add(value_bytes)); + Ok(Expansion { + rows, + bytes: limits.check_bytes(bytes)?, + fields, + keyed_fields: fields, + field_bytes, + empty_rows, + }) +} + +const NULL_JSON_LEN: usize = "null".len(); + +// Length of a scalar serialized as compact JSON +fn scalar_json_len(value: &Value) -> Option { + match value { + Value::Null => Some(NULL_JSON_LEN), + Value::Bool(true) => Some("true".len()), + Value::Bool(false) => Some("false".len()), + Value::String(s) => string_json_len(s), + Value::Number(number) => Some(number_json_len(number)), + Value::Array(_) | Value::Object(_) => None, + } +} + +// Length of a number serialized as compact JSON. Integers are written as plain decimal +// digits; floats are counted by formatting them the way serde_json does +fn number_json_len(number: &serde_json::Number) -> usize { + let digits = |n: u64| n.checked_ilog10().map_or(1, |log| log as usize + 1); + if let Some(n) = number.as_u64() { + digits(n) + } else if let Some(n) = number.as_i64() { + // Only negative integers get here, as `as_u64` takes all others + 1 + digits(n.unsigned_abs()) + } else { + let mut counter = ByteCounter(0); + // Writing to a counter cannot fail, and a finite float is a few dozen bytes + serde_json::to_writer(&mut counter, number).expect("counting never fails"); + counter.0 + } +} + +/// Length of `s` serialized as a JSON string, quotes included, escaped the way serde_json +/// does it. `None` only when the length overflows. +pub(crate) fn string_json_len(s: &str) -> Option { + let quoted = s.len().checked_add(2)?; + // Escapes are rare, and this check (unlike the count below) vectorizes + let needs_escape = |byte: u8| (byte < 0x20) | (byte == b'"') | (byte == b'\\'); + let escaped = s.as_bytes().chunks(64).any(|chunk| { + chunk + .iter() + .fold(false, |found, &byte| found | needs_escape(byte)) + }); + if !escaped { + return Some(quoted); + } + s.bytes().try_fold(quoted, |len, byte| { + len.checked_add(match byte { + b'"' | b'\\' | b'\x08' | b'\x0c' | b'\n' | b'\r' | b'\t' => 1, + 0x00..=0x1f => 5, + _ => 0, + }) + }) +} + +struct ByteCounter(usize); + +impl std::io::Write for ByteCounter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0 += buf.len(); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + fn generic_flattening_inner( value: &Value, normalize_empty_array: bool, ) -> Result, JsonFlattenError> { match value { - // Without arrays the fold below rebuilds an exact copy of the object + // Without arrays the loop below rebuilds an exact copy of the object Value::Object(map) if !map_contains_array(map) => Ok(vec![value.clone()]), Value::Array(arr) if arr.is_empty() && normalize_empty_array => Ok(vec![Value::Null]), - Value::Array(arr) => Ok(arr - .iter() - .flat_map(|flatten_item| { - generic_flattening_inner(flatten_item, normalize_empty_array).unwrap_or_default() - }) - .collect()), + Value::Array(arr) => { + let mut flattened = Vec::new(); + for item in arr { + flattened.extend(generic_flattening_inner(item, normalize_empty_array)?); + } + Ok(flattened) + } Value::Object(map) => { - let results = map - .iter() - .fold(vec![Map::new()], |results, (key, val)| match val { - Value::Array(arr) => { - if arr.is_empty() { - // Generic flattening expands non-empty arrays into scalar rows. - // Keep empty arrays consistent with that representation: they - // contain no concrete value and must not create `List(Null)` - // schema fields. - results - .into_iter() - .map(|mut result| { - result.insert(key.clone(), Value::Null); - result - }) - .collect() - } else { - arr.iter() - .flat_map(|flatten_item| { - generic_flattening_inner(flatten_item, true).unwrap_or_default() - }) - .flat_map(|flattened_item| { - results.iter().map(move |result| { - let mut new_obj = result.clone(); - new_obj.insert(key.clone(), flattened_item.clone()); - new_obj - }) - }) - .collect() - } + let mut results = vec![Map::new()]; + for (key, val) in map { + // A non-empty array expands into the rows of its items, an object into its + // own rows; every row so far is repeated for each of them + let items = match val { + // Generic flattening expands non-empty arrays into scalar rows. Keep + // empty arrays consistent with that representation: they contain no + // concrete value and must not create `List(Null)` schema fields. + Value::Array(arr) if arr.is_empty() => { + set_in_rows(&mut results, key, Value::Null); + continue; } - Value::Object(_) => generic_flattening(val) - .unwrap_or_default() - .iter() - .flat_map(|nested_result| { - results.iter().map(move |result| { - let mut new_obj = result.clone(); - new_obj.insert(key.clone(), nested_result.clone()); - new_obj - }) - }) - .collect(), - _ => results - .into_iter() - .map(|mut result| { - result.insert(key.clone(), val.clone()); - result - }) - .collect(), - }); + Value::Array(arr) => arr.as_slice(), + Value::Object(_) => std::slice::from_ref(val), + _ => { + set_in_rows(&mut results, key, val.clone()); + continue; + } + }; + let Some((first, rest)) = items.split_first() else { + continue; + }; + let mut alternatives = generic_flattening_inner(first, true)?; + if rest.is_empty() && alternatives.len() == 1 { + // The only alternative goes into the rows as they are + set_in_rows( + &mut results, + key, + alternatives.pop().expect("one alternative"), + ); + continue; + } + // Items are expanded one at a time rather than all up front + let mut expanded = Vec::new(); + repeat_rows(&mut expanded, &results, key, alternatives); + for item in rest { + repeat_rows( + &mut expanded, + &results, + key, + generic_flattening_inner(item, true)?, + ); + } + results = expanded; + } Ok(results.into_iter().map(Value::Object).collect()) } @@ -398,8 +1050,42 @@ fn generic_flattening_inner( } } +// Sets `key` to `value` in every row +fn set_in_rows(rows: &mut [Map], key: &str, value: Value) { + if let Some((last, rest)) = rows.split_last_mut() { + for row in rest { + row.insert(key.to_owned(), value.clone()); + } + last.insert(key.to_owned(), value); + } +} + +// Appends to `expanded`, for each alternative in turn, a copy of every row in `rows` +// with `key` set to that alternative +fn repeat_rows( + expanded: &mut Vec>, + rows: &[Map], + key: &str, + alternatives: Vec, +) { + let Some((last, rest)) = rows.split_last() else { + return; + }; + expanded.reserve(alternatives.len().saturating_mul(rows.len())); + for alternative in alternatives { + for row in rest { + let mut row = row.clone(); + row.insert(key.to_owned(), alternative.clone()); + expanded.push(row); + } + let mut row = last.clone(); + row.insert(key.to_owned(), alternative); + expanded.push(row); + } +} + // Returns true if any value nested anywhere below this map is an array -pub(crate) fn map_contains_array(map: &Map) -> bool { +fn map_contains_array(map: &Map) -> bool { map.values().any(|value| match value { Value::Array(_) => true, Value::Object(nested) => map_contains_array(nested), @@ -407,45 +1093,27 @@ pub(crate) fn map_contains_array(map: &Map) -> bool { }) } -/// Returns `Ok` with the result of `convert_to_array(generic_flattening(&value)?)` when -/// generic flattening would leave every record unchanged, consuming `value` instead of -/// cloning it. That holds for an object, or an array of objects, with no arrays nested -/// anywhere inside. Returns `Err` with `value` untouched for every other shape, which -/// must take the generic path (arrays expand into rows, nested top-level arrays are -/// spliced, empty arrays become null and non-objects are errors). -pub fn generic_flattening_noop(value: Value) -> Result { - match value { - Value::Object(map) if !map_contains_array(&map) => { - Ok(Value::Array(vec![Value::Object(map)])) - } - Value::Array(arr) - if arr.iter().all(|item| match item { - Value::Object(map) => !map_contains_array(map), - _ => false, - }) => - { - Ok(Value::Array(arr)) - } - other => Err(other), - } -} - /// recursively checks the level of nesting for the serde Value /// if Value has more than configured `P_MAX_FLATTEN_LEVEL` levels of hierarchy, returns true /// example - if `P_MAX_FLATTEN_LEVEL` is 4, and the JSON is /// 1. `{"a":{"b":{"c":{"d":{"e":["a","b"]}}}}}` ~> returns true /// 2. `{"a": [{"b": 1}, {"c": 2}], "d": {"e": 4}}` ~> returns false pub fn has_more_than_max_allowed_levels(value: &Value, current_level: usize) -> bool { - if current_level > PARSEABLE.options.event_flatten_level { + has_more_than_levels(value, current_level, PARSEABLE.options.event_flatten_level) +} + +/// `has_more_than_max_allowed_levels` with the maximum level given explicitly +pub fn has_more_than_levels(value: &Value, current_level: usize, max_level: usize) -> bool { + if current_level > max_level { return true; } match value { Value::Array(arr) => arr .iter() - .any(|item| has_more_than_max_allowed_levels(item, current_level)), + .any(|item| has_more_than_levels(item, current_level, max_level)), Value::Object(map) => map .values() - .any(|val| has_more_than_max_allowed_levels(val, current_level + 1)), + .any(|val| has_more_than_levels(val, current_level + 1, max_level)), _ => false, } } @@ -459,6 +1127,15 @@ pub fn convert_to_array(flattened: Vec) -> Result Result, JsonFlattenError>; + type Borrowed = fn(&Value, &mut FlattenBudget) -> Result, JsonFlattenError>; + #[test] - fn generic_flattening_noop_matches_legacy() { - let mut fast = 0; + fn generic_flattening_owned_matches_generic_flattening() { + let modes: [(Owned, Borrowed); 2] = [ + (generic_flattening_owned, generic_flattening), + (pipeline_flattening_owned, pipeline_flattening), + ]; + for (owned_flattening, flattening) in modes { + owned_matches_borrowed(owned_flattening, flattening); + } + } + + fn owned_matches_borrowed(owned_flattening: Owned, flattening: Borrowed) { + let tight = [ + GENEROUS_LIMITS, + limits(1, 1, 1 << 20), + limits(2, 3, 1 << 20), + limits(4, 1 << 20, 64), + limits(1 << 20, 1 << 20, 16), + limits(3, 5, 40), + ]; + let mut unchanged = 0; for value in corpus() { - let expected = super::legacy::generic_flattening(&value) - .and_then(super::convert_to_array) - .map(|v| serde_json::to_string(&v).unwrap()); - match super::generic_flattening_noop(value.clone()) { - Ok(result) => { - fast += 1; + for limits in tight { + // Also start from a partly used budget, so that the request limits matter + for used in [0, 1] { + let mut owned_budget = FlattenBudget::new(limits); + let mut budget = FlattenBudget::new(limits); + for _ in 0..used { + owned_budget + .charge_event(FlattenCost { rows: 1, bytes: 8 }) + .unwrap(); + budget + .charge_event(FlattenCost { rows: 1, bytes: 8 }) + .unwrap(); + } + let owned = owned_flattening(value.clone(), &mut owned_budget); + if limits == GENEROUS_LIMITS + && used == 0 + && owned.as_ref().is_ok_and(|rows| match &value { + Value::Array(arr) => rows == arr, + object => rows.as_slice() == std::slice::from_ref(object), + }) + { + unchanged += 1; + } assert_eq!( - Ok(serde_json::to_string(&result).unwrap()), - expected.map_err(|e| e.to_string()), - "input: {value}" + ordered(owned), + ordered(flattening(&value, &mut budget)), + "input: {value} limits: {limits:?}" ); + assert_eq!(owned_budget.used(), budget.used(), "input: {value}"); } - Err(returned) => assert_eq!(returned, value), } } - assert!(fast > 1000, "fast path exercised only {fast} times"); + assert!( + unchanged > 1000, + "unchanged fast path exercised only {unchanged} times" + ); + } + + #[test] + fn generic_flattening_owned_keeps_or_expands_shapes() { + for (value, expected) in [ + (json!({"a": {"b": {}}}), json!([{"a": {"b": {}}}])), + (json!([]), json!([])), + (json!([{"a": 1}, {"b": "x"}]), json!([{"a": 1}, {"b": "x"}])), + (json!({"a": []}), json!([{"a": null}])), + ( + json!({"a": {"b": [1, 2]}}), + json!([{"a": {"b": 1}}, {"a": {"b": 2}}]), + ), + (json!([{"a": 1}, [{"b": 2}]]), json!([{"a": 1}, {"b": 2}])), + ( + json!([{"a": 1}, {"b": [1, 2]}]), + json!([{"a": 1}, {"b": 1}, {"b": 2}]), + ), + (json!([{"a": 1}, 1]), json!([{"a": 1}, 1])), + (json!(1), json!([1])), + (json!(null), json!([null])), + ] { + let rows = + generic_flattening_owned(value.clone(), &mut FlattenBudget::new(GENEROUS_LIMITS)) + .unwrap(); + assert_eq!(Value::Array(rows), expected, "input: {value}"); + } + } + + #[test] + fn scalar_json_len_matches_serde_json() { + let mut strings: Vec = (0u8..=0x7f).map(|b| char::from(b).to_string()).collect(); + strings.extend( + [ + "", + "é", + "日本", + "\u{10ffff}", + "\u{2028}", + "a\"b\\c", + "\u{7f}", + ] + .map(str::to_owned), + ); + // Escapes on either side of the 64 byte chunks the fast check scans + for len in [63, 64, 65, 127, 128, 129] { + strings.push("x".repeat(len)); + for at in [0, len / 2, len - 1] { + let mut s = "x".repeat(len); + s.replace_range(at..at + 1, "\n"); + strings.push(s.clone()); + s.replace_range(at..at + 1, "\u{1}"); + strings.push(s); + } + } + let mut values: Vec = strings.into_iter().map(Value::String).collect(); + values.extend( + [ + 0, + 1, + 9, + 10, + 99, + 100, + 999_999, + 1_000_000, + u64::MAX / 10, + u64::MAX, + ] + .map(|n| json!(n)), + ); + values.extend([-1, -9, -10, -99, -100, i64::MIN + 1, i64::MIN, i64::MAX].map(|n| json!(n))); + values.extend( + [ + 0.0, + -0.0, + 1.5, + -1.5, + 0.1, + 1e15, + 1e16, + 1e21, + 1e22, + 1e-7, + 123_456_789.125, + 1.0 / 3.0, + 0.1 + 0.2, + f64::MAX, + f64::MIN, + f64::MIN_POSITIVE, + 5e-324, + -5e-324, + 9_007_199_254_740_992.0, + ] + .map(|n| json!(n)), + ); + values.extend([json!(true), json!(false), Value::Null]); + for value in values { + assert_eq!( + super::scalar_json_len(&value), + Some(serde_json::to_vec(&value).unwrap().len()), + "value: {value}" + ); + if let Value::String(s) = &value { + assert_eq!( + super::string_json_len(s), + Some(serde_json::to_vec(s).unwrap().len()) + ); + } + } } #[test] @@ -959,25 +1792,6 @@ mod tests { )); } - #[test] - fn generic_flattening_noop_rejects_expanding_shapes() { - for value in [ - json!({"a": []}), - json!({"a": {"b": [1]}}), - json!([[{"a": 1}]]), - json!([{"a": 1}, 1]), - json!(1), - json!(null), - ] { - assert_eq!(super::generic_flattening_noop(value.clone()), Err(value)); - } - assert_eq!( - super::generic_flattening_noop(json!({"a": {"b": {}}})), - Ok(json!([{"a": {"b": {}}}])) - ); - assert_eq!(super::generic_flattening_noop(json!([])), Ok(json!([]))); - } - #[test] fn flatten_single_key_string() { let mut obj = json!({"key": "value"}); @@ -1236,7 +2050,10 @@ mod tests { fn flatten_json() { let value = json!({"a":{"b":{"e":["a","b"]}}}); let expected = vec![json!({"a":{"b":{"e":"a"}}}), json!({"a":{"b":{"e":"b"}}})]; - assert_eq!(generic_flattening(&value).unwrap(), expected); + assert_eq!( + generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap(), + expected + ); } #[test] @@ -1244,7 +2061,10 @@ mod tests { let value = json!({"a":{"b":{"e":[]}}}); let expected = vec![json!({"a":{"b":{"e":null}}})]; - assert_eq!(generic_flattening(&value).unwrap(), expected); + assert_eq!( + generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap(), + expected + ); } #[test] @@ -1252,7 +2072,7 @@ mod tests { let value = json!({"a": [[]]}); let expected = vec![json!({"a": null})]; - let flattened = generic_flattening(&value).unwrap(); + let flattened = generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap(); assert!(flattened.iter().all(Value::is_object)); assert_eq!(flattened, expected); } @@ -1261,6 +2081,507 @@ mod tests { fn generic_flattening_preserves_empty_top_level_array() { let value = json!([]); - assert_eq!(generic_flattening(&value).unwrap(), Vec::::new()); + assert_eq!( + generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap(), + Vec::::new() + ); + } + + fn limits(rows_per_event: usize, rows_per_request: usize, bytes: usize) -> FlattenLimits { + FlattenLimits { + max_rows_per_event: rows_per_event, + max_rows_per_request: rows_per_request, + max_bytes_per_request: bytes, + } + } + + fn rows_cost(rows: &[Value]) -> FlattenCost { + FlattenCost { + rows: rows.len(), + bytes: rows + .iter() + .map(|row| serde_json::to_vec(row).unwrap().len()) + .sum(), + } + } + + // Splits a value into events the way `charge_generic_flattening` does + fn events(value: &Value) -> Vec<&Value> { + match value { + Value::Array(arr) => arr.iter().flat_map(events).collect(), + event => vec![event], + } + } + + fn cube(len: usize) -> Value { + let arr: Vec = (0..len).map(|i| json!(i)).collect(); + json!({"a": arr, "b": arr, "c": arr}) + } + + #[test] + fn preflight_cost_matches_built_rows() { + let mut values = corpus(); + values.extend([ + json!({"k\"": "a\nb\u{1}\\\u{7f}é\t\r\u{8}\u{c}/", "a": [1, 2]}), + json!({"a": [1.5, -0.0, 1e300, u64::MAX, i64::MIN], "b": [true, false, null]}), + json!({"a": [[[[1, 2], [3]], [4]]], "b": [[], [[]]]}), + json!({"a": [1, 2, 3], "b": {"c": [1, 2], "d": [{"e": [1, 2]}, {"e": [3]}]}}), + json!([[{"a": [1, 2]}], [], [[]], {"b": [{"c": []}, null]}]), + json!({"a": {"b": {"c": {}}}, "": [""]}), + ]); + for value in values { + let built = generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap(); + let mut budget = FlattenBudget::new(GENEROUS_LIMITS); + charge_generic_flattening(&value, &mut budget).unwrap(); + assert_eq!(budget.used(), rows_cost(&built), "input: {value}"); + for event in events(&value) { + let built = generic_flattening_with_limits(event, GENEROUS_LIMITS).unwrap(); + assert_eq!( + generic_flattening_cost(event, &GENEROUS_LIMITS).unwrap(), + rows_cost(&built), + "event: {event}" + ); + } + } + } + + #[test] + fn preflight_cost_of_empty_null_and_scalar_values() { + let cases = [ + (json!({}), 1, "{}"), + (json!({"a": []}), 1, r#"{"a":null}"#), + (json!({"a": [[]]}), 1, r#"{"a":null}"#), + (json!({"a": [null]}), 1, r#"{"a":null}"#), + (json!({"a": null}), 1, r#"{"a":null}"#), + (json!(null), 1, "null"), + (json!("s"), 1, r#""s""#), + (json!(12), 1, "12"), + (json!([]), 0, ""), + (json!([[]]), 0, ""), + ]; + for (value, rows, json) in cases { + let mut budget = FlattenBudget::new(GENEROUS_LIMITS); + charge_generic_flattening(&value, &mut budget).unwrap(); + assert_eq!( + budget.used(), + FlattenCost { + rows, + bytes: json.len() + }, + "input: {value}" + ); + } + } + + #[test] + fn rows_per_event_limit_is_exact() { + // 10 x 10 sibling arrays + let value = + json!({"a": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]}); + assert_eq!( + generic_flattening_with_limits(&value, limits(100, 100, 1 << 20)) + .unwrap() + .len(), + 100 + ); + assert!(matches!( + generic_flattening_with_limits(&value, limits(99, 100, 1 << 20)), + Err(JsonFlattenError::RowsPerEventLimitExceeded(99)) + )); + + // 2 x (2 x 3) nested product + let value = json!({"a": [1, 2], "b": {"c": [1, 2], "d": [{"e": [1, 2]}, {"e": [3]}]}}); + assert_eq!( + generic_flattening_cost(&value, &limits(12, 12, 1 << 20)) + .unwrap() + .rows, + 12 + ); + assert!(matches!( + generic_flattening_cost(&value, &limits(11, 12, 1 << 20)), + Err(JsonFlattenError::RowsPerEventLimitExceeded(11)) + )); + + // Deeply nested arrays are spliced into the rows of their parent key + let value = json!({"a": [[[[1, 2], [3]], [4]]], "b": [[[5]], 6]}); + assert_eq!( + generic_flattening_cost(&value, &limits(8, 8, 1 << 20)) + .unwrap() + .rows, + 8 + ); + assert!(generic_flattening_cost(&value, &limits(7, 8, 1 << 20)).is_err()); + } + + #[test] + fn rows_per_request_limit_adds_up_individually_legal_events() { + let event = json!({"a": [1, 2, 3, 4, 5], "b": [1, 2]}); + let batch = Value::Array(vec![event.clone(); 3]); + assert_eq!( + generic_flattening_with_limits(&batch, limits(10, 30, 1 << 20)) + .unwrap() + .len(), + 30 + ); + assert!(matches!( + generic_flattening_with_limits(&batch, limits(10, 29, 1 << 20)), + Err(JsonFlattenError::RowsPerRequestLimitExceeded(29)) + )); + + // Nested top level arrays are spliced, so their elements are events of their own + let batch = json!([[event.clone(), [event.clone()]], event]); + assert!(generic_flattening_with_limits(&batch, limits(10, 30, 1 << 20)).is_ok()); + assert!(matches!( + generic_flattening_with_limits(&batch, limits(29, 29, 1 << 20)), + Err(JsonFlattenError::RowsPerRequestLimitExceeded(29)) + )); + } + + #[test] + fn budget_is_shared_and_unchanged_by_a_rejected_event() { + let event = json!({"a": [1, 2, 3]}); + let cost = generic_flattening_cost(&event, &GENEROUS_LIMITS).unwrap(); + let mut budget = FlattenBudget::new(limits(3, 7, 1 << 20)); + generic_flattening(&event, &mut budget).unwrap(); + generic_flattening(&event, &mut budget).unwrap(); + assert_eq!(budget.used().rows, 6); + assert!(matches!( + generic_flattening(&event, &mut budget), + Err(JsonFlattenError::RowsPerRequestLimitExceeded(7)) + )); + assert_eq!(budget.used().rows, 6); + // A smaller event still fits in what is left + generic_flattening(&json!({"a": 1}), &mut budget).unwrap(); + assert_eq!(budget.used().rows, 7); + assert_eq!(budget.used().bytes, 2 * cost.bytes + r#"{"a":1}"#.len()); + + let mut budget = FlattenBudget::new(limits(3, 7, 2 * cost.bytes)); + budget.charge_event(cost).unwrap(); + assert!(matches!( + budget.charge_event(FlattenCost { + rows: 1, + bytes: cost.bytes + 1 + }), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(_)) + )); + assert_eq!(budget.used(), cost); + budget.charge_event(cost).unwrap(); + assert_eq!(budget.used().bytes, 2 * cost.bytes); + } + + #[test] + fn bytes_limit_counts_repeated_and_escaped_strings_exactly() { + let long = "x".repeat(1000); + let value = json!({"s": long, "a": [1, 2, 3]}); + let row_len = serde_json::to_vec(&json!({"s": long, "a": 1})) + .unwrap() + .len(); + let bytes = 3 * row_len; + assert_eq!( + generic_flattening_cost(&value, &limits(3, 3, bytes)).unwrap(), + FlattenCost { rows: 3, bytes } + ); + assert!(matches!( + generic_flattening_with_limits(&value, limits(3, 3, bytes - 1)), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(limit)) if limit == bytes - 1 + )); + + // Escapes take more room than the characters they stand for + let value = json!({"k\"\n": "\"\\\u{1}\u{1f}é", "a": [1, 2]}); + let bytes = + rows_cost(&generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap()).bytes; + assert_eq!(bytes, 2 * r#"{"k\"\n":"\"\\\u0001\u001fé","a":1}"#.len()); + assert!(generic_flattening_cost(&value, &limits(2, 2, bytes)).is_ok()); + assert!(generic_flattening_cost(&value, &limits(2, 2, bytes - 1)).is_err()); + } + + #[test] + fn cartesian_product_is_rejected_before_it_is_built() { + // 1000 x 1000 x 1000 = 10^9 rows from an 11 KiB event + let value = cube(1000); + assert!(matches!( + generic_flattening_cost(&value, &FlattenLimits::DEFAULT), + Err(JsonFlattenError::RowsPerEventLimitExceeded(1000)) + )); + let mut budget = FlattenBudget::new(FlattenLimits::DEFAULT); + assert!(generic_flattening(&value, &mut budget).is_err()); + assert_eq!(budget.used(), FlattenCost::default()); + + // A cube of 10 x 10 x 10 is exactly at the default limit + assert_eq!( + generic_flattening_with_limits(&cube(10), FlattenLimits::DEFAULT) + .unwrap() + .len(), + 1000 + ); + assert!(generic_flattening_with_limits(&cube(11), FlattenLimits::DEFAULT).is_err()); + } + + #[test] + fn overflowing_counts_are_limit_errors() { + let unlimited = limits(usize::MAX, usize::MAX, usize::MAX); + let pairs = |keys: usize| { + Value::Object( + (0..keys) + .map(|key| (key.to_string(), json!([0, 1]))) + .collect(), + ) + }; + // The values in the 2^60 rows of 60 keys take more than usize::MAX bytes, and since + // every row takes at least two bytes these overflow before the rows do + assert!(matches!( + generic_flattening_cost(&pairs(60), &unlimited), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(usize::MAX)) + )); + // Rows are checked against the event limit as they are multiplied + assert!(matches!( + generic_flattening_cost(&pairs(65), &limits(1 << 20, usize::MAX, usize::MAX)), + Err(JsonFlattenError::RowsPerEventLimitExceeded(_)) + )); + // And the same for the structurally flattened fields + assert!(matches!( + flatten_pipeline_cost(&pairs(60), &unlimited), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(usize::MAX)) + )); + assert!(matches!( + flatten_pipeline_cost(&pairs(65), &limits(1 << 20, usize::MAX, usize::MAX)), + Err(JsonFlattenError::RowsPerEventLimitExceeded(_)) + )); + } + + // Compact JSON bytes and fields of `rows` once structurally flattened, as the callers + // of `pipeline_flattening` do it. Rows that are not objects stay as they are. + fn structural_cost(rows: &[Value]) -> (usize, usize) { + let mut bytes = 0; + let mut fields = 0; + for row in rows { + let mut row = row.clone(); + if let Value::Object(_) = row { + flatten(&mut row, "_", None, None, None, false).unwrap(); + fields += row.as_object().unwrap().len(); + } + bytes += serde_json::to_vec(&row).unwrap().len(); + } + (bytes, fields) + } + + fn pipeline_fixtures() -> Vec { + vec![ + long_parent("p", true), + long_parent("p\"\n\\\u{1}é", true), + json!({"p\"\n": {"a": 1, "b\\": {"c": "x"}}}), + json!({"a": {"b": {"c": 1, "d": {"e": "x"}}}, "f": null}), + json!({"a": {}, "b": {"c": {}}}), + json!({"a": {}, "x": [1, 2]}), + json!({"": {"": [1, {"": 2}]}}), + json!({"a": [{"x": 1}, {"y": 2, "z": {"w": 3}}, {}, null, 4, []]}), + json!({"a": [{"b": 1}, {"c": 2}], "d": {"e": [1, 2], "f": {}}}), + json!({"a": [{"b": {}}, {"b": {"c": [{}, {"d": 1}]}}], "e": [[], [[{"f": 1}]]]}), + json!([{"a": {"b": 1}}, [{"c": {"d": [1, 2]}}], 3, null, [[]]]), + ] + } + + // `parent` holding 32 fields, in 4 rows when `repeated` + fn long_parent(parent: &str, repeated: bool) -> Value { + let fields: Map = (0..32).map(|i| (format!("f{i}"), json!(i))).collect(); + let mut map = Map::new(); + map.insert(parent.to_owned(), Value::Object(fields)); + if repeated { + map.insert("r".to_owned(), json!([0, 1, 2, 3])); + } + Value::Object(map) + } + + #[test] + fn pipeline_cost_matches_generic_and_structural_output() { + let (mut exact, mut collided) = (0, 0); + for value in corpus().into_iter().chain(pipeline_fixtures()) { + let mut charged = FlattenCost::default(); + for event in events(&value) { + let built = generic_flattening_with_limits(event, GENEROUS_LIMITS).unwrap(); + let generic = rows_cost(&built); + let cost = flatten_pipeline_cost(event, &GENEROUS_LIMITS).unwrap(); + // The generic part is exactly `generic_flattening_cost` + assert_eq!( + (cost.rows, cost.generic_bytes), + (generic.rows, generic.bytes), + "event: {event}" + ); + let (bytes, fields) = structural_cost(&built); + if fields == cost.structural_fields { + assert_eq!(cost.structural_bytes, bytes, "event: {event}"); + exact += 1; + } else { + // Fields whose joined names collide are overwritten, but counted + assert!( + fields < cost.structural_fields && bytes < cost.structural_bytes, + "event: {event}" + ); + collided += 1; + } + assert_eq!( + cost.charged(), + FlattenCost { + rows: generic.rows, + bytes: generic.bytes.max(cost.structural_bytes) + } + ); + charged = charged.add(cost.charged(), &GENEROUS_LIMITS).unwrap(); + } + let mut budget = FlattenBudget::new(GENEROUS_LIMITS); + charge_flatten_pipeline(&value, &mut budget).unwrap(); + assert_eq!(budget.used(), charged, "input: {value}"); + // The rows are those of generic flattening + assert_eq!( + ordered(pipeline_flattening( + &value, + &mut FlattenBudget::new(GENEROUS_LIMITS) + )), + ordered(generic_flattening_with_limits(&value, GENEROUS_LIMITS)), + ); + } + assert!(exact > 1000, "exact for {exact} events"); + assert!(collided > 10, "collisions in {collided} events"); + } + + #[test] + fn structural_fields_are_named_after_their_parents() { + for (value, expected) in [ + ( + json!({"p\"\n": {"a": 1, "b\\": {"c": "x"}}}), + json!([{"p\"\n_a": 1, "p\"\n_b\\_c": "x"}]), + ), + ( + json!({"a": {"b": {"c": 1, "d": {"e": "x"}}}, "f": null}), + json!([{"a_b_c": 1, "a_b_d_e": "x", "f": null}]), + ), + (json!({"a": {}, "b": {"c": {}}}), json!([{}])), + (json!({"a": {}, "x": [1, 2]}), json!([{"x": 1}, {"x": 2}])), + ( + json!({"": {"": [1, {"": 2}]}}), + json!([{"_": 1}, {"__": 2}]), + ), + ( + json!({"a": [{"x": 1}, {"y": 2, "z": {"w": 3}}, {}, null, 4, []]}), + json!([{"a_x": 1}, {"a_y": 2, "a_z_w": 3}, {}, {"a": null}, {"a": 4}, {"a": null}]), + ), + ] { + let built = generic_flattening_with_limits(&value, GENEROUS_LIMITS).unwrap(); + let mut flattened = Value::Array(built.clone()); + flatten(&mut flattened, "_", None, None, None, false).unwrap(); + assert_eq!(flattened, expected, "input: {value}"); + let cost = flatten_pipeline_cost(&value, &GENEROUS_LIMITS).unwrap(); + let expected = expected.as_array().unwrap(); + assert_eq!( + cost.structural_bytes, + rows_cost(expected).bytes, + "input: {value}" + ); + assert_eq!( + cost.structural_fields, + expected + .iter() + .map(|row| row.as_object().unwrap().len()) + .sum::(), + "input: {value}" + ); + } + } + + #[test] + fn long_parent_names_repeated_in_fields_are_rejected_before_expansion() { + // The configuration of the HTTP smoke test + let small = limits(4, 6, 4096); + for repeated in [true, false] { + let value = long_parent(&"p".repeat(128), repeated); + // The generic rows fit, the structurally flattened ones do not + let generic = generic_flattening_cost(&value, &small).unwrap(); + assert!(generic.bytes <= 4096, "{generic:?}"); + assert!(matches!( + flatten_pipeline_cost(&value, &small), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + )); + let mut budget = FlattenBudget::new(small); + assert!(matches!( + pipeline_flattening_owned(value.clone(), &mut budget), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + )); + assert!(matches!( + pipeline_flattening(&value, &mut budget), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + )); + assert!(charge_flatten_pipeline(&value, &mut budget).is_err()); + assert_eq!(budget.used(), FlattenCost::default()); + // Generic flattening alone is charged its own rows only + generic_flattening_owned(value, &mut budget).unwrap(); + + // A short parent name fits, charged for the larger of both representations + let value = long_parent("p", repeated); + let mut budget = FlattenBudget::new(small); + let rows = pipeline_flattening_owned(value, &mut budget).unwrap(); + let generic = rows_cost(&rows); + let (structural, _) = structural_cost(&rows); + assert!(structural > generic.bytes); + assert_eq!( + budget.used(), + FlattenCost { + rows: generic.rows, + bytes: structural + } + ); + } + } + + #[test] + fn structural_bytes_add_up_over_events() { + let event = long_parent("p", true); + let cost = flatten_pipeline_cost(&event, &GENEROUS_LIMITS) + .unwrap() + .charged(); + assert!( + cost.bytes + > generic_flattening_cost(&event, &GENEROUS_LIMITS) + .unwrap() + .bytes + ); + // Also through the array-free fast path + let flat = long_parent("p", false); + let flat_cost = flatten_pipeline_cost(&flat, &GENEROUS_LIMITS) + .unwrap() + .charged(); + let batch = json!([flat, event]); + let bytes = flat_cost.bytes + cost.bytes; + + let mut budget = FlattenBudget::new(limits(4, 8, bytes - 1)); + assert!(matches!( + pipeline_flattening_owned(batch.clone(), &mut budget), + Err(JsonFlattenError::BytesPerRequestLimitExceeded(_)) + )); + // The first event stays charged, the rejected one is not + assert_eq!(budget.used(), flat_cost); + + let mut budget = FlattenBudget::new(limits(4, 8, bytes)); + assert_eq!( + pipeline_flattening_owned(batch, &mut budget).unwrap().len(), + 5 + ); + assert_eq!(budget.used().bytes, bytes); + } + + #[test] + fn limit_errors_are_identified() { + assert!(JsonFlattenError::RowsPerEventLimitExceeded(1).is_limit_exceeded()); + assert!(JsonFlattenError::RowsPerRequestLimitExceeded(1).is_limit_exceeded()); + assert!(JsonFlattenError::BytesPerRequestLimitExceeded(1).is_limit_exceeded()); + assert!(!JsonFlattenError::CannotFlatten.is_limit_exceeded()); + assert!(!JsonFlattenError::ExpectedObjectInArray.is_limit_exceeded()); + let message = JsonFlattenError::RowsPerEventLimitExceeded(1000).to_string(); + assert!(message.contains("1000"), "{message}"); + assert!( + message.contains("P_MAX_FLATTEN_ROWS_PER_EVENT"), + "{message}" + ); + let message = JsonFlattenError::RowsPerRequestLimitExceeded(10).to_string(); + assert!(message.contains("Split the batch"), "{message}"); } } diff --git a/src/utils/json/mod.rs b/src/utils/json/mod.rs index 2ffc8f9a3..1a7f72ee5 100644 --- a/src/utils/json/mod.rs +++ b/src/utils/json/mod.rs @@ -20,7 +20,7 @@ use std::fmt; use std::num::NonZeroU32; use flatten::{ - convert_to_array, generic_flattening, generic_flattening_noop, has_more_than_max_allowed_levels, + FlattenBudget, convert_to_array, has_more_than_max_allowed_levels, pipeline_flattening_owned, }; use serde::de::Visitor; use serde_json; @@ -36,6 +36,9 @@ pub mod strict; /// calls the function `flatten_json` which results Vec or Error /// in case when Vec is returned, converts the Vec to Value of Array /// this is to ensure recursive flattening does not happen for heavily nested jsons +/// Generic flattening is charged to `budget`, together with the structural flattening of its +/// rows that follows, see `flatten::charge_flatten_pipeline` +#[allow(clippy::too_many_arguments)] #[cfg_attr(feature = "hotpath", hotpath::measure)] pub fn flatten_json_body( body: Value, @@ -45,8 +48,9 @@ pub fn flatten_json_body( schema_version: SchemaVersion, validation_required: bool, log_source: &LogSource, + budget: &mut FlattenBudget, ) -> Result { - // Flatten the json body only if new schema and has less than 4 levels of nesting + // Flatten the json body only if new schema and not nested deeper than P_MAX_FLATTEN_LEVEL let apply_generic_flattening = should_apply_generic_flattening(&body, schema_version, log_source); flatten_json_body_inner( @@ -56,6 +60,7 @@ pub fn flatten_json_body( time_partition_limit, custom_partition, validation_required, + budget, ) } @@ -66,15 +71,10 @@ fn flatten_json_body_inner( time_partition_limit: Option, custom_partition: Option<&String>, validation_required: bool, + budget: &mut FlattenBudget, ) -> Result { let mut nested_value = if apply_generic_flattening { - match generic_flattening_noop(body) { - Ok(unchanged) => unchanged, - Err(body) => { - let flattened_json = generic_flattening(&body)?; - convert_to_array(flattened_json)? - } - } + convert_to_array(pipeline_flattening_owned(body, budget)?)? } else { body }; @@ -100,18 +100,16 @@ fn should_apply_generic_flattening( && matches!(log_source, LogSource::Json | LogSource::Custom(_)) } -/// Applies generic flattening and handles the result for partitioned processing +/// Applies generic flattening, charged to `budget` together with the structural flattening +/// below, and handles the result for partitioned processing pub fn apply_generic_flattening_for_partition( element: Value, time_partition: Option<&String>, time_partition_limit: Option, custom_partition: Option<&String>, + budget: &mut FlattenBudget, ) -> Result, anyhow::Error> { - let flattened_json = match element { - // Generic flattening returns an array-free object unchanged as its only row - Value::Object(ref map) if !flatten::map_contains_array(map) => vec![element], - element => generic_flattening(&element)?, - }; + let flattened_json = pipeline_flattening_owned(element, budget)?; if flattened_json.len() == 1 { // Single result - process normally @@ -152,6 +150,7 @@ fn process_partitioned_element( custom_partition: Option<&String>, schema_version: SchemaVersion, log_source: &LogSource, + budget: &mut FlattenBudget, ) -> Result, anyhow::Error> { if should_apply_generic_flattening(&element, schema_version, log_source) { apply_generic_flattening_for_partition( @@ -159,6 +158,7 @@ fn process_partitioned_element( time_partition, time_partition_limit, custom_partition, + budget, ) } else { let mut nested_value = element; @@ -182,6 +182,7 @@ fn process_partitioned_array( custom_partition: Option<&String>, schema_version: SchemaVersion, log_source: &LogSource, + budget: &mut FlattenBudget, ) -> Result, anyhow::Error> { let mut result = Vec::new(); @@ -193,6 +194,7 @@ fn process_partitioned_array( custom_partition, schema_version, log_source, + budget, )?; result.extend(processed_elements); } @@ -208,6 +210,7 @@ fn process_partitioned_non_array( custom_partition: Option<&String>, schema_version: SchemaVersion, log_source: &LogSource, + budget: &mut FlattenBudget, ) -> Result, anyhow::Error> { // convert to an array for processing let arr = vec![body]; @@ -218,6 +221,7 @@ fn process_partitioned_non_array( custom_partition, schema_version, log_source, + budget, )?; Ok(processed_elements) } @@ -231,6 +235,7 @@ fn process_non_partitioned( custom_partition: Option<&String>, schema_version: SchemaVersion, log_source: &LogSource, + budget: &mut FlattenBudget, ) -> Result, anyhow::Error> { let data = flatten_json_body( body, @@ -240,6 +245,7 @@ fn process_non_partitioned( schema_version, true, log_source, + budget, )?; // For non-partitioned processing, return the flattened data as a single item @@ -255,6 +261,7 @@ pub fn convert_array_to_object( custom_partition: Option<&String>, schema_version: SchemaVersion, log_source: &LogSource, + budget: &mut FlattenBudget, ) -> Result, anyhow::Error> { let _span = info_span!("convert_array_to_object").entered(); if time_partition.is_some() || custom_partition.is_some() { @@ -266,6 +273,7 @@ pub fn convert_array_to_object( custom_partition, schema_version, log_source, + budget, ), _ => process_partitioned_non_array( body, @@ -274,6 +282,7 @@ pub fn convert_array_to_object( custom_partition, schema_version, log_source, + budget, ), } } else { @@ -284,6 +293,7 @@ pub fn convert_array_to_object( custom_partition, schema_version, log_source, + budget, ) } } @@ -343,6 +353,7 @@ where #[cfg(test)] mod tests { use super::*; + use flatten::{FlattenLimits, GENEROUS_LIMITS, JsonFlattenError}; use serde::{Deserialize, Serialize}; use serde_json::json; #[derive(Serialize, Deserialize)] @@ -389,6 +400,10 @@ mod tests { Ok(result) } + fn generous() -> FlattenBudget { + FlattenBudget::new(GENEROUS_LIMITS) + } + fn ordered(result: Result) -> Result { result .map(|value| serde_json::to_string(&value).unwrap()) @@ -421,6 +436,7 @@ mod tests { None, partition.as_ref(), validation_required, + &mut generous(), )), ordered(legacy_flatten_json_body( value.clone(), @@ -447,6 +463,7 @@ mod tests { None, None, partition.as_ref(), + &mut generous(), )), ordered(legacy_apply_generic_flattening_for_partition( value.clone(), @@ -462,17 +479,26 @@ mod tests { fn flatten_json_body_wraps_flat_object() { let body = json!({"a": {"b": 2}, "c": {}, "d": 1}); assert_eq!( - flatten_json_body_inner(body.clone(), true, None, None, None, false).unwrap(), + flatten_json_body_inner(body.clone(), true, None, None, None, false, &mut generous()) + .unwrap(), json!([{"a_b": 2, "d": 1}]) ); assert_eq!( - flatten_json_body_inner(json!({"b": 1, "a": 2}), true, None, None, None, false) - .unwrap(), + flatten_json_body_inner( + json!({"b": 1, "a": 2}), + true, + None, + None, + None, + false, + &mut generous() + ) + .unwrap(), json!([{"b": 1, "a": 2}]) ); // Too deeply nested bodies skip generic flattening and are not wrapped assert_eq!( - flatten_json_body_inner(body, false, None, None, None, false).unwrap(), + flatten_json_body_inner(body, false, None, None, None, false, &mut generous()).unwrap(), json!({"a_b": 2, "d": 1}) ); } @@ -568,7 +594,8 @@ mod tests { None, SchemaVersion::V0, false, - &crate::event::format::LogSource::default() + &crate::event::format::LogSource::default(), + &mut generous(), ) .is_err() ) @@ -604,6 +631,7 @@ mod tests { SchemaVersion::V0, false, &crate::event::format::LogSource::default(), + &mut generous(), ) .unwrap(); @@ -653,6 +681,7 @@ mod tests { None, SchemaVersion::V0, &crate::event::format::LogSource::default(), + &mut generous(), ); assert!(result.is_ok()); @@ -662,4 +691,235 @@ mod tests { assert_eq!(objects.len(), 1); assert_eq!(objects[0], json); } + + fn limits(rows_per_event: usize, rows_per_request: usize, bytes: usize) -> FlattenBudget { + FlattenBudget::new(FlattenLimits { + max_rows_per_event: rows_per_event, + max_rows_per_request: rows_per_request, + max_bytes_per_request: bytes, + }) + } + + fn limit_error(result: Result) -> JsonFlattenError { + match result { + Ok(_) => panic!("expected a flatten limit error"), + Err(err) => match err.downcast::() { + Ok(err) if err.is_limit_exceeded() => err, + other => panic!("expected a flatten limit error, got {other:?}"), + }, + } + } + + #[test] + fn flatten_json_body_rejects_expansion_over_limits() { + let body = json!([{"a": [1, 2, 3], "b": [1, 2]}, {"c": [1, 2, 3, 4, 5, 6]}]); + let rows = flatten_json_body_inner( + body.clone(), + true, + None, + None, + None, + false, + &mut limits(6, 12, 1 << 20), + ) + .unwrap(); + assert_eq!(rows.as_array().unwrap().len(), 12); + assert!(matches!( + limit_error(flatten_json_body_inner( + body.clone(), + true, + None, + None, + None, + false, + &mut limits(5, 12, 1 << 20), + )), + JsonFlattenError::RowsPerEventLimitExceeded(5) + )); + assert!(matches!( + limit_error(flatten_json_body_inner( + body.clone(), + true, + None, + None, + None, + false, + &mut limits(6, 11, 1 << 20), + )), + JsonFlattenError::RowsPerRequestLimitExceeded(11) + )); + // Bodies too deeply nested for generic flattening are not charged + let mut budget = limits(1, 1, 1); + flatten_json_body_inner(body, false, None, None, None, false, &mut budget).unwrap(); + assert_eq!(budget.used(), flatten::FlattenCost::default()); + } + + #[test] + fn noop_fast_path_is_charged() { + let body = json!([{"a": 1}, {"b": {"c": "x"}}, {}]); + let bytes = r#"{"a":1}{"b":{"c":"x"}}{}"#.len(); + let mut budget = limits(1, 3, bytes); + let flattened = + flatten_json_body_inner(body.clone(), true, None, None, None, false, &mut budget) + .unwrap(); + assert_eq!(flattened, json!([{"a": 1}, {"b_c": "x"}, {}])); + assert_eq!(budget.used(), flatten::FlattenCost { rows: 3, bytes }); + assert!(matches!( + limit_error(flatten_json_body_inner( + body.clone(), + true, + None, + None, + None, + false, + &mut limits(1, 2, bytes), + )), + JsonFlattenError::RowsPerRequestLimitExceeded(2) + )); + assert!(matches!( + limit_error(flatten_json_body_inner( + body.clone(), + true, + None, + None, + None, + false, + &mut limits(1, 3, bytes - 1), + )), + JsonFlattenError::BytesPerRequestLimitExceeded(_) + )); + // A single array-free object takes the same fast path + assert!(matches!( + limit_error(flatten_json_body_inner( + json!({"a": "xyz"}), + true, + None, + None, + None, + false, + &mut limits(1, 1, r#"{"a":"xyz"}"#.len() - 1), + )), + JsonFlattenError::BytesPerRequestLimitExceeded(_) + )); + } + + #[test] + fn partitioned_elements_share_the_request_budget() { + let partition = Some("a".to_owned()); + let element = json!({"a": "x", "b": [1, 2, 3]}); + let flat_element = json!({"a": "x", "b": 1}); + // Partitioned processing charges every element of the request to one budget + let mut budget = limits(3, 8, 1 << 20); + for element in [&element, &flat_element, &element] { + apply_generic_flattening_for_partition( + element.clone(), + None, + None, + partition.as_ref(), + &mut budget, + ) + .unwrap(); + } + assert_eq!(budget.used().rows, 7); + // The fast path for array-free elements is charged too + apply_generic_flattening_for_partition( + flat_element.clone(), + None, + None, + partition.as_ref(), + &mut budget, + ) + .unwrap(); + assert!(matches!( + limit_error(apply_generic_flattening_for_partition( + flat_element, + None, + None, + partition.as_ref(), + &mut budget, + )), + JsonFlattenError::RowsPerRequestLimitExceeded(8) + )); + assert!(matches!( + limit_error(apply_generic_flattening_for_partition( + element, + None, + None, + partition.as_ref(), + &mut limits(2, 8, 1 << 20), + )), + JsonFlattenError::RowsPerEventLimitExceeded(2) + )); + } + + // 32 fields below a parent named `parent`, in 4 rows when `repeated` + fn long_parent(parent: &str, repeated: bool) -> Value { + let fields: serde_json::Map = + (0..32).map(|i| (format!("f{i}"), json!(i))).collect(); + let mut map = serde_json::Map::new(); + map.insert(parent.to_owned(), Value::Object(fields)); + if repeated { + map.insert("r".to_owned(), json!([0, 1, 2, 3])); + } + Value::Object(map) + } + + fn flatten_body(body: Value, budget: &mut FlattenBudget) -> Result { + flatten_json_body_inner(body, true, None, None, None, false, budget) + } + + #[test] + fn structural_field_names_are_charged_before_expansion() { + // The configuration of the HTTP smoke test + let small = FlattenLimits { + max_rows_per_event: 4, + max_rows_per_request: 6, + max_bytes_per_request: 4096, + }; + let long = long_parent(&"p".repeat(128), true); + // The generic rows alone fit + assert!( + flatten::generic_flattening_cost(&long, &small) + .unwrap() + .bytes + <= 4096 + ); + let mut budget = FlattenBudget::new(small); + assert!(matches!( + limit_error(flatten_body(long.clone(), &mut budget)), + JsonFlattenError::BytesPerRequestLimitExceeded(4096) + )); + let partition = Some("r".to_owned()); + assert!(matches!( + limit_error(apply_generic_flattening_for_partition( + long, + None, + None, + partition.as_ref(), + &mut budget, + )), + JsonFlattenError::BytesPerRequestLimitExceeded(4096) + )); + assert_eq!(budget.used(), flatten::FlattenCost::default()); + + // A short parent name fits, charged for the structurally flattened rows + let rows = flatten_body(long_parent("p", true), &mut budget).unwrap(); + let rows = rows.as_array().unwrap(); + assert_eq!(rows.len(), 4); + assert_eq!(rows[0]["p_f31"], json!(31)); + let bytes: usize = rows + .iter() + .map(|row| serde_json::to_vec(row).unwrap().len()) + .sum(); + assert_eq!(budget.used(), flatten::FlattenCost { rows: 4, bytes }); + + // An array-free event that fits on its own, but not in what is left of the request + let flat = long_parent(&"q".repeat(100), false); + flatten_body(flat.clone(), &mut FlattenBudget::new(small)).unwrap(); + assert!(matches!( + limit_error(flatten_body(flat, &mut budget)), + JsonFlattenError::BytesPerRequestLimitExceeded(4096) + )); + assert_eq!(budget.used(), flatten::FlattenCost { rows: 4, bytes }); + } } diff --git a/tools/flatten-bench/.gitignore b/tools/flatten-bench/.gitignore new file mode 100644 index 000000000..df611ec85 --- /dev/null +++ b/tools/flatten-bench/.gitignore @@ -0,0 +1,2 @@ +/target +__pycache__/ diff --git a/tools/flatten-bench/Cargo.lock b/tools/flatten-bench/Cargo.lock new file mode 100644 index 000000000..69d1fc360 --- /dev/null +++ b/tools/flatten-bench/Cargo.lock @@ -0,0 +1,436 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "anyhow" +version = "1.0.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "bumpalo" +version = "3.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" + +[[package]] +name = "cc" +version = "1.2.51" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a0aeaff4ff1a90589618835a598e545176939b97874f7abc7851caa0618f203" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "find-msvc-tools" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645cbb3a84e60b7531617d5ae4e57f7e27308f6445f5abf653209ea76dec8dff" + +[[package]] +name = "flatten-bench" +version = "0.0.0" +dependencies = [ + "anyhow", + "chrono", + "serde", + "serde_json", + "thiserror", + "tracing", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "iana-time-zone" +version = "0.1.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33e57f83510bb73707521ebaffa789ec8caf86f9657cad665b092b581d40e9fb" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "itoa" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" + +[[package]] +name = "js-sys" +version = "0.3.83" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "pin-project-lite" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "indexmap", + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f63587ca0f12b72a0600bcba1d40081f830876000bb46dd2337a3051618f4fc8" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ff15c8ecd7de3849db632e14d18d2571fa09dfc5ed93479bc4485c7a517c913" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "unicode-ident" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" + +[[package]] +name = "wasm-bindgen" +version = "0.2.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "zmij" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f4a4e8e9dc5c62d159f04fcdbe07f4c3fb710415aab4754bf11505501e3251d" diff --git a/tools/flatten-bench/Cargo.toml b/tools/flatten-bench/Cargo.toml new file mode 100644 index 000000000..f6224b332 --- /dev/null +++ b/tools/flatten-bench/Cargo.toml @@ -0,0 +1,32 @@ +# Standalone harness that compiles the production `src/utils/json` sources (and a +# git-extracted baseline copy of them) without building the Parseable server. +# It has its own empty workspace so it never joins the root crate's build graph. +[package] +name = "flatten-bench" +version = "0.0.0" +edition = "2024" +publish = false + +[workspace] + +[lib] +# The production files carry `#[cfg(test)]` modules that reference server-only +# paths, so the library itself is never compiled in test mode. +test = false +doctest = false + +[dependencies] +anyhow = "1.0" +chrono = "0.4" +serde = { version = "1.0.228", features = ["derive"] } +# The server binary gets `preserve_order` through feature unification (serde_json +# depends on indexmap in the root Cargo.lock), so Map behaves the same here. +serde_json = { version = "1.0.150", features = ["preserve_order"] } +thiserror = "2.0" +tracing = "0.1.44" + +[profile.release] +debug = 1 + +[lints.rust] +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(feature, values("hotpath"))'] } diff --git a/tools/flatten-bench/README.md b/tools/flatten-bench/README.md new file mode 100644 index 000000000..ec093ba7e --- /dev/null +++ b/tools/flatten-bench/README.md @@ -0,0 +1,79 @@ +# flatten-bench + +Measured results against 547c5118 are in [RESULTS.md](RESULTS.md). + +Standalone harness for the generic flattening limits in `src/utils/json`. It compiles the +production sources unmodified, so it measures the real code without building the server +(own `[workspace]`, depends only on serde/serde_json/chrono/thiserror/anyhow/tracing): + +- `current_json`: the working tree's `src/utils/json/{mod,flatten,strict}.rs`, included by path. +- `baseline_json`: the same files at `547c5118` (the merge-base, including its no-op fast + paths), extracted byte for byte with `git show` by `build.rs` into `OUT_DIR`. Override + with `FLATTEN_BENCH_BASELINE_REV=`. This is NOT the `#[cfg(test)] legacy` module in + flatten.rs, which predates those fast paths; `[lib] test = false` keeps production + `#[cfg(test)]` code out entirely. The bench prints the revision and flatten.rs blob id. + +`src/lib.rs` stubs the few server items those files use (`crate::cli`, `PARSEABLE`, +`LogSource`, `SchemaVersion`); `stub_defaults_match_server` checks the stubbed +`DEFAULT_MAX_FLATTEN_*` against `src/cli.rs`. `src/api.rs` is the only file that names +production signatures. + +## Run + +```sh +export PATH=$HOME/.cargo/bin:$PATH +cd tools/flatten-bench +RUST_LIB_BACKTRACE=0 CARGO_BUILD_JOBS=4 cargo test --release -- --nocapture +cargo run --release --bin flatten_bench | tee /tmp/flatten-bench-$(git rev-parse --short HEAD)-$(date +%s).log +``` + +Settings: `FLATTEN_BENCH_SAMPLES` (31), `FLATTEN_BENCH_WARMUP` (5), +`FLATTEN_BENCH_SAMPLE_NS` (2_000_000 ns per sample). For a reproducible report, record +`git rev-parse HEAD`, `git diff --stat -- src/utils/json`, `rustc -V`, and the loadavg +lines the bench prints; run on a quiet machine (no parallel cargo builds) 2-3 times and +compare medians, since same-host noise is often larger than small deltas. + +## What is timed + +- Parsing is never timed: inputs are built once as `serde_json::Value`. +- Baseline safety: every accepted fixture is checked ONCE (`BaselineInput::check`, a + saturating row upper bound <= 1000) before any baseline call and outside every timed or + counted region. Timed baseline calls take the checked `BaselineInput` and do no extra + walk; the current side does only its real preflight. Hostile inputs are never given to + the baseline. `baseline_convert`/`baseline_generic` re-check per call and are for tests only. +- Consuming entry point (`convert_array_to_object`): each sample pre-clones `batch` inputs + outside the timer; the timer covers the call plus dropping its output and the consumed + input, the same for both sides. Borrowed APIs reuse one input; output drop is timed. +- Baseline/current samples alternate AB, BA, ... on one thread; median and p95 per call. + Outputs are asserted identical (order-sensitive) before timing. A checksum is printed and + results go through `black_box`. +- Sections: entry point; `generic_flattening` only (bounded vs unbounded; the current side + also contains the rewritten expansion); preflight only (`charge_flatten_pipeline`, which + counts generic and structurally flattened bytes, on a fresh budget, per event like the + entry point); rejections (current only). + +## Allocation numbers + +From one separate counted call per implementation (never the timed ones), on the calling +thread only, in requested bytes (not RSS): `allocs` alloc/realloc calls, `alloc_bytes` +cumulative bytes requested, `peak_delta` highest live bytes above the level at call start, +`net_delta` live bytes at the end minus the start. `peak_delta` is a delta, not an absolute +peak: for consuming calls, freeing the input during the call lowers it. The borrowed +`generic_flattening` section shows the extra memory without that offset. + +## HTTP smoke test + +`http_smoke.py` (Python 3 standard library only) starts a built server binary with +`local-store` on an ephemeral 127.0.0.1 port, in a temporary directory, with small flatten +caps (rows per event and bytes from the environment, rows per request from the CLI flag, +overriding a different environment value). It checks accepted requests at the limits and +`400` rejections over them, with unchanged event counts and schemas, for JSON (on +`/api/v1/ingest` and `/api/v1/logstream/{name}`), a known log format, custom partitions, +Kinesis, OTEL JSON logs and schema detection, including events whose long or quote-escaped +parent names only exceed the byte limit once structural flattening repeats them in every +field name. + +``` +P_BUILD_NO_UI=1 cargo build --bin parseable +python3 tools/flatten-bench/http_smoke.py target/debug/parseable +``` diff --git a/tools/flatten-bench/RESULTS.md b/tools/flatten-bench/RESULTS.md new file mode 100644 index 000000000..c8285a1fc --- /dev/null +++ b/tools/flatten-bench/RESULTS.md @@ -0,0 +1,155 @@ +# flatten-bench results: generic flattening limits vs 547c5118 + +Three back-to-back runs of the benchmark (`/tmp/flatten-row-cap-benchmark-r{6,7,8}.log`) on +the current sources, which count structurally flattened rows in the byte budget. Runs on +earlier revisions are not reported here. + +## What was measured + +| | | +|---|---| +| Baseline | `547c5118583b8e3884a2b6920c3b5a0ffa675f87`, `src/utils/json/flatten.rs` blob `55b4b9d4cb33e5cfa1fd29cf0ce440b3c90af666`, extracted with `git show` (not the `#[cfg(test)] legacy` module) | +| Current | Sources measured on `fix/flatten-row-cap` before publishing the commit, identified by sha256: `src/utils/json/flatten.rs` `b8a5008aa30722f0f6e1261f185a76ca95ae422e93aff271d7870d5eba020a99`, `src/utils/json/mod.rs` `d9346345fb670395e14a72e3a6cf64c9a1d632f6649aae73033e7e7fac0b47ba` | +| Toolchain | Rust 1.96.0 (`ac68faa20`), release profile, `debug_assertions=false` | +| Sampling | 31 samples, 5 warmup, ~2 ms per sample, baseline/current alternated AB/BA on one thread | +| Timed region | the call plus dropping its output (and, for the consuming entry point, its input). JSON parsing and input clones are outside it. Baseline inputs are checked to be <= 1000 rows once, before and outside any timed or counted region | +| Correctness | accepted outputs asserted identical (order-sensitive) to the baseline before timing | +| Host | a shared developer machine with other work running. The 1-minute loadavg was 3.3 to 4.6 during the runs (printed in each log). Each ratio below is a median, listed for all three runs, so small differences are noise | + +Limits are the defaults: 1000 rows per event, 10 000 rows per request, and 64 MiB +(67 108 864 bytes) per request. Each event is charged the larger of its generic +intermediate rows and its structurally flattened output rows, both as compact JSON +(repeated keys, values, escapes and delimiters included), plus Kinesis metadata and OTEL +body-key prefixes on those routes. Taking the larger of the two over-counts when +structural flattening merges colliding field names, so the charge errs high. It does not +measure RSS. + +## Ingestion entry point (`convert_array_to_object`, V1 JSON, includes the no-op fast paths) + +cur/base median ratio for runs r6 / r7 / r8 (below 1 means the new code is faster): + +| case | rows | cur/base | notes | +|---|---:|---|---| +| flat_single | 1 | 1.368 / 1.430 / 1.434 | about +165 to 198 ns per event (433–457 ns to 614–655 ns) | +| flat_batch_1000_noop | 1000 | 1.452 / 1.453 / 1.414 | +41–45%, about 133–137 µs per 1000 events (295–322 µs to 428–455 µs) | +| nested_no_arrays_200 | 200 | 1.109 / 1.083 / 1.109 | +8–11% | +| moderate_array_100 | 100 | 0.948 / 0.972 / 0.965 | 3–5% faster. alloc bytes 644 640 → 555 450 (−14%), allocs 4324 → 3524 | +| product_10x10x10 | 1000 | 0.990 / 0.999 / 0.985 | within noise, so no speedup is claimed | +| partitioned_batch_50x4 | 200 | 0.988 / 0.994 / 0.991 | within noise, so no speedup is claimed. alloc bytes 1 244 166 → 1 170 966 (−6%), allocs 6857 → 6057 | + +Allocation counts are identical for the flat and nested no-array cases. Peak memory does not +regress on the array cases: `peak_delta` is 270 834 → 269 942 for moderate_array_100, +2 382 594 → 2 382 592 for product_10x10x10, and 197 971 → 197 605 for +partitioned_batch_50x4. For a consuming call, `peak_delta` is offset by freeing the input; +the borrowed section below has no such offset and shows the same or lower peaks. + +## `generic_flattening` only (bounded vs unbounded, no no-op fast path) + +| case | cur/base r6 / r7 / r8 | +|---|---| +| flat_single | 1.380 / 1.348 / 1.372 | +| flat_batch_1000_noop | 1.269 / 1.315 / 1.269 | +| nested_no_arrays_200 | 1.250 / 1.227 / 1.248 | +| moderate_array_100 | 0.927 / 0.926 / 0.926 | +| product_10x10x10 | 1.003 / 0.999 / 0.983 | +| partitioned_batch_50x4 | 0.974 / 0.992 / 0.988 | + +Preflight alone (`charge_flatten_pipeline`, generic and structural bytes, on a fresh budget +per event, current only) makes 0 allocations for every case. Medians over the three runs: +flat_single 249–254 ns, flat_batch_1000 184–188 µs, nested_no_arrays_200 48.6–49.7 µs, +moderate_array_100 6.7 µs, product 438–442 ns, partitions 9.6–9.8 µs. The entry point's +flat overhead (about 135 µs per 1000 events) is lower than the standalone flat preflight +because array-free events are checked and charged in the same walk that would otherwise +detect the no-op. + +## Rejections (current only; the baseline is never run on these inputs) + +| input | entry-point median r6 / r7 / r8 | +|---|---| +| hostile_cube_1000^3 (10^9 rows) | 22.9 / 22.7 / 22.6 µs | +| deep_product_nested_arrays (~1M rows) | 15.6 / 14.2 / 15.7 µs | +| deep_product_nested_objects (10^6 rows) | 2.6 / 2.4 / 2.5 µs | +| large_repeated_string (600 rows, over 64 MiB) | 20.0 / 20.1 / 19.5 µs | +| escaped_repeated_string (500 rows, over 64 MiB after escaping) | 62.8 / 58.7 / 62.5 µs | +| just_over_event_limit (1001 rows) | 9.6 / 10.0 / 7.5 µs | + +- **Memory:** every rejection at the entry point allocates exactly once, 112 bytes (the + `anyhow` error), with `peak_delta=0`. In `/tmp/flatten-row-cap-harness-tests-r5.log` the + core paths (preflight, per-event `charge`, `pipeline charge`, and bounded + `generic_flattening`) report 0 allocations, 0 bytes and `peak_delta=0` for every hostile + input, and `core charge, remaining rows/bytes insufficient` reports the same for a request + budget whose remainder is too small. For an event whose generic rows fit but whose + structurally flattened field names do not (a 128-character parent copied into 32 field + names over 4 rows), `structural preflight, long parent` and `structural charge, long + parent` report 0 allocations, and the entry point rejects it with 1 allocation of 112 + bytes and `peak_delta=0` (2 allocations, 184 bytes, `peak_delta=72` with a partition + field). `tests/alloc_limits.rs` asserts these. They are measured by the counting + allocator, not as full-process RSS. +- **Time:** these times include freeing the consumed input. They grow with the input's size + and with how much of it is walked before the limit trips, not with the rows the input would + expand to. They are not constant across inputs. +- **Before this change:** the baseline would materialize all of the rows: 10^9 for the cube, + and more than 64 MiB of rows for the string cases. + +## Verification + +Results for the sources identified by the hashes above: + +| check | result | +|---|---| +| focused server unit tests (`utils::json`, `otel::logs`, `handlers::http::kinesis`, `option::flatten_limit_tests`, `storage::field_stats` flatten-limit tests, `handlers::http::ingest` flatten test), one invocation | 104 pass (`/tmp/flatten-row-cap-focused-tests-r3.log`) | +| allocation/harness tests (`tools/flatten-bench`, release) | 13 pass (`/tmp/flatten-row-cap-harness-tests-r5.log`) | +| HTTP smoke test against a freshly built `parseable` binary | 67/67 pass (`/tmp/flatten-row-cap-http-smoke-r3.log`) | +| `cargo clippy --all-targets --offline` | Pass, with existing unrelated warnings (`/tmp/flatten-row-cap-clippy-r3.log`) | +| `cargo fmt --all -- --check` | Pass (`/tmp/flatten-row-cap-fmt-final.log`) | +| Final Sol 6.1 core and integration reviews | No actionable findings; the escaped-parent HTTP coverage suggestion is included in the passing smoke run | + +The HTTP smoke test checks that late failures leave no schema sentinel and do not change the +event count, covers OTEL prefix and Kinesis metadata attack inputs, rejects events whose long +or quote-escaped parent names only exceed the byte limit once structurally flattened (on the +JSON, logstream, custom partition and Kinesis routes), and verifies environment flag +precedence. + +```sh +cargo test --lib --offline -- --test-threads=1 utils::json otel::logs handlers::http::kinesis \ + option::flatten_limit_tests storage::field_stats::tests::stats_payload_flatten_limits_are_bounded \ + storage::field_stats::tests::widest_stats_payload_fits_its_flatten_limits \ + handlers::http::ingest::tests::flatten +cargo test --release --manifest-path tools/flatten-bench/Cargo.toml --offline -- --nocapture +cargo build --bin parseable --offline +python3 tools/flatten-bench/http_smoke.py target/debug/parseable +``` + +## Interpretation + +- **Remaining cost:** flat and nested inputs without arrays are 8–45% slower (about 0.13 µs + per flat event in a batch, 0.17–0.20 µs for a single event). This is the price of the + added safety: the event has to be walked to count its generic and structural rows and + bytes before anything is built. The fast paths still return the input + unchanged without copying it, and allocation counts are unchanged. +- **Inputs with arrays:** these break even or improve slightly, with fewer and smaller + allocations. Only moderate_array_100 is measurably faster; product_10x10x10 and + partitioned_batch_50x4 are within noise. +- **Limits:** none of the optimizations relax them. A limit that would be exceeded is + rejected before any row is built. +- **Scope:** these are microbenchmarks of the flattening code only, not end-to-end ingestion + throughput. Parsing, Arrow conversion, and I/O are not included. + +## Reproduce + +Run from the repository root. The repository must have the baseline commit available +(`git cat-file -e 547c5118`), because `build.rs` extracts it with `git show`. + +```sh +export PATH=$HOME/.cargo/bin:$PATH P_BUILD_NO_UI=1 CARGO_BUILD_JOBS=4 +sha256sum src/utils/json/flatten.rs src/utils/json/mod.rs; rustc -V +RUST_LIB_BACKTRACE=0 cargo test --release --manifest-path tools/flatten-bench/Cargo.toml -- --nocapture +for r in 1 2 3; do + cargo run --release --manifest-path tools/flatten-bench/Cargo.toml --bin flatten_bench \ + | tee /tmp/flatten-bench-run$r.log +done +``` + +Run on an otherwise quiet machine with no parallel cargo builds. Compare the medians across +runs, and record the loadavg lines. The printed checksum adds up output row counts over +batches whose sizes are calibrated at run time, so it is expected to differ between runs. diff --git a/tools/flatten-bench/build.rs b/tools/flatten-bench/build.rs new file mode 100644 index 000000000..d2011937a --- /dev/null +++ b/tools/flatten-bench/build.rs @@ -0,0 +1,71 @@ +//! Extracts the baseline `src/utils/json` sources with `git show :` into +//! OUT_DIR, byte for byte, and generates a one-line module declaration pointing at +//! them. Nothing in the extracted files is edited; the crate root provides the same +//! `crate::parseable`, `crate::event` and `crate::metadata` stubs to both copies. +//! +//! Override the revision with FLATTEN_BENCH_BASELINE_REV (default 547c5118, the +//! merge-base of fix/flatten-row-cap with main). + +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::{env, fs}; + +const DEFAULT_REV: &str = "547c5118"; +const FILES: &[&str] = &["mod.rs", "flatten.rs", "strict.rs"]; + +fn main() { + println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rerun-if-env-changed=FLATTEN_BENCH_BASELINE_REV"); + + let rev = env::var("FLATTEN_BENCH_BASELINE_REV").unwrap_or_else(|_| DEFAULT_REV.to_owned()); + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let repo_root = manifest_dir.join("../.."); + let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); + let json_dir = out_dir.join("baseline").join("json"); + fs::create_dir_all(&json_dir).unwrap(); + + for file in FILES { + let spec = format!("{rev}:src/utils/json/{file}"); + let output = Command::new("git") + .arg("-C") + .arg(&repo_root) + .args(["show", &spec]) + .output() + .expect("failed to run git"); + assert!( + output.status.success(), + "git show {spec} failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + fs::write(json_dir.join(file), &output.stdout).unwrap(); + } + + let commit = git(&repo_root, &["rev-parse", &format!("{rev}^{{commit}}")]); + let blob = git( + &repo_root, + &["rev-parse", &format!("{rev}:src/utils/json/flatten.rs")], + ); + println!("cargo:rustc-env=FLATTEN_BENCH_BASELINE_COMMIT={commit}"); + println!("cargo:rustc-env=FLATTEN_BENCH_BASELINE_FLATTEN_BLOB={blob}"); + + let mod_path = json_dir.join("mod.rs"); + fs::write( + out_dir.join("baseline_modules.rs"), + format!( + "#[path = {:?}]\npub mod baseline_json;\n", + mod_path.to_str().expect("OUT_DIR must be UTF-8") + ), + ) + .unwrap(); +} + +fn git(repo_root: &Path, args: &[&str]) -> String { + let output = Command::new("git") + .arg("-C") + .arg(repo_root) + .args(args) + .output() + .expect("failed to run git"); + assert!(output.status.success(), "git {args:?} failed"); + String::from_utf8(output.stdout).unwrap().trim().to_owned() +} diff --git a/tools/flatten-bench/http_smoke.py b/tools/flatten-bench/http_smoke.py new file mode 100644 index 000000000..85c093206 --- /dev/null +++ b/tools/flatten-bench/http_smoke.py @@ -0,0 +1,706 @@ +#!/usr/bin/env python3 +"""HTTP smoke test of the generic flattening limits against a real Parseable binary. + +Starts ` local-store` on 127.0.0.1 (ephemeral port, temporary data and staging +directories, small flatten caps), then for plain JSON (on `/api/v1/ingest` and +`/api/v1/logstream/{name}`), a known log format, a custom partitioned stream, Kinesis, OTEL +JSON logs and schema detection checks that: + - a request at the limits is accepted (200) and its rows are counted in the stream stats + - a request over a limit, including one that only fails on a late event, gets a 400 + naming the limit, the stream's event count does not change, and a sentinel field of an + earlier, valid event of that request does not show up in the stream schema. +Over-limit bytes include those of the structurally flattened rows, whose field names repeat +the names of their parents (`{"": {"f0": 0, ...}}` becomes +`{"_f0": 0, ...}`), even when the generic rows are small. + +The row-per-event and byte caps are set through the environment. The row-per-request cap is +set to 999 in the environment and to 6 on the command line, so every request-rows check +below also proves that the flag takes precedence. + +Python 3 standard library only. Build the server first, e.g. + P_BUILD_NO_UI=1 cargo build --bin parseable + python3 tools/flatten-bench/http_smoke.py target/debug/parseable +Exit status 0 when every check passed. +""" + +import argparse +import base64 +import json +import os +import socket +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +import uuid + +ROWS_PER_EVENT = 4 +ROWS_PER_REQUEST = 6 +BYTES_PER_REQUEST = 4096 +# Set in the environment and overridden by the command line flag +ROWS_PER_REQUEST_ENV = 999 +AUTH = "Basic " + base64.b64encode(b"admin:admin").decode() + +EVENT_MSG = f"more than {ROWS_PER_EVENT} rows (P_MAX_FLATTEN_ROWS_PER_EVENT)" +REQUEST_MSG = f"more than {ROWS_PER_REQUEST} rows (P_MAX_FLATTEN_ROWS_PER_REQUEST)" +BYTES_MSG = f"more than {BYTES_PER_REQUEST} bytes of JSON (P_MAX_FLATTEN_BYTES_PER_REQUEST)" + + +class Server: + def __init__(self, base): + self.base = base + + def request(self, method, path, body=None, headers=None): + data = None + all_headers = {"Authorization": AUTH, **(headers or {})} + if body is not None: + data = body if isinstance(body, bytes) else json.dumps(body).encode() + all_headers.setdefault("Content-Type", "application/json") + req = urllib.request.Request(self.base + path, data, all_headers, method=method) + try: + with urllib.request.urlopen(req, timeout=30) as resp: + return resp.status, resp.read().decode(errors="replace") + except urllib.error.HTTPError as e: + return e.code, e.read().decode(errors="replace") + + def schema(self, stream): + """Schema of the stream as text, empty when it does not exist yet""" + status, text = self.request("GET", f"/api/v1/logstream/{stream}/schema") + if status == 404: + return "" + if status != 200: + raise AssertionError(f"schema of {stream}: {status} {text}") + return text + + def events(self, stream): + status, text = self.request("GET", f"/api/v1/logstream/{stream}/stats") + if status == 404: + return 0 + if status != 200: + raise AssertionError(f"stats of {stream}: {status} {text}") + return json.loads(text)["ingestion"]["count"] + + +def free_port(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def start(binary, workdir, startup_timeout): + port = free_port() + for name in ("data", "staging"): + os.makedirs(os.path.join(workdir, name)) + # Nothing from the caller's P_* settings, so the caps below are the ones in effect + env = {k: v for k, v in os.environ.items() if not k.startswith("P_")} + env.update( + P_ADDR=f"127.0.0.1:{port}", + P_USERNAME="admin", + P_PASSWORD="admin", + P_FS_DIR=os.path.join(workdir, "data"), + P_STAGING_DIR=os.path.join(workdir, "staging"), + P_CHECK_UPDATE="false", + P_SEND_ANONYMOUS_USAGE_DATA="false", + P_RESOURCE_CHECK_ENABLED="false", + P_MAX_FLATTEN_ROWS_PER_EVENT=str(ROWS_PER_EVENT), + P_MAX_FLATTEN_ROWS_PER_REQUEST=str(ROWS_PER_REQUEST_ENV), + P_MAX_FLATTEN_BYTES_PER_REQUEST=str(BYTES_PER_REQUEST), + RUST_LOG=env.get("RUST_LOG", "warn"), + ) + log = open(os.path.join(workdir, "server.log"), "wb") + proc = subprocess.Popen( + [ + os.path.abspath(binary), + "local-store", + "--max-flatten-rows-per-request", + str(ROWS_PER_REQUEST), + ], + cwd=workdir, + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + server = Server(f"http://127.0.0.1:{port}") + deadline = time.monotonic() + startup_timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + break + try: + if server.request("GET", "/api/v1/liveness")[0] == 200: + return proc, log, server + except OSError: + pass + time.sleep(0.25) + stop(proc, log) + raise RuntimeError(f"server did not start (exit code {proc.returncode})") + + +def stop(proc, log): + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + log.close() + + +class Checks: + def __init__(self, server): + self.server = server + self.failures = [] + + def check(self, name, ok, detail=""): + print(f"{'PASS' if ok else 'FAIL'} {name}{'' if ok else ' ' + detail}") + if not ok: + self.failures.append(name) + + def accepted(self, name, stream, rows, method, path, body, headers=None, field=None): + """`field`: a field of the request that must show up in the schema""" + before = self.server.events(stream) + status, text = self.server.request(method, path, body, headers) + after = self.server.events(stream) + in_schema = field is None or field in self.server.schema(stream) + self.check( + f"{name}: accepted, {rows} events", + status == 200 and after - before == rows and in_schema, + f"status {status}, events {before} -> {after}, {field} in schema: {in_schema}, " + f"body {text[:300]!r}", + ) + + def rejected( + self, name, stream, message, method, path, body, headers=None, sentinel=None + ): + """`sentinel`: a field of an earlier, valid event that must not reach the schema""" + before = self.server.events(stream) if stream else 0 + status, text = self.server.request(method, path, body, headers) + after = self.server.events(stream) if stream else 0 + leaked = sentinel is not None and sentinel in self.server.schema(stream) + self.check( + f"{name}: 400, no events", + status == 400 and message in text and after == before and not leaked, + f"status {status}, events {before} -> {after}, {sentinel} in schema: {leaked}, " + f"body {text[:300]!r}", + ) + + +def unique(prefix): + return f"{prefix}{uuid.uuid4().hex[:10]}" + + +def arr(n): + return list(range(n)) + + +def compact_len(value): + return len(json.dumps(value, separators=(",", ":"))) + + +# Three 1-row events of ~1.5 KiB each: two fit in BYTES_PER_REQUEST, three do not +BIG = "x" * 1500 +assert 2 * compact_len({"s": BIG}) < BYTES_PER_REQUEST < 3 * compact_len({"s": BIG}) + + +# 32 small fields below a parent. Structural flattening names every one of them after the +# parent, `{parent}_f0` and so on +FIELDS = {f"f{i}": i for i in range(32)} +LONG_PARENT, MID_PARENT = "p" * 128, "q" * 64 + + +def parent_event(parent, repeated=False, **extra): + """32 fields below `parent`, in 4 rows when `repeated`""" + event = {parent: dict(FIELDS)} + if repeated: + event["r"] = arr(4) + event.update(extra) + return event + + +def generic_rows(event): + """Generic rows of an event whose only arrays are top level arrays of scalars""" + rows = [{}] + for key, value in event.items(): + alternatives = value if isinstance(value, list) else [value] + rows = [{**row, key: alt} for alt in alternatives for row in rows] + return rows + + +def structural(row, parent=None): + """`row` with nested objects flattened into fields joined with `_`""" + out = {} + for key, value in row.items(): + name = key if parent is None else f"{parent}_{key}" + if isinstance(value, dict): + out.update(structural(value, name)) + else: + out[name] = value + return out + + +def rows_len(rows): + return sum(compact_len(row) for row in rows) + + +def structural_len(event): + return rows_len([structural(row) for row in generic_rows(event)]) + + +# The generic rows of the long parent fit in the byte limit, the structurally flattened ones +# do not. One mid-parent event fits, two do not, though their generic rows would. +_long = parent_event(LONG_PARENT, True) +assert len(generic_rows(_long)) == ROWS_PER_EVENT +assert rows_len(generic_rows(_long)) < BYTES_PER_REQUEST < structural_len(_long) +_short = parent_event("p", True) +assert max(rows_len(generic_rows(_short)), structural_len(_short)) < BYTES_PER_REQUEST +_mid = parent_event(MID_PARENT) +assert 2 * rows_len(generic_rows(_mid)) < BYTES_PER_REQUEST +assert structural_len(_mid) < BYTES_PER_REQUEST < 2 * structural_len(_mid) +# A parent of quotes, escaped in every field name: one small generic row, over the byte +# limit only once structurally flattened +ESCAPED_PARENT = '"' * 64 +_escaped = parent_event(ESCAPED_PARENT) +assert len(generic_rows(_escaped)) == 1 +assert rows_len(generic_rows(_escaped)) < BYTES_PER_REQUEST < structural_len(_escaped) + + +def structural_checks(c, label, s, post, h, sentinel, wrap=lambda event: event): + """Long parent names, which only structural flattening repeats, on one ingestion path. + `wrap` turns a list of events into a request body""" + c.accepted( + f"{label} short parent", s, 4, *post, wrap([parent_event("p", True)]), h, "p_f31" + ) + c.rejected( + f"{label} long parent copied into every field", + s, + BYTES_MSG, + *post, + wrap([parent_event(LONG_PARENT, True)]), + h, + ) + c.rejected( + f"{label} escaped parent copied into every field", + s, + BYTES_MSG, + *post, + wrap([parent_event(ESCAPED_PARENT)]), + h, + ) + c.accepted( + f"{label} parent under byte limit", s, 1, *post, wrap([parent_event(MID_PARENT)]), h + ) + c.rejected( + f"{label} parent bytes, late event", + s, + BYTES_MSG, + *post, + wrap([parent_event(MID_PARENT, **{sentinel: 1}), parent_event(MID_PARENT)]), + h, + sentinel, + ) + + +def run_logstream_route(c): + # The same JSON ingestion, on the route of an existing stream + s = unique("flatroute") + status, text = c.server.request("PUT", f"/api/v1/logstream/{s}") + c.check("logstream route stream created", status == 200, f"{status} {text[:300]!r}") + post = ("POST", f"/api/v1/logstream/{s}") + control, sentinel = unique("control"), unique("sentinel") + c.accepted( + "logstream route event at row limit", + s, + 4, + *post, + {"a": arr(4), control: 1}, + None, + control, + ) + c.rejected("logstream route event over row limit", s, EVENT_MSG, *post, {"a": arr(5)}) + c.rejected( + "logstream route request rows, late event", + s, + REQUEST_MSG, + *post, + [{"a": arr(3), sentinel: 1}, {"a": 1}, {"a": 1}, {"b": arr(2)}], + None, + sentinel, + ) + c.rejected( + "logstream route request bytes, late event", + s, + BYTES_MSG, + *post, + [{"s": BIG, sentinel: 1}, {"s": BIG}, {"s": BIG}], + None, + sentinel, + ) + structural_checks(c, "logstream route", s, post, None, unique("sentinel")) + + +def run_json(c): + s = unique("flatjson") + h = {"X-P-Stream": s} + post = ("POST", "/api/v1/ingest") + control, sentinel = unique("control"), unique("sentinel") + c.accepted("json event at row limit", s, 4, *post, {"a": arr(4), control: 1}, h, control) + c.accepted("json request at row limit", s, 6, *post, [{"a": arr(2)}, {"b": arr(4)}], h) + c.rejected("json event over row limit", s, EVENT_MSG, *post, {"a": arr(5)}, h) + c.rejected( + "json sibling array product", s, EVENT_MSG, *post, {"a": arr(2), "b": arr(3)}, h + ) + c.rejected( + "json nested array product", s, EVENT_MSG, *post, {"a": [{"b": arr(3)}, {"c": arr(2)}]}, h + ) + c.rejected( + "json request rows, late event", + s, + REQUEST_MSG, + *post, + [{"a": arr(3), sentinel: 1}, {"a": 1}, {"a": 1}, {"b": arr(2)}], + h, + sentinel, + ) + c.accepted("json request under byte limit", s, 2, *post, [{"s": BIG}, {"s": BIG}], h) + c.rejected( + "json request bytes, late event", + s, + BYTES_MSG, + *post, + [{"s": BIG, sentinel: 1}, {"s": BIG}, {"s": BIG}], + h, + sentinel, + ) + # Escapes count: 700 quotes serialize to 1400 bytes + c.rejected( + "json request bytes with escapes", s, BYTES_MSG, *post, [{"s": '"' * 700}] * 3, h + ) + structural_checks(c, "json", s, post, h, unique("sentinel")) + + +def run_partitioned(c): + s = unique("flatpart") + status, text = c.server.request( + "PUT", f"/api/v1/logstream/{s}", headers={"X-P-Custom-Partition": "tenant"} + ) + c.check("partitioned stream created", status == 200, f"{status} {text[:300]!r}") + h = {"X-P-Stream": s} + post = ("POST", "/api/v1/ingest") + sentinel = unique("sentinel") + ev = lambda t: {"tenant": t, "x": arr(3)} + c.accepted("partitioned request at row limit", s, 6, *post, [ev("a"), ev("b")], h) + c.rejected( + "partitioned request rows, late partition", + s, + REQUEST_MSG, + *post, + [{**ev("a"), sentinel: 1}, ev("b"), ev("c")], + h, + sentinel, + ) + big = lambda t: {"tenant": t, "s": BIG} + c.accepted("partitioned request under byte limit", s, 2, *post, [big("a"), big("b")], h) + c.rejected( + "partitioned request bytes, late partition", + s, + BYTES_MSG, + *post, + [{**big("a"), sentinel: 1}, big("b"), big("c")], + h, + sentinel, + ) + c.rejected( + "partitioned event over row limit", s, EVENT_MSG, *post, [{"tenant": "a", "x": arr(5)}], h + ) + structural_checks( + c, + "partitioned", + s, + post, + h, + unique("sentinel"), + lambda events: [{"tenant": f"t{i}", **event} for i, event in enumerate(events)], + ) + + +def run_known_format(c): + # `block_log` events that already hold its fields (`body`, `timestamp`) are accepted as + # they are, and then flattened like plain JSON + s = unique("flatformat") + h = {"X-P-Stream": s, "X-P-Log-Source": "block_log"} + post = ("POST", "/api/v1/ingest") + sentinel = unique("sentinel") + ev = lambda n: {"body": "line", "timestamp": "2024-01-11T09:08:34.290Z", "a": arr(n)} + c.accepted("known format event at row limit", s, 4, *post, ev(4), h) + c.rejected("known format event over row limit", s, EVENT_MSG, *post, ev(5), h) + c.rejected( + "known format request rows, late event", + s, + REQUEST_MSG, + *post, + [{**ev(3), sentinel: 1}, ev(1), ev(3)], + h, + sentinel, + ) + + +def kinesis(records, request_id="req"): + return { + "requestId": request_id, + "timestamp": 1704964113659, + "records": [ + {"data": base64.b64encode(json.dumps(r).encode()).decode()} for r in records + ], + } + + +def run_kinesis(c): + s = unique("flatkinesis") + sentinel = unique("sentinel") + h = {"X-P-Stream": s, "X-P-Log-Source": "kinesis"} + post = ("POST", "/api/v1/ingest") + c.accepted("kinesis request at row limit", s, 6, *post, kinesis([{"a": arr(3)}] * 2), h) + # A decoded array is one event + c.rejected( + "kinesis decoded array over event limit", + s, + EVENT_MSG, + *post, + kinesis([[{"a": arr(3)}, {"b": arr(2)}]]), + h, + ) + c.rejected( + "kinesis request rows, late record", + s, + REQUEST_MSG, + *post, + kinesis([{"a": arr(3), sentinel: 1}, {"a": arr(3)}, {"a": arr(3)}]), + h, + sentinel, + ) + # The rows are tiny, the 1000 byte requestId copied into each of them is not + long_id = "r" * 1000 + c.accepted( + "kinesis long requestId under byte limit", s, 3, *post, kinesis([{"a": arr(3)}], long_id), h + ) + c.rejected( + "kinesis requestId bytes, late record", + s, + BYTES_MSG, + *post, + kinesis([{"a": arr(3), sentinel: 1}, {"a": arr(3)}], long_id), + h, + sentinel, + ) + # Each record is one event, and its metadata fields take ~50 bytes per row + structural_checks(c, "kinesis", s, post, h, unique("sentinel"), kinesis) + + +def otel_body(body): + """A string body, or a dict of key -> string for a key-value list body""" + if isinstance(body, str): + return {"stringValue": body} + values = [{"key": k, "value": {"stringValue": v}} for k, v in body.items()] + return {"kvlistValue": {"values": values}} + + +def otel(resources): + """resources: list of resources, each a list of scopes, each a list of bodies""" + now = str(time.time_ns()) + return { + "resourceLogs": [ + { + "scopeLogs": [ + { + "logRecords": [ + { + "timeUnixNano": now, + "severityNumber": 9, + "body": otel_body(body), + } + for body in scope + ] + } + for scope in resource + ] + } + for resource in resources + ] + } + + +def run_otel(c): + s = unique("flatotel") + h = {"X-P-Stream": s} + post = ("POST", "/v1/logs") + three = json.dumps({"a": arr(3)}) + control, sentinel = unique("control"), unique("sentinel") + # A record is one event; each of these bodies flattens to 3 rows merged into 1 record + c.accepted( + "otel request at row limit", + s, + 2, + *post, + otel([[[{"b": json.dumps({"a": arr(3), control: 1})}], [three]]]), + h, + control, + ) + c.rejected( + "otel body over event limit", s, EVENT_MSG, *post, otel([[[json.dumps({"a": arr(5)})]]]), h + ) + c.rejected( + "otel request rows, late resource", + s, + REQUEST_MSG, + *post, + otel([[[json.dumps({"a": arr(3), sentinel: 1})], [three]], [[three]]]), + h, + sentinel, + ) + c.accepted("otel body that is not JSON", s, 1, *post, otel([[["plain text"]]]), h) + + # A single row of 32 small fields. Every field is renamed to `{key}_{field}` when it is + # merged into the record, so a long key costs 32 times its length + wide = json.dumps({f"f{i}": i for i in range(32)}) + prefixed = lambda key, body: compact_len({f"{key}_{f}": v for f, v in json.loads(body).items()}) + long_key, mid_key = "k" * 128, "m" * 64 + assert prefixed("k", wide) < BYTES_PER_REQUEST < 32 * (len(long_key) + 1) + assert prefixed(mid_key, wide) < BYTES_PER_REQUEST < 2 * prefixed(mid_key, wide) + c.accepted("otel short body key", s, 1, *post, otel([[[{"k": wide}]]]), h, "k_f31") + c.rejected( + "otel long body key copied into every field", + s, + BYTES_MSG, + *post, + otel([[[{long_key: wide}]]]), + h, + ) + c.accepted("otel body key under byte limit", s, 1, *post, otel([[[{mid_key: wide}]]]), h) + c.rejected( + "otel body key bytes, late resource", + s, + BYTES_MSG, + *post, + otel([[[{mid_key: wide, "s": json.dumps({sentinel: 1})}]], [[{mid_key: wide}]]]), + h, + sentinel, + ) + + # Bodies with a parent object: structurally flattening the record repeats the body key + # and the parent in the name of every field below it, `{key}_{parent}_{field}` + nested = lambda parent, repeated=False, **extra: json.dumps( + parent_event(parent, repeated, **extra) + ) + c.accepted( + "otel nested body", s, 1, *post, otel([[[nested("p", True)]]]), h, "body_p_f31" + ) + c.rejected( + "otel nested body, long parent", + s, + BYTES_MSG, + *post, + otel([[[nested(LONG_PARENT, True)]]]), + h, + ) + # With a long key copied into each of the 32 nested fields; generic rows and the key + # copied into the root field alone are ~400 bytes + nested_key = lambda key: structural_len({key: parent_event("p")}) + assert compact_len({f"{long_key}_p": FIELDS}) < BYTES_PER_REQUEST < nested_key(long_key) + c.rejected( + "otel long body key copied into nested fields", + s, + BYTES_MSG, + *post, + otel([[[{long_key: nested("p")}]]]), + h, + ) + c.rejected( + "otel nested body bytes, late resource", + s, + BYTES_MSG, + *post, + otel([[[nested(MID_PARENT, **{sentinel: 1})]], [[nested(MID_PARENT)]]]), + h, + sentinel, + ) + + +def run_schema_detect(c): + path = ("POST", "/api/v1/logstream/schema/detect") + status, text = c.server.request(*path, [{"a": arr(2)}, {"b": arr(4)}]) + c.check("schema detect at row limit: 200", status == 200, f"{status} {text[:300]!r}") + c.rejected("schema detect event over row limit", None, EVENT_MSG, *path, {"a": arr(5)}) + c.rejected( + "schema detect request rows", None, REQUEST_MSG, *path, [{"a": arr(4)}, {"b": arr(3)}] + ) + c.rejected( + "schema detect request bytes", None, BYTES_MSG, *path, [{"s": BIG}, {"s": BIG}, {"s": BIG}] + ) + status, text = c.server.request(*path, parent_event("p", True)) + c.check( + "schema detect short parent: 200", + status == 200 and "p_f31" in text, + f"{status} {text[:300]!r}", + ) + c.rejected( + "schema detect long parent copied into every field", + None, + BYTES_MSG, + *path, + parent_event(LONG_PARENT, True), + ) + c.rejected( + "schema detect parent bytes, late event", + None, + BYTES_MSG, + *path, + [parent_event(MID_PARENT), parent_event(MID_PARENT)], + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("binary", help="path to the parseable server binary") + parser.add_argument("--startup-timeout", type=float, default=120, help="seconds") + parser.add_argument("--keep", action="store_true", help="keep the temporary directory") + args = parser.parse_args() + + workdir = tempfile.mkdtemp(prefix="parseable-flatten-smoke-") + proc = log = None + ok = False + try: + proc, log, server = start(args.binary, workdir, args.startup_timeout) + c = Checks(server) + runs = ( + run_json, + run_logstream_route, + run_known_format, + run_partitioned, + run_kinesis, + run_otel, + ) + for run in runs + (run_schema_detect,): + run(c) + ok = not c.failures + print(f"\n{'all checks passed' if ok else f'{len(c.failures)} check(s) failed'}") + finally: + if proc is not None: + stop(proc, log) + if not ok: + log_path = os.path.join(workdir, "server.log") + if os.path.exists(log_path): + with open(log_path, errors="replace") as f: + print("--- server.log (tail) ---\n" + "".join(f.readlines()[-40:])) + if args.keep or not ok: + print(f"kept {workdir}") + else: + import shutil + + shutil.rmtree(workdir, ignore_errors=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/flatten-bench/src/alloc_counter.rs b/tools/flatten-bench/src/alloc_counter.rs new file mode 100644 index 000000000..96d3865fc --- /dev/null +++ b/tools/flatten-bench/src/alloc_counter.rs @@ -0,0 +1,134 @@ +//! A global allocator wrapper that counts allocations made by the current thread while a +//! [`Measure`] guard is armed on it. Other threads (the test harness, stdout capture) are +//! never counted, so measurements stay isolated without `--test-threads=1`. +//! +//! Install it in each binary or integration test with +//! `#[global_allocator] static A: CountingAlloc = CountingAlloc;`. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; +use std::sync::{Mutex, MutexGuard}; + +pub struct CountingAlloc; + +thread_local! { + static ARMED: Cell = const { Cell::new(false) }; +} + +static COUNT: AtomicU64 = AtomicU64::new(0); +static TOTAL_BYTES: AtomicU64 = AtomicU64::new(0); +// Live bytes relative to arming; goes negative when pre-existing memory (e.g. a consumed +// input) is freed while armed +static CURRENT: AtomicI64 = AtomicI64::new(0); +static PEAK: AtomicI64 = AtomicI64::new(0); +// Serializes guards process-wide, since the counters are shared +static LOCK: Mutex<()> = Mutex::new(()); + +fn armed() -> bool { + ARMED.try_with(Cell::get).unwrap_or(false) +} + +fn on_alloc(size: usize) { + COUNT.fetch_add(1, Ordering::Relaxed); + TOTAL_BYTES.fetch_add(size as u64, Ordering::Relaxed); + let now = CURRENT.fetch_add(size as i64, Ordering::Relaxed) + size as i64; + PEAK.fetch_max(now, Ordering::Relaxed); +} + +fn on_dealloc(size: usize) { + CURRENT.fetch_sub(size as i64, Ordering::Relaxed); +} + +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() && armed() { + on_alloc(layout.size()); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc_zeroed(layout) }; + if !ptr.is_null() && armed() { + on_alloc(layout.size()); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) }; + if armed() { + on_dealloc(layout.size()); + } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new = unsafe { System.realloc(ptr, layout, new_size) }; + if !new.is_null() && armed() { + // Counted as a fresh allocation of `new_size` replacing the old block + on_dealloc(layout.size()); + on_alloc(new_size); + } + new + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AllocStats { + /// Number of alloc/alloc_zeroed/realloc calls + pub count: u64, + /// Cumulative requested bytes over those calls (a realloc counts its full new size); + /// an upper bound on new memory, not live memory + pub total_bytes: u64, + /// Highest live-bytes DELTA above the arming point: extra memory needed on top of what + /// was live before. Not an absolute process peak, and frees of memory allocated before + /// arming (e.g. a consumed input) lower it, so it can understate a consuming call's need + pub peak_bytes: i64, + /// Live-bytes delta when the guard was finished (negative if pre-existing memory was freed) + pub net_bytes: i64, +} + +/// Arms counting on the current thread until [`Measure::finish`] (or drop). +/// Blocks while another thread holds a guard. +pub struct Measure { + _lock: MutexGuard<'static, ()>, +} + +impl Measure { + pub fn start() -> Self { + let lock = LOCK.lock().unwrap_or_else(|p| p.into_inner()); + COUNT.store(0, Ordering::SeqCst); + TOTAL_BYTES.store(0, Ordering::SeqCst); + CURRENT.store(0, Ordering::SeqCst); + PEAK.store(0, Ordering::SeqCst); + ARMED.with(|a| a.set(true)); + Self { _lock: lock } + } + + pub fn finish(self) -> AllocStats { + ARMED.with(|a| a.set(false)); + let stats = AllocStats { + count: COUNT.load(Ordering::SeqCst), + total_bytes: TOTAL_BYTES.load(Ordering::SeqCst), + peak_bytes: PEAK.load(Ordering::SeqCst), + net_bytes: CURRENT.load(Ordering::SeqCst), + }; + drop(self); + stats + } +} + +impl Drop for Measure { + fn drop(&mut self) { + ARMED.with(|a| a.set(false)); + } +} + +/// Runs `f` with counting armed and returns its result with the stats. +pub fn measure(f: impl FnOnce() -> R) -> (R, AllocStats) { + let guard = Measure::start(); + let result = f(); + (result, guard.finish()) +} diff --git a/tools/flatten-bench/src/api.rs b/tools/flatten-bench/src/api.rs new file mode 100644 index 000000000..ac8fc0115 --- /dev/null +++ b/tools/flatten-bench/src/api.rs @@ -0,0 +1,175 @@ +//! The only place that names production signatures, so an API change is a one-file fix. +//! Baseline calls go through `baseline_json` (547c5118 source), current calls through +//! `current_json` (working tree source). Both use the ingestion configuration the HTTP +//! JSON path uses: schema V1, `LogSource::Json`, no time partition. + +use serde_json::Value; + +use crate::baseline_json as base; +use crate::current_json as cur; +use crate::event::format::LogSource; +use crate::fixtures::{BASELINE_MAX_ROWS, generic_row_upper_bound}; +use crate::metadata::SchemaVersion; + +pub use cur::flatten::{FlattenBudget, FlattenCost, FlattenLimits, JsonFlattenError, PipelineCost}; + +pub const UNLIMITED: FlattenLimits = FlattenLimits { + max_rows_per_event: usize::MAX, + max_rows_per_request: usize::MAX, + max_bytes_per_request: usize::MAX, +}; + +fn assert_baseline_safe(body: &Value) { + let bound = generic_row_upper_bound(body); + assert!( + bound <= BASELINE_MAX_ROWS, + "refusing to run the unbounded baseline on an input of up to {bound} rows" + ); +} + +/// An input checked once to expand to at most `BASELINE_MAX_ROWS` rows: the only form the +/// unbounded baseline accepts without walking it again. Clones keep the guarantee (same +/// value), so the benchmark checks each fixture once, outside any timed or counted region, +/// and the timed baseline calls do no extra work the current code does not do. +#[derive(Clone)] +pub struct BaselineInput(Value); + +impl BaselineInput { + /// Panics unless `body` is small enough for the unbounded baseline + pub fn check(body: Value) -> Self { + assert_baseline_safe(&body); + Self(body) + } + + pub fn value(&self) -> &Value { + &self.0 + } +} + +/// Baseline `convert_array_to_object`, the HTTP ingestion entry point, on a prevalidated input. +pub fn baseline_convert_checked( + input: BaselineInput, + custom_partition: Option<&String>, +) -> anyhow::Result> { + base::convert_array_to_object( + input.0, + None, + None, + custom_partition, + SchemaVersion::V1, + &LogSource::Json, + ) +} + +/// Baseline `generic_flattening` (unbounded) on a prevalidated input. +pub fn baseline_generic_checked(input: &BaselineInput) -> Vec { + base::flatten::generic_flattening(&input.0).expect("baseline generic_flattening never errors") +} + +/// Guarded convenience wrapper for tests: checks `body` on every call. Never time this. +pub fn baseline_convert( + body: Value, + custom_partition: Option<&String>, +) -> anyhow::Result> { + baseline_convert_checked(BaselineInput::check(body), custom_partition) +} + +/// Guarded convenience wrapper for tests: checks `value` on every call. Never time this. +pub fn baseline_generic(value: &Value) -> Vec { + assert_baseline_safe(value); + base::flatten::generic_flattening(value).expect("baseline generic_flattening never errors") +} + +/// Current `convert_array_to_object` with a fresh per-request budget. +pub fn current_convert( + body: Value, + custom_partition: Option<&String>, + limits: FlattenLimits, +) -> anyhow::Result> { + let mut budget = FlattenBudget::new(limits); + current_convert_with(body, custom_partition, &mut budget) +} + +pub fn current_convert_with( + body: Value, + custom_partition: Option<&String>, + budget: &mut FlattenBudget, +) -> anyhow::Result> { + cur::convert_array_to_object( + body, + None, + None, + custom_partition, + SchemaVersion::V1, + &LogSource::Json, + budget, + ) +} + +pub fn current_generic( + value: &Value, + limits: FlattenLimits, +) -> Result, JsonFlattenError> { + cur::flatten::generic_flattening_with_limits(value, limits) +} + +/// Generic-only preflight (`charge_generic_flattening`): a fresh request budget charged per +/// event (each element of a top level array is its own event), nothing expanded. The entry +/// point charges [`current_pipeline_charge`] instead. +pub fn current_charge( + value: &Value, + limits: FlattenLimits, +) -> Result { + let mut budget = FlattenBudget::new(limits); + cur::flatten::charge_generic_flattening(value, &mut budget)?; + Ok(budget.used()) +} + +/// `generic_flattening_cost`, which takes `value` as ONE event even when it is an array; +/// the entry point instead charges batches per element, see [`current_charge`]. +pub fn current_cost( + value: &Value, + limits: &FlattenLimits, +) -> Result { + cur::flatten::generic_flattening_cost(value, limits) +} + +/// `flatten_pipeline_cost`: generic and structurally flattened bytes of `value` taken as ONE +/// event, like [`current_cost`]. +pub fn current_pipeline_cost( + value: &Value, + limits: &FlattenLimits, +) -> Result { + cur::flatten::flatten_pipeline_cost(value, limits) +} + +/// Preflight exactly as the entry point runs it (`charge_flatten_pipeline`): per event like +/// [`current_charge`], each charged the larger of its generic and structurally flattened bytes. +pub fn current_pipeline_charge( + value: &Value, + limits: FlattenLimits, +) -> Result { + let mut budget = FlattenBudget::new(limits); + cur::flatten::charge_flatten_pipeline(value, &mut budget)?; + Ok(budget.used()) +} + +/// Structural flattening the entry point applies to the generic rows: objects get their +/// nested fields joined with `_`, other rows stay as they are. +pub fn current_structural(rows: &[Value]) -> Vec { + rows.iter() + .map(|row| { + let mut row = row.clone(); + if row.is_object() { + cur::flatten::flatten(&mut row, "_", None, None, None, false) + .expect("generic rows have no arrays"); + } + row + }) + .collect() +} + +/// Order-sensitive comparison key: serialization keeps map insertion order. +pub fn ordered(values: &[Value]) -> String { + serde_json::to_string(values).unwrap() +} diff --git a/tools/flatten-bench/src/bin/flatten_bench.rs b/tools/flatten-bench/src/bin/flatten_bench.rs new file mode 100644 index 000000000..3a617bb1e --- /dev/null +++ b/tools/flatten-bench/src/bin/flatten_bench.rs @@ -0,0 +1,384 @@ +//! Baseline (547c5118 source) vs current (working tree source) flattening benchmark. +//! +//! Run with `cargo run --release --bin flatten_bench` from tools/flatten-bench. +//! Env: FLATTEN_BENCH_SAMPLES (default 31), FLATTEN_BENCH_WARMUP (default 5), +//! FLATTEN_BENCH_SAMPLE_NS (target time per sample, default 2_000_000). +//! +//! Method: inputs are built once (no parsing is timed). Every baseline input is checked +//! ONCE to expand to <= 1000 rows (`BaselineInput::check`) before any baseline call and +//! outside every timed/counted region; timed baseline calls take the checked input and do +//! no extra walk. For the consuming entry point each sample pre-clones a batch of inputs +//! outside the timer, then times flattening plus dropping the output and the consumed +//! input, identically for both sides; borrowed APIs reuse one input and drop the output +//! inside the timer. Baseline and current samples alternate (AB, BA, ...) on the same +//! thread with identical inputs. Allocation stats come from one separate counted call per +//! implementation, never from timed ones; see `fmt_alloc` for what each number means. +//! Rejection cases run only against the current code. + +use std::hint::black_box; +use std::time::Instant; + +use flatten_bench::alloc_counter::{AllocStats, CountingAlloc, measure}; +use flatten_bench::api::{ + BaselineInput, FlattenLimits, baseline_convert_checked, baseline_generic_checked, + current_convert, current_generic, current_pipeline_charge, ordered, +}; +use flatten_bench::fixtures::{self, Case, accepted_cases}; +use serde_json::Value; + +#[global_allocator] +static ALLOC: CountingAlloc = CountingAlloc; + +const DEFAULT: FlattenLimits = FlattenLimits::DEFAULT; + +struct Config { + samples: usize, + warmup: usize, + sample_ns: u128, +} + +fn env_or(name: &str, default: u128) -> u128 { + std::env::var(name) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +fn loadavg() -> String { + std::fs::read_to_string("/proc/loadavg") + .map(|s| s.trim().to_owned()) + .unwrap_or_else(|_| "n/a".into()) +} + +/// An implementation under test: `Owned` consumes a pre-cloned input per call (the +/// ingestion entry point), `Borrowed` reads the same input every call. +/// `T` is `BaselineInput` for the baseline (so it can only ever see checked inputs) and +/// `Value` for the current code. +enum Imp<'a, T> { + Owned(&'a dyn Fn(T) -> usize), + Borrowed(&'a dyn Fn(&T) -> usize), +} + +impl Clone for Imp<'_, T> { + fn clone(&self) -> Self { + *self + } +} +impl Copy for Imp<'_, T> {} + +impl Imp<'_, T> { + /// One counted call, with any input clone made before counting starts + fn alloc_stats(self, input: &T) -> AllocStats { + match self { + Imp::Owned(f) => { + let input = input.clone(); + measure(move || f(input)).1 + } + Imp::Borrowed(f) => measure(|| f(input)).1, + } + } +} + +/// Runs `imp` `batch` times and returns the mean ns per call. Clones are made before the +/// timer starts; the timed region is the call, including dropping its output and, for +/// `Owned`, the consumed input. +fn time_batch(input: &T, batch: usize, checksum: &mut u64, imp: Imp) -> f64 { + match imp { + Imp::Owned(f) => { + let mut inputs: Vec = (0..batch).map(|_| input.clone()).collect(); + let start = Instant::now(); + // `drain` keeps the Vec buffer itself alive, so only per-call work is timed + for input in inputs.drain(..) { + *checksum = checksum.wrapping_add(black_box(f(black_box(input))) as u64); + } + let ns = start.elapsed().as_nanos() as f64 / batch as f64; + drop(inputs); + ns + } + Imp::Borrowed(f) => { + let start = Instant::now(); + for _ in 0..batch { + *checksum = checksum.wrapping_add(black_box(f(black_box(input))) as u64); + } + start.elapsed().as_nanos() as f64 / batch as f64 + } + } +} + +/// Batch size giving about `sample_ns` per sample, capped so the pre-cloned inputs stay +/// around 256 MiB of serialized size at most. +fn batch_size(cfg: &Config, input: &Value, one_ns: f64) -> usize { + let input_bytes = serde_json::to_vec(input).unwrap().len().max(1); + let memory_cap = ((256 << 20) / input_bytes).max(1); + ((cfg.sample_ns as f64 / one_ns.max(1.0)) as usize).clamp(1, 2000.min(memory_cap)) +} + +fn percentile(sorted: &[f64], p: f64) -> f64 { + let idx = ((sorted.len() - 1) as f64 * p).round() as usize; + sorted[idx] +} + +struct Summary { + median: f64, + p95: f64, +} + +fn summarize(mut samples: Vec) -> Summary { + samples.sort_by(f64::total_cmp); + Summary { + median: percentile(&samples, 0.5), + p95: percentile(&samples, 0.95), + } +} + +fn fmt_ns(ns: f64) -> String { + if ns >= 1e6 { + format!("{:.3}ms", ns / 1e6) + } else if ns >= 1e3 { + format!("{:.2}us", ns / 1e3) + } else { + format!("{ns:.0}ns") + } +} + +/// All numbers are for one call, measured on this thread only, in requested bytes (not RSS): +/// `allocs` alloc+realloc calls; `alloc_bytes` cumulative bytes requested (not live); +/// `peak_delta` highest live bytes above the level at call start, i.e. extra memory the +/// call needed on top of its input (freeing a consumed input during the call offsets it, +/// so it is not an absolute peak); `net_delta` live bytes at the end minus at the start +/// (negative when the call freed its consumed input and dropped its output). +fn fmt_alloc(stats: AllocStats) -> String { + format!( + "allocs={} alloc_bytes={} peak_delta={} net_delta={}", + stats.count, stats.total_bytes, stats.peak_bytes, stats.net_bytes + ) +} + +/// Times two implementations on the same input with alternating order. +fn compare( + cfg: &Config, + label: &str, + base_input: &BaselineInput, + input: &Value, + base: Imp, + cur: Imp, + checksum: &mut u64, +) { + debug_assert!(base_input.value() == input); + // Calibrate a batch size so one sample takes about `sample_ns` + let one = time_batch(base_input, 1, checksum, base).max(time_batch(input, 1, checksum, cur)); + let batch = batch_size(cfg, input, one); + + for _ in 0..cfg.warmup { + time_batch(base_input, batch, checksum, base); + time_batch(input, batch, checksum, cur); + } + let (mut b, mut c) = (Vec::new(), Vec::new()); + for i in 0..cfg.samples { + if i % 2 == 0 { + b.push(time_batch(base_input, batch, checksum, base)); + c.push(time_batch(input, batch, checksum, cur)); + } else { + c.push(time_batch(input, batch, checksum, cur)); + b.push(time_batch(base_input, batch, checksum, base)); + } + } + let (bs, cs) = (summarize(b), summarize(c)); + + let base_alloc = base.alloc_stats(base_input); + let cur_alloc = cur.alloc_stats(input); + + println!( + "{label:<44} batch={batch:<5} base median={:>10} p95={:>10} | cur median={:>10} p95={:>10} | cur/base={:.3}", + fmt_ns(bs.median), + fmt_ns(bs.p95), + fmt_ns(cs.median), + fmt_ns(cs.p95), + cs.median / bs.median, + ); + println!( + "{:<44} base {} | cur {}", + "", + fmt_alloc(base_alloc), + fmt_alloc(cur_alloc), + ); +} + +/// Times only the current implementation (rejections and preflight). +fn time_current(cfg: &Config, label: &str, input: &Value, cur: Imp, checksum: &mut u64) { + let one = time_batch(input, 1, checksum, cur); + let batch = batch_size(cfg, input, one); + for _ in 0..cfg.warmup { + time_batch(input, batch, checksum, cur); + } + let samples = (0..cfg.samples) + .map(|_| time_batch(input, batch, checksum, cur)) + .collect(); + let s = summarize(samples); + let alloc = cur.alloc_stats(input); + println!( + "{label:<44} batch={batch:<5} cur median={:>10} p95={:>10} | {}", + fmt_ns(s.median), + fmt_ns(s.p95), + fmt_alloc(alloc) + ); +} + +fn main() { + let cfg = Config { + samples: env_or("FLATTEN_BENCH_SAMPLES", 31) as usize, + warmup: env_or("FLATTEN_BENCH_WARMUP", 5) as usize, + sample_ns: env_or("FLATTEN_BENCH_SAMPLE_NS", 2_000_000), + }; + println!( + "baseline commit {} (flatten.rs blob {})", + flatten_bench::BASELINE_COMMIT, + flatten_bench::BASELINE_FLATTEN_BLOB + ); + println!( + "samples={} warmup={} sample_target={}ns debug_assertions={}", + cfg.samples, + cfg.warmup, + cfg.sample_ns, + cfg!(debug_assertions) + ); + println!("loadavg at start: {}", loadavg()); + if cfg!(debug_assertions) { + println!("WARNING: not a release build; numbers are meaningless"); + } + + let mut checksum = 0u64; + let cases: Vec = accepted_cases(); + // The one safety check per fixture, before any baseline call and outside all timing + let checked: Vec = cases + .iter() + .map(|case| BaselineInput::check(case.body.clone())) + .collect(); + + println!( + "\n== convert_array_to_object (ingestion entry, V1 Json, includes no-op fast paths) ==" + ); + for (case, checked) in cases.iter().zip(&checked) { + let partition = case.custom_partition.clone(); + let base_rows = baseline_convert_checked(checked.clone(), partition.as_ref()).unwrap(); + let cur_rows = current_convert(case.body.clone(), partition.as_ref(), DEFAULT).unwrap(); + assert_eq!( + ordered(&base_rows), + ordered(&cur_rows), + "{} output differs", + case.name + ); + let rows: usize = cur_rows + .iter() + .map(|v| v.as_array().map_or(1, Vec::len)) + .sum(); + let p = partition.clone(); + let base = move |v: BaselineInput| { + let out = baseline_convert_checked(v, p.as_ref()).unwrap(); + let n = out.len(); + drop(black_box(out)); + n + }; + let p = partition.clone(); + let cur = move |v: Value| { + let out = current_convert(v, p.as_ref(), DEFAULT).unwrap(); + let n = out.len(); + drop(black_box(out)); + n + }; + compare( + &cfg, + &format!("{} ({} rows, outputs equal)", case.name, rows), + checked, + &case.body, + Imp::Owned(&base), + Imp::Owned(&cur), + &mut checksum, + ); + } + + println!( + "\n== generic_flattening only (bounded vs unbounded expansion, no no-op fast path) ==" + ); + for (case, checked) in cases.iter().zip(&checked) { + let base_rows = baseline_generic_checked(checked); + let cur_rows = current_generic(&case.body, DEFAULT).unwrap(); + assert_eq!( + ordered(&base_rows), + ordered(&cur_rows), + "{} output differs", + case.name + ); + let base = |v: &BaselineInput| { + let out = baseline_generic_checked(v); + let n = out.len(); + drop(black_box(out)); + n + }; + let cur = |v: &Value| { + let out = current_generic(v, DEFAULT).unwrap(); + let n = out.len(); + drop(black_box(out)); + n + }; + compare( + &cfg, + &format!("{} ({} rows, outputs equal)", case.name, cur_rows.len()), + checked, + &case.body, + Imp::Borrowed(&base), + Imp::Borrowed(&cur), + &mut checksum, + ); + } + + println!( + "\n== preflight only, current: charge_flatten_pipeline (generic + structural bytes) on a fresh budget, per event as the entry point charges ==" + ); + for case in &cases { + let cost = |v: &Value| match current_pipeline_charge(v, DEFAULT) { + Ok(c) => c.rows ^ c.bytes, + Err(_) => 0, + }; + time_current( + &cfg, + case.name, + &case.body, + Imp::Borrowed(&cost), + &mut checksum, + ); + } + + println!("\n== rejections, current only (baseline would expand these unbounded) =="); + let rejections: Vec<(&str, Value)> = vec![ + ("hostile_cube_1000^3", fixtures::hostile_cube()), + ( + "deep_product_nested_arrays", + fixtures::deep_product_nested_arrays(), + ), + ( + "deep_product_nested_objects", + fixtures::deep_product_nested_objects(), + ), + ( + "large_repeated_string (600 rows)", + fixtures::large_repeated_string(), + ), + ( + "escaped_repeated_string (500 rows)", + fixtures::escaped_repeated_string(), + ), + ("just_over_event_limit", fixtures::just_over_event_limit()), + ]; + for (name, input) in &rejections { + let reject = |v: Value| { + let err = current_convert(v, None, DEFAULT).expect_err("must be rejected"); + drop(black_box(err)); + 1 + }; + time_current(&cfg, name, input, Imp::Owned(&reject), &mut checksum); + } + + println!("\nloadavg at end: {}", loadavg()); + println!("checksum {checksum}"); +} diff --git a/tools/flatten-bench/src/fixtures.rs b/tools/flatten-bench/src/fixtures.rs new file mode 100644 index 000000000..1071e2350 --- /dev/null +++ b/tools/flatten-bench/src/fixtures.rs @@ -0,0 +1,225 @@ +//! Deterministic inputs shared by the benchmark and the allocation tests. +//! +//! Accepted cases expand to at most 1000 rows, so they are safe to run through the +//! unbounded baseline; [`generic_row_upper_bound`] re-checks that before any baseline +//! call. Rejection cases must only ever be passed to the bounded current code. + +use serde_json::{Map, Value, json}; + +/// Configured defaults under test. +pub const EVENT_ROW_LIMIT: u64 = 1000; +pub const REQUEST_ROW_LIMIT: u64 = 10_000; +pub const EXPANDED_BYTES_LIMIT: u64 = 67_108_864; + +/// Hard cap for anything handed to the baseline. +pub const BASELINE_MAX_ROWS: u64 = 1000; + +pub struct Case { + pub name: &'static str, + pub body: Value, + /// Custom partition field, which routes through the per-element partitioned path + pub custom_partition: Option, +} + +/// Saturating upper bound on the rows generic flattening can produce for `value`: +/// sibling arrays multiply, array elements add, and nesting multiplies through. Used only +/// as a guard against accidentally handing an enormous expansion to the baseline. +pub fn generic_row_upper_bound(value: &Value) -> u64 { + match value { + Value::Array(arr) if arr.is_empty() => 1, + Value::Array(arr) => arr.iter().fold(0u64, |acc, item| { + acc.saturating_add(generic_row_upper_bound(item)) + }), + Value::Object(map) => map.values().fold(1u64, |acc, item| { + acc.saturating_mul(generic_row_upper_bound(item)) + }), + _ => 1, + } +} + +fn scalars(map: &mut Map, prefix: &str, n: usize, seed: usize) { + for i in 0..n { + let value = match (i + seed) % 5 { + 0 => json!(format!("value-{prefix}-{i}-{seed}")), + 1 => json!((i * 7919 + seed) as u64), + 2 => json!(((i + seed) as f64) * 1.25), + 3 => json!((i + seed) % 2 == 0), + _ => Value::Null, + }; + map.insert(format!("{prefix}field_{i}"), value); + } +} + +fn flat_object(fields: usize, seed: usize) -> Value { + let mut map = Map::new(); + map.insert("timestamp".into(), json!("2026-10-26T12:00:00.000Z")); + map.insert("tenant".into(), json!(format!("t{}", seed % 4))); + scalars(&mut map, "", fields, seed); + Value::Object(map) +} + +fn nested_object(seed: usize) -> Value { + let mut map = Map::new(); + map.insert("timestamp".into(), json!("2026-10-26T12:00:00.000Z")); + scalars(&mut map, "", 4, seed); + let mut l1 = Map::new(); + scalars(&mut l1, "l1_", 3, seed); + let mut l2 = Map::new(); + scalars(&mut l2, "l2_", 3, seed); + let mut l3 = Map::new(); + scalars(&mut l3, "l3_", 3, seed); + l2.insert("deeper".into(), Value::Object(l3)); + l1.insert("inner".into(), Value::Object(l2)); + map.insert("outer".into(), Value::Object(l1)); + let mut side = Map::new(); + scalars(&mut side, "s_", 2, seed); + map.insert("side".into(), Value::Object(side)); + Value::Object(map) +} + +/// Inputs that every implementation must accept; each expands to <= 1000 rows. +pub fn accepted_cases() -> Vec { + let mut cases = Vec::new(); + + cases.push(Case { + name: "flat_single", + body: flat_object(20, 1), + custom_partition: None, + }); + + cases.push(Case { + name: "flat_batch_1000_noop", + body: Value::Array((0..1000).map(|i| flat_object(12, i)).collect()), + custom_partition: None, + }); + + cases.push(Case { + name: "nested_no_arrays_200", + body: Value::Array((0..200).map(nested_object).collect()), + custom_partition: None, + }); + + let mut moderate = Map::new(); + scalars(&mut moderate, "", 8, 3); + moderate.insert( + "items".into(), + Value::Array( + (0..100) + .map(|i| { + let mut item = Map::new(); + scalars(&mut item, "item_", 5, i); + Value::Object(item) + }) + .collect(), + ), + ); + cases.push(Case { + name: "moderate_array_100", + body: Value::Object(moderate), + custom_partition: None, + }); + + let mut product = Map::new(); + scalars(&mut product, "", 6, 5); + product.insert( + "a".into(), + Value::Array((0..10).map(|i| json!(i)).collect()), + ); + product.insert( + "b".into(), + Value::Array( + (0..10) + .map(|i| json!({"x": i, "y": format!("y{i}")})) + .collect(), + ), + ); + product.insert( + "c".into(), + Value::Array((0..10).map(|i| json!(format!("c{i}"))).collect()), + ); + cases.push(Case { + name: "product_10x10x10", + body: Value::Object(product), + custom_partition: None, + }); + + cases.push(Case { + name: "partitioned_batch_50x4", + body: Value::Array( + (0..50) + .map(|i| { + let mut obj = flat_object(6, i).as_object().unwrap().clone(); + obj.insert( + "items".into(), + Value::Array((0..4).map(|j| json!({"k": j, "v": i * j})).collect()), + ); + Value::Object(obj) + }) + .collect(), + ), + custom_partition: Some("tenant".into()), + }); + + for case in &cases { + assert!( + generic_row_upper_bound(&case.body) <= BASELINE_MAX_ROWS, + "{} exceeds the baseline row cap", + case.name + ); + } + cases +} + +fn ints(n: u64) -> Value { + Value::Array((0..n).map(|i| json!(i)).collect()) +} + +/// `{"a": [0..1000], "b": [0..1000], "c": [0..1000]}`: 10^9 rows if expanded. +pub fn hostile_cube() -> Value { + json!({"a": ints(1000), "b": ints(1000), "c": ints(1000)}) +} + +/// Rows multiply through nesting: 32 elements of `a`, each a 32x32 product, times a +/// sibling `d` of 32, about 1M rows from a ~2K-value input. +pub fn deep_product_nested_arrays() -> Value { + let inner: Vec = (0..32) + .map(|_| json!({"b": ints(32), "c": ints(32)})) + .collect(); + json!({"a": inner, "d": ints(32)}) +} + +/// Arrays under nested objects multiply like siblings: 100^3 rows. +pub fn deep_product_nested_objects() -> Value { + json!({"x": {"y": {"z": ints(100)}}, "w": {"v": ints(100)}, "u": ints(100)}) +} + +/// Legal row count (600) whose expanded compact JSON exceeds 64 MiB because a 128 KiB +/// string repeats on every row. +pub fn large_repeated_string() -> Value { + json!({"blob": "x".repeat(128 * 1024), "id": ints(600)}) +} + +/// Legal row count (500) where the raw string is only ~34 MiB repeated, but escaping each +/// `"` doubles it to ~68 MiB of compact JSON, so only an escape-aware estimate rejects it. +pub fn escaped_repeated_string() -> Value { + json!({"blob": "\"".repeat(70_000), "id": ints(500)}) +} + +/// One row over the per-event limit. +pub fn just_over_event_limit() -> Value { + json!({"a": ints(EVENT_ROW_LIMIT + 1)}) +} + +/// Limits of the HTTP smoke test, small enough for [`long_parent`] to matter. +pub const SMALL_HTTP_LIMITS: (usize, usize, usize) = (4, 6, 4096); + +/// `{parent: {"f0": 0, ..., "f31": 31}, "r": [0, 1, 2, 3]}`: 4 small generic rows, but +/// structural flattening names all 32 fields of each after `parent`. With a 128 character +/// parent the generic rows take ~1.6 KiB and the structurally flattened ones ~17 KiB. +pub fn long_parent(parent: &str) -> Value { + let fields: Map = (0..32).map(|i| (format!("f{i}"), json!(i))).collect(); + let mut map = Map::new(); + map.insert(parent.to_owned(), Value::Object(fields)); + map.insert("r".to_owned(), ints(4)); + Value::Object(map) +} diff --git a/tools/flatten-bench/src/lib.rs b/tools/flatten-bench/src/lib.rs new file mode 100644 index 000000000..373553b99 --- /dev/null +++ b/tools/flatten-bench/src/lib.rs @@ -0,0 +1,86 @@ +//! Compiles the production JSON flattening sources twice, unmodified: +//! +//! - `current_json`: the working tree's `src/utils/json/mod.rs` (and its `flatten`/`strict` +//! children), included by path. +//! - `baseline_json`: the same files at the baseline revision, extracted by `build.rs`. +//! +//! Both resolve `crate::parseable::PARSEABLE`, `crate::event::format::LogSource` and +//! `crate::metadata::SchemaVersion` to the stubs below, which mirror the server's types +//! and CLI defaults closely enough for the flattening paths and nothing more. + +#![allow(dead_code, unused_imports, clippy::all)] + +pub mod alloc_counter; +pub mod api; +pub mod fixtures; + +#[path = "../../../src/utils/json/mod.rs"] +pub mod current_json; + +include!(concat!(env!("OUT_DIR"), "/baseline_modules.rs")); + +/// Revision and blob the baseline copy was extracted from. +pub const BASELINE_COMMIT: &str = env!("FLATTEN_BENCH_BASELINE_COMMIT"); +pub const BASELINE_FLATTEN_BLOB: &str = env!("FLATTEN_BENCH_BASELINE_FLATTEN_BLOB"); + +/// Subset of `crate::cli` read by `src/utils/json`. +pub mod cli { + pub const DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT: usize = 1000; + pub const DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST: usize = 10_000; + pub const DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST: usize = 64 * 1024 * 1024; + + /// Fields of `crate::cli::Options` read by `src/utils/json`, with the CLI defaults. + pub struct Options { + pub event_flatten_level: usize, + pub event_max_chunk_age: u64, + pub max_flatten_rows_per_event: usize, + pub max_flatten_rows_per_request: usize, + pub max_flatten_bytes_per_request: usize, + } +} + +pub mod parseable { + use crate::cli::*; + + pub struct Parseable { + pub options: Options, + } + + pub static PARSEABLE: Parseable = Parseable { + options: Options { + event_flatten_level: 10, + event_max_chunk_age: 1, + max_flatten_rows_per_event: DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT, + max_flatten_rows_per_request: DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST, + max_flatten_bytes_per_request: DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST, + }, + }; +} + +pub mod event { + pub mod format { + /// Mirror of `crate::event::format::LogSource` (serde attributes omitted). + #[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] + pub enum LogSource { + Kinesis, + OtelLogs, + OtelMetrics, + OtelTraces, + Pmeta, + #[default] + Json, + Custom(String), + } + } +} + +pub mod metadata { + /// Mirror of `crate::metadata::SchemaVersion`. + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] + #[non_exhaustive] + pub enum SchemaVersion { + #[default] + V0, + V1, + } +} diff --git a/tools/flatten-bench/tests/alloc_limits.rs b/tools/flatten-bench/tests/alloc_limits.rs new file mode 100644 index 000000000..d0b5d01bb --- /dev/null +++ b/tools/flatten-bench/tests/alloc_limits.rs @@ -0,0 +1,477 @@ +//! Allocation bounds for rejected expansions, and preflight/actual agreement for accepted +//! ones, against the working tree's `src/utils/json` sources. +//! +//! Inputs are built before counting starts; only allocations made by the test's own +//! thread while a guard is armed are counted. Run with RUST_LIB_BACKTRACE=0 so anyhow +//! does not capture (and allocate) a backtrace for the rejected requests. + +use flatten_bench::alloc_counter::{AllocStats, CountingAlloc, measure}; +use flatten_bench::api::{ + FlattenBudget, FlattenCost, FlattenLimits, JsonFlattenError, UNLIMITED, baseline_generic, + current_charge, current_convert, current_convert_with, current_cost, current_generic, + current_pipeline_charge, current_pipeline_cost, current_structural, ordered, +}; +use flatten_bench::fixtures::{self, accepted_cases, generic_row_upper_bound}; +use serde_json::{Value, json}; + +#[global_allocator] +static ALLOC: CountingAlloc = CountingAlloc; + +const MAX_REJECT_BYTES: u64 = 64 * 1024; +const DEFAULT: FlattenLimits = FlattenLimits::DEFAULT; + +fn report(label: &str, stats: AllocStats) { + eprintln!( + "{label}: allocs={} alloc_bytes={} peak_delta={} net_delta={}", + stats.count, stats.total_bytes, stats.peak_bytes, stats.net_bytes + ); +} + +fn assert_small(label: &str, stats: AllocStats) { + report(label, stats); + assert!( + stats.total_bytes < MAX_REJECT_BYTES, + "{label} allocated {} bytes while rejecting", + stats.total_bytes + ); +} + +/// The core (non-anyhow) rejection paths build nothing at all: not even the error. +fn assert_zero(label: &str, stats: AllocStats) { + report(label, stats); + assert_eq!( + (stats.count, stats.total_bytes, stats.peak_bytes), + (0, 0, 0), + "{label} allocated while rejecting" + ); +} + +fn limit_error(err: &anyhow::Error) -> &JsonFlattenError { + let err = err + .downcast_ref::() + .unwrap_or_else(|| panic!("expected a JsonFlattenError, got {err:#}")); + assert!(err.is_limit_exceeded(), "expected a limit error, got {err}"); + err +} + +fn hostile_inputs() -> Vec<(&'static str, Value)> { + vec![ + ("hostile_cube_1000^3", fixtures::hostile_cube()), + ( + "deep_product_nested_arrays", + fixtures::deep_product_nested_arrays(), + ), + ( + "deep_product_nested_objects", + fixtures::deep_product_nested_objects(), + ), + ("large_repeated_string", fixtures::large_repeated_string()), + ( + "escaped_repeated_string", + fixtures::escaped_repeated_string(), + ), + ("just_over_event_limit", fixtures::just_over_event_limit()), + ] +} + +#[test] +fn preflight_rejects_without_allocating() { + for (name, input) in hostile_inputs() { + let (result, stats) = measure(|| current_cost(&input, &DEFAULT)); + assert!( + matches!(&result, Err(err) if err.is_limit_exceeded()), + "{name}: {result:?}" + ); + assert_zero(&format!("preflight {name}"), stats); + + // Per-event charging on a fresh request budget, as the entry point does it + let (result, stats) = measure(|| current_charge(&input, DEFAULT)); + assert!( + matches!(&result, Err(err) if err.is_limit_exceeded()), + "{name}: {result:?}" + ); + assert_zero(&format!("charge {name}"), stats); + + let (result, stats) = measure(|| current_pipeline_charge(&input, DEFAULT)); + assert!( + matches!(&result, Err(err) if err.is_limit_exceeded()), + "{name}: {result:?}" + ); + assert_zero(&format!("pipeline charge {name}"), stats); + } +} + +#[test] +fn bounded_generic_flattening_rejects_cheaply() { + for (name, input) in hostile_inputs() { + let mut budget = FlattenBudget::new(DEFAULT); + let (result, stats) = measure(|| { + let result = + flatten_bench::current_json::flatten::generic_flattening(&input, &mut budget); + result.map(|rows| rows.len()) + }); + assert!( + matches!(&result, Err(err) if err.is_limit_exceeded()), + "{name}: {result:?}" + ); + assert_zero(&format!("generic_flattening {name}"), stats); + assert_eq!(budget.used(), FlattenCost::default(), "{name}"); + } +} + +#[test] +fn ingestion_entry_point_rejects_cheaply() { + for (name, input) in hostile_inputs() { + // The consumed input is freed inside the guard, which only lowers `net` + let (result, stats) = + measure(|| current_convert(input, None, DEFAULT).map(|rows| rows.len())); + let err = result.expect_err(name); + limit_error(&err); + // `err` (an anyhow box, 112 bytes today) is allocated inside the guard, so + // this entry point is held to the 64 KiB bound rather than to zero + assert_small(&format!("convert_array_to_object {name}"), stats); + } +} + +#[test] +fn expected_error_kinds() { + let kind = |value: &Value| match current_cost(value, &DEFAULT) { + Err(JsonFlattenError::RowsPerEventLimitExceeded(_)) => "rows/event", + Err(JsonFlattenError::BytesPerRequestLimitExceeded(_)) => "bytes/request", + other => panic!("unexpected {other:?}"), + }; + assert_eq!(kind(&fixtures::hostile_cube()), "rows/event"); + assert_eq!(kind(&fixtures::deep_product_nested_arrays()), "rows/event"); + assert_eq!(kind(&fixtures::deep_product_nested_objects()), "rows/event"); + assert_eq!(kind(&fixtures::just_over_event_limit()), "rows/event"); + // 600 and 500 rows are legal; only the byte estimate (with escapes) rejects these + assert_eq!(kind(&fixtures::large_repeated_string()), "bytes/request"); + assert_eq!(kind(&fixtures::escaped_repeated_string()), "bytes/request"); +} + +#[test] +fn insufficient_remaining_request_budget() { + let product = accepted_cases() + .into_iter() + .find(|case| case.name == "product_10x10x10") + .unwrap() + .body; + + // 9_995 rows already used by earlier events of the same request + let mut budget = FlattenBudget::new(DEFAULT); + for _ in 0..9 { + budget + .charge_event(FlattenCost { + rows: 1000, + bytes: 0, + }) + .unwrap(); + } + budget + .charge_event(FlattenCost { + rows: 995, + bytes: 0, + }) + .unwrap(); + let (result, stats) = measure(|| { + flatten_bench::current_json::flatten::charge_generic_flattening(&product, &mut budget) + }); + assert!( + matches!( + result, + Err(JsonFlattenError::RowsPerRequestLimitExceeded(_)) + ), + "{result:?}" + ); + assert_zero("core charge, remaining rows insufficient", stats); + let (result, stats) = + measure(|| current_convert_with(product, None, &mut budget).map(|r| r.len())); + let err = result.expect_err("1000 rows cannot fit in the remaining 5"); + assert!( + matches!( + limit_error(&err), + JsonFlattenError::RowsPerRequestLimitExceeded(_) + ), + "{err:#}" + ); + assert_small("remaining rows insufficient", stats); + assert_eq!(budget.used().rows, 9_995, "a failed charge must not commit"); + + // Bytes nearly exhausted by earlier events + let mut budget = FlattenBudget::new(DEFAULT); + let filler = DEFAULT.max_bytes_per_request - 16; + budget + .charge_event(FlattenCost { + rows: 1, + bytes: filler, + }) + .unwrap(); + let small = json!({"message": "this row is longer than sixteen bytes"}); + let (result, stats) = measure(|| { + flatten_bench::current_json::flatten::charge_generic_flattening(&small, &mut budget) + }); + assert!( + matches!( + result, + Err(JsonFlattenError::BytesPerRequestLimitExceeded(_)) + ), + "{result:?}" + ); + assert_zero("core charge, remaining bytes insufficient", stats); + let (result, stats) = + measure(|| current_convert_with(small, None, &mut budget).map(|r| r.len())); + let err = result.expect_err("row does not fit in the remaining 16 bytes"); + assert!( + matches!( + limit_error(&err), + JsonFlattenError::BytesPerRequestLimitExceeded(_) + ), + "{err:#}" + ); + assert_small("remaining bytes insufficient", stats); + assert_eq!(budget.used().bytes, filler); +} + +#[test] +fn request_budget_is_shared_across_partitions() { + // 11 partitioned elements of 1000 rows each: every event is legal, the request is not + let body = Value::Array( + (0..11) + .map(|i| json!({"tenant": format!("t{i}"), "a": (0..1000).collect::>()})) + .collect(), + ); + let partition = "tenant".to_owned(); + let (result, stats) = + measure(|| current_convert(body, Some(&partition), DEFAULT).map(|r| r.len())); + let err = result.expect_err("11_000 rows exceed the 10_000 request limit"); + assert!( + matches!( + limit_error(&err), + JsonFlattenError::RowsPerRequestLimitExceeded(_) + ), + "{err:#}" + ); + // Not asserted: earlier partitions may be materialized before the 11th is charged + report("partitioned request over limit", stats); +} + +/// Small deterministic generator of mixed shapes, bounded far below 1000 rows. +fn generated(seed: u64, count: usize) -> Vec { + let mut state = seed | 1; + let mut next = move |n: u64| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state % n + }; + fn build(next: &mut dyn FnMut(u64) -> u64, depth: u32) -> Value { + const KEYS: &[&str] = &["a", "b", "a_b", "c", "", "q\"\\\n", "ü"]; + let roll = next(10); + if depth == 0 || roll < 4 { + return match next(8) { + 0 => Value::Null, + 1 => json!(true), + 2 => json!(-3), + 3 => json!(1.5e300), + 4 => json!(""), + 5 => json!("esc\"\\\u{1}\t ünï"), + 6 => json!(u64::MAX), + _ => json!(i64::MIN), + }; + } + if roll < 7 { + let mut map = serde_json::Map::new(); + for _ in 0..next(4) { + let key = KEYS[next(KEYS.len() as u64) as usize]; + map.insert(key.to_owned(), build(&mut *next, depth - 1)); + } + Value::Object(map) + } else { + let len = next(4); + let mut arr = Vec::new(); + for _ in 0..len { + arr.push(build(&mut *next, depth - 1)); + } + Value::Array(arr) + } + } + (0..count).map(|_| build(&mut next, 5)).collect() +} + +#[test] +fn preflight_matches_actual_output() { + let mut inputs: Vec = accepted_cases().into_iter().map(|case| case.body).collect(); + inputs.extend(generated(0x5eed, 2000)); + inputs.push(json!([[{"a": 1}], {"b": [1, 2]}, [], 3])); + let mut checked = 0; + for input in inputs { + if generic_row_upper_bound(&input) > 1000 { + continue; + } + let actual = current_generic(&input, UNLIMITED).expect("unlimited"); + let cost = current_cost(&input, &UNLIMITED).expect("unlimited"); + let bytes: usize = actual + .iter() + .map(|row| serde_json::to_vec(row).unwrap().len()) + .sum(); + assert_eq!( + (cost.rows, cost.bytes), + (actual.len(), bytes), + "cost mismatch for {input}" + ); + // The pipeline cost has the same generic part, and covers the structurally + // flattened rows: exactly, unless joined field names collide + let pipeline = current_pipeline_cost(&input, &UNLIMITED).expect("unlimited"); + assert_eq!( + (pipeline.rows, pipeline.generic_bytes), + (cost.rows, cost.bytes) + ); + let structural = current_structural(&actual); + let structural_bytes: usize = structural + .iter() + .map(|row| serde_json::to_vec(row).unwrap().len()) + .sum(); + let structural_fields: usize = structural + .iter() + .filter_map(Value::as_object) + .map(|row| row.len()) + .sum(); + if structural_fields == pipeline.structural_fields { + assert_eq!(pipeline.structural_bytes, structural_bytes, "{input}"); + } else { + assert!(pipeline.structural_bytes > structural_bytes, "{input}"); + } + assert!(pipeline.charged().bytes >= bytes.max(structural_bytes)); + // Bounded output is identical to the baseline 547c5118 output + assert_eq!( + ordered(&actual), + ordered(&baseline_generic(&input)), + "{input}" + ); + checked += 1; + } + eprintln!("checked {checked} inputs"); + assert!(checked > 1000); +} + +#[test] +fn accepted_cases_match_baseline_at_entry_point() { + for case in accepted_cases() { + let partition = case.custom_partition.as_ref(); + let base = flatten_bench::api::baseline_convert(case.body.clone(), partition).unwrap(); + let cur = current_convert(case.body, partition, DEFAULT).unwrap(); + assert_eq!(ordered(&base), ordered(&cur), "{}", case.name); + } +} + +#[test] +fn counter_sees_allocations() { + let (len, stats) = measure(|| vec![0u8; 100_000].len()); + assert_eq!(len, 100_000); + assert!( + stats.total_bytes >= 100_000 && stats.peak_bytes >= 100_000, + "{stats:?}" + ); +} + +/// The harness stubs `crate::cli`; make sure its defaults are the server's, and the +/// approved ones. +#[test] +fn stub_defaults_match_server() { + let cli = include_str!("../../../src/cli.rs"); + for line in [ + "pub const DEFAULT_MAX_FLATTEN_ROWS_PER_EVENT: usize = 1000;", + "pub const DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST: usize = 10_000;", + "pub const DEFAULT_MAX_FLATTEN_BYTES_PER_REQUEST: usize = 64 * 1024 * 1024;", + ] { + assert!(cli.contains(line), "src/cli.rs no longer has `{line}`"); + } + assert_eq!( + FlattenLimits::DEFAULT, + FlattenLimits { + max_rows_per_event: fixtures::EVENT_ROW_LIMIT as usize, + max_rows_per_request: fixtures::REQUEST_ROW_LIMIT as usize, + max_bytes_per_request: fixtures::EXPANDED_BYTES_LIMIT as usize, + } + ); +} + +#[test] +#[should_panic(expected = "refusing to run the unbounded baseline")] +fn baseline_input_check_refuses_hostile_input() { + // Only walks the input to bound its rows; nothing reaches the baseline + flatten_bench::api::BaselineInput::check(flatten_bench::fixtures::hostile_cube()); +} + +#[test] +fn baseline_input_check_accepts_every_benchmark_case() { + for case in flatten_bench::fixtures::accepted_cases() { + let checked = flatten_bench::api::BaselineInput::check(case.body.clone()); + assert!(checked.value() == &case.body, "{}", case.name); + } +} + +fn small_limits() -> FlattenLimits { + let (rows_per_event, rows_per_request, bytes) = fixtures::SMALL_HTTP_LIMITS; + FlattenLimits { + max_rows_per_event: rows_per_event, + max_rows_per_request: rows_per_request, + max_bytes_per_request: bytes, + } +} + +#[test] +fn structural_field_names_are_rejected_without_expanding() { + let small = small_limits(); + let input = fixtures::long_parent(&"p".repeat(128)); + // The generic rows alone fit + assert!(current_cost(&input, &small).is_ok()); + + let (result, stats) = measure(|| current_pipeline_cost(&input, &small)); + assert!( + matches!( + result, + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + ), + "{result:?}" + ); + assert_zero("structural preflight, long parent", stats); + let (result, stats) = measure(|| current_pipeline_charge(&input, small)); + assert!( + matches!( + result, + Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) + ), + "{result:?}" + ); + assert_zero("structural charge, long parent", stats); + + for partition in [None, Some("r".to_owned())] { + let input = input.clone(); + let mut budget = FlattenBudget::new(small); + let (result, stats) = measure(|| { + current_convert_with(input, partition.as_ref(), &mut budget).map(|rows| rows.len()) + }); + let err = result.expect_err("long parent names exceed 4096 bytes"); + assert!( + matches!( + limit_error(&err), + JsonFlattenError::BytesPerRequestLimitExceeded(4096) + ), + "{err:#}" + ); + // No row was expanded: the structural rows alone would take ~17 KiB + assert_small( + &format!("convert_array_to_object long parent, partition {partition:?}"), + stats, + ); + assert!(stats.total_bytes < 1024, "{stats:?}"); + assert_eq!(budget.used(), FlattenCost::default()); + } + + // A short parent is accepted, as by the baseline + let short = fixtures::long_parent("p"); + let cur = current_convert(short.clone(), None, small).unwrap(); + let base = flatten_bench::api::baseline_convert(short, None).unwrap(); + assert_eq!(ordered(&cur), ordered(&base)); +} From 7fe8a17e73e752493da46f970b63764aa413f016 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:03:43 -0500 Subject: [PATCH 2/2] fix(stats): stream bounded rows and clarify ingest limits --- docs/json-flattening-limits.md | 117 ++++- src/cli.rs | 2 +- src/handlers/http/kinesis.rs | 2 + src/handlers/http/modal/utils/ingest_utils.rs | 6 +- src/otel/logs.rs | 60 +-- src/storage/field_stats.rs | 417 ++++++++++++++---- src/storage/object_storage.rs | 8 +- tools/flatten-bench/README.md | 31 +- tools/flatten-bench/RESULTS.md | 46 +- tools/flatten-bench/build.rs | 6 +- tools/flatten-bench/http_smoke.py | 99 ++++- tools/flatten-bench/src/alloc_counter.rs | 96 +++- tools/flatten-bench/tests/alloc_limits.rs | 124 +++++- 13 files changed, 816 insertions(+), 198 deletions(-) diff --git a/docs/json-flattening-limits.md b/docs/json-flattening-limits.md index 04bc08b91..073168c34 100644 --- a/docs/json-flattening-limits.md +++ b/docs/json-flattening-limits.md @@ -3,9 +3,17 @@ Generic flattening turns every nested array of an event into rows: an event becomes one row for each combination of the values of its arrays. Sibling arrays multiply, so `{"a": [0..999], "b": [0..999], "c": [0..999]}`, about 11 KiB of JSON, would expand into -10^9 rows. Three limits cap that expansion. Parseable checks them before it builds any row, -and rejects the whole request with `400 Bad Request` when one is exceeded. Nothing from a -rejected request is ingested. +10^9 rows. Three limits cap that expansion. + +Parseable checks each event against the limits before it builds the rows of that event: it +counts the rows and bytes the event expands into without building them. When an event would +go over a limit, the whole request is rejected with `400 Bad Request`. None of its events +are stored, and none of its fields are added to the dataset schema. Events that came before +it in the same request may already have been expanded in memory by then, within the +limits, and are dropped with the rest: in a dataset with a custom or time partition, and +for Kinesis and OTEL, the events of a request are expanded one after the other. An +endpoint that creates the dataset on first use, such as `/api/v1/ingest`, does that before +it flattens the request, so a rejected request can still leave a new, empty dataset. | Environment variable | Flag | Default | Limit | | --------------------------------- | --------------------------------- | -------------------- | --------------------------------------------------------- | @@ -15,10 +23,16 @@ rejected request is ingested. Each value must be a positive integer. A limit cannot be set to zero or switched off. +These limits are new, and they reject some requests that earlier versions accepted, +including batches of small events without any arrays. Read +[Upgrading](#upgrading-requests-that-are-now-rejected) before you upgrade. + The byte limit counts the size of the expanded rows as compact JSON, including keys, quotes, escapes and delimiters. It also counts what is copied into every row afterwards: Kinesis request metadata and OTEL body key prefixes (see below). It is not a measure of the -memory a request uses. +memory a request uses. It is independent of `P_MAX_EVENT_PAYLOAD_SIZE` (10 MiB by default, +at most 100 MiB), which limits the size of the request body as it is received: a request +within one of them can be over the other. After generic flattening, every row is flattened structurally: nested objects become fields named after the keys on their path, joined with `_`. A parent's name is repeated in every @@ -26,26 +40,30 @@ field below it, so `{"<128 character name>": {"f0": 0, ..., "f31": 31}, "r": [0, expands into 4 rows of about 400 bytes, but each of them structurally flattened takes about 4.4 KiB. Every event is therefore charged the larger of the two: the compact JSON of its generic rows, or of those rows structurally flattened. Both are computed in one pass over -the event, before any row or field name is built. The structurally flattened size counts -every scalar as a field of its own, even when its joined name collides with another field -(for example `{"a_b": 1, "a": {"b": 2}}`, where one `a_b` overwrites the other), so with -such collisions it is an overestimate. Empty objects add no field. +the event, before any row or field name of it is built. The structurally flattened size +counts every scalar as a field of its own, even when its joined name collides with another +field (for example `{"a_b": 1, "a": {"b": 2}}`, where one `a_b` overwrites the other), so +with such collisions it is an overestimate. Empty objects add no field. ## What counts as an event -The limits apply only where generic flattening runs: +The limits apply only where generic flattening runs. There, every event takes at least one +row, including an event without arrays, which generic flattening leaves as it is: - **JSON and custom sources** (`/api/v1/ingest`, `/api/v1/logstream/{name}`) on datasets using schema version V1. Every object in the request body is one event, after any nested - top-level arrays are spliced in. Bodies nested deeper than `P_MAX_FLATTEN_LEVEL` skip - generic flattening, so these limits do not apply to them. + top-level arrays are spliced in, so a batch (a top-level array) of 10,001 objects takes at + least 10,001 rows. Bodies nested deeper than `P_MAX_FLATTEN_LEVEL` skip generic + flattening, so these limits do not apply to them. - **Kinesis**: every decoded `data` record is one event, even when it holds a JSON array. Each of its rows also gets the request's `requestId` and `timestamp`, so their compact JSON size, with keys and delimiters, is counted once per row, on top of the larger of the generic and structurally flattened sizes. It is counted in full even for rows that - already have these fields. + already have these fields. A record nested deeper than `P_MAX_FLATTEN_LEVEL` is ingested + as it is, and not counted. - **OTEL logs**: a log record whose body (or a body field) is a string holding a JSON object - has that object flattened. All flattened bodies of one log record together are one event. + has that object flattened. All flattened bodies of one log record together are one event, + of at least one row per flattened body. A log record without such a body takes no rows. This counts every intermediate row, even though the rows end up merged into the single log record. Each top-level field of a flattened body is renamed to `{body key}_{field}` before the merge, so the body key, with its escapes and the `_`, is counted again for @@ -61,15 +79,78 @@ The limits apply only where generic flattening runs: Every request has a budget of its own. All events of a request share that budget, across custom partition elements, Kinesis records, and OTEL resources and scopes. OTEL metrics and traces, Kafka, and V0 datasets do not use generic flattening, so these limits do not apply -to them. +to them. Internal dataset statistics are bounded differently, see +[below](#internal-dataset-statistics-pstats). + +## Upgrading: requests that are now rejected + +Earlier versions flattened every event without a limit. With the defaults, these requests +are now rejected: + +- **More than 10,000 events in one request.** Every event counts toward + `P_MAX_FLATTEN_ROWS_PER_REQUEST`, including an event without arrays. A JSON batch of + 10,001 `{"a":1}` events is about 80 KB, far below the 10 MiB default of + `P_MAX_EVENT_PAYLOAD_SIZE`, and it is rejected. So is a Kinesis request of more than + 10,000 records, an OTEL logs request of more than 10,000 log records with JSON object + bodies, and schema detection of more than 10,000 events. Shippers that batch by size or + by time, rather than by a number of events, can send such requests. +- **More than 1,000 values in an array of one event.** An event such as `{"items": [...]}` + with 1,001 items expands into 1,001 rows, over `P_MAX_FLATTEN_ROWS_PER_EVENT`. That + includes an event that wraps its records in an array, like `{"records": [...]}`, and a + Kinesis record that decodes to an array of more than 1,000 objects. The elements of a + top-level array in a request body are events of their own, so this limit does not apply + to a top-level array (the limit per request does). +- **Rows over 64 MiB in one request.** For a JSON batch of events without arrays or nested + objects, the rows take about the size of the batch without whitespace. With the 10 MiB + default of `P_MAX_EVENT_PAYLOAD_SIZE`, such batches stay well below + `P_MAX_FLATTEN_BYTES_PER_REQUEST`; they can reach it once `P_MAX_EVENT_PAYLOAD_SIZE` is + raised close to or above 64 MiB. Nested arrays, and parent names repeated in the + structurally flattened field names, make the rows larger than the request. + +Nothing is truncated: a request over a limit is rejected whole, with a `400` that names the +limit, and no part of it is stored. Many shippers treat a `400` as permanent and drop the +batch instead of retrying it, so compare these limits with your largest requests before you +upgrade, and look for these errors afterwards. + +To migrate, do one of these: + +- Configure each shipper to send at most 10,000 events per request (Kinesis records, or + OTEL log records with JSON object bodies), and to send each record as an event of its own + (an element of a top-level array, or a Kinesis record of its own) instead of wrapping + many of them in the array of one event. +- Raise the limits to fit your largest requests. `P_MAX_FLATTEN_ROWS_PER_REQUEST` must be at + least the number of events in a request, plus the rows their arrays add. + `P_MAX_FLATTEN_ROWS_PER_EVENT` must be at least the length of the longest array of one + event (the product of the lengths, for sibling arrays). The request size limit and the + byte limit are independent. Raising `P_MAX_EVENT_PAYLOAD_SIZE`, which can go up to + 100 MiB, lets larger requests in, but the rows they expand into are still held to + `P_MAX_FLATTEN_BYTES_PER_REQUEST` (64 MiB by default) and to the row limits, so raise + those too for larger flat batches. Higher limits let a request use more memory. + +## Internal dataset statistics (`pstats`) + +Parseable computes statistics for every parquet file it uploads (unless +`P_COLLECT_DATASET_STATS` is off) and ingests them into the internal `pstats` dataset, one +row for every distinct value of every field (up to `P_MAX_FIELD_STATISTICS` per field), and +one row for a field without distinct values. These rows are not generically flattened. +They are built straight from the statistics, with the fields and values that flattening +gave them before, one row at a time, and each row is ingested before the next is built. So +the row limits do not apply, and the statistics of a file are ingested in full however many +rows they take. -Internal dataset statistics (`pstats`) use the configured byte limit. Their row limits come -from the size of the stats payload, which is the field count times `P_MAX_FIELD_STATISTICS`. +The byte limit applies to each row on its own. A row is measured before it is built. A row +larger than `P_MAX_FLATTEN_BYTES_PER_REQUEST` as compact JSON, which takes a distinct value +or a field name of about that size, is skipped with a warning that names the stream, the +file and the field (cut to 128 characters), but no value. The other rows of the file are +still ingested. A failure to compute or ingest the statistics of a file is logged as a +warning, with the stream, the tenant and the file. ## When a request is rejected The error message names the limit that was exceeded and its value. To get under the limit, reduce the number or length of the nested arrays in each event, or split the batch into -smaller requests. You can also raise the limit. +smaller requests. You can also raise the limit, see +[Upgrading](#upgrading-requests-that-are-now-rejected). -`tools/flatten-bench/http_smoke.py` checks these responses against a built server binary. +`tools/flatten-bench/http_smoke.py` checks these responses against a built server binary, +including flat batches at and just over the rows per request. diff --git a/src/cli.rs b/src/cli.rs index 08a7c5a3b..847aae9a6 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -659,7 +659,7 @@ pub struct Options { env = "P_MAX_FLATTEN_ROWS_PER_REQUEST", default_value_t = DEFAULT_MAX_FLATTEN_ROWS_PER_REQUEST, value_parser = validation::validate_flatten_limit, - help = "Maximum number of rows all events of an ingest request may expand into when their nested arrays are flattened" + help = "Maximum number of rows all events of an ingest request may expand into when their nested arrays are flattened. Every event takes at least one row, including an event without arrays" )] pub max_flatten_rows_per_request: usize, diff --git a/src/handlers/http/kinesis.rs b/src/handlers/http/kinesis.rs index 38b845dbf..7625a41a3 100644 --- a/src/handlers/http/kinesis.rs +++ b/src/handlers/http/kinesis.rs @@ -53,6 +53,8 @@ struct Data { // The data field is base64 encoded JSON (there can be multiple data fields), and there is a requestId and timestamp field. // Each decoded data field is one event for `budget`, even when it holds an array, charged // for its generic rows or, when larger, those rows structurally flattened, plus the metadata. +// It is charged before its rows are built; when that fails, the rows built for the records +// before it are dropped with the error. // Kinesis logs are flattened to the following format: // { // "CHANGE": 3.16, diff --git a/src/handlers/http/modal/utils/ingest_utils.rs b/src/handlers/http/modal/utils/ingest_utils.rs index d5acb7ec4..f3dab0ef0 100644 --- a/src/handlers/http/modal/utils/ingest_utils.rs +++ b/src/handlers/http/modal/utils/ingest_utils.rs @@ -446,7 +446,9 @@ pub fn push_logs( /// Zero-allocation JSON byte size counter. /// Uses serde_json::to_writer with a writer that only counts bytes. -pub fn json_byte_size(value: &Value) -> u64 { +/// Exact for values that always serialize, such as a `Value` or a struct of strings and +/// numbers, which is what it is used for. +pub fn json_byte_size(value: &(impl serde::Serialize + ?Sized)) -> u64 { struct CountingWriter(u64); impl std::io::Write for CountingWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { @@ -458,7 +460,7 @@ pub fn json_byte_size(value: &Value) -> u64 { } } let mut counter = CountingWriter(0); - // serde_json::to_writer won't fail on Value types + // serde_json::to_writer won't fail on those let _ = serde_json::to_writer(&mut counter, value); counter.0 } diff --git a/src/otel/logs.rs b/src/otel/logs.rs index 2d7af8379..14acac4c6 100644 --- a/src/otel/logs.rs +++ b/src/otel/logs.rs @@ -413,8 +413,9 @@ pub fn flatten_otel_protobuf( /// this function performs the custom flattening of the otel logs /// and returns a `Vec` of `Value::Object` of the flattened json -/// Fails, with nothing returned, once a log body is over the limits of `budget`, which is -/// shared by all records of the request +/// Every record is charged to `budget`, which all records of the request share, before the +/// rows of its bodies are built. Fails once a record would go over its limits, dropping the +/// records converted before it pub fn flatten_otel_logs( message: &LogsData, tenant_id: &str, @@ -663,26 +664,36 @@ mod tests { )); } + // Logs collected for `tenant` over all dates, so that a test running across midnight (UTC) + // still counts what it collected. Each test uses a tenant of its own, as the counter is + // global. + fn logs_collected(tenant: &str) -> u64 { + use prometheus::core::Collector; + + crate::metrics::TOTAL_LOGS_COLLECTED_BY_DATE + .collect() + .iter() + .flat_map(|family| family.get_metric()) + .filter(|metric| { + metric + .get_label() + .iter() + .any(|label| label.get_name() == "tenant_id" && label.get_value() == tenant) + }) + .map(|metric| metric.get_counter().get_value() as u64) + .sum() + } + #[test] fn logs_collected_metric_counts_only_converted_requests() { - use crate::metrics::TOTAL_LOGS_COLLECTED_BY_DATE; - - // A tenant of its own, as the counter is global let tenant = "flatten-limit-metric-tenant"; - let collected = || { - let date = chrono::Utc::now().date_naive().to_string(); - TOTAL_LOGS_COLLECTED_BY_DATE - .with_label_values(&["all", &date, tenant]) - .get() - }; - - let before = collected(); + let before = logs_collected(tenant); // Fails on the third record, after the scopes of the first resource were converted assert!(flatten_otel_logs(&logs(SIX_ROW_BODY), tenant, &mut limits(6, 17)).is_err()); - assert_eq!(collected(), before); + assert_eq!(logs_collected(tenant), before); flatten_otel_logs(&logs(SIX_ROW_BODY), tenant, &mut limits(6, 18)).unwrap(); - assert_eq!(collected(), before + 3); + assert_eq!(logs_collected(tenant), before + 3); } fn byte_limit(bytes: usize) -> FlattenBudget { @@ -808,15 +819,7 @@ mod tests { #[test] fn logs_collected_metric_is_unchanged_when_a_body_key_is_over_the_budget() { - use crate::metrics::TOTAL_LOGS_COLLECTED_BY_DATE; - let tenant = "flatten-limit-prefix-metric-tenant"; - let collected = || { - let date = chrono::Utc::now().date_naive().to_string(); - TOTAL_LOGS_COLLECTED_BY_DATE - .with_label_values(&["all", &date, tenant]) - .get() - }; let body = wide_body(); let long_key = "k".repeat(128); let scope = |key: &str| ScopeLogs { @@ -837,12 +840,19 @@ mod tests { ], }; - let before = collected(); + let before = logs_collected(tenant); assert!(matches!( flatten_otel_logs(&logs, tenant, &mut byte_limit(4096)), Err(JsonFlattenError::BytesPerRequestLimitExceeded(4096)) )); - assert_eq!(collected(), before); + assert_eq!(logs_collected(tenant), before); + + // Without the second resource, the request is converted and counted + let fitting = LogsData { + resource_logs: logs.resource_logs[..1].to_vec(), + }; + flatten_otel_logs(&fitting, tenant, &mut byte_limit(4096)).unwrap(); + assert_eq!(logs_collected(tenant), before + 1); } // Compact JSON bytes of the rows of `body` merged into the record under `key`, once diff --git a/src/storage/field_stats.rs b/src/storage/field_stats.rs index 141453856..fdf804546 100644 --- a/src/storage/field_stats.rs +++ b/src/storage/field_stats.rs @@ -26,6 +26,7 @@ use crate::handlers::http::cluster::send_query_request; use crate::handlers::http::ingest::PostError; use crate::handlers::http::middleware::CLUSTER_SECRET; use crate::handlers::http::middleware::CLUSTER_SECRET_HEADER; +use crate::handlers::http::modal::utils::ingest_utils::json_byte_size; use crate::handlers::http::query::Query; use crate::handlers::http::query::QueryError; use crate::handlers::http::query::query; @@ -39,9 +40,8 @@ use crate::tenants::TENANT_METADATA; use crate::utils::create_intracluster_auth_headermap; use crate::utils::get_tenant_id_from_request; use crate::utils::get_user_from_request; -use crate::utils::json::apply_generic_flattening_for_partition; -use crate::utils::json::flatten::FlattenBudget; use crate::utils::json::flatten::FlattenLimits; +use crate::utils::json::flatten::validate_custom_partition; use actix_web::HttpRequest; use actix_web::HttpResponse; use actix_web::Responder; @@ -118,6 +118,61 @@ struct DatasetStats { field_stats: Vec, } +impl DatasetStats { + /// The rows of the stats dataset for these stats, borrowed from them: what generic and + /// then structural flattening (with `_`) make of the stats as JSON. A row for every + /// distinct value of every field, in order, and a row for a field without any. + fn rows(&self) -> impl Iterator> { + self.field_stats.iter().flat_map(move |field| { + // Generic flattening makes an empty `distinct_stats` array a null + let null = DistinctColumns::Null { + field_stats_distinct_stats: (), + }; + let no_distinct_values = field.distinct_stats.is_empty().then_some(null); + field + .distinct_stats + .iter() + .map(|stat| DistinctColumns::Value { + field_stats_distinct_stats_distinct_value: &stat.distinct_value, + field_stats_distinct_stats_count: stat.count, + }) + .chain(no_distinct_values) + .map(move |distinct| StatsRow { + dataset_name: &self.dataset_name, + stats_id: &self.stats_id, + field_stats_field_name: &field.field_name, + field_stats_count: field.count, + field_stats_distinct_count: field.distinct_count, + distinct, + }) + }) + } +} + +/// A row of the stats dataset, see `DatasetStats::rows`. The fields are in the order of the +/// flattened row, which the JSON of the row keeps. +#[derive(Serialize, Debug)] +struct StatsRow<'a> { + dataset_name: &'a str, + stats_id: &'a str, + field_stats_field_name: &'a str, + field_stats_count: i64, + field_stats_distinct_count: i64, + #[serde(flatten)] + distinct: DistinctColumns<'a>, +} + +#[derive(Serialize, Debug)] +#[serde(untagged)] +enum DistinctColumns<'a> { + Value { + field_stats_distinct_stats_distinct_value: &'a str, + field_stats_distinct_stats_count: i64, + }, + /// `"field_stats_distinct_stats": null`, for a field without distinct values + Null { field_stats_distinct_stats: () }, +} + /// Calculates field statistics for the stream and pushes them to the internal stats dataset. /// This function creates a new internal stream for stats if it doesn't exist. /// It collects statistics for each field in the stream @@ -173,28 +228,14 @@ async fn calculate_field_stats_inner( return Ok(stats_calculated); } stats_calculated = true; - let stats_value = - serde_json::to_value(&stats).map_err(|e| ObjectStorageError::Invalid(e.into()))?; ensure_dataset_stats_stream(tenant_id).await?; - let limits = field_stats_flatten_limits( - stats.field_stats.len(), - max_field_statistics, - FlattenLimits::from_config().max_bytes_per_request, - ); - let vec_json = apply_generic_flattening_for_partition( - stats_value, - None, - None, - Some(&DATASET_STATS_CUSTOM_PARTITION.to_string()), - &mut FlattenBudget::new(limits), - )?; let mut p_custom_fields = HashMap::new(); p_custom_fields.insert(USER_AGENT_KEY.to_string(), "parseable".to_string()); let schema = PARSEABLE .get_stream(DATASET_STATS_STREAM_NAME, tenant_id)? .get_schema_raw(); - for json in vec_json { - let origin_size = crate::handlers::http::modal::utils::ingest_utils::json_byte_size(&json); + let max_row_bytes = FlattenLimits::from_config().max_bytes_per_request; + for_each_stats_row(&stats, max_row_bytes, |json, origin_size| { json::Event { json, p_timestamp: parquet_ts, @@ -214,26 +255,52 @@ async fn calculate_field_stats_inner( false, )? .process()?; - } + Ok(()) + })?; Ok(stats_calculated) } -/// Generic flattening limits for one stats payload. It expands into a row per distinct value -/// of each field, and there are at most `max_field_statistics` of those per field, so the -/// row limits follow from its shape instead of the limits for ingested events, which a wide -/// dataset would exceed -fn field_stats_flatten_limits( - field_count: usize, - max_field_statistics: usize, - max_bytes: usize, -) -> FlattenLimits { - let max_rows = field_count - .max(1) - .saturating_mul(max_field_statistics.max(1)); - FlattenLimits { - max_rows_per_event: max_rows, - max_rows_per_request: max_rows, - max_bytes_per_request: max_bytes, +/// Hands the rows of the stats dataset for `stats` (see `DatasetStats::rows`) to `process` +/// one at a time, each with its size as compact JSON. A row is measured before it is built, +/// and only built when it is at most `max_row_bytes`, so one row of at most that size is all +/// that is held at a time, however many fields and distinct values the stats have. A row over +/// it, which takes a distinct value or a field name of about that size, is skipped with a +/// warning, and the rows after it are still processed. Returns the number of rows skipped. +fn for_each_stats_row( + stats: &DatasetStats, + max_row_bytes: usize, + mut process: impl FnMut(serde_json::Value, u64) -> Result<(), PostError>, +) -> Result { + let custom_partition = DATASET_STATS_CUSTOM_PARTITION.to_string(); + let mut skipped = 0; + for row in stats.rows() { + let size = json_byte_size(&row); + if size > max_row_bytes as u64 { + warn!( + "Skipping a field stats row of field {} for parquet file {} of stream {}: its {size} bytes of JSON are over P_MAX_FLATTEN_BYTES_PER_REQUEST ({max_row_bytes})", + shortened(row.field_stats_field_name), + stats.stats_id, + stats.dataset_name + ); + skipped += 1; + continue; + } + let json = serde_json::to_value(&row).map_err(|e| ObjectStorageError::Invalid(e.into()))?; + // As flattening validated the custom partition of every row + validate_custom_partition( + json.as_object().expect("a stats row is a JSON object"), + Some(&custom_partition), + )?; + process(json, size)?; + } + Ok(skipped) +} + +/// The first 128 characters of `name`, which can be of any length, for a log message +fn shortened(name: &str) -> Cow<'_, str> { + match name.char_indices().nth(128) { + Some((end, _)) => Cow::Owned(format!("{}...", &name[..end])), + None => Cow::Borrowed(name), } } @@ -253,8 +320,9 @@ fn log_stats_calculation_time( "Field stats calculation completed for parquet file {parquet_file} in {elapsed_ms} ms. stream={stream_name}" ); } else { + // The caller warns about the error, with the stream, tenant and path let err = result.as_ref().expect_err("result checked as error"); - warn!( + debug!( "Field stats calculation failed for parquet file {parquet_file} after {elapsed_ms} ms. stream={stream_name}: {err}" ); } @@ -1135,15 +1203,17 @@ mod tests { }; use arrow_schema::{DataType, Field, Schema, TimeUnit}; use parquet::{arrow::ArrowWriter, file::properties::WriterProperties}; + use serde_json::json; use temp_dir::TempDir; + use crate::handlers::http::modal::utils::ingest_utils::json_byte_size; use crate::storage::field_stats::{ DATASET_STATS_CUSTOM_PARTITION, DatasetStats, DistinctStat, FieldCountState, FieldStat, NULL_VALUE, build_stats_sql, collect_all_field_stats_from_parquet_blocking, - collect_field_counts, field_stats_flatten_limits, + collect_field_counts, for_each_stats_row, }; use crate::utils::json::apply_generic_flattening_for_partition; - use crate::utils::json::flatten::FlattenBudget; + use crate::utils::json::flatten::{FlattenBudget, FlattenLimits, GENEROUS_LIMITS, legacy}; async fn create_test_parquet_with_data() -> (TempDir, std::path::PathBuf) { let temp_dir = TempDir::new().unwrap(); @@ -1525,83 +1595,244 @@ mod tests { ); } - fn dataset_stats(distinct_per_field: &[usize]) -> serde_json::Value { - let field_stats = distinct_per_field - .iter() - .enumerate() - .map(|(field, &distinct)| FieldStat { - field_name: format!("field{field}"), - count: 10, - distinct_count: distinct as i64, - distinct_stats: (0..distinct) - .map(|value| DistinctStat { - distinct_value: format!("value{value}"), - count: 1, - }) - .collect(), - }) - .collect(); - serde_json::to_value(DatasetStats { + fn field_stat(name: &str, distinct_values: &[&str]) -> FieldStat { + FieldStat { + field_name: name.to_owned(), + count: 10, + distinct_count: distinct_values.len() as i64, + distinct_stats: distinct_values + .iter() + .enumerate() + .map(|(i, value)| DistinctStat { + distinct_value: (*value).to_owned(), + count: (distinct_values.len() - i) as i64, + }) + .collect(), + } + } + + fn dataset_stats(field_stats: Vec) -> DatasetStats { + DatasetStats { dataset_name: "dataset".to_owned(), stats_id: "stats".to_owned(), field_stats, + } + } + + // `field{i}` with the distinct values `value0`, `value1`, ... up to `distinct_per_field[i]` + fn uniform_stats(distinct_per_field: &[usize]) -> DatasetStats { + dataset_stats( + distinct_per_field + .iter() + .enumerate() + .map(|(field, &distinct)| FieldStat { + field_name: format!("field{field}"), + count: 10, + distinct_count: distinct as i64, + distinct_stats: (0..distinct) + .map(|value| DistinctStat { + distinct_value: format!("value{value}"), + count: 1, + }) + .collect(), + }) + .collect(), + ) + } + + // The rows `for_each_stats_row` hands over, with their sizes, and how many it skipped + fn streamed_rows( + stats: &DatasetStats, + max_row_bytes: usize, + ) -> (Vec<(serde_json::Value, u64)>, usize) { + let mut rows = Vec::new(); + let skipped = for_each_stats_row(stats, max_row_bytes, |row, size| { + rows.push((row, size)); + Ok(()) }) - .unwrap() + .unwrap(); + (rows, skipped) } - fn flatten_stats( - stats: serde_json::Value, - field_count: usize, - max_field_statistics: usize, + // The rows the stats used to be ingested as: the stats as JSON, generically and then + // structurally flattened by the oracle copies of those functions + fn legacy_rows(stats: &DatasetStats) -> Vec { + let partition = DATASET_STATS_CUSTOM_PARTITION.to_string(); + legacy::generic_flattening(&serde_json::to_value(stats).unwrap()) + .unwrap() + .into_iter() + .map(|mut row| { + legacy::flatten(&mut row, "_", Some(&partition), true).unwrap(); + row + }) + .collect() + } + + fn flatten_stats_payload( + stats: &DatasetStats, + limits: FlattenLimits, ) -> anyhow::Result> { apply_generic_flattening_for_partition( - stats, + serde_json::to_value(stats).unwrap(), None, None, Some(&DATASET_STATS_CUSTOM_PARTITION.to_string()), - &mut FlattenBudget::new(field_stats_flatten_limits( - field_count, - max_field_statistics, - 1 << 30, - )), + &mut FlattenBudget::new(limits), ) } #[test] - fn widest_stats_payload_fits_its_flatten_limits() { - // Far more rows than the limits for ingested events allow - let fields = 300; - let max_field_statistics = 50; - let rows = flatten_stats( - dataset_stats(&[max_field_statistics; 300]), - fields, - max_field_statistics, - ) - .unwrap(); - assert_eq!(rows.len(), fields * max_field_statistics); + fn stats_rows_are_the_flattened_stats() { + let cases = [ + dataset_stats(vec![field_stat("a", &["x", "y"])]), + // A field without distinct values gets a row with a null `distinct_stats` + dataset_stats(vec![ + field_stat("empty", &[]), + field_stat("b", &["1"]), + field_stat("also_empty", &[]), + ]), + dataset_stats(vec![field_stat( + "esc\"aped\n", + &["quo\"te", "back\\slash", "\u{1}ctl", "ünï", ""], + )]), + uniform_stats(&[3, 0, 2, 1]), + ]; + for stats in cases { + let (rows, skipped) = + streamed_rows(&stats, FlattenLimits::DEFAULT.max_bytes_per_request); + assert_eq!(skipped, 0); + for (row, size) in &rows { + assert_eq!(*size, json_byte_size(row)); + } + let rows: Vec<_> = rows.into_iter().map(|(row, _)| row).collect(); + // Serialized, so that the order of the keys is compared too + let expected = serde_json::to_string(&legacy_rows(&stats)).unwrap(); + assert_eq!(serde_json::to_string(&rows).unwrap(), expected); + // The same as the stats flattened as one payload, which they were before + let flattened = flatten_stats_payload(&stats, GENEROUS_LIMITS).unwrap(); + assert_eq!(serde_json::to_string(&flattened).unwrap(), expected); + } + + let (rows, _) = streamed_rows( + &dataset_stats(vec![field_stat("empty", &[]), field_stat("b", &["v"])]), + FlattenLimits::DEFAULT.max_bytes_per_request, + ); + let rows: Vec<_> = rows + .iter() + .map(|(row, _)| serde_json::to_string(row).unwrap()) + .collect(); assert_eq!( - rows[0]["field_stats_distinct_stats_distinct_value"], - "value0" + rows, + [ + r#"{"dataset_name":"dataset","stats_id":"stats","field_stats_field_name":"empty","field_stats_count":10,"field_stats_distinct_count":0,"field_stats_distinct_stats":null}"#, + r#"{"dataset_name":"dataset","stats_id":"stats","field_stats_field_name":"b","field_stats_count":10,"field_stats_distinct_count":1,"field_stats_distinct_stats_distinct_value":"v","field_stats_distinct_stats_count":1}"#, + ] ); + } - // A field without distinct values still takes a row - let rows = flatten_stats(dataset_stats(&[0, 2]), 2, 2).unwrap(); - assert_eq!(rows.len(), 3); - let rows = flatten_stats(dataset_stats(&[0]), 1, 0).unwrap(); - assert_eq!(rows.len(), 1); + #[test] + fn stats_rows_keep_the_custom_partition_validation() { + let mut stats = dataset_stats(vec![field_stat("a", &["x"])]); + stats.dataset_name = "with.period".to_owned(); + let result = for_each_stats_row( + &stats, + FlattenLimits::DEFAULT.max_bytes_per_request, + |_, _| unreachable!("a row with an invalid custom partition is not processed"), + ); + assert!(result.is_err()); + assert!(flatten_stats_payload(&stats, GENEROUS_LIMITS).is_err()); } #[test] - fn stats_payload_flatten_limits_are_bounded() { - assert_eq!(field_stats_flatten_limits(3, 4, 7).max_rows_per_event, 12); - assert_eq!(field_stats_flatten_limits(3, 4, 7).max_rows_per_request, 12); - assert_eq!(field_stats_flatten_limits(3, 4, 7).max_bytes_per_request, 7); - assert_eq!(field_stats_flatten_limits(0, 0, 7).max_rows_per_event, 1); + fn stats_rows_over_one_byte_limit_are_all_processed() { + // 1000 rows of about 225 bytes each, 225 KB together, against a limit of 512 bytes + let stats = uniform_stats(&[25; 40]); + let max_row_bytes = 512; + let (rows, skipped) = streamed_rows(&stats, max_row_bytes); + assert_eq!((rows.len(), skipped), (1000, 0)); + assert!(rows.iter().all(|(_, size)| *size <= max_row_bytes as u64)); + let total: u64 = rows.iter().map(|(_, size)| size).sum(); + assert!(total > 400 * max_row_bytes as u64, "{total}"); assert_eq!( - field_stats_flatten_limits(usize::MAX, 2, 7).max_rows_per_event, - usize::MAX + serde_json::to_string(&rows.into_iter().map(|(row, _)| row).collect::>()) + .unwrap(), + serde_json::to_string(&legacy_rows(&stats)).unwrap() + ); + // Flattened as one payload under the same limit, nothing was ingested + let limits = FlattenLimits { + max_rows_per_event: 1000, + max_rows_per_request: 1000, + max_bytes_per_request: max_row_bytes, + }; + assert!(flatten_stats_payload(&stats, limits).is_err()); + } + + #[test] + fn stats_rows_skip_only_a_row_over_the_byte_limit() { + let long_value = "v".repeat(600); + let long_name = "f".repeat(600); + let stats = dataset_stats(vec![ + field_stat("a", &["x", long_value.as_str(), "y"]), + field_stat(&long_name, &["z"]), + field_stat("b", &[]), + ]); + let (rows, skipped) = streamed_rows(&stats, 512); + assert_eq!(skipped, 2); + let kept: Vec<_> = rows + .iter() + .map(|(row, _)| { + ( + row["field_stats_field_name"].clone(), + row.get("field_stats_distinct_stats_distinct_value") + .cloned(), + row["field_stats_distinct_count"].clone(), + ) + }) + .collect(); + assert_eq!( + kept, + [ + (json!("a"), Some(json!("x")), json!(3)), + (json!("a"), Some(json!("y")), json!(3)), + (json!("b"), None, json!(0)), + ] + ); + // The rows kept are those the stats were flattened into + let legacy: Vec<_> = legacy_rows(&stats) + .into_iter() + .filter(|row| json_byte_size(row) <= 512) + .collect(); + assert_eq!( + serde_json::to_string(&rows.into_iter().map(|(row, _)| row).collect::>()) + .unwrap(), + serde_json::to_string(&legacy).unwrap() + ); + } + + #[test] + fn stats_rows_of_the_widest_payload_fit_the_default_byte_limit() { + // Far more rows than the limits for ingested events allow, each processed on its own + let fields = 300; + let max_field_statistics = 50; + let stats = uniform_stats(&[max_field_statistics; 300]); + let mut rows = 0; + let mut first = None; + let skipped = for_each_stats_row( + &stats, + FlattenLimits::DEFAULT.max_bytes_per_request, + |row, _| { + if first.is_none() { + first = Some(row); + } + rows += 1; + Ok(()) + }, + ) + .unwrap(); + assert_eq!((rows, skipped), (fields * max_field_statistics, 0)); + assert_eq!( + first.unwrap()["field_stats_distinct_stats_distinct_value"], + "value0" ); - // More distinct values than the payload can hold are rejected - assert!(flatten_stats(dataset_stats(&[3, 2]), 2, 2).is_err()); } } diff --git a/src/storage/object_storage.rs b/src/storage/object_storage.rs index d41ea7d2f..51984dc01 100644 --- a/src/storage/object_storage.rs +++ b/src/storage/object_storage.rs @@ -265,10 +265,10 @@ async fn calculate_stats_if_enabled( if let Err(err) = calculate_field_stats(stream_name, path, schema, max_field_statistics, &tenant_id).await { - tracing::trace!( - "Error calculating field stats for stream {}: {}", - stream_name, - err + warn!( + "Error calculating field stats for stream {stream_name} (tenant {}) from parquet file {}: {err}", + tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), + path.display() ); } } diff --git a/tools/flatten-bench/README.md b/tools/flatten-bench/README.md index ec093ba7e..74642dc4d 100644 --- a/tools/flatten-bench/README.md +++ b/tools/flatten-bench/README.md @@ -7,9 +7,11 @@ production sources unmodified, so it measures the real code without building the (own `[workspace]`, depends only on serde/serde_json/chrono/thiserror/anyhow/tracing): - `current_json`: the working tree's `src/utils/json/{mod,flatten,strict}.rs`, included by path. -- `baseline_json`: the same files at `547c5118` (the merge-base, including its no-op fast - paths), extracted byte for byte with `git show` by `build.rs` into `OUT_DIR`. Override - with `FLATTEN_BENCH_BASELINE_REV=`. This is NOT the `#[cfg(test)] legacy` module in +- `baseline_json`: the same files at `547c5118`, the fixed comparison baseline (including + its no-op fast paths), extracted byte for byte with `git show` by `build.rs` into + `OUT_DIR`. The branch was first based on that revision and later rebased onto `04b73b57`, + where these three files are the same blobs. Override with + `FLATTEN_BENCH_BASELINE_REV=`. This is NOT the `#[cfg(test)] legacy` module in flatten.rs, which predates those fast paths; `[lib] test = false` keeps production `#[cfg(test)]` code out entirely. The bench prints the revision and flatten.rs blob id. @@ -59,19 +61,30 @@ thread only, in requested bytes (not RSS): `allocs` alloc/realloc calls, `alloc_ cumulative bytes requested, `peak_delta` highest live bytes above the level at call start, `net_delta` live bytes at the end minus the start. `peak_delta` is a delta, not an absolute peak: for consuming calls, freeing the input during the call lowers it. The borrowed -`generic_flattening` section shows the extra memory without that offset. +`generic_flattening` section shows the extra memory without that offset. These counted +calls use `measure`, which has no cap. + +`tests/alloc_limits.rs` measures rejections with `measure_bounded` instead, which also caps +the bytes requested at 64 KiB, the allowance those tests assert. An allocation that would go +over the cap is never handed to the system allocator: the test process prints a line +starting with `alloc_counter:` to stderr and exits with status 88 +(`CAP_EXCEEDED_EXIT_CODE`), so a regression that started expanding a hostile input stops +there instead of using more memory. `allocation_over_the_cap_exits` checks this in a child +process. ## HTTP smoke test `http_smoke.py` (Python 3 standard library only) starts a built server binary with -`local-store` on an ephemeral 127.0.0.1 port, in a temporary directory, with small flatten -caps (rows per event and bytes from the environment, rows per request from the CLI flag, -overriding a different environment value). It checks accepted requests at the limits and -`400` rejections over them, with unchanged event counts and schemas, for JSON (on +`local-store` in a temporary directory, on ephemeral 127.0.0.1 ports for HTTP and for the +gRPC, Arrow Flight and query gRPC servers (which otherwise take 8001 to 8003), with small +flatten caps (rows per event and bytes from the environment, rows per request from the CLI +flag, overriding a different environment value). It checks accepted requests at the limits +and `400` rejections over them, with unchanged event counts and schemas, for JSON (on `/api/v1/ingest` and `/api/v1/logstream/{name}`), a known log format, custom partitions, Kinesis, OTEL JSON logs and schema detection, including events whose long or quote-escaped parent names only exceed the byte limit once structural flattening repeats them in every -field name. +field name, and batches of events without arrays (a row each) at and just over the rows per +request. ``` P_BUILD_NO_UI=1 cargo build --bin parseable diff --git a/tools/flatten-bench/RESULTS.md b/tools/flatten-bench/RESULTS.md index c8285a1fc..4fc4102ef 100644 --- a/tools/flatten-bench/RESULTS.md +++ b/tools/flatten-bench/RESULTS.md @@ -73,18 +73,20 @@ detect the no-op. | escaped_repeated_string (500 rows, over 64 MiB after escaping) | 62.8 / 58.7 / 62.5 µs | | just_over_event_limit (1001 rows) | 9.6 / 10.0 / 7.5 µs | -- **Memory:** every rejection at the entry point allocates exactly once, 112 bytes (the - `anyhow` error), with `peak_delta=0`. In `/tmp/flatten-row-cap-harness-tests-r5.log` the - core paths (preflight, per-event `charge`, `pipeline charge`, and bounded - `generic_flattening`) report 0 allocations, 0 bytes and `peak_delta=0` for every hostile - input, and `core charge, remaining rows/bytes insufficient` reports the same for a request - budget whose remainder is too small. For an event whose generic rows fit but whose - structurally flattened field names do not (a 128-character parent copied into 32 field - names over 4 rows), `structural preflight, long parent` and `structural charge, long - parent` report 0 allocations, and the entry point rejects it with 1 allocation of 112 - bytes and `peak_delta=0` (2 allocations, 184 bytes, `peak_delta=72` with a partition - field). `tests/alloc_limits.rs` asserts these. They are measured by the counting - allocator, not as full-process RSS. +- **Memory:** the entry point rejects each of these single-event inputs with exactly one + allocation, 112 bytes (the `anyhow` error), with `peak_delta=0`. A request of several + events may have built the rows of earlier events before a later one is rejected (see the + partitioned case in `tests/alloc_limits.rs`), so this is not a bound for every request. + In `/tmp/flatten-row-cap-harness-tests-r5.log` the core paths (preflight, per-event + `charge`, `pipeline charge`, and bounded `generic_flattening`) report 0 allocations, 0 + bytes and `peak_delta=0` for every hostile input, and `core charge, remaining rows/bytes + insufficient` reports the same for a request budget whose remainder is too small. For an + event whose generic rows fit but whose structurally flattened field names do not (a + 128-character parent copied into 32 field names over 4 rows), `structural preflight, + long parent` and `structural charge, long parent` report 0 allocations, and the entry + point rejects it with 1 allocation of 112 bytes and `peak_delta=0` (2 allocations, 184 + bytes, `peak_delta=72` with a partition field). `tests/alloc_limits.rs` asserts these. + They are measured by the counting allocator, not as full-process RSS. - **Time:** these times include freeing the consumed input. They grow with the input's size and with how much of it is walked before the limit trips, not with the rows the input would expand to. They are not constant across inputs. @@ -93,7 +95,12 @@ detect the no-op. ## Verification -Results for the sources identified by the hashes above: +These are the results of the first validation of the sources identified by the hashes +above, before the branch was rebased onto `04b73b57`. The rebase did not change those two +files (the rebased commit has the same sha256 for both). The later validation of the +rebased commit, with its own test selection and counts, is reported in the pull request +description, not here, and this file records no validation of later changes either. The +counts below belong to the first validation only: | check | result | |---|---| @@ -120,18 +127,25 @@ cargo build --bin parseable --offline python3 tools/flatten-bench/http_smoke.py target/debug/parseable ``` +These are the commands of the first validation. The two `storage::field_stats` tests named in +them have since been replaced by the `storage::field_stats::tests::stats_rows` tests, as the +stats rows are no longer flattened as one payload, and the harness and HTTP smoke test have +more cases than the counts above. + ## Interpretation - **Remaining cost:** flat and nested inputs without arrays are 8–45% slower (about 0.13 µs per flat event in a batch, 0.17–0.20 µs for a single event). This is the price of the added safety: the event has to be walked to count its generic and structural rows and - bytes before anything is built. The fast paths still return the input + bytes before any of its rows is built. The fast paths still return the input unchanged without copying it, and allocation counts are unchanged. - **Inputs with arrays:** these break even or improve slightly, with fewer and smaller allocations. Only moderate_array_100 is measurably faster; product_10x10x10 and partitioned_batch_50x4 are within noise. -- **Limits:** none of the optimizations relax them. A limit that would be exceeded is - rejected before any row is built. +- **Limits:** none of the optimizations relax them. An event that would exceed a limit is + rejected before any of its rows is built. Earlier events of the same request may have + been built by then (one at a time with a custom or time partition, for Kinesis records + and OTEL log records), within the request's budget; the rejected request stores none of them. - **Scope:** these are microbenchmarks of the flattening code only, not end-to-end ingestion throughput. Parsing, Arrow conversion, and I/O are not included. diff --git a/tools/flatten-bench/build.rs b/tools/flatten-bench/build.rs index d2011937a..8a4b0cc06 100644 --- a/tools/flatten-bench/build.rs +++ b/tools/flatten-bench/build.rs @@ -3,8 +3,10 @@ //! them. Nothing in the extracted files is edited; the crate root provides the same //! `crate::parseable`, `crate::event` and `crate::metadata` stubs to both copies. //! -//! Override the revision with FLATTEN_BENCH_BASELINE_REV (default 547c5118, the -//! merge-base of fix/flatten-row-cap with main). +//! Override the revision with FLATTEN_BENCH_BASELINE_REV (default 547c5118, the fixed +//! comparison baseline: the revision fix/flatten-row-cap was first based on. Its +//! `src/utils/json` files are the same blobs as at 04b73b57, which the branch was rebased +//! onto). use std::path::{Path, PathBuf}; use std::process::Command; diff --git a/tools/flatten-bench/http_smoke.py b/tools/flatten-bench/http_smoke.py index 85c093206..0984a01c9 100644 --- a/tools/flatten-bench/http_smoke.py +++ b/tools/flatten-bench/http_smoke.py @@ -1,14 +1,16 @@ #!/usr/bin/env python3 """HTTP smoke test of the generic flattening limits against a real Parseable binary. -Starts ` local-store` on 127.0.0.1 (ephemeral port, temporary data and staging -directories, small flatten caps), then for plain JSON (on `/api/v1/ingest` and -`/api/v1/logstream/{name}`), a known log format, a custom partitioned stream, Kinesis, OTEL -JSON logs and schema detection checks that: +Starts ` local-store` on 127.0.0.1 (ephemeral ports for HTTP, gRPC, Arrow Flight +and query gRPC, temporary data and staging directories, small flatten caps), then for plain +JSON (on `/api/v1/ingest` and `/api/v1/logstream/{name}`), a known log format, a custom +partitioned stream, Kinesis, OTEL JSON logs and schema detection checks that: - a request at the limits is accepted (200) and its rows are counted in the stream stats - a request over a limit, including one that only fails on a late event, gets a 400 naming the limit, the stream's event count does not change, and a sentinel field of an earlier, valid event of that request does not show up in the stream schema. +Events without arrays are a row each, so a batch of more flat events than the rows per +request is rejected too (OTEL counts the records with a JSON object body). Over-limit bytes include those of the structurally flattened rows, whose field names repeat the names of their parents (`{"": {"f0": 0, ...}}` becomes `{"_f0": 0, ...}`), even when the generic rows are small. @@ -83,20 +85,31 @@ def events(self, stream): return json.loads(text)["ingestion"]["count"] -def free_port(): - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] +def free_ports(count): + """`count` distinct ports that were free on 127.0.0.1 a moment ago""" + sockets = [socket.socket() for _ in range(count)] + try: + for s in sockets: + s.bind(("127.0.0.1", 0)) + return [s.getsockname()[1] for s in sockets] + finally: + for s in sockets: + s.close() def start(binary, workdir, startup_timeout): - port = free_port() + # The gRPC, Arrow Flight and query gRPC servers default to fixed ports (8001-8003), + # which another local server may hold + port, grpc_port, flight_port, query_grpc_port = free_ports(4) for name in ("data", "staging"): os.makedirs(os.path.join(workdir, name)) # Nothing from the caller's P_* settings, so the caps below are the ones in effect env = {k: v for k, v in os.environ.items() if not k.startswith("P_")} env.update( P_ADDR=f"127.0.0.1:{port}", + P_GRPC_PORT=str(grpc_port), + P_FLIGHT_PORT=str(flight_port), + P_QUERY_GRPC_PORT=str(query_grpc_port), P_USERNAME="admin", P_PASSWORD="admin", P_FS_DIR=os.path.join(workdir, "data"), @@ -199,6 +212,14 @@ def compact_len(value): return len(json.dumps(value, separators=(",", ":"))) +def flat(n, sentinel=None): + """`n` events without arrays, a row each. The first one holds `sentinel` when given""" + events = [{"f": i} for i in range(n)] + if sentinel is not None: + events[0][sentinel] = 1 + return events + + # Three 1-row events of ~1.5 KiB each: two fit in BYTES_PER_REQUEST, three do not BIG = "x" * 1500 assert 2 * compact_len({"s": BIG}) < BYTES_PER_REQUEST < 3 * compact_len({"s": BIG}) @@ -347,6 +368,17 @@ def run_json(c): control, sentinel = unique("control"), unique("sentinel") c.accepted("json event at row limit", s, 4, *post, {"a": arr(4), control: 1}, h, control) c.accepted("json request at row limit", s, 6, *post, [{"a": arr(2)}, {"b": arr(4)}], h) + flat_sentinel = unique("sentinel") + c.accepted("json flat batch at request row limit", s, 6, *post, flat(6), h) + c.rejected( + "json flat batch over request row limit", + s, + REQUEST_MSG, + *post, + flat(7, flat_sentinel), + h, + flat_sentinel, + ) c.rejected("json event over row limit", s, EVENT_MSG, *post, {"a": arr(5)}, h) c.rejected( "json sibling array product", s, EVENT_MSG, *post, {"a": arr(2), "b": arr(3)}, h @@ -391,6 +423,20 @@ def run_partitioned(c): sentinel = unique("sentinel") ev = lambda t: {"tenant": t, "x": arr(3)} c.accepted("partitioned request at row limit", s, 6, *post, [ev("a"), ev("b")], h) + # Every partitioned event is flattened on its own, the earlier ones before the last one + # goes over the limit, and still none of them is written + flat_part = lambda events: [{"tenant": f"t{i}", **e} for i, e in enumerate(events)] + flat_sentinel = unique("sentinel") + c.accepted("partitioned flat batch at request row limit", s, 6, *post, flat_part(flat(6)), h) + c.rejected( + "partitioned flat batch over request row limit", + s, + REQUEST_MSG, + *post, + flat_part(flat(7, flat_sentinel)), + h, + flat_sentinel, + ) c.rejected( "partitioned request rows, late partition", s, @@ -462,6 +508,17 @@ def run_kinesis(c): h = {"X-P-Stream": s, "X-P-Log-Source": "kinesis"} post = ("POST", "/api/v1/ingest") c.accepted("kinesis request at row limit", s, 6, *post, kinesis([{"a": arr(3)}] * 2), h) + flat_sentinel = unique("sentinel") + c.accepted("kinesis flat records at request row limit", s, 6, *post, kinesis(flat(6)), h) + c.rejected( + "kinesis flat records over request row limit", + s, + REQUEST_MSG, + *post, + kinesis(flat(7, flat_sentinel)), + h, + flat_sentinel, + ) # A decoded array is one event c.rejected( "kinesis decoded array over event limit", @@ -560,6 +617,30 @@ def run_otel(c): sentinel, ) c.accepted("otel body that is not JSON", s, 1, *post, otel([[["plain text"]]]), h) + # A record whose body is a JSON object without arrays is a row; one whose body is not + # JSON is none + flat_sentinel = unique("sentinel") + flat_bodies = lambda events: [json.dumps(event) for event in events] + c.accepted( + "otel flat JSON bodies at request row limit", s, 6, *post, otel([[flat_bodies(flat(6))]]), h + ) + c.rejected( + "otel flat JSON bodies over request row limit", + s, + REQUEST_MSG, + *post, + otel([[flat_bodies(flat(7, flat_sentinel))]]), + h, + flat_sentinel, + ) + c.accepted( + "otel bodies that are not JSON over request row limit", + s, + 7, + *post, + otel([[["plain text"] * 7]]), + h, + ) # A single row of 32 small fields. Every field is renamed to `{key}_{field}` when it is # merged into the record, so a long key costs 32 times its length diff --git a/tools/flatten-bench/src/alloc_counter.rs b/tools/flatten-bench/src/alloc_counter.rs index 96d3865fc..6b21d1fb5 100644 --- a/tools/flatten-bench/src/alloc_counter.rs +++ b/tools/flatten-bench/src/alloc_counter.rs @@ -2,14 +2,27 @@ //! [`Measure`] guard is armed on it. Other threads (the test harness, stdout capture) are //! never counted, so measurements stay isolated without `--test-threads=1`. //! +//! A guard started with [`Measure::start_bounded`] (or [`measure_bounded`]) also caps the +//! bytes requested on the armed thread. An allocation that would take them over the cap is +//! stopped before it reaches the system allocator: the process writes a line starting with +//! `alloc_counter:` to stderr and exits with status [`CAP_EXCEEDED_EXIT_CODE`]. An allocator +//! must not panic or unwind, so this is not an ordinary test failure; it keeps a rejection +//! test whose input started to be expanded from using more memory than the cap. Exiting +//! flushes stdout, so the code measured under a cap must not print to stdout. Guards from +//! [`Measure::start`] and [`measure`], which the benchmark uses, have no cap. +//! //! Install it in each binary or integration test with //! `#[global_allocator] static A: CountingAlloc = CountingAlloc;`. use std::alloc::{GlobalAlloc, Layout, System}; use std::cell::Cell; +use std::io::Write; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; use std::sync::{Mutex, MutexGuard}; +/// Exit status of a process in which an allocation went over the cap of a bounded [`Measure`] +pub const CAP_EXCEEDED_EXIT_CODE: i32 = 88; + pub struct CountingAlloc; thread_local! { @@ -18,6 +31,8 @@ thread_local! { static COUNT: AtomicU64 = AtomicU64::new(0); static TOTAL_BYTES: AtomicU64 = AtomicU64::new(0); +// Cap on TOTAL_BYTES while armed; u64::MAX for none +static CAP_BYTES: AtomicU64 = AtomicU64::new(u64::MAX); // Live bytes relative to arming; goes negative when pre-existing memory (e.g. a consumed // input) is freed while armed static CURRENT: AtomicI64 = AtomicI64::new(0); @@ -29,6 +44,49 @@ fn armed() -> bool { ARMED.try_with(Cell::get).unwrap_or(false) } +// Called before an armed allocation of `size` bytes is handed to the system allocator +fn check_cap(size: usize) { + let total = TOTAL_BYTES + .load(Ordering::Relaxed) + .saturating_add(size as u64); + let cap = CAP_BYTES.load(Ordering::Relaxed); + if total > cap { + cap_exceeded(size, total, cap); + } +} + +#[cold] +fn cap_exceeded(size: usize, total: u64, cap: u64) -> ! { + // Nothing from here on is counted or capped. The message is formatted on the stack and + // written to a duplicate of fd 2, as `Stderr` takes a lock the measured code may hold. + let _ = ARMED.try_with(|armed| armed.set(false)); + let mut message = [0u8; 512]; + let mut cursor = std::io::Cursor::new(&mut message[..]); + let _ = writeln!( + cursor, + "alloc_counter: an allocation of {size} bytes on thread {:?} would take the bytes \ + requested while measuring to {total}, over the cap of {cap}; exiting with status \ + {CAP_EXCEEDED_EXIT_CODE}", + std::thread::current().name().unwrap_or(""), + ); + let len = cursor.position() as usize; + write_stderr(&message[..len]); + std::process::exit(CAP_EXCEEDED_EXIT_CODE) +} + +#[cfg(unix)] +fn write_stderr(bytes: &[u8]) { + use std::os::fd::AsFd; + if let Ok(fd) = std::io::stderr().as_fd().try_clone_to_owned() { + let _ = std::fs::File::from(fd).write_all(bytes); + } +} + +#[cfg(not(unix))] +fn write_stderr(bytes: &[u8]) { + let _ = std::io::stderr().write_all(bytes); +} + fn on_alloc(size: usize) { COUNT.fetch_add(1, Ordering::Relaxed); TOTAL_BYTES.fetch_add(size as u64, Ordering::Relaxed); @@ -42,16 +100,24 @@ fn on_dealloc(size: usize) { unsafe impl GlobalAlloc for CountingAlloc { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let armed = armed(); + if armed { + check_cap(layout.size()); + } let ptr = unsafe { System.alloc(layout) }; - if !ptr.is_null() && armed() { + if !ptr.is_null() && armed { on_alloc(layout.size()); } ptr } unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let armed = armed(); + if armed { + check_cap(layout.size()); + } let ptr = unsafe { System.alloc_zeroed(layout) }; - if !ptr.is_null() && armed() { + if !ptr.is_null() && armed { on_alloc(layout.size()); } ptr @@ -65,8 +131,12 @@ unsafe impl GlobalAlloc for CountingAlloc { } unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let armed = armed(); + if armed { + check_cap(new_size); + } let new = unsafe { System.realloc(ptr, layout, new_size) }; - if !new.is_null() && armed() { + if !new.is_null() && armed { // Counted as a fresh allocation of `new_size` replacing the old block on_dealloc(layout.size()); on_alloc(new_size); @@ -97,12 +167,24 @@ pub struct Measure { } impl Measure { + /// Counts without a cap pub fn start() -> Self { + Self::start_with_cap(u64::MAX) + } + + /// Counts, and exits the process (see the module docs) before an allocation that would + /// take the bytes requested on this thread over `cap_bytes` is made + pub fn start_bounded(cap_bytes: u64) -> Self { + Self::start_with_cap(cap_bytes) + } + + fn start_with_cap(cap_bytes: u64) -> Self { let lock = LOCK.lock().unwrap_or_else(|p| p.into_inner()); COUNT.store(0, Ordering::SeqCst); TOTAL_BYTES.store(0, Ordering::SeqCst); CURRENT.store(0, Ordering::SeqCst); PEAK.store(0, Ordering::SeqCst); + CAP_BYTES.store(cap_bytes, Ordering::SeqCst); ARMED.with(|a| a.set(true)); Self { _lock: lock } } @@ -123,6 +205,7 @@ impl Measure { impl Drop for Measure { fn drop(&mut self) { ARMED.with(|a| a.set(false)); + CAP_BYTES.store(u64::MAX, Ordering::SeqCst); } } @@ -132,3 +215,10 @@ pub fn measure(f: impl FnOnce() -> R) -> (R, AllocStats) { let result = f(); (result, guard.finish()) } + +/// [`measure`] under the cap of [`Measure::start_bounded`] +pub fn measure_bounded(cap_bytes: u64, f: impl FnOnce() -> R) -> (R, AllocStats) { + let guard = Measure::start_bounded(cap_bytes); + let result = f(); + (result, guard.finish()) +} diff --git a/tools/flatten-bench/tests/alloc_limits.rs b/tools/flatten-bench/tests/alloc_limits.rs index d0b5d01bb..34ab46da8 100644 --- a/tools/flatten-bench/tests/alloc_limits.rs +++ b/tools/flatten-bench/tests/alloc_limits.rs @@ -5,7 +5,9 @@ //! thread while a guard is armed are counted. Run with RUST_LIB_BACKTRACE=0 so anyhow //! does not capture (and allocate) a backtrace for the rejected requests. -use flatten_bench::alloc_counter::{AllocStats, CountingAlloc, measure}; +use flatten_bench::alloc_counter::{ + AllocStats, CAP_EXCEEDED_EXIT_CODE, CountingAlloc, measure, measure_bounded, +}; use flatten_bench::api::{ FlattenBudget, FlattenCost, FlattenLimits, JsonFlattenError, UNLIMITED, baseline_generic, current_charge, current_convert, current_convert_with, current_cost, current_generic, @@ -17,6 +19,9 @@ use serde_json::{Value, json}; #[global_allocator] static ALLOC: CountingAlloc = CountingAlloc; +/// What a rejection may allocate. Rejections are asserted to stay under it, and measured +/// with it as a hard cap (`measure_bounded`): a rejection that started to expand its input +/// exits the test process with `CAP_EXCEEDED_EXIT_CODE` before it allocates more. const MAX_REJECT_BYTES: u64 = 64 * 1024; const DEFAULT: FlattenLimits = FlattenLimits::DEFAULT; @@ -77,7 +82,7 @@ fn hostile_inputs() -> Vec<(&'static str, Value)> { #[test] fn preflight_rejects_without_allocating() { for (name, input) in hostile_inputs() { - let (result, stats) = measure(|| current_cost(&input, &DEFAULT)); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || current_cost(&input, &DEFAULT)); assert!( matches!(&result, Err(err) if err.is_limit_exceeded()), "{name}: {result:?}" @@ -85,14 +90,16 @@ fn preflight_rejects_without_allocating() { assert_zero(&format!("preflight {name}"), stats); // Per-event charging on a fresh request budget, as the entry point does it - let (result, stats) = measure(|| current_charge(&input, DEFAULT)); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || current_charge(&input, DEFAULT)); assert!( matches!(&result, Err(err) if err.is_limit_exceeded()), "{name}: {result:?}" ); assert_zero(&format!("charge {name}"), stats); - let (result, stats) = measure(|| current_pipeline_charge(&input, DEFAULT)); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { + current_pipeline_charge(&input, DEFAULT) + }); assert!( matches!(&result, Err(err) if err.is_limit_exceeded()), "{name}: {result:?}" @@ -105,7 +112,7 @@ fn preflight_rejects_without_allocating() { fn bounded_generic_flattening_rejects_cheaply() { for (name, input) in hostile_inputs() { let mut budget = FlattenBudget::new(DEFAULT); - let (result, stats) = measure(|| { + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { let result = flatten_bench::current_json::flatten::generic_flattening(&input, &mut budget); result.map(|rows| rows.len()) @@ -123,8 +130,9 @@ fn bounded_generic_flattening_rejects_cheaply() { fn ingestion_entry_point_rejects_cheaply() { for (name, input) in hostile_inputs() { // The consumed input is freed inside the guard, which only lowers `net` - let (result, stats) = - measure(|| current_convert(input, None, DEFAULT).map(|rows| rows.len())); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { + current_convert(input, None, DEFAULT).map(|rows| rows.len()) + }); let err = result.expect_err(name); limit_error(&err); // `err` (an anyhow box, 112 bytes today) is allocated inside the guard, so @@ -173,7 +181,7 @@ fn insufficient_remaining_request_budget() { bytes: 0, }) .unwrap(); - let (result, stats) = measure(|| { + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { flatten_bench::current_json::flatten::charge_generic_flattening(&product, &mut budget) }); assert!( @@ -184,8 +192,9 @@ fn insufficient_remaining_request_budget() { "{result:?}" ); assert_zero("core charge, remaining rows insufficient", stats); - let (result, stats) = - measure(|| current_convert_with(product, None, &mut budget).map(|r| r.len())); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { + current_convert_with(product, None, &mut budget).map(|r| r.len()) + }); let err = result.expect_err("1000 rows cannot fit in the remaining 5"); assert!( matches!( @@ -207,7 +216,7 @@ fn insufficient_remaining_request_budget() { }) .unwrap(); let small = json!({"message": "this row is longer than sixteen bytes"}); - let (result, stats) = measure(|| { + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { flatten_bench::current_json::flatten::charge_generic_flattening(&small, &mut budget) }); assert!( @@ -218,8 +227,9 @@ fn insufficient_remaining_request_budget() { "{result:?}" ); assert_zero("core charge, remaining bytes insufficient", stats); - let (result, stats) = - measure(|| current_convert_with(small, None, &mut budget).map(|r| r.len())); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { + current_convert_with(small, None, &mut budget).map(|r| r.len()) + }); let err = result.expect_err("row does not fit in the remaining 16 bytes"); assert!( matches!( @@ -255,6 +265,49 @@ fn request_budget_is_shared_across_partitions() { report("partitioned request over limit", stats); } +/// Events without arrays count toward the request limit too, a row each, though generic +/// flattening leaves them as they are: 10 000 of them are accepted and 10 001 rejected, in +/// under 100 KiB of JSON +#[test] +fn flat_batch_counts_a_row_per_event() { + let limit = DEFAULT.max_rows_per_request; + let batch = |events: usize| Value::Array((0..events).map(|i| json!({"t": i % 4})).collect()); + let over = batch(limit + 1); + assert!(serde_json::to_vec(&over).unwrap().len() < 100 * 1024); + + // Without a partition the batch comes back as a single array, unchanged + let rows = current_convert(batch(limit), None, DEFAULT).unwrap(); + assert_eq!(rows[0], batch(limit)); + let input = over.clone(); + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { + current_convert(input, None, DEFAULT).map(|rows| rows.len()) + }); + let err = result.expect_err("10 001 events are over the 10 000 rows of a request"); + assert!( + matches!( + limit_error(&err), + JsonFlattenError::RowsPerRequestLimitExceeded(10_000) + ), + "{err:#}" + ); + // Charged in the walk that finds them free of arrays, nothing built + assert_small("flat batch over request rows", stats); + + // With a partition every event is flattened on its own, so the events before the + // 10 001st are built before it is charged; the request is rejected all the same + let partition = "t".to_owned(); + let rows = current_convert(batch(limit), Some(&partition), DEFAULT).unwrap(); + assert_eq!(rows.len(), limit); + let err = current_convert(over, Some(&partition), DEFAULT).expect_err("over the limit"); + assert!( + matches!( + limit_error(&err), + JsonFlattenError::RowsPerRequestLimitExceeded(10_000) + ), + "{err:#}" + ); +} + /// Small deterministic generator of mixed shapes, bounded far below 1000 rows. fn generated(seed: u64, count: usize) -> Vec { let mut state = seed | 1; @@ -372,6 +425,43 @@ fn counter_sees_allocations() { stats.total_bytes >= 100_000 && stats.peak_bytes >= 100_000, "{stats:?}" ); + // Under the cap, a bounded measurement counts the same + let (len, bounded) = measure_bounded(MAX_REJECT_BYTES * 2, || vec![0u8; 100_000].len()); + assert_eq!(len, 100_000); + assert_eq!(bounded, stats); +} + +const CAP_CHILD_ENV: &str = "FLATTEN_BENCH_ALLOCATION_CAP_CHILD"; + +/// An allocation over the cap of a bounded measurement exits the process instead of being +/// made, so this runs `allocation_cap_child` in a process of its own +#[test] +fn allocation_over_the_cap_exits() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "allocation_cap_child", "--ignored"]) + .env(CAP_CHILD_ENV, "1") + .output() + .unwrap(); + let stderr = String::from_utf8_lossy(&output.stderr); + assert_eq!( + output.status.code(), + Some(CAP_EXCEEDED_EXIT_CODE), + "{stderr}" + ); + assert!( + stderr.contains("alloc_counter: an allocation of 1048576 bytes"), + "{stderr}" + ); +} + +#[test] +#[ignore = "exits the process; allocation_over_the_cap_exits runs it"] +fn allocation_cap_child() { + if std::env::var_os(CAP_CHILD_ENV).is_none() { + return; + } + let (len, stats) = measure_bounded(MAX_REJECT_BYTES, || vec![0u8; 1 << 20].len()); + panic!("{len} bytes allocated over the cap: {stats:?}"); } /// The harness stubs `crate::cli`; make sure its defaults are the server's, and the @@ -427,7 +517,8 @@ fn structural_field_names_are_rejected_without_expanding() { // The generic rows alone fit assert!(current_cost(&input, &small).is_ok()); - let (result, stats) = measure(|| current_pipeline_cost(&input, &small)); + let (result, stats) = + measure_bounded(MAX_REJECT_BYTES, || current_pipeline_cost(&input, &small)); assert!( matches!( result, @@ -436,7 +527,8 @@ fn structural_field_names_are_rejected_without_expanding() { "{result:?}" ); assert_zero("structural preflight, long parent", stats); - let (result, stats) = measure(|| current_pipeline_charge(&input, small)); + let (result, stats) = + measure_bounded(MAX_REJECT_BYTES, || current_pipeline_charge(&input, small)); assert!( matches!( result, @@ -449,7 +541,7 @@ fn structural_field_names_are_rejected_without_expanding() { for partition in [None, Some("r".to_owned())] { let input = input.clone(); let mut budget = FlattenBudget::new(small); - let (result, stats) = measure(|| { + let (result, stats) = measure_bounded(MAX_REJECT_BYTES, || { current_convert_with(input, partition.as_ref(), &mut budget).map(|rows| rows.len()) }); let err = result.expect_err("long parent names exceed 4096 bytes");