From 2e4ed3ec43a8fb65023046ecfd61b97ba8d491ef Mon Sep 17 00:00:00 2001 From: Tony Le Date: Thu, 8 Oct 2026 05:05:21 -0400 Subject: [PATCH] feat(kafka): Add Arroyo consumer backend --- .github/workflows/ci.yml | 64 +++++- src/config/mod.rs | 7 + src/kafka/consumer.rs | 153 +++++++++----- src/kafka/consumer_backend.rs | 387 ++++++++++++++++++++++++++++++++++ src/kafka/mod.rs | 1 + src/main.rs | 1 + 6 files changed, 554 insertions(+), 59 deletions(-) create mode 100644 src/kafka/consumer_backend.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 21af567e..01a2025f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -266,8 +266,14 @@ jobs: token: ${{ secrets.CODECOV_TOKEN }} rebalance-integration-test: - name: Rebalance integration test + name: Rebalance integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + consumer_backend: [rdkafka, arroyo] + env: + TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -304,8 +310,14 @@ jobs: make test-rebalance worker-processing-integration-test: - name: Worker processing integration test + name: Worker processing integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + consumer_backend: [rdkafka, arroyo] + env: + TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -342,8 +354,14 @@ jobs: make test-worker-processing upkeep-retry-integration-test: - name: Upkeep retry integration test + name: Upkeep retry integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + consumer_backend: [rdkafka, arroyo] + env: + TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -380,8 +398,14 @@ jobs: make test-upkeep-retry upkeep-expiry-integration-test: - name: Upkeep expiry integration test + name: Upkeep expiry integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + consumer_backend: [rdkafka, arroyo] + env: + TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -418,8 +442,14 @@ jobs: make test-upkeep-expiry upkeep-delay-integration-test: - name: Upkeep delay integration test + name: Upkeep delay integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + consumer_backend: [rdkafka, arroyo] + env: + TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -456,8 +486,14 @@ jobs: make test-upkeep-delay failed-tasks-integration-test: - name: Failed tasks integration test + name: Failed tasks integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + consumer_backend: [rdkafka, arroyo] + env: + TASKBROKER_USE_ARROYO_CONSUMER: ${{ matrix.consumer_backend == 'arroyo' }} steps: - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v2 @@ -494,8 +530,14 @@ jobs: make test-failed-tasks multi-topic-integration-test: - name: Multi-topic integration test + name: Multi-topic integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + 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,8 +574,14 @@ jobs: make test-multi-topic push-postgres-integration-test: - name: Push mode postgres integration test + name: Push mode postgres integration test${{ matrix.consumer_backend == 'arroyo' && ' (arroyo)' || '' }} runs-on: ubuntu-latest + 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 3c5289f3..ab6f37fb 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -151,6 +151,9 @@ 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, @@ -297,6 +300,7 @@ 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, @@ -1042,6 +1046,7 @@ 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] @@ -1236,6 +1241,7 @@ 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, @@ -1246,6 +1252,7 @@ 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 75480515..520c8b7a 100644 --- a/src/kafka/consumer.rs +++ b/src/kafka/consumer.rs @@ -31,6 +31,7 @@ use tracing::{debug, error, info, instrument, warn}; use crate::store::traits::ActivationStore; use crate::store::types::TopicPartition; +use super::consumer_backend::{Assignment, ConsumerBackend}; use super::message::MessageBackend; pub async fn start_no_consume_mode( @@ -92,11 +93,9 @@ 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>, - &BTreeSet<(String, i32)>, - ) -> 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(); @@ -104,18 +103,19 @@ pub async fn start_consumer( // 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> = Arc::new( - kafka_client_config - .create_with_context(context) - .expect("Consumer creation failed"), - ); - - consumer - .subscribe(topics) - .expect("Can't subscribe to specified topics"); + let consumer = Arc::new(ConsumerBackend::new( + topics, + kafka_client_config, + context, + use_arroyo, + )?); handle_shutdown_signals(event_sender.clone()); - poll_consumer_client(consumer.clone(), client_shutdown_receiver); + // 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(), @@ -123,14 +123,18 @@ pub async fn start_consumer( ) .set(0); handle_events( - consumer, + consumer.clone(), event_receiver, activation_store, client_shutdown_sender, spawn_actors, topics_tag, ) - .await + .await?; + // Dropping the consumer blocks while it closes, and Arroyo also makes its final commit then, + // so keep it off the async workers. + task::spawn_blocking(move || drop(consumer)).await?; + Ok(()) } pub fn handle_shutdown_signals(event_sender: UnboundedSender<(Event, SyncSender<()>)>) { @@ -176,9 +180,9 @@ pub fn poll_consumer_client( #[derive(Debug)] pub struct KafkaContext { - event_sender: UnboundedSender<(Event, SyncSender<()>)>, + pub(super) event_sender: UnboundedSender<(Event, SyncSender<()>)>, /// The topic(s) this consumer is subscribed to, used as a metric tag. - topics_tag: String, + pub(super) topics_tag: String, } impl KafkaContext { @@ -207,7 +211,10 @@ impl ConsumerContext for KafkaContext { return; } let _ = self.event_sender.send(( - Event::Assign(tpl.to_topic_map().keys().cloned().collect()), + 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..."); @@ -258,7 +265,7 @@ impl ConsumerContext for KafkaContext { #[derive(Debug)] pub enum Event { - Assign(BTreeSet<(String, i32)>), + Assign(Assignment), Revoke(BTreeSet<(String, i32)>), Shutdown, } @@ -272,7 +279,7 @@ pub struct ActorHandles { impl ActorHandles { #[instrument(skip(self))] - async fn shutdown(mut self, deadline: Duration) { + async fn shutdown(mut self, deadline: Duration) -> JoinSet> { debug!("Signaling shutdown to actors..."); self.shutdown.cancel(); info!("Actor shutdown signaled, waiting for rendezvous..."); @@ -286,9 +293,11 @@ impl ActorHandles { "Unable to rendezvous within callback deadline, \ aborting all tasks within JoinSet" ); - self.join_set.abort_all(); } } + // Returned so the Arroyo path can join them before committing or closing the consumer. + self.join_set.abort_all(); + self.join_set } async fn join_next(&mut self) -> Option, JoinError>> { @@ -342,8 +351,8 @@ macro_rules! processing_strategy { reduce: $reduce_first:expr $(,$reduce_rest:expr)*, } ) => {{ - |consumer: Arc>, - tpl: &std::collections::BTreeSet<(String, i32)>| + |consumer: Arc<$crate::kafka::consumer_backend::ConsumerBackend>, + assignment: $crate::kafka::consumer_backend::Assignment| -> $crate::kafka::consumer::ActorHandles { let start = std::time::Instant::now(); @@ -355,11 +364,10 @@ macro_rules! processing_strategy { let (map_sender, reduce_receiver) = tokio::sync::mpsc::channel(1); let (err_sender, err_receiver) = tokio::sync::mpsc::channel(1); - for (topic, partition) in tpl.iter() { - let queue = consumer - .split_partition_queue(topic, *partition) - .expect("Unable to split topic by Partition"); - + let queues = consumer + .partition_queues(assignment) + .expect("Unable to split topic by Partition"); + for queue in queues { handles.spawn($crate::kafka::consumer::map( queue, $map_fn, @@ -410,14 +418,11 @@ 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>, - &BTreeSet<(String, i32)>, - ) -> ActorHandles, + mut spawn_actors: impl FnMut(Arc, Assignment) -> ActorHandles, topics_tag: String, ) -> Result<(), anyhow::Error> { const CALLBACK_DURATION: Duration = Duration::from_secs(4); @@ -444,7 +449,8 @@ pub async fn handle_events( }; info!("Received event: {:?}", event); state = match (state, event) { - (ConsumerState::Ready, Event::Assign(tpl)) => { + (ConsumerState::Ready, Event::Assign(assignment)) => { + let tpl = assignment.partitions.clone(); metrics::gauge!( "arroyo.consumer.current_partitions", "topic" => topics_tag.clone(), @@ -452,7 +458,7 @@ pub async fn handle_events( ) .set(tpl.len() as f64); activation_store.assign_partitions(&mut tpl.iter().map(TopicPartition::from)); - ConsumerState::Consuming(spawn_actors(consumer.clone(), &tpl), tpl) + ConsumerState::Consuming(spawn_actors(consumer.clone(), assignment), tpl) } (ConsumerState::Ready, Event::Revoke(_)) => { unreachable!("Got partition revocation before the consumer has started") @@ -467,7 +473,12 @@ 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)); - handles.shutdown(CALLBACK_DURATION).await; + 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; + } metrics::gauge!( "arroyo.consumer.current_partitions", "topic" => topics_tag.clone(), @@ -478,7 +489,10 @@ pub async fn handle_events( } (ConsumerState::Consuming(handles, tpl), Event::Shutdown) => { activation_store.revoke_partitions(&mut tpl.iter().map(TopicPartition::from)); - handles.shutdown(CALLBACK_DURATION).await; + let mut actors = handles.shutdown(CALLBACK_DURATION).await; + if consumer.is_arroyo() { + actors.shutdown().await; + } debug!("Signaling shutdown to client..."); shutdown_client.take(); metrics::gauge!( @@ -834,13 +848,7 @@ pub async fn reduce_err( } pub trait CommitClient { - fn store_offsets(&self, tpl: &TopicPartitionList) -> KafkaResult<()>; -} - -impl CommitClient for StreamConsumer { - fn store_offsets(&self, tpl: &TopicPartitionList) -> KafkaResult<()> { - Consumer::store_offsets(self, tpl) - } + fn store_offsets(&self, tpl: &TopicPartitionList) -> Result<(), Error>; } #[derive(Default)] @@ -904,20 +912,21 @@ mod tests { use std::time::Duration; use anyhow::{Error, anyhow}; - use futures::Stream; - use rdkafka::error::{KafkaError, KafkaResult}; + use futures::{Stream, future}; + use rdkafka::error::KafkaError; use rdkafka::message::OwnedMessage; use rdkafka::{Offset, Timestamp, TopicPartitionList}; use tokio::sync::{broadcast, mpsc, oneshot}; + use tokio::task::JoinSet; use tokio::time::sleep; use tokio_stream::wrappers::BroadcastStream; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tokio_util::sync::CancellationToken; use crate::kafka::consumer::{ - CommitClient, KafkaMessage, MessageQueue, ReduceConfig, ReduceShutdownBehaviour, - ReduceShutdownCondition, Reducer, ReducerWhenFullBehaviour, commit, map, reduce, - reduce_err, + ActorHandles, CommitClient, KafkaMessage, MessageQueue, ReduceConfig, + ReduceShutdownBehaviour, ReduceShutdownCondition, Reducer, ReducerWhenFullBehaviour, + commit, map, reduce, reduce_err, }; use crate::kafka::message::MessageBackend; use crate::kafka::os_stream_writer::{OsStream, OsStreamWriter}; @@ -927,7 +936,7 @@ mod tests { } impl CommitClient for MockCommitClient { - fn store_offsets(&self, tpl: &TopicPartitionList) -> KafkaResult<()> { + fn store_offsets(&self, tpl: &TopicPartitionList) -> Result<(), Error> { self.offsets.write().unwrap().push(tpl.clone()); Ok(()) } @@ -1088,6 +1097,48 @@ mod tests { } } + #[rstest::rstest] + #[case::rendezvous(true)] + #[case::deadline(false)] + #[tokio::test] + async fn test_native_actor_shutdown_defers_cleanup(#[case] rendezvous_completed: bool) { + let consumer = Arc::new(()); + let mut join_set = JoinSet::new(); + let (started_sender, started) = oneshot::channel(); + join_set.spawn({ + let consumer = consumer.clone(); + async move { + let _consumer = consumer; + started_sender.send(()).unwrap(); + future::pending().await + } + }); + started.await.unwrap(); + assert_eq!(Arc::strong_count(&consumer), 2); + + let (rendezvous_sender, rendezvous) = oneshot::channel(); + let _rendezvous_guard = if rendezvous_completed { + drop(rendezvous_sender); + None + } else { + Some(rendezvous_sender) + }; + let handles = ActorHandles { + join_set, + shutdown: CancellationToken::new(), + rendezvous, + }; + let deadline = if rendezvous_completed { + Duration::from_secs(1) + } else { + Duration::ZERO + }; + let mut cleanup = handles.shutdown(deadline).await; + assert_eq!(Arc::strong_count(&consumer), 2); + cleanup.shutdown().await; + assert_eq!(Arc::strong_count(&consumer), 1); + } + #[tokio::test] async fn test_commit() { let offsets = Arc::new(RwLock::new(Vec::new())); diff --git a/src/kafka/consumer_backend.rs b/src/kafka/consumer_backend.rs new file mode 100644 index 00000000..a5817396 --- /dev/null +++ b/src/kafka/consumer_backend.rs @@ -0,0 +1,387 @@ +use std::collections::{BTreeSet, HashMap}; +use std::fmt::Debug; +use std::sync::Arc; +use std::sync::mpsc::sync_channel; + +use rdkafka::consumer::stream_consumer::StreamPartitionQueue; +use rdkafka::consumer::{Consumer, StreamConsumer}; +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 tracing::{info, warn}; + +use super::consumer::{CommitClient, Event, KafkaContext, 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()), + } + } +} + +impl CommitClient for ConsumerBackend { + 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") + } + } + } +} + +fn arroyo_partition(topic: &str, partition: i32) -> Result { + Ok(Partition::new( + Topic::new(topic), + u16::try_from(partition).context("Partition index does not fit Arroyo's partition type")?, + )) +} + +pub enum PartitionQueueBackend { + Rdkafka(StreamPartitionQueue), + Arroyo(AsyncPartitionQueue), +} + +impl MessageQueue for PartitionQueueBackend { + 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(), + } + } +} + +impl KafkaContext { + /// Arroyo's counterpart to `pre_rebalance`: blocks until `handle_events` has processed + /// the event. + fn handle_rebalance(&self, event: Event) { + let (partition_count, metric) = match &event { + Event::Assign(assignment) => ( + assignment.partitions.len(), + "arroyo.consumer.partitions_assigned.count", + ), + Event::Revoke(partitions) => { + (partitions.len(), "arroyo.consumer.partitions_revoked.count") + } + Event::Shutdown => unreachable!("Shutdown is not a rebalance event"), + }; + if partition_count == 0 { + warn!("Got rebalance event with no partitions"); + return; + } + let (sender, receiver) = sync_channel(0); + let _ = self.event_sender.send((event, sender)); + info!("Partition rebalance event sent, waiting for rendezvous..."); + let _ = receiver.recv(); + info!("Rendezvous complete"); + metrics::counter!( + metric, + "topic" => self.topics_tag.clone(), + "application" => "taskbroker", + ) + .increment(partition_count as u64); + } +} + +impl AsyncAssignmentCallbacks for KafkaContext { + fn on_assign(&self, queues: Vec) { + let partitions = queues + .iter() + .map(|queue| { + let partition = queue.partition(); + ( + partition.topic.as_str().to_owned(), + i32::from(partition.index), + ) + }) + .collect(); + self.handle_rebalance(Event::Assign(Assignment { + partitions, + arroyo_queues: queues, + })); + } + + fn on_revoke(&self, partitions: Vec) { + let partitions = partitions + .iter() + .map(|partition| { + ( + partition.topic.as_str().to_owned(), + i32::from(partition.index), + ) + }) + .collect(); + self.handle_rebalance(Event::Revoke(partitions)); + } + + /// Arroyo has already logged the error. Shut down like the rdkafka client does 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. +pub struct Assignment { + pub partitions: BTreeSet<(String, i32)>, + pub(super) arroyo_queues: Vec, +} + +impl Debug for Assignment { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("Assignment").field(&self.partitions).finish() + } +} + +impl KafkaMessage for Result { + fn into_message(self) -> Result { + self + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use rdkafka::consumer::{BaseConsumer, CommitMode, 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 crate::kafka::consumer::{CommitClient, Event, KafkaContext, KafkaMessage, MessageQueue}; + + use super::ConsumerBackend; + + #[rstest::rstest] + #[case::rdkafka(false)] + #[case::arroyo(true)] + #[tokio::test] + async fn test_partition_receive_and_commit(#[case] use_arroyo: bool) { + let cluster = MockCluster::new(1).unwrap(); + cluster.create_topic("test", 1, 1).unwrap(); + let mut config = ClientConfig::new(); + config + .set("bootstrap.servers", cluster.bootstrap_servers()) + .set("group.id", "test-consumer-backend") + .set("auto.offset.reset", "earliest") + .set("enable.auto.commit", "false") + .set("enable.auto.offset.store", "false"); + let producer: FutureProducer = ClientConfig::new() + .set("bootstrap.servers", cluster.bootstrap_servers()) + .create() + .unwrap(); + producer + .send( + FutureRecord::to("test").key("key").payload("payload"), + Duration::from_secs(5), + ) + .await + .unwrap(); + + let (sender, mut events) = mpsc::unbounded_channel(); + let consumer = Arc::new( + ConsumerBackend::new( + &["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(); + pin_mut!(stream); + let message = stream + .next() + .await + .context("Queue closed")? + .into_message()?; + ensure!(message.topic() == "test" && message.partition() == 0); + ensure!(message.key() == Some(b"key".as_slice())); + ensure!(message.payload() == Some(b"payload".as_slice())); + + 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)?; + 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 || -> Result<(), Error> { + let observer: BaseConsumer = config.create()?; + let committed = observer.committed_offsets(offsets, Duration::from_secs(5))?; + ensure!(committed.find_partition("test", 0).unwrap().offset() == next_offset); + Ok(()) + }) + .await? + }) + .await; + + 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(); + } + + #[rstest::rstest] + #[case::smallest("smallest")] + #[case::beginning("beginning")] + #[case::largest("largest")] + #[case::end("end")] + #[tokio::test] + async fn test_consumer_accepts_offset_reset_aliases(#[case] policy: &str) { + let mut config = ClientConfig::new(); + config + .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(); + } + } +} diff --git a/src/kafka/mod.rs b/src/kafka/mod.rs index 1ef184e2..54c8fb7a 100644 --- a/src/kafka/mod.rs +++ b/src/kafka/mod.rs @@ -2,6 +2,7 @@ pub mod activation_batcher; pub mod activation_writer; pub mod admin; pub mod consumer; +pub mod consumer_backend; pub mod deserialize; pub mod deserialize_activation; pub mod deserialize_raw; diff --git a/src/main.rs b/src/main.rs index 30436b4c..fb324055 100644 --- a/src/main.rs +++ b/src/main.rs @@ -208,6 +208,7 @@ 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: