Merge pull request #165 from influxdata/dom/pre-warm

feat: support pre-warming PartitionClient

Author: domodwyer

benches/throughput.rs

@@ -448,7 +448,7 @@ async fn setup_rskafka(connection: Vec<String>) -> PartitionClient {
         .await
         .unwrap();
 
-    client.partition_client(topic_name, 0).unwrap()
+    client.partition_client(topic_name, 0).await.unwrap()
 }
 
 static LOG_SETUP: Once = Once::new();

src/client/mod.rs

@@ -116,16 +116,12 @@ impl Client {
     }
 
     /// Returns a client for performing operations on a specific partition
-    pub fn partition_client(
+    pub async fn partition_client(
         &self,
         topic: impl Into<String> + Send,
         partition: i32,
     ) -> Result<PartitionClient> {
-        Ok(PartitionClient::new(
-            topic.into(),
-            partition,
-            Arc::clone(&self.brokers),
-        ))
+        PartitionClient::new(topic.into(), partition, Arc::clone(&self.brokers)).await
     }
 
     /// Returns a list of topics in the cluster

src/client/partition.rs

@@ -22,8 +22,10 @@ use crate::{
     validation::ExactlyOne,
 };
 use async_trait::async_trait;
-use std::ops::{ControlFlow, Deref, Range};
-use std::sync::Arc;
+use std::{
+    ops::{ControlFlow, Deref, Range},
+    sync::Arc,
+};
 use time::OffsetDateTime;
 use tokio::sync::Mutex;
 use tracing::{error, info};
@@ -95,14 +97,35 @@ impl std::fmt::Debug for PartitionClient {
 }
 
 impl PartitionClient {
-    pub(super) fn new(topic: String, partition: i32, brokers: Arc<BrokerConnector>) -> Self {
-        Self {
+    pub(super) async fn new(
+        topic: String,
+        partition: i32,
+        brokers: Arc<BrokerConnector>,
+    ) -> Result<Self> {
+        let p = Self {
             topic,
             partition,
-            brokers,
+            brokers: Arc::clone(&brokers),
             backoff_config: Default::default(),
             current_broker: Mutex::new(None),
-        }
+        };
+
+        // Force discover and establish a cached connection to the leader
+        let scope = &p;
+        maybe_retry(
+            &Default::default(),
+            &*brokers,
+            "leader_detection",
+            || async move {
+                scope
+                    .get_leader(MetadataLookupMode::CachedArbitrary)
+                    .await?;
+                Ok(())
+            },
+        )
+        .await?;
+
+        Ok(p)
     }
 
     /// Topic

tests/client.rs

@@ -88,7 +88,10 @@ async fn test_partition_client() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(topic_name.clone(), 0).unwrap();
+    let partition_client = client
+        .partition_client(topic_name.clone(), 0)
+        .await
+        .unwrap();
     assert_eq!(partition_client.topic(), &topic_name);
     assert_eq!(partition_client.partition(), 0);
 }
@@ -162,7 +165,7 @@ async fn test_socks5() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(topic_name, 0).unwrap();
+    let partition_client = client.partition_client(topic_name, 0).await.unwrap();
 
     let record = record(b"");
     partition_client
@@ -194,7 +197,7 @@ async fn test_produce_empty() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(&topic_name, 1).unwrap();
+    let partition_client = client.partition_client(&topic_name, 1).await.unwrap();
     partition_client
         .produce(vec![], Compression::NoCompression)
         .await
@@ -216,7 +219,7 @@ async fn test_consume_empty() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(&topic_name, 1).unwrap();
+    let partition_client = client.partition_client(&topic_name, 1).await.unwrap();
     let (records, watermark) = partition_client
         .fetch_records(0, 1..10_000, 1_000)
         .await
@@ -240,7 +243,7 @@ async fn test_consume_offset_out_of_range() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(&topic_name, 1).unwrap();
+    let partition_client = client.partition_client(&topic_name, 1).await.unwrap();
     let record = record(b"");
     let offsets = partition_client
         .produce(vec![record], Compression::NoCompression)
@@ -279,7 +282,10 @@ async fn test_get_offset() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(topic_name.clone(), 0).unwrap();
+    let partition_client = client
+        .partition_client(topic_name.clone(), 0)
+        .await
+        .unwrap();
 
     assert_eq!(
         partition_client
@@ -340,7 +346,7 @@ async fn test_produce_consume_size_cutoff() {
         .await
         .unwrap();
 
-    let partition_client = Arc::new(client.partition_client(&topic_name, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic_name, 0).await.unwrap());
 
     let record_1 = large_record();
     let record_2 = large_record();
@@ -413,7 +419,7 @@ async fn test_consume_midbatch() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(&topic_name, 0).unwrap();
+    let partition_client = client.partition_client(&topic_name, 0).await.unwrap();
 
     // produce two records into a single batch
     let record_1 = record(b"x");
@@ -458,7 +464,7 @@ async fn test_delete_records() {
         .await
         .unwrap();
 
-    let partition_client = client.partition_client(&topic_name, 0).unwrap();
+    let partition_client = client.partition_client(&topic_name, 0).await.unwrap();
 
     // produce the following record batches:
     // - record_1

tests/consumer.rs

@@ -34,7 +34,7 @@ async fn test_stream_consumer_start_at_0() {
 
     let record = record(b"x");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
     partition_client
         .produce(vec![record.clone()], Compression::NoCompression)
         .await
@@ -85,7 +85,7 @@ async fn test_stream_consumer_start_at_1() {
     let record_1 = record(b"x");
     let record_2 = record(b"y");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
     partition_client
         .produce(
             vec![record_1.clone(), record_2.clone()],
@@ -121,7 +121,7 @@ async fn test_stream_consumer_offset_out_of_range() {
         .await
         .unwrap();
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
 
     let mut stream = StreamConsumerBuilder::new(partition_client, StartOffset::At(1)).build();
 
@@ -155,7 +155,7 @@ async fn test_stream_consumer_start_at_earliest() {
     let record_1 = record(b"x");
     let record_2 = record(b"y");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
     partition_client
         .produce(vec![record_1.clone()], Compression::NoCompression)
         .await
@@ -204,7 +204,7 @@ async fn test_stream_consumer_start_at_earliest_empty() {
 
     let record = record(b"x");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
 
     let mut stream =
         StreamConsumerBuilder::new(Arc::clone(&partition_client), StartOffset::Earliest)
@@ -245,7 +245,7 @@ async fn test_stream_consumer_start_at_earliest_after_deletion() {
     let record_1 = record(b"x");
     let record_2 = record(b"y");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
     partition_client
         .produce(
             vec![record_1.clone(), record_2.clone()],
@@ -287,7 +287,7 @@ async fn test_stream_consumer_start_at_latest() {
     let record_1 = record(b"x");
     let record_2 = record(b"y");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
     partition_client
         .produce(vec![record_1.clone()], Compression::NoCompression)
         .await
@@ -330,7 +330,7 @@ async fn test_stream_consumer_start_at_latest_empty() {
 
     let record = record(b"x");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
 
     let mut stream = StreamConsumerBuilder::new(Arc::clone(&partition_client), StartOffset::Latest)
         .with_max_wait_ms(50)

tests/produce_consume.rs

@@ -258,7 +258,12 @@ async fn assert_produce_consume<F1, G1, F2, G2>(
         .create_topic(&topic_name, n_partitions, 1, 5_000)
         .await
         .unwrap();
-    let partition_client = Arc::new(client.partition_client(topic_name.clone(), 1).unwrap());
+    let partition_client = Arc::new(
+        client
+            .partition_client(topic_name.clone(), 1)
+            .await
+            .unwrap(),
+    );
 
     // timestamps for records. We'll reorder the messages though to ts2, ts1, ts3
     let ts1 = now();

tests/producer.rs

@@ -25,7 +25,7 @@ async fn test_batch_producer() {
 
     let record = record(b"");
 
-    let partition_client = Arc::new(client.partition_client(&topic, 0).unwrap());
+    let partition_client = Arc::new(client.partition_client(&topic, 0).await.unwrap());
     let producer = BatchProducerBuilder::new(partition_client)
         .with_linger(Duration::from_secs(5))
         .build(RecordAggregator::new(record.approximate_size() * 2 + 1));