From 3f0a3af1ee8dc7f23fa52530a851a5b105c5935c Mon Sep 17 00:00:00 2001 From: Tony Le Date: Wed, 21 Oct 2026 14:44:09 -0400 Subject: [PATCH] ref(consumer): Remove the legacy rdkafka consumer path --- .github/workflows/ci.yml | 64 +------- src/config/mod.rs | 7 - src/kafka/consumer.rs | 217 +++------------------------ src/kafka/consumer_backend.rs | 270 +++++++++++----------------------- src/main.rs | 1 - 5 files changed, 117 insertions(+), 442 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d9047bbf..fe3e97ec 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -268,14 +268,8 @@ jobs: token: ${{ secrets.CODECOV_TOKEN }} rebalance-integration-test: - name: Rebalance integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Rebalance integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -312,14 +306,8 @@ jobs: make test-rebalance worker-processing-integration-test: - name: Worker processing integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Worker processing integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -356,14 +344,8 @@ jobs: make test-worker-processing upkeep-retry-integration-test: - name: Upkeep retry integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Upkeep retry integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -400,14 +382,8 @@ jobs: make test-upkeep-retry upkeep-expiry-integration-test: - name: Upkeep expiry integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Upkeep expiry integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -444,14 +420,8 @@ jobs: make test-upkeep-expiry upkeep-delay-integration-test: - name: Upkeep delay integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Upkeep delay integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -488,14 +458,8 @@ jobs: make test-upkeep-delay failed-tasks-integration-test: - name: Failed tasks integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Failed tasks integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -532,14 +496,8 @@ jobs: make test-failed-tasks multi-topic-integration-test: - name: Multi-topic integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Multi-topic integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -576,14 +534,8 @@ jobs: make test-multi-topic push-postgres-integration-test: - name: Push mode postgres integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} + name: Push mode postgres integration test runs-on: ubuntu-24.04 - strategy: - fail-fast: false - matrix: - consumer_backend: [rdkafka, arroyo] - env: - TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 diff --git a/src/config/mod.rs b/src/config/mod.rs index ab6f37fb..3c5289f3 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -151,9 +151,6 @@ pub struct Config { /// The path to the runtime config file pub runtime_config_path: Option, - /// Consume with Arroyo's async Kafka consumer instead of rdkafka's stream consumer. - pub use_arroyo_consumer: bool, - /// The frequency at which upkeep tasks /// (discarding, retrying activations, etc.) are executed. pub upkeep_task_interval_ms: u64, @@ -300,7 +297,6 @@ impl Default for Config { kafka_send_timeout_ms: 500, store: StoreConfig::default(), runtime_config_path: None, - use_arroyo_consumer: false, upkeep_task_interval_ms: 1000, upkeep_unhealthy_interval_ms: 5000, health_check_killswitched: false, @@ -1046,7 +1042,6 @@ mod tests { assert_eq!(config.store.sqlite.vacuum_page_count, None); assert_eq!(config.canary_tasks, 0); assert!(config.worker_map.is_empty()); - assert!(!config.use_arroyo_consumer); } #[test] @@ -1241,7 +1236,6 @@ mod tests { jail.set_env("TASKBROKER_DATABASE_ADAPTER", "postgres"); jail.set_env("TASKBROKER_MAX_PROCESSING_ATTEMPTS", "5"); jail.set_env("TASKBROKER_CANARY_TASKS", "2"); - jail.set_env("TASKBROKER_USE_ARROYO_CONSUMER", "true"); let args = Args { run: Run::Broker, @@ -1252,7 +1246,6 @@ mod tests { assert_eq!(config.store.adapter, DatabaseAdapter::Postgres); assert_eq!(config.store.max_processing_attempts, 5); assert_eq!(config.canary_tasks, 2); - assert!(config.use_arroyo_consumer); Ok(()) }); diff --git a/src/kafka/consumer.rs b/src/kafka/consumer.rs index 520c8b7a..629a3502 100644 --- a/src/kafka/consumer.rs +++ b/src/kafka/consumer.rs @@ -7,14 +7,11 @@ use std::sync::mpsc::{SyncSender, sync_channel}; use std::time::Duration; use std::{cmp, iter}; -use rdkafka::consumer::stream_consumer::StreamPartitionQueue; -use rdkafka::consumer::{BaseConsumer, Consumer, ConsumerContext, Rebalance, StreamConsumer}; -use rdkafka::error::{KafkaError, KafkaResult}; -use rdkafka::message::BorrowedMessage; +use rdkafka::consumer::{BaseConsumer, Consumer}; +use rdkafka::error::KafkaError; use rdkafka::types::RDKafkaErrorCode; -use rdkafka::{ClientConfig, ClientContext, Offset, TopicPartitionList}; +use rdkafka::{ClientConfig, Offset, TopicPartitionList}; -use tokio::runtime::Handle; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}; use tokio::sync::{mpsc, oneshot}; use tokio::task::{JoinError, JoinSet}; @@ -26,26 +23,22 @@ use tokio_util::sync::CancellationToken; use anyhow::{Error, anyhow}; use futures::{Stream, StreamExt, future, pin_mut}; -use tracing::{debug, error, info, instrument, warn}; +use sentry_arroyo::backends::kafka::async_consumer::AsyncKafkaConsumer; +use tracing::{debug, error, info, instrument}; use crate::store::traits::ActivationStore; use crate::store::types::TopicPartition; -use super::consumer_backend::{Assignment, ConsumerBackend}; +use super::consumer_backend::{Assignment, KafkaContext, create_consumer}; use super::message::MessageBackend; pub async fn start_no_consume_mode( topics: &[&str], kafka_client_config: &ClientConfig, ) -> Result<(), Error> { - let topics_tag = topics.join(","); - let (event_sender, _) = unbounded_channel(); - let context = KafkaContext::new(event_sender.clone(), topics_tag.clone()); - let consumer: Arc> = Arc::new( - kafka_client_config - .create_with_context(context) - .expect("Consumer creation failed"), - ); + let consumer: BaseConsumer = kafka_client_config + .create() + .expect("Consumer creation failed"); // Validates: broker reachability, SASL/SSL auth handshake, // topic existence, and partition count — no group join. @@ -93,29 +86,17 @@ pub async fn start_no_consume_mode( pub async fn start_consumer( topics: &[&str], kafka_client_config: &ClientConfig, - use_arroyo: bool, activation_store: Arc, - spawn_actors: impl FnMut(Arc, Assignment) -> ActorHandles, + spawn_actors: impl FnMut(Arc, Assignment) -> ActorHandles, ) -> Result<(), Error> { - let (client_shutdown_sender, client_shutdown_receiver) = oneshot::channel(); let (event_sender, event_receiver) = unbounded_channel(); // Metrics tag only. Ownership keys come per-topic from the assignment `tpl`, // never from this string — so a multi-topic consumer stays correct. let topics_tag = topics.join(","); let context = KafkaContext::new(event_sender.clone(), topics_tag.clone()); - let consumer = Arc::new(ConsumerBackend::new( - topics, - kafka_client_config, - context, - use_arroyo, - )?); + let consumer = Arc::new(create_consumer(topics, kafka_client_config, context)?); handle_shutdown_signals(event_sender.clone()); - // Arroyo polls the consumer on its own thread and commits when it is dropped below, so only - // rdkafka needs a poller to drive rebalances and commit on shutdown. - if let ConsumerBackend::Rdkafka { consumer } = consumer.as_ref() { - poll_consumer_client(consumer.clone(), client_shutdown_receiver); - } metrics::gauge!( "arroyo.consumer.current_partitions", "topic" => topics_tag.clone(), @@ -126,7 +107,6 @@ pub async fn start_consumer( consumer.clone(), event_receiver, activation_store, - client_shutdown_sender, spawn_actors, topics_tag, ) @@ -147,122 +127,6 @@ pub fn handle_shutdown_signals(event_sender: UnboundedSender<(Event, SyncSender< }); } -#[instrument(skip_all)] -pub fn poll_consumer_client( - consumer: Arc>, - mut shutdown: oneshot::Receiver<()>, -) { - task::spawn_blocking(|| { - Handle::current().block_on(async move { - let _guard = elegant_departure::get_shutdown_guard().shutdown_on_drop(); - loop { - select! { - biased; - _ = &mut shutdown => { - debug!("Received shutdown signal, commiting state in sync mode..."); - let _ = consumer.commit_consumer_state(rdkafka::consumer::CommitMode::Sync); - break; - } - msg = consumer.recv() => { - if let Err(KafkaError::MessageConsumption(RDKafkaErrorCode::BrokerTransportFailure)) = msg { - error!("Failed to connect to broker, retrying...") - } else { - error!("Got unexpected status from consumer client: {:?}", msg); - break - } - } - } - } - debug!("Shutdown complete"); - }); - }); -} - -#[derive(Debug)] -pub struct KafkaContext { - pub(super) event_sender: UnboundedSender<(Event, SyncSender<()>)>, - /// The topic(s) this consumer is subscribed to, used as a metric tag. - pub(super) topics_tag: String, -} - -impl KafkaContext { - pub fn new(event_sender: UnboundedSender<(Event, SyncSender<()>)>, topics_tag: String) -> Self { - Self { - event_sender, - topics_tag, - } - } -} - -impl ClientContext for KafkaContext {} - -impl ConsumerContext for KafkaContext { - #[instrument(skip_all)] - fn pre_rebalance(&self, _: &BaseConsumer, rebalance: &Rebalance) { - let (rendezvous_sender, rendezvous_receiver) = sync_channel(0); - match rebalance { - Rebalance::Assign(tpl) => { - debug!("Got pre-rebalance callback, kind: Assign"); - if tpl.count() == 0 { - warn!( - "Got partition assignment with no partitions, \ - this is likely due to there being more consumers than partitions" - ); - return; - } - let _ = self.event_sender.send(( - Event::Assign(Assignment { - partitions: tpl.to_topic_map().keys().cloned().collect(), - arroyo_queues: Vec::new(), - }), - rendezvous_sender, - )); - info!("Partition assignment event sent, waiting for rendezvous..."); - let _ = rendezvous_receiver.recv(); - info!("Rendezvous complete"); - metrics::counter!( - "arroyo.consumer.partitions_assigned.count", - "topic" => self.topics_tag.clone(), - "application" => "taskbroker", - ) - .increment(tpl.count() as u64); - } - Rebalance::Revoke(tpl) => { - debug!("Got pre-rebalance callback, kind: Revoke"); - if tpl.count() == 0 { - warn!( - "Got partition revocation with no partitions, \ - this is likely due to there being more consumers than partitions" - ); - return; - } - let _ = self.event_sender.send(( - Event::Revoke(tpl.to_topic_map().keys().cloned().collect()), - rendezvous_sender, - )); - info!("Partition revocation event sent, waiting for rendezvous..."); - let _ = rendezvous_receiver.recv(); - info!("Rendezvous complete"); - metrics::counter!( - "arroyo.consumer.partitions_revoked.count", - "topic" => self.topics_tag.clone(), - "application" => "taskbroker", - ) - .increment(tpl.count() as u64); - } - Rebalance::Error(err) => { - debug!("Got pre-rebalance callback, kind: Error"); - error!("Got rebalance error: {}", err); - } - } - } - - #[instrument(skip(self))] - fn commit_callback(&self, result: KafkaResult<()>, _offsets: &TopicPartitionList) { - debug!("Got commit callback"); - } -} - #[derive(Debug)] pub enum Event { Assign(Assignment), @@ -279,7 +143,7 @@ pub struct ActorHandles { impl ActorHandles { #[instrument(skip(self))] - async fn shutdown(mut self, deadline: Duration) -> JoinSet> { + async fn shutdown(mut self, deadline: Duration) { debug!("Signaling shutdown to actors..."); self.shutdown.cancel(); info!("Actor shutdown signaled, waiting for rendezvous..."); @@ -295,9 +159,9 @@ impl ActorHandles { ); } } - // Returned so the Arroyo path can join them before committing or closing the consumer. - self.join_set.abort_all(); - self.join_set + // Wait for the aborted actors to exit, so none stores offsets or holds the consumer after + // the partitions are revoked or the consumer is closed. + self.join_set.shutdown().await; } async fn join_next(&mut self) -> Option, JoinError>> { @@ -351,7 +215,7 @@ macro_rules! processing_strategy { reduce: $reduce_first:expr $(,$reduce_rest:expr)*, } ) => {{ - |consumer: Arc<$crate::kafka::consumer_backend::ConsumerBackend>, + |consumer: Arc, assignment: $crate::kafka::consumer_backend::Assignment| -> $crate::kafka::consumer::ActorHandles { let start = std::time::Instant::now(); @@ -364,10 +228,7 @@ macro_rules! processing_strategy { let (map_sender, reduce_receiver) = tokio::sync::mpsc::channel(1); let (err_sender, err_receiver) = tokio::sync::mpsc::channel(1); - let queues = consumer - .partition_queues(assignment) - .expect("Unable to split topic by Partition"); - for queue in queues { + for queue in assignment.queues { handles.spawn($crate::kafka::consumer::map( queue, $map_fn, @@ -418,17 +279,15 @@ enum ConsumerState { #[instrument(skip_all)] pub async fn handle_events( - consumer: Arc, + consumer: Arc, events: UnboundedReceiver<(Event, SyncSender<()>)>, activation_store: Arc, - shutdown_client: oneshot::Sender<()>, - mut spawn_actors: impl FnMut(Arc, Assignment) -> ActorHandles, + mut spawn_actors: impl FnMut(Arc, Assignment) -> ActorHandles, topics_tag: String, ) -> Result<(), anyhow::Error> { const CALLBACK_DURATION: Duration = Duration::from_secs(4); let _guard = elegant_departure::get_shutdown_guard().shutdown_on_drop(); - let mut shutdown_client = Some(shutdown_client); let mut events_stream = UnboundedReceiverStream::new(events); let mut state = ConsumerState::Ready; @@ -473,12 +332,7 @@ pub async fn handle_events( "Revoked TPL should be equal to the subset of TPL we're consuming from" ); activation_store.revoke_partitions(&mut revoked.iter().map(TopicPartition::from)); - let mut actors = handles.shutdown(CALLBACK_DURATION).await; - if consumer.is_arroyo() { - // Join the aborted actors so none still holds the consumer when it - // is closed. - actors.shutdown().await; - } + handles.shutdown(CALLBACK_DURATION).await; metrics::gauge!( "arroyo.consumer.current_partitions", "topic" => topics_tag.clone(), @@ -489,12 +343,7 @@ pub async fn handle_events( } (ConsumerState::Consuming(handles, tpl), Event::Shutdown) => { activation_store.revoke_partitions(&mut tpl.iter().map(TopicPartition::from)); - let mut actors = handles.shutdown(CALLBACK_DURATION).await; - if consumer.is_arroyo() { - actors.shutdown().await; - } - debug!("Signaling shutdown to client..."); - shutdown_client.take(); + handles.shutdown(CALLBACK_DURATION).await; metrics::gauge!( "arroyo.consumer.current_partitions", "topic" => topics_tag.clone(), @@ -518,28 +367,10 @@ pub trait KafkaMessage { fn into_message(self) -> Result; } -impl KafkaMessage for Result, KafkaError> { - fn into_message(self) -> Result { - match self { - Ok(borrowed_msg) => Ok(borrowed_msg.detach().into()), - Err(err) => Err(anyhow!( - "Cannot detach message, got error from kafka: {:?}", - err - )), - } - } -} - pub trait MessageQueue { fn stream(&self) -> impl Stream; } -impl MessageQueue for StreamPartitionQueue { - fn stream(&self) -> impl Stream { - self.stream() - } -} - #[instrument(skip_all)] pub async fn map( queue: impl MessageQueue, @@ -1101,7 +932,7 @@ mod tests { #[case::rendezvous(true)] #[case::deadline(false)] #[tokio::test] - async fn test_native_actor_shutdown_defers_cleanup(#[case] rendezvous_completed: bool) { + async fn test_actor_shutdown_joins_actors(#[case] rendezvous_completed: bool) { let consumer = Arc::new(()); let mut join_set = JoinSet::new(); let (started_sender, started) = oneshot::channel(); @@ -1133,9 +964,7 @@ mod tests { } else { Duration::ZERO }; - let mut cleanup = handles.shutdown(deadline).await; - assert_eq!(Arc::strong_count(&consumer), 2); - cleanup.shutdown().await; + handles.shutdown(deadline).await; assert_eq!(Arc::strong_count(&consumer), 1); } diff --git a/src/kafka/consumer_backend.rs b/src/kafka/consumer_backend.rs index a5817396..fb1b664b 100644 --- a/src/kafka/consumer_backend.rs +++ b/src/kafka/consumer_backend.rs @@ -1,119 +1,58 @@ use std::collections::{BTreeSet, HashMap}; use std::fmt::Debug; -use std::sync::Arc; -use std::sync::mpsc::sync_channel; +use std::sync::mpsc::{SyncSender, sync_channel}; -use rdkafka::consumer::stream_consumer::StreamPartitionQueue; -use rdkafka::consumer::{Consumer, StreamConsumer}; +use anyhow::{Context, Error}; +use futures::Stream; use rdkafka::error::KafkaError; use rdkafka::{ClientConfig, TopicPartitionList}; - -use anyhow::{Context, Error}; -use futures::{Stream, StreamExt}; use sentry_arroyo::backends::kafka::async_consumer::{ AsyncAssignmentCallbacks, AsyncKafkaConsumer, AsyncPartitionQueue, }; use sentry_arroyo::backends::kafka::config::KafkaConfig; use sentry_arroyo::types::{Partition, Topic}; +use tokio::sync::mpsc::UnboundedSender; use tracing::{info, warn}; -use super::consumer::{CommitClient, Event, KafkaContext, KafkaMessage, MessageQueue}; +use super::consumer::{CommitClient, Event, KafkaMessage, MessageQueue}; use super::message::MessageBackend; -pub enum ConsumerBackend { - Rdkafka { - consumer: Arc>, - }, - Arroyo { - consumer: AsyncKafkaConsumer, - }, -} - -impl ConsumerBackend { - pub fn is_arroyo(&self) -> bool { - matches!(self, Self::Arroyo { .. }) - } - - pub fn new( - topics: &[&str], - config: &ClientConfig, - context: KafkaContext, - use_arroyo: bool, - ) -> Result { - if use_arroyo { - let arroyo_config = KafkaConfig::new_config( - vec![ - config - .get("bootstrap.servers") - .context("Missing brokers")? - .to_owned(), - ], - Some(config.config_map().clone()), - ); - let topics: Vec<_> = topics.iter().map(|topic| Topic::new(topic)).collect(); - Ok(Self::Arroyo { - consumer: AsyncKafkaConsumer::new(arroyo_config, &topics, context) - .context("Could not create Arroyo consumer")?, - }) - } else { - let consumer: StreamConsumer = config - .create_with_context(context) - .expect("Consumer creation failed"); - consumer - .subscribe(topics) - .expect("Can't subscribe to specified topics"); - Ok(Self::Rdkafka { - consumer: Arc::new(consumer), - }) - } - } - - pub fn partition_queues( - &self, - assignment: Assignment, - ) -> Result, Error> { - match self { - Self::Rdkafka { consumer } => assignment - .partitions - .iter() - .map(|(topic, partition)| { - consumer - .split_partition_queue(topic, *partition) - .map(PartitionQueueBackend::Rdkafka) - .context("Unable to split partition queue") - }) - .collect(), - Self::Arroyo { .. } => Ok(assignment - .arroyo_queues - .into_iter() - .map(PartitionQueueBackend::Arroyo) - .collect()), - } - } +pub fn create_consumer( + topics: &[&str], + config: &ClientConfig, + context: KafkaContext, +) -> Result { + let arroyo_config = KafkaConfig::new_config( + vec![ + config + .get("bootstrap.servers") + .context("Missing brokers")? + .to_owned(), + ], + Some(config.config_map().clone()), + ); + let topics: Vec<_> = topics.iter().map(|topic| Topic::new(topic)).collect(); + AsyncKafkaConsumer::new(arroyo_config, &topics, context) + .context("Could not create Arroyo consumer") } -impl CommitClient for ConsumerBackend { +impl CommitClient for AsyncKafkaConsumer { fn store_offsets(&self, tpl: &TopicPartitionList) -> Result<(), Error> { - match self { - Self::Rdkafka { consumer } => consumer.store_offsets(tpl).map_err(Error::from), - Self::Arroyo { consumer } => { - let mut offsets = HashMap::with_capacity(tpl.count()); - for element in tpl.elements() { - let offset = element - .offset() - .to_raw() - .and_then(|offset| u64::try_from(offset).ok()) - .context("Cannot stage an invalid offset")?; - offsets.insert( - arroyo_partition(element.topic(), element.partition())?, - offset, - ); - } - consumer - .store_offsets(offsets) - .context("Arroyo consumer offset storage failed") - } + let mut offsets = HashMap::with_capacity(tpl.count()); + for element in tpl.elements() { + let offset = element + .offset() + .to_raw() + .and_then(|offset| u64::try_from(offset).ok()) + .context("Cannot stage an invalid offset")?; + offsets.insert( + arroyo_partition(element.topic(), element.partition())?, + offset, + ); } + // The inherent method, not this trait method. + AsyncKafkaConsumer::store_offsets(self, offsets) + .context("Arroyo consumer offset storage failed") } } @@ -124,33 +63,36 @@ fn arroyo_partition(topic: &str, partition: i32) -> Result { )) } -pub enum PartitionQueueBackend { - Rdkafka(StreamPartitionQueue), - Arroyo(AsyncPartitionQueue), -} - -impl MessageQueue for PartitionQueueBackend { +impl MessageQueue for AsyncPartitionQueue { fn stream(&self) -> impl Stream { - match self { - Self::Rdkafka(queue) => MessageQueue::stream(queue) - .map(KafkaMessage::into_message) - .left_stream(), - Self::Arroyo(queue) => futures::stream::unfold(queue, |queue| async move { - let message = queue - .recv() - .await? - .map(MessageBackend::from) - .context("Arroyo partition receive failed"); - Some((message, queue)) - }) - .right_stream(), - } + futures::stream::unfold(self, |queue| async move { + let message = queue + .recv() + .await? + .map(MessageBackend::from) + .context("Arroyo partition receive failed"); + Some((message, queue)) + }) } } +#[derive(Debug)] +pub struct KafkaContext { + event_sender: UnboundedSender<(Event, SyncSender<()>)>, + /// The topic(s) this consumer is subscribed to, used as a metric tag. + topics_tag: String, +} + impl KafkaContext { - /// Arroyo's counterpart to `pre_rebalance`: blocks until `handle_events` has processed - /// the event. + pub fn new(event_sender: UnboundedSender<(Event, SyncSender<()>)>, topics_tag: String) -> Self { + Self { + event_sender, + topics_tag, + } + } + + /// Called from Arroyo's assignment callbacks: blocks until `handle_events` has processed the + /// event. fn handle_rebalance(&self, event: Event) { let (partition_count, metric) = match &event { Event::Assign(assignment) => ( @@ -192,10 +134,7 @@ impl AsyncAssignmentCallbacks for KafkaContext { ) }) .collect(); - self.handle_rebalance(Event::Assign(Assignment { - partitions, - arroyo_queues: queues, - })); + self.handle_rebalance(Event::Assign(Assignment { partitions, queues })); } fn on_revoke(&self, partitions: Vec) { @@ -211,17 +150,16 @@ impl AsyncAssignmentCallbacks for KafkaContext { self.handle_rebalance(Event::Revoke(partitions)); } - /// Arroyo has already logged the error. Shut down like the rdkafka client does on - /// unexpected consumer errors. + /// Arroyo has already logged the error. Shut down on unexpected consumer errors. fn on_error(&self, _error: KafkaError) { drop(elegant_departure::shutdown()); } } -/// The assigned partitions, with their queues when Arroyo has already split them. +/// The assigned partitions and their queues. pub struct Assignment { pub partitions: BTreeSet<(String, i32)>, - pub(super) arroyo_queues: Vec, + pub queues: Vec, } impl Debug for Assignment { @@ -241,28 +179,23 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use rdkafka::consumer::{BaseConsumer, CommitMode, Consumer}; + use rdkafka::consumer::{BaseConsumer, Consumer}; use rdkafka::mocking::MockCluster; use rdkafka::producer::{FutureProducer, FutureRecord}; use rdkafka::{ClientConfig, Offset, TopicPartitionList}; - use tokio::runtime::Handle; - use tokio::sync::{mpsc, oneshot}; - use tokio::task; - use tokio::time::timeout; - use anyhow::{Context, Error, ensure}; use futures::{StreamExt, pin_mut}; + use tokio::sync::mpsc; + use tokio::task; + use tokio::time::timeout; - use crate::kafka::consumer::{CommitClient, Event, KafkaContext, KafkaMessage, MessageQueue}; + use crate::kafka::consumer::{CommitClient, Event, KafkaMessage, MessageQueue}; - use super::ConsumerBackend; + use super::{KafkaContext, create_consumer}; - #[rstest::rstest] - #[case::rdkafka(false)] - #[case::arroyo(true)] #[tokio::test] - async fn test_partition_receive_and_commit(#[case] use_arroyo: bool) { + async fn test_partition_receive_and_commit() { let cluster = MockCluster::new(1).unwrap(); cluster.create_topic("test", 1, 1).unwrap(); let mut config = ClientConfig::new(); @@ -286,39 +219,21 @@ mod tests { let (sender, mut events) = mpsc::unbounded_channel(); let consumer = Arc::new( - ConsumerBackend::new( + create_consumer( &["test"], &config, KafkaContext::new(sender, "test".to_owned()), - use_arroyo, ) .unwrap(), ); - let (shutdown_sender, shutdown) = oneshot::channel::<()>(); - let poller = task::spawn_blocking({ - let consumer = consumer.clone(); - move || { - Handle::current().block_on(async { - // Arroyo polls the consumer on its own thread. - let ConsumerBackend::Rdkafka { consumer } = consumer.as_ref() else { - return None; - }; - tokio::select! { - _ = shutdown => None, - message = consumer.recv() => Some(message.map(|message| message.detach())), - } - }) - } - }); let result = timeout(Duration::from_secs(30), async { let (event, rendezvous) = events.recv().await.context("Missing assignment")?; let Event::Assign(assignment) = event else { anyhow::bail!("Expected assignment, got {event:?}"); }; - let queues = consumer.partition_queues(assignment)?; drop(rendezvous); - ensure!(queues.len() == 1); - let stream = queues[0].stream(); + ensure!(assignment.queues.len() == 1); + let stream = assignment.queues[0].stream(); pin_mut!(stream); let message = stream .next() @@ -332,15 +247,9 @@ mod tests { let next_offset = Offset::Offset(message.offset() + 1); let mut offsets = TopicPartitionList::new(); offsets.add_partition_offset("test", 0, next_offset)?; - consumer.store_offsets(&offsets)?; + CommitClient::store_offsets(consumer.as_ref(), &offsets)?; let committer = consumer.clone(); - task::spawn_blocking(move || match committer.as_ref() { - ConsumerBackend::Rdkafka { consumer } => { - consumer.commit_consumer_state(CommitMode::Sync) - } - ConsumerBackend::Arroyo { consumer } => consumer.commit_consumer_state(), - }) - .await??; + task::spawn_blocking(move || committer.commit_consumer_state()).await??; task::spawn_blocking(move || -> Result<(), Error> { let observer: BaseConsumer = config.create()?; let committed = observer.committed_offsets(offsets, Duration::from_secs(5))?; @@ -351,14 +260,9 @@ mod tests { }) .await; + // Dropping the consumer revokes its partitions, which would wait on `events` forever. drop(events); - drop(shutdown_sender); - let poll_result = poller.await.unwrap(); task::spawn_blocking(move || drop(consumer)).await.unwrap(); - assert!( - poll_result.is_none(), - "Unexpected main queue result: {poll_result:?}" - ); result.unwrap().unwrap(); } @@ -374,14 +278,12 @@ mod tests { .set("bootstrap.servers", "127.0.0.1:1") .set("group.id", "test-offset-reset-aliases") .set("auto.offset.reset", policy); - for use_arroyo in [false, true] { - let (events, _) = mpsc::unbounded_channel(); - let context = KafkaContext::new(events, "test".to_owned()); - let consumer = ConsumerBackend::new(&["test"], &config, context, use_arroyo) - .unwrap_or_else(|error| panic!("{policy}, use_arroyo={use_arroyo}: {error}")); - tokio::task::spawn_blocking(move || drop(consumer)) - .await - .unwrap(); - } + let (events, _) = mpsc::unbounded_channel(); + let context = KafkaContext::new(events, "test".to_owned()); + let consumer = create_consumer(&["test"], &config, context) + .unwrap_or_else(|error| panic!("{policy}: {error}")); + tokio::task::spawn_blocking(move || drop(consumer)) + .await + .unwrap(); } } diff --git a/src/main.rs b/src/main.rs index fb324055..30436b4c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -208,7 +208,6 @@ async fn main() -> Result<(), Error> { start_consumer( &topic_refs, &consumer_config.kafka_consumer_config_for(&task_topic), - consumer_config.use_arroyo_consumer, consumer_store.clone(), processing_strategy!({ err: