diff --git a/src/connectors/kafka/config.rs b/src/connectors/kafka/config.rs index 98d9b855e..ffd7aaf87 100644 --- a/src/connectors/kafka/config.rs +++ b/src/connectors/kafka/config.rs @@ -82,6 +82,12 @@ pub struct KafkaConfig { #[derive(Debug, Clone, Args)] pub struct ConsumerConfig { + /// Trusted destination for every consumed topic. Headers never grant tenant access. + #[arg(long = "consumer-tenant-id", env = "P_KAFKA_CONSUMER_TENANT_ID", + help = "Destination tenant for all consumed topics; requires --multi-tenancy. Use a separate consumer group or separate topics per tenant. In multi-tenant mode, mismatching x-p-tenant headers are dropped; when unset, all x-p-tenant headers (including DEFAULT_TENANT) are dropped. Single-tenant mode ignores headers. Leave P_KAFKA_CONSUMER_TENANT_ID unset, not empty", + value_parser = clap::builder::NonEmptyStringValueParser::new())] + pub tenant_id: Option, + #[arg( long = "consumer-topics", env = "P_KAFKA_CONSUMER_TOPICS", @@ -585,6 +591,26 @@ pub struct SecurityConfig { } impl KafkaConfig { + /// Validate before creating a consumer: tenant IDs become object-store paths. + pub fn validate_tenant_mode(&self, multi_tenant: bool) -> anyhow::Result<()> { + if let Some(tenant) = self + .consumer + .as_ref() + .and_then(|consumer| consumer.tenant_id.as_deref()) + { + anyhow::ensure!( + multi_tenant, + "--consumer-tenant-id / P_KAFKA_CONSUMER_TENANT_ID requires --multi-tenancy / P_MULTI_TENANCY" + ); + crate::parseable::validate_tenant_id(tenant).map_err(|error| { + anyhow::anyhow!( + "Invalid --consumer-tenant-id / P_KAFKA_CONSUMER_TENANT_ID: {error}" + ) + })?; + } + Ok(()) + } + pub fn to_rdkafka_consumer_config(&self) -> ClientConfig { let mut config = ClientConfig::new(); @@ -1135,6 +1161,7 @@ impl Default for ProducerConfig { impl Default for ConsumerConfig { fn default() -> Self { Self { + tenant_id: None, topics: vec![], group_id: "parseable-connect-cg".to_string(), buffer_size: 10_000, @@ -1511,4 +1538,67 @@ mod tests { assert_eq!(config.get("sasl.oauthbearer.client.id"), Some("unused")); assert_eq!(config.get("sasl.oauthbearer.client.secret"), Some("unused")); } + + #[test] + fn empty_consumer_tenant_id_is_rejected_by_cli() { + #[derive(Debug, Parser)] + struct ConsumerCli { + #[command(flatten)] + consumer: ConsumerConfig, + } + + // A valid value proves this CLI fixture can otherwise parse the option. + let cli = ConsumerCli::try_parse_from(["test", "--consumer-tenant-id=tenant-a"]) + .expect("a nonempty tenant ID should parse"); + assert_eq!(cli.consumer.tenant_id.as_deref(), Some("tenant-a")); + let error = ConsumerCli::try_parse_from(["test", "--consumer-tenant-id="]) + .expect_err("an explicitly empty tenant ID must fail parsing"); + assert_eq!(error.kind(), clap::error::ErrorKind::InvalidValue); + assert!(error.to_string().contains("--consumer-tenant-id")); + } +} + +#[cfg(test)] +mod tenant_env_tests { + use super::*; + + #[test] + fn empty_consumer_tenant_env_is_rejected_in_child_process() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--ignored", + "--exact", + "connectors::kafka::config::tenant_env_tests::tenant_env_child", + "--nocapture", + ]) + .env("P_KAFKA_CONSUMER_TENANT_ID", "") + .env("PARSEABLE_TENANT_ENV_TEST_CHILD", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert!(String::from_utf8_lossy(&output.stdout).contains("1 passed")); + } + + #[test] + #[ignore = "run only as a child of empty_consumer_tenant_env_is_rejected_in_child_process"] + fn tenant_env_child() { + if std::env::var_os("PARSEABLE_TENANT_ENV_TEST_CHILD").is_none() { + return; + } + #[derive(Debug, clap::Parser)] + struct ConsumerCli { + #[command(flatten)] + consumer: ConsumerConfig, + } + use clap::Parser; + let error = ConsumerCli::try_parse_from(["test"]) + .expect_err("empty tenant env must fail startup parsing"); + assert_eq!(error.kind(), clap::error::ErrorKind::InvalidValue); + assert!(error.to_string().contains("--consumer-tenant-id")); + } } diff --git a/src/connectors/kafka/consumer.rs b/src/connectors/kafka/consumer.rs index 9b4deb544..97baaf266 100644 --- a/src/connectors/kafka/consumer.rs +++ b/src/connectors/kafka/consumer.rs @@ -24,9 +24,9 @@ use crate::connectors::kafka::{ ConsumerRecord, KafkaContext, StreamConsumer, TopicPartition, partition_stream, }; use futures_util::FutureExt; +use rdkafka::Message; use rdkafka::Statistics; use rdkafka::consumer::Consumer; -use rdkafka::message::BorrowedMessage; use std::sync::Arc; use std::time::Duration; use tokio::sync::{RwLock, mpsc}; @@ -38,6 +38,7 @@ pub struct KafkaStreams { stream_state: Arc>, statistics: Arc>, shutdown_handle: Shutdown, + multi_tenant: bool, } impl KafkaStreams { @@ -47,6 +48,20 @@ impl KafkaStreams { shutdown_handle: Shutdown, ) -> anyhow::Result { info!("Initializing KafkaStreams..."); + let multi_tenant = crate::parseable::PARSEABLE.options.is_multi_tenant(); + context.config.validate_tenant_mode(multi_tenant)?; + if multi_tenant + && context + .config + .consumer + .as_ref() + .and_then(|config| config.tenant_id.as_ref()) + .is_none() + { + warn!( + "Multi-tenant Kafka consumer has no destination tenant; set --consumer-tenant-id / P_KAFKA_CONSUMER_TENANT_ID. Headerless records use the default destination; records carrying tenant headers are dropped" + ); + } let statistics = Arc::clone(&context.statistics); let consumer = KafkaStreams::create_consumer(context); info!("KafkaStreams initialized successfully."); @@ -56,6 +71,7 @@ impl KafkaStreams { stream_state, statistics, shutdown_handle, + multi_tenant, }) } @@ -102,6 +118,13 @@ impl KafkaStreams { pub fn partitioned(&self) -> ReceiverStream { let (stream_tx, stream_rx) = mpsc::channel(100); let consumer = self.consumer(); + let multi_tenant = self.multi_tenant; + let configured_tenant = consumer + .context() + .config + .consumer + .as_ref() + .and_then(|config| config.tenant_id.clone()); let stream_state = self.state(); let tokio_handle = tokio::runtime::Handle::current(); let shutdown_handle = self.shutdown_handle.clone(); @@ -112,7 +135,7 @@ impl KafkaStreams { let result: Result<(), ConnectorError> = tokio::select! { result = consumer.recv() => { match result { - Ok(msg) => KafkaStreams::handle_message(msg, &stream_state, &stream_tx).await.map_err(Into::into), + Ok(msg) => KafkaStreams::handle_message(msg, &stream_state, &stream_tx, configured_tenant.as_deref(), multi_tenant).await.map_err(Into::into), Err(e) => Err(e.into()) } } @@ -152,14 +175,16 @@ impl KafkaStreams { } /// Handle individual Kafka message and route it to the proper partition stream - async fn handle_message( - msg: BorrowedMessage<'_>, + async fn handle_message( + msg: M, stream_state: &RwLock, stream_tx: &mpsc::Sender, + configured_tenant: Option<&str>, + multi_tenant: bool, ) -> anyhow::Result<()> { + let tp = TopicPartition::new(msg.topic().to_owned(), msg.partition()); + let consumer_record = ConsumerRecord::from_message(msg, configured_tenant, multi_tenant); let mut state = stream_state.write().await; - let tp = TopicPartition::from_kafka_msg(&msg); - let consumer_record = ConsumerRecord::from_borrowed_msg(msg); let partition_stream_tx = KafkaStreams::get_or_create_partition_stream(&mut state, stream_tx, tp).await; @@ -244,3 +269,52 @@ impl KafkaStreams { }; } } + +#[cfg(test)] +mod routing_tests { + use super::*; + use futures_util::StreamExt; + + #[tokio::test] + async fn rejected_record_is_delivered_to_partition_stream_receiver() { + let state = RwLock::new(StreamState::new(1)); + let (tx, mut rx) = mpsc::channel(1); + KafkaStreams::handle_message( + crate::connectors::kafka::message_conversion_tests::message("tenant-b"), + &state, + &tx, + Some("tenant-a"), + true, + ) + .await + .unwrap(); + let partition = tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .expect("rejected record must create a partition stream") + .unwrap(); + assert_eq!( + partition.topic_partition(), + &TopicPartition::new("trusted-topic".into(), 3) + ); + partition + .run_drain(|mut records| async move { + let record = tokio::time::timeout(Duration::from_secs(1), records.next()) + .await + .expect("marker must be routed to the partition receiver") + .unwrap(); + assert_eq!( + record, + ConsumerRecord::new( + None, + None, + "trusted-topic".into(), + 3, + 42, + Some(123456), + Some("tenant-a".into()), + ) + ); + }) + .await; + } +} diff --git a/src/connectors/kafka/metrics.rs b/src/connectors/kafka/metrics.rs index 4ea509908..5f85544c6 100644 --- a/src/connectors/kafka/metrics.rs +++ b/src/connectors/kafka/metrics.rs @@ -24,6 +24,16 @@ use prometheus::{ use rdkafka::Statistics; use std::sync::{Arc, RwLock}; +/// Rejected records, without attacker-controlled labels. +pub static TENANT_MISMATCH_DROPPED: once_cell::sync::Lazy = + once_cell::sync::Lazy::new(|| { + prometheus::IntCounter::new( + "parseable_kafka_tenant_mismatch_dropped_total", + "Kafka records dropped because a tenant header disagrees with trusted configuration", + ) + .expect("valid Kafka tenant counter") + }); + #[derive(Debug)] pub struct KafkaMetricsCollector { stats: Arc>, diff --git a/src/connectors/kafka/mod.rs b/src/connectors/kafka/mod.rs index 59404dd6d..1d9dd8026 100644 --- a/src/connectors/kafka/mod.rs +++ b/src/connectors/kafka/mod.rs @@ -24,12 +24,11 @@ use aws_types::region::Region; use rdkafka::client::OAuthToken; use rdkafka::consumer::{ConsumerContext, Rebalance}; use rdkafka::error::KafkaResult; -use rdkafka::message::{BorrowedMessage, Headers}; +use rdkafka::message::BorrowedMessage; use rdkafka::producer::ProducerContext; use rdkafka::topic_partition_list::TopicPartitionListElem; use rdkafka::{ClientContext, Message, Offset, Statistics}; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; use std::error::Error; use std::sync::{Arc, OnceLock, RwLock}; use std::time::Duration; @@ -170,7 +169,8 @@ pub struct ConsumerRecord { pub partition: i32, pub offset: i64, pub timestamp: Option, - pub tenant_id: Option, // pub headers: Option>>, + /// Destination resolved exclusively from trusted consumer configuration. + pub tenant_id: Option, } impl ConsumerRecord { @@ -194,23 +194,35 @@ impl ConsumerRecord { } } - pub fn from_borrowed_msg(msg: BorrowedMessage) -> Self { - let tenant_id = if let Some(headers) = extract_headers(&msg) - && let Some(tenant_id) = headers.get(TENANT_ID) - { - tenant_id.clone() - } else { - None - }; + pub fn from_message( + msg: M, + configured_tenant: Option<&str>, + multi_tenant: bool, + ) -> Self { + let tenant_id = tenant_for_ingestion( + multi_tenant, + configured_tenant, + tenant_header_values(msg.headers()), + &metrics::TENANT_MISMATCH_DROPPED, + ); + + if tenant_id.is_none() { + warn_tenant_mismatch(msg.topic(), msg.partition(), configured_tenant.is_some()); + } + // Rejected records retain their position in the partition queue, but + // never retain producer payloads or keys. Tombstones use the same shape. + let accepted = tenant_id.is_some(); Self { - key: msg.key().map(|k| k.to_vec()), - payload: msg.payload().map(|p| p.to_vec()), + key: accepted.then(|| msg.key().map(|k| k.to_vec())).flatten(), + payload: accepted + .then(|| msg.payload().map(|p| p.to_vec())) + .flatten(), topic: msg.topic().to_owned(), partition: msg.partition(), offset: msg.offset(), timestamp: msg.timestamp().to_millis(), - tenant_id, // headers: extract_headers(&msg), + tenant_id: tenant_id.unwrap_or_else(|| configured_tenant.map(str::to_owned)), } } @@ -232,19 +244,93 @@ impl ConsumerRecord { } } -#[allow(unused)] -fn extract_headers(msg: &BorrowedMessage<'_>) -> Option>> { - msg.headers().map(|headers| { - headers - .iter() - .map(|header| { - ( - header.key.to_string(), - header.value.map(|v| String::from_utf8_lossy(v).to_string()), - ) - }) - .collect() - }) +/// Extract only exact tenant-header keys, preserving null values for rejection. +fn tenant_header_values( + headers: Option<&H>, +) -> impl Iterator> + '_ { + headers + .into_iter() + .flat_map(|headers| (0..headers.count()).map(move |index| headers.get(index))) + .filter(|header| header.key == TENANT_ID) + .map(|header| header.value) +} + +/// Headers are assertions only: every supplied value must match the trusted +/// consumer tenant. Reject null, invalid UTF-8, and conflicting duplicate values. +/// Single-tenant mode ignores all headers. In multi-tenant mode, without +/// configuration only headerless records use the existing default destination. +fn resolve_tenant<'a>( + multi_tenant: bool, + configured: Option<&str>, + mut headers: impl Iterator>, +) -> Result, ()> { + if !multi_tenant { + return Ok(None); + } + if headers.any(|value| match (configured, value) { + (Some(tenant), Some(value)) => tenant.as_bytes() != value, + _ => true, + }) { + return Err(()); + } + Ok(configured.map(str::to_owned)) +} + +fn tenant_for_ingestion<'a>( + multi_tenant: bool, + configured: Option<&str>, + headers: impl Iterator>, + dropped: &prometheus::IntCounter, +) -> Option> { + match resolve_tenant(multi_tenant, configured, headers) { + Ok(tenant) => Some(tenant), + Err(()) => { + dropped.inc(); + None + } + } +} + +/// A monotonic, process-wide warning window. The counter remains unlabeled; +/// each warning identifies the topic/partition of the drop that opens a window. +#[derive(Default)] +struct TenantDropWarnings { + last_warning: Option, + dropped: u64, +} + +impl TenantDropWarnings { + fn record_drop(&mut self, now: std::time::Instant) -> Option { + self.dropped = self.dropped.saturating_add(1); + if self + .last_warning + .is_some_and(|last| now.duration_since(last) < Duration::from_secs(60)) + { + return None; + } + self.last_warning = Some(now); + Some(std::mem::take(&mut self.dropped)) + } +} + +fn warn_tenant_mismatch(topic: &str, partition: i32, tenant_configured: bool) { + static WARNINGS: std::sync::Mutex = + std::sync::Mutex::new(TenantDropWarnings { + last_warning: None, + dropped: 0, + }); + let mut warnings = WARNINGS + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(dropped) = warnings.record_drop(std::time::Instant::now()) { + warn!( + topic, + partition, + dropped, + tenant_configured, + "Dropping Kafka records with rejected x-p-tenant headers; drops since previous warning across all partitions; see parseable_kafka_tenant_mismatch_dropped_total" + ); + } } impl ConsumerContext for KafkaContext { @@ -401,3 +487,349 @@ impl ClientContext for KafkaContext { }) } } + +#[cfg(test)] +mod tenant_tests { + use super::*; + + #[test] + fn single_tenant_headers_are_ignored_without_drops() { + let counter = + prometheus::IntCounter::new("test_ignored_tenant_drops", "test drops").unwrap(); + for header in [Some(b"tenant-b".as_slice()), None, Some(b"\xff".as_slice())] { + assert_eq!( + tenant_for_ingestion(false, None, [header].into_iter(), &counter), + Some(None) + ); + } + assert_eq!(counter.get(), 0); + } + + #[test] + fn multi_tenant_without_config_or_header_keeps_base_default() { + assert_eq!(resolve_tenant(true, None, std::iter::empty()), Ok(None)); + } + + fn config_with_tenant(tenant: Option<&str>) -> KafkaConfig { + let mut config = KafkaConfig::default(); + config.consumer.as_mut().unwrap().tenant_id = tenant.map(str::to_owned); + config + } + + #[test] + fn tenant_config_requires_multi_tenancy() { + let error = config_with_tenant(Some("tenant-a")) + .validate_tenant_mode(false) + .unwrap_err(); + assert!(error.to_string().contains("requires --multi-tenancy")); + } + + #[test] + fn unsafe_configured_tenant_ids_are_rejected() { + for tenant in [ + "", + ".", + "..", + "a/b", + "a\\b", + " a", + "a b", + "a\t", + "a\n", + "a..b", + "-tenant", + "tenant!", + "abcdefghijklmnopqrstuvwxyz01234567890", + ] { + assert!( + config_with_tenant(Some(tenant)) + .validate_tenant_mode(true) + .is_err(), + "accepted {tenant:?}" + ); + } + } + + #[test] + fn configured_tenant_uses_server_validation_including_reserved_default() { + let tenant = crate::parseable::DEFAULT_TENANT; + assert!( + config_with_tenant(Some(tenant)) + .validate_tenant_mode(true) + .is_err() + ); + } + + #[test] + fn valid_and_unconfigured_modes_are_allowed() { + assert!( + config_with_tenant(Some("tenant-a")) + .validate_tenant_mode(true) + .is_ok() + ); + assert!(config_with_tenant(None).validate_tenant_mode(false).is_ok()); + assert!(config_with_tenant(None).validate_tenant_mode(true).is_ok()); + } + + #[test] + fn mismatch_is_dropped_and_counted() { + let counter = prometheus::IntCounter::new("test_tenant_drops", "test drops").unwrap(); + let record = tenant_for_ingestion( + true, + Some("tenant-a"), + [Some(b"tenant-b".as_slice())].into_iter(), + &counter, + ); + assert_eq!(record, None); + assert_eq!(counter.get(), 1); + } + + #[test] + fn no_config_no_header_keeps_single_tenant_default() { + assert_eq!(resolve_tenant(false, None, std::iter::empty()), Ok(None)); + } + + #[test] + fn configured_tenant_without_header_is_used() { + assert_eq!( + resolve_tenant(true, Some("tenant-a"), std::iter::empty()), + Ok(Some("tenant-a".into())) + ); + } + + #[test] + fn matching_header_is_accepted() { + assert_eq!( + resolve_tenant( + true, + Some("tenant-a"), + [Some(b"tenant-a".as_slice())].into_iter() + ), + Ok(Some("tenant-a".into())) + ); + } + + #[test] + fn unconfigured_header_cannot_select_tenant() { + assert_eq!( + resolve_tenant(true, None, [Some(b"tenant-b".as_slice())].into_iter()), + Err(()) + ); + } + + #[test] + fn duplicate_headers_cannot_hide_mismatch() { + assert_eq!( + resolve_tenant( + true, + Some("tenant-a"), + [Some(b"tenant-b".as_slice()), Some(b"tenant-a".as_slice())].into_iter() + ), + Err(()) + ); + } + + #[test] + fn null_and_invalid_headers_are_rejected() { + for header in [None, Some(b"\xff".as_slice())] { + assert_eq!( + resolve_tenant(true, Some("tenant-a"), [header].into_iter()), + Err(()) + ); + } + } + + #[test] + fn owned_unrelated_headers_do_not_hide_tenant_mismatch() { + use rdkafka::message::{Header, OwnedHeaders}; + + let headers = OwnedHeaders::new() + .insert(Header { + key: "trace-id", + value: Some("trace"), + }) + .insert(Header { + key: "x-p-tenant", + value: Some("tenant-b"), + }) + .insert(Header { + key: "content-type", + value: Some("application/json"), + }); + let counter = prometheus::IntCounter::new("test_header_drops", "test drops").unwrap(); + assert_eq!( + tenant_for_ingestion( + true, + Some("tenant-a"), + tenant_header_values(Some(&headers)), + &counter + ), + None + ); + assert_eq!(counter.get(), 1); + } + + #[test] + fn owned_case_variant_key_does_not_select_tenant() { + use rdkafka::message::{Header, OwnedHeaders}; + + let headers = OwnedHeaders::new().insert(Header { + key: "X-P-Tenant", + value: Some("tenant-b"), + }); + let counter = prometheus::IntCounter::new("test_header_drops", "test drops").unwrap(); + assert_eq!( + tenant_for_ingestion( + true, + Some("tenant-a"), + tenant_header_values(Some(&headers)), + &counter + ), + Some(Some("tenant-a".into())) + ); + assert_eq!(counter.get(), 0); + } + + #[test] + fn owned_matching_tenant_header_is_accepted() { + use rdkafka::message::{Header, OwnedHeaders}; + + let headers = OwnedHeaders::new().insert(Header { + key: "x-p-tenant", + value: Some("tenant-a"), + }); + let counter = prometheus::IntCounter::new("test_header_drops", "test drops").unwrap(); + assert_eq!( + tenant_for_ingestion( + true, + Some("tenant-a"), + tenant_header_values(Some(&headers)), + &counter + ), + Some(Some("tenant-a".into())) + ); + assert_eq!(counter.get(), 0); + } + + #[test] + fn owned_null_tenant_header_is_dropped() { + use rdkafka::message::{Header, OwnedHeaders}; + + let headers = OwnedHeaders::new().insert(Header { + key: "x-p-tenant", + value: None::<&str>, + }); + let counter = prometheus::IntCounter::new("test_header_drops", "test drops").unwrap(); + assert_eq!( + tenant_for_ingestion( + true, + Some("tenant-a"), + tenant_header_values(Some(&headers)), + &counter + ), + None + ); + assert_eq!(counter.get(), 1); + } + + #[test] + fn absent_headers_preserve_configured_or_default_destination() { + use rdkafka::message::OwnedHeaders; + + let counter = + prometheus::IntCounter::new("test_absent_header_drops", "test drops").unwrap(); + for configured in [Some("tenant-a"), None] { + assert_eq!( + tenant_for_ingestion( + true, + configured, + tenant_header_values(None::<&OwnedHeaders>), + &counter + ), + Some(configured.map(str::to_owned)) + ); + } + assert_eq!(counter.get(), 0); + } +} + +#[cfg(test)] +mod tenant_warning_tests { + use super::*; + + #[test] + fn tenant_drop_warnings_repeat_after_sixty_seconds_with_interval_count() { + let now = std::time::Instant::now(); + let mut warnings = TenantDropWarnings::default(); + assert_eq!(warnings.record_drop(now), Some(1)); + assert_eq!(warnings.record_drop(now), None); + assert_eq!(warnings.record_drop(now + Duration::from_secs(59)), None); + assert_eq!(warnings.record_drop(now + Duration::from_secs(60)), Some(3)); + assert_eq!(warnings.record_drop(now + Duration::from_secs(119)), None); + assert_eq!( + warnings.record_drop(now + Duration::from_secs(120)), + Some(2) + ); + } +} + +#[cfg(test)] +mod message_conversion_tests { + use super::*; + use rdkafka::message::{Header, OwnedHeaders, OwnedMessage, Timestamp}; + + pub(super) fn message(tenant: &str) -> OwnedMessage { + OwnedMessage::new( + Some(br#"{"value":1}"#.to_vec()), + Some(b"record-key".to_vec()), + "trusted-topic".into(), + Timestamp::CreateTime(123456), + 3, + 42, + Some(OwnedHeaders::new().insert(Header { + key: TENANT_ID, + value: Some(tenant), + })), + ) + } + + #[test] + fn rejected_message_becomes_marker_with_trusted_tenant_and_broker_metadata() { + let record = ConsumerRecord::from_message(message("tenant-b"), Some("tenant-a"), true); + assert_eq!( + record, + ConsumerRecord::new( + None, + None, + "trusted-topic".into(), + 3, + 42, + Some(123456), + Some("tenant-a".into()), + ) + ); + let tpl = record.offset_to_commit().unwrap(); + assert_eq!( + tpl.find_partition("trusted-topic", 3).unwrap().offset(), + Offset::Offset(43) + ); + } + + #[test] + fn accepted_message_preserves_payload_and_key() { + let msg = message("tenant-a"); + let record = ConsumerRecord::from_message(msg, Some("tenant-a"), true); + assert_eq!( + record, + ConsumerRecord::new( + Some(br#"{"value":1}"#.to_vec()), + Some(b"record-key".to_vec()), + "trusted-topic".into(), + 3, + 42, + Some(123456), + Some("tenant-a".into()), + ) + ); + } +} diff --git a/src/connectors/kafka/processor.rs b/src/connectors/kafka/processor.rs index 76a65c866..7b16209fa 100644 --- a/src/connectors/kafka/processor.rs +++ b/src/connectors/kafka/processor.rs @@ -277,7 +277,6 @@ fn build_events<'a>( for record in records { let _span = record_span(record).entered(); let Some(payload) = record.payload.as_ref() else { - warn!("Skipping Kafka record without a payload"); continue; }; total_payload_size += payload.len() as u64; @@ -340,6 +339,11 @@ impl Processor, ()> for ParseableSinkProcessor { let len = records.len(); debug!("Processing {len} records"); + // Offset-only markers and tombstones need no stream or metadata access. + if !chunk_needs_stream(&records) { + return Ok(()); + } + let stream_name = records.first().map(|r| r.topic.clone()).unwrap_or_default(); let tenant_id = records.first().and_then(|r| r.tenant_id.clone()); @@ -397,30 +401,15 @@ where self.buffer_config.buffer_timeout, ); - chunked_stream - .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 {:?}; offsets up to {:?} will still be committed and these records are dropped: {:?}", - record_count, tp, tpl, e - ); - } - - //CommitMode::Async race condition. - //@see https://github.com/confluentinc/librdkafka/issues/4534 - //@see https://github.com/confluentinc/librdkafka/issues/4059 - if let Err(e) = self.consumer.commit(&tpl, CommitMode::Sync) { - error!(error = %e, "Failed to commit offsets for {:?}", tpl); - } else { - debug!("Committed offsets for {:?}", tpl); - } - } - }) - .await; + process_partition_chunks(self.processor.as_ref(), tp, chunked_stream, |tpl| { + // Synchronous commits avoid librdkafka's asynchronous commit race. + if let Err(e) = self.consumer.commit(tpl, CommitMode::Sync) { + error!(error = %e, "Failed to commit offsets for {:?}", tpl); + } else { + debug!("Committed offsets for {:?}", tpl); + } + }) + .await; self.processor.post_stream().await?; @@ -428,6 +417,48 @@ where } } +/// Eight chunks allow useful processing overlap while bounding memory and CPU +/// work per partition. Ordered buffering holds commits behind earlier chunks. +const MAX_IN_FLIGHT_CHUNKS_PER_PARTITION: usize = 8; + +fn chunk_needs_stream(records: &[ConsumerRecord]) -> bool { + records.iter().any(|record| record.payload.is_some()) +} + +/// Process chunks concurrently, yielding commit offsets in partition input order. +async fn process_partition_chunks( + processor: &P, + tp: TopicPartition, + chunks: S, + mut commit: impl FnMut(&rdkafka::TopicPartitionList), +) where + P: Processor, ()>, + S: futures_util::Stream>, +{ + let completed = chunks.map(|records| { + let tp = &tp; + async move { + let 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) = processor.process(records).await { + // Preserve the connector's existing drop-on-processing-error policy. + error!( + "Failed to process {} records for {:?}; offsets up to {:?} will still be committed and these records are dropped: {:?}", + record_count, tp, tpl, e + ); + } + Some(tpl) + } + }).buffered(MAX_IN_FLIGHT_CHUNKS_PER_PARTITION); + futures_util::pin_mut!(completed); + while let Some(tpl) = completed.next().await { + if let Some(tpl) = tpl { + commit(&tpl); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -711,10 +742,13 @@ mod tests { tracing::subscriber::with_default(subscriber, check); let output = String::from_utf8(output.lock().unwrap().clone()).unwrap(); let warnings: Vec<_> = output.lines().collect(); - assert_eq!(warnings.len(), 12, "{output}"); + assert_eq!(warnings.len(), 11, "{output}"); + // Offset markers and tombstones are intentionally silent. + assert!(!output.contains("offset=107"), "{output}"); + assert!(!output.contains("without a payload"), "{output}"); for (warning, offset) in warnings .iter() - .zip([100, 101, 102, 103, 106, 107, 105, 200, 201, 202, 203, 204]) + .zip([100, 101, 102, 103, 106, 105, 200, 201, 202, 203, 204]) { for field in [ "topic=test".to_owned(), @@ -990,3 +1024,137 @@ mod tests { } } } + +#[cfg(test)] +mod offset_marker_tests { + use super::*; + use std::sync::Mutex; + + fn marker(offset: i64) -> ConsumerRecord { + ConsumerRecord::new(None, None, "rejected".into(), 2, offset, Some(123), None) + } + + fn committed_offset(tpl: &rdkafka::TopicPartitionList) -> i64 { + assert_eq!(tpl.count(), 1); + let entry = tpl.find_partition("rejected", 2).unwrap(); + match entry.offset() { + rdkafka::Offset::Offset(offset) => offset, + other => panic!("unexpected offset: {other:?}"), + } + } + + #[tokio::test] + async fn all_rejected_partition_commits_next_offset() { + let mut commits = Vec::new(); + process_partition_chunks( + &GatedProcessor { + release: Mutex::new(None), + processed: Mutex::new(Vec::new()), + }, + TopicPartition::new("rejected".into(), 2), + futures_util::stream::iter([vec![marker(10), marker(11)], vec![marker(12)]]), + |tpl| commits.push(committed_offset(tpl)), + ) + .await; + assert_eq!(commits, [12, 13]); + } + + struct GatedProcessor { + release: Mutex>>, + processed: Mutex>, + } + + #[async_trait] + impl Processor, ()> for GatedProcessor { + async fn process(&self, records: Vec) -> anyhow::Result<()> { + let release = self.release.lock().unwrap().take(); + if let Some(release) = release { + release.await?; + } + self.processed.lock().unwrap().extend( + records + .iter() + .filter(|r| r.payload.is_some()) + .map(|r| r.offset), + ); + Ok(()) + } + } + + #[tokio::test] + async fn rejected_tail_waits_for_accepted_records_before_commit() { + let (release, wait) = tokio::sync::oneshot::channel(); + let processor = GatedProcessor { + release: Mutex::new(Some(wait)), + processed: Mutex::new(Vec::new()), + }; + let mut accepted = marker(10); + accepted.payload = Some(br#"{"value":1}"#.to_vec()); + let commits = Mutex::new(Vec::new()); + let processing = process_partition_chunks( + &processor, + TopicPartition::new("rejected".into(), 2), + futures_util::stream::iter([vec![accepted, marker(11)], vec![marker(12), marker(13)]]), + |tpl| { + assert_eq!(*processor.processed.lock().unwrap(), [10]); + commits.lock().unwrap().push(committed_offset(tpl)); + }, + ); + tokio::pin!(processing); + assert!(futures_util::poll!(&mut processing).is_pending()); + assert!(commits.lock().unwrap().is_empty()); + release.send(()).unwrap(); + processing.await; + assert_eq!(*commits.lock().unwrap(), [12, 14]); + } + + #[test] + fn chunk_needs_stream_only_for_payload_records() { + assert!(!chunk_needs_stream(&[])); + assert!(!chunk_needs_stream(&[marker(1), marker(2)])); + let mut accepted = marker(3); + accepted.payload = Some(b"{}".to_vec()); + assert!(chunk_needs_stream(&[marker(2), accepted])); + } + + struct OverlapProcessor { + release: Mutex>>, + started: Mutex>, + } + + #[async_trait] + impl Processor, ()> for OverlapProcessor { + async fn process(&self, records: Vec) -> anyhow::Result<()> { + let offset = records[0].offset; + self.started.lock().unwrap().push(offset); + if offset == 10 { + let wait = self.release.lock().unwrap().take().unwrap(); + wait.await?; + } + Ok(()) + } + } + + #[tokio::test] + async fn later_chunk_starts_while_earlier_chunk_blocks_commits() { + let (release, wait) = tokio::sync::oneshot::channel(); + let processor = OverlapProcessor { + release: Mutex::new(Some(wait)), + started: Mutex::new(Vec::new()), + }; + let commits = Mutex::new(Vec::new()); + let processing = process_partition_chunks( + &processor, + TopicPartition::new("rejected".into(), 2), + futures_util::stream::iter([vec![marker(10)], vec![marker(11)]]), + |tpl| commits.lock().unwrap().push(committed_offset(tpl)), + ); + tokio::pin!(processing); + assert!(futures_util::poll!(&mut processing).is_pending()); + assert_eq!(*processor.started.lock().unwrap(), [10, 11]); + assert!(commits.lock().unwrap().is_empty()); + release.send(()).unwrap(); + processing.await; + assert_eq!(*commits.lock().unwrap(), [11, 12]); + } +} diff --git a/src/connectors/mod.rs b/src/connectors/mod.rs index 094039420..40a5016c7 100644 --- a/src/connectors/mod.rs +++ b/src/connectors/mod.rs @@ -94,6 +94,7 @@ where let kafka_streams = KafkaStreams::init(kafka_context, stream_state, shutdown_handle.clone())?; + registry.register(Box::new(kafka::metrics::TENANT_MISMATCH_DROPPED.clone()))?; let stats = kafka_streams.statistics(); registry.register(Box::new(KafkaMetricsCollector::new(stats)?))?;