ref(consumer): Add batch write timeout to Rust consumer (#6043)

* add new batch write timeout parameter

* linting + add main CLI entry arg

* linting

* linitng again

* fix linting again

* add CLI arg to rust consumer entrypoint

* add tests

* add snapshots ? and add constraint for timeout

* clean up timeout constraint

* remove snapshots?

* change test to should panic

rust_snuba/benches/processors.rs

@@ -97,6 +97,7 @@ fn create_factory(
             broker_config: BrokerConfig::default(),
         },
         stop_at_timestamp: None,
+        batch_write_timeout: None,
     };
     Box::new(factory)
 }

rust_snuba/src/consumer.rs

@@ -41,6 +41,7 @@ pub fn consumer(
     python_max_queue_depth: Option<usize>,
     health_check_file: Option<&str>,
     stop_at_timestamp: Option<i64>,
+    batch_write_timeout_ms: Option<u64>,
 ) {
     py.allow_threads(|| {
         consumer_impl(
@@ -57,6 +58,7 @@ pub fn consumer(
             python_max_queue_depth,
             health_check_file,
             stop_at_timestamp,
+            batch_write_timeout_ms,
         )
     });
 }
@@ -76,13 +78,25 @@ pub fn consumer_impl(
     python_max_queue_depth: Option<usize>,
     health_check_file: Option<&str>,
     stop_at_timestamp: Option<i64>,
+    batch_write_timeout_ms: Option<u64>,
 ) -> usize {
     setup_logging();
 
     let consumer_config = config::ConsumerConfig::load_from_str(consumer_config_raw).unwrap();
     let max_batch_size = consumer_config.max_batch_size;
     let max_batch_time = Duration::from_millis(consumer_config.max_batch_time_ms);
 
+    let batch_write_timeout = match batch_write_timeout_ms {
+        Some(timeout_ms) => {
+            if timeout_ms >= consumer_config.max_batch_time_ms {
+                Some(Duration::from_millis(timeout_ms))
+            } else {
+                None
+            }
+        }
+        None => None,
+    };
+
     for storage in &consumer_config.storages {
         tracing::info!(
             "Storage: {}, ClickHouse Table Name: {}, Message Processor: {:?}, ClickHouse host: {}, ClickHouse port: {}, ClickHouse HTTP port: {}, ClickHouse database: {}",
@@ -232,6 +246,7 @@ pub fn consumer_impl(
         physical_topic_name: Topic::new(&consumer_config.raw_topic.physical_topic_name),
         accountant_topic_config: consumer_config.accountant_topic,
         stop_at_timestamp,
+        batch_write_timeout,
     };
 
     let topic = Topic::new(&consumer_config.raw_topic.physical_topic_name);

rust_snuba/src/factory.rs

@@ -56,6 +56,7 @@ pub struct ConsumerStrategyFactory {
     pub physical_topic_name: Topic,
     pub accountant_topic_config: config::TopicConfig,
     pub stop_at_timestamp: Option<i64>,
+    pub batch_write_timeout: Option<Duration>,
 }
 
 impl ProcessingStrategyFactory<KafkaPayload> for ConsumerStrategyFactory {
@@ -118,6 +119,7 @@ impl ProcessingStrategyFactory<KafkaPayload> for ConsumerStrategyFactory {
             &self.storage_config.clickhouse_cluster.user,
             &self.storage_config.clickhouse_cluster.password,
             self.async_inserts,
+            self.batch_write_timeout,
         );
 
         let accumulator = Arc::new(

rust_snuba/src/strategies/clickhouse/batch.rs

@@ -42,6 +42,7 @@ impl BatchFactory {
         clickhouse_user: &str,
         clickhouse_password: &str,
         async_inserts: bool,
+        batch_write_timeout: Option<Duration>,
     ) -> Self {
         let mut headers = HeaderMap::with_capacity(5);
         headers.insert(CONNECTION, HeaderValue::from_static("keep-alive"));
@@ -72,10 +73,18 @@ impl BatchFactory {
         let url = format!("http://{hostname}:{http_port}?{query_params}");
         let query = format!("INSERT INTO {table} FORMAT JSONEachRow");
 
-        let client = ClientBuilder::new()
-            .default_headers(headers)
-            .build()
-            .unwrap();
+        let client = if let Some(timeout_duration) = batch_write_timeout {
+            ClientBuilder::new()
+                .default_headers(headers)
+                .timeout(timeout_duration)
+                .build()
+                .unwrap()
+        } else {
+            ClientBuilder::new()
+                .default_headers(headers)
+                .build()
+                .unwrap()
+        };
 
         BatchFactory {
             client,
@@ -252,6 +261,7 @@ mod tests {
             "default",
             "",
             false,
+            None,
         );
 
         let mut batch = factory.new_batch();
@@ -286,6 +296,7 @@ mod tests {
             "default",
             "",
             true,
+            None,
         );
 
         let mut batch = factory.new_batch();
@@ -319,6 +330,7 @@ mod tests {
             "default",
             "",
             false,
+            None,
         );
 
         let mut batch = factory.new_batch();
@@ -350,6 +362,7 @@ mod tests {
             "default",
             "",
             false,
+            None,
         );
 
         let mut batch = factory.new_batch();
@@ -366,4 +379,72 @@ mod tests {
         // ensure there has not been any HTTP request
         mock.assert_hits(0);
     }
+
+    #[test]
+    #[should_panic]
+    fn test_write_no_time() {
+        crate::testutils::initialize_python();
+        let server = MockServer::start();
+
+        let concurrency = ConcurrencyConfig::new(1);
+        let factory = BatchFactory::new(
+            &server.host(),
+            server.port(),
+            "testtable",
+            "testdb",
+            &concurrency,
+            "default",
+            "",
+            true,
+            // pass in an unreasonably short timeout
+            // which prevents the client request from reaching Clickhouse
+            Some(Duration::from_millis(0)),
+        );
+
+        let mut batch = factory.new_batch();
+
+        batch
+            .write_rows(&RowData::from_encoded_rows(vec![
+                br#"{"hello": "world"}"#.to_vec()
+            ]))
+            .unwrap();
+
+        concurrency.handle().block_on(batch.finish()).unwrap();
+    }
+
+    #[test]
+    fn test_write_enough_time() {
+        crate::testutils::initialize_python();
+        let server = MockServer::start();
+        let mock = server.mock(|when, then| {
+            when.method(POST).path("/").body("{\"hello\": \"world\"}\n");
+            then.status(200).body("hi");
+        });
+
+        let concurrency = ConcurrencyConfig::new(1);
+        let factory = BatchFactory::new(
+            &server.host(),
+            server.port(),
+            "testtable",
+            "testdb",
+            &concurrency,
+            "default",
+            "",
+            true,
+            // pass in a reasonable timeout
+            Some(Duration::from_millis(1000)),
+        );
+
+        let mut batch = factory.new_batch();
+
+        batch
+            .write_rows(&RowData::from_encoded_rows(vec![
+                br#"{"hello": "world"}"#.to_vec()
+            ]))
+            .unwrap();
+
+        concurrency.handle().block_on(batch.finish()).unwrap();
+
+        mock.assert();
+    }
 }

snuba/cli/rust_consumer.py

@@ -158,6 +158,12 @@
     type=int,
     help="Unix timestamp after which to stop processing messages",
 )
+@click.option(
+    "--batch-write-timeout-ms",
+    type=int,
+    default=None,
+    help="Optional timeout for batch writer client connecting and sending request to Clickhouse",
+)
 def rust_consumer(
     *,
     storage_names: Sequence[str],
@@ -186,6 +192,7 @@ def rust_consumer(
     health_check_file: Optional[str],
     enforce_schema: bool,
     stop_at_timestamp: Optional[int],
+    batch_write_timeout_ms: Optional[int]
 ) -> None:
     """
     Experimental alternative to `snuba consumer`
@@ -235,6 +242,7 @@ def rust_consumer(
         python_max_queue_depth,
         health_check_file,
         stop_at_timestamp,
+        batch_write_timeout_ms,
     )
 
     sys.exit(exitcode)