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, 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(()) +}