?

shotover-proxy/benches/windsock/kafka.rs

@@ -16,6 +16,8 @@ use shotover::transforms::TransformConfig;
 use std::sync::Arc;
 use std::{collections::HashMap, time::Duration};
 use test_helpers::docker_compose::docker_compose;
+use test_helpers::rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication};
+use test_helpers::rdkafka::client::DefaultClientContext;
 use test_helpers::rdkafka::config::ClientConfig;
 use test_helpers::rdkafka::consumer::{Consumer, StreamConsumer};
 use test_helpers::rdkafka::producer::{FutureProducer, FutureRecord};
@@ -347,6 +349,24 @@ impl Bench for KafkaBench {
  // only one string field so we just directly store the value in resources
  let broker_address = resources;
 
+ let admin: AdminClient<DefaultClientContext> = ClientConfig::new()
+ .set("bootstrap.servers", broker_address)
+ .create()
+ .unwrap();
+ admin
+ .create_topics(
+ &[NewTopic {
+ name: "topic_foo",
+ num_partitions: 1,
+ replication: TopicReplication::Fixed(1),
+ config: vec![],
+ }],
+ &AdminOptions::new()
+ .operation_timeout(Some(Timeout::After(Duration::from_secs(60)))),
+ )
+ .await
+ .unwrap();
+
  let producer: FutureProducer = ClientConfig::new()
  .set("bootstrap.servers", broker_address)
  .set("message.timeout.ms", "5000")
@@ -413,14 +433,10 @@ struct BenchTaskProducerKafka {
 #[async_trait]
 impl BenchTaskProducer for BenchTaskProducerKafka {
  async fn produce_one(&self) -> Result<(), String> {
- let key = rand::random::<[u8; 4]>();
+ // key is set to None which will result in round robin routing between all brokers
+ let record: FutureRecord<(), _> = FutureRecord::to("topic_foo").payload(&self.message);
  self.producer
- .send(
- FutureRecord::to("topic_foo")
- .payload(&self.message)
- .key(&key),
- Timeout::Never,
- )
+ .send(record, Timeout::Never)
  .await
  // Take just the error, ignoring the message contents because large messages result in unreadable noise in the logs.
  .map_err(|e| format!("{:?}", e.0))