diff --git a/changelog.md b/changelog.md index 03b69f046..5fc3cd027 100644 --- a/changelog.md +++ b/changelog.md @@ -4,7 +4,10 @@ See also the [rdkafka-sys changelog](rdkafka-sys/changelog.md). ## Unreleased -None +* Spawn the periodic fallback wakeup for queues created via + `StreamConsumer::split_partition_queue` too. Previously only the main queue + had it, so a lost edge-triggered wakeup could stall a single partition's + stream indefinitely while other partitions kept consuming (#665). ## 0.39.0 (2026-01-25) diff --git a/src/consumer/stream_consumer.rs b/src/consumer/stream_consumer.rs index 469c7e1bb..beebc0d8a 100644 --- a/src/consumer/stream_consumer.rs +++ b/src/consumer/stream_consumer.rs @@ -199,6 +199,7 @@ where { base: Arc>, wakers: Arc, + poll_interval: Duration, _shutdown_trigger: oneshot::Sender<()>, _runtime: PhantomData, } @@ -265,6 +266,7 @@ where Ok(StreamConsumer { base, wakers, + poll_interval, _shutdown_trigger: shutdown_trigger, _runtime: PhantomData, }) @@ -356,12 +358,33 @@ where self: &Arc, topic: &str, partition: i32, - ) -> Option> { + ) -> Option> + where + R: AsyncRuntime, + { self.base .split_partition_queue(topic, partition) .map(|queue| { let wakers = Arc::new(WakerSlab::new()); unsafe { enable_nonempty_callback(&queue.queue, &wakers) }; + + // Same periodic fallback wakeup as the main queue (see + // `from_config_and_context`): the non-empty callback is + // edge-triggered, so a lost wakeup would otherwise park this + // partition's stream forever. Holds a `Weak` so the loop ends + // once the queue is dropped and its `WakerSlab` is freed. + let poll_interval = self.poll_interval; + let weak_wakers = Arc::downgrade(&wakers); + R::spawn(async move { + loop { + R::delay_for(poll_interval / 2).await; + match weak_wakers.upgrade() { + Some(wakers) => wakers.wake_all(), + None => break, + } + } + }); + StreamPartitionQueue { queue, wakers, diff --git a/tests/test_partition_queue_wakeup.rs b/tests/test_partition_queue_wakeup.rs new file mode 100644 index 000000000..552286ba3 --- /dev/null +++ b/tests/test_partition_queue_wakeup.rs @@ -0,0 +1,80 @@ +//! Deterministic regression test for the split-partition-queue fallback wakeup. +//! +//! No broker required: a parked partition stream with no data and no non-empty +//! callback edge can only be woken by the periodic fallback loop. With the fix +//! it is woken within ~max.poll.interval.ms; without it, never. + +use std::pin::pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll, Wake, Waker}; + +use futures::stream::Stream; +use tokio::time::{sleep, Duration, Instant}; + +use rdkafka::config::ClientConfig; +use rdkafka::consumer::{Consumer, StreamConsumer}; +use rdkafka::topic_partition_list::{Offset, TopicPartitionList}; + +// Waker that records when it is woken, so we can observe the fallback wake. +struct SignalWaker(AtomicBool); + +impl Wake for SignalWaker { + fn wake(self: Arc) { + self.0.store(true, Ordering::SeqCst); + } + fn wake_by_ref(self: &Arc) { + self.0.store(true, Ordering::SeqCst); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_split_partition_queue_fallback_wakeup() { + let _r = env_logger::try_init(); + + // poll_interval = 500ms => fallback fires every 250ms. Dummy bootstrap: + // we never connect, the test is purely about local waker mechanics. + let consumer: Arc = Arc::new( + ClientConfig::new() + .set("group.id", "test-fallback-wakeup") + .set("bootstrap.servers", "localhost:9092") + .set("session.timeout.ms", "1000") + .set("max.poll.interval.ms", "1000") + .set("enable.auto.commit", "false") + .create() + .expect("consumer creation failed"), + ); + + let mut tpl = TopicPartitionList::new(); + tpl.add_partition_offset("nonexistent-topic", 0, Offset::Beginning) + .unwrap(); + consumer.assign(&tpl).unwrap(); + + let pq = consumer + .split_partition_queue("nonexistent-topic", 0) + .expect("split_partition_queue returned None"); + + let signal = Arc::new(SignalWaker(AtomicBool::new(false))); + let waker = Waker::from(signal.clone()); + + // Park the stream: first poll registers our waker and returns Pending + // (empty queue, no broker -> no data). + let mut stream = pin!(pq.stream()); + let mut cx = Context::from_waker(&waker); + match stream.as_mut().poll_next(&mut cx) { + Poll::Pending => {} + Poll::Ready(_) => panic!("expected Pending on empty queue"), + } + + // The only thing that can wake us now is the fallback loop. With the fix it + // fires within ~250ms; without it, this never trips and we time out. + let deadline = Instant::now() + Duration::from_secs(3); + while Instant::now() < deadline { + if signal.0.load(Ordering::SeqCst) { + return; // woken by the fallback loop — pass + } + sleep(Duration::from_millis(50)).await; + } + + panic!("partition queue stream was never woken — split queue is not covered by the fallback wakeup"); +}