diff --git a/src/connectors/kafka/processor.rs b/src/connectors/kafka/processor.rs index a892691b2..cbac3e66f 100644 --- a/src/connectors/kafka/processor.rs +++ b/src/connectors/kafka/processor.rs @@ -111,6 +111,9 @@ fn build_and_process_event( let mut p_custom_fields = HashMap::new(); p_custom_fields.insert(USER_AGENT_KEY.to_string(), "kafka".to_string()); + // Offsets are committed whether or not the chunk is ingested and the + // producer gets no rejection signal, so the dataset field limit is not + // enforced here: rejecting would silently drop the whole chunk. json::Event::new(Value::Array(json_vec), Utc::now()) .into_event( stream_name, @@ -126,7 +129,7 @@ fn build_and_process_event( &tenant_id, meta.infer_timestamp, )? - .process()?; + .process_without_field_limit()?; Ok(()) } @@ -198,9 +201,13 @@ where .for_each_concurrent(None, |records| async { if let Some(last_record) = records.iter().max_by_key(|r| r.offset) { let tpl = last_record.offset_to_commit().unwrap(); + let record_count = records.len(); if let Err(e) = self.processor.process(records).await { - error!("Failed to process records for {:?}: {:?}", tp, e); + error!( + "Failed to process {} records for {:?}; offsets up to {:?} will still be committed and these records are dropped: {:?}", + record_count, tp, tpl, e + ); } //CommitMode::Async race condition. diff --git a/src/event/mod.rs b/src/event/mod.rs index 4baa8f8c7..cfd075b0c 100644 --- a/src/event/mod.rs +++ b/src/event/mod.rs @@ -20,9 +20,10 @@ pub mod format; use arrow_array::RecordBatch; -use arrow_schema::{Field, Fields, Schema}; +use arrow_schema::{ArrowError, Field, Fields, Schema, SchemaRef}; use itertools::Itertools; -use std::sync::Arc; +use std::collections::HashSet; +use std::sync::{Arc, RwLock}; use tracing::{info_span, instrument}; @@ -30,7 +31,7 @@ use self::error::EventError; use crate::{ LOCK_EXPECT, handlers::TelemetryType, - metadata::update_stats, + metadata::{LogStreamMetadata, update_stats}, metrics::{increment_events_ingested_by_date, increment_events_ingested_size_by_date}, parseable::{DEFAULT_TENANT, PARSEABLE, StagingError, StreamNotFound}, storage::StreamType, @@ -62,6 +63,26 @@ pub struct Event { // Events holds the schema related to a each event for a single log stream impl Event { + /// Processes the event, holding user-defined streams to the configured + /// dataset field limit. + pub fn process(self) -> Result<(), EventError> { + let limit = PARSEABLE.options.dataset_fields_allowed_limit; + self.process_with_field_limit(limit) + } + + /// Processes the event without enforcing the dataset field limit. + /// + /// For connector ingestion (such as Kafka), which commits source offsets + /// regardless of the outcome and cannot report a rejection back to the + /// producer; enforcing the limit there would silently drop data. + /// + /// Also used for events whose schemas were already checked and committed + /// together with `commit_ingested_schemas`, so that a request yielding + /// several events cannot fail on the limit after some were staged. + pub fn process_without_field_limit(self) -> Result<(), EventError> { + self.process_with_field_limit(usize::MAX) + } + #[instrument( name = "event_process", level = "info", @@ -73,7 +94,7 @@ impl Event { ) )] #[cfg_attr(feature = "hotpath", hotpath::measure)] - pub fn process(self) -> Result<(), EventError> { + fn process_with_field_limit(self, field_limit: usize) -> Result<(), EventError> { let mut key = get_schema_key(&self.rb.schema().fields); if self.time_partition.is_some() { let parsed_timestamp_to_min = self.parsed_timestamp.format("%Y%m%dT%H%M").to_string(); @@ -88,11 +109,24 @@ impl Event { let stream = PARSEABLE.get_or_create_stream(&self.stream_name, &self.tenant_id); - if self.is_first_event { + // User ingestion is held to the given field limit; internal + // streams have fixed schemas and are exempt. + let schema_committed = if self.stream_type == StreamType::UserDefined { + commit_ingested_schema( + &self.stream_name, + &stream.metadata, + &self.rb.schema(), + field_limit, + self.is_first_event, + )? + } else if self.is_first_event { commit_schema(&self.stream_name, self.rb.schema(), &self.tenant_id)?; - if !stream.get_static_schema_flag() { - stream.stage_schema_file(stream.get_schema().as_ref().clone())?; - } + true + } else { + false + }; + if schema_committed && !stream.get_static_schema_flag() { + stream.stage_schema_file(stream.get_schema().as_ref().clone())?; } stream.push( @@ -172,8 +206,216 @@ pub fn commit_schema( .write() .expect(LOCK_EXPECT) .schema; + merge_into_schema(map, &schema)?; + Ok(()) +} + +/// Checks an ingested event against the stream's field limit and commits its +/// schema to the stream's in-memory schema when needed. +/// +/// The event is rejected if the stream's fields together with the event's +/// would exceed `limit`, which includes streams already over the limit. A first +/// event's schema is merged in. Any other event on a dynamic schema has its +/// missing fields, such as new header fields, added without merging the rest. +/// Events on a static schema (first or not) never change the stored schema, so +/// only the stored schema is checked against the limit; the generated `p_` and +/// header fields such events carry are not counted. +/// The final count and commit happen under one write lock on the stream +/// metadata, so concurrent events cannot each pass the check and together +/// exceed the limit; events adding nothing only take a read lock. Startup and +/// query-time schema refreshes use `commit_schema`, which has no limit. +/// +/// Returns whether the schema was committed. +pub(crate) fn commit_ingested_schema( + stream_name: &str, + metadata: &RwLock, + schema: &Schema, + limit: usize, + is_first_event: bool, +) -> Result { + { + let metadata = metadata.read().expect(LOCK_EXPECT); + if metadata.static_schema_flag { + // The stored schema cannot grow, so the event's extra fields + // (generated and header columns) must not count. + check_field_limit(stream_name, metadata.schema.len(), limit)?; + return Ok(false); + } + if !is_first_event && !has_new_fields(&metadata.schema, schema) { + // Common case: nothing to count or commit, so avoid building the + // set of new field names. + check_field_limit(stream_name, metadata.schema.len(), limit)?; + return Ok(false); + } + let new_fields = new_field_names(&metadata.schema, [schema]).len(); + check_field_limit(stream_name, metadata.schema.len() + new_fields, limit)?; + } + + let _span = info_span!("commit_ingested_schema", stream_name).entered(); + let mut metadata = metadata.write().expect(LOCK_EXPECT); + if metadata.static_schema_flag { + check_field_limit(stream_name, metadata.schema.len(), limit)?; + return Ok(false); + } + let new_fields = new_field_names(&metadata.schema, [schema]).len(); + check_field_limit(stream_name, metadata.schema.len() + new_fields, limit)?; + if is_first_event { + merge_into_schema(&mut metadata.schema, schema).map_err(StagingError::from)?; + } else { + if new_fields == 0 { + return Ok(false); + } + // Only add the missing fields: a header field sharing a name with a + // typed stream field was never merged and must not fail the event. + add_missing_fields(&mut metadata.schema, schema); + } + Ok(true) +} + +/// Checks the union of several events' schemas against the stream's field +/// limit and commits all of them under a single write lock, applying the same +/// per-event rules as `commit_ingested_schema` (merge first events, add only +/// missing fields for others, never grow a static schema). +/// +/// Either every schema is committed or none is, so a request yielding several +/// events reserves its fields before staging any event. The events can then +/// be processed without a further limit check, so the request cannot fail on +/// the limit partway even if concurrent requests add fields meanwhile. +/// +/// Returns whether the schema was committed. +pub(crate) fn commit_ingested_schemas( + stream_name: &str, + metadata: &RwLock, + schemas: &[(SchemaRef, bool)], + limit: usize, +) -> Result { + let any_first_event = schemas.iter().any(|(_, is_first)| *is_first); + { + let metadata = metadata.read().expect(LOCK_EXPECT); + let new_fields = check_schemas_locked( + stream_name, + &metadata, + schemas.iter().map(|(schema, _)| &**schema), + limit, + )?; + if metadata.static_schema_flag || (new_fields == 0 && !any_first_event) { + return Ok(false); + } + } + + let _span = info_span!("commit_ingested_schemas", stream_name).entered(); + let mut metadata = metadata.write().expect(LOCK_EXPECT); + let new_fields = check_schemas_locked( + stream_name, + &metadata, + schemas.iter().map(|(schema, _)| &**schema), + limit, + )?; + if metadata.static_schema_flag || (new_fields == 0 && !any_first_event) { + return Ok(false); + } + // Build the result on a copy so a merge failure leaves the schema as is. + let mut merged = metadata.schema.clone(); + for (schema, is_first_event) in schemas { + if *is_first_event { + merge_into_schema(&mut merged, schema).map_err(StagingError::from)?; + } else { + add_missing_fields(&mut merged, schema); + } + } + metadata.schema = merged; + Ok(true) +} + +/// Checks that the stream's fields together with every schema in `schemas` +/// stay within `limit`, counting each new field name once. +/// +/// For static-schema streams only the stored schema is checked, since +/// ingestion cannot grow it. +/// +/// It does not reserve room; use `commit_ingested_schemas` for that. +#[cfg_attr(not(test), allow(dead_code))] +pub(crate) fn check_ingested_schemas( + stream_name: &str, + metadata: &RwLock, + schemas: &[SchemaRef], + limit: usize, +) -> Result<(), EventError> { + let metadata = metadata.read().expect(LOCK_EXPECT); + check_schemas_locked( + stream_name, + &metadata, + schemas.iter().map(|schema| &**schema), + limit, + ) + .map(|_| ()) +} + +/// Checks `schemas` against the limit given already locked stream metadata, +/// returning the number of new field names (always 0 for static schemas). +fn check_schemas_locked<'a>( + stream_name: &str, + metadata: &LogStreamMetadata, + schemas: impl IntoIterator, + limit: usize, +) -> Result { + if metadata.static_schema_flag { + check_field_limit(stream_name, metadata.schema.len(), limit)?; + return Ok(0); + } + let new_fields = new_field_names(&metadata.schema, schemas).len(); + check_field_limit(stream_name, metadata.schema.len() + new_fields, limit)?; + Ok(new_fields) +} + +/// Whether `schema` has any field name that `current` lacks. +fn has_new_fields(current: &HashMap>, schema: &Schema) -> bool { + schema + .fields() + .iter() + .any(|field| !current.contains_key(field.name())) +} + +/// Adds the fields of `schema` whose names `map` lacks, leaving existing +/// fields untouched. +fn add_missing_fields(map: &mut HashMap>, schema: &Schema) { + for field in schema.fields() { + if !map.contains_key(field.name()) { + map.insert(field.name().clone(), field.clone()); + } + } +} + +/// Names of the fields in `schemas` that `current` lacks, without duplicates. +fn new_field_names<'a>( + current: &HashMap>, + schemas: impl IntoIterator, +) -> HashSet<&'a str> { + schemas + .into_iter() + .flat_map(|schema| schema.fields().iter()) + .map(|field| field.name().as_str()) + .filter(|name| !current.contains_key(*name)) + .collect() +} + +fn check_field_limit(stream_name: &str, total: usize, limit: usize) -> Result<(), EventError> { + if total > limit { + return Err(EventError::FieldsCountLimitExceeded( + stream_name.to_owned(), + total, + limit, + )); + } + Ok(()) +} + +fn merge_into_schema( + map: &mut HashMap>, + schema: &Schema, +) -> Result<(), ArrowError> { let current_schema = Schema::new(map.values().cloned().collect::()); - let schema = Schema::try_merge(vec![current_schema, schema.as_ref().clone()])?; + let schema = Schema::try_merge(vec![current_schema, schema.clone()])?; map.clear(); map.extend(schema.fields.iter().map(|f| (f.name().clone(), f.clone()))); Ok(()) @@ -189,5 +431,465 @@ pub mod error { Staging(#[from] StagingError), #[error("ObjectStorage Error: {0}")] ObjectStorage(#[from] ObjectStorageError), + #[error( + "Failed to ingest events in dataset {0}. Total number of fields {1} exceeds the permissible limit of {2}." + )] + FieldsCountLimitExceeded(String, usize, usize), + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + fmt::Debug, + sync::{Arc, Barrier, RwLock}, + thread, + }; + + use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; + use chrono::Utc; + use serde_json::{Value, json}; + + use super::{ + DEFAULT_TIMESTAMP_KEY, check_ingested_schemas, commit_ingested_schema, + commit_ingested_schemas, error::EventError, merge_into_schema, + }; + use crate::{ + event::format::{EventFormat, json}, + metadata::{LogStreamMetadata, SchemaVersion}, + }; + + const STREAM: &str = "logs"; + + fn p_timestamp() -> Field { + Field::new( + DEFAULT_TIMESTAMP_KEY, + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ) + } + + fn utf8_fields(prefix: &str, count: usize) -> Vec { + (0..count) + .map(|i| Field::new(format!("{prefix}{i}"), DataType::Utf8, true)) + .collect() + } + + fn metadata(fields: Vec) -> RwLock { + RwLock::new(LogStreamMetadata { + schema: fields + .into_iter() + .map(|field| (field.name().clone(), Arc::new(field))) + .collect(), + ..Default::default() + }) + } + + fn field_names(metadata: &RwLock) -> Vec { + let mut names: Vec<_> = metadata.read().unwrap().schema.keys().cloned().collect(); + names.sort(); + names + } + + /// A record with a string value for each of the given keys. + fn record(prefix: &str, count: usize) -> Value { + Value::Object( + (0..count) + .map(|i| (format!("{prefix}{i}"), json!("v"))) + .collect(), + ) + } + + /// The schema and first-event flag ingestion derives for `json` against + /// the stream's current schema, including the generated `p_` and header + /// fields. + fn event_schema( + metadata: &RwLock, + json: Value, + headers: &[(&str, &str)], + ) -> (SchemaRef, bool) { + let stream_schema = metadata.read().unwrap().schema.clone(); + let static_schema_flag = metadata.read().unwrap().static_schema_flag; + let headers: HashMap = headers + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + let (rb, is_first_event) = json::Event::new(json, Utc::now()) + .into_recordbatch( + &stream_schema, + static_schema_flag, + None, + SchemaVersion::V1, + &headers, + false, + ) + .unwrap(); + (rb.schema(), is_first_event) + } + + /// Commits `json` as ingestion would, returning whether the schema changed. + fn ingest( + metadata: &RwLock, + json: Value, + headers: &[(&str, &str)], + limit: usize, + ) -> Result { + let (schema, is_first_event) = event_schema(metadata, json, headers); + commit_ingested_schema(STREAM, metadata, &schema, limit, is_first_event) + } + + fn assert_limit_error(result: Result, total: usize, limit: usize) { + match result { + Err(EventError::FieldsCountLimitExceeded(name, t, l)) => { + assert_eq!((name.as_str(), t, l), (STREAM, total, limit)); + } + other => panic!("expected a field limit error, got {other:?}"), + } + } + + #[test] + fn accepts_event_reaching_exactly_the_limit() { + // 9 fields plus the generated p_timestamp + let stream = metadata(vec![]); + assert!(ingest(&stream, record("f", 9), &[], 10).unwrap()); + assert_eq!(field_names(&stream).len(), 10); + assert!(field_names(&stream).contains(&DEFAULT_TIMESTAMP_KEY.to_owned())); + + // Known fields at exactly the limit keep being accepted + assert!(!ingest(&stream, record("f", 9), &[], 10).unwrap()); + assert_eq!(field_names(&stream).len(), 10); + } + + #[test] + fn rejects_event_one_field_over_the_limit() { + let stream = metadata(vec![]); + assert_limit_error(ingest(&stream, record("f", 9), &[], 9), 10, 9); + assert!(field_names(&stream).is_empty()); + } + + #[test] + fn rejects_huge_single_event_on_empty_stream() { + let stream = metadata(vec![]); + assert_limit_error(ingest(&stream, record("f", 1500), &[], 1000), 1501, 1000); + assert!(field_names(&stream).is_empty()); + } + + #[test] + fn unlimited_commit_accepts_wide_and_over_limit_streams() { + // The connector path commits with no effective limit + let stream = metadata(vec![]); + assert!(ingest(&stream, record("f", 1500), &[], usize::MAX).unwrap()); + assert_eq!(field_names(&stream).len(), 1501); + assert!(ingest(&stream, record("g", 5), &[], usize::MAX).unwrap()); + assert_eq!(field_names(&stream).len(), 1506); + } + + #[test] + fn rejects_disjoint_batches_once_they_add_up_past_the_limit() { + let stream = metadata(vec![]); + // p_timestamp + 3 + 3 + 3 = 10 + for prefix in ["a", "b", "c"] { + assert!(ingest(&stream, record(prefix, 3), &[], 10).unwrap()); + } + let before = field_names(&stream); + assert_eq!(before.len(), 10); + + assert_limit_error(ingest(&stream, record("d", 3), &[], 10), 13, 10); + assert_eq!(field_names(&stream), before); + } + + #[test] + fn counts_existing_and_repeated_fields_once() { + let mut fields = utf8_fields("f", 5); + fields.push(p_timestamp()); + let stream = metadata(fields); + // Overlaps the stream on f0..f4 and p_timestamp; only f5..f8 are new + assert!(ingest(&stream, record("f", 9), &[], 10).unwrap()); + assert_eq!(field_names(&stream).len(), 10); + + // The same names repeated across schemas count once + let schema: SchemaRef = Arc::new(Schema::new(utf8_fields("g", 3))); + let schemas = [schema.clone(), schema.clone(), schema]; + let stream = metadata(vec![p_timestamp()]); + check_ingested_schemas(STREAM, &stream, &schemas, 4).unwrap(); + assert_limit_error(check_ingested_schemas(STREAM, &stream, &schemas, 3), 4, 3); + } + + #[test] + fn counts_generated_header_and_typed_sibling_fields() { + let stream = metadata(vec![ + p_timestamp(), + Field::new("count", DataType::Int64, true), + ]); + let json = json!({"count": "not a number", "msg": "hello"}); + let headers = [("p_user_agent", "curl"), ("env", "prod")]; + + let (schema, is_first_event) = event_schema(&stream, json, &headers); + assert!(is_first_event); + let names: Vec<_> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + let sibling = names + .iter() + .find(|name| name.starts_with("count_")) + .unwrap_or_else(|| panic!("no typed sibling for count in {names:?}")); + assert!(!names.contains(&"count"), "{names:?}"); + // p_timestamp and count already exist; the sibling, msg and both + // header fields are new + let new_fields = [*sibling, "msg", "p_user_agent", "env"]; + assert!( + new_fields.iter().all(|name| names.contains(name)), + "{names:?}" + ); + assert_eq!(names.len(), 5, "{names:?}"); + + assert_limit_error( + commit_ingested_schema(STREAM, &stream, &schema, 5, is_first_event), + 6, + 5, + ); + assert_eq!(field_names(&stream).len(), 2); + + assert!(commit_ingested_schema(STREAM, &stream, &schema, 6, is_first_event).unwrap()); + assert_eq!(field_names(&stream).len(), 6); + // The original field keeps its type + assert_eq!( + stream.read().unwrap().schema["count"].data_type(), + &DataType::Int64 + ); + } + + #[test] + fn known_schema_event_counts_and_keeps_new_header_fields() { + let stream = metadata(vec![p_timestamp(), Field::new("msg", DataType::Utf8, true)]); + let json = json!({"msg": "hello"}); + let headers = [("env", "prod")]; + + let (_, is_first_event) = event_schema(&stream, json.clone(), &headers); + assert!( + !is_first_event, + "a known-schema event should not be a first event" + ); + + assert_limit_error(ingest(&stream, json.clone(), &headers, 2), 3, 2); + assert_eq!(field_names(&stream).len(), 2); + + assert!(ingest(&stream, json.clone(), &headers, 3).unwrap()); + assert!(field_names(&stream).contains(&"env".to_owned())); + + // Nothing new the second time, so nothing is committed + assert!(!ingest(&stream, json, &headers, 3).unwrap()); + } + + #[test] + fn known_schema_event_only_adds_missing_fields() { + // A header named like a typed stream field is not merged with it + let stream = metadata(vec![ + p_timestamp(), + Field::new("level", DataType::Int64, true), + ]); + let schema = Schema::new(vec![ + p_timestamp(), + Field::new("level", DataType::Utf8, true), + Field::new("env", DataType::Utf8, true), + ]); + + assert!(commit_ingested_schema(STREAM, &stream, &schema, 10, false).unwrap()); + let stream = stream.read().unwrap(); + assert_eq!(stream.schema["level"].data_type(), &DataType::Int64); + assert_eq!(stream.schema["env"].data_type(), &DataType::Utf8); + } + + #[test] + fn static_schema_known_event_ignores_uncommitted_header_fields() { + let stream = metadata(vec![p_timestamp(), Field::new("msg", DataType::Utf8, true)]); + stream.write().unwrap().static_schema_flag = true; + let json = json!({"msg": "hello"}); + let headers = [("env", "prod"), ("p_user_agent", "curl")]; + + // The stored schema is exactly at the limit; header fields are never + // committed to a static schema, so they must not count, every time. + assert!(!ingest(&stream, json.clone(), &headers, 2).unwrap()); + assert!(!ingest(&stream, json.clone(), &headers, 2).unwrap()); + assert_eq!(field_names(&stream).len(), 2); + + // The whole-request check ignores them too + let (schema, _) = event_schema(&stream, json.clone(), &headers); + check_ingested_schemas(STREAM, &stream, &[schema], 2).unwrap(); + + // A static schema already over the limit is still rejected + assert_limit_error(ingest(&stream, json, &headers, 1), 2, 1); + assert_eq!(field_names(&stream).len(), 2); + } + + #[test] + fn static_schema_first_event_neither_counts_nor_merges_extra_fields() { + let stream = metadata(vec![p_timestamp(), Field::new("msg", DataType::Utf8, true)]); + stream.write().unwrap().static_schema_flag = true; + let schema = Schema::new(vec![ + p_timestamp(), + Field::new("msg", DataType::Utf8, true), + Field::new("p_user_agent", DataType::Utf8, true), + Field::new("env", DataType::Utf8, true), + ]); + + // Treated like a non-first event: the stored schema is at the limit + // and must not grow. + assert!(!commit_ingested_schema(STREAM, &stream, &schema, 2, true).unwrap()); + assert!(!commit_ingested_schema(STREAM, &stream, &schema, 2, false).unwrap()); + assert!( + !commit_ingested_schemas(STREAM, &stream, &[(Arc::new(schema), true)], 2).unwrap() + ); + assert_eq!(field_names(&stream).len(), 2); + } + + #[test] + fn over_limit_stream_rejects_ingestion_but_stays_recoverable() { + let mut wide = utf8_fields("f", 19); + wide.push(p_timestamp()); + let wide_schema = Schema::new(wide); + + // Startup load and query refresh merge without a limit + let stream = metadata(vec![]); + merge_into_schema(&mut stream.write().unwrap().schema, &wide_schema).unwrap(); + let before = field_names(&stream); + assert_eq!(before.len(), 20); + + // Known fields only, on both event paths + let known = Schema::new(vec![p_timestamp(), Field::new("f0", DataType::Utf8, true)]); + assert_limit_error( + commit_ingested_schema(STREAM, &stream, &known, 10, false), + 20, + 10, + ); + assert_limit_error( + commit_ingested_schema(STREAM, &stream, &known, 10, true), + 20, + 10, + ); + // A new field + assert_limit_error(ingest(&stream, record("g", 1), &[], 10), 21, 10); + assert_limit_error( + check_ingested_schemas(STREAM, &stream, &[Arc::new(known)], 10), + 20, + 10, + ); + + assert_eq!(field_names(&stream), before); + } + + #[test] + fn whole_request_check_rejects_union_of_events() { + let stream = metadata(vec![p_timestamp()]); + let a: SchemaRef = Arc::new(Schema::new(utf8_fields("a", 6))); + let b: SchemaRef = Arc::new(Schema::new(utf8_fields("b", 6))); + + // Each event fits on its own; together they do not + check_ingested_schemas(STREAM, &stream, std::slice::from_ref(&a), 12).unwrap(); + check_ingested_schemas(STREAM, &stream, std::slice::from_ref(&b), 12).unwrap(); + assert_limit_error(check_ingested_schemas(STREAM, &stream, &[a, b], 12), 13, 12); + assert_eq!(field_names(&stream).len(), 1); + } + + #[test] + fn whole_request_commit_is_all_or_nothing() { + let stream = metadata(vec![p_timestamp()]); + let a: SchemaRef = Arc::new(Schema::new(utf8_fields("a", 6))); + let b: SchemaRef = Arc::new(Schema::new(utf8_fields("b", 6))); + + // The union does not fit: nothing is committed + assert_limit_error( + commit_ingested_schemas(STREAM, &stream, &[(a.clone(), true), (b.clone(), true)], 12), + 13, + 12, + ); + assert_eq!(field_names(&stream).len(), 1); + + // The union fits: everything is committed at once, so processing the + // events afterwards adds nothing and cannot hit the limit + assert!( + commit_ingested_schemas(STREAM, &stream, &[(a.clone(), true), (b.clone(), false)], 13) + .unwrap() + ); + assert_eq!(field_names(&stream).len(), 13); + assert!(!commit_ingested_schema(STREAM, &stream, &b, 13, false).unwrap()); + // Nothing new: no commit + assert!(!commit_ingested_schemas(STREAM, &stream, &[(b, false)], 13).unwrap()); + + // A type conflict in any event leaves the schema untouched + let before = field_names(&stream); + let c: SchemaRef = Arc::new(Schema::new(utf8_fields("c", 1))); + let conflict: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "a0", + DataType::Int64, + true, + )])); + let result = + commit_ingested_schemas(STREAM, &stream, &[(c, true), (conflict, true)], 100); + assert!(matches!(result, Err(EventError::Staging(_))), "{result:?}"); + assert_eq!(field_names(&stream), before); + } + + #[test] + fn rejected_type_conflict_leaves_schema_unmodified() { + let stream = metadata(vec![p_timestamp(), Field::new("a", DataType::Int64, true)]); + let before = field_names(&stream); + let schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + ]); + + let result = commit_ingested_schema(STREAM, &stream, &schema, 10, true); + assert!(matches!(result, Err(EventError::Staging(_))), "{result:?}"); + assert_eq!(field_names(&stream), before); + } + + #[test] + fn concurrent_writers_cannot_exceed_the_limit_together() { + const WRITERS: usize = 16; + // Room for exactly 4 writers adding 3 fields each + const LIMIT: usize = 1 + 3 * 4; + + let stream = metadata(vec![p_timestamp()]); + let barrier = Barrier::new(WRITERS); + let results: Vec<_> = thread::scope(|scope| { + let handles: Vec<_> = (0..WRITERS) + .map(|writer| { + let (stream, barrier) = (&stream, &barrier); + scope.spawn(move || { + let schema = Schema::new(utf8_fields(&format!("w{writer}_"), 3)); + barrier.wait(); + commit_ingested_schema(STREAM, stream, &schema, LIMIT, true) + }) + }) + .collect(); + handles.into_iter().map(|h| h.join().unwrap()).collect() + }); + + let accepted = results.iter().filter(|r| r.is_ok()).count(); + assert_eq!(accepted, 4); + for result in results.into_iter().filter(Result::is_err) { + assert_limit_error(result, LIMIT + 3, LIMIT); + } + assert_eq!(field_names(&stream).len(), LIMIT); + } + + #[test] + fn concurrent_writers_of_the_same_fields_all_fit() { + const WRITERS: usize = 8; + const LIMIT: usize = 13; + + let stream = metadata(vec![p_timestamp()]); + let barrier = Barrier::new(WRITERS); + thread::scope(|scope| { + for _ in 0..WRITERS { + let (stream, barrier) = (&stream, &barrier); + scope.spawn(move || { + let schema = Schema::new(utf8_fields("f", LIMIT - 1)); + barrier.wait(); + commit_ingested_schema(STREAM, stream, &schema, LIMIT, true).unwrap(); + }); + } + }); + assert_eq!(field_names(&stream).len(), LIMIT); } } diff --git a/src/handlers/http/ingest.rs b/src/handlers/http/ingest.rs index 746ae8a98..31b75d819 100644 --- a/src/handlers/http/ingest.rs +++ b/src/handlers/http/ingest.rs @@ -494,7 +494,7 @@ pub enum PostError { #[error("Header Error: {0}")] Header(#[from] ParseHeaderError), #[error("Event Error: {0}")] - Event(#[from] EventError), + Event(#[source] EventError), #[error("Invalid Request: {0}")] Invalid(#[from] anyhow::Error), #[error("{0}")] @@ -544,6 +544,18 @@ pub enum PostError { StreamBeingDeleted(String), } +impl From for PostError { + fn from(err: EventError) -> Self { + match err { + // A client error, so it must not surface as a generic event failure + EventError::FieldsCountLimitExceeded(stream_name, count, limit) => { + PostError::FieldsCountLimitExceeded(stream_name, count, limit) + } + err => PostError::Event(err), + } + } +} + impl actix_web::ResponseError for PostError { fn status_code(&self) -> StatusCode { use PostError::*; @@ -1138,6 +1150,47 @@ mod tests { ); } + #[test] + fn field_limit_event_error_is_bad_request() { + use actix_web::{ResponseError, http::StatusCode}; + + use super::PostError; + use crate::event::error::EventError; + + let err = PostError::from(EventError::FieldsCountLimitExceeded( + "logs".to_owned(), + 1001, + 1000, + )); + + assert!(matches!( + &err, + PostError::FieldsCountLimitExceeded(name, 1001, 1000) if name == "logs" + )); + assert_eq!(err.status_code(), StatusCode::BAD_REQUEST); + let message = err.to_string(); + assert!(message.contains("logs"), "{message}"); + assert!(message.contains("1001"), "{message}"); + assert!(message.contains("1000"), "{message}"); + } + + #[test] + fn other_event_errors_stay_server_errors() { + use actix_web::{ResponseError, http::StatusCode}; + use arrow_schema::ArrowError; + + use super::PostError; + use crate::{event::error::EventError, parseable::StagingError}; + + let err = PostError::from(EventError::Staging(StagingError::Arrow( + ArrowError::SchemaError("conflict".to_owned()), + ))); + + assert!(matches!(err, PostError::Event(EventError::Staging(_)))); + assert_eq!(err.status_code(), StatusCode::INTERNAL_SERVER_ERROR); + assert!(std::error::Error::source(&err).is_some()); + } + #[test] fn flatten_limit_errors_are_bad_requests() { use actix_web::{ResponseError, http::StatusCode}; diff --git a/src/handlers/http/modal/utils/ingest_utils.rs b/src/handlers/http/modal/utils/ingest_utils.rs index dd261af5f..66936516e 100644 --- a/src/handlers/http/modal/utils/ingest_utils.rs +++ b/src/handlers/http/modal/utils/ingest_utils.rs @@ -31,7 +31,8 @@ use tracing::{error, instrument, warn}; use crate::{ event::{ - FORMAT_KEY, SOURCE_IP_KEY, USER_AGENT_KEY, + FORMAT_KEY, SOURCE_IP_KEY, USER_AGENT_KEY, commit_ingested_schemas, + error::EventError, format::{EventFormat, LogSource, json}, }, handlers::{ @@ -437,8 +438,29 @@ pub fn push_logs( ) }) .collect(); - for event in events.map_err(PostError::Invalid)? { - event.process()?; + let events = events.map_err(PostError::Invalid)?; + // Check the union of all events' fields against the limit and commit + // it under one write lock before processing any event. The fields are + // then reserved, so processing cannot fail on the limit partway through + // (leaving earlier events staged) even if concurrent requests add + // fields meanwhile. + let schemas: Vec<_> = events + .iter() + .map(|event| (event.rb.schema(), event.is_first_event)) + .collect(); + let committed = commit_ingested_schemas( + stream_name, + &stream.metadata, + &schemas, + PARSEABLE.options.dataset_fields_allowed_limit, + )?; + if committed && !stream.get_static_schema_flag() { + stream + .stage_schema_file(stream.get_schema().as_ref().clone()) + .map_err(EventError::from)?; + } + for event in events { + event.process_without_field_limit()?; } } Ok(()) diff --git a/src/handlers/http/query.rs b/src/handlers/http/query.rs index e0b9a0be5..f23582fd1 100644 --- a/src/handlers/http/query.rs +++ b/src/handlers/http/query.rs @@ -60,6 +60,52 @@ use crate::utils::time::{TimeParseError, TimeRange, parse_time_expression}; use crate::utils::{get_tenant_id_from_request, user_auth_for_datasets}; pub const TIME_ELAPSED_HEADER: &str = "p-time-elapsed"; + +/// Maximum number of bins accepted by the manifest-based counts path. +const MAX_COUNTS_BINS: u64 = 10_000; + +/// Mirrors the default bin sizing used by `CountsRequest::get_bounds` when +/// `numBins` is omitted. +fn default_counts_bins(total_minutes: u64) -> u64 { + if total_minutes <= 60 * 5 { + // till 5 hours, 1 bin = 1 min + total_minutes + } else if total_minutes <= 60 * 24 { + // till 1 day, 1 bin = 5 min + total_minutes.div_ceil(5) + } else if total_minutes <= 60 * 24 * 10 { + // till 10 days, 1 bin = 1 hour + total_minutes.div_ceil(60) + } else { + // > 10 days, 1 bin = 1 day + total_minutes.div_ceil(1440) + } +} + +/// Returns an explicit bin count only when the server-derived default would +/// fall outside the accepted range. In-range defaults are left to +/// `get_bounds` so existing bin layouts are unchanged; out-of-range defaults +/// are clamped so callers that omit `numBins` are never rejected for it. +fn clamped_default_counts_bins(total_minutes: u64) -> Option { + let num_bins = default_counts_bins(total_minutes); + (!(1..=MAX_COUNTS_BINS).contains(&num_bins)).then(|| num_bins.clamp(1, MAX_COUNTS_BINS)) +} + +/// Computes a clamped default bin count for a counts request that omitted +/// `numBins`. Unparseable or inverted ranges are left untouched so the +/// existing error handling in `get_bin_density` still applies. +fn clamped_default_counts_bins_for_range(start_time: &str, end_time: &str) -> Option { + let time_range = TimeRange::parse_human_time(start_time, end_time).ok()?; + let total_minutes = time_range + .end + .signed_duration_since(time_range.start) + .num_minutes(); + if total_minutes < 0 { + return None; + } + clamped_default_counts_bins(total_minutes as u64) +} + /// Query Request through http endpoint. #[derive(Debug, Deserialize, Serialize, Clone)] #[serde(rename_all = "camelCase")] @@ -422,11 +468,18 @@ pub async fn get_counts( req: HttpRequest, counts_request: Json, ) -> Result { + let mut body = counts_request.into_inner(); + if body.conditions.is_some() { + // The SQL/grouped counts path derives its bin interval from the time + // range and never reads numBins, so values it ignores are not rejected. + body.num_bins = None; + } else { + body.validate_num_bins()?; + } + let creds = extract_session_key_from_req(&req)?; let permissions = Users.get_permissions(&creds); let tenant_id = get_tenant_id_from_request(&req)?; - let body = counts_request.into_inner(); - // does user have access to table? user_auth_for_datasets(&permissions, std::slice::from_ref(&body.stream), &tenant_id).await?; // Track billing metrics for query calls @@ -539,6 +592,12 @@ pub async fn get_counts( } } + if body.num_bins.is_none() { + // The caller did not send numBins; clamp an out-of-range server + // default instead of rejecting the request for a parameter it omitted. + body.num_bins = clamped_default_counts_bins_for_range(&body.start_time, &body.end_time); + } + let records = body.get_bin_density(&tenant_id).await?; let res = json!({ "fields": vec!["start_time", "end_time", "count"], @@ -780,6 +839,62 @@ mod response_tests { use actix_web::{ResponseError, body::to_bytes, http::header::CONTENT_TYPE}; use arrow_array::{Float64Array, TimestampMillisecondArray}; + #[actix_web::test] + async fn counts_reject_invalid_num_bins_with_bad_request_before_stream_access() { + let app = actix_web::test::init_service( + actix_web::App::new().route("/counts", web::post().to(get_counts)), + ) + .await; + + // Only the manifest-based (unfiltered) path uses numBins, so only it + // rejects out-of-range values. + for num_bins in [0, 10_001, 1_000_000_000, u64::MAX] { + let payload = json!({ + "stream": "logs", + "startTime": "2024-01-01T00:00:00Z", + "endTime": "2024-01-01T01:00:00Z", + "numBins": num_bins, + }); + let request = actix_web::test::TestRequest::post() + .uri("/counts") + .set_json(payload) + .to_request(); + let response = actix_web::test::call_service(&app, request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + response.headers().get(CONTENT_TYPE).unwrap(), + "text/plain; charset=utf-8" + ); + assert_eq!( + actix_web::test::read_body(response).await.as_ref(), + b"CustomError: numBins must be between 1 and 10000 (inclusive)" + ); + } + } + + #[test] + fn counts_default_bins_are_clamped_instead_of_rejected() { + // In-range defaults are left to get_bounds unchanged. + assert_eq!(clamped_default_counts_bins(60), None); + assert_eq!(clamped_default_counts_bins(60 * 24 * 10), None); + assert_eq!(clamped_default_counts_bins(MAX_COUNTS_BINS * 1440), None); + // Out-of-range defaults are clamped into the accepted range. + assert_eq!(clamped_default_counts_bins(0), Some(1)); + assert_eq!( + clamped_default_counts_bins(MAX_COUNTS_BINS * 1440 + 1), + Some(MAX_COUNTS_BINS) + ); + // A ~30 year range without numBins no longer fails on numBins. + assert_eq!( + clamped_default_counts_bins_for_range("2000-01-01T00:00:00Z", "2030-01-01T00:00:00Z"), + Some(MAX_COUNTS_BINS) + ); + assert_eq!( + clamped_default_counts_bins_for_range("2024-01-01T00:00:00Z", "2024-01-01T01:00:00Z"), + None + ); + } + #[test] fn streaming_batch_bytes_preserve_legacy_rows_and_newline() { let batch = RecordBatch::try_from_iter(vec![ diff --git a/src/query/mod.rs b/src/query/mod.rs index 873a52c8e..15f6bc57e 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -90,6 +90,7 @@ type BoxedBatchStream = SendableRecordBatchStream; /// Result type returned by query execution: either collected batches or a streaming adapter, plus field names. type QueryResult = Result<(Either, BoxedBatchStream>, Vec), ExecuteError>; const DEFAULT_COUNTS_TOP_K: usize = 10; +const MAX_COUNT_BINS: u64 = 10_000; pub static SCHEMA_PROVIDER: OnceCell> = OnceCell::new(); @@ -733,7 +734,7 @@ pub struct CountsRequest { pub start_time: String, /// Excluded end time for counts query pub end_time: String, - /// optional number of bins to divide the time range into + /// Optional number of bins to divide the time range into (1 through 10,000). pub num_bins: Option, /// Conditions pub conditions: Option, @@ -743,7 +744,24 @@ fn quote_identifier(value: &str) -> String { format!("\"{}\"", value.replace('"', "\"\"")) } +fn validate_count_bins(num_bins: u64) -> Result<(), QueryError> { + if !(1..=MAX_COUNT_BINS).contains(&num_bins) { + return Err(QueryError::CustomError(format!( + "numBins must be between 1 and {MAX_COUNT_BINS} (inclusive)" + ))); + } + Ok(()) +} + impl CountsRequest { + /// Validate explicit bin counts before stream access or query planning. + pub(crate) fn validate_num_bins(&self) -> Result<(), QueryError> { + if let Some(num_bins) = self.num_bins { + validate_count_bins(num_bins)?; + } + Ok(()) + } + /// This function is supposed to read maninfest files for the given stream, /// get the sum of `num_rows` between the `startTime` and `endTime`, /// divide that by number of bins and return in a manner acceptable for the console @@ -751,16 +769,16 @@ impl CountsRequest { &self, tenant_id: &Option, ) -> Result, QueryError> { + self.validate_num_bins()?; + let time_range = TimeRange::parse_human_time(&self.start_time, &self.end_time)?; + let counts = self.get_bounds(&time_range)?; + let time_partition = PARSEABLE .get_stream(&self.stream, tenant_id) .map_err(|err| anyhow::Error::msg(err.to_string()))? .get_time_partition() .unwrap_or_else(|| DEFAULT_TIMESTAMP_KEY.to_owned()); - // get time range - let time_range = TimeRange::parse_human_time(&self.start_time, &self.end_time)?; let all_manifest_files = get_manifest_list(&self.stream, &time_range, tenant_id).await?; - // get bounds - let counts = self.get_bounds(&time_range)?; // we have start and end times for each bin // we also have all the manifest files for the given time range @@ -828,54 +846,32 @@ impl CountsRequest { } }; - if num_bins == 0 { - return Err(QueryError::CustomError( - "numBins must be greater than 0".to_string(), - )); - } + validate_count_bins(num_bins)?; // divide minutes by num bins to get minutes per bin let quotient = total_minutes / num_bins; - let remainder = total_minutes % num_bins; - let have_remainder = remainder > 0; // now create multiple bounds [startTime, endTime) - // Should we exclude the last one??? - let mut bounds = vec![]; - + let mut bounds = Vec::with_capacity(num_bins as usize); let mut start = time_range.start; - - let loop_end = if have_remainder { - num_bins - } else { - num_bins - 1 - }; - - // Create bins for all but the last date - for _ in 0..loop_end { - let end = start + Duration::minutes(quotient as i64); + for bin in 0..num_bins { + // Keep the existing interval for preceding bins and include any + // remainder in the final bin instead of allocating an extra bin. + let end = if bin == num_bins - 1 { + time_range.end + } else { + start + Duration::minutes(quotient as i64) + }; bounds.push(TimeBounds { start, end }); start = end; } - // Add the last bin, accounting for any remainder, should we include it? - if have_remainder { - bounds.push(TimeBounds { - start, - end: start + Duration::minutes(remainder as i64), - }); - } else { - bounds.push(TimeBounds { - start, - end: start + Duration::minutes(quotient as i64), - }); - } - Ok(bounds) } /// This function will get executed only if self.conditions is some pub async fn get_df_sql(&self, time_column: String) -> Result { + self.validate_num_bins()?; // unwrap because we have asserted that it is some let count_conditions = self.conditions.as_ref().unwrap(); @@ -1301,8 +1297,8 @@ mod tests { use serde_json::json; use super::{ - AbortQueryOnDrop, DataFusionError, RecordBatch, RecordBatchStreamAdapter, - forward_partition_batches, + AbortQueryOnDrop, DataFusionError, DateTime, Duration, MAX_COUNT_BINS, RecordBatch, + RecordBatchStreamAdapter, TimeRange, Utc, forward_partition_batches, }; use crate::query::{ CountConditions, CountsRequest, flatten_objects_for_count, resolve_stream_names, @@ -1375,6 +1371,152 @@ mod tests { assert_eq!(conditions.top_k, Some(3)); } + fn counts_request_for_minutes(num_bins: Option, minutes: i64) -> CountsRequest { + let start = DateTime::parse_from_rfc3339("2024-01-01T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + CountsRequest { + stream: "logs".to_string(), + start_time: start.to_rfc3339(), + end_time: (start + Duration::minutes(minutes)).to_rfc3339(), + num_bins, + conditions: None, + } + } + + #[test] + fn test_counts_bounds_respect_bin_limits_and_cover_range() { + for (num_bins, minutes) in [ + (1, 11), + (1, 0), + (3, 10), + (MAX_COUNT_BINS, 10_000), + (MAX_COUNT_BINS, 10_001), + (MAX_COUNT_BINS, 1), + ] { + let request = counts_request_for_minutes(Some(num_bins), minutes); + let range = + TimeRange::parse_human_time(&request.start_time, &request.end_time).unwrap(); + let bounds = request.get_bounds(&range).unwrap(); + + assert_eq!(bounds.len(), num_bins as usize); + assert!(bounds.len() <= MAX_COUNT_BINS as usize); + assert_eq!(bounds.first().unwrap().start, range.start); + assert_eq!(bounds.last().unwrap().end, range.end); + assert!(bounds.iter().all(|bin| bin.start <= bin.end)); + assert!(bounds.windows(2).all(|pair| pair[0].end == pair[1].start)); + // Earlier boundaries remain unchanged; the last bin absorbs the remainder. + let quotient = minutes / num_bins as i64; + for bin in &bounds[..bounds.len() - 1] { + assert_eq!(bin.end - bin.start, Duration::minutes(quotient)); + } + } + } + + #[test] + fn test_counts_bounds_reject_invalid_num_bins() { + for num_bins in [0, MAX_COUNT_BINS + 1, 1_000_000_000, u64::MAX] { + let request = counts_request_for_minutes(Some(num_bins), 60); + let range = + TimeRange::parse_human_time(&request.start_time, &request.end_time).unwrap(); + let error = request.get_bounds(&range).err().unwrap(); + assert_eq!( + error.to_string(), + "CustomError: numBins must be between 1 and 10000 (inclusive)" + ); + } + } + + #[test] + fn test_counts_default_bounds_preserve_bin_counts() { + for (minutes, expected_bins) in [ + (60, 60), + (300, 300), + (301, 61), + (1440, 288), + (1441, 25), + (14_400, 240), + (14_401, 11), + (10_000 * 1440, 10_000), + ] { + let request = counts_request_for_minutes(None, minutes); + let range = + TimeRange::parse_human_time(&request.start_time, &request.end_time).unwrap(); + let bounds = request.get_bounds(&range).unwrap(); + assert_eq!(bounds.len(), expected_bins); + assert_eq!(bounds.first().unwrap().start, range.start); + assert_eq!(bounds.last().unwrap().end, range.end); + assert!(bounds.windows(2).all(|pair| pair[0].end == pair[1].start)); + } + } + + #[tokio::test] + async fn test_counts_manifest_rejects_oversized_default_before_stream_access() { + let request = counts_request_for_minutes(None, 10_001 * 1440); + let range = TimeRange::parse_human_time(&request.start_time, &request.end_time).unwrap(); + assert!(request.get_bounds(&range).is_err()); + let error = request.get_bin_density(&None).await.unwrap_err(); + assert_eq!( + error.to_string(), + "CustomError: numBins must be between 1 and 10000 (inclusive)" + ); + } + + #[tokio::test] + async fn test_counts_methods_reject_invalid_num_bins_before_stream_access() { + for num_bins in [0, MAX_COUNT_BINS + 1, 1_000_000_000, u64::MAX] { + let mut request = counts_request_for_minutes(Some(num_bins), 60); + // Invalid times prove explicit limits are checked before any other work. + request.start_time = "invalid".to_string(); + request.end_time = "invalid".to_string(); + let error = request.get_bin_density(&None).await.unwrap_err(); + assert_eq!( + error.to_string(), + "CustomError: numBins must be between 1 and 10000 (inclusive)" + ); + + request.conditions = Some(CountConditions { + conditions: None, + group_by: None, + top_k: None, + }); + let error = request + .get_df_sql("p_timestamp".to_string()) + .await + .unwrap_err(); + assert_eq!( + error.to_string(), + "CustomError: numBins must be between 1 and 10000 (inclusive)" + ); + } + } + + #[tokio::test] + async fn test_counts_sql_preserves_default_intervals_for_valid_num_bins() { + for (minutes, interval) in [ + (60, "1m"), + (300, "1m"), + (301, "5m"), + (1440, "5m"), + (1441, "1h"), + (14_399, "1h"), + (14_400, "1d"), + (14_401, "1d"), + ] { + for num_bins in [None, Some(1), Some(MAX_COUNT_BINS)] { + let mut request = counts_request_for_minutes(num_bins, minutes); + request.conditions = Some(CountConditions { + conditions: None, + group_by: None, + top_k: None, + }); + let sql = request.get_df_sql("p_timestamp".to_string()).await.unwrap(); + assert!(sql.contains(&format!("DATE_BIN('{interval}',"))); + assert!(sql.contains(&format!("INTERVAL '{interval}'"))); + } + } + } + #[tokio::test] async fn test_counts_sql_applies_top_k_to_group_by() { let request = CountsRequest {