working?

john-z-yang · john-z-yang · commit a20edb35544a · 2024-12-16T06:38:40.000-08:00
diff --git a/python/integration_tests/test_consumer_rebalancing.py b/python/integration_tests/test_consumer_rebalancing.py
@@ -210,8 +210,13 @@ def test_tasks_written_once_during_rebalancing() -> None:
             for log_line_index, line in enumerate(lines):
                 if "[31mERROR" in line:
                     # If there is an error in log file, capture 10 lines before and after the error line
-                    consumer_error_logs.append(f"Error found in consumer_{i}. Logging 10 lines before and after the error line:")
-                    for j in range(max(0, log_line_index - 10), min(len(lines) - 1, log_line_index + 10)):
+                    consumer_error_logs.append(
+                        f"Error found in consumer_{i}. Logging 10 lines before and after the error line:"
+                    )
+                    for j in range(
+                        max(0, log_line_index - 10),
+                        min(len(lines) - 1, log_line_index + 10),
+                    ):
                         consumer_error_logs.append(lines[j].strip())
                     consumer_error_logs.append("")
 
diff --git a/src/config.rs b/src/config.rs
@@ -136,6 +136,7 @@ impl Config {
                 "session.timeout.ms",
                 self.kafka_session_timeout_ms.to_string(),
             )
+            .set("partition.assignment.strategy", "cooperative-sticky")
             .set("enable.partition.eof", "false")
             .set("enable.auto.commit", "true")
             .set(
diff --git a/src/consumer/kafka.rs b/src/consumer/kafka.rs
@@ -32,7 +32,7 @@ use tokio::{
         mpsc::{self, unbounded_channel, UnboundedReceiver, UnboundedSender},
         oneshot,
     },
-    task::{self, JoinError, JoinSet},
+    task::{self, JoinError, JoinHandle, JoinSet},
     time::{self, sleep, MissedTickBehavior},
 };
 use tokio_stream::wrappers::UnboundedReceiverStream;
@@ -61,14 +61,16 @@ pub async fn start_consumer(
         .expect("Can't subscribe to specified topics");
 
     handle_os_signals(event_sender.clone());
-    poll_consumer_client(consumer.clone(), client_shutdown_receiver);
+    let rdkafka_driver = poll_consumer_client(consumer.clone(), client_shutdown_receiver);
     handle_events(
         consumer,
         event_receiver,
         client_shutdown_sender,
         spawn_actors,
     )
-    .await
+    .await?;
+    rdkafka_driver.await?;
+    Ok(())
 }
 
 pub fn handle_os_signals(event_sender: UnboundedSender<(Event, SyncSender<()>)>) {
@@ -85,23 +87,28 @@ pub fn handle_os_signals(event_sender: UnboundedSender<(Event, SyncSender<()>)>)
 pub fn poll_consumer_client(
     consumer: Arc<StreamConsumer<KafkaContext>>,
     shutdown: oneshot::Receiver<()>,
-) {
+) -> JoinHandle<()> {
     task::spawn_blocking(|| {
         Handle::current().block_on(async move {
             let _guard = elegant_departure::get_shutdown_guard().shutdown_on_drop();
             select! {
                 biased;
-                _ = shutdown => {
-                    debug!("Received shutdown signal, commiting state in sync mode...");
-                    let _ = consumer.commit_consumer_state(rdkafka::consumer::CommitMode::Sync);
-                }
+                _ = shutdown => {}
                 msg = consumer.recv() => {
                     error!("Got unexpected message from consumer client: {:?}", msg);
                 }
-            }
+
+            };
+
+            select! {
+                biased;
+                _ = consumer.recv() => {}
+                _ = sleep(Duration::from_secs(8)) => {}
+            };
+
             debug!("Shutdown complete");
         });
-    });
+    })
 }
 
 #[derive(Debug)]
@@ -118,8 +125,20 @@ impl KafkaContext {
 impl ClientContext for KafkaContext {}
 
 impl ConsumerContext for KafkaContext {
-    #[instrument(skip_all)]
-    fn pre_rebalance(&self, _: &BaseConsumer<Self>, rebalance: &Rebalance) {
+    #[instrument(skip(self, base_consumer))]
+    fn pre_rebalance(&self, base_consumer: &BaseConsumer<Self>, rebalance: &Rebalance) {
+        if let Rebalance::Assign(tpl) = rebalance {
+            if tpl.count() == 0 {
+                return;
+            }
+        }
+        base_consumer
+            .pause(
+                &base_consumer
+                    .assignment()
+                    .expect("Unable to fetch assigned TPL"),
+            )
+            .expect("Unable to pause consumer");
         let (rendezvous_sender, rendezvous_receiver) = sync_channel(0);
         match rebalance {
             Rebalance::Assign(tpl) => {
@@ -149,6 +168,31 @@ impl ConsumerContext for KafkaContext {
         }
     }
 
+    #[instrument(skip(self, base_consumer))]
+    fn post_rebalance(&self, base_consumer: &BaseConsumer<Self>, rebalance: &Rebalance) {
+        if let Rebalance::Assign(tpl) = rebalance {
+            if tpl.count() == 0 {
+                return;
+            }
+        }
+        let assignment = base_consumer
+            .assignment()
+            .expect("Failed to get assigned TPL");
+        if assignment.count() != 0 {
+            base_consumer
+                .seek_partitions(
+                    base_consumer
+                        .committed(rdkafka::util::Timeout::Never)
+                        .expect("Failed to get commited TPL"),
+                    rdkafka::util::Timeout::Never,
+                )
+                .expect("Failed to seek to commited offset");
+            base_consumer
+                .resume(&assignment)
+                .expect("Failed to resume consumer");
+        }
+    }
+
     #[instrument(skip(self))]
     fn commit_callback(&self, result: KafkaResult<()>, _offsets: &TopicPartitionList) {
         debug!("Got commit callback");
@@ -336,7 +380,7 @@ pub async fn handle_events(
 
     let mut state = ConsumerState::Ready;
 
-    while let ConsumerState::Ready { .. } | ConsumerState::Consuming { .. } = state {
+    while let ConsumerState::Ready | ConsumerState::Consuming { .. } = state {
         select! {
             res = match state {
                 ConsumerState::Consuming(ref mut handles, _) => Either::Left(handles.join_next()),
@@ -352,8 +396,8 @@ pub async fn handle_events(
                 };
                 info!("Received event: {:?}", event);
                 state = match (state, event) {
-                    (ConsumerState::Ready, Event::Assign(assigned)) => {
-                        ConsumerState::Consuming(spawn_actors(consumer.clone(), &assigned), assigned)
+                    (ConsumerState::Ready, Event::Assign(tpl)) => {
+                        ConsumerState::Consuming(spawn_actors(consumer.clone(), &tpl), tpl)
                     }
                     (ConsumerState::Ready, Event::Revoke(_)) => {
                         unreachable!("Got partition revocation before the consumer has started")
@@ -364,17 +408,17 @@ pub async fn handle_events(
                             tpl.is_disjoint(&assigned),
                             "Newly assigned TPL should be disjoint from TPL we're consuming from"
                         );
-                        tpl.append(&mut assigned);
                         debug!(
                             "{} additional topic partitions added after assignment",
                             assigned.len()
                         );
+                        tpl.append(&mut assigned);
                         handles.shutdown(CALLBACK_DURATION).await;
                         ConsumerState::Consuming(spawn_actors(consumer.clone(), &tpl), tpl)
                     }
                     (ConsumerState::Consuming(handles, mut tpl), Event::Revoke(revoked)) => {
                         assert!(
-                            tpl.is_subset(&revoked),
+                            revoked.is_subset(&tpl),
                             "Revoked TPL should be a subset of TPL we're consuming from"
                         );
                         tpl.retain(|e| !revoked.contains(e));
@@ -750,7 +794,7 @@ impl CommitClient for StreamConsumer<KafkaContext> {
     }
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 struct HighwaterMark {
     data: HashMap<(String, i32), i64>,
 }
@@ -795,6 +839,7 @@ pub async fn commit(
     while let Some(msgs) = receiver.recv().await {
         let mut highwater_mark = HighwaterMark::new();
         msgs.0.iter().for_each(|msg| highwater_mark.track(msg));
+        debug!("Store: {:?}", highwater_mark);
         consumer.store_offsets(&highwater_mark.into()).unwrap();
     }
     debug!("Shutdown complete");
diff --git a/src/main.rs b/src/main.rs
@@ -1,25 +1,17 @@
-use anyhow::{anyhow, Error};
+use anyhow::Error;
 use clap::Parser;
 use std::{sync::Arc, time::Duration};
-use taskbroker::grpc_middleware::MetricsLayer;
-use taskbroker::upkeep::upkeep;
-use tokio::select;
 use tokio::signal::unix::SignalKind;
 use tokio::task::JoinHandle;
-use tonic::transport::Server;
-use tonic_health::server::health_reporter;
 use tracing::{error, info};
 
-use sentry_protos::sentry::v1::consumer_service_server::ConsumerServiceServer;
-
 use taskbroker::config::Config;
 use taskbroker::consumer::{
     deserialize_activation::{self, DeserializeConfig},
     inflight_activation_writer::{ActivationWriterConfig, InflightActivationWriter},
     kafka::start_consumer,
     os_stream_writer::{OsStream, OsStreamWriter},
 };
-use taskbroker::grpc_server::MyConsumerService;
 use taskbroker::inflight_activation_store::InflightActivationStore;
 use taskbroker::logging;
 use taskbroker::metrics;
@@ -50,14 +42,14 @@ async fn main() -> Result<(), Error> {
     let store = Arc::new(InflightActivationStore::new(&config.db_path).await?);
 
     // Upkeep thread
-    let upkeep_task = tokio::spawn({
-        let upkeep_store = store.clone();
-        let upkeep_config = config.clone();
-        async move {
-            upkeep(upkeep_config, upkeep_store).await;
-            Ok(())
-        }
-    });
+    // let upkeep_task = tokio::spawn({
+    //     let upkeep_store = store.clone();
+    //     let upkeep_config = config.clone();
+    //     async move {
+    //         upkeep(upkeep_config, upkeep_store).await;
+    //         Ok(())
+    //     }
+    // });
 
     // Consumer from kafka
     let consumer_task = tokio::spawn({
@@ -89,66 +81,66 @@ async fn main() -> Result<(), Error> {
     });
 
     // GRPC server
-    let grpc_server_task = tokio::spawn({
-        let grpc_store = store.clone();
-        let grpc_config = config.clone();
-        async move {
-            let guard = elegant_departure::get_shutdown_guard().shutdown_on_drop();
-            let (mut health_reporter_fn, health_service) = health_reporter();
-            health_reporter_fn
-                .set_serving::<ConsumerServiceServer<MyConsumerService>>()
-                .await;
-            let addr = format!("{}:{}", grpc_config.grpc_addr, grpc_config.grpc_port)
-                .parse()
-                .expect("Failed to parse address");
-
-            let layers = tower::ServiceBuilder::new()
-                .layer(MetricsLayer::default())
-                .into_inner();
-
-            let server = Server::builder()
-                .layer(layers)
-                .add_service(ConsumerServiceServer::new(MyConsumerService {
-                    store: grpc_store,
-                }))
-                .add_service(health_service)
-                .serve(addr);
-
-            info!("GRPC server listening on {}", addr);
-            select! {
-                biased;
-
-                res = server => {
-                    info!("GRPC server task failed, shutting down");
-                    health_reporter_fn.set_not_serving::<ConsumerServiceServer<MyConsumerService>>().await;
-
-                    // Wait for any running requests to drain
-                    tokio::time::sleep(Duration::from_secs(5)).await;
-                    match res {
-                        Ok(()) => Ok(()),
-                        Err(e) => Err(anyhow!("GRPC server task failed: {:?}", e)),
-                    }
-                }
-                _ = guard.wait() => {
-                    info!("Cancellation token received, shutting down GRPC server");
-                    health_reporter_fn.set_not_serving::<ConsumerServiceServer<MyConsumerService>>().await;
-
-                    // Wait for any running requests to drain
-                    tokio::time::sleep(Duration::from_secs(5)).await;
-                    Ok(())
-                }
-            }
-        }
-    });
+    // let grpc_server_task = tokio::spawn({
+    //     let grpc_store = store.clone();
+    //     let grpc_config = config.clone();
+    //     async move {
+    //         let guard = elegant_departure::get_shutdown_guard().shutdown_on_drop();
+    //         let (mut health_reporter_fn, health_service) = health_reporter();
+    //         health_reporter_fn
+    //             .set_serving::<ConsumerServiceServer<MyConsumerService>>()
+    //             .await;
+    //         let addr = format!("{}:{}", grpc_config.grpc_addr, grpc_config.grpc_port)
+    //             .parse()
+    //             .expect("Failed to parse address");
+
+    //         let layers = tower::ServiceBuilder::new()
+    //             .layer(MetricsLayer::default())
+    //             .into_inner();
+
+    //         let server = Server::builder()
+    //             .layer(layers)
+    //             .add_service(ConsumerServiceServer::new(MyConsumerService {
+    //                 store: grpc_store,
+    //             }))
+    //             .add_service(health_service)
+    //             .serve(addr);
+
+    //         info!("GRPC server listening on {}", addr);
+    //         select! {
+    //             biased;
+
+    //             res = server => {
+    //                 info!("GRPC server task failed, shutting down");
+    //                 health_reporter_fn.set_not_serving::<ConsumerServiceServer<MyConsumerService>>().await;
+
+    //                 // Wait for any running requests to drain
+    //                 tokio::time::sleep(Duration::from_secs(5)).await;
+    //                 match res {
+    //                     Ok(()) => Ok(()),
+    //                     Err(e) => Err(anyhow!("GRPC server task failed: {:?}", e)),
+    //                 }
+    //             }
+    //             _ = guard.wait() => {
+    //                 info!("Cancellation token received, shutting down GRPC server");
+    //                 health_reporter_fn.set_not_serving::<ConsumerServiceServer<MyConsumerService>>().await;
+
+    //                 // Wait for any running requests to drain
+    //                 tokio::time::sleep(Duration::from_secs(5)).await;
+    //                 Ok(())
+    //             }
+    //         }
+    //     }
+    // });
 
     elegant_departure::tokio::depart()
         .on_termination()
         .on_sigint()
         .on_signal(SignalKind::hangup())
         .on_signal(SignalKind::quit())
         .on_completion(log_task_completion("consumer", consumer_task))
-        .on_completion(log_task_completion("grpc_server", grpc_server_task))
-        .on_completion(log_task_completion("upkeep_task", upkeep_task))
+        // .on_completion(log_task_completion("grpc_server", grpc_server_task))
+        // .on_completion(log_task_completion("upkeep_task", upkeep_task))
         .await;
 
     Ok(())

Original file line number	Diff line number	Diff line change
`@@ -136,6 +136,7 @@ impl Config {`
`136`	`136`	`"session.timeout.ms",`
`137`	`137`	`self.kafka_session_timeout_ms.to_string(),`
`138`	`138`	`)`
	`139`	`+ .set("partition.assignment.strategy", "cooperative-sticky")`
`139`	`140`	`.set("enable.partition.eof", "false")`
`140`	`141`	`.set("enable.auto.commit", "true")`
`141`	`142`	`.set(`