From f471c037a0c107073b8b8bddc50dbecd2dacc69a Mon Sep 17 00:00:00 2001 From: Daniel Nelson Date: Wed, 12 Aug 2026 19:01:51 -0700 Subject: [PATCH 1/2] dekaf: don't close connection immediately on invalid schema Collection schemas are refreshed every spec_ttl. Currently if a document does not validate against its collection schema the session would log an error and close the connection. This would interrupt all partitions on this connection and the client would need to reconnect. If the error is corrected by the inferred schema you it should be updated on the next spec_ttl refresh, but if the client reconnects before then it would result in further disconnections. With this change, when a partition fails to validate against the schema, it is placed into "cooldown". LEADER_NOT_AVAILABLE will be reported for this partition. This will indicate to the client that it should retry, but will not affect the ability for other partitions to make progress. The partition remains in cooldown until the collection schema is updated where on the next fetch a new Read is created, or until a timeout (4m default, 2x the spec_ttl) where the connection will be closed with an error like before. --- crates/dekaf/src/lib.rs | 2 +- crates/dekaf/src/main.rs | 6 + crates/dekaf/src/session.rs | 297 +++++++++++++++++++++++++---------- crates/dekaf/src/topology.rs | 13 ++ 4 files changed, 235 insertions(+), 83 deletions(-) diff --git a/crates/dekaf/src/lib.rs b/crates/dekaf/src/lib.rs index 4e8a02aa39b..e75f6b56382 100644 --- a/crates/dekaf/src/lib.rs +++ b/crates/dekaf/src/lib.rs @@ -346,7 +346,7 @@ impl App { /// Dispatch a read request `frame` of the current session, writing its response into `out`. /// `raw_sasl_auth` is the state of SASL "raw" mode authentication, /// and conditions the interpretation of request frames. -#[tracing::instrument(level = "trace", err(level = "warn"), skip_all)] +#[tracing::instrument(level = "trace", skip_all)] pub async fn dispatch_request_frame( session: &mut Session, raw_sasl_auth: &mut bool, diff --git a/crates/dekaf/src/main.rs b/crates/dekaf/src/main.rs index 57f893ec3b6..efb7facd26c 100644 --- a/crates/dekaf/src/main.rs +++ b/crates/dekaf/src/main.rs @@ -115,6 +115,11 @@ pub struct Cli { #[arg(long, env = "SPEC_TTL", value_parser = humantime::parse_duration, default_value = "2m")] spec_ttl: std::time::Duration, + /// How long a connection will keep returning LeaderNotAvailable for a partition + /// stuck failing schema validation before giving up and closing the connection. + #[arg(long, env = "SCHEMA_ERROR_COOLDOWN_TIMEOUT", value_parser = humantime::parse_duration, default_value = "4m")] + schema_error_cooldown_timeout: std::time::Duration, + /// Timeout for TLS handshake completion #[arg(long, env = "TLS_HANDSHAKE_TIMEOUT", value_parser = humantime::parse_duration, default_value = "10s")] tls_handshake_timeout: std::time::Duration, @@ -434,6 +439,7 @@ async fn async_main(cli: Cli) -> anyhow::Result<()> { upstream_auth.clone(), cli.read_buffer_chunk_limit, cli.combined_partition_fetch_limit, + cli.schema_error_cooldown_timeout, ), socket, tls_acceptor.clone(), diff --git a/crates/dekaf/src/session.rs b/crates/dekaf/src/session.rs index 76441cae7bd..474085ede4e 100644 --- a/crates/dekaf/src/session.rs +++ b/crates/dekaf/src/session.rs @@ -32,7 +32,32 @@ struct PendingRead { // Last time this read was included in a Fetch request. Used to reap reads for // partitions the client has stopped fetching. last_accessed: std::time::Instant, - handle: tokio_util::task::AbortOnDropHandle>, + // Identifies the journal, used to look for entries in cooldown. + journal_name: String, + // The Collection::schema_hash this read was started against, used to + // detect if the read is still current. + schema_hash: String, + // None while the partition is cooling down after a schema error. + handle: Option>>, +} + +/// Tracks a journal that's been cooling down since a schema-validation +/// error, so repeated Fetch requests don't pay for rebuilding a Read against +/// a binding that's known to still be broken. +struct CooldownEntry { + // When this journal first started failing schema validation. + first_failed_at: std::time::Instant, + // The Collection::schema_hash in effect when the failure occurred. + schema_hash: String, +} + +fn is_schema_validation_error(err: &anyhow::Error) -> bool { + err.chain().any(|e| { + matches!( + e.downcast_ref::(), + Some(avro::Error::NotMatched { .. } | avro::Error::ParseFloat(..)) + ) + }) } /// Resolve a current, non-regressing high watermark for a Fetch response. @@ -100,6 +125,10 @@ pub struct Session { read_buffer_size: usize, // Total byte budget for a Fetch, divided evenly across partitions to cap per-partition reads. combined_partition_fetch_limit: usize, + // Journals currently cooling down after a schema-validation error, keyed by journal name. + cooldown: HashMap, + // How long a journal can stay in cooldown before this connection gives up and closes. + schema_error_hard_fail_after: std::time::Duration, } impl Session { @@ -110,6 +139,7 @@ impl Session { upstream_auth: KafkaClientAuth, read_buffer_size: usize, combined_partition_fetch_limit: usize, + schema_error_hard_fail_after: std::time::Duration, ) -> Self { Self { app, @@ -119,6 +149,8 @@ impl Session { read_buffer_size, combined_partition_fetch_limit, reads: HashMap::new(), + cooldown: HashMap::new(), + schema_error_hard_fail_after, auth: None, secret, data_preview_state: SessionDataPreviewState::Unknown, @@ -612,6 +644,18 @@ impl Session { if reaped > 0 { tracing::info!(reaped, remaining = self.reads.len(), "reaped stale reads"); } + + // Drop cooldown entries for journals no client is fetching anymore. + // A cooldown otherwise only clears when a partition is fetched again + // with an updated schema hash, so an abandoned partition's entry + // would persist for the life of the connection. + let journals_in_use: std::collections::HashSet<&str> = self + .reads + .values() + .map(|pending| pending.journal_name.as_str()) + .collect(); + self.cooldown + .retain(|journal_name, _| journals_in_use.contains(journal_name.as_str())); } /// Fetch records from select "partitions" (journals) and "topics" (collections). @@ -739,10 +783,15 @@ impl Session { } }; - let pending_info = self.reads.get(&key).map(|p| (p.offset, p.leader_epoch)); + let pending_info = self + .reads + .get(&key) + .map(|p| (p.offset, p.leader_epoch, p.handle.is_some())); match pending_info { - Some((pending_offset, pending_epoch)) if pending_offset == fetch_offset => { + Some((pending_offset, pending_epoch, has_handle)) + if pending_offset == fetch_offset && has_handle => + { // Validate pending read's epoch is still current let auth = self.auth.as_ref().unwrap(); let current_epoch = Collection::new(&auth, &key.0) @@ -911,6 +960,35 @@ impl Session { tracing::debug!(collection = ?&key.0, partition=partition_request.partition, "Partition doesn't exist!"); continue; // Partition doesn't exist. }; + + let journal_name = partition.spec.name.clone(); + + if let Some(cooled) = self.cooldown.get(&journal_name) { + // If the schema hasn't changed, no need to retry this partition. + if cooled.schema_hash == collection.schema_hash { + self.reads.insert( + key.clone(), + PendingRead { + offset: fetch_offset, + last_write_head: fetch_offset, + leader_epoch: collection.binding_backfill_counter as i32, + last_accessed: std::time::Instant::now(), + journal_name, + schema_hash: collection.schema_hash.clone(), + handle: None, + }, + ); + continue; + } + + // Otherwise since we have a new schema remove it from cooldown. + tracing::info!( + journal = journal_name, + "partition received updated schema, exiting cooldown" + ); + self.cooldown.remove(&journal_name); + } + let (key_schema_id, value_schema_id) = collection.registered_schema_ids(&pg_client).await?; let pending = PendingRead { @@ -918,76 +996,80 @@ impl Session { last_write_head: fetch_offset, leader_epoch: collection.binding_backfill_counter as i32, last_accessed: std::time::Instant::now(), - handle: tokio_util::task::AbortOnDropHandle::new(match data_preview_params { - // Startree: 0, Tinybird: 12 - Some(PartitionOffset { - fragment_start, - offset: latest_offset, - .. - }) if latest_offset - fetch_offset <= 12 => { - let diff = latest_offset - fetch_offset; - metrics::counter!( - "dekaf_fetch_requests", - "topic_name" => key.0.to_string(), - "partition_index" => key.1.to_string(), - "task_name" => task_name.to_string(), - "state" => "new_data_preview_read" - ) - .increment(1); - tokio::spawn(propagate_task_forwarder( - Read::new( - self.app.task_manager.get_listener(task_name.as_str()), - &collection, - partition, - fragment_start, - key_schema_id, - value_schema_id, - Some(partition_request.fetch_offset - 1), - &auth, - self.read_buffer_size, + journal_name, + schema_hash: collection.schema_hash.clone(), + handle: Some(tokio_util::task::AbortOnDropHandle::new( + match data_preview_params { + // Startree: 0, Tinybird: 12 + Some(PartitionOffset { + fragment_start, + offset: latest_offset, + .. + }) if latest_offset - fetch_offset <= 12 => { + let diff = latest_offset - fetch_offset; + metrics::counter!( + "dekaf_fetch_requests", + "topic_name" => key.0.to_string(), + "partition_index" => key.1.to_string(), + "task_name" => task_name.to_string(), + "state" => "new_data_preview_read" ) - .await? - .next_batch( - // Have to read at least 2 docs, as the very last doc - // will probably be a control document and will be - // ignored by the consumer, looking like 0 docs were read - crate::read::ReadTarget::Docs(max(diff as usize, 2)), - timeout, - ), - )) - } - _ => { - metrics::counter!( - "dekaf_fetch_requests", - "topic_name" => key.0.to_string(), - "partition_index" => key.1.to_string(), - "task_name" => task_name.to_string(), - "state" => "new_regular_read" - ) - .increment(1); - tokio::spawn(propagate_task_forwarder( - Read::new( - self.app.task_manager.get_listener(task_name.as_str()), - &collection, - partition, - fetch_offset, - key_schema_id, - value_schema_id, - None, - &auth, - self.read_buffer_size, + .increment(1); + tokio::spawn(propagate_task_forwarder( + Read::new( + self.app.task_manager.get_listener(task_name.as_str()), + &collection, + partition, + fragment_start, + key_schema_id, + value_schema_id, + Some(partition_request.fetch_offset - 1), + &auth, + self.read_buffer_size, + ) + .await? + .next_batch( + // Have to read at least 2 docs, as the very last doc + // will probably be a control document and will be + // ignored by the consumer, looking like 0 docs were read + crate::read::ReadTarget::Docs(max(diff as usize, 2)), + timeout, + ), + )) + } + _ => { + metrics::counter!( + "dekaf_fetch_requests", + "topic_name" => key.0.to_string(), + "partition_index" => key.1.to_string(), + "task_name" => task_name.to_string(), + "state" => "new_regular_read" ) - .await? - .next_batch( - crate::read::ReadTarget::Bytes( - (partition_request.partition_max_bytes as usize) - .min(per_partition_limit), + .increment(1); + tokio::spawn(propagate_task_forwarder( + Read::new( + self.app.task_manager.get_listener(task_name.as_str()), + &collection, + partition, + fetch_offset, + key_schema_id, + value_schema_id, + None, + &auth, + self.read_buffer_size, + ) + .await? + .next_batch( + crate::read::ReadTarget::Bytes( + (partition_request.partition_max_bytes as usize) + .min(per_partition_limit), + ), + timeout, ), - timeout, - ), - )) - } - }), + )) + } + }, + )), }; tracing::info!( @@ -999,14 +1081,16 @@ impl Session { ); if let Some(old) = self.reads.insert(key.clone(), pending) { - tracing::warn!( - topic = topic_request.topic.as_str(), - partition = partition_request.partition, - old_offset = old.offset, - new_offset = fetch_offset, - read_lifetime = ?old.last_accessed.elapsed(), - "discarding pending read due to offset jump", - ); + if old.offset != fetch_offset { + tracing::warn!( + topic = topic_request.topic.as_str(), + partition = partition_request.partition, + old_offset = old.offset, + new_offset = fetch_offset, + read_lifetime = ?old.last_accessed.elapsed(), + "discarding pending read due to offset jump", + ); + } } } } @@ -1103,7 +1187,56 @@ impl Session { continue; } - let (read, batch) = (&mut pending.handle).await??; + let Some(handle) = &mut pending.handle else { + // Still cooling down after a schema-validation error. + let cooled = self + .cooldown + .get(&pending.journal_name) + .context("cooling-down PendingRead has no cooldown entry")?; + + if cooled.first_failed_at.elapsed() > self.schema_error_hard_fail_after { + anyhow::bail!( + "unable to validate schema and no updated schema received for journal: {}", + pending.journal_name, + ); + } + + partition_responses.push( + PartitionData::default() + .with_partition_index(partition_request.partition) + .with_error_code(ResponseError::LeaderNotAvailable.code()), + ); + continue; + }; + + let (read, batch) = match handle.await? { + Ok(ok) => ok, + Err(err) => { + if !is_schema_validation_error(&err) { + return Err(err); + } + + tracing::warn!( + journal = %pending.journal_name, + error = ?err, + "partition failed schema validation, entering cooldown" + ); + + self.cooldown.entry(pending.journal_name.clone()).or_insert( + CooldownEntry { + first_failed_at: std::time::Instant::now(), + schema_hash: pending.schema_hash.clone(), + }, + ); + partition_responses.push( + PartitionData::default() + .with_partition_index(partition_request.partition) + .with_error_code(ResponseError::LeaderNotAvailable.code()), + ); + self.reads.remove(&key); + continue; + } + }; let batch = match batch { BatchResult::TargetExceededBeforeTimeout(b) => Some(b), @@ -1144,8 +1277,8 @@ impl Session { pending.offset = read.offset; pending.last_write_head = read.last_write_head; pending.last_accessed = std::time::Instant::now(); - pending.handle = tokio_util::task::AbortOnDropHandle::new(tokio::spawn( - propagate_task_forwarder( + pending.handle = Some(tokio_util::task::AbortOnDropHandle::new( + tokio::spawn(propagate_task_forwarder( read.next_batch( crate::read::ReadTarget::Bytes( (partition_request.partition_max_bytes as usize) @@ -1153,7 +1286,7 @@ impl Session { ), timeout, ), - ), + )), )); let response_high_watermark = match resolve_high_watermark( diff --git a/crates/dekaf/src/topology.rs b/crates/dekaf/src/topology.rs index 85610459af3..4fb8d32ddfb 100644 --- a/crates/dekaf/src/topology.rs +++ b/crates/dekaf/src/topology.rs @@ -58,6 +58,7 @@ pub struct Collection { pub spec: flow::CollectionSpec, pub uuid_ptr: json::Pointer, pub value_schema: avro::Schema, + pub schema_hash: String, pub extractors: Vec<(avro::Schema, utils::CustomizableExtractor)>, pub binding_backfill_counter: u32, } @@ -261,6 +262,17 @@ impl Collection { let key_schema = avro::key_to_avro(&key_ptr, collection_schema_shape); + // Content-addresses the derived key/value schemas so callers can + // cheaply detect when a binding's effective schema has changed, + // without a network round trip (unlike `registered_schema_id`'s + // `avro_schema_md5`, which addresses the same schemas against the + // control plane's schema registry table). + let schema_hash = { + let key_json = serde_json::to_value(&key_schema).unwrap().to_string(); + let value_json = serde_json::to_value(&value_schema).unwrap().to_string(); + format!("{:x}", md5::compute(format!("{key_json}{value_json}"))) + }; + let (mut not_before, not_after) = ( binding.not_before.map(|b| { uuid::Clock::from_unix(b.seconds.try_into().unwrap(), b.nanos.try_into().unwrap()) @@ -315,6 +327,7 @@ impl Collection { spec: collection_spec, uuid_ptr, value_schema, + schema_hash, extractors, // Start the backfill counter (which will map to the topic leader epoch) at 1, not 0. // Kafka consumers don't seem to handle going from epoch 0 to epoch 1 gracefully. Specifically, From d5d4c743779df147039f71e050449221d3a89299 Mon Sep 17 00:00:00 2001 From: Daniel Nelson Date: Thu, 13 Aug 2026 18:30:21 -0700 Subject: [PATCH 2/2] dekaf: add e2e test for schema cooldown --- crates/dekaf/tests/e2e/fetch_offsets.rs | 38 +---- .../e2e/fixtures/schema_cooldown.flow.yaml | 57 +++++++ crates/dekaf/tests/e2e/harness.rs | 30 ++++ crates/dekaf/tests/e2e/main.rs | 1 + crates/dekaf/tests/e2e/raw_kafka.rs | 71 +++++++++ crates/dekaf/tests/e2e/schema_cooldown.rs | 149 ++++++++++++++++++ 6 files changed, 311 insertions(+), 35 deletions(-) create mode 100644 crates/dekaf/tests/e2e/fixtures/schema_cooldown.flow.yaml create mode 100644 crates/dekaf/tests/e2e/schema_cooldown.rs diff --git a/crates/dekaf/tests/e2e/fetch_offsets.rs b/crates/dekaf/tests/e2e/fetch_offsets.rs index caa8e1f96e6..f6b9b02201e 100644 --- a/crates/dekaf/tests/e2e/fetch_offsets.rs +++ b/crates/dekaf/tests/e2e/fetch_offsets.rs @@ -1,9 +1,7 @@ use super::DekafTestEnv; -use super::raw_kafka::TestKafkaClient; +use super::raw_kafka::{TestKafkaClient, decode_fetch_records}; use anyhow::Context; -use bytes::Buf; -use kafka_protocol::messages; -use kafka_protocol::records::{Record, RecordBatchDecoder}; +use kafka_protocol::records::Record; use serde_json::json; use std::time::Duration; @@ -11,36 +9,6 @@ const FIXTURE: &str = include_str!("fixtures/basic.flow.yaml"); const TOPIC: &str = "test_topic"; const PARTITION: i32 = 0; -/// Decode all records, including control records, from a raw FetchResponse. -/// -/// Unlike the rdkafka consumer, this surfaces exactly what Dekaf put on the -/// wire: librdkafka silently filters out control records and records below -/// the fetch offset, which masks bugs in which documents a fetch serves. -fn decode_fetch_records(resp: &messages::FetchResponse) -> anyhow::Result> { - let partition = resp - .responses - .iter() - .find(|t| t.topic.as_str() == TOPIC) - .and_then(|t| t.partitions.iter().find(|p| p.partition_index == PARTITION)) - .context("missing partition in fetch response")?; - - anyhow::ensure!( - partition.error_code == 0, - "fetch returned error code {}", - partition.error_code - ); - - let Some(mut buf) = partition.records.clone() else { - return Ok(Vec::new()); - }; - - let mut records = Vec::new(); - while buf.has_remaining() { - records.extend(RecordBatchDecoder::decode(&mut buf)?.records); - } - Ok(records) -} - async fn fetch_records_at( client: &mut TestKafkaClient, offset: i64, @@ -48,7 +16,7 @@ async fn fetch_records_at( let resp = client .fetch_with_epoch(TOPIC, PARTITION, offset, -1) .await?; - decode_fetch_records(&resp) + decode_fetch_records(&resp, TOPIC, PARTITION) } /// Fetch at `offset`, retrying empty responses (e.g. while the server-side diff --git a/crates/dekaf/tests/e2e/fixtures/schema_cooldown.flow.yaml b/crates/dekaf/tests/e2e/fixtures/schema_cooldown.flow.yaml new file mode 100644 index 00000000000..2da0b6a4324 --- /dev/null +++ b/crates/dekaf/tests/e2e/fixtures/schema_cooldown.flow.yaml @@ -0,0 +1,57 @@ +collections: + test_data_a: + writeSchema: + type: object + properties: + id: { type: string } + value: { type: [integer, string] } + required: [id] + readSchema: + type: object + properties: + id: { type: string } + value: { type: integer } + required: [id] + key: [/id] + + test_data_b: + schema: + type: object + properties: + id: { type: string } + value: { type: integer } + required: [id] + key: [/id] + +captures: + source_ingest: + endpoint: + connector: + image: ghcr.io/estuary/source-http-ingest:dev + config: + paths: ["/data_a", "/data_b"] + bindings: + - resource: { path: "/data_a", stream: "/data_a" } + target: test_data_a + - resource: { path: "/data_b", stream: "/data_b" } + target: test_data_b + +materializations: + dekaf_test: + endpoint: + dekaf: + variant: testing + config: + token: "test-token-12345" + strict_topic_names: false + bindings: + - source: test_data_a + resource: { topic_name: topic_a } + fields: + recommended: true + exclude: [flow_published_at] + - source: test_data_b + resource: { topic_name: topic_b } + fields: + recommended: true + exclude: [flow_published_at] diff --git a/crates/dekaf/tests/e2e/harness.rs b/crates/dekaf/tests/e2e/harness.rs index a139ecd5abd..3a7a639917b 100644 --- a/crates/dekaf/tests/e2e/harness.rs +++ b/crates/dekaf/tests/e2e/harness.rs @@ -526,6 +526,36 @@ impl DekafTestEnv { self.publish_catalog(&catalog).await } + /// Republish `collection_name` with a new `read_schema`, leaving its + /// `write_schema` untouched. The fixture must already declare separate + /// `writeSchema`/`readSchema` (not a single `schema`) for this to apply. + pub async fn set_collection_read_schema( + &self, + collection_name: &str, + read_schema: serde_json::Value, + ) -> anyhow::Result<()> { + let collection = models::Collection::new(collection_name); + let mut coll_def = self + .catalog + .collections + .get(&collection) + .context("collection not in fixture")? + .clone(); + + coll_def.read_schema = Some(models::Schema::new(models::RawValue::from_value( + &read_schema, + ))); + + tracing::info!(%collection_name, "Updating collection read schema"); + + let catalog = models::Catalog { + collections: [(collection, coll_def)].into(), + ..Default::default() + }; + + self.publish_catalog(&catalog).await + } + /// Cleanup test specs synchronously. fn cleanup_sync(&self) { let cmd = match flowctl_command() { diff --git a/crates/dekaf/tests/e2e/main.rs b/crates/dekaf/tests/e2e/main.rs index 9fa3be34ca8..d98575bf716 100644 --- a/crates/dekaf/tests/e2e/main.rs +++ b/crates/dekaf/tests/e2e/main.rs @@ -12,6 +12,7 @@ mod list_offsets; mod migration; mod not_ready; mod partition_eofs; +mod schema_cooldown; pub use harness::{ ConnectionInfo, DekafTestEnv, cluster_name, cluster_name_2, connection_info_for_dataplane, diff --git a/crates/dekaf/tests/e2e/raw_kafka.rs b/crates/dekaf/tests/e2e/raw_kafka.rs index 0b9959943c6..42e00fb4104 100644 --- a/crates/dekaf/tests/e2e/raw_kafka.rs +++ b/crates/dekaf/tests/e2e/raw_kafka.rs @@ -1,7 +1,10 @@ +use anyhow::Context; +use bytes::Buf; use dekaf::{KafkaApiClient, KafkaClientAuth}; use kafka_protocol::{ messages::{self, offset_fetch_response::OffsetFetchResponsePartition}, protocol::StrBytes, + records::{Record, RecordBatchDecoder}, }; /// Protocol versions to use for test requests. @@ -66,6 +69,40 @@ impl TestKafkaClient { self.inner.send_request(req, Some(header)).await } + /// Fetch several `(topic, partition, fetch_offset)` requests in a single + /// call, on one connection — for asserting that a broken partition's + /// error doesn't affect a sibling partition's Fetch response. + pub async fn fetch_multi( + &mut self, + requests: &[(&str, i32, i64)], + ) -> anyhow::Result { + let req = messages::FetchRequest::default() + .with_max_wait_ms(1000) + .with_min_bytes(1) + .with_max_bytes(1024 * 1024) + .with_topics( + requests + .iter() + .map(|&(topic, partition, offset)| { + messages::fetch_request::FetchTopic::default() + .with_topic(topic_name(topic)) + .with_partitions(vec![ + messages::fetch_request::FetchPartition::default() + .with_partition(partition) + .with_fetch_offset(offset) + .with_partition_max_bytes(1024 * 1024), + ]) + }) + .collect(), + ); + + let header = messages::RequestHeader::default() + .with_request_api_key(messages::ApiKey::Fetch as i16) + .with_request_api_version(protocol_versions::FETCH); + + self.inner.send_request(req, Some(header)).await + } + /// ListOffsets with explicit `current_leader_epoch`. /// /// - `timestamp = -2`: earliest offset @@ -195,6 +232,40 @@ impl TestKafkaClient { } } +/// Decode all records, including control records, from a raw FetchResponse. +/// +/// Unlike the rdkafka consumer, this surfaces exactly what Dekaf put on the +/// wire: librdkafka silently filters out control records and records below +/// the fetch offset, which masks bugs in which documents a fetch serves. +pub fn decode_fetch_records( + resp: &messages::FetchResponse, + topic: &str, + partition: i32, +) -> anyhow::Result> { + let partition_data = resp + .responses + .iter() + .find(|t| t.topic.as_str() == topic) + .and_then(|t| t.partitions.iter().find(|p| p.partition_index == partition)) + .context("missing partition in fetch response")?; + + anyhow::ensure!( + partition_data.error_code == 0, + "fetch returned error code {}", + partition_data.error_code + ); + + let Some(mut buf) = partition_data.records.clone() else { + return Ok(Vec::new()); + }; + + let mut records = Vec::new(); + while buf.has_remaining() { + records.extend(RecordBatchDecoder::decode(&mut buf)?.records); + } + Ok(records) +} + /// Extract the error code from a FetchResponse for a specific topic/partition. pub fn fetch_partition_error( resp: &messages::FetchResponse, diff --git a/crates/dekaf/tests/e2e/schema_cooldown.rs b/crates/dekaf/tests/e2e/schema_cooldown.rs new file mode 100644 index 00000000000..fcf73de0539 --- /dev/null +++ b/crates/dekaf/tests/e2e/schema_cooldown.rs @@ -0,0 +1,149 @@ +use super::DekafTestEnv; +use crate::raw_kafka::{TestKafkaClient, decode_fetch_records, fetch_partition_error}; +use kafka_protocol::ResponseError; +use kafka_protocol::records::Record; +use serde_json::json; +use std::time::Duration; + +const FIXTURE: &str = include_str!("fixtures/schema_cooldown.flow.yaml"); + +// The task_manager caches the MaterializationSpec and refreshes every `spec_ttl` (2m by default). +// We may need to wait this long after updating the schema for it to get picked up for use in a +// fetch. +const SCHEMA_PROPAGATION_TIMEOUT: Duration = Duration::from_secs(150); + +fn widened_value_read_schema() -> serde_json::Value { + json!({ + "type": "object", + "properties": {"id": {"type": "string"}, "value": {"type": ["integer", "string"]}}, + "required": ["id"], + }) +} + +/// A partition whose documents fail Avro schema validation reports +/// LeaderNotAvailable instead of erroring the whole connection, and a +/// sibling partition on the same connection is unaffected. Once the +/// collection's schema catches up, the partition recovers without a +/// reconnect, and the previously-unreadable document is served. +#[tokio::test] +async fn test_schema_error_cooldown_isolates_partition_and_recovers() -> anyhow::Result<()> { + super::init_tracing(); + + let env = DekafTestEnv::setup("schema_cooldown", FIXTURE).await?; + let collection_a = format!("{}/test_data_a", env.namespace); + + // Write a valid document to both topics. + env.inject_documents("data_a", [json!({"id": "1", "value": 1})]) + .await?; + env.inject_documents("data_b", [json!({"id": "1", "value": 1})]) + .await?; + + let info = env.connection_info().await?; + let token = env.dekaf_token()?; + let mut client = TestKafkaClient::connect(&info.broker, &info.username, &token).await?; + + let mut records_a: Vec = vec![]; + let mut records_b: Vec = vec![]; + + // As a baseline check that everything is working by reading these documents back. + let deadline = std::time::Instant::now() + Duration::from_secs(30); + loop { + let resp = client + .fetch_multi(&[("topic_a", 0, 0), ("topic_b", 0, 0)]) + .await?; + + assert!(fetch_partition_error(&resp, "topic_a", 0) == Some(0)); + records_a.extend(decode_fetch_records(&resp, "topic_a", 0)?); + + assert!(fetch_partition_error(&resp, "topic_b", 0) == Some(0)); + records_b.extend(decode_fetch_records(&resp, "topic_b", 0)?); + + if records_a.iter().filter(|r| !r.control).count() == 1 + && records_b.iter().filter(|r| !r.control).count() == 1 + { + break; + } + if std::time::Instant::now() > deadline { + anyhow::bail!("baseline fetch never succeeded for both topics"); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + + // Write an document that does not validate to the schema of topic_a, topic_b gets another + // valid document. + env.inject_documents("data_a", [json!({"id": "2", "value": "a string"})]) + .await?; + env.inject_documents("data_b", [json!({"id": "2", "value": 2})]) + .await?; + + // Now topic_a should begin returning LeaderNotAvailable, while topic_b is unaffected. + let leader_not_available = ResponseError::LeaderNotAvailable.code(); + let deadline = std::time::Instant::now() + Duration::from_secs(30); + loop { + let resp = client + .fetch_multi(&[ + ("topic_a", 0, records_a.iter().last().unwrap().offset + 1), + ("topic_b", 0, records_b.iter().last().unwrap().offset + 1), + ]) + .await?; + + // It is possible that the new document with the schema error hasn't yet landed, so we + // still receive an success code. + let err_a = fetch_partition_error(&resp, "topic_a", 0); + anyhow::ensure!( + err_a == Some(0) || err_a == Some(leader_not_available), + "topic_a should be no error or LeaderNotAvailable, received {err_a:?}" + ); + + assert!(fetch_partition_error(&resp, "topic_b", 0) == Some(0)); + records_b.extend(decode_fetch_records(&resp, "topic_b", 0)?); + + if err_a == Some(leader_not_available) + && records_b.iter().filter(|r| !r.control).count() == 2 + { + break; + } + + if std::time::Instant::now() > deadline { + anyhow::bail!( + "topic_a never entered schema cooldown after writing an invalid document" + ); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + + // Widen topic_a readSchema to accept the string too, simulating + // an inferred schema catching up with the data actually being written. + env.set_collection_read_schema(&collection_a, widened_value_read_schema()) + .await?; + + // Once the schema propagates, topic_a should recover and + // serve the previously-missing document; topic_b remains fine throughout. + let deadline = std::time::Instant::now() + SCHEMA_PROPAGATION_TIMEOUT; + loop { + let resp = client + .fetch_multi(&[ + ("topic_a", 0, records_a.iter().last().unwrap().offset + 1), + ("topic_b", 0, records_b.iter().last().unwrap().offset + 1), + ]) + .await?; + + if fetch_partition_error(&resp, "topic_a", 0) == Some(0) { + let records = decode_fetch_records(&resp, "topic_a", 0)?; + records_a.extend(records); + } + + assert!(fetch_partition_error(&resp, "topic_b", 0) == Some(0)); + records_b.extend(decode_fetch_records(&resp, "topic_b", 0)?); + + if records_a.iter().filter(|r| !r.control).count() == 2 { + break; + } + if std::time::Instant::now() > deadline { + anyhow::bail!("topic_a never recovered its missing document after the schema update"); + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + + Ok(()) +}