Fix BLQ concurrency lifetime and deduplicate router setup

Author: cursoragent

rust_snuba/benches/processors.rs

@@ -103,6 +103,7 @@ fn create_factory(
         join_timeout_ms: None,
         health_check: "arroyo".to_string(),
         use_row_binary: false,
+        blq_concurrency: ConcurrencyConfig::with_runtime(concurrency, RUNTIME.handle().to_owned()),
         blq_producer_config: None,
         blq_topic: None,
     };

rust_snuba/src/consumer.rs

@@ -285,6 +285,7 @@ pub fn consumer_impl(
         join_timeout_ms,
         health_check: health_check.to_string(),
         use_row_binary,
+        blq_concurrency: ConcurrencyConfig::new(10),
         blq_producer_config: dlq_producer_config.clone(),
         blq_topic: dlq_topic,
     };

rust_snuba/src/factory_v2.rs

@@ -69,6 +69,7 @@ pub struct ConsumerStrategyFactoryV2 {
     pub join_timeout_ms: Option<u64>,
     pub health_check: String,
     pub use_row_binary: bool,
+    pub blq_concurrency: ConcurrencyConfig,
     pub blq_producer_config: Option<KafkaConfig>,
     pub blq_topic: Option<Topic>,
 }
@@ -267,43 +268,7 @@ impl ProcessingStrategyFactory<KafkaPayload> for ConsumerStrategyFactoryV2 {
             Some(Duration::from_millis(self.join_timeout_ms.unwrap_or(0))),
         );
 
-        let blq_enabled_flag = options("snuba")
-            .ok()
-            .and_then(|o| o.get("consumer.blq_enabled").ok())
-            .and_then(|v| v.as_bool())
-            .unwrap_or(false);
-        let next_step: Box<dyn ProcessingStrategy<KafkaPayload>> =
-            if let (true, Some(blq_producer_config), Some(blq_topic)) =
-                (blq_enabled_flag, &self.blq_producer_config, self.blq_topic)
-            {
-                let stale_threshold = TimeDelta::minutes(30);
-                let static_friction = TimeDelta::minutes(2);
-                tracing::info!(
-                "Routing all messages older than {:?} to the topic {:?} with static_friction {:?}",
-                stale_threshold,
-                self.blq_topic,
-                static_friction
-            );
-                let concurrency = ConcurrencyConfig::new(10);
-                let blq_producer = Produce::new(
-                    CommitOffsets::new(Duration::from_millis(250)),
-                    KafkaProducer::new(blq_producer_config.clone()),
-                    &concurrency,
-                    TopicOrPartition::Topic(blq_topic),
-                );
-                Box::new(
-                    BLQRouter::new(
-                        next_step,
-                        blq_producer,
-                        stale_threshold,
-                        Some(static_friction),
-                    )
-                    .expect("invalid BLQRouter config"),
-                )
-            } else {
-                tracing::info!("Not using a backlog-queue",);
-                Box::new(next_step)
-            };
+        let next_step = self.wrap_with_blq_router(Box::new(next_step));
 
         if let Some(path) = &self.health_check_file {
             {
@@ -324,6 +289,49 @@ impl ProcessingStrategyFactory<KafkaPayload> for ConsumerStrategyFactoryV2 {
 }
 
 impl ConsumerStrategyFactoryV2 {
+    fn wrap_with_blq_router(
+        &self,
+        next_step: Box<dyn ProcessingStrategy<KafkaPayload>>,
+    ) -> Box<dyn ProcessingStrategy<KafkaPayload>> {
+        let blq_enabled_flag = options("snuba")
+            .ok()
+            .and_then(|o| o.get("consumer.blq_enabled").ok())
+            .and_then(|v| v.as_bool())
+            .unwrap_or(false);
+
+        if let (true, Some(blq_producer_config), Some(blq_topic)) =
+            (blq_enabled_flag, &self.blq_producer_config, self.blq_topic)
+        {
+            let stale_threshold = TimeDelta::minutes(30);
+            let static_friction = TimeDelta::minutes(2);
+            tracing::info!(
+                "Routing all messages older than {:?} to the topic {:?} with static_friction {:?}",
+                stale_threshold,
+                self.blq_topic,
+                static_friction
+            );
+
+            let blq_producer = Produce::new(
+                CommitOffsets::new(Duration::from_millis(250)),
+                KafkaProducer::new(blq_producer_config.clone()),
+                &self.blq_concurrency,
+                TopicOrPartition::Topic(blq_topic),
+            );
+            Box::new(
+                BLQRouter::new(
+                    next_step,
+                    blq_producer,
+                    stale_threshold,
+                    Some(static_friction),
+                )
+                .expect("invalid BLQRouter config"),
+            )
+        } else {
+            tracing::info!("Not using a backlog-queue",);
+            next_step
+        }
+    }
+
     fn create_row_binary_pipeline<
         T: clickhouse::Row + serde::Serialize + Clone + Send + Sync + 'static,
     >(
@@ -420,43 +428,7 @@ impl ConsumerStrategyFactoryV2 {
             Some(Duration::from_millis(self.join_timeout_ms.unwrap_or(0))),
         );
 
-        let blq_enabled_flag = options("snuba")
-            .ok()
-            .and_then(|o| o.get("consumer.blq_enabled").ok())
-            .and_then(|v| v.as_bool())
-            .unwrap_or(false);
-        let next_step: Box<dyn ProcessingStrategy<KafkaPayload>> =
-            if let (true, Some(blq_producer_config), Some(blq_topic)) =
-                (blq_enabled_flag, &self.blq_producer_config, self.blq_topic)
-            {
-                let stale_threshold = TimeDelta::minutes(30);
-                let static_friction = TimeDelta::minutes(2);
-                tracing::info!(
-                "Routing all messages older than {:?} to the topic {:?} with static_friction {:?}",
-                stale_threshold,
-                self.blq_topic,
-                static_friction
-            );
-                let concurrency = ConcurrencyConfig::new(10);
-                let blq_producer = Produce::new(
-                    CommitOffsets::new(Duration::from_millis(250)),
-                    KafkaProducer::new(blq_producer_config.clone()),
-                    &concurrency,
-                    TopicOrPartition::Topic(blq_topic),
-                );
-                Box::new(
-                    BLQRouter::new(
-                        next_step,
-                        blq_producer,
-                        stale_threshold,
-                        Some(static_friction),
-                    )
-                    .expect("invalid BLQRouter config"),
-                )
-            } else {
-                tracing::info!("Not using a backlog-queue",);
-                Box::new(next_step)
-            };
+        let next_step = self.wrap_with_blq_router(Box::new(next_step));
 
         if let Some(path) = &self.health_check_file {
             if self.health_check == "snuba" {